[mlir][VectorType] Allow arbitrary dimensions to be scalable

[awarzynski]

Summary

At the moment, only the trailing dimensions in the vector type can be
scalable, i.e. this is supported:

vector<2x[4]xf32>

and this is not allowed:

vector<[2]x4xf32>

This patch extends the vector type so that arbitrary dimensions can be
scalable. To this end, an array of bool values is added to every vector
type to denote whether the corresponding dimensions are scalable or not.
For example, for this vector:

vector<[2]x[3]x4xf32>

the following array would be created:

{true, true, false}.

Additionally, the current syntax:

vector<[2x3]x4xf32>

is replaced with:

vector<[2]x[3]x4xf32>

This is primarily to simplify parsing (this way, the parser can easily
process one dimension at a time rather than e.g. tracking whether
"scalable block" has been entered/left).

NOTE: The isScalableDim parameter of VectorType (introduced in this patch) makes numScalableDims redundant. For the time being, numScalableDims is preserved to facilitate the transition between the two parameters. numScalableDims will be removed in one of the subsequent patches.

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/

[jsetoain]

mlir/lib/AsmParser/Parser.h

214–215

Is there a reason to keep this?

[jsetoain]

mlir/lib/AsmParser/Parser.h

214–215

Never mind. The blue colour in the "NOTE" tripped my mental parser 🙂

[dcaballe]

This is great! Thanks for taking care of this! A few comments

mlir/include/mlir/Bytecode/BytecodeImplementation.h

257

blob -> bool?

mlir/include/mlir/IR/BuiltinTypes.td

1085

Also replying to:

NOTE: The isScalableDim parameter of VectorType (introduced in this patch) makes numScalableDims redundant. For the time being, numScalableDims is preserved to facilitate the transition between the two parameters. numScalableDims will be removed in one of the subsequent patches.

Is this transition really that painful? The new representation is a superset of the previous one it should be a simple API change.
I'm worried about the complexity of having both representations in place and potential errors derived from that.

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

995–1001

I would make sure this is explicitly initialized to false

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

2799–2807

Not sure I follow the use case here but I think we shouldn't have any specific constraint at this level of abstraction regarding the number of scalable dimensions allowed. We should fail when we try to legalize this, if the scalable dimensions haven't been removed, unrolled (that's interesting, how do we unrolled a scalable vector dims?) or replaced in some way and there is no target representation for that (e.g., llvm intrinsics, SVE/SME ops, etc.)

3534

Ha, this is the part that required the bitmap representation (see comment in 3541).

I think we just have to apply the same permutation as to the shape so... we can do either...

invPermMap.compose(vecType.getScalableDims())  // not sure if we can pass a bool SmallVector.

or

applyPermutationMap(invPerMap, vecType.getScalableDims())

mlir/lib/IR/AsmPrinter.cpp

2463

commented code

[awarzynski]

Addressing Diego's comments, thanks for reviewing!

[awarzynski]

Thank you for reviewing!

mlir/include/mlir/IR/BuiltinTypes.td

1085

Is this transition really that painful?

Keeping both really helped with debugging and sanity-checking. Also, this is a rather intrusive change and I just like the idea of doing things like this in smaller, incremental steps.

Separately, I am worried that once this lands we'll see some "untested" code breaking (because, for example, I am making the verifiers a bit stricter). And smaller patches are easier to debug :)

The new representation is a superset of the previous one it should be a simple API change.

Yes, but the previous representation imposes some rather strong assumptions.

I'm worried about the complexity of having both representations in place and potential errors derived from that.

Fear not, I intend to remove it very soon :) I will try that now and see how complex that would be.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

2799–2807

Tl;Dr Please check the updated version :)

Not sure I follow the use case here

It's triggered by this test: vector-scalable-outerproduct.mlir. It leads to numScalableDims = 3.

I could just ignore that, but my changes to VectorType::verify mean that VectorType::get below (L2814) fails:

if (numScale != numScalableDims)
  return emitError() << "number of scalable dims must match, explicit: "
                     << numScalableDims << ", and bools:" << numScale;

This is why I'd like to keep numScalableDims for now - it helps capturing these cases. And I do have a "nice" fix for this specific example :)

I think we shouldn't have any specific constraint at this level of abstraction regarding the number of scalable dimensions allowed. We should fail when we try to legalize this, if the scalable dimensions haven't been removed, unrolled (that's interesting, how do we unrolled a scalable vector dims?) or replaced in some way and there is no target representation for that (e.g., llvm intrinsics, SVE/SME ops, etc.)

Agreed. In this particular case, we were setting numScalableDims to 3 for this result vector shape: vector<[4]x[4]xf32. That's incorrect and VectorType::verify should complain. At the moment, that would be ignored.

3534

Thanks! I'm inclined to preserve the assert from L3542 for now. WDYT?

[awarzynski]

mlir/include/mlir/IR/BuiltinTypes.td

1085

Uploaded here: https://reviews.llvm.org/D153412. Happy to merge with this patch if that's preferred. Both approaches (1 vs 2 patches) have their merits.

[awarzynski]

Update comments/documentation

[kuhar]

Just some nits

mlir/include/mlir/IR/BuiltinTypes.td

1078

its

1081–1082

nit: the { should be go to the previous line

1083

nit: you can insert into the existing vector with isScalabeVec.resize(shape.size(), false) or isScalableVec.append(shape.size(), false)

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

998–1001

Would return reder.parseByte(result) work here?

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

1002

nit: prefer C++ casts (static_cast)

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

60–62

nit: if you need an array ref, you can construct one from the single element (i.e., without any container)

[awarzynski]

Addressing comments from Jakub

Thank you @kuhar, your pointers are much appreciated - I always learn a lot from them :)

[dcaballe]

LGTM % pending comments. Thanks!

mlir/include/mlir/IR/BuiltinTypes.td

1085

Fair enough. We can fo with two patches then if both are landing soon

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

3534

SG but then we shouldn't need the permutation as the map would always be an identity?

[awarzynski]

Rename isScalableDim as scalableDims. This was suggested in https://reviews.llvm.org/D153412

Also updated "mlir/include/mlir/IR/BuiltinTypes.h" with some fairly straightforward changes that I forgot to include earlier.

[awarzynski]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

998–1001

It would, parseByte also returns LogicalResult. Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

3534

Good point. I took another look and decided to remove the assert. It should be safe, though it's tricky to trigger this code path just now. IIUC, first I need to update the transform dialect to allow something like this vector_sizes [2, [4], 8].

[kuhar]

LGTM

[c-rhodes]

mlir/include/mlir/IR/BuiltinTypes.h

314

nit: this variable name is a bit confusing at first I thought new meant this is temporary until numScalableDims is removed, but then realised it's copied from setShape? The latter is doing an update so it is new, but that's not the case during construction. How about scalableDims but use this to refer to the class member?

317–322

if/else like below?

326–330

nit: reverse condition

mlir/include/mlir/IR/BuiltinTypes.td

1028

not sure this is right, it means dim must always be between square brackets?

1063

1070

nit: this can be inferred from the type

mlir/lib/AsmParser/Parser.h

215

mlir/lib/AsmParser/TypeParser.cpp

474

As above, not sure this is correct.

485–487

490–492

nit: drop braces

[awarzynski]

Thank you both for reviewing. I believe that this addresses all the outstanding comments. Do let me know if I missed anything. Otherwise, I will merge this shortly.

mlir/include/mlir/IR/BuiltinTypes.h

314

Fair point, let me use this->scalableDims.

mlir/include/mlir/IR/BuiltinTypes.td

1028

Good catch! I will update in the next revision. Let me know whether it makes sense.

1070

Yeah, it was meant to be some clever comment :)

[awarzynski]

Address comments from Cullen

Thanks for reviewing @c-rhodes, lots of great suggestions!

[c-rhodes]

I've left a few more minor nits but otherwise LGTM, cheers.

mlir/include/mlir/IR/BuiltinTypes.h

317–321

nit: drop braces

326–330

nit: drop braces

mlir/include/mlir/IR/BuiltinTypes.td

1080–1084

can this not be written as?

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

998–1000

nit: could be written as

[awarzynski]

Thanks Cullen, I will include your suggestions in the final version that I will merge shortly.

mlir/include/mlir/IR/BuiltinTypes.td

1080–1084

IIUC, you'd be creating an ArrayRef from a temporary, which is disabled.

mlir/include/mlir/Bytecode/BytecodeImplementation.h

public:
   readAPFloatWithKnownSemantics(const llvm::fltSemantics &semantics) = 0;
 
   /// Read a string from the bytecode.
   virtual LogicalResult readString(StringRef &result) = 0;
 
   /// Read a blob from the bytecode.
   virtual LogicalResult readBlob(ArrayRef<char> &result) = 0;
 
+  /// Read a bool from the bytecode.
+  virtual LogicalResult readBool(bool &result) = 0;
+
 private:
   /// Read a handle to a dialect resource.
   virtual FailureOr<AsmDialectResourceHandle> readResourceHandle() = 0;
 };
 
 //===----------------------------------------------------------------------===//
 // DialectBytecodeWriter
 //===----------------------------------------------------------------------===//
public:
   /// owned by an attribute or type.
   virtual void writeOwnedString(StringRef str) = 0;
 
   /// Write a blob to the bytecode, which is owned by the caller and is
   /// guaranteed to not die before the end of the bytecode process. The blob is
   /// written as-is, with no additional compression or compaction.
   virtual void writeOwnedBlob(ArrayRef<char> blob) = 0;
 
+  /// Write a bool to the output stream.
+  virtual void writeOwnedBool(bool value) = 0;
+
   /// Return the bytecode version being emitted for.
   virtual int64_t getBytecodeVersion() const = 0;
 };
 
 //===--------------------------------------------------------------------===//
 // Dialect Version Interface.
 //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/BuiltinDialectBytecode.td

 def VectorType : DialectType<(type
   Array<SignedVarIntList>:$shape,
   Type:$elementType
 )> {
   let printerPredicate = "!$_val.getNumScalableDims()";
 }
 
 def VectorTypeWithScalableDims : DialectType<(type
+  Array<BoolList>:$scalableDims,
   VarInt:$numScalableDims,
   Array<SignedVarIntList>:$shape,
   Type:$elementType
 )> {
   let printerPredicate = "$_val.getNumScalableDims()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, shape, elementType, numScalableDims)";
+  let cBuilder = "get<$_resultType>(context, shape, elementType, numScalableDims, scalableDims)";
 }
 }
 
 /// This enum contains marker codes used to indicate which attribute is
 /// currently being decoded, and how it should be decoded. The order of these
 /// codes should generally be unchanged, as any changes will inevitably break
 /// compatibility with older bytecode.
 

mlir/include/mlir/IR/BuiltinTypes.h

 
 /// This is a builder type that keeps local references to arguments. Arguments
 /// that are passed into the builder must outlive the builder.
 class VectorType::Builder {
 public:
   /// Build from another VectorType.
   explicit Builder(VectorType other)
       : shape(other.getShape()), elementType(other.getElementType()),
-        numScalableDims(other.getNumScalableDims()) {}
+        numScalableDims(other.getNumScalableDims()),
+        scalableDims(other.getScalableDims()) {}
 
   /// Build from scratch.
   Builder(ArrayRef<int64_t> shape, Type elementType,
-          unsigned numScalableDims = 0)
+          unsigned numScalableDims = 0, ArrayRef<bool> scalableDims = {})
       : shape(shape), elementType(elementType),
-        numScalableDims(numScalableDims) {}
+        numScalableDims(numScalableDims) {
+    if (scalableDims.empty())
+      scalableDims = SmallVector<bool>(shape.size(), false);
+    else
+      this->scalableDims = scalableDims;
+  }
 
-  Builder &setShape(ArrayRef<int64_t> newShape,
-                    unsigned newNumScalableDims = 0) {
+  Builder &setShape(ArrayRef<int64_t> newShape, unsigned newNumScalableDims = 0,
+                    ArrayRef<bool> newIsScalableDim = {}) {
     numScalableDims = newNumScalableDims;
+    if (newIsScalableDim.empty())
+      scalableDims = SmallVector<bool>(shape.size(), false);
+    else
+      scalableDims = newIsScalableDim;
+
     shape = newShape;
     return *this;
   }
 
   Builder &setElementType(Type newElementType) {
     elementType = newElementType;
     return *this;
   }
 
   /// Erase a dim from shape @pos.
   Builder &dropDim(unsigned pos) {
     assert(pos < shape.size() && "overflow");
     if (pos >= shape.size() - numScalableDims)
       numScalableDims--;
     if (storage.empty())
       storage.append(shape.begin(), shape.end());
+    if (storageScalableDims.empty())
+      storageScalableDims.append(scalableDims.begin(), scalableDims.end());
     storage.erase(storage.begin() + pos);
+    storageScalableDims.erase(storageScalableDims.begin() + pos);
     shape = {storage.data(), storage.size()};
+    scalableDims =
+        ArrayRef<bool>(storageScalableDims.data(), storageScalableDims.size());
     return *this;
   }
 
   /// In the particular case where the vector has a single dimension that we
   /// drop, return the scalar element type.
   // TODO: unify once we have a VectorType that supports 0-D.
   operator Type() {
     if (shape.empty())
       return elementType;
-    return VectorType::get(shape, elementType, numScalableDims);
+    return VectorType::get(shape, elementType, numScalableDims, scalableDims);
   }
 
 private:
   ArrayRef<int64_t> shape;
   // Owning shape data for copy-on-write operations.
   SmallVector<int64_t> storage;
   Type elementType;
   unsigned numScalableDims;
+  ArrayRef<bool> scalableDims;
+  // Owning scalableDims data for copy-on-write operations.
+  SmallVector<bool> storageScalableDims;
 };
 
 /// Given an `originalShape` and a `reducedShape` assumed to be a subset of
 /// `originalShape` with some `1` entries erased, return the set of indices
 /// that specifies which of the entries of `originalShape` are dropped to obtain
 /// `reducedShape`. The returned mask can be applied as a projection to
 /// `originalShape` to obtain the `reducedShape`. This mask is useful to track
 /// which dimensions must be kept when e.g. compute MemRef strides under

mlir/include/mlir/IR/BuiltinTypes.td

 def Builtin_Vector : Builtin_Type<"Vector", [ShapedTypeInterface], "Type"> {
   let summary = "Multi-dimensional SIMD vector type";
   let description = [{
     Syntax:
 
     ```
     vector-type ::= `vector` `<` vector-dim-list vector-element-type `>`
     vector-element-type ::= float-type | integer-type | index-type
-    vector-dim-list := (static-dim-list `x`)? (`[` static-dim-list `]` `x`)?
-    static-dim-list ::= decimal-literal (`x` decimal-literal)*
+    vector-dim-list := (static-dim-list `x`)?
+    static-dim-list ::= static-dim (`x` static-dim)*
+    static-dim ::= (decimal-literal | `[` decimal-literal `]`)
     ```
 
     The vector type represents a SIMD style vector used by target-specific
     operation sets like AVX or SVE. While the most common use is for 1D
     vectors (e.g. vector<16 x f32>) we also support multidimensional registers
     on targets that support them (like TPUs). The dimensions of a vector type
     can be fixed-length, scalable, or a combination of the two. The 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.
+    dimensions in a vector are indicated between square brackets ([ ]).
 
     Vector shapes must be positive decimal integers. 0D vectors are allowed by
     omitting the dimension: `vector<f32>`.
 
     Note: hexadecimal integer literals are not allowed in vector type
     declarations, `vector<0x42xi32>` is invalid because it is interpreted as a
     2D vector with shape `(0, 42)` and zero shapes are not allowed.
 
     Examples:
 
     ```mlir
     // A 2D fixed-length vector of 3x42 i32 elements.
     vector<3x42xi32>
 
     // A 1D scalable-length vector that contains a multiple of 4 f32 elements.
     vector<[4]xf32>
 
     // A 2D scalable-length vector that contains a multiple of 2x8 f32 elements.
-    vector<[2x8]xf32>
+    vector<[2]x[8]xf32>
 
     // A 2D mixed fixed/scalable vector that contains 4 scalable vectors of 4 f32 elements.
     vector<4x[4]xf32>
+
+    // A 3D mixed fixed/scalable vector in which only the inner dimension is
+    // scalable.
+    vector<2x[4]x8xf32>
     ```
   }];
   let parameters = (ins
     ArrayRefParameter<"int64_t">:$shape,
     "Type":$elementType,
-    "unsigned":$numScalableDims
+    "unsigned":$numScalableDims,
+    ArrayRefParameter<"bool">:$scalableDims
   );
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape, "Type":$elementType,
-      CArg<"unsigned", "0">:$numScalableDims
+      CArg<"unsigned", "0">:$numScalableDims,
+      CArg<"ArrayRef<bool>", "{}">:$scalableDims
     ), [{
+      // While `scalableDims` is optional, its default value should be
+      // `false` for every dim in `shape`.
+      SmallVector<bool> isScalableVec;
+      if (scalableDims.empty()) {
+        isScalableVec.resize(shape.size(), false);
+        scalableDims = isScalableVec;
+      }
       return $_get(elementType.getContext(), shape, elementType,
-                   numScalableDims);
+                   numScalableDims, scalableDims);
     }]>
   ];
   let extraClassDeclaration = [{
     /// This is a builder type that keeps local references to arguments.
     /// Arguments that are passed into the builder must outlive the builder.
     class Builder;
 
     /// Returns true if the given type can be used as an element of a vector

mlir/include/mlir/IR/BytecodeBase.td

def SignedVarInt :
   WithBuilder<"$_args",
   WithPrinter<"$_writer.writeSignedVarInt($_getter)",
   WithType   <"int64_t">>>>;
 def Blob :
   WithParser <"succeeded($_reader.readBlob($_var))",
   WithBuilder<"$_args",
   WithPrinter<"$_writer.writeOwnedBlob($_getter)",
   WithType   <"ArrayRef<char>">>>>;
+def Bool :
+  WithParser <"succeeded($_reader.readBool($_var))",
+  WithBuilder<"$_args",
+  WithPrinter<"$_writer.writeOwnedBool($_getter)",
+  WithType   <"bool">>>>;
 class KnownWidthAPInt<string s> :
   WithParser <"succeeded(readAPIntWithKnownWidth($_reader, " # s # ", $_var))",
   WithBuilder<"$_args",
   WithPrinter<"$_writer.writeAPIntWithKnownWidth($_getter)",
   WithType   <"FailureOr<APInt>">>>>;
 class KnownSemanticsAPFloat<string s> :
   WithParser <"succeeded(readAPFloatWithKnownSemantics($_reader, " # s # ", $_var))",
   WithBuilder<"$_args",
class Array<Bytecode t> {
   Bytecode elemT = t;
 
   string cBuilder = "$_args";
 }
 // - Array elements currently needs a different bytecode type to accommodate
 //   for the list print/parsing.
 class List<Bytecode t> : WithGetter<"$_member", t>;
 def SignedVarIntList : List<SignedVarInt>;
+def BoolList : List<Bool>;
 
 // Define dialect attribute or type.
 class DialectAttrOrType<dag d> {
   // Any members starting with underscore is not fed to create function but
   // treated as purely local variable.
   dag members = d;
 
   // When needing to specify a custom return type.

mlir/lib/AsmParser/Parser.h

public:
   Type parseTensorType();
 
   /// Parse a tuple type.
   Type parseTupleType();
 
   /// Parse a vector type.
   VectorType parseVectorType();
   ParseResult parseVectorDimensionList(SmallVectorImpl<int64_t> &dimensions,
-                                       unsigned &numScalableDims);
+                                       unsigned &numScalableDims,
+                                       SmallVectorImpl<bool> &scalableDims);
   ParseResult parseDimensionListRanked(SmallVectorImpl<int64_t> &dimensions,
                                        bool allowDynamic = true,
                                        bool withTrailingX = true);
   ParseResult parseIntegerInDimensionList(int64_t &value);
   ParseResult parseXInDimensionList();
 
   //===--------------------------------------------------------------------===//
   // Attribute Parsing

mlir/lib/AsmParser/TypeParser.cpp

 ///
 VectorType Parser::parseVectorType() {
   consumeToken(Token::kw_vector);
 
   if (parseToken(Token::less, "expected '<' in vector type"))
     return nullptr;
 
   SmallVector<int64_t, 4> dimensions;
+  SmallVector<bool, 4> scalableDims;
   unsigned numScalableDims;
-  if (parseVectorDimensionList(dimensions, numScalableDims))
+  if (parseVectorDimensionList(dimensions, numScalableDims, scalableDims))
     return nullptr;
   if (any_of(dimensions, [](int64_t i) { return i <= 0; }))
     return emitError(getToken().getLoc(),
                      "vector types must have positive constant sizes"),
            nullptr;
 
   // Parse the element type.
   auto typeLoc = getToken().getLoc();
   auto elementType = parseType();
   if (!elementType || parseToken(Token::greater, "expected '>' in vector type"))
     return nullptr;
 
   if (!VectorType::isValidElementType(elementType))
     return emitError(typeLoc, "vector elements must be int/index/float type"),
            nullptr;
 
-  return VectorType::get(dimensions, elementType, numScalableDims);
+  return VectorType::get(dimensions, elementType, numScalableDims,
+                         scalableDims);
 }
 
-/// Parse a dimension list in a vector type. This populates the dimension list,
-/// and returns the number of scalable dimensions in `numScalableDims`.
-///
-/// vector-dim-list := (static-dim-list `x`)? (`[` static-dim-list `]` `x`)?
-/// static-dim-list ::= decimal-literal (`x` decimal-literal)*
+/// Parse a dimension list in a vector type. This populates the dimension list.
+/// For i-th dimension, `scalableDims[i]` contains either:
+///   * `false` for a non-scalable dimension (e.g. `4`),
+///   * `true` for a scalable dimension (e.g. `[4]`).
+/// This method also returns the number of scalable dimensions in
+/// `numScalableDims`.
+///
+/// vector-dim-list := (static-dim-list `x`)?
+/// static-dim-list ::= static-dim (`x` static-dim)*
+/// static-dim ::= (decimal-literal | `[` decimal-literal `]`)
 ///
 ParseResult
 Parser::parseVectorDimensionList(SmallVectorImpl<int64_t> &dimensions,
-                                 unsigned &numScalableDims) {
+                                 unsigned &numScalableDims,
+                                 SmallVectorImpl<bool> &scalableDims) {
   numScalableDims = 0;
   // If there is a set of fixed-length dimensions, consume it
-  while (getToken().is(Token::integer)) {
-    int64_t value;
-    if (parseIntegerInDimensionList(value))
-      return failure();
-    dimensions.push_back(value);
-    // Make sure we have an 'x' or something like 'xbf32'.
-    if (parseXInDimensionList())
-      return failure();
-  }
-  // If there is a set of scalable dimensions, consume it
-  if (consumeIf(Token::l_square)) {
-    while (getToken().is(Token::integer)) {
-      int64_t value;
-      if (parseIntegerInDimensionList(value))
-        return failure();
-      dimensions.push_back(value);
+  while (getToken().is(Token::integer) || getToken().is(Token::l_square)) {
+    bool scalable = consumeIf(Token::l_square);
+    if (scalable) {
+      if (!consumeIf(Token::r_square))
+        return emitWrongTokenError("missing ']' closing scalable dimension");
       numScalableDims++;
-      // Check if we have reached the end of the scalable dimension list
-      if (consumeIf(Token::r_square)) {
-        // Make sure we have something like 'xbf32'.
-        return parseXInDimensionList();
-      }
-      // Make sure we have an 'x'
-      if (parseXInDimensionList())
-        return failure();
-    }
-    // If we make it here, we've finished parsing the dimension list
-    // without finding ']' closing the set of scalable dimensions
-    return emitWrongTokenError(
-        "missing ']' closing set of scalable dimensions");
-  }
+    scalableDims.push_back(scalable);
+    // Make sure we have an 'x' or something like 'xbf32'.
 
   return success();
 }
 
 /// Parse a dimension list of a tensor or memref type.  This populates the
 /// dimension list, using ShapedType::kDynamic for the `?` dimensions if
 /// `allowDynamic` is set and errors out on `?` otherwise. Parsing the trailing
 /// `x` is configurable.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

LogicalResult readBlob(ArrayRef<char> &result) override {
     if (failed(reader.parseVarInt(dataSize)) ||
         failed(reader.parseBytes(dataSize, data)))
       return failure();
     result = llvm::ArrayRef(reinterpret_cast<const char *>(data.data()),
                             data.size());
     return success();
   }
 
+  LogicalResult readBool(bool &result) override {
+    return reader.parseByte(result);
+  }
+
 private:
   AttrTypeReader &attrTypeReader;
   StringSectionReader &stringReader;
   ResourceSectionReader &resourceReader;
   EncodingReader &reader;
 };
 
 /// Wraps the properties section and handles reading properties out of it.
 class PropertiesSectionReader {

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

public:
   }
 
   void writeOwnedBlob(ArrayRef<char> blob) override {
     emitter.emitVarInt(blob.size());
     emitter.emitOwnedBlob(ArrayRef<uint8_t>(
         reinterpret_cast<const uint8_t *>(blob.data()), blob.size()));
   }
 
+  void writeOwnedBool(bool value) override { emitter.emitByte(value); }
+
   int64_t getBytecodeVersion() const override { return bytecodeVersion; }
 
 private:
   int64_t bytecodeVersion;
   EncodingEmitter &emitter;
   IRNumberingState &numberingState;
   StringSectionBuilder &stringSection;
 };

mlir/lib/Bytecode/Writer/IRNumbering.cpp

struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
   void writeAPIntWithKnownWidth(const APInt &value) override {}
   void writeAPFloatWithKnownSemantics(const APFloat &value) override {}
   void writeOwnedString(StringRef) override {
     // TODO: It might be nice to prenumber strings and sort by the number of
     // references. This could potentially be useful for optimizing things like
     // file locations.
   }
   void writeOwnedBlob(ArrayRef<char> blob) override {}
+  void writeOwnedBool(bool value) override {}
 
   int64_t getBytecodeVersion() const override {
     llvm_unreachable("unexpected querying of version in IRNumbering");
   }
 
   /// The parent numbering state that is populated by this writer.
   IRNumberingState &state;
 };

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

 ///  * n>1 `vector<ax...xkxT>` convert via an (n-1)-D array type to
 ///    `!llvm.array<ax...array<jxvector<kxT>>>`.
 Type LLVMTypeConverter::convertVectorType(VectorType type) {
   auto elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
   if (type.getShape().empty())
     return VectorType::get({1}, elementType);
-  Type vectorType = VectorType::get(type.getShape().back(), elementType,
-                                    type.getNumScalableDims());
+  Type vectorType =
+      VectorType::get(type.getShape().back(), elementType,
+                      type.getNumScalableDims(), type.getScalableDims().back());
   assert(LLVM::isCompatibleVectorType(vectorType) &&
          "expected vector type compatible with the LLVM dialect");
   auto shape = type.getShape();
   for (int i = shape.size() - 2; i >= 0; --i)
     vectorType = LLVM::LLVMArrayType::get(vectorType, shape[i]);
   return vectorType;
 }
 

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 static Type getI1SameShape(Type type) {
   auto i1Type = IntegerType::get(type.getContext(), 1);
   if (auto tensorType = llvm::dyn_cast<RankedTensorType>(type))
     return RankedTensorType::get(tensorType.getShape(), i1Type);
   if (llvm::isa<UnrankedTensorType>(type))
     return UnrankedTensorType::get(i1Type);
   if (auto vectorType = llvm::dyn_cast<VectorType>(type))
     return VectorType::get(vectorType.getShape(), i1Type,
-                           vectorType.getNumScalableDims());
+                           vectorType.getNumScalableDims(),
+                           vectorType.getScalableDims());
   return i1Type;
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 void arith::ConstantOp::getAsmResultNames(

mlir/lib/Dialect/ArmSVE/IR/ArmSVEDialect.cpp

 // ScalableVector versions of general helpers for comparison ops
 //===----------------------------------------------------------------------===//
 
 /// Return the scalable vector of the same shape and containing i1.
 static Type getI1SameShape(Type type) {
   auto i1Type = IntegerType::get(type.getContext(), 1);
   if (auto sVectorType = llvm::dyn_cast<VectorType>(type))
     return VectorType::get(sVectorType.getShape(), i1Type,
-                           sVectorType.getNumScalableDims());
+                           sVectorType.getNumScalableDims(),
+                           sVectorType.getScalableDims());
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // Tablegen Definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.cpp.inc"

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

Type mlir::LLVM::getVectorType(Type elementType, unsigned numElements,
   (void)useBuiltIn;
   assert((useLLVM ^ useBuiltIn) && "expected LLVM-compatible fixed-vector type "
                                    "to be either builtin or LLVM dialect type");
   if (useLLVM) {
     if (isScalable)
       return LLVMScalableVectorType::get(elementType, numElements);
     return LLVMFixedVectorType::get(elementType, numElements);
   }
-  return VectorType::get(numElements, elementType, (unsigned)isScalable);
+
+  // LLVM vectors are always 1-D, hence only 1 bool is required to mark it as
+  // scalable/non-scalable.
+  SmallVector<bool> scalableDims(1, isScalable);
+
+  return VectorType::get(numElements, elementType,
+                         static_cast<unsigned>(isScalable), scalableDims);
 }
 
 Type mlir::LLVM::getVectorType(Type elementType,
                                const llvm::ElementCount &numElements) {
   if (numElements.isScalable())
     return getVectorType(elementType, numElements.getKnownMinValue(),
                          /*isScalable=*/true);
   return getVectorType(elementType, numElements.getFixedValue(),
                        /*isScalable=*/false);

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

VectorType getCanonicalVecType(
     // Make sure we don't end up with unsupported scalable vector dimensions
     // after the permutation. If so, we should bail out on that operation in the
     // scalable preconditions.
     assert(areValidScalableVecDims(scalableDims) &&
            "Permuted scalable vector dimensions are not supported");
 
     // TODO: Extend scalable vector type to support a bit map.
     bool numScalableDims = !scalableVecDims.empty() && scalableVecDims.back();
-    return VectorType::get(vectorShape, elementType, numScalableDims);
+    return VectorType::get(vectorShape, elementType, numScalableDims,
+                           scalableDims);
   }
 
   /// Masks an operation with the canonical vector mask if the operation needs
   /// masking. Returns the masked operation or the original operation if masking
   /// is not needed. If provided, the canonical mask for this operation is
   /// permuted using `maybeMaskingMap`.
   Operation *
   maskOperation(RewriterBase &rewriter, Operation *opToMask, LinalgOp linalgOp,
vectorizeOneOp(RewriterBase &rewriter, VectorizationState &state,
   SmallVector<Value> vecOperands;
   for (Value scalarOperand : op->getOperands()) {
     Value vecOperand = bvm.lookup(scalarOperand);
     assert(vecOperand && "Vector operand couldn't be found");
 
     if (firstMaxRankedType) {
       auto vecType = VectorType::get(firstMaxRankedType.getShape(),
                                      getElementTypeOrSelf(vecOperand.getType()),
-                                     firstMaxRankedType.getNumScalableDims());
+                                     firstMaxRankedType.getNumScalableDims(),
+                                     firstMaxRankedType.getScalableDims());
       vecOperands.push_back(broadcastIfNeeded(rewriter, vecOperand, vecType));
     } else {
       vecOperands.push_back(vecOperand);
     }
   }
   //   c. for elementwise, the result is the vector with the firstMaxRankedShape
   SmallVector<Type> resultTypes;
   for (Type resultType : op->getResultTypes()) {
     resultTypes.push_back(
         firstMaxRankedType
             ? VectorType::get(firstMaxRankedType.getShape(), resultType,
-                              firstMaxRankedType.getNumScalableDims())
+                              firstMaxRankedType.getNumScalableDims(),
+                              firstMaxRankedType.getScalableDims())
             : resultType);
   }
   //   d. Build and return the new op.
   return VectorizationResult{
       VectorizationStatus::NewOp,
       rewriter.create(op->getLoc(), op->getName().getIdentifier(), vecOperands,
                       resultTypes, op->getAttrs())};
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

 /// Helper test for invariant argument (defined outside given block).
 static bool isInvariantArg(BlockArgument arg, Block *block) {
   return arg.getOwner() != block;
 }
 
 /// Constructs vector type for element type.
 static VectorType vectorType(VL vl, Type etp) {
   unsigned numScalableDims = vl.enableVLAVectorization;
-  return VectorType::get(vl.vectorLength, etp, numScalableDims);
+  return VectorType::get(vl.vectorLength, etp, numScalableDims,
+                         vl.enableVLAVectorization);
 }
 
 /// Constructs vector type from a memref value.
 static VectorType vectorType(VL vl, Value mem) {
   return vectorType(vl, getMemRefType(mem).getElementType());
 }
 
 /// Constructs vector iteration mask.
 static Value genVectorMask(PatternRewriter &rewriter, Location loc, VL vl,

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

return emitOpError() << "destination type " << getType()
                          << getSourceVectorType();
 
   return success();
 }
 
 /// Returns the mask type expected by this operation.
 Type MultiDimReductionOp::getExpectedMaskType() {
   auto vecType = getSourceVectorType();
-  return VectorType::get(vecType.getShape(),
-                         IntegerType::get(vecType.getContext(), /*width=*/1),
-                         vecType.getNumScalableDims());
+  return VectorType::get(
+      vecType.getShape(), IntegerType::get(vecType.getContext(), /*width=*/1),
+      vecType.getNumScalableDims(), vecType.getScalableDims());
 }
 
 namespace {
 // Only unit dimensions that are being reduced are folded. If the dimension is
 // unit, but not reduced, it is not folded, thereby keeping the output type the
 // same. If not all dimensions which are reduced are of unit dimension, this
 // transformation does nothing. This is just a generalization of
 // ElideSingleElementReduction for ReduceOp.
void ReductionOp::print(OpAsmPrinter &p) {
   p << " : " << getVector().getType() << " into " << getDest().getType();
 }
 
 // MaskableOpInterface methods.
 
 /// Returns the mask type expected by this operation.
 Type ReductionOp::getExpectedMaskType() {
   auto vecType = getSourceVectorType();
-  return VectorType::get(vecType.getShape(),
-                         IntegerType::get(vecType.getContext(), /*width=*/1),
-                         vecType.getNumScalableDims());
+  return VectorType::get(
+      vecType.getShape(), IntegerType::get(vecType.getContext(), /*width=*/1),
+      vecType.getNumScalableDims(), vecType.getScalableDims());
 }
 
 Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
                                          OpBuilder &builder, Location loc,
                                          Value vector) {
   switch (op) {
   case arith::AtomicRMWKind::addf:
   case arith::AtomicRMWKind::addi:
if (operandsInfo.size() < 2)
     return parser.emitError(parser.getNameLoc(),
                             "expected at least 2 operands");
   VectorType vLHS = llvm::dyn_cast<VectorType>(tLHS);
   VectorType vRHS = llvm::dyn_cast<VectorType>(tRHS);
   if (!vLHS)
     return parser.emitError(parser.getNameLoc(),
                             "expected vector type for operand #1");
 
-  unsigned numScalableDims = vLHS.getNumScalableDims();
   VectorType resType;
   if (vRHS) {
-    numScalableDims += vRHS.getNumScalableDims();
+    SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0],
+                                      vRHS.getScalableDims()[0]};
+    auto numScalableDims =
+        count_if(scalableDimsRes, [](bool isScalable) { return isScalable; });
     resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
-                              vLHS.getElementType(), numScalableDims);
+                              vLHS.getElementType(), numScalableDims,
+                              scalableDimsRes);
   } else {
     // Scalar RHS operand
+    SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0]};
+    auto numScalableDims =
+        count_if(scalableDimsRes, [](bool isScalable) { return isScalable; });
     resType = VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType(),
-                              numScalableDims);
+                              numScalableDims, scalableDimsRes);
   }
 
   if (!result.attributes.get(OuterProductOp::getKindAttrStrName())) {
     result.attributes.append(
         OuterProductOp::getKindAttrStrName(),
         CombiningKindAttr::get(result.getContext(),
                                OuterProductOp::getDefaultKind()));
   }
 

 }
 
 // MaskableOpInterface methods.
 
 /// Returns the mask type expected by this operation. Mostly used for
 /// verification purposes. It requires the operation to be vectorized."
 Type OuterProductOp::getExpectedMaskType() {
   auto vecType = this->getResultVectorType();
-  return VectorType::get(vecType.getShape(),
-                         IntegerType::get(vecType.getContext(), /*width=*/1),
-                         vecType.getNumScalableDims());
+  return VectorType::get(
+      vecType.getShape(), IntegerType::get(vecType.getContext(), /*width=*/1),
+      vecType.getNumScalableDims(), vecType.getScalableDims());
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ReshapeOp::verify() {
   // Verify that rank(numInputs/outputs) + numFixedVec dim matches vec rank.

 /// Infers the mask type for a transfer op given its vector type and
 /// permutation map. The mask in a transfer op operation applies to the
 /// tensor/buffer part of it and its type should match the vector shape
 /// *before* any permutation or broadcasting.
 static VectorType inferTransferOpMaskType(VectorType vecType,
                                           AffineMap permMap) {
   auto i1Type = IntegerType::get(permMap.getContext(), 1);
   AffineMap invPermMap = inversePermutation(compressUnusedDims(permMap));
-  // TODO: Extend the scalable vector type representation with a bit map.
-  assert((permMap.isMinorIdentity() || vecType.getNumScalableDims() == 0) &&
-         "Scalable vectors are not supported yet");
   assert(invPermMap && "Inversed permutation map couldn't be computed");
   SmallVector<int64_t, 8> maskShape = invPermMap.compose(vecType.getShape());
-  return VectorType::get(maskShape, i1Type, vecType.getNumScalableDims());
+
+  SmallVector<bool> scalableDims =
+      applyPermutationMap(invPermMap, vecType.getScalableDims());
+
+  return VectorType::get(maskShape, i1Type, vecType.getNumScalableDims(),
+                         scalableDims);
 }
 
 ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
   auto &builder = parser.getBuilder();
   SMLoc typesLoc;
   OpAsmParser::UnresolvedOperand sourceInfo;
   SmallVector<OpAsmParser::UnresolvedOperand, 8> indexInfo;
   OpAsmParser::UnresolvedOperand paddingInfo;
   SmallVector<Type, 2> types;
   OpAsmParser::UnresolvedOperand maskInfo;

 }
 
 // MaskableOpInterface methods.
 
 /// Returns the mask type expected by this operation. Mostly used for
 /// verification purposes. It requires the operation to be vectorized."
 Type GatherOp::getExpectedMaskType() {
   auto vecType = this->getIndexVectorType();
-  return VectorType::get(vecType.getShape(),
-                         IntegerType::get(vecType.getContext(), /*width=*/1),
-                         vecType.getNumScalableDims());
+  return VectorType::get(
+      vecType.getShape(), IntegerType::get(vecType.getContext(), /*width=*/1),
+      vecType.getNumScalableDims(), vecType.getScalableDims());
 }
 
 std::optional<SmallVector<int64_t, 4>> GatherOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getVectorType().getShape());
 }
 
 namespace {
 class GatherFolder final : public OpRewritePattern<GatherOp> {

mlir/lib/IR/AsmPrinter.cpp

TypeSwitch<Type>(type)
           printType(results[0]);
         } else {
           os << '(';
           interleaveComma(results, [&](Type ty) { printType(ty); });
           os << ')';
         }
       })
       .Case<VectorType>([&](VectorType vectorTy) {
+        auto scalableDims = vectorTy.getScalableDims();
         os << "vector<";
         auto vShape = vectorTy.getShape();
         unsigned lastDim = vShape.size();
-        unsigned lastFixedDim = lastDim - vectorTy.getNumScalableDims();
         unsigned dimIdx = 0;
-        for (dimIdx = 0; dimIdx < lastFixedDim; dimIdx++)
-          os << vShape[dimIdx] << 'x';
-        if (vectorTy.isScalable()) {
+        for (dimIdx = 0; dimIdx < lastDim; dimIdx++) {
+          if (!scalableDims.empty() && scalableDims[dimIdx])
             os << '[';
-          unsigned secondToLastDim = lastDim - 1;
-          for (; dimIdx < secondToLastDim; dimIdx++)
-            os << vShape[dimIdx] << 'x';
-          os << vShape[dimIdx] << "]x";
+          os << vShape[dimIdx];
+          if (!scalableDims.empty() && scalableDims[dimIdx])
+            os << ']';
+          os << 'x';
         }
         printType(vectorTy.getElementType());
         os << '>';
       })
       .Case<RankedTensorType>([&](RankedTensorType tensorTy) {
         os << "tensor<";
         for (int64_t dim : tensorTy.getShape()) {
           if (ShapedType::isDynamic(dim))

mlir/lib/IR/BuiltinTypes.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult VectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  ArrayRef<int64_t> shape, Type elementType,
-                                 unsigned numScalableDims) {
+                                 unsigned numScalableDims,
+                                 ArrayRef<bool> scalableDims) {
   if (!isValidElementType(elementType))
     return emitError()
            << "vector elements must be int/index/float type but got "
            << elementType;
 
   if (any_of(shape, [](int64_t i) { return i <= 0; }))
     return emitError()
            << "vector types must have positive constant sizes but got "
            << shape;
 
+  if (numScalableDims > shape.size())
+    return emitError()
+           << "number of scalable dims cannot exceed the number of dims"
+           << " (" << numScalableDims << " vs " << shape.size() << ")";
+
+  if (scalableDims.size() != shape.size())
+    return emitError() << "number of dims must match, got "
+                       << scalableDims.size() << " and " << shape.size();
+
+  auto numScale =
+      count_if(scalableDims, [](bool isScalable) { return isScalable; });
+  if (numScale != numScalableDims)
+    return emitError() << "number of scalable dims must match, explicit: "
+                       << numScalableDims << ", and bools:" << numScale;
+
   return success();
 }
 
 VectorType VectorType::scaleElementBitwidth(unsigned scale) {
   if (!scale)
     return VectorType();
   if (auto et = llvm::dyn_cast<IntegerType>(getElementType()))
     if (auto scaledEt = et.scaleElementBitwidth(scale))

mlir/test/Dialect/Builtin/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 //===----------------------------------------------------------------------===//
 // UnrealizedConversionCastOp
 //===----------------------------------------------------------------------===//
 
 // expected-error@+1 {{expected at least one result for cast operation}}
 "builtin.unrealized_conversion_cast"() : () -> ()
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
-// expected-error@+1 {{missing ']' closing set of scalable dimensions}}
+// expected-error@+1 {{missing ']' closing scalable dimension}}
 func.func @scalable_vector_arg(%arg0: vector<[4xf32>) { }
 
 // -----
+
+// expected-error@+1 {{missing ']' closing scalable dimension}}
+func.func @scalable_vector_arg(%arg0: vector<[4x4]xf32>) { }

mlir/test/Dialect/Builtin/ops.mlir

 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
 // A basic 1D scalable vector
 %scalable_vector_1d = "foo.op"() : () -> vector<[4]xi32>
 
 // A 2D scalable vector
-%scalable_vector_2d = "foo.op"() : () -> vector<[2x2]xf64>
+%scalable_vector_2d = "foo.op"() : () -> vector<[2]x[2]xf64>
 
 // A 2D scalable vector with fixed-length dimensions
 %scalable_vector_2d_mixed = "foo.op"() : () -> vector<2x[4]xbf16>
 
 // A multi-dimensional vector with mixed scalable and fixed-length dimensions
-%scalable_vector_multi_mixed = "foo.op"() : () -> vector<2x2x[4x4]xi8>
+%scalable_vector_multi_mixed = "foo.op"() : () -> vector<2x2x[4]x[4]xi8>

mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s | mlir-opt
 
 func.func @scalable_outerproduct(%src : memref<?xf32>) {
   %idx = arith.constant 0 : index
   %cst = arith.constant 1.0 : f32
   %0 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
   %1 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
 
   %op = vector.outerproduct %0, %1 : vector<[4]xf32>, vector<[4]xf32>
-  vector.store %op, %src[%idx] : memref<?xf32>, vector<[4x4]xf32>
+  vector.store %op, %src[%idx] : memref<?xf32>, vector<[4]x[4]xf32>
 
   %op2 = vector.outerproduct %0, %cst : vector<[4]xf32>, f32
   vector.store %op2, %src[%idx] : memref<?xf32>, vector<[4]xf32>
   return
 }
 
 // -----
 
 func.func @invalid_outerproduct(%src : memref<?xf32>) {
   %idx = arith.constant 0 : index
   %0 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
   %1 = vector.load %src[%idx] : memref<?xf32>, vector<4xf32>
 
   // expected-error @+1 {{expected either all or none of vector operands #1 and #2 to be scalable}}
   %op = vector.outerproduct %0, %1 : vector<[4]xf32>, vector<4xf32>
 }
 // -----
 
 func.func @invalid_outerproduct1(%src : memref<?xf32>) {
   %idx = arith.constant 0 : index
-  %0 = vector.load %src[%idx] : memref<?xf32>, vector<[4x4]xf32>
+  %0 = vector.load %src[%idx] : memref<?xf32>, vector<[4]x[4]xf32>
   %1 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
 
-  // expected-error @+1 {{expected 1-d vector for operand #1}}
-  %op = vector.outerproduct %0, %1 : vector<[4x4]xf32>, vector<[4]xf32>
+  // expected-error @+1 {{'vector.outerproduct' op expected 1-d vector for operand #1}}
+  %op = vector.outerproduct %0, %1 : vector<[4]x[4]xf32>, vector<[4]xf32>
 }