triton-lang · zhuxinguang33 · Jul 20, 2026 · Aug 3, 2026 · Jul 29, 2026 · Jul 31, 2026
@@ -37,43 +37,48 @@ 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;
   // The variables that need to be controlled as a consumer
   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
+  // unique counter value for each ifblock scf.for only.
   llvm::DenseMap<scf::IfOp, Value> cntArgs;
 
   // DAG for if block cross-core dependencies
@@ -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

@@ -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"; }

@@ -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);

@@ -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();

@@ -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

@@ -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();

@@ -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,17 +51,17 @@ 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);
 
   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