From e0b8a61001ec6f1c216adef49e98dae5c0ac153c Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Mon, 21 Sep 2026 21:55:53 +0800 Subject: [PATCH 1/4] fix(orcjit): load addon through TVM-FFI --- .../python/tvm_ffi_orcjit/__init__.py | 21 +++-------- .../python/tvm_ffi_orcjit/session.py | 4 +- .../tests/test_library_visibility.py | 37 +++++++++++++++++++ python/tvm_ffi/cpp/extension.py | 4 +- python/tvm_ffi/module.py | 2 +- src/ffi/extra/module.cc | 8 +++- 6 files changed, 53 insertions(+), 23 deletions(-) diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py index e1650303c..b56570a62 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py @@ -29,8 +29,6 @@ """ -import ctypes -import os import platform from pathlib import Path @@ -54,8 +52,8 @@ _lib_path = None for path in _LIB_PATH: if path.exists(): - _ = load_module(str(path)) _lib_path = path + break if _lib_path is None: raise RuntimeError( f"Could not find {_LIB_NAME}. " @@ -63,19 +61,10 @@ f"Please ensure the package is installed correctly." ) -# Keep a second, process-lifetime local handle. RTLD_NODELETE is important for -# modules pinned by keep_module_alive: their object deleters point into JIT code -# owned by this DSO and may run during interpreter shutdown, after Python module -# globals have otherwise released their handles. This does not promote the DSO -# or its statically linked LLVM into the process-global symbol namespace. -if os.name == "posix": - _c_lib = ctypes.CDLL( - str(_lib_path), - mode=ctypes.RTLD_LOCAL | getattr(os, "RTLD_NODELETE", 0), - ) -else: - _dll_directory = os.add_dll_directory(str(_lib_path.parent)) - _c_lib = ctypes.CDLL(str(_lib_path)) +# The TVM-FFI loader uses local symbol scope. Retain the returned module here +# and in TVM-FFI's process-lifetime module registry so registered functions and +# JIT-owned object deleters remain valid through interpreter shutdown. +_lib_module = load_module(_lib_path, keep_module_alive=True) from .session import ExecutionSession, default_session diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py index e00e91866..2bb1098fd 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py @@ -190,8 +190,8 @@ def load_module( ``keep_module_alive`` mirrors :func:`tvm_ffi.load_module`'s option of the same name. When True, the module is inserted into the runtime's process-global module registry, so its JITDylib — and every function - pointer, deleter, and static allocation it owns — stays mapped until - the interpreter unloads ``libtvm_ffi``. Use this when Objects produced + pointer, deleter, and static allocation it owns — stays mapped for the + duration of the process. Use this when Objects produced by the module may outlive the local ``mod`` reference (e.g., a JIT-allocated ``String`` or ``Array`` returned to Python and held past ``del mod``). When False (default), the caller owns the module's diff --git a/addons/tvm_ffi_orcjit/tests/test_library_visibility.py b/addons/tvm_ffi_orcjit/tests/test_library_visibility.py index 8f853313b..bd8833aa7 100644 --- a/addons/tvm_ffi_orcjit/tests/test_library_visibility.py +++ b/addons/tvm_ffi_orcjit/tests/test_library_visibility.py @@ -22,9 +22,11 @@ import platform import shutil import subprocess +import sys import pytest import tvm_ffi_orcjit +from utils import build_test_objects @pytest.mark.skipif(platform.system() == "Windows", reason="RTLD_DEFAULT is POSIX-only") @@ -52,3 +54,38 @@ def test_addon_exports_only_initializer() -> None: output = subprocess.run(command, check=True, capture_output=True, text=True).stdout exported = {line.split()[0].split("@@", 1)[0] for line in output.splitlines() if line.strip()} assert exported == expected + + +def test_tvm_ffi_loader_survives_process_shutdown() -> None: + """The TVM-FFI keep-alive registry must retain the addon through shutdown.""" + obj_dir = build_test_objects() + candidates = [ + "cc-gcc/test_funcs.o", + "cc/test_funcs.o", + "cc-appleclang/test_funcs.o", + "c-msvc/test_funcs.o", + "c-clang-cl/test_funcs.o", + "c/test_funcs.o", + "c-gcc/test_funcs.o", + ] + obj_path = next( + (obj_dir / candidate for candidate in candidates if (obj_dir / candidate).is_file()), None + ) + if obj_path is None: + pytest.skip("no test object is available") + + script = """ +import sys +from tvm_ffi_orcjit import ExecutionSession + +module = ExecutionSession().load_module(sys.argv[1], keep_module_alive=True) +assert module.test_add(2, 3) == 5 +del module +""" + result = subprocess.run( + [sys.executable, "-c", script, str(obj_path.resolve())], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py index 09bb2f6a6..f11b3284a 100644 --- a/python/tvm_ffi/cpp/extension.py +++ b/python/tvm_ffi/cpp/extension.py @@ -1116,7 +1116,7 @@ def load_inline( # noqa: PLR0913 keep_module_alive Whether to keep the module alive. If True, the module will be kept alive - for the duration of the program until libtvm_ffi.so is unloaded. + for the duration of the process. backend The GPU backend to use. It can be "cuda" or "hip". @@ -1441,7 +1441,7 @@ def load( # noqa: PLR0913 keep_module_alive Whether to keep the module alive. If True, the module will be kept alive - for the duration of the program until libtvm_ffi.so is unloaded. + for the duration of the process. backend The GPU backend to use. It can be "cuda" or "hip". diff --git a/python/tvm_ffi/module.py b/python/tvm_ffi/module.py index 8c69eaebe..ee795dd8a 100644 --- a/python/tvm_ffi/module.py +++ b/python/tvm_ffi/module.py @@ -445,7 +445,7 @@ def load_module(path: str | PathLike, keep_module_alive: bool = True) -> Module: keep_module_alive Whether to keep the module alive. If True, the module will be kept alive - for the duration of the program until libtvm_ffi.so is unloaded. + for the duration of the process. Returns ------- diff --git a/src/ffi/extra/module.cc b/src/ffi/extra/module.cc index 31bb95bb3..f8c2cd3c6 100644 --- a/src/ffi/extra/module.cc +++ b/src/ffi/extra/module.cc @@ -49,8 +49,12 @@ class ModuleGlobals { } static ModuleGlobals* Get() { - static ModuleGlobals instance; - return &instance; + // Process-lifetime by design. A pinned DSO may have registered functions + // whose callable contexts and deleters live in that DSO. Destroying this + // registry during static teardown can unload the DSO before the global + // function registry and other dependents release those callbacks. + static auto* instance = new ModuleGlobals(); + return instance; } private: From c9877919cc937446617479e11e237dac175b1aa3 Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Tue, 22 Sep 2026 00:01:52 +0800 Subject: [PATCH 2/4] fix(orcjit): harden lifecycle and evaluation --- addons/tvm_ffi_orcjit/CMakeLists.txt | 13 ++ addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md | 48 +++-- addons/tvm_ffi_orcjit/README.md | 15 +- .../python/tvm_ffi_orcjit/__init__.py | 9 +- .../python/tvm_ffi_orcjit/session.py | 23 +-- .../src/ffi/llvm_patches/README.md | 19 +- .../src/ffi/llvm_patches/init_fini_plugin.h | 10 +- .../ffi/llvm_patches/macho_cxa_atexit_shim.cc | 4 +- .../ffi/llvm_patches/macho_cxa_atexit_shim.h | 3 +- .../llvm_patches/win_dll_import_generator.cc | 27 +-- addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc | 165 +++++++++++++----- addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h | 34 +++- .../src/ffi/orcjit_memory_manager.cc | 9 +- .../src/ffi/orcjit_memory_manager.h | 15 +- .../tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 22 ++- .../tvm_ffi_orcjit/src/ffi/orcjit_session.h | 22 +-- addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc | 48 +++-- addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h | 26 ++- addons/tvm_ffi_orcjit/tests/CMakeLists.txt | 4 +- addons/tvm_ffi_orcjit/tests/README.md | 4 +- .../tests/sources/c/test_context.c | 24 +++ .../tests/sources/c/test_link_order_base.c | 2 +- .../tests/sources/c/test_link_order_caller.c | 6 +- .../tests/sources/cc/test_link_order_base.cc | 2 +- .../sources/cc/test_link_order_caller.cc | 2 +- addons/tvm_ffi_orcjit/tests/test_basic.py | 73 +++++++- .../tests/test_session_load_module.py | 8 +- 27 files changed, 457 insertions(+), 180 deletions(-) diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt b/addons/tvm_ffi_orcjit/CMakeLists.txt index af79db900..9ced5a168 100644 --- a/addons/tvm_ffi_orcjit/CMakeLists.txt +++ b/addons/tvm_ffi_orcjit/CMakeLists.txt @@ -22,6 +22,9 @@ endif () # ---- Find packages ---- find_package(LLVM REQUIRED CONFIG) +if (LLVM_PACKAGE_VERSION VERSION_LESS 22.0) + message(FATAL_ERROR "tvm_ffi_orcjit requires LLVM 22 or newer; found ${LLVM_PACKAGE_VERSION}") +endif () message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION} in ${LLVM_DIR}") find_package( @@ -76,12 +79,22 @@ execute_process( COMMAND ${LLVM_CONFIG_EXE} --link-static --libs Core OrcJIT Support native OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE _llvm_libs + ERROR_VARIABLE _llvm_libs_error + RESULT_VARIABLE _llvm_libs_rc ) +if (NOT _llvm_libs_rc EQUAL 0) + message(FATAL_ERROR "llvm-config --link-static --libs failed: ${_llvm_libs_error}") +endif () execute_process( COMMAND ${LLVM_CONFIG_EXE} --link-static --ldflags OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE _llvm_ldflags + ERROR_VARIABLE _llvm_ldflags_error + RESULT_VARIABLE _llvm_ldflags_rc ) +if (NOT _llvm_ldflags_rc EQUAL 0) + message(FATAL_ERROR "llvm-config --link-static --ldflags failed: ${_llvm_ldflags_error}") +endif () separate_arguments(_llvm_libs_list NATIVE_COMMAND "${_llvm_libs}") separate_arguments(_llvm_ldflags_list NATIVE_COMMAND "${_llvm_ldflags}") diff --git a/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md b/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md index 01472b983..9ab675ee2 100644 --- a/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md +++ b/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md @@ -155,8 +155,9 @@ symbol resolution must be deferred until all relevant objects are present. ## 3. LLVM ORC JIT v2 -**ORC** stands for *On Request Compilation*. LLVM ORC JIT v2 (introduced in LLVM 9, -stabilized in LLVM 13+) is a complete redesign of LLVM's JIT infrastructure. It is +**ORC** stands for *On Request Compilation*. The ORCv2 APIs have been available since +LLVM 7 and replaced ORCv1 completely in LLVM 12. They redesign LLVM's JIT infrastructure +around a composable, concurrent symbol/materialization model. ORCv2 is designed to be composable, asynchronous, and correct for production use (unlike the older `MCJIT` which had several fundamental limitations around multi-module linking). @@ -312,7 +313,7 @@ The three platform objects in LLVM are: | Platform | OS | Init section driven | | --- | --- | --- | | `MachOPlatform` | macOS / iOS | `__DATA,__mod_init_func` | -| `ELFNativePlatform` | Linux / ELF | `.init_array`, TLS | +| `ELFNixPlatform` | Linux / ELF | `.init_array`, TLS | | `COFFPlatform` | Windows | `.CRT$XC*` init, `__cxa_atexit` interop | `ExecutorNativePlatform` is a convenience builder that auto-selects the right platform @@ -322,17 +323,18 @@ for the host OS and loads the ORC runtime from a given path. The addon takes a different approach on each platform: -- **macOS**: ORC platform support is *optional*. When the caller passes an ORC runtime - path to `ExecutionSession`, `ExecutorNativePlatform` activates `MachOPlatform`. - `jit_->initialize(dylib)` and `jit_->deinitialize(dylib)` then drive `__mod_init_func` - and `__cxa_atexit` teardown natively. Without the path, the addon falls back to its - own `InitFiniPlugin`. +- **macOS**: the addon deliberately does not configure an ORC platform. It uses + `InitFiniPlugin` for `__mod_init_func` / `__mod_term_func` and a per-dylib + `__cxa_atexit` shim. This avoids a compact-unwind address-delta failure in the + current `MachOPlatform` integration. The `orc_rt` argument is ignored. - **Windows**: `COFFPlatform` is skipped entirely because it requires MSVC CRT symbols (`_CxxThrowException`, RTTI vtables, iostream objects) that are not resolvable in the JIT context. Instead, `InitFiniPlugin` manually handles `.CRT$XC*` / `.CRT$XT*` - init/fini sections. -- **Linux**: `ELFNativePlatform` is not used. `InitFiniPlugin` handles `.init_array` / - `.fini_array` / `.ctors` / `.dtors` directly, without the ORC runtime. + init/fini sections. The `orc_rt` argument is ignored. +- **Linux**: the default embedded `liborc_rt` configures `ExecutorNativePlatform` + (and therefore `ELFNixPlatform`). A custom archive may be supplied by path or bytes, + and `None` disables the platform. `InitFiniPlugin` still collects and invokes ELF + init/fini arrays to work around current upstream behavior. --- @@ -360,21 +362,28 @@ unit: load all its objects at once, then look up functions on the result. ### 4.2 Loading and lookup -`load_module` adds each object to the JITDylib, and JITLink parses it into a -`LinkGraph`, resolves relocations, and allocates executable JIT memory. -`get_function` looks the symbol up (materializing lazily), then wraps the raw +`load_module` registers each object with the JITDylib. On the first relevant lookup, +JITLink parses it into a `LinkGraph`, resolves relocations, and allocates executable +JIT memory. `get_function` looks the symbol up (materializing lazily), then wraps the raw pointer as a `tvm_ffi::Function`. Symbols resolve against the dylib's own -default link order (this dylib → Platform → process/runtime symbols); objects +link order (this dylib → LLJIT defaults, with a compiler-selected C++ runtime +inserted before process symbols on Linux); objects that reference each other must be loaded together, since there is no linking between separate `load_module` results. -Two addon-specific pieces sit in this pipeline: +Several addon-specific pieces sit in this pipeline: - **`InitFiniPlugin`** — a JITLink pass plugin that keeps init/fini sections (`.init_array`/`.ctors`/`.fini_array`/`.dtors`, `__mod_init_func`, `.CRT$XC*`) live, then collects their function pointers after fixup. The addon runs them in priority order at first lookup and at teardown, replacing the ORC platform's initializer machinery. See `llvm_patches/init_fini_plugin.h`. +- **Linux slab-pool memory manager** — reserves contiguous virtual-address regions, + separates executable from non-executable allocations, recycles freed regions, and + grows by adding slabs. This keeps 32-bit PC-relative relocations in range and offers + explicit reclamation through `clear_free_slabs()`. +- **GOTPCRELX correction (Linux/x86-64)** — repairs or reverses unsafe JITLink + relaxations before fixup. - **Windows DLL import stubs** — `DLLImportDefinitionGenerator` resolves `__imp_XXX` references to host-DLL functions by emitting JIT-memory pointer + trampoline stubs, keeping `PCRel32` fixups within ±2 GB of the JIT code. See @@ -391,12 +400,12 @@ import tvm_ffi_orcjit as oj sess = oj.default_session() # 2. Load a compiled object file into a fresh JITDylib -# → object parsed, JITLink links it, InitFiniPlugin collects ctors, -# context symbols injected eagerly, embedded binary (if any) expanded +# → objects registered; context and embedded-binary probes may materialize +# objects immediately, and context slots are populated before ctors run mod = sess.load_module("add.o") # returns a tvm_ffi.Module # 3. Look up and call a function -# → LLVM resolves "__tvm_ffi_add"; pending constructors fire on first lookup +# → LLVM resolves "__tvm_ffi_add"; any pending constructors run before return result = mod.add(3, 4) # → 7 ``` @@ -426,6 +435,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); | `JITDylib` | Symbol namespace / virtual shared library | `ORCJITDynamicLibraryObj::dylib_` | | `JITLink` | LLVM's JIT-aware linker | Used inside `ObjectLinkingLayer` | | JITLink pass pipeline | Pre-prune → post-alloc → post-fixup hooks | Where `InitFiniPlugin` runs | +| Slab pool | Contiguous, growable Linux JIT memory arena | `SlabPoolMemoryManager` / `Slab` | | `DefinitionGenerator` | Fallback symbol provider | `DLLImportDefinitionGenerator` (Win) | | Link order | Search path across JITDylibs for symbol resolution | LLJIT default (Main → Platform → ProcessSymbols) | | `__tvm_ffi_` prefix | Namespace for TVM-FFI exported functions | Used in `GetFunction()` | diff --git a/addons/tvm_ffi_orcjit/README.md b/addons/tvm_ffi_orcjit/README.md index a8be74695..88ea74ffb 100644 --- a/addons/tvm_ffi_orcjit/README.md +++ b/addons/tvm_ffi_orcjit/README.md @@ -26,8 +26,9 @@ TVM-FFI exported functions. - **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2 - **High-Level Loading**: `default_session().load_module(...)` mirrors `tvm_ffi.load_module`, returning a plain `tvm_ffi.Module` - **Unified Input**: Load from a file path, in-memory object bytes, or a list mixing both -- **Shared Session**: A process-wide session so multiple callers share one JIT environment (process symbols, arena, linking) +- **Shared Session**: A process-wide session so callers share process-symbol resolution and the Linux slab pool while loaded modules remain isolated - **Symbol Isolation**: Separate `load_module` calls define independent symbol namespaces, so they can define the same symbol without conflicts +- **Bounded-Range Memory**: A growable Linux slab pool keeps JIT code/data close enough for 32-bit PC-relative relocations and supports explicit drained-slab reclamation - **Init/Fini Support**: Handles static constructors/destructors across ELF (`.init_array`/`.ctors`), Mach-O (`__mod_init_func`), and COFF (`.CRT$XC*`/`.CRT$XT*`) - **Cross-Platform**: Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64) - **Multi-Compiler**: Tested with LLVM Clang, GCC, Apple Clang, MSVC, and clang-cl @@ -206,7 +207,8 @@ Compile: `clang -O2 -c -o example.o example.c` relocations from COFF objects before JITLink graph building, working around a JITLink limitation with COMDAT section symbols. -Please refers to [ORCJIT_PRIMER.md](./ORCJIT_PRIMER.md) to learn more about object file, linking, llvm orcjit v2, and how the addon works. +Refer to [ORCJIT_PRIMER.md](./ORCJIT_PRIMER.md) for background on object files, +linking, LLVM ORC JIT v2, and the addon's architecture. ## Project Structure @@ -219,6 +221,9 @@ tvm_ffi_orcjit/ │ ├── orcjit_session.h │ ├── orcjit_dylib.cc # JIT dylib module (object loading, symbol lookup) │ ├── orcjit_dylib.h +│ ├── orcjit_memory_manager.* # Growable Linux slab pool +│ ├── orcjit_slab.* # Contiguous-VA allocator and page protection +│ ├── llvm_patches/ # Isolated upstream LLVM workarounds │ └── orcjit_utils.h # LLVM error handling utilities ├── python/tvm_ffi_orcjit/ │ ├── __init__.py # Module exports and library loading @@ -261,6 +266,12 @@ The package requires LLVM 22+. Set `LLVM_PREFIX` to the LLVM install prefix: export LLVM_PREFIX=/path/to/llvm ``` +### Reclaiming unused JIT memory on Linux + +Dropping a module returns its regions to the session's slab pool for reuse. To +return fully drained slabs to the operating system, call +`session.clear_free_slabs()`. The call is synchronized with concurrent JIT work. + ## License Apache License 2.0 diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py index b56570a62..0b86584b0 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py @@ -51,14 +51,13 @@ ] _lib_path = None for path in _LIB_PATH: - if path.exists(): + if path.is_file(): _lib_path = path break if _lib_path is None: raise RuntimeError( - f"Could not find {_LIB_NAME}. " - f"Searched in {_LIB_PATH} and site-packages. " - f"Please ensure the package is installed correctly." + f"Could not find {_LIB_NAME}. Searched {_LIB_PATH}. " + "Please ensure the package is installed correctly." ) # The TVM-FFI loader uses local symbol scope. Retain the returned module here @@ -66,7 +65,7 @@ # JIT-owned object deleters remain valid through interpreter shutdown. _lib_module = load_module(_lib_path, keep_module_alive=True) -from .session import ExecutionSession, default_session +from .session import ExecutionSession, default_session # noqa: E402 __all__ = ["ExecutionSession", "default_session"] diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py index 2bb1098fd..bc1c1c3a0 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py @@ -119,10 +119,11 @@ def __init__( slab_size : int Per-slab capacity in bytes for the JIT memory manager. Linux only — ignored on macOS and Windows, where the slab allocator is compiled - out. 0 = arch default (64 MB; initial slab halves on mmap failure - down to 8 MB under RLIMIT_AS / container limits), >0 = custom size, - <0 = disable slab allocator (LLJIT uses its default scattered-mmap - allocator). + out. 0 = 64 MB default (the initial slab halves on mmap failure + down to 8 MB under RLIMIT_AS / container limits), >=4 MB = custom + size, <0 = disable the slab allocator (LLJIT uses its default + scattered-mmap allocator). Positive values below 4 MB are rejected + because both allocation pools need a 2 MB commit chunk. The session holds a growable pool of slabs: a fresh slab is mmap'd on demand when no existing one can fit a graph. Graphs that don't @@ -243,10 +244,9 @@ def clear_free_slabs(self) -> int: Fresh slabs that have never been allocated on are preserved, so the session remains ready to accept new work. - Safety: call when no JIT work is in flight on another thread. From - single-threaded Python this is always safe; once ``del lib`` has - returned, the C++ destructor has finished and the slab's live count - reflects the drop. + The operation is serialized with JIT allocation and module teardown, + so it is safe to call while other host threads use the same session. + Only fully drained slabs are reclaimed. Returns ------- @@ -267,9 +267,10 @@ def default_session() -> ExecutionSession: """Return the process-wide shared execution session. A single leaked, never-destroyed session shared by all callers in the - process, so they share one LLVM ``ExecutionSession`` — hence process - symbols, the slab arena, and cross-library linking. Created on first call - and cached for the lifetime of the process. + process, so they share one LLVM ``ExecutionSession`` — hence process-symbol + resolution, the slab pool, and synchronization infrastructure. Separate + loaded modules remain isolated symbol namespaces. Created on first call and + cached for the lifetime of the process. The session uses the ORC runtime embedded in the extension (no on-disk path lookup). For an isolated session or a tuned arena, construct an diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md index c2706a71a..2b43c8a4b 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md @@ -45,8 +45,9 @@ Each patch file opens with a fixed-shape header describing: - **ELF init/fini** (Linux branch of `init_fini_plugin.{h,cc}`) LLVM issue: [llvm/llvm-project#175981](https://github.com/llvm/llvm-project/issues/175981). - Upstream status: open, patch submitted. - Remove when: LLVM floor bumps past the release that contains the fix. + Upstream status: merged; included in LLVM 23.1.1, not LLVM 22.1.0. + Remove when: the LLVM floor reaches 23 and the addon switches its lifecycle + calls to the upstream `ELFNixPlatform` path. - **COFF ctor/dtor** (Windows branch of `init_fini_plugin.{h,cc}`) LLVM issue: COFFPlatform stalled. @@ -54,8 +55,18 @@ Each patch file opens with a fixed-shape header describing: Remove when: COFFPlatform becomes usable end-to-end with clang-cl / MSVC objects. -macOS already has working `MachOPlatform`, so no patch file is needed -for that platform. +- **Mach-O `__cxa_atexit` scoping** (`macho_cxa_atexit_shim.{h,cc}`) + Upstream dependency: re-enabling `MachOPlatform` after the compact-unwind + per-graph DSO-base fix is available. + Remove when: the addon can use `MachOPlatform` end-to-end. + +- **COFF unwind-data stripping** (`win_coff_pdata_strip.{h,cc}`) + Upstream dependency: usable `COFFPlatform` support with SEH registration. + Remove when: `.pdata` / `.xdata` can be registered and relocated normally. + +- **Windows DLL import stubs** (`win_dll_import_generator.{h,cc}`) + Upstream dependency: usable `COFFPlatform` and in-range DLL call stubs. + Remove when: host DLL imports work end-to-end without the custom generator. ## Removal checklist diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h index 21262e796..ee01a628e 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h @@ -52,11 +52,11 @@ * * ## Removal — Linux * - * LLVM issue: https://github.com/llvm/llvm-project/issues/175981 - * When the upstream fix lands and the project's minimum LLVM version - * bumps past the first release containing it, replace this plugin's - * Linux usage with `ELFNixPlatform` and delete the ELF handling path - * from this file. Concretely: + * LLVM issue: https://github.com/llvm/llvm-project/pull/175981 + * The upstream fix is included in LLVM 23.1.1 but not LLVM 22.1.0, the + * addon's current CI baseline. When the project's minimum LLVM version + * reaches 23 and its lifecycle calls use `ELFNixPlatform`, delete the ELF + * handling path from this file. Concretely: * - Remove the ELF-section branches (`.init_array`, `.ctors`, * `.fini_array`, `.dtors`) from `InitFiniPlugin::modifyPassConfig`. * - If no platform still needs this plugin, delete this file outright diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc index 0e1597ac2..f2cd48246 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc @@ -67,12 +67,12 @@ CxaAtexitRecordsScope::CxaAtexitRecordsScope(CxaAtexitRecords* records) } CxaAtexitRecordsScope::~CxaAtexitRecordsScope() { g_active_cxa_records = prev_; } -void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd) { +llvm::Error InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd) { llvm::orc::SymbolMap shim_syms; shim_syms[ES.intern("___cxa_atexit")] = { llvm::orc::ExecutorAddr::fromPtr(reinterpret_cast(&tvm_ffi_cxa_atexit_shim)), llvm::JITSymbolFlags::Exported | llvm::JITSymbolFlags::Callable}; - llvm::cantFail(jd.define(llvm::orc::absoluteSymbols(std::move(shim_syms)))); + return jd.define(llvm::orc::absoluteSymbols(std::move(shim_syms))); } void DrainCxaAtexit(CxaAtexitRecords& records) { diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h index 8aee61f06..f215e2097 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h @@ -61,6 +61,7 @@ #ifdef __APPLE__ #include +#include #include #include @@ -100,7 +101,7 @@ class CxaAtexitRecordsScope { * libSystem fallback — JITDylib::define-time symbols are searched before * the link order. */ -void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd); +llvm::Error InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd); /*! \brief Drain captured `(fn, arg)` records LIFO, running each dtor. * diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc index 59f86004c..901f627f5 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc @@ -52,6 +52,7 @@ #include // clang-format on +#include #include namespace tvm { @@ -61,22 +62,26 @@ namespace orcjit { void* DLLImportDefinitionGenerator::FindInProcessModules(const std::string& Name) { // Try specific runtime DLLs first, then tvm_ffi.dll (loaded by Python), // then all process modules, then LLVM's search. - static const char* kRuntimeDLLs[] = { + static constexpr std::array kRuntimeDLLs = { "vcruntime140.dll", "vcruntime140_1.dll", "ucrtbase.dll", "msvcp140.dll", }; - // NOTE: We intentionally do not call FreeLibrary() here. These runtime DLLs - // (vcruntime140, ucrtbase, etc.) are already loaded by the process and will - // remain loaded for its lifetime. LoadLibraryA merely increments the refcount; - // the extra refcount is harmless and avoids the overhead of balancing - // Get/FreeLibrary for every symbol lookup. - for (const char* dll : kRuntimeDLLs) { - if (HMODULE hMod = LoadLibraryA(dll)) { - if (auto addr = GetProcAddress(hMod, Name.c_str())) { - return reinterpret_cast(addr); - } + // Resolve the runtime handles once. GetModuleHandle avoids a refcount change + // for DLLs already present; LoadLibrary supplies a process-lifetime handle for + // an optional runtime that was not yet loaded. + static const std::array kRuntimeHandles = []() { + std::array handles{}; + for (std::size_t i = 0; i < kRuntimeDLLs.size(); ++i) { + handles[i] = GetModuleHandleA(kRuntimeDLLs[i]); + if (handles[i] == nullptr) handles[i] = LoadLibraryA(kRuntimeDLLs[i]); + } + return handles; + }(); + for (HMODULE module : kRuntimeHandles) { + if (module != nullptr) { + if (auto addr = GetProcAddress(module, Name.c_str())) return reinterpret_cast(addr); } } // Also check tvm_ffi.dll (host process symbol provider) diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc index a9f1dc1ee..b09252149 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc @@ -281,7 +281,76 @@ Module ORCJITDynamicLibraryObj::Finalize() { return GetRef(this); } -void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { +void ORCJITDynamicLibraryObj::WaitForInitializers() { + std::unique_lock lock(initializer_mutex_); + const std::thread::id current = std::this_thread::get_id(); + initializer_cv_.wait(lock, [this, current]() { + return initializer_depth_ == 0 || initializer_thread_ == current; + }); +} + +bool ORCJITDynamicLibraryObj::InitializerTurnAvailable() { + std::lock_guard lock(initializer_mutex_); + return initializer_depth_ == 0 || initializer_thread_ == std::this_thread::get_id(); +} + +void ORCJITDynamicLibraryObj::BeginInitializerRun() { + std::lock_guard lock(initializer_mutex_); + const std::thread::id current = std::this_thread::get_id(); + TVM_FFI_CHECK(initializer_depth_ == 0 || initializer_thread_ == current, InternalError) + << "Initializer ownership changed unexpectedly"; + if (initializer_depth_ == 0) initializer_thread_ = current; + ++initializer_depth_; +} + +void ORCJITDynamicLibraryObj::EndInitializerRun() { + bool notify = false; + { + std::lock_guard lock(initializer_mutex_); + TVM_FFI_CHECK(initializer_depth_ != 0 && initializer_thread_ == std::this_thread::get_id(), + InternalError) + << "Initializer completion does not match its owner"; + if (--initializer_depth_ == 0) { + initializer_thread_ = std::thread::id(); + notify = true; + } + } + if (notify) initializer_cv_.notify_all(); +} + +void ORCJITDynamicLibraryObj::RunInitializers( + const std::vector& entries) { + if (entries.empty()) return; + try { +#ifdef __APPLE__ + // Route any __cxa_atexit registrations made during init to this dylib's + // records; see llvm_patches/macho_cxa_atexit_shim.h. + CxaAtexitRecordsScope scope(&cxa_atexit_records_); +#endif + ORCJITExecutionSessionObj::RunInitFiniEntries(entries); + } catch (...) { + EndInitializerRun(); + throw; + } + EndInitializerRun(); +} + +void ORCJITDynamicLibraryObj::RunPendingInitializers() { + std::vector init; + while (true) { + WaitForInitializers(); + std::unique_lock session_lock(session_->mutex_); + // Another thread may have started initialization after the wait but before + // this session lock was acquired. Re-check under the session→gate lock order. + if (!InitializerTurnAvailable()) continue; + init = session_->DrainPendingInitializers(GetJITDylib()); + if (!init.empty()) BeginInitializerRun(); + break; + } + RunInitializers(init); +} + +void* ORCJITDynamicLibraryObj::GetSymbol(const String& name, bool run_initializers) { // Search this dylib only. Its JITDylib link order (set at creation) already // chains to Main → Platform → ProcessSymbols for host/runtime symbols, so a // single-entry search order resolves everything a self-contained module needs. @@ -293,42 +362,50 @@ void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { // Drain under the lock; run constructors after release (see below). llvm::Expected symbol_or_err = llvm::orc::ExecutorSymbolDef(); std::vector init; - { - std::lock_guard lock(session_->mutex_); + while (true) { + WaitForInitializers(); + std::unique_lock session_lock(session_->mutex_); + if (!InitializerTurnAvailable()) continue; symbol_or_err = jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str())); - init = session_->DrainPendingInitializers(GetJITDylib()); + if (symbol_or_err && run_initializers) { + init = session_->DrainPendingInitializers(GetJITDylib()); + if (!init.empty()) BeginInitializerRun(); + } else if (!symbol_or_err && run_initializers) { + // Never execute entries collected by a failed materialization: their + // target memory may already have been abandoned. + session_->DrainPendingInitializers(GetJITDylib()); + } + break; } - // Run this dylib's constructors (drained above) with the lock released. -#ifdef __APPLE__ - // Route any __cxa_atexit registrations made during init to this dylib's - // records; see llvm_patches/macho_cxa_atexit_shim.h. - CxaAtexitRecordsScope scope(&cxa_atexit_records_); -#endif - ORCJITExecutionSessionObj::RunInitFiniEntries(init); - if (!symbol_or_err) { llvm::Error remaining = llvm::handleErrors(symbol_or_err.takeError(), [](const llvm::orc::SymbolsNotFound&) {}); if (remaining) TVM_FFI_ORCJIT_LLVM_CALL(std::move(remaining)); return nullptr; } + // Run this dylib's constructors with the session lock released. The + // per-dylib initializer gate blocks other threads until completion while + // allowing same-thread constructor callbacks to re-enter the dylib. + RunInitializers(init); return symbol_or_err->getAddress().toPtr(); } void ORCJITDynamicLibraryObj::InitContextSymbols() { // Called once from Finalize before the dylib is published, so no guard is - // needed. Point the library-context slot at this module and inject any - // registered context symbols. - if (void** ctx_addr = reinterpret_cast(GetSymbol(symbol::tvm_ffi_library_ctx))) { + // needed. Resolve every context slot without running initializers, populate + // the slots, then run the collected initializers. Constructors can therefore + // safely use the library context on their first instruction. + if (void** ctx_addr = reinterpret_cast(GetSymbol(symbol::tvm_ffi_library_ctx, false))) { *ctx_addr = this; } Module::VisitContextSymbols([this](const String& name, void* symbol) { - if (void** ctx_addr = reinterpret_cast(GetSymbol(name))) { + if (void** ctx_addr = reinterpret_cast(GetSymbol(name, false))) { *ctx_addr = symbol; } }); + RunPendingInitializers(); } llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { @@ -337,9 +414,9 @@ llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { } Optional ORCJITDynamicLibraryObj::GetFunction(const String& name) { - // Pure symbol lookup. Context symbols were injected once at load time (see - // Finalize), so this holds no lock and does no refresh — the returned - // Function, once resolved, is invoked lock-free on the hot path. + // Context symbols were injected once at load time (see Finalize). Resolution + // is serialized by GetSymbol, but the returned Function is invoked lock-free + // on the hot path. // // TVM-FFI exports have the __tvm_ffi_ prefix. std::string symbol_name = symbol::tvm_ffi_symbol_prefix + std::string(name); @@ -356,31 +433,31 @@ Optional ORCJITDynamicLibraryObj::GetFunction(const String& name) { //------------------------------------- static void RegisterOrcJITFunctions() { - static bool registered = false; - if (registered) return; - registered = true; - - namespace refl = tvm::ffi::reflection; - - refl::ObjectDef(); - - refl::GlobalDef() - .def("tvm_ffi_orcjit.ExecutionSession", - [](const Optional>& orc_rt, int64_t slab_size_bytes) { - return ORCJITExecutionSession(orc_rt, slab_size_bytes); - }) - .def("tvm_ffi_orcjit.GlobalDefaultSession", - []() { return ORCJITExecutionSessionObj::GlobalDefault(); }) - .def("tvm_ffi_orcjit.SessionLoadModule", - [](const ORCJITExecutionSession& session, const Array>& objects, - const String& name, const Optional& cxx_runtime_path, - const Optional& libstdcxx_nonshared_path) -> Module { - return session->LoadModule(objects, name, cxx_runtime_path, libstdcxx_nonshared_path); - }) - .def("tvm_ffi_orcjit.SessionClearFreeSlabs", - [](const ORCJITExecutionSession& session) -> int64_t { - return session->ClearFreeSlabs(); - }); + static std::once_flag once; + std::call_once(once, []() { + namespace refl = tvm::ffi::reflection; + + refl::ObjectDef(); + + refl::GlobalDef() + .def("tvm_ffi_orcjit.ExecutionSession", + [](const Optional>& orc_rt, int64_t slab_size_bytes) { + return ORCJITExecutionSession(orc_rt, slab_size_bytes); + }) + .def("tvm_ffi_orcjit.GlobalDefaultSession", + []() { return ORCJITExecutionSessionObj::GlobalDefault(); }) + .def("tvm_ffi_orcjit.SessionLoadModule", + [](const ORCJITExecutionSession& session, const Array>& objects, + const String& name, const Optional& cxx_runtime_path, + const Optional& libstdcxx_nonshared_path) -> Module { + return session->LoadModule(objects, name, cxx_runtime_path, + libstdcxx_nonshared_path); + }) + .def("tvm_ffi_orcjit.SessionClearFreeSlabs", + [](const ORCJITExecutionSession& session) -> int64_t { + return session->ClearFreeSlabs(); + }); + }); } TVM_FFI_STATIC_INIT_BLOCK() { RegisterOrcJITFunctions(); } diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h index b458b1271..370dffdd5 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h @@ -31,7 +31,12 @@ #include #include +#include +#include #include +#include +#include +#include #include "llvm_patches/macho_cxa_atexit_shim.h" #include "orcjit_session.h" @@ -105,7 +110,25 @@ class ORCJITDynamicLibraryObj : public ModuleObj { * \param name The symbol name to look up * \return Pointer to the symbol, or nullptr if not found */ - void* GetSymbol(const String& name); + void* GetSymbol(const String& name, bool run_initializers = true); + + /*! \brief Wait until another thread has finished running this dylib's initializers. */ + void WaitForInitializers(); + + /*! \brief Check the initializer gate while holding the session mutex. */ + bool InitializerTurnAvailable(); + + /*! \brief Mark a (possibly nested) initializer run owned by the current thread. */ + void BeginInitializerRun(); + + /*! \brief Finish an initializer run and wake blocked lookup threads when outermost. */ + void EndInitializerRun(); + + /*! \brief Run entries outside the session mutex while preserving lookup ordering. */ + void RunInitializers(const std::vector& entries); + + /*! \brief Drain and run every initializer collected by raw context lookups. */ + void RunPendingInitializers(); /*! * \brief Get the underlying LLVM JITDylib @@ -128,6 +151,15 @@ class ORCJITDynamicLibraryObj : public ModuleObj { /*! \brief Whether Finalize has run; guards against double-finalizing. */ bool finalized_{false}; + // Constructors run without the session mutex to permit re-entry. These + // fields prevent a second thread from observing callable code before the + // first thread completes initialization; same-thread nested lookup remains + // allowed for constructor callbacks. + std::mutex initializer_mutex_; + std::condition_variable initializer_cv_; + std::thread::id initializer_thread_; + std::size_t initializer_depth_{0}; + #ifdef __APPLE__ /*! \brief Per-dylib __cxa_atexit registry. * diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc index a970c9760..279d54788 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc @@ -27,6 +27,7 @@ #include #include +#include #include #include @@ -60,8 +61,8 @@ SlabPoolMemoryManager::SlabPoolMemoryManager(std::size_t page_size, std::size_t } cap /= 2; } - llvm::report_fatal_error("SlabPoolMemoryManager: failed to reserve at least " + - llvm::Twine(floor / (1024 * 1024)) + " MB of virtual address space"); + TVM_FFI_THROW(RuntimeError) << "SlabPoolMemoryManager: failed to reserve at least " + << floor / (1024 * 1024) << " MB of virtual address space"; } std::unique_ptr SlabPoolMemoryManager::createSlab(std::size_t capacity) { @@ -79,8 +80,8 @@ void SlabPoolMemoryManager::allocate(const llvm::jitlink::JITLinkDylib* /*JD*/, // user callback, since the LLJIT linker issues nested lookups (and // thus re-entrant allocate() calls via materialization) from inside // OnAllocated and a coarse lock would deadlock. Snapshot raw pointers - // under the lock; slabs are guaranteed to outlive this call because - // clearFreeSlabs() is only safe when the session is quiescent. + // under the lock. Slabs outlive this call because the enclosing execution + // session serializes allocation, teardown, and clearFreeSlabs(). // // Slab::allocate is synchronous (invokes its callback inline on every // code path), so a captured std::optional observes the result before diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h index a71487dc6..6ec89707b 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h @@ -23,7 +23,7 @@ * * `SlabPoolMemoryManager` implements `JITLinkMemoryManager` on top of a * per-session `std::vector>`. On each `allocate` - * it picks the first `Slab` that can fit the graph; if none do, it + * it picks the first `Slab` that can fit the graph; if none does, it * `mmap`s a fresh slab sized to fit (`Slab::capacityForFootprint`) and * appends it. Normal-size graphs land on a `slab_size`-sized slab; * skewed or oversize graphs land on a power-of-2 larger slab whose @@ -63,8 +63,8 @@ namespace orcjit { /*! * \brief `JITLinkMemoryManager` backed by a growable pool of `Slab`s. * - * The constructor reserves one initial slab (halving its capacity down - * to `kMinSlabSize` if `mmap` fails under RLIMIT_AS). Subsequent + * The constructor reserves one initial slab, halving on `mmap` failure down to + * `min(slab_size, kMinSlabSize)`. Subsequent * slabs are added on demand by `allocate()` at a capacity chosen by * `Slab::capacityForFootprint` — `slab_size_` for normal graphs, the * next power of two up for skewed / oversize graphs. No retry, no @@ -78,8 +78,9 @@ class SlabPoolMemoryManager : public llvm::jitlink::JITLinkMemoryManager { // granule. Small enough that a pinned slab only wastes 64 MB of RSS. static constexpr std::size_t kDefaultSlabSize = std::size_t{64} << 20; // 64 MB - // Lower bound on initial-slab reservation. If the first `mmap` - // fails and halving drops below this, the constructor aborts. + // Lower bound for initial-slab fallback when the requested capacity is at + // least this large. Smaller valid custom capacities use their own size. If + // every reservation fails, construction raises a RuntimeError. // 8 MB is enough for a minimal JITDylib setup under very tight // RLIMIT_AS. static constexpr std::size_t kMinSlabSize = std::size_t{8} << 20; // 8 MB @@ -108,8 +109,8 @@ class SlabPoolMemoryManager : public llvm::jitlink::JITLinkMemoryManager { * prior allocation — back to the OS via `munmap`. * * Returns the number of slabs reclaimed. Safe to call any time the - * session is quiescent (no concurrent JIT work in flight). A typical - * pattern is to call this after dropping a batch of libraries: + * The enclosing execution session serializes this operation with JIT work. + * A typical pattern is to call it after dropping a batch of libraries: * * for lib in libs: del lib * session.clear_free_slabs() # Python API diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc index cbe17859c..bb56eee4f 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc @@ -24,6 +24,7 @@ #include "orcjit_session.h" +#include #include #include #include @@ -40,6 +41,7 @@ #include #include #include +#include #if defined(__linux__) && defined(__GLIBCXX__) #include @@ -198,7 +200,7 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const Optional0 = custom size, <0 = disable arena (LLJIT uses its // default allocator — scattered mmap, no PC-rel guarantee). // The parameter is Linux-only; on macOS/Windows the arena is compiled out @@ -226,6 +228,10 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const Optional 0) { + constexpr std::size_t kMinCustomSlabSize = 2 * Slab::kCommitGranularity; + TVM_FFI_CHECK(static_cast(slab_size_bytes) >= kMinCustomSlabSize, ValueError) + << "slab_size must be 0, negative (disabled), or at least " << kMinCustomSlabSize + << " bytes, but got " << slab_size_bytes; slab_size = static_cast(slab_size_bytes); } else { slab_size = SlabPoolMemoryManager::kDefaultSlabSize; @@ -402,6 +408,9 @@ ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary( llvm::orc::JITDylib& jit_dylib = TVM_FFI_ORCJIT_LLVM_CALL(jit_->getExecutionSession().createJITDylib(lib_name.c_str())); + // If any subsequent link-order/generator setup fails, remove the partially + // configured dylib while the session mutex is still held. + llvm::scope_exit cleanup_dylib([this, &jit_dylib]() { RemoveDylib(&jit_dylib); }); #if defined(__linux__) && defined(__GLIBCXX__) llvm::orc::JITDylib* cxx_runtime_dylib = nullptr; if (cxx_runtime) { @@ -452,16 +461,16 @@ ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary( } #endif - auto dylib_obj = make_object(GetRef(this), - &jit_dylib, jit_.get(), lib_name); - #ifdef __APPLE__ // Inject ___cxa_atexit on the user JITDylib so it wins over 's // fallback (which resolves to libSystem's and would orphan dtors from // our drop-time drain). See llvm_patches/macho_cxa_atexit_shim.h. - InstallCxaAtexitShim(jit_->getExecutionSession(), jit_dylib); + TVM_FFI_ORCJIT_LLVM_CALL(InstallCxaAtexitShim(jit_->getExecutionSession(), jit_dylib)); #endif + auto dylib_obj = make_object(GetRef(this), + &jit_dylib, jit_.get(), lib_name); + cleanup_dylib.release(); return ORCJITDynamicLibrary(std::move(dylib_obj)); } @@ -528,6 +537,9 @@ void ORCJITExecutionSessionObj::AddPendingDeinitializer(llvm::orc::JITDylib* jit int64_t ORCJITExecutionSessionObj::ClearFreeSlabs() { #ifdef __linux__ if (memory_manager_) { + // Synchronize with lookup/materialization and dylib teardown. This makes + // the public API safe even when another host thread is using the session. + std::lock_guard lock(mutex_); return static_cast(memory_manager_->clearFreeSlabs()); } #endif diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h index 7c2ea4efb..b08b7c614 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h @@ -34,7 +34,6 @@ #include #include -#include #include #include #include @@ -79,9 +78,10 @@ class ORCJITExecutionSessionObj : public Object { * \brief Get the process-wide shared execution session. * * A leaked, never-destroyed singleton so multiple callers in one process - * share one LLVM ExecutionSession — hence process symbols, the slab arena, - * and cross-library linking. Never torn down (interpreter finalization could - * otherwise call back into the host language during teardown). + * share one LLVM ExecutionSession — hence process-symbol resolution, the slab + * pool, and synchronization infrastructure. Loaded modules remain isolated + * symbol namespaces. Never torn down (interpreter finalization could otherwise + * call back into the host language during teardown). * * Always uses the ORC runtime embedded in this extension; deliberately not * user-configurable (a shared process-wide singleton with a hidden runtime @@ -111,9 +111,9 @@ class ORCJITExecutionSessionObj : public Object { * \c String path or in-memory \c Bytes image), injects context symbols * eagerly, and — if the objects embed a library binary — reconstructs the * import tree so the result behaves like a normally-loaded tvm-ffi module. - * The whole sequence runs under one recursive session lock, and the dylib is - * never exposed in a partially-loaded state, so finalization happens exactly - * once and cannot be repeated or interleaved. + * The dylib is not exposed until finalization completes. Leaf operations are + * serialized by the session mutex, while JIT constructors run with that mutex + * released so they can safely re-enter the session. * * \param objects Array whose elements are each a \c String path or \c Bytes * object-file image. @@ -195,9 +195,9 @@ class ORCJITExecutionSessionObj : public Object { /*! * \brief Release drained slabs (no live JIT allocations) back to the OS. * - * Returns the number of slabs reclaimed. No-op on macOS/Windows - * where the slab pool is compiled out, or when the pool has been - * disabled via `slab_size < 0`. + * Returns the number of slabs reclaimed. No-op on macOS/Windows, where the + * slab pool is compiled out, or when the pool has been disabled via + * `slab_size < 0`. Serialized with allocation and teardown by \ref mutex_. */ int64_t ClearFreeSlabs(); @@ -213,7 +213,7 @@ class ORCJITExecutionSessionObj : public Object { std::unique_ptr jit_; /*! \brief Counter for auto-generating library names */ - std::atomic dylib_counter_{0}; + int dylib_counter_{0}; /*! \brief Compiler-selected C++ runtime search JITDylibs, keyed by shared-library path. */ std::unordered_map cxx_runtime_dylibs_; diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc index 5c2d629cb..b24e37b4f 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc @@ -37,10 +37,12 @@ #include #include +#include #include #include #include #include +#include namespace tvm { namespace ffi { @@ -167,20 +169,38 @@ Error Slab::commitPages(void* addr, std::size_t size) { std::size_t last_chunk = (offset + size - 1) / kCommitGranularity; for (std::size_t i = first_chunk; i <= last_chunk; ++i) { - if (committed_[i].load(std::memory_order_acquire) != 0) continue; - std::size_t chunk_offset = i * kCommitGranularity; - std::size_t chunk_len = std::min(kCommitGranularity, arena_capacity_ - chunk_offset); - // mprotect is idempotent, so a concurrent racer calling it on the same chunk - // is harmless. Only flip the flag after success — otherwise a failed commit - // followed by freeRegion() would leave committed_[i] == 1, causing a - // later allocation to skip mprotect and write into PROT_NONE memory. - if (::mprotect(arena_base_ + chunk_offset, chunk_len, PROT_READ | PROT_WRITE) != 0) { - return make_error("Slab: mprotect(RW) failed for chunk at offset " + - formatv("{0:x}", chunk_offset) + ": " + - std::strerror(errno), - inconvertibleErrorCode()); + constexpr std::uint8_t kUncommitted = 0; + constexpr std::uint8_t kCommitting = 1; + constexpr std::uint8_t kCommitted = 2; + auto& state = committed_[i]; + while (true) { + std::uint8_t observed = state.load(std::memory_order_acquire); + if (observed == kCommitted) break; + if (observed == kCommitting) { + std::this_thread::yield(); + continue; + } + if (!state.compare_exchange_weak(observed, kCommitting, std::memory_order_acq_rel, + std::memory_order_acquire)) { + continue; + } + + // Only the thread that claimed this chunk may change its protection. + // A duplicate mprotect(RW) racing after another allocation finalized to + // RX/R would silently remove that allocation's execute/read-only state. + std::size_t chunk_offset = i * kCommitGranularity; + std::size_t chunk_len = std::min(kCommitGranularity, arena_capacity_ - chunk_offset); + if (::mprotect(arena_base_ + chunk_offset, chunk_len, PROT_READ | PROT_WRITE) != 0) { + int commit_errno = errno; + state.store(kUncommitted, std::memory_order_release); + return make_error("Slab: mprotect(RW) failed for chunk at offset " + + formatv("{0:x}", chunk_offset) + ": " + + std::strerror(commit_errno), + inconvertibleErrorCode()); + } + state.store(kCommitted, std::memory_order_release); + break; } - committed_[i].store(1, std::memory_order_release); } return Error::success(); } @@ -411,7 +431,7 @@ Expected Slab::bumpAllocate(std::size_t size, bool is_exec) { } // Bump allocate within the pool's limit. - if (bump + size > limit) { + if (size > limit - bump) { return make_error(is_exec ? "exec" : "non-exec", bump, size, limit); } diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h index 0029dc513..e7f1d2158 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h @@ -25,11 +25,9 @@ * bump-allocates from it, keeping all JIT allocations within range of * PC-relative relocations (±2 GB on x86_64, ±4 GB on AArch64). * - * The `Slab` is the unit-of-VA-reservation for the OrcJIT memory manager. - * Today it is used as a single-slab arena owned by - * `ArenaJITLinkMemoryManager`. Stage B of the refactor will introduce a - * `SlabPoolMemoryManager` that holds multiple Slabs and grows by mmap-ing - * new ones on demand. + * The `Slab` is the unit of VA reservation for `SlabPoolMemoryManager`, + * which owns a growable set of slabs and adds one when no existing slab can + * satisfy a graph's per-pool footprint. * * ## Page commit + Transparent Huge Page (THP) support * @@ -37,8 +35,9 @@ * size matches the Linux huge-page granule on both x86_64 and AArch64, * enabling THP promotion via `madvise(MADV_HUGEPAGE)` on the full * reservation. Each 2 MB commit-chunk is `mprotect`-ed to RW exactly once - * via an atomic bitmap flag (`committed_`), avoiding lock contention with - * the per-pool allocator mutex. + * via an atomic state bitmap (`committed_`), avoiding lock contention with + * the per-pool allocator mutex while ensuring only one thread changes the + * protection of a commit chunk. * * ## Dual-pool exec / non-exec split * @@ -111,8 +110,8 @@ class SlabPoolExhaustedError : public llvm::ErrorInfo { * Zero-sized sub-regions indicate no allocation from that pool. * * \p owner points to the Slab that handed out this allocation. With one - * slab per session today this is redundant, but stamping it now makes - * Stage B's pool-manager routing O(1) without address comparison. + * Stamping the owning slab makes pool-manager deallocation routing O(1) + * without address comparison. */ struct FinalizedAllocInfo { Slab* owner; ///< Slab that owns these offsets. @@ -160,8 +159,8 @@ class Slab { /*! \brief Construct a Slab and reserve \p capacity bytes of VA. * * On reservation failure, returns with \c base() == nullptr — the - * caller is expected to retry at a smaller capacity or - * \c report_fatal_error. + * caller is expected to retry at a smaller capacity or propagate a + * recoverable allocation error. */ Slab(std::size_t page_size, std::size_t capacity); @@ -302,9 +301,8 @@ class Slab { std::vector free_list_non_exec_; std::vector free_list_exec_; - /*! \brief Per-commit-chunk flags (0 = uncommitted, 1 = committed). - * Lock-free: each chunk is mprotect'd exactly once via - * compare_exchange. */ + /*! \brief Per-commit-chunk states (0 = uncommitted, 1 = committing, + * 2 = committed). Each chunk is mprotect'd exactly once. */ std::unique_ptr[]> committed_; std::size_t num_commit_chunks_ = 0; diff --git a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt index 92bca5e5a..e5e0f338f 100644 --- a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt +++ b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt @@ -88,7 +88,7 @@ if (NOT WIN32) endif () # Pure C object files — built on all platforms (no C++ runtime deps) -enablelanguage(C) +enable_language(C) add_test_object(sources/c/test_funcs.c) add_test_object(sources/c/test_funcs2.c) add_test_object(sources/c/test_funcs_conflict.c) @@ -103,7 +103,7 @@ add_test_object(sources/c/test_ctor_dtor.c) # CUDA object files — optional find_package(CUDAToolkit) if (CUDAToolkit_FOUND) - enablelanguage(CUDA) + enable_language(CUDA) message(STATUS "CUDA found: ${CUDAToolkit_VERSION}") add_test_object(sources/cuda/test_funcs.cu) endif () diff --git a/addons/tvm_ffi_orcjit/tests/README.md b/addons/tvm_ffi_orcjit/tests/README.md index beb9f51b7..379946bca 100644 --- a/addons/tvm_ffi_orcjit/tests/README.md +++ b/addons/tvm_ffi_orcjit/tests/README.md @@ -87,9 +87,9 @@ variant subdirectories (c/, cc/, c-gcc/, etc.). | `test_funcs2` | More arithmetic (subtract, divide) | | `test_funcs_conflict` | Symbol conflict testing (duplicate `add`) | | `test_call_global` | Callbacks into Python-registered global functions | -| `test_context` | First-lookup library-context injection | +| `test_context` | Eager library-context injection, including constructor ordering | | `test_types` | Zero-arg, multi-arg, float, void return types | -| `test_link_order_base` / `test_link_order_caller` | Cross-library symbol resolution | +| `test_link_order_base` / `test_link_order_caller` | Cross-object symbol resolution within one module | | `test_error` | Error propagation from JIT'd code | | `test_ctor_dtor` | Constructor/destructor and init/fini sections | diff --git a/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c index d4a407604..12aaa01c9 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c +++ b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c @@ -20,6 +20,19 @@ #include TVM_FFI_DLL_EXPORT void* __tvm_ffi__library_ctx = NULL; +static int context_was_set_during_init = 0; + +static void record_context_during_init(void) { + context_was_set_during_init = __tvm_ffi__library_ctx != NULL; +} + +#ifdef _MSC_VER +typedef void(__cdecl* ctor_t)(void); +#pragma section(".CRT$XCU", read) +__declspec(allocate(".CRT$XCU")) ctor_t __tvm_test_context_init = record_context_during_init; +#else +__attribute__((constructor)) static void context_init(void) { record_context_during_init(); } +#endif TVM_FFI_DLL_EXPORT int __tvm_ffi_context_is_set(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { @@ -31,3 +44,14 @@ TVM_FFI_DLL_EXPORT int __tvm_ffi_context_is_set(void* self, const TVMFFIAny* arg result->v_int64 = __tvm_ffi__library_ctx != NULL; return 0; } + +TVM_FFI_DLL_EXPORT int __tvm_ffi_context_was_set_during_init(void* self, const TVMFFIAny* args, + int32_t num_args, TVMFFIAny* result) { + (void)self; + (void)args; + (void)num_args; + result->type_index = kTVMFFIInt; + result->zero_padding = 0; + result->v_int64 = context_was_set_during_init; + return 0; +} diff --git a/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_base.c b/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_base.c index 21f26df9f..377d3f8ca 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_base.c +++ b/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_base.c @@ -18,7 +18,7 @@ */ /* - * Base library for cross-library linking test. + * Base object for the intra-module linking test. * Exports helper_add which is called by test_link_order_caller.c. */ #include diff --git a/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_caller.c b/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_caller.c index 215a267be..9214ed75d 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_caller.c +++ b/addons/tvm_ffi_orcjit/tests/sources/c/test_link_order_caller.c @@ -18,17 +18,17 @@ */ /* - * Caller library for cross-library linking test. + * Caller object for the intra-module linking test. * References __tvm_ffi_helper_add from test_link_order_base.c via extern * declaration, and exports cross_lib_add which forwards to it. */ #include -/* Declare external symbol from the base library */ +/* Declare the external symbol from the base object. */ extern int __tvm_ffi_helper_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result); -/* cross_lib_add: forwards to helper_add in the base library */ +/* cross_lib_add: forwards to helper_add in the base object. */ TVM_FFI_DLL_EXPORT int __tvm_ffi_cross_lib_add(void* self, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* result) { return __tvm_ffi_helper_add(self, args, num_args, result); diff --git a/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_base.cc b/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_base.cc index c878c6ea0..a5df5815c 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_base.cc +++ b/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_base.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -// Base library for cross-library linking test (C++ version). +// Base object for the intra-module linking test (C++ version). // Exports helper_add which is called by test_link_order_caller.cc. #include diff --git a/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_caller.cc b/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_caller.cc index 85dde579b..63a5b8244 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_caller.cc +++ b/addons/tvm_ffi_orcjit/tests/sources/cc/test_link_order_caller.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -// Caller library for cross-library linking test (C++ version). +// Caller object for the intra-module linking test (C++ version). // References __tvm_ffi_helper_add from test_link_order_base.cc and // exports cross_lib_add which forwards to it. diff --git a/addons/tvm_ffi_orcjit/tests/test_basic.py b/addons/tvm_ffi_orcjit/tests/test_basic.py index c7b9e822e..570e41511 100644 --- a/addons/tvm_ffi_orcjit/tests/test_basic.py +++ b/addons/tvm_ffi_orcjit/tests/test_basic.py @@ -106,6 +106,14 @@ def call_global_obj(self) -> str: """Return path prefix for test_call_global object.""" return f"{self.subdir}/test_call_global" + def link_order_base_obj(self) -> str: + """Return path prefix for the intra-module dependency object.""" + return f"{self.subdir}/test_link_order_base" + + def link_order_caller_obj(self) -> str: + """Return path prefix for the object that calls the dependency.""" + return f"{self.subdir}/test_link_order_caller" + def types_obj(self) -> str: """Return path prefix for test_types object.""" return f"{self.subdir}/test_types" @@ -191,6 +199,17 @@ def test_multiple_objects_in_one_module(v: Variant) -> None: assert mod.get_function(v.fn("test_divide"))(20, 4) == 5 +@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) +@pytest.mark.parametrize("reverse", [False, True], ids=["caller-first", "base-first"]) +def test_intra_module_dependency_is_input_order_independent(v: Variant, reverse: bool) -> None: + """Undefined symbols resolve between objects regardless of input order.""" + objects = [v.link_order_caller_obj(), v.link_order_base_obj()] + if reverse: + objects.reverse() + mod = load(*objects) + assert mod.get_function(v.fn("cross_lib_add"))(17, 25) == 42 + + @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_function_not_found(v: Variant) -> None: """Raise AttributeError for a missing function name.""" @@ -285,8 +304,9 @@ def test_call_global(v: Variant) -> None: def test_context_injected_at_load() -> None: - """Context is injected eagerly at load, so concurrent lookups all observe it.""" + """Context is injected before constructors and concurrent lookups observe it.""" mod = load("c/test_context") + assert mod.context_was_set_during_init() == 1 num_workers = 8 barrier = threading.Barrier(num_workers) @@ -461,6 +481,42 @@ def _append_ctor_log(x: str) -> None: assert "" not in log +def test_concurrent_first_lookup_waits_for_initializers() -> None: + """A second thread cannot call newly materialized code before init completes.""" + ctor_entered = threading.Event() + release_ctor = threading.Event() + second_started = threading.Event() + second_returned = threading.Event() + + @tvm_ffi.register_global_func("append_log", override=True) + def _block_first_ctor(_value: str) -> None: + if not ctor_entered.is_set(): + ctor_entered.set() + release_ctor.wait(timeout=5) + + mod = load("c/test_ctor_dtor") + + def first_call() -> None: + mod.main() + + def second_call() -> None: + second_started.set() + mod.main() + second_returned.set() + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(first_call) + assert ctor_entered.wait(timeout=5), "constructor did not start" + second = pool.submit(second_call) + assert second_started.wait(timeout=5), "second lookup did not start" + assert not second_returned.wait(timeout=0.25), ( + "second call returned before the first thread completed initialization" + ) + release_ctor.set() + first.result(timeout=5) + second.result(timeout=5) + + # --------------------------------------------------------------------------- # Module drop — dropping a loaded Module while its session is still alive. # @@ -630,7 +686,7 @@ def test_containers_tuple(v: Variant) -> None: # --------------------------------------------------------------------------- -# Slab-pool growth (Stage B). +# Slab-pool growth. # # A session holds a growable pool of Slabs, each `slab_size` bytes. When a # JITLink graph won't fit in any existing slab, the pool mmap's a new one; @@ -640,10 +696,8 @@ def test_containers_tuple(v: Variant) -> None: # --------------------------------------------------------------------------- -# 8 MB is the practical floor — Slab::kCommitGranularity is 2 MB and the -# dual-pool midpoint needs at least two commit chunks of headroom above it, -# so smaller capacities break the pool layout. See SlabPoolMemoryManager -# kMinSlabSize. +# Use 8 MB to force frequent growth while leaving more headroom than the 4 MB +# structural minimum (one 2 MB commit chunk per allocation pool). _SMALL_SLAB = 8 * 1024 * 1024 @@ -850,3 +904,10 @@ def test_clear_free_slabs_disabled_pool() -> None: """When the slab pool is disabled, clear_free_slabs is a no-op (returns 0).""" session = ExecutionSession(slab_size=-1) assert session.clear_free_slabs() == 0 + + +@pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") +def test_rejects_too_small_custom_slab() -> None: + """A slab needs at least one 2 MB commit chunk per allocation pool.""" + with pytest.raises(ValueError, match="slab_size must be"): + ExecutionSession(slab_size=2 * 1024 * 1024) diff --git a/addons/tvm_ffi_orcjit/tests/test_session_load_module.py b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py index 1b11b753b..562ede31c 100644 --- a/addons/tvm_ffi_orcjit/tests/test_session_load_module.py +++ b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py @@ -311,8 +311,8 @@ def test_load_module_expands_embedded_library_bin(tmp_path: Path) -> None: """An embedded library binary is deserialized and its imports are wired.""" # A custom module kind whose loader returns a real orcjit module. This also - # forces the loader to re-enter load_module while the outer call holds the - # session lock — exercising the recursive session lock. + # forces load_module to re-enter the shared session during outer-module + # finalization, exercising nested loading without deadlock. @tvm_ffi.register_global_func("ffi.Module.load_from_bytes.orcjit_test_probe", override=True) def _load_probe(_data: bytes) -> tvm_ffi.Module: return default_session().load_module(obj("c/test_funcs2")) @@ -342,8 +342,8 @@ def _load_probe(_data: bytes) -> tvm_ffi.Module: # --------------------------------------------------------------------------- # Concurrency — the shared session driven by many threads at once. # -# Exercises the recursive session lock: overlapping create / add / lookup / -# drop on one shared ExecutionSession must not corrupt linker state. +# Exercises the session lock and per-dylib initializer gate: overlapping +# create / add / lookup / drop must not corrupt linker state. # --------------------------------------------------------------------------- From b37ff54658ce29ce38bd88dbae1285184be50c79 Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 20:12:45 +0000 Subject: [PATCH 3/4] [FIX] Remove unused ORC JIT lint suppression (#809) --- addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py index 0b86584b0..bc3cdee7b 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py @@ -65,7 +65,7 @@ # JIT-owned object deleters remain valid through interpreter shutdown. _lib_module = load_module(_lib_path, keep_module_alive=True) -from .session import ExecutionSession, default_session # noqa: E402 +from .session import ExecutionSession, default_session __all__ = ["ExecutionSession", "default_session"] From d2e19d2a642e97fa485f244b65ecb054129714b8 Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Tue, 22 Sep 2026 19:35:54 +0800 Subject: [PATCH 4/4] test(orcjit): preserve injected context read --- addons/tvm_ffi_orcjit/tests/sources/c/test_context.c | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c index 12aaa01c9..8e6b0a5c0 100644 --- a/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c +++ b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c @@ -19,7 +19,9 @@ #include -TVM_FFI_DLL_EXPORT void* __tvm_ffi__library_ctx = NULL; +// The host writes this slot after compilation and before constructors run. +// Volatile keeps the constructor's read from being folded to NULL. +TVM_FFI_DLL_EXPORT void* volatile __tvm_ffi__library_ctx = NULL; static int context_was_set_during_init = 0; static void record_context_during_init(void) {