[Relax] Add Python function support and BasePyModule for PyTorch integration (#18229)

### **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

Author: tlopex

python/tvm/ir/module.py

@@ -67,6 +67,7 @@ def __init__(self, functions=None, attrs=None, global_infos=None):
             attrs,
             global_infos,
         )
+        self.pyfuncs = {}
 
     def clone(self) -> "IRModule":
         return _ffi_api.Module_Clone(self)

python/tvm/relax/__init__.py

@@ -98,6 +98,9 @@
 # utils
 from .utils import convert_to_expr
 
+# BasePyModule
+from .base_py_module import BasePyModule
+
 # Import submodules in the last to avoid dependency
 from . import exec_builder
 from . import expr

python/tvm/relax/base_py_module.py

@@ -19,6 +19,7 @@
 from typing import Any, Dict, Union
 
 import tvm
+from tvm.relax import ExternFunc
 from ....ir.module import IRModule
 from ...ir_builder import IRBuilder
 from . import doc
@@ -86,12 +87,14 @@ def parse(
         extra_vars = _default_globals()
 
     ann = {}
+    all_pyfuncs = {}
     if inspect.isfunction(program):
         ann = {program.__name__: program.__annotations__}
     elif inspect.isclass(program):
         for name, func in program.__dict__.items():
             if inspect.isfunction(func):
                 ann[name] = func.__annotations__
+                all_pyfuncs[name] = func
 
     source = Source(program)
     parser = Parser(source, ann)
@@ -101,6 +104,10 @@ def parse(
         except ParserError as err:
             parser.report_error(err.node, err.args[0])
     ret = builder.get()
+    # Attach pyfuncs to the IRModule
+    if inspect.isclass(program) and isinstance(ret, IRModule):
+        _attach_pyfuncs_to_irmodule(ret, all_pyfuncs)
+
     # check well-formedness in both Relax and TIR
     if check_well_formed:
         check_ret = ret
@@ -122,3 +129,65 @@ def parse(
                 err=f"{WELL_FORMED_ERROR_MESSAGE}\n\nTraceback: {str(err)}",
             )
     return ret
+
+
+def _create_python_packed_func(pyfunc):
+    """Create a PackedFunc wrapper for a Python function.
+
+    This function creates a PackedFunc that can be called from TVM runtime
+    and will execute the original Python function.
+
+    Parameters
+    ----------
+    pyfunc : Callable
+        The Python function to wrap.
+
+    Returns
+    -------
+    PackedFunc
+        A PackedFunc that wraps the Python function.
+    """
+
+    def packed_func_wrapper(*args, **kwargs):
+        """Wrapper function that calls the original Python function."""
+        try:
+            result = pyfunc(*args, **kwargs)
+            return result
+        except Exception as error:
+            print(f"Error calling Python function {pyfunc.__name__}: {error}")
+            raise
+
+    return packed_func_wrapper
+
+
+def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
+    """Attach Python functions to IRModule with reduced nesting."""
+    if not all_pyfuncs:
+        return
+
+    if not hasattr(irmodule, "pyfuncs"):
+        irmodule.pyfuncs = {}
+
+    for global_var, func in irmodule.functions_items():
+        if not isinstance(func, ExternFunc):
+            continue
+        if not func.attrs.get("is_pyfunc", False):
+            continue
+
+        pyfunc_name = global_var.name_hint
+        if pyfunc_name not in all_pyfuncs:
+            continue
+
+        pyfunc = all_pyfuncs[pyfunc_name]
+        irmodule.pyfuncs[pyfunc_name] = pyfunc
+
+        try:
+            source_code = inspect.getsource(pyfunc)
+            func = func.with_attr("python_source", source_code)
+        except (OSError, TypeError):
+            func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")
+
+        packed_func = _create_python_packed_func(pyfunc)
+        func = func.with_attr("python_packed_func", packed_func)
+
+        irmodule[global_var] = func

python/tvm/script/parser/core/entry.py

@@ -343,6 +343,8 @@ class Parser(doc.NodeVisitor):
     function_annotations: Optional[Dict[str, Dict[str, Any]]]
     var_table: VarTable
     inside_function: bool  # whether we are within a function
+    current_class: Optional[str] = None  # current class being parsed
+    base_py_module_context: bool = False  # whether current class inherits from BasePyModule
 
     def __init__(
         self,
@@ -414,6 +416,39 @@ def pop_token():
 
         return _deferred(pop_token)
 
+    def set_class_context(self, class_name: str, is_base_py_module: bool = False):
+        """Set the current class context for parsing.
+
+        Parameters
+        ----------
+        class_name : str
+            The name of the current class being parsed.
+        is_base_py_module : bool
+            Whether the current class inherits from BasePyModule.
+        """
+        self.current_class = class_name
+        self.base_py_module_context = is_base_py_module
+
+    def _get_current_class_context(self) -> Optional[str]:
+        """Get the current class context.
+
+        Returns
+        -------
+        Optional[str]
+            The name of the current class, or None if not in a class context.
+        """
+        return self.current_class
+
+    def _is_base_py_module_context(self) -> bool:
+        """Check if the current class context allows Python functions.
+
+        Returns
+        -------
+        bool
+            True if Python functions are allowed in the current context.
+        """
+        return self.base_py_module_context
+
     def with_diag_source(self, source: Source):
         """Add a new source as with statement.
 

python/tvm/script/parser/core/parser.py

@@ -18,7 +18,7 @@
 from tvm.ir import Range
 from ...ir_builder.ir import *  # pylint: disable=redefined-builtin
 from . import parser as _parser
-from .entry import ir_module
+from .entry import ir_module, pyfunc
 
 
 __all__ = [
@@ -28,5 +28,6 @@
     "dummy_global_info",
     "Range",
     "lookup_vdevice",
+    "pyfunc",
     "vdevice",
 ]

python/tvm/script/parser/ir/__init__.py

@@ -17,9 +17,12 @@
 """The entry point of TVM parser for ir module."""
 
 import inspect
-from typing import Optional, Type
+from typing import Callable, Optional, Type
 
-from tvm.ir import IRModule
+from tvm.ir import IRModule, GlobalVar
+from tvm.relax.expr import ExternFunc
+from tvm.relax.base_py_module import BasePyModule
+from tvm import cpu, ir
 
 from .._core import parse, utils
 
@@ -47,7 +50,86 @@ def ir_module(mod: Optional[Type] = None, check_well_formed: bool = True) -> IRM
     def decorator_wrapper(mod):
         if not inspect.isclass(mod):
             raise TypeError(f"Expect a class, but got: {mod}")
+
+        # Check BasePyModule inheritance
+        base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)
+
         m = parse(mod, utils.inspect_class_capture(mod), check_well_formed=check_well_formed)
+
+        if base_py_module_inherited:
+            # Collect pyfunc methods
+            pyfunc_methods = [
+                name
+                for name, attr in mod.__dict__.items()
+                if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
+            ]
+
+            mod._pyfunc_methods = pyfunc_methods
+
+            # Create ExternFunc nodes
+
+            for method_name in pyfunc_methods:
+                try:
+                    existing_gvars = [
+                        global_var
+                        for global_var in m.get_global_vars()
+                        if global_var.name_hint == method_name
+                    ]
+
+                    extern_func = ExternFunc(method_name)
+                    extern_func = extern_func.with_attr("is_pyfunc", True)
+                    extern_func = extern_func.with_attr("function_type", "python")
+                    extern_func = extern_func.with_attr("python_function_name", method_name)
+                    extern_func = extern_func.with_attr(
+                        "python_source", f"# Source for {method_name}"
+                    )
+                    extern_func = extern_func.with_attr("python_packed_func", None)
+
+                    if existing_gvars:
+                        m[existing_gvars[0]] = extern_func
+                    else:
+                        m[GlobalVar(method_name)] = extern_func
+
+                except Exception:  # pylint: disable=broad-exception-caught
+                    continue
+
+            class ModuleFactory:
+                """Factory class for creating BasePyModule instances with Python functions."""
+
+                def __init__(self, module, pyfunc_methods, original_class):
+                    self.ir_module = module
+                    self.pyfunc_methods = pyfunc_methods
+                    self.original_class = original_class
+
+                def __call__(self, device=None, target=None):
+
+                    if device is None:
+                        device = cpu(0)
+
+                    instance_ir_mod = ir.IRModule()
+                    for global_var, func in self.ir_module.functions_items():
+                        instance_ir_mod[global_var] = func
+
+                    instance = BasePyModule(instance_ir_mod, device, target)
+
+                    for method_name in self.pyfunc_methods:
+                        if hasattr(self.original_class, method_name):
+                            method = getattr(self.original_class, method_name)
+                            instance.add_python_function(method_name, method)
+
+                    return instance
+
+                def __getattr__(self, name):
+                    if hasattr(self.ir_module, name):
+                        return getattr(self.ir_module, name)
+                    raise AttributeError(
+                        f"'{self.__class__.__name__}' object has no attribute '{name}'"
+                    )
+
+            factory = ModuleFactory(m, pyfunc_methods, mod)
+            setattr(factory, "__name__", mod.__name__)
+            return factory
+
         setattr(m, "__name__", mod.__name__)
         return m
 
@@ -61,4 +143,10 @@ def decorator_wrapper(mod):
         return decorator_wrapper
 
 
-setattr(ir_module, "dispatch_token", "ir")
+def pyfunc(func: Callable):
+    # Set the dispatch_token on the decorated function
+    setattr(func, "dispatch_token", "pyfunc")
+    return func
+
+
+setattr(pyfunc, "dispatch_token", "pyfunc")