当前位置：首页 > news >正文

Triton Linalg - WrapFuncBodyWithSingleBlockPass

news 2025/9/5 22:35:49

功能

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"];
}

Pass.cpp

WrapFuncBodyWithSingleBlock的cpp部分实现在：triton-linalg/triton-linalg/lib/Dialect/Triton/Transforms/WrapFuncBodyWithSingleBlock.cpp

runOnOperation

runOnOperation实现了对Module内每个实现了FunctionOpInterface的Operation（即Function）的遍历，并调用encapsulateMultiBlock。

struct WrapFuncBodyWithSingleBlockPass: public WrapFuncBodyWithSingleBlockBase<WrapFuncBodyWithSingleBlockPass> {void runOnOperation() override {getOperation()->walk([&](FunctionOpInterface func) { encapsulateMultiBlock(func); });}
};

encapsulateMultiBlock

创建Block

首先使用Function的Entry Block的参数，创新一个newBlock，插入Function的第一个Block之前，即变成新的Entry Block:

static void encapsulateMultiBlock(FunctionOpInterface funcOp) {...auto &entryBlock = body.front();auto blockArgTypes = entryBlock.getArgumentTypes();SmallVector<Location> blockArgLocs(blockArgTypes.size(), loc);Block *newBlock = builder.createBlock(&body, body.begin(), blockArgTypes, blockArgLocs);...
}

上述示例foo：

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):
^bb1(%0: i1, %1: i32, %2: i32):  // no predecessors"cf.cond_br"(%0)[^bb2, ^bb3] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()
^bb2:  // pred: ^bb1"tt.return"(%1) : (i32) -> ()
^bb3:  // pred: ^bb1"tt.return"(%2) : (i32) -> ()
}) : () -> ()

插入scf.execute_region

接着在newBlock（新的Entry Block）的Operation List头部插入scf.execute_region（返回值类型和Function一致）：

static void encapsulateMultiBlock(FunctionOpInterface funcOp) {...builder.setInsertionPointToStart(newBlock);FunctionType funcType = cast<FunctionType>(funcOp.getFunctionType());auto containerOp = builder.create<scf::ExecuteRegionOp>(loc, funcType.getResults());...
}

上述示例foo：

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):%0 = "scf.execute_region"() ({}) : () -> i32
^bb1(%1: i1, %2: i32, %3: i32):  // no predecessors"cf.cond_br"(%1)[^bb2, ^bb3] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()
^bb2:  // pred: ^bb1"tt.return"(%2) : (i32) -> ()
^bb3:  // pred: ^bb1"tt.return"(%3) : (i32) -> ()
}) : () -> ()

拷贝blocks

将原Function Body的Blocks追加到scf.execute_region的Region（跳过第一个，即newBlock）：

static void encapsulateMultiBlock(FunctionOpInterface funcOp) {...auto &containerRegion = containerOp.getRegion();auto &blocksToMove = body.getBlocks();containerRegion.getBlocks().splice(containerRegion.end(), blocksToMove,std::next(blocksToMove.begin()),blocksToMove.end());...
}

上述示例foo:

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):%0 = "scf.execute_region"() ({^bb0(%arg3: i1, %arg4: i32, %arg5: i32):"cf.cond_br"(%arg3)[^bb1, ^bb2] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()^bb1:  // pred: ^bb0"tt.return"(%arg4) : (i32) -> ()^bb2:  // pred: ^bb0"tt.return"(%arg5) : (i32) -> ()}) : () -> i32
}) : () -> ()

替换Block Arguments

删除scf.execute_region的Entry Block中的参数，并将后续Block使用到的地方替换为newBlock的参数（按照设计scf.execute_region不能有Block Arguments）：

  auto &containerEntryBlock = containerRegion.getBlocks().front();for (auto &arg : containerEntryBlock.getArguments()) {for (OpOperand &use : llvm::make_early_inc_range(arg.getUses())) {auto newValue = newBlock->getArgument(arg.getArgNumber());use.getOwner()->setOperand(use.getOperandNumber(), newValue);}}containerEntryBlock.eraseArguments(0, containerEntryBlock.getNumArguments());

上述示例foo：

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):%0 = "scf.execute_region"() ({"cf.cond_br"(%arg0)[^bb1, ^bb2] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()^bb1:  // pred: ^bb0"tt.return"(%arg1) : (i32) -> ()^bb2:  // pred: ^bb0"tt.return"(%arg2) : (i32) -> ()}) : () -> i32
}) : () -> ()

插入Yield Block

在scf.execute_region包含的block中追加一个yieldBlock（参数与包含的Block中的ReturnLike Op一致）：

  auto &yieldBlock = containerRegion.emplaceBlock();auto operandTypes = returnOps.front()->getOperandTypes();auto unknownLoc = UnknownLoc::get(ctx);SmallVector<Location> unknownLocs(operandTypes.size(), unknownLoc);yieldBlock.addArguments(operandTypes, unknownLocs);builder.setInsertionPointToEnd(&yieldBlock);builder.create<scf::YieldOp>(unknownLoc, yieldBlock.getArguments());

上述示例foo：

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):%0 = "scf.execute_region"() ({"cf.cond_br"(%arg0)[^bb1, ^bb2] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()^bb1:  // pred: ^bb0"tt.return"(%arg1) : (i32) -> ()^bb2:  // pred: ^bb0"tt.return"(%arg2) : (i32) -> ()^bb3(%1: i32):  // no predecessors"scf.yield"(%1) : (i32) -> ()}) : () -> i32
}) : () -> ()

插入branch

在sch.execute_region直接包含的ReturnLike Op后面插入一个branch语句，跳转到上面插入的yield Block(scf.execute_region的Bufferization只支持1个Yield)：

  SmallVector<Operation *> returnOps;containerOp.walk([&](Operation *op) {if (op->hasTrait<OpTrait::ReturnLike>() &&op->getParentOp() == containerOp.getOperation())returnOps.push_back(op);});for (auto *returnOp : returnOps) {builder.setInsertionPoint(returnOp);builder.create<cf::BranchOp>(returnOp->getLoc(), &yieldBlock,returnOp->getOperands());}

上述示例foo：

"tt.func"() <{function_type = (i1, i32, i32) -> i32, sym_name = "wrap_multi_block_triton_func"}> ({
^bb0(%arg0: i1, %arg1: i32, %arg2: i32):%0 = "scf.execute_region"() ({"cf.cond_br"(%arg0)[^bb1, ^bb2] <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (i1) -> ()^bb1:  // pred: ^bb0"cf.br"(%arg1)[^bb3] : (i32) -> ()"tt.return"(%arg1) : (i32) -> ()^bb2:  // pred: ^bb0"cf.br"(%arg2)[^bb3] : (i32) -> ()"tt.return"(%arg2) : (i32) -> ()^bb3(%1: i32):  // 2 preds: ^bb1, ^bb2"scf.yield"(%1) : (i32) -> ()}) : () -> i32
}) : () -> ()

删除ReturnLike Op

让子Block中的最后一个ReturnLike Op(任选一个)返回scf.execute_region的Result，并作为newBlock的terminator，删除其他的ReturnLike Op：

  Operation *terminator = returnOps.pop_back_val();for (auto *returnOp : returnOps)returnOp->erase();terminator->setOperands(containerOp.getResults());builder.setInsertionPointToEnd(&body.back());terminator->remove();builder.insert(terminator);

上述示例foo：

tt.func @wrap_multi_block_triton_func(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {%0 = scf.execute_region -> i32 {cf.cond_br %arg0, ^bb1, ^bb2^bb1:  // pred: ^bb0cf.br ^bb3(%arg1 : i32)^bb2:  // pred: ^bb0cf.br ^bb3(%arg2 : i32)^bb3(%1: i32):  // 2 preds: ^bb1, ^bb2scf.yield %1 : i32}tt.return %0 : i32
}

至此就完成了将原Function的所有Block封装到scf.execute_region的全部工作，新的Function只包含1个Block，即newBlock。