As far as I know, before Relax officially take out we have to implement relay_to_tir plugin to “bring your own lowering”. This could be just a relay pass to rewrite those relay calls which we are interested into call_lowered.
Basically I think we need:
-
Provide your own relay op strategy for TensorIR. Currently relay strategy machanism is only about TE compute and schedule yet.
-
Interact with your own op strategy in
relay_to_tirpass, getPrimFuncfor each relay call you want to lower via the strategy, then rewrite original expr tocall_loweredof the primfunc. -
Also if you want to lower fused functions, you have to determine how to compose
PrimFuncs of sub-ops into a singlePrimFunc. That could be a lot of things, but generally we can follow howTECompilerchain TE expressions of fused sub-ops.