[mlir][BuiltinTypes] Return VectorType from VectorType::Builder conversion operator

[benmxwl-arm]

Summary

0-D vectors are now supported, so the special case of returning the just
the element type can now be removed.

A few callers that relied on the old behaviour have been updated.

[nicolasvasilache]

thanks!

[awarzynski]

LGTM, thanks!

Btw, shouldn't this be updated too: https://github.com/llvm/llvm-project/blob/6862f0fab12176c7e4cdc095409deea7c26dde04/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L688-L690 ?

[benmxwl-arm]

In D159122#4626174, @awarzynski wrote:

LGTM, thanks!

Btw, shouldn't this be updated too: https://github.com/llvm/llvm-project/blob/6862f0fab12176c7e4cdc095409deea7c26dde04/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L688-L690 ?

I don't think it needs to be updated (does not look like that depends on rank 0 vectors decomposing to element types)

mlir/include/mlir/IR/BuiltinTypes.h

Builder &dropDim(unsigned pos) {
     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;
+  operator VectorType() {
     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;

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

struct Canonicalize0DShuffleOp : public OpRewritePattern<ShuffleOp> {
   LogicalResult matchAndRewrite(ShuffleOp shuffleOp,
                                 PatternRewriter &rewriter) const override {
     VectorType v1VectorType = shuffleOp.getV1VectorType();
     ArrayAttr mask = shuffleOp.getMask();
     if (v1VectorType.getRank() > 0)
       return failure();
     if (mask.size() != 1)
       return failure();
-    Type resType = VectorType::Builder(v1VectorType).setShape({1});
+    VectorType resType = VectorType::Builder(v1VectorType).setShape({1});
     if (llvm::cast<IntegerAttr>(mask[0]).getInt() == 0)
       rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
                                                        shuffleOp.getV1());
     else
       rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
                                                        shuffleOp.getV2());
     return success();
   }

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

 
 // Helper method to possibly drop a dimension in a load.
 // TODO
 static Value reshapeLoad(Location loc, Value val, VectorType type,
                          int64_t index, int64_t pos,
                          PatternRewriter &rewriter) {
   if (index == -1)
     return val;
-  Type lowType = VectorType::Builder(type).dropDim(0);
+  Type lowType = type.getRank() > 1 ? VectorType::Builder(type).dropDim(0)
+                                    : type.getElementType();
   // At extraction dimension?
   if (index == 0)
     return rewriter.create<vector::ExtractOp>(loc, lowType, val, pos);
   // Unroll leading dimensions.
   VectorType vType = cast<VectorType>(lowType);
-  Type resType = VectorType::Builder(type).dropDim(index);
-  auto resVectorType = cast<VectorType>(resType);
+  VectorType resType = VectorType::Builder(type).dropDim(index);
   Value result = rewriter.create<arith::ConstantOp>(
-      loc, resVectorType, rewriter.getZeroAttr(resVectorType));
-  for (int64_t d = 0, e = resVectorType.getDimSize(0); d < e; d++) {
+      loc, resType, rewriter.getZeroAttr(resType));
+  for (int64_t d = 0, e = resType.getDimSize(0); d < e; d++) {
     Value ext = rewriter.create<vector::ExtractOp>(loc, vType, val, d);
     Value load = reshapeLoad(loc, ext, vType, index - 1, pos, rewriter);
-    result =
-        rewriter.create<vector::InsertOp>(loc, resVectorType, load, result, d);
+    result = rewriter.create<vector::InsertOp>(loc, resType, load, result, d);
   }
   return result;
 }
 
 // Helper method to possibly drop a dimension in a store.
 // TODO
 static Value reshapeStore(Location loc, Value val, Value result,
                           VectorType type, int64_t index, int64_t pos,
                           PatternRewriter &rewriter) {
   // Unmodified?
   if (index == -1)
     return val;
   // At insertion dimension?
   if (index == 0)
     return rewriter.create<vector::InsertOp>(loc, type, val, result, pos);
   // Unroll leading dimensions.
-  Type lowType = VectorType::Builder(type).dropDim(0);
-  VectorType vType = cast<VectorType>(lowType);
-  Type insType = VectorType::Builder(vType).dropDim(0);
+  VectorType lowType = VectorType::Builder(type).dropDim(0);
+  Type insType = lowType.getRank() > 1 ? VectorType::Builder(lowType).dropDim(0)
+                                       : lowType.getElementType();
   for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
-    Value ext = rewriter.create<vector::ExtractOp>(loc, vType, result, d);
+    Value ext = rewriter.create<vector::ExtractOp>(loc, lowType, result, d);
     Value ins = rewriter.create<vector::ExtractOp>(loc, insType, val, d);
-    Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
+    Value sto = reshapeStore(loc, ins, ext, lowType, index - 1, pos, rewriter);
     result = rewriter.create<vector::InsertOp>(loc, type, sto, result, d);
   }
   return result;
 }
 
 /// Helper to create arithmetic operation associated with a kind of contraction.
 static std::optional<Value>
 createContractArithOp(Location loc, Value x, Value y, Value acc,