How to lower relay built modules?

A tutorial shows that tensor expression can be lower to specific codes. How can I apply similar lowering to relay built models? I believe there should be an API perform the transformation from relay to tir but just cannot find. Appreciate if any one can give an referernce

In the default compilation flow the lowering from relay → te → tir is implemented using a relay operator strategy.

For a tutorial on relay operator strategies have a look at: Relay Operator Strategy — tvm 0.9.dev0 documentation

You will find the default strategies in tvm/relay/op/strategy.

Thanks for the pointer. The tutorial is for how to bind platform specific te/tir with relay operator. If now I have a relay model,

x = relay.var("x", shape=[1, 10])
w = relay.var("w", shape=[20, 10])
y = relay.dense(x, w)
fn = relay.Function([x, w], y)
mod = tvm.IRModule.from_expr(fn)

How can I lower such expression / module to platform specific TE/TIR, for example, targeting CPU and use “llvm” as the platform.

As far as I know this is still not directly possible. If your main goal is to look at the generated tir have a look at my previous answer to: TVM terms: relay, topi, tir, te - #6 by jack-willturner

If you really need to implement something like a relay.lower you could probably return the TIR modules here: https://github.com/apache/tvm/blob/b5f1dabce45e4d4a68a2ecadeee811bac93b1c3c/python/tvm/relay/build_module.py#L479

1 Like

Hi @cgerum ,

Thanks for the pointer. However, I notice that

  • lowered_ir_mods in the provided link is tvm.ir.container.Map which is not the expected TIR module.
  • The customized print_tir in your previous ansower now raises following error.
  Check failed: (it != functions.end()) is false: There is no definition of @tvmgen_default_fused_layout_transform

Any workthrough to this situation?

full code

x = relay.var("x", shape=[1, 10])
w = relay.var("w", shape=[20, 10])
y = relay.nn.dense(x, w)
fn = relay.Function([x, w], y)
fmod = tvm.IRModule.from_expr(fn)

@tvm.tir.transform.prim_func_pass(opt_level=0)
def print_tir(f, mod, ctx):
    print(f)

with tvm.transform.PassContext(
        opt_level=3, config={"tir.add_lower_pass": [(3, print_tir)]}
    ):
        lib = relay.build(fmod, target="llvm")

Sorry, my bad it should be:

@tvm.tir.transform.prim_func_pass(opt_level=0)
def print_tir(f, mod, ctx):
    print(f)
    return f

Considering lowered_ir_mods, the map contains a mapping from tvm.target.Target to the corresponding IR-Modules. So if you want to get the IR just do:

for mod in lowered_ir_mods.values():
    print(mod)
1 Like