[mlir] Fix unintentional mutation by VectorType/RankedTensorType::Builder dropDim

[jpienaar]

mlir/include/mlir/IR/BuiltinTypes.h

286–287

Why not have this return *this still? dropDim is different than all the others here.

[nicolasvasilache]

mlir/include/mlir/IR/BuiltinTypes.h

286–287

Because dropping a dim would break ArrayRef contiguity requirements.

[mehdi_amini]

This LG, but this does not entirely address the quirks of this builder: we have a "mostly flex" API but not for the "dropDim" which is somehow always "terminal" in the chain? It also means you can only drop a single dimension really.

[nicolasvasilache]

In D113933#3133765, @mehdi_amini wrote:

This LG, but this does not entirely address the quirks of this builder: we have a "mostly flex" API but not for the "dropDim" which is somehow always "terminal" in the chain? It also means you can only drop a single dimension really.

Since it is non-owning, it has to be terminal unless we change the underlying design of MemRefType::Builder, VectorType::Builder and TensorType::Builder consistently to all take an owning SmallVector<int64_t> shape which I don't think is a good idea. I wouldn't go as far as creating a new Builder class just for this purpose, this would be more confusing. Maybe a better naming scheme would help?

Yes, atm this only drops a single dimension, I have not had needs for more than this so far. As usual with MLIR we can extend when the need arises.
Generally cleaning up all the verbose multiline type creations in MLIR is a cleanup step that I'd like to turn into intro tasks.

[mehdi_amini]

I don't know the history of all these Builders actually, it just gives me a feel of a fairly "ad-hoc" API that makes it a bit specialized for a subset of the users, whereas I'd have thought of APIs exposed in mlir/include/mlir/IR/BuiltinTypes.h to be a bit more generic.
I wonder how @rriddle would see all this?

[nicolasvasilache]

MemRefType::Builder was introduced here: https://reviews.llvm.org/D73296

[ftynse]

In D113933#3133770, @nicolasvasilache wrote:

In D113933#3133765, @mehdi_amini wrote:

This LG, but this does not entirely address the quirks of this builder: we have a "mostly flex" API but not for the "dropDim" which is somehow always "terminal" in the chain? It also means you can only drop a single dimension really.

Since it is non-owning, it has to be terminal unless we change the underlying design of MemRefType::Builder, VectorType::Builder and TensorType::Builder consistently to all take an owning SmallVector<int64_t> shape which I don't think is a good idea. I wouldn't go as far as creating a new Builder class just for this purpose, this would be more confusing.

I don't see a problem of having an owning vector there, especially the SmallVector that keeps ~6 elements on stack (IIRC, it's 64 bytes total storage, and 16 are used for size and capacity). If we were using the regular type construction API, we would have used SmallVector anyway.

Maybe a better naming scheme would help?

We can also go Java-style and add an explicit .build() method that constructs the type, and have the .dropDimensionsAndBuild() counterpart.

In D113933#3133792, @mehdi_amini wrote:

I don't know the history of all these Builders actually, it just gives me a feel of a fairly "ad-hoc" API that makes it a bit specialized for a subset of the users, whereas I'd have thought of APIs exposed in mlir/include/mlir/IR/BuiltinTypes.h to be a bit more generic.
I wonder how @rriddle would see all this?

The original MemRefBuilder was added to remove the boilerplate of creating memref type that is exactly the same as another memref except for one property (address space or layout), abundant in conversions. It looks much cleaner to write MemRefType::Builder(original).setMemorySpace(42) than to write MemRefType::get(original.getContext(), original.getShape(), original.getLayout(), /*memorySpace=*/42). Both ways are equally generic and equally BuiltinTypes.h.

[mehdi_amini]

Both ways are equally generic

I agree, but only in absence of terminal APIs (ones that don't return *this) in these fluent style Builder classes, because they are non-composable. These APIs makes it look to me like they are motivated by specific call sites / uses, and hence are not really generic anymore.
Having something like MemRefType::Builder(original).setMemorySpace(42) which returns a builder and rely on the conversion operator to finalize (or with an extra build() or finalize() method as you suggested) would seems fine to me on the other hand.

There is some cost as Nicolas mentioned: you always have to copy the shape into a SmallVector when creating the builder instead of copying only if the user is using dropDim. I don't know if this is something to optimize for here though (memcpy of a few int64 to the stack), we may be able to do a copy-on-write as well, but that may not be trivial / worth it.

[rriddle]

In D113933#3136339, @mehdi_amini wrote:

Both ways are equally generic

I agree, but only in absence of terminal APIs (ones that don't return *this) in these fluent style Builder classes, because they are non-composable. These APIs makes it look to me like they are motivated by specific call sites / uses, and hence are not really generic anymore.
Having something like MemRefType::Builder(original).setMemorySpace(42) which returns a builder and rely on the conversion operator to finalize (or with an extra build() or finalize() method as you suggested) would seems fine to me on the other hand.

There is some cost as Nicolas mentioned: you always have to copy the shape into a SmallVector when creating the builder instead of copying only if the user is using dropDim. I don't know if this is something to optimize for here though (memcpy of a few int64 to the stack), we may be able to do a copy-on-write as well, but that may not be trivial / worth it.

Can't we do copy-on-write for the shape? I don't see why we need to copy the shape if we haven't modified it. These builders are short lived on the stack anyways, so I don't see a big deal about making the builder slightly bigger.

[nicolasvasilache]

How would copy-on-write be implemented?
More specifically who has the ownership of the copied vector and guarantees that it is live until the builder dies?

Atm the vector is local to dropDim and I don't have a good way to let it escape.
Is there maybe a specific CopyOnWriteArrayRef<int64_t> shape; in LLVM I could use in place of the ArrayRef<int64_t> shape member in the Builder ?

[mehdi_amini]

In D113933#3137055, @nicolasvasilache wrote:

How would copy-on-write be implemented?

With the builder having a SmallVector member *and* an ArrayRef.

SmallVector<int64_t> storage;
ArrayRef<int64_t> shape;

Type dropDim(unsigned pos) const {
  if (storage.empty()) storage = shape;
  assert(pos < shape.size() && "overflow");
  storage.erase(storage.begin() + pos);
  shape = {storage.data(), storage.size()};
  return *this;
}

[nicolasvasilache]

Implement CoW + rebase.

[nicolasvasilache]

Fix bug: assign -> append

mlir/include/mlir/IR/BuiltinTypes.h

Builder &setElementType(Type newElementType) {
     return *this;
   }
 
   Builder &setEncoding(Attribute newEncoding) {
     encoding = newEncoding;
     return *this;
   }
 
-  /// Create a new RankedTensor by erasing a dim from shape @pos.
-  RankedTensorType dropDim(unsigned pos) {
+  /// Erase a dim from shape @pos.
+  Builder &dropDim(unsigned pos) {
     assert(pos < shape.size() && "overflow");
-    SmallVector<int64_t, 4> newShape(shape.begin(), shape.end());
-    newShape.erase(newShape.begin() + pos);
-    return setShape(newShape);
+    if (storage.empty())
+      storage.append(shape.begin(), shape.end());
+    storage.erase(storage.begin() + pos);
+    shape = {storage.data(), storage.size()};
+    return *this;
   }
 
   operator RankedTensorType() {
     return RankedTensorType::get(shape, elementType, encoding);
   }
 
 private:
   ArrayRef<int64_t> shape;
+  // Owning shape data for copy-on-write operations.
+  SmallVector<int64_t> storage;
   Type elementType;
   Attribute encoding;
 };
 
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
Builder &setShape(ArrayRef<int64_t> newShape) {
     return *this;
   }
 
   Builder &setElementType(Type newElementType) {
     elementType = newElementType;
     return *this;
   }
 
-  /// Create a new VectorType by erasing a dim from shape @pos.
+  /// Erase a dim from shape @pos.
+  Builder &dropDim(unsigned pos) {
+    assert(pos < shape.size() && "overflow");
+    if (storage.empty())
+      storage.append(shape.begin(), shape.end());
+    storage.erase(storage.begin() + pos);
+    shape = {storage.data(), storage.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.
-  Type dropDim(unsigned pos) {
-    assert(pos < shape.size() && "overflow");
-    if (shape.size() == 1)
+  operator Type() {
+    if (shape.empty())
       return elementType;
-    SmallVector<int64_t, 4> newShape(shape.begin(), shape.end());
-    newShape.erase(newShape.begin() + pos);
-    return setShape(newShape);
+    return VectorType::get(shape, elementType);
   }
 
-  operator VectorType() { return VectorType::get(shape, elementType); }
-
 private:
   ArrayRef<int64_t> shape;
+  // Owning shape data for copy-on-write operations.
+  SmallVector<int64_t> storage;
   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

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

LogicalResult matchAndRewrite(linalg::Conv2DNhwcHwcfOp convOp,
     bool removeH = (fhSize == 1 && ohSize == 1);
     bool removeW = (fwSize == 1 && owSize == 1);
     if (!removeH && !removeW)
       return failure();
 
     // Get new shapes and types for all operands by removing the size-1
     // dimension.
     using RTTBuilder = RankedTensorType::Builder;
-    auto newInputType = RTTBuilder(inputType).dropDim((removeH ? 1 : 2));
-    auto newFilterType = RTTBuilder(filterType).dropDim((removeH ? 0 : 1));
-    auto newOutputType = RTTBuilder(outputType).dropDim(removeH ? 1 : 2);
+    RankedTensorType newInputType =
+        RTTBuilder(inputType).dropDim((removeH ? 1 : 2));
+    RankedTensorType newFilterType =
+        RTTBuilder(filterType).dropDim((removeH ? 0 : 1));
+    RankedTensorType newOutputType =
+        RTTBuilder(outputType).dropDim(removeH ? 1 : 2);
 
     // Rank-reduce operands.
     Location loc = convOp.getLoc();
     Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(
         rewriter, loc, input, newInputType);
     Value newFilter = tensor::createCanonicalRankReducingExtractSliceOp(
         rewriter, loc, filter, newFilterType);
     Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(
LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
     bool removeH = (khSize == 1 && ohSize == 1);
     bool removeW = (kwSize == 1 && owSize == 1);
     if (!removeH && !removeW)
       return failure();
 
     // Get new shapes and types for all operands by removing the size-1
     // dimension.
     using RTTBuilder = RankedTensorType::Builder;
-    auto newInputType = RTTBuilder(inputType).dropDim((removeH ? 1 : 2));
-    auto newKernelType = RTTBuilder(kernelType).dropDim((removeH ? 0 : 1));
-    auto newOutputType = RTTBuilder(outputType).dropDim(removeH ? 1 : 2);
+    RankedTensorType newInputType =
+        RTTBuilder(inputType).dropDim((removeH ? 1 : 2));
+    RankedTensorType newKernelType =
+        RTTBuilder(kernelType).dropDim((removeH ? 0 : 1));
+    RankedTensorType newOutputType =
+        RTTBuilder(outputType).dropDim(removeH ? 1 : 2);
 
     // Rank-reduce operands.
     Location loc = convOp.getLoc();
     Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(
         rewriter, loc, input, newInputType);
     Value newKernel = tensor::createCanonicalRankReducingExtractSliceOp(
         rewriter, loc, kernel, newKernelType);
     Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(

mlir/lib/Dialect/Vector/VectorTransforms.cpp

static Value reshapeLoad(Location loc, Value val, VectorType type,
   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>();
-  auto resType = VectorType::Builder(type).dropDim(index).cast<VectorType>();
+  Type resType = VectorType::Builder(type).dropDim(index);
+  auto resVectorType = resType.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++) {
+      loc, resVectorType, rewriter.getZeroAttr(resVectorType));
+  for (int64_t d = 0, e = resVectorType.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);
+    result = rewriter.create<vector::InsertOp>(loc, resVectorType, load, result,
+                                               posAttr);
   }
   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,