[mlir] Add VectorType::Builder to more easily build vector types from existing ones

[nicolasvasilache]

Add NFC.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.h

342

Is this right? Seems like the newShape SmallVector lifetime ends with this function, but the Builder keeps a reference to it.

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.h

342

Oh is it that this method returns a Type which forces to invoke the VectorType operator below?
So the returned object is OK but the Builder is basically busted?
This seems like a fragile API to me.

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.h

342

Yeah, this doesn't look ideal. Can we just create a copy of this builder and set the shape?

[nicolasvasilache]

ughh, sorry you're right I mess up the builder for nothing .. will send a followup shortly

mlir/include/mlir/IR/BuiltinTypes.h

Builder &setElementType(Type newElementType) {
     return *this;
   }
 
   Builder &setEncoding(Attribute newEncoding) {
     encoding = newEncoding;
     return *this;
   }
 
-  /// Create a new RankedTensorType by erasing a dim from shape.
-  RankedTensorType dropDim(unsigned dim) {
+  /// Create a new RankedTensor by erasing a dim from shape @pos.
+  RankedTensorType dropDim(unsigned pos) {
+    assert(pos < shape.size() && "overflow");
     SmallVector<int64_t, 4> newShape(shape.begin(), shape.end());
-    newShape.erase(newShape.begin() + dim);
+    newShape.erase(newShape.begin() + pos);
     return setShape(newShape);
   }
 
   operator RankedTensorType() {
     return RankedTensorType::get(shape, elementType, encoding);
   }
 
 private:
   ArrayRef<int64_t> shape;
   Type elementType;
   Attribute encoding;
 };
 
+//===----------------------------------------------------------------------===//
+// VectorType
+//===----------------------------------------------------------------------===//
+
+/// 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()) {}
+
+  /// Build from scratch.
+  Builder(ArrayRef<int64_t> shape, Type elementType)
+      : shape(shape), elementType(elementType) {}
+
+  Builder &setShape(ArrayRef<int64_t> newShape) {
+    shape = newShape;
+    return *this;
+  }
+
+  Builder &setElementType(Type newElementType) {
+    elementType = newElementType;
+    return *this;
+  }
+
+  /// Create a new VectorType by erasing a dim from shape @pos.
+  /// 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.
+  Type dropDim(unsigned pos) {
+    assert(pos < shape.size() && "overflow");
+    if (shape.size() == 1)
+      return elementType;
+    SmallVector<int64_t, 4> newShape(shape.begin(), shape.end());
+    newShape.erase(newShape.begin() + pos);
+    return setShape(newShape);
+  }
+
+  operator VectorType() { return VectorType::get(shape, elementType); }
+
+private:
+  ArrayRef<int64_t> shape;
+  Type elementType;
+};
+
 /// 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
 /// `reducedShape`. The returned mask can be applied as a projection to
 /// `originalShape` to obtain the `reducedShape`. This mask is useful to track
 /// which dimensions must be kept when e.g. compute MemRef strides under
 /// rank-reducing operations. Return None if reducedShape cannot be obtained
 /// by dropping only `1` entries in `originalShape`.

mlir/include/mlir/IR/BuiltinTypes.td

def Builtin_Vector : Builtin_Type<"Vector", [
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape, "Type":$elementType
     ), [{
       return $_get(elementType.getContext(), shape, elementType);
     }]>
   ];
   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 of 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 t.isa<IntegerType, IndexType, FloatType>();
     }
 
     /// Get or create a new VectorType with the same shape as `this` and an

mlir/lib/Dialect/Vector/VectorOps.cpp

if (masksInfo.empty())
     return success();
   if (masksInfo.size() != 2)
     return parser.emitError(parser.getNameLoc(),
                             "expected zero or exactly 2 vector mask operands");
   auto lhsType = types[0].cast<VectorType>();
   auto rhsType = types[1].cast<VectorType>();
   auto maskElementType = parser.getBuilder().getI1Type();
   std::array<Type, 2> maskTypes = {
-      VectorType::get(lhsType.getShape(), maskElementType),
-      VectorType::get(rhsType.getShape(), maskElementType)};
+      VectorType::Builder(lhsType).setElementType(maskElementType),
+      VectorType::Builder(rhsType).setElementType(maskElementType)};
   if (parser.resolveOperands(masksInfo, maskTypes, loc, result.operands))
     return failure();
   return success();
 }
 
 static void print(OpAsmPrinter &p, ContractionOp op) {
   // TODO: Unify printing code with linalg ops.
   auto attrNames = op.getTraitAttrNames();

mlir/lib/Dialect/Vector/VectorTransforms.cpp

for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
     // Re-insert remaining indices, but renamed when occurring
     // after the removed index.
     auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
     results.push_back(targetExpr);
   }
   return AffineMap::get(map.getNumDims() - 1, 0, results, ctx);
 }
 
-// Helper to drop dimension from vector type.
-static Type adjustType(VectorType tp, int64_t index) {
-  int64_t rank = tp.getRank();
-  Type eltType = tp.getElementType();
-  if (rank == 1) {
-    assert(index == 0 && "index for scalar result out of bounds");
-    return eltType;
-  }
-  SmallVector<int64_t, 4> adjustedShape;
-  for (int64_t i = 0; i < rank; ++i) {
-    // Omit dimension at the given index.
-    if (i == index)
-      continue;
-    // Otherwise, add dimension back.
-    adjustedShape.push_back(tp.getDimSize(i));
-  }
-  return VectorType::get(adjustedShape, eltType);
-}
-
 // 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 = adjustType(type, 0);
+  Type lowType = VectorType::Builder(type).dropDim(0);
   // At extraction dimension?
   if (index == 0) {
     auto posAttr = rewriter.getI64ArrayAttr(pos);
     return rewriter.create<vector::ExtractOp>(loc, lowType, val, posAttr);
   }
   // Unroll leading dimensions.
   VectorType vType = lowType.cast<VectorType>();
-  VectorType resType = adjustType(type, index).cast<VectorType>();
+  auto resType = VectorType::Builder(type).dropDim(index).cast<VectorType>();
   Value result = rewriter.create<arith::ConstantOp>(
       loc, resType, rewriter.getZeroAttr(resType));
   for (int64_t d = 0, e = resType.getDimSize(0); d < e; d++) {
     auto posAttr = rewriter.getI64ArrayAttr(d);
     Value ext = rewriter.create<vector::ExtractOp>(loc, vType, val, posAttr);
     Value load = reshapeLoad(loc, ext, vType, index - 1, pos, rewriter);
     result =
         rewriter.create<vector::InsertOp>(loc, resType, load, result, posAttr);
static Value reshapeStore(Location loc, Value val, Value result,
   if (index == -1)
     return val;
   // At insertion dimension?
   if (index == 0) {
     auto posAttr = rewriter.getI64ArrayAttr(pos);
     return rewriter.create<vector::InsertOp>(loc, type, val, result, posAttr);
   }
   // Unroll leading dimensions.
-  Type lowType = adjustType(type, 0);
+  Type lowType = VectorType::Builder(type).dropDim(0);
   VectorType vType = lowType.cast<VectorType>();
-  Type insType = adjustType(vType, 0);
+  Type insType = VectorType::Builder(vType).dropDim(0);
   for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
     auto posAttr = rewriter.getI64ArrayAttr(d);
     Value ext = rewriter.create<vector::ExtractOp>(loc, vType, result, posAttr);
     Value ins = rewriter.create<vector::ExtractOp>(loc, insType, val, posAttr);
     Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
     result = rewriter.create<vector::InsertOp>(loc, type, sto, result, posAttr);
   }
   return result;