[Mlir-commits] [mlir] [MLIR][LLVM] Add Continuous Loop Peeling transform to SCF (PR #71555)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Nov 7 08:11:49 PST 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-scf
Author: None (muneebkhan85)
<details>
<summary>Changes</summary>
This patch adds continuous loop peeling to scf loop transforms in the MLIR backend. This transforms the target loop into a chain of loops, with step sizes that are powers of two and decrease exponentially across subsequent loops.
The transform is similar to loop.peel in the effect that it creates a loop with a step (that is power of 2) to divide the range evenly, with the difference that the remaining iterations are spread across similar loops with exponentially decreasing step sizes, with the last loop with step size of 2^0 = 1.
Originally authored by Litu Zhou litu.zhou@<!-- -->huawei.com.
---
Full diff: https://github.com/llvm/llvm-project/pull/71555.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td (+36)
- (modified) mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp (+147)
- (added) mlir/test/Dialect/SCF/loop-continuous-peel.mlir (+98)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index 14df7e23a430fb1..e3d79a7f0ae40f3 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -147,6 +147,42 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
}];
}
+def LoopContinuousPeelOp : Op<Transform_Dialect, "loop.loop_continuous_peel",
+ [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+ TransformOpInterface, TransformEachOpTrait]> {
+ let description = [{
+ Transforms the loop into a chain of loops, with step sizes that are
+ powers of two and decrease exponetially across subsequent loops.
+ The transform is similar to loop.peel in the effect that it creates a loop
+ with a step (that is power of 2) to divide the range evenly, with the
+ difference that the remaining iterations are spread across similar loops
+ with exponentially decreasing step sizes, with the last loop with step size
+ of 2^0 = 1.
+
+ #### Return modes
+
+ This operation consumes the `target` handles and produces the
+ continuously-peeled loop.
+ }];
+
+ let arguments =
+ (ins TransformHandleTypeInterface:$target,
+ DefaultValuedAttr<BoolAttr, "false">:$single_iter_opt);
+ // TODO: Return both the peeled loop and the remainder loop.
+ let results = (outs TransformHandleTypeInterface:$transformed);
+
+ let assemblyFormat =
+ "$target attr-dict `:` functional-type(operands, results)";
+
+ let extraClassDeclaration = [{
+ ::mlir::DiagnosedSilenceableFailure applyToOne(
+ ::mlir::transform::TransformRewriter &rewriter,
+ ::mlir::Operation *target,
+ ::mlir::transform::ApplyToEachResultList &results,
+ ::mlir::transform::TransformState &state);
+ }];
+}
+
def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
TransformOpInterface, TransformEachOpTrait]> {
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 62370604142cd5b..dcba6a8b406b21f 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -206,6 +206,153 @@ transform::LoopPeelOp::applyToOne(transform::TransformRewriter &rewriter,
return DiagnosedSilenceableFailure::success();
}
+//===---------------------------------------------------------------------===//
+// LoopContinuousPeelOp
+//===---------------------------------------------------------------------===//
+
+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;
+}
+
+DiagnosedSilenceableFailure transform::LoopContinuousPeelOp::applyToOne(
+ transform::TransformRewriter &rewriter, Operation *target,
+ transform::ApplyToEachResultList &results,
+ transform::TransformState &state) {
+ scf::ForOp loop, currentLoop, partialLoop;
+ loop = dyn_cast<scf::ForOp>(target);
+ auto lbInt = getConstantIntValue(loop.getLowerBound());
+ auto stepInt = getConstantIntValue(loop.getStep());
+ if (!stepInt.has_value() || *stepInt <= 0)
+ return DiagnosedSilenceableFailure::
+ definiteFailure(); // step size must be a known positive constant
+ Value initialUb = loop.getUpperBound();
+ Value initialStep = loop.getStep();
+ uint64_t loopStep = *stepInt;
+ currentLoop = loop;
+ AffineExpr sym0, sym1, sym2;
+ bindSymbols(rewriter.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) {
+ rewriter.setInsertionPoint(currentLoop);
+ auto constStepOp =
+ rewriter.create<arith::ConstantIndexOp>(currentLoop.getLoc(), loopStep);
+ currentLoop.getStepMutable().assign(constStepOp);
+ rewriter.setInsertionPoint(currentLoop);
+ Value splitBound = rewriter.createOrFold<affine::AffineApplyOp>(
+ currentLoop.getLoc(), splitMap,
+ ValueRange{currentLoop.getLowerBound(), currentLoop.getUpperBound(),
+ currentLoop.getStep()});
+ LogicalResult status =
+ splitLoopHelper(rewriter, currentLoop, partialLoop, 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) {
+ rewriter.setInsertionPoint(affineOp);
+ auto clonedOp = cast<affine::AffineMinOp>(rewriter.clone(*affineOp));
+ LogicalResult result = scf::rewritePeeledMinMaxOp(
+ rewriter, clonedOp, currentLoop.getInductionVar(), initialUb,
+ initialStep,
+ /*insideLoop=*/true);
+ if (result.succeeded())
+ rewriter.replaceOp(affineOp, currentLoop.getStep());
+ else
+ rewriter.eraseOp(clonedOp); // to avoid infinite walk
+ });
+ currentLoop.walk([&](affine::AffineMaxOp affineOp) {
+ rewriter.setInsertionPoint(affineOp);
+ auto clonedOp = cast<affine::AffineMaxOp>(rewriter.clone(*affineOp));
+ LogicalResult result = scf::rewritePeeledMinMaxOp(
+ rewriter, clonedOp, currentLoop.getInductionVar(), initialUb,
+ initialStep,
+ /*insideLoop=*/true);
+ if (result.succeeded())
+ rewriter.replaceOp(affineOp, currentLoop.getStep());
+ else
+ rewriter.eraseOp(clonedOp); // to avoid infinite walk
+ });
+
+ // Prepare for the next iteration
+ loops.push_back(currentLoop);
+ if (failed(status))
+ break;
+ currentLoop = partialLoop;
+ uint64_t maxPower = llvm::bit_floor(loopStep);
+ loopStep = maxPower == loopStep ? maxPower >> 1 : maxPower;
+ }
+ assert(loops.size() > 0 && "There should be at least one loop available");
+ if (getSingleIterOpt()) {
+ for (size_t i = 1; i < loops.size(); ++i) {
+ convertSingleIterFor(rewriter, loops[i]);
+ }
+ }
+
+ results.push_back(loops.front());
+ return DiagnosedSilenceableFailure::success();
+}
+
//===----------------------------------------------------------------------===//
// LoopPipelineOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SCF/loop-continuous-peel.mlir b/mlir/test/Dialect/SCF/loop-continuous-peel.mlir
new file mode 100644
index 000000000000000..752e1b1efed92ac
--- /dev/null
+++ b/mlir/test/Dialect/SCF/loop-continuous-peel.mlir
@@ -0,0 +1,98 @@
+// RUN: mlir-opt %s --transform-interpreter -split-input-file | FileCheck %s
+
+#map = affine_map<(d0) -> ()>
+#map1 = affine_map<(d0) -> (d0)>
+module {
+ func.func @foo(%arg0: f32, %arg1: tensor<?xf32>) -> tensor<?xf32> {
+ %0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel"]} ins(%arg0 : f32) outs(%arg1 : tensor<?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %3 = arith.mulf %in, %out : f32
+ linalg.yield %3 : f32
+ } -> tensor<?xf32>
+ return %0 : tensor<?xf32>
+ }
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops = transform.structured.tile_using_for %0[8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %2 = transform.cast %loops : !transform.any_op to !transform.op<"scf.for">
+ %3 = transform.loop.loop_continuous_peel %2 {single_iter_opt = true} : (!transform.op<"scf.for">) -> (!transform.any_op)
+ transform.yield
+ }
+}
+
+// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1, s2] -> (s1 - s1 mod s2)>
+// CHECK: #[[MAP1:.*]] = affine_map<() -> (8)>
+// CHECK: #[[MAP2:.*]] = affine_map<(d0) -> (d0 - 1)>
+// CHECK: #[[MAP3:.*]] = affine_map<(d0) -> ()>
+// CHECK: #[[MAP4:.*]] = affine_map<(d0) -> (d0)>
+
+// CHECK: func.func @foo(%[[S:.*]]: f32, %[[INVEC1:.*]]: tensor<?xf32>) -> tensor<?xf32> {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM:.*]] = tensor.dim %[[INVEC1]], %[[C0]] : tensor<?xf32>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %{{.*}} = arith.constant 8 : index
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[IDX0:.*]] = affine.apply #[[MAP]]()[%[[C0]], %[[DIM]], %[[C8]]]
+// CHECK: %[[INS1:.*]] = scf.for %[[IDX:.*]] = %[[C0]] to %[[IDX0]] step %[[C8]] iter_args(%[[AINVEC1:.*]] = %[[INVEC1]]) -> (tensor<?xf32>) {
+// CHECK: %{{.*}} = affine.apply #[[MAP2]](%[[C8]])
+// CHECK: %[[XS8:.*]] = tensor.extract_slice %[[AINVEC1]][%[[IDX]]] [%[[C8]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK: %[[MUL:.*]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel"]} ins(%{{.*}} : f32) outs(%[[XS8]] : tensor<?xf32>) {
+// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
+// CHECK: %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : f32
+// CHECK: linalg.yield %{{.*}} : f32
+// CHECK: } -> tensor<?xf32>
+// CHECK: %[[INS:.*]] = tensor.insert_slice %[[MUL]] into %[[AINVEC1]][%[[IDX]]] [%[[C8]]] [1] : tensor<?xf32> into tensor<?xf32>
+// CHECK: scf.yield %[[INS]] : tensor<?xf32>
+// CHECK: }
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[IDX2:.*]] = affine.apply #[[MAP]]()[%[[IDX0]], %[[DIM]], %[[C4]]]
+// CHECK: %[[CMP3:.*]] = arith.cmpi slt, %[[IDX0]], %[[IDX2]] : index
+// CHECK: %[[INS2:.*]] = scf.if %[[CMP3]] -> (tensor<?xf32>) {
+// CHECK: %{{.*}} = affine.apply #[[MAP2]](%[[C4]])
+// CHECK: %[[XS4:.*]] = tensor.extract_slice %[[INS1]][%[[IDX0]]] [%[[C4]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK: %[[MUL:.*]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel"]} ins(%[[S]] : f32) outs(%[[XS4]] : tensor<?xf32>) {
+// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
+// CHECK: %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : f32
+// CHECK: linalg.yield %{{.*}} : f32
+// CHECK: } -> tensor<?xf32>
+// CHECK: %[[INS:.*]] = tensor.insert_slice %[[MUL]] into %[[INS1]][%[[IDX0]]] [%[[C4]]] [1] : tensor<?xf32> into tensor<?xf32>
+// CHECK: scf.yield %[[INS]] : tensor<?xf32>
+// CHECK: } else {
+// CHECK: scf.yield %[[INS1]] : tensor<?xf32>
+// CHECK: }
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[IDX3:.*]] = affine.apply #[[MAP]]()[%[[IDX2]], %[[DIM]], %[[C2]]]
+// CHECK: %[[CMP4:.*]] = arith.cmpi slt, %[[IDX2]], %[[IDX3]] : index
+// CHECK: %[[INS3:.*]] = scf.if %[[CMP4]] -> (tensor<?xf32>) {
+// CHECK: %{{.*}} = affine.apply #[[MAP2]](%[[C2]])
+// CHECK: %[[XS2:.*]] = tensor.extract_slice %[[INS2]][%[[IDX2]]] [%[[C2]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK: %[[MUL:.*]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel"]} ins(%[[S]] : f32) outs(%[[XS2]] : tensor<?xf32>) {
+// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
+// CHECK: %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : f32
+// CHECK: linalg.yield %{{.*}} : f32
+// CHECK: } -> tensor<?xf32>
+// CHECK: %[[INS:.*]] = tensor.insert_slice %[[MUL]] into %[[INS2]][%[[IDX2]]] [%[[C2]]] [1] : tensor<?xf32> into tensor<?xf32>
+// CHECK: scf.yield %[[INS]] : tensor<?xf32>
+// CHECK: } else {
+// CHECK: scf.yield %[[INS2]] : tensor<?xf32>
+// CHECK: }
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %{{.*}} = affine.apply #[[MAP]]()[%[[IDX3]], %[[DIM]], %[[C1]]]
+// CHECK: %[[CMP5:.*]] = arith.cmpi slt, %[[IDX3]], %[[DIM]] : index
+// CHECK: %[[INS4:.*]] = scf.if %[[CMP5]] -> (tensor<?xf32>) {
+// CHECK: %{{.*}} = affine.apply #[[MAP2]](%[[C1]])
+// CHECK: %[[XS1:.*]] = tensor.extract_slice %[[INS3]][%[[IDX3]]] [%[[C1]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK: %[[MUL:.*]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel"]} ins(%[[S]] : f32) outs(%[[XS1]] : tensor<?xf32>) {
+// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
+// CHECK: %{{.*}} = arith.mulf %{{.*}}, %{{.*}} : f32
+// CHECK: linalg.yield %{{.*}} : f32
+// CHECK: } -> tensor<?xf32>
+// CHECK: %[[INS:.*]] = tensor.insert_slice %[[MUL]] into %[[INS3]][%[[IDX3]]] [%[[C1]]] [1] : tensor<?xf32> into tensor<?xf32>
+// CHECK: scf.yield %[[INS]] : tensor<?xf32>
+// CHECK: } else {
+// CHECK: scf.yield %[[INS3]] : tensor<?xf32>
+// CHECK: }
+// CHECK: return %[[INS4]] : tensor<?xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/71555
More information about the Mlir-commits
mailing list