Relay to relax translator details

[sanirudh]

Hi,

I’ve been playing around with relax infrastructure for a while and I wanted to ask a couple questions related to the relay to relax translator and its design/future directions.

1. First of all, I see that the relay_translator.py is in the tvm/relax/testing directory. Is this because the relay_translator is a temporary solution until relax brings its own frontend converters and adopts most/if not all the relay optimizations.

2. If its not meant for testing, is there a plan to make it a core part of relax and port it to a c++ implementation since relax is planned as an optional module for now and might stay that way for a while.

One last related doubt. I see that the PassPrefix passes are run before translating to relax, but FuseOps is not run. Does that mean there’s a separate fusion pass in relax and if so, is there a plan to only work with relax fusion.

Thanks,
Anirudh

[Hzfengsy]

Thanks @sanirudh for your interest.

@sanirudh

Is this because the relay_translator is a temporary solution until relax brings its own frontend converters and adopts most/if not all the relay optimizations.

You are right! relay_translator is designed to only be a testing solution in a long-term view. Relax is going to have its own front-end importer. There is already an FX importer for PyTorch at the MLC fork and will upstream later.

@sanirudh

is there a plan to make it a core part of relax and port it to a c++ implementation since relax is planned as an optional module for now and might stay that way for a while.

Although I think I’ve answered the question that it’s only for testing and a temporary solution, I would ask another question: what’s the causation between porting it to C++ and relay being an optional module?

@sanirudh

Does that mean there’s a separate fusion pass in relax and if so, is there a plan to only work with relax fusion.

Yes, we have FuseOps and FuseTIR passes in Relax, as we would not plan to use Relay’s passes in a long-time view.

[sanirudh]

Thanks a lot for the detailed answer @Hzfengsy.

There is already an FX importer for PyTorch at the MLC fork and will upstream later.

A quick follow-up question here, are there other frontends under development. We heavily depend on onnx frontend, so just wanted to ask if there already is development work going on for the onnx frontend and if not perhaps we could try to help out here as it might be of major interest to us.

what’s the causation between porting it to C++ and relay being an optional module?

I don’t think there’s any specific connection between the C++ port and relax being optional. I should have phrased that question differently, sorry about that.

My major interest was that, since there are quite a few things (like frontends, op support, etc.) that need improvements in Relax, and it might be a while to get there, we would need to work with a combination of Relay + Relax to get performance till we reach a point where we don’t need relay.

So the question about a c++ implementation of the relay translator should have been about the stability of the relay translator until we reach a point where we can go directly to Relax.

Yes, we have FuseOps and FuseTIR passes in Relax, as we would not plan to use Relay’s passes in a long-time view.

The question about Relay vs Relax fusion was because we need to use Relay frontend right now and then relay translator, but relay translator does not generate Relax ops (it only generates call_tir calls to lowered primfuncs), so how would fusion go in this workflow?

IIUC, Relax FuseOps only works on Relax operators, not the lowered TIR PrimFuncs. I went through Extract Op Info from Primfunc, and was curious if there’s any plan to do some sort of fusion directly on prim_funcs based on tir::PatternKindAnalyzer?

[Hzfengsy]

@sanirudh

so just wanted to ask if there already is development work going on for the onnx frontend and if not perhaps we could try to help out here as it might be of major interest to us.

I’m sorry that we don’t have enough bandwidth to support it. I faithfully appreciate it if you could help.

@sanirudh

So the question about a c++ implementation of the relay translator should have been about the stability of the relay translator until we reach a point where we can go directly to Relax.

You are right. It needs a while to get the full support of frontends and ops. It’s a good thing if we can move the translator to C++ and I don’t think there is any technical blocker. However, the Relax module that is translated from Relay may miss some features (e.g. symbolic shape deduction)

@sanirudh

IIUC, Relax FuseOps only works on Relax operators, not the lowered TIR PrimFuncs.

Relax FuseOps is based on Graph-TIR module. In other words, it analysis the fusion patterns of the TIR functions instead of operators. That’s why we can do fusion after the relay translator.

[tqchen]

@jwfromm have some efforts in prototyping onnx support and perhaps he can chime in a bit here

[sanirudh]

Thanks for the reply.

However, the Relax module that is translated from Relay may miss some features (e.g. symbolic shape deduction)

Yes that makes sense, and I agree, for symbolic shape support, the relay translator workflow might not be the best choice.

Relax FuseOps is based on Graph-TIR module. In other words, it analysis the fusion patterns of the TIR functions instead of operators. That’s why we can do fusion after the relay translator.

I did not realize that by just glancing over the pass as it was just checking for OpPatternKind attribute similar to the Relay fusion pass. I’ll try out the relax FuseOps, thanks for this update.

[sanirudh]

Great. Would love to know some details anytime, thanks.

[jwfromm]

Hi @sanirudh, sorry for the late reply on this, we have been working on a direct onnx to relax importer and have made good progress, it’s currently on our branch here. We plan on upstreaming once we do a little clean up and documentation but it’s in pretty reasonable shape

[sanirudh]

Thanks for the reply @jwfromm. This looks great, looking forward to the PR upstream