[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

Author: tqchen

.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 &&

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
 )
 

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"

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);

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.
    *

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)

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 |

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;
+}

addons/tvm_ffi_orcjit/tests/test_basic.py

@@ -21,6 +21,8 @@
 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)
 # ---------------------------------------------------------------------------