# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: F401, F821 """BasePyModule: Base class for IRModules with Python function support.""" import inspect import os from typing import Any, Optional, Union import numpy as np from tvm_ffi import Function import tvm from tvm import relax, tirx from tvm.ir import IRModule from tvm.runtime import Device, Tensor from tvm.target import Target try: from torch.utils.dlpack import to_dlpack as to_dlpack_legacy except ImportError: to_dlpack_legacy = None try: from tvm_ffi._optional_torch_c_dlpack import load_torch_c_dlpack_extension _FASTER_DLPACK_EXTENSION = load_torch_c_dlpack_extension() except ImportError: _FASTER_DLPACK_EXTENSION = None 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 self.compiled_tir_funcs: dict[str, Function] = {} self.extern_funcs: dict[str, Function] = {} self.tir_func_names: list[str] = [] self.relax_func_names: list[str] = [] self.relax_vm: relax.VirtualMachine | None = None self.pyfuncs: dict[str, Any] = {} if target is None: target = Target.from_device(device) elif isinstance(target, str): target = Target(target) self.target = target self._collect_function_names() self._compile_functions() self._wrap_tir_functions() self._wrap_relax_functions() self._register_python_functions() def _collect_function_names(self): """Collect names of TIR and Relax functions from IRModule.""" for global_var, func in self.ir_mod.functions_items(): if isinstance(func, tirx.PrimFunc): self.tir_func_names.append(global_var.name_hint) elif isinstance(func, relax.Function): self.relax_func_names.append(global_var.name_hint) def _compile_functions(self): """Compile TIR and Relax functions using JIT compilation.""" # Compile TIR functions first tir_mod = tvm.IRModule( { gv: func for gv, func in self.ir_mod.functions_items() if isinstance(func, tirx.PrimFunc) } ) if tir_mod: try: tir_exec_mod = tvm.compile(tir_mod, target=self.target) for func_name in self.tir_func_names: self.compiled_tir_funcs[func_name] = tir_exec_mod[func_name] # pylint: disable=broad-exception-caught except Exception as error: print(f"Warning: Failed to compile one or more TIR functions: {error}") if self.relax_func_names: try: exec_mod = tvm.compile(self.ir_mod, target=self.target) self.relax_vm = relax.VirtualMachine(exec_mod, self.device) # pylint: disable=broad-exception-caught except Exception as error: print(f"Warning: Failed to compile Relax VM: {error}") self.relax_vm = None def _wrap_tir_functions(self): """Wrap TIR functions to make them accessible as instance attributes.""" for func_name, func in self.compiled_tir_funcs.items(): setattr(self, func_name, func) def _wrap_relax_functions(self): """Wrap Relax functions to be callable from Python with auto conversion.""" for func_name in self.relax_func_names: def _create_relax_wrapper(name): def wrapper(*args, **kwargs): """Wrapper for Relax function with automatic tensor conversion.""" if hasattr(self.ir_mod, "pyfuncs") and name in self.ir_mod.pyfuncs: return self.ir_mod.pyfuncs[name](*args, **kwargs) if self.relax_vm is not None: converted_args = self._convert_pytorch_to_tvm(list(args)) converted_kwargs = { k: self._convert_pytorch_to_tvm(v) for k, v in kwargs.items() } result = self.relax_vm[name](*converted_args, **converted_kwargs) return self._convert_tvm_to_pytorch(result) raise RuntimeError( f"Neither converted Python function nor Relax VM available for {name}" ) wrapper.__name__ = name wrapper.__doc__ = f"Wrapped Relax function: {name}" return wrapper setattr(self, func_name, _create_relax_wrapper(func_name)) def _register_python_functions(self): """Register Python functions with the VM runtime for call_py_func support.""" if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs: return try: register_py_func = tvm.get_global_func("vm.builtin.register_py_func") except ValueError: return for func_name, py_func in self.ir_mod.pyfuncs.items(): def create_py_func_wrapper(name, original_func): def wrapper(*args, **kwargs): converted_args = [self._convert_tvm_to_pytorch(arg) for arg in args] converted_kwargs = { k: self._convert_tvm_to_pytorch(v) for k, v in kwargs.items() } result = original_func(self, *converted_args, **converted_kwargs) return self._convert_pytorch_to_tvm(result) wrapper.__name__ = name return wrapper wrapped_func = create_py_func_wrapper(func_name, py_func) register_py_func(func_name, wrapped_func) def call_tir(self, tir_func, args, out_ty): """Call a TIR function with PyTorch tensors.""" # Try to get function name from different sources if isinstance(tir_func, str): func_name = tir_func elif hasattr(tir_func, "name"): func_name = tir_func.name elif hasattr(tir_func, "__name__"): func_name = tir_func.__name__ else: # Try to find by function object reference for name, func in self.compiled_tir_funcs.items(): if func == tir_func: func_name = name break else: func_name = None if not func_name or func_name not in self.compiled_tir_funcs: available_funcs = list(self.compiled_tir_funcs.keys()) raise ValueError( f"Could not resolve or find compiled TIR function: {tir_func}. " f"Available functions: {available_funcs}" ) func = self.compiled_tir_funcs[func_name] out = self._create_output_tensors(out_ty, args) tvm_args = self._convert_pytorch_to_tvm(args) tvm_out = self._convert_pytorch_to_tvm(out) func(*tvm_args, *tvm_out) result = self._convert_tvm_to_pytorch(tvm_out) return result[0] if len(result) == 1 else result def call_dps_packed(self, func_name: str, args, out_ty): """Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack.""" if hasattr(self, func_name) and callable(getattr(self, func_name)): return getattr(self, func_name)(*args) if func_name not in self.extern_funcs: try: self.extern_funcs[func_name] = tvm.get_global_func(func_name) except ValueError as error: raise ValueError( f"Function '{func_name}' not found as a global function. " f"Please implement it as a method or register it." ) from error func = self.extern_funcs[func_name] out = self._create_output_tensors(out_ty, args) tvm_args = self._convert_pytorch_to_tvm(args) tvm_out = self._convert_pytorch_to_tvm(out) func(*tvm_args, *tvm_out) return out[0] if len(out) == 1 else out def call_py_func(self, func_name: str, args): """Call a Python function stored in the module's pyfuncs.""" if func_name not in self.pyfuncs: raise ValueError(f"Python function '{func_name}' not found in module pyfuncs") py_func = self.pyfuncs[func_name] return py_func(self, *args) def _create_output_tensors(self, out_ty, in_args=None): # pylint: disable=import-outside-toplevel import torch ty_list = out_ty if isinstance(out_ty, list) else [out_ty] out_tensors = [] for ty in ty_list: if isinstance(ty, tuple | list) and all(isinstance(x, int | np.integer) for x in ty): out_tensors.append(torch.zeros(list(map(int, ty)), dtype=torch.float32)) continue if hasattr(ty, "shape") and hasattr(ty, "dtype"): concrete_shape = self._infer_concrete_shape_from_args(ty.shape, in_args) torch_dtype = self._convert_tvm_dtype_to_torch(ty.dtype) out_tensors.append(torch.zeros(concrete_shape, dtype=torch_dtype)) continue out_tensors.append(torch.zeros((1,), dtype=torch.float32)) return out_tensors def _infer_concrete_shape_from_args(self, shape, in_args): concrete = [] symbolic_positions = [] for idx, dim in enumerate(shape): if isinstance(dim, int | np.integer): concrete.append(int(dim)) elif isinstance(dim, tirx.IntImm): concrete.append(int(dim.value)) else: concrete.append(None) symbolic_positions.append(idx) if not symbolic_positions: return concrete candidates = [] if in_args is not None: if not isinstance(in_args, list | tuple): in_args = [in_args] for obj in in_args: if hasattr(obj, "shape") and isinstance(obj.shape, tuple | list): try: candidates.append(tuple(int(x) for x in obj.shape)) continue except (ValueError, TypeError): # Skip objects with invalid shapes pass target_ndim = len(shape) for cand in candidates: if len(cand) == target_ndim: for pos in symbolic_positions: concrete[pos] = cand[pos] if all(x is not None for x in concrete): return concrete raise ValueError( "Cannot infer concrete output shape from symbolic shape and inputs. " "Please provide a concrete `out_ty` (e.g., a tuple/list of ints) " "or ensure input tensors carry shapes that determine output extents." ) def _convert_tvm_dtype_to_torch(self, tvm_dtype: str) -> "torch.dtype": """Convert TVM dtype string to PyTorch dtype.""" # pylint: disable=import-outside-toplevel import torch dtype_mapping = { "float32": torch.float32, "float64": torch.float64, "int32": torch.int32, "int64": torch.int64, "bool": torch.bool, } return dtype_mapping.get(str(tvm_dtype), torch.float32) def _convert_pytorch_to_tvm( self, tensors: Any | list[Any] | tuple[Any, ...] ) -> Tensor | list[Tensor]: """Convert PyTorch tensors to TVM Tensors using DLPack.""" # pylint: disable=import-outside-toplevel import torch if isinstance(tensors, list | tuple): return [self._convert_single_pytorch_to_tvm(t) for t in tensors] return self._convert_single_pytorch_to_tvm(tensors) def _convert_single_pytorch_to_tvm(self, tensor: Any) -> Tensor: """Convert a single PyTorch tensor to TVM Tensor with faster DLPack converter.""" # pylint: disable=import-outside-toplevel import torch if isinstance(tensor, Tensor): return tensor if isinstance(tensor, torch.Tensor): # 1. Try faster C++ DLPack converter if _FASTER_DLPACK_EXTENSION is not None: try: dlpack = torch.to_dlpack(tensor) return tvm.runtime.from_dlpack(dlpack) except (AttributeError, ValueError): pass # Fall through to the next method # 2. Try modern `torch.to_dlpack` (preferred for PyTorch >= 1.7) try: dlpack = torch.to_dlpack(tensor) return tvm.runtime.from_dlpack(dlpack) except (AttributeError, ValueError): pass # Fall through to the next method # 3. Try legacy `torch.utils.dlpack.to_dlpack` if to_dlpack_legacy: try: dlpack = to_dlpack_legacy(tensor) return tvm.runtime.from_dlpack(dlpack) except (AttributeError, ValueError) as error_legacy: print( f"Warning: Legacy DLPack conversion failed ({error_legacy}), " f"using numpy fallback." ) # 4. If all DLPack methods fail, use numpy fallback numpy_array = tensor.detach().cpu().numpy() return tvm.runtime.tensor(numpy_array, device=self.device) # For other types (like scalars, lists), convert to numpy first try: numpy_array = np.array(tensor, dtype=np.float32) return tvm.runtime.tensor(numpy_array, device=self.device) except (TypeError, ValueError) as error: raise TypeError( f"Unsupported type for conversion to TVM Tensor: {type(tensor)}" ) from error def _convert_tvm_to_pytorch( self, tvm_tensors: Any | list[Any] ) -> Union["torch.Tensor", list["torch.Tensor"]]: """Convert TVM Tensors to PyTorch tensors using DLPack.""" if isinstance(tvm_tensors, list | tuple): return [self._convert_single_tvm_to_pytorch(tensor) for tensor in tvm_tensors] return self._convert_single_tvm_to_pytorch(tvm_tensors) def _convert_single_tvm_to_pytorch(self, tvm_tensor: Any) -> "torch.Tensor": """Convert a single TVM Tensor to PyTorch tensor using faster DLPack converter.""" # pylint: disable=import-outside-toplevel import torch if isinstance(tvm_tensor, torch.Tensor): return tvm_tensor if not isinstance(tvm_tensor, Tensor): return torch.tensor(tvm_tensor) # 1. Try faster C++ DLPack converter if _FASTER_DLPACK_EXTENSION is not None: try: return torch.from_dlpack(tvm_tensor) except (AttributeError, ValueError): pass # Fall through to the next method # 2. Try standard DLPack conversion try: return torch.from_dlpack(tvm_tensor) # pylint: disable=broad-exception-caught except Exception as error: print(f"Warning: DLPack conversion from TVM failed ({error}), using numpy fallback") numpy_array = tvm_tensor.numpy() return torch.from_numpy(numpy_array) def get_function(self, name: str) -> Function | None: """Get a compiled function by name.""" if name in self.compiled_tir_funcs: return self.compiled_tir_funcs[name] if name in self.extern_funcs: return self.extern_funcs[name] if self.relax_vm and name in self.relax_func_names: try: if hasattr(self, name): return getattr(self, name) return self.relax_vm[name] except AttributeError as error: print(f"Warning: Failed to get Relax function '{name}': {error}") return None def list_functions(self) -> dict[str, list[str]]: """List all available functions.""" return { "tirx": self.tir_func_names, "relax": self.relax_func_names, "extern": list(self.extern_funcs.keys()), } def add_python_function(self, name: str, func: callable): """Add a Python function to the module.""" self.pyfuncs[name] = func # Create a wrapper that handles both instance methods and static functions # pylint: disable=import-outside-toplevel import functools @functools.wraps(func) def wrapper(*args, **kwargs): sig = inspect.signature(func) params = list(sig.parameters.keys()) if params and params[0] == "self": return func(self, *args, **kwargs) else: return func(*args, **kwargs) # Set the wrapper as an instance attribute setattr(self, name, wrapper) def script( self, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = "I", module_alias: str = "cls", int_dtype: str = "int32", float_dtype: str = "void", verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, show_object_address: bool = False, show_all_ty: bool = True, extra_config: dict | None = None, ) -> str: """Print TVM IR into TVMScript text format with Python function support. This method extends the standard IRModule script() method to handle Python functions stored in the IRModule's pyfuncs attribute. """ # First get the standard IRModule script base_script = self.ir_mod.script( name=name, show_meta=show_meta, ir_prefix=ir_prefix, module_alias=module_alias, int_dtype=int_dtype, float_dtype=float_dtype, verbose_expr=verbose_expr, indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, show_object_address=show_object_address, show_all_ty=show_all_ty, extra_config=extra_config, ) # If there are no Python functions, return the base script if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs: return base_script # Insert Python functions into the script return self._insert_python_functions(base_script, indent_spaces) def _insert_python_functions(self, base_script: str, indent_spaces: int) -> str: """Insert Python functions into the TVMScript output.""" lines = base_script.split("\n") result_lines = [] # Find the class definition line and insert Python functions after it class_found = False class_indent = 0 for line in lines: result_lines.append(line) # Look for class definition if not class_found and line.strip().startswith("class "): class_found = True class_indent = len(line) - len(line.lstrip()) # Insert Python functions after the class definition if hasattr(self.ir_mod, "pyfuncs") and self.ir_mod.pyfuncs: for func_name, func in self.ir_mod.pyfuncs.items(): # Get the function source code func_source = self._get_function_source(func) if func_source: # Format the function with proper indentation formatted_func = self._format_python_function( func_name, func_source, class_indent + indent_spaces ) result_lines.append(formatted_func) result_lines.append("") # Add empty line for separation return "\n".join(result_lines) def _get_function_source(self, func: callable) -> str | None: """Get the source code of a Python function.""" try: source = inspect.getsource(func) return source except (OSError, TypeError): # If we can't get the source, return None return None def _format_python_function(self, _func_name: str, func_source: str, indent: int) -> str: """Format a Python function with proper indentation for TVMScript.""" lines = func_source.split("\n") formatted_lines = [] for line in lines: # Skip the function definition line if it's already properly indented if line.strip().startswith("def ") or line.strip().startswith("@"): # Keep decorators and function definition as is formatted_lines.append(" " * indent + line.strip()) else: # Add proper indentation for the function body formatted_lines.append(" " * indent + line.strip()) return "\n".join(formatted_lines) def show(self, style: str | None = None, black_format: bool | None = None, **kwargs) -> None: """A sugar for print highlighted TVM script with Python function support. This method extends the standard IRModule show() method to handle Python functions stored in the IRModule's pyfuncs attribute. """ from tvm.script.highlight import cprint # pylint: disable=import-outside-toplevel if black_format is None: env = os.environ.get("TVM_BLACK_FORMAT") black_format = env and int(env) script_content = self.script(**kwargs) cprint(script_content, style=style, black_format=black_format)