Demonstrate using the new APIs in scalable-aware code :)

This is not a complete change, this just updates a few examples found
by grepping for getScalableDims().

Author: MacDue

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

@@ -490,12 +490,11 @@ FailureOr<Type> LLVMTypeConverter::convertVectorType(VectorType type) const {
     return {};
   if (type.getShape().empty())
     return VectorType::get({1}, elementType);
-  Type vectorType = VectorType::get(type.getShape().back(), elementType,
-                                    type.getScalableDims().back());
+  Type vectorType = VectorType::get(elementType, type.getDims().takeBack());
   assert(LLVM::isCompatibleVectorType(vectorType) &&
          "expected vector type compatible with the LLVM dialect");
   // Only the trailing dimension can be scalable.
-  if (llvm::is_contained(type.getScalableDims().drop_back(), true))
+  if (type.getDims().dropBack().hasScalableDims())
     return failure();
   auto shape = type.getShape();
   for (int i = shape.size() - 2; i >= 0; --i)

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

@@ -37,8 +37,7 @@ using namespace mlir::vector;
 // Helper to reduce vector type by *all* but one rank at back.
 static VectorType reducedVectorTypeBack(VectorType tp) {
   assert((tp.getRank() > 1) && "unlowerable vector type");
-  return VectorType::get(tp.getShape().take_back(), tp.getElementType(),
-                         tp.getScalableDims().take_back());
+  return VectorType::get(tp.getElementType(), tp.getDims().takeBack());
 }
 
 // Helper that picks the proper sequence for inserting.

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

@@ -319,12 +319,13 @@ static FailureOr<MemRefType> unpackOneDim(MemRefType type) {
   auto vectorType = dyn_cast<VectorType>(type.getElementType());
   // Vectors with leading scalable dims are not supported.
   // It may be possible to support these in future by using dynamic memref dims.
-  if (vectorType.getScalableDims().front())
+  VectorDim leadingDim = vectorType.getDims().front();
+  if (leadingDim.isScalable())
     return failure();
   auto memrefShape = type.getShape();
   SmallVector<int64_t, 8> newMemrefShape;
   newMemrefShape.append(memrefShape.begin(), memrefShape.end());
-  newMemrefShape.push_back(vectorType.getDimSize(0));
+  newMemrefShape.push_back(leadingDim.getFixedSize());
   return MemRefType::get(newMemrefShape,
                          VectorType::Builder(vectorType).dropDim(0));
 }
@@ -1091,18 +1092,17 @@ struct UnrollTransferReadConversion
     auto vecType = dyn_cast<VectorType>(vec.getType());
     auto xferVecType = xferOp.getVectorType();
 
-    if (xferVecType.getScalableDims()[0]) {
+    VectorDim dim = xferVecType.getDim(0);
+    if (dim.isScalable()) {
       // Cannot unroll a scalable dimension at compile time.
       return failure();
     }
 
     VectorType newXferVecType = VectorType::Builder(xferVecType).dropDim(0);
 
-    int64_t dimSize = xferVecType.getShape()[0];
-
     // Generate fully unrolled loop of transfer ops.
     Location loc = xferOp.getLoc();
-    for (int64_t i = 0; i < dimSize; ++i) {
+    for (int64_t i = 0; i < dim.getFixedSize(); ++i) {
       Value iv = rewriter.create<arith::ConstantIndexOp>(loc, i);
 
       vec = generateInBoundsCheck(

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

@@ -30,8 +30,7 @@ using namespace mlir::arm_sve;
 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.getScalableDims());
+    return VectorType::get(i1Type, sVectorType.getDims());
   return nullptr;
 }
 

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

@@ -1217,9 +1217,8 @@ vectorizeOneOp(RewriterBase &rewriter, VectorizationState &state,
     assert(vecOperand && "Vector operand couldn't be found");
 
     if (firstMaxRankedType) {
-      auto vecType = VectorType::get(firstMaxRankedType.getShape(),
-                                     getElementTypeOrSelf(vecOperand.getType()),
-                                     firstMaxRankedType.getScalableDims());
+      auto vecType = VectorType::get(getElementTypeOrSelf(vecOperand.getType()),
+                                     firstMaxRankedType.getDims());
       vecOperands.push_back(broadcastIfNeeded(rewriter, vecOperand, vecType));
     } else {
       vecOperands.push_back(vecOperand);
@@ -1230,8 +1229,7 @@ vectorizeOneOp(RewriterBase &rewriter, VectorizationState &state,
   for (Type resultType : op->getResultTypes()) {
     resultTypes.push_back(
         firstMaxRankedType
-            ? VectorType::get(firstMaxRankedType.getShape(), resultType,
-                              firstMaxRankedType.getScalableDims())
+            ? VectorType::get(resultType, firstMaxRankedType.getDims())
             : resultType);
   }
   //   d. Build and return the new op.

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

@@ -422,23 +422,21 @@ MultiDimReductionOp::getShapeForUnroll() {
 }
 
 LogicalResult MultiDimReductionOp::verify() {
-  SmallVector<int64_t> targetShape;
-  SmallVector<bool> scalableDims;
+  SmallVector<VectorDim> targetDims;
   Type inferredReturnType;
-  auto sourceScalableDims = getSourceVectorType().getScalableDims();
-  for (auto it : llvm::enumerate(getSourceVectorType().getShape()))
-    if (!llvm::any_of(getReductionDims().getValue(), [&](Attribute attr) {
-          return llvm::cast<IntegerAttr>(attr).getValue() == it.index();
-        })) {
-      targetShape.push_back(it.value());
-      scalableDims.push_back(sourceScalableDims[it.index()]);
+  for (auto [idx, dim] : llvm::enumerate(getSourceVectorType().getDims()))
+    if (!llvm::any_of(getReductionDims().getValue(),
+                      [idx = idx](Attribute attr) {
+                        return llvm::cast<IntegerAttr>(attr).getValue() == idx;
+                      })) {
+      targetDims.push_back(dim);
     }
   // TODO: update to also allow 0-d vectors when available.
-  if (targetShape.empty())
+  if (targetDims.empty())
     inferredReturnType = getSourceVectorType().getElementType();
   else
-    inferredReturnType = VectorType::get(
-        targetShape, getSourceVectorType().getElementType(), scalableDims);
+    inferredReturnType =
+        VectorType::get(getSourceVectorType().getElementType(), targetDims);
   if (getType() != inferredReturnType)
     return emitOpError() << "destination type " << getType()
                          << " is incompatible with source type "
@@ -450,9 +448,8 @@ LogicalResult MultiDimReductionOp::verify() {
 /// 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.getScalableDims());
+  return VectorType::get(IntegerType::get(vecType.getContext(), /*width=*/1),
+                         vecType.getDims());
 }
 
 namespace {
@@ -491,8 +488,7 @@ struct ElideUnitDimsInMultiDimReduction
     if (auto dstVecType = dyn_cast<VectorType>(reductionOp.getDestType())) {
       if (mask) {
         VectorType newMaskType =
-            VectorType::get(dstVecType.getShape(), rewriter.getI1Type(),
-                            dstVecType.getScalableDims());
+            VectorType::get(rewriter.getI1Type(), dstVecType.getDims());
         mask = rewriter.create<vector::ShapeCastOp>(loc, newMaskType, mask);
       }
       cast = rewriter.create<vector::ShapeCastOp>(
@@ -559,9 +555,8 @@ LogicalResult ReductionOp::verify() {
 /// 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.getScalableDims());
+  return VectorType::get(IntegerType::get(vecType.getContext(), /*width=*/1),
+                         vecType.getDims());
 }
 
 Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
@@ -1252,8 +1247,7 @@ ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
     auto n = std::min<size_t>(adaptor.getStaticPosition().size(),
                               vectorType.getRank());
     inferredReturnTypes.push_back(VectorType::get(
-        vectorType.getShape().drop_front(n), vectorType.getElementType(),
-        vectorType.getScalableDims().drop_front(n)));
+        vectorType.getElementType(), vectorType.getDims().dropFront(n)));
   }
   return success();
 }
@@ -3040,15 +3034,11 @@ ParseResult OuterProductOp::parse(OpAsmParser &parser, OperationState &result) {
 
   VectorType resType;
   if (vRHS) {
-    SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0],
-                                      vRHS.getScalableDims()[0]};
-    resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
-                              vLHS.getElementType(), scalableDimsRes);
+    resType = VectorType::get(vLHS.getElementType(),
+                              {vLHS.getDim(0), vRHS.getDim(0)});
   } else {
     // Scalar RHS operand
-    SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0]};
-    resType = VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType(),
-                              scalableDimsRes);
+    resType = VectorType::get(vLHS.getElementType(), {vLHS.getDim(0)});
   }
 
   if (!result.attributes.get(OuterProductOp::getKindAttrName(result.name))) {
@@ -3115,9 +3105,8 @@ LogicalResult OuterProductOp::verify() {
 /// 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.getScalableDims());
+  return VectorType::get(IntegerType::get(vecType.getContext(), /*width=*/1),
+                         vecType.getDims());
 }
 
 //===----------------------------------------------------------------------===//
@@ -5064,25 +5053,13 @@ class ShapeCastConstantFolder final : public OpRewritePattern<ShapeCastOp> {
 ///   vector<4x1x1xi1> --> vector<4x1>
 ///
 static VectorType trimTrailingOneDims(VectorType oldType) {
-  ArrayRef<int64_t> oldShape = oldType.getShape();
-  ArrayRef<int64_t> newShape = oldShape;
-
-  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
-  ArrayRef<bool> newScalableDims = oldScalableDims;
-
-  while (!newShape.empty() && newShape.back() == 1 && !newScalableDims.back()) {
-    newShape = newShape.drop_back(1);
-    newScalableDims = newScalableDims.drop_back(1);
-  }
-
-  // Make sure we have at least 1 dimension.
+  // Note: This will always keep at least one dim (even if it's a unit dim).
   // TODO: Add support for 0-D vectors.
-  if (newShape.empty()) {
-    newShape = oldShape.take_back();
-    newScalableDims = oldScalableDims.take_back();
-  }
+  VectorDims newDims = oldType.getDims();
+  while (newDims.size() > 1 && newDims.back() == VectorDim::getFixed(1))
+    newDims = newDims.dropBack();
 
-  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+  return VectorType::get(oldType.getElementType(), newDims);
 }
 
 /// Folds qualifying shape_cast(create_mask) into a new create_mask

mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp

@@ -113,13 +113,11 @@ struct TransferReadPermutationLowering
     permutationMap = inversePermutation(permutationMap);
     AffineMap newMap = permutationMap.compose(map);
     // Apply the reverse transpose to deduce the type of the transfer_read.
-    ArrayRef<int64_t> originalShape = op.getVectorType().getShape();
-    SmallVector<int64_t> newVectorShape(originalShape.size());
-    ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();
-    SmallVector<bool> newScalableDims(originalShape.size());
+    auto originalDims = op.getVectorType().getDims();
+    SmallVector<VectorDim> newVectorDims(op.getVectorType().getRank(),
+                                         VectorDim::getFixed(0));
     for (const auto &pos : llvm::enumerate(permutation)) {
-      newVectorShape[pos.value()] = originalShape[pos.index()];
-      newScalableDims[pos.value()] = originalScalableDims[pos.index()];
+      newVectorDims[pos.value()] = originalDims[pos.index()];
     }
 
     // Transpose in_bounds attribute.
@@ -129,8 +127,8 @@ struct TransferReadPermutationLowering
                          : ArrayAttr();
 
     // Generate new transfer_read operation.
-    VectorType newReadType = VectorType::get(
-        newVectorShape, op.getVectorType().getElementType(), newScalableDims);
+    VectorType newReadType =
+        VectorType::get(op.getVectorType().getElementType(), newVectorDims);
     Value newRead = rewriter.create<vector::TransferReadOp>(
         op.getLoc(), newReadType, op.getSource(), op.getIndices(),
         AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),

mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp

@@ -344,10 +344,8 @@ class TransposeOpLowering : public OpRewritePattern<vector::TransposeOp> {
       // Source with leading unit dim (inverse) is also replaced. Unit dim must
       // be fixed. Non-unit can be scalable.
       if (resType.getRank() == 2 &&
-          ((resType.getShape().front() == 1 &&
-            !resType.getScalableDims().front()) ||
-           (resType.getShape().back() == 1 &&
-            !resType.getScalableDims().back())) &&
+          (resType.getDims().front() == VectorDim::getFixed(1) ||
+           resType.getDims().back() == VectorDim::getFixed(1)) &&
           transp == ArrayRef<int64_t>({1, 0})) {
         rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, input);
         return success();

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

@@ -25,24 +25,15 @@ using namespace mlir::vector;
 // Trims leading one dimensions from `oldType` and returns the result type.
 // Returns `vector<1xT>` if `oldType` only has one element.
 static VectorType trimLeadingOneDims(VectorType oldType) {
-  ArrayRef<int64_t> oldShape = oldType.getShape();
-  ArrayRef<int64_t> newShape = oldShape;
-
-  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
-  ArrayRef<bool> newScalableDims = oldScalableDims;
-
-  while (!newShape.empty() && newShape.front() == 1 &&
-         !newScalableDims.front()) {
-    newShape = newShape.drop_front(1);
-    newScalableDims = newScalableDims.drop_front(1);
-  }
+  VectorDims oldDims = oldType.getDims();
+  VectorDims newDims = oldDims.dropWhile(
+      [](VectorDim dim) { return dim == VectorDim::getFixed(1); });
 
   // Make sure we have at least 1 dimension per vector type requirements.
-  if (newShape.empty()) {
-    newShape = oldShape.take_back();
-    newScalableDims = oldType.getScalableDims().take_back();
-  }
-  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+  if (newDims.empty())
+    newDims = oldDims.takeBack(1);
+
+  return VectorType::get(oldType.getElementType(), newDims);
 }
 
 /// Return a smallVector of size `rank` containing all zeros.

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

@@ -316,15 +316,11 @@ static int getReducedRank(ArrayRef<int64_t> shape) {
 /// Trims non-scalable one dimensions from `oldType` and returns the result
 /// type.
 static VectorType trimNonScalableUnitDims(VectorType oldType) {
-  SmallVector<int64_t> newShape;
-  SmallVector<bool> newScalableDims;
-  for (auto [dimIdx, dimSize] : llvm::enumerate(oldType.getShape())) {
-    if (dimSize == 1 && !oldType.getScalableDims()[dimIdx])
-      continue;
-    newShape.push_back(dimSize);
-    newScalableDims.push_back(oldType.getScalableDims()[dimIdx]);
-  }
-  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+  auto newDims = llvm::to_vector(
+      llvm::make_filter_range(oldType.getDims(), [](VectorDim dim) {
+        return dim != VectorDim::getFixed(1);
+      }));
+  return VectorType::get(oldType.getElementType(), newDims);
 }
 
 // Rewrites vector.create_mask 'op' to drop non-scalable one dimensions.
@@ -337,9 +333,9 @@ createMaskDropNonScalableUnitDims(PatternRewriter &rewriter, Location loc,
     return failure();
 
   SmallVector<Value> reducedOperands;
-  for (auto [dim, dimIsScalable, operand] : llvm::zip_equal(
-           type.getShape(), type.getScalableDims(), op.getOperands())) {
-    if (dim == 1 && !dimIsScalable) {
+  for (auto [dim, operand] :
+       llvm::zip_equal(type.getDims(), op.getOperands())) {
+    if (dim == VectorDim::getFixed(1)) {
       // If the mask for the unit dim is not a constant of 1, do nothing.
       auto constant = operand.getDefiningOp<arith::ConstantIndexOp>();
       if (!constant || (constant.value() != 1))