[mlir][scf] Fix a bug in scf::ForOp loop unroll with an epilogue

[tatianashp]

Summary

Fixes a bug in formation and simplification of an epilogue loop generated
during loop unroll of scf::ForOp (https://bugs.llvm.org/show_bug.cgi?id=46689)

[mehdi_amini]

mlir/test/Transforms/scf-loop-unroll.mlir

2

Can you remove the -allow-unregistered-dialect here and use std.add (or another test operation), and same for types?

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

217

Nit: consider using getTerminator() instead of back() for better readability.

mlir/test/Transforms/scf-loop-unroll.mlir

9

Could you make it 11 iterations to make the test stronger (or an additional one if you prefer)? Is coverage missing for a symbolic upper bound in the presence of return values?

[mehdi_amini]

mlir/test/Transforms/scf-loop-unroll.mlir

9

I believe the crash from PR46689 only exists when the epilogue has a single iteration

[tatianashp]

Update: Fix handling of loop carried variables inside the loop.
This update addresses reviewer comments and fixes another, previously undetected bug:

• Loop carried variables are now correctly propagated during unrolling the body of the loop.
• Added another test with symbolic upper bound and multiple loop-carried variables.

[tatianashp]

mlir/test/Transforms/scf-loop-unroll.mlir

9

I added another test with a symbolic upper bound and multiple carried values to test the case when epilogue loop is not unrolled. It does not trigger crash from PR46689, but good to have as additional coverage.

[mehdi_amini]

Seems like you uploaded the same diff as before, you may want to double-check and retry here?

[tatianashp]

Uploaded changes.

[tatianashp]

Remove changes pulled in from unrelated updates to mlir/lib/Transforms/Utils/LoopUtils.cpp

[mehdi_amini]

LGTM again, thanks. You probably should wait one or two days to give a change to Uday to check again.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1052–1053

Nit: These can be made triple /// as well while at this.

1089

Nit: please use i = 0, e = ...; i < e form.

1156

Nit: /*iterArgs=*/{}, /*yieldedValues=*/{};

[tatianashp]

A few minor changes addressing Uday's comments.

[tatianashp]

Mehdi, Uday - would you like to commit this diff? I don't have commit priveledges.

[mehdi_amini]

In D87583#2319842, @tatianashp wrote:

Mehdi, Uday - would you like to commit this diff? I don't have commit priveledges.

Sure! But it needs to be rebased right now I believe.

[tatianashp]

Rebased

mlir/lib/Transforms/Utils/LoopUtils.cpp

if (!lbCstOp || !ubCstOp || !stepCstOp || lbCstOp.getValue() < 0 ||
     return failure();
   int64_t tripCount = mlir::ceilDiv(ubCstOp.getValue() - lbCstOp.getValue(),
                                     stepCstOp.getValue());
   if (tripCount != 1)
     return failure();
   auto iv = forOp.getInductionVar();
   iv.replaceAllUsesWith(lbCstOp);
 
+  // Replace uses of iterArgs with iterOperands.
+  auto iterOperands = forOp.getIterOperands();
+  auto iterArgs = forOp.getRegionIterArgs();
+  for (auto e : llvm::zip(iterOperands, iterArgs))
+    std::get<1>(e).replaceAllUsesWith(std::get<0>(e));
+
+  // Replace uses of loop results with the values yielded by the loop.
+  auto outerResults = forOp.getResults();
+  auto innerResults = forOp.getBody()->getTerminator()->getOperands();
+  for (auto e : llvm::zip(outerResults, innerResults))
+    std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
+
   // Move the loop body operations, except for its terminator, to the loop's
   // containing block.
   auto *parentBlock = forOp.getOperation()->getBlock();
-  forOp.getBody()->back().erase();
+  forOp.getBody()->getTerminator()->erase();
   parentBlock->getOperations().splice(Block::iterator(forOp),
                                       forOp.getBody()->getOperations());
   forOp.erase();
   return success();
 }
 
 /// Promotes all single iteration 'for' ops in `f`, i.e., moves
 /// their body into the containing Block.

 LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
                                          uint64_t unrollFactor) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollFactor)
     return loopUnrollByFactor(forOp, mayBeConstantTripCount.getValue());
   return loopUnrollByFactor(forOp, unrollFactor);
 }
 
-// Generates unrolled copies of AffineForOp or scf::ForOp 'loopBodyBlock', with
-// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
-// 'forOpIV' for each unrolled body.
-static void generateUnrolledLoop(
-    Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
-    function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn) {
+/// Generates unrolled copies of AffineForOp or scf::ForOp 'loopBodyBlock', with
+/// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
+/// 'forOpIV' for each unrolled body.
+static void
+generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
+                     function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
+                     ValueRange iterArgs, ValueRange yieldedValues) {
   // Builder to insert unrolled bodies just before the terminator of the body of
   // 'forOp'.
   auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
 
   // Keep a pointer to the last non-terminator operation in the original block
   // so that we know what to clone (since we are doing this in-place).
   Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);
 
   // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
+  SmallVector<Value, 4> lastYielded(yieldedValues);
+
   for (unsigned i = 1; i < unrollFactor; i++) {
     BlockAndValueMapping operandMap;
 
+    // Prepare operand map.
+    operandMap.map(iterArgs, lastYielded);
+
     // If the induction variable is used, create a remapping to the value for
     // this unrolled instance.
     if (!forOpIV.use_empty()) {
       Value ivUnroll = ivRemapFn(i, forOpIV, builder);
       operandMap.map(forOpIV, ivUnroll);
     }
 
     // Clone the original body of 'forOp'.
     for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
       builder.clone(*it, operandMap);
+
+    // Update yielded values.
+    for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
+      lastYielded[i] = operandMap.lookup(yieldedValues[i]);
   }
+
+  // Update operands of the yield statement.
+  loopBodyBlock->getTerminator()->setOperands(lastYielded);
 }
 
 /// Unrolls this loop by the specified factor. Returns success if the loop
 /// is successfully unrolled.
 LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
                                        uint64_t unrollFactor) {
   assert(unrollFactor > 0 && "unroll factor should be positive");
 
LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
   forOp.setStep(step * unrollFactor);
   generateUnrolledLoop(forOp.getBody(), forOp.getInductionVar(), unrollFactor,
                        [&](unsigned i, Value iv, OpBuilder b) {
                          // iv' = iv + i * step
                          auto d0 = b.getAffineDimExpr(0);
                          auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
                          return b.create<AffineApplyOp>(forOp.getLoc(), bumpMap,
                                                         iv);
-                       });
+                       },
+                       /*iterArgs=*/{}, /*yieldedValues=*/{});
 
   // Promote the loop body up if this has turned into a single iteration loop.
   promoteIfSingleIteration(forOp);
   return success();
 }
 
 /// Unrolls 'forOp' by 'unrollFactor', returns success if the loop is unrolled.
 LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
   }
 
   // Create epilogue clean up loop starting at 'upperBoundUnrolled'.
   if (generateEpilogueLoop) {
     OpBuilder epilogueBuilder(forOp.getOperation()->getBlock(),
                               std::next(Block::iterator(forOp)));
     auto epilogueForOp = cast<scf::ForOp>(epilogueBuilder.clone(*forOp));
     epilogueForOp.setLowerBound(upperBoundUnrolled);
+
+    // Update uses of loop results.
+    auto results = forOp.getResults();
+    auto epilogueResults = epilogueForOp.getResults();
+    auto epilogueIterOperands = epilogueForOp.getIterOperands();
+
+    for (auto e : llvm::zip(results, epilogueResults, epilogueIterOperands)) {
+      std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
+      epilogueForOp.getOperation()->replaceUsesOfWith(std::get<2>(e),
+                                                      std::get<0>(e));
+    }
     promoteIfSingleIteration(epilogueForOp);
   }
 
   // Create unrolled loop.
   forOp.setUpperBound(upperBoundUnrolled);
   forOp.setStep(stepUnrolled);
-  generateUnrolledLoop(forOp.getBody(), forOp.getInductionVar(), unrollFactor,
-                       [&](unsigned i, Value iv, OpBuilder b) {
-                         // iv' = iv + step * i;
-                         auto stride = b.create<MulIOp>(
-                             loc, step, b.create<ConstantIndexOp>(loc, i));
-                         return b.create<AddIOp>(loc, iv, stride);
-                       });
+
+  auto iterArgs = ValueRange(forOp.getRegionIterArgs());
+  auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
+
+  generateUnrolledLoop(
+      forOp.getBody(), forOp.getInductionVar(), unrollFactor,
+        auto stride =
+            b.create<MulIOp>(loc, step, b.create<ConstantIndexOp>(loc, i));
+      },
+      iterArgs, yieldedValues);
   // Promote the loop body up if this has turned into a single iteration loop.
   promoteIfSingleIteration(forOp);
   return success();
 }
 
 LogicalResult mlir::loopUnrollJamUpToFactor(AffineForOp forOp,
                                             uint64_t unrollJamFactor) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);

mlir/test/Transforms/scf-loop-unroll.mlir

+// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=3" -split-input-file -canonicalize | FileCheck %s
+
+// CHECK-LABEL: scf_loop_unroll_single
+func @scf_loop_unroll_single(%arg0 : f32, %arg1 : f32) -> f32 {
+  %from = constant 0 : index
+  %to = constant 10 : index
+  %step = constant 1 : index
+  %sum = scf.for %iv = %from to %to step %step iter_args(%sum_iter = %arg0) -> (f32) {
+    %next = addf %sum_iter, %arg1 : f32
+    scf.yield %next : f32
+  }
+  // CHECK:      %[[SUM:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[V0:.*]] =
+  // CHECK-NEXT:   %[[V1:.*]] = addf %[[V0]]
+  // CHECK-NEXT:   %[[V2:.*]] = addf %[[V1]]
+  // CHECK-NEXT:   %[[V3:.*]] = addf %[[V2]]
+  // CHECK-NEXT:   scf.yield %[[V3]]
+  // CHECK-NEXT: }
+  // CHECK-NEXT: %[[RES:.*]] = addf %[[SUM]],
+  // CHECK-NEXT: return %[[RES]]
+  return %sum : f32
+}
+
+// CHECK-LABEL: scf_loop_unroll_double_symbolic_ub
+// CHECK-SAME:     (%{{.*}}: f32, %{{.*}}: f32, %[[N:.*]]: index)
+func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : index) -> (f32,f32) {
+  %from = constant 0 : index
+  %step = constant 1 : index
+  %sum:2 = scf.for %iv = %from to %n step %step iter_args(%i0 = %arg0, %i1 = %arg1) -> (f32, f32) {
+    %sum0 = addf %i0, %arg0 : f32
+    %sum1 = addf %i1, %arg1 : f32
+    scf.yield %sum0, %sum1 : f32, f32
+  }
+  return %sum#0, %sum#1 : f32, f32
+  // CHECK:      %[[C0:.*]] = constant 0 : index
+  // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+  // CHECK-NEXT: %[[C3:.*]] = constant 3 : index
+  // CHECK-NEXT: %[[REM:.*]] = remi_signed %[[N]], %[[C3]]
+  // CHECK-NEXT: %[[UB:.*]] = subi %[[N]], %[[REM]]
+  // CHECK-NEXT: %[[SUM:.*]]:2 = scf.for {{.*}} = %[[C0]] to %[[UB]] step %[[C3]] iter_args
+  // CHECK:      }
+  // CHECK-NEXT: %[[SUM1:.*]]:2 = scf.for {{.*}} = %[[UB]] to %[[N]] step %[[C1]] iter_args(%[[V1:.*]] = %[[SUM]]#0, %[[V2:.*]] = %[[SUM]]#1)
+  // CHECK:      }
+  // CHECK-NEXT: return %[[SUM1]]#0, %[[SUM1]]#1
+}