[RFC] Changing base types for tensors and memrefs from C++ base classes to (type) interfaces

MLIR
1

This RFC is to discuss the current state of what “tensor” and “memref” are in MLIR, their strong coupling to ranked/unranked builtin types, and is to propose a change to the base types so that a user-defined tensor / memref could also be considered “tensor” / “memref” in MLIR’s upstream infrastructure (unless the infrastructure explicitly works with just a specific type).

Current state

There are 2 primary base classes, one for “tensor” and one for “memref”:

  • TensorType - is a base class of Ranked- / Unranked- TensorType
  • BaseMemRefType - is a base class of Ranked- / Unranked- MemRefType

They themselves are fairly simple:

  • Inherit mlir::Type
  • Attach SharedTypeInterface to derived classes
  • Provide a couple of functions (::classof() and static isValidElementType)

In the general sense, these are base classes - they declare something derived respectively as “tensor” or “memref”. The class hierarchy looks roughly the following way:
MyTensorType → TensorType → Type, ShapedType(Interface)
MyMemrefType → BaseMemRefType → Type, ShapedType(Interface)

Conveniently, one can avoid dealing with “my” tensor / memref directly by using the respective base class.

The problem

In reality, the hand-wavy class hierarchy above is not possible. In fact, this is due to TensorType and BaseMemRefType being “intrusive” types: they actually know what they can hold. That is,

inline bool BaseMemRefType::classof(Type type) {
  // thus: mlir::isa<BaseMemRefType> == mlir::isa<MemRefType, UnrankedMemRefType>
  return llvm::isa<MemRefType, UnrankedMemRefType>(type);
}

inline bool TensorType::classof(Type type) {
  // thus: mlir::isa<TensorType> == mlir::isa<RankedTensorType, UnrankedTensorType>
  return llvm::isa<RankedTensorType, UnrankedTensorType>(type);
}

(comments are mine, source code is roughly here)

Thus, in case one has a perfectly compatible tensor type that happens to be none of {RankedTensorType, UnrankedTensorType} then their tensor type is not considered a “tensor”. Well, at least when it comes down to mlir::isa<TensorType>(myTensor). Same reasoning applies as well to memrefs.

Worthy to mention that there’s an “extension” mechanism for tensors called encoding that allows one to store arbitrary information in a particular way into the builtin tensor. The memref’s counterpart to this would be layout. Yet, when one needs to have a custom tensor altogether (e.g. at least for the sake of parsing / printing / internal interactions), they may end up having unexpected issues. For instance, this is what we’ve faced in our downstream when trying to befriend one-shot bufferization and our own tensors / memrefs.

The proposal

The problem seems somewhat artificial. What exists are base classes that are way too knowledgeable about their derived counterparts. Luckily, there is a solution to this in the present day MLIR - type interfaces! The overall motivation for the change is (seems to be) similar to the motivation behind ShapedType changes.

Converting (base) TensorType and BaseMemRefType to type interfaces is almost a NFC, in fact, since the classes themselves do not implement any new APIs. Yet, being type interfaces, allows users to write custom tensor / memref types against these interfaces and ensure that the generic code relying on “tensor” / “memref” would still function.

The implementation PR: [mlir] Convert TensorType and BaseMemRefType to interfaces by andrey-golubev · Pull Request #133053 · llvm/llvm-project · GitHub

Potential problems

With base tensor / memref types becoming interfaces, any generic logic in MLIR is subject to potential new bugs and issues when non-builtin tensors / memrefs are involved.
Yet, in my opinion, this is fine: if a downstream project suffers with custom tensor / memref type support, they are free to suggest improvements (kind of what I do here now!).
I think overall the general infrastructure would also evolve to be better suited for various demands of downstream projects.

2

This is a perfectly normal setup for LLVM-style RTTI: How to set up LLVM-style RTTI for your class hierarchy — LLVM 21.0.0git documentation.

Could you please elaborate what specific passes you expect to work transparently on custom downstream memref or tensor types with no change? This sounds rather extreme as a change and if a downstream needs custom tensor type, I would encourage it to have custom everything, there are too many assumptions baked in about tensors.

1 Like
3

I won’t remember all the details anymore unfortunately (our downstream started looking at this about a year ago), but we’ve had issues primarily with one-shot bufferization. We wanted to be able to use general infrastructure with “my tensor” → “my memref” conversion working. There are some extension points (callbacks, etc.) that allow to specify custom conversion within the general framework but that wasn’t enough, unfortunately.
I believe the problems were around operation boundaries (when one has to insert to_tensor / to_memref, etc.) as well as in the general infrastructure since it seems to just “hardcode” the usage of BaseMemRefType (see completely random example I found just now).

So, in a sense, one-shot bufferization needs a concept of “tensor” and “memref” that does not mean exactly builtin tensor / memref. I figured that TensorType / BaseMemRefType are suitable exactly for this purpose.

Update: perhaps worth mentioning explicitly. I am not talking about tensor<...> or memref<...> - these remain “as is”: builtin types. The main interest is to promote their base classes to type interfaces.

4

I am looking at this from the perspective of a code cleanup.

Whether it’s useful/desirable to have custom memref/tensor types – I don’t know. Personally, I don’t have a use case for it.

What I like about the refactoring is that we can get rid of the hand-written C++ definitions of TensorType and BaseMemRefType. And instead use the “default” way of defining stuff in TableGen. I think this reduces complexity from the code base. Also in terms of LoC, it’s reduced by 60 lines if you exclude the additional testing code.

Also TensorType/BaseMemRefType already look like an interface to me: It’s an “abstract” type. In a sense that you are supposed to use a concrete implementation. I think nothing prevents users from defining their own types that inherit from TensorType/BaseMemRefType, apart from the hand-written BaseMemRefType::classof(Type type).

Overall I think this is a nice cleanup.

1 Like
5

@ftynse did you have time to look at my reply (and also at the reply from Matthias?). Are you strongly against the change? So far, apart from one-shot bufferization that definitely would benefit from a more “allowing” concept of tensors / memrefs, I see that XeGPU has its own “tensor” like type - at present it subclasses ::mlir::TensorType, similarly to builtin (un)ranked tensor types.

p.s.: sorry for the direct approach, but I would love to progress further with this.

6

It is perfectly normal for RFCs to stay up for O(weeks) and for other contributors to answer them at their preferred pace. We are not on-call here.

Yes, I am opposed to changes to the built-in type hierarchy with extremely thin motivation for upstream. The change may seem innocuous to you, but it has long-running conceptual implications that will be very difficult to revert (i.e., interfaces can be added to the type without changing the type, so one will be able to add “MemRefTypeInterface” to a pointer).

If we want bufferization specifically to work with non-builtin types, we should generalize bufferization.

1 Like
7

In general, I don’t really have any strong reasoning on “adding MemRefTypeInterface to a pointer” situation (honestly, I would be appalled if that happened), this is a valid point. My only comment would be: since both Tensor/MemRefType attach ShapedTypeInterface, the tensor/memref-wannabe has to also implement interface methods of that (e.g. getShape() / getElementType()) - due to transitivity - so such a type really has to be special?

If we want bufferization specifically to work with non-builtin types, we should generalize bufferization.

Yes. In which case, bufferization has to introduce its own “tensor” and “memref” concepts on top of pretty much ShapedType. I would assume these concepts would look pretty much like the base classes from builtins.

Thinking a bit more generically, what bufferization seems to need (and, to be fair, MLIR in general?) is a “shaped type with a tag” where that tag says “tensor” / “memref”. At least within bufferization logic this is the bare minimum change I can think of. Since MLIR already kind of has it (through the aforementioned base classes), it seems reasonable to just extend them to larger use.

Also, I vaguely remember sparse tensor discussion where the general premise was something along the lines of “it would be really cool to have sparse tensor as its own thing”. I believe having “tensor type interface” of sorts aids this as well? (As mentioned in the proposal, I believe that dealing with “tensors” without knowing what they are is a powerful concept). I guess this also boils down to what MLIR tries to be: so far I interpret it as a general framework and to me it seems to miss an abstraction around tensors / memrefs.

It is perfectly normal for RFCs to stay up for O(weeks) and for other contributors to answer them at their preferred pace. We are not on-call here.

Sorry!

8

Dwelling on the topic I would want to clarify something:
@ftynse would it be a reasonable compromise to instead introduce a “new tensor type interface” (name TBD) that is a type interface that exists as a separate entity from the current TensorType base class that users could attach to custom tensors? In a sense, the type hierarchy for tensors becomes:

  • builtin case: RankedTensor → TensorType (remains “as is”) → NewTensorTypeInterface → ShapedTypeInterface (or: TensorType is not attaching anything itself anymore, instead: RankedTensor → TensorType , NewTensorTypeInterface → ShapedTypeInterface)
  • custom case: MyTensor → NewTensorTypeInterface → ShapedTypeInterface

and ditto for memrefs.

This at least avoids the problem of breaking existing code anyhow? (Since “new tensor type interface” is not used anywhere). Trying to understand whether you are against modifying specifically TensorType / BaseMemRefType or against the idea in general.

So, generally, the job is this:

  1. Add new interfaces, use them in builtins
  2. Go through one-shot bufferization and replace TensorType / BaseMemRefType with their new counterparts
  3. Long-term, if new interfaces supersede the base classes, this would be the point at which we do the “jump” (e.g. TensorType/BaseMemRefType removed, only new type interfaces remain) yet this goes its own way
  4. New type interfaces also become something MLIR can advertise to users when they wish to have their custom tensors / memrefs (e.g. sparse, tiled)

Also summoning @mehdi_amini as he was the one who asked to start this RFC and I would love to capture his feedback here as well!

The reasoning: I actually find it rather complicated from the one-shot bufferization perspective to manage without a dedicated type interface: the logic there returns BaseMemRefType from (even free-standing) functions for example, stores memrefs internally, etc. Maybe I am missing some alternative design that LLVM has, but it seems that - for the sake of simplicity - we just need a type interface that could be used as an “abstraction”.
I also wouldn’t want it to be bufferization-specific in a sense (i.e. defined in bufferization utilities), since this prevents using such a thing in builtins? (there’s a possibility of dynamically attaching an interface to a type afair, but I find it to be more of a loophole due to lack of a better design?).

9

I’m in line with @ftynse feedback so far, so I didn’t feel the need to chime in. I actually asked for this RFC for this reason. While I’m not opposing to it completely, this isn’t just a simple refactoring or a “cleanup”. It opens up the entire type system in a new way and seems to me like a fairly fundamental change.

If the only motivation is the OneShot bufferization right now, can this be solved with a “config” struct defining callbacks used by OneShot bufferization? You need something to map a tensor type to the desired buffered equivalent anyway right?

10

I think one-shot bufferization would need a larger-scope change. For instance,

  • Mapping tensor to buffer assumes the output is BaseMemRefType. Without any new entities, the best we could do is replace memref with ShapedTypeInterface - but this breaks the “semantics”
  • Bufferization logic also relies on BaseMemRefType rather heavily (e.g. it forces the produced IR to hold base memref type)
  • Other parts of logic operate on TensorType / BaseMemRefType

So, overall, (one-shot) bufferization needs a concept “tensor” and “memref” (otherwise, the code is ambiguous as one cannot distinguish the one from the other). Fwiw, I wouldn’t have come up with the idea to change base classes to interfaces if everything worked out of the box with config callbacks.

(FYI: original old discussion - [mlir][One-Shot Bufferizer] Customizable tensor <-> memref conversions)

11

Right, but I don’t see how introducing a new type interface avoids it?

I agree, but this can be provided as callbacks.

1 Like
12

Assuming new type interface, replacing BaseMemRefType with “new memref type interface” and same for TensorType is sufficient? Then any type that implements “tensor” or “memref” interface could be stored / used “natively”. I guess there would be complications still as I only did this experiment with the aforementioned switch from TensorType → TensorTypeInterface.

Perhaps, but I guess this is where I’m missing how it could be done differently. I mean, technically, bufferization can operate on a mlir::Type and then have additional checks (e.g. asserts) to ensure correct semantics - is this what you’re (kind of) suggesting?
I’m sticking to type interfaces here because they seem to be introduced for this reason: abstract away implementation details via the API (as polymorphism primitive) and this automatically enforces “correct” semantics.

Having this around built-ins just seems friendlier? Anyhow, these interfaces could be placed in bufferization files if touching builtins is prohibited. (Then, maybe there’s a way for builtins and non-builtins to co-exist without too much boilerplate).

13

A type interface does not immediately allow you to create a new type: for a given tensor type, how do you know during bufferization what kind of “memref” to create?

Yes: isa<TensorType>(type) can easily be replaced by isaTensorTypeLikeCallbackFn(type) for example.

Right, but you’re trying to change the core builtin types for the sake of only bufferization, if this is controversial because of the widespread ramifications to consider, I’m just providing an alternative specific to the case of Bufferization to unblock this if this is an easier path.

14

Sorry, I didn’t explain it in too many details: most of the building blocks are already there - see BufferizationOptions. So one can already provide things like “allocate new buffer”, “map tensor to buffer”, etc. etc.
Except for some minor leftovers (e.g. one-shot bufferization currently uses buffer -> tensor roundtrip conversion that is not customizable - I have some ideas for it also, but it’s outside of this discussion), the existing logic is powerful enough.

Fair point. I kind of thought that it might actually be a good idea! :slight_smile:
(Again referring to sparse tensor in MLIR). Our downstream also implements own sparse tensor for example (also memref). So long-term, I thought that doing this initial uphill battle would pay off (also for upstream if sparse tensor is ever to become a thing - and it would be a thing if there’s a large enough corporate need for it :wink: ).

On a personal note, I think “like-tensor” / “like-memref” are neat concepts in bufferization. So I’d rather maybe have LikeTensorTypeInterface in bufferization utils with some hacky ::classof that also supports builtin tensors (and, as usual, ditto for memrefs).

Chaging:

  using FunctionArgTypeConverterFn =
      std::function<BaseMemRefType(TensorType, Attribute memorySpace,
                                   func::FuncOp, const BufferizationOptions &)>;

to:

  using FunctionArgTypeConverterFn =
      std::function<Type(Type, Attribute memorySpace,
                         func::FuncOp, const BufferizationOptions &)>;

looses a lot of semantics for almost nothing?
Incidentally, this also increases the probability of bugs: what if I forget that I need a buffer type and accidentally put a tensor? And if there’s a missing isaTensorTypeLikeCallbackFn(type) assert somewhere I’m doomed to debug this for a while. Having type interfaces just avoids this altogether.

15

The type interface could also be defined in the bufferization dialect. You can attach it to existing types via external model. That would make it clear that we are talking about tensor-like / buffer-like types for bufferization purposes. There could be interface methods on the type interfaces to:

  • build tensor ↔ buffer conversion ops (to_memref / to_tensor)
  • default implementation to build buffer copy op (memref.copy), which would only be used if not overridden in the bufferization options
  • default implementation to build buffer allocation op (memref.alloc), which would only be used if not overridden in the bufferization options

I haven’t fully thought this through yet. Not sure if it’s worth the overhead over just adding a few more lambdas to the bufferization options.

How about adding a type interface verifier that checks that only RankedTensorType / UnrankedTensorType implement the TensorType interface? That way we can have the cleanup without opening up the type system. It would also make clear in the code base that we thought about this and decided not to do that (for now).

1 Like
16

Yes, exactly. I am moving in this direction as well now (see the paragraph starting at “On a personal note, I think “like-tensor” / “like-memref”…” - well, except instead of ::classof we could attach interfaces dynamically, yes).
From the standpoint of “tensor / memref to store in a variable”, type interfaces seem an awesome simplification over callbacks acting as customization points.

For this part (tensor ↔ buffer conversion and other things) I don’t have a say. I think current design works and redefining it for the sake of something else that may or may not work is for now just overhead?
Again, current infrastructure in fact almost works out of the box. I have a simple test in MLIR that shows that it does and our downstream uses the current infrastructure successfully with custom tensors and memrefs (we have 2 distinct new types and also tensor with custom encoding / memref with custom layout).
The adjustments required were: a) some trickery with TensorType / BaseMemRefType (the proper design for which we’re discussing here); b) some revision of certain verification procedures (once I’m done with this and something is merged, I plan to move on with that other one).
Edit: actually, there were some additional adjustments required (i.e. the --one-shot-bufferize pass couldn’t be used directly, one need to change default options still) but this is outside of the scope of this RFC.

17

To summarize,

  1. Doing the original change: “convert TensorType / BaseMemRefType to type interfaces” - No [1]
  2. Updating bufferization dialect / utils to allow custom type usages - Yes, two options:
    a) Bufferization-specific type interfaces (can also be attached to builtins)
    b) More customization points (i.e. callbacks in options) + moving from BaseMemRefType to Type or some similar “higher-order” abstraction

Does this work?
Extra: Is there any strong preference for 2.a) over 2.b) or vice versa?

[1]: Should we consider Matthias’ proposal?

18

I’m still just not sure why you think that using an interface in this condition is a “cleanup” compared to the current situation?

19

As a general note, my take on upstream design/abstraction changes is that they must be beneficial for multiple use cases, preferably upstream, or across multiple visible downstreams. The deeper or closer to “core IR” is the change, the broader the benefit should be. Scoping down the change sort of lowers this breadth requirement. Another way to think about it is how one would explain why we have a certain abstraction upstream when we only have one concrete implementation thereof.

For the interfaces specifically, the tricky point is not implementing the ODS/C++ logic, but ironing out the English specification of the contract the interface defines. For the case of a hypothetical TensorLikeTypeInterface, we will have to define “what is a tensor” (including as opposed to vector, memref-style buffer). We have been grappling with this definition for a better half of a decade, I’m afraid. There is a concurrent thread on relaxing the vector type that hits this same definition problem… I will certainly welcome ideas around this definition though, but I have a feeling this RFC really does not try to think about that.

Normally, I’d say there is the overhead of requiring users not to forget attaching the interface prior to using the bufferization pass. But the pass already heavily relies on the external models, so also having type interface models wouldn’t make much of a difference IMO. On the upside, it allows for the more conventional and thus less error-prone API usage with isa and interface methods.

We would be back to what we already have but using MLIR-specific infrastructure instead of LLVM infrastructure. Is it really worth the effort?

2 Likes
20

I do not wish to wake the dead, especially those I have killed, but the “cleanup” portion of this is very similar to Named Constraints.

Wanting to provide a TableGen specified “base class” for a set of types that should not be extensible is exactly such a case where a named constraint beats an interface. So, I also think the currently proposed way of doing things isn’t really a cleanup.

But I do like where the thread is going. I would much rather have a specification of what a memref and a tensor are, although I would like those specs to not include the assumption that they have layouts that remain fixed :slight_smile: . If we enabled more users to plumb their dialects to bufferization, maybe we’d finally get around to ironing out the remaining weirdness and conservative copies.