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 329ec838 [ORCJIT] Embed liborc_rt into the extension and add orc_rt
selector (#660)
329ec838 is described below
commit 329ec838400c18d364062f5ec75dbe4263461d37
Author: Yaxing Cai <[email protected]>
AuthorDate: Sat Jul 11 15:33:27 2026 +0800
[ORCJIT] Embed liborc_rt into the extension and add orc_rt selector (#660)
## Motivation
The ORC runtime shipped as a separate `liborc_rt*.a` data file next to
the extension `.so` and was resolved at startup via a filesystem glob
(through the `tvm_ffi_orcjit.DefaultOrcRuntimePath` Python hook). That
couples two artifacts, lets the runtime be lost/relocated by
repackaging, and does a filesystem glob to find its own runtime on first
session.
## What this does
Bakes `liborc_rt.a` into the extension `.so`'s `.rodata` via a generated
`.incbin` stub (`orc_rt_embed.S.in`) and hands the bytes to
`ExecutorNativePlatform` as a **zero-copy** `MemoryBuffer`. The wheel
now ships **one artifact**, startup does **no path lookup**, and the
runtime cannot be separated from the extension. The embedded
`start`/`end` symbols are `.hidden`, so they add no interposition
surface and don't affect the existing `--exclude-libs` symbol hiding.
This is **Solution A** of the embedded-orcrt design.
### `default_session()`
Always uses the embedded runtime and is **deliberately not
user-configurable**: a shared process-wide singleton with a hidden
runtime knob would let whichever caller runs first silently pick the
platform for everyone. The old `DefaultOrcRuntimePath` discovery hook is
dropped.
### User-created `ExecutionSession` — new `orc_rt` selector
Renamed `orc_rt_path` → `orc_rt`, now a 4-state selector:
| `orc_rt=` | meaning |
|---|---|
| `"auto"` (default) | the embedded runtime |
| `str` / `Path` | a custom `liborc_rt` archive on disk |
| `bytes` | a custom `liborc_rt` archive in memory |
| `None` | no ORC platform |
C++ takes `Optional<Variant<String, Bytes>>`; empty `String` ==
`"auto"`, `nullopt` == `None`. Linux/ELF only, matching where the ORC
platform is wired up today; macOS/Windows ignore the selector as before.
A custom runtime must match the LLVM/compiler-rt the extension was built
against.
## Testing
- New `orc_rt` tests: custom path (str + `Path`), in-memory bytes,
`None` (no platform), and bad-type rejection. Path/bytes cases are
`@skipif` non-Linux (the selector only takes effect on ELF).
- Verified locally (aarch64 Linux, LLVM 22): the `.so` carries the
hidden `orc_rt_archive` symbols in `.rodata`, absent from the dynamic
table; no `liborc_rt*.a` bundled; C++ exception unwinding across a JIT
frame works under `auto`/path/bytes (proving `ELFNixPlatform` is
genuinely installed).
- Full addon suite: 164 passed, 11 skipped. clang-format / ruff /
cmake-format / ASF-header all clean.
---
addons/tvm_ffi_orcjit/CMakeLists.txt | 23 +++--
.../python/tvm_ffi_orcjit/session.py | 88 +++++++++-----------
addons/tvm_ffi_orcjit/src/ffi/orc_rt_embed.S.in | 34 ++++++++
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc | 4 +-
addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 97 +++++++++++++++-------
addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h | 32 +++++--
.../tests/test_session_load_module.py | 61 +++++++++++++-
7 files changed, 239 insertions(+), 100 deletions(-)
diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt
b/addons/tvm_ffi_orcjit/CMakeLists.txt
index fb33741f..bdb2952f 100644
--- a/addons/tvm_ffi_orcjit/CMakeLists.txt
+++ b/addons/tvm_ffi_orcjit/CMakeLists.txt
@@ -170,10 +170,12 @@ if (CMAKE_SYSTEM_NAME MATCHES
"Linux|Android|FreeBSD|NetBSD|OpenBSD" AND CMAKE_C
target_link_options(tvm_ffi_orcjit PRIVATE
"-Wl,--exclude-libs=${_exclude_libs_arg}")
endif ()
-# ---- Find and bundle liborc_rt ----
-# Platform notes:
+# ---- Find and embed liborc_rt (Linux/ELF only) ----
+# The archive bytes are baked into the .so's .rodata (see
src/ffi/orc_rt_embed.S.in) and handed to
+# ExecutorNativePlatform as a zero-copy MemoryBuffer at runtime, so the wheel
ships one artifact
+# with no on-disk path lookup. Platform notes:
#
-# * Linux: liborc_rt is used for ExecutorNativePlatform → ELFNixPlatform.
+# * Linux: liborc_rt drives 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.
@@ -193,9 +195,18 @@ if (NOT APPLE AND NOT WIN32)
endif ()
if (ORC_RT_PATH)
- message(STATUS "Found liborc_rt: ${ORC_RT_PATH}")
- file(COPY "${ORC_RT_PATH}" DESTINATION "${CMAKE_BINARY_DIR}")
- install(FILES "${ORC_RT_PATH}" DESTINATION lib)
+ message(STATUS "Embedding liborc_rt: ${ORC_RT_PATH}")
+ enable_language(ASM)
+ configure_file(
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/orc_rt_embed.S.in"
+ "${CMAKE_CURRENT_BINARY_DIR}/orc_rt_embed.S" @ONLY
+ )
+ target_sources(tvm_ffi_orcjit PRIVATE
"${CMAKE_CURRENT_BINARY_DIR}/orc_rt_embed.S")
+ target_compile_definitions(tvm_ffi_orcjit PRIVATE
TVM_FFI_ORCJIT_EMBED_ORC_RT=1)
+ # Re-run .incbin if the archive itself changes (path is stable across
builds).
+ set_property(
+ SOURCE "${CMAKE_CURRENT_BINARY_DIR}/orc_rt_embed.S" PROPERTY
OBJECT_DEPENDS "${ORC_RT_PATH}"
+ )
endif ()
# ---- Install ----
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 89e6b58e..66c7789f 100644
--- a/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py
+++ b/addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py
@@ -18,7 +18,6 @@
from __future__ import annotations
-import sys
import threading
from pathlib import Path
from typing import TYPE_CHECKING
@@ -26,48 +25,12 @@ from typing import TYPE_CHECKING
import tvm_ffi
from tvm_ffi import Module, Object, register_object
-from . import _ffi_api, _lib_dir
+from . import _ffi_api
if TYPE_CHECKING:
from collections.abc import Sequence
-def _find_orc_rt_library() -> str | None:
- """Find the bundled liborc_rt library in the same directory as the
.so/.dll."""
- # Windows: skip ORC runtime entirely. LLVM's COFFPlatform (loaded via
- # ExecutorNativePlatform with liborc_rt) depends on MSVC C++ runtime
symbols
- # that are not available in the JIT environment. On Windows, ORC JIT uses a
- # C-only strategy: JIT objects are compiled as pure C (TVMFFISafeCallType
ABI),
- # 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.
- #
- # 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:
- for lib_path in _lib_dir.glob(pattern):
- return str(lib_path)
- return None
-
-
-@tvm_ffi.register_global_func("tvm_ffi_orcjit.DefaultOrcRuntimePath",
override=True)
-def _default_orc_runtime_path() -> str:
- """Resolve the ORC runtime path for the shared session (C++ hook).
-
- Called once by the C++ ``GlobalDefault`` singleton on first use. Returns
the
- runtime bundled next to this extension, or an empty string to run with no
- ORC platform (macOS/Windows, or when no bundled runtime is found).
- """
- return _find_orc_rt_library() or ""
-
-
@register_object("tvm_ffi_orcjit.ExecutionSession")
class ExecutionSession(Object):
"""ORC JIT Execution Session.
@@ -86,14 +49,26 @@ class ExecutionSession(Object):
"""
- def __init__(self, orc_rt_path: str | None = None, slab_size: int = 0) ->
None:
+ def __init__(
+ self,
+ orc_rt: str | Path | bytes | bytearray | None = "auto",
+ slab_size: int = 0,
+ ) -> None:
"""Initialize ExecutionSession.
Parameters
----------
- orc_rt_path : str, optional
- Optional path to the liborc_rt library. If not provided, it will be
- automatically discovered using clang.
+ orc_rt : str or Path or bytes or None
+ Which ORC runtime to install. Linux/ELF only — ignored on macOS and
+ Windows, which never configure an ORC platform.
+
+ - ``"auto"`` (default): the runtime embedded in this extension.
+ - a path (``str`` or ``Path``): a custom liborc_rt archive on disk.
+ - ``bytes``: a custom liborc_rt archive held in memory.
+ - ``None``: no ORC platform at all.
+
+ A custom runtime (path or bytes) must match the LLVM/compiler-rt
this
+ extension was built against; ``"auto"`` is almost always what you
want.
slab_size : int
Per-slab capacity in bytes for the JIT memory manager. Linux only —
ignored on macOS and Windows, where the slab allocator is compiled
@@ -109,11 +84,23 @@ class ExecutionSession(Object):
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, slab_size) # type: ignore
+ # Normalize to what the C++ ctor (Optional<Variant<String, Bytes>>)
reads:
+ # None -> no platform; "" -> auto/embedded; str -> path; bytes ->
in-memory.
+ rt: str | bytes | None
+ if orc_rt is None:
+ rt = None
+ elif orc_rt == "auto":
+ rt = ""
+ elif isinstance(orc_rt, (bytes, bytearray)):
+ rt = bytes(orc_rt)
+ elif isinstance(orc_rt, (str, Path)):
+ rt = str(orc_rt)
+ else:
+ raise TypeError(
+ "orc_rt must be 'auto', a path (str or Path), liborc_rt bytes,
or None, "
+ f"but got {type(orc_rt).__name__}"
+ )
+ self.__init_handle_by_constructor__(_ffi_api.ExecutionSession, rt,
slab_size) # type: ignore
def load_module(
self,
@@ -223,10 +210,9 @@ def default_session() -> ExecutionSession:
symbols, the slab arena, and cross-library linking. Created on first call
and cached for the lifetime of the process.
- The ORC runtime path is resolved once, on first use, by the C++ singleton
- calling the registered ``tvm_ffi_orcjit.DefaultOrcRuntimePath`` hook (which
- returns the runtime bundled next to this extension, or none). For an
isolated
- session or a tuned arena, construct an :class:`ExecutionSession` directly.
+ The session uses the ORC runtime embedded in the extension (no on-disk path
+ lookup). For an isolated session or a tuned arena, construct an
+ :class:`ExecutionSession` directly.
Returns
-------
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orc_rt_embed.S.in
b/addons/tvm_ffi_orcjit/src/ffi/orc_rt_embed.S.in
new file mode 100644
index 00000000..6937b1fb
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/src/ffi/orc_rt_embed.S.in
@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+/*
+ * Embed liborc_rt.a verbatim into .rodata (Solution A). CMake configure_file
+ * substitutes @ORC_RT_PATH@; the assembler copies the bytes with .incbin, so
+ * no 5 MB literal is compiled. The start/end symbols are .hidden to match the
+ * target's hidden visibility (no dynamic-table entry, no interposition).
+ */
+ .section .rodata
+ .balign 8
+ .hidden orc_rt_archive_start
+ .global orc_rt_archive_start
+orc_rt_archive_start:
+ .incbin "@ORC_RT_PATH@"
+ .hidden orc_rt_archive_end
+ .global orc_rt_archive_end
+orc_rt_archive_end:
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
index 6e98e863..1de989df 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
@@ -363,8 +363,8 @@ static void RegisterOrcJITFunctions() {
refl::GlobalDef()
.def("tvm_ffi_orcjit.ExecutionSession",
- [](const std::string& orc_rt_path, int64_t slab_size_bytes) {
- return ORCJITExecutionSession(orc_rt_path, slab_size_bytes);
+ [](const Optional<Variant<String, Bytes>>& orc_rt, int64_t
slab_size_bytes) {
+ return ORCJITExecutionSession(orc_rt, slab_size_bytes);
})
.def("tvm_ffi_orcjit.GlobalDefaultSession",
[]() { return ORCJITExecutionSessionObj::GlobalDefault(); })
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
index 24e4cc0c..4f534dc0 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc
@@ -27,15 +27,16 @@
#include <llvm/ExecutionEngine/Orc/LLJIT.h>
#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
#include <llvm/Support/Error.h>
+#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/Process.h>
#include <llvm/Support/TargetSelect.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/error.h>
-#include <tvm/ffi/function.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/reflection/registry.h>
#include <cstddef>
+#include <cstdint>
#include <mutex>
#include "orcjit_dylib.h"
@@ -71,7 +72,57 @@ struct LLVMInitializer {
static LLVMInitializer llvm_initializer;
-ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string&
orc_rt_path,
+#ifdef TVM_FFI_ORCJIT_EMBED_ORC_RT
+// liborc_rt.a embedded in .rodata by orc_rt_embed.S.in. C linkage binds to the
+// assembler symbols; the linkage spec must be at namespace scope, not in-fn.
+extern "C" const char orc_rt_archive_start[];
+extern "C" const char orc_rt_archive_end[];
+#endif
+
+namespace {
+#ifdef TVM_FFI_ORCJIT_EMBED_ORC_RT
+// Zero-copy view of the embedded archive; its bytes live for the image
lifetime.
+// Size via uintptr_t: the two symbols are distinct objects, so subtracting the
+// pointers directly would be UB.
+std::unique_ptr<llvm::MemoryBuffer> GetEmbeddedOrcRuntimeBuffer() {
+ return llvm::MemoryBuffer::getMemBuffer(
+ llvm::StringRef(orc_rt_archive_start,
+ reinterpret_cast<std::uintptr_t>(orc_rt_archive_end) -
+
reinterpret_cast<std::uintptr_t>(orc_rt_archive_start)),
+ "liborc_rt.a", /*RequiresNullTerminator=*/false);
+}
+#endif
+
+// Install ExecutorNativePlatform per the `orc_rt` selector (see the ctor doc).
+// A no-op except on Linux/ELF: macOS skips the platform (compact-unwind bug)
+// and Windows never wires up COFFPlatform, so both ignore the selector.
+void SetUpOrcPlatform(llvm::orc::LLJITBuilder& builder,
+ const Optional<Variant<String, Bytes>>& orc_rt) {
+#if defined(__APPLE__) || defined(_WIN32)
+ (void)builder;
+ (void)orc_rt;
+#else
+ if (!orc_rt.has_value()) return; // None -> no platform
+ const Variant<String, Bytes>& sel = orc_rt.value();
+ if (auto opt_path = sel.as<String>()) {
+ const String& path = *opt_path;
+ if (path.empty()) { // "auto" -> embedded (or nothing if compiled out)
+#ifdef TVM_FFI_ORCJIT_EMBED_ORC_RT
+
builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(GetEmbeddedOrcRuntimeBuffer()));
+#endif
+ } else {
+ builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(path.operator
std::string()));
+ }
+ } else { // Bytes: ExecutorNativePlatform takes ownership of the copy.
+ const Bytes& bytes = sel.get<Bytes>();
+
builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(llvm::MemoryBuffer::getMemBufferCopy(
+ llvm::StringRef(bytes.data(), bytes.size()), "liborc_rt.a")));
+ }
+#endif
+}
+} // namespace
+
+ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const
Optional<Variant<String, Bytes>>& orc_rt,
int64_t slab_size_bytes)
: jit_(nullptr) {
// Create slab-backed memory manager — pre-reserves a contiguous VA region
@@ -156,21 +207,15 @@
ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_p
};
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()) {
- builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform(orc_rt_path));
- }
-#else
- (void)orc_rt_path;
-#endif
+ // Configure the ORC platform from `orc_rt` (a no-op off Linux/ELF; see
+ // SetUpOrcPlatform). macOS in particular must 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).
+ // InitFiniPlugin below handles __mod_init_func / __mod_term_func instead;
+ // tradeoff: no C++ exception unwinding across JIT frames on macOS.
+ SetUpOrcPlatform(builder, orc_rt);
setup_builder(builder);
jit_ = TVM_FFI_ORCJIT_LLVM_CALL(builder.create());
#ifdef _WIN32
@@ -198,26 +243,20 @@
ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const std::string& orc_rt_p
#endif
}
-ORCJITExecutionSession::ORCJITExecutionSession(const std::string& orc_rt_path,
+ORCJITExecutionSession::ORCJITExecutionSession(const Optional<Variant<String,
Bytes>>& orc_rt,
int64_t slab_size_bytes) {
ObjectPtr<ORCJITExecutionSessionObj> obj =
- make_object<ORCJITExecutionSessionObj>(orc_rt_path, slab_size_bytes);
+ make_object<ORCJITExecutionSessionObj>(orc_rt, slab_size_bytes);
data_ = std::move(obj);
}
ORCJITExecutionSession ORCJITExecutionSessionObj::GlobalDefault() {
// Leaked, never-destroyed: the owned LLJIT and slab arena must not be torn
// down during interpreter finalization, where teardown could call back into
- // the host language. Resolve the ORC runtime path once, from the registered
- // global tvm_ffi_orcjit.DefaultOrcRuntimePath (a Python hook that returns
the
- // bundled runtime); absent or empty means no platform.
- static ORCJITExecutionSession* inst = [] {
- std::string orc_rt_path;
- if (auto fpath =
Function::GetGlobal("tvm_ffi_orcjit.DefaultOrcRuntimePath")) {
- orc_rt_path = (*fpath)().cast<String>().operator std::string();
- }
- return new ORCJITExecutionSession(orc_rt_path, 0);
- }();
+ // the host language. Always "auto" (empty String) and not user-configurable
+ // (see the header) — the embedded runtime's .rodata outlives it trivially.
+ static ORCJITExecutionSession* inst =
+ new ORCJITExecutionSession(Variant<String, Bytes>(String("")), 0);
return *inst;
}
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
index 7e9b8998..6988490f 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.h
@@ -31,6 +31,7 @@
#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/object.h>
+#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <atomic>
@@ -58,10 +59,21 @@ class ORCJITExecutionSession;
class ORCJITExecutionSessionObj : public Object {
public:
/*!
- * \brief Default constructor (for make_object)
+ * \brief Construct a session, selecting the ORC runtime.
+ *
+ * \param orc_rt Runtime selector (Linux/ELF only; ignored on macOS/Windows,
+ * which never configure an ORC platform):
+ * - empty \c String (default, "auto"): the runtime embedded in this
+ * extension (or no platform if the embed was compiled out);
+ * - non-empty \c String: a custom liborc_rt archive on disk;
+ * - \c Bytes: a custom liborc_rt archive held in memory;
+ * - \c nullopt: no ORC platform at all.
+ * A custom runtime must match the LLVM this extension was built
against.
+ * \param slab_size_bytes Per-slab capacity for the JIT memory arena.
*/
- explicit ORCJITExecutionSessionObj(const std::string& orc_rt_path = "",
- int64_t slab_size_bytes = 0);
+ explicit ORCJITExecutionSessionObj(
+ const Optional<Variant<String, Bytes>>& orc_rt = Variant<String,
Bytes>(String("")),
+ int64_t slab_size_bytes = 0);
/*!
* \brief Get the process-wide shared execution session.
@@ -71,9 +83,10 @@ class ORCJITExecutionSessionObj : public Object {
* and cross-library linking. Never torn down (interpreter finalization could
* otherwise call back into the host language during teardown).
*
- * The ORC runtime path is resolved on first call from the registered global
- * \c tvm_ffi_orcjit.DefaultOrcRuntimePath (a Python-side hook that returns
the
- * bundled runtime, or empty to disable); absent or empty means no platform.
+ * Always uses the ORC runtime embedded in this extension; deliberately not
+ * user-configurable (a shared process-wide singleton with a hidden runtime
+ * knob would let the first caller pick the platform for everyone). Construct
+ * \ref ORCJITExecutionSession directly for a custom runtime.
*
* \return The shared execution session.
*/
@@ -221,9 +234,12 @@ class ORCJITExecutionSession : public ObjectRef {
public:
/*!
* \brief Create a new ExecutionSession
- * \return The created execution session instance
+ * \param orc_rt ORC runtime selector; see \ref ORCJITExecutionSessionObj.
+ * \param slab_size_bytes Per-slab capacity for the JIT memory arena.
*/
- explicit ORCJITExecutionSession(const std::string& orc_rt_path = "", int64_t
slab_size_bytes = 0);
+ explicit ORCJITExecutionSession(
+ const Optional<Variant<String, Bytes>>& orc_rt = Variant<String,
Bytes>(String("")),
+ 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/tests/test_session_load_module.py
b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py
index 461e3340..1b11b753 100644
--- a/addons/tvm_ffi_orcjit/tests/test_session_load_module.py
+++ b/addons/tvm_ffi_orcjit/tests/test_session_load_module.py
@@ -31,6 +31,7 @@ limitations unrelated to this API.
from __future__ import annotations
import gc
+import sys
import threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
@@ -51,6 +52,24 @@ def obj(name: str) -> str:
return str(path)
+def _find_orc_rt_archive() -> str:
+ """Locate a liborc_rt archive under LLVM_PREFIX, or skip.
+
+ Used only by the custom-``orc_rt`` override tests. There is no bundled copy
+ next to the extension anymore (the runtime is embedded), so these tests
fall
+ back to the archive shipped with the build's LLVM.
+ """
+ import os # noqa: PLC0415
+
+ prefix = os.environ.get("LLVM_PREFIX")
+ if not prefix:
+ pytest.skip("LLVM_PREFIX not set; cannot locate a liborc_rt archive")
+ matches = sorted(Path(prefix).glob("lib/clang/*/lib/*/liborc_rt*.a"))
+ if not matches:
+ pytest.skip("no liborc_rt archive found under LLVM_PREFIX")
+ return str(matches[0])
+
+
# ---------------------------------------------------------------------------
# default_session()
# ---------------------------------------------------------------------------
@@ -92,10 +111,44 @@ def test_default_session_concurrent_first_call() -> None:
assert all(s is sessions[0] for s in sessions)
-def test_default_orc_runtime_path_hook_registered() -> None:
- """The C++ singleton's ORC-path hook is registered and returns a string."""
- path = tvm_ffi.get_global_func("tvm_ffi_orcjit.DefaultOrcRuntimePath")()
- assert isinstance(path, str) # bundled runtime path, or "" when
none/disabled
+# ---------------------------------------------------------------------------
+# orc_rt override — custom runtime selectors on a user-created session.
+# (The "auto" default is covered via ExecutionSession() elsewhere.)
+# ---------------------------------------------------------------------------
+
+# A custom ORC runtime is only installed on Linux/ELF; elsewhere it is ignored.
+elf_only = pytest.mark.skipif(
+ sys.platform != "linux", reason="ORC platform (custom liborc_rt) is
Linux/ELF only"
+)
+
+
+@elf_only
+def test_orc_rt_explicit_path() -> None:
+ """A custom on-disk liborc_rt archive is accepted (str and Path)."""
+ archive = _find_orc_rt_archive()
+ for spec in (archive, Path(archive)):
+ mod = ExecutionSession(orc_rt=spec).load_module(obj("c/test_funcs"))
+ assert mod.test_add(4, 5) == 9
+
+
+@elf_only
+def test_orc_rt_in_memory_bytes() -> None:
+ """A liborc_rt archive supplied as in-memory bytes is accepted."""
+ data = Path(_find_orc_rt_archive()).read_bytes()
+ mod = ExecutionSession(orc_rt=data).load_module(obj("c/test_funcs"))
+ assert mod.test_add(6, 7) == 13
+
+
+def test_orc_rt_none_no_platform() -> None:
+ """orc_rt=None runs with no ORC platform; C-ABI objects still load."""
+ mod = ExecutionSession(orc_rt=None).load_module(obj("c/test_funcs"))
+ assert mod.test_add(8, 9) == 17
+
+
+def test_orc_rt_rejects_bad_type() -> None:
+ """A non-str, non-bytes, non-None selector raises a clear TypeError."""
+ with pytest.raises(TypeError, match=r"orc_rt must be"):
+ ExecutionSession(orc_rt=12345)
# ---------------------------------------------------------------------------
