[ORCJIT] Add shared session and high-level load_module

[cyx-6]

Adds a process-wide shared execution session and a unified load_module on the tvm_ffi_orcjit addon, mirroring core tvm_ffi.load_module.

Changes

• default_session() — leaked, process-wide shared ExecutionSession; settable ORC-runtime-path registry resolved on first use.
• ExecutionSession.load_module(objects, name) — loads one or more objects (path str/Path or in-memory bytes; single item or list) into one fresh JITDylib, injects context symbols eagerly, expands any embedded library binary into an import tree, and returns a plain tvm_ffi.Module.
• Recursive session lock serializing all compound JITDylib ops (create / add / lookup+init / set-link-order / teardown) for safe shared-session use across threads; the resolved-call hot path stays lock-free.
• Namespace rename — registered globals and the object key move to ffi_orcjit; set_link_order kept but dropped from the advertised surface.
• cpp.build — copy a single object through instead of ld -r, which some system linkers use to stamp a nonzero sh_addr on zero-size sections (e.g. .note.GNU-stack), producing a JITLink block that overlaps in EHFrameEdgeFixer and breaks C++ (eh_frame) loads.

Also folds in earlier addon cleanups (hide static-linked archive symbols; numpydoc docstrings).

Testing

New tests/test_session_load_module.py covers input shapes, eager context, drop, embedded-binary expansion, and concurrent shared-session use. Full addon suite green (154 passed, 3 skipped) on Linux aarch64.

[gemini-code-assist]

Code Review

This pull request introduces a high-level loading API for the TVM-FFI ORCJIT addon, adding a process-wide shared default_session() and a unified load_module() method that accepts both file paths and in-memory bytes. It also implements support for deserializing embedded library binaries to reconstruct import trees, introduces recursive mutex synchronization to serialize compound JITDylib operations, and updates the C++ build system to bypass partial linking for single objects. The review feedback highlights several critical security and robustness issues in the newly added C++ binary blob parser, such as potential integer overflows, excessive memory allocation, out-of-bounds reads in CSR parsing, and null-pointer dereferences. Additionally, the reviewer noted a thread-safety issue in the lazy initialization of the Python default_session().

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The check cursor_ + sizeof(uint64_t) <= size_ is susceptible to integer overflow if size_ is extremely large. To prevent integer overflow, rewrite the check using subtraction: size_ - cursor_ >= sizeof(uint64_t).

  uint64_t ReadU64() {
    TVM_FFI_CHECK(size_ - cursor_ >= sizeof(uint64_t), RuntimeError)
        << "Corrupt library binary: unexpected end of blob";

[gemini-code-assist]

The check cursor_ + nbytes <= size_ is susceptible to integer overflow if nbytes is extremely large (e.g., 0xFFFFFFFFFFFFFFFF). If an overflow occurs, the check can be bypassed, leading to out-of-bounds reads or excessive memory allocation. To prevent integer overflow, rewrite the check using subtraction: nbytes <= size_ - cursor_.

  std::string ReadString() {
    uint64_t nbytes = ReadU64();
    TVM_FFI_CHECK(nbytes <= size_ - cursor_, RuntimeError)
        << "Corrupt library binary: string length exceeds blob";

[gemini-code-assist]

In ReadU64Vector(), out.reserve(static_cast<size_t>(count)) is called with a count read directly from the binary blob. If count is extremely large or corrupt, this will attempt to allocate a massive amount of memory, leading to an out-of-memory crash (DoS).

Since each element in the vector is a uint64_t (8 bytes), the remaining size in the blob must be at least count * sizeof(uint64_t). We should validate this before reserving memory to prevent excessive allocation and out-of-bounds reads.

  std::vector<uint64_t> ReadU64Vector() {
    uint64_t count = ReadU64();
    TVM_FFI_CHECK(count <= (size_ - cursor_) / sizeof(uint64_t), RuntimeError)
        << "Corrupt library binary: vector size exceeds remaining blob size";
    std::vector<uint64_t> out;
    out.reserve(static_cast<size_t>(count));
    for (uint64_t i = 0; i < count; ++i) {
      out.push_back(ReadU64());
    }
    return out;
  }

[gemini-code-assist]

The CSR (Compressed Sparse Row) import tree is parsed without validating that indptr elements are monotonically non-decreasing, that indptr[0] == 0, or that indptr.back() <= child_indices.size(). If any of these properties are violated, the loop for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j) can perform out-of-bounds reads on child_indices.

Add validation checks for the CSR structure to ensure safety.

  std::vector<uint64_t> indptr = reader.ReadU64Vector();
  std::vector<uint64_t> child_indices = reader.ReadU64Vector();
  TVM_FFI_CHECK(!indptr.empty(), RuntimeError) << "Corrupt library binary: empty import tree";
  TVM_FFI_CHECK(indptr[0] == 0, RuntimeError)
      << "Corrupt library binary: indptr[0] must be 0";
  size_t num_modules = indptr.size() - 1;
  for (size_t i = 0; i < num_modules; ++i) {
    TVM_FFI_CHECK(indptr[i] <= indptr[i + 1], RuntimeError)
        << "Corrupt library binary: indptr must be non-decreasing";
  }
  TVM_FFI_CHECK(indptr.back() <= child_indices.size(), RuntimeError)
      << "Corrupt library binary: indptr.back() exceeds child_indices size";

[gemini-code-assist]

If a custom module loader fails or returns an uninitialized/null Module, calling ImportModule on it will cause a segmentation fault. We should defensively check that both the parent and child modules are defined before importing.

  // Wire the import tree (CSR) using the public ModuleObj::ImportModule.
  for (size_t i = 0; i < modules.size(); ++i) {
    for (uint64_t j = indptr[i]; j < indptr[i + 1]; ++j) {
      uint64_t child = child_indices[j];
      TVM_FFI_CHECK(child < modules.size(), RuntimeError)
          << "Corrupt library binary: child index out of range";
      TVM_FFI_CHECK(modules[i].defined(), RuntimeError)
          << "Corrupt library binary: module at index " << i << " is null";
      TVM_FFI_CHECK(modules[static_cast<size_t>(child)].defined(), RuntimeError)
          << "Corrupt library binary: child module at index " << child << " is null";
      modules[i]->ImportModule(modules[static_cast<size_t>(child)]);
    }
  }

[gemini-code-assist]

In default_session(), the lazy initialization of _default_session is not thread-safe. Since _ffi_api.GetDefaultOrcRuntimePath() and _ffi_api.SetDefaultOrcRuntimePath() are FFI calls that may release the Python GIL, multiple threads calling default_session() concurrently during the first invocation could race, leading to multiple initializations of the process-wide shared session.

Use a threading.Lock to ensure thread-safe lazy initialization.

_default_session: ExecutionSession | None = None
_session_lock = __import__("threading").Lock()


def default_session() -> ExecutionSession:
    """Return the process-wide shared execution session.

    A single leaked, never-destroyed session shared by all callers in the
    process, so they share one LLVM ``ExecutionSession`` — hence process
    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 in order: an embedder-registered default
    (via ``ffi_orcjit.SetDefaultOrcRuntimePath``), then the runtime bundled next
    to this extension, then none. For an isolated session or a tuned arena,
    construct an :class:`ExecutionSession` directly instead.

    Returns
    -------
    ExecutionSession
        The shared execution session.

    Examples
    --------
    >>> import tvm_ffi_orcjit as oj
    >>> session = oj.default_session()
    >>> mod = session.load_module("dylib0.o")

    """
    global _default_session
    if _default_session is None:
        with _session_lock:
            if _default_session is None:
                # Honor an embedder-registered default; otherwise fall back to the
                # runtime bundled next to the extension. This must happen before the
                # first DefaultSession() call, which creates the leaked singleton.
                if not _ffi_api.GetDefaultOrcRuntimePath() and (bundled := _find_orc_rt_library()):
                    _ffi_api.SetDefaultOrcRuntimePath(bundled)
                _default_session = _ffi_api.DefaultSession()
    return _default_session

[tqchen]

ffi::Variant<String, Bytes>

[tqchen]

this is not needed, maybe we can simply have python register the tvm_ffi_orcjit.DefaultOrcRuntimePath function

[tqchen]

GlobalDefault

[tqchen]

do we need recursive muted here? seems normal one is fine

[tqchen]

likely no need to expose

[tqchen]

consider write tvm_ffi_orcjit as full path name

[tqchen]

i think we can further simplify if we assume ORCJITDynamicLibrary is unit operation, we explicitly do refresh in LoadModule, then GetFunction do not need to do that, avoid exposing functions like SetLinkOrder AddObjectBytes, leave those as internal

[cyx-6]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the tvm_ffi_orcjit addon to introduce a high-level load_module API and a process-wide shared default_session(), replacing the previous low-level incremental library loading mechanism. It also adds support for deserializing embedded library binaries and implements thread-safe locking for JIT operations. Additionally, single object files now bypass partial linking in extension.py to prevent section address rewriting issues. One critical issue was identified in the deserialization parser where an empty or single-element indptr vector can cause an out-of-bounds access when returning the root module; adding a size check is recommended to prevent undefined behavior.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[cyx-6]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the tvm_ffi_orcjit addon to introduce a high-level, process-wide shared session API (default_session()) and a unified module loading API (load_module(...)), replacing the previous low-level incremental library APIs. It also implements support for deserializing and wiring embedded library binaries (__tvm_ffi__library_bin), improves thread safety with robust locking, and optimizes single-object compilation by bypassing partial linking. Feedback on the changes highlights a security concern in ProcessEmbeddedLibraryBin where the untrusted nbytes header from the embedded binary is read without an upper-bound check, which could lead to integer overflow or out-of-memory crashes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The nbytes header read from the embedded library binary is untrusted and can be arbitrarily large (e.g., due to corruption or malicious input). Relying on it directly to size the BlobReader and subsequently guard vector/string allocations can lead to integer overflow bypasses and process crashes via out-of-memory (OOM) or allocation failures (e.g., in reserve()). Enforcing a sane upper limit on nbytes (such as 2 GB) prevents these issues.

Module ProcessEmbeddedLibraryBin(const char* library_bin, const Module& lib_module) {
  uint64_t nbytes = 0;
  for (size_t i = 0; i < sizeof(nbytes); ++i) {
    nbytes |= static_cast<uint64_t>(static_cast<unsigned char>(library_bin[i])) << (i * 8);
  }
  // Sanity check to prevent integer overflow or absurd memory allocation if the header is corrupt.
  constexpr uint64_t kMaxLibraryBinSize = 2ULL * 1024 * 1024 * 1024; // 2 GB
  TVM_FFI_CHECK(nbytes < kMaxLibraryBinSize, RuntimeError)
      << "Corrupt library binary: size header " << nbytes << " exceeds maximum allowed size (2 GB)";
  BlobReader reader(library_bin + sizeof(nbytes), static_cast<size_t>(nbytes));

[cyx-6]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the tvm_ffi_orcjit addon to introduce a high-level, unified loading API. It replaces the low-level incremental library API (create_library and add) with ExecutionSession.load_module, which accepts file paths, in-memory object bytes, or a sequence of both, returning a plain tvm_ffi.Module. It also introduces default_session(), a process-wide shared execution session, and implements support for deserializing embedded library binaries (__tvm_ffi__library_bin) to reconstruct import trees. Additionally, a copy_object rule is added to the C++ builder to bypass ld -r for single object files, preventing section address rewriting issues. I have no feedback to provide as there are no review comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.