How to lower relay built modules?

Questions
#1

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

1 Like
#2

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.

#3

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.

1 Like
#4

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
#5

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")
1 Like
#6

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