[mlir][transform] Add support for expressing scalable tile sizes

[awarzynski]

Summary

This patch enables specifying scalable tile sizes when using the
Transform dialect to drive tiling, e.g.:

%1, %loop = transform.structured.tile %0 [[4]]

This is implemented by extending the TileOp with a dedicated attribute
for "scalability" and by updating various parsing hooks.

This change is a part of larger effort to enable scalable vectorisation
in Linalg. See this RFC for more context:

• https://discourse.llvm.org/t/rfc-scalable-vectorisation-in-linalg/

[awarzynski]

I'm new to most of this code, so wanted to send this for early feedback to check whether I'm heading in the right direction.

I do realise that I am yet to test/cater for the following 2 cases (as discussed in the RFC):

• using SSA values for tile sizes,
• multi-dim tiling (e.g. [4, [8]]

The former should just work, the latter will require tweaking the parser.

Thank you for taking a look!

[ftynse]

The extension itself is fine, but I have a suggestion on how to rework the op parser to avoid this change being intrusive into core IR parser and other dialects.

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

1532

Do we assume that the scalable dimension is always the last? If not, I'd rather model this as an optional integer attribute that indicates the position of the scalable dimension (and is absent when no dimensions are scalable).

mlir/include/mlir/IR/OpImplementation.h

701 ↗

(On Diff #523693)

I don't think we want to change the core IR parser to support what feels like a special case. It looks possible to craft a parseElementFn that will do what is needed here along the following lines:

bool isScalable= false;
if (parseOptionalLBracket().succeeded())
  isScalable = true;
// do the usual parsing with optional SSA value and integer
if (isScalable && parseRBracket().failed())
  return failure();

mlir/lib/AsmParser/Parser.cpp

122 ↗

(On Diff #523693)

This function is not limited to sizes, it's an arbitrary comma-separated list that may be applied to SSA values, value/type pairs, etc.

[awarzynski]

Thank you for taking a look, @ftynse !

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

1532

Do we assume that the scalable dimension is always the last?

Good question. Looking at scalable vectors in MLIR, it would always be the trailing dim:

he scalable dimensions in a vector are indicated between square brackets ([ ]), and all fixed-length dimensions, if present, must precede the set of scalable dimensions. That is, a vector<2x[4]xf32> is valid, but vector<[4]x2xf32> is not.

However, it doesn't have to be only 1 dim. So, I'd limit this to the trailing dimensions. Personally I'd focus on 1 dim to begin with, but don't mind much.

mlir/include/mlir/IR/OpImplementation.h

701 ↗

(On Diff #523693)

I don't think we want to change the core IR parser to support what feels like a special case.

Good point, but we still need a way to communicate "scalability" all the way to TileOp::parse. Updating parseElementFnalone won't be enough. We'll still need to tweak parseDynamicIndexList, but I guess that would already be less intrusive.

Let me demonstrate first and then we can discuss whether that's OK or whether to refine more.

[awarzynski]

Revert changes in parseCommaSeparatedList

[ftynse]

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

1532

Ok. I would suggest renaming the attribute to something like last_dim_scalable for clarity.

mlir/include/mlir/Interfaces/ViewLikeInterface.h

71–82

Please update the documentation to reflect the new argument.

mlir/test/Dialect/Linalg/transform-op-tile.mlir

170–182

How much of this do we actually need to test the new functionality? The test looks like the entire IR with SSA values changed to captures. We prefer more minimalist tests to avoid them being spuriously broken by, e.g., changes in the tensor.dim pretty-printer syntax.

184

Don't hardcode generated alias names such as #map1, capture them instead of necessary.

191

Ditto.

205–207

Let's not use !pdl.operation anymore, prefer !transform.any_op or a more specific type.

208

It would be nice to also have a test with at least one fixed-length dimension and a scalable dimension.

[awarzynski]

Addressing comments from Alex (thank you!)

• Renamed the attribute (scalable --> last_tile_size_scalable)
• Trimmed the test and added another case with a fixed-length tile size
• Added documentation
• Refined transform::TileOp::apply to make sure that only the trailing dim is made scalable

[awarzynski]

mlir/test/Dialect/Linalg/transform-op-tile.mlir

170–182

Good point, this has now been trimmed.

[ftynse]

One small change in tests is needed, otherwise good to go.

mlir/test/Dialect/Linalg/transform-op-tile.mlir

178

This (and below) still hardcodes #map, #map2`, etc. These alias names are similar to SSA value names and have no guarantee of being preserved. You can use -mlir-print-local-scope to print maps inline or capture them with more CHECKs.

[c-rhodes]

mlir/include/mlir/Interfaces/ViewLikeInterface.h

78–79

nit: scalable indices are not allowed.

mlir/lib/Interfaces/ViewLikeInterface.cpp

128–174

should this check it's the last index that's scalable? It will currently parse:

transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
  %1, %loops:3 = transform.structured.tile %0 [4, [4], [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
}

or

transform.sequence failures(propagate) {
^bb0(%arg1: !transform.any_op):
  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
  %1, %loops:3 = transform.structured.tile %0 [[4], [4], [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
}

are there any negative tests for the parser for things like this?

[awarzynski]

mlir/include/mlir/Interfaces/ViewLikeInterface.h

78–79

Thank you :)

mlir/lib/Interfaces/ViewLikeInterface.cpp

128–174

Thanks, good catch! I will be sending an update shortly that addresses both your comments.

mlir/test/Dialect/Linalg/transform-op-tile.mlir

178

Argh, sorry about that! I wasn't aware of -mlir-print-local-scope and it looks like a perfect solution, thank you!

I will ignore #map2 below as that's far less relevant to "scalability" than this one (i.e. I'll replace it with a generic regex).

[awarzynski]

Address the remaining comments

• Make sure that only the trailing tile size can be scalable (and add a negative test that verifies this)
• Refine the tests (remove references to #map)

[c-rhodes]

One minor nit but otherwise LGTM, nice work!

mlir/test/Dialect/Linalg/transform-op-tile.mlir

226

nit: for specifying

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

let description = [{
 
     If the internal implementation of tiling for any of the operations fails,
     produces a definite failure.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                    Variadic<TransformParamTypeOrAnyHandle>:$dynamic_sizes,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sizes,
-                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$interchange);
+                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$interchange,
+                   DefaultValuedOptionalAttr<BoolAttr, "false">:$last_tile_size_scalable);
   let results = (outs TransformHandleTypeInterface:$tiled_linalg_op,
                       Variadic<TransformHandleTypeInterface>:$loops);
   let builders = [
     OpBuilder<(ins "TypeRange":$loopTypes,
                    "Value":$target,
                    "ArrayRef<int64_t>":$staticTileSizes,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange)>,
     OpBuilder<(ins "TypeRange":$loopTypes,

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 ///
 /// where `values` is of ODS type `Variadic<*>` and `integers` is of ODS
 /// type `I64ArrayAttr`. Parse a mixed list with either (1) static integer
 /// values or (2) SSA values. Fill `integers` with the integer ArrayAttr, where
 /// `kDynamic` encodes the position of SSA values. Add the parsed SSA values
 /// to `values` in-order. If `valueTypes` is non-null, fill it with types
 /// corresponding to values; otherwise the caller must handle the types.
 ///
 /// E.g. after parsing "[%arg0 : index, 7, 42, %arg42 : i32]":
 ///   1. `result` is filled with the i64 ArrayAttr "[`kDynamic`, 7, 42,
 ///   `kDynamic`]"
 ///   2. `ssa` is filled with "[%arg0, %arg1]".
+///
+/// Trailing indices can be scalable. For example, "42" in "[7, [42]]" is
+/// scalable. This notation is similar to how scalable dims are marked when
+/// defining Vectors. If /p isTrailingIdxScalable is null, scalable indices are
+/// not allowed/expected. When it's not null, this hook will set the
+/// corresponding value to:
+///   * true if the trailing idx is scalable,
+///   * false otherwise.
 ParseResult parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
-    DenseI64ArrayAttr &integers, SmallVectorImpl<Type> *valueTypes = nullptr,
+    DenseI64ArrayAttr &integers, bool *isTrailingIdxScalable = nullptr,
+    SmallVectorImpl<Type> *valueTypes = nullptr,
     AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square);
 inline ParseResult parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
     DenseI64ArrayAttr &integers, SmallVectorImpl<Type> &valueTypes,
     AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square) {
-  return parseDynamicIndexList(parser, values, integers, &valueTypes,
+  return parseDynamicIndexList(parser, values, integers,
+                               /*isTrailingIdxScalable=*/nullptr, &valueTypes,
                                delimiter);
 }
 
 /// Verify that a the `values` has as many elements as the number of entries in
 /// `attr` for which `isDynamic` evaluates to true.
 LogicalResult verifyListOfOperandsOrIntegers(Operation *op, StringRef name,
                                              unsigned expectedNumElements,
                                              ArrayRef<int64_t> attr,
                                              ValueRange values);
 
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE_H_

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

for (Operation *op : dynamicSizeProducers.back()) {
       diag.attachNote(transformValue.getLoc()) << "for this handle";
       return diag;
     }
   }
 
   SmallVector<Operation *> tiled;
   SmallVector<SmallVector<Operation *, 4>, 4> loops;
   loops.resize(getLoops().size());
+  bool scalable = getLastTileSizeScalable();
   for (auto [i, op] : llvm::enumerate(targets)) {
     auto tilingInterface = dyn_cast<TilingInterface>(op);
     auto dpsInterface = dyn_cast<DestinationStyleOpInterface>(op);
     if (!tilingInterface || !dpsInterface) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "only ops implementing TilingInterface and "
                                     "DestinationStyleOpInterface are supported";
       diag.attachNote(op->getLoc()) << "target op";
       return diag;
     }
 
     scf::SCFTilingOptions tilingOptions;
     if (!tileSizes.empty()) {
       tilingOptions.setTileSizeComputationFunction([&, index = i](OpBuilder &b,
                                                                   Operation *) {
         SmallVector<Value, 4> sizes;
         sizes.reserve(tileSizes.size());
         unsigned dynamicIdx = 0;
-        for (OpFoldResult ofr : getMixedSizes()) {
+        unsigned trailingIdx = getMixedSizes().size() - 1;
+
+        for (auto [ofrIdx, ofr] : llvm::enumerate(getMixedSizes())) {
           if (auto attr = llvm::dyn_cast_if_present<Attribute>(ofr)) {
+            // Only the trailing tile size is allowed to be scalable atm.
+            if (scalable && (ofrIdx == trailingIdx)) {
+              auto val = b.create<arith::ConstantIndexOp>(
+                  getLoc(), attr.cast<IntegerAttr>().getInt());
+              Value vscale =
+                  b.create<vector::VectorScaleOp>(getLoc(), b.getIndexType());
+              sizes.push_back(b.create<arith::MulIOp>(getLoc(), val, vscale));
+            } else {
               sizes.push_back(b.create<arith::ConstantIndexOp>(
                   getLoc(), cast<IntegerAttr>(attr).getInt()));
+            }
             continue;
           }
           ArrayRef<Operation *> dynamicSizes = dynamicSizeProducers[dynamicIdx];
           ArrayRef<int64_t> params = paramSizes[dynamicIdx];
           ++dynamicIdx;
           assert((dynamicSizes.empty() ^ params.empty()) &&
                  "expected either dynamic sizes or parameters");
           if (!params.empty()) {

 
 ParseResult transform::TileOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
   OpAsmParser::UnresolvedOperand target;
   SmallVector<OpAsmParser::UnresolvedOperand> dynamicSizes;
   DenseI64ArrayAttr staticSizes;
   FunctionType functionalType;
   llvm::SMLoc operandLoc;
+  bool scalable = false;
   if (parser.parseOperand(target) || parser.getCurrentLocation(&operandLoc) ||
-      parseDynamicIndexList(parser, dynamicSizes, staticSizes) ||
+      parseDynamicIndexList(parser, dynamicSizes, staticSizes, &scalable) ||
       parseOptionalInterchange(parser, result) ||
       parser.parseColonType(functionalType))
     return ParseResult::failure();
 
   size_t numExpectedLoops =
       staticSizes.size() - llvm::count(staticSizes.asArrayRef(), 0);
   if (functionalType.getNumResults() != numExpectedLoops + 1) {
     return parser.emitError(parser.getNameLoc())
            << "expected " << (numExpectedLoops + 1) << " result type(s)";
   }
   if (functionalType.getNumInputs() != dynamicSizes.size() + 1) {
     return parser.emitError(operandLoc)
            << "expected " << dynamicSizes.size() + 1 << " operand type(s)";
   }
   if (parser.resolveOperand(target, functionalType.getInputs().front(),
                             result.operands) ||
       parser.resolveOperands(dynamicSizes,
                              functionalType.getInputs().drop_front(),
                              operandLoc, result.operands)) {
     return failure();
   }
 
+  auto scalableAttr = parser.getBuilder().getBoolAttr(scalable);
+  result.addAttribute(getLastTileSizeScalableAttrName(result.name),
+                      scalableAttr);
+
   result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
   result.addTypes(functionalType.getResults());
   return success();
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ' << getTarget();
   printDynamicIndexList(p, getOperation(), getDynamicSizes(), getStaticSizes());

mlir/lib/Dialect/SCF/IR/SCF.cpp

ParseResult ForallOp::parse(OpAsmParser &parser, OperationState &result) {
   if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren))
     return failure();
 
   DenseI64ArrayAttr staticLbs, staticUbs, staticSteps;
   SmallVector<OpAsmParser::UnresolvedOperand> dynamicLbs, dynamicUbs,
       dynamicSteps;
   if (succeeded(parser.parseOptionalKeyword("in"))) {
     // Parse upper bounds.
-    if (parseDynamicIndexList(parser, dynamicUbs, staticUbs,
-                              /*valueTypes=*/nullptr,
-                              OpAsmParser::Delimiter::Paren) ||
+    if (parseDynamicIndexList(
+            parser, dynamicUbs, staticUbs, /*scalable=*/nullptr,
+            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicUbs, indexType, result.operands))
       return failure();
 
     unsigned numLoops = ivs.size();
     staticLbs = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 0));
     staticSteps = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 1));
   } else {
     // Parse lower bounds.
     if (parser.parseEqual() ||
-        parseDynamicIndexList(parser, dynamicLbs, staticLbs,
-                              /*valueTypes=*/nullptr,
-                              OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(
+            parser, dynamicLbs, staticLbs, /*scalable=*/nullptr,
+            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
 
         parser.resolveOperands(dynamicLbs, indexType, result.operands))
       return failure();
 
     // Parse upper bounds.
     if (parser.parseKeyword("to") ||
-        parseDynamicIndexList(parser, dynamicUbs, staticUbs,
-                              /*valueTypes=*/nullptr,
-                              OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(
+            parser, dynamicUbs, staticUbs, /*scalable=*/nullptr,
+            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicUbs, indexType, result.operands))
       return failure();
 
     // Parse step values.
     if (parser.parseKeyword("step") ||
-        parseDynamicIndexList(parser, dynamicSteps, staticSteps,
-                              /*valueTypes=*/nullptr,
-                              OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(
+            parser, dynamicSteps, staticSteps, /*scalable=*/nullptr,
+            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicSteps, indexType, result.operands))
       return failure();
   }
 
   // Parse out operands and results.
   SmallVector<OpAsmParser::Argument, 4> regionOutArgs;
   SmallVector<OpAsmParser::UnresolvedOperand, 4> outOperands;
   SMLoc outOperandsLoc = parser.getCurrentLocation();

mlir/lib/Dialect/Transform/Utils/Utils.cpp

if (parser.parseLParen().failed() ||
         parser.parseRParen().failed()) {
       return failure();
     }
     packed.emplace(packedOperand);
     integers = parser.getBuilder().getDenseI64ArrayAttr({});
     return success();
   }
 
-  return parseDynamicIndexList(parser, values, integers, &valueTypes);
+  return parseDynamicIndexList(parser, values, integers, /*scalable=*/nullptr,
+                               &valueTypes);
 }

mlir/lib/Interfaces/ViewLikeInterface.cpp

if (ShapedType::isDynamic(integer)) {
       ++idx;
     } else {
       printer << integer;
     }
   });
   printer << rightDelimiter;
 }
 
 ParseResult mlir::parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
-    DenseI64ArrayAttr &integers, SmallVectorImpl<Type> *valueTypes,
-    AsmParser::Delimiter delimiter) {
+    DenseI64ArrayAttr &integers, bool *isTrailingIdxScalable,
+    SmallVectorImpl<Type> *valueTypes, AsmParser::Delimiter delimiter) {
 
   SmallVector<int64_t, 4> integerVals;
+  bool foundScalable = false;
   auto parseIntegerOrValue = [&]() {
     OpAsmParser::UnresolvedOperand operand;
     auto res = parser.parseOptionalOperand(operand);
+
+    // If `foundScalable` has already been set to `true` then a non-trailing
+    // tile size was identified as scalable.
+    if (foundScalable) {
+      parser.emitError(parser.getNameLoc())
+          << "non-trailing tile size cannot be scalable";
+      return failure();
+    }
+
+    if (isTrailingIdxScalable && parser.parseOptionalLSquare().succeeded())
+      foundScalable = true;
+
     if (res.has_value() && succeeded(res.value())) {
       values.push_back(operand);
       integerVals.push_back(ShapedType::kDynamic);
       if (valueTypes && parser.parseColonType(valueTypes->emplace_back()))
         return failure();
     } else {
       int64_t integer;
       if (failed(parser.parseInteger(integer)))
         return failure();
       integerVals.push_back(integer);
     }
+    if (foundScalable && parser.parseOptionalRSquare().failed())
+      return failure();
     return success();
   };
   if (parser.parseCommaSeparatedList(delimiter, parseIntegerOrValue,
                                      " in dynamic index list"))
     return parser.emitError(parser.getNameLoc())
            << "expected SSA value or integer";
   integers = parser.getBuilder().getDenseI64ArrayAttr(integerVals);
+  if (isTrailingIdxScalable)
+    *isTrailingIdxScalable = foundScalable;
   return success();
 }
 
 bool mlir::detail::sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
     llvm::function_ref<bool(OpFoldResult, OpFoldResult)> cmp) {
   if (a.static_offsets().size() != b.static_offsets().size())
     return false;
   if (a.static_sizes().size() != b.static_sizes().size())
     return false;

mlir/test/Dialect/Linalg/transform-op-tile.mlir

-// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file --verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --mlir-print-local-scope --split-input-file --verify-diagnostics %s | FileCheck %s
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1, %loops:3 = transform.structured.tile %0 [4, 4, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
 }
 
 // CHECK-LABEL: func @tile_linalg_matmul(

 }
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.structured.tile_to_forall_op %0 tile_sizes[1, 1]
          : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+
+module {
+  func.func @scalable_tile(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>, %arg3: f32) -> tensor<?xf32> {
+    %0 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) outs(%arg2 : tensor<?xf32>) {
+    ^bb0(%in_1: f32, %in_2: f32, %out: f32):
+      %1 = arith.addf %in_1, %in_2 : f32
+      %2 = arith.mulf %arg3, %1 : f32
+      linalg.yield %2 : f32
+    } -> tensor<?xf32>
+    return %0 : tensor<?xf32>
+  }
+}
+
+// CHECK-LABEL:   func.func @scalable_tile(
+// CHECK-SAME:      %[[ARG_0:.*]]: tensor<?xf32>, %[[ARG_1:.*]]: tensor<?xf32>, %[[ARG_2:.*]]: tensor<?xf32>,
+// CHECK:           %[[C4:.*]] = arith.constant 0 : index
+// CHECK:           %[[DIM:.*]] = tensor.dim %[[ARG_0]], %[[C4]] : tensor<?xf32>
+// CHECK:           %[[VEC_SIZE:.*]] = arith.constant 4 : index
+// CHECK:           %[[VS:.*]] = vector.vscale
+// CHECK:           %[[STEP:.*]] = arith.muli %[[VEC_SIZE]], %[[VS]] : index
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           scf.for %[[IV:.*]] = %[[C0]] to %[[DIM]] step %[[STEP]] iter_args(%[[VAL:.*]] = %[[ARG_2]]) -> (tensor<?xf32>) {
+// CHECK:             %[[SIZE:.*]] = affine.min affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>(%[[IV]])[%[[STEP]], %[[DIM]]]
+// CHECK:             %[[SLICE_ARG0:.*]] = tensor.extract_slice %[[ARG_0]][%[[IV]]] [%[[SIZE]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK:             %[[SLICE_ARG1:.*]] = tensor.extract_slice %[[ARG_1]][%[[IV]]] [%[[SIZE]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK:             %[[SLICE_ARG2:.*]] = tensor.extract_slice %[[VAL]][%[[IV]]] [%[[SIZE]]] [1] : tensor<?xf32> to tensor<?xf32>
+// CHECK:             linalg.generic {indexing_maps = {{.*}}, iterator_types = ["parallel"]} ins(%[[SLICE_ARG0]], %[[SLICE_ARG1]] : tensor<?xf32>, tensor<?xf32>) outs(%[[SLICE_ARG2]] : tensor<?xf32>) {
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !transform.any_op):
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loop = transform.structured.tile %0 [[4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @scalable_and_fixed_length_tile
+// CHECK:           %[[STEP_0:.*]] = arith.constant 4 : index
+// CHECK:           %[[STEP_1:.*]] = arith.constant 4 : index
+// CHECK:           %[[C4:.*]] = arith.constant 4 : index
+// CHECK:           %[[VS:.*]] = vector.vscale
+// CHECK:           %[[STEP_2:.*]] = arith.muli %[[C4]], %[[VS]] : index
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           %[[C128:.*]] = arith.constant 128 : index
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[C0]] to %[[C128]] step %[[STEP_0]]
+// CHECK:             %[[C0_1:.*]] = arith.constant 0 : index
+// CHECK:             %[[C128_1:.*]] = arith.constant 128 : index
+// CHECK:             scf.for %[[VAL_16:.*]] = %[[C0_1]] to %[[C128_1]] step %[[STEP_1]]
+// CHECK:               %[[C0_2:.*]] = arith.constant 0 : index
+// CHECK:               %[[C128_2:.*]] = arith.constant 128 : index
+// CHECK:               scf.for %{{.*}} = %[[C0_2]] to %[[C128_2]] step %[[STEP_2]]
+
+func.func @scalable_and_fixed_length_tile(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32> {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+
+  return %0 : tensor<128x128xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1, %loops:3 = transform.structured.tile %0 [4, 4, [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
+// TODO: Add support for for specyfying more than one scalable tile size
+
+func.func @scalable_and_fixed_length_tile(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32> {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+
+  return %0 : tensor<128x128xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-error @below {{non-trailing tile size cannot be scalable}}
+  // expected-error @below {{expected SSA value or integer}}
+  %1, %loops:3 = transform.structured.tile %0 [4, [4], [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}