Triton Linalg - WrapFuncBodyWithSingleBlockPass

原创

已于 2025-09-03 14:45:14 修改 · 914 阅读

18 ·

CC 4.0 BY-SA版权

文章标签：

#ai编程 #mlir #开发语言 #人工智能

于 2025-09-03 11:36:11 首次发布

功能

WrapFuncBodyWithSingleBlockPass用于将函数体中有多个Block的Function，用scf.execute_region封起来，让Function变为单Block的形式。

示例：

转换前：

 tt.func @foo(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
     cf.cond_br %arg0, ^bb1, ^bb2
   ^bb1:
     tt.return %arg1: i32
   ^bb2:
     tt.return %arg2: i32
 }

转换后：

tt.func @foo(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
   %0 = scf.execute_region -> i32 {
       cf.cond_br %arg0, ^bb1, ^bb2
     ^bb1:
       cb.br ^bb3(%arg1 : i32)
     ^bb2:
       cb.br ^bb3(%arg2 : i32)
     ^bb3(%1 : i32):
       scf.yield %1 : i32
     }
     tt.return %0 : i32
   }

WrapFuncBodyWithSingleBlockPass简化了Function的控制流，解决了在scf.for 循环中内联函数调用时由于多块（multi-block）问题导致的内联问题。

如对foo调用的代码：

scf.for %iv = %0 to %10 step %1 {
  %result = call @foo(%arg0, %arg1, %iv) : (i1, i32, i32) -> i32
}

假如inline优化需要把foo做内联，则原始的多Block的foo会导致生成的IR不正确或难以优化。

实现

Pass.td

WrapFuncBodyWithSingleBlock的td部分实现在triton-linalg/triton-linalg/include/triton-linalg/Dialect/Triton/Transforms/Passes.td：

WrapFuncBodyWithSingleBlock是一个ModuleOp范围的Pass。
constructor为createWrapFuncBodyWithSingleBlockPass。

def WrapFuncBodyWithSingleBlock : Pass<"wrap-func-body-with-single-block", "::mlir::ModuleOp"> {
  let summary = "Wrap function body with single block";
  let description = [{
    This pass wraps function body into a block by moving body to a `scf.execute_region`.
  }];

  let constructor = "mlir::triton::createWrapFuncBodyWithSingleBlockPass()";

  let dependentDialects = ["triton::TritonDialect", "scf::SCFDialect"];
}