[Relax] Best way to implement small graph transformations

Hi everyone,

I’m facing the following simple need at the Relax graph level: I’d like to transform, in a given Relax module and only for its main function, a specific sequence of ops into another sequence of ops.

Things like “if op1 is followed by op2, then only keep op2”, or “if op1 is followed by op2, then remove both them completely”. So the notion here of “pattern” is very simple, as it would be based only on the name of the ops.

For this specific needs, it won’t be things much more complicated than that. Of course, we could generalize it to deal with “given a pattern (expressed as a graph), transform the original module/function by applying a given functor to any matching sub-graph”.

Before I start to write something, either general or specific for our needs, do you know if anything in the TVM project could be reused easily for doing that (or if anything does that already, reuse is best!). I saw for instance the files dataflow_matcher.h/cc, dataflow_pattern.h/cc and dataflow_pattern_functor.h/cc. Would you advise to use that? Is there any good example of how it’s being used?

Also, we could consider writing that at the C++ or Python level. Would you have any advice for this kind of need, regarding where to implement that?

Any idea, advise or pointer is greatly appreciated, as I’d like to keep things well organized, so if something has already been designed for this kind of use, then I’d want to reuse that for sure.

Many thanks and Happy Thanksgiving to you all!

The use case sounds like a job for the relax pattern matcher and so I would think that’s the way to go.

As for the question about python vs C++, I think either way should be fine. The way I normally think is that if the pass is going to be added as default to all targets, C++ is a better choice, otherwise either one is fine (programmer choice). Others can probably correct me on this as this is mostly just what I follow on my own.

Thanks for the response @sanirudh ! :slight_smile: