Any way to replace a certain 'relay op' in the model with the calculation I defined using 'te'?

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?

Looking forward to your answer :smiley:

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.

1 Like

Thank you so much for your answer.:heart:

How should transfer the relay level model to Relax (with unity branch) level. Are there any docs or test code files that I can refer to?

You can use the relay_translator to convert your relay model to relax. Take a look at the examples in https://github.com/apache/tvm/blob/unity/tests/python/relax/test_relay_translator.py.

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.

1 Like

Thank you very much, it really helps a lot.