Page MenuHomePhabricator

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

Authored by nicolasvasilache on Mon, Nov 15, 12:33 PM.

Diff Detail

Event Timeline

nicolasvasilache requested review of this revision.Mon, Nov 15, 12:33 PM
Herald added a project: Restricted Project. · View Herald TranscriptMon, Nov 15, 12:33 PM
jpienaar added inline comments.
mlir/include/mlir/IR/BuiltinTypes.h
286–287

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

nicolasvasilache marked an inline comment as done.Mon, Nov 15, 12:46 PM
nicolasvasilache added inline comments.
mlir/include/mlir/IR/BuiltinTypes.h
286–287

Because dropping a dim would break ArrayRef contiguity requirements.

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 marked an inline comment as done.Mon, Nov 15, 10:49 PM

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.

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?

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

ftynse added a subscriber: ftynse.Tue, Nov 16, 12:42 AM

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.

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.

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.

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 added a comment.EditedWed, Nov 17, 1:53 AM

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 ?

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;
}

Implement CoW + rebase.

gysit accepted this revision.Mon, Nov 22, 2:39 AM
This revision is now accepted and ready to land.Mon, Nov 22, 2:39 AM

Fix bug: assign -> append

This revision was landed with ongoing or failed builds.Mon, Nov 22, 2:55 AM
This revision was automatically updated to reflect the committed changes.