AI Compiler Drills/ 提示与讲解
T5 提示与讲解 · Track T · Triton(kernel compiler) · 对照 solutions/t5_filecheck.py 逐行读

T5 讲解:FileCheck 是怎样“对答案”的

拿 coalesce.mlir 风格的一段输出,一条一条 CHECK 地走一遍参考解:pos 指针怎样前进、NOT 什么时候被检查、LABEL 怎样把输入切块、[[row_layout]] 这种变量怎样把前后两条 CHECK 绑在一起;最后讲怎样在 Triton 仓库里读、改、写一个 lit 测试。

怎么读这一页

对照 drillst5_filecheck.py,以及测试文件里的 OUT(一个 @transpose 函数和一个 @other 函数)。

① 一次完整的匹配

检查串(来自 test_real_coalesce_style_checks_pass):

// CHECK: [[row:#.*]] = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16]
// CHECK: [[col:#.*]] = #ttg.blocked<{sizePerThread = [4, 1]
// CHECK-LABEL: @transpose
// CHECK: [[ptr:%.*]] = ttg.convert_layout {{.*}} -> tensor<64x64x!tt.ptr<f32>, [[row]]>
// CHECK-NEXT: [[val:%[0-9]+]] = tt.load [[ptr]]
// CHECK-SAME: [[row]]>
// CHECK-NOT: tt.dot
// CHECK: tt.store {{.*}}, [[col]]>
步 指令 在哪里找 结果
0 LABEL @transpose 先扫一遍全部 LABEL 块 A = 文件头 … @transpose 之前;块 B = @transpose 之后
1 CHECK row 块 A,从 0 开始 匹配第 1 行,row = "#blocked"
2 CHECK col 块 A,从第 1 行之后 第 2 行,col = "#blocked1"
3 CHECK ptr 块 B %8 = ttg.convert_layout ... #blocked>,ptr = "%8"
4 NEXT 下一行 %9 = tt.load %8 ...,val = "%9"
5 SAME [[row]]> 同一行、%8 之后 #blocked> ✔
6 NOT tt.dot 挂起 —
7 CHECK store 从第 4 行之后 tt.store %11, %10 : ... #blocked1>;此时检查 6 的区间:第 4 行末 … store 之前,没有 tt.dot ✔
pos 指针与挂起的 NOT悬停节点查看连接
类比 · 读书时用书签

pos 是一枚书签:每找到一条 CHECK,书签就夹到那里,下一条只能往后翻。NOT 像“这两枚书签之间不许出现某个词”——你得先放好下一枚书签,才知道要检查哪一段。LABEL 像章节:先把每一章的起点都找出来,每条 CHECK 只能在自己那一章里翻。

② compile_pattern:三种片段

drillscompile_pattern(46–74 行) 逐字符扫描:

片段 例子 生成的正则
字面 #ttg.blocked<{ \#ttg\.blocked<\{(re.escape)
{{…}} {{.*}} (?:.*)
捕获 [[row:#.*]] (?P<row>#.*)
引用 [[row]] re.escape(env["row"]) → \#blocked
空白 连续空格 \s+

一个容易忽略的细节:[[row:#.*]] 里的 .* 是贪婪的,但它后面紧跟着字面的 = #ttg.blocked<{,正则会回溯到正确的位置,row 最后只是 #blocked。

③ NEXT / SAME:行边界

drills_line_bounds(86–89 行) 给出某个位置所在行的 [start, end)。SAME 在“上一个匹配的末尾 … 行尾”里搜;NEXT 在“上一个匹配所在行的下一行”里搜。测试 CHECK: tt.func @transpose + CHECK-NEXT: tt.load 会失败:func 的下一行是 convert_layout。

④ DAG 与 COUNT

⑤ 变量的作用域

默认情况下变量在整个文件里有效;--enable-var-scope 时,普通变量在每个 LABEL 处清空,只有 $ 开头的全局变量保留。Triton 默认打开了它:tritontest/lit.cfg.py 里设置了 FILECHECK_OPTS=--enable-var-scope(“enabled by default in MLIR test”)。所以 accelerate-matmul.mlir 在文件头捕获 layout 时写成 [[$MMA:.+]](这个文件里有 72 处 [[$),否则进入第一个 CHECK-LABEL 块后变量就失效了。写新测试时忘了 $ 是很常见的失败原因;测试 test_var_scope 演示的就是这一点。

⑥ 在 Triton 仓库里:读、改、写

一个 lit 测试的生命周期悬停节点查看连接
lit -v build/cmake.*/test/TritonGPU/coalesce.mlir      # 一个文件
triton-opt test/TritonGPU/coalesce.mlir -split-input-file -tritongpu-coalesce | less   # 看原始输出

写新测试时:输入 IR 尽量小(T4 的 ddmin);只 CHECK 你的改动影响的东西;layout 用捕获变量而不是 #blocked3 这样的名字;// ----- 分隔的每一段都要有自己的 CHECK-LABEL。

⑦ 源码与资料索引

⑧ 自测

为什么 `CHECK-LABEL: @transpose` + `CHECK: arith.constant` + `CHECK-LABEL: @other` 会失败,而 arith.constant 明明在文件里?

LABEL 先把输入切成块:@transpose 块从它的位置一直到 @other 之前。arith.constant 只出现在 @other 之后,属于下一个块,第一个块里找不到。这正是 LABEL 的意义:防止一条 CHECK “偷看”到别的函数里的同名内容,让失败信息准确地指向出问题的函数。

CHECK 和 CHECK-NEXT 该选哪个?

默认用 CHECK:它容忍中间插入无关的行,测试更稳。只有“紧挨着”本身就是要检查的性质时才用 NEXT(比如两条指令必须相邻、没有 barrier 插在中间)。过多的 NEXT 会让任何无害的输出变化都打破测试,reviewer 常会要求放宽。