-
Notifications
You must be signed in to change notification settings - Fork 14.3k
[mlir][Interfaces] LoopLikeOpInterface
: Add helper to get yielded values
#67305
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[mlir][Interfaces] LoopLikeOpInterface
: Add helper to get yielded values
#67305
Conversation
@llvm/pr-subscribers-flang-fir-hlfir @llvm/pr-subscribers-mlir-linalg ChangesAdd a new interface method that returns the yielded values. Also add a verifier that checks the number of inits/iter_args/yielded values. Most of the checked invariants (but not all of them) are already covered by the Depends on #67121. Only review the top commit. Patch is 52.29 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/67305.diff 21 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 704c2704536d20a..56b4a609e62c001 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -490,20 +490,6 @@ void buildAffineLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
function_ref<void(OpBuilder &, Location, ValueRange)>
bodyBuilderFn = nullptr);
-/// Replace `loop` with a new loop where `newIterOperands` are appended with
-/// new initialization values and `newYieldedValues` are added as new yielded
-/// values. The returned ForOp has `newYieldedValues.size()` new result values.
-/// Additionally, if `replaceLoopResults` is true, all uses of
-/// `loop.getResults()` are replaced with the first `loop.getNumResults()`
-/// return values of the original loop respectively. The original loop is
-/// deleted and the new loop returned.
-/// Prerequisite: `newIterOperands.size() == newYieldedValues.size()`.
-AffineForOp replaceForOpWithNewYields(OpBuilder &b, AffineForOp loop,
- ValueRange newIterOperands,
- ValueRange newYieldedValues,
- ValueRange newIterArgs,
- bool replaceLoopResults = true);
-
/// AffineBound represents a lower or upper bound in the for operation.
/// This class does not own the underlying operands. Instead, it refers
/// to the operands stored in the AffineForOp. Its life span should not exceed
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 5a1baaf4e1611c8..6eb0a4468fd91b5 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -120,7 +120,8 @@ def AffineForOp : Affine_Op<"for",
[AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound"]>,
+ "getSingleUpperBound", "getYieldedValues",
+ "replaceWithAdditionalYields"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>]> {
let summary = "for operation";
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0c93989ca99a4eb..c63703327f5bb34 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -121,7 +121,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
def ForOp : SCF_Op<"for",
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInits", "getSingleInductionVar", "getSingleLowerBound",
- "getSingleStep", "getSingleUpperBound", "promoteIfSingleIteration"]>,
+ "getSingleStep", "getSingleUpperBound", "getYieldedValues",
+ "promoteIfSingleIteration", "replaceWithAdditionalYields"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
ConditionallySpeculatable,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
@@ -241,9 +242,11 @@ def ForOp : SCF_Op<"for",
function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
Value getInductionVar() { return getBody()->getArgument(0); }
+
Block::BlockArgListType getRegionIterArgs() {
return getBody()->getArguments().drop_front(getNumInductionVars());
}
+
/// Return the `index`-th region iteration argument.
BlockArgument getRegionIterArg(unsigned index) {
assert(index < getNumRegionIterArgs() &&
@@ -1080,6 +1083,11 @@ def WhileOp : SCF_Op<"while",
ConditionOp getConditionOp();
YieldOp getYieldOp();
+
+ /// Return the values that are yielded from the "before" region (by the
+ /// ConditionOp).
+ ValueRange getYieldedValues();
+
Block::BlockArgListType getBeforeArguments();
Block::BlockArgListType getAfterArguments();
Block *getBeforeBody() { return &getBefore().front(); }
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index bde30c9c3528dbc..9bdd6eb833876f0 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -34,39 +34,6 @@ class CallOp;
class FuncOp;
} // namespace func
-/// Replace the `loop` with `newIterOperands` added as new initialization
-/// values. `newYieldValuesFn` is a callback that can be used to specify
-/// the additional values to be yielded by the loop. The number of
-/// values returned by the callback should match the number of new
-/// initialization values. This function
-/// - Moves (i.e. doesnt clone) operations from the `loop` to the newly created
-/// loop
-/// - Replaces the uses of `loop` with the new loop.
-/// - `loop` isnt erased, but is left in a "no-op" state where the body of the
-/// loop just yields the basic block arguments that correspond to the
-/// initialization values of a loop. The loop is dead after this method.
-/// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
-/// `newIterOperands` within the generated new loop are replaced
-/// with the corresponding `BlockArgument` in the loop body.
-using NewYieldValueFn = std::function<SmallVector<Value>(
- OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs)>;
-scf::ForOp replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
- ValueRange newIterOperands,
- const NewYieldValueFn &newYieldValuesFn,
- bool replaceIterOperandsUsesInLoop = true);
-// Simpler API if the new yields are just a list of values that can be
-// determined ahead of time.
-inline scf::ForOp
-replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
- ValueRange newIterOperands, ValueRange newYields,
- bool replaceIterOperandsUsesInLoop = true) {
- auto fn = [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
- return SmallVector<Value>(newYields.begin(), newYields.end());
- };
- return replaceLoopWithNewYields(builder, loop, newIterOperands, fn,
- replaceIterOperandsUsesInLoop);
-}
-
/// Update a perfectly nested loop nest to yield new values from the innermost
/// loop and propagating it up through the loop nest. This function
/// - Expects `loopNest` to be a perfectly nested loop with outer most loop
@@ -82,11 +49,10 @@ replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
/// - If `replaceIterOperandsUsesInLoop` is true, all uses of the
/// `newIterOperands` within the generated new loop are replaced with the
/// corresponding `BlockArgument` in the loop body.
-SmallVector<scf::ForOp>
-replaceLoopNestWithNewYields(OpBuilder &builder, ArrayRef<scf::ForOp> loopNest,
- ValueRange newIterOperands,
- const NewYieldValueFn &newYieldValueFn,
- bool replaceIterOperandsUsesInLoop = true);
+SmallVector<scf::ForOp> replaceLoopNestWithNewYields(
+ RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
+ ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
+ bool replaceIterOperandsUsesInLoop = true);
/// Outline a region with a single block into a new FuncOp.
/// Assumes the FuncOp result types is the type of the yielded operands of the
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 9d81a61fac88566..7c7d378d0590ab1 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -17,6 +17,18 @@
namespace mlir {
class RewriterBase;
+
+/// A function that returns the additional yielded values during
+/// `replaceWithAdditionalYields`. `newBbArgs` are the newly added region
+/// iter_args. This function should return as many values as there are block
+/// arguments in `newBbArgs`.
+using NewYieldValuesFn = std::function<SmallVector<Value>(
+ OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs)>;
+
+namespace detail {
+/// Verify invariants of the LoopLikeOpInterface.
+LogicalResult verifyLoopLikeOpInterface(Operation *op);
+} // namespace detail
} // namespace mlir
/// Include the generated interface declarations.
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index cb6b2f4ed4ae8b5..23049669ec9348f 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -141,6 +141,42 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
return ::mlir::Block::BlockArgListType();
}]
>,
+ InterfaceMethod<[{
+ TODO
+ }],
+ /*retTy=*/"::mlir::ValueRange",
+ /*methodName=*/"getYieldedValues",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::ValueRange();
+ }]
+ >,
+ InterfaceMethod<[{
+ Append the specified additional "init" operands: replace this loop with
+ a new loop that has the additional init operands. The loop body of
+ this loop is moved over to the new loop.
+
+ `newInitOperands` specifies the additional "init" operands.
+ `newYieldValuesFn` is a function that returns the yielded values (which
+ can be computed based on the additional region iter_args). If
+ `replaceInitOperandUsesInLoop` is set, all uses of the additional init
+ operands inside of this loop are replaced with the corresponding, newly
+ added region iter_args.
+
+ Note: Loops that do not support init/iter_args should return "failure".
+ }],
+ /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+ /*methodName=*/"replaceWithAdditionalYields",
+ /*args=*/(ins "::mlir::RewriterBase &":$rewriter,
+ "::mlir::ValueRange":$newInitOperands,
+ "bool":$replaceInitOperandUsesInLoop,
+ "const ::mlir::NewYieldValuesFn &":$newYieldValuesFn),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::failure();
+ }]
+ >,
];
let extraClassDeclaration = [{
@@ -149,6 +185,30 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
/// because the control flow graph is cyclic
static bool blockIsInLoop(Block *block);
}];
+
+ let extraSharedClassDeclaration = [{
+ /// Append the specified additional "init" operands: replace this loop with
+ /// a new loop that has the additional init operands. The loop body of this
+ /// loop is moved over to the new loop.
+ ///
+ /// The newly added region iter_args are yielded from the loop.
+ ::mlir::FailureOr<::mlir::LoopLikeOpInterface>
+ replaceWithAdditionalIterOperands(::mlir::RewriterBase &rewriter,
+ ::mlir::ValueRange newInitOperands,
+ bool replaceInitOperandUsesInLoop) {
+ return $_op.replaceWithAdditionalYields(
+ rewriter, newInitOperands, replaceInitOperandUsesInLoop,
+ [](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
+ return SmallVector<Value>(newBBArgs);
+ });
+ }
+ }];
+
+ let verifyWithRegions = 1;
+
+ let verify = [{
+ return detail::verifyLoopLikeOpInterface($_op);
+ }];
}
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 9ecc568883a3b0b..4109167de44e860 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2193,6 +2193,10 @@ unsigned AffineForOp::getNumIterOperands() {
return getNumOperands() - lbMap.getNumInputs() - ubMap.getNumInputs();
}
+ValueRange AffineForOp::getYieldedValues() {
+ return cast<AffineYieldOp>(getBody()->getTerminator()).getOperands();
+}
+
void AffineForOp::print(OpAsmPrinter &p) {
p << ' ';
p.printRegionArgument(getBody()->getArgument(0), /*argAttrs=*/{},
@@ -2575,6 +2579,58 @@ std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
}
+FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
+ RewriterBase &rewriter, ValueRange newInitOperands,
+ bool replaceInitOperandUsesInLoop,
+ const NewYieldValuesFn &newYieldValuesFn) {
+ // Create a new loop before the existing one, with the extra operands.
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(getOperation());
+ auto inits = llvm::to_vector(getInits());
+ inits.append(newInitOperands.begin(), newInitOperands.end());
+ AffineForOp newLoop = rewriter.create<AffineForOp>(
+ getLoc(), getLowerBoundOperands(), getLowerBoundMap(),
+ getUpperBoundOperands(), getUpperBoundMap(), getStep(), inits);
+
+ // Generate the new yield values and append them to the scf.yield operation.
+ auto yieldOp = cast<AffineYieldOp>(getBody()->getTerminator());
+ ArrayRef<BlockArgument> newIterArgs =
+ newLoop.getBody()->getArguments().take_back(newInitOperands.size());
+ {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(yieldOp);
+ SmallVector<Value> newYieldedValues =
+ newYieldValuesFn(rewriter, getLoc(), newIterArgs);
+ assert(newInitOperands.size() == newYieldedValues.size() &&
+ "expected as many new yield values as new iter operands");
+ rewriter.updateRootInPlace(yieldOp, [&]() {
+ yieldOp.getOperandsMutable().append(newYieldedValues);
+ });
+ }
+
+ // Move the loop body to the new op.
+ rewriter.mergeBlocks(getBody(), newLoop.getBody(),
+ newLoop.getBody()->getArguments().take_front(
+ getBody()->getNumArguments()));
+
+ if (replaceInitOperandUsesInLoop) {
+ // Replace all uses of `newInitOperands` with the corresponding basic block
+ // arguments.
+ for (auto it : llvm::zip(newInitOperands, newIterArgs)) {
+ rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),
+ [&](OpOperand &use) {
+ Operation *user = use.getOwner();
+ return newLoop->isProperAncestor(user);
+ });
+ }
+ }
+
+ // Replace the old loop.
+ rewriter.replaceOp(getOperation(),
+ newLoop->getResults().take_front(getNumResults()));
+ return cast<LoopLikeOpInterface>(newLoop.getOperation());
+}
+
Speculation::Speculatability AffineForOp::getSpeculatability() {
// `affine.for (I = Start; I < End; I += 1)` terminates for all values of
// Start and End.
@@ -2725,50 +2781,6 @@ void mlir::affine::buildAffineLoopNest(
buildAffineLoopFromValues);
}
-AffineForOp mlir::affine::replaceForOpWithNewYields(OpBuilder &b,
- AffineForOp loop,
- ValueRange newIterOperands,
- ValueRange newYieldedValues,
- ValueRange newIterArgs,
- bool replaceLoopResults) {
- assert(newIterOperands.size() == newYieldedValues.size() &&
- "newIterOperands must be of the same size as newYieldedValues");
- // Create a new loop before the existing one, with the extra operands.
- OpBuilder::InsertionGuard g(b);
- b.setInsertionPoint(loop);
- auto operands = llvm::to_vector<4>(loop.getInits());
- operands.append(newIterOperands.begin(), newIterOperands.end());
- SmallVector<Value, 4> lbOperands(loop.getLowerBoundOperands());
- SmallVector<Value, 4> ubOperands(loop.getUpperBoundOperands());
- SmallVector<Value, 4> steps(loop.getStep());
- auto lbMap = loop.getLowerBoundMap();
- auto ubMap = loop.getUpperBoundMap();
- AffineForOp newLoop =
- b.create<AffineForOp>(loop.getLoc(), lbOperands, lbMap, ubOperands, ubMap,
- loop.getStep(), operands);
- // Take the body of the original parent loop.
- newLoop.getRegion().takeBody(loop.getRegion());
- for (Value val : newIterArgs)
- newLoop.getRegion().addArgument(val.getType(), val.getLoc());
-
- // Update yield operation with new values to be added.
- if (!newYieldedValues.empty()) {
- auto yield = cast<AffineYieldOp>(newLoop.getBody()->getTerminator());
- b.setInsertionPoint(yield);
- auto yieldOperands = llvm::to_vector<4>(yield.getOperands());
- yieldOperands.append(newYieldedValues.begin(), newYieldedValues.end());
- b.create<AffineYieldOp>(yield.getLoc(), yieldOperands);
- yield.erase();
- }
- if (replaceLoopResults) {
- for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
- loop.getNumResults()))) {
- std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
- }
- }
- return newLoop;
-}
-
//===----------------------------------------------------------------------===//
// AffineIfOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 3ecb8664e3fd765..5053b08ee0834cd 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -19,6 +19,7 @@
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>
@@ -361,16 +362,22 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
if (!tripCount || *tripCount != 1)
return failure();
- auto iterOperands = forOp.getInits();
auto *parentOp = forOp->getParentOp();
if (!isa<AffineForOp>(parentOp))
return failure();
- auto newOperands = forOp.getBody()->getTerminator()->getOperands();
- OpBuilder b(parentOp);
+ SmallVector<Value> newOperands;
+ llvm::append_range(newOperands,
+ forOp.getBody()->getTerminator()->getOperands());
+ IRRewriter rewriter(parentOp->getContext());
+ int64_t parentOpNumResults = parentOp->getNumResults();
// Replace the parent loop and add iteroperands and results from the `forOp`.
AffineForOp parentForOp = forOp->getParentOfType<AffineForOp>();
- AffineForOp newLoop = replaceForOpWithNewYields(
- b, parentForOp, iterOperands, newOperands, forOp.getRegionIterArgs());
+ AffineForOp newLoop =
+ cast<AffineForOp>(*parentForOp.replaceWithAdditionalYields(
+ rewriter, forOp.getInits(), /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs) {
+ return newOperands;
+ }));
// For sibling-fusion users, collect operations that use the results of the
// `forOp` outside the new parent loop that has absorbed all its iter args
@@ -387,7 +394,7 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
// Update the results of the `forOp` in the new loop.
for (unsigned i = 0, e = forOp.getNumResults(); i != e; ++i) {
forOp.getResult(i).replaceAllUsesWith(
- newLoop.getResult(i + parentOp->getNumResults()));
+ newLoop.getResult(i + parentOpNumResults));
}
// For sibling-fusion users, move operations that use the results of the
// `forOp` outside the new parent loop
@@ -412,7 +419,6 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations());
forOp.erase();
- parentForOp.erase();
return success();
}
diff --git a/mlir/lib/Dialect/Affine/Utils/Loo...
[truncated]
|
5a05ece
to
cbf5444
Compare
cbf5444
to
8780fcc
Compare
…alues Add a new interface method that returns the yielded values. Also add a verifier that checks the number of inits/iter_args/yielded values. Most of the checked invariants (but not all of them) are already covered by the `RegionBranchOpInterface`, but the `LoopLikeOpInterface` now provides (additional) error messages that are easier to read.
8780fcc
to
54844eb
Compare
@@ -1972,6 +1972,12 @@ mlir::Value fir::IterWhileOp::blockArgToSourceOp(unsigned blockArgNum) { | |||
return {}; | |||
} | |||
|
|||
mlir::ValueRange fir::IterWhileOp::getYieldedValues() { | |||
auto *term = getRegion().front().getTerminator(); | |||
return getFinalValue() ? term->getOperands().drop_front() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just double checking, the drop_front
here drops the final value, is that correct?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes that's correct. (There is similar code in the op verifier: auto opResults = getFinalValue() ? getResults().drop_front() : getResults();
)
Add a new interface method that returns the yielded values.
Also add a verifier that checks the number of inits/iter_args/yielded values. Most of the checked invariants (but not all of them) are already covered by the
RegionBranchOpInterface
, but theLoopLikeOpInterface
now provides (additional) error messages that are easier to read.