Try not using llvm::function_ref.

Author: MaheshRavishankar

mlir/include/mlir/Interfaces/LoopLikeInterface.h

@@ -22,31 +22,9 @@ class RewriterBase;
 /// `replaceWithAdditionalYields`. `newBbArgs` are the newly added region
 /// iter_args. This function should return as many values as there are block
 /// arguments in `newBbArgs`.
-using NewYieldValuesFn = llvm::function_ref<SmallVector<Value>(
+using NewYieldValuesFn = std::function<SmallVector<Value>(
     OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs)>;
 
-/// 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`.
-/// - `resultStrides` is of the same size as `tiledValues` and represents
-///   the strides to use when inserting corresponding element from `tiledValues`
-///   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)>;
-
 namespace detail {
 /// Verify invariants of the LoopLikeOpInterface.
 LogicalResult verifyLoopLikeOpInterface(Operation *op);

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -99,6 +99,25 @@ static OpFoldResult getBoundedTileSize(OpBuilder &b, Location loc,
       b, loc, minMap, SmallVector<OpFoldResult>{iv, tileSize, size});
 }
 
+/// 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 = std::function<LogicalResult(
+    RewriterBase &rewriter, Location loc, ValueRange ivs, ValueRange newBbArgs,
+    SmallVector<Value> &tiledValues,
+    SmallVector<SmallVector<OpFoldResult>> &resultOffsets,
+    SmallVector<SmallVector<OpFoldResult>> &resultSizes)>;
+
 /// Clones the operation and updates the destination if the operation
 /// implements the `DestinationStyleOpInterface`.
 static Operation *cloneOpAndUpdateDestinationArgs(RewriterBase &rewriter,
@@ -288,25 +307,6 @@ 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
@@ -441,8 +441,8 @@ FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForallOp>(
 FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop(
     LoopLikeOpInterface loopLikeOp, RewriterBase &rewriter,
     ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn) {
-  return TypeSwitch<LoopLikeOpInterface, FailureOr<LoopLikeOpInterface>>(
-             loopLikeOp)
+  return TypeSwitch<Operation *, FailureOr<LoopLikeOpInterface>>(
+             loopLikeOp.getOperation())
       .Case<scf::ForOp, scf::ForallOp>(
           [&](auto loopOp) -> FailureOr<LoopLikeOpInterface> {
             return yieldTiledValuesAndReplaceLoop(
@@ -460,7 +460,7 @@ FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop(
 /// the additional values to yield form the innermost loop.
 static LogicalResult addInitOperandsToLoopNest(
     RewriterBase &rewriter, MutableArrayRef<LoopLikeOpInterface> loops,
-    ValueRange newInitValues, const YieldTiledValuesFn &getNewTiledYieldsFn) {
+    ValueRange newInitValues, YieldTiledValuesFn getNewTiledYieldsFn) {
   SmallVector<scf::ForOp> newLoops;
   if (loops.empty())
     return success();