Drop yieldTiledValuesAndReplace as an interface method.

MaheshRavishankar · MaheshRavishankar · commit 20271d9fd2ee · 2024-01-18T17:23:06.000-08:00
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -220,55 +220,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*defaultImplementation=*/[{
         return ::mlir::failure();
       }]
-    >,
-    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.
-
-        This method is similar to `replaceWithAdditionalYields` but instead of
-        yielding a value from within the loop, it allows each loop construct
-        implementing this method to handle the result of each iteration
-        appropriately. This allows for unified handling of operations
-        like `scf.forall` which don't yield a value from the loop, but instead
-        the terminator specifies where to insert the tile computed by the body of
-        the loop. For example,
-        
-        ```mlir
-        %0 = scf.forall ... shared_outs(%arg0 = %arg1) {
-          ...
-          %tiled_value = ...
-          scf.forall.in_parallel {
-            tensor.parallel_insert_slice %tiled_value into %arg0[%o1, %o2]...
-          }
-        }
-        ```
-
-        For an `scf.for` the same computation would be represented as
-        ```mlir
-        %0 = scf.for ... iter_args(%arg0 = %arg1) {
-          ...
-          %tiled_value = ...
-          %insert = tensor.insert_slice %tiled_value into %arg0[%o1, %o2]...
-          scf.yield %insert
-        }
-        ```
-
-        So for the caller, the tiled value (`%tiled_value`) and the offsets
-        `(%o1, %o2)` and sizes (not shown) are generated the same way, but
-        the implementation method for the different loop constructs handles
-        the difference in representation.
-      }],
-      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
-      /*methodName=*/"yieldTiledValuesAndReplace",
-      /*args=*/(ins "::mlir::RewriterBase &":$rewriter,
-                    "::mlir::ValueRange":$newInitOperands,
-                    "::mlir::YieldTiledValuesFn":$yieldTiledValuesFn),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        return ::mlir::failure();
-      }]
-    >,
+    >
   ];
 
   let extraClassDeclaration = [{
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -588,64 +588,6 @@ ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,
   return cast<LoopLikeOpInterface>(newLoop.getOperation());
 }
 
-FailureOr<LoopLikeOpInterface>
-ForOp::yieldTiledValuesAndReplace(RewriterBase &rewriter,
-                                  ValueRange newInitOperands,
-                                  YieldTiledValuesFn yieldTiledValuesFn) {
-  OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(getOperation());
-
-  auto inits = llvm::to_vector(getInitArgs());
-  inits.append(newInitOperands.begin(), newInitOperands.end());
-  auto newLoop = rewriter.create<ForOp>(
-      getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
-      [](OpBuilder &, Location, Value, ValueRange) {});
-
-  // Move the loop body to the new op.
-  rewriter.mergeBlocks(getBody(), newLoop.getBody(),
-                       newLoop.getBody()->getArguments().take_front(
-                           getBody()->getNumArguments()));
-
-  auto yieldOp = cast<scf::YieldOp>(newLoop.getBody()->getTerminator());
-  rewriter.setInsertionPoint(yieldOp);
-
-  SmallVector<Value> tiledValues;
-  SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
-  ValueRange newRegionIterArgs =
-      newLoop.getRegionIterArgs().take_back(newInitOperands.size());
-  if (failed(yieldTiledValuesFn(rewriter, getLoc(), newLoop.getInductionVar(),
-                                newRegionIterArgs, tiledValues, resultOffsets,
-                                resultSizes))) {
-    return rewriter.notifyMatchFailure(getOperation(),
-                                       "failed to get tiled values");
-  }
-
-  if (tiledValues.size() != resultOffsets.size() ||
-      tiledValues.size() != resultSizes.size()) {
-    return rewriter.notifyMatchFailure(
-        getOperation(),
-        "expected number of tiled values returned, the number of offset "
-        "vectors and number of size vectors to be the same");
-  }
-
-  SmallVector<Value> newYieldValues = llvm::to_vector(yieldOp.getOperands());
-  for (auto [tiledValue, regionIterArg, resultOffset, resultSize] :
-       llvm::zip_equal(tiledValues, newRegionIterArgs, resultOffsets,
-                       resultSizes)) {
-    SmallVector<OpFoldResult> resultStride(resultOffset.size(),
-                                           rewriter.getIndexAttr(1));
-    Value insert = rewriter.create<tensor::InsertSliceOp>(
-        yieldOp->getLoc(), tiledValue, regionIterArg, resultOffset, resultSize,
-        resultStride);
-    newYieldValues.push_back(insert);
-  }
-
-  rewriter.replaceOpWithNewOp<scf::YieldOp>(yieldOp, newYieldValues);
-  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)
@@ -692,54 +634,6 @@ MutableArrayRef<OpOperand> ForallOp::getInitsMutable() {
   return getOutputsMutable();
 }
 
-FailureOr<LoopLikeOpInterface>
-ForallOp::yieldTiledValuesAndReplace(RewriterBase &rewriter,
-                                     ValueRange newInitOperands,
-                                     YieldTiledValuesFn yieldTiledValuesFn) {
-  OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(getOperation());
-  auto inits = llvm::to_vector(getOutputs());
-  inits.append(newInitOperands.begin(), newInitOperands.end());
-  auto newLoop = rewriter.create<scf::ForallOp>(
-      getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
-      inits, getMapping(), [](OpBuilder &, Location, ValueRange) {});
-
-  // Move the region of the current block to the newly created op.
-  Block *newLoopBody = newLoop.getBody();
-  rewriter.mergeBlocks(
-      getBody(), newLoopBody,
-      newLoopBody->getArguments().take_front(getBody()->getNumArguments()));
-
-  auto terminator = cast<scf::InParallelOp>(newLoopBody->getTerminator());
-  rewriter.setInsertionPoint(terminator);
-  SmallVector<Value> tiledValues;
-  SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
-  ValueRange regionIterArgs =
-      newLoop.getRegionIterArgs().take_back(newInitOperands.size());
-  if (failed(yieldTiledValuesFn(rewriter, getLoc(), newLoop.getInductionVars(),
-                                regionIterArgs, tiledValues, resultOffsets,
-                                resultSizes))) {
-    return rewriter.notifyMatchFailure(getOperation(),
-                                       "failed to get yielded tiled values");
-  }
-
-  // Update the terminator.
-  rewriter.setInsertionPointToEnd(terminator.getBody());
-
-  for (auto [tiledValue, iterArg, resultOffset, resultSize] : llvm::zip_equal(
-           tiledValues, regionIterArgs, resultOffsets, resultSizes)) {
-    SmallVector<OpFoldResult> resultStride(resultOffset.size(),
-                                           rewriter.getIndexAttr(1));
-    rewriter.create<tensor::ParallelInsertSliceOp>(
-        terminator.getLoc(), tiledValue, iterArg, resultOffset, resultSize,
-        resultStride);
-  }
-
-  rewriter.replaceOp(getOperation(),
-                     newLoop->getResults().take_front(getNumResults()));
-  return cast<LoopLikeOpInterface>(newLoop.getOperation());
-}
-
 /// Promotes the loop body of a scf::ForallOp to its containing block.
 void mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp) {
   OpBuilder::InsertionGuard g(rewriter);
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -23,6 +23,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
@@ -287,6 +288,171 @@ static LogicalResult generateLoopNest(RewriterBase &rewriter, Location loc,
   return rewriter.notifyMatchFailure(loc, "unhandled loop type");
 }
 
+/// A function that allows returning additional yielded values during
+/// `yieldTiledValuesAndReplace`.
+/// - `ivs` induction variable for the loop.
+/// - `newBbArgs` basic block arguments corresponding to newly added iter_args.
+/// - `tiledValues` the tiled values to return. Must be of same size as
+///   `newbbArgs`, each element of this array is inserted into the corresponding
+///   element in `newbbArgs`.
+/// - `resultOffsets` is of the same size as `tiledValues` and represents
+///   the offsets to use when inserting corresponding element from `tiledValues`
+///   into the element from `newBbArgs`.
+/// - `resultSizes` is of the same size as `tiledValues` and represents
+///   the size of the corresponding element from `tiledValues` inserted into
+///   the element from `newBbArgs`.
+using YieldTiledValuesFn = llvm::function_ref<LogicalResult(
+    RewriterBase &rewriter, Location loc, ValueRange ivs, ValueRange newBbArgs,
+    SmallVector<Value> &tiledValues,
+    SmallVector<SmallVector<OpFoldResult>> &resultOffsets,
+    SmallVector<SmallVector<OpFoldResult>> &resultSizes)>;
+
+/// Append the specified additional `newInitOperands` operands to the
+/// loops existing `init` operands (or similar), and replace `loopOp` with
+/// the new loop that has the additional init operands. The loop body of
+/// this loop is moved over to the new loop. `yieldTiledValuesFn`
+/// is called to get the new tiled values returned, and the offset
+/// and sizes at which the tiled value is inserted into the
+/// new region iter_args that correspond to the newly added init operands.
+template <typename LoopType>
+FailureOr<LoopLikeOpInterface>
+yieldTiledValuesAndReplaceLoop(LoopType loopOp, RewriterBase &rewriter,
+                               ValueRange newInitOperands,
+                               YieldTiledValuesFn yieldTiledValuesFn) {
+  return rewriter.notifyMatchFailure(loopOp, "unhandled loop type");
+}
+
+/// Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.for`.
+template <>
+FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForOp>(
+    scf::ForOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands,
+    YieldTiledValuesFn yieldTiledValuesFn) {
+  OpBuilder::InsertionGuard g(rewriter);
+  Location loc = loopOp.getLoc();
+  rewriter.setInsertionPoint(loopOp);
+
+  auto inits = llvm::to_vector(loopOp.getInitArgs());
+  inits.append(newInitOperands.begin(), newInitOperands.end());
+  auto newLoop = rewriter.create<scf::ForOp>(
+      loc, loopOp.getLowerBound(), loopOp.getUpperBound(), loopOp.getStep(),
+      inits, [](OpBuilder &, Location, Value, ValueRange) {});
+
+  // Move the loop body to the new op.
+  Block *loopBody = loopOp.getBody();
+  Block *newLoopBody = newLoop.getBody();
+  rewriter.mergeBlocks(
+      loopBody, newLoopBody,
+      newLoopBody->getArguments().take_front(loopBody->getNumArguments()));
+
+  auto yieldOp = cast<scf::YieldOp>(newLoopBody->getTerminator());
+  rewriter.setInsertionPoint(yieldOp);
+
+  SmallVector<Value> tiledValues;
+  SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
+  ValueRange newRegionIterArgs =
+      newLoop.getRegionIterArgs().take_back(newInitOperands.size());
+  if (failed(yieldTiledValuesFn(rewriter, loc, newLoop.getInductionVar(),
+                                newRegionIterArgs, tiledValues, resultOffsets,
+                                resultSizes))) {
+    return rewriter.notifyMatchFailure(loopOp, "failed to get tiled values");
+  }
+
+  if (tiledValues.size() != resultOffsets.size() ||
+      tiledValues.size() != resultSizes.size()) {
+    return rewriter.notifyMatchFailure(
+        loopOp,
+        "expected number of tiled values returned, the number of offset "
+        "vectors and number of size vectors to be the same");
+  }
+
+  SmallVector<Value> newYieldValues = llvm::to_vector(yieldOp.getOperands());
+  for (auto [tiledValue, regionIterArg, resultOffset, resultSize] :
+       llvm::zip_equal(tiledValues, newRegionIterArgs, resultOffsets,
+                       resultSizes)) {
+    SmallVector<OpFoldResult> resultStride(resultOffset.size(),
+                                           rewriter.getIndexAttr(1));
+    Value insert = rewriter.create<tensor::InsertSliceOp>(
+        yieldOp->getLoc(), tiledValue, regionIterArg, resultOffset, resultSize,
+        resultStride);
+    newYieldValues.push_back(insert);
+  }
+
+  rewriter.replaceOpWithNewOp<scf::YieldOp>(yieldOp, newYieldValues);
+  rewriter.replaceOp(loopOp,
+                     newLoop->getResults().take_front(loopOp.getNumResults()));
+  return cast<LoopLikeOpInterface>(newLoop.getOperation());
+}
+
+/// Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.forall`
+template <>
+FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForallOp>(
+    scf::ForallOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands,
+    YieldTiledValuesFn yieldTiledValuesFn) {
+  OpBuilder::InsertionGuard g(rewriter);
+  Location loc = loopOp.getLoc();
+  rewriter.setInsertionPoint(loopOp);
+  auto inits = llvm::to_vector(loopOp.getOutputs());
+  inits.append(newInitOperands.begin(), newInitOperands.end());
+  auto newLoop = rewriter.create<scf::ForallOp>(
+      loc, loopOp.getMixedLowerBound(), loopOp.getMixedUpperBound(),
+      loopOp.getMixedStep(), inits, loopOp.getMapping(),
+      [](OpBuilder &, Location, ValueRange) {});
+
+  // Move the region of the current block to the newly created op.
+  Block *loopBody = loopOp.getBody();
+  Block *newLoopBody = newLoop.getBody();
+  rewriter.mergeBlocks(
+      loopBody, newLoopBody,
+      newLoopBody->getArguments().take_front(loopBody->getNumArguments()));
+
+  auto terminator = cast<scf::InParallelOp>(newLoopBody->getTerminator());
+  rewriter.setInsertionPoint(terminator);
+  SmallVector<Value> tiledValues;
+  SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
+  ValueRange regionIterArgs =
+      newLoop.getRegionIterArgs().take_back(newInitOperands.size());
+  if (failed(yieldTiledValuesFn(rewriter, loc, newLoop.getInductionVars(),
+                                regionIterArgs, tiledValues, resultOffsets,
+                                resultSizes))) {
+    return rewriter.notifyMatchFailure(loopOp,
+                                       "failed to get yielded tiled values");
+  }
+
+  // Update the terminator.
+  rewriter.setInsertionPointToEnd(terminator.getBody());
+
+  for (auto [tiledValue, iterArg, resultOffset, resultSize] : llvm::zip_equal(
+           tiledValues, regionIterArgs, resultOffsets, resultSizes)) {
+    SmallVector<OpFoldResult> resultStride(resultOffset.size(),
+                                           rewriter.getIndexAttr(1));
+    rewriter.create<tensor::ParallelInsertSliceOp>(
+        terminator.getLoc(), tiledValue, iterArg, resultOffset, resultSize,
+        resultStride);
+  }
+
+  rewriter.replaceOp(loopOp,
+                     newLoop->getResults().take_front(loopOp.getNumResults()));
+  return cast<LoopLikeOpInterface>(newLoop.getOperation());
+}
+
+/// Implementation of `yieldTiledValuesAndReplaceLoop` for
+/// `LoopLikeOpInterface`, that just dispatches to the implementation for each
+/// supported loop type.
+FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop(
+    LoopLikeOpInterface loopLikeOp, RewriterBase &rewriter,
+    ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn) {
+  return TypeSwitch<LoopLikeOpInterface, FailureOr<LoopLikeOpInterface>>(
+             loopLikeOp)
+      .Case<scf::ForOp, scf::ForallOp>(
+          [&](auto loopOp) -> FailureOr<LoopLikeOpInterface> {
+            return yieldTiledValuesAndReplaceLoop(
+                loopOp, rewriter, newInitOperands, yieldTiledValuesFn);
+          })
+      .Default([&](auto loopOp) -> FailureOr<LoopLikeOpInterface> {
+        return rewriter.notifyMatchFailure(loopOp, "unhandled loop type");
+      });
+}
+
 /// Method to add new init values to a loop nest. Updates `loops` in-place with
 /// new loops that use the `newInitValues`.
 /// The outer-loops are updated to yield the new result values of the inner
@@ -334,8 +500,8 @@ static LogicalResult addInitOperandsToLoopNest(
   // Update the loop body of the innermost loop to get new yield values.
   LoopLikeOpInterface innerMostLoop = loops.back();
   FailureOr<LoopLikeOpInterface> newInnerMostLoop =
-      innerMostLoop.yieldTiledValuesAndReplace(rewriter, newInitValues,
-                                               getNewTiledYieldsFn);
+      yieldTiledValuesAndReplaceLoop(innerMostLoop, rewriter, newInitValues,
+                                     getNewTiledYieldsFn);
 
   if (failed(newInnerMostLoop))
     return innerMostLoop.emitOpError("failed to return additional yields");
