[mlir][Interfaces] `LoopLikeOpInterface`: Support ops with multiple regions #66754

matthias-springer · 2023-09-19T09:37:05Z

This commit implements LoopLikeOpInterface on scf.while. This enables LICM (and potentially other transforms) on scf.while.

LoopLikeOpInterface::getLoopBody() is renamed to getLoopRegions and can now return multiple regions.

Also fix a bug in the default implementation of LoopLikeOpInterface::isDefinedOutsideOfLoop(), which returned "false" for some values that are defined outside of the loop (in a nested op, in such a way that the value does not dominate the loop). This interface is currently only used for LICM and there is no way to trigger this bug, so no test is added.

llvmbot · 2023-09-19T09:38:11Z

@llvm/pr-subscribers-flang-fir-hlfir
@llvm/pr-subscribers-mlir-tosa
@llvm/pr-subscribers-mlir-scf
@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-memref
@llvm/pr-subscribers-mlir-spirv
@llvm/pr-subscribers-mlir-async
@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir-sparse
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir-affine

Changes

This commit implements LoopLikeOpInterface on scf.while. This enables LICM (and potentially other transforms) on scf.while.

LoopLikeOpInterface::getLoopBody() can now return multiple regions.

Also fix a bug in the default implementation of LoopLikeOpInterface::isDefinedOutsideOfLoop(), which returned "false" for some values that are defined outside of the loop (in a nested op, in such a way that the value does not dominate the loop). This interface is currently only used for LICM and there is no way to trigger this bug, so no test is added.

Patch is 28.42 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/66754.diff

22 Files Affected:

(modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (+1)
(modified) mlir/include/mlir/Interfaces/LoopLikeInterface.td (+3-3)
(modified) mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h (+3-2)
(modified) mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp (+1-1)
(modified) mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp (+3-4)
(modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+8-10)
(modified) mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp (+2-2)
(modified) mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp (+2-2)
(modified) mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp (+1-2)
(modified) mlir/lib/Dialect/Linalg/Transforms/Loops.cpp (+3-3)
(modified) mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp (+1-1)
(modified) mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp (+5-3)
(modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+12-8)
(modified) mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp (+3-6)
(modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+6-10)
(modified) mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp (+3-3)
(modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+2-2)
(modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp (+2-2)
(modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+1-1)
(modified) mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp (+2-2)
(modified) mlir/test/Transforms/loop-invariant-code-motion.mlir (+21)
(modified) mlir/test/lib/Dialect/Test/TestOps.td (+1-1)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 8a9ce949a750d43..cbc9449f3c36bb2 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -958,6 +958,7 @@ def ReduceReturnOp :
 def WhileOp : SCF_Op<"while",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getEntrySuccessorOperands"]>,
+     DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      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 9ccc97251e7e669..61f99cee29b593f 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -36,14 +36,14 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*args=*/(ins "::mlir::Value ":$value),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return value.getParentRegion()->isProperAncestor(&$_op.getLoopBody());
+        return !$_op->isAncestor(value.getParentRegion()->getParentOp());
       }]
     >,
     InterfaceMethod<[{
-        Returns the region that makes up the body of the loop and should be
+        Returns the regions that make up the body of the loop and should be
         inspected for loop-invariant operations.
       }],
-      /*retTy=*/"::mlir::Region &",
+      /*retTy=*/"::llvm::SmallVector<::mlir::Region *>",
       /*methodName=*/"getLoopBody"
     >,
     InterfaceMethod<[{
diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index e54675967d82552..d824884ccb752a3 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -11,12 +11,13 @@
 
 #include "mlir/Support/LLVM.h"
 
+#include "llvm/ADT/SmallVector.h"
+
 namespace mlir {
 
 class LoopLikeOpInterface;
 class Operation;
 class Region;
-class RegionRange;
 class Value;
 
 /// Given a list of regions, perform loop-invariant code motion. An operation is
@@ -61,7 +62,7 @@ class Value;
 ///
 /// Returns the number of operations moved.
 size_t moveLoopInvariantCode(
-    RegionRange regions,
+    SmallVector<Region *> regions,
     function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
     function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
     function_ref<void(Operation *, Region *)> moveOutOfRegion);
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index 81a6378fc7e49e3..f6e3053b8ae6a96 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -163,7 +163,7 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
     signatureConverter.remapInput(0, newIndVar);
     for (unsigned i = 1, e = body->getNumArguments(); i < e; i++)
       signatureConverter.remapInput(i, header->getArgument(i));
-    body = rewriter.applySignatureConversion(&forOp.getLoopBody(),
+    body = rewriter.applySignatureConversion(&forOp.getRegion(),
                                              signatureConverter);
 
     // Move the blocks from the forOp into the loopOp. This is the body of the
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index c8871c945cbe759..f0412648608a6e4 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -1103,7 +1103,7 @@ convertBroadcastOp(RewriterBase &rewriter, vector::BroadcastOp op,
 }
 
 // Replace ForOp with a new ForOp with extra operands. The YieldOp is not
-// updated and needs to be updated separatly for the loop to be correct.
+// updated and needs to be updated separately for the loop to be correct.
 static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter,
                                                scf::ForOp loop,
                                                ValueRange newInitArgs) {
@@ -1119,9 +1119,8 @@ static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter,
       operands);
   newLoop.getBody()->erase();
 
-  newLoop.getLoopBody().getBlocks().splice(
-      newLoop.getLoopBody().getBlocks().begin(),
-      loop.getLoopBody().getBlocks());
+  newLoop.getRegion().getBlocks().splice(
+      newLoop.getRegion().getBlocks().begin(), loop.getRegion().getBlocks());
   for (Value operand : newInitArgs)
     newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
 
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 4455cc88e65e55d..e1bdac49ecb2297 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2380,8 +2380,7 @@ void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
 /// induction variable. AffineForOp only has one region, so zero is the only
 /// valid value for `index`.
 OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  assert((point.isParent() || point == getLoopBody()) &&
-         "invalid region point");
+  assert((point.isParent() || point == getRegion()) && "invalid region 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.
@@ -2395,8 +2394,7 @@ OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
 /// not a constant.
 void AffineForOp::getSuccessorRegions(
     RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
-  assert((point.isParent() || point == getLoopBody()) &&
-         "expected loop region");
+  assert((point.isParent() || point == getRegion()) && "expected loop region");
   // The loop may typically branch back to its body or to the parent operation.
   // If the predecessor is the parent op and the trip count is known to be at
   // least one, branch into the body using the iterator arguments. And in cases
@@ -2404,7 +2402,7 @@ void AffineForOp::getSuccessorRegions(
   std::optional<uint64_t> tripCount = getTrivialConstantTripCount(*this);
   if (point.isParent() && tripCount.has_value()) {
     if (tripCount.value() > 0) {
-      regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+      regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
       return;
     }
     if (tripCount.value() == 0) {
@@ -2422,7 +2420,7 @@ void AffineForOp::getSuccessorRegions(
 
   // In all other cases, the loop may branch back to itself or the parent
   // operation.
-  regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+  regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
@@ -2561,7 +2559,7 @@ bool AffineForOp::matchingBoundOperandList() {
   return true;
 }
 
-Region &AffineForOp::getLoopBody() { return getRegion(); }
+SmallVector<Region *> AffineForOp::getLoopBody() { return {&getRegion()}; }
 
 std::optional<Value> AffineForOp::getSingleInductionVar() {
   return getInductionVar();
@@ -2758,9 +2756,9 @@ AffineForOp mlir::affine::replaceForOpWithNewYields(OpBuilder &b,
       b.create<AffineForOp>(loop.getLoc(), lbOperands, lbMap, ubOperands, ubMap,
                             loop.getStep(), operands);
   // Take the body of the original parent loop.
-  newLoop.getLoopBody().takeBody(loop.getLoopBody());
+  newLoop.getRegion().takeBody(loop.getRegion());
   for (Value val : newIterArgs)
-    newLoop.getLoopBody().addArgument(val.getType(), val.getLoc());
+    newLoop.getRegion().addArgument(val.getType(), val.getLoc());
 
   // Update yield operation with new values to be added.
   if (!newYieldedValues.empty()) {
@@ -3848,7 +3846,7 @@ void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
     ensureTerminator(*bodyRegion, builder, result.location);
 }
 
-Region &AffineParallelOp::getLoopBody() { return getRegion(); }
+SmallVector<Region *> AffineParallelOp::getLoopBody() { return {&getRegion()}; }
 
 unsigned AffineParallelOp::getNumDims() { return getSteps().size(); }
 
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index c9b7f25c545cd1e..6b3776533bed4ae 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -85,11 +85,11 @@ static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
                                       opsToHoist))
       return false;
   } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
-    if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs,
+    if (!areAllOpsInTheBlockListInvariant(forOp.getRegion(), indVar, iterArgs,
                                           opsWithUsers, opsToHoist))
       return false;
   } else if (auto parOp = dyn_cast<AffineParallelOp>(op)) {
-    if (!areAllOpsInTheBlockListInvariant(parOp.getLoopBody(), indVar, iterArgs,
+    if (!areAllOpsInTheBlockListInvariant(parOp.getRegion(), indVar, iterArgs,
                                           opsWithUsers, opsToHoist))
       return false;
   } else if (!isMemoryEffectFree(&op) &&
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
index 7bd25dbbca6365a..12a28c2e23b221a 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
@@ -429,7 +429,7 @@ static ParallelComputeFunction createParallelComputeFunction(
       mapping.map(op.getInductionVars(), computeBlockInductionVars);
       mapping.map(computeFuncType.captures, captures);
 
-      for (auto &bodyOp : op.getLoopBody().getOps())
+      for (auto &bodyOp : op.getRegion().getOps())
         b.clone(bodyOp, mapping);
     };
   };
@@ -732,7 +732,7 @@ AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
 
   // Make sure that all constants will be inside the parallel operation body to
   // reduce the number of parallel compute function arguments.
-  cloneConstantsIntoTheRegion(op.getLoopBody(), rewriter);
+  cloneConstantsIntoTheRegion(op.getRegion(), rewriter);
 
   // Compute trip count for each loop induction variable:
   //   tripCount = ceil_div(upperBound - lowerBound, step);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 2ee21099cfb14c7..7c6639304d97c58 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -219,8 +219,7 @@ void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
             // Replace all uses of the `transferRead` with the corresponding
             // basic block argument.
             transferRead.getVector().replaceUsesWithIf(
-                newForOp.getLoopBody().getArguments().back(),
-                [&](OpOperand &use) {
+                newForOp.getBody()->getArguments().back(), [&](OpOperand &use) {
                   Operation *user = use.getOwner();
                   return newForOp->isProperAncestor(user);
                 });
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
index 5a44a85c95c75a2..f67a3bc595915c0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
@@ -199,9 +199,9 @@ static void replaceIndexOpsByInductionVariables(RewriterBase &rewriter,
   // Replace the index operations in the body of the innermost loop op.
   if (!loopOps.empty()) {
     auto loopOp = cast<LoopLikeOpInterface>(loopOps.back());
-    for (IndexOp indexOp :
-         llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
-      rewriter.replaceOp(indexOp, allIvs[indexOp.getDim()]);
+    for (Region *r : loopOp.getLoopBody())
+      for (IndexOp indexOp : llvm::make_early_inc_range(r->getOps<IndexOp>()))
+        rewriter.replaceOp(indexOp, allIvs[indexOp.getDim()]);
   }
 }
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
index f4556787668d45b..7ab4ea41a2cd89d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp
@@ -303,7 +303,7 @@ static Operation *isTensorChunkAccessedByUnknownOp(Operation *writeOp,
       // pass-through tensor arguments left from previous level of
       // hoisting.
       if (auto forUser = dyn_cast<scf::ForOp>(user)) {
-        Value arg = forUser.getLoopBody().getArgument(
+        Value arg = forUser.getBody()->getArgument(
             use.getOperandNumber() - forUser.getNumControlOperands() +
             /*iv value*/ 1);
         uses.push_back(arg.getUses());
diff --git a/mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp b/mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
index eb1df2a87b99a7d..a528bd6cd1ff493 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
@@ -152,8 +152,9 @@ mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
   std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
   std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
   std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
-  if (!inductionVar || !lowerBound || !singleStep) {
-    LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb or step\n");
+  if (!inductionVar || !lowerBound || !singleStep ||
+      !llvm::hasSingleElement(candidateLoop.getLoopBody())) {
+    LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb, step or region\n");
     return failure();
   }
 
@@ -184,7 +185,8 @@ mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
 
   // 3. Within the loop, build the modular leading index (i.e. each loop
   // iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).
-  rewriter.setInsertionPointToStart(&candidateLoop.getLoopBody().front());
+  rewriter.setInsertionPointToStart(
+      &candidateLoop.getLoopBody().front()->front());
   Value ivVal = *inductionVar;
   Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);
   Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 2a760c76d2f6867..138412ffc70726d 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -530,7 +530,7 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
   return success();
 }
 
-Region &ForOp::getLoopBody() { return getRegion(); }
+SmallVector<Region *> ForOp::getLoopBody() { return {&getRegion()}; }
 
 ForOp mlir::scf::getForInductionVarOwner(Value val) {
   auto ivArg = llvm::dyn_cast<BlockArgument>(val);
@@ -558,11 +558,11 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,
   // Both the operation itself and the region may be branching into the body or
   // back into the operation itself. It is possible for loop not to enter the
   // body.
-  regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+  regions.push_back(RegionSuccessor(&getRegion(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
-Region &ForallOp::getLoopBody() { return getRegion(); }
+SmallVector<Region *> ForallOp::getLoopBody() { return {&getRegion()}; }
 
 /// Promotes the loop body of a forallOp to its containing block if it can be
 /// determined that the loop has a single iteration.
@@ -894,7 +894,7 @@ struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
       blockArgs.reserve(op.getInitArgs().size() + 1);
       blockArgs.push_back(op.getLowerBound());
       llvm::append_range(blockArgs, op.getInitArgs());
-      replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
+      replaceOpWithRegion(rewriter, op, op.getRegion(), blockArgs);
       return success();
     }
 
@@ -2872,7 +2872,7 @@ void ParallelOp::print(OpAsmPrinter &p) {
       /*elidedAttrs=*/ParallelOp::getOperandSegmentSizeAttr());
 }
 
-Region &ParallelOp::getLoopBody() { return getRegion(); }
+SmallVector<Region *> ParallelOp::getLoopBody() { return {&getRegion()}; }
 
 ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = llvm::dyn_cast<BlockArgument>(val);
@@ -2926,7 +2926,7 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
       // loop body and nested ReduceOp's
       SmallVector<Value> results;
       results.reserve(op.getInitVals().size());
-      for (auto &bodyOp : op.getLoopBody().front().without_terminator()) {
+      for (auto &bodyOp : op.getBody()->without_terminator()) {
         auto reduce = dyn_cast<ReduceOp>(bodyOp);
         if (!reduce) {
           rewriter.clone(bodyOp, mapping);
@@ -2965,7 +2965,7 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
 
   LogicalResult matchAndRewrite(ParallelOp op,
                                 PatternRewriter &rewriter) const override {
-    Block &outerBody = op.getLoopBody().front();
+    Block &outerBody = *op.getBody();
     if (!llvm::hasSingleElement(outerBody.without_terminator()))
       return failure();
 
@@ -2985,7 +2985,7 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
 
     auto bodyBuilder = [&](OpBuilder &builder, Location /*loc*/,
                            ValueRange iterVals, ValueRange) {
-      Block &innerBody = innerOp.getLoopBody().front();
+      Block &innerBody = *innerOp.getBody();
       assert(iterVals.size() ==
              (outerBody.getNumArguments() + innerBody.getNumArguments()));
       IRMapping mapping;
@@ -3203,6 +3203,10 @@ void WhileOp::getSuccessorRegions(RegionBranchPoint point,
   regions.emplace_back(&getAfter(), getAfter().getArguments());
 }
 
+SmallVector<Region *> WhileOp::getLoopBody() {
+  return {&getBefore(), &getAfter()};
+}
+
 /// Parses a `while` op.
 ///
 /// op ::= `scf.while` assignments `:` function-type region `do` region
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 47f6da25b325144..88c6f3da656f3ba 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -35,11 +35,8 @@ struct ForOpInterface
 
     // An EQ constraint can be added if the yielded value (dimension size)
     // equals the corresponding block argument (dimension size).
-    assert(forOp.getLoopBody().hasOneBlock() &&
-           "multiple blocks not supported");
-    Value yieldedValue =
-        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator())
-            .getOperand(iterArgIdx);
+    Value yieldedValue = cast<scf::YieldOp>(forOp.getBody()->getTerminator())
+                             .getOperand(iterArgIdx);
     Value iterArg = forOp.getRegionIterArg(iterArgIdx);
     Value initArg = forOp.getInitArgs()[iterArgIdx];
 
@@ -68,7 +65,7 @@ struct ForOpInterface
           // Stop when reaching a value that is defined outside of the loop. It
           // is impossible to reach an iter_arg from there.
           Operation *op = v.getDefiningOp();
-          return forOp.getLoopBody().findAncestorOpInRegion(*op) == nullptr;
+          return forOp.getRegion().findAncestorOpInRegion(*op) == nullptr;
         });
     if (failed(status))
       return;
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index e09e14dbeb2c7f3..bcbc693a9742ccc 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -489,8 +489,7 @@ struct ForOpInterface
     auto forOp = cast<scf::ForOp>(op);
     OpOperand &forOperand = forOp.getOpOperandForResult(opResult);
     auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
-    auto yieldOp =
-        cast<scf::YieldOp>(forOp.getLoopBod...
[truncated]

zero9178

I added two comments that are rather bikesheddy, so feel free to ignore them if you don't see the point in them.
Otherwise, LGTM as this is a mostly mechanical change.

Do note that there is a seemingly real build failure in Flang according to CI.

mlir/test/Transforms/loop-invariant-code-motion.mlir

mlir/include/mlir/Interfaces/LoopLikeInterface.td

…egions This commit implements `LoopLikeOpInterface` on `scf.while`. This enables LICM (and potentially other transforms) on `scf.while`. `LoopLikeOpInterface::getLoopBody()` can now return multiple regions. Also fix a bug in the default implementation of `LoopLikeOpInterface::isDefinedOutsideOfLoop()`, which returned "false" for some values that are defined outside of the loop (in a nested op, in such a way that the value does not dominate the loop). This interface is currently only used for LICM and there is no way to trigger this bug, so no test is added. BEGIN_PUBLIC No public commit message needed for presubmit. END_PUBLIC

llvm/llvm-project#66754 extends the `LoopLikeOpInterface`, so the signature of `getLoopBody` has changed. `ForOp::getRegion` can be used instead.

llvm/llvm-project#66754 extends the `LoopLikeOpInterface`: the signature of `getLoopBody` has changed. `ForOp::getRegion` can be used instead. This change works with and without llvm/llvm-project#66754.

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

llvm/llvm-project#66754 extends the `LoopLikeOpInterface`: the signature of `getLoopBody` has changed. `ForOp::getRegion` can be used instead. This change works with and without llvm/llvm-project#66754.

matthias-springer requested review from zero9178, nicolasvasilache and Mogball September 19, 2023 09:37

llvmbot added mlir:core MLIR Core Infrastructure mlir:linalg mlir:gpu mlir:spirv mlir:affine mlir:sparse Sparse compiler in MLIR mlir mlir:tosa mlir:memref mlir:scf mlir:async mlir:afine labels Sep 19, 2023

matthias-springer force-pushed the loop_like_op_interface_multiple_regions branch 2 times, most recently from 58c78c9 to d21cb79 Compare September 19, 2023 09:52

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir labels Sep 19, 2023

zero9178 approved these changes Sep 19, 2023

View reviewed changes

mlir/test/Transforms/loop-invariant-code-motion.mlir Outdated Show resolved Hide resolved

mlir/include/mlir/Interfaces/LoopLikeInterface.td Outdated Show resolved Hide resolved

matthias-springer force-pushed the loop_like_op_interface_multiple_regions branch 2 times, most recently from 50abe0e to 19fabf7 Compare September 19, 2023 13:34

matthias-springer force-pushed the loop_like_op_interface_multiple_regions branch from 19fabf7 to 5f22e1a Compare September 19, 2023 13:42

matthias-springer merged commit 9b5ef2b into llvm:main Sep 19, 2023

matthias-springer mentioned this pull request Sep 19, 2023

Integration fixes for llvm/llvm-project#66754 triton-lang/triton#2338

Merged

This was referenced Sep 19, 2023

[SelectionDAG] [NFC] Add pre-commit test for PR66701. srpande/llvm-project#2

Closed

[SelectionDAG] [NFC] Add pre-commit test for PR66701. srpande/llvm-project#3

Closed

This was referenced Sep 20, 2023

sve gather scatter offset sinking #66932

Closed

sve gather scatter offset sinking #66994

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`: Support ops with multiple regions #66754

[mlir][Interfaces] `LoopLikeOpInterface`: Support ops with multiple regions #66754

matthias-springer commented Sep 19, 2023 •

edited

Loading

llvmbot commented Sep 19, 2023 •

edited

Loading

zero9178 left a comment

[mlir][Interfaces] LoopLikeOpInterface: Support ops with multiple regions #66754

[mlir][Interfaces] LoopLikeOpInterface: Support ops with multiple regions #66754

Conversation

matthias-springer commented Sep 19, 2023 • edited Loading

llvmbot commented Sep 19, 2023 • edited Loading

zero9178 left a comment

Choose a reason for hiding this comment

[mlir][Interfaces] `LoopLikeOpInterface`: Support ops with multiple regions #66754

[mlir][Interfaces] `LoopLikeOpInterface`: Support ops with multiple regions #66754

matthias-springer commented Sep 19, 2023 •

edited

Loading

llvmbot commented Sep 19, 2023 •

edited

Loading