Class BasePyModule

python/tvm/relax/base_py_module.py:46–616

Base class that allows Python functions in IRModule with DLPack conversion. This class provides the infrastructure for: 1. JIT compilation of TIR and Relax functions. 2. DLPack-based conversion between PyTorch tensors and TVM Tensors. 3. Wrapping Relax functions for easy Python call

class BasePyModule:
    """Base class that allows Python functions in IRModule with DLPack conversion.

    This class provides the infrastructure for:
    1. JIT compilation of TIR and Relax functions.
    2. DLPack-based conversion between PyTorch tensors and TVM Tensors.
    3. Wrapping Relax functions for easy Python calling.
    4. Cross-function calls between Python, TIR, and Relax functions.

    Only IRModules that inherit from this class are allowed to contain Python functions.
    """

    def __del__(self):
        """Clean up registered Python functions on module destruction."""
        try:
            clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
            clear_func()
        except (ValueError, AttributeError):
            pass

    def __init__(
        self,
        ir_mod: IRModule,
        device: Device,
        target: Target | None = None,
    ):
        """Initialize BasePyModule with JIT compilation and DLPack conversion."""
        self.device = device
        self.ir_mod = ir_mod

        # Delegate IRModule operations
        self.functions = ir_mod.functions
        self.attrs = ir_mod.attrs
        self.global_infos = ir_mod.global_infos
        self.__getitem__ = ir_mod.__getitem__
        self.__setitem__ = ir_mod.__setitem__
        self.functions_items = ir_mod.functions_items
        self.with_attr = ir_mod.with_attr
        self.get_attr = ir_mod.get_attr
        self.update_global_info = ir_mod.update_global_info

        def _getattr_python_function(name: str) -> Any:
            """Support direct attribute access to funcs and IRModule methods."""
            if name in self.pyfuncs:
                return self.pyfuncs[name]
            if name in self.compiled_tir_funcs:
                return self.compiled_tir_funcs[name]
            if self.relax_vm and name in self.relax_func_names:
                try:
                    return self.relax_vm[name]
                except AttributeError:  # More specific exception
                    return None
            if hasattr(self.ir_mod, name):
                return getattr(self.ir_mod, name)
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

        self.__getattr__ = _getattr_python_function