[mlir] Support iterator arguments in the loop unrolling

[rriddle]

mlir/lib/Transforms/Utils/LoopUtils.cpp

491

Use ValueRange instead of ArrayRefs here, it removes the need to create a smallvector below.

[liufengdb]

fix some test failures

[liufengdb]

remove the redundant comments

[liufengdb]

remove unused

[ftynse]

mlir/lib/Transforms/Utils/LoopUtils.cpp

486–487

Can this use ValueRange instead of ArrayRef?

506

Nit: camelCase please

507

lookupOrDefault(terminator->getOperand(..)) will remove the need for if below

514

Can we just bail out of transformation in this case, instead of producing incorrect code?

mlir/test/Dialect/SCF/loop-unroll.mlir

263

Let's not pattern-match SSA value names

[bondhugula]

@liufengdb Can you add a commit summary please even if brief?

mlir/lib/Transforms/Utils/LoopUtils.cpp

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.
+// 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,
+    Block *loopBodyBlock, Value forOpIV, ArrayRef<BlockArgument> iterArgs,
+    uint64_t unrollFactor,
     function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn) {
   // 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).
+  // Keep a pointer to the 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);
+  Block::iterator terminator = std::next(srcBlockEnd);
+
+  BlockAndValueMapping operandMap;
 
   // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
   for (unsigned i = 1; i < unrollFactor; i++) {
-    BlockAndValueMapping operandMap;
+    // If there are iterator arguments, map them to the return results from the
+    // previous iteration.
+    for (auto iterArg : llvm::enumerate(iterArgs)) {
+      Value iterArgValue = iterArg.value();
+      Value mapped_returned =
+          operandMap.lookupOrNull(terminator->getOperand(iterArg.index()));
+      if (!mapped_returned)
+        mapped_returned = terminator->getOperand(iterArg.index());
+      operandMap.map(iterArgValue, mapped_returned);
+    }
+
+    // TODO: Seems like the operandMap needs to be cleared for nested loops.
+    // In other words, the nested loop with iteration argument isn't supported.
+    if (iterArgs.empty())
+      operandMap.clear();
 
     // 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'.
+    // Clone the original body of 'forOp' without the terminator.
     for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
       builder.clone(*it, operandMap);
   }
+
+  // Finally, adjust the operands of the terminator to be the unrolled ones.
+  for (auto operand : llvm::enumerate(terminator->getOperands())) {
+    terminator->setOperand(operand.index(), operandMap.lookup(operand.value()));
+  }
 }
 
 /// 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");
 
if (getLargestDivisorOfTripCount(forOp) % unrollFactor != 0) {
     // Adjust upper bound of the original loop; this is the same as the lower
     // bound of the cleanup loop.
     forOp.setUpperBound(cleanupOperands, cleanupMap);
   }
 
   // Scale the step of loop being unrolled by unroll factor.
   int64_t step = forOp.getStep();
   forOp.setStep(step * unrollFactor);
-  generateUnrolledLoop(forOp.getBody(), forOp.getInductionVar(), unrollFactor,
+  generateUnrolledLoop(forOp.getBody(), forOp.getInductionVar(),
+                       /*iterArgs=*/{}, 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);
                        });
 
if (generateEpilogueLoop) {
     auto epilogueForOp = cast<scf::ForOp>(epilogueBuilder.clone(*forOp));
     epilogueForOp.setLowerBound(upperBoundUnrolled);
     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);
-                       });
+  SmallVector<Value, 4> initArgs(forOp.initArgs().begin(),
+                                 forOp.initArgs().end());
+  generateUnrolledLoop(
+      forOp.getBody(), forOp.getInductionVar(), forOp.getRegionIterArgs(),
+      unrollFactor, [&](unsigned i, Value iv, OpBuilder b) {
+        auto stride =
+            b.create<MulIOp>(loc, step, b.create<ConstantIndexOp>(loc, i));
+
   // 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/Dialect/SCF/loop-unroll.mlir

 //  UNROLL-BY-3-NEXT:    store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:    %[[C2_IV:.*]] = constant 2 : index
 //  UNROLL-BY-3-NEXT:    %[[V2:.*]] = muli %[[C1]], %[[C2_IV]] : index
 //  UNROLL-BY-3-NEXT:    %[[V3:.*]] = addi %[[IV]], %[[V2]] : index
 //  UNROLL-BY-3-NEXT:    store %{{.*}}, %[[MEM]][%[[V3]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:  }
 //  UNROLL-BY-3-NEXT:  store %{{.*}}, %[[MEM]][%[[C9]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:  return
+
+func @reduce_loop_unroll(%buffer: memref<1024xf32>, %lb: index, %ub: index, %step: index) -> (f32) {
+  %sum_0 = constant 0.0 : f32
+  %sum = scf.for %iv = %lb to %ub step %step iter_args(%sum_iter = %sum_0) -> (f32) {
+    %t = load %buffer[%iv] : memref<1024xf32>
+    %sum_next = addf %sum_iter, %t : f32
+    scf.yield %sum_next : f32
+  }
+  return %sum : f32
+
+// UNROLL-BY-2-LABEL: func @reduce_loop_unroll
+// UNROLL-BY-2: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[AR:.*]] = %{{.*}}) -> (f32) {
+// UNROLL-BY-2: %[[ADD0:.*]] = addf %arg5, %11 : f32
+// UNROLL-BY-2: %[[ADD1:.*]] = addf %[[ADD0]], %15 : f32
+// UNROLL-BY-2: scf.yield %[[ADD1]] : f32
+// UNROLL-BY-2: }
+
+// UNROLL-BY-3-LABEL: func @reduce_loop_unroll
+// UNROLL-BY-3: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[AR:.*]] = %{{.*}}) -> (f32) {
+// UNROLL-BY-3: %[[ADD0:.*]] = addf %arg5, %11 : f32
+// UNROLL-BY-3: %[[ADD1:.*]] = addf %[[ADD0]], %15 : f32
+// UNROLL-BY-3: %[[ADD2:.*]] = addf %[[ADD1]], %19 : f32
+// UNROLL-BY-3: scf.yield %[[ADD2]] : f32
+// UNROLL-BY-3: }
+}
 No newline at end of file