Introduce grant_always operation

include/NeuraDialect/NeuraOps.td

@@ -285,3 +285,19 @@ def Neura_GrantOnceOp : Op<NeuraDialect, "grant_once"> {
 
   // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
 }
+
+def Neura_GrantAlwaysOp : Op<NeuraDialect, "grant_always"> {
+  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<f32, i1> -> !neura.data<f32, i1>
+  }];
+
+  let arguments = (ins AnyType:$value);
+  let results = (outs AnyType:$result);
+
+  // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
+}

lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp

@@ -14,40 +14,73 @@ 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<Value, Value> &granted_once_map) {
-  SmallVector<Value> live_out_values;
+void GrantPredicateInEntryBlock(Block *entry_block, OpBuilder &builder) {
+  SmallVector<Value> live_out_arg_values;
+  SmallVector<Value> live_out_non_arg_values;
 
   // Step 1: Collects all live-out values first.
   for (Operation &op : *entry_block) {
     for (Value result : op.getResults()) {
       if (!isa<neura::PredicatedValue>(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<neura::Br, neura::CondBr>(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<neura::Br, neura::CondBr>(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<neura::GrantOnceOp>(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<neura::Br, neura::CondBr>(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<neura::GrantAlwaysOp>(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<neura::Br, neura::CondBr>(user)) {
+      if (user->getBlock() != entry_block &&
+          !isa<neura::Br, neura::CondBr>(user)) {
         use.set(granted.getResult());
       }
     }
@@ -283,8 +316,7 @@ struct TransformCtrlToDataFlowPass
     module.walk([&](func::FuncOp func) {
 
       OpBuilder builder(func.getContext());
-      DenseMap<Value, Value> granted_once_map;
-      insertGrantOnceInEntryBlock(&func.getBody().front(), builder, granted_once_map);
+      GrantPredicateInEntryBlock(&func.getBody().front(), builder);
 
       // Get blocks in post-order
       SmallVector<Block *> postOrder;

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<i64, i1>
-// CTRL2DATA-NEXT:   %1 = "neura.grant_once"(%0) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %1 = "neura.grant_always"(%0) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:   %2 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:   %3 = "neura.grant_once"(%2) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:   %4 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:   %5 = "neura.grant_once"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %5 = "neura.grant_always"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:   %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:   %7 = "neura.grant_once"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %7 = "neura.grant_always"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %8 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %9 = "neura.grant_once"(%8) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %10 = neura.reserve : !neura.data<i64, i1>

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<i64, i1>
+// CHECK-NEXT:   %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CHECK-NEXT:   %2 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CHECK-NEXT:   %3 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CHECK-NEXT:   %4 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// 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
+// CHECK-NEXT: ^bb1(%5: !neura.data<f32, i1>, %6: !neura.data<f32, i1>):  // pred: ^bb0
+// CHECK-NEXT:   %7 = "neura.fadd"(%5, %6) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CHECK-NEXT:   neura.br %7 : !neura.data<f32, i1> to ^bb3
+// CHECK-NEXT: ^bb2(%8: !neura.data<f32, i1>):  // pred: ^bb0
+// CHECK-NEXT:   %9 = "neura.fmul"(%8, %2) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CHECK-NEXT:   neura.br %9 : !neura.data<f32, i1> to ^bb3
+// CHECK-NEXT: ^bb3(%10: !neura.data<f32, i1>):  // 2 preds: ^bb1, ^bb2
+// CHECK-NEXT:   "neura.return"(%10) : (!neura.data<f32, i1>) -> ()
+// CHECK-NEXT: }
+
+// CTRL2DATA:      func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura"} {
+// CTRL2DATA-NEXT:   %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %2 = "neura.grant_once"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %4 = "neura.grant_always"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %5 = "neura.grant_once"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %6 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %7 = "neura.grant_once"(%6) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %8 = "neura.icmp"(%arg0, %0) <{cmpType = "eq"}> : (i64, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %9 = "neura.grant_once"(%8) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %10 = neura.grant_predicate %7, %9 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %11 = neura.grant_predicate %4, %9 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %12 = "neura.not"(%9) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %13 = neura.grant_predicate %2, %12 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %14 = "neura.not"(%9) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %15 = neura.grant_predicate %5, %14 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %16 = "neura.fadd"(%13, %15) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %17 = "neura.fmul"(%10, %11) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %18 = "neura.phi"(%16, %17) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   "neura.return"(%18) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT: }

test/neura/ctrl/branch_no_arg.mlir → 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<i64, i1>
 // CTRL2DATA-NEXT:   %1 = "neura.constant"() <{predicate = true, value = 1.000000e+00 : f32}> : () -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:   %2 = "neura.grant_once"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %2 = "neura.grant_always"(%1) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %3 = "neura.constant"() <{predicate = true, value = 2.000000e+00 : f32}> : () -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:   %4 = "neura.grant_once"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:   %4 = "neura.grant_always"(%3) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %5 = "neura.constant"() <{predicate = true, value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %6 = "neura.grant_once"(%5) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:   %7 = "neura.constant"() <{predicate = true, value = 4.000000e+00 : f32}> : () -> !neura.data<f32, i1>