This is an automated email from the ASF dual-hosted git repository.
cyx-6 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new e251c6ab [FIX] Load ORC JIT addon through TVM-FFI (#809)
e251c6ab is described below
commit e251c6abee7a111c6467f84d854be512e33e2db0
Author: Yaxing Cai <[email protected]>
AuthorDate: Tue Sep 22 20:31:48 2026 +0800
[FIX] Load ORC JIT addon through TVM-FFI (#809)
## Why
#807 localized the ORCJIT addon and its linked runtimes, but retained a
second raw `ctypes.CDLL(..., RTLD_NODELETE)` handle to keep the addon
mapped through process shutdown. The addon should use the standard
`tvm_ffi.load_module` path exclusively.
Simply deleting the raw handle exposed a SIGSEGV during process exit
when a JIT module was pinned with `keep_module_alive=True`. TVM-FFI's
`ModuleGlobals` describes its entries as modules that should not be
unloaded, but the registry itself was a destructible function-local
static. Static teardown could therefore unload a DSO while the
process-global function table and other dependents still held callbacks
whose code belonged to it.
## What changed
- load and retain the ORCJIT addon only with `tvm_ffi.load_module(...,
keep_module_alive=True)`
- remove the duplicate `ctypes.CDLL`, `RTLD_NODELETE`, and Windows
DLL-directory handling
- stop library discovery after the first valid install/editable
candidate
- make TVM-FFI's keep-alive module registry genuinely process-lifetime,
consistent with its contract and with the existing process-lifetime
global function table
- update `keep_module_alive` documentation to state that pinned modules
live for the duration of the process
- add a child-process regression that pins a JIT module, drops the local
reference, and verifies clean shutdown
## Validation
- before the `ModuleGlobals` lifetime fix, the new shutdown reproducer
exited with SIGSEGV (`returncode == -11`)
- after the fix, the shutdown reproducer exits cleanly
- native `manylinux_2_28` ARM64 container with LLVM 22.1.0
- ORCJIT addon tests: 92 passed, 23 skipped in both serial and
pytest-xdist runs
- TVM-FFI build/load regression tests: 8 passed
- C and C++ ORCJIT quick-start examples
- `nm -D --defined-only`: exactly `TVMFFIOrcJITInitialize`
- all repository pre-commit hooks on changed files
- hosted CI pending
---------
Co-authored-by: tqchen <[email protected]>
---
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 | 28 +---
.../python/tvm_ffi_orcjit/session.py | 27 ++--
.../tvm_ffi_orcjit/src/ffi/llvm_patches/README.md | 19 ++-
.../src/ffi/llvm_patches/init_fini_plugin.h | 10 +-
.../src/ffi/llvm_patches/macho_cxa_atexit_shim.cc | 4 +-
.../src/ffi/llvm_patches/macho_cxa_atexit_shim.h | 3 +-
.../ffi/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 +-
.../tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h | 15 +-
addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 22 ++-
addons/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 +-
.../tvm_ffi_orcjit/tests/sources/c/test_context.c | 28 +++-
.../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 +-
.../tests/sources/cc/test_link_order_caller.cc | 2 +-
addons/tvm_ffi_orcjit/tests/test_basic.py | 73 ++++++++-
.../tests/test_library_visibility.py | 37 +++++
.../tests/test_session_load_module.py | 8 +-
python/tvm_ffi/cpp/extension.py | 4 +-
python/tvm_ffi/module.py | 2 +-
src/ffi/extra/module.cc | 8 +-
31 files changed, 512 insertions(+), 203 deletions(-)
diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt
b/addons/tvm_ffi_orcjit/CMakeLists.txt
index af79db90..9ced5a16 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 01472b98..9ab675ee 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 a8be7469..88ea74ff 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 e1650303..bc3cdee7 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 @@ Examples
"""
-import ctypes
-import os
import platform
from pathlib import Path
@@ -53,29 +51,19 @@ _LIB_PATH = [
]
_lib_path = None
for path in _LIB_PATH:
- if path.exists():
- _ = load_module(str(path))
+ 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."
)
-# 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 e00e9186..bc1c1c3a 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 @@ class ExecutionSession(Object):
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
@@ -190,8 +191,8 @@ class ExecutionSession(Object):
``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
@@ -243,10 +244,9 @@ class ExecutionSession(Object):
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 c2706a71..2b43c8a4 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 21262e79..ee01a628 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 0e1597ac..f2cd4824 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<void*>(&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 8aee61f0..f215e209 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 <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/Support/Error.h>
#include <utility>
#include <vector>
@@ -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 59f86004..901f627f 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 <psapi.h>
// clang-format on
+#include <array>
#include <string>
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<const char*, 4> 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<void*>(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<HMODULE, kRuntimeDLLs.size()> kRuntimeHandles = []()
{
+ std::array<HMODULE, kRuntimeDLLs.size()> 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<void*>(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 a9f1dc1e..b0925214 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<Module>(this);
}
-void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) {
+void ORCJITDynamicLibraryObj::WaitForInitializers() {
+ std::unique_lock<std::mutex> 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<std::mutex> lock(initializer_mutex_);
+ return initializer_depth_ == 0 || initializer_thread_ ==
std::this_thread::get_id();
+}
+
+void ORCJITDynamicLibraryObj::BeginInitializerRun() {
+ std::lock_guard<std::mutex> 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<std::mutex> 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<ORCJITExecutionSessionObj::InitFiniEntry>& 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<ORCJITExecutionSessionObj::InitFiniEntry> init;
+ while (true) {
+ WaitForInitializers();
+ std::unique_lock<std::mutex> 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<llvm::orc::ExecutorSymbolDef> symbol_or_err =
llvm::orc::ExecutorSymbolDef();
std::vector<ORCJITExecutionSessionObj::InitFiniEntry> init;
- {
- std::lock_guard<std::mutex> lock(session_->mutex_);
+ while (true) {
+ WaitForInitializers();
+ std::unique_lock<std::mutex> 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*>();
}
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<void**>(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<void**>(GetSymbol(symbol::tvm_ffi_library_ctx, false))) {
*ctx_addr = this;
}
Module::VisitContextSymbols([this](const String& name, void* symbol) {
- if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name))) {
+ if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name, false))) {
*ctx_addr = symbol;
}
});
+ RunPendingInitializers();
}
llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() {
@@ -337,9 +414,9 @@ llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib()
{
}
Optional<Function> 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<Function>
ORCJITDynamicLibraryObj::GetFunction(const String& name) {
//-------------------------------------
static void RegisterOrcJITFunctions() {
- static bool registered = false;
- if (registered) return;
- registered = true;
-
- namespace refl = tvm::ffi::reflection;
-
- refl::ObjectDef<ORCJITExecutionSessionObj>();
-
- refl::GlobalDef()
- .def("tvm_ffi_orcjit.ExecutionSession",
- [](const Optional<Variant<String, Bytes>>& 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<Variant<String, Bytes>>& objects,
- const String& name, const Optional<String>& cxx_runtime_path,
- const Optional<String>& 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<ORCJITExecutionSessionObj>();
+
+ refl::GlobalDef()
+ .def("tvm_ffi_orcjit.ExecutionSession",
+ [](const Optional<Variant<String, Bytes>>& 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<Variant<String, Bytes>>& objects,
+ const String& name, const Optional<String>& cxx_runtime_path,
+ const Optional<String>& 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 b458b127..370dffdd 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 <tvm/ffi/object.h>
#include <tvm/ffi/string.h>
+#include <condition_variable>
+#include <cstddef>
#include <memory>
+#include <mutex>
+#include <thread>
+#include <vector>
#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<ORCJITExecutionSessionObj::InitFiniEntry>& 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 a970c976..279d5478 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 <llvm/Support/Error.h>
#include <llvm/Support/FormatVariadic.h>
+#include <tvm/ffi/error.h>
#include <algorithm>
#include <optional>
@@ -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<Slab> 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 a71487dc..6ec89707 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<std::unique_ptr<Slab>>`. 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 cbe17859..bb56eee4 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 <llvm/ADT/ScopeExit.h>
#include <llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h>
#include <llvm/ExecutionEngine/Orc/ExecutionUtils.h>
#include <llvm/ExecutionEngine/Orc/LLJIT.h>
@@ -40,6 +41,7 @@
#include <cstddef>
#include <cstdint>
#include <mutex>
+#include <sstream>
#if defined(__linux__) && defined(__GLIBCXX__)
#include <bits/functexcept.h>
@@ -198,7 +200,7 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const
Optional<Variant<Stri
// x86_64, ±4 GB AArch64). Eliminates scattered-mmap relocation overflow
// (LLVM #173269).
//
- // slab_size_bytes: 0 = arch default (1 GB x86_64 / AArch64, with fallback),
+ // slab_size_bytes: 0 = default (64 MB, with fallback),
// >0 = 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<Variant<Stri
auto page_size = llvm::sys::Process::getPageSizeEstimate();
size_t slab_size;
if (slab_size_bytes > 0) {
+ constexpr std::size_t kMinCustomSlabSize = 2 * Slab::kCommitGranularity;
+ TVM_FFI_CHECK(static_cast<uint64_t>(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<size_t>(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<ORCJITDynamicLibraryObj>(GetRef<ORCJITExecutionSession>(this),
- &jit_dylib,
jit_.get(), lib_name);
-
#ifdef __APPLE__
// Inject ___cxa_atexit on the user JITDylib so it wins over <Platform>'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<ORCJITDynamicLibraryObj>(GetRef<ORCJITExecutionSession>(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<std::mutex> lock(mutex_);
return static_cast<int64_t>(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 7c2ea4ef..b08b7c61 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 <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
-#include <atomic>
#include <memory>
#include <mutex>
#include <string>
@@ -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<llvm::orc::LLJIT> jit_;
/*! \brief Counter for auto-generating library names */
- std::atomic<int> dylib_counter_{0};
+ int dylib_counter_{0};
/*! \brief Compiler-selected C++ runtime search JITDylibs, keyed by
shared-library path. */
std::unordered_map<std::string, llvm::orc::JITDylib*> 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 5c2d629c..b24e37b4 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 <sys/mman.h>
#include <algorithm>
+#include <cassert>
#include <cerrno>
#include <cstdio>
#include <cstring>
#include <limits>
+#include <thread>
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<StringError>("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<StringError>("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<std::size_t> 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<SlabPoolExhaustedError>(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 0029dc51..e7f1d215 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<SlabPoolExhaustedError> {
* 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<FreeBlock> free_list_non_exec_;
std::vector<FreeBlock> 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<std::atomic<std::uint8_t>[]> 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 92bca5e5..e5e0f338 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 beb9f51b..379946bc 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 d4a40760..8e6b0a5c 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,22 @@
#include <tvm/ffi/c_api.h>
-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) {
+ 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 +46,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 21f26df9..377d3f8c 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 <tvm/ffi/c_api.h>
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 215a267b..9214ed75 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 <tvm/ffi/c_api.h>
-/* 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 c878c6ea..a5df5815 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 <tvm/ffi/function.h>
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 85dde579..63a5b824 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 c7b9e822..570e4151 100644
--- a/addons/tvm_ffi_orcjit/tests/test_basic.py
+++ b/addons/tvm_ffi_orcjit/tests/test_basic.py
@@ -106,6 +106,14 @@ class Variant:
"""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
[email protected]("v", _all_variants, ids=_variant_id)
[email protected]("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 test_ctor_dtor(v: Variant) -> None:
assert "<dtors>" 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
+
+
[email protected](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_library_visibility.py
b/addons/tvm_ffi_orcjit/tests/test_library_visibility.py
index 8f853313..bd8833aa 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 ctypes
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/addons/tvm_ffi_orcjit/tests/test_session_load_module.py
b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py
index 1b11b753..562ede31 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 test_load_module_expands_embedded_library_bin(tmp_path:
Path) -> None:
# ---------------------------------------------------------------------------
# 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.
# ---------------------------------------------------------------------------
diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py
index 09bb2f6a..f11b3284 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 8c69eaeb..ee795dd8 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 31bb95bb..f8c2cd3c 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:
