Reading attribute stored in Relay IR at scheduling time

Questions
1

Hello,

Here I am, reaching out to the amazing TVM community with my second question.

I have been looking at ways to read attributes stored in Relay IR during scheduling time (e.g. when tiling )

This is an example of what I mean with Relay IR attribute:

%1 = fn (%p02: Tensor[(1, 3, 1080, 1920), int8], %p1: Tensor[(16, 3, 3, 3), int8], myAttribute="valueToRead") -> Tensor[(1, 16, 1080, 1920), float32] { ... }

The attribute I would like to read is called “myAttribute” and I am trying to retrieve its value when I am scheduling that particular function.

I couldn’t find a way to access that attribute at scheduling time. I would like to access that value without having to change any of the internals of TVM.

Is that possible? Do you have any pointers on how to do this?

2

As a quick summary, there are 3 levels of ‘IR’ when lowering in TVM: Relay, TE (Tensor Expression language) and TIR. The scheduling language is currently implemented to work on TE, so we must lower Relay -> TE in order to start scheduling. Unfortunately, that lowering step is part of the internals of TVM and it’s not particularly modular, so it isn’t simple to pass an attribute through.

Could you elaborate a bit on your use case? If you want to change the way things are tiled then AutoTVM might provide an alternative route.

3

Thanks @mbaret for your quick summary and explanation!

There are some custom attributes attached to the Relay IR that would be useful for me to access in TE and/or TIR. I have written some custom TIR passes that transform loops in different ways (e.g. one transformation is loop peeling). The information contained in those Relay IR attributes would be helpful to make more informed choices during those transformations.

4

Just quickly with respect to loop peeling, in case you haven’t seen it already, TIR has a loop partitioning pass which kind of does this, so it may be worth a look.

The way I’ve done something similar to what you want involved creating custom Relay operators and associating special TE computes with them that I could manually attach Relay information to. This is pretty contrived though, so I wouldn’t necessarily recommend it.

5

I took a look at the TIR loop partitioning pass, thanks for the pointer!

Does your solution involve changing the internals of TVM?

6

Not exactly changing them, but replicating them (essentially a custom compile_engine). I’m using it for more than just passing custom attributes across though, so it sounds like your use case may not be enough to motivate that large of a change.