Skip to content

Commit cf5983b

Browse files
tanchengtancheng
authored
Merge pull request #49 from coredac/introduce_grant_always
Introduce grant_always operation
2 parents 67ecd5f + 492c728 commit cf5983b

File tree

5 files changed

+139
-19
lines changed

5 files changed

+139
-19
lines changed

‎include/NeuraDialect/NeuraOps.td‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,3 +285,19 @@ def Neura_GrantOnceOp : Op<NeuraDialect, "grant_once"> {
285285

286286
// let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
287287
}
288+
289+
def Neura_GrantAlwaysOp : Op<NeuraDialect, "grant_always"> {
290+
let summary = "Marks a value as valid always.";
291+
let description = [{
292+
Grants a value always-valid predicate: the resulting value is considered valid
293+
during the entire application lifetime.
294+
295+
Example:
296+
%v = neura.grant_always %init : !neura.data<f32, i1> -> !neura.data<f32, i1>
297+
}];
298+
299+
let arguments = (ins AnyType:$value);
300+
let results = (outs AnyType:$result);
301+
302+
// let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
303+
}

‎lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp‎

Lines changed: 46 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,40 +14,73 @@ using namespace mlir;
1414

1515
// Inserts `grant_once` for every predicated value defined in the entry block
1616
// that is used outside of the block (i.e., a live-out).
17-
void insertGrantOnceInEntryBlock(Block *entry_block, OpBuilder &builder,
18-
DenseMap<Value, Value> &granted_once_map) {
19-
SmallVector<Value> live_out_values;
17+
void GrantPredicateInEntryBlock(Block *entry_block, OpBuilder &builder) {
18+
SmallVector<Value> live_out_arg_values;
19+
SmallVector<Value> live_out_non_arg_values;
2020

2121
// Step 1: Collects all live-out values first.
2222
for (Operation &op : *entry_block) {
2323
for (Value result : op.getResults()) {
2424
if (!isa<neura::PredicatedValue>(result.getType()))
2525
continue;
2626

27-
bool is_live_out = llvm::any_of(result.getUses(), [&](OpOperand &use) {
27+
bool used_in_branch = false;
28+
bool used_elsewhere = false;
29+
30+
for (OpOperand &use : result.getUses()) {
2831
Operation *user = use.getOwner();
29-
return user->getBlock() != entry_block || isa<neura::Br, neura::CondBr>(user);
30-
});
3132

32-
if (is_live_out && !granted_once_map.contains(result))
33-
live_out_values.push_back(result);
33+
// Case 1: Operand of a branch/cond_br → grant_once
34+
if (isa<neura::Br, neura::CondBr>(user)) {
35+
used_in_branch = true;
36+
}
37+
38+
// Case 2: Used directly in other blocks → grant_always
39+
if (user->getBlock() != entry_block) {
40+
used_elsewhere = true;
41+
}
42+
}
43+
44+
if (used_in_branch)
45+
live_out_arg_values.push_back(result);
46+
if (used_elsewhere)
47+
live_out_non_arg_values.push_back(result);
3448
}
3549
}
3650

3751
// Step 2: Inserts grant_once for each candidate.
38-
for (Value val : live_out_values) {
52+
// Inserts grant_once.
53+
for (Value val : live_out_arg_values) {
3954
Operation *def_op = val.getDefiningOp();
4055
if (!def_op)
4156
continue;
4257

4358
builder.setInsertionPointAfter(def_op);
4459
auto granted = builder.create<neura::GrantOnceOp>(def_op->getLoc(), val.getType(), val);
45-
granted_once_map[val] = granted.getResult();
4660

47-
// Replaces external uses with granted result.
61+
// Replaces uses in branch ops.
62+
for (OpOperand &use : llvm::make_early_inc_range(val.getUses())) {
63+
Operation *user = use.getOwner();
64+
if (isa<neura::Br, neura::CondBr>(user)) {
65+
use.set(granted.getResult());
66+
}
67+
}
68+
}
69+
70+
// Inserts grant_always.
71+
for (Value val : live_out_non_arg_values) {
72+
Operation *def_op = val.getDefiningOp();
73+
if (!def_op)
74+
continue;
75+
76+
builder.setInsertionPointAfter(def_op);
77+
auto granted = builder.create<neura::GrantAlwaysOp>(def_op->getLoc(), val.getType(), val);
78+
79+
// Replaces direct external uses (not in entry block, not in branch ops).
4880
for (OpOperand &use : llvm::make_early_inc_range(val.getUses())) {
4981
Operation *user = use.getOwner();
50-
if (user->getBlock() != entry_block || isa<neura::Br, neura::CondBr>(user)) {
82+
if (user->getBlock() != entry_block &&
83+
!isa<neura::Br, neura::CondBr>(user)) {
5184
use.set(granted.getResult());
5285
}
5386
}
@@ -283,8 +316,7 @@ struct TransformCtrlToDataFlowPass
283316
module.walk([&](func::FuncOp func) {
284317

285318
OpBuilder builder(func.getContext());
286-
DenseMap<Value, Value> granted_once_map;
287-
insertGrantOnceInEntryBlock(&func.getBody().front(), builder, granted_once_map);
319+
GrantPredicateInEntryBlock(&func.getBody().front(), builder);
288320

289321
// Get blocks in post-order
290322
SmallVector<Block *> postOrder;

‎test/neura/ctrl/branch_for.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,13 @@ func.func @loop_test() -> f32 {
5656

5757
// CTRL2DATA: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
5858
// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data<i64, i1>
59-
// CTRL2DATA-NEXT: %1 = "neura.grant_once"(%0) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
59+
// CTRL2DATA-NEXT: %1 = "neura.grant_always"(%0) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
6060
// CTRL2DATA-NEXT: %2 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
6161
// CTRL2DATA-NEXT: %3 = "neura.grant_once"(%2) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
6262
// CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
63-
// CTRL2DATA-NEXT: %5 = "neura.grant_once"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
63+
// CTRL2DATA-NEXT: %5 = "neura.grant_always"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
6464
// CTRL2DATA-NEXT: %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
65-
// CTRL2DATA-NEXT: %7 = "neura.grant_once"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
65+
// CTRL2DATA-NEXT: %7 = "neura.grant_always"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
6666
// CTRL2DATA-NEXT: %8 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
6767
// CTRL2DATA-NEXT: %9 = "neura.grant_once"(%8) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
6868
// CTRL2DATA-NEXT: %10 = neura.reserve : !neura.data<i64, i1>

‎test/neura/ctrl/branch_with_and_without_arg.mlir‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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 @test(%in: i64) -> f32 {
15+
%c0 = llvm.mlir.constant(0 : i64) : i64
16+
%c1 = llvm.mlir.constant(1.0 : f32) : f32
17+
%c2 = llvm.mlir.constant(2.0 : f32) : f32
18+
%c3 = llvm.mlir.constant(3.0 : f32) : f32
19+
%cond = llvm.icmp "eq" %in, %c0 : i64
20+
llvm.cond_br %cond, ^bb2(%c3 : f32), ^bb1(%c1, %c2 : f32, f32)
21+
22+
^bb1(%ca: f32, %cb: f32):
23+
%a = llvm.fadd %ca, %cb : f32
24+
llvm.br ^bb3(%a : f32)
25+
26+
^bb2(%cc: f32):
27+
%b = llvm.fmul %cc, %c2 : f32
28+
llvm.br ^bb3(%b : f32)
29+
30+
^bb3(%v: f32):
31+
return %v : f32
32+
}
33+
34+
// CHECK: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} {
35+
// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
36+
// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
37+
// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
38+
// CHECK-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
39+
// CHECK-NEXT: %4 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data<i64, i1>) -> !neura.data<i1, i1>
40+
// CHECK-NEXT: neura.cond_br %4 : !neura.data<i1, i1> then %3 : !neura.data<f32, i1> to ^bb2 else %1, %2 : !neura.data<f32, i1>, !neura.data<f32, i1> to ^bb1
41+
// CHECK-NEXT: ^bb1(%5: !neura.data<f32, i1>, %6: !neura.data<f32, i1>): // pred: ^bb0
42+
// CHECK-NEXT: %7 = "neura.fadd"(%5, %6) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
43+
// CHECK-NEXT: neura.br %7 : !neura.data<f32, i1> to ^bb3
44+
// CHECK-NEXT: ^bb2(%8: !neura.data<f32, i1>): // pred: ^bb0
45+
// CHECK-NEXT: %9 = "neura.fmul"(%8, %2) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
46+
// CHECK-NEXT: neura.br %9 : !neura.data<f32, i1> to ^bb3
47+
// CHECK-NEXT: ^bb3(%10: !neura.data<f32, i1>): // 2 preds: ^bb1, ^bb2
48+
// CHECK-NEXT: "neura.return"(%10) : (!neura.data<f32, i1>) -> ()
49+
// CHECK-NEXT: }
50+
51+
// CTRL2DATA: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} {
52+
// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
53+
// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
54+
// CTRL2DATA-NEXT: %2 = "neura.grant_once"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
55+
// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
56+
// CTRL2DATA-NEXT: %4 = "neura.grant_always"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
57+
// CTRL2DATA-NEXT: %5 = "neura.grant_once"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
58+
// CTRL2DATA-NEXT: %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
59+
// CTRL2DATA-NEXT: %7 = "neura.grant_once"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
60+
// CTRL2DATA-NEXT: %8 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data<i64, i1>) -> !neura.data<i1, i1>
61+
// CTRL2DATA-NEXT: %9 = "neura.grant_once"(%8) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
62+
// CTRL2DATA-NEXT: %10 = neura.grant_predicate %7, %9 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
63+
// CTRL2DATA-NEXT: %11 = neura.grant_predicate %4, %9 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
64+
// CTRL2DATA-NEXT: %12 = "neura.not"(%9) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
65+
// CTRL2DATA-NEXT: %13 = neura.grant_predicate %2, %12 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
66+
// CTRL2DATA-NEXT: %14 = "neura.not"(%9) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
67+
// CTRL2DATA-NEXT: %15 = neura.grant_predicate %5, %14 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
68+
// CTRL2DATA-NEXT: %16 = "neura.fadd"(%13, %15) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
69+
// CTRL2DATA-NEXT: %17 = "neura.fmul"(%10, %11) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
70+
// CTRL2DATA-NEXT: %18 = "neura.phi"(%16, %17) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
71+
// CTRL2DATA-NEXT: "neura.return"(%18) : (!neura.data<f32, i1>) -> ()
72+
// CTRL2DATA-NEXT: }

test/neura/ctrl/branch_no_arg.mlir renamed to test/neura/ctrl/branch_without_arg.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,9 +53,9 @@ func.func @test(%in: i64) -> f32 {
5353
// CTRL2DATA: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} {
5454
// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
5555
// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
56-
// CTRL2DATA-NEXT: %2 = "neura.grant_once"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
56+
// CTRL2DATA-NEXT: %2 = "neura.grant_always"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
5757
// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
58-
// CTRL2DATA-NEXT: %4 = "neura.grant_once"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
58+
// CTRL2DATA-NEXT: %4 = "neura.grant_always"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
5959
// CTRL2DATA-NEXT: %5 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
6060
// CTRL2DATA-NEXT: %6 = "neura.grant_once"(%5) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
6161
// CTRL2DATA-NEXT: %7 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data<f32, i1>

0 commit comments

Comments
 (0)