Add support for having parallel dimension *after* reduction dimension in split reduction

[vmurali]

Summary

Previously, splitReduction transformation added the split parallel dimension *before* the reduction dimension, leading to tiling for reduction. This commit creates an option to create the parallel dimension *after* the reduction dimension, allowing us to transform the op into vertical reduction with SIMD parallelism.

[vmurali]

Fixed a bug - was earlier conflating control.index with reductionDimension.

[ThomasRaoux]

Looks good to me. Please address the couple style nits.

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

1366–1378

nit: comments should end with a period here and below.

1368–1373

nit: could you move the comments above the associated members to match the style in the rest of the code.

[dcaballe]

LGTM, thanks!

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

468

nit: formatting.

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

1005

nit: Add comment with /*innerParallel=*/false

[vmurali]

Addressed the nit issues.

[vmurali]

Added a period

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

let description = [{
     reduction into a parallel and reduction dimension. 
     A new `linalg.generic` op is created to perform the rest of the reduction. 
     
     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.
+      - inner_parallel: specifies whether the parallel dimension is before or
+        after the reduction dimension in the splitting op.
       - 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)
 
     #### Return modes
     
     This operation ignores non-Linalg ops and drops them in the return.
let description = [{
 
      return %4 : tensor<16x32xf32>
     ```
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64Attr, "{}">:$split_factor,
                    DefaultValuedAttr<I64Attr, "{}">:$insert_split_dimension,
+                   UnitAttr:$inner_parallel,
                    UnitAttr:$use_scaling_algorithm,
                    UnitAttr:$use_alloc);
   let results = (outs PDL_Operation:$init_or_alloc_op,
                       PDL_Operation:$fill_op,
                       PDL_Operation:$split_linalg_op,
                       PDL_Operation:$combining_linalg_op);
 
   let assemblyFormat = "$target attr-dict";

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

static void insert(RewritePatternSet &patterns,
                      const LinalgTilingOptions &options,
                      const LinalgTransformationFilter &f) {
     patterns.add<LinalgTilingPattern>(OpTy::getOperationName(),
                                       patterns.getContext(), options, f);
     TilingPatterns<OpTypes...>::insert(patterns, options, f);
   }
 };
 
-/// 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.
+/// Split Reduction options.
+struct SplitReductionOptions {
+  // Ratio used to split the reduction dimension.  If the ratio is <= 1, nothing
+  // will be done.
+  int64_t ratio = 0;
+  // Index where the extra dimension is added to the intermediate tensor shape.
+  unsigned index = 0;
+  // If the inner dimension after splitting is parallel or reduction.
+  bool innerParallel = false;
+};
+
+/// Function signature to control reduction splitting. This returns
+/// `SplitReductionOptions`.
 // TODO: don't use unsigned unless doing bit manipulation.
 using ControlSplitReductionFn =
-    std::function<std::pair<int64_t, unsigned>(LinalgOp op)>;
+    std::function<SplitReductionOptions(LinalgOp op)>;
 
 /// Patterns to apply `splitReduction` below.
 void populateSplitReductionPattern(
     RewritePatternSet &patterns,
     const ControlSplitReductionFn &controlSplitReductionFn,
     const LinalgTransformationFilter &f = LinalgTransformationFilter(),
     bool useAlloc = false);
 

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

 // SplitReductionOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::SplitReductionOp::applyToOne(linalg::LinalgOp target,
                                         SmallVectorImpl<Operation *> &results,
                                         transform::TransformState &state) {
   ControlSplitReductionFn splitFn = [&](LinalgOp) {
-    return std::pair<int64_t, unsigned>(getSplitFactor(),
-                                        getInsertSplitDimension());
+    return linalg::SplitReductionOptions{int64_t(getSplitFactor()),
+                                         unsigned(getInsertSplitDimension()),
+                                         /*innerParallel=*/false};
   };
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<SplitReductionResult> splitResult =
       (getUseScalingAlgorithm())
           ? splitReductionByScaling(rewriter, target, splitFn, getUseAlloc())
           : splitReduction(rewriter, target, splitFn, getUseAlloc());
   if (failed(splitResult))

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

 }
 
 FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
     PatternRewriter &b, LinalgOp op,
     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;
+  SplitReductionOptions control = controlSplitReductionFn(op);
+  int64_t ratio = control.ratio;
+  unsigned insertSplitDimension = control.index;
   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();
for (OpOperand *operand : op.getInputOperands()) {
     AffineMap map = op.getTiedIndexingMap(operand);
     SmallVector<int64_t> newShape;
     SmallVector<AffineExpr> exprs;
     SmallVector<ReassociationIndices> reassociation;
     unsigned index = 0;
     for (unsigned idx : llvm::seq<unsigned>(0, map.getNumResults())) {
       unsigned dim = map.getDimPosition(idx);
       if (reductionDim == dim) {
+        if (control.innerParallel) {
+          newShape.push_back(op.getShape(operand)[idx] / ratio);
+          newShape.push_back(ratio);
+        } else {
           newShape.push_back(ratio);
           newShape.push_back(op.getShape(operand)[idx] / ratio);
+        }
         reassociation.push_back({index++, index++});
+        if (control.innerParallel) {
+          exprs.push_back(b.getAffineDimExpr(reductionDim));
+          exprs.push_back(b.getAffineDimExpr(reductionDim + 1));
+        } else {
           exprs.push_back(b.getAffineDimExpr(insertSplitDimension));
           exprs.push_back(
               b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1));
+        }
         continue;
       }
       newShape.push_back(op.getShape(operand)[idx]);
+      if (control.innerParallel) {
+        exprs.push_back(
+            b.getAffineDimExpr(dim <= reductionDim ? dim : dim + 1));
+      } else {
         exprs.push_back(
             b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1));
+      }
       reassociation.push_back({index++});
     }
     newMaps.push_back(
         AffineMap::get(map.getNumDims() + 1, 0, exprs, op.getContext()));
     // If the shape is unchanged the input doesn't change.
     if (newShape == op.getShape(operand)) {
       newInputs.push_back(operand->get());
       continue;
FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
   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)) {
     if (idx == insertSplitDimension) {
       newOutputShape.push_back(ratio);
+      if (control.innerParallel) {
+        outputExpr.push_back(b.getAffineDimExpr(reductionDim + 1));
+      } else {
         outputExpr.push_back(b.getAffineDimExpr(insertSplitDimension));
+      }
       continue;
     }
-    unsigned oldDim = idx < insertSplitDimension ? idx : idx - 1;
-    newOutputShape.push_back(oldShape[oldDim]);
-    unsigned dim = oldOutputMap.getDimPosition(oldDim);
+    unsigned oldIdx = idx < insertSplitDimension ? idx : idx - 1;
+    newOutputShape.push_back(oldShape[oldIdx]);
+    unsigned dim = oldOutputMap.getDimPosition(oldIdx);
+    if (control.innerParallel) {
+      outputExpr.push_back(
+          b.getAffineDimExpr(dim <= reductionDim ? dim : dim + 1));
+    } else {
       outputExpr.push_back(
           b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1));
     }
+  }
   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, 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())
+    if (insertSplitDimension == it.index() && !control.innerParallel)
       newIteratorTypes.push_back(getParallelIteratorTypeName());
     newIteratorTypes.push_back(it.value().cast<StringAttr>().getValue());
+    if (insertSplitDimension == it.index() && control.innerParallel)
+      newIteratorTypes.push_back(getParallelIteratorTypeName());
   }
   // Create the new op matching the original op with an extra parallel
   // dimension.
   GenericOp genericOp = b.create<GenericOp>(
       loc, TypeRange({initOrAllocTensor.getType()}), newInputs,
       ValueRange({identityTensor}), newMaps, newIteratorTypes);
   b.inlineRegionBefore(op->getRegion(0), genericOp.getRegion(),
                        genericOp.getRegion().begin());

 /// Core rewrite implementation.
 FailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
     PatternRewriter &b, LinalgOp op,
     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;
+  SplitReductionOptions control = controlSplitReductionFn(op);
+  if (control.innerParallel)
+    return b.notifyMatchFailure(op, "innerParallel not supported");
+
+  int64_t splitFactor = control.ratio;
+  unsigned insertSplitDimension = control.index;
   if (splitFactor <= 1)
     return b.notifyMatchFailure(op, "split factor needs to be greater than 1");
 
   SmallVector<unsigned> dims;
   op.getReductionDims(dims);
   if (dims.empty())
     return b.notifyMatchFailure(op, "needs at least 1 reduction dimension");
 

mlir/test/Dialect/Linalg/split_reduction.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-split-reduction  -split-input-file  | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-split-reduction-inner-parallel  -split-input-file  | FileCheck %s --check-prefix=INNERPARALLELCHECK
 
 func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
   %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
                     outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
   return %0: tensor<16x32xf32>
 }
 
 //  CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>

 //      CHECK: } -> tensor<16x32x4xf32>
 //      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]],
 // CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) {
 //      CHECK:   arith.addf
 //      CHECK:   linalg.yield %{{.*}} : f32
 //      CHECK: } -> tensor<16x32xf32>
 //      CHECK: return %[[R]] : tensor<16x32xf32>
 
+//  INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+//  INNERPARALLELCHECK-LABEL: @matmul_split
+//  INNERPARALLELCHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32
+//  INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x64x4xf32>
+//  INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<64x4x32xf32>
+//  INNERPARALLELCHECK-DAG: %[[INI:.*]] = linalg.init_tensor [16, 32, 4] : tensor<16x32x4xf32>
+//      INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32>
+//      INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
+// INNERPARALLELCHECK-SAME:   , iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
+// INNERPARALLELCHECK-SAME:   ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) {
+//      INNERPARALLELCHECK:   arith.mulf
+//      INNERPARALLELCHECK:   arith.addf
+//      INNERPARALLELCHECK:   linalg.yield
+//      INNERPARALLELCHECK: } -> tensor<16x32x4xf32>
+//      INNERPARALLELCHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]],
+// INNERPARALLELCHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) {
+//      INNERPARALLELCHECK:   arith.addf
+//      INNERPARALLELCHECK:   linalg.yield %{{.*}} : f32
+//      INNERPARALLELCHECK: } -> tensor<16x32xf32>
+//      INNERPARALLELCHECK: return %[[R]] : tensor<16x32xf32>
+
 // -----
 
 func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> {
   %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
                                           affine_map<(d0) -> ()>,
                                           affine_map<(d0) -> ()>],
    iterator_types = ["reduction"]}
    ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>)

 //      CHECK:   linalg.yield
 //      CHECK: } -> tensor<4xf32>
 //      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {
 //      CHECK:   arith.mulf
 //      CHECK:   linalg.yield
 //      CHECK: } -> tensor<f32>
 //      CHECK: return %[[R]] : tensor<f32>
 
+//  INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()>
+//  INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d1)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
+//INNERPARALLELCHECK-LABEL: @generic_split_1d
+//      INNERPARALLELCHECK: %[[ID:.*]] = arith.constant 1.000000e+00 : f32
+//      INNERPARALLELCHECK: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<8x4xf32>
+//      INNERPARALLELCHECK: %[[INI:.*]] = linalg.init_tensor [4] : tensor<4xf32>
+//      INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>
+//      INNERPARALLELCHECK: %[[G:.*]] = linalg.generic
+//      INNERPARALLELCHECK:   {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],
+//      INNERPARALLELCHECK:   iterator_types = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) {
+//      INNERPARALLELCHECK:   arith.subf
+//      INNERPARALLELCHECK:   math.exp
+//      INNERPARALLELCHECK:   arith.mulf
+//      INNERPARALLELCHECK:   linalg.yield
+//      INNERPARALLELCHECK: } -> tensor<4xf32>
+//      INNERPARALLELCHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {
+//      INNERPARALLELCHECK:   arith.mulf
+//      INNERPARALLELCHECK:   linalg.yield
+//      INNERPARALLELCHECK: } -> tensor<f32>
+//      INNERPARALLELCHECK: return %[[R]] : tensor<f32>
+
 // -----
 
 func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)
   -> tensor<5x2xf32>
 {
   %0 = linalg.generic {
       indexing_maps = [
         affine_map<(d0, d1, d2) -> (d1, d0)>,

 //      CHECK:   linalg.yield
 //      CHECK: } -> tensor<5x2x4xf32>
 //      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
 // CHECK-SAME:   ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
 //      CHECK:   arith.maxf
 //      CHECK:   linalg.yield
 //      CHECK:  } -> tensor<5x2xf32>
 //      CHECK: return %[[R]] : tensor<5x2xf32>
+
+//  INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//  INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// INNERPARALLELCHECK-LABEL:  func @generic_split_3d
+//      INNERPARALLELCHECK: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32
+//  INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32>
+//  INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32>
+//      INNERPARALLELCHECK: %[[INI:.*]] = linalg.init_tensor [5, 2, 4] : tensor<5x2x4xf32>
+//      INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>
+//      INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
+// INNERPARALLELCHECK-SAME:   ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {
+//      INNERPARALLELCHECK:   arith.addf
+//      INNERPARALLELCHECK:   arith.maxf
+//      INNERPARALLELCHECK:   linalg.yield
+//      INNERPARALLELCHECK: } -> tensor<5x2x4xf32>
+//      INNERPARALLELCHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
+// INNERPARALLELCHECK-SAME:   ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
+//      INNERPARALLELCHECK:   arith.maxf
+//      INNERPARALLELCHECK:   linalg.yield
+//      INNERPARALLELCHECK:  } -> tensor<5x2xf32>
+//      INNERPARALLELCHECK: return %[[R]] : tensor<5x2xf32>

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

Option<bool> testSwapSubTensorPadTensor{
       *this, "test-swap-subtensor-padtensor",
       llvm::cl::desc("Test rewrite of subtensor(tensor.pad) into "
                      "tensor.pad(subtensor)"),
       llvm::cl::init(false)};
   Option<bool> testSplitReduction{
       *this, "test-split-reduction",
       llvm::cl::desc("Test split reduction transformation"),
       llvm::cl::init(false)};
+  Option<bool> testSplitReductionInnerParallel{
+      *this, "test-split-reduction-inner-parallel",
+      llvm::cl::desc("Test split reduction with inner parallel transformation"),
+      llvm::cl::init(false)};
   ListOption<int64_t> peeledLoops{
       *this, "peeled-loops",
       llvm::cl::desc("Loops to be peeled when test-tile-pattern")};
   ListOption<int64_t> tileSizes{
       *this, "tile-sizes",
       llvm::cl::desc("Linalg tile sizes for test-tile-pattern")};
   Option<bool> skipPartial{
       *this, "skip-partial",

 }
 
 static void applySplitReduction(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   linalg::populateSplitReductionPattern(
       patterns,
       [](LinalgOp op) {
         unsigned insertDimIndex = op.getNumLoops() - 1;
-        return std::make_pair(4, insertDimIndex);
+        return SplitReductionOptions{4, insertDimIndex, false};
+      },
+      LinalgTransformationFilter(
+          ArrayRef<StringAttr>{},
+          StringAttr::get(funcOp.getContext(), "SPLIT")));
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
+static void applySplitReductionInnerParallel(func::FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  linalg::populateSplitReductionPattern(
+      patterns,
+      [](LinalgOp op) {
+        unsigned insertDimIndex = op.getNumLoops() - 1;
+        return SplitReductionOptions{4, insertDimIndex, true};
       },
       LinalgTransformationFilter(
           ArrayRef<StringAttr>{},
           StringAttr::get(funcOp.getContext(), "SPLIT")));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyBubbleUpExtractSliceOpPattern(func::FuncOp funcOp) {
void TestLinalgTransforms::runOnOperation() {
   if (testTilePattern)
     return applyTilePattern(getOperation(), loopType, tileSizes, peeledLoops,
                             /*scalarizeDynamicDims=*/false);
   if (testTileScalarizeDynamicDims)
     return applyTilePattern(getOperation(), loopType, tileSizes,
                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
   if (testSplitReduction)
     return applySplitReduction(getOperation());
+  if (testSplitReductionInnerParallel)
+    return applySplitReductionInnerParallel(getOperation());
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
   if (testSwapExtractSliceWithFill)
     return applySwapExtractSliceWithFillPattern(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir