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

[yaoyaoding]

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]

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.

Footnotes

1. 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]

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]

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]

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]

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]

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.

[tqchen]

Refer to existing LibraryModule impl.

[tqchen]

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

[yzh119]

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

[tqchen]

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