We are starting to experiment with tensor expressions, aimed at creating custom-generated code for an accelerator. I can see what’s happening at the tensor expression level, but am trying to see the same thing when coming from Relay.
I have a toy example with element-wise matrix add:
A = te.placeholder(shape=(160, 640), name='A') B = te.placeholder(shape=(160, 640), name='B') C = te.compute(A.shape, lambda i,j: A[i,j] + B[i,j], name='C') s = my_custom_schedule(C)
I can then do all of the following:
# see the schedule! tedd.viz_schedule_tree(s, output_dot_string = True) # create an IR module and look at it! irmod = tvm.lower(s, [A, B, C], name="fadd") # compile it to C code! cmod = tvm.build(irmod, [A, B, C], target="c", name="fadd") cmod.save("./fadd.c", "c")
So far all good.
But what I really want is to do all this from Relay, using the strategy mechanism to plug in my custom schedule. I’ve managed to plug in the custom schedule, create a Relay IRModule, compile it to a runtime module, and see the schedule reflected in the generated C code.
@schedule_injective.register(["mydevice"]) def my_custom_schedule(T) ... x = relay.var('x', shape=(160, 640)) y = relay.var('y', shape=(160, 640)) z = relay.add(x, y) func : relay.function.Function = relay.Function([x, y], z) mod : tvm.ir.module.IRModule = tvm.IRModule.from_expr(func) _, lib, _ = relay.build(mod, target="c --device=mydevice")
But in the relay.build flow I can’t figure out how to intercept and examine the schedule and the TIR-level module as in the TE flow above so that we can see what’s going on.
I’ve looked at Relay’s build_module and there does not seem to be an API to get just the intermediate TIR-level representation including the schedule – the output is a runtime module.
I’ve what seem to be related posts on this question, and seen some recent posts on the AOT compiler and revising the compilation flow, but it’s still not clear if there is a way to do this or if maybe I’m thinking about it the wrong way.
Thanks for any help.