Revise loop_conrtol and replace grant_always operations

include/NeuraDialect/NeuraOps.td

@@ -341,14 +341,17 @@ def Neura_GrantPredicateOp : Op<NeuraDialect, "grant_predicate"> {
 def Neura_GrantOnceOp : Op<NeuraDialect, "grant_once"> {
   let summary = "Marks a value as valid once.";
   let description = [{
-    Grants a value a one-time predicate: the resulting value is considered valid
+    Either grants a value a one-time predicate, or creates a constant value with a one-time predicate. The resulting value is considered valid
     only during its first activation. Used to initialize recurrence cycles.
 
     Example:
       %v = neura.grant_once %init : !neura.data<f32, i1> -> !neura.data<f32, i1>
+      %c = neura.grant_once <{value = 42 : i64}> : !neura.data<i64, i1>
   }];
 
-  let arguments = (ins AnyType:$value);
+  let arguments = (ins 
+  Optional<AnyType>:$value,
+  OptionalAttr<AnyAttr>:$constant_value);
   let results = (outs AnyType:$result);
 
   // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
@@ -364,7 +367,9 @@ def Neura_GrantAlwaysOp : Op<NeuraDialect, "grant_always"> {
       %v = neura.grant_always %init : !neura.data<f32, i1> -> !neura.data<f32, i1>
   }];
 
-  let arguments = (ins AnyType:$value);
+  let arguments = (ins 
+  Optional<AnyType>:$value,
+  OptionalAttr<AnyAttr>:$constant_value);
   let results = (outs AnyType:$result);
 
   // let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)";
@@ -373,31 +378,43 @@ def Neura_GrantAlwaysOp : Op<NeuraDialect, "grant_always"> {
 // ----------------------------------------------------
 // Defines fused control flow operations.
 
-def Neura_LoopControllerOp : Op<NeuraDialect, "loop_controller">{
+def Neura_LoopControlOp : Op<NeuraDialect, "loop_control">{
   let summary = "Generates loop indicies and valid predicates.";
   let description = [{
-    Manages a single level of loop execution based on cycle counting.
-    Each loop_controller outputs a current index value and a valid predicate.
+    Controls loop execution with minimal recurrence cycle length (1).
+    Takes the current index and produces the next index and validity predicate.
 
-    The loop_controller uses dynamic loop bounds (start, end, step),
-    allowing for variable-length loops and runtime-determined bounds.
+    This operation combines comparison and increment logic into a single operation,
+    enabling efficient loop execution on dataflow architectures with modulo scheduling.
 
-    The execution is conditioned on the parent_valid input, creating an
-    efficient hierarchical structure for nested loops.
+    The first iteration automatically uses the start value because the reserve operation's
+    predicate bit is initially 0, while the start value's predicate bit is 1.
+
+    Example:
+      // Loop control that calculates next index and validity in one step
+      %next_idx, %loop_valid = neura.loop_control(
+        parent_valid = %parent_valid,
+        start = %start_val,
+        end = %end_val,
+        step = %step_val
+      ) <{iteration_type = "increment"}> : !neura.data<i1, i1>,
+                                           !neura.data<i64, i1>, !neura.data<i64, i1>, !neura.data<i64, i1>
+                                           -> !neura.data<i64, i1>, !neura.data<i1, i1>
   }];
 
   let arguments = (ins
-    AnyType:$parent_valid,  // Valid predicate from the parent loop
-    AnyType:$start,         // Start index of the loop
-    AnyType:$end,           // End index of the loop
-    AnyType:$step           // Step size for the loop
+    AnyType:$parentValid,  // Valid predicate from the parent loop.
+    StrAttr:$iterationType, // Type of the loop iteration (e.g., "increment", "decrement").
+    AnyType:$start,         // Start index of the loop (optional if startValue attr is presented).
+    AnyType:$end,           // End index of the loop (optional if endValue attr is presented).
+    AnyType:$step           // Step size for the loop (optional if stepValue attr is presented).
   );
 
   let results = (outs
-    AnyType:$index,         // Current loop index
+    AnyType:$nextindex,         // Current loop index
     AnyType:$valid          // Valid predicate for the current index
   );
 
   let assemblyFormat =
-    "$parent_valid `(` $start `,` $end `,` $step `)` attr-dict `:` type($parent_valid) `,` type($start) `,` type($end) `,` type($step) `->` type($index) `,` type($valid)";
+    " `(``parent_valid` `=` $parentValid `,` `start` `=` $start `,` `end` `=` $end `,` `step` `=` $step`)` attr-dict `:` type($parentValid) `,` type($start) `,` type($end) `,` type($step) `->` type($nextindex) `,` type($valid)";
 }

lib/NeuraDialect/Mapping/mapping_util.cpp

@@ -16,10 +16,14 @@ using namespace mlir::neura;
 namespace {
 
 inline OperationKind getOperationKindFromMlirOp(Operation *op) {
-  if (isa<neura::AddOp>(op))   return IAdd;
-  if (isa<neura::MulOp>(op))   return IMul;
-  if (isa<neura::FAddOp>(op))  return FAdd;
-  if (isa<neura::FMulOp>(op))  return FMul;
+  if (isa<neura::AddOp>(op))
+    return IAdd;
+  if (isa<neura::MulOp>(op))
+    return IMul;
+  if (isa<neura::FAddOp>(op))
+    return FAdd;
+  if (isa<neura::FMulOp>(op))
+    return FMul;
   // TODO: Complete the list here.
   // @Jackcuii, https://github.com/coredac/dataflow/issues/82.
   return IAdd;
@@ -127,9 +131,8 @@ int mlir::neura::calculateResMii(Operation *func_op,
   // Count all "compute" operations (non-terminators, non-block ops).
   func_op->walk([&](Operation *op) {
     // Skips non-materialized ops.
-    if (isa<func::FuncOp>(op) || isa<neura::CtrlMovOp,
-                                     neura::DataMovOp,
-                                     neura::ReserveOp>(op)) {
+    if (isa<func::FuncOp>(op) ||
+        isa<neura::CtrlMovOp, neura::DataMovOp, neura::ReserveOp>(op)) {
       return;
     }
     ++num_ops;
@@ -384,13 +387,15 @@ bool mlir::neura::tryRouteBackwardMove(Operation *mov_op, MappingLoc src_loc,
   return tryRouteDataMove(mov_op, src_loc, dst_loc, true, state, path_out);
 }
 
-Register *mlir::neura::getAvailableRegister(
-    const MappingState &state, Tile *tile, int start_time, int exclusive_end_time) {
+Register *mlir::neura::getAvailableRegister(const MappingState &state,
+                                            Tile *tile, int start_time,
+                                            int exclusive_end_time) {
   for (Register *reg : tile->getRegisters()) {
     // FIXME: We may need constrain the register availability to the conflicting
-    // input channel (either the input channel or a register file on the same input
-    // direction could be active at one time).
-    if (state.isAvailableAcrossTimeInRange(reg, start_time, exclusive_end_time)) {
+    // input channel (either the input channel or a register file on the same
+    // input direction could be active at one time).
+    if (state.isAvailableAcrossTimeInRange(reg, start_time,
+                                           exclusive_end_time)) {
       return reg;
     }
   }
@@ -460,7 +465,8 @@ bool mlir::neura::tryRouteDataMove(Operation *mov_op, MappingLoc src_loc,
         //          << " at tile: " << dst_tile->getId()
         //          << " from time step: " << current_step
         //          << " till deadline step: " << deadline_step << "\n";
-        // Register is available, so we can occupy the specific register across remaining time steps.
+        // Register is available, so we can occupy the specific register across
+        // remaining time steps.
         std::vector<MappingLoc> register_occupyings;
         for (int t = current_step; t < deadline_step; ++t) {
           MappingLoc register_loc{available_register, t};
@@ -597,10 +603,10 @@ bool mlir::neura::canReachLocInTime(const MappingLoc &src_loc,
     // // Explores all next step tiles from the current location.
     // for (const MappingLoc &next_loc_tile :
     //      mapping_state.getNextStepTiles(current_loc)) {
-    
+
     // Explores all next step tiles from the current location.
     for (const MappingLoc &current_loc_out_link :
-             mapping_state.getCurrentStepLinks(current_loc)) {
+         mapping_state.getCurrentStepLinks(current_loc)) {
 
       // Makes sure the link is not occupied.
       if (!mapping_state.isAvailableAcrossTime(current_loc_out_link)) {
@@ -614,7 +620,8 @@ bool mlir::neura::canReachLocInTime(const MappingLoc &src_loc,
       }
 
       // Records the tile for further exploration.
-      Tile *next_tile = llvm::dyn_cast<Link>(current_loc_out_link.resource)->getDstTile();
+      Tile *next_tile =
+          llvm::dyn_cast<Link>(current_loc_out_link.resource)->getDstTile();
       assert(next_tile && "Next location must be a Tile");
       if (visited.contains(next_tile)) {
         continue;
@@ -687,6 +694,7 @@ std::vector<MappingLoc> mlir::neura::calculateAward(Operation *op,
 
   llvm::errs() << "[calculateAward] Operation: " << *op
                << "; Producers: " << producers.size() << "\n";
+
   for (Tile *tile : architecture.getAllTiles()) {
     if (!tile->canSupportOperation(getOperationKindFromMlirOp(op))) {
       llvm::errs() << "[calculateAward] Tile: " << tile->getType()
@@ -815,8 +823,8 @@ bool mlir::neura::placeAndRoute(Operation *op, const MappingLoc &target_loc,
   if (mapping_state.bindOp(target_loc, op)) {
     std::vector<Operation *> routed_operands;
     std::vector<Operation *> routed_ctrl_movs;
-    llvm::errs() << "[DEBUG] Schedule op " << *op << " onto loc: "
-                 << target_loc.resource->getType() << "#"
+    llvm::errs() << "[DEBUG] Schedule op " << *op
+                 << " onto loc: " << target_loc.resource->getType() << "#"
                  << target_loc.resource->getId()
                  << " @t=" << target_loc.time_step << "\n";
     // Tries to route the data move operations.
@@ -886,10 +894,12 @@ bool mlir::neura::placeAndRoute(Operation *op, const MappingLoc &target_loc,
       }
       llvm::errs() << "[DEBUG] Failed to route ctrl_mov: " << *ctrl_mov
                    << " from " << target_loc.resource->getType() << "#"
-                   << target_loc.resource->getId() << " @t=" << target_loc.time_step
-                   << " to " << backward_loc.resource->getType() << "#"
+                   << target_loc.resource->getId()
+                   << " @t=" << target_loc.time_step << " to "
+                   << backward_loc.resource->getType() << "#"
                    << backward_loc.resource->getId()
-                   << " @t=" << backward_loc.time_step << "; so unschedule op\n";
+                   << " @t=" << backward_loc.time_step
+                   << "; so unschedule op\n";
       mapping_state.unbindOp(op);
       for (Operation *routed_ctrl_mov : routed_ctrl_movs) {
         llvm::errs() << "[DEBUG] Releasing route for routed ctrl_mov: "