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 f64ddbe6 [ORCJIT] Reinject contexts on first function lookup (#655)
f64ddbe6 is described below
commit f64ddbe6db5002aa737336696f67a91decdc7ed2
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Jul 6 18:16:50 2026 +0800
[ORCJIT] Reinject contexts on first function lookup (#655)
## Summary
- defer context injection from empty-library construction to the first
ordinary function lookup
- cache each successful refresh; after LLVM accepts an object,
AddObjectFile publishes the refresh flag as false
- keep AddObjectFile outside the thread-safety contract: it must not
race GetFunction, and a failed add naturally leaves the prior flag
unchanged
- keep concurrent GetFunction calls thread-safe with an acquire fast
path and a mutex-protected, double-checked refresh on the false path,
followed by one ordinary symbol lookup
- preserve ordinary GetSymbol initializer execution and LLVM error
propagation
- keep the ORCJIT wheel test environment aligned with the
repository-wide parallel pytest configuration by installing pytest-xdist
- bump the addon package to 0.1.1
If callers violate the AddObjectFile/GetFunction contract, LLVM may
expose a newly added symbol and GetFunction may return its stable
allocated address before that object's context slots are refreshed. The
symbol remains visible, but context initialization is not guaranteed for
that unsupported overlap.
## Regression
One C object defines only __tvm_ffi__library_ctx and a probe.
Barrier-synchronized first lookups verify that every returned function
observes the injected pointer without adding another fixture.
## Validation
- GCC- and Clang-built addon suites: 139 passed, 3 platform skips each
- fresh CI-equivalent wheel environment with local root package, default
parallel pytest, and C/C++ quick starts
- 200 rounds with eight concurrent first GetFunction callers after
object addition
- focused context, invalid-object, constructor, and destructor tests
- fresh 0.1.1 wheel build, isolated install, focused regression, and
quick starts
- applicable formatting and lint hooks
---
.github/actions/build-orcjit-wheel/action.yml | 2 +-
addons/tvm_ffi_orcjit/CMakeLists.txt | 2 +-
addons/tvm_ffi_orcjit/pyproject.toml | 2 +-
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc | 40 ++++++++++++++++------
addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h | 12 +++++++
addons/tvm_ffi_orcjit/tests/CMakeLists.txt | 1 +
addons/tvm_ffi_orcjit/tests/README.md | 1 +
.../tvm_ffi_orcjit/tests/sources/c/test_context.c | 33 ++++++++++++++++++
addons/tvm_ffi_orcjit/tests/test_basic.py | 24 +++++++++++++
9 files changed, 103 insertions(+), 14 deletions(-)
diff --git a/.github/actions/build-orcjit-wheel/action.yml
b/.github/actions/build-orcjit-wheel/action.yml
index 4d0fda58..7cefbed5 100644
--- a/.github/actions/build-orcjit-wheel/action.yml
+++ b/.github/actions/build-orcjit-wheel/action.yml
@@ -102,7 +102,7 @@ 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"
- CIBW_TEST_REQUIRES: pytest ninja
+ CIBW_TEST_REQUIRES: pytest pytest-xdist ninja
CIBW_TEST_COMMAND: >-
pip install --force-reinstall {project} &&
pytest {project}/addons/tvm_ffi_orcjit/tests -v &&
diff --git a/addons/tvm_ffi_orcjit/CMakeLists.txt
b/addons/tvm_ffi_orcjit/CMakeLists.txt
index b0ef5cba..d25b06f7 100644
--- a/addons/tvm_ffi_orcjit/CMakeLists.txt
+++ b/addons/tvm_ffi_orcjit/CMakeLists.txt
@@ -1,7 +1,7 @@
cmake_minimum_required(VERSION 3.20)
project(
tvm_ffi_orcjit
- VERSION 0.1.0
+ VERSION 0.1.1
LANGUAGES C CXX
)
diff --git a/addons/tvm_ffi_orcjit/pyproject.toml
b/addons/tvm_ffi_orcjit/pyproject.toml
index 49bab100..c2221890 100644
--- a/addons/tvm_ffi_orcjit/pyproject.toml
+++ b/addons/tvm_ffi_orcjit/pyproject.toml
@@ -21,7 +21,7 @@ build-backend = "scikit_build_core.build"
[project]
name = "apache-tvm_ffi_orcjit"
-version = "0.1.0"
+version = "0.1.1"
description = "Load TVM-FFI exported object files using LLVM ORC JIT v2"
readme = "README.md"
requires-python = ">=3.10"
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
index 8db8c781..f6bd329a 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.cc
@@ -68,14 +68,6 @@
ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ORCJITExecutionSession session,
: session_(std::move(session)), dylib_(dylib), jit_(jit),
name_(std::move(name)) {
TVM_FFI_CHECK(dylib_ != nullptr, ValueError) << "JITDylib cannot be null";
TVM_FFI_CHECK(jit_ != nullptr, ValueError) << "LLJIT cannot be null";
- if (void** ctx_addr =
reinterpret_cast<void**>(GetSymbol(ffi::symbol::tvm_ffi_library_ctx))) {
- *ctx_addr = this;
- }
- Module::VisitContextSymbols([this](const ffi::String& name, void* symbol) {
- if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name))) {
- *ctx_addr = symbol;
- }
- });
}
ORCJITDynamicLibraryObj::~ORCJITDynamicLibraryObj() {
@@ -111,8 +103,9 @@ void ORCJITDynamicLibraryObj::AddObjectFile(const String&
path) {
TVM_FFI_THROW(IOError) << "Failed to read object file: " << path;
}
- // Add object file to this JITDylib
+ // AddObjectFile is not thread-safe and must not race GetFunction.
TVM_FFI_ORCJIT_LLVM_CALL(jit_->addObjectFile(*dylib_,
std::move(*buffer_or_err)));
+ context_symbol_refreshed_.store(false, std::memory_order_release);
}
void ORCJITDynamicLibraryObj::SetLinkOrder(const
std::vector<llvm::orc::JITDylib*>& dylibs) {
@@ -157,8 +150,29 @@ void* ORCJITDynamicLibraryObj::GetSymbol(const String&
name) {
CxaAtexitRecordsScope scope(&cxa_atexit_records_);
#endif
session_->RunPendingInitializers(GetJITDylib());
- // Convert ExecutorAddr to pointer
- return symbol_or_err ? symbol_or_err->getAddress().toPtr<void*>() : nullptr;
+
+ if (!symbol_or_err) {
+ llvm::Error remaining =
+ llvm::handleErrors(symbol_or_err.takeError(), [](const
llvm::orc::SymbolsNotFound&) {});
+ if (remaining) TVM_FFI_ORCJIT_LLVM_CALL(std::move(remaining));
+ return nullptr;
+ }
+ return symbol_or_err->getAddress().toPtr<void*>();
+}
+
+void ORCJITDynamicLibraryObj::InitContextSymbols() {
+ std::lock_guard<std::mutex> lock(context_symbol_refresh_mutex_);
+ if (context_symbol_refreshed_.load(std::memory_order_acquire)) return;
+
+ if (void** ctx_addr =
reinterpret_cast<void**>(GetSymbol(symbol::tvm_ffi_library_ctx))) {
+ *ctx_addr = this;
+ }
+ Module::VisitContextSymbols([this](const String& name, void* symbol) {
+ if (void** ctx_addr = reinterpret_cast<void**>(GetSymbol(name))) {
+ *ctx_addr = symbol;
+ }
+ });
+ context_symbol_refreshed_.store(true, std::memory_order_release);
}
llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() {
@@ -184,6 +198,10 @@ Optional<Function>
ORCJITDynamicLibraryObj::GetFunction(const String& name) {
// TVM-FFI exports have __tvm_ffi_ prefix
std::string symbol_name = symbol::tvm_ffi_symbol_prefix + std::string(name);
+ if (!context_symbol_refreshed_.load(std::memory_order_acquire)) {
+ InitContextSymbols();
+ }
+
// Try to get the symbol - return NullOpt if not found
if (void* symbol = GetSymbol(symbol_name)) {
TVMFFISafeCallType c_func = reinterpret_cast<TVMFFISafeCallType>(symbol);
diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
index 0ab2bdc0..7221a938 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_dylib.h
@@ -30,6 +30,9 @@
#include <tvm/ffi/object.h>
#include <tvm/ffi/string.h>
+#include <atomic>
+#include <mutex>
+
#include "llvm_patches/macho_cxa_atexit_shim.h"
#include "orcjit_session.h"
@@ -73,6 +76,9 @@ class ORCJITDynamicLibraryObj : public ModuleObj {
*/
void SetLinkOrder(const std::vector<llvm::orc::JITDylib*>& dylibs);
+ /*! \brief Refresh context slots after an object add when needed. */
+ void InitContextSymbols();
+
/*!
* \brief Look up a symbol in this library
* \param name The symbol name to look up
@@ -107,6 +113,12 @@ class ORCJITDynamicLibraryObj : public ModuleObj {
/*! \brief Link order tracking (to support incremental linking) */
llvm::orc::JITDylibSearchOrder link_order_;
+ /*! \brief Whether context slots have been refreshed since the last object
add. */
+ std::atomic<bool> context_symbol_refreshed_{false};
+
+ /*! \brief Serializes the false-path context refresh between function
lookups. */
+ std::mutex context_symbol_refresh_mutex_;
+
#ifdef __APPLE__
/*! \brief Per-dylib __cxa_atexit registry.
*
diff --git a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt
b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt
index 9259e1e7..950a0730 100644
--- a/addons/tvm_ffi_orcjit/tests/CMakeLists.txt
+++ b/addons/tvm_ffi_orcjit/tests/CMakeLists.txt
@@ -92,6 +92,7 @@ add_test_object(sources/c/test_funcs.c)
add_test_object(sources/c/test_funcs2.c)
add_test_object(sources/c/test_funcs_conflict.c)
add_test_object(sources/c/test_call_global.c)
+add_test_object(sources/c/test_context.c)
add_test_object(sources/c/test_types.c)
add_test_object(sources/c/test_link_order_base.c)
add_test_object(sources/c/test_link_order_caller.c)
diff --git a/addons/tvm_ffi_orcjit/tests/README.md
b/addons/tvm_ffi_orcjit/tests/README.md
index 0ab9a2d2..d2207924 100644
--- a/addons/tvm_ffi_orcjit/tests/README.md
+++ b/addons/tvm_ffi_orcjit/tests/README.md
@@ -87,6 +87,7 @@ variant subdirectories (c/, cc/, c-gcc/, etc.).
| `test_funcs2` | More arithmetic (subtract, divide) |
| `test_funcs_conflict` | Symbol conflict testing (duplicate `add`) |
| `test_call_global` | Callbacks into Python-registered global functions |
+| `test_context` | First-lookup library-context injection |
| `test_types` | Zero-arg, multi-arg, float, void return types |
| `test_link_order_base` / `test_link_order_caller` | Cross-library symbol
resolution |
| `test_error` | Error propagation from JIT'd code |
diff --git a/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c
b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c
new file mode 100644
index 00000000..d4a40760
--- /dev/null
+++ b/addons/tvm_ffi_orcjit/tests/sources/c/test_context.c
@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+
+#include <tvm/ffi/c_api.h>
+
+TVM_FFI_DLL_EXPORT void* __tvm_ffi__library_ctx = NULL;
+
+TVM_FFI_DLL_EXPORT int __tvm_ffi_context_is_set(void* self, const TVMFFIAny*
args, int32_t num_args,
+ TVMFFIAny* result) {
+ (void)self;
+ (void)args;
+ (void)num_args;
+ result->type_index = kTVMFFIInt;
+ result->zero_padding = 0;
+ result->v_int64 = __tvm_ffi__library_ctx != NULL;
+ return 0;
+}
diff --git a/addons/tvm_ffi_orcjit/tests/test_basic.py
b/addons/tvm_ffi_orcjit/tests/test_basic.py
index c125e6b6..17939497 100644
--- a/addons/tvm_ffi_orcjit/tests/test_basic.py
+++ b/addons/tvm_ffi_orcjit/tests/test_basic.py
@@ -21,6 +21,8 @@ from __future__ import annotations
import gc
import sys
import tempfile
+import threading
+from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import pytest
@@ -295,6 +297,28 @@ def test_call_global(v: Variant) -> None:
assert mul_func(11, 11) == 121
+# ---------------------------------------------------------------------------
+# Library-context injection
+# ---------------------------------------------------------------------------
+
+
+def test_context_injected_after_object_add() -> None:
+ """Concurrent first lookups after an object add observe the injected
context pointer."""
+ session = ExecutionSession()
+ lib = session.create_library("context_after_add")
+ lib.add(obj("c/test_context"))
+
+ num_workers = 8
+ barrier = threading.Barrier(num_workers)
+
+ def lookup(_worker: int) -> int:
+ barrier.wait()
+ return lib.get_function("context_is_set")()
+
+ with ThreadPoolExecutor(max_workers=num_workers) as pool:
+ assert list(pool.map(lookup, range(num_workers))) == [1] * num_workers
+
+
# ---------------------------------------------------------------------------
# Error handling — pure Python (Group 1)
# ---------------------------------------------------------------------------
