E2e model compilation with Relax

Development unity
#1

Hi,

We’re just starting to look into Relax for e2e model compilation and I’ve some basic questions. I would greatly appreciate if you could share some insights on the these:

  1. AFAIK, when we translate Relay into Relax, it uses OpStrategy to find the OP implementations, however, it appears to use only compute and ignore schedule. This means that if we want to make any scheduling changes to the generated STIR for relax, it can be done only via a TIR pass or metascheduler/auto-tuner. Could you please confirm if this understanding is correct?

  2. Can we use Relax to compile a model e2e without using MetaScheduler and execute it on VM? I’ve found some examples in the tests/python/relax using VM (without using auto-tuning) but these’re very small, single op examples. Just wondering if someone has tried using this workflow for a full model and can share their observations

Thank you so much! -Jyotsna

1 Like
#2

Hi @jverma,

Glad that you’re looking into Relax code! To answer your questions:

  1. The key theme in TVM Unity and Relax is everything is an IRModule (Relax + TIR functions) to IRModule transformation, and transforming TIR Primfuncs is part of that. Writing custom TIR Primfuncs to replace default ones or applying scheduling to them is possible, and it does not have to be via TIR passes or MetaSchedule. @junrushao and @vinx13 may have hands-on experience with it, feel free to add more or any examples.

  2. Yes, totally. MetaSchedule is a pass, and is not part of the minimum build pipeline. For example, this e2e Resnet example is run on Relax VM without Metaschedule tuning.

Happy to answer any additional questions you may have! -Yuchen

2 Likes
#3

MetaSchedule allows custom manual schedule to be registered, there are some examples here https://github.com/apache/tvm/blob/main/src/meta_schedule/schedule/cpu/winograd.cc

#4

Hi @yuchenj,

That’s good to know. Thanks for sharing the example!

#5

Hello @yuchenj: could you please share an example for this. We’re trying to do something similar where we want to replace the default TIR for an operator with the custom TIR primfunc.

#6

Hi @jverma,

Sure, the relay translator from_relay api has an optional translate_op_with_tir parameter which allows user to replace the default TIR with a custom TIR primfunc implementation.

Example: https://github.com/tlc-pack/relax/blob/relax/tests/python/relax/test_relay_translator.py#L272 Implementation: https://github.com/tlc-pack/relax/blob/relax/python/tvm/relax/testing/relay_translator.py#L170

#7

Thanks a lot! It’s much appreciated.

1 Like
#8

there are demos about compiling model with relax frontend in tvm unity, hope to help you. relax/test_frontend_dynamo.py at relax · octoml/relax · GitHub

relax/test_frontend_from_fx.py at relax · octoml/relax · GitHub