[mlir] Extract forall_to_for logic into reusable function and add pass #89636
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write If you have received no comments on your PR for a week, you can request a review If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-scf Author: Jorn Tuyls (jtuyls) ChangesThis PR extracts the existing cc @ftynse @matthias-springer @MaheshRavishankar Could you help with reviewing this PR? Full diff: https://github.com/llvm/llvm-project/pull/89636.diff 7 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h b/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h
index 90b315e83a8cfd..31c3d0eb629d28 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/Passes.h
@@ -59,6 +59,9 @@ createParallelLoopTilingPass(llvm::ArrayRef<int64_t> tileSize = {},
/// loop range.
std::unique_ptr<Pass> createForLoopRangeFoldingPass();
+/// Creates a pass that converts SCF forall loops to SCF for loops.
+std::unique_ptr<Pass> createForallToForLoopPass();
+
// Creates a pass which lowers for loops into while loops.
std::unique_ptr<Pass> createForToWhileLoopPass();
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Passes.td b/mlir/include/mlir/Dialect/SCF/Transforms/Passes.td
index 350611ad86873d..a7aeb42d60c0e9 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/Passes.td
@@ -120,6 +120,11 @@ def SCFForLoopRangeFolding : Pass<"scf-for-loop-range-folding"> {
let constructor = "mlir::createForLoopRangeFoldingPass()";
}
+def SCFForallToForLoop : Pass<"scf-forall-to-for"> {
+ let summary = "Convert SCF forall loops to SCF for loops";
+ let constructor = "mlir::createForallToForLoopPass()";
+}
+
def SCFForToWhileLoop : Pass<"scf-for-to-while"> {
let summary = "Convert SCF for loops to SCF while loops";
let constructor = "mlir::createForToWhileLoopPass()";
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
index 220dcb35571d27..9a9731161ddf7d 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h
@@ -28,10 +28,15 @@ class Value;
namespace scf {
class IfOp;
+class ForallOp;
class ForOp;
class ParallelOp;
class WhileOp;
+/// Try converting scf.forall into a set of nested scf.for loops.
+LogicalResult forallToForLoop(RewriterBase &rewriter, ForallOp forallOp,
+ SmallVector<Operation *> *results);
+
/// Fuses all adjacent scf.parallel operations with identical bounds and step
/// into one scf.parallel operations. Uses a naive aliasing and dependency
/// analysis.
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 7e4faf8b73afbb..0e3bc8ad4cacee 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -69,16 +69,7 @@ transform::ForallToForOp::apply(transform::TransformRewriter &rewriter,
return diag;
}
- rewriter.setInsertionPoint(target);
-
- if (!target.getOutputs().empty()) {
- return emitSilenceableError()
- << "unsupported shared outputs (didn't bufferize?)";
- }
-
SmallVector<OpFoldResult> lbs = target.getMixedLowerBound();
- SmallVector<OpFoldResult> ubs = target.getMixedUpperBound();
- SmallVector<OpFoldResult> steps = target.getMixedStep();
if (getNumResults() != lbs.size()) {
DiagnosedSilenceableFailure diag =
@@ -89,28 +80,15 @@ transform::ForallToForOp::apply(transform::TransformRewriter &rewriter,
return diag;
}
- auto loc = target.getLoc();
- SmallVector<Value> ivs;
- for (auto &&[lb, ub, step] : llvm::zip(lbs, ubs, steps)) {
- Value lbValue = getValueOrCreateConstantIndexOp(rewriter, loc, lb);
- Value ubValue = getValueOrCreateConstantIndexOp(rewriter, loc, ub);
- Value stepValue = getValueOrCreateConstantIndexOp(rewriter, loc, step);
- auto loop = rewriter.create<scf::ForOp>(
- loc, lbValue, ubValue, stepValue, ValueRange(),
- [](OpBuilder &, Location, Value, ValueRange) {});
- ivs.push_back(loop.getInductionVar());
- rewriter.setInsertionPointToStart(loop.getBody());
- rewriter.create<scf::YieldOp>(loc);
- rewriter.setInsertionPointToStart(loop.getBody());
+ SmallVector<Operation *> opResults;
+ if (failed(scf::forallToForLoop(rewriter, target, &opResults))) {
+ DiagnosedSilenceableFailure diag = emitSilenceableError()
+ << "failed to convert forall into for";
+ return diag;
}
- rewriter.eraseOp(target.getBody()->getTerminator());
- rewriter.inlineBlockBefore(target.getBody(), &*rewriter.getInsertionPoint(),
- ivs);
- rewriter.eraseOp(target);
-
- for (auto &&[i, iv] : llvm::enumerate(ivs)) {
- results.set(cast<OpResult>(getTransformed()[i]),
- {iv.getParentBlock()->getParentOp()});
+
+ for (auto [i, res] : llvm::enumerate(opResults)) {
+ results.set(cast<OpResult>(getTransformed()[i]), {res});
}
return DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
index a2925aef17ca78..e7671c9cc28f8b 100644
--- a/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRSCFTransforms
BufferDeallocationOpInterfaceImpl.cpp
BufferizableOpInterfaceImpl.cpp
Bufferize.cpp
+ ForallToFor.cpp
ForToWhile.cpp
LoopCanonicalization.cpp
LoopPipelining.cpp
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
new file mode 100644
index 00000000000000..2a4e97564fbe2c
--- /dev/null
+++ b/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
@@ -0,0 +1,92 @@
+//===- ForallToFor.cpp - scf.forall to scf.for loop conversion ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Transforms SCF.ForallOp's into SCF.ForOp's.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/Transforms/Passes.h"
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_SCFFORALLTOFORLOOP
+#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace llvm;
+using namespace mlir;
+using scf::ForallOp;
+using scf::ForOp;
+
+LogicalResult
+mlir::scf::forallToForLoop(RewriterBase &rewriter, scf::ForallOp forallOp,
+ SmallVector<Operation *> *results = nullptr) {
+ rewriter.setInsertionPoint(forallOp);
+
+ if (!forallOp.getOutputs().empty()) {
+ return forallOp.emitOpError()
+ << "unsupported shared outputs (didn't bufferize?)";
+ }
+
+ SmallVector<OpFoldResult> lbs = forallOp.getMixedLowerBound();
+ SmallVector<OpFoldResult> ubs = forallOp.getMixedUpperBound();
+ SmallVector<OpFoldResult> steps = forallOp.getMixedStep();
+
+ auto loc = forallOp.getLoc();
+ SmallVector<Value> ivs;
+ for (auto &&[lb, ub, step] : llvm::zip(lbs, ubs, steps)) {
+ Value lbValue = getValueOrCreateConstantIndexOp(rewriter, loc, lb);
+ Value ubValue = getValueOrCreateConstantIndexOp(rewriter, loc, ub);
+ Value stepValue = getValueOrCreateConstantIndexOp(rewriter, loc, step);
+ auto loop =
+ rewriter.create<ForOp>(loc, lbValue, ubValue, stepValue, ValueRange(),
+ [](OpBuilder &, Location, Value, ValueRange) {});
+ if (results)
+ results->push_back(loop);
+ ivs.push_back(loop.getInductionVar());
+ rewriter.setInsertionPointToStart(loop.getBody());
+ rewriter.create<scf::YieldOp>(loc);
+ rewriter.setInsertionPointToStart(loop.getBody());
+ }
+ rewriter.eraseOp(forallOp.getBody()->getTerminator());
+ rewriter.inlineBlockBefore(forallOp.getBody(), &*rewriter.getInsertionPoint(),
+ ivs);
+ rewriter.eraseOp(forallOp);
+ return success();
+}
+
+namespace {
+struct ForallToForLoopLoweringPattern : public OpRewritePattern<ForallOp> {
+ using OpRewritePattern<ForallOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ForallOp forallOp,
+ PatternRewriter &rewriter) const override {
+ if (failed(scf::forallToForLoop(rewriter, forallOp)))
+ return failure();
+ return success();
+ }
+};
+
+struct ForallToForLoop : public impl::SCFForallToForLoopBase<ForallToForLoop> {
+ void runOnOperation() override {
+ auto *parentOp = getOperation();
+ MLIRContext *ctx = parentOp->getContext();
+ RewritePatternSet patterns(ctx);
+ patterns.add<ForallToForLoopLoweringPattern>(ctx);
+ (void)applyPatternsAndFoldGreedily(parentOp, std::move(patterns));
+ }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createForallToForLoopPass() {
+ return std::make_unique<ForallToForLoop>();
+}
diff --git a/mlir/test/Dialect/SCF/forall-to-for.mlir b/mlir/test/Dialect/SCF/forall-to-for.mlir
new file mode 100644
index 00000000000000..e7d183fb9d2b54
--- /dev/null
+++ b/mlir/test/Dialect/SCF/forall-to-for.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(scf-forall-to-for))' -split-input-file | FileCheck %s
+
+func.func private @callee(%i: index, %j: index)
+
+// CHECK-LABEL: @two_iters
+// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index
+func.func @two_iters(%ub1: index, %ub2: index) {
+ scf.forall (%i, %j) in (%ub1, %ub2) {
+ func.call @callee(%i, %j) : (index, index) -> ()
+ }
+ // CHECK: scf.for %[[IV1:.+]] = %{{.*}} to %[[UB1]]
+ // CHECK: scf.for %[[IV2:.+]] = %{{.*}} to %[[UB2]]
+ // CHECK: func.call @callee(%[[IV1]], %[[IV2]])
+ return
+}
+
+// -----
+
+func.func private @callee(%i: index, %j: index)
+
+// CHECK-LABEL: @repeated
+// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index
+func.func @repeated(%ub1: index, %ub2: index) {
+ scf.forall (%i, %j) in (%ub1, %ub2) {
+ func.call @callee(%i, %j) : (index, index) -> ()
+ }
+ // CHECK: scf.for %[[IV1:.+]] = %{{.*}} to %[[UB1]]
+ // CHECK: scf.for %[[IV2:.+]] = %{{.*}} to %[[UB2]]
+ // CHECK: func.call @callee(%[[IV1]], %[[IV2]])
+ scf.forall (%i, %j) in (%ub1, %ub2) {
+ func.call @callee(%i, %j) : (index, index) -> ()
+ }
+ // CHECK: scf.for %[[IV1:.+]] = %{{.*}} to %[[UB1]]
+ // CHECK: scf.for %[[IV2:.+]] = %{{.*}} to %[[UB2]]
+ // CHECK: func.call @callee(%[[IV1]], %[[IV2]])
+ return
+}
+
+// -----
+
+func.func private @callee(%i: index, %j: index, %k: index, %l: index)
+
+// CHECK-LABEL: @nested
+// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index, %[[UB3:.+]]: index, %[[UB4:.+]]: index
+func.func @nested(%ub1: index, %ub2: index, %ub3: index, %ub4: index) {
+ // CHECK: scf.for %[[IV1:.+]] = %{{.*}} to %[[UB1]]
+ // CHECK: scf.for %[[IV2:.+]] = %{{.*}} to %[[UB2]]
+ // CHECK: scf.for %[[IV3:.+]] = %{{.*}} to %[[UB3]]
+ // CHECK: scf.for %[[IV4:.+]] = %{{.*}} to %[[UB4]]
+ // CHECK: func.call @callee(%[[IV1]], %[[IV2]], %[[IV3]], %[[IV4]])
+ scf.forall (%i, %j) in (%ub1, %ub2) {
+ scf.forall (%k, %l) in (%ub3, %ub4) {
+ func.call @callee(%i, %j, %k, %l) : (index, index, index, index) -> ()
+ }
+ }
+ return
+}
|
MaheshRavishankar
requested changes
Apr 23, 2024
jtuyls
force-pushed
the
scf-forall-to-for
branch
from
f506d72 to
45f5065
Compare
ftynse
requested changes
Apr 23, 2024
| << "unsupported shared outputs (didn't bufferize?)"; | ||
| } | ||
|
|
||
| auto loc = forallOp.getLoc(); |
There was a problem hiding this comment.
Please expand auto unless the type is obvious from context. Here and below.
https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable
Comment on lines
50
to
52
| if (results) { | ||
| llvm::copy(loopNest.loops, std::back_inserter(*results)); | ||
| } |
There was a problem hiding this comment.
Suggested change
| if (results) { | |
| llvm::copy(loopNest.loops, std::back_inserter(*results)); | |
| } | |
| if (results) | |
| llvm::append_range(*results, loopNest.loops); |
There was a problem hiding this comment.
Also, if you postpone this until the end of the function, it becomes possible to move the vector elements.
jtuyls
force-pushed
the
scf-forall-to-for
branch
from
45f5065 to
34da233
Compare
jtuyls
force-pushed
the
scf-forall-to-for
branch
from
34da233 to
b467dae
Compare
ftynse
approved these changes
Apr 24, 2024
MaheshRavishankar
approved these changes
Apr 24, 2024
|
@jtuyls Congratulations on having your first Pull Request (PR) merged into the LLVM Project! Your changes will be combined with recent changes from other authors, then tested Please check whether problems have been caused by your change specifically, as How to do this, and the rest of the post-merge process, is covered in detail here. If your change does cause a problem, it may be reverted, or you can revert it yourself. If you don't get any reports, no action is required from you. Your changes are working as expected, well done! |
jtuyls
added a commit
to nod-ai/iree-amd-aie
that referenced
this pull request
Apr 26, 2024
Bump to latest IREE. Update for deprecated `mlir::Value::isa/dyn_cast/cast/` member functions: llvm/llvm-project#89238. See also: https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443. This brings in a standalone `forallToFor` transformation which will avoid duplication in this repo: llvm/llvm-project#89636
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR extracts the existing
scf.foralltoscf.forconversion logic inside a transform op (#65474) into a standalone function which can be used in other transformations and adds ascf-forall-to-forpass.cc @ftynse @matthias-springer @MaheshRavishankar Could you help with reviewing this PR?