diff --git a/include/NeuraDialect/NeuraOps.td b/include/NeuraDialect/NeuraOps.td index 223eee9f..852e02c1 100644 --- a/include/NeuraDialect/NeuraOps.td +++ b/include/NeuraDialect/NeuraOps.td @@ -285,3 +285,19 @@ def Neura_GrantOnceOp : Op { // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)"; } + +def Neura_GrantAlwaysOp : Op { + let summary = "Marks a value as valid always."; + let description = [{ + Grants a value always-valid predicate: the resulting value is considered valid + during the entire application lifetime. + + Example: + %v = neura.grant_always %init : !neura.data -> !neura.data + }]; + + let arguments = (ins AnyType:$value); + let results = (outs AnyType:$result); + + // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)"; +} \ No newline at end of file diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp index 367456bc..095aedec 100644 --- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp +++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp @@ -14,9 +14,9 @@ using namespace mlir; // Inserts `grant_once` for every predicated value defined in the entry block // that is used outside of the block (i.e., a live-out). -void insertGrantOnceInEntryBlock(Block *entry_block, OpBuilder &builder, - DenseMap &granted_once_map) { - SmallVector live_out_values; +void GrantPredicateInEntryBlock(Block *entry_block, OpBuilder &builder) { + SmallVector live_out_arg_values; + SmallVector live_out_non_arg_values; // Step 1: Collects all live-out values first. for (Operation &op : *entry_block) { @@ -24,30 +24,63 @@ void insertGrantOnceInEntryBlock(Block *entry_block, OpBuilder &builder, if (!isa(result.getType())) continue; - bool is_live_out = llvm::any_of(result.getUses(), [&](OpOperand &use) { + bool used_in_branch = false; + bool used_elsewhere = false; + + for (OpOperand &use : result.getUses()) { Operation *user = use.getOwner(); - return user->getBlock() != entry_block || isa(user); - }); - if (is_live_out && !granted_once_map.contains(result)) - live_out_values.push_back(result); + // Case 1: Operand of a branch/cond_br → grant_once + if (isa(user)) { + used_in_branch = true; + } + + // Case 2: Used directly in other blocks → grant_always + if (user->getBlock() != entry_block) { + used_elsewhere = true; + } + } + + if (used_in_branch) + live_out_arg_values.push_back(result); + if (used_elsewhere) + live_out_non_arg_values.push_back(result); } } // Step 2: Inserts grant_once for each candidate. - for (Value val : live_out_values) { + // Inserts grant_once. + for (Value val : live_out_arg_values) { Operation *def_op = val.getDefiningOp(); if (!def_op) continue; builder.setInsertionPointAfter(def_op); auto granted = builder.create(def_op->getLoc(), val.getType(), val); - granted_once_map[val] = granted.getResult(); - // Replaces external uses with granted result. + // Replaces uses in branch ops. + for (OpOperand &use : llvm::make_early_inc_range(val.getUses())) { + Operation *user = use.getOwner(); + if (isa(user)) { + use.set(granted.getResult()); + } + } + } + + // Inserts grant_always. + for (Value val : live_out_non_arg_values) { + Operation *def_op = val.getDefiningOp(); + if (!def_op) + continue; + + builder.setInsertionPointAfter(def_op); + auto granted = builder.create(def_op->getLoc(), val.getType(), val); + + // Replaces direct external uses (not in entry block, not in branch ops). for (OpOperand &use : llvm::make_early_inc_range(val.getUses())) { Operation *user = use.getOwner(); - if (user->getBlock() != entry_block || isa(user)) { + if (user->getBlock() != entry_block && + !isa(user)) { use.set(granted.getResult()); } } @@ -283,8 +316,7 @@ struct TransformCtrlToDataFlowPass module.walk([&](func::FuncOp func) { OpBuilder builder(func.getContext()); - DenseMap granted_once_map; - insertGrantOnceInEntryBlock(&func.getBody().front(), builder, granted_once_map); + GrantPredicateInEntryBlock(&func.getBody().front(), builder); // Get blocks in post-order SmallVector postOrder; diff --git a/test/neura/ctrl/branch_for.mlir b/test/neura/ctrl/branch_for.mlir index eac44582..024c6012 100644 --- a/test/neura/ctrl/branch_for.mlir +++ b/test/neura/ctrl/branch_for.mlir @@ -56,13 +56,13 @@ func.func @loop_test() -> f32 { // CTRL2DATA: func.func @loop_test() -> f32 attributes {accelerator = "neura"} { // CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 10 : i64}> : () -> !neura.data -// CTRL2DATA-NEXT: %1 = "neura.grant_once"(%0) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %1 = "neura.grant_always"(%0) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %2 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data // CTRL2DATA-NEXT: %3 = "neura.grant_once"(%2) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %4 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data -// CTRL2DATA-NEXT: %5 = "neura.grant_once"(%4) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %5 = "neura.grant_always"(%4) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data -// CTRL2DATA-NEXT: %7 = "neura.grant_once"(%6) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %7 = "neura.grant_always"(%6) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %8 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data // CTRL2DATA-NEXT: %9 = "neura.grant_once"(%8) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %10 = neura.reserve : !neura.data diff --git a/test/neura/ctrl/branch_with_and_without_arg.mlir b/test/neura/ctrl/branch_with_and_without_arg.mlir new file mode 100644 index 00000000..04bc1751 --- /dev/null +++ b/test/neura/ctrl/branch_with_and_without_arg.mlir @@ -0,0 +1,72 @@ +// 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 @test(%in: i64) -> f32 { + %c0 = llvm.mlir.constant(0 : i64) : i64 + %c1 = llvm.mlir.constant(1.0 : f32) : f32 + %c2 = llvm.mlir.constant(2.0 : f32) : f32 + %c3 = llvm.mlir.constant(3.0 : f32) : f32 + %cond = llvm.icmp "eq" %in, %c0 : i64 + llvm.cond_br %cond, ^bb2(%c3 : f32), ^bb1(%c1, %c2 : f32, f32) + +^bb1(%ca: f32, %cb: f32): + %a = llvm.fadd %ca, %cb : f32 + llvm.br ^bb3(%a : f32) + +^bb2(%cc: f32): + %b = llvm.fmul %cc, %c2 : f32 + llvm.br ^bb3(%b : f32) + +^bb3(%v: f32): + return %v : f32 +} + +// CHECK: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} { +// CHECK-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data +// CHECK-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %2 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data +// CHECK-NEXT: %4 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data) -> !neura.data +// CHECK-NEXT: neura.cond_br %4 : !neura.data then %3 : !neura.data to ^bb2 else %1, %2 : !neura.data, !neura.data to ^bb1 +// CHECK-NEXT: ^bb1(%5: !neura.data, %6: !neura.data): // pred: ^bb0 +// CHECK-NEXT: %7 = "neura.fadd"(%5, %6) : (!neura.data, !neura.data) -> !neura.data +// CHECK-NEXT: neura.br %7 : !neura.data to ^bb3 +// CHECK-NEXT: ^bb2(%8: !neura.data): // pred: ^bb0 +// CHECK-NEXT: %9 = "neura.fmul"(%8, %2) : (!neura.data, !neura.data) -> !neura.data +// CHECK-NEXT: neura.br %9 : !neura.data to ^bb3 +// CHECK-NEXT: ^bb3(%10: !neura.data): // 2 preds: ^bb1, ^bb2 +// CHECK-NEXT: "neura.return"(%10) : (!neura.data) -> () +// CHECK-NEXT: } + +// CTRL2DATA: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} { +// CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data +// CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %2 = "neura.grant_once"(%1) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %4 = "neura.grant_always"(%3) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %5 = "neura.grant_once"(%3) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data +// CTRL2DATA-NEXT: %7 = "neura.grant_once"(%6) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %8 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %9 = "neura.grant_once"(%8) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %10 = neura.grant_predicate %7, %9 : !neura.data, !neura.data -> !neura.data +// CTRL2DATA-NEXT: %11 = neura.grant_predicate %4, %9 : !neura.data, !neura.data -> !neura.data +// CTRL2DATA-NEXT: %12 = "neura.not"(%9) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %13 = neura.grant_predicate %2, %12 : !neura.data, !neura.data -> !neura.data +// CTRL2DATA-NEXT: %14 = "neura.not"(%9) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %15 = neura.grant_predicate %5, %14 : !neura.data, !neura.data -> !neura.data +// CTRL2DATA-NEXT: %16 = "neura.fadd"(%13, %15) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %17 = "neura.fmul"(%10, %11) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: %18 = "neura.phi"(%16, %17) : (!neura.data, !neura.data) -> !neura.data +// CTRL2DATA-NEXT: "neura.return"(%18) : (!neura.data) -> () +// CTRL2DATA-NEXT: } diff --git a/test/neura/ctrl/branch_no_arg.mlir b/test/neura/ctrl/branch_without_arg.mlir similarity index 95% rename from test/neura/ctrl/branch_no_arg.mlir rename to test/neura/ctrl/branch_without_arg.mlir index 5aa9808f..385fe20f 100644 --- a/test/neura/ctrl/branch_no_arg.mlir +++ b/test/neura/ctrl/branch_without_arg.mlir @@ -53,9 +53,9 @@ func.func @test(%in: i64) -> f32 { // CTRL2DATA: func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} { // CTRL2DATA-NEXT: %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data // CTRL2DATA-NEXT: %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data -// CTRL2DATA-NEXT: %2 = "neura.grant_once"(%1) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %2 = "neura.grant_always"(%1) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data -// CTRL2DATA-NEXT: %4 = "neura.grant_once"(%3) : (!neura.data) -> !neura.data +// CTRL2DATA-NEXT: %4 = "neura.grant_always"(%3) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %5 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data // CTRL2DATA-NEXT: %6 = "neura.grant_once"(%5) : (!neura.data) -> !neura.data // CTRL2DATA-NEXT: %7 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data