[Mlir-commits] [mlir] [MLIR][LLVM] Add Continuous Loop Peeling transform to SCF (PR #71555)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Nov 15 02:52:17 PST 2023
================
@@ -105,6 +106,168 @@ static void specializeForLoopForUnrolling(ForOp op) {
op.erase();
}
+static LogicalResult splitLoopHelper(RewriterBase &b, scf::ForOp &forOp,
+ scf::ForOp &partialIteration,
+ Value &splitBound) {
+ RewriterBase::InsertionGuard guard(b);
+ auto lbInt = getConstantIntValue(forOp.getLowerBound());
+ auto ubInt = getConstantIntValue(forOp.getUpperBound());
+ auto stepInt = getConstantIntValue(forOp.getStep());
+
+ // No specialization necessary if step already divides upper bound evenly.
+ if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) % *stepInt == 0)
+ return failure();
+ // No specialization necessary if step size is 1.
+ if (stepInt == static_cast<int64_t>(1))
+ return failure();
+
+ // Create ForOp for partial iteration.
+ b.setInsertionPointAfter(forOp);
+ partialIteration = cast<scf::ForOp>(b.clone(*forOp.getOperation()));
+ partialIteration.getLowerBoundMutable().assign(splitBound);
+ forOp.replaceAllUsesWith(partialIteration->getResults());
+ partialIteration.getInitArgsMutable().assign(forOp->getResults());
+
+ // Set new upper loop bound.
+ b.updateRootInPlace(
+ forOp, [&]() { forOp.getUpperBoundMutable().assign(splitBound); });
+
+ return success();
+}
+
+static scf::IfOp convertSingleIterFor(RewriterBase &b, scf::ForOp &forOp) {
+ Location loc = forOp->getLoc();
+ IRMapping mapping;
+ mapping.map(forOp.getInductionVar(), forOp.getLowerBound());
+ for (auto [arg, operand] :
+ llvm::zip(forOp.getRegionIterArgs(), forOp.getInitsMutable())) {
+ mapping.map(arg, operand.get());
+ }
+ b.setInsertionPoint(forOp);
+ auto cond =
+ b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
+ forOp.getLowerBound(), forOp.getUpperBound());
+ auto ifOp = b.create<scf::IfOp>(loc, forOp->getResultTypes(), cond, true);
+ // then branch
+ b.setInsertionPointToStart(ifOp.thenBlock());
+ for (Operation &op : forOp.getBody()->getOperations()) {
+ b.clone(op, mapping);
+ }
+ // else branch
+ b.setInsertionPointToStart(ifOp.elseBlock());
+ if (!forOp->getResultTypes().empty()) {
+ b.create<scf::YieldOp>(loc, forOp.getInits());
+ }
+ b.replaceOp(forOp, ifOp->getResults());
+ return ifOp;
+}
+
+/// Rewrite a for loop with bounds/step that potentially do not divide the
+/// iteration space evenly into a chain of for loops where the step is a
+/// power of 2 and decreases exponentially across subsequent loops. Helps
+/// divide the iteration space across all resulting peeled loops evenly.
+///
+/// Optionally, convert all single iteration for loops to if-else
+/// blocks when convert_single_iter_loops_to_if attribute is set to true or
+/// alternatively with the convert-single-iter-loops-to-if option for the
+/// scf-for-loop-continuous-peeling pass.
+static LogicalResult continuousPeelForLoop(RewriterBase &b, ForOp forOp,
+ ForOp &partialIteration,
+ bool convertSingleIterLoopsToIf) {
+
+ scf::ForOp currentLoop;
+ auto lbInt = getConstantIntValue(forOp.getLowerBound());
+ auto stepInt = getConstantIntValue(forOp.getStep());
+
+ // Step size must be a known positive constant greater than 1.
+ if (stepInt && stepInt <= static_cast<int64_t>(1))
+ return failure();
+
+ Value initialUb = forOp.getUpperBound();
+ Value initialStep = forOp.getStep();
+ uint64_t loopStep = *stepInt;
+ currentLoop = forOp;
+ AffineExpr sym0, sym1, sym2;
+ bindSymbols(b.getContext(), sym0, sym1, sym2);
+ AffineMap defaultSplitMap =
+ AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2)});
+ AffineMap powerSplitMap = AffineMap::get(0, 3, {sym1 - (sym1 % sym2)});
+ bool usePowerSplit = (lbInt.has_value()) &&
+ (*lbInt % *stepInt == static_cast<int64_t>(0)) &&
+ (loopStep == llvm::bit_floor(loopStep));
+ AffineMap splitMap = usePowerSplit ? powerSplitMap : defaultSplitMap;
+ SmallVector<scf::ForOp> loops;
+ while (loopStep) {
+ b.setInsertionPoint(currentLoop);
+ auto constStepOp =
+ b.create<arith::ConstantIndexOp>(currentLoop.getLoc(), loopStep);
+ currentLoop.getStepMutable().assign(constStepOp);
+ b.setInsertionPoint(currentLoop);
+ Value splitBound = b.createOrFold<affine::AffineApplyOp>(
+ currentLoop.getLoc(), splitMap,
+ ValueRange{currentLoop.getLowerBound(), currentLoop.getUpperBound(),
+ currentLoop.getStep()});
+ LogicalResult status =
+ splitLoopHelper(b, currentLoop, partialIteration, splitBound);
+
+ // Canonicalize min/max affine operations
+ // It uses scf::rewritePeeledMinMaxOp to identify operations to be replaced,
+ // they are then replaced by the current step size.
+ // TODO: Alternative method - update affine map to reflect the loop step
+ // Example: min(ub - iv, 8) -> min(ub - iv, 4)
+ currentLoop.walk([&](affine::AffineMinOp affineOp) {
+ b.setInsertionPoint(affineOp);
+ auto clonedOp = cast<affine::AffineMinOp>(b.clone(*affineOp));
+ LogicalResult result = scf::rewritePeeledMinMaxOp(
+ b, clonedOp, currentLoop.getInductionVar(), initialUb, initialStep,
+ /*insideLoop=*/true);
+ if (result.succeeded())
+ b.replaceOp(affineOp, currentLoop.getStep());
+ else
+ b.eraseOp(clonedOp); // to avoid infinite walk
+ });
+ currentLoop.walk([&](affine::AffineMaxOp affineOp) {
+ b.setInsertionPoint(affineOp);
+ auto clonedOp = cast<affine::AffineMaxOp>(b.clone(*affineOp));
+ LogicalResult result = scf::rewritePeeledMinMaxOp(
+ b, clonedOp, currentLoop.getInductionVar(), initialUb, initialStep,
+ /*insideLoop=*/true);
+ if (result.succeeded())
+ b.replaceOp(affineOp, currentLoop.getStep());
----------------
muneebkhan85 wrote:
The step needs to be corrected once the rewritePeeledMinMaxOp has rewritten the affineOp. If not corrected, the old step from the original loop is propagated which is not correct in case of continuously peeling. This is what happens in the case of using rewriteAffineOpAfterPeeling. The generated code for the test case looks as shown below (see commented lines to spot the problem). Also, affineOp if used in scf::rewritePeeledMinMaxOp is not accessible after the function returns successfully, as the operation is rewritten by canonicalizeMinMaxOp. That is the reason for applying the rewriting on the clonedOp and then correcting the step using affineOp. Couldn't think if a better way of doing this.
```
#map = affine_map<()[s0, s1, s2] -> (s1 - s1 mod s2)>
#map1 = affine_map<(d0, d1)[s0] -> (8, d0 - d1)> // step size 8 propogated as a constant. Incorrect assumption
// for continuous loop peeling
module {
func.func @foo(%arg0: index) -> index {
%c8 = arith.constant 8 : index
%c4 = arith.constant 4 : index
%c2 = arith.constant 2 : index
%c0 = arith.constant 0 : index
%0 = affine.apply #map()[%c0, %arg0, %c8]
%1 = scf.for %arg1 = %c0 to %0 step %c8 iter_args(%arg2 = %c0) -> (index) {
%10 = affine.min #map1(%arg0, %arg1)[%c8]
%11 = index.add %10, %arg2
scf.yield %11 : index
}
%2 = affine.apply #map()[%0, %arg0, %c4]
%3 = arith.cmpi slt, %0, %2 : index
%4 = scf.if %3 -> (index) {
%10 = affine.min #map1(%arg0, %0)[%c8] // step size not corrected!
%11 = index.add %10, %1
scf.yield %11 : index
} else {
scf.yield %1 : index
}
%5 = affine.apply #map()[%2, %arg0, %c2]
%6 = arith.cmpi slt, %2, %5 : index
%7 = scf.if %6 -> (index) {
%10 = affine.min #map1(%arg0, %2)[%c8] // step size not corrected!
%11 = index.add %10, %4
scf.yield %11 : index
} else {
scf.yield %4 : index
}
%8 = arith.cmpi slt, %5, %arg0 : index
%9 = scf.if %8 -> (index) {
%10 = affine.min #map1(%arg0, %5)[%c8] // step size not corrected!
%11 = index.add %10, %7
scf.yield %11 : index
} else {
scf.yield %7 : index
}
return %9 : index
}
}
```
https://github.com/llvm/llvm-project/pull/71555
More information about the Mlir-commits
mailing list