You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ya...@apache.org on 2026/05/12 13:01:55 UTC
(tvm-ffi) branch main updated: [REFACTOR][OrcJIT] Per-session slab-pool memory manager (#574)
This is an automated email from the ASF dual-hosted git repository.
cyx-6 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new fb5ad97 [REFACTOR][OrcJIT] Per-session slab-pool memory manager (#574)
fb5ad97 is described below
commit fb5ad97ad629763f893d7fe593f6b096d0baa737
Author: Yaxing Cai <ca...@gmail.com>
AuthorDate: Tue May 12 21:01:49 2026 +0800
[REFACTOR][OrcJIT] Per-session slab-pool memory manager (#574)
## Summary
Refactors the `tvm_ffi_orcjit` addon along three axes — memory
lifecycle,
LLVM-workaround isolation, and macOS/Windows enablement. 14 commits,
each
independently revertible.
## What changed
**Memory lifecycle.** The fixed ~1 GB per-session arena becomes a
growable
pool of fixed-size `Slab`s. Dropping a `DynamicLibrary` now calls
`removeJITDylib` and actually frees its JIT pages (previously code
accumulated until session destruction). `Session.clear_free_slabs()`
releases drained slabs back to the OS. `orcjit_memory_manager.cc` shrank
~740 → 174 lines with `Slab` pulled into its own `orcjit_slab.{h,cc}`.
**LLVM-workaround isolation.** Every upstream-LLVM workaround — formerly
inlined in `orcjit_session.cc` — moved into `llvm_patches/` (5 modules:
`gotpcrelx_fix`, `init_fini_plugin`, `macho_cxa_atexit_shim`,
`win_coff_pdata_strip`, `win_dll_import_generator`). Each has a
fixed-shape
header with issue links, trigger, symptom, and a `## Removal` checklist
so
it's mechanically deletable once upstream fixes land.
`orcjit_session.cc`
shrank ~870 → 320 lines of pure orchestration.
**macOS + Windows enablement.** macOS: skip `MachOPlatform` to sidestep
the
compact-unwind 32-bit-delta bug, and interpose `___cxa_atexit`
per-JITDylib
with a TLS-routed shim that drains at dylib drop. Windows: stub DLL
imports
through JIT-allocated PLT entries (keeps PCRel32 in range) and strip
`.pdata` / `.xdata` relocations pre-graph-build. The three-platform
strategy is documented inline in each patch header.
**Function lifetime + perf.** Captured `Function` values hold a ref to
the
owning dylib so JIT pages survive as long as the handle.
`FunctionObj::safe_call` binds directly to the JIT'd thunk, eliminating
a
per-call indirection.
## API
- **Break**: `ExecutionSession(arena_size=…)` →
`ExecutionSession(slab_size=…)`.
Size is per-slab, default 64 MB; `-1` disables the pool (LLJIT falls
back
to scattered mmap).
- **New**: `ExecutionSession.clear_free_slabs() -> int` — reclaims
drained
slabs, returns the count.
## Test plan
- [x] aarch64 Linux: `pytest addons/tvm_ffi_orcjit/tests`
- [x] macOS arm64: 99 passed, 42 skipped
- [ ] CI: lint / clang-tidy / Windows / x86_64 Linux
---
addons/tvm_ffi_orcjit/CMakeLists.txt | 40 +-
.../python/tvm_ffi_orcjit/session.py | 54 +-
.../tvm_ffi_orcjit/src/ffi/llvm_patches/README.md | 70 ++
.../src/ffi/llvm_patches/gotpcrelx_fix.cc | 162 +++++
.../src/ffi/llvm_patches/gotpcrelx_fix.h | 92 +++
.../src/ffi/llvm_patches/init_fini_plugin.cc | 258 +++++++
.../src/ffi/llvm_patches/init_fini_plugin.h | 123 ++++
.../src/ffi/llvm_patches/macho_cxa_atexit_shim.cc | 91 +++
.../src/ffi/llvm_patches/macho_cxa_atexit_shim.h | 119 ++++
.../src/ffi/llvm_patches/win_coff_pdata_strip.cc | 131 ++++
.../src/ffi/llvm_patches/win_coff_pdata_strip.h | 83 +++
.../ffi/llvm_patches/win_dll_import_generator.cc | 154 +++++
.../ffi/llvm_patches/win_dll_import_generator.h | 93 +++
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc | 92 ++-
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h | 13 +
.../src/ffi/orcjit_memory_manager.cc | 738 +++-----------------
.../tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h | 249 +++----
addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 756 +++------------------
addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h | 34 +-
.../{orcjit_memory_manager.cc => orcjit_slab.cc} | 499 ++++++++------
addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h | 324 +++++++++
addons/tvm_ffi_orcjit/tests/test_basic.py | 397 +++++++++++
.../{test_arena.py => test_memory_manager.py} | 227 ++++---
23 files changed, 2962 insertions(+), 1837 deletions(-)
diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt b/addons/tvm_ffi_orcjit/CMakeLists.txt
index 9a86bac..b0ef5cb 100644
--- a/addons/tvm_ffi_orcjit/CMakeLists.txt
+++ b/addons/tvm_ffi_orcjit/CMakeLists.txt
@@ -38,8 +38,16 @@ find_package(tvm_ffi CONFIG REQUIRED)
# ---- Build shared library ----
add_library(
- tvm_ffi_orcjit SHARED src/ffi/orcjit_session.cc src/ffi/orcjit_dylib.cc
- src/ffi/orcjit_memory_manager.cc
+ tvm_ffi_orcjit SHARED
+ src/ffi/orcjit_session.cc
+ src/ffi/orcjit_dylib.cc
+ src/ffi/orcjit_memory_manager.cc
+ src/ffi/orcjit_slab.cc
+ src/ffi/llvm_patches/gotpcrelx_fix.cc
+ src/ffi/llvm_patches/init_fini_plugin.cc
+ src/ffi/llvm_patches/macho_cxa_atexit_shim.cc
+ src/ffi/llvm_patches/win_coff_pdata_strip.cc
+ src/ffi/llvm_patches/win_dll_import_generator.cc
)
set_target_properties(
tvm_ffi_orcjit PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON
@@ -145,27 +153,27 @@ elseif (UNIX AND NOT WIN32)
endif ()
# ---- Find and bundle liborc_rt ----
+# Platform notes:
+#
+# * Linux: liborc_rt is used for ExecutorNativePlatform → ELFNixPlatform.
+# * macOS: MachOPlatform is skipped to sidestep the compact-unwind delta bug (see
+# src/ffi/orcjit_session.cc); liborc_rt is unused.
+# * Windows: COFFPlatform is not hooked up; liborc_rt is unused.
set(ORC_RT_PATH "")
-if (DEFINED ENV{ORC_RT_PATH})
- set(ORC_RT_PATH "$ENV{ORC_RT_PATH}")
-else ()
- if (WIN32)
- file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/lib/clang/*/lib/orc_rt*.lib")
- if (NOT _orc_rt_candidates)
- file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/Library/lib/clang/*/lib/orc_rt*.lib")
- endif ()
+if (NOT APPLE AND NOT WIN32)
+ if (DEFINED ENV{ORC_RT_PATH})
+ set(ORC_RT_PATH "$ENV{ORC_RT_PATH}")
else ()
file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/lib/clang/*/lib/liborc_rt*.a")
+ if (_orc_rt_candidates)
+ list(GET _orc_rt_candidates 0 ORC_RT_PATH)
+ endif ()
endif ()
- if (_orc_rt_candidates)
- list(GET _orc_rt_candidates 0 ORC_RT_PATH)
+ if (NOT ORC_RT_PATH)
+ message(WARNING "Could not find liborc_rt. ORC runtime features will be disabled.")
endif ()
endif ()
-if (NOT ORC_RT_PATH)
- message(WARNING "Could not find liborc_rt. ORC runtime features will be disabled.")
-endif ()
-
if (ORC_RT_PATH)
message(STATUS "Found liborc_rt: ${ORC_RT_PATH}")
file(COPY "${ORC_RT_PATH}" DESTINATION "${CMAKE_BINARY_DIR}")
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 02fccaf..391fd35 100644
--- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py
+++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py
@@ -35,7 +35,14 @@ def _find_orc_rt_library() -> str | None:
# avoiding all C++ runtime dependencies (magic statics, RTTI, sized delete,
# SEH, COMDAT). Our custom InitFiniPlugin handles .CRT$XC*/.CRT$XT* init/fini
# sections, and DLLImportDefinitionGenerator resolves __imp_ DLL import stubs.
- if sys.platform == "win32":
+ #
+ # macOS: skip ORC runtime too. ExecutorNativePlatform would install
+ # MachOPlatform, which triggers a compact-unwind 32-bit-delta bug in
+ # JITLink when a user graph mmaps below the per-JITDylib Mach-O header
+ # (see repo-root fix-machoplatform-libunwind-dso-base.patch). Our
+ # InitFiniPlugin handles __mod_init_func / __mod_term_func instead.
+ # Tradeoff: no C++ exception unwinding across JIT frames on macOS.
+ if sys.platform in ("win32", "darwin"):
return None
patterns = ["liborc_rt*.a"]
for pattern in patterns:
@@ -60,25 +67,34 @@ class ExecutionSession(Object):
"""
- def __init__(self, orc_rt_path: str | None = None, arena_size: int = 0) -> None:
+ 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.
- arena_size: Arena size in bytes for the JIT memory manager.
- Linux only — ignored on macOS and Windows, where the
- arena is compiled out.
- 0 = arch default (4GB x86_64, 8GB AArch64; falls back to
- smaller sizes under RLIMIT_AS), >0 = custom size,
- <0 = disable arena.
+ 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.
"""
if orc_rt_path is None:
orc_rt_path = _find_orc_rt_library()
if orc_rt_path is None:
orc_rt_path = ""
- self.__init_handle_by_constructor__(_ffi_api.ExecutionSession, orc_rt_path, arena_size) # type: ignore
+ 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.
@@ -94,3 +110,23 @@ class ExecutionSession(Object):
lib = DynamicLibrary.__new__(DynamicLibrary)
lib.__move_handle_from__(handle)
return lib
+
+ def clear_free_slabs(self) -> int:
+ """Release drained slabs (no live JIT allocations) back to the OS.
+
+ Call this after dropping a batch of libraries to reclaim RSS.
+ 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.
+
+ Returns:
+ 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
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md
new file mode 100644
index 0000000..c2706a7
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/README.md
@@ -0,0 +1,70 @@
+<!--- 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. -->
+
+# LLVM patches
+
+This directory holds workarounds for upstream LLVM bugs and missing
+features. **A file belongs here if and only if its entire reason to
+exist is an LLVM defect, and we would delete the whole file once
+upstream catches up.**
+
+Features that happen to coexist with a bug (e.g. the arena memory
+manager, which we keep for THP and contiguous layout regardless of
+LLVM state) live at the top level of `src/ffi/`, not here.
+
+Each patch file opens with a fixed-shape header describing:
+
+- which LLVM issue it addresses (link or "not yet filed"),
+- affected version range,
+- exact trigger conditions,
+- symptom without the patch, and
+- a `## Removal` section listing the `#include` and plugin-registration
+ line(s) to delete when the upstream fix lands and the project's
+ minimum LLVM version bumps past it.
+
+## Index
+
+- **GOTPCRELX relaxation** (`gotpcrelx_fix.{h,cc}`)
+ LLVM issue: TBD — issue not yet filed.
+ Upstream status: open.
+ Remove when: LLVM floor bumps past the release that contains the fix.
+
+- **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.
+
+- **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.
+
+## Removal checklist
+
+When deleting a patch file:
+
+1. Delete the `.h` and `.cc` pair.
+2. Remove the matching `#include "llvm_patches/<file>.h"` in
+ `orcjit_session.cc`.
+3. Remove the plugin-registration line(s) in `orcjit_session.cc`
+ identified in the file's `## Removal` header block.
+4. Remove the corresponding sources from `addons/tvm_ffi_orcjit/CMakeLists.txt`.
+5. Update the index above.
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.cc
new file mode 100644
index 0000000..7f5593c
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.cc
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file gotpcrelx_fix.cc
+ * \brief LLVM JITLink GOTPCRELX relaxation bug workaround (x86_64).
+ *
+ * See gotpcrelx_fix.h for the trigger, symptom, and removal procedure.
+ */
+#include "gotpcrelx_fix.h"
+
+#if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64))
+
+#include <llvm/ADT/DenseMap.h>
+#include <llvm/ExecutionEngine/JITLink/JITLink.h>
+#include <llvm/ExecutionEngine/JITLink/x86_64.h>
+#include <llvm/Support/Error.h>
+#include <llvm/Support/MathExtras.h>
+
+#include <cstdint>
+#include <string>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+namespace {
+
+/*! \brief Correct broken GOTPCRELX relaxations produced by
+ * optimizeGOTAndStubAccesses().
+ *
+ * Strategy:
+ * 1. Build target-symbol → GOT-entry-symbol map (O(B+S) up front).
+ * 2. For every Pointer32 edge whose preceding bytes are 67 e8
+ * (relaxed call) or e9 (relaxed jmp):
+ * - If the target is reachable via a signed 32-bit PC-relative
+ * displacement, change the edge to BranchPCRel32.
+ * - Otherwise revert the relaxation: restore the original
+ * indirect-call/jmp opcode bytes (ff 15 / ff 25), retarget
+ * the edge to the GOT entry, and use PCRel32 with addend 0
+ * (JITLink normalises GOTPCRELX addends to 0).
+ */
+llvm::Error fixBrokenGOTPCRELXRelaxation(llvm::jitlink::LinkGraph& G) {
+ using namespace llvm::jitlink;
+ // Build block → first symbol at offset 0 (for GOT entry symbol lookup).
+ llvm::DenseMap<Block*, Symbol*> BlockToSym;
+ for (auto* Sym : G.defined_symbols()) {
+ if (Sym->getOffset() == 0 && !BlockToSym.count(&Sym->getBlock())) {
+ BlockToSym[&Sym->getBlock()] = Sym;
+ }
+ }
+
+ // Build target symbol → GOT entry symbol map.
+ // GOT entries are pointer-sized blocks with exactly one Pointer64 edge.
+ llvm::DenseMap<Symbol*, Symbol*> SymToGOTSym;
+ for (auto* B : G.blocks()) {
+ if (B->getSize() != G.getPointerSize()) continue;
+ if (B->edges_size() != 1) continue;
+ auto& E = *B->edges().begin();
+ if (E.getKind() == x86_64::Pointer64) {
+ auto It = BlockToSym.find(B);
+ if (It != BlockToSym.end()) {
+ SymToGOTSym[&E.getTarget()] = It->second;
+ }
+ }
+ }
+
+ for (auto* B : G.blocks()) {
+ for (auto& E : B->edges()) {
+ if (E.getKind() != x86_64::Pointer32) continue;
+ if (E.getOffset() < 2) continue;
+
+ auto MutableContent = B->getMutableContent(G);
+ auto* FixupData = reinterpret_cast<uint8_t*>(MutableContent.data()) + E.getOffset();
+ uint8_t Prev2 = FixupData[-2];
+ uint8_t Prev1 = FixupData[-1];
+
+ bool isRelaxedCall = (Prev2 == 0x67 && Prev1 == 0xe8);
+ bool isRelaxedJmp = (Prev1 == 0xe9);
+ if (!isRelaxedCall && !isRelaxedJmp) continue;
+
+ // Check if PC-relative displacement would fit.
+ auto TargetAddr = E.getTarget().getAddress();
+ auto FixupAddr = B->getFixupAddress(E);
+ int64_t Displacement = TargetAddr.getValue() - (FixupAddr.getValue() + 4) + E.getAddend();
+ if (llvm::isInt<32>(Displacement)) {
+ E.setKind(x86_64::BranchPCRel32);
+ continue;
+ }
+
+ // Distance doesn't fit — revert to indirect call/jmp through GOT.
+ auto It = SymToGOTSym.find(&E.getTarget());
+ if (It == SymToGOTSym.end()) {
+ return llvm::make_error<llvm::StringError>(
+ "Cannot revert GOTPCRELX relaxation: no GOT entry for " +
+ (E.getTarget().hasName() ? std::string(*E.getTarget().getName())
+ : std::string("<anon>")),
+ llvm::inconvertibleErrorCode());
+ }
+
+ Symbol* GOTSym = It->second;
+ if (isRelaxedCall) {
+ // Restore: 67 e8 → ff 15 (call *[rip+disp32])
+ FixupData[-2] = 0xff;
+ FixupData[-1] = 0x15;
+ } else {
+ // Restore: e9 XX XX XX XX 90 → ff 25 XX XX XX XX
+ FixupData[-1] = 0xff;
+ FixupData[0] = 0x25;
+ // For jmp, the optimization shifted offset by -1; shift back.
+ E.setOffset(E.getOffset() + 1);
+ }
+ E.setKind(x86_64::PCRel32);
+ E.setTarget(*GOTSym);
+ E.setAddend(0);
+ }
+ }
+ return llvm::Error::success();
+}
+
+} // namespace
+
+void GOTPCRELXFixPlugin::modifyPassConfig(llvm::orc::MaterializationResponsibility& MR,
+ llvm::jitlink::LinkGraph& G,
+ llvm::jitlink::PassConfiguration& Config) {
+ Config.PreFixupPasses.emplace_back(fixBrokenGOTPCRELXRelaxation);
+}
+
+llvm::Error GOTPCRELXFixPlugin::notifyFailed(llvm::orc::MaterializationResponsibility& MR) {
+ return llvm::Error::success();
+}
+
+llvm::Error GOTPCRELXFixPlugin::notifyRemovingResources(llvm::orc::JITDylib& JD,
+ llvm::orc::ResourceKey K) {
+ return llvm::Error::success();
+}
+
+void GOTPCRELXFixPlugin::notifyTransferringResources(llvm::orc::JITDylib& JD,
+ llvm::orc::ResourceKey DstKey,
+ llvm::orc::ResourceKey SrcKey) {}
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // __linux__ && __x86_64__
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.h
new file mode 100644
index 0000000..2880ca6
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/gotpcrelx_fix.h
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file gotpcrelx_fix.h
+ * \brief LLVM JITLink GOTPCRELX relaxation bug workaround (x86_64).
+ *
+ * LLVM issue: not yet filed, internal TODO.
+ * Affected versions: observed on LLVM 20.x and 21.x; likely all versions
+ * since `optimizeGOTAndStubAccesses` landed in
+ * JITLink/x86_64.cpp.
+ * Trigger: x86_64 JITLink with external symbols whose resolved addresses
+ * fit in uint32 (e.g. libc PLT entries in a non-PIE executable,
+ * or any low-VA process symbol) while JIT code is at high
+ * addresses (as produced by our arena memory manager).
+ * Symptom: `optimizeGOTAndStubAccesses` relaxes
+ * `call *foo@GOTPCREL(%rip)` (ff 15)
+ * into
+ * `addr32 call foo` (67 e8)
+ * and sets the edge kind to `Pointer32` (absolute 32-bit). But
+ * `call rel32` is always PC-relative, so the absolute fixup
+ * produces a garbage displacement. Result: SIGSEGV at JIT
+ * execution or during ORC-runtime teardown.
+ *
+ * The `GOTPCRELXFixPlugin` registers a PreFixupPass that runs *after*
+ * `optimizeGOTAndStubAccesses`, detects `Pointer32` edges preceded by
+ * `67 e8` / `e9` bytes, and either converts them to `BranchPCRel32` (if
+ * the PC-relative displacement fits in int32) or reverts the relaxation
+ * to an indirect call/jmp through the GOT (`ff 15` / `ff 25`, edge kind
+ * `PCRel32`, addend 0).
+ *
+ * ## Removal
+ *
+ * When the upstream fix lands and the project's minimum LLVM version
+ * bumps past the first release containing it, delete this file and
+ * remove:
+ * - the `#include "llvm_patches/gotpcrelx_fix.h"` in orcjit_session.cc
+ * - the `OLL->addPlugin(std::make_unique<GOTPCRELXFixPlugin>())` call
+ * inside the `setObjectLinkingLayerCreator` lambda in orcjit_session.cc.
+ * - the `llvm_patches/gotpcrelx_fix.cc` entry in
+ * addons/tvm_ffi_orcjit/CMakeLists.txt.
+ */
+#ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_
+#define TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_
+
+#if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64))
+
+#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+/*! \brief PreFixupPass plugin that corrects broken GOTPCRELX
+ * relaxations produced by JITLink's `optimizeGOTAndStubAccesses`.
+ *
+ * See the file-level docstring above for the trigger, symptom, and
+ * removal procedure.
+ */
+class GOTPCRELXFixPlugin : public llvm::orc::ObjectLinkingLayer::Plugin {
+ public:
+ void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G,
+ llvm::jitlink::PassConfiguration& Config) override;
+ llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override;
+ llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override;
+ void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey,
+ llvm::orc::ResourceKey SrcKey) override;
+};
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // __linux__ && __x86_64__
+
+#endif // TVM_FFI_ORCJIT_LLVM_PATCHES_GOTPCRELX_FIX_H_
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.cc
new file mode 100644
index 0000000..9d44ffb
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.cc
@@ -0,0 +1,258 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file init_fini_plugin.cc
+ * \brief Init/fini section handling for ELF, MachO, and COFF JIT objects.
+ *
+ * See init_fini_plugin.h for the trigger, symptom, and removal
+ * procedure for each platform.
+ */
+#include "init_fini_plugin.h"
+
+#include <llvm/ExecutionEngine/JITLink/JITLink.h>
+#include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+void InitFiniPlugin::modifyPassConfig(llvm::orc::MaterializationResponsibility& MR,
+ llvm::jitlink::LinkGraph& G,
+ llvm::jitlink::PassConfiguration& Config) {
+ auto& jit_dylib = MR.getTargetJITDylib();
+ // Mark all init/fini section blocks and their edge targets as live
+ // so they survive dead-stripping.
+ Config.PrePrunePasses.emplace_back([](llvm::jitlink::LinkGraph& G) {
+ for (auto& Section : G.sections()) {
+ auto section_name = Section.getName();
+ // ELF: .init_array*, .fini_array*, .ctors*, .dtors*
+ // Mach-O: __DATA,__mod_init_func, __DATA,__mod_term_func
+ // COFF: .CRT$XC* (ctors), .CRT$XT* (dtors)
+ if (section_name.starts_with(".init_array") || section_name.starts_with(".fini_array") ||
+ section_name.starts_with(".ctors") || section_name.starts_with(".dtors") ||
+ section_name == "__DATA,__mod_init_func" || section_name == "__DATA,__mod_term_func" ||
+ section_name.starts_with(".CRT$XC") || section_name.starts_with(".CRT$XT")) {
+ for (auto* Block : Section.blocks()) {
+ bool has_live_sym = false;
+ for (auto* Sym : G.defined_symbols()) {
+ if (&Sym->getBlock() == Block) {
+ Sym->setLive(true);
+ has_live_sym = true;
+ }
+ }
+ // MSVC may emit .CRT$XC* blocks with data but no symbol table
+ // entries (static variables in __declspec(allocate) sections).
+ // Add an anonymous symbol so the block survives dead-stripping.
+ if (!has_live_sym) {
+ G.addAnonymousSymbol(*Block, 0, Block->getSize(), false, true).setLive(true);
+ }
+ for (auto& Edge : Block->edges()) {
+ Edge.getTarget().setLive(true);
+ }
+ }
+ }
+ }
+ return llvm::Error::success();
+ });
+#ifdef _WIN32
+ // Without COFFPlatform, __ImageBase (used by IMAGE_REL_AMD64_ADDR32NB /
+ // Pointer32NB relocations) defaults to 0. This causes all Pointer32 fixups
+ // to overflow since JIT addresses don't fit in 32 bits.
+ //
+ // Fix: set __ImageBase to the lowest block address in the graph after
+ // allocation. This makes all intra-graph Pointer32NB offsets small.
+ //
+ // Also strip .pdata/.xdata edges: SEH unwind data references external
+ // handlers (e.g., __CxxFrameHandler3) in DLLs that may be >4GB from
+ // __ImageBase. Since we don't call RtlAddFunctionTable, SEH data is
+ // unused anyway.
+ Config.PostAllocationPasses.emplace_back([](llvm::jitlink::LinkGraph& G) {
+ // Set __ImageBase to the lowest allocated block address.
+ auto ImageBaseName = G.intern("__ImageBase");
+ llvm::jitlink::Symbol* ImageBase = nullptr;
+ for (auto* Sym : G.external_symbols()) {
+ if (Sym->getName() == ImageBaseName) {
+ ImageBase = Sym;
+ break;
+ }
+ }
+ if (ImageBase) {
+ llvm::orc::ExecutorAddr BaseAddr;
+ for (auto* B : G.blocks()) {
+ if (!BaseAddr || B->getAddress() < BaseAddr) {
+ BaseAddr = B->getAddress();
+ }
+ }
+ ImageBase->getAddressable().setAddress(BaseAddr);
+ }
+ // Strip .pdata/.xdata edges: external handlers may be >4GB from
+ // __ImageBase, and we don't register SEH data anyway.
+ for (auto& Sec : G.sections()) {
+ if (Sec.getName() == ".pdata" || Sec.getName().starts_with(".xdata")) {
+ for (auto* B : Sec.blocks()) {
+ while (!B->edges_empty()) {
+ B->removeEdge(B->edges().begin());
+ }
+ }
+ }
+ }
+ return llvm::Error::success();
+ });
+#endif
+ // After fixups, read resolved function pointers from all init/fini data sections.
+ // Handles ELF (.init_array, .ctors, .fini_array, .dtors),
+ // Mach-O (__DATA,__mod_init_func, __DATA,__mod_term_func),
+ // and COFF (.CRT$XC*, .CRT$XT*) section conventions.
+ Config.PostFixupPasses.emplace_back([this, &jit_dylib](llvm::jitlink::LinkGraph& G) {
+ using Entry = ORCJITExecutionSessionObj::InitFiniEntry;
+ for (auto& Sec : G.sections()) {
+ auto section_name = Sec.getName();
+
+ // --- ELF sections ---
+ bool is_init_array = section_name.starts_with(".init_array");
+ bool is_ctors = section_name.starts_with(".ctors");
+ bool is_fini_array = section_name.starts_with(".fini_array");
+ bool is_dtors = section_name.starts_with(".dtors");
+ // --- Mach-O sections ---
+ bool is_mod_init = (section_name == "__DATA,__mod_init_func");
+ bool is_mod_term = (section_name == "__DATA,__mod_term_func");
+ // --- COFF sections ---
+ bool is_crt_xc = section_name.starts_with(".CRT$XC");
+ bool is_crt_xt = section_name.starts_with(".CRT$XT");
+
+ if (!is_init_array && !is_ctors && !is_fini_array && !is_dtors && !is_mod_init &&
+ !is_mod_term && !is_crt_xc && !is_crt_xt)
+ continue;
+
+ int priority = 0;
+ Entry::Section sec;
+ bool is_init;
+ // ELF default priority for sections without a numeric suffix is 65535.
+ // Lower priority numbers run first for .init_array; .fini_array and .ctors
+ // negate so that higher-numbered entries run first (reverse order).
+ if (is_init_array) {
+ if (section_name.consume_front(".init_array.")) {
+ section_name.getAsInteger(10, priority);
+ } else {
+ priority = 65535;
+ }
+ sec = Entry::Section::kInitArray;
+ is_init = true;
+ } else if (is_ctors) {
+ if (section_name.consume_front(".ctors.") && !section_name.getAsInteger(10, priority)) {
+ priority = -priority;
+ }
+ sec = Entry::Section::kCtors;
+ is_init = true;
+ } else if (is_fini_array) {
+ if (section_name.consume_front(".fini_array.") &&
+ !section_name.getAsInteger(10, priority)) {
+ priority = -priority;
+ } else {
+ priority = -65535;
+ }
+ sec = Entry::Section::kFiniArray;
+ is_init = false;
+ } else if (is_dtors) {
+ if (section_name.consume_front(".dtors.")) {
+ section_name.getAsInteger(10, priority);
+ }
+ sec = Entry::Section::kDtors;
+ is_init = false;
+ } else if (is_mod_init) {
+ // Mach-O __mod_init_func: no priority system, treated as init_array
+ sec = Entry::Section::kInitArray;
+ is_init = true;
+ } else if (is_mod_term) {
+ // Mach-O __mod_term_func: no priority system, treated as fini_array
+ sec = Entry::Section::kFiniArray;
+ is_init = false;
+ } else if (is_crt_xc) {
+ // COFF .CRT$XC[suffix]: C++ constructors, sorted alphabetically by suffix.
+ // Convert suffix to integer priority that preserves alphabetical ordering.
+ // E.g., .CRT$XCA → 'A'*100000=6500000, .CRT$XCU → 'U'*100000=8500000,
+ // .CRT$XCT00200 → 'T'*100000+200=8400200
+ sec = Entry::Section::kInitArray;
+ is_init = true;
+ auto suffix = section_name.substr(7); // after ".CRT$XC"
+ if (!suffix.empty()) {
+ priority = static_cast<int>(suffix[0]) * 100000;
+ if (suffix.size() > 1) {
+ int num = 0;
+ suffix.substr(1).getAsInteger(10, num);
+ priority += num;
+ }
+ }
+ } else {
+ // COFF .CRT$XT[suffix]: C++ destructors, same suffix-to-priority scheme.
+ sec = Entry::Section::kFiniArray;
+ is_init = false;
+ auto suffix = section_name.substr(7); // after ".CRT$XT"
+ if (!suffix.empty()) {
+ priority = static_cast<int>(suffix[0]) * 100000;
+ if (suffix.size() > 1) {
+ int num = 0;
+ suffix.substr(1).getAsInteger(10, num);
+ priority += num;
+ }
+ }
+ }
+
+ for (auto* Block : Sec.blocks()) {
+ auto Content = Block->getContent();
+ size_t PtrSize = G.getPointerSize();
+ for (size_t Offset = 0; Offset + PtrSize <= Content.size(); Offset += PtrSize) {
+ uint64_t FnAddr = 0;
+ memcpy(&FnAddr, Content.data() + Offset, PtrSize);
+ if (FnAddr != 0) {
+ Entry entry{llvm::orc::ExecutorAddr(FnAddr), sec, priority};
+ if (is_init) {
+ session_->AddPendingInitializer(&jit_dylib, entry);
+ } else {
+ session_->AddPendingDeinitializer(&jit_dylib, entry);
+ }
+ }
+ }
+ }
+ }
+ return llvm::Error::success();
+ });
+}
+
+llvm::Error InitFiniPlugin::notifyFailed(llvm::orc::MaterializationResponsibility& MR) {
+ return llvm::Error::success();
+}
+
+llvm::Error InitFiniPlugin::notifyRemovingResources(llvm::orc::JITDylib& JD,
+ llvm::orc::ResourceKey K) {
+ return llvm::Error::success();
+}
+
+void InitFiniPlugin::notifyTransferringResources(llvm::orc::JITDylib& JD,
+ llvm::orc::ResourceKey DstKey,
+ llvm::orc::ResourceKey SrcKey) {}
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h
new file mode 100644
index 0000000..21262e7
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/init_fini_plugin.h
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file init_fini_plugin.h
+ * \brief Init/fini section handling for ELF, MachO, and COFF JIT objects.
+ *
+ * Emulates the missing/broken/bypassed LLVM ORC platform support for
+ * init/fini sections on all three host platforms. Collects function
+ * pointers from `.init_array` / `.fini_array` / `.ctors` / `.dtors`
+ * (ELF), `__DATA,__mod_init_func` / `__DATA,__mod_term_func` (MachO),
+ * and `.CRT$XC*` / `.CRT$XT*` (COFF); ties them to the containing
+ * `JITDylib`; and runs them in priority order through
+ * `ORCJITExecutionSessionObj::Run{Pending{Init,De}initializers}`.
+ *
+ * On Windows the plugin additionally patches `__ImageBase` (set to the
+ * lowest block address so `IMAGE_REL_AMD64_ADDR32NB` fixups don't
+ * overflow) and strips `.pdata` / `.xdata` edges (SEH handlers live in
+ * DLLs > 4 GB from `__ImageBase` and are never registered with
+ * `RtlAddFunctionTable` anyway). Those pieces also disappear once
+ * `COFFPlatform` becomes usable.
+ *
+ * Trigger: any JIT module on any platform containing constructors,
+ * destructors, or `__attribute__((constructor))` /
+ * MSVC `#pragma init_seg` equivalents.
+ * Symptom without the patch:
+ * - ELF: constructors/destructors never run (`ELFNixPlatform`
+ * enumerates but does not invoke them before
+ * llvm/llvm-project#175981).
+ * - MachO: we skip `MachOPlatform` entirely to sidestep the
+ * compact-unwind 32-bit delta bug (see `orcjit_session.cc`), so no
+ * platform runs init/fini; this plugin is the only mechanism.
+ * - COFF: relocation overflow / unresolved-SEH crashes
+ * (`COFFPlatform` is not hooked up because its MSVC CRT symbol
+ * requirements cannot be satisfied).
+ *
+ * ## 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:
+ * - 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
+ * and follow the checklist in `llvm_patches/README.md`.
+ *
+ * ## Removal — macOS
+ *
+ * Tied to re-enabling `MachOPlatform`. That requires the compact-unwind
+ * per-graph `dso_base` fix (see `fix-machoplatform-libunwind-dso-base.patch`
+ * in the repo root) to land in our LLVM. Until then we skip MachOPlatform
+ * and this plugin handles `__mod_init_func` / `__mod_term_func`. When
+ * MachOPlatform is re-enabled, delete the MachO-section branches from
+ * `InitFiniPlugin::modifyPassConfig` and drop the macOS side of the
+ * `addPlugin` call in `orcjit_session.cc`.
+ *
+ * ## Removal — Windows
+ *
+ * LLVM status: `COFFPlatform` has been stalled for 2+ years because it
+ * requires MSVC CRT symbols (`_CxxThrowException`, RTTI vtables, ...)
+ * that LLVM's COFF ORC runtime cannot provide.
+ * When `COFFPlatform` becomes usable end-to-end with clang-cl / MSVC
+ * objects, replace this plugin's Windows usage with it and delete the
+ * COFF handling path (including the `__ImageBase` fixup and the
+ * `.pdata` / `.xdata` edge stripping).
+ */
+#ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_
+#define TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_
+
+#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
+
+#include "../orcjit_session.h"
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+/*! \brief Init/fini section collector and runner for ELF, MachO, and COFF.
+ *
+ * See the file-level docstring above for the three-platform strategy
+ * and the removal procedure for each platform.
+ */
+class InitFiniPlugin : public llvm::orc::ObjectLinkingLayer::Plugin {
+ // Store a raw pointer to avoid a reference cycle:
+ // Session → LLJIT → ObjectLinkingLayer → Plugin → Session
+ // The plugin's lifetime is bounded by the ObjectLinkingLayer which is
+ // owned by LLJIT which is owned by the session, so the pointer is always valid.
+ ORCJITExecutionSessionObj* session_;
+
+ public:
+ explicit InitFiniPlugin(ORCJITExecutionSessionObj* session) : session_(session) {}
+
+ void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G,
+ llvm::jitlink::PassConfiguration& Config) override;
+ llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override;
+ llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override;
+ void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey,
+ llvm::orc::ResourceKey SrcKey) override;
+};
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // TVM_FFI_ORCJIT_LLVM_PATCHES_INIT_FINI_PLUGIN_H_
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc
new file mode 100644
index 0000000..0e1597a
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.cc
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file macho_cxa_atexit_shim.cc
+ * \brief Per-JITDylib `__cxa_atexit` interposer implementation.
+ *
+ * See macho_cxa_atexit_shim.h for the trigger, symptom, and removal
+ * procedure.
+ */
+
+#include "macho_cxa_atexit_shim.h"
+
+#ifdef __APPLE__
+
+#include <llvm/ExecutionEngine/Orc/AbsoluteSymbols.h>
+#include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h>
+#include <llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h>
+#include <llvm/Support/Error.h>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+namespace {
+// TLS pointer to the currently-active dylib's records vector.
+// Each CxaAtexitRecordsScope saves the previous pointer and restores it
+// on exit, so nested scopes compose correctly across re-entrant init/fini.
+//
+// We cannot override ___dso_handle (LLJIT's Platform has already defined
+// it in every user JITDylib), so we don't rely on the `dso_handle` arg
+// passed to the shim. Instead, each ORCJITDynamicLibraryObj publishes its
+// own records vector via this TLS slot, scoped around any JIT entry point
+// that may run ctors / dtors. The shim pushes (fn, arg) into the
+// TLS-pointed vector, or silently drops if no scope is active (which
+// would be a stray call from outside any JIT execution — acceptable
+// degradation).
+thread_local CxaAtexitRecords* g_active_cxa_records = nullptr;
+
+extern "C" int tvm_ffi_cxa_atexit_shim(void (*fn)(void*), void* arg,
+ void* /*dso_handle*/) noexcept {
+ if (fn == nullptr || g_active_cxa_records == nullptr) return 0;
+ g_active_cxa_records->emplace_back(fn, arg);
+ return 0;
+}
+} // namespace
+
+CxaAtexitRecordsScope::CxaAtexitRecordsScope(CxaAtexitRecords* records)
+ : prev_(g_active_cxa_records) {
+ g_active_cxa_records = records;
+}
+CxaAtexitRecordsScope::~CxaAtexitRecordsScope() { g_active_cxa_records = prev_; }
+
+void 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))));
+}
+
+void DrainCxaAtexit(CxaAtexitRecords& records) {
+ CxaAtexitRecordsScope scope(&records);
+ while (!records.empty()) {
+ auto [fn, arg] = records.back();
+ records.pop_back();
+ fn(arg);
+ }
+}
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // __APPLE__
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h
new file mode 100644
index 0000000..8aee61f
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/macho_cxa_atexit_shim.h
@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file macho_cxa_atexit_shim.h
+ * \brief Per-JITDylib `__cxa_atexit` interposer for macOS JIT.
+ *
+ * We skip `MachOPlatform` entirely on macOS to sidestep the
+ * compact-unwind 32-bit-delta bug in JITLink's `CompactUnwindSupport`
+ * (see the analysis in orcjit_session.cc and
+ * fix-machoplatform-libunwind-dso-base.patch at the repo root). With
+ * no Platform in the picture, clang-lowered
+ * `__attribute__((destructor))` and C++ global dtors — which register
+ * through `__cxa_atexit(fn, arg, &__dso_handle)` during init — would
+ * fall through to libSystem's `___cxa_atexit`, orphaning those
+ * callbacks from our drop-time drain.
+ *
+ * This shim:
+ * 1. Installs an absolute symbol for `___cxa_atexit` on each user
+ * JITDylib that points at our own capture function
+ * (`InstallCxaAtexitShim`).
+ * 2. Publishes the owning dylib's `CxaAtexitRecords` vector in TLS
+ * (`CxaAtexitRecordsScope`) so the capture function knows where
+ * to push `(fn, arg)` pairs.
+ * 3. Drains the captured records LIFO at dylib destruction time
+ * (`DrainCxaAtexit`).
+ *
+ * Trigger: any macOS JIT object containing static destructors,
+ * `__attribute__((destructor))` functions, or C++ global
+ * objects with non-trivial dtors.
+ * Symptom without the shim: destructors never run; any resource
+ * held by a JIT global leaks for the lifetime of the host
+ * process.
+ *
+ * ## Removal
+ *
+ * Tied to re-enabling `MachOPlatform`. See the macOS removal notes
+ * in init_fini_plugin.h. When `MachOPlatform` is restored, delete
+ * this file, drop the `#include` in orcjit_session.cc / orcjit_dylib.h,
+ * and remove the `cxa_atexit_records_` field from `ORCJITDynamicLibraryObj`.
+ */
+#ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_
+#define TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_
+
+#ifdef __APPLE__
+
+#include <llvm/ExecutionEngine/Orc/Core.h>
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+/*! \brief `(dtor, arg)` pairs captured from per-dylib `__cxa_atexit` calls. */
+using CxaAtexitRecords = std::vector<std::pair<void (*)(void*), void*>>;
+
+/*! \brief RAII scope that publishes a per-dylib `CxaAtexitRecords` vector
+ * to the TLS slot consulted by the `___cxa_atexit` shim.
+ *
+ * The shim reads the TLS pointer and pushes `(fn, arg)` into the pointed-to
+ * vector; outside any active scope the shim silently drops registrations.
+ * Wrap any JIT entry point that can run ctors / dtors (e.g. `GetSymbol`'s
+ * initialize call, the dtor's drain loop) with one of these, constructed
+ * with `&cxa_atexit_records_` from the owning `ORCJITDynamicLibraryObj`.
+ */
+class CxaAtexitRecordsScope {
+ public:
+ explicit CxaAtexitRecordsScope(CxaAtexitRecords* records);
+ ~CxaAtexitRecordsScope();
+ CxaAtexitRecordsScope(const CxaAtexitRecordsScope&) = delete;
+ CxaAtexitRecordsScope& operator=(const CxaAtexitRecordsScope&) = delete;
+
+ private:
+ CxaAtexitRecords* prev_;
+};
+
+/*! \brief Install the `___cxa_atexit` absolute-symbol shim on \p jd.
+ *
+ * Must be called once per user JITDylib, before any JIT code on that dylib
+ * materializes. Placing the definition on the dylib itself (rather than
+ * injecting into the link order) ensures it wins over `<Platform>`'s
+ * libSystem fallback — JITDylib::define-time symbols are searched before
+ * the link order.
+ */
+void InstallCxaAtexitShim(llvm::orc::ExecutionSession& ES, llvm::orc::JITDylib& jd);
+
+/*! \brief Drain captured `(fn, arg)` records LIFO, running each dtor.
+ *
+ * Pop-and-call order handles re-entrant registrations from within a dtor —
+ * the internal `CxaAtexitRecordsScope` keeps the TLS pointer live so any
+ * `___cxa_atexit` call from inside a dtor also lands in \p records.
+ */
+void DrainCxaAtexit(CxaAtexitRecords& records);
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // __APPLE__
+
+#endif // TVM_FFI_ORCJIT_LLVM_PATCHES_MACHO_CXA_ATEXIT_SHIM_H_
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.cc
new file mode 100644
index 0000000..270a1cb
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.cc
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file win_coff_pdata_strip.cc
+ * \brief COFF `.pdata` / `.xdata` relocation stripper implementation.
+ *
+ * See win_coff_pdata_strip.h for the trigger, symptom, and removal
+ * procedure.
+ */
+
+#include "win_coff_pdata_strip.h"
+
+#ifdef _WIN32
+
+#include <llvm/ADT/SmallVector.h>
+#include <llvm/ADT/StringRef.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>> StripCoffPdataXdata(
+ std::unique_ptr<llvm::MemoryBuffer> Buf) {
+ const char* Data = Buf->getBufferStart();
+ size_t Size = Buf->getBufferSize();
+ if (Size < 20) return std::move(Buf);
+
+ // Parse COFF header (regular or bigobj format)
+ uint16_t w0, w1;
+ std::memcpy(&w0, Data, 2);
+ std::memcpy(&w1, Data + 2, 2);
+ bool bigobj = (w0 == 0 && w1 == 0xFFFF);
+
+ uint16_t machine;
+ uint32_t num_sections, ptr_to_symtab, num_symbols;
+ size_t sec_hdr_start, sym_entry_size;
+ if (bigobj) {
+ if (Size < 56) return std::move(Buf);
+ std::memcpy(&machine, Data + 6, 2);
+ std::memcpy(&num_sections, Data + 44, 4);
+ std::memcpy(&ptr_to_symtab, Data + 48, 4);
+ std::memcpy(&num_symbols, Data + 52, 4);
+ sec_hdr_start = 56;
+ sym_entry_size = 20;
+ } else {
+ machine = w0;
+ uint16_t ns, opt_hdr_size;
+ std::memcpy(&ns, Data + 2, 2);
+ std::memcpy(&opt_hdr_size, Data + 16, 2);
+ std::memcpy(&ptr_to_symtab, Data + 8, 4);
+ std::memcpy(&num_symbols, Data + 12, 4);
+ num_sections = ns;
+ sec_hdr_start = 20 + opt_hdr_size;
+ sym_entry_size = 18;
+ }
+ if (machine != 0x8664) return std::move(Buf);
+
+ // String table follows the symbol table
+ size_t strtab_start = ptr_to_symtab + static_cast<size_t>(num_symbols) * sym_entry_size;
+
+ // Resolve a section name (inline 8-byte or "/offset" string table ref)
+ constexpr size_t kSecHdrSize = 40;
+ auto resolve_name = [&](size_t hdr_off) -> llvm::StringRef {
+ const char* raw = Data + hdr_off;
+ if (raw[0] == '/' && raw[1] >= '0' && raw[1] <= '9') {
+ uint32_t offset = 0;
+ for (int j = 1; j < 8 && raw[j] >= '0' && raw[j] <= '9'; ++j)
+ offset = offset * 10 + (raw[j] - '0');
+ size_t pos = strtab_start + offset;
+ if (pos < Size) {
+ size_t len = 0;
+ while (pos + len < Size && Data[pos + len]) ++len;
+ return {Data + pos, len};
+ }
+ }
+ size_t len = 0;
+ while (len < 8 && raw[len]) ++len;
+ return {raw, len};
+ };
+
+ // Collect section header offsets needing relocation stripping
+ llvm::SmallVector<size_t, 8> strip_offsets;
+ for (uint32_t i = 0; i < num_sections; ++i) {
+ size_t off = sec_hdr_start + i * kSecHdrSize;
+ if (off + kSecHdrSize > Size) break;
+ auto name = resolve_name(off);
+ if (name.starts_with(".pdata") || name.starts_with(".xdata")) {
+ uint16_t num_relocs;
+ std::memcpy(&num_relocs, Data + off + 32, 2);
+ if (num_relocs > 0) strip_offsets.push_back(off);
+ }
+ }
+ if (strip_offsets.empty()) return std::move(Buf);
+
+ // Create mutable copy, zero out PointerToRelocations and NumberOfRelocations
+ llvm::SmallVector<char> MutableBuf(Data, Data + Size);
+ for (auto off : strip_offsets) {
+ std::memset(&MutableBuf[off + 24], 0, 4); // PointerToRelocations
+ std::memset(&MutableBuf[off + 32], 0, 2); // NumberOfRelocations
+ }
+ return llvm::MemoryBuffer::getMemBufferCopy(llvm::StringRef(MutableBuf.data(), MutableBuf.size()),
+ Buf->getBufferIdentifier());
+}
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // _WIN32
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.h
new file mode 100644
index 0000000..6e87ae4
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_coff_pdata_strip.h
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file win_coff_pdata_strip.h
+ * \brief COFF `.pdata` / `.xdata` relocation stripper for Windows JIT.
+ *
+ * Installed as an `ObjTransformLayer` transform on Windows. clang-cl
+ * places static functions in COMDAT sections, and the `.pdata` SEH
+ * unwind data has relocations targeting COMDAT leader symbols.
+ * JITLink's `COFFLinkGraphBuilder` doesn't register COMDAT leaders in
+ * its symbol table when the second COMDAT symbol is `CLASS_STATIC`
+ * (not `CLASS_EXTERNAL`), causing "Could not find symbol" errors at
+ * graph-build time.
+ *
+ * `init_fini_plugin.cc` also strips `.pdata` / `.xdata` edges in a
+ * `PostAllocationPass`; this transform moves the stripping earlier
+ * (pre-graph-build) to prevent the graph builder error. Both pieces
+ * disappear once `COFFPlatform` becomes usable.
+ *
+ * Trigger: Windows x86_64 JIT of clang-cl objects that contain
+ * static-linkage functions (i.e. almost any C++ input).
+ * Symptom without the transform: JITLink graph-build failure with
+ * "Could not find symbol" referencing a COMDAT leader.
+ *
+ * We avoid `llvm/Object/COFF.h` because `windows.h` (included
+ * transitively by `LLJIT.h`) defines `IMAGE_*` macros that conflict
+ * with LLVM's COFF enums; we parse the COFF header with raw
+ * `memcpy`s instead.
+ *
+ * ## Removal
+ *
+ * Tied to enabling `COFFPlatform`. See the Windows removal notes in
+ * init_fini_plugin.h.
+ */
+#ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_
+#define TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_
+
+#ifdef _WIN32
+
+#include <llvm/Support/Error.h>
+#include <llvm/Support/MemoryBuffer.h>
+
+#include <memory>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+/*! \brief Zero out `PointerToRelocations` / `NumberOfRelocations` on
+ * `.pdata` and `.xdata` section headers before the COFF object
+ * reaches JITLink's graph builder.
+ *
+ * Returns the (possibly copy-modified) buffer. Safe to install as the
+ * `ObjTransformLayer` transform — invokes no LLVM object parsers and
+ * short-circuits for any non-x86_64 or too-small input.
+ */
+llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>> StripCoffPdataXdata(
+ std::unique_ptr<llvm::MemoryBuffer> Buf);
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // _WIN32
+
+#endif // TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_COFF_PDATA_STRIP_H_
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc
new file mode 100644
index 0000000..59f8600
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file win_dll_import_generator.cc
+ * \brief Windows DLL-import symbol generator implementation.
+ *
+ * See win_dll_import_generator.h for the trigger, symptom, and removal
+ * procedure.
+ */
+
+#include "win_dll_import_generator.h"
+
+#ifdef _WIN32
+
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN
+#endif
+
+#include <llvm/ADT/DenseMap.h>
+#include <llvm/ExecutionEngine/JITLink/JITLink.h>
+#include <llvm/ExecutionEngine/JITLink/x86_64.h>
+#include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h>
+#include <llvm/Support/DynamicLibrary.h>
+#include <llvm/Support/Error.h>
+#include <llvm/TargetParser/SubtargetFeature.h>
+
+// windows.h must precede psapi.h — psapi.h uses SIZE_T / DWORD typedefs
+// defined in windows.h. Left to its own devices clang-format sorts these
+// alphabetically, landing psapi.h first and breaking the MSVC build.
+// clang-format off
+#include <windows.h>
+#include <psapi.h>
+// clang-format on
+
+#include <string>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+void* DLLImportDefinitionGenerator::FindInProcessModules(const std::string& Name) {
+ // Try specific runtime DLLs first, then tvm_ffi.dll (loaded by Python),
+ // then all process modules, then LLVM's search.
+ static const char* kRuntimeDLLs[] = {
+ "vcruntime140.dll",
+ "vcruntime140_1.dll",
+ "ucrtbase.dll",
+ "msvcp140.dll",
+ };
+ // NOTE: We intentionally do not call FreeLibrary() here. These runtime DLLs
+ // (vcruntime140, ucrtbase, etc.) are already loaded by the process and will
+ // remain loaded for its lifetime. LoadLibraryA merely increments the refcount;
+ // the extra refcount is harmless and avoids the overhead of balancing
+ // Get/FreeLibrary for every symbol lookup.
+ for (const char* dll : kRuntimeDLLs) {
+ if (HMODULE hMod = LoadLibraryA(dll)) {
+ if (auto addr = GetProcAddress(hMod, Name.c_str())) {
+ return reinterpret_cast<void*>(addr);
+ }
+ }
+ }
+ // Also check tvm_ffi.dll (host process symbol provider)
+ if (HMODULE hTvmFfi = GetModuleHandleA("tvm_ffi.dll")) {
+ if (auto addr = GetProcAddress(hTvmFfi, Name.c_str())) {
+ return reinterpret_cast<void*>(addr);
+ }
+ }
+ HMODULE hMods[1024];
+ DWORD cbNeeded;
+ if (EnumProcessModules(GetCurrentProcess(), hMods, sizeof(hMods), &cbNeeded)) {
+ DWORD count = cbNeeded / sizeof(HMODULE);
+ if (count > 1024) count = 1024;
+ for (DWORD i = 0; i < count; ++i) {
+ if (auto addr = GetProcAddress(hMods[i], Name.c_str())) {
+ return reinterpret_cast<void*>(addr);
+ }
+ }
+ }
+ if (void* addr = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(Name)) {
+ return addr;
+ }
+ return nullptr;
+}
+
+llvm::Error DLLImportDefinitionGenerator::tryToGenerate(
+ llvm::orc::LookupState& LS, llvm::orc::LookupKind K, llvm::orc::JITDylib& JD,
+ llvm::orc::JITDylibLookupFlags JDLookupFlags, const llvm::orc::SymbolLookupSet& LookupSet) {
+ // Step 1: Collect unique base names (strip __imp_ prefix) and resolve addresses.
+ llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorAddr> Resolved;
+ for (auto& [Name, Flags] : LookupSet) {
+ llvm::StringRef NameStr = *Name;
+ std::string BaseName =
+ NameStr.starts_with("__imp_") ? NameStr.drop_front(6).str() : NameStr.str();
+ if (BaseName == "__ImageBase") continue;
+ auto InternedBase = ES_.intern(BaseName);
+ if (Resolved.count(InternedBase)) continue;
+ void* Addr = FindInProcessModules(BaseName);
+ if (Addr) {
+ Resolved[InternedBase] = llvm::orc::ExecutorAddr::fromPtr(Addr);
+ }
+ }
+ if (Resolved.empty()) return llvm::Error::success();
+
+ // Step 2: Build a LinkGraph with __imp_ pointers and PLT jump stubs.
+ auto G = std::make_unique<llvm::jitlink::LinkGraph>(
+ "<DLL_IMPORT_STUBS>", ES_.getSymbolStringPool(), ES_.getTargetTriple(),
+ llvm::SubtargetFeatures(), llvm::jitlink::getGenericEdgeKindName);
+ auto Prot = static_cast<llvm::orc::MemProt>(static_cast<unsigned>(llvm::orc::MemProt::Read) |
+ static_cast<unsigned>(llvm::orc::MemProt::Exec));
+ auto& Sec = G->createSection("__dll_stubs", Prot);
+
+ for (auto& [InternedName, Addr] : Resolved) {
+ // Absolute symbol at the real address (local to this graph)
+ auto& Target = G->addAbsoluteSymbol(G->intern(("__real_" + *InternedName).str()), Addr,
+ G->getPointerSize(), llvm::jitlink::Linkage::Strong,
+ llvm::jitlink::Scope::Local, false);
+ // __imp_XXX pointer (GOT-like entry)
+ auto& Ptr = llvm::jitlink::x86_64::createAnonymousPointer(*G, Sec, &Target);
+ Ptr.setName(G->intern(("__imp_" + *InternedName).str()));
+ Ptr.setLinkage(llvm::jitlink::Linkage::Strong);
+ Ptr.setScope(llvm::jitlink::Scope::Default);
+ // XXX jump stub (PLT-like entry) for direct calls
+ auto& StubBlock = llvm::jitlink::x86_64::createPointerJumpStubBlock(*G, Sec, Ptr);
+ G->addDefinedSymbol(StubBlock, 0, *InternedName, StubBlock.getSize(),
+ llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Default, true, false);
+ }
+ return L_.add(JD, std::move(G));
+}
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // _WIN32
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h
new file mode 100644
index 0000000..3b0d34c
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file win_dll_import_generator.h
+ * \brief Windows DLL-import symbol generator for JIT code (COFF/x86_64).
+ *
+ * On Windows with the MSVC ABI, COFF objects reference DLL-imported
+ * symbols via `__imp_XXX` pointer stubs and direct calls. Without
+ * `COFFPlatform` (which we skip because its MSVC CRT dependencies
+ * cannot be satisfied — see init_fini_plugin.h for the Windows
+ * removal notes), JITLink has no way to resolve those references.
+ *
+ * `DLLImportDefinitionGenerator` is a `DefinitionGenerator` that, for
+ * each requested symbol:
+ * 1. Looks up the real address by walking the MSVC runtime DLLs,
+ * `tvm_ffi.dll`, all loaded process modules, and finally LLVM's
+ * `SearchForAddressOfSymbol`.
+ * 2. Emits a fresh `LinkGraph` containing a GOT-like `__imp_XXX`
+ * pointer holding the real address, plus a PLT-like jump stub
+ * (`jmp [__imp_XXX]`) so direct calls stay in range. All stubs
+ * land in JIT memory, keeping every PCRel32 fixup within ±2 GB
+ * of the JIT arena.
+ *
+ * Trigger: any Windows x86_64 JIT graph that references a DLL-imported
+ * symbol (which, for MSVC-compiled objects, is effectively
+ * any C runtime or libtvm_ffi entry point).
+ * Symptom without the generator: unresolved `__imp_XXX` externals
+ * at JITLink time, or PCRel32 overflow for direct DLL calls.
+ *
+ * ## Removal
+ *
+ * Tied to enabling `COFFPlatform`. When `COFFPlatform` becomes usable
+ * end-to-end with clang-cl / MSVC objects (blocked upstream on MSVC
+ * CRT symbol requirements), delete this file and the corresponding
+ * `addGenerator` call in `orcjit_session.cc`.
+ */
+#ifndef TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_
+#define TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_
+
+#ifdef _WIN32
+
+#include <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+/*! \brief JIT-allocated `__imp_XXX` pointer stubs + PLT jumps for DLL imports.
+ *
+ * See the file-level docstring for the trigger, symptom, and removal
+ * procedure.
+ */
+class DLLImportDefinitionGenerator : public llvm::orc::DefinitionGenerator {
+ public:
+ DLLImportDefinitionGenerator(llvm::orc::ExecutionSession& ES, llvm::orc::ObjectLinkingLayer& L)
+ : ES_(ES), L_(L) {}
+
+ llvm::Error tryToGenerate(llvm::orc::LookupState& LS, llvm::orc::LookupKind K,
+ llvm::orc::JITDylib& JD, llvm::orc::JITDylibLookupFlags JDLookupFlags,
+ const llvm::orc::SymbolLookupSet& LookupSet) override;
+
+ private:
+ static void* FindInProcessModules(const std::string& Name);
+
+ llvm::orc::ExecutionSession& ES_;
+ llvm::orc::ObjectLinkingLayer& L_;
+};
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // _WIN32
+
+#endif // TVM_FFI_ORCJIT_LLVM_PATCHES_WIN_DLL_IMPORT_GENERATOR_H_
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
index 11a6d52..8db8c78 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
@@ -45,6 +45,23 @@ namespace tvm {
namespace ffi {
namespace orcjit {
+namespace {
+// When JIT thunks start reading their ctx/handle arg, grow this wrapper with
+// a DylibFnContext prefix whose first field is the thunk pointer, and flip
+// `safe_call` at a slab-emitted redirect that dispatches via ctx[0]. Example:
+//
+// struct DylibFnContext { TVMFFISafeCallType fn; /* + closure fields */ };
+// struct DylibFnContextWithModule {
+// DylibFnContext ctx; // first — pointer-interconvertible with wrapper
+// Module module_ref;
+// };
+struct DylibFnContextWithModule {
+ Module module_ref; // keeps the owning dylib (and its slab) alive
+};
+
+void DeleteDylibFnContextWithModule(void* p) { delete static_cast<DylibFnContextWithModule*>(p); }
+} // namespace
+
ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session,
llvm::orc::JITDylib* dylib, llvm::orc::LLJIT* jit,
String name)
@@ -62,15 +79,29 @@ ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session,
}
ORCJITDynamicLibraryObj::~ORCJITDynamicLibraryObj() {
- // See InitFiniPlugin class doc in orcjit_session.cc for the three-platform
- // init/fini strategy (macOS native vs. Linux/Windows plugin).
-#if defined(__linux__) || defined(_WIN32)
+ // 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());
-#else
- if (auto err = jit_->deinitialize(*dylib_)) {
- llvm::consumeError(std::move(err));
- }
+#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.
+ 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_);
+ dylib_ = nullptr;
}
void ORCJITDynamicLibraryObj::AddObjectFile(const String& path) {
@@ -113,23 +144,19 @@ void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) {
auto symbol_or_err =
jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str()));
- // Run pending initializers. Both paths are idempotent and support incremental
- // loading (new object files added between GetSymbol calls).
- // See InitFiniPlugin class doc in orcjit_session.cc for three-platform strategy.
-#if defined(__linux__) || defined(_WIN32)
- // Linux/Windows: RunPendingInitializers drains and erases the map entry;
- // subsequent calls are no-ops until new object files add fresh entries.
- session_->RunPendingInitializers(GetJITDylib());
-#else
- // macOS: LLJIT tracks an InitializedDylib set (LLJIT.h:624). First call
- // invokes dlopen (refcount 0→1, runs initializers). Subsequent calls
- // invoke dlupdate (no refcount bump, runs only new initializers from
- // newly-added object files). The single deinitialize() in the destructor
- // calls dlclose (refcount 1→0), so the count stays balanced.
- if (auto err = jit_->initialize(*dylib_)) {
- llvm::consumeError(std::move(err));
- }
+ // 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).
+#ifdef __APPLE__
+ CxaAtexitRecordsScope scope(&cxa_atexit_records_);
#endif
+ session_->RunPendingInitializers(GetJITDylib());
// Convert ExecutorAddr to pointer
return symbol_or_err ? symbol_or_err->getAddress().toPtr<void*>() : nullptr;
}
@@ -159,16 +186,9 @@ Optional<Function> ORCJITDynamicLibraryObj::GetFunction(const String& name) {
// Try to get the symbol - return NullOpt if not found
if (void* symbol = GetSymbol(symbol_name)) {
- // Wrap C function pointer as tvm-ffi Function.
- // Capture a strong self-ref so the module (and its JIT code pages) stays
- // alive as long as the returned Function is alive.
TVMFFISafeCallType c_func = reinterpret_cast<TVMFFISafeCallType>(symbol);
- Module self_strong_ref = GetRef<Module>(this);
- return Function::FromPacked([c_func, self_strong_ref](PackedArgs args, Any* rv) {
- TVM_FFI_ICHECK_LT(rv->type_index(), ffi::TypeIndex::kTVMFFIStaticObjectBegin);
- TVM_FFI_CHECK_SAFE_CALL((*c_func)(nullptr, reinterpret_cast<const TVMFFIAny*>(args.data()),
- args.size(), reinterpret_cast<TVMFFIAny*>(rv)));
- });
+ auto* wrapper = new DylibFnContextWithModule{GetRef<Module>(this)};
+ return Function::FromExternC(wrapper, c_func, DeleteDylibFnContextWithModule);
}
return std::nullopt;
}
@@ -188,12 +208,16 @@ static void RegisterOrcJITFunctions() {
refl::GlobalDef()
.def("orcjit.ExecutionSession",
- [](const std::string& orc_rt_path, int64_t arena_size_bytes) {
- return ORCJITExecutionSession(orc_rt_path, arena_size_bytes);
+ [](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",
+ [](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 43bd688..0ab2bdc 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
@@ -30,6 +30,7 @@
#include <tvm/ffi/object.h>
#include <tvm/ffi/string.h>
+#include "llvm_patches/macho_cxa_atexit_shim.h"
#include "orcjit_session.h"
namespace tvm {
@@ -105,6 +106,18 @@ class ORCJITDynamicLibraryObj : public ModuleObj {
/*! \brief Link order tracking (to support incremental linking) */
llvm::orc::JITDylibSearchOrder link_order_;
+
+#ifdef __APPLE__
+ /*! \brief Per-dylib __cxa_atexit registry.
+ *
+ * Without MachOPlatform, clang-lowered \c __attribute__((destructor))
+ * and C++ global dtors register through \c __cxa_atexit during init.
+ * We interpose \c ___cxa_atexit per-JITDylib (see
+ * \c orcjit_session.cc); the shim pushes \c (fn, arg) into the vector
+ * published via \c CxaAtexitRecordsScope around each JIT entry. The
+ * destructor drains LIFO before \c RemoveDylib. */
+ CxaAtexitRecords cxa_atexit_records_;
+#endif // __APPLE__
};
/*!
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc
index 26be4b8..a970c97 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc
@@ -19,678 +19,154 @@
/*!
* \file orcjit_memory_manager.cc
- * \brief Arena-based JITLinkMemoryManager implementation.
- *
- * Follows the InProcessMemoryManager pattern from LLVM but replaces
- * per-object mmap with bump allocation from a pre-reserved arena.
- * Pages are committed in 2 MB slabs to enable Transparent Huge Page
- * (THP) promotion — see the class docstring in orcjit_memory_manager.h.
+ * \brief Growable per-session pool of `Slab`s.
*/
#include "orcjit_memory_manager.h"
#ifdef __linux__
-#include <llvm/ADT/DenseSet.h>
-#include <llvm/ADT/SmallVector.h>
-#include <llvm/ExecutionEngine/JITLink/JITLink.h>
-#include <llvm/ExecutionEngine/JITLink/aarch64.h>
-#include <llvm/ExecutionEngine/JITLink/x86_64.h>
-#include <llvm/ExecutionEngine/Orc/Shared/AllocationActions.h>
-#include <llvm/Support/Alignment.h>
+#include <llvm/Support/Error.h>
#include <llvm/Support/FormatVariadic.h>
-#include <llvm/Support/Memory.h>
-#include <sys/mman.h>
#include <algorithm>
-#include <cerrno>
-#include <cstdio>
-#include <cstring>
+#include <optional>
+#include <utility>
namespace tvm {
namespace ffi {
namespace orcjit {
-using namespace llvm;
-using namespace llvm::jitlink;
-using namespace llvm::orc;
-
-// ── Overflow section edge classification ───────────────────────────
-//
-// Conservative whitelist: only known absolute relocation kinds return true.
-// Unknown or future edge kinds default to PC-relative → sections stay in
-// the arena (safe: never breaks relocations, just forgoes the overflow
-// optimization for unknown kinds).
-
-namespace {
-
-bool isAbsoluteEdge(const Triple& TT, Edge::Kind K) {
- if (K < Edge::FirstRelocation) return true; // KeepAlive, Invalid — not a relocation constraint
- if (TT.isAArch64()) {
- using namespace llvm::jitlink::aarch64;
- switch (K) {
- case Pointer64:
- case Pointer32:
- case Pointer64Authenticated:
- case MoveWide16:
- return true;
- default:
- return false;
- }
- }
- if (TT.isX86()) {
- using namespace llvm::jitlink::x86_64;
- switch (K) {
- case Pointer64:
- case Pointer32:
- case Pointer32Signed:
- case Pointer16:
- case Pointer8:
- case Size64:
- case Size32:
- return true;
- default:
- return false;
- }
- }
- return false; // Unknown arch — treat as PC-relative (safe)
-}
-
-} // namespace
-
-// ── Platform abstraction ────────────────────────────────────────────
-
-void* ArenaJITLinkMemoryManager::reserveVA(size_t size) {
- void* p = ::mmap(nullptr, size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
- if (p == MAP_FAILED) return nullptr;
- return p;
-}
-
-void ArenaJITLinkMemoryManager::releaseVA(void* addr, size_t size) {
- int rc = ::munmap(addr, size);
- assert(rc == 0 && "munmap failed in arena destructor");
- (void)rc;
-}
-
-Error ArenaJITLinkMemoryManager::commitPages(void* addr, size_t size) {
- if (size == 0) return Error::success();
- // Commit at slab (2 MB) granularity for THP promotion.
- size_t offset = static_cast<char*>(addr) - arena_base_;
- size_t first_slab = offset / kSlabSize;
- size_t last_slab = (offset + size - 1) / kSlabSize;
-
- for (size_t i = first_slab; i <= last_slab; ++i) {
- if (slab_committed_[i].load(std::memory_order_acquire) != 0) continue;
- size_t slab_offset = i * kSlabSize;
- size_t slab_len = std::min(kSlabSize, arena_capacity_ - slab_offset);
- // mprotect is idempotent, so a concurrent racer calling it on the same slab
- // is harmless. Only flip the flag after success — otherwise a failed commit
- // followed by freeRegion() would leave slab_committed_[i] == 1, causing a
- // later allocation to skip mprotect and write into PROT_NONE memory.
- if (::mprotect(arena_base_ + slab_offset, slab_len, PROT_READ | PROT_WRITE) != 0) {
- return make_error<StringError>(
- "ArenaJITLinkMemoryManager: mprotect(RW) failed for slab at offset " +
- formatv("{0:x}", slab_offset) + ": " + std::strerror(errno),
- inconvertibleErrorCode());
- }
- slab_committed_[i].store(1, std::memory_order_release);
- }
- return Error::success();
-}
-
-void ArenaJITLinkMemoryManager::decommitPages(void* addr, size_t size) {
- // Intentionally a no-op for arena pages. The ORC runtime may still reference
- // deallocated JIT memory during session teardown (e.g., ELFNixPlatform
- // deinitializers run after some allocations are freed). Decommitting
- // (MADV_DONTNEED or mprotect PROT_NONE) would cause segfaults or illegal
- // instructions during shutdown.
- //
- // Physical pages stay committed but are returned to the free list for reuse.
- // The arena destructor releases all VA and physical memory via munmap.
- (void)addr;
- (void)size;
-}
-
-Error ArenaJITLinkMemoryManager::protectPages(void* addr, size_t size, MemProt Prot) {
- int prot = PROT_NONE;
- if ((Prot & MemProt::Read) != MemProt::None) prot |= PROT_READ;
- if ((Prot & MemProt::Write) != MemProt::None) prot |= PROT_WRITE;
- if ((Prot & MemProt::Exec) != MemProt::None) prot |= PROT_EXEC;
- if (::mprotect(addr, size, prot) != 0) {
- return make_error<StringError>("ArenaJITLinkMemoryManager: mprotect failed at " +
- formatv("{0:x}", addr) + " size " + formatv("{0:x}", size) +
- ": " + std::strerror(errno),
- inconvertibleErrorCode());
- }
- if ((Prot & MemProt::Exec) != MemProt::None) {
- sys::Memory::InvalidateInstructionCache(addr, size);
- }
- return Error::success();
-}
-
-// ── ArenaInFlightAlloc ──────────────────────────────────────────────
-
-class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManager::InFlightAlloc {
- public:
- // A contiguous region within one pool: [offset, offset + standard_size + finalize_size).
- // Standard-lifetime bytes come first; Finalize-lifetime bytes follow and are freed
- // at the end of finalize(). Any field may be 0 to indicate no allocation from
- // that pool on this call.
- struct PoolRegion {
- size_t offset;
- size_t standard_size;
- size_t finalize_size;
- };
-
- ArenaInFlightAlloc(ArenaJITLinkMemoryManager& MM, LinkGraph& G, BasicLayout BL,
- PoolRegion non_exec, PoolRegion exec,
- std::vector<OverflowBlock> overflow_blocks)
- : MM(MM),
- G(&G),
- BL(std::move(BL)),
- non_exec_(non_exec),
- exec_(exec),
- overflow_blocks_(std::move(overflow_blocks)) {}
-
- ~ArenaInFlightAlloc() override {
- assert(!G && "ArenaInFlightAlloc destroyed without finalize or abandon");
- }
-
- void finalize(OnFinalizedFunction OnFinalized) override {
- // Apply target protections for each arena segment.
- if (auto Err = applyProtections()) {
- OnFinalized(std::move(Err));
- return;
- }
-
- // Apply target protections for overflow blocks.
- for (auto& ob : overflow_blocks_) {
- if (auto Err = MM.protectPages(ob.addr, ob.size, ob.prot)) {
- OnFinalized(std::move(Err));
- return;
- }
- }
-
- // Run finalization actions (e.g., register EH frames).
- auto DeallocActions = shared::runFinalizeActions(BL.graphAllocActions());
- if (!DeallocActions) {
- OnFinalized(DeallocActions.takeError());
- return;
- }
-
- // Decommit finalize-lifetime pages in each pool — they're no longer needed.
- for (auto& R : {non_exec_, exec_}) {
- if (R.finalize_size > 0) {
- MM.decommitPages(MM.arena_base_ + R.offset + R.standard_size, R.finalize_size);
- MM.freeRegion(R.offset + R.standard_size, R.finalize_size);
- }
- }
-
-#ifndef NDEBUG
- G = nullptr;
-#endif
-
- // Create finalized allocation handle. LLVM's FinalizedAlloc stores an
- // opaque ExecutorAddr (integer), so we must use raw new here. Ownership
- // transfers to deallocate(), which LLVM guarantees is called for every
- // finalized allocation.
- auto* FA = new FinalizedAllocInfo{
- non_exec_.offset, non_exec_.standard_size, exec_.offset,
- exec_.standard_size, std::move(*DeallocActions), std::move(overflow_blocks_)};
- OnFinalized(FinalizedAlloc(ExecutorAddr::fromPtr(FA)));
- }
-
- void abandon(OnAbandonedFunction OnAbandoned) override {
- // Decommit and return each pool's full region to the appropriate free list.
- for (auto& R : {non_exec_, exec_}) {
- size_t total = R.standard_size + R.finalize_size;
- if (total > 0) {
- MM.decommitPages(MM.arena_base_ + R.offset, total);
- MM.freeRegion(R.offset, total);
- }
- }
-
- // Release overflow blocks.
- for (auto& ob : overflow_blocks_) {
- ::munmap(ob.addr, ob.size);
- }
-
-#ifndef NDEBUG
- G = nullptr;
-#endif
-
- OnAbandoned(Error::success());
- }
-
- private:
- Error applyProtections() {
- for (auto& KV : BL.segments()) {
- const auto& AG = KV.first;
- auto& Seg = KV.second;
-
- auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, MM.page_size_);
- if (auto Err = MM.protectPages(Seg.WorkingMem, SegSize, AG.getMemProt())) return Err;
- }
- return Error::success();
- }
-
- ArenaJITLinkMemoryManager& MM;
- LinkGraph* G;
- BasicLayout BL;
- PoolRegion non_exec_;
- PoolRegion exec_;
- std::vector<OverflowBlock> overflow_blocks_;
-};
-
-// ── ArenaJITLinkMemoryManager ───────────────────────────────────────
-
-ArenaJITLinkMemoryManager::ArenaJITLinkMemoryManager(size_t page_size, size_t arena_capacity)
- : arena_base_(nullptr),
- arena_capacity_(arena_capacity),
- page_size_(page_size),
- midpoint_(0),
- exec_bump_limit_(0),
- non_exec_bump_(0),
- exec_bump_(0) {
- // Try requested capacity, halve on failure down to a minimum floor.
- // The floor is the smaller of kMinArenaCapacity and the requested size,
- // so explicit small arenas (e.g. 16 MB for tests) are honoured.
- // mmap(PROT_NONE | MAP_NORESERVE) can still fail under RLIMIT_AS or
- // extreme VA fragmentation.
- size_t floor = std::min(arena_capacity_, kMinArenaCapacity);
- size_t cap = arena_capacity_;
+using llvm::Error;
+using llvm::Expected;
+
+SlabPoolMemoryManager::SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size)
+ : page_size_(page_size), slab_size_(slab_size) {
+ // Reserve the initial slab. Halving retry only applies here: if the
+ // very first mmap fails (RLIMIT_AS, container limits), we halve the
+ // requested size down to kMinSlabSize before giving up. Subsequent
+ // slabs added during allocate() use exactly slab_size_ and propagate
+ // errors on mmap failure.
+ std::size_t floor = std::min(slab_size_, kMinSlabSize);
+ std::size_t cap = slab_size_;
while (cap >= floor) {
- arena_base_ = static_cast<char*>(reserveVA(cap));
- if (arena_base_) {
- arena_capacity_ = cap;
- // Partition the arena into two pools at a 2 MB-aligned midpoint.
- // The exec pool starts at midpoint_, which is therefore on a 2 MB
- // boundary — r-x segments pack into a minimum number of 2 MB pages.
- //
- // Constraint: cross-pool displacements (e.g. .text → .rodata via
- // ADRP+ADD on aarch64) must fit in ±kPCRelReach. The farthest pair
- // of bytes is (end of exec, start of non-exec), separated by at most
- // `exec_bump_limit_`, so we cap the exec pool's upper bound at
- // kPCRelReach even when the VA reservation is larger.
- exec_bump_limit_ = std::min(cap, kPCRelReach);
- size_t raw_midpoint = static_cast<size_t>(exec_bump_limit_ * kDefaultNonExecFraction);
- midpoint_ = (raw_midpoint / kSlabSize) * kSlabSize;
- if (midpoint_ == 0) midpoint_ = kSlabSize;
- if (midpoint_ >= exec_bump_limit_) midpoint_ = exec_bump_limit_ - kSlabSize;
- non_exec_bump_ = 0;
- exec_bump_ = midpoint_;
- // Initialize slab commit tracking. make_unique<T[]>(n) value-initializes
- // the array to zero in C++17.
- num_slabs_ = (cap + kSlabSize - 1) / kSlabSize;
- slab_committed_ = std::make_unique<std::atomic<uint8_t>[]>(num_slabs_);
- // Hint THP promotion for the entire arena. Intentionally unchecked —
- // MADV_HUGEPAGE is advisory and may fail if THP is disabled system-wide.
- (void)::madvise(arena_base_, cap, MADV_HUGEPAGE);
+ auto slab = std::make_unique<Slab>(page_size_, cap);
+ if (slab->isValid()) {
+ // Pin the actual initial-slab size to whatever we succeeded with.
+ // If RLIMIT_AS forced us to 8 MB, we keep 8 MB as the working slab
+ // size; growing later at 64 MB would just fail again.
+ slab_size_ = cap;
+ slabs_.push_back(std::move(slab));
return;
}
cap /= 2;
}
- llvm::report_fatal_error("ArenaJITLinkMemoryManager: failed to reserve at least " +
- Twine(floor / (1024 * 1024)) + " MB of virtual address space");
+ llvm::report_fatal_error("SlabPoolMemoryManager: failed to reserve at least " +
+ llvm::Twine(floor / (1024 * 1024)) + " MB of virtual address space");
}
-ArenaJITLinkMemoryManager::~ArenaJITLinkMemoryManager() {
- if (arena_base_) {
- releaseVA(arena_base_, arena_capacity_);
- }
+std::unique_ptr<Slab> SlabPoolMemoryManager::createSlab(std::size_t capacity) {
+ auto slab = std::make_unique<Slab>(page_size_, capacity);
+ if (!slab->isValid()) return nullptr;
+ return slab;
}
-Expected<size_t> ArenaJITLinkMemoryManager::bumpAllocate(size_t size, bool is_exec) {
- std::lock_guard<std::mutex> Lock(mu_);
-
- auto& free_list = is_exec ? free_list_exec_ : free_list_non_exec_;
- auto& bump = is_exec ? exec_bump_ : non_exec_bump_;
- size_t limit = is_exec ? exec_bump_limit_ : midpoint_;
-
- // Try free list first (best-fit). O(n) scan — acceptable for the expected
- // workload of tens of JIT allocations, not thousands.
- size_t best_idx = free_list.size();
- size_t best_waste = std::numeric_limits<size_t>::max();
- for (size_t i = 0; i < free_list.size(); ++i) {
- if (free_list[i].size >= size && free_list[i].size - size < best_waste) {
- best_idx = i;
- best_waste = free_list[i].size - size;
- if (best_waste == 0) break;
- }
- }
-
- if (best_idx < free_list.size()) {
- size_t offset = free_list[best_idx].offset;
- if (free_list[best_idx].size == size) {
- free_list.erase(free_list.begin() + best_idx);
- } else {
- free_list[best_idx].offset += size;
- free_list[best_idx].size -= size;
- }
- return offset;
- }
-
- // Bump allocate within the pool's limit.
- if (bump + size > limit) {
- return make_error<StringError>(
- std::string("ArenaJITLinkMemoryManager: ") + (is_exec ? "exec" : "non-exec") +
- " pool exhausted (used " + formatv("{0:x}", bump).str() + " + requested " +
- formatv("{0:x}", size).str() + " > limit " + formatv("{0:x}", limit).str() + ")",
- inconvertibleErrorCode());
- }
-
- size_t offset = bump;
- bump += size;
- return offset;
-}
-
-void ArenaJITLinkMemoryManager::freeRegion(size_t offset, size_t size) {
- if (size == 0) return;
- std::lock_guard<std::mutex> Lock(mu_);
-
- // Route to the correct pool's free list based on offset.
- auto& free_list = (offset >= midpoint_) ? free_list_exec_ : free_list_non_exec_;
-
- // Insert into free list in sorted order.
- auto it = std::lower_bound(free_list.begin(), free_list.end(), offset,
- [](const FreeBlock& fb, size_t off) { return fb.offset < off; });
- it = free_list.insert(it, FreeBlock{offset, size});
-
- // Coalesce with next.
- auto next = it + 1;
- if (next != free_list.end() && it->offset + it->size == next->offset) {
- it->size += next->size;
- free_list.erase(next);
- }
-
- // Coalesce with previous.
- if (it != free_list.begin()) {
- auto prev = it - 1;
- if (prev->offset + prev->size == it->offset) {
- prev->size += it->size;
- free_list.erase(it);
- }
- }
-}
-
-void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
- OnAllocatedFunction OnAllocated) {
- // ── Overflow section classification ──
- //
- // Sections matching known overflow names (e.g. .nv_fatbin — large GPU
- // device blobs referenced only by absolute relocations) are allocated
- // outside the arena via separate mmap(), keeping the arena compact for
- // code + small rodata.
- //
- // Two-phase check:
- // Phase 1 — Name-based candidate selection (.nv_fatbin).
- // Phase 2 — Edge validation: any PC-relative cross-section edge
- // targeting a candidate section disqualifies it (the
- // section stays in the arena). This handles cases where
- // the compiler generates ADRP/RIP-relative refs even for
- // data sections.
+void SlabPoolMemoryManager::allocate(const llvm::jitlink::JITLinkDylib* /*JD*/,
+ llvm::jitlink::LinkGraph& G, OnAllocatedFunction OnAllocated) {
+ using AllocResult = Expected<std::unique_ptr<InFlightAlloc>>;
+
+ // Step 1: first-fit over existing slabs. `pool_mu_` only protects
+ // the slabs_ vector — never held across a Slab::allocate call or a
+ // user callback, since the LLJIT linker issues nested lookups (and
+ // thus re-entrant allocate() calls via materialization) from inside
+ // OnAllocated and a coarse lock would deadlock. Snapshot raw pointers
+ // under the lock; slabs are guaranteed to outlive this call because
+ // clearFreeSlabs() is only safe when the session is quiescent.
//
- // Validated candidates are temporarily set to NoAlloc so BasicLayout
- // skips them, then immediately restored before returning. By the time
- // JITLink's fixUpBlocks runs, sections are back to Standard — avoiding
- // the debug assert that prohibits edges from allocated sections to
- // NoAlloc sections.
- DenseSet<Section*> overflow_candidates;
- for (auto& Sec : G.sections()) {
- if (Sec.getMemLifetime() == MemLifetime::NoAlloc) continue;
- StringRef Name = Sec.getName();
- if (Name.starts_with(".nv_fatbin")) {
- overflow_candidates.insert(&Sec);
- }
- }
-
- // Phase 2: edge validation — disqualify candidates with incoming PC-relative edges.
- if (!overflow_candidates.empty()) {
- const auto& TT = G.getTargetTriple();
- for (auto& Sec : G.sections()) {
- for (auto* B : Sec.blocks()) {
- for (auto& E : B->edges()) {
- if (!E.isRelocation()) continue;
- if (isAbsoluteEdge(TT, E.getKind())) continue;
- // PC-relative edge — if it targets a candidate, disqualify.
- if (!E.getTarget().isDefined()) continue;
- auto* TargetSec = &E.getTarget().getBlock().getSection();
- overflow_candidates.erase(TargetSec);
- }
- }
- if (overflow_candidates.empty()) break;
- }
- }
-
- // Apply: temporarily hide validated overflow sections from BasicLayout.
- SmallVector<std::pair<Section*, MemLifetime>, 4> overflow_sections;
- for (auto* Sec : overflow_candidates) {
- overflow_sections.push_back({Sec, Sec->getMemLifetime()});
- Sec->setMemLifetime(MemLifetime::NoAlloc);
- }
-
- BasicLayout BL(G);
-
- // Restore overflow sections to their original lifetime immediately.
- // BasicLayout has already captured its segment list; subsequent LLVM
- // passes (fixUpBlocks) will see the sections as normal Standard sections.
- for (auto& [Sec, OrigLifetime] : overflow_sections) {
- Sec->setMemLifetime(OrigLifetime);
- }
-
- // Compute total sizes grouped by lifetime.
- auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(page_size_);
- if (!SegsSizes) {
- OnAllocated(SegsSizes.takeError());
- return;
- }
-
- if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
- OnAllocated(make_error<llvm::jitlink::JITLinkError>(
- "Total requested size " + formatv("{0:x}", SegsSizes->total()) + " for graph " +
- G.getName() + " exceeds address space"));
- return;
- }
-
- auto TotalSize = static_cast<size_t>(SegsSizes->total());
- if (TotalSize == 0 && overflow_sections.empty()) {
- // Empty graph — return a no-op allocation.
- OnAllocated(std::make_unique<ArenaInFlightAlloc>(
- *this, G, std::move(BL), ArenaInFlightAlloc::PoolRegion{0, 0, 0},
- ArenaInFlightAlloc::PoolRegion{midpoint_, 0, 0}, std::vector<OverflowBlock>{}));
- return;
- }
-
- // ── Dual-pool split ──
- //
- // Partition each segment into one of four buckets based on (Prot, Lifetime):
- // non-exec × Standard / Finalize → non-exec pool (below midpoint_)
- // exec × Standard / Finalize → exec pool (at/above midpoint_)
- //
- // Within each pool, Standard segments come first and Finalize segments
- // second, so the Finalize tail of each pool can be freed after finalize().
- size_t ne_std_size = 0, ne_fin_size = 0;
- size_t e_std_size = 0, e_fin_size = 0;
- for (auto& KV : BL.segments()) {
- auto& AG = KV.first;
- auto& Seg = KV.second;
- auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size_);
- bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None;
- bool is_finalize = AG.getMemLifetime() == MemLifetime::Finalize;
- if (is_exec) {
- (is_finalize ? e_fin_size : e_std_size) += SegSize;
- } else {
- (is_finalize ? ne_fin_size : ne_std_size) += SegSize;
- }
- }
- size_t ne_total = ne_std_size + ne_fin_size;
- size_t e_total = e_std_size + e_fin_size;
-
- ArenaInFlightAlloc::PoolRegion ne_region{0, 0, 0};
- ArenaInFlightAlloc::PoolRegion e_region{midpoint_, 0, 0};
-
- auto allocPool = [&](size_t req, bool is_exec) -> Expected<size_t> {
- if (req == 0) return size_t{0};
- auto off = bumpAllocate(req, is_exec);
- if (!off) return off.takeError();
- if (auto Err = commitPages(arena_base_ + *off, req)) {
- freeRegion(*off, req);
- return std::move(Err);
- }
- std::memset(arena_base_ + *off, 0, req);
- return *off;
- };
-
- if (ne_total > 0) {
- auto off = allocPool(ne_total, /*is_exec=*/false);
- if (!off) {
- OnAllocated(off.takeError());
+ // Slab::allocate is synchronous (invokes its callback inline on every
+ // code path), so a captured std::optional observes the result before
+ // the call returns.
+ std::vector<Slab*> snapshot;
+ {
+ std::lock_guard<std::mutex> lock(pool_mu_);
+ snapshot.reserve(slabs_.size());
+ for (auto& s : slabs_) snapshot.push_back(s.get());
+ }
+ for (Slab* slab : snapshot) {
+ std::optional<AllocResult> observed;
+ slab->allocate(G, [&](AllocResult R) { observed.emplace(std::move(R)); });
+ AllocResult result = std::move(*observed);
+ if (result) {
+ OnAllocated(std::move(result));
return;
}
- ne_region = {*off, ne_std_size, ne_fin_size};
- }
- if (e_total > 0) {
- auto off = allocPool(e_total, /*is_exec=*/true);
- if (!off) {
- // Unwind non-exec allocation on failure to keep the pools consistent.
- if (ne_total > 0) {
- decommitPages(arena_base_ + ne_region.offset, ne_total);
- freeRegion(ne_region.offset, ne_total);
- }
- OnAllocated(off.takeError());
- return;
+ Error E = result.takeError();
+ if (E.isA<SlabPoolExhaustedError>()) {
+ // Retriable: this graph didn't fit this slab's per-pool budget.
+ llvm::consumeError(std::move(E));
+ continue;
}
- e_region = {*off, e_std_size, e_fin_size};
- }
-
- // Assign addresses to segments from four cursors. Standard comes first in
- // each pool, then Finalize.
- auto NeStdCursor = ExecutorAddr::fromPtr(arena_base_ + ne_region.offset);
- auto NeFinCursor = ExecutorAddr::fromPtr(arena_base_ + ne_region.offset + ne_std_size);
- auto EStdCursor = ExecutorAddr::fromPtr(arena_base_ + e_region.offset);
- auto EFinCursor = ExecutorAddr::fromPtr(arena_base_ + e_region.offset + e_std_size);
-
- for (auto& KV : BL.segments()) {
- auto& AG = KV.first;
- auto& Seg = KV.second;
- bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None;
- bool is_finalize = AG.getMemLifetime() == MemLifetime::Finalize;
- auto& Cursor = is_exec ? (is_finalize ? EFinCursor : EStdCursor)
- : (is_finalize ? NeFinCursor : NeStdCursor);
- Seg.WorkingMem = Cursor.toPtr<char*>();
- Seg.Addr = Cursor;
- auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size_);
- Cursor += SegSize;
+ // Terminal (mmap, mprotect, JITLink, BasicLayout).
+ OnAllocated(std::move(E));
+ return;
}
- // Apply layout — copies content and assigns block addresses for arena segments.
- if (auto Err = BL.apply()) {
- // On error: decommit and free both pool regions.
- if (ne_total > 0) {
- decommitPages(arena_base_ + ne_region.offset, ne_total);
- freeRegion(ne_region.offset, ne_total);
- }
- if (e_total > 0) {
- decommitPages(arena_base_ + e_region.offset, e_total);
- freeRegion(e_region.offset, e_total);
- }
- OnAllocated(std::move(Err));
+ // Step 2: grow. Size the fresh slab to fit this graph's per-pool
+ // footprint — normal graphs fall through at slab_size_, skewed or
+ // oversize graphs double up until both pools can host them (see
+ // Slab::capacityForFootprint). A single growth branch replaces the
+ // pre-gate "normal vs oversize" split; we trade one `N × failed
+ // Slab::allocate` scan (step 1) for the extra pre-filter — negligible
+ // since each failed call is a BasicLayout build with no mmap.
+ auto fp = Slab::computeGraphFootprint(G, page_size_);
+ std::size_t cap = Slab::capacityForFootprint(fp, slab_size_);
+ auto slab = createSlab(cap);
+ if (!slab) {
+ OnAllocated(
+ llvm::make_error<llvm::StringError>("SlabPoolMemoryManager: mmap failed for new slab of " +
+ llvm::formatv("{0:x}", cap).str() + " bytes",
+ llvm::inconvertibleErrorCode()));
return;
}
-
- // ── Allocate overflow sections via mmap() outside the arena ──
- std::vector<OverflowBlock> overflow_allocs;
-
- for (auto& [Sec, _] : overflow_sections) {
- // Compute total size for this section's blocks.
- size_t total_sec_size = 0;
- for (auto* B : Sec->blocks()) {
- total_sec_size = alignTo(total_sec_size, B->getAlignment());
- total_sec_size += B->getSize();
- }
- if (total_sec_size == 0) continue;
- total_sec_size = alignTo(total_sec_size, page_size_);
-
- // mmap outside the arena.
- void* addr =
- ::mmap(nullptr, total_sec_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
- if (addr == MAP_FAILED) {
- // Clean up prior overflow allocs, free both pool regions, report error.
- for (auto& ob : overflow_allocs) ::munmap(ob.addr, ob.size);
- if (ne_total > 0) {
- decommitPages(arena_base_ + ne_region.offset, ne_total);
- freeRegion(ne_region.offset, ne_total);
- }
- if (e_total > 0) {
- decommitPages(arena_base_ + e_region.offset, e_total);
- freeRegion(e_region.offset, e_total);
- }
- OnAllocated(
- make_error<StringError>("ArenaJITLinkMemoryManager: overflow mmap failed for section " +
- Sec->getName() + ": " + std::strerror(errno),
- inconvertibleErrorCode()));
- return;
- }
-
- // Layout blocks within the mmap'd region.
- char* ptr = static_cast<char*>(addr);
- for (auto* B : Sec->blocks()) {
- uint64_t align = B->getAlignment();
- ptr = reinterpret_cast<char*>(alignTo(reinterpret_cast<uintptr_t>(ptr), align));
- size_t bsize = B->getSize();
- // Copy content and redirect block's mutable content pointer.
- if (!B->isZeroFill()) {
- auto content = B->getContent();
- std::memcpy(ptr, content.data(), content.size());
- B->setMutableContent(MutableArrayRef<char>(ptr, bsize));
- }
- // Assign block address (working mem == executor addr for in-process JIT).
- B->setAddress(ExecutorAddr::fromPtr(ptr));
- ptr += bsize;
- }
-
- overflow_allocs.push_back({addr, total_sec_size, Sec->getMemProt()});
+ Slab* raw = slab.get();
+ {
+ std::lock_guard<std::mutex> lock(pool_mu_);
+ slabs_.push_back(std::move(slab));
}
-
- OnAllocated(std::make_unique<ArenaInFlightAlloc>(*this, G, std::move(BL), ne_region, e_region,
- std::move(overflow_allocs)));
+ raw->allocate(G, std::move(OnAllocated));
}
-void ArenaJITLinkMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
- OnDeallocatedFunction OnDeallocated) {
+void SlabPoolMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
+ OnDeallocatedFunction OnDeallocated) {
Error DeallocErr = Error::success();
-
for (auto& Alloc : Allocs) {
- // Reclaim ownership of the FinalizedAllocInfo created in finalize().
auto* FA = Alloc.release().toPtr<FinalizedAllocInfo*>();
-
- // Run deallocation actions in reverse order.
- while (!FA->DeallocActions.empty()) {
- if (auto Err = FA->DeallocActions.back().runWithSPSRetErrorMerged())
- DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
- FA->DeallocActions.pop_back();
- }
-
- // Decommit and free each pool's Standard region.
- if (FA->non_exec_standard_size > 0) {
- decommitPages(arena_base_ + FA->non_exec_offset, FA->non_exec_standard_size);
- freeRegion(FA->non_exec_offset, FA->non_exec_standard_size);
- }
- if (FA->exec_standard_size > 0) {
- decommitPages(arena_base_ + FA->exec_offset, FA->exec_standard_size);
- freeRegion(FA->exec_offset, FA->exec_standard_size);
- }
-
- // Release overflow blocks.
- for (auto& ob : FA->overflow_blocks) {
- ::munmap(ob.addr, ob.size);
- }
-
+ FA->owner->deallocateOne(FA, DeallocErr);
delete FA;
}
-
OnDeallocated(std::move(DeallocErr));
}
+std::size_t SlabPoolMemoryManager::clearFreeSlabs() {
+ // Partition under the lock, move discards to a local vector, drop the
+ // lock, then let ~Slab (which calls munmap) run outside the lock.
+ // Keeping munmap outside pool_mu_ matches the rest of allocate/deallocate,
+ // which also never hold the lock across syscalls that might block.
+ std::vector<std::unique_ptr<Slab>> discard;
+ {
+ std::lock_guard<std::mutex> lock(pool_mu_);
+ auto keep_end = std::partition(slabs_.begin(), slabs_.end(),
+ [](const auto& s) { return !s->isReclaimable(); });
+ discard.reserve(static_cast<std::size_t>(slabs_.end() - keep_end));
+ for (auto it = keep_end; it != slabs_.end(); ++it) {
+ discard.push_back(std::move(*it));
+ }
+ slabs_.erase(keep_end, slabs_.end());
+ }
+ std::size_t reclaimed = discard.size();
+ // discard goes out of scope — Slab destructors munmap each reservation.
+ return reclaimed;
+}
+
} // namespace orcjit
} // namespace ffi
} // namespace tvm
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h
index 8f83c38..a71487d 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h
@@ -19,207 +19,116 @@
/*!
* \file orcjit_memory_manager.h
- * \brief Arena-based JITLinkMemoryManager for LLVM ORC JIT.
+ * \brief Per-session growable slab pool.
*
- * Pre-reserves a contiguous virtual address region and bump-allocates
- * from it, keeping all JIT allocations within range of PC-relative
- * relocations (±2GB on x86_64, ±4GB on AArch64).
+ * `SlabPoolMemoryManager` implements `JITLinkMemoryManager` on top of a
+ * per-session `std::vector<std::unique_ptr<Slab>>`. On each `allocate`
+ * it picks the first `Slab` that can fit the graph; if none do, it
+ * `mmap`s a fresh slab sized to fit (`Slab::capacityForFootprint`) and
+ * appends it. Normal-size graphs land on a `slab_size`-sized slab;
+ * skewed or oversize graphs land on a power-of-2 larger slab whose
+ * per-pool budgets cover the graph.
*
- * This eliminates relocation overflow caused by scattered mmap
- * allocations under ASLR (LLVM issue #173269).
+ * ## Lifecycle
*
- * ## GOTPCRELX relaxation workaround
- *
- * The arena triggers a latent bug in LLVM JITLink's
- * `optimizeGOTAndStubAccesses()` (x86_64.cpp). That pass relaxes
- * `call *foo@GOTPCREL(%rip)` (ff 15) → `addr32 call foo` (67 e8) and
- * sets the edge kind to `Pointer32` (absolute 32-bit address). However
- * the `call rel32` instruction is always **PC-relative** — the `67`
- * prefix is just padding — so the fixup should be PC-relative too
- * (matching the static linker's `R_X86_64_PC32`).
+ * Once a slab is added to the pool it stays mapped until it is
+ * reclaimed or until the pool (and its enclosing session) is
+ * destroyed. Individual graphs are deallocated via
+ * `FA->owner->deallocateOne(...)`, returning bytes to the slab's free
+ * list. Drained slabs can be returned to the OS via
+ * `clearFreeSlabs()`.
*
- * The bug is latent because the relaxation only fires when the target
- * address fits in 32 bits (`isUInt<32>`). On PIE executables every
- * resolved symbol is at a high address, so the guard is never true and
- * the relaxation never runs. On **non-PIE** executables the PLT
- * entries for libc functions (malloc, free, …) live near 0x400000, the
- * guard passes, and the wrong fixup produces a garbage displacement →
- * SIGSEGV during ORC-runtime teardown.
+ * ## GOTPCRELX relaxation workaround
*
- * `GOTPCRELXFixPlugin` in orcjit_session.cc works around this: a
- * PreFixupPass that runs *after* `optimizeGOTAndStubAccesses` detects
- * `Pointer32` edges on `67 e8` / `e9` instructions and either
- * (a) converts to `BranchPCRel32` when the PC-relative displacement
- * fits in int32, or
- * (b) reverts the relaxation entirely — restores the `ff 15` /
- * `ff 25` opcode bytes and retargets the edge to the GOT entry
- * with `PCRel32` + addend 0.
+ * See `llvm_patches/gotpcrelx_fix.cc`. The plugin is added per-session
+ * to the `ObjectLinkingLayer` alongside this memory manager and is
+ * orthogonal to pool growth.
*/
#ifndef TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_
#define TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_
#include <llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h>
-#include <llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h>
-#include <atomic>
+#include <cstddef>
#include <memory>
#include <mutex>
#include <vector>
+#include "orcjit_slab.h"
+
namespace tvm {
namespace ffi {
namespace orcjit {
-/*! \brief Arena-based memory manager for JITLink.
- *
- * Reserves a large contiguous VA region at construction time using
- * PROT_NONE (zero physical memory cost). Each allocate() call
- * bump-allocates from this region, commits pages as RW, and assigns
- * addresses. On finalization, pages receive their target protections.
- * On deallocation, pages are decommitted and returned to a free list.
- *
- * The default arena size is strictly larger than the architecture's
- * PC-relative relocation limit (4 GB on x86_64, 8 GB on AArch64) so
- * the arena is never the bottleneck — JITLink's own relocation overflow
- * checker fires first, matching dlopen/ld.so failure semantics. If the
- * initial reservation fails (RLIMIT_AS, container limits), the
- * constructor halves the capacity down to kMinArenaCapacity (256 MB).
- *
- * ## Slab-based commit with Transparent Huge Page (THP) support
- *
- * Arena pages are committed in 2 MB slabs (kSlabSize) rather than
- * per-allocation. Each slab is committed exactly once via an atomic
- * flag (lock-free, no contention with the allocator mutex).
+/*!
+ * \brief `JITLinkMemoryManager` backed by a growable pool of `Slab`s.
*
- * Benefits:
- * - Batches up to 512 page faults into a single sequential mprotect
- * per slab, reducing kernel trap overhead.
- * - 2 MB matches the Linux huge page size on both x86_64 and AArch64.
- * Combined with madvise(MADV_HUGEPAGE) applied at construction, the
- * kernel can promote each fully-faulted slab into a single TLB
- * entry (replacing 512 x 4 KB entries), reducing TLB misses during
- * JIT code execution.
- * - Worst-case waste is <2 MB in the last partially-used slab —
- * negligible for typical ML workloads.
+ * The constructor reserves one initial slab (halving its capacity down
+ * to `kMinSlabSize` if `mmap` fails under RLIMIT_AS). Subsequent
+ * slabs are added on demand by `allocate()` at a capacity chosen by
+ * `Slab::capacityForFootprint` — `slab_size_` for normal graphs, the
+ * next power of two up for skewed / oversize graphs. No retry, no
+ * halving on growth; mmap errors propagate.
*/
-class ArenaJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
+class SlabPoolMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
public:
- // Default arena: strictly larger than the relocation limit so the arena
- // is never the bottleneck. JITLink's own overflow check fires first,
- // matching dlopen/ld.so failure semantics.
- //
- // x86_64 PC32: ±2GB → 4GB default (2× headroom)
- // AArch64 ADRP: ±4GB → 8GB default (2× headroom)
- static constexpr size_t kDefaultArenaCapacity_x86_64 = size_t{4} << 30; // 4 GB
- static constexpr size_t kDefaultArenaCapacity_AArch64 = size_t{8} << 30; // 8 GB
- static constexpr size_t kMinArenaCapacity = size_t{256} << 20; // 256 MB floor
- // Slab commit granularity. Matches Linux huge page size (2 MB) on both
- // x86_64 and AArch64, enabling THP promotion via madvise(MADV_HUGEPAGE).
- static constexpr size_t kSlabSize = size_t{2} << 20; // 2 MB
- // PC-relative relocation reach (tightest binding fixup). Cross-pool
- // references (.text → .rodata, .eh_frame → .text, etc.) must fit within
- // this signed displacement. The binding constraint on both x86_64 and
- // aarch64 is the signed 32-bit Delta32 used in .eh_frame unwind records
- // (±2 GB), not the wider ADRP+ADD / RIP-rel reach. The dual-pool allocator
- // keeps both pools inside kPCRelReach bytes of each other even when the VA
- // reservation is larger, so cross-pool Delta32 fixups always resolve.
- static constexpr size_t kPCRelReach = (size_t{1} << 31) - kSlabSize; // ~2 GB
-
- // Default fraction of the arena reserved for non-exec segments (r--, rw-).
- // The remainder holds exec segments (r-x). Picked by splitting the arena
- // at a 2 MB-aligned boundary (midpoint_); the exec pool thus starts on a
- // 2 MB page boundary, maximizing r-x page packing.
- // Typical CUDA binding objects: ~2 parts rodata+data to 1 part text.
- static constexpr double kDefaultNonExecFraction = 2.0 / 3.0;
-
- explicit ArenaJITLinkMemoryManager(size_t page_size, size_t arena_capacity);
- ~ArenaJITLinkMemoryManager() override;
-
- ArenaJITLinkMemoryManager(const ArenaJITLinkMemoryManager&) = delete;
- ArenaJITLinkMemoryManager& operator=(const ArenaJITLinkMemoryManager&) = delete;
- ArenaJITLinkMemoryManager(ArenaJITLinkMemoryManager&&) = delete;
- ArenaJITLinkMemoryManager& operator=(ArenaJITLinkMemoryManager&&) = delete;
+ // Default per-slab capacity. 64 MB is above the p99 size of typical
+ // ML JIT graphs (single-kernel bindings, fused kernels), below the
+ // PC-relative relocation limit, and a multiple of the 2 MB THP
+ // granule. Small enough that a pinned slab only wastes 64 MB of RSS.
+ static constexpr std::size_t kDefaultSlabSize = std::size_t{64} << 20; // 64 MB
+
+ // Lower bound on initial-slab reservation. If the first `mmap`
+ // fails and halving drops below this, the constructor aborts.
+ // 8 MB is enough for a minimal JITDylib setup under very tight
+ // RLIMIT_AS.
+ static constexpr std::size_t kMinSlabSize = std::size_t{8} << 20; // 8 MB
+
+ explicit SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size);
+ ~SlabPoolMemoryManager() override = default;
+
+ SlabPoolMemoryManager(const SlabPoolMemoryManager&) = delete;
+ SlabPoolMemoryManager& operator=(const SlabPoolMemoryManager&) = delete;
+ SlabPoolMemoryManager(SlabPoolMemoryManager&&) = delete;
+ SlabPoolMemoryManager& operator=(SlabPoolMemoryManager&&) = delete;
void allocate(const llvm::jitlink::JITLinkDylib* JD, llvm::jitlink::LinkGraph& G,
OnAllocatedFunction OnAllocated) override;
void deallocate(std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) override;
- private:
- class ArenaInFlightAlloc;
+ /*! \brief Number of slabs currently held (test introspection). */
+ std::size_t numSlabs() const {
+ std::lock_guard<std::mutex> lock(pool_mu_);
+ return slabs_.size();
+ }
- /*! \brief A section allocated outside the arena via separate mmap().
+ /*!
+ * \brief Release drained slabs — zero live allocations and at least one
+ * prior allocation — back to the OS via `munmap`.
*
- * Sections whose only cross-section references use absolute relocations
- * (e.g. .nv_fatbin) are placed here to keep the arena compact. */
- struct OverflowBlock {
- void* addr; // mmap'd base address
- size_t size; // mmap'd size (page-aligned)
- llvm::orc::MemProt prot; // target protection for finalize
- };
-
- /*! \brief Metadata for a finalized allocation, stored via FinalizedAlloc handle.
+ * Returns the number of slabs reclaimed. Safe to call any time the
+ * session is quiescent (no concurrent JIT work in flight). A typical
+ * pattern is to call this after dropping a batch of libraries:
*
- * The arena is split into two pools at midpoint_. Each allocate() call may
- * consume a region from either or both pools. Standard-lifetime pages remain
- * committed after finalize(); Finalize-lifetime pages are decommitted at the
- * end of finalize(). Zero-sized sub-regions indicate no allocation from that
- * pool. */
- struct FinalizedAllocInfo {
- size_t non_exec_offset; // offset of non-exec Standard region (or 0 if unused)
- size_t non_exec_standard_size; // bytes retained in non-exec pool after finalize
- size_t exec_offset; // offset of exec Standard region (or midpoint_ if unused)
- size_t exec_standard_size; // bytes retained in exec pool after finalize
- std::vector<llvm::orc::shared::WrapperFunctionCall> DeallocActions;
- std::vector<OverflowBlock> overflow_blocks;
- };
-
- /*! \brief Bump-allocate from the selected pool. Returns offset within arena. */
- llvm::Expected<size_t> bumpAllocate(size_t size, bool is_exec);
-
- /*! \brief Return a region to the appropriate free list (coalesces adjacent blocks).
- * Pool is identified by comparing offset against midpoint_. */
- void freeRegion(size_t offset, size_t size);
-
- // ── Platform abstraction ──
- static void* reserveVA(size_t size);
- static void releaseVA(void* addr, size_t size);
- llvm::Error commitPages(void* addr, size_t size);
- static void decommitPages(void* addr, size_t size);
- static llvm::Error protectPages(void* addr, size_t size, llvm::orc::MemProt Prot);
-
- char* arena_base_;
- size_t arena_capacity_;
- size_t page_size_;
-
- // ── Dual-pool split ──
- // The arena is partitioned at midpoint_ (a 2 MB-aligned offset) into:
- // non-exec pool = [arena_base_, arena_base_ + midpoint_ )
- // exec pool = [arena_base_ + midpoint_, arena_base_ + exec_bump_limit_)
- // Both pools grow upward from their base. The exec pool starts on a 2 MB
- // boundary so r-x segments can pack as tightly as possible into 2 MB pages.
- //
- // exec_bump_limit_ = min(arena_capacity_, kPCRelReach). Bytes beyond this
- // limit stay reserved (VA only, no commit) but are not used for allocation
- // so cross-pool references always fit within the PC-relative reach.
- size_t midpoint_;
- size_t exec_bump_limit_;
-
- std::mutex mu_;
- size_t non_exec_bump_; // next free offset in non-exec pool ∈ [0, midpoint_]
- size_t exec_bump_; // next free offset in exec pool ∈ [midpoint_, arena_capacity_]
-
- struct FreeBlock {
- size_t offset;
- size_t size;
- };
- std::vector<FreeBlock> free_list_non_exec_;
- std::vector<FreeBlock> free_list_exec_;
-
- /*! \brief Per-slab commit flags (0 = uncommitted, 1 = committed).
- * Lock-free: each slab is committed exactly once via compare_exchange. */
- std::unique_ptr<std::atomic<uint8_t>[]> slab_committed_;
- size_t num_slabs_ = 0;
+ * for lib in libs: del lib
+ * session.clear_free_slabs() # Python API
+ *
+ * Fresh slabs that have never been allocated on are preserved — the
+ * session remains ready to accept new JIT work.
+ */
+ std::size_t clearFreeSlabs();
+
+ private:
+ /*! \brief Reserve a fresh slab at exactly \p capacity bytes. Returns
+ * nullptr on mmap failure (caller reports the error). */
+ std::unique_ptr<Slab> createSlab(std::size_t capacity);
+
+ std::size_t page_size_;
+ std::size_t slab_size_;
+
+ mutable std::mutex pool_mu_;
+ std::vector<std::unique_ptr<Slab>> slabs_;
};
} // namespace orcjit
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
index 0bca531..3470980 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
@@ -24,43 +24,35 @@
#include "orcjit_session.h"
-#include <llvm/ADT/DenseMap.h>
-#include <llvm/ExecutionEngine/JITLink/JITLink.h>
-#include <llvm/ExecutionEngine/JITLink/x86_64.h>
-#include <llvm/ExecutionEngine/Orc/AbsoluteSymbols.h>
#include <llvm/ExecutionEngine/Orc/LLJIT.h>
#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
-#include <llvm/ExecutionEngine/Orc/ObjectTransformLayer.h>
-#include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h>
-#include <llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h>
-#include <llvm/Support/DynamicLibrary.h>
#include <llvm/Support/Error.h>
#include <llvm/Support/Process.h>
#include <llvm/Support/TargetSelect.h>
-#include <llvm/TargetParser/SubtargetFeature.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/reflection/registry.h>
#include <cstddef>
-#include <cstring>
+#include "orcjit_dylib.h"
#include "orcjit_memory_manager.h"
+#include "orcjit_utils.h"
-#ifdef _WIN32
-#ifndef NOMINMAX
-#define NOMINMAX
-#endif
-#ifndef WIN32_LEAN_AND_MEAN
-#define WIN32_LEAN_AND_MEAN
+#if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64))
+#include "llvm_patches/gotpcrelx_fix.h"
#endif
-#include <psapi.h>
-#include <windows.h>
+#include "llvm_patches/init_fini_plugin.h"
+#ifdef __APPLE__
+#include "llvm_patches/macho_cxa_atexit_shim.h"
#endif
+#ifdef _WIN32
+#include <llvm/ExecutionEngine/Orc/ObjectTransformLayer.h>
-#include "orcjit_dylib.h"
-#include "orcjit_utils.h"
+#include "llvm_patches/win_coff_pdata_strip.h"
+#include "llvm_patches/win_dll_import_generator.h"
+#endif
namespace tvm {
namespace ffi {
@@ -77,537 +69,47 @@ struct LLVMInitializer {
static LLVMInitializer llvm_initializer;
-// Custom ObjectLinkingLayer plugin for init/fini section handling.
-//
-// Collects function pointers from init/fini sections (.init_array, .fini_array,
-// .ctors, .dtors, .CRT$XC*, .CRT$XT*) and runs them in priority order.
-//
-// Three-platform init/fini strategy:
-//
-// macOS: MachOPlatform (via orc_rt) handles __mod_init_func/__mod_term_func
-// and __cxa_atexit natively. We delegate to jit_->initialize() and
-// deinitialize(). No InitFiniPlugin needed.
-//
-// Windows: COFFPlatform is unusable -- it requires MSVC CRT symbols
-// (_CxxThrowException, RTTI vtables, etc.) that LLVM's COFF ORC
-// runtime cannot provide (stalled for 2+ years). Our plugin handles
-// .CRT$XC*/.CRT$XT* sections instead.
-//
-// Linux: ELFNixPlatform does not handle .init_array/.fini_array correctly
-// prior to llvm/llvm-project#175981. Once a new LLVM release includes
-// that patch, we can switch Linux to ELFNixPlatform and remove the
-// plugin for that platform.
-//
-// The plugin is compiled only on Linux/Windows and can be removed per-platform
-// as LLVM's native platform support matures.
-class InitFiniPlugin : public llvm::orc::ObjectLinkingLayer::Plugin {
- // Store a raw pointer to avoid a reference cycle:
- // Session → LLJIT → ObjectLinkingLayer → Plugin → Session
- // The plugin's lifetime is bounded by the ObjectLinkingLayer which is
- // owned by LLJIT which is owned by the session, so the pointer is always valid.
- ORCJITExecutionSessionObj* session_;
-
- public:
- explicit InitFiniPlugin(ORCJITExecutionSessionObj* session) : session_(session) {}
-
- void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G,
- llvm::jitlink::PassConfiguration& Config) override {
- auto& jit_dylib = MR.getTargetJITDylib();
- // Mark all init/fini section blocks and their edge targets as live
- // so they survive dead-stripping.
- Config.PrePrunePasses.emplace_back([](llvm::jitlink::LinkGraph& G) {
- for (auto& Section : G.sections()) {
- auto section_name = Section.getName();
- // ELF: .init_array*, .fini_array*, .ctors*, .dtors*
- // Mach-O: __DATA,__mod_init_func, __DATA,__mod_term_func
- // COFF: .CRT$XC* (ctors), .CRT$XT* (dtors)
- if (section_name.starts_with(".init_array") || section_name.starts_with(".fini_array") ||
- section_name.starts_with(".ctors") || section_name.starts_with(".dtors") ||
- section_name == "__DATA,__mod_init_func" || section_name == "__DATA,__mod_term_func" ||
- section_name.starts_with(".CRT$XC") || section_name.starts_with(".CRT$XT")) {
- for (auto* Block : Section.blocks()) {
- bool has_live_sym = false;
- for (auto* Sym : G.defined_symbols()) {
- if (&Sym->getBlock() == Block) {
- Sym->setLive(true);
- has_live_sym = true;
- }
- }
- // MSVC may emit .CRT$XC* blocks with data but no symbol table
- // entries (static variables in __declspec(allocate) sections).
- // Add an anonymous symbol so the block survives dead-stripping.
- if (!has_live_sym) {
- G.addAnonymousSymbol(*Block, 0, Block->getSize(), false, true).setLive(true);
- }
- for (auto& Edge : Block->edges()) {
- Edge.getTarget().setLive(true);
- }
- }
- }
- }
- return llvm::Error::success();
- });
-#ifdef _WIN32
- // Without COFFPlatform, __ImageBase (used by IMAGE_REL_AMD64_ADDR32NB /
- // Pointer32NB relocations) defaults to 0. This causes all Pointer32 fixups
- // to overflow since JIT addresses don't fit in 32 bits.
- //
- // Fix: set __ImageBase to the lowest block address in the graph after
- // allocation. This makes all intra-graph Pointer32NB offsets small.
- //
- // Also strip .pdata/.xdata edges: SEH unwind data references external
- // handlers (e.g., __CxxFrameHandler3) in DLLs that may be >4GB from
- // __ImageBase. Since we don't call RtlAddFunctionTable, SEH data is
- // unused anyway.
- Config.PostAllocationPasses.emplace_back([](llvm::jitlink::LinkGraph& G) {
- // Set __ImageBase to the lowest allocated block address.
- auto ImageBaseName = G.intern("__ImageBase");
- llvm::jitlink::Symbol* ImageBase = nullptr;
- for (auto* Sym : G.external_symbols()) {
- if (Sym->getName() == ImageBaseName) {
- ImageBase = Sym;
- break;
- }
- }
- if (ImageBase) {
- llvm::orc::ExecutorAddr BaseAddr;
- for (auto* B : G.blocks()) {
- if (!BaseAddr || B->getAddress() < BaseAddr) {
- BaseAddr = B->getAddress();
- }
- }
- ImageBase->getAddressable().setAddress(BaseAddr);
- }
- // Strip .pdata/.xdata edges: external handlers may be >4GB from
- // __ImageBase, and we don't register SEH data anyway.
- for (auto& Sec : G.sections()) {
- if (Sec.getName() == ".pdata" || Sec.getName().starts_with(".xdata")) {
- for (auto* B : Sec.blocks()) {
- while (!B->edges_empty()) {
- B->removeEdge(B->edges().begin());
- }
- }
- }
- }
- return llvm::Error::success();
- });
-#endif
- // After fixups, read resolved function pointers from all init/fini data sections.
- // Handles ELF (.init_array, .ctors, .fini_array, .dtors),
- // Mach-O (__DATA,__mod_init_func, __DATA,__mod_term_func),
- // and COFF (.CRT$XC*, .CRT$XT*) section conventions.
- Config.PostFixupPasses.emplace_back([this, &jit_dylib](llvm::jitlink::LinkGraph& G) {
- using Entry = ORCJITExecutionSessionObj::InitFiniEntry;
- for (auto& Sec : G.sections()) {
- auto section_name = Sec.getName();
-
- // --- ELF sections ---
- bool is_init_array = section_name.starts_with(".init_array");
- bool is_ctors = section_name.starts_with(".ctors");
- bool is_fini_array = section_name.starts_with(".fini_array");
- bool is_dtors = section_name.starts_with(".dtors");
- // --- Mach-O sections ---
- bool is_mod_init = (section_name == "__DATA,__mod_init_func");
- bool is_mod_term = (section_name == "__DATA,__mod_term_func");
- // --- COFF sections ---
- bool is_crt_xc = section_name.starts_with(".CRT$XC");
- bool is_crt_xt = section_name.starts_with(".CRT$XT");
-
- if (!is_init_array && !is_ctors && !is_fini_array && !is_dtors && !is_mod_init &&
- !is_mod_term && !is_crt_xc && !is_crt_xt)
- continue;
-
- int priority = 0;
- Entry::Section sec;
- bool is_init;
- // ELF default priority for sections without a numeric suffix is 65535.
- // Lower priority numbers run first for .init_array; .fini_array and .ctors
- // negate so that higher-numbered entries run first (reverse order).
- if (is_init_array) {
- if (section_name.consume_front(".init_array.")) {
- section_name.getAsInteger(10, priority);
- } else {
- priority = 65535;
- }
- sec = Entry::Section::kInitArray;
- is_init = true;
- } else if (is_ctors) {
- if (section_name.consume_front(".ctors.") && !section_name.getAsInteger(10, priority)) {
- priority = -priority;
- }
- sec = Entry::Section::kCtors;
- is_init = true;
- } else if (is_fini_array) {
- if (section_name.consume_front(".fini_array.") &&
- !section_name.getAsInteger(10, priority)) {
- priority = -priority;
- } else {
- priority = -65535;
- }
- sec = Entry::Section::kFiniArray;
- is_init = false;
- } else if (is_dtors) {
- if (section_name.consume_front(".dtors.")) {
- section_name.getAsInteger(10, priority);
- }
- sec = Entry::Section::kDtors;
- is_init = false;
- } else if (is_mod_init) {
- // Mach-O __mod_init_func: no priority system, treated as init_array
- sec = Entry::Section::kInitArray;
- is_init = true;
- } else if (is_mod_term) {
- // Mach-O __mod_term_func: no priority system, treated as fini_array
- sec = Entry::Section::kFiniArray;
- is_init = false;
- } else if (is_crt_xc) {
- // COFF .CRT$XC[suffix]: C++ constructors, sorted alphabetically by suffix.
- // Convert suffix to integer priority that preserves alphabetical ordering.
- // E.g., .CRT$XCA → 'A'*100000=6500000, .CRT$XCU → 'U'*100000=8500000,
- // .CRT$XCT00200 → 'T'*100000+200=8400200
- sec = Entry::Section::kInitArray;
- is_init = true;
- auto suffix = section_name.substr(7); // after ".CRT$XC"
- if (!suffix.empty()) {
- priority = static_cast<int>(suffix[0]) * 100000;
- if (suffix.size() > 1) {
- int num = 0;
- suffix.substr(1).getAsInteger(10, num);
- priority += num;
- }
- }
- } else {
- // COFF .CRT$XT[suffix]: C++ destructors, same suffix-to-priority scheme.
- sec = Entry::Section::kFiniArray;
- is_init = false;
- auto suffix = section_name.substr(7); // after ".CRT$XT"
- if (!suffix.empty()) {
- priority = static_cast<int>(suffix[0]) * 100000;
- if (suffix.size() > 1) {
- int num = 0;
- suffix.substr(1).getAsInteger(10, num);
- priority += num;
- }
- }
- }
-
- for (auto* Block : Sec.blocks()) {
- auto Content = Block->getContent();
- size_t PtrSize = G.getPointerSize();
- for (size_t Offset = 0; Offset + PtrSize <= Content.size(); Offset += PtrSize) {
- uint64_t FnAddr = 0;
- memcpy(&FnAddr, Content.data() + Offset, PtrSize);
- if (FnAddr != 0) {
- Entry entry{llvm::orc::ExecutorAddr(FnAddr), sec, priority};
- if (is_init) {
- session_->AddPendingInitializer(&jit_dylib, entry);
- } else {
- session_->AddPendingDeinitializer(&jit_dylib, entry);
- }
- }
- }
- }
- }
- return llvm::Error::success();
- });
- }
-
- llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override {
- return llvm::Error::success();
- }
-
- llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override {
- return llvm::Error::success();
- }
-
- void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey,
- llvm::orc::ResourceKey SrcKey) override {}
-};
-
-#ifdef _WIN32
-/*!
- * \brief Custom definition generator for Windows DLL import symbols.
- *
- * On Windows with the MSVC ABI, COFF objects reference DLL-imported symbols
- * via __imp_XXX pointer stubs and direct calls. Without COFFPlatform (which
- * we skip due to MSVC CRT dependency issues), JITLink cannot resolve these.
- *
- * For each resolved symbol, this generator creates a JIT-allocated LinkGraph
- * containing:
- * - __imp_XXX: a GOT-like pointer entry holding the real address
- * - XXX: a PLT-like jump stub (`jmp [__imp_XXX]`) for direct calls
- *
- * By allocating stubs in JIT memory (rather than using absoluteSymbols at
- * host-process addresses), all PCRel32 fixups from JIT code reach safely.
- *
- * Symbol search order:
- * 1. Specific MSVC runtime DLLs (vcruntime140, ucrtbase, msvcp140)
- * 2. All loaded process modules (EnumProcessModules)
- * 3. LLVM's SearchForAddressOfSymbol
- */
-class DLLImportDefinitionGenerator : public llvm::orc::DefinitionGenerator {
- llvm::orc::ExecutionSession& ES_;
- llvm::orc::ObjectLinkingLayer& L_;
-
- static void* FindInProcessModules(const std::string& Name) {
- // Try specific runtime DLLs first, then tvm_ffi.dll (loaded by Python),
- // then all process modules, then LLVM's search.
- static const char* kRuntimeDLLs[] = {
- "vcruntime140.dll",
- "vcruntime140_1.dll",
- "ucrtbase.dll",
- "msvcp140.dll",
- };
- // NOTE: We intentionally do not call FreeLibrary() here. These runtime DLLs
- // (vcruntime140, ucrtbase, etc.) are already loaded by the process and will
- // remain loaded for its lifetime. LoadLibraryA merely increments the refcount;
- // the extra refcount is harmless and avoids the overhead of balancing
- // Get/FreeLibrary for every symbol lookup.
- for (const char* dll : kRuntimeDLLs) {
- if (HMODULE hMod = LoadLibraryA(dll)) {
- if (auto addr = GetProcAddress(hMod, Name.c_str())) {
- return reinterpret_cast<void*>(addr);
- }
- }
- }
- // Also check tvm_ffi.dll (host process symbol provider)
- if (HMODULE hTvmFfi = GetModuleHandleA("tvm_ffi.dll")) {
- if (auto addr = GetProcAddress(hTvmFfi, Name.c_str())) {
- return reinterpret_cast<void*>(addr);
- }
- }
- HMODULE hMods[1024];
- DWORD cbNeeded;
- if (EnumProcessModules(GetCurrentProcess(), hMods, sizeof(hMods), &cbNeeded)) {
- DWORD count = cbNeeded / sizeof(HMODULE);
- if (count > 1024) count = 1024;
- for (DWORD i = 0; i < count; ++i) {
- if (auto addr = GetProcAddress(hMods[i], Name.c_str())) {
- return reinterpret_cast<void*>(addr);
- }
- }
- }
- if (void* addr = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol(Name)) {
- return addr;
- }
- return nullptr;
- }
-
- public:
- DLLImportDefinitionGenerator(llvm::orc::ExecutionSession& ES, llvm::orc::ObjectLinkingLayer& L)
- : ES_(ES), L_(L) {}
-
- llvm::Error tryToGenerate(llvm::orc::LookupState& LS, llvm::orc::LookupKind K,
- llvm::orc::JITDylib& JD, llvm::orc::JITDylibLookupFlags JDLookupFlags,
- const llvm::orc::SymbolLookupSet& LookupSet) override {
- // Step 1: Collect unique base names (strip __imp_ prefix) and resolve addresses.
- llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorAddr> Resolved;
- for (auto& [Name, Flags] : LookupSet) {
- llvm::StringRef NameStr = *Name;
- std::string BaseName =
- NameStr.starts_with("__imp_") ? NameStr.drop_front(6).str() : NameStr.str();
- if (BaseName == "__ImageBase") continue;
- auto InternedBase = ES_.intern(BaseName);
- if (Resolved.count(InternedBase)) continue;
- void* Addr = FindInProcessModules(BaseName);
- if (Addr) {
- Resolved[InternedBase] = llvm::orc::ExecutorAddr::fromPtr(Addr);
- }
- }
- if (Resolved.empty()) return llvm::Error::success();
-
- // Step 2: Build a LinkGraph with __imp_ pointers and PLT jump stubs.
- auto G = std::make_unique<llvm::jitlink::LinkGraph>(
- "<DLL_IMPORT_STUBS>", ES_.getSymbolStringPool(), ES_.getTargetTriple(),
- llvm::SubtargetFeatures(), llvm::jitlink::getGenericEdgeKindName);
- auto Prot = static_cast<llvm::orc::MemProt>(static_cast<unsigned>(llvm::orc::MemProt::Read) |
- static_cast<unsigned>(llvm::orc::MemProt::Exec));
- auto& Sec = G->createSection("__dll_stubs", Prot);
-
- for (auto& [InternedName, Addr] : Resolved) {
- // Absolute symbol at the real address (local to this graph)
- auto& Target = G->addAbsoluteSymbol(G->intern(("__real_" + *InternedName).str()), Addr,
- G->getPointerSize(), llvm::jitlink::Linkage::Strong,
- llvm::jitlink::Scope::Local, false);
- // __imp_XXX pointer (GOT-like entry)
- auto& Ptr = llvm::jitlink::x86_64::createAnonymousPointer(*G, Sec, &Target);
- Ptr.setName(G->intern(("__imp_" + *InternedName).str()));
- Ptr.setLinkage(llvm::jitlink::Linkage::Strong);
- Ptr.setScope(llvm::jitlink::Scope::Default);
- // XXX jump stub (PLT-like entry) for direct calls
- auto& StubBlock = llvm::jitlink::x86_64::createPointerJumpStubBlock(*G, Sec, Ptr);
- G->addDefinedSymbol(StubBlock, 0, *InternedName, StubBlock.getSize(),
- llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Default, true,
- false);
- }
- return L_.add(JD, std::move(G));
- }
-};
-#endif // _WIN32
-
-#if defined(__linux__) && (defined(__x86_64__) || defined(_M_X64))
-/*! \brief Fix LLVM JITLink GOTPCRELX relaxation bug (x86_64).
- *
- * optimizeGOTAndStubAccesses() relaxes `call *foo@GOTPCREL(%rip)` (ff 15)
- * to `addr32 call foo` (67 e8) and sets the edge to Pointer32 (absolute
- * 32-bit). But `call rel32` is always PC-relative — the CPU computes
- * target = RIP + imm32, not target = imm32. The Pointer32 fixup writes
- * the absolute address, producing a wrong displacement.
- *
- * This only manifests when external symbols resolve to low addresses
- * (< 4 GB, e.g. PLT entries in a non-PIE executable) while JIT code is
- * at high addresses (the arena at 0x7f...). The optimization fires
- * because isUInt<32>(target) is true, but the resulting fixup is wrong.
- *
- * The PreFixupPass reverts broken relaxations back to indirect calls
- * through the GOT. See orcjit_memory_manager.h for full context.
- */
-class GOTPCRELXFixPlugin : public llvm::orc::ObjectLinkingLayer::Plugin {
- public:
- void modifyPassConfig(llvm::orc::MaterializationResponsibility& MR, llvm::jitlink::LinkGraph& G,
- llvm::jitlink::PassConfiguration& Config) override {
- Config.PreFixupPasses.emplace_back(fixBrokenGOTPCRELXRelaxation);
- }
- llvm::Error notifyFailed(llvm::orc::MaterializationResponsibility& MR) override {
- return llvm::Error::success();
- }
- llvm::Error notifyRemovingResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey K) override {
- return llvm::Error::success();
- }
- void notifyTransferringResources(llvm::orc::JITDylib& JD, llvm::orc::ResourceKey DstKey,
- llvm::orc::ResourceKey SrcKey) override {}
-
- private:
- /*! \brief Correct broken GOTPCRELX relaxations produced by
- * optimizeGOTAndStubAccesses().
- *
- * Strategy:
- * 1. Build target-symbol → GOT-entry-symbol map (O(B+S) up front).
- * 2. For every Pointer32 edge whose preceding bytes are 67 e8
- * (relaxed call) or e9 (relaxed jmp):
- * - If the target is reachable via a signed 32-bit PC-relative
- * displacement, change the edge to BranchPCRel32.
- * - Otherwise revert the relaxation: restore the original
- * indirect-call/jmp opcode bytes (ff 15 / ff 25), retarget
- * the edge to the GOT entry, and use PCRel32 with addend 0
- * (JITLink normalises GOTPCRELX addends to 0).
- */
- static llvm::Error fixBrokenGOTPCRELXRelaxation(llvm::jitlink::LinkGraph& G) {
- using namespace llvm::jitlink;
- // Build block → first symbol at offset 0 (for GOT entry symbol lookup).
- llvm::DenseMap<Block*, Symbol*> BlockToSym;
- for (auto* Sym : G.defined_symbols()) {
- if (Sym->getOffset() == 0 && !BlockToSym.count(&Sym->getBlock())) {
- BlockToSym[&Sym->getBlock()] = Sym;
- }
- }
-
- // Build target symbol → GOT entry symbol map.
- // GOT entries are pointer-sized blocks with exactly one Pointer64 edge.
- llvm::DenseMap<Symbol*, Symbol*> SymToGOTSym;
- for (auto* B : G.blocks()) {
- if (B->getSize() != G.getPointerSize()) continue;
- if (B->edges_size() != 1) continue;
- auto& E = *B->edges().begin();
- if (E.getKind() == x86_64::Pointer64) {
- auto It = BlockToSym.find(B);
- if (It != BlockToSym.end()) {
- SymToGOTSym[&E.getTarget()] = It->second;
- }
- }
- }
-
- for (auto* B : G.blocks()) {
- for (auto& E : B->edges()) {
- if (E.getKind() != x86_64::Pointer32) continue;
- if (E.getOffset() < 2) continue;
-
- auto MutableContent = B->getMutableContent(G);
- auto* FixupData = reinterpret_cast<uint8_t*>(MutableContent.data()) + E.getOffset();
- uint8_t Prev2 = FixupData[-2];
- uint8_t Prev1 = FixupData[-1];
-
- bool isRelaxedCall = (Prev2 == 0x67 && Prev1 == 0xe8);
- bool isRelaxedJmp = (Prev1 == 0xe9);
- if (!isRelaxedCall && !isRelaxedJmp) continue;
-
- // Check if PC-relative displacement would fit.
- auto TargetAddr = E.getTarget().getAddress();
- auto FixupAddr = B->getFixupAddress(E);
- int64_t Displacement = TargetAddr.getValue() - (FixupAddr.getValue() + 4) + E.getAddend();
- if (llvm::isInt<32>(Displacement)) {
- E.setKind(x86_64::BranchPCRel32);
- continue;
- }
-
- // Distance doesn't fit — revert to indirect call/jmp through GOT.
- auto It = SymToGOTSym.find(&E.getTarget());
- if (It == SymToGOTSym.end()) {
- return llvm::make_error<llvm::StringError>(
- "Cannot revert GOTPCRELX relaxation: no GOT entry for " +
- (E.getTarget().hasName() ? std::string(*E.getTarget().getName())
- : std::string("<anon>")),
- llvm::inconvertibleErrorCode());
- }
-
- Symbol* GOTSym = It->second;
- if (isRelaxedCall) {
- // Restore: 67 e8 → ff 15 (call *[rip+disp32])
- FixupData[-2] = 0xff;
- FixupData[-1] = 0x15;
- } else {
- // Restore: e9 XX XX XX XX 90 → ff 25 XX XX XX XX
- FixupData[-1] = 0xff;
- FixupData[0] = 0x25;
- // For jmp, the optimization shifted offset by -1; shift back.
- E.setOffset(E.getOffset() + 1);
- }
- E.setKind(x86_64::PCRel32);
- E.setTarget(*GOTSym);
- E.setAddend(0);
- }
- }
- return llvm::Error::success();
- }
-};
-#endif // __linux__ && __x86_64__
-
ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_path,
- int64_t arena_size_bytes)
+ int64_t slab_size_bytes)
: jit_(nullptr) {
- // Create arena memory manager — pre-reserves contiguous VA region so all
- // JIT allocations stay within PC-relative relocation range (±2GB x86_64,
- // ±4GB AArch64). Eliminates scattered-mmap relocation overflow (LLVM #173269).
+ // Create slab-backed memory manager — pre-reserves a contiguous VA region
+ // so all JIT allocations stay within PC-relative relocation range (±2 GB
+ // x86_64, ±4 GB AArch64). Eliminates scattered-mmap relocation overflow
+ // (LLVM #173269).
//
- // arena_size_bytes: 0 = arch default (4GB x86_64, 8GB AArch64, with fallback),
- // >0 = custom size, <0 = disable arena.
+ // slab_size_bytes: 0 = arch default (1 GB x86_64 / AArch64, with fallback),
+ // >0 = custom size, <0 = disable arena (LLJIT uses its
+ // default allocator — scattered mmap, no PC-rel guarantee).
// The parameter is Linux-only; on macOS/Windows the arena is compiled out
// entirely (see #ifdef below) and the value is ignored.
//
- // The default is strictly larger than the relocation limit so the arena is
- // never the bottleneck — JITLink's own overflow check fires first, matching
- // dlopen/ld.so semantics. The constructor halves capacity on mmap failure
- // (RLIMIT_AS, containers) down to 256 MB.
+ // `slab_size_bytes` is the per-slab capacity for the growable pool.
+ // Session memory grows in slab-sized increments; graphs that don't
+ // fit a normal slab trigger a power-of-2 larger slab sized to fit
+ // (see `Slab::capacityForFootprint`).
+ //
+ // The default (64 MB) is above typical ML JIT graph sizes while well
+ // under the PC-relative relocation limit. The initial-slab constructor
+ // halves its capacity on mmap failure (RLIMIT_AS, containers) down to
+ // 8 MB; subsequent slabs are reserved at the size returned by
+ // `capacityForFootprint` (>= slab_size) and mmap errors propagate.
//
// LLJIT auto-configures ObjectLinkingLayer (JITLink) on x86_64 and aarch64
// Linux (see LLJITBuilderState::prepareForConstruction). We override
- // the layer creator to pass our arena MM. macOS/Windows are excluded:
- // macOS MachOPlatform teardown crashes with the arena; Windows needs
- // further testing.
+ // the layer creator to pass our memory manager. macOS/Windows are gated
+ // off pending testing. (The historical "MachOPlatform teardown crashes
+ // with the arena" concern is moot now that we skip MachOPlatform below,
+ // but enabling the slab on macOS still needs a validation pass.)
#ifdef __linux__
- if (arena_size_bytes >= 0) {
+ if (slab_size_bytes >= 0) {
auto page_size = llvm::sys::Process::getPageSizeEstimate();
- size_t capacity;
- if (arena_size_bytes > 0) {
- capacity = static_cast<size_t>(arena_size_bytes);
+ size_t slab_size;
+ if (slab_size_bytes > 0) {
+ slab_size = static_cast<size_t>(slab_size_bytes);
} else {
-#if defined(__aarch64__)
- capacity = ArenaJITLinkMemoryManager::kDefaultArenaCapacity_AArch64;
-#else
- capacity = ArenaJITLinkMemoryManager::kDefaultArenaCapacity_x86_64;
-#endif
+ slab_size = SlabPoolMemoryManager::kDefaultSlabSize;
}
- memory_manager_ = std::make_unique<ArenaJITLinkMemoryManager>(page_size, capacity);
+ memory_manager_ = std::make_unique<SlabPoolMemoryManager>(page_size, slab_size);
}
#endif
@@ -625,8 +127,11 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_p
});
} // if (memory_manager_)
#elif defined(__APPLE__) || defined(_WIN32)
- // macOS: MachOPlatform (via ExecutorNativePlatform) requires ObjectLinkingLayer.
- // Windows: need ObjectLinkingLayer for InitFiniPlugin and DLLImportDefinitionGenerator.
+ // Force ObjectLinkingLayer (JITLink) so we can attach InitFiniPlugin.
+ // macOS: LLJIT already defaults to JITLink for Darwin, but the explicit
+ // creator keeps the static_cast in the addPlugin site below type-safe.
+ // Windows: LLJIT defaults to RTDyld; we need JITLink for InitFiniPlugin
+ // and DLLImportDefinitionGenerator.
builder.setObjectLinkingLayerCreator(
[](llvm::orc::ExecutionSession& ES)
-> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
@@ -648,119 +153,36 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_p
(void)builder;
};
+ auto builder = llvm::orc::LLJITBuilder();
+#ifndef __APPLE__
+ // macOS: always skip ExecutorNativePlatform / MachOPlatform to sidestep
+ // the compact-unwind 32-bit-delta bug in JITLink's CompactUnwindSupport
+ // (personality delta against a per-JITDylib header base wraps `uint64_t`
+ // and fails `isUInt<32>` when a later user graph mmaps below the header;
+ // see the repo-root fix-machoplatform-libunwind-dso-base.patch for the
+ // full analysis). InitFiniPlugin below handles __mod_init_func /
+ // __mod_term_func instead. Tradeoff: no C++ exception unwinding across
+ // JIT frames on macOS.
if (!orc_rt_path.empty()) {
- auto builder = llvm::orc::LLJITBuilder();
builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(orc_rt_path));
- setup_builder(builder);
- jit_ = TVM_FFI_ORCJIT_LLVM_CALL(builder.create());
- } else {
- auto builder = llvm::orc::LLJITBuilder();
- setup_builder(builder);
- jit_ = TVM_FFI_ORCJIT_LLVM_CALL(builder.create());
}
+#else
+ (void)orc_rt_path;
+#endif
+ setup_builder(builder);
+ jit_ = TVM_FFI_ORCJIT_LLVM_CALL(builder.create());
#ifdef _WIN32
- // Strip .pdata/.xdata relocations from COFF objects before JITLink graph building.
- // clang-cl puts static functions in COMDAT sections, and .pdata SEH unwind data
- // has relocations targeting COMDAT leader symbols. JITLink's COFFLinkGraphBuilder
- // doesn't register COMDAT leaders in its symbol table when the second COMDAT symbol
- // is CLASS_STATIC (not CLASS_EXTERNAL), causing "Could not find symbol" errors.
- // We already strip .pdata/.xdata edges in a PostAllocationPass; this moves the
- // stripping earlier to prevent the graph builder error.
- // Strip .pdata/.xdata relocations using raw COFF binary manipulation.
- // We avoid llvm/Object/COFF.h because windows.h (included transitively by
- // LLJIT.h) defines IMAGE_* macros that conflict with LLVM's COFF enums.
- jit_->getObjTransformLayer().setTransform(
- [](std::unique_ptr<llvm::MemoryBuffer> Buf)
- -> llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>> {
- const char* Data = Buf->getBufferStart();
- size_t Size = Buf->getBufferSize();
- if (Size < 20) return std::move(Buf);
-
- // Parse COFF header (regular or bigobj format)
- uint16_t w0, w1;
- std::memcpy(&w0, Data, 2);
- std::memcpy(&w1, Data + 2, 2);
- bool bigobj = (w0 == 0 && w1 == 0xFFFF);
-
- uint16_t machine;
- uint32_t num_sections, ptr_to_symtab, num_symbols;
- size_t sec_hdr_start, sym_entry_size;
- if (bigobj) {
- if (Size < 56) return std::move(Buf);
- std::memcpy(&machine, Data + 6, 2);
- std::memcpy(&num_sections, Data + 44, 4);
- std::memcpy(&ptr_to_symtab, Data + 48, 4);
- std::memcpy(&num_symbols, Data + 52, 4);
- sec_hdr_start = 56;
- sym_entry_size = 20;
- } else {
- machine = w0;
- uint16_t ns, opt_hdr_size;
- std::memcpy(&ns, Data + 2, 2);
- std::memcpy(&opt_hdr_size, Data + 16, 2);
- std::memcpy(&ptr_to_symtab, Data + 8, 4);
- std::memcpy(&num_symbols, Data + 12, 4);
- num_sections = ns;
- sec_hdr_start = 20 + opt_hdr_size;
- sym_entry_size = 18;
- }
- if (machine != 0x8664) return std::move(Buf);
-
- // String table follows the symbol table
- size_t strtab_start = ptr_to_symtab + static_cast<size_t>(num_symbols) * sym_entry_size;
-
- // Resolve a section name (inline 8-byte or "/offset" string table ref)
- constexpr size_t kSecHdrSize = 40;
- auto resolve_name = [&](size_t hdr_off) -> llvm::StringRef {
- const char* raw = Data + hdr_off;
- if (raw[0] == '/' && raw[1] >= '0' && raw[1] <= '9') {
- uint32_t offset = 0;
- for (int j = 1; j < 8 && raw[j] >= '0' && raw[j] <= '9'; ++j)
- offset = offset * 10 + (raw[j] - '0');
- size_t pos = strtab_start + offset;
- if (pos < Size) {
- size_t len = 0;
- while (pos + len < Size && Data[pos + len]) ++len;
- return {Data + pos, len};
- }
- }
- size_t len = 0;
- while (len < 8 && raw[len]) ++len;
- return {raw, len};
- };
-
- // Collect section header offsets needing relocation stripping
- llvm::SmallVector<size_t, 8> strip_offsets;
- for (uint32_t i = 0; i < num_sections; ++i) {
- size_t off = sec_hdr_start + i * kSecHdrSize;
- if (off + kSecHdrSize > Size) break;
- auto name = resolve_name(off);
- if (name.starts_with(".pdata") || name.starts_with(".xdata")) {
- uint16_t num_relocs;
- std::memcpy(&num_relocs, Data + off + 32, 2);
- if (num_relocs > 0) strip_offsets.push_back(off);
- }
- }
- if (strip_offsets.empty()) return std::move(Buf);
-
- // Create mutable copy, zero out PointerToRelocations and NumberOfRelocations
- llvm::SmallVector<char> MutableBuf(Data, Data + Size);
- for (auto off : strip_offsets) {
- std::memset(&MutableBuf[off + 24], 0, 4); // PointerToRelocations
- std::memset(&MutableBuf[off + 32], 0, 2); // NumberOfRelocations
- }
- return llvm::MemoryBuffer::getMemBufferCopy(
- llvm::StringRef(MutableBuf.data(), MutableBuf.size()), Buf->getBufferIdentifier());
- });
+ // Strip .pdata/.xdata relocations from COFF objects before JITLink graph
+ // building. See llvm_patches/win_coff_pdata_strip.h for the rationale.
+ jit_->getObjTransformLayer().setTransform(&StripCoffPdataXdata);
#endif
-#if defined(__linux__) || defined(_WIN32)
- // Linux/Windows: use our custom InitFiniPlugin for init/fini section
- // collection and priority-ordered execution. See InitFiniPlugin class
- // documentation for the three-platform init/fini strategy.
+ // Use our custom InitFiniPlugin on every platform for init/fini section
+ // collection and priority-ordered execution (ELF .init_array/.fini_array,
+ // MachO __mod_init_func/__mod_term_func, COFF .CRT$XC*/.CRT$XT*). See
+ // llvm_patches/init_fini_plugin.h for per-platform removal criteria.
auto& objlayer = jit_->getObjLinkingLayer();
static_cast<llvm::orc::ObjectLinkingLayer&>(objlayer).addPlugin(
std::make_unique<InitFiniPlugin>(this));
-#endif
#ifdef _WIN32
// On Windows, the default process-symbol generator only searches the main
// exe module via GetProcAddress(GetModuleHandle(NULL), ...). Add a
@@ -775,9 +197,9 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_p
}
ORCJITExecutionSession::ORCJITExecutionSession(const std::string& orc_rt_path,
- int64_t arena_size_bytes) {
+ int64_t slab_size_bytes) {
ObjectPtr<ORCJITExecutionSessionObj> obj =
- make_object<ORCJITExecutionSessionObj>(orc_rt_path, arena_size_bytes);
+ make_object<ORCJITExecutionSessionObj>(orc_rt_path, slab_size_bytes);
data_ = std::move(obj);
}
@@ -803,10 +225,17 @@ ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary(const Strin
jit_dylib.addToLinkOrder(*kv.first, kv.second);
}
- auto dylib = ORCJITDynamicLibrary(make_object<ORCJITDynamicLibraryObj>(
- GetRef<ORCJITExecutionSession>(this), &jit_dylib, jit_.get(), lib_name));
+ auto dylib_obj = make_object<ORCJITDynamicLibraryObj>(GetRef<ORCJITExecutionSession>(this),
+ &jit_dylib, jit_.get(), lib_name);
+
+#ifdef __APPLE__
+ // Inject ___cxa_atexit on the user JITDylib so it wins over <Platform>'s
+ // fallback (which resolves to libSystem's and would orphan dtors from
+ // our drop-time drain). See llvm_patches/macho_cxa_atexit_shim.h.
+ InstallCxaAtexitShim(jit_->getExecutionSession(), jit_dylib);
+#endif
- return dylib;
+ return ORCJITDynamicLibrary(std::move(dylib_obj));
}
llvm::orc::ExecutionSession& ORCJITExecutionSessionObj::GetLLVMExecutionSession() {
@@ -859,6 +288,33 @@ void ORCJITExecutionSessionObj::AddPendingDeinitializer(llvm::orc::JITDylib* jit
pending_deinitializers_[jit_dylib].push_back(entry);
}
+int64_t ORCJITExecutionSessionObj::ClearFreeSlabs() {
+#ifdef __linux__
+ if (memory_manager_) {
+ return static_cast<int64_t>(memory_manager_->clearFreeSlabs());
+ }
+#endif
+ return 0;
+}
+
+void ORCJITExecutionSessionObj::RemoveDylib(llvm::orc::JITDylib* jit_dylib) {
+ if (jit_dylib == nullptr) return;
+ // Drop any pending init/fini records keyed by this JITDylib*. After removal
+ // the address may be recycled for a freshly-created JITDylib; leftover
+ // entries would then be attributed to the wrong dylib.
+ pending_initializers_.erase(jit_dylib);
+ pending_deinitializers_.erase(jit_dylib);
+
+ if (jit_ == nullptr) return;
+ // removeJITDylib is best-effort at destruction time: the session may already
+ // be tearing down, the platform may report an error during clear(), etc.
+ // Swallow errors rather than throwing from a destructor; the session
+ // destructor will munmap everything when it runs.
+ if (auto err = jit_->getExecutionSession().removeJITDylib(*jit_dylib)) {
+ llvm::consumeError(std::move(err));
+ }
+}
+
} // namespace orcjit
} // namespace ffi
} // namespace tvm
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
index e64c17b..a35983a 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
@@ -56,7 +56,7 @@ class ORCJITExecutionSessionObj : public Object {
* \brief Default constructor (for make_object)
*/
explicit ORCJITExecutionSessionObj(const std::string& orc_rt_path = "",
- int64_t arena_size_bytes = 0);
+ int64_t slab_size_bytes = 0);
/*!
* \brief Create a new DynamicLibrary (JITDylib) in this session
@@ -98,9 +98,34 @@ class ORCJITExecutionSessionObj : public Object {
void AddPendingInitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry);
void AddPendingDeinitializer(llvm::orc::JITDylib* jd, const InitFiniEntry& entry);
+ /*!
+ * \brief Remove a JITDylib from the ExecutionSession, releasing its JIT
+ * 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
+ * further use of the \c JITDylib* after this call — it becomes "Closed" and
+ * its address may be reused by a subsequent \c createJITDylib.
+ *
+ * Also erases any pending init/fini map entries keyed by \p jd so that a
+ * subsequent \c JITDylib allocated at the same address starts with a clean
+ * slate.
+ */
+ void RemoveDylib(llvm::orc::JITDylib* jd);
+
+ /*!
+ * \brief Release drained slabs (no live JIT allocations) back to the OS.
+ *
+ * Returns the number of slabs reclaimed. No-op on macOS/Windows
+ * where the slab pool is compiled out, or when the pool has been
+ * disabled via `slab_size < 0`.
+ */
+ int64_t ClearFreeSlabs();
+
private:
- /*! \brief Arena memory manager — must be declared before jit_ for destruction order */
- std::unique_ptr<ArenaJITLinkMemoryManager> memory_manager_;
+ /*! \brief Slab-pool memory manager — must be declared before jit_ for destruction order */
+ std::unique_ptr<SlabPoolMemoryManager> memory_manager_;
/*! \brief The LLVM ORC JIT instance */
std::unique_ptr<llvm::orc::LLJIT> jit_;
@@ -122,8 +147,7 @@ class ORCJITExecutionSession : public ObjectRef {
* \brief Create a new ExecutionSession
* \return The created execution session instance
*/
- explicit ORCJITExecutionSession(const std::string& orc_rt_path = "",
- int64_t arena_size_bytes = 0);
+ explicit ORCJITExecutionSession(const std::string& orc_rt_path = "", int64_t slab_size_bytes = 0);
// Required: define object reference methods
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ORCJITExecutionSession, ObjectRef,
ORCJITExecutionSessionObj);
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc
similarity index 51%
copy from addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc
copy to addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc
index 26be4b8..5c2d629 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.cc
@@ -18,15 +18,11 @@
*/
/*!
- * \file orcjit_memory_manager.cc
- * \brief Arena-based JITLinkMemoryManager implementation.
- *
- * Follows the InProcessMemoryManager pattern from LLVM but replaces
- * per-object mmap with bump allocation from a pre-reserved arena.
- * Pages are committed in 2 MB slabs to enable Transparent Huge Page
- * (THP) promotion — see the class docstring in orcjit_memory_manager.h.
+ * \file orcjit_slab.cc
+ * \brief Slab implementation — bump allocator + commit bitmap + free list.
*/
-#include "orcjit_memory_manager.h"
+
+#include "orcjit_slab.h"
#ifdef __linux__
@@ -35,7 +31,6 @@
#include <llvm/ExecutionEngine/JITLink/JITLink.h>
#include <llvm/ExecutionEngine/JITLink/aarch64.h>
#include <llvm/ExecutionEngine/JITLink/x86_64.h>
-#include <llvm/ExecutionEngine/Orc/Shared/AllocationActions.h>
#include <llvm/Support/Alignment.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/Support/Memory.h>
@@ -45,6 +40,7 @@
#include <cerrno>
#include <cstdio>
#include <cstring>
+#include <limits>
namespace tvm {
namespace ffi {
@@ -54,11 +50,28 @@ using namespace llvm;
using namespace llvm::jitlink;
using namespace llvm::orc;
+// ── SlabPoolExhaustedError ──────────────────────────────────────────
+
+char SlabPoolExhaustedError::ID = 0;
+
+SlabPoolExhaustedError::SlabPoolExhaustedError(const char* pool, std::size_t used,
+ std::size_t requested, std::size_t limit)
+ : pool_(pool), used_(used), requested_(requested), limit_(limit) {}
+
+void SlabPoolExhaustedError::log(raw_ostream& os) const {
+ os << "Slab: " << pool_ << " pool exhausted (used " << formatv("{0:x}", used_) << " + requested "
+ << formatv("{0:x}", requested_) << " > limit " << formatv("{0:x}", limit_) << ")";
+}
+
+std::error_code SlabPoolExhaustedError::convertToErrorCode() const {
+ return inconvertibleErrorCode();
+}
+
// ── Overflow section edge classification ───────────────────────────
//
// Conservative whitelist: only known absolute relocation kinds return true.
// Unknown or future edge kinds default to PC-relative → sections stay in
-// the arena (safe: never breaks relocations, just forgoes the overflow
+// the slab (safe: never breaks relocations, just forgoes the overflow
// optimization for unknown kinds).
namespace {
@@ -95,70 +108,104 @@ bool isAbsoluteEdge(const Triple& TT, Edge::Kind K) {
return false; // Unknown arch — treat as PC-relative (safe)
}
+/*! \brief Identify sections eligible for the overflow (separate-mmap)
+ * path.
+ *
+ * Name-based candidate selection followed by edge validation: any
+ * PC-relative cross-section edge targeting a candidate disqualifies
+ * it (the section must live inside the slab so the fixup can reach).
+ * Returns the surviving candidate set.
+ */
+DenseSet<Section*> classifyOverflowSections(LinkGraph& G) {
+ DenseSet<Section*> candidates;
+ for (auto& Sec : G.sections()) {
+ if (Sec.getMemLifetime() == MemLifetime::NoAlloc) continue;
+ StringRef Name = Sec.getName();
+ if (Name.starts_with(".nv_fatbin")) {
+ candidates.insert(&Sec);
+ }
+ }
+ if (candidates.empty()) return candidates;
+
+ const auto& TT = G.getTargetTriple();
+ for (auto& Sec : G.sections()) {
+ for (auto* B : Sec.blocks()) {
+ for (auto& E : B->edges()) {
+ if (!E.isRelocation()) continue;
+ if (isAbsoluteEdge(TT, E.getKind())) continue;
+ if (!E.getTarget().isDefined()) continue;
+ auto* TargetSec = &E.getTarget().getBlock().getSection();
+ candidates.erase(TargetSec);
+ }
+ }
+ if (candidates.empty()) break;
+ }
+ return candidates;
+}
+
} // namespace
// ── Platform abstraction ────────────────────────────────────────────
-void* ArenaJITLinkMemoryManager::reserveVA(size_t size) {
+void* Slab::reserveVA(std::size_t size) {
void* p = ::mmap(nullptr, size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
if (p == MAP_FAILED) return nullptr;
return p;
}
-void ArenaJITLinkMemoryManager::releaseVA(void* addr, size_t size) {
+void Slab::releaseVA(void* addr, std::size_t size) {
int rc = ::munmap(addr, size);
- assert(rc == 0 && "munmap failed in arena destructor");
+ assert(rc == 0 && "munmap failed in Slab destructor");
(void)rc;
}
-Error ArenaJITLinkMemoryManager::commitPages(void* addr, size_t size) {
+Error Slab::commitPages(void* addr, std::size_t size) {
if (size == 0) return Error::success();
- // Commit at slab (2 MB) granularity for THP promotion.
- size_t offset = static_cast<char*>(addr) - arena_base_;
- size_t first_slab = offset / kSlabSize;
- size_t last_slab = (offset + size - 1) / kSlabSize;
-
- for (size_t i = first_slab; i <= last_slab; ++i) {
- if (slab_committed_[i].load(std::memory_order_acquire) != 0) continue;
- size_t slab_offset = i * kSlabSize;
- size_t slab_len = std::min(kSlabSize, arena_capacity_ - slab_offset);
- // mprotect is idempotent, so a concurrent racer calling it on the same slab
+ // Commit at commit-chunk (2 MB) granularity for THP promotion.
+ std::size_t offset = static_cast<char*>(addr) - arena_base_;
+ std::size_t first_chunk = offset / kCommitGranularity;
+ std::size_t last_chunk = (offset + size - 1) / kCommitGranularity;
+
+ for (std::size_t i = first_chunk; i <= last_chunk; ++i) {
+ if (committed_[i].load(std::memory_order_acquire) != 0) continue;
+ std::size_t chunk_offset = i * kCommitGranularity;
+ std::size_t chunk_len = std::min(kCommitGranularity, arena_capacity_ - chunk_offset);
+ // mprotect is idempotent, so a concurrent racer calling it on the same chunk
// is harmless. Only flip the flag after success — otherwise a failed commit
- // followed by freeRegion() would leave slab_committed_[i] == 1, causing a
+ // followed by freeRegion() would leave committed_[i] == 1, causing a
// later allocation to skip mprotect and write into PROT_NONE memory.
- if (::mprotect(arena_base_ + slab_offset, slab_len, PROT_READ | PROT_WRITE) != 0) {
- return make_error<StringError>(
- "ArenaJITLinkMemoryManager: mprotect(RW) failed for slab at offset " +
- formatv("{0:x}", slab_offset) + ": " + std::strerror(errno),
- inconvertibleErrorCode());
+ if (::mprotect(arena_base_ + chunk_offset, chunk_len, PROT_READ | PROT_WRITE) != 0) {
+ return make_error<StringError>("Slab: mprotect(RW) failed for chunk at offset " +
+ formatv("{0:x}", chunk_offset) + ": " +
+ std::strerror(errno),
+ inconvertibleErrorCode());
}
- slab_committed_[i].store(1, std::memory_order_release);
+ committed_[i].store(1, std::memory_order_release);
}
return Error::success();
}
-void ArenaJITLinkMemoryManager::decommitPages(void* addr, size_t size) {
- // Intentionally a no-op for arena pages. The ORC runtime may still reference
+void Slab::decommitPages(void* addr, std::size_t size) {
+ // Intentionally a no-op for slab pages. The ORC runtime may still reference
// deallocated JIT memory during session teardown (e.g., ELFNixPlatform
// deinitializers run after some allocations are freed). Decommitting
// (MADV_DONTNEED or mprotect PROT_NONE) would cause segfaults or illegal
// instructions during shutdown.
//
// Physical pages stay committed but are returned to the free list for reuse.
- // The arena destructor releases all VA and physical memory via munmap.
+ // The slab destructor releases all VA and physical memory via munmap.
(void)addr;
(void)size;
}
-Error ArenaJITLinkMemoryManager::protectPages(void* addr, size_t size, MemProt Prot) {
+Error Slab::protectPages(void* addr, std::size_t size, MemProt Prot) {
int prot = PROT_NONE;
if ((Prot & MemProt::Read) != MemProt::None) prot |= PROT_READ;
if ((Prot & MemProt::Write) != MemProt::None) prot |= PROT_WRITE;
if ((Prot & MemProt::Exec) != MemProt::None) prot |= PROT_EXEC;
if (::mprotect(addr, size, prot) != 0) {
- return make_error<StringError>("ArenaJITLinkMemoryManager: mprotect failed at " +
- formatv("{0:x}", addr) + " size " + formatv("{0:x}", size) +
- ": " + std::strerror(errno),
+ return make_error<StringError>("Slab: mprotect failed at " + formatv("{0:x}", addr) + " size " +
+ formatv("{0:x}", size) + ": " + std::strerror(errno),
inconvertibleErrorCode());
}
if ((Prot & MemProt::Exec) != MemProt::None) {
@@ -167,36 +214,35 @@ Error ArenaJITLinkMemoryManager::protectPages(void* addr, size_t size, MemProt P
return Error::success();
}
-// ── ArenaInFlightAlloc ──────────────────────────────────────────────
+// ── InFlightAlloc ───────────────────────────────────────────────────
-class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManager::InFlightAlloc {
+class Slab::InFlightAlloc : public JITLinkMemoryManager::InFlightAlloc {
public:
// A contiguous region within one pool: [offset, offset + standard_size + finalize_size).
// Standard-lifetime bytes come first; Finalize-lifetime bytes follow and are freed
// at the end of finalize(). Any field may be 0 to indicate no allocation from
// that pool on this call.
struct PoolRegion {
- size_t offset;
- size_t standard_size;
- size_t finalize_size;
+ std::size_t offset;
+ std::size_t standard_size;
+ std::size_t finalize_size;
};
- ArenaInFlightAlloc(ArenaJITLinkMemoryManager& MM, LinkGraph& G, BasicLayout BL,
- PoolRegion non_exec, PoolRegion exec,
- std::vector<OverflowBlock> overflow_blocks)
- : MM(MM),
+ InFlightAlloc(Slab& S, LinkGraph& G, BasicLayout BL, PoolRegion non_exec, PoolRegion exec,
+ std::vector<FinalizedAllocInfo::OverflowBlock> overflow_blocks)
+ : S(S),
G(&G),
BL(std::move(BL)),
non_exec_(non_exec),
exec_(exec),
overflow_blocks_(std::move(overflow_blocks)) {}
- ~ArenaInFlightAlloc() override {
- assert(!G && "ArenaInFlightAlloc destroyed without finalize or abandon");
+ ~InFlightAlloc() override {
+ assert(!G && "Slab::InFlightAlloc destroyed without finalize or abandon");
}
void finalize(OnFinalizedFunction OnFinalized) override {
- // Apply target protections for each arena segment.
+ // Apply target protections for each slab segment.
if (auto Err = applyProtections()) {
OnFinalized(std::move(Err));
return;
@@ -204,7 +250,7 @@ class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManage
// Apply target protections for overflow blocks.
for (auto& ob : overflow_blocks_) {
- if (auto Err = MM.protectPages(ob.addr, ob.size, ob.prot)) {
+ if (auto Err = S.protectPages(ob.addr, ob.size, ob.prot)) {
OnFinalized(std::move(Err));
return;
}
@@ -220,8 +266,8 @@ class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManage
// Decommit finalize-lifetime pages in each pool — they're no longer needed.
for (auto& R : {non_exec_, exec_}) {
if (R.finalize_size > 0) {
- MM.decommitPages(MM.arena_base_ + R.offset + R.standard_size, R.finalize_size);
- MM.freeRegion(R.offset + R.standard_size, R.finalize_size);
+ S.decommitPages(S.arena_base_ + R.offset + R.standard_size, R.finalize_size);
+ S.freeRegion(R.offset + R.standard_size, R.finalize_size);
}
}
@@ -233,19 +279,28 @@ class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManage
// opaque ExecutorAddr (integer), so we must use raw new here. Ownership
// transfers to deallocate(), which LLVM guarantees is called for every
// finalized allocation.
- auto* FA = new FinalizedAllocInfo{
- non_exec_.offset, non_exec_.standard_size, exec_.offset,
- exec_.standard_size, std::move(*DeallocActions), std::move(overflow_blocks_)};
- OnFinalized(FinalizedAlloc(ExecutorAddr::fromPtr(FA)));
+ auto* FA = new FinalizedAllocInfo{&S,
+ non_exec_.offset,
+ non_exec_.standard_size,
+ exec_.offset,
+ exec_.standard_size,
+ std::move(*DeallocActions),
+ std::move(overflow_blocks_)};
+ // Bump the slab's live-alloc counter before publishing the handle so
+ // a concurrent `clearFreeSlabs` in another thread cannot see this
+ // slab as reclaimable between handle publication and the FA reaching
+ // LLJIT's bookkeeping.
+ S.noteAllocated();
+ OnFinalized(JITLinkMemoryManager::FinalizedAlloc(ExecutorAddr::fromPtr(FA)));
}
void abandon(OnAbandonedFunction OnAbandoned) override {
// Decommit and return each pool's full region to the appropriate free list.
for (auto& R : {non_exec_, exec_}) {
- size_t total = R.standard_size + R.finalize_size;
+ std::size_t total = R.standard_size + R.finalize_size;
if (total > 0) {
- MM.decommitPages(MM.arena_base_ + R.offset, total);
- MM.freeRegion(R.offset, total);
+ S.decommitPages(S.arena_base_ + R.offset, total);
+ S.freeRegion(R.offset, total);
}
}
@@ -267,90 +322,76 @@ class ArenaJITLinkMemoryManager::ArenaInFlightAlloc : public JITLinkMemoryManage
const auto& AG = KV.first;
auto& Seg = KV.second;
- auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, MM.page_size_);
- if (auto Err = MM.protectPages(Seg.WorkingMem, SegSize, AG.getMemProt())) return Err;
+ auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, S.page_size_);
+ if (auto Err = S.protectPages(Seg.WorkingMem, SegSize, AG.getMemProt())) return Err;
}
return Error::success();
}
- ArenaJITLinkMemoryManager& MM;
+ Slab& S;
LinkGraph* G;
BasicLayout BL;
PoolRegion non_exec_;
PoolRegion exec_;
- std::vector<OverflowBlock> overflow_blocks_;
+ std::vector<FinalizedAllocInfo::OverflowBlock> overflow_blocks_;
};
-// ── ArenaJITLinkMemoryManager ───────────────────────────────────────
+// ── Slab ────────────────────────────────────────────────────────────
-ArenaJITLinkMemoryManager::ArenaJITLinkMemoryManager(size_t page_size, size_t arena_capacity)
+Slab::Slab(std::size_t page_size, std::size_t capacity)
: arena_base_(nullptr),
- arena_capacity_(arena_capacity),
+ arena_capacity_(capacity),
page_size_(page_size),
midpoint_(0),
exec_bump_limit_(0),
non_exec_bump_(0),
exec_bump_(0) {
- // Try requested capacity, halve on failure down to a minimum floor.
- // The floor is the smaller of kMinArenaCapacity and the requested size,
- // so explicit small arenas (e.g. 16 MB for tests) are honoured.
- // mmap(PROT_NONE | MAP_NORESERVE) can still fail under RLIMIT_AS or
- // extreme VA fragmentation.
- size_t floor = std::min(arena_capacity_, kMinArenaCapacity);
- size_t cap = arena_capacity_;
- while (cap >= floor) {
- arena_base_ = static_cast<char*>(reserveVA(cap));
- if (arena_base_) {
- arena_capacity_ = cap;
- // Partition the arena into two pools at a 2 MB-aligned midpoint.
- // The exec pool starts at midpoint_, which is therefore on a 2 MB
- // boundary — r-x segments pack into a minimum number of 2 MB pages.
- //
- // Constraint: cross-pool displacements (e.g. .text → .rodata via
- // ADRP+ADD on aarch64) must fit in ±kPCRelReach. The farthest pair
- // of bytes is (end of exec, start of non-exec), separated by at most
- // `exec_bump_limit_`, so we cap the exec pool's upper bound at
- // kPCRelReach even when the VA reservation is larger.
- exec_bump_limit_ = std::min(cap, kPCRelReach);
- size_t raw_midpoint = static_cast<size_t>(exec_bump_limit_ * kDefaultNonExecFraction);
- midpoint_ = (raw_midpoint / kSlabSize) * kSlabSize;
- if (midpoint_ == 0) midpoint_ = kSlabSize;
- if (midpoint_ >= exec_bump_limit_) midpoint_ = exec_bump_limit_ - kSlabSize;
- non_exec_bump_ = 0;
- exec_bump_ = midpoint_;
- // Initialize slab commit tracking. make_unique<T[]>(n) value-initializes
- // the array to zero in C++17.
- num_slabs_ = (cap + kSlabSize - 1) / kSlabSize;
- slab_committed_ = std::make_unique<std::atomic<uint8_t>[]>(num_slabs_);
- // Hint THP promotion for the entire arena. Intentionally unchecked —
- // MADV_HUGEPAGE is advisory and may fail if THP is disabled system-wide.
- (void)::madvise(arena_base_, cap, MADV_HUGEPAGE);
- return;
- }
- cap /= 2;
- }
- llvm::report_fatal_error("ArenaJITLinkMemoryManager: failed to reserve at least " +
- Twine(floor / (1024 * 1024)) + " MB of virtual address space");
+ arena_base_ = static_cast<char*>(reserveVA(capacity));
+ if (!arena_base_) return; // Caller inspects isValid() and retries.
+
+ // Partition the slab into two pools at a 2 MB-aligned midpoint. The
+ // exec pool starts at midpoint_, which is therefore on a 2 MB
+ // boundary — r-x segments pack into a minimum number of 2 MB pages.
+ //
+ // Constraint: cross-pool displacements (e.g. .text → .rodata via
+ // ADRP+ADD on aarch64) must fit in ±kPCRelReach. The farthest pair of
+ // bytes is (end of exec, start of non-exec), separated by at most
+ // `exec_bump_limit_`, so we cap the exec pool's upper bound at
+ // kPCRelReach even when the VA reservation is larger.
+ exec_bump_limit_ = std::min(capacity, kPCRelReach);
+ std::size_t raw_midpoint = static_cast<std::size_t>(exec_bump_limit_ * kDefaultNonExecFraction);
+ midpoint_ = (raw_midpoint / kCommitGranularity) * kCommitGranularity;
+ if (midpoint_ == 0) midpoint_ = kCommitGranularity;
+ if (midpoint_ >= exec_bump_limit_) midpoint_ = exec_bump_limit_ - kCommitGranularity;
+ non_exec_bump_ = 0;
+ exec_bump_ = midpoint_;
+ // Initialize commit tracking. make_unique<T[]>(n) value-initializes
+ // the array to zero in C++17.
+ num_commit_chunks_ = (capacity + kCommitGranularity - 1) / kCommitGranularity;
+ committed_ = std::make_unique<std::atomic<std::uint8_t>[]>(num_commit_chunks_);
+ // Hint THP promotion for the entire slab. Intentionally unchecked —
+ // MADV_HUGEPAGE is advisory and may fail if THP is disabled system-wide.
+ (void)::madvise(arena_base_, capacity, MADV_HUGEPAGE);
}
-ArenaJITLinkMemoryManager::~ArenaJITLinkMemoryManager() {
+Slab::~Slab() {
if (arena_base_) {
releaseVA(arena_base_, arena_capacity_);
}
}
-Expected<size_t> ArenaJITLinkMemoryManager::bumpAllocate(size_t size, bool is_exec) {
+Expected<std::size_t> Slab::bumpAllocate(std::size_t size, bool is_exec) {
std::lock_guard<std::mutex> Lock(mu_);
auto& free_list = is_exec ? free_list_exec_ : free_list_non_exec_;
auto& bump = is_exec ? exec_bump_ : non_exec_bump_;
- size_t limit = is_exec ? exec_bump_limit_ : midpoint_;
+ std::size_t limit = is_exec ? exec_bump_limit_ : midpoint_;
// Try free list first (best-fit). O(n) scan — acceptable for the expected
// workload of tens of JIT allocations, not thousands.
- size_t best_idx = free_list.size();
- size_t best_waste = std::numeric_limits<size_t>::max();
- for (size_t i = 0; i < free_list.size(); ++i) {
+ std::size_t best_idx = free_list.size();
+ std::size_t best_waste = std::numeric_limits<std::size_t>::max();
+ for (std::size_t i = 0; i < free_list.size(); ++i) {
if (free_list[i].size >= size && free_list[i].size - size < best_waste) {
best_idx = i;
best_waste = free_list[i].size - size;
@@ -359,7 +400,7 @@ Expected<size_t> ArenaJITLinkMemoryManager::bumpAllocate(size_t size, bool is_ex
}
if (best_idx < free_list.size()) {
- size_t offset = free_list[best_idx].offset;
+ std::size_t offset = free_list[best_idx].offset;
if (free_list[best_idx].size == size) {
free_list.erase(free_list.begin() + best_idx);
} else {
@@ -371,19 +412,15 @@ Expected<size_t> ArenaJITLinkMemoryManager::bumpAllocate(size_t size, bool is_ex
// Bump allocate within the pool's limit.
if (bump + size > limit) {
- return make_error<StringError>(
- std::string("ArenaJITLinkMemoryManager: ") + (is_exec ? "exec" : "non-exec") +
- " pool exhausted (used " + formatv("{0:x}", bump).str() + " + requested " +
- formatv("{0:x}", size).str() + " > limit " + formatv("{0:x}", limit).str() + ")",
- inconvertibleErrorCode());
+ return make_error<SlabPoolExhaustedError>(is_exec ? "exec" : "non-exec", bump, size, limit);
}
- size_t offset = bump;
+ std::size_t offset = bump;
bump += size;
return offset;
}
-void ArenaJITLinkMemoryManager::freeRegion(size_t offset, size_t size) {
+void Slab::freeRegion(std::size_t offset, std::size_t size) {
if (size == 0) return;
std::lock_guard<std::mutex> Lock(mu_);
@@ -392,7 +429,7 @@ void ArenaJITLinkMemoryManager::freeRegion(size_t offset, size_t size) {
// Insert into free list in sorted order.
auto it = std::lower_bound(free_list.begin(), free_list.end(), offset,
- [](const FreeBlock& fb, size_t off) { return fb.offset < off; });
+ [](const FreeBlock& fb, std::size_t off) { return fb.offset < off; });
it = free_list.insert(it, FreeBlock{offset, size});
// Coalesce with next.
@@ -412,54 +449,74 @@ void ArenaJITLinkMemoryManager::freeRegion(size_t offset, size_t size) {
}
}
-void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
- OnAllocatedFunction OnAllocated) {
+Slab::GraphFootprint Slab::computeGraphFootprint(LinkGraph& G, std::size_t page_size) {
+ // Overflow sections live outside any slab (separate mmap at finalize
+ // time) — exclude them from the in-slab footprint. See
+ // classifyOverflowSections() for the candidate-selection rules.
+ DenseSet<Section*> overflow_candidates = classifyOverflowSections(G);
+ SmallVector<std::pair<Section*, MemLifetime>, 4> hidden;
+ for (auto* Sec : overflow_candidates) {
+ hidden.push_back({Sec, Sec->getMemLifetime()});
+ Sec->setMemLifetime(MemLifetime::NoAlloc);
+ }
+ BasicLayout BL(G);
+ for (auto& [Sec, OrigLifetime] : hidden) {
+ Sec->setMemLifetime(OrigLifetime);
+ }
+
+ std::size_t ne_total = 0, e_total = 0;
+ for (auto& KV : BL.segments()) {
+ auto& AG = KV.first;
+ auto& Seg = KV.second;
+ auto SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, page_size);
+ bool is_exec = (AG.getMemProt() & MemProt::Exec) != MemProt::None;
+ if (is_exec) {
+ e_total += SegSize;
+ } else {
+ ne_total += SegSize;
+ }
+ }
+ return GraphFootprint{ne_total, e_total};
+}
+
+std::size_t Slab::capacityForFootprint(GraphFootprint fp, std::size_t base_size) {
+ // Mirrors the split formula in Slab::Slab (see the ctor for the
+ // rationale behind each clamp). Kept local rather than sharing a
+ // helper with the ctor so that changes to the split policy require
+ // a deliberate touch in both places.
+ auto budgets = [](std::size_t cap) {
+ std::size_t exec_limit = std::min(cap, kPCRelReach);
+ std::size_t raw_mid = static_cast<std::size_t>(exec_limit * kDefaultNonExecFraction);
+ std::size_t mid = (raw_mid / kCommitGranularity) * kCommitGranularity;
+ if (mid == 0) mid = kCommitGranularity;
+ if (mid >= exec_limit) mid = exec_limit - kCommitGranularity;
+ return std::pair<std::size_t, std::size_t>{mid, exec_limit - mid};
+ };
+ std::size_t cap = base_size;
+ while (true) {
+ auto [ne_budget, e_budget] = budgets(cap);
+ if (fp.non_exec <= ne_budget && fp.exec <= e_budget) return cap;
+ // Budgets plateau once exec_limit saturates; further VA doesn't help.
+ if (cap >= kPCRelReach) return cap;
+ cap *= 2;
+ }
+}
+
+void Slab::allocate(LinkGraph& G, JITLinkMemoryManager::OnAllocatedFunction OnAllocated) {
// ── Overflow section classification ──
//
// Sections matching known overflow names (e.g. .nv_fatbin — large GPU
// device blobs referenced only by absolute relocations) are allocated
- // outside the arena via separate mmap(), keeping the arena compact for
- // code + small rodata.
- //
- // Two-phase check:
- // Phase 1 — Name-based candidate selection (.nv_fatbin).
- // Phase 2 — Edge validation: any PC-relative cross-section edge
- // targeting a candidate section disqualifies it (the
- // section stays in the arena). This handles cases where
- // the compiler generates ADRP/RIP-relative refs even for
- // data sections.
+ // outside the slab via separate mmap(), keeping the slab compact for
+ // code + small rodata. See classifyOverflowSections() for the
+ // candidate-selection + edge-validation rules.
//
// Validated candidates are temporarily set to NoAlloc so BasicLayout
// skips them, then immediately restored before returning. By the time
// JITLink's fixUpBlocks runs, sections are back to Standard — avoiding
// the debug assert that prohibits edges from allocated sections to
// NoAlloc sections.
- DenseSet<Section*> overflow_candidates;
- for (auto& Sec : G.sections()) {
- if (Sec.getMemLifetime() == MemLifetime::NoAlloc) continue;
- StringRef Name = Sec.getName();
- if (Name.starts_with(".nv_fatbin")) {
- overflow_candidates.insert(&Sec);
- }
- }
-
- // Phase 2: edge validation — disqualify candidates with incoming PC-relative edges.
- if (!overflow_candidates.empty()) {
- const auto& TT = G.getTargetTriple();
- for (auto& Sec : G.sections()) {
- for (auto* B : Sec.blocks()) {
- for (auto& E : B->edges()) {
- if (!E.isRelocation()) continue;
- if (isAbsoluteEdge(TT, E.getKind())) continue;
- // PC-relative edge — if it targets a candidate, disqualify.
- if (!E.getTarget().isDefined()) continue;
- auto* TargetSec = &E.getTarget().getBlock().getSection();
- overflow_candidates.erase(TargetSec);
- }
- }
- if (overflow_candidates.empty()) break;
- }
- }
+ DenseSet<Section*> overflow_candidates = classifyOverflowSections(G);
// Apply: temporarily hide validated overflow sections from BasicLayout.
SmallVector<std::pair<Section*, MemLifetime>, 4> overflow_sections;
@@ -484,19 +541,20 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
return;
}
- if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
+ if (SegsSizes->total() > std::numeric_limits<std::size_t>::max()) {
OnAllocated(make_error<llvm::jitlink::JITLinkError>(
"Total requested size " + formatv("{0:x}", SegsSizes->total()) + " for graph " +
G.getName() + " exceeds address space"));
return;
}
- auto TotalSize = static_cast<size_t>(SegsSizes->total());
+ auto TotalSize = static_cast<std::size_t>(SegsSizes->total());
if (TotalSize == 0 && overflow_sections.empty()) {
// Empty graph — return a no-op allocation.
- OnAllocated(std::make_unique<ArenaInFlightAlloc>(
- *this, G, std::move(BL), ArenaInFlightAlloc::PoolRegion{0, 0, 0},
- ArenaInFlightAlloc::PoolRegion{midpoint_, 0, 0}, std::vector<OverflowBlock>{}));
+ OnAllocated(std::make_unique<InFlightAlloc>(*this, G, std::move(BL),
+ InFlightAlloc::PoolRegion{0, 0, 0},
+ InFlightAlloc::PoolRegion{midpoint_, 0, 0},
+ std::vector<FinalizedAllocInfo::OverflowBlock>{}));
return;
}
@@ -508,8 +566,8 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
//
// Within each pool, Standard segments come first and Finalize segments
// second, so the Finalize tail of each pool can be freed after finalize().
- size_t ne_std_size = 0, ne_fin_size = 0;
- size_t e_std_size = 0, e_fin_size = 0;
+ std::size_t ne_std_size = 0, ne_fin_size = 0;
+ std::size_t e_std_size = 0, e_fin_size = 0;
for (auto& KV : BL.segments()) {
auto& AG = KV.first;
auto& Seg = KV.second;
@@ -522,20 +580,32 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
(is_finalize ? ne_fin_size : ne_std_size) += SegSize;
}
}
- size_t ne_total = ne_std_size + ne_fin_size;
- size_t e_total = e_std_size + e_fin_size;
+ std::size_t ne_total = ne_std_size + ne_fin_size;
+ std::size_t e_total = e_std_size + e_fin_size;
- ArenaInFlightAlloc::PoolRegion ne_region{0, 0, 0};
- ArenaInFlightAlloc::PoolRegion e_region{midpoint_, 0, 0};
+ InFlightAlloc::PoolRegion ne_region{0, 0, 0};
+ InFlightAlloc::PoolRegion e_region{midpoint_, 0, 0};
- auto allocPool = [&](size_t req, bool is_exec) -> Expected<size_t> {
- if (req == 0) return size_t{0};
+ auto allocPool = [&](std::size_t req, bool is_exec) -> Expected<std::size_t> {
+ if (req == 0) return std::size_t{0};
auto off = bumpAllocate(req, is_exec);
if (!off) return off.takeError();
if (auto Err = commitPages(arena_base_ + *off, req)) {
freeRegion(*off, req);
return std::move(Err);
}
+ // Recycled-region protection reset. commitPages() only mprotects a 2 MB
+ // commit-chunk the first time it's touched (guarded by committed_).
+ // When the region was previously handed out, finalized to r-x/r--/rw-
+ // and later returned to the free list by deallocateOne(), those finalize
+ // protections are still in effect — the memset(0) below would fault on
+ // read-only pages. mprotect(RW) has no decommit effect, so the
+ // physical pages stay resident; we're only restoring write access for
+ // the upcoming zero-fill and subsequent JITLink content writes.
+ if (auto Err = protectPages(arena_base_ + *off, req, MemProt::Read | MemProt::Write)) {
+ freeRegion(*off, req);
+ return std::move(Err);
+ }
std::memset(arena_base_ + *off, 0, req);
return *off;
};
@@ -582,7 +652,7 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
Cursor += SegSize;
}
- // Apply layout — copies content and assigns block addresses for arena segments.
+ // Apply layout — copies content and assigns block addresses for slab segments.
if (auto Err = BL.apply()) {
// On error: decommit and free both pool regions.
if (ne_total > 0) {
@@ -597,12 +667,12 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
return;
}
- // ── Allocate overflow sections via mmap() outside the arena ──
- std::vector<OverflowBlock> overflow_allocs;
+ // ── Allocate overflow sections via mmap() outside the slab ──
+ std::vector<FinalizedAllocInfo::OverflowBlock> overflow_allocs;
for (auto& [Sec, _] : overflow_sections) {
// Compute total size for this section's blocks.
- size_t total_sec_size = 0;
+ std::size_t total_sec_size = 0;
for (auto* B : Sec->blocks()) {
total_sec_size = alignTo(total_sec_size, B->getAlignment());
total_sec_size += B->getSize();
@@ -610,7 +680,7 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
if (total_sec_size == 0) continue;
total_sec_size = alignTo(total_sec_size, page_size_);
- // mmap outside the arena.
+ // mmap outside the slab.
void* addr =
::mmap(nullptr, total_sec_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
if (addr == MAP_FAILED) {
@@ -624,10 +694,9 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
decommitPages(arena_base_ + e_region.offset, e_total);
freeRegion(e_region.offset, e_total);
}
- OnAllocated(
- make_error<StringError>("ArenaJITLinkMemoryManager: overflow mmap failed for section " +
- Sec->getName() + ": " + std::strerror(errno),
- inconvertibleErrorCode()));
+ OnAllocated(make_error<StringError>(
+ "Slab: overflow mmap failed for section " + Sec->getName() + ": " + std::strerror(errno),
+ inconvertibleErrorCode()));
return;
}
@@ -636,7 +705,7 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
for (auto* B : Sec->blocks()) {
uint64_t align = B->getAlignment();
ptr = reinterpret_cast<char*>(alignTo(reinterpret_cast<uintptr_t>(ptr), align));
- size_t bsize = B->getSize();
+ std::size_t bsize = B->getSize();
// Copy content and redirect block's mutable content pointer.
if (!B->isZeroFill()) {
auto content = B->getContent();
@@ -651,44 +720,46 @@ void ArenaJITLinkMemoryManager::allocate(const JITLinkDylib* JD, LinkGraph& G,
overflow_allocs.push_back({addr, total_sec_size, Sec->getMemProt()});
}
- OnAllocated(std::make_unique<ArenaInFlightAlloc>(*this, G, std::move(BL), ne_region, e_region,
- std::move(overflow_allocs)));
+ OnAllocated(std::make_unique<InFlightAlloc>(*this, G, std::move(BL), ne_region, e_region,
+ std::move(overflow_allocs)));
}
-void ArenaJITLinkMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
- OnDeallocatedFunction OnDeallocated) {
- Error DeallocErr = Error::success();
-
- for (auto& Alloc : Allocs) {
- // Reclaim ownership of the FinalizedAllocInfo created in finalize().
- auto* FA = Alloc.release().toPtr<FinalizedAllocInfo*>();
-
- // Run deallocation actions in reverse order.
- while (!FA->DeallocActions.empty()) {
- if (auto Err = FA->DeallocActions.back().runWithSPSRetErrorMerged())
- DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
- FA->DeallocActions.pop_back();
+void Slab::deallocateOne(FinalizedAllocInfo* FA, Error& err_out) {
+ // Run deallocation actions in reverse order.
+ while (!FA->DeallocActions.empty()) {
+ if (auto Err = FA->DeallocActions.back().runWithSPSRetErrorMerged()) {
+ err_out = joinErrors(std::move(err_out), std::move(Err));
}
+ FA->DeallocActions.pop_back();
+ }
- // Decommit and free each pool's Standard region.
- if (FA->non_exec_standard_size > 0) {
- decommitPages(arena_base_ + FA->non_exec_offset, FA->non_exec_standard_size);
- freeRegion(FA->non_exec_offset, FA->non_exec_standard_size);
- }
- if (FA->exec_standard_size > 0) {
- decommitPages(arena_base_ + FA->exec_offset, FA->exec_standard_size);
- freeRegion(FA->exec_offset, FA->exec_standard_size);
- }
+ // Decommit and free each pool's Standard region.
+ //
+ // We intentionally do *not* reset page protection here. During session
+ // teardown the ORC runtime may still execute (or read) deallocated JIT
+ // pages from other dylibs while their DeallocActions unwind — same
+ // rationale as decommitPages being a no-op. Protection is reset lazily
+ // in allocate() when the region is re-handed out (see the
+ // "recycled-region protection reset" block in allocPool).
+ if (FA->non_exec_standard_size > 0) {
+ decommitPages(arena_base_ + FA->non_exec_offset, FA->non_exec_standard_size);
+ freeRegion(FA->non_exec_offset, FA->non_exec_standard_size);
+ }
+ if (FA->exec_standard_size > 0) {
+ decommitPages(arena_base_ + FA->exec_offset, FA->exec_standard_size);
+ freeRegion(FA->exec_offset, FA->exec_standard_size);
+ }
- // Release overflow blocks.
- for (auto& ob : FA->overflow_blocks) {
- ::munmap(ob.addr, ob.size);
- }
+ // Decrement the live-alloc counter. After this point the slab may be
+ // observed as reclaimable by `SlabPoolMemoryManager::clearFreeSlabs`.
+ // Safe: all DeallocActions have already run and the region has been
+ // returned to the free list.
+ noteDeallocated();
- delete FA;
+ // Release overflow blocks.
+ for (auto& ob : FA->overflow_blocks) {
+ ::munmap(ob.addr, ob.size);
}
-
- OnDeallocated(std::move(DeallocErr));
}
} // namespace orcjit
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h
new file mode 100644
index 0000000..0029dc5
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_slab.h
@@ -0,0 +1,324 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file orcjit_slab.h
+ * \brief Single contiguous-VA region + dual-pool bump allocator.
+ *
+ * A `Slab` owns one `mmap(PROT_NONE)` reservation of fixed capacity and
+ * bump-allocates from it, keeping all JIT allocations within range of
+ * PC-relative relocations (±2 GB on x86_64, ±4 GB on AArch64).
+ *
+ * The `Slab` is the unit-of-VA-reservation for the OrcJIT memory manager.
+ * Today it is used as a single-slab arena owned by
+ * `ArenaJITLinkMemoryManager`. Stage B of the refactor will introduce a
+ * `SlabPoolMemoryManager` that holds multiple Slabs and grows by mmap-ing
+ * new ones on demand.
+ *
+ * ## Page commit + Transparent Huge Page (THP) support
+ *
+ * Pages are committed in 2 MB chunks (`kCommitGranularity`) — the 2 MB
+ * size matches the Linux huge-page granule on both x86_64 and AArch64,
+ * enabling THP promotion via `madvise(MADV_HUGEPAGE)` on the full
+ * reservation. Each 2 MB commit-chunk is `mprotect`-ed to RW exactly once
+ * via an atomic bitmap flag (`committed_`), avoiding lock contention with
+ * the per-pool allocator mutex.
+ *
+ * ## Dual-pool exec / non-exec split
+ *
+ * The slab is partitioned at a 2 MB-aligned `midpoint_` into two bump
+ * pools:
+ *
+ * non-exec pool = [base, base + midpoint_ )
+ * exec pool = [base + midpoint_, base + exec_bump_limit_ )
+ *
+ * Both pools grow upward; cross-pool displacements (.text → .rodata etc.)
+ * must fit in ±`kPCRelReach` — we cap `exec_bump_limit_` at that reach
+ * even when the VA reservation is larger.
+ */
+#ifndef TVM_FFI_ORCJIT_ORCJIT_SLAB_H_
+#define TVM_FFI_ORCJIT_ORCJIT_SLAB_H_
+
+#include <llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h>
+#include <llvm/ExecutionEngine/Orc/Shared/AllocationActions.h>
+#include <llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h>
+#include <llvm/Support/Error.h>
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <vector>
+
+namespace tvm {
+namespace ffi {
+namespace orcjit {
+
+class Slab; // forward-declared for FinalizedAllocInfo.
+
+/*!
+ * \brief Retriable "slab is out of room for this graph" error.
+ *
+ * Raised from `Slab::bumpAllocate` when the requested bytes would
+ * overflow the pool bump cursor (and no free-list region fits).
+ * `SlabPoolMemoryManager::allocate` catches this error type
+ * specifically, consumes it, and retries on the next slab — or
+ * creates a new one.
+ *
+ * All other `Slab::allocate` failures (mmap, mprotect, JITLink) stay
+ * as `StringError` / `JITLinkError` and are propagated to the caller
+ * without retry.
+ */
+class SlabPoolExhaustedError : public llvm::ErrorInfo<SlabPoolExhaustedError> {
+ public:
+ static char ID;
+ SlabPoolExhaustedError(const char* pool, std::size_t used, std::size_t requested,
+ std::size_t limit);
+ void log(llvm::raw_ostream& os) const override;
+ std::error_code convertToErrorCode() const override;
+
+ private:
+ const char* pool_;
+ std::size_t used_;
+ std::size_t requested_;
+ std::size_t limit_;
+};
+
+/*!
+ * \brief Metadata for a finalized allocation, stored via FinalizedAlloc
+ * handle.
+ *
+ * Each allocate() call may consume a region from either or both pools.
+ * Standard-lifetime pages remain committed after finalize();
+ * Finalize-lifetime pages are decommitted at the end of finalize().
+ * Zero-sized sub-regions indicate no allocation from that pool.
+ *
+ * \p owner points to the Slab that handed out this allocation. With one
+ * slab per session today this is redundant, but stamping it now makes
+ * Stage B's pool-manager routing O(1) without address comparison.
+ */
+struct FinalizedAllocInfo {
+ Slab* owner; ///< Slab that owns these offsets.
+ std::size_t non_exec_offset; ///< offset of non-exec Standard region (or 0 if unused).
+ std::size_t non_exec_standard_size;
+ std::size_t exec_offset; ///< offset of exec Standard region (or midpoint_ if unused).
+ std::size_t exec_standard_size;
+ std::vector<llvm::orc::shared::WrapperFunctionCall> DeallocActions;
+ struct OverflowBlock {
+ void* addr; ///< separately-mmap'd base (outside the slab).
+ std::size_t size; ///< mapping size (page-aligned).
+ llvm::orc::MemProt prot; ///< target protection for finalize.
+ };
+ std::vector<OverflowBlock> overflow_blocks;
+};
+
+/*!
+ * \brief One contiguous JIT memory reservation.
+ *
+ * Exposes the per-graph `allocate` entry point (matching
+ * `JITLinkMemoryManager::allocate`'s callback signature) and the
+ * per-FinalizedAlloc `deallocateOne` used by the outer memory manager to
+ * route deallocation.
+ */
+class Slab {
+ public:
+ // Commit / THP granularity. Every 2 MB chunk is mprotect'd RW exactly
+ // once via `committed_`; `madvise(MADV_HUGEPAGE)` can then promote a
+ // fully-faulted chunk into a single huge page.
+ static constexpr std::size_t kCommitGranularity = std::size_t{2} << 20; // 2 MB
+
+ // PC-relative relocation reach (tightest binding fixup). Cross-pool
+ // references must fit in a signed 32-bit displacement. The binding
+ // constraint on both x86_64 and aarch64 is the signed 32-bit Delta32
+ // used in .eh_frame unwind records (±2 GB), not the wider ADRP+ADD /
+ // RIP-rel reach. `exec_bump_limit_` is capped at this reach so
+ // cross-pool Delta32 fixups always resolve.
+ static constexpr std::size_t kPCRelReach = (std::size_t{1} << 31) - kCommitGranularity; // ~2 GB
+
+ // Fraction of the slab reserved for non-exec segments (r--, rw-). The
+ // remainder holds exec (r-x). Typical CUDA binding objects: ~2 parts
+ // rodata+data to 1 part text.
+ static constexpr double kDefaultNonExecFraction = 2.0 / 3.0;
+
+ /*! \brief Construct a Slab and reserve \p capacity bytes of VA.
+ *
+ * On reservation failure, returns with \c base() == nullptr — the
+ * caller is expected to retry at a smaller capacity or
+ * \c report_fatal_error.
+ */
+ Slab(std::size_t page_size, std::size_t capacity);
+
+ ~Slab();
+
+ Slab(const Slab&) = delete;
+ Slab& operator=(const Slab&) = delete;
+ Slab(Slab&&) = delete;
+ Slab& operator=(Slab&&) = delete;
+
+ /*! \brief True iff the reservation succeeded. */
+ bool isValid() const noexcept { return arena_base_ != nullptr; }
+
+ /*!
+ * \brief Page-aligned per-pool byte totals for a LinkGraph.
+ *
+ * Does not include sections routed to the overflow (separate-mmap)
+ * path. `SlabPoolMemoryManager` uses this to size a fresh slab —
+ * see `capacityForFootprint`.
+ */
+ struct GraphFootprint {
+ std::size_t non_exec;
+ std::size_t exec;
+ std::size_t total() const noexcept { return non_exec + exec; }
+ };
+
+ /*!
+ * \brief Compute pool footprint for \p G without mutating bookkeeping.
+ *
+ * Temporarily reclassifies overflow-section lifetimes while a
+ * `BasicLayout` is built, then restores them — same prefix as
+ * `allocate()`. Safe to call repeatedly; no state is retained.
+ */
+ static GraphFootprint computeGraphFootprint(llvm::jitlink::LinkGraph& G, std::size_t page_size);
+
+ /*!
+ * \brief Power-of-2 capacity that fits \p fp in both pools.
+ *
+ * Starts at \p base_size (the pool's nominal slab size) and doubles
+ * until the Slab split formula yields per-pool budgets ≥ \p fp. The
+ * returned value, when fed back to the Slab ctor, produces a slab
+ * whose non-exec and exec pools are both large enough to host the
+ * graph — no separate "oversize path" needed at the pool layer.
+ *
+ * Stops doubling once the capacity reaches \c kPCRelReach : beyond
+ * that point \c exec_bump_limit_ saturates, so the budgets no longer
+ * grow with VA. For truly gigantic graphs the caller will observe a
+ * subsequent \c Slab::allocate failure — rare in practice because
+ * per-pool budgets at the plateau are ~1.3 GB / ~0.7 GB.
+ */
+ static std::size_t capacityForFootprint(GraphFootprint fp, std::size_t base_size);
+
+ /*! \brief Single-graph JIT allocation entry point. */
+ void allocate(llvm::jitlink::LinkGraph& G,
+ llvm::jitlink::JITLinkMemoryManager::OnAllocatedFunction OnAllocated);
+
+ /*! \brief Per-FA teardown: run DeallocActions, free pool regions,
+ * release overflow blocks. Caller deletes the FA afterwards.
+ *
+ * Errors are joined into \p err_out; never throws.
+ */
+ void deallocateOne(FinalizedAllocInfo* FA, llvm::Error& err_out);
+
+ /*! \brief Address-range ownership check. */
+ bool contains(const void* addr) const noexcept {
+ auto* p = static_cast<const char*>(addr);
+ return p >= arena_base_ && p < arena_base_ + arena_capacity_;
+ }
+
+ char* base() const noexcept { return arena_base_; }
+ std::size_t capacity() const noexcept { return arena_capacity_; }
+ std::size_t page_size() const noexcept { return page_size_; }
+
+ /*!
+ * \brief Record that a new FinalizedAlloc was published from this slab.
+ * Called by `InFlightAlloc::finalize` just before returning the
+ * FinalizedAlloc handle to the caller.
+ */
+ void noteAllocated() noexcept {
+ live_count_.fetch_add(1, std::memory_order_relaxed);
+ ever_used_.store(true, std::memory_order_release);
+ }
+
+ /*!
+ * \brief Record that a FinalizedAlloc on this slab has been released.
+ * Called by `deallocateOne` after the region is returned to the
+ * free list and DeallocActions have run.
+ */
+ void noteDeallocated() noexcept { live_count_.fetch_sub(1, std::memory_order_acq_rel); }
+
+ /*!
+ * \brief True iff this slab has ever hosted a FinalizedAlloc and
+ * currently has zero live ones.
+ *
+ * Used by `SlabPoolMemoryManager::clearFreeSlabs` to decide which
+ * slabs can be munmap'd. A fresh slab that has never been used is
+ * *not* reclaimable — the caller still expects to allocate on it.
+ */
+ bool isReclaimable() const noexcept {
+ return ever_used_.load(std::memory_order_acquire) &&
+ live_count_.load(std::memory_order_acquire) == 0;
+ }
+
+ private:
+ class InFlightAlloc; // defined in orcjit_slab.cc
+
+ /*! \brief Bump-allocate from the selected pool. Returns offset within
+ * the slab's VA reservation. */
+ llvm::Expected<std::size_t> bumpAllocate(std::size_t size, bool is_exec);
+
+ /*! \brief Return a region to the appropriate free list. Pool is
+ * identified by comparing offset against midpoint_. */
+ void freeRegion(std::size_t offset, std::size_t size);
+
+ // ── Platform abstraction (all implemented in orcjit_slab.cc) ──
+ static void* reserveVA(std::size_t size);
+ static void releaseVA(void* addr, std::size_t size);
+ llvm::Error commitPages(void* addr, std::size_t size);
+ static void decommitPages(void* addr, std::size_t size);
+ static llvm::Error protectPages(void* addr, std::size_t size, llvm::orc::MemProt Prot);
+
+ char* arena_base_;
+ std::size_t arena_capacity_;
+ std::size_t page_size_;
+
+ // Dual-pool split. See class docstring.
+ std::size_t midpoint_;
+ std::size_t exec_bump_limit_;
+
+ std::mutex mu_;
+ std::size_t non_exec_bump_; // next free offset in non-exec pool ∈ [0, midpoint_]
+ std::size_t exec_bump_; // next free offset in exec pool ∈ [midpoint_, exec_bump_limit_]
+
+ struct FreeBlock {
+ std::size_t offset;
+ std::size_t size;
+ };
+ std::vector<FreeBlock> free_list_non_exec_;
+ std::vector<FreeBlock> free_list_exec_;
+
+ /*! \brief Per-commit-chunk flags (0 = uncommitted, 1 = committed).
+ * Lock-free: each chunk is mprotect'd exactly once via
+ * compare_exchange. */
+ std::unique_ptr<std::atomic<std::uint8_t>[]> committed_;
+ std::size_t num_commit_chunks_ = 0;
+
+ /*! \brief Count of live FinalizedAllocs held on this slab. */
+ std::atomic<std::size_t> live_count_{0};
+
+ /*! \brief Becomes true on the first `noteAllocated` call; never reset.
+ * A slab that has never been used is not a reclaim candidate
+ * (callers may still plan to allocate on it). */
+ std::atomic<bool> ever_used_{false};
+};
+
+} // namespace orcjit
+} // namespace ffi
+} // namespace tvm
+
+#endif // TVM_FFI_ORCJIT_ORCJIT_SLAB_H_
diff --git a/addons/tvm_ffi_orcjit/tests/test_basic.py b/addons/tvm_ffi_orcjit/tests/test_basic.py
index ca5ba18..c125e6b 100644
--- a/addons/tvm_ffi_orcjit/tests/test_basic.py
+++ b/addons/tvm_ffi_orcjit/tests/test_basic.py
@@ -18,8 +18,10 @@
from __future__ import annotations
+import gc
import sys
import tempfile
+from pathlib import Path
import pytest
import tvm_ffi
@@ -473,3 +475,398 @@ def test_ctor_dtor(v: Variant) -> None:
assert post.index("<fini_array.102>") < post.index("<fini_array.101>")
assert "<ctors>" not in log
assert "<dtors>" not in log
+
+
+# ---------------------------------------------------------------------------
+# Dylib removal — dropping a DynamicLibrary while its session is still alive.
+#
+# Dropping the last Python reference to a DynamicLibrary 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
+# per-dylib deallocation actually happens at drop time.
+# ---------------------------------------------------------------------------
+
+
+@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."""
+ 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
+
+
+@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
+
+
+@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.
+
+ 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
+
+
+@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."""
+ session = ExecutionSession()
+
+ keep = session.create_library("keep")
+ keep.add(obj(v.funcs_obj()))
+
+ drop = session.create_library("drop")
+ drop.add(obj(v.funcs2_obj()))
+ 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.
+ assert keep.get_function(v.fn("test_add"))(5, 5) == 10
+ fresh = session.create_library("fresh")
+ fresh.add(obj(v.funcs2_obj()))
+ 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."""
+ session = ExecutionSession()
+ lib = session.create_library("hold")
+ lib.add(obj(v.funcs_obj()))
+
+ captured = lib.get_function(v.fn("test_add"))
+ del lib
+ gc.collect()
+
+ assert captured(11, 31) == 42
+
+
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_drop_runs_static_destructors(v: Variant) -> None:
+ """Drop runs static destructors immediately rather than at session teardown."""
+ log: list[str] = []
+
+ @tvm_ffi.register_global_func("append_log", override=True)
+ 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"))()
+
+ ctor_len = len(log)
+ assert "<main>" in log, f"main not observed: {log}"
+
+ del lib
+ gc.collect()
+
+ post = log[ctor_len:]
+ # Platform-specific fini markers: ELF (.fini_array / .dtors), Mach-O
+ # (__mod_term_func surfaces as fini_array), COFF (.CRT$XT*).
+ if sys.platform == "win32":
+ markers = ("crt.XT",)
+ else:
+ markers = ("dtors", "fini_array")
+ assert any(m in entry for entry in post for m in markers), (
+ f"no static destructors on drop (platform={sys.platform}): pre={log[:ctor_len]} post={post}"
+ )
+
+
+@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."""
+ session = ExecutionSession()
+
+ 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
+
+ del caller
+ gc.collect()
+
+ 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_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()
+ gc.collect()
+ # Session destructor runs at end-of-test; a regression in removal would
+ # surface as a crash under pytest.
+
+
+# ---------------------------------------------------------------------------
+# 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).
+# ---------------------------------------------------------------------------
+
+
+# 8 MB is the practical floor — Slab::kCommitGranularity is 2 MB and the
+# dual-pool midpoint needs at least two commit chunks of headroom above it,
+# so smaller capacities break the pool layout. See SlabPoolMemoryManager
+# kMinSlabSize.
+_SMALL_SLAB = 8 * 1024 * 1024
+
+
+@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."""
+ session = ExecutionSession(slab_size=_SMALL_SLAB)
+
+ libs = []
+ 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)
+
+ # 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
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only")
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_small_slab_recycles_after_drop(v: Variant) -> None:
+ """Drop returns bytes to the slab's free list for subsequent reuse.
+
+ If the free list weren't working, 32 iterations of load/drop under an
+ 8 MB slab would either exhaust VA or force many new slabs; this test
+ passing on CI containers is the recycling evidence.
+ """
+ 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
+ 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."""
+ 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()))
+
+ # 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]
+ 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)
+
+ 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
+
+
+# ---------------------------------------------------------------------------
+# 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.
+# ---------------------------------------------------------------------------
+
+
+def _build_big_object(tmp_path: Path, byte_size: int, name: str = "big_object") -> str:
+ """Compile a .c file with a large `.rodata` blob to force the oversize path.
+
+ Returns the path to the generated .o file. Skips the test if the build
+ toolchain is unavailable. ``name`` lets callers produce multiple
+ distinctly-sized objects within the same ``tmp_path``.
+ """
+ try:
+ import tvm_ffi.cpp # noqa: PLC0415
+ except ImportError:
+ pytest.skip("tvm_ffi.cpp not available for on-the-fly object build")
+
+ src = tmp_path / f"{name}.c"
+ # Global BSS array — volatile + a runtime read in big_probe() prevents
+ # the compiler from folding sizeof() out and eliding the array. The
+ # array becomes a ZeroFill segment of `byte_size` bytes in the
+ # LinkGraph, which counts toward the slab footprint
+ # (see Slab::computeGraphFootprint).
+ src.write_text(
+ f"""
+ #include <tvm/ffi/c_api.h>
+ #include <stdint.h>
+ volatile unsigned char big_data[{byte_size}];
+ TVM_FFI_DLL_EXPORT int __tvm_ffi_big_probe(void* self,
+ const TVMFFIAny* args,
+ int32_t num_args,
+ TVMFFIAny* result) {{
+ (void)self; (void)args; (void)num_args;
+ /* Reads big_data[0] (always zero for BSS) so the compiler must
+ emit the array; returns the array's size for verification. */
+ result->type_index = kTVMFFIInt;
+ result->zero_padding = 0;
+ result->v_int64 = (int64_t)sizeof(big_data) + (int64_t)big_data[0];
+ return 0;
+ }}
+ """
+ )
+ try:
+ return tvm_ffi.cpp.build(
+ name=name,
+ sources=[str(src)],
+ output=f"{name}.o",
+ extra_cflags=["-O2"],
+ build_directory=str(tmp_path / f".build_{name}"),
+ )
+ except Exception as exc:
+ pytest.skip(f"could not build big object for reclaim test: {exc}")
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only")
+def test_clear_free_slabs_no_drained() -> None:
+ """Calling clear_free_slabs with nothing to reclaim returns 0."""
+ 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)
+ 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.
+
+ 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
+ dedicated oversize slab then hosts exactly one graph — dropping it
+ returns the slab to fully-drained state, so it is reclaimable.
+ """
+ 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
+ gc.collect()
+ # The oversize slab holds only the now-dropped lib'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}"
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only")
+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
+ gc.collect()
+
+ assert session.clear_free_slabs() >= 1
+ assert session.clear_free_slabs() == 0
+
+
+@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.
+
+ Uses two different oversize blob sizes so the two libs 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
+ 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
+
+ # Keep lib — 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
+
+ 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
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="slab pool is Linux-only")
+def test_clear_free_slabs_disabled_pool() -> None:
+ """When the slab pool is disabled, clear_free_slabs is a no-op (returns 0)."""
+ session = ExecutionSession(slab_size=-1)
+ assert session.clear_free_slabs() == 0
diff --git a/addons/tvm_ffi_orcjit/tests/test_arena.py b/addons/tvm_ffi_orcjit/tests/test_memory_manager.py
similarity index 74%
rename from addons/tvm_ffi_orcjit/tests/test_arena.py
rename to addons/tvm_ffi_orcjit/tests/test_memory_manager.py
index a5ad22a..b5cb1ba 100644
--- a/addons/tvm_ffi_orcjit/tests/test_arena.py
+++ b/addons/tvm_ffi_orcjit/tests/test_memory_manager.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Tests for JIT memory arena — verifies co-location and relocation safety.
+"""Tests for JIT memory manager — verifies co-location and relocation safety.
Background
----------
@@ -40,11 +40,13 @@ This matters for **PC-relative relocations with limited range**:
calls use ADRP (page-relative) which has the same scatter problem
at a wider threshold.
-The **arena memory manager** solves this by pre-reserving a contiguous
-VA region (default 4 GB x86_64 / 8 GB AArch64, with fallback to smaller
-sizes) via ``mmap(PROT_NONE)`` and bump-allocating within it, guaranteeing
-all JIT allocations stay within relocation range regardless of external
-VA pressure.
+Our ``SlabPoolMemoryManager`` solves this by pre-reserving one or more
+contiguous VA slabs via ``mmap(PROT_NONE)`` and bump-allocating within
+them, guaranteeing allocations that land on the same slab stay within
+relocation range regardless of external VA pressure. These tests pin
+``slab_size`` large enough that every graph they exercise fits on the
+initial slab, so the "co-located within one slab" property becomes
+observable end-to-end.
Note on ``-fPIC`` vs ``-fpie``
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -59,22 +61,24 @@ problematic relocation type.
Test structure
--------------
-1. **Co-location** (test 1): arena keeps objects within 16 MB.
-2. **Scatter baseline** (test 2): without arena, VA blocker pushes
- objects >128 MB apart — proves arena is responsible for co-location.
+1. **Co-location** (test 1): a single slab keeps objects within its
+ size.
+2. **Scatter baseline** (test 2): with the slab pool disabled, a VA
+ blocker pushes objects far apart — proves the slab pool is
+ responsible for co-location.
3. **Hidden-symbol calls** (test 3): ADRP/PC32 cross-object calls
- succeed under VA pressure with arena.
+ succeed under VA pressure with a slab.
4. **Large data section** (test 4): 4 MB ``.nv_fatbin`` section loads
- correctly within the arena.
+ correctly within the slab.
5. **Overflow section** (test 5): ``.nv_fatbin`` data is allocated
- outside the arena via separate mmap.
+ outside the slab via separate mmap.
6. **Leaked materialization** (test 6): ``__dso_handle`` resolves after
prior sessions leaked mmap slabs from failed materializations.
7. **Delta32 overflow** (test 7): ``-fpie`` GCC objects + 3 GB VA
- blocker. With arena → PASSES; without arena → Delta32 overflow.
+ blocker. With the slab pool → PASSES; without → Delta32 overflow.
-All tests use a small arena (16 MB) and 256 MB-3 GB VA blockers -- safe
-for CI containers.
+All tests use a 256 MB slab and 256 MB-3 GB VA blockers — safe for CI
+containers.
"""
from __future__ import annotations
@@ -170,8 +174,15 @@ _all_variants = _c_variants + _cpp_variants
_is_linux = sys.platform == "linux"
_is_x86_64 = platform.machine() in ("x86_64", "AMD64")
-# Arena test parameters
-_ARENA_SIZE = 16 * 1024 * 1024 # 16MB — small arena for testing
+# Slab test parameters.
+#
+# Under the slab-pool design, `slab_size` is the per-slab capacity — each
+# allocation that exceeds what an existing slab can hold spawns a new one.
+# These tests assert single-slab invariants (objects within slab_size,
+# one contiguous VA region), so `_SLAB_SIZE` must be large enough that
+# every graph we load fits in the first slab. 256 MB comfortably covers
+# the test objects (all < 5 MB each) plus their overhead.
+_SLAB_SIZE = 256 * 1024 * 1024 # 256MB — single-slab headroom for these tests
_BLOCK_RADIUS = 256 * 1024 * 1024 # 256MB — safe for CI containers
_DSO_BLOCK_RADIUS = 3 * 1024 * 1024 * 1024 # 3GB — needed to overflow PC32 (±2GB)
@@ -262,22 +273,22 @@ def free_blockers(blockers: list[tuple[int, int]]) -> None:
# ---------------------------------------------------------------------------
-# Test 1: Arena co-location — objects stay within arena range
+# Test 1: Slab co-location — objects stay within slab range
# ---------------------------------------------------------------------------
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.parametrize("variant", _all_variants)
-def test_arena_colocation(variant: str) -> None:
- """With arena, objects in separate libraries have close code addresses.
+def test_slab_colocation(variant: str) -> None:
+ """With slab, objects in separate libraries have close code addresses.
- Uses a 16MB arena and inserts a 256MB VA blocker between object loads.
- Without the arena, the blocker would push the second object far away.
- With the arena, both objects land within the 16MB region.
+ Uses a 16MB slab and inserts a 256MB VA blocker between object loads.
+ Without the slab, the blocker would push the second object far away.
+ With the slab, both objects land within the 16MB region.
"""
maps_before = set(_parse_maps())
- session = ExecutionSession(arena_size=_ARENA_SIZE)
+ session = ExecutionSession(slab_size=_SLAB_SIZE)
lib1 = session.create_library("lib1")
lib1.add(obj(f"{variant}/test_addr"))
addr1 = lib1.get_function("code_address")()
@@ -296,19 +307,19 @@ def test_arena_colocation(variant: str) -> None:
free_blockers(blockers)
distance = abs(addr1 - addr2)
- assert distance < _ARENA_SIZE, (
- f"Objects should be within {_ARENA_SIZE} bytes, "
+ assert distance < _SLAB_SIZE, (
+ f"Objects should be within {_SLAB_SIZE} bytes, "
f"but distance is {distance} ({distance / (1024**2):.1f} MB)"
)
# ---------------------------------------------------------------------------
-# Test 2: Arena effect — compare with-arena vs without-arena under VA pressure
+# Test 2: Slab effect — compare with-slab vs without-slab under VA pressure
# ---------------------------------------------------------------------------
def _measure_distance_under_pressure(
- variant: str, arena_size: int, radius: int = _BLOCK_RADIUS
+ variant: str, slab_size: int, radius: int = _BLOCK_RADIUS
) -> tuple[int | None, bool]:
"""Load two objects under VA pressure and return (distance, overflowed).
@@ -318,7 +329,7 @@ def _measure_distance_under_pressure(
"""
maps_before = set(_parse_maps())
- session = ExecutionSession(arena_size=arena_size)
+ session = ExecutionSession(slab_size=slab_size)
lib1 = session.create_library("lib1")
lib1.add(obj(f"{variant}/test_addr"))
addr1 = lib1.get_function("code_address")()
@@ -341,74 +352,74 @@ def _measure_distance_under_pressure(
return abs(addr1 - addr2), False
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.parametrize("variant", _all_variants)
-def test_arena_keeps_objects_close(variant: str) -> None:
- """Arena co-locates objects that would otherwise scatter or overflow.
+def test_slab_keeps_objects_close(variant: str) -> None:
+ """Slab co-locates objects that would otherwise scatter or overflow.
Runs the same workload twice under identical VA pressure — once with
- the arena and once without — and compares the outcomes:
+ the slab and once without — and compares the outcomes:
- - **With arena**: both objects must land within the arena size (16 MB).
- - **Without arena**: the blocker should either cause a relocation
+ - **With slab**: both objects must land within the slab size (16 MB).
+ - **Without slab**: the blocker should either cause a relocation
overflow (proving scatter beyond relocation range) or produce a
measurably larger distance.
- The test proves the arena is responsible for co-location by showing a
+ The test proves the slab is responsible for co-location by showing a
strictly better outcome with it enabled. If the VA blocker happens to
be ineffective (e.g., LLVM slab reuse on 64k-page kernels), the test
- still passes as long as the arena keeps objects within range.
+ still passes as long as the slab keeps objects within range.
"""
- # Phase 1: with arena — must always succeed and be within arena range
- arena_dist, arena_overflow = _measure_distance_under_pressure(variant, arena_size=_ARENA_SIZE)
- assert not arena_overflow, "Arena session should not overflow"
- assert arena_dist is not None
- assert arena_dist < _ARENA_SIZE, (
- f"With arena, objects should be within {_ARENA_SIZE} bytes, "
- f"but distance is {arena_dist} ({arena_dist / (1024**2):.1f} MB)"
+ # Phase 1: with slab — must always succeed and be within slab range
+ slab_dist, slab_overflow = _measure_distance_under_pressure(variant, slab_size=_SLAB_SIZE)
+ assert not slab_overflow, "Slab session should not overflow"
+ assert slab_dist is not None
+ assert slab_dist < _SLAB_SIZE, (
+ f"With slab, objects should be within {_SLAB_SIZE} bytes, "
+ f"but distance is {slab_dist} ({slab_dist / (1024**2):.1f} MB)"
)
- # Phase 2: without arena — expect scatter or overflow
- no_arena_dist, no_arena_overflow = _measure_distance_under_pressure(variant, arena_size=-1)
+ # Phase 2: without slab — expect scatter or overflow
+ no_slab_dist, no_slab_overflow = _measure_distance_under_pressure(variant, slab_size=-1)
- if no_arena_overflow:
- # Relocation overflow without arena proves the blocker forced
- # scatter beyond relocation range — arena prevented this.
+ if no_slab_overflow:
+ # Relocation overflow without slab proves the blocker forced
+ # scatter beyond relocation range — slab prevented this.
return
- assert no_arena_dist is not None
- if no_arena_dist > arena_dist:
- # Without arena produced a larger distance — arena effect shown.
+ assert no_slab_dist is not None
+ if no_slab_dist > slab_dist:
+ # Without slab produced a larger distance — slab effect shown.
return
- # Blocker was ineffective (both distances are small). The arena
+ # Blocker was ineffective (both distances are small). The slab
# assertion above already passed, which is the key property. We
- # cannot distinguish arena effect from lucky placement here.
+ # cannot distinguish slab effect from lucky placement here.
pytest.skip(
- f"VA blocker ineffective: arena={arena_dist / 1024:.0f} KB, "
- f"no-arena={no_arena_dist / 1024:.0f} KB — "
- f"cannot demonstrate arena effect on this kernel"
+ f"VA blocker ineffective: slab={slab_dist / 1024:.0f} KB, "
+ f"no-slab={no_slab_dist / 1024:.0f} KB — "
+ f"cannot demonstrate slab effect on this kernel"
)
# ---------------------------------------------------------------------------
-# Test 3: Hidden-symbol ADRP/PC32 relocation with arena + blocker
+# Test 3: Hidden-symbol ADRP/PC32 relocation with slab + blocker
# ---------------------------------------------------------------------------
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.parametrize("variant", _all_variants)
-def test_arena_hidden_symbol_with_blocker(variant: str) -> None:
- """Arena prevents hidden-visibility relocation overflow under VA pressure.
+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 arena, the blocker would push objects apart causing overflow.
- With the arena, both objects are co-located and the call succeeds.
+ Without slab, the blocker would push objects apart causing overflow.
+ With the slab, both objects are co-located and the call succeeds.
"""
maps_before = set(_parse_maps())
- session = ExecutionSession(arena_size=_ARENA_SIZE)
+ session = ExecutionSession(slab_size=_SLAB_SIZE)
lib = session.create_library("hidden_test")
# Load helper and force materialization
@@ -434,14 +445,14 @@ def test_arena_hidden_symbol_with_blocker(variant: str) -> None:
# ---------------------------------------------------------------------------
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.parametrize("variant", _all_variants)
def test_large_data_section(variant: str) -> None:
"""Load object with a 4MB .nv_fatbin section — basic correctness.
The .nv_fatbin section is referenced only by absolute relocations,
so it can live anywhere. This test verifies the object loads and
- the function works. The 4MB section fits in the arena.
+ the function works. The 4MB section fits in the slab.
"""
session = ExecutionSession()
lib = session.create_library("fatbin")
@@ -451,25 +462,25 @@ def test_large_data_section(variant: str) -> None:
# ---------------------------------------------------------------------------
-# Test 5: Overflow section — .nv_fatbin lands outside the arena
+# Test 5: Overflow section — .nv_fatbin lands outside the slab
# ---------------------------------------------------------------------------
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.parametrize("variant", _all_variants)
-def test_overflow_section_outside_arena(variant: str) -> None:
- """Overflow sections (.nv_fatbin) are allocated outside the arena.
+def test_overflow_section_outside_slab(variant: str) -> None:
+ """Overflow sections (.nv_fatbin) are allocated outside the slab.
- The arena memory manager detects sections named .nv_fatbin and
- allocates them via a separate mmap() outside the arena. This keeps
- the arena compact for code + small rodata, reducing 2MB THP region
+ The slab memory manager detects sections named .nv_fatbin and
+ allocates them via a separate mmap() outside the slab. This keeps
+ the slab compact for code + small rodata, reducing 2MB THP region
count and iTLB pressure.
- Verification: get the fatbin data address and the arena VA range
+ Verification: get the fatbin data address and the slab VA range
from /proc/self/maps, then assert the fatbin address is NOT within
- the arena region.
+ the slab region.
"""
- session = ExecutionSession(arena_size=_ARENA_SIZE)
+ session = ExecutionSession(slab_size=_SLAB_SIZE)
lib = session.create_library("fatbin_overflow")
lib.add(obj(f"{variant}/fake_fatbin"))
@@ -479,16 +490,16 @@ def test_overflow_section_outside_arena(variant: str) -> None:
# Get the actual address of the fatbin data in memory.
fatbin_addr = lib.get_function("get_fatbin_addr")()
- # Find the arena mapping: a single large region matching the arena size.
- # The arena is reserved as PROT_NONE and then committed in slabs, so
- # look for the contiguous region that spans _ARENA_SIZE.
+ # 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
+ # look for the contiguous region that spans _SLAB_SIZE.
maps = _parse_maps()
- arena_regions = [(s, e) for s, e in maps if (e - s) >= _ARENA_SIZE]
+ slab_regions = [(s, e) for s, e in maps if (e - s) >= _SLAB_SIZE]
- # The fatbin address must not fall within any arena-sized region.
- for start, end in arena_regions:
+ # The fatbin address must not fall within any slab-sized region.
+ for start, end in slab_regions:
assert not (start <= fatbin_addr < end), (
- f"Fatbin data at {fatbin_addr:#x} should be OUTSIDE the arena "
+ f"Fatbin data at {fatbin_addr:#x} should be OUTSIDE the slab "
f"[{start:#x}, {end:#x}) but landed inside"
)
@@ -520,14 +531,14 @@ def test_dso_handle_relocation_after_failed_materialization(variant: str) -> Non
leaked slabs accumulate and push subsequent ``mmap`` allocations to
higher addresses.
- The arena prevents overflow because all allocations — both
+ The slab prevents overflow because all allocations — both
``__dso_handle``'s ``LinkGraph`` and the code ``LinkGraph`` — land
within the same contiguous pre-reserved VA region.
- Without arena: may FAIL on x86_64 with GCC PIE objects after
+ Without slab: may FAIL on x86_64 with GCC PIE objects after
repeated leaked materializations push slabs >2 GB
apart.
- With arena: PASSES (all allocations in same arena).
+ With slab: PASSES (all allocations in same slab).
"""
# Step 1: Trigger leaked materializations to consume low VA space.
leaked_sessions = []
@@ -554,22 +565,22 @@ def test_dso_handle_relocation_after_failed_materialization(variant: str) -> Non
# ---------------------------------------------------------------------------
-# Test 6: __dso_handle Delta32 overflow — arena prevents it (x86_64 PIE)
+# Test 6: __dso_handle Delta32 overflow — slab prevents it (x86_64 PIE)
#
# GCC -fpie objects use R_X86_64_PC32 (±2GB) for __dso_handle.
# ELFNixPlatform's DSOHandleMaterializationUnit allocates __dso_handle
# in a separate LinkGraph from the code. Under VA pressure, these two
# allocations can land >2GB apart, overflowing the Delta32 fixup.
-# The arena keeps them co-located within relocation range.
+# The slab keeps them co-located within relocation range.
# ---------------------------------------------------------------------------
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.skipif(not _is_x86_64, reason="Delta32 overflow requires x86_64")
@pytest.mark.skipif(not _pie_cpp_variants, reason="No GCC PIE C++ variants built")
@pytest.mark.parametrize("variant", _pie_cpp_variants or ["skip"])
-def test_dso_handle_delta32_with_arena(variant: str) -> None:
- """Arena prevents __dso_handle Delta32 overflow under VA pressure.
+def test_dso_handle_delta32_with_slab(variant: str) -> None:
+ """Slab prevents __dso_handle Delta32 overflow under VA pressure.
Root cause
----------
@@ -589,22 +600,22 @@ def test_dso_handle_delta32_with_arena(variant: str) -> None:
Test strategy
-------------
- 1. Create a session with arena enabled (16 MB).
+ 1. Create a session with slab enabled (16 MB).
2. Load PIE GCC objects into lib1 — this triggers materialization of
both ``__dso_handle`` (via ``DSOHandleMaterializationUnit``) and
- the code (via ``lib.add``), all within the arena.
- 3. Block 3 GB of VA around the first allocation — without arena this
+ the code (via ``lib.add``), all within the slab.
+ 3. Block 3 GB of VA around the first allocation — without slab this
would force the next ``mmap`` to land >2 GB away.
- 4. Load a second PIE GCC object into lib2 — with arena, this still
+ 4. Load a second PIE GCC object into lib2 — with slab, this still
lands within the 16 MB region.
5. Assert the function call succeeds — proves Delta32 is in range.
- See ``test_dso_handle_delta32_overflow_without_arena`` for the
- counterpart proving the overflow occurs without arena.
+ See ``test_dso_handle_delta32_overflow_without_slab`` for the
+ counterpart proving the overflow occurs without slab.
"""
maps_before = set(_parse_maps())
- session = ExecutionSession(arena_size=_ARENA_SIZE)
+ 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
@@ -623,15 +634,15 @@ def test_dso_handle_delta32_with_arena(variant: str) -> None:
free_blockers(blockers)
-@pytest.mark.skipif(not _is_linux, reason="Arena is Linux-only")
+@pytest.mark.skipif(not _is_linux, reason="Slab is Linux-only")
@pytest.mark.skipif(not _is_x86_64, reason="Delta32 overflow requires x86_64")
@pytest.mark.skipif(not _pie_cpp_variants, reason="No GCC PIE C++ variants built")
@pytest.mark.parametrize("variant", _pie_cpp_variants or ["skip"])
-def test_dso_handle_delta32_overflow_without_arena(variant: str) -> None:
- """Without arena, PIE __dso_handle PC32 overflows under VA pressure.
+def test_dso_handle_delta32_overflow_without_slab(variant: str) -> None:
+ """Without slab, PIE __dso_handle PC32 overflows under VA pressure.
- Same setup as ``test_dso_handle_delta32_with_arena`` but with arena
- disabled (``arena_size=-1``).
+ Same setup as ``test_dso_handle_delta32_with_slab`` but with slab
+ disabled (``slab_size=-1``).
The 3 GB VA blocker fills all free gaps within ±3 GB of the first
session's JIT allocations. When lib2 is loaded, ``InProcessMemoryMapper``
@@ -643,14 +654,14 @@ def test_dso_handle_delta32_overflow_without_arena(variant: str) -> None:
±2 GB → JITLink reports ``Delta32 fixup ... is out of range``.
The test accepts both outcomes:
- - **Exception** (PC32 overflow): proves the arena is needed.
+ - **Exception** (PC32 overflow): proves the slab is needed.
- **Success** (GOTPCRELX used): GCC chose GOT-relative despite
- ``-fpie`` — no overflow possible, but the arena is still
+ ``-fpie`` — no overflow possible, but the slab is still
beneficial for other relocation types.
"""
maps_before = set(_parse_maps())
- session = ExecutionSession(arena_size=-1) # arena disabled
+ session = ExecutionSession(slab_size=-1) # slab disabled
lib1 = session.create_library("lib1")
lib1.add(obj(f"{variant}/test_addr"))
lib1.get_function("code_address")()
@@ -668,7 +679,7 @@ def test_dso_handle_delta32_overflow_without_arena(variant: str) -> None:
# If we get here, GCC used GOTPCRELX — no overflow.
assert result == 1030
except Exception:
- # R_X86_64_PC32 overflow as expected — proves arena is needed.
+ # R_X86_64_PC32 overflow as expected — proves slab is needed.
pass
finally:
free_blockers(blockers)