I want to replace part of the calculations in the model with operators that I define in tvm.te level. What should I do?
I have read the documents, including the method of optimizing the operator at the ‘te level’, including optimizing and build model, but I did not see how to apply the optimized ‘te level operator’ to the model.
Do I have to define a new relay op and follow this discussion to replace?Or do i have any other way?
In Relay, I believe you have to define a new relay op and then replace it as you mentioned.
In Relax (with unity branch), you can define your computation with a PrimFunc and call it directly with call_tir. That way you don’t need to define a custom operator for special cases, and instead directly call into the implementation function.
The relay translator uses the relay op strategy to generate PrimFunc implementations of each of the relay operators and inserts a call to those PrimFuncs using call_tir that I mentioned earlier.
So once you’ve converted to relax, you can replace any of the generated PrimFuncs with a custom impelementation.