Insert Relay Pass After Optimizations

Development
#1

When implementing a module pass, I decided to have it run after the default Relay optimization passes, but before transition to TIR. This allowed me to benefit from the simple and canonical Relay representation, while already able to reason about the likely temporary buffers between fused operations.

However, in my top level code that uses relay.build, I did not see a way to register my pass to be run after the default Relay passes.

My solution is a simple optional hook that allows the user to register a function that takes a module and params (seem to be gone in current TVM), and returns a module:

const runtime::PackedFunc* pfPostPass = runtime::Registry::Get("relay.backend.PostOptPass");
if (pfPostPass) {
  Map<String, Constant> argParams;
  for (const auto& param : params) {
    argParams.Set(param.first, Constant(param.second));
  }
  relay_module = (*pfPostPass)(relay_module, argParams);
}

It is inserted right after the OptimizeImpl call:

It would be used as follows:

@tvm.register_func("relay.backend.PostOptPass")
def _post_pass(mod, params):
    return MyCustomPass()(mod)

Is this something that is useful to anyone else? Would you be open to include this in the code base? Is there a better way to achive this?

1 Like
#2

Hi @r.stahl, we have been aware of this challenge of adding a custom pass to the current compilation pipeline (see the challange C3 in @sunggg’s Relax Pass Infrastructure).

In the design of Relax (Relay Next), we follow the following design principles:

  1. Each pass is a IRModule → IRModule transformation.
  2. Decouple the optimization passes from the build system. For example, the AutoTIR(MetaSchedule) tuning is an optimization pass in Relax, and it is outside of the build.
  3. A minimum and universal build that is able to build every valid IRModule to runtime.Module. The build can be invoked at any stage during the compilation, and can be invoked inside a pass for example to do performance measurement for some tuning passes.

These design principles enable flexible and customizable compilation pipelines without the need to hack into the core of the compiler, and allow developers and researchers to explore new spaces. Feel free to check out @sunggg’s discussion thread for more details. :slight_smile:

#3

@yuchenj Thank you for making me aware of this! This should definitely give sufficient flexibility to integrate a pass after other optimization passes.

However, since Relax is a larger effort that may take quite some time to upstream, I was wondering whether my proposed simple hook could be useful to anyone else and find support for a quick upstreaming.

#4

I have a feeling that we already have various kinds of “hook”, but not sure if they are relevant here. cc @Mousius @lhutton1 @mbs-octoml

#5

@masahi Thanks for your input. I’m aware of the RelayToTIR hook (https://github.com/apache/tvm/pull/8423). That would fit the requirement of “Relay After Optimization”, but would add a burden to make the translation to TIR. Would there be a way to easily call into the default TIR conversion? If yes, this also seems like a good approach.

#6

Hi @r.stahl, apologies for chiming in late, I’ve been away the past few days.

It sounds as though you’d like to insert a pass into the standard flow of TVM - after the default relay optimizations, but before TIR. I’m not convinced the RelayToTIR hook would work for your use case, since this targets functions considered “external”. Your functions would need to be marked with a Compiler attribute for the relevant RelayToTIR hook to then be run.

I’m curious, what’s the reason for wanting to implement this using a hook, rather than just inserting the pass where you desire? If the concern is that the pass shouldn’t be enabled by default, I think we could make use PassContext to manually turn it on.

Hope this helps :slight_smile:

#7

@lhutton1 Thank you for clarifying.

In my opinion, the pass makes more sense as an external project, because it carries quite a few dependencies with it, while not being suitable for the default optimization.