From 23d136b4e79e5c1e6d816dd8ae4da154105fd118 Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Sun, 28 Dec 2025 15:48:37 +0800
Subject: [PATCH 1/7] add return void op & transform pass

---
 include/NeuraDialect/NeuraOps.td              |  18 ++
 include/NeuraDialect/NeuraPasses.h            |   1 +
 include/NeuraDialect/NeuraPasses.td           |  12 +
 lib/NeuraDialect/Transforms/CMakeLists.txt    |   1 +
 .../Transforms/CanonicalizeReturnPass.cpp     | 205 ++++++++++++++++++
 .../TransformCtrlToDataFlowPass.cpp           |   2 +
 6 files changed, 239 insertions(+)
 create mode 100644 lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp

diff --git a/include/NeuraDialect/NeuraOps.td b/include/NeuraDialect/NeuraOps.td
index f51d97ff..e0f8b012 100644
--- a/include/NeuraDialect/NeuraOps.td
+++ b/include/NeuraDialect/NeuraOps.td
@@ -375,6 +375,24 @@ def Neura_ReturnOp : Op<NeuraDialect, "return", [Terminator]> {
   // let assemblyFormat = "($values^)? `,` attr-dict";
 }
 
+// Defines a return operation for void functions.
+def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", [Terminator]>{
+  let summary = "Return operation for void functions with execution trigger.";
+  let description = [{
+    Terminates a void function. The trigger operand provides the execution
+    condition in dataflow mode - the return executes when the trigger becomes valid.
+    
+    Without a trigger (empty operands), this represents a static return that
+    always executes. After canonicalization, void returns should have triggers.
+
+    Example:
+      neura.return_void %trigger : i1       // Triggered return in ctrl mode
+      neura.return_void %cond : !neura.data<i1, i1>  // In dataflow mode
+  }];
+  let arguments = (ins Optional<AnyType>:$trigger);
+  let assemblyFormat = "($trigger^ `:` type($trigger))? attr-dict";
+}
+
 // Defines a cast operation for type conversion.
 def Neura_CastOp : Op<NeuraDialect, "cast">{
   let summary = "Generic type conversion operation";
diff --git a/include/NeuraDialect/NeuraPasses.h b/include/NeuraDialect/NeuraPasses.h
index 4e6e38ec..72e97d72 100644
--- a/include/NeuraDialect/NeuraPasses.h
+++ b/include/NeuraDialect/NeuraPasses.h
@@ -25,6 +25,7 @@ std::unique_ptr<mlir::Pass> createTransformCtrlToDataFlowPass();
 std::unique_ptr<mlir::Pass> createLeveragePredicatedValuePass();
 std::unique_ptr<mlir::Pass> createMapToAcceleratorPass();
 std::unique_ptr<mlir::Pass> createGenerateCodePass();
+std::unique_ptr<mlir::Pass> createCanonicalizeReturnPass();
 std::unique_ptr<mlir::Pass> createCanonicalizeLiveInPass();
 std::unique_ptr<mlir::Pass> createPromoteFuncArgToConstPass();
 std::unique_ptr<mlir::Pass> createTransformToSteerControlPass();
diff --git a/include/NeuraDialect/NeuraPasses.td b/include/NeuraDialect/NeuraPasses.td
index 35888719..de1dcc67 100644
--- a/include/NeuraDialect/NeuraPasses.td
+++ b/include/NeuraDialect/NeuraPasses.td
@@ -75,6 +75,18 @@ def FuseLoopControl: Pass<"fuse-loop-control", "ModuleOp">{
   let constructor = "neura::createFuseLoopControlPass()";
 }
 
+def CanonicalizeReturn : Pass<"canonicalize-return", "ModuleOp">{
+  let summary = "Adds execution conditions to void return operations";
+  let description = [{
+    This pass processes void functions and adds trigger conditions to their
+    return operations. For void functions, the return operation needs an
+    execution condition to properly trigger in dataflow mode.
+    
+    This pass should run before CanonicalizeLiveIn and TransformCtrlToDataFlow.
+  }];
+  let constructor = "neura::createCanonicalizeReturnPass()";
+}
+
 def CanonicalizeLiveIn : Pass<"canonicalize-live-in", "ModuleOp"> {
   let summary = "Canonicalizes live-in values/operations in each basic block.";
   let description = [{
diff --git a/lib/NeuraDialect/Transforms/CMakeLists.txt b/lib/NeuraDialect/Transforms/CMakeLists.txt
index 83a4f847..e043bcce 100644
--- a/lib/NeuraDialect/Transforms/CMakeLists.txt
+++ b/lib/NeuraDialect/Transforms/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_library(
     LeveragePredicatedValuePass.cpp
     MapToAcceleratorPass.cpp
     GenerateCodePass.cpp
+    CanonicalizeReturnPass.cpp
     CanonicalizeLiveInPass.cpp
     CanonicalizeCastPass.cpp
     PromoteFuncArgToConstPass.cpp
diff --git a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
new file mode 100644
index 00000000..528653ff
--- /dev/null
+++ b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
@@ -0,0 +1,205 @@
+#include "NeuraDialect/NeuraDialect.h"
+#include "NeuraDialect/NeuraOps.h"
+#include "NeuraDialect/NeuraPasses.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+#include <memory>
+
+using namespace mlir;
+
+#define GEN_PASS_DEF_CANONICALIZERETURN
+#include "NeuraDialect/NeuraPasses.h.inc"
+
+namespace {
+// Checks if a function is a void function (i.e., has no return values).
+static bool isVoidFunction(func::FuncOp func_op) {
+  if (func_op.getNumResults() == 0) {
+    return true;
+  }
+  if (func_op.getNumResults() == 1) {
+    if (isa<LLVM::LLVMVoidType>(func_op.getResultTypes()[0])) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// Converts neura.return with no operands to neura.return_void.
+static void convertEmptyReturnToReturnVoid(Region &region, OpBuilder &builder) {
+  SmallVector<neura::ReturnOp> empty_returns;
+
+  region.walk([&](neura::ReturnOp ret_op) {
+    if (ret_op.getNumOperands() == 0) {
+      empty_returns.push_back(ret_op);
+    }
+  });
+
+  for (neura::ReturnOp ret_op : empty_returns) {
+    builder.setInsertionPoint(ret_op);
+    builder.create<neura::ReturnVoidOp>(ret_op.getLoc(), Value{});
+    ret_op.erase();
+  }
+}
+
+static void processEmptyReturnVoidBlock(Block *ret_block,
+                                        neura::ReturnVoidOp ret_void_op,
+                                        OpBuilder &builder) {
+  SmallVector<Block *> predecessor_blocks(ret_block->getPredecessors());
+  // Entry bolock with return_void is unreachable; no action needed.
+  if (predecessor_blocks.empty()) {
+    assert(false && "Entry block with neura.return_void is unreachable.");
+  }
+
+  // Adds a block argument for the trigger value.
+  BlockArgument trigger_arg =
+      ret_block->addArgument(builder.getI1Type(), ret_void_op.getLoc());
+
+  // Updates each predecessor's terminator to pass the trigger value.
+  for (Block *pred_block : predecessor_blocks) {
+    Operation *terminator = pred_block->getTerminator();
+    if (auto cond_br = dyn_cast<neura::CondBr>(terminator)) {
+      Value cond = cond_br.getCondition();
+      Value trigger_value = nullptr;
+
+      bool is_true_branch = (cond_br.getTrueDest() == ret_block);
+      bool is_false_branch = (cond_br.getFalseDest() == ret_block);
+
+      if (is_true_branch && !is_false_branch) {
+        // True branck leads to return_void, uses condition directly.
+        trigger_value = cond;
+      } else if (!is_true_branch && is_false_branch) {
+        // False branch leads to return_void, uses negated condition.
+        builder.setInsertionPoint(terminator);
+        Value negated_cond = builder.create<neura::NotOp>(terminator->getLoc(),
+                                                          cond.getType(), cond);
+        trigger_value = negated_cond;
+      } else {
+        assert(false &&
+               "Unsupported case: Both branches lead to neura.return_void.");
+      }
+
+      if (trigger_value) {
+        SmallVector<Value> true_args(cond_br.getTrueArgs());
+        SmallVector<Value> false_args(cond_br.getFalseArgs());
+
+        if (is_true_branch) {
+          true_args.push_back(trigger_value);
+        }
+        if (is_false_branch) {
+          false_args.push_back(trigger_value);
+        }
+
+        builder.setInsertionPoint(cond_br);
+        builder.create<neura::CondBr>(
+            cond_br.getLoc(), cond_br.getCondition(), true_args, false_args,
+            cond_br.getTrueDest(), cond_br.getFalseDest());
+        cond_br.erase();
+      }
+    } else if (auto br = dyn_cast<neura::Br>(terminator)) {
+      Value trigger_value;
+
+      // Looks for any suitable value in the predecessor block to use as
+      // trigger.
+      for (Operation &op : llvm::reverse(*pred_block)) {
+        if (&op == terminator) {
+          continue;
+        }
+        if (op.getNumResults() > 0) {
+          trigger_value = op.getResult(0);
+          break;
+        }
+      }
+
+      // If no suitable value is found, reports an error.
+      if (!trigger_value) {
+        assert(false && "No suitable value found for trigger.");
+      }
+
+      SmallVector<Value> args(br.getArgs());
+      args.push_back(trigger_value);
+
+      builder.setInsertionPoint(br);
+      builder.create<neura::Br>(br.getLoc(), args, br.getDest());
+      br.erase();
+    }
+  }
+  // Updates the return_void operation to use the block argument as trigger.
+  builder.setInsertionPoint(ret_void_op);
+  builder.create<neura::ReturnVoidOp>(ret_void_op.getLoc(), trigger_arg);
+  ret_void_op.erase();
+}
+
+struct CanonicalizeReturnPass
+    : public PassWrapper<CanonicalizeReturnPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CanonicalizeReturnPass)
+
+  StringRef getArgument() const override { return "canonicalize-return"; }
+  StringRef getDescription() const override {
+    return "Canonicalizes return operations in Neura functions.";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<neura::NeuraDialect, func::FuncDialect>();
+  }
+
+  void runOnOperation() override {
+    func::FuncOp func_op = getOperation();
+    // Checks for neura accelerator attribute.
+    auto accel_attr = func_op->getAttrOfType<StringAttr>("accelerator");
+    if (!accel_attr) {
+      return;
+    }
+
+    // Skips non-void functions.
+    if (!isVoidFunction(func_op)) {
+      return;
+    }
+
+    Region &region = func_op.getBody();
+    if (region.empty()) {
+      return;
+    }
+
+    OpBuilder builder(func_op.getContext());
+
+    // Step 1: Converts empty neura.return to neura.return_void.
+    convertEmptyReturnToReturnVoid(region, builder);
+
+    // Step 2: Collects all return_void operations without triggers.
+    SmallVector<neura::ReturnVoidOp> ret_void_ops;
+    region.walk([&](neura::ReturnVoidOp ret_void_op) {
+      if (!ret_void_op.getTrigger()) {
+        ret_void_ops.push_back(ret_void_op);
+      }
+    });
+
+    // Step 3: Processes each return_void block.
+    for (neura::ReturnVoidOp ret_void_op : ret_void_ops) {
+      Block *ret_block = ret_void_op->getBlock();
+
+      // Checks if ret_block only contains the return_void operation.
+      bool is_empty_block = (ret_block->getOperations().size() == 1);
+
+      if (is_empty_block) {
+        processEmptyReturnVoidBlock(ret_block, ret_void_op, builder);
+      } else {
+        // TODO: Handle non-empty return blocks.
+        // The basic idea is to create a new block that only contains the
+        // return_void operation, and redirect the original return block to this
+        // new block.
+        assert(false && "Unsupported case: return block is not empty.");
+      }
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::neura::createCanonicalizeReturnPass() {
+  return std::make_unique<CanonicalizeReturnPass>();
+}
\ No newline at end of file
diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
index 02d0e8df..e73a6ffa 100644
--- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
+++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
@@ -262,6 +262,8 @@ void buildControlFlowInfo(Region &region, ControlFlowInfo &ctrl_info,
 
     } else if (auto rt = dyn_cast<neura::ReturnOp>(terminator)) {
       llvm::errs() << "[ctrl2data] ReturnOp found: " << *rt << "\n";
+    } else if (auto rt = dyn_cast<neura::ReturnVoidOp>(terminator)) {
+      llvm::errs() << "[ctrl2data] ReturnVoidOp found: " << *rt << "\n";
     } else {
       llvm::errs() << "[ctrl2data] Unknown terminator: " << *terminator << "\n";
       assert(false);

From 5074b8e77d93db65908575d5797ad3b74532cd7b Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Mon, 29 Dec 2025 13:55:40 +0800
Subject: [PATCH 2/7] refactor the return logic

---
 .../Transforms/CanonicalizeReturnPass.cpp     | 149 ++++++++++--------
 .../TransformCtrlToDataFlowPass.cpp           |  24 ++-
 2 files changed, 106 insertions(+), 67 deletions(-)

diff --git a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
index 528653ff..871d47e2 100644
--- a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
+++ b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
@@ -56,77 +56,102 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
     assert(false && "Entry block with neura.return_void is unreachable.");
   }
 
-  // Adds a block argument for the trigger value.
-  BlockArgument trigger_arg =
-      ret_block->addArgument(builder.getI1Type(), ret_void_op.getLoc());
+  // Seperates predecessor blocks into cond_br and br blocks.
+  SmallVector<Block *> cond_br_preds;
+  SmallVector<Block *> br_preds;
 
-  // Updates each predecessor's terminator to pass the trigger value.
   for (Block *pred_block : predecessor_blocks) {
     Operation *terminator = pred_block->getTerminator();
-    if (auto cond_br = dyn_cast<neura::CondBr>(terminator)) {
-      Value cond = cond_br.getCondition();
-      Value trigger_value = nullptr;
-
-      bool is_true_branch = (cond_br.getTrueDest() == ret_block);
-      bool is_false_branch = (cond_br.getFalseDest() == ret_block);
-
-      if (is_true_branch && !is_false_branch) {
-        // True branck leads to return_void, uses condition directly.
-        trigger_value = cond;
-      } else if (!is_true_branch && is_false_branch) {
-        // False branch leads to return_void, uses negated condition.
-        builder.setInsertionPoint(terminator);
-        Value negated_cond = builder.create<neura::NotOp>(terminator->getLoc(),
-                                                          cond.getType(), cond);
-        trigger_value = negated_cond;
-      } else {
-        assert(false &&
-               "Unsupported case: Both branches lead to neura.return_void.");
-      }
+    if (isa<neura::CondBr>(terminator)) {
+      cond_br_preds.push_back(pred_block);
+    } else if (isa<neura::Br>(terminator)) {
+      br_preds.push_back(pred_block);
+    }
+  }
 
-      if (trigger_value) {
-        SmallVector<Value> true_args(cond_br.getTrueArgs());
-        SmallVector<Value> false_args(cond_br.getFalseArgs());
-
-        if (is_true_branch) {
-          true_args.push_back(trigger_value);
-        }
-        if (is_false_branch) {
-          false_args.push_back(trigger_value);
-        }
-
-        builder.setInsertionPoint(cond_br);
-        builder.create<neura::CondBr>(
-            cond_br.getLoc(), cond_br.getCondition(), true_args, false_args,
-            cond_br.getTrueDest(), cond_br.getFalseDest());
-        cond_br.erase();
-      }
-    } else if (auto br = dyn_cast<neura::Br>(terminator)) {
-      Value trigger_value;
-
-      // Looks for any suitable value in the predecessor block to use as
-      // trigger.
-      for (Operation &op : llvm::reverse(*pred_block)) {
-        if (&op == terminator) {
-          continue;
-        }
-        if (op.getNumResults() > 0) {
-          trigger_value = op.getResult(0);
-          break;
-        }
+  // Handles br_preds: copies return_void to pred_block, and utilizes a suitable
+  // value to trigger it.
+  for (Block *pred_block : br_preds) {
+    neura::Br br = cast<neura::Br>(pred_block->getTerminator());
+
+    // Finds a suitable trigger value in the predecessor block.
+    Value trigger_value = nullptr;
+
+    // Iterates through operations in reverse order to find the last suitable
+    // value.
+    for (Operation &op : llvm::reverse(*pred_block)) {
+      // Skips the terminator itself.
+      if (&op == br) {
+        continue;
       }
 
-      // If no suitable value is found, reports an error.
-      if (!trigger_value) {
-        assert(false && "No suitable value found for trigger.");
+      // Looks for any suitable value in the predecessor block.
+      if (op.getNumResults() > 0) {
+        trigger_value = op.getResult(0);
+        break;
       }
+    }
+
+    if (!trigger_value) {
+      assert(false && "No suitable value found in predecessor block.");
+    }
 
-      SmallVector<Value> args(br.getArgs());
-      args.push_back(trigger_value);
+    builder.setInsertionPoint(br);
+    builder.create<neura::ReturnVoidOp>(br.getLoc(), trigger_value);
+    br.erase();
+  }
+
+  // If there are no cond_br predecessors, removes the return_void block.
+  if (cond_br_preds.empty()) {
+    ret_void_op.erase();
+    ret_block->erase();
+    return;
+  }
+
+  // Handles cond_preds: adds a block argument for the trigger value, and
+  // updates each predecessor's terminator to pass the trigger value.
+  BlockArgument trigger_arg =
+      ret_block->addArgument(builder.getI1Type(), ret_void_op.getLoc());
+
+  // Updates each cond_pred block's terminator to pass the trigger value.
+  for (Block *pred_block : cond_br_preds) {
+    neura::CondBr cond_br = cast<neura::CondBr>(pred_block->getTerminator());
+    Value cond = cond_br.getCondition();
+    Value trigger_value = nullptr;
+
+    bool is_true_branch = (cond_br.getTrueDest() == ret_block);
+    bool is_false_branch = (cond_br.getFalseDest() == ret_block);
+
+    if (is_true_branch && !is_false_branch) {
+      // True branck leads to return_void, uses condition directly.
+      trigger_value = cond;
+    } else if (!is_true_branch && is_false_branch) {
+      // False branch leads to return_void, uses negated condition.
+      builder.setInsertionPoint(cond_br);
+      Value negated_cond =
+          builder.create<neura::NotOp>(cond_br.getLoc(), cond.getType(), cond);
+      trigger_value = negated_cond;
+    } else {
+      assert(false &&
+             "Unsupported case: Both branches lead to neura.return_void.");
+    }
+
+    if (trigger_value) {
+      SmallVector<Value> true_args(cond_br.getTrueArgs());
+      SmallVector<Value> false_args(cond_br.getFalseArgs());
+
+      if (is_true_branch) {
+        true_args.push_back(trigger_value);
+      }
+      if (is_false_branch) {
+        false_args.push_back(trigger_value);
+      }
 
-      builder.setInsertionPoint(br);
-      builder.create<neura::Br>(br.getLoc(), args, br.getDest());
-      br.erase();
+      builder.setInsertionPoint(cond_br);
+      builder.create<neura::CondBr>(
+          cond_br.getLoc(), cond_br.getCondition(), true_args, false_args,
+          cond_br.getTrueDest(), cond_br.getFalseDest());
+      cond_br.erase();
     }
   }
   // Updates the return_void operation to use the block argument as trigger.
diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
index e73a6ffa..dd6963db 100644
--- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
+++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
@@ -567,15 +567,29 @@ void transformControlFlowToDataFlow(Region &region, ControlFlowInfo &ctrl_info,
     }
   }
 
-  if (return_ops.size() > 1) {
+  SmallVector<neura::ReturnVoidOp> return_void_ops;
+  for (Operation &op : *entry_block) {
+    if (neura::ReturnVoidOp return_void_op =
+            dyn_cast<neura::ReturnVoidOp>(op)) {
+      return_void_ops.push_back(return_void_op);
+    }
+  }
+
+  if (return_ops.size() + return_void_ops.size() > 1) {
     llvm::errs() << "[ctrl2data] Error: Multiple ReturnOps found in the entry "
                     "block after flattening.\n";
     assert(false &&
            "Multiple ReturnOps found in the entry block after flattening.");
-  } else if (return_ops.size() == 1) {
-    neura::ReturnOp last_return = return_ops.back();
-    last_return->moveAfter(&entry_block->back());
-  } else {
+  } else if ((return_ops.size() == 1 && return_void_ops.size() == 0) ||
+             (return_ops.size() == 0 && return_void_ops.size() == 1)) {
+    if (return_void_ops.size() == 1) {
+      neura::ReturnVoidOp last_return_void = return_void_ops.back();
+      last_return_void->moveAfter(&entry_block->back());
+    } else if (return_ops.size() == 1) {
+      neura::ReturnOp last_return = return_ops.back();
+      last_return->moveAfter(&entry_block->back());
+    }
+  } else if (return_ops.size() == 0 && return_void_ops.size() == 0) {
     llvm::errs() << "[ctrl2data] Error: No ReturnOp found in the entry "
                     "block after flattening.\n";
     assert(false && "No ReturnOp found in the entry block after flattening.");

From c62e4da9ba9b364b6515fa9b11ac64d23fb64d00 Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Mon, 29 Dec 2025 16:55:19 +0800
Subject: [PATCH 3/7] distinguish return in ctrl & dataflow mode

---
 include/NeuraDialect/NeuraOps.td              | 30 ++++++--
 lib/NeuraDialect/NeuraOps.cpp                 | 23 ++++++
 .../Transforms/CanonicalizeReturnPass.cpp     | 64 ++++++++++-----
 .../TransformCtrlToDataFlowPass.cpp           | 77 ++++++++++---------
 .../TransformToSteerControlPass.cpp           |  2 +-
 5 files changed, 136 insertions(+), 60 deletions(-)

diff --git a/include/NeuraDialect/NeuraOps.td b/include/NeuraDialect/NeuraOps.td
index e0f8b012..32cf2acb 100644
--- a/include/NeuraDialect/NeuraOps.td
+++ b/include/NeuraDialect/NeuraOps.td
@@ -59,7 +59,7 @@ def Neura_KernelOp : Op<NeuraDialect, "kernel", [RecursiveMemoryEffects, SingleB
 }
 
 // Yield operation for fused_op and kernel regions.
-def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike, HasParent<"FusedOp, KernelOp">]> {
+def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike]> {
   let summary = "Yield values from a neura.kernel or neura.fused_op region.";
   let description = [{
     Returns values from a neura.kernel or neura.fused_op region to the parent operation.
@@ -79,6 +79,8 @@ def Neura_YieldOp : Op<NeuraDialect, "yield", [Terminator, Pure, ReturnLike, Has
   ];
 
   let assemblyFormat = [{($values^ `:` type($values))? attr-dict}];
+
+  let hasVerifier = 1;
 }
 
 
@@ -375,9 +377,9 @@ def Neura_ReturnOp : Op<NeuraDialect, "return", [Terminator]> {
   // let assemblyFormat = "($values^)? `,` attr-dict";
 }
 
-// Defines a return operation for void functions.
-def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", [Terminator]>{
-  let summary = "Return operation for void functions with execution trigger.";
+// Defines a return operation for void functions in dataflow mode.
+def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", []>{
+  let summary = "Return operation for void functions in dataflow mode.";
   let description = [{
     Terminates a void function. The trigger operand provides the execution
     condition in dataflow mode - the return executes when the trigger becomes valid.
@@ -386,13 +388,31 @@ def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", [Terminator]>{
     always executes. After canonicalization, void returns should have triggers.
 
     Example:
-      neura.return_void %trigger : i1       // Triggered return in ctrl mode
       neura.return_void %cond : !neura.data<i1, i1>  // In dataflow mode
+
+    This is NOT a terminator - the block ends with neura.yield instead.
   }];
   let arguments = (ins Optional<AnyType>:$trigger);
   let assemblyFormat = "($trigger^ `:` type($trigger))? attr-dict";
 }
 
+// Defines a return operation for non-void fnctions in dataflow mode.
+def Neura_ReturnValueOp : Op<NeuraDialect, "return_value", []>{
+  let summary = "Return operation for non-void functions in dataflow mode.";
+  let description = [{
+    Represents a return point for non-void functions in dataflow mode.
+    The value operand provides the return value and the trigger operand.
+
+    This is NOT a terminator - the block ends with neura.yield instead.
+
+    Example:
+      neura.return_value %ret_val : !neura.data<i32, i1> // In dataflow mode
+  }];
+  let arguments = (ins Variadic<AnyType>:$values);
+  let assemblyFormat = "($values^ `:` type($values))? attr-dict";
+}
+
+
 // Defines a cast operation for type conversion.
 def Neura_CastOp : Op<NeuraDialect, "cast">{
   let summary = "Generic type conversion operation";
diff --git a/lib/NeuraDialect/NeuraOps.cpp b/lib/NeuraDialect/NeuraOps.cpp
index 4242633a..3c30f8cd 100644
--- a/lib/NeuraDialect/NeuraOps.cpp
+++ b/lib/NeuraDialect/NeuraOps.cpp
@@ -1,5 +1,7 @@
 #include "NeuraDialect/NeuraOps.h"
 #include "NeuraDialect/NeuraDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/Support/LogicalResult.h"
@@ -7,6 +9,27 @@
 using namespace mlir;
 using namespace mlir::neura;
 
+LogicalResult YieldOp::verify() {
+  Operation *parent_op = (*this)->getParentOp();
+
+  if (!parent_op) {
+    return emitOpError("must have a parent operation.");
+  }
+
+  // Allows yield in FusedOp and KernelOp
+  if (isa<FusedOp>(parent_op) || isa<KernelOp>(parent_op)) {
+    return success();
+  }
+
+  // Allows yield in func.func (for dataflow mode)
+  if (isa<func::FuncOp>(parent_op)) {
+    return success();
+  }
+
+  return emitOpError("expects parent op to be one of 'neura.fused_op', "
+                     "'neura.kernel', or 'func.func'");
+}
+
 LogicalResult PhiStartOp::verify() {
   // Checks if this phi_start is inside a fused_op.
   Operation *parent_op = getOperation()->getParentOp();
diff --git a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
index 871d47e2..b803ac40 100644
--- a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
+++ b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
@@ -4,6 +4,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Pass/Pass.h"
@@ -17,6 +18,11 @@ using namespace mlir;
 #include "NeuraDialect/NeuraPasses.h.inc"
 
 namespace {
+// Return type attribute values.
+constexpr const char *kReturnTypeAttr = "return_type";
+constexpr const char *kReturnTypeVoid = "void";
+constexpr const char *kReturnTypeValue = "value";
+
 // Checks if a function is a void function (i.e., has no return values).
 static bool isVoidFunction(func::FuncOp func_op) {
   if (func_op.getNumResults() == 0) {
@@ -47,8 +53,25 @@ static void convertEmptyReturnToReturnVoid(Region &region, OpBuilder &builder) {
   }
 }
 
+// Marks empty returns with "is_void" attribute and adds trigger values.
+static void processVoidReturns(Region &region, OpBuilder &builder) {
+  SmallVector<neura::ReturnOp> empty_returns;
+
+  region.walk([&](neura::ReturnOp ret_op) {
+    if (ret_op.getNumOperands() == 0) {
+      empty_returns.push_back(ret_op);
+    } else {
+      ret_op->setAttr(kReturnTypeAttr, builder.getStringAttr(kReturnTypeValue));
+    }
+  });
+
+  for (neura::ReturnOp ret_op : empty_returns) {
+    ret_op->setAttr(kReturnTypeAttr, builder.getStringAttr(kReturnTypeVoid));
+  }
+}
+
 static void processEmptyReturnVoidBlock(Block *ret_block,
-                                        neura::ReturnVoidOp ret_void_op,
+                                        neura::ReturnOp void_ret_op,
                                         OpBuilder &builder) {
   SmallVector<Block *> predecessor_blocks(ret_block->getPredecessors());
   // Entry bolock with return_void is unreachable; no action needed.
@@ -97,13 +120,14 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
     }
 
     builder.setInsertionPoint(br);
-    builder.create<neura::ReturnVoidOp>(br.getLoc(), trigger_value);
+    auto new_ret = builder.create<neura::ReturnOp>(br.getLoc(), trigger_value);
+    new_ret->setAttr(kReturnTypeAttr, builder.getStringAttr(kReturnTypeVoid));
     br.erase();
   }
 
   // If there are no cond_br predecessors, removes the return_void block.
   if (cond_br_preds.empty()) {
-    ret_void_op.erase();
+    void_ret_op.erase();
     ret_block->erase();
     return;
   }
@@ -111,7 +135,7 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
   // Handles cond_preds: adds a block argument for the trigger value, and
   // updates each predecessor's terminator to pass the trigger value.
   BlockArgument trigger_arg =
-      ret_block->addArgument(builder.getI1Type(), ret_void_op.getLoc());
+      ret_block->addArgument(builder.getI1Type(), void_ret_op.getLoc());
 
   // Updates each cond_pred block's terminator to pass the trigger value.
   for (Block *pred_block : cond_br_preds) {
@@ -132,8 +156,7 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
           builder.create<neura::NotOp>(cond_br.getLoc(), cond.getType(), cond);
       trigger_value = negated_cond;
     } else {
-      assert(false &&
-             "Unsupported case: Both branches lead to neura.return_void.");
+      assert(false && "Unsupported case: Both branches lead to neura.return.");
     }
 
     if (trigger_value) {
@@ -155,9 +178,11 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
     }
   }
   // Updates the return_void operation to use the block argument as trigger.
-  builder.setInsertionPoint(ret_void_op);
-  builder.create<neura::ReturnVoidOp>(ret_void_op.getLoc(), trigger_arg);
-  ret_void_op.erase();
+  builder.setInsertionPoint(void_ret_op);
+  auto new_ret =
+      builder.create<neura::ReturnOp>(void_ret_op.getLoc(), trigger_arg);
+  new_ret->setAttr(kReturnTypeAttr, builder.getStringAttr(kReturnTypeVoid));
+  void_ret_op.erase();
 }
 
 struct CanonicalizeReturnPass
@@ -193,19 +218,22 @@ struct CanonicalizeReturnPass
 
     OpBuilder builder(func_op.getContext());
 
-    // Step 1: Converts empty neura.return to neura.return_void.
-    convertEmptyReturnToReturnVoid(region, builder);
-
-    // Step 2: Collects all return_void operations without triggers.
-    SmallVector<neura::ReturnVoidOp> ret_void_ops;
-    region.walk([&](neura::ReturnVoidOp ret_void_op) {
-      if (!ret_void_op.getTrigger()) {
-        ret_void_ops.push_back(ret_void_op);
+    // Step 1: Marks empty returns with "is_void" attribute.
+    processVoidReturns(region, builder);
+
+    // Step 2: Collects all return operations with "is_void" attribute.
+    SmallVector<neura::ReturnOp> ret_void_ops;
+    region.walk([&](neura::ReturnOp ret_op) {
+      if (ret_op->hasAttr(kReturnTypeAttr)) {
+        if (dyn_cast<StringAttr>(ret_op->getAttr(kReturnTypeAttr)).getValue() ==
+            kReturnTypeVoid) {
+          ret_void_ops.push_back(ret_op);
+        }
       }
     });
 
     // Step 3: Processes each return_void block.
-    for (neura::ReturnVoidOp ret_void_op : ret_void_ops) {
+    for (neura::ReturnOp ret_void_op : ret_void_ops) {
       Block *ret_block = ret_void_op->getBlock();
 
       // Checks if ret_block only contains the return_void operation.
diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
index dd6963db..06d4f5a7 100644
--- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
+++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
@@ -262,27 +262,38 @@ void buildControlFlowInfo(Region &region, ControlFlowInfo &ctrl_info,
 
     } else if (auto rt = dyn_cast<neura::ReturnOp>(terminator)) {
       llvm::errs() << "[ctrl2data] ReturnOp found: " << *rt << "\n";
-    } else if (auto rt = dyn_cast<neura::ReturnVoidOp>(terminator)) {
-      llvm::errs() << "[ctrl2data] ReturnVoidOp found: " << *rt << "\n";
     } else {
-      llvm::errs() << "[ctrl2data] Unknown terminator: " << *terminator << "\n";
-      assert(false);
+      assert(false && "Unknown terminator operation in control flow graph.");
     }
   }
 }
 
-Value getPrecessedCondition(Value condition, bool is_not_condition,
+Value getProcessedCondition(Value condition, bool is_not_condition,
                             llvm::MapVector<Value, Value> &condition_cache,
                             OpBuilder &builder) {
   if (!is_not_condition) {
     return condition;
   }
 
+  // First, checks if we already have a cached negated condition.
   auto it = condition_cache.find(condition);
   if (it != condition_cache.end()) {
     return it->second;
   }
 
+  // Second, searches for an existing not operation that uses this condition.
+  // This handles the not operations created by CanonicalizeReturnPass.
+  for (OpOperand &use : condition.getUses()) {
+    if (neura::NotOp not_op = dyn_cast<neura::NotOp>(use.getOwner())) {
+      if (not_op.getOperand() == condition) {
+        Value not_result = not_op.getResult();
+        condition_cache[condition] = not_result;
+        return not_result;
+      }
+    }
+  }
+
+  // Third, creates a new not operation to negate the condition.
   Block *source = condition.getDefiningOp()->getBlock();
   builder.setInsertionPoint(source->getTerminator());
   Value not_condition = builder.create<neura::NotOp>(
@@ -384,7 +395,7 @@ void createReserveAndPhiOps(
 
       Value predicated_val = val;
       if (edge->condition) {
-        Value processed_condition = getPrecessedCondition(
+        Value processed_condition = getProcessedCondition(
             edge->condition, edge->is_not_condition, condition_cache, builder);
         predicated_val = builder.create<neura::GrantPredicateOp>(
             edge->condition.getLoc(), val.getType(), predicated_val,
@@ -412,7 +423,7 @@ void createReserveAndPhiOps(
       Value predicated_val = val;
       if (edge->condition) {
         builder.setInsertionPoint(edge->source->getTerminator());
-        Value processed_condition = getPrecessedCondition(
+        Value processed_condition = getProcessedCondition(
             edge->condition, edge->is_not_condition, condition_cache, builder);
 
         predicated_val = builder.create<neura::GrantPredicateOp>(
@@ -559,43 +570,37 @@ void transformControlFlowToDataFlow(Region &region, ControlFlowInfo &ctrl_info,
     block->erase();
   }
 
-  // Handles return operations in the entry block.
+  // Converts neura.return to return_void or return_value.
   SmallVector<neura::ReturnOp> return_ops;
-  for (Operation &op : *entry_block) {
-    if (neura::ReturnOp return_op = dyn_cast<neura::ReturnOp>(op)) {
+  for (Operation &op : llvm::make_early_inc_range(*entry_block)) {
+    if (auto return_op = dyn_cast<neura::ReturnOp>(op)) {
       return_ops.push_back(return_op);
     }
   }
 
-  SmallVector<neura::ReturnVoidOp> return_void_ops;
-  for (Operation &op : *entry_block) {
-    if (neura::ReturnVoidOp return_void_op =
-            dyn_cast<neura::ReturnVoidOp>(op)) {
-      return_void_ops.push_back(return_void_op);
-    }
-  }
-
-  if (return_ops.size() + return_void_ops.size() > 1) {
-    llvm::errs() << "[ctrl2data] Error: Multiple ReturnOps found in the entry "
-                    "block after flattening.\n";
-    assert(false &&
-           "Multiple ReturnOps found in the entry block after flattening.");
-  } else if ((return_ops.size() == 1 && return_void_ops.size() == 0) ||
-             (return_ops.size() == 0 && return_void_ops.size() == 1)) {
-    if (return_void_ops.size() == 1) {
-      neura::ReturnVoidOp last_return_void = return_void_ops.back();
-      last_return_void->moveAfter(&entry_block->back());
-    } else if (return_ops.size() == 1) {
-      neura::ReturnOp last_return = return_ops.back();
-      last_return->moveAfter(&entry_block->back());
+  // Adds neura.yield at the end of the entry block as terminator.
+  builder.setInsertionPointToEnd(entry_block);
+  builder.create<neura::YieldOp>(entry_block->getTerminator()->getLoc());
+
+  for (neura::ReturnOp return_op : return_ops) {
+    builder.setInsertionPoint(return_op);
+
+    if (return_op->hasAttr("return_type") &&
+        dyn_cast<StringAttr>(return_op->getAttr("return_type")).getValue() ==
+            "void") {
+      Value trigger = return_op->getOperand(0);
+      builder.create<neura::ReturnVoidOp>(return_op.getLoc(), trigger);
+    } else if (return_op->hasAttr("return_type") &&
+               dyn_cast<StringAttr>(return_op->getAttr("return_type"))
+                       .getValue() == "value") {
+      builder.create<neura::ReturnValueOp>(return_op.getLoc(),
+                                           return_op.getValues());
     }
-  } else if (return_ops.size() == 0 && return_void_ops.size() == 0) {
-    llvm::errs() << "[ctrl2data] Error: No ReturnOp found in the entry "
-                    "block after flattening.\n";
-    assert(false && "No ReturnOp found in the entry block after flattening.");
+    return_op.erase();
   }
 
-  // Sets the "dataflow_mode" attribute to "predicate" for the parent function.
+  // Sets the "dataflow_mode" attribute to "predicate" for the parent
+  // function.
   if (auto func = dyn_cast<func::FuncOp>(region.getParentOp())) {
     if (!func->hasAttr("dataflow_mode")) {
       func->setAttr("dataflow_mode",
diff --git a/lib/NeuraDialect/Transforms/TransformToSteerControlPass.cpp b/lib/NeuraDialect/Transforms/TransformToSteerControlPass.cpp
index b2389616..8ce98592 100644
--- a/lib/NeuraDialect/Transforms/TransformToSteerControlPass.cpp
+++ b/lib/NeuraDialect/Transforms/TransformToSteerControlPass.cpp
@@ -723,7 +723,7 @@ struct TransformToSteerControlPass
       }
     });
 
-    // 从后向前删除，避免依赖问题
+    // Deletes the collected operations in reverse order.
     for (auto it = to_erase.rbegin(); it != to_erase.rend(); ++it) {
       (*it)->erase();
     }

From 327c343b8fc94e5679b93ee1b2b0368f03684b29 Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Tue, 30 Dec 2025 00:15:02 +0800
Subject: [PATCH 4/7] modify tests

---
 lib/NeuraDialect/Mapping/mapping_util.cpp     |   3 +-
 lib/NeuraDialect/NeuraPasses.cpp              |   1 +
 .../Transforms/CanonicalizeReturnPass.cpp     |  38 +-
 .../TransformCtrlToDataFlowPass.cpp           |  14 +-
 .../bert/bert_node1/bert_node1.mlir           | 106 +--
 .../bert/bert_node28/bert_node28.mlir         | 235 ++++---
 test/arith2neura/add.mlir                     |  11 -
 test/c2llvm2mlir/nested_loop/test.mlir        |   6 +-
 test/c2llvm2mlir/simple_loop/test.mlir        |   2 +
 test/code_gen/test_code_generate.mlir         | 104 +--
 .../complex_nested/complex_nested.mlir        | 655 ++++++------------
 .../non_perfect_nested.mlir                   | 557 +++++++--------
 .../perfect_nested/perfect_nested.mlir        | 109 +--
 .../perfect_reduction/perfect_reduction.mlir  |   6 +-
 .../simple_loop/simple_loop.mlir              |  71 +-
 .../simple_loop_reduction.mlir                |  11 +-
 test/e2e/bicg/bicg_kernel.mlir                | 162 ++---
 test/e2e/fir/fir_kernel.mlir                  |  83 +--
 test/e2e/fir/fir_kernel_vec.mlir              |  92 +--
 test/e2e/histogram/histogram_kernel.mlir      | 198 +++---
 test/e2e/relu/relu_kernel.mlir                | 130 ++--
 test/honor_arch/fir_removed_tiles_test.mlir   |   1 +
 test/mapping_quality/branch_for.mlir          |  16 +-
 test/mapping_quality/tiny_loop.mlir           |   2 +
 test/neura/ctrl/branch.mlir                   |   7 +-
 test/neura/ctrl/branch_for.mlir               |  18 +-
 .../ctrl/branch_with_and_without_arg.mlir     |   4 +-
 test/neura/ctrl/branch_without_arg.mlir       |   4 +-
 test/neura/ctrl/for_with_if.mlir              |   4 +-
 test/neura/ctrl/nested_branch.mlir            |   7 +-
 test/neura/for_loop/kernel_test.mlir          |  13 +-
 test/neura/for_loop/relu_test.mlir            |  12 +-
 test/neura/fusion/test.mlir                   |  20 +-
 .../loop_convert_controlflow_to_dataflow.mlir |   1 +
 test/neura/steer_ctrl/for_with_if.mlir        |   4 +-
 .../steer_ctrl/loop_with_return_value.mlir    |   5 +-
 .../steer_ctrl/loop_without_return_value.mlir |  18 +-
 37 files changed, 1277 insertions(+), 1453 deletions(-)
 delete mode 100644 test/arith2neura/add.mlir

diff --git a/lib/NeuraDialect/Mapping/mapping_util.cpp b/lib/NeuraDialect/Mapping/mapping_util.cpp
index fe2e4c81..df4de446 100644
--- a/lib/NeuraDialect/Mapping/mapping_util.cpp
+++ b/lib/NeuraDialect/Mapping/mapping_util.cpp
@@ -126,7 +126,8 @@ OperationKind getOperationKindFromMlirOp(Operation *op) {
 // Returns true if the operation does not need CGRA tile placement.
 bool is_non_materialized(Operation *op) {
   // Returns true if the operation does not need CGRA tile placement.
-  return mlir::isa<neura::ReserveOp, neura::CtrlMovOp, neura::DataMovOp>(op);
+  return mlir::isa<neura::ReserveOp, neura::CtrlMovOp, neura::DataMovOp,
+                   neura::YieldOp>(op);
 }
 
 } // namespace neura
diff --git a/lib/NeuraDialect/NeuraPasses.cpp b/lib/NeuraDialect/NeuraPasses.cpp
index 958b3f46..72a78ec7 100644
--- a/lib/NeuraDialect/NeuraPasses.cpp
+++ b/lib/NeuraDialect/NeuraPasses.cpp
@@ -25,6 +25,7 @@ void mlir::neura::registerNeuraConversionPassPipeline() {
         pm.addPass(mlir::createLowerLlvmToNeuraPass());
         pm.addPass(mlir::createPrintOpGraphPass(os));
 
+        pm.addPass(mlir::neura::createCanonicalizeReturnPass());
         pm.addPass(mlir::neura::createCanonicalizeCastPass());
         pm.addPass(mlir::neura::createPromoteFuncArgToConstPass());
         pm.addPass(mlir::neura::createFoldConstantPass());
diff --git a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
index b803ac40..d6054c40 100644
--- a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
+++ b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/Dominance.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/raw_ostream.h"
 #include <memory>
 
@@ -36,31 +37,17 @@ static bool isVoidFunction(func::FuncOp func_op) {
   return false;
 }
 
-// Converts neura.return with no operands to neura.return_void.
-static void convertEmptyReturnToReturnVoid(Region &region, OpBuilder &builder) {
-  SmallVector<neura::ReturnOp> empty_returns;
-
-  region.walk([&](neura::ReturnOp ret_op) {
-    if (ret_op.getNumOperands() == 0) {
-      empty_returns.push_back(ret_op);
-    }
-  });
-
-  for (neura::ReturnOp ret_op : empty_returns) {
-    builder.setInsertionPoint(ret_op);
-    builder.create<neura::ReturnVoidOp>(ret_op.getLoc(), Value{});
-    ret_op.erase();
-  }
-}
-
 // Marks empty returns with "is_void" attribute and adds trigger values.
-static void processVoidReturns(Region &region, OpBuilder &builder) {
+static void processReturns(Region &region, OpBuilder &builder) {
   SmallVector<neura::ReturnOp> empty_returns;
 
   region.walk([&](neura::ReturnOp ret_op) {
+    llvm::errs() << "[ctrl2data] Processing neura.return operation...\n";
+    llvm::errs() << ret_op << "\n";
     if (ret_op.getNumOperands() == 0) {
       empty_returns.push_back(ret_op);
     } else {
+      llvm::errs() << "[ctrl2data] Marking neura.return with value...\n";
       ret_op->setAttr(kReturnTypeAttr, builder.getStringAttr(kReturnTypeValue));
     }
   });
@@ -206,11 +193,6 @@ struct CanonicalizeReturnPass
       return;
     }
 
-    // Skips non-void functions.
-    if (!isVoidFunction(func_op)) {
-      return;
-    }
-
     Region &region = func_op.getBody();
     if (region.empty()) {
       return;
@@ -218,8 +200,14 @@ struct CanonicalizeReturnPass
 
     OpBuilder builder(func_op.getContext());
 
-    // Step 1: Marks empty returns with "is_void" attribute.
-    processVoidReturns(region, builder);
+    // Step 1: Marks empty returns with "void" attribute.
+    processReturns(region, builder);
+
+    if (!isVoidFunction(func_op)) {
+      llvm::errs() << "[ctrl2data] Function is not void, no further action "
+                      "needed.\n";
+      return;
+    }
 
     // Step 2: Collects all return operations with "is_void" attribute.
     SmallVector<neura::ReturnOp> ret_void_ops;
diff --git a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
index 06d4f5a7..8601663b 100644
--- a/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
+++ b/lib/NeuraDialect/Transforms/TransformCtrlToDataFlowPass.cpp
@@ -578,9 +578,8 @@ void transformControlFlowToDataFlow(Region &region, ControlFlowInfo &ctrl_info,
     }
   }
 
-  // Adds neura.yield at the end of the entry block as terminator.
-  builder.setInsertionPointToEnd(entry_block);
-  builder.create<neura::YieldOp>(entry_block->getTerminator()->getLoc());
+  llvm::errs() << "[ctrl2data] Converting neura.return operations to "
+                  "return_void/value...\n";
 
   for (neura::ReturnOp return_op : return_ops) {
     builder.setInsertionPoint(return_op);
@@ -588,6 +587,7 @@ void transformControlFlowToDataFlow(Region &region, ControlFlowInfo &ctrl_info,
     if (return_op->hasAttr("return_type") &&
         dyn_cast<StringAttr>(return_op->getAttr("return_type")).getValue() ==
             "void") {
+      llvm::errs() << "[ctrl2data] Converting to neura.return_void.\n";
       Value trigger = return_op->getOperand(0);
       builder.create<neura::ReturnVoidOp>(return_op.getLoc(), trigger);
     } else if (return_op->hasAttr("return_type") &&
@@ -595,10 +595,18 @@ void transformControlFlowToDataFlow(Region &region, ControlFlowInfo &ctrl_info,
                        .getValue() == "value") {
       builder.create<neura::ReturnValueOp>(return_op.getLoc(),
                                            return_op.getValues());
+    } else {
+      assert(false && "Unknown return type attribute in neura.return.");
     }
     return_op.erase();
   }
 
+  llvm::errs()
+      << "[ctrl2data] All neura.return operations converted successfully.\n";
+  // Adds neura.yield at the end of the entry block as terminator.
+  builder.setInsertionPointToEnd(entry_block);
+  builder.create<neura::YieldOp>(builder.getUnknownLoc());
+
   // Sets the "dataflow_mode" attribute to "predicate" for the parent
   // function.
   if (auto func = dyn_cast<func::FuncOp>(region.getParentOp())) {
diff --git a/test/affine2neura/bert/bert_node1/bert_node1.mlir b/test/affine2neura/bert/bert_node1/bert_node1.mlir
index dedb0c14..64610dfd 100644
--- a/test/affine2neura/bert/bert_node1/bert_node1.mlir
+++ b/test/affine2neura/bert/bert_node1/bert_node1.mlir
@@ -15,6 +15,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -86,57 +87,60 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %20 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %21 = neura.phi_start %9, %20 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %22 = "neura.icmp"(%21, %19) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %23 = neura.grant_predicate %17, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %24 = neura.grant_predicate %19, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %25 = neura.grant_predicate %15, %22 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %26 = neura.grant_predicate %13, %22 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %27 = neura.grant_predicate %21, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %28 = neura.grant_predicate %11, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %29 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %30 = neura.phi_start %28, %29 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %23 = "neura.not"(%22) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %24 = neura.grant_predicate %17, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %25 = neura.grant_predicate %19, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %26 = neura.grant_predicate %15, %22 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %27 = neura.grant_predicate %13, %22 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %28 = neura.grant_predicate %21, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %29 = neura.grant_predicate %11, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %30 = neura.grant_predicate %23, %23 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %30 : !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %31 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %32 = neura.phi_start %27, %31 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %33 = neura.reserve : !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %34 = neura.phi_start %26, %33 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<memref<?x1x128x1x1x128xi8>, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %35 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %36 = neura.phi_start %23, %35 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %37 = neura.reserve : !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %38 = neura.phi_start %25, %37 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<memref<?x1x1x1x1x128xi8>, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %39 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %40 = neura.phi_start %24, %39 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %32 = neura.phi_start %29, %31 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %33 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %34 = neura.phi_start %28, %33 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %35 = neura.reserve : !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %36 = neura.phi_start %27, %35 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<memref<?x1x128x1x1x128xi8>, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %37 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %38 = neura.phi_start %24, %37 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %39 = neura.reserve : !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %40 = neura.phi_start %26, %39 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<memref<?x1x1x1x1x128xi8>, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
 // CTRL2DATA-NEXT:     %41 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %42 = neura.phi_start %23, %41 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %43 = "neura.icmp"(%42, %40) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %44 = neura.grant_predicate %38, %43 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %45 = neura.grant_predicate %36, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %42, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %34, %43 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %32, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %49 = neura.grant_predicate %30, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %50 = neura.grant_predicate %40, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %51 = "neura.not"(%43) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %52 = neura.grant_predicate %32, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %53 = neura.grant_predicate %30, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %54 = neura.grant_predicate %40, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %55 = neura.grant_predicate %36, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %56 = neura.grant_predicate %38, %51 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %57 = neura.grant_predicate %34, %51 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %58 = "neura.add"(%52, %53) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %58 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %54 -> %18 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %55 -> %16 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %56 -> %14 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %57 -> %12 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %53 -> %10 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %59 = neura.load_indexed %44[%45, %45, %45, %45, %45, %46 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x1x1x1x128xi8>, i1> : !neura.data<i8, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %59 to %47[%45, %45, %48, %45, %45, %46 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x128x1x1x128xi8>, i1> : !neura.data<i8, i1>
-// CTRL2DATA-NEXT:     %60 = "neura.add"(%46, %49) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %60 -> %41 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %50 -> %39 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %44 -> %37 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %45 -> %35 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %47 -> %33 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %48 -> %31 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %49 -> %29 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
+// CTRL2DATA-NEXT:     %42 = neura.phi_start %25, %41 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %43 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %44 = neura.phi_start %24, %43 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %45 = "neura.icmp"(%44, %42) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %40, %45 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %38, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %44, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %49 = neura.grant_predicate %36, %45 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %50 = neura.grant_predicate %34, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %51 = neura.grant_predicate %32, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %52 = neura.grant_predicate %42, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %53 = "neura.not"(%45) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %54 = neura.grant_predicate %34, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %55 = neura.grant_predicate %32, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %56 = neura.grant_predicate %42, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %57 = neura.grant_predicate %38, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %58 = neura.grant_predicate %40, %53 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %59 = neura.grant_predicate %36, %53 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %60 = "neura.add"(%54, %55) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %60 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %56 -> %18 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %57 -> %16 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %58 -> %14 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %59 -> %12 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %55 -> %10 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %61 = neura.load_indexed %46[%47, %47, %47, %47, %47, %48 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x1x1x1x128xi8>, i1> : !neura.data<i8, i1>
+// CTRL2DATA-NEXT:     neura.store_indexed %61 to %49[%47, %47, %50, %47, %47, %48 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x128x1x1x128xi8>, i1> : !neura.data<i8, i1>
+// CTRL2DATA-NEXT:     %62 = "neura.add"(%48, %51) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %62 -> %43 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %52 -> %41 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %46 -> %39 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %47 -> %37 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %49 -> %35 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %50 -> %33 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %51 -> %31 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/affine2neura/bert/bert_node28/bert_node28.mlir b/test/affine2neura/bert/bert_node28/bert_node28.mlir
index bff2def5..65494bb4 100644
--- a/test/affine2neura/bert/bert_node28/bert_node28.mlir
+++ b/test/affine2neura/bert/bert_node28/bert_node28.mlir
@@ -15,10 +15,12 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
 // RUN: | FileCheck %s --check-prefix=CTRL2DATA
+
 module attributes {} {
   func.func @_Z11bert_node28PA128_A768_KfPA768_S0_PA128_A768_f(%arg0: memref<?x128x768xf32>, %arg1: memref<?x768x768xf32>, %arg2: memref<?x128x768xf32>) attributes {} {
     affine.for %arg3 = 0 to 128 {
@@ -114,127 +116,130 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %28 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %29 = neura.phi_start %13, %28 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %30 = "neura.icmp"(%29, %27) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %31 = neura.grant_predicate %25, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %32 = neura.grant_predicate %23, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %33 = neura.grant_predicate %29, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %34 = neura.grant_predicate %21, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %35 = neura.grant_predicate %27, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %36 = neura.grant_predicate %19, %30 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %37 = neura.grant_predicate %17, %30 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %38 = neura.grant_predicate %15, %30 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %39 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %40 = neura.phi_start %38, %39 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %41 = neura.reserve : !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %42 = neura.phi_start %37, %41 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<memref<?x768x768xf32>, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %43 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %44 = neura.phi_start %36, %43 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %45 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %46 = neura.phi_start %35, %45 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %31 = "neura.not"(%30) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %32 = neura.grant_predicate %25, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %33 = neura.grant_predicate %23, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %34 = neura.grant_predicate %29, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %35 = neura.grant_predicate %21, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %36 = neura.grant_predicate %27, %30 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %37 = neura.grant_predicate %19, %30 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %38 = neura.grant_predicate %17, %30 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %39 = neura.grant_predicate %15, %30 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %40 = neura.grant_predicate %31, %31 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %40 : !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %41 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %42 = neura.phi_start %39, %41 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %43 = neura.reserve : !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %44 = neura.phi_start %38, %43 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<memref<?x768x768xf32>, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %45 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %46 = neura.phi_start %37, %45 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
 // CTRL2DATA-NEXT:     %47 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %48 = neura.phi_start %34, %47 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %48 = neura.phi_start %36, %47 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %49 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %50 = neura.phi_start %33, %49 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %50 = neura.phi_start %35, %49 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %51 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %52 = neura.phi_start %31, %51 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %52 = neura.phi_start %34, %51 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %53 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %54 = neura.phi_start %32, %53 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %55 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %56 = neura.phi_start %31, %55 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %57 = "neura.icmp"(%56, %54) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %58 = neura.grant_predicate %52, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %59 = neura.grant_predicate %54, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %60 = neura.grant_predicate %44, %57 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %61 = neura.grant_predicate %50, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %62 = neura.grant_predicate %42, %57 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %63 = neura.grant_predicate %56, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %64 = neura.grant_predicate %40, %57 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %65 = neura.grant_predicate %48, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %66 = neura.grant_predicate %46, %57 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %67 = "neura.not"(%57) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %68 = neura.grant_predicate %50, %67 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %69 = neura.grant_predicate %48, %67 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %70 = neura.grant_predicate %46, %67 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %71 = neura.grant_predicate %52, %67 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %72 = neura.grant_predicate %54, %67 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %73 = neura.grant_predicate %44, %67 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %74 = neura.grant_predicate %42, %67 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %75 = neura.grant_predicate %40, %67 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %76 = "neura.add"(%68, %69) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %76 -> %28 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %70 -> %26 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %71 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %72 -> %22 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %69 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %73 -> %18 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %74 -> %16 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %75 -> %14 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %77 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %78 = neura.phi_start %66, %77 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %56 = neura.phi_start %33, %55 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %57 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %58 = neura.phi_start %32, %57 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %59 = "neura.icmp"(%58, %56) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %60 = neura.grant_predicate %54, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %61 = neura.grant_predicate %56, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %62 = neura.grant_predicate %46, %59 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %63 = neura.grant_predicate %52, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %64 = neura.grant_predicate %44, %59 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %65 = neura.grant_predicate %58, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %66 = neura.grant_predicate %42, %59 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %67 = neura.grant_predicate %50, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %68 = neura.grant_predicate %48, %59 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %69 = "neura.not"(%59) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %70 = neura.grant_predicate %52, %69 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %71 = neura.grant_predicate %50, %69 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %72 = neura.grant_predicate %48, %69 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %73 = neura.grant_predicate %54, %69 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %74 = neura.grant_predicate %56, %69 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %75 = neura.grant_predicate %46, %69 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %76 = neura.grant_predicate %44, %69 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %77 = neura.grant_predicate %42, %69 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %78 = "neura.add"(%70, %71) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %78 -> %28 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %72 -> %26 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %73 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %74 -> %22 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %71 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %75 -> %18 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %76 -> %16 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %77 -> %14 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
 // CTRL2DATA-NEXT:     %79 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %80 = neura.phi_start %65, %79 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %81 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %82 = neura.phi_start %64, %81 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %83 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %84 = neura.phi_start %63, %83 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %85 = neura.reserve : !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %86 = neura.phi_start %62, %85 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<memref<?x768x768xf32>, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %87 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %88 = neura.phi_start %61, %87 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %80 = neura.phi_start %68, %79 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %81 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %82 = neura.phi_start %67, %81 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %83 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %84 = neura.phi_start %66, %83 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %85 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %86 = neura.phi_start %65, %85 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %87 = neura.reserve : !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %88 = neura.phi_start %64, %87 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<memref<?x768x768xf32>, i1> -> !neura.data<memref<?x768x768xf32>, i1>
 // CTRL2DATA-NEXT:     %89 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %90 = neura.phi_start %58, %89 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %91 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %92 = neura.phi_start %60, %91 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %93 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %94 = neura.phi_start %59, %93 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %90 = neura.phi_start %63, %89 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %91 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %92 = neura.phi_start %60, %91 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %93 = neura.reserve : !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %94 = neura.phi_start %62, %93 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<memref<?x128x768xf32>, i1> -> !neura.data<memref<?x128x768xf32>, i1>
 // CTRL2DATA-NEXT:     %95 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %96 = neura.phi_start %58, %95 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %97 = "neura.icmp"(%96, %94) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %98 = neura.grant_predicate %92, %97 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %99 = neura.grant_predicate %90, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %100 = neura.grant_predicate %88, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %101 = neura.grant_predicate %96, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %102 = neura.grant_predicate %86, %97 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %103 = neura.grant_predicate %84, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %104 = neura.grant_predicate %82, %97 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %105 = neura.grant_predicate %80, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %106 = neura.grant_predicate %94, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %107 = neura.grant_predicate %78, %97 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %108 = "neura.not"(%97) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %109 = neura.grant_predicate %84, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %110 = neura.grant_predicate %80, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %111 = neura.grant_predicate %94, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %112 = neura.grant_predicate %90, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %113 = neura.grant_predicate %88, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %114 = neura.grant_predicate %78, %108 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %115 = neura.grant_predicate %92, %108 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %116 = neura.grant_predicate %86, %108 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     %117 = neura.grant_predicate %82, %108 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %118 = "neura.add"(%109, %110) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %118 -> %55 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %111 -> %53 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %112 -> %51 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %113 -> %49 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %110 -> %47 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %114 -> %45 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %115 -> %43 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %116 -> %41 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %117 -> %39 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     %119 = neura.load_indexed %98[%99, %100, %101 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     %120 = neura.load_indexed %102[%99, %101, %103 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x768x768xf32>, i1> : !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     %121 = neura.load_indexed %104[%99, %100, %103 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     %122 = "neura.fmul"(%119, %120) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     %123 = "neura.fadd"(%121, %122) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %123 to %104[%99, %100, %103 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     %124 = "neura.add"(%101, %105) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %124 -> %95 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %106 -> %93 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %98 -> %91 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %99 -> %89 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %100 -> %87 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %102 -> %85 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %103 -> %83 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %104 -> %81 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %105 -> %79 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %107 -> %77 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
-// CTRL2DATA-NEXT:   }
+// CTRL2DATA-NEXT:     %96 = neura.phi_start %61, %95 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %97 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %98 = neura.phi_start %60, %97 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %99 = "neura.icmp"(%98, %96) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %100 = neura.grant_predicate %94, %99 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %101 = neura.grant_predicate %92, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %102 = neura.grant_predicate %90, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %103 = neura.grant_predicate %98, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %104 = neura.grant_predicate %88, %99 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %105 = neura.grant_predicate %86, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %106 = neura.grant_predicate %84, %99 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %107 = neura.grant_predicate %82, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %108 = neura.grant_predicate %96, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %109 = neura.grant_predicate %80, %99 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %110 = "neura.not"(%99) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %111 = neura.grant_predicate %86, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %112 = neura.grant_predicate %82, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %113 = neura.grant_predicate %96, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %114 = neura.grant_predicate %92, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %115 = neura.grant_predicate %90, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %116 = neura.grant_predicate %80, %110 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %117 = neura.grant_predicate %94, %110 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %118 = neura.grant_predicate %88, %110 : !neura.data<memref<?x768x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     %119 = neura.grant_predicate %84, %110 : !neura.data<memref<?x128x768xf32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %120 = "neura.add"(%111, %112) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %120 -> %57 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %113 -> %55 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %114 -> %53 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %115 -> %51 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %112 -> %49 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %116 -> %47 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %117 -> %45 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %118 -> %43 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %119 -> %41 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     %121 = neura.load_indexed %100[%101, %102, %103 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     %122 = neura.load_indexed %104[%101, %103, %105 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x768x768xf32>, i1> : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     %123 = neura.load_indexed %106[%101, %102, %105 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     %124 = "neura.fmul"(%121, %122) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     %125 = "neura.fadd"(%123, %124) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.store_indexed %125 to %106[%101, %102, %105 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128x768xf32>, i1> : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     %126 = "neura.add"(%103, %107) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %126 -> %97 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %108 -> %95 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %100 -> %93 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %101 -> %91 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %102 -> %89 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %104 -> %87 : !neura.data<memref<?x768x768xf32>, i1> !neura.data<memref<?x768x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %105 -> %85 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %106 -> %83 : !neura.data<memref<?x128x768xf32>, i1> !neura.data<memref<?x128x768xf32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %107 -> %81 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %109 -> %79 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.yield
+// CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/arith2neura/add.mlir b/test/arith2neura/add.mlir
deleted file mode 100644
index 71e1a995..00000000
--- a/test/arith2neura/add.mlir
+++ /dev/null
@@ -1,11 +0,0 @@
-// RUN: neura-compiler --neura-conversion %s | FileCheck %s --check-prefix=COMPILER
-// RUN: mlir-neura-opt --assign-accelerator --lower-arith-to-neura %s | FileCheck %s --check-prefix=OPT
-
-func.func @test(%a: f32) -> f32 {
-  %b = arith.constant 2.0 : f32
-  %res = arith.addf %a, %b : f32
-  return %res : f32
-}
-
-// COMPILER: neura.fadd
-// OPT:      neura.fadd
\ No newline at end of file
diff --git a/test/c2llvm2mlir/nested_loop/test.mlir b/test/c2llvm2mlir/nested_loop/test.mlir
index d73a6d60..ee5b9131 100644
--- a/test/c2llvm2mlir/nested_loop/test.mlir
+++ b/test/c2llvm2mlir/nested_loop/test.mlir
@@ -4,6 +4,7 @@
 // RUN: mlir-neura-opt --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -14,6 +15,7 @@
 // RUN: mlir-neura-opt --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -25,7 +27,7 @@
 // CHECK-LLVM2NEURA: accelerator = "neura"
 // CHECK-LLVM2NEURA: %25 = neura.alloca %24 : !neura.data<i32, i1> -> !neura.data<!llvm.ptr, i1>
 // CHECK-LLVM2NEURA: %38 = neura.phi_start %37, %36 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CHECK-LLVM2NEURA: %182 = neura.sext %181 : !neura.data<i32, i1> -> !neura.data<i64, i1>
-// CHECK-LLVM2NEURA: %201 = "neura.mul"(%199, %200) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CHECK-LLVM2NEURA: %188 = neura.sext %187 : !neura.data<i32, i1> -> !neura.data<i64, i1>
+// CHECK-LLVM2NEURA: %207 = "neura.mul"(%205, %206) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
 
 // CHECK-LLVM2NEURA-MAP: func.func @_Z6kernelPiS_S_(%arg0: !llvm.ptr {llvm.noundef}, %arg1: !llvm.ptr {llvm.noundef}, %arg2: !llvm.ptr {llvm.noundef}) -> !llvm.void attributes {CConv = #llvm.cconv<ccc>, accelerator = "neura", dataflow_mode = "predicate", frame_pointer = #llvm.framePointerKind<all>, linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 13 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 9 : i32, res_mii = 6 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}, no_inline, no_unwind, optimize_none, passthrough = ["mustprogress", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic", unnamed_addr = 0 : i64, visibility_ = 0 : i64} {
\ No newline at end of file
diff --git a/test/c2llvm2mlir/simple_loop/test.mlir b/test/c2llvm2mlir/simple_loop/test.mlir
index bb89c67e..09285607 100644
--- a/test/c2llvm2mlir/simple_loop/test.mlir
+++ b/test/c2llvm2mlir/simple_loop/test.mlir
@@ -24,6 +24,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -37,6 +38,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/code_gen/test_code_generate.mlir b/test/code_gen/test_code_generate.mlir
index 6353b259..c8fbb422 100644
--- a/test/code_gen/test_code_generate.mlir
+++ b/test/code_gen/test_code_generate.mlir
@@ -1,6 +1,7 @@
 // RUN: mlir-neura-opt %s \
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -69,60 +70,59 @@ func.func @loop_test() -> f32 {
 //   written into local register $20.
 //
 
-// YAML: array_config:
-// YAML:   columns: 4
-// YAML:   rows: 4
-// YAML:   compiled_ii: 5
-// YAML:   cores:
-// YAML:     - column: 1
-// YAML:       row: 0
-// YAML:       core_id: "1"
-// YAML:       entries:
-// YAML:         - entry_id: "entry0"
-// YAML:           instructions:
-// YAML:             - index_per_ii: 3
-// YAML:               operations:
-// YAML:                 - opcode: "RETURN"
-// YAML:                   id: 65
-// YAML:                   time_step: 8
-// YAML:                   invalid_iterations: 1
-// YAML:                   src_operands:
-// YAML:                     - operand: "NORTH"
-// YAML:                       color: "RED"
-// YAML:     - column: 0
-// YAML:       row: 1
-// YAML:       core_id: "4"
-// YAML:       entries:
-// YAML:         - entry_id: "entry0"
-// YAML:           instructions:
-// YAML:             - index_per_ii: 0
-// YAML:               operations:
-// YAML:                 - opcode: "DATA_MOV"
-// YAML:                   id: 36
-// YAML:                   time_step: 5
-// YAML:                   invalid_iterations: 1
-// YAML:                   src_operands:
-// YAML:                     - operand: "EAST"
-// YAML:                       color: "RED"
-// YAML:                   dst_operands:
-// YAML:                     - operand: "$0"
-// YAML:                       color: "RED"
-// YAML:                 - opcode: "DATA_MOV"
-// YAML:                   id: 49
-// YAML:                   time_step: 5
-// YAML:                   invalid_iterations: 1
-// YAML:                   src_operands:
-// YAML:                     - operand: "EAST"
-// YAML:                       color: "RED"
-// YAML:                   dst_operands:
-// YAML:                     - operand: "$3"
-// YAML:                       color: "RED"
-// YAML:             - index_per_ii: 1
-// YAML:               operations:
-// ASM: # Compiled II: 5
+// YAML:      array_config:
+// YAML-NEXT:   columns: 4
+// YAML-NEXT:   rows: 4
+// YAML-NEXT:   compiled_ii: 5
+// YAML-NEXT:   cores:
+// YAML-NEXT:     - column: 1
+// YAML-NEXT:       row: 0
+// YAML-NEXT:       core_id: "1"
+// YAML-NEXT:       entries:
+// YAML-NEXT:         - entry_id: "entry0"
+// YAML-NEXT:           instructions:
+// YAML-NEXT:             - index_per_ii: 3
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "RETURN_VALUE"
+// YAML-NEXT:                   id: 66
+// YAML-NEXT:                   time_step: 8
+// YAML-NEXT:                   invalid_iterations: 1
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "NORTH"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:     - column: 0
+// YAML-NEXT:       row: 1
+// YAML-NEXT:       core_id: "4"
+// YAML-NEXT:       entries:
+// YAML-NEXT:         - entry_id: "entry0"
+// YAML-NEXT:           instructions:
+// YAML-NEXT:             - index_per_ii: 0
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "DATA_MOV"
+// YAML-NEXT:                   id: 37
+// YAML-NEXT:                   time_step: 5
+// YAML-NEXT:                   invalid_iterations: 1
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "EAST"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                 - opcode: "DATA_MOV"
+// YAML-NEXT:                   id: 50
+// YAML-NEXT:                   time_step: 5
+// YAML-NEXT:                   invalid_iterations: 1
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "EAST"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$3"
+// YAML-NEXT:                       color: "RED"
+
+// ASM:      # Compiled II: 5
 // ASM:      PE(1,0):
 // ASM-NEXT: {
-// ASM-NEXT:   RETURN, [NORTH, RED] (t=8, inv_iters=1)
+// ASM-NEXT:   RETURN_VALUE, [NORTH, RED] (t=8, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=3)
 // ASM:      PE(0,1):
 // ASM-NEXT: {
diff --git a/test/controflow_fuse/complex_nested/complex_nested.mlir b/test/controflow_fuse/complex_nested/complex_nested.mlir
index e8ecf0ae..0fa153b4 100644
--- a/test/controflow_fuse/complex_nested/complex_nested.mlir
+++ b/test/controflow_fuse/complex_nested/complex_nested.mlir
@@ -15,6 +15,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -216,457 +217,235 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %36 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %37 = neura.phi_start %17, %36 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %38 = "neura.icmp"(%37, %35) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %39 = neura.grant_predicate %33, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %40 = neura.grant_predicate %35, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %41 = neura.grant_predicate %31, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %42 = neura.grant_predicate %29, %38 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %43 = neura.grant_predicate %37, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %44 = neura.grant_predicate %27, %38 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %45 = neura.grant_predicate %25, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %23, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %21, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %19, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %49 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %50 = neura.phi_start %48, %49 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %39 = "neura.not"(%38) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %40 = neura.grant_predicate %33, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %41 = neura.grant_predicate %35, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %42 = neura.grant_predicate %31, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %43 = neura.grant_predicate %29, %38 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %44 = neura.grant_predicate %37, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %45 = neura.grant_predicate %27, %38 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %25, %38 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %23, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %21, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %49 = neura.grant_predicate %19, %38 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %50 = neura.grant_predicate %39, %39 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %50 : !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %51 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %52 = neura.phi_start %47, %51 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %52 = neura.phi_start %49, %51 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %53 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %54 = neura.phi_start %46, %53 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %55 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %56 = neura.phi_start %45, %55 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %57 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %58 = neura.phi_start %44, %57 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %59 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %60 = neura.phi_start %39, %59 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %54 = neura.phi_start %48, %53 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %55 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %56 = neura.phi_start %47, %55 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %57 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %58 = neura.phi_start %46, %57 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %59 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %60 = neura.phi_start %45, %59 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
 // CTRL2DATA-NEXT:     %61 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %62 = neura.phi_start %43, %61 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %63 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %64 = neura.phi_start %42, %63 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %65 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %66 = neura.phi_start %41, %65 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %67 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %68 = neura.phi_start %40, %67 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %62 = neura.phi_start %40, %61 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %63 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %64 = neura.phi_start %44, %63 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %65 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %66 = neura.phi_start %43, %65 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %67 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %68 = neura.phi_start %42, %67 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %69 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %70 = neura.phi_start %39, %69 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %71 = "neura.icmp"(%70, %68) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %72 = neura.grant_predicate %66, %71 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %73 = neura.grant_predicate %64, %71 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %74 = neura.grant_predicate %62, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %75 = neura.grant_predicate %70, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %76 = neura.grant_predicate %60, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %77 = neura.grant_predicate %68, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %78 = neura.grant_predicate %58, %71 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %79 = neura.grant_predicate %56, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %80 = neura.grant_predicate %54, %71 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %81 = neura.grant_predicate %52, %71 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %82 = neura.grant_predicate %50, %71 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %83 = "neura.not"(%71) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %84 = neura.grant_predicate %60, %83 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %85 = neura.grant_predicate %66, %83 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %86 = neura.grant_predicate %68, %83 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %87 = neura.grant_predicate %64, %83 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %88 = neura.grant_predicate %62, %83 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %89 = neura.grant_predicate %56, %83 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %90 = neura.grant_predicate %54, %83 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %91 = neura.grant_predicate %52, %83 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %92 = neura.grant_predicate %58, %83 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %93 = neura.grant_predicate %50, %83 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %94 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %95 = neura.phi_start %85, %94 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %70 = neura.phi_start %41, %69 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %71 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %72 = neura.phi_start %40, %71 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %73 = "neura.icmp"(%72, %70) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %74 = neura.grant_predicate %68, %73 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %75 = neura.grant_predicate %66, %73 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %76 = neura.grant_predicate %64, %73 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %77 = neura.grant_predicate %72, %73 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %78 = neura.grant_predicate %62, %73 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %79 = neura.grant_predicate %70, %73 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %80 = neura.grant_predicate %60, %73 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %81 = neura.grant_predicate %58, %73 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %82 = neura.grant_predicate %56, %73 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %83 = neura.grant_predicate %54, %73 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %84 = neura.grant_predicate %52, %73 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %85 = "neura.not"(%73) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %86 = neura.grant_predicate %62, %85 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %87 = neura.grant_predicate %68, %85 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %88 = neura.grant_predicate %70, %85 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %89 = neura.grant_predicate %66, %85 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %90 = neura.grant_predicate %64, %85 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %91 = neura.grant_predicate %58, %85 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %92 = neura.grant_predicate %56, %85 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %93 = neura.grant_predicate %54, %85 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %94 = neura.grant_predicate %60, %85 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %95 = neura.grant_predicate %52, %85 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %96 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %97 = neura.phi_start %93, %96 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %98 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %99 = neura.phi_start %92, %98 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %100 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %101 = neura.phi_start %91, %100 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %102 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %103 = neura.phi_start %84, %102 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %104 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %105 = neura.phi_start %90, %104 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %106 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %107 = neura.phi_start %89, %106 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %97 = neura.phi_start %87, %96 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %98 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %99 = neura.phi_start %95, %98 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %100 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %101 = neura.phi_start %94, %100 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %102 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %103 = neura.phi_start %93, %102 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %104 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %105 = neura.phi_start %86, %104 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %106 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %107 = neura.phi_start %92, %106 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %108 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %109 = neura.phi_start %88, %108 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %110 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %111 = neura.phi_start %87, %110 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %112 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %113 = neura.phi_start %86, %112 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %114 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %115 = neura.phi_start %85, %114 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %116 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %117 = neura.phi_start %84, %116 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %118 = "neura.icmp"(%117, %113) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %119 = neura.grant_predicate %111, %118 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %120 = neura.grant_predicate %109, %118 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %121 = neura.grant_predicate %117, %118 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %122 = neura.grant_predicate %115, %118 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %123 = neura.grant_predicate %107, %118 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %124 = neura.grant_predicate %113, %118 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %125 = neura.grant_predicate %105, %118 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %126 = neura.grant_predicate %103, %118 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %127 = neura.grant_predicate %101, %118 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %128 = neura.grant_predicate %99, %118 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %129 = neura.grant_predicate %97, %118 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %130 = neura.grant_predicate %95, %118 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %131 = "neura.not"(%118) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %132 = neura.grant_predicate %115, %131 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %133 = neura.grant_predicate %105, %131 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %134 = neura.grant_predicate %103, %131 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %135 = neura.grant_predicate %113, %131 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %136 = neura.grant_predicate %101, %131 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %137 = neura.grant_predicate %99, %131 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %138 = neura.grant_predicate %109, %131 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %139 = neura.grant_predicate %107, %131 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %140 = neura.grant_predicate %111, %131 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %141 = neura.grant_predicate %97, %131 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %142 = neura.grant_predicate %95, %131 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %143 = "neura.div"(%132, %133) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %144 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %145 = neura.phi_start %133, %144 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %109 = neura.phi_start %91, %108 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %110 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %111 = neura.phi_start %90, %110 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %112 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %113 = neura.phi_start %89, %112 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %114 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %115 = neura.phi_start %88, %114 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %116 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %117 = neura.phi_start %87, %116 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %118 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %119 = neura.phi_start %86, %118 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %120 = "neura.icmp"(%119, %115) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %121 = neura.grant_predicate %113, %120 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %122 = neura.grant_predicate %111, %120 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %123 = neura.grant_predicate %119, %120 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %124 = neura.grant_predicate %117, %120 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %125 = neura.grant_predicate %109, %120 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %126 = neura.grant_predicate %115, %120 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %127 = neura.grant_predicate %107, %120 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %128 = neura.grant_predicate %105, %120 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %129 = neura.grant_predicate %103, %120 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %130 = neura.grant_predicate %101, %120 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %131 = neura.grant_predicate %99, %120 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %132 = neura.grant_predicate %97, %120 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %133 = "neura.not"(%120) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %134 = neura.grant_predicate %117, %133 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %135 = neura.grant_predicate %107, %133 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %136 = neura.grant_predicate %105, %133 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %137 = neura.grant_predicate %115, %133 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %138 = neura.grant_predicate %103, %133 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %139 = neura.grant_predicate %101, %133 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %140 = neura.grant_predicate %111, %133 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %141 = neura.grant_predicate %109, %133 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %142 = neura.grant_predicate %113, %133 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %143 = neura.grant_predicate %99, %133 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %144 = neura.grant_predicate %97, %133 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %145 = "neura.div"(%134, %135) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %146 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %147 = neura.phi_start %142, %146 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %147 = neura.phi_start %135, %146 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %148 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %149 = neura.phi_start %143, %148 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %149 = neura.phi_start %144, %148 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %150 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %151 = neura.phi_start %141, %150 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %152 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %153 = neura.phi_start %140, %152 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %154 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %155 = neura.phi_start %139, %154 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %151 = neura.phi_start %145, %150 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %152 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %153 = neura.phi_start %143, %152 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %154 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %155 = neura.phi_start %142, %154 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
 // CTRL2DATA-NEXT:     %156 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %157 = neura.phi_start %138, %156 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %158 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %159 = neura.phi_start %137, %158 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %160 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %161 = neura.phi_start %136, %160 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %162 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %163 = neura.phi_start %134, %162 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %157 = neura.phi_start %141, %156 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %158 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %159 = neura.phi_start %140, %158 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %160 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %161 = neura.phi_start %139, %160 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %162 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %163 = neura.phi_start %138, %162 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %164 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %165 = neura.phi_start %135, %164 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %165 = neura.phi_start %136, %164 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %166 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %167 = neura.phi_start %134, %166 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %168 = "neura.icmp"(%167, %165) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %169 = neura.grant_predicate %163, %168 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %170 = neura.grant_predicate %161, %168 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %171 = neura.grant_predicate %165, %168 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %172 = neura.grant_predicate %159, %168 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %173 = neura.grant_predicate %167, %168 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %174 = neura.grant_predicate %157, %168 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %175 = neura.grant_predicate %155, %168 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %176 = neura.grant_predicate %153, %168 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %177 = neura.grant_predicate %151, %168 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %178 = neura.grant_predicate %149, %168 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %179 = neura.grant_predicate %147, %168 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %180 = neura.grant_predicate %145, %168 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %181 = "neura.not"(%168) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %182 = neura.grant_predicate %163, %181 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %183 = neura.grant_predicate %165, %181 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %184 = neura.grant_predicate %153, %181 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %185 = neura.grant_predicate %157, %181 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %186 = neura.grant_predicate %149, %181 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %187 = neura.grant_predicate %155, %181 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %188 = neura.grant_predicate %147, %181 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %189 = neura.grant_predicate %159, %181 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %190 = neura.grant_predicate %145, %181 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %191 = neura.grant_predicate %161, %181 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %192 = neura.grant_predicate %151, %181 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %193 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %194 = neura.phi_start %192, %193 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %167 = neura.phi_start %137, %166 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %168 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %169 = neura.phi_start %136, %168 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %170 = "neura.icmp"(%169, %167) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %171 = neura.grant_predicate %165, %170 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %172 = neura.grant_predicate %163, %170 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %173 = neura.grant_predicate %167, %170 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %174 = neura.grant_predicate %161, %170 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %175 = neura.grant_predicate %169, %170 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %176 = neura.grant_predicate %159, %170 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %177 = neura.grant_predicate %157, %170 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %178 = neura.grant_predicate %155, %170 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %179 = neura.grant_predicate %153, %170 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %180 = neura.grant_predicate %151, %170 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %181 = neura.grant_predicate %149, %170 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %182 = neura.grant_predicate %147, %170 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %183 = "neura.not"(%170) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %184 = neura.grant_predicate %165, %183 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %185 = neura.grant_predicate %167, %183 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %186 = neura.grant_predicate %155, %183 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %187 = neura.grant_predicate %159, %183 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %188 = neura.grant_predicate %151, %183 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %189 = neura.grant_predicate %157, %183 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %190 = neura.grant_predicate %149, %183 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %191 = neura.grant_predicate %161, %183 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %192 = neura.grant_predicate %147, %183 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %193 = neura.grant_predicate %163, %183 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %194 = neura.grant_predicate %153, %183 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %195 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %196 = neura.phi_start %191, %195 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %196 = neura.phi_start %194, %195 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %197 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %198 = neura.phi_start %190, %197 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %199 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %200 = neura.phi_start %189, %199 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %201 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %202 = neura.phi_start %188, %201 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %203 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %204 = neura.phi_start %182, %203 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %198 = neura.phi_start %193, %197 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %199 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %200 = neura.phi_start %192, %199 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %201 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %202 = neura.phi_start %191, %201 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %203 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %204 = neura.phi_start %190, %203 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %205 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %206 = neura.phi_start %187, %205 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %207 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %208 = neura.phi_start %186, %207 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %209 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %210 = neura.phi_start %185, %209 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %211 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %212 = neura.phi_start %184, %211 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %213 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %214 = neura.phi_start %183, %213 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %206 = neura.phi_start %184, %205 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %207 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %208 = neura.phi_start %189, %207 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %209 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %210 = neura.phi_start %188, %209 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %211 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %212 = neura.phi_start %187, %211 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %213 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %214 = neura.phi_start %186, %213 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
 // CTRL2DATA-NEXT:     %215 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %216 = neura.phi_start %182, %215 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %217 = "neura.icmp"(%216, %214) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %218 = neura.grant_predicate %212, %217 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %219 = neura.grant_predicate %210, %217 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %220 = neura.grant_predicate %216, %217 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %221 = neura.grant_predicate %208, %217 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %222 = neura.grant_predicate %206, %217 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %223 = neura.grant_predicate %214, %217 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %224 = neura.grant_predicate %204, %217 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %225 = neura.grant_predicate %202, %217 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %226 = neura.grant_predicate %200, %217 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %227 = neura.grant_predicate %198, %217 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %228 = neura.grant_predicate %196, %217 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %229 = neura.grant_predicate %194, %217 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %230 = "neura.not"(%217) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %231 = neura.grant_predicate %210, %230 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %232 = neura.grant_predicate %206, %230 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %233 = neura.grant_predicate %214, %230 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %234 = neura.grant_predicate %204, %230 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %235 = neura.grant_predicate %202, %230 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %236 = neura.grant_predicate %212, %230 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %237 = neura.grant_predicate %200, %230 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %238 = neura.grant_predicate %198, %230 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %239 = neura.grant_predicate %196, %230 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %240 = neura.grant_predicate %194, %230 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %241 = "neura.add"(%231, %232) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %241 -> %36 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %233 -> %34 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %234 -> %32 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %235 -> %30 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %236 -> %28 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %237 -> %26 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %232 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %238 -> %22 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %239 -> %20 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %240 -> %18 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %242 = neura.load_indexed %218[%219, %220 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %243 = "neura.icmp"(%242, %221) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %244 = neura.grant_predicate %221, %243 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %245 = neura.grant_predicate %218, %243 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %246 = neura.grant_predicate %219, %243 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %247 = neura.grant_predicate %220, %243 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %248 = neura.grant_predicate %222, %243 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %249 = neura.grant_predicate %223, %243 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %250 = neura.grant_predicate %224, %243 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %251 = neura.grant_predicate %225, %243 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %252 = neura.grant_predicate %226, %243 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %253 = neura.grant_predicate %227, %243 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %254 = neura.grant_predicate %228, %243 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %255 = neura.grant_predicate %229, %243 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %256 = "neura.not"(%243) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %257 = neura.grant_predicate %220, %256 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %258 = neura.grant_predicate %222, %256 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %259 = neura.grant_predicate %223, %256 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %260 = neura.grant_predicate %218, %256 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %261 = neura.grant_predicate %219, %256 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %262 = neura.grant_predicate %221, %256 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %263 = neura.grant_predicate %224, %256 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %264 = neura.grant_predicate %225, %256 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %265 = neura.grant_predicate %226, %256 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %266 = neura.grant_predicate %227, %256 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %267 = neura.grant_predicate %228, %256 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %268 = neura.grant_predicate %229, %256 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %244 to %245[%246, %247 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %269 = "neura.phi"(%268, %255) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %270 = "neura.phi"(%267, %254) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %271 = "neura.phi"(%266, %253) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %272 = "neura.phi"(%265, %252) : (!neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1>) -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %273 = "neura.phi"(%264, %251) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %274 = "neura.phi"(%263, %250) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %275 = "neura.phi"(%262, %244) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %276 = "neura.phi"(%261, %246) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %277 = "neura.phi"(%260, %245) : (!neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1>) -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %278 = "neura.phi"(%259, %249) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %279 = "neura.phi"(%258, %248) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %280 = "neura.phi"(%257, %247) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %281 = "neura.add"(%280, %279) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %281 -> %215 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %278 -> %213 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %277 -> %211 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %276 -> %209 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %275 -> %207 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %279 -> %205 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %274 -> %203 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %273 -> %201 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %272 -> %199 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %271 -> %197 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %270 -> %195 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %269 -> %193 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %282 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %283 = neura.phi_start %180, %282 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %284 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %285 = neura.phi_start %179, %284 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %286 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %287 = neura.phi_start %178, %286 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %288 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %289 = neura.phi_start %170, %288 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %290 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %291 = neura.phi_start %169, %290 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %292 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %293 = neura.phi_start %177, %292 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %294 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %295 = neura.phi_start %176, %294 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %296 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %297 = neura.phi_start %175, %296 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %298 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %299 = neura.phi_start %174, %298 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %300 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %301 = neura.phi_start %173, %300 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %302 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %303 = neura.phi_start %172, %302 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %304 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %305 = neura.phi_start %171, %304 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %306 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %307 = neura.phi_start %170, %306 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %308 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %309 = neura.phi_start %169, %308 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %310 = "neura.icmp"(%309, %305) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %311 = neura.grant_predicate %303, %310 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %312 = neura.grant_predicate %309, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %313 = neura.grant_predicate %301, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %314 = neura.grant_predicate %299, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %315 = neura.grant_predicate %307, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %316 = neura.grant_predicate %297, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %317 = neura.grant_predicate %305, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %318 = neura.grant_predicate %295, %310 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %319 = neura.grant_predicate %293, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %320 = neura.grant_predicate %291, %310 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %321 = neura.grant_predicate %289, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %322 = neura.grant_predicate %287, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %323 = neura.grant_predicate %285, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %324 = neura.grant_predicate %283, %310 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %325 = "neura.not"(%310) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %326 = neura.grant_predicate %295, %325 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %327 = neura.grant_predicate %299, %325 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %328 = neura.grant_predicate %301, %325 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %329 = neura.grant_predicate %307, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %330 = neura.grant_predicate %293, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %331 = neura.grant_predicate %297, %325 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %332 = neura.grant_predicate %305, %325 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %333 = neura.grant_predicate %291, %325 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %334 = neura.grant_predicate %289, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %335 = neura.grant_predicate %303, %325 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %336 = neura.grant_predicate %287, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %337 = neura.grant_predicate %285, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %338 = neura.grant_predicate %283, %325 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %339 = neura.load_indexed %326[%327, %328 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %340 = "neura.mul"(%339, %329) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %341 = "neura.div"(%340, %330) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %341 to %326[%327, %328 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %342 = "neura.add"(%328, %331) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %342 -> %166 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %332 -> %164 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %333 -> %162 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %334 -> %160 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %335 -> %158 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %327 -> %156 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %331 -> %154 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %326 -> %152 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %330 -> %150 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %336 -> %148 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %337 -> %146 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %338 -> %144 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %343 = neura.load_indexed %311[%312, %313, %314 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %344 = "neura.icmp"(%343, %315) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %345 = "neura.sel"(%344, %343, %315) : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %346 = "neura.add"(%312, %316) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %346 -> %308 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %345 -> %306 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %317 -> %304 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %311 -> %302 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %313 -> %300 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %314 -> %298 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %316 -> %296 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %318 -> %294 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %319 -> %292 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %320 -> %290 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %321 -> %288 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %322 -> %286 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %323 -> %284 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %324 -> %282 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %347 = neura.load_indexed %119[%120, %121 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %348 = "neura.add"(%122, %347) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %349 = "neura.add"(%121, %123) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %349 -> %116 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %348 -> %114 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %124 -> %112 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %119 -> %110 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %120 -> %108 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %123 -> %106 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %125 -> %104 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %126 -> %102 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %127 -> %100 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %128 -> %98 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %129 -> %96 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %130 -> %94 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %72 to %73[%74, %75 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %350 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %351 = neura.phi_start %82, %350 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %352 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %353 = neura.phi_start %81, %352 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %354 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %355 = neura.phi_start %80, %354 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %356 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %357 = neura.phi_start %76, %356 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %358 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %359 = neura.phi_start %72, %358 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %360 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %361 = neura.phi_start %79, %360 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %362 = neura.reserve : !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %363 = neura.phi_start %73, %362 : !neura.data<memref<?x32xi32>, i1>, !neura.data<memref<?x32xi32>, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %364 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %365 = neura.phi_start %75, %364 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %366 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %367 = neura.phi_start %74, %366 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %368 = neura.reserve : !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %369 = neura.phi_start %78, %368 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<memref<?x32x32xi32>, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %370 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %371 = neura.phi_start %77, %370 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %372 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %373 = neura.phi_start %76, %372 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %374 = "neura.icmp"(%373, %371) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %375 = neura.grant_predicate %369, %374 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %376 = neura.grant_predicate %367, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %377 = neura.grant_predicate %365, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %378 = neura.grant_predicate %373, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %379 = neura.grant_predicate %363, %374 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %380 = neura.grant_predicate %361, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %381 = neura.grant_predicate %371, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %382 = neura.grant_predicate %359, %374 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %383 = neura.grant_predicate %357, %374 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %384 = neura.grant_predicate %355, %374 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %385 = neura.grant_predicate %353, %374 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %386 = neura.grant_predicate %351, %374 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %387 = "neura.not"(%374) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %388 = neura.grant_predicate %365, %387 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %389 = neura.grant_predicate %361, %387 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %390 = neura.grant_predicate %371, %387 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %391 = neura.grant_predicate %359, %387 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %392 = neura.grant_predicate %363, %387 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     %393 = neura.grant_predicate %367, %387 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %394 = neura.grant_predicate %357, %387 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %395 = neura.grant_predicate %369, %387 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     %396 = neura.grant_predicate %355, %387 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %397 = neura.grant_predicate %353, %387 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %398 = neura.grant_predicate %351, %387 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %399 = "neura.add"(%388, %389) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %399 -> %69 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %390 -> %67 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %391 -> %65 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %392 -> %63 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %393 -> %61 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %394 -> %59 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %395 -> %57 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %389 -> %55 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %396 -> %53 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %397 -> %51 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %398 -> %49 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %400 = neura.load_indexed %375[%376, %377, %378 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %401 = neura.load_indexed %379[%376, %377 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %402 = "neura.add"(%401, %400) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %402 to %379[%376, %377 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %403 = "neura.add"(%378, %380) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %403 -> %372 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %381 -> %370 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %375 -> %368 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %376 -> %366 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %377 -> %364 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %379 -> %362 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %380 -> %360 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %382 -> %358 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %383 -> %356 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %384 -> %354 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %385 -> %352 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %386 -> %350 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
-// CTRL2DATA-NEXT:   }
+// CTRL2DATA-NEXT:     %216 = neura.phi_start %185, %215 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %217 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %218 = neura.phi_start %184, %217 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %219 = "neura.icmp"(%218, %216) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %220 = neura.grant_predicate %214, %219 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %221 = neura.grant_predicate %212, %219 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %222 = neura.grant_predicate %218, %219 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %223 = neura.grant_predicate %210, %219 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %224 = neura.grant_predicate %208, %219 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %225 = neura.grant_predicate %216, %219 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %226 = neura.grant_predicate %206, %219 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %227 = neura.grant_predicate %204, %219 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %228 = neura.grant_predicate %202, %219 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %229 = neura.grant_predicate %200, %219 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %230 = neura.grant_predicate %198, %219 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %231 = neura.grant_predicate %196, %219 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %232 = "neura.not"(%219) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %233 = neura.grant_predicate %212, %232 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %234 = neura.grant_predicate %208, %232 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %235 = neura.grant_predicate %216, %232 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %236 = neura.grant_predicate %206, %232 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %237 = neura.grant_predicate %204, %232 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %238 = neura.grant_predicate %214, %232 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %239 = neura.grant_predicate %202, %232 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %240 = neura.grant_predicate %200, %232 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %241 = neura.grant_predicate %198, %232 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %242 = neura.grant_predicate %196, %232 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %243 = "neura.add"(%233, %234) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %243 -> %36 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %235 -> %34 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %236 -> %32 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %237 -> %30 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %238 -> %28 : !neura.data<memref<?x32xi32>, i1> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %239 -> %26 : !neura.data<memref<?x32x32xi32>, i1> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %234 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %240 -> %22 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %241 -> %20 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %242 -> %18 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %244 = neura.load_indexed %220[%221, %222 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x32xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %245 = "neura.icmp"(%244, %223) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %246 = neura.grant_predicate %223, %245 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %247 = neura.grant_predicate %220, %245 : !neura.data<memref<?x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32xi32>, i1>
+// CTRL2DATA-NEXT:     %248 = neura.grant_predicate %221, %245 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %249 = neura.grant_predicate %222, %245 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %250 = neura.grant_predicate %224, %245 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %251 = neura.grant_predicate %225, %245 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %252 = neura.grant_predicate %226, %245 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %253 = neura.grant_predicate %227, %245 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %254 = neura.grant_predicate %228, %245 : !neura.data<memref<?x32x32xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x32x32xi32>, i1>
+// CTRL2DATA-NEXT:     %255 = neura.grant_predicate %229, %245 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %256 = neura.grant_predicate %230, %245 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %257 = neura.grant_predicate %231, %245 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %258 = "neura.not"(%245) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %259 = neura.grant_predicate %222, %258 : !neura.data<i64, i1
diff --git a/test/controflow_fuse/non_perfect_nested/non_perfect_nested.mlir b/test/controflow_fuse/non_perfect_nested/non_perfect_nested.mlir
index 760ae0ce..598ac289 100644
--- a/test/controflow_fuse/non_perfect_nested/non_perfect_nested.mlir
+++ b/test/controflow_fuse/non_perfect_nested/non_perfect_nested.mlir
@@ -15,6 +15,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -139,295 +140,267 @@ module attributes {} {
 // CHECK-NEXT:   }
 
 
-// CTRL2DATA:        func.func @_Z29non_perfect_extra_computationPA128_iS0_(%arg0: memref<?x128xi32>, %arg1: memref<?x128xi32>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
-// CTRL2DATA-NEXT:     %0 = "neura.constant"() <{value = "%arg0"}> : () -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %1 = "neura.grant_once"(%0) : (!neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %2 = "neura.constant"() <{value = "%arg1"}> : () -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %3 = "neura.grant_once"(%2) : (!neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %4 = "neura.constant"() <{value = 4 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %5 = "neura.grant_once"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %6 = "neura.constant"() <{value = 3 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %7 = "neura.grant_once"(%6) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %8 = "neura.constant"() <{value = 2 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %9 = "neura.grant_once"(%8) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %10 = "neura.constant"() <{value = 1 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %11 = "neura.grant_once"(%10) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %12 = "neura.constant"() <{value = 128 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %13 = "neura.grant_once"(%12) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %14 = "neura.constant"() <{value = 100 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %15 = "neura.grant_once"(%14) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %16 = "neura.constant"() <{value = 3 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %17 = "neura.grant_once"(%16) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %18 = "neura.constant"() <{value = 2 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %19 = "neura.grant_once"(%18) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %20 = "neura.constant"() <{value = 1000 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %21 = "neura.grant_once"(%20) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %22 = "neura.constant"() <{value = -1000 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %23 = "neura.grant_once"(%22) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %24 = "neura.constant"() <{value = 128 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %25 = "neura.grant_once"(%24) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %26 = "neura.constant"() <{value = 0 : i32}> : () -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %27 = "neura.grant_once"(%26) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %28 = "neura.constant"() <{value = 0 : i64}> : () -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %29 = "neura.grant_once"(%28) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %30 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %31 = neura.phi_start %5, %30 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %32 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %33 = neura.phi_start %7, %32 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %34 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %35 = neura.phi_start %9, %34 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %36 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %37 = neura.phi_start %15, %36 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %38 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %39 = neura.phi_start %25, %38 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %40 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %41 = neura.phi_start %11, %40 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %42 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %43 = neura.phi_start %3, %42 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %44 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %45 = neura.phi_start %1, %44 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %46 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %47 = neura.phi_start %23, %46 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %48 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %49 = neura.phi_start %21, %48 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %50 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %51 = neura.phi_start %29, %50 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %52 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %53 = neura.phi_start %17, %52 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %54 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %55 = neura.phi_start %27, %54 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %56 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %57 = neura.phi_start %19, %56 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %58 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %59 = neura.phi_start %13, %58 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %60 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %61 = neura.phi_start %29, %60 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %62 = "neura.icmp"(%61, %59) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %63 = neura.grant_predicate %61, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %64 = neura.grant_predicate %57, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %65 = neura.grant_predicate %55, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %66 = neura.grant_predicate %53, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %67 = neura.grant_predicate %51, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %68 = neura.grant_predicate %49, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %69 = neura.grant_predicate %47, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %70 = neura.grant_predicate %59, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %71 = neura.grant_predicate %45, %62 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %72 = neura.grant_predicate %43, %62 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %73 = neura.grant_predicate %41, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %74 = neura.grant_predicate %39, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %75 = neura.grant_predicate %37, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %76 = neura.grant_predicate %35, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %77 = neura.grant_predicate %33, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %78 = neura.grant_predicate %31, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %79 = "neura.cast"(%63) <{cast_type = "i64_to_i32"}> : (!neura.data<i64, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %80 = "neura.div"(%79, %64) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %81 = "neura.mul"(%64, %80) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %82 = "neura.sub"(%79, %81) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %83 = "neura.icmp"(%82, %65) <{cmpType = "eq"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %84 = "neura.sel"(%83, %64, %66) : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %85 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %86 = neura.phi_start %69, %85 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %87 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %88 = neura.phi_start %68, %87 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %89 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %90 = neura.phi_start %66, %89 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %91 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %92 = neura.phi_start %64, %91 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %93 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %94 = neura.phi_start %78, %93 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %95 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %96 = neura.phi_start %77, %95 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %97 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %98 = neura.phi_start %76, %97 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %99 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %100 = neura.phi_start %67, %99 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %101 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %102 = neura.phi_start %75, %101 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %103 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %104 = neura.phi_start %65, %103 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %105 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %106 = neura.phi_start %74, %105 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %107 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %108 = neura.phi_start %73, %107 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %109 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %110 = neura.phi_start %72, %109 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %111 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %112 = neura.phi_start %84, %111 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %113 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %114 = neura.phi_start %63, %113 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %115 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %116 = neura.phi_start %71, %115 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %117 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %118 = neura.phi_start %70, %117 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %119 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %120 = neura.phi_start %65, %119 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %121 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %122 = neura.phi_start %69, %121 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %123 = neura.reserve : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %124 = neura.phi_start %68, %123 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %125 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %126 = neura.phi_start %67, %125 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %127 = "neura.icmp"(%126, %118) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %128 = neura.grant_predicate %116, %127 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %129 = neura.grant_predicate %114, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %130 = neura.grant_predicate %126, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %131 = neura.grant_predicate %112, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %132 = neura.grant_predicate %110, %127 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %133 = neura.grant_predicate %120, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %134 = neura.grant_predicate %122, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %135 = neura.grant_predicate %124, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %136 = neura.grant_predicate %108, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %137 = neura.grant_predicate %118, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %138 = neura.grant_predicate %106, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %139 = neura.grant_predicate %104, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %140 = neura.grant_predicate %102, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %141 = neura.grant_predicate %100, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %142 = neura.grant_predicate %98, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %143 = neura.grant_predicate %96, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %144 = neura.grant_predicate %94, %127 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %145 = neura.grant_predicate %92, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %146 = neura.grant_predicate %90, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %147 = neura.grant_predicate %88, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %148 = neura.grant_predicate %86, %127 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %149 = "neura.not"(%127) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %150 = neura.grant_predicate %120, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %151 = neura.grant_predicate %106, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %152 = neura.grant_predicate %122, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %153 = neura.grant_predicate %124, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %154 = neura.grant_predicate %104, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %155 = neura.grant_predicate %102, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %156 = neura.grant_predicate %110, %149 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %157 = neura.grant_predicate %114, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %158 = neura.grant_predicate %100, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %159 = neura.grant_predicate %108, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %160 = neura.grant_predicate %98, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %161 = neura.grant_predicate %96, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %162 = neura.grant_predicate %94, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %163 = neura.grant_predicate %118, %149 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %164 = neura.grant_predicate %92, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %165 = neura.grant_predicate %90, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %166 = neura.grant_predicate %88, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %167 = neura.grant_predicate %86, %149 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %168 = neura.grant_predicate %116, %149 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %169 = "neura.div"(%150, %151) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %170 = "neura.sub"(%152, %153) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %171 = "neura.icmp"(%170, %154) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %172 = neura.grant_predicate %169, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %173 = neura.grant_predicate %155, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %174 = neura.grant_predicate %170, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %175 = neura.grant_predicate %156, %171 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %176 = neura.grant_predicate %157, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %177 = neura.grant_predicate %158, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %178 = neura.grant_predicate %152, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %179 = neura.grant_predicate %159, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %180 = neura.grant_predicate %153, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %181 = neura.grant_predicate %160, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %182 = neura.grant_predicate %161, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %183 = neura.grant_predicate %162, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %184 = neura.grant_predicate %163, %171 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %185 = neura.grant_predicate %164, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %186 = neura.grant_predicate %154, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %187 = neura.grant_predicate %165, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %188 = neura.grant_predicate %166, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %189 = neura.grant_predicate %167, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %190 = neura.grant_predicate %168, %171 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %191 = neura.grant_predicate %151, %171 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %192 = "neura.not"(%171) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %193 = neura.grant_predicate %169, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %194 = neura.grant_predicate %156, %192 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %195 = neura.grant_predicate %157, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %196 = neura.grant_predicate %158, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %197 = neura.grant_predicate %152, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %198 = neura.grant_predicate %159, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %199 = neura.grant_predicate %153, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %200 = neura.grant_predicate %160, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %201 = neura.grant_predicate %161, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %202 = neura.grant_predicate %170, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %203 = neura.grant_predicate %162, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %204 = neura.grant_predicate %163, %192 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %205 = neura.grant_predicate %164, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %206 = neura.grant_predicate %154, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %207 = neura.grant_predicate %165, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %208 = neura.grant_predicate %166, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %209 = neura.grant_predicate %167, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %210 = neura.grant_predicate %168, %192 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %211 = neura.grant_predicate %151, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %212 = neura.grant_predicate %155, %192 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %213 = "neura.mul"(%172, %173) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %214 = "neura.div"(%213, %174) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %215 = "neura.phi"(%173, %212) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %216 = "neura.phi"(%191, %211) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %217 = "neura.phi"(%190, %210) : (!neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %218 = "neura.phi"(%189, %209) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %219 = "neura.phi"(%188, %208) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %220 = "neura.phi"(%187, %207) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %221 = "neura.phi"(%186, %206) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %222 = "neura.phi"(%185, %205) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %223 = "neura.phi"(%184, %204) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %224 = "neura.phi"(%183, %203) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %225 = "neura.phi"(%174, %202) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %226 = "neura.phi"(%182, %201) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %227 = "neura.phi"(%181, %200) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %228 = "neura.phi"(%180, %199) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %229 = "neura.phi"(%179, %198) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %230 = "neura.phi"(%178, %197) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %231 = "neura.phi"(%177, %196) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %232 = "neura.phi"(%176, %195) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %233 = "neura.phi"(%175, %194) : (!neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     %234 = "neura.phi"(%172, %193) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %235 = "neura.phi"(%214, %193) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %234 to %233[%232, %231 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %230 to %233[%232, %229 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %228 to %233[%232, %227 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %235 to %233[%232, %226 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %225 to %233[%232, %224 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %236 = "neura.add"(%232, %229) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %236 -> %60 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %223 -> %58 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %222 -> %56 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %221 -> %54 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %220 -> %52 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %231 -> %50 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %219 -> %48 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %218 -> %46 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %217 -> %44 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %233 -> %42 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %229 -> %40 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %216 -> %38 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %215 -> %36 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %227 -> %34 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %226 -> %32 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %224 -> %30 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %237 = neura.load_indexed %128[%129, %130 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %238 = "neura.mul"(%237, %131) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %238 to %132[%129, %130 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %239 = neura.load_indexed %128[%129, %130 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %240 = "neura.add"(%133, %239) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %241 = "neura.icmp"(%239, %134) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %242 = "neura.sel"(%241, %239, %134) : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %243 = "neura.icmp"(%239, %135) <{cmpType = "slt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %244 = "neura.sel"(%243, %239, %135) : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %245 = "neura.add"(%130, %136) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %245 -> %125 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %244 -> %123 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %242 -> %121 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %240 -> %119 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %137 -> %117 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %128 -> %115 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %129 -> %113 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %131 -> %111 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %132 -> %109 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %136 -> %107 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %138 -> %105 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %139 -> %103 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %140 -> %101 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %141 -> %99 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %142 -> %97 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %143 -> %95 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %144 -> %93 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %145 -> %91 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %146 -> %89 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %147 -> %87 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %148 -> %85 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
-// CTRL2DATA-NEXT:   }
\ No newline at end of file
+// CTRL2DATA:      func.func @_Z29non_perfect_extra_computationPA128_iS0_(%arg0: memref<?x128xi32>, %arg1: memref<?x128xi32>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
+// CTRL2DATA-NEXT:   %0 = "neura.constant"() <{value = "%arg0"}> : () -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %1 = "neura.grant_once"(%0) : (!neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %2 = "neura.constant"() <{value = "%arg1"}> : () -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %3 = "neura.grant_once"(%2) : (!neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %4 = "neura.constant"() <{value = 4 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %5 = "neura.grant_once"(%4) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %6 = "neura.constant"() <{value = 3 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %7 = "neura.grant_once"(%6) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %8 = "neura.constant"() <{value = 2 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %9 = "neura.grant_once"(%8) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %10 = "neura.constant"() <{value = 1 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %11 = "neura.grant_once"(%10) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %12 = "neura.constant"() <{value = 128 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %13 = "neura.grant_once"(%12) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %14 = "neura.constant"() <{value = 100 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %15 = "neura.grant_once"(%14) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %16 = "neura.constant"() <{value = 3 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %17 = "neura.grant_once"(%16) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %18 = "neura.constant"() <{value = 2 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %19 = "neura.grant_once"(%18) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %20 = "neura.constant"() <{value = 1000 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %21 = "neura.grant_once"(%20) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %22 = "neura.constant"() <{value = -1000 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %23 = "neura.grant_once"(%22) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %24 = "neura.constant"() <{value = 128 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %25 = "neura.grant_once"(%24) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %26 = "neura.constant"() <{value = 0 : i32}> : () -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %27 = "neura.grant_once"(%26) : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %28 = "neura.constant"() <{value = 0 : i64}> : () -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %29 = "neura.grant_once"(%28) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %30 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %31 = neura.phi_start %5, %30 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %32 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %33 = neura.phi_start %7, %32 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %34 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %35 = neura.phi_start %9, %34 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %36 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %37 = neura.phi_start %15, %36 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %38 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %39 = neura.phi_start %25, %38 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %40 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %41 = neura.phi_start %11, %40 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %42 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %43 = neura.phi_start %3, %42 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %44 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %45 = neura.phi_start %1, %44 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %46 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %47 = neura.phi_start %23, %46 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %48 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %49 = neura.phi_start %21, %48 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %50 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %51 = neura.phi_start %29, %50 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %52 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %53 = neura.phi_start %17, %52 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %54 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %55 = neura.phi_start %27, %54 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %56 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %57 = neura.phi_start %19, %56 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %58 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %59 = neura.phi_start %13, %58 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %60 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %61 = neura.phi_start %29, %60 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %62 = "neura.icmp"(%61, %59) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %63 = "neura.not"(%62) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %64 = neura.grant_predicate %61, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %65 = neura.grant_predicate %57, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %66 = neura.grant_predicate %55, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %67 = neura.grant_predicate %53, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %68 = neura.grant_predicate %51, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %69 = neura.grant_predicate %49, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %70 = neura.grant_predicate %47, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %71 = neura.grant_predicate %59, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %72 = neura.grant_predicate %45, %62 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %73 = neura.grant_predicate %43, %62 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %74 = neura.grant_predicate %41, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %75 = neura.grant_predicate %39, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %76 = neura.grant_predicate %37, %62 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %77 = neura.grant_predicate %35, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %78 = neura.grant_predicate %33, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %79 = neura.grant_predicate %31, %62 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %80 = neura.grant_predicate %63, %63 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   neura.return_void %80 : !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %81 = "neura.cast"(%64) <{cast_type = "i64_to_i32"}> : (!neura.data<i64, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %82 = "neura.div"(%81, %65) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %83 = "neura.mul"(%65, %82) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %84 = "neura.sub"(%81, %83) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %85 = "neura.icmp"(%84, %66) <{cmpType = "eq"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %86 = "neura.sel"(%85, %65, %67) : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %87 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %88 = neura.phi_start %70, %87 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %89 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %90 = neura.phi_start %69, %89 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %91 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %92 = neura.phi_start %67, %91 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %93 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %94 = neura.phi_start %65, %93 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %95 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %96 = neura.phi_start %79, %95 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %97 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %98 = neura.phi_start %78, %97 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %99 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %100 = neura.phi_start %77, %99 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %101 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %102 = neura.phi_start %68, %101 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %103 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %104 = neura.phi_start %76, %103 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %105 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %106 = neura.phi_start %66, %105 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %107 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %108 = neura.phi_start %75, %107 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %109 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %110 = neura.phi_start %74, %109 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %111 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %112 = neura.phi_start %73, %111 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %113 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %114 = neura.phi_start %86, %113 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %115 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %116 = neura.phi_start %64, %115 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %117 = neura.reserve : !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %118 = neura.phi_start %72, %117 : !neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %119 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %120 = neura.phi_start %71, %119 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %121 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %122 = neura.phi_start %66, %121 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %123 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %124 = neura.phi_start %70, %123 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %125 = neura.reserve : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %126 = neura.phi_start %69, %125 : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %127 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %128 = neura.phi_start %68, %127 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %129 = "neura.icmp"(%128, %120) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %130 = neura.grant_predicate %118, %129 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %131 = neura.grant_predicate %116, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %132 = neura.grant_predicate %128, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %133 = neura.grant_predicate %114, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %134 = neura.grant_predicate %112, %129 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %135 = neura.grant_predicate %122, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %136 = neura.grant_predicate %124, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %137 = neura.grant_predicate %126, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %138 = neura.grant_predicate %110, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %139 = neura.grant_predicate %120, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %140 = neura.grant_predicate %108, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %141 = neura.grant_predicate %106, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %142 = neura.grant_predicate %104, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %143 = neura.grant_predicate %102, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %144 = neura.grant_predicate %100, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %145 = neura.grant_predicate %98, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %146 = neura.grant_predicate %96, %129 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %147 = neura.grant_predicate %94, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %148 = neura.grant_predicate %92, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %149 = neura.grant_predicate %90, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %150 = neura.grant_predicate %88, %129 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %151 = "neura.not"(%129) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %152 = neura.grant_predicate %122, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %153 = neura.grant_predicate %108, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %154 = neura.grant_predicate %124, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %155 = neura.grant_predicate %126, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %156 = neura.grant_predicate %106, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %157 = neura.grant_predicate %104, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %158 = neura.grant_predicate %112, %151 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %159 = neura.grant_predicate %116, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %160 = neura.grant_predicate %102, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %161 = neura.grant_predicate %110, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %162 = neura.grant_predicate %100, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %163 = neura.grant_predicate %98, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %164 = neura.grant_predicate %96, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %165 = neura.grant_predicate %120, %151 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %166 = neura.grant_predicate %94, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %167 = neura.grant_predicate %92, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %168 = neura.grant_predicate %90, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %169 = neura.grant_predicate %88, %151 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %170 = neura.grant_predicate %118, %151 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %171 = "neura.div"(%152, %153) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %172 = "neura.sub"(%154, %155) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %173 = "neura.icmp"(%172, %156) <{cmpType = "sgt"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %174 = neura.grant_predicate %171, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %175 = neura.grant_predicate %157, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %176 = neura.grant_predicate %172, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %177 = neura.grant_predicate %158, %173 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %178 = neura.grant_predicate %159, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %179 = neura.grant_predicate %160, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %180 = neura.grant_predicate %154, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %181 = neura.grant_predicate %161, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %182 = neura.grant_predicate %155, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %183 = neura.grant_predicate %162, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %184 = neura.grant_predicate %163, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %185 = neura.grant_predicate %164, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %186 = neura.grant_predicate %165, %173 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %187 = neura.grant_predicate %166, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %188 = neura.grant_predicate %156, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %189 = neura.grant_predicate %167, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %190 = neura.grant_predicate %168, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %191 = neura.grant_predicate %169, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %192 = neura.grant_predicate %170, %173 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %193 = neura.grant_predicate %153, %173 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %194 = "neura.not"(%173) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:   %195 = neura.grant_predicate %171, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %196 = neura.grant_predicate %158, %194 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %197 = neura.grant_predicate %159, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %198 = neura.grant_predicate %160, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %199 = neura.grant_predicate %154, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %200 = neura.grant_predicate %161, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %201 = neura.grant_predicate %155, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %202 = neura.grant_predicate %162, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %203 = neura.grant_predicate %163, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %204 = neura.grant_predicate %172, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %205 = neura.grant_predicate %164, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %206 = neura.grant_predicate %165, %194 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %207 = neura.grant_predicate %166, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %208 = neura.grant_predicate %156, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %209 = neura.grant_predicate %167, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %210 = neura.grant_predicate %168, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %211 = neura.grant_predicate %169, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %212 = neura.grant_predicate %170, %194 : !neura.data<memref<?x128xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %213 = neura.grant_predicate %153, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %214 = neura.grant_predicate %157, %194 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %215 = "neura.mul"(%174, %175) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %216 = "neura.div"(%215, %176) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %217 = "neura.phi"(%175, %214) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %218 = "neura.phi"(%193, %213) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %219 = "neura.phi"(%192, %212) : (!neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %220 = "neura.phi"(%191, %211) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %221 = "neura.phi"(%190, %210) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %222 = "neura.phi"(%189, %209) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %223 = "neura.phi"(%188, %208) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %224 = "neura.phi"(%187, %207) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %225 = "neura.phi"(%186, %206) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %226 = "neura.phi"(%185, %205) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %227 = "neura.phi"(%176, %204) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %228 = "neura.phi"(%184, %203) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %229 = "neura.phi"(%183, %202) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %230 = "neura.phi"(%182, %201) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %231 = "neura.phi"(%181, %200) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %232 = "neura.phi"(%180, %199) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %233 = "neura.phi"(%179, %198) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %234 = "neura.phi"(%178, %197) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %235 = "neura.phi"(%177, %196) : (!neura.data<memref<?x128xi32>, i1>, !neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   %236 = "neura.phi"(%174, %195) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %237 = "neura.phi"(%216, %195) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.store_indexed %236 to %235[%234, %233 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.store_indexed %232 to %235[%234, %231 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.store_indexed %230 to %235[%234, %229 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.store_indexed %237 to %235[%234, %228 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.store_indexed %227 to %235[%234, %226 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %238 = "neura.add"(%234, %231) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %238 -> %60 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %225 -> %58 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %224 -> %56 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %223 -> %54 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %222 -> %52 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %233 -> %50 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %221 -> %48 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %220 -> %46 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %219 -> %44 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %235 -> %42 : !neura.data<memref<?x128xi32>, i1> !neura.data<memref<?x128xi32>, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %231 -> %40 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %218 -> %38 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %217 -> %36 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %229 -> %34 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %228 -> %32 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   neura.ctrl_mov %226 -> %30 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:   %239 = neura.load_indexed %130[%131, %132 : !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x128xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:   %240 = "neura.mul"(%239, %133) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data
diff --git a/test/controflow_fuse/perfect_nested/perfect_nested.mlir b/test/controflow_fuse/perfect_nested/perfect_nested.mlir
index 22c6f982..0e76c4fe 100644
--- a/test/controflow_fuse/perfect_nested/perfect_nested.mlir
+++ b/test/controflow_fuse/perfect_nested/perfect_nested.mlir
@@ -24,6 +24,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -38,6 +39,7 @@
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -116,7 +118,7 @@ module attributes {} {
 // CAST-NEXT:   }
 
 
-// CTRL2DATA:        func.func @_Z10bert_node1PA1_A1_A1_A1_A128_bPA1_A128_S1_(%arg0: memref<?x1x1x1x1x128xi8>, %arg1: memref<?x1x128x1x1x128xi8>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
+// CTRL2DATA:      func.func @_Z10bert_node1PA1_A1_A1_A1_A128_bPA1_A128_S1_(%arg0: memref<?x1x1x1x1x128xi8>, %arg1: memref<?x1x128x1x1x128xi8>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
 // CTRL2DATA-NEXT:     %0 = "neura.constant"() <{value = "%arg0"}> : () -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
 // CTRL2DATA-NEXT:     %1 = "neura.grant_once"(%0) : (!neura.data<memref<?x1x1x1x1x128xi8>, i1>) -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
 // CTRL2DATA-NEXT:     %2 = "neura.constant"() <{value = "%arg1"}> : () -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
@@ -140,59 +142,62 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %20 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %21 = neura.phi_start %9, %20 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %22 = "neura.icmp"(%21, %19) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %23 = neura.grant_predicate %17, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %24 = neura.grant_predicate %19, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %25 = neura.grant_predicate %15, %22 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %26 = neura.grant_predicate %13, %22 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %27 = neura.grant_predicate %21, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %28 = neura.grant_predicate %11, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %29 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %30 = neura.phi_start %28, %29 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %23 = "neura.not"(%22) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %24 = neura.grant_predicate %17, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %25 = neura.grant_predicate %19, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %26 = neura.grant_predicate %15, %22 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %27 = neura.grant_predicate %13, %22 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %28 = neura.grant_predicate %21, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %29 = neura.grant_predicate %11, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %30 = neura.grant_predicate %23, %23 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %30 : !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %31 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %32 = neura.phi_start %27, %31 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %33 = neura.reserve : !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %34 = neura.phi_start %26, %33 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<memref<?x1x128x1x1x128xi8>, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %35 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %36 = neura.phi_start %23, %35 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %37 = neura.reserve : !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %38 = neura.phi_start %25, %37 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<memref<?x1x1x1x1x128xi8>, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %39 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %40 = neura.phi_start %24, %39 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %32 = neura.phi_start %29, %31 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %33 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %34 = neura.phi_start %28, %33 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %35 = neura.reserve : !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %36 = neura.phi_start %27, %35 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<memref<?x1x128x1x1x128xi8>, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %37 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %38 = neura.phi_start %24, %37 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %39 = neura.reserve : !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %40 = neura.phi_start %26, %39 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<memref<?x1x1x1x1x128xi8>, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
 // CTRL2DATA-NEXT:     %41 = neura.reserve : !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %42 = neura.phi_start %23, %41 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %43 = "neura.icmp"(%42, %40) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %44 = neura.grant_predicate %38, %43 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %45 = neura.grant_predicate %36, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %42, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %34, %43 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %32, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %49 = neura.grant_predicate %30, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %50 = neura.grant_predicate %40, %43 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %51 = "neura.not"(%43) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %52 = neura.grant_predicate %32, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %53 = neura.grant_predicate %30, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %54 = neura.grant_predicate %40, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %55 = neura.grant_predicate %36, %51 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %56 = neura.grant_predicate %38, %51 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %57 = neura.grant_predicate %34, %51 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     %58 = "neura.add"(%52, %53) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %58 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %54 -> %18 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %55 -> %16 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %56 -> %14 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %57 -> %12 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %53 -> %10 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %59 = neura.load_indexed %44[%45, %45, %45, %45, %45, %46 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x1x1x1x128xi8>, i1> : !neura.data<i8, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %59 to %47[%45, %45, %48, %45, %45, %46 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x128x1x1x128xi8>, i1> : !neura.data<i8, i1>
-// CTRL2DATA-NEXT:     %60 = "neura.add"(%46, %49) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %60 -> %41 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %50 -> %39 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %44 -> %37 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %45 -> %35 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %47 -> %33 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %48 -> %31 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %49 -> %29 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
+// CTRL2DATA-NEXT:     %42 = neura.phi_start %25, %41 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %43 = neura.reserve : !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %44 = neura.phi_start %24, %43 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %45 = "neura.icmp"(%44, %42) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %46 = neura.grant_predicate %40, %45 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %47 = neura.grant_predicate %38, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %48 = neura.grant_predicate %44, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %49 = neura.grant_predicate %36, %45 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %50 = neura.grant_predicate %34, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %51 = neura.grant_predicate %32, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %52 = neura.grant_predicate %42, %45 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %53 = "neura.not"(%45) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %54 = neura.grant_predicate %34, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %55 = neura.grant_predicate %32, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %56 = neura.grant_predicate %42, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %57 = neura.grant_predicate %38, %53 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %58 = neura.grant_predicate %40, %53 : !neura.data<memref<?x1x1x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %59 = neura.grant_predicate %36, %53 : !neura.data<memref<?x1x128x1x1x128xi8>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     %60 = "neura.add"(%54, %55) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %60 -> %20 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %56 -> %18 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %57 -> %16 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %58 -> %14 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %59 -> %12 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %55 -> %10 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %61 = neura.load_indexed %46[%47, %47, %47, %47, %47, %48 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x1x1x1x128xi8>, i1> : !neura.data<i8, i1>
+// CTRL2DATA-NEXT:     neura.store_indexed %61 to %49[%47, %47, %50, %47, %47, %48 : !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>] !neura.data<memref<?x1x128x1x1x128xi8>, i1> : !neura.data<i8, i1>
+// CTRL2DATA-NEXT:     %62 = "neura.add"(%48, %51) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %62 -> %43 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %52 -> %41 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %46 -> %39 : !neura.data<memref<?x1x1x1x1x128xi8>, i1> !neura.data<memref<?x1x1x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %47 -> %37 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %49 -> %35 : !neura.data<memref<?x1x128x1x1x128xi8>, i1> !neura.data<memref<?x1x128x1x1x128xi8>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %50 -> %33 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %51 -> %31 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 // MAPPING:      func.func @_Z10bert_node1PA1_A1_A1_A1_A128_bPA1_A128_S1_(%arg0: memref<?x1x1x1x1x128xi8>, %arg1: memref<?x1x128x1x1x128xi8>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 10 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 8 : i32, res_mii = 2 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
\ No newline at end of file
diff --git a/test/controflow_fuse/perfect_reduction/perfect_reduction.mlir b/test/controflow_fuse/perfect_reduction/perfect_reduction.mlir
index 7917c16c..77c41c7b 100644
--- a/test/controflow_fuse/perfect_reduction/perfect_reduction.mlir
+++ b/test/controflow_fuse/perfect_reduction/perfect_reduction.mlir
@@ -25,6 +25,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -107,7 +108,7 @@ module attributes {} {
 // CAST-NEXT:   }
 
 
-// CTRL2DATA:   func.func @_Z27perfect_nested_reduction_2dPA128_i(%arg0: memref<?x128xi32>) -> i32 attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
+// CTRL2DATA:        func.func @_Z27perfect_nested_reduction_2dPA128_i(%arg0: memref<?x128xi32>) -> i32 attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
 // CTRL2DATA-NEXT:     %0 = "neura.constant"() <{value = "%arg0"}> : () -> !neura.data<memref<?x128xi32>, i1>
 // CTRL2DATA-NEXT:     %1 = "neura.grant_once"(%0) : (!neura.data<memref<?x128xi32>, i1>) -> !neura.data<memref<?x128xi32>, i1>
 // CTRL2DATA-NEXT:     %2 = "neura.constant"() <{value = 1 : i64}> : () -> !neura.data<i64, i1>
@@ -139,6 +140,7 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %28 = neura.grant_predicate %11, %22 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %29 = "neura.not"(%22) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %30 = neura.grant_predicate %19, %29 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.return_value %30 : !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %31 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %32 = neura.phi_start %23, %31 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %33 = neura.reserve : !neura.data<i64, i1>
@@ -185,5 +187,5 @@ module attributes {} {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %47 -> %35 : !neura.data<i64, i1> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %50 -> %33 : !neura.data<i64, i1> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %52 -> %31 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%30) : (!neura.data<i32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/controflow_fuse/simple_loop/simple_loop.mlir b/test/controflow_fuse/simple_loop/simple_loop.mlir
index 7d191587..13b2e91f 100644
--- a/test/controflow_fuse/simple_loop/simple_loop.mlir
+++ b/test/controflow_fuse/simple_loop/simple_loop.mlir
@@ -29,6 +29,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -43,6 +44,7 @@
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -60,6 +62,7 @@
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -69,7 +72,8 @@
 // RUN: --insert-data-mov \
 // RUN: --map-to-accelerator="mapping-strategy=heuristic backtrack-config=customized" \
 // RUN: --architecture-spec=../../arch_spec/architecture.yaml \
-// RUN: | FileCheck %s --check-prefix=FUSE-MAPPING
+// RUN: -o %t-fused-mapped.mlir
+// RUN: FileCheck %s --input-file=%t-fused-mapped.mlir --check-prefix=FUSE-MAPPING
 
 module attributes {} {
   func.func @_Z11simple_loopPiS_(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
@@ -163,37 +167,46 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %26 = neura.reserve : !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %27 = neura.phi_start %13, %26 : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %28 = "neura.icmp"(%27, %25) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
-// CTRL2DATA-NEXT:     %29 = neura.grant_predicate %23, %28 : !neura.data<memref<?xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?xi32>, i1>
-// CTRL2DATA-NEXT:     %30 = neura.grant_predicate %27, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %31 = neura.grant_predicate %21, %28 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %32 = neura.grant_predicate %19, %28 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %33 = neura.grant_predicate %17, %28 : !neura.data<memref<?xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?xi32>, i1>
-// CTRL2DATA-NEXT:     %34 = neura.grant_predicate %15, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %35 = neura.grant_predicate %25, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     %36 = neura.load_indexed %29[%30 : !neura.data<i64, i1>] !neura.data<memref<?xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %37 = "neura.mul"(%36, %31) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %38 = "neura.add"(%37, %32) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.store_indexed %38 to %33[%30 : !neura.data<i64, i1>] !neura.data<memref<?xi32>, i1> : !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     %39 = "neura.add"(%30, %34) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %39 -> %26 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %35 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %29 -> %22 : !neura.data<memref<?xi32>, i1> !neura.data<memref<?xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %31 -> %20 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %32 -> %18 : !neura.data<i32, i1> !neura.data<i32, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %33 -> %16 : !neura.data<memref<?xi32>, i1> !neura.data<memref<?xi32>, i1>
-// CTRL2DATA-NEXT:     neura.ctrl_mov %34 -> %14 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
+// CTRL2DATA-NEXT:     %29 = "neura.not"(%28) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %30 = neura.grant_predicate %23, %28 : !neura.data<memref<?xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?xi32>, i1>
+// CTRL2DATA-NEXT:     %31 = neura.grant_predicate %27, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %32 = neura.grant_predicate %21, %28 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %33 = neura.grant_predicate %19, %28 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %34 = neura.grant_predicate %17, %28 : !neura.data<memref<?xi32>, i1>, !neura.data<i1, i1> -> !neura.data<memref<?xi32>, i1>
+// CTRL2DATA-NEXT:     %35 = neura.grant_predicate %15, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %36 = neura.grant_predicate %25, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     %37 = neura.grant_predicate %29, %29 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %37 : !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     %38 = neura.load_indexed %30[%31 : !neura.data<i64, i1>] !neura.data<memref<?xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %39 = "neura.mul"(%38, %32) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %40 = "neura.add"(%39, %33) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.store_indexed %40 to %34[%31 : !neura.data<i64, i1>] !neura.data<memref<?xi32>, i1> : !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     %41 = "neura.add"(%31, %35) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %41 -> %26 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %36 -> %24 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %30 -> %22 : !neura.data<memref<?xi32>, i1> !neura.data<memref<?xi32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %32 -> %20 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %33 -> %18 : !neura.data<i32, i1> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %34 -> %16 : !neura.data<memref<?xi32>, i1> !neura.data<memref<?xi32>, i1>
+// CTRL2DATA-NEXT:     neura.ctrl_mov %35 -> %14 : !neura.data<i64, i1> !neura.data<i64, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 
 // FUSE:        func.func @_Z11simple_loopPiS_(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
-// FUSE-NEXT:     %0 = "neura.grant_always"() <{constant_value = true}> : () -> !neura.data<i1, i1>
-// FUSE-NEXT:     %nextindex, %valid = "neura.loop_control"(%0) <{end = 128 : i64, iterationType = "increment", start = 0 : i64, step = 1 : i64}> : (!neura.data<i1, i1>) -> (!neura.data<i64, i1>, !neura.data<i1, i1>)
-// FUSE-NEXT:     %1 = neura.load_indexed [%nextindex : !neura.data<i64, i1>]  {lhs_value = "%arg0"} : !neura.data<i32, i1>
-// FUSE-NEXT:     %2 = "neura.mul"(%1) {rhs_value = 2 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// FUSE-NEXT:     %3 = "neura.add"(%2) {rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// FUSE-NEXT:     neura.store_indexed %3 to [%nextindex : !neura.data<i64, i1>]  {rhs_value = "%arg1"} : !neura.data<i32, i1>
-// FUSE-NEXT:     "neura.return"() : () -> ()
+// FUSE-NEXT:    %0 = "neura.grant_always"() <{constant_value = true}> : () -> !neura.data<i1, i1>
+// FUSE-NEXT:    %nextindex, %valid = "neura.loop_control"(%0) <{end = 128 : i64, iterationType = "increment", start = 0 : i64, step = 1 : i64}> : (!neura.data<i1, i1>) -> (!neura.data<i64, i1>, !neura.data<i1, i1>)
+// FUSE-NEXT:    %1 = "neura.not"(%valid) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// FUSE-NEXT:    %2 = neura.grant_predicate %1, %1 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// FUSE-NEXT:    neura.return_void %2 : !neura.data<i1, i1>
+// FUSE-NEXT:    %3 = neura.load_indexed [%nextindex : !neura.data<i64, i1>]  {lhs_value = "%arg0"} : !neura.data<i32, i1>
+// FUSE-NEXT:    %4 = "neura.mul"(%3) {rhs_value = 2 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// FUSE-NEXT:    %5 = "neura.add"(%4) {rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// FUSE-NEXT:    neura.store_indexed %5 to [%nextindex : !neura.data<i64, i1>]  {rhs_value = "%arg1"} : !neura.data<i32, i1>
+// FUSE-NEXT:    neura.yield
 // FUSE-NEXT:   }
 
-// FUSE-MAPPING:  func.func @_Z11simple_loopPiS_(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
\ No newline at end of file
+// FUSE-MAPPING:  func.func @_Z11simple_loopPiS_
+// FUSE-MAPPING-SAME: accelerator = "neura"
+// FUSE-MAPPING-SAME: dataflow_mode = "predicate"
+// FUSE-MAPPING-SAME: mapping_info = {compiled_ii = 2 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}
\ No newline at end of file
diff --git a/test/controflow_fuse/simple_loop_reduction/simple_loop_reduction.mlir b/test/controflow_fuse/simple_loop_reduction/simple_loop_reduction.mlir
index be127d82..5eabc8a7 100644
--- a/test/controflow_fuse/simple_loop_reduction/simple_loop_reduction.mlir
+++ b/test/controflow_fuse/simple_loop_reduction/simple_loop_reduction.mlir
@@ -30,6 +30,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -44,6 +45,7 @@
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -61,6 +63,7 @@
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -122,7 +125,7 @@ module attributes {} {
 // CANONICALIZE-NEXT:     "neura.return"(%16) : (i32) -> ()
 // CANONICALIZE-NEXT:   }
 
-// CTRL2DATA:      func.func @_Z10simpleloopv() -> i32 attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
+// CTRL2DATA:        func.func @_Z10simpleloopv() -> i32 attributes {accelerator = "neura", dataflow_mode = "predicate", llvm.linkage = #llvm.linkage<external>} {
 // CTRL2DATA-NEXT:     %0 = "neura.constant"() <{value = 1 : i64}> : () -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %1 = "neura.grant_once"(%0) : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %2 = "neura.constant"() <{value = 128 : i64}> : () -> !neura.data<i64, i1>
@@ -146,6 +149,7 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %20 = neura.grant_predicate %11, %16 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %21 = "neura.not"(%16) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %22 = neura.grant_predicate %13, %21 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.return_value %22 : !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %23 = "neura.cast"(%17) <{cast_type = "i64_to_i32"}> : (!neura.data<i64, i1>) -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %24 = "neura.add"(%18, %23) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %25 = "neura.add"(%17, %19) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
@@ -153,7 +157,7 @@ module attributes {} {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %24 -> %12 : !neura.data<i32, i1> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %20 -> %10 : !neura.data<i64, i1> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %19 -> %8 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%22) : (!neura.data<i32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 
@@ -166,10 +170,11 @@ module attributes {} {
 // FUSE-NEXT:     %4 = "neura.not"(%valid) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // FUSE-NEXT:     %5 = neura.grant_predicate %2, %valid : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
 // FUSE-NEXT:     %6 = neura.grant_predicate %2, %4 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// FUSE-NEXT:     neura.return_value %6 : !neura.data<i32, i1>
 // FUSE-NEXT:     %7 = "neura.cast"(%nextindex) <{cast_type = "i64_to_i32"}> : (!neura.data<i64, i1>) -> !neura.data<i32, i1>
 // FUSE-NEXT:     %8 = "neura.add"(%5, %7) : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
 // FUSE-NEXT:     neura.ctrl_mov %8 -> %1 : !neura.data<i32, i1> !neura.data<i32, i1>
-// FUSE-NEXT:     "neura.return"(%6) : (!neura.data<i32, i1>) -> ()
+// FUSE-NEXT:     neura.yield
 // FUSE-NEXT:   }
 
 
diff --git a/test/e2e/bicg/bicg_kernel.mlir b/test/e2e/bicg/bicg_kernel.mlir
index fe2b0f9b..46c78188 100644
--- a/test/e2e/bicg/bicg_kernel.mlir
+++ b/test/e2e/bicg/bicg_kernel.mlir
@@ -32,6 +32,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -161,135 +162,91 @@
 
 //MAPPING: func.func @kernel
 //MAPPING-SAME: accelerator = "neura", dataflow_mode = "predicate"
-//MAPPING-SAME: mapping_info = {compiled_ii = 11 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 9 : i32, res_mii = 5 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}
+//MAPPING-SAME: mapping_info = {compiled_ii = 13 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 9 : i32, res_mii = 6 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}
 
-// YAML: array_config:
-// YAML:   columns: 4
-// YAML:   rows: 4
-// YAML:   compiled_ii: 11
-// YAML:   cores:
-// YAML:     - column: 0
-// YAML:       row: 0
-// YAML:       core_id: "0"
-// YAML:       entries:
-// YAML:         - entry_id: "entry0"
-// YAML:           instructions:
-// YAML:             - index_per_ii: 0
-// YAML:               operations:
-// YAML:                 - opcode: "CONSTANT"
-// YAML:                   id: 1
-// YAML:                   time_step: 0
-// YAML:                   invalid_iterations: 0
-// YAML:                   src_operands:
-// YAML:                     - operand: "arg0"
-// YAML:                       color: "RED"
-// YAML:                   dst_operands:
-// YAML:                     - operand: "$0"
-// YAML:                       color: "RED"
-// YAML:             - index_per_ii: 1
-// YAML:               operations:
-// YAML:                 - opcode: "ICMP_SGT"
-// YAML:                   id: 24
-// YAML:                   time_step: 1
-// YAML:                   invalid_iterations: 0
-// YAML:                   src_operands:
-// YAML:                     - operand: "$0"
-// YAML:                       color: "RED"
-// YAML:                     - operand: "#0"
-// YAML:                       color: "RED"
-// YAML:                   dst_operands:
-// YAML:                     - operand: "$0"
-// YAML:                       color: "RED"
-// YAML:             - index_per_ii: 2
-// YAML:               operations:
-// YAML:                 - opcode: "GRANT_ONCE"
-// YAML:                   id: 2
-// YAML:                   time_step: 2
-// YAML:                   invalid_iterations: 0
-// YAML:                   src_operands:
-// YAML:                     - operand: "arg1"
-// YAML:                       color: "RED"
-// YAML:                   dst_operands:
-// YAML:                     - operand: "$1"
-// YAML:                       color: "RED"
-// YAML:                     - operand: "$3"
+// YAML:      array_config:
+// YAML-NEXT:   columns: 4
+// YAML-NEXT:   rows: 4
+// YAML-NEXT:   compiled_ii: 13
+// YAML-NEXT:   cores:
+// YAML-NEXT:     - column: 0
+// YAML-NEXT:       row: 0
+// YAML-NEXT:       core_id: "0"
+// YAML-NEXT:       entries:
+// YAML-NEXT:         - entry_id: "entry0"
+// YAML-NEXT:           instructions:
+// YAML-NEXT:             - index_per_ii: 0
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "CONSTANT"
+// YAML-NEXT:                   id: 1
+// YAML-NEXT:                   time_step: 0
+// YAML-NEXT:                   invalid_iterations: 0
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "arg0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:             - index_per_ii: 1
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "ICMP_SGT"
+// YAML-NEXT:                   id: 24
+// YAML-NEXT:                   time_step: 1
+// YAML-NEXT:                   invalid_iterations: 0
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "$0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                     - operand: "#0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "NORTH"
+// YAML-NEXT:                       color: "RED"
 
 
-// ASM:      # Compiled II: 11
+// ASM:      # Compiled II: 13
 // ASM:      PE(0,0):
 // ASM-NEXT: {
 // ASM-NEXT:   CONSTANT, [arg0] -> [$0] (t=0, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=0)
 // ASM-NEXT: {
-// ASM-NEXT:   ICMP_SGT, [$0], [#0] -> [$0] (t=1, inv_iters=0)
+// ASM-NEXT:   ICMP_SGT, [$0], [#0] -> [NORTH, RED] (t=1, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=1)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_ONCE, [arg1] -> [$1], [$3] (t=2, inv_iters=0)
+// ASM-NEXT:   GRANT_ONCE, [arg1] -> [NORTH, RED], [$0] (t=2, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=2)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_ONCE, [$0] -> [EAST, RED], [NORTH, RED], [$0], [$2] (t=3, inv_iters=0)
+// ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [EAST, RED] (t=3, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=3)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$1], [$0] -> [$0], [$1] (t=4, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=4)
-// ASM-NEXT: {
-// ASM-NEXT:   ICMP_SGT, [$0], [#0] -> [$0], [EAST, RED], [NORTH, RED] (t=5, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=5)
-// ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [NORTH, RED], [EAST, RED] -> [NORTH, RED] (t=17, inv_iters=1)
-// ASM-NEXT: } (idx_per_ii=6)
-// ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$1], [$0] -> [EAST, RED] (t=7, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=7)
-// ASM-NEXT: {
-// ASM-NEXT:   NOT, [$2] -> [$0], [EAST, RED] (t=8, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=8)
+// ASM-NEXT:   GRANT_PREDICATE, [$0], [NORTH, RED] -> [$0], [$1] (t=10, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=10)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$3], [$0] -> [$0], [NORTH, RED] (t=9, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=9)
+// ASM-NEXT:   ICMP_SGT, [$0], [#0] -> [EAST, RED], [$0], [NORTH, RED] (t=11, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=11)
 // ASM-NEXT: {
-// ASM-NEXT:   ICMP_SGT, [$0], [#0] -> [NORTH, RED], [EAST, RED] (t=10, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=10)
+// ASM-NEXT:   GRANT_PREDICATE, [$1], [$0] -> [EAST, RED] (t=12, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=12)
 // ASM:      PE(1,0):
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_ONCE, [arg4] -> [$2] (t=11, inv_iters=1)
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [EAST, RED] (t=11, inv_iters=1)
-// ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [$1] (t=11, inv_iters=1)
+// ASM-NEXT:   ZEXT, [WEST, RED] -> [EAST, RED] (t=13, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=0)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [EAST, RED], [$5] (t=12, inv_iters=1)
+// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [NORTH, RED] (t=14, inv_iters=1)
+// ASM-NEXT:   DATA_MOV, [EAST, RED] -> [$0] (t=14, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=1)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$2], [$3] -> [$0] (t=13, inv_iters=1)
-// ASM-NEXT:   DATA_MOV, [EAST, RED] -> [$2] (t=13, inv_iters=1)
+// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [EAST, RED] (t=15, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=2)
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [NORTH, RED] (t=14, inv_iters=1)
-// ASM-NEXT: } (idx_per_ii=3)
-// ASM-NEXT: {
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [NORTH, RED] (t=4, inv_iters=0)
-// ASM-NEXT:   GRANT_PREDICATE, [$4], [$2] -> [$1] (t=15, inv_iters=1)
+// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [EAST, RED] (t=4, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=4)
 // ASM-NEXT: {
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [EAST, RED] (t=5, inv_iters=0)
-// ASM-NEXT:   GRANT_PREDICATE, [$5], [NORTH, RED] -> [$2] (t=16, inv_iters=1)
-// ASM-NEXT:   DATA_MOV, [EAST, RED] -> [WEST, RED] (t=16, inv_iters=1)
-// ASM-NEXT: } (idx_per_ii=5)
-// ASM-NEXT: {
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [NORTH, RED] (t=6, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=6)
+// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [$1] (t=11, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=11)
 // ASM-NEXT: {
-// ASM-NEXT:   ZEXT, [WEST, RED] -> [$0] (t=8, inv_iters=0)
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [$3] (t=8, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=8)
-// ASM-NEXT: {
-// ASM-NEXT:   PHI_START, [$0], [$2] -> [$0] (t=9, inv_iters=0)
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [EAST, RED] (t=9, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=9)
-// ASM-NEXT: {
-// ASM-NEXT:   PHI_START, [$0], [$1] -> [$4], [$0] (t=10, inv_iters=0)
-// ASM-NEXT:   DATA_MOV, [WEST, RED] -> [$1] (t=10, inv_iters=0)
-// ASM-NEXT: } (idx_per_ii=10)
+// ASM-NEXT:   NOT, [WEST, RED] -> [$0], [$1] (t=12, inv_iters=0)
+// ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [EAST, RED] (t=12, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=12)
 
 
 // RUN: mlir-neura-opt %t-kernel.mlir \
@@ -297,6 +254,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/e2e/fir/fir_kernel.mlir b/test/e2e/fir/fir_kernel.mlir
index ab9ecd9f..8377e386 100644
--- a/test/e2e/fir/fir_kernel.mlir
+++ b/test/e2e/fir/fir_kernel.mlir
@@ -8,6 +8,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -31,44 +32,45 @@
 // MAPPING-NEXT:     %0 = "neura.grant_once"() <{constant_value = 0 : i64}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i64, i1>
 // MAPPING-NEXT:     %1 = "neura.grant_once"() <{constant_value = 0 : i32}> {dfg_id = 1 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 1 : i32}]} : () -> !neura.data<i32, i1>
 // MAPPING-NEXT:     %2 = neura.reserve {dfg_id = 2 : i32} : !neura.data<i32, i1>
-// MAPPING-NEXT:     %3 = "neura.data_mov"(%1) {dfg_id = 5 : i32, mapping_locs = [{id = 128 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %4 = neura.phi_start %3, %2 {dfg_id = 7 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 0 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %3 = "neura.data_mov"(%1) {dfg_id = 6 : i32, mapping_locs = [{id = 128 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %4 = neura.phi_start %3, %2 {dfg_id = 8 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 0 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
 // MAPPING-NEXT:     %5 = neura.reserve {dfg_id = 3 : i32} : !neura.data<i64, i1>
-// MAPPING-NEXT:     %6 = "neura.data_mov"(%0) {dfg_id = 4 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %7 = neura.phi_start %6, %5 {dfg_id = 6 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %8 = "neura.data_mov"(%7) {dfg_id = 10 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %9 = "neura.gep"(%8) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 14 : i32, lhs_value = "%arg0", mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %10 = "neura.data_mov"(%9) {dfg_id = 18 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %11 = "neura.load"(%10) {dfg_id = 21 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %12 = "neura.data_mov"(%7) {dfg_id = 9 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %13 = "neura.gep"(%12) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 13 : i32, lhs_value = "%arg2", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %14 = "neura.data_mov"(%13) {dfg_id = 17 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %15 = "neura.load"(%14) {dfg_id = 20 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %16 = "neura.data_mov"(%15) {dfg_id = 24 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %17 = "neura.data_mov"(%11) {dfg_id = 25 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %18 = "neura.mul"(%16, %17) {dfg_id = 27 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %19 = "neura.data_mov"(%18) {dfg_id = 30 : i32, mapping_locs = [{id = 43 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %20 = "neura.data_mov"(%4) {dfg_id = 11 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %21 = "neura.add"(%19, %20) {dfg_id = 32 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %22 = "neura.data_mov"(%7) {dfg_id = 8 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %23 = "neura.add"(%22) {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %24 = "neura.data_mov"(%23) {dfg_id = 16 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %25 = "neura.icmp"(%24) <{cmpType = "eq"}> {dfg_id = 19 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 32 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %26 = "neura.data_mov"(%25) {dfg_id = 23 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %27 = "neura.not"(%26) {dfg_id = 26 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %28 = "neura.data_mov"(%23) {dfg_id = 15 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %29 = "neura.data_mov"(%27) {dfg_id = 29 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %30 = neura.grant_predicate %28, %29 {dfg_id = 31 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:     neura.ctrl_mov %30 -> %5 {dfg_id = 33 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
-// MAPPING-NEXT:     %31 = "neura.data_mov"(%21) {dfg_id = 35 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %32 = "neura.data_mov"(%27) {dfg_id = 28 : i32, mapping_locs = [{id = 31 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 160 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %33 = neura.grant_predicate %31, %32 {dfg_id = 37 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:     neura.ctrl_mov %33 -> %2 {dfg_id = 39 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 7 : i32}, {id = 129 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
-// MAPPING-NEXT:     %34 = "neura.data_mov"(%21) {dfg_id = 34 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %35 = "neura.data_mov"(%25) {dfg_id = 22 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %36 = neura.grant_predicate %34, %35 {dfg_id = 36 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %37 = "neura.data_mov"(%36) {dfg_id = 38 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     "neura.return"(%37) {dfg_id = 40 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>) -> ()
+// MAPPING-NEXT:     %6 = "neura.data_mov"(%0) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %7 = neura.phi_start %6, %5 {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %8 = "neura.data_mov"(%7) {dfg_id = 11 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %9 = "neura.gep"(%8) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 15 : i32, lhs_value = "%arg0", mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %10 = "neura.data_mov"(%9) {dfg_id = 19 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %11 = "neura.load"(%10) {dfg_id = 22 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %12 = "neura.data_mov"(%7) {dfg_id = 10 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %13 = "neura.gep"(%12) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 14 : i32, lhs_value = "%arg2", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %14 = "neura.data_mov"(%13) {dfg_id = 18 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %15 = "neura.load"(%14) {dfg_id = 21 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %16 = "neura.data_mov"(%15) {dfg_id = 25 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %17 = "neura.data_mov"(%11) {dfg_id = 26 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %18 = "neura.mul"(%16, %17) {dfg_id = 28 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %19 = "neura.data_mov"(%18) {dfg_id = 31 : i32, mapping_locs = [{id = 43 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %20 = "neura.data_mov"(%4) {dfg_id = 12 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %21 = "neura.add"(%19, %20) {dfg_id = 33 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %22 = "neura.data_mov"(%7) {dfg_id = 9 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %23 = "neura.add"(%22) {dfg_id = 13 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %24 = "neura.data_mov"(%23) {dfg_id = 17 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %25 = "neura.icmp"(%24) <{cmpType = "eq"}> {dfg_id = 20 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 32 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %26 = "neura.data_mov"(%25) {dfg_id = 24 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %27 = "neura.not"(%26) {dfg_id = 27 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %28 = "neura.data_mov"(%23) {dfg_id = 16 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %29 = "neura.data_mov"(%27) {dfg_id = 30 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %30 = neura.grant_predicate %28, %29 {dfg_id = 32 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     neura.ctrl_mov %30 -> %5 {dfg_id = 34 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
+// MAPPING-NEXT:     %31 = "neura.data_mov"(%21) {dfg_id = 36 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %32 = "neura.data_mov"(%27) {dfg_id = 29 : i32, mapping_locs = [{id = 31 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 160 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %33 = neura.grant_predicate %31, %32 {dfg_id = 38 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     neura.ctrl_mov %33 -> %2 {dfg_id = 40 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 7 : i32}, {id = 129 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
+// MAPPING-NEXT:     %34 = "neura.data_mov"(%21) {dfg_id = 35 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %35 = "neura.data_mov"(%25) {dfg_id = 23 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %36 = neura.grant_predicate %34, %35 {dfg_id = 37 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %37 = "neura.data_mov"(%36) {dfg_id = 39 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     neura.return_value %37 : !neura.data<i32, i1> {dfg_id = 41 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 2 : i32}]}
+// MAPPING-NEXT:     neura.yield {dfg_id = 4 : i32}
 // MAPPING-NEXT:   }
 
 // YAML:      array_config:
@@ -95,7 +97,7 @@
 // YAML-NEXT:                     - operand: "$0"
 // YAML-NEXT:                       color: "RED"
 // YAML-NEXT:                 - opcode: "CTRL_MOV"
-// YAML-NEXT:                   id: 39
+// YAML-NEXT:                   id: 40
 // YAML-NEXT:                   time_step: 8
 // YAML-NEXT:                   invalid_iterations: 1
 // YAML-NEXT:                   src_operands:
@@ -107,7 +109,7 @@
 // YAML-NEXT:             - index_per_ii: 4
 // YAML-NEXT:               operations:
 // YAML-NEXT:                 - opcode: "PHI_START"
-// YAML-NEXT:                   id: 7
+// YAML-NEXT:                   id: 8
 // YAML-NEXT:                   time_step: 4
 // YAML-NEXT:                   invalid_iterations: 0
 // YAML-NEXT:                   src_operands:
@@ -151,7 +153,7 @@
 // ASM-NEXT: } (idx_per_ii=2)
 // ASM-NEXT: {
 // ASM-NEXT:   DATA_MOV, [EAST, RED] -> [$1] (t=3, inv_iters=0)
-// ASM-NEXT:   RETURN, [$0] (t=8, inv_iters=1)
+// ASM-NEXT:   RETURN_VALUE, [$0] (t=8, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=3)
 // ASM:      PE(2,2):
 // ASM-NEXT: {
@@ -174,6 +176,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/e2e/fir/fir_kernel_vec.mlir b/test/e2e/fir/fir_kernel_vec.mlir
index 3dbb678a..f8d7eb14 100644
--- a/test/e2e/fir/fir_kernel_vec.mlir
+++ b/test/e2e/fir/fir_kernel_vec.mlir
@@ -8,6 +8,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -30,49 +31,50 @@
 // MAPPING-SAME:   x_tiles = 4
 // MAPPING-SAME:   y_tiles = 4
 //
-// MAPPING:          %0 = "neura.grant_once"() <{constant_value = 0 : i64}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %1 = "neura.grant_once"() <{constant_value = dense<0> : vector<4xi32>}> {dfg_id = 1 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 1 : i32}]} : () -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %2 = neura.reserve {dfg_id = 2 : i32} : !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %3 = "neura.data_mov"(%1) {dfg_id = 5 : i32, mapping_locs = [{id = 128 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %4 = neura.phi_start %3, %2 {dfg_id = 7 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 0 : i32, y = 1 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1> -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %5 = neura.reserve {dfg_id = 3 : i32} : !neura.data<i64, i1>
-// MAPPING-NEXT:     %6 = "neura.data_mov"(%0) {dfg_id = 4 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %7 = neura.phi_start %6, %5 {dfg_id = 6 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %8 = "neura.data_mov"(%7) {dfg_id = 10 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %9 = "neura.gep"(%8) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 14 : i32, lhs_value = "%arg0", mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %10 = "neura.data_mov"(%9) {dfg_id = 18 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %11 = "neura.load"(%10) {dfg_id = 21 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %12 = "neura.data_mov"(%7) {dfg_id = 9 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %13 = "neura.gep"(%12) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 13 : i32, lhs_value = "%arg2", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %14 = "neura.data_mov"(%13) {dfg_id = 17 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:     %15 = "neura.load"(%14) {dfg_id = 20 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %16 = "neura.data_mov"(%15) {dfg_id = 24 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %17 = "neura.data_mov"(%11) {dfg_id = 25 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %18 = "neura.vmul"(%16, %17) {dfg_id = 27 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %19 = "neura.data_mov"(%18) {dfg_id = 30 : i32, mapping_locs = [{id = 43 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %20 = "neura.data_mov"(%4) {dfg_id = 11 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %21 = "neura.vadd"(%19, %20) {dfg_id = 32 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %22 = "neura.data_mov"(%7) {dfg_id = 8 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %23 = "neura.add"(%22) {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 4 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %24 = "neura.data_mov"(%23) {dfg_id = 16 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %25 = "neura.icmp"(%24) <{cmpType = "eq"}> {dfg_id = 19 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 32 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %26 = "neura.data_mov"(%25) {dfg_id = 23 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %27 = "neura.not"(%26) {dfg_id = 26 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %28 = "neura.data_mov"(%23) {dfg_id = 15 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:     %29 = "neura.data_mov"(%27) {dfg_id = 29 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %30 = neura.grant_predicate %28, %29 {dfg_id = 31 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:     neura.ctrl_mov %30 -> %5 {dfg_id = 33 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
-// MAPPING-NEXT:     %31 = "neura.data_mov"(%21) {dfg_id = 35 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}, {id = 160 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %32 = "neura.data_mov"(%27) {dfg_id = 28 : i32, mapping_locs = [{id = 31 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 161 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 161 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %33 = neura.grant_predicate %31, %32 {dfg_id = 37 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 1 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<i1, i1> -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     neura.ctrl_mov %33 -> %2 {dfg_id = 39 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}]} : !neura.data<vector<4xi32>, i1> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %34 = "neura.data_mov"(%21) {dfg_id = 34 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %35 = "neura.data_mov"(%25) {dfg_id = 22 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:     %36 = neura.grant_predicate %34, %35 {dfg_id = 36 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<i1, i1> -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %37 = "neura.data_mov"(%36) {dfg_id = 38 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
-// MAPPING-NEXT:     %38 = "neura.vector.reduce.add"(%37) {dfg_id = 40 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     %39 = "neura.data_mov"(%38) {dfg_id = 41 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:     "neura.return"(%39) {dfg_id = 42 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>) -> ()
+// MAPPING:      %0 = "neura.grant_once"() <{constant_value = 0 : i64}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i64, i1>
+// MAPPING-NEXT: %1 = "neura.grant_once"() <{constant_value = dense<0> : vector<4xi32>}> {dfg_id = 1 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 0 : i32, y = 1 : i32}]} : () -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %2 = neura.reserve {dfg_id = 2 : i32} : !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %3 = "neura.data_mov"(%1) {dfg_id = 6 : i32, mapping_locs = [{id = 128 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %4 = neura.phi_start %3, %2 {dfg_id = 8 : i32, mapping_locs = [{id = 4 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 0 : i32, y = 1 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1> -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %5 = neura.reserve {dfg_id = 3 : i32} : !neura.data<i64, i1>
+// MAPPING-NEXT: %6 = "neura.data_mov"(%0) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %7 = neura.phi_start %6, %5 {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT: %8 = "neura.data_mov"(%7) {dfg_id = 11 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %9 = "neura.gep"(%8) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 15 : i32, lhs_value = "%arg0", mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT: %10 = "neura.data_mov"(%9) {dfg_id = 19 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT: %11 = "neura.load"(%10) {dfg_id = 22 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %12 = "neura.data_mov"(%7) {dfg_id = 10 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %13 = "neura.gep"(%12) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 14 : i32, lhs_value = "%arg2", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT: %14 = "neura.data_mov"(%13) {dfg_id = 18 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT: %15 = "neura.load"(%14) {dfg_id = 21 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %16 = "neura.data_mov"(%15) {dfg_id = 25 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %17 = "neura.data_mov"(%11) {dfg_id = 26 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %18 = "neura.vmul"(%16, %17) {dfg_id = 28 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %19 = "neura.data_mov"(%18) {dfg_id = 31 : i32, mapping_locs = [{id = 43 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 42 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %20 = "neura.data_mov"(%4) {dfg_id = 12 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 16 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %21 = "neura.vadd"(%19, %20) {dfg_id = 33 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<vector<4xi32>, i1>, !neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %22 = "neura.data_mov"(%7) {dfg_id = 9 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %23 = "neura.add"(%22) {dfg_id = 13 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 4 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %24 = "neura.data_mov"(%23) {dfg_id = 17 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %25 = "neura.icmp"(%24) <{cmpType = "eq"}> {dfg_id = 20 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 32 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %26 = "neura.data_mov"(%25) {dfg_id = 24 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %27 = "neura.not"(%26) {dfg_id = 27 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %28 = "neura.data_mov"(%23) {dfg_id = 16 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT: %29 = "neura.data_mov"(%27) {dfg_id = 30 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %30 = neura.grant_predicate %28, %29 {dfg_id = 32 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT: neura.ctrl_mov %30 -> %5 {dfg_id = 34 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
+// MAPPING-NEXT: %31 = "neura.data_mov"(%21) {dfg_id = 36 : i32, mapping_locs = [{id = 29 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}, {id = 160 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %32 = "neura.data_mov"(%27) {dfg_id = 29 : i32, mapping_locs = [{id = 31 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 29 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 161 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 161 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %33 = neura.grant_predicate %31, %32 {dfg_id = 38 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 1 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<i1, i1> -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: neura.ctrl_mov %33 -> %2 {dfg_id = 40 : i32, mapping_locs = [{id = 13 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}]} : !neura.data<vector<4xi32>, i1> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %34 = "neura.data_mov"(%21) {dfg_id = 35 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %35 = "neura.data_mov"(%25) {dfg_id = 23 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT: %36 = neura.grant_predicate %34, %35 {dfg_id = 37 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<vector<4xi32>, i1>, !neura.data<i1, i1> -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %37 = "neura.data_mov"(%36) {dfg_id = 39 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<vector<4xi32>, i1>
+// MAPPING-NEXT: %38 = "neura.vector.reduce.add"(%37) {dfg_id = 41 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 1 : i32, y = 2 : i32}]} : (!neura.data<vector<4xi32>, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT: %39 = "neura.data_mov"(%38) {dfg_id = 42 : i32, mapping_locs = [{id = 288 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT: neura.return_value %39 : !neura.data<i32, i1> {dfg_id = 43 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 1 : i32, y = 2 : i32}]}
+// MAPPING-NEXT: neura.yield {dfg_id = 4 : i32}
 
 // YAML:      array_config:
 // YAML-NEXT:   columns: 4
@@ -100,7 +102,7 @@
 // YAML-NEXT:             - index_per_ii: 4
 // YAML-NEXT:               operations:
 // YAML-NEXT:                 - opcode: "PHI_START"
-// YAML-NEXT:                   id: 7
+// YAML-NEXT:                   id: 8
 // YAML-NEXT:                   time_step: 4
 // YAML-NEXT:                   invalid_iterations: 0
 // YAML-NEXT:                   src_operands:
@@ -149,7 +151,7 @@
 // ASM-NEXT:   VECTOR.REDUCE.ADD, [$0] -> [$0] (t=8, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=3)
 // ASM-NEXT: {
-// ASM-NEXT:   RETURN, [$0] (t=9, inv_iters=1)
+// ASM-NEXT:   RETURN_VALUE, [$0] (t=9, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=4)
 // ASM:      PE(2,2):
 // ASM-NEXT: {
diff --git a/test/e2e/histogram/histogram_kernel.mlir b/test/e2e/histogram/histogram_kernel.mlir
index a381e8e0..8aee3420 100644
--- a/test/e2e/histogram/histogram_kernel.mlir
+++ b/test/e2e/histogram/histogram_kernel.mlir
@@ -10,6 +10,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -31,96 +32,126 @@
 // MAPPING-SAME: res_mii = 2
 // MAPPING-SAME: x_tiles = 4
 // MAPPING-SAME: y_tiles = 4
-// MAPPING-NEXT:    %0 = "neura.grant_once"() <{constant_value = 0 : i64}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %1 = neura.reserve {dfg_id = 1 : i32} : !neura.data<i64, i1>
-// MAPPING-NEXT:    %2 = "neura.data_mov"(%0) {dfg_id = 3 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %3 = neura.phi_start %2, %1 {dfg_id = 4 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %4 = "neura.data_mov"(%3) {dfg_id = 6 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %5 = "neura.gep"(%4) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 8 : i32, lhs_value = "%arg0", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:    %6 = "neura.data_mov"(%5) {dfg_id = 11 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:    %7 = "neura.load"(%6) {dfg_id = 13 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %8 = "neura.data_mov"(%7) {dfg_id = 15 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %9 = "neura.mul"(%8) {dfg_id = 17 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 5 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %10 = "neura.data_mov"(%9) {dfg_id = 19 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %11 = "neura.add"(%10) {dfg_id = 21 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 3 : i32}], rhs_value = -5 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %12 = "neura.data_mov"(%11) {dfg_id = 23 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %13 = "neura.div"(%12) {dfg_id = 24 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 3 : i32}], rhs_value = 18 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %14 = "neura.data_mov"(%13) {dfg_id = 25 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %15 = neura.sext %14 {dfg_id = 26 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 3 : i32, y = 3 : i32}]} : !neura.data<i32, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %16 = "neura.data_mov"(%15) {dfg_id = 27 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %17 = "neura.gep"(%16) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 28 : i32, lhs_value = "%arg1", mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:    %18 = "neura.data_mov"(%17) {dfg_id = 30 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:    %19 = "neura.load"(%18) {dfg_id = 31 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %20 = "neura.data_mov"(%19) {dfg_id = 32 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %21 = "neura.add"(%20) {dfg_id = 33 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 10 : i32, x = 2 : i32, y = 3 : i32}], rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %22 = "neura.data_mov"(%21) {dfg_id = 34 : i32, mapping_locs = [{id = 448 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:    %23 = "neura.data_mov"(%17) {dfg_id = 29 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 449 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}, {id = 449 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:    "neura.store"(%22, %23) {dfg_id = 35 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i32, i1>, !neura.data<!llvm.ptr, i1>) -> ()
-// MAPPING-NEXT:    %24 = "neura.data_mov"(%3) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %25 = "neura.add"(%24) {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %26 = "neura.data_mov"(%25) {dfg_id = 10 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %27 = "neura.icmp"(%26) <{cmpType = "eq"}> {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 20 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:    %28 = "neura.data_mov"(%27) {dfg_id = 14 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:    %29 = "neura.not"(%28) {dfg_id = 16 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:    %30 = "neura.data_mov"(%25) {dfg_id = 9 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:    %31 = "neura.data_mov"(%29) {dfg_id = 18 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:    %32 = neura.grant_predicate %30, %31 {dfg_id = 20 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:    neura.ctrl_mov %32 -> %1 {dfg_id = 22 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
-// MAPPING-NEXT:    "neura.return"() {dfg_id = 2 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 2 : i32, y = 2 : i32}]} : () -> ()
+// MAPPING-NEXT:     %0 = "neura.grant_once"() <{constant_value = 0 : i64}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %1 = neura.reserve {dfg_id = 1 : i32} : !neura.data<i64, i1>
+// MAPPING-NEXT:     %2 = "neura.data_mov"(%0) {dfg_id = 3 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %3 = neura.phi_start %2, %1 {dfg_id = 4 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i64, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %4 = "neura.data_mov"(%3) {dfg_id = 6 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %5 = "neura.gep"(%4) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 8 : i32, lhs_value = "%arg0", mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %6 = "neura.data_mov"(%5) {dfg_id = 11 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %7 = "neura.load"(%6) {dfg_id = 13 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %8 = "neura.data_mov"(%7) {dfg_id = 17 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %9 = "neura.mul"(%8) {dfg_id = 20 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 5 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %10 = "neura.data_mov"(%9) {dfg_id = 23 : i32, mapping_locs = [{id = 37 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %11 = "neura.add"(%10) {dfg_id = 26 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 3 : i32}], rhs_value = -5 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %12 = "neura.data_mov"(%11) {dfg_id = 28 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 5 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %13 = "neura.div"(%12) {dfg_id = 29 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 3 : i32, y = 3 : i32}], rhs_value = 18 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %14 = "neura.data_mov"(%13) {dfg_id = 30 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %15 = neura.sext %14 {dfg_id = 31 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 3 : i32, y = 3 : i32}]} : !neura.data<i32, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %16 = "neura.data_mov"(%15) {dfg_id = 32 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %17 = "neura.gep"(%16) <{operandSegmentSizes = array<i32: 0, 1>}> {dfg_id = 33 : i32, lhs_value = "%arg1", mapping_locs = [{id = 15 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %18 = "neura.data_mov"(%17) {dfg_id = 35 : i32, mapping_locs = [{id = 480 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %19 = "neura.load"(%18) {dfg_id = 36 : i32, mapping_locs = [{id = 15 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 3 : i32, y = 3 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %20 = "neura.data_mov"(%19) {dfg_id = 37 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %21 = "neura.add"(%20) {dfg_id = 38 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 10 : i32, x = 2 : i32, y = 3 : i32}], rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %22 = "neura.data_mov"(%21) {dfg_id = 39 : i32, mapping_locs = [{id = 448 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %23 = "neura.data_mov"(%17) {dfg_id = 34 : i32, mapping_locs = [{id = 46 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 449 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}, {id = 449 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     "neura.store"(%22, %23) {dfg_id = 40 : i32, mapping_locs = [{id = 14 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i32, i1>, !neura.data<!llvm.ptr, i1>) -> ()
+// MAPPING-NEXT:     %24 = "neura.data_mov"(%3) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %25 = "neura.add"(%24) {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i64} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %26 = "neura.data_mov"(%25) {dfg_id = 10 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %27 = "neura.icmp"(%26) <{cmpType = "eq"}> {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 20 : i64} : (!neura.data<i64, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %28 = "neura.data_mov"(%27) {dfg_id = 16 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %29 = "neura.not"(%28) {dfg_id = 19 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %30 = "neura.data_mov"(%25) {dfg_id = 9 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %31 = "neura.data_mov"(%29) {dfg_id = 22 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %32 = neura.grant_predicate %30, %31 {dfg_id = 25 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     neura.ctrl_mov %32 -> %1 {dfg_id = 27 : i32, mapping_locs = [{id = 32 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i64, i1> !neura.data<i64, i1>
+// MAPPING-NEXT:     %33 = "neura.data_mov"(%27) {dfg_id = 14 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}, {id = 224 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}, {id = 224 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 5 : i32}, {id = 224 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}, {id = 224 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}, {id = 224 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}, {id = 224 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %34 = "neura.data_mov"(%27) {dfg_id = 15 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 36 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 225 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 5 : i32}, {id = 225 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 225 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}, {id = 225 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}, {id = 225 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %35 = neura.grant_predicate %33, %34 {dfg_id = 18 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 10 : i32, x = 3 : i32, y = 1 : i32}]} : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %36 = "neura.data_mov"(%35) {dfg_id = 21 : i32, mapping_locs = [{id = 226 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     neura.return_void %36 : !neura.data<i1, i1> {dfg_id = 24 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 3 : i32, y = 1 : i32}]}
+// MAPPING-NEXT:     neura.yield {dfg_id = 2 : i32}
 
-// YAML:     array_config:
-// YAML-NEXT:  columns: 4
-// YAML-NEXT:  rows: 4
-// YAML-NEXT:  compiled_ii: 5
-// YAML-NEXT:  cores:
-// YAML-NEXT:    - column: 2
-// YAML-NEXT:      row: 2
-// YAML-NEXT:      core_id: "10"
-// YAML-NEXT:      entries:
-// YAML-NEXT:        - entry_id: "entry0"
-// YAML-NEXT:          instructions:
-// YAML-NEXT:            - index_per_ii: 0
-// YAML-NEXT:              operations:
-// YAML-NEXT:                - opcode: "GRANT_PREDICATE"
-// YAML-NEXT:                  id: 20
-// YAML-NEXT:                  time_step: 5
-// YAML-NEXT:                  invalid_iterations: 1
-// YAML-NEXT:                  src_operands:
-// YAML-NEXT:                    - operand: "$0"
-// YAML-NEXT:                      color: "RED"
-// YAML-NEXT:                    - operand: "$1"
-// YAML-NEXT:                      color: "RED"
-// YAML-NEXT:                  dst_operands:
-// YAML-NEXT:                    - operand: "EAST"
-// YAML-NEXT:                      color: "RED"
-// YAML-NEXT:            - index_per_ii: 1
-// YAML-NEXT:              operations:
-// YAML-NEXT:                - opcode: "RETURN"
-// YAML-NEXT:                  id: 2
-// YAML-NEXT:                  time_step: 11
-// YAML-NEXT:                  invalid_iterations: 2
-// YAML-NEXT:            - index_per_ii: 2
-// YAML-NEXT:              operations:
-// YAML-NEXT:                - opcode: "GEP"
-// YAML-NEXT:                  id: 8
-// YAML-NEXT:                  time_step: 2
-// YAML-NEXT:                  invalid_iterations: 0
-// YAML-NEXT:                  src_operands:
-// YAML-NEXT:                    - operand: "EAST"
-// YAML-NEXT:                      color: "RED"
-// YAML-NEXT:                  dst_operands:
-// YAML-NEXT:                    - operand: "$0"
-// YAML-NEXT:                      color: "RED"
+// YAML:      array_config:
+// YAML-NEXT:   columns: 4
+// YAML-NEXT:   rows: 4
+// YAML-NEXT:   compiled_ii: 5
+// YAML-NEXT:   cores:
+// YAML-NEXT:     - column: 3
+// YAML-NEXT:       row: 1
+// YAML-NEXT:       core_id: "7"
+// YAML-NEXT:       entries:
+// YAML-NEXT:         - entry_id: "entry0"
+// YAML-NEXT:           instructions:
+// YAML-NEXT:             - index_per_ii: 0
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "GRANT_PREDICATE"
+// YAML-NEXT:                   id: 18
+// YAML-NEXT:                   time_step: 10
+// YAML-NEXT:                   invalid_iterations: 2
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "$0"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                     - operand: "$1"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$2"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:             - index_per_ii: 1
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "RETURN_VOID"
+// YAML-NEXT:                   id: 24
+// YAML-NEXT:                   time_step: 11
+// YAML-NEXT:                   invalid_iterations: 2
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "$2"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:             - index_per_ii: 3
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "DATA_MOV"
+// YAML-NEXT:                   id: 15
+// YAML-NEXT:                   time_step: 3
+// YAML-NEXT:                   invalid_iterations: 0
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "NORTH"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$1"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:             - index_per_ii: 4
+// YAML-NEXT:               operations:
+// YAML-NEXT:                 - opcode: "DATA_MOV"
+// YAML-NEXT:                   id: 14
+// YAML-NEXT:                   time_step: 4
+// YAML-NEXT:                   invalid_iterations: 0
+// YAML-NEXT:                   src_operands:
+// YAML-NEXT:                     - operand: "NORTH"
+// YAML-NEXT:                       color: "RED"
+// YAML-NEXT:                   dst_operands:
+// YAML-NEXT:                     - operand: "$0"
+// YAML-NEXT:                       color: "RED"
 
 // ASM:      # Compiled II: 5
-// ASM:      PE(2,2):
+// ASM:      PE(3,1):
 // ASM-NEXT: {
-// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [EAST, RED] (t=5, inv_iters=1)
+// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [$2] (t=10, inv_iters=2)
 // ASM-NEXT: } (idx_per_ii=0)
 // ASM-NEXT: {
-// ASM-NEXT:   RETURN (t=11, inv_iters=2)
+// ASM-NEXT:   RETURN_VOID, [$2] (t=11, inv_iters=2)
 // ASM-NEXT: } (idx_per_ii=1)
 // ASM-NEXT: {
+// ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [$1] (t=3, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=3)
+// ASM-NEXT: {
+// ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [$0] (t=4, inv_iters=0)
+// ASM-NEXT: } (idx_per_ii=4)
+// ASM:      PE(2,2):
+// ASM-NEXT: {
+// ASM-NEXT:   GRANT_PREDICATE, [$0], [$1] -> [EAST, RED] (t=5, inv_iters=1)
+// ASM-NEXT: } (idx_per_ii=0)
+// ASM-NEXT: {
 // ASM-NEXT:   GEP, [EAST, RED] -> [$0] (t=2, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=2)
 // ASM-NEXT: {
@@ -141,7 +172,7 @@
 // ASM-NEXT:   ADD, [$0], [#1] -> [$0], [WEST, RED] (t=2, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=2)
 // ASM-NEXT: {
-// ASM-NEXT:   ICMP_EQ, [$0], [#20] -> [WEST, RED] (t=3, inv_iters=0)
+// ASM-NEXT:   ICMP_EQ, [$0], [#20] -> [WEST, RED], [SOUTH, RED] (t=3, inv_iters=0)
 // ASM-NEXT: } (idx_per_ii=3)
 // ASM-NEXT: {
 // ASM-NEXT:   MUL, [WEST, RED], [#5] -> [NORTH, RED] (t=4, inv_iters=0)
@@ -152,6 +183,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/e2e/relu/relu_kernel.mlir b/test/e2e/relu/relu_kernel.mlir
index 45d2248b..717dd7a6 100644
--- a/test/e2e/relu/relu_kernel.mlir
+++ b/test/e2e/relu/relu_kernel.mlir
@@ -16,6 +16,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -31,56 +32,60 @@
 // Check the mapped MLIR contains proper structure and neura operations.
 // RUN: FileCheck %s --input-file=%t-mapping.mlir -check-prefix=MAPPING
 // MAPPING:      func.func @kernel(%arg0: i32 {llvm.noundef}, %arg1: i32 {llvm.noundef}, %arg2: i32 {llvm.noundef}, %arg3: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.writeonly}, %arg4: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, %arg5: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readnone}) -> !llvm.void attributes {CConv = #llvm.cconv<ccc>, accelerator = "neura", dataflow_mode = "predicate", linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 5 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 5 : i32, res_mii = 2 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}, memory_effects = #llvm.memory_effects<other = none, argMem = readwrite, inaccessibleMem = none>, no_unwind, passthrough = ["nofree", "norecurse", "nosync", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic", unnamed_addr = 1 : i64, visibility_ = 0 : i64} {
-// MAPPING-NEXT:   %0 = "neura.grant_once"() <{constant_value = 0 : i32}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %1 = neura.reserve {dfg_id = 1 : i32} : !neura.data<i32, i1>
-// MAPPING-NEXT:   %2 = "neura.data_mov"(%0) {dfg_id = 4 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 31 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}, {id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}, {id = 288 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 288 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}, {id = 288 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 5 : i32}, {id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %3 = neura.phi_start %2, %1 {dfg_id = 6 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %4 = neura.reserve {dfg_id = 2 : i32} : !neura.data<i32, i1>
-// MAPPING-NEXT:   %5 = "neura.data_mov"(%0) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %6 = neura.phi_start %5, %4 {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %7 = "neura.data_mov"(%6) {dfg_id = 11 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %8 = "neura.cast"(%7) <{cast_type = "trunc"}> {dfg_id = 13 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %9 = "neura.data_mov"(%8) {dfg_id = 17 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 2 : i32}, {id = 321 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %10 = "neura.div"(%9) {dfg_id = 20 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}], rhs_value = 70 : i16} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %11 = "neura.data_mov"(%8) {dfg_id = 16 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %12 = "neura.rem"(%11) {dfg_id = 19 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}], rhs_value = 70 : i16} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %13 = "neura.data_mov"(%10) {dfg_id = 23 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}, {id = 321 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 5 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %14 = neura.zext %13 {dfg_id = 26 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i16, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %15 = "neura.data_mov"(%12) {dfg_id = 22 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
-// MAPPING-NEXT:   %16 = neura.zext %15 {dfg_id = 25 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i16, i1> -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %17 = "neura.data_mov"(%14) {dfg_id = 32 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %18 = "neura.data_mov"(%16) {dfg_id = 30 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 192 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %19 = "neura.gep"(%17, %18) <{operandSegmentSizes = array<i32: 0, 2>}> {dfg_id = 36 : i32, lhs_value = "%arg4", mapping_locs = [{id = 6 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:   %20 = "neura.data_mov"(%19) {dfg_id = 40 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:   %21 = "neura.load"(%20) {dfg_id = 41 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %22 = "neura.data_mov"(%21) {dfg_id = 43 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %23 = "neura.icmp"(%22) <{cmpType = "sge"}> {dfg_id = 44 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 2 : i32, y = 1 : i32}], rhs_value = 0 : i32} : (!neura.data<i32, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %24 = "neura.data_mov"(%23) {dfg_id = 45 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %25 = "neura.data_mov"(%21) {dfg_id = 42 : i32, mapping_locs = [{id = 193 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}, {id = 193 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %26 = "neura.data_mov"(%3) {dfg_id = 9 : i32, mapping_locs = [{id = 28 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 7 : i32}, {id = 33 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 194 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %27 = "neura.sel"(%24, %25, %26) {dfg_id = 46 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 10 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %28 = "neura.data_mov"(%14) {dfg_id = 31 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}, {id = 448 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %29 = "neura.data_mov"(%16) {dfg_id = 29 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 449 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 449 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
-// MAPPING-NEXT:   %30 = "neura.gep"(%28, %29) <{operandSegmentSizes = array<i32: 0, 2>}> {dfg_id = 35 : i32, lhs_value = "%arg3", mapping_locs = [{id = 14 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:   %31 = "neura.data_mov"(%27) {dfg_id = 47 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %32 = "neura.data_mov"(%30) {dfg_id = 39 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 33 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 9 : i32}, {id = 193 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
-// MAPPING-NEXT:   "neura.store"(%31, %32) {dfg_id = 48 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i32, i1>, !neura.data<!llvm.ptr, i1>) -> ()
-// MAPPING-NEXT:   %33 = "neura.data_mov"(%6) {dfg_id = 10 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %34 = "neura.add"(%33) {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %35 = "neura.data_mov"(%34) {dfg_id = 15 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %36 = "neura.icmp"(%35) <{cmpType = "eq"}> {dfg_id = 18 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 4200 : i32} : (!neura.data<i32, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %37 = "neura.data_mov"(%36) {dfg_id = 21 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %38 = "neura.not"(%37) {dfg_id = 24 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %39 = "neura.data_mov"(%34) {dfg_id = 14 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 224 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 224 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %40 = "neura.data_mov"(%38) {dfg_id = 28 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %41 = neura.grant_predicate %39, %40 {dfg_id = 34 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:   neura.ctrl_mov %41 -> %4 {dfg_id = 38 : i32, mapping_locs = [{id = 23 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
-// MAPPING-NEXT:   %42 = "neura.data_mov"(%3) {dfg_id = 8 : i32, mapping_locs = [{id = 289 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}, {id = 289 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
-// MAPPING-NEXT:   %43 = "neura.data_mov"(%38) {dfg_id = 27 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 290 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 6 : i32}, {id = 290 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 7 : i32}, {id = 290 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
-// MAPPING-NEXT:   %44 = neura.grant_predicate %42, %43 {dfg_id = 33 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
-// MAPPING-NEXT:   neura.ctrl_mov %44 -> %1 {dfg_id = 37 : i32, mapping_locs = [{id = 289 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}, {id = 289 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 11 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
-// MAPPING-NEXT:   "neura.return"() {dfg_id = 3 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 2 : i32, y = 0 : i32}]} : () -> ()
-// MAPPING-NEXT: }
+// MAPPING:          %0 = "neura.grant_once"() <{constant_value = 0 : i32}> {dfg_id = 0 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 3 : i32, y = 2 : i32}]} : () -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %1 = neura.reserve {dfg_id = 1 : i32} : !neura.data<i32, i1>
+// MAPPING-NEXT:     %2 = "neura.data_mov"(%0) {dfg_id = 4 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}, {id = 31 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}, {id = 288 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}, {id = 288 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 288 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}, {id = 288 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 5 : i32}, {id = 288 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %3 = neura.phi_start %2, %1 {dfg_id = 6 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %4 = neura.reserve {dfg_id = 2 : i32} : !neura.data<i32, i1>
+// MAPPING-NEXT:     %5 = "neura.data_mov"(%0) {dfg_id = 5 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 0 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %6 = neura.phi_start %5, %4 {dfg_id = 7 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 1 : i32, x = 3 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i32, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %7 = "neura.data_mov"(%6) {dfg_id = 11 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 1 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %8 = "neura.cast"(%7) <{cast_type = "trunc"}> {dfg_id = 13 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 2 : i32, y = 2 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %9 = "neura.data_mov"(%8) {dfg_id = 17 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 2 : i32}, {id = 321 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %10 = "neura.div"(%9) {dfg_id = 20 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 2 : i32, y = 2 : i32}], rhs_value = 70 : i16} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %11 = "neura.data_mov"(%8) {dfg_id = 16 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %12 = "neura.rem"(%11) {dfg_id = 19 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 2 : i32, y = 2 : i32}], rhs_value = 70 : i16} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %13 = "neura.data_mov"(%10) {dfg_id = 25 : i32, mapping_locs = [{id = 321 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}, {id = 321 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 5 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %14 = neura.zext %13 {dfg_id = 29 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 6 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i16, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %15 = "neura.data_mov"(%12) {dfg_id = 24 : i32, mapping_locs = [{id = 320 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 320 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i16, i1>) -> !neura.data<i16, i1>
+// MAPPING-NEXT:     %16 = neura.zext %15 {dfg_id = 28 : i32, mapping_locs = [{id = 10 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 2 : i32, y = 2 : i32}]} : !neura.data<i16, i1> -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %17 = "neura.data_mov"(%14) {dfg_id = 36 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %18 = "neura.data_mov"(%16) {dfg_id = 34 : i32, mapping_locs = [{id = 33 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 192 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %19 = "neura.gep"(%17, %18) <{operandSegmentSizes = array<i32: 0, 2>}> {dfg_id = 41 : i32, lhs_value = "%arg4", mapping_locs = [{id = 6 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 7 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %20 = "neura.data_mov"(%19) {dfg_id = 45 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %21 = "neura.load"(%20) {dfg_id = 46 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %22 = "neura.data_mov"(%21) {dfg_id = 48 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %23 = "neura.icmp"(%22) <{cmpType = "sge"}> {dfg_id = 49 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 2 : i32, y = 1 : i32}], rhs_value = 0 : i32} : (!neura.data<i32, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %24 = "neura.data_mov"(%23) {dfg_id = 50 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %25 = "neura.data_mov"(%21) {dfg_id = 47 : i32, mapping_locs = [{id = 193 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}, {id = 193 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %26 = "neura.data_mov"(%3) {dfg_id = 9 : i32, mapping_locs = [{id = 28 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 7 : i32}, {id = 33 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 194 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 9 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %27 = "neura.sel"(%24, %25, %26) {dfg_id = 51 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 10 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i1, i1>, !neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %28 = "neura.data_mov"(%14) {dfg_id = 35 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 6 : i32}, {id = 448 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %29 = "neura.data_mov"(%16) {dfg_id = 33 : i32, mapping_locs = [{id = 34 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 449 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 449 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i64, i1>) -> !neura.data<i64, i1>
+// MAPPING-NEXT:     %30 = "neura.gep"(%28, %29) <{operandSegmentSizes = array<i32: 0, 2>}> {dfg_id = 40 : i32, lhs_value = "%arg3", mapping_locs = [{id = 14 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 2 : i32, y = 3 : i32}]} : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     %31 = "neura.data_mov"(%27) {dfg_id = 52 : i32, mapping_locs = [{id = 192 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %32 = "neura.data_mov"(%30) {dfg_id = 44 : i32, mapping_locs = [{id = 45 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 8 : i32}, {id = 33 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 9 : i32}, {id = 193 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}]} : (!neura.data<!llvm.ptr, i1>) -> !neura.data<!llvm.ptr, i1>
+// MAPPING-NEXT:     "neura.store"(%31, %32) {dfg_id = 53 : i32, mapping_locs = [{id = 6 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, resource = "tile", time_step = 11 : i32, x = 2 : i32, y = 1 : i32}]} : (!neura.data<i32, i1>, !neura.data<!llvm.ptr, i1>) -> ()
+// MAPPING-NEXT:     %33 = "neura.data_mov"(%6) {dfg_id = 10 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 1 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 1 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %34 = "neura.add"(%33) {dfg_id = 12 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 2 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 1 : i32} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %35 = "neura.data_mov"(%34) {dfg_id = 15 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 2 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %36 = "neura.icmp"(%35) <{cmpType = "eq"}> {dfg_id = 18 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 3 : i32, x = 3 : i32, y = 2 : i32}], rhs_value = 4200 : i32} : (!neura.data<i32, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %37 = "neura.data_mov"(%36) {dfg_id = 23 : i32, mapping_locs = [{id = 352 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %38 = "neura.not"(%37) {dfg_id = 27 : i32, mapping_locs = [{id = 11 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 4 : i32, x = 3 : i32, y = 2 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %39 = "neura.data_mov"(%34) {dfg_id = 14 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 2 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 2 : i32}, {id = 224 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 3 : i32}, {id = 224 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 4 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %40 = "neura.data_mov"(%38) {dfg_id = 32 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %41 = neura.grant_predicate %39, %40 {dfg_id = 39 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 5 : i32, x = 3 : i32, y = 1 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     neura.ctrl_mov %41 -> %4 {dfg_id = 43 : i32, mapping_locs = [{id = 23 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
+// MAPPING-NEXT:     %42 = "neura.data_mov"(%3) {dfg_id = 8 : i32, mapping_locs = [{id = 289 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}, {id = 289 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i32, i1>) -> !neura.data<i32, i1>
+// MAPPING-NEXT:     %43 = "neura.data_mov"(%38) {dfg_id = 31 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 31 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 290 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 6 : i32}, {id = 290 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 7 : i32}, {id = 290 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 2 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %44 = neura.grant_predicate %42, %43 {dfg_id = 38 : i32, mapping_locs = [{id = 9 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 1 : i32, y = 2 : i32}]} : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// MAPPING-NEXT:     neura.ctrl_mov %44 -> %1 {dfg_id = 42 : i32, mapping_locs = [{id = 289 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 9 : i32}, {id = 289 : i32, index_per_ii = 0 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 10 : i32}, {id = 289 : i32, index_per_ii = 1 : i32, invalid_iterations = 2 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 11 : i32}]} : !neura.data<i32, i1> !neura.data<i32, i1>
+// MAPPING-NEXT:     %45 = "neura.data_mov"(%36) {dfg_id = 21 : i32, mapping_locs = [{id = 36 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}, {id = 225 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 4 : i32}, {id = 225 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 5 : i32}, {id = 225 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 6 : i32}, {id = 225 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %46 = "neura.data_mov"(%36) {dfg_id = 22 : i32, mapping_locs = [{id = 35 : i32, index_per_ii = 3 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 3 : i32}, {id = 32 : i32, index_per_ii = 4 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 4 : i32}, {id = 36 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "link", time_step = 5 : i32}, {id = 224 : i32, index_per_ii = 1 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 6 : i32}, {id = 224 : i32, index_per_ii = 2 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 0 : i32, resource = "register", time_step = 7 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %47 = neura.grant_predicate %45, %46 {dfg_id = 26 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 8 : i32, x = 3 : i32, y = 1 : i32}]} : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// MAPPING-NEXT:     %48 = "neura.data_mov"(%47) {dfg_id = 30 : i32, mapping_locs = [{id = 225 : i32, index_per_ii = 3 : i32, invalid_iterations = 1 : i32, per_tile_register_id = 1 : i32, resource = "register", time_step = 8 : i32}]} : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// MAPPING-NEXT:     neura.return_void %48 : !neura.data<i1, i1> {dfg_id = 37 : i32, mapping_locs = [{id = 7 : i32, index_per_ii = 4 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 9 : i32, x = 3 : i32, y = 1 : i32}]}
+// MAPPING-NEXT:     neura.yield {dfg_id = 3 : i32}
 
 
 // YAML:      array_config:
@@ -89,18 +94,6 @@
 // YAML-NEXT:   compiled_ii: 5
 // YAML-NEXT:   cores:
 // YAML-NEXT:     - column: 2
-// YAML-NEXT:       row: 0
-// YAML-NEXT:       core_id: "2"
-// YAML-NEXT:       entries:
-// YAML-NEXT:         - entry_id: "entry0"
-// YAML-NEXT:           instructions:
-// YAML-NEXT:             - index_per_ii: 4
-// YAML-NEXT:               operations:
-// YAML-NEXT:                 - opcode: "RETURN"
-// YAML-NEXT:                   id: 3
-// YAML-NEXT:                   time_step: 9
-// YAML-NEXT:                   invalid_iterations: 1
-// YAML-NEXT:     - column: 2
 // YAML-NEXT:       row: 1
 // YAML-NEXT:       core_id: "6"
 // YAML-NEXT:       entries:
@@ -109,7 +102,7 @@
 // YAML-NEXT:             - index_per_ii: 0
 // YAML-NEXT:               operations:
 // YAML-NEXT:                 - opcode: "SEL"
-// YAML-NEXT:                   id: 46
+// YAML-NEXT:                   id: 51
 // YAML-NEXT:                   time_step: 10
 // YAML-NEXT:                   invalid_iterations: 2
 // YAML-NEXT:                   src_operands:
@@ -123,7 +116,7 @@
 // YAML-NEXT:                     - operand: "$0"
 // YAML-NEXT:                       color: "RED"
 // YAML-NEXT:                 - opcode: "DATA_MOV"
-// YAML-NEXT:                   id: 39
+// YAML-NEXT:                   id: 44
 // YAML-NEXT:                   time_step: 10
 // YAML-NEXT:                   invalid_iterations: 2
 // YAML-NEXT:                   src_operands:
@@ -131,14 +124,9 @@
 // YAML-NEXT:                       color: "RED"
 // YAML-NEXT:                   dst_operands:
 // YAML-NEXT:                     - operand: "$1"
-// YAML-NEXT:                       color: "RED"
 
 
 // ASM: # Compiled II: 5
-// ASM: PE(2,0):
-// ASM-NEXT: {
-// ASM-NEXT:   RETURN (t=9, inv_iters=1)
-// ASM-NEXT: } (idx_per_ii=4)
 // ASM:      PE(2,1):
 // ASM-NEXT: {
 // ASM-NEXT:   SEL, [$0], [$1], [$2] -> [$0] (t=10, inv_iters=2)
@@ -163,5 +151,9 @@
 // ASM-NEXT:   GRANT_PREDICATE, [$0], [NORTH, RED] -> [NORTH, RED] (t=5, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=0)
 // ASM-NEXT: {
+// ASM-NEXT:  DATA_MOV, [NORTH, RED] -> [$0] (t=6, inv_iters=1)
+// ASM-NEXT: } (idx_per_ii=1)
+// ASM-NEXT: {
 // ASM-NEXT:   DATA_MOV, [NORTH, RED] -> [$0] (t=3, inv_iters=0)
+// ASM-NEXT:   GRANT_PREDICATE, [$1], [$0] -> [$1] (t=8, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=3)
\ No newline at end of file
diff --git a/test/honor_arch/fir_removed_tiles_test.mlir b/test/honor_arch/fir_removed_tiles_test.mlir
index 9a9abb17..a1d94188 100644
--- a/test/honor_arch/fir_removed_tiles_test.mlir
+++ b/test/honor_arch/fir_removed_tiles_test.mlir
@@ -10,6 +10,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/mapping_quality/branch_for.mlir b/test/mapping_quality/branch_for.mlir
index 403f6b0b..48358227 100644
--- a/test/mapping_quality/branch_for.mlir
+++ b/test/mapping_quality/branch_for.mlir
@@ -15,6 +15,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -24,6 +25,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -34,6 +36,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -45,6 +48,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -57,6 +61,7 @@
 // RUN: mlir-neura-opt %s \
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -134,7 +139,8 @@ func.func @loop_test() -> f32 {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %12 -> %4 : !neura.data<f32, i1> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %13 = "neura.not"(%10) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %14 = neura.grant_predicate %8, %13 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%14) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %14 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 // FUSE:        func.func @loop_test() -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate"} {
@@ -153,7 +159,8 @@ func.func @loop_test() -> f32 {
 // FUSE-NEXT:     neura.ctrl_mov %10 -> %2 : !neura.data<f32, i1> !neura.data<f32, i1>
 // FUSE-NEXT:     %11 = "neura.not"(%8) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // FUSE-NEXT:     %12 = neura.grant_predicate %6, %11 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
-// FUSE-NEXT:     "neura.return"(%12) : (!neura.data<f32, i1>) -> ()
+// FUSE-NEXT:     neura.return_value %12 : !neura.data<f32, i1>
+// FUSE-NEXT:     neura.yield
 // FUSE-NEXT:   }
 
 // MOV:      func.func @loop_test() -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate"} {
@@ -185,7 +192,8 @@ func.func @loop_test() -> f32 {
 // MOV-NEXT:   %23 = "neura.data_mov"(%21) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // MOV-NEXT:   %24 = neura.grant_predicate %22, %23 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
 // MOV-NEXT:   %25 = "neura.data_mov"(%24) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
-// MOV-NEXT:   "neura.return"(%25) : (!neura.data<f32, i1>) -> ()
+// MOV-NEXT:     neura.return_value %25 : !neura.data<f32, i1>
+// MOV-NEXT:     neura.yield
 
 // MAPPING:   func.func @loop_test() -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate", mapping_info = {compiled_ii = 4 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 4 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
 
@@ -228,7 +236,7 @@ func.func @loop_test() -> f32 {
 // YAML-NEXT:             - index_per_ii: 3
 // YAML-NEXT:               operations:
 // YAML-NEXT:                 - opcode: "DATA_MOV"
-// YAML-NEXT:                   id: 200001
+// YAML-NEXT:                   id: 210001
 // YAML-NEXT:                   time_step: 3
 // YAML-NEXT:                   invalid_iterations: 0
 // YAML-NEXT:                   src_operands:
diff --git a/test/mapping_quality/tiny_loop.mlir b/test/mapping_quality/tiny_loop.mlir
index 6795a176..64f3edd4 100644
--- a/test/mapping_quality/tiny_loop.mlir
+++ b/test/mapping_quality/tiny_loop.mlir
@@ -17,6 +17,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -36,6 +37,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --fold-constant \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
diff --git a/test/neura/ctrl/branch.mlir b/test/neura/ctrl/branch.mlir
index d5634f03..eba379e0 100644
--- a/test/neura/ctrl/branch.mlir
+++ b/test/neura/ctrl/branch.mlir
@@ -2,6 +2,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   | FileCheck %s
@@ -10,6 +11,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -53,7 +55,7 @@ func.func @test(%in: i64) -> f32 {
 // CHECK-NEXT:     %12 = "neura.fmul"(%10, %11) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CHECK-NEXT:     neura.br %12 : !neura.data<f32, i1> to ^bb3
 // CHECK-NEXT:   ^bb3(%13: !neura.data<f32, i1>):  // 2 preds: ^bb1, ^bb2
-// CHECK-NEXT:     "neura.return"(%13) : (!neura.data<f32, i1>) -> ()
+// CHECK-NEXT:     "neura.return"(%13) {return_type = "value"} : (!neura.data<f32, i1>) -> ()
 // CHECK-NEXT:   }
 
 // CTRL2DATA:   func.func @test(%arg0: i64) -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate"} {
@@ -77,5 +79,6 @@ func.func @test(%in: i64) -> f32 {
 // CTRL2DATA-NEXT:     %17 = "neura.fadd"(%15, %16) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %18 = "neura.fmul"(%12, %13) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %19 = "neura.phi"(%17, %18) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%19) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %19 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
diff --git a/test/neura/ctrl/branch_for.mlir b/test/neura/ctrl/branch_for.mlir
index 11cd2344..bb8aff85 100644
--- a/test/neura/ctrl/branch_for.mlir
+++ b/test/neura/ctrl/branch_for.mlir
@@ -17,6 +17,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -27,6 +28,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -38,6 +40,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -50,6 +53,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -64,6 +68,7 @@
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -141,7 +146,8 @@ func.func @loop_test() -> f32 {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %12 -> %4 : !neura.data<f32, i1> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %13 = "neura.not"(%10) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %14 = neura.grant_predicate %8, %13 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%14) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %14 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 
@@ -161,7 +167,8 @@ func.func @loop_test() -> f32 {
 // FUSE-NEXT:     neura.ctrl_mov %10 -> %2 : !neura.data<f32, i1> !neura.data<f32, i1>
 // FUSE-NEXT:     %11 = "neura.not"(%8) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // FUSE-NEXT:     %12 = neura.grant_predicate %6, %11 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
-// FUSE-NEXT:     "neura.return"(%12) : (!neura.data<f32, i1>) -> ()
+// FUSE-NEXT:     neura.return_value %12 : !neura.data<f32, i1>
+// FUSE-NEXT:     neura.yield
 // FUSE-NEXT:   }
 
 // MOV:        func.func @loop_test() -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate"} {
@@ -193,7 +200,8 @@ func.func @loop_test() -> f32 {
 // MOV-NEXT:     %23 = "neura.data_mov"(%21) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // MOV-NEXT:     %24 = neura.grant_predicate %22, %23 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
 // MOV-NEXT:     %25 = "neura.data_mov"(%24) : (!neura.data<f32, i1>) -> !neura.data<f32, i1>
-// MOV-NEXT:     "neura.return"(%25) : (!neura.data<f32, i1>) -> ()
+// MOV-NEXT:     neura.return_value %25 : !neura.data<f32, i1>
+// MOV-NEXT:     neura.yield
 // MOV-NEXT:   }
 
 // MAPPING:        func.func @loop_test() -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate", mapping_info = {compiled_ii = 4 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 4 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
@@ -222,7 +230,7 @@ func.func @loop_test() -> f32 {
 // YAML-NEXT:                     - operand: "EAST"
 // YAML-NEXT:                       color: "RED"
 // YAML-NEXT:                 - opcode: "DATA_MOV"
-// YAML-NEXT:                   id: 180001
+// YAML-NEXT:                   id: 190001
 // YAML-NEXT:                   time_step: 4
 // YAML-NEXT:                   invalid_iterations: 1
 // YAML-NEXT:                   src_operands:
@@ -269,5 +277,5 @@ func.func @loop_test() -> f32 {
 // ASM-NEXT:   GRANT_PREDICATE, [WEST, RED], [$0] -> [$0] (t=5, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=1)
 // ASM-NEXT: {
-// ASM-NEXT:   RETURN, [$0] (t=6, inv_iters=1)
+// ASM-NEXT:   RETURN_VALUE, [$0] (t=6, inv_iters=1)
 // ASM-NEXT: } (idx_per_ii=2)
diff --git a/test/neura/ctrl/branch_with_and_without_arg.mlir b/test/neura/ctrl/branch_with_and_without_arg.mlir
index 4cee357e..d861d1d5 100644
--- a/test/neura/ctrl/branch_with_and_without_arg.mlir
+++ b/test/neura/ctrl/branch_with_and_without_arg.mlir
@@ -10,6 +10,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -72,5 +73,6 @@ func.func @test(%in: i64) -> f32 {
 // CTRL2DATA-NEXT:     %15 = "neura.fadd"(%13, %14) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %16 = "neura.fmul"(%10, %11) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %17 = "neura.phi"(%15, %16) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%17) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %17 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/neura/ctrl/branch_without_arg.mlir b/test/neura/ctrl/branch_without_arg.mlir
index a94c826a..e505afda 100644
--- a/test/neura/ctrl/branch_without_arg.mlir
+++ b/test/neura/ctrl/branch_without_arg.mlir
@@ -10,6 +10,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -76,5 +77,6 @@ func.func @test(%in: i64) -> f32 {
 // CTRL2DATA-NEXT:     %17 = "neura.fadd"(%15, %16) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %18 = "neura.fmul"(%12, %13) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %19 = "neura.phi"(%17, %18) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%19) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %19 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/neura/ctrl/for_with_if.mlir b/test/neura/ctrl/for_with_if.mlir
index 9f856356..ad8ba343 100644
--- a/test/neura/ctrl/for_with_if.mlir
+++ b/test/neura/ctrl/for_with_if.mlir
@@ -19,6 +19,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -138,6 +139,7 @@ module attributes {} {
 // CTRL2DATA-NEXT:     %45 = neura.grant_predicate %31, %36 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %46 = "neura.not"(%36) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %47 = neura.grant_predicate %33, %46 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
+// CTRL2DATA-NEXT:     neura.return_value %47 : !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %48 = "neura.icmp"(%37, %38) <{cmpType = "sge"}> : (!neura.data<i32, i1>, !neura.data<i32, i1>) -> !neura.data<i1, i1>
 // CTRL2DATA-NEXT:     %49 = neura.grant_predicate %37, %48 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     %50 = neura.grant_predicate %39, %48 : !neura.data<i32, i1>, !neura.data<i1, i1> -> !neura.data<i32, i1>
@@ -182,5 +184,5 @@ module attributes {} {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %74 -> %22 : !neura.data<i32, i1> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %73 -> %20 : !neura.data<i32, i1> !neura.data<i32, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %72 -> %18 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%47) : (!neura.data<i32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/neura/ctrl/nested_branch.mlir b/test/neura/ctrl/nested_branch.mlir
index 7be88a4a..92fe6975 100644
--- a/test/neura/ctrl/nested_branch.mlir
+++ b/test/neura/ctrl/nested_branch.mlir
@@ -2,6 +2,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   | FileCheck %s
@@ -10,6 +11,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -59,7 +61,7 @@ func.func @complex_test(%in: i64) -> f32 {
 // CHECK-NEXT:     %17 = "neura.fmul"(%15, %16) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CHECK-NEXT:     neura.br %17 : !neura.data<f32, i1> to ^bb4
 // CHECK-NEXT:   ^bb4(%18: !neura.data<f32, i1>):  // 2 preds: ^bb1, ^bb3
-// CHECK-NEXT:     "neura.return"(%18) : (!neura.data<f32, i1>) -> ()
+// CHECK-NEXT:     "neura.return"(%18) {return_type = "value"} : (!neura.data<f32, i1>) -> ()
 // CHECK-NEXT:   }
 
 // CTRL2DATA:      func.func @complex_test(%arg0: i64) -> f32 attributes {accelerator = "neura", dataflow_mode = "predicate"} {
@@ -99,5 +101,6 @@ func.func @complex_test(%in: i64) -> f32 {
 // CTRL2DATA-NEXT:     neura.ctrl_mov %26 -> %18 : !neura.data<f32, i1> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %31 = "neura.fmul"(%12, %13) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // CTRL2DATA-NEXT:     %32 = "neura.phi"(%29, %31) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
-// CTRL2DATA-NEXT:     "neura.return"(%32) : (!neura.data<f32, i1>) -> ()
+// CTRL2DATA-NEXT:     neura.return_value %32 : !neura.data<f32, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
\ No newline at end of file
diff --git a/test/neura/for_loop/kernel_test.mlir b/test/neura/for_loop/kernel_test.mlir
index 51a15b07..1c00b1d5 100644
--- a/test/neura/for_loop/kernel_test.mlir
+++ b/test/neura/for_loop/kernel_test.mlir
@@ -16,6 +16,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -27,6 +28,7 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
@@ -107,7 +109,9 @@
 // CHECK-FUSED-NEXT:     neura.ctrl_mov %34 -> %11 : !neura.data<i64, i1> !neura.data<i64, i1>
 // CHECK-FUSED-NEXT:     %35 = neura.grant_predicate %10, %28 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CHECK-FUSED-NEXT:     neura.ctrl_mov %35 -> %9 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CHECK-FUSED-NEXT:     "neura.return"() : () -> ()
+// CHECK-FUSED-NEXT:     %36 = neura.grant_predicate %27, %27 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CHECK-FUSED-NEXT:     neura.return_void %36 : !neura.data<i1, i1>
+// CHECK-FUSED-NEXT:     neura.yield
 // CHECK-FUSED-NEXT:   }
 
 // CHECK-MOV:        func.func @_Z6kernelPfS_S_
@@ -193,5 +197,10 @@
 // CHECK-MOV-NEXT:     %71 = "neura.data_mov"(%51) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // CHECK-MOV-NEXT:     %72 = neura.grant_predicate %70, %71 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CHECK-MOV-NEXT:     neura.ctrl_mov %72 -> %11 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CHECK-MOV-NEXT:     "neura.return"() : () -> ()
+// CHECK-MOV-NEXT:     %73 = "neura.data_mov"(%49) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CHECK-MOV-NEXT:     %74 = "neura.data_mov"(%49) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CHECK-MOV-NEXT:     %75 = neura.grant_predicate %73, %74 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CHECK-MOV-NEXT:     %76 = "neura.data_mov"(%75) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
+// CHECK-MOV-NEXT:     neura.return_void %76 : !neura.data<i1, i1>
+// CHECK-MOV-NEXT:     neura.yield
 // CHECK-MOV-NEXT:   }
\ No newline at end of file
diff --git a/test/neura/for_loop/relu_test.mlir b/test/neura/for_loop/relu_test.mlir
index 951f89c4..366083d6 100644
--- a/test/neura/for_loop/relu_test.mlir
+++ b/test/neura/for_loop/relu_test.mlir
@@ -13,23 +13,25 @@
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
 // RUN:  | FileCheck %s --check-prefix=CTRL2DATA
 
-// RUN: mlir-neura-opt %t-relu.mlir\
+// RUN: mlir-neura-opt %t-relu.mlir \
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
 // RUN:   --promote-func-arg-to-const \
 // RUN:   --fold-constant \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
 // RUN:   --fold-constant \
 // RUN:   --insert-data-mov \
 // RUN:   --map-to-accelerator="mapping-strategy=heuristic backtrack-config=customized" \
-// RUN:  | FileCheck %s --check-prefix=MAPPING
+// RUN:   | FileCheck %s --check-prefix=MAPPING
 
 // CHECK:      func.func @_Z6kernelPiS_
 // CHECK-SAME: accelerator = "neura"
@@ -127,8 +129,10 @@
 // CTRL2DATA-NEXT:     neura.ctrl_mov %55 -> %14 : !neura.data<i64, i1> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     %56 = neura.grant_predicate %46, %50 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
 // CTRL2DATA-NEXT:     neura.ctrl_mov %56 -> %12 : !neura.data<i64, i1> !neura.data<i64, i1>
-// CTRL2DATA-NEXT:     "neura.return"() : () -> ()
+// CTRL2DATA-NEXT:     %57 = neura.grant_predicate %49, %49 : !neura.data<i1, i1>, !neura.data<i1, i1> -> !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.return_void %57 : !neura.data<i1, i1>
+// CTRL2DATA-NEXT:     neura.yield
 // CTRL2DATA-NEXT:   }
 
 
-// MAPPING: func.func @_Z6kernelPiS_(%arg0: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, %arg1: !llvm.ptr {llvm.nocapture, llvm.noundef}) -> !llvm.void attributes {CConv = #llvm.cconv<ccc>, accelerator = "neura", dataflow_mode = "predicate", linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 5 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 5 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}, memory_effects = #llvm.memory_effects<other = none, argMem = readwrite, inaccessibleMem = none>, no_unwind, passthrough = ["mustprogress", "nofree", "norecurse", "nosync", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic", unnamed_addr = 1 : i64, visibility_ = 0 : i64} {
\ No newline at end of file
+// MAPPING: func.func @_Z6kernelPiS_(%arg0: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, %arg1: !llvm.ptr {llvm.nocapture, llvm.noundef}) -> !llvm.void attributes {CConv = #llvm.cconv<ccc>, accelerator = "neura", dataflow_mode = "predicate", linkage = #llvm.linkage<external>, mapping_info = {compiled_ii = 5 : i32, mapping_mode = "spatial-temporal", mapping_strategy = "heuristic", rec_mii = 5 : i32, res_mii = 2 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}, memory_effects = #llvm.memory_effects<other = none, argMem = readwrite, inaccessibleMem = none>, no_unwind, passthrough = ["mustprogress", "nofree", "norecurse", "nosync", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic", unnamed_addr = 1 : i64, visibility_ = 0 : i64} {
\ No newline at end of file
diff --git a/test/neura/fusion/test.mlir b/test/neura/fusion/test.mlir
index f3b1dc32..d7d6711d 100644
--- a/test/neura/fusion/test.mlir
+++ b/test/neura/fusion/test.mlir
@@ -3,6 +3,7 @@
 // RUN: mlir-neura-opt --architecture-spec=%S/../../arch_spec/architecture.yaml --assign-accelerator \
 // RUN:           --lower-llvm-to-neura \
 // RUN:           --promote-func-arg-to-const \
+// RUN:           --canonicalize-return \
 // RUN:           --canonicalize-live-in \
 // RUN:           --leverage-predicated-value \
 // RUN:           --fold-constant \
@@ -16,6 +17,7 @@
 // RUN: mlir-neura-opt --architecture-spec=%S/../../arch_spec/architecture.yaml --assign-accelerator \
 // RUN:           --lower-llvm-to-neura \
 // RUN:           --promote-func-arg-to-const \
+// RUN:           --canonicalize-return \
 // RUN:           --canonicalize-live-in \
 // RUN:           --leverage-predicated-value \
 // RUN:           --fold-constant \
@@ -38,6 +40,7 @@
 // RUN:           --lower-llvm-to-neura \
 // RUN:           --promote-func-arg-to-const \
 // RUN:           --canonicalize-cast \
+// RUN:           --canonicalize-return \
 // RUN:           --canonicalize-live-in \
 // RUN:           --leverage-predicated-value \
 // RUN:           --fold-constant \
@@ -73,24 +76,26 @@
 // RUN:           --lower-llvm-to-neura \
 // RUN:           --promote-func-arg-to-const \
 // RUN:           --canonicalize-cast \
+// RUN:           --canonicalize-return \
 // RUN:           --canonicalize-live-in \
 // RUN:           --leverage-predicated-value \
 // RUN:           --fold-constant \
 // RUN:           --transform-ctrl-to-data-flow \
 // RUN:           --fold-constant \
-// RUN:           --init-pattern %t-kernel.mlir | FileCheck %s --check-prefix=CHECK-INIT-PATTERN
+// RUN:           --init-pattern %t-kernel.mlir \
+// RUN:           | FileCheck %s --check-prefix=CHECK-INIT-PATTERN
 
 // CHECK-INIT-PATTERN:         %21:2 = "neura.fused_op"(%16, %20) <{frequency = 6 : i64, pattern_id = 2 : i64, pattern_name = "gep->load"}> ({
 // CHECK-INIT-PATTERN-NEXT:    ^bb0(%arg5: !neura.data<!llvm.ptr, i1>, %arg6: !neura.data<i64, i1>):
-// CHECK-INIT-PATTERN-NEXT:      %74 = "neura.gep"(%arg5, %arg6) <{operandSegmentSizes = array<i32: 1, 1>}> : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// CHECK-INIT-PATTERN-NEXT:      %75 = "neura.load"(%74) : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// CHECK-INIT-PATTERN-NEXT:      neura.yield %74, %75 : !neura.data<!llvm.ptr, i1>, !neura.data<i32, i1>
+// CHECK-INIT-PATTERN-NEXT:      %75 = "neura.gep"(%arg5, %arg6) <{operandSegmentSizes = array<i32: 1, 1>}> : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// CHECK-INIT-PATTERN-NEXT:      %76 = "neura.load"(%75) : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// CHECK-INIT-PATTERN-NEXT:      neura.yield %75, %76 : !neura.data<!llvm.ptr, i1>, !neura.data<i32, i1>
 // CHECK-INIT-PATTERN-NEXT:    }) : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> (!neura.data<!llvm.ptr, i1>, !neura.data<i32, i1>)
 // CHECK-INIT-PATTERN-NEXT:    %22 = "neura.fused_op"(%18, %20) <{frequency = 6 : i64, pattern_id = 2 : i64, pattern_name = "gep->load"}> ({
 // CHECK-INIT-PATTERN-NEXT:    ^bb0(%arg5: !neura.data<!llvm.ptr, i1>, %arg6: !neura.data<i64, i1>):
-// CHECK-INIT-PATTERN-NEXT:      %74 = "neura.gep"(%arg5, %arg6) <{operandSegmentSizes = array<i32: 1, 1>}> : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
-// CHECK-INIT-PATTERN-NEXT:      %75 = "neura.load"(%74) : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
-// CHECK-INIT-PATTERN-NEXT:      neura.yield %75 : !neura.data<i32, i1>
+// CHECK-INIT-PATTERN-NEXT:      %75 = "neura.gep"(%arg5, %arg6) <{operandSegmentSizes = array<i32: 1, 1>}> : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> !neura.data<!llvm.ptr, i1>
+// CHECK-INIT-PATTERN-NEXT:      %76 = "neura.load"(%75) : (!neura.data<!llvm.ptr, i1>) -> !neura.data<i32, i1>
+// CHECK-INIT-PATTERN-NEXT:      neura.yield %76 : !neura.data<i32, i1>
 // CHECK-INIT-PATTERN-NEXT:    }) : (!neura.data<!llvm.ptr, i1>, !neura.data<i64, i1>) -> !neura.data<i32, i1>
 
 // RUN: mlir-neura-opt --architecture-spec=%S/../../arch_spec/architecture.yaml --verify-each=true --mlir-print-ir-after-failure \
@@ -98,6 +103,7 @@
 // RUN:           --lower-llvm-to-neura \
 // RUN:           --promote-func-arg-to-const \
 // RUN:           --canonicalize-cast \
+// RUN:           --canonicalize-return \
 // RUN:           --canonicalize-live-in \
 // RUN:           --leverage-predicated-value \
 // RUN:           --fold-constant \
diff --git a/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir b/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
index 3c2271f7..90df022d 100644
--- a/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
+++ b/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
@@ -1,6 +1,7 @@
 // RUN: mlir-neura-opt \
 // RUN:   --assign-accelerator \
 // RUN:   --lower-llvm-to-neura \
+// RUN:   --canonicalize-return \
 // RUN:   --canonicalize-live-in \
 // RUN:   --leverage-predicated-value \
 // RUN:   --transform-ctrl-to-data-flow \
diff --git a/test/neura/steer_ctrl/for_with_if.mlir b/test/neura/steer_ctrl/for_with_if.mlir
index 6fc68701..fe145f99 100644
--- a/test/neura/steer_ctrl/for_with_if.mlir
+++ b/test/neura/steer_ctrl/for_with_if.mlir
@@ -7,6 +7,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -74,6 +75,7 @@ module attributes {} {
 // CHECK-NEXT:   %28 = "neura.icmp"(%27, %25) <{cmpType = "slt"}> : (i64, i64) -> i1
 // CHECK-NEXT:   neura.ctrl_mov %28 -> %9 : i1 i1
 // CHECK-NEXT:   %29 = neura.false_steer %26, %28 : i32, i1 -> i32
+// CHECK-NEXT:   neura.return_value %29 : i32
 // CHECK-NEXT:   %30 = "neura.icmp"(%26, %24) <{cmpType = "sge"}> : (i32, i32) -> i1
 // CHECK-NEXT:   %31 = neura.true_steer %26, %30 : i32, i1 -> i32
 // CHECK-NEXT:   %32 = neura.true_steer %23, %30 : i32, i1 -> i32
@@ -102,4 +104,4 @@ module attributes {} {
 // CHECK-NEXT:   neura.ctrl_mov %47 -> %7 : i32 i32
 // CHECK-NEXT:   %48 = "neura.add"(%41, %38) : (i64, i64) -> i64
 // CHECK-NEXT:   neura.ctrl_mov %48 -> %8 : i64 i64
-// CHECK-NEXT:   "neura.return"(%29) : (i32) -> ()
+// CHECK-NEXT:   neura.yield
diff --git a/test/neura/steer_ctrl/loop_with_return_value.mlir b/test/neura/steer_ctrl/loop_with_return_value.mlir
index dd034790..b71e0984 100644
--- a/test/neura/steer_ctrl/loop_with_return_value.mlir
+++ b/test/neura/steer_ctrl/loop_with_return_value.mlir
@@ -7,6 +7,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -22,6 +23,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -61,11 +63,12 @@ module {
 // CHECK-NEXT:   %11 = "neura.icmp"(%10, %8) <{cmpType = "slt"}> : (i64, i64) -> i1
 // CHECK-NEXT:   neura.ctrl_mov %11 -> %2 : i1 i1
 // CHECK-NEXT:   %12 = neura.false_steer %9, %11 : i64, i1 -> i64
+// CHECK-NEXT:   neura.return_value %12 : i64
 // CHECK-NEXT:   %13 = "neura.add"(%9, %9) : (i64, i64) -> i64
 // CHECK-NEXT:   neura.ctrl_mov %13 -> %0 : i64 i64
 // CHECK-NEXT:   %14 = "neura.add"(%10, %7) : (i64, i64) -> i64
 // CHECK-NEXT:   neura.ctrl_mov %14 -> %1 : i64 i64
-// CHECK-NEXT:   "neura.return"(%12) : (i64) -> ()
+// CHECK-NEXT:   neura.yield
 // CHECK-NEXT: }
 
 // MAPPING:      func.func @simple_add_loop() -> i64 attributes {accelerator = "neura", dataflow_mode = "steering", mapping_info = {compiled_ii = 4 : i32, mapping_mode = "spatial-only", mapping_strategy = "heuristic", rec_mii = 2 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} {
diff --git a/test/neura/steer_ctrl/loop_without_return_value.mlir b/test/neura/steer_ctrl/loop_without_return_value.mlir
index b20ceb53..55e28f08 100644
--- a/test/neura/steer_ctrl/loop_without_return_value.mlir
+++ b/test/neura/steer_ctrl/loop_without_return_value.mlir
@@ -7,6 +7,7 @@
 // RUN: --lower-llvm-to-neura \
 // RUN: --canonicalize-cast \
 // RUN: --promote-func-arg-to-const \
+// RUN: --canonicalize-return \
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
@@ -47,10 +48,13 @@ module attributes {} {
 // CHECK-NEXT:   %15 = neura.carry %8, %1, %0 : i64, i1, i64 -> i64
 // CHECK-NEXT:   %16 = "neura.icmp"(%15, %14) <{cmpType = "slt"}> : (i64, i64) -> i1
 // CHECK-NEXT:   neura.ctrl_mov %16 -> %1 : i1 i1
-// CHECK-NEXT:   %17 = neura.load_indexed %13[%15 : i64] memref<?xi32> : i32
-// CHECK-NEXT:   %18 = "neura.mul"(%17, %12) : (i32, i32) -> i32
-// CHECK-NEXT:   %19 = "neura.add"(%18, %11) : (i32, i32) -> i32
-// CHECK-NEXT:   neura.store_indexed %19 to %10[%15 : i64] memref<?xi32> : i32
-// CHECK-NEXT:   %20 = "neura.add"(%15, %9) : (i64, i64) -> i64
-// CHECK-NEXT:   neura.ctrl_mov %20 -> %0 : i64 i64
-// CHECK-NEXT:   "neura.return"() : () -> ()
+// CHECK-NEXT:   %17 = "neura.not"(%16) : (i1) -> i1
+// CHECK-NEXT:   %18 = neura.false_steer %17, %16 : i1, i1 -> i1
+// CHECK-NEXT:   neura.return_void %18 : i1
+// CHECK-NEXT:   %19 = neura.load_indexed %13[%15 : i64] memref<?xi32> : i32
+// CHECK-NEXT:   %20 = "neura.mul"(%19, %12) : (i32, i32) -> i32
+// CHECK-NEXT:   %21 = "neura.add"(%20, %11) : (i32, i32) -> i32
+// CHECK-NEXT:   neura.store_indexed %21 to %10[%15 : i64] memref<?xi32> : i32
+// CHECK-NEXT:   %22 = "neura.add"(%15, %9) : (i64, i64) -> i64
+// CHECK-NEXT:   neura.ctrl_mov %22 -> %0 : i64 i64
+// CHECK-NEXT:   neura.yield
\ No newline at end of file

From 50c968094246b1bfcebd04a81a0a1ff56dc49cdc Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Tue, 30 Dec 2025 00:22:52 +0800
Subject: [PATCH 5/7] disable interpreter test for return

---
 .../loop_convert_controlflow_to_dataflow.mlir            | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir b/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
index 90df022d..a3d9a94c 100644
--- a/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
+++ b/test/neura/interpreter/loop_convert_controlflow_to_dataflow.mlir
@@ -8,11 +8,11 @@
 // RUN:   --fold-constant \
 // RUN:   %s -o %t_dataflow.mlir 
 
-// RUN: neura-interpreter %t_dataflow.mlir --verbose --dataflow > %t_output.txt
+// RU: neura-interpreter %t_dataflow.mlir --verbose --dataflow > %t_output.txt
 
-// RUN: FileCheck %s --check-prefix=DATAFLOW_IR --input-file=%t_dataflow.mlir
+// RU: FileCheck %s --check-prefix=DATAFLOW_IR --input-file=%t_dataflow.mlir
 
-// RUN: FileCheck %s --check-prefix=INTERPRETER_OUTPUT --input-file=%t_output.txt
+// RU: FileCheck %s --check-prefix=INTERPRETER_OUTPUT --input-file=%t_output.txt
 
 func.func @loop_sum() -> f32 {
   %c0 = "neura.constant"() <{predicate = true, value = 0.000000e+00 : f32}> : () -> f32   // init_i / init_sum
@@ -62,6 +62,7 @@ func.func @loop_sum() -> f32 {
 // DATAFLOW_IR-NEXT:     %19 = neura.grant_predicate %9, %14 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     %20 = "neura.not"(%14) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
 // DATAFLOW_IR-NEXT:     %21 = neura.grant_predicate %11, %20 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
+// DATAFLOW_IR-NEXT:     neura.return_value %21 : !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     %22 = "neura.fadd"(%16, %17) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     %23 = "neura.fadd"(%15, %18) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     neura.ctrl_mov %23 -> %12 : !neura.data<f32, i1> !neura.data<f32, i1>
@@ -69,7 +70,7 @@ func.func @loop_sum() -> f32 {
 // DATAFLOW_IR-NEXT:     neura.ctrl_mov %19 -> %8 : !neura.data<f32, i1> !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     neura.ctrl_mov %17 -> %6 : !neura.data<f32, i1> !neura.data<f32, i1>
 // DATAFLOW_IR-NEXT:     neura.ctrl_mov %18 -> %4 : !neura.data<f32, i1> !neura.data<f32, i1>
-// DATAFLOW_IR-NEXT:     "neura.return"(%21) : (!neura.data<f32, i1>) -> ()
+// DATAFLOW_IR-NEXT:     neura.yield
 // DATAFLOW_IR-NEXT:   }
 
 // INTERPRETER_OUTPUT: [neura-interpreter]  DFG Iteration 5 | Topological Level 6 | ready_to_execute_ops 3

From 26b918781876d43feb6451220b5c4940b310980d Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Tue, 30 Dec 2025 12:26:22 +0800
Subject: [PATCH 6/7] modify comments & add test

---
 .../Transforms/CanonicalizeReturnPass.cpp             |  3 +++
 test/arith2neura/add.mlir                             | 11 +++++++++++
 2 files changed, 14 insertions(+)
 create mode 100644 test/arith2neura/add.mlir

diff --git a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
index d6054c40..bd90765b 100644
--- a/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
+++ b/lib/NeuraDialect/Transforms/CanonicalizeReturnPass.cpp
@@ -85,6 +85,9 @@ static void processEmptyReturnVoidBlock(Block *ret_block,
     neura::Br br = cast<neura::Br>(pred_block->getTerminator());
 
     // Finds a suitable trigger value in the predecessor block.
+    // In dataflow semantics, even void returns need a value to trigger them.
+    // We use the result of the last operation in the predecessor block as the
+    // trigger signal.
     Value trigger_value = nullptr;
 
     // Iterates through operations in reverse order to find the last suitable
diff --git a/test/arith2neura/add.mlir b/test/arith2neura/add.mlir
new file mode 100644
index 00000000..71e1a995
--- /dev/null
+++ b/test/arith2neura/add.mlir
@@ -0,0 +1,11 @@
+// RUN: neura-compiler --neura-conversion %s | FileCheck %s --check-prefix=COMPILER
+// RUN: mlir-neura-opt --assign-accelerator --lower-arith-to-neura %s | FileCheck %s --check-prefix=OPT
+
+func.func @test(%a: f32) -> f32 {
+  %b = arith.constant 2.0 : f32
+  %res = arith.addf %a, %b : f32
+  return %res : f32
+}
+
+// COMPILER: neura.fadd
+// OPT:      neura.fadd
\ No newline at end of file

From e20077740010b5b0fd1a72dad7eff96010852788 Mon Sep 17 00:00:00 2001
From: ShangkunLI <shangkun.li@connect.ust.hk>
Date: Tue, 30 Dec 2025 12:31:27 +0800
Subject: [PATCH 7/7] modify tablegen file

---
 include/NeuraDialect/NeuraOps.td | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/include/NeuraDialect/NeuraOps.td b/include/NeuraDialect/NeuraOps.td
index 32cf2acb..1804c9f7 100644
--- a/include/NeuraDialect/NeuraOps.td
+++ b/include/NeuraDialect/NeuraOps.td
@@ -389,6 +389,8 @@ def Neura_ReturnVoidOp : Op<NeuraDialect, "return_void", []>{
 
     Example:
       neura.return_void %cond : !neura.data<i1, i1>  // In dataflow mode
+      ...
+      neura.yield                       // Function ends here.
 
     This is NOT a terminator - the block ends with neura.yield instead.
   }];
@@ -407,6 +409,8 @@ def Neura_ReturnValueOp : Op<NeuraDialect, "return_value", []>{
 
     Example:
       neura.return_value %ret_val : !neura.data<i32, i1> // In dataflow mode
+      ...
+      neura.yield                       // Function ends here.
   }];
   let arguments = (ins Variadic<AnyType>:$values);
   let assemblyFormat = "($values^ `:` type($values))? attr-dict";