From b70ff45ec18e7f651bf15d7ad84300b1f6473afb Mon Sep 17 00:00:00 2001
From: cxtverygood123 <chenxiaotian12@huawei.com>
Date: Mon, 20 Jul 2026 18:01:45 +0800
Subject: [PATCH 01/11] [ssbuffer](feat) support whileOP in innerscope

fix cleancode
---
 .../include/DynamicCVPipeline/Common/Utils.h  |  40 ++
 .../AddMultiBufferInnerScope.cpp              | 428 +++++++++++-------
 .../lib/DynamicCVPipeline/Common/Utils.cpp    |  52 +++
 .../Inner-scope-memref-dep.mlir               |  12 +-
 .../AllocMultiCache/Inner-scope-whileop.mlir  | 280 ++++++++++++
 .../AllocMultiCache/Inner_scope_i1_test.mlir  |  12 +-
 6 files changed, 659 insertions(+), 165 deletions(-)
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir

diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
index 1b39284f24..20f0a036c8 100644
--- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
+++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
@@ -23,6 +23,7 @@
 #ifndef ADD_AUTO_SCHEDULING_COMMON_UTILS_H
 #define ADD_AUTO_SCHEDULING_COMMON_UTILS_H
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
@@ -61,6 +62,7 @@ inline constexpr llvm::StringLiteral kCrossCoreDeps = "ssbuffer.crossCoreDeps";
 inline constexpr llvm::StringLiteral kIntraDeps = "ssbuffer.intraDeps";
 inline constexpr llvm::StringLiteral kMemCrossDeps = "ssbuffer.memCrossDeps";
 inline constexpr llvm::StringLiteral kMayNotExec = "ssbuffer.may_not_exec";
+inline constexpr llvm::StringLiteral kIterCounter = "ssbuffer.iterCounter";
 inline constexpr llvm::StringLiteral kClone = "ssbuffer.clone";
 inline constexpr llvm::StringLiteral kEnableUbRefineOpt =
     "ssbuffer.enable_ub_refine_opt";
@@ -121,6 +123,44 @@ bool isScfOp(Operation *op);
 bool isOnlyDirectlyUse(Operation *preOp, Operation *nextOp,
                        const CVPipeline::MemoryDependenceGraph &memGraph);
 
+// Wrapper around a "main loop" — either scf.for or scf.while carrying the
+// ssbuffer.main_loop attribute. Lets downstream code treat both uniformly.
+class MainLoop {
+public:
+  Operation *op = nullptr;
+  Block *body = nullptr;
+  Value iterCounter;
+
+  Block *getBody() const;
+  Operation *getOperation() const;
+  MLIRContext *getContext() const;
+  Location getLoc() const;
+  Block *getBlock() const;
+  Block::iterator getIterator() const;
+  Operation *operator->() const;
+  bool isWhile() const;
+
+  // Iter args carried across loop iterations, as BlockArguments.
+  // forOp:   getRegionIterArgs().
+  // whileOp: after-body args.
+  SmallVector<Value> getIterArgs() const;
+
+  // Only meaningful for whileOp (before-body args); forOp returns empty.
+  // Same count/types as getIterArgs() on whileOp, distinct Value identity.
+  SmallVector<Value> getBeforeIterArgs() const;
+
+  explicit MainLoop(Operation *loopOp);
+
+  // Returns the scf.yield terminator of a forOp's body / whileOp's after
+  // body. Returns {} if `loopOp` is neither.
+  static scf::YieldOp getLoopYieldOp(Operation *loopOp);
+};
+
+// True when `op` is a main_loop loop op (forOp / whileOp carrying the tag).
+inline bool isMainLoopOp(Operation *op) {
+  return op && isa<scf::ForOp, scf::WhileOp>(op) && op->hasAttr(kMainLoop);
+}
+
 inline bool isCubeOp(Operation *op) {
   return !isScfOp(op) && CVPipeline::getOpCoreType(op) == CoreType::CUBE_ONLY;
 }
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp
index d1ca2bf8c9..7b6aed113f 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp
@@ -57,63 +57,31 @@ using BufferMap = DenseMap<Value, SmallVector<BufferPair>>;
 // Buffer count constants
 constexpr int kBufferCountOne = 1;
 
-// Toggle: when true, producer / consumer buffer chains are inserted at the
-// boundary of the contiguous `ssbuffer.block_id = X` region instead of right
-// after the dep def op / right before the dep user op.
-//
-//   - producer chain: inserted AFTER the last op in depDefinedOp's block that
-//     carries the same `ssbuffer.block_id` as depDefinedOp. If depDefinedOp has
-//     no direct block_id attribute, falls back to "right after depDefinedOp".
-//   - consumer chain: inserted BEFORE the first op in depUser's block that
-//     carries the same `ssbuffer.block_id` as depUser. If depUser has no direct
-//     block_id attribute, falls back to "right before depUser".
-//
-// This keeps the producer/consumer chains grouped with their block_id region
-// rather than interleaved with intermediate compute ops.
-//
-// The toggle is read from a module-level attribute
-// `CVPipeline::kInsertionOptimization` (i.e. "ssbuffer.insertionOptimization").
-// Python callers set it via `set_enable_buffer_insert_optimization` in
-// `triton_ascend.cc`, which writes the attribute onto the ModuleOp. The
-// attribute is checked inline at each `processDepVal` call site below.
 namespace mlir {
 namespace triton {
 
-// Check if forOp has main_loop attribute
-static bool hasMainLoopAttr(scf::ForOp forOp) {
-  if (forOp->hasAttr(kMainLoop)) {
-    return true;
-  }
-  if (auto *term = forOp.getBody()->getTerminator())
-    return term->hasAttr(kMainLoop);
-  return false;
-}
-
-// Collect main_loop forOps in a single block
+// Collect main_loop loops (forOp / whileOp) in a single block
 static int collectMainLoopsInBlock(Block &block,
-                                   SmallVector<scf::ForOp> &mainLoopForOps) {
+                                   SmallVector<Operation *> &mainLoops) {
   int count = 0;
   for (Operation &op : block) {
-    if (auto forOp = dyn_cast<scf::ForOp>(&op)) {
-      if (hasMainLoopAttr(forOp)) {
-        mainLoopForOps.push_back(forOp);
-        count++;
-      }
+    if (isMainLoopOp(&op)) {
+      mainLoops.push_back(&op);
+      count++;
     }
   }
   return count;
 }
 
-// Recursively collect main_loop forOps, returns count of collected items
-static int
-collectMainLoopsRecursively(Region &region,
-                            SmallVector<scf::ForOp> &mainLoopForOps) {
+// Recursively collect main_loop loops, returns count of collected items
+static int collectMainLoopsRecursively(Region &region,
+                                       SmallVector<Operation *> &mainLoops) {
   int totalCount = 0;
   for (Block &block : region) {
-    totalCount += collectMainLoopsInBlock(block, mainLoopForOps);
+    totalCount += collectMainLoopsInBlock(block, mainLoops);
     for (Operation &op : block) {
       for (auto &nestedRegion : op.getRegions())
-        totalCount += collectMainLoopsRecursively(nestedRegion, mainLoopForOps);
+        totalCount += collectMainLoopsRecursively(nestedRegion, mainLoops);
     }
   }
   return totalCount;
@@ -124,23 +92,15 @@ struct InnerBlockInfo {
   SmallVector<Operation *> ops;
 };
 
-// Effective block_id for cross-block dep judgment: prefer the enclosing
-// multi-region op's block_id (inner ops are attributed to the op itself,
-// not their own innermost id), else the innermost recorded block_id walking
-// up to the main_loop boundary.
 static std::optional<int64_t> getOutermostSsbufferId(Operation *op) {
   std::optional<int64_t> result;
   for (Operation *current = op; current; current = current->getParentOp()) {
-    // Any multi-region op (scf.if, scf.while, ...) acts as a logical
-    // block boundary: its block_id overrides anything inside it.
-    if (current->getNumRegions() >= 2)
-      return getOpBlockId(current);
-
-    // main_loop is an attribute, not exclusive to forOp. Take the
-    // boundary's id only if nothing was recorded on the way up.
     if (current->hasAttr(kMainLoop))
       return result.has_value() ? result : -1;
 
+    if (current->getNumRegions() >= 2)
+      return getOpBlockId(current);
+
     // Otherwise remember the deepest id seen; the parent walk will
     // overwrite it if a closer-to-boundary op carries one.
     if (auto curId = getOpBlockId(current); curId.has_value())
@@ -242,6 +202,7 @@ static void collectDepValue(Value operand, Block *body, Operation *currentOp,
 
   auto currentOutermost = getOutermostSsbufferId(currentOp);
   auto operandOutermost = getOutermostSsbufferId(operand.getDefiningOp());
+
   if (currentOutermost.has_value() && currentOutermost == operandOutermost)
     return;
 
@@ -261,41 +222,27 @@ static void collectDepValue(Value operand, Block *body, Operation *currentOp,
     depValueMap[groupKey].push_back(operand);
 }
 
-// Recursively find nested main_loop
-static scf::ForOp findNestedMainloopInForOp(scf::ForOp forOp) {
+// Recursively find a nested main_loop (forOp / whileOp) inside `loop`'s body
+static Operation *findNestedMainloop(const MainLoop &loop) {
   SmallVector<Operation *> allOps;
-  collectNestedOps(forOp.getBody(), allOps);
+  collectNestedOps(loop.getBody(), allOps);
 
   for (Operation *op : allOps) {
-    auto nestedFor = dyn_cast<scf::ForOp>(op);
-    if (!nestedFor)
-      continue;
-    if (nestedFor->hasAttr(kMainLoop))
-      return nestedFor;
+    if (isa<scf::ForOp, scf::WhileOp>(op) && op->hasAttr(kMainLoop))
+      return op;
   }
   return {};
 }
 
 bool isInsideMainLoopForOp(Operation *op) {
-  Operation *parent = op->getParentOp();
-  if (!parent) {
-    return false;
-  }
-  if (auto forOp = dyn_cast<scf::ForOp>(parent)) {
-    return forOp->hasAttr(kMainLoop);
-  }
-  return false;
+  return isMainLoopOp(op->getParentOp());
 }
 
 bool isInsideMainLoopForOpTraverse(Operation *op) {
-  Operation *parent = op->getParentOp();
-  while (parent) {
-    if (auto forOp = dyn_cast<scf::ForOp>(parent)) {
-      if (forOp->hasAttr(kMainLoop)) {
-        return true;
-      }
-    }
-    parent = parent->getParentOp();
+  for (Operation *parent = op->getParentOp(); parent;
+       parent = parent->getParentOp()) {
+    if (isMainLoopOp(parent))
+      return true;
   }
   return false;
 }
@@ -373,11 +320,12 @@ forEachYieldedCrossBlockDep(Operation *op,
 // dep collected here has element type i1; the caller is expected to abort and
 // trigger fallback in that case.
 static int
-collectInnerBlockInfo(scf::ForOp forOp, DenseMap<Value, InnerBlockInfo> &blocks,
+collectInnerBlockInfo(const MainLoop &loop,
+                      DenseMap<Value, InnerBlockInfo> &blocks,
                       DenseMap<Value, SmallVector<Value>> &depValueMap,
                       SmallVector<Operation *> &allOps, bool &i1Found) {
   depValueMap.clear();
-  Block *body = forOp.getBody();
+  Block *body = loop.getBody();
   if (!body)
     return 0;
 
@@ -589,7 +537,7 @@ collectScalarDeps(DenseMap<Value, SmallVector<Value>> &depValueMap,
 }
 
 // True if op is nested strictly inside the main loop.
-static bool isOpInMainLoop(Operation *op, scf::ForOp mainLoop) {
+static bool isOpInMainLoop(Operation *op, const MainLoop &mainLoop) {
   return op && mainLoop.getOperation()->isProperAncestor(op);
 }
 
@@ -610,7 +558,7 @@ static void collectOpDependencies(Operation *op, SmallVector<Value> &deps) {
 
 // Depth-first build of the scalar op slice feeding `root`. Recursion stops at
 // tensor operands
-static void buildScalarSlice(Value root, scf::ForOp mainLoop,
+static void buildScalarSlice(Value root, const MainLoop &mainLoop,
                              SmallVector<Operation *> &sliceInOrder,
                              DenseSet<Operation *> &visited,
                              llvm::SetVector<Value> &boundaryTensors) {
@@ -651,7 +599,7 @@ static Operation *getAncestorInBlock(Operation *op, Block *block) {
 // consumer blocks and rewire those consumers to the local copy. Returns true on
 // rewrite.
 static bool
-rematerializeScalarDep(Value root, int producerId, scf::ForOp mainLoop,
+rematerializeScalarDep(Value root, int producerId, const MainLoop &mainLoop,
                        const SmallVector<Operation *> &sliceInOrder) {
   Block *body = mainLoop.getBody();
 
@@ -719,7 +667,7 @@ rematerializeScalarDep(Value root, int producerId, scf::ForOp mainLoop,
 // Scan the main loop for cross-block scalar dependencies whose data originates
 // from a tensor, and rematerialize the scalar portion into each consumer block
 // so the tensor part can use the normal tensor-dependency buffering.
-static void rematerializeTensorRootedScalarDeps(scf::ForOp mainLoop) {
+static void rematerializeTensorRootedScalarDeps(const MainLoop &mainLoop) {
   Block *body = mainLoop.getBody();
   if (!body) {
     return;
@@ -769,12 +717,18 @@ static void rematerializeTensorRootedScalarDeps(scf::ForOp mainLoop) {
   }
 }
 
-// Compute iteration index: (iv - lb) / step, used for buffer selection in
-// double buffering
-static Value getIterCount(OpBuilder &builder, mlir::scf::ForOp forOp,
+static Value getIterCount(OpBuilder &builder, const MainLoop &loop,
                           Location loc, SmallVector<Operation *> *newOps,
                           int blockId = -1) {
   auto i32Type = builder.getI32Type();
+
+  if (loop.isWhile()) {
+    assert(loop.iterCounter &&
+           "whileOp main_loop requires a global iteration counter");
+    return loop.iterCounter;
+  }
+
+  auto forOp = cast<scf::ForOp>(loop.getOperation());
   Value iv = forOp.getInductionVar();
   Value lb = forOp.getLowerBound();
   Value step = forOp.getStep();
@@ -1032,11 +986,11 @@ buildIfChain(OpBuilder &builder, Location loc, Value indexVal,
 }
 
 // Compute buffer index: iterCount % N
-static Value computeBufferIndex(OpBuilder &builder, mlir::scf::ForOp forOp,
+static Value computeBufferIndex(OpBuilder &builder, const MainLoop &loop,
                                 Location loc, int N,
                                 SmallVector<Operation *> *newOps,
                                 int blockId = -1) {
-  Value iterCount = getIterCount(builder, forOp, loc, newOps, blockId);
+  Value iterCount = getIterCount(builder, loop, loc, newOps, blockId);
   Value Nval = builder.create<mlir::arith::ConstantIntOp>(loc, N, 32);
   Value bufIdx = builder.create<mlir::arith::RemSIOp>(loc, iterCount, Nval);
   if (newOps) {
@@ -1055,7 +1009,7 @@ static Value computeBufferIndex(OpBuilder &builder, mlir::scf::ForOp forOp,
 
 static SmallVector<Operation *>
 insertProducerLogic(OpBuilder &builder, Value depVal,
-                    SmallVector<BufferPair> &buffers, mlir::scf::ForOp forOp,
+                    SmallVector<BufferPair> &buffers, const MainLoop &loop,
                     int groupId = -1) {
   SmallVector<Operation *> newOps;
   int N = buffers.size();
@@ -1074,7 +1028,7 @@ insertProducerLogic(OpBuilder &builder, Value depVal,
     return newOps;
   }
 
-  Value bufIdx = computeBufferIndex(builder, forOp, loc, N, &newOps);
+  Value bufIdx = computeBufferIndex(builder, loop, loc, N, &newOps);
   SmallVector<Operation *> outIfOps;
   if (buildIfChain(
           builder, loc, bufIdx, buffers, newOps, outIfOps,
@@ -1121,7 +1075,7 @@ static mlir::bufferization::ToTensorOp createToTensorOp(OpBuilder &builder,
 
 static int insertConsumerLogic(OpBuilder &builder, Value depVal,
                                SmallVector<BufferPair> &buffers,
-                               mlir::scf::ForOp forOp,
+                               const MainLoop &loop,
                                SmallVector<Operation *> &outIfOps,
                                int groupId = -1, int blockId = -1) {
   SmallVector<Operation *> newOps;
@@ -1137,7 +1091,7 @@ static int insertConsumerLogic(OpBuilder &builder, Value depVal,
     return 0;
   }
 
-  Value readIdx = computeBufferIndex(builder, forOp, loc, N, &newOps, blockId);
+  Value readIdx = computeBufferIndex(builder, loop, loc, N, &newOps, blockId);
   auto memrefType = mlir::cast<mlir::MemRefType>(buffers[0].second.getType());
   auto tensorType = mlir::RankedTensorType::get(memrefType.getShape(),
                                                 memrefType.getElementType());
@@ -1217,7 +1171,7 @@ collectCrossBlockUsers(Value depVal, int producerId,
 static Operation *
 insertBufferSelectionInRegion(OpBuilder &builder, Region &region, Location loc,
                               Value depVal, SmallVector<BufferPair> &buffers,
-                              mlir::scf::ForOp forOp, int blockId) {
+                              const MainLoop &loop, int blockId) {
   auto memrefType = mlir::cast<mlir::MemRefType>(buffers[0].second.getType());
   auto tensorType = mlir::RankedTensorType::get(memrefType.getShape(),
                                                 memrefType.getElementType());
@@ -1227,7 +1181,7 @@ insertBufferSelectionInRegion(OpBuilder &builder, Region &region, Location loc,
 
   // Compute buffer index
   Value readIdx =
-      computeBufferIndex(builder, forOp, loc, buffers.size(), nullptr, blockId);
+      computeBufferIndex(builder, loop, loc, buffers.size(), nullptr, blockId);
 
   // Build buffer selection if-else chain
   SmallVector<Operation *> newIfOps;
@@ -1308,12 +1262,12 @@ static bool isMultiRegionConsumerFromYield(Operation *depUser, Value depVal) {
 // among all ops in the block
 static int processNormalConsumerBlock(OpBuilder &consumedBuilder, Value depVal,
                                       SmallVector<BufferPair> &buffers,
-                                      mlir::scf::ForOp mainLoopForOp,
+                                      const MainLoop &loop,
                                       SmallVector<Operation *> &opsInBlock,
                                       int userBlockId, int groupId,
                                       OpBuilder &globalBuilder) {
   SmallVector<Operation *> resultIfOps;
-  int ret = insertConsumerLogic(consumedBuilder, depVal, buffers, mainLoopForOp,
+  int ret = insertConsumerLogic(consumedBuilder, depVal, buffers, loop,
                                 resultIfOps, groupId, userBlockId);
   if (ret != 0)
     return -1;
@@ -1352,9 +1306,8 @@ static int processNormalConsumerBlock(OpBuilder &consumedBuilder, Value depVal,
 // from index 1 onwards
 static int processMultiRegionAllYields(OpBuilder &consumedBuilder, Value depVal,
                                        SmallVector<BufferPair> &buffers,
-                                       mlir::scf::ForOp mainLoopForOp,
-                                       Operation *depUser, int userBlockId,
-                                       int groupId) {
+                                       const MainLoop &loop, Operation *depUser,
+                                       int userBlockId, int groupId) {
   // Generic: check if op has >= 2 regions
   if (depUser->getNumRegions() < 2)
     return 0;
@@ -1374,8 +1327,8 @@ static int processMultiRegionAllYields(OpBuilder &consumedBuilder, Value depVal,
         continue;
 
       Operation *selectIf = insertBufferSelectionInRegion(
-          consumedBuilder, region, yieldOp.getLoc(), depVal, buffers,
-          mainLoopForOp, userBlockId);
+          consumedBuilder, region, yieldOp.getLoc(), depVal, buffers, loop,
+          userBlockId);
       if (!selectIf)
         return -1;
 
@@ -1427,7 +1380,7 @@ static Operation *findFirstOpWithBlockIdInBlock(Operation *anchorOp,
 }
 
 // Process producer and consumer for a single dependency value
-static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
+static int processDepVal(Value depVal, const MainLoop &loop,
                          BufferMap &bufferMap,
                          DenseMap<Value, SmallVector<Operation *>> &depUserMap,
                          OpBuilder &globalBuilder, int producerId,
@@ -1447,11 +1400,11 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
   // processDepVal can be called multiple times in the same pass run and stay
   // in sync with whatever the Python caller last wrote onto the ModuleOp.
   bool enableOpt = false;
-  if (mlir::ModuleOp mod = mainLoopForOp->getParentOfType<mlir::ModuleOp>())
+  if (mlir::ModuleOp mod = loop->getParentOfType<mlir::ModuleOp>())
     enableOpt = mod->hasAttr(CVPipeline::kInsertionOptimization);
 
   // Create producer
-  OpBuilder producedBuffers(mainLoopForOp.getContext());
+  OpBuilder producedBuffers(loop.getContext());
   // When enable_buffer_insert_optimization is on, place the producer chain at
   // the end of depDefinedOp's block_id=X region (after the last op with that
   // block_id). Otherwise keep the original "right after depDefinedOp" anchor.
@@ -1464,8 +1417,8 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
     }
   }
   producedBuffers.setInsertionPointAfter(producerAnchor);
-  SmallVector<Operation *> producerNewOps = insertProducerLogic(
-      producedBuffers, depVal, buffers, mainLoopForOp, groupId);
+  SmallVector<Operation *> producerNewOps =
+      insertProducerLogic(producedBuffers, depVal, buffers, loop, groupId);
   addBlockAttrForOps(producerNewOps, producerId, globalBuilder);
   if (buffers.size() > kBufferCountOne) {
     for (auto *op : producerNewOps) {
@@ -1490,7 +1443,7 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
 
     if (isMultiRegionConsumerFromYield(depUser, depVal)) {
       // Multi-region op: process independently
-      OpBuilder consumedBuilder(mainLoopForOp.getContext());
+      OpBuilder consumedBuilder(loop.getContext());
       // When enable_buffer_insert_optimization is on, place the consumer chain
       // at the start of depUser's block_id=X region (before the first op with
       // that block_id). Otherwise keep "right before depUser".
@@ -1504,9 +1457,9 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
       }
       consumedBuilder.setInsertionPoint(consumerAnchor);
 
-      if (int ret = processMultiRegionAllYields(consumedBuilder, depVal,
-                                                buffers, mainLoopForOp, depUser,
-                                                *userBlockId, groupId))
+      if (int ret =
+              processMultiRegionAllYields(consumedBuilder, depVal, buffers,
+                                          loop, depUser, *userBlockId, groupId))
         return ret;
     } else {
       // Normal op: collect by block_id for batch processing
@@ -1538,7 +1491,7 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
         continue;
 
       Operation *firstOp = opsInRegion.front();
-      OpBuilder consumedBuilder(mainLoopForOp.getContext());
+      OpBuilder consumedBuilder(loop.getContext());
       // When enable_buffer_insert_optimization is on, place the consumer chain
       // at the start of the dep user's block_id=X region (before the first op
       // with that block_id). Otherwise keep "right before firstOp".
@@ -1552,9 +1505,9 @@ static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp,
       }
       consumedBuilder.setInsertionPoint(consumerAnchor);
 
-      if (int ret = processNormalConsumerBlock(
-              consumedBuilder, depVal, buffers, mainLoopForOp, opsInRegion,
-              userBlockId, groupId, globalBuilder))
+      if (int ret = processNormalConsumerBlock(consumedBuilder, depVal, buffers,
+                                               loop, opsInRegion, userBlockId,
+                                               groupId, globalBuilder))
         return ret;
     }
   }
@@ -1668,7 +1621,7 @@ cloneEmptyFillToConsumers(Value depVal, int producerId,
 //
 // Returns 0 on success, -1 on failure.
 static int
-cloneEmptyFillsInBlocks(scf::ForOp mainLoopForOp,
+cloneEmptyFillsInBlocks(const MainLoop &loop,
                         DenseMap<Value, InnerBlockInfo> &blocks,
                         DenseMap<Value, SmallVector<Value>> &depValueMap,
                         DenseMap<Value, SmallVector<Operation *>> &depUserMap,
@@ -1692,9 +1645,8 @@ cloneEmptyFillsInBlocks(scf::ForOp mainLoopForOp,
       if (!isEmptyFillPattern(depVal))
         continue;
 
-      // Skip if parentOp is not the main_loop forOp (clone logic
-      // currently expects the empty/fill to be inside main_loop).
-      if (defOp->getParentOp() != mainLoopForOp.getOperation())
+      // Skip if parentOp is not the main_loop
+      if (defOp->getParentOp() != loop.getOperation())
         continue;
 
       auto producerId = getOpBlockId(defOp);
@@ -1712,7 +1664,7 @@ cloneEmptyFillsInBlocks(scf::ForOp mainLoopForOp,
 
 // Process cross-block tensor dependencies for double buffering
 static int processTensorDependencies(
-    mlir::scf::ForOp mainLoopForOp, DenseMap<Value, InnerBlockInfo> &blocks,
+    const MainLoop &loop, DenseMap<Value, InnerBlockInfo> &blocks,
     DenseMap<Value, SmallVector<Value>> &depValueMap,
     DenseMap<Value, SmallVector<Operation *>> &depUserMap, BufferMap &bufferMap,
     OpBuilder &globalBuilder, int &groupId) {
@@ -1743,9 +1695,9 @@ static int processTensorDependencies(
       if (isa<tensor::EmptyOp>(depVal.getDefiningOp()))
         continue;
 
-      // Check if definingOp's parentOp is the main_loop forOp
+      // Check if definingOp's parentOp is the main_loop
       auto *parentOp = depVal.getDefiningOp()->getParentOp();
-      if (parentOp != mainLoopForOp.getOperation())
+      if (parentOp != loop.getOperation())
         continue;
 
       // The empty+fill pattern has already been cloned by
@@ -1777,8 +1729,8 @@ static int processTensorDependencies(
         continue;
 
       // Process cross-block dependency with double buffering
-      if (processDepVal(depVal, mainLoopForOp, bufferMap, depUserMap,
-                        globalBuilder, *producerId, groupId) != 0)
+      if (processDepVal(depVal, loop, bufferMap, depUserMap, globalBuilder,
+                        *producerId, groupId) != 0)
         return -1;
       groupId++;
     }
@@ -1786,15 +1738,15 @@ static int processTensorDependencies(
   return 0;
 }
 
-static BufferMap insertBuffersBeforeFor(mlir::scf::ForOp forOp,
-                                        SmallVector<Value> &valueList,
-                                        OpBuilder &builder, int groupId) {
+static BufferMap insertBuffersBeforeLoop(const MainLoop &loop,
+                                         SmallVector<Value> &valueList,
+                                         OpBuilder &builder, int groupId) {
   BufferMap bufferMap;
-  Block *parentBlock = forOp->getBlock();
+  Block *parentBlock = loop.getBlock();
   OpBuilder insertedBuffers(builder.getContext());
-  insertedBuffers.setInsertionPoint(parentBlock, forOp->getIterator());
+  insertedBuffers.setInsertionPoint(parentBlock, loop.getIterator());
 
-  BufferCountManager bufferCountMgr(forOp);
+  BufferCountManager bufferCountMgr(loop.getOperation());
   int bufNum = bufferCountMgr.getBufferCountByType(
       BufferCountManager::DepType::IntraCore);
 
@@ -1810,13 +1762,13 @@ static BufferMap insertBuffersBeforeFor(mlir::scf::ForOp forOp,
           AddressSpaceAttr::get(insertedBuffers.getContext(), addrSpace));
 
       auto allocOp =
-          insertedBuffers.create<memref::AllocOp>(forOp.getLoc(), memrefType);
+          insertedBuffers.create<memref::AllocOp>(loop.getLoc(), memrefType);
 
       auto genericType = MemRefType::get(shapedType.getShape(), elemType,
                                          MemRefLayoutAttrInterface{}, 0u);
 
       auto casted = insertedBuffers.create<memref::MemorySpaceCastOp>(
-          forOp.getLoc(), genericType, allocOp.getResult());
+          loop.getLoc(), genericType, allocOp.getResult());
 
       buffers.push_back({casted.getResult(), casted.getResult()});
     }
@@ -1839,10 +1791,144 @@ hasMemrefDepValue(DenseMap<Value, SmallVector<Value>> &depValueMap) {
   return false;
 }
 
-static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp,
-                               OpBuilder &builder, scope::ScopeOp vectorScope,
-                               int &groupId, bool &i1Found) {
-  OpBuilder globalBuilder(mainLoopForOp.getContext());
+// Build the before-region of the new whileOp
+static void buildBeforeRegion(scf::WhileOp oldWhile, OpBuilder &bb, Location bl,
+                              ValueRange iterArgs) {
+  Block *oldBefore = oldWhile.getBeforeBody();
+  IRMapping mapper;
+  unsigned numOrig = oldBefore->getNumArguments();
+  for (unsigned i = 0; i < numOrig; ++i)
+    mapper.map(oldBefore->getArgument(i), iterArgs[i]);
+
+  Operation *oldCond = nullptr;
+  for (Operation &op : *oldBefore) {
+    if (isa<scf::ConditionOp>(&op)) {
+      oldCond = &op;
+      continue;
+    }
+    bb.clone(op, mapper);
+  }
+  if (!oldCond)
+    return;
+
+  SmallVector<Value> newCondOps;
+  for (Value operand : oldCond->getOperands()) {
+    Value mapped = mapper.lookupOrNull(operand);
+    newCondOps.push_back(mapped ? mapped : operand);
+  }
+  newCondOps.push_back(iterArgs[numOrig]);
+  Value condValue = newCondOps.front();
+  ArrayRef<Value> carriedValues = ArrayRef<Value>(newCondOps).drop_front();
+  bb.create<scf::ConditionOp>(bl, condValue, carriedValues);
+}
+
+// Build the after-region of the new whileOp
+static void buildAfterRegion(scf::WhileOp oldWhile, OpBuilder &ab, Location al,
+                             ValueRange iterArgs, Value &counterIterArgOut) {
+  Block *oldAfter = oldWhile.getAfterBody();
+  unsigned numOrig = oldAfter->getNumArguments();
+  counterIterArgOut = iterArgs[numOrig];
+
+  IRMapping mapper;
+  for (unsigned i = 0; i < numOrig; ++i)
+    mapper.map(oldAfter->getArgument(i), iterArgs[i]);
+
+  Operation *oldYield = nullptr;
+  for (Operation &op : *oldAfter) {
+    if (isa<scf::YieldOp>(&op)) {
+      oldYield = &op;
+      continue;
+    }
+    ab.clone(op, mapper);
+  }
+  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);
+  ab.create<scf::YieldOp>(al, newYieldOps);
+}
+
+// Create a pass-managed global iteration counter for a whileOp main_loop
+static std::pair<Value, scf::WhileOp>
+setupWhileIterArgCounter(const MainLoop &loop, OpBuilder &builder) {
+  auto oldWhile = cast<scf::WhileOp>(loop.getOperation());
+  Location loc = loop.getLoc();
+  MLIRContext *ctx = loop.getContext();
+  Type i32Type = builder.getI32Type();
+
+  // Init 0 for the new counter iter_arg, inserted before oldWhile so the
+  // new whileOp can replace it in-place.
+  OpBuilder preBuilder(ctx);
+  preBuilder.setInsertionPoint(oldWhile);
+  Value zero = preBuilder.create<arith::ConstantIntOp>(loc, 0, 32);
+
+  // Old inits/result-types + i32 counter appended at the end.
+  SmallVector<Value> newInits(oldWhile.getInits().begin(),
+                              oldWhile.getInits().end());
+  newInits.push_back(zero);
+  SmallVector<Type> newResultTypes(oldWhile.getResultTypes().begin(),
+                                   oldWhile.getResultTypes().end());
+  newResultTypes.push_back(i32Type);
+
+  // Captured by the after-builder when the do-region is constructed;
+  // returned to the caller as the live iteration count.
+  Value counterIterArg;
+
+  OpBuilder cb(ctx);
+  cb.setInsertionPoint(oldWhile);
+  auto newWhile = cb.create<scf::WhileOp>(
+      loc, newResultTypes, newInits,
+      [&](OpBuilder &bb, Location bl, ValueRange iterArgs) {
+        buildBeforeRegion(oldWhile, bb, bl, iterArgs);
+      },
+      [&](OpBuilder &ab, Location al, ValueRange iterArgs) {
+        buildAfterRegion(oldWhile, ab, al, iterArgs, counterIterArg);
+      });
+
+  // Move attrs
+  for (auto attr : oldWhile->getAttrs())
+    newWhile->setAttr(attr.getName(), attr.getValue());
+  newWhile->setAttr(kIterCounter, cb.getUnitAttr());
+
+  for (unsigned i = 0, e = oldWhile.getNumResults(); i < e; ++i)
+    oldWhile.getResult(i).replaceAllUsesWith(newWhile.getResult(i));
+  oldWhile.erase();
+
+  return {counterIterArg, newWhile};
+}
+
+static int addInnerMultiBuffer(MainLoop mainLoop, OpBuilder &builder,
+                               scope::ScopeOp vectorScope, int &groupId,
+                               bool &i1Found) {
+  OpBuilder globalBuilder(mainLoop.getContext());
 
   // Two-phase dep collection for empty+fill cloning:
   //   Phase 1 (initial): collect deps, build user map, then clone the
@@ -1865,8 +1951,26 @@ static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp,
   DenseMap<Value, InnerBlockInfo> blocks;
   DenseMap<Value, SmallVector<Value>> depValueMap;
   SmallVector<Operation *> allOps;
-  if (collectInnerBlockInfo(mainLoopForOp, blocks, depValueMap, allOps,
-                            i1Found) != 0)
+
+  // whileOp: bufNum>1 needs pre-created counter (no implicit iter count);
+  // bufNum==1 skips to avoid dead iter_arg.
+  if (mainLoop.isWhile()) {
+    BufferCountManager bufferCountMgr(mainLoop.getOperation());
+    int bufNum = bufferCountMgr.getBufferCountByType(
+        BufferCountManager::DepType::IntraCore);
+    if (bufNum > kBufferCountOne) {
+      auto [counter, newWhile] =
+          setupWhileIterArgCounter(mainLoop, globalBuilder);
+      // Update mainLoop so subsequent collections target the new whileOp
+      // (the old one was erased inside setupWhileIterArgCounter).
+      mainLoop.op = newWhile;
+      mainLoop.body = newWhile.getAfterBody();
+      mainLoop.iterCounter = counter;
+    }
+  }
+
+  if (collectInnerBlockInfo(mainLoop, blocks, depValueMap, allOps, i1Found) !=
+      0)
     return -1;
 
   if (blocks.empty())
@@ -1877,11 +1981,11 @@ static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp,
   // consumer-block users; the cloned fills will rewrite those users' uses.
   DenseMap<Value, SmallVector<Operation *>> initialDepUserMap =
       buildDepUserMap(blocks, allOps, depValueMap);
-  if (cloneEmptyFillsInBlocks(mainLoopForOp, blocks, depValueMap,
-                              initialDepUserMap, globalBuilder) != 0)
+  if (cloneEmptyFillsInBlocks(mainLoop, blocks, depValueMap, initialDepUserMap,
+                              globalBuilder) != 0)
     return -1;
 
-  rematerializeTensorRootedScalarDeps(mainLoopForOp);
+  rematerializeTensorRootedScalarDeps(mainLoop);
 
   // Phase 2: re-collect deps now that cloned ops (and rematerialized scalar
   // chains) have created new cross-block references. depValueMap and allOps
@@ -1890,8 +1994,8 @@ static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp,
   blocks.clear();
   depValueMap.clear();
   allOps.clear();
-  if (collectInnerBlockInfo(mainLoopForOp, blocks, depValueMap, allOps,
-                            i1Found) != 0)
+  if (collectInnerBlockInfo(mainLoop, blocks, depValueMap, allOps, i1Found) !=
+      0)
     return -1;
 
   // Phase 2 may surface i1 tensor deps that the clone introduced (e.g. a
@@ -1907,25 +2011,38 @@ static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp,
   // Memref-type dep values are not supported here; fail loudly so downstream
   // passes don't see an unmarked-but-skipped scope.
   if (hasMemrefDepValue(depValueMap)) {
-    LDBG("ERROR: Memref type dependent values found!");
+    LDBG("Falling Back: Memref type dependent values found!");
     return -1;
   }
 
   auto depUserMap = buildDepUserMap(blocks, allOps, depValueMap);
 
+  LLVM_DEBUG(
+      llvm::dbgs() << "[addInnerMultiBuffer] before collectBufferValues\n");
   auto valueList = collectBufferValues(depValueMap);
+  LLVM_DEBUG(
+      llvm::dbgs()
+      << "[addInnerMultiBuffer] before insertBuffersBeforeLoop, valueList.size="
+      << valueList.size() << "\n");
+
   auto bufferMap =
-      insertBuffersBeforeFor(mainLoopForOp, valueList, builder, groupId);
+      insertBuffersBeforeLoop(mainLoop, valueList, builder, groupId);
 
+  LLVM_DEBUG(
+      llvm::dbgs() << "[addInnerMultiBuffer] before collectScalarDeps\n");
   auto scalarValueList = collectScalarDeps(depValueMap, depUserMap);
 
+  LLVM_DEBUG(llvm::dbgs() << "[addInnerMultiBuffer] before markScalarDeps\n");
   markScalarDeps(scalarValueList, depUserMap, globalBuilder, 1);
 
-  if (processTensorDependencies(mainLoopForOp, blocks, depValueMap, depUserMap,
+  LLVM_DEBUG(llvm::dbgs()
+             << "[addInnerMultiBuffer] before processTensorDependencies\n");
+  if (processTensorDependencies(mainLoop, blocks, depValueMap, depUserMap,
                                 bufferMap, globalBuilder, groupId) != 0) {
     return -1;
   }
 
+  LLVM_DEBUG(llvm::dbgs() << "[addInnerMultiBuffer] DONE\n");
   return 0;
 }
 
@@ -1963,10 +2080,10 @@ void AddMultiBufferInnerScopePass::runOnOperation() {
       return WalkResult::advance();
     }
 
-    // Step 3: Collect all forOps with main_loop attribute
-    SmallVector<scf::ForOp> mainLoopForOps;
+    // Step 3: Collect all main_loop loops (forOp / whileOp) in the scope
+    SmallVector<Operation *> mainLoops;
     int foundCount =
-        collectMainLoopsRecursively(scope.getBodyRegion(), mainLoopForOps);
+        collectMainLoopsRecursively(scope.getBodyRegion(), mainLoops);
     if (foundCount < 0) {
       LDBG("collectMainLoopsRecursively failed");
       return WalkResult::interrupt();
@@ -1974,24 +2091,19 @@ void AddMultiBufferInnerScopePass::runOnOperation() {
     if (foundCount == 0)
       return WalkResult::advance();
 
-    // Step 4: Process each main_loop forOp
+    // Step 4: Process each main_loop
     int groupId = 0;
-    for (scf::ForOp mainLoopForOp : mainLoopForOps) {
-      scf::ForOp nestedMainloop = findNestedMainloopInForOp(mainLoopForOp);
-      if (nestedMainloop) {
+    for (Operation *loopOp : mainLoops) {
+      MainLoop mainLoop(loopOp);
+      if (findNestedMainloop(mainLoop)) {
         LDBG("Nested main_loop found, this is not allowed");
         return WalkResult::interrupt();
       }
       // i1Found is reset per main_loop so it only triggers fallback for
       // the current scope's deps.
       bool i1Found = false;
-      int ret =
-          addInnerMultiBuffer(mainLoopForOp, builder, scope, groupId, i1Found);
+      int ret = addInnerMultiBuffer(mainLoop, builder, scope, groupId, i1Found);
       if (i1Found) {
-        // i1 tensor deps are not safe to multi-buffer; mark the module
-        // with ERRCODE_IGNORED and bail out so downstream passes see the
-        // fallback attribute. Mirrors the AnalyzeName pass pattern:
-        // setFallbackAttr(module) + signalPassFailure() + return.
         LDBG("i1 tensor dep found, setting fallback attribute");
         CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_IGNORED);
         return WalkResult::interrupt();
diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
index 37959eeb2f..6d39f18da6 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
@@ -293,5 +293,57 @@ int64_t getBTSizeFromValidBroadcastOp(linalg::BroadcastOp broadcastOp) {
   return sizeBytes;
 }
 
+Block *MainLoop::getBody() const { return body; }
+
+Operation *MainLoop::getOperation() const { return op; }
+
+MLIRContext *MainLoop::getContext() const { return op->getContext(); }
+
+Location MainLoop::getLoc() const { return op->getLoc(); }
+
+Block *MainLoop::getBlock() const { return op->getBlock(); }
+
+Block::iterator MainLoop::getIterator() const { return op->getIterator(); }
+
+Operation *MainLoop::operator->() const { return op; }
+
+bool MainLoop::isWhile() const { return isa<scf::WhileOp>(op); }
+
+SmallVector<Value> MainLoop::getIterArgs() const {
+  SmallVector<Value> result;
+  if (auto f = dyn_cast<scf::ForOp>(op)) {
+    result.append(f.getRegionIterArgs().begin(), f.getRegionIterArgs().end());
+  } else if (auto w = dyn_cast<scf::WhileOp>(op)) {
+    Block::BlockArgListType args = w.getAfterBody()->getArguments();
+    result.append(args.begin(), args.end());
+  }
+  return result;
+}
+
+SmallVector<Value> MainLoop::getBeforeIterArgs() const {
+  SmallVector<Value> result;
+  if (auto w = dyn_cast<scf::WhileOp>(op)) {
+    Block::BlockArgListType args = w.getBeforeBody()->getArguments();
+    result.append(args.begin(), args.end());
+  }
+  return result;
+}
+
+MainLoop::MainLoop(Operation *loopOp) {
+  op = loopOp;
+  if (auto f = dyn_cast<scf::ForOp>(loopOp))
+    body = f.getBody();
+  else if (auto w = dyn_cast<scf::WhileOp>(loopOp))
+    body = w.getAfterBody();
+}
+
+scf::YieldOp MainLoop::getLoopYieldOp(Operation *loopOp) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loopOp))
+    return dyn_cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp))
+    return dyn_cast<scf::YieldOp>(whileOp.getAfter().front().getTerminator());
+  return {};
+}
+
 } // namespace CVPipeline
 } // namespace mlir
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-memref-dep.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-memref-dep.mlir
index f2b3ca34a8..4382cf5e8d 100644
--- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-memref-dep.mlir
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-memref-dep.mlir
@@ -1,12 +1,16 @@
-// RUN: (triton-opt --add_multi_buffer_inner_scope %s 2>&1 || echo "PASS") | FileCheck %s
-// CHECK: PASS
+// RUN: triton-opt --add_multi_buffer_inner_scope %s 2>&1 | FileCheck %s
+// Pass signals fallback via triton_ascend.dynamic_cv_pipeline.rc = 1
+// (ERRCODE_FAILED); the IR is otherwise unchanged.
+
+// CHECK-LABEL: module attributes
+// CHECK-SAME: triton_ascend.dynamic_cv_pipeline.rc = 1
 
 // T26: Memref Type Dependency Triggers Fallback
 // Test: When memref.alloc with block_id=X produces a memref, and later
 //       bufferization.to_tensor with block_id=Y uses that memref (X != Y),
 //       the memref is a cross-block dependency and hasMemrefDepValue returns
 //       true, causing pass failure.
-// Key Check: Pass should fail (exit code != 0)
+// Key Check: Pass falls back via the ERRCODE_FAILED attribute.
 
 module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
   func.func @test_t26_memref_dep_fallback() {
@@ -24,7 +28,7 @@ module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
         %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32} : memref<128xf32>
         // bufferization.to_tensor in block_id = 10 (consumer block, different from 9)
         // This uses %alloc directly, creating a cross-block memref dependency
-        %tensor_from_alloc = bufferization.to_tensor %alloc restrict writable {ssbuffer.block_id = 10 : i32} : memref<128xf32>
+        %tensor_from_alloc = bufferization.to_tensor %alloc {ssbuffer.block_id = 10 : i32} : memref<128xf32> to tensor<128xf32>
         // Use the tensor in block_id = 10
         %consumed = arith.addf %tensor_from_alloc, %tensor_from_alloc {ssbuffer.block_id = 10 : i32} : tensor<128xf32>
         // Producer continuation in block_id = 5
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir
new file mode 100644
index 0000000000..bbd43eb386
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner-scope-whileop.mlir
@@ -0,0 +1,280 @@
+// RUN: triton-opt --add_multi_buffer_inner_scope %s | FileCheck %s
+
+// T-while-A: whileOp as main_loop, INTRA bufNum == 1 (single-buffer scope).
+//   Verifies:
+//     - whileOp carrying ssbuffer.main_loop on its terminator IS recognized
+//       (legacy shape supported by hasMainLoopAttr).
+//     - setupWhileIterArgCounter is SKIPPED (bufNum==1 → no dead iter_arg).
+//     - Cross-block tensor dep still gets single-buffer treatment
+//       (single memref.alloc + single hivm.hir.copy + single to_tensor).
+//     - NO scf.if dispatch is emitted (N==1 fast path).
+//     - The original whileOp is preserved verbatim (no extra i32 iter_arg).
+//   We pin the buffer count to 1 via the module-level
+//   `ssbuffer.intra_buf_count` attribute so the default of 2 doesn't
+//   trigger the counter-setup branch.
+
+// CHECK-LABEL: func.func @test_while_mainloop_bufnum_one
+// Exactly one UB alloc before the whileOp.
+// CHECK-DAG:   memref.alloc() : memref<128xf32, #hivm.address_space<ub>>
+// CHECK-NOT:   memref.alloc() {{.*}}: memref<128xf32, #hivm.address_space<ub>>
+// Original whileOp do-region bb0 has only 2 block-args (no i32 counter).
+// CHECK:       ^bb0(%{{.*}}: tensor<128xf32>, %{{.*}}: i32):
+// Single producer-side hivm.hir.copy.
+// CHECK:       hivm.hir.copy ins({{.*}} : tensor<128xf32>) outs({{.*}} : memref<128xf32>)
+// Single consumer-side bufferization.to_tensor (the readback).
+// CHECK:       bufferization.to_tensor {{.*}}: memref<128xf32> to tensor<128xf32>
+// main_loop attribute survives on the new whileOp.
+// CHECK:       } {{.*}}ssbuffer.main_loop = 1 : i64
+
+// T-while-B: whileOp as main_loop, INTRA bufNum == 2 (multi-buffer scope).
+//   Verifies:
+//     - setupWhileIterArgCounter IS called (bufNum>1).
+//     - whileOp is REPLACED with a new one that has an extra i32 iter_arg
+//       (init=0, yielded as counter+1 at end of do-region).
+//     - The new do-region bb0 has ONE extra block-arg compared to the input
+//       (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.
+
+// 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).
+// CHECK:       scf.if
+// CHECK:         hivm.hir.copy
+// 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
+
+// T-while-C: whileOp main_loop with scf.if inside the do-region (multi-region
+// consumer pattern). The whileOp carries `ssbuffer.main_loop` on itself (NOT
+// on the terminator). This exercises `hasMainLoopAttr`'s "op has attr" path,
+// in contrast to T-while-A/B which use the terminator-attr legacy shape.
+// Verifies the cross-block tensor dep flowing through scf.if is buffered.
+
+// CHECK-LABEL: func.func @test_while_mainloop_attr_on_op
+// 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>>
+// 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
+
+// T-while-D: Regression guard for the getOutermostSsbufferId priority fix.
+//   Before the fix, walking up from an op inside a whileOp's do-region would
+//   hit the multi-region check (`numRegions >= 2`) BEFORE the kMainLoop check
+//   and return the whileOp's own block_id, breaking cross-block dep
+//   classification.
+//   Setup: whileOp (block_id = 5) wrapping a producer block_id = 8 inside the
+//   do-region and a consumer block_id = 12 also inside the do-region.
+//   Cross-block judgment must use 8 vs 12 (NOT 5 vs 12).
+//   Result: a multi-buffer MUST be emitted for the producer (8 → 12 is cross-
+//   block). If the bug were present, getOutermostSsbufferId would return 5
+//   for both, classify them as same-block, and SKIP the multi-buffer.
+
+// CHECK-LABEL: func.func @test_while_outermost_id_priority
+// Multi-buffer must be emitted (proves cross-block judgment saw 8 != 12).
+// CHECK-DAG:   memref.alloc() : memref<32xf32, #hivm.address_space<ub>>
+// CHECK-DAG:   memref.alloc() : memref<32xf32, #hivm.address_space<ub>>
+// 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
+// Counter-aware whileOp carries ssbuffer.iterCounter alongside main_loop.
+// CHECK:       } {{.*}}ssbuffer.iterCounter, {{.*}}ssbuffer.main_loop = 1 : i64
+
+// T-while-E: whileOp main_loop with a tensor::EmptyOp + linalg::FillOp pattern
+//   inside the do-region. The empty+fill is cloned into each consumer block.
+//   Verifies cloneEmptyFillsInBlocks works on a whileOp main_loop (the body
+//   is `getAfterBody()`, not `getBody()`).
+//   The cloned fill's ins stays the same (arith.constant, lives outside the
+//   main_loop) — no clone of the scalar chain is needed.
+
+// CHECK-LABEL: func.func @test_while_clone_empty_fill
+// Original fill at block_id = 8 (producer block) is preserved.
+// CHECK:       linalg.fill {{.*}}{ssbuffer.block_id = 8 : i32} {{.*}}outs({{.*}} : tensor<32x1xf32>)
+// Original fill at block_id = 14 (second producer block) is preserved.
+// CHECK:       linalg.fill {{.*}}{ssbuffer.block_id = 14 : i32} {{.*}}outs({{.*}} : tensor<32x1xf32>)
+// Cloned fills land in block_id = 12 (consumer block) — there are two of them,
+// one per original producer. CHECK matches at least once.
+// CHECK:       linalg.fill {{.*}}{ssbuffer.block_id = 12 : i32} {{.*}}outs({{.*}} : tensor<32x1xf32>)
+// 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
+
+// T-while-F (TODO): Negative — two whileOps in the same scope, neither is the
+//   "outer" main_loop; one is nested in the other's do-region. The pass must
+//   refuse via findNestedMainloop (returns WalkResult::interrupt + fallback).
+//   Not implemented in this initial drop — to be added when the negative case
+//   is needed for a regression guard.
+
+// ---- Inputs ----
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">,
+                   ssbuffer.intra_buf_count = 1 : i32} {
+  // T-while-A: whileOp with main_loop on the terminator, single-buffer.
+  func.func @test_while_mainloop_bufnum_one() {
+    %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<128xf32>
+    %carry = linalg.fill ins(%cst_zero : f32) outs(%init : tensor<128xf32>) -> tensor<128xf32>
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32)
+          : (tensor<128xf32>, i32) -> (tensor<128xf32>, i32) {
+        %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : tensor<128xf32>, i32
+      } do {
+      ^bb0(%arg0: tensor<128xf32>, %arg1: i32):
+        // Producer block_id = 7.
+        %alloc = memref.alloc() {ssbuffer.block_id = 7 : i32} : memref<128xf32>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32>
+        // Consumer block_id = 10 (cross-block).
+        %consumed = arith.addf %prod, %prod {ssbuffer.block_id = 10 : i32} : tensor<128xf32>
+        %next = arith.addi %arg1, %c1_i32 : i32
+        scf.yield %consumed, %next : tensor<128xf32>, i32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 23 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">,
+                   ssbuffer.intra_buf_count = 2 : i32} {
+  // T-while-B: whileOp with main_loop on the terminator, multi-buffer (bufNum=2).
+  func.func @test_while_mainloop_bufnum_two() {
+    %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<128xf32>
+    %carry = linalg.fill ins(%cst_zero : f32) outs(%init : tensor<128xf32>) -> tensor<128xf32>
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32)
+          : (tensor<128xf32>, i32) -> (tensor<128xf32>, i32) {
+        %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : tensor<128xf32>, i32
+      } do {
+      ^bb0(%arg0: tensor<128xf32>, %arg1: i32):
+        %alloc = memref.alloc() {ssbuffer.block_id = 7 : i32} : memref<128xf32>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32>
+        %consumed = arith.addf %prod, %prod {ssbuffer.block_id = 10 : i32} : tensor<128xf32>
+        %next = arith.addi %arg1, %c1_i32 : i32
+        scf.yield %consumed, %next : tensor<128xf32>, i32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 23 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+
+  // T-while-C: whileOp with main_loop on the OP itself (not terminator).
+  func.func @test_while_mainloop_attr_on_op() {
+    %c0_i32 = arith.constant 0 : i32
+    %c10_i32 = arith.constant 10 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %cst_one = arith.constant 1.0 : f16
+    %init = tensor.empty() : tensor<64xf16>
+    %carry = linalg.fill ins(%cst_one : f16) outs(%init : tensor<64xf16>) -> tensor<64xf16>
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32)
+          : (tensor<64xf16>, i32) -> (tensor<64xf16>, i32) {
+        %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : tensor<64xf16>, i32
+      } do {
+      ^bb0(%arg0: tensor<64xf16>, %arg1: i32):
+        %alloc = memref.alloc() {ssbuffer.block_id = 7 : i32} : memref<64xf16>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 7 : i32} : memref<64xf16> to tensor<64xf16>
+        // Multi-region consumer (scf.if inside the do-region).
+        %cnd = arith.cmpi slt, %arg1, %c1_i32 {ssbuffer.block_id = 11 : i32} : i32
+        %consumed = scf.if %cnd -> (tensor<64xf16>) {
+          %a = arith.addf %prod, %prod {ssbuffer.block_id = 11 : i32} : tensor<64xf16>
+          scf.yield %a : tensor<64xf16>
+        } else {
+          %b = arith.mulf %prod, %prod {ssbuffer.block_id = 12 : i32} : tensor<64xf16>
+          scf.yield %b : tensor<64xf16>
+        } {ssbuffer.block_id = 12 : i32}
+        %next = arith.addi %arg1, %c1_i32 : i32
+        scf.yield %consumed, %next : tensor<64xf16>, i32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 23 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+
+  // T-while-D: whileOp main_loop with cross-block dep that tests
+  // getOutermostSsbufferId's main_loop-first priority.
+  // Producer at block_id=8, consumer at block_id=12, both inside do-region.
+  // whileOp itself is at block_id=5 (so the buggy code would have collapsed
+  // both to 5, hiding the cross-block dep).
+  func.func @test_while_outermost_id_priority() {
+    %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<32xf32>
+    %carry = linalg.fill ins(%cst_zero : f32) outs(%init : tensor<32xf32>) -> tensor<32xf32>
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %carry, %arg1 = %c0_i32)
+          : (tensor<32xf32>, i32) -> (tensor<32xf32>, i32) {
+        %cmp = arith.cmpi slt, %arg1, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : tensor<32xf32>, i32
+      } do {
+      ^bb0(%arg0: tensor<32xf32>, %arg1: i32):
+        %alloc = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<32xf32>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 8 : i32} : memref<32xf32> to tensor<32xf32>
+        %consumed = arith.addf %prod, %prod {ssbuffer.block_id = 12 : i32} : tensor<32xf32>
+        %next = arith.addi %arg1, %c1_i32 : i32
+        scf.yield %consumed, %next : tensor<32xf32>, i32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 5 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+
+  // T-while-E: whileOp main_loop with empty+fill pattern.
+  // Two pairs of empty+fill, both cloned to consumer block 12.
+  func.func @test_while_clone_empty_fill() {
+    %c0_i32 = arith.constant 0 : i32
+    %c10_i32 = arith.constant 10 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %cst_one = arith.constant 1.0 : f32
+    scope.scope : () -> () {
+      %result:2 = scf.while (%arg0 = %c0_i32, %arg1 = %cst_one)
+          : (i32, f32) -> (i32, f32) {
+        %cmp = arith.cmpi slt, %arg0, %c10_i32 {ssbuffer.block_id = 16 : i32} : i32
+        scf.condition(%cmp) %arg0, %arg1 : i32, f32
+      } do {
+      ^bb0(%arg0: i32, %arg1: f32):
+        // Producer block_id = 8: empty+fill (cloneable).
+        %empty8 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<32x1xf32>
+        %fill8 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_one : f32)
+            outs(%empty8 : tensor<32x1xf32>) -> tensor<32x1xf32>
+        // Producer block_id = 14: empty+fill (cloneable).
+        %empty14 = tensor.empty() {ssbuffer.block_id = 14 : i32} : tensor<32x1xf32>
+        %fill14 = linalg.fill {ssbuffer.block_id = 14 : i32} ins(%cst_one : f32)
+            outs(%empty14 : tensor<32x1xf32>) -> tensor<32x1xf32>
+        // Consumer block_id = 12: uses BOTH fills cross-block.
+        %consumed = arith.addf %fill8, %fill14 {ssbuffer.block_id = 12 : i32} : tensor<32x1xf32>
+        %next = arith.addi %arg0, %c1_i32 : i32
+        scf.yield %next, %arg1 : i32, f32
+      } attributes {ssbuffer.main_loop = 1 : i64, ssbuffer.block_id = 23 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner_scope_i1_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner_scope_i1_test.mlir
index fdbaf0e843..78e511c8b7 100644
--- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner_scope_i1_test.mlir
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Inner_scope_i1_test.mlir
@@ -1,5 +1,11 @@
-// RUN: (triton-opt --add_multi_buffer_inner_scope %s 2>&1 || echo "PASS") | FileCheck %s
-// CHECK: PASS
+// RUN: triton-opt --add_multi_buffer_inner_scope %s 2>&1 | FileCheck %s
+// Pass sets triton_ascend.dynamic_cv_pipeline.rc = 1 (ERRCODE_FAILED) to
+// signal fallback; the outer runOnOperation wrapper at
+// AddMultiBufferInnerScope.cpp:2180-2183 overwrites the more specific
+// ERRCODE_IGNORED=2 that addInnerMultiBuffer sets for i1 deps.
+
+// CHECK-LABEL: module attributes
+// CHECK-SAME: triton_ascend.dynamic_cv_pipeline.rc = 1
 
 // T-i1: i1 Tensor Dependency Triggers Fallback
 // Test: When a tensor dep with element type i1 is produced in one block and
@@ -28,7 +34,7 @@ module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
       scf.for %i = %c0_i32 to %c100_i32 step %c1_i32  : i32 {
         // Producer: tensor<128xi1> in block 7 (NOT empty+fill pattern)
         %alloc = memref.alloc() {ssbuffer.block_id = 7 : i32} : memref<128xi1>
-        %prod = bufferization.to_tensor %alloc restrict writable {ssbuffer.block_id = 7 : i32} : memref<128xi1>
+        %prod = bufferization.to_tensor %alloc {ssbuffer.block_id = 7 : i32} : memref<128xi1> to tensor<128xi1>
         // Consumer in block 10 (cross-block)
         %consumed = arith.ori %prod, %prod {ssbuffer.block_id = 10 : i32} : tensor<128xi1>
       } {ssbuffer.main_loop = 1 : i64}

From 12ad08d7f630b440b27ce392a11e93133db73476 Mon Sep 17 00:00:00 2001
From: Four1er <fqf@users.noreply.github.com>
Date: Mon, 3 Aug 2026 15:28:07 +0800
Subject: [PATCH 02/11] [ssbuffer](feat) add scf.while op support in
 OpClassifier

---
 .../include/DynamicCVPipeline/Common/Utils.h  |  56 +++++
 .../PlanComputeBlock/OpClassifier.h           |   3 +-
 .../PlanComputeBlock/OpClassifier.cpp         | 211 ++++++++++--------
 3 files changed, 175 insertions(+), 95 deletions(-)

diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
index 20f0a036c8..7ecff19b6d 100644
--- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
+++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
@@ -165,6 +165,62 @@ inline bool isCubeOp(Operation *op) {
   return !isScfOp(op) && CVPipeline::getOpCoreType(op) == CoreType::CUBE_ONLY;
 }
 
+// ============================================================================
+// Unified Loop Helpers: abstract ForOp/WhileOp differences
+// ============================================================================
+// Get the body block of a loop (ForOp's body or WhileOp's after-body block)
+inline Block *getLoopBodyBlock(Operation *loop) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loop))
+    return forOp.getBody();
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loop))
+    return whileOp.getAfterBody()->getNextNode();
+  return nullptr;
+}
+
+// Get the init values of a loop (ForOp's initArgs or WhileOp's inits)
+inline ValueRange getLoopInitValues(Operation *loop) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loop))
+    return forOp.getInitArgs();
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loop))
+    return whileOp.getInits();
+  return {};
+}
+
+// Get the yield terminator of a loop's body
+inline Operation *getLoopYieldOp(Operation *loop) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loop))
+    return forOp.getBody()->getTerminator();
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loop))
+    return whileOp.getAfterBody()->getTerminator();
+  return nullptr;
+}
+
+// Check if a block argument is an iter_arg of a loop (ForOp body or WhileOp
+// after-body)
+inline bool isLoopIterArg(BlockArgument blockArg) {
+  Operation *parentOp = blockArg.getOwner()->getParentOp();
+  if (isa<scf::ForOp>(parentOp))
+    return true;
+  if (auto whileOp = dyn_cast<scf::WhileOp>(parentOp))
+    return blockArg.getOwner() == whileOp.getAfterBody()->getNextNode();
+  return false;
+}
+
+// Helper: Check if a value is a scalar (not a tensor type)
+inline bool isScalarType(Value value) {
+  return !isa<RankedTensorType>(value.getType());
+}
+
+// Helper: Check if a value is a scalar iter_arg from scf.for or scf.while
+inline bool isScalarIterArgOp(Value iterArg) {
+  auto blockArg = dyn_cast<BlockArgument>(iterArg);
+  if (!blockArg)
+    return false;
+  if (!isLoopIterArg(blockArg))
+    return false;
+  return isScalarType(iterArg);
+}
+
 bool isVectorOnlyOp(Operation *op);
 
 bool isScalarLike(Value value);
diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h
index ece019cb7f..8f95e956ca 100644
--- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h
+++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h
@@ -134,8 +134,9 @@ class OpClassifierPass
   // results are consumed exclusively by CUBE ops.
   int penetrateCubeIntoForLoops();
 
-  // Helper: decide whether an scf.for is a pure cube-loader loop
+  // Helper: decide whether an scf.for or scf.while is a pure cube-loader loop
   bool isCubeLoaderForOp(scf::ForOp forOp);
+  bool isCubeLoaderForWhileOp(scf::WhileOp whileOp);
 
   // Initialize the pass
   void initializePass(ModuleOp module);
diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
index f2043fb625..5b52de064b 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
@@ -44,6 +44,7 @@
 #include "bishengir/Dialect/Utils/Util.h"
 
 using namespace mlir;
+using namespace mlir::CVPipeline;
 static constexpr const char *DEBUG_TYPE = "op-classifier";
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LOG_DEBUG(...)                                                         \
@@ -590,30 +591,9 @@ int OpClassifierPass::patternMatchCUBE() {
   return 0;
 }
 
-// Helper: Check if a value is a scalar (not a tensor type)
-static bool isScalarType(Value value) {
-  return !isa<RankedTensorType>(value.getType());
-}
-
-// Helper: Check if a value is a scalar iter_arg from scf.for
-// An iter_arg is a BlockArgument of scf.for's loop body, and it must be scalar
-// type
-static bool isScalarIterArgOp(Value iterArg) {
-  // iter_arg is a BlockArgument of scf.for's body
-  auto blockArg = dyn_cast<BlockArgument>(iterArg);
-  if (!blockArg)
-    return false;
-  // Check if parent is scf.for
-  Operation *parentOp = blockArg.getOwner()->getParentOp();
-  if (!isa<scf::ForOp>(parentOp))
-    return false;
-  // Check if the iter_arg is scalar type (not tensor)
-  return isScalarType(iterArg);
-}
-
 // Helper: Find iter_arg initialization op and yield-assigning op for scf.for
-// loop-carried scalar When a def comes from an scf.for iter_arg and is a scalar
-// compute op, we need to:
+// or scf.while loop-carried scalar. When a def comes from an scf.for/scf.while
+// iter_arg and is a scalar compute op, we need to:
 // 1. Find the iter_arg's initialization op (the op that provides the initial
 // value)
 // 2. Find the yieldOp, then trace to the op that provides the yielded value
@@ -626,42 +606,42 @@ findIterArgUpstreamOps(Value def,
   if (!blockArg)
     return;
 
-  // Check if parent is scf.for
   Operation *parentOp = blockArg.getOwner()->getParentOp();
-  auto forOp = dyn_cast<scf::ForOp>(parentOp);
-  if (!forOp)
+  if (!isa<scf::ForOp, scf::WhileOp>(parentOp))
     return;
 
   // Get the iter_arg index from block argument
   unsigned argIdx = blockArg.getArgNumber();
-
-  // Get the iter_arg and check its type - must be scalar (not tensor)
-  // The init value is at forOp.getInitArgs()[argIdx]
-  if (argIdx > forOp.getInitArgs().size() || argIdx == 0)
+  auto inits = getLoopInitValues(parentOp);
+  // For scf.for: argIdx 0 is lb (loop lower bound), not an iter_arg, skip it.
+  // For scf.while: argIdx 0 can be a valid iter_arg, don't skip.
+  if (argIdx >= inits.size() || (isa<scf::ForOp>(parentOp) && argIdx == 0))
     return;
-  Value initValue = forOp.getInitArgs()[argIdx - 1];
+
+  Value initValue = inits[argIdx - (isa<scf::ForOp>(parentOp) ? 1 : 0)];
   if (!isScalarType(initValue))
     return;
 
-  // Find the initialization op for this iter_arg
   Operation *initDef = initValue.getDefiningOp();
-  if (initDef && initDef != forOp) {
+  if (initDef && initDef != parentOp) {
     LLVM_DEBUG(DBGS() << "[findIterArgUpstreamOps] init def: " << *initDef
                       << "\n");
     upstreamOps.push_back(initDef);
   }
 
-  // Find the yieldOp and the op that provides the yielded value
-  // The yieldOp has operands corresponding to the iteration results
-  // For iter_arg i, yieldOp.getOperand(i) is the value yielded for that
-  // iter_arg
-  Operation *yieldOp = forOp.getBody()->getTerminator();
-  if (!isa<scf::YieldOp>(yieldOp) || argIdx > yieldOp->getNumOperands())
+  // For scf.while, iter args in after region are 0-indexed in both inits and
+  // yield For scf.for, iter args are 1-indexed in initArgs (arg 0 is lb) and
+  // 1-indexed in yield (arg 0 is lb)
+  unsigned yieldOperandIdx = (isa<scf::ForOp>(parentOp)) ? argIdx - 1 : argIdx;
+
+  Operation *yieldOp = getLoopYieldOp(parentOp);
+  if (!yieldOp || !isa<scf::YieldOp>(yieldOp) ||
+      yieldOperandIdx >= yieldOp->getNumOperands())
     return;
 
-  Value yieldedValue = yieldOp->getOperand(argIdx - 1);
+  Value yieldedValue = yieldOp->getOperand(yieldOperandIdx);
   Operation *yieldedDef = yieldedValue.getDefiningOp();
-  if (yieldedDef && yieldedDef != forOp) {
+  if (yieldedDef && yieldedDef != parentOp) {
     LLVM_DEBUG(DBGS() << "[findIterArgUpstreamOps] yielded def: " << *yieldedDef
                       << "\n");
     upstreamOps.push_back(yieldedDef);
@@ -838,13 +818,14 @@ void OpClassifierPass::markFillOpsAsCube() {
       outsIsCube = true;
       LLVM_DEBUG(DBGS() << "\tfill outs defined by CUBE op: "
                         << outsDef->getName().getStringRef() << "\n");
-    } else if (!outsDef) { // Case 2: outs is a BlockArgument (scf.for/scf.if
-                           // iter_arg)
+    } else if (!outsDef) { // Case 2: outs is a BlockArgument
+                           // (scf.for/scf.if/scf.while iter_arg)
       auto blockArg = dyn_cast<BlockArgument>(outs);
       if (blockArg) {
         Operation *parentOp = blockArg.getOwner()->getParentOp();
-        // Check if it's an scf.for or scf.if iter_arg that is CUBE
-        if ((isa<scf::ForOp>(parentOp) || isa<scf::IfOp>(parentOp)) &&
+        // Check if it's an scf.for, scf.if, or scf.while iter_arg that is CUBE
+        if ((isa<scf::ForOp>(parentOp) || isa<scf::IfOp>(parentOp) ||
+             isa<scf::WhileOp>(parentOp)) &&
             opCoreTypes[parentOp] == OP_CUBE_ONLY) {
           outsIsCube = true;
           LLVM_DEBUG(DBGS() << "\tfill outs is CUBE iter_arg of: "
@@ -1025,28 +1006,42 @@ bool isDisqualifyingLoaderOp(Operation *op) {
 } // namespace
 
 bool OpClassifierPass::isCubeLoaderForOp(scf::ForOp forOp) {
-  // Every result must be live and used only by CUBE consumers. A single
-  // non-cube (or scf.yield) user disqualifies the loop.
-  bool hasCubeConsumer = false;
+  // Every result must be live and used only by CUBE consumers
   for (Value result : forOp.getResults()) {
     for (Operation *user : result.getUsers()) {
-      if (isa<linalg::MatmulOp>(user) || getCoreType(user) == OP_CUBE_ONLY) {
-        hasCubeConsumer = true;
+      if (isa<linalg::MatmulOp>(user) || getCoreType(user) == OP_CUBE_ONLY)
         continue;
-      }
       return false;
     }
   }
-  if (!hasCubeConsumer) {
-    return false;
-  }
-
-  // The body (including nested regions) must be pure data movement.
+  // Body must be pure data movement
   bool disqualified = false;
   forOp.getBody()->walk([&](Operation *op) {
-    if (op == forOp.getOperation()) {
+    if (op == forOp.getOperation())
       return WalkResult::advance();
+    if (isDisqualifyingLoaderOp(op)) {
+      disqualified = true;
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  return !disqualified;
+}
+
+bool OpClassifierPass::isCubeLoaderForWhileOp(scf::WhileOp whileOp) {
+  // Every result must be live and used only by CUBE consumers
+  for (Value result : whileOp.getResults()) {
+    for (Operation *user : result.getUsers()) {
+      if (isa<linalg::MatmulOp>(user) || getCoreType(user) == OP_CUBE_ONLY)
+        continue;
+      return false;
     }
+  }
+  // After body must be pure data movement
+  bool disqualified = false;
+  whileOp.getAfterBody()->walk([&](Operation *op) {
+    if (op == whileOp.getOperation())
+      return WalkResult::advance();
     if (isDisqualifyingLoaderOp(op)) {
       disqualified = true;
       return WalkResult::interrupt();
@@ -1056,28 +1051,37 @@ bool OpClassifierPass::isCubeLoaderForOp(scf::ForOp forOp) {
   return !disqualified;
 }
 
+// Helper: mark all ops in a loop body as CUBE_ONLY
+static void
+markLoopBodyAsCube(Operation *loop, Block *body,
+                   llvm::DenseMap<Operation *, OpCoreType> &opCoreTypes) {
+  body->walk([&](Operation *op) {
+    if (op == loop || isa<scf::SCFDialect>(op->getDialect()))
+      return;
+    auto it = opCoreTypes.find(op);
+    if (it != opCoreTypes.end())
+      it->second = OP_CUBE_ONLY;
+  });
+  opCoreTypes[loop] = OP_CUBE_ONLY;
+}
+
 int OpClassifierPass::penetrateCubeIntoForLoops() {
   // Collect first so recoloring earlier loops cannot perturb the scan.
-  llvm::SmallVector<scf::ForOp> loaderLoops;
+  // Use a combined walk to collect both ForOps and WhileOps
+  llvm::SmallVector<Operation *> loaderLoops;
   getOperation().walk([&](scf::ForOp forOp) {
-    if (isCubeLoaderForOp(forOp)) {
+    if (isCubeLoaderForOp(forOp))
       loaderLoops.push_back(forOp);
-    }
+  });
+  getOperation().walk([&](scf::WhileOp whileOp) {
+    if (isCubeLoaderForWhileOp(whileOp))
+      loaderLoops.push_back(whileOp);
   });
 
-  for (scf::ForOp forOp : loaderLoops) {
-    // Set core_type to OP_CUBE_ONLY for scf.for.
-    forOp.getBody()->walk([&](Operation *op) {
-      if (op == forOp.getOperation() ||
-          isa<scf::SCFDialect>(op->getDialect())) {
-        return;
-      }
-      auto it = opCoreTypes.find(op);
-      if (it != opCoreTypes.end()) {
-        it->second = OP_CUBE_ONLY;
-      }
-    });
-    opCoreTypes[forOp] = OP_CUBE_ONLY;
+  for (Operation *loop : loaderLoops) {
+    Block *body = getLoopBodyBlock(loop);
+    if (body)
+      markLoopBodyAsCube(loop, body, opCoreTypes);
   }
 
   return 0;
@@ -1409,25 +1413,36 @@ int OpClassifierPass::handleSCFYield() {
 }
 
 OpCoreType OpClassifierPass::getForInitCoreType(OpOperand *operand) const {
-  auto forOp = llvm::dyn_cast<scf::ForOp>(operand->getOwner());
-  if (!forOp) {
-    return OP_UNDETERMINED;
-  }
-  auto iterArg = forOp.getTiedLoopRegionIterArg(operand);
-  if (!iterArg) {
-    // the result is used as lower/upper bound or step
-    return OP_UNDETERMINED;
+  // Unified handling for scf.for and scf.while using the tied loop interface
+  Operation *owner = operand->getOwner();
+
+  if (auto forOp = dyn_cast<scf::ForOp>(owner)) {
+    auto iterArg = forOp.getTiedLoopRegionIterArg(operand);
+    if (!iterArg)
+      return OP_UNDETERMINED;
+    auto sourceOperand = forOp.getTiedLoopYieldedValue(iterArg);
+    if (!sourceOperand)
+      return OP_UNDETERMINED;
+    auto defOp = sourceOperand->get().getDefiningOp();
+    if (!defOp)
+      return OP_UNDETERMINED;
+    return getCoreType(defOp);
+  }
+
+  if (auto whileOp = dyn_cast<scf::WhileOp>(owner)) {
+    auto iterArg = whileOp.getTiedLoopRegionIterArg(operand);
+    if (!iterArg)
+      return OP_UNDETERMINED;
+    auto sourceOperand = whileOp.getTiedLoopYieldedValue(iterArg);
+    if (!sourceOperand)
+      return OP_UNDETERMINED;
+    auto defOp = sourceOperand->get().getDefiningOp();
+    if (!defOp)
+      return OP_UNDETERMINED;
+    return getCoreType(defOp);
   }
-  auto sourceOperand = forOp.getTiedLoopYieldedValue(iterArg);
-  if (!sourceOperand) {
-    return OP_UNDETERMINED;
-  }
-  auto defOp = sourceOperand->get().getDefiningOp();
-  if (!defOp) {
-    // might be a blockarg, not necessary for our use-case
-    return OP_UNDETERMINED;
-  }
-  return getCoreType(defOp);
+
+  return OP_UNDETERMINED;
 }
 
 // ============================================================================
@@ -1532,7 +1547,7 @@ void OpClassifierPass::splitOperationForCubeAndVector(
         continue;
       }
       OpCoreType coreType = getCoreType(user);
-      if (llvm::isa<scf::ForOp>(user)) {
+      if (llvm::isa<scf::ForOp, scf::WhileOp>(user)) {
         coreType = getForInitCoreType(&use);
       }
       if (coreType == OP_VECTOR_ONLY) {
@@ -1600,10 +1615,18 @@ int OpClassifierPass::stampToIR() {
       continue;
     OpCoreType coreType = it->second;
 
-    // Skip scf dialect operations
-    if (llvm::isa<scf::SCFDialect>(op->getDialect()))
+    // Skip most scf dialect operations except scf.condition (terminator in
+    // while's before region)
+    if (llvm::isa<scf::SCFDialect>(op->getDialect()) &&
+        !llvm::isa<scf::ConditionOp>(op))
       continue;
 
+    // scf.condition is always VECTOR (it's the condition check in while's
+    // before region)
+    if (llvm::isa<scf::ConditionOp>(op)) {
+      coreType = OP_VECTOR_ONLY;
+    }
+
     // Skip linalg operations' internal block operations
     Operation *parent = op->getParentOp();
     bool isInsideLinalgBlock = false;

From c985f3e829a64e4d6d5b8d395c728aeedeebf5ef Mon Sep 17 00:00:00 2001
From: shaoyiyang <shaoyiyang1@huawei.com>
Date: Wed, 29 Jul 2026 11:51:42 +0800
Subject: [PATCH 03/11] [ssbuffer](feat) Adapter whileOp in plan vector block

---
 .../include/DynamicCVPipeline/Common/Utils.h  |  2 +
 .../lib/DynamicCVPipeline/Common/Utils.cpp    | 29 +++++++++
 .../ComputeBlockOpt/UBUsageOptPass.cpp        |  8 ++-
 .../PlanComputeBlock/PlanCubeBlock.cpp        | 23 ++++---
 .../PlanComputeBlock/PlanVectorBlockPass.cpp  | 37 +++++------
 .../SplitDataflow/RefineArgsBlockId.cpp       | 62 +++++++++----------
 6 files changed, 97 insertions(+), 64 deletions(-)

diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
index 7ecff19b6d..620161185a 100644
--- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
+++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
@@ -246,6 +246,8 @@ bool allResultHasOneUser(Operation *op);
 
 int64_t getBTSizeFromValidBroadcastOp(linalg::BroadcastOp broadcastOp);
 
+int getLoopCarriedArgIndex(Value operand, Block *block);
+
 } // namespace CVPipeline
 } // namespace mlir
 
diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
index 6d39f18da6..39445be371 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -345,5 +346,33 @@ scf::YieldOp MainLoop::getLoopYieldOp(Operation *loopOp) {
   return {};
 }
 
+int getLoopCarriedArgIndex(Value operand, Block *block) {
+  auto barg = dyn_cast<BlockArgument>(operand);
+  if (!barg || barg.getOwner() != block) {
+    return -1;
+  }
+
+  auto parentOp = block->getParentOp();
+  if (!isa<scf::ForOp, scf::WhileOp>(parentOp)) {
+    return -1;
+  }
+
+  auto *terminator = block->getTerminator();
+  if (!terminator || !isa<scf::YieldOp>(terminator)) {
+    return -1;
+  }
+
+  int numArgs = block->getNumArguments();
+  int numYieldOperands = terminator->getNumOperands();
+  int offset = numArgs - numYieldOperands;
+  int argIdx = barg.getArgNumber() - offset;
+
+  if (argIdx < 0 || argIdx >= numYieldOperands) {
+    return -1;
+  }
+
+  return argIdx;
+}
+
 } // namespace CVPipeline
 } // namespace mlir
diff --git a/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/UBUsageOptPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/UBUsageOptPass.cpp
index 89ffd2255c..d875894665 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/UBUsageOptPass.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/ComputeBlockOpt/UBUsageOptPass.cpp
@@ -221,7 +221,7 @@ void UBUsageOptPass::buildUBUsageGraph(
   };
 
   Operation *terminator = block->getTerminator();
-  if (terminator) {
+  if (terminator && isa<scf::YieldOp>(terminator)) {
     unsigned maxArgIdx = std::min<unsigned>(block->getNumArguments(),
                                             terminator->getNumOperands());
     for (unsigned argIdx = 0; argIdx < maxArgIdx; ++argIdx) {
@@ -253,7 +253,8 @@ void UBUsageOptPass::buildUBUsageGraph(
             continue;
           }
         } else if (auto blockArg = dyn_cast<BlockArgument>(operand)) {
-          if (blockArg.getOwner() == block && terminator) {
+          if (blockArg.getOwner() == block && terminator &&
+              isa<scf::YieldOp>(terminator)) {
             unsigned numArgs = block->getNumArguments();
             unsigned numYieldOperands = terminator->getNumOperands();
             // for op offset=1, while op offset=0
@@ -738,7 +739,8 @@ bool applyRecordChange(DenseMap<int, int> &recordChange,
 llvm::LogicalResult UBUsageOptPass::UBUsageOptimization(
     Block *block, const CVPipeline::MemoryDependenceGraph &memGraph,
     CVPipeline::ComputeBlockIdManager &bm) {
-  if (!isa<scf::ForOp>(block->getParentOp())) {
+  if (!(isa<scf::ForOp>(block->getParentOp()) ||
+        isa<scf::WhileOp>(block->getParentOp()))) {
     return llvm::success();
   }
   DenseMap<Operation *, int> op2nodeId;
diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp
index de6518ea41..58ffb585aa 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp
@@ -209,14 +209,23 @@ void SeedRegionPlanner::run() {
       if (auto *def = iop.getDefiningOp()) {
         tryAddToGroup(def);
       }
-      // Check loop-carried dependencies (SCF ForOp iter_args)
+      // Check loop-carried dependencies (SCF ForOp/WhileOp iter_args)
       if (auto barg = dyn_cast<BlockArgument>(iop)) {
-        if (barg.getOwner() == block && isa<scf::ForOp>(block->getParentOp()) &&
-            barg.getArgNumber() > 0) {
-          auto *yieldOp = barg.getOwner()->getTerminator();
-          if (auto *yieldedValDef = yieldOp->getOperand(barg.getArgNumber() - 1)
-                                        .getDefiningOp()) {
-            tryAddToGroup(yieldedValDef);
+        if (barg.getOwner() == block) {
+          auto *terminator = block->getTerminator();
+          if (terminator && isa<scf::YieldOp>(terminator) &&
+              isa<scf::ForOp, scf::WhileOp>(block->getParentOp())) {
+            unsigned numArgs = block->getNumArguments();
+            unsigned numYieldOperands = terminator->getNumOperands();
+            int offset = (int)numArgs - (int)numYieldOperands;
+            int argIdx = (int)barg.getArgNumber() - offset;
+
+            if (argIdx >= 0 && argIdx < (int)numYieldOperands) {
+              Value yielded = terminator->getOperand(argIdx);
+              if (auto *yieldDefOp = yielded.getDefiningOp()) {
+                tryAddToGroup(yieldDefOp);
+              }
+            }
           }
         }
       }
diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp
index 20f38cea90..f21b2f3884 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp
@@ -245,19 +245,6 @@ findOpsAdjacentToCube(Block *block, const SmallVector<Operation *> &fuseGroup,
   return toProcess;
 }
 
-static int getLoopCarriedArgIndex(Value operand, Block *block) {
-  auto barg = dyn_cast<BlockArgument>(operand);
-  if (!barg || barg.getOwner() != block ||
-      !isa<scf::ForOp>(block->getParentOp())) {
-    return -1;
-  }
-  unsigned argIdx = barg.getArgNumber();
-  if (argIdx == 0) {
-    return -1;
-  }
-  return argIdx;
-}
-
 static SetVector<Operation *>
 collectKeepOps(Block *block, SmallVector<Operation *> toProcess,
                const SmallVector<Operation *> &fuseGroup,
@@ -281,14 +268,20 @@ collectKeepOps(Block *block, SmallVector<Operation *> toProcess,
       }
 
       // Loop-carried dependency: block argument -> yielded value
-      int argIdx = getLoopCarriedArgIndex(operand, block);
+      int argIdx = CVPipeline::getLoopCarriedArgIndex(operand, block);
       if (argIdx == -1) {
         continue;
       }
       auto barg = cast<BlockArgument>(operand);
-      auto *yieldOp = barg.getOwner()->getTerminator();
-      auto *yieldedDef = yieldOp->getOperand(argIdx - 1).getDefiningOp();
-      if (!keepOps.contains(yieldedDef) &&
+      auto *terminator = barg.getOwner()->getTerminator();
+      if (!terminator || !isa<scf::YieldOp>(terminator)) {
+        continue;
+      }
+
+      Value yielded = terminator->getOperand(argIdx);
+      Operation *yieldedDef = yielded.getDefiningOp();
+
+      if (yieldedDef && !keepOps.contains(yieldedDef) &&
           llvm::is_contained(fuseGroup, yieldedDef)) {
         toProcess.push_back(yieldedDef);
       }
@@ -382,16 +375,16 @@ extractToProcessFromFuseGroup(Block *block,
   }
 
   SetVector<Operation *> toRemove;
-  auto forOp = dyn_cast<scf::ForOp>(block->getParentOp());
-  if (forOp) {
+  auto *terminator = block->getTerminator();
+  if (terminator && isa<scf::YieldOp>(terminator) &&
+      isa<scf::ForOp, scf::WhileOp>(block->getParentOp())) {
     for (auto op : nowFuseGroup) {
       for (auto operand : op->getOperands()) {
         int argIdx = getLoopCarriedArgIndex(operand, block);
-        if (argIdx <= 0) {
+        if (argIdx == -1) {
           continue;
         }
-        auto *yieldOp = block->getTerminator();
-        auto yieldOperand = yieldOp->getOperand(argIdx - 1);
+        auto yieldOperand = terminator->getOperand(argIdx);
         auto *defOp = yieldOperand.getDefiningOp();
         if (defOp && bm.getBlockIdByOp(defOp) == -1 &&
             !llvm::is_contained(nowFuseGroup, defOp)) {
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp
index 0e835386fe..9a9b113cd4 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp
@@ -39,20 +39,7 @@ static constexpr const char *DEBUG_TYPE = "refine-args-block-id";
 
 using namespace mlir::triton;
 
-int getLoopCarriedArgIndex(Value operand, Block *block) {
-  auto barg = dyn_cast<BlockArgument>(operand);
-  if (!barg || barg.getOwner() != block ||
-      !isa<scf::ForOp>(block->getParentOp())) {
-    return -1;
-  }
-  unsigned argIdx = barg.getArgNumber();
-  if (argIdx == 0) {
-    return -1;
-  }
-  return argIdx;
-}
-
-int findFirstUser(BlockArgument iterArg, Block *forBlock,
+int findFirstUser(Value iterArg, Block *forBlock,
                   CVPipeline::ComputeBlockIdManager &bm) {
   llvm::SetVector<Value> visited;
   SmallVector<std::pair<Value, int>> worklist;
@@ -90,9 +77,10 @@ bool isDependenceOther(Operation *yieldDefOp, Block *forBlock, int argsId,
       }
     } else {
       // if have block argument from for block. Skip;
-      if (getLoopCarriedArgIndex(operand, forBlock) != argsId + 1) {
+      if (CVPipeline::getLoopCarriedArgIndex(operand, forBlock) != argsId) {
         LOG_DEBUG("Yield def op depends on other arg:"
-                  << getLoopCarriedArgIndex(operand, forBlock) << "\n");
+                  << CVPipeline::getLoopCarriedArgIndex(operand, forBlock)
+                  << "\n");
         return true;
       }
     }
@@ -109,19 +97,27 @@ bool isDependenceOther(Operation *yieldDefOp, Block *forBlock, int argsId,
   return false;
 }
 
-void processOnefor(scf::ForOp forOp, CVPipeline::ComputeBlockIdManager &bm,
-                   const CVPipeline::MemoryDependenceGraph &memGraph) {
-
-  Block *forBlock = &forOp.getRegion().front();
-  auto yieldOp = dyn_cast<scf::YieldOp>(forBlock->getTerminator());
-  if (!yieldOp) {
-    LOG_DEBUG("No yield op found in for block\n");
+void processOneLoop(Operation *loopOp, CVPipeline::ComputeBlockIdManager &bm,
+                    const CVPipeline::MemoryDependenceGraph &memGraph) {
+  auto ml = CVPipeline::MainLoop(loopOp);
+  Block *loopBlock = nullptr;
+  for (Region &region : loopOp->getRegions()) {
+    if (auto *terminator = region.front().getTerminator()) {
+      if (isa<scf::YieldOp>(terminator)) {
+        loopBlock = &region.front();
+        break;
+      }
+    }
+  }
+  if (!loopBlock) {
+    LOG_DEBUG("No yield op found in loop block\n");
     return;
   }
-  ArrayRef<BlockArgument> iterArgs = forOp.getRegionIterArgs();
+  auto yieldOp = cast<scf::YieldOp>(loopBlock->getTerminator());
 
+  SmallVector<Value> iterArgs = ml.getIterArgs();
   for (size_t i = 0; i < iterArgs.size(); ++i) {
-    BlockArgument argsi = iterArgs[i];
+    auto argsi = iterArgs[i];
 
     Value yieldOperand = yieldOp.getOperand(i);
     Operation *yieldDefOp = yieldOperand.getDefiningOp();
@@ -132,14 +128,14 @@ void processOnefor(scf::ForOp forOp, CVPipeline::ComputeBlockIdManager &bm,
     }
     LOG_DEBUG("yieldDefOp: " << *yieldDefOp << "\n"
                              << "idx: " << i << "\n");
-    if (isDependenceOther(yieldDefOp, forBlock, i, memGraph)) {
+    if (isDependenceOther(yieldDefOp, loopBlock, i, memGraph)) {
       continue;
     }
 
     int updateBlockId = bm.getBlockIdByOp(yieldDefOp);
     LOG_DEBUG("Update block id for yield def op: " << updateBlockId << "\n");
 
-    int firstUserBlockId = findFirstUser(argsi, forBlock, bm);
+    int firstUserBlockId = findFirstUser(argsi, loopBlock, bm);
     LOG_DEBUG("First user block id: " << firstUserBlockId << "\n");
 
     if (firstUserBlockId != -1 && updateBlockId != firstUserBlockId) {
@@ -150,7 +146,7 @@ void processOnefor(scf::ForOp forOp, CVPipeline::ComputeBlockIdManager &bm,
       auto firstUserOps = bm.getOpsByBlockId(firstUserBlockId);
       Operation *lastOpInFirstUserBlock = nullptr;
       for (Operation *op : firstUserOps) {
-        if (op->getBlock() != forBlock)
+        if (op->getBlock() != loopBlock)
           continue;
         if (!lastOpInFirstUserBlock ||
             op->isBeforeInBlock(lastOpInFirstUserBlock)) {
@@ -175,11 +171,13 @@ void RefineArgsBlockIdPass::runOnOperation() {
   CVPipeline::ComputeBlockIdManager bm(moduleOp);
   auto &aa = getAnalysis<AliasAnalysis>();
   LOG_DEBUG(*moduleOp);
-  moduleOp.walk([&](scf::ForOp forOp) {
-    if (forOp->hasAttr("ssbuffer.main_loop")) {
-      auto memDepGraph = CVPipeline::MemoryDependenceGraph(forOp, aa);
-      processOnefor(forOp, bm, memDepGraph);
+  moduleOp.walk([&](Operation *op) {
+    if (!op->hasAttr(CVPipeline::kMainLoop) ||
+        !isa<scf::ForOp, scf::WhileOp>(op)) {
+      return;
     }
+    auto memDepGraph = CVPipeline::MemoryDependenceGraph(op, aa);
+    processOneLoop(op, bm, memDepGraph);
   });
 
   LOG_DEBUG("--- exit RefineArgsBlockIdPass --->\n");

From 644d9e2d0cd13356999cd64ff0a4c6827b3855b4 Mon Sep 17 00:00:00 2001
From: dingyi <dingyi@huawei.com>
Date: Fri, 31 Jul 2026 15:36:26 +0800
Subject: [PATCH 04/11] [ssbuffer](feat) support scf.while in SplitDataflow

---
 .../SplitDataflow/DataDependencyAnalysis.h    |  11 +-
 .../SplitDataflow/AddBlockIdForControlOps.cpp |  25 +--
 .../SplitDataflow/DataDependencyAnalysis.cpp  | 148 ++++++++++--------
 .../InterCoreTransferAndSync.cpp              |   2 +-
 .../SplitDataflow/MarkMainLoop.cpp            |  31 ++--
 .../SplitDataflow/SeparateCVScope.cpp         | 140 +++++++++++------
 .../separate_cv_scope_canonicalize_ut.mlir    |  68 ++++++++
 .../separate_cv_scope_ut.mlir                 |  45 ++++++
 8 files changed, 329 insertions(+), 141 deletions(-)

diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h
index 15b120df57..6074f4ff7a 100644
--- a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h
+++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h
@@ -164,16 +164,17 @@ class DataDependencyAnalysisPass
   collectDiffCoreTypeUsers(mlir::BlockArgument iterArg,
                            llvm::StringRef initCoreType);
   void
-  insertProducerAndRecordDeps(scf::ForOp forOp, mlir::BlockArgument iterArg,
+  insertProducerAndRecordDeps(mlir::LoopLikeOpInterface loopOp,
+                              mlir::BlockArgument loopArg,
                               llvm::StringRef initCoreType,
                               llvm::SmallVector<mlir::Operation *> &diffUsers,
                               DataDependencyInfo &info);
-  void insertConsumerAndRecordDeps(scf::ForOp forOp, mlir::Value yieldedValue,
-                                   int iterArgIndex,
+  void insertConsumerAndRecordDeps(mlir::LoopLikeOpInterface loopOp,
+                                   mlir::Value yieldedValue, int iterArgIndex,
                                    llvm::StringRef initCoreType,
                                    DataDependencyInfo &info);
-  void recordInitValueDeps(scf::ForOp forOp, mlir::Value initValue,
-                           llvm::StringRef yieldCoreType,
+  void recordInitValueDeps(mlir::LoopLikeOpInterface loopOp,
+                           mlir::Value initValue, llvm::StringRef yieldCoreType,
                            DataDependencyInfo &info);
   void updateCoreTypeAtIndex(Operation *op, int index,
                              llvm::StringRef newCoreType);
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp
index c01a70351f..18f3c0fd6b 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp
@@ -55,26 +55,29 @@ void AddBlockIdForControlOpsPass::runOnOperation() {
       return;
     }
 
-    if (isa<scf::ForOp, scf::IfOp>(op)) {
+    if (isa<scf::ForOp, scf::IfOp, scf::WhileOp>(op)) {
       maxBlockId++;
       setOpBlockId(op, maxBlockId);
       LOG_DEBUG("Added block_id " << maxBlockId << " to " << *op << "\n");
     }
 
-    // coretype of scf.yield may not be the same with defining op
-    if (isa<scf::YieldOp>(op) && isa<scf::IfOp>(op->getParentOp())) {
-      Operation *parentOp = op->getParentOp();
-      auto ifBlockIdOpt = CVPipeline::getOpBlockId(parentOp);
+    Operation *parentOp = op->getParentOp();
+    bool isControlTerminator =
+        (isa<scf::YieldOp>(op) && isa<scf::IfOp, scf::WhileOp>(parentOp)) ||
+        (isa<scf::ConditionOp>(op) && isa<scf::WhileOp>(parentOp));
+    if (isControlTerminator) {
+      auto parentBlockIdOpt = CVPipeline::getOpBlockId(parentOp);
 
-      int yieldBlockId;
-      if (ifBlockIdOpt) {
-        yieldBlockId = *ifBlockIdOpt;
+      int terminatorBlockId;
+      if (parentBlockIdOpt) {
+        terminatorBlockId = *parentBlockIdOpt;
       } else {
         maxBlockId++;
-        yieldBlockId = maxBlockId;
+        terminatorBlockId = maxBlockId;
       }
-      setOpBlockId(op, yieldBlockId);
-      LOG_DEBUG("Added block_id " << yieldBlockId << " to " << *op << "\n");
+      setOpBlockId(op, terminatorBlockId);
+      LOG_DEBUG("Added block_id " << terminatorBlockId << " to " << *op
+                                  << "\n");
     }
   });
 
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp
index 74aacd3a6d..01b7cfef98 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp
@@ -127,7 +127,8 @@ void DataDependencyAnalysisPass::updateCoreTypeAtIndex(
 bool DataDependencyAnalysisPass::isControlFlowOp(mlir::Operation *op) {
   if (!op)
     return false;
-  return isa<scf::ForOp>(op) || isa<scf::IfOp>(op) || isa<scf::YieldOp>(op);
+  return isa<scf::ForOp, scf::WhileOp>(op) || isa<scf::IfOp>(op) ||
+         isa<scf::YieldOp, scf::ConditionOp>(op);
 }
 
 bool DataDependencyAnalysisPass::isCubeOrVectorOp(mlir::Operation *op) {
@@ -218,11 +219,7 @@ bool DataDependencyAnalysisPass::isOuterOpArg(mlir::Value value) {
   return false;
 }
 
-// Resolve nested scf.for iterArg init value and return its defining op
-// If `initValue` is a BlockArgument of an outer scf.for iterArg, walk up
-// the enclosing for-loops and return the defining op of the real init value
-// (the corresponding `forOp.getInitArgs()[argIndex]`) until a non-BlockArgument
-// is found.
+// Resolve nested loop iterArg init values until reaching a non-BlockArgument.
 mlir::Value DataDependencyAnalysisPass::resolveNestedIterArgInitValue(
     mlir::Value initValue) {
   llvm::DenseSet<mlir::Value> visited;
@@ -233,18 +230,27 @@ mlir::Value DataDependencyAnalysisPass::resolveNestedIterArgInitValue(
     auto blockArg = dyn_cast<mlir::BlockArgument>(currentValue);
     if (!blockArg)
       break;
-    // Parent op of the block containing this argument
     mlir::Operation *parentOp = blockArg.getOwner()->getParentOp();
     LOG_DEBUG("parentOp: " << *parentOp << "\n");
-    auto outerFor = dyn_cast<scf::ForOp>(parentOp);
-    if (!outerFor)
+
+    unsigned argIndex = blockArg.getArgNumber();
+    ValueRange initArgs;
+    if (auto outerFor = dyn_cast<scf::ForOp>(parentOp)) {
+      if (argIndex == 0)
+        break;
+      --argIndex;
+      initArgs = outerFor.getInitArgs();
+    } else if (auto outerWhile = dyn_cast<scf::WhileOp>(parentOp)) {
+      // Triton keeps while before/after block args aligned 1:1 with init args.
+      initArgs = outerWhile.getInits();
+    } else {
       break;
-    unsigned argIndex = blockArg.getArgNumber() - 1;
+    }
+
     LOG_DEBUG("argIndex: " << argIndex << "\n");
-    if (argIndex >= outerFor.getInitArgs().size())
+    if (argIndex >= initArgs.size())
       break;
-    // Move to the init value corresponding to this iterArg
-    currentValue = outerFor.getInitArgs()[argIndex];
+    currentValue = initArgs[argIndex];
   }
   LOG_DEBUG("currentValue: " << currentValue << "\n");
   return currentValue;
@@ -402,19 +408,20 @@ DataDependencyAnalysisPass::collectDiffCoreTypeUsers(
   return diffUsers;
 }
 
-// Inserts a producer block at the beginning of the for loop body and records
+// Inserts a producer block at the beginning of the loop body and records
 // cross-core-type dependencies for each user in diffUsers.
 void DataDependencyAnalysisPass::insertProducerAndRecordDeps(
-    scf::ForOp forOp, mlir::BlockArgument iterArg, llvm::StringRef initCoreType,
+    mlir::LoopLikeOpInterface loopOp, mlir::BlockArgument loopArg,
+    llvm::StringRef initCoreType,
     llvm::SmallVector<mlir::Operation *> &diffUsers, DataDependencyInfo &info) {
   auto &v2cDependencies = info.getV2CDependencies();
   auto &c2vDependencies = info.getC2VDependencies();
-  auto &blockInfoMap = info.getBlockInfoMap();
 
-  OpBuilder builder(forOp);
-  Block &bodyBlock = forOp.getRegion().front();
-  builder.setInsertionPointToStart(&bodyBlock);
-  Location loc = forOp.getLoc();
+  Operation *loopOperation = loopOp.getOperation();
+  CVPipeline::MainLoop loop(loopOperation);
+  OpBuilder builder(loopOperation);
+  builder.setInsertionPointToStart(loop.getBody());
+  Location loc = loopOperation->getLoc();
   auto constOp = createBlockInfoConstOp(builder, loc, initCoreType, info);
   int newId = *CVPipeline::getOpBlockId(constOp);
 
@@ -440,7 +447,7 @@ void DataDependencyAnalysisPass::insertProducerAndRecordDeps(
     auto &targetDeps = (depType == DependencyType::VectorToCube)
                            ? v2cDependencies
                            : c2vDependencies;
-    if (!collectDepInfo(iterArg, depType, targetDeps, newId, userBlockId,
+    if (!collectDepInfo(loopArg, depType, targetDeps, newId, userBlockId,
                         info)) {
       continue;
     }
@@ -456,13 +463,13 @@ void DataDependencyAnalysisPass::insertProducerAndRecordDeps(
 }
 
 void DataDependencyAnalysisPass::insertConsumerAndRecordDeps(
-    scf::ForOp forOp, mlir::Value yieldedValue, int iterArgIndex,
-    llvm::StringRef initCoreType, DataDependencyInfo &info) {
+    mlir::LoopLikeOpInterface loopOp, mlir::Value yieldedValue,
+    int iterArgIndex, llvm::StringRef initCoreType, DataDependencyInfo &info) {
   auto &v2cDependencies = info.getV2CDependencies();
   auto &c2vDependencies = info.getC2VDependencies();
-  auto &blockInfoMap = info.getBlockInfoMap();
 
-  auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+  Operation *loopOperation = loopOp.getOperation();
+  scf::YieldOp yieldOp = CVPipeline::MainLoop::getLoopYieldOp(loopOperation);
   OpBuilder builder(yieldOp);
   Location loc = yieldOp.getLoc();
   auto constOp = createBlockInfoConstOp(builder, loc, initCoreType, info);
@@ -495,7 +502,7 @@ void DataDependencyAnalysisPass::insertConsumerAndRecordDeps(
   targetDeps.back().consumerYieldOp = yieldOp;
 
   updateCoreTypeAtIndex(yieldOp, iterArgIndex, initCoreType);
-  updateCoreTypeAtIndex(forOp, iterArgIndex, initCoreType);
+  updateCoreTypeAtIndex(loopOperation, iterArgIndex, initCoreType);
 
   LOG_DEBUG("Recorded yield producer dependency: "
             << initCoreType << ", iniProducerBlockId=" << yieldedDefBlockId
@@ -503,8 +510,8 @@ void DataDependencyAnalysisPass::insertConsumerAndRecordDeps(
 }
 
 void DataDependencyAnalysisPass::recordInitValueDeps(
-    scf::ForOp forOp, mlir::Value initValue, llvm::StringRef yieldCoreType,
-    DataDependencyInfo &info) {
+    mlir::LoopLikeOpInterface loopOp, mlir::Value initValue,
+    llvm::StringRef yieldCoreType, DataDependencyInfo &info) {
   auto &v2cDependencies = info.getV2CDependencies();
   auto &c2vDependencies = info.getC2VDependencies();
 
@@ -516,12 +523,12 @@ void DataDependencyAnalysisPass::recordInitValueDeps(
   }
   int initDefBlockId = *initDefBlockIdOpt;
 
-  auto forOpBlockIdOpt = CVPipeline::getOpBlockId(forOp);
-  if (!forOpBlockIdOpt) {
-    LOG_DEBUG("Warning: ForOp block ID not found.\n");
+  auto loopBlockIdOpt = CVPipeline::getOpBlockId(loopOp.getOperation());
+  if (!loopBlockIdOpt) {
+    LOG_DEBUG("Warning: Loop block ID not found.\n");
     return;
   }
-  int forOpBlockId = *forOpBlockIdOpt;
+  int loopBlockId = *loopBlockIdOpt;
 
   DependencyType depType;
   if (yieldCoreType == ssbufferCoreTypeVectorAttr) {
@@ -534,27 +541,27 @@ void DataDependencyAnalysisPass::recordInitValueDeps(
                          ? v2cDependencies
                          : c2vDependencies;
   LOG_DEBUG("iniProducerBlockId=" << initDefBlockId << ", iniConsumerBlockId="
-                                  << forOpBlockId << "\n");
+                                  << loopBlockId << "\n");
   if (!collectDepInfo(initValue, depType, targetDeps, initDefBlockId,
-                      forOpBlockId, info)) {
+                      loopBlockId, info)) {
     return;
   }
 
   LOG_DEBUG("Recorded init value dependency: "
             << yieldCoreType << ", iniProducerBlockId=" << initDefBlockId
-            << ", iniConsumerBlockId=" << forOpBlockId << "\n");
+            << ", iniConsumerBlockId=" << loopBlockId << "\n");
 }
 
-bool checkYieldCoreType(mlir::Operation *yieldOp) {
-  if (!isa<scf::YieldOp>(yieldOp)) {
+static bool checkLoopYieldCoreType(scf::YieldOp yieldOp) {
+  if (!yieldOp) {
     return false;
   }
-  for (unsigned index = 0; index < yieldOp->getNumOperands(); ++index) {
-    mlir::Value value = yieldOp->getOperand(index);
+  for (unsigned index = 0; index < yieldOp.getNumOperands(); ++index) {
+    mlir::Value value = yieldOp.getOperand(index);
     llvm::StringRef yieldCoreType = getCoreTypeWithIndex(yieldOp, index);
 
     mlir::Operation *definingOp = value.getDefiningOp();
-    if (!definingOp || !isa<scf::ForOp>(definingOp)) {
+    if (!definingOp || !isa<scf::ForOp, scf::WhileOp>(definingOp)) {
       continue;
     }
     auto defResult = dyn_cast<mlir::OpResult>(value);
@@ -569,31 +576,43 @@ bool checkYieldCoreType(mlir::Operation *yieldOp) {
   return true;
 }
 
-// Process iterArg dependencies for all scf.for operations in the module.
-// This function iterates through all for loops and checks each iterArg to
-// determine if there are cross-core-type data dependencies.
+// Process iterArg dependencies for all supported loop-like operations.
 void DataDependencyAnalysisPass::processIterArgDependencies() {
   auto &info = getAnalysis<DataDependencyInfo>();
 
-  // Step1: Collect all scf.for operations in the module
-  llvm::SmallVector<scf::ForOp> forOps;
-  module.walk([&](scf::ForOp forOp) { forOps.push_back(forOp); });
-  LOG_DEBUG("Processing iterArg dependencies, found " << forOps.size()
-                                                      << " scf.for ops\n");
-
-  // Step2: Process each iterArg of each scf.for operation
-  for (scf::ForOp forOp : forOps) {
-    size_t numIterArgs = forOp.getInitArgs().size();
-    mlir::Operation *yieldOp = forOp.getBody()->getTerminator();
-    if (!checkYieldCoreType(yieldOp)) {
+  llvm::SmallVector<mlir::LoopLikeOpInterface> loopOps;
+  module.walk(
+      [&](mlir::LoopLikeOpInterface loopOp) { loopOps.push_back(loopOp); });
+  LOG_DEBUG("Processing iterArg dependencies, found " << loopOps.size()
+                                                      << " loop ops\n");
+
+  for (mlir::LoopLikeOpInterface loopOp : loopOps) {
+    Operation *loopOperation = loopOp.getOperation();
+    if (!isa<scf::ForOp, scf::WhileOp>(loopOperation)) {
+      continue;
+    }
+
+    CVPipeline::MainLoop loop(loopOperation);
+    scf::YieldOp yieldOp = CVPipeline::MainLoop::getLoopYieldOp(loopOperation);
+    if (!checkLoopYieldCoreType(yieldOp)) {
       LOG_DEBUG("[ERROR]: Yield core type mismatch defining op\n");
       CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
       return;
     }
+
+    size_t numIterArgs = loopOp.getInits().size();
+    SmallVector<Value> iterArgs = loop.getIterArgs();
+    if (iterArgs.size() != numIterArgs) {
+      LOG_DEBUG("[ERROR]: Loop iter_arg count does not match init count\n");
+      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+      return;
+    }
+
     for (int iterArgIndex = 0; iterArgIndex < numIterArgs; ++iterArgIndex) {
-      mlir::Value initValue = forOp.getInits()[iterArgIndex];
-      mlir::BlockArgument iterArg = forOp.getRegionIterArg(iterArgIndex);
-      mlir::Value yieldedValue = forOp.getYieldedValues()[iterArgIndex];
+      mlir::Value initValue = loopOp.getInits()[iterArgIndex];
+      mlir::BlockArgument iterArg =
+          cast<mlir::BlockArgument>(iterArgs[iterArgIndex]);
+      mlir::Value yieldedValue = loopOp.getYieldedValues()[iterArgIndex];
       LOG_DEBUG("initValue" << initValue << "\n");
       LOG_DEBUG("yieldedValue" << yieldedValue << "\n");
 
@@ -610,7 +629,7 @@ void DataDependencyAnalysisPass::processIterArgDependencies() {
       if (!yieldedDefOp) {
         continue;
       }
-      auto yieldCoreType = getCoreTypeWithIndex(forOp, iterArgIndex);
+      auto yieldCoreType = getCoreTypeWithIndex(loopOperation, iterArgIndex);
 
       if (!initDefOp) {
         auto realInitValue = resolveNestedIterArgInitValue(initValue);
@@ -638,7 +657,7 @@ void DataDependencyAnalysisPass::processIterArgDependencies() {
       if (initCoreType == yieldCoreType || isCubeOrVectorOp(initDefOp)) {
         auto diffUsers = collectDiffCoreTypeUsers(iterArg, yieldCoreType);
         if (!diffUsers.empty()) {
-          insertProducerAndRecordDeps(forOp, iterArg, yieldCoreType, diffUsers,
+          insertProducerAndRecordDeps(loopOp, iterArg, yieldCoreType, diffUsers,
                                       info);
         }
       } else {
@@ -661,14 +680,14 @@ void DataDependencyAnalysisPass::processIterArgDependencies() {
           }
         }
         if (!initCoreTypeUsers.empty() && !yieldCoreTypeUsers.empty()) {
-          recordInitValueDeps(forOp, initValue, yieldCoreType, info);
-          insertProducerAndRecordDeps(forOp, iterArg, yieldCoreType,
+          recordInitValueDeps(loopOp, initValue, yieldCoreType, info);
+          insertProducerAndRecordDeps(loopOp, iterArg, yieldCoreType,
                                       yieldCoreTypeUsers, info);
         } else if (!initCoreTypeUsers.empty() && yieldCoreTypeUsers.empty()) {
-          insertConsumerAndRecordDeps(forOp, yieldedValue, iterArgIndex,
+          insertConsumerAndRecordDeps(loopOp, yieldedValue, iterArgIndex,
                                       initCoreType, info);
         } else if (initCoreTypeUsers.empty() && !yieldCoreTypeUsers.empty()) {
-          recordInitValueDeps(forOp, initValue, yieldCoreType, info);
+          recordInitValueDeps(loopOp, initValue, yieldCoreType, info);
         } else {
           LOG_DEBUG("no dependencies with: " << iterArg << "\n");
         }
@@ -1048,6 +1067,9 @@ void DataDependencyAnalysisPass::runOnOperation() {
 
   // Step 2: Analyze iter_args dependencies
   processIterArgDependencies();
+  if (CVPipeline::hasFallbackAttr(module)) {
+    return;
+  }
   createBlockInfoMap(info);
 
   // Step 3: Analyze dependencies (populate v2c, c2v lists)
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp
index 36ef8df735..d4b6ea2ee6 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp
@@ -437,7 +437,7 @@ InterCoreTransferAndSyncPass::findMainLoopforTransfer(Operation *endOp,
   }
   Operation *current = lca;
   while (current) {
-    if (isa<scf::ForOp>(current)) {
+    if (isa<scf::ForOp, scf::WhileOp>(current)) {
       return current;
     }
     current = current->getParentOp();
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp
index cf789b0810..7850232b4c 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp
@@ -44,21 +44,24 @@ void MarkMainLoopPass::runOnOperation() {
   }
 
   int mainLoopIdCounter = 0;
-  SmallVector<scf::ForOp> mainLoops;
+  SmallVector<Operation *> mainLoops;
 
-  // Find all candidate main loops
+  // Find all candidate main loops (ForOp + WhileOp)
   module.walk([&](Operation *op) {
     if (isa<hivm::FixpipeOp, hivm::CopyOp>(op)) {
       if (auto forOp = op->getParentOfType<scf::ForOp>()) {
         mainLoops.push_back(forOp);
       }
+      if (auto whileOp = op->getParentOfType<scf::WhileOp>()) {
+        mainLoops.push_back(whileOp);
+      }
     }
   });
 
-  for (scf::ForOp forOp : mainLoops) {
-    if (!forOp->hasAttr(CVPipeline::kMainLoop)) {
+  for (Operation *loopOp : mainLoops) {
+    if (!loopOp->hasAttr(CVPipeline::kMainLoop)) {
       // Add attribute with integer value (current counter ID)
-      forOp->setAttr(
+      loopOp->setAttr(
           CVPipeline::kMainLoop,
           Builder(module.getContext()).getI32IntegerAttr(mainLoopIdCounter));
       mainLoopIdCounter++;
@@ -67,24 +70,24 @@ void MarkMainLoopPass::runOnOperation() {
 
   // Remove main_loop attribute from outer loops if nested loops both have it
   // Keep only the innermost main_loop
-  SmallVector<scf::ForOp> allMainLoops;
-  module.walk([&](scf::ForOp forOp) {
-    if (forOp->hasAttr(CVPipeline::kMainLoop)) {
-      allMainLoops.push_back(forOp);
+  SmallVector<Operation *> allMainLoops;
+  module.walk([&](Operation *loopOp) {
+    if (CVPipeline::isMainLoopOp(loopOp)) {
+      allMainLoops.push_back(loopOp);
     }
   });
 
-  for (scf::ForOp forOp : allMainLoops) {
-    // Check if there's any nested for loop with main_loop attribute
+  for (Operation *loopOp : allMainLoops) {
+    // Check if there's any nested loop with main_loop attribute
     bool hasNestedMainLoop = false;
-    forOp.walk([&](scf::ForOp nestedForOp) {
-      if (nestedForOp != forOp && nestedForOp->hasAttr(CVPipeline::kMainLoop)) {
+    loopOp->walk([&](Operation *nestedLoopOp) {
+      if (nestedLoopOp != loopOp && CVPipeline::isMainLoopOp(nestedLoopOp)) {
         hasNestedMainLoop = true;
       }
     });
     // Remove attribute from outer loop if inner loop also has it
     if (hasNestedMainLoop) {
-      forOp->removeAttr(CVPipeline::kMainLoop);
+      loopOp->removeAttr(CVPipeline::kMainLoop);
     }
   }
 
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp
index 46505d853c..4c71999b42 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp
@@ -353,7 +353,8 @@ static bool canSkipForwardingUse(OpOperand &use, StringRef scopeType) {
     }
     unsigned resultIndex = operandIndex - kForOpOperandPrefixCount;
     return resultIndex < forOp.getNumResults() &&
-           info->getResultType(resultIndex) != scopeType;
+           info->getResultType(resultIndex) != scopeType &&
+           !needsLoopCarryPreserve(forOp, resultIndex, scopeType);
   }
 
   if (auto whileOp = dyn_cast<scf::WhileOp>(user)) {
@@ -367,7 +368,9 @@ static bool canSkipForwardingUse(OpOperand &use, StringRef scopeType) {
     if (needsLoopCarryPreserve(whileOp, slotIdx, scopeType)) {
       return false;
     }
-    return true;
+    return llvm::none_of(whileOp->getUsers(), [&](Operation *user) {
+      return matchesScope(user, scopeType);
+    });
   }
   return false;
 }
@@ -556,6 +559,10 @@ static UseCheckResult checkConditionUse(OpOperand &use, Operation *owner,
   unsigned idx = use.getOperandNumber();
   if (conditionOp->getParentOp() != owner || idx == 0 || idx != slotIndex + 1) {
     Operation *parentOp = conditionOp->getParentOp();
+    if (parentOp && idx == 0 &&
+        controlFlowOpHasScopeContent(parentOp, scopeType)) {
+      return UseCheckResult::Active;
+    }
     if (parentOp && !matchesScope(parentOp, scopeType)) {
       return UseCheckResult::Continue;
     }
@@ -684,10 +691,52 @@ static bool isProducedByForeignScope(Value operand, StringRef scopeType) {
   return !matchesScope(producer, scopeType);
 }
 
-static LogicalResult neutralizeYieldInRegion(Operation *op,
-                                             const CoreTypeInfo &info,
-                                             StringRef scopeType,
-                                             Location loc) {
+static LogicalResult
+neutralizeCarriedTerminatorOperand(Operation *op, const CoreTypeInfo &info,
+                                   StringRef scopeType, Location loc,
+                                   unsigned slotIndex, OpOperand &operand) {
+  if (info.getResultType(slotIndex) == scopeType) {
+    return success();
+  }
+
+  // Preserve values read by an in-loop consumer through an iter_arg.
+  if (needsLoopCarryPreserve(op, slotIndex, scopeType)) {
+    return success();
+  }
+
+  Value oldOperand = operand.get();
+
+  // A foreign-scope producer must not be kept alive solely by the parent
+  // result-user check.
+  bool isLoopOp = isa<scf::ForOp, scf::WhileOp>(op);
+  if ((!isLoopOp || !isProducedByForeignScope(oldOperand, scopeType)) &&
+      slotIndex < op->getNumResults()) {
+    if (Operation *resultUser =
+            findLiveUser(op->getResult(slotIndex), scopeType)) {
+      logDebug("skip neutralizing carried operand #", slotIndex, " for scope ",
+               scopeType, " because parent result #", slotIndex,
+               " still has live user '", resultUser->getName().getStringRef(),
+               "'");
+      return success();
+    }
+  }
+
+  OpBuilder builder(operand.getOwner());
+  Value replacement = buildNeutralValue(builder, oldOperand, loc, scopeType);
+  if (!replacement) {
+    logDebug("neutralizeRegionTerminators failed for op '",
+             op->getName().getStringRef(), "' at carried operand #", slotIndex,
+             " in scope ", scopeType);
+    return failure();
+  }
+  operand.set(replacement);
+  return success();
+}
+
+static LogicalResult neutralizeRegionTerminators(Operation *op,
+                                                 const CoreTypeInfo &info,
+                                                 StringRef scopeType,
+                                                 Location loc) {
   if (op->getNumRegions() == 0) {
     return success();
   }
@@ -698,49 +747,23 @@ static LogicalResult neutralizeYieldInRegion(Operation *op,
     }
 
     for (Block &block : region) {
-      auto yieldOp = dyn_cast<scf::YieldOp>(block.getTerminator());
-      if (!yieldOp) {
+      Operation *terminator = block.getTerminator();
+      unsigned carriedOperandOffset = 0;
+      if (isa<scf::ConditionOp>(terminator)) {
+        // Operand 0 is the while condition; carried args start at operand 1.
+        carriedOperandOffset = 1;
+      } else if (!isa<scf::YieldOp>(terminator)) {
         continue;
       }
 
-      OpBuilder builder(yieldOp);
-      for (unsigned i = 0; i < yieldOp.getNumOperands(); ++i) {
-        if (info.getResultType(i) == scopeType) {
-          continue;
-        }
-
-        // First defense: skip neutralization when an in-loop consumer reads the
-        // carried value through an iter_arg.
-        if (needsLoopCarryPreserve(op, i, scopeType)) {
-          continue;
-        }
-
-        Value oldOperand = yieldOp.getOperand(i);
-
-        // Second defense: skip the result-user check when the value is produced
-        // by a foreign-scope op to prevent it from being trapped.
-        bool isLoopOp = isa<scf::ForOp, scf::WhileOp>(op);
-        if ((!isLoopOp || !isProducedByForeignScope(oldOperand, scopeType)) &&
-            i < op->getNumResults()) {
-          if (Operation *resultUser =
-                  findLiveUser(op->getResult(i), scopeType)) {
-            logDebug("skip neutralizing yield operand #", i, " for scope ",
-                     scopeType, " because parent result #", i,
-                     " still has live user '",
-                     resultUser->getName().getStringRef(), "'");
-            continue;
-          }
-        }
-
-        Value replacement =
-            buildNeutralValue(builder, oldOperand, loc, scopeType);
-        if (!replacement) {
-          logDebug("neutralizeYieldInRegion failed for op '",
-                   op->getName().getStringRef(), "' at operand #", i,
-                   " in scope ", scopeType);
+      unsigned numCarriedOperands =
+          terminator->getNumOperands() - carriedOperandOffset;
+      for (unsigned i = 0; i < numCarriedOperands; ++i) {
+        if (failed(neutralizeCarriedTerminatorOperand(
+                op, info, scopeType, loc, i,
+                terminator->getOpOperand(i + carriedOperandOffset)))) {
           return failure();
         }
-        yieldOp.setOperand(i, replacement);
       }
     }
   }
@@ -789,6 +812,15 @@ static LogicalResult neutralizeTerminatorUses(Operation *op,
 static LogicalResult executeActions(SmallVector<PendingAction> &actions,
                                     StringRef scopeType);
 
+static Operation *findLiveResultUser(Operation *op, StringRef scopeType) {
+  for (Value result : op->getResults()) {
+    if (Operation *user = findLiveUser(result, scopeType)) {
+      return user;
+    }
+  }
+  return nullptr;
+}
+
 static LogicalResult normalizeRegionOp(Operation *op, StringRef scopeType) {
   auto infoOpt = parseCoreTypeInfo(op);
   if (!infoOpt) {
@@ -807,10 +839,24 @@ static LogicalResult normalizeRegionOp(Operation *op, StringRef scopeType) {
 
   debugDumpOperation("before normalizeRegionOp", op);
 
+  // Keep a loop intact when its body has no op for this scope but its results
+  // are still consumed here; neutralizing carried values would change them.
+  if (isa<scf::ForOp, scf::WhileOp>(op) &&
+      !controlFlowOpHasScopeContent(op, scopeType)) {
+    if (Operation *resultUser = findLiveResultUser(op, scopeType)) {
+      logDebug("preserving complete loop '", op->getName().getStringRef(),
+               "' in scope ", scopeType,
+               " because its result still has live user '",
+               resultUser->getName().getStringRef(), "'");
+      return success();
+    }
+  }
+
   if (op->getNumRegions() > 0) {
-    if (failed(neutralizeYieldInRegion(op, info, scopeType, loc))) {
-      logDebug("normalizeRegionOp failed while neutralizing yields for op '",
-               op->getName().getStringRef(), "' in scope ", scopeType);
+    if (failed(neutralizeRegionTerminators(op, info, scopeType, loc))) {
+      logDebug(
+          "normalizeRegionOp failed while neutralizing terminators for op '",
+          op->getName().getStringRef(), "' in scope ", scopeType);
       return failure();
     }
 
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_canonicalize_ut.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_canonicalize_ut.mlir
index e705b641eb..e0191e1cff 100644
--- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_canonicalize_ut.mlir
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_canonicalize_ut.mlir
@@ -661,7 +661,9 @@ module {
 // CHECK-NEXT:     %[[SELC:.*]] = arith.select
 // CHECK-NEXT:     scf.condition(%[[SELC]])
 // CHECK-NEXT:   } do {
+// CHECK-NEXT:   ^bb0
 // CHECK-NEXT:     memref.store
+// CHECK-NEXT:     arith.addi
 // CHECK-NEXT:     scf.yield
 // CHECK-NEXT:   }
 // CHECK-NEXT:   scope.return
@@ -842,3 +844,69 @@ module {
     func.return
   }
 }
+
+// -----
+
+// A VECTOR for can initialize a VECTOR while whose scalar results are consumed
+// directly by CUBE operations. The CUBE clone must retain both complete state
+// transitions; keeping either loop shell with neutral yields changes the live
+// while results.
+// CHECK-LABEL: func.func @vector_while_results_drive_cube(
+// CHECK:      scope.scope : () -> () {
+// CHECK:        %[[INIT:.*]]:2 = scf.for
+// CHECK:          arith.index_cast
+// CHECK-NEXT:     arith.addi
+// CHECK-NEXT:     arith.addi
+// CHECK-NEXT:     scf.yield
+// CHECK:        %[[WHILE:.*]]:2 = scf.while
+// CHECK:          arith.cmpi
+// CHECK-NEXT:     scf.condition
+// CHECK:        } do {
+// CHECK:          arith.addi
+// CHECK-NEXT:     arith.addi
+// CHECK-NEXT:     scf.yield
+// CHECK:        arith.muli %[[WHILE]]#0, %[[WHILE]]#1
+// CHECK-NEXT:   memref.store
+// CHECK-NEXT:   scope.return
+// CHECK-NEXT: } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<CUBE>}
+module {
+  func.func @vector_while_results_drive_cube(
+      %index_init: i64, %offset_init: i64, %limit: i64,
+      %out: memref<1xi64>) {
+    %idx = arith.constant {ssbuffer.core_type = "CUBE"} 0 : index
+    %c0 = arith.constant {ssbuffer.core_type = "VECTOR"} 0 : index
+    %c2 = arith.constant {ssbuffer.core_type = "VECTOR"} 2 : index
+    %c1 = arith.constant {ssbuffer.core_type = "VECTOR"} 1 : index
+    %step = arith.constant {ssbuffer.core_type = "VECTOR"} 1 : i64
+    %0:2 = scf.for %i = %c0 to %c2 step %c1
+        iter_args(%index = %index_init, %offset = %offset_init)
+        -> (i64, i64) {
+      %i64 = arith.index_cast %i {ssbuffer.core_type = "VECTOR"}
+          : index to i64
+      %next_index = arith.addi %index, %i64
+          {ssbuffer.core_type = "VECTOR"} : i64
+      %next_offset = arith.addi %offset, %next_index
+          {ssbuffer.core_type = "VECTOR"} : i64
+      scf.yield {ssbuffer.core_type = "VECTOR, VECTOR"}
+          %next_index, %next_offset : i64, i64
+    } {ssbuffer.core_type = "VECTOR, VECTOR"}
+    %1:2 = scf.while (%index = %0#0, %offset = %0#1)
+        : (i64, i64) -> (i64, i64) {
+      %continue = arith.cmpi slt, %index, %limit
+          {ssbuffer.core_type = "VECTOR"} : i64
+      scf.condition(%continue) %index, %offset : i64, i64
+    } do {
+    ^bb0(%index: i64, %offset: i64):
+      %next_index = arith.addi %index, %step
+          {ssbuffer.core_type = "VECTOR"} : i64
+      %next_offset = arith.addi %offset, %next_index
+          {ssbuffer.core_type = "VECTOR"} : i64
+      scf.yield {ssbuffer.core_type = "VECTOR, VECTOR"}
+          %next_index, %next_offset : i64, i64
+    } attributes {ssbuffer.core_type = "VECTOR, VECTOR"}
+    %product = arith.muli %1#0, %1#1 {ssbuffer.core_type = "CUBE"} : i64
+    memref.store %product, %out[%idx] {ssbuffer.core_type = "CUBE"}
+        : memref<1xi64>
+    func.return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_ut.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_ut.mlir
index 8cbd1b1f2d..4305a2ec2e 100644
--- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_ut.mlir
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/separate_cv_scope_ut.mlir
@@ -52,4 +52,49 @@ module {
     memref.store %1#2, %outc[%idxc] {ssbuffer.core_type = "CUBE"} : memref<1xi32>
     func.return
   }
+
+  // A while induction variable can belong to VECTOR while still controlling
+  // the trip count of CUBE work. Both separated scopes must preserve the
+  // predicate-carried value and its update.
+
+  // CHECK-LABEL: func.func @while_predicate_controls_mixed_scope(
+  // CHECK-SAME: %[[UB:.*]]: i32
+  // CHECK: scope.scope : () -> () {
+  // CHECK: scf.while (%[[V_IV:.*]] = %{{.*}}) : (i32) -> i32 {
+  // CHECK: %[[V_COND:.*]] = arith.cmpi slt, %[[V_IV]], %[[UB]]
+  // CHECK: scf.condition(%[[V_COND]]) %[[V_IV]] : i32
+  // CHECK: ^bb0(%[[V_BODY_IV:.*]]: i32):
+  // CHECK: %[[V_NEXT:.*]] = arith.addi %[[V_BODY_IV]], %{{.*}}
+  // CHECK: scf.yield %[[V_NEXT]] : i32
+  // CHECK: } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+  // CHECK: scope.scope : () -> () {
+  // CHECK: scf.while (%[[C_IV:.*]] = %{{.*}}) : (i32) -> i32 {
+  // CHECK: %[[C_COND:.*]] = arith.cmpi slt, %[[C_IV]], %[[UB]]
+  // CHECK: scf.condition(%[[C_COND]]) %[[C_IV]] : i32
+  // CHECK: ^bb0(%[[C_BODY_IV:.*]]: i32):
+  // CHECK: memref.store
+  // CHECK: %[[C_NEXT:.*]] = arith.addi %[[C_BODY_IV]], %{{.*}}
+  // CHECK: scf.yield %[[C_NEXT]] : i32
+  // CHECK: } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<CUBE>}
+  func.func @while_predicate_controls_mixed_scope(
+      %ub: i32, %outv: memref<1xi32>, %outc: memref<1xi32>) {
+    %idxv = arith.constant {ssbuffer.core_type = "VECTOR"} 0 : index
+    %idxc = arith.constant {ssbuffer.core_type = "CUBE"} 0 : index
+    %c0 = arith.constant {ssbuffer.core_type = "VECTOR"} 0 : i32
+    %c1 = arith.constant {ssbuffer.core_type = "VECTOR"} 1 : i32
+    %cube_value = arith.constant {ssbuffer.core_type = "CUBE"} 7 : i32
+
+    %result = scf.while (%iv = %c0) : (i32) -> i32 {
+      %cond = arith.cmpi slt, %iv, %ub {ssbuffer.core_type = "VECTOR"} : i32
+      scf.condition(%cond) {ssbuffer.core_type = "VECTOR"} %iv : i32
+    } do {
+    ^bb0(%iv: i32):
+      memref.store %cube_value, %outc[%idxc] {ssbuffer.core_type = "CUBE"} : memref<1xi32>
+      %next = arith.addi %iv, %c1 {ssbuffer.core_type = "VECTOR"} : i32
+      scf.yield {ssbuffer.core_type = "VECTOR"} %next : i32
+    } attributes {ssbuffer.core_type = "VECTOR"}
+
+    memref.store %result, %outv[%idxv] {ssbuffer.core_type = "VECTOR"} : memref<1xi32>
+    func.return
+  }
 }

From b00c0885694da0621e7350d94468f58ddcd69275 Mon Sep 17 00:00:00 2001
From: sxm <sunxianmiao@h-partners.com>
Date: Fri, 31 Jul 2026 17:13:38 +0800
Subject: [PATCH 05/11] [ssbuffer](fix) alias through to_tensor and fix
 MayImplicitTranspose

---
 .../Common/MemoryEffectsTracker.cpp           | 32 ++++++++++++++++---
 .../PlanComputeBlock/OpClassifier.cpp         |  5 ++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp
index a240134523..676dd5f058 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp
@@ -50,6 +50,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
@@ -79,9 +80,22 @@ bool isDefinedInside(Value v, Operation *op) {
   return op->isProperAncestor(defOp);
 }
 
+static Value getAliasSource(Value val) {
+  auto *op = val.getDefiningOp();
+  if (!op) {
+    return nullptr;
+  }
+  return llvm::TypeSwitch<Operation *, Value>(op)
+      .Case([](ViewLikeOpInterface viewOp) { return viewOp.getViewSource(); })
+      .Case([](bufferization::ToTensorOp totensorOp) {
+        return totensorOp.getBuffer();
+      })
+      .Default([](auto) { return nullptr; });
+}
+
 Value getViewSource(Value val) {
-  while (auto viewLike = val.getDefiningOp<ViewLikeOpInterface>()) {
-    val = viewLike.getViewSource();
+  while (auto source = getAliasSource(val)) {
+    val = source;
   }
   return val;
 }
@@ -332,15 +346,25 @@ MemoryDependenceGraph::collectOuterEffects(Operation *op, bool &unknown,
 }
 
 AliasResult MemoryDependenceGraph::queryAlias(Value lhs, Value rhs) {
+  auto lhsSource = getViewSource(lhs);
+  auto rhsSource = getViewSource(rhs);
+  if (!rhsSource) {
+    rhsSource = rhs;
+  }
   auto isFuncEntryArg = [](const Value &val) -> bool {
     auto arg = llvm::dyn_cast<BlockArgument>(val);
-    return arg && arg.getOwner()->isEntryBlock();
+    if (!arg) {
+      return false;
+    }
+    auto *block = arg.getOwner();
+    return block->isEntryBlock() &&
+           llvm::isa<func::FuncOp>(block->getParentOp());
   };
   if (isFuncEntryArg(getViewSource(lhs)) &&
       isFuncEntryArg(getViewSource(rhs))) {
     return lhs == rhs ? AliasResult::MustAlias : AliasResult::NoAlias;
   }
-  return aa.alias(lhs, rhs);
+  return aa.alias(lhsSource, rhsSource);
 }
 
 SmallVector<MemoryDependenceGraph::MemSlot *>
diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
index 5b52de064b..2be5415ff3 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
@@ -41,6 +41,7 @@
 #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h"
 
 #include "bishengir/Dialect/Annotation/IR/Annotation.h"
+#include "bishengir/Dialect/HIVM/Utils/Utils.h"
 #include "bishengir/Dialect/Utils/Util.h"
 
 using namespace mlir;
@@ -247,7 +248,9 @@ void OpClassifierPass::matchTransposePattern(Operation *def) {
   // Check input tensor
   auto operands = transposeOp->getOperands();
   for (const auto &op : operands) {
-    if (shouldMarkCubeSeed(op.getDefiningOp())) {
+    if (shouldMarkCubeSeed(op.getDefiningOp()) &&
+        !utils::getAnnotateOpWithAttr(
+            op, hivm::kMayImplicitTransposeWithLastAxis)) {
       markCube(op.getDefiningOp());
       cubeSeeds.push_back(op.getDefiningOp());
       break; // No need to check other operands, one is enough to seed the

From 2afca011b3a00ba0643b4a63ed9f778cd0afca66 Mon Sep 17 00:00:00 2001
From: sxm <sunxianmiao@h-partners.com>
Date: Wed, 12 Aug 2026 14:17:55 +0800
Subject: [PATCH 06/11] [ssbuffer](feat) AnalyzeScope support whileOp

---
 .../AnalyzeDataFlow/AnalyzeName.cpp           |  3 +-
 .../AnalyzeDataFlow/AnalyzeScope.cpp          | 71 ++++++++++---------
 .../PreCheckAvailable/PreCheckBlacklist.cpp   |  1 -
 .../PreCheckAvailable/test_have_whileop.mlir  | 28 --------
 .../test_pcb13_mlir_while_loop_c_nonzero.py   |  4 +-
 ...cb14_mlir_while_loop_c_broadcast_scalar.py |  4 +-
 6 files changed, 42 insertions(+), 69 deletions(-)
 delete mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PreCheckAvailable/test_have_whileop.mlir

diff --git a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp
index 80e407255c..57c8f6a7f0 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp
@@ -40,7 +40,8 @@ using namespace triton;
 
 namespace {
 
-static constexpr llvm::StringLiteral interceptrFunc[]{""};
+static constexpr llvm::StringLiteral interceptrFunc[]{
+    "_parallel_hstu_attn_bwd"};
 
 static LogicalResult verifyFuncNames(ModuleOp module) {
   bool intercepted = false;
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeScope.cpp
index c5c752d26e..031fdadde3 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeScope.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeScope.cpp
@@ -82,31 +82,31 @@ static bool checkTransferInteraction(mlir::Operation *op) {
 }
 
 static bool checkVecScopeMainLoop(ModuleOp module) {
-  bool hasMainLoopFor = false;
-  bool allForSatisfy = true;
+  bool hasMainLoop = false;
+  bool allMainLoopsSatisfy = true;
 
   module.walk([&](scope::ScopeOp scopeOp) -> WalkResult {
     if (!isVectorScope(scopeOp)) {
       return WalkResult::advance();
     }
 
-    scopeOp.walk([&](scf::ForOp forOp) -> WalkResult {
-      if (!forOp->hasAttr(CVPipeline::kMainLoop)) {
+    scopeOp.walk([&](Operation *op) -> WalkResult {
+      if (!isMainLoopOp(op)) {
         return WalkResult::advance();
       }
 
-      hasMainLoopFor = true;
+      hasMainLoop = true;
       bool hasCVInteraction = false;
 
-      forOp.walk([&](mlir::Operation *op) -> WalkResult {
-        if (op == forOp) {
+      op->walk([&](mlir::Operation *innerOp) -> WalkResult {
+        if (innerOp == op) {
           return WalkResult::advance();
         }
 
         // ops with "ssbuffer.transfer_id" are injected in SplitDataflowPass for
         // data transfer
-        if (op->hasAttr(CVPipeline::kTransferId)) {
-          hasCVInteraction = checkTransferInteraction(op);
+        if (innerOp->hasAttr(CVPipeline::kTransferId)) {
+          hasCVInteraction = checkTransferInteraction(innerOp);
           if (hasCVInteraction) {
             return WalkResult::interrupt();
           }
@@ -115,42 +115,44 @@ static bool checkVecScopeMainLoop(ModuleOp module) {
         return WalkResult::advance();
       });
 
-      // As long as there is a forOp with "ssbuffer.main_loop" not a real
-      // mainloop, the processing conditions are not met, need to skip.
+      // As long as there is a for/while with "ssbuffer.main_loop" that is not a
+      // real mainloop, the processing conditions are not met, need to skip.
       if (!hasCVInteraction) {
-        allForSatisfy = false;
+        allMainLoopsSatisfy = false;
         return WalkResult::interrupt();
       }
 
       return WalkResult::advance();
     });
 
-    if (!allForSatisfy) {
+    if (!allMainLoopsSatisfy) {
       return WalkResult::interrupt();
     }
 
     return WalkResult::advance();
   });
 
-  return hasMainLoopFor && allForSatisfy;
+  return hasMainLoop && allMainLoopsSatisfy;
 }
 
-// For every main_loop id, gather all forOps sharing that id and count the
-// hivm.hir.copy and hivm.hir.fixpipe ops within them. Only when ALL main_loop
-// ids have either count equal to zero (every id has only copy or only
-// fixpipe, none has both), the dynamic CV pipeline cannot be applied and we
-// fall back to the original workflow.
+// For every main_loop id, gather all for/while ops sharing that id and count
+// the hivm.hir.copy and hivm.hir.fixpipe ops within them. Only when ALL
+// main_loop ids have either count equal to zero (every id has only copy or
+// only fixpipe, none has both), the dynamic CV pipeline cannot be applied and
+// we fall back to the original workflow.
 //   - hivm::CopyOp    typically appears in VECTOR scope main_loops
 //   - hivm::FixpipeOp typically appears in CUBE scope main_loops
-// Nested regions inside the main_loop forOp are also walked, and scf.yield
+// Nested regions inside the main_loop op are also walked, and scf.yield
 // terminators are skipped.
 static bool isMainLoopOnlyCopyOrFixpipe(ModuleOp module) {
   // main_loop id -> (countCopy, countFixpipe)
   llvm::DenseMap<int, std::pair<int, int>> idToCounts;
 
-  module.walk([&](scf::ForOp forOp) -> WalkResult {
-    auto mainLoopAttr =
-        forOp->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
+  module.walk([&](Operation *op) -> WalkResult {
+    if (!isa<scf::ForOp, scf::WhileOp>(op)) {
+      return WalkResult::advance();
+    }
+    auto mainLoopAttr = op->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
     if (!mainLoopAttr) {
       return WalkResult::advance();
     }
@@ -158,18 +160,18 @@ static bool isMainLoopOnlyCopyOrFixpipe(ModuleOp module) {
     int id = mainLoopAttr.getInt();
     auto &counts = idToCounts[id];
 
-    forOp.walk([&](mlir::Operation *op) -> WalkResult {
-      // Skip the forOp itself (the walk visits it first)
-      if (op == forOp) {
+    op->walk([&](mlir::Operation *innerOp) -> WalkResult {
+      // Skip the loop op itself (the walk visits it first)
+      if (innerOp == op) {
         return WalkResult::advance();
       }
       // Skip yield terminators (they are not real ops)
-      if (isa<scf::YieldOp>(op)) {
+      if (isa<scf::YieldOp>(innerOp)) {
         return WalkResult::advance();
       }
-      if (isa<hivm::CopyOp>(op)) {
+      if (isa<hivm::CopyOp>(innerOp)) {
         ++counts.first;
-      } else if (isa<hivm::FixpipeOp>(op)) {
+      } else if (isa<hivm::FixpipeOp>(innerOp)) {
         ++counts.second;
       }
       return WalkResult::advance();
@@ -195,15 +197,14 @@ static bool isMainLoopOnlyCopyOrFixpipe(ModuleOp module) {
 }
 
 static LogicalResult verifyMainLoop(ModuleOp module) {
-  // Only skip if ALL forOps lack main_loop attr
-  bool hasMainLoopForOp = false;
-  module.walk([&](scf::ForOp forOp) {
-    if (forOp->hasAttr("ssbuffer.main_loop")) {
-      hasMainLoopForOp = true;
+  bool hasMainLoopOp = false;
+  module.walk([&](Operation *op) {
+    if (isMainLoopOp(op)) {
+      hasMainLoopOp = true;
     }
   });
 
-  if (!hasMainLoopForOp) {
+  if (!hasMainLoopOp) {
     LDBG("[INFO]: No cycle of multiple iterations, the DynamicCVPipeline pass "
          "will be interrupted, and resumed to the original workflow.");
     CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_IGNORED);
diff --git a/third_party/ascend/lib/DynamicCVPipeline/PreCheckAvailable/PreCheckBlacklist.cpp b/third_party/ascend/lib/DynamicCVPipeline/PreCheckAvailable/PreCheckBlacklist.cpp
index 6daba7e900..5c45979576 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/PreCheckAvailable/PreCheckBlacklist.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/PreCheckAvailable/PreCheckBlacklist.cpp
@@ -34,7 +34,6 @@ using namespace triton;
 // The blacklist operations that should skip SSBUFFER
 static const llvm::SmallVector<llvm::StringRef> kBlacklistOpNames = {
     "scope.scope",
-    "scf.while",
 };
 
 static constexpr const char *DEBUG_TYPE = "pre-check-blacklist";
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PreCheckAvailable/test_have_whileop.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PreCheckAvailable/test_have_whileop.mlir
deleted file mode 100644
index a267fc2be4..0000000000
--- a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PreCheckAvailable/test_have_whileop.mlir
+++ /dev/null
@@ -1,28 +0,0 @@
-// RUN: triton-opt --pre-check-blacklist --debug %s 2>&1 | FileCheck %s
-
-// Test Case: Module with scf.while operation (blacklist op)
-// Should be detected by PreCheckBlacklistPass
-
-//CHECK: SSBUFFER will be skipped because scf.while operation was found
-module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
-  func.func @test_while_blacklist() {
-    %c0_i32 = arith.constant 0 : i32
-    %c1_i32 = arith.constant 1 : i32
-    %c32_i32 = arith.constant 32 : i32
-
-    %cst = arith.constant 0.000000e+00 : f32
-    %empty = tensor.empty() : tensor<32x32xf32>
-    %init = linalg.fill ins(%cst : f32) outs(%empty : tensor<32x32xf32>) -> tensor<32x32xf32>
-
-    %result:2 = scf.while (%arg0 = %init, %arg1 = %c0_i32) : (tensor<32x32xf32>, i32) -> (tensor<32x32xf32>, i32) {
-      %cond = arith.cmpi slt, %arg1, %c32_i32 : i32
-      scf.condition(%cond) %arg0, %arg1 : tensor<32x32xf32>, i32
-    } do {
-    ^bb0(%arg0: tensor<32x32xf32>, %arg1: i32):
-      %next_idx = arith.addi %arg1, %c1_i32 : i32
-      scf.yield %arg0, %next_idx : tensor<32x32xf32>, i32
-    }
-
-    return
-  }
-}
diff --git a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb13_mlir_while_loop_c_nonzero.py b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb13_mlir_while_loop_c_nonzero.py
index 0ca74a3094..e3f3dbd0e2 100644
--- a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb13_mlir_while_loop_c_nonzero.py
+++ b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb13_mlir_while_loop_c_nonzero.py
@@ -224,7 +224,7 @@ def test_pcb13_tc01():
     assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty"
     assert "func.func @pcb13_tc01_while_matmul_add(" in mlir, \
         "Kernel function definition not found in MLIR code"
-    assert "scope" not in mlir, "MLIR code does not contain the 'scope' keyword"
+    assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword"
 
     # Output MLIR code to the specified path
 
@@ -247,7 +247,7 @@ def test_pcb13_tc02():
     assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty"
     assert "func.func @pcb13_tc02_while_matmul_add(" in mlir, \
         "Kernel function definition not found in MLIR code"
-    assert "scope" not in mlir, "MLIR code does not contain the 'scope' keyword"
+    assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword"
 
     # Output MLIR code to the specified path
 
diff --git a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb14_mlir_while_loop_c_broadcast_scalar.py b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb14_mlir_while_loop_c_broadcast_scalar.py
index 31f556d8a0..2142315d2a 100644
--- a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb14_mlir_while_loop_c_broadcast_scalar.py
+++ b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_pcb14_mlir_while_loop_c_broadcast_scalar.py
@@ -223,7 +223,7 @@ def test_pcb14_tc01():
     assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty"
     assert "func.func @pcb14_tc01_while_matmul_scalar(" in mlir, \
         "Kernel function definition not found in MLIR code"
-    assert "scope" not in mlir, "Fallback scenario: MLIR code unexpectedly contains the 'scope' keyword"
+    assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword"
 
     # Output MLIR code to the specified path
 
@@ -246,7 +246,7 @@ def test_pcb14_tc02():
     assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty"
     assert "func.func @pcb14_tc02_while_matmul_scalar(" in mlir, \
         "Kernel function definition not found in MLIR code"
-    assert "scope" not in mlir, "Fallback scenario: MLIR code unexpectedly contains the 'scope' keyword"
+    assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword"
 
     # Output MLIR code to the specified path
 

From e00006bc6bc8227f05df65cc0520733f5b4e8aae Mon Sep 17 00:00:00 2001
From: 1801ZDL <1241272204@qq.com>
Date: Fri, 31 Jul 2026 16:56:43 +0800
Subject: [PATCH 07/11] [ssbuffer](feat) support whileOP in outerscope

Re-adapt OuterScope whileop support on top of cxt's MainLoop API:
- parentOpHasMainLoopAttr now uses CVPipeline::isMainLoopOp (handles whileOp)
- collectBufferAllocs identifies cross-core buffer from transfer op operands
- inject i32 iteration counter into main_loop whileOps (preInjectWhileOpToggles)
- prepareLoopPolling unifies polling control flow for forOp / whileOp
- add Outer-scope-whileop.mlir UT

Co-Authored-By: DeLong code
---
 .../AddMultiBufferOuterScope.cpp              | 339 +++++++++++++-----
 .../AllocMultiCache/Outer-scope-whileop.mlir  | 102 ++++++
 2 files changed, 352 insertions(+), 89 deletions(-)
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Outer-scope-whileop.mlir

diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp
index a2f166db3f..d791afaf80 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp
@@ -76,28 +76,13 @@ static bool isInVectorScope(Operation *op) {
 
 // --- main_loop attribute helpers ---
 
-/// Check if forOp (or its terminator) has ssbuffer.main_loop attribute
-static bool forOpHasMainLoopAttr(scf::ForOp forOp) {
-  if (forOp->hasAttr("ssbuffer.main_loop")) {
-    return true;
-  }
-  Operation *terminator = forOp.getBody()->getTerminator();
-  return terminator && terminator->hasAttr("ssbuffer.main_loop");
-}
-
-/// Check if a sync op's direct parent has ssbuffer.main_loop attribute
+/// Check if a sync op's direct parent is a main_loop op (forOp / whileOp
+/// carrying the ssbuffer.main_loop attribute)
 static bool parentOpHasMainLoopAttr(Operation *syncOp) {
   if (!syncOp) {
     return false;
   }
-  Operation *parent = syncOp->getParentOp();
-  if (!parent) {
-    return false;
-  }
-  if (auto forOp = dyn_cast<scf::ForOp>(parent)) {
-    return forOpHasMainLoopAttr(forOp);
-  }
-  return false;
+  return CVPipeline::isMainLoopOp(syncOp->getParentOp());
 }
 
 // --- Operation search helpers ---
@@ -196,37 +181,82 @@ collectOpsByTransferId(ModuleOp module,
   return 0;
 }
 
-/// Collect alloc/mark pairs (independent of block_id and main_loop)
+/// Collect alloc/mark pairs from transfer ops in the group.
+/// Identifies the correct cross-core buffer (ub/cbuf) used by each transfer op,
+/// ignoring local buffers (cc on CUBE side) that are not part of the data
+/// transfer.
 static int collectBufferAllocs(const SmallVector<Operation *> &ops,
-                               BufferAllocInfo &info) {
-  SmallVector<Operation *> allocs;
-  SmallVector<Operation *> marks;
+                               TransferGroupInfo &info) {
+  // Helper: find the annotation.mark for a given alloc op
+  auto findMarkForAlloc = [](Operation *allocOp) -> Operation * {
+    Value allocResult = allocOp->getResult(0);
+    for (auto *user : allocResult.getUsers()) {
+      if (isa<annotation::MarkOp>(user))
+        return user;
+    }
+    return nullptr;
+  };
 
-  for (Operation *op : ops) {
-    if (isa<memref::AllocOp>(op)) {
-      allocs.push_back(op);
-    } else if (isa<annotation::MarkOp>(op)) {
-      marks.push_back(op);
+  // Identify sender's cross-core buffer from transferOp's outs operand
+  if (info.senderChain.transferOp) {
+    Operation *transferOp = info.senderChain.transferOp;
+    // fixpipe / hir.copy: cross-core buffer is the last operand (outs)
+    Value crossCoreBuf =
+        transferOp->getOperand(transferOp->getNumOperands() - 1);
+    if (auto *defOp = crossCoreBuf.getDefiningOp()) {
+      if (isa<memref::AllocOp>(defOp)) {
+        info.senderBuf.allocOp = defOp;
+        info.senderBuf.markOp = findMarkForAlloc(defOp);
+        LDBG("Sender cross-core buffer: alloc from transferOp outs");
+      }
     }
   }
 
-  LDBG("collectBufferAllocs: allocs=" << allocs.size()
-                                      << ", marks=" << marks.size());
-
-  // Pair in order: sender first, receiver second
-  if (!allocs.empty()) {
-    info.sender.allocOp = allocs[0];
-  }
-  if (allocs.size() > 1) {
-    info.receiver.allocOp = allocs[1];
+  // Identify receiver's cross-core buffer from transferOp's input operand
+  if (info.receiverChain.transferOp) {
+    Operation *transferOp = info.receiverChain.transferOp;
+    // memref.memory_space_cast / hivm.convert_layout: cross-core buffer is
+    // the first operand
+    Value crossCoreBuf = transferOp->getOperand(0);
+    if (auto *defOp = crossCoreBuf.getDefiningOp()) {
+      if (isa<memref::AllocOp>(defOp)) {
+        info.receiverBuf.allocOp = defOp;
+        info.receiverBuf.markOp = findMarkForAlloc(defOp);
+        LDBG("Receiver cross-core buffer: alloc from transferOp input");
+      }
+    }
   }
-  if (!marks.empty()) {
-    info.sender.markOp = marks[0];
+
+  // Collect alloc/mark for the OTHER side if not yet found.
+  // Some transfer ops (e.g. fixpipe) have both a local input (cc) and a
+  // cross-core output (ub). The receiver side's buffer is the cross-core one.
+  // Walk all allocs in the group to find any remaining unassigned buffer.
+  SmallVector<Operation *> allocs;
+  for (Operation *op : ops) {
+    if (isa<memref::AllocOp>(op))
+      allocs.push_back(op);
   }
-  if (marks.size() > 1) {
-    info.receiver.markOp = marks[1];
+
+  // Fill missing side from remaining allocs (prefer allocs with marks)
+  for (auto *allocOp : allocs) {
+    if (allocOp == info.senderBuf.allocOp ||
+        allocOp == info.receiverBuf.allocOp)
+      continue;
+    Operation *mark = findMarkForAlloc(allocOp);
+    if (!info.senderBuf.allocOp) {
+      info.senderBuf.allocOp = allocOp;
+      info.senderBuf.markOp = mark;
+    } else if (!info.receiverBuf.allocOp) {
+      info.receiverBuf.allocOp = allocOp;
+      info.receiverBuf.markOp = mark;
+    }
   }
 
+  LDBG("Sender buffer: " << (info.senderBuf.allocOp ? "alloc" : "none") << " + "
+                         << (info.senderBuf.markOp ? "mark" : "none"));
+  LDBG("Receiver buffer: " << (info.receiverBuf.allocOp ? "alloc" : "none")
+                           << " + "
+                           << (info.receiverBuf.markOp ? "mark" : "none"));
   return 0;
 }
 
@@ -394,20 +424,7 @@ static int buildTransferGroupData(int tid, const SmallVector<Operation *> &ops,
 
   LDBG("Building group tid=" << tid << ", ops=" << ops.size());
 
-  // 1. Collect buffer alloc/mark pairs
-  BufferAllocInfo bufInfo;
-  if (collectBufferAllocs(ops, bufInfo)) {
-    return -1;
-  }
-  info.senderBuf = bufInfo.sender;
-  info.receiverBuf = bufInfo.receiver;
-  LDBG("Sender buffer: " << (info.senderBuf.allocOp ? "alloc" : "none") << " + "
-                         << (info.senderBuf.markOp ? "mark" : "none"));
-  LDBG("Receiver buffer: " << (info.receiverBuf.allocOp ? "alloc" : "none")
-                           << " + "
-                           << (info.receiverBuf.markOp ? "mark" : "none"));
-
-  // 2. Determine original flag
+  // 1. Determine original flag
   for (Operation *op : ops) {
     if ((isa<hivm::SyncBlockSetOp>(op) || isa<hivm::SyncBlockWaitOp>(op))) {
       int f = getFlagFromSyncOp(op);
@@ -418,7 +435,7 @@ static int buildTransferGroupData(int tid, const SmallVector<Operation *> &ops,
     }
   }
 
-  // 3. Collect extra sync (parent has no main_loop)
+  // 2. Collect extra sync (parent has no main_loop)
   ExtraSyncInfo extraInfo;
   if (collectExtraSync(ops, info.originalFlag, extraInfo)) {
     return -1;
@@ -433,7 +450,7 @@ static int buildTransferGroupData(int tid, const SmallVector<Operation *> &ops,
     LDBG("Extra sync: not found");
   }
 
-  // 4. Collect transfer chain (parent has main_loop)
+  // 3. Collect transfer chain (parent has main_loop)
   TransferChainInfo chainInfo;
   if (collectTransferChains(ops, info.originalFlag, chainInfo)) {
     return -1;
@@ -441,7 +458,7 @@ static int buildTransferGroupData(int tid, const SmallVector<Operation *> &ops,
   info.senderChain = chainInfo.sender;
   info.receiverChain = chainInfo.receiver;
 
-  // 5. Determine direction
+  // 4. Determine direction
   if (info.senderChain.transferOp) {
     if (isa<hivm::FixpipeOp>(info.senderChain.transferOp)) {
       info.isCtoV = true;
@@ -450,10 +467,12 @@ static int buildTransferGroupData(int tid, const SmallVector<Operation *> &ops,
     }
   }
 
-  // For C→V transfer, sender uses receiver's buffer (the second alloc)
-  if (info.isCtoV && info.senderBuf.allocOp && info.receiverBuf.allocOp) {
-    LDBG("C→V transfer: swapping sender/receiver buffers");
-    std::swap(info.senderBuf, info.receiverBuf);
+  // 5. Collect buffer alloc/mark pairs from transfer ops
+  //    Must run after transfer chain collection to identify the correct
+  //    cross-core buffer (ub/cbuf) from each transfer op's operands,
+  //    ignoring local buffers (e.g. cc on CUBE side).
+  if (collectBufferAllocs(ops, info)) {
+    return -1;
   }
 
   // 6. Acquire output flag
@@ -741,6 +760,87 @@ static int setSsbufferTags(Operation *op, OpBuilder &builder, int blockId,
   return 0;
 }
 
+/// Ensure a WhileOp has an i32 iteration counter loop-carried variable.
+/// Returns the counter Value; polling condition is (counter % 2) == 0.
+/// Reuses an existing counter (e.g. one injected by InnerScope, detected via
+/// ssbuffer.iterCounter); injects a new one only when absent.
+static Value ensureWhileOpHasCounter(scf::WhileOp whileOp) {
+  if (whileOp->hasAttr(CVPipeline::kIterCounter)) {
+    Block &after = whileOp.getAfter().front();
+    return after.getArgument(after.getNumArguments() - 1);
+  }
+
+  OpBuilder builder(whileOp);
+  Location loc = whileOp.getLoc();
+  auto oldWhile = whileOp;
+  Type i32Type = builder.getI32Type();
+
+  // Init counter = 0
+  Value zero = builder.create<arith::ConstantIntOp>(loc, 0, 32);
+
+  SmallVector<Value> newInits(oldWhile.getInits());
+  newInits.push_back(zero);
+  SmallVector<Type> newResultTypes(oldWhile.getResultTypes());
+  newResultTypes.push_back(i32Type);
+
+  Value counterIterArg;
+
+  // Rebuild via the Builder callback API (matching InnerScope's
+  // setupWhileIterArgCounter)
+  auto newWhile = builder.create<scf::WhileOp>(
+      loc, newResultTypes, newInits,
+      [&](OpBuilder &bb, Location bl, ValueRange iterArgs) {
+        Block *oldBefore = oldWhile.getBeforeBody();
+        unsigned n = oldBefore->getNumArguments();
+        IRMapping map;
+        for (unsigned i = 0; i < n; ++i)
+          map.map(oldBefore->getArgument(i), iterArgs[i]);
+
+        for (Operation &op : oldBefore->without_terminator())
+          bb.clone(op, map);
+
+        auto oldCond = cast<scf::ConditionOp>(oldBefore->getTerminator());
+        SmallVector<Value> condArgs;
+        for (Value a : oldCond.getArgs())
+          condArgs.push_back(map.lookupOrDefault(a));
+        condArgs.push_back(iterArgs[n]); // counter
+        bb.create<scf::ConditionOp>(
+            bl, map.lookupOrDefault(oldCond.getCondition()), condArgs);
+      },
+      [&](OpBuilder &ab, Location al, ValueRange iterArgs) {
+        Block *oldAfter = oldWhile.getAfterBody();
+        unsigned n = oldAfter->getNumArguments();
+        counterIterArg = iterArgs[n];
+        IRMapping map;
+        for (unsigned i = 0; i < n; ++i)
+          map.map(oldAfter->getArgument(i), iterArgs[i]);
+
+        for (Operation &op : oldAfter->without_terminator())
+          ab.clone(op, map);
+
+        auto oldYield = cast<scf::YieldOp>(oldAfter->getTerminator());
+        Value one = ab.create<arith::ConstantIntOp>(al, 1, 32);
+        Value nextCounter = ab.create<arith::AddIOp>(al, counterIterArg, one);
+        SmallVector<Value> yOps;
+        for (Value v : oldYield.getOperands())
+          yOps.push_back(map.lookupOrDefault(v));
+        yOps.push_back(nextCounter);
+        ab.create<scf::YieldOp>(al, yOps);
+      });
+
+  // Copy attrs (must include ssbuffer.main_loop) and mark as processed
+  for (auto attr : oldWhile->getAttrs())
+    newWhile->setAttr(attr.getName(), attr.getValue());
+  newWhile->setAttr(CVPipeline::kIterCounter, builder.getUnitAttr());
+
+  // Replace results (exclude counter result)
+  for (unsigned i = 0, e = oldWhile.getNumResults(); i < e; ++i)
+    oldWhile.getResult(i).replaceAllUsesWith(newWhile.getResult(i));
+  oldWhile.erase();
+
+  return counterIterArg;
+}
+
 /// Create polling condition: (iter / step) % 2 == 0 (true=input, false=output)
 static Value createPollingCondition(scf::ForOp forOp, OpBuilder &builder,
                                     int blockId, int tid) {
@@ -1108,24 +1208,53 @@ static int processTransferChain(TransferOpChain &chain, Value cond,
   return 0;
 }
 
+/// Create polling condition and builder for a loop op (ForOp or WhileOp).
+/// Returns the condition Value; `builderOut` is set to the insertion point
+/// for subsequent wrapping ops (before the loop terminator).
+static Value prepareLoopPolling(Operation *loopOp, Operation *waitOp,
+                                OpBuilder &builderOut) {
+  int bid = getBlockId(waitOp);
+  int tid = getTransferId(waitOp);
+
+  if (auto forOp = dyn_cast<scf::ForOp>(loopOp)) {
+    OpBuilder condBuilder(forOp.getBody(), Block::iterator(waitOp));
+    Value cond = createPollingCondition(forOp, condBuilder, bid, tid);
+    builderOut.setInsertionPoint(forOp.getBody()->getTerminator());
+    return cond;
+  }
+
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp)) {
+    // Counter was already injected in preprocessing. Polling condition:
+    // (counter % 2) == 0
+    Block &after = whileOp.getAfter().front();
+    Value counter = after.getArgument(after.getNumArguments() - 1);
+    builderOut.setInsertionPoint(after.getTerminator());
+    OpBuilder condBuilder(builderOut);
+    Value c2 =
+        condBuilder.create<arith::ConstantIntOp>(whileOp.getLoc(), 2, 32);
+    Value rem =
+        condBuilder.create<arith::RemSIOp>(whileOp.getLoc(), counter, c2);
+    Value c0 =
+        condBuilder.create<arith::ConstantIntOp>(whileOp.getLoc(), 0, 32);
+    return condBuilder.create<arith::CmpIOp>(whileOp.getLoc(),
+                                             arith::CmpIPredicate::eq, rem, c0);
+  }
+
+  llvm_unreachable("unexpected loop op type");
+}
+
 /// Add polling control flow for all transfer groups
 static int addPollingControlFlow(DenseMap<int, TransferGroupInfo> &groups) {
   for (auto &p : groups) {
     TransferGroupInfo &g = p.second;
 
-    // Get sender's scf.for
+    // Get sender's loop op (ForOp or WhileOp)
     Operation *senderWaitParent = g.senderChain.waitOp->getParentOp();
-    scf::ForOp senderForOp = cast<scf::ForOp>(senderWaitParent);
-
-    int senderBid = getBlockId(g.senderChain.waitOp);
-    int senderTid = getTransferId(g.senderChain.waitOp);
 
-    // Insert polling condition at sender waitOp's position
-    OpBuilder senderCondBuilderForInsert(senderForOp.getBody(),
-                                         Block::iterator(g.senderChain.waitOp));
-    Value senderCond = createPollingCondition(
-        senderForOp, senderCondBuilderForInsert, senderBid, senderTid);
-    OpBuilder senderBuilder(senderForOp.getBody()->getTerminator());
+    // Prepare polling condition and builder for sender loop
+    OpBuilder senderBuilder(senderWaitParent->getContext());
+    Value senderCond = prepareLoopPolling(senderWaitParent,
+                                          g.senderChain.waitOp, senderBuilder);
 
     // Process sender chain (isProducer=true)
     if (processTransferChain(g.senderChain, senderCond, g.senderInputBuffer,
@@ -1134,28 +1263,22 @@ static int addPollingControlFlow(DenseMap<int, TransferGroupInfo> &groups) {
       return -1;
     }
 
-    // Process receiver chain (may use different scf.for) (isProducer=false)
+    // Process receiver chain (may use different loop op) (isProducer=false)
     if (g.receiverChain.waitOp) {
       Operation *receiverWaitParent = g.receiverChain.waitOp->getParentOp();
 
       if (receiverWaitParent == senderWaitParent) {
-        // Use the same cond
+        // Use the same cond and builder
         if (processTransferChain(g.receiverChain, senderCond,
                                  g.receiverInputBuffer, g.receiverOutputBuffer,
                                  g.outputFlag, false, senderBuilder) != 0) {
           return -1;
         }
       } else {
-        // Receiver uses a different scf.for, create new cond
-        scf::ForOp receiverForOp = cast<scf::ForOp>(receiverWaitParent);
-        int receiverBid = getBlockId(g.receiverChain.waitOp);
-        int receiverTid = getTransferId(g.receiverChain.waitOp);
-        OpBuilder receiverCondBuilderForInsert(
-            receiverForOp.getBody(), Block::iterator(g.receiverChain.waitOp));
-        Value receiverCond =
-            createPollingCondition(receiverForOp, receiverCondBuilderForInsert,
-                                   receiverBid, receiverTid);
-        OpBuilder receiverBuilder(receiverForOp.getBody()->getTerminator());
+        // Receiver uses a different loop op, prepare new cond and builder
+        OpBuilder receiverBuilder(receiverWaitParent->getContext());
+        Value receiverCond = prepareLoopPolling(
+            receiverWaitParent, g.receiverChain.waitOp, receiverBuilder);
         if (processTransferChain(g.receiverChain, receiverCond,
                                  g.receiverInputBuffer, g.receiverOutputBuffer,
                                  g.outputFlag, false, receiverBuilder) != 0) {
@@ -1167,6 +1290,36 @@ static int addPollingControlFlow(DenseMap<int, TransferGroupInfo> &groups) {
   return 0;
 }
 
+// ============================================================================
+// Preprocessing: inject iteration counter into WhileOps with main_loop
+// ============================================================================
+
+/// Inject an i32 iteration counter loop-carried variable into every WhileOp
+/// that has main_loop and contains transfer_id ops. Must run BEFORE Step 1 so
+/// subsequent data collection sees the already-modified IR.
+static void preInjectWhileOpToggles(ModuleOp module) {
+  SmallVector<scf::WhileOp> whileOps;
+  module.walk([&](scf::WhileOp whileOp) {
+    if (!CVPipeline::isMainLoopOp(whileOp))
+      return;
+    bool hasTransferOps = false;
+    whileOp.walk([&](Operation *op) {
+      if (op->hasAttr(mlir::CVPipeline::kTransferId)) {
+        hasTransferOps = true;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+    if (hasTransferOps)
+      whileOps.push_back(whileOp);
+  });
+
+  for (auto whileOp : whileOps)
+    ensureWhileOpHasCounter(whileOp);
+
+  LDBG("Preprocessed " << whileOps.size() << " WhileOps with toggle injection");
+}
+
 // ============================================================================
 // Pass entry point
 // ============================================================================
@@ -1182,6 +1335,19 @@ void AddMultiBufferOuterScopePass::runOnOperation() {
   LDBG("[AddMultiBufferOuterScope] ENTER");
   LDBG("============================================================");
 
+  // Determine buffer mode early; only inject toggle for double-buffer
+  int interCoreBufNum = BufferCountManager(module).getBufferCountByType(
+      BufferCountManager::DepType::InterCore);
+  bool isDoubleBuf = (interCoreBufNum > 1);
+  LDBG("[BufferCount] interCoreBufNum=" << interCoreBufNum
+                                        << " doubleBuf=" << isDoubleBuf);
+
+  // Preprocessing: inject iteration counter into WhileOps before data
+  // collection (only needed for double-buffer polling)
+  if (isDoubleBuf) {
+    preInjectWhileOpToggles(module);
+  }
+
   // Step 1: Collect transfer group information
   LDBG("[Step 1/3] Start: transfer group collection");
   FlagIdManager flagIdMgr(module);
@@ -1195,11 +1361,6 @@ void AddMultiBufferOuterScopePass::runOnOperation() {
   }
   LDBG("[Step 1/3] Done: " << groups.size() << " transfer groups");
 
-  int interCoreBufNum = BufferCountManager(module).getBufferCountByType(
-      BufferCountManager::DepType::InterCore);
-  bool isDoubleBuf = (interCoreBufNum > 1);
-  LDBG("[BufferCount] interCoreBufNum=" << interCoreBufNum
-                                        << " doubleBuf=" << isDoubleBuf);
   if (isDoubleBuf) {
     // Tag llvm.load/store volatile ops with crossDeps
     DenseMap<int, SmallVector<Operation *>> loadStoreByTid;
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Outer-scope-whileop.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Outer-scope-whileop.mlir
new file mode 100644
index 0000000000..844041d01b
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AllocMultiCache/Outer-scope-whileop.mlir
@@ -0,0 +1,102 @@
+// RUN: triton-opt --add_multi_buffer_outer_scope %s | FileCheck %s
+
+// Test: C→V transfer with scf.while (main_loop) in single-buffer mode.
+// Pass must not crash; IR structure (whileOp, main_loop, TCB marks) preserved.
+
+// CHECK-LABEL: func.func @tc_while_ctov_sender
+// CHECK: scf.while
+// CHECK: ssbuffer.main_loop
+// CHECK: tightly_coupled_buffer
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+func.func @tc_while_ctov_sender() {
+  %c0_i32 = arith.constant 0 : i32
+  %c100_i32 = arith.constant 100 : i32
+  %c1_i32 = arith.constant 1 : i32
+  // --- VECTOR scope (receiver for C→V: memspace_cast reads from ub) ---
+  scope.scope : () -> () {
+    %buf_ub = memref.alloc() {ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    annotation.mark %buf_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    scf.while (%iter = %c0_i32) : (i32) -> i32 {
+      %cond = arith.cmpi slt, %iter, %c100_i32 : i32
+      scf.condition(%cond) %iter : i32
+    } do {
+    ^bb0(%iter: i32):
+      hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_FIX>, <PIPE_V>] flag = 3
+      %buf = memref.memory_space_cast %buf_ub {ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<ub>> to memref<128xf16>
+      %t = bufferization.to_tensor %buf restrict writable : memref<128xf16> to tensor<128xf16>
+      hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 3
+      %next_iter = arith.addi %iter, %c1_i32 : i32
+      scf.yield %next_iter : i32
+    } attributes {ssbuffer.main_loop = 1 : i64}
+    scope.return
+  } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+  hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 10 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_FIX>, <PIPE_V>] flag = 3
+  hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_V>, <PIPE_FIX>] flag = 3
+  // --- CUBE scope (sender for C→V: fixpipe writes to ub) ---
+  scope.scope : () -> () {
+    %buf_cc = memref.alloc() {ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<cc>>
+    %buf_ub = memref.alloc() {ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    annotation.mark %buf_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args() -> () : i32 {
+      hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_V>, <PIPE_FIX>] flag = 3
+      hivm.hir.fixpipe {ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32} ins(%buf_cc : memref<128xf16, #hivm.address_space<cc>>) outs(%buf_ub : memref<128xf16, #hivm.address_space<ub>>)
+      hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 20 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_FIX>, <PIPE_V>] flag = 3
+      scf.yield
+    } {ssbuffer.main_loop = 1 : i64}
+    scope.return
+  } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+  return
+}
+
+
+// Test: Both sender and receiver use scf.while with main_loop
+// CHECK-LABEL: func.func @tc_while_both_sides
+// CHECK: scf.while
+// CHECK: ssbuffer.main_loop
+
+func.func @tc_while_both_sides() {
+  %c0_i32 = arith.constant 0 : i32
+  %c100_i32 = arith.constant 100 : i32
+  %c1_i32 = arith.constant 1 : i32
+  // --- VECTOR scope ---
+  scope.scope : () -> () {
+    %buf_ub = memref.alloc() {ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    annotation.mark %buf_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    scf.while (%iter = %c0_i32) : (i32) -> i32 {
+      %cond = arith.cmpi slt, %iter, %c100_i32 : i32
+      scf.condition(%cond) %iter : i32
+    } do {
+    ^bb0(%iter: i32):
+      hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32}[<VECTOR>, <PIPE_FIX>, <PIPE_V>] flag = 6
+      %buf = memref.memory_space_cast %buf_ub {ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<ub>> to memref<128xf16>
+      %t = bufferization.to_tensor %buf restrict writable : memref<128xf16> to tensor<128xf16>
+      hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 6
+      %next_iter = arith.addi %iter, %c1_i32 : i32
+      scf.yield %next_iter : i32
+    } attributes {ssbuffer.main_loop = 1 : i64}
+    scope.return
+  } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+  hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 30 : i32, ssbuffer.transfer_id = 2 : i32}[<VECTOR>, <PIPE_FIX>, <PIPE_V>] flag = 6
+  hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32}[<CUBE>, <PIPE_V>, <PIPE_FIX>] flag = 6
+  // --- CUBE scope ---
+  scope.scope : () -> () {
+    %buf_cc = memref.alloc() {ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<cc>>
+    %buf_ub = memref.alloc() {ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    annotation.mark %buf_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128xf16, #hivm.address_space<ub>>
+    scf.while (%iter = %c0_i32) : (i32) -> i32 {
+      %cond = arith.cmpi slt, %iter, %c100_i32 : i32
+      scf.condition(%cond) %iter : i32
+    } do {
+    ^bb0(%iter: i32):
+      hivm.hir.sync_block_wait {ssbuffer.analyze_flag_id, ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32}[<CUBE>, <PIPE_V>, <PIPE_FIX>] flag = 6
+      hivm.hir.fixpipe {ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32} ins(%buf_cc : memref<128xf16, #hivm.address_space<cc>>) outs(%buf_ub : memref<128xf16, #hivm.address_space<ub>>)
+      hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 40 : i32, ssbuffer.transfer_id = 2 : i32}[<CUBE>, <PIPE_FIX>, <PIPE_V>] flag = 6
+      %next_iter = arith.addi %iter, %c1_i32 : i32
+      scf.yield %next_iter : i32
+    } attributes {ssbuffer.main_loop = 1 : i64}
+    scope.return
+  } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+  return
+}
+}

From 5af6a84832d0a6e47fd4bc722bacf1d8633805f5 Mon Sep 17 00:00:00 2001
From: m-everglow <2276518549@qq.com>
Date: Mon, 6 Jul 2026 11:37:22 +0800
Subject: [PATCH 08/11] [ssbuffer](feat) support scf.while op in
 AddControlFlowCondition

---
 .../AddControlFlowCondition.h                 |  40 +-
 .../AddControlFlowCondition/CloneOps.h        |   5 +-
 .../AddControlFlowCondition/CreateIfOps.h     |   6 +-
 .../AddControlFlowCondition/ProcessArgs.h     |  57 +-
 .../{UpdateForOps.h => UpdateLoopOps.h}       |  20 +-
 .../AddControlFlowCondition/Utils.h           |  81 +-
 .../include/DynamicCVPipeline/Common/Utils.h  |   2 +
 .../AddControlFlowCondition.cpp               |  10 +-
 .../AddControlFlowCondition/CloneOps.cpp      | 255 +++---
 .../AddControlFlowCondition/CreateIfOps.cpp   |  85 +-
 .../InitDependentMap.cpp                      | 128 ++-
 .../AddControlFlowCondition/ProcessArgs.cpp   | 689 +++++++++++-----
 .../AddControlFlowCondition/UpdateForOps.cpp  | 641 ---------------
 .../UpdateLoopIterTimes.cpp                   |   7 +-
 .../AddControlFlowCondition/UpdateLoopOps.cpp | 740 ++++++++++++++++++
 .../AddControlFlowCondition/Utils.cpp         | 268 ++++++-
 .../MarkGMLoadPass.cpp                        |   2 +-
 .../while-clone-ops.mlir                      | 110 +++
 .../while-create-if-ops.mlir                  | 130 +++
 .../while-process-args.mlir                   | 117 +++
 .../while-update-loop-ops.mlir                | 146 ++++
 21 files changed, 2387 insertions(+), 1152 deletions(-)
 rename third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/{UpdateForOps.h => UpdateLoopOps.h} (80%)
 delete mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp
 create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-clone-ops.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-create-if-ops.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-process-args.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-loop-ops.mlir

diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
index 1a6011a2fb..9d5681c132 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
@@ -37,16 +37,15 @@ namespace triton {
 constexpr int CROSS_CORE_BUFFER_COUNT_THRESHOLD = 1;
 constexpr int INTRA_CORE_BUFFER_COUNT_THRESHOLD = 2;
 
-// Indicates the relationship between a tensor iter_arg and ssbuffer.if in the
-// main_loop
+// Relationship between a tensor iter_arg and ssbuffer.if in the main_loop
 struct TensorIterArgIfOpRelation {
   Value iterArg;
   scf::IfOp producer;
   llvm::SmallVector<scf::IfOp> consumers;
 };
 
-// Indicates the variables that need to be controlled when an ifOp is both a
-// producer and consumer of a tensor iter_args
+// Variables to control when an ifOp is both producer and consumer of a tensor
+// iter_args
 struct TensorIterArgIfOpVars {
   // The variables that need to be controlled as a producer
   llvm::SmallVector<Value> producerVars;
@@ -54,23 +53,29 @@ struct TensorIterArgIfOpVars {
   llvm::SmallVector<Value> consumerVars;
 };
 
+// Per scf.while block-arg map: whileOp -> block_id -> (new_arg_idx ->
+// old_arg_idx).
+using WhileBlockArgMap =
+    llvm::DenseMap<scf::WhileOp, llvm::DenseMap<int, llvm::DenseMap<int, int>>>;
+
 struct ControlFlowConditionInfo {
-  llvm::DenseMap<scf::ForOp, SmallVector<int>> blockCounters;
-  llvm::DenseMap<scf::ForOp, int> blockCounterNums;
-  llvm::DenseMap<scf::ForOp, SmallVector<int>> innerDepConds;
+  // Keys: main-loop op (scf.for/scf.while carrying ssbuffer.main_loop)
+  llvm::DenseMap<Operation *, SmallVector<int>> blockCounters;
+  llvm::DenseMap<Operation *, int> blockCounterNums;
+  llvm::DenseMap<Operation *, SmallVector<int>> innerDepConds;
 
   llvm::DenseMap<Operation *, SmallVector<Operation *>> crossCoreDependentMap;
-  llvm::DenseMap<scf::ForOp,
+  llvm::DenseMap<Operation *,
                  llvm::DenseMap<Operation *, SmallVector<Operation *>>>
       intraCoreDependentMap;
-  // Used to store the producer/consumer relationship between the tensor type
-  // iter_args in the main_loop and ssbuffer.if Note: vector index corresponds
-  // to iter arg index in the for op
-  llvm::DenseMap<scf::ForOp, llvm::SmallVector<TensorIterArgIfOpRelation>>
+  // Stores producer/consumer relationship between tensor iter_args in main_loop
+  // and ssbuffer.if; vector index corresponds to iter arg index in the
+  // main-loop op
+  llvm::DenseMap<Operation *, llvm::SmallVector<TensorIterArgIfOpRelation>>
       tensorIterArgDepsMap;
-  // Used to record the index of the control condition variable for the newly
-  // created iter_args for tensor iter_args
-  llvm::DenseMap<scf::ForOp, llvm::DenseMap<Value, SmallVector<int>>>
+  // Records control condition variable index for newly created iter_args of
+  // tensor iter_args
+  llvm::DenseMap<Operation *, llvm::DenseMap<Value, SmallVector<int>>>
       tensorIterArgIndicesMap;
 
   // unique counter value for each ifblock
@@ -83,6 +88,11 @@ struct ControlFlowConditionInfo {
   // Buffer counts for flowOpt condition
   int intraCoreBufferCount = 0;
   int crossCoreBufferCount = 0;
+
+  // Per scf.while (with main_loop attr): records per-block new iter_args
+  // mirroring iter_args used in scf.condition. Keys: whileOp -> block_id -> new
+  // iter_arg index. Value: original iter_arg index.
+  WhileBlockArgMap whileBlockArgMap;
 };
 
 class AddControlFlowConditionPass
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h
index d57203532d..baa187e1a0 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h
@@ -22,7 +22,6 @@
 
 #ifndef TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H
 #define TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Pass/Pass.h"
@@ -37,8 +36,8 @@ class CloneOpsPass : public PassWrapper<CloneOpsPass, OperationPass<ModuleOp>> {
   void runOnOperation() override;
 
   LogicalResult validateBlockIdsConsecutive(ModuleOp module);
-  LogicalResult cloneOpsInMainLoop(scf::ForOp forOp);
-  LogicalResult cleanupClonedOpsInMainLoop(scf::ForOp forOp);
+  LogicalResult cloneOpsInMainLoop(Operation *op);
+  LogicalResult cleanupClonedOpsInMainLoop(Operation *op);
   LogicalResult validateClonedOpsInVector(ModuleOp module);
 
   llvm::StringRef getArgument() const override { return "clone-ops"; }
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h
index 6a59ae9de3..54c1892948 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h
@@ -39,14 +39,16 @@ class CreateIfOpsPass
 
   void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; }
 
+  // `op` is the main-loop op (scf.for or scf.while) carrying
+  // ssbuffer.main_loop.
   LogicalResult computeYieldValues(
-      scf::ForOp forOp,
+      Operation *op,
       const llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
       llvm::DenseMap<int, SmallVector<Value>> &thenYieldValues,
       llvm::DenseMap<int, SmallVector<Value>> &elseYieldValues);
 
   LogicalResult createIfInMainLoop(
-      scf::ForOp forOp,
+      Operation *op,
       const llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
       const llvm::DenseMap<int, SmallVector<Value>> &thenYieldValues,
       const llvm::DenseMap<int, SmallVector<Value>> &elseYieldValues);
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.h
index 4ea08cc31d..44f072f029 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.h
@@ -27,15 +27,15 @@
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Pass/Pass.h"
 
+#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h"
+
 namespace mlir {
 namespace triton {
 
 struct ControlFlowConditionInfo;
 
-// For each shared iter_arg, we need to track:
-// - Which block_ids use it
-// - Who is the owner (first block_id in order)
-// - For each non-owner block, what new iter_arg index to use
+// For each shared iter_arg tracks: which block_ids use it, who is owner (first
+// block_id in order), and what new iter_arg index each non-owner block uses
 struct SharedArgInfo {
   int argIndex;
   Value iterArg;
@@ -48,6 +48,21 @@ struct SharedArgInfo {
         newArgIndex(newIdx), nonOwnerBlockId(nonOwner) {}
 };
 
+// Per-whileOp state for cloning cond-used iter_arg update chains into the new
+// scf.while's after body. Populated by planWhileIterArgDescriptors; consumed by
+// cloneWhileBlockChains, buildNewWhileYield, recordWhileBlockArgMap.
+struct WhileIterArgClonePlan {
+  // Per-origIdx metadata, keyed on the original iter_arg index.
+  llvm::DenseMap<unsigned, Operation *> compOp;
+  llvm::DenseMap<unsigned, llvm::DenseSet<Operation *>> chainOps;
+  llvm::DenseMap<unsigned, unsigned> posInClonedVec;
+  // (blockId, newArgIdx, origIdx) triples in planning order, one per (blockId,
+  // cond-used iter_arg) pair
+  SmallVector<std::tuple<int, unsigned, unsigned>> newArgDescriptors;
+  // Output: cloned compOp results per blockId, indexed by posInClonedVec
+  llvm::DenseMap<int, SmallVector<Value>> clonedPerBlock;
+};
+
 class ProcessArgsPass
     : public PassWrapper<ProcessArgsPass, OperationPass<ModuleOp>> {
 public:
@@ -57,11 +72,45 @@ class ProcessArgsPass
 
   LogicalResult processSharedIterArgs(ModuleOp module);
 
+  // Snapshots whileOp iter_args; clones cond-used update chain per block (same
+  // ssbuffer.block_id run); records (new_arg_idx, old_arg_idx) in
+  // ControlFlowConditionInfo.
+  LogicalResult updateIndependentCondsInWhileBlocks(ModuleOp module);
+
+  // Per-whileOp driver for updateIndependentCondsInWhileBlocks.
+  LogicalResult processWhileIterArgsInWhileOp(scf::WhileOp whileOp,
+                                              ControlFlowConditionInfo *info);
+
+  // Per-op driver for shared-iter_args processing.
+  LogicalResult processSharedIterArgsInLoop(Operation *op,
+                                            ControlFlowConditionInfo *info);
+
+  // Completes the scf.while path: migrate before/after bodies, rebuild
+  // yield/condition, transfer maps.
+  LogicalResult processSharedArgsInWhileOp(
+      scf::WhileOp whileOp, scf::WhileOp newWhileOp,
+      SmallVector<SharedArgInfo> &sharedArgsInfo,
+      const llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
+      const llvm::DenseMap<int, llvm::DenseSet<Operation *>>
+          &sharedArgToChainOps,
+      ControlFlowConditionInfo *info);
+
   void setConditionInfo(ControlFlowConditionInfo *info_) { info = info_; }
 
   llvm::StringRef getArgument() const override { return "process-args"; }
 
   ControlFlowConditionInfo *info = nullptr;
+
+  // Original iter_args of every scf.while op with main_loop attr, captured at
+  // start of ProcessArgs. Used to identify iter_args referenced by
+  // scf.condition.
+  llvm::DenseMap<scf::WhileOp, SmallVector<unsigned>>
+      originalWhileIterArgIndices;
+
+  // Local copy of whileBlockArgMap; also mirrored to info->whileBlockArgMap
+  // when info is set, so the mapping is observable when --process-args runs
+  // standalone (info may be null).
+  WhileBlockArgMap localWhileBlockArgMap;
 };
 
 std::unique_ptr<OperationPass<ModuleOp>> createProcessArgsPass();
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.h
similarity index 80%
rename from third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h
rename to third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.h
index 75a355fc9a..3f87c5f6de 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.h
@@ -20,8 +20,8 @@
  * THE SOFTWARE.
  */
 
-#ifndef TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H
-#define TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H
+#ifndef TRITON_ASCEND_SSBUF_UPDATE_LOOP_OPS_FOR_CONTROL_FLOW_H
+#define TRITON_ASCEND_SSBUF_UPDATE_LOOP_OPS_FOR_CONTROL_FLOW_H
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 
@@ -30,16 +30,16 @@
 namespace mlir {
 namespace triton {
 
-class UpdateForOpsPass
-    : public PassWrapper<UpdateForOpsPass, OperationPass<ModuleOp>> {
+class UpdateLoopOpsPass
+    : public PassWrapper<UpdateLoopOpsPass, OperationPass<ModuleOp>> {
 public:
-  UpdateForOpsPass() = default;
+  UpdateLoopOpsPass() = default;
 
   void runOnOperation() override;
 
   void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; }
 
-  llvm::StringRef getArgument() const override { return "update-for-ops"; }
+  llvm::StringRef getArgument() const override { return "update-loop-ops"; }
 
 private:
   LogicalResult
@@ -51,8 +51,8 @@ class UpdateForOpsPass
 
   LogicalResult insertInterCorePipeS(ModuleOp module);
 
-  // Analyze the dependencies of the tensor type iter_args in the main_loop with
-  // the ssbuffer.if ops
+  // Analyze tensor type iter_args dependencies in main_loop with ssbuffer.if
+  // ops
   LogicalResult
   analyzeTensorIterArgDependencies(ModuleOp module,
                                    ControlFlowConditionInfo *info);
@@ -60,8 +60,8 @@ class UpdateForOpsPass
   ControlFlowConditionInfo *info = nullptr;
 };
 
-std::unique_ptr<OperationPass<ModuleOp>> createUpdateForOpsPass();
+std::unique_ptr<OperationPass<ModuleOp>> createUpdateLoopOpsPass();
 
 } // namespace triton
 } // namespace mlir
-#endif // TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H
+#endif // TRITON_ASCEND_SSBUF_UPDATE_LOOP_OPS_FOR_CONTROL_FLOW_H
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h
index b43c7f105a..076b7fde1d 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h
@@ -29,21 +29,19 @@
 #include "llvm/ADT/SmallVector.h"
 #include <optional>
 
-namespace mlir {
-namespace triton {
+#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h"
+#include "ascend/include/DynamicCVPipeline/Common/Utils.h"
 
-// Attribute names for DynamicCV pipeline
-inline constexpr llvm::StringLiteral kSSBufferIfAttr = "ssbuffer.if";
-inline constexpr llvm::StringLiteral kHIVMMatmulLimitedInCubeAttr =
-    "hivm.matmul_limited_in_cube";
+namespace mlir {
 
 // Collect all nested ops within an operation's regions
 LogicalResult collectAllNestedOps(Operation *op,
                                   llvm::DenseSet<Operation *> &regionOps);
 
-// Group operations by their block_id attribute
+// Group body ops of a main-loop op (scf.for or scf.while carrying
+// ssbuffer.main_loop) by block_id.
 LogicalResult
-collectOpsByBlockId(scf::ForOp forOp,
+collectOpsByBlockId(Operation *op,
                     llvm::DenseMap<int, SmallVector<Operation *>> &blockOps);
 
 // Topological sort of operations based on operand dependencies
@@ -53,11 +51,18 @@ LogicalResult topologicalSort(llvm::DenseSet<Operation *> &ops,
 
 LogicalResult topologicalSort(SmallVector<Operation *> &ops);
 
-// Get block_ids in order of appearance in for loop body
-SmallVector<int> getBlockIdsInOrder(scf::ForOp forOp);
+// Get block_ids in order of appearance in the main-loop body (forOp body or
+// whileOp after-region body). Returns empty if `op` is neither.
+SmallVector<int> getBlockIdsInOrder(Operation *op);
+
+// Count unique ssbuffer.if values inside a main-loop op (scf.for or scf.while
+// carrying ssbuffer.main_loop), walking all nested ops. Returns 0 if none.
+int countUniqueIfBlockIds(Operation *loopOp);
 
-// Get the block_id of the immediate child of scf.for that contains op
-std::optional<int> getForDirectChildBlockId(Operation *op);
+// Get block_id of immediate child of main-loop (scf.for/scf.while carrying
+// ssbuffer.main_loop) that contains op. For scf.while, "body" means the
+// after-region block.
+std::optional<int> getLoopDirectChildBlockId(Operation *op);
 
 // Find the tcb group id that contains value v
 int findTcbGroupId(
@@ -68,11 +73,55 @@ int findTcbGroupId(
 // Returns failure if scopeOp does not have tcore_type attribute
 LogicalResult getScopeType(Operation *scopeOp, bool &isCube, bool &isVector);
 
-// Check if op is a scf.if whose body only contains hivm.hir.sync_block_wait,
-// hivm.hir.sync_block_set and hivm.fixpipe ops (excluding terminators).
-// Returns false if op is not a scf.if or contains any other op.
+// Check if op is scf.if whose body only contains hivm.hir.sync_block_wait,
+// hivm.hir.sync_block_set and hivm.fixpipe ops (excluding terminators). Returns
+// false if op is not scf.if or contains any other op.
 bool isIfOpWithOnlySyncOps(Operation *op);
 
-} // namespace triton
+// Migrate ops from oldBlock to newBlock; replaceAllUsesWith on oldBlock's
+// args to newBlock's args (same index). Used for both branches of scf.while
+// (before/after) and for scf.for body replacement.
+void migrateBody(Block *oldBlock, Block *newBlock);
+
+// Migrate both before and after regions of a scf.while op. Does not touch
+// terminators — the caller is expected to build a new scf.condition and
+// scf.yield in the new regions.
+void migrateWhileBodies(scf::WhileOp oldWhileOp, scf::WhileOp newWhileOp);
+
+// Build new scf.yield at end of `newBlock`: copies oldBlock's yield operands,
+// appends `extraYieldValues`, creates new scf::YieldOp, erases old yield.
+LogicalResult buildNewYieldOp(Block *oldBlock, Block *newBlock,
+                              Operation *newOp,
+                              ArrayRef<Value> extraYieldValues);
+
+// Replace all uses of `oldOp`'s results with `newOp`'s matching results.
+// No-op when `oldOp` is result-less.
+void replaceOpResultUses(Operation *oldOp, Operation *newOp);
+
+// Build new scf.condition in `newWhileOp`'s before region. Condition preserved
+// from `whileOp`; forwarded values = new before-block args (incl. extras).
+void buildNewWhileCondition(scf::WhileOp whileOp, scf::WhileOp newWhileOp);
+
+// Creates a new scf.for with `extraInitArgs` appended to the original init
+// args. Returns `oldForOp` unchanged when `extraInitArgs` is empty.
+scf::ForOp createNewForOpWithExtras(scf::ForOp oldForOp,
+                                    ArrayRef<Value> extraInitArgs);
+
+// Creates a new scf.while with `extraInitArgs` appended to the original inits
+// and empty before/after blocks. Returns `oldWhileOp` unchanged when empty.
+scf::WhileOp createNewWhileOpWithExtras(scf::WhileOp oldWhileOp,
+                                        ArrayRef<Value> extraInitArgs);
+
+// Dispatches createNewForOpWithExtras / createNewWhileOpWithExtras by op type.
+// Returns nullptr if `oldOp` is neither scf.for nor scf.while.
+Operation *createMainLoopOpWithExtras(Operation *oldOp,
+                                      ArrayRef<Value> extraInitArgs);
+
+// Prints whileBlockArgMap (whileOp -> block_id -> (new_arg_idx -> old_arg_idx))
+// to the debug stream, gated by LLVM_DEBUG. `header` is logged once before the
+// iteration.
+void dumpWhileBlockArgMap(const triton::WhileBlockArgMap &map,
+                          llvm::StringRef header);
+
 } // namespace mlir
 #endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_UTILS_H
diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
index 620161185a..6efd4613c6 100644
--- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
+++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h
@@ -68,6 +68,8 @@ inline constexpr llvm::StringLiteral kEnableUbRefineOpt =
     "ssbuffer.enable_ub_refine_opt";
 inline constexpr llvm::StringLiteral kInsertionOptimization =
     "ssbuffer.insertionOptimization";
+inline constexpr llvm::StringLiteral kArg = "ssbuffer.arg";
+inline constexpr llvm::StringLiteral kWhileArg = "ssbuffer.while_arg";
 static constexpr llvm::StringLiteral kInlinableQuantScaleAttr =
     "enable_fast_tf32_mul";
 inline constexpr llvm::StringLiteral kGMLoadMultiBufferHintAttr = "gm_load";
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp
index b4e33f9489..34262f45e3 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp
@@ -26,8 +26,8 @@
 #include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h"
 #include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.h"
 #include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h"
-#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h"
 #include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h"
+#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.h"
 #include "ascend/include/DynamicCVPipeline/Common/Utils.h"
 #include "bishengir/Dialect/HIVM/IR/HIVM.h"
 #include "bishengir/Dialect/Scope/IR/Scope.h"
@@ -104,9 +104,9 @@ void AddControlFlowConditionPass::runOnOperation() {
 
   // Step4: Update for ops with block counters and inner dependency conditions,
   // and insert PIPE_S inter-core synchronization
-  std::unique_ptr<UpdateForOpsPass> updateForOpsPass(new UpdateForOpsPass());
-  updateForOpsPass->setConditionInfo(&info);
-  pm.addPass(std::move(updateForOpsPass));
+  std::unique_ptr<UpdateLoopOpsPass> updateLoopOpsPass(new UpdateLoopOpsPass());
+  updateLoopOpsPass->setConditionInfo(&info);
+  pm.addPass(std::move(updateLoopOpsPass));
 
   // Step5:Update the conditions of ifOp based on the intraCoreDependentMap and
   // crossCoreDependentMap
@@ -141,7 +141,7 @@ void registerAddControlFlowConditionPasses() {
   registerPass(createCloneOpsPass);
   registerPass(createCreateIfOpsPass);
   registerPass(createProcessArgsPass);
-  registerPass(createUpdateForOpsPass);
+  registerPass(createUpdateLoopOpsPass);
   registerPass(createAddControlFlowConditionPass);
 }
 } // namespace triton
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp
index 4755cf99c5..0eeb79a17b 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp
@@ -38,18 +38,20 @@ static constexpr const char *DEBUG_TYPE = "CloneOps";
 #define LDBG(...)                                                              \
   LLVM_DEBUG({                                                                 \
     DBGS();                                                                    \
-    llvm::outs() << __VA_ARGS__;                                               \
-    llvm::outs() << "\n";                                                      \
+    llvm::dbgs() << __VA_ARGS__;                                               \
+    llvm::dbgs() << "\n";                                                      \
   })
 
 using namespace mlir;
 using namespace triton;
+using namespace CVPipeline;
 using namespace hivm;
 
 using MemDepGraph = std::unique_ptr<CVPipeline::MemoryDependenceGraph>;
 using MemDepGraphT = CVPipeline::MemoryDependenceGraph;
 
-// Update op operands using value mapping, skip yield values of forOp
+// Updates op operands via value mapping; skips main-loop body yield values
+// (yieldValues collected once per main-loop from the body's terminator).
 static LogicalResult
 updateCloneMapping(Operation *op, llvm::DenseMap<Value, Value> &valueMap,
                    const llvm::DenseSet<Value> &yieldValues) {
@@ -58,8 +60,8 @@ updateCloneMapping(Operation *op, llvm::DenseMap<Value, Value> &valueMap,
   }
 
   for (OpOperand &operand : op->getOpOperands()) {
-    // Only skip if this operand is a yield value from the main_loop forOp
-    // Nested yield ops (in scf.if/scf.for) should have their operands updated
+    // Only skip if this operand is a yield value from the main_loop body.
+    // Nested yield ops (in scf.if/scf.for) should have their operands updated.
     Value v = operand.get();
     if (yieldValues.contains(v)) {
       continue;
@@ -69,7 +71,7 @@ updateCloneMapping(Operation *op, llvm::DenseMap<Value, Value> &valueMap,
     if (it != valueMap.end()) {
       if (it->second.getType() != v.getType()) {
         LDBG("[Error]: type mismatch in value mapping: "
-             << v.getType() << " vs " << it->second.getType() << "\n");
+             << v.getType() << " vs " << it->second.getType());
         return failure();
       }
       operand.set(it->second);
@@ -107,12 +109,13 @@ static Operation *cloneOpWithMapping(Operation *op, OpBuilder &builder,
   return cloned;
 }
 
-// Clone ops for a single block in vector/cube mode
+// Clones ops for a single block in vector/cube mode. `bodyBlock` (forOp body
+// or whileOp after-region) only supplies iter-arg yields not to be remapped.
 static LogicalResult
 cloneOpsForBlock(int curId, SmallVector<Operation *> &curOps,
                  const SmallVector<int> &earlierIds,
                  const llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
-                 scf::ForOp forOp) {
+                 Block *bodyBlock) {
   if (curOps.empty() || earlierIds.empty()) {
     return success();
   }
@@ -144,9 +147,11 @@ cloneOpsForBlock(int curId, SmallVector<Operation *> &curOps,
   curOps.insert(curOps.begin(), clonedOps.begin(), clonedOps.end());
 
   llvm::DenseSet<Value> yieldValues;
-  if (auto yieldOp = dyn_cast<scf::YieldOp>(forOp.getBody()->getTerminator())) {
-    for (Value operand : yieldOp.getOperands()) {
-      yieldValues.insert(operand);
+  if (bodyBlock) {
+    if (auto yieldOp = dyn_cast<scf::YieldOp>(bodyBlock->getTerminator())) {
+      for (Value operand : yieldOp.getOperands()) {
+        yieldValues.insert(operand);
+      }
     }
   }
 
@@ -163,21 +168,27 @@ cloneOpsForBlock(int curId, SmallVector<Operation *> &curOps,
   return success();
 }
 
-// Clone ops for all blocks in main loop
-LogicalResult CloneOpsPass::cloneOpsInMainLoop(scf::ForOp forOp) {
-  llvm::DenseMap<int, SmallVector<Operation *>> blockOps;
-  if (failed(collectOpsByBlockId(forOp, blockOps))) {
+// Clones ops for all blocks in the main loop op (scf.for or scf.while).
+LogicalResult CloneOpsPass::cloneOpsInMainLoop(Operation *op) {
+  Block *bodyBlock = MainLoop(op).getBody();
+  if (!bodyBlock) {
+    LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp or "
+         "scf::WhileOp");
     return failure();
   }
 
-  SmallVector<int> idsInOrder = getBlockIdsInOrder(forOp);
+  llvm::DenseMap<int, SmallVector<Operation *>> blockOps;
+  if (failed(collectOpsByBlockId(op, blockOps))) {
+    return failure();
+  }
+  SmallVector<int> idsInOrder = getBlockIdsInOrder(op);
 
   for (int i = idsInOrder.size() - 1; i >= 0; --i) {
     int curId = idsInOrder[i];
     SmallVector<int> earlierIds(idsInOrder.begin(), idsInOrder.begin() + i);
 
     if (failed(cloneOpsForBlock(curId, blockOps[curId], earlierIds, blockOps,
-                                forOp))) {
+                                bodyBlock))) {
       return failure();
     }
   }
@@ -185,15 +196,8 @@ LogicalResult CloneOpsPass::cloneOpsInMainLoop(scf::ForOp forOp) {
   return success();
 }
 
-// Check if an op should be erased during cleanup (for cube)
-// sameBlockIdExecAfter: precomputed map from cloned op to its same-block-id
-// exec-after ops
-//   (already filtered to skip SyncBlockWaitOp/SyncBlockSetOp), built once per
-//   block from a single MemoryDependenceGraph.
-// erasedOps: set of cloned ops already erased in this cleanup pass; an
-// exec-after entry
-//   that has been erased no longer pins the current op (equivalent to the
-//   original per-op graph rebuild reflecting the post-erasure IR).
+// Should an op be erased during cube cleanup? sameBlockIdExecAfter: per-block
+// same-block-id exec-after ops (sync filtered); erasedOps no longer pin ops.
 static bool shouldEraseOpForCube(
     Operation *op,
     const llvm::DenseMap<Operation *, llvm::SmallPtrSet<Operation *, 4>>
@@ -205,11 +209,10 @@ static bool shouldEraseOpForCube(
     return true;
   }
 
-  auto opBlockId = getForDirectChildBlockId(op);
+  auto opBlockId = getLoopDirectChildBlockId(op);
 
-  // Rule 2: If op has results, check via SSA if result is used by later ops in
-  // same block_id Use for-direct-child block_id (the immediate child of
-  // scf.for) for comparison
+  // Rule 2: if op has results, check via SSA whether any result is used by a
+  // later op with the same loop-direct-child block_id (main-loop's own child).
   if (op->getNumResults() > 0) {
     for (auto result : op->getResults()) {
       if (result.use_empty()) {
@@ -220,27 +223,24 @@ static bool shouldEraseOpForCube(
         // No block_id but result is used, be conservative and keep
         return false;
       }
-      // Check if any user is in the same for-direct-child block_id
+      // Check if any user is in the same loop-direct-child block_id
       bool usedInSameBlockId =
           llvm::any_of(result.getUsers(), [&](Operation *user) {
-            auto userBlockId = getForDirectChildBlockId(user);
+            auto userBlockId = getLoopDirectChildBlockId(user);
             return userBlockId && *userBlockId == *opBlockId;
           });
       if (usedInSameBlockId) {
-        // Result used in same for-direct-child block, cannot erase
+        // Result used in same loop-direct-child block, cannot erase
         return false;
       }
     }
-    // All results are either unused or not used in same for-direct-child block,
-    // can erase
+    // All results are either unused or not used in same loop-direct-child
+    // block, can erase
     return true;
   }
 
-  // Rule 3: If op has no results, consult the precomputed same-block-id
-  // exec-after set. An exec-after op that has already been erased in this pass
-  // is no longer live and does not require keeping this op. This mirrors the
-  // original per-iteration memGraph rebuild, where already-erased ops simply
-  // disappear from getExecAfter().
+  // Rule 3: no results -> consult precomputed same-block-id exec-after set; ops
+  // already erased in this pass are not live (mirrors per-iteration rebuild).
   auto it = sameBlockIdExecAfter.find(op);
   if (it != sameBlockIdExecAfter.end()) {
     for (Operation *execOp : it->second) {
@@ -248,9 +248,8 @@ static bool shouldEraseOpForCube(
       if (erasedOps.contains(execOp)) {
         continue;
       }
-      // scf.if whose body only contains sync_block_wait/sync_block_set ops
-      // will be cleaned up in its own turn; skip it here so it does not
-      // block the current op's erasure.
+      // scf.if bodies holding only sync_block_wait/set are cleaned up in their
+      // own turn; skip so they don't block this op's erasure.
       if (isIfOpWithOnlySyncOps(execOp)) {
         continue;
       }
@@ -271,14 +270,34 @@ static bool shouldEraseOpForVector(Operation *op) {
                        [](auto result) { return !result.use_empty(); });
 }
 
-// Cleanup for cloned ops in a forOp
-// memGraphFactory: callable that rebuilds MemoryDependenceGraph for current IR
-// state
+// Verify cloned sync/fixpipe ops were erased after cleanup; nullptr bodyBlock
+// means nothing to check.
+static LogicalResult validateClonedSyncOpsErased(Block *bodyBlock) {
+  if (!bodyBlock) {
+    return success();
+  }
+  for (Operation &op : bodyBlock->without_terminator()) {
+    if (!op.hasAttr(CVPipeline::kClone)) {
+      continue;
+    }
+    if (isa<SyncBlockWaitOp>(&op) || isa<SyncBlockSetOp>(&op) ||
+        isa<hivm::FixpipeOp>(&op)) {
+      LDBG("[ERROR]: Cloned sync/fixpipe op should have been erased: "
+           << op.getName());
+      return failure();
+    }
+  }
+
+  return success();
+}
+
+// Cleans up cloned ops in a main-loop op. `bodyBlock` is only scanned after
+// cleanup; `mainLoopOp` goes to `memGraphFactory` to rebuild the mem graph.
 static LogicalResult
-cleanupClonedOps(scf::ForOp forOp,
+cleanupClonedOps(Operation *mainLoopOp, Block *bodyBlock,
                  llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
                  const SmallVector<int> &idsInOrder, bool isCube,
-                 std::function<MemDepGraph(scf::ForOp)> memGraphFactory) {
+                 std::function<MemDepGraph(Operation *)> memGraphFactory) {
   for (int i = idsInOrder.size() - 1; i >= 0; --i) {
     auto &curOps = blockOps[idsInOrder[i]];
     if (curOps.empty()) {
@@ -297,10 +316,8 @@ cleanupClonedOps(scf::ForOp forOp,
       continue;
     }
 
-    // Locate the start of the contiguous cloned-op suffix. The original cleanup
-    // loop broke on the first non-cloned op; preserve that behavior so we only
-    // consider ops that the previous implementation would have considered,
-    // regardless of how topologicalSort interleaves cloned and non-cloned ops.
+    // Locate the start of the contiguous cloned-op suffix; the original
+    // cleanup broke at the first non-cloned op, so keep that same set.
     int firstClonedIdx = 0;
     for (int j = startIdx - 1; j >= 0; --j) {
       if (!curOps[j]->hasAttr(CVPipeline::kClone)) {
@@ -316,23 +333,19 @@ cleanupClonedOps(scf::ForOp forOp,
       clonedOps.push_back(curOps[j]);
     }
 
-    // Build the MemoryDependenceGraph at most once per block, and only if at
-    // least one cloned op has no results (Rule 3 is the only path that needs
-    // it; Rule 2 is a pure SSA check and Rule 1 is type-based).
-    std::unique_ptr<MemDepGraphT> memGraph;
+    // Build the MemoryDependenceGraph at most once per block, only if a cloned
+    // op has no results (only Rule 3 needs it; Rule 2 SSA, Rule 1 type-based).
+    MemDepGraph memGraph;
     if (isCube) {
       bool needsMemGraph = llvm::any_of(
           clonedOps, [](Operation *o) { return o->getNumResults() == 0; });
       if (needsMemGraph) {
-        memGraph = memGraphFactory(forOp);
+        memGraph = memGraphFactory(mainLoopOp);
       }
     }
 
-    // Precompute, for each cloned op, the set of its exec-after ops that share
-    // its block_id. Building the memGraph was the expensive part; we now do
-    // getExecAfter once per op up front and reuse the result during cleanup.
-    // An erased op is filtered out at check time (see shouldEraseOpForCube),
-    // which reproduces the original "rebuild after each erasure" effect.
+    // Precompute each cloned op's exec-after ops sharing its block_id once and
+    // reuse it; erased ops are filtered at check time (as if rebuilt).
     llvm::DenseMap<Operation *, llvm::SmallPtrSet<Operation *, 4>>
         sameBlockIdExecAfter;
     if (memGraph) {
@@ -341,13 +354,13 @@ cleanupClonedOps(scf::ForOp forOp,
           // Rule 2 path: memgraph is not consulted.
           continue;
         }
-        auto opBlockId = getForDirectChildBlockId(op);
+        auto opBlockId = getLoopDirectChildBlockId(op);
         if (!opBlockId) {
           continue;
         }
         for (Operation *execOp : memGraph->getExecAfter(op)) {
-          // sync_block_wait/sync_block_set ops are not memory side effects in
-          // analyzing cleanup ops, therefore we skip them.
+          // sync_block_wait/sync_block_set are not memory side effects for
+          // cleanup analysis, so skip them.
           if (isa<SyncBlockWaitOp>(execOp) || isa<SyncBlockSetOp>(execOp)) {
             continue;
           }
@@ -359,10 +372,8 @@ cleanupClonedOps(scf::ForOp forOp,
       }
     }
 
-    // Erase cloned ops from bottom to top, matching the original ordering.
-    // erasedOps mirrors the effect of the per-iteration graph rebuild: an op
-    // already erased in this pass is treated as absent from the precomputed
-    // exec-after sets.
+    // Erase cloned ops bottom-to-top (original ordering); erasedOps makes an
+    // already-erased op count as absent from the precomputed exec-after sets.
     llvm::DenseSet<Operation *> erasedOps;
     for (int j = startIdx; j >= firstClonedIdx; --j) {
       Operation *op = curOps[j];
@@ -376,25 +387,21 @@ cleanupClonedOps(scf::ForOp forOp,
     }
   }
 
-  // Check whether new forOp is valid after cleanup
-  for (Operation &op : forOp.getBody()->without_terminator()) {
-    if (op.hasAttr(CVPipeline::kClone)) {
-      if (isa<SyncBlockWaitOp>(op) || isa<SyncBlockSetOp>(op) ||
-          isa<hivm::FixpipeOp>(op)) {
-        LDBG("[ERROR]: Cloned sync/fixpipe op should have been erased: "
-             << op.getName() << "\n");
-        return failure();
-      }
-    }
-  }
-
-  return success();
+  return validateClonedSyncOpsErased(bodyBlock);
 }
 
-// Cleanup cloned ops for a single main loop
-LogicalResult CloneOpsPass::cleanupClonedOpsInMainLoop(scf::ForOp forOp) {
+// Cleans up cloned ops for one main loop (scf.for or scf.while); body block
+// comes from MainLoop::getBody, the op feeds collect/order helpers+factory.
+LogicalResult CloneOpsPass::cleanupClonedOpsInMainLoop(Operation *op) {
+  Block *bodyBlock = MainLoop(op).getBody();
+  if (!bodyBlock) {
+    LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp or "
+         "scf::WhileOp");
+    return failure();
+  }
+
   ModuleOp module = getOperation();
-  scope::ScopeOp scopeOp = forOp->getParentOfType<scope::ScopeOp>();
+  scope::ScopeOp scopeOp = op->getParentOfType<scope::ScopeOp>();
   if (!scopeOp) {
     return success();
   }
@@ -409,20 +416,20 @@ LogicalResult CloneOpsPass::cleanupClonedOpsInMainLoop(scf::ForOp forOp) {
                                                   hivm::TCoreType::CUBE));
 
   llvm::DenseMap<int, SmallVector<Operation *>> blockOps;
-  if (failed(collectOpsByBlockId(forOp, blockOps))) {
+  if (failed(collectOpsByBlockId(op, blockOps))) {
     return failure();
   }
+  SmallVector<int> idsInOrder = getBlockIdsInOrder(op);
 
-  SmallVector<int> idsInOrder = getBlockIdsInOrder(forOp);
-  if (failed(cleanupClonedOps(forOp, blockOps, idsInOrder, isCube,
-                              [&](scf::ForOp forOp) -> MemDepGraph {
+  if (failed(cleanupClonedOps(op, bodyBlock, blockOps, idsInOrder, isCube,
+                              [&](Operation *loopOp) -> MemDepGraph {
                                 if (!isCube) {
                                   return nullptr;
                                 }
                                 auto &aliasAnalysis =
                                     getAnalysis<mlir::AliasAnalysis>();
                                 return std::make_unique<MemDepGraphT>(
-                                    forOp, aliasAnalysis);
+                                    loopOp, aliasAnalysis);
                               }))) {
     return failure();
   }
@@ -430,14 +437,14 @@ LogicalResult CloneOpsPass::cleanupClonedOpsInMainLoop(scf::ForOp forOp) {
   return success();
 }
 
-// Validate that each block_id's ops form contiguous ranges (not interleaved
-// with other ids) e.g., [1,1,2,2] is valid, but [1,2,1,2] is invalid
-static bool areBlockIdsConsecutive(scf::ForOp forOp) {
+// Validate each block_id's ops form contiguous ranges, not interleaved (e.g.
+// [1,1,2,2] valid, [1,2,1,2] invalid), on a main-loop body block.
+static bool areBlockIdsConsecutive(Block *bodyBlock) {
   SmallVector<int> idsInOrder;
-  for (Operation &op : forOp.getBody()->without_terminator()) {
+  for (Operation &op : bodyBlock->without_terminator()) {
     auto blockIdOpt = CVPipeline::getOpBlockId(&op);
     if (!blockIdOpt) {
-      LDBG("[ERROR]: Op missing ssbuffer.block_id: " << op.getName() << "\n");
+      LDBG("[ERROR]: Op missing ssbuffer.block_id: " << op.getName());
       return false;
     }
 
@@ -455,7 +462,7 @@ static bool areBlockIdsConsecutive(scf::ForOp forOp) {
 
     for (size_t k = j; k < idsInOrder.size(); ++k) {
       if (idsInOrder[k] == currentId) {
-        LDBG("[ERROR]: block_id: " << currentId << " is interleaved\n");
+        LDBG("[ERROR]: block_id: " << currentId << " is interleaved");
         return false;
       }
     }
@@ -468,15 +475,16 @@ static bool areBlockIdsConsecutive(scf::ForOp forOp) {
 
 LogicalResult CloneOpsPass::validateBlockIdsConsecutive(ModuleOp module) {
   WalkResult result = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr(CVPipeline::kMainLoop)) {
+    if (!isMainLoopOp(op)) {
       return WalkResult::advance();
     }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
+    Block *bodyBlock = MainLoop(op).getBody();
+    if (!bodyBlock) {
+      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp or "
+           "scf::WhileOp");
       return WalkResult::interrupt();
     }
-    if (!areBlockIdsConsecutive(forOp)) {
+    if (!areBlockIdsConsecutive(bodyBlock)) {
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
@@ -487,20 +495,21 @@ LogicalResult CloneOpsPass::validateBlockIdsConsecutive(ModuleOp module) {
   return success();
 }
 
-// Check that no op in a VECTOR scope's main_loop forOp has a tensor result \
-// carrying the ssbuffer.clone attribute.
+// Checks that no op in a VECTOR scope's main_loop op (scf.for or scf.while)
+// has a tensor result carrying ssbuffer.clone.
 LogicalResult CloneOpsPass::validateClonedOpsInVector(ModuleOp module) {
   WalkResult result = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr(CVPipeline::kMainLoop)) {
+    if (!isMainLoopOp(op)) {
       return WalkResult::advance();
     }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
+    Block *bodyBlock = MainLoop(op).getBody();
+    if (!bodyBlock) {
+      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp or "
+           "scf::WhileOp");
       return WalkResult::interrupt();
     }
 
-    scope::ScopeOp scopeOp = forOp->getParentOfType<scope::ScopeOp>();
+    scope::ScopeOp scopeOp = op->getParentOfType<scope::ScopeOp>();
     if (!scopeOp) {
       return WalkResult::advance();
     }
@@ -516,11 +525,11 @@ LogicalResult CloneOpsPass::validateClonedOpsInVector(ModuleOp module) {
       return WalkResult::advance();
     }
 
-    for (Operation &bodyOp : forOp.getBody()->without_terminator()) {
+    for (Operation &bodyOp : bodyBlock->without_terminator()) {
       if (!bodyOp.hasAttr(CVPipeline::kClone)) {
         continue;
       }
-      if (isa<tensor::EmptyOp>(bodyOp)) {
+      if (isa<tensor::EmptyOp>(&bodyOp)) {
         continue;
       }
       bool hasTensorDep = llvm::any_of(bodyOp.getResults(), [](Value result) {
@@ -528,7 +537,7 @@ LogicalResult CloneOpsPass::validateClonedOpsInVector(ModuleOp module) {
       });
       if (hasTensorDep) {
         LDBG("[Error]: VECTOR main_loop contains cloned op with tensor type: "
-             << bodyOp.getName() << "\n");
+             << bodyOp.getName());
         return WalkResult::interrupt();
       }
     }
@@ -549,7 +558,7 @@ void CloneOpsPass::runOnOperation() {
     return;
   }
 
-  LDBG("before cloneOps:\n" << module << "\n");
+  LDBG("before cloneOps:\n" << module);
 
   // Validate block_ids are consecutive before cloning
   if (failed(validateBlockIdsConsecutive(module))) {
@@ -557,35 +566,31 @@ void CloneOpsPass::runOnOperation() {
     return;
   }
 
-  // Clone ops in vector/cube to ensure that each block_id has its own
-  // ops without sharing
+  // Clone ops in vector/cube so each block_id owns its ops (no sharing); entry
+  // points take any main-loop op (scf.for/scf.while) and dispatch internally.
   auto walkResult = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr(CVPipeline::kMainLoop)) {
+    if (!isMainLoopOp(op)) {
       return WalkResult::advance();
     }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
-      return WalkResult::interrupt();
-    }
 
-    if (failed(cloneOpsInMainLoop(forOp))) {
+    if (failed(cloneOpsInMainLoop(op))) {
+      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
       return WalkResult::interrupt();
     }
 
-    if (failed(cleanupClonedOpsInMainLoop(forOp))) {
+    if (failed(cleanupClonedOpsInMainLoop(op))) {
+      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
   });
   if (walkResult.wasInterrupted()) {
-    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
     return;
   }
 
-  LDBG("after cloneOps:\n" << module << "\n");
+  LDBG("after cloneOps:\n" << module);
 
-  // Validate no cloned tensor ops remaining in VECTOR main_loop forOp
+  // Validate no cloned tensor ops remaining in VECTOR main_loop op
   if (failed(validateClonedOpsInVector(module))) {
     CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
     return;
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp
index aa0e9be22c..d027ab7dcc 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp
@@ -22,11 +22,6 @@
 
 #include "llvm/Support/Debug.h"
 
-#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h"
-#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h"
-#include "ascend/include/DynamicCVPipeline/Common/Utils.h"
-#include "bishengir/Dialect/HIVM/IR/HIVM.h"
-#include "bishengir/Dialect/Scope/IR/Scope.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -45,12 +40,13 @@ static constexpr const char *DEBUG_TYPE = "CreateIfOps";
 #define LDBG(...)                                                              \
   LLVM_DEBUG({                                                                 \
     DBGS();                                                                    \
-    llvm::outs() << __VA_ARGS__;                                               \
-    llvm::outs() << "\n";                                                      \
+    llvm::dbgs() << __VA_ARGS__;                                               \
+    llvm::dbgs() << "\n";                                                      \
   })
 
 using namespace mlir;
 using namespace triton;
+using namespace CVPipeline;
 
 // Check if a value is used outside the given region ops
 static bool isUsedOutsideRegion(Value v,
@@ -64,7 +60,7 @@ static bool isUsedOutsideRegion(Value v,
   return false;
 }
 
-// Find the iteration argument in main loop that corresponds to the given value
+// Find iter_arg in main loop for `v` (yield operand idx = body arg position).
 static Value findIterArgInMainLoop(Value v, mlir::Type t) {
   for (Operation *user : v.getUsers()) {
     auto yieldOp = dyn_cast<scf::YieldOp>(user);
@@ -72,14 +68,21 @@ static Value findIterArgInMainLoop(Value v, mlir::Type t) {
       continue;
     }
 
-    auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-    if (!forOp) {
+    SmallVector<Value> iterArgs =
+        MainLoop(yieldOp->getParentOp()).getIterArgs();
+    if (iterArgs.empty()) {
       continue;
     }
 
     for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) {
       if (operand.getAsOpaquePointer() == v.getAsOpaquePointer()) {
-        Value iterArg = forOp.getRegionIterArgs()[idx];
+        if (idx >= iterArgs.size()) {
+          LDBG("[Error]: yield operand index "
+               << idx << " out of range for iter_args size "
+               << iterArgs.size());
+          continue;
+        }
+        Value iterArg = iterArgs[idx];
         if (iterArg.getType() == t) {
           return iterArg;
         }
@@ -87,7 +90,7 @@ static Value findIterArgInMainLoop(Value v, mlir::Type t) {
     }
   }
 
-  LDBG("[Error]: else yield value not found in forOp iter_args: " << v << "\n");
+  LDBG("[Error]: else yield value not found in forOp/whileOp iter_args: " << v);
   return nullptr;
 }
 
@@ -99,21 +102,20 @@ static LogicalResult replaceExternalIfOpUses(scf::IfOp ifOp,
     Value oldVal = oldYieldValues[i];
 
     if (!oldVal) {
-      LDBG("[Error]: oldVal is null at index " << i << "\n");
+      LDBG("[Error]: oldVal is null at index " << i);
       return failure();
     }
 
     if (i >= ifOp.getNumResults()) {
       LDBG("[Error]: index " << i << " exceeds ifOp results count "
-                             << ifOp.getNumResults() << "\n");
+                             << ifOp.getNumResults());
       return failure();
     }
 
     Value newVal = ifOp.getResult(i);
     if (oldVal.getType() != newVal.getType()) {
       LDBG("[Error]: type mismatch at index " << i << ": " << oldVal.getType()
-                                              << " vs " << newVal.getType()
-                                              << "\n");
+                                              << " vs " << newVal.getType());
       return failure();
     }
 
@@ -144,13 +146,14 @@ static LogicalResult replaceExternalIfOpUses(scf::IfOp ifOp,
   return success();
 }
 
-// Compute yield values for each block: values that need to be yielded from the
-// if
+// Computes yield values per block (values yielded from the if); op-agnostic as
+// findIterArgInMainLoop dispatches on the yield's parent op (forOp/whileOp).
 LogicalResult CreateIfOpsPass::computeYieldValues(
-    scf::ForOp forOp,
+    Operation *loopOp,
     const llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
     llvm::DenseMap<int, SmallVector<Value>> &thenYieldValues,
     llvm::DenseMap<int, SmallVector<Value>> &elseYieldValues) {
+  (void)loopOp; // unused; findIterArgInMainLoop walks the yield's parent.
   for (auto &p : blockOps) {
     int id = p.first;
     const SmallVector<Operation *> &ops = p.second;
@@ -212,7 +215,7 @@ static scf::IfOp createIfOpForBlock(OpBuilder &builder, Location loc,
   // Check size consistency
   if (needsYield && thenValues.size() != elseValues.size()) {
     LDBG("[Error]: then/else yield count mismatch: "
-         << thenValues.size() << " vs " << elseValues.size() << "\n");
+         << thenValues.size() << " vs " << elseValues.size());
     return scf::IfOp();
   }
 
@@ -222,7 +225,7 @@ static scf::IfOp createIfOpForBlock(OpBuilder &builder, Location loc,
       if (thenValues[i].getType() != elseValues[i].getType()) {
         LDBG("[Error]: then/else yield type mismatch at index "
              << i << ": " << thenValues[i].getType() << " vs "
-             << elseValues[i].getType() << "\n");
+             << elseValues[i].getType());
         return scf::IfOp();
       }
     }
@@ -240,7 +243,7 @@ static scf::IfOp createIfOpForBlock(OpBuilder &builder, Location loc,
     ifOp = builder.create<scf::IfOp>(loc, TypeRange{}, trueVal, false);
   }
 
-  ifOp->setAttr(kSSBufferIfAttr, builder.getI32IntegerAttr(blockId));
+  ifOp->setAttr(CVPipeline::kIf, builder.getI32IntegerAttr(blockId));
 
   // notify npuir that of the scenario
   ifOp->setAttr(CVPipeline::kHIVMMatmulLimitedInCubeAttr,
@@ -256,7 +259,7 @@ static LogicalResult moveOpsToThenBranch(scf::IfOp ifOp,
                                          const SmallVector<Value> &elseValues,
                                          Location loc) {
   if (ops.empty() && !thenValues.empty()) {
-    LDBG("[Error]: moving empty ops but thenValues not empty\n");
+    LDBG("[Error]: moving empty ops but thenValues not empty");
     return failure();
   }
 
@@ -278,13 +281,19 @@ static LogicalResult moveOpsToThenBranch(scf::IfOp ifOp,
   return success();
 }
 
-// Create if ops (scf.if %true) for each block_id in the main loop
+// Create if ops (scf.if %true) for each block_id in the main loop.
+// `op` is the main-loop op (scf.for or scf.while).
 LogicalResult CreateIfOpsPass::createIfInMainLoop(
-    scf::ForOp forOp,
+    Operation *op,
     const llvm::DenseMap<int, SmallVector<Operation *>> &blockOps,
     const llvm::DenseMap<int, SmallVector<Value>> &thenYieldValues,
     const llvm::DenseMap<int, SmallVector<Value>> &elseYieldValues) {
-  SmallVector<int> ids = getBlockIdsInOrder(forOp);
+  SmallVector<int> ids = getBlockIdsInOrder(op);
+  if (ids.empty() && !MainLoop(op).getBody()) {
+    LDBG("[Error]: op with ssbuffer.main_loop is neither scf::ForOp nor "
+         "scf::WhileOp");
+    return failure();
+  }
 
   for (int id : ids) {
     const SmallVector<Operation *> &ops = blockOps.lookup(id);
@@ -324,37 +333,35 @@ void CreateIfOpsPass::runOnOperation() {
     return;
   }
 
-  LDBG("before createIfOps:\n" << module << "\n");
+  LDBG("before createIfOps:\n" << module);
 
   auto walkResult = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr("ssbuffer.main_loop")) {
+    if (!isMainLoopOp(op)) {
       return WalkResult::advance();
     }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
-      return WalkResult::interrupt();
-    }
 
-    // Create if ops (scf.if %true) by block_id
     llvm::DenseMap<int, SmallVector<Operation *>> blockOps;
-    if (failed(collectOpsByBlockId(forOp, blockOps))) {
+    if (failed(collectOpsByBlockId(op, blockOps))) {
+      LDBG("[Error]: op with ssbuffer.main_loop is neither scf::ForOp nor "
+           "scf::WhileOp, or a body op is missing ssbuffer.block_id\n");
       return WalkResult::interrupt();
     }
 
+    // blockCounterNums is keyed on the main-loop op (scf.for or scf.while).
+    // Record it for both forOp and whileOp.
     if (info) {
-      info->blockCounterNums[forOp] = blockOps.size();
+      info->blockCounterNums[op] = blockOps.size();
     }
 
     llvm::DenseMap<int, SmallVector<Value>> thenYieldValues;
     llvm::DenseMap<int, SmallVector<Value>> elseYieldValues;
 
-    if (failed(computeYieldValues(forOp, blockOps, thenYieldValues,
+    if (failed(computeYieldValues(op, blockOps, thenYieldValues,
                                   elseYieldValues))) {
       return WalkResult::interrupt();
     }
 
-    if (failed(createIfInMainLoop(forOp, blockOps, thenYieldValues,
+    if (failed(createIfInMainLoop(op, blockOps, thenYieldValues,
                                   elseYieldValues))) {
       return WalkResult::interrupt();
     }
@@ -365,7 +372,7 @@ void CreateIfOpsPass::runOnOperation() {
     return;
   }
 
-  LDBG("after createIfOps:\n" << module << "\n");
+  LDBG("after createIfOps:\n" << module);
 }
 
 namespace mlir {
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp
index e189711379..3a187598c1 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp
@@ -22,6 +22,8 @@
 
 #include "third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h"
 #include "ascend/include/DynamicCVPipeline/Common/BufferCountManager.h"
+#include "bishengir/Dialect/Annotation/IR/Annotation.h"
+#include "bishengir/Dialect/HIVM/IR/HIVM.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h"
@@ -41,26 +43,21 @@ static constexpr const char *DEBUG_TYPE = "InitDependentMap";
 #define LDBG(...) LLVM_DEBUG(DBGS() << __VA_ARGS__ << "\n")
 
 using namespace mlir;
-using namespace hivm;
 using namespace triton;
+using namespace CVPipeline;
+using namespace hivm;
 
-// Function: Check if a consumer op is inside a given mainLoop
-//            but not inside any nested mainloop, and push to consumers if true
-// Input: Consumer operation, target mainLoop forOp, reference to consumers
-// vector Output: consumers - push consumer to this vector if it's inside
-// mainLoop (not in nested mainloop) Return: 0 for success (consumer pushed), -1
-// for failure (not in mainLoop or in nested mainLoop)
-static int isConsumerInMainLoop(Operation *consumer, scf::ForOp mainLoop,
+// Returns 0 if `consumer` is inside `mainLoop` (and pushes it to `consumers`)
+// or inside a nested mainloop (skip), or -1 on error.
+static int isConsumerInMainLoop(Operation *consumer, Operation *mainLoop,
                                 SmallVector<Operation *> &consumers) {
   Operation *current = consumer->getParentOp();
 
   // Traverse up the parent chain until we reach the top (nullptr)
   while (current != nullptr) {
-    if (auto forOp = dyn_cast<scf::ForOp>(current)) {
-      if (forOp->hasAttr(CVPipeline::kMainLoop) && forOp != mainLoop) {
-        // comsumer Op not in the current mainloop
-        return 0;
-      }
+    if (isMainLoopOp(current) && current != mainLoop) {
+      // consumer Op not in the current mainloop
+      return 0;
     }
     // If we reach the target mainLoop, consumer is inside it
     if (current == mainLoop) {
@@ -74,12 +71,8 @@ static int isConsumerInMainLoop(Operation *consumer, scf::ForOp mainLoop,
   return -1;
 }
 
-// Function: Collect ops with dependency attributes, grouped by group ID
-// Input: Root operation to traverse (module or forOp), attribute name
-// Output: depsByGroup - Ops grouped by group ID, format: group -> [(op, role),
-// ...]
-//         Attribute format: [group, role], role: 1=producer, 0=consumer
-// Return: 0 for success, -1 for failure
+// Collect ops with dependency attr `attrName` into depsByGroup (group ->
+// [(op, role)], attr = [group, role], 1=producer/0=consumer). 0 ok, -1 fail.
 static int
 collectDepsByGroup(Operation *rootOp, const char *attrName,
                    llvm::DenseMap<int, SmallVector<std::pair<Operation *, int>>>
@@ -113,16 +106,12 @@ collectDepsByGroup(Operation *rootOp, const char *attrName,
   return ret;
 }
 
-// Function: Build mapping from consumer Operation to producer Operation
-// Input: Ops grouped by group ID, format: group -> [(op, role), ...]
-//        role: 1=producer, 0=consumer
-//        mainLoop: if not nullptr, only include consumers inside this mainLoop
-// Output: result - Mapping from consumer Operation* to list of producer
-// Operation* Return: 0 for success, -1 for failure
+// Build consumer -> producers mapping from depsByGroup (role 1=producer,
+// 0=consumer); if mainLoop != nullptr only consumers inside it. 0 ok, -1 fail.
 static int buildProducerConsumerMapping(
     llvm::DenseMap<int, SmallVector<std::pair<Operation *, int>>> &depsByGroup,
     llvm::DenseMap<Operation *, SmallVector<Operation *>> &result,
-    scf::ForOp mainLoop = nullptr) {
+    Operation *mainLoop = nullptr) {
   for (auto &groupEntry : depsByGroup) {
     auto &ops = groupEntry.second;
 
@@ -165,16 +154,23 @@ static int buildProducerConsumerMapping(
   return 0;
 }
 
+// Collects every scf.for/scf.while op tagged CVPipeline::kMainLoop, stored
+// uniformly as Operation* so downstream lookups ignore the op kind.
 static int collectMainLoopById(ModuleOp module,
-                               llvm::DenseMap<scf::ForOp, int> &mainLoopById) {
+                               llvm::DenseMap<Operation *, int> &mainLoopById) {
   int ret = 0;
-  module.walk([&](scf::ForOp forOp) {
-    if (!forOp->hasAttr(CVPipeline::kMainLoop))
+  module.walk([&](Operation *op) {
+    if (!op->hasAttr(CVPipeline::kMainLoop))
       return;
-    auto mainLoopIdAttr =
-        forOp->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
+    if (!isMainLoopOp(op)) {
+      LDBG("Do not support mainloop op other than scf.for or scf.while: "
+           << op->getName());
+      ret = -1;
+      return;
+    }
+    auto mainLoopIdAttr = op->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
     if (mainLoopIdAttr) {
-      mainLoopById[forOp] = mainLoopIdAttr.getInt();
+      mainLoopById[op] = mainLoopIdAttr.getInt();
     }
   });
   return ret;
@@ -182,7 +178,7 @@ static int collectMainLoopById(ModuleOp module,
 
 static int
 findMainLoopIdContainingOp(Operation *op,
-                           llvm::DenseMap<scf::ForOp, int> &mainLoopById) {
+                           llvm::DenseMap<Operation *, int> &mainLoopById) {
   for (auto &entry : mainLoopById) {
     if (entry.first->isAncestor(op)) {
       return entry.second;
@@ -197,8 +193,8 @@ static int filterMemCrossCoreDepsByMainLoop(
     llvm::DenseMap<Operation *, SmallVector<Operation *>> &filteredDepsMap) {
   LDBG("memCrossCore dependencies before filter: " << initialDepsMap.size());
 
-  // Step 1: Collect all main_loop forOps and their ids
-  llvm::DenseMap<scf::ForOp, int> mainLoopById;
+  // Step 1: Collect all main_loop ops (scf.for or scf.while) and their ids
+  llvm::DenseMap<Operation *, int> mainLoopById;
   if (collectMainLoopById(module, mainLoopById) != 0) {
     LDBG("collectMainLoopById Failed!");
     return -1;
@@ -258,12 +254,8 @@ static int filterMemCrossCoreDepsByMainLoop(
   return 0;
 }
 
-// Initialize crossCoreDependentMap (cross-core data dependency)
-// Find ops with ssbuffer.crossDeps attribute
-// Attribute value is a list: [group, role], role: 1=producer, 0=consumer
-// Map key is consumer, value is list of all producers in the same group
-// Constraint: consumer and producer must be in the same main_loop (with same
-// id) Return: 0 for success, -1 for failure
+// Init crossCoreDependentMap from ssbuffer.crossDeps ([group, role]; 1=producer
+// 0=consumer): consumer -> same-group producers, same main_loop. 0 ok, -1 fail.
 int initCrossCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) {
   // Step 1: Collect all crossDeps by group (including memCrossDeps)
   llvm::DenseMap<int, SmallVector<std::pair<Operation *, int>>>
@@ -294,11 +286,8 @@ int initCrossCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) {
   return 0;
 }
 
-// Initialize intraCoreDependentMap (intra-core data dependency)
-// Find forOp with ssbuffer.main_loop attribute
-// Collect all intra-core deps from module (producers may be outside the loop)
-// For each mainLoop, filter consumers that are inside it (not in nested
-// mainloops) Return: 0 for success, -1 for failure
+// Initializes intraCoreDependentMap: per main_loop op (scf.for/scf.while) keeps
+// consumers inside it but not in nested mainloops. 0 on success, -1 on error.
 int initIntraCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) {
   // Collect all intra-core deps from the entire module
   llvm::DenseMap<int, SmallVector<std::pair<Operation *, int>>>
@@ -309,20 +298,21 @@ int initIntraCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) {
     return -1;
   }
 
-  // For each mainLoop, build mapping with consumers inside it
+  // For each mainLoop, build the mapping with consumers inside it; both scf.for
+  // and scf.while can carry ssbuffer.main_loop (map keyed on Operation*).
   int ret = 0;
   module.walk([&](Operation *op) {
     if (!op->hasAttr(CVPipeline::kMainLoop))
       return;
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("Do not support other mainloop except forOp!");
+    if (!isMainLoopOp(op)) {
+      LDBG("Do not support mainloop op other than scf.for or scf.while: "
+           << op->getName());
       ret = -1;
       return;
     }
 
     llvm::DenseMap<Operation *, SmallVector<Operation *>> depMap;
-    if (buildProducerConsumerMapping(allIntraDepsByGroup, depMap, forOp) != 0) {
+    if (buildProducerConsumerMapping(allIntraDepsByGroup, depMap, op) != 0) {
       LDBG("buildProducerConsumerMapping on intraDeps Failed!");
       ret = -1;
       return;
@@ -330,7 +320,7 @@ int initIntraCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) {
 
     // Only insert if there are dependencies for this mainLoop
     if (!depMap.empty()) {
-      info->intraCoreDependentMap[forOp] = depMap;
+      info->intraCoreDependentMap[op] = depMap;
     }
   });
   return ret;
@@ -355,13 +345,10 @@ static void printDependentMaps(ControlFlowConditionInfo *info) {
   LDBG("intraCoreDependentMap size: " << info->intraCoreDependentMap.size());
   LDBG("intraCoreDependentMap contents:");
   for (auto &forEntry : info->intraCoreDependentMap) {
-    scf::ForOp forOp = forEntry.first;
+    Operation *loopOp = forEntry.first;
     auto &depMap = forEntry.second;
-    LDBG("  ForOp (depMap size: " << depMap.size() << "):");
-    LDBG("    ");
-    LLVM_DEBUG(llvm::dbgs() << '[' << DEBUG_TYPE << "] ";
-               forOp->print(llvm::dbgs(), OpPrintingFlags().skipRegions());
-               llvm::dbgs() << "\n";);
+    LDBG("  MainLoopOp (depMap size: " << depMap.size() << "):");
+    LDBG("    " << OpWithFlags(loopOp, OpPrintingFlags().skipRegions()));
 
     for (auto &entry : depMap) {
       Operation *consumer = entry.first;
@@ -404,9 +391,8 @@ static scf::IfOp findIfOpContainingOp(Operation *op) {
   return nullptr;
 }
 
-// Compute producer buffer count from dependency maps
-// Rule: traverse cross-core and intra-core maps, assign max size
-// If intra-core map is empty, use BufferCountManager's IntraCore value
+// Compute producer buffer counts (max map size) from cross/intra-core maps;
+// falls back to BufferCountManager IntraCore when the intra-core map is empty.
 static void computeProducerBufferCount(ControlFlowConditionInfo *info,
                                        ModuleOp module) {
   // Get cross-core buffer count (max size in the map)
@@ -417,16 +403,16 @@ static void computeProducerBufferCount(ControlFlowConditionInfo *info,
   }
   LDBG("Cross-core buffer count (max): " << info->crossCoreBufferCount);
 
-  // Get intra-core buffer count (max size across all forOps)
+  // Get intra-core buffer count (max size across all main loops)
   info->intraCoreBufferCount = 0;
-  for (auto &forOpEntry : info->intraCoreDependentMap) {
-    auto &intraDepMap = forOpEntry.second;
+  for (auto &loopEntry : info->intraCoreDependentMap) {
+    auto &intraDepMap = loopEntry.second;
     for (auto &entry : intraDepMap) {
       info->intraCoreBufferCount =
           std::max(info->intraCoreBufferCount, (int)entry.second.size());
     }
   }
-  LDBG("Intra-core buffer count (max across all forOps): "
+  LDBG("Intra-core buffer count (max across all main loops): "
        << info->intraCoreBufferCount);
 
   // If intra-core map is empty, use BufferCountManager's IntraCore value
@@ -486,9 +472,8 @@ static int buildIfBlockCrossCoreDAG(ModuleOp module,
   return 0;
 }
 
-// Detect cross-core cycle in the if-block DAG using DFS.
-// All edges in this DAG are cross-core (CUBE↔VECTOR), so any cycle
-// indicates a deadlock-prone bidirectional data dependency.
+// Detect cross-core cycle in the if-block DAG via DFS; all edges are cross-core
+// (CUBE<->VECTOR), so any cycle is a deadlock-prone bidirectional dependency.
 enum class DfsState : uint8_t { Unvisited, Visiting, Done };
 
 static bool dfsCycle(scf::IfOp node,
@@ -562,9 +547,8 @@ static void dfsFindNodesAtDistance(
   }
 }
 
-// Collect flowOpt if block pairs from DAG using DFS
-// Find all start nodes (in-degree = 0), then use DFS to find nodes at distance
-// 2
+// Collect flowOpt if-block pairs from the DAG: find start nodes (in-degree 0),
+// then DFS for nodes at distance 2.
 static int collectFlowOptIfOpPairs(ModuleOp module,
                                    ControlFlowConditionInfo *info) {
   // Step 1: Calculate in-degree for each node
@@ -654,7 +638,7 @@ void InitDependentMapPass::runOnOperation() {
   // Print all dependent maps for verification
   LLVM_DEBUG(printDependentMaps(info));
 
-  // Step 3: Compute producer buffer count for flowOpt condition
+  // Step 4: Compute producer buffer count for flowOpt condition
   computeProducerBufferCount(info, module);
 
   // Step 4: Build if block DAG from crossCoreDependentMap (always)
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.cpp
index 631a420a33..8fbe41e432 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/ProcessArgs.cpp
@@ -42,35 +42,40 @@ static constexpr const char *DEBUG_TYPE = "ProcessArgs";
   })
 
 using namespace mlir;
+using namespace CVPipeline;
 using namespace triton;
 
-// Collects mapping from iter_arg index to block_ids that use it.
-// For each iter_arg, tracks which block_ids reference it in their operations.
+// Collects iter_arg index -> block_ids that use it. `ivOffset` is 1 for
+// scf.for (IV at arg 0) or 0 for scf.while. Uses body args directly
+// (block-local).
 static LogicalResult collectArgIndexToBlockIds(
-    scf::ForOp forOp,
+    Block *body, unsigned ivOffset,
     llvm::DenseMap<int, llvm::DenseSet<int>> &argIndexToBlockIds) {
-  Block *body = forOp.getBody();
   if (!body || !body->mightHaveTerminator()) {
-    LDBG("[Error]: forOp body is invalid or has no terminator\n");
+    LDBG("[Error]: loop body is invalid or has no terminator");
     return failure();
   }
 
   for (Operation &op : body->without_terminator()) {
-    auto blockIdAttr = op.getAttrOfType<IntegerAttr>("ssbuffer.block_id");
+    auto blockIdAttr = op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId);
     if (!blockIdAttr)
       continue;
     int blockId = blockIdAttr.getInt();
 
     for (OpOperand &operand : op.getOpOperands()) {
       Value v = operand.get();
-      for (unsigned i = 0; i < forOp.getNumRegionIterArgs(); ++i) {
-        Value iterArg = forOp.getRegionIterArgs()[i];
+      for (BlockArgument iterArg : body->getArguments()) {
+        int argIdx = iterArg.getArgNumber();
+        if (argIdx < (int)ivOffset) {
+          // scf.for's IV at block arg 0 — never an iter_arg.
+          continue;
+        }
         // Skip tensor-type iter_args, only process scalar and index types
         if (mlir::isa<TensorType>(iterArg.getType())) {
           continue;
         }
         if (v == iterArg) {
-          argIndexToBlockIds[i].insert(blockId);
+          argIndexToBlockIds[argIdx - (int)ivOffset].insert(blockId);
         }
       }
     }
@@ -78,9 +83,8 @@ static LogicalResult collectArgIndexToBlockIds(
   return success();
 }
 
-// Finds iter_args used by multiple block_ids (shared args).
-// Determines owner block (first in order) and creates SharedArgInfo for each
-// non-owner. Each non-owner block gets its own extra iter_arg.
+// Finds iter_args used by multiple block_ids. Owner = first block in order;
+// each non-owner gets its own extra iter_arg via SharedArgInfo.
 static LogicalResult findSharedArgs(
     const llvm::DenseMap<int, llvm::DenseSet<int>> &argIndexToBlockIds,
     const SmallVector<int> &idsInOrder,
@@ -105,37 +109,29 @@ static LogicalResult findSharedArgs(
 
     // Each non-owner block for this argIndex gets its own extra iter_arg
     for (int bid : blockIds) {
-      if (bid != ownerBlockId) {
-        sharedArgsInfo.push_back(
-            SharedArgInfo(argIndex, ownerBlockId, extraArgCount, bid));
-        extraArgCount++;
-      }
+      if (bid == ownerBlockId)
+        continue;
+      sharedArgsInfo.push_back(
+          SharedArgInfo(argIndex, ownerBlockId, extraArgCount, bid));
+      extraArgCount++;
     }
   }
   return success();
 }
 
-// Finds the computation operation in owner block that produces the iter_arg
-// value. compOp is the defining op of the iter_arg in the scf.yield operand
-// list.
-static LogicalResult findCompOpInOwnerBlock(scf::ForOp forOp, Block *body,
-                                            const SharedArgInfo &info,
-                                            Operation *&compOp) {
+// Returns the op defining iter_arg at `argIndex` in `body`'s scf.yield (top
+// of update chain), or nullptr if out of bounds. Shared by both paths.
+static Operation *findYieldDefiningOp(Block *body, unsigned argIndex) {
   auto yieldOp = cast<scf::YieldOp>(body->getTerminator());
-  Value yieldArg = yieldOp.getOperand(info.argIndex);
-
-  if (auto *defOp = yieldArg.getDefiningOp()) {
-    compOp = defOp;
-    return success();
+  if (argIndex >= yieldOp.getNumOperands()) {
+    return nullptr;
   }
-
-  return failure();
+  return yieldOp.getOperand(argIndex).getDefiningOp();
 }
 
-// Collects all operations in the computation chain by backward traversal from
-// compOp. Builds the dependency graph needed to clone the computation for
-// non-owner blocks.
-static void collectChainOps(scf::ForOp forOp, Operation *compOp,
+// Collects all operations in the computation chain by backward traversal
+// from compOp, scoped to ops inside `loopOp`'s body.
+static void collectChainOps(Operation *loopOp, Operation *compOp,
                             llvm::DenseSet<Operation *> &chainOps) {
   SmallVector<Operation *> worklist;
   worklist.push_back(compOp);
@@ -148,7 +144,7 @@ static void collectChainOps(scf::ForOp forOp, Operation *compOp,
 
     for (Value operand : op->getOperands()) {
       if (auto *defOp = operand.getDefiningOp()) {
-        if (defOp->getParentOp() == forOp && !chainOps.contains(defOp)) {
+        if (defOp->getParentOp() == loopOp && !chainOps.contains(defOp)) {
           worklist.push_back(defOp);
         }
       }
@@ -156,9 +152,10 @@ static void collectChainOps(scf::ForOp forOp, Operation *compOp,
   }
 }
 
-// Builds computation info (compOp and chainOps) for each shared arg.
+// Builds compOp + chainOps for each shared arg. `body` is the loop body
+// (forOp body or whileOp after-body) used to locate scf.yield.
 static LogicalResult buildCompInfoForSharedArgs(
-    scf::ForOp forOp, Block *body, SmallVector<SharedArgInfo> &sharedArgsInfo,
+    Operation *loopOp, Block *body, SmallVector<SharedArgInfo> &sharedArgsInfo,
     llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
     llvm::DenseMap<int, llvm::DenseSet<Operation *>> &sharedArgToChainOps) {
   for (auto &info : sharedArgsInfo) {
@@ -166,64 +163,23 @@ static LogicalResult buildCompInfoForSharedArgs(
     if (sharedArgToCompOp.contains(argIndex))
       continue;
 
-    Operation *compOp = nullptr;
-    if (failed(findCompOpInOwnerBlock(forOp, body, info, compOp))) {
+    Operation *compOp = findYieldDefiningOp(body, argIndex);
+    if (!compOp) {
       continue;
     }
 
     sharedArgToCompOp[argIndex] = compOp;
 
     llvm::DenseSet<Operation *> chainOps;
-    collectChainOps(forOp, compOp, chainOps);
+    collectChainOps(loopOp, compOp, chainOps);
     sharedArgToChainOps[argIndex] = chainOps;
   }
   return success();
 }
 
-// Creates a new scf.for op with extra iter_args for shared arguments.
-// Copies attributes from the original for op.
-// Each SharedArgInfo entry (non-owner block) gets its own extra iter_arg.
-static scf::ForOp
-createNewForOp(scf::ForOp forOp,
-               const SmallVector<SharedArgInfo> &sharedArgsInfo) {
-  OpBuilder builder(forOp);
-  SmallVector<Value> newInitArgs(forOp.getInitArgs().begin(),
-                                 forOp.getInitArgs().end());
-
-  // Each non-owner block gets its own extra iter_arg
-  for (auto &info : sharedArgsInfo) {
-    newInitArgs.push_back(forOp.getInitArgs()[info.argIndex]);
-  }
-
-  scf::ForOp newForOp = builder.create<scf::ForOp>(
-      forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
-      forOp.getStep(), newInitArgs);
-
-  for (auto &attr : forOp->getAttrs()) {
-    newForOp->setAttr(attr.getName(), attr.getValue());
-  }
-  return newForOp;
-}
-
-// Migrates operations from old block to new block.
-// Redirects block arguments to new block arguments and moves all ops.
-static void migrateBody(Block *oldBlock, Block *newBlock) {
-  for (unsigned i = 0; i < oldBlock->getNumArguments(); ++i) {
-    oldBlock->getArgument(i).replaceAllUsesWith(newBlock->getArgument(i));
-  }
-
-  for (Operation &op :
-       llvm::make_early_inc_range(oldBlock->without_terminator())) {
-    op.moveBefore(newBlock, newBlock->end());
-  }
-}
-
-// Clones the computation chain for a non-owner block.
-// Topologically sorts the chain and clones each op with remapped operands.
-// argRemapping: maps migrated iter_arg Value -> new extra iter_arg Value.
-// resultMapper: maps original op results -> cloned op results.
-// clonedArgIdx: unique index for this non-owner block's clone (used as
-// ssbuffer.arg).
+// Clones the computation chain for a non-owner block (topo-sorted, operands
+// remapped). argRemapping: iter_arg -> new extra. clonedArgIdx: this block's
+// arg.
 static LogicalResult
 cloneChainForBlock(SharedArgInfo &info, Operation *compOp,
                    const llvm::DenseSet<Operation *> &chainOps, Block *newBlock,
@@ -257,9 +213,9 @@ cloneChainForBlock(SharedArgInfo &info, Operation *compOp,
       continue;
 
     Operation *cloned = cloneBuilder.clone(*op, opMapper);
-    cloned->setAttr("ssbuffer.block_id",
+    cloned->setAttr(CVPipeline::kBlockId,
                     cloneBuilder.getI32IntegerAttr(info.nonOwnerBlockId));
-    cloned->setAttr("ssbuffer.arg",
+    cloned->setAttr(CVPipeline::kArg,
                     cloneBuilder.getI32IntegerAttr(info.argIndex));
 
     resultMapper.map(op->getResult(0), cloned->getResult(0));
@@ -275,7 +231,7 @@ static LogicalResult replaceIterArgsInBlock(SharedArgInfo &info,
                                             IRMapping &argRemapping,
                                             OpBuilder &cloneBuilder) {
   for (Operation &op : newBlock->without_terminator()) {
-    auto blockIdAttr = op.getAttrOfType<IntegerAttr>("ssbuffer.block_id");
+    auto blockIdAttr = op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId);
     if (!blockIdAttr || blockIdAttr.getInt() != info.nonOwnerBlockId)
       continue;
 
@@ -284,7 +240,7 @@ static LogicalResult replaceIterArgsInBlock(SharedArgInfo &info,
       if (argRemapping.contains(operand)) {
         Value newVal = argRemapping.lookup(operand);
         op.setOperand(i, newVal);
-        op.setAttr("ssbuffer.arg",
+        op.setAttr(CVPipeline::kArg,
                    cloneBuilder.getI32IntegerAttr(info.argIndex));
       }
     }
@@ -293,24 +249,22 @@ static LogicalResult replaceIterArgsInBlock(SharedArgInfo &info,
 }
 
 // Processes each shared arg: finds insertion point, clones chain, replaces
-// iter_args. Collects cloned results for building new yield operands.
+// iter_args. `ivOffset` is 1 for scf.for (IV at arg 0), 0 for scf.while.
 static LogicalResult processSharedArgsIteration(
-    scf::ForOp forOp, Block *newBlock,
-    SmallVector<SharedArgInfo> &sharedArgsInfo,
+    Block *newBlock, SmallVector<SharedArgInfo> &sharedArgsInfo,
     const llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
     const llvm::DenseMap<int, llvm::DenseSet<Operation *>> &sharedArgToChainOps,
-    const SmallVector<Value> &oldBlockArgs, SmallVector<Value> &clonedResults) {
-  unsigned numOriginalIterArgs = forOp.getNumRegionIterArgs();
-  unsigned extraIterArgsBase =
-      1 + numOriginalIterArgs; // block arg index where extra iter_args start
+    ValueRange iterArgs, unsigned ivOffset, SmallVector<Value> &clonedResults) {
+  unsigned numOriginalIterArgs = iterArgs.size();
+  unsigned extraIterArgsBase = ivOffset + numOriginalIterArgs;
 
   int clonedArgIdx = clonedResults.size();
   for (auto &info : sharedArgsInfo) {
     int argIndex = info.argIndex;
-    info.iterArg = forOp.getRegionIterArgs()[argIndex];
+    info.iterArg = iterArgs[argIndex];
 
     // The migrated iter_arg (original iter_arg moved to new block)
-    Value migratedIterArg = newBlock->getArgument(argIndex + 1);
+    Value migratedIterArg = newBlock->getArgument(argIndex + ivOffset);
     // The new extra iter_arg added for this shared arg
     unsigned newExtraBlockArgIdx = extraIterArgsBase + info.newArgIndex;
     Value newExtraIterArg = newBlock->getArgument(newExtraBlockArgIdx);
@@ -321,7 +275,7 @@ static LogicalResult processSharedArgsIteration(
 
     Operation *lastOpInBlock = nullptr;
     for (Operation &op : newBlock->without_terminator()) {
-      auto blockIdAttr = op.getAttrOfType<IntegerAttr>("ssbuffer.block_id");
+      auto blockIdAttr = op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId);
       if (blockIdAttr && blockIdAttr.getInt() == info.nonOwnerBlockId) {
         lastOpInBlock = &op;
       }
@@ -356,21 +310,27 @@ static LogicalResult processSharedArgsIteration(
 // Prepares all shared args data: collects arg->blockId mapping, finds shared
 // args, and builds computation info for each shared arg.
 static LogicalResult prepareSharedArgsData(
-    scf::ForOp forOp, SmallVector<SharedArgInfo> &sharedArgsInfo,
+    Operation *loopOp, Block *body, SmallVector<SharedArgInfo> &sharedArgsInfo,
     llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
     llvm::DenseMap<int, llvm::DenseSet<Operation *>> &sharedArgToChainOps) {
-  Block *body = forOp.getBody();
   if (!body || !body->mightHaveTerminator()) {
-    LDBG("[Error]: forOp body is invalid or has no terminator\n");
+    LDBG("[Error]: loop body is invalid or has no terminator");
     return failure();
   }
 
+  // ivOffset: 1 for scf.for (IV at block arg 0), 0 for scf.while (no IV).
+  unsigned ivOffset = isa<scf::ForOp>(loopOp) ? 1 : 0;
+
   llvm::DenseMap<int, llvm::DenseSet<int>> argIndexToBlockIds;
-  if (failed(collectArgIndexToBlockIds(forOp, argIndexToBlockIds))) {
+  if (failed(collectArgIndexToBlockIds(body, ivOffset, argIndexToBlockIds))) {
     return failure();
   }
 
-  SmallVector<int> idsInOrder = getBlockIdsInOrder(forOp);
+  SmallVector<int> idsInOrder = getBlockIdsInOrder(loopOp);
+  if (idsInOrder.empty() && !MainLoop(loopOp).getBody()) {
+    LDBG("[Error]: loopOp is neither scf::ForOp nor scf::WhileOp");
+    return failure();
+  }
   if (failed(findSharedArgs(argIndexToBlockIds, idsInOrder, sharedArgsInfo))) {
     return failure();
   }
@@ -380,9 +340,9 @@ static LogicalResult prepareSharedArgsData(
   }
 
   LDBG("[INFO]: Found " << sharedArgsInfo.size()
-                        << " shared iter_args to process\n");
+                        << " shared iter_args to process");
 
-  if (failed(buildCompInfoForSharedArgs(forOp, body, sharedArgsInfo,
+  if (failed(buildCompInfoForSharedArgs(loopOp, body, sharedArgsInfo,
                                         sharedArgToCompOp,
                                         sharedArgToChainOps))) {
     return failure();
@@ -391,117 +351,480 @@ static LogicalResult prepareSharedArgsData(
   return success();
 }
 
-// Builds new yield op with original operands plus cloned results.
-static LogicalResult buildNewYieldOp(Block *oldBlock, Block *newBlock,
-                                     scf::ForOp newForOp,
-                                     const SmallVector<Value> &clonedResults) {
-  auto oldYield = cast<scf::YieldOp>(oldBlock->getTerminator());
-  SmallVector<Value> yieldOperands;
+// Replaces all uses of the old main-loop op with the new op's results, erases
+// the old op, and transfers its intraCoreDependentMap entry to the new op.
+static LogicalResult replaceMainLoopOpAndErase(Operation *oldOp,
+                                               Operation *newOp,
+                                               ControlFlowConditionInfo *info) {
+  replaceOpResultUses(oldOp, newOp);
 
-  for (unsigned i = 0; i < oldYield.getNumOperands(); ++i) {
-    yieldOperands.push_back(oldYield.getOperand(i));
-  }
-  for (auto &result : clonedResults) {
-    yieldOperands.push_back(result);
+  // Transfer intraCoreDependentMap entry from oldOp to newOp.
+  if (info && info->intraCoreDependentMap.count(oldOp)) {
+    info->intraCoreDependentMap[newOp] = info->intraCoreDependentMap[oldOp];
+    info->intraCoreDependentMap.erase(oldOp);
   }
 
-  OpBuilder builder = OpBuilder::atBlockEnd(newBlock);
-  builder.create<scf::YieldOp>(newForOp.getLoc(), yieldOperands);
-  oldYield.erase();
+  oldOp->erase();
   return success();
 }
 
-// Replaces all uses of old for op with new for op results and erases old for
-// op. Also transfers intraCoreDependentMap entry from oldForOp to newForOp.
-static LogicalResult replaceForOpAndErase(scf::ForOp oldForOp,
-                                          scf::ForOp newForOp,
-                                          ControlFlowConditionInfo *info) {
-  if (oldForOp.getNumResults() > 0) {
-    SmallVector<Value> newResults;
-    for (unsigned i = 0; i < oldForOp.getNumResults(); ++i) {
-      newResults.push_back(newForOp.getResult(i));
-    }
-    oldForOp.replaceAllUsesWith(newResults);
-  }
+// Derives the body block to inspect for shared-iter_arg analysis and the IV
+// offset (1 for scf.for, 0 for scf.while). Returns false if `op` is neither.
+static bool getOpIterParams(Operation *op, Block *&inspectBody,
+                            unsigned &ivOffset) {
+  MainLoop ml = MainLoop(op);
+  if (ml.getBody()) {
+    inspectBody = ml.getBody();
+    ivOffset = ml.isWhile() ? 0 : 1;
+    return true;
+  }
+  return false;
+}
+
+// Completes the scf.for path: migrate body, clone per-block chains, rebuild
+// scf.yield, swap old op out.
+static LogicalResult processSharedArgsInForOp(
+    scf::ForOp forOp, scf::ForOp newForOp,
+    SmallVector<SharedArgInfo> &sharedArgsInfo,
+    const llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
+    const llvm::DenseMap<int, llvm::DenseSet<Operation *>> &sharedArgToChainOps,
+    ControlFlowConditionInfo *info) {
+  Block *oldBlock = forOp.getBody();
+  Block *newBlock = newForOp.getBody();
+  migrateBody(oldBlock, newBlock);
 
-  // Transfer intraCoreDependentMap entry from oldForOp to newForOp
-  if (info && info->intraCoreDependentMap.count(oldForOp)) {
-    info->intraCoreDependentMap[newForOp] =
-        info->intraCoreDependentMap[oldForOp];
-    info->intraCoreDependentMap.erase(oldForOp);
+  SmallVector<Value> clonedResults;
+  if (failed(processSharedArgsIteration(
+          newBlock, sharedArgsInfo, sharedArgToCompOp, sharedArgToChainOps,
+          MainLoop(forOp).getIterArgs(), 1, clonedResults))) {
+    return failure();
+  }
+  if (failed(buildNewYieldOp(oldBlock, newBlock, newForOp, clonedResults))) {
+    return failure();
   }
+  return replaceMainLoopOpAndErase(forOp, newForOp, info);
+}
 
-  oldForOp.erase();
-  return success();
+// Completes the scf.while path: migrate before/after bodies, clone per-block
+// chains, rebuild scf.yield/condition, transfer maps, swap old op out.
+LogicalResult ProcessArgsPass::processSharedArgsInWhileOp(
+    scf::WhileOp whileOp, scf::WhileOp newWhileOp,
+    SmallVector<SharedArgInfo> &sharedArgsInfo,
+    const llvm::DenseMap<int, Operation *> &sharedArgToCompOp,
+    const llvm::DenseMap<int, llvm::DenseSet<Operation *>> &sharedArgToChainOps,
+    ControlFlowConditionInfo *info) {
+  migrateWhileBodies(whileOp, newWhileOp);
+
+  SmallVector<Value> clonedResults;
+  if (failed(processSharedArgsIteration(
+          newWhileOp.getAfterBody(), sharedArgsInfo, sharedArgToCompOp,
+          sharedArgToChainOps, MainLoop(whileOp).getIterArgs(), 0,
+          clonedResults))) {
+    return failure();
+  }
+  buildNewWhileCondition(whileOp, newWhileOp);
+  if (failed(buildNewYieldOp(whileOp.getAfterBody(), newWhileOp.getAfterBody(),
+                             newWhileOp, clonedResults))) {
+    return failure();
+  }
+  // Transfer originalWhileIterArgIndices so downstream still resolves
+  // newWhileOp->old arg indices.
+  if (originalWhileIterArgIndices.count(whileOp)) {
+    originalWhileIterArgIndices[newWhileOp] =
+        originalWhileIterArgIndices[whileOp];
+    originalWhileIterArgIndices.erase(whileOp);
+  }
+  // Transfer whileBlockArgMap entries (local + info mirror) so the (block_id,
+  // new_arg_idx) -> old_arg_idx mapping survives the old-op erase.
+  if (localWhileBlockArgMap.count(whileOp)) {
+    localWhileBlockArgMap[newWhileOp] =
+        std::move(localWhileBlockArgMap[whileOp]);
+    localWhileBlockArgMap.erase(whileOp);
+  }
+  if (info && info->whileBlockArgMap.count(whileOp)) {
+    info->whileBlockArgMap[newWhileOp] =
+        std::move(info->whileBlockArgMap[whileOp]);
+    info->whileBlockArgMap.erase(whileOp);
+  }
+  return replaceMainLoopOpAndErase(whileOp, newWhileOp, info);
 }
 
-// Main entry point for processing shared iter_args in a single for op.
-// Orchestrates data preparation, new for op creation, body migration, and
-// cloning.
-static LogicalResult
-processSharedIterArgsInForOp(scf::ForOp forOp, ControlFlowConditionInfo *info) {
+// Single entry point for processing shared iter_args in a main-loop op
+// (scf.for or scf.while). whileOp builds new scf.condition, transfers indices.
+LogicalResult
+ProcessArgsPass::processSharedIterArgsInLoop(Operation *op,
+                                             ControlFlowConditionInfo *info) {
+  Block *inspectBody = nullptr;
+  unsigned ivOffset = 0;
+  if (!getOpIterParams(op, inspectBody, ivOffset)) {
+    LDBG("[Error]: op with ssbuffer.main_loop is neither scf::ForOp nor "
+         "scf::WhileOp");
+    return failure();
+  }
+
   SmallVector<SharedArgInfo> sharedArgsInfo;
   llvm::DenseMap<int, Operation *> sharedArgToCompOp;
   llvm::DenseMap<int, llvm::DenseSet<Operation *>> sharedArgToChainOps;
-
-  if (failed(prepareSharedArgsData(forOp, sharedArgsInfo, sharedArgToCompOp,
-                                   sharedArgToChainOps))) {
+  if (failed(prepareSharedArgsData(op, inspectBody, sharedArgsInfo,
+                                   sharedArgToCompOp, sharedArgToChainOps))) {
     return failure();
   }
-
   if (sharedArgsInfo.empty()) {
     return success();
   }
 
-  scf::ForOp newForOp = createNewForOp(forOp, sharedArgsInfo);
-  Block *oldBlock = forOp.getBody();
-  Block *newBlock = newForOp.getBody();
+  // Build extra init values from original arg values at shared indices. Each
+  // new iter_arg shadows the existing one at the same index.
+  SmallVector<Value> extraInitArgs;
+  extraInitArgs.reserve(sharedArgsInfo.size());
+  ValueRange origInits;
+  if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+    origInits = forOp.getInitArgs();
+  } else if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
+    origInits = whileOp.getInits();
+  }
+  for (const auto &info : sharedArgsInfo) {
+    extraInitArgs.push_back(origInits[info.argIndex]);
+  }
 
-  SmallVector<Value> oldBlockArgs;
-  for (unsigned i = 0; i < oldBlock->getNumArguments(); ++i) {
-    oldBlockArgs.push_back(oldBlock->getArgument(i));
+  Operation *newOp = createMainLoopOpWithExtras(op, extraInitArgs);
+  if (!newOp) {
+    return failure();
   }
 
-  migrateBody(oldBlock, newBlock);
+  if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+    return processSharedArgsInForOp(forOp, cast<scf::ForOp>(newOp),
+                                    sharedArgsInfo, sharedArgToCompOp,
+                                    sharedArgToChainOps, info);
+  }
+  return processSharedArgsInWhileOp(
+      cast<scf::WhileOp>(op), cast<scf::WhileOp>(newOp), sharedArgsInfo,
+      sharedArgToCompOp, sharedArgToChainOps, info);
+}
 
-  SmallVector<Value> clonedResults;
-  if (failed(processSharedArgsIteration(forOp, newBlock, sharedArgsInfo,
-                                        sharedArgToCompOp, sharedArgToChainOps,
-                                        oldBlockArgs, clonedResults))) {
+// Walks module to find for/while ops with ssbuffer.main_loop attribute and
+// dispatches each into processSharedIterArgsInLoop.
+LogicalResult ProcessArgsPass::processSharedIterArgs(ModuleOp module) {
+  WalkResult result = module.walk([&](Operation *op) -> WalkResult {
+    if (!isMainLoopOp(op)) {
+      return WalkResult::advance();
+    }
+    if (failed(processSharedIterArgsInLoop(op, info))) {
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+
+  if (result.wasInterrupted()) {
     return failure();
   }
+  return success();
+}
 
-  if (failed(buildNewYieldOp(oldBlock, newBlock, newForOp, clonedResults))) {
-    return failure();
+// ----- Per-block update chain for scf.while iter_args used in scf.condition
+// For each scf.while op with main_loop, every block_id's run keeps its own
+// iter_arg value; clones update chain into each block, extends scf.while with
+// one extra iter_arg per (block_id, original iter_arg), and records
+// (new_arg_idx -> old_arg_idx) in info->whileBlockArgMap.
+
+// Snapshots original iter_args of every scf.while with main_loop; per whileOp
+// clones cond-used update chains per block_id. Recorded in
+// info->whileBlockArgMap.
+LogicalResult
+ProcessArgsPass::updateIndependentCondsInWhileBlocks(ModuleOp module) {
+  // Collect whileOps in a worklist (must NOT mutate IR during the walk —
+  // processWhileIterArgsInWhileOp replaces the old whileOp). Snapshot
+  // original iter_args in the same pass so we can later identify which
+  // iter_args were referenced by scf.condition in the input.
+  SmallVector<scf::WhileOp> worklist;
+  module.walk([&](scf::WhileOp whileOp) {
+    if (!isMainLoopOp(whileOp))
+      return;
+    SmallVector<unsigned> indices;
+    for (unsigned i = 0; i < whileOp.getNumOperands(); ++i)
+      indices.push_back(i);
+    originalWhileIterArgIndices[whileOp] = indices;
+    worklist.push_back(whileOp);
+  });
+  for (scf::WhileOp whileOp : worklist) {
+    if (failed(processWhileIterArgsInWhileOp(whileOp, info)))
+      return failure();
   }
 
-  if (failed(replaceForOpAndErase(forOp, newForOp, info))) {
+  // Dump whileBlockArgMap (new_whileop -> block_id -> (new_arg_idx ->
+  // old_arg_idx)); pass-local map so observable when --process-args runs
+  // standalone.
+  dumpWhileBlockArgMap(localWhileBlockArgMap,
+                       "whileBlockArgMap contents (new_whileop -> block_id -> "
+                       "(new_arg_idx -> old_arg_idx))");
+  return success();
+}
+
+// Returns original iter_arg indices contributing to cond value of scf.condition
+// for `whileOp`. Walks def-chain to before-block BlockArguments (cond-used).
+static llvm::DenseSet<unsigned> collectConditionUsedIterArgIndices(
+    scf::WhileOp whileOp, const SmallVector<unsigned> &originalIndices) {
+  llvm::DenseSet<unsigned> used;
+  auto cond = whileOp.getConditionOp();
+  Value condValue = cond.getCondition();
+  Block *beforeBlock = whileOp.getBeforeBody();
+
+  // BFS over def-chain of condValue; stops at before-block BlockArgument
+  // (collected) or no defining op in region (parent BlockArgument — ignored).
+  llvm::SmallPtrSet<Value, 16> visited;
+  SmallVector<Value> worklist;
+  worklist.push_back(condValue);
+  while (!worklist.empty()) {
+    Value v = worklist.pop_back_val();
+    if (!visited.insert(v).second)
+      continue;
+
+    if (auto blockArg = dyn_cast<BlockArgument>(v)) {
+      if (blockArg.getOwner() == beforeBlock) {
+        unsigned idx = blockArg.getArgNumber();
+        if (idx < originalIndices.size()) {
+          used.insert(idx);
+        }
+      }
+      continue;
+    }
+
+    Operation *defOp = v.getDefiningOp();
+    if (!defOp)
+      continue;
+    for (Value operand : defOp->getOperands()) {
+      worklist.push_back(operand);
+    }
+  }
+  return used;
+}
+
+// (findYieldCompOpForWhile folded into findYieldDefiningOp — see top of
+// file. Callers now pass `whileOp.getAfterBody()` directly.)
+
+// Returns the last op in `body` whose `ssbuffer.block_id` matches `blockId`
+// (i.e., end of the block's run of consecutive ops).
+static Operation *findLastOpWithBlockId(Block *body, int blockId) {
+  Operation *last = nullptr;
+  for (Operation &op : body->without_terminator()) {
+    auto attr = op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId);
+    if (attr && attr.getInt() == blockId) {
+      last = &op;
+    }
+  }
+  return last;
+}
+
+// Builds a clone of the update chain after the last op with `blockId` in
+// `body`. Annotates with ssbuffer.while_arg/block_id; remaps operands.
+static LogicalResult cloneUpdateChainForWhileBlock(
+    scf::WhileOp whileOp, Block *body, Operation *compOp, int blockId,
+    unsigned originalArgIndex, unsigned newArgIndex,
+    const llvm::DenseSet<Operation *> &chainOps, Value &cloned) {
+  SmallVector<Operation *> sortedChain(chainOps.begin(), chainOps.end());
+  if (failed(topologicalSort(sortedChain))) {
+    cloned = Value();
     return failure();
   }
 
+  // Find insertion point: after the last op with this block_id.
+  Operation *lastOp = findLastOpWithBlockId(body, blockId);
+  OpBuilder builder(body, body->end());
+  if (lastOp) {
+    builder.setInsertionPointAfter(lastOp);
+  }
+  // If no op has this block_id, the builder is at body->end() (set by ctor);
+  // avoid setInsertionPoint(terminator) since it may not exist yet.
+
+  IRMapping resultMapper;
+  for (Operation *op : sortedChain) {
+    IRMapping opMapper;
+    for (OpOperand &operand : op->getOpOperands()) {
+      Value oldVal = operand.get();
+      Value newVal = oldVal;
+      // Remap the original iter_arg to the new iter_arg for this block;
+      // chain-internal operands remap to the corresponding cloned value.
+      auto blockArg = dyn_cast<BlockArgument>(oldVal);
+      if (blockArg && blockArg.getOwner() == body &&
+          (unsigned)blockArg.getArgNumber() == originalArgIndex) {
+        newVal = body->getArgument(newArgIndex);
+      } else if (resultMapper.contains(oldVal)) {
+        newVal = resultMapper.lookup(oldVal);
+      }
+      opMapper.map(oldVal, newVal);
+    }
+
+    if (resultMapper.contains(op->getResult(0)))
+      continue;
+
+    Operation *cloned = builder.clone(*op, opMapper);
+    cloned->setAttr(CVPipeline::kBlockId, builder.getI32IntegerAttr(blockId));
+    cloned->setAttr(CVPipeline::kWhileArg,
+                    builder.getI32IntegerAttr(originalArgIndex));
+    resultMapper.map(op->getResult(0), cloned->getResult(0));
+    builder.setInsertionPointAfter(cloned);
+  }
+
+  cloned = resultMapper.lookup(compOp->getResult(0));
   return success();
 }
 
-// Walks module to find for ops with ssbuffer.main_loop attribute.
-// Processes each main loop to handle shared iter_args.
-LogicalResult ProcessArgsPass::processSharedIterArgs(ModuleOp module) {
-  WalkResult result = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr("ssbuffer.main_loop")) {
-      return WalkResult::advance();
+// For each cond-used iter_arg, plans one new iter_arg per (block_id, origIdx).
+// compOp + def-chain must be computed BEFORE migrateBody (caused lookup crash).
+static WhileIterArgClonePlan
+planWhileIterArgDescriptors(scf::WhileOp whileOp,
+                            const SmallVector<unsigned> &originalIndices,
+                            const llvm::DenseSet<unsigned> &condUsed) {
+  WhileIterArgClonePlan plan;
+  SmallVector<int> blockIdsInOrder = getBlockIdsInOrder(whileOp);
+  if (blockIdsInOrder.empty())
+    return plan;
+  unsigned nextNewArgIdx = whileOp.getNumOperands();
+  unsigned descIdx = 0;
+  for (unsigned origIdx : originalIndices) {
+    if (!condUsed.contains(origIdx))
+      continue;
+    plan.posInClonedVec[origIdx] = descIdx++;
+
+    Operation *compOp = findYieldDefiningOp(whileOp.getAfterBody(), origIdx);
+    if (!compOp) {
+      LDBG("[WARN]: no compOp for while iter_arg idx=" << origIdx);
+      continue;
     }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
-      return WalkResult::interrupt();
+
+    llvm::DenseSet<Operation *> chainOps;
+    collectChainOps(whileOp, compOp, chainOps);
+    plan.compOp[origIdx] = compOp;
+    plan.chainOps[origIdx] = chainOps;
+
+    for (int blockId : blockIdsInOrder) {
+      plan.newArgDescriptors.push_back({blockId, nextNewArgIdx++, origIdx});
     }
+  }
+  // Pre-size clonedPerBlock so cloneWhileBlockChains can assign by index
+  // without a placeholder. Slots for cond-used origIdx with no compOp stay
+  // null and are never read (newArgDescriptors skips them).
+  unsigned numOrigIds = descIdx;
+  for (int blockId : blockIdsInOrder) {
+    plan.clonedPerBlock[blockId].resize(numOrigIds);
+  }
+  return plan;
+}
 
-    if (failed(processSharedIterArgsInForOp(forOp, info))) {
-      return WalkResult::interrupt();
+// Clones update chains into newAfter per (block_id, origIdx); reuses compOp/
+// chainOps from before migrateBody; records into plan.clonedPerBlock.
+static LogicalResult cloneWhileBlockChains(scf::WhileOp newWhileOp,
+                                           WhileIterArgClonePlan &plan) {
+  Block *newAfter = newWhileOp.getAfterBody();
+  // Derive (blockId, origIdx) -> newArgIdx and the ordered blockId list from
+  // newArgDescriptors (first-appearance order matches getBlockIdsInOrder).
+  llvm::DenseMap<std::pair<int, unsigned>, unsigned> blockOrigToNewArg;
+  llvm::DenseSet<int> seenBlocks;
+  SmallVector<int> blockIdsInOrder;
+  for (auto &desc : plan.newArgDescriptors) {
+    int blockId;
+    unsigned newArgIdx, origIdx;
+    std::tie(blockId, newArgIdx, origIdx) = desc;
+    blockOrigToNewArg[{blockId, origIdx}] = newArgIdx;
+    if (seenBlocks.insert(blockId).second)
+      blockIdsInOrder.push_back(blockId);
+  }
+
+  // Iterate unique origIdx in first-appearance order;
+  // planWhileIterArgDescriptors pushes (blockId, origIdx) in originalIndices x
+  // blockIdsInOrder order.
+  llvm::DenseSet<unsigned> seenOrig;
+  for (auto &desc : plan.newArgDescriptors) {
+    unsigned origIdx;
+    std::tie(std::ignore, std::ignore, origIdx) = desc;
+    if (!seenOrig.insert(origIdx).second)
+      continue;
+    Operation *compOp = plan.compOp.lookup(origIdx);
+    if (!compOp)
+      continue;
+    const llvm::DenseSet<Operation *> &chainOps = plan.chainOps.lookup(origIdx);
+
+    for (int blockId : blockIdsInOrder) {
+      unsigned newArgIdx = blockOrigToNewArg.lookup({blockId, origIdx});
+      Value cloned;
+      if (failed(cloneUpdateChainForWhileBlock(newWhileOp, newAfter, compOp,
+                                               blockId, origIdx, newArgIdx,
+                                               chainOps, cloned))) {
+        return failure();
+      }
+      plan.clonedPerBlock[blockId][plan.posInClonedVec.lookup(origIdx)] =
+          cloned;
     }
-    return WalkResult::advance();
-  });
+  }
+  return success();
+}
 
-  if (result.wasInterrupted()) {
+// Per (block_id, cond-used iter_arg) pair: clone the update chain into the
+// new after body, extend the scf.while, and record the new arg mapping.
+LogicalResult
+ProcessArgsPass::processWhileIterArgsInWhileOp(scf::WhileOp whileOp,
+                                               ControlFlowConditionInfo *info) {
+  auto it = originalWhileIterArgIndices.find(whileOp);
+  if (it == originalWhileIterArgIndices.end())
+    return success();
+  const SmallVector<unsigned> &originalIndices = it->second;
+  llvm::DenseSet<unsigned> condUsed =
+      collectConditionUsedIterArgIndices(whileOp, originalIndices);
+  if (condUsed.empty())
+    return success();
+
+  WhileIterArgClonePlan plan =
+      planWhileIterArgDescriptors(whileOp, originalIndices, condUsed);
+  if (plan.newArgDescriptors.empty()) {
+    return success();
+  }
+
+  // Gather origIdx from plan.newArgDescriptors in order — same shadowing rule
+  // as shared-iter-args (one new init per orig to shadow), routes through
+  // createMainLoopOpWithExtras.
+  SmallVector<Value> extraInitArgs;
+  extraInitArgs.reserve(plan.newArgDescriptors.size());
+  ValueRange origInits = whileOp.getInits();
+  for (auto &desc : plan.newArgDescriptors) {
+    unsigned origIdx;
+    std::tie(std::ignore, std::ignore, origIdx) = desc;
+    extraInitArgs.push_back(origInits[origIdx]);
+  }
+  auto newWhileOp =
+      cast<scf::WhileOp>(createMainLoopOpWithExtras(whileOp, extraInitArgs));
+  migrateWhileBodies(whileOp, newWhileOp);
+  if (failed(cloneWhileBlockChains(newWhileOp, plan))) {
+    return failure();
+  }
+  buildNewWhileCondition(whileOp, newWhileOp);
+  // Pre-compute extra yield values from plan.clonedPerBlock in
+  // newArgDescriptors order; appended to new scf.yield operands by
+  // buildNewYieldOp.
+  SmallVector<Value> extraYieldValues;
+  extraYieldValues.reserve(plan.newArgDescriptors.size());
+  for (auto &desc : plan.newArgDescriptors) {
+    int blockId;
+    unsigned newArgIdx, origIdx;
+    std::tie(blockId, newArgIdx, origIdx) = desc;
+    extraYieldValues.push_back(plan.clonedPerBlock.lookup(
+        blockId)[plan.posInClonedVec.lookup(origIdx)]);
+  }
+  if (failed(buildNewYieldOp(whileOp.getAfterBody(), newWhileOp.getAfterBody(),
+                             newWhileOp, extraYieldValues))) {
+    return failure();
+  }
+  // Mirror (block_id, new_arg_idx) -> orig_idx into localWhileBlockArgMap and
+  // info->whileBlockArgMap so downstream passes can resolve newWhileOp's args.
+  for (auto &desc : plan.newArgDescriptors) {
+    int blockId;
+    unsigned newArgIdx, origIdx;
+    std::tie(blockId, newArgIdx, origIdx) = desc;
+    localWhileBlockArgMap[newWhileOp][blockId][newArgIdx] = (int)origIdx;
+    if (info) {
+      info->whileBlockArgMap[newWhileOp][blockId][newArgIdx] = (int)origIdx;
+    }
+  }
+  if (failed(replaceMainLoopOpAndErase(whileOp, newWhileOp, info))) {
     return failure();
   }
   return success();
@@ -514,14 +837,24 @@ void ProcessArgsPass::runOnOperation() {
     return;
   }
 
-  LDBG("before processArgs:\n" << module << "\n");
+  LDBG("before processArgs:\n" << module);
+
+  // 1. While-specific decoupling: snapshot original iter_args; per scf.while,
+  //    clone cond-used iter_arg update chains per block_id; record in
+  //    info->whileBlockArgMap.
+  if (failed(updateIndependentCondsInWhileBlocks(module))) {
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+    return;
+  }
 
+  // 2. Process shared iter_args (adds per-block clones for args shared
+  //    across block_ids). Uses originalWhileIterArgIndices captured above.
   if (failed(processSharedIterArgs(module))) {
     CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
     return;
   }
 
-  LDBG("after processArgs:\n" << module << "\n");
+  LDBG("after processArgs:\n" << module);
 }
 
 namespace mlir {
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp
deleted file mode 100644
index 7f72f2e59e..0000000000
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp
+++ /dev/null
@@ -1,641 +0,0 @@
-/*
- * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
- * THE SOFTWARE.
- */
-
-#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h"
-#include "ascend/include/DynamicCVPipeline/Common/Utils.h"
-#include "bishengir/Dialect/HIVM/IR/HIVM.h"
-#include "bishengir/Dialect/Scope/IR/Scope.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/Debug.h"
-
-static constexpr const char *DEBUG_TYPE = "UpdateForOps";
-static constexpr int kPipeSFlagId = 15;
-static constexpr const char *kSsbufferMainLoop = "ssbuffer.main_loop";
-static constexpr const char *kSsbufferIf = "ssbuffer.if";
-#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
-#define LDBG(...)                                                              \
-  LLVM_DEBUG({                                                                 \
-    DBGS();                                                                    \
-    llvm::outs() << __VA_ARGS__;                                               \
-    llvm::outs() << "\n";                                                      \
-  })
-
-using llvm::SmallVector;
-using namespace mlir;
-using namespace triton;
-using namespace hivm;
-
-// Replace old block arguments with new ones
-static LogicalResult replaceBlockArguments(Block *oldBlock, Block *newBlock) {
-  if (!oldBlock || !newBlock) {
-    LDBG("[Error]: oldBlock or newBlock is null\n");
-    return failure();
-  }
-
-  unsigned totalArgs = oldBlock->getNumArguments();
-
-  for (unsigned i = 0; i < totalArgs; ++i) {
-    oldBlock->getArgument(i).replaceAllUsesWith(newBlock->getArgument(i));
-  }
-
-  return success();
-}
-
-// Collect forOps that need processing based on info
-static SmallVector<scf::ForOp>
-collectForOpsToProcess(ModuleOp module,
-                       const llvm::DenseMap<scf::ForOp, int> &numInfo) {
-  SmallVector<scf::ForOp> forOps;
-
-  module.walk([&](scf::ForOp forOp) {
-    if (numInfo.count(forOp)) {
-      forOps.push_back(forOp);
-    }
-  });
-
-  return forOps;
-}
-
-// Create new yield operands: original yield ops + extra args from new block
-static SmallVector<Value> createNewYieldOperands(scf::YieldOp oldYield,
-                                                 unsigned oldNumArgs,
-                                                 Block *newBlock,
-                                                 int numExtraArgs) {
-  SmallVector<Value> newYieldOperands;
-
-  for (unsigned i = 0; i < oldNumArgs; ++i) {
-    newYieldOperands.push_back(oldYield.getOperand(i));
-  }
-
-  for (int i = 0; i < numExtraArgs; ++i) {
-    newYieldOperands.push_back(newBlock->getArgument(1 + oldNumArgs + i));
-  }
-
-  return newYieldOperands;
-}
-
-// Derive block counters from ssbuffer.if attributes when info is not
-// pre-populated
-LogicalResult
-UpdateForOpsPass::deriveBlockCountersFromIfOps(ModuleOp module,
-                                               ControlFlowConditionInfo *info) {
-  if (!info) {
-    LDBG("[Error]: info is null\n");
-    return failure();
-  }
-
-  module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr("ssbuffer.main_loop")) {
-      return WalkResult::advance();
-    }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
-      return WalkResult::interrupt();
-    }
-
-    // Collect unique ssbuffer.if values in this for op's body
-    llvm::DenseSet<int> ifBlockIds;
-
-    forOp.walk([&](Operation *innerOp) {
-      if (auto ifAttr = innerOp->getAttrOfType<IntegerAttr>("ssbuffer.if")) {
-        ifBlockIds.insert(ifAttr.getInt());
-      }
-    });
-
-    if (!ifBlockIds.empty()) {
-      info->blockCounterNums[forOp] = ifBlockIds.size();
-    }
-    return WalkResult::advance();
-  });
-
-  return success();
-}
-
-// Create new for op with extra iter args and migrate body
-static scf::ForOp
-createForOpAndMigrateBody(scf::ForOp oldForOp, int numExtraArgs,
-                          const SmallVector<Value> &extraInitArgs) {
-  if (numExtraArgs < 0) {
-    LDBG("[Error]: invalid numExtraArgs " << numExtraArgs << "\n");
-    return scf::ForOp();
-  }
-  if (numExtraArgs == 0)
-    return oldForOp;
-  if ((int)extraInitArgs.size() != numExtraArgs) {
-    LDBG("[Error]: extraInitArgs size " << extraInitArgs.size()
-                                        << " != numExtraArgs " << numExtraArgs
-                                        << "\n");
-    return scf::ForOp();
-  }
-
-  OpBuilder builder(oldForOp);
-  // Create new for op with extra iter args
-  SmallVector<Value> newInitArgs(oldForOp.getInitArgs().begin(),
-                                 oldForOp.getInitArgs().end());
-  llvm::append_range(newInitArgs, extraInitArgs);
-
-  scf::ForOp newForOp = builder.create<scf::ForOp>(
-      oldForOp.getLoc(), oldForOp.getLowerBound(), oldForOp.getUpperBound(),
-      oldForOp.getStep(), newInitArgs);
-
-  for (auto &attr : oldForOp->getAttrs())
-    newForOp->setAttr(attr.getName(), attr.getValue());
-
-  // Migrate body
-  Block *oldBlock = oldForOp.getBody();
-  Block *newBlock = newForOp.getBody();
-
-  if (failed(replaceBlockArguments(oldBlock, newBlock))) {
-    newForOp.erase();
-    return scf::ForOp();
-  }
-
-  for (Operation &op :
-       llvm::make_early_inc_range(oldBlock->without_terminator()))
-    op.moveBefore(newBlock, newBlock->end());
-
-  auto oldYield = cast<scf::YieldOp>(oldBlock->getTerminator());
-  SmallVector<Value> newYieldOperands = createNewYieldOperands(
-      oldYield, oldForOp.getNumRegionIterArgs(), newBlock, numExtraArgs);
-
-  builder.setInsertionPointToEnd(newBlock);
-  builder.create<scf::YieldOp>(newForOp.getLoc(), newYieldOperands);
-  oldYield.erase();
-
-  return newForOp;
-}
-
-static LogicalResult replaceForOpUsesAndErase(scf::ForOp oldForOp,
-                                              scf::ForOp newForOp) {
-  if (oldForOp.getNumResults() > 0) {
-    SmallVector<Value> newResults;
-    for (unsigned i = 0; i < oldForOp.getNumResults(); ++i) {
-      if (oldForOp.getResult(i).getType() != newForOp.getResult(i).getType()) {
-        LDBG("[Error]: result type mismatch at index " << i << "\n");
-        return failure();
-      }
-      newResults.push_back(newForOp.getResult(i));
-    }
-    oldForOp.replaceAllUsesWith(newResults);
-  }
-
-  oldForOp.erase();
-  return success();
-}
-
-LogicalResult extendForOpWithExtraArgs(scf::ForOp oldForOp,
-                                       ControlFlowConditionInfo *info) {
-  int numBlockCounters = info->blockCounterNums[oldForOp];
-  int numInnerDepConds = info->intraCoreDependentMap[oldForOp].size();
-  int totalExtraArgs = numBlockCounters + numInnerDepConds;
-
-  int numTensorIterArgs = 0;
-  // Record the number of consumers for each tensor iter_args (the number of
-  // parameters to be created)
-  llvm::DenseMap<Value, int> tensorIterArgNumConsumers;
-  // First, copy the depsVec out to avoid iterator invalidation later
-  llvm::SmallVector<TensorIterArgIfOpRelation> depsVecCopy;
-  auto tensorIterArgDepsIt = info->tensorIterArgDepsMap.find(oldForOp);
-  if (tensorIterArgDepsIt != info->tensorIterArgDepsMap.end()) {
-    depsVecCopy = tensorIterArgDepsIt->second; // Make a copy
-    for (auto &entry : depsVecCopy) {
-      Value iterArg = entry.iterArg;
-      int numConsumers = entry.consumers.size();
-      tensorIterArgNumConsumers[iterArg] = numConsumers;
-      numTensorIterArgs += numConsumers;
-    }
-  }
-
-  totalExtraArgs += numTensorIterArgs;
-  if (totalExtraArgs == 0) {
-    return success();
-  }
-
-  OpBuilder builder(oldForOp);
-  SmallVector<Value> extraInitArgs;
-  for (int i = 0; i < numBlockCounters; ++i)
-    extraInitArgs.push_back(oldForOp.getLowerBound());
-  for (int i = 0; i < numInnerDepConds; ++i)
-    extraInitArgs.push_back(builder.create<arith::ConstantOp>(
-        oldForOp.getLoc(), builder.getI32Type(), builder.getI32IntegerAttr(0)));
-  // Add an initial value (1) for the new parameter iter_arg of tensor
-  for (int i = 0; i < numTensorIterArgs; ++i)
-    extraInitArgs.push_back(builder.create<arith::ConstantOp>(
-        oldForOp.getLoc(), builder.getI32Type(), builder.getI32IntegerAttr(1)));
-
-  scf::ForOp newForOp =
-      createForOpAndMigrateBody(oldForOp, totalExtraArgs, extraInitArgs);
-  if (!newForOp) {
-    return failure();
-  }
-
-  unsigned baseIdx = oldForOp.getNumRegionIterArgs();
-  if (numBlockCounters > 0) {
-    SmallVector<int> indices;
-    for (int j = 0; j < numBlockCounters; ++j)
-      indices.push_back(baseIdx + j);
-    info->blockCounters.erase(oldForOp);
-    info->blockCounters[newForOp] = indices;
-  }
-
-  if (numInnerDepConds > 0) {
-    SmallVector<int> indices;
-    for (int j = 0; j < numInnerDepConds; ++j)
-      indices.push_back(baseIdx + numBlockCounters + j);
-    info->innerDepConds.erase(oldForOp);
-    info->innerDepConds[newForOp] = indices;
-  }
-
-  // Record the index of the new parameter iter_arg for the tensor and update
-  // the corresponding map
-  if (numTensorIterArgs > 0) {
-    unsigned tensorBaseIdx = baseIdx + numBlockCounters + numInnerDepConds;
-    auto &newIndicesMap = info->tensorIterArgIndicesMap[newForOp];
-
-    unsigned currentIdx = tensorBaseIdx;
-    for (auto &entry : depsVecCopy) {
-      Value iterArg = entry.iterArg;
-      int numConsumers = entry.consumers.size();
-      SmallVector<int> indices;
-      for (int j = 0; j < numConsumers; ++j) {
-        indices.push_back(currentIdx++);
-      }
-      newIndicesMap[iterArg] = indices;
-    }
-
-    info->tensorIterArgIndicesMap.erase(oldForOp);
-    info->tensorIterArgDepsMap[newForOp] = std::move(depsVecCopy);
-    info->tensorIterArgDepsMap.erase(oldForOp);
-  }
-
-  if (info->intraCoreDependentMap.count(oldForOp)) {
-    info->intraCoreDependentMap[newForOp] =
-        info->intraCoreDependentMap[oldForOp];
-    info->intraCoreDependentMap.erase(oldForOp);
-  }
-
-  return replaceForOpUsesAndErase(oldForOp, newForOp);
-}
-
-// Add block counter and inner dependency condition iter args to for ops
-LogicalResult UpdateForOpsPass::addBlockCountersAndInnerDepConds(
-    ModuleOp module, ControlFlowConditionInfo *info) {
-  llvm::DenseSet<scf::ForOp> forOpsToProcess;
-
-  for (auto &p : info->blockCounterNums) {
-    if (p.second < 0) {
-      LDBG("[Error]: invalid blockCounterNum " << p.second << "\n");
-      return failure();
-    }
-    forOpsToProcess.insert(p.first);
-  }
-  for (auto &p : info->tensorIterArgDepsMap) {
-    forOpsToProcess.insert(p.first);
-  }
-
-  for (scf::ForOp forOp : forOpsToProcess) {
-    if (failed(extendForOpWithExtraArgs(forOp, info)))
-      return failure();
-  }
-
-  return success();
-}
-
-// Insert sync ops inside a forOp: wait at start, set before yield
-static LogicalResult insertSyncOpsInsideForOp(Block *forBody, Location loc,
-                                              hivm::TCoreTypeAttr coreType,
-                                              PipeAttr setPipe,
-                                              PipeAttr waitPipe, int waitFlagId,
-                                              int setFlagId) {
-  Operation *forTerminator = forBody->getTerminator();
-  if (!forTerminator) {
-    return failure();
-  }
-
-  // Insert wait at for loop start
-  OpBuilder insertionBuilder(&forBody->front());
-  auto waitFlagAttr = insertionBuilder.getIntegerAttr(
-      insertionBuilder.getI64Type(), waitFlagId);
-  insertionBuilder.create<SyncBlockWaitOp>(loc, coreType, setPipe, waitPipe,
-                                           waitFlagAttr);
-
-  // Insert set before yield
-  OpBuilder setBuilder(forTerminator);
-  auto setFlagAttr =
-      setBuilder.getIntegerAttr(setBuilder.getI64Type(), setFlagId);
-  setBuilder.setInsertionPoint(forTerminator);
-  setBuilder.create<SyncBlockSetOp>(loc, coreType, setPipe, waitPipe,
-                                    setFlagAttr);
-
-  return success();
-}
-
-// Insert set before or wait after a forOp
-static LogicalResult insertSetOrWaitForForOp(scf::ForOp forOp, Location loc,
-                                             hivm::TCoreTypeAttr coreType,
-                                             PipeAttr setPipe,
-                                             PipeAttr waitPipe, int flagId,
-                                             bool isBefore) {
-  OpBuilder builder(forOp);
-  auto flagAttr = builder.getIntegerAttr(builder.getI64Type(), flagId);
-  if (isBefore) {
-    builder.create<SyncBlockSetOp>(loc, coreType, setPipe, waitPipe, flagAttr);
-  } else {
-    builder.setInsertionPointAfter(forOp);
-    builder.create<SyncBlockWaitOp>(loc, coreType, setPipe, waitPipe, flagAttr);
-  }
-  return success();
-}
-
-// Insert PIPE_S for a main_loop forOp based on forOp type and scope type
-static LogicalResult
-insertPipeSForMainLoopForOp(scf::ForOp forOp, scope::ScopeOp scopeOp,
-                            bool isScopeCube, bool isScopeVector,
-                            PipeAttr setPipe, PipeAttr waitPipe, int flagId) {
-  Block *forBody = &forOp.getRegion().front();
-  Location loc = forOp.getLoc();
-  bool isVectorFirst = forOp->hasAttr("ssbuffer.vector_first");
-  auto cubeType =
-      hivm::TCoreTypeAttr::get(forOp.getContext(), hivm::TCoreType::CUBE);
-  auto vectorType =
-      hivm::TCoreTypeAttr::get(forOp.getContext(), hivm::TCoreType::VECTOR);
-
-  if (isVectorFirst) {
-    if (isScopeCube) {
-      // vector_first + CUBE: before forop (SET), inside (WAIT/SET)
-      if (failed(insertSetOrWaitForForOp(forOp, loc, cubeType, setPipe,
-                                         waitPipe, flagId, true))) {
-        return failure();
-      }
-      if (failed(insertSyncOpsInsideForOp(forBody, loc, cubeType, setPipe,
-                                          waitPipe, flagId, flagId))) {
-        return failure();
-      }
-    } else if (isScopeVector) {
-      // vector_first + VECTOR: inside (WAIT/SET), after forop (WAIT)
-      if (failed(insertSyncOpsInsideForOp(forBody, loc, vectorType, setPipe,
-                                          waitPipe, flagId, flagId))) {
-        return failure();
-      }
-      if (failed(insertSetOrWaitForForOp(forOp, loc, vectorType, setPipe,
-                                         waitPipe, flagId, false))) {
-        return failure();
-      }
-    }
-  } else {
-    // cube_first (including default when neither attribute is present)
-    if (isScopeCube) {
-      // cube_first + CUBE: inside (WAIT/SET), after forop (WAIT)
-      if (failed(insertSyncOpsInsideForOp(forBody, loc, cubeType, setPipe,
-                                          waitPipe, flagId, flagId))) {
-        return failure();
-      }
-      if (failed(insertSetOrWaitForForOp(forOp, loc, cubeType, setPipe,
-                                         waitPipe, flagId, false))) {
-        return failure();
-      }
-    } else if (isScopeVector) {
-      // cube_first + VECTOR: before forop (SET), inside (WAIT/SET)
-      if (failed(insertSetOrWaitForForOp(forOp, loc, vectorType, setPipe,
-                                         waitPipe, flagId, true))) {
-        return failure();
-      }
-      if (failed(insertSyncOpsInsideForOp(forBody, loc, vectorType, setPipe,
-                                          waitPipe, flagId, flagId))) {
-        return failure();
-      }
-    }
-  }
-  return success();
-}
-
-LogicalResult UpdateForOpsPass::insertInterCorePipeS(ModuleOp module) {
-  auto cubeCoreType =
-      hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE);
-  auto vectorCoreType =
-      hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::VECTOR);
-  auto setPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S);
-  auto waitPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S);
-
-  WalkResult result = module.walk([&](scope::ScopeOp scopeOp) -> WalkResult {
-    auto scopeTypeAttr =
-        scopeOp->getAttrOfType<hivm::TCoreTypeAttr>("hivm.tcore_type");
-    if (!scopeTypeAttr) {
-      return WalkResult::advance();
-    }
-
-    bool isScopeCube = (scopeTypeAttr == cubeCoreType);
-    bool isScopeVector = (scopeTypeAttr == vectorCoreType);
-
-    WalkResult innerResult = scopeOp.walk([&](scf::ForOp forOp) -> WalkResult {
-      if (!forOp->hasAttr("ssbuffer.main_loop")) {
-        return WalkResult::advance();
-      }
-      if (failed(insertPipeSForMainLoopForOp(forOp, scopeOp, isScopeCube,
-                                             isScopeVector, setPipeType,
-                                             waitPipeType, kPipeSFlagId))) {
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
-    });
-
-    if (innerResult.wasInterrupted()) {
-      return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  });
-
-  return result.wasInterrupted() ? failure() : success();
-}
-
-// Analyze the producer/consumer relationship between the tensor type iter_args
-// in the main_loop and ssbuffer.if
-LogicalResult UpdateForOpsPass::analyzeTensorIterArgDependencies(
-    ModuleOp module, ControlFlowConditionInfo *info) {
-  bool failed = false;
-  module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr(kSsbufferMainLoop)) {
-      return WalkResult::advance();
-    }
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n");
-      failed = true;
-      return WalkResult::interrupt();
-    }
-
-    LDBG("Analyzing main_loop forOp: " << forOp << "\n");
-
-    for (auto iterArg : forOp.getRegionIterArgs()) {
-      if (!mlir::isa<TensorType>(iterArg.getType())) {
-        continue;
-      }
-
-      LDBG("Found tensor type iter_arg: " << iterArg << "\n");
-      scf::IfOp producerIfOp = nullptr;
-      llvm::SmallVector<scf::IfOp> consumerIfOps;
-
-      for (auto &use : iterArg.getUses()) {
-        Operation *user = use.getOwner();
-        scf::IfOp ifOp = nullptr;
-        Operation *curr = user;
-        while (curr && curr != forOp.getOperation()) {
-          if (auto currIf = dyn_cast<scf::IfOp>(curr)) {
-            if (currIf->hasAttr(kSsbufferIf)) {
-              ifOp = currIf;
-              break;
-            }
-          }
-          curr = curr->getParentOp();
-        }
-
-        if (!ifOp) {
-          LDBG("Use of tensor iter_arg "
-               << iterArg << " is not inside any ssbuffer.if op." << "\n");
-          continue;
-        }
-
-        // Only the direct terminator of this ssbuffer.if is a producer.
-        // Nested if/for/while yields that forward iter_arg are consumers.
-        bool isProducer = isa<scf::YieldOp>(user) &&
-                          user->getParentOp() == ifOp.getOperation();
-
-        // Check and update status of ifOp
-        if (isProducer) {
-          if (producerIfOp && producerIfOp != ifOp) {
-            // Found a different producer ifOp! This is an error.
-            LDBG("[Error]: tensor iter_arg "
-                 << iterArg << " has multiple different producers!\n");
-            LDBG("Existing producer: " << producerIfOp << "\n");
-            LDBG("New producer: " << ifOp << "\n");
-            failed = true;
-            return WalkResult::interrupt();
-          }
-          if (!producerIfOp) {
-            // First producer, or upgrade from consumer
-            auto it = llvm::find(consumerIfOps, ifOp);
-            if (it != consumerIfOps.end()) {
-              consumerIfOps.erase(it);
-              LDBG("This ifOp was consumer, now updated to producer: " << ifOp
-                                                                       << "\n");
-            } else {
-              LDBG("Found producer ifOp (first time): " << ifOp << "\n");
-            }
-            producerIfOp = ifOp;
-          }
-          // Else: already is this producer, do nothing
-        } else {
-          // isConsumer
-          if (producerIfOp == ifOp) {
-            // Already a producer, even if current use is consumer, do nothing
-            continue;
-          }
-          if (!llvm::is_contained(consumerIfOps, ifOp)) {
-            consumerIfOps.push_back(ifOp);
-            LDBG("Found consumer ifOp (first time): " << ifOp << "\n");
-          }
-          // Else: already a consumer, do nothing
-        }
-      }
-      // Check: must have both producers AND consumers
-      if (!producerIfOp || consumerIfOps.empty()) {
-        LDBG("tensor iter_arg " << iterArg << " has only "
-                                << (!producerIfOp ? "consumers" : "producers")
-                                << ", skipped\n");
-        continue;
-      }
-      TensorIterArgIfOpRelation relation;
-      relation.iterArg = iterArg;
-      relation.producer = producerIfOp;
-      relation.consumers = consumerIfOps;
-
-      info->tensorIterArgDepsMap[forOp].push_back(relation);
-      LDBG("Recorded tensor iter_arg dependency: "
-           << iterArg << " has 1 producer, " << relation.consumers.size()
-           << " consumers\n");
-    }
-
-    return WalkResult::advance();
-  });
-
-  return failed ? failure() : success();
-}
-
-void UpdateForOpsPass::runOnOperation() {
-  ModuleOp module = getOperation();
-
-  if (CVPipeline::hasFallbackAttr(module)) {
-    return;
-  }
-
-  LDBG("before updateForOps:\n" << module << "\n");
-
-  // Use provided info, or create a local one if not available
-  ControlFlowConditionInfo localInfo;
-  ControlFlowConditionInfo *infoToUse = info ? info : &localInfo;
-
-  // Analyze the dependencies of the tensor type iter_args in the main_loop with
-  // the ssbuffer.if ops
-  if (failed(analyzeTensorIterArgDependencies(module, infoToUse))) {
-    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
-    return;
-  }
-
-  // Derive block counters from ssbuffer.if if blockCounterNums is empty
-  if (infoToUse->blockCounterNums.empty()) {
-    if (failed(deriveBlockCountersFromIfOps(module, infoToUse))) {
-      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
-      return;
-    }
-  }
-
-  // Update for ops iter_args for block counters and inner dependency conditions
-  if (infoToUse && (!infoToUse->blockCounterNums.empty() ||
-                    !infoToUse->intraCoreDependentMap.empty() ||
-                    !infoToUse->tensorIterArgDepsMap.empty()))
-    if (failed(addBlockCountersAndInnerDepConds(module, infoToUse))) {
-      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
-      return;
-    }
-
-  // Insert PIPE_S inter-core synchronization
-  if (failed(insertInterCorePipeS(module))) {
-    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
-    return;
-  }
-
-  LDBG("after updateForOps:\n" << module << "\n");
-}
-
-namespace mlir {
-namespace triton {
-
-std::unique_ptr<OperationPass<ModuleOp>> createUpdateForOpsPass() {
-  return std::make_unique<UpdateForOpsPass>();
-}
-
-} // namespace triton
-} // namespace mlir
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
index 084b792f09..dcea9545f6 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
@@ -83,7 +83,7 @@ static scf::ForOp getOtherScopeMainloop(ModuleOp module, bool currentIsCube,
     // Determine this scope's type
     bool scopeIsCube = false;
     bool scopeIsVector = false;
-    if (failed(triton::getScopeType(scopeOp, scopeIsCube, scopeIsVector))) {
+    if (failed(getScopeType(scopeOp, scopeIsCube, scopeIsVector))) {
       ret = -1;
       LDBG("failed to get ScopeOp core type!");
       return mlir::WalkResult::interrupt();
@@ -479,8 +479,7 @@ std::pair<int, int> UpdateLoopIterTimesPass::calculateCrossDepsFactor(
     }
     currentScope = currentScope->getParentOp();
   }
-  if (failed(
-          triton::getScopeType(currentScope, currentIsCube, currentIsVector))) {
+  if (failed(getScopeType(currentScope, currentIsCube, currentIsVector))) {
     LDBG("Current forOp is not in a valid cube or vector scope!");
     return {-1, -1};
   }
@@ -895,7 +894,7 @@ int UpdateLoopIterTimesPass::GetMainLoopIdToLoopOpMap(
     // Determine if it's CUBE or VECTOR
     bool isCube = false;
     bool isVector = false;
-    if (failed(triton::getScopeType(scopeOp, isCube, isVector))) {
+    if (failed(getScopeType(scopeOp, isCube, isVector))) {
       ret = -1;
       LDBG("mlir do not processed by split mix kernel!");
       return mlir::WalkResult::interrupt();
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
new file mode 100644
index 0000000000..838d4339a2
--- /dev/null
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
@@ -0,0 +1,740 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.h"
+#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h"
+#include "ascend/include/DynamicCVPipeline/Common/Utils.h"
+#include "bishengir/Dialect/HIVM/IR/HIVM.h"
+#include "bishengir/Dialect/Scope/IR/Scope.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+
+static constexpr const char *DEBUG_TYPE = "UpdateLoopOps";
+static constexpr int kPipeSFlagId = 15;
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define LDBG(...)                                                              \
+  LLVM_DEBUG({                                                                 \
+    DBGS();                                                                    \
+    llvm::dbgs() << __VA_ARGS__;                                               \
+    llvm::dbgs() << "\n";                                                      \
+  })
+
+using namespace mlir;
+using namespace triton;
+using namespace hivm;
+using namespace CVPipeline;
+
+// Derive block counters from ssbuffer.if for the main-loop op (scf.for or
+// scf.while with ssbuffer.main_loop). blockCounterNums keyed on Operation*.
+LogicalResult UpdateLoopOpsPass::deriveBlockCountersFromIfOps(
+    ModuleOp module, ControlFlowConditionInfo *info) {
+  if (!info) {
+    LDBG("[Error]: info is null");
+    return failure();
+  }
+
+  module.walk([&](Operation *op) -> WalkResult {
+    if (!isMainLoopOp(op)) {
+      return WalkResult::advance();
+    }
+
+    int numIfBlockIds = countUniqueIfBlockIds(op);
+    if (numIfBlockIds > 0) {
+      info->blockCounterNums[op] = numIfBlockIds;
+    }
+    return WalkResult::advance();
+  });
+
+  return success();
+}
+
+// Creates new scf.for with extras and migrates old body. Yields each new
+// iter_arg forward (counts on themselves, deps/tensors via buildNewYieldOp).
+static scf::ForOp
+createForOpAndMigrateBody(scf::ForOp oldForOp,
+                          const llvm::SmallVector<Value> &extraInitArgs) {
+  scf::ForOp newForOp = createNewForOpWithExtras(oldForOp, extraInitArgs);
+  if (newForOp == oldForOp) {
+    return oldForOp;
+  }
+
+  Block *oldBlock = oldForOp.getBody();
+  Block *newBlock = newForOp.getBody();
+  migrateBody(oldBlock, newBlock);
+
+  // Append extra iter_arg yields (new block args at [1+oldNumArgs, +numExtra)).
+  unsigned oldNumArgs = oldForOp.getNumRegionIterArgs();
+  llvm::SmallVector<Value> extras;
+  extras.reserve(extraInitArgs.size());
+  for (size_t i = 0; i < extraInitArgs.size(); ++i) {
+    extras.push_back(newBlock->getArgument(1 + oldNumArgs + i));
+  }
+  if (failed(buildNewYieldOp(oldBlock, newBlock, newForOp, extras))) {
+    newForOp.erase();
+    return scf::ForOp();
+  }
+
+  return newForOp;
+}
+
+// Splices the new main-loop op in place of the old one: type-checks results,
+// redirects external uses to new results, erases old op.
+static LogicalResult replaceMainLoopOpUsesAndErase(Operation *oldOp,
+                                                   Operation *newOp) {
+  if (oldOp->getNumResults() > 0) {
+    for (unsigned i = 0; i < oldOp->getNumResults(); ++i) {
+      if (oldOp->getResult(i).getType() != newOp->getResult(i).getType()) {
+        LDBG("[Error]: main_loop op result type mismatch at index " << i);
+        return failure();
+      }
+    }
+  }
+
+  replaceOpResultUses(oldOp, newOp);
+  oldOp->erase();
+  return success();
+}
+
+// Appends the new iter_args (new before-block args at
+// [numOriginal..numOriginal+numExtra-1]) to the scf.condition's forwarded
+// values.
+static void extendWhileCondition(scf::ConditionOp oldCond, Block *newBefore,
+                                 unsigned numOriginalIterArgs,
+                                 unsigned numExtraArgs) {
+  llvm::SmallVector<Value> newCondValues(oldCond.getArgs().begin(),
+                                         oldCond.getArgs().end());
+  for (unsigned i = 0; i < numExtraArgs; ++i) {
+    newCondValues.push_back(newBefore->getArgument(numOriginalIterArgs + i));
+  }
+  oldCond.getArgsMutable().assign(newCondValues);
+}
+
+// Appends new iter_args (new after-block args at
+// [numOriginal..numOriginal+numExtra-1]) to scf.yield. Counters forward
+// themselves.
+static void extendWhileYield(scf::YieldOp oldYield, Block *newAfter,
+                             unsigned numOriginalIterArgs,
+                             unsigned numExtraArgs) {
+  llvm::SmallVector<Value> newYieldOperands(oldYield.getOperands().begin(),
+                                            oldYield.getOperands().end());
+  for (unsigned i = 0; i < numExtraArgs; ++i) {
+    newYieldOperands.push_back(newAfter->getArgument(numOriginalIterArgs + i));
+  }
+  oldYield.getResultsMutable().assign(newYieldOperands);
+}
+
+// Migrates both before/after regions: migrateBody args+ops, move terminator,
+// extend with new iter_args.
+static LogicalResult migrateWhileRegions(scf::WhileOp oldWhileOp,
+                                         scf::WhileOp newWhileOp,
+                                         unsigned numOriginalIterArgs,
+                                         unsigned numExtraArgs) {
+  Block *oldBefore = oldWhileOp.getBeforeBody();
+  Block *newBefore = newWhileOp.getBeforeBody();
+  migrateBody(oldBefore, newBefore);
+  auto oldCond = cast<scf::ConditionOp>(oldBefore->getTerminator());
+  oldCond->moveBefore(newBefore, newBefore->end());
+  extendWhileCondition(oldCond, newBefore, numOriginalIterArgs, numExtraArgs);
+
+  Block *oldAfter = oldWhileOp.getAfterBody();
+  Block *newAfter = newWhileOp.getAfterBody();
+  migrateBody(oldAfter, newAfter);
+  auto oldYield = cast<scf::YieldOp>(oldAfter->getTerminator());
+  oldYield->moveBefore(newAfter, newAfter->end());
+  extendWhileYield(oldYield, newAfter, numOriginalIterArgs, numExtraArgs);
+
+  return success();
+}
+
+// Creates new scf.while with extras, then migrates old before/after regions.
+// condition/yield forward new iter_args (mirroring forOp layout).
+static scf::WhileOp
+createWhileOpAndMigrateBody(scf::WhileOp oldWhileOp,
+                            const llvm::SmallVector<Value> &extraInitArgs) {
+  scf::WhileOp newWhileOp =
+      createNewWhileOpWithExtras(oldWhileOp, extraInitArgs);
+  if (newWhileOp == oldWhileOp) {
+    return oldWhileOp;
+  }
+
+  unsigned numOriginalIterArgs = oldWhileOp.getInits().size();
+  if (failed(migrateWhileRegions(oldWhileOp, newWhileOp, numOriginalIterArgs,
+                                 extraInitArgs.size()))) {
+    newWhileOp.erase();
+    return scf::WhileOp();
+  }
+
+  return newWhileOp;
+}
+
+// Dispatches createForOpAndMigrateBody / createWhileOpAndMigrateBody by op
+// type. Returns nullptr if op is neither scf.for nor scf.while.
+static Operation *
+createMainLoopOpAndMigrateBody(Operation *oldLoopOp,
+                               const llvm::SmallVector<Value> &extraInitArgs) {
+  if (auto forOp = dyn_cast<scf::ForOp>(oldLoopOp))
+    return createForOpAndMigrateBody(forOp, extraInitArgs);
+  if (auto whileOp = dyn_cast<scf::WhileOp>(oldLoopOp))
+    return createWhileOpAndMigrateBody(whileOp, extraInitArgs);
+  return nullptr;
+}
+
+// Computes extra-arg counts (block counters / inner dep conds / tensor
+// iter_args) for `oldLoopOp`. Returns depsVecCopy so the caller can move it
+// into info later.
+static void computeMainLoopExtraArgs(
+    Operation *oldLoopOp, ControlFlowConditionInfo *info, int &numBlockCounters,
+    int &numInnerDepConds, int &numTensorIterArgs,
+    llvm::SmallVector<TensorIterArgIfOpRelation> &depsVecCopy) {
+  numBlockCounters = info->blockCounterNums.lookup(oldLoopOp);
+  numInnerDepConds = info->intraCoreDependentMap.count(oldLoopOp)
+                         ? (int)info->intraCoreDependentMap[oldLoopOp].size()
+                         : 0;
+
+  numTensorIterArgs = 0;
+  auto tensorIterArgDepsIt = info->tensorIterArgDepsMap.find(oldLoopOp);
+  if (tensorIterArgDepsIt != info->tensorIterArgDepsMap.end()) {
+    depsVecCopy = tensorIterArgDepsIt->second;
+    for (auto &entry : depsVecCopy) {
+      numTensorIterArgs += entry.consumers.size();
+    }
+  }
+}
+
+// Returns the index of the first extra iter_arg in the new op. ForOp excludes
+// IV; WhileOp shares arg layout between before/after.
+static unsigned getMainLoopBaseIdx(Operation *oldLoopOp, bool isWhile) {
+  return isWhile ? (unsigned)cast<scf::WhileOp>(oldLoopOp).getInits().size()
+                 : cast<scf::ForOp>(oldLoopOp).getNumRegionIterArgs();
+}
+
+// Appends initial values for extra iter_args: block counters from
+// blockCounterInitFn (forOp reuses getLowerBound(); whileOp creates a new
+// arith.constant per counter so each new iter_arg has a distinct SSA value),
+// dep conds from i32(0), tensor iter_args from i32(1).
+static void
+buildMainLoopExtraInitArgs(OpBuilder &builder, Location loc,
+                           llvm::function_ref<Value()> blockCounterInitFn,
+                           int numBlockCounters, int numInnerDepConds,
+                           int numTensorIterArgs,
+                           llvm::SmallVector<Value> &extraInitArgs) {
+  for (int i = 0; i < numBlockCounters; ++i) {
+    extraInitArgs.push_back(blockCounterInitFn());
+  }
+  for (int i = 0; i < numInnerDepConds; ++i) {
+    extraInitArgs.push_back(builder.create<arith::ConstantOp>(
+        loc, builder.getI32Type(), builder.getI32IntegerAttr(0)));
+  }
+  for (int i = 0; i < numTensorIterArgs; ++i) {
+    extraInitArgs.push_back(builder.create<arith::ConstantOp>(
+        loc, builder.getI32Type(), builder.getI32IntegerAttr(1)));
+  }
+}
+
+// Migrates block-counter / inner-dep-cond index ranges from oldLoopOp to newOp.
+static void recordMainLoopBlockCountersAndConds(
+    Operation *oldLoopOp, Operation *newOp, ControlFlowConditionInfo *info,
+    unsigned baseIdx, int numBlockCounters, int numInnerDepConds) {
+  if (numBlockCounters > 0) {
+    llvm::SmallVector<int> indices;
+    for (int j = 0; j < numBlockCounters; ++j)
+      indices.push_back(baseIdx + j);
+    info->blockCounters.erase(oldLoopOp);
+    info->blockCounters[newOp] = indices;
+  }
+
+  if (numInnerDepConds > 0) {
+    llvm::SmallVector<int> indices;
+    for (int j = 0; j < numInnerDepConds; ++j)
+      indices.push_back(baseIdx + numBlockCounters + j);
+    info->innerDepConds.erase(oldLoopOp);
+    info->innerDepConds[newOp] = indices;
+  }
+}
+
+// Migrates tensor iter_arg index ranges and moves depsVecCopy into
+// info->tensorIterArgDepsMap.
+static void recordMainLoopTensorIterArgs(
+    Operation *oldLoopOp, Operation *newOp, ControlFlowConditionInfo *info,
+    unsigned baseIdx, int numBlockCounters, int numInnerDepConds,
+    int numTensorIterArgs,
+    llvm::SmallVector<TensorIterArgIfOpRelation> &depsVecCopy) {
+  if (numTensorIterArgs == 0) {
+    return;
+  }
+  unsigned tensorBaseIdx = baseIdx + numBlockCounters + numInnerDepConds;
+  auto &newIndicesMap = info->tensorIterArgIndicesMap[newOp];
+
+  unsigned currentIdx = tensorBaseIdx;
+  for (auto &entry : depsVecCopy) {
+    llvm::SmallVector<int> indices;
+    for (int j = 0; j < (int)entry.consumers.size(); ++j) {
+      indices.push_back(currentIdx++);
+    }
+    newIndicesMap[entry.iterArg] = indices;
+  }
+
+  info->tensorIterArgIndicesMap.erase(oldLoopOp);
+  info->tensorIterArgDepsMap[newOp] = std::move(depsVecCopy);
+  info->tensorIterArgDepsMap.erase(oldLoopOp);
+}
+
+// Transfers intraCoreDependentMap (and WhileOp-only whileBlockArgMap) from
+// oldLoopOp to newOp so the maps survive the old op being erased.
+static void transferMainLoopInfoMaps(Operation *oldLoopOp, Operation *newOp,
+                                     ControlFlowConditionInfo *info,
+                                     bool isWhile) {
+  if (info->intraCoreDependentMap.count(oldLoopOp)) {
+    info->intraCoreDependentMap[newOp] = info->intraCoreDependentMap[oldLoopOp];
+    info->intraCoreDependentMap.erase(oldLoopOp);
+  }
+  if (isWhile) {
+    auto oldWhileOp = cast<scf::WhileOp>(oldLoopOp);
+    auto newWhileOp = cast<scf::WhileOp>(newOp);
+    if (info->whileBlockArgMap.count(oldWhileOp)) {
+      info->whileBlockArgMap[newWhileOp] =
+          std::move(info->whileBlockArgMap[oldWhileOp]);
+      info->whileBlockArgMap.erase(oldWhileOp);
+    }
+  }
+}
+
+// Extends scf.for / scf.while with iter_args for block counters, inner dep
+// conds, and tensor iter_args (seeded from maps).
+static LogicalResult
+extendMainLoopOpWithExtraArgs(Operation *oldLoopOp,
+                              ControlFlowConditionInfo *info) {
+  int numBlockCounters, numInnerDepConds, numTensorIterArgs;
+  llvm::SmallVector<TensorIterArgIfOpRelation> depsVecCopy;
+  computeMainLoopExtraArgs(oldLoopOp, info, numBlockCounters, numInnerDepConds,
+                           numTensorIterArgs, depsVecCopy);
+
+  int totalExtraArgs = numBlockCounters + numInnerDepConds + numTensorIterArgs;
+  if (totalExtraArgs == 0) {
+    return success();
+  }
+
+  // ForOp block-counter init uses lowerBound (single Value reused for every
+  // counter); WhileOp creates a fresh arith.constant per counter so each new
+  // iter_arg has a distinct SSA value (expected by tests and downstream
+  // passes).
+  OpBuilder builder(oldLoopOp);
+  Value forOpLowerBound;
+  llvm::function_ref<Value()> blockCounterInitFn;
+  bool isWhile = false;
+  if (auto forOp = dyn_cast<scf::ForOp>(oldLoopOp)) {
+    forOpLowerBound = forOp.getLowerBound();
+    blockCounterInitFn = [&]() { return forOpLowerBound; };
+  } else if (auto whileOp = dyn_cast<scf::WhileOp>(oldLoopOp)) {
+    blockCounterInitFn = [&]() {
+      return builder.create<arith::ConstantOp>(
+          whileOp.getLoc(), builder.getI32Type(), builder.getI32IntegerAttr(0));
+    };
+    isWhile = true;
+  } else {
+    LDBG("[Error]: main_loop op is neither scf::ForOp nor scf::WhileOp");
+    return failure();
+  }
+
+  llvm::SmallVector<Value> extraInitArgs;
+  buildMainLoopExtraInitArgs(builder, oldLoopOp->getLoc(), blockCounterInitFn,
+                             numBlockCounters, numInnerDepConds,
+                             numTensorIterArgs, extraInitArgs);
+
+  Operation *newOp = createMainLoopOpAndMigrateBody(oldLoopOp, extraInitArgs);
+  if (!newOp) {
+    return failure();
+  }
+
+  unsigned baseIdx = getMainLoopBaseIdx(oldLoopOp, isWhile);
+  recordMainLoopBlockCountersAndConds(oldLoopOp, newOp, info, baseIdx,
+                                      numBlockCounters, numInnerDepConds);
+  recordMainLoopTensorIterArgs(oldLoopOp, newOp, info, baseIdx,
+                               numBlockCounters, numInnerDepConds,
+                               numTensorIterArgs, depsVecCopy);
+  transferMainLoopInfoMaps(oldLoopOp, newOp, info, isWhile);
+
+  return replaceMainLoopOpUsesAndErase(oldLoopOp, newOp);
+}
+
+// Add block counter and inner dep cond iter args to for ops (and whileOps).
+// Processed from info-driven `mainLoopOpsToProcess` set (keyed on Operation*).
+LogicalResult UpdateLoopOpsPass::addBlockCountersAndInnerDepConds(
+    ModuleOp module, ControlFlowConditionInfo *info) {
+  llvm::DenseSet<Operation *> mainLoopOpsToProcess;
+
+  for (auto &p : info->blockCounterNums) {
+    if (p.second < 0) {
+      LDBG("[Error]: invalid blockCounterNum " << p.second);
+      return failure();
+    }
+    mainLoopOpsToProcess.insert(p.first);
+  }
+  for (auto &p : info->tensorIterArgDepsMap) {
+    mainLoopOpsToProcess.insert(p.first);
+  }
+
+  for (Operation *loopOp : mainLoopOpsToProcess) {
+    if (failed(extendMainLoopOpWithExtraArgs(loopOp, info)))
+      return failure();
+  }
+
+  return success();
+}
+
+// Insert sync ops inside a main-loop body (scf.for or scf.while after-region):
+// wait at start, set before yield. Body block is resolved via getMainLoopBody.
+static LogicalResult
+insertSyncOpsInsideMainLoop(Block *loopBody, Location loc,
+                            hivm::TCoreTypeAttr coreType, PipeAttr setPipe,
+                            PipeAttr waitPipe, int waitFlagId, int setFlagId) {
+  Operation *forTerminator = loopBody->getTerminator();
+  if (!forTerminator) {
+    return failure();
+  }
+
+  // Insert wait at loop body start
+  OpBuilder insertionBuilder(&loopBody->front());
+  auto waitFlagAttr = insertionBuilder.getIntegerAttr(
+      insertionBuilder.getI64Type(), waitFlagId);
+  insertionBuilder.create<SyncBlockWaitOp>(loc, coreType, setPipe, waitPipe,
+                                           waitFlagAttr);
+
+  // Insert set before yield
+  OpBuilder setBuilder(forTerminator);
+  auto setFlagAttr =
+      setBuilder.getIntegerAttr(setBuilder.getI64Type(), setFlagId);
+  setBuilder.setInsertionPoint(forTerminator);
+  setBuilder.create<SyncBlockSetOp>(loc, coreType, setPipe, waitPipe,
+                                    setFlagAttr);
+
+  return success();
+}
+
+// Insert a sync op (SET before, or WAIT after) outside a main-loop op.
+// op-agnostic. `isBefore`: true -> SET, false -> WAIT.
+static LogicalResult
+insertSyncOpsOutsideMainLoop(Operation *loopOp, Location loc,
+                             hivm::TCoreTypeAttr coreType, PipeAttr setPipe,
+                             PipeAttr waitPipe, int flagId, bool isBefore) {
+  OpBuilder builder(loopOp);
+  auto flagAttr = builder.getIntegerAttr(builder.getI64Type(), flagId);
+  if (isBefore) {
+    builder.create<SyncBlockSetOp>(loc, coreType, setPipe, waitPipe, flagAttr);
+  } else {
+    builder.setInsertionPointAfter(loopOp);
+    builder.create<SyncBlockWaitOp>(loc, coreType, setPipe, waitPipe, flagAttr);
+  }
+  return success();
+}
+
+// Returns the loop body block for a main-loop op. scf.for has the single
+// region; scf.while uses the after region (where the actual loop body lives).
+static Block *getMainLoopBody(Operation *loopOp) {
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp)) {
+    return &whileOp.getAfter().front();
+  }
+  return &loopOp->getRegion(0).front();
+}
+
+// Insert PIPE_S for a main_loop op (scf.for or scf.while) by loop/scope type.
+// Body resolved via getMainLoopBody — scf.while uses the after region block.
+static LogicalResult
+insertPipeSForMainLoopOp(Operation *loopOp, scope::ScopeOp scopeOp,
+                         bool isScopeCube, bool isScopeVector, PipeAttr setPipe,
+                         PipeAttr waitPipe, int flagId) {
+  Block *loopBody = getMainLoopBody(loopOp);
+  Location loc = loopOp->getLoc();
+  bool isVectorFirst = loopOp->hasAttr(CVPipeline::kVectorFirst);
+  auto cubeType =
+      hivm::TCoreTypeAttr::get(loopOp->getContext(), hivm::TCoreType::CUBE);
+  auto vectorType =
+      hivm::TCoreTypeAttr::get(loopOp->getContext(), hivm::TCoreType::VECTOR);
+
+  if (isVectorFirst) {
+    if (isScopeCube) {
+      // vector_first + CUBE: before loop op (SET), inside (WAIT/SET)
+      if (failed(insertSyncOpsOutsideMainLoop(loopOp, loc, cubeType, setPipe,
+                                              waitPipe, flagId, true))) {
+        return failure();
+      }
+      if (failed(insertSyncOpsInsideMainLoop(loopBody, loc, cubeType, setPipe,
+                                             waitPipe, flagId, flagId))) {
+        return failure();
+      }
+    } else if (isScopeVector) {
+      // vector_first + VECTOR: inside (WAIT/SET), after loop op (WAIT)
+      if (failed(insertSyncOpsInsideMainLoop(loopBody, loc, vectorType, setPipe,
+                                             waitPipe, flagId, flagId))) {
+        return failure();
+      }
+      if (failed(insertSyncOpsOutsideMainLoop(loopOp, loc, vectorType, setPipe,
+                                              waitPipe, flagId, false))) {
+        return failure();
+      }
+    }
+  } else {
+    // cube_first (including default when neither attribute is present)
+    if (isScopeCube) {
+      // cube_first + CUBE: inside (WAIT/SET), after loop op (WAIT)
+      if (failed(insertSyncOpsInsideMainLoop(loopBody, loc, cubeType, setPipe,
+                                             waitPipe, flagId, flagId))) {
+        return failure();
+      }
+      if (failed(insertSyncOpsOutsideMainLoop(loopOp, loc, cubeType, setPipe,
+                                              waitPipe, flagId, false))) {
+        return failure();
+      }
+    } else if (isScopeVector) {
+      // cube_first + VECTOR: before loop op (SET), inside (WAIT/SET)
+      if (failed(insertSyncOpsOutsideMainLoop(loopOp, loc, vectorType, setPipe,
+                                              waitPipe, flagId, true))) {
+        return failure();
+      }
+      if (failed(insertSyncOpsInsideMainLoop(loopBody, loc, vectorType, setPipe,
+                                             waitPipe, flagId, flagId))) {
+        return failure();
+      }
+    }
+  }
+  return success();
+}
+
+LogicalResult UpdateLoopOpsPass::insertInterCorePipeS(ModuleOp module) {
+  auto cubeCoreType =
+      hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE);
+  auto vectorCoreType =
+      hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::VECTOR);
+  auto setPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S);
+  auto waitPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S);
+
+  WalkResult result = module.walk([&](scope::ScopeOp scopeOp) -> WalkResult {
+    auto scopeTypeAttr =
+        scopeOp->getAttrOfType<hivm::TCoreTypeAttr>("hivm.tcore_type");
+    if (!scopeTypeAttr) {
+      return WalkResult::advance();
+    }
+
+    bool isScopeCube = (scopeTypeAttr == cubeCoreType);
+    bool isScopeVector = (scopeTypeAttr == vectorCoreType);
+
+    // Walk both scf.for and scf.while with ssbuffer.main_loop; both need
+    // same PIPE_S sync. insertPipeSForMainLoopOp dispatches via
+    // getMainLoopBody.
+    WalkResult innerResult = scopeOp.walk([&](Operation *op) -> WalkResult {
+      if (!isMainLoopOp(op)) {
+        return WalkResult::advance();
+      }
+      if (failed(insertPipeSForMainLoopOp(op, scopeOp, isScopeCube,
+                                          isScopeVector, setPipeType,
+                                          waitPipeType, kPipeSFlagId))) {
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+
+    if (innerResult.wasInterrupted()) {
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+
+  return result.wasInterrupted() ? failure() : success();
+}
+
+// Analyzes producer/consumer relationship between tensor iter_args in
+// main_loop and ssbuffer.if. Downstream only acts on forOp (whileOp future
+// use).
+LogicalResult UpdateLoopOpsPass::analyzeTensorIterArgDependencies(
+    ModuleOp module, ControlFlowConditionInfo *info) {
+  bool failed = false;
+  module.walk([&](Operation *op) -> WalkResult {
+    if (!isMainLoopOp(op)) {
+      return WalkResult::advance();
+    }
+
+    LDBG("Analyzing main_loop op: " << op);
+
+    // Get iter_args to analyze. forOp: getRegionIterArgs() (body args - IV).
+    // whileOp: getAfterArguments() (before-block args aren't visible in body).
+    llvm::SmallVector<Value> iterArgsVec = MainLoop(op).getIterArgs();
+
+    for (auto iterArg : iterArgsVec) {
+      if (!mlir::isa<TensorType>(iterArg.getType())) {
+        continue;
+      }
+
+      LDBG("Found tensor type iter_arg: " << iterArg);
+      scf::IfOp producerIfOp = nullptr;
+      llvm::SmallVector<scf::IfOp> consumerIfOps;
+
+      for (auto &use : iterArg.getUses()) {
+        Operation *user = use.getOwner();
+        scf::IfOp ifOp = nullptr;
+        Operation *curr = user;
+        while (curr && curr != op) {
+          if (auto currIf = dyn_cast<scf::IfOp>(curr)) {
+            if (currIf->hasAttr(CVPipeline::kIf)) {
+              ifOp = currIf;
+              break;
+            }
+          }
+          curr = curr->getParentOp();
+        }
+
+        if (!ifOp) {
+          LDBG("Use of tensor iter_arg "
+               << iterArg << " is not inside any ssbuffer.if op.");
+          continue;
+        }
+
+        // Only the direct terminator of this ssbuffer.if is a producer.
+        // Nested if/for/while yields that forward iter_arg are consumers.
+        bool isProducer = isa<scf::YieldOp>(user) &&
+                          user->getParentOp() == ifOp.getOperation();
+
+        // Check and update status of ifOp
+        if (isProducer) {
+          if (producerIfOp && producerIfOp != ifOp) {
+            // Found a different producer ifOp! This is an error.
+            LDBG("[Error]: tensor iter_arg "
+                 << iterArg << " has multiple different producers!");
+            LDBG("Existing producer: " << producerIfOp);
+            LDBG("New producer: " << ifOp);
+            failed = true;
+            return WalkResult::interrupt();
+          }
+          if (!producerIfOp) {
+            // First producer, or upgrade from consumer
+            auto it = llvm::find(consumerIfOps, ifOp);
+            if (it != consumerIfOps.end()) {
+              consumerIfOps.erase(it);
+              LDBG("This ifOp was consumer, now updated to producer: " << ifOp);
+            } else {
+              LDBG("Found producer ifOp (first time): " << ifOp);
+            }
+            producerIfOp = ifOp;
+          }
+          // Else: already is this producer, do nothing
+        } else {
+          // isConsumer
+          if (producerIfOp == ifOp) {
+            // Already a producer, even if current use is consumer, do nothing
+            continue;
+          }
+          if (!llvm::is_contained(consumerIfOps, ifOp)) {
+            consumerIfOps.push_back(ifOp);
+            LDBG("Found consumer ifOp (first time): " << ifOp << "\n");
+          }
+          // Else: already a consumer, do nothing
+        }
+      }
+      // Check: must have both producers AND consumers
+      if (!producerIfOp || consumerIfOps.empty()) {
+        LDBG("tensor iter_arg " << iterArg << " has only "
+                                << (!producerIfOp ? "consumers" : "producers")
+                                << ", skipped");
+        continue;
+      }
+      TensorIterArgIfOpRelation relation;
+      relation.iterArg = iterArg;
+      relation.producer = producerIfOp;
+      relation.consumers = consumerIfOps;
+
+      // Record into the map (keyed on Operation*; both scf.for and scf.while
+      // participate).
+      info->tensorIterArgDepsMap[op].push_back(relation);
+      LDBG("Recorded tensor iter_arg dependency: "
+           << iterArg << " has 1 producer, " << relation.consumers.size()
+           << " consumers");
+    }
+
+    return WalkResult::advance();
+  });
+
+  return failed ? failure() : success();
+}
+
+void UpdateLoopOpsPass::runOnOperation() {
+  ModuleOp module = getOperation();
+
+  if (CVPipeline::hasFallbackAttr(module)) {
+    return;
+  }
+
+  LDBG("before updateLoopOps:\n" << module);
+
+  // Use provided info, or create a local one if not available
+  ControlFlowConditionInfo localInfo;
+  ControlFlowConditionInfo *infoToUse = info ? info : &localInfo;
+
+  // Analyze the dependencies of the tensor type iter_args in the main_loop with
+  // the ssbuffer.if ops
+  if (failed(analyzeTensorIterArgDependencies(module, infoToUse))) {
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+    return;
+  }
+
+  // Derive block counters from ssbuffer.if if blockCounterNums is empty
+  if (infoToUse->blockCounterNums.empty()) {
+    if (failed(deriveBlockCountersFromIfOps(module, infoToUse))) {
+      CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+      return;
+    }
+  }
+
+  // Update for/while ops iter_args for block counters and inner dep conds.
+  // forOp from info-driven sets (no-ops when empty); whileOp unconditionally.
+  if (infoToUse &&
+      (failed(addBlockCountersAndInnerDepConds(module, infoToUse)))) {
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+    return;
+  }
+
+  // Insert PIPE_S inter-core synchronization
+  if (failed(insertInterCorePipeS(module))) {
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+    return;
+  }
+
+  // Dump whileBlockArgMap (whileOp -> block_id -> (new_arg_idx -> old_arg_idx))
+  // after all whileOp replacements; verify it survived
+  // replaceMainLoopOpUsesAndErase.
+  dumpWhileBlockArgMap(infoToUse->whileBlockArgMap,
+                       "whileBlockArgMap contents after updateLoopOps "
+                       "(whileOp -> block_id -> (new_arg_idx -> old_arg_idx))");
+
+  LDBG("after updateLoopOps:\n" << module);
+}
+
+namespace mlir {
+namespace triton {
+
+std::unique_ptr<OperationPass<ModuleOp>> createUpdateLoopOpsPass() {
+  return std::make_unique<UpdateLoopOpsPass>();
+}
+
+} // namespace triton
+} // namespace mlir
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp
index ef94712068..64109efdc5 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp
@@ -23,15 +23,25 @@
 #include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h"
 #include "ascend/include/DynamicCVPipeline/Common/Utils.h"
 #include "bishengir/Dialect/HIVM/IR/HIVM.h"
+#include "llvm/Support/Debug.h"
 #include <optional>
 
-using namespace mlir;
+static constexpr const char *DEBUG_TYPE = "AddControlFlowConditionUtils";
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define LDBG(...)                                                              \
+  LLVM_DEBUG({                                                                 \
+    DBGS();                                                                    \
+    llvm::dbgs() << __VA_ARGS__;                                               \
+    llvm::dbgs() << "\n";                                                      \
+  })
+
+namespace mlir {
 using namespace llvm;
+using namespace CVPipeline;
 
 // Collect all nested ops within an operation's regions
-LogicalResult
-triton::collectAllNestedOps(Operation *op,
-                            llvm::DenseSet<Operation *> &regionOps) {
+LogicalResult collectAllNestedOps(Operation *op,
+                                  llvm::DenseSet<Operation *> &regionOps) {
   if (!op) {
     return failure();
   }
@@ -44,7 +54,7 @@ triton::collectAllNestedOps(Operation *op,
   for (Region &region : op->getRegions()) {
     for (Block &block : region) {
       for (Operation &nestedOp : block) {
-        if (failed(triton::collectAllNestedOps(&nestedOp, regionOps))) {
+        if (failed(collectAllNestedOps(&nestedOp, regionOps))) {
           return failure();
         }
       }
@@ -54,10 +64,17 @@ triton::collectAllNestedOps(Operation *op,
   return success();
 }
 
-// Group operations by their block_id attribute
-LogicalResult triton::collectOpsByBlockId(
-    scf::ForOp forOp, llvm::DenseMap<int, SmallVector<Operation *>> &blockOps) {
-  for (Operation &op : forOp.getBody()->without_terminator()) {
+// Group operations by their block_id attribute. `op` must be scf.for or
+// scf.while (see MainLoop::getBody).
+LogicalResult
+collectOpsByBlockId(Operation *op,
+                    llvm::DenseMap<int, SmallVector<Operation *>> &blockOps) {
+  Block *bodyBlock = MainLoop(op).getBody();
+  if (!bodyBlock) {
+    return failure();
+  }
+
+  for (Operation &op : bodyBlock->without_terminator()) {
     if (auto attr = op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId)) {
       blockOps[attr.getInt()].push_back(&op);
     } else {
@@ -121,9 +138,9 @@ dfsTopologicalSort(Operation *op, llvm::DenseSet<Operation *> &visited,
 }
 
 // Topological sort of operations based on operand dependencies
-LogicalResult triton::topologicalSort(llvm::DenseSet<Operation *> &ops,
-                                      llvm::DenseMap<Operation *, int> *opOrder,
-                                      SmallVectorImpl<Operation *> &sorted) {
+LogicalResult topologicalSort(llvm::DenseSet<Operation *> &ops,
+                              llvm::DenseMap<Operation *, int> *opOrder,
+                              SmallVectorImpl<Operation *> &sorted) {
   llvm::DenseSet<Operation *> visited;
   llvm::DenseSet<Operation *> inStack;
 
@@ -143,23 +160,29 @@ LogicalResult triton::topologicalSort(llvm::DenseSet<Operation *> &ops,
   return success();
 }
 
-LogicalResult triton::topologicalSort(SmallVector<Operation *> &ops) {
+LogicalResult topologicalSort(SmallVector<Operation *> &ops) {
   llvm::DenseSet<Operation *> opSet(ops.begin(), ops.end());
   SmallVector<Operation *> sorted;
 
-  if (failed(triton::topologicalSort(opSet, nullptr, sorted))) {
-    return failure();
+  if (succeeded(topologicalSort(opSet, nullptr, sorted))) {
+    ops.assign(sorted.begin(), sorted.end());
+    return success();
   }
-  ops.assign(sorted.begin(), sorted.end());
-  return success();
+  return failure();
 }
 
-// Get block_ids in order of appearance in for loop body
-SmallVector<int> triton::getBlockIdsInOrder(scf::ForOp forOp) {
+// Get block_ids in order of appearance in the main-loop body (forOp body
+// or whileOp after-region body). Returns empty if `op` is neither.
+SmallVector<int> getBlockIdsInOrder(Operation *op) {
+  Block *bodyBlock = MainLoop(op).getBody();
+  if (!bodyBlock) {
+    return {};
+  }
+
   SmallVector<int> idsInOrder;
   llvm::DenseSet<int> seenIds;
 
-  for (Operation &op : forOp.getBody()->without_terminator()) {
+  for (Operation &op : bodyBlock->without_terminator()) {
     if (auto blockIdAttr =
             op.getAttrOfType<IntegerAttr>(CVPipeline::kBlockId)) {
       int id = blockIdAttr.getInt();
@@ -171,21 +194,17 @@ SmallVector<int> triton::getBlockIdsInOrder(scf::ForOp forOp) {
   return idsInOrder;
 }
 
-// Get the block_id of the immediate child of scf.for that contains op
-// For nested ops inside scf.if/scf.for, returns the block_id of the immediate
-// child of scf.for Only considers scf.for ops that have ssbuffer.main_loop
-// attribute
-std::optional<int> triton::getForDirectChildBlockId(Operation *op) {
+// Get block_id of the immediate child of main-loop op (scf.for or scf.while
+// carrying ssbuffer.main_loop) containing op. For scf.while "body" is the
+// after-region block.
+std::optional<int> getLoopDirectChildBlockId(Operation *op) {
   if (!op) {
     return std::nullopt;
   }
   Operation *parent = op->getParentOp();
   while (parent) {
-    // Found the main_loop forOp, op is its direct child
-    if (auto forOp = dyn_cast<scf::ForOp>(parent)) {
-      if (forOp->hasAttr(CVPipeline::kMainLoop)) {
-        return CVPipeline::getOpBlockId(op);
-      }
+    if (CVPipeline::isMainLoopOp(parent)) {
+      return CVPipeline::getOpBlockId(op);
     }
     op = parent;
     parent = parent->getParentOp();
@@ -193,8 +212,20 @@ std::optional<int> triton::getForDirectChildBlockId(Operation *op) {
   return std::nullopt;
 }
 
+// Counts unique ssbuffer.if values inside a main-loop op (scf.for or
+// scf.while), walking all nested ops. Returns 0 if none.
+int countUniqueIfBlockIds(Operation *loopOp) {
+  llvm::DenseSet<int> ifBlockIds;
+  loopOp->walk([&](Operation *innerOp) {
+    if (auto ifAttr = innerOp->getAttrOfType<IntegerAttr>(CVPipeline::kIf)) {
+      ifBlockIds.insert(ifAttr.getInt());
+    }
+  });
+  return static_cast<int>(ifBlockIds.size());
+}
+
 // Find the tcb group id that contains value v
-int triton::findTcbGroupId(
+int findTcbGroupId(
     Value v,
     llvm::DenseMap<int, SmallVector<Value>> &tightlyCoupledBufferGroups) {
   for (auto &tcbEntry : tightlyCoupledBufferGroups) {
@@ -205,11 +236,9 @@ int triton::findTcbGroupId(
   return -1;
 }
 
-// Get isCube/isVector based on the scope's tcore_type attribute
-// Returns failure if scopeOp does not have tcore_type attribute or it's not
-// CUBE/VECTOR
-LogicalResult triton::getScopeType(Operation *scopeOp, bool &isCube,
-                                   bool &isVector) {
+// Get isCube/isVector from scopeOp's tcore_type attribute.
+// Returns failure if attribute is missing or not CUBE/VECTOR.
+LogicalResult getScopeType(Operation *scopeOp, bool &isCube, bool &isVector) {
   isCube = false;
   isVector = false;
 
@@ -237,7 +266,7 @@ LogicalResult triton::getScopeType(Operation *scopeOp, bool &isCube,
 
 // Check if op is a scf.if whose body only contains hivm.hir.sync_block_wait,
 // hivm.hir.sync_block_set and hivm.fixpipe ops (excluding terminators).
-bool triton::isIfOpWithOnlySyncOps(Operation *op) {
+bool isIfOpWithOnlySyncOps(Operation *op) {
   auto ifOp = dyn_cast<scf::IfOp>(op);
   if (!ifOp) {
     return false;
@@ -259,3 +288,168 @@ bool triton::isIfOpWithOnlySyncOps(Operation *op) {
 
   return !result.wasInterrupted();
 }
+
+// Migrate ops from oldBlock to newBlock; replaceAllUsesWith on oldBlock's args
+// to newBlock's args (same index).
+void migrateBody(Block *oldBlock, Block *newBlock) {
+  for (unsigned i = 0; i < oldBlock->getNumArguments(); ++i) {
+    oldBlock->getArgument(i).replaceAllUsesWith(newBlock->getArgument(i));
+  }
+
+  for (Operation &op :
+       llvm::make_early_inc_range(oldBlock->without_terminator())) {
+    op.moveBefore(newBlock, newBlock->end());
+  }
+}
+
+// Migrate both before and after regions of a scf.while op. Does not touch
+// terminators (the caller builds new condition/yield in the new regions).
+void migrateWhileBodies(scf::WhileOp oldWhileOp, scf::WhileOp newWhileOp) {
+  migrateBody(oldWhileOp.getBeforeBody(), newWhileOp.getBeforeBody());
+  migrateBody(oldWhileOp.getAfterBody(), newWhileOp.getAfterBody());
+}
+
+// Build new scf.yield at end of `newBlock`: copies oldBlock's yield operands,
+// appends `extraYieldValues`, creates new scf::YieldOp, erases old yield.
+LogicalResult buildNewYieldOp(Block *oldBlock, Block *newBlock,
+                              Operation *newOp,
+                              ArrayRef<Value> extraYieldValues) {
+  auto oldYield = cast<scf::YieldOp>(oldBlock->getTerminator());
+  SmallVector<Value> yieldOperands;
+  for (unsigned i = 0; i < oldYield.getNumOperands(); ++i) {
+    yieldOperands.push_back(oldYield.getOperand(i));
+  }
+  for (Value v : extraYieldValues) {
+    yieldOperands.push_back(v);
+  }
+  OpBuilder builder = OpBuilder::atBlockEnd(newBlock);
+  builder.create<scf::YieldOp>(newOp->getLoc(), yieldOperands);
+  oldYield.erase();
+  return success();
+}
+
+// Replace all uses of `oldOp`'s results with `newOp`'s matching results.
+// No-op when `oldOp` is result-less.
+void replaceOpResultUses(Operation *oldOp, Operation *newOp) {
+  if (oldOp->getNumResults() == 0)
+    return;
+
+  SmallVector<Value> newResults;
+  for (unsigned i = 0; i < oldOp->getNumResults(); ++i) {
+    newResults.push_back(newOp->getResult(i));
+  }
+  oldOp->replaceAllUsesWith(newResults);
+}
+
+// Build new scf.condition in `newWhileOp`'s before region. Condition preserved
+// from `whileOp`; forwarded values = new before-block args (incl. extras).
+void buildNewWhileCondition(scf::WhileOp whileOp, scf::WhileOp newWhileOp) {
+  auto oldCond = whileOp.getConditionOp();
+  Value origCond = oldCond.getCondition();
+
+  OpBuilder beforeBuilder(newWhileOp.getBeforeBody(),
+                          newWhileOp.getBeforeBody()->end());
+  SmallVector<Value> forwardedValues;
+  for (BlockArgument arg : newWhileOp.getBeforeArguments()) {
+    forwardedValues.push_back(arg);
+  }
+  beforeBuilder.create<scf::ConditionOp>(whileOp.getLoc(), origCond,
+                                         forwardedValues);
+  oldCond.erase();
+}
+
+// Creates a new scf.for with `extraInitArgs` appended to the original init
+// args. Returns `oldForOp` unchanged when `extraInitArgs` is empty.
+scf::ForOp createNewForOpWithExtras(scf::ForOp oldForOp,
+                                    ArrayRef<Value> extraInitArgs) {
+  if (extraInitArgs.empty()) {
+    return oldForOp;
+  }
+
+  OpBuilder builder(oldForOp);
+  SmallVector<Value> newInitArgs(oldForOp.getInitArgs().begin(),
+                                 oldForOp.getInitArgs().end());
+  llvm::append_range(newInitArgs, extraInitArgs);
+
+  scf::ForOp newForOp = builder.create<scf::ForOp>(
+      oldForOp.getLoc(), oldForOp.getLowerBound(), oldForOp.getUpperBound(),
+      oldForOp.getStep(), newInitArgs);
+
+  for (auto &attr : oldForOp->getAttrs()) {
+    newForOp->setAttr(attr.getName(), attr.getValue());
+  }
+  return newForOp;
+}
+
+// Creates a new scf.while with `extraInitArgs` appended to the original inits
+// and empty before/after blocks. Returns `oldWhileOp` unchanged when empty.
+scf::WhileOp createNewWhileOpWithExtras(scf::WhileOp oldWhileOp,
+                                        ArrayRef<Value> extraInitArgs) {
+  if (extraInitArgs.empty()) {
+    return oldWhileOp;
+  }
+
+  OpBuilder builder(oldWhileOp);
+
+  SmallVector<Value> newInits(oldWhileOp.getInits().begin(),
+                              oldWhileOp.getInits().end());
+  llvm::append_range(newInits, extraInitArgs);
+
+  SmallVector<Type> newResultTypes(oldWhileOp->getResultTypes().begin(),
+                                   oldWhileOp->getResultTypes().end());
+  for (Value v : extraInitArgs) {
+    newResultTypes.push_back(v.getType());
+  }
+
+  scf::WhileOp newWhileOp = builder.create<scf::WhileOp>(
+      oldWhileOp.getLoc(), newResultTypes, newInits);
+
+  for (auto &attr : oldWhileOp->getAttrs()) {
+    newWhileOp->setAttr(attr.getName(), attr.getValue());
+  }
+
+  SmallVector<Type> argTypes;
+  argTypes.reserve(newInits.size());
+  for (Value v : newInits) {
+    argTypes.push_back(v.getType());
+  }
+  SmallVector<Location> argLocs(newInits.size(), oldWhileOp.getLoc());
+
+  builder.createBlock(&newWhileOp.getBefore(), {}, argTypes, argLocs);
+  builder.createBlock(&newWhileOp.getAfter(), {}, argTypes, argLocs);
+
+  return newWhileOp;
+}
+
+// Dispatches createNewForOpWithExtras / createNewWhileOpWithExtras by op type.
+// Returns nullptr if `oldOp` is neither scf.for nor scf.while.
+Operation *createMainLoopOpWithExtras(Operation *oldOp,
+                                      ArrayRef<Value> extraInitArgs) {
+  if (auto forOp = dyn_cast<scf::ForOp>(oldOp)) {
+    return createNewForOpWithExtras(forOp, extraInitArgs);
+  }
+  if (auto whileOp = dyn_cast<scf::WhileOp>(oldOp)) {
+    return createNewWhileOpWithExtras(whileOp, extraInitArgs);
+  }
+  return nullptr;
+}
+
+// Prints whileBlockArgMap (whileOp -> block_id -> (new_arg_idx ->
+// old_arg_idx)) to the debug stream, gated by LLVM_DEBUG.
+void dumpWhileBlockArgMap(const triton::WhileBlockArgMap &map,
+                          llvm::StringRef header) {
+  LLVM_DEBUG({
+    LDBG("[INFO]: " << header);
+    for (auto &[whileOp, blockArgMap] : map) {
+      LDBG("  whileOp @" << whileOp->getLoc());
+      for (auto &[blockId, argIdxMap] : blockArgMap) {
+        for (auto &[newArgIdx, oldArgIdx] : argIdxMap) {
+          LDBG("    block_id=" << blockId << " new_arg_idx=" << newArgIdx
+                               << " -> old_arg_idx=" << oldArgIdx);
+        }
+      }
+    }
+  });
+}
+
+} // namespace mlir
diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MarkGMLoadPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MarkGMLoadPass.cpp
index 6ad3eae8f7..bdd293f428 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MarkGMLoadPass.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MarkGMLoadPass.cpp
@@ -183,7 +183,7 @@ static int resolveBufferCount(scope::ScopeOp scopeOp) {
   }
   bool isCube = false;
   bool isVector = false;
-  if (failed(triton::getScopeType(scopeOp, isCube, isVector))) {
+  if (failed(getScopeType(scopeOp, isCube, isVector))) {
     return buffer_num;
   }
   if (isVector) {
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-clone-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-clone-ops.mlir
new file mode 100644
index 0000000000..136c12f1f5
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-clone-ops.mlir
@@ -0,0 +1,110 @@
+// RUN: triton-opt --clone-ops %s --allow-unregistered-dialect | FileCheck %s
+
+// Test CloneOps handling of an scf.while main loop with a cross-block update
+// chain (block 9 owns the addition; block 7 builds on it via muli/addi and
+// also adds extra vector-only index ops) and the vector-only index chain in
+// block 7 that feeds the cube block (block 8). The pass clones the cross-block
+// source into block 7 (`clone = 9`) and clones the vector index ops into block
+// 8 (`clone = 7`).
+
+// CHECK: func.func @pcb12_tc01_while_matmul_fill
+// CHECK: scf.while
+// The cross-block update chain source in block 9 stays in place.
+// CHECK: arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32} : i32
+// It is cloned into block 7 carrying ssbuffer.clone = 9, so block 7 can
+// build its own muli/addi chain without crossing blocks.
+// CHECK: arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32} : i32
+// Block 7's vector-only index ops land in block 8 with clone = 7.
+// CHECK: arith.index_cast %arg7 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+// CHECK: arith.maxsi {{.*}} {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+// CHECK: arith.minsi {{.*}} {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+// CHECK: arith.cmpi slt, {{.*}} {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+// CHECK: arith.index_cast %arg6 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+// The terminator yields the chain's top result for the loop's single iter_arg.
+// CHECK: scf.yield %{{.*}} : i32
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @pcb12_tc01_while_matmul_fill(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32, %arg16: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %cst = arith.constant {ssbuffer.block_id = 11 : i32} 0.000000e+00 : f16
+    %c1_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 1 : i32
+    %c0_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : i32
+    %c0 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : index
+    %c32 = arith.constant {ssbuffer.block_id = 10 : i32} 32 : index
+    %c64 = arith.constant {ssbuffer.block_id = 10 : i32} 64 : index
+    %cst_0 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[4, 1, 16, 16]> : tensor<4xi64>
+    %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[16, 4, 16]> : tensor<3xi64>
+    %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[1, 2, 16, 16]> : tensor<4xi64>
+    %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[32, 1, 16]> : tensor<3xi64>
+    scope.scope : () -> () {
+      %alloc_16 = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_16 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_17 = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_17 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+
+      %0 = scf.while (%arg17 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.cmpi slt, %arg17, %arg7 {Undefined, ssbuffer.block_id = 4 : i32} : i32
+        scf.condition(%1) %arg17 : i32
+      } do {
+      ^bb0(%arg17: i32):
+        %1 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32} : i32
+        %2 = arith.index_cast %arg7 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %3 = arith.maxsi %2, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %4 = arith.minsi %3, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %5 = arith.cmpi slt, %4, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %6 = arith.index_cast %arg6 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %7 = arith.maxsi %6, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %8 = arith.minsi %7, %c64 {ssbuffer.block_id = 7 : i32} : index
+        %9 = arith.cmpi slt, %8, %c64 {ssbuffer.block_id = 7 : i32} : index
+
+        %mj_1 = arith.muli %1, %c1_i32 {ssbuffer.block_id = 7 : i32} : i32
+        %mj_2 = arith.addi %mj_1, %arg17 {ssbuffer.block_id = 7 : i32} : i32
+
+        %alloc = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<32xf16>
+        %alloc_5 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<64xf16>
+        scf.if %9 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_5 : memref<64xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        scf.if %5 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc : memref<32xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        %10 = arith.muli %arg17, %arg8 {ssbuffer.block_id = 8 : i32} : i32
+        %11 = arith.index_cast %10 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%11], sizes: [32], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<32xf16, strided<[1], offset: ?>>
+        %subview = memref.subview %reinterpret_cast[0] [%4] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_6 = memref.subview %alloc[0] [%4] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview, %subview_6 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %12 = bufferization.to_tensor %alloc restrict writable {ssbuffer.block_id = 8 : i32} : memref<32xf16> to tensor<32xf16>
+        %13 = arith.muli %arg17, %arg9 {ssbuffer.block_id = 8 : i32} : i32
+        %14 = arith.index_cast %13 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast_7 = memref.reinterpret_cast %arg3 to offset: [%14], sizes: [64], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<64xf16, strided<[1], offset: ?>>
+        %subview_8 = memref.subview %reinterpret_cast_7[0] [%8] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_9 = memref.subview %alloc_5[0] [%8] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview_8, %subview_9 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %15 = bufferization.to_tensor %alloc_5 restrict writable {ssbuffer.block_id = 8 : i32} : memref<64xf16> to tensor<64xf16>
+        %expanded = tensor.expand_shape %12 [[0, 1]] output_shape [32, 1] {ssbuffer.block_id = 8 : i32} : tensor<32xf16> into tensor<32x1xf16>
+        %16 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<32x16xf16>
+        %17 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%16 : tensor<32x16xf16>) -> tensor<32x16xf16>
+        %inserted_slice = tensor.insert_slice %expanded into %17[0, 0] [32, 1] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<32x1xf16> into tensor<32x16xf16>
+        %expanded_10 = tensor.expand_shape %15 [[0, 1]] output_shape [1, 64] {ssbuffer.block_id = 8 : i32} : tensor<64xf16> into tensor<1x64xf16>
+        %18 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<16x64xf16>
+        %19 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%18 : tensor<16x64xf16>) -> tensor<16x64xf16>
+        %inserted_slice_11 = tensor.insert_slice %expanded_10 into %19[0, 0] [1, 64] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<1x64xf16> into tensor<16x64xf16>
+        %reshape = tensor.reshape %inserted_slice(%cst_3) {ssbuffer.block_id = 8 : i32} : (tensor<32x16xf16>, tensor<3xi64>) -> tensor<32x1x16xf16>
+        %20 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<1x32x16xf16>
+        %transposed = linalg.transpose ins(%reshape : tensor<32x1x16xf16>) outs(%20 : tensor<1x32x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_12 = tensor.reshape %transposed(%cst_2) {ssbuffer.block_id = 8 : i32} : (tensor<1x32x16xf16>, tensor<4xi64>) -> tensor<1x2x16x16xf16>
+        %reshape_13 = tensor.reshape %inserted_slice_11(%cst_1) {ssbuffer.block_id = 8 : i32} : (tensor<16x64xf16>, tensor<3xi64>) -> tensor<16x4x16xf16>
+        %21 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<4x16x16xf16>
+        %transposed_14 = linalg.transpose ins(%reshape_13 : tensor<16x4x16xf16>) outs(%21 : tensor<4x16x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_15 = tensor.reshape %transposed_14(%cst_0) {ssbuffer.block_id = 8 : i32} : (tensor<4x16x16xf16>, tensor<4xi64>) -> tensor<4x1x16x16xf16>
+        hivm.hir.copy ins(%reshape_12 : tensor<1x2x16x16xf16>) outs(%alloc_16 : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 1
+        hivm.hir.copy ins(%reshape_15 : tensor<4x1x16x16xf16>) outs(%alloc_17 : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        scf.yield %mj_2 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-create-if-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-create-if-ops.mlir
new file mode 100644
index 0000000000..9139025963
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-create-if-ops.mlir
@@ -0,0 +1,130 @@
+// RUN: triton-opt --create-if-ops %s --allow-unregistered-dialect | FileCheck %s
+
+// Test CreateIfOps on an scf.while main loop after ProcessArgs. The loop
+// already carries 5 iter_args. Each block in the do region is wrapped in an
+// scf.if guarded by a boolean, and the per-block while_arg update chains (an
+// addi/muli/addi triple tagged `ssbuffer.while_arg = 0`) are threaded through
+// the if results so they still reach scf.yield.
+
+// CHECK: func.func @pcb12_tc01_while_matmul_fill
+// CHECK: scf.while
+// CHECK: ^bb0(%arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32, %arg21: i32):
+// Block 9 wrapped into an scf.if returning its single while_arg update.
+// CHECK: scf.if %{{.*}} -> (i32) {
+// CHECK: arith.muli %{{.*}}, %c1_i32 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: arith.addi %{{.*}}, %arg19 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: } {hivm.matmul_limited_in_cube, ssbuffer.if = 9 : i32}
+// Block 7 wrapped into an scf.if returning its condition update and while_arg.
+// CHECK: scf.if %{{.*}} -> (i32, i32) {
+// CHECK: arith.addi %{{.*}}, %arg20 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: } {hivm.matmul_limited_in_cube, ssbuffer.if = 7 : i32}
+// Block 8 wrapped into an scf.if returning its arg update and while_arg.
+// CHECK: scf.if %{{.*}} -> (i32, i32) {
+// CHECK: arith.addi %{{.*}}, %arg21 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: } {hivm.matmul_limited_in_cube, ssbuffer.if = 8 : i32}
+// The terminator wires the if results back into the 5 yielded values.
+// CHECK: scf.yield %{{.*}}#0, %{{.*}}#0, %{{.*}}, %{{.*}}#1, %{{.*}}#1 : i32, i32, i32, i32, i32
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @pcb12_tc01_while_matmul_fill(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32, %arg16: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %cst = arith.constant {ssbuffer.block_id = 11 : i32} 0.000000e+00 : f16
+    %c1_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 1 : i32
+    %c0_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : i32
+    %c0 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : index
+    %c32 = arith.constant {ssbuffer.block_id = 10 : i32} 32 : index
+    %c64 = arith.constant {ssbuffer.block_id = 10 : i32} 64 : index
+    %cst_0 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[4, 1, 16, 16]> : tensor<4xi64>
+    %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[16, 4, 16]> : tensor<3xi64>
+    %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[1, 2, 16, 16]> : tensor<4xi64>
+    %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[32, 1, 16]> : tensor<3xi64>
+    scope.scope : () -> () {
+      %alloc = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_4 = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_4 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      %0:5 = scf.while (%arg17 = %c0_i32, %arg18 = %c0_i32, %arg19 = %c0_i32, %arg20 = %c0_i32, %arg21 = %c0_i32) : (i32, i32, i32, i32, i32) -> (i32, i32, i32, i32, i32) {
+        %1 = arith.cmpi slt, %arg17, %arg7 {Undefined, ssbuffer.block_id = 4 : i32} : i32
+        scf.condition(%1) %arg17, %arg18, %arg19, %arg20, %arg21 : i32, i32, i32, i32, i32
+      } do {
+      ^bb0(%arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32, %arg21: i32):
+        %1 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32} : i32
+        %2 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %3 = arith.muli %2, %c1_i32 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %4 = arith.addi %3, %arg19 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %5 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32} : i32
+        %6 = arith.index_cast %arg7 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %7 = arith.maxsi %6, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %8 = arith.minsi %7, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %9 = arith.cmpi slt, %8, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %10 = arith.index_cast %arg6 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %11 = arith.maxsi %10, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %12 = arith.minsi %11, %c64 {ssbuffer.block_id = 7 : i32} : index
+        %13 = arith.cmpi slt, %12, %c64 {ssbuffer.block_id = 7 : i32} : index
+        %14 = arith.muli %5, %c1_i32 {ssbuffer.block_id = 7 : i32} : i32
+        %15 = arith.addi %14, %arg17 {ssbuffer.block_id = 7 : i32} : i32
+        %16 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %17 = arith.muli %16, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %18 = arith.addi %17, %arg20 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %19 = arith.index_cast %arg7 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+        %20 = arith.maxsi %19, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %21 = arith.minsi %20, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %22 = arith.cmpi slt, %21, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %23 = arith.index_cast %arg6 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+        %24 = arith.maxsi %23, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %25 = arith.minsi %24, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %26 = arith.cmpi slt, %25, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %alloc_5 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<32xf16>
+        %alloc_6 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<64xf16>
+        scf.if %26 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_6 : memref<64xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        scf.if %22 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_5 : memref<32xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        %27 = arith.muli %arg18, %arg8 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+        %28 = arith.index_cast %27 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%28], sizes: [32], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<32xf16, strided<[1], offset: ?>>
+        %subview = memref.subview %reinterpret_cast[0] [%21] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_7 = memref.subview %alloc_5[0] [%21] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview, %subview_7 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %29 = bufferization.to_tensor %alloc_5 restrict writable {ssbuffer.block_id = 8 : i32} : memref<32xf16> to tensor<32xf16>
+        %30 = arith.muli %arg18, %arg9 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+        %31 = arith.index_cast %30 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast_8 = memref.reinterpret_cast %arg3 to offset: [%31], sizes: [64], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<64xf16, strided<[1], offset: ?>>
+        %subview_9 = memref.subview %reinterpret_cast_8[0] [%25] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_10 = memref.subview %alloc_6[0] [%25] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview_9, %subview_10 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %32 = bufferization.to_tensor %alloc_6 restrict writable {ssbuffer.block_id = 8 : i32} : memref<64xf16> to tensor<64xf16>
+        %expanded = tensor.expand_shape %29 [[0, 1]] output_shape [32, 1] {ssbuffer.block_id = 8 : i32} : tensor<32xf16> into tensor<32x1xf16>
+        %33 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<32x16xf16>
+        %34 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%33 : tensor<32x16xf16>) -> tensor<32x16xf16>
+        %inserted_slice = tensor.insert_slice %expanded into %34[0, 0] [32, 1] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<32x1xf16> into tensor<32x16xf16>
+        %expanded_11 = tensor.expand_shape %32 [[0, 1]] output_shape [1, 64] {ssbuffer.block_id = 8 : i32} : tensor<64xf16> into tensor<1x64xf16>
+        %35 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<16x64xf16>
+        %36 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%35 : tensor<16x64xf16>) -> tensor<16x64xf16>
+        %inserted_slice_12 = tensor.insert_slice %expanded_11 into %36[0, 0] [1, 64] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<1x64xf16> into tensor<16x64xf16>
+        %reshape = tensor.reshape %inserted_slice(%cst_3) {ssbuffer.block_id = 8 : i32} : (tensor<32x16xf16>, tensor<3xi64>) -> tensor<32x1x16xf16>
+        %37 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<1x32x16xf16>
+        %transposed = linalg.transpose ins(%reshape : tensor<32x1x16xf16>) outs(%37 : tensor<1x32x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_13 = tensor.reshape %transposed(%cst_2) {ssbuffer.block_id = 8 : i32} : (tensor<1x32x16xf16>, tensor<4xi64>) -> tensor<1x2x16x16xf16>
+        %reshape_14 = tensor.reshape %inserted_slice_12(%cst_1) {ssbuffer.block_id = 8 : i32} : (tensor<16x64xf16>, tensor<3xi64>) -> tensor<16x4x16xf16>
+        %38 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<4x16x16xf16>
+        %transposed_15 = linalg.transpose ins(%reshape_14 : tensor<16x4x16xf16>) outs(%38 : tensor<4x16x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_16 = tensor.reshape %transposed_15(%cst_0) {ssbuffer.block_id = 8 : i32} : (tensor<4x16x16xf16>, tensor<4xi64>) -> tensor<4x1x16x16xf16>
+        hivm.hir.copy ins(%reshape_13 : tensor<1x2x16x16xf16>) outs(%alloc : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 1
+        hivm.hir.copy ins(%reshape_16 : tensor<4x1x16x16xf16>) outs(%alloc_4 : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}
+        %39 = arith.addi %c1_i32, %c1_i32 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32, ssbuffer.clone = 9 : i32} : i32
+        %40 = arith.muli %39, %c1_i32 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+        %41 = arith.addi %40, %arg18 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+        %42 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %43 = arith.muli %42, %c1_i32 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        %44 = arith.addi %43, %arg21 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        scf.yield %15, %41, %4, %18, %44 : i32, i32, i32, i32, i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-process-args.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-process-args.mlir
new file mode 100644
index 0000000000..432ddb9cb3
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-process-args.mlir
@@ -0,0 +1,117 @@
+// RUN: triton-opt --process-args %s --allow-unregistered-dialect | FileCheck %s
+
+// Test ProcessArgs adaptation of an scf.while main loop. The iter_arg used by
+// scf.condition (index 0) is referenced by three separate blocks, so the pass
+// clones its update chain (an addi/muli/addi triple) once per block and grows
+// the loop from 1 to 5 iter_args. Each per-block clone is tagged with
+// `ssbuffer.while_arg = 0`, and the buffer-offset ops that consume the induction
+// arg are rewritten onto the new pipeline arg and tagged with `ssbuffer.arg = 0`.
+
+// CHECK: func.func @pcb12_tc01_while_matmul_fill
+// The while op now carries 5 iter_args and 5 results.
+// CHECK: %{{.*}}:5 = scf.while (%arg17 = %c0_i32, %arg18 = %c0_i32, %arg19 = %c0_i32, %arg20 = %c0_i32, %arg21 = %c0_i32) : (i32, i32, i32, i32, i32) -> (i32, i32, i32, i32, i32)
+// scf.condition forwards all 5 iter_args.
+// CHECK: scf.condition(%{{.*}}) %arg17, %arg18, %arg19, %arg20, %arg21 : i32, i32, i32, i32, i32
+// CHECK: ^bb0(%arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32, %arg21: i32):
+// The tail of block 9's per-block clone chain feeds iter_arg %arg18.
+// CHECK: arith.addi %{{.*}}, %arg18 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// Block 7's per-block clone chain feeds iter_arg %arg19.
+// CHECK: arith.addi %{{.*}}, %arg19 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// The buffer-offset op is rewritten onto pipeline arg %arg21 and tagged ssbuffer.arg.
+// CHECK: arith.muli %arg21, %arg8 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+// Block 8's per-block clone chain feeds iter_arg %arg20.
+// CHECK: arith.addi %{{.*}}, %arg20 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// The terminator yields the original update plus the three per-block updates.
+// CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32, i32, i32
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @pcb12_tc01_while_matmul_fill(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32, %arg16: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %cst = arith.constant {ssbuffer.block_id = 11 : i32} 0.000000e+00 : f16
+    %c1_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 1 : i32
+    %c0_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : i32
+    %c0 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : index
+    %c32 = arith.constant {ssbuffer.block_id = 10 : i32} 32 : index
+    %c64 = arith.constant {ssbuffer.block_id = 10 : i32} 64 : index
+    %cst_0 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[4, 1, 16, 16]> : tensor<4xi64>
+    %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[16, 4, 16]> : tensor<3xi64>
+    %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[1, 2, 16, 16]> : tensor<4xi64>
+    %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[32, 1, 16]> : tensor<3xi64>
+    scope.scope : () -> () {
+      %alloc = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_4 = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_4 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      %0 = scf.while (%arg17 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.cmpi slt, %arg17, %arg7 {Undefined, ssbuffer.block_id = 4 : i32} : i32
+        scf.condition(%1) %arg17 : i32
+      } do {
+      ^bb0(%arg17: i32):
+        %1 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32} : i32
+        %2 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32} : i32
+        %3 = arith.index_cast %arg7 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %4 = arith.maxsi %3, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %5 = arith.minsi %4, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %6 = arith.cmpi slt, %5, %c32 {ssbuffer.block_id = 7 : i32} : index
+        %7 = arith.index_cast %arg6 {ssbuffer.block_id = 7 : i32} : i32 to index
+        %8 = arith.maxsi %7, %c0 {ssbuffer.block_id = 7 : i32} : index
+        %9 = arith.minsi %8, %c64 {ssbuffer.block_id = 7 : i32} : index
+        %10 = arith.cmpi slt, %9, %c64 {ssbuffer.block_id = 7 : i32} : index
+        %11 = arith.muli %2, %c1_i32 {ssbuffer.block_id = 7 : i32} : i32
+        %12 = arith.addi %11, %arg17 {ssbuffer.block_id = 7 : i32} : i32
+        %13 = arith.index_cast %arg7 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+        %14 = arith.maxsi %13, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %15 = arith.minsi %14, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %16 = arith.cmpi slt, %15, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %17 = arith.index_cast %arg6 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+        %18 = arith.maxsi %17, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %19 = arith.minsi %18, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %20 = arith.cmpi slt, %19, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+        %alloc_5 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<32xf16>
+        %alloc_6 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<64xf16>
+        scf.if %20 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_6 : memref<64xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        scf.if %16 {
+          linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_5 : memref<32xf16>)
+        } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+        %21 = arith.muli %arg17, %arg8 {ssbuffer.block_id = 8 : i32} : i32
+        %22 = arith.index_cast %21 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%22], sizes: [32], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<32xf16, strided<[1], offset: ?>>
+        %subview = memref.subview %reinterpret_cast[0] [%15] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_7 = memref.subview %alloc_5[0] [%15] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview, %subview_7 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %23 = bufferization.to_tensor %alloc_5 restrict writable {ssbuffer.block_id = 8 : i32} : memref<32xf16> to tensor<32xf16>
+        %24 = arith.muli %arg17, %arg9 {ssbuffer.block_id = 8 : i32} : i32
+        %25 = arith.index_cast %24 {ssbuffer.block_id = 8 : i32} : i32 to index
+        %reinterpret_cast_8 = memref.reinterpret_cast %arg3 to offset: [%25], sizes: [64], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<64xf16, strided<[1], offset: ?>>
+        %subview_9 = memref.subview %reinterpret_cast_8[0] [%19] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+        %subview_10 = memref.subview %alloc_6[0] [%19] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16> to memref<?xf16, strided<[1]>>
+        memref.copy %subview_9, %subview_10 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+        %26 = bufferization.to_tensor %alloc_6 restrict writable {ssbuffer.block_id = 8 : i32} : memref<64xf16> to tensor<64xf16>
+        %expanded = tensor.expand_shape %23 [[0, 1]] output_shape [32, 1] {ssbuffer.block_id = 8 : i32} : tensor<32xf16> into tensor<32x1xf16>
+        %27 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<32x16xf16>
+        %28 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%27 : tensor<32x16xf16>) -> tensor<32x16xf16>
+        %inserted_slice = tensor.insert_slice %expanded into %28[0, 0] [32, 1] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<32x1xf16> into tensor<32x16xf16>
+        %expanded_11 = tensor.expand_shape %26 [[0, 1]] output_shape [1, 64] {ssbuffer.block_id = 8 : i32} : tensor<64xf16> into tensor<1x64xf16>
+        %29 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<16x64xf16>
+        %30 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%29 : tensor<16x64xf16>) -> tensor<16x64xf16>
+        %inserted_slice_12 = tensor.insert_slice %expanded_11 into %30[0, 0] [1, 64] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<1x64xf16> into tensor<16x64xf16>
+        %reshape = tensor.reshape %inserted_slice(%cst_3) {ssbuffer.block_id = 8 : i32} : (tensor<32x16xf16>, tensor<3xi64>) -> tensor<32x1x16xf16>
+        %31 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<1x32x16xf16>
+        %transposed = linalg.transpose ins(%reshape : tensor<32x1x16xf16>) outs(%31 : tensor<1x32x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_13 = tensor.reshape %transposed(%cst_2) {ssbuffer.block_id = 8 : i32} : (tensor<1x32x16xf16>, tensor<4xi64>) -> tensor<1x2x16x16xf16>
+        %reshape_14 = tensor.reshape %inserted_slice_12(%cst_1) {ssbuffer.block_id = 8 : i32} : (tensor<16x64xf16>, tensor<3xi64>) -> tensor<16x4x16xf16>
+        %32 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<4x16x16xf16>
+        %transposed_15 = linalg.transpose ins(%reshape_14 : tensor<16x4x16xf16>) outs(%32 : tensor<4x16x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+        %reshape_16 = tensor.reshape %transposed_15(%cst_0) {ssbuffer.block_id = 8 : i32} : (tensor<4x16x16xf16>, tensor<4xi64>) -> tensor<4x1x16x16xf16>
+        hivm.hir.copy ins(%reshape_13 : tensor<1x2x16x16xf16>) outs(%alloc : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 1
+        hivm.hir.copy ins(%reshape_16 : tensor<4x1x16x16xf16>) outs(%alloc_4 : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}
+        hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        scf.yield %12 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-loop-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-loop-ops.mlir
new file mode 100644
index 0000000000..5978208e7b
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-loop-ops.mlir
@@ -0,0 +1,146 @@
+// RUN: triton-opt --update-loop-ops %s --allow-unregistered-dialect | FileCheck %s
+
+// Test UpdateLoopOps on an scf.while main loop after CreateIfOps. The loop's do
+// region already holds the per-block scf.if wrappers with their while_arg
+// update chains. The pass emits a prologue sync_block_set, grows the loop from
+// 5 to 8 iter_args (adding 3 pipeline sync args), and brackets the body with a
+// sync_block_wait / sync_block_set pair while preserving the guarded ifs.
+
+// CHECK: func.func @pcb12_tc01_while_matmul_fill
+// A prologue sync_block_set is emitted before the loop.
+// CHECK: hivm.hir.sync_block_set[<VECTOR>, <PIPE_S>, <PIPE_S>] flag = 15
+// The while op grows to 8 iter_args (5 from ProcessArgs + 3 sync args).
+// CHECK: %{{.*}}:8 = scf.while (%arg17 = %c0_i32, %arg18 = %c0_i32, %arg19 = %c0_i32, %arg20 = %c0_i32, %arg21 = %c0_i32, %arg22 = %c0_i32_5, %arg23 = %c0_i32_6, %arg24 = %c0_i32_7) : (i32, i32, i32, i32, i32, i32, i32, i32) -> (i32, i32, i32, i32, i32, i32, i32, i32)
+// CHECK: scf.condition(%{{.*}}) %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23, %arg24
+// CHECK: ^bb0(%arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32, %arg21: i32, %arg22: i32, %arg23: i32, %arg24: i32):
+// The body opens with a sync_block_wait matching the prologue set.
+// CHECK: hivm.hir.sync_block_wait[<VECTOR>, <PIPE_S>, <PIPE_S>] flag = 15
+// Per-block while_arg update chains survive inside the guarded ifs.
+// CHECK: arith.addi %{{.*}}, %arg19 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: arith.addi %{{.*}}, %arg20 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// CHECK: arith.addi %{{.*}}, %arg21 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+// A trailing sync_block_set closes the body before the terminator.
+// CHECK: hivm.hir.sync_block_set[<VECTOR>, <PIPE_S>, <PIPE_S>] flag = 15
+// CHECK: scf.yield %{{.*}}#0, %{{.*}}#0, %{{.*}}, %{{.*}}#1, %{{.*}}#1, %arg22, %arg23, %arg24 : i32, i32, i32, i32, i32, i32, i32, i32
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @pcb12_tc01_while_matmul_fill(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32 {tt.divisibility = 16 : i32}, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32, %arg16: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %cst = arith.constant {ssbuffer.block_id = 11 : i32} 0.000000e+00 : f16
+    %c1_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 1 : i32
+    %c0_i32 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : i32
+    %c0 = arith.constant {ssbuffer.block_id = 10 : i32} 0 : index
+    %c32 = arith.constant {ssbuffer.block_id = 10 : i32} 32 : index
+    %c64 = arith.constant {ssbuffer.block_id = 10 : i32} 64 : index
+    %cst_0 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[4, 1, 16, 16]> : tensor<4xi64>
+    %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[16, 4, 16]> : tensor<3xi64>
+    %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[1, 2, 16, 16]> : tensor<4xi64>
+    %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} dense<[32, 1, 16]> : tensor<3xi64>
+    scope.scope : () -> () {
+      %alloc = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32} : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_4 = memref.alloc() {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_4 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32} : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>
+      %0:5 = scf.while (%arg17 = %c0_i32, %arg18 = %c0_i32, %arg19 = %c0_i32, %arg20 = %c0_i32, %arg21 = %c0_i32) : (i32, i32, i32, i32, i32) -> (i32, i32, i32, i32, i32) {
+        %1 = arith.cmpi slt, %arg17, %arg7 {Undefined, ssbuffer.block_id = 4 : i32} : i32
+        scf.condition(%1) %arg17, %arg18, %arg19, %arg20, %arg21 : i32, i32, i32, i32, i32
+      } do {
+      ^bb0(%arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32, %arg21: i32):
+        %true = arith.constant true
+        %1 = scf.if %true -> (i32) {
+          %4 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32} : i32
+          %5 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 9 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %6 = arith.muli %5, %c1_i32 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %7 = arith.addi %6, %arg19 {ssbuffer.block_id = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          scf.yield %7 : i32
+        } else {
+          scf.yield %arg19 : i32
+        } {hivm.matmul_limited_in_cube, ssbuffer.if = 9 : i32}
+        %true_5 = arith.constant true
+        %2:2 = scf.if %true_5 -> (i32, i32) {
+          %4 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32} : i32
+          %5 = arith.index_cast %arg7 {ssbuffer.block_id = 7 : i32} : i32 to index
+          %6 = arith.maxsi %5, %c0 {ssbuffer.block_id = 7 : i32} : index
+          %7 = arith.minsi %6, %c32 {ssbuffer.block_id = 7 : i32} : index
+          %8 = arith.cmpi slt, %7, %c32 {ssbuffer.block_id = 7 : i32} : index
+          %9 = arith.index_cast %arg6 {ssbuffer.block_id = 7 : i32} : i32 to index
+          %10 = arith.maxsi %9, %c0 {ssbuffer.block_id = 7 : i32} : index
+          %11 = arith.minsi %10, %c64 {ssbuffer.block_id = 7 : i32} : index
+          %12 = arith.cmpi slt, %11, %c64 {ssbuffer.block_id = 7 : i32} : index
+          %13 = arith.muli %4, %c1_i32 {ssbuffer.block_id = 7 : i32} : i32
+          %14 = arith.addi %13, %arg17 {ssbuffer.block_id = 7 : i32} : i32
+          %15 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %16 = arith.muli %15, %c1_i32 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %17 = arith.addi %16, %arg20 {ssbuffer.block_id = 7 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          scf.yield %14, %17 : i32, i32
+        } else {
+          scf.yield %arg17, %arg20 : i32, i32
+        } {hivm.matmul_limited_in_cube, ssbuffer.if = 7 : i32}
+        %true_6 = arith.constant true
+        %3:2 = scf.if %true_6 -> (i32, i32) {
+          %4 = arith.index_cast %arg7 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+          %5 = arith.maxsi %4, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %6 = arith.minsi %5, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %7 = arith.cmpi slt, %6, %c32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %8 = arith.index_cast %arg6 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : i32 to index
+          %9 = arith.maxsi %8, %c0 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %10 = arith.minsi %9, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %11 = arith.cmpi slt, %10, %c64 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 7 : i32} : index
+          %alloc_7 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<32xf16>
+          %alloc_8 = memref.alloc() {ssbuffer.block_id = 8 : i32} : memref<64xf16>
+          scf.if %11 {
+            linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_8 : memref<64xf16>)
+          } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+          scf.if %7 {
+            linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%alloc_7 : memref<32xf16>)
+          } {hivm.unlikely_condition, ssbuffer.block_id = 8 : i32}
+          %12 = arith.muli %arg18, %arg8 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+          %13 = arith.index_cast %12 {ssbuffer.block_id = 8 : i32} : i32 to index
+          %reinterpret_cast = memref.reinterpret_cast %arg2 to offset: [%13], sizes: [32], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<32xf16, strided<[1], offset: ?>>
+          %subview = memref.subview %reinterpret_cast[0] [%6] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+          %subview_9 = memref.subview %alloc_7[0] [%6] [1] {ssbuffer.block_id = 8 : i32} : memref<32xf16> to memref<?xf16, strided<[1]>>
+          memref.copy %subview, %subview_9 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+          %14 = bufferization.to_tensor %alloc_7 restrict writable {ssbuffer.block_id = 8 : i32} : memref<32xf16> to tensor<32xf16>
+          %15 = arith.muli %arg18, %arg9 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+          %16 = arith.index_cast %15 {ssbuffer.block_id = 8 : i32} : i32 to index
+          %reinterpret_cast_10 = memref.reinterpret_cast %arg3 to offset: [%16], sizes: [64], strides: [1] {ssbuffer.block_id = 8 : i32} : memref<?xf16> to memref<64xf16, strided<[1], offset: ?>>
+          %subview_11 = memref.subview %reinterpret_cast_10[0] [%10] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1], offset: ?>>
+          %subview_12 = memref.subview %alloc_8[0] [%10] [1] {ssbuffer.block_id = 8 : i32} : memref<64xf16> to memref<?xf16, strided<[1]>>
+          memref.copy %subview_11, %subview_12 {ssbuffer.block_id = 8 : i32} : memref<?xf16, strided<[1], offset: ?>> to memref<?xf16, strided<[1]>>
+          %17 = bufferization.to_tensor %alloc_8 restrict writable {ssbuffer.block_id = 8 : i32} : memref<64xf16> to tensor<64xf16>
+          %expanded = tensor.expand_shape %14 [[0, 1]] output_shape [32, 1] {ssbuffer.block_id = 8 : i32} : tensor<32xf16> into tensor<32x1xf16>
+          %18 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<32x16xf16>
+          %19 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%18 : tensor<32x16xf16>) -> tensor<32x16xf16>
+          %inserted_slice = tensor.insert_slice %expanded into %19[0, 0] [32, 1] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<32x1xf16> into tensor<32x16xf16>
+          %expanded_13 = tensor.expand_shape %17 [[0, 1]] output_shape [1, 64] {ssbuffer.block_id = 8 : i32} : tensor<64xf16> into tensor<1x64xf16>
+          %20 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<16x64xf16>
+          %21 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst : f16) outs(%20 : tensor<16x64xf16>) -> tensor<16x64xf16>
+          %inserted_slice_14 = tensor.insert_slice %expanded_13 into %21[0, 0] [1, 64] [1, 1] {ssbuffer.block_id = 8 : i32} : tensor<1x64xf16> into tensor<16x64xf16>
+          %reshape = tensor.reshape %inserted_slice(%cst_3) {ssbuffer.block_id = 8 : i32} : (tensor<32x16xf16>, tensor<3xi64>) -> tensor<32x1x16xf16>
+          %22 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<1x32x16xf16>
+          %transposed = linalg.transpose ins(%reshape : tensor<32x1x16xf16>) outs(%22 : tensor<1x32x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+          %reshape_15 = tensor.reshape %transposed(%cst_2) {ssbuffer.block_id = 8 : i32} : (tensor<1x32x16xf16>, tensor<4xi64>) -> tensor<1x2x16x16xf16>
+          %reshape_16 = tensor.reshape %inserted_slice_14(%cst_1) {ssbuffer.block_id = 8 : i32} : (tensor<16x64xf16>, tensor<3xi64>) -> tensor<16x4x16xf16>
+          %23 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<4x16x16xf16>
+          %transposed_17 = linalg.transpose ins(%reshape_16 : tensor<16x4x16xf16>) outs(%23 : tensor<4x16x16xf16>) permutation = [1, 0, 2]  {ssbuffer.block_id = 8 : i32}
+          %reshape_18 = tensor.reshape %transposed_17(%cst_0) {ssbuffer.block_id = 8 : i32} : (tensor<4x16x16xf16>, tensor<4xi64>) -> tensor<4x1x16x16xf16>
+          hivm.hir.copy ins(%reshape_15 : tensor<1x2x16x16xf16>) outs(%alloc : memref<1x2x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}
+          hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 0 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 1
+          hivm.hir.copy ins(%reshape_18 : tensor<4x1x16x16xf16>) outs(%alloc_4 : memref<4x1x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}
+          %24 = arith.addi %c1_i32, %c1_i32 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32, ssbuffer.clone = 9 : i32} : i32
+          %25 = arith.muli %24, %c1_i32 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+          %26 = arith.addi %25, %arg18 {ssbuffer.arg = 0 : i32, ssbuffer.block_id = 8 : i32} : i32
+          %27 = arith.addi %c1_i32, %c1_i32 {ssbuffer.block_id = 8 : i32, ssbuffer.clone = 9 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %28 = arith.muli %27, %c1_i32 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          %29 = arith.addi %28, %arg21 {ssbuffer.block_id = 8 : i32, ssbuffer.while_arg = 0 : i32} : i32
+          hivm.hir.sync_block_set {ssbuffer.analyze_flag_id, ssbuffer.block_id = 8 : i32, ssbuffer.transfer_id = 1 : i32}[<VECTOR>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+          scf.yield %26, %29 : i32, i32
+        } else {
+          scf.yield %arg18, %arg21 : i32, i32
+        } {hivm.matmul_limited_in_cube, ssbuffer.if = 8 : i32}
+        scf.yield %2#0, %3#0, %1, %2#1, %3#1 : i32, i32, i32, i32, i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.matmul_limited_in_cube, hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+    return
+  }
+}

From 6a3520907a9046044960eeae6d797a127805cc2d Mon Sep 17 00:00:00 2001
From: sxm <sunxianmiao@h-partners.com>
Date: Wed, 5 Aug 2026 18:08:29 +0800
Subject: [PATCH 09/11] [ssbuffer](fix) add condtion support whileop

---
 .../AddControlFlowCondition.h                 |   2 +-
 .../UpdateConditionInfo.h                     |  33 +-
 .../UpdateConditionInfo.cpp                   | 653 +++++++++++++-----
 .../while-update-condition-counter.mlir       |  94 +++
 .../while-update-condition-crosscore.mlir     | 112 +++
 .../while-update-condition-intracore.mlir     |  90 +++
 ...while-update-condition-tensor-iterarg.mlir | 106 +++
 7 files changed, 897 insertions(+), 193 deletions(-)
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-counter.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-crosscore.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-intracore.mlir
 create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-tensor-iterarg.mlir

diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
index 9d5681c132..d13330df77 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h
@@ -78,7 +78,7 @@ struct ControlFlowConditionInfo {
   llvm::DenseMap<Operation *, llvm::DenseMap<Value, SmallVector<int>>>
       tensorIterArgIndicesMap;
 
-  // unique counter value for each ifblock
+  // unique counter value for each ifblock scf.for only.
   llvm::DenseMap<scf::IfOp, Value> cntArgs;
 
   // DAG for if block cross-core dependencies
diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h
index 20129dcbe7..72a77f02f8 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h
@@ -64,6 +64,12 @@ class UpdateConditionInfoPass
   int updateIfConds(ModuleOp module,
                     SmallVector<SmallVector<Value>> ssbufferPtrs);
 
+  // Collect ssbuffer ifOps: for walks body; while-do walks after-region only.
+  int collectSSBufferIfOps(Operation *loopOp, SmallVector<scf::IfOp> &ifOps);
+
+  // Validate blockCounters for for; while skips (counters are for-only).
+  int validateBlockCounters(Operation *loopOp, size_t ifOpCount);
+
   void updateForIterTimes(ModuleOp module);
 
   scf::ForOp extendForOpIterationCount(scf::ForOp oldForOp, int ifCount,
@@ -75,18 +81,18 @@ class UpdateConditionInfoPass
                                   scf::ForOp oldForOp, scf::ForOp newForOp,
                                   IRMapping &mapper);
 
-  Value getVarValue(scf::ForOp forOp, int varIndex);
+  Value getVarValue(Operation *loopOp, int varIndex);
 
   void collectDependencyBuffers(
-      ModuleOp module, SmallVector<scf::ForOp> &mainLoopForOps,
+      ModuleOp module, SmallVector<Operation *> &mainLoopOps,
       DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
           &crossCoreBuffers,
-      DenseMap<scf::ForOp,
+      DenseMap<Operation *,
                DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>>
           &intraCoreBuffersMap);
 
   int buildIdxToVarMap(
-      scf::ForOp forOp,
+      Operation *loopOp,
       const DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
           &intraCoreBuffers,
       DenseMap<int, Value> &idxToVar);
@@ -124,7 +130,7 @@ class UpdateConditionInfoPass
       DenseMap<Value, VarUpdateType> &varUpdateTypes);
 
   // Build the ifOp variable mapping for the tensor iter_args
-  int buildTensorIterArgIfOpVarMap(scf::ForOp forOp);
+  int buildTensorIterArgIfOpVarMap(Operation *loopOp);
 
   // Collect the consumption conditions of the tensor iter_args consumer
   void collectTensorIterArgInputConditions(
@@ -151,8 +157,7 @@ class UpdateConditionInfoPass
                             bool hasCounter, Value counter, Value step);
 
   void populateNewElseBlock(scf::IfOp newIfOp, scf::IfOp oldIfOp,
-                            bool needsYield, bool oldHasElse, bool hasCounter,
-                            Value counter);
+                            bool oldHasElse, bool hasCounter, Value counter);
 
   scf::IfOp
   createNewIfOpWithBlocks(scf::IfOp oldIfOp, Value combinedCond,
@@ -172,9 +177,16 @@ class UpdateConditionInfoPass
 
   int updateForOpYield(scf::ForOp forOp);
 
+  // Update after-region yield for while when control vars were rewritten.
+  int updateWhileOpYield(scf::WhileOp whileOp);
+
+  // Dispatch yield update for scf.for / scf.while main_loop.
+  int updateLoopYield(Operation *loopOp);
+
+  // loopOp is scf.for or scf.while main_loop.
   int combineConditions(ModuleOp module, Value crossCoreCond,
                         Value intraCoreCond, Value flowOptCond, scf::IfOp ifOp,
-                        scf::ForOp forOp, size_t &usedCounterNum,
+                        Operation *loopOp, size_t &usedCounterNum,
                         DenseMap<Value, VarUpdateType> &varUpdateTypes);
 
   int setCrossCoreCondition(
@@ -185,8 +197,9 @@ class UpdateConditionInfoPass
       scf::IfOp ifOp, SmallVector<SmallVector<Value>> ssbufferPtrs,
       Value &crossCoreCond);
 
-  // Set the FlowOpt extra condition for the third if block in the DAG
-  int setFlowOptCondition(scf::IfOp currentIfOp, scf::ForOp forOp,
+  // Set the FlowOpt extra condition for the third if block in the DAG.
+  // Needs lb/ub/step from scf.for; scf.while leaves flowOptCond null.
+  int setFlowOptCondition(scf::IfOp currentIfOp, Operation *loopOp,
                           Value &flowOptCond);
 
   // Update DAG nodes after ifOp replacement
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp
index 03cbd4aef7..3f17209eee 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp
@@ -21,6 +21,7 @@
  */
 
 #include <functional>
+#include <optional>
 #include <string>
 #include <type_traits>
 
@@ -34,6 +35,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/ValueRange.h"
 
@@ -89,6 +91,148 @@ static void logOutputGroupValues(llvm::StringRef label,
   LDBG(os.str());
 }
 
+// Read block id from ssbuffer.if on ifOp. Missing attr is unexpected.
+static int getIfBlockId(scf::IfOp ifOp, int &outBlockId) {
+  auto attr = ifOp->getAttrOfType<IntegerAttr>(kIf);
+  if (!attr) {
+    LDBG("ssbuffer.if missing block id on ifOp: " << ifOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+  outBlockId = static_cast<int>(attr.getInt());
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+// for: body region iter args; while-do: after-region args.
+static int getLoopRegionIterArgs(Operation *loopOp,
+                                 MutableArrayRef<BlockArgument> &outArgs) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loopOp)) {
+    outArgs = forOp.getRegionIterArgs();
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp)) {
+    outArgs = whileOp.getAfterArguments();
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+  LDBG("getLoopRegionIterArgs expects scf.for or scf.while, got " << *loopOp
+                                                                  << "\n");
+  return UPDATE_CONDITION_INFO_FAILED;
+}
+
+static int getLoopRegionIterArg(Operation *loopOp, int argIdx, Value &outArg) {
+  MutableArrayRef<BlockArgument> args;
+  if (getLoopRegionIterArgs(loopOp, args) == UPDATE_CONDITION_INFO_FAILED)
+    return UPDATE_CONDITION_INFO_FAILED;
+  if (argIdx < 0 || argIdx >= static_cast<int>(args.size())) {
+    LDBG("Invalid loop region iter arg index: "
+         << argIdx << ", iter args " << args.size() << ", loopOp=" << *loopOp
+         << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+  outArg = args[argIdx];
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+// Clone only the SSA def-chain that produces `condition` when it lives in
+// beforeRegion.
+static int cloneConditionDefChain(Value value, Region &beforeRegion,
+                                  IRMapping &mapping, OpBuilder &builder,
+                                  Value &outValue) {
+  if (Value mapped = mapping.lookupOrNull(value)) {
+    outValue = mapped;
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
+  if (auto blockArg = dyn_cast<BlockArgument>(value)) {
+    if (blockArg.getParentRegion() == &beforeRegion) {
+      LDBG("Before-region arg is not in whileBlockArgMap mapping: " << value
+                                                                    << "\n");
+      return UPDATE_CONDITION_INFO_FAILED;
+    }
+    outValue = value;
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
+  Operation *defOp = value.getDefiningOp();
+  if (!defOp) {
+    outValue = value;
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+  // Defined outside before-region: reuse as-is (constants, outer values).
+  if (defOp->getParentRegion() != &beforeRegion) {
+    outValue = value;
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
+  for (Value operand : defOp->getOperands()) {
+    Value remappedOperand;
+    if (cloneConditionDefChain(operand, beforeRegion, mapping, builder,
+                               remappedOperand) == UPDATE_CONDITION_INFO_FAILED)
+      return UPDATE_CONDITION_INFO_FAILED;
+    if (!mapping.lookupOrNull(operand))
+      mapping.map(operand, remappedOperand);
+  }
+
+  Operation *cloned = builder.clone(*defOp, mapping);
+  outValue = cloned->getResult(cast<OpResult>(value).getResultNumber());
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+// Remap scf.condition(x): clone only ops that produce x into the after region,
+// remapping before args via whileBlockArgMap[while][blockId]:
+// {new_arg_idx(after) -> old_arg_idx(before)}.
+static int buildWhileCounterCondition(
+    scf::WhileOp whileOp, scf::IfOp ifOp, ControlFlowConditionInfo *info,
+    OpBuilder &builder, const DenseMap<Value, Value> &controlVarToLatestValue,
+    Value &outCond) {
+  int blockId;
+  if (getIfBlockId(ifOp, blockId) == UPDATE_CONDITION_INFO_FAILED)
+    return UPDATE_CONDITION_INFO_FAILED;
+
+  auto whileIt = info->whileBlockArgMap.find(whileOp);
+  if (whileIt == info->whileBlockArgMap.end()) {
+    LDBG("whileBlockArgMap has no entry for whileOp=" << whileOp << ", ifOp="
+                                                      << ifOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+  auto blockIt = whileIt->second.find(blockId);
+  if (blockIt == whileIt->second.end()) {
+    LDBG("whileBlockArgMap has no entry for blockId "
+         << blockId << ", whileOp=" << whileOp << ", ifOp=" << ifOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+
+  const DenseMap<int, int> &argIdxMap = blockIt->second;
+  auto beforeArgs = whileOp.getBeforeArguments();
+  auto afterArgs = whileOp.getAfterArguments();
+
+  IRMapping mapping;
+  for (auto [newArgIdx, oldArgIdx] : argIdxMap) {
+    if (oldArgIdx < 0 || oldArgIdx >= static_cast<int>(beforeArgs.size()) ||
+        newArgIdx < 0 || newArgIdx >= static_cast<int>(afterArgs.size())) {
+      LDBG("Invalid whileBlockArgMap entry (new="
+           << newArgIdx << ", old=" << oldArgIdx << "), whileOp=" << whileOp
+           << ", ifOp=" << ifOp << "\n");
+      return UPDATE_CONDITION_INFO_FAILED;
+    }
+    Value afterArg = afterArgs[newArgIdx];
+    auto latestIt = controlVarToLatestValue.find(afterArg);
+    if (latestIt != controlVarToLatestValue.end())
+      afterArg = latestIt->second;
+    mapping.map(beforeArgs[oldArgIdx], afterArg);
+  }
+
+  // Only take condition value x from scf.condition(x);
+  scf::ConditionOp condOp = whileOp.getConditionOp();
+  Value beforeCond = condOp.getCondition();
+  if (cloneConditionDefChain(beforeCond, whileOp.getBefore(), mapping, builder,
+                             outCond) == UPDATE_CONDITION_INFO_FAILED) {
+    LDBG("Failed to remap while before-region condition expression, whileOp="
+         << whileOp << ", ifOp=" << ifOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
 // Allocate the SSBuffer pointer
 SmallVector<SmallVector<Value>>
 UpdateConditionInfoPass::allocSSBuffer(ModuleOp module) {
@@ -140,10 +284,10 @@ UpdateConditionInfoPass::allocSSBuffer(ModuleOp module) {
 
 // Collect dependency buffer
 void UpdateConditionInfoPass::collectDependencyBuffers(
-    ModuleOp module, SmallVector<scf::ForOp> &mainLoopForOps,
+    ModuleOp module, SmallVector<Operation *> &mainLoopOps,
     DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
         &crossCoreBuffers,
-    DenseMap<scf::ForOp,
+    DenseMap<Operation *,
              DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>>
         &intraCoreBuffersMap) {
   // Collect crossCoreBuffers by traversing module in deterministic order
@@ -159,18 +303,18 @@ void UpdateConditionInfoPass::collectDependencyBuffers(
     return WalkResult::advance();
   });
 
-  // Collect intraCoreBuffers for all forOps
-  for (scf::ForOp forOp : mainLoopForOps) {
-    if (info->intraCoreDependentMap.count(forOp)) {
-      auto &forOpDeps = info->intraCoreDependentMap[forOp];
+  // Collect intraCoreBuffers for all main_loop for/while ops
+  for (Operation *loopOp : mainLoopOps) {
+    if (info->intraCoreDependentMap.count(loopOp)) {
+      auto &loopDeps = info->intraCoreDependentMap[loopOp];
       DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
           intraCoreBuffers;
       int intraCoreIdx = 0;
-      for (auto &entry : forOpDeps) {
+      for (auto &entry : loopDeps) {
         intraCoreBuffers[intraCoreIdx][entry.first] = entry.second;
         intraCoreIdx++;
       }
-      intraCoreBuffersMap[forOp] = intraCoreBuffers;
+      intraCoreBuffersMap[loopOp] = intraCoreBuffers;
     }
   }
 }
@@ -202,18 +346,22 @@ int addEquivalentOps(Operation *op, SmallVector<Operation *> &tcbOps,
 }
 
 int UpdateConditionInfoPass::buildIdxToVarMap(
-    scf::ForOp forOp,
+    Operation *loopOp,
     const DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
         &intraCoreBuffers,
     DenseMap<int, Value> &idxToVar) {
   int varIdx = 0;
-  int iterArgNum = static_cast<int>(forOp.getNumRegionIterArgs());
+  MutableArrayRef<BlockArgument> regionIterArgs;
+  if (getLoopRegionIterArgs(loopOp, regionIterArgs) ==
+      UPDATE_CONDITION_INFO_FAILED)
+    return UPDATE_CONDITION_INFO_FAILED;
+  int iterArgNum = static_cast<int>(regionIterArgs.size());
 
-  const auto &innerDepIndices = info->innerDepConds[forOp];
+  const auto &innerDepIndices = info->innerDepConds[loopOp];
   if (innerDepIndices.size() < intraCoreBuffers.size()) {
     LDBG("Not enough inner dependency condition indices: assigned "
          << innerDepIndices.size() << ", expected " << intraCoreBuffers.size()
-         << "\n");
+         << ", loopOp=" << *loopOp << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -222,12 +370,13 @@ int UpdateConditionInfoPass::buildIdxToVarMap(
 
     int argIdx = innerDepIndices[varIdx];
     if (argIdx < 0 || argIdx >= iterArgNum) {
-      LDBG("Invalid inner dependency arg index: " << argIdx << ", iter args "
-                                                  << iterArgNum << "\n");
+      LDBG("Invalid inner dependency arg index: "
+           << argIdx << ", iter args " << iterArgNum << ", loopOp=" << *loopOp
+           << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
 
-    idxToVar[idx] = forOp.getRegionIterArgs()[argIdx];
+    idxToVar[idx] = regionIterArgs[argIdx];
     LDBG("Assign intraCore buffer group " << idx << " to iter arg index "
                                           << argIdx << "\n");
     varIdx++;
@@ -367,13 +516,17 @@ int UpdateConditionInfoPass::getInputOutputValues(
   return UPDATE_CONDITION_INFO_SUCCESS;
 }
 
-Value UpdateConditionInfoPass::getVarValue(scf::ForOp forOp, int varIndex) {
-  if (!info->innerDepConds.count(forOp))
+Value UpdateConditionInfoPass::getVarValue(Operation *loopOp, int varIndex) {
+  if (!info->innerDepConds.count(loopOp))
     return Value();
-  SmallVector<int> &innerDepIndices = info->innerDepConds[forOp];
+  SmallVector<int> &innerDepIndices = info->innerDepConds[loopOp];
   if (varIndex < (int)innerDepIndices.size()) {
     int argIdx = innerDepIndices[varIndex];
-    return forOp.getRegionIterArgs()[argIdx];
+    Value arg;
+    if (getLoopRegionIterArg(loopOp, argIdx, arg) ==
+        UPDATE_CONDITION_INFO_FAILED)
+      return Value();
+    return arg;
   }
   return Value();
 }
@@ -776,19 +929,19 @@ int UpdateConditionInfoPass::collectIntraCoreOutputConditions(
 }
 
 // Build the ifOp variable mapping for the tensor iter_args
-int UpdateConditionInfoPass::buildTensorIterArgIfOpVarMap(scf::ForOp forOp) {
+int UpdateConditionInfoPass::buildTensorIterArgIfOpVarMap(Operation *loopOp) {
   // Clear any previous data
   tensorIterArgIfOpVars.clear();
 
-  if (!info->tensorIterArgDepsMap.count(forOp) ||
-      !info->tensorIterArgIndicesMap.count(forOp)) {
-    LDBG("Skip buildTensorIterArgIfOpVarMap: no tensor iter_args info for this "
-         "forOp\n");
+  if (!info->tensorIterArgDepsMap.count(loopOp) ||
+      !info->tensorIterArgIndicesMap.count(loopOp)) {
+    LDBG("Skip buildTensorIterArgIfOpVarMap: no tensor iter_args info for "
+         << *loopOp << "\n");
     return UPDATE_CONDITION_INFO_SUCCESS;
   }
 
-  auto &depsVec = info->tensorIterArgDepsMap[forOp];
-  auto &indicesMap = info->tensorIterArgIndicesMap[forOp];
+  auto &depsVec = info->tensorIterArgDepsMap[loopOp];
+  auto &indicesMap = info->tensorIterArgIndicesMap[loopOp];
 
   llvm::DenseMap<scf::IfOp, llvm::DenseSet<Value>> producerVars;
   llvm::DenseMap<scf::IfOp, llvm::DenseSet<Value>> consumerVars;
@@ -798,22 +951,32 @@ int UpdateConditionInfoPass::buildTensorIterArgIfOpVarMap(scf::ForOp forOp) {
     TensorIterArgIfOpRelation &relation = depEntry;
 
     if (!indicesMap.count(origIterArg)) {
-      LDBG("[Error]: origIterArg not found in indicesMap\n");
+      LDBG("[Error]: origIterArg not found in indicesMap, loopOp=" << *loopOp
+                                                                   << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
     SmallVector<int> &argIndices = indicesMap[origIterArg];
 
     if (relation.consumers.size() != argIndices.size()) {
       LDBG("[Error]: consumers size mismatch: "
-           << relation.consumers.size() << " vs " << argIndices.size() << "\n");
+           << relation.consumers.size() << " vs " << argIndices.size()
+           << ", loopOp=" << *loopOp << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
 
     // Establish a mapping (one-to-one) from consumers to variables
+    // for: region iter args; while-do: after-region args.
     llvm::DenseMap<scf::IfOp, Value> consumerToVar;
     for (size_t i = 0; i < relation.consumers.size(); ++i) {
       scf::IfOp consumer = relation.consumers[i];
-      Value var = forOp.getRegionIterArg(argIndices[i]);
+      Value var;
+      if (getLoopRegionIterArg(loopOp, argIndices[i], var) ==
+          UPDATE_CONDITION_INFO_FAILED) {
+        LDBG("[Error]: invalid tensor iter_arg index "
+             << argIndices[i] << ", loopOp=" << *loopOp
+             << ", consumer=" << consumer << "\n");
+        return UPDATE_CONDITION_INFO_FAILED;
+      }
       consumerToVar[consumer] = var;
     }
 
@@ -960,8 +1123,14 @@ int UpdateConditionInfoPass::setIntraCoreCondition(
 
 // Set the FlowOpt extra condition for the third if block in the DAG
 int UpdateConditionInfoPass::setFlowOptCondition(scf::IfOp currentIfOp,
-                                                 scf::ForOp forOp,
+                                                 Operation *loopOp,
                                                  Value &flowOptCond) {
+  auto forOp = dyn_cast<scf::ForOp>(loopOp);
+  if (!forOp) {
+    flowOptCond = nullptr;
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
   // Check if current ifOp is a target node (third node) in flowOptIfOpPairs
   if (!info->flowOptIfOpPairs.count(currentIfOp)) {
     LDBG("Current ifOp is not a flowOpt target node, skip.");
@@ -982,8 +1151,9 @@ int UpdateConditionInfoPass::setFlowOptCondition(scf::IfOp currentIfOp,
   scf::IfOp sourceIfOp = info->flowOptIfOpPairs[currentIfOp];
   if (!info->cntArgs.count(sourceIfOp)) {
     LDBG("[Error] Start node has no counter in cntArgs, cannot build flowOpt "
-         "condition. "
-         << "sourceIfOp: " << *sourceIfOp);
+         "condition. currentIfOp="
+         << currentIfOp << ", sourceIfOp=" << sourceIfOp << ", forOp=" << forOp
+         << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -1089,10 +1259,9 @@ void UpdateConditionInfoPass::updateControlVarToLatestValue(scf::IfOp newIfOp,
 
 // Update the yield in the forOp
 int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) {
-  LDBG("Enter update forOp yield " << "\n");
   if (controlVarToLatestValue.empty()) {
-    LDBG("Failed to update forOp yield: no latest control variable values."
-         << "\n");
+    LDBG("Failed to update forOp yield: no latest control variable values, "
+         << "forOp=" << forOp << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -1100,7 +1269,8 @@ int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) {
   Block *forBody = forOp.getBody();
   auto yieldOp = dyn_cast<scf::YieldOp>(forBody->getTerminator());
   if (!yieldOp) {
-    LDBG("Failed to update forOp yield: terminator is not scf.yield." << "\n");
+    LDBG("Failed to update forOp yield: terminator is not scf.yield, forOp="
+         << forOp << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -1109,7 +1279,7 @@ int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) {
   if (newYieldOperands.size() != forOp.getNumRegionIterArgs()) {
     LDBG("Failed to update forOp yield: yield operands "
          << newYieldOperands.size() << ", iter args "
-         << forOp.getNumRegionIterArgs() << "\n");
+         << forOp.getNumRegionIterArgs() << ", forOp=" << forOp << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -1124,8 +1294,8 @@ int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) {
     auto it = iterArgToIndex.find(origVar);
     if (it == iterArgToIndex.end()) {
       LDBG("Failed to update forOp yield: control variable is not a region "
-           "iter arg."
-           << "\n");
+           "iter arg, forOp="
+           << forOp << ", var=" << origVar << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
     newYieldOperands[it->second] = latestValue;
@@ -1137,7 +1307,130 @@ int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) {
   yieldOp.erase();
   LDBG("Updated forOp yield with " << controlVarToLatestValue.size()
                                    << " latest control values." << "\n");
-  LDBG("Exit update forOp yield " << "\n");
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+int UpdateConditionInfoPass::updateWhileOpYield(scf::WhileOp whileOp) {
+  LDBG("Enter update whileOp yield " << "\n");
+  if (controlVarToLatestValue.empty()) {
+    LDBG("Skip update whileOp yield: no latest control variable values, "
+         << "whileOp=" << whileOp << "\n");
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
+  Location loc = whileOp.getLoc();
+  Block *afterBody = whileOp.getAfterBody();
+  auto yieldOp = dyn_cast<scf::YieldOp>(afterBody->getTerminator());
+  if (!yieldOp) {
+    LDBG("Failed to update whileOp yield: terminator is not scf.yield, "
+         << "whileOp=" << whileOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+
+  SmallVector<Value> newYieldOperands(yieldOp.getOperands().begin(),
+                                      yieldOp.getOperands().end());
+  auto afterArgs = whileOp.getAfterArguments();
+  if (newYieldOperands.size() != afterArgs.size()) {
+    LDBG("Failed to update whileOp yield: yield operands "
+         << newYieldOperands.size() << ", after args " << afterArgs.size()
+         << ", whileOp=" << whileOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+
+  DenseMap<Value, unsigned> iterArgToIndex;
+  for (auto [idx, arg] : llvm::enumerate(afterArgs))
+    iterArgToIndex[arg] = idx;
+
+  for (auto &entry : controlVarToLatestValue) {
+    Value origVar = entry.first;
+    Value latestValue = entry.second;
+    auto it = iterArgToIndex.find(origVar);
+    if (it == iterArgToIndex.end()) {
+      LDBG("Failed to update whileOp yield: control variable is not an after "
+           "region iter arg, whileOp="
+           << whileOp << ", var=" << origVar << "\n");
+      return UPDATE_CONDITION_INFO_FAILED;
+    }
+    newYieldOperands[it->second] = latestValue;
+    LDBG("Update whileOp yield operand index " << it->second << "\n");
+  }
+
+  OpBuilder yieldBuilder(yieldOp);
+  yieldBuilder.create<scf::YieldOp>(loc, newYieldOperands);
+  yieldOp.erase();
+  LDBG("Updated whileOp yield with " << controlVarToLatestValue.size()
+                                     << " latest control values." << "\n");
+  LDBG("Exit update whileOp yield " << "\n");
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+int UpdateConditionInfoPass::updateLoopYield(Operation *loopOp) {
+  if (auto forOp = dyn_cast<scf::ForOp>(loopOp))
+    return updateForOpYield(forOp);
+  if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp))
+    return updateWhileOpYield(whileOp);
+  LDBG("updateLoopYield expects scf.for or scf.while, got " << *loopOp << "\n");
+  return UPDATE_CONDITION_INFO_FAILED;
+}
+
+int UpdateConditionInfoPass::collectSSBufferIfOps(
+    Operation *loopOp, SmallVector<scf::IfOp> &ifOps) {
+  auto collectIfOps = [&](Operation *op) -> WalkResult {
+    if (!op->hasAttr(kIf)) {
+      return WalkResult::advance();
+    }
+
+    auto ifOp = dyn_cast<scf::IfOp>(op);
+    if (!ifOp) {
+      LDBG("Found unsupported ssbuffer if op: "
+           << op->getName() << " under loopOp=" << *loopOp << "\n");
+      return WalkResult::interrupt();
+    }
+
+    ifOps.push_back(ifOp);
+    return WalkResult::advance();
+  };
+
+  WalkResult ifWalkResult;
+  if (auto forOp = dyn_cast<scf::ForOp>(loopOp)) {
+    ifWalkResult = forOp.walk(collectIfOps);
+  } else if (auto whileOp = dyn_cast<scf::WhileOp>(loopOp)) {
+    // while-do: only after-region.
+    ifWalkResult = whileOp.getAfterBody()->walk(collectIfOps);
+  } else {
+    LDBG("collectSSBufferIfOps expects scf.for or scf.while, got " << *loopOp
+                                                                   << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+
+  if (ifWalkResult.wasInterrupted()) {
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+  return UPDATE_CONDITION_INFO_SUCCESS;
+}
+
+int UpdateConditionInfoPass::validateBlockCounters(Operation *loopOp,
+                                                   size_t ifOpCount) {
+  auto forOp = dyn_cast<scf::ForOp>(loopOp);
+  if (!forOp) {
+    // blockCounters / cntArgs are for-only.
+    return UPDATE_CONDITION_INFO_SUCCESS;
+  }
+
+  auto counterIt = info->blockCounters.find(forOp);
+  if (counterIt == info->blockCounters.end()) {
+    LDBG("Failed to assign counters for ssbuffer if ops: no counters for "
+         << "forOp=" << forOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
+
+  size_t counterNum = counterIt->second.size();
+  if (ifOpCount > counterNum) {
+    LDBG("Failed to assign counters for all ssbuffer if ops: if ops "
+         << ifOpCount << ", counters " << counterNum << ", forOp=" << forOp
+         << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
+  }
   return UPDATE_CONDITION_INFO_SUCCESS;
 }
 
@@ -1234,15 +1527,11 @@ void UpdateConditionInfoPass::populateNewThenBlock(
   thenBuilder.create<scf::YieldOp>(loc, thenYieldOperands);
 }
 
-void UpdateConditionInfoPass::populateNewElseBlock(
-    scf::IfOp newIfOp, scf::IfOp oldIfOp, bool needsYield, bool oldHasElse,
-    bool hasCounter, Value counter) {
-  if (!needsYield && !oldHasElse) {
-    LDBG("Skip populating else block: no yield needed and old if has no else."
-         << "\n");
-    return;
-  }
-
+void UpdateConditionInfoPass::populateNewElseBlock(scf::IfOp newIfOp,
+                                                   scf::IfOp oldIfOp,
+                                                   bool oldHasElse,
+                                                   bool hasCounter,
+                                                   Value counter) {
   Location loc = newIfOp.getLoc();
   Block &newElseBlock = newIfOp.getElseRegion().front();
   SmallVector<Value> oldElseYieldOperands;
@@ -1260,43 +1549,40 @@ void UpdateConditionInfoPass::populateNewElseBlock(
          << oldElseYieldOperands.size() << " old else yield operands." << "\n");
   }
 
-  if (needsYield) {
-    OpBuilder elseBuilder(&newElseBlock, newElseBlock.end());
-    SmallVector<Value> elseYieldOperands;
-    for (Value operand : oldElseYieldOperands) {
-      Value newOperand = operand;
-      auto it = controlVarToLatestValue.find(operand);
-      if (it != controlVarToLatestValue.end()) {
-        newOperand = it->second;
-      }
-      elseYieldOperands.push_back(newOperand);
+  OpBuilder elseBuilder(&newElseBlock, newElseBlock.end());
+  SmallVector<Value> elseYieldOperands;
+  for (Value operand : oldElseYieldOperands) {
+    Value newOperand = operand;
+    auto it = controlVarToLatestValue.find(operand);
+    if (it != controlVarToLatestValue.end()) {
+      newOperand = it->second;
     }
+    elseYieldOperands.push_back(newOperand);
+  }
 
-    for (Value var : currentUsedVars) {
-      Value varToUse = var;
-      auto it = controlVarToLatestValue.find(var);
-      if (it != controlVarToLatestValue.end()) {
-        varToUse = it->second;
-      }
-      elseYieldOperands.push_back(varToUse);
+  for (Value var : currentUsedVars) {
+    Value varToUse = var;
+    auto it = controlVarToLatestValue.find(var);
+    if (it != controlVarToLatestValue.end()) {
+      varToUse = it->second;
     }
+    elseYieldOperands.push_back(varToUse);
+  }
 
-    if (hasCounter) {
-      Value counterToUse = counter;
-      auto it = controlVarToLatestValue.find(counter);
-      if (it != controlVarToLatestValue.end()) {
-        counterToUse = it->second;
-      }
-      elseYieldOperands.push_back(counterToUse);
+  if (hasCounter) {
+    Value counterToUse = counter;
+    auto it = controlVarToLatestValue.find(counter);
+    if (it != controlVarToLatestValue.end()) {
+      counterToUse = it->second;
     }
+    elseYieldOperands.push_back(counterToUse);
+  }
 
-    LDBG("Create else yield with " << elseYieldOperands.size() << " operands."
-                                   << "\n");
-    elseBuilder.create<scf::YieldOp>(loc, elseYieldOperands);
-  } else if (oldElseYieldOp) {
+  LDBG("Create else yield with " << elseYieldOperands.size() << " operands."
+                                 << "\n");
+  elseBuilder.create<scf::YieldOp>(loc, elseYieldOperands);
+  if (oldElseYieldOp) {
     oldElseYieldOp->erase();
-    LDBG("Erase old else yield because new if does not need yield values."
-         << "\n");
   }
 }
 
@@ -1310,9 +1596,12 @@ scf::IfOp UpdateConditionInfoPass::createNewIfOpWithBlocks(
 
   bool needsYield = !currentUsedVars.empty() || hasCounter;
   bool oldHasElse = oldIfOp.getElseRegion().hasOneBlock();
+  // Only create else when we will populate it;
+  bool withElse = needsYield || oldHasElse;
   LDBG("Create replacement if op: needs yield "
-       << needsYield << ", old has else " << oldHasElse
-       << ", current used vars " << currentUsedVars.size() << "." << "\n");
+       << needsYield << ", old has else " << oldHasElse << ", with else "
+       << withElse << ", current used vars " << currentUsedVars.size() << "."
+       << "\n");
 
   Block &oldThenBlock = oldIfOp.getThenRegion().front();
   Operation *oldThenYieldOp = nullptr;
@@ -1321,7 +1610,7 @@ scf::IfOp UpdateConditionInfoPass::createNewIfOpWithBlocks(
   SmallVector<Type> resultTypes =
       buildNewIfResultTypes(oldIfOp, hasCounter, counter);
   scf::IfOp newIfOp =
-      builder.create<scf::IfOp>(loc, resultTypes, combinedCond, true);
+      builder.create<scf::IfOp>(loc, resultTypes, combinedCond, withElse);
   LDBG("Created replacement if op with " << resultTypes.size() << " results."
                                          << "\n");
 
@@ -1331,8 +1620,9 @@ scf::IfOp UpdateConditionInfoPass::createNewIfOpWithBlocks(
 
   populateNewThenBlock(newIfOp, oldThenBlock, oldThenYieldOp, oldYieldOperands,
                        varUpdateTypes, hasCounter, counter, step);
-  populateNewElseBlock(newIfOp, oldIfOp, needsYield, oldHasElse, hasCounter,
-                       counter);
+  if (withElse) {
+    populateNewElseBlock(newIfOp, oldIfOp, oldHasElse, hasCounter, counter);
+  }
 
   for (size_t i = 0; i < oldIfOp.getNumResults(); ++i) {
     oldIfOp.getResult(i).replaceAllUsesWith(newIfOp.getResult(i));
@@ -1346,12 +1636,13 @@ scf::IfOp UpdateConditionInfoPass::createNewIfOpWithBlocks(
 // condition + flowOpt condition
 int UpdateConditionInfoPass::combineConditions(
     ModuleOp module, Value crossCoreCond, Value intraCoreCond,
-    Value flowOptCond, scf::IfOp ifOp, scf::ForOp forOp, size_t &usedCounterNum,
-    DenseMap<Value, VarUpdateType> &varUpdateTypes) {
+    Value flowOptCond, scf::IfOp ifOp, Operation *loopOp,
+    size_t &usedCounterNum, DenseMap<Value, VarUpdateType> &varUpdateTypes) {
   Location loc = ifOp.getLoc();
   SmallVector<Value> validConditions;
   Value counter;
-  bool hasCounter = false;
+  // Only for updates counter args inside if then  (+step).
+  bool updateCounterArg = false;
 
   if (crossCoreCond) {
     validConditions.push_back(crossCoreCond);
@@ -1363,55 +1654,74 @@ int UpdateConditionInfoPass::combineConditions(
     validConditions.push_back(flowOptCond);
   }
 
-  if (!info->blockCounters.count(forOp)) {
-    LDBG("Missing block counters for forOp." << "\n");
-    return UPDATE_CONDITION_INFO_FAILED;
-  }
-
-  SmallVector<int> &counterIndices = info->blockCounters[forOp];
+  auto forOp = dyn_cast<scf::ForOp>(loopOp);
+  auto whileOp = dyn_cast<scf::WhileOp>(loopOp);
+  OpBuilder condBuilder(ifOp);
 
-  if (info->cntArgs.count(ifOp)) {
-    counter = info->cntArgs[ifOp];
-    hasCounter = true;
-  } else {
-    if (usedCounterNum >= counterIndices.size()) {
-      LDBG("Not enough counters for ssbuffer if ops: used "
-           << usedCounterNum << ", counters " << counterIndices.size() << "\n");
+  if (forOp) {
+    if (!info->blockCounters.count(forOp)) {
+      LDBG("Missing block counters for forOp=" << forOp << ", ifOp=" << ifOp
+                                               << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
 
-    int argIdx = counterIndices[usedCounterNum];
-    int iterArgNum = static_cast<int>(forOp.getNumRegionIterArgs());
-    if (argIdx < 0 || argIdx >= iterArgNum) {
-      LDBG("Invalid counter arg index: " << argIdx << ", iter args "
-                                         << iterArgNum << "\n");
-      return UPDATE_CONDITION_INFO_FAILED;
-    }
+    SmallVector<int> &counterIndices = info->blockCounters[forOp];
 
-    counter = forOp.getRegionIterArgs()[argIdx];
-    hasCounter = true;
-    info->cntArgs[ifOp] = counter;
-    usedCounterNum++;
-    LDBG("Assign counter iter arg index " << argIdx << " to ssbuffer if op."
-                                          << "\n");
-  }
+    if (info->cntArgs.count(ifOp)) {
+      counter = info->cntArgs[ifOp];
+      updateCounterArg = true;
+    } else {
+      if (usedCounterNum >= counterIndices.size()) {
+        LDBG("Not enough counters for ssbuffer if ops: used "
+             << usedCounterNum << ", counters " << counterIndices.size()
+             << ", forOp=" << forOp << ", ifOp=" << ifOp << "\n");
+        return UPDATE_CONDITION_INFO_FAILED;
+      }
 
-  LDBG("this ifop used counter is: " << counter << "\n");
-  if (hasCounter) {
-    OpBuilder builder(ifOp);
+      int argIdx = counterIndices[usedCounterNum];
+      int iterArgNum = static_cast<int>(forOp.getNumRegionIterArgs());
+      if (argIdx < 0 || argIdx >= iterArgNum) {
+        LDBG("Invalid counter arg index: " << argIdx << ", iter args "
+                                           << iterArgNum << ", forOp=" << forOp
+                                           << ", ifOp=" << ifOp << "\n");
+        return UPDATE_CONDITION_INFO_FAILED;
+      }
+
+      counter = forOp.getRegionIterArgs()[argIdx];
+      updateCounterArg = true;
+      info->cntArgs[ifOp] = counter;
+      usedCounterNum++;
+      LDBG("Assign counter iter arg index " << argIdx << " to ssbuffer if op."
+                                            << "\n");
+    }
+
+    LDBG("this ifop used counter is: " << counter << "\n");
     Value upperBound = forOp.getUpperBound();
     Value counterToUse = counter;
     auto latestIt = controlVarToLatestValue.find(counter);
     if (latestIt != controlVarToLatestValue.end()) {
       counterToUse = latestIt->second;
     }
-    Value counterCond = builder.create<arith::CmpIOp>(
+    Value counterCond = condBuilder.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::slt, counterToUse, upperBound);
     validConditions.push_back(counterCond);
+  } else if (whileOp) {
+    // Only add remapped scf.condition(x) expression;
+    Value counterCond;
+    if (buildWhileCounterCondition(whileOp, ifOp, info, condBuilder,
+                                   controlVarToLatestValue,
+                                   counterCond) == UPDATE_CONDITION_INFO_FAILED)
+      return UPDATE_CONDITION_INFO_FAILED;
+    validConditions.push_back(counterCond);
+  } else {
+    LDBG("Unsupported loop op for counter condition, loopOp="
+         << *loopOp << ", ifOp=" << ifOp << "\n");
+    return UPDATE_CONDITION_INFO_FAILED;
   }
 
   if (validConditions.empty()) {
-    LDBG("Failed to build any condition for ssbuffer if op." << "\n");
+    LDBG("Failed to build any condition for ssbuffer if op, ifOp="
+         << ifOp << ", loopOp=" << *loopOp << "\n");
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
@@ -1424,13 +1734,14 @@ int UpdateConditionInfoPass::combineConditions(
         builder.create<arith::AndIOp>(loc, combinedCond, validConditions[i]);
   }
 
+  Value step = updateCounterArg ? forOp.getStep() : Value();
   scf::IfOp newIfOp = createNewIfOpWithBlocks(
-      ifOp, combinedCond, varUpdateTypes, hasCounter, counter, forOp.getStep());
+      ifOp, combinedCond, varUpdateTypes, updateCounterArg, counter, step);
 
   // Update DAG nodes
   updateDAGAfterIfOpReplacement(ifOp, newIfOp);
 
-  if (hasCounter) {
+  if (updateCounterArg) {
     info->cntArgs.erase(ifOp);
     info->cntArgs[newIfOp] = counter;
   }
@@ -1444,8 +1755,8 @@ int UpdateConditionInfoPass::combineConditions(
   }
 
   // Update tensorIterArgDepsMap with new ifOp
-  if (info->tensorIterArgDepsMap.count(forOp)) {
-    auto &depsVec = info->tensorIterArgDepsMap[forOp];
+  if (info->tensorIterArgDepsMap.count(loopOp)) {
+    auto &depsVec = info->tensorIterArgDepsMap[loopOp];
     for (auto &relation : depsVec) {
       // Update producer ifOp - use pointer comparison
       if (relation.producer.getOperation() == ifOp.getOperation()) {
@@ -1460,7 +1771,7 @@ int UpdateConditionInfoPass::combineConditions(
     }
   }
 
-  updateControlVarToLatestValue(newIfOp, ifOp, hasCounter, counter);
+  updateControlVarToLatestValue(newIfOp, ifOp, updateCounterArg, counter);
 
   ifOp.erase();
   return UPDATE_CONDITION_INFO_SUCCESS;
@@ -1469,96 +1780,72 @@ int UpdateConditionInfoPass::combineConditions(
 // Update the conditions of ifOp.
 int UpdateConditionInfoPass::updateIfConds(
     ModuleOp module, SmallVector<SmallVector<Value>> ssbufferPtrs) {
-  // Walk the forOp in the module to update the conditions of ifOp
-  SmallVector<scf::ForOp> mainLoopForOps;
+  // Walk main_loop for/while ops to update the conditions of ifOp
+  SmallVector<Operation *> mainLoopOps;
   WalkResult walkResult = module.walk([&](Operation *op) -> WalkResult {
-    if (!op->hasAttr(kMainLoop)) {
+    if (CVPipeline::isMainLoopOp(op)) {
+      mainLoopOps.push_back(op);
       return WalkResult::advance();
     }
 
-    auto forOp = dyn_cast<scf::ForOp>(op);
-    if (!forOp) {
-      LDBG("Found unsupported main loop op: " << op->getName() << "\n");
+    if (op->hasAttr(kMainLoop)) {
+      LDBG("Found unsupported main loop op: " << *op << "\n");
       return WalkResult::interrupt();
     }
-    mainLoopForOps.push_back(forOp);
     return WalkResult::advance();
   });
   if (walkResult.wasInterrupted()) {
     return UPDATE_CONDITION_INFO_FAILED;
   }
 
-  // Step0: Collect dependency buffers once outside the for loop
+  // Step0: Collect dependency buffers once outside the loop
   DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
       crossCoreBuffers;
-  DenseMap<scf::ForOp,
+  DenseMap<Operation *,
            DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>>
       intraCoreBuffersMap;
-  collectDependencyBuffers(module, mainLoopForOps, crossCoreBuffers,
+  collectDependencyBuffers(module, mainLoopOps, crossCoreBuffers,
                            intraCoreBuffersMap);
 
-  for (scf::ForOp forOp : mainLoopForOps) {
+  for (Operation *loopOp : mainLoopOps) {
     controlVarToLatestValue.clear();
 
-    // Step 0: Build the ifOp variable mapping for the tensor iter_args
-    if (buildTensorIterArgIfOpVarMap(forOp) == UPDATE_CONDITION_INFO_FAILED) {
-      return UPDATE_CONDITION_INFO_FAILED;
-    }
-
     // Step1: Get intraCoreBuffers from pre-collected map
     DenseMap<int, DenseMap<Operation *, SmallVector<Operation *>>>
         intraCoreBuffers;
-    if (intraCoreBuffersMap.count(forOp)) {
-      intraCoreBuffers = intraCoreBuffersMap[forOp];
+    if (intraCoreBuffersMap.count(loopOp)) {
+      intraCoreBuffers = intraCoreBuffersMap[loopOp];
     }
 
+    // for/while requires at least one of cross/intra deps;
     if (crossCoreBuffers.empty() && intraCoreBuffers.empty()) {
-      LDBG("crossCoreBuffers and intraCoreBuffers are all empty!" << "\n");
+      LDBG("crossCoreBuffers and intraCoreBuffers are all empty! loopOp="
+           << *loopOp << "\n");
       return UPDATE_CONDITION_INFO_FAILED;
     }
 
-    // Step2:Assign a variable to each inputValue of this forOp
+    // Step2: Assign control variables (intraCore inputs / tensor iter_args)
     DenseMap<int, Value> idxToVar;
-    if (buildIdxToVarMap(forOp, intraCoreBuffers, idxToVar) ==
+    if (buildIdxToVarMap(loopOp, intraCoreBuffers, idxToVar) ==
         UPDATE_CONDITION_INFO_FAILED) {
       return UPDATE_CONDITION_INFO_FAILED;
     }
+    if (buildTensorIterArgIfOpVarMap(loopOp) == UPDATE_CONDITION_INFO_FAILED) {
+      return UPDATE_CONDITION_INFO_FAILED;
+    }
+
     size_t usedCounterNum = 0;
     SmallVector<scf::IfOp> ifOps;
-    WalkResult ifWalkResult = forOp.walk([&](Operation *op) -> WalkResult {
-      if (!op->hasAttr(kIf)) {
-        return WalkResult::advance();
-      }
-
-      auto ifOp = dyn_cast<scf::IfOp>(op);
-      if (!ifOp) {
-        LDBG("Found unsupported ssbuffer if op: " << op->getName() << "\n");
-        return WalkResult::interrupt();
-      }
-
-      ifOps.push_back(ifOp);
-      return WalkResult::advance();
-    });
-    if (ifWalkResult.wasInterrupted()) {
+    if (collectSSBufferIfOps(loopOp, ifOps) == UPDATE_CONDITION_INFO_FAILED) {
       return UPDATE_CONDITION_INFO_FAILED;
     }
-    auto counterIt = info->blockCounters.find(forOp);
-    if (counterIt == info->blockCounters.end()) {
-      LDBG("Failed to assign counters for ssbuffer if ops: no counters for "
-           "forOp."
-           << "\n");
+    if (validateBlockCounters(loopOp, ifOps.size()) ==
+        UPDATE_CONDITION_INFO_FAILED) {
       return UPDATE_CONDITION_INFO_FAILED;
     }
 
-    size_t counterNum = counterIt->second.size();
-    if (ifOps.size() > counterNum) {
-      LDBG("Failed to assign counters for all ssbuffer if ops: if ops "
-           << ifOps.size() << ", counters " << counterNum << "\n");
-      return UPDATE_CONDITION_INFO_FAILED;
-    }
-    // Update the conditions of ifOp in this forOp.
+    // Update the conditions of ifOp in this loop op.
     for (scf::IfOp ifOp : ifOps) {
-      // Walk the ifOp in this forOp to update the conditions of ifOp
       SmallVector<int> crossCoreInputValues;
       SmallVector<int> crossCoreOutputValues;
       SmallVector<int> intraCoreInputValues;
@@ -1568,7 +1855,8 @@ int UpdateConditionInfoPass::updateIfConds(
                                crossCoreInputValues, crossCoreOutputValues,
                                intraCoreInputValues,
                                intraCoreOutputValues) != 0) {
-        LDBG("getInputOutputValues failed!" << "\n");
+        LDBG("getInputOutputValues failed! ifOp=" << ifOp << ", loopOp="
+                                                  << *loopOp << "\n");
         return UPDATE_CONDITION_INFO_FAILED;
       }
 
@@ -1577,7 +1865,8 @@ int UpdateConditionInfoPass::updateIfConds(
       if (setCrossCoreCondition(crossCoreInputValues, crossCoreOutputValues,
                                 crossCoreBuffers, ifOp, ssbufferPtrs,
                                 crossCoreCond) != 0) {
-        LDBG("setCrossCoreCondition failed!" << "\n");
+        LDBG("setCrossCoreCondition failed! ifOp=" << ifOp << ", loopOp="
+                                                   << *loopOp << "\n");
         return UPDATE_CONDITION_INFO_FAILED;
       }
       // Step4:Set the intraCore condition
@@ -1589,22 +1878,22 @@ int UpdateConditionInfoPass::updateIfConds(
           UPDATE_CONDITION_INFO_FAILED) {
         return UPDATE_CONDITION_INFO_FAILED;
       }
-      // Step5:Set the flowOpt condition
+      // Step5:Set the flowOpt condition (for only; needs lb/ub/step)
       Value flowOptCond;
-      if (setFlowOptCondition(ifOp, forOp, flowOptCond) ==
+      if (setFlowOptCondition(ifOp, loopOp, flowOptCond) ==
           UPDATE_CONDITION_INFO_FAILED) {
         return UPDATE_CONDITION_INFO_FAILED;
       }
-      // Step6:Combine the conditions: crossCore condition + intraCore condition
-      // + counter condition + flowOpt condition
+      // Step6:Combine the conditions: crossCore + intraCore + counter +
+      // flowOpt
       if (combineConditions(module, crossCoreCond, intraCoreCond, flowOptCond,
-                            ifOp, forOp, usedCounterNum,
+                            ifOp, loopOp, usedCounterNum,
                             varUpdateTypes) == UPDATE_CONDITION_INFO_FAILED) {
         return UPDATE_CONDITION_INFO_FAILED;
       }
     }
-    // Step6:Update the yield variable of the forOp
-    if (updateForOpYield(forOp) == UPDATE_CONDITION_INFO_FAILED) {
+    // Step7: Update loop yield with latest control values
+    if (updateLoopYield(loopOp) == UPDATE_CONDITION_INFO_FAILED) {
       return UPDATE_CONDITION_INFO_FAILED;
     }
   }
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-counter.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-counter.mlir
new file mode 100644
index 0000000000..4be4179329
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-counter.mlir
@@ -0,0 +1,94 @@
+// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s
+
+// UpdateConditionInfo on scf.while: clone scf.condition(x) def-chain
+// (extsi + cmpi) into each ssbuffer.if via whileBlockArgMap.
+
+// CHECK-LABEL: func.func @while_update_condition_counter
+
+// While-do remains; before-region still has extsi/cmpi i64 predicates.
+// CHECK: scf.while
+// CHECK: arith.extsi %{{.*}} : i32 to i64
+// CHECK: arith.cmpi slt, %{{.*}} : i64
+// CHECK: scf.condition
+
+// Vector ssbuffer.if = 5: cross-core cond and-ed with cloned while counter.
+// CHECK: %[[EXT5:.*]] = arith.extsi %{{.*}} : i32 to i64
+// CHECK: %[[CMP5:.*]] = arith.cmpi slt, %[[EXT5]], %{{.*}} : i64
+// CHECK: %[[AND5:.*]] = arith.andi %{{.*}}, %[[CMP5]]
+// CHECK: scf.if %[[AND5]]
+// CHECK: } {{.*}}ssbuffer.if = 5
+
+// Vector ssbuffer.if = 6: same cloned counter shape.
+// CHECK: %[[EXT6:.*]] = arith.extsi %{{.*}} : i32 to i64
+// CHECK: %[[CMP6:.*]] = arith.cmpi slt, %[[EXT6]], %{{.*}} : i64
+// CHECK: %[[AND6:.*]] = arith.andi %{{.*}}, %[[CMP6]]
+// CHECK: scf.if %[[AND6]]
+// CHECK: } {{.*}}ssbuffer.if = 6
+
+// CHECK: scf.yield
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @while_update_condition_counter(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %c0_i32 = arith.constant 0 : i32
+    %c16_i32 = arith.constant 16 : i32
+    %bound = arith.extsi %arg7 : i32 to i64
+
+    scope.scope : () -> () {
+      %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub0 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub0 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      %alloc_ub1 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub1 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %memspacecast = memref.memory_space_cast %alloc_ub0 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space<ub>> to memref<128x128xf32>
+        %empty = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x8x16x16xf16>
+        hivm.hir.copy ins(%empty : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [2 : i32, 1 : i32]}
+        %next5 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 5 : i32} : i32
+
+        hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %memspacecast_1 = memref.memory_space_cast %alloc_ub1 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossCoreDeps = [1 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space<ub>> to memref<128x128xf32>
+        %next6 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 6 : i32} : i32
+        scf.yield %next6 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+
+    scope.scope : () -> () {
+      %alloc_c = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_c {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub2 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub2 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      %alloc_ub3 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub3 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %t0 = tensor.empty() {ssbuffer.block_id = 0 : i32} : tensor<128x128xf32>
+        hivm.hir.fixpipe {dma_mode = #hivm.dma_mode<nz2nd>, ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 1 : i32]} ins(%t0 : tensor<128x128xf32>) outs(%alloc_ub2 : memref<128x128xf32, #hivm.address_space<ub>>)
+        %next0 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 0 : i32} : i32
+
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %conv = hivm.hir.convert_layout %alloc_c output_shape [128, 128] {dstLayout = #hivm.data_layout<ND>, srcLayout = #hivm.data_layout<nZ>, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [2 : i32, 0 : i32]} : (memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) -> memref<128x128xf16, #hivm.address_space<cbuf>>
+        %t1 = tensor.empty() {ssbuffer.block_id = 1 : i32} : tensor<128x128xf32>
+        hivm.hir.fixpipe {dma_mode = #hivm.dma_mode<nz2nd>, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossCoreDeps = [1 : i32, 1 : i32]} ins(%t1 : tensor<128x128xf32>) outs(%alloc_ub3 : memref<128x128xf32, #hivm.address_space<ub>>)
+        %next1 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 1 : i32} : i32
+        scf.yield %next1 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-crosscore.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-crosscore.mlir
new file mode 100644
index 0000000000..106f4a17d3
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-crosscore.mlir
@@ -0,0 +1,112 @@
+// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s
+
+// UpdateConditionInfo on scf.while: cross-core (核间) buffer conditions.
+// Same dual-scope crossCoreDeps skeleton as while-update-condition-counter;
+// CHECKs focus on ssbuffer memref init, memref.load/cmpi before ssbuffer.if,
+// and load/addi|subi/store updates inside then (post-#1359 MLIR ssbuffer).
+
+// CHECK-LABEL: func.func @while_update_condition_crosscore
+
+// SSBuffer slots initialized to 0 via pointer_cast + memref.store.
+// CHECK: hivm.hir.pointer_cast(%{{.*}}) : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: memref.store %{{.*}}, %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: memref.store %{{.*}}, %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+
+// CHECK: scf.while
+
+// Vector if=5: consumer sgt 0 + producer slt limit, then and-ed with while counter.
+// CHECK: %[[LD_IN5:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: annotation.mark %[[LD_IN5]] {memref_ext.volatile}
+// CHECK: %[[SGT5:.*]] = arith.cmpi sgt, %[[LD_IN5]], %{{.*}} : i32
+// CHECK: %[[LD_OUT5:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: annotation.mark %[[LD_OUT5]] {memref_ext.volatile}
+// CHECK: %[[SLT5:.*]] = arith.cmpi slt, %[[LD_OUT5]], %{{.*}} : i32
+// CHECK: %[[CROSS5:.*]] = arith.andi %[[SGT5]], %[[SLT5]]
+// CHECK: %[[AND5:.*]] = arith.andi %[[CROSS5]], %{{.*}}
+// CHECK: scf.if %[[AND5]]
+// CHECK: %[[LD_DEC5:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: %[[SUB5:.*]] = arith.subi %[[LD_DEC5]], %{{.*}} : i32
+// CHECK: memref.store %[[SUB5]], %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: %[[LD_INC5:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: %[[ADD5:.*]] = arith.addi %[[LD_INC5]], %{{.*}} : i32
+// CHECK: memref.store %[[ADD5]], %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: } {{.*}}ssbuffer.if = 5
+
+// Vector if=6: consumer-only sgt 0, then and-ed with while counter.
+// CHECK: %[[LD_IN6:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: annotation.mark %[[LD_IN6]] {memref_ext.volatile}
+// CHECK: %[[SGT6:.*]] = arith.cmpi sgt, %[[LD_IN6]], %{{.*}} : i32
+// CHECK: %[[AND6:.*]] = arith.andi %[[SGT6]], %{{.*}}
+// CHECK: scf.if %[[AND6]]
+// CHECK: %[[LD_DEC6:.*]] = memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: %[[SUB6:.*]] = arith.subi %[[LD_DEC6]], %{{.*}} : i32
+// CHECK: memref.store %[[SUB6]], %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: } {{.*}}ssbuffer.if = 6
+
+// CHECK: scf.yield
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @while_update_condition_crosscore(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %c0_i32 = arith.constant 0 : i32
+    %c16_i32 = arith.constant 16 : i32
+    %bound = arith.extsi %arg7 : i32 to i64
+
+    scope.scope : () -> () {
+      %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub0 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub0 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      %alloc_ub1 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub1 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %memspacecast = memref.memory_space_cast %alloc_ub0 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space<ub>> to memref<128x128xf32>
+        %empty = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x8x16x16xf16>
+        hivm.hir.copy ins(%empty : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [2 : i32, 1 : i32]}
+        %next5 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 5 : i32} : i32
+
+        hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %memspacecast_1 = memref.memory_space_cast %alloc_ub1 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossCoreDeps = [1 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space<ub>> to memref<128x128xf32>
+        %next6 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 6 : i32} : i32
+        scf.yield %next6 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+
+    scope.scope : () -> () {
+      %alloc_c = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_c {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub2 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub2 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      %alloc_ub3 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub3 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %t0 = tensor.empty() {ssbuffer.block_id = 0 : i32} : tensor<128x128xf32>
+        hivm.hir.fixpipe {dma_mode = #hivm.dma_mode<nz2nd>, ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 1 : i32]} ins(%t0 : tensor<128x128xf32>) outs(%alloc_ub2 : memref<128x128xf32, #hivm.address_space<ub>>)
+        %next0 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 0 : i32} : i32
+
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %conv = hivm.hir.convert_layout %alloc_c output_shape [128, 128] {dstLayout = #hivm.data_layout<ND>, srcLayout = #hivm.data_layout<nZ>, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [2 : i32, 0 : i32]} : (memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) -> memref<128x128xf16, #hivm.address_space<cbuf>>
+        %t1 = tensor.empty() {ssbuffer.block_id = 1 : i32} : tensor<128x128xf32>
+        hivm.hir.fixpipe {dma_mode = #hivm.dma_mode<nz2nd>, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossCoreDeps = [1 : i32, 1 : i32]} ins(%t1 : tensor<128x128xf32>) outs(%alloc_ub3 : memref<128x128xf32, #hivm.address_space<ub>>)
+        %next1 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 1 : i32} : i32
+        scf.yield %next1 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-intracore.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-intracore.mlir
new file mode 100644
index 0000000000..61d2ac9469
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-intracore.mlir
@@ -0,0 +1,90 @@
+// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s
+
+// UpdateConditionInfo on scf.while: intra-core (核内) conditions.
+// VECTOR while has producer (intraDeps role=1) then consumer (role=0).
+// Actual shape (with optional cross gate on the same if):
+//   producer: cmpi slt control_var, 1 ; then addi control_var, 1
+//   consumer: cmpi sgt latest, 0      ; then subi latest, 1
+
+// CHECK-LABEL: func.func @while_update_condition_intracore
+
+// CHECK: scf.while
+
+// Producer if=5: skip cross memref.load/slt, then intra control_var < 1.
+// CHECK: memref.load %{{.*}}[] : memref<i32, #hivm.address_space<ssbuf>>
+// CHECK: arith.cmpi slt, %{{.*}}, %{{.*}} : i32
+// CHECK: %[[LIM:.*]] = arith.constant 1 : i32
+// CHECK: %[[SLT:.*]] = arith.cmpi slt, %{{.*}}, %[[LIM]] : i32
+// CHECK: %[[AND_P0:.*]] = arith.andi %{{.*}}, %[[SLT]]
+// CHECK: %[[COND5:.*]] = arith.andi %[[AND_P0]], %{{.*}}
+// CHECK: scf.if %[[COND5]]
+// CHECK: %[[ONE_P:.*]] = arith.constant 1 : i32
+// CHECK: %[[INC:.*]] = arith.addi %{{.*}}, %[[ONE_P]] : i32
+// CHECK: scf.yield %{{.*}}, %[[INC]]
+// CHECK: } {{.*}}ssbuffer.if = 5
+
+// Consumer if=7: uses producer if result (latest), sgt 0 then -1.
+// CHECK: %[[SGT:.*]] = arith.cmpi sgt, %{{.*}}, %{{.*}} : i32
+// CHECK: %[[COND7:.*]] = arith.andi %[[SGT]], %{{.*}}
+// CHECK: scf.if %[[COND7]]
+// CHECK: %[[ONE_C:.*]] = arith.constant 1 : i32
+// CHECK: %[[DEC:.*]] = arith.subi %{{.*}}, %[[ONE_C]] : i32
+// CHECK: } {{.*}}ssbuffer.if = 7
+
+// CHECK: scf.yield
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">, ssbuffer.intra_buf_count = 2 : i32} {
+  func.func @while_update_condition_intracore(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %c0_i32 = arith.constant 0 : i32
+    %c16_i32 = arith.constant 16 : i32
+    %bound = arith.extsi %arg7 : i32 to i64
+
+    scope.scope : () -> () {
+      %alloc_ub0 = memref.alloc() : memref<128xf32, #hivm.address_space<ub>>
+      %mem0 = memref.memory_space_cast %alloc_ub0 : memref<128xf32, #hivm.address_space<ub>> to memref<128xf32>
+      %alloc_ub1 = memref.alloc() : memref<128xf32, #hivm.address_space<ub>>
+      %mem1 = memref.memory_space_cast %alloc_ub1 : memref<128xf32, #hivm.address_space<ub>> to memref<128xf32>
+      %alloc_cbuf = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<8x4x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_cbuf {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<8x4x16x16xf16, #hivm.address_space<cbuf>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        // Producer block (intraDeps role = 1).
+        %t_prod = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<128xf32>
+        hivm.hir.copy ins(%t_prod : tensor<128xf32>) outs(%mem0 : memref<128xf32>) {ssbuffer.block_id = 5 : i32, ssbuffer.intraDeps = [0 : i32, 1 : i32]}
+        %cbuf_t = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x4x16x16xf16>
+        hivm.hir.copy ins(%cbuf_t : tensor<8x4x16x16xf16>) outs(%alloc_cbuf : memref<8x4x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [1 : i32, 1 : i32]}
+        %next5 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 5 : i32} : i32
+
+        // Consumer block (intraDeps role = 0).
+        %t_cons = bufferization.to_tensor %mem0 restrict writable {ssbuffer.block_id = 7 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32> to tensor<128xf32>
+        %next7 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 7 : i32} : i32
+        scf.yield %next7 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+
+    scope.scope : () -> () {
+      %alloc_cbuf = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<8x4x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_cbuf {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<8x4x16x16xf16, #hivm.address_space<cbuf>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 2 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %conv = hivm.hir.convert_layout %alloc_cbuf output_shape [64, 128] {dstLayout = #hivm.data_layout<ND>, srcLayout = #hivm.data_layout<nZ>, ssbuffer.block_id = 2 : i32, ssbuffer.crossCoreDeps = [1 : i32, 0 : i32], ssbuffer.transfer_id = 1 : i32} : (memref<8x4x16x16xf16, #hivm.address_space<cbuf>>) -> memref<64x128xf16, #hivm.address_space<cbuf>>
+        %next2 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 2 : i32} : i32
+        scf.yield %next2 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+    return
+  }
+}
diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-tensor-iterarg.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-tensor-iterarg.mlir
new file mode 100644
index 0000000000..348e230dea
--- /dev/null
+++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/while-update-condition-tensor-iterarg.mlir
@@ -0,0 +1,106 @@
+// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s
+
+// UpdateConditionInfo on scf.while: tensor-typed iter_arg control vars.
+// while carries (tensor, iv); block 5 consumes the tensor, block 6 produces a
+// new tensor for yield. UpdateLoopOps appends i32=1 control args; conditions:
+//   consumer: control_var == 1, then -1
+//   producer: control_var == 0, then +1
+// (same semantics as for-path iter_args_deps_add_conditions)
+
+// CHECK-LABEL: func.func @while_update_condition_tensor_iterarg
+
+// Tensor-iter control arg is seeded with 1 and appears on the while.
+// CHECK: %[[TINIT:.*]] = arith.constant 1 : i32
+// CHECK: scf.while
+// CHECK-SAME: %[[TINIT]]
+
+// Consumer if=5: after cross sgt, eq 1 is and-ed into cond; then -1 inside if.
+// IR order: c1 (cross) ; load ; sgt ; c1 (tensor) ; eq ; andi ; andi ; scf.if
+// Bind ONE only after sgt so it is the tensor eq constant, not the cross-core one.
+// CHECK: arith.cmpi sgt
+// CHECK: %[[ONE:.*]] = arith.constant 1 : i32
+// CHECK: arith.cmpi eq, %{{.*}}, %[[ONE]] : i32
+// CHECK: arith.andi
+// CHECK: %[[COND5:.*]] = arith.andi
+// CHECK: scf.if %[[COND5]]
+// CHECK: arith.subi
+// CHECK: } {{.*}}ssbuffer.if = 5
+
+// Producer if=6: after cross slt, eq 0 is and-ed into cond; then +1 inside if.
+// CHECK: arith.cmpi slt
+// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32
+// CHECK: arith.cmpi eq, %{{.*}}, %[[ZERO]] : i32
+// CHECK: arith.andi
+// CHECK: %[[COND6:.*]] = arith.andi
+// CHECK: scf.if %[[COND6]]
+// CHECK: arith.addi
+// CHECK: } {{.*}}ssbuffer.if = 6
+
+// CHECK: scf.yield
+
+module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} {
+  func.func @while_update_condition_tensor_iterarg(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xf16> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32 {tt.divisibility = 16 : i32}, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} {
+    %c0_i32 = arith.constant 0 : i32
+    %c16_i32 = arith.constant 16 : i32
+    %cst = arith.constant 0.000000e+00 : f32
+    %bound = arith.extsi %arg7 : i32 to i64
+
+    scope.scope : () -> () {
+      %empty0 = tensor.empty() : tensor<16x16xf32>
+      %init_t = linalg.fill ins(%cst : f32) outs(%empty0 : tensor<16x16xf32>) -> tensor<16x16xf32>
+      %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0:2 = scf.while (%arg14 = %init_t, %arg15 = %c0_i32) : (tensor<16x16xf32>, i32) -> (tensor<16x16xf32>, i32) {
+        %1 = arith.extsi %arg15 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14, %arg15 : tensor<16x16xf32>, i32
+      } do {
+      ^bb0(%arg14: tensor<16x16xf32>, %arg15: i32):
+        // Consumer of tensor iter_arg (non-yield use) + cross-core input.
+        hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %sum = arith.addf %arg14, %arg14 {ssbuffer.block_id = 5 : i32} : tensor<16x16xf32>
+        %memspacecast = memref.memory_space_cast %alloc_ub {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space<ub>> to memref<128x128xf32>
+
+        // Producer of tensor iter_arg (loop yield comes from this block) + cross-core output.
+        // CreateIfOps else-yields %arg14 → analyzeTensorIterArgDependencies marks producer.
+        hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 3 : i32}[<VECTOR>, <PIPE_V>, <PIPE_FIX>] flag = 4
+        %empty1 = tensor.empty() {ssbuffer.block_id = 6 : i32} : tensor<16x16xf32>
+        %new_t = linalg.fill {ssbuffer.block_id = 6 : i32} ins(%cst : f32) outs(%empty1 : tensor<16x16xf32>) -> tensor<16x16xf32>
+        %cbuf_t = tensor.empty() {ssbuffer.block_id = 6 : i32} : tensor<8x8x16x16xf16>
+        hivm.hir.copy ins(%cbuf_t : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [0 : i32, 1 : i32]}
+        %next6 = arith.addi %arg15, %c16_i32 {ssbuffer.block_id = 6 : i32} : i32
+        scf.yield %new_t, %next6 : tensor<16x16xf32>, i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<VECTOR>}
+
+    scope.scope : () -> () {
+      %alloc_c = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      annotation.mark %alloc_c {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space<cbuf>>
+      %alloc_ub = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+      annotation.mark %alloc_ub {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space<ub>>
+
+      %0 = scf.while (%arg14 = %c0_i32) : (i32) -> i32 {
+        %1 = arith.extsi %arg14 {Undefined, ssbuffer.block_id = 4 : i32} : i32 to i64
+        %2 = arith.cmpi slt, %1, %bound {Undefined, ssbuffer.block_id = 4 : i32} : i64
+        scf.condition(%2) %arg14 : i32
+      } do {
+      ^bb0(%arg14: i32):
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %t0 = tensor.empty() {ssbuffer.block_id = 0 : i32} : tensor<128x128xf32>
+        hivm.hir.fixpipe {dma_mode = #hivm.dma_mode<nz2nd>, ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossCoreDeps = [0 : i32, 1 : i32]} ins(%t0 : tensor<128x128xf32>) outs(%alloc_ub : memref<128x128xf32, #hivm.address_space<ub>>)
+        %next0 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 0 : i32} : i32
+
+        hivm.hir.sync_block_wait {ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 1 : i32}[<CUBE>, <PIPE_MTE3>, <PIPE_MTE1>] flag = 2
+        %conv = hivm.hir.convert_layout %alloc_c output_shape [128, 128] {dstLayout = #hivm.data_layout<ND>, srcLayout = #hivm.data_layout<nZ>, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossCoreDeps = [0 : i32, 0 : i32]} : (memref<8x8x16x16xf16, #hivm.address_space<cbuf>>) -> memref<128x128xf16, #hivm.address_space<cbuf>>
+        %next1 = arith.addi %arg14, %c16_i32 {ssbuffer.block_id = 1 : i32} : i32
+        scf.yield %next1 : i32
+      } attributes {Undefined, ssbuffer.main_loop = 0 : i32}
+      scope.return
+    } {hivm.tcore_type = #hivm.tcore_type<CUBE>}
+    return
+  }
+}

From 75b808ae3da2d3c9e20994624db624ac13aed48d Mon Sep 17 00:00:00 2001
From: fishofnanqi <1074959344@qq.com>
Date: Wed, 29 Jul 2026 10:07:07 +0800
Subject: [PATCH 10/11] [ssbuffer](feat) adapter whileOp in UpdateLoopIterTimes

Signed-off-by: fishofnanqi <1074959344@qq.com>
---
 .../UpdateLoopIterTimes.h                     |   8 +
 .../UpdateLoopIterTimes.cpp                   | 229 ++++++++++++++++--
 2 files changed, 221 insertions(+), 16 deletions(-)

diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h
index c95a0c6643..666702793a 100644
--- a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h
+++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h
@@ -73,6 +73,9 @@ class UpdateLoopIterTimesPass
 
   int replaceForOpCounterInIfOps();
 
+  int UpdateWhileLoopCondition(
+      DenseMap<int, SmallVector<Operation *>> &mainLoopIdMap);
+
   // Calculate factor = requiredBuffers / x
   std::pair<int, int> calculateFactor(scf::ForOp forOp);
 
@@ -125,6 +128,11 @@ class UpdateLoopIterTimesPass
 
   int updateCntArgsAfterClone(scf::ForOp oldForOp, IRMapping &mapper,
                               SmallVector<scf::IfOp> &ifOpsInThisFor);
+
+  void updateMainLoopMaps(Operation *oldForOp, Operation *newForOp,
+                          DenseMap<int, SmallVector<Operation *>> &cmap,
+                          DenseMap<int, SmallVector<Operation *>> &vmap,
+                          DenseMap<Operation *, IterationTimesInfo> &infoMap);
 };
 
 std::unique_ptr<OperationPass<ModuleOp>> createUpdateLoopIterTimesPass();
diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
index dcea9545f6..4d3aade93b 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp
@@ -27,6 +27,7 @@
 #include "bishengir/Dialect/Annotation/IR/Annotation.h"
 #include "bishengir/Dialect/HIVM/IR/HIVM.h"
 #include "bishengir/Dialect/Scope/IR/Scope.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "llvm/ADT/DenseMap.h"
@@ -101,11 +102,21 @@ static scf::ForOp getOtherScopeMainloop(ModuleOp module, bool currentIsCube,
     // id
     scopeOp.walk([&](Operation *op) {
       if (op->hasAttr(CVPipeline::kMainLoop)) {
+        // Support both ForOp and WhileOp for mainloop
+        bool isValidLoop = isa<scf::ForOp>(op) || isa<scf::WhileOp>(op);
+        if (!isValidLoop) {
+          LDBG("mainloop must be ForOp or WhileOp!");
+          return WalkResult::advance();
+        }
+
+        // Only cast to ForOp for compatibility with current interface
+        // WhileOp will be skipped in later steps
         auto targetForOp = dyn_cast<scf::ForOp>(op);
         if (!targetForOp) {
-          LDBG("do not support other mainloop op except ForOp");
+          // WhileOp case: skip for now, will be handled separately
           return WalkResult::advance();
         }
+
         auto targetMainLoopId =
             targetForOp->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
         if (targetMainLoopId && targetMainLoopId.getInt() == mainLoopId) {
@@ -836,7 +847,137 @@ scf::ForOp UpdateLoopIterTimesPass::extendForOpIterationCount(
   return newForOp;
 }
 
-// step4: Replace loop counter by if blocks' counter
+// Update WhileOp condition based on ifblock conditions
+// Uses info->whileBlockArgMap: {WhileOp: {block_id: {new_arg_idx:
+// old_arg_idx}}} For each IfOp in WhileOp, get its block_id Get new args and
+// mapping from info->whileBlockArgMap Copy entire beforeRegion and replace old
+// args with new args Combine all ifblock conditions with OR operation
+int UpdateLoopIterTimesPass::UpdateWhileLoopCondition(
+    DenseMap<int, SmallVector<Operation *>> &mainLoopIdMap) {
+  int ret = 0;
+
+  for (auto &entry : mainLoopIdMap) {
+    for (Operation *loopOp : entry.second) {
+      if (!isa<scf::WhileOp>(loopOp)) {
+        continue;
+      }
+
+      auto whileOp = dyn_cast<scf::WhileOp>(loopOp);
+      if (!info->whileBlockArgMap.count(whileOp)) {
+        LDBG("WhileOp not found in whileBlockArgMap!");
+        return -1;
+      }
+      auto &blockArgMap = info->whileBlockArgMap[whileOp];
+      if (blockArgMap.empty()) {
+        LDBG("blockArgMap is empty for WhileOp!");
+        return -1;
+      }
+
+      // Get the 'before' region which contains the condition
+      Region &beforeRegion = whileOp.getBefore();
+      Block &beforeBlock = beforeRegion.front();
+      Operation *terminator = beforeBlock.getTerminator();
+      if (!terminator || !isa<scf::ConditionOp>(terminator)) {
+        LDBG("Before block has no valid ConditionOp terminator!");
+        return -1;
+      }
+      auto conditionOp = dyn_cast<scf::ConditionOp>(terminator);
+      Value originalCondition = conditionOp.getCondition();
+      OpBuilder builder(conditionOp);
+      Location loc = whileOp.getLoc();
+
+      // Collect all operations to clone
+      SmallVector<Operation *> opsToClone;
+      for (Operation &op : beforeBlock.without_terminator()) {
+        opsToClone.push_back(&op);
+      }
+
+      // Traverse all IfOps with ssbuffer.if and
+      // build condition for each IfOp directly
+      Value combinedCondition;
+      whileOp.walk([&](scf::IfOp ifOp) {
+        if (!ifOp->hasAttr(CVPipeline::kIf)) {
+          return WalkResult::advance();
+        }
+
+        auto blockIdAttr = ifOp->getAttrOfType<IntegerAttr>(CVPipeline::kIf);
+        if (!blockIdAttr) {
+          ret = -1;
+          return WalkResult::interrupt();
+        }
+        int blockId = blockIdAttr.getInt();
+        if (!blockArgMap.count(blockId)) {
+          ret = -1;
+          LDBG("blockId " << blockId << " not found in blockArgMap!");
+          return WalkResult::interrupt();
+        }
+
+        auto &argMap = blockArgMap[blockId]; // {new_arg_idx: old_arg_idx}
+
+        // Step 1: Build arg replacement mapping
+        IRMapping mapper;
+        for (auto &argEntry : argMap) {
+          int newArgIdx = argEntry.first;
+          int oldArgIdx = argEntry.second;
+
+          BlockArgument oldArg = beforeBlock.getArgument(oldArgIdx);
+          BlockArgument newArg = beforeBlock.getArgument(newArgIdx);
+
+          mapper.map(oldArg, newArg);
+        }
+
+        // Step 2: Clone entire beforeRegion
+        // Set insertion point before conditionOp
+        builder.setInsertionPoint(conditionOp);
+
+        // Clone all operations
+        Value newCondition;
+        for (Operation *op : opsToClone) {
+          Operation *clonedOp = builder.clone(*op, mapper);
+
+          // If this op defines the original condition, get the corresponding
+          // result
+          if (op->getResultTypes().size() > 0) {
+            for (unsigned i = 0; i < op->getNumResults(); ++i) {
+              if (op->getResult(i) == originalCondition) {
+                newCondition = clonedOp->getResult(i);
+                break;
+              }
+            }
+          }
+        }
+
+        if (!newCondition) {
+          ret = -1;
+          LDBG("Failed to create new condition for blockId: " << blockId);
+          return WalkResult::interrupt();
+        }
+
+        // Step 3: Combine with existing conditions using OR
+        if (!combinedCondition) {
+          combinedCondition = newCondition;
+        } else {
+          combinedCondition = builder.create<arith::OrIOp>(
+              loc, combinedCondition, newCondition);
+        }
+
+        return WalkResult::advance();
+      });
+
+      // Step 4: Update the condition in the before block
+      if (combinedCondition) {
+        conditionOp.getConditionMutable().assign(combinedCondition);
+        LDBG("Updated WhileOp condition for: " << whileOp);
+      } else {
+        LDBG("Failed to generate valid condition!");
+        return -1;
+      }
+    }
+  }
+
+  return ret;
+}
+
 // Traverse each mainloop, find ifOp with ssbuffer.if attribute inside,
 // and replace the mainloop's induction variable with the counter in cntArgs
 int UpdateLoopIterTimesPass::replaceForOpCounterInIfOps() {
@@ -844,6 +985,12 @@ int UpdateLoopIterTimesPass::replaceForOpCounterInIfOps() {
   // Traverse all mainloops in the module
   getOperation().walk([&](Operation *op) {
     if (op->hasAttr(CVPipeline::kMainLoop)) {
+      // Skip WhileOp, only process ForOp in this step
+      if (isa<scf::WhileOp>(op)) {
+        LDBG("Skip WhileOp in replaceForOpCounterInIfOps.");
+        return WalkResult::advance();
+      }
+
       auto forOp = dyn_cast<scf::ForOp>(op);
       if (!forOp) {
         ret = -1;
@@ -903,21 +1050,21 @@ int UpdateLoopIterTimesPass::GetMainLoopIdToLoopOpMap(
     // Walk for loops inside the scope
     scopeOp.walk([&](Operation *op) {
       if (op->hasAttr(CVPipeline::kMainLoop)) {
-        auto forOp = dyn_cast<scf::ForOp>(op);
-        if (!forOp) {
+        // Support both ForOp and WhileOp for mainloop
+        bool isValidLoop = isa<scf::ForOp>(op) || isa<scf::WhileOp>(op);
+        if (!isValidLoop) {
           ret = -1;
-          LDBG("do not surpport other loop op temprarily!");
+          LDBG("mainloop must be ForOp or WhileOp!");
           return mlir::WalkResult::interrupt();
         }
 
-        auto mainLoopId =
-            forOp->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
+        auto mainLoopId = op->getAttrOfType<IntegerAttr>(CVPipeline::kMainLoop);
         if (mainLoopId) {
           int id = mainLoopId.getInt();
           if (isCube) {
-            cmap[id].push_back(forOp.getOperation());
+            cmap[id].push_back(op);
           } else if (isVector) {
-            vmap[id].push_back(forOp.getOperation());
+            vmap[id].push_back(op);
           }
         }
       }
@@ -936,6 +1083,10 @@ int UpdateLoopIterTimesPass::ComputeMainLoopTimes(
     DenseMap<Operation *, IterationTimesInfo> &infoMap) {
   for (auto &entry : loopMap) {
     for (Operation *loopOp : entry.second) {
+      if (isa<scf::WhileOp>(loopOp)) {
+        continue;
+      }
+
       scf::ForOp forOp = dyn_cast<scf::ForOp>(loopOp);
       if (!forOp) {
         LDBG("currently only support forOp!");
@@ -992,6 +1143,10 @@ int UpdateLoopIterTimesPass::collectForOpsAndUpdateMax(
     DenseMap<Operation *, IterationTimesInfo> &infoMap) {
   if (map.count(id)) {
     for (Operation *loopOp : map[id]) {
+      if (!isa<scf::ForOp>(loopOp)) {
+        continue;
+      }
+
       allForOps.push_back(loopOp);
       if (infoMap.count(loopOp)) {
         IterationTimesInfo &iterInfo = infoMap[loopOp];
@@ -1013,6 +1168,36 @@ int UpdateLoopIterTimesPass::collectForOpsAndUpdateMax(
   return 0;
 }
 
+void UpdateLoopIterTimesPass::updateMainLoopMaps(
+    Operation *oldForOp, Operation *newForOp,
+    DenseMap<int, SmallVector<Operation *>> &cmap,
+    DenseMap<int, SmallVector<Operation *>> &vmap,
+    DenseMap<Operation *, IterationTimesInfo> &infoMap) {
+  // Update cmap
+  for (auto &entry : cmap) {
+    for (Operation *&op : entry.second) {
+      if (op == oldForOp) {
+        op = newForOp;
+      }
+    }
+  }
+  // Update vmap
+  for (auto &entry : vmap) {
+    for (Operation *&op : entry.second) {
+      if (op == oldForOp) {
+        op = newForOp;
+      }
+    }
+  }
+  // Update infoMap
+  auto it = infoMap.find(oldForOp);
+  if (it != infoMap.end()) {
+    IterationTimesInfo iterInfo = std::move(it->second);
+    infoMap.erase(it);
+    infoMap[newForOp] = std::move(iterInfo);
+  }
+}
+
 int UpdateLoopIterTimesPass::UpdateForLoopIteration(
     DenseMap<int, SmallVector<Operation *>> &cmap,
     DenseMap<int, SmallVector<Operation *>> &vmap,
@@ -1045,13 +1230,13 @@ int UpdateLoopIterTimesPass::UpdateForLoopIteration(
     if (ret != 0)
       return -1;
 
-    if (maxIfCount == 0) {
-      LDBG("no ifblock in mainloop!");
-      return -1;
-    }
-
     // Update all loops with this id using the same max values
     for (Operation *loopOp : sameIdForOps) {
+      if (maxIfCount == 0) {
+        LDBG("no ifblock in mainloop!");
+        return -1;
+      }
+
       scf::ForOp oldForOp = dyn_cast<scf::ForOp>(loopOp);
       if (!oldForOp) {
         LDBG("do not surpport other loop op except forOp!");
@@ -1067,12 +1252,12 @@ int UpdateLoopIterTimesPass::UpdateForLoopIteration(
         LDBG("extendForOpIterationCount failed!");
         return -1;
       }
+      updateMainLoopMaps(loopOp, newForOp.getOperation(), cmap, vmap, infoMap);
       allForOps.push_back(loopOp);
     }
   }
 
-  // Delete old forOp after cntArgs has been updated in
-  // extendForOpIterationCount
+  // Delete old forOp after cntArgs and maps have been updated
   for (Operation *loopOp : allForOps) {
     if (!loopOp) {
       LDBG("erasing error: loopOp is nullptr, there are nested mainloop!");
@@ -1134,6 +1319,18 @@ void UpdateLoopIterTimesPass::runOnOperation() {
     CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
   }
 
+  // step5: Update WhileOp condition based on ifblock conditions
+  ret = UpdateWhileLoopCondition(cmap);
+  if (ret != 0) {
+    LDBG("UpdateWhileLoopCondition from cube Failed!");
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+  }
+  ret = UpdateWhileLoopCondition(vmap);
+  if (ret != 0) {
+    LDBG("UpdateWhileLoopCondition from vector Failed!");
+    CVPipeline::setFallbackAttr(module, CVPipeline::ERRCODE_FAILED);
+  }
+
   LDBG("after updateloopitertimes:\n" << module);
   LDBG("\nExit UpdateLoopIterTimes pass.");
 }

From f90f9fbd9a4ebcc1e6488706cea08360f9c14e07 Mon Sep 17 00:00:00 2001
From: sxm <sunxianmiao@h-partners.com>
Date: Wed, 5 Aug 2026 21:39:33 +0800
Subject: [PATCH 11/11] [ssbuffer](fix) fix updateloopops bug

---
 .../AddControlFlowCondition/UpdateLoopOps.cpp | 38 +++++++++----------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
index 838d4339a2..0f91812028 100644
--- a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
+++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopOps.cpp
@@ -230,18 +230,22 @@ static unsigned getMainLoopBaseIdx(Operation *oldLoopOp, bool isWhile) {
                  : cast<scf::ForOp>(oldLoopOp).getNumRegionIterArgs();
 }
 
-// Appends initial values for extra iter_args: block counters from
-// blockCounterInitFn (forOp reuses getLowerBound(); whileOp creates a new
-// arith.constant per counter so each new iter_arg has a distinct SSA value),
-// dep conds from i32(0), tensor iter_args from i32(1).
-static void
-buildMainLoopExtraInitArgs(OpBuilder &builder, Location loc,
-                           llvm::function_ref<Value()> blockCounterInitFn,
-                           int numBlockCounters, int numInnerDepConds,
-                           int numTensorIterArgs,
-                           llvm::SmallVector<Value> &extraInitArgs) {
+// Appends initial values for extra iter_args:
+// - block counters: forOp reuses lowerBound; whileOp creates a fresh i32(0)
+//   per counter (distinct SSA values)
+// - inner dep conds: i32(0)
+// - tensor iter_args: i32(1)
+static void buildMainLoopExtraInitArgs(
+    OpBuilder &builder, Location loc, Value forOpLowerBound, bool isWhile,
+    int numBlockCounters, int numInnerDepConds, int numTensorIterArgs,
+    llvm::SmallVector<Value> &extraInitArgs) {
   for (int i = 0; i < numBlockCounters; ++i) {
-    extraInitArgs.push_back(blockCounterInitFn());
+    if (isWhile) {
+      extraInitArgs.push_back(builder.create<arith::ConstantOp>(
+          loc, builder.getI32Type(), builder.getI32IntegerAttr(0)));
+    } else {
+      extraInitArgs.push_back(forOpLowerBound);
+    }
   }
   for (int i = 0; i < numInnerDepConds; ++i) {
     extraInitArgs.push_back(builder.create<arith::ConstantOp>(
@@ -342,16 +346,10 @@ extendMainLoopOpWithExtraArgs(Operation *oldLoopOp,
   // passes).
   OpBuilder builder(oldLoopOp);
   Value forOpLowerBound;
-  llvm::function_ref<Value()> blockCounterInitFn;
   bool isWhile = false;
   if (auto forOp = dyn_cast<scf::ForOp>(oldLoopOp)) {
     forOpLowerBound = forOp.getLowerBound();
-    blockCounterInitFn = [&]() { return forOpLowerBound; };
-  } else if (auto whileOp = dyn_cast<scf::WhileOp>(oldLoopOp)) {
-    blockCounterInitFn = [&]() {
-      return builder.create<arith::ConstantOp>(
-          whileOp.getLoc(), builder.getI32Type(), builder.getI32IntegerAttr(0));
-    };
+  } else if (isa<scf::WhileOp>(oldLoopOp)) {
     isWhile = true;
   } else {
     LDBG("[Error]: main_loop op is neither scf::ForOp nor scf::WhileOp");
@@ -359,8 +357,8 @@ extendMainLoopOpWithExtraArgs(Operation *oldLoopOp,
   }
 
   llvm::SmallVector<Value> extraInitArgs;
-  buildMainLoopExtraInitArgs(builder, oldLoopOp->getLoc(), blockCounterInitFn,
-                             numBlockCounters, numInnerDepConds,
+  buildMainLoopExtraInitArgs(builder, oldLoopOp->getLoc(), forOpLowerBound,
+                             isWhile, numBlockCounters, numInnerDepConds,
                              numTensorIterArgs, extraInitArgs);
 
   Operation *newOp = createMainLoopOpAndMigrateBody(oldLoopOp, extraInitArgs);