[mlir][Linalg] Implement the tiling interface for softmax

[qcolombet]

Summary

Original implementation from @harsh

[rengolin]

mlir/test/Dialect/Linalg/tile-softmax.mlir

25

The semantics of this is to take the max of each tile, which is not always the same as the softmax of the original dimension.

If the tile is a whole head, this may be what you want. If not, you'll get different results.

I'm not sure how to restrict this in a meaningful way, or perhaps this is up to the compiler to "do the right thing".

At the very least, this should be documented somewhere, perhaps in the op description?

40

I'm assuming this also works with the tile-and-fuse pass. Could there be a test for that, too?

[qcolombet]

mlir/test/Dialect/Linalg/tile-softmax.mlir

25

Good point.
In this example I went with the easiest transformation to apply, but you're right this may not be correct.
Documenting this somewhere is sensible, but it may make more sense to reject invalid tiling.
At the same time maybe the compiler/ir author knows something we don't and we should let this go through...

@ftynse, @nicolasvasilache do you have a recommendation on what we should do here?

Technically getTiledImplementation can return a failure, but I don't know if that's doable/desirable to do that check here.

40

Let me give a try!

[qcolombet]

mlir/test/Dialect/Linalg/tile-softmax.mlir

25

Spoke with @nicolasvasilache offline and the conclusion was the compiler has to do the right thing.
I'll add some wording around that.

[rengolin]

Should be fine with just transform tests for now. Once we incorporate on existing infra we'll see how it plays with passes.

[jpienaar]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

53

This is the deprecated form https://mlir.llvm.org/deprecation/

2260

s/mlir:://

[qcolombet]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

53

Thanks I didn't know about that.

[qcolombet]

• Add a tile and fuse test
• Remove deprecated use of cast
• Fix the way we generate the dimension (value vs. attribute)

[qcolombet]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

55

FYI, the "fix" I mentioned in my previous update is here. We use an attribute instead of a value here and that allows the fusing to work properly. I'll highlight what was broken without that.

65

@springerm Do you know if such utility function already exists somewhere?

2260

I had to keep that otherwise clang grab the clone from Operation:: which doesn't take any argument.

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

563

This change is not technically required for this PR, but when I "screwed up" my tiling by returning a value instead of an attribute, this led to slice ops with dynamic dimensions and these were breaking ExtractSliceOp::rankReduceIfNeeded.

By changing this assert into a note, this makes the compiler more robust in case someone does the same mistake as I did. I.e., it will not fuse instead of crashing the compiler.

mlir/test/Dialect/Linalg/tile-softmax.mlir

40

@rengolin I added a test for tile and fuse and it exposed a weakness in the fuse algorithm, which I fixed here as well (it was asserting.)

[akuegel]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

20

This introduces a cycle between LinalgUtils and Linalg. Is this intended?

[qcolombet]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

20

Ah good catch, I forgot to remove that line.
I saw this problem and in the end I moved the utility functions directly in this file, so the include is actually not needed.

Let me fix that.

[qcolombet]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

20

Fixed at cf244b55c3d2

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

 #define LINALG_OPS
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/TilingInterface.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
 class Linalg_Op<string mnemonic, list<Trait> traits = []> :
     Op<Linalg_Dialect, mnemonic, traits>;
 
 def Linalg_YieldOp : Linalg_Op<"yield", [Pure, ReturnLike, Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
def Linalg_IndexOp : Linalg_Op<"index", [Pure]>,
 
   let assemblyFormat = [{ $dim attr-dict `:` type($result) }];
   let hasVerifier = 1;
 }
 
 def Linalg_SoftmaxOp : Linalg_Op<"softmax",
     [DestinationStyleOpInterface,
      PredOpTrait<"input and output have same element type", TCopVTEtIsSameAs<0, 1>>,
-     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+     DeclareOpInterfaceMethods<TilingInterface,
+      ["getIterationDomain",
+       "getLoopIteratorTypes",
+       "getResultTilePosition",
+       "getTiledImplementation"]>]> {
   let summary = "Softmax operator";
   let description = [{
     linalg.softmax computes a numerically stable version of softmax.
 
     For a given input tensor and a specified dimension `d`, compute:
       1. the max `m` along that dimension `d`
       2. f(x) = exp(x - m)
       3. sum f(x) along dimension d to get l(x).
       4. compute the final result f(x) / l(x).
 
     This is an aggregate linalg operation that further reduces to a small DAG of
     structured operations.
+
+    Warning: Regarding the tiling capabilities, the implementation doesn't
+    check that the provided dimensions make sense. This is the responsability
+    of the transformation calling the tiling to ensure that the provided
+    sizes for each dimension make sense with respect to the semantic of
+    softmax.
   }];
 
   let arguments = (ins AnyShaped:$input,
                        AnyShaped:$output,
                        I64Attr:$dimension
   );
 
   let results = (outs Variadic<AnyRankedTensor>:$result);

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"

 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::linalg;
 
+/// Return a `memref.dim` or `tensor.dim` for the shape of `v` at `dim`.
+static OpFoldResult getDimValue(OpBuilder &builder, Location loc, Value v,
+                                int64_t dim) {
+  auto type = cast<ShapedType>(v.getType());
+  if (!type.isDynamicDim(dim))
+    return builder.getIndexAttr(type.getDimSize(dim));
+
+  return getAsOpFoldResult(
+      TypeSwitch<Type, Value>(v.getType())
+          .Case<RankedTensorType>([&](RankedTensorType t) -> Value {
+            return builder.create<tensor::DimOp>(loc, v, dim);
+          })
+          .Case<MemRefType>([&](MemRefType t) -> Value {
+            return builder.create<memref::DimOp>(loc, v, dim);
+          }));
+}
+
+/// Returns a memref.subview or a tensor.extract_slice based on the type of the
+/// `source`.
+static Value getSlice(OpBuilder &b, Location loc, Value source,
+                      ArrayRef<OpFoldResult> offsets,
+                      ArrayRef<OpFoldResult> sizes,
+                      ArrayRef<OpFoldResult> strides) {
+  return TypeSwitch<Type, Value>(source.getType())
+      .Case<RankedTensorType>([&](RankedTensorType t) -> Value {
+        return b.create<tensor::ExtractSliceOp>(loc, source, offsets, sizes,
+                                                strides);
+      })
+      .Case<MemRefType>([&](MemRefType type) -> Value {
+        return b.create<memref::SubViewOp>(loc, source, offsets, sizes,
+                                           strides);
+      })
+      .Default([&](Type t) { return nullptr; });
+}
+
 //===----------------------------------------------------------------------===//
 // Helper functions
 //===----------------------------------------------------------------------===//
 
 Value linalg::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                 int64_t dim) {
   if (llvm::isa<UnrankedMemRefType, MemRefType>(source.getType()))
     return b.createOrFold<memref::DimOp>(loc, source, dim);
LogicalResult SoftmaxOp::verify() {
   int64_t inputRank = getInputOperandRank();
   int64_t dimension = getDimension();
   if ((dimension < 0) || (dimension >= inputRank))
     return emitOpError("incorrect dimension specified");
 
   return success();
 }
 
+SmallVector<Range> SoftmaxOp::getIterationDomain(OpBuilder &builder) {
+  int64_t operandRank = getInputOperandRank();
+  SmallVector<Range> loopBounds(operandRank);
+  Location loc = getLoc();
+  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+  Value one = builder.create<arith::ConstantIndexOp>(loc, 1);
+  Value source = getInput();
+  for (auto dim : llvm::seq<int64_t>(0, operandRank)) {
+    loopBounds[dim].offset = zero;
+    loopBounds[dim].size = getDimValue(builder, loc, source, dim);
+    loopBounds[dim].stride = one;
+  }
+  return loopBounds;
+}
+
+SmallVector<utils::IteratorType> SoftmaxOp::getLoopIteratorTypes() {
+  SmallVector<utils::IteratorType> iteratorTypes(getInputOperandRank(),
+                                                 utils::IteratorType::parallel);
+  iteratorTypes[getDimension()] = utils::IteratorType::reduction;
+  return iteratorTypes;
+}
+
+FailureOr<TilingResult>
+SoftmaxOp::getTiledImplementation(OpBuilder &builder,
+                                  ArrayRef<OpFoldResult> offsets,
+                                  ArrayRef<OpFoldResult> sizes) {
+  int64_t rank = getInputOperandRank();
+  auto oneAttr = builder.getI64IntegerAttr(1);
+  SmallVector<OpFoldResult> strides(rank, oneAttr);
+  SmallVector<Value> tiledOperands;
+  tiledOperands.emplace_back(
+      getSlice(builder, getLoc(), getInput(), offsets, sizes, strides));
+  tiledOperands.emplace_back(
+      getSlice(builder, getLoc(), getOutput(), offsets, sizes, strides));
+
+  SmallVector<Type, 4> resultTypes;
+  if (hasTensorSemantics())
+    resultTypes.push_back(tiledOperands[1].getType());
+  Operation *tiledOp =
+      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
+
+  return TilingResult{{tiledOp}, SmallVector<Value>(tiledOp->getResults())};
+}
+
+LogicalResult SoftmaxOp::getResultTilePosition(
+    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
+    SmallVector<OpFoldResult> &resultSizes) {
+  if (resultNumber == 0) {
+    resultOffsets.assign(offsets.begin(), offsets.end());
+    resultSizes.assign(sizes.begin(), sizes.end());
+    return success();
+  }
+  return failure();
+}
+
 // cast(dynamic) -> static.
 LogicalResult SoftmaxOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
   return memref::foldMemRefCast(*this);
 }
 
 LogicalResult
 SoftmaxOp::reifyResultShapes(OpBuilder &b,
                              ReifiedRankedShapedTypeDims &reifiedReturnShapes) {

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

for (auto tiledOp : tileAndFuseResult->tiledOps) {
     LLVM_DEBUG(DBGS() << "tiledProducer: " << *tiledOp << "\n");
   }
 #endif
 
   // Replace the extract op.
   auto maybeRankReduced = tensor::ExtractSliceOp::rankReduceIfNeeded(
       rewriter, sliceOpToTile->getLoc(), tileAndFuseResult->tiledValues[0],
       cast<RankedTensorType>(sliceOpToTile->getResult(0).getType()).getShape());
-  assert(succeeded(maybeRankReduced) && "unexpected shape");
+  if (failed(maybeRankReduced)) {
+    diag.attachNote(producerOp->getLoc())
+        << "shape types don't match (missing canonicalization?):\nTiledOp: "
+        << tileAndFuseResult->tiledValues[0]
+        << "\nSliceOp: " << sliceOpToTile.getOperation() << '\n';
+    return {};
+  }
   rewriter.replaceOp(sliceOpToTile, *maybeRankReduced);
 
   // Add new outputs to containing op, if required
   Operation *newContainingOp = replaceForAllWithNewSignature(
       rewriter, diag, producerOp, containingOp, *tileAndFuseResult,
       resultNumber, offsets, sizes);
 
   return std::make_tuple(tileAndFuseResult->tiledOps, newContainingOp);

mlir/test/Dialect/Linalg/tile-softmax.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize --split-input-file | FileCheck %s
+
+// Check that we can tile softmax on tensors.
+// The tiling here is 2x3.
+// So the shape used in the inner loop should be 2x3x256, however since 3
+// doesn't divide the second dimension (64), we should see a '?' in the shape.
+// The actual size, used through extract_slice/insert_slice, should come from a
+// `min(64 - current iteration index, 3)`
+
+// CHECK: #[[$MIN_MAP:.*]] = affine_map<(d0) -> (-d0 + 64, 3)>
+// CHECK-LABEL:   func.func @softmax(
+// CHECK-SAME:                       %[[VAL_0:.*]]: tensor<16x64x256xf32>) -> tensor<16x64x256xf32> {
+// CHECK-DAG:       %[[C3:.*]] = arith.constant 3 : index
+// CHECK-DAG:       %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG:       %[[C64:.*]] = arith.constant 64 : index
+// CHECK-DAG:       %[[C16:.*]] = arith.constant 16 : index
+// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           %[[TENSOR_EMPTY:.*]] = tensor.empty() : tensor<16x64x256xf32>
+// CHECK:           %[[VAL_7:.*]] = scf.for %[[VAL_8:.*]] = %[[C0]] to %[[C16]] step %[[C2]] iter_args(%[[VAL_9:.*]] = %[[TENSOR_EMPTY]]) -> (tensor<16x64x256xf32>) {
+// CHECK:             %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[C0]] to %[[C64]] step %[[C3]] iter_args(%[[VAL_12:.*]] = %[[VAL_9]]) -> (tensor<16x64x256xf32>) {
+// CHECK:               %[[VAL_13:.*]] = affine.min #[[$MIN_MAP]](%[[VAL_11]])
+// CHECK:               %[[VAL_14:.*]] = tensor.extract_slice %[[VAL_0]]{{\[}}%[[VAL_8]], %[[VAL_11]], 0] [2, %[[VAL_13]], 256] [1, 1, 1] : tensor<16x64x256xf32> to tensor<2x?x256xf32>
+// CHECK:               %[[VAL_15:.*]] = tensor.extract_slice %[[VAL_12]]{{\[}}%[[VAL_8]], %[[VAL_11]], 0] [2, %[[VAL_13]], 256] [1, 1, 1] : tensor<16x64x256xf32> to tensor<2x?x256xf32>
+// CHECK:               %[[VAL_16:.*]] = linalg.softmax dimension(1) ins(%[[VAL_14]] : tensor<2x?x256xf32>) outs(%[[VAL_15]] : tensor<2x?x256xf32>) -> tensor<2x?x256xf32>
+// CHECK:               %[[VAL_17:.*]] = tensor.insert_slice %[[VAL_16]] into %[[VAL_12]]{{\[}}%[[VAL_8]], %[[VAL_11]], 0] [2, %[[VAL_13]], 256] [1, 1, 1] : tensor<2x?x256xf32> into tensor<16x64x256xf32>
+// CHECK:               scf.yield %[[VAL_17]] : tensor<16x64x256xf32>
+// CHECK:             }
+// CHECK:             scf.yield %[[VAL_18:.*]] : tensor<16x64x256xf32>
+// CHECK:           }
+// CHECK:           return %[[VAL_19:.*]] : tensor<16x64x256xf32>
+// CHECK:         }
+func.func @softmax(%arg0: tensor<16x64x256xf32>) -> tensor<16x64x256xf32> {
+  %0 = tensor.empty() : tensor<16x64x256xf32>
+  %1 = linalg.softmax
+         dimension(1) ins(%arg0 : tensor<16x64x256xf32>) outs(%0 : tensor<16x64x256xf32>) -> tensor<16x64x256xf32>
+  return %1 : tensor<16x64x256xf32>
+}
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !transform.any_op):
+    %0 = transform.structured.match ops{["linalg.softmax"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loop:2 = transform.structured.tile %0 [2, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
+// Test the softmax tiling interface with the tile_to_forall_op transform and
+// check that it composes properly with the fuse transform.
+// This should sink the linalg.generic inside the scf.forall and run that
+// generic on 2x4x256 tensors (2==16/8, 4==64/16).
+
+// CHECK: #[[$TIMES2_MAP:.*]] = affine_map<(d0) -> (d0 * 2)>
+// CHECK: #[[$TIMES4_MAP:.*]] = affine_map<(d0) -> (d0 * 4)>
+// CHECK-LABEL:   func.func @softmax_tile_n_fuse(
+// CHECK-SAME:                       %[[VAL_0:.*]]: tensor<16x64x256xf32>) -> tensor<16x64x256xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_2:.*]] = tensor.empty() : tensor<16x64x256xf32>
+// CHECK:           %[[VAL_3:.*]] = tensor.empty() : tensor<16x64x256xf32>
+// CHECK:           %[[VAL_4:.*]] = scf.forall (%[[VAL_5:.*]], %[[VAL_6:.*]]) in (8, 16) shared_outs(%[[VAL_7:.*]] = %[[VAL_3]]) -> (tensor<16x64x256xf32>) {
+// CHECK:             %[[VAL_8:.*]] = affine.apply #[[$TIMES2_MAP]](%[[VAL_5]])
+// CHECK:             %[[VAL_9:.*]] = affine.apply #[[$TIMES4_MAP]](%[[VAL_6]])
+// CHECK:             %[[VAL_10:.*]] = affine.apply #[[$TIMES2_MAP]](%[[VAL_5]])
+// CHECK:             %[[VAL_11:.*]] = affine.apply #[[$TIMES4_MAP]](%[[VAL_6]])
+// CHECK:             %[[VAL_12:.*]] = affine.apply #[[$TIMES2_MAP]](%[[VAL_5]])
+// CHECK:             %[[VAL_13:.*]] = affine.apply #[[$TIMES4_MAP]](%[[VAL_6]])
+// CHECK:             %[[VAL_14:.*]] = tensor.extract_slice %[[VAL_0]]{{\[}}%[[VAL_10]], %[[VAL_11]], 0] [2, 4, 256] [1, 1, 1] : tensor<16x64x256xf32> to tensor<2x4x256xf32>
+// CHECK:             %[[VAL_15:.*]] = tensor.extract_slice %[[VAL_2]]{{\[}}%[[VAL_12]], %[[VAL_13]], 0] [2, 4, 256] [1, 1, 1] : tensor<16x64x256xf32> to tensor<2x4x256xf32>
+// CHECK:             %[[VAL_16:.*]] = linalg.generic {indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[VAL_14]] : tensor<2x4x256xf32>) outs(%[[VAL_15]] : tensor<2x4x256xf32>) {
+// CHECK:             ^bb0(%[[VAL_17:.*]]: f32, %[[VAL_18:.*]]: f32):
+// CHECK:               %[[VAL_19:.*]] = arith.addf %[[VAL_18]], %[[VAL_1]] : f32
+// CHECK:               linalg.yield %[[VAL_19]] : f32
+// CHECK:             } -> tensor<2x4x256xf32>
+// CHECK:             %[[VAL_20:.*]] = tensor.extract_slice %[[VAL_7]]{{\[}}%[[VAL_8]], %[[VAL_9]], 0] [2, 4, 256] [1, 1, 1] : tensor<16x64x256xf32> to tensor<2x4x256xf32>
+// CHECK:             %[[VAL_21:.*]] = linalg.softmax dimension(1) ins(%[[VAL_22:.*]] : tensor<2x4x256xf32>) outs(%[[VAL_20]] : tensor<2x4x256xf32>) -> tensor<2x4x256xf32>
+// CHECK:             scf.forall.in_parallel {
+// CHECK:               tensor.parallel_insert_slice %[[VAL_21]] into %[[VAL_7]]{{\[}}%[[VAL_8]], %[[VAL_9]], 0] [2, 4, 256] [1, 1, 1] : tensor<2x4x256xf32> into tensor<16x64x256xf32>
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return %[[VAL_23:.*]] : tensor<16x64x256xf32>
+// CHECK:         }
+
+func.func @softmax_tile_n_fuse(%arg0: tensor<16x64x256xf32>) -> tensor<16x64x256xf32> {
+  %empty = tensor.empty() : tensor<16x64x256xf32>
+  %cst = arith.constant 1.000000e+00 : f32
+  %eltwise = linalg.generic
+      {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                        affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+       iterator_types = ["parallel", "parallel", "parallel"]
+      }
+      ins(%arg0 : tensor<16x64x256xf32>)
+      outs(%empty : tensor<16x64x256xf32>) {
+    ^bb0(%arg2: f32, %arg3: f32):
+      %arg3Plus1 = arith.addf %arg3, %cst : f32
+      linalg.yield %arg3Plus1 : f32
+    } -> tensor<16x64x256xf32>
+
+  %0 = tensor.empty() : tensor<16x64x256xf32>
+  %1 = linalg.softmax
+         dimension(1) ins(%eltwise : tensor<16x64x256xf32>) outs(%0 : tensor<16x64x256xf32>) -> tensor<16x64x256xf32>
+  return %1 : tensor<16x64x256xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.softmax"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+
+  // Tile the root.
+  %forall_op, %tiled_op = transform.structured.tile_to_forall_op %0 num_threads [8, 16]
+       : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+  // Fuse all producers.
+  %1 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.structured.fuse_into_containing_op %1 into %forall_op
+    : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+// -----
+
+// Same as the previous test but on memrefs.
+
+// CHECK: #[[$MIN_MAP:.*]] = affine_map<(d0) -> (-d0 + 64, 3)>
+// CHECK-LABEL:   func.func @softmax_memref(
+// CHECK-SAME:                              %[[VAL_0:.*]]: memref<16x64x256xf32>,
+// CHECK-SAME:                              %[[VAL_1:.*]]: memref<16x64x256xf32>) {
+// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[C16:.*]] = arith.constant 16 : index
+// CHECK-DAG:       %[[C64:.*]] = arith.constant 64 : index
+// CHECK-DAG:       %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG:       %[[C3:.*]] = arith.constant 3 : index
+// CHECK:           scf.for %[[VAL_7:.*]] = %[[C0]] to %[[C16]] step %[[C2]] {
+// CHECK:             scf.for %[[VAL_8:.*]] = %[[C0]] to %[[C64]] step %[[C3]] {
+// CHECK:               %[[VAL_9:.*]] = affine.min #[[$MIN_MAP]](%[[VAL_8]])
+// CHECK:               %[[VAL_10:.*]] = memref.subview %[[VAL_0]]{{\[}}%[[VAL_7]], %[[VAL_8]], 0] [2, %[[VAL_9]], 256] [1, 1, 1] : memref<16x64x256xf32> to memref<2x?x256xf32, strided<[16384, 256, 1], offset: ?>>
+// CHECK:               %[[VAL_11:.*]] = memref.subview %[[VAL_1]]{{\[}}%[[VAL_7]], %[[VAL_8]], 0] [2, %[[VAL_9]], 256] [1, 1, 1] : memref<16x64x256xf32> to memref<2x?x256xf32, strided<[16384, 256, 1], offset: ?>>
+// CHECK:               linalg.softmax dimension(1) ins(%[[VAL_10]] : memref<2x?x256xf32, strided<[16384, 256, 1], offset: ?>>) outs(%[[VAL_11]] : memref<2x?x256xf32, strided<[16384, 256, 1], offset: ?>>)
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+func.func @softmax_memref(%arg0: memref<16x64x256xf32>, %arg1: memref<16x64x256xf32>) {
+  linalg.softmax
+    dimension(1) ins(%arg0 : memref<16x64x256xf32>) outs(%arg1 : memref<16x64x256xf32>)
+  return
+}
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !transform.any_op):
+    %0 = transform.structured.match ops{["linalg.softmax"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loop:2 = transform.structured.tile %0 [2, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}