[Mlir-commits] [mlir] 9909ef2 - [mlir][scf] Fix a bug in scf::ForOp loop unroll with an epilogue
Uday Bondhugula
llvmlistbot at llvm.org
Sat Oct 10 01:49:13 PDT 2020
Author: Tatiana Shpeisman
Date: 2020-10-10T14:18:25+05:30
New Revision: 9909ef292daddccbd3b1154cec173014b847880a
URL: https://github.com/llvm/llvm-project/commit/9909ef292daddccbd3b1154cec173014b847880a
DIFF: https://github.com/llvm/llvm-project/commit/9909ef292daddccbd3b1154cec173014b847880a.diff
LOG: [mlir][scf] Fix a bug in scf::ForOp loop unroll with an epilogue
Fixes a bug in formation and simplification of an epilogue loop generated
during loop unroll of scf::ForOp (https://bugs.llvm.org/show_bug.cgi?id=46689)
Differential Revision: https://reviews.llvm.org/D87583
Added:
mlir/test/Transforms/scf-loop-unroll.mlir
Modified:
mlir/lib/Transforms/Utils/LoopUtils.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index cf79e267fb8a..8f4b0f800c05 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -206,10 +206,22 @@ LogicalResult mlir::promoteIfSingleIteration(scf::ForOp forOp) {
auto iv = forOp.getInductionVar();
iv.replaceAllUsesWith(lbCstOp);
+ // Replace uses of iterArgs with iterOperands.
+ auto iterOperands = forOp.getIterOperands();
+ auto iterArgs = forOp.getRegionIterArgs();
+ for (auto e : llvm::zip(iterOperands, iterArgs))
+ std::get<1>(e).replaceAllUsesWith(std::get<0>(e));
+
+ // Replace uses of loop results with the values yielded by the loop.
+ auto outerResults = forOp.getResults();
+ auto innerResults = forOp.getBody()->getTerminator()->getOperands();
+ for (auto e : llvm::zip(outerResults, innerResults))
+ std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
+
// Move the loop body operations, except for its terminator, to the loop's
// containing block.
auto *parentBlock = forOp.getOperation()->getBlock();
- forOp.getBody()->back().erase();
+ forOp.getBody()->getTerminator()->erase();
parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations());
forOp.erase();
@@ -1038,12 +1050,13 @@ LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
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) {
+/// 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) {
// Builder to insert unrolled bodies just before the terminator of the body of
// 'forOp'.
auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
@@ -1053,9 +1066,14 @@ static void generateUnrolledLoop(
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()) {
@@ -1066,7 +1084,14 @@ static void generateUnrolledLoop(
// Clone the original body of 'forOp'.
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
builder.clone(*it, operandMap);
+
+ // Update yielded values.
+ for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
+ lastYielded[i] = operandMap.lookup(yieldedValues[i]);
}
+
+ // Update operands of the yield statement.
+ loopBodyBlock->getTerminator()->setOperands(lastYielded);
}
/// Unrolls this loop by the specified factor. Returns success if the loop
@@ -1127,7 +1152,8 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
return b.create<AffineApplyOp>(forOp.getLoc(), bumpMap,
iv);
- });
+ },
+ /*iterArgs=*/{}, /*yieldedValues=*/{});
// Promote the loop body up if this has turned into a single iteration loop.
promoteIfSingleIteration(forOp);
@@ -1212,19 +1238,36 @@ LogicalResult mlir::loopUnrollByFactor(scf::ForOp forOp,
std::next(Block::iterator(forOp)));
auto epilogueForOp = cast<scf::ForOp>(epilogueBuilder.clone(*forOp));
epilogueForOp.setLowerBound(upperBoundUnrolled);
+
+ // Update uses of loop results.
+ auto results = forOp.getResults();
+ auto epilogueResults = epilogueForOp.getResults();
+ auto epilogueIterOperands = epilogueForOp.getIterOperands();
+
+ for (auto e : llvm::zip(results, epilogueResults, epilogueIterOperands)) {
+ std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
+ epilogueForOp.getOperation()->replaceUsesOfWith(std::get<2>(e),
+ std::get<0>(e));
+ }
promoteIfSingleIteration(epilogueForOp);
}
// Create unrolled loop.
forOp.setUpperBound(upperBoundUnrolled);
forOp.setStep(stepUnrolled);
- 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);
- });
+
+ auto iterArgs = ValueRange(forOp.getRegionIterArgs());
+ auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
+
+ 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);
// Promote the loop body up if this has turned into a single iteration loop.
promoteIfSingleIteration(forOp);
return success();
diff --git a/mlir/test/Transforms/scf-loop-unroll.mlir b/mlir/test/Transforms/scf-loop-unroll.mlir
new file mode 100644
index 000000000000..3698185ad4cc
--- /dev/null
+++ b/mlir/test/Transforms/scf-loop-unroll.mlir
@@ -0,0 +1,44 @@
+// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=3" -split-input-file -canonicalize | FileCheck %s
+
+// CHECK-LABEL: scf_loop_unroll_single
+func @scf_loop_unroll_single(%arg0 : f32, %arg1 : f32) -> f32 {
+ %from = constant 0 : index
+ %to = constant 10 : index
+ %step = constant 1 : index
+ %sum = scf.for %iv = %from to %to step %step iter_args(%sum_iter = %arg0) -> (f32) {
+ %next = addf %sum_iter, %arg1 : f32
+ scf.yield %next : f32
+ }
+ // CHECK: %[[SUM:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[V0:.*]] =
+ // CHECK-NEXT: %[[V1:.*]] = addf %[[V0]]
+ // CHECK-NEXT: %[[V2:.*]] = addf %[[V1]]
+ // CHECK-NEXT: %[[V3:.*]] = addf %[[V2]]
+ // CHECK-NEXT: scf.yield %[[V3]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: %[[RES:.*]] = addf %[[SUM]],
+ // CHECK-NEXT: return %[[RES]]
+ return %sum : f32
+}
+
+// CHECK-LABEL: scf_loop_unroll_double_symbolic_ub
+// CHECK-SAME: (%{{.*}}: f32, %{{.*}}: f32, %[[N:.*]]: index)
+func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : index) -> (f32,f32) {
+ %from = constant 0 : index
+ %step = constant 1 : index
+ %sum:2 = scf.for %iv = %from to %n step %step iter_args(%i0 = %arg0, %i1 = %arg1) -> (f32, f32) {
+ %sum0 = addf %i0, %arg0 : f32
+ %sum1 = addf %i1, %arg1 : f32
+ scf.yield %sum0, %sum1 : f32, f32
+ }
+ return %sum#0, %sum#1 : f32, f32
+ // CHECK: %[[C0:.*]] = constant 0 : index
+ // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+ // CHECK-NEXT: %[[C3:.*]] = constant 3 : index
+ // CHECK-NEXT: %[[REM:.*]] = remi_signed %[[N]], %[[C3]]
+ // CHECK-NEXT: %[[UB:.*]] = subi %[[N]], %[[REM]]
+ // CHECK-NEXT: %[[SUM:.*]]:2 = scf.for {{.*}} = %[[C0]] to %[[UB]] step %[[C3]] iter_args
+ // CHECK: }
+ // CHECK-NEXT: %[[SUM1:.*]]:2 = scf.for {{.*}} = %[[UB]] to %[[N]] step %[[C1]] iter_args(%[[V1:.*]] = %[[SUM]]#0, %[[V2:.*]] = %[[SUM]]#1)
+ // CHECK: }
+ // CHECK-NEXT: return %[[SUM1]]#0, %[[SUM1]]#1
+}
More information about the Mlir-commits
mailing list