mlir中cse pass通过mlir::createcsepass()调用,作用于ssa形式的operation树,仅在同一基本块内识别无副作用的等价表达式并复用首次计算结果,不跨block、不跨函数、不跨dialect,且受fastmath、别名属性及动态操作限制。

MLIR里直接调用createCSEPass就行
MLIR本身不自己实现CSE逻辑,而是复用LLVM基础设施里的createCSEPass。这个Pass作用在SSA形式的Operation树上,自动识别同一基本块(Block)内语义等价、无副作用干扰的重复表达式,并把后续出现替换成第一次计算的结果变量。
常见使用方式是往PassManager里加一行:
pm.addPass(mlir::createCSEPass());
注意它只对当前作用域生效——比如加在func::FuncOp层级,就只优化该函数内部;加在Module层级,则会跨函数尝试(但实际受限于Symbol可见性,通常仍限于单个函数)。
容易踩的坑:
-
createCSEPass默认不跨Block做消除,循环体外和循环体内算两次,不算公共子表达式 - 如果两个计算之间有写内存、调用外部函数、或任何标记为
hasSideEffects的Op,CSE会主动跳过 - 浮点运算受
fastmath属性影响:没加arithmetic.fastmath时,addf x, y和addf y, x可能不被识别为等价(尽管数学上成立)
LLVM IR阶段用opt -enable-cse或-O2
LLVM IR层面的CSE由MachineCSE(针对机器码)和GVN(Global Value Numbering,更激进)共同覆盖,但最直接对应传统CSE的是-enable-cse选项。不过日常几乎不用手动开它,因为-O2及以上已默认包含。
验证是否生效,可以用:
opt -O2 -print-after=cse input.ll
观察输出中重复的%x = add i32 %a, %b是否只剩一个定义,其余全变成use %x。
关键限制:
- LLVM CSE只作用于Function内,不跨函数
- 对指针运算、数组索引这类带
getelementptr的表达式,需确保地址计算路径完全一致,否则不合并 - 调试信息(
!dbg)不影响CSE判断,但若某条指令带!noalias或!invariant.load,可能改变别名分析结果,间接影响CSE范围
Triton这类DSL如何触发CSE
Triton在make_ttir阶段通过passes.common.add_cse(pm)插入CSE Pass,本质就是调MLIR的createCSEPass。但它加在TTIR(Triton Tensor IR)层,比原始MLIR更贴近GPU张量语义。
你写的Triton kernel里出现多次a + b,只要满足:
- 都在同一个
for循环块内(或同级Block) - 中间没穿插
tl.store、tl.atomic_add等有副作用的操作 - 操作数类型、精度、fastmath标志一致
编译时就会被合并。比如这段kernel:
z = x + y w = x + y # 这行会被优化掉,w直接复用z
但如果写成:
z = x + y tl.store(ptr, z) w = x + y # 这里不会复用z,因为store引入了副作用
CSE不是万能的:哪些重复根本消不掉
有些“看起来重复”的计算,CSE天然无能为力:
- 跨基本块的重复,比如if分支里各算一次
a * b,CSE不合并(得靠GVN或LoopInvariantCodeMotion) - 不同Dialect的Op,比如
arith.addi和linalg.generic里封装的加法,CSE不认为等价 - 含非常量索引的tensor load:
load A[i]和load A[j]即使i == j运行时成立,编译期也无法证明,故不消除 - 涉及
tensor.extract_slice等带动态尺寸的Op,shape未折叠前,CSE无法判定等价性
真正要压榨性能,得配合canonicalizer先规一化表达式形态,再让CSE进场——顺序错了,CSE看到的就不是“重复”,而是“长得像的两坨不同东西”。











