diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h --- a/mlir/include/mlir/Transforms/LoopUtils.h +++ b/mlir/include/mlir/Transforms/LoopUtils.h @@ -40,9 +40,14 @@ /// 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 annotateFn = nullptr); +LogicalResult loopUnrollByFactor( + scf::ForOp forOp, uint64_t unrollFactor, + function_ref annotateFn = nullptr); /// Unrolls this loop by the specified unroll factor or its trip count, /// whichever is lower. diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -1072,15 +1072,20 @@ /// 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 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 ivRemapFn, + function_ref annotateFn, + ValueRange iterArgs, ValueRange yieldedValues) { // 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); @@ -1102,22 +1107,30 @@ } // 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 annotateFn) { assert(unrollFactor > 0 && "unroll factor should be positive"); if (unrollFactor == 1) @@ -1186,6 +1199,7 @@ auto bumpMap = AffineMap::get(1, 0, d0 + i * step); return b.create(forOp.getLoc(), bumpMap, iv); }, + /*annotateFn=*/annotateFn, /*iterArgs=*/iterArgs, /*yieldedValues=*/yieldedValues); // Promote the loop body up if this has turned into a single iteration loop. @@ -1194,8 +1208,9 @@ } /// 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 annotateFn) { assert(unrollFactor > 0 && "expected positive unroll factor"); if (unrollFactor == 1) return promoteIfSingleIteration(forOp); @@ -1300,7 +1315,7 @@ b.create(loc, step, b.create(loc, i)); return b.create(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(); diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir --- a/mlir/test/Dialect/SCF/loop-unroll.mlir +++ b/mlir/test/Dialect/SCF/loop-unroll.mlir @@ -2,6 +2,7 @@ // 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, @@ -179,6 +180,10 @@ // 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 +// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM]][%[[V1]]] {unrolled_iteration = 1 : ui32} : memref + // 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) { @@ -269,4 +274,5 @@ // UNROLL-UP-TO-NEXT: store %{{.*}}, %[[MEM]][%[[V0]]] : memref // UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}} // UNROLL-UP-TO-NEXT: affine.store %{{.*}}, %[[MEM]][%[[V1]]] : memref -// UNROLL-UP-TO-NEXT: return \ No newline at end of file +// UNROLL-UP-TO-NEXT: return + diff --git a/mlir/test/lib/Transforms/TestLoopUnrolling.cpp b/mlir/test/lib/Transforms/TestLoopUnrolling.cpp --- a/mlir/test/lib/Transforms/TestLoopUnrolling.cpp +++ b/mlir/test/lib/Transforms/TestLoopUnrolling.cpp @@ -40,9 +40,11 @@ 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 { @@ -52,12 +54,20 @@ 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 unrollFactor{*this, "unroll-factor", llvm::cl::desc("Loop unroll factor."), llvm::cl::init(1)}; + Option annotateLoop{*this, "annotate", + llvm::cl::desc("Annotate unrolled iterations."), + llvm::cl::init(false)}; Option unrollUpToFactor{*this, "unroll-up-to-factor", llvm::cl::desc("Loop unroll up to factor."), llvm::cl::init(false)};