[mlir] Support iterator arguments in the loop unrolling

[rriddle]

mlir/lib/Transforms/Utils/LoopUtils.cpp

485

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

481–482

Can this use ValueRange instead of ArrayRef?

510

Nit: camelCase please

511

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

518

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

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

273

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

LogicalResult mlir::loopUnrollUpToFactor(AffineForOp 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) {
+static void
+generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
+                     function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
+                     ArrayRef<BlockArgument> iterArgs,
+                     ArrayRef<Value> iterInits) {
   // 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);
+  // 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(), 1);
+
+  BlockAndValueMapping operandMap;
+
+  // If there are iterator arguments, create a remapping to the initial values.
+  for (auto iterArg : llvm::zip(iterArgs, iterInits)) {
+    Value iterArgValue = std::get<0>(iterArg);
+    operandMap.map(iterArgValue, std::get<1>(iterArg));
+  }
 
   // Unroll the contents of 'forOp' (append unrollFactor - 1 additional copies).
+  Operation *last_cloned = nullptr;
   for (unsigned i = 1; i < unrollFactor; i++) {
-    BlockAndValueMapping operandMap;
+    // `last_cloned` points to the terminator op of the last iteration
+    // instance. Remap the iterator arguments to results of that iteration.
+    if (last_cloned && last_cloned->isKnownTerminator()) {
+      for (auto iterArg : llvm::enumerate(iterArgs)) {
+        Value iterArgValue = iterArg.value();
+        auto mapped_returned =
+            operandMap.lookup(last_cloned->getOperand(iterArg.index()));
+        operandMap.map(iterArgValue, mapped_returned);
+      }
+      // Remove the terminator op of the last iteration instance.
+      last_cloned->erase();
+    }
 
     // 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', including the terminator op. The
+    // terminator op will be erased if there is anther unrolling iteration.
     for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
-      builder.clone(*it, operandMap);
+      last_cloned = builder.clone(*it, operandMap);
   }
 }
 
 /// 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=*/{}, /*iterInits=*/{});
 
   // 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,
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(), unrollFactor,
+        auto stride =
+            b.create<MulIOp>(loc, step, b.create<ConstantIndexOp>(loc, i));
+      },
+      forOp.getRegionIterArgs(), initArgs);
+
   // 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

 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2' | FileCheck %s --check-prefix UNROLL-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-3
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
 
+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
+}
+
 func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
                           %arg3: memref<?xf32>) {
   %0 = constant 7.0 : f32
   scf.for %i0 = %arg0 to %arg1 step %arg2 {
     store %0, %arg3[%i0] : memref<?xf32>
   }
   return
 }

 //  UNROLL-BY-3-NEXT:    %[[V1:.*]] = addi %[[IV]], %[[V0]] : index
 //  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