[mlir] Add first-class support for scalability in VectorType dims

mlir/include/mlir/IR/BuiltinTypes.h

@@ -12,6 +12,7 @@
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/Support/ADTExtras.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace llvm {
 class BitVector;
@@ -181,6 +182,239 @@ class BaseMemRefType : public Type, public ShapedType::Trait<BaseMemRefType> {
   operator ShapedType() const { return llvm::cast<ShapedType>(*this); }
 };
 
+//===----------------------------------------------------------------------===//
+// VectorDim
+//===----------------------------------------------------------------------===//
+
+/// This class represents a dimension of a vector type. Unlike other ShapedTypes
+/// vector dimensions can have scalable quantities, which means the dimension
+/// has a known minimum size, which is scaled by a constant that is only
+/// known at runtime.
+class VectorDim {
+public:
+  explicit constexpr VectorDim(int64_t quantity, bool scalable)
+      : quantity(quantity), scalable(scalable){};
+
+  /// Constructs a new fixed dimension.
+  constexpr static VectorDim getFixed(int64_t quantity) {
+    return VectorDim(quantity, false);
+  }
+
+  /// Constructs a new scalable dimension.
+  constexpr static VectorDim getScalable(int64_t quantity) {
+    return VectorDim(quantity, true);
+  }
+
+  /// Returns true if this dimension is scalable;
+  constexpr bool isScalable() const { return scalable; }
+
+  /// Returns true if this dimension is fixed.
+  constexpr bool isFixed() const { return !isScalable(); }
+
+  /// Returns the minimum number of elements this dimension can contain.
+  constexpr int64_t getMinSize() const { return quantity; }
+
+  /// If this dimension is fixed returns the number of elements, otherwise
+  /// aborts.
+  constexpr int64_t getFixedSize() const {
+    assert(isFixed());
+    return quantity;
+  }
+
+  constexpr bool operator==(VectorDim const &dim) const {
+    return quantity == dim.quantity && scalable == dim.scalable;
+  }
+
+  constexpr bool operator!=(VectorDim const &dim) const {
+    return !(*this == dim);
+  }
+
+  /// Print the dim.
+  void print(raw_ostream &os) {
+    if (isScalable())
+      os << '[';
+    os << getMinSize();
+    if (isScalable())
+      os << ']';
+  }
+
+  /// Helper class for indexing into a list of sizes (and possibly empty) list
+  /// of scalable dimensions, extracting VectorDim elements.
+  struct Indexer {
+    explicit Indexer(ArrayRef<int64_t> sizes, ArrayRef<bool> scalableDims)
+        : sizes(sizes), scalableDims(scalableDims) {
+      assert(
+          scalableDims.empty() ||
+          sizes.size() == scalableDims.size() &&
+              "expected `scalableDims` to be empty or match `sizes` in length");
+    }
+
+    VectorDim operator[](size_t idx) const {
+      int64_t size = sizes[idx];
+      bool scalable = scalableDims.empty() ? false : scalableDims[idx];
+      return VectorDim(size, scalable);
+    }
+
+    ArrayRef<int64_t> sizes;
+    ArrayRef<bool> scalableDims;
+  };
+
+private:
+  int64_t quantity;
+  bool scalable;
+};
+
+inline raw_ostream &operator<<(raw_ostream &os, VectorDim dim) {
+  dim.print(os);
+  return os;
+}
+
+//===----------------------------------------------------------------------===//
+// VectorDims
+//===----------------------------------------------------------------------===//
+
+/// Represents a non-owning list of vector dimensions. The underlying dimension
+/// sizes and scalability flags are stored a two seperate lists to match the
+/// storage of a VectorType.
+class VectorDims : public VectorDim::Indexer {
+public:
+  using VectorDim::Indexer::Indexer;
+
+  class Iterator : public llvm::iterator_facade_base<
+                       Iterator, std::random_access_iterator_tag, VectorDim,
+                       std::ptrdiff_t, VectorDim, VectorDim> {
+  public:
+    Iterator(VectorDim::Indexer indexer, size_t index)
+        : indexer(indexer), index(index){};
+
+    // Iterator boilerplate.
+    ptrdiff_t operator-(const Iterator &rhs) const { return index - rhs.index; }
+    bool operator==(const Iterator &rhs) const { return index == rhs.index; }
+    bool operator<(const Iterator &rhs) const { return index < rhs.index; }
+    Iterator &operator+=(ptrdiff_t offset) {
+      index += offset;
+      return *this;
+    }
+    Iterator &operator-=(ptrdiff_t offset) {
+      index -= offset;
+      return *this;
+    }
+    VectorDim operator*() const { return indexer[index]; }
+
+    VectorDim::Indexer getIndexer() const { return indexer; }
+    ptrdiff_t getIndex() const { return index; }
+
+  private:
+    VectorDim::Indexer indexer;
+    ptrdiff_t index;
+  };
+
+  // Generic definitions.
+  using value_type = VectorDim;
+  using iterator = Iterator;
+  using const_iterator = Iterator;
+  using reverse_iterator = std::reverse_iterator<iterator>;
+  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
+  using size_type = size_t;
+  using difference_type = ptrdiff_t;
+
+  /// Construct from iterator pair.
+  VectorDims(Iterator begin, Iterator end)
+      : VectorDims(VectorDims(begin.getIndexer())
+                       .slice(begin.getIndex(), end - begin)) {}
+
+  VectorDims(VectorDim::Indexer indexer) : VectorDim::Indexer(indexer){};
+
+  Iterator begin() const { return Iterator(*this, 0); }
+  Iterator end() const { return Iterator(*this, size()); }
+
+  /// Check if the dims are empty.
+  bool empty() const { return sizes.empty(); }
+
+  /// Get the number of dims.
+  size_t size() const { return sizes.size(); }
+
+  /// Return the first dim.
+  VectorDim front() const { return (*this)[0]; }
+
+  /// Return the last dim.
+  VectorDim back() const { return (*this)[size() - 1]; }
+
+  /// Chop of thie first \p n dims, and keep the remaining \p m
+  /// dims.
+  VectorDims slice(size_t n, size_t m) const {
+    ArrayRef<int64_t> newSizes = sizes.slice(n, m);
+    ArrayRef<bool> newScalableDims =
+        scalableDims.empty() ? ArrayRef<bool>{} : scalableDims.slice(n, m);
+    return VectorDims(newSizes, newScalableDims);
+  }
+
+  /// Drop the first \p n dims.
+  VectorDims dropFront(size_t n = 1) const { return slice(n, size() - n); }
+
+  /// Drop the last \p n dims.
+  VectorDims dropBack(size_t n = 1) const { return slice(0, size() - n); }
+
+  /// Return a copy of *this with only the first \p n elements.
+  VectorDims takeFront(size_t n = 1) const {
+    if (n >= size())
+      return *this;
+    return dropBack(size() - n);
+  }
+
+  /// Return a copy of *this with only the last \p n elements.
+  VectorDims takeBack(size_t n = 1) const {
+    if (n >= size())
+      return *this;
+    return dropFront(size() - n);
+  }
+
+  /// Return copy of *this with the first n dims matching the predicate removed.
+  template <class PredicateT>
+  VectorDims dropWhile(PredicateT predicate) const {
+    return VectorDims(llvm::find_if_not(*this, predicate), end());
+  }
+
+  /// Returns true if one or more of the dims are scalable.
+  bool hasScalableDims() const {
+    return llvm::is_contained(getScalableDims(), true);
+  }
+
+  /// Check for dim equality.
+  bool equals(VectorDims rhs) const {
+    if (size() != rhs.size())
+      return false;
+    return std::equal(begin(), end(), rhs.begin());
+  }
+
+  /// Check for dim equality.
+  bool equals(ArrayRef<VectorDim> rhs) const {
+    if (size() != rhs.size())
+      return false;
+    return std::equal(begin(), end(), rhs.begin());
+  }
+
+  /// Return the underlying sizes.
+  ArrayRef<int64_t> getSizes() const { return sizes; }
+
+  /// Return the underlying scalable dims.
+  ArrayRef<bool> getScalableDims() const { return scalableDims; }
+};
+
+inline bool operator==(VectorDims lhs, VectorDims rhs) {
+  return lhs.equals(rhs);
+}
+
+inline bool operator!=(VectorDims lhs, VectorDims rhs) { return !(lhs == rhs); }
+
+inline bool operator==(VectorDims lhs, ArrayRef<VectorDim> rhs) {
+  return lhs.equals(rhs);
+}
+
+inline bool operator!=(VectorDims lhs, ArrayRef<VectorDim> rhs) {
+  return !(lhs == rhs);
+}
+
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/BuiltinTypes.td

@@ -1114,13 +1114,36 @@ def Builtin_Vector : Builtin_Type<"Vector", "vector", [ShapedTypeInterface], "Ty
         scalableDims = isScalableVec;
       }
       return $_get(elementType.getContext(), shape, elementType, scalableDims);
+    }]>,
+    TypeBuilderWithInferredContext<(ins "Type":$elementType, "ArrayRef<VectorDim>": $shape), [{
+      SmallVector<int64_t> sizes;
+      SmallVector<bool> scalableDims;
+      for (VectorDim dim : shape) {
+        sizes.push_back(dim.getMinSize());
+        scalableDims.push_back(dim.isScalable());
+      }
+      return get(sizes, elementType, scalableDims);
+    }]>,
+    TypeBuilderWithInferredContext<(ins "Type":$elementType, "VectorDims": $shape), [{
+      return get(shape.getSizes(), elementType, shape.getScalableDims());
     }]>
   ];
   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 the value of the specified dimension (including scalability).
+    VectorDim getDim(unsigned idx) const {
+      assert(idx < getRank() && "invalid dim index for vector type");
+      return getDims()[idx];
+    }
+
+    /// Returns the dimensions of this vector type (including scalability).
+    VectorDims getDims() const {
+      return VectorDims(getShape(), getScalableDims());
+    }
+
     /// 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.

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.