[Mlir-commits] [mlir] [mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits (PR #66925)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 20 09:07:02 PDT 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
<details>
<summary>Changes</summary>
`AffineForOp::getInitOperands` is renamed to `getInits` to be consistent with MLIR operand getter naming conventions. ("get" + operand name)
---
Full diff: https://github.com/llvm/llvm-project/pull/66925.diff
11 Files Affected:
- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1-1)
- (modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (+4-4)
- (modified) mlir/include/mlir/Interfaces/LoopLikeInterface.td (+25)
- (modified) mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp (+1-1)
- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+8-8)
- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+3-3)
- (modified) mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp (+1-1)
- (modified) mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp (+3-3)
- (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+1-1)
- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+8)
- (modified) mlir/lib/Interfaces/LoopLikeInterface.cpp (+1)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 69798905f7a15fa..5a1baaf4e1611c8 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -259,7 +259,7 @@ def AffineForOp : Affine_Op<"for",
Block::BlockArgListType getRegionIterArgs() {
return getBody()->getArguments().drop_front();
}
- Operation::operand_range getIterOperands() {
+ Operation::operand_range getInits() {
return getOperands().drop_front(getNumControlOperands());
}
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 08b71e20a2bc079..b8343abcb5cdc28 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -120,8 +120,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
def ForOp : SCF_Op<"for",
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound", "promoteIfSingleIteration"]>,
+ ["getInits", "getSingleInductionVar", "getSingleLowerBound",
+ "getSingleStep", "getSingleUpperBound", "promoteIfSingleIteration"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
ConditionallySpeculatable,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
@@ -329,7 +329,7 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["promoteIfSingleIteration"]>,
+ ["getInits", "promoteIfSingleIteration"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
@@ -948,7 +948,7 @@ def ReduceReturnOp :
def WhileOp : SCF_Op<"while",
[DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>,
- DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+ DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getRegionIterArgs"]>,
RecursiveMemoryEffects, SingleBlock]> {
let summary = "a generic 'while' loop";
let description = [{
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 44d32dd609fc9d5..cb6b2f4ed4ae8b5 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -116,6 +116,31 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
return std::nullopt;
}]
>,
+ InterfaceMethod<[{
+ Return the "init" operands that are used as initialization values for
+ the region "iter_args" of this loop.
+ }],
+ /*retTy=*/"::mlir::OperandRange",
+ /*methodName=*/"getInits",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::OperandRange($_op->operand_end(), $_op->operand_end());
+ }]
+ >,
+ InterfaceMethod<[{
+ Return the region "iter_args" (block arguments) that correspond to the
+ "init" operands. If the op has multiple regions, return the
+ corresponding block arguments of the entry region.
+ }],
+ /*retTy=*/"::mlir::Block::BlockArgListType",
+ /*methodName=*/"getRegionIterArgs",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::Block::BlockArgListType();
+ }]
+ >,
];
let extraClassDeclaration = [{
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 783745a51822e5e..456fd487352fb5d 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -158,7 +158,7 @@ class AffineForLowering : public OpRewritePattern<AffineForOp> {
Value upperBound = lowerAffineUpperBound(op, rewriter);
Value step = rewriter.create<arith::ConstantIndexOp>(loc, op.getStep());
auto scfForOp = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound,
- step, op.getIterOperands());
+ step, op.getInits());
rewriter.eraseBlock(scfForOp.getBody());
rewriter.inlineRegionBefore(op.getRegion(), scfForOp.getRegion(),
scfForOp.getRegion().end());
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 5899c198b703b5e..45c65e04e635a2e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2209,7 +2209,7 @@ void AffineForOp::print(OpAsmPrinter &p) {
if (getNumIterOperands() > 0) {
p << " iter_args(";
auto regionArgs = getRegionIterArgs();
- auto operands = getIterOperands();
+ auto operands = getInits();
llvm::interleaveComma(llvm::zip(regionArgs, operands), p, [&](auto it) {
p << std::get<0>(it) << " = " << std::get<1>(it);
@@ -2331,7 +2331,7 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
if (tripCount && *tripCount == 0) {
// The initial values of the iteration arguments would be the op's
// results.
- rewriter.replaceOp(forOp, forOp.getIterOperands());
+ rewriter.replaceOp(forOp, forOp.getInits());
return success();
}
SmallVector<Value, 4> replacements;
@@ -2352,7 +2352,7 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
unsigned pos = std::distance(iterArgs.begin(), iterArgIt);
if (pos != i)
iterArgsNotInOrder = true;
- replacements.push_back(forOp.getIterOperands()[pos]);
+ replacements.push_back(forOp.getInits()[pos]);
}
}
// Bail out when the trip count is unknown and the loop returns any value
@@ -2384,7 +2384,7 @@ OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
// The initial operands map to the loop arguments after the induction
// variable or are forwarded to the results when the trip count is zero.
- return getIterOperands();
+ return getInits();
}
/// Given the region at `index`, or the parent operation if `index` is None,
@@ -2440,7 +2440,7 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
// does not return any results. Since ops that do not return results cannot
// be folded away, we would enter an infinite loop of folds on the same
// affine.for op.
- results.assign(getIterOperands().begin(), getIterOperands().end());
+ results.assign(getInits().begin(), getInits().end());
folded = true;
}
return success(folded);
@@ -2466,7 +2466,7 @@ void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
auto ubOperands = getUpperBoundOperands();
newOperands.append(ubOperands.begin(), ubOperands.end());
- auto iterOperands = getIterOperands();
+ auto iterOperands = getInits();
newOperands.append(iterOperands.begin(), iterOperands.end());
(*this)->setOperands(newOperands);
@@ -2479,7 +2479,7 @@ void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) {
SmallVector<Value, 4> newOperands(getLowerBoundOperands());
newOperands.append(ubOperands.begin(), ubOperands.end());
- auto iterOperands = getIterOperands();
+ auto iterOperands = getInits();
newOperands.append(iterOperands.begin(), iterOperands.end());
(*this)->setOperands(newOperands);
@@ -2745,7 +2745,7 @@ AffineForOp mlir::affine::replaceForOpWithNewYields(OpBuilder &b,
// 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.getIterOperands());
+ 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());
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 072e858220feae3..85c2602aa266d67 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1315,13 +1315,13 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
// Vectorize 'iter_args'.
SmallVector<Value, 8> vecIterOperands;
if (!isLoopVecDim) {
- for (auto operand : forOp.getIterOperands())
+ for (auto operand : forOp.getInits())
vecIterOperands.push_back(vectorizeOperand(operand, state));
} else {
// For reduction loops we need to pass a vector of neutral elements as an
// initial value of the accumulator. We will add the original initial value
// later.
- for (auto redAndOperand : llvm::zip(reductions, forOp.getIterOperands())) {
+ for (auto redAndOperand : llvm::zip(reductions, forOp.getInits())) {
vecIterOperands.push_back(createInitialVector(
std::get<0>(redAndOperand).kind, std::get<1>(redAndOperand), state));
}
@@ -1450,7 +1450,7 @@ static Operation *vectorizeAffineYieldOp(AffineYieldOp yieldOp,
assert(reducedVal && "expect non-null value for parallel reduction loop");
assert(combinerOps.size() == 1 && "expect only one combiner op");
// IterOperands are neutral element vectors.
- Value neutralVal = cast<AffineForOp>(newParentOp).getIterOperands()[i];
+ Value neutralVal = cast<AffineForOp>(newParentOp).getInits()[i];
state.builder.setInsertionPoint(combinerOps.back());
Value maskedReducedVal = state.builder.create<arith::SelectOp>(
reducedVal.getLoc(), mask, reducedVal, neutralVal);
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 70aba4cedc7f30a..3ecb8664e3fd765 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -361,7 +361,7 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
if (!tripCount || *tripCount != 1)
return failure();
- auto iterOperands = forOp.getIterOperands();
+ auto iterOperands = forOp.getInits();
auto *parentOp = forOp->getParentOp();
if (!isa<AffineForOp>(parentOp))
return failure();
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index b54c2b359ba6b79..e6c4b2f8447470c 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -114,7 +114,7 @@ getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
/// yield values while promoting single iteration affine.for ops.
static void replaceIterArgsAndYieldResults(AffineForOp forOp) {
// Replace uses of iter arguments with iter operands (initial values).
- auto iterOperands = forOp.getIterOperands();
+ auto iterOperands = forOp.getInits();
auto iterArgs = forOp.getRegionIterArgs();
for (auto e : llvm::zip(iterOperands, iterArgs))
std::get<1>(e).replaceAllUsesWith(std::get<0>(e));
@@ -987,7 +987,7 @@ static LogicalResult generateCleanupLoopForUnroll(AffineForOp forOp,
// and produce results for the original users of `forOp` results.
auto results = forOp.getResults();
auto cleanupResults = cleanupForOp.getResults();
- auto cleanupIterOperands = cleanupForOp.getIterOperands();
+ auto cleanupIterOperands = cleanupForOp.getInits();
for (auto e : llvm::zip(results, cleanupResults, cleanupIterOperands)) {
std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
@@ -1200,7 +1200,7 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
OpBuilder builder(forOp.getContext());
for (AffineForOp oldForOp : loopsWithIterArgs) {
SmallVector<Value, 4> dupIterOperands, dupIterArgs, dupYieldOperands;
- ValueRange oldIterOperands = oldForOp.getIterOperands();
+ ValueRange oldIterOperands = oldForOp.getInits();
ValueRange oldIterArgs = oldForOp.getRegionIterArgs();
ValueRange oldYieldOperands =
cast<AffineYieldOp>(oldForOp.getBody()->getTerminator()).getOperands();
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 88bfc5789f2d64b..cef5080daad29bc 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -377,7 +377,7 @@ mlir::affine::affineParallelize(AffineForOp forOp,
SmallVector<Value> newResults;
newResults.reserve(numReductions);
for (unsigned i = 0; i < numReductions; ++i) {
- Value init = forOp.getIterOperands()[i];
+ Value init = forOp.getInits()[i];
// This works because we are only handling single-op reductions at the
// moment. A switch on reduction kind or a mechanism to collect operations
// participating in the reduction will be necessary for multi-op reductions.
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 45e68e23a71d60e..a884b599cbd7d8e 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -532,6 +532,8 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
SmallVector<Region *> ForOp::getLoopRegions() { return {&getRegion()}; }
+OperandRange ForOp::getInits() { return getInitArgs(); }
+
ForOp mlir::scf::getForInductionVarOwner(Value val) {
auto ivArg = llvm::dyn_cast<BlockArgument>(val);
if (!ivArg)
@@ -564,6 +566,8 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,
SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }
+OperandRange ForallOp::getInits() { return getOutputs(); }
+
/// Promotes the loop body of a forallOp to its containing block if it can be
/// determined that the loop has a single iteration.
LogicalResult scf::ForallOp::promoteIfSingleIteration(RewriterBase &rewriter) {
@@ -3183,6 +3187,10 @@ Block::BlockArgListType WhileOp::getAfterArguments() {
return getAfterBody()->getArguments();
}
+Block::BlockArgListType WhileOp::getRegionIterArgs() {
+ return getBeforeArguments();
+}
+
void WhileOp::getSuccessorRegions(RegionBranchPoint point,
SmallVectorImpl<RegionSuccessor> ®ions) {
// The parent op always branches to the condition region.
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 85197f95c47e14a..781a21bb3ecd3a0 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Interfaces/LoopLikeInterface.h"
+
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/ADT/DenseSet.h"
``````````
</details>
https://github.com/llvm/llvm-project/pull/66925
More information about the Mlir-commits
mailing list