[Mlir-commits] [mlir] 63086d6 - [mlir][Interfaces] `LoopLikeOpInterface`: Add `replaceWithAdditionalYields` (#67121)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 26 22:53:43 PDT 2023
Author: Matthias Springer
Date: 2023-09-27T07:53:39+02:00
New Revision: 63086d6aa0af9bb7fc73c670d680191ae646f7d8
URL: https://github.com/llvm/llvm-project/commit/63086d6aa0af9bb7fc73c670d680191ae646f7d8
DIFF: https://github.com/llvm/llvm-project/commit/63086d6aa0af9bb7fc73c670d680191ae646f7d8.diff
LOG: [mlir][Interfaces] `LoopLikeOpInterface`: Add `replaceWithAdditionalYields` (#67121)
`affine::replaceForOpWithNewYields` and `replaceLoopWithNewYields` (for
"scf.for") are now interface methods and additional loop-carried
variables can now be added to "scf.for"/"affine.for" uniformly. (No more
`TypeSwitch` needed.)
Note: `scf.while` and other loops with loop-carried variables can
implement `replaceWithAdditionalYields`, but to keep this commit small,
that is not done in this commit.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
mlir/include/mlir/Dialect/SCF/Utils/Utils.h
mlir/include/mlir/Interfaces/LoopLikeInterface.h
mlir/include/mlir/Interfaces/LoopLikeInterface.td
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
mlir/lib/Dialect/SCF/IR/SCF.cpp
mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
mlir/lib/Dialect/SCF/Utils/Utils.cpp
mlir/test/Transforms/scf-replace-with-new-yields.mlir
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
Removed:
################################################################################
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..d8ef0506d0822d7 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -120,7 +120,7 @@ def AffineForOp : Affine_Op<"for",
[AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound"]>,
+ "getSingleUpperBound", "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 6130f031ca6ab2d..e1a604a88715f0e 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -122,7 +122,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
def ForOp : SCF_Op<"for",
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInits", "getSingleInductionVar", "getSingleLowerBound",
- "getSingleStep", "getSingleUpperBound", "promoteIfSingleIteration"]>,
+ "getSingleStep", "getSingleUpperBound", "promoteIfSingleIteration",
+ "replaceWithAdditionalYields"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
ConditionallySpeculatable,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
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..0eebb984e5897ae 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -17,6 +17,13 @@
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 mlir
/// Include the generated interface declarations.
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index cb6b2f4ed4ae8b5..ded0a29292ff6f0 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -141,6 +141,31 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
return ::mlir::Block::BlockArgListType();
}]
>,
+ 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 +174,24 @@ 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);
+ });
+ }
+ }];
}
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 9ecc568883a3b0b..6c060c90e24af82 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2575,6 +2575,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 +2777,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/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index e6c4b2f8447470c..9d8ed9b4ac93387 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1197,9 +1197,9 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// `unrollJamFactor` copies of its iterOperands, iter_args and yield
// operands.
SmallVector<AffineForOp, 4> newLoopsWithIterArgs;
- OpBuilder builder(forOp.getContext());
+ IRRewriter rewriter(forOp.getContext());
for (AffineForOp oldForOp : loopsWithIterArgs) {
- SmallVector<Value, 4> dupIterOperands, dupIterArgs, dupYieldOperands;
+ SmallVector<Value> dupIterOperands, dupYieldOperands;
ValueRange oldIterOperands = oldForOp.getInits();
ValueRange oldIterArgs = oldForOp.getRegionIterArgs();
ValueRange oldYieldOperands =
@@ -1208,19 +1208,21 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// fix iterOperands and yield operands after cloning of sub-blocks.
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end());
- dupIterArgs.append(oldIterArgs.begin(), oldIterArgs.end());
dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end());
}
// Create a new loop with additional iterOperands, iter_args and yield
// operands. This new loop will take the loop body of the original loop.
- AffineForOp newForOp = affine::replaceForOpWithNewYields(
- builder, oldForOp, dupIterOperands, dupYieldOperands, dupIterArgs);
+ bool forOpReplaced = oldForOp == forOp;
+ AffineForOp newForOp =
+ cast<AffineForOp>(*oldForOp.replaceWithAdditionalYields(
+ rewriter, dupIterOperands, /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs) {
+ return dupYieldOperands;
+ }));
newLoopsWithIterArgs.push_back(newForOp);
// `forOp` has been replaced with a new loop.
- if (oldForOp == forOp)
+ if (forOpReplaced)
forOp = newForOp;
- assert(oldForOp.use_empty() && "old for op should not have any user");
- oldForOp.erase();
// Update `operandMaps` for `newForOp` iterArgs and results.
ValueRange newIterArgs = newForOp.getRegionIterArgs();
unsigned oldNumIterArgs = oldIterArgs.size();
@@ -1294,7 +1296,7 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// into one value. For example, for %0:2 = affine.for ... and addf, we add
// %1 = arith.addf %0#0, %0#1, and replace the following uses of %0#0 with
// %1.
- builder.setInsertionPointAfter(forOp);
+ rewriter.setInsertionPointAfter(forOp);
auto loc = forOp.getLoc();
unsigned oldNumResults = forOp.getNumResults() / unrollJamFactor;
for (LoopReduction &reduction : reductions) {
@@ -1305,7 +1307,7 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
rhs = forOp.getResult(i * oldNumResults + pos);
// Create ops based on reduction type.
- lhs = arith::getReductionOp(reduction.kind, builder, loc, lhs, rhs);
+ lhs = arith::getReductionOp(reduction.kind, rewriter, loc, lhs, rhs);
if (!lhs)
return failure();
Operation *op = lhs.getDefiningOp();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index a9debb7bbc489a4..72bd2b409f5d52b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -842,8 +842,11 @@ padThroughLoopIterArg(RewriterBase &rewriter, Value paddedValueBeforeHoisting,
outerSliceOp.getMixedStrides());
rewriter.replaceAllUsesWith(forOp.getResult(iterArgNumber), extracted);
}
- scf::ForOp newForOp =
- replaceLoopWithNewYields(rewriter, forOp, initArgs, yieldOperands);
+ scf::ForOp newForOp = cast<scf::ForOp>(*forOp.replaceWithAdditionalYields(
+ rewriter, initArgs, /*replaceInitOperandUsesInLoop=*/true,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
+ return yieldOperands;
+ }));
LLVM_DEBUG(DBGS() << "newForOp results: " << newForOp.getNumResults()
<< "\n");
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 7c6639304d97c58..1c68ca49725effb 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -191,47 +191,24 @@ void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
transferWrite->moveAfter(loop);
// Rewrite `loop` with new yields by cloning and erase the original loop.
- OpBuilder b(transferRead);
- NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
- ArrayRef<BlockArgument> newBBArgs) {
+ IRRewriter rewriter(transferRead.getContext());
+ NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
+ ArrayRef<BlockArgument> newBBArgs) {
return SmallVector<Value>{transferWrite.getVector()};
};
- // Transfer write has been hoisted, need to update the written vector by
- // the value yielded by the newForOp.
- return TypeSwitch<Operation *, WalkResult>(loop)
- .Case<scf::ForOp>([&](scf::ForOp scfForOp) {
- auto newForOp = replaceLoopWithNewYields(
- b, scfForOp, transferRead.getVector(), yieldFn);
- transferWrite.getVectorMutable().assign(
- newForOp.getResults().back());
- changed = true;
- loop.erase();
- // Need to interrupt and restart because erasing the loop messes up
- // the walk.
- return WalkResult::interrupt();
- })
- .Case<affine::AffineForOp>([&](affine::AffineForOp affineForOp) {
- auto newForOp = replaceForOpWithNewYields(
- b, affineForOp, transferRead.getVector(),
- SmallVector<Value>{transferWrite.getVector()},
- transferWrite.getVector());
- // Replace all uses of the `transferRead` with the corresponding
- // basic block argument.
- transferRead.getVector().replaceUsesWithIf(
- newForOp.getBody()->getArguments().back(), [&](OpOperand &use) {
- Operation *user = use.getOwner();
- return newForOp->isProperAncestor(user);
- });
- transferWrite.getVectorMutable().assign(
- newForOp.getResults().back());
- changed = true;
- loop.erase();
- // Need to interrupt and restart because erasing the loop messes up
- // the walk.
- return WalkResult::interrupt();
- })
- .Default([](Operation *) { return WalkResult::interrupt(); });
+ auto maybeNewLoop = loop.replaceWithAdditionalYields(
+ rewriter, transferRead.getVector(),
+ /*replaceInitOperandUsesInLoop=*/true, yieldFn);
+ if (failed(maybeNewLoop))
+ return WalkResult::interrupt();
+
+ transferWrite.getVectorMutable().assign(
+ maybeNewLoop->getOperation()->getResults().back());
+ changed = true;
+ // Need to interrupt and restart because erasing the loop messes up
+ // the walk.
+ return WalkResult::interrupt();
});
}
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
index 7ab4ea41a2cd89d..91e0d139ec5c2f0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
@@ -363,13 +363,13 @@ static scf::ForOp hoistTransferReadWrite(
// 2. Rewrite `loop` with an additional yield. This is the quantity that is
// computed iteratively but whose storage has become loop-invariant.
- NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
- ArrayRef<BlockArgument> newBBArgs) {
+ NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
+ ArrayRef<BlockArgument> newBBArgs) {
return SmallVector<Value>{transferWriteOp.getVector()};
};
- auto newForOp = replaceLoopWithNewYields(
- rewriter, forOp, {transferReadOp.getVector()}, yieldFn);
- rewriter.eraseOp(forOp);
+ auto newForOp = cast<scf::ForOp>(*forOp.replaceWithAdditionalYields(
+ rewriter, {transferReadOp.getVector()},
+ /*replaceInitOperandUsesInLoop=*/true, yieldFn));
// 3. Update the yield. Invariant: initArgNumber is the destination tensor.
auto yieldOp =
@@ -425,13 +425,13 @@ static scf::ForOp hoistExtractInsertSlice(RewriterBase &rewriter,
// 2. Rewrite `loop` with an additional yield. This is the quantity that is
// computed iteratively but whose storage has become loop-invariant.
- NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
- ArrayRef<BlockArgument> newBBArgs) {
+ NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,
+ ArrayRef<BlockArgument> newBBArgs) {
return SmallVector<Value>{insertSliceOp.getSource()};
};
- auto newForOp = replaceLoopWithNewYields(rewriter, forOp,
- extractSliceOp.getResult(), yieldFn);
- rewriter.eraseOp(forOp);
+ auto newForOp = cast<scf::ForOp>(*forOp.replaceWithAdditionalYields(
+ rewriter, extractSliceOp.getResult(),
+ /*replaceInitOperandUsesInLoop=*/true, yieldFn));
// 3. Update the yield. Invariant: initArgNumber is the destination tensor.
auto yieldOp =
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8788597a1cefcfa..8d8481421e18d57 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -529,6 +529,59 @@ SmallVector<Region *> ForOp::getLoopRegions() { return {&getRegion()}; }
OperandRange ForOp::getInits() { return getInitArgs(); }
+FailureOr<LoopLikeOpInterface>
+ForOp::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(getInitArgs());
+ inits.append(newInitOperands.begin(), newInitOperands.end());
+ scf::ForOp newLoop = rewriter.create<scf::ForOp>(
+ getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
+ [](OpBuilder &, Location, Value, ValueRange) {});
+
+ // Generate the new yield values and append them to the scf.yield operation.
+ auto yieldOp = cast<scf::YieldOp>(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.getResultsMutable().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());
+}
+
ForOp mlir::scf::getForInductionVarOwner(Value val) {
auto ivArg = llvm::dyn_cast<BlockArgument>(val);
if (!ivArg)
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index bc913e94a2837b2..72fb30388c631cf 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -197,7 +197,7 @@ yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
ArrayRef<SmallVector<OpFoldResult>> tileOffsetsList,
ArrayRef<SmallVector<OpFoldResult>> tileSizesList,
MutableArrayRef<scf::ForOp> loops) {
- NewYieldValueFn yieldValueFn =
+ NewYieldValuesFn yieldValueFn =
[&](OpBuilder &b, Location loc,
ArrayRef<BlockArgument> newBBArgs) -> SmallVector<Value> {
SmallVector<Value> inserts;
@@ -219,7 +219,6 @@ yieldTiledValues(RewriterBase &rewriter, ValueRange initValues,
replaceLoopNestWithNewYields(rewriter, loops, initValues, yieldValueFn,
/*replaceIterOperandsUsesInLoop =*/false);
for (const auto &loop : llvm::enumerate(loops)) {
- rewriter.eraseOp(loop.value());
loops[loop.index()] = newLoops[loop.index()];
}
return llvm::to_vector(llvm::map_range(
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 411503700eb01c3..5360c493f8f8d71 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -38,77 +38,9 @@ struct LoopParams {
};
} // namespace
-scf::ForOp
-mlir::replaceLoopWithNewYields(OpBuilder &builder, scf::ForOp loop,
- ValueRange newIterOperands,
- const NewYieldValueFn &newYieldValuesFn,
- bool replaceIterOperandsUsesInLoop) {
- // Create a new loop before the existing one, with the extra operands.
- OpBuilder::InsertionGuard g(builder);
- builder.setInsertionPoint(loop);
- auto operands = llvm::to_vector(loop.getInitArgs());
- llvm::append_range(operands, newIterOperands);
- scf::ForOp newLoop = builder.create<scf::ForOp>(
- loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
- operands, [](OpBuilder &, Location, Value, ValueRange) {});
-
- Block *loopBody = loop.getBody();
- Block *newLoopBody = newLoop.getBody();
-
- // Move the body of the original loop to the new loop.
- newLoopBody->getOperations().splice(newLoopBody->end(),
- loopBody->getOperations());
-
- // Generate the new yield values to use by using the callback and append the
- // yield values to the scf.yield operation.
- auto yield = cast<scf::YieldOp>(newLoopBody->getTerminator());
- ArrayRef<BlockArgument> newBBArgs =
- newLoopBody->getArguments().take_back(newIterOperands.size());
- {
- OpBuilder::InsertionGuard g(builder);
- builder.setInsertionPoint(yield);
- SmallVector<Value> newYieldedValues =
- newYieldValuesFn(builder, loop.getLoc(), newBBArgs);
- assert(newIterOperands.size() == newYieldedValues.size() &&
- "expected as many new yield values as new iter operands");
- yield.getResultsMutable().append(newYieldedValues);
- }
-
- // Remap the BlockArguments from the original loop to the new loop
- // BlockArguments.
- ArrayRef<BlockArgument> bbArgs = loopBody->getArguments();
- for (auto it :
- llvm::zip(bbArgs, newLoopBody->getArguments().take_front(bbArgs.size())))
- std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
-
- if (replaceIterOperandsUsesInLoop) {
- // Replace all uses of `newIterOperands` with the corresponding basic block
- // arguments.
- for (auto it : llvm::zip(newIterOperands, newBBArgs)) {
- std::get<0>(it).replaceUsesWithIf(std::get<1>(it), [&](OpOperand &use) {
- Operation *user = use.getOwner();
- return newLoop->isProperAncestor(user);
- });
- }
- }
-
- // Replace all uses of the original loop with corresponding values from the
- // new loop.
- loop.replaceAllUsesWith(
- newLoop.getResults().take_front(loop.getNumResults()));
-
- // Add a fake yield to the original loop body that just returns the
- // BlockArguments corresponding to the iter_args. This makes it a no-op loop.
- // The loop is dead. The caller is expected to erase it.
- builder.setInsertionPointToEnd(loopBody);
- builder.create<scf::YieldOp>(loop->getLoc(), loop.getRegionIterArgs());
-
- return newLoop;
-}
-
SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
- OpBuilder &builder, ArrayRef<scf::ForOp> loopNest,
- ValueRange newIterOperands, const NewYieldValueFn &newYieldValueFn,
+ RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
+ ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
bool replaceIterOperandsUsesInLoop) {
if (loopNest.empty())
return {};
@@ -146,31 +78,32 @@ SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
// }
// ```
//
- // The inner most loop is handled using the `replaceLoopWithNewYields`
+ // The inner most loop is handled using the `replaceWithAdditionalYields`
// that works on a single loop.
if (loopNest.size() == 1) {
- auto innerMostLoop = replaceLoopWithNewYields(
- builder, loopNest.back(), newIterOperands, newYieldValueFn,
- replaceIterOperandsUsesInLoop);
+ auto innerMostLoop =
+ cast<scf::ForOp>(*loopNest.back().replaceWithAdditionalYields(
+ rewriter, newIterOperands, replaceIterOperandsUsesInLoop,
+ newYieldValuesFn));
return {innerMostLoop};
}
// The outer loops are modified by calling this method recursively
// - The return value of the inner loop is the value yielded by this loop.
// - The region iter args of this loop are the init_args for the inner loop.
SmallVector<scf::ForOp> newLoopNest;
- NewYieldValueFn fn =
+ NewYieldValuesFn fn =
[&](OpBuilder &innerBuilder, Location loc,
ArrayRef<BlockArgument> innerNewBBArgs) -> SmallVector<Value> {
- newLoopNest = replaceLoopNestWithNewYields(builder, loopNest.drop_front(),
- innerNewBBArgs, newYieldValueFn,
+ newLoopNest = replaceLoopNestWithNewYields(rewriter, loopNest.drop_front(),
+ innerNewBBArgs, newYieldValuesFn,
replaceIterOperandsUsesInLoop);
return llvm::to_vector(llvm::map_range(
newLoopNest.front().getResults().take_back(innerNewBBArgs.size()),
[](OpResult r) -> Value { return r; }));
};
scf::ForOp outerMostLoop =
- replaceLoopWithNewYields(builder, loopNest.front(), newIterOperands, fn,
- replaceIterOperandsUsesInLoop);
+ cast<scf::ForOp>(*loopNest.front().replaceWithAdditionalYields(
+ rewriter, newIterOperands, replaceIterOperandsUsesInLoop, fn));
newLoopNest.insert(newLoopNest.begin(), outerMostLoop);
return newLoopNest;
}
diff --git a/mlir/test/Transforms/scf-replace-with-new-yields.mlir b/mlir/test/Transforms/scf-replace-with-new-yields.mlir
index 802f86e6806b318..d7af15f52e72a9b 100644
--- a/mlir/test/Transforms/scf-replace-with-new-yields.mlir
+++ b/mlir/test/Transforms/scf-replace-with-new-yields.mlir
@@ -15,7 +15,4 @@ func.func @doubleup(%lb: index, %ub: index, %step: index, %extra_arg: f32) -> f3
// CHECK: %[[DOUBLE:.+]] = arith.addf %[[INIT1]], %[[INIT1]]
// CHECK: %[[DOUBLE2:.+]] = arith.addf %[[DOUBLE]], %[[DOUBLE]]
// CHECK: scf.yield %[[DOUBLE]], %[[DOUBLE2]]
-// CHECK: %[[OLDLOOP:.+]] = scf.for
-// CHECK-SAME: iter_args(%[[INIT:.+]] = %[[ARG]])
-// CHECK: scf.yield %[[INIT]]
// CHECK: return %[[NEWLOOP]]#0
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index 455c9234b8c93de..1d40615305c02c4 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -50,19 +50,23 @@ struct TestSCFForUtilsPass
auto newInitValues = forOp.getInitArgs();
if (newInitValues.empty())
return;
- NewYieldValueFn fn = [&](OpBuilder &b, Location loc,
- ArrayRef<BlockArgument> newBBArgs) {
- Block *block = newBBArgs.front().getOwner();
+ auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+ SmallVector<Value> oldYieldValues(yieldOp.getResults().begin(),
+ yieldOp.getResults().end());
+ NewYieldValuesFn fn = [&](OpBuilder &b, Location loc,
+ ArrayRef<BlockArgument> newBBArgs) {
SmallVector<Value> newYieldValues;
- for (auto yieldVal :
- cast<scf::YieldOp>(block->getTerminator()).getResults()) {
+ for (auto yieldVal : oldYieldValues) {
newYieldValues.push_back(
b.create<arith::AddFOp>(loc, yieldVal, yieldVal));
}
return newYieldValues;
};
- OpBuilder b(forOp);
- replaceLoopWithNewYields(b, forOp, newInitValues, fn);
+ IRRewriter rewriter(forOp.getContext());
+ if (failed(forOp.replaceWithAdditionalYields(
+ rewriter, newInitValues, /*replaceInitOperandUsesInLoop=*/true,
+ fn)))
+ signalPassFailure();
});
}
}
More information about the Mlir-commits
mailing list