Skip to content

Enable loop ctrl with placeholder for SSA and backward ctrl move #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tancheng merged 2 commits into from
Jun 16, 2025
Merged

Enable loop ctrl with placeholder for SSA and backward ctrl move #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca6e5df
Enable loop ctrl with placeholder for SSA and backward ctrl move
tancheng Jun 15, 2025
7287abe
[test] Provide for loop test
tancheng Jun 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions include/NeuraDialect/NeuraOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@ def Neura_ConstantOp : Op<NeuraDialect, "constant"> {
def Neura_AddOp : Op<NeuraDialect, "add"> {
let summary = "Integer addition operation";
let opName = "add";
let arguments = (ins AnyInteger:$lhs, AnyInteger:$rhs, Optional<AnyType>:$predicate);
let results = (outs AnyInteger:$result);
let arguments = (ins AnyType:$lhs, AnyType:$rhs, Optional<AnyType>:$predicate);
let results = (outs AnyType:$result);
// let assemblyFormat = "$lhs `,` $rhs `,` $predicate attr-dict `:` type($result)";
let traits = [SameOperandsAndResultElementType];
}
Expand Down
45 changes: 33 additions & 12 deletions lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ void getBlocksInPostOrder(Block *startBlock, SmallVectorImpl<Block *> &postOrder

// Creates phi nodes for all live-in values in the given block.
void createPhiNodesForBlock(Block *block, OpBuilder &builder,
DenseMap<Value, Value> &value_map) {
DenseMap<Value, Value> &value_map,
SmallVectorImpl<std::tuple<Value, Value, Block *>> &deferred_ctrl_movs) {
if (block->hasNoPredecessors()) {
// Skips phi insertion for entry block.
return;
Expand All @@ -41,17 +42,9 @@ void createPhiNodesForBlock(Block *block, OpBuilder &builder,
// Identifies operands defined in other blocks.
if (operand.getDefiningOp() &&
operand.getDefiningOp()->getBlock() != block) {
// Checks if the live-in is a block argument. SSA form forces this rule.
bool found_in_block_argument = false;
for (BlockArgument arg : block->getArguments()) {
if (arg == operand) {
found_in_block_argument = true;
break;
}
}
live_ins.push_back(operand);
continue;
}

// Collects all block arguments.
if (auto blockArg = llvm::dyn_cast<BlockArgument>(operand)) {
live_ins.push_back(operand);
Expand Down Expand Up @@ -107,7 +100,19 @@ void createPhiNodesForBlock(Block *block, OpBuilder &builder,
llvm::errs() << "Unknown branch terminator in block: " << *pred << "\n";
continue;
}
phi_operands.push_back(incoming);

// If the incoming value is defined in the same block, inserts a `neura.reserve`
// and defer a backward ctrl move.
if (incoming.getDefiningOp() && incoming.getDefiningOp()->getBlock() == block) {
builder.setInsertionPointToStart(block);
auto placeholder = builder.create<neura::ReserveOp>(loc, incoming.getType());
phi_operands.push_back(placeholder.getResult());
// Defers the backward ctrl move operation to be inserted after all phi operands
// are defined. Inserted: (real_defined_value, just_created_reserve, current_block).
deferred_ctrl_movs.emplace_back(incoming, placeholder.getResult(), block);
} else {
phi_operands.push_back(incoming);
}
}
}

Expand Down Expand Up @@ -181,6 +186,13 @@ struct TransformCtrlToDataFlowPass
void runOnOperation() override {
ModuleOp module = getOperation();

// Declares a vector to hold deferred backward ctrl move operations.
// This is useful when a live-in value is defined within the same block.
// The tuple contains:
// - real value (the one that is defined in the same block, after the placeholder)
// - placeholder value (the one that will be used in the phi node)
// - block where the backward ctrl move should be inserted
SmallVector<std::tuple<Value, Value, Block *>, 4> deferred_ctrl_movs;
module.walk([&](func::FuncOp func) {
// Get blocks in post-order
SmallVector<Block *> postOrder;
Expand All @@ -194,7 +206,7 @@ struct TransformCtrlToDataFlowPass
// Process blocks bottom-up
for (Block *block : postOrder) {
// Creates phi nodes for live-ins.
createPhiNodesForBlock(block, builder, value_map);
createPhiNodesForBlock(block, builder, value_map, deferred_ctrl_movs);
}

// Flattens blocks into the entry block.
Expand Down Expand Up @@ -234,6 +246,15 @@ struct TransformCtrlToDataFlowPass
block->erase();
}
});

// Inserts the deferred backward ctrl move operations after phi operands
// are defined.
for (auto &[realVal, placeholder, block] : deferred_ctrl_movs) {
Operation *defOp = realVal.getDefiningOp();
assert(defOp && "Backward ctrl move's source must be an op result");
OpBuilder movBuilder(defOp->getBlock(), ++Block::iterator(defOp));
movBuilder.create<neura::CtrlMovOp>(defOp->getLoc(), realVal, placeholder);
}
}
};
} // namespace
Expand Down
65 changes: 65 additions & 0 deletions test/neura/ctrl/branch_for.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
// RUN: mlir-neura-opt %s \
// RUN: --assign-accelerator \
// RUN: --lower-llvm-to-neura \
// RUN: --leverage-predicated-value \
// RUN: | FileCheck %s

// RUN: mlir-neura-opt %s \
// RUN: --assign-accelerator \
// RUN: --lower-llvm-to-neura \
// RUN: --leverage-predicated-value \
// RUN: --transform-ctrl-to-data-flow \
// RUN: | FileCheck %s -check-prefix=CTRL2DATA

func.func @loop_test() -> f32 {
%n = llvm.mlir.constant(10 : i64) : i64
%c0 = llvm.mlir.constant(0 : i64) : i64
%c1 = llvm.mlir.constant(1 : i64) : i64
%c1f = llvm.mlir.constant(3.0 : f32) : f32
%acc_init = llvm.mlir.constant(0.0 : f32) : f32

llvm.br ^bb1(%c0, %acc_init : i64, f32)

^bb1(%i: i64, %acc: f32): // loop body + check + increment
%next_acc = llvm.fadd %acc, %c1f : f32
%i_next = llvm.add %i, %c1 : i64
%cmp = llvm.icmp "slt" %i_next, %n : i64
llvm.cond_br %cmp, ^bb1(%i_next, %next_acc : i64, f32), ^exit(%next_acc : f32)

^exit(%result: f32):
return %result : f32
}

// CHECK: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data<i64, i1>
// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
// CHECK-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
// CHECK-NEXT: %4 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
// CHECK-NEXT: neura.br %1, %4 : !neura.data<i64, i1>, !neura.data<f32, i1> to ^bb1
// CHECK-NEXT: ^bb1(%5: !neura.data<i64, i1>, %6: !neura.data<f32, i1>): // 2 preds: ^bb0, ^bb1
// CHECK-NEXT: %7 = "neura.fadd"(%6, %3) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
// CHECK-NEXT: %8 = "neura.add"(%5, %2) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
// CHECK-NEXT: %9 = "neura.icmp"(%8, %0) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
// CHECK-NEXT: neura.cond_br %9 : !neura.data<i1, i1> then %8, %7 : !neura.data<i64, i1>, !neura.data<f32, i1> to ^bb1 else %7 : !neura.data<f32, i1> to ^bb2
// CHECK-NEXT: ^bb2(%10: !neura.data<f32, i1>): // pred: ^bb1
// CHECK-NEXT: "neura.return"(%10) : (!neura.data<f32, i1>) -> ()
// CHECK-NEXT: }

// CTRL2DATA: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: %2 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
// CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
// CTRL2DATA-NEXT: %5 = neura.reserve : !neura.data<i64, i1>
// CTRL2DATA-NEXT: %6 = "neura.phi"(%1, %5) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: %7 = neura.reserve : !neura.data<f32, i1>
// CTRL2DATA-NEXT: %8 = "neura.phi"(%4, %7) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
// CTRL2DATA-NEXT: %9 = "neura.fadd"(%8, %3) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
// CTRL2DATA-NEXT: neura.ctrl_mov %9 -> %7 : !neura.data<f32, i1> !neura.data<f32, i1>
// CTRL2DATA-NEXT: %10 = "neura.add"(%6, %2) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
// CTRL2DATA-NEXT: neura.ctrl_mov %10 -> %5 : !neura.data<i64, i1> !neura.data<i64, i1>
// CTRL2DATA-NEXT: %11 = "neura.icmp"(%10, %0) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
// CTRL2DATA-NEXT: "neura.return"(%9) : (!neura.data<f32, i1>) -> ()
// CTRL2DATA-NEXT: }