This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx/standalone-package-cleanup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 0bb1535c9c628d0ec28fcbc157152cecaf346eab Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 15 01:36:34 2026 +0000 [REFACTOR][Python] Load core and parent extensions from their owning packages --- .../tvm_tirx/runtime => python/tvm}/__init__.py | 34 +++------ {tirx/python/tvm_tirx => python/tvm}/base.py | 41 +++-------- {tirx/python/tvm_tirx => python/tvm}/libinfo.py | 12 ++-- python/tvm/relax/backend/gpu_generic/cumsum.py | 2 +- python/tvm/relax/backend/gpu_generic/sampling.py | 4 +- python/tvm/relax/backend/patterns.py | 2 +- .../tvm/relax/frontend/nn/llm/_decode_kernels.py | 2 +- python/tvm/relax/frontend/nn/llm/_kernel_common.py | 6 +- python/tvm/relax/frontend/nn/llm/_page_kernels.py | 2 +- .../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 4 +- .../tvm/relax/frontend/nn/llm/_prefill_parallel.py | 2 +- .../relax/frontend/nn/llm/position_embedding.py | 4 +- python/tvm/relax/frontend/nn/llm/tree_attn.py | 4 +- python/tvm/relax/frontend/nn/op.py | 7 +- .../tvm/relax/frontend/tflite/tflite_frontend.py | 82 ++++++++++++++++------ python/tvm/relax/testing/attention.py | 4 +- python/tvm/relax/testing/matmul.py | 4 +- .../tvm/relax/transform/legalize_ops/inspect_op.py | 2 +- python/tvm/rpc/minrpc.py | 30 +++++--- .../tvm_tirx => python/tvm}/runtime/__init__.py | 27 +------ python/tvm/s_tir/tensor_intrin/arm_cpu.py | 4 +- .../tvm/s_tir/tensor_intrin/dot_product_common.py | 2 +- python/tvm/s_tir/tensor_intrin/hexagon.py | 2 +- python/tvm/s_tir/tensor_intrin/metal.py | 4 +- python/tvm/s_tir/tensor_intrin/riscv_cpu.py | 2 +- python/tvm/s_tir/tensor_intrin/rocm.py | 2 +- python/tvm/s_tir/tensor_intrin/x86.py | 2 +- tirx/python/tvm_tirx/__init__.py | 12 ---- tirx/python/tvm_tirx/backend/_autoload_backends.py | 21 +----- tirx/python/tvm_tirx/backend/cuda/__init__.py | 2 +- tirx/python/tvm_tirx/backend/metal/__init__.py | 2 +- tirx/python/tvm_tirx/backend/opencl/__init__.py | 2 +- tirx/python/tvm_tirx/backend/rocm/__init__.py | 2 +- tirx/python/tvm_tirx/backend/vulkan/__init__.py | 2 +- tirx/python/tvm_tirx/base.py | 15 ++-- tirx/python/tvm_tirx/libinfo.py | 27 ++++--- tirx/python/tvm_tirx/runtime/__init__.py | 7 -- tirx/python/tvm_tirx/script/parser/entry.py | 1 + web/tests/python/prepare_test_libs.py | 4 +- web/tests/python/relax_rpc_test.py | 5 +- 40 files changed, 178 insertions(+), 217 deletions(-) diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/python/tvm/__init__.py similarity index 50% copy from tirx/python/tvm_tirx/runtime/__init__.py copy to python/tvm/__init__.py index b57d3af1b2..7320967ffb 100644 --- a/tirx/python/tvm_tirx/runtime/__init__.py +++ b/python/tvm/__init__.py @@ -15,28 +15,16 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""TVM runtime namespace.""" +# pylint: disable=redefined-builtin, wildcard-import +"""TVM compiler dialects and runtime extensions built on TIRx.""" -from tvm_ffi import convert, Object -from tvm_ffi._dtype import dtype as DataType, DataTypeCode +from .libinfo import __version__ +from .base import _RUNTIME_ONLY +from . import runtime +from . import rpc -# class exposures -from .device import Device -from .memory_scope import MemoryScope -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 - -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 +if not _RUNTIME_ONLY: + from . import te + from . import s_tir + from .driver import compile + from . import relax diff --git a/tirx/python/tvm_tirx/base.py b/python/tvm/base.py similarity index 54% copy from tirx/python/tvm_tirx/base.py copy to python/tvm/base.py index ca65863d5b..40469d954a 100644 --- a/tirx/python/tvm_tirx/base.py +++ b/python/tvm/base.py @@ -16,49 +16,30 @@ # under the License. # coding: utf-8 # pylint: disable=invalid-name, import-outside-toplevel -"""Base library for TVM.""" +"""Load parent TVM extensions after their standalone TIRx dependency.""" -import os +from pathlib import Path from tvm_ffi import libinfo as tvm_ffi_libinfo from tvm_ffi.libinfo import load_lib_ctypes -from . import libinfo - -# ---------------------------- -# library loading -# ---------------------------- - -# Whether only the runtime library is loaded (runtime-only wheel, or -# ``TVM_USE_RUNTIME_LIB=1``). Set during library loading below. -_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1" - -# Handles of the core libraries actually loaded, keyed by basename -# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_compiler": <CDLL>}``). Downstream / -# autoloaded extensions can inspect this to skip duplicate libraries -# (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle. -_LOADED_LIBS = {} +from tvm_tirx.base import _RUNTIME_ONLY as _CORE_RUNTIME_ONLY +from . import libinfo -# runtime is loaded RTLD_GLOBAL to expose its symbols to subsequent loads; -# compiler is loaded RTLD_LOCAL. -_LOADED_LIBS["tvm_runtime"] = load_lib_ctypes( - "tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths() -) - +_RUNTIME_ONLY = _CORE_RUNTIME_ONLY +_LOADED_LIBS = { + "tvm_runtime_extra": load_lib_ctypes( + "tvm", "tvm_runtime_extra", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths() + ) +} if not _RUNTIME_ONLY: - _lib_paths = libinfo.package_lib_paths() + _lib_paths = [Path(_LOADED_LIBS["tvm_runtime_extra"]._name).resolve().parent] try: tvm_ffi_libinfo._find_library_by_basename("tvm", "tvm_compiler", extra_lib_paths=_lib_paths) except RuntimeError: - # A runtime-only installation need not contain the compiler library. _RUNTIME_ONLY = True else: _LOADED_LIBS["tvm_compiler"] = load_lib_ctypes( "tvm", "tvm_compiler", "RTLD_LOCAL", extra_lib_paths=_lib_paths ) - -if _RUNTIME_ONLY: - from tvm_ffi import registry as _tvm_ffi_registry - - _tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True diff --git a/tirx/python/tvm_tirx/libinfo.py b/python/tvm/libinfo.py similarity index 89% copy from tirx/python/tvm_tirx/libinfo.py copy to python/tvm/libinfo.py index 312d769f7e..18bc4c2daf 100644 --- a/tirx/python/tvm_tirx/libinfo.py +++ b/python/tvm/libinfo.py @@ -41,7 +41,7 @@ def package_lib_paths() -> list[Path]: covers the wheel-install layout (``python/tvm/lib/``) and the in-tree dev build layouts (``<worktree>/build/lib/`` and ``<worktree>/lib/``). ``TVM_LIBRARY_PATH`` is prepended when set so it takes priority. Callers - pick the basenames they want (e.g. ``libtvm_runtime.so``) and the load + pick the basenames they want (e.g. ``libtvm_runtime_extra.so``) and the load mode; this function only returns the search path. """ pkg = _rel_top_directory() # python/tvm/ @@ -56,8 +56,8 @@ def package_lib_paths() -> list[Path]: return paths -def find_libtvm_runtime() -> str: - """Return the path to the ``libtvm_runtime`` shared library. +def find_libtvm_runtime_extra() -> str: + """Return the path to the ``libtvm_runtime_extra`` shared library. Returns ------- @@ -70,11 +70,11 @@ def find_libtvm_runtime() -> str: If the runtime library cannot be found. """ candidate = tvm_ffi_libinfo._find_library_by_basename( - "tvm", "tvm_runtime", extra_lib_paths=package_lib_paths() + "tvm", "tvm_runtime_extra", extra_lib_paths=package_lib_paths() ) if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda _: True): return ret - raise RuntimeError("Cannot find libtvm_runtime") + raise RuntimeError("Cannot find libtvm_runtime_extra") def find_include_path() -> str: @@ -84,7 +84,7 @@ def find_include_path() -> str: _rel_top_directory() / "include", _dev_top_directory() / "include", ], - cond=lambda p: (p / "tvm" / "runtime").is_dir(), + cond=lambda p: (p / "tvm" / "runtime" / "vm").is_dir(), ): return ret raise RuntimeError("Cannot find TVM include path.") diff --git a/python/tvm/relax/backend/gpu_generic/cumsum.py b/python/tvm/relax/backend/gpu_generic/cumsum.py index 2fb20bdb54..7610912f69 100644 --- a/python/tvm/relax/backend/gpu_generic/cumsum.py +++ b/python/tvm/relax/backend/gpu_generic/cumsum.py @@ -19,7 +19,7 @@ import math -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from tvm_tirx.tirx import Function diff --git a/python/tvm/relax/backend/gpu_generic/sampling.py b/python/tvm/relax/backend/gpu_generic/sampling.py index 6e9d2fe499..0c28923e99 100644 --- a/python/tvm/relax/backend/gpu_generic/sampling.py +++ b/python/tvm/relax/backend/gpu_generic/sampling.py @@ -20,8 +20,8 @@ import math from collections.abc import Callable -import tvm, tvm_tirx -from tvm_tirx.script import s_tir as Ts +import tvm_tirx +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from tvm_tirx.tirx import Function diff --git a/python/tvm/relax/backend/patterns.py b/python/tvm/relax/backend/patterns.py index 75de6ae575..f5cff44054 100644 --- a/python/tvm/relax/backend/patterns.py +++ b/python/tvm/relax/backend/patterns.py @@ -19,6 +19,7 @@ from collections.abc import Mapping +from tvm.relax import script as R from tvm.relax.dpl.pattern import ( DFPattern, GlobalVarPattern, @@ -28,7 +29,6 @@ from tvm.relax.dpl.pattern import ( is_tuple_get_item, wildcard, ) -from tvm_tirx.script import relax as R from tvm_tirx.script import tirx as T diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py index ec3ce3e670..a51b56fb5f 100644 --- a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py +++ b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py @@ -31,7 +31,7 @@ Contents: import math from typing import Any -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py index e2f32620ae..d26010eefd 100644 --- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py +++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py @@ -37,11 +37,11 @@ import enum import math from typing import Any -import tvm, tvm_tirx -from tvm_tirx import tirx +import tvm from tvm import s_tir +from tvm.s_tir import script as Ts +from tvm_tirx import tirx from tvm_tirx.runtime import DataType -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/relax/frontend/nn/llm/_page_kernels.py b/python/tvm/relax/frontend/nn/llm/_page_kernels.py index 8c7f2a248d..95273278b6 100644 --- a/python/tvm/relax/frontend/nn/llm/_page_kernels.py +++ b/python/tvm/relax/frontend/nn/llm/_page_kernels.py @@ -32,7 +32,7 @@ This module contains: from __future__ import annotations # pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py index d0ee8e155a..36530e3dea 100644 --- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py +++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py @@ -29,9 +29,9 @@ K/V loading path that is specific to its storage layout. import math from typing import Any, Literal -import tvm, tvm_tirx +import tvm +from tvm.s_tir import script as Ts from tvm_tirx import tirx -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py index eba66a746c..46253c9bf4 100644 --- a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py +++ b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py @@ -16,7 +16,7 @@ # under the License. """Parallel dot-product prefill for wide FP32 WebGPU attention heads.""" -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from ._kernel_common import _get_kv_chunk_len, _get_seq_offset, _length_info_buffer, _rope diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py index cddf2462ae..99ea24bedc 100644 --- a/python/tvm/relax/frontend/nn/llm/position_embedding.py +++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py @@ -22,9 +22,9 @@ from collections.abc import Callable from functools import partial from typing import Any -from tvm_tirx import tirx from tvm.relax.frontend.nn import Tensor, op -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts +from tvm_tirx import tirx from tvm_tirx.script import tirx as T # pylint: disable=invalid-name diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py b/python/tvm/relax/frontend/nn/llm/tree_attn.py index 1b52f319fd..777b1d638e 100644 --- a/python/tvm/relax/frontend/nn/llm/tree_attn.py +++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py @@ -22,9 +22,9 @@ import math from typing import Any -from tvm_tirx import tirx from tvm import s_tir -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts +from tvm_tirx import tirx from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py index 10c925ed4f..e21e91f755 100644 --- a/python/tvm/relax/frontend/nn/op.py +++ b/python/tvm/relax/frontend/nn/op.py @@ -25,10 +25,10 @@ from typing import Any, TypeVar import numpy as np -import tvm, tvm_tirx +import tvm_tirx from tvm import te +from tvm.s_tir import script as Ts from tvm_tirx import tirx as _tir -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from ... import expr as rx @@ -2181,7 +2181,8 @@ def extern( out: OutType, ) -> OutType: """Invoke an extern function during runtime. The extern function must be registered with the " - TVM runtime using `reflection::GlobalDef().def` (C++), or `tvm_tirx.register_global_func` (Python). + TVM runtime using `reflection::GlobalDef().def` (C++), or + `tvm_tirx.register_global_func` (Python). Parameters ---------- diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 0a66475452..27d5c21498 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -27,13 +27,13 @@ import math import numpy as np -import tvm, tvm_tirx -from tvm_tirx import tirx +import tvm_tirx from tvm import relax -from tvm_tirx.ir import TupleType as _TupleType from tvm.relax import ExternFunc as _ExternFunc from tvm.relax import op as _op -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts +from tvm_tirx import tirx +from tvm_tirx.ir import TupleType as _TupleType from .tflite_flexbuffer import FlexBufferDecoder @@ -751,16 +751,22 @@ class OperatorConverter: offsets = np.frombuffer(data, dtype="<i4", count=count + 1, offset=4).astype(np.int64) if np.any(offsets < header_size) or np.any(offsets > len(data)): - raise tvm_tirx.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets") + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} has out-of-bounds string tensor offsets" + ) if np.any(offsets[:-1] > offsets[1:]): - raise tvm_tirx.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets") + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} has non-monotonic string tensor offsets" + ) try: values = [ data[int(offsets[i]) : int(offsets[i + 1])].decode("utf-8") for i in range(count) ] except UnicodeDecodeError as e: - raise tvm_tirx.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} has invalid UTF-8 string data: {e}" + ) from e shape = self._get_tensor_shape_tuple(tensor_wrapper) expected_count = math.prod(shape) if shape else 1 if expected_count != count: @@ -811,11 +817,15 @@ class OperatorConverter: key_shape = self._get_tensor_shape_tuple(key_tensor) value_shape = self._get_tensor_shape_tuple(value_tensor) if key_shape != value_shape: - raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires keys and values same shape") + raise tvm_tirx.error.OpNotImplemented( + "HASHTABLE_IMPORT requires keys and values same shape" + ) if not self._has_tensor_buffer_data(key_tensor) or not self._has_tensor_buffer_data( value_tensor ): - raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values") + raise tvm_tirx.error.OpNotImplemented( + "HASHTABLE_IMPORT requires constant keys and values" + ) if self._is_tflite_string_type(table_info["key_dtype"]): keys = self._get_string_tensor_value(key_tensor, "HASHTABLE_IMPORT") @@ -889,7 +899,9 @@ class OperatorConverter: query_shape = self._get_tensor_shape_tuple(query_tensor) output_shape = self._get_tensor_shape_tuple(output_tensor) if output_shape != query_shape: - raise tvm_tirx.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape") + raise tvm_tirx.error.OpNotImplemented( + "HASHTABLE_FIND output shape must match query shape" + ) query_values = self._get_string_tensor_value(query_tensor, "HASHTABLE_FIND") default_values = self.get_tensor_value(default_tensor) @@ -1554,7 +1566,9 @@ class OperatorConverter: from tflite.LocalResponseNormalizationOptions import LocalResponseNormalizationOptions if self.is_quantized(op): - raise tvm_tirx.error.OpNotImplemented("TFlite quantized LRN operator is not supported yet.") + raise tvm_tirx.error.OpNotImplemented( + "TFlite quantized LRN operator is not supported yet." + ) input_tensors = self.get_input_tensors(op) assert len(input_tensors) == 1, "input tensors length should be 1" @@ -2096,7 +2110,9 @@ class OperatorConverter: elif dtype.startswith(("int", "uint")): op_fn = _op.bitwise_and else: - raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_AND with dtype {dtype} is not supported") + raise tvm_tirx.error.OpNotImplemented( + f"STABLEHLO_AND with dtype {dtype} is not supported" + ) return self.bb.normalize(op_fn(lhs, rhs)) def _convert_stablehlo_or(self, op): @@ -2114,7 +2130,9 @@ class OperatorConverter: elif dtype.startswith(("int", "uint")): op_fn = _op.bitwise_or else: - raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_OR with dtype {dtype} is not supported") + raise tvm_tirx.error.OpNotImplemented( + f"STABLEHLO_OR with dtype {dtype} is not supported" + ) return self.bb.normalize(op_fn(lhs, rhs)) def _convert_stablehlo_ternary(self, op, relax_op): @@ -2561,7 +2579,9 @@ class OperatorConverter: elif direction == _DIR.STABLEHLO_COMPARISON_DIRECTION_GT: descending = True else: - raise tvm_tirx.error.OpNotImplemented("STABLEHLO_SORT only supports LT or GT comparators") + raise tvm_tirx.error.OpNotImplemented( + "STABLEHLO_SORT only supports LT or GT comparators" + ) data = self.get_tensor_expr(input_tensors[0]) return self.bb.normalize( @@ -2602,7 +2622,9 @@ class OperatorConverter: return self.get_tensor_expr(input_tensors[0]) target = call_target_name or "<empty>" - raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_CUSTOM_CALL target {target} is not supported") + raise tvm_tirx.error.OpNotImplemented( + f"STABLEHLO_CUSTOM_CALL target {target} is not supported" + ) def _convert_stablehlo_rng_bit_generator(self, op): """Convert STABLEHLO_RNG_BIT_GENERATOR to a bit-exact call_tir kernel.""" @@ -2638,7 +2660,9 @@ class OperatorConverter: if algorithm == "threefry" and state_len != 2: raise tvm_tirx.error.OpNotImplemented(f"{op_name} THREEFRY requires a u64[2] state") if algorithm == "philox" and state_len not in (2, 3): - raise tvm_tirx.error.OpNotImplemented(f"{op_name} PHILOX requires a u64[2] or u64[3] state") + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} PHILOX requires a u64[2] or u64[3] state" + ) out_state_tensor, out_tensor = output_tensors if self.get_tensor_type_str(out_state_tensor.tensor.Type()) != "uint64": @@ -2650,7 +2674,9 @@ class OperatorConverter: ) out_dtype = self.get_tensor_type_str(out_tensor.tensor.Type()) if out_dtype not in ("int32", "int64", "uint32", "uint64"): - raise tvm_tirx.error.OpNotImplemented(f"{op_name} output dtype {out_dtype} is not supported") + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} output dtype {out_dtype} is not supported" + ) out_shape = tuple(self._get_static_tensor_shape(out_tensor, op_name)) function = _build_stablehlo_rng_bit_generator_function( @@ -2779,7 +2805,9 @@ class OperatorConverter: def _check_tensor_metadata_match(self, actual, expected, op_name, tensor_role): """Validate that two TFLite tensors have matching static metadata.""" if self._get_tensor_metadata(actual) != self._get_tensor_metadata(expected): - raise tvm_tirx.error.OpNotImplemented(f"{op_name} {tensor_role} tensor metadata mismatch") + raise tvm_tirx.error.OpNotImplemented( + f"{op_name} {tensor_role} tensor metadata mismatch" + ) def _check_subgraph_tensor_metadata( self, subgraph, op_name, tensor_role, subgraph_indices, expected_tensors @@ -3087,7 +3115,9 @@ class OperatorConverter: input_tensors = self.get_input_tensors(op) output_tensors = self.get_output_tensors(op) if len(input_tensors) != 0 or len(output_tensors) != 0: - raise tvm_tirx.error.OpNotImplemented("CALL_ONCE with inputs or outputs is not supported") + raise tvm_tirx.error.OpNotImplemented( + "CALL_ONCE with inputs or outputs is not supported" + ) init_subgraph = self._get_subgraph(init_subgraph_index, "CALL_ONCE") if init_subgraph.InputsLength() != 0 or init_subgraph.OutputsLength() != 0: @@ -3603,7 +3633,9 @@ class OperatorConverter: data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_CONVOLUTION") kernel_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_CONVOLUTION") if len(data_shape) != 4 or len(kernel_shape) != 4: - raise tvm_tirx.error.OpNotImplemented("STABLEHLO_CONVOLUTION only supports 2D convolution") + raise tvm_tirx.error.OpNotImplemented( + "STABLEHLO_CONVOLUTION only supports 2D convolution" + ) if data_shape[3] != kernel_shape[2]: raise tvm_tirx.error.OpNotImplemented( "STABLEHLO_CONVOLUTION input channels must match kernel input channels" @@ -3691,7 +3723,9 @@ class OperatorConverter: "STABLEHLO_GATHER only supports collapsed_slice_dims matching the gather axis" ) if len(slice_sizes) != len(data_shape): - raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER slice_sizes must match operand rank") + raise tvm_tirx.error.OpNotImplemented( + "STABLEHLO_GATHER slice_sizes must match operand rank" + ) for i, (size, dim) in enumerate(zip(slice_sizes, data_shape)): expected = 1 if i == axis else dim if size != expected: @@ -3703,7 +3737,9 @@ class OperatorConverter: "STABLEHLO_GATHER only supports trailing index_vector_dim" ) if not indices_shape or indices_shape[index_vector_dim] != 1: - raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER only supports index vector size 1") + raise tvm_tirx.error.OpNotImplemented( + "STABLEHLO_GATHER only supports index vector size 1" + ) indices_batch_shape = indices_shape[:index_vector_dim] expected_offset_dims = list(range(axis)) + list( @@ -8515,7 +8551,7 @@ def _build_tflite_rfft2d_fft_function(input_shape, output_pair_shape): # this kernel. function_source = ( "from tvm_tirx.script.parser import tirx as T\n" - "from tvm_tirx.script import s_tir as Ts\n" + "from tvm.s_tir import script as Ts\n" "@Ts.function(private=True, check_well_formed=False)\n" "def kernel(\n" f" data: T.Tensor({tuple(int(x) for x in input_shape)}, 'float32'),\n" diff --git a/python/tvm/relax/testing/attention.py b/python/tvm/relax/testing/attention.py index 6aa1f4a946..67bdd2bc79 100644 --- a/python/tvm/relax/testing/attention.py +++ b/python/tvm/relax/testing/attention.py @@ -17,9 +17,9 @@ """Relax script for attention module.""" -import tvm, tvm_tirx +import tvm_tirx +from tvm.relax import script as R from tvm.relax.script import ir_builder as relax_builder -from tvm_tirx.script import relax as R from tvm_tirx.script import tirx as T from tvm_tirx.script.ir_builder import IRBuilder diff --git a/python/tvm/relax/testing/matmul.py b/python/tvm/relax/testing/matmul.py index 5ebbf3e60e..3251ccbbef 100644 --- a/python/tvm/relax/testing/matmul.py +++ b/python/tvm/relax/testing/matmul.py @@ -17,9 +17,9 @@ # ruff: noqa: RUF005 """Utilities to construct matmul workloads.""" -import tvm, tvm_tirx +import tvm_tirx +from tvm.relax import script as R from tvm.relax.script import ir_builder as relax_builder -from tvm_tirx.script import relax as R from tvm_tirx.script.ir_builder import IRBuilder diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py index 068982e2ce..8bdf1931e8 100644 --- a/python/tvm/relax/transform/legalize_ops/inspect_op.py +++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py @@ -19,8 +19,8 @@ import enum +from tvm.s_tir import script as Ts from tvm_tirx.ir import Call -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from ... import op diff --git a/python/tvm/rpc/minrpc.py b/python/tvm/rpc/minrpc.py index 7e5c812e22..5885f75005 100644 --- a/python/tvm/rpc/minrpc.py +++ b/python/tvm/rpc/minrpc.py @@ -20,7 +20,8 @@ import os import tvm_ffi -from tvm_tirx import libinfo +from tvm import libinfo +from tvm_tirx import libinfo as core_libinfo from tvm_tirx.support import cc @@ -39,13 +40,15 @@ def find_minrpc_server_libpath(server="posix_popen_server"): """ curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__))) source_dir = os.path.abspath(os.path.join(curr_dir, "..", "..", "..")) - minrpc_dir = os.path.join(source_dir, "src", "runtime", "rpc", "minrpc") - path = os.path.join(minrpc_dir, server, f"{server}.cc") - - candidates = [path] - if not os.path.isfile(path): - raise RuntimeError(f"Cannot find minserver {server}, in candidates {candidates}") - return minrpc_dir, path + roots = [os.path.join(curr_dir, ".."), source_dir] + candidates = [] + for root in roots: + minrpc_dir = os.path.join(root, "src", "runtime", "rpc", "minrpc") + path = os.path.join(minrpc_dir, server, f"{server}.cc") + candidates.append(path) + if os.path.isfile(path): + return minrpc_dir, path + raise RuntimeError(f"Cannot find minserver {server}, in candidates {candidates}") def with_minrpc(compile_func, server="posix_popen_server"): @@ -65,7 +68,8 @@ def with_minrpc(compile_func, server="posix_popen_server"): The return compilation. """ minrpc_dir, server_path = find_minrpc_server_libpath(server) - runtime_path = libinfo.find_libtvm_runtime() + runtime_path = libinfo.find_libtvm_runtime_extra() + core_runtime_path = core_libinfo.find_libtvm_runtime() tvm_ffi_path = tvm_ffi.libinfo.find_libtvm_ffi() runtime_dir = os.path.abspath(os.path.dirname(runtime_path)) @@ -75,23 +79,27 @@ def with_minrpc(compile_func, server="posix_popen_server"): # Note that however, this approach won't work on remote. # Always recommend to link statically. options += ["-Wl,-rpath=" + runtime_dir] + options += ["-Wl,-rpath=" + os.path.dirname(core_runtime_path)] options += ["-Wl,-rpath=" + tvm_ffi_dir] # Force the linker to retain the runtime shared lib dependency even if no # symbol from it is directly referenced. The minrpc server binary only # touches tvm_ffi APIs at link time but relies on global functions (e.g. # ``rpc.CreateEventDrivenServer``) registered via static initializers - # inside libtvm_runtime; with the default ``--as-needed`` those + # inside libtvm_runtime_extra; with the default ``--as-needed`` those # registrations would never run. options += ["-Wl,--no-as-needed"] default_include_paths = [ libinfo.find_include_path(), + core_libinfo.find_include_path(), tvm_ffi.libinfo.find_include_path(), tvm_ffi.libinfo.find_dlpack_include_path(), ] options += ["-I" + path for path in default_include_paths] options += ["-I" + minrpc_dir] fcompile = cc.cross_compiler( - compile_func, options=options, add_files=[server_path, runtime_path, tvm_ffi_path] + compile_func, + options=options, + add_files=[server_path, runtime_path, core_runtime_path, tvm_ffi_path], ) fcompile.__name__ = "with_minrpc" fcompile.need_system_lib = True diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/python/tvm/runtime/__init__.py similarity index 50% copy from tirx/python/tvm_tirx/runtime/__init__.py copy to python/tvm/runtime/__init__.py index b57d3af1b2..bfaec7febf 100644 --- a/tirx/python/tvm_tirx/runtime/__init__.py +++ b/python/tvm/runtime/__init__.py @@ -15,28 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""TVM runtime namespace.""" +# pylint: disable=redefined-builtin, wildcard-import +"""Parent runtime extensions: virtual machines and distributed execution.""" -from tvm_ffi import convert, Object -from tvm_ffi._dtype import dtype as DataType, DataTypeCode - -# class exposures -from .device import Device -from .memory_scope import MemoryScope -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 - -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 +from . import disco, vm diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py index d566967fd3..8de05107fa 100644 --- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py +++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py @@ -18,10 +18,10 @@ # ruff: noqa: E501, F401 """Intrinsics for ARM tensorization.""" +from tvm.s_tir import script as Ts +from tvm.s_tir.script.ir_builder import function as build_function from tvm_tirx import tirx from tvm_tirx.codegen import llvm_version_major -from tvm.s_tir.script.ir_builder import function as build_function -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from tvm_tirx.script.ir_builder import IRBuilder diff --git a/python/tvm/s_tir/tensor_intrin/dot_product_common.py b/python/tvm/s_tir/tensor_intrin/dot_product_common.py index e3c8bc9ab2..7204ea6ce1 100644 --- a/python/tvm/s_tir/tensor_intrin/dot_product_common.py +++ b/python/tvm/s_tir/tensor_intrin/dot_product_common.py @@ -17,7 +17,7 @@ # pylint: disable=invalid-name,missing-function-docstring """Dot product related intrinsics.""" -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from .. import TensorIntrin diff --git a/python/tvm/s_tir/tensor_intrin/hexagon.py b/python/tvm/s_tir/tensor_intrin/hexagon.py index 34c2183330..74e4c9e9ce 100644 --- a/python/tvm/s_tir/tensor_intrin/hexagon.py +++ b/python/tvm/s_tir/tensor_intrin/hexagon.py @@ -17,7 +17,7 @@ # pylint: disable=invalid-name,missing-function-docstring """Intrinsics for Hexagon tensorization.""" -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from .. import TensorIntrin diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py index bbee2ab12b..b0ca2411ff 100644 --- a/python/tvm/s_tir/tensor_intrin/metal.py +++ b/python/tvm/s_tir/tensor_intrin/metal.py @@ -19,9 +19,9 @@ from typing import Literal -from tvm_tirx.runtime import MemoryScope from tvm.s_tir import TensorIntrin -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts +from tvm_tirx.runtime import MemoryScope from tvm_tirx.script import tirx as T from tvm_tirx.tirx import Expr, Function, Var diff --git a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py index 4c970aa422..b31f411e78 100644 --- a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py +++ b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py @@ -21,9 +21,9 @@ import logging import tvm_ffi +from tvm.s_tir import script as Ts from tvm_tirx.codegen import llvm_get_vector_width, target_has_features from tvm_tirx.runtime import DataType -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from tvm_tirx.target import Target diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py index 4d84fddbde..ccfaa9d075 100644 --- a/python/tvm/s_tir/tensor_intrin/rocm.py +++ b/python/tvm/s_tir/tensor_intrin/rocm.py @@ -26,8 +26,8 @@ lane dimension also lets CompactTensorAllocation remove the logical 64-lane extent without a warp-memory lowering pass. """ +from tvm.s_tir import script as Ts from tvm_tirx.runtime import convert -from tvm_tirx.script import s_tir as Ts from tvm_tirx.script import tirx as T from tvm_tirx.tirx.expr import Cast, IntImm diff --git a/python/tvm/s_tir/tensor_intrin/x86.py b/python/tvm/s_tir/tensor_intrin/x86.py index 959648b64a..e73370ab86 100644 --- a/python/tvm/s_tir/tensor_intrin/x86.py +++ b/python/tvm/s_tir/tensor_intrin/x86.py @@ -17,7 +17,7 @@ # pylint: disable=invalid-name,missing-function-docstring """Intrinsics for x86 tensorization.""" -from tvm_tirx.script import s_tir as Ts +from tvm.s_tir import script as Ts from tvm_tirx.script import tirx as T from .. import TensorIntrin diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index e6381cbe3b..efbdb05a5e 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -60,12 +60,6 @@ from .target import device_from_target # Shared code generation helpers from . import codegen -# tvm.te -from . import te - -# tvm.driver -from .driver import compile - # others from . import sym @@ -75,12 +69,6 @@ from . import support # Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup from .support import rocm as _rocm, nvcc as _nvcc -# Relax contain modules that are only available in compiler package -# Do not import them if TVM is built with runtime only -if not _RUNTIME_ONLY: - # tvm.relax — registers itself via tvm.script.register_dialect in its __init__ - from . import relax - # NOTE: This file should be python2 compatible so we can # raise proper error message when user run the package using # an older version of the python diff --git a/tirx/python/tvm_tirx/backend/_autoload_backends.py b/tirx/python/tvm_tirx/backend/_autoload_backends.py index a348b7ec00..600075640b 100644 --- a/tirx/python/tvm_tirx/backend/_autoload_backends.py +++ b/tirx/python/tvm_tirx/backend/_autoload_backends.py @@ -21,12 +21,8 @@ from __future__ import annotations import os import warnings from importlib.metadata import entry_points -from pathlib import Path - -from tvm_ffi.libinfo import load_lib_ctypes from tvm_tirx.backend.loader import load -from tvm_tirx.base import _LOADED_LIBS _BUILTIN_BACKENDS = ( "cuda", @@ -48,17 +44,6 @@ def _load_builtin_backends() -> None: for name in _BUILTIN_BACKENDS: load(name) - runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent - try: - # Load libtvm_runtime_extra if available for registration side effects. - _LOADED_LIBS["tvm_runtime_extra"] = load_lib_ctypes( - package="tvm", - target_name="tvm_runtime_extra", - extra_lib_paths=[runtime_dir], - mode="RTLD_LOCAL", - ) - except (OSError, FileNotFoundError, RuntimeError): - pass return None @@ -78,10 +63,10 @@ def _autoload_backends() -> None: _load_builtin_backends() # Out-of-tree extensions opt into being loaded automatically at ``import tvm`` time - # by declaring an entry point in the ``tvm.backends`` group: - # [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload". + # by declaring an entry point in the ``tvm_tirx.backends`` group: + # [project.entry-points."tvm_tirx.backends"] tvm_foo = "tvm_foo:_autoload". # Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``. - for entry_pt in entry_points(group="tvm.backends"): + for entry_pt in entry_points(group="tvm_tirx.backends"): try: entry_pt.load()() except Exception as e: # pylint: disable=broad-except diff --git a/tirx/python/tvm_tirx/backend/cuda/__init__.py b/tirx/python/tvm_tirx/backend/cuda/__init__.py index 01bd9316f6..ee50466435 100644 --- a/tirx/python/tvm_tirx/backend/cuda/__init__.py +++ b/tirx/python/tvm_tirx/backend/cuda/__init__.py @@ -72,7 +72,7 @@ def register_backend(): try: # Runtime sidecars only need registration side effects; libtvm_runtime is global. _LOADED_LIBS["tvm_runtime_cuda"] = load_lib_ctypes( - package="tvm", + package="tvm_tirx", target_name="tvm_runtime_cuda", extra_lib_paths=[runtime_dir], mode="RTLD_LOCAL", diff --git a/tirx/python/tvm_tirx/backend/metal/__init__.py b/tirx/python/tvm_tirx/backend/metal/__init__.py index 46afaa5f13..222d01c78c 100644 --- a/tirx/python/tvm_tirx/backend/metal/__init__.py +++ b/tirx/python/tvm_tirx/backend/metal/__init__.py @@ -50,7 +50,7 @@ def register_backend(): try: # Runtime sidecars only need registration side effects; libtvm_runtime is global. _LOADED_LIBS["tvm_runtime_metal"] = load_lib_ctypes( - package="tvm", + package="tvm_tirx", target_name="tvm_runtime_metal", extra_lib_paths=[runtime_dir], mode="RTLD_LOCAL", diff --git a/tirx/python/tvm_tirx/backend/opencl/__init__.py b/tirx/python/tvm_tirx/backend/opencl/__init__.py index 15d2277b67..633d540593 100644 --- a/tirx/python/tvm_tirx/backend/opencl/__init__.py +++ b/tirx/python/tvm_tirx/backend/opencl/__init__.py @@ -46,7 +46,7 @@ def register_backend(): try: # Runtime sidecars only need registration side effects; libtvm_runtime is global. _LOADED_LIBS["tvm_runtime_opencl"] = load_lib_ctypes( - package="tvm", + package="tvm_tirx", target_name="tvm_runtime_opencl", extra_lib_paths=[runtime_dir], mode="RTLD_LOCAL", diff --git a/tirx/python/tvm_tirx/backend/rocm/__init__.py b/tirx/python/tvm_tirx/backend/rocm/__init__.py index 20ff0f9ea1..ffd3ccfccd 100644 --- a/tirx/python/tvm_tirx/backend/rocm/__init__.py +++ b/tirx/python/tvm_tirx/backend/rocm/__init__.py @@ -45,7 +45,7 @@ def register_backend(): try: # Runtime sidecars only need registration side effects; libtvm_runtime is global. _LOADED_LIBS["tvm_runtime_rocm"] = load_lib_ctypes( - package="tvm", + package="tvm_tirx", target_name="tvm_runtime_rocm", extra_lib_paths=[runtime_dir], mode="RTLD_LOCAL", diff --git a/tirx/python/tvm_tirx/backend/vulkan/__init__.py b/tirx/python/tvm_tirx/backend/vulkan/__init__.py index 476fb716af..c13d17c574 100644 --- a/tirx/python/tvm_tirx/backend/vulkan/__init__.py +++ b/tirx/python/tvm_tirx/backend/vulkan/__init__.py @@ -53,7 +53,7 @@ def register_backend(): try: # Runtime sidecars only need registration side effects; libtvm_runtime is global. _LOADED_LIBS["tvm_runtime_vulkan"] = load_lib_ctypes( - package="tvm", + package="tvm_tirx", target_name="tvm_runtime_vulkan", extra_lib_paths=[runtime_dir], mode="RTLD_LOCAL", diff --git a/tirx/python/tvm_tirx/base.py b/tirx/python/tvm_tirx/base.py index ca65863d5b..8bc62431c2 100644 --- a/tirx/python/tvm_tirx/base.py +++ b/tirx/python/tvm_tirx/base.py @@ -19,6 +19,7 @@ """Base library for TVM.""" import os +from pathlib import Path from tvm_ffi import libinfo as tvm_ffi_libinfo from tvm_ffi.libinfo import load_lib_ctypes @@ -34,7 +35,7 @@ from . import libinfo _RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1" # Handles of the core libraries actually loaded, keyed by basename -# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_compiler": <CDLL>}``). Downstream / +# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_tirx": <CDLL>}``). Downstream / # autoloaded extensions can inspect this to skip duplicate libraries # (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle. _LOADED_LIBS = {} @@ -43,19 +44,21 @@ _LOADED_LIBS = {} # runtime is loaded RTLD_GLOBAL to expose its symbols to subsequent loads; # compiler is loaded RTLD_LOCAL. _LOADED_LIBS["tvm_runtime"] = load_lib_ctypes( - "tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths() + "tvm_tirx", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths() ) if not _RUNTIME_ONLY: - _lib_paths = libinfo.package_lib_paths() + _lib_paths = [Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent] try: - tvm_ffi_libinfo._find_library_by_basename("tvm", "tvm_compiler", extra_lib_paths=_lib_paths) + tvm_ffi_libinfo._find_library_by_basename( + "tvm_tirx", "tvm_tirx", extra_lib_paths=_lib_paths + ) except RuntimeError: # A runtime-only installation need not contain the compiler library. _RUNTIME_ONLY = True else: - _LOADED_LIBS["tvm_compiler"] = load_lib_ctypes( - "tvm", "tvm_compiler", "RTLD_LOCAL", extra_lib_paths=_lib_paths + _LOADED_LIBS["tvm_tirx"] = load_lib_ctypes( + "tvm_tirx", "tvm_tirx", "RTLD_LOCAL", extra_lib_paths=_lib_paths ) if _RUNTIME_ONLY: diff --git a/tirx/python/tvm_tirx/libinfo.py b/tirx/python/tvm_tirx/libinfo.py index 312d769f7e..79345f706b 100644 --- a/tirx/python/tvm_tirx/libinfo.py +++ b/tirx/python/tvm_tirx/libinfo.py @@ -25,7 +25,7 @@ from tvm_ffi import libinfo as tvm_ffi_libinfo def _rel_top_directory() -> Path: - """Top directory in the installed (wheel) layout: ``python/tvm/``.""" + """Top directory in the installed (wheel) layout: ``python/tvm_tirx/``.""" return Path(__file__).parent @@ -37,22 +37,25 @@ def _dev_top_directory() -> Path: def package_lib_paths() -> list[Path]: """Return search directories for TVM's shared libraries. - Anchored on this file's location (``python/tvm/libinfo.py``), the list - covers the wheel-install layout (``python/tvm/lib/``) and the in-tree dev + Anchored on this file's location (``python/tvm_tirx/libinfo.py``), the list + covers the wheel-install layout (``python/tvm_tirx/lib/``) and the in-tree dev build layouts (``<worktree>/build/lib/`` and ``<worktree>/lib/``). - ``TVM_LIBRARY_PATH`` is prepended when set so it takes priority. Callers + ``TVM_TIRX_LIBRARY_PATH`` is prepended when set so it takes priority. Callers pick the basenames they want (e.g. ``libtvm_runtime.so``) and the load mode; this function only returns the search path. """ - pkg = _rel_top_directory() # python/tvm/ + pkg = _rel_top_directory() # python/tvm_tirx/ paths: list[Path] = [] - if os.environ.get("TVM_LIBRARY_PATH"): - paths.append(Path(os.environ["TVM_LIBRARY_PATH"])) + if os.environ.get("TVM_TIRX_LIBRARY_PATH"): + paths.append(Path(os.environ["TVM_TIRX_LIBRARY_PATH"])) paths += [ pkg / "lib", # wheel layout _dev_top_directory() / "build" / "lib", # dev: <worktree>/build/lib _dev_top_directory() / "lib", # dev: <worktree>/lib ] + dev = _dev_top_directory().resolve() + if dev.name == "tirx" and (dev.parent / "CMakeLists.txt").is_file(): + paths.append(dev.parent / "build" / "lib") return paths @@ -70,7 +73,7 @@ def find_libtvm_runtime() -> str: If the runtime library cannot be found. """ candidate = tvm_ffi_libinfo._find_library_by_basename( - "tvm", "tvm_runtime", extra_lib_paths=package_lib_paths() + "tvm_tirx", "tvm_runtime", extra_lib_paths=package_lib_paths() ) if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda _: True): return ret @@ -90,10 +93,4 @@ def find_include_path() -> str: raise RuntimeError("Cannot find TVM include path.") -# The version is written by setuptools_scm into _version.py at build time -# (see [tool.setuptools_scm] in pyproject.toml). The fallback keeps a source -# checkout with no build run importable. -try: - from ._version import version as __version__ -except ImportError: # pragma: no cover - source tree without a build - __version__ = "0.28.dev0" +__version__ = "0.1.0a0" diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py index b57d3af1b2..b9eebb95ff 100644 --- a/tirx/python/tvm_tirx/runtime/__init__.py +++ b/tirx/python/tvm_tirx/runtime/__init__.py @@ -32,11 +32,4 @@ 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 -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 diff --git a/tirx/python/tvm_tirx/script/parser/entry.py b/tirx/python/tvm_tirx/script/parser/entry.py index 60302557e6..00d529b920 100644 --- a/tirx/python/tvm_tirx/script/parser/entry.py +++ b/tirx/python/tvm_tirx/script/parser/entry.py @@ -570,6 +570,7 @@ def _prepare_transpiler( No builder frame or expression is created here. """ namespace = { + "tvm": sys.modules.get("tvm"), "tvm_tirx": sys.modules.get("tvm_tirx"), **_NAMESPACES, **environment, diff --git a/web/tests/python/prepare_test_libs.py b/web/tests/python/prepare_test_libs.py index c6dc645b95..bdc70184fd 100644 --- a/web/tests/python/prepare_test_libs.py +++ b/web/tests/python/prepare_test_libs.py @@ -18,11 +18,11 @@ import os -import tvm, tvm_tirx +import tvm_tirx from tvm import relax, te from tvm.contrib import emcc +from tvm.relax import script as R from tvm_tirx.script import ir as I -from tvm_tirx.script import relax as R def prepare_relax_lib(base_path): diff --git a/web/tests/python/relax_rpc_test.py b/web/tests/python/relax_rpc_test.py index 1444558ba6..e3d9cdbe48 100644 --- a/web/tests/python/relax_rpc_test.py +++ b/web/tests/python/relax_rpc_test.py @@ -18,10 +18,11 @@ import numpy as np -import tvm, tvm_tirx +import tvm +import tvm_tirx from tvm import relax, rpc from tvm.contrib import emcc -from tvm_tirx.script import relax as R +from tvm.relax import script as R from tvm_tirx.support import utils proxy_host = "127.0.0.1"
