chore: import upstream snapshot with attribution

python/tvm/relax/base_py_module.py

@@ -0,0 +1,614 @@
+# 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)