[mlir][transform] Allow arbitrary indices to be scalable

[awarzynski]

Summary

This change lifts the limitation that only the trailing dimensions/sizes
in dynamic index lists can be scalable. It allows us to extend
MaskedVectorizeOp and TileOp from the Transform dialect so that the
following is allowed:

%1, %loops:3 = transform.structured.tile %0 [4, [4], [4]]

This is also a follow up for https://reviews.llvm.org/D153372
that will enable the following (middle vector dimension is scalable):

transform.structured.masked_vectorize %0 vector_sizes [2, [4], 8]

To facilate this change, the hooks for parsing and printing dynamic
index lists are updated accordingly (printDynamicIndexList and
parseDynamicIndexList, respectively). MaskedVectorizeOp and TileOp
are updated to include an array of attribute of bools that captures
whether the corresponding vector dimension/tile size, respectively, are
scalable or not.

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

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

[c-rhodes]

I've left a few bits but otherwise LGTM, cheers

mlir/lib/Interfaces/ViewLikeInterface.cpp

105

nit: I count 5 different variable names now for the same concept:

• scalableDims
• scalable_sizes
• scalableVals
• scalableIdxs
• scalables

I think it would be clearer to use consistent naming.

149–150

150–153

166

If the intention here is to check there are scalableVals then not empty would be clearer.

[awarzynski]

Thanks Cullen, I will make the suggested adjustments before merging.

mlir/lib/Interfaces/ViewLikeInterface.cpp

105

Meh, naming is hard. I agree, but IMHO it's a wider problem and not directly related to this patch.

• The definition of MaskedVectorizeOp contains $vector_sizes and $static_sizes This patch adds scalable_sizes, which is consistent with the naming scheme in LinalgTransformOps.td
• The signature of printDynamicIndexList contains values and integers. This patch adds scalables, which is consistent with what's already there.

You are right that the names in *.td and *.cpp files are inconsistent. However, this inconsistency predates this change. Having said that, scalableIdxs should be replaced with scalables (for consistency). I will update that.

scalableVals is a local function variable and is consistent with integerVals (per-exisiting function variable serving similar purpose).

scalableDims is only used in the context of VectorType in which the concept of a "dimension" makes sense (it does not make sense in the context of e.g. tiles).

Thanks for raising this! I will make a few adjustments that hopefully will reduce the confusion.

150–153

I swear you suggested something similar recently in one of my patches. I need to pay more attention - thank you!

[mehdi_amini]

https://lab.llvm.org/buildbot/#/builders/61/builds/45451

The bot seems broken by this, can you have a look please?

[pifon2a]

I reverted it for now https://github.com/llvm/llvm-project/commit/14d073b50f960674a62ef8ad2c34f6fc1e9b0061

[awarzynski]

In D154336#4470333, @pifon2a wrote:

I reverted it for now https://github.com/llvm/llvm-project/commit/14d073b50f960674a62ef8ad2c34f6fc1e9b0061

Apologies for the breakage and thanks for reverting!

I keep MLIR_ENABLE_BINDINGS_PYTHON as default (off), so that's why I didn't see that :( will be AFK most of today, so reverting was the right call, thanks!

[awarzynski]

So it looks like the failure was due to missing changes in the python bindings - I am about to send an update.

I will go ahead and re-apply this patch sometime tomorrow unless there are new comments.

[awarzynski]

Update Python bindings

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<BoolAttr, "false">:$last_tile_size_scalable);
+                   DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes);
   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,
let description = [{
     values) do not satisfy the constraints mentioned above. It produces a
     silenceable failure if at least one target op is not a Linalg op or fails to
     vectorize.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        Variadic<TransformHandleTypeInterface>:$vector_sizes,
                        UnitAttr:$vectorize_nd_extract,
+                       DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:
+                          $scalable_sizes,
                        DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
-                          $static_vector_sizes,
-                       DefaultValuedOptionalAttr<BoolAttr, "false">:$last_vector_size_scalable);
+                          $static_vector_sizes);
 
   let results = (outs);
   let assemblyFormat = [{
       $target
       `vector_sizes` custom<DynamicIndexList>($vector_sizes,
                                               $static_vector_sizes,
                                               type($vector_sizes),
-                                              $last_vector_size_scalable)
+                                              $scalable_sizes)
       attr-dict
       `:` type($target)
   }];
 
   let extraClassDeclaration = [{
     // TODO: applyToOne.
     ::mlir::DiagnosedSilenceableFailure apply(
         ::mlir::transform::TransformRewriter &rewriter,

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 /// where `values` is of ODS type `Variadic<*>` and `integers` is of ODS
 /// type `I64ArrayAttr`. Prints a list with either (1) the static integer value
 /// in `integers` is `kDynamic` or (2) the next value otherwise. If `valueTypes`
 /// is non-empty, it is expected to contain as many elements as `values`
 /// indicating their types. This allows idiomatic printing of mixed value and
 /// integer attributes in a list. E.g.
 /// `[%arg0 : index, 7, 42, %arg42 : i32]`.
 ///
-/// If  `isTrailingIdxScalable` is true, then wrap the trailing index with
-/// square brackets, e.g. `[42]`, to denote scalability. This would normally be
-/// used for scalable tile or vector sizes.
+/// Indices can be scalable. For example, "4" in "[2, [4], 8]" is scalable.
+/// This notation is similar to how scalable dims are marked when defining
+/// Vectors. For each value in `integers`, the corresponding `bool` in
+/// `scalables` encodes whether it's a scalable index. If `scalableVals` is
+/// empty then assume that all indices are non-scalable.
 void printDynamicIndexList(
     OpAsmPrinter &printer, Operation *op, OperandRange values,
     ArrayRef<int64_t> integers, TypeRange valueTypes = TypeRange(),
-    BoolAttr isTrailingIdxScalable = {},
+    ArrayRef<bool> scalables = {},
     AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square);
 
 /// Parser hook for custom directive in assemblyFormat.
 ///
 ///   custom<DynamicIndexList>($values, $integers)
 ///   custom<DynamicIndexList>($values, $integers, type($values))
 ///
 /// 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.
+/// Indices can be scalable. For example, "4" in "[2, [4], 8]" is scalable.
+/// This notation is similar to how scalable dims are marked when defining
+/// Vectors. For each value in `integers`, the corresponding `bool` in
+/// `scalableVals` encodes whether it's a scalable index.
 ParseResult parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
-    DenseI64ArrayAttr &integers, bool *isTrailingIdxScalable = nullptr,
+    DenseI64ArrayAttr &integers, DenseBoolArrayAttr &scalableVals,
     SmallVectorImpl<Type> *valueTypes = nullptr,
     AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square);
 inline ParseResult parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+    DenseI64ArrayAttr &integers, SmallVectorImpl<Type> *valueTypes = nullptr,
+    AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square) {
+  DenseBoolArrayAttr scalableVals = {};
+  return parseDynamicIndexList(parser, values, integers, scalableVals,
+                               valueTypes, delimiter);
+}
+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,
-                               /*isTrailingIdxScalable=*/nullptr, &valueTypes,
-                               delimiter);
+  DenseBoolArrayAttr scalableVals = {};
+  return parseDynamicIndexList(parser, values, integers, scalableVals,
+                               &valueTypes, delimiter);
 }
 inline ParseResult parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
     DenseI64ArrayAttr &integers, SmallVectorImpl<Type> &valueTypes,
-    BoolAttr &isTrailingIdxScalable,
+    DenseBoolArrayAttr &scalableVals,
     AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square) {
 
-  bool scalable = false;
-  auto res = parseDynamicIndexList(parser, values, integers, &scalable,
-                                   &valueTypes, delimiter);
-  auto scalableAttr = parser.getBuilder().getBoolAttr(scalable);
-  isTrailingIdxScalable = scalableAttr;
-  return res;
+  return parseDynamicIndexList(parser, values, integers, scalableVals,
 }
 
 /// 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();
+  auto scalableSizes = getScalableSizes();
   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;
-        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)) {
+            if (scalableSizes[ofrIdx]) {
               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()));

 
 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;
+  DenseBoolArrayAttr scalableVals;
+
   if (parser.parseOperand(target) || parser.getCurrentLocation(&operandLoc) ||
-      parseDynamicIndexList(parser, dynamicSizes, staticSizes, &scalable) ||
+      parseDynamicIndexList(parser, dynamicSizes, staticSizes, scalableVals) ||
       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(getScalableSizesAttrName(result.name), scalableVals);
 
   result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
   result.addTypes(functionalType.getResults());
   return success();
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ' << getTarget();
   printDynamicIndexList(p, getOperation(), getDynamicSizes(), getStaticSizes(),
-                        /*valueTypes=*/{}, getLastTileSizeScalableAttr(),
+                        /*valueTypes=*/{}, getScalableSizesAttr(),
                         OpAsmParser::Delimiter::Square);
   printOptionalInterchange(p, getInterchange());
   p << " : ";
   p.printFunctionalType(getOperands().getTypes(), getResults().getTypes());
 }
 
 void transform::TileOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
if (!matchPattern(szPayloadOp->getResult(0), m_Constant(&attr))) {
       diag.attachNote(szPayloadOp->getLoc()) << "vector size payload op";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
 
     vectorSizes.push_back(attr.getInt());
   }
 
   // TODO: Check that the correct number of vectorSizes was provided.
-  SmallVector<bool> scalableVecDims(vectorSizes.size(), false);
-  scalableVecDims.back() = getLastVectorSizeScalable();
   for (Operation *target : targets) {
     if (!isa<linalg::LinalgOp, tensor::PadOp>(target)) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Unsupported Op, cannot vectorize";
     }
 
-    if (failed(linalg::vectorize(rewriter, target, vectorSizes, scalableVecDims,
+    if (failed(linalg::vectorize(rewriter, target, vectorSizes,
+                                 getScalableSizes(),
                                  getVectorizeNdExtract()))) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Attempted to vectorize, but failed";
     }
   }
 
   return DiagnosedSilenceableFailure::success();
 }

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

 }
 
 void ForallOp::print(OpAsmPrinter &p) {
   Operation *op = getOperation();
   p << " (" << getInductionVars();
   if (isNormalized()) {
     p << ") in ";
     printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),
-                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          /*valueTypes=*/{}, /*scalables=*/{},
                           OpAsmParser::Delimiter::Paren);
   } else {
     p << ") = ";
     printDynamicIndexList(p, op, getDynamicLowerBound(), getStaticLowerBound(),
-                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          /*valueTypes=*/{}, /*scalables=*/{},
                           OpAsmParser::Delimiter::Paren);
     p << " to ";
     printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),
-                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          /*valueTypes=*/{}, /*scalables=*/{},
                           OpAsmParser::Delimiter::Paren);
     p << " step ";
     printDynamicIndexList(p, op, getDynamicStep(), getStaticStep(),
-                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          /*valueTypes=*/{}, /*scalable=*/{},
                           OpAsmParser::Delimiter::Paren);
   }
   printInitializationList(p, getRegionOutArgs(), getOutputs(), " shared_outs");
   p << " ";
   if (!getRegionOutArgs().empty())
     p << "-> (" << getResultTypes() << ") ";
   p.printRegion(getRegion(),
                 /*printEntryBlockArgs=*/false,
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, /*isTrailingIdxScalable=*/nullptr,
-            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
+    if (parseDynamicIndexList(parser, dynamicUbs, staticUbs,
+                              /*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, /*isTrailingIdxScalable=*/nullptr,
-            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(parser, dynamicLbs, staticLbs,
+                              /*valueTypes=*/nullptr,
+                              OpAsmParser::Delimiter::Paren) ||
 
         parser.resolveOperands(dynamicLbs, indexType, result.operands))
       return failure();
 
     // Parse upper bounds.
     if (parser.parseKeyword("to") ||
-        parseDynamicIndexList(
-            parser, dynamicUbs, staticUbs, /*isTrailingIdxScalable=*/nullptr,
-            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(parser, dynamicUbs, staticUbs,
+                              /*valueTypes=*/nullptr,
+                              OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicUbs, indexType, result.operands))
       return failure();
 
     // Parse step values.
     if (parser.parseKeyword("step") ||
-        parseDynamicIndexList(
-            parser, dynamicSteps, staticSteps, /*scalable=*/nullptr,
-            /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
+        parseDynamicIndexList(parser, dynamicSteps, staticSteps,
+                              /*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,
-                               /*isTrailingIdxScalable=*/nullptr, &valueTypes);
+  return parseDynamicIndexList(parser, values, integers, &valueTypes);
 }

mlir/lib/Interfaces/ViewLikeInterface.cpp

static char getRightDelimiter(AsmParser::Delimiter delimiter) {
   default:
     llvm_unreachable("unsupported delimiter");
   }
 }
 
 void mlir::printDynamicIndexList(OpAsmPrinter &printer, Operation *op,
                                  OperandRange values,
                                  ArrayRef<int64_t> integers,
-                                 TypeRange valueTypes,
-                                 BoolAttr isTrailingIdxScalable,
+                                 TypeRange valueTypes, ArrayRef<bool> scalables,
                                  AsmParser::Delimiter delimiter) {
   char leftDelimiter = getLeftDelimiter(delimiter);
   char rightDelimiter = getRightDelimiter(delimiter);
   printer << leftDelimiter;
   if (integers.empty()) {
     printer << rightDelimiter;
     return;
   }
 
-  int64_t trailingScalableInteger;
-  if (isTrailingIdxScalable && isTrailingIdxScalable.getValue()) {
-    // ATM only the trailing idx can be scalable
-    trailingScalableInteger = integers.back();
-    integers = integers.drop_back();
-  }
-
-  unsigned idx = 0;
+  unsigned dynamicValIdx = 0;
+  unsigned scalableIndexIdx = 0;
   llvm::interleaveComma(integers, printer, [&](int64_t integer) {
+    if (not scalables.empty() && scalables[scalableIndexIdx])
+      printer << "[";
     if (ShapedType::isDynamic(integer)) {
-      printer << values[idx];
+      printer << values[dynamicValIdx];
       if (!valueTypes.empty())
-        printer << " : " << valueTypes[idx];
-      ++idx;
+        printer << " : " << valueTypes[dynamicValIdx];
+      ++dynamicValIdx;
     } else {
       printer << integer;
     }
-  });
-
-  // Print the trailing scalable index
-  if (isTrailingIdxScalable && isTrailingIdxScalable.getValue()) {
-    if (!integers.empty())
-      printer << ", ";
-    printer << "[";
-    printer << trailingScalableInteger;
+    if (!scalables.empty() && scalables[scalableIndexIdx])
       printer << "]";
-  }
+
+    scalableIndexIdx++;
+  });
 
   printer << rightDelimiter;
 }
 
 ParseResult mlir::parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
-    DenseI64ArrayAttr &integers, bool *isTrailingIdxScalable,
+    DenseI64ArrayAttr &integers, DenseBoolArrayAttr &scalables,
     SmallVectorImpl<Type> *valueTypes, AsmParser::Delimiter delimiter) {
 
   SmallVector<int64_t, 4> integerVals;
-  bool foundScalable = false;
+  SmallVector<bool, 4> scalableVals;
   auto parseIntegerOrValue = [&]() {
     OpAsmParser::UnresolvedOperand operand;
     auto res = parser.parseOptionalOperand(operand);
 
-    // If `foundScalable` has already been set to `true` then a non-trailing
-    // index was identified as scalable.
-    if (foundScalable) {
-      parser.emitError(parser.getNameLoc())
-          << "non-trailing index cannot be scalable";
-      return failure();
-    }
-
-    if (isTrailingIdxScalable && parser.parseOptionalLSquare().succeeded())
-      foundScalable = true;
+    // When encountering `[`, assume that this is a scalable index.
+    scalableVals.push_back(parser.parseOptionalLSquare().succeeded());
 
     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())
+
+    // If this is assumed to be a scalable index, verify that there's a closing
+    // `]`.
+    if (scalableVals.back() && 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;
+  scalables = parser.getBuilder().getDenseBoolArrayAttr(scalableVals);
   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;

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
     from ..ir import *
     from ._ods_common import get_op_result_or_value as _get_op_result_or_value
     from ..dialects import pdl, transform
 except ImportError as e:
     raise RuntimeError("Error loading imports from extension module") from e
 
 from typing import List, Optional, Sequence, Union, overload
 
 IntOrAttrList = Sequence[Union[IntegerAttr, int]]
 OptionalIntList = Optional[Union[ArrayAttr, IntOrAttrList]]
 
+BoolOrAttrList = Sequence[Union[BoolAttr, bool]]
+OptionalBoolList = Optional[Union[ArrayAttr, BoolOrAttrList]]
+
 
 def _get_int_int_array_attr(
     values: Optional[Union[ArrayAttr, Sequence[Union[ArrayAttr, IntOrAttrList]]]]
 ) -> ArrayAttr:
     """Creates an array attribute containing array attributes of integers.
 
     If the operand is already an array attribute, forwards it. Otherwise treats
     the operand as a list of attributes or integers, potentially interpserced, to
def __init__(
         self,
         loop_types: Union[Type, List[Type]],
         target: Union[Operation, Value],
         *,
         sizes: Optional[
             Union[Sequence[Union[int, IntegerAttr, Operation, Value]], ArrayAttr]
         ] = None,
         interchange: OptionalIntList = None,
+        scalable_sizes: OptionalBoolList = None,
         loc=None,
         ip=None,
     ):
         ...
 
     @overload
     def __init__(
         self,
         target: Union[Operation, Value, OpView],
         *,
         sizes: Optional[
             Union[Sequence[Union[int, IntegerAttr, Operation, Value]], ArrayAttr]
         ] = None,
         interchange: OptionalIntList = None,
+        scalable_sizes: OptionalBoolList = None,
         loc=None,
         ip=None,
     ):
         ...
 
     def __init__(
         self,
         loop_types_or_target: Union[Type, List[Type], Operation, Value],
         target_or_none: Optional[Union[Operation, Value, OpView]] = None,
         *,
         sizes: Optional[
             Union[Sequence[Union[int, IntegerAttr, Operation, Value]], ArrayAttr]
         ] = None,
         interchange: OptionalIntList = None,
+        scalable_sizes: OptionalBoolList = None,
         loc=None,
         ip=None,
     ):
         if interchange is None:
             interchange = []
         if sizes is None:
             sizes = []
+        if scalable_sizes is None:
+            scalable_sizes = []
 
         static_sizes = []
         dynamic_sizes = []
         if isinstance(sizes, ArrayAttr):
             sizes_attr = sizes
         else:
             for size in sizes:
                 if isinstance(size, int):
):
 
         super().__init__(
             target.type,
             loop_types,
             target,
             dynamic_sizes=dynamic_sizes,
             static_sizes=sizes_attr,
             interchange=interchange,
+            scalable_sizes=scalable_sizes,
             loc=loc,
             ip=ip,
         )
 
     def __extract_values(self, attr: Optional[DenseI64ArrayAttr]) -> List[int]:
         if not attr:
             return []
         return [element for element in attr]

mlir/python/mlir/ir.py

 def _f64ArrayAttr(x, context):
     return ArrayAttr.get([_f64Attr(v, context) for v in x])
 
 
 @register_attribute_builder("DenseI64ArrayAttr")
 def _denseI64ArrayAttr(x, context):
     return DenseI64ArrayAttr.get(x, context=context)
 
+@register_attribute_builder("DenseBoolArrayAttr")
+def _denseBoolArrayAttr(x, context):
+    return DenseBoolArrayAttr.get(x, context=context)
+
 
 @register_attribute_builder("TypeAttr")
 def _typeAttr(x, context):
     return TypeAttr.get(x, context=context)
 
 
 @register_attribute_builder("TypeArrayAttr")
 def _typeArrayAttr(x, context):

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

func.func @scalable_and_fixed_length_tile(
   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 index 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)
-}

mlir/test/Dialect/Transform/ops.mlir

 
 // CHECK: transform.sequence
 // CHECK: transform.structured.tile %0[4, 4, [4]]
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.structured.tile %0 [4, 4, [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
 }
+
+// CHECK: transform.sequence
+// CHECK: transform.structured.tile %0{{\[}}[2], 4, 8]
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.structured.tile %0 [[2], 4, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}