[mlir][VectorType] Remove `numScalableDims` from the vector type

[awarzynski]

Summary

This is a follow-up of https://reviews.llvm.org/D153372 in which
numScalableDims (single integer) was effectively replaced with
isScalableDim bitmask.

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/

Depends on D153372

[dcaballe]

LGTM. Only a renaming suggesting that I think it's important.

mlir/include/mlir/IR/BuiltinTypes.td

1083

nit: space after comma :)

mlir/lib/IR/BuiltinTypes.cpp

263

Should we rename the isScalableDim field to scalableDims? isScalableDim sounds like a query to know if a specific dimension is scalable

[awarzynski]

mlir/lib/IR/BuiltinTypes.cpp

263

+1 Updated in https://reviews.llvm.org/D153372 (base patch for this change)

[awarzynski]

Rebase after the updates in D153372

[kuhar]

mlir/include/mlir/IR/BuiltinTypes.td

1100–1107

this can be: return llvm::is_contained(getScalableDims(), true);

1105–1107

Is this the same as checking if all values are true? If yes, I'd use return !llvm::is_contained(getScalableDims(), false);

[awarzynski]

Incorporate suggestions from @kuhar - thanks!

[kuhar]

mlir/include/mlir/IR/BuiltinTypes.td

1102–1104

Reading this definition again, should we consider a 0-d vector scalable? Logically I don't see a contradiction, but I can see that it could be confusing when one has to check that a vector is *both* scalable and has all dims scalable.

Perhaps a more useful/less error prone function would be something like isFullyScalable that combines both checks?

[awarzynski]

mlir/include/mlir/IR/BuiltinTypes.td

1102–1104

You are raising very good points.

Reading this definition again, should we consider a 0-d vector scalable?

For 0-d vectors there are no "dimensions", so there's no mechanism that would allow us to distinguish fixed-width from scalable cases. For now I'd treat them as fixed-width. Simply as the safer, better established default option. WDYT?

In general, IMHO, we need a bit more time to see how these 0-d vectors get created/processed in the context of scalable vectorisation etc. Without that it's hard to make an informed decision.

[awarzynski]

Make sure that 0-d vectors are treated as fixed-width

[kuhar]

LGTM, thanks for the discussion!

mlir/include/mlir/IR/BuiltinTypes.td

1102–1104

SGTM

1103

nit: Use // is for local comments and end the sentence with .

mlir/include/mlir/IR/BuiltinDialectBytecode.td

def UnrankedTensorType : DialectType<(type
   Type:$elementType
 )>;
 
 let cType = "VectorType" in {
 def VectorType : DialectType<(type
   Array<SignedVarIntList>:$shape,
   Type:$elementType
 )> {
-  let printerPredicate = "!$_val.getNumScalableDims()";
+  let printerPredicate = "!$_val.isScalable()";
 }
 
 def VectorTypeWithScalableDims : DialectType<(type
   Array<BoolList>:$scalableDims,
-  VarInt:$numScalableDims,
   Array<SignedVarIntList>:$shape,
   Type:$elementType
 )> {
-  let printerPredicate = "$_val.getNumScalableDims()";
+  let printerPredicate = "$_val.isScalable()";
   // Note: order of serialization does not match order of builder.
-  let cBuilder = "get<$_resultType>(context, shape, elementType, numScalableDims, scalableDims)";
+  let cBuilder = "get<$_resultType>(context, shape, elementType, 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()),
         scalableDims(other.getScalableDims()) {}
 
   /// Build from scratch.
   Builder(ArrayRef<int64_t> shape, Type elementType,
           unsigned numScalableDims = 0, ArrayRef<bool> scalableDims = {})
-      : shape(shape), elementType(elementType),
-        numScalableDims(numScalableDims) {
+      : shape(shape), elementType(elementType) {
     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,
                     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, scalableDims);
+    return VectorType::get(shape, elementType, 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

mlir/include/mlir/IR/BuiltinTypes.td

let description = [{
     // 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,
     ArrayRefParameter<"bool">:$scalableDims
   );
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape, "Type":$elementType,
-      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, scalableDims);
+      return $_get(elementType.getContext(), shape, elementType, 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
     /// type. In particular, vectors can consist of integer, index, or float
     /// primitives.
     static bool isValidElementType(Type t) {
       return ::llvm::isa<IntegerType, IndexType, FloatType>(t);
     }
 
     /// Returns true if the vector contains scalable dimensions.
     bool isScalable() const {
-      return getNumScalableDims() > 0;
+      return llvm::is_contained(getScalableDims(), true);
+    }
+    bool allDimsScalable() const {
+      /// Treat 0-d vectors as fixed size
+      if (getRank() == 0)
+        return false;
+      return !llvm::is_contained(getScalableDims(), false);
     }
 
     /// Get or create a new VectorType with the same shape as `this` and an
     /// element type of bitwidth scaled by `scale`.
     /// Return null if the scaled element type cannot be represented.
     VectorType scaleElementBitwidth(unsigned scale);
 
     /// Returns if this type is ranked (always true).
     bool hasRank() const { return true; }

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,
                                        SmallVectorImpl<bool> &scalableDims);
   ParseResult parseDimensionListRanked(SmallVectorImpl<int64_t> &dimensions,
                                        bool allowDynamic = true,
                                        bool withTrailingX = true);
   ParseResult parseIntegerInDimensionList(int64_t &value);
   ParseResult parseXInDimensionList();
 
   //===--------------------------------------------------------------------===//

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, scalableDims))
+  if (parseVectorDimensionList(dimensions, 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,
-                         scalableDims);
+  return VectorType::get(dimensions, elementType, scalableDims);
 }
 
 /// 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,
                                  SmallVectorImpl<bool> &scalableDims) {
-  numScalableDims = 0;
   // If there is a set of fixed-length dimensions, consume it
   while (getToken().is(Token::integer) || getToken().is(Token::l_square)) {
     int64_t value;
     bool scalable = consumeIf(Token::l_square);
     if (parseIntegerInDimensionList(value))
       return failure();
     dimensions.push_back(value);
     if (scalable) {
       if (!consumeIf(Token::r_square))
         return emitWrongTokenError("missing ']' closing scalable dimension");
-      numScalableDims++;
     }
     scalableDims.push_back(scalable);
     // Make sure we have an 'x' or something like 'xbf32'.
     if (parseXInDimensionList())
       return failure();
   }
 
   return success();

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.getScalableDims().back());
+  Type vectorType = VectorType::get(type.getShape().back(), elementType,
+                                    type.getScalableDims().back());
   assert(LLVM::isCompatibleVectorType(vectorType) &&
          "expected vector type compatible with the LLVM dialect");
+  assert((type.isScalable() == type.allDimsScalable()) &&
+         "expected scalable vector with all dims scalable");
   auto shape = type.getShape();
   for (int i = shape.size() - 2; i >= 0; --i)
     vectorType = LLVM::LLVMArrayType::get(vectorType, shape[i]);
   return vectorType;
 }
 
 /// Convert a type in the context of the default or bare pointer calling
 /// convention. Calling convention sensitive types, such as MemRefType and

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 #include <optional>
 
 using namespace mlir;
 using namespace mlir::vector;
 
 // Helper to reduce vector type by one rank at front.
 static VectorType reducedVectorTypeFront(VectorType tp) {
   assert((tp.getRank() > 1) && "unlowerable vector type");
-  unsigned numScalableDims = tp.getNumScalableDims();
-  if (tp.getShape().size() == numScalableDims)
-    --numScalableDims;
   return VectorType::get(tp.getShape().drop_front(), tp.getElementType(),
-                         numScalableDims);
+                         tp.getScalableDims().drop_front());
 }
 
 // Helper to reduce vector type by *all* but one rank at back.
 static VectorType reducedVectorTypeBack(VectorType tp) {
   assert((tp.getRank() > 1) && "unlowerable vector type");
-  unsigned numScalableDims = tp.getNumScalableDims();
-  if (numScalableDims > 0)
-    --numScalableDims;
   return VectorType::get(tp.getShape().take_back(), tp.getElementType(),
-                         numScalableDims);
+                         tp.getScalableDims().take_back());
 }
 
 // Helper that picks the proper sequence for inserting.
 static Value insertOne(ConversionPatternRewriter &rewriter,
                        LLVMTypeConverter &typeConverter, Location loc,
                        Value val1, Value val2, Type llvmType, int64_t rank,
                        int64_t pos) {
   assert(rank > 0 && "0-D vector corner case should have been handled already");

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.getScalableDims());
   return i1Type;
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 

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.getScalableDims());
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // Tablegen Definitions
 //===----------------------------------------------------------------------===//
 

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

Type mlir::LLVM::getVectorType(Type elementType, unsigned numElements,
   if (useLLVM) {
     if (isScalable)
       return LLVMScalableVectorType::get(elementType, numElements);
     return LLVMFixedVectorType::get(elementType, numElements);
   }
 
   // LLVM vectors are always 1-D, hence only 1 bool is required to mark it as
   // scalable/non-scalable.
-  SmallVector<bool> scalableDims(1, false);
-  if (isScalable)
-    scalableDims[0] = true;
-
-  return VectorType::get(numElements, elementType,
-                         static_cast<unsigned>(isScalable), scalableDims);
+  return VectorType::get(numElements, elementType, {isScalable});
 }
 
 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(),
Type mlir::LLVM::getScalableVectorType(Type elementType, unsigned numElements) {
   bool useLLVM = LLVMScalableVectorType::isValidElementType(elementType);
   bool useBuiltIn = VectorType::isValidElementType(elementType);
   (void)useBuiltIn;
   assert((useLLVM ^ useBuiltIn) && "expected LLVM-compatible scalable-vector "
                                    "type to be either builtin or LLVM dialect "
                                    "type");
   if (useLLVM)
     return LLVMScalableVectorType::get(elementType, numElements);
-  return VectorType::get(numElements, elementType, /*numScalableDims=*/1);
+
+  // LLVM vectors are always 1-D, hence only 1 bool is required to mark it as
+  // scalable/non-scalable.
+  return VectorType::get(numElements, elementType, /*scalableDims=*/true);
 }
 
 llvm::TypeSize mlir::LLVM::getPrimitiveTypeSizeInBits(Type type) {
   assert(isCompatibleType(type) &&
          "expected a type compatible with the LLVM dialect");
 
   return llvm::TypeSwitch<Type, llvm::TypeSize>(type)
       .Case<BFloat16Type, Float16Type>(

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,
-                           scalableDims);
+    return VectorType::get(vectorShape, elementType, 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.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.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,
-                         vl.enableVLAVectorization);
+  return VectorType::get(vl.vectorLength, etp, 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.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

 }
 
 Type Merger::inferType(ExprId e, Value src) const {
   // Obtain the destination type from the cast node.
   Type dtp = exp(e).val.getType();
   // Inspect source type. For vector types, apply the same
   // vectorization to the destination type.
   if (auto vtp = dyn_cast<VectorType>(src.getType()))
-    return VectorType::get(vtp.getNumElements(), dtp, vtp.getNumScalableDims());
+    return VectorType::get(vtp.getNumElements(), dtp, vtp.getScalableDims());
   return dtp;
 }
 
 /// Ensures that sparse compiler can generate code for expression.
 static bool isAdmissibleBranchExp(Operation *op, Block *block, Value v) {
   // Arguments are always admissible.
   if (isa<BlockArgument>(v))
     return true;

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(), vecType.getScalableDims());
+  return VectorType::get(vecType.getShape(),
+                         IntegerType::get(vecType.getContext(), /*width=*/1),
+                         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(), vecType.getScalableDims());
+  return VectorType::get(vecType.getShape(),
+                         IntegerType::get(vecType.getContext(), /*width=*/1),
+                         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:
Type ContractionOp::getExpectedMaskType() {
   for (auto [dimIdx, dimSize] : llvm::enumerate(lhsType.getShape()))
     maskShape[lhsIdxMap.getDimPosition(dimIdx)] = dimSize;
   for (auto [dimIdx, dimSize] : llvm::enumerate(rhsType.getShape()))
     maskShape[rhsIdxMap.getDimPosition(dimIdx)] = dimSize;
 
   assert(!ShapedType::isDynamicShape(maskShape) &&
          "Mask shape couldn't be computed");
   // TODO: Extend the scalable vector type representation with a bit map.
-  assert(lhsType.getNumScalableDims() == 0 &&
-         rhsType.getNumScalableDims() == 0 &&
+  assert(!lhsType.isScalable() && !rhsType.isScalable() &&
          "Scalable vectors are not supported yet");
 
   return VectorType::get(maskShape,
                          IntegerType::get(lhsType.getContext(), /*width=*/1));
 }
 
 SmallVector<StringRef> ContractionOp::getTraitAttrNames() {
   return SmallVector<StringRef>{getIndexingMapsAttrName(),
ParseResult OuterProductOp::parse(OpAsmParser &parser, OperationState &result) {
   if (!vLHS)
     return parser.emitError(parser.getNameLoc(),
                             "expected vector type for operand #1");
 
   VectorType resType;
   if (vRHS) {
     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,
-                              scalableDimsRes);
+                              vLHS.getElementType(), 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, scalableDimsRes);
+                              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(), vecType.getScalableDims());
+  return VectorType::get(vecType.getShape(),
+                         IntegerType::get(vecType.getContext(), /*width=*/1),
+                         vecType.getScalableDims());
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ReshapeOp::verify() {
   // Verify that rank(numInputs/outputs) + numFixedVec dim matches vec rank.
static VectorType inferTransferOpMaskType(VectorType vecType,
   auto i1Type = IntegerType::get(permMap.getContext(), 1);
   AffineMap invPermMap = inversePermutation(compressUnusedDims(permMap));
   assert(invPermMap && "Inversed permutation map couldn't be computed");
   SmallVector<int64_t, 8> maskShape = invPermMap.compose(vecType.getShape());
 
   SmallVector<bool> scalableDims =
       applyPermutationMap(invPermMap, vecType.getScalableDims());
 
-  return VectorType::get(maskShape, i1Type, vecType.getNumScalableDims(),
-                         scalableDims);
+  return VectorType::get(maskShape, i1Type, 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;

 }
 
 // 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(), vecType.getScalableDims());
+  return VectorType::get(vecType.getShape(),
+                         IntegerType::get(vecType.getContext(), /*width=*/1),
+                         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/Dialect/Vector/Transforms/VectorTransforms.cpp

LogicalResult matchAndRewrite(ConcreteOp xferOp,
     unsigned lastIndex = llvm::size(xferOp.getIndices()) - 1;
     Value off = xferOp.getIndices()[lastIndex];
     Value dim =
         vector::createOrFoldDimOp(rewriter, loc, xferOp.getSource(), lastIndex);
     Value b = rewriter.create<arith::SubIOp>(loc, dim.getType(), dim, off);
     Value mask = rewriter.create<vector::CreateMaskOp>(
         loc,
         VectorType::get(vtp.getShape(), rewriter.getI1Type(),
-                        vtp.getNumScalableDims()),
+                        vtp.getScalableDims()),
         b);
     if (xferOp.getMask()) {
       // Intersect the in-bounds with the mask specified as an op parameter.
       mask = rewriter.create<arith::AndIOp>(loc, mask, xferOp.getMask());
     }
 
     rewriter.updateRootInPlace(xferOp, [&]() {
       xferOp.getMaskMutable().assign(mask);

mlir/lib/IR/BuiltinTypes.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult VectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  ArrayRef<int64_t> shape, Type elementType,
-                                 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))
-      return VectorType::get(getShape(), scaledEt, getNumScalableDims());
+      return VectorType::get(getShape(), scaledEt, getScalableDims());
   if (auto et = llvm::dyn_cast<FloatType>(getElementType()))
     if (auto scaledEt = et.scaleElementBitwidth(scale))
-      return VectorType::get(getShape(), scaledEt, getNumScalableDims());
+      return VectorType::get(getShape(), scaledEt, getScalableDims());
   return VectorType();
 }
 
 VectorType VectorType::cloneWith(std::optional<ArrayRef<int64_t>> shape,
                                  Type elementType) const {
   return VectorType::get(shape.value_or(getShape()), elementType,
-                         getNumScalableDims());
+                         getScalableDims());
 }
 
 //===----------------------------------------------------------------------===//
 // TensorType
 //===----------------------------------------------------------------------===//
 
 Type TensorType::getElementType() const {
   return llvm::TypeSwitch<TensorType, Type>(*this)