Support post-processing Ops in unrolled loop iterations

[nimiwio]

Summary

This can be useful when one needs to know which unrolled iteration an Op belongs to, for example, conveying noalias information among memory-affecting ops in parallel-access loops.

[mehdi_amini]

LGTM overall

mlir/include/mlir/Transforms/LoopUtils.h

44

Not sure about this TODO, can you just fill the doc here?

mlir/lib/Transforms/Utils/LoopUtils.cpp

1088

Nit: remove spurious braces

mlir/test/lib/Transforms/TestLoopUnrolling.cpp

50

I'm not sure I get this TODO?

[nimiwio]

Addressed comments.

mlir/test/lib/Transforms/TestLoopUnrolling.cpp

50

Removed.

[nimiwio]

Thank you for the review! Can this be committed?

[mehdi_amini]

In D107789#2937367, @nimiwio wrote:

Thank you for the review! Can this be committed?

To clarify: are you asking me to commit for you (don't have commit access) or do you ask confirmation that you can land it now?

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

198

One nit: in general I prefer targeted testing.
Meaning we could write minimally:

// UNROLL-BY-2-ANNOTATE: memref.store {{.*}} unrolled_iteration = 0
// UNROLL-BY-2-ANNOTATE: memref.store {{.*}} unrolled_iteration = 1

The rest is already tested in other tests.

[nimiwio]

This comment was removed by nimiwio.

[nimiwio]

Minified test

[nimiwio]

In D107789#2938910, @mehdi_amini wrote:

In D107789#2937367, @nimiwio wrote:

Thank you for the review! Can this be committed?

To clarify: are you asking me to commit for you (don't have commit access) or do you ask confirmation that you can land it now?

I guess both.. This is my first commit (so I definitely don't have commit access) and I'm also unsure of what amount of review/consensus sufficient to land a change. If it's okay to commit now, that would be good with me.

[nimiwio]

Fix test

mlir/include/mlir/Transforms/LoopUtils.h

 } // end namespace scf
 
 /// Unrolls this for operation completely if the trip count is known to be
 /// constant. Returns failure otherwise.
 LogicalResult loopUnrollFull(AffineForOp forOp);
 
 /// Unrolls this for operation by the specified unroll factor. Returns failure
 /// if the loop cannot be unrolled either due to restrictions or due to invalid
-/// unroll factors. Requires positive loop bounds and step.
-LogicalResult loopUnrollByFactor(AffineForOp forOp, uint64_t unrollFactor);
-LogicalResult loopUnrollByFactor(scf::ForOp forOp, uint64_t unrollFactor);
+/// unroll factors. Requires positive loop bounds and step. If specified,
+/// annotates the Ops in each unrolled iteration by applying `annotateFn`.
+LogicalResult loopUnrollByFactor(
+    AffineForOp forOp, uint64_t unrollFactor,
+    function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
+LogicalResult loopUnrollByFactor(
+    scf::ForOp forOp, uint64_t unrollFactor,
+    function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
 
 /// Unrolls this loop by the specified unroll factor or its trip count,
 /// whichever is lower.
 LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor);
 
 /// Returns true if `loops` is a perfectly nested loop nest, where loops appear
 /// in it from outermost to innermost.
 bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef<AffineForOp> loops);

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,
-                     ValueRange iterArgs, ValueRange yieldedValues) {
+/// 'forOpIV' for each unrolled body. If specified, annotates the Ops in each
+/// unrolled iteration using annotateFn.
+static void generateUnrolledLoop(
+    Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
+    function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
   // Builder to insert unrolled bodies just before the terminator of the body of
   // 'forOp'.
   auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
 
+  if (!annotateFn)
+    annotateFn = [](unsigned, Operation *, OpBuilder) {};
+
   // 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);
+    for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
+      Operation *clonedOp = builder.clone(*it, operandMap);
+      annotateFn(i, clonedOp, builder);
+    }
 
     // Update yielded values.
     for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
       lastYielded[i] = operandMap.lookup(yieldedValues[i]);
   }
 
+  // Make sure we annotate the Ops in the original body. We do this last so that
+  // any annotations are not copied into the cloned Ops above.
+  for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
+    annotateFn(0, &*it, builder);
+
   // 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) {
+LogicalResult mlir::loopUnrollByFactor(
+    AffineForOp forOp, uint64_t unrollFactor,
+    function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "unroll factor should be positive");
 
   if (unrollFactor == 1)
     return promoteIfSingleIteration(forOp);
 
   // Nothing in the loop body other than the terminator.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
LogicalResult mlir::loopUnrollByFactor(
   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);
       },
+      /*annotateFn=*/annotateFn,
       /*iterArgs=*/iterArgs, /*yieldedValues=*/yieldedValues);
 
   // Promote the loop body up if this has turned into a single iteration loop.
   (void)promoteIfSingleIteration(forOp);
   return success();
 }
 
 /// Unrolls 'forOp' by 'unrollFactor', returns success if the loop is unrolled.
-LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
-                                       uint64_t unrollFactor) {
+LogicalResult mlir::loopUnrollByFactor(
+    scf::ForOp forOp, uint64_t unrollFactor,
+    function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "expected positive unroll factor");
   if (unrollFactor == 1)
     return promoteIfSingleIteration(forOp);
 
   // Return if the loop body is empty.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
LogicalResult mlir::loopUnrollByFactor(
   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);
       },
-      iterArgs, yieldedValues);
+      annotateFn, iterArgs, yieldedValues);
   // Promote the loop body up if this has turned into a single iteration loop.
   (void)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
+// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE
 // RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
 
 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 {
     memref.store %0, %arg3[%i0] : memref<?xf32>
   }

 //  UNROLL-BY-2-NEXT:    memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
 //  UNROLL-BY-2-NEXT:    %[[C1_IV:.*]] = constant 1 : index
 //  UNROLL-BY-2-NEXT:    %[[V0:.*]] = muli %[[C1]], %[[C1_IV]] : index
 //  UNROLL-BY-2-NEXT:    %[[V1:.*]] = addi %[[IV]], %[[V0]] : index
 //  UNROLL-BY-2-NEXT:    memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
 //  UNROLL-BY-2-NEXT:  }
 //  UNROLL-BY-2-NEXT:  return
 
+// UNROLL-BY-2-ANNOTATE-LABEL: func @static_loop_unroll_by_2
+// UNROLL-BY-2-ANNOTATE:    memref.store %{{.*}}, %[[MEM]][%[[IV]]] {unrolled_iteration = 0 : ui32} : memref<?xf32>
+// UNROLL-BY-2-ANNOTATE:    memref.store %{{.*}}, %[[MEM]][%[[V1]]] {unrolled_iteration = 1 : ui32} : memref<?xf32>
+
 // Test that epilogue clean up loop is generated (trip count is not
 // a multiple of unroll factor).
 func @static_loop_unroll_by_3(%arg0 : memref<?xf32>) {
   %0 = constant 7.0 : f32
   %lb = constant 0 : index
   %ub = constant 20 : index
   %step = constant 1 : index
   scf.for %i0 = %lb to %ub step %step {
     memref.store %0, %arg0[%i0] : memref<?xf32>
   }
   return
 }
 
 // UNROLL-BY-3-LABEL: func @static_loop_unroll_by_3
 //  UNROLL-BY-3-SAME:  %[[MEM:.*0]]: memref<?xf32>
 //
 //   UNROLL-BY-3-DAG:  %[[C0:.*]] = constant 0 : index
 //   UNROLL-BY-3-DAG:  %[[C1:.*]] = constant 1 : index
 //   UNROLL-BY-3-DAG:  %[[C20:.*]] = constant 20 : index
 //   UNROLL-BY-3-DAG:  %[[C18:.*]] = constant 18 : index

 //  UNROLL-UP-TO-SAME:  %[[MEM:.*0]]: memref<?xf32>
 //
 //   UNROLL-UP-TO-DAG:  %[[C0:.*]] = constant 0 : index
 //   UNROLL-UP-TO-DAG:  %[[C2:.*]] = constant 2 : index
 //   UNROLL-UP-TO-NEXT: %[[V0:.*]] = affine.apply {{.*}}
 //   UNROLL-UP-TO-NEXT: store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>
 //   UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}}
 //   UNROLL-UP-TO-NEXT: affine.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
-//   UNROLL-UP-TO-NEXT: return
+//   UNROLL-UP-TO-NEXT: return
+

mlir/test/lib/Transforms/TestLoopUnrolling.cpp

 public:
   StringRef getArgument() const final { return "test-loop-unrolling"; }
   StringRef getDescription() const final {
     return "Tests loop unrolling transformation";
   }
   TestLoopUnrollingPass() = default;
   TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
   explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
-                                 unsigned loopDepthParam) {
+                                 unsigned loopDepthParam,
+                                 bool annotateLoopParam) {
     unrollFactor = unrollFactorParam;
     loopDepth = loopDepthParam;
+    annotateLoop = annotateLoopParam;
   }
 
   void runOnFunction() override {
     FuncOp func = getFunction();
     SmallVector<scf::ForOp, 4> loops;
     func.walk([&](scf::ForOp forOp) {
       if (getNestingDepth(forOp) == loopDepth)
         loops.push_back(forOp);
     });
+    auto annotateFn = [this](unsigned i, Operation *op, OpBuilder b) {
+      if (annotateLoop) {
+        op->setAttr("unrolled_iteration", b.getUI32IntegerAttr(i));
+      }
+    };
     for (auto loop : loops)
-      (void)loopUnrollByFactor(loop, unrollFactor);
+      (void)loopUnrollByFactor(loop, unrollFactor, annotateFn);
   }
   Option<uint64_t> unrollFactor{*this, "unroll-factor",
                                 llvm::cl::desc("Loop unroll factor."),
                                 llvm::cl::init(1)};
+  Option<bool> annotateLoop{*this, "annotate",
+                            llvm::cl::desc("Annotate unrolled iterations."),
+                            llvm::cl::init(false)};
   Option<bool> unrollUpToFactor{*this, "unroll-up-to-factor",
                                 llvm::cl::desc("Loop unroll up to factor."),
                                 llvm::cl::init(false)};
   Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
                              llvm::cl::init(0)};
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLoopUnrollingPass() {
   PassRegistration<TestLoopUnrollingPass>();
 }
 } // namespace test
 } // namespace mlir