@@ -1072,15 +1072,20 @@ LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
1072
1072
1073
1073
// / Generates unrolled copies of AffineForOp or scf::ForOp 'loopBodyBlock', with
1074
1074
// / 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) {
1080
1082
// Builder to insert unrolled bodies just before the terminator of the body of
1081
1083
// 'forOp'.
1082
1084
auto builder = OpBuilder::atBlockTerminator (loopBodyBlock);
1083
1085
1086
+ if (!annotateFn)
1087
+ annotateFn = [](unsigned , Operation *, OpBuilder) {};
1088
+
1084
1089
// Keep a pointer to the last non-terminator operation in the original block
1085
1090
// so that we know what to clone (since we are doing this in-place).
1086
1091
Block::iterator srcBlockEnd = std::prev (loopBodyBlock->end (), 2 );
@@ -1102,22 +1107,30 @@ generateUnrolledLoop(Block *loopBodyBlock, Value forOpIV, uint64_t unrollFactor,
1102
1107
}
1103
1108
1104
1109
// 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
+ }
1107
1114
1108
1115
// Update yielded values.
1109
1116
for (unsigned i = 0 , e = lastYielded.size (); i < e; i++)
1110
1117
lastYielded[i] = operandMap.lookup (yieldedValues[i]);
1111
1118
}
1112
1119
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
+
1113
1125
// Update operands of the yield statement.
1114
1126
loopBodyBlock->getTerminator ()->setOperands (lastYielded);
1115
1127
}
1116
1128
1117
1129
// / Unrolls this loop by the specified factor. Returns success if the loop
1118
1130
// / 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) {
1121
1134
assert (unrollFactor > 0 && " unroll factor should be positive" );
1122
1135
1123
1136
if (unrollFactor == 1 )
@@ -1186,6 +1199,7 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
1186
1199
auto bumpMap = AffineMap::get (1 , 0 , d0 + i * step);
1187
1200
return b.create <AffineApplyOp>(forOp.getLoc (), bumpMap, iv);
1188
1201
},
1202
+ /* annotateFn=*/ annotateFn,
1189
1203
/* iterArgs=*/ iterArgs, /* yieldedValues=*/ yieldedValues);
1190
1204
1191
1205
// Promote the loop body up if this has turned into a single iteration loop.
@@ -1194,8 +1208,9 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
1194
1208
}
1195
1209
1196
1210
// / 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) {
1199
1214
assert (unrollFactor > 0 && " expected positive unroll factor" );
1200
1215
if (unrollFactor == 1 )
1201
1216
return promoteIfSingleIteration (forOp);
@@ -1300,7 +1315,7 @@ LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
1300
1315
b.create <MulIOp>(loc, step, b.create <ConstantIndexOp>(loc, i));
1301
1316
return b.create <AddIOp>(loc, iv, stride);
1302
1317
},
1303
- iterArgs, yieldedValues);
1318
+ annotateFn, iterArgs, yieldedValues);
1304
1319
// Promote the loop body up if this has turned into a single iteration loop.
1305
1320
(void )promoteIfSingleIteration (forOp);
1306
1321
return success ();
0 commit comments