[Utility] fix pass-by-reference addIterArgsToLoop API (#7029)

This API with the `scf::ForOp&` passed by reference is a foot-gun
so this PR refactors to return the new `ForOp` instead.

Author: makslevental

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -141,8 +141,8 @@ scf::ForOp replaceForOpWithNewSignature(
     SmallVectorImpl<std::tuple<Value, Value>> &replacements);
 scf::ForOp replaceForOpWithNewSignature(OpBuilder &rewriter, scf::ForOp loop,
                                         ValueRange newIterOperands);
-Block::BlockArgListType addIterArgsToLoop(OpBuilder &rewriter, scf::ForOp &loop,
-                                          ValueRange newIterOperands);
+[[nodiscard]] scf::ForOp addIterArgsToLoop(OpBuilder &rewriter, scf::ForOp loop,
+                                           ValueRange newIterOperands);
 
 // Replace WhileOp with a new WhileOp with extra operands. The YieldOp is not
 // updated and needs to be updated separately for the loop to be correct.

lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp

@@ -337,7 +337,8 @@ ttng::TMEMAllocOp hoistTMEMAlloc(TMEMTokenAllocOp alloc, scf::ForOp &forOp) {
   // By hoisting the allocation out of the loop, we need to turn the underlying
   // memory variable into a loop-carried depdendency.
   auto tokType = builder.getType<AsyncTokenType>();
-  Value newTok = addIterArgsToLoop(builder, forOp, newAlloc.getToken()).front();
+  forOp = addIterArgsToLoop(builder, forOp, newAlloc.getToken());
+  Value newTok = forOp.getRegionIterArgs().back();
   appendToForOpYield(forOp, joinLastMemoryUses(builder, alloc.getToken()));
 
   if (src != nullptr) {

lib/Dialect/TritonGPU/Transforms/OptimizeAccumulatorInit.cpp

@@ -249,7 +249,7 @@ class OptimizeAccumulatorInitPass
       }
 
       Value loopArgFlagValue = loopArgIsZero ? vFalse : vTrue;
-      (void)addIterArgsToLoop(rewriter, forOp, {loopArgFlagValue});
+      forOp = addIterArgsToLoop(rewriter, forOp, {loopArgFlagValue});
       loopArgFlagValue =
           forOp.getRegionIterArg(forOp.getNumRegionIterArgs() - 1);
 

lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp

@@ -556,7 +556,7 @@ scf::ForOp lowerLoads(scf::ForOp forOp, CoarseSchedule &schedule) {
   }
 
   // Patch the loop to add the new loop carried dependencies.
-  (void)addIterArgsToLoop(builder, forOp, newOperands);
+  forOp = addIterArgsToLoop(builder, forOp, newOperands);
 
   // Update yield op with temporary yield values
   auto forYield = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
@@ -750,7 +750,7 @@ scf::ForOp lowerTMADescriptors(scf::ForOp forOp, CoarseSchedule &schedule) {
     newOperands.push_back(zero);
   }
 
-  (void)addIterArgsToLoop(builder, forOp, newOperands);
+  forOp = addIterArgsToLoop(builder, forOp, newOperands);
 
   auto tmaCounters = ArrayRef<BlockArgument>(forOp.getBody()->getArguments())
                          .slice(tmaCounterArgsStartIdx);

lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp

@@ -106,7 +106,7 @@ Value triton::sinkValueRedefinition(RewriterBase &rewriter, Value in, Value out,
     // `in` is live into the loop body. `out` becomes the live-out if the
     // loop executes at least once.
     if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-      (void)addIterArgsToLoop(rewriter, forOp, in);
+      forOp = addIterArgsToLoop(rewriter, forOp, in);
       appendToForOpYield(forOp, out);
       out = forOp.getResults().back();
       continue;

lib/Dialect/TritonGPU/Transforms/Utility.cpp

@@ -682,17 +682,15 @@ scf::ForOp replaceForOpWithNewSignature(OpBuilder &rewriter, scf::ForOp loop,
   return newForOp;
 }
 
-Block::BlockArgListType addIterArgsToLoop(OpBuilder &rewriter, scf::ForOp &loop,
-                                          ValueRange newIterOperands) {
-  unsigned curArgIdx = loop.getNumRegionIterArgs();
+scf::ForOp addIterArgsToLoop(OpBuilder &rewriter, scf::ForOp loop,
+                             ValueRange newIterOperands) {
   scf::ForOp newLoop =
       replaceForOpWithNewSignature(rewriter, loop, newIterOperands);
   // Save the caller from insertion point invalidation.
   if (rewriter.getInsertionPoint() == loop->getIterator())
     rewriter.setInsertionPoint(newLoop);
   loop.erase();
-  loop = newLoop;
-  return loop.getRegionIterArgs().slice(curArgIdx);
+  return newLoop;
 }
 
 scf::WhileOp replaceWhileOpWithNewSignature(

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/LoadMMASpecialization.cpp

@@ -118,8 +118,8 @@ addIndexAndPhase(PartitionBuilder &b, scf::ForOp &loop, unsigned numStages,
   b.setInsertionPoint(loop);
 
   // Index and phase both start at 0.
-  unsigned curArgIdx = loop.getNumRegionIterArgs();
-  auto newArgs = addIterArgsToLoop(b, loop, {b.intCst(0), b.intCst(0)});
+  loop = addIterArgsToLoop(b, loop, {b.intCst(0), b.intCst(0)});
+  auto newArgs = loop.getRegionIterArgs().take_back(2);
   BlockArgument index = newArgs[0];
   BlockArgument phase = newArgs[1];
 
@@ -488,7 +488,8 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
       createTMemAlloc(b, oldAllocOp, /*multiBuffered=*/true, numMmaStages);
 
   // Use placeholder values for the indices in the loop.
-  auto indexPhase = addIterArgsToLoop(b, loop, {b.intCst(0), b.intCst(0)});
+  loop = addIterArgsToLoop(b, loop, {b.intCst(0), b.intCst(0)});
+  auto indexPhase = loop.getRegionIterArgs().take_back(2);
   BlockArgument index = indexPhase[0];
   BlockArgument phase = indexPhase[1];
 

third_party/amd/lib/TritonAMDGPUTransforms/StreamPipeline.cpp

@@ -893,7 +893,7 @@ void StreamPipeliner::createStreamOps() {
 
   unsigned newOperandIndex = forOp.getBody()->getNumArguments();
   // Patch the loop to add the new loop carried dependencies.
-  (void)addIterArgsToLoop(builder, forOp, {extractIdx});
+  forOp = addIterArgsToLoop(builder, forOp, {extractIdx});
 
   // Create one counter for the extract indices to avoid creating long
   // live range.