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
Hi @cgerum ,
Thanks for the pointer. However, I notice that
-
lowered_ir_mods
in the provided link istvm.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)