Skip to content

Commit 3a2ff98

Browse files
nimiwiojoker-eph
authored andcommitted
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
1 parent 856a6a5 commit 3a2ff98

File tree

4 files changed

+54
-18
lines changed

4 files changed

+54
-18
lines changed

mlir/include/mlir/Transforms/LoopUtils.h

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,14 @@ LogicalResult loopUnrollFull(AffineForOp forOp);
4040

4141
/// Unrolls this for operation by the specified unroll factor. Returns failure
4242
/// if the loop cannot be unrolled either due to restrictions or due to invalid
43-
/// unroll factors. Requires positive loop bounds and step.
44-
LogicalResult loopUnrollByFactor(AffineForOp forOp, uint64_t unrollFactor);
45-
LogicalResult loopUnrollByFactor(scf::ForOp forOp, uint64_t unrollFactor);
43+
/// unroll factors. Requires positive loop bounds and step. If specified,
44+
/// annotates the Ops in each unrolled iteration by applying `annotateFn`.
45+
LogicalResult loopUnrollByFactor(
46+
AffineForOp forOp, uint64_t unrollFactor,
47+
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
48+
LogicalResult loopUnrollByFactor(
49+
scf::ForOp forOp, uint64_t unrollFactor,
50+
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn = nullptr);
4651

4752
/// Unrolls this loop by the specified unroll factor or its trip count,
4853
/// whichever is lower.

mlir/lib/Transforms/Utils/LoopUtils.cpp

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1072,15 +1072,20 @@ LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
10721072

10731073
/// Generates unrolled copies of AffineForOp or scf::ForOp 'loopBodyBlock', with
10741074
/// associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap
1075-
/// 'forOpIV' for each unrolled body.
1076-
static void
1077-
generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
1078-
function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
1079-
ValueRange iterArgs, ValueRange yieldedValues) {
1075+
/// 'forOpIV' for each unrolled body. If specified, annotates the Ops in each
1076+
/// unrolled iteration using annotateFn.
1077+
static void generateUnrolledLoop(
1078+
Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
1079+
function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
1080+
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
1081+
ValueRange iterArgs, ValueRange yieldedValues) {
10801082
// Builder to insert unrolled bodies just before the terminator of the body of
10811083
// 'forOp'.
10821084
auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
10831085

1086+
if (!annotateFn)
1087+
annotateFn = [](unsigned, Operation *, OpBuilder) {};
1088+
10841089
// Keep a pointer to the last non-terminator operation in the original block
10851090
// so that we know what to clone (since we are doing this in-place).
10861091
Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);
@@ -1102,22 +1107,30 @@ generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
11021107
}
11031108

11041109
// Clone the original body of 'forOp'.
1105-
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
1106-
builder.clone(*it, operandMap);
1110+
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
1111+
Operation *clonedOp = builder.clone(*it, operandMap);
1112+
annotateFn(i, clonedOp, builder);
1113+
}
11071114

11081115
// Update yielded values.
11091116
for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
11101117
lastYielded[i] = operandMap.lookup(yieldedValues[i]);
11111118
}
11121119

1120+
// Make sure we annotate the Ops in the original body. We do this last so that
1121+
// any annotations are not copied into the cloned Ops above.
1122+
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
1123+
annotateFn(0, &*it, builder);
1124+
11131125
// Update operands of the yield statement.
11141126
loopBodyBlock->getTerminator()->setOperands(lastYielded);
11151127
}
11161128

11171129
/// Unrolls this loop by the specified factor. Returns success if the loop
11181130
/// is successfully unrolled.
1119-
LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
1120-
uint64_t unrollFactor) {
1131+
LogicalResult mlir::loopUnrollByFactor(
1132+
AffineForOp forOp, uint64_t unrollFactor,
1133+
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
11211134
assert(unrollFactor > 0 && "unroll factor should be positive");
11221135

11231136
if (unrollFactor == 1)
@@ -1186,6 +1199,7 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
11861199
auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
11871200
return b.create<AffineApplyOp>(forOp.getLoc(), bumpMap, iv);
11881201
},
1202+
/*annotateFn=*/annotateFn,
11891203
/*iterArgs=*/iterArgs, /*yieldedValues=*/yieldedValues);
11901204

11911205
// Promote the loop body up if this has turned into a single iteration loop.
@@ -1194,8 +1208,9 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
11941208
}
11951209

11961210
/// Unrolls 'forOp' by 'unrollFactor', returns success if the loop is unrolled.
1197-
LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
1198-
uint64_t unrollFactor) {
1211+
LogicalResult mlir::loopUnrollByFactor(
1212+
scf::ForOp forOp, uint64_t unrollFactor,
1213+
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
11991214
assert(unrollFactor > 0 && "expected positive unroll factor");
12001215
if (unrollFactor == 1)
12011216
return promoteIfSingleIteration(forOp);
@@ -1300,7 +1315,7 @@ LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
13001315
b.create<MulIOp>(loc, step, b.create<ConstantIndexOp>(loc, i));
13011316
return b.create<AddIOp>(loc, iv, stride);
13021317
},
1303-
iterArgs, yieldedValues);
1318+
annotateFn, iterArgs, yieldedValues);
13041319
// Promote the loop body up if this has turned into a single iteration loop.
13051320
(void)promoteIfSingleIteration(forOp);
13061321
return success();

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-3
33
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
44
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
5+
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE
56
// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
67

78
func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
@@ -179,6 +180,10 @@ func @static_loop_unroll_by_2(%arg0 : memref<?xf32>) {
179180
// UNROLL-BY-2-NEXT: }
180181
// UNROLL-BY-2-NEXT: return
181182

183+
// UNROLL-BY-2-ANNOTATE-LABEL: func @static_loop_unroll_by_2
184+
// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM:.*0]][%{{.*}}] {unrolled_iteration = 0 : ui32} : memref<?xf32>
185+
// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM]][%{{.*}}] {unrolled_iteration = 1 : ui32} : memref<?xf32>
186+
182187
// Test that epilogue clean up loop is generated (trip count is not
183188
// a multiple of unroll factor).
184189
func @static_loop_unroll_by_3(%arg0 : memref<?xf32>) {
@@ -269,4 +274,5 @@ func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {
269274
// UNROLL-UP-TO-NEXT: store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>
270275
// UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}}
271276
// UNROLL-UP-TO-NEXT: affine.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
272-
// UNROLL-UP-TO-NEXT: return
277+
// UNROLL-UP-TO-NEXT: return
278+

mlir/test/lib/Transforms/TestLoopUnrolling.cpp

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,11 @@ class TestLoopUnrollingPass
4040
TestLoopUnrollingPass() = default;
4141
TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
4242
explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
43-
unsigned loopDepthParam) {
43+
unsigned loopDepthParam,
44+
bool annotateLoopParam) {
4445
unrollFactor = unrollFactorParam;
4546
loopDepth = loopDepthParam;
47+
annotateLoop = annotateLoopParam;
4648
}
4749

4850
void runOnFunction() override {
@@ -52,12 +54,20 @@ class TestLoopUnrollingPass
5254
if (getNestingDepth(forOp) == loopDepth)
5355
loops.push_back(forOp);
5456
});
57+
auto annotateFn = [this](unsigned i, Operation *op, OpBuilder b) {
58+
if (annotateLoop) {
59+
op->setAttr("unrolled_iteration", b.getUI32IntegerAttr(i));
60+
}
61+
};
5562
for (auto loop : loops)
56-
(void)loopUnrollByFactor(loop, unrollFactor);
63+
(void)loopUnrollByFactor(loop, unrollFactor, annotateFn);
5764
}
5865
Option<uint64_t> unrollFactor{*this, "unroll-factor",
5966
llvm::cl::desc("Loop unroll factor."),
6067
llvm::cl::init(1)};
68+
Option<bool> annotateLoop{*this, "annotate",
69+
llvm::cl::desc("Annotate unrolled iterations."),
70+
llvm::cl::init(false)};
6171
Option<bool> unrollUpToFactor{*this, "unroll-up-to-factor",
6272
llvm::cl::desc("Loop unroll up to factor."),
6373
llvm::cl::init(false)};

0 commit comments

Comments
 (0)