Any way to extract tir function from relay?

Questions
1

@Wheest for exmaple, build the relay, and the relay trans to relay ir.

first to get lib.

    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod_bert, target=target, params=params_bert)
    print(lib.ir_mod)

the ir contains below layer structure,

%22 = reshape(%21, newshape=[-1, 1024]) /* ty=Tensor[(256, 1024), float32] */;
%23 = cast(%bert.encoder.layer.attention.self.query.weight, dtype="float16") /* ty=Tensor[(1024, 1024), float16] */;

the tir func may like below

@tvmgen_default_fused_reshape_cast = primfn(args_7: handle, arg_type_ids_7: handle, num_args_7: int32, out_ret_value_7: handle, out_ret_tcode_7: handle, resource_handle_7: handle) -> int32

how can i get the tir func and feed other input, only test this kernel for some reason?

1 Like
2

I’m also interested in this. @chenugray , did you find a way to get the TIR from a Relay op?

Thanks!

1 Like
3

i wanna find way to automative single op accuracy testing. I use another way to extract single op ir definition, then run the single op ir func def. not using tir.

1 Like
4

me too.did you find a way to get the TIR from a Relay op? i use this interface

tec = relay.backend.te_compiler.get()

lower_module = tec.lower(opt_mode[“main”], “llvm”)

but get errors…

1 Like
5

I’m also interested in this.

6

It is possible to insert your own tir pass into the pass context to print out the tir at the end of the transformation phase.

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

Then when issuing the relay build command, make sure to append your new pass onto the end of the pass sequence

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

The “tir.add_lower_pass”: [{3, dump_tir)]} part is just stating which phase of the tir transformation passes you want to insert your pass into, i.e. after which phase of optimization you’d like to output the tir.

You can find more information in gallery/how_to/extend_tvm/low_level_custom_pass.py

1 Like
7

it works,thank you. But I print to the log

@tvm.tir.transform.prim_func_pass(opt_level=0)

def dump_tir(f, mod, ctx):

 with open("temp_tir.log", "a") as file:

     file.write(f.script())

 return f