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 <[email protected]> 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 + + [email protected]("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 + + [email protected] +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 + + [email protected](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); +""" + + [email protected](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(), + ] + + [email protected] +def source_module(native_source): + import tvm_tirx + + return tvm_tirx.runtime._ffi_api.CSourceModuleCreate( + native_source, "cc", ["main", "answer", "callback"], None + ) + + [email protected](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 + + [email protected](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 [email protected](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); -""" - - [email protected](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(), - ] - - [email protected] -def source_module(native_source): - return tvm_tirx.runtime._ffi_api.CSourceModuleCreate( - native_source, "cc", ["main", "answer", "callback"], None - ) - - [email protected](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 - - [email protected](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 - - [email protected]("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") [email protected]("name", ["tvm_tirx", "tvm_runtime"]) [email protected]("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",
