From ff5485ab0c5c273fa81293d2fc83d4108b4cbeb3 Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Thu, 9 Jul 2026 06:55:47 +0000 Subject: [PATCH] [ORCJIT] Add shared session and high-level load_module Introduce a process-wide shared ExecutionSession and a high-level load_module unit operation on the ORCJIT addon, replacing the low-level create_library / add / set_link_order surface. - Session: GlobalDefault() singleton resolving the ORC runtime path from a Python-registered tvm_ffi_orcjit.DefaultOrcRuntimePath hook; plain (non-recursive) mutex; low-level dylib ops (add object / lookup / finalize) are internal, composed by LoadModule. - load_module(objects, name, keep_module_alive): links one or more object files/images into a fresh JITDylib, injects context symbols eagerly, and expands any embedded library binary. keep_module_alive pins the module in the runtime's global registry so JIT-allocated Objects may outlive the local handle. - Reduced in-addon library-binary parser with adversarial bounds checks. - CI: build the wheel against the in-tree core (--no-build-isolation) so the addon's tvm_ffi ABI matches the core the tests reinstall. - Link libstdc++/libgcc dynamically and hide static-archive (LLVM/zlib/zstd) symbols via --exclude-libs,ALL, so the JIT resolves the C++ runtime from the process without leaking LLVM symbols that could interpose with the host. - Tests for session/load_module, malformed blobs, concurrency, and container returns; docs and examples updated for the unit-op API. --- .github/actions/build-orcjit-wheel/action.yml | 4 + addons/tvm_ffi_orcjit/CMakeLists.txt | 18 + addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md | 177 ++---- addons/tvm_ffi_orcjit/README.md | 79 ++- .../examples/quick-start/add.cc | 4 - .../examples/quick-start/run.py | 32 +- .../python/tvm_ffi_orcjit/__init__.py | 18 +- .../python/tvm_ffi_orcjit/_ffi_api.py | 2 +- .../python/tvm_ffi_orcjit/dylib.py | 85 --- .../python/tvm_ffi_orcjit/session.py | 194 +++++-- addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc | 304 +++++++--- addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h | 53 +- .../tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 70 ++- .../tvm_ffi_orcjit/src/ffi/orcjit_session.h | 88 ++- addons/tvm_ffi_orcjit/tests/CMakeLists.txt | 1 + addons/tvm_ffi_orcjit/tests/README.md | 2 +- .../tests/sources/cc/test_containers.cc | 71 +++ addons/tvm_ffi_orcjit/tests/test_basic.py | 542 ++++++++---------- .../tests/test_memory_manager.py | 103 ++-- .../tests/test_session_load_module.py | 437 ++++++++++++++ python/tvm_ffi/cpp/extension.py | 15 +- 21 files changed, 1482 insertions(+), 817 deletions(-) delete mode 100644 addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/dylib.py create mode 100644 addons/tvm_ffi_orcjit/tests/sources/cc/test_containers.cc create mode 100644 addons/tvm_ffi_orcjit/tests/test_session_load_module.py diff --git a/.github/actions/build-orcjit-wheel/action.yml b/.github/actions/build-orcjit-wheel/action.yml index 7cefbed5b..bc2a2ad52 100644 --- a/.github/actions/build-orcjit-wheel/action.yml +++ b/.github/actions/build-orcjit-wheel/action.yml @@ -102,6 +102,10 @@ runs: CIBW_ENVIRONMENT: LLVM_PREFIX=/opt/llvm CIBW_ENVIRONMENT_WINDOWS: LLVM_PREFIX="C:/opt/llvm" CIBW_CONTAINER_ENGINE: "docker; create_args: --volume /opt/llvm:/opt/llvm" + # Build against the in-tree core (not PyPI) so the addon's tvm_ffi ABI + # matches the core the tests reinstall; a mismatch segfaults the JIT. + CIBW_BEFORE_BUILD: pip install "scikit-build-core>=0.10.0" cmake ninja "{project}" + CIBW_BUILD_FRONTEND: "pip; args: --no-build-isolation" CIBW_TEST_REQUIRES: pytest pytest-xdist ninja CIBW_TEST_COMMAND: >- pip install --force-reinstall {project} && diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt b/addons/tvm_ffi_orcjit/CMakeLists.txt index d25b06f7f..fb33741f8 100644 --- a/addons/tvm_ffi_orcjit/CMakeLists.txt +++ b/addons/tvm_ffi_orcjit/CMakeLists.txt @@ -149,9 +149,27 @@ if (APPLE) COMMENT "Fixing libc++ rpath to use system library" ) elseif (UNIX AND NOT WIN32) + # Static libstdc++/libgcc: the conda LLVM toolchain's libstdc++ is newer than the manylinux ABI + # floor, so linking it dynamically would import too-recent GLIBCXX_* symbols and fail auditwheel. target_link_options(tvm_ffi_orcjit PRIVATE -static-libstdc++ -static-libgcc) endif () +# Hide symbols pulled from static archives so they are not re-exported and cannot interpose with a +# host process's own copies (e.g. PyTorch's bundled LLVM). Exclude LLVM/zlib/zstd by name -- NOT +# --exclude-libs,ALL, which would also hide the static libstdc++ symbols (operator new, __throw_*, +# ...) that liborc_rt and JIT'd C++ code resolve from this .so at run time. The LLVM list is derived +# from llvm-config so it tracks the toolchain version. +if (CMAKE_SYSTEM_NAME MATCHES "Linux|Android|FreeBSD|NetBSD|OpenBSD" AND CMAKE_CXX_COMPILER_ID + MATCHES "GNU|Clang" +) + # --exclude-libs needs exact archive filenames (no globs), so turn each -lLLVM* from llvm-config + # into libLLVM*.a and add zlib/zstd. One colon-separated list -- a -Wl comma list would be split. + set(_exclude_archives ${_llvm_libs_list} "-lz" "-lzstd") + list(TRANSFORM _exclude_archives REPLACE "^-l(.+)$" "lib\\1.a") + list(JOIN _exclude_archives ":" _exclude_libs_arg) + target_link_options(tvm_ffi_orcjit PRIVATE "-Wl,--exclude-libs=${_exclude_libs_arg}") +endif () + # ---- Find and bundle liborc_rt ---- # Platform notes: # diff --git a/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md b/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md index af3a7e7e8..01472b983 100644 --- a/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md +++ b/addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md @@ -345,128 +345,40 @@ With the background above, here is how the addon maps onto the concepts. ```text Python / C++ API LLVM ORC v2 concept ───────────────────────────────────────────────────────────────── +default_session() shared (leaked) LLJIT + ExecutionSession ExecutionSession LLJIT + ExecutionSession -DynamicLibrary JITDylib -dylib.add("foo.o") ObjectLinkingLayer.add(buffer) -dylib.get_function("add") ExecutionSession.lookup("__tvm_ffi_add") -dylib.set_link_order([a, b]) JITDylib.setLinkOrder([a, b, main, ...]) +session.load_module(objs) createJITDylib + addObjectFile(s) + wire imports +(returned) Module JITDylib, wrapped as a tvm_ffi.Module +module.get_function("add") ExecutionSession.lookup("__tvm_ffi_add") ``` -### 4.2 Object File Loading Pipeline - -```text -dylib.add("foo.o") - │ - ▼ ORCJITDynamicLibraryObj::AddObjectFile() - │ jit_->addObjectFile(*dylib_, MemoryBuffer) - │ - ▼ ObjectTransformLayer (macOS: skipped; Linux/Win: may strip .pdata) - │ Windows: strip .pdata/.xdata avoids JITLink COMDAT limitation - │ fix __ImageBase fixes Pointer32NB relocations - │ - ▼ ObjectLinkingLayer (JITLink) - │ Parse object → LinkGraph - │ Run InitFiniPlugin passes: - │ PrePrunePasses: mark .init_array / .ctors blocks as live - │ PostAllocationPasses: (Windows) set __ImageBase, strip SEH sections - │ PostFixupPasses: extract resolved init/fini function pointers - │ → session.AddPendingInitializer(dylib, entry) - │ Resolve relocations via ExecutionSession.lookup() - │ Allocate JIT memory, write code, mark executable - │ - ▼ JITDylib symbol table updated - (symbols are defined but constructors not yet run) -``` - -### 4.3 Symbol Lookup and Initialization - -```text -dylib.get_function("add") - │ - ▼ ORCJITDynamicLibraryObj::GetFunction("add") - │ → GetSymbol("__tvm_ffi_add") - │ build JITDylibSearchOrder: [this, linked dylibs, LLJIT default] - │ jit_->getExecutionSession().lookup(order, "__tvm_ffi_add") - │ ← ExecutorAddr - │ - ▼ Linux/Windows: RunPendingInitializers() - │ Sort entries by priority - │ Call each function pointer (C++ ctors, .init_array entries) - │ - ▼ macOS: jit_->initialize(*dylib_) - │ ORC MachOPlatform calls __mod_init_func pointers - │ - ▼ Wrap raw function pointer as tvm_ffi::Function - via Function::FromPacked lambda (marshals AnyView args ↔ TVMFFIAny) -``` - -### 4.4 InitFiniPlugin Detail - -`InitFiniPlugin` is an `ObjectLinkingLayer::Plugin` — it receives callbacks during -JITLink's pass pipeline for every object being linked. - -```text -JITLink pass pipeline for each object file: -───────────────────────────────────────────────────────────── -PrePrunePasses - → InitFiniPlugin::modifyPassConfig (PrePrunePasses) - For each section named .init_array* / .ctors* / .fini_array* / .dtors* - (ELF) or __DATA,__mod_init_func (Mach-O) or .CRT$XC* (COFF): - Mark all blocks in that section as "keep" (live) - → Prevents dead-code elimination from removing constructors - -PostAllocationPasses (Windows only) - → InitFiniPlugin::modifyPassConfig (PostAllocationPasses) - Set __ImageBase = lowest allocated block address - Strip .pdata / .xdata exception handler sections - -PostFixupPasses - → InitFiniPlugin::modifyPassConfig (PostFixupPasses) - Iterate all blocks in init/fini sections - Read each 8-byte slot as an ExecutorAddr - Parse section name to determine: priority, is_init vs is_fini - Call session->AddPendingInitializer(dylib, InitFiniEntry{addr, section, priority}) - (or AddPendingDeinitializer for dtors/fini_array) -``` - -### 4.5 Cross-Library Symbol Resolution - -```text -libA depends on a symbol defined in libB: - -libA.set_link_order([libB]) - → JITDylibSearchOrder for libA: [libA, libB, main(ProcessSymbols)] - -When libA's object file has an unresolved symbol "foo": - JITLink asks ExecutionSession.lookup([libA, libB, main], "foo") - → found in libB → returns libB's ExecutorAddr for "foo" - → relocation in libA patched with that address -``` - -### 4.6 Windows DLL Import Stubs - -Windows MSVC objects reference DLL functions through `__imp_XXX` pointer stubs (the -Import Address Table pattern). At static link time the linker creates these stubs. In -JIT mode there is no linker, so `DLLImportDefinitionGenerator` creates them on demand: - -```text -JITLink encounters undefined symbol "__imp_malloc" - │ - ▼ DLLImportDefinitionGenerator::tryToGenerate() - │ Search ucrtbase.dll, msvcrt.dll, then all process modules - │ for the real address of "malloc" - │ - ▼ Allocate two JIT-memory stubs: - │ __imp_malloc → 8-byte slot containing &malloc (host address) - │ malloc → x86_64 jmp [__imp_malloc] trampoline - │ - ▼ Define both symbols in the JITDylib - → JITLink can now apply PCRel32 reloc to __imp_malloc (stub is close in JIT memory) -``` - -The stubs must live in JIT-allocated memory (not at the host process address) because -x86_64 `PCRel32` relocations can only reach ±2 GB. The host's `malloc` may be farther -than 2 GB from the JIT allocation. +`session.load_module(objects)` is the sole public loader: it creates one +`JITDylib`, adds every object (path or in-memory bytes), injects context symbols +eagerly, expands any embedded library binary into an import tree, and returns a +plain `tvm_ffi.Module`. There is no incremental library API — a module is a +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 +pointer as a `tvm_ffi::Function`. Symbols resolve against the dylib's own +default link order (this dylib → Platform → process/runtime symbols); 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: + +- **`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`. +- **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 + `llvm_patches/win_dll_import_generator.h`. --- @@ -475,22 +387,17 @@ than 2 GB from the JIT allocation. ```python import tvm_ffi_orcjit as oj -# 1. Create an ExecutionSession (wraps LLJIT) -sess = oj.ExecutionSession() - -# 2. Create a JITDylib -lib = sess.create_dynamic_library() - -# 3. Load a compiled object file -# → object parsed, JITLink links it, InitFiniPlugin collects ctors -lib.add("add.o") +# 1. Get the shared ExecutionSession (wraps LLJIT; created once per process) +sess = oj.default_session() -# 4. Look up a function -# → LLVM resolves "__tvm_ffi_add", RunPendingInitializers() fires ctors -add = lib.get_function("add") # returns tvm_ffi.Function +# 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 +mod = sess.load_module("add.o") # returns a tvm_ffi.Module -# 5. Call it -result = add(3, 4) # → 7 +# 3. Look up and call a function +# → LLVM resolves "__tvm_ffi_add"; pending constructors fire on first lookup +result = mod.add(3, 4) # → 7 ``` The corresponding C++ side of `add.o`: @@ -511,7 +418,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); | Concept | What it is | Where in the addon | | --- | --- | --- | -| Object file | Container of machine code, data, symbols, relocations | Input to `dylib.add()` | +| Object file | Container of machine code, data, symbols, relocations | Input to `session.load_module()` | | Relocation | Recipe to patch a code address at link/JIT time | Applied by JITLink | | `.init_array` / `.ctors` | Array of C++ constructor pointers in ELF objects | Collected by `InitFiniPlugin` | | `ExecutionSession` | Root of the ORC JIT environment | `ORCJITExecutionSessionObj` | @@ -520,7 +427,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); | `JITLink` | LLVM's JIT-aware linker | Used inside `ObjectLinkingLayer` | | JITLink pass pipeline | Pre-prune → post-alloc → post-fixup hooks | Where `InitFiniPlugin` runs | | `DefinitionGenerator` | Fallback symbol provider | `DLLImportDefinitionGenerator` (Win) | -| Link order | Search path across JITDylibs for symbol resolution | `SetLinkOrder()` | +| 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 c68d19331..ce08ad8b8 100644 --- a/addons/tvm_ffi_orcjit/README.md +++ b/addons/tvm_ffi_orcjit/README.md @@ -24,9 +24,10 @@ TVM-FFI exported functions. ## Features - **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2 -- **Multiple Libraries**: Create separate dynamic libraries with independent symbol namespaces -- **Incremental Loading**: Add multiple object files to the same library incrementally -- **Symbol Isolation**: Different libraries can define the same symbol without conflicts +- **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) +- **Symbol Isolation**: Separate `load_module` calls define independent symbol namespaces, so they can define the same symbol without conflicts - **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 @@ -105,55 +106,54 @@ auto-discover it and `LLVM_PREFIX` is not needed. ### Basic Example -```python -from tvm_ffi_orcjit import ExecutionSession +The high-level API mirrors `tvm_ffi.load_module`: a process-wide shared session +plus a `load_module` that accepts a path, in-memory object bytes, or a list of +either, and returns a plain `tvm_ffi.Module`. -# Create an execution session -session = ExecutionSession() +```python +import tvm_ffi_orcjit as oj -# Create a dynamic library -lib = session.create_library() +# Shared process-wide session (created once, cached). +session = oj.default_session() -# Load an object file -lib.add("example.o") +# Load a single object file by path. +mod = session.load_module("example.o") -# Get and call a function -add_func = lib.get_function("add") -result = add_func(1, 2) +# Call an exported function. +result = mod.add(1, 2) print(f"Result: {result}") # Output: Result: 3 ``` -### Multiple Libraries with Symbol Isolation - -```python -session = ExecutionSession() +### Loading Multiple Objects and In-Memory Bytes -lib1 = session.create_library("lib1") -lib2 = session.create_library("lib2") +Objects passed together are linked into one module (the same way a multi-object +shared library links). Each element may be a path or an object-file image in +memory. -lib1.add("implementation_v1.o") -lib2.add("implementation_v2.o") +```python +from pathlib import Path -add_v1 = lib1.get_function("add") -add_v2 = lib2.get_function("add") +session = oj.default_session() -print(add_v1(5, 3)) # Uses implementation from lib1 -print(add_v2(5, 3)) # Uses implementation from lib2 +mod = session.load_module( + [ + "math_ops.o", # path + Path("kernel.o").read_bytes(), # in-memory object bytes + ] +) +result = mod.call_math(10, 20) ``` -### Cross-Library Linking - -```python -session = ExecutionSession() +### Isolated Sessions -base_lib = session.create_library("base") -base_lib.add("math_ops.o") +`default_session()` is shared across the process. For an isolated symbol +namespace or a tuned memory arena, construct an `ExecutionSession` directly: -caller_lib = session.create_library("caller") -caller_lib.set_link_order(base_lib) # Can resolve symbols from base_lib -caller_lib.add("caller.o") +```python +from tvm_ffi_orcjit import ExecutionSession -result = caller_lib.get_function("call_math")(10, 20) +session = ExecutionSession() # independent LLVM ExecutionSession +mod = session.load_module("impl.o") ``` ## Writing Functions for OrcJIT @@ -211,13 +211,12 @@ tvm_ffi_orcjit/ ├── src/ffi/ │ ├── orcjit_session.cc # ExecutionSession (LLJIT setup, plugins) │ ├── orcjit_session.h -│ ├── orcjit_dylib.cc # DynamicLibrary (object loading, symbol lookup) +│ ├── orcjit_dylib.cc # JIT dylib module (object loading, symbol lookup) │ ├── orcjit_dylib.h │ └── orcjit_utils.h # LLVM error handling utilities ├── python/tvm_ffi_orcjit/ │ ├── __init__.py # Module exports and library loading -│ ├── session.py # Python ExecutionSession wrapper -│ └── dylib.py # Python DynamicLibrary wrapper +│ └── session.py # Python ExecutionSession + default_session ├── tests/ # See tests/README.md └── examples/quick-start/ # Complete example with CMake ``` @@ -226,7 +225,7 @@ tvm_ffi_orcjit/ Runs on Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64) via `cibuildwheel`. Each platform builds test objects with multiple compilers -and runs the full test suite. See `.github/workflows/tvm_ffi_orcjit.yml`. +and runs the full test suite. See the `orcjit` job in `.github/workflows/ci_test.yml`. ## Troubleshooting diff --git a/addons/tvm_ffi_orcjit/examples/quick-start/add.cc b/addons/tvm_ffi_orcjit/examples/quick-start/add.cc index dd3c3e2e8..58187f84d 100644 --- a/addons/tvm_ffi_orcjit/examples/quick-start/add.cc +++ b/addons/tvm_ffi_orcjit/examples/quick-start/add.cc @@ -40,7 +40,3 @@ int fib_impl(int n) { return fib_impl(n - 1) + fib_impl(n - 2); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(fibonacci, fib_impl); - -// String concatenation example -std::string concat_impl(std::string a, std::string b) { return a + b; } -TVM_FFI_DLL_EXPORT_TYPED_FUNC(concat, concat_impl); diff --git a/addons/tvm_ffi_orcjit/examples/quick-start/run.py b/addons/tvm_ffi_orcjit/examples/quick-start/run.py index 3a648a1d3..000848cd9 100755 --- a/addons/tvm_ffi_orcjit/examples/quick-start/run.py +++ b/addons/tvm_ffi_orcjit/examples/quick-start/run.py @@ -20,9 +20,8 @@ This script demonstrates how to: 1. Compile C/C++ source files to object files using tvm_ffi.cpp.build -2. Load them into an ORC JIT ExecutionSession -3. Get functions by name -4. Call them like regular Python functions +2. Load them via the shared ORC JIT session (default_session().load_module) +3. Call the resulting module's functions like regular Python functions Usage: python run.py # Load C++ object file (add.o) @@ -37,7 +36,7 @@ from pathlib import Path import tvm_ffi.cpp -from tvm_ffi_orcjit import ExecutionSession +import tvm_ffi_orcjit as oj SCRIPT_DIR = Path(__file__).resolve().parent @@ -71,47 +70,34 @@ def _build_object(lang: str) -> str: def _run_tests(obj_file: str, lang: str) -> None: """Load object file and run all test assertions. - All JIT references (functions, lib, session) are released automatically + All JIT references (functions, module) are released automatically when this function returns. """ print(f"Loading object file: {obj_file} (lang={lang})") - # Create execution session and dynamic library - session = ExecutionSession() - lib = session.create_library() - lib.add(obj_file) + # Load the object into a module on the shared session. + mod = oj.default_session().load_module(obj_file) print("Object file loaded successfully\n") # Get and call the 'add' function print("=== Testing add function ===") - add = lib.get_function("add") - result = add(10, 20) + result = mod.add(10, 20) print(f"add(10, 20) = {result}") assert result == 30, f"Expected 30, got {result}" # Get and call the 'multiply' function print("\n=== Testing multiply function ===") - multiply = lib.get_function("multiply") - result = multiply(7, 6) + result = mod.multiply(7, 6) print(f"multiply(7, 6) = {result}") assert result == 42, f"Expected 42, got {result}" # Get and call the 'fibonacci' function print("\n=== Testing fibonacci function ===") - fibonacci = lib.get_function("fibonacci") - result = fibonacci(10) + result = mod.fibonacci(10) print(f"fibonacci(10) = {result}") assert result == 55, f"Expected 55, got {result}" - if lang == "cpp": - # String concatenation only available in C++ variant (uses std::string) - print("\n=== Testing concat function ===") - concat = lib.get_function("concat") - result = concat("Hello, ", "World!") - print(f"concat('Hello, ', 'World!') = '{result}'") - assert result == "Hello, World!", f"Expected 'Hello, World!', got '{result}'" - print("\n" + "=" * 50) print("All tests passed successfully!") print("=" * 50) 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 1a94d8922..d555efd25 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py @@ -20,13 +20,12 @@ This module provides functionality to load object files (.o) compiled with TVM-FFI exports using LLVM ORC JIT v2. -Example: - >>> from tvm_ffi_orcjit import ExecutionSession - >>> session = ExecutionSession() - >>> lib = session.create_library() - >>> lib.add("example.o") - >>> func = lib.get_function("my_function") - >>> result = func(arg1, arg2) +Examples +-------- +>>> import tvm_ffi_orcjit as oj +>>> session = oj.default_session() +>>> mod = session.load_module("example.o") +>>> result = mod.my_function(arg1, arg2) """ @@ -81,10 +80,9 @@ warnings.warn(f"Failed to explicitly initialize orcjit library: {e}") -from .dylib import DynamicLibrary -from .session import ExecutionSession +from .session import ExecutionSession, default_session -__all__ = ["DynamicLibrary", "ExecutionSession"] +__all__ = ["ExecutionSession", "default_session"] try: from importlib.metadata import version diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/_ffi_api.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/_ffi_api.py index fb2464b40..1ed88e3ee 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/_ffi_api.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/_ffi_api.py @@ -18,4 +18,4 @@ import tvm_ffi -tvm_ffi.init_ffi_api("orcjit", __name__) +tvm_ffi.init_ffi_api("tvm_ffi_orcjit", __name__) diff --git a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/dylib.py b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/dylib.py deleted file mode 100644 index 6cfda6cfb..000000000 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/dylib.py +++ /dev/null @@ -1,85 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""ORC JIT Dynamic Library.""" - -from __future__ import annotations - -from pathlib import Path - -from tvm_ffi import Module - - -class DynamicLibrary(Module): - """ORC JIT Dynamic Library (JITDylib). - - Represents a collection of symbols that can be loaded from object files and linked - against other dynamic libraries. Supports JIT compilation and symbol resolution. - - Examples - -------- - >>> session = ExecutionSession() - >>> lib = session.create_library() - >>> lib.add("add.o") - >>> lib.add("multiply.o") - >>> add_func = lib.get_function("add") - >>> result = add_func(1, 2) - - """ - - def add(self, object_file: str | Path) -> None: - """Add an object file to this dynamic library. - - Parameters - ---------- - object_file : str or Path - Path to the object file to load. - - Examples - -------- - >>> lib.add("add.o") - >>> lib.add(Path("multiply.o")) - - """ - if isinstance(object_file, Path): - object_file = str(object_file) - self.get_function("orcjit.add_object_file")(object_file) - - def set_link_order(self, *libraries: DynamicLibrary) -> None: - """Set the link order for symbol resolution. - - When resolving symbols, this library will search in the specified libraries - in the order provided. This replaces any previous link order. - - Parameters - ---------- - *libraries : DynamicLibrary - One or more dynamic libraries to search for symbols (in order). - - Examples - -------- - >>> session = ExecutionSession() - >>> lib_utils = session.create_library() - >>> lib_utils.add("utils.o") - >>> lib_core = session.create_library() - >>> lib_core.add("core.o") - >>> lib_main = session.create_library() - >>> lib_main.add("main.o") - >>> # main can call symbols from utils and core (utils searched first) - >>> lib_main.set_link_order(lib_utils, lib_core) - - """ - self.get_function("orcjit.set_link_order")(list(libraries)) 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 391fd35d4..89e6b58ec 100644 --- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py +++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py @@ -19,11 +19,17 @@ from __future__ import annotations import sys +import threading +from pathlib import Path +from typing import TYPE_CHECKING -from tvm_ffi import Object, register_object +import tvm_ffi +from tvm_ffi import Module, Object, register_object from . import _ffi_api, _lib_dir -from .dylib import DynamicLibrary + +if TYPE_CHECKING: + from collections.abc import Sequence def _find_orc_rt_library() -> str | None: @@ -51,43 +57,56 @@ def _find_orc_rt_library() -> str | None: return None -@register_object("orcjit.ExecutionSession") +@tvm_ffi.register_global_func("tvm_ffi_orcjit.DefaultOrcRuntimePath", override=True) +def _default_orc_runtime_path() -> str: + """Resolve the ORC runtime path for the shared session (C++ hook). + + Called once by the C++ ``GlobalDefault`` singleton on first use. Returns the + runtime bundled next to this extension, or an empty string to run with no + ORC platform (macOS/Windows, or when no bundled runtime is found). + """ + return _find_orc_rt_library() or "" + + +@register_object("tvm_ffi_orcjit.ExecutionSession") class ExecutionSession(Object): """ORC JIT Execution Session. Manages the LLVM ORC JIT execution environment and creates dynamic libraries (JITDylibs). This is the top-level context for JIT compilation and symbol management. + Prefer :func:`default_session` for a process-wide shared session; construct + an ``ExecutionSession`` directly only for an isolated session or a tuned arena. + Examples -------- >>> session = ExecutionSession() - >>> lib = session.create_library(name="main") - >>> lib.add("add.o") - >>> add_func = lib.get_function("add") + >>> mod = session.load_module("add.o") + >>> add_func = mod.get_function("add") """ def __init__(self, orc_rt_path: str | None = None, slab_size: int = 0) -> None: """Initialize ExecutionSession. - Args: - orc_rt_path: Optional path to the liborc_rt library. If not provided, - it will be automatically discovered using clang. - slab_size: 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). - - 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 fit a normal slab trigger - a power-of-2 larger slab (slab_size, 2*slab_size, ...) - sized to fit. Drained slabs stay mapped until the - session is destroyed or ``clear_free_slabs()`` is - called. + Parameters + ---------- + orc_rt_path : str, optional + Optional path to the liborc_rt library. If not provided, it will be + automatically discovered using clang. + 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). + + 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 + fit a normal slab trigger a power-of-2 larger slab (slab_size, + 2*slab_size, ...) sized to fit. Drained slabs stay mapped until the + session is destroyed or ``clear_free_slabs()`` is called. """ if orc_rt_path is None: @@ -96,20 +115,78 @@ def __init__(self, orc_rt_path: str | None = None, slab_size: int = 0) -> None: orc_rt_path = "" self.__init_handle_by_constructor__(_ffi_api.ExecutionSession, orc_rt_path, slab_size) # type: ignore - def create_library(self, name: str = "") -> DynamicLibrary: - """Create a new dynamic library associated with this execution session. - - Args: - name: Optional name for the library. If empty, a unique name will be generated. - - Returns: - A new DynamicLibrary instance. + def load_module( + self, + objects: str | Path | bytes | bytearray | Sequence[str | Path | bytes | bytearray], + name: str = "", + keep_module_alive: bool = False, + ) -> Module: + """Load one or more object files into a fresh module. + + All objects are linked into one dynamic library (JITDylib), context + symbols are wired up, and any embedded library binary is expanded — so + the result behaves like a module loaded by :func:`tvm_ffi.load_module`. + + Parameters + ---------- + objects : str or Path or bytes or bytearray, or a sequence of them + Object files to load, given as paths and/or in-memory object-file + images. A single item is accepted as shorthand for a one-element list. + name : str + Optional name for the underlying library. Auto-generated if empty. + keep_module_alive : bool + If True, pin the module in the runtime's process-global registry + (see Notes). Defaults to False. + + Returns + ------- + Module + A :class:`tvm_ffi.Module` whose imports and library context are fully + wired. + + Notes + ----- + ``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 + 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 + lifetime and its JIT memory is reclaimed on drop; callers must ensure + no JIT-produced Object outlives the returned module. + + Pinning is transitive: the module holds a strong reference to its + :class:`ExecutionSession`, so pinning one module also keeps that + session (and its slab-pool memory manager) alive for the process + lifetime. Reclaim is slab-granular — a slab shared between a pinned + and an unpinned module stays mapped until the pinned module is gone. + + Examples + -------- + >>> session = default_session() + >>> mod = session.load_module(["a.o", "b.o", Path("c.o").read_bytes()]) + >>> mod.my_function(1, 2) """ - handle = _ffi_api.ExecutionSessionCreateDynamicLibrary(self, name) # type: ignore - lib = DynamicLibrary.__new__(DynamicLibrary) - lib.__move_handle_from__(handle) - return lib + if isinstance(objects, (str, Path, bytes, bytearray)): + objects = [objects] + normalized: list[str | bytes] = [] + for obj in objects: + if isinstance(obj, (bytes, bytearray)): + normalized.append(bytes(obj)) + elif isinstance(obj, (str, Path)): + normalized.append(str(obj)) + else: + raise TypeError( + "load_module objects must be a path (str or Path) or object-file " + f"bytes, but got {type(obj).__name__}" + ) + mod = _ffi_api.SessionLoadModule(self, normalized, name) # type: ignore + if keep_module_alive: + tvm_ffi._ffi_api.ModuleGlobalsAdd(mod) # type: ignore + return mod def clear_free_slabs(self) -> int: """Release drained slabs (no live JIT allocations) back to the OS. @@ -123,10 +200,51 @@ def clear_free_slabs(self) -> int: returned, the C++ destructor has finished and the slab's live count reflects the drop. - Returns: + Returns + ------- + int Number of slabs actually munmap'd. Returns 0 on macOS/Windows (slab pool compiled out) or when the pool is disabled via ``slab_size=-1``. """ - return int(_ffi_api.ExecutionSessionClearFreeSlabs(self)) # type: ignore + return int(_ffi_api.SessionClearFreeSlabs(self)) # type: ignore + + +_default_session: ExecutionSession | None = None +_default_session_lock = threading.Lock() + + +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. + + The ORC runtime path is resolved once, on first use, by the C++ singleton + calling the registered ``tvm_ffi_orcjit.DefaultOrcRuntimePath`` hook (which + returns the runtime bundled next to this extension, or none). For an isolated + session or a tuned arena, construct an :class:`ExecutionSession` directly. + + Returns + ------- + ExecutionSession + The shared execution session. + + Examples + -------- + >>> import tvm_ffi_orcjit as oj + >>> session = oj.default_session() + >>> mod = session.load_module("dylib0.o") + + """ + global _default_session # noqa: PLW0603 — module-level singleton cache + # Double-checked locking: the FFI call releases the GIL, so guard the + # first-call caching to avoid two threads racing to fetch it. + if _default_session is None: + with _default_session_lock: + if _default_session is None: + _default_session = _ffi_api.GlobalDefaultSession() + return _default_session diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc index f6bd329a1..6e98e863b 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc @@ -32,12 +32,18 @@ #include #include #include +#include #include #include #include #include #include +#include +#include +#include +#include + #include "orcjit_session.h" #include "orcjit_utils.h" @@ -60,6 +66,122 @@ struct DylibFnContextWithModule { }; void DeleteDylibFnContextWithModule(void* p) { delete static_cast(p); } + +// Minimal little-endian reader for the embedded library-binary blob. The addon +// links only against public tvm-ffi headers, so it re-implements the reduced +// subset of the core parser it needs (see ProcessEmbeddedLibraryBin). All +// supported targets are little-endian (x86_64, aarch64, arm64). +class BlobReader { + public: + BlobReader(const char* data, size_t size) : data_(data), size_(size) {} + + // All bounds checks use subtraction (size_ - cursor_ never underflows because + // cursor_ <= size_ is an invariant) to avoid cursor_ + n overflowing. + uint64_t ReadU64() { + TVM_FFI_CHECK(size_ - cursor_ >= sizeof(uint64_t), RuntimeError) + << "Corrupt library binary: unexpected end of blob"; + uint64_t value = 0; + for (size_t i = 0; i < sizeof(uint64_t); ++i) { + value |= static_cast(static_cast(data_[cursor_ + i])) << (i * 8); + } + cursor_ += sizeof(uint64_t); + return value; + } + + std::string ReadString() { + uint64_t nbytes = ReadU64(); + TVM_FFI_CHECK(nbytes <= size_ - cursor_, RuntimeError) + << "Corrupt library binary: string length exceeds blob"; + std::string out(data_ + cursor_, static_cast(nbytes)); + cursor_ += static_cast(nbytes); + return out; + } + + std::vector ReadU64Vector() { + uint64_t count = ReadU64(); + // Guard reserve() against a corrupt count: each element needs 8 more bytes, + // so count can't exceed the remaining byte budget. Prevents an OOM abort. + TVM_FFI_CHECK(count <= (size_ - cursor_) / sizeof(uint64_t), RuntimeError) + << "Corrupt library binary: vector size exceeds remaining blob"; + std::vector out; + out.reserve(static_cast(count)); + for (uint64_t i = 0; i < count; ++i) { + out.push_back(ReadU64()); + } + return out; + } + + private: + const char* data_; + size_t size_; + size_t cursor_{0}; +}; + +// Reduced re-implementation of the core ProcessLibraryBin parser. +// Deserializes the modules embedded in a __tvm_ffi__library_bin blob, wires up +// the import tree, and returns the root module. The "_lib" placeholder is +// filled by `lib_module` (the JIT dylib module itself). Custom modules are +// deserialized through the public ffi.Module.load_from_bytes. registry. +// +// Blob layout (little-endian): +// > > +// [] [] ... +// where vec = , str = , +// and the import tree is a CSR: module i imports child_indices[indptr[i]..]. +Module ProcessEmbeddedLibraryBin(const char* library_bin, const Module& lib_module) { + uint64_t nbytes = 0; + for (size_t i = 0; i < sizeof(nbytes); ++i) { + nbytes |= static_cast(static_cast(library_bin[i])) << (i * 8); + } + BlobReader reader(library_bin + sizeof(nbytes), static_cast(nbytes)); + + std::vector indptr = reader.ReadU64Vector(); + std::vector child_indices = reader.ReadU64Vector(); + // Validate the CSR shape up front so the wiring loop below cannot read out of + // bounds: indptr has n+1 entries for n>=1 modules, starts at 0, is + // non-decreasing, and its tail stays within child_indices. + TVM_FFI_CHECK(indptr.size() >= 2, RuntimeError) + << "Corrupt library binary: import tree indptr must have at least 2 entries"; + TVM_FFI_CHECK(indptr.front() == 0, RuntimeError) + << "Corrupt library binary: import tree indptr must start at 0"; + for (size_t i = 0; i + 1 < indptr.size(); ++i) { + TVM_FFI_CHECK(indptr[i] <= indptr[i + 1], RuntimeError) + << "Corrupt library binary: import tree indptr must be non-decreasing"; + } + TVM_FFI_CHECK(indptr.back() <= child_indices.size(), RuntimeError) + << "Corrupt library binary: import tree indptr exceeds child index count"; + size_t num_modules = indptr.size() - 1; + + std::vector modules; + modules.reserve(num_modules); + for (size_t i = 0; i < num_modules; ++i) { + std::string kind = reader.ReadString(); + if (kind == "_lib") { + // Placeholder for the symbol source — the JIT dylib module itself. + modules.push_back(lib_module); + } else { + std::string module_bytes = reader.ReadString(); + auto floader = Function::GetGlobal("ffi.Module.load_from_bytes." + kind); + TVM_FFI_CHECK(floader.has_value(), RuntimeError) + << "Library binary embeds a {" << kind << "} module but loader " + << "ffi.Module.load_from_bytes." << kind << " is not registered"; + // cast throws if the loader returned a non-Module (e.g. None), and + // Module is non-nullable, so the result is always a valid module. + modules.push_back((*floader)(Bytes(module_bytes)).cast()); + } + } + + // Wire the import tree (CSR) using the public ModuleObj::ImportModule. + for (size_t i = 0; i < modules.size(); ++i) { + for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j) { + uint64_t child = child_indices[j]; + TVM_FFI_CHECK(child < modules.size(), RuntimeError) + << "Corrupt library binary: child index out of range"; + modules[i]->ImportModule(modules[static_cast(child)]); + } + } + return modules[0]; +} } // namespace ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session, @@ -71,85 +193,117 @@ ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session, } ORCJITDynamicLibraryObj::~ORCJITDynamicLibraryObj() { - // Step 1: run static destructors for code in this JITDylib via our - // InitFiniPlugin (uniform across Linux / macOS / Windows — see the - // plugin docstring in llvm_patches/init_fini_plugin.h). - session_->RunPendingDeinitializers(GetJITDylib()); + // Step 1: run this dylib's static destructors. Drain the entries under the + // lock but run them released — a JIT'd dtor may re-enter the session on this + // thread, which the plain mutex could not survive while held. + std::vector deinit; + { + std::lock_guard lock(session_->mutex_); + deinit = session_->DrainPendingDeinitializers(GetJITDylib()); + } + ORCJITExecutionSessionObj::RunInitFiniEntries(deinit); #ifdef __APPLE__ - // Step 1b (macOS only): drain per-dylib __cxa_atexit registrations in - // LIFO. Clang on Darwin lowers __attribute__((destructor)) and C++ - // global dtors as __cxa_atexit(fn, arg, &__dso_handle) registrations - // during init; the shim in llvm_patches/macho_cxa_atexit_shim.cc - // captured them into cxa_atexit_records_ via the TLS pointer - // published by the GetSymbol-time scope. + // Drain per-dylib __cxa_atexit registrations (LIFO) captured during init; see + // llvm_patches/macho_cxa_atexit_shim.h. DrainCxaAtexit(cxa_atexit_records_); #endif - // Step 2: remove the JITDylib from the ExecutionSession. Triggers - // JITDylib::clear(), which releases all tracked linker resources — in - // particular, invokes the ObjectLinkingLayer's ResourceManager which calls - // our memory manager's deallocate() for every FinalizedAlloc belonging to - // this dylib. Without this call JIT code pages accumulate in the arena - // until the whole session is destroyed. No Platform is registered in any - // of our LLJIT configurations (see orcjit_session.cc), so there is no - // Platform::teardownJITDylib callback to worry about. - session_->RemoveDylib(dylib_); + // Step 2: remove the JITDylib, releasing its JIT memory via the memory + // manager's deallocate() (see orcjit_session.cc — no Platform teardown here). + { + std::lock_guard lock(session_->mutex_); + session_->RemoveDylib(dylib_); + } dylib_ = nullptr; } void ORCJITDynamicLibraryObj::AddObjectFile(const String& path) { - // Read object file + // Read object file from disk into an owned MemoryBuffer. auto buffer_or_err = llvm::MemoryBuffer::getFile(path.c_str()); if (!buffer_or_err) { TVM_FFI_THROW(IOError) << "Failed to read object file: " << path; } + AddObjectBuffer(std::move(*buffer_or_err)); +} - // AddObjectFile is not thread-safe and must not race GetFunction. - TVM_FFI_ORCJIT_LLVM_CALL(jit_->addObjectFile(*dylib_, std::move(*buffer_or_err))); - context_symbol_refreshed_.store(false, std::memory_order_release); +void ORCJITDynamicLibraryObj::AddObjectBytes(const Bytes& bytes) { + // Copy the bytes into an owned MemoryBuffer: LLVM takes ownership and the + // source bytes may not outlive linking. + auto buffer = llvm::MemoryBuffer::getMemBufferCopy(llvm::StringRef(bytes.data(), bytes.size()), + name_.operator std::string()); + AddObjectBuffer(std::move(buffer)); } -void ORCJITDynamicLibraryObj::SetLinkOrder(const std::vector& dylibs) { - // Rebuild the link order: user-specified libraries first, then the LLJIT - // default link order (Main → Platform → ProcessSymbols). Preserving the - // default link order is essential — without ProcessSymbols, C++ objects - // that need host-process symbols (runtime, libtvm_ffi) would fail to link. - link_order_.clear(); +void ORCJITDynamicLibraryObj::AddObjectBuffer(std::unique_ptr buffer) { + // Leaf op (holds mutex_): addObjectFile only defines the object's + // materialization unit — no JIT code runs here. + std::lock_guard lock(session_->mutex_); + TVM_FFI_ORCJIT_LLVM_CALL(jit_->addObjectFile(*dylib_, std::move(buffer))); +} - for (auto* lib : dylibs) { - link_order_.emplace_back(lib, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); - } - for (auto& kv : jit_->defaultLinkOrder()) { - link_order_.emplace_back(kv.first, kv.second); +Module ORCJITExecutionSessionObj::LoadModule(const Array>& objects, + const String& name) { + // Hold no lock here: the callees each lock mutex_ at the leaf, and the fresh + // dylib is unpublished until this returns, so no other thread can race it. + ORCJITDynamicLibrary dylib = CreateDynamicLibrary(name); + ORCJITDynamicLibraryObj* self = dylib.get(); + + for (const Variant& object : objects) { + if (auto opt_path = object.as()) { + self->AddObjectFile(*opt_path); + } else { + self->AddObjectBytes(object.get()); + } } - // Set the link order in the LLVM JITDylib - dylib_->setLinkOrder(link_order_, false); + return self->Finalize(); +} + +Module ORCJITDynamicLibraryObj::Finalize() { + // Called once, from LoadModule, holding no session lock (its callees lock at + // the leaf). The guard prevents a second pass from appending duplicate imports + // (ImportModule does not dedup). + TVM_FFI_CHECK(!finalized_, InternalError) << "ORCJIT dynamic library already finalized"; + finalized_ = true; + + // Inject context symbols eagerly rather than deferring to first lookup. + InitContextSymbols(); + + // If the objects embed a library binary, reconstruct the import tree so the + // result behaves like a normally-loaded tvm-ffi module. Otherwise this is a + // plain JIT dylib module. + if (const char* library_bin = + reinterpret_cast(GetSymbol(symbol::tvm_ffi_library_bin))) { + return ProcessEmbeddedLibraryBin(library_bin, GetRef(this)); + } + return GetRef(this); } void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { - // Build search order: this dylib first, then all linked dylibs + // 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. llvm::orc::JITDylibSearchOrder search_order; search_order.emplace_back(dylib_, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); - // Append linked libraries - search_order.insert(search_order.end(), link_order_.begin(), link_order_.end()); - // Look up symbol using the full search order - auto symbol_or_err = - jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str())); + // Hold the lock across the lookup: materialization runs on this thread and, + // via InitFiniPlugin, populates the session's pending maps (lock-guarded). + // 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_); + symbol_or_err = + jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str())); + init = session_->DrainPendingInitializers(GetJITDylib()); + } - // Run pending initializers via InitFiniPlugin. RunPendingInitializers - // drains and erases the map entry; subsequent calls are no-ops until new - // object files add fresh entries — supports incremental loading. - // - // On macOS the scope publishes this dylib's __cxa_atexit records vector - // so the ___cxa_atexit shim (see orcjit_session.cc) can route dtor - // registrations here for our destructor to drain. Static init on C - // and C++ code typically happens here (first GetSymbol resolves the - // `main` entry point, which materializes and fires __mod_init_func). + // 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 - session_->RunPendingInitializers(GetJITDylib()); + ORCJITExecutionSessionObj::RunInitFiniEntries(init); if (!symbol_or_err) { llvm::Error remaining = @@ -161,9 +315,9 @@ void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { } void ORCJITDynamicLibraryObj::InitContextSymbols() { - std::lock_guard lock(context_symbol_refresh_mutex_); - if (context_symbol_refreshed_.load(std::memory_order_acquire)) return; - + // 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))) { *ctx_addr = this; } @@ -172,7 +326,6 @@ void ORCJITDynamicLibraryObj::InitContextSymbols() { *ctx_addr = symbol; } }); - context_symbol_refreshed_.store(true, std::memory_order_release); } llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { @@ -181,28 +334,12 @@ llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { } Optional ORCJITDynamicLibraryObj::GetFunction(const String& name) { - if (name == "orcjit.add_object_file") { - return Function::FromTyped([this](const String& path) { AddObjectFile(path); }); - } - if (name == "orcjit.set_link_order") { - return Function::FromTyped([this](const Array& libraries) { - std::vector libs; - libs.reserve(libraries.size()); - for (const ORCJITDynamicLibrary& lib : libraries) { - libs.push_back(&lib->GetJITDylib()); - } - SetLinkOrder(libs); - }); - } - - // TVM-FFI exports have __tvm_ffi_ prefix + // 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. + // + // TVM-FFI exports have the __tvm_ffi_ prefix. std::string symbol_name = symbol::tvm_ffi_symbol_prefix + std::string(name); - - if (!context_symbol_refreshed_.load(std::memory_order_acquire)) { - InitContextSymbols(); - } - - // Try to get the symbol - return NullOpt if not found if (void* symbol = GetSymbol(symbol_name)) { TVMFFISafeCallType c_func = reinterpret_cast(symbol); auto* wrapper = new DylibFnContextWithModule{GetRef(this)}; @@ -225,15 +362,16 @@ static void RegisterOrcJITFunctions() { refl::ObjectDef(); refl::GlobalDef() - .def("orcjit.ExecutionSession", + .def("tvm_ffi_orcjit.ExecutionSession", [](const std::string& orc_rt_path, int64_t slab_size_bytes) { return ORCJITExecutionSession(orc_rt_path, slab_size_bytes); }) - .def("orcjit.ExecutionSessionCreateDynamicLibrary", - [](const ORCJITExecutionSession& session, const String& name) -> Module { - return session->CreateDynamicLibrary(name); - }) - .def("orcjit.ExecutionSessionClearFreeSlabs", + .def("tvm_ffi_orcjit.GlobalDefaultSession", + []() { return ORCJITExecutionSessionObj::GlobalDefault(); }) + .def("tvm_ffi_orcjit.SessionLoadModule", + [](const ORCJITExecutionSession& session, const Array>& objects, + const String& name) -> Module { return session->LoadModule(objects, name); }) + .def("tvm_ffi_orcjit.SessionClearFreeSlabs", [](const ORCJITExecutionSession& session) -> int64_t { return session->ClearFreeSlabs(); }); diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h index 7221a9383..b458b1271 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h @@ -25,13 +25,13 @@ #define TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ #include +#include #include #include #include #include -#include -#include +#include #include "llvm_patches/macho_cxa_atexit_shim.h" #include "orcjit_session.h" @@ -44,6 +44,11 @@ class ORCJITExecutionSession; class ORCJITDynamicLibraryObj : public ModuleObj { public: + // The session is this dylib's factory and lifetime owner (create / teardown); + // it also drives the high-level load path (ORCJITExecutionSessionObj::LoadModule) + // through the private add and finalize helpers below. + friend class ORCJITExecutionSessionObj; + /*! * \brief Constructor * \param session The parent execution session @@ -68,15 +73,31 @@ class ORCJITDynamicLibraryObj : public ModuleObj { void AddObjectFile(const String& path); /*! - * \brief Set the link order for symbol resolution - * \param dylibs Vector of libraries to search for symbols (in order) + * \brief Add an in-memory object-file image to this library. + * \param bytes The object-file bytes. Copied into an owned MemoryBuffer, so + * the caller's bytes need not outlive linking. + */ + void AddObjectBytes(const Bytes& bytes); + + /*! + * \brief Add an object-file MemoryBuffer to this library (shared back end). + * \param buffer The object-file image. LLVM takes ownership. + */ + void AddObjectBuffer(std::unique_ptr buffer); + + /*! + * \brief Finalize the high-level load: inject context symbols and, if the + * objects embed a library binary, reconstruct the import tree. * - * When resolving symbols, this library will search in the specified libraries - * in the order provided. This replaces any previous link order. + * Must run exactly once per dylib (guarded by \c finalized_); a second call + * would append duplicate imports, since \c ModuleObj::ImportModule does not + * deduplicate. Only \c LoadModule calls this. + * + * \return The root module (this dylib when there is no embedded binary). */ - void SetLinkOrder(const std::vector& dylibs); + Module Finalize(); - /*! \brief Refresh context slots after an object add when needed. */ + /*! \brief Inject library-context symbols; runs once per dylib (see Finalize). */ void InitContextSymbols(); /*! @@ -92,12 +113,6 @@ class ORCJITDynamicLibraryObj : public ModuleObj { */ llvm::orc::JITDylib& GetJITDylib(); - /*! - * \brief Get the name of this library - * \return The library name - */ - String GetName() const { return name_; } - /*! \brief Parent execution session (for lifetime management) */ ORCJITExecutionSession session_; @@ -110,14 +125,8 @@ class ORCJITDynamicLibraryObj : public ModuleObj { /*! \brief Library name */ String name_; - /*! \brief Link order tracking (to support incremental linking) */ - llvm::orc::JITDylibSearchOrder link_order_; - - /*! \brief Whether context slots have been refreshed since the last object add. */ - std::atomic context_symbol_refreshed_{false}; - - /*! \brief Serializes the false-path context refresh between function lookups. */ - std::mutex context_symbol_refresh_mutex_; + /*! \brief Whether Finalize has run; guards against double-finalizing. */ + bool finalized_{false}; #ifdef __APPLE__ /*! \brief Per-dylib __cxa_atexit registry. diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc index 3470980d4..24e4cc0c3 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc @@ -31,10 +31,12 @@ #include #include #include +#include #include #include #include +#include #include "orcjit_dylib.h" #include "orcjit_memory_manager.h" @@ -203,9 +205,29 @@ ORCJITExecutionSession::ORCJITExecutionSession(const std::string& orc_rt_path, data_ = std::move(obj); } +ORCJITExecutionSession ORCJITExecutionSessionObj::GlobalDefault() { + // Leaked, never-destroyed: the owned LLJIT and slab arena must not be torn + // down during interpreter finalization, where teardown could call back into + // the host language. Resolve the ORC runtime path once, from the registered + // global tvm_ffi_orcjit.DefaultOrcRuntimePath (a Python hook that returns the + // bundled runtime); absent or empty means no platform. + static ORCJITExecutionSession* inst = [] { + std::string orc_rt_path; + if (auto fpath = Function::GetGlobal("tvm_ffi_orcjit.DefaultOrcRuntimePath")) { + orc_rt_path = (*fpath)().cast().operator std::string(); + } + return new ORCJITExecutionSession(orc_rt_path, 0); + }(); + return *inst; +} + ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary(const String& name) { TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; + // Compound topology op — serialize against concurrent create / add / lookup / + // teardown on this shared session (lock order: session lock first). + std::lock_guard lock(mutex_); + // Generate name if not provided String lib_name = name; if (lib_name.empty()) { @@ -250,31 +272,41 @@ llvm::orc::LLJIT& ORCJITExecutionSessionObj::GetLLJIT() { using CtorDtor = void (*)(); -void ORCJITExecutionSessionObj::RunPendingInitializers(llvm::orc::JITDylib& jit_dylib) { - auto it = pending_initializers_.find(&jit_dylib); - if (it != pending_initializers_.end()) { - llvm::sort(it->second, [](const InitFiniEntry& a, const InitFiniEntry& b) { +namespace { +// Remove a dylib's entries from `pending` and return them sorted into run order +// (by section, then priority). Caller holds the session lock. +std::vector DrainSorted( + std::unordered_map>& + pending, + llvm::orc::JITDylib& jit_dylib) { + std::vector entries; + auto it = pending.find(&jit_dylib); + if (it != pending.end()) { + entries = std::move(it->second); + pending.erase(it); + llvm::sort(entries, [](const ORCJITExecutionSessionObj::InitFiniEntry& a, + const ORCJITExecutionSessionObj::InitFiniEntry& b) { if (a.section != b.section) return static_cast(a.section) < static_cast(b.section); return a.priority < b.priority; }); - for (const auto& entry : it->second) { - entry.address.toPtr()(); - } - pending_initializers_.erase(it); } + return entries; } +} // namespace -void ORCJITExecutionSessionObj::RunPendingDeinitializers(llvm::orc::JITDylib& jit_dylib) { - auto it = pending_deinitializers_.find(&jit_dylib); - if (it != pending_deinitializers_.end()) { - llvm::sort(it->second, [](const InitFiniEntry& a, const InitFiniEntry& b) { - if (a.section != b.section) return static_cast(a.section) < static_cast(b.section); - return a.priority < b.priority; - }); - for (const auto& entry : it->second) { - entry.address.toPtr()(); - } - pending_deinitializers_.erase(it); +std::vector +ORCJITExecutionSessionObj::DrainPendingInitializers(llvm::orc::JITDylib& jit_dylib) { + return DrainSorted(pending_initializers_, jit_dylib); +} + +std::vector +ORCJITExecutionSessionObj::DrainPendingDeinitializers(llvm::orc::JITDylib& jit_dylib) { + return DrainSorted(pending_deinitializers_, jit_dylib); +} + +void ORCJITExecutionSessionObj::RunInitFiniEntries(const std::vector& entries) { + for (const auto& entry : entries) { + entry.address.toPtr()(); } } diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h index a35983a67..7e9b89989 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h @@ -27,11 +27,15 @@ #include #include #include +#include +#include +#include #include #include #include #include +#include #include #include @@ -41,8 +45,9 @@ namespace tvm { namespace ffi { namespace orcjit { -// Forward declaration +// Forward declarations class ORCJITDynamicLibrary; +class ORCJITExecutionSession; /*! * \brief ExecutionSession object for LLVM ORC JIT v2 @@ -58,6 +63,22 @@ class ORCJITExecutionSessionObj : public Object { explicit ORCJITExecutionSessionObj(const std::string& orc_rt_path = "", int64_t slab_size_bytes = 0); + /*! + * \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). + * + * The ORC runtime path is resolved on first call from the registered global + * \c tvm_ffi_orcjit.DefaultOrcRuntimePath (a Python-side hook that returns the + * bundled runtime, or empty to disable); absent or empty means no platform. + * + * \return The shared execution session. + */ + static ORCJITExecutionSession GlobalDefault(); + /*! * \brief Create a new DynamicLibrary (JITDylib) in this session * \param name Optional name for the library (for debugging) @@ -65,6 +86,26 @@ class ORCJITExecutionSessionObj : public Object { */ ORCJITDynamicLibrary CreateDynamicLibrary(const String& name); + /*! + * \brief Load a set of objects into one fresh dynamic library and return the + * resulting module, fully wired. + * + * End-to-end high-level entry: creates a JITDylib, adds every object (a + * \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. + * + * \param objects Array whose elements are each a \c String path or \c Bytes + * object-file image. + * \param name Optional JITDylib name (auto-generated when empty). + * \return The root module, with imports and library context fully wired (the + * dylib itself when there is no embedded library binary). + */ + Module LoadModule(const Array>& objects, const String& name); + /*! * \brief Get the underlying LLVM ExecutionSession * \return Reference to the LLVM ExecutionSession @@ -78,7 +119,8 @@ class ORCJITExecutionSessionObj : public Object { llvm::orc::LLJIT& GetLLJIT(); static constexpr bool _type_mutable = true; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("orcjit.ExecutionSession", ORCJITExecutionSessionObj, Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tvm_ffi_orcjit.ExecutionSession", ORCJITExecutionSessionObj, + Object); struct InitFiniEntry { enum class Section { @@ -92,9 +134,24 @@ class ORCJITExecutionSessionObj : public Object { int priority; }; - void RunPendingInitializers(llvm::orc::JITDylib& jit_dylib); - void RunPendingDeinitializers(llvm::orc::JITDylib& jit_dylib); + /*! + * \brief Remove a dylib's pending init/fini entries, sorted into run order. + * + * Separates collection from execution: the caller drains under \ref mutex_, + * then runs the entries with the lock released (see \ref mutex_). Caller must + * hold \ref mutex_. + * + * \param jit_dylib The dylib whose entries to drain. + * \return The entries in execution order (empty if none). + */ + std::vector DrainPendingInitializers(llvm::orc::JITDylib& jit_dylib); + std::vector DrainPendingDeinitializers(llvm::orc::JITDylib& jit_dylib); + + /*! \brief Run drained init/fini entries in order. Call with the lock released. */ + static void RunInitFiniEntries(const std::vector& entries); + // Called by InitFiniPlugin during materialization, which the triggering + // GetSymbol runs while already holding mutex_ — so these do not lock. void AddPendingInitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry); void AddPendingDeinitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry); @@ -103,8 +160,8 @@ class ORCJITExecutionSessionObj : public Object { * memory and dropping it from the session's dylib list. * * Invoked by \c ORCJITDynamicLibraryObj's destructor after any required - * static-destructor sequence (\c RunPendingDeinitializers on Linux/Windows, - * \c LLJIT::deinitialize on macOS) has completed. The caller must ensure no + * static-destructor sequence (drained via \c DrainPendingDeinitializers on + * Linux/Windows) has completed. The caller must ensure no * further use of the \c JITDylib* after this call — it becomes "Closed" and * its address may be reused by a subsequent \c createJITDylib. * @@ -124,6 +181,11 @@ class ORCJITExecutionSessionObj : public Object { int64_t ClearFreeSlabs(); private: + // The dylib serializes its compound operations (add object / lookup+init / + // teardown) on this session's mutex_, and drains its pending init/fini + // entries under it. See the locking discipline on mutex_. + friend class ORCJITDynamicLibraryObj; + /*! \brief Slab-pool memory manager — must be declared before jit_ for destruction order */ std::unique_ptr memory_manager_; /*! \brief The LLVM ORC JIT instance */ @@ -132,6 +194,20 @@ class ORCJITExecutionSessionObj : public Object { /*! \brief Counter for auto-generating library names */ std::atomic dylib_counter_{0}; + /*! + * \brief Serializes compound JITDylib operations on this shared session + * (create / add object / symbol lookup+init / teardown) and guards the + * pending init/fini maps below. + * + * Plain, non-recursive: acquired exactly once per operation, in the leaf op + * (CreateDynamicLibrary / AddObjectBuffer / GetSymbol / the dylib destructor); + * composers (LoadModule / Finalize / GetFunction) hold no lock. JIT'd + * constructors/destructors run after the lock is released (entries are drained + * under it), so a JIT'd ctor re-entering the session on the same thread does + * not deadlock. The resolved-call hot path never takes it. + */ + std::mutex mutex_; + std::unordered_map> pending_initializers_; std::unordered_map> pending_deinitializers_; }; diff --git a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt index 950a07300..92bca5e5a 100644 --- a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt +++ b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt @@ -84,6 +84,7 @@ if (NOT WIN32) add_test_object(sources/cc/test_link_order_base.cc) add_test_object(sources/cc/test_link_order_caller.cc) add_test_object(sources/cc/test_error.cc) + add_test_object(sources/cc/test_containers.cc) endif () # Pure C object files — built on all platforms (no C++ runtime deps) diff --git a/addons/tvm_ffi_orcjit/tests/README.md b/addons/tvm_ffi_orcjit/tests/README.md index d2207924b..beb9f51b7 100644 --- a/addons/tvm_ffi_orcjit/tests/README.md +++ b/addons/tvm_ffi_orcjit/tests/README.md @@ -106,5 +106,5 @@ static initialization and finalization: ## CI -The CI workflow (`.github/workflows/tvm_ffi_orcjit.yml`) runs `pytest` and the +The CI `orcjit` job (in `.github/workflows/ci_test.yml`) runs `pytest` and the quick-start example via `cibuildwheel`. diff --git a/addons/tvm_ffi_orcjit/tests/sources/cc/test_containers.cc b/addons/tvm_ffi_orcjit/tests/sources/cc/test_containers.cc new file mode 100644 index 000000000..541533cca --- /dev/null +++ b/addons/tvm_ffi_orcjit/tests/sources/cc/test_containers.cc @@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// One JIT-allocated FFI container per test. Every returned Object is freshly +// constructed inside this translation unit, so its header.deleter points into +// JIT-mapped memory — the escape hazard that keep_module_alive=True pins the +// JITDylib against. + +#include +#include +#include +#include +#include + +#include + +using tvm::ffi::Array; +using tvm::ffi::Map; +using tvm::ffi::String; +using tvm::ffi::Tuple; + +// Character replace in a String. +String test_string_replace_impl(String s, String from, String to) { + std::string out(s); + char c_from = std::string(from)[0]; + char c_to = std::string(to)[0]; + for (char& c : out) { + if (c == c_from) c = c_to; + } + return String(out); +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_string_replace, test_string_replace_impl); + +// Concatenate two Arrays. +Array test_array_concat_impl(Array a, Array b) { + Array out; + for (int64_t x : a) out.push_back(x); + for (int64_t x : b) out.push_back(x); + return out; +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_array_concat, test_array_concat_impl); + +// Hand-built Map. +Map test_map_build_impl() { + Map out; + out.Set(String("a"), 1); + out.Set(String("b"), 2); + return out; +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_map_build, test_map_build_impl); + +// Reverse a four-element Tuple. +Tuple test_tuple_reverse_impl( + Tuple t) { + return Tuple(t.get<3>(), t.get<2>(), t.get<1>(), t.get<0>()); +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_tuple_reverse, test_tuple_reverse_impl); diff --git a/addons/tvm_ffi_orcjit/tests/test_basic.py b/addons/tvm_ffi_orcjit/tests/test_basic.py index 17939497a..c7b9e822e 100644 --- a/addons/tvm_ffi_orcjit/tests/test_basic.py +++ b/addons/tvm_ffi_orcjit/tests/test_basic.py @@ -14,7 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Basic tests for tvm_ffi_orcjit functionality.""" +"""Basic tests for tvm_ffi_orcjit functionality. + +Everything loads through the high-level ``ExecutionSession.load_module`` API +(there is no low-level incremental library API). Each ``load_module`` links its +objects into one fresh JITDylib and returns a plain ``tvm_ffi.Module``; dropping +the module tears the dylib down. +""" from __future__ import annotations @@ -27,8 +33,8 @@ import pytest import tvm_ffi +from tvm_ffi import Module from tvm_ffi_orcjit import ExecutionSession -from tvm_ffi_orcjit.dylib import DynamicLibrary from utils import build_test_objects # --------------------------------------------------------------------------- @@ -57,16 +63,18 @@ def obj(name: str) -> str: return str(path) -def make_lib( - *obj_names: str, session: ExecutionSession | None = None, name: str = "" -) -> tuple[ExecutionSession, DynamicLibrary]: - """Create a library and load one or more object files into it.""" +def load( + *obj_names: str, + session: ExecutionSession | None = None, + name: str = "", + keep_module_alive: bool = False, +) -> Module: + """Load one or more object files into a fresh module on ``session``.""" if session is None: session = ExecutionSession() - lib = session.create_library(name) - for o in obj_names: - lib.add(obj(o)) - return session, lib + return session.load_module( + [obj(o) for o in obj_names], name, keep_module_alive=keep_module_alive + ) # --------------------------------------------------------------------------- @@ -102,14 +110,6 @@ def types_obj(self) -> str: """Return path prefix for test_types object.""" return f"{self.subdir}/test_types" - def link_order_base_obj(self) -> str: - """Return path prefix for test_link_order_base object.""" - return f"{self.subdir}/test_link_order_base" - - def link_order_caller_obj(self) -> str: - """Return path prefix for test_link_order_caller object.""" - return f"{self.subdir}/test_link_order_caller" - def error_obj(self) -> str: """Return path prefix for test_error object.""" return f"{self.subdir}/test_error" @@ -118,6 +118,16 @@ def ctor_dtor_obj(self) -> str: """Return path prefix for test_ctor_dtor object.""" return f"{self.subdir}/test_ctor_dtor" + def containers_obj(self) -> str: + """Return path prefix for test_containers object. + + C++-only source; C variants return an empty prefix so ``obj()`` sees + a missing path and skips the parametrized case for that variant. + """ + if not self.subdir.startswith("cc"): + return "" + return f"{self.subdir}/test_containers" + def fn(self, base_name: str) -> str: """Return the function name for a given base name.""" return base_name @@ -133,7 +143,6 @@ def _discover_variants() -> list[Variant]: _all_variants = _discover_variants() -_cpp_only = [v for v in _all_variants if v.subdir.startswith("cc")] def _variant_id(v: Variant) -> str: @@ -141,7 +150,7 @@ def _variant_id(v: Variant) -> str: # --------------------------------------------------------------------------- -# Session / library creation +# Session creation # --------------------------------------------------------------------------- @@ -151,22 +160,6 @@ def test_create_session() -> None: assert session is not None -def test_create_library() -> None: - """Create a DynamicLibrary from an ExecutionSession.""" - session = ExecutionSession() - lib = session.create_library() - assert lib is not None - - -def test_multiple_libraries() -> None: - """Create multiple named libraries in one session.""" - session = ExecutionSession() - lib1 = session.create_library("lib1") - lib2 = session.create_library("lib2") - assert lib1 is not None - assert lib2 is not None - - # --------------------------------------------------------------------------- # Load & execute — parametrized over C / C++ # --------------------------------------------------------------------------- @@ -175,67 +168,58 @@ def test_multiple_libraries() -> None: @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_load_and_execute(v: Variant) -> None: """Load test_funcs and verify add/multiply.""" - _, lib = make_lib(v.funcs_obj()) - assert lib.get_function(v.fn("test_add"))(10, 20) == 30 - assert lib.get_function(v.fn("test_multiply"))(7, 6) == 42 + mod = load(v.funcs_obj()) + assert mod.get_function(v.fn("test_add"))(10, 20) == 30 + assert mod.get_function(v.fn("test_multiply"))(7, 6) == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_load_and_execute_second_set(v: Variant) -> None: """Load test_funcs2 and verify subtract/divide.""" - _, lib = make_lib(v.funcs2_obj()) - assert lib.get_function(v.fn("test_subtract"))(10, 3) == 7 - assert lib.get_function(v.fn("test_divide"))(20, 4) == 5 + mod = load(v.funcs2_obj()) + assert mod.get_function(v.fn("test_subtract"))(10, 3) == 7 + assert mod.get_function(v.fn("test_divide"))(20, 4) == 5 + + +@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) +def test_multiple_objects_in_one_module(v: Variant) -> None: + """Multiple objects linked into one module expose all their functions.""" + mod = load(v.funcs_obj(), v.funcs2_obj()) + assert mod.get_function(v.fn("test_add"))(5, 3) == 8 + assert mod.get_function(v.fn("test_multiply"))(4, 5) == 20 + assert mod.get_function(v.fn("test_subtract"))(10, 3) == 7 + assert mod.get_function(v.fn("test_divide"))(20, 4) == 5 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_function_not_found(v: Variant) -> None: """Raise AttributeError for a missing function name.""" - _, lib = make_lib(v.funcs_obj()) + mod = load(v.funcs_obj()) with pytest.raises(AttributeError, match="Module has no function"): - lib.get_function("nonexistent_function") + mod.get_function("nonexistent_function") # --------------------------------------------------------------------------- -# Multi-object / multi-library — parametrized over C / C++ +# Multi-module — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_gradually_add_objects(v: Variant) -> None: - """Add multiple objects incrementally and verify all functions work.""" - _session, lib = make_lib(v.funcs_obj()) - - add_func = lib.get_function(v.fn("test_add")) - mul_func = lib.get_function(v.fn("test_multiply")) - assert add_func(5, 3) == 8 - assert mul_func(4, 5) == 20 - - lib.add(obj(v.funcs2_obj())) - assert lib.get_function(v.fn("test_subtract"))(10, 3) == 7 - assert lib.get_function(v.fn("test_divide"))(20, 4) == 5 - - # First object's functions still work - assert add_func(10, 20) == 30 - assert mul_func(7, 6) == 42 - - -@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_two_separate_libraries(v: Variant) -> None: - """Separate libraries expose only their own functions.""" +def test_two_separate_modules(v: Variant) -> None: + """Separate modules on one session expose only their own functions.""" session = ExecutionSession() - _, lib1 = make_lib(v.funcs_obj(), session=session, name="lib1") - _, lib2 = make_lib(v.funcs2_obj(), session=session, name="lib2") + mod1 = load(v.funcs_obj(), session=session, name="lib1") + mod2 = load(v.funcs2_obj(), session=session, name="lib2") - assert lib1.get_function(v.fn("test_add"))(5, 3) == 8 - assert lib1.get_function(v.fn("test_multiply"))(4, 5) == 20 - assert lib2.get_function(v.fn("test_subtract"))(10, 3) == 7 - assert lib2.get_function(v.fn("test_divide"))(20, 4) == 5 + assert mod1.get_function(v.fn("test_add"))(5, 3) == 8 + assert mod1.get_function(v.fn("test_multiply"))(4, 5) == 20 + assert mod2.get_function(v.fn("test_subtract"))(10, 3) == 7 + assert mod2.get_function(v.fn("test_divide"))(20, 4) == 5 with pytest.raises(AttributeError, match="Module has no function"): - lib1.get_function(v.fn("test_subtract")) + mod1.get_function(v.fn("test_subtract")) with pytest.raises(AttributeError, match="Module has no function"): - lib2.get_function(v.fn("test_add")) + mod2.get_function(v.fn("test_add")) # --------------------------------------------------------------------------- @@ -244,25 +228,23 @@ def test_two_separate_libraries(v: Variant) -> None: @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_symbol_conflict_same_library(v: Variant) -> None: - """Duplicate symbol in the same library raises an error.""" - _, lib = make_lib(v.funcs_obj()) - assert lib.get_function(v.fn("test_add"))(10, 20) == 30 +def test_symbol_conflict_same_module(v: Variant) -> None: + """Duplicate symbol across objects in one module raises an error.""" with pytest.raises(Exception): - lib.add(obj(v.conflict_obj())) + load(v.funcs_obj(), v.conflict_obj()) @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_symbol_conflict_different_libraries(v: Variant) -> None: - """Same symbol in different libraries resolves independently.""" +def test_symbol_conflict_different_modules(v: Variant) -> None: + """Same symbol in different modules resolves independently.""" session = ExecutionSession() - _, lib1 = make_lib(v.funcs_obj(), session=session, name="lib1") - _, lib2 = make_lib(v.conflict_obj(), session=session, name="lib2") + mod1 = load(v.funcs_obj(), session=session, name="lib1") + mod2 = load(v.conflict_obj(), session=session, name="lib2") - assert lib1.get_function(v.fn("test_add"))(10, 20) == 30 - assert lib2.get_function(v.fn("test_add"))(10, 20) == 1030 - assert lib1.get_function(v.fn("test_multiply"))(5, 6) == 30 - assert lib2.get_function(v.fn("test_multiply"))(5, 6) == 60 + assert mod1.get_function(v.fn("test_add"))(10, 20) == 30 + assert mod2.get_function(v.fn("test_add"))(10, 20) == 1030 + assert mod1.get_function(v.fn("test_multiply"))(5, 6) == 30 + assert mod2.get_function(v.fn("test_multiply"))(5, 6) == 60 # --------------------------------------------------------------------------- @@ -287,12 +269,12 @@ def _host_mul(a: int, b: int) -> int: @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_call_global(v: Variant) -> None: """JIT code calls back into Python-registered global functions.""" - _, lib = make_lib(v.call_global_obj()) - add_func = lib.get_function(v.fn("test_call_global_add")) + mod = load(v.call_global_obj()) + add_func = mod.get_function(v.fn("test_call_global_add")) assert add_func(10, 20) == 30 assert add_func(100, 200) == 300 - mul_func = lib.get_function(v.fn("test_call_global_mul")) + mul_func = mod.get_function(v.fn("test_call_global_mul")) assert mul_func(7, 6) == 42 assert mul_func(11, 11) == 121 @@ -302,136 +284,114 @@ def test_call_global(v: Variant) -> None: # --------------------------------------------------------------------------- -def test_context_injected_after_object_add() -> None: - """Concurrent first lookups after an object add observe the injected context pointer.""" - session = ExecutionSession() - lib = session.create_library("context_after_add") - lib.add(obj("c/test_context")) +def test_context_injected_at_load() -> None: + """Context is injected eagerly at load, so concurrent lookups all observe it.""" + mod = load("c/test_context") num_workers = 8 barrier = threading.Barrier(num_workers) def lookup(_worker: int) -> int: barrier.wait() - return lib.get_function("context_is_set")() + return mod.get_function("context_is_set")() with ThreadPoolExecutor(max_workers=num_workers) as pool: assert list(pool.map(lookup, range(num_workers))) == [1] * num_workers # --------------------------------------------------------------------------- -# Error handling — pure Python (Group 1) +# Error handling # --------------------------------------------------------------------------- -def test_empty_library() -> None: - """get_function on an empty library raises AttributeError.""" - session = ExecutionSession() - lib = session.create_library("empty") +def test_missing_function_on_empty_module() -> None: + """get_function for a missing name raises AttributeError.""" + mod = load(_all_variants[0].funcs_obj()) with pytest.raises(AttributeError, match="Module has no function"): - lib.get_function("nonexistent") + mod.get_function("nonexistent") def test_invalid_object_file_path() -> None: - """Adding a nonexistent object file raises an error.""" + """Loading a nonexistent object file raises an error.""" session = ExecutionSession() - lib = session.create_library("bad_path") with pytest.raises(Exception): - lib.add("/nonexistent_path/does_not_exist.o") + session.load_module("/nonexistent_path/does_not_exist.o") def test_invalid_object_file_content() -> None: - """Adding a file with garbage content raises an error.""" + """Loading a file with garbage content raises an error.""" with tempfile.NamedTemporaryFile(suffix=".o", delete=False) as f: f.write(b"this is not a valid object file") f.flush() session = ExecutionSession() - lib = session.create_library("bad_content") with pytest.raises(Exception): - lib.add(f.name) + session.load_module(f.name) def test_multiple_independent_sessions() -> None: """Two independent sessions don't interfere with each other.""" session1 = ExecutionSession() session2 = ExecutionSession() - _, lib1 = make_lib(_all_variants[0].funcs_obj(), session=session1) - _, lib2 = make_lib(_all_variants[0].funcs2_obj(), session=session2) + mod1 = load(_all_variants[0].funcs_obj(), session=session1) + mod2 = load(_all_variants[0].funcs2_obj(), session=session2) - assert lib1.get_function("test_add")(5, 3) == 8 - assert lib2.get_function("test_subtract")(10, 3) == 7 + assert mod1.get_function("test_add")(5, 3) == 8 + assert mod2.get_function("test_subtract")(10, 3) == 7 - # Each session's library doesn't see the other's functions + # Each session's module doesn't see the other's functions with pytest.raises(AttributeError, match="Module has no function"): - lib1.get_function("test_subtract") + mod1.get_function("test_subtract") with pytest.raises(AttributeError, match="Module has no function"): - lib2.get_function("test_add") + mod2.get_function("test_add") # --------------------------------------------------------------------------- -# Type variety — parametrized over C / C++ (Group 2) +# Type variety — parametrized over C / C++ # --------------------------------------------------------------------------- @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_zero_arg_function(v: Variant) -> None: """Zero-arg function returns constant 42.""" - _, lib = make_lib(v.types_obj()) - assert lib.get_function(v.fn("test_zero_arg"))() == 42 + mod = load(v.types_obj()) + assert mod.get_function(v.fn("test_zero_arg"))() == 42 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_four_arg_function(v: Variant) -> None: """Four integer arguments summed.""" - _, lib = make_lib(v.types_obj()) - assert lib.get_function(v.fn("test_four_args"))(1, 2, 3, 4) == 10 - assert lib.get_function(v.fn("test_four_args"))(100, 200, 300, 400) == 1000 + mod = load(v.types_obj()) + assert mod.get_function(v.fn("test_four_args"))(1, 2, 3, 4) == 10 + assert mod.get_function(v.fn("test_four_args"))(100, 200, 300, 400) == 1000 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_float_function(v: Variant) -> None: """Float multiply returns approximate result.""" - _, lib = make_lib(v.types_obj()) - result = lib.get_function(v.fn("test_float_multiply"))(3.14, 2.0) + mod = load(v.types_obj()) + result = mod.get_function(v.fn("test_float_multiply"))(3.14, 2.0) assert result == pytest.approx(6.28) @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_void_function(v: Variant) -> None: """Void function returns None.""" - _, lib = make_lib(v.types_obj()) - result = lib.get_function(v.fn("test_void_function"))() + mod = load(v.types_obj()) + result = mod.get_function(v.fn("test_void_function"))() assert result is None # --------------------------------------------------------------------------- -# Advanced — cross-library linking and error propagation (Group 3) +# Error propagation # --------------------------------------------------------------------------- -@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_set_link_order(v: Variant) -> None: - """Cross-library symbol resolution via set_link_order.""" - session = ExecutionSession() - # Base library exports helper_add - lib_base = session.create_library("base") - lib_base.add(obj(v.link_order_base_obj())) - # Caller library references helper_add from base - lib_caller = session.create_library("caller") - lib_caller.set_link_order(lib_base) - lib_caller.add(obj(v.link_order_caller_obj())) - - cross_add = lib_caller.get_function(v.fn("cross_lib_add")) - assert cross_add(10, 20) == 30 - assert cross_add(100, 200) == 300 - - @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_error_propagation(v: Variant) -> None: """JIT function that signals an error raises a Python exception.""" - _, lib = make_lib(v.error_obj()) + mod = load(v.error_obj()) with pytest.raises(Exception, match="test error"): - lib.get_function(v.fn("test_error"))() + mod.get_function(v.fn("test_error"))() # --------------------------------------------------------------------------- @@ -441,9 +401,9 @@ def test_error_propagation(v: Variant) -> None: def test_load_and_execute_cuda_function() -> None: """Load and execute CUDA-compiled test objects.""" - _, lib = make_lib("cuda/test_funcs") - assert lib.get_function("test_add")(10, 20) == 30 - assert lib.get_function("test_multiply")(7, 6) == 42 + mod = load("cuda/test_funcs") + assert mod.get_function("test_add")(10, 20) == 30 + assert mod.get_function("test_multiply")(7, 6) == 42 # --------------------------------------------------------------------------- @@ -461,9 +421,9 @@ def _append_ctor_log(x: str) -> None: nonlocal log log += x - _, lib = make_lib(v.ctor_dtor_obj()) - lib.get_function(v.fn("main"))() - del lib + mod = load(v.ctor_dtor_obj()) + mod.get_function(v.fn("main"))() + del mod main_idx = log.index("
") pre = log[:main_idx] @@ -502,9 +462,9 @@ def _append_ctor_log(x: str) -> None: # --------------------------------------------------------------------------- -# Dylib removal — dropping a DynamicLibrary while its session is still alive. +# Module drop — dropping a loaded Module while its session is still alive. # -# Dropping the last Python reference to a DynamicLibrary removes the JITDylib +# Dropping the last reference to a load_module result removes the JITDylib # from the ExecutionSession (releasing its JIT memory) in addition to running # any pending static destructors. These tests pin that contract and guard # against regressions in the slab-pool memory-manager refactor, which assumes @@ -513,78 +473,59 @@ def _append_ctor_log(x: str) -> None: @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_drop_empty_library(v: Variant) -> None: - """Drop an empty library, then reuse the session.""" +def test_drop_loaded_module_then_reload(v: Variant) -> None: + """Drop a module with live JIT code, then load and use a fresh one.""" session = ExecutionSession() - lib = session.create_library("empty") - del lib - # Session still works. - lib2 = session.create_library("after") - lib2.add(obj(v.funcs_obj())) - assert lib2.get_function(v.fn("test_add"))(1, 2) == 3 + mod1 = load(v.funcs_obj(), session=session, name="lib1") + assert mod1.get_function(v.fn("test_add"))(3, 4) == 7 + del mod1 -@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_drop_loaded_library_then_recreate(v: Variant) -> None: - """Drop a library with live JIT code, then create and use a fresh one.""" - session = ExecutionSession() - - lib1 = session.create_library("lib1") - lib1.add(obj(v.funcs_obj())) - assert lib1.get_function(v.fn("test_add"))(3, 4) == 7 - del lib1 - - lib2 = session.create_library("lib2") - lib2.add(obj(v.funcs_obj())) - assert lib2.get_function(v.fn("test_add"))(100, 1) == 101 + mod2 = load(v.funcs_obj(), session=session, name="lib2") + assert mod2.get_function(v.fn("test_add"))(100, 1) == 101 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_repeated_create_drop_many_iterations(v: Variant) -> None: - """Long create / load / call / drop cycle must not degrade or crash. +def test_repeated_load_drop_many_iterations(v: Variant) -> None: + """Long load / call / drop cycle must not degrade or crash. Exercises the memory manager's recycled-region path: each iteration frees JIT memory that a later iteration may be handed back by the arena. """ session = ExecutionSession() for i in range(32): - lib = session.create_library(f"iter_{i}") - lib.add(obj(v.funcs_obj())) - assert lib.get_function(v.fn("test_multiply"))(i + 1, 2) == (i + 1) * 2 - del lib + mod = load(v.funcs_obj(), session=session, name=f"iter_{i}") + assert mod.get_function(v.fn("test_multiply"))(i + 1, 2) == (i + 1) * 2 + del mod @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_drop_one_library_does_not_affect_another(v: Variant) -> None: - """Dropping one library must leave unrelated sibling libraries working.""" +def test_drop_one_module_does_not_affect_another(v: Variant) -> None: + """Dropping one module must leave unrelated sibling modules working.""" session = ExecutionSession() - keep = session.create_library("keep") - keep.add(obj(v.funcs_obj())) + keep = load(v.funcs_obj(), session=session, name="keep") - drop = session.create_library("drop") - drop.add(obj(v.funcs2_obj())) + drop = load(v.funcs2_obj(), session=session, name="drop") assert drop.get_function(v.fn("test_subtract"))(10, 3) == 7 del drop gc.collect() - # Untouched sibling still works; room now exists for a fresh library. + # Untouched sibling still works; room now exists for a fresh module. assert keep.get_function(v.fn("test_add"))(5, 5) == 10 - fresh = session.create_library("fresh") - fresh.add(obj(v.funcs2_obj())) + fresh = load(v.funcs2_obj(), session=session, name="fresh") assert fresh.get_function(v.fn("test_divide"))(20, 4) == 5 @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_captured_function_keeps_library_alive(v: Variant) -> None: - """A captured Function keeps the library alive after the lib handle is dropped.""" +def test_captured_function_keeps_module_alive(v: Variant) -> None: + """A captured Function keeps the module alive after the handle is dropped.""" session = ExecutionSession() - lib = session.create_library("hold") - lib.add(obj(v.funcs_obj())) + mod = load(v.funcs_obj(), session=session, name="hold") - captured = lib.get_function(v.fn("test_add")) - del lib + captured = mod.get_function(v.fn("test_add")) + del mod gc.collect() assert captured(11, 31) == 42 @@ -600,14 +541,13 @@ def _append(x: str) -> None: log.append(x) session = ExecutionSession() - lib = session.create_library("ctor_dtor") - lib.add(obj(v.ctor_dtor_obj())) - lib.get_function(v.fn("main"))() + mod = load(v.ctor_dtor_obj(), session=session, name="ctor_dtor") + mod.get_function(v.fn("main"))() ctor_len = len(log) assert "
" in log, f"main not observed: {log}" - del lib + del mod gc.collect() post = log[ctor_len:] @@ -623,53 +563,80 @@ def _append(x: str) -> None: @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_drop_caller_leaves_base_usable(v: Variant) -> None: - """Dropping a set_link_order caller leaves its base library usable.""" +def test_drop_all_modules_then_session(v: Variant) -> None: + """Dropping every module before the session still leaves clean teardown.""" session = ExecutionSession() + mods = [load(v.funcs_obj(), session=session, name=f"lib_{i}") for i in range(4)] + for mod in reversed(mods): + del mod + mods.clear() + gc.collect() + # Session destructor runs at end-of-test; a regression in removal would + # surface as a crash under pytest. + + +# --------------------------------------------------------------------------- +# FFI-container returns under keep_module_alive. +# +# JIT-emitted code allocates each returned Object through inline templates +# (make_object / make_inplace_array_object), so the object's header.deleter +# points into JIT-mapped memory. If the owning Module were torn down before +# the escaped Object, the deleter would dangle. keep_module_alive=True pins +# the module in the runtime's process-global registry, keeping the JITDylib +# mapped for the process lifetime. Each test loads with the flag, drops the +# local mod reference, forces GC, and then reads the returned container — +# proving the deleter is still reachable after the local mod is gone. +# --------------------------------------------------------------------------- - base = session.create_library("base") - base.add(obj(v.link_order_base_obj())) - caller = session.create_library("caller") - caller.set_link_order(base) - caller.add(obj(v.link_order_caller_obj())) - assert caller.get_function(v.fn("cross_lib_add"))(1, 2) == 3 +@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) +def test_containers_string(v: Variant) -> None: + """String return survives del mod when pinned.""" + mod = load(v.containers_obj(), keep_module_alive=True) + result = mod.test_string_replace("hello world", "l", "L") + del mod + gc.collect() + assert str(result) == "heLLo worLd" + - del caller +@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) +def test_containers_array(v: Variant) -> None: + """Array return survives del mod when pinned.""" + mod = load(v.containers_obj(), keep_module_alive=True) + result = mod.test_array_concat([1, 2, 3], [4, 5, 6]) + del mod gc.collect() + assert list(result) == [1, 2, 3, 4, 5, 6] - caller2 = session.create_library("caller2") - caller2.set_link_order(base) - caller2.add(obj(v.link_order_caller_obj())) - assert caller2.get_function(v.fn("cross_lib_add"))(10, 20) == 30 + +@pytest.mark.parametrize("v", _all_variants, ids=_variant_id) +def test_containers_map(v: Variant) -> None: + """Map return survives del mod when pinned.""" + mod = load(v.containers_obj(), keep_module_alive=True) + result = mod.test_map_build() + del mod + gc.collect() + assert dict(result) == {"a": 1, "b": 2} @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_drop_all_libraries_then_session(v: Variant) -> None: - """Dropping every library before the session still leaves clean teardown.""" - session = ExecutionSession() - libs = [] - for i in range(4): - lib = session.create_library(f"lib_{i}") - lib.add(obj(v.funcs_obj())) - libs.append(lib) - for lib in reversed(libs): - del lib - libs.clear() +def test_containers_tuple(v: Variant) -> None: + """Tuple<...> return survives del mod when pinned.""" + mod = load(v.containers_obj(), keep_module_alive=True) + result = mod.test_tuple_reverse((1, "a", 2, "b")) + del mod gc.collect() - # Session destructor runs at end-of-test; a regression in removal would - # surface as a crash under pytest. + assert tuple(result) == ("b", 2, "a", 1) # --------------------------------------------------------------------------- # Slab-pool growth (Stage B). # -# A session now 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; graphs larger than a single slab go to a dedicated -# oversize slab. These tests pin the behavioral contract: small -# slab_size sessions must still link many libraries (by growing) and -# must reuse drained bytes within a slab (via the free list). +# 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; +# graphs larger than a single slab go to a dedicated oversize slab. These tests +# pin the behavioral contract: small slab_size sessions must still link many +# modules (by growing) and must reuse drained bytes within a slab (free list). # --------------------------------------------------------------------------- @@ -683,20 +650,19 @@ def test_drop_all_libraries_then_session(v: Variant) -> None: @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) def test_pool_grows_under_small_slab(v: Variant) -> None: - """Tight slab_size forces the pool to grow as more libraries load.""" + """Tight slab_size forces the pool to grow as more modules load.""" session = ExecutionSession(slab_size=_SMALL_SLAB) - libs = [] + mods = [] for i in range(16): - lib = session.create_library(f"grow_{i}") - lib.add(obj(v.funcs_obj())) - assert lib.get_function(v.fn("test_add"))(i, i + 1) == 2 * i + 1 - libs.append(lib) + mod = load(v.funcs_obj(), session=session, name=f"grow_{i}") + assert mod.get_function(v.fn("test_add"))(i, i + 1) == 2 * i + 1 + mods.append(mod) - # All libraries remain callable after growth. Catches bugs where the - # pool routes deallocate to the wrong Slab via FA->owner. - for i, lib in enumerate(libs): - assert lib.get_function(v.fn("test_multiply"))(i, 2) == i * 2 + # All modules remain callable after growth. Catches bugs where the pool + # routes deallocate to the wrong Slab via FA->owner. + for i, mod in enumerate(mods): + assert mod.get_function(v.fn("test_multiply"))(i, 2) == i * 2 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @@ -710,48 +676,43 @@ def test_small_slab_recycles_after_drop(v: Variant) -> None: """ session = ExecutionSession(slab_size=_SMALL_SLAB) for i in range(32): - lib = session.create_library(f"recycle_{i}") - lib.add(obj(v.funcs_obj())) - assert lib.get_function(v.fn("test_add"))(i, 1) == i + 1 - del lib + mod = load(v.funcs_obj(), session=session, name=f"recycle_{i}") + assert mod.get_function(v.fn("test_add"))(i, 1) == i + 1 + del mod gc.collect() @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") @pytest.mark.parametrize("v", _all_variants, ids=_variant_id) -def test_pool_survives_mixed_load_drop_create(v: Variant) -> None: - """Interleaved load / drop / create exercises growth + free-list together.""" +def test_pool_survives_mixed_load_drop(v: Variant) -> None: + """Interleaved load / drop exercises growth + free-list together.""" session = ExecutionSession(slab_size=_SMALL_SLAB) - # Ramp up to 4 concurrent libraries. - live = [session.create_library(f"base_{i}") for i in range(4)] - for lib in live: - lib.add(obj(v.funcs_obj())) + # Ramp up to 4 concurrent modules. + live: list[Module | None] = [ + load(v.funcs_obj(), session=session, name=f"base_{i}") for i in range(4) + ] - # Drop two, add three, verify remaining two still work, verify new - # three work. Mixing drop/create in one session exercises slab - # growth alongside free-list reuse from the dropped bytes. - live[0] = None # type: ignore[assignment] - live[2] = None # type: ignore[assignment] + # Drop two, load three, verify remaining two still work, verify new three + # work. Mixing drop/load in one session exercises slab growth alongside + # free-list reuse from the dropped bytes. + live[0] = None + live[2] = None gc.collect() - new_libs = [] - for i in range(3): - lib = session.create_library(f"after_drop_{i}") - lib.add(obj(v.funcs2_obj())) - new_libs.append(lib) + new_mods = [load(v.funcs2_obj(), session=session, name=f"after_drop_{i}") for i in range(3)] assert live[1].get_function(v.fn("test_add"))(10, 5) == 15 assert live[3].get_function(v.fn("test_multiply"))(7, 6) == 42 - for i, lib in enumerate(new_libs): - assert lib.get_function(v.fn("test_subtract"))(i + 10, i) == 10 + for i, mod in enumerate(new_mods): + assert mod.get_function(v.fn("test_subtract"))(i + 10, i) == 10 # --------------------------------------------------------------------------- # Manual slab reclamation via session.clear_free_slabs() (Stage C). # -# After dropping a batch of libraries, the user can call clear_free_slabs() -# to release any drained Slabs (zero live JIT allocations) back to the OS. +# After dropping a batch of modules, the user can call clear_free_slabs() to +# release any drained Slabs (zero live JIT allocations) back to the OS. # --------------------------------------------------------------------------- @@ -810,16 +771,15 @@ def test_clear_free_slabs_no_drained() -> None: session = ExecutionSession() # Fresh session — the initial slab was never used, so not reclaimable. assert session.clear_free_slabs() == 0 - # Load + keep a library: live allocations pin the slab. - lib = session.create_library("alive") - lib.add(obj(_all_variants[0].funcs_obj())) - lib.get_function("test_add")(1, 2) + # Load + keep a module: live allocations pin the slab. + mod = load(_all_variants[0].funcs_obj(), session=session) + mod.get_function("test_add")(1, 2) assert session.clear_free_slabs() == 0 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_reclaims_oversize(tmp_path: Path) -> None: - """Dropping an oversize-path library lets clear_free_slabs reclaim its slab. + """Dropping an oversize-path module lets clear_free_slabs reclaim its slab. Uses a 3 MB rodata blob that exceeds the 4 MB slab's usable-per-slab threshold (~2 MB after midpoint slack), forcing the oversize path. The @@ -829,12 +789,11 @@ def test_clear_free_slabs_reclaims_oversize(tmp_path: Path) -> None: big_obj = _build_big_object(tmp_path, 3 * 1024 * 1024) # 4 MB slab → usable = slab_size / 2 = 2 MB → 3 MB blob forces oversize. session = ExecutionSession(slab_size=4 * 1024 * 1024) - lib = session.create_library("oversize") - lib.add(big_obj) - assert lib.get_function("big_probe")() == 3 * 1024 * 1024 - del lib + mod = session.load_module(big_obj) + assert mod.get_function("big_probe")() == 3 * 1024 * 1024 + del mod gc.collect() - # The oversize slab holds only the now-dropped lib's graph, so + # The oversize slab holds only the now-dropped module's graph, so # clearFreeSlabs reclaims it. The initial (never-used) slab stays. reclaimed = session.clear_free_slabs() assert reclaimed >= 1, f"expected to reclaim the drained oversize slab, got {reclaimed}" @@ -845,10 +804,9 @@ def test_clear_free_slabs_idempotent(tmp_path: Path) -> None: """A second call after everything has been reclaimed returns 0.""" big_obj = _build_big_object(tmp_path, 3 * 1024 * 1024) session = ExecutionSession(slab_size=4 * 1024 * 1024) - lib = session.create_library("x") - lib.add(big_obj) - lib.get_function("big_probe")() - del lib + mod = session.load_module(big_obj) + mod.get_function("big_probe")() + del mod gc.collect() assert session.clear_free_slabs() >= 1 @@ -857,36 +815,34 @@ def test_clear_free_slabs_idempotent(tmp_path: Path) -> None: @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") def test_clear_free_slabs_preserves_live_pool(tmp_path: Path) -> None: - """Reclaim runs only on drained slabs; live libraries keep working. + """Reclaim runs only on drained slabs; live modules keep working. - Uses two different oversize blob sizes so the two libs land on + Uses two different oversize blob sizes so the two modules land on separately-sized slabs: the 3 MB blob forces the pool to grow to 8 MB (next power of two covering a 4 MB pool's non-exec budget), - while the 5 MB blob forces 16 MB. Dropping the 3 MB lib then + while the 5 MB blob forces 16 MB. Dropping the 3 MB module then leaves its 8 MB slab drained while the 16 MB slab stays live. """ small_obj = _build_big_object(tmp_path, 3 * 1024 * 1024, name="small_obj") large_obj = _build_big_object(tmp_path, 5 * 1024 * 1024, name="large_obj") session = ExecutionSession(slab_size=4 * 1024 * 1024) - # Drop lib — lands on the 8 MB grown slab. - drop_lib = session.create_library("drop") - drop_lib.add(small_obj) - drop_lib.get_function("big_probe")() - del drop_lib + # Drop module — lands on the 8 MB grown slab. + drop_mod = session.load_module(small_obj) + drop_mod.get_function("big_probe")() + del drop_mod - # Keep lib — doesn't fit the 8 MB slab (5 MB > its ~4 MB non-exec + # Keep module — doesn't fit the 8 MB slab (5 MB > its ~4 MB non-exec # budget), so the pool grows to a 16 MB slab for this one. - keep_lib = session.create_library("keep") - keep_lib.add(large_obj) - assert keep_lib.get_function("big_probe")() == 5 * 1024 * 1024 + keep_mod = session.load_module(large_obj) + assert keep_mod.get_function("big_probe")() == 5 * 1024 * 1024 gc.collect() reclaimed = session.clear_free_slabs() assert reclaimed >= 1, f"expected drained slab to be reclaimed, got {reclaimed}" - # The kept lib's slab was untouched; it still executes correctly. - assert keep_lib.get_function("big_probe")() == 5 * 1024 * 1024 + # The kept module's slab was untouched; it still executes correctly. + assert keep_mod.get_function("big_probe")() == 5 * 1024 * 1024 @pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only") diff --git a/addons/tvm_ffi_orcjit/tests/test_memory_manager.py b/addons/tvm_ffi_orcjit/tests/test_memory_manager.py index b5cb1ba60..046059b46 100644 --- a/addons/tvm_ffi_orcjit/tests/test_memory_manager.py +++ b/addons/tvm_ffi_orcjit/tests/test_memory_manager.py @@ -289,9 +289,8 @@ def test_slab_colocation(variant: str) -> None: maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=_SLAB_SIZE) - lib1 = session.create_library("lib1") - lib1.add(obj(f"{variant}/test_addr")) - addr1 = lib1.get_function("code_address")() + mod1 = session.load_module(obj(f"{variant}/test_addr"), "lib1") + addr1 = mod1.get_function("code_address")() # Find where LLVM placed the first allocation and block nearby VA maps_after = _parse_maps() @@ -300,9 +299,8 @@ def test_slab_colocation(variant: str) -> None: blockers = block_nearby_va(jit_center) try: - lib2 = session.create_library("lib2") - lib2.add(obj(f"{variant}/test_addr")) - addr2 = lib2.get_function("code_address")() + mod2 = session.load_module(obj(f"{variant}/test_addr"), "lib2") + addr2 = mod2.get_function("code_address")() finally: free_blockers(blockers) @@ -330,9 +328,8 @@ def _measure_distance_under_pressure( maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=slab_size) - lib1 = session.create_library("lib1") - lib1.add(obj(f"{variant}/test_addr")) - addr1 = lib1.get_function("code_address")() + mod1 = session.load_module(obj(f"{variant}/test_addr"), "lib1") + addr1 = mod1.get_function("code_address")() maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) @@ -340,10 +337,9 @@ def _measure_distance_under_pressure( blockers = block_nearby_va(jit_center, radius=radius) try: - lib2 = session.create_library("lib2") try: - lib2.add(obj(f"{variant}/test_addr")) - addr2 = lib2.get_function("code_address")() + mod2 = session.load_module(obj(f"{variant}/test_addr"), "lib2") + addr2 = mod2.get_function("code_address")() except Exception: return None, True finally: @@ -413,29 +409,32 @@ def test_slab_hidden_symbol_with_blocker(variant: str) -> None: """Slab prevents hidden-visibility relocation overflow under VA pressure. Loads two objects with hidden-visibility cross-references (ADRP+ADD - on AArch64, PC32 on x86_64) with a VA blocker between them. - Without slab, the blocker would push objects apart causing overflow. - With the slab, both objects are co-located and the call succeeds. + on AArch64, PC32 on x86_64) into one module under VA pressure. A VA + blocker is placed before the load; without slab the two objects could + scatter and overflow, with the slab they are co-located and the call + succeeds. """ maps_before = set(_parse_maps()) + # Warm the session so a slab is reserved, then block nearby VA before the + # real load to force any fresh scatter to land far away. session = ExecutionSession(slab_size=_SLAB_SIZE) - lib = session.create_library("hidden_test") + warmup = session.load_module(obj(f"{variant}/test_hidden_helper"), "warmup") + assert warmup.get_function("hidden_add")(1, 2) == 3 - # Load helper and force materialization - lib.add(obj(f"{variant}/test_hidden_helper")) - assert lib.get_function("hidden_add")(1, 2) == 3 - - # Block nearby VA to force scatter maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) jit_center = max(s for s, e in new_maps) if new_maps else 0xFFFF00000000 blockers = block_nearby_va(jit_center) try: - lib.add(obj(f"{variant}/test_hidden_caller")) - fn = lib.get_function("call_hidden_add") - assert fn(10, 20) == 30 + # Helper + caller cross-reference via hidden visibility, so they must + # live in one module (one JITDylib) to resolve. + mod = session.load_module( + [obj(f"{variant}/test_hidden_helper"), obj(f"{variant}/test_hidden_caller")], + "hidden_test", + ) + assert mod.get_function("call_hidden_add")(10, 20) == 30 finally: free_blockers(blockers) @@ -455,9 +454,8 @@ def test_large_data_section(variant: str) -> None: the function works. The 4MB section fits in the slab. """ session = ExecutionSession() - lib = session.create_library("fatbin") - lib.add(obj(f"{variant}/fake_fatbin")) - fn = lib.get_function("get_fatbin_size") + mod = session.load_module(obj(f"{variant}/fake_fatbin"), "fatbin") + fn = mod.get_function("get_fatbin_size") assert fn() == 4 * 1024 * 1024 @@ -481,14 +479,13 @@ def test_overflow_section_outside_slab(variant: str) -> None: the slab region. """ session = ExecutionSession(slab_size=_SLAB_SIZE) - lib = session.create_library("fatbin_overflow") - lib.add(obj(f"{variant}/fake_fatbin")) + mod = session.load_module(obj(f"{variant}/fake_fatbin"), "fatbin_overflow") # Verify the function still works correctly. - assert lib.get_function("get_fatbin_size")() == 4 * 1024 * 1024 + assert mod.get_function("get_fatbin_size")() == 4 * 1024 * 1024 # Get the actual address of the fatbin data in memory. - fatbin_addr = lib.get_function("get_fatbin_addr")() + fatbin_addr = mod.get_function("get_fatbin_addr")() # Find the slab mapping: a single large region matching the slab size. # The slab is reserved as PROT_NONE and then committed in slabs, so @@ -540,28 +537,26 @@ def test_dso_handle_relocation_after_failed_materialization(variant: str) -> Non apart. With slab: PASSES (all allocations in same slab). """ - # Step 1: Trigger leaked materializations to consume low VA space. + # Step 1: Trigger leaked materializations to consume low VA space. Each + # failed load (duplicate symbol across the two objects) leaks its slab. leaked_sessions = [] for _ in range(3): s0 = ExecutionSession() - lib0 = s0.create_library("warmup") - lib0.add(obj(f"{variant}/test_funcs")) - lib0.get_function("test_add")(10, 20) + warm = s0.load_module(obj(f"{variant}/test_funcs"), "warmup") + warm.get_function("test_add")(10, 20) try: - lib0.add(obj(f"{variant}/test_funcs_conflict")) + s0.load_module([obj(f"{variant}/test_funcs"), obj(f"{variant}/test_funcs_conflict")]) except Exception: pass - leaked_sessions.append((s0, lib0)) + leaked_sessions.append((s0, warm)) - # Step 2: Fresh session — cross-library resolution must still work. + # Step 2: Fresh session — resolution must still work after the leaks. session = ExecutionSession() - lib1 = session.create_library("lib1") - lib1.add(obj(f"{variant}/test_funcs")) - assert lib1.get_function("test_add")(10, 20) == 30 + mod1 = session.load_module(obj(f"{variant}/test_funcs"), "lib1") + assert mod1.get_function("test_add")(10, 20) == 30 - lib2 = session.create_library("lib2") - lib2.add(obj(f"{variant}/test_funcs_conflict")) - assert lib2.get_function("test_add")(10, 20) == 1030 + mod2 = session.load_module(obj(f"{variant}/test_funcs_conflict"), "lib2") + assert mod2.get_function("test_add")(10, 20) == 1030 # --------------------------------------------------------------------------- @@ -616,9 +611,8 @@ def test_dso_handle_delta32_with_slab(variant: str) -> None: maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=_SLAB_SIZE) - lib1 = session.create_library("lib1") - lib1.add(obj(f"{variant}/test_funcs")) - assert lib1.get_function("test_add")(10, 20) == 30 + mod1 = session.load_module(obj(f"{variant}/test_funcs"), "lib1") + assert mod1.get_function("test_add")(10, 20) == 30 # Block 3GB of VA around the first allocation to force scatter maps_after = _parse_maps() @@ -627,9 +621,8 @@ def test_dso_handle_delta32_with_slab(variant: str) -> None: blockers = block_nearby_va(jit_center, radius=_DSO_BLOCK_RADIUS) try: - lib2 = session.create_library("lib2") - lib2.add(obj(f"{variant}/test_funcs_conflict")) - assert lib2.get_function("test_add")(10, 20) == 1030 + mod2 = session.load_module(obj(f"{variant}/test_funcs_conflict"), "lib2") + assert mod2.get_function("test_add")(10, 20) == 1030 finally: free_blockers(blockers) @@ -662,9 +655,8 @@ def test_dso_handle_delta32_overflow_without_slab(variant: str) -> None: maps_before = set(_parse_maps()) session = ExecutionSession(slab_size=-1) # slab disabled - lib1 = session.create_library("lib1") - lib1.add(obj(f"{variant}/test_addr")) - lib1.get_function("code_address")() + mod1 = session.load_module(obj(f"{variant}/test_addr"), "lib1") + mod1.get_function("code_address")() maps_after = _parse_maps() new_maps = _find_new_mappings(maps_before, maps_after) @@ -672,10 +664,9 @@ def test_dso_handle_delta32_overflow_without_slab(variant: str) -> None: blockers = block_nearby_va(jit_center, radius=_DSO_BLOCK_RADIUS) try: - lib2 = session.create_library("lib2") try: - lib2.add(obj(f"{variant}/test_funcs_conflict")) - result = lib2.get_function("test_add")(10, 20) + mod2 = session.load_module(obj(f"{variant}/test_funcs_conflict"), "lib2") + result = mod2.get_function("test_add")(10, 20) # If we get here, GCC used GOTPCRELX — no overflow. assert result == 1030 except Exception: diff --git a/addons/tvm_ffi_orcjit/tests/test_session_load_module.py b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py new file mode 100644 index 000000000..461e3340d --- /dev/null +++ b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py @@ -0,0 +1,437 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for the high-level session ``load_module`` API and shared session. + +These cover the refactor's new surface: + +- ``default_session()`` — the process-wide shared execution session. +- ``ExecutionSession.load_module`` — unified path/bytes input, single item or + list, returning a plain ``tvm_ffi.Module``. +- Concurrent create / load / call / drop on the shared session. + +They deliberately use pure-C objects: the C ABI path is uniform across +platforms, whereas some C++ object layouts hit environment-specific JITLink +limitations unrelated to this API. +""" + +from __future__ import annotations + +import gc +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest +import tvm_ffi +from tvm_ffi_orcjit import ExecutionSession, default_session +from utils import build_test_objects + +OBJ_DIR = build_test_objects() + + +def obj(name: str) -> str: + """Return path to a pre-built test object file, or skip if missing.""" + path = OBJ_DIR / f"{name}.o" + if not path.exists(): + pytest.skip(f"{path.name} not found (not built)") + return str(path) + + +# --------------------------------------------------------------------------- +# default_session() +# --------------------------------------------------------------------------- + + +def test_default_session_is_cached() -> None: + """default_session() returns the same shared instance every time.""" + assert default_session() is default_session() + + +def test_default_session_is_execution_session() -> None: + """default_session() returns an ExecutionSession.""" + assert isinstance(default_session(), ExecutionSession) + + +def test_default_session_load_and_call() -> None: + """A module loaded on the shared session is callable.""" + mod = default_session().load_module(obj("c/test_funcs")) + assert mod.test_add(10, 20) == 30 + assert mod.test_multiply(7, 6) == 42 + + +def test_default_session_concurrent_first_call() -> None: + """Many threads racing the first default_session() all get one shared instance. + + Guards the double-checked-locking init: the underlying FFI call releases the + GIL, so without the lock concurrent first callers could build separate + sessions. + """ + num_workers = 16 + barrier = threading.Barrier(num_workers) + + def get(_worker: int) -> ExecutionSession: + barrier.wait() + return default_session() + + with ThreadPoolExecutor(max_workers=num_workers) as pool: + sessions = list(pool.map(get, range(num_workers))) + assert all(s is sessions[0] for s in sessions) + + +def test_default_orc_runtime_path_hook_registered() -> None: + """The C++ singleton's ORC-path hook is registered and returns a string.""" + path = tvm_ffi.get_global_func("tvm_ffi_orcjit.DefaultOrcRuntimePath")() + assert isinstance(path, str) # bundled runtime path, or "" when none/disabled + + +# --------------------------------------------------------------------------- +# load_module — input shapes +# --------------------------------------------------------------------------- + + +def test_load_module_single_path() -> None: + """A single path (not wrapped in a list) loads correctly.""" + mod = default_session().load_module(obj("c/test_funcs")) + assert isinstance(mod, tvm_ffi.Module) + assert mod.test_add(1, 2) == 3 + + +def test_load_module_single_bytes() -> None: + """A single in-memory object image loads correctly.""" + blob = Path(obj("c/test_funcs")).read_bytes() + mod = default_session().load_module(blob) + assert mod.test_add(4, 5) == 9 + + +def test_load_module_list_of_paths() -> None: + """A list of object paths is linked into one module.""" + mod = default_session().load_module([obj("c/test_funcs"), obj("c/test_funcs2")]) + assert mod.test_add(10, 20) == 30 + assert mod.test_subtract(10, 3) == 7 + assert mod.test_divide(20, 4) == 5 + + +def test_load_module_mixed_path_and_bytes() -> None: + """A list mixing a path and in-memory bytes links into one module.""" + blob = Path(obj("c/test_funcs2")).read_bytes() + mod = default_session().load_module([obj("c/test_funcs"), blob]) + assert mod.test_multiply(7, 6) == 42 + assert mod.test_subtract(9, 4) == 5 + + +def test_load_module_returns_plain_module_kind() -> None: + """The returned module reports the orcjit kind.""" + mod = default_session().load_module(obj("c/test_funcs")) + assert mod.kind == "orcjit" + + +def test_load_module_rejects_bad_element_type() -> None: + """A non-path, non-bytes element raises a clear TypeError.""" + with pytest.raises(TypeError, match=r"path .* or object-file bytes"): + default_session().load_module([12345]) + + +def test_load_module_empty_list() -> None: + """An empty object list yields a valid, empty module (no functions).""" + mod = default_session().load_module([]) + assert isinstance(mod, tvm_ffi.Module) + with pytest.raises(AttributeError, match="Module has no function"): + mod.get_function("anything") + + +def test_load_module_named() -> None: + """An explicit name is accepted.""" + mod = default_session().load_module(obj("c/test_funcs"), name="my_named_lib") + assert mod.test_add(2, 2) == 4 + + +def test_load_module_context_symbol_injected_eagerly() -> None: + """Eager load wires the library-context pointer before any get_function.""" + mod = default_session().load_module(obj("c/test_context")) + # The context slot is populated at load time, so the first (and only) + # lookup already observes a non-null __tvm_ffi__library_ctx. + assert mod.context_is_set() == 1 + + +def test_load_module_drop_frees() -> None: + """Dropping every reference to a loaded module tears down its dylib.""" + session = ExecutionSession() # isolated session so teardown is observable + mod = session.load_module(obj("c/test_funcs")) + assert mod.test_add(1, 1) == 2 + del mod + gc.collect() + # Session still usable after the drop. + mod2 = session.load_module(obj("c/test_funcs2")) + assert mod2.test_subtract(5, 2) == 3 + + +# --------------------------------------------------------------------------- +# Embedded library binary — the __tvm_ffi__library_bin path (§3.3). +# +# When a loaded object embeds a serialized import tree, load_module must +# deserialize each embedded module by kind (via ffi.Module.load_from_bytes.*), +# wire the imports onto the "_lib" placeholder (the JIT dylib itself), and +# return the root module — so cross-module lookups resolve through imports. +# --------------------------------------------------------------------------- + + +def _u64(value: int) -> bytes: + return int(value).to_bytes(8, "little") + + +def _blob_str(data: bytes) -> bytes: + return _u64(len(data)) + data + + +def _u64_vec(values: list[int]) -> bytes: + return _u64(len(values)) + b"".join(_u64(v) for v in values) + + +def _make_library_bin( + modules: list[tuple[str, bytes]], indptr: list[int], children: list[int] +) -> bytes: + """Serialize a library-bin blob matching the core ProcessLibraryBin format.""" + stream = _u64_vec(indptr) + _u64_vec(children) + for kind, payload in modules: + stream += _blob_str(kind.encode()) + if kind != "_lib": + stream += _blob_str(payload) + return _u64(len(stream)) + stream + + +def _raw_library_bin(stream: bytes) -> bytes: + """Wrap a hand-crafted (possibly malformed) inner stream with its length header.""" + return _u64(len(stream)) + stream + + +def _build_library_bin_object(tmp_path: Path, blob: bytes, extra_export: str) -> str: + """Compile a C object exposing __tvm_ffi__library_bin plus one function.""" + try: + import tvm_ffi.cpp # noqa: PLC0415 + except ImportError: + pytest.skip("tvm_ffi.cpp not available for on-the-fly object build") + + array_body = ", ".join(str(b) for b in blob) + src = tmp_path / "libbin.c" + src.write_text( + f""" + #include + #include + TVM_FFI_DLL_EXPORT unsigned char __tvm_ffi__library_bin[] = {{ {array_body} }}; + TVM_FFI_DLL_EXPORT int {extra_export}(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 = 777; + return 0; + }} + """ + ) + try: + return tvm_ffi.cpp.build( + name="libbin", + sources=[str(src)], + output="libbin.o", + extra_cflags=["-O2"], + build_directory=str(tmp_path / ".build_libbin"), + ) + except Exception as exc: + pytest.skip(f"could not build library-bin object: {exc}") + + +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. + @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")) + + # 2-module import tree: module 0 = "_lib" (the JIT dylib), module 1 = the + # custom probe module. Module 0 imports module 1. + # indptr[i]..indptr[i+1] indexes children[]: module 0 -> child_indices[0:1] = [1] + blob = _make_library_bin( + modules=[("_lib", b""), ("orcjit_test_probe", b"")], + indptr=[0, 1, 1], + children=[1], + ) + libbin_obj = _build_library_bin_object(tmp_path, blob, extra_export="__tvm_ffi_probe_marker") + + # Load the JIT dylib's own function object alongside the library-bin object. + root = default_session().load_module([obj("c/test_funcs"), libbin_obj]) + + # The dylib's own exported function resolves directly. + assert root.get_function("test_add")(10, 20) == 30 + assert root.get_function("probe_marker")() == 777 + + # The embedded module's function resolves through the wired import tree. + assert root.get_function("test_subtract", query_imports=True)(10, 3) == 7 + assert len(root.imports) == 1 + + +# --------------------------------------------------------------------------- +# 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. +# --------------------------------------------------------------------------- + + +def test_shared_session_concurrent_load_call_drop() -> None: + """Many threads load + call + drop on the shared session simultaneously.""" + session = default_session() + num_workers = 8 + barrier = threading.Barrier(num_workers) + + def worker(i: int) -> int: + barrier.wait() + total = 0 + for _ in range(4): + mod = session.load_module(obj("c/test_funcs"), name=f"conc_{i}") + total += mod.test_add(i, i) + total += mod.test_multiply(i, 2) + del mod + return total + + with ThreadPoolExecutor(max_workers=num_workers) as pool: + results = list(pool.map(worker, range(num_workers))) + + # Each worker: 4 * (i+i + i*2) = 4 * 4i = 16i + assert results == [16 * i for i in range(num_workers)] + + +def test_shared_session_concurrent_distinct_modules() -> None: + """Concurrent loads of different objects resolve independently.""" + session = default_session() + num_workers = 6 + barrier = threading.Barrier(num_workers) + + def worker(i: int) -> tuple[int, int]: + barrier.wait() + m_funcs = session.load_module(obj("c/test_funcs")) + m_funcs2 = session.load_module(obj("c/test_funcs2")) + return m_funcs.test_add(i, 1), m_funcs2.test_subtract(i + 5, i) + + with ThreadPoolExecutor(max_workers=num_workers) as pool: + results = list(pool.map(worker, range(num_workers))) + + assert results == [(i + 1, 5) for i in range(num_workers)] + + +# --------------------------------------------------------------------------- +# Malformed embedded library binary — the parser must reject corrupt blobs +# with a clean RuntimeError, never crash. Each case targets one guard in +# ProcessEmbeddedLibraryBin / BlobReader. +# --------------------------------------------------------------------------- + + +def _load_bad_blob(tmp_path: Path, blob: bytes) -> tvm_ffi.Module: + """Build an object carrying `blob` as __tvm_ffi__library_bin and load it.""" + libbin_obj = _build_library_bin_object(tmp_path, blob, extra_export="__tvm_ffi_probe_marker") + return default_session().load_module(libbin_obj) + + +def test_malformed_indptr_single_entry(tmp_path: Path) -> None: + """Indptr with one entry (num_modules == 0) must be rejected, not return modules[0].""" + blob = _make_library_bin(modules=[], indptr=[0], children=[]) + with pytest.raises(Exception, match="at least 2 entries"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_indptr_empty(tmp_path: Path) -> None: + """An empty indptr vector is rejected.""" + blob = _make_library_bin(modules=[], indptr=[], children=[]) + with pytest.raises(Exception, match="at least 2 entries"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_indptr_nonzero_start(tmp_path: Path) -> None: + """Indptr must start at 0.""" + blob = _make_library_bin(modules=[("_lib", b"")], indptr=[1, 1], children=[]) + with pytest.raises(Exception, match="must start at 0"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_indptr_non_monotonic(tmp_path: Path) -> None: + """Indptr must be non-decreasing.""" + # 3 entries -> 2 modules; indptr goes 0 -> 2 -> 1 (decreasing). + blob = _make_library_bin( + modules=[("_lib", b""), ("_lib", b"")], indptr=[0, 2, 1], children=[0, 0] + ) + with pytest.raises(Exception, match="non-decreasing"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_indptr_tail_exceeds_children(tmp_path: Path) -> None: + """indptr.back() must not exceed the child index count.""" + blob = _make_library_bin(modules=[("_lib", b"")], indptr=[0, 5], children=[]) + with pytest.raises(Exception, match="exceeds child index count"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_child_index_out_of_range(tmp_path: Path) -> None: + """A child index >= module count must be rejected.""" + blob = _make_library_bin(modules=[("_lib", b"")], indptr=[0, 1], children=[9]) + with pytest.raises(Exception, match="child index out of range"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_truncated_stream(tmp_path: Path) -> None: + """A length header larger than the actual bytes triggers a bounds error.""" + # Inner stream claims a u64 vector of 4 entries but supplies no bytes for them. + inner = _u64(4) # count=4, then nothing + blob = _raw_library_bin(inner) + with pytest.raises(Exception, match=r"end of blob|exceeds remaining|exceeds blob"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_oversize_vector_count(tmp_path: Path) -> None: + """A huge vector count is rejected before reserve(), not OOM.""" + inner = _u64(2**60) # absurd count with no backing bytes + blob = _raw_library_bin(inner) + with pytest.raises(Exception, match="exceeds remaining blob"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_unregistered_kind(tmp_path: Path) -> None: + """An embedded module of unknown kind reports the missing loader.""" + blob = _make_library_bin( + modules=[("_lib", b""), ("no_such_kind", b"payload")], + indptr=[0, 1, 1], + children=[1], + ) + with pytest.raises(Exception, match=r"load_from_bytes\.no_such_kind is not registered"): + _load_bad_blob(tmp_path, blob) + + +def test_malformed_loader_returns_non_module(tmp_path: Path) -> None: + """A loader returning a non-Module (e.g. None) is rejected by the cast.""" + + @tvm_ffi.register_global_func("ffi.Module.load_from_bytes.orcjit_null_probe", override=True) + def _load_null(_data: bytes): # noqa: ANN202 + return None + + blob = _make_library_bin( + modules=[("_lib", b""), ("orcjit_null_probe", b"x")], + indptr=[0, 1, 1], + children=[1], + ) + with pytest.raises(Exception, match=r"Cannot convert.*ffi\.Module"): + _load_bad_blob(tmp_path, blob) diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py index f27a48fe3..09bb2f6a6 100644 --- a/python/tvm_ffi/cpp/extension.py +++ b/python/tvm_ffi/cpp/extension.py @@ -518,6 +518,16 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912 ninja.append(" command = ld -r -o $out $in") ninja.append("") + # Copy a single object straight through, bypassing `ld -r`. A partial + # link of one input is a pointless round-trip, and some system linkers + # (e.g. GNU ld 2.45) assign a nonzero sh_addr to zero-size, non-SHF_ALLOC + # sections like `.note.GNU-stack` during it. JITLink then materializes a + # zero-length block at that address, which collides with a real block in + # EHFrameEdgeFixer's block map and aborts C++ (eh_frame-bearing) loads. + ninja.append("rule copy_object") + ninja.append(" command = cp -f $in $out") + ninja.append("") + if embed_cubin: ninja.append("rule embed_cubin") ninja.append( @@ -562,7 +572,10 @@ def _generate_ninja_build( # noqa: PLR0915, PLR0912 if object_mode: # Object-only output: merge all object files into the target. if not IS_WINDOWS: - ninja.append(f"build {output_name}: merge_objects {' '.join(obj_files)}") + # A single object needs no partial link — copy it through to avoid + # `ld -r` rewriting section addresses (see the copy_object rule). + rule = "copy_object" if len(obj_files) == 1 else "merge_objects" + ninja.append(f"build {output_name}: {rule} {' '.join(obj_files)}") ninja.append("") ninja.append(f"default {output_name}") else: