triton-lang · m-everglow · Sep 1, 2026
@@ -1849,7 +1849,9 @@ static void buildBeforeRegion(scf::WhileOp oldWhile, OpBuilder &bb, Location bl,
   bb.create<scf::ConditionOp>(bl, condValue, carriedValues);
 }
 
-// Build the after-region of the new whileOp
+// Build the after-region of the new whileOp. The counter add-1 is not
+// created here; it is inserted later by insertWhileCounterOps at the right
+// position. counterIterArgOut is yielded as a placeholder and gets replaced.
 static void buildAfterRegion(scf::WhileOp oldWhile, OpBuilder &ab, Location al,
                              ValueRange iterArgs, Value &counterIterArgOut) {
   Block *oldAfter = oldWhile.getAfterBody();
@@ -1871,35 +1873,12 @@ static void buildAfterRegion(scf::WhileOp oldWhile, OpBuilder &ab, Location al,
   if (!oldYield)
     return;
 
-  std::optional<int> counterBlockId;
-  if (Block *doBlock = ab.getInsertionBlock()) {
-    for (Operation &op : llvm::reverse(*doBlock)) {
-      if (auto id = getOpBlockId(&op); id.has_value()) {
-        counterBlockId = id;
-        break;
-      }
-    }
-  }
-  if (!counterBlockId)
-    counterBlockId = getOpBlockId(oldWhile);
-
-  Value one = ab.create<arith::ConstantIntOp>(al, 1, 32);
-  Value nextCounter = ab.create<arith::AddIOp>(al, counterIterArgOut, one);
-  nextCounter.getDefiningOp()->setAttr(kIterCounter, ab.getUnitAttr());
-
-  if (counterBlockId) {
-    one.getDefiningOp()->setAttr(kBlockId,
-                                 ab.getI32IntegerAttr(*counterBlockId));
-    nextCounter.getDefiningOp()->setAttr(kBlockId,
-                                         ab.getI32IntegerAttr(*counterBlockId));
-  }
-
   SmallVector<Value> newYieldOps;
   for (Value operand : oldYield->getOperands()) {
     Value mapped = mapper.lookupOrNull(operand);
     newYieldOps.push_back(mapped ? mapped : operand);
   }
-  newYieldOps.push_back(nextCounter);
+  newYieldOps.push_back(counterIterArgOut);
   ab.create<scf::YieldOp>(al, newYieldOps);
 }
 
@@ -1952,6 +1931,77 @@ setupWhileIterArgCounter(const MainLoop &loop, OpBuilder &builder) {
   return {counterIterArg, newWhile};
 }
 
+// Insert the iterCounter add-1 (and its constant 1) right after the last op
+// sharing the block_id of the first op that consumes mainLoop.iterCounter.
+// Must run after processTensorDependencies (dispatch ops land there later).
+// Fallback when no counter user exists: last op's block_id, mirroring the
+// legacy buildAfterRegion placement.
+static void insertWhileCounterOps(const MainLoop &mainLoop) {
+  if (!mainLoop.isWhile() || !mainLoop.iterCounter)
+    return;
+
+  auto whileOp = cast<scf::WhileOp>(mainLoop.getOperation());
+  Block *doBlock = whileOp.getAfterBody();
+  if (!doBlock)
+    return;
+
+  std::optional<int> counterBlockId;
+  for (Operation &op : *doBlock) {
+    if (llvm::is_contained(op.getOperands(), mainLoop.iterCounter)) {
+      counterBlockId = getOpBlockId(&op);
+      break;
+    }
+  }
+  if (!counterBlockId) {
+    for (Operation &op : llvm::reverse(*doBlock)) {
+      if (auto id = getOpBlockId(&op); id.has_value()) {
+        counterBlockId = id;
+        break;
+      }
+    }
+    if (!counterBlockId)
+      return;
+  }
+
+  Operation *lastWithBlockId = nullptr;
+  for (Operation &op : *doBlock) {
+    auto id = getOpBlockId(&op);
+    if (id.has_value() && *id == *counterBlockId)
+      lastWithBlockId = &op;
+  }
+
+  Location loc = lastWithBlockId->getLoc();
+  OpBuilder builder(mainLoop.getContext());
+  builder.setInsertionPointAfter(lastWithBlockId);
+  IntegerAttr blockIdAttr = builder.getI32IntegerAttr(*counterBlockId);
+
+  Value one = builder.create<arith::ConstantIntOp>(
+      loc, 1, mainLoop.iterCounter.getType().getIntOrFloatBitWidth());
+  one.getDefiningOp()->setAttr(kBlockId, blockIdAttr);
+
+  Value nextCounter =
+      builder.create<arith::AddIOp>(loc, mainLoop.iterCounter, one);
+  Operation *iterAddOp = nextCounter.getDefiningOp();
+  iterAddOp->setAttr(kBlockId, blockIdAttr);
+  iterAddOp->setAttr(kIterCounter, builder.getUnitAttr());
+
+  // Replace the placeholder counterIterArgOut operand in the yield with
+  // nextCounter so the iter_arg increments each iteration.
+  Operation *yieldOp = doBlock->getTerminator();
+  SmallVector<Value> newOperands(yieldOp->getOperands().begin(),
+                                 yieldOp->getOperands().end());
+  bool replaced = false;
+  for (Value &operand : newOperands) {
+    if (operand == mainLoop.iterCounter) {
+      operand = nextCounter;
+      replaced = true;
+      break;
+    }
+  }
+  if (replaced)
+    yieldOp->setOperands(newOperands);
+}
+
 static int addInnerMultiBuffer(MainLoop mainLoop, OpBuilder &builder,
                                scope::ScopeOp vectorScope, int &groupId,
                                bool &i1Found) {
@@ -2087,6 +2137,10 @@ static int addInnerMultiBuffer(MainLoop mainLoop, OpBuilder &builder,
     return -1;
   }
 
+  // WhileOp only: insert the counter add-1 now that the dispatch ops that
+  // use mainLoop.iterCounter have been emitted.
+  insertWhileCounterOps(mainLoop);
+
   LLVM_DEBUG(llvm::dbgs() << "[addInnerMultiBuffer] DONE\n");
   return 0;
 }

@@ -42,8 +42,7 @@ namespace {
 
 static constexpr llvm::StringLiteral interceptrFunc[]{
     "chunk_abc_bwd_kernel_dh",      "flash_varlen_fwd_kernel",
-    "chunk_gsa_bwd_k_kernel_dqkvg", "_jagged_flash_attention_bwd_basic_kernel",
-    "_sparse_decode_kernel",        "chunk_gsa_fwd_k_kernel_intra",
+    "chunk_gsa_bwd_k_kernel_dqkvg", "_sparse_decode_kernel",
     "_sparse_decode_model1_kernel", "sparse_flash_attention_grad_kernel"};
 
 static LogicalResult verifyFuncNames(ModuleOp module) {

@@ -35,21 +35,24 @@
 //       (3 block-args instead of 2; the new arg is i32).
 //     - arith.addi counter, 1 is present in the do-region for the yield.
 //     - Multi-buffer producer/consumer with scf.if dispatch works as forOp.
+//     - The counter add-1 op (and its constant-1) is relocated to the
+//       block_id of the first op that consumes the counter arg (here: the
+//       producer scf.if at block_id = 7), not its old fallback position.
 
 // CHECK-LABEL: func.func @test_while_mainloop_bufnum_two
 // Two UB allocs (ping/pong).
 // CHECK-DAG:   memref.alloc() : memref<128xf32, #hivm.address_space<ub>>
 // CHECK-DAG:   memref.alloc() : memref<128xf32, #hivm.address_space<ub>>
 // WhileOp's do-region bb0 has 3 block-args now (the new i32 counter is the last one).
 // CHECK:       ^bb0(%{{.*}}: tensor<128xf32>, %{{.*}}: i32, %{{.*}}: i32):
-// Producer scf.if dispatch (the original counter increment lives inside this region).
+// Producer scf.if dispatch.
 // CHECK:       scf.if
 // CHECK:         hivm.hir.copy
+// arith.addi increment for the multi-buffer counter (block_id=7, the first user's block).
+// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 7 : i32, ssbuffer.iterCounter} : i32
 // Consumer scf.if dispatch returning tensor.
 // CHECK:       scf.if {{.*}} -> (tensor<128xf32>)
 // CHECK:         bufferization.to_tensor
-// arith.addi increment for the multi-buffer counter is present (block_id=10), tagged iterCounter.
-// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 10 : i32, ssbuffer.iterCounter} : i32
 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop.
 // CHECK:       } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64
 
@@ -63,11 +66,11 @@
 // Producer alloc + copy before the whileOp.
 // CHECK-DAG:   memref.alloc() : memref<64xf16, #hivm.address_space<ub>>
 // CHECK-DAG:   memref.alloc() : memref<64xf16, #hivm.address_space<ub>>
+// arith.addi increment for the multi-buffer counter (block_id=7, the first user's block).
+// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 7 : i32, ssbuffer.iterCounter} : i32
 // Consumer scf.if + to_tensor inside the do-region.
 // CHECK:       scf.if {{.*}} -> (tensor<64xf16>)
 // CHECK:         bufferization.to_tensor
-// arith.addi increment for the multi-buffer counter is present (block_id=12), tagged iterCounter.
-// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 12 : i32, ssbuffer.iterCounter} : i32
 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop.
 // CHECK:       } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64
 
@@ -90,8 +93,8 @@
 // Producer-side dispatch.
 // CHECK:       scf.if
 // CHECK:         hivm.hir.copy
-// arith.addi increment for the multi-buffer counter is present (block_id=12), tagged iterCounter.
-// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 12 : i32, ssbuffer.iterCounter} : i32
+// arith.addi increment for the multi-buffer counter (block_id=8, the first user's block).
+// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 8 : i32, ssbuffer.iterCounter} : i32
 // Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop.
 // CHECK:       } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64
 
@@ -121,6 +124,34 @@
 //   Not implemented in this initial drop — to be added when the negative case
 //   is needed for a regression guard.
 
+// T-while-G: Regression guard for relocateWhileCounterOps.
+//   The counter add-1 (and its constant 1) used to be appended at the end of
+//   the do-region, carrying the whileOp's own block_id (= 99 here, distinct
+//   from both producer and consumer). The pass now relocates them to the
+//   block_id of the first op that consumes the counter iter_arg — the
+//   producer scf.if at block_id = 5.
+//   We assert:
+//     - The counter add-1 carries block_id = 5 (NOT 99, NOT 50).
+//     - The constant 1 used by the addi also carries block_id = 5.
+//     - The scf.yield that consumes the add-1's SSA value still references
+//     // the same %result (SSA is preserved across the move).
+
+// CHECK-LABEL: func.func @test_while_counter_relocation
+// Multi-buffer producer (two allocs at block 5) emits the dispatch.
+// CHECK-DAG:   memref.alloc() : memref<16xf32, #hivm.address_space<ub>>
+// CHECK-DAG:   memref.alloc() : memref<16xf32, #hivm.address_space<ub>>
+// CHECK:       scf.if
+// CHECK:         hivm.hir.copy
+// Constant 1 carrying block_id = 5 (relocated with the addi).
+// CHECK:       %{{.+}} = arith.constant {ssbuffer.block_id = 5 : i32} 1 : i32
+// Counter add-1 carries block_id = 5 and is tagged iterCounter.
+// CHECK:       %{{.+}} = arith.addi %{{.+}}, %{{.+}} {ssbuffer.block_id = 5 : i32, ssbuffer.iterCounter} : i32
+// scf.yield forwards the counter (referenced by the addi's %addi-res), proving
+// SSA is preserved.
+// CHECK:       scf.yield %{{.*}}, %{{.*}}, %{{.+}}
+// Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop.
+// CHECK:       } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64
+
 // ---- Inputs ----
 
 module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">,
@@ -277,4 +308,37 @@ module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">,
     } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
     return
   }
+
+  // T-while-G: dedicated regression test for relocateWhileCounterOps.
+  // Producer at block_id = 5, consumer at block_id = 50, whileOp at block_id = 99.
+  // Before the fix: counter add-1 would land at block_id = 99 (whileOp's own id).
+  // After the fix: counter add-1 (and its constant 1) land at block_id = 5
+  // (the block_id of the first op that consumes the counter iter_arg).
+  func.func @test_while_counter_relocation() {
+    %c0_i32 = arith.constant 0 : i32
+    %c10_i32 = arith.constant 10 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %cst_zero = arith.constant 0.0 : f32
+    %init = tensor.empty() : tensor<16xf32>
+    %carry = linalg.fill ins(%cst_zero : f32) outs(%init : tensor<16xf32>) -> tensor<16xf32>
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32)
+          : (tensor<16xf32>, i32) -> (tensor<16xf32>, i32) {
+        %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : tensor<16xf32>, i32
+      } do {
+      ^bb0(%arg0: tensor<16xf32>, %arg1: i32):
+        // Producer block_id = 5.
+        %alloc = memref.alloc() {ssbuffer.block_id = 5 : i32} : memref<16xf32>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 5 : i32} : memref<16xf32> to tensor<16xf32>
+        // Consumer block_id = 50 (very different from 5, to make the relocation
+        // verdict unambiguous).
+        %consumed = arith.addf %prod, %prod {ssbuffer.block_id = 50 : i32} : tensor<16xf32>
+        %next = arith.addi %arg1, %c1_i32 : i32
+        scf.yield %consumed, %next : tensor<16xf32>, i32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 99 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
 }