Custom Op base interfaces (or, adding methods to ops via traits, or ...)

[Ben Vanik]

I'd like to be able to add methods to ops in my dialect. I can do this today by defining a custom OpTrait with my method and wiring it up via tblgen NativeOpTrait. It's great because I can compose the traits really nicely but is limited because there's no way to dynamically detect if a generic op has a trait and no way to then dispatch the method.

This was working, but now I'm finding myself needing the ability to get type-erased access to these methods similar to what AbstractOperation provides: I'm writing serialization code that would like to generically process Operation* pointers. I'm having trouble with that as there's no vtable/no isa/etc to help make this work.

I found Nicolas' code here:

https://github.com/tensorflow/mlir/blob/master/include/mlir/Linalg/IR/LinalgOps.h#L305

It seems to do exactly what I want but is a decent amount of boilerplate that I'd like to avoid copy/pasting. I remember there being some mention of another approach at a previous team meeting but forgot what the outcome was.

Does anyone have any pointers/tips/etc? I'm also open to a "you're holding it wrong" answer - I may be trying to apply some OOP to where there shouldn't be any :)

Thanks!

[Mehdi AMINI]

On Wed, May 22, 2019 at 3:37 PM 'Ben Vanik' via MLIR <ml...@tensorflow.org> wrote:

I'd like to be able to add methods to ops in my dialect. I can do this today by defining a custom OpTrait with my method and wiring it up via tblgen NativeOpTrait. It's great because I can compose the traits really nicely but is limited because there's no way to dynamically detect if a generic op has a trait and no way to then dispatch the method.

This was working, but now I'm finding myself needing the ability to get type-erased access to these methods similar to what AbstractOperation provides: I'm writing serialization code that would like to generically process Operation* pointers. I'm having trouble with that as there's no vtable/no isa/etc to help make this work.

I found Nicolas' code here:

https://github.com/tensorflow/mlir/blob/master/include/mlir/Linalg/IR/LinalgOps.h#L305

It seems to do exactly what I want but is a decent amount of boilerplate that I'd like to avoid copy/pasting. I remember there being some mention of another approach at a previous team meeting but forgot what the outcome was.

The approach I displayed during the team meeting is exactly what Nicolas ended up adopting in Linalg

There is some boilerplate involved initially, but once the structure is set-up, adding a new Op does not require changing anything (TableGen add it to the list).

Adding a new method to dispatch requires to add 4 lines (1 declaration in the concept and 3 in the model to forward).

I haven't found a better solution: I tried 3 different approaches, and there always was a similar amount of boilerplate involved.

(actually I know one area of improvement: the current implementation has one extra pointer dereferencing than strictly necessary, but the only way I could see to avoid it is "inlining" the vtable and it would require more boilerplate and/or template meta-programming)

Someone may have a better idea though!

-- 

Mehdi

 

[Ben Vanik]

Fantastic! Thanks for the response!

I'll move forward with this approach!

[Lei Zhang]

I remembered that one possible way discussed is to have another embedded word in AbstractOperation for dialect-specific properties/traits like what we have right now there as opProperties. Then dialects can set properties/traits into this word and probe them when necessarily. If we find common properties across dialects, we can the promote it into the "core" one to make it generally available. Possibly you are mentioning this?

I'm not exactly sure about your use case; but is it possible to have a special field in op definition spec as the marker and then have a TableGen backend to generate code according to it? I'm looking into having TableGen to generate some visitor code to alleviate the boilerplate in general. Another workaround to unblock the progress I think might be to abuse attributes. You can set special marker unit attributes on operations and then query them. Not sure it is the appropriate way to go in the long-run though.

Thanks,

Lei

--
You received this message because you are subscribed to the Google Groups "MLIR" group.
To unsubscribe from this group and stop receiving emails from it, send an email to mlir+uns...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/CAOB_hFS8o%3DGOdwtVzOhnBcYmKNgoAns-yd9KNi7p0b97YsJ9oQ%40mail.gmail.com.

[Ben Vanik]

Ahh yeah that sounds similar to what I was trying to do at first. To make serialization code cleaner and keep definitions with the opdefs I wanted to add traits that denoted how each op was to be serialized. I tried using extraClassDeclaration to add methods when required but naturally there's no way to dynamically call those. This may be relevant to SPIR-V and other formats that would be nice to serialize directly from MLIR without needing a separate serialization table/system.

To work around things I made a string map of operation name to custom std::function handling the logic. It's nasty but is simple to understand for the moment. I'd like to have something with more compile-time safety like the method Mehdi came up with but want to try to understand what I want before I bake it out like that :) If tablegen could generate all the Concept and nested interface stuff for me based on declared methods that'd be awesome, too.

Visitor code would be nice - especially as things get more involved with regions!

[Lei Zhang]

On Thu, May 23, 2019 at 11:44 AM Ben Vanik <benv...@google.com> wrote:

Ahh yeah that sounds similar to what I was trying to do at first. To make serialization code cleaner and keep definitions with the opdefs I wanted to add traits that denoted how each op was to be serialized. I tried using extraClassDeclaration to add methods when required but naturally there's no way to dynamically call those. This may be relevant to SPIR-V and other formats that would be nice to serialize directly from MLIR without needing a separate serialization table/system.

Yep, agreed. 

[Chris Lattner]

On May 23, 2019, at 8:44 AM, 'Ben Vanik' via MLIR <ml...@tensorflow.org> wrote:
> To work around things I made a string map of operation name to custom std::function handling the logic. It's nasty but is simple to understand for the moment. I'd like to have something with more compile-time safety like the method Mehdi came up with but want to try to understand what I want before I bake it out like that :) If tablegen could generate all the Concept and nested interface stuff for me based on declared methods that'd be awesome, too.
>
> Visitor code would be nice - especially as things get more involved with regions!

I’d really love to see tblgen start producing visitors for dialects that want them. Setting a “I want a visitor” field on the dialect class in tblgen could enable generation.

-Chris

[Lei Zhang]

Sure thing, I will work on it next week! :)