Skip to content

Commit a80ea51

Browse files
tanchengtancheng
authored
Merge pull request #36 from coredac/ctrl_to_data_for_loop
Enable loop ctrl with placeholder for SSA and backward ctrl move
2 parents 44d6cae + 3f02ac3 commit a80ea51

File tree

3 files changed

+100
-14
lines changed

3 files changed

+100
-14
lines changed

‎include/NeuraDialect/NeuraOps.td‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ def Neura_ConstantOp : Op<NeuraDialect, "constant"> {
1818
def Neura_AddOp : Op<NeuraDialect, "add"> {
1919
let summary = "Integer addition operation";
2020
let opName = "add";
21-
let arguments = (ins AnyInteger:$lhs, AnyInteger:$rhs, Optional<AnyType>:$predicate);
22-
let results = (outs AnyInteger:$result);
21+
let arguments = (ins AnyType:$lhs, AnyType:$rhs, Optional<AnyType>:$predicate);
22+
let results = (outs AnyType:$result);
2323
// let assemblyFormat = "$lhs `,` $rhs `,` $predicate attr-dict `:` type($result)";
2424
let traits = [SameOperandsAndResultElementType];
2525
}

‎lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp‎

Lines changed: 33 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ void getBlocksInPostOrder(Block *startBlock, SmallVectorImpl<Block *> &postOrder
2828

2929
// Creates phi nodes for all live-in values in the given block.
3030
void createPhiNodesForBlock(Block *block, OpBuilder &builder,
31-
DenseMap<Value, Value> &value_map) {
31+
DenseMap<Value, Value> &value_map,
32+
SmallVectorImpl<std::tuple<Value, Value, Block *>> &deferred_ctrl_movs) {
3233
if (block->hasNoPredecessors()) {
3334
// Skips phi insertion for entry block.
3435
return;
@@ -41,17 +42,9 @@ void createPhiNodesForBlock(Block *block, OpBuilder &builder,
4142
// Identifies operands defined in other blocks.
4243
if (operand.getDefiningOp() &&
4344
operand.getDefiningOp()->getBlock() != block) {
44-
// Checks if the live-in is a block argument. SSA form forces this rule.
45-
bool found_in_block_argument = false;
46-
for (BlockArgument arg : block->getArguments()) {
47-
if (arg == operand) {
48-
found_in_block_argument = true;
49-
break;
50-
}
51-
}
5245
live_ins.push_back(operand);
46+
continue;
5347
}
54-
5548
// Collects all block arguments.
5649
if (auto blockArg = llvm::dyn_cast<BlockArgument>(operand)) {
5750
live_ins.push_back(operand);
@@ -107,7 +100,19 @@ void createPhiNodesForBlock(Block *block, OpBuilder &builder,
107100
llvm::errs() << "Unknown branch terminator in block: " << *pred << "\n";
108101
continue;
109102
}
110-
phi_operands.push_back(incoming);
103+
104+
// If the incoming value is defined in the same block, inserts a `neura.reserve`
105+
// and defer a backward ctrl move.
106+
if (incoming.getDefiningOp() && incoming.getDefiningOp()->getBlock() == block) {
107+
builder.setInsertionPointToStart(block);
108+
auto placeholder = builder.create<neura::ReserveOp>(loc, incoming.getType());
109+
phi_operands.push_back(placeholder.getResult());
110+
// Defers the backward ctrl move operation to be inserted after all phi operands
111+
// are defined. Inserted: (real_defined_value, just_created_reserve, current_block).
112+
deferred_ctrl_movs.emplace_back(incoming, placeholder.getResult(), block);
113+
} else {
114+
phi_operands.push_back(incoming);
115+
}
111116
}
112117
}
113118

@@ -181,6 +186,13 @@ struct TransformCtrlToDataFlowPass
181186
void runOnOperation() override {
182187
ModuleOp module = getOperation();
183188

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

200212
// Flattens blocks into the entry block.
@@ -234,6 +246,15 @@ struct TransformCtrlToDataFlowPass
234246
block->erase();
235247
}
236248
});
249+
250+
// Inserts the deferred backward ctrl move operations after phi operands
251+
// are defined.
252+
for (auto &[realVal, placeholder, block] : deferred_ctrl_movs) {
253+
Operation *defOp = realVal.getDefiningOp();
254+
assert(defOp && "Backward ctrl move's source must be an op result");
255+
OpBuilder movBuilder(defOp->getBlock(), ++Block::iterator(defOp));
256+
movBuilder.create<neura::CtrlMovOp>(defOp->getLoc(), realVal, placeholder);
257+
}
237258
}
238259
};
239260
} // namespace

‎test/neura/ctrl/branch_for.mlir‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// RUN: mlir-neura-opt %s \
2+
// RUN: --assign-accelerator \
3+
// RUN: --lower-llvm-to-neura \
4+
// RUN: --leverage-predicated-value \
5+
// RUN: | FileCheck %s
6+
7+
// RUN: mlir-neura-opt %s \
8+
// RUN: --assign-accelerator \
9+
// RUN: --lower-llvm-to-neura \
10+
// RUN: --leverage-predicated-value \
11+
// RUN: --transform-ctrl-to-data-flow \
12+
// RUN: | FileCheck %s -check-prefix=CTRL2DATA
13+
14+
func.func @loop_test() -> f32 {
15+
%n = llvm.mlir.constant(10 : i64) : i64
16+
%c0 = llvm.mlir.constant(0 : i64) : i64
17+
%c1 = llvm.mlir.constant(1 : i64) : i64
18+
%c1f = llvm.mlir.constant(3.0 : f32) : f32
19+
%acc_init = llvm.mlir.constant(0.0 : f32) : f32
20+
21+
llvm.br ^bb1(%c0, %acc_init : i64, f32)
22+
23+
^bb1(%i: i64, %acc: f32): // loop body + check + increment
24+
%next_acc = llvm.fadd %acc, %c1f : f32
25+
%i_next = llvm.add %i, %c1 : i64
26+
%cmp = llvm.icmp "slt" %i_next, %n : i64
27+
llvm.cond_br %cmp, ^bb1(%i_next, %next_acc : i64, f32), ^exit(%next_acc : f32)
28+
29+
^exit(%result: f32):
30+
return %result : f32
31+
}
32+
33+
// CHECK: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
34+
// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data<i64, i1>
35+
// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
36+
// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
37+
// CHECK-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
38+
// CHECK-NEXT: %4 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
39+
// CHECK-NEXT: neura.br %1, %4 : !neura.data<i64, i1>, !neura.data<f32, i1> to ^bb1
40+
// CHECK-NEXT: ^bb1(%5: !neura.data<i64, i1>, %6: !neura.data<f32, i1>): // 2 preds: ^bb0, ^bb1
41+
// CHECK-NEXT: %7 = "neura.fadd"(%6, %3) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
42+
// CHECK-NEXT: %8 = "neura.add"(%5, %2) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
43+
// CHECK-NEXT: %9 = "neura.icmp"(%8, %0) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
44+
// 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
45+
// CHECK-NEXT: ^bb2(%10: !neura.data<f32, i1>): // pred: ^bb1
46+
// CHECK-NEXT: "neura.return"(%10) : (!neura.data<f32, i1>) -> ()
47+
// CHECK-NEXT: }
48+
49+
// CTRL2DATA: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
50+
// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data<i64, i1>
51+
// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
52+
// CTRL2DATA-NEXT: %2 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
53+
// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
54+
// CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
55+
// CTRL2DATA-NEXT: %5 = neura.reserve : !neura.data<i64, i1>
56+
// CTRL2DATA-NEXT: %6 = "neura.phi"(%1, %5) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
57+
// CTRL2DATA-NEXT: %7 = neura.reserve : !neura.data<f32, i1>
58+
// CTRL2DATA-NEXT: %8 = "neura.phi"(%4, %7) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
59+
// CTRL2DATA-NEXT: %9 = "neura.fadd"(%8, %3) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
60+
// CTRL2DATA-NEXT: neura.ctrl_mov %9 -> %7 : !neura.data<f32, i1> !neura.data<f32, i1>
61+
// CTRL2DATA-NEXT: %10 = "neura.add"(%6, %2) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
62+
// CTRL2DATA-NEXT: neura.ctrl_mov %10 -> %5 : !neura.data<i64, i1> !neura.data<i64, i1>
63+
// CTRL2DATA-NEXT: %11 = "neura.icmp"(%10, %0) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
64+
// CTRL2DATA-NEXT: "neura.return"(%9) : (!neura.data<f32, i1>) -> ()
65+
// CTRL2DATA-NEXT: }

0 commit comments

Comments
 (0)