chore: import upstream snapshot with attribution

python/tvm/runtime/__init__.py

@@ -0,0 +1,54 @@
+# isort: skip_file
+# 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.
+"""TVM runtime namespace."""
+
+from tvm_ffi import convert, Object
+from tvm_ffi._dtype import dtype as DataType, DataTypeCode
+
+# Import _ffi_node_api for its side effect of installing AsRepr as
+# tvm_ffi.core.__object_repr__.
+from . import _ffi_node_api
+
+# class exposures
+from .script_printer import Scriptable
+from .object_generic import ObjectConvertible
+from .device import Device
+from ._tensor import Tensor, tensor, empty
+from .module import Module
+from .executable import Executable
+
+# function exposures
+from ._tensor import device, cpu, cuda, opencl, vulkan, metal
+from ._tensor import vpi, rocm, ext_dev, from_dlpack
+from .module import load_module, enabled, system_lib, load_static_library, num_threads
+from .object_generic import const
+from .params import (
+    save_param_dict,
+    load_param_dict,
+    save_param_dict_to_file,
+    load_param_dict_from_file,
+)
+
+try:
+    from . import disco
+except (ImportError, ValueError):
+    # disco C++ runtime is in libtvm_runtime_extra which may not be present.
+    # Make the disco module optional.
+    disco = None  # type: ignore[assignment]
+
+from tvm_ffi import Shape as ShapeTuple

python/tvm/runtime/_ffi_api.py

@@ -0,0 +1,22 @@
+# 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.
+"""FFI APIs for tvm.runtime"""
+
+import tvm_ffi
+
+# Exports functions registered in runtime namespace.
+tvm_ffi.init_ffi_api("runtime", __name__)

python/tvm/runtime/_ffi_node_api.py

@@ -0,0 +1,44 @@
+# 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.
+
+# pylint: disable=invalid-name, unused-argument
+"""FFI for tvm.node"""
+
+import tvm_ffi
+import tvm_ffi.core
+
+
+# The implementations below are default ones when the corresponding
+# functions are not available in the runtime only mode.
+# They will be overriden via tvm_ffi.init_ffi_api to the ones registered
+def AsRepr(obj):
+    return type(obj).__name__ + "(" + obj.__ctypes_handle__().value + ")"
+
+
+def SaveJSON(obj):
+    raise RuntimeError("Do not support object serialization in runtime only mode")
+
+
+def LoadJSON(json_str):
+    raise RuntimeError("Do not support object serialization in runtime only mode")
+
+
+# Exports functions registered in node namespace.
+tvm_ffi.init_ffi_api("node", __name__)
+
+# Override the default repr function for tvm_ffi.core.Object.
+tvm_ffi.core.__object_repr__ = AsRepr

python/tvm/runtime/_tensor.py

@@ -0,0 +1,522 @@
+# 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.
+# pylint: disable=invalid-name, unused-import, redefined-outer-name
+# ruff: noqa: F401, RUF005
+"""Runtime Tensor API"""
+
+import ctypes
+import warnings
+from typing import Optional
+
+import numpy as np
+
+try:
+    import ml_dtypes
+except ImportError:
+    ml_dtypes = None
+
+import tvm_ffi
+from tvm_ffi import DLDeviceType, device
+
+import tvm
+from tvm.runtime import Device
+
+from . import _ffi_api
+
+
+def from_dlpack(ext_tensor):
+    """
+    Convert an external tensor to an Tensor.
+
+    Parameters
+    ----------
+    ext_tensor : object
+        The external tensor to convert.
+
+    require_alignment : int
+        The minimum required alignment to check for the tensor.
+
+    require_contiguous : bool
+        Whether to check for contiguous memory.
+    """
+    # TODO(tvm-team): change to require_alignment=0 and require_contiguous=False
+    # once we update the compiler generated code to guard against misaligned access.
+    return tvm_ffi.from_dlpack(
+        ext_tensor,
+        require_alignment=64,
+        require_contiguous=True,
+    )
+
+
+@tvm_ffi.register_object("ffi.Tensor")
+class Tensor(tvm_ffi.core.Tensor):
+    """Lightweight Tensor class of TVM runtime.
+
+    Strictly this is only an Array Container (a buffer object)
+    No arthimetic operations are defined.
+    All operations are performed by TVM functions.
+
+    The goal is not to re-build yet another array library.
+    Instead, this is a minimal data structure to demonstrate
+    how can we use TVM in existing project which might have their own array containers.
+    """
+
+    def __setitem__(self, in_slice, value):
+        """Set ndarray value"""
+        if (
+            not isinstance(in_slice, slice)
+            or in_slice.start is not None
+            or in_slice.stop is not None
+        ):
+            raise ValueError("Array only support set from numpy array")
+        if isinstance(value, Tensor):
+            if not value.same_as(self):
+                value.copyto(self)
+        elif isinstance(value, np.ndarray | np.generic):
+            self.copyfrom(value)
+        else:
+            raise TypeError(f"type {type(value)} not supported")
+
+    def copyfrom(self, source_array):
+        """Perform a synchronous copy from the array.
+
+        Parameters
+        ----------
+        source_array : array_like
+            The data source we should like to copy from.
+
+        Returns
+        -------
+        arr : Tensor
+            Reference to self.
+        """
+        if isinstance(source_array, Tensor):
+            source_array.copyto(self)
+            return self
+
+        if not isinstance(source_array, np.ndarray):
+            try:
+                source_array = np.array(source_array, dtype=self.dtype)
+            except Exception:
+                raise TypeError(
+                    f"array must be an array_like data, type {type(source_array)} is not supported"
+                )
+
+        t = tvm_ffi.dtype(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t = t.with_lanes(1)
+            dtype = str(t)
+
+        if source_array.shape != shape:
+            raise ValueError(
+                f"array shape do not match the shape of Tensor {source_array.shape} vs {shape}"
+            )
+        numpy_str_map = tvm_ffi.dtype._NUMPY_DTYPE_TO_STR
+        np_dtype_str = (
+            numpy_str_map[source_array.dtype]
+            if source_array.dtype in numpy_str_map
+            else str(source_array.dtype)
+        )
+        if (not source_array.flags["C_CONTIGUOUS"]) or (
+            dtype == "bfloat16" or dtype != np_dtype_str
+        ):
+            if dtype == "bfloat16":
+                source_array = np.frombuffer(source_array.tobytes(), "uint16")
+            source_array = np.ascontiguousarray(
+                source_array, dtype="uint16" if dtype == "bfloat16" else dtype
+            )
+        if self.dtype.startswith("float4_e2m1fn"):
+            # we need to pack the input data when converting to float4_e2m1fn type,
+            data_bits = source_array.view(dtype="uint8").flatten()
+            if data_bits.size % 2:
+                data_bits = np.pad(data_bits, (0, 1), mode="constant", constant_values=0)
+            data_bits = data_bits.reshape(-1, 2)
+            packed = ((data_bits[:, 0] & 0x0F) << 4) | (data_bits[:, 1] & 0x0F)
+            source_array = packed.astype(np.int8)
+        assert source_array.flags["C_CONTIGUOUS"]
+        data = source_array.ctypes.data_as(ctypes.c_void_p)
+        nbytes = source_array.size * source_array.dtype.itemsize
+        _ffi_api.TVMTensorCopyFromBytes(self, data, nbytes)
+        return self
+
+    def __repr__(self):
+        # exception safety handling for chandle=None
+        if self.__chandle__() == 0:
+            return type(self).__name__ + "(chandle=None)"
+        res = f"<tvm.runtime.Tensor shape={self.shape}, {self.device}>\n"
+        res += self.numpy().__repr__()
+        return res
+
+    def __str__(self):
+        return str(self.numpy())
+
+    def numpy(self):
+        """Convert this array to numpy array
+
+        Returns
+        -------
+        np_arr : numpy.ndarray
+            The corresponding numpy array.
+        """
+        t = tvm_ffi.dtype(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        old_dtype = dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t = t.with_lanes(1)
+            dtype = str(t)
+        if dtype == "int4":
+            dtype = "int8"
+        if dtype in [
+            "bfloat16",
+            "float8_e3m4",
+            "float8_e4m3",
+            "float8_e4m3b11fnuz",
+            "float8_e4m3fn",
+            "float8_e4m3fnuz",
+            "float8_e5m2",
+            "float8_e5m2fnuz",
+            "float8_e8m0fnu",
+            "float6_e2m3fn",
+            "float6_e3m2fn",
+            "float4_e2m1fn",
+        ]:
+            if ml_dtypes is None:
+                raise RuntimeError(
+                    f"ml_dtypes is not installed, cannot convert {dtype} array to numpy."
+                )
+            try:
+                dtype = getattr(ml_dtypes, dtype)
+            except AttributeError:
+                raise RuntimeError(f"ml_dtypes has no attribute '{dtype}', cannot convert array.")
+        np_arr = np.empty(shape, dtype=dtype)
+        assert np_arr.flags["C_CONTIGUOUS"]
+        data = np_arr.ctypes.data_as(ctypes.c_void_p)
+        # TODO(kathy): revisit and get a mirrored function of ffi::GetDataSize
+        # in Python to replace line below
+        nbytes = np_arr.size if dtype == "bool" else (np_arr.size * old_dtype.bits + 7) // 8
+        _ffi_api.TVMTensorCopyToBytes(self, data, nbytes)
+
+        if old_dtype == "int4" or old_dtype.startswith("float4_e2m1fn"):
+            length = np_arr.size
+            np_arr = np_arr.view("int8")
+            np_arr_ret = np.empty((length,), dtype="int8")
+            np_arr = np_arr.reshape((length,))
+            odd_index = np.bitwise_and(np_arr, 0x0F)
+            even_index = np.bitwise_and(np_arr >> 4, 0x0F)
+            np_arr_ret[1::2] = odd_index[0 : length // 2]
+            np_arr_ret[0::2] = even_index[0 : (length + 1) // 2]
+            return np_arr_ret.reshape(shape).view(dtype)
+
+        return np_arr
+
+    def copyto(self, target, mem_scope=None):
+        """Copy array to target
+
+        Parameters
+        ----------
+        target : Tensor
+            The target array to be copied, must have same shape as this array.
+
+        mem_scope : Optional[str]
+            The memory scope of the array.
+        """
+        if isinstance(target, Tensor):
+            return self._copyto(target)
+        if isinstance(target, tvm_ffi.core.Device):
+            res = empty(self.shape, self.dtype, target, mem_scope)
+            return self._copyto(res)
+        raise ValueError(f"Unsupported target type {type(target)}")
+
+    def _copyto(self, target_nd):
+        """Internal function that implements copy to target ndarray."""
+        _ffi_api.TVMTensorCopyFromTo(self, target_nd)
+        return target_nd
+
+    def _create_view(self, shape, dtype: str | None = None, relative_byte_offset: int = 0):
+        """Create a view into an existing array.
+
+        The view shares the same allocation and datatype as the
+        existing array, but can have a different array shape.  This is
+        useful for runtimes that support non-flat memory, where both
+        the physical shape of an allocation and the logical shape of
+        the tensor it represents may need to be independently
+        specified.
+
+        Warning: This function should not be used outside of low-level
+        manipulations, as it breaks non-aliasing assumptions made by
+        TVM.  This function may also be removed/replaced in the
+        future.
+
+        Parameters
+        ----------
+        shape: Union[tvm_ffi.Shape, Sequence[typing.SupportsInt]]
+
+            The shape of the view.
+
+        dtype: Optional[str]
+
+            The datatype of the view.  If None (default), the view
+            will be the same data type as the current array.
+
+        relative_byte_offset: int
+
+            The location of the view, relative to the location of the current
+            array.
+
+            Note: While the `DLTensor.byte_offset` field of the returned view
+            is usually the same as `relative_byte_offset`, this is not
+            guaranteed.  The `DLTensor.byte_offset` field is relative to the
+            start of the backing allocation, while the `relative_byte_offset`
+            is relative to the start of `self`.
+
+        """
+
+        if not isinstance(shape, tvm_ffi.Shape):
+            shape = tvm_ffi.Shape([int(dim) for dim in shape])
+
+        if dtype is None:
+            dtype = self.dtype
+
+        return _ffi_api.TVMTensorCreateView(self, shape, dtype, relative_byte_offset)
+
+
+def empty(shape, dtype="float32", device=None, mem_scope=None):
+    """Create an empty array given shape and device
+
+    Parameters
+    ----------
+    shape : Union[tvm_ffi.Shape, Sequence[typing.SupportsInt]]
+        The shape of the array.
+
+    dtype : type or str
+        The data type of the array.
+
+    device : Device
+        The device of the array.
+
+    mem_scope : Optional[str]
+        The memory scope of the array.
+
+    Returns
+    -------
+    arr : tvm.runtime.Tensor
+        The array tvm supported.
+    """
+    device = device or cpu()
+    if not isinstance(shape, tvm_ffi.Shape):
+        shape = tvm_ffi.Shape([int(dim) for dim in shape])
+    dtype = tvm_ffi.dtype(dtype)
+    arr = _ffi_api.TVMTensorAllocWithScope(shape, dtype, device, mem_scope)
+    return arr
+
+
+def tensor(arr, device=None, mem_scope=None):
+    """Create an tensor from source arr.
+
+    Parameters
+    ----------
+    arr : numpy.ndarray
+        The array to be copied from
+
+    device : Device, optional
+        The device to create the array
+
+    mem_scope : Optional[str]
+        The memory scope of the array
+
+    Returns
+    -------
+    ret : Tensor
+        The created array
+    """
+    device = device or cpu()
+
+    if not isinstance(arr, np.ndarray | Tensor):
+        arr = np.asarray(arr)
+    return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)
+
+
+def cpu(dev_id=0):
+    """Construct a CPU device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLCPU, dev_id)
+
+
+def cuda(dev_id=0):
+    """Construct a CUDA GPU device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLCUDA, dev_id)
+
+
+def rocm(dev_id=0):
+    """Construct a ROCM device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLROCM, dev_id)
+
+
+def opencl(dev_id=0):
+    """Construct a OpenCL device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLOpenCL, dev_id)
+
+
+def metal(dev_id=0):
+    """Construct a metal device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLMetal, dev_id)
+
+
+def vpi(dev_id=0):
+    """Construct a VPI simulated device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLVPI, dev_id)
+
+
+def vulkan(dev_id=0):
+    """Construct a Vulkan device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLVulkan, dev_id)
+
+
+def ext_dev(dev_id=0):
+    """Construct a extension device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+
+    Note
+    ----
+    This API is reserved for quick testing of new
+    device by plugin device API as ext_dev.
+    """
+    return device(DLDeviceType.kDLExtDev, dev_id)
+
+
+def hexagon(dev_id=0):
+    """Construct a Hexagon device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLHexagon, dev_id)
+
+
+def webgpu(dev_id=0):
+    """Construct a webgpu device.
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    dev : Device
+        The created device
+    """
+    return device(DLDeviceType.kDLWebGPU, dev_id)
+
+
+# Register back to FFI
+tvm_ffi.core._set_class_tensor(Tensor)

python/tvm/runtime/device.py

@@ -0,0 +1,330 @@
+# 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.
+"""Common runtime ctypes."""
+
+# pylint: disable=invalid-name
+import json
+
+import tvm_ffi
+
+from . import _ffi_api
+
+RPC_SESS_MASK = 128
+
+
+class Device(tvm_ffi.core.Device):
+    """TVM device strucure."""
+
+    def _GetDeviceAttr(self, device_type, device_id, attr_id):
+        """Internal helper function to invoke runtime.GetDeviceAttr"""
+        # pylint: disable=import-outside-toplevel
+        return _ffi_api.GetDeviceAttr(device_type, device_id, attr_id)
+
+    @property
+    def exist(self):
+        """Whether this device exists.
+
+        Returns True if TVM has support for the device, if the
+        physical device is present, and the device is accessible
+        through appropriate drivers (e.g. CUDA/Vulkan).
+
+        Returns
+        -------
+        exist : bool
+            True if the device exists
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 0) != 0
+
+    @property
+    def max_threads_per_block(self):
+        """Maximum number of threads on each block.
+
+        Returns device value for CUDA, Metal, ROCm, OpenCL, and Vulkan
+        devices.  Returns remote device value for RPC devices.
+        Returns None for all other devices.
+
+        Returns
+        -------
+        max_threads_per_block : int or None
+            The number of threads on each block
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 1)
+
+    @property
+    def warp_size(self):
+        """Number of threads that execute concurrently.
+
+        Returns device value for CUDA, ROCm, and Vulkan.  Returns
+        1 for Metal and OpenCL devices, regardless of the physical
+        device.  Returns remote device value for RPC devices.  Returns
+        None for all other devices.
+
+        Returns
+        -------
+        warp_size : int or None
+            Number of threads that execute concurrently
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 2)
+
+    @property
+    def max_shared_memory_per_block(self):
+        """Total amount of shared memory per block in bytes.
+
+        Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
+        Returns remote device value for RPC devices.  Returns None for
+        all other devices.
+
+        Returns
+        -------
+        max_shared_memory_per_block : int or None
+            Total amount of shared memory per block in bytes
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 3)
+
+    @property
+    def compute_version(self):
+        """Get compute version number as string.
+
+        Returns maximum API version (e.g. CUDA/OpenCL/Vulkan)
+        supported by the device.
+
+        Returns device value for CUDA, ROCm, OpenCL, and
+        Vulkan. Returns remote device value for RPC devices.  Returns
+        None for all other devices.
+
+        Returns
+        -------
+        version : str or None
+            The version string in `major.minor` format.
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 4)
+
+    @property
+    def device_name(self):
+        """Return the vendor-specific name of device.
+
+        Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
+        Returns remote device value for RPC devices.  Returns None for
+        all other devices.
+
+        Returns
+        -------
+        device_name : str or None
+            The name of the device.
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 5)
+
+    @property
+    def max_clock_rate(self):
+        """Return the max clock frequency of device (kHz).
+
+        Returns device value for CUDA, ROCm, and OpenCL.  Returns
+        remote device value for RPC devices.  Returns None for all
+        other devices.
+
+        Returns
+        -------
+        max_clock_rate : int or None
+            The maximum clock frequency of the device (kHz)
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 6)
+
+    @property
+    def multi_processor_count(self):
+        """Return the number of compute units in the device.
+
+        Returns device value for CUDA, ROCm, and OpenCL.  Returns
+        remote device value for RPC devices.  Returns None for all
+        other devices.
+
+        Returns
+        -------
+        multi_processor_count : int or None
+            Thee number of compute units in the device
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 7)
+
+    @property
+    def max_thread_dimensions(self):
+        """Return the maximum size of each thread axis
+
+        Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
+        Returns remote device value for RPC devices.  Returns None for
+        all other devices.
+
+        Returns
+        -------
+        dims: List of int, or None
+            The maximum length of threadIdx.x, threadIdx.y, threadIdx.z
+
+        """
+        return json.loads(self._GetDeviceAttr(self.dlpack_device_type(), self.index, 8))
+
+    @property
+    def api_version(self):
+        """Returns version number of the SDK used to compile TVM.
+
+        For example, CUDA_VERSION for CUDA or VK_HEADER_VERSION for
+        Vulkan.
+
+        Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
+        Returns remote device value for RPC devices.  Returns None for
+        all other devices.
+
+        Returns
+        -------
+        version : int or None
+            The version of the SDK
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 11)
+
+    @property
+    def driver_version(self):
+        """Returns version number of the driver
+
+        Returns driver vendor's internal version number.
+        (e.g. "450.408.256" for nvidia-driver-450)
+
+        Returns device value for opencl and vulkan.  Returns remote
+        device value for RPC devices.  Returns None for all other
+        devices.
+
+        Returns
+        -------
+        version : str or None
+            The version string in `major.minor.patch` format.
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 12)
+
+    @property
+    def l2_cache_size_bytes(self):
+        """Return the size of the device L2 cache in bytes
+
+        Supported devices include CUDA/ROCM/OpenCL.
+
+        Returns
+        -------
+        l2_cache_size_bytes : int or None
+            The size of the device L2 cache in bytes returned by device runtime API.
+            Return None if the device does not support this feature.
+
+        Note
+        ----
+        The value returned by opencl's API is smaller than actual device L2 cache size.
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 13)
+
+    @property
+    def total_global_memory(self):
+        """Return size of the total global memory.
+
+        Supported devices include CUDA/ROCm/Metal/OpenCL.
+
+        Returns
+        -------
+        total_global_memory : int or None
+            Return the total size of global memory on device in bytes.
+            Return None if the device does not support this feature.
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 14)
+
+    @property
+    def available_global_memory(self):
+        """Return size of the available global memory.
+
+        Supported devices include CUDA.
+
+        Returns
+        -------
+        available_global_memory : int or None
+            Return the amount of unallocated global memory on device in bytes.
+            Return None if the device does not support this feature.
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 15)
+
+    def texture_spatial_limit(self):
+        """Returns limits for textures by spatial dimensions
+
+        Returns
+        -------
+        limit : int or None
+            Maximum size of the texture by spatial dimensions
+
+        """
+        return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 12)
+
+    def create_raw_stream(self):
+        """Create a new runtime stream at the context.
+
+        User should free the stream after use.
+
+        Returns
+        -------
+        stream : TVMStreamHandle
+            The created runtime stream.
+        """
+        return _ffi_api.Device_StreamCreate(self)
+
+    def free_raw_stream(self, stream):
+        """Free a created stream handle.
+
+        Parameters
+        ----------
+        stream : TVMStreamHandle
+            The stream which should to be released.
+        """
+        _ffi_api.Device_StreamFree(self, stream)
+
+    def set_raw_stream(self, stream):
+        """Set a created stream handle.
+
+        Parameters
+        ----------
+        stream : TVMStreamHandle
+            The stream which should to be set to the device.
+        """
+        _ffi_api.Device_SetStream(self, stream)
+
+    def sync(self, stream=None):
+        """Synchronize until jobs finished at the context.
+
+        Parameters
+        ----------
+        stream : TVMStreamHandle
+            Jobs in this stream should be finished.
+        """
+        _ffi_api.Device_StreamSync(self, stream or 0)
+
+    def __device_type_name__(self):
+        if self.dlpack_device_type() >= RPC_SESS_MASK:
+            tbl_id = self.dlpack_device_type() / RPC_SESS_MASK - 1
+            dev_type = self.dlpack_device_type() % RPC_SESS_MASK
+            return f"remote[{tbl_id}]:{Device._DEVICE_TYPE_TO_NAME[dev_type]}"
+        return Device._DEVICE_TYPE_TO_NAME[self.dlpack_device_type()]
+
+
+tvm_ffi.core._set_class_device(Device)

python/tvm/runtime/disco/__init__.py

@@ -0,0 +1,28 @@
+# isort: skip_file
+# 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.
+"""TVM distributed runtime API."""
+
+from .session import (
+    DModule,
+    DPackedFunc,
+    DRef,
+    ProcessSession,
+    Session,
+    SocketSession,
+    ThreadedSession,
+)

python/tvm/runtime/disco/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs from C++"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("runtime.disco", __name__)

python/tvm/runtime/disco/process_pool.py

@@ -0,0 +1,193 @@
+# 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.
+# pylint: disable=invalid-name
+"""Pipe worker for multi-processing."""
+
+import os
+import subprocess
+import sys
+
+from tvm_ffi import Shape, register_global_func
+
+
+class DiscoPopenWorker:
+    """A subprocess worker via Popen.
+
+    PopenWorker provides a low-level
+    API to interact with a separate process via Popen.
+
+    Parameters
+    ----------
+    worker_id : int
+        The worker id of the current worker.
+
+    num_workers : int
+        The total number of workers.
+
+    num_groups : int
+        The total number of worker groups.
+
+    stdout: Union[None, int, IO[Any]]
+        The standard output streams handler specified for the popen process.
+
+    stderr: Union[None, int, IO[Any]]
+        The standard error streams handler specified for the popen process.
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        worker_id: int,
+        num_workers: int,
+        num_groups: int,
+        entrypoint: str = "tvm.exec.disco_worker",
+        stdout=None,
+        stderr=None,
+    ):
+        self.worker_id = worker_id
+        self.num_workers = num_workers
+        self.num_groups = num_groups
+        self.entrypoint = entrypoint
+        self._proc = None
+        self._stdout = stdout
+        self._stderr = stderr
+
+    def __del__(self):
+        try:
+            self.kill()
+        except ImportError:
+            pass
+
+    def kill(self):
+        """Kill the current running process and cleanup.
+
+        Note
+        ----
+        The worker can start a new process when send is called again.
+        """
+        if self._proc is not None:
+            # kill all child processes recursively
+            try:
+                _kill_child_processes(self._proc.pid)
+            except TypeError:
+                pass
+            try:
+                self._proc.kill()
+            except OSError:
+                pass
+
+            # Join the child process to avoid zombie processes
+            self.join(timeout=1.0)
+            self._proc = None
+
+    def join(self, timeout=None):
+        """Join the current process worker before it terminates.
+
+        Parameters
+        ----------
+        timeout: Optional[number]
+            Timeout value, block at most timeout seconds if it
+            is a positive number.
+        """
+        if self._proc:
+            try:
+                self._proc.wait(timeout)
+            except subprocess.TimeoutExpired:
+                pass
+
+    def start(self):
+        """Start a new subprocess if nothing is available"""
+        if self._proc is not None:
+            return None, None
+
+        # connect subprocess with a pair of pipes
+        main_read, worker_write = os.pipe()
+        worker_read, main_write = os.pipe()
+
+        cmd = [
+            sys.executable,
+            "-m",
+            self.entrypoint,
+            str(self.worker_id),
+            str(self.num_workers),
+            str(self.num_groups),
+        ]
+        if sys.platform == "win32":
+            import msvcrt  # pylint: disable=import-error,import-outside-toplevel
+
+            worker_read_handle = msvcrt.get_osfhandle(worker_read)
+            worker_write_handle = msvcrt.get_osfhandle(worker_write)
+            os.set_handle_inheritable(worker_read_handle, True)
+            os.set_handle_inheritable(worker_write_handle, True)
+            cmd += [str(worker_read_handle), str(worker_write_handle)]
+            self._proc = subprocess.Popen(
+                cmd,
+                close_fds=False,
+                stdout=self._stdout,
+                stderr=self._stderr,
+            )
+        else:
+            cmd += [str(worker_read), str(worker_write)]
+            self._proc = subprocess.Popen(  # pylint: disable=consider-using-with
+                cmd,
+                pass_fds=(worker_read, worker_write),
+                stdout=self._stdout,
+                stderr=self._stderr,
+            )
+
+        # close worker side of the pipe
+        os.close(worker_read)
+        os.close(worker_write)
+        return main_read, main_write
+
+
+def _kill_child_processes(pid):
+    """Kill all child processes recursively for a given pid.
+
+    Parameters
+    ----------
+    pid : int
+        The given parameter id.
+    """
+    import psutil  # pylint: disable=import-outside-toplevel
+
+    try:
+        parent = psutil.Process(pid)
+        children = parent.children(recursive=True)
+    except psutil.NoSuchProcess:
+        return
+
+    for process in children:
+        try:
+            process.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+
+@register_global_func("runtime.disco.create_process_pool")
+def _create_process_pool(num_workers: int, num_groups: int, entrypoint: str):
+    """Create a process pool where the workers' are [1, num_workers)."""
+    pool = [DiscoPopenWorker(i, num_workers, num_groups, entrypoint) for i in range(1, num_workers)]
+
+    def result_func(worker_id: int):
+        nonlocal pool
+        if worker_id != 0:
+            read_fd, write_fd = pool[worker_id - 1].start()
+            return Shape([read_fd, write_fd])
+        del pool
+        return None
+
+    return result_func

python/tvm/runtime/disco/session.py

@@ -0,0 +1,673 @@
+# 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
+"""This module defines a Session in Disco. Session is the primary interface that users interact
+with the distributed runtime.
+"""
+
+import logging
+import os
+import pickle
+from collections.abc import Callable, Sequence
+from typing import Any, Optional, Union
+
+import numpy as np
+from tvm_ffi import Object, Shape, get_global_func, register_global_func, register_object
+
+from .._tensor import Tensor
+from .._tensor import tensor as _as_Tensor
+from ..device import Device
+from . import _ffi_api, process_pool  # pylint: disable=unused-import
+
+
+@register_object("runtime.disco.DRef")
+class DRef(Object):
+    """An object that exists on all workers. The controller process assigns a unique "register id"
+    to each object, and the worker process uses this id to refer to the object residing on itself.
+    """
+
+    def debug_get_from_remote(self, worker_id: int) -> Any:
+        """Get the value of a DRef from a remote worker. It is only used for debugging purposes.
+
+        Parameters
+        ----------
+        worker_id : int
+            The id of the worker to be fetched from.
+
+        Returns
+        -------
+        value : object
+            The value of the register.
+        """
+        return _ffi_api.DRefDebugGetFromRemote(self, worker_id)  # type: ignore # pylint: disable=no-member
+
+    def debug_copy_from(
+        self,
+        worker_id: int,
+        value: np.ndarray | Tensor,
+    ) -> None:
+        """Copy an Tensor value to remote for debugging purposes.
+
+        Parameters
+        ----------
+        worker_id : int
+            The id of the worker to be copied to.
+
+        value : Union[numpy.ndarray, Tensor]
+            The value to be copied.
+        """
+        if not isinstance(value, Tensor):
+            value = _as_Tensor(value)
+        return _ffi_api.DRefDebugCopyFrom(self, worker_id, value)  # type: ignore # pylint: disable=no-member
+
+
+class DPackedFunc(DRef):
+    """A PackedFunc in a Disco session."""
+
+    # tvm_ffi Object subclasses cannot store Python attributes by default
+    # (the metaclass sets `__slots__ = ()`); list the field(s) we store here.
+    __slots__ = ("session",)
+
+    def __init__(self, dref: DRef, session: "Session") -> None:
+        self.__move_handle_from__(dref)
+        self.session = session
+
+    def __call__(self, *args) -> DRef:
+        return self.session.call_packed(self, *args)
+
+
+class DModule(DRef):
+    """A Module in a Disco session."""
+
+    # tvm_ffi Object subclasses cannot store Python attributes by default
+    # (the metaclass sets `__slots__ = ()`); list the field(s) we store here.
+    __slots__ = ("session",)
+
+    def __init__(self, dref: DRef, session: "Session") -> None:
+        self.__move_handle_from__(dref)
+        self.session = session
+
+    def __getitem__(self, name: str) -> DPackedFunc:
+        func = self.session._get_cached_method("ffi.ModuleGetFunction")
+        return DPackedFunc(func(self, name, False), self.session)
+
+
+@register_object("runtime.disco.Session")
+class Session(Object):
+    """A Disco interactive session. It allows users to interact with the Disco command queue with
+    various PackedFunc calling convention."""
+
+    # tvm_ffi Object subclasses cannot store Python attributes by default
+    # (the metaclass sets `__slots__ = ()`); list the fields we store here:
+    # the method-lookup cache and the lazily bound import helper.
+    __slots__ = ("_cache", "_import_python_module")
+
+    def _get_cached_method(self, name: str) -> Callable:
+        if not hasattr(self, "_cache"):
+            cache = self._cache = {}  # pylint: disable=attribute-defined-outside-init
+        else:
+            cache = self._cache
+        if name not in cache:
+            func = cache[name] = self.get_global_func(name)
+        else:
+            func = cache[name]
+        return func
+
+    def empty(
+        self,
+        shape: Sequence[int],
+        dtype: str,
+        device: Device | None = None,
+        worker0_only: bool = False,
+        in_group: bool = True,
+    ) -> DRef:
+        """Create an empty Tensor on all workers and attach them to a DRef.
+
+        Parameters
+        ----------
+        shape : tuple of int
+            The shape of the Tensor.
+
+        dtype : str
+            The data type of the Tensor.
+
+        device : Optional[Device] = None
+            The device of the Tensor.
+
+        worker0_only: bool
+            If False (default), allocate an array on each worker.  If
+            True, only allocate an array on worker0.
+
+        in_group: bool
+            Take effective when `worker0_only` is True. If True (default),
+            allocate an array on each first worker in each group. If
+            False, only allocate an array on worker0 globally.
+
+        Returns
+        -------
+        array : DRef
+            The created Tensor.
+
+        """
+        func = self._get_cached_method("runtime.disco.empty")
+        return func(Shape(shape), dtype, device, worker0_only, in_group)
+
+    def shutdown(self):
+        """Shut down the Disco session"""
+        _ffi_api.SessionShutdown(self)  # type: ignore # pylint: disable=no-member
+
+    @property
+    def num_workers(self) -> int:
+        """Return the number of workers in the session"""
+        return _ffi_api.SessionGetNumWorkers(self)  # type: ignore # pylint: disable=no-member
+
+    def get_global_func(self, name: str) -> DRef:
+        """Get a global function on workers.
+
+        Parameters
+        ----------
+        name : str
+            The name of the global function.
+
+        Returns
+        -------
+        func : DRef
+            The global packed function
+        """
+        return DPackedFunc(_ffi_api.SessionGetGlobalFunc(self, name), self)  # type: ignore # pylint: disable=no-member
+
+    def import_python_module(self, module_name: str) -> None:
+        """Import a python module in each worker
+
+        This may be required before call
+
+        Parameters
+        ----------
+        module_name: str
+
+            The python module name, as it would be used in a python
+            `import` statement.
+        """
+        if not hasattr(self, "_import_python_module"):
+            self._import_python_module = self.get_global_func("runtime.disco._import_python_module")
+
+        self._import_python_module(module_name)
+
+    def call_packed(self, func: DRef, *args) -> DRef:
+        """Call a PackedFunc on workers providing variadic arguments.
+
+        Parameters
+        ----------
+        func : PackedFunc
+            The function to be called.
+        *args : various types
+            In the variadic arguments, the supported types include:
+            - integers and floating point numbers;
+            - DLDataType;
+            - DLDevice;
+            - str (std::string in C++);
+            - DRef.
+
+        Returns
+        -------
+        return_value : various types
+            The return value of the function call.
+
+        Notes
+        -----
+        Examples of unsupported types:
+        - Tensor, DLTensor,;
+        - TVM Objects, including PackedFunc, Module and String.
+        """
+        return _ffi_api.SessionCallPacked(self, 0, 0, func, *args)  # type: ignore # pylint: disable=no-member
+
+    def _sync_worker(self, worker_id: int) -> None:
+        """Synchronize the controller with a worker, and it will wait until the worker finishes
+        executing all the existing instructions. This function is usually used for worker-0, because
+        it is the only worker that is assumed to collocate with the controller. Syncing with other
+        workers may not be supported and should only be used for debugging purposes.
+
+        Parameters
+        ----------
+        worker_id : int
+            The id of the worker to be synced with.
+        """
+        return _ffi_api.SessionSyncWorker(self, worker_id)  # type: ignore # pylint: disable=no-member
+
+    def _sync_all(self) -> None:
+        """Synchronize the controller with all workers in the current session, and it will
+        wait until all workers finish executing all the existing instructions."""
+        for i in range(self.num_workers):
+            self._sync_worker(i)
+
+    def sync_worker_0(self) -> None:
+        """Synchronize the controller with worker-0, and it will wait until the worker-0 finishes
+        executing all the existing instructions."""
+        return self._sync_worker(0)
+
+    def copy_from_worker_0(self, host_array: Tensor, remote_array: DRef) -> None:
+        """Copy an Tensor from worker-0 to the controller-side Tensor.
+
+        Parameters
+        ----------
+        host_array : numpy.ndarray
+            The array to be copied to worker-0.
+
+        remote_array : Tensor
+            The Tensor on worker-0.
+        """
+        return _ffi_api.SessionCopyFromWorker0(self, host_array, remote_array)  # type: ignore # pylint: disable=no-member
+
+    def copy_to_worker_0(self, host_array: Tensor, remote_array: DRef | None = None) -> DRef:
+        """Copy the controller-side Tensor to worker-0.
+
+        Parameters
+        ----------
+        host_array : Tensor
+            The array to be copied to worker-0.
+
+        remote_array : Optiona[DRef]
+            The destination Tensor on worker-0.
+
+        Returns
+        -------
+        output_array: DRef
+
+            The DRef containing the copied data on worker0, and
+            std::nullopt on all other workers.  If `remote_array` was
+            provided, this return value is the same as `remote_array`.
+            Otherwise, it is the newly allocated space.
+
+        """
+        if remote_array is None:
+            remote_array = self.empty(host_array.shape, host_array.dtype, worker0_only=True)
+
+        _ffi_api.SessionCopyToWorker0(self, host_array, remote_array)  # type: ignore # pylint: disable=no-member
+        return remote_array
+
+    def load_vm_module(
+        self,
+        path: str,
+        device: Device | None = None,
+    ) -> DModule:
+        """Load a VM module from a file.
+
+        Parameters
+        ----------
+        path : str
+            The path to the VM module file.
+
+        device : Optional[Device] = None
+            The device to load the VM module to. Default to the default device of each worker.
+
+        Returns
+        -------
+        module : DModule
+            The loaded VM module.
+        """
+        func = self._get_cached_method("runtime.disco.load_vm_module")
+        return DModule(func(path, device), self)
+
+    def init_ccl(self, ccl: str, *device_ids):
+        """Initialize the underlying communication collective library.
+
+        Parameters
+        ----------
+        ccl : str
+            The name of the communication collective library. Currently supported libraries are:
+            - nccl
+            - rccl
+            - mpi
+
+        *device_ids : int
+            The device IDs to be used by the underlying communication library.
+        """
+        assert ccl in ("nccl", "rccl"), f"Unsupported CCL backend: {ccl}"
+        _ffi_api.SessionInitCCL(self, ccl, Shape(device_ids))  # type: ignore # pylint: disable=no-member
+        self._clear_ipc_memory_pool()
+
+    def broadcast(
+        self,
+        src: np.ndarray | Tensor,
+        dst: DRef | None = None,
+        in_group: bool = True,
+    ) -> DRef:
+        """Broadcast an array to all workers
+
+        Parameters
+        ----------
+        src: Union[np.ndarray, Tensor]
+            The array to be broadcasted.
+
+        dst: Optional[DRef]
+            The output array.  If None, an array matching the shape
+            and dtype of `src` will be allocated on each worker.
+
+        in_group: bool
+            Whether the broadcast operation performs globally or in group as default.
+
+        Returns
+        -------
+        output_array: DRef
+
+            The DRef containing the broadcasted data on all workers.
+            If `dst` was provided, this return value is the same as
+            `dst`.  Otherwise, it is the newly allocated space.
+
+        """
+        if not isinstance(src, Tensor):
+            src = _as_Tensor(src)
+
+        if dst is None:
+            dst = self.empty(src.shape, src.dtype)
+
+        src_dref = self.copy_to_worker_0(src)
+        self.broadcast_from_worker0(src_dref, dst, in_group)
+
+        return dst
+
+    def broadcast_from_worker0(self, src: DRef, dst: DRef, in_group: bool = True) -> DRef:
+        """Broadcast an array from worker-0 to all other workers.
+
+        Parameters
+        ----------
+        src: Union[np.ndarray, Tensor]
+            The array to be broadcasted.
+
+        dst: Optional[DRef]
+            The output array.  If None, an array matching the shape
+            and dtype of `src` will be allocated on each worker.
+
+        in_group: bool
+            Whether the broadcast operation performs globally or in group as default.
+        """
+        func = self._get_cached_method("runtime.disco.broadcast_from_worker0")
+        func(src, in_group, dst)
+
+    def scatter(
+        self,
+        src: np.ndarray | Tensor,
+        dst: DRef | None = None,
+        in_group: bool = True,
+    ) -> DRef:
+        """Scatter an array across all workers
+
+        Parameters
+        ----------
+        src: Union[np.ndarray, Tensor]
+            The array to be scattered.  The first dimension of this
+            array, `src.shape[0]`, must be equal to the number of
+            workers.
+
+        dst: Optional[DRef]
+            The output array.  If None, an array with compatible shape
+            and the same dtype as `src` will be allocated on each
+            worker.
+
+        in_group: bool
+            Whether the scatter operation performs globally or in group as default.
+
+        Returns
+        -------
+        output_array: DRef
+
+            The DRef containing the scattered data on all workers.
+            If `dst` was provided, this return value is the same as
+            `dst`.  Otherwise, it is the newly allocated space.
+
+        """
+        assert src.shape[0] == self.num_workers
+
+        if not isinstance(src, Tensor):
+            src = _as_Tensor(src)
+
+        if dst is None:
+            dst = self.empty(src.shape[1:], src.dtype)
+
+        src_dref = self.copy_to_worker_0(src)
+        self.scatter_from_worker0(src_dref, dst, in_group)
+
+        return dst
+
+    def scatter_from_worker0(self, from_array: DRef, to_array: DRef, in_group: bool = True) -> None:
+        """Scatter an array from worker-0 to all other workers.
+
+        Parameters
+        ----------
+        src: Union[np.ndarray, Tensor]
+            The array to be scattered.  The first dimension of this
+            array, `src.shape[0]`, must be equal to the number of
+            workers.
+
+        dst: Optional[DRef]
+            The output array.  If None, an array with compatible shape
+            and the same dtype as `src` will be allocated on each
+            worker.
+
+        in_group: bool
+            Whether the scatter operation performs globally or in group as default.
+        """
+        func = self._get_cached_method("runtime.disco.scatter_from_worker0")
+        func(from_array, in_group, to_array)
+
+    def gather_to_worker0(self, from_array: DRef, to_array: DRef, in_group: bool = True) -> None:
+        """Gather an array from all other workers to worker-0.
+
+        Parameters
+        ----------
+        from_array : DRef
+            The array to be gathered from.
+
+        to_array : DRef
+            The array to be gathered to.
+
+        in_group: bool
+            Whether the gather operation performs globally or in group as default.
+        """
+        func = self._get_cached_method("runtime.disco.gather_to_worker0")
+        func(from_array, in_group, to_array)
+
+    def allreduce(
+        self,
+        src: DRef,
+        dst: DRef,
+        op: str = "sum",  # pylint: disable=invalid-name
+        in_group: bool = True,
+    ) -> DRef:
+        """Perform an allreduce operation on an array.
+
+        Parameters
+        ----------
+        array : DRef
+            The array to be reduced.
+
+        op : str = "sum"
+            The reduce operation to be performed. Available options are:
+            - "sum"
+            - "prod"
+            - "min"
+            - "max"
+            - "avg"
+
+        in_group : bool
+            Whether the reduce operation performs globally or in group as default.
+        """
+        if op not in REDUCE_OPS:
+            raise ValueError(f"Unsupported reduce op: {op}. Available ops are: {REDUCE_OPS.keys()}")
+        op = Shape([REDUCE_OPS[op]])
+        func = self._get_cached_method("runtime.disco.allreduce")
+        func(src, op, in_group, dst)
+
+    def allgather(
+        self,
+        src: DRef,
+        dst: DRef,
+        in_group: bool = True,
+    ) -> DRef:
+        """Perform an allgather operation on an array.
+
+        Parameters
+        ----------
+        src : DRef
+            The array to be gathered from.
+
+        dst : DRef
+            The array to be gathered to.
+
+        in_group : bool
+            Whether the reduce operation performs globally or in group as default.
+        """
+        func = self._get_cached_method("runtime.disco.allgather")
+        func(src, in_group, dst)
+
+    def _clear_ipc_memory_pool(self):
+        # Clear the IPC memory allocator when the allocator exists.
+        name = "runtime.disco.cuda_ipc.cuda_ipc_memory_allocator_clear"
+        if get_global_func(name, allow_missing=True) is not None:
+            self.call_packed(self.get_global_func(name))
+
+
+@register_object("runtime.disco.ThreadedSession")
+class ThreadedSession(Session):
+    """A Disco session backed by multi-threading."""
+
+    def __init__(self, num_workers: int, num_groups: int = 1) -> None:
+        """Create a disco session backed by multiple threads in the same process."""
+        self.__init_handle_by_constructor__(
+            _ffi_api.SessionThreaded,  # type: ignore # pylint: disable=no-member
+            num_workers,
+            num_groups,
+        )
+
+
+@register_object("runtime.disco.ProcessSession")
+class ProcessSession(Session):
+    """A Disco session backed by pipe-based multi-processing."""
+
+    def __init__(
+        self,
+        num_workers: int,
+        num_groups: int = 1,
+        entrypoint: str = "tvm.exec.disco_worker",
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.SessionProcess,  # type: ignore # pylint: disable=no-member
+            num_workers,
+            num_groups,
+            "runtime.disco.create_process_pool",
+            entrypoint,
+        )
+        self._configure_structlog()
+
+    def _configure_structlog(self) -> None:
+        try:
+            import structlog  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            return
+
+        root_logger = logging.getLogger()
+        if len(root_logger.handlers) == 1 and isinstance(
+            root_logger.handlers[0].formatter, structlog.stdlib.ProcessorFormatter
+        ):
+            stdlib_formatter = root_logger.handlers[0].formatter
+        else:
+            stdlib_formatter = None
+
+        stdlib_level = root_logger.level
+
+        full_config = (structlog.get_config(), stdlib_formatter, stdlib_level)
+
+        config = pickle.dumps(full_config)
+        func = self.get_global_func("runtime.disco._configure_structlog")
+        func(config, os.getpid())
+
+
+@register_global_func("runtime.disco.create_socket_session_local_workers")
+def _create_socket_session_local_workers(num_workers) -> Session:
+    """Create the local session for each distributed node over socket session."""
+    return ProcessSession(num_workers)
+
+
+@register_object("runtime.disco.SocketSession")
+class SocketSession(Session):
+    """A Disco session backed by socket-based multi-node communication."""
+
+    def __init__(
+        self,
+        num_nodes: int,
+        num_workers_per_node: int,
+        num_groups: int,
+        host: str,
+        port: int,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.SocketSession,  # type: ignore # pylint: disable=no-member
+            num_nodes,
+            num_workers_per_node,
+            num_groups,
+            host,
+            port,
+        )
+
+
+@register_global_func("runtime.disco._configure_structlog")
+def _configure_structlog(pickled_config: bytes, parent_pid: int) -> None:
+    """Configure structlog for all disco workers
+
+    The child processes
+
+    Parameters
+    ----------
+    pickled_config: bytes
+
+        The pickled configuration for structlog
+
+    parent_pid: int
+
+        The PID of the main process.  This is used to restrict the
+    """
+    if os.getpid() == parent_pid:
+        return
+
+    import structlog  # pylint: disable=import-outside-toplevel
+
+    full_config = pickle.loads(pickled_config)
+    structlog_config, stdlib_formatter, stdlib_level = full_config
+
+    root_logger = logging.getLogger()
+
+    root_logger.setLevel(stdlib_level)
+    if stdlib_formatter is not None:
+        handler = logging.StreamHandler()
+        handler.setFormatter(stdlib_formatter)
+        root_logger.addHandler(handler)
+
+    structlog.configure(**structlog_config)
+
+
+@register_global_func("runtime.disco._import_python_module")
+def _import_python_module(module_name: str) -> None:
+    __import__(module_name)
+
+
+REDUCE_OPS = {
+    "sum": 0,
+    "prod": 1,
+    "min": 2,
+    "max": 3,
+    "avg": 4,
+}

python/tvm/runtime/executable.py

@@ -0,0 +1,175 @@
+# 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.
+# pylint: disable=invalid-name, no-member
+
+"""Executable object for TVM Runtime"""
+
+from collections.abc import Callable
+from typing import Any
+
+from tvm_ffi import Function
+
+import tvm
+from tvm.support import utils as _utils
+
+from . import Module
+
+
+class Executable:
+    """The executable object generated by `tvm.compile`."""
+
+    def __init__(self, mod: Module):
+        """Initialize the Executable object."""
+        self.mod: Module = mod
+        self._jitted_mod: Module | None = None
+
+    def __getitem__(self, name: str) -> Function:
+        """Get the Function from the jitted module."""
+        return self.jit().get_function(name, query_imports=True)
+
+    def __call__(self, *args, **kwargs) -> Any:
+        """Call the executable."""
+        return self.jit().main(*args, **kwargs)
+
+    def jit(
+        self,
+        *,
+        fcompile: Callable[[str, list[str], dict[str, Any]], None] | None = None,
+        addons: list[str] | None = None,
+        force_recompile: bool = False,
+        **kwargs,
+    ) -> Module:
+        """Just-in-time compile and link the modules.
+
+        The Executable returned by tvm.compile may not be directly
+        runnable as they may contain CUDA source files and objects that
+        are yet to be compiled and linked.
+        This function helps to create a runtime.Module for these cases.
+
+        Parameters
+        ----------
+        fcompile : function(target, file_list, kwargs), optional
+            The compilation function to use create the final library object during
+
+        addons : list of str, optional
+            Additional object files to link against.
+
+        force_recompile : bool, optional
+            If True, force a recompile of the module.
+
+        kwargs : dict, optional
+            Additional arguments passed to fcompile
+
+        Returns
+        -------
+        rt_mod: tvm.runtime.Module
+            A runnable runtime module that can be passed to VirtualMachine.
+
+        Examples
+        --------
+        .. code:: python
+
+            ex = tvm.compile(mod, target)
+            rt_mod = ex.jit()
+
+        """
+
+        # If the module is already jitted and we don't want to force a recompile,
+        # return the cached module
+        if self._jitted_mod is not None and not force_recompile:
+            return self._jitted_mod
+
+        # TODO(tvm-team): Update runtime.Module interface
+        # to query these properties as bitmask.
+        def _not_runnable(x):
+            return x.kind in ("c", "static_library")
+
+        # pylint:disable = protected-access
+        not_runnable_list = self.mod._collect_from_import_tree(_not_runnable)
+
+        # everything is runnable, directly return mod.
+        if len(not_runnable_list) == 0:
+            return self.mod
+
+        # found source module, or other not runnable modules need to be export and load
+        # TODO(tvm-team): Support runnable but not exportable module.
+        # by collecting the link and allow export_library skip those modules.
+        workspace_dir = _utils.tempdir()
+        dso_path = workspace_dir.relpath("exported.so")
+        self.export_library(dso_path, fcompile=fcompile, addons=addons, **kwargs)
+        self._jitted_mod = tvm.runtime.load_module(dso_path)
+        return self._jitted_mod
+
+    def export_library(
+        self,
+        file_name,
+        *,
+        fcompile=None,
+        addons=None,
+        workspace_dir=None,
+        **kwargs,
+    ):
+        """
+        Export the module and all imported modules into a single device library.
+
+        This function only works on host LLVM modules, other runtime::Module
+        subclasses will work with this API but they must support implement
+        the save and load mechanisms of modules completely including saving
+        from streams and files. This will pack your non-shared library module
+        into a single shared library which can later be loaded by TVM.
+
+        Parameters
+        ----------
+        file_name : str
+            The name of the shared library.
+
+        fcompile : function(target, file_list, kwargs), optional
+            The compilation function to use create the final library object during
+            export.
+
+            For example, when fcompile=_cc.create_shared, or when it is not supplied but
+            module is "llvm," this is used to link all produced artifacts
+            into a final dynamic library.
+
+            This behavior is controlled by the type of object exported.
+            If fcompile has attribute object_format, will compile host library
+            to that format. Otherwise, will use default format "o".
+
+        addons : list of str, optional
+            Additional object files to link against.
+
+        workspace_dir : str, optional
+            The path of the directory used to create the intermediate
+            artifacts when exporting the module.
+            If this is not provided a temporary dir will be created.
+
+        kwargs : dict, optional
+            Additional arguments passed to fcompile
+
+        Returns
+        -------
+        result of fcompile()  : unknown, optional
+            If the compilation function returns an artifact it would be returned via
+            export_library, if any.
+        """
+        return self.mod.export_library(
+            file_name,
+            fcompile=fcompile,
+            addons=addons,
+            workspace_dir=workspace_dir,
+            **kwargs,
+        )

python/tvm/runtime/module.py

@@ -0,0 +1,511 @@
+# 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
+
+# pylint: disable=invalid-name, unused-import, import-outside-toplevel, inconsistent-return-statements
+"""Runtime Module namespace."""
+
+import json
+import os
+import struct
+from collections.abc import Sequence
+
+import numpy as np
+from tvm_ffi import (
+    Module as _Module,
+)
+from tvm_ffi import libinfo as tvm_ffi_libinfo
+from tvm_ffi import (
+    load_module as _load_module,
+)
+from tvm_ffi import (
+    register_object as _register_object,
+)
+from tvm_ffi import (
+    system_lib,
+)
+
+import tvm.libinfo
+from tvm.base import _RUNTIME_ONLY
+
+from . import _ffi_api
+
+
+class BenchmarkResult:
+    """Runtimes from benchmarking"""
+
+    def __init__(self, results: Sequence[float]):
+        """Construct a new BenchmarkResult from a sequence of runtimes.
+
+        Parameters
+        ----------
+        results : Sequence[float]
+            Raw times from benchmarking
+
+        Attributes
+        ----------
+        min : float
+            Minimum runtime in seconds of all results.
+        mean : float
+            Mean runtime in seconds of all results. If py:meth:`Module.time_evaluator` or
+            `benchmark` is called with `number` > 0, then each result is already the mean of a
+            `number` of runtimes, so this becomes the mean of means.
+        median : float
+            Median runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called
+            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
+            this becomes the median of means.
+        max : float
+            Maximum runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called
+            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
+            this becomes the maximum of those means.
+        std : float
+            Standard deviation in seconds of runtimes. If py:meth:`Module.time_evaluator` is called
+            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
+            this becomes the standard deviation of means.
+        results : Sequence[float]
+            The collected runtimes (in seconds). This may be a series of mean runtimes if
+            py:meth:`Module.time_evaluator` or `benchmark` was run with `number` > 1.
+        """
+        self.results = results
+        self.mean = np.mean(self.results)
+        self.std = np.std(self.results)
+        self.median = np.median(self.results)
+        self.min = np.min(self.results)
+        self.max = np.max(self.results)
+
+    def __repr__(self):
+        return (
+            f"BenchmarkResult(min={self.min}, mean={self.mean}, median={self.median}, "
+            f"max={self.max}, std={self.std}, results={self.results})"
+        )
+
+    def __str__(self):
+        return (
+            f"Execution time summary:\n"
+            f"{'mean (ms)':^12} {'median (ms)':^12} {'max (ms)':^12} "
+            f"{'min (ms)':^12} {'std (ms)':^12}\n"
+            f"{self.mean * 1000:^12.4f} {self.median * 1000:^12.4f} {self.max * 1000:^12.4f} "
+            f"{self.min * 1000:^12.4f} {self.std * 1000:^12.4f}"
+            "               "
+        )
+
+
+# override the Module class in ffi.Module
+@_register_object("ffi.Module")
+class Module(_Module):
+    """Runtime Module."""
+
+    def _collect_from_import_tree(self, filter_func):
+        """Helper function to collect modules from the tree matching a filter_func, then return it.
+
+        Parameters
+        ----------
+        filter_func : Callable[[Module], bool]
+            A function which is invoked for each Module discovered in the import tree (including
+            self).
+
+        Returns
+        -------
+        list[Module] :
+            A list of matching Module.
+        """
+        visited, stack, dso_modules = set(), [], []
+        # append root module
+        visited.add(self)
+        stack.append(self)
+        while stack:
+            module = stack.pop()
+            assert module.is_compilation_exportable() or module.is_binary_serializable(), (
+                f"Module {module.kind} should be either dso exportable or binary serializable."
+            )
+
+            if filter_func(module):
+                dso_modules.append(module)
+            for m in module.imports:
+                if m not in visited:
+                    visited.add(m)
+                    stack.append(m)
+        return dso_modules
+
+    def _collect_dso_modules(self):
+        """Collect all compilation exportable modules from the import tree."""
+        return self._collect_from_import_tree(lambda m: m.is_compilation_exportable())
+
+    def export_library(
+        self,
+        file_name,
+        *,
+        fcompile=None,
+        fpack_imports=None,
+        addons=None,
+        workspace_dir=None,
+        **kwargs,
+    ):
+        """
+        Export the module and all imported modules into a single device library.
+
+        This function only works on host LLVM modules, other runtime::Module
+        subclasses will work with this API but they must support implement
+        the save and load mechanisms of modules completely including saving
+        from streams and files. This will pack your non-shared library module
+        into a single shared library which can later be loaded by TVM.
+
+        Parameters
+        ----------
+        file_name : str
+            The name of the shared library.
+
+        fcompile : function(target, file_list, kwargs), optional
+            The compilation function to use create the final library object during
+            export.
+
+            For example, when fcompile=_cc.create_shared, or when it is not supplied but
+            module is "llvm," this is used to link all produced artifacts
+            into a final dynamic library.
+
+            This behavior is controlled by the type of object exported.
+            If fcompile has attribute object_format, will compile host library
+            to that format. Otherwise, will use default format "o".
+
+        fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str,
+                                workspace_dir: str) -> str
+            Function used to pack imported modules from `mod` into a file suitable for passing
+            to fcompile as an input file. The result can be a C source, or an .o object file,
+            or any other file that the fcompile function can handle. The function returns the
+            name of the created file.
+
+            If not provided, the imported modules will be serialized either via packing to an
+            LLVM module, or to a C source file.
+
+        workspace_dir : str, optional
+            The path of the directory used to create the intermediate
+            artifacts when exporting the module.
+            If this is not provided a temporary dir will be created.
+
+        kwargs : dict, optional
+            Additional arguments passed to fcompile
+
+        Returns
+        -------
+        result of fcompile()  : unknown, optional
+            If the compilation function returns an artifact it would be returned via
+            export_library, if any.
+        """
+        # NOTE: this function depends on contrib library features
+        # which are only available in when TVM function is available.
+        if _RUNTIME_ONLY:
+            raise RuntimeError("Cannot call export_library in runtime only mode")
+
+        # Extra dependencies during runtime.
+        from pathlib import Path
+
+        from tvm.contrib import tvmjs as _tvmjs
+        from tvm.support import cc as _cc
+        from tvm.support import tar as _tar
+        from tvm.support import utils as _utils
+
+        if isinstance(file_name, Path):
+            file_name = str(file_name)
+
+        modules = self._collect_dso_modules()
+        if workspace_dir is None:
+            temp = _utils.tempdir()
+            workspace_dir = temp.temp_dir
+        files = addons if addons else []
+        is_system_lib = False
+        has_c_module = False
+        system_lib_prefix = None
+        llvm_target = None
+        global_object_format = "o"
+
+        def get_source_format_from_module(module):
+            for fmt in module.get_write_formats():
+                if fmt in ["c", "cc", "cpp", "cu"]:
+                    return fmt
+            raise ValueError(f"Module {module.kind} does not exporting to c, cc, cpp or cu.")
+
+        for index, module in enumerate(modules):
+            if fcompile is not None and hasattr(fcompile, "object_format"):
+                if module.kind == "c":
+                    object_format = get_source_format_from_module(module)
+                    has_c_module = True
+                else:
+                    global_object_format = object_format = fcompile.object_format
+            else:
+                if module.kind == "c":
+                    if len(module.get_write_formats()) > 0:
+                        object_format = get_source_format_from_module(module)
+                    else:
+                        object_format = "c"
+                    if "cc" in kwargs:
+                        if kwargs["cc"] == "nvcc":
+                            object_format = "cu"
+                    has_c_module = True
+                else:
+                    assert module.is_compilation_exportable()
+                    global_object_format = object_format = "o"
+
+            path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}")
+            module.write_to_file(path_obj)
+            files.append(path_obj)
+            if module.kind == "llvm":
+                is_system_lib = module.get_function("__tvm_is_system_module")()
+                llvm_target = module.get_function("_get_target_string")()
+                system_lib_prefix = module.get_function("__tvm_get_system_lib_prefix")()
+
+        if not fcompile:
+            if file_name.endswith(".tar"):
+                fcompile = _tar.tar
+            elif file_name.endswith(".wasm"):
+                fcompile = _tvmjs.create_tvmjs_wasm
+            else:
+                fcompile = _cc.create_shared
+
+        if llvm_target is None and hasattr(fcompile, "get_target_triple"):
+            triple = fcompile.get_target_triple()
+            assert triple, "Target triple should not be empty"
+            llvm_target = json.dumps({"kind": "llvm", "mtriple": triple.strip()})
+
+        if getattr(fcompile, "need_system_lib", False) and not is_system_lib:
+            raise ValueError(f"{fcompile!s} need --system-lib option")
+
+        if self.imports:
+            pack_lib_prefix = system_lib_prefix if system_lib_prefix else ""
+
+            if fpack_imports is not None:
+                path_out = fpack_imports(self, is_system_lib, pack_lib_prefix, workspace_dir)
+                files.append(path_out)
+            elif _ffi_api.RuntimeEnabled("llvm") and llvm_target:
+                path_obj = os.path.join(
+                    workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}"
+                )
+                m = _ffi_api.ModulePackImportsToLLVM(
+                    self, is_system_lib, llvm_target, pack_lib_prefix
+                )
+                m.write_to_file(path_obj)
+                files.append(path_obj)
+            else:
+                path_cc = os.path.join(workspace_dir, f"{pack_lib_prefix}devc.c")
+                with open(path_cc, "w") as f:
+                    f.write(_ffi_api.ModulePackImportsToC(self, is_system_lib, pack_lib_prefix))
+                files.append(path_cc)
+
+        # The imports could contain a c module but the object format could be tar
+        # Thus, it would not recognize the following include paths as options
+        # which are there assuming a c compiler is the fcompile.
+        if has_c_module and not file_name.endswith(".tar"):
+            options = []
+            if "options" in kwargs:
+                opts = kwargs["options"]
+                options = opts if isinstance(opts, list | tuple) else [opts]
+            default_include_paths = [
+                tvm.libinfo.find_include_path(),
+                tvm_ffi_libinfo.find_include_path(),
+                tvm_ffi_libinfo.find_dlpack_include_path(),
+            ]
+            opts = options + ["-I" + path for path in default_include_paths]
+            kwargs.update({"options": opts})
+
+        return fcompile(file_name, files, **kwargs)
+
+    def time_evaluator(
+        self,
+        func_name,
+        dev,
+        number=10,
+        repeat=1,
+        min_repeat_ms=0,
+        limit_zero_time_iterations=100,
+        cooldown_interval_ms=0,
+        repeats_to_cooldown=1,
+        cache_flush_bytes=0,
+        f_preproc="",
+    ):
+        """Get an evaluator that measures time cost of running function.
+
+        Parameters
+        ----------
+        func_name: str
+            The name of the function in the module.
+
+        dev: Device
+            The device we should run this function on.
+
+        number: int
+            The number of times to run this function for taking average.
+            We call these runs as one `repeat` of measurement.
+
+        repeat: int, optional
+            The number of times to repeat the measurement.
+            In total, the function will be invoked (1 + number x repeat) times,
+            where the first one is warm up and will be discarded.
+            The returned result contains `repeat` costs,
+            each of which is an average of `number` costs.
+
+        min_repeat_ms: int, optional
+            The minimum duration of one `repeat` in milliseconds.
+            By default, one `repeat` contains `number` runs. If this parameter is set,
+            the parameters `number` will be dynamically adjusted to meet the
+            minimum duration requirement of one `repeat`.
+            i.e., When the run time of one `repeat` falls below this time, the `number` parameter
+            will be automatically increased.
+
+        limit_zero_time_iterations: int, optional
+            The maximum number of repeats when measured time is equal to 0.
+            It helps to avoid hanging during measurements.
+
+        cooldown_interval_ms: int, optional
+            The cooldown interval in milliseconds between the number of repeats defined by
+            `repeats_to_cooldown`.
+
+        repeats_to_cooldown: int, optional
+            The number of repeats before the cooldown is activated.
+
+        cache_flush_bytes: int, optional
+            The number of bytes to flush from the cache before each repeat.
+
+        f_preproc: str, optional
+            The preprocess function name we want to execute before executing the time evaluator.
+
+        Note
+        ----
+        The function will be invoked  (1 + number x repeat) times,
+        with the first call discarded in case there is lazy initialization.
+
+        Returns
+        -------
+        ftimer : function
+            The function that takes same argument as func and returns a BenchmarkResult.
+            The ProfileResult reports `repeat` time costs in seconds.
+        """
+        try:
+            feval = _ffi_api.RPCTimeEvaluator(
+                self,
+                func_name,
+                dev.dlpack_device_type(),
+                dev.index,
+                number,
+                repeat,
+                min_repeat_ms,
+                limit_zero_time_iterations,
+                cooldown_interval_ms,
+                repeats_to_cooldown,
+                cache_flush_bytes,
+                f_preproc,
+            )
+
+            def evaluator(*args):
+                """Internal wrapped evaluator."""
+                # Wrap feval so we can add more stats in future.
+                blob = feval(*args)
+                fmt = "@" + ("d" * repeat)
+                results = struct.unpack(fmt, blob)
+                return BenchmarkResult(results)
+
+            return evaluator
+        except NameError:
+            raise NameError("time_evaluator is only supported when RPC is enabled")
+
+
+def load_module(path):
+    """Load module from file.
+
+    Parameters
+    ----------
+    path : str
+        The path to the module file.
+
+    Returns
+    -------
+    module : runtime.Module
+        The loaded module
+
+    Note
+    ----
+    This function will automatically call
+    cc.create_shared if the path is in format .o or .tar
+    """
+    if os.path.isfile(path):
+        path = os.path.realpath(path)
+    else:
+        raise ValueError(f"cannot find file {path}")
+
+    # High level handling for .o and .tar file.
+    # We support this to be consistent with RPC module load.
+    if path.endswith(".o"):
+        # Extra dependencies during runtime.
+        from tvm.support import cc as _cc
+
+        _cc.create_shared(path + ".so", path)
+        path += ".so"
+    elif path.endswith(".tar"):
+        # Extra dependencies during runtime.
+        from tvm.support import cc as _cc
+        from tvm.support import tar as _tar
+        from tvm.support import utils as _utils
+
+        tar_temp = _utils.tempdir(custom_path=path.replace(".tar", ""))
+        _tar.untar(path, tar_temp.temp_dir)
+        files = [tar_temp.relpath(x) for x in tar_temp.listdir()]
+        _cc.create_shared(path + ".so", files)
+        path += ".so"
+    # Redirect to the load API
+    return _load_module(path)
+
+
+def load_static_library(path, func_names):
+    """Load the .o library at path which implements functions with func_names.
+    Unlike the generic load_module the result will remain as a static_library
+    and will not be relinked on-the-fly into a .so library."""
+    return _ffi_api.ModuleLoadStaticLibrary(path, func_names)
+
+
+def enabled(target):
+    """Whether module runtime is enabled for target
+
+    Parameters
+    ----------
+    target : str or Dict[str, Any] or tvm.target.Target
+        The target device type.
+
+    Returns
+    -------
+    enabled : bool
+        Whether runtime is enabled.
+
+    Examples
+    --------
+    The following code checks if gpu is enabled.
+
+    >>> tvm.runtime.enabled("gpu")
+    """
+    if isinstance(target, dict):
+        target = target.get("kind", "")
+    elif hasattr(target, "kind"):
+        target = target.kind.name
+    return _ffi_api.RuntimeEnabled(target)
+
+
+def num_threads() -> int:
+    """Get the number of threads in use by the TVM runtime.
+
+    Returns
+    -------
+    int
+        Number of threads in use.
+    """
+    return _ffi_api.NumThreads()

python/tvm/runtime/object_generic.py

@@ -0,0 +1,70 @@
+# 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
+"""Common implementation of object generic related logic"""
+
+# pylint: disable=unused-import, invalid-name
+from tvm_ffi import ObjectConvertible
+
+from . import _ffi_node_api
+
+
+def _scalar_type_inference(value):
+    if hasattr(value, "dtype"):
+        return str(value.dtype)
+    elif isinstance(value, bool):
+        return "bool"
+    elif isinstance(value, float):
+        # We intentionally prefer convert the float to float32 since it's more common in DL.
+        if -3.40282347e38 <= value <= 3.40282347e38:
+            return "float32"
+        else:
+            return "float64"
+    elif isinstance(value, int):
+        # We intentionally prefer convert the python int to int32 since it's more common in DL.
+        if -2147483648 <= value <= 2147483647:
+            return "int32"
+        else:
+            return "int64"
+    else:
+        raise NotImplementedError(f"Cannot automatically inference the type. value={value}")
+
+
+def const(value, dtype=None, span=None):
+    """construct a constant
+
+    Parameters
+    ----------
+    value : number
+        The content of the constant number.
+
+    dtype : str or None, optional
+        The data type.
+
+    span : Optional[Span]
+        The location of the constant value in the source.
+
+    Returns
+    -------
+    const_val: tvm.Expr
+        The result expression.
+    """
+    if dtype is None:
+        dtype = _scalar_type_inference(value)
+    if dtype == "uint64" and value >= (1 << 63):
+        return _ffi_node_api.LargeUIntImm(dtype, value & ((1 << 32) - 1), value >> 32, span)
+    return _ffi_node_api._const(value, dtype, span)

python/tvm/runtime/params.py

@@ -0,0 +1,111 @@
+# 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.
+# pylint: disable=invalid-name
+"""Helper utility to save and load parameter dicts."""
+
+from . import Tensor, _ffi_api, tensor
+
+
+def _to_tensor(params):
+    transformed = {}
+
+    for k, v in params.items():
+        if not isinstance(v, Tensor):
+            transformed[k] = tensor(v)
+        else:
+            transformed[k] = v
+
+    return transformed
+
+
+def save_param_dict(params):
+    """Save parameter dictionary to binary bytes.
+
+    The result binary bytes can be loaded by the
+    GraphModule with API "load_params".
+
+    Parameters
+    ----------
+    params : dict of str to Tensor
+        The parameter dictionary.
+
+    Returns
+    -------
+    param_bytes: bytearray
+        Serialized parameters.
+
+    Examples
+    --------
+    .. code-block:: python
+
+       # set up the parameter dict
+       params = {"param0": arr0, "param1": arr1}
+       # save the parameters as byte array
+       param_bytes = tvm.runtime.save_param_dict(params)
+       # We can serialize the param_bytes and load it back later.
+       # Pass in byte array to module to directly set parameters
+       tvm.runtime.load_param_dict(param_bytes)
+    """
+    return _ffi_api.SaveParams(_to_tensor(params))
+
+
+def save_param_dict_to_file(params, path):
+    """Save parameter dictionary to file.
+
+    Parameters
+    ----------
+    params : dict of str to Tensor
+        The parameter dictionary.
+
+    path: str
+        The path to the parameter file.
+    """
+    return _ffi_api.SaveParamsToFile(_to_tensor(params), path)
+
+
+def load_param_dict(param_bytes):
+    """Load parameter dictionary from binary bytes.
+
+    Parameters
+    ----------
+    param_bytes: bytearray
+        Serialized parameters.
+
+    Returns
+    -------
+    params : dict of str to Tensor
+        The parameter dictionary.
+    """
+    if isinstance(param_bytes, bytes | str):
+        param_bytes = bytearray(param_bytes)
+    return _ffi_api.LoadParams(param_bytes)
+
+
+def load_param_dict_from_file(path):
+    """Load parameter dictionary from file.
+
+    Parameters
+    ----------
+    path: str
+        The path to the parameter file to load from.
+
+    Returns
+    -------
+    params : dict of str to Tensor
+        The parameter dictionary.
+    """
+    return _ffi_api.LoadParamsFromFile(path)

python/tvm/runtime/script_printer.py

@@ -0,0 +1,418 @@
+# 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.
+"""Configuration of TVMScript printer"""
+
+import os
+from collections.abc import Sequence
+
+from tvm_ffi import get_global_func, register_object
+from tvm_ffi.access_path import AccessPath
+
+from tvm.runtime import Object
+
+from . import _ffi_node_api
+
+
+@register_object("script.PrinterConfig")
+class PrinterConfig(Object):
+    """Configuration of TVMScript printer"""
+
+    binding_names: Sequence[str]
+    show_meta: bool
+    ir_prefix: str
+    module_alias: str
+    buffer_dtype: str
+    int_dtype: str
+    float_dtype: str
+    verbose_expr: bool
+    indent_spaces: int
+    print_line_numbers: bool
+    num_context_lines: int
+    syntax_sugar: bool
+    show_object_address: bool
+    extra_config: dict
+    path_to_underline: list[AccessPath] | None
+    path_to_annotate: dict[AccessPath, str] | None
+    obj_to_underline: list[AccessPath] | None
+    obj_to_annotate: dict[AccessPath, str] | None
+
+    def __init__(
+        self,
+        *,
+        name: str | None = None,
+        show_meta: bool = False,
+        ir_prefix: str = "I",
+        module_alias: str = "cls",
+        buffer_dtype: str = "float32",
+        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 | None = None,
+        syntax_sugar: bool = True,
+        show_object_address: bool = False,
+        show_all_ty: bool = True,
+        extra_config: dict | None = None,
+        path_to_underline: list[AccessPath] | None = None,
+        path_to_annotate: dict[AccessPath, str] | None = None,
+        obj_to_underline: list[Object] | None = None,
+        obj_to_annotate: dict[Object, str] | None = None,
+    ) -> None:
+        if num_context_lines is None:
+            num_context_lines = -1
+        cfg: dict = {
+            "show_meta": show_meta,
+            "ir_prefix": ir_prefix,
+            "module_alias": module_alias,
+            "buffer_dtype": buffer_dtype,
+            "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,
+            "path_to_underline": path_to_underline,
+            "path_to_annotate": path_to_annotate,
+            "obj_to_underline": obj_to_underline,
+            "obj_to_annotate": obj_to_annotate,
+            # Dialect-specific config via dotted keys in extra_config
+            "relax.show_all_ty": show_all_ty,
+        }
+
+        if name is not None:
+            cfg["name"] = name
+        if extra_config is not None:
+            cfg["extra_config"] = extra_config
+        self.__init_handle_by_constructor__(
+            _ffi_node_api.PrinterConfig,
+            cfg,  # type: ignore # pylint: disable=no-member
+        )
+
+
+def _script(obj: Object, config: PrinterConfig) -> str:
+    return _ffi_node_api.TVMScriptPrinterScript(obj, config)  # type: ignore # pylint: disable=no-member
+
+
+def _relax_script(obj: Object, config: PrinterConfig) -> str:
+    func = get_global_func("script.printer.ReprPrintRelax")
+    return func(obj, config)
+
+
+class Scriptable:
+    """A base class that enables the script() and show() method."""
+
+    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,
+        path_to_underline: list[AccessPath] | None = None,
+        path_to_annotate: dict[AccessPath, str] | None = None,
+        obj_to_underline: list[Object] | None = None,
+        obj_to_annotate: dict[Object, str] | None = None,
+    ) -> str:
+        """Print TVM IR into TVMScript text format
+
+        Parameters
+        ----------
+        name : Optional[str] = None
+            The name of the object
+        show_meta : bool = False
+            Whether to print the meta data of the object
+        ir_prefix : str = "I"
+            The prefix of AST nodes from tvm.ir
+        module_alias : str = "cls"
+            The alias of the current module at cross-function call,
+            Directly use module name if it's empty.
+        int_dtype : str = "int32"
+            The default data type of integer
+        float_dtype : str = "void"
+            The default data type of float
+        verbose_expr : bool = False
+            Whether to print the detailed definition of each variable in the expression
+        indent_spaces : int = 4
+            The number of spaces for indentation
+        print_line_numbers : bool = False
+            Whether to print line numbers
+        num_context_lines : int = -1
+            The number of lines of context to print before and after the line to underline.
+        syntax_sugar: bool = True
+            Whether to output with syntax sugar, set false for complete printing.
+        show_object_address: bool = False
+            Whether to include the object's address as part of the TVMScript name
+        show_all_ty: bool = True
+            If True (default), annotate all variable bindings with the struct
+            info of that variable.  If False, only add annotations where
+            required for unambiguous round-trip of Relax -> TVMScript -> Relax.
+        extra_config : Optional[dict] = None
+            Dialect-specific configuration passed through to PrinterConfig.extra_config.
+            Keys are conventionally namespaced as "<dialect>.<knob>", e.g.
+            ``{"tirx.prefix": "T"}``.
+        path_to_underline : Optional[List[AccessPath]] = None
+            Object path to be underlined
+        path_to_annotate : Optional[Dict[AccessPath, str]] = None
+            Object path to be annotated
+        obj_to_underline : Optional[List[Object]] = None
+            Object to be underlined
+        obj_to_annotate : Optional[Dict[Object, str]] = None
+            Object to be annotated
+
+        Returns
+        -------
+        script : str
+            The TVM Script of the given TVM IR
+
+        """
+        # Auto-switch to tirx (`T`/`tirx`) flavor only when explicitly
+        # printing a PrimFunc / IRModule that has no s_tir-tagged content.
+        # Free objects (Buffer, BufferRegion, ...) keep the default `T`/`tir`
+        # flavor -- they have no enclosing function to indicate tirx vs s_tir.
+        merged_extra: dict = {}
+        if extra_config is not None:
+            merged_extra.update(extra_config)
+
+        # Only auto-switch if the caller has not already set a tirx.prefix override.
+        if "tirx.prefix" not in merged_extra:
+            from tvm.ir import IRModule  # pylint: disable=import-outside-toplevel
+            from tvm.tirx import PrimFunc  # pylint: disable=import-outside-toplevel
+
+            switch_to_tirx = False
+            if isinstance(self, PrimFunc):
+                attrs = getattr(self, "attrs", None)
+                if attrs is None or not attrs.get("s_tir", False):
+                    switch_to_tirx = True
+            elif isinstance(self, IRModule):
+                any_prim = False
+                any_s_tir = False
+                for _, base_func in self.functions.items():
+                    if isinstance(base_func, PrimFunc):
+                        any_prim = True
+                        if getattr(base_func, "attrs", None) and base_func.attrs.get(
+                            "s_tir", False
+                        ):
+                            any_s_tir = True
+                            break
+                if any_prim and not any_s_tir:
+                    switch_to_tirx = True
+            if switch_to_tirx:
+                merged_extra["tirx.prefix"] = "T"
+
+        return _script(
+            self,
+            PrinterConfig(
+                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=merged_extra if merged_extra else None,
+                path_to_underline=path_to_underline,
+                path_to_annotate=path_to_annotate,
+                obj_to_underline=obj_to_underline,
+                obj_to_annotate=obj_to_annotate,
+            ),
+        )
+
+    def _relax_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,
+        extra_config: dict | None = None,
+        path_to_underline: list[AccessPath] | None = None,
+        path_to_annotate: dict[AccessPath, str] | None = None,
+        obj_to_underline: list[Object] | None = None,
+        obj_to_annotate: dict[Object, str] | None = None,
+    ) -> str:
+        return _relax_script(
+            self,
+            PrinterConfig(
+                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,
+                extra_config=extra_config,
+                path_to_underline=path_to_underline,
+                path_to_annotate=path_to_annotate,
+                obj_to_underline=obj_to_underline,
+                obj_to_annotate=obj_to_annotate,
+            ),
+        )
+
+    def show(
+        self,
+        style: str | None = None,
+        black_format: bool | None = None,
+        *,
+        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,
+        path_to_underline: list[AccessPath] | None = None,
+        path_to_annotate: dict[AccessPath, str] | None = None,
+        obj_to_underline: list[Object] | None = None,
+        obj_to_annotate: dict[Object, str] | None = None,
+    ) -> None:
+        """A sugar for print highlighted TVM script.
+
+        Parameters
+        ----------
+        style : str, optional
+            Pygmentize printing style, auto-detected if None.  See
+            `tvm.script.highlight.cprint` for more details.
+
+        black_format: Optional[bool]
+
+            If true, use the formatter Black to format the TVMScript.
+            If false, do not apply the auto-formatter.
+
+            If None (default), determine the behavior based on the
+            environment variable "TVM_BLACK_FORMAT".  If this
+            environment variable is unset, set to the empty string, or
+            set to the integer zero, black auto-formatting will be
+            disabled.  If the environment variable is set to a
+            non-zero integer, black auto-formatting will be enabled.
+
+            Note that the "TVM_BLACK_FORMAT" environment variable only
+            applies to the `.show()` method, and not the underlying
+            `.script()` method.  The `.show()` method is intended for
+            human-readable output based on individual user
+            preferences, while the `.script()` method is intended to
+            provided a consistent output regardless of environment.
+
+        name : Optional[str] = None
+            The name of the object
+        show_meta : bool = False
+            Whether to print the meta data of the object
+        ir_prefix : str = "I"
+            The prefix of AST nodes from tvm.ir
+        module_alias : str = "cls"
+            The alias of the current module at cross-function call,
+            Directly use module name if it's empty.
+        int_dtype : str = "int32"
+            The default data type of integer
+        float_dtype : str = "void"
+            The default data type of float
+        verbose_expr : bool = False
+            Whether to print the detailed definition of each variable in the expression
+        indent_spaces : int = 4
+            The number of spaces for indentation
+        print_line_numbers : bool = False
+            Whether to print line numbers
+        num_context_lines : int = -1
+            The number of lines of context to print before and after the line to underline.
+        syntax_sugar: bool = True
+            Whether to output with syntax sugar, set false for complete printing.
+        show_object_address: bool = False
+            Whether to include the object's address as part of the TVMScript name
+        show_all_ty: bool = True
+            If True (default), annotate all variable bindings with the struct
+            info of that variable.  If False, only add annotations where
+            required for unambiguous round-trip of Relax -> TVMScript -> Relax.
+        extra_config : Optional[dict] = None
+            Dialect-specific configuration passed through to PrinterConfig.extra_config.
+        path_to_underline : Optional[List[AccessPath]] = None
+            Object path to be underlined
+        path_to_annotate : Optional[Dict[AccessPath, str]] = None
+            Object path to be annotated
+        obj_to_underline : Optional[List[Object]] = None
+            Object to be underlined
+        obj_to_annotate : Optional[Dict[Object, str]] = None
+            Object to be annotated
+
+        """
+        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)
+
+        cprint(
+            self.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,
+                path_to_underline=path_to_underline,
+                path_to_annotate=path_to_annotate,
+                obj_to_underline=obj_to_underline,
+                obj_to_annotate=obj_to_annotate,
+            ),
+            style=style,
+            black_format=black_format,
+        )

python/tvm/runtime/vm.py

@@ -0,0 +1,477 @@
+# 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.
+# pylint: disable=invalid-name, redefined-builtin, no-else-return, consider-using-dict-items
+# ruff: noqa: RUF005
+"""The Relax virtual machine."""
+
+from collections.abc import Callable
+from enum import IntEnum
+from numbers import Integral, Number
+from typing import Any
+
+import numpy as np  # type: ignore
+from tvm_ffi import Function, register_global_func
+
+import tvm
+from tvm.runtime import Device, Object
+
+from ..rpc.base import RPC_SESS_MASK
+
+
+class VMInstrumentReturnKind(IntEnum):
+    NO_OP = 0
+    # skip the following call, only valid in before
+    SKIP_RUN = 1
+
+
+class VirtualMachine:
+    """Relax VM runtime."""
+
+    NAIVE_ALLOCATOR = 1
+    POOLED_ALLOCATOR = 2
+
+    def __init__(
+        self,
+        rt_mod: tvm.runtime.Module | tvm.runtime.Executable,
+        device: Device | list[Device],
+        memory_cfg: str | dict[Device, str] | None = None,
+    ) -> None:
+        """
+        Construct a VirtualMachine wrapper object.
+
+        Parameters
+        ----------
+        rt_mod: Union[tvm.runtime.Module, tvm.runtime.Executable]
+            Runtime module exported by the result of build.
+
+        device : Union[Device, List[Device]]
+            The device to deploy the module.
+
+        memory_cfg : Optional[Union[str, Dict[Device, str]]]
+            Config the type of memory allocator. The allocator type can be ["naive",
+            "pooled"]. If memory_cfg is None, all devices will use pooled allocator
+            by default. If memory_cfg is string, all devices will use the specified
+            allocator type. If memory_cfg is a dict, each device uses the allocator
+            type specified in the dict, or pooled allocator if not specified in the
+            dict.
+        """
+        if not isinstance(rt_mod, tvm.runtime.Module):
+            if isinstance(rt_mod, tvm.runtime.Executable):
+                rt_mod = rt_mod.jit()
+            else:
+                raise ValueError("Expect the rt_mod to be an runtime.Module")
+
+        self.module = rt_mod["vm_load_executable"]()
+        self._invoke_closure = self.module["invoke_closure"]
+        self._save_function = self.module["save_function"]
+        self._set_input = self.module["set_input"]
+        self._invoke_stateful = self.module["invoke_stateful"]
+        self._get_output = self.module["get_output"]
+        self._get_output_arity = self.module["get_output_arity"]
+        self._get_function_arity = self.module["get_function_arity"]
+        self._get_function_param_name = self.module["get_function_param_name"]
+        self._set_instrument = self.module["set_instrument"]
+        self._setup_device(device, memory_cfg)
+
+    def _setup_device(self, dev: Device, memory_cfg: str | dict[Device, str]) -> None:
+        """init devices and allocators."""
+        devs = dev
+        if not isinstance(dev, list | tuple):
+            if not isinstance(dev, tvm.runtime.Device):
+                raise TypeError("dev is expected to be Device or List[Device]")
+            devs = [dev]
+
+        # CPU is required for executing shape functions
+        if devs[-1].dlpack_device_type() % RPC_SESS_MASK != tvm.cpu().dlpack_device_type():
+            devs.append(tvm.cpu())
+
+        default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
+        if memory_cfg is None:
+            memory_cfg = {}
+        elif isinstance(memory_cfg, str):
+            assert memory_cfg in ["naive", "pooled"]
+            if memory_cfg == "naive":
+                default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR
+            memory_cfg = {}
+        elif not isinstance(memory_cfg, dict):
+            raise TypeError(
+                "memory_cfg is expected be string or dictionary, "
+                + f"but received {type(memory_cfg)}"
+            )
+        init_args = []
+        for device in devs:
+            init_args.append(device.dlpack_device_type() % RPC_SESS_MASK)
+            init_args.append(device.index)
+            alloc_type = memory_cfg[device] if device in memory_cfg else default_alloc_type
+            init_args.append(alloc_type)
+        self.module["vm_initialization"](*init_args)
+
+    def __getitem__(self, key: str) -> Function:
+        return self.module[key]
+
+    def invoke_closure(self, closure: Object, *args: Any) -> Object:
+        """Invoke a closure.
+
+        Parameters
+        ----------
+        closure : Object
+            The VMClosure Object.
+
+        args : list[tvm.runtime.Tensor] or list[np.ndarray]
+            The arguments to the closure.
+
+        Returns
+        -------
+        result : Object
+            The output.
+        """
+        return self._invoke_closure(closure, *args)
+
+    def save_function(
+        self,
+        func_name: str,
+        saved_name: str,
+        *args: list[Any],
+        include_return: bool = True,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        """
+        Convenience function. Takes a function from the module and saves
+        a `Function` that, when called, will invoke the function with the given arguments.
+        The `Function` can be accessed from the module using `saved_name`.
+        This is included to facilitate timing trials:
+        Invoking the returned `Function` will have less overhead from dictionary lookups
+        than normally running through the VM.
+
+        If the saved name is taken, it can be overridden, though it cannot override
+        the name of a function defined in the Relax source.
+
+        This is really creating a closure, but the function has a different name
+        to avoid confusion with `invoke_closure` (they are not meant to be used together).
+
+        Parameters
+        ----------
+        func_name : str
+            The function that should be packaged up.
+
+        saved_name : str
+            The name that the resulting closure should be saved under.
+
+        include_return : bool
+            Whether the saved Function should return its output.
+            If timing over RPC, it may not be desirable to send output
+            between machines.
+
+        args : List[Any]
+            The arguments to package up with the function.
+
+        kwargs : Dict[str, Any]
+            Any named arguments to package up with the function
+        """
+        cargs: list[Any] = []
+        if kwargs:
+            args = self._convert_func_named_args(func_name, args, **kwargs)
+        for arg in args:
+            self._convert(arg, cargs)
+        self._save_function(func_name, saved_name, int(include_return), *cargs)
+
+    def _convert(self, arg: Any, cargs: list) -> None:
+        """helper function to convert arguments to vm function."""
+
+        def _gettype(arg):
+            if isinstance(arg, np.float16):
+                return "float16"
+            elif isinstance(arg, Integral | bool):
+                return "int32"
+            else:
+                return "float32"
+
+        if isinstance(arg, Object):
+            cargs.append(arg)
+        elif isinstance(arg, np.ndarray):
+            nd_arr = tvm.runtime.tensor(arg, device=tvm.cpu(0))
+            cargs.append(nd_arr)
+        elif isinstance(arg, tvm.runtime.Tensor):
+            cargs.append(arg)
+        elif isinstance(arg, tuple | list):
+            field_args: list[Any] = []
+            for field in arg:
+                self._convert(field, field_args)
+            cargs.append(tuple(field_args))
+        elif isinstance(arg, Number | bool):
+            dtype = _gettype(arg)
+            value = tvm.runtime.tensor(np.array(arg, dtype=dtype), device=tvm.cpu(0))
+            cargs.append(value)
+        elif isinstance(arg, str):
+            cargs.append(arg)
+        else:
+            raise TypeError(f"Unsupported type: {type(arg)}")
+
+    def _convert_func_named_args(self, func_name: str, args: Any, **kwargs: Any) -> Any:
+        """
+        Takes named function parameters and returns a list of those needed,
+        in the order they should appear
+        """
+        # kwargs can be a super set of the required function parameters.
+        # We only find the ones that are needed.
+        func_arity = self._get_function_arity(func_name)
+        func_params = [self._get_function_param_name(func_name, i) for i in range(func_arity)]
+        new_args = [None] * len(func_params)
+        cnt = 0
+        for k in kwargs:
+            if k in func_params:
+                idx = func_params.index(k)
+                new_args[idx] = kwargs[k]
+                cnt += 1
+            else:
+                print(f'Warning: Keyword argument "{k}" is unused in {func_name}')
+        assert len(args) + cnt == len(func_params)
+        idx = 0
+        for i, arg in enumerate(new_args):
+            if arg is None:
+                new_args[i] = args[idx]
+                idx += 1
+        return new_args
+
+    def set_input(self, func_name: str, *args: Any, **kwargs: Any) -> None:
+        """Set the inputs to a function.
+        This interface works when using VM over RPC by internally converting Tensor in
+        the arguments to DLTensor, which is supported in RPC where remote could only
+        have a minimal C runtime.
+
+        Note: If `set_input` is used, the function *must* be called using `invoke_stateful`
+        and the results must be obtained using `get_outputs`.
+
+        Parameters
+        ----------
+        func_name : str
+            The name of the function.
+        args: List[tvm.runtime.Tensor] or List[np.ndarray]
+            The arguments to the function.
+        kwargs: dict of str to tvm.runtime.Tensor or np.ndarray
+            Named arguments to the function.
+        """
+        cargs: list[Any] = []
+
+        if kwargs:
+            args = self._convert_func_named_args(func_name, args, **kwargs)
+
+        for arg in args:
+            self._convert(arg, cargs)
+
+        self._set_input(func_name, *cargs)
+
+    def invoke_stateful(self, func_name: str) -> None:
+        """
+        Call the named function from the VM module using the arguments set using `set_input`.
+        It is an error to call `invoke_stateful` without using `set_input` first
+        (even if it's to set 0 inputs); conversely, if `set_input` has been called,
+        it is an error to call the function without using `invoke_stateful`.
+
+        The results of the call can be obtained by calling `get_outputs`.
+
+        Parameters
+        ----------
+        func_name: str
+            The name of the function to call.
+        """
+        self._invoke_stateful(func_name)
+
+    def get_outputs(self, func_name: str) -> tvm.Object | tuple[Any]:
+        """
+        Get the value output by the function by the given name
+        after a call of `invoke_stateful`.
+
+        It is an error to call this function without first calling `invoke_stateful`.
+
+        Parameters
+        ----------
+        func_name: str
+            The name of the function whose output should be fetched.
+
+        Returns
+        -------
+        ret: Union[tvm.Object, Tuple[Any]]
+            The result of the earlier call to the function via `invoke_stateful`.
+            If the result is a tuple, it returns a list of the fields.
+            The fields are potentially also tuples, so these can be arbitrily nested.
+        """
+
+        # to deal with potentially nested tuples, we need to query for arity recursively
+        def get_output_rec(func_name, *idx):
+            arity = self._get_output_arity(func_name, *idx)
+            if arity == -1:
+                return self._get_output(func_name, *idx)
+            # otherwise we need to specify more indices
+            idx_list = list(idx)
+            return tuple(get_output_rec(func_name, *(idx_list + [i])) for i in range(arity))
+
+        return get_output_rec(func_name)
+
+    def set_instrument(self, instrument: Function) -> None:
+        """Set an instrumentation function.
+
+        If instrument is present, the function will be called
+        before/after each Call instruction. The function have
+        the following signature:
+
+        .. code:: python
+
+            def instrument(
+                func: Union[VMClosure, Function],
+                func_symbol: str,
+                before_run: bool,
+                ret_value: any,
+                *args) -> bool:
+                pass
+
+        The instrument takes the following parameters:
+        - func: function object to be called.
+        - func_symbol: the symbol name of the function.
+        - before_run: whether it is before or after call.
+        - ret_value: the return value of the call, only valid after run.
+        - args: the arguments being passed to call.
+
+        The instrument function can choose an integer,
+        which corresponds to action direction for the
+        following run. See VMInstrumentReturnKind for
+        more details.
+
+        Parameters
+        ----------
+        instrument: tvm_ffi.Function
+            A instrumentation function that get invoked every VM call instr.
+
+        See Also
+        --------
+        VMInstrumentReturnKind: the possible return values in VM.
+        """
+        self._set_instrument(instrument)
+
+    def time_evaluator(
+        self,
+        func_name: str,
+        dev: Device,
+        number: int = 10,
+        repeat: int = 1,
+        min_repeat_ms: int = 0,
+        cooldown_interval_ms: int = 0,
+        repeats_to_cooldown: int = 1,
+        f_preproc: str = "",
+    ) -> Callable[..., tvm.runtime.module.BenchmarkResult]:
+        """
+        Returns an evaluator that times a function in the module.
+        This follows the same convention as time_evaluator in tvm.runtime.module.
+        This can be used in combination with save_function() so that the
+        timings avoid extra dictionary lookups.
+
+        Parameters
+        ----------
+        func_name: str
+            The name of the function in the module.
+
+        dev: Device
+            The device we should run this function on.
+
+        number: int
+            The number of times to run this function for taking average.
+            We call these runs as one `repeat` of measurement.
+
+        repeat: int, optional
+            The number of times to repeat the measurement.
+            In total, the function will be invoked (1 + number x repeat) times,
+            where the first one is warm up and will be discarded.
+            The returned result contains `repeat` costs,
+            each of which is an average of `number` costs.
+
+        min_repeat_ms: int, optional
+            The minimum duration of one `repeat` in milliseconds.
+            By default, one `repeat` contains `number` runs. If this parameter is set,
+            the parameters `number` will be dynamically adjusted to meet the
+            minimum duration requirement of one `repeat`.
+            i.e., When the run time of one `repeat` falls below this time, the `number` parameter
+            will be automatically increased.
+
+        cooldown_interval_ms: int, optional
+            The cooldown interval in milliseconds between the number of repeats defined by
+            `repeats_to_cooldown`.
+
+        repeats_to_cooldown: int, optional
+            The number of repeats before the cooldown is activated.
+
+        f_preproc: str, optional
+            The preprocess function name we want to execute before executing the time evaluator.
+
+        Note
+        ----
+        The function will be invoked  (1 + number x repeat) times,
+        with the first call discarded in case there is lazy initialization.
+
+        Example
+        -------
+        Normal use with a VM function (may not work over RPC if the function returns a tuple):
+
+        .. code-block:: python
+
+            target = tvm.target.Target("llvm", host="llvm")
+            ex = tvm.compile(TestTimeEvaluator, target)
+            vm = relax.VirtualMachine(mod, tvm.cpu())
+            timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)
+
+        Use with the stateful API:
+
+        .. code-block:: python
+
+            target = tvm.target.Target("llvm", host="llvm")
+            ex = tvm.compile(TestTimeEvaluator, target)
+            vm = relax.VirtualMachine(mod, tvm.cpu())
+            vm.set_input("func_name", arg0, arg1, ..., argn)
+            timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")
+
+        With saved closures via `save_function` (this results in
+        fewer dictionary lookups in the timed portion):
+
+        .. code-block:: python
+
+            target = tvm.target.Target("llvm", host="llvm")
+            ex = tvm.compile(TestTimeEvaluator, target)
+            vm = relax.VirtualMachine(mod, tvm.cpu())
+            vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn)
+            timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()
+
+        Returns
+        -------
+        ftimer : function
+            The function that takes same argument as func and returns a BenchmarkResult.
+            The ProfileResult reports `repeat` time costs in seconds.
+
+        """
+        return self.module.time_evaluator(
+            func_name,
+            dev,
+            number=number,
+            repeat=repeat,
+            min_repeat_ms=min_repeat_ms,
+            cooldown_interval_ms=cooldown_interval_ms,
+            repeats_to_cooldown=repeats_to_cooldown,
+            f_preproc=f_preproc,
+        )
+
+
+@register_global_func("vm.builtin.debug_print")
+def _print(lineo: str, array) -> None:
+    print(f"{lineo}: shape = {array.shape}, dtype = {array.dtype}, data =\n{array}")