[mlir][Linalg] Split reduction transform op

[nicolasvasilache]

Summary

This revision separates the LinalgSplitReduction pattern, whose application is based on attributes,
from its implementation.
A transform dialect op extension is added to control the application of the transformation at a finer granularity.

[nicolasvasilache]

Add CHECKs to test

[ftynse]

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

161

I'd expand a bit on what the transformation does.

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

404

Nit: can't one just omit the name?

410

Payload ops shouldn't be null, if they are, this sounds like a bug. Note that .front() will fail on the empty list, not return null.

412

Can't we just make it apply to each operand via the TransformEachOpTrait?

412

Otherwise, please add a test for this user-visible error message.

423

Nit: I'd rather do getSplitLinalgOp().cast<OpResult>()

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

68

Nit: expand auto

110

Not for this diff: it would be nice to have this reported through DiagnosedSilenceableFailure. I don't know if we can connect the rewriter to that.

[nicolasvasilache]

Implementation that does not require ExpandShapeOp.

[nicolasvasilache]

It turns out that the transformation had a few issues, in particular related to dynamic dimensions and somewhat hardcoded assumptions.
I spent some time writing a version that does not require ExpandShapeOp and that instead rewrites the dimension with a "convolution" flavor.
This has certain advantages over the ExpandShapeOp version for some of the work I am doing but may be generally undesirable for @ThomasRaoux 's use cases.

Let's see what Thomas thinks of this version.

[nicolasvasilache]

Add more comments to the test.

[nicolasvasilache]

Address review.

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

412

Extended TransformEachOpTrait to work with multi-result.

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

110

Ack

[nicolasvasilache]

In D128165#3596813, @nicolasvasilache wrote:

It turns out that the transformation had a few issues, in particular related to dynamic dimensions and somewhat hardcoded assumptions.
I spent some time writing a version that does not require ExpandShapeOp and that instead rewrites the dimension with a "convolution" flavor.
This has certain advantages over the ExpandShapeOp version for some of the work I am doing but may be generally undesirable for @ThomasRaoux 's use cases.

Let's see what Thomas thinks of this version.

Reverted back to an NFC PR from the point of view of the transformation itself.
I'll send the proposed "SplitReduction without ExpanShapeOp" separately.

[ftynse]

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

763–764

Something is broken with this signature

792

This looks fishy

809–811

I'd do something like SmallVector<SmallVector<Operation *>, 1> here to avoid over-allocating on stack in the common case of there being only one result.

818–838

const &, we don't want to copy here

831

Ditto

[nicolasvasilache]

Add test

[nicolasvasilache]

Avoid spurious deletions.

[nicolasvasilache]

Fix build

[nicolasvasilache]

Rebase

[nicolasvasilache]

Fix comment

[nicolasvasilache]

I see a bunch of spurious stderr: 'fatal: Unable to create '/var/lib/buildkite-agent/builds/llvm-project-fork/.git/index.lock': File exists. in the build bots.
Tests pass locally, landing as is.

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

def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::FailureOr<::mlir::linalg::LinalgOp> applyToOne(
         ::mlir::linalg::LinalgOp target);
   }];
 }
 
+def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
+       [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+        TransformEachOpTrait, TransformOpInterface]> {
+  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:
+    
+    ```
+      %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:
+    
+    ```
+      %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);
+  }];
+}
+
 def TileOp : Op<Transform_Dialect, "structured.tile",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
   let description = [{
     Indicates that the given `target` op should be tiled with the options
     provided as attributes. This transform generates a loop nest with a smaller
     ("tiled") target operation in its body. Currently limited to LinalgOps.
 

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

public:
   }
 };
 
 /// Function signature to control reduction splitting. This returns a pair
 /// containing a ratio and a dimension index. The ratio is used to split the
 /// reduction dimension. The dimension index is used to control where the extra
 /// dimension is added to the intermediate tensor shape. If the ratio value is
 /// less or equal to 1 then nothing will be done.
+// 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());

 ///   linalg.yield %5 : f32
 /// } -> tensor<f32>
 /// ```
 FailureOr<LinalgOp>
 splitReduction(PatternRewriter &b, LinalgOp op,
                const ControlSplitReductionFn &controlSplitReductionFn,
                const LinalgTransformationFilter &f);
 
+/// 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 {
+  FillOp fillOp;
+  LinalgOp splitLinalgOp;
+  LinalgOp resultCombiningLinalgOp;
+};
+FailureOr<SplitReductionResult>
+splitReduction(PatternRewriter &b, LinalgOp op,
+               const ControlSplitReductionFn &controlSplitReductionFn);
+
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

if (diagnostic) {
       diagnostic.reset();
       result = ::mlir::success();
     }
     return result;
   }
 
   /// Streams the given values into the diagnotic. Expects this object to be a
   /// silencable failure.
-  template <typename T> DiagnosedSilenceableFailure &operator<<(T &&value) & {
+  template <typename T>
+  DiagnosedSilenceableFailure &operator<<(T &&value) & {
     assert(isSilenceableFailure() &&
            "can only append output in silencable failure state");
     *diagnostic << std::forward<T>(value);
     return *this;
   }
-  template <typename T> DiagnosedSilenceableFailure &&operator<<(T &&value) && {
+  template <typename T>
+  DiagnosedSilenceableFailure &&operator<<(T &&value) && {
     return std::move(this->operator<<(std::forward<T>(value)));
   }
 
   /// Attaches a note to the diagnostic. Expects this object to be a silencable
   /// failure.
   Diagnostic &attachNote(Optional<Location> loc = llvm::None) {
     assert(isSilenceableFailure() &&
            "can only attach notes to silencable failures");
LogicalResult mapBlockArguments(TransformState &state) {
     assert(
         this->getOperation()->getNumRegions() == 1 &&
         "must indicate the region to map if the operation has more than one");
     return mapBlockArguments(state, this->getOperation()->getRegion(0));
   }
 };
 
 /// Trait implementing the TransformOpInterface for operations applying a
-/// transformation to a single operation handle and producing a single operation
-/// handle. The op must implement a method with one of the following signatures:
+/// transformation to a single operation handle and producing one or multiple
+/// operation handles.
+/// The op must implement a method with one of the following signatures:
 ///   - FailureOr<convertible-to-Operation*> applyToOne(OpTy)
+///   - FailureOr<SmallVector<convertible-to-Operation*>> applyToOne(OpTy)
 ///   - LogicalResult applyToOne(OpTy)
 /// to perform a transformation that is applied in turn to all payload IR
 /// operations that correspond to the handle of the transform IR operation.
 /// In the functions above, OpTy is either Operation * or a concrete payload IR
 /// Op class that the transformation is applied to (NOT the class of the
-/// transform IR op). The op is expected to have one operand and zero or one
-/// results.
+/// transform IR op). The op is expected to have a single operand.
 template <typename OpTy>
 class TransformEachOpTrait
     : public OpTrait::TraitBase<OpTy, TransformEachOpTrait> {
 public:
   /// Calls `applyToOne` for every payload operation associated with the operand
   /// of this transform IR op. If `applyToOne` returns ops, associates them with
   /// the result of this transform op.
   DiagnosedSilenceableFailure apply(TransformResults &transformResults,

 
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h.inc"
 
 namespace mlir {
 namespace transform {
 namespace detail {
 /// Appends `result` to the vector assuming it corresponds to the success state
 /// in `FailureOr<convertible-to-Operation*>`. If `result` is just a
-/// `LogicalResult`, does nothing.
+/// `LogicalResult`, appends an empy vector.
 template <typename Ty>
 std::enable_if_t<std::is_same<Ty, LogicalResult>::value, LogicalResult>
-appendTransformResultToVector(Ty result,
-                              SmallVectorImpl<Operation *> &results) {
+appendTransformResultToVector(
+    Ty result, SmallVectorImpl<SmallVector<Operation *>> &results) {
+  results.push_back(SmallVector<Operation *>());
   return result;
 }
+
 template <typename Ty>
-std::enable_if_t<!std::is_same<Ty, LogicalResult>::value, LogicalResult>
-appendTransformResultToVector(Ty result,
-                              SmallVectorImpl<Operation *> &results) {
-  static_assert(
-      std::is_convertible<typename Ty::value_type, Operation *>::value,
-      "expected transform function to return operations");
+std::enable_if_t<
+    llvm::conjunction<
+        llvm::negation<std::is_same<Ty, LogicalResult>>,
+        std::is_convertible<typename Ty::value_type, Operation *>>::value,
+    LogicalResult>
+appendTransformResultToVector(
+    Ty result, SmallVectorImpl<SmallVector<Operation *>> &results) {
   if (failed(result))
     return failure();
-
-  results.push_back(*result);
+  results.push_back(SmallVector<Operation *>{*result});
   return success();
 }
 
-/// Applies a one-to-one transform to each of the given targets. Puts the
-/// results of transforms, if any, in `results` in the same order. Fails if any
-/// of the application fails. Individual transforms must be callable with
-/// one of the following signatures:
+template <typename ContainerTy>
+std::enable_if_t<
+    llvm::conjunction<
+        llvm::negation<std::is_same<ContainerTy, LogicalResult>>,
+        llvm::negation<std::is_convertible<typename ContainerTy::value_type,
+                                           Operation *>>>::value,
+    LogicalResult>
+appendTransformResultToVector(
+    ContainerTy resultContainer,
+    SmallVectorImpl<SmallVector<Operation *>> &results) {
+  if (failed(resultContainer))
+    return failure();
+  results.push_back(*resultContainer);
+  return success();
+}
+/// Applies a one-to-one or a one-to-many transform to each of the given
+/// targets. Puts the results of transforms, if any, in `results` in the same
+/// order. Fails if any of the application fails. Individual transforms must be
+/// callable with one of the following signatures:
 ///   - FailureOr<convertible-to-Operation*>(OpTy)
 ///   - LogicalResult(OpTy)
+///   - FailureOr<SmallVectorImpl<convertible-to-Operation*>>(
+///       SmallVectorImpl<OpTy>)
+///   - LogicalResult(SmallVectorImpl<OpTy>)
 /// where OpTy is either
 ///   - Operation *, in which case the transform is always applied;
 ///   - a concrete Op class, in which case a check is performed whether
 ///   `targets` contains operations of the same class and a silencable failure
 ///   is reported if it does not.
 template <typename FnTy>
 DiagnosedSilenceableFailure
 applyTransformToEach(ArrayRef<Operation *> targets,
-                     SmallVectorImpl<Operation *> &results, FnTy transform) {
+                     SmallVectorImpl<SmallVector<Operation *>> &results,
+                     FnTy transform) {
   using OpTy = typename llvm::function_traits<FnTy>::template arg_t<0>;
   static_assert(std::is_convertible<OpTy, Operation *>::value,
                 "expected transform function to take an operation");
   using RetTy = typename llvm::function_traits<FnTy>::result_t;
   static_assert(std::is_convertible<RetTy, LogicalResult>::value,
                 "expected transform function to return LogicalResult or "
                 "FailureOr<convertible-to-Operation*>");
   for (Operation *target : targets) {
     auto specificOp = dyn_cast<OpTy>(target);
     if (!specificOp) {
       Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error);
       diag << "attempted to apply transform to the wrong op kind";
       return DiagnosedSilenceableFailure::silencableFailure(std::move(diag));
     }
 
     auto result = transform(specificOp);
     if (failed(appendTransformResultToVector(result, results)))
       return DiagnosedSilenceableFailure::definiteFailure();
   }
   return DiagnosedSilenceableFailure::success();
 }
 } // namespace detail
 } // namespace transform
 } // namespace mlir
 
 template <typename OpTy>
 mlir::DiagnosedSilenceableFailure
 mlir::transform::TransformEachOpTrait<OpTy>::apply(
     TransformResults &transformResults, TransformState &state) {
   using TransformOpType = typename llvm::function_traits<
       decltype(&OpTy::applyToOne)>::template arg_t<0>;
   ArrayRef<Operation *> targets =
       state.getPayloadOps(this->getOperation()->getOperand(0));
-  SmallVector<Operation *> results;
+  SmallVector<SmallVector<Operation *>, 1> results;
+  // In the multi-result case, collect the number of results each transform
+  // produced.
   DiagnosedSilenceableFailure result = detail::applyTransformToEach(
       targets, results, [&](TransformOpType specificOp) {
         return static_cast<OpTy *>(this)->applyToOne(specificOp);
       });
   if (!result.succeeded())
     return result;
-
+  for (const SmallVector<Operation *> &oneTargetResults : results) {
+    if (OpTy::template hasTrait<OpTrait::ZeroResults>())
+      continue;
     if (OpTy::template hasTrait<OpTrait::OneResult>()) {
       transformResults.set(
-        this->getOperation()->getResult(0).template cast<OpResult>(), results);
+          this->getOperation()->getResult(0).template cast<OpResult>(),
+          oneTargetResults);
+      continue;
+    }
+    if (this->getOperation()->getNumResults() != oneTargetResults.size()) {
+      Diagnostic diag(this->getOperation()->getLoc(),
+                      DiagnosticSeverity::Error);
+      diag << "unexpected number of results (got " << oneTargetResults.size()
+           << " expected " << this->getOperation()->getNumResults() << ")";
+      return DiagnosedSilenceableFailure::silencableFailure(std::move(diag));
+    }
+    for (const auto &it :
+         llvm::zip(this->getOperation()->getResults(), oneTargetResults)) {
+      transformResults.set(std::get<0>(it).template cast<OpResult>(),
+                           std::get<1>(it));
+    }
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 template <typename OpTy>
 mlir::LogicalResult
 mlir::transform::TransformEachOpTrait<OpTy>::verifyTrait(Operation *op) {
   static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
                 "expected single-operand op");
-  static_assert(OpTy::template hasTrait<OpTrait::OneResult>() ||
-                    OpTy::template hasTrait<OpTrait::ZeroResults>(),
-                "expected zero- or single-result op");
   if (!op->getName().getInterface<TransformOpInterface>()) {
     return op->emitError() << "TransformEachOpTrait should only be attached to "
                               "ops that implement TransformOpInterface";
   }
 
   return success();
 }
 
 #endif // DIALECT_TRANSFORM_IR_TRANSFORMINTERFACES_H

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

FailureOr<TiledLinalgOp> result =
       pattern.returningMatchAndRewrite(target, rewriter);
   if (failed(result))
     return failure();
 
   return result->op;
 }
 
 //===----------------------------------------------------------------------===//
+// SplitReductionOp
+//===----------------------------------------------------------------------===//
+
+FailureOr<SmallVector<Operation *>>
+transform::SplitReductionOp::applyToOne(LinalgOp target) {
+  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};
+}
+
+//===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::TileOp::apply(TransformResults &transformResults,
                          TransformState &state) {
   LinalgTilingOptions tilingOptions;
   SmallVector<int64_t> tileSizes = extractI64Array(getSizes());
 
   if (!tileSizes.empty())
     tilingOptions.setTileSizes(tileSizes);
   tilingOptions.setInterchange(extractUIntArray(getInterchange()));
   LinalgTilingPattern pattern(getContext(), tilingOptions);

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

 FailureOr<LinalgOp> mlir::linalg::splitReduction(
     PatternRewriter &b, LinalgOp op,
     const ControlSplitReductionFn &controlSplitReductionFn,
     const LinalgTransformationFilter &filter) {
   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);
+  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) {
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPoint(op);
+
   std::pair<int64_t, unsigned> control = controlSplitReductionFn(op);
   int64_t ratio = control.first;
   unsigned insertDimIndex = control.second;
   if (ratio <= 1)
     return b.notifyMatchFailure(op, "split ratio needs to be greater than 1");
+
   SmallVector<unsigned> dims;
   op.getReductionDims(dims);
   assert(dims.size() == 1);
   unsigned reductionDim = dims[0];
   SmallVector<int64_t, 4> loopRanges = op.getStaticLoopRanges();
   int64_t reductionDimSize = loopRanges[reductionDim];
   if (reductionDimSize == ShapedType::kDynamicSize ||
       reductionDimSize % ratio != 0 || insertDimIndex >= loopRanges.size())
     return b.notifyMatchFailure(
         op, "Reduction dimension not divisible by split ratio");
+
   SmallVector<Operation *, 4> combinerOps;
   if (!matchReduction(op.getRegionOutputArgs(), 0, combinerOps) ||
       combinerOps.size() != 1)
     return b.notifyMatchFailure(op, "Cannot match the reduction pattern");
+
   Operation *reductionOp = combinerOps[0];
   Optional<Attribute> identity = getIdentity(reductionOp);
   if (!identity)
-    return b.notifyMatchFailure(op, "Unknown identity value for the redution");
+    return b.notifyMatchFailure(op, "Unknown identity value for the reduction");
 
   Location loc = op->getLoc();
   SmallVector<Value> newInputs;
   SmallVector<AffineMap> newMaps;
   // Calculate the new shapes and indexing maps of the input operands.
   for (OpOperand *operand : op.getInputOperands()) {
     AffineMap map = op.getTiedIndexingMap(operand);
     SmallVector<int64_t> newShape;
for (OpOperand *operand : op.getInputOperands()) {
     }
     Type newType = RankedTensorType::get(
         newShape,
         operand->get().getType().cast<RankedTensorType>().getElementType());
     Value newInput = b.create<tensor::ExpandShapeOp>(
         loc, newType, operand->get(), reassociation);
     newInputs.push_back(newInput);
   }
+
   // Calculate the new output map and shape, we insert the new dimension based
   // on the index returned by `controlSplitReductionFn`.
   SmallVector<int64_t> newOutputShape;
   AffineMap oldOutputMap = op.getTiedIndexingMap(op.getOutputOperand(0));
   ArrayRef<int64_t> oldShape = op.getShape(op.getOutputOperand(0));
   SmallVector<AffineExpr> outputExpr;
   for (unsigned idx :
        llvm::seq<unsigned>(0, oldOutputMap.getNumResults() + 1)) {
FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
   // Create the new op matching the original op with an extra parallel
   // dimension.
   GenericOp genericOp = b.create<GenericOp>(
       loc, TypeRange({initTensor.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.
+  // 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);
   SmallVector<Value> outputOperands = op.getOutputOperands();
   SmallVector<StringRef> reductionIteratorTypes;
   SmallVector<AffineExpr> exprs;
   for (unsigned i : llvm::seq<unsigned>(0, intermRank)) {
     if (insertDimIndex == i) {
       reductionIteratorTypes.push_back(getReductionIteratorTypeName());
auto reduction = b.create<GenericOp>(
       outputOperands, reductionMaps, reductionIteratorTypes,
       [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());
-  filter.replaceLinalgTransformationFilter(b, genericOp);
-  filter.replaceLinalgTransformationFilter(b, reduction);
-  return cast<LinalgOp>(genericOp.getOperation());
+
+  return SplitReductionResult{identityTensor.getDefiningOp<FillOp>(),
+                              cast<LinalgOp>(genericOp.getOperation()),
+                              reduction};
 }
 
 namespace {
 
 struct LinalgSplitReduction : public OpInterfaceRewritePattern<LinalgOp> {
   /// Construct a generic pattern applied to all LinalgOp that verify `filter`.
   LinalgSplitReduction(MLIRContext *context,
                        ControlSplitReductionFn controlSplitReductionFn,

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

+// RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
+
+// CHECK-LABEL: func.func @matmul_split
+func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+
+  //      CHECK: linalg.generic 
+  // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+  // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}} : tensor<16x4x64xf32>, tensor<4x64x32xf32>)
+  // CHECK-SAME: outs(%{{[a-zA-Z0-9]*}} : tensor<16x32x4xf32>) {
+
+  //      CHECK: linalg.generic 
+  // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+  // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}} : tensor<16x32x4xf32>)
+  // CHECK-SAME: outs(%{{[a-zA-Z0-9]*}} : tensor<16x32xf32>) {
+  %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
+                    outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+  return %0: tensor<16x32xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %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 %0 { split_factor = 4, insert_split_dimension = 2}
+  }
+}

mlir/test/Dialect/Transform/test-interpreter.mlir

^bb1(%arg1: !pdl.operation):
     %0 = transform.pdl_match @match_const in %arg1
     %1 = transform.loop.get_parent_for %0
     // expected-error @below {{only isolated-from-above ops can be alternative scopes}}
     alternatives %1 {
     ^bb2(%arg2: !pdl.operation):
     }
   }
 }
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{unexpected number of results (got 0 expected 3)}}
+  transform.test_wrong_number_of_results %arg0
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTTRANSFORMDIALECTEXTENSION_H
 #define MLIR_TESTTRANSFORMDIALECTEXTENSION_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
-#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 class DialectRegistry;
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "TestTransformDialectExtension.h.inc"
 
 namespace test {
 /// Registers the test extension to the Transform dialect.
 void registerTestTransformDialectExtension(::mlir::DialectRegistry &registry);
 } // namespace test
 
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_H

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

DiagnosedSilenceableFailure mlir::test::TestEmitRemarkAndEraseOperandOp::apply(
   for (Operation *op : state.getPayloadOps(getTarget()))
     op->erase();
 
   if (getFailAfterErase())
     return emitSilenceableError() << "silencable error";
   return DiagnosedSilenceableFailure::success();
 }
 
+FailureOr<SmallVector<Operation *>>
+mlir::test::TestWrongNumberOfResultsOp::applyToOne(Operation *) {
+  return SmallVector<Operation *>{};
+}
+
 namespace {
 /// Test extension of the Transform dialect. Registers additional ops and
 /// declares PDL as dependent dialect since the additional ops are using PDL
 /// types for operands and results.
 class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

: Op<Transform_Dialect, "test_emit_remark_and_erase_operand",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]> {
   let arguments = (ins PDL_Operation:$target, StrAttr:$remark,
                    UnitAttr:$fail_after_erase);
   let assemblyFormat = "$target `,` $remark attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
+def TestWrongNumberOfResultsOp
+  : Op<Transform_Dialect, "test_wrong_number_of_results",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface, 
+     TransformEachOpTrait, TransformOpInterface]> {
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$a,
+                      PDL_Operation:$b,
+                      PDL_Operation:$c);
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::test";
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<::llvm::SmallVector<::mlir::Operation *>> applyToOne(
+        ::mlir::Operation *target);
+  }];
+}
+
 #endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD