[mlir][BuiltinTypes] Add VectorType::sliceDims() and drop[Front/Back]Dims()

[benmxwl-arm]

Summary

These aim to be slightly safer helpers for chopping up vector types.

Before this, if you wanted to drop the first dim of a vector type you
would do something like:

VectorType::get(vectorType.getShape().drop_front(),
                vectorType.getElementType(),
                vectorType.getScalableDims().drop_front());

However, it's easy to forget the scalable dims and do:

VectorType::get(vectorType.getShape().drop_front(),
                vectorType.getElementType());

Which drops all the scalable dims. There are still many other
error-prone cases with the current design of scalable vector types
(which may mean a larger refactor is needed), but these helpers make
things little nicer.

[dcaballe]

mlir/include/mlir/IR/BuiltinTypes.td

1127

If you are into improving the VectorType, you could give a try at removing the cloneWith method in a separate patch :). It has been pointed out to me multiple times that it's not very helpful as it could just be replaced by a just a VectorType::get. I kind of agree. @nicolasvasilache

1145

These methods seem useful for any shaped type, right? Would it make sense to move them to ShapedType?

mlir/lib/IR/BuiltinTypes.cpp

278

n, m -> start/begin, size? (we should have similar APIs in MLIR/LLVM so we may want to match those)

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.td

1127

Both RankedTensorType and MemRefType have a cloneWith as well!

1145

Makes sense to me.

I'd say though that these methods are encouraging somehow inefficient idioms because chaining them would trigger context-wide creation of the intermediate entities!

[nicolasvasilache]

Generally, I'd love something with significantly more consistency across Vector/MemRef/RankedTensor types than adding a few methods to just VectorType.

For building ShapedType types from an initial type + applying some modifications, we have the fluent APIs here: https://github.com/llvm/llvm-project/blob/985e399647d591d6130ba6fe08c5b5f6cb87d9f6/mlir/include/mlir/IR/BuiltinTypes.h#L312
Can we refactor/extend those consistently instead of adding new one-off builder-like functionality in new places ?

Alternatively, we could also collectively decide that these ShapedType::Builder have to go and instead extend the interfaces; but this should be done in a consistent fashion across all uses to reduce surprises and increase API uniformity.
Given the current state of the codebase it should be much simpler to extend ShapedType::Builder for small increases in functionality.

[nicolasvasilache]

Before this, if you wanted to drop the first dim of a vector type you
would do something like:

VectorType::get(vectorType.getShape().drop_front(),
                vectorType.getElementType(),
                vectorType.getScalableDims().drop_front());

I would go:

VectorType::Builder(vectorType).dropDim(0);

[benmxwl-arm]

VectorType::Builder looks like what I want. I'm happy to drop these changes in favour of using the builder more widely :)
I hadn't seen it before; it seems to be only used in a few places (I could only find 3 functions). Is there a reason why it's not used more widely?

If you are into improving the VectorType, you could give a try at removing the cloneWith method in a separate patch :)

I think removing cloneWith would be the opposite of what I want to achieve with this patch.
For just changing the element type vectorType.cloneWith({}, newElType) seems safer (from the perspective of not dropping scalable dims), than VectorType::get(vectorType.getShape(), newElType, vectorType.getScalableDims()).

[nicolasvasilache]

In D158751#4625313, @benmxwl-arm wrote:

VectorType::Builder looks like what I want. I'm happy to drop these changes in favour of using the builder more widely :)
I hadn't seen it before; it seems to be only used in a few places (I could only find 3 functions). Is there a reason why it's not used more widely?

Great! Probably because it is quite newer than the original Tensor/MemRef/Vector::get methods that were implemented within the first few weeks of the project.
Feel free to send a PSA on discourse to raise awareness.

[benmxwl-arm]

Abandoning in favour of VectorType::Builder :)

I may add methods similar to these there (and to the other builders) if needed.

mlir/include/mlir/IR/BuiltinTypes.td

let extraClassDeclaration = [{
     VectorType scaleElementBitwidth(unsigned scale);
 
     /// Returns if this type is ranked (always true).
     bool hasRank() const { return true; }
 
     /// Clone this vector type with the given shape and element type. If the
     /// provided shape is `std::nullopt`, the current shape of the type is used.
     VectorType cloneWith(std::optional<ArrayRef<int64_t>> shape,
                          Type elementType) const;
+
+    /// Return a new vector type which results from dropping the first \p n
+    /// dimensions of this vector type.
+    VectorType dropFrontDims(size_t n) const;
+
+    /// Return a new vector type which results from dropping the last \p n
+    /// dimensions of this vector type.
+    VectorType dropBackDims(size_t n) const;
+
+    /// Alias for dropFrontDims(1)
+    VectorType dropFrontDim() const { return dropFrontDims(1); }
+
+    /// Alias for dropBackDims(1)
+    VectorType dropBackDim() const { return dropBackDims(1); }
+
+    /// Return a new vector type which results from dropping the first \p n
+    /// dimensions and keeping \p m following dimensions.
+    VectorType sliceDims(size_t n, size_t m) const;
   }];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 #endif // BUILTIN_TYPES

mlir/lib/IR/BuiltinTypes.cpp

 }
 
 VectorType VectorType::cloneWith(std::optional<ArrayRef<int64_t>> shape,
                                  Type elementType) const {
   return VectorType::get(shape.value_or(getShape()), elementType,
                          getScalableDims());
 }
 
+VectorType VectorType::dropFrontDims(size_t n) const {
+  assert(size_t(getRank()) >= n && "Dropping more dims than exist");
+  return sliceDims(n, size_t(getRank()) - n);
+}
+
+VectorType VectorType::dropBackDims(size_t n) const {
+  assert(size_t(getRank()) >= n && "Dropping more dims than exist");
+  return sliceDims(0, size_t(getRank()) - n);
+}
+
+VectorType VectorType::sliceDims(size_t n, size_t m) const {
+  return VectorType::get(getShape().slice(n, m), getElementType(),
+                         getScalableDims().slice(n, m));
+}
+
 //===----------------------------------------------------------------------===//
 // TensorType
 //===----------------------------------------------------------------------===//
 
 Type TensorType::getElementType() const {
   return llvm::TypeSwitch<TensorType, Type>(*this)
       .Case<RankedTensorType, UnrankedTensorType>(
           [](auto type) { return type.getElementType(); });