New Relay Operator based on TensorIR

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_tir pass, get PrimFunc for each relay call you want to lower via the strategy, then rewrite original expr to call_lowered of the primfunc.

  • Also if you want to lower fused functions, you have to determine how to compose PrimFuncs of sub-ops into a single PrimFunc. That could be a lot of things, but generally we can follow how TECompiler chain TE expressions of fused sub-ops.

2 Likes