[ORCJIT] Add shared session and high-level load_module

.github/actions/build-orcjit-wheel/action.yml

@@ -102,6 +102,10 @@ runs:
         CIBW_ENVIRONMENT: LLVM_PREFIX=/opt/llvm
         CIBW_ENVIRONMENT_WINDOWS: LLVM_PREFIX="C:/opt/llvm"
         CIBW_CONTAINER_ENGINE: "docker; create_args: --volume /opt/llvm:/opt/llvm"
+        # Build against the in-tree core (not PyPI) so the addon's tvm_ffi ABI
+        # matches the core the tests reinstall; a mismatch segfaults the JIT.
+        CIBW_BEFORE_BUILD: pip install "scikit-build-core>=0.10.0" cmake ninja "{project}"
+        CIBW_BUILD_FRONTEND: "pip; args: --no-build-isolation"
         CIBW_TEST_REQUIRES: pytest pytest-xdist ninja
         CIBW_TEST_COMMAND: >-
           pip install --force-reinstall {project} &&

addons/tvm_ffi_orcjit/CMakeLists.txt

@@ -149,9 +149,27 @@ if (APPLE)
     COMMENT "Fixing libc++ rpath to use system library"
   )
 elseif (UNIX AND NOT WIN32)
+  # Static libstdc++/libgcc: the conda LLVM toolchain's libstdc++ is newer than the manylinux ABI
+  # floor, so linking it dynamically would import too-recent GLIBCXX_* symbols and fail auditwheel.
   target_link_options(tvm_ffi_orcjit PRIVATE -static-libstdc++ -static-libgcc)
 endif ()
 
+# Hide symbols pulled from static archives so they are not re-exported and cannot interpose with a
+# host process's own copies (e.g. PyTorch's bundled LLVM). Exclude LLVM/zlib/zstd by name -- NOT
+# --exclude-libs,ALL, which would also hide the static libstdc++ symbols (operator new, __throw_*,
+# ...) that liborc_rt and JIT'd C++ code resolve from this .so at run time. The LLVM list is derived
+# from llvm-config so it tracks the toolchain version.
+if (CMAKE_SYSTEM_NAME MATCHES "Linux|Android|FreeBSD|NetBSD|OpenBSD" AND CMAKE_CXX_COMPILER_ID
+                                                                         MATCHES "GNU|Clang"
+)
+  # --exclude-libs needs exact archive filenames (no globs), so turn each -lLLVM* from llvm-config
+  # into libLLVM*.a and add zlib/zstd. One colon-separated list -- a -Wl comma list would be split.
+  set(_exclude_archives ${_llvm_libs_list} "-lz" "-lzstd")
+  list(TRANSFORM _exclude_archives REPLACE "^-l(.+)$" "lib\\1.a")
+  list(JOIN _exclude_archives ":" _exclude_libs_arg)
+  target_link_options(tvm_ffi_orcjit PRIVATE "-Wl,--exclude-libs=${_exclude_libs_arg}")
+endif ()
+
 # ---- Find and bundle liborc_rt ----
 # Platform notes:
 #

addons/tvm_ffi_orcjit/ORCJIT_PRIMER.md

@@ -345,128 +345,40 @@ With the background above, here is how the addon maps onto the concepts.
 ```text
 Python / C++ API                   LLVM ORC v2 concept
 ─────────────────────────────────────────────────────────────────
+default_session()                  shared (leaked) LLJIT + ExecutionSession
 ExecutionSession                   LLJIT + ExecutionSession
-DynamicLibrary                     JITDylib
-dylib.add("foo.o")                 ObjectLinkingLayer.add(buffer)
-dylib.get_function("add")          ExecutionSession.lookup("__tvm_ffi_add")
-dylib.set_link_order([a, b])       JITDylib.setLinkOrder([a, b, main, ...])
+session.load_module(objs)          createJITDylib + addObjectFile(s) + wire imports
+(returned) Module                  JITDylib, wrapped as a tvm_ffi.Module
+module.get_function("add")         ExecutionSession.lookup("__tvm_ffi_add")
 ```
 
-### 4.2 Object File Loading Pipeline
-
-```text
-dylib.add("foo.o")
-   │
-   ▼ ORCJITDynamicLibraryObj::AddObjectFile()
-   │   jit_->addObjectFile(*dylib_, MemoryBuffer)
-   │
-   ▼ ObjectTransformLayer                    (macOS: skipped; Linux/Win: may strip .pdata)
-   │   Windows: strip .pdata/.xdata           avoids JITLink COMDAT limitation
-   │            fix __ImageBase               fixes Pointer32NB relocations
-   │
-   ▼ ObjectLinkingLayer (JITLink)
-   │   Parse object → LinkGraph
-   │   Run InitFiniPlugin passes:
-   │     PrePrunePasses:    mark .init_array / .ctors blocks as live
-   │     PostAllocationPasses: (Windows) set __ImageBase, strip SEH sections
-   │     PostFixupPasses:   extract resolved init/fini function pointers
-   │                        → session.AddPendingInitializer(dylib, entry)
-   │   Resolve relocations via ExecutionSession.lookup()
-   │   Allocate JIT memory, write code, mark executable
-   │
-   ▼ JITDylib symbol table updated
-      (symbols are defined but constructors not yet run)
-```
-
-### 4.3 Symbol Lookup and Initialization
-
-```text
-dylib.get_function("add")
-   │
-   ▼ ORCJITDynamicLibraryObj::GetFunction("add")
-   │   → GetSymbol("__tvm_ffi_add")
-   │       build JITDylibSearchOrder: [this, linked dylibs, LLJIT default]
-   │       jit_->getExecutionSession().lookup(order, "__tvm_ffi_add")
-   │       ← ExecutorAddr
-   │
-   ▼ Linux/Windows: RunPendingInitializers()
-   │   Sort entries by priority
-   │   Call each function pointer (C++ ctors, .init_array entries)
-   │
-   ▼ macOS: jit_->initialize(*dylib_)
-   │   ORC MachOPlatform calls __mod_init_func pointers
-   │
-   ▼ Wrap raw function pointer as tvm_ffi::Function
-      via Function::FromPacked lambda (marshals AnyView args ↔ TVMFFIAny)
-```
-
-### 4.4 InitFiniPlugin Detail
-
-`InitFiniPlugin` is an `ObjectLinkingLayer::Plugin` — it receives callbacks during
-JITLink's pass pipeline for every object being linked.
-
-```text
-JITLink pass pipeline for each object file:
-─────────────────────────────────────────────────────────────
-PrePrunePasses
-  → InitFiniPlugin::modifyPassConfig (PrePrunePasses)
-      For each section named .init_array* / .ctors* / .fini_array* / .dtors*
-      (ELF) or __DATA,__mod_init_func (Mach-O) or .CRT$XC* (COFF):
-        Mark all blocks in that section as "keep" (live)
-        → Prevents dead-code elimination from removing constructors
-
-PostAllocationPasses (Windows only)
-  → InitFiniPlugin::modifyPassConfig (PostAllocationPasses)
-      Set __ImageBase = lowest allocated block address
-      Strip .pdata / .xdata exception handler sections
-
-PostFixupPasses
-  → InitFiniPlugin::modifyPassConfig (PostFixupPasses)
-      Iterate all blocks in init/fini sections
-      Read each 8-byte slot as an ExecutorAddr
-      Parse section name to determine: priority, is_init vs is_fini
-      Call session->AddPendingInitializer(dylib, InitFiniEntry{addr, section, priority})
-      (or AddPendingDeinitializer for dtors/fini_array)
-```
-
-### 4.5 Cross-Library Symbol Resolution
-
-```text
-libA depends on a symbol defined in libB:
-
-libA.set_link_order([libB])
-    → JITDylibSearchOrder for libA: [libA, libB, main(ProcessSymbols)]
-
-When libA's object file has an unresolved symbol "foo":
-    JITLink asks ExecutionSession.lookup([libA, libB, main], "foo")
-    → found in libB → returns libB's ExecutorAddr for "foo"
-    → relocation in libA patched with that address
-```
-
-### 4.6 Windows DLL Import Stubs
-
-Windows MSVC objects reference DLL functions through `__imp_XXX` pointer stubs (the
-Import Address Table pattern). At static link time the linker creates these stubs. In
-JIT mode there is no linker, so `DLLImportDefinitionGenerator` creates them on demand:
-
-```text
-JITLink encounters undefined symbol "__imp_malloc"
-   │
-   ▼ DLLImportDefinitionGenerator::tryToGenerate()
-   │   Search ucrtbase.dll, msvcrt.dll, then all process modules
-   │   for the real address of "malloc"
-   │
-   ▼ Allocate two JIT-memory stubs:
-   │   __imp_malloc  → 8-byte slot containing &malloc (host address)
-   │   malloc        → x86_64 jmp [__imp_malloc] trampoline
-   │
-   ▼ Define both symbols in the JITDylib
-      → JITLink can now apply PCRel32 reloc to __imp_malloc (stub is close in JIT memory)
-```
-
-The stubs must live in JIT-allocated memory (not at the host process address) because
-x86_64 `PCRel32` relocations can only reach ±2 GB. The host's `malloc` may be farther
-than 2 GB from the JIT allocation.
+`session.load_module(objects)` is the sole public loader: it creates one
+`JITDylib`, adds every object (path or in-memory bytes), injects context symbols
+eagerly, expands any embedded library binary into an import tree, and returns a
+plain `tvm_ffi.Module`. There is no incremental library API — a module is a
+unit: load all its objects at once, then look up functions on the result.
+
+### 4.2 Loading and lookup
+
+`load_module` adds each object to the JITDylib, and JITLink parses it into a
+`LinkGraph`, resolves relocations, and allocates executable JIT memory.
+`get_function` looks the symbol up (materializing lazily), then wraps the raw
+pointer as a `tvm_ffi::Function`. Symbols resolve against the dylib's own
+default link order (this dylib → Platform → process/runtime symbols); objects
+that reference each other must be loaded together, since there is no linking
+between separate `load_module` results.
+
+Two addon-specific pieces sit in this pipeline:
+
+- **`InitFiniPlugin`** — a JITLink pass plugin that keeps init/fini sections
+  (`.init_array`/`.ctors`/`.fini_array`/`.dtors`, `__mod_init_func`, `.CRT$XC*`)
+  live, then collects their function pointers after fixup. The addon runs them
+  in priority order at first lookup and at teardown, replacing the ORC
+  platform's initializer machinery. See `llvm_patches/init_fini_plugin.h`.
+- **Windows DLL import stubs** — `DLLImportDefinitionGenerator` resolves
+  `__imp_XXX` references to host-DLL functions by emitting JIT-memory pointer +
+  trampoline stubs, keeping `PCRel32` fixups within ±2 GB of the JIT code. See
+  `llvm_patches/win_dll_import_generator.h`.
 
 ---
 
@@ -475,22 +387,17 @@ than 2 GB from the JIT allocation.
 ```python
 import tvm_ffi_orcjit as oj
 
-# 1. Create an ExecutionSession (wraps LLJIT)
-sess = oj.ExecutionSession()
-
-# 2. Create a JITDylib
-lib = sess.create_dynamic_library()
-
-# 3. Load a compiled object file
-#    → object parsed, JITLink links it, InitFiniPlugin collects ctors
-lib.add("add.o")
+# 1. Get the shared ExecutionSession (wraps LLJIT; created once per process)
+sess = oj.default_session()
 
-# 4. Look up a function
-#    → LLVM resolves "__tvm_ffi_add", RunPendingInitializers() fires ctors
-add = lib.get_function("add")   # returns tvm_ffi.Function
+# 2. Load a compiled object file into a fresh JITDylib
+#    → object parsed, JITLink links it, InitFiniPlugin collects ctors,
+#      context symbols injected eagerly, embedded binary (if any) expanded
+mod = sess.load_module("add.o")   # returns a tvm_ffi.Module
 
-# 5. Call it
-result = add(3, 4)   # → 7
+# 3. Look up and call a function
+#    → LLVM resolves "__tvm_ffi_add"; pending constructors fire on first lookup
+result = mod.add(3, 4)   # → 7
 ```
 
 The corresponding C++ side of `add.o`:
@@ -511,7 +418,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl);
 
 | Concept | What it is | Where in the addon |
 | --- | --- | --- |
-| Object file | Container of machine code, data, symbols, relocations | Input to `dylib.add()` |
+| Object file | Container of machine code, data, symbols, relocations | Input to `session.load_module()` |
 | Relocation | Recipe to patch a code address at link/JIT time | Applied by JITLink |
 | `.init_array` / `.ctors` | Array of C++ constructor pointers in ELF objects | Collected by `InitFiniPlugin` |
 | `ExecutionSession` | Root of the ORC JIT environment | `ORCJITExecutionSessionObj` |
@@ -520,7 +427,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl);
 | `JITLink` | LLVM's JIT-aware linker | Used inside `ObjectLinkingLayer` |
 | JITLink pass pipeline | Pre-prune → post-alloc → post-fixup hooks | Where `InitFiniPlugin` runs |
 | `DefinitionGenerator` | Fallback symbol provider | `DLLImportDefinitionGenerator` (Win) |
-| Link order | Search path across JITDylibs for symbol resolution | `SetLinkOrder()` |
+| Link order | Search path across JITDylibs for symbol resolution | LLJIT default (Main → Platform → ProcessSymbols) |
 | `__tvm_ffi_` prefix | Namespace for TVM-FFI exported functions | Used in `GetFunction()` |
 
 ---

addons/tvm_ffi_orcjit/README.md

@@ -24,9 +24,10 @@ TVM-FFI exported functions.
 ## Features
 
 - **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2
-- **Multiple Libraries**: Create separate dynamic libraries with independent symbol namespaces
-- **Incremental Loading**: Add multiple object files to the same library incrementally
-- **Symbol Isolation**: Different libraries can define the same symbol without conflicts
+- **High-Level Loading**: `default_session().load_module(...)` mirrors `tvm_ffi.load_module`, returning a plain `tvm_ffi.Module`
+- **Unified Input**: Load from a file path, in-memory object bytes, or a list mixing both
+- **Shared Session**: A process-wide session so multiple callers share one JIT environment (process symbols, arena, linking)
+- **Symbol Isolation**: Separate `load_module` calls define independent symbol namespaces, so they can define the same symbol without conflicts
 - **Init/Fini Support**: Handles static constructors/destructors across ELF (`.init_array`/`.ctors`), Mach-O (`__mod_init_func`), and COFF (`.CRT$XC*`/`.CRT$XT*`)
 - **Cross-Platform**: Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64)
 - **Multi-Compiler**: Tested with LLVM Clang, GCC, Apple Clang, MSVC, and clang-cl
@@ -105,55 +106,54 @@ auto-discover it and `LLVM_PREFIX` is not needed.
 
 ### Basic Example
 
-```python
-from tvm_ffi_orcjit import ExecutionSession
+The high-level API mirrors `tvm_ffi.load_module`: a process-wide shared session
+plus a `load_module` that accepts a path, in-memory object bytes, or a list of
+either, and returns a plain `tvm_ffi.Module`.
 
-# Create an execution session
-session = ExecutionSession()
+```python
+import tvm_ffi_orcjit as oj
 
-# Create a dynamic library
-lib = session.create_library()
+# Shared process-wide session (created once, cached).
+session = oj.default_session()
 
-# Load an object file
-lib.add("example.o")
+# Load a single object file by path.
+mod = session.load_module("example.o")
 
-# Get and call a function
-add_func = lib.get_function("add")
-result = add_func(1, 2)
+# Call an exported function.
+result = mod.add(1, 2)
 print(f"Result: {result}")  # Output: Result: 3
 ```
 
-### Multiple Libraries with Symbol Isolation
-
-```python
-session = ExecutionSession()
+### Loading Multiple Objects and In-Memory Bytes
 
-lib1 = session.create_library("lib1")
-lib2 = session.create_library("lib2")
+Objects passed together are linked into one module (the same way a multi-object
+shared library links). Each element may be a path or an object-file image in
+memory.
 
-lib1.add("implementation_v1.o")
-lib2.add("implementation_v2.o")
+```python
+from pathlib import Path
 
-add_v1 = lib1.get_function("add")
-add_v2 = lib2.get_function("add")
+session = oj.default_session()
 
-print(add_v1(5, 3))  # Uses implementation from lib1
-print(add_v2(5, 3))  # Uses implementation from lib2
+mod = session.load_module(
+    [
+        "math_ops.o",                    # path
+        Path("kernel.o").read_bytes(),   # in-memory object bytes
+    ]
+)
+result = mod.call_math(10, 20)
 ```
 
-### Cross-Library Linking
-
-```python
-session = ExecutionSession()
+### Isolated Sessions
 
-base_lib = session.create_library("base")
-base_lib.add("math_ops.o")
+`default_session()` is shared across the process. For an isolated symbol
+namespace or a tuned memory arena, construct an `ExecutionSession` directly:
 
-caller_lib = session.create_library("caller")
-caller_lib.set_link_order(base_lib)  # Can resolve symbols from base_lib
-caller_lib.add("caller.o")
+```python
+from tvm_ffi_orcjit import ExecutionSession
 
-result = caller_lib.get_function("call_math")(10, 20)
+session = ExecutionSession()          # independent LLVM ExecutionSession
+mod = session.load_module("impl.o")
 ```
 
 ## Writing Functions for OrcJIT
@@ -211,13 +211,12 @@ tvm_ffi_orcjit/
 ├── src/ffi/
 │   ├── orcjit_session.cc       # ExecutionSession (LLJIT setup, plugins)
 │   ├── orcjit_session.h
-│   ├── orcjit_dylib.cc         # DynamicLibrary (object loading, symbol lookup)
+│   ├── orcjit_dylib.cc         # JIT dylib module (object loading, symbol lookup)
 │   ├── orcjit_dylib.h
 │   └── orcjit_utils.h          # LLVM error handling utilities
 ├── python/tvm_ffi_orcjit/
 │   ├── __init__.py             # Module exports and library loading
-│   ├── session.py              # Python ExecutionSession wrapper
-│   └── dylib.py                # Python DynamicLibrary wrapper
+│   └── session.py              # Python ExecutionSession + default_session
 ├── tests/                      # See tests/README.md
 └── examples/quick-start/       # Complete example with CMake
 ```
@@ -226,7 +225,7 @@ tvm_ffi_orcjit/
 
 Runs on Linux (x86_64, aarch64), macOS (arm64), Windows (AMD64) via
 `cibuildwheel`. Each platform builds test objects with multiple compilers
-and runs the full test suite. See `.github/workflows/tvm_ffi_orcjit.yml`.
+and runs the full test suite. See the `orcjit` job in `.github/workflows/ci_test.yml`.
 
 ## Troubleshooting
 

addons/tvm_ffi_orcjit/examples/quick-start/add.cc

@@ -40,7 +40,3 @@ int fib_impl(int n) {
   return fib_impl(n - 1) + fib_impl(n - 2);
 }
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fibonacci, fib_impl);
-
-// String concatenation example
-std::string concat_impl(std::string a, std::string b) { return a + b; }
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(concat, concat_impl);