[Mlir-commits] [mlir] 3a2ff98 - Support post-processing Ops in unrolled loop iterations
Mehdi Amini
llvmlistbot at llvm.org
Wed Aug 11 16:11:20 PDT 2021
Author: Tyler Augustine
Date: 2021-08-11T23:11:10Z
New Revision: 3a2ff982d78fe20cc34454f5b28a95eecac9c142
URL: https://github.com/llvm/llvm-project/commit/3a2ff982d78fe20cc34454f5b28a95eecac9c142
DIFF: https://github.com/llvm/llvm-project/commit/3a2ff982d78fe20cc34454f5b28a95eecac9c142.diff
LOG: Support post-processing Ops in unrolled loop iterations
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.
Reviewed By: mehdi_amini
Differential Revision: https://reviews.llvm.org/D107789
Added:
Modified:
mlir/include/mlir/Transforms/LoopUtils.h
mlir/lib/Transforms/Utils/LoopUtils.cpp
mlir/test/Dialect/SCF/loop-unroll.mlir
mlir/test/lib/Transforms/TestLoopUnrolling.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
index c49a333d2791b..f6aa469e42530 100644
--- a/mlir/include/mlir/Transforms/LoopUtils.h
+++ b/mlir/include/mlir/Transforms/LoopUtils.h
@@ -40,9 +40,14 @@ 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.
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index b12532d7f316f..15aea84c57dbc 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1072,15 +1072,20 @@ LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
/// 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<Value(unsigned, Value, OpBuilder)> ivRemapFn,
+ function_ref<void(unsigned, Operation *, OpBuilder)> 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 @@ generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
}
// 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)
@@ -1186,6 +1199,7 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
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.
@@ -1194,8 +1208,9 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
}
/// 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);
@@ -1300,7 +1315,7 @@ LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
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();
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index a824a7df51ced..e7c0b32eb1cb6 100644
--- 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 @@ func @static_loop_unroll_by_2(%arg0 : 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:.*0]][%{{.*}}] {unrolled_iteration = 0 : ui32} : memref<?xf32>
+// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM]][%{{.*}}] {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>) {
@@ -269,4 +274,5 @@ func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {
// 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
\ 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
index e0e94bb6ba073..e779688f45c07 100644
--- a/mlir/test/lib/Transforms/TestLoopUnrolling.cpp
+++ b/mlir/test/lib/Transforms/TestLoopUnrolling.cpp
@@ -40,9 +40,11 @@ class TestLoopUnrollingPass
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 @@ class TestLoopUnrollingPass
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)};
More information about the Mlir-commits
mailing list