From 822a4644980f7fb0a8101741556a700f493cff1e Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Thu, 21 Sep 2023 10:26:28 +0200 Subject: [PATCH] [mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits `AffineForOp::getInitOperands` is renamed to `getInits` to be consistent with MLIR operand getter naming conventions. ("get" + operand name) BEGIN_PUBLIC No public commit message needed for presubmit. END_PUBLIC --- .../mlir/Dialect/Affine/IR/AffineOps.td | 2 +- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 6 ++--- .../mlir/Interfaces/LoopLikeInterface.td | 25 +++++++++++++++++++ .../AffineToStandard/AffineToStandard.cpp | 2 +- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 16 ++++++------ .../Affine/Transforms/SuperVectorize.cpp | 6 ++--- .../Dialect/Affine/Utils/LoopFusionUtils.cpp | 2 +- mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 6 ++--- mlir/lib/Dialect/Affine/Utils/Utils.cpp | 2 +- mlir/lib/Dialect/SCF/IR/SCF.cpp | 6 +++++ mlir/lib/Interfaces/LoopLikeInterface.cpp | 1 + 11 files changed, 53 insertions(+), 21 deletions(-) diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index 69798905f7a15..5a1baaf4e1611 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 08b71e20a2bc0..0c93989ca99a4 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, + ["getInits", "getSingleInductionVar", "getSingleLowerBound", + "getSingleStep", "getSingleUpperBound", "promoteIfSingleIteration"]>, AllTypesMatch<["lowerBound", "upperBound", "step"]>, ConditionallySpeculatable, DeclareOpInterfaceMethods, - DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, 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 44d32dd609fc9..cb6b2f4ed4ae8 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 783745a51822e..456fd487352fb 100644 --- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp +++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp @@ -158,7 +158,7 @@ class AffineForLowering : public OpRewritePattern { Value upperBound = lowerAffineUpperBound(op, rewriter); Value step = rewriter.create(loc, op.getStep()); auto scfForOp = rewriter.create(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 5899c198b703b..45c65e04e635a 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 { 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 replacements; @@ -2352,7 +2352,7 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern { 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 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 lbOperands(loop.getLowerBoundOperands()); SmallVector ubOperands(loop.getUpperBoundOperands()); diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp index 072e858220fea..85c2602aa266d 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 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(newParentOp).getIterOperands()[i]; + Value neutralVal = cast(newParentOp).getInits()[i]; state.builder.setInsertionPoint(combinerOps.back()); Value maskedReducedVal = state.builder.create( reducedVal.getLoc(), mask, reducedVal, neutralVal); diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp index 70aba4cedc7f3..3ecb8664e3fd7 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 tripCount = getConstantTripCount(forOp); if (!tripCount || *tripCount != 1) return failure(); - auto iterOperands = forOp.getIterOperands(); + auto iterOperands = forOp.getInits(); auto *parentOp = forOp->getParentOp(); if (!isa(parentOp)) return failure(); diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp index b54c2b359ba6b..e6c4b2f844747 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 dupIterOperands, dupIterArgs, dupYieldOperands; - ValueRange oldIterOperands = oldForOp.getIterOperands(); + ValueRange oldIterOperands = oldForOp.getInits(); ValueRange oldIterArgs = oldForOp.getRegionIterArgs(); ValueRange oldYieldOperands = cast(oldForOp.getBody()->getTerminator()).getOperands(); diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp index 88bfc5789f2d6..cef5080daad29 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 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 45e68e23a71d6..fe3d62ccceca5 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 ForOp::getLoopRegions() { return {&getRegion()}; } +OperandRange ForOp::getInits() { return getInitArgs(); } + ForOp mlir::scf::getForInductionVarOwner(Value val) { auto ivArg = llvm::dyn_cast(val); if (!ivArg) @@ -3183,6 +3185,10 @@ Block::BlockArgListType WhileOp::getAfterArguments() { return getAfterBody()->getArguments(); } +Block::BlockArgListType WhileOp::getRegionIterArgs() { + return getBeforeArguments(); +} + void WhileOp::getSuccessorRegions(RegionBranchPoint point, SmallVectorImpl ®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 85197f95c47e1..781a21bb3ecd3 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"