You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/16 20:55:25 UTC
(tvm) 04/12: [REFACTOR] Separate TIRx compiler and runtime deployment APIs (#20763)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx-standalone-root
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 341e3fd7f5fc779f3bd216472c262d8d17428bf3
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Fri Oct 16 14:12:53 2026 -0400
[REFACTOR] Separate TIRx compiler and runtime deployment APIs (#20763)
Expose component-based build and link configuration through
`tvm-tirx-config` and name the compiler library `libtvm_tirx_compiler`.
Move export and JIT orchestration into `codegen.Executable`, use FFI
modules directly in the runtime, and keep object/archive linking
explicit in compiler and RPC consumers. Move benchmarking policy to the
parent runtime while retaining the small device timer interface and
backend providers in TIRx.
---
tirx/CMakeLists.txt | 31 +-
tirx/cmake/modules/ClangFlags.cmake | 2 +-
tirx/cmake/modules/Logging.cmake | 4 +-
tirx/cmake/tvm_tirx-config.cmake | 6 +-
tirx/include/tvm/runtime/base.h | 2 +-
tirx/include/tvm/runtime/timer.h | 64 +---
tirx/pyproject.toml | 3 +
tirx/python/tvm_tirx/__init__.py | 3 +-
tirx/python/tvm_tirx/base.py | 8 +-
.../tvm_tirx/{codegen.py => codegen/__init__.py} | 4 +-
.../{runtime/module.py => codegen/executable.py} | 351 ++++++---------------
tirx/python/tvm_tirx/config.py | 209 +++++++++---
tirx/python/tvm_tirx/runtime/__init__.py | 43 ++-
tirx/python/tvm_tirx/runtime/executable.py | 175 ----------
tirx/python/tvm_tirx/support/cc.py | 6 +-
tirx/python/tvm_tirx/support/ndk.py | 4 +-
tirx/python/tvm_tirx/tirx/_compile.py | 4 +-
tirx/src/runtime/timer.cc | 276 ----------------
tirx/tests/cpp/CMakeLists.txt | 2 +-
.../test_validate_runtime_library.py | 43 ++-
.../python/backend/cuda/codegen/test_fastmath.py | 2 +-
tirx/tests/python/codegen/c/test_basic.py | 6 +-
tirx/tests/python/codegen/test_executable.py | 145 +++++++++
tirx/tests/python/conftest.py | 78 +++++
tirx/tests/python/runtime/test_module.py | 152 +--------
tirx/tests/python/runtime/test_tirx_config.py | 46 ++-
tirx/tests/python/support/test_config.py | 130 +++++++-
tirx/tests/python/tirx/iket/test_profiler.py | 5 +-
28 files changed, 784 insertions(+), 1020 deletions(-)
diff --git a/tirx/CMakeLists.txt b/tirx/CMakeLists.txt
index 683377e4da..4001e955d9 100644
--- a/tirx/CMakeLists.txt
+++ b/tirx/CMakeLists.txt
@@ -223,31 +223,31 @@ include(cmake/modules/ROCM.cmake)
include(cmake/modules/LLVM.cmake)
include(cmake/modules/contrib/Z3.cmake)
include(cmake/modules/Git.cmake)
-# ######### Target: `tvm_tirx` ##########
+# ######### Target: `tvm_tirx_compiler` ##########
-add_library(tvm_tirx_objs OBJECT ${COMPILER_SRCS})
-target_link_libraries(tvm_tirx_objs PUBLIC tvm_tirx_header)
+add_library(tvm_tirx_compiler_objs OBJECT ${COMPILER_SRCS})
+target_link_libraries(tvm_tirx_compiler_objs PUBLIC tvm_tirx_header)
if(TARGET tvm_llvm_header)
- target_link_libraries(tvm_tirx_objs PRIVATE tvm_llvm_header)
+ target_link_libraries(tvm_tirx_compiler_objs PRIVATE tvm_llvm_header)
endif()
if(MSVC)
- target_compile_definitions(tvm_tirx_objs PRIVATE TVM_TIRX_EXPORTS TVM_FFI_EXPORTS)
+ target_compile_definitions(tvm_tirx_compiler_objs PRIVATE TVM_TIRX_EXPORTS TVM_FFI_EXPORTS)
endif()
-add_library(tvm_tirx SHARED $<TARGET_OBJECTS:tvm_tirx_objs>)
-add_library(tvm_tirx::shared ALIAS tvm_tirx)
-target_link_libraries(tvm_tirx PUBLIC tvm_runtime tvm_tirx_header)
-target_link_libraries(tvm_tirx PRIVATE ${TVM_LINKER_LIBS})
+add_library(tvm_tirx_compiler SHARED $<TARGET_OBJECTS:tvm_tirx_compiler_objs>)
+add_library(tvm_tirx_compiler::shared ALIAS tvm_tirx_compiler)
+target_link_libraries(tvm_tirx_compiler PUBLIC tvm_runtime tvm_tirx_header)
+target_link_libraries(tvm_tirx_compiler PRIVATE ${TVM_LINKER_LIBS})
# Avoid GNU ld relaxation of LLVM's GOT calls into incorrect direct calls.
if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT "${TVM_TIRX_USE_LLVM}" MATCHES ${IS_FALSE_PATTERN})
- target_link_options(tvm_tirx PRIVATE "-Wl,--no-relax" "-Wl,--as-needed")
+ target_link_options(tvm_tirx_compiler PRIVATE "-Wl,--no-relax" "-Wl,--as-needed")
endif()
include(cmake/modules/Logging.cmake)
include(cmake/modules/ClangFlags.cmake)
# ######### Library configuration ##########
-foreach(_target tvm_runtime_objs tvm_tirx_objs tvm_runtime tvm_runtime_static tvm_tirx)
+foreach(_target tvm_runtime_objs tvm_tirx_compiler_objs tvm_runtime tvm_runtime_static tvm_tirx_compiler)
tvm_tirx_configure_target(${_target})
tvm_ffi_add_prefix_map(${_target} "${CMAKE_CURRENT_SOURCE_DIR}")
endforeach()
@@ -255,7 +255,7 @@ foreach(_target IN LISTS TVM_RUNTIME_BACKEND_LIBS)
tvm_ffi_add_prefix_map(${_target}_objs "${CMAKE_CURRENT_SOURCE_DIR}")
tvm_ffi_add_prefix_map(${_target} "${CMAKE_CURRENT_SOURCE_DIR}")
endforeach()
-foreach(_target tvm_runtime tvm_tirx ${TVM_RUNTIME_BACKEND_LIBS})
+foreach(_target tvm_runtime tvm_tirx_compiler ${TVM_RUNTIME_BACKEND_LIBS})
if(TVM_TIRX_BUILD_PYTHON_MODULE)
if(APPLE)
set_property(TARGET ${_target} APPEND PROPERTY INSTALL_RPATH "@loader_path/../../tvm_ffi/lib")
@@ -284,7 +284,7 @@ foreach(_header IN LISTS _headers)
endforeach()
# Libraries
-install(TARGETS tvm_tirx_header tvm_tirx tvm_runtime tvm_runtime_static
+install(TARGETS tvm_tirx_header tvm_tirx_compiler tvm_runtime tvm_runtime_static
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx
RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx)
@@ -297,6 +297,11 @@ install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/"
DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx/utils"
COMPONENT tvm_tirx FILES_MATCHING PATTERN "Find*.cmake")
+if(NOT TVM_TIRX_BUILD_PYTHON_MODULE)
+ install(PROGRAMS "${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_tirx/config.py"
+ DESTINATION "${CMAKE_INSTALL_BINDIR}" RENAME tvm-tirx-config COMPONENT tvm_tirx)
+endif()
+
# Python package artifacts and licenses
if(TVM_TIRX_BUILD_PYTHON_MODULE)
foreach(_extra_lib IN LISTS TVM_TIRX_PACKAGE_EXTRA_LIBS)
diff --git a/tirx/cmake/modules/ClangFlags.cmake b/tirx/cmake/modules/ClangFlags.cmake
index f296165726..b8a6cddb60 100644
--- a/tirx/cmake/modules/ClangFlags.cmake
+++ b/tirx/cmake/modules/ClangFlags.cmake
@@ -82,7 +82,7 @@ if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang")
-Wno-gnu-anonymous-struct
-Wno-nested-anon-types
)
- target_compile_options(tvm_tirx_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
+ target_compile_options(tvm_tirx_compiler_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
target_compile_options(tvm_runtime_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
diff --git a/tirx/cmake/modules/Logging.cmake b/tirx/cmake/modules/Logging.cmake
index df7f2a7b7f..df8b3955d7 100644
--- a/tirx/cmake/modules/Logging.cmake
+++ b/tirx/cmake/modules/Logging.cmake
@@ -21,9 +21,9 @@ include(FindPackageHandleStandardArgs)
if(TVM_TIRX_USE_CUSTOM_LOGGING)
# Set and propagate TVM_LOG_CUSTOMIZE flag is custom logging has been requested
- target_compile_definitions(tvm_tirx_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
+ target_compile_definitions(tvm_tirx_compiler_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
target_compile_definitions(tvm_runtime_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_tirx PUBLIC TVM_LOG_CUSTOMIZE=1)
+ target_compile_definitions(tvm_tirx_compiler PUBLIC TVM_LOG_CUSTOMIZE=1)
target_compile_definitions(tvm_runtime PUBLIC TVM_LOG_CUSTOMIZE=1)
target_compile_definitions(tvm_runtime_static INTERFACE TVM_LOG_CUSTOMIZE=1)
endif()
diff --git a/tirx/cmake/tvm_tirx-config.cmake b/tirx/cmake/tvm_tirx-config.cmake
index 4ef057cd43..118342b0ef 100644
--- a/tirx/cmake/tvm_tirx-config.cmake
+++ b/tirx/cmake/tvm_tirx-config.cmake
@@ -38,7 +38,7 @@ if(NOT TARGET tvm_tirx::header)
target_include_directories(tvm_tirx::header INTERFACE "${tvm_tirx_INCLUDE_DIR}")
endif()
-foreach(_library tvm_runtime tvm_tirx)
+foreach(_library tvm_runtime tvm_tirx_compiler)
if(TARGET ${_library}::shared)
continue()
endif()
@@ -50,7 +50,7 @@ foreach(_library tvm_runtime tvm_tirx)
set_target_properties(${_library}::shared PROPERTIES IMPORTED_IMPLIB "${_implib}")
endif()
target_link_libraries(${_library}::shared INTERFACE tvm_tirx::header)
- if(_library STREQUAL "tvm_tirx")
- target_link_libraries(tvm_tirx::shared INTERFACE tvm_runtime::shared)
+ if(_library STREQUAL "tvm_tirx_compiler")
+ target_link_libraries(tvm_tirx_compiler::shared INTERFACE tvm_runtime::shared)
endif()
endforeach()
diff --git a/tirx/include/tvm/runtime/base.h b/tirx/include/tvm/runtime/base.h
index d4e7e6e20c..0848fff40c 100644
--- a/tirx/include/tvm/runtime/base.h
+++ b/tirx/include/tvm/runtime/base.h
@@ -77,7 +77,7 @@
#define TVM_DLL_EXPORT __attribute__((visibility("default")))
#endif
-// Core IR, TIRx and shared compiler infrastructure live in libtvm_tirx.
+// Core IR, TIRx and shared compiler infrastructure live in libtvm_tirx_compiler.
#ifndef TVM_TIRX_DLL
#if defined(__EMSCRIPTEN__)
#define TVM_TIRX_DLL EMSCRIPTEN_KEEPALIVE
diff --git a/tirx/include/tvm/runtime/timer.h b/tirx/include/tvm/runtime/timer.h
index 07abab3b0d..684098c32d 100644
--- a/tirx/include/tvm/runtime/timer.h
+++ b/tirx/include/tvm/runtime/timer.h
@@ -19,15 +19,12 @@
/*!
* \file tvm/runtime/timer.h
- * \brief Runtime timer primitives: Timer, TimerObj.
+ * \brief Device timer interface and object reference.
*/
#ifndef TVM_RUNTIME_TIMER_H_
#define TVM_RUNTIME_TIMER_H_
-#include <tvm/ffi/function.h>
-#include <tvm/runtime/base.h>
-#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/tensor.h>
+#include <tvm/ffi/object.h>
namespace tvm {
namespace runtime {
@@ -76,63 +73,6 @@ class TimerObj : public ffi::Object {
*/
class Timer : public ffi::ObjectRef {
public:
- /*!
- * \brief Get a device specific timer.
- * \param dev The device to time.
- * \return A `Timer` that has already been started.
- *
- * Use this function to time runtime of arbitrary regions of code on a specific
- * device. The code that you want to time should be running on the device
- * otherwise the timer will not return correct results. This is a lower level
- * interface than TimeEvaluator and only runs the timed code once
- * (TimeEvaluator runs the code multiple times).
- *
- * A default timer is used if a device specific one does not exist. This
- * timer performs synchronization between the device and CPU, which can lead
- * to overhead in the reported results.
- *
- * Example usage:
- * \code{.cpp}
- * Timer t = Timer::Start(Device::cpu());
- * my_long_running_function();
- * t->Stop();
- * ... // some more computation
- * int64_t nanosecs = t->SyncAndGetElapsedNanos() // elapsed time in nanoseconds
- * \endcode
- *
- * To add a new device-specific timer, register a new function
- * "runtime.timer.my_device" (where `my_device` is the `DeviceName` of your
- * device). This function should accept a `Device` and return a new `Timer`
- * that has already been started.
- *
- * For example, this is how the CPU timer is implemented:
- * \code{.cpp}
- * class CPUTimerObj : public TimerObj {
- * public:
- * virtual void Start() { start_ = std::chrono::high_resolution_clock::now(); }
- * virtual void Stop() { duration_ = std::chrono::high_resolution_clock::now() - start_; }
- * virtual int64_t SyncAndGetElapsedNanos() { return duration_.count(); }
- * virtual ~CPUTimerObj() {}
- *
- * static constexpr const char* _type_key = "runtime.CPUTimerObj";
- * TVM_FFI_DECLARE_OBJECT_INFO_FINAL(CPUTimerObj, TimerObj);
- *
- * private:
- * std::chrono::high_resolution_clock::time_point start_;
- * std::chrono::duration<int64_t, std::nano> duration_;
- * };
- *
- *
- * TVM_FFI_STATIC_INIT_BLOCK() {
- * namespace refl = tvm::ffi::reflection;
- * refl::GlobalDef().def("runtime.timer.cpu", [](Device dev) {
- * return Timer(ffi::UnsafeInit{}, ffi::make_object<CPUTimerObj>());
- * });
- * }
- * \endcode
- */
- static TVM_RUNTIME_DLL Timer Start(Device dev);
-
explicit Timer(ffi::UnsafeInit tag, ffi::ObjectPtr<TimerObj> n) : ffi::ObjectRef(tag) {
data_ = std::move(n);
}
diff --git a/tirx/pyproject.toml b/tirx/pyproject.toml
index 9043a5163e..4ef79dcb0c 100644
--- a/tirx/pyproject.toml
+++ b/tirx/pyproject.toml
@@ -51,6 +51,9 @@ dependencies = [
"packaging",
]
+[project.scripts]
+tvm-tirx-config = "tvm_tirx.config:__main__"
+
[project.optional-dependencies]
torch = ["torch"]
cuda = ["cuda-bindings"]
diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py
index a66f56d844..fcdeb1f0bd 100644
--- a/tirx/python/tvm_tirx/__init__.py
+++ b/tirx/python/tvm_tirx/__init__.py
@@ -19,13 +19,14 @@
"""TVM: Open Deep Learning Compiler Stack."""
import sys
+import os
from typing import TYPE_CHECKING
def _is_config_mode():
"""Keep the configuration CLI independent of native library loading."""
argv = getattr(sys, "orig_argv", ())
- return any(
+ return os.path.basename(sys.argv[0]) in ("tvm-tirx-config", "tvm-tirx-config.exe") or any(
arg == "-m" and argv[index + 1 : index + 2] == ["tvm_tirx.config"]
for index, arg in enumerate(argv)
)
diff --git a/tirx/python/tvm_tirx/base.py b/tirx/python/tvm_tirx/base.py
index 8bc62431c2..17c4a05b7b 100644
--- a/tirx/python/tvm_tirx/base.py
+++ b/tirx/python/tvm_tirx/base.py
@@ -35,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_tirx": <CDLL>}``). Downstream /
+# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_tirx_compiler": <CDLL>}``). Downstream /
# autoloaded extensions can inspect this to skip duplicate libraries
# (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle.
_LOADED_LIBS = {}
@@ -51,14 +51,14 @@ if not _RUNTIME_ONLY:
_lib_paths = [Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent]
try:
tvm_ffi_libinfo._find_library_by_basename(
- "tvm_tirx", "tvm_tirx", extra_lib_paths=_lib_paths
+ "tvm_tirx", "tvm_tirx_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_tirx"] = load_lib_ctypes(
- "tvm_tirx", "tvm_tirx", "RTLD_LOCAL", extra_lib_paths=_lib_paths
+ _LOADED_LIBS["tvm_tirx_compiler"] = load_lib_ctypes(
+ "tvm_tirx", "tvm_tirx_compiler", "RTLD_LOCAL", extra_lib_paths=_lib_paths
)
if _RUNTIME_ONLY:
diff --git a/tirx/python/tvm_tirx/codegen.py b/tirx/python/tvm_tirx/codegen/__init__.py
similarity index 98%
rename from tirx/python/tvm_tirx/codegen.py
rename to tirx/python/tvm_tirx/codegen/__init__.py
index ba3b2036f0..afc5fc7c3f 100644
--- a/tirx/python/tvm_tirx/codegen.py
+++ b/tirx/python/tvm_tirx/codegen/__init__.py
@@ -18,7 +18,9 @@
from tvm_ffi import Array, get_global_func
-from .target import Target
+from tvm_tirx.target import Target
+
+from .executable import Executable, load_static_library
def build_module(mod, target):
diff --git a/tirx/python/tvm_tirx/runtime/module.py b/tirx/python/tvm_tirx/codegen/executable.py
similarity index 52%
rename from tirx/python/tvm_tirx/runtime/module.py
rename to tirx/python/tvm_tirx/codegen/executable.py
index 0743a45f50..674b47c127 100644
--- a/tirx/python/tvm_tirx/runtime/module.py
+++ b/tirx/python/tvm_tirx/codegen/executable.py
@@ -14,100 +14,106 @@
# 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."""
+"""Compiler-owned export and JIT orchestration for ordinary FFI modules."""
import json
import os
-import struct
-from collections.abc import Sequence
+from collections.abc import Callable
+from typing import Any
-import numpy as np
-from tvm_ffi import (
- Module as _Module,
-)
+from tvm_ffi import Function, Module, load_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_tirx.libinfo
from tvm_tirx.base import _RUNTIME_ONLY
+from tvm_tirx.runtime import _ffi_api
+from tvm_tirx.support import utils as _utils
+
+
+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
-from . import _ffi_api
+ 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)
-class BenchmarkResult:
- """Runtimes from benchmarking"""
+ 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.
- def __init__(self, results: Sequence[float]):
- """Construct a new BenchmarkResult from a sequence of runtimes.
+ 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
----------
- results : Sequence[float]
- Raw times from benchmarking
+ 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_ffi.Module
+ A runnable runtime module that can be passed to VirtualMachine.
+
+ Examples
+ --------
+ .. code:: python
+
+ ex = tvm.compile(mod, target)
+ rt_mod = ex.jit()
- 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."""
+
+ # 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._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 = load_module(dso_path)
+ return self._jitted_mod
def _collect_from_import_tree(self, filter_func):
"""Helper function to collect modules from the tree matching a filter_func, then return it.
@@ -125,8 +131,8 @@ class Module(_Module):
"""
visited, stack, dso_modules = set(), [], []
# append root module
- visited.add(self)
- stack.append(self)
+ visited.add(self.mod)
+ stack.append(self.mod)
while stack:
module = stack.pop()
assert module.is_compilation_exportable() or module.is_binary_serializable(), (
@@ -226,7 +232,7 @@ class Module(_Module):
if workspace_dir is None:
temp = _utils.tempdir()
workspace_dir = temp.temp_dir
- files = addons if addons else []
+ files = list(addons) if addons else []
is_system_lib = False
has_c_module = False
system_lib_prefix = None
@@ -282,24 +288,24 @@ class Module(_Module):
if getattr(fcompile, "need_system_lib", False) and not is_system_lib:
raise ValueError(f"{fcompile!s} need --system-lib option")
- if self.imports:
+ if self.mod.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)
+ path_out = fpack_imports(self.mod, 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}"
)
_ffi_api.ModulePackImportsToObjectFile(
- self, is_system_lib, llvm_target, path_obj, pack_lib_prefix
+ self.mod, is_system_lib, llvm_target, path_obj, pack_lib_prefix
)
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))
+ f.write(_ffi_api.ModulePackImportsToC(self.mod, is_system_lib, pack_lib_prefix))
files.append(path_cc)
# The imports could contain a c module but the object format could be tar
@@ -320,186 +326,9 @@ class Module(_Module):
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.
- """
- feval = _ffi_api.ModuleTimeEvaluator(
- self,
- func_name,
- dev,
- 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
-
-
-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_tirx.support import cc as _cc
-
- _cc.create_shared(path + ".so", path)
- path += ".so"
- elif path.endswith(".tar"):
- # Extra dependencies during runtime.
- from tvm_tirx.support import cc as _cc
- from tvm_tirx.support import tar as _tar
- from tvm_tirx.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_tirx.target.Target
- The target device type.
-
- Returns
- -------
- enabled : bool
- Whether runtime is enabled.
-
- Examples
- --------
- The following code checks if gpu is enabled.
-
- >>> tvm_tirx.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()
diff --git a/tirx/python/tvm_tirx/config.py b/tirx/python/tvm_tirx/config.py
index de55764a37..466a30fe14 100644
--- a/tirx/python/tvm_tirx/config.py
+++ b/tirx/python/tvm_tirx/config.py
@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
# 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
@@ -21,13 +22,26 @@ there with an exact library directory, independently of Python package search.
"""
import argparse
+import importlib.machinery
+import importlib.util
import json
+import os
+import shlex
import sys
from pathlib import Path
__version__ = "0.1.0a0"
-_LIBRARY_NAMES = ("tvm_tirx", "tvm_runtime")
+_LIBRARY_NAMES = ("tvm_tirx_compiler", "tvm_runtime")
+_COMPONENTS = {
+ "compiler": "tvm_tirx_compiler",
+ "runtime": "tvm_runtime",
+ "cuda": "tvm_runtime_cuda",
+ "metal": "tvm_runtime_metal",
+ "opencl": "tvm_runtime_opencl",
+ "vulkan": "tvm_runtime_vulkan",
+ "rocm": "tvm_runtime_rocm",
+}
def _find_prefix(start: Path) -> Path:
@@ -38,19 +52,25 @@ def _find_prefix(start: Path) -> Path:
def _paths(prefix: Path | None, library_dir: Path | None) -> tuple[Path, Path]:
+ explicit_prefix = prefix is not None
if library_dir is not None:
library_dir = library_dir.resolve()
prefix = prefix.resolve() if prefix is not None else _find_prefix(library_dir)
else:
prefix = prefix.resolve() if prefix is not None else _find_prefix(Path(__file__).parent)
library_dir = prefix / "lib"
- for candidate in (prefix / "lib", prefix / "lib64", prefix / "build/lib"):
+ candidates = [prefix / "lib", prefix / "lib64", prefix / "build/lib"]
+ if prefix.name == "tirx":
+ candidates.append(prefix.parent / "build/lib")
+ if not explicit_prefix and os.environ.get("TVM_TIRX_LIBRARY_PATH"):
+ candidates.insert(0, Path(os.environ["TVM_TIRX_LIBRARY_PATH"]))
+ for candidate in candidates:
if all(
any(
(candidate / filename).is_file()
for filename in (f"lib{name}.so", f"lib{name}.dylib", f"{name}.dll")
)
- for name in _LIBRARY_NAMES
+ for name in ("tvm_runtime",)
):
library_dir = candidate
break
@@ -59,11 +79,11 @@ def _paths(prefix: Path | None, library_dir: Path | None) -> tuple[Path, Path]:
return prefix, library_dir
-def _libraries(library_dir: Path, platform: str) -> dict[str, dict[str, str]]:
+def _libraries(library_dir: Path, platform: str, names=_LIBRARY_NAMES) -> dict[str, dict[str, str]]:
libraries = {}
windows = platform in ("Windows", "win32")
apple = platform in ("Darwin", "iOS", "darwin")
- for name in _LIBRARY_NAMES:
+ for name in names:
if windows:
filenames = (f"{name}.dll",)
elif apple:
@@ -102,47 +122,162 @@ def configuration(
}
+def _ffi_info():
+ """Read the selected FFI package's pure Python helper without importing its runtime."""
+ spec = importlib.machinery.PathFinder.find_spec("tvm_ffi")
+ if spec is None or spec.origin is None:
+ raise RuntimeError("Cannot find the matching apache-tvm-ffi Python package")
+ helper = Path(spec.origin).parent / "libinfo.py"
+ spec = importlib.util.spec_from_file_location("_tvm_ffi_libinfo", helper)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+def _available(library_dir, platform):
+ available = []
+ for component, name in _COMPONENTS.items():
+ try:
+ _libraries(library_dir, platform, (name,))
+ except RuntimeError:
+ continue
+ available.append(component)
+ return available
+
+
+def _link_inputs(library_dir, platform, components):
+ available = _available(library_dir, platform)
+ selected = list(dict.fromkeys(components or ["all"]))
+ if "all" in selected:
+ selected = [*available, *(name for name in selected if name != "all")]
+ for name in selected:
+ if name != "ffi" and name not in available:
+ raise RuntimeError(
+ f"Component '{name}' is unavailable; available: {' '.join(available)} ffi"
+ )
+ # Dependents precede their dependencies even if the command line lists runtime first.
+ ordered = [name for name in selected if name not in ("runtime", "ffi")]
+ if any(name != "ffi" for name in selected):
+ ordered.append("runtime")
+ libraries = _libraries(library_dir, platform, [_COMPONENTS[name] for name in ordered])
+ ffi = _ffi_info()
+ ffi_library = (
+ ffi.find_windows_implib() if platform in ("Windows", "win32") else ffi.find_libtvm_ffi()
+ )
+ return [
+ *(Path(lib["implib"] or lib["location"]) for lib in libraries.values()),
+ Path(ffi_library),
+ ]
+
+
def __main__() -> None:
- """Print installed include, library, or CMake package paths."""
- parser = argparse.ArgumentParser(description=__doc__)
- parser.add_argument("--prefix", type=Path, help="Use this exact installation prefix")
- parser.add_argument("--library-dir", type=Path, help="Use this exact native library directory")
- parser.add_argument("--platform", default=sys.platform, help=argparse.SUPPRESS)
- outputs = parser.add_mutually_exclusive_group()
- outputs.add_argument(
- "--includedir", action="store_true", help="Print the public include directory"
+ """Print LLVM-style component configuration without loading native libraries."""
+ parser = argparse.ArgumentParser(
+ description=__doc__,
+ epilog="Components default to all available components. Examples: --libs runtime; "
+ "--libs runtime cuda; --libs compiler. Shared linking is the default.",
)
- outputs.add_argument("--libdir", action="store_true", help="Print the native library directory")
- outputs.add_argument(
- "--libfiles", action="store_true", help="Print library files, one per line"
+ parser.add_argument("components", nargs="*", help="Components to link (default: all)")
+ parser.add_argument("--install-prefix", type=Path, help="Query this exact installation prefix")
+ parser.add_argument(
+ "--library-dir", type=Path, help="Query this exact native library directory"
)
- outputs.add_argument(
- "--cmakedir", action="store_true", help="Print the CMake package directory"
+ parser.add_argument("--platform", default=sys.platform, help=argparse.SUPPRESS)
+ outputs = (
+ "version",
+ "prefix",
+ "includedir",
+ "libdir",
+ "cmakedir",
+ "components",
+ "cflags",
+ "cxxflags",
+ "ldflags",
+ "libs",
+ "libnames",
+ "libfiles",
+ "json",
)
- outputs.add_argument(
- "--json", action="store_true", help="Print all consumer configuration as JSON"
+ for output in outputs:
+ parser.add_argument("--" + output, dest="show_" + output, action="store_true")
+ mode = parser.add_mutually_exclusive_group()
+ mode.add_argument("--link-shared", action="store_true", help="Use shared libraries (default)")
+ mode.add_argument(
+ "--link-static", action="store_true", help="Require a complete static closure"
)
args = parser.parse_args()
- if not any((args.includedir, args.libdir, args.libfiles, args.cmakedir, args.json)):
+ if not any(getattr(args, "show_" + output) for output in outputs):
parser.print_help()
return
try:
- if args.json or args.libfiles:
- result = configuration(args.prefix, args.library_dir, args.platform)
- if args.json:
- print(json.dumps(result))
- else:
- for library in result["libraries"].values():
- print(library["implib"] or library["location"])
- else:
- prefix, library_dir = _paths(args.prefix, args.library_dir)
- if args.includedir:
- print(prefix / "include")
- elif args.libdir:
- print(library_dir)
- else:
- print(library_dir / "cmake/tvm_tirx")
- except RuntimeError as error:
+ if args.link_static:
+ raise RuntimeError(
+ "Static linking is unavailable: this package does not provide a "
+ "complete static dependency closure (including TVM FFI and its dependencies)"
+ )
+ if args.show_version:
+ print(__version__)
+ if all(not getattr(args, "show_" + output) for output in outputs if output != "version"):
+ return
+ prefix, library_dir = _paths(args.install_prefix, args.library_dir)
+ paths = {
+ "prefix": prefix,
+ "includedir": prefix / "include",
+ "libdir": library_dir,
+ "cmakedir": library_dir / "cmake/tvm_tirx",
+ }
+ if not paths["cmakedir"].is_dir() and (prefix / "cmake/tvm_tirx-config.cmake").is_file():
+ paths["cmakedir"] = prefix / "cmake"
+ for output, path in paths.items():
+ if getattr(args, "show_" + output):
+ print(path)
+ if args.show_components:
+ print(" ".join([*_available(library_dir, args.platform), "ffi"]))
+ if args.show_json:
+ print(json.dumps(configuration(prefix, library_dir, args.platform)))
+ if args.show_cflags or args.show_cxxflags:
+ ffi = _ffi_info()
+ includes = dict.fromkeys(
+ [str(prefix / "include"), ffi.find_include_path(), ffi.find_dlpack_include_path()]
+ )
+ flags = [*("-I" + path for path in includes), "-DTVM_INDEX_DEFAULT_I64=1"]
+ if args.show_cflags:
+ print(shlex.join(flags))
+ if args.show_cxxflags:
+ print(shlex.join([*flags, "-std=c++20"]))
+ if any((args.show_ldflags, args.show_libs, args.show_libnames, args.show_libfiles)):
+ libraries = _link_inputs(library_dir, args.platform, args.components)
+ apple = args.platform in ("Darwin", "iOS", "darwin")
+ if args.show_ldflags:
+ print(
+ shlex.join(
+ [
+ "-L" + str(path)
+ for path in dict.fromkeys(lib.parent for lib in libraries)
+ ]
+ )
+ )
+ if args.show_libnames:
+ print(shlex.join([lib.name for lib in libraries]))
+ if args.show_libfiles:
+ print(shlex.join([str(lib) for lib in libraries]))
+ if args.show_libs:
+ if apple:
+ flags = [
+ flag
+ for lib in libraries
+ for flag in ("-Xlinker", "-needed_library", "-Xlinker", str(lib))
+ ]
+ elif args.platform.startswith("linux") or args.platform == "Linux":
+ flags = [
+ "-Wl,--push-state,--no-as-needed",
+ *map(str, libraries),
+ "-Wl,--pop-state",
+ ]
+ else:
+ raise RuntimeError(f"--libs retention flags are unsupported on {args.platform}")
+ print(shlex.join(flags))
+ except (RuntimeError, OSError) as error:
parser.exit(1, f"{error}\n")
diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py
index b9eebb95ff..aa8f098b9d 100644
--- a/tirx/python/tvm_tirx/runtime/__init__.py
+++ b/tirx/python/tvm_tirx/runtime/__init__.py
@@ -17,19 +17,54 @@
# under the License.
"""TVM runtime namespace."""
-from tvm_ffi import convert, Object
+from tvm_ffi import convert, Object, Module, load_module, system_lib
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
+from . import _ffi_api
from tvm_ffi import Shape as ShapeTuple
+
+
+def enabled(target):
+ """Whether module runtime is enabled for target
+
+ Parameters
+ ----------
+ target : str or Dict[str, Any] or tvm_tirx.target.Target
+ The target device type.
+
+ Returns
+ -------
+ enabled : bool
+ Whether runtime is enabled.
+
+ Examples
+ --------
+ The following code checks if gpu is enabled.
+
+ >>> tvm_tirx.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()
diff --git a/tirx/python/tvm_tirx/runtime/executable.py b/tirx/python/tvm_tirx/runtime/executable.py
deleted file mode 100644
index 56e1a3f482..0000000000
--- a/tirx/python/tvm_tirx/runtime/executable.py
+++ /dev/null
@@ -1,175 +0,0 @@
-# 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_tirx
-from tvm_tirx.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_tirx.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_tirx.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,
- )
diff --git a/tirx/python/tvm_tirx/support/cc.py b/tirx/python/tvm_tirx/support/cc.py
index 476a9b9b97..5080bc308c 100644
--- a/tirx/python/tvm_tirx/support/cc.py
+++ b/tirx/python/tvm_tirx/support/cc.py
@@ -296,14 +296,16 @@ def cross_compiler(
from tvm_tirx.support import cc, ndk
# export using arm gcc
+ from tvm_tirx.codegen import Executable
+
mod = build_runtime_module()
- mod.export_library(path_dso,
+ Executable(mod).export_library(path_dso,
fcompile=cc.cross_compiler("arm-linux-gnueabihf-gcc"))
# specialize ndk compilation options.
specialized_ndk = cc.cross_compiler(
ndk.create_shared,
["--sysroot=/path/to/sysroot", "-shared", "-fPIC", "-lm"])
- mod.export_library(path_dso, fcompile=specialized_ndk)
+ Executable(mod).export_library(path_dso, fcompile=specialized_ndk)
"""
base_options = [] if options is None else options
kwargs = {}
diff --git a/tirx/python/tvm_tirx/support/ndk.py b/tirx/python/tvm_tirx/support/ndk.py
index 85b6ecb92f..b525abef2a 100644
--- a/tirx/python/tvm_tirx/support/ndk.py
+++ b/tirx/python/tvm_tirx/support/ndk.py
@@ -159,8 +159,10 @@ def get_global_symbol_section_map(path, *, nm=None) -> dict[str, str]:
@register_global_func("s_tir.meta_schedule.builder.export_ndk")
def _ndk_export(mod):
+ from tvm_tirx.codegen import Executable # pylint: disable=import-outside-toplevel
+
tmp_dir = tempfile.mkdtemp()
binary_name = "tmp_binary.so"
binary_path = Path(tmp_dir) / binary_name
- mod.export_library(binary_path, fcompile=create_shared)
+ Executable(mod).export_library(binary_path, fcompile=create_shared)
return str(binary_path)
diff --git a/tirx/python/tvm_tirx/tirx/_compile.py b/tirx/python/tvm_tirx/tirx/_compile.py
index 6d2a6f18b1..f88c114ee2 100644
--- a/tirx/python/tvm_tirx/tirx/_compile.py
+++ b/tirx/python/tvm_tirx/tirx/_compile.py
@@ -22,8 +22,8 @@ from tvm_ffi import DLDeviceType
import tvm_tirx
from tvm_tirx import ir
+from tvm_tirx.codegen import Executable
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.runtime import Executable
from tvm_tirx.target import Target
from tvm_tirx.tirx import Function
@@ -180,7 +180,7 @@ def compile( # pylint: disable=redefined-builtin
Returns
-------
- tvm_tirx.runtime.Executable
+ tvm_tirx.codegen.Executable
An executable combining both host and device code.
"""
# Convert Function to IRModule
diff --git a/tirx/src/runtime/timer.cc b/tirx/src/runtime/timer.cc
deleted file mode 100644
index 16014401b8..0000000000
--- a/tirx/src/runtime/timer.cc
+++ /dev/null
@@ -1,276 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file src/runtime/timer.cc
- * \brief Runtime timer primitives: Timer, WrapTimeEvaluator.
- */
-
-#include <tvm/ffi/extra/module.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/c_backend_api.h>
-#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/runtime/timer.h>
-
-#include <chrono>
-#include <mutex>
-#include <set>
-#include <sstream>
-#include <thread>
-#if defined(_M_X64) || defined(__x86_64__)
-#include <immintrin.h>
-#endif
-
-namespace tvm {
-namespace runtime {
-
-class DefaultTimerObj : public TimerObj {
- public:
- virtual void Start() {
- DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
- start_ = std::chrono::high_resolution_clock::now();
- }
- virtual void Stop() {
- DeviceAPI::Get(device_)->StreamSync(device_, nullptr);
- duration_ = std::chrono::high_resolution_clock::now() - start_;
- }
- virtual int64_t SyncAndGetElapsedNanos() { return duration_.count(); }
- virtual ~DefaultTimerObj() {}
-
- explicit DefaultTimerObj(Device dev) : device_(dev) {}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("runtime.DefaultTimer", DefaultTimerObj, TimerObj);
-
- private:
- std::chrono::high_resolution_clock::time_point start_;
- std::chrono::duration<int64_t, std::nano> duration_;
- Device device_;
-};
-
-static Timer DefaultTimer(Device dev) {
- return Timer(ffi::UnsafeInit{}, ffi::make_object<DefaultTimerObj>(dev));
-}
-
-class CPUTimerObj : public TimerObj {
- public:
- virtual void Start() { start_ = std::chrono::high_resolution_clock::now(); }
- virtual void Stop() { duration_ = std::chrono::high_resolution_clock::now() - start_; }
- virtual int64_t SyncAndGetElapsedNanos() { return duration_.count(); }
- virtual ~CPUTimerObj() {}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("runtime.CPUTimer", CPUTimerObj, TimerObj);
-
- private:
- std::chrono::high_resolution_clock::time_point start_;
- std::chrono::duration<int64_t, std::nano> duration_;
-};
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("runtime.timer.cpu", [](Device dev) {
- return Timer(ffi::UnsafeInit{}, ffi::make_object<CPUTimerObj>());
- });
-}
-
-Timer Timer::Start(Device dev) {
- // Function-local statics: thread-safe lazy init (C++11 magic statics),
- // visible only to this function.
- static std::set<DLDeviceType> seen_devices;
- static std::mutex seen_devices_lock;
- auto f = tvm::ffi::Function::GetGlobal(std::string("runtime.timer.") +
- DLDeviceType2Str(dev.device_type));
- if (!f.has_value()) {
- {
- std::lock_guard<std::mutex> lock(seen_devices_lock);
- if (seen_devices.find(dev.device_type) == seen_devices.end()) {
- LOG(WARNING)
- << "No timer implementation for " << DLDeviceType2Str(dev.device_type)
- << ", using default timer instead. It may be inaccurate or have extra overhead.";
- seen_devices.insert(dev.device_type);
- }
- }
- Timer t = DefaultTimer(dev);
- t->Start();
- return t;
- } else {
- Timer t = f->operator()(dev).cast<Timer>();
- t->Start();
- return t;
- }
-}
-
-static ffi::Function WrapTimeEvaluator(ffi::Function pf, Device dev, int number, int repeat,
- int min_repeat_ms, int limit_zero_time_iterations,
- int cooldown_interval_ms, int repeats_to_cooldown,
- int cache_flush_bytes, ffi::Function f_preproc) {
- TVM_FFI_ICHECK(pf != nullptr);
-
- auto ftimer = [pf, dev, number, repeat, min_repeat_ms, limit_zero_time_iterations,
- cooldown_interval_ms, repeats_to_cooldown, cache_flush_bytes,
- f_preproc](const ffi::AnyView* args, int num_args, ffi::Any* rv) mutable {
- ffi::Any temp;
- std::ostringstream os;
- // skip first time call, to activate lazy compilation components.
- pf.CallPacked(args, num_args, &temp);
-
- // allocate two large arrays to flush L2 cache
- ffi::Optional<ffi::Tensor> arr1, arr2;
- if (cache_flush_bytes > 0) {
- arr1 = tvm::runtime::Empty({cache_flush_bytes / 4}, {kDLInt, 32, 1}, dev);
- arr2 = tvm::runtime::Empty({cache_flush_bytes / 4}, {kDLInt, 32, 1}, dev);
- }
-
- DeviceAPI::Get(dev)->StreamSync(dev, nullptr);
-
- for (int i = 0; i < repeat; ++i) {
- if (f_preproc != nullptr) {
- f_preproc.CallPacked(args, num_args, &temp);
- }
- double duration_ms = 0.0;
- int absolute_zero_times = 0;
- do {
- if (duration_ms > 0.0) {
- const double golden_ratio = 1.618;
- number = static_cast<int>(
- std::max((min_repeat_ms / (duration_ms / number) + 1), number * golden_ratio));
- }
- if (cache_flush_bytes > 0) {
- tvm::runtime::CopyFromTo(arr2.value().GetDLTensorPtr(), arr1.value().GetDLTensorPtr());
- }
- DeviceAPI::Get(dev)->StreamSync(dev, nullptr);
- // start timing
- Timer t = Timer::Start(dev);
- for (int j = 0; j < number; ++j) {
- pf.CallPacked(args, num_args, &temp);
- }
- t->Stop();
- int64_t t_nanos = t->SyncAndGetElapsedNanos();
- if (t_nanos == 0) absolute_zero_times++;
- duration_ms = t_nanos / 1e6;
- } while (duration_ms < min_repeat_ms && absolute_zero_times < limit_zero_time_iterations);
-
- double speed = duration_ms / 1e3 / number;
- os.write(reinterpret_cast<char*>(&speed), sizeof(speed));
-
- if (cooldown_interval_ms > 0 && (i % repeats_to_cooldown) == 0) {
- std::this_thread::sleep_for(std::chrono::milliseconds(cooldown_interval_ms));
- }
- }
-
- std::string blob = os.str();
- // return the time.
- *rv = ffi::Bytes(std::move(blob));
- };
- return ffi::Function::FromPacked(ftimer);
-}
-
-/*!
- * \brief Flush the cache.
- * \param addr The address of data we want to flush
- * \param len The length of data
- */
-/*
- * When we are in the tuning of TVM, we will make TVM occupy
- * the cache fully and doesn't flush it during iteration.
- * This has problems then in e2e testing, since arrays that
- * we assume exist in cache (ie. weights) are evicted during e2e runs,
- * which leads to lower performance.
- */
-inline void CPUCacheFlushImpl(const char* addr, unsigned int len) {
-#if (defined(_M_X64) || defined(__x86_64__) || defined(__aarch64__))
-
-#if defined(__aarch64__)
- size_t ctr_el0 = 0;
- asm volatile("mrs %0, ctr_el0" : "=r"(ctr_el0));
- const size_t cache_line = 4 << ((ctr_el0 >> 16) & 15);
-#else
- const size_t cache_line = 64;
-#endif
-
- if (addr == nullptr || len <= 0) {
- return;
- }
-
- for (uintptr_t uptr = (uintptr_t)addr & ~(cache_line - 1); uptr < (uintptr_t)addr + len;
- uptr += cache_line) {
-#if defined(__aarch64__)
- asm volatile("dc civac, %0\n\t" : : "r"(reinterpret_cast<const void*>(uptr)) : "memory");
-#else
- _mm_clflush(reinterpret_cast<const void*>(uptr));
-#endif
- }
-
-#if defined(__aarch64__)
- asm volatile("dmb ishst" : : : "memory");
-#endif
-
-#endif
-}
-
-inline void CPUCacheFlush(int begin_index, const ffi::PackedArgs& args) {
- for (int i = begin_index; i < args.size(); i++) {
- CPUCacheFlushImpl(static_cast<char*>((args[i].cast<DLTensor*>()->data)),
- ffi::GetDataSize(*(args[i].cast<DLTensor*>())));
- }
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("runtime.LocalTimeEvaluator",
- [](ffi::Optional<ffi::Module> opt_mod, std::string name, int device_type, int device_id,
- int number, int repeat, int min_repeat_ms, int limit_zero_time_iterations,
- int cooldown_interval_ms, int repeats_to_cooldown, int cache_flush_bytes,
- std::string f_preproc_name) {
- Device dev{static_cast<DLDeviceType>(device_type), device_id};
- auto pf = opt_mod.has_value() ? opt_mod.value()->GetFunction(name)
- : ffi::Function::GetGlobal(name);
- TVM_FFI_ICHECK(pf.has_value()) << "Cannot find " << name;
- ffi::Function f_preproc;
- if (!f_preproc_name.empty()) {
- f_preproc = ffi::Function::GetGlobalRequired(f_preproc_name);
- }
- return WrapTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
- limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc);
- })
- .def("runtime.ModuleTimeEvaluator",
- [](ffi::Module mod, std::string name, ffi::AnyView device, int number, int repeat,
- int min_repeat_ms, int limit_zero_time_iterations, int cooldown_interval_ms,
- int repeats_to_cooldown, int cache_flush_bytes, std::string f_preproc_name) {
- if (std::string(mod->kind()) == "rpc") {
- return ffi::Function::GetGlobalRequired("rpc.ModuleTimeEvaluator")(
- mod, name, device, number, repeat, min_repeat_ms,
- limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
- cache_flush_bytes, f_preproc_name)
- .cast<ffi::Function>();
- }
- auto dev = device.cast<DLDevice>();
- return ffi::Function::GetGlobalRequired("runtime.LocalTimeEvaluator")(
- mod, name, static_cast<int>(dev.device_type), dev.device_id, number, repeat,
- min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc_name)
- .cast<ffi::Function>();
- })
- .def_packed("runtime.CacheFlushCPUArguments",
- [](ffi::PackedArgs args, ffi::Any* rv) { CPUCacheFlush(1, args); });
-}
-
-} // namespace runtime
-} // namespace tvm
diff --git a/tirx/tests/cpp/CMakeLists.txt b/tirx/tests/cpp/CMakeLists.txt
index f1fae95c49..95802d605b 100644
--- a/tirx/tests/cpp/CMakeLists.txt
+++ b/tirx/tests/cpp/CMakeLists.txt
@@ -38,7 +38,7 @@ endif()
include(GoogleTest)
file(GLOB_RECURSE _test_sources CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/*.cc")
add_executable(tirx_cpptest ${_test_sources})
-target_link_libraries(tirx_cpptest PRIVATE tvm_tirx ${_gtest_targets} ${CMAKE_DL_LIBS})
+target_link_libraries(tirx_cpptest PRIVATE tvm_tirx_compiler ${_gtest_targets} ${CMAKE_DL_LIBS})
if(TVM_TIRX_USE_THREADS)
target_link_libraries(tirx_cpptest PRIVATE Threads::Threads)
endif()
diff --git a/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py b/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
index d076481def..ea92272ea6 100644
--- a/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
+++ b/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
@@ -14,21 +14,24 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Post-install checks for a built TVM wheel.
+"""Post-install runtime and AOT checks.
These live in ``tests/python/all-platform-minimal-test`` so the standard suite and
the cibuildwheel ``test-command`` run a single pytest invocation. The assertions
-here are wheel-specific things the rest of the suite cannot check -- that LLVM is
+for wheel packaging check things the rest of the suite cannot check -- that LLVM is
enabled (the other LLVM test merely *skips* when LLVM is absent) and that the CUDA
runtime library got bundled -- so each is gated behind a ``TVM_WHEEL_EXPECT_*`` env
var and SKIPS unless that var is set. cibuildwheel sets the vars (see
``CIBW_TEST_ENVIRONMENT`` in ``.github/actions/build-wheel-for-publish``); ordinary
source-build CI (e.g. ``main.yml``) leaves them unset, so these tests skip there and
-never fail a non-wheel / non-LLVM / non-CUDA build.
+never fail a non-wheel / non-LLVM / non-CUDA build. The CPU AOT case also runs
+on Linux/macOS source builds with LLVM enabled.
"""
import glob
import os
+import subprocess
+import sys
from pathlib import Path
import pytest
@@ -54,3 +57,37 @@ def test_cuda_runtime_present():
str(libdir / "tvm_runtime_cuda.*")
)
assert present, "CUDA runtime expected but not bundled in tvm/lib"
+
+
+def test_cpu_aot(tmp_path):
+ if sys.platform not in ("linux", "darwin") or not tvm_tirx.runtime.enabled("llvm"):
+ pytest.skip("CPU AOT requires Linux/macOS and LLVM")
+ from tvm_tirx.script import tirx as T
+
+ @T.function
+ def add_one(A: T.Tensor((16,), "int32", align=4), B: T.Tensor((16,), "int32", align=4)):
+ for i in range(16):
+ B[i] = A[i] + 1
+
+ module = tvm_tirx.tirx.compile(add_one, target="llvm")
+ path = tmp_path / "model.so"
+ module.export_library(path)
+ subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ """
+import sys
+import numpy as np
+import tvm_ffi
+module = tvm_ffi.load_module(sys.argv[1])
+a = np.arange(16, dtype="int32")
+b = np.empty_like(a)
+module["add_one"](tvm_ffi.from_dlpack(a), tvm_ffi.from_dlpack(b))
+np.testing.assert_array_equal(b, a + 1)
+assert "tvm" not in sys.modules and "tvm_tirx" not in sys.modules
+""",
+ str(path),
+ ],
+ check=True,
+ )
diff --git a/tirx/tests/python/backend/cuda/codegen/test_fastmath.py b/tirx/tests/python/backend/cuda/codegen/test_fastmath.py
index c8bfdc2858..ad420eb38b 100644
--- a/tirx/tests/python/backend/cuda/codegen/test_fastmath.py
+++ b/tirx/tests/python/backend/cuda/codegen/test_fastmath.py
@@ -26,8 +26,8 @@ import pytest
import tvm_tirx
import tvm_tirx.testing
import tvm_tirx.tirx as tirx
+from tvm_tirx.codegen import Executable
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.runtime.executable import Executable
from tvm_tirx.script import tirx as T
from tvm_tirx.support.nvcc import have_fp16
from tvm_tirx.testing import env
diff --git a/tirx/tests/python/codegen/c/test_basic.py b/tirx/tests/python/codegen/c/test_basic.py
index 28349fb7d9..9dd738f428 100644
--- a/tirx/tests/python/codegen/c/test_basic.py
+++ b/tirx/tests/python/codegen/c/test_basic.py
@@ -230,7 +230,7 @@ def test_workspace_allocation_cast():
assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
temp = utils.tempdir()
- built.export_library(temp.relpath("workspace.so"))
+ tvm_tirx.codegen.Executable(built).export_library(temp.relpath("workspace.so"))
def test_local_alloc_buffer_uses_plain_c_pointer():
@@ -248,7 +248,7 @@ def test_local_alloc_buffer_uses_plain_c_pointer():
temp = utils.tempdir()
path_dso = temp.relpath("local_alloc.so")
- built.export_library(path_dso)
+ tvm_tirx.codegen.Executable(built).export_library(path_dso)
loaded = tvm_tirx.runtime.load_module(path_dso)
data = tvm_tirx.runtime.tensor(np.array([1.0], dtype="float32"))
@@ -288,7 +288,7 @@ def test_if_then_else_avoids_extraneous_parentheses():
temp = utils.tempdir()
path_dso = temp.relpath("if_then_else.so")
- built.export_library(path_dso)
+ tvm_tirx.codegen.Executable(built).export_library(path_dso)
loaded = tvm_tirx.runtime.load_module(path_dso)
a = tvm_tirx.runtime.tensor(np.arange(8, dtype="int32"))
diff --git a/tirx/tests/python/codegen/test_executable.py b/tirx/tests/python/codegen/test_executable.py
new file mode 100644
index 0000000000..c34036065b
--- /dev/null
+++ b/tirx/tests/python/codegen/test_executable.py
@@ -0,0 +1,145 @@
+# 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.
+"""Executable export, JIT, and static import coverage."""
+
+import shlex
+import subprocess
+import sys
+
+import pytest
+
+import tvm_tirx
+import tvm_tirx.testing
+from tvm_tirx.codegen import Executable
+from tvm_tirx.script import tirx as T
+from tvm_tirx.support import cc
+from tvm_tirx.testing import env
+
+
+def test_import_static_library(source_module, native_compile_options, tmp_path):
+ source = tmp_path / "other.cc"
+ source.write_text(
+ "#include <tvm/ffi/function.h>\n"
+ "int other() { return 17; }\n"
+ "TVM_FFI_DLL_EXPORT_TYPED_FUNC(other, other);\n"
+ )
+ obj = tmp_path / "other.o"
+ cc.create_shared(str(obj), [str(source)], options=["-c", *native_compile_options])
+ imported = tvm_tirx.codegen.load_static_library(str(obj), ["other"])
+ assert imported.implements_function("other")
+ assert imported.is_compilation_exportable()
+ source_module.import_module(imported)
+ path = tmp_path / "combined.so"
+ Executable(source_module).export_library(str(path), options=native_compile_options)
+ loaded = tvm_tirx.runtime.load_module(str(path))
+ assert loaded.kind == "library"
+ assert len(loaded.imports) == 0
+ for name, expected in (("answer", 42), ("other", 17)):
+ assert loaded.implements_function(name)
+ assert loaded.get_function(name)() == expected
+ assert not loaded.is_compilation_exportable()
+
+
+def test_executable_init(source_module):
+ executable = Executable(source_module)
+ assert executable.mod is source_module
+ assert executable._jitted_mod is None
+
+
+def test_executable_getitem(source_module, native_compile_options):
+ executable = Executable(source_module)
+ executable.jit(options=native_compile_options)
+ assert executable["answer"]() == 42
+
+
+def test_executable_jit_already_runnable(native_library):
+ module = tvm_tirx.runtime.load_module(str(native_library))
+ executable = Executable(module)
+ assert executable.jit() is module
+ assert executable.jit(force_recompile=True) is module
+ assert executable["answer"]() == 42
+
+
+def test_executable_jit_force_recompile(source_module, native_compile_options):
+ executable = Executable(source_module)
+ first = executable.jit(options=native_compile_options)
+ assert executable.jit() is first
+ rebuilt = executable.jit(force_recompile=True, options=native_compile_options)
+ assert rebuilt is not first
+ assert first["answer"]() == rebuilt["answer"]() == 42
+
+
+@pytest.mark.parametrize("use_workspace", [False, True])
+def test_executable_export_library(source_module, native_compile_options, tmp_path, use_workspace):
+ kwargs = {"options": native_compile_options}
+ if use_workspace:
+ workspace = tmp_path / "workspace"
+ workspace.mkdir()
+ kwargs["workspace_dir"] = str(workspace)
+ path = tmp_path / "module.so"
+ Executable(source_module).export_library(str(path), **kwargs)
+ assert path.is_file()
+ assert tvm_tirx.runtime.load_module(str(path))["answer"]() == 42
+
+
+@pytest.mark.gpu
+def test_cuda_aot(tmp_path):
+ if sys.platform != "linux" or not env.has_cuda() or not tvm_tirx.runtime.enabled("llvm"):
+ pytest.skip("CUDA AOT requires Linux CUDA and an LLVM host compiler")
+ pytest.importorskip("torch")
+
+ @T.function
+ def add_one(A: T.Tensor((16,), "int32"), B: T.Tensor((16,), "int32")):
+ with T.launch_thread("blockIdx.x", 1):
+ with T.launch_thread("threadIdx.x", 16) as i:
+ B[i] = A[i] + 1
+
+ module = tvm_tirx.tirx.compile(add_one, target=tvm_tirx.target.Target("cuda", host="llvm"))
+ path = tmp_path / "model.so"
+ options = shlex.split(
+ subprocess.check_output(
+ [sys.executable, "-m", "tvm_tirx.config", "--ldflags", "--libs", "runtime", "cuda"],
+ text=True,
+ )
+ )
+ library_dir = subprocess.check_output(
+ [sys.executable, "-m", "tvm_tirx.config", "--libdir"], text=True
+ ).strip()
+ module.export_library(path, options=[*options, "-Wl,-rpath," + library_dir])
+ subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ """
+import sys
+import torch
+import tvm_ffi
+module = tvm_ffi.load_module(sys.argv[1])
+a = torch.arange(16, dtype=torch.int32, device="cuda")
+b = torch.empty_like(a)
+module["add_one"](a, b)
+torch.testing.assert_close(b, a + 1)
+assert "tvm" not in sys.modules and "tvm_tirx" not in sys.modules
+""",
+ str(path),
+ ],
+ check=True,
+ )
+
+
+if __name__ == "__main__":
+ tvm_tirx.testing.main()
diff --git a/tirx/tests/python/conftest.py b/tirx/tests/python/conftest.py
new file mode 100644
index 0000000000..42a3cb04b1
--- /dev/null
+++ b/tirx/tests/python/conftest.py
@@ -0,0 +1,78 @@
+# 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.
+"""Shared native module fixtures for runtime and code generation coverage."""
+
+import pytest
+
+
+@pytest.fixture(scope="module")
+def native_source():
+ return r"""
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/function.h>
+void add_one(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
+ for (int64_t i = 0; i < a.size(0); ++i) {
+ static_cast<float*>(b.data_ptr())[i] = static_cast<float*>(a.data_ptr())[i] + 1;
+ }
+}
+int64_t answer() { return 42; }
+void callback(tvm::ffi::Function func) { func(); }
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(main, add_one);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(answer, answer);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(callback, callback);
+"""
+
+
+@pytest.fixture(scope="module")
+def native_compile_options():
+ import tvm_ffi
+
+ return [
+ "-std=c++17",
+ "-I" + tvm_ffi.libinfo.find_include_path(),
+ "-I" + tvm_ffi.libinfo.find_dlpack_include_path(),
+ ]
+
+
+@pytest.fixture
+def source_module(native_source):
+ import tvm_tirx
+
+ return tvm_tirx.runtime._ffi_api.CSourceModuleCreate(
+ native_source, "cc", ["main", "answer", "callback"], None
+ )
+
+
+@pytest.fixture(scope="module")
+def native_object(tmp_path_factory, native_source, native_compile_options):
+ from tvm_tirx.support import cc
+
+ directory = tmp_path_factory.mktemp("runtime_module")
+ source = directory / "module.cc"
+ source.write_text(native_source)
+ path = directory / "module.o"
+ cc.create_shared(str(path), [str(source)], options=["-c", *native_compile_options])
+ return path
+
+
+@pytest.fixture(scope="module")
+def native_library(native_object):
+ from tvm_tirx.support import cc
+
+ path = native_object.with_suffix(".so")
+ cc.create_shared(str(path), [str(native_object)])
+ return path
diff --git a/tirx/tests/python/runtime/test_module.py b/tirx/tests/python/runtime/test_module.py
index c1dc2bc86b..dbcc77d73a 100644
--- a/tirx/tests/python/runtime/test_module.py
+++ b/tirx/tests/python/runtime/test_module.py
@@ -14,12 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Runtime module loading, properties, exports, timing, and Executable wrappers."""
+"""Runtime and FFI module loading and properties."""
import os
import subprocess
import sys
-import time
import numpy as np
import pytest
@@ -27,62 +26,9 @@ import tvm_ffi
import tvm_tirx
import tvm_tirx.testing
-from tvm_tirx.runtime import Executable
-from tvm_tirx.runtime.module import BenchmarkResult
from tvm_tirx.support import cc
-@pytest.fixture(scope="module")
-def native_source():
- return r"""
-#include <tvm/ffi/container/tensor.h>
-#include <tvm/ffi/function.h>
-void add_one(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
- for (int64_t i = 0; i < a.size(0); ++i) {
- static_cast<float*>(b.data_ptr())[i] = static_cast<float*>(a.data_ptr())[i] + 1;
- }
-}
-int64_t answer() { return 42; }
-void callback(tvm::ffi::Function func) { func(); }
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(main, add_one);
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(answer, answer);
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(callback, callback);
-"""
-
-
-@pytest.fixture(scope="module")
-def native_compile_options():
- return [
- "-std=c++17",
- "-I" + tvm_ffi.libinfo.find_include_path(),
- "-I" + tvm_ffi.libinfo.find_dlpack_include_path(),
- ]
-
-
-@pytest.fixture
-def source_module(native_source):
- return tvm_tirx.runtime._ffi_api.CSourceModuleCreate(
- native_source, "cc", ["main", "answer", "callback"], None
- )
-
-
-@pytest.fixture(scope="module")
-def native_object(tmp_path_factory, native_source, native_compile_options):
- directory = tmp_path_factory.mktemp("runtime_module")
- source = directory / "module.cc"
- source.write_text(native_source)
- path = directory / "module.o"
- cc.create_shared(str(path), [str(source)], options=["-c", *native_compile_options])
- return path
-
-
-@pytest.fixture(scope="module")
-def native_library(native_object):
- path = native_object.with_suffix(".so")
- cc.create_shared(str(path), [str(native_object)])
- return path
-
-
@pytest.fixture(scope="module")
def system_library_source(native_source):
return (
@@ -101,6 +47,11 @@ static RegisterSystemLibrary register_system_library;
def test_module_properties(source_module, native_library):
+ assert tvm_tirx.runtime.Module is tvm_ffi.Module
+ assert tvm_tirx.runtime.load_module is tvm_ffi.load_module
+ assert tvm_tirx.runtime.system_lib is tvm_ffi.system_lib
+ assert not hasattr(source_module, "export_library")
+ assert not hasattr(source_module, "time_evaluator")
assert source_module.is_binary_serializable()
assert not source_module.is_runnable()
assert source_module.is_compilation_exportable()
@@ -141,30 +92,6 @@ np.testing.assert_equal(b.numpy(), a.numpy() + 1)
)
-def test_import_static_library(source_module, native_compile_options, tmp_path):
- source = tmp_path / "other.cc"
- source.write_text(
- "#include <tvm/ffi/function.h>\n"
- "int other() { return 17; }\n"
- "TVM_FFI_DLL_EXPORT_TYPED_FUNC(other, other);\n"
- )
- obj = tmp_path / "other.o"
- cc.create_shared(str(obj), [str(source)], options=["-c", *native_compile_options])
- imported = tvm_tirx.runtime.load_static_library(str(obj), ["other"])
- assert imported.implements_function("other")
- assert imported.is_compilation_exportable()
- source_module.import_module(imported)
- path = tmp_path / "combined.so"
- source_module.export_library(str(path), options=native_compile_options)
- loaded = tvm_tirx.runtime.load_module(str(path))
- assert loaded.kind == "library"
- assert len(loaded.imports) == 0
- for name, expected in (("answer", 42), ("other", 17)):
- assert loaded.implements_function(name)
- assert loaded.get_function(name)() == expected
- assert not loaded.is_compilation_exportable()
-
-
@pytest.mark.skipif(sys.platform == "win32", reason="system library registration on Unix")
def test_system_library(system_library_source, native_compile_options, tmp_path):
source = tmp_path / "system.cc"
@@ -194,72 +121,5 @@ np.testing.assert_equal(b.numpy(), a.numpy() + 1)
)
-def test_min_repeat_ms(native_library):
- calls = []
-
- def callback():
- time.sleep(0.1)
- calls.append(1)
-
- module = tvm_tirx.runtime.load_module(str(native_library))
- module.time_evaluator("callback", tvm_tirx.cpu(), number=1, repeat=1)(callback)
- assert len(calls) == 2
- module.time_evaluator("callback", tvm_tirx.cpu(), number=1, repeat=1, min_repeat_ms=1000)(
- callback
- )
- assert len(calls) > 12
-
-
-def test_benchmark_result():
- result = BenchmarkResult([1, 2, 2, 5])
- assert result.mean == 2.5
- assert result.median == 2.0
- assert result.min == 1
- assert result.max == 5
- assert result.std == 1.5
-
-
-def test_executable_init(source_module):
- executable = Executable(source_module)
- assert executable.mod is source_module
- assert executable._jitted_mod is None
-
-
-def test_executable_getitem(source_module, native_compile_options):
- executable = Executable(source_module)
- executable.jit(options=native_compile_options)
- assert executable["answer"]() == 42
-
-
-def test_executable_jit_already_runnable(native_library):
- module = tvm_tirx.runtime.load_module(str(native_library))
- executable = Executable(module)
- assert executable.jit() is module
- assert executable.jit(force_recompile=True) is module
- assert executable["answer"]() == 42
-
-
-def test_executable_jit_force_recompile(source_module, native_compile_options):
- executable = Executable(source_module)
- first = executable.jit(options=native_compile_options)
- assert executable.jit() is first
- rebuilt = executable.jit(force_recompile=True, options=native_compile_options)
- assert rebuilt is not first
- assert first["answer"]() == rebuilt["answer"]() == 42
-
-
-@pytest.mark.parametrize("use_workspace", [False, True])
-def test_executable_export_library(source_module, native_compile_options, tmp_path, use_workspace):
- kwargs = {"options": native_compile_options}
- if use_workspace:
- workspace = tmp_path / "workspace"
- workspace.mkdir()
- kwargs["workspace_dir"] = str(workspace)
- path = tmp_path / "module.so"
- Executable(source_module).export_library(str(path), **kwargs)
- assert path.is_file()
- assert tvm_tirx.runtime.load_module(str(path))["answer"]() == 42
-
-
if __name__ == "__main__":
tvm_tirx.testing.main()
diff --git a/tirx/tests/python/runtime/test_tirx_config.py b/tirx/tests/python/runtime/test_tirx_config.py
index 648ef0309a..ae68e8799f 100644
--- a/tirx/tests/python/runtime/test_tirx_config.py
+++ b/tirx/tests/python/runtime/test_tirx_config.py
@@ -17,8 +17,11 @@
"""The config command must work without importing native runtime dependencies."""
import os
+import shlex
+import shutil
import subprocess
import sys
+import sysconfig
from pathlib import Path
import tvm_tirx
@@ -35,7 +38,7 @@ def test_config_without_native_imports(tmp_path):
" raise AssertionError(fullname)\n"
"sys.meta_path.insert(0, BlockNative())\n"
"def audit(event, args):\n"
- " if event == 'ctypes.dlopen': raise AssertionError(args)\n"
+ " if event == 'ctypes.dlopen' and args[0] is not None: raise AssertionError(args)\n"
"sys.addaudithook(audit)\n"
)
env = dict(os.environ)
@@ -45,14 +48,39 @@ def test_config_without_native_imports(tmp_path):
str(Path(tvm_tirx.__file__).parent.parent),
)
)
- command = [sys.executable, "-m", "tvm_tirx.config"]
- include = subprocess.check_output([*command, "--includedir"], env=env, text=True).strip()
- assert (Path(include) / "tvm" / "runtime").is_dir()
- libraries = subprocess.check_output([*command, "--libfiles"], env=env, text=True).splitlines()
- assert len(libraries) == 2
- assert all(Path(library).is_file() for library in libraries)
- cmake_dir = subprocess.check_output([*command, "--cmakedir"], env=env, text=True).strip()
- assert (Path(cmake_dir) / "tvm_tirx-config.cmake").is_file()
+ for entry in ("module", "console"):
+ command = [sys.executable, "-m", "tvm_tirx.config"]
+ if entry == "console":
+ executable = Path(sysconfig.get_path("scripts")) / "tvm-tirx-config"
+ if not executable.is_file():
+ continue
+ command = [str(executable)]
+ include = subprocess.check_output([*command, "--includedir"], env=env, text=True).strip()
+ assert (Path(include) / "tvm" / "runtime").is_dir()
+ libraries = subprocess.check_output(
+ [*command, "--libfiles", "runtime"], env=env, text=True
+ ).strip()
+ libraries = shlex.split(libraries)
+ assert len(libraries) == 2
+ assert all(Path(library).is_file() for library in libraries)
+ cmake_dir = subprocess.check_output([*command, "--cmakedir"], env=env, text=True).strip()
+ assert (Path(cmake_dir) / "tvm_tirx-config.cmake").is_file()
+ if shutil.which("c++") and sys.platform != "win32":
+ _check_config_public_headers_compile()
+
+
+def _check_config_public_headers_compile():
+ flags = shlex.split(
+ subprocess.check_output(
+ [sys.executable, "-m", "tvm_tirx.config", "--cxxflags", "runtime"], text=True
+ )
+ )
+ subprocess.run(
+ ["c++", *flags, "-x", "c++", "-fsyntax-only", "-"],
+ input="#include <tvm/runtime/base.h>\n#include <tvm/ffi/container/tensor.h>\n",
+ text=True,
+ check=True,
+ )
if __name__ == "__main__":
diff --git a/tirx/tests/python/support/test_config.py b/tirx/tests/python/support/test_config.py
index 76788bc50f..009847d985 100644
--- a/tirx/tests/python/support/test_config.py
+++ b/tirx/tests/python/support/test_config.py
@@ -19,10 +19,12 @@
import importlib.util
import json
import os
+import shlex
import shutil
import subprocess
import sys
from pathlib import Path
+from types import SimpleNamespace
import pytest
@@ -39,14 +41,21 @@ def _installation(prefix, libdir="lib", windows=False):
library_dir = prefix / libdir
cmake_dir = library_dir / "cmake/tvm_tirx"
cmake_dir.mkdir(parents=True)
- for name in ("tvm_tirx", "tvm_runtime"):
+ for name in ("tvm_tirx_compiler", "tvm_runtime"):
filenames = (f"{name}.dll", f"{name}.lib") if windows else (f"lib{name}.so",)
for filename in filenames:
(library_dir / filename).touch()
return library_dir
-def test_exact_library_directory_and_relocation(tmp_path):
+def test_exact_library_directory_and_relocation(tmp_path, monkeypatch, capsys):
+ _check_component_closure_and_artifacts(tmp_path / "closure", monkeypatch)
+ for platform in ("linux", "darwin"):
+ _check_cli_link_flags_and_raw_inputs(tmp_path / platform, monkeypatch, capsys, platform)
+ _check_unavailable_static_is_not_silently_shared(tmp_path / "static")
+ _check_runtime_only_lib64_and_explicit_prefix_ignore_environment(
+ tmp_path / "runtime", monkeypatch
+ )
prefix = tmp_path / "native install"
libdir = _installation(prefix, "lib/x86_64-linux-gnu")
(libdir / "libtvm_runtime_static.a").touch()
@@ -54,7 +63,7 @@ def test_exact_library_directory_and_relocation(tmp_path):
(libdir / "libtvm_ffi.so").touch()
result = config.configuration(library_dir=libdir, platform="Linux")
assert result["include_dir"] == str(prefix / "include")
- assert set(result["libraries"]) == {"tvm_tirx", "tvm_runtime"}
+ assert set(result["libraries"]) == {"tvm_tirx_compiler", "tvm_runtime"}
assert set(result) == {"version", "include_dir", "library_dir", "cmake_dir", "libraries"}
moved = tmp_path / "relocated"
shutil.copytree(prefix, moved)
@@ -70,7 +79,7 @@ def test_prefix_selects_installed_lib64(tmp_path):
result = config.configuration(prefix=prefix, platform="Linux")
assert result["library_dir"] == str(libdir)
result = subprocess.run(
- [sys.executable, str(_CONFIG_PATH), "--prefix", str(prefix), "--cmakedir"],
+ [sys.executable, str(_CONFIG_PATH), "--install-prefix", str(prefix), "--cmakedir"],
check=True,
text=True,
capture_output=True,
@@ -80,14 +89,14 @@ def test_prefix_selects_installed_lib64(tmp_path):
config.configuration(library_dir=prefix / "lib", platform="Linux")
-@pytest.mark.parametrize("name", ["tvm_tirx", "tvm_runtime"])
+@pytest.mark.parametrize("name", ["tvm_tirx_compiler", "tvm_runtime"])
def test_windows_library_pairing(tmp_path, name):
libdir = _installation(tmp_path / "windows", windows=True)
(libdir / "tvm_runtime_static.lib").touch()
(libdir / "tvm_runtime_cuda.dll").touch()
(libdir / "tvm_runtime_cuda.lib").touch()
result = config.configuration(library_dir=libdir, platform="Windows")
- assert set(result["libraries"]) == {"tvm_tirx", "tvm_runtime"}
+ assert set(result["libraries"]) == {"tvm_tirx_compiler", "tvm_runtime"}
assert result["libraries"][name] == {
"location": str(libdir / f"{name}.dll"),
"implib": str(libdir / f"{name}.lib"),
@@ -115,7 +124,7 @@ def test_module_cli_without_native_dependencies(tmp_path):
env = os.environ.copy()
env["PYTHONPATH"] = str(_CONFIG_PATH.parents[1])
result = subprocess.run(
- [sys.executable, "-S", "-m", "tvm_tirx.config", "--prefix", str(prefix), "--json"],
+ [sys.executable, "-S", "-m", "tvm_tirx.config", "--install-prefix", str(prefix), "--json"],
cwd=tmp_path,
env=env,
check=True,
@@ -143,10 +152,10 @@ find_package(tvm_tirx CONFIG REQUIRED)
get_property(imported DIRECTORY PROPERTY IMPORTED_TARGETS)
list(FILTER imported EXCLUDE REGEX "^Python::")
list(SORT imported)
-if(NOT "${imported}" STREQUAL "tvm_runtime::shared;tvm_tirx::header;tvm_tirx::shared")
+if(NOT "${imported}" STREQUAL "tvm_runtime::shared;tvm_tirx::header;tvm_tirx_compiler::shared")
message(FATAL_ERROR "Unexpected imported targets: ${imported}")
endif()
-foreach(target tvm_tirx::header tvm_tirx::shared tvm_runtime::shared)
+foreach(target tvm_tirx::header tvm_tirx_compiler::shared tvm_runtime::shared)
get_target_property(links ${target} INTERFACE_LINK_LIBRARIES)
if(links MATCHES "ffi|OpenMP|Threads")
message(FATAL_ERROR "Unexpected implicit dependency: ${links}")
@@ -170,3 +179,106 @@ endforeach()
text=True,
capture_output=True,
)
+
+
+def _check_component_closure_and_artifacts(tmp_path, monkeypatch):
+ libdir = _installation(tmp_path / "install")
+ for filename in ("libtvm_runtime_cuda.so", "libtvm_runtime_metal.so", "libtvm_ffi.so"):
+ (libdir / filename).touch()
+ monkeypatch.setattr(
+ config,
+ "_ffi_info",
+ lambda: SimpleNamespace(find_libtvm_ffi=lambda: str(libdir / "libtvm_ffi.so")),
+ )
+
+ def names(components):
+ return [path.name for path in config._link_inputs(libdir, "linux", components)]
+
+ assert names(["runtime", "cuda", "runtime", "ffi"]) == [
+ "libtvm_runtime_cuda.so",
+ "libtvm_runtime.so",
+ "libtvm_ffi.so",
+ ]
+ assert names(["compiler"]) == ["libtvm_tirx_compiler.so", "libtvm_runtime.so", "libtvm_ffi.so"]
+ assert names([]) == names(["all"])
+ assert len(names([])) == len(set(names([])))
+ assert names(["ffi"]) == ["libtvm_ffi.so"]
+ with pytest.raises(RuntimeError, match="Component 'missing' is unavailable"):
+ names(["missing"])
+
+
+def _check_cli_link_flags_and_raw_inputs(tmp_path, monkeypatch, capsys, platform):
+ prefix = tmp_path / "native prefix"
+ libdir = _installation(prefix)
+ suffix = "dylib" if platform == "darwin" else "so"
+ for name in ("tvm_runtime", "tvm_tirx_compiler", "tvm_ffi"):
+ (libdir / f"lib{name}.{suffix}").touch()
+ monkeypatch.setattr(
+ config,
+ "_ffi_info",
+ lambda: SimpleNamespace(find_libtvm_ffi=lambda: str(libdir / f"libtvm_ffi.{suffix}")),
+ )
+
+ def run(*options):
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "tvm-tirx-config",
+ "--install-prefix",
+ str(prefix),
+ "--platform",
+ platform,
+ *options,
+ "runtime",
+ ],
+ )
+ config.__main__()
+ return shlex.split(capsys.readouterr().out)
+
+ inputs = run("--libfiles")
+ assert inputs == [
+ str(libdir / f"libtvm_runtime.{suffix}"),
+ str(libdir / f"libtvm_ffi.{suffix}"),
+ ]
+ assert run("--libnames") == [Path(path).name for path in inputs]
+ assert run("--ldflags") == ["-L" + str(libdir)]
+ flags = run("--libs", "--link-shared")
+ if platform == "linux":
+ assert flags == ["-Wl,--push-state,--no-as-needed", *inputs, "-Wl,--pop-state"]
+ else:
+ assert flags == [
+ flag for path in inputs for flag in ("-Xlinker", "-needed_library", "-Xlinker", path)
+ ]
+
+
+def _check_unavailable_static_is_not_silently_shared(tmp_path):
+ prefix = tmp_path / "install"
+ libdir = _installation(prefix)
+ (libdir / "libtvm_runtime_static.a").touch()
+ result = subprocess.run(
+ [
+ sys.executable,
+ str(_CONFIG_PATH),
+ "--install-prefix",
+ str(prefix),
+ "--link-static",
+ "--libs",
+ "runtime",
+ ],
+ capture_output=True,
+ text=True,
+ )
+ assert result.returncode != 0
+ assert "complete static dependency closure" in result.stderr
+ assert not result.stdout
+
+
+def _check_runtime_only_lib64_and_explicit_prefix_ignore_environment(tmp_path, monkeypatch):
+ prefix = tmp_path / "runtime-only"
+ libdir = _installation(prefix, "lib64")
+ (libdir / "libtvm_tirx_compiler.so").unlink()
+ unrelated = _installation(tmp_path / "unrelated")
+ monkeypatch.setenv("TVM_TIRX_LIBRARY_PATH", str(unrelated))
+ assert config._paths(prefix, None) == (prefix, libdir)
+ assert config._available(libdir, "linux") == ["runtime"]
diff --git a/tirx/tests/python/tirx/iket/test_profiler.py b/tirx/tests/python/tirx/iket/test_profiler.py
index 8669194de2..c4a549a3c7 100644
--- a/tirx/tests/python/tirx/iket/test_profiler.py
+++ b/tirx/tests/python/tirx/iket/test_profiler.py
@@ -33,6 +33,7 @@ import pytest
import tvm_tirx
from tvm_tirx.backend.cuda import transforms as cuda_transforms
+from tvm_tirx.codegen import Executable
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
from tvm_tirx.tirx.cuda.iket import IketProfiler
@@ -305,7 +306,7 @@ def _compile(func, *, target=TARGET):
def _cuda_source(executable):
- modules = executable.mod._collect_from_import_tree( # pylint: disable=protected-access
+ modules = Executable(executable.mod)._collect_from_import_tree( # pylint: disable=protected-access
lambda module: module.kind == "cuda"
)
assert len(modules) == 1
@@ -789,7 +790,7 @@ def test_multi_kernel_module_has_no_tvm_control_plane():
assert "__iket_evt_decl_a_1_attrs" in source
assert "__iket_evt_decl_b_2_attrs" in source
assert "plain_entry_kernel" in source
- for module in executable.mod._collect_from_import_tree(lambda _module: True):
+ for module in Executable(executable.mod)._collect_from_import_tree(lambda _module: True):
for control_name in (
"__tvm_iket_get_metadata",
"__tvm_iket_set_context",