tlopex opened a new pull request, #18326:
URL: https://github.com/apache/tvm/pull/18326
This PR implements the `R.call_py_func` operator that allows compiled TVM
Relax modules to call Python functions at runtime. This enables integration
between TVM's compiled code and Python through a robust VM backend
implementation.
#### Simple Usage with BasePyModule
```python
@I.ir_module
class MyModule(BasePyModule):
@I.pyfunc
def torch_relu(self, x):
return torch.relu(x)
@R.function
def forward(x: R.Tensor((10,), "float32")) -> R.Tensor((10,), "float32"):
return R.call_py_func("torch_relu", (x,), out_sinfo=R.Tensor((10,),
"float32"))
```
#### Direct VM Backend Usage (Manual)
```python
# Manually register Python function with VM backend
register_func = tvm.get_global_func("vm.builtin.register_py_func")
register_func("my_func", my_python_function)
# Use in Relax function (compiled to VM backend)
@R.function
def test(x: R.Tensor((5,), "float32")) -> R.Tensor((5,), "float32"):
return R.call_py_func("my_func", (x,), out_sinfo=R.Tensor((5,),
"float32"))
# Manual cleanup (required for direct VM backend usage)
clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
clear_func()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
