[OrcJIT] Add LLVM ORC JIT v2 dynamic object loader addon #254

yaoyaoding · 2025-11-10T23:35:39Z

This PR introduces a new addon package tvm-ffi-orcjit that enables dynamic loading of TVM-FFI exported object files (.o) at runtime using LLVM's ORC JIT v2 engine.

The addon provides a Python API for loading compiled object files, load the tvm-ffi functions defined in the object files.

The API is organized around three main concepts:

ExecutionSession: A JIT compilation context that manages the lifetime of dynamic libraries.
DynamicLibrary: Represents a shared library that links multiple object files.
Object Files: Compiled .o files containing TVM-FFI exported functions.

Usage Example

from tvm_ffi_orcjit import create_session

# Create an execution session (JIT compilation context)
session = create_session()

# Load an object file into the session, returns a DynamicLibrary handle
dylib = session.load("example.o")

# Get and call a function from the loaded object file
add_func = dylib.get_function("simple_add")
result = add_func(1, 2)
print(f"Result: {result}")  # Output: Result: 3

For incremental loading, you can add multiple object files to the same session:

session = create_session()

# Load first object file
dylib = session.load("math_ops.o")
add = dylib.get_function("simple_add")

# Incrementally load more object files into the same library
dylib.load("string_ops.o")

# Now both functions are accessible from the same library
concat = dylib.get_function("string_concat")

See the test for more example.

TODO

Support Windows platform (currently Linux and macOS only)
Customize memory allocation for dynamic libraries (allow user-provided allocators)

gemini-code-assist · 2025-11-10T23:36:02Z

Summary of Changes

Hello @yaoyaoding, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant new capability to TVM-FFI by integrating LLVM's ORC JIT v2. This integration enables the dynamic loading and execution of C++ object files directly from Python, offering a flexible and efficient way to extend TVM-FFI's functionality at runtime. The new addon provides a clear Python API for managing JIT compilation contexts and dynamic libraries, making it easier for developers to incorporate custom C++ logic without requiring a full recompile of their applications.

Highlights

New Addon Package: Introduces tvm-ffi-orcjit, a new addon package designed to enable dynamic loading of TVM-FFI exported object files (.o) at runtime.
LLVM ORC JIT v2 Integration: Leverages LLVM's ORC JIT v2 engine for Just-In-Time compilation and execution of C++ code, providing a powerful and flexible runtime environment.
Python API: Provides a Pythonic API centered around ExecutionSession (JIT compilation context) and DynamicLibrary (manages object files and symbols) for seamless interaction.
Incremental Loading: Supports incremental loading, allowing multiple object files to be added to the same DynamicLibrary instance, with functions from all loaded files remaining accessible.
Comprehensive Examples and Tests: Includes a quick-start example demonstrating basic usage and a robust test suite covering core functionalities, including symbol resolution and conflict handling.

Ignored Files

Ignored by pattern: .github/workflows/** (3)
- .github/workflows/orcjit-publish.yml
- .github/workflows/orcjit-tests.yml
- .github/workflows/tvm-ffi-orcjit/ci_test.yml

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a new addon package, tvm-ffi-orcjit, to enable dynamic loading of object files using LLVM's ORC JIT v2. The overall design is robust, with a clear separation between the C++ backend and the Python frontend. The implementation demonstrates a deep understanding of LLVM's JIT capabilities, including a clever workaround for the __dso_handle issue. The tests are comprehensive and cover important scenarios like symbol conflicts.

I have identified a few areas for improvement, primarily concerning inconsistencies in documentation, build scripts, and a potential bug in the CMake configuration. There is also a mismatch between a Python API and its C++ backend. Addressing these points will enhance the quality and usability of this new addon.

gemini-code-assist · 2025-11-10T23:38:24Z

addons/tvm-ffi-orcjit/CMakeLists.txt

+target_link_libraries(
+  tvm_ffi_orcjit
+  PUBLIC tvm_ffi
+  PRIVATE LLVM


The find_package(LLVM ... CONFIG) command does not create a single LLVM target. Instead, llvm_map_components_to_libnames populates the LLVM_LIBS variable with a list of component library targets. You should link against ${LLVM_LIBS}.

PRIVATE ${LLVM_LIBS}

gemini-code-assist · 2025-11-10T23:38:24Z

addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py

+    def link_against(self, *libraries: DynamicLibrary) -> None:
+        """Link this library against other dynamic libraries.
+
+        Sets the search order for symbol resolution. Symbols not found in this library
+        will be searched in the linked libraries in the order specified.
+
+        Parameters
+        ----------
+        *libraries : DynamicLibrary
+            One or more dynamic libraries to link against.
+
+        Examples
+        --------
+        >>> session = create_session()
+        >>> lib_utils = session.create_library()
+        >>> lib_utils.add("utils.o")
+        >>> lib_main = session.create_library()
+        >>> lib_main.add("main.o")
+        >>> lib_main.link_against(lib_utils)  # main can call utils symbols
+
+        """
+        handles = [lib._handle for lib in libraries]
+        self._link_func(self._handle, *handles)


The link_against method accepts *libraries, which implies it can link against multiple libraries. However, the C++ backend function orcjit.DynamicLibraryLinkAgainst only supports linking against a single library at a time. This will cause a TypeError if more than one library is passed.

To fix this, you should either update the Python method to accept only one library or modify the C++ backend to handle multiple libraries. Given the current C++ implementation, changing the Python API is the most direct solution.

def link_against(self, library: DynamicLibrary) -> None: """Link this library against another dynamic library. Sets the search order for symbol resolution. Symbols not found in this library will be searched in the linked library. Parameters ---------- library : DynamicLibrary The dynamic library to link against. Examples -------- >>> session = create_session() >>> lib_utils = session.create_library() >>> lib_utils.add("utils.o") >>> lib_main = session.create_library() >>> lib_main.add("main.o") >>> lib_main.link_against(lib_utils) # main can call utils symbols """ self._link_func(self._handle, library._handle)

gemini-code-assist · 2025-11-10T23:38:24Z

addons/tvm-ffi-orcjit/CMakeLists.txt

+separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS})
+add_definitions(${LLVM_DEFINITIONS_LIST})


These add_definitions are redundant as target_compile_definitions is used for the tvm_ffi_orcjit target on line 93. It's better to use target_compile_definitions for target-specific definitions. I suggest removing these lines to avoid redundancy.

gemini-code-assist · 2025-11-10T23:38:24Z

addons/tvm-ffi-orcjit/README.md

+from tvm_ffi_orcjit import ObjectLoader
+
+# Create a loader instance
+loader = ObjectLoader()
+
+# Load an object file
+loader.load("example.o")
+
+# Get and call a function
+add_func = loader.get_function("simple_add")
+result = add_func(1, 2)
+print(f"Result: {result}")  # Output: Result: 3
+```
+
+### Incremental Loading
+
+Load multiple object files and access functions from all of them:
+
+```python
+from tvm_ffi_orcjit import ObjectLoader
+
+loader = ObjectLoader()
+
+# Load first object file
+loader.load("math_ops.o")
+add = loader.get_function("simple_add")
+
+# Load second object file - functions from first remain accessible
+loader.load("string_ops.o")
+concat = loader.get_function("string_concat")
+
+# Both functions work
+print(add(10, 20))           # From math_ops.o
+print(concat("Hello", "World"))  # From string_ops.o
+```


The usage examples in this README appear to be outdated. They refer to an ObjectLoader class which is not present in the current implementation. The correct API, as demonstrated in examples/quick-start/run.py, uses create_session, session.create_library, and lib.add(). Please update the examples to reflect the current API to avoid confusion for new users.

gemini-code-assist · 2025-11-10T23:38:25Z

addons/tvm-ffi-orcjit/README.md

+### Direct Module Access
+
+You can also use TVM-FFI's `load_module` directly (`.o` files are automatically handled):
+
+```python
+import tvm_ffi
+
+# Load object file as a module
+module = tvm_ffi.load_module("example.o")
+
+# Get function
+func = module.get_function("my_function")
+result = func(arg1, arg2)
+```


This section on "Direct Module Access" suggests that tvm_ffi.load_module("example.o") is supported. However, the implementation does not seem to register a module loader for .o files. If this feature is not yet implemented, it would be best to remove this section or clearly mark it as a future capability to prevent user confusion.

addons/tvm-ffi-orcjit/README.md

addons/tvm-ffi-orcjit/examples/quick-start/compile.sh

addons/tvm-ffi-orcjit/pyproject.toml

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

tqchen · 2025-11-11T00:09:15Z

addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.cc

+        int (*)(void* handle, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* rv);
+    auto c_func = reinterpret_cast<TVMFFISafeCallType>(symbol);
+
+    return Function::FromPacked([c_func, name](PackedArgs args, Any* rv) {


Refer to existing LibraryModule impl.

https://github.com/apache/tvm-ffi/blob/main/src/ffi/extra/library_module.cc#L53

tqchen · 2025-11-11T00:10:13Z

addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.cc

+    // Try to get the symbol - return NullOpt if not found
+    void* symbol = nullptr;
+    try {
+      symbol = dylib_->GetSymbol(symbol_name);


return nullptr instead when symbol is not found

tqchen · 2025-11-11T00:11:01Z

addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py

+    from .session import ExecutionSession
+
+
+class DynamicLibrary:


Make it wrap ffi.Module, so we can directly do lib.foo

tqchen · 2025-11-11T00:11:22Z

addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py

+        self._link_func = get_global_func("orcjit.DynamicLibraryLinkAgainst")
+        self._to_module_func = get_global_func("orcjit.DynamicLibraryToModule")
+
+    def add(self, object_file: str | Path) -> None:


lib.add_object_file(path) => lib.mod["__add_object_file__"](path)

lib.foo

lib.function

tqchen · 2025-11-11T00:12:41Z

addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py

+            object_file = str(object_file)
+        self._add_func(self._handle, object_file)
+
+    def link_against(self, *libraries: DynamicLibrary) -> None:


lib.set_link_order(list[libraries])

tqchen · 2025-11-11T00:24:08Z

addons/tvm-ffi-orcjit/MANIFEST.in

+include LICENSE
+include pyproject.toml
+include CMakeLists.txt
+recursive-include include *.h


manifest is not needed and only need in pyproject

yzh119 · 2025-11-11T04:10:06Z

One general question, how do we plan to manage the versions of these add ons?

tqchen · 2025-11-11T12:12:34Z

I think most likely they can evolve independently for now and optional

yaoyaoding added 8 commits November 10, 2025 15:43

add orcjit

d87c16a

add more tests

501b635

fix tests

6ce56a1

wip

d27a5b0

fix

cf2da5d

fix cmakelists.txt

76d13d5

fix

b57756d

fix tests on macos

331404b

gemini-code-assist bot reviewed Nov 10, 2025

View reviewed changes

yaoyaoding and others added 3 commits November 10, 2025 18:40

Update addons/tvm-ffi-orcjit/pyproject.toml

c15f037

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update addons/tvm-ffi-orcjit/examples/quick-start/compile.sh

1bc971c

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update addons/tvm-ffi-orcjit/README.md

c9dc320

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

tqchen reviewed Nov 11, 2025

View reviewed changes

yaoyaoding added 4 commits November 10, 2025 19:35

wip

b855a9e

use the right way to define object

0dfc2d1

move include to src

5af6c77

update cmakelists.txt

b9eb2aa

		separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS})
		add_definitions(${LLVM_DEFINITIONS_LIST})

[OrcJIT] Add LLVM ORC JIT v2 dynamic object loader addon #254

Are you sure you want to change the base?

[OrcJIT] Add LLVM ORC JIT v2 dynamic object loader addon #254

Uh oh!

Conversation

yaoyaoding commented Nov 10, 2025

Usage Example

TODO

Uh oh!

gemini-code-assist bot commented Nov 10, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

tqchen Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 commented Nov 11, 2025

Uh oh!

tqchen commented Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tqchen commented Nov 11, 2025 •

edited

Loading