[mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits #66925

matthias-springer · 2023-09-20T16:05:51Z

AffineForOp::getInitOperands is renamed to getInits to be consistent with MLIR operand getter naming conventions. ("get" + operand name)

llvmbot · 2023-09-20T16:06:58Z

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-scf

@llvm/pr-subscribers-mlir-affine

Changes

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 --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> &regions) {
   // 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"

zero9178 · 2023-09-20T18:04:57Z

mlir/lib/Dialect/SCF/IR/SCF.cpp

@@ -564,6 +566,8 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,

 SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }

+OperandRange ForallOp::getInits() { return getOutputs(); }


The name of this operand makes this line look very confusing out of context 😅

Ahh yes, this looks confusing indeed, but there's not much we can do (maybe renaming the outputs operand).

Actually, the "outputs" and bbarg of scf.forall is not an iter_arg like scf.for. Because there are no iterations, just parallelism. Maybe we want to treat them like iter_args at some point, but I'm removing it again for now.

…ter_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

There were three breaking changes in LLVM: 1: In the LLVM dialect, NullOp was replaced with a more general ZeroOp. See: llvm/llvm-project#67183. 2: AffineForOp's getInitOperands was renamed to getInits. See: llvm/llvm-project#66925. 3: In the SCF dialect, the WhileOp now implements LoopLikeOpInterface, which now supports multiple loop regions. The getLoopBody member-function was removed in favour of LoopLike's getLoopRegions. The While Op only has a single region, so we can replace calls to getLoopBody with getRegion. See: llvm/llvm-project#66754

matthias-springer requested review from ftynse, zero9178 and nicolasvasilache September 20, 2023 16:05

llvmbot added mlir:affine mlir mlir:scf mlir:afine labels Sep 20, 2023

zero9178 approved these changes Sep 21, 2023

View reviewed changes

matthias-springer force-pushed the loop_like_op_interface_inits branch from cf5d6b7 to 822a464 Compare September 21, 2023 08:28

matthias-springer merged commit 3bd7a9b into llvm:main Sep 21, 2023

kstoimenov mentioned this pull request Sep 22, 2023

Add memcpm test kstoimenov/llvm-project#13

Closed

rwy7 mentioned this pull request Sep 27, 2023

Bump LLVM llvm/circt#6207

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits #66925

[mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits #66925

matthias-springer commented Sep 20, 2023

llvmbot commented Sep 20, 2023 •

edited

Loading

zero9178 Sep 20, 2023

matthias-springer Sep 21, 2023 •

edited

Loading

matthias-springer Sep 21, 2023

		@@ -564,6 +566,8 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,

		SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }

		OperandRange ForallOp::getInits() { return getOutputs(); }

[mlir][Interfaces] LoopLikeOpInterface: Add helpers to get region iter_args and inits #66925

[mlir][Interfaces] LoopLikeOpInterface: Add helpers to get region iter_args and inits #66925

Conversation

matthias-springer commented Sep 20, 2023

llvmbot commented Sep 20, 2023 • edited Loading

zero9178 Sep 20, 2023

Choose a reason for hiding this comment

matthias-springer Sep 21, 2023 • edited Loading

Choose a reason for hiding this comment

matthias-springer Sep 21, 2023

Choose a reason for hiding this comment

[mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits #66925

[mlir][Interfaces] `LoopLikeOpInterface`: Add helpers to get region iter_args and inits #66925

llvmbot commented Sep 20, 2023 •

edited

Loading

matthias-springer Sep 21, 2023 •

edited

Loading