[FIX] Load ORC JIT addon through TVM-FFI

addons/tvm_ffi_orcjit/CMakeLists.txt

@@ -22,6 +22,9 @@ endif ()
 
 # ---- Find packages ----
 find_package(LLVM REQUIRED CONFIG)
+if (LLVM_PACKAGE_VERSION VERSION_LESS 22.0)
+  message(FATAL_ERROR "tvm_ffi_orcjit requires LLVM 22 or newer; found ${LLVM_PACKAGE_VERSION}")
+endif ()
 message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION} in ${LLVM_DIR}")
 
 find_package(
@@ -76,12 +79,22 @@ execute_process(
   COMMAND ${LLVM_CONFIG_EXE} --link-static --libs Core OrcJIT Support native
   OUTPUT_STRIP_TRAILING_WHITESPACE
   OUTPUT_VARIABLE _llvm_libs
+  ERROR_VARIABLE _llvm_libs_error
+  RESULT_VARIABLE _llvm_libs_rc
 )
+if (NOT _llvm_libs_rc EQUAL 0)
+  message(FATAL_ERROR "llvm-config --link-static --libs failed: ${_llvm_libs_error}")
+endif ()
 execute_process(
   COMMAND ${LLVM_CONFIG_EXE} --link-static --ldflags
   OUTPUT_STRIP_TRAILING_WHITESPACE
   OUTPUT_VARIABLE _llvm_ldflags
+  ERROR_VARIABLE _llvm_ldflags_error
+  RESULT_VARIABLE _llvm_ldflags_rc
 )
+if (NOT _llvm_ldflags_rc EQUAL 0)
+  message(FATAL_ERROR "llvm-config --link-static --ldflags failed: ${_llvm_ldflags_error}")
+endif ()
 separate_arguments(_llvm_libs_list NATIVE_COMMAND "${_llvm_libs}")
 separate_arguments(_llvm_ldflags_list NATIVE_COMMAND "${_llvm_ldflags}")
 

addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md

@@ -155,8 +155,9 @@ symbol resolution must be deferred until all relevant objects are present.
 
 ## 3. LLVM ORC JIT v2
 
-**ORC** stands for *On Request Compilation*. LLVM ORC JIT v2 (introduced in LLVM 9,
-stabilized in LLVM 13+) is a complete redesign of LLVM's JIT infrastructure. It is
+**ORC** stands for *On Request Compilation*. The ORCv2 APIs have been available since
+LLVM 7 and replaced ORCv1 completely in LLVM 12. They redesign LLVM's JIT infrastructure
+around a composable, concurrent symbol/materialization model. ORCv2 is
 designed to be composable, asynchronous, and correct for production use (unlike the
 older `MCJIT` which had several fundamental limitations around multi-module linking).
 
@@ -312,7 +313,7 @@ The three platform objects in LLVM are:
 | Platform | OS | Init section driven |
 | --- | --- | --- |
 | `MachOPlatform` | macOS / iOS | `__DATA,__mod_init_func` |
-| `ELFNativePlatform` | Linux / ELF | `.init_array`, TLS |
+| `ELFNixPlatform` | Linux / ELF | `.init_array`, TLS |
 | `COFFPlatform` | Windows | `.CRT$XC*` init, `__cxa_atexit` interop |
 
 `ExecutorNativePlatform` is a convenience builder that auto-selects the right platform
@@ -322,17 +323,18 @@ for the host OS and loads the ORC runtime from a given path.
 
 The addon takes a different approach on each platform:
 
-- **macOS**: ORC platform support is *optional*. When the caller passes an ORC runtime
-  path to `ExecutionSession`, `ExecutorNativePlatform` activates `MachOPlatform`.
-  `jit_->initialize(dylib)` and `jit_->deinitialize(dylib)` then drive `__mod_init_func`
-  and `__cxa_atexit` teardown natively. Without the path, the addon falls back to its
-  own `InitFiniPlugin`.
+- **macOS**: the addon deliberately does not configure an ORC platform. It uses
+  `InitFiniPlugin` for `__mod_init_func` / `__mod_term_func` and a per-dylib
+  `__cxa_atexit` shim. This avoids a compact-unwind address-delta failure in the
+  current `MachOPlatform` integration. The `orc_rt` argument is ignored.
 - **Windows**: `COFFPlatform` is skipped entirely because it requires MSVC CRT symbols
   (`_CxxThrowException`, RTTI vtables, iostream objects) that are not resolvable in
   the JIT context. Instead, `InitFiniPlugin` manually handles `.CRT$XC*` / `.CRT$XT*`
-  init/fini sections.
-- **Linux**: `ELFNativePlatform` is not used. `InitFiniPlugin` handles `.init_array` /
-  `.fini_array` / `.ctors` / `.dtors` directly, without the ORC runtime.
+  init/fini sections. The `orc_rt` argument is ignored.
+- **Linux**: the default embedded `liborc_rt` configures `ExecutorNativePlatform`
+  (and therefore `ELFNixPlatform`). A custom archive may be supplied by path or bytes,
+  and `None` disables the platform. `InitFiniPlugin` still collects and invokes ELF
+  init/fini arrays to work around current upstream behavior.
 
 ---
 
@@ -360,21 +362,28 @@ unit: load all its objects at once, then look up functions on the result.
 
 ### 4.2 Loading and lookup
 
-`load_module` adds each object to the JITDylib, and JITLink parses it into a
-`LinkGraph`, resolves relocations, and allocates executable JIT memory.
-`get_function` looks the symbol up (materializing lazily), then wraps the raw
+`load_module` registers each object with the JITDylib. On the first relevant lookup,
+JITLink parses it into a `LinkGraph`, resolves relocations, and allocates executable
+JIT memory. `get_function` looks the symbol up (materializing lazily), then wraps the raw
 pointer as a `tvm_ffi::Function`. Symbols resolve against the dylib's own
-default link order (this dylib → Platform → process/runtime symbols); objects
+link order (this dylib → LLJIT defaults, with a compiler-selected C++ runtime
+inserted before process symbols on Linux); objects
 that reference each other must be loaded together, since there is no linking
 between separate `load_module` results.
 
-Two addon-specific pieces sit in this pipeline:
+Several addon-specific pieces sit in this pipeline:
 
 - **`InitFiniPlugin`** — a JITLink pass plugin that keeps init/fini sections
   (`.init_array`/`.ctors`/`.fini_array`/`.dtors`, `__mod_init_func`, `.CRT$XC*`)
   live, then collects their function pointers after fixup. The addon runs them
   in priority order at first lookup and at teardown, replacing the ORC
   platform's initializer machinery. See `llvm_patches/init_fini_plugin.h`.
+- **Linux slab-pool memory manager** — reserves contiguous virtual-address regions,
+  separates executable from non-executable allocations, recycles freed regions, and
+  grows by adding slabs. This keeps 32-bit PC-relative relocations in range and offers
+  explicit reclamation through `clear_free_slabs()`.
+- **GOTPCRELX correction (Linux/x86-64)** — repairs or reverses unsafe JITLink
+  relaxations before fixup.
 - **Windows DLL import stubs** — `DLLImportDefinitionGenerator` resolves
   `__imp_XXX` references to host-DLL functions by emitting JIT-memory pointer +
   trampoline stubs, keeping `PCRel32` fixups within ±2 GB of the JIT code. See
@@ -391,12 +400,12 @@ import tvm_ffi_orcjit as oj
 sess = oj.default_session()
 
 # 2. Load a compiled object file into a fresh JITDylib
-#    → object parsed, JITLink links it, InitFiniPlugin collects ctors,
-#      context symbols injected eagerly, embedded binary (if any) expanded
+#    → objects registered; context and embedded-binary probes may materialize
+#      objects immediately, and context slots are populated before ctors run
 mod = sess.load_module("add.o")   # returns a tvm_ffi.Module
 
 # 3. Look up and call a function
-#    → LLVM resolves "__tvm_ffi_add"; pending constructors fire on first lookup
+#    → LLVM resolves "__tvm_ffi_add"; any pending constructors run before return
 result = mod.add(3, 4)   # → 7
 ```
 
@@ -426,6 +435,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl);
 | `JITDylib` | Symbol namespace / virtual shared library | `ORCJITDynamicLibraryObj::dylib_` |
 | `JITLink` | LLVM's JIT-aware linker | Used inside `ObjectLinkingLayer` |
 | JITLink pass pipeline | Pre-prune → post-alloc → post-fixup hooks | Where `InitFiniPlugin` runs |
+| Slab pool | Contiguous, growable Linux JIT memory arena | `SlabPoolMemoryManager` / `Slab` |
 | `DefinitionGenerator` | Fallback symbol provider | `DLLImportDefinitionGenerator` (Win) |
 | Link order | Search path across JITDylibs for symbol resolution | LLJIT default (Main → Platform → ProcessSymbols) |
 | `__tvm_ffi_` prefix | Namespace for TVM-FFI exported functions | Used in `GetFunction()` |

addons/tvm_ffi_orcjit/README.md

@@ -26,8 +26,9 @@ TVM-FFI exported functions.
 - **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2
 - **High-Level Loading**: `default_session().load_module(...)` mirrors `tvm_ffi.load_module`, returning a plain `tvm_ffi.Module`
 - **Unified Input**: Load from a file path, in-memory object bytes, or a list mixing both
-- **Shared Session**: A process-wide session so multiple callers share one JIT environment (process symbols, arena, linking)
+- **Shared Session**: A process-wide session so callers share process-symbol resolution and the Linux slab pool while loaded modules remain isolated
 - **Symbol Isolation**: Separate `load_module` calls define independent symbol namespaces, so they can define the same symbol without conflicts
+- **Bounded-Range Memory**: A growable Linux slab pool keeps JIT code/data close enough for 32-bit PC-relative relocations and supports explicit drained-slab reclamation
 - **Init/Fini Support**: Handles static constructors/destructors across ELF (`.init_array`/`.ctors`), Mach-O (`__mod_init_func`), and COFF (`.CRT$XC*`/`.CRT$XT*`)
 - **Cross-Platform**: Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64)
 - **Multi-Compiler**: Tested with LLVM Clang, GCC, Apple Clang, MSVC, and clang-cl
@@ -206,7 +207,8 @@ Compile: `clang -O2 -c -o example.o example.c`
   relocations from COFF objects before JITLink graph building, working around a
   JITLink limitation with COMDAT section symbols.
 
-Please refers to [ORCJIT_PRIMER.md](./ORCJIT_PRIMER.md) to learn more about object file, linking, llvm orcjit v2, and how the addon works.
+Refer to [ORCJIT_PRIMER.md](./ORCJIT_PRIMER.md) for background on object files,
+linking, LLVM ORC JIT v2, and the addon's architecture.
 
 ## Project Structure
 
@@ -219,6 +221,9 @@ tvm_ffi_orcjit/
 │   ├── orcjit_session.h
 │   ├── orcjit_dylib.cc         # JIT dylib module (object loading, symbol lookup)
 │   ├── orcjit_dylib.h
+│   ├── orcjit_memory_manager.*  # Growable Linux slab pool
+│   ├── orcjit_slab.*            # Contiguous-VA allocator and page protection
+│   ├── llvm_patches/            # Isolated upstream LLVM workarounds
 │   └── orcjit_utils.h          # LLVM error handling utilities
 ├── python/tvm_ffi_orcjit/
 │   ├── __init__.py             # Module exports and library loading
@@ -261,6 +266,12 @@ The package requires LLVM 22+. Set `LLVM_PREFIX` to the LLVM install prefix:
 export LLVM_PREFIX=/path/to/llvm
 ```
 
+### Reclaiming unused JIT memory on Linux
+
+Dropping a module returns its regions to the session's slab pool for reuse. To
+return fully drained slabs to the operating system, call
+`session.clear_free_slabs()`. The call is synchronized with concurrent JIT work.
+
 ## License
 
 Apache License 2.0

addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/__init__.py

@@ -29,8 +29,6 @@
 
 """
 
-import ctypes
-import os
 import platform
 from pathlib import Path
 
@@ -53,29 +51,19 @@
 ]
 _lib_path = None
 for path in _LIB_PATH:
-    if path.exists():
-        _ = load_module(str(path))
+    if path.is_file():
         _lib_path = path
+        break
 if _lib_path is None:
     raise RuntimeError(
-        f"Could not find {_LIB_NAME}. "
-        f"Searched in {_LIB_PATH} and site-packages. "
-        f"Please ensure the package is installed correctly."
+        f"Could not find {_LIB_NAME}. Searched {_LIB_PATH}. "
+        "Please ensure the package is installed correctly."
     )
 
-# Keep a second, process-lifetime local handle. RTLD_NODELETE is important for
-# modules pinned by keep_module_alive: their object deleters point into JIT code
-# owned by this DSO and may run during interpreter shutdown, after Python module
-# globals have otherwise released their handles. This does not promote the DSO
-# or its statically linked LLVM into the process-global symbol namespace.
-if os.name == "posix":
-    _c_lib = ctypes.CDLL(
-        str(_lib_path),
-        mode=ctypes.RTLD_LOCAL | getattr(os, "RTLD_NODELETE", 0),
-    )
-else:
-    _dll_directory = os.add_dll_directory(str(_lib_path.parent))
-    _c_lib = ctypes.CDLL(str(_lib_path))
+# The TVM-FFI loader uses local symbol scope. Retain the returned module here
+# and in TVM-FFI's process-lifetime module registry so registered functions and
+# JIT-owned object deleters remain valid through interpreter shutdown.
+_lib_module = load_module(_lib_path, keep_module_alive=True)
 
 from .session import ExecutionSession, default_session
 

addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py

@@ -119,10 +119,11 @@ def __init__(
         slab_size : int
             Per-slab capacity in bytes for the JIT memory manager. Linux only —
             ignored on macOS and Windows, where the slab allocator is compiled
-            out. 0 = arch default (64 MB; initial slab halves on mmap failure
-            down to 8 MB under RLIMIT_AS / container limits), >0 = custom size,
-            <0 = disable slab allocator (LLJIT uses its default scattered-mmap
-            allocator).
+            out. 0 = 64 MB default (the initial slab halves on mmap failure
+            down to 8 MB under RLIMIT_AS / container limits), >=4 MB = custom
+            size, <0 = disable the slab allocator (LLJIT uses its default
+            scattered-mmap allocator). Positive values below 4 MB are rejected
+            because both allocation pools need a 2 MB commit chunk.
 
             The session holds a growable pool of slabs: a fresh slab is mmap'd
             on demand when no existing one can fit a graph. Graphs that don't
@@ -190,8 +191,8 @@ def load_module(
         ``keep_module_alive`` mirrors :func:`tvm_ffi.load_module`'s option of
         the same name. When True, the module is inserted into the runtime's
         process-global module registry, so its JITDylib — and every function
-        pointer, deleter, and static allocation it owns — stays mapped until
-        the interpreter unloads ``libtvm_ffi``. Use this when Objects produced
+        pointer, deleter, and static allocation it owns — stays mapped for the
+        duration of the process. Use this when Objects produced
         by the module may outlive the local ``mod`` reference (e.g., a
         JIT-allocated ``String`` or ``Array`` returned to Python and held past
         ``del mod``). When False (default), the caller owns the module's
@@ -243,10 +244,9 @@ def clear_free_slabs(self) -> int:
         Fresh slabs that have never been allocated on are preserved, so
         the session remains ready to accept new work.
 
-        Safety: call when no JIT work is in flight on another thread. From
-        single-threaded Python this is always safe; once ``del lib`` has
-        returned, the C++ destructor has finished and the slab's live count
-        reflects the drop.
+        The operation is serialized with JIT allocation and module teardown,
+        so it is safe to call while other host threads use the same session.
+        Only fully drained slabs are reclaimed.
 
         Returns
         -------
@@ -267,9 +267,10 @@ def default_session() -> ExecutionSession:
     """Return the process-wide shared execution session.
 
     A single leaked, never-destroyed session shared by all callers in the
-    process, so they share one LLVM ``ExecutionSession`` — hence process
-    symbols, the slab arena, and cross-library linking. Created on first call
-    and cached for the lifetime of the process.
+    process, so they share one LLVM ``ExecutionSession`` — hence process-symbol
+    resolution, the slab pool, and synchronization infrastructure. Separate
+    loaded modules remain isolated symbol namespaces. Created on first call and
+    cached for the lifetime of the process.
 
     The session uses the ORC runtime embedded in the extension (no on-disk path
     lookup). For an isolated session or a tuned arena, construct an

addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md

@@ -45,17 +45,28 @@ Each patch file opens with a fixed-shape header describing:
 
 - **ELF init/fini** (Linux branch of `init_fini_plugin.{h,cc}`)
   LLVM issue: [llvm/llvm-project#175981](https://github.com/llvm/llvm-project/issues/175981).
-  Upstream status: open, patch submitted.
-  Remove when: LLVM floor bumps past the release that contains the fix.
+  Upstream status: merged; included in LLVM 23.1.1, not LLVM 22.1.0.
+  Remove when: the LLVM floor reaches 23 and the addon switches its lifecycle
+  calls to the upstream `ELFNixPlatform` path.
 
 - **COFF ctor/dtor** (Windows branch of `init_fini_plugin.{h,cc}`)
   LLVM issue: COFFPlatform stalled.
   Upstream status: stalled 2+ years.
   Remove when: COFFPlatform becomes usable end-to-end with clang-cl /
   MSVC objects.
 
-macOS already has working `MachOPlatform`, so no patch file is needed
-for that platform.
+- **Mach-O `__cxa_atexit` scoping** (`macho_cxa_atexit_shim.{h,cc}`)
+  Upstream dependency: re-enabling `MachOPlatform` after the compact-unwind
+  per-graph DSO-base fix is available.
+  Remove when: the addon can use `MachOPlatform` end-to-end.
+
+- **COFF unwind-data stripping** (`win_coff_pdata_strip.{h,cc}`)
+  Upstream dependency: usable `COFFPlatform` support with SEH registration.
+  Remove when: `.pdata` / `.xdata` can be registered and relocated normally.
+
+- **Windows DLL import stubs** (`win_dll_import_generator.{h,cc}`)
+  Upstream dependency: usable `COFFPlatform` and in-range DLL call stubs.
+  Remove when: host DLL imports work end-to-end without the custom generator.
 
 ## Removal checklist
 

addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h

@@ -52,11 +52,11 @@
  *
  * ## Removal — Linux
  *
- * LLVM issue: https://github.com/llvm/llvm-project/issues/175981
- * When the upstream fix lands and the project's minimum LLVM version
- * bumps past the first release containing it, replace this plugin's
- * Linux usage with `ELFNixPlatform` and delete the ELF handling path
- * from this file.  Concretely:
+ * LLVM issue: https://github.com/llvm/llvm-project/pull/175981
+ * The upstream fix is included in LLVM 23.1.1 but not LLVM 22.1.0, the
+ * addon's current CI baseline. When the project's minimum LLVM version
+ * reaches 23 and its lifecycle calls use `ELFNixPlatform`, delete the ELF
+ * handling path from this file. Concretely:
  *   - Remove the ELF-section branches (`.init_array`, `.ctors`,
  *     `.fini_array`, `.dtors`) from `InitFiniPlugin::modifyPassConfig`.
  *   - If no platform still needs this plugin, delete this file outright

addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc

@@ -67,12 +67,12 @@ CxaAtexitRecordsScope::CxaAtexitRecordsScope(CxaAtexitRecords* records)
 }
 CxaAtexitRecordsScope::~CxaAtexitRecordsScope() { g_active_cxa_records = prev_; }
 
-void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd) {
+llvm::Error InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd) {
   llvm::orc::SymbolMap shim_syms;
   shim_syms[ES.intern("___cxa_atexit")] = {
       llvm::orc::ExecutorAddr::fromPtr(reinterpret_cast<void*>(&tvm_ffi_cxa_atexit_shim)),
       llvm::JITSymbolFlags::Exported | llvm::JITSymbolFlags::Callable};
-  llvm::cantFail(jd.define(llvm::orc::absoluteSymbols(std::move(shim_syms))));
+  return jd.define(llvm::orc::absoluteSymbols(std::move(shim_syms)));
 }
 
 void DrainCxaAtexit(CxaAtexitRecords& records) {

addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h

@@ -61,6 +61,7 @@
 #ifdef __APPLE__
 
 #include <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/Support/Error.h>
 
 #include <utility>
 #include <vector>
@@ -100,7 +101,7 @@ class CxaAtexitRecordsScope {
  *  libSystem fallback — JITDylib::define-time symbols are searched before
  *  the link order.
  */
-void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd);
+llvm::Error InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd);
 
 /*! \brief Drain captured `(fn, arg)` records LIFO, running each dtor.
  *