### **Overview**
This PR implements native Python function support in TVM Relax through
the `@I.pyfunc` decorator and `BasePyModule`, which enable seamless
integration between TVM's compilation pipeline and Python/PyTorch runtime
environments. This enhancement allows users to write Python functions
directly in TVMScript that can interoperate with Relax and TIR functions
that provides enhanced debugging capabilities and leveraging existing
PyTorch operator libraries.
### **Key Features**
**TVMScript Parser Enhancement**
- `@I.pyfunc` decorator: Marks Python functions for integration into IRModules
- Dual storage format: Stores both raw string representation (for TVMScript
printing) and captured PackedFunc (for runtime execution)
- ExternFunc representation: Each Python function is represented as an
ExternFunc node with attributes storing source code and runtime wrapper
**Complete BasePyModule Implementation**
- DLPack-based tensor conversion: Seamless conversion between PyTorch
tensors and TVM NDArrays
- Cross-function interoperability: Python functions can call Relax/TIR
functions and vice versa
- JIT compilation: Delays compilation until module instantiation for flexible
late-stage modifications
- Dynamic function registration: Supports runtime addition of Python functions
### Future Work
- TVMScript printer for IRModules with Python functions: Print IRModules
in proper format with high-level operator mapping from Relax ops to PyTorch
ops, handling symbolic shapes
- R.call_py_func primitive: Introduce Relax primitive to invoke corresponding
PackedFunc of specified Python functions at runtime
* [Unity][TVMScript] Produce var = R.ExternFunc("") statements
Prior to this commit, any `ExternFunc` usage in a relax function would
print the string name of the function on its own line, omitting any
variable definition, and later use of the variable would occur without
a definition. This commit updates the printing of `R.ExternFunc` to
appear as a normal relax variable.
* Preserve special handling as callee, test round-trip
* Updated parser to handle `var = R.ExternFunc(...)` in IRModule
Since this is now a representation that may be produced by the
TVMScript printer, it must also be handled at the parser.
The current cross-function calls in TVMScript will cause PyLint warnings,
since the GlobalVar will be marked as undefined vars, e.g.:
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
gv0 = R.call_tir(tir_func, x, R.Tensor((128,), dtype="float32")) # <= `tir_func` is not defined in Python syntax.
return gv0
```
This PR changes the behavior into `TestModule.tir_func` instead of direct `tir_func`
```python
@I.ir_module
class TestModule:
@T.prim_func
def tir_func(
x: T.Buffer((T.int64(128),), "float32"), y: T.Buffer((T.int64(128),), "float32")
):
T.evaluate(0)
@R.function
def foo(x: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
cls = TestModule # Use `cls` to refer the current Module
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128,), dtype="float32"))
return gv0
```
NOTE: It's a breaking change, the old style is deprecated.
Additionally, this PR contains the following minor fixes:
- mark `R.function` as staticmethod as what we do for `T.prim_func`
- make `I`, `R`, `T`, `cls` be the builtin keywords for the printer
- define names for functions, modules to prevent naming conflict
- checking the var names is valid via regex expression
- fix typos
This PR enables context-aware parsing for TVMScript. It means that the parser has full control of the statements in the specific context/namespace.
For example, we can override the global var `__call__` method in Relax function to make sure to generate Relax Calls instead of Relay Calls.
Multi-line strings might make less sense to be printed out by default,
as they could be LLVM snippets, CUDA source code and anything hard to
comprehend but easy to mess up with the TVMScript itself. Therefore,
this PR is introduced to print them as metadata by default.
This PR introduces some minor restructuring of the `python/tvm/script`
folder structure to make it more convenient for future upstreaming.
Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
Tianqi Chen