[mlir][Linalg] Make SplitReduction use bufferize::AllocTensor

[nicolasvasilache]

Summary

SplitReduction is known to require a concrete additional buffer
to store tempoaray information.
Modify it to introduce a bufferization.alloc_tensor.
This behaves better with subset-based tiling and bufferization.

[nicolasvasilache]

Drop spurious files.

[nicolasvasilache]

Cleanup.

[nicolasvasilache]

Fix missing paren.

[mravishankar]

I am not sure introducing an allocation semantics op at this level makes sense for all use cases. AFAICS, the bufferization allocation op should only be introduced "close to bufferization" and that might not be the case for all use cases. Definitely is not the case for within IREE. At the very least it should be optional to use this op.

[nicolasvasilache]

Refactor transforms and make alloc_tensor optional.

[nicolasvasilache]

In D128722#3615830, @mravishankar wrote:

I am not sure introducing an allocation semantics op at this level makes sense for all use cases. AFAICS, the bufferization allocation op should only be introduced "close to bufferization" and that might not be the case for all use cases. Definitely is not the case for within IREE. At the very least it should be optional to use this op.

Turned it into an option,.

[nicolasvasilache]

Rebase.

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
   let description = [{
     Indicates that the given `target` op should be transformed with the 
     `splitReduction` transformation and split factor provided as attribute.
 
     The `splitReduction` transformation splits the first single linalg op 
     reduction into a parallel and reduction dimension. 
     A new `linalg.generic` op is created to perform the rest of the reduction. 
     
-    Example:
-    
+    The transformation supports different configurations attributes:
+      - split_factor: the factor by which to split (i.e. the size of the 
+        remaining reduction after splitting).
+      - insert_split_dimension: the dimension in the temporary tensor into 
+        which the new parallel dimension is inserted.
+      - use_scaling_algorithm: whether to use a scaling based formulation that 
+        does not create an ExpandShapeOp (default: do not use scaling)
+      - use_alloc: whether to use an alloc op to allocate the temporary 
+        tensor (default: do not use alloc op)
+
+    This op returns 4 handles to:
+      - the init op (or tensor_alloc op if use_alloc = true), 
+      - the fill op used to initialize the neutral element, 
+      - the split op and 
+      - the result-combining op.
+    Example (default: use_scaling_algorithm = false, use_alloc = false):
+    ====================================================================
     ```
       %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
                                             affine_map<(d0) -> ()>],
             iterator_types = ["reduction"]}
       ins(%in : tensor<32xf32>)
       outs(%out : tensor<f32>) {
       ^bb0(%arg1: f32, %arg2: f32):
         %y = arith.addf %arg1, %arg2 : f32
         linalg.yield %y : f32
       } -> tensor<f32>
     ```
     
-    To:
+    is split into:
     
     ```
       %cst = arith.constant 0.000000e+00 : f32
       %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
       %1 = linalg.init_tensor [4] : tensor<4xf32>
       %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32>
       %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                                             affine_map<(d0, d1) -> (d0)>],
         iterator_types = ["parallel", "reduction"]}
         ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
         ^bb0(%arg3: f32, %arg5: f32):
         %5 = arith.addf %arg3, %arg4 : f32
         linalg.yield %5 : f32
       } -> tensor<4xf32>
       %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
                                             affine_map<(d0) -> ()>],
         iterator_types = ["reduction"]}
         ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
         ^bb0(%arg3: f32, %arg4: f32):
         %5 = arith.addf %arg3, %arg4 : f32
         linalg.yield %5 : f32
       } -> tensor<f32>
     ```
 
-    This op returns handles to the fill op used to initialize the neutral 
-    element, the split op and the result-combining op.
-  }];
-
-  let arguments = (ins PDL_Operation:$target,
-                   DefaultValuedAttr<I64Attr, "{}">:$split_factor,
-                   DefaultValuedAttr<I64Attr, "{}">:$insert_split_dimension);
-  let results = (outs PDL_Operation:$fill_op,
-                      PDL_Operation:$split_linalg_op,
-                      PDL_Operation:$combining_linalg_op);
-
-  let assemblyFormat = "$target attr-dict";
-
-  let extraClassDeclaration = [{
-    ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
-        ::mlir::linalg::LinalgOp target, TransformState &state);
-  }];
-}
-
-def SplitReductionByScalingOp : 
-  Op<Transform_Dialect, "structured.split_reduction_by_scaling",
-       [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-        TransformEachOpTrait, TransformOpInterface]> {
-  let description = [{
-    Indicates that the given `target` op should be transformed with the 
-    `splitReductionByScaling` transformation and split factor provided as 
-    attribute.
-
+    Example (use_scaling_algorithm = true, use_alloc = true):
+    =========================================================
     Instead of introducing an ExpandShapeOp, this scaling-based implementation 
     rewrites a reduction dimension `k` into `k * split_factor + kk`.
     The dimension `kk` is added as an extra parallel dimension to the 
     intermediate output tensor at position `insert_split_dimension`.
 
     Consider a minimal example where `k` is reduced: 
         O(i, j) += I(i, j, k)
     Assume i=3, j=5, k=128, split_factor=16 and insert_split_dimension=0.
%4 = linalg.generic {indexing_maps = [#map4, #map5],
        outs(%C : tensor<16x32xf32>) {
          ^bb0(%arg3: f32, %arg4: f32):
            %5 = arith.addf %arg3, %arg4 : f32
            linalg.yield %5 : f32
      } -> tensor<16x32xf32>
 
      return %4 : tensor<16x32xf32>
     ```
-
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64Attr, "{}">:$split_factor,
-                   DefaultValuedAttr<I64Attr, "{}">:$insert_split_dimension);
+                   DefaultValuedAttr<I64Attr, "{}">:$insert_split_dimension,
+                   UnitAttr:$use_scaling_algorithm,
+                   UnitAttr:$use_alloc);
   let results = (outs PDL_Operation:$fill_op,
                       PDL_Operation:$split_linalg_op,
                       PDL_Operation:$combining_linalg_op);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 // TODO: don't use unsigned unless doing bit manipulation.
 using ControlSplitReductionFn =
     std::function<std::pair<int64_t, unsigned>(LinalgOp op)>;
 
 /// Patterns to apply `splitReduction` below.
 void populateSplitReductionPattern(
     RewritePatternSet &patterns,
     const ControlSplitReductionFn &controlSplitReductionFn,
-    const LinalgTransformationFilter &f = LinalgTransformationFilter());
+    const LinalgTransformationFilter &f = LinalgTransformationFilter(),
+    bool useAlloc = false);
 
 /// Apply transformation to split the single linalg op reduction into a parallel
 /// and reduction dimension. Then create a new linalg.generic op doing the rest
 /// of the reduction. Return the new linalg op with an extra parallel dimension
 /// or failure if the transformation didn't happen.
 /// Example:
 /// ```
 ///  %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,

 ///   ^bb0(%arg3: f32, %arg4: f32):
 ///   %5 = arith.addf %arg3, %arg4 : f32
 ///   linalg.yield %5 : f32
 /// } -> tensor<f32>
 /// ```
 FailureOr<LinalgOp>
 splitReduction(PatternRewriter &b, LinalgOp op,
                const ControlSplitReductionFn &controlSplitReductionFn,
-               const LinalgTransformationFilter &f);
+               const LinalgTransformationFilter &f, bool useAlloc = false);
 
 /// Filterless version of the above.
 /// Returns both the new linalg ops as well as the fillOp needed to initialize
 /// the temporary expanded tensor with the proper neutral element.
 struct SplitReductionResult {
+  Operation *initOrAlloc;
   FillOp fillOp;
   LinalgOp splitLinalgOp;
   LinalgOp resultCombiningLinalgOp;
 };
 FailureOr<SplitReductionResult>
 splitReduction(PatternRewriter &b, LinalgOp op,
-               const ControlSplitReductionFn &controlSplitReductionFn);
+               const ControlSplitReductionFn &controlSplitReductionFn,
+               bool useAlloc = false);
 
 /// Scaling-based implementation of the split reduction transformation.
 /// Instead of introducing an ExpandShapeOp, this rewrites a reduction dimension
 /// `k` into `k * scale + kk`.
 ///
 /// Example:
 /// ```
 ///  %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)

 ///        %5 = arith.addf %arg3, %arg4 : f32
 ///        linalg.yield %5 : f32
 ///  } -> tensor<16x32xf32>
 ///
 ///  return %4 : tensor<16x32xf32>
 /// ```
 FailureOr<SplitReductionResult>
 splitReductionByScaling(PatternRewriter &b, LinalgOp op,
-                        const ControlSplitReductionFn &controlSplitReductionFn);
+                        const ControlSplitReductionFn &controlSplitReductionFn,
+                        bool useAlloc = false);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::SplitReductionOp::applyToOne(LinalgOp target,
                                         TransformState &state) {
   ControlSplitReductionFn splitFn = [&](LinalgOp) {
     return std::pair<int64_t, unsigned>(getSplitFactor(),
                                         getInsertSplitDimension());
   };
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<SplitReductionResult> splitResult =
-      splitReduction(rewriter, target, splitFn);
-  if (failed(splitResult))
-    return getOperation()->emitError("failed to apply");
-  return SmallVector<Operation *>{splitResult->fillOp,
-                                  splitResult->splitLinalgOp,
-                                  splitResult->resultCombiningLinalgOp};
-}
-
-//===----------------------------------------------------------------------===//
-// SplitReductionByScalingOp
-//===----------------------------------------------------------------------===//
-
-FailureOr<SmallVector<Operation *>>
-transform::SplitReductionByScalingOp::applyToOne(LinalgOp target,
-                                                 TransformState &state) {
-  ControlSplitReductionFn splitFn = [&](LinalgOp) {
-    return std::pair<int64_t, unsigned>(getSplitFactor(),
-                                        getInsertSplitDimension());
-  };
-  SimpleRewriter rewriter(getContext());
-  rewriter.setInsertionPoint(target);
-  FailureOr<SplitReductionResult> splitResult =
-      splitReductionByScaling(rewriter, target, splitFn);
+      (getUseScalingAlgorithm())
+          ? splitReductionByScaling(rewriter, target, splitFn, getUseAlloc())
+          : splitReduction(rewriter, target, splitFn, getUseAlloc());
   if (failed(splitResult))
     return getOperation()->emitError("failed to apply");
   return SmallVector<Operation *>{splitResult->fillOp,
                                   splitResult->splitLinalgOp,
                                   splitResult->resultCombiningLinalgOp};
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp

 // between a parallel and a reduction dimension.
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
static Attribute getNeutralElement(Operation *op) {
   if (isa<arith::MulIOp>(op))
     return b.getIntegerAttr(resultType, 1);
   return Attribute();
 }
 
 FailureOr<LinalgOp> mlir::linalg::splitReduction(
     PatternRewriter &b, LinalgOp op,
     const ControlSplitReductionFn &controlSplitReductionFn,
-    const LinalgTransformationFilter &filter) {
+    const LinalgTransformationFilter &filter, bool useAlloc) {
   if (failed(filter.checkAndNotify(b, op)) || !op.hasTensorSemantics() ||
       op.getNumReductionLoops() != 1 || op.getNumOutputs() != 1 ||
       !op.hasOnlyProjectedPermutations())
     return b.notifyMatchFailure(op, "precondition not met");
 
   FailureOr<SplitReductionResult> res =
-      splitReduction(b, op, controlSplitReductionFn);
+      splitReduction(b, op, controlSplitReductionFn, useAlloc);
   if (failed(res))
     return failure();
 
   filter.replaceLinalgTransformationFilter(b, res->splitLinalgOp);
   filter.replaceLinalgTransformationFilter(b, res->resultCombiningLinalgOp);
 
   return res->splitLinalgOp;
 }
 
 FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
     PatternRewriter &b, LinalgOp op,
-    const ControlSplitReductionFn &controlSplitReductionFn) {
+    const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc) {
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPoint(op);
 
   std::pair<int64_t, unsigned> control = controlSplitReductionFn(op);
   int64_t ratio = control.first;
   unsigned insertSplitDimension = control.second;
   if (ratio <= 1)
     return b.notifyMatchFailure(op, "split ratio needs to be greater than 1");
if (idx == insertSplitDimension) {
       continue;
     }
     unsigned oldDim = idx < insertSplitDimension ? idx : idx - 1;
     newOutputShape.push_back(oldShape[oldDim]);
     unsigned dim = oldOutputMap.getDimPosition(oldDim);
     outputExpr.push_back(
         b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1));
   }
-  Value initTensor = b.create<linalg::InitTensorOp>(
+  Value initOrAllocTensor;
+  if (useAlloc) {
+    initOrAllocTensor = b.create<bufferization::AllocTensorOp>(
+        loc,
+        RankedTensorType::get(newOutputShape,
+                              op.getRegionOutputArgs()[0].getType()),
+        ValueRange{});
+  } else {
+    initOrAllocTensor = b.create<linalg::InitTensorOp>(
         loc, newOutputShape, op.getRegionOutputArgs()[0].getType());
+  }
   Value constantOp = b.create<arith::ConstantOp>(loc, identity);
   Value identityTensor =
-      b.create<linalg::FillOp>(op->getLoc(), constantOp, initTensor)
+      b.create<linalg::FillOp>(op->getLoc(), constantOp, initOrAllocTensor)
           .getResult(0);
 
   newMaps.push_back(AffineMap::get(oldOutputMap.getNumDims() + 1, 0, outputExpr,
                                    op.getContext()));
   SmallVector<StringRef> newIteratorTypes;
   for (auto &it : llvm::enumerate(op.iterator_types())) {
     if (insertSplitDimension == it.index())
       newIteratorTypes.push_back(getParallelIteratorTypeName());
     newIteratorTypes.push_back(it.value().cast<StringAttr>().getValue());
   }
   // Create the new op matching the original op with an extra parallel
   // dimension.
   GenericOp genericOp = b.create<GenericOp>(
-      loc, TypeRange({initTensor.getType()}), newInputs,
+      loc, TypeRange({initOrAllocTensor.getType()}), newInputs,
       ValueRange({identityTensor}), newMaps, newIteratorTypes);
   b.inlineRegionBefore(op->getRegion(0), genericOp.region(),
                        genericOp.region().begin());
 
   // Then create a new reduction that only reduce the newly added dimension
   // from the previous op.
   unsigned intermRank = newOutputShape.size();
   AffineMap inputMap = b.getMultiDimIdentityMap(intermRank);
auto reduction = b.create<GenericOp>(
       [reductionOp](OpBuilder &b, Location loc, ValueRange inputs) {
         Operation *clonedReductionOp = b.clone(*reductionOp);
         clonedReductionOp->setOperand(0, inputs[0]);
         clonedReductionOp->setOperand(1, inputs[1]);
         b.create<linalg::YieldOp>(loc, clonedReductionOp->getResult(0));
       });
   b.replaceOp(op, reduction.getResults());
 
-  return SplitReductionResult{identityTensor.getDefiningOp<FillOp>(),
-                              cast<LinalgOp>(genericOp.getOperation()),
-                              reduction};
+  return SplitReductionResult{
+      initOrAllocTensor.getDefiningOp(), identityTensor.getDefiningOp<FillOp>(),
+      cast<LinalgOp>(genericOp.getOperation()), reduction};
 }
 
 /// Rewrite f(i, j, k, ...) into f(i, j, k * ratio + kk, ...)
 /// TODO: Additional pattern to rewrite f(i, j, k * ratio + kk, ...) into
 /// f(i, j, k, kk, ...) with a proper ExpandShapeOp. This is probably better
 /// done as a transform to enable better vectorization.
 static AffineMap scaleReductionDim(LinalgOp op, OpOperand &opOperand,
                                    unsigned reductionDimPos,
AffineMap idMap =
       AffineMap::getMultiDimIdentityMap(map.getNumDims(), op.getContext());
   AffineMap shiftedIdMap = idMap.shiftDims(1, /*offset=*/reductionDimPos + 1);
   return map.compose(shiftedIdMap).insertResult(reductionDim, reductionDimPos);
 }
 
 /// Core rewrite implementation.
 FailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
     PatternRewriter &b, LinalgOp op,
-    const ControlSplitReductionFn &controlSplitReductionFn) {
+    const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc) {
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPoint(op);
 
   // Matcher part, enforce preconditions.
   std::pair<int64_t, unsigned> control = controlSplitReductionFn(op);
   int64_t splitFactor = control.first;
   unsigned insertSplitDimension = control.second;
   if (splitFactor <= 1)
FailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
   SmallVector<Attribute> neutralElements = llvm::to_vector<4>(
       llvm::map_range(combinerOps, [&](Operation *reductionOp) {
         return getNeutralElement(reductionOp);
       }));
   if (!llvm::all_of(neutralElements, [](Attribute attr) { return attr; }))
     return b.notifyMatchFailure(op, "unknown reduction neutral");
 
   // TODO: relax this when multi-reduction support is available.
-  if (op.getNumOutputs() != (int)neutralElements.size())
+  if (op.getNumOutputs() != static_cast<int64_t>(neutralElements.size()))
     return b.notifyMatchFailure(op, "expect one reduction per output");
 
   // Rewrite part.
   // Step 1. Build the intermediate outputs filled with the proper
   // neutralElements. Such outputs are of the same shape with an extra dimension
   // inserted at `insertSplitDimension`.
   //
   // Consider a minimal example where `k` is reduced:
   //     O(i, j) += I(i, j, k)
   // Assume i=3, j=5, k=128, splitFactor=16 and insertSplitDimension=0.
   // The compute is rewritten as:
   //   a. O_i(kk, i, j) += I(i, j, 16 * k + kk)
   //   b. O(i, j) += O_i(kk, i, j)
   // The intermediate tensor O_i is of shape (128/16)x3x5 == 8x3x5.
   Location loc = op->getLoc();
   MLIRContext *context = op.getContext();
   // For now assume outputs are 1-1 with reduction neutralElements.
   // TODO: generalize when multi-reduction support is available.
   SmallVector<Value> newOutputs;
   newOutputs.reserve(op.getNumOutputs());
+  SmallVector<Operation *> initOrAllocTensorOps;
   SmallVector<linalg::FillOp> fillOps;
   fillOps.reserve(op.getNumOutputs());
   for (auto it : llvm::zip(op.outputs(), neutralElements)) {
     Value rankedTensor = std::get<0>(it);
     auto t = rankedTensor.getType().cast<RankedTensorType>();
     RankedTensorType newT = RankedTensorType::Builder(t).insertDim(
         reductionDimSize / splitFactor, insertSplitDimension);
     SmallVector<Value> dims =
         tensor::createDynamicDimValues(b, loc, rankedTensor);
-    Value initTensor = b.create<linalg::InitTensorOp>(
+    Value initOrAllocTensor;
+    if (useAlloc) {
+      initOrAllocTensor =
+          b.create<bufferization::AllocTensorOp>(loc, newT, dims);
+    } else {
+      initOrAllocTensor = b.create<linalg::InitTensorOp>(
           loc, dims, newT.getShape(), t.getElementType());
+    }
     Value constantOp = b.create<arith::ConstantOp>(loc, std::get<1>(it));
     fillOps.push_back(
-        b.create<linalg::FillOp>(op->getLoc(), constantOp, initTensor));
+        b.create<linalg::FillOp>(op->getLoc(), constantOp, initOrAllocTensor));
     newOutputs.push_back(fillOps.back().getResult(0));
+    initOrAllocTensorOps.push_back(initOrAllocTensor.getDefiningOp());
   }
 
   // Step 2. Reindex / expand indexing maps.
   // Reindex existing input indexings: k -> k * splitFactor + k'.
   SmallVector<AffineMap> newMaps;
   newMaps.reserve(op.getNumInputsAndOutputs() + 1);
   for (OpOperand *o : op.getInputOperands())
     newMaps.push_back(scaleReductionDim(op, *o, reductionDimPos, splitFactor));
for (auto it :
     // clang-format on
 
     results.push_back(reductionOp);
   }
 
   // TODO: extend when multi-reduction support is available.
   assert(fillOps.size() == results.size() && results.size() == 1);
   b.replaceOp(op, results.front()->getResults());
-  return SplitReductionResult{fillOps.front(),
+  return SplitReductionResult{initOrAllocTensorOps.front(), fillOps.front(),
                               cast<LinalgOp>(genericOp.getOperation()),
                               results.front()};
 }
 
 namespace {
 
 struct LinalgSplitReduction : public OpInterfaceRewritePattern<LinalgOp> {
   /// Construct a generic pattern applied to all LinalgOp that verify `filter`.
   LinalgSplitReduction(MLIRContext *context,
                        ControlSplitReductionFn controlSplitReductionFn,
-                       LinalgTransformationFilter f, PatternBenefit benefit = 1)
+                       LinalgTransformationFilter f, bool useAlloc = false,
+                       PatternBenefit benefit = 1)
       : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
         controlSplitReductionFn(std::move(controlSplitReductionFn)),
-        filter(std::move(f)) {}
+        useAlloc(useAlloc), filter(std::move(f)) {}
 
   LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    return splitReduction(rewriter, op, controlSplitReductionFn, filter);
+    return splitReduction(rewriter, op, controlSplitReductionFn, filter,
+                          useAlloc);
   }
 
 private:
   ControlSplitReductionFn controlSplitReductionFn;
+  bool useAlloc;
   LinalgTransformationFilter filter;
 };
 
 } // namespace
 
 void linalg::populateSplitReductionPattern(
     RewritePatternSet &patterns,
     const ControlSplitReductionFn &controlSplitReductionFn,
-    const LinalgTransformationFilter &f) {
+    const LinalgTransformationFilter &f, bool useAlloc) {
   patterns.add<LinalgSplitReduction>(patterns.getContext(),
-                                     controlSplitReductionFn, f);
+                                     controlSplitReductionFn, f, useAlloc);
 }

mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir

 // RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
 
 // CHECK-LABEL: func.func @matmul_split
 func.func @matmul_split(%A : tensor<?x256xf32>, %B: tensor<256x32xf32>, %C: tensor<?x32xf32>) -> tensor<?x32xf32> {
 
+  //      CHECK: bufferization.alloc_tensor({{.*}}) : tensor<?x32x64xf32>
   //      CHECK: linalg.generic 
   // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]
   // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}} : tensor<?x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>)
   // CHECK-SAME: outs(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>) {
 
   //      CHECK: linalg.generic 
   // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
   // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>)
pdl.pattern @pdl_target : benefit(1) {
     %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
     // TODO: we don't want this, but it is the required terminator for pdl.pattern
     rewrite %0 with "transform.dialect"
   }
 
   transform.sequence %arg0 {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @pdl_target in %arg1
-    %1:3 = transform.structured.split_reduction_by_scaling %0 { split_factor = 4, insert_split_dimension = 2}
+    %1:3 = transform.structured.split_reduction %0 
+      { split_factor = 4, insert_split_dimension = 2, use_scaling_algorithm, use_alloc}
   }
 }

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
struct TestLinalgTransforms
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgTransforms)
 
   TestLinalgTransforms() = default;
   TestLinalgTransforms(const TestLinalgTransforms &pass) : PassWrapper(pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     // clang-format off
     registry.insert<AffineDialect,
+                    bufferization::BufferizationDialect,
                     memref::MemRefDialect,
                     scf::SCFDialect,
                     linalg::LinalgDialect,
                     vector::VectorDialect,
                     gpu::GPUDialect>();
     // clang-format on
   }
   StringRef getArgument() const final {