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

[awarzynski]

Summary

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

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

This is implemented by extending the MaskedVectorizeOp with a dedicated
attribute for "scalability" and by creating a custom parser. At the
moment, only the trailing vec size can be scalable. The following is not
yet supported:

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

As the vectoriser does not support scalable vectorisation just yet, it's
tricky to test this change. In the meantime, use the debug output,
--debug-only=linalg-vectorization, to check whether scalable
vectorisation has been switched on.

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/

Similar patch for tiling: https://reviews.llvm.org/D150944

[dcaballe]

Thanks! It looks like a good starting point to me. I should be able to add support to the vectorizer with this and the proper public API. Leaving the TD approval to Alex or Nicolas. I'm ok with landing this in the meantime.

[ftynse]

Can we please have a parser roundtrip test for the changes? Printer and parser don't seem to match.

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

3100

Nit: don't hardcode 4 here.

3123–3127

Why is this necessary?

3133

Why can't this just parseType?

3143–3145

Shouldn't this be a property?

3153

I think this needs to account for isLastVectorSizeScalable.

3154

This should also print the attribute dictionary.

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

1538

If you are looking for a way to test this before it's fully supported, how about op->emitWarning("scalable vectorization not supported yet") that can be checked with --verify-diagnostics?

[ftynse]

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

3113–3116

Would it be possible to use the custom directive for this, additionally taking the "scalable" attribute? The code looks long and boilerplaty.

[awarzynski]

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

3113–3116

Sure, thanks for the suggestion! I will upload a new version shortly.

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

1538

Thanks, that's a great suggestion!

[awarzynski]

Switch back to using custom directive, add a test.

I was able to get rid of the custom printer/parser, but had to tweak a few other things. Nothing major.

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

let description = [{
     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<DenseI64ArrayAttr, "{}">:
-                          $static_vector_sizes);
+                          $static_vector_sizes,
+                       DefaultValuedOptionalAttr<BoolAttr, "false">:$last_vector_size_scalable);
+
   let results = (outs);
   let assemblyFormat = [{
       $target
       `vector_sizes` custom<DynamicIndexList>($vector_sizes,
                                               $static_vector_sizes,
-                                              type($vector_sizes))
+                                              type($vector_sizes),
+                                              $last_vector_size_scalable)
       attr-dict
       `:` type($target)
   }];
 
   let extraClassDeclaration = [{
     // TODO: applyToOne.
     ::mlir::DiagnosedSilenceableFailure apply(
         ::mlir::transform::TransformResults &transformResults,

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

 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
 /// must match the rank of the iteration space of the operation and the sizes
 /// must be smaller or equal than their counterpart interation space sizes, if
 /// static. `inputVectorShapes` also allows the vectorization of operations with
 /// dynamic shapes.
 LogicalResult vectorize(RewriterBase &rewriter, Operation *op,
                         ArrayRef<int64_t> inputVectorSizes = {},
-                        bool vectorizeNDExtract = false);
+                        bool vectorizeNDExtract = false,
+                        bool lastVectorSizeScalable = false);
 
 /// Emit a suitable vector form for a Copy op with fully static shape.
 LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
 
 /// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToLoops(RewriterBase &rewriter,
                                        LinalgOp linalgOp);
 

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 /// `[%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.
 void printDynamicIndexList(
     OpAsmPrinter &printer, Operation *op, OperandRange values,
     ArrayRef<int64_t> integers, TypeRange valueTypes = TypeRange(),
-    AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square,
-    bool isTrailingIdxScalable = false);
+    BoolAttr isTrailingIdxScalable = {},
+    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
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);
 }
+inline ParseResult parseDynamicIndexList(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+    DenseI64ArrayAttr &integers, SmallVectorImpl<Type> &valueTypes,
+    BoolAttr &isTrailingIdxScalable,
+    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;
+}
 
 /// 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

ParseResult transform::TileOp::parse(OpAsmParser &parser,
   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=*/{}, OpAsmParser::Delimiter::Square,
-                        getLastTileSizeScalable());
+                        /*valueTypes=*/{}, getLastTileSizeScalableAttr(),
+                        OpAsmParser::Delimiter::Square);
   printOptionalInterchange(p, getInterchange());
   p << " : ";
   p.printFunctionalType(getOperands().getTypes(), getResults().getTypes());
 }
 
 void transform::TileOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
transform::VectorizeOp::applyToOne(Operation *target,
 
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // MaskedVectorizeOp
 //===----------------------------------------------------------------------===//
-
 DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
     mlir::transform::TransformResults &transformResults,
     mlir::transform::TransformState &state) {
   TrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
   auto targets = state.getPayloadOps(getTarget());
   if (std::empty(targets))
     return DiagnosedSilenceableFailure::success();
 
   SmallVector<int64_t> vectorSizes;
   for (OpFoldResult sz : getMixedVectorSizes()) {
     if (sz.is<Attribute>()) {
       auto attr = sz.get<Attribute>();
       vectorSizes.push_back(cast<IntegerAttr>(attr).getInt());
       continue;
     }
 
     auto szPayloads = state.getPayloadOps(sz.get<Value>());
     if (!llvm::hasSingleElement(szPayloads)) {
       auto diag = this->emitOpError(
           "requires vector size handle that is mapped to 1 payload op");
       diag.attachNote(sz.get<Value>().getLoc())
           << "mapped to " << llvm::range_size(szPayloads) << " payload ops";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
 
     Operation *szPayloadOp = *szPayloads.begin();
     if (szPayloadOp->getNumResults() != 1 ||
         !szPayloadOp->getResult(0).getType().isIndex()) {
       auto diag = this->emitOpError(
           "requires vector size payload op with 1 index result");
       diag.attachNote(szPayloadOp->getLoc()) << "vector size payload op";
       return DiagnosedSilenceableFailure::definiteFailure();
     }
 
     IntegerAttr attr;
     if (!matchPattern(szPayloadOp->getResult(0), m_Constant(&attr))) {
       auto diag = this->emitOpError("requires constant vector size");
       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.
 
   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,
-                                 getVectorizeNdExtract()))) {
+                                 getVectorizeNdExtract(),
+                                 getLastVectorSizeScalable()))) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Attempted to vectorize, but failed";
     }
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::MaskedVectorizeOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getVectorSizes(), effects);
   modifiesPayload(effects);

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

 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
 /// must match the rank of the iteration space of the operation and the input
 /// vector sizes must be greater than or equal to their counterpart iteration
 /// space sizes, if static. `inputVectorShapes` also allows the vectorization of
 /// operations with dynamic shapes.
 LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter, Operation *op,
                                       ArrayRef<int64_t> inputVectorSizes,
-                                      bool vectorizeNDExtract) {
+                                      bool vectorizeNDExtract,
+                                      bool lastVectorSizeScalable) {
   LDBG("Attempting to vectorize:\n" << *op << "\n");
   LDBG("Input vector sizes: ");
   LLVM_DEBUG(llvm::interleaveComma(inputVectorSizes, llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << "\n");
+  LDBG("Scalable vectorisation: " << lastVectorSizeScalable << "\n");
+
+  if (lastVectorSizeScalable)
+    op->emitWarning("Scalable vectorization is not supported yet");
 
   if (failed(
           vectorizeOpPrecondition(op, inputVectorSizes, vectorizeNDExtract))) {
     LDBG("Vectorization pre-conditions failed\n");
     return failure();
   }
 
   // Initialize vectorization state.

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=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          OpAsmParser::Delimiter::Paren);
   } else {
     p << ") = ";
     printDynamicIndexList(p, op, getDynamicLowerBound(), getStaticLowerBound(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          OpAsmParser::Delimiter::Paren);
     p << " to ";
     printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          OpAsmParser::Delimiter::Paren);
     p << " step ";
     printDynamicIndexList(p, op, getDynamicStep(), getStaticStep(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, /*=isTrailingIdxScalable=*/{},
+                          OpAsmParser::Delimiter::Paren);
   }
   printInitializationList(p, getRegionOutArgs(), getOutputs(), " shared_outs");
   p << " ";
   if (!getRegionOutArgs().empty())
     p << "-> (" << getResultTypes() << ") ";
   p.printRegion(getRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/getNumResults() > 0);

mlir/lib/Interfaces/ViewLikeInterface.cpp

default:
     llvm_unreachable("unsupported delimiter");
   }
 }
 
 void mlir::printDynamicIndexList(OpAsmPrinter &printer, Operation *op,
                                  OperandRange values,
                                  ArrayRef<int64_t> integers,
                                  TypeRange valueTypes,
-                                 AsmParser::Delimiter delimiter,
-                                 bool isTrailingIdxScalable) {
+                                 BoolAttr isTrailingIdxScalable,
+                                 AsmParser::Delimiter delimiter) {
   char leftDelimiter = getLeftDelimiter(delimiter);
   char rightDelimiter = getRightDelimiter(delimiter);
   printer << leftDelimiter;
   if (integers.empty()) {
     printer << rightDelimiter;
     return;
   }
 
   int64_t trailingScalableInteger;
-  if (isTrailingIdxScalable) {
+  if (isTrailingIdxScalable && isTrailingIdxScalable.getValue()) {
     // ATM only the trailing idx can be scalable
     trailingScalableInteger = integers.back();
     integers = integers.drop_back();
   }
 
   unsigned idx = 0;
   llvm::interleaveComma(integers, printer, [&](int64_t integer) {
     if (ShapedType::isDynamic(integer)) {
       printer << values[idx];
       if (!valueTypes.empty())
         printer << " : " << valueTypes[idx];
       ++idx;
     } else {
       printer << integer;
     }
   });
 
   // Print the trailing scalable index
-  if (isTrailingIdxScalable) {
+  if (isTrailingIdxScalable && isTrailingIdxScalable.getValue()) {
+    if (!integers.empty())
       printer << ", ";
     printer << "[";
     printer << trailingScalableInteger;
     printer << "]";
   }
 
   printer << rightDelimiter;
 }
 
 ParseResult mlir::parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
     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.
+    // index was identified as scalable.
     if (foundScalable) {
       parser.emitError(parser.getNameLoc())
-          << "non-trailing tile size cannot be scalable";
+          << "non-trailing index cannot be scalable";
       return failure();
     }
 
     if (isTrailingIdxScalable && parser.parseOptionalLSquare().succeeded())
       foundScalable = true;
 
     if (res.has_value() && succeeded(res.value())) {
       values.push_back(operand);

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

%0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, 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 {{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/Linalg/vectorization-masked.mlir

-// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file --verify-diagnostics | FileCheck %s
 
 func.func @vectorize_dynamic_identity(%arg0: tensor<?xf32>,
                                       %arg1: tensor<?xf32>,
                                       %arg2: tensor<?xf32>) -> tensor<?xf32> {
   %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,
                                          affine_map<(d0) -> (d0)>,
                                          affine_map<(d0) -> (d0)>],
                    iterator_types = ["parallel"] }

 // CHECK:           return
 // CHECK:         }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.structured.masked_vectorize %0 vector_sizes [8, 16, 4] : !transform.any_op
 }
+
+// -----
+
+func.func @vectorize_dynamic_matmul_scalable(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
+  // expected-warning @+1 {{Scalable vectorization is not supported yet}}
+  linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
+            outs(%C: memref<?x?xf32>)
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.structured.masked_vectorize %0 vector_sizes [8, 16, [4]] : !transform.any_op
+}