diff --git a/.github/workflows/orcjit-publish.yml b/.github/workflows/orcjit-publish.yml new file mode 100644 index 00000000..6ed0d4d6 --- /dev/null +++ b/.github/workflows/orcjit-publish.yml @@ -0,0 +1,164 @@ +# 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. + +name: Publish TVM-FFI-OrcJIT + +on: + workflow_dispatch: + inputs: + branch: + description: "Branch or tag to publish (manual run)" + required: true + default: "main" + pypi_repository: + description: "PyPI repository (pypi or testpypi)" + required: true + default: "testpypi" + type: choice + options: + - pypi + - testpypi + +jobs: + build_wheels: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + include: + - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux2014} + - {os: ubuntu-latest, arch: x86_64, linux_image: manylinux_2_28} + - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux2014} + - {os: ubuntu-24.04-arm, arch: aarch64, linux_image: manylinux_2_28} + - {os: macos-14, arch: arm64, linux_image: ""} + + steps: + # Special handling for macOS arm64 + python 3.8/3.9 + - uses: actions/setup-python@v5 + with: + python-version: 3.9 + if: runner.os == 'macOS' && runner.arch == 'ARM64' + + - uses: astral-sh/setup-uv@v6 + if: matrix.os != 'macos-14' + + - uses: actions/checkout@v5 + with: + ref: ${{ inputs.branch }} + submodules: recursive + fetch-depth: 0 + fetch-tags: true + + - name: Print current commit + run: git log -1 --oneline + + - name: Install LLVM (Linux) + if: runner.os == 'Linux' + run: | + wget https://apt.llvm.org/llvm.sh + chmod +x llvm.sh + sudo ./llvm.sh 18 + sudo apt-get install -y llvm-18-dev + + - name: Install LLVM (macOS) + if: runner.os == 'macOS' + run: | + brew install llvm@18 + echo "LLVM_DIR=$(brew --prefix llvm@18)" >> $GITHUB_ENV + + - name: Build wheels + uses: pypa/cibuildwheel@v3.1.4 + env: + CIBW_ARCHS_MACOS: ${{ matrix.arch }} + CIBW_ARCHS_LINUX: ${{ matrix.arch }} + CIBW_MANYLINUX_X86_64_IMAGE: ${{ matrix.linux_image }} + CIBW_MANYLINUX_AARCH64_IMAGE: ${{ matrix.linux_image }} + CIBW_BUILD_VERBOSITY: 1 + CIBW_BUILD: cp39-* cp310-* cp311-* cp312-* + CIBW_SKIP: "*-musllinux_*" + # Install LLVM in the manylinux container + CIBW_BEFORE_ALL_LINUX: | + yum install -y wget + wget https://github.com/llvm/llvm-project/releases/download/llvmorg-18.1.8/clang+llvm-18.1.8-x86_64-linux-gnu-ubuntu-18.04.tar.xz + tar -xf clang+llvm-18.1.8-x86_64-linux-gnu-ubuntu-18.04.tar.xz + export PATH="$PWD/clang+llvm-18.1.8-x86_64-linux-gnu-ubuntu-18.04/bin:$PATH" + export LLVM_DIR="$PWD/clang+llvm-18.1.8-x86_64-linux-gnu-ubuntu-18.04" + # Placeholder: may need to install tvm-ffi first + CIBW_BEFORE_BUILD: | + pip install apache-tvm-ffi || echo "apache-tvm-ffi not yet published" + + with: + package-dir: addons/tvm-ffi-orcjit + output-dir: wheelhouse + + - uses: actions/upload-artifact@v4 + with: + name: cibw-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }} + path: ./wheelhouse/*.whl + + build_sdist: + name: Build source distribution + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + with: + ref: ${{ inputs.branch }} + submodules: recursive + fetch-depth: 0 + fetch-tags: true + + - uses: astral-sh/setup-uv@v6 + + - name: Build sdist + working-directory: addons/tvm-ffi-orcjit + run: pipx run build --sdist --outdir dist . + + - name: Check metadata + working-directory: addons/tvm-ffi-orcjit + run: pipx run twine check dist/* + + - uses: actions/upload-artifact@v4 + with: + name: cibw-sdist + path: addons/tvm-ffi-orcjit/dist/*.tar.gz + + upload_pypi: + needs: [build_wheels, build_sdist] + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write + attestations: write + steps: + - uses: actions/download-artifact@v4 + with: + # unpacks all CIBW artifacts into dist/ + pattern: cibw-* + path: dist + merge-multiple: true + + - name: Generate artifact attestation for sdist and wheels + uses: actions/attest-build-provenance@v1 + with: + subject-path: dist/* + + - name: Publish package distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + attestations: true + verbose: true + repository-url: ${{ inputs.pypi_repository == 'testpypi' && 'https://test.pypi.org/legacy/' || '' }} diff --git a/.github/workflows/orcjit-tests.yml b/.github/workflows/orcjit-tests.yml new file mode 100644 index 00000000..33bf486b --- /dev/null +++ b/.github/workflows/orcjit-tests.yml @@ -0,0 +1,104 @@ +# 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. + +name: TVM-FFI-OrcJIT CI Tests + +on: + push: + branches: + - main + - dev + - orcjit + paths: + - 'addons/tvm-ffi-orcjit/**' + - '.github/workflows/tvm-ffi-orcjit/ci_test.yml' + pull_request: + paths: + - 'addons/tvm-ffi-orcjit/**' + - '.github/workflows/tvm-ffi-orcjit/ci_test.yml' + workflow_dispatch: + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] + python-version: ['3.9'] + + steps: + - uses: actions/checkout@v5 + with: + submodules: recursive + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install LLVM (Ubuntu) + if: runner.os == 'Linux' + run: | + wget https://apt.llvm.org/llvm.sh + chmod +x llvm.sh + sudo ./llvm.sh 18 + sudo apt-get install -y llvm-18-dev clang-18 + echo "CC=clang-18" >> $GITHUB_ENV + echo "CXX=clang++-18" >> $GITHUB_ENV + + - name: Install LLVM (macOS) + if: runner.os == 'macOS' + run: | + brew install llvm@18 + echo "LLVM_DIR=$(brew --prefix llvm@18)" >> $GITHUB_ENV + echo "CC=$(brew --prefix llvm@18)/bin/clang" >> $GITHUB_ENV + echo "CXX=$(brew --prefix llvm@18)/bin/clang++" >> $GITHUB_ENV + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Install tvm-ffi (core package) + run: | + uv pip install -e . --system + + - name: Build and install tvm-ffi-orcjit + working-directory: addons/tvm-ffi-orcjit + run: | + uv pip install -e . --system + + - name: Install test dependencies + run: | + uv pip install pytest --system + + - name: Build test object files + working-directory: addons/tvm-ffi-orcjit/tests + run: | + cmake -B build + cmake --build build --target install + + - name: Run tests + working-directory: addons/tvm-ffi-orcjit + run: | + pytest tests/ -v + + - name: Run example + working-directory: addons/tvm-ffi-orcjit/examples/quick-start + run: | + cmake -B build + cmake --build build + python run.py diff --git a/.github/workflows/tvm-ffi-orcjit/ci_test.yml b/.github/workflows/tvm-ffi-orcjit/ci_test.yml new file mode 100644 index 00000000..df798038 --- /dev/null +++ b/.github/workflows/tvm-ffi-orcjit/ci_test.yml @@ -0,0 +1,104 @@ +# 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. + +name: TVM-FFI-OrcJIT CI Tests + +on: + push: + branches: + - main + - dev + - orcjit + paths: + - 'addons/tvm-ffi-orcjit/**' + - '.github/workflows/tvm-ffi-orcjit/ci_test.yml' + pull_request: + paths: + - 'addons/tvm-ffi-orcjit/**' + - '.github/workflows/tvm-ffi-orcjit/ci_test.yml' + workflow_dispatch: + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] + python-version: ['3.9', '3.10', '3.11', '3.12'] + + steps: + - uses: actions/checkout@v5 + with: + submodules: recursive + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install LLVM (Ubuntu) + if: runner.os == 'Linux' + run: | + wget https://apt.llvm.org/llvm.sh + chmod +x llvm.sh + sudo ./llvm.sh 18 + sudo apt-get install -y llvm-18-dev clang-18 + echo "CC=clang-18" >> $GITHUB_ENV + echo "CXX=clang++-18" >> $GITHUB_ENV + + - name: Install LLVM (macOS) + if: runner.os == 'macOS' + run: | + brew install llvm@18 + echo "LLVM_DIR=$(brew --prefix llvm@18)" >> $GITHUB_ENV + echo "CC=$(brew --prefix llvm@18)/bin/clang" >> $GITHUB_ENV + echo "CXX=$(brew --prefix llvm@18)/bin/clang++" >> $GITHUB_ENV + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Install tvm-ffi (core package) + run: | + uv pip install -e . --system + + - name: Build and install tvm-ffi-orcjit + working-directory: addons/tvm-ffi-orcjit + run: | + uv pip install -e . --system + + - name: Install test dependencies + run: | + uv pip install pytest --system + + - name: Build test objects + working-directory: addons/tvm-ffi-orcjit/tests + run: | + cmake -B build + cmake --build build + + - name: Run tests + working-directory: addons/tvm-ffi-orcjit + run: | + pytest tests/ -v + + - name: Run example + working-directory: addons/tvm-ffi-orcjit/examples/quick-start + run: | + cmake -B build + cmake --build build + python run.py diff --git a/addons/tvm-ffi-orcjit/CMakeLists.txt b/addons/tvm-ffi-orcjit/CMakeLists.txt new file mode 100644 index 00000000..7c5e575a --- /dev/null +++ b/addons/tvm-ffi-orcjit/CMakeLists.txt @@ -0,0 +1,71 @@ +cmake_minimum_required(VERSION 3.20) +project( + tvm_ffi_orcjit + VERSION 0.1.0 + LANGUAGES C CXX +) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +# Find dependencies +find_package(LLVM REQUIRED CONFIG) +message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}") +message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}") + +# Find tvm-ffi package using the same method as quickstart example +find_package( + Python + COMPONENTS Interpreter + REQUIRED +) + +# Get the cmake directory for tvm-ffi +execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE tvm_ffi_ROOT +) +message(STATUS "tvm_ffi_ROOT: ${tvm_ffi_ROOT}") + +find_package(tvm_ffi CONFIG REQUIRED) + +# LLVM components needed for ORC JIT v2 +llvm_map_components_to_libnames(LLVM_LIBS Core OrcJIT Support native) + +# Filter out non-existent targets +set(LLVM_LIBS_FILTERED) +foreach (lib ${LLVM_LIBS}) + if (TARGET ${lib}) + list(APPEND LLVM_LIBS_FILTERED ${lib}) + else () + message(STATUS "Skipping non-existent LLVM target: ${lib}") + endif () +endforeach () +set(LLVM_LIBS ${LLVM_LIBS_FILTERED}) + +# Source files +set(SOURCES src/ffi/orcjit_session.cc src/ffi/orcjit_dylib.cc) + +# Build shared library +add_library(tvm_ffi_orcjit SHARED ${SOURCES}) + +target_include_directories( + tvm_ffi_orcjit PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${LLVM_INCLUDE_DIRS} +) + +target_link_libraries(tvm_ffi_orcjit PRIVATE tvm_ffi_header tvm_ffi_shared LLVM) + +# Compile definitions +separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS}) +target_compile_definitions(tvm_ffi_orcjit PRIVATE ${LLVM_DEFINITIONS_LIST}) + +# Installation rules +install( + TARGETS tvm_ffi_orcjit + LIBRARY DESTINATION lib + ARCHIVE DESTINATION lib + RUNTIME DESTINATION bin +) diff --git a/addons/tvm-ffi-orcjit/README.md b/addons/tvm-ffi-orcjit/README.md new file mode 100644 index 00000000..bbc3b102 --- /dev/null +++ b/addons/tvm-ffi-orcjit/README.md @@ -0,0 +1,349 @@ + + + + + + + + + + + + + + + + + +# TVM-FFI OrcJIT + +A Python package that enables dynamic loading of compiled object files (`.o`) using LLVM ORC JIT v2, providing a flexible JIT execution environment for TVM-FFI exported functions. + +## Features + +- **JIT Execution**: Load and execute compiled object files at runtime using LLVM's ORC JIT v2 +- **Multiple Libraries**: Create separate dynamic libraries with independent symbol namespaces +- **Incremental Loading**: Add multiple object files to the same library incrementally +- **Symbol Isolation**: Different libraries can define the same symbol without conflicts +- **TVM-FFI Integration**: Seamlessly works with TVM-FFI's stable C ABI +- **Cross-Platform**: Supports Linux, macOS, and Windows (on the plan) +- **Python API**: Simple Pythonic interface for JIT compilation and execution + +## Installation + +### Prerequisites + +- Python 3.9+ +- CMake 3.18+ +- LLVM 18+ (with ORC JIT support) +- C/C++ compiler with C++17 support +- Ninja build system (optional, recommended) + +### Installing LLVM + +**Ubuntu/Debian:** + +```bash +sudo apt-get install -y llvm-18-dev +``` + +**macOS:** + +```bash +brew install llvm@18 +export LLVM_DIR=$(brew --prefix llvm@18) +export CC=$(brew --prefix llvm@18)/bin/clang +export CXX=$(brew --prefix llvm@18)/bin/clang++ +``` + +### Build from Source + +1. Install the core TVM-FFI package: + + ```bash + pip install apache-tvm-ffi + ``` + +1. Install tvm-ffi-orcjit: + + ```bash + pip install tvm-ffi-orcjit + ``` + +### Development Installation + +For development, install from source: + +```bash +git clone --recursive https://github.com/yaoyaoding/tvm-ffi.git +cd tvm-ffi/addons/tvm-ffi-orcjit +pip install -e . +``` + +## Usage + +### Basic Example + +```python +from tvm_ffi_orcjit import create_session + +# Create an execution session +session = create_session() + +# Create a dynamic library +lib = session.create_library() + +# Load an object file +lib.add("example.o") + +# Get and call a function +add_func = lib.get_function("test_add") +result = add_func(1, 2) +print(f"Result: {result}") # Output: Result: 3 +``` + +### Multiple Libraries with Symbol Isolation + +Create separate libraries to avoid symbol conflicts: + +```python +from tvm_ffi_orcjit import create_session + +session = create_session() + +# Create two separate libraries +lib1 = session.create_library("lib1") +lib2 = session.create_library("lib2") + +# Each library can have its own version of the same symbol +lib1.add("implementation_v1.o") # Contains test_add +lib2.add("implementation_v2.o") # Contains test_add with different behavior + +# Get functions from different libraries +add_v1 = lib1.get_function("test_add") +add_v2 = lib2.get_function("test_add") + +print(add_v1(5, 3)) # Uses implementation from lib1 +print(add_v2(5, 3)) # Uses implementation from lib2 +``` + +### Incremental Loading + +Add multiple object files to the same library: + +```python +from tvm_ffi_orcjit import create_session + +session = create_session() +lib = session.create_library() + +# Load multiple object files incrementally +lib.add("math_ops.o") +lib.add("string_ops.o") +lib.add("utils.o") + +# Access functions from any loaded object file +add = lib.get_function("test_add") +subtract = lib.get_function("test_subtract") +concat = lib.get_function("string_concat") + +print(add(10, 5)) # From math_ops.o +print(subtract(10, 5)) # From math_ops.o +print(concat("Hello", " World")) # From string_ops.o +``` + +## How It Works + +1. **ExecutionSession**: Manages the LLVM ORC JIT execution session and multiple dynamic libraries +2. **DynamicLibrary**: Represents a JITDylib that can load object files and resolve symbols +3. **Symbol Resolution**: Uses LLVM's ORC JIT v2 symbol lookup with proper linkage semantics +4. **Memory Management**: Allocates `__dso_handle` in JIT memory to ensure proper relocations +5. **TVM-FFI Integration**: Functions are exposed through TVM-FFI's PackedFunc interface + +### Technical Details + +- **ORC JIT v2**: Uses LLVM's modern JIT infrastructure (LLJIT) +- **Weak Linkage**: Each library gets its own `__dso_handle` with weak linkage +- **IR-based Allocation**: Creates LLVM IR modules for runtime symbols to ensure JIT memory allocation +- **Cross-Platform**: Correctly handles `.so` (Linux), `.dylib` (macOS), and `.dll` (Windows) + +## Development + +### Building Tests + +The project includes comprehensive tests with CMake-built test objects: + +```bash +cd tests +cmake -B build +cmake --build build +pytest -v +``` + +### Project Structure + +```text +tvm-ffi-orcjit/ +├── CMakeLists.txt # CMake build configuration +├── pyproject.toml # Python package metadata +├── README.md # This file +├── include/ +│ └── tvm/ffi/orcjit/ +│ ├── orcjit_session.h # ExecutionSession C++ header +│ └── orcjit_dylib.h # DynamicLibrary C++ header +├── src/ +│ └── ffi/ +│ ├── orcjit_session.cc # ExecutionSession implementation +│ └── orcjit_dylib.cc # DynamicLibrary implementation +├── python/ +│ └── tvm_ffi_orcjit/ +│ ├── __init__.py # Module exports and library loading +│ ├── session.py # Python ExecutionSession wrapper +│ └── dylib.py # Python DynamicLibrary wrapper +├── tests/ +│ ├── CMakeLists.txt # Test object file builds +│ ├── test_basic.py # Python tests +│ └── sources/ +│ ├── test_funcs.cc # Test functions +│ ├── test_funcs2.cc # Additional test functions +│ └── test_funcs_conflict.cc # Conflicting symbols for testing +└── examples/ + └── quick-start/ # Complete example with CMake +``` + +## Examples + +Complete examples are available in the `examples/` directory: + +- **quick-start/**: Demonstrates basic usage with a simple add function + - Shows how to compile C++ code with TVM-FFI exports + - Loads and executes the compiled object file + - Uses CMake for building the example + +## Writing Functions for OrcJIT + +Functions must use TVM-FFI's export macros: + +```cpp +#include + +// Simple function +TVM_FFI_DLL_EXPORT_TYPED_FUNC(simple_add, [](int a, int b) { + return a + b; +}); + +// Function with implementation +static int multiply_impl(int a, int b) { + return a * b; +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(simple_multiply, multiply_impl); +``` + +Compile with C++17: + +```bash +clang++ -std=c++17 -fPIC -c -o example.o example.cc +``` + +## Requirements + +The package depends on: + +- `apache-tvm-ffi>=0.1.0` - TVM-FFI core library +- LLVM 18+ (linked at build time) - For ORC JIT v2 functionality +- Python 3.9+ - For runtime + +## Known Limitations + +### Optimized Code and Relocations + +When compiling object files with optimization enabled (`-O2`, `-O3`), ensure your code doesn't generate PC-relative relocations that exceed ±2GB range. The package allocates `__dso_handle` in JIT memory to mitigate this, but extremely large programs may still encounter issues. + +**Workaround**: Compile test objects with `-O0` if you encounter "relocation out of range" errors during sequential test runs. + +## License + +Apache License 2.0 + +## Contributing + +Contributions are welcome! Please ensure that: + +1. Code follows the existing C++17 and Python style +2. New features include tests in `tests/test_basic.py` +3. Documentation is updated (README and docstrings) +4. CI tests pass on all platforms (Linux, macOS) + +## Troubleshooting + +### "Cannot find global function" error + +The shared library wasn't loaded. This usually means: + +- The library file extension doesn't match your platform +- The library wasn't installed correctly +- Python can't find the library file + +**Solution**: Reinstall the package: + +```bash +pip install --force-reinstall tvm-ffi-orcjit +``` + +### "Duplicate definition of symbol" error + +You're adding multiple object files with the same symbol to the same library. + +**Solution**: Use separate libraries for different implementations: + +```python +lib1 = session.create_library("lib1") +lib2 = session.create_library("lib2") +``` + +### "Symbol not found" error + +The symbol wasn't exported with TVM-FFI macros. + +**Solution**: Use `TVM_FFI_DLL_EXPORT_TYPED_FUNC`: + +```cpp +TVM_FFI_DLL_EXPORT_TYPED_FUNC(my_function, impl); +``` + +### Relocation errors with optimized code + +Object files compiled with `-O2` or higher may fail with "relocation out of range" in some scenarios. + +**Solution**: + +- Use `-O0` for test/development builds +- Run tests in separate processes (using `pytest-xdist`) +- This limitation primarily affects test scenarios, not production use + +### LLVM version mismatch + +The package requires LLVM 18+. Using older versions will cause build failures. + +**Solution**: Install LLVM 18: + +```bash +# Ubuntu +sudo ./llvm.sh 18 + +# macOS +brew install llvm@18 +``` + +### CMake can't find LLVM + +Set the `LLVM_DIR` environment variable: + +```bash +# macOS +export LLVM_DIR=$(brew --prefix llvm@18)/lib/cmake/llvm + +# Linux +export LLVM_DIR=/usr/lib/llvm-18/cmake +``` diff --git a/addons/tvm-ffi-orcjit/examples/quick-start/CMakeLists.txt b/addons/tvm-ffi-orcjit/examples/quick-start/CMakeLists.txt new file mode 100644 index 00000000..7fb1dfed --- /dev/null +++ b/addons/tvm-ffi-orcjit/examples/quick-start/CMakeLists.txt @@ -0,0 +1,54 @@ +# 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. + +cmake_minimum_required(VERSION 3.18) +project(tvm_ffi_orcjit_example) + +# Set C++ standard +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# Run `python -m tvm_ffi.config --cmakedir` to find tvm-ffi package +find_package( + Python + COMPONENTS Interpreter + REQUIRED +) +execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE tvm_ffi_ROOT +) +find_package(tvm_ffi CONFIG REQUIRED) + +# Create object library (not a shared library, just object files) +add_library(add_obj OBJECT add.cc) +target_link_libraries(add_obj PRIVATE tvm_ffi_header) + +# Set compiler flags +target_compile_options(add_obj PRIVATE -fPIC -O2) + +# Custom target to copy the object file to the example directory +add_custom_target( + copy_obj_file ALL + COMMAND ${CMAKE_COMMAND} -E copy $ ${CMAKE_CURRENT_SOURCE_DIR}/add.o + COMMAND ${CMAKE_COMMAND} -E echo "Successfully compiled add.o" + COMMAND ${CMAKE_COMMAND} -E echo "" + COMMAND ${CMAKE_COMMAND} -E echo "You can now run: python run.py" + DEPENDS add_obj + COMMENT "Copying add.o to example directory" +) diff --git a/addons/tvm-ffi-orcjit/examples/quick-start/README.md b/addons/tvm-ffi-orcjit/examples/quick-start/README.md new file mode 100644 index 00000000..bdb6b6b5 --- /dev/null +++ b/addons/tvm-ffi-orcjit/examples/quick-start/README.md @@ -0,0 +1,129 @@ + + + + + + + + + + + + + + + + + +# Quick Start Example + +This example demonstrates the basic usage of tvm-ffi-orcjit to compile C++ functions and load them dynamically at runtime. + +## What's Included + +- `add.cc` - Simple C++ source file with math functions exported via TVM-FFI +- `run.py` - Python script that loads and calls the compiled functions +- `CMakeLists.txt` - CMake configuration to compile the C++ code into an object file +- `compile.sh` - Legacy shell script (CMake is recommended for cross-platform support) + +## Prerequisites + +- Python 3.8+ +- CMake 3.18+ +- C++ compiler (g++, clang++, or MSVC) +- TVM-FFI and tvm-ffi-orcjit packages + +## Installation + +First, install the required packages: + +```bash +# Navigate to the repository root +cd ../../.. + +# Install TVM-FFI in editable mode +pip install -e . + +# Install tvm-ffi-orcjit in editable mode +pip install -e addons/tvm-ffi-orcjit + +# Return to the example directory +cd addons/tvm-ffi-orcjit/examples/quick-start +``` + +After installation, `tvm-ffi-config` will be available in your PATH and used by the compile script to get the correct include directories and compiler flags. + +## Steps + +### 1. Compile the C++ code + +Using CMake (recommended for cross-platform): + +```bash +cmake -B build +cmake --build build +``` + +Or using the legacy shell script (Unix-like systems only): + +```bash +./compile.sh +``` + +Both methods will create `add.o` - a compiled object file with exported functions. + +### 2. Run the Python loader + +```bash +python run.py +``` + +This will: + +- Load the object file using tvm-ffi-orcjit +- Call the exported functions +- Print the results + +## Expected Output + +```text +Loading object file: add.o +✓ Object file loaded successfully + +=== Testing add function === +add(10, 20) = 30 + +=== Testing multiply function === +multiply(7, 6) = 42 + +=== Testing fibonacci function === +fibonacci(10) = 55 + +=== Testing concat function === +concat('Hello, ', 'World!') = 'Hello, World!' + +================================================== +✓ All tests passed successfully! +================================================== +``` + +## How It Works + +1. **C++ Side** (`add.cc`): + - Functions are exported using `TVM_FFI_DLL_EXPORT_TYPED_FUNC` macro + - The macro registers functions with TVM-FFI's global function registry + +2. **Python Side** (`run.py`): + - `create_session()` creates an ORC JIT execution session + - `session.create_library()` creates a dynamic library (JITDylib) + - `lib.add()` loads the `.o` file into the JIT + - `lib.get_function()` looks up symbols in the JIT-compiled code + - Functions are called like normal Python functions + +## Next Steps + +- Modify `add.cc` to add your own functions +- Recompile with CMake: `cmake --build build` +- Load and test in Python + +For more details on the API, see the main package documentation. diff --git a/addons/tvm-ffi-orcjit/examples/quick-start/add.cc b/addons/tvm-ffi-orcjit/examples/quick-start/add.cc new file mode 100644 index 00000000..443638aa --- /dev/null +++ b/addons/tvm-ffi-orcjit/examples/quick-start/add.cc @@ -0,0 +1,46 @@ +/* + * 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. + */ + +/* + * Quick Start Example - Simple Math Functions + * + * This file demonstrates how to export C++ functions using TVM-FFI + * so they can be loaded dynamically at runtime with tvm-ffi-orcjit. + */ + +#include + +// Simple addition function +int add_impl(int a, int b) { return a + b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, add_impl); + +// Multiplication function +int multiply_impl(int a, int b) { return a * b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(multiply, multiply_impl); + +// Fibonacci function (recursive) +int fib_impl(int n) { + if (n <= 1) return n; + return fib_impl(n - 1) + fib_impl(n - 2); +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fibonacci, fib_impl); + +// String concatenation example +std::string concat_impl(std::string a, std::string b) { return a + b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(concat, concat_impl); diff --git a/addons/tvm-ffi-orcjit/examples/quick-start/run.py b/addons/tvm-ffi-orcjit/examples/quick-start/run.py new file mode 100755 index 00000000..5077f26d --- /dev/null +++ b/addons/tvm-ffi-orcjit/examples/quick-start/run.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +# 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. + +"""Quick Start Example - Load and call functions from add.o. + +This script demonstrates how to: +1. Create an ExecutionSession instance +2. Create a DynamicLibrary +3. Load a compiled object file +4. Get functions by name +5. Call them like regular Python functions +""" + +import sys +from pathlib import Path + +# Add the parent python directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "python")) + + +from tvm_ffi_orcjit import ExecutionSession + + +def main() -> int: + """Run the quick start example.""" + # Check if object file exists + obj_file = Path("add.o") + if not obj_file.exists(): + print(f"Error: {obj_file} not found!") + print("Please run ./compile.sh first to compile the C++ code.") + return 1 + + print(f"Loading object file: {obj_file}") + + # Create execution session and dynamic library + session = ExecutionSession() + lib = session.create_library() + lib.add(str(obj_file)) + + print("✓ Object file loaded successfully\n") + + # Get and call the 'add' function + print("=== Testing add function ===") + add = lib.get_function("add") + result = add(10, 20) + print(f"add(10, 20) = {result}") + assert result == 30, f"Expected 30, got {result}" + + # Get and call the 'multiply' function + print("\n=== Testing multiply function ===") + multiply = lib.get_function("multiply") + result = multiply(7, 6) + print(f"multiply(7, 6) = {result}") + assert result == 42, f"Expected 42, got {result}" + + # Get and call the 'fibonacci' function + print("\n=== Testing fibonacci function ===") + fibonacci = lib.get_function("fibonacci") + result = fibonacci(10) + print(f"fibonacci(10) = {result}") + assert result == 55, f"Expected 55, got {result}" + + # Get and call the 'concat' function + print("\n=== Testing concat function ===") + concat = lib.get_function("concat") + result = concat("Hello, ", "World!") + print(f"concat('Hello, ', 'World!') = '{result}'") + assert result == "Hello, World!", f"Expected 'Hello, World!', got '{result}'" + + print("\n" + "=" * 50) + print("✓ All tests passed successfully!") + print("=" * 50) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/addons/tvm-ffi-orcjit/pyproject.toml b/addons/tvm-ffi-orcjit/pyproject.toml new file mode 100644 index 00000000..f02b7103 --- /dev/null +++ b/addons/tvm-ffi-orcjit/pyproject.toml @@ -0,0 +1,71 @@ +# 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. + +[build-system] +requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"] +build-backend = "scikit_build_core.build" + +[project] +name = "tvm-ffi-orcjit" +version = "0.1.0" +description = "Load TVM-FFI exported object files using LLVM ORC JIT v2" +readme = "README.md" +requires-python = ">=3.8" +license = { text = "Apache-2.0" } +authors = [{ name = "TVM-FFI OrcJIT Contributors" }] +keywords = ["tvm-ffi", "llvm", "jit", "orcjit"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: C++", +] +dependencies = ["apache-tvm-ffi>=0.1.0"] + +[project.urls] +Homepage = "https://github.com/apache/tvm-ffi" +Repository = "https://github.com/apache/tvm-ffi" + +[tool.scikit-build] +cmake.version = "CMakeLists.txt" +cmake.build-type = "Release" +wheel.py-api = "py3" +build-dir = "build" +build.verbose = true +editable.rebuild = false +editable.verbose = true +wheel.packages = ["python/tvm_ffi_orcjit"] +wheel.install-dir = "tvm_ffi_orcjit" +# Don't set install-dir, let it use the default python/ directory +sdist.include = [ + "/README.md", + "/LICENSE", + "/pyproject.toml", + "/CMakeLists.txt", + "/include/**/*.h", + "/src/**/*.cc", + "/src/**/*.cpp", + "/python/**/*.py", +] + +[tool.scikit-build.cmake.define] +CMAKE_EXPORT_COMPILE_COMMANDS = "ON" diff --git a/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/__init__.py b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/__init__.py new file mode 100644 index 00000000..566c7e70 --- /dev/null +++ b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/__init__.py @@ -0,0 +1,81 @@ +# 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. + +"""TVM-FFI OrcJIT. + +This module provides functionality to load object files (.o) compiled with TVM-FFI +exports using LLVM ORC JIT v2. + +Example: + >>> from tvm_ffi_orcjit import create_session + >>> session = create_session() + >>> lib = session.create_library() + >>> lib.add("example.o") + >>> func = lib.get_function("my_function") + >>> result = func(arg1, arg2) + +""" + +import ctypes +import platform +from pathlib import Path + +from tvm_ffi import load_module + +# Determine the library extension based on platform +if platform.system() == "Darwin": + _LIB_EXT = "dylib" +elif platform.system() == "Windows": + _LIB_EXT = "dll" +else: + _LIB_EXT = "so" + +# Load the orcjit extension library +_LIB_PATH = [ + Path(__file__).parent.parent.parent / "lib" / f"libtvm_ffi_orcjit.{_LIB_EXT}", + Path(__file__).parent.parent.parent / "build" / f"libtvm_ffi_orcjit.{_LIB_EXT}", +] +_lib_path_str = None +for path in _LIB_PATH: + if path.exists(): + _ = load_module(str(path)) + _lib_path_str = str(path) +if _lib_path_str is None: + raise RuntimeError( + f"Could not find libtvm_ffi_orcjit.{_LIB_EXT}. " + f"Searched in {_LIB_PATH} and site-packages. " + f"Please ensure the package is installed correctly." + ) + +# Explicitly initialize the library to register functions +# This is needed because static initializers may not run when loaded via dlopen +try: + # Load the library with ctypes and call the initialization function + c_lib = ctypes.CDLL(_lib_path_str, mode=ctypes.RTLD_GLOBAL) + init_func = c_lib.TVMFFIOrcJITInitialize + init_func.restype = None + init_func() +except Exception as e: + import warnings + + warnings.warn(f"Failed to explicitly initialize orcjit library: {e}") + +from .dylib import DynamicLibrary +from .session import ExecutionSession + +__all__ = ["DynamicLibrary", "ExecutionSession"] +__version__ = "0.1.0" diff --git a/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/_ffi_api.py b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/_ffi_api.py new file mode 100644 index 00000000..fb2464b4 --- /dev/null +++ b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/_ffi_api.py @@ -0,0 +1,21 @@ +# 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. +"""FFI APIs for orcjit.""" + +import tvm_ffi + +tvm_ffi.init_ffi_api("orcjit", __name__) diff --git a/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py new file mode 100644 index 00000000..446b8a40 --- /dev/null +++ b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/dylib.py @@ -0,0 +1,85 @@ +# 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. +"""ORC JIT Dynamic Library.""" + +from __future__ import annotations + +from pathlib import Path + +from tvm_ffi import Module + + +class DynamicLibrary(Module): + """ORC JIT Dynamic Library (JITDylib). + + Represents a collection of symbols that can be loaded from object files and linked + against other dynamic libraries. Supports JIT compilation and symbol resolution. + + Examples + -------- + >>> session = create_session() + >>> lib = session.create_library() + >>> lib.add("add.o") + >>> lib.add("multiply.o") + >>> add_func = lib.get_function("add") + >>> result = add_func(1, 2) + + """ + + def add(self, object_file: str | Path) -> None: + """Add an object file to this dynamic library. + + Parameters + ---------- + object_file : str or Path + Path to the object file to load. + + Examples + -------- + >>> lib.add("add.o") + >>> lib.add(Path("multiply.o")) + + """ + if isinstance(object_file, Path): + object_file = str(object_file) + self.get_function("orcjit.add_object_file")(object_file) + + def set_link_order(self, *libraries: DynamicLibrary) -> None: + """Set the link order for symbol resolution. + + When resolving symbols, this library will search in the specified libraries + in the order provided. This replaces any previous link order. + + Parameters + ---------- + *libraries : DynamicLibrary + One or more dynamic libraries to search for symbols (in order). + + Examples + -------- + >>> session = create_session() + >>> lib_utils = session.create_library() + >>> lib_utils.add("utils.o") + >>> lib_core = session.create_library() + >>> lib_core.add("core.o") + >>> lib_main = session.create_library() + >>> lib_main.add("main.o") + >>> # main can call symbols from utils and core (utils searched first) + >>> lib_main.set_link_order(lib_utils, lib_core) + + """ + self.get_function("orcjit.set_link_order")(libraries) diff --git a/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/session.py b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/session.py new file mode 100644 index 00000000..1269c7b3 --- /dev/null +++ b/addons/tvm-ffi-orcjit/python/tvm_ffi_orcjit/session.py @@ -0,0 +1,60 @@ +# 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. +"""ORC JIT Execution Session.""" + +from __future__ import annotations + +from tvm_ffi import Object, register_object + +from . import _ffi_api +from .dylib import DynamicLibrary + + +@register_object("orcjit.ExecutionSession") +class ExecutionSession(Object): + """ORC JIT Execution Session. + + Manages the LLVM ORC JIT execution environment and creates dynamic libraries (JITDylibs). + This is the top-level context for JIT compilation and symbol management. + + Examples + -------- + >>> session = ExecutionSession() + >>> lib = session.create_library(name="main") + >>> lib.add("add.o") + >>> add_func = lib.get_function("add") + + """ + + def __init__(self) -> None: + """Initialize ExecutionSession from a handle.""" + self.__init_handle_by_constructor__(_ffi_api.ExecutionSession) # type: ignore + + def create_library(self, name: str = "") -> DynamicLibrary: + """Create a new dynamic library associated with this execution session. + + Args: + name: Optional name for the library. If empty, a unique name will be generated. + + Returns: + A new DynamicLibrary instance. + + """ + handle = _ffi_api.ExecutionSessionCreateDynamicLibrary(self, name) # type: ignore + lib = DynamicLibrary.__new__(DynamicLibrary) + lib.__move_handle_from__(handle) + return lib diff --git a/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.cc b/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.cc new file mode 100644 index 00000000..4c716e01 --- /dev/null +++ b/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.cc @@ -0,0 +1,180 @@ +/* + * 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. + */ + +/*! + * \file orcjit_dylib.cc + * \brief LLVM ORC JIT DynamicLibrary implementation + */ + +#include "orcjit_dylib.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "orcjit_session.h" +#include "tvm/ffi/function.h" + +namespace tvm { +namespace ffi { +namespace orcjit { + +ORCJITDynamicLibraryObj::ORCJITDynamicLibraryObj(ObjectPtr session, + llvm::orc::JITDylib* dylib, llvm::orc::LLJIT* jit, + String name) + : session_(std::move(session)), dylib_(dylib), jit_(jit), name_(std::move(name)) { + if (void** ctx_addr = reinterpret_cast(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(GetSymbol(ffi::symbol::tvm_ffi_library_ctx))) { + *ctx_addr = symbol; + } + }); + TVM_FFI_CHECK(dylib_ != nullptr, ValueError) << "JITDylib cannot be null"; + TVM_FFI_CHECK(jit_ != nullptr, ValueError) << "LLJIT cannot be null"; +} + +void ORCJITDynamicLibraryObj::AddObjectFile(const String& path) { + // Read object file + auto buffer_or_err = llvm::MemoryBuffer::getFile(path.c_str()); + if (!buffer_or_err) { + TVM_FFI_THROW(IOError) << "Failed to read object file: " << path; + } + + // Add object file to this JITDylib + auto err = jit_->addObjectFile(*dylib_, std::move(*buffer_or_err)); + if (err) { + std::string err_msg; + llvm::handleAllErrors(std::move(err), + [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); + TVM_FFI_THROW(ValueError) << "Failed to add object file '" << path << "': " << err_msg; + } +} + +void ORCJITDynamicLibraryObj::SetLinkOrder(const std::vector& dylibs) { + // Clear and rebuild the link order + link_order_.clear(); + + for (auto* lib : dylibs) { + link_order_.emplace_back(lib, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); + } + + // Set the link order in the LLVM JITDylib + dylib_->setLinkOrder(link_order_, false); +} + +void* ORCJITDynamicLibraryObj::GetSymbol(const String& name) { + // Build search order: this dylib first, then all linked dylibs + llvm::orc::JITDylibSearchOrder search_order; + search_order.emplace_back(dylib_, llvm::orc::JITDylibLookupFlags::MatchAllSymbols); + // Append linked libraries + search_order.insert(search_order.end(), link_order_.begin(), link_order_.end()); + + // Look up symbol using the full search order + auto symbol_or_err = + jit_->getExecutionSession().lookup(search_order, jit_->mangleAndIntern(name.c_str())); + + // Convert ExecutorAddr to pointer + return symbol_or_err ? symbol_or_err->getAddress().toPtr() : nullptr; +} + +llvm::orc::JITDylib& ORCJITDynamicLibraryObj::GetJITDylib() { + TVM_FFI_CHECK(dylib_ != nullptr, InternalError) << "JITDylib is null"; + return *dylib_; +} + +Optional ORCJITDynamicLibraryObj::GetFunction(const String& name) { + if (name == "orcjit.add_object_file") { + return Function::FromTyped([this](const String& path) { AddObjectFile(path); }); + } + if (name == "orcjit.set_link_order") { + return Function::FromTyped([this](const Array& libraries) { + std::vector libs; + libs.reserve(libraries.size()); + for (const ORCJITDynamicLibrary& lib : libraries) { + libs.push_back(&GetJITDylib()); + } + SetLinkOrder(libs); + }); + } + + // TVM-FFI exports have __tvm_ffi_ prefix + std::string symbol_name = symbol::tvm_ffi_symbol_prefix + std::string(name); + + // Try to get the symbol - return NullOpt if not found + if (void* symbol = GetSymbol(symbol_name)) { + cantFail(jit_->initialize(*dylib_)); + // Wrap C function pointer as tvm-ffi Function + TVMFFISafeCallType c_func = reinterpret_cast(symbol); + + return Function::FromPacked([c_func, name](PackedArgs args, Any* rv) { + TVM_FFI_ICHECK_LT(rv->type_index(), ffi::TypeIndex::kTVMFFIStaticObjectBegin); + TVM_FFI_CHECK_SAFE_CALL((*c_func)(nullptr, reinterpret_cast(args.data()), + args.size(), reinterpret_cast(rv))); + }); + } + return std::nullopt; +} + +//------------------------------------- +// Registration +//------------------------------------- + +static void RegisterOrcJITFunctions() { + static bool registered = false; + if (registered) return; + registered = true; + + namespace refl = tvm::ffi::reflection; + + refl::ObjectDef(); + + refl::GlobalDef() + .def("orcjit.ExecutionSession", []() { return ORCJITExecutionSession(); }) + .def("orcjit.ExecutionSessionCreateDynamicLibrary", + [](const ORCJITExecutionSession& session, const String& name) -> Module { + return session->CreateDynamicLibrary(name); + }); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + // This block may not execute when loaded via dlopen on some platforms. + // Call TVMFFIOrcJITInitialize() explicitly if functions are not registered. + RegisterOrcJITFunctions(); +} + +} // namespace orcjit +} // namespace ffi +} // namespace tvm + +// C API for explicit initialization +extern "C" { + +TVM_FFI_DLL_EXPORT void TVMFFIOrcJITInitialize() { tvm::ffi::orcjit::RegisterOrcJITFunctions(); } +} diff --git a/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.h b/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.h new file mode 100644 index 00000000..1d1eec13 --- /dev/null +++ b/addons/tvm-ffi-orcjit/src/ffi/orcjit_dylib.h @@ -0,0 +1,126 @@ +/* + * 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. + */ + +/*! + * \file orcjit_dylib.h + * \brief LLVM ORC JIT DynamicLibrary (JITDylib) wrapper + */ +#ifndef TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ +#define TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ + +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace ffi { +namespace orcjit { + +class ORCJITExecutionSessionObj; + +class ORCJITDynamicLibraryObj : public ModuleObj { + public: + /*! + * \brief Constructor + * \param session The parent execution session + * \param dylib The LLVM JITDylib + * \param jit The LLJIT instance + * \param name The library name + */ + ORCJITDynamicLibraryObj(ObjectPtr session, llvm::orc::JITDylib* dylib, + llvm::orc::LLJIT* jit, String name); + + const char* kind() const final { return "orcjit"; } + + Optional GetFunction(const String& name) override; + + private: + /*! + * \brief Add an object file to this library + * \param path Path to the object file to load + */ + void AddObjectFile(const String& path); + + /*! + * \brief Set the link order for symbol resolution + * \param dylibs Vector of libraries to search for symbols (in order) + * + * When resolving symbols, this library will search in the specified libraries + * in the order provided. This replaces any previous link order. + */ + void SetLinkOrder(const std::vector& dylibs); + + /*! + * \brief Look up a symbol in this library + * \param name The symbol name to look up + * \return Pointer to the symbol, or nullptr if not found + */ + void* GetSymbol(const String& name); + + /*! + * \brief Get the underlying LLVM JITDylib + * \return Reference to the LLVM JITDylib + */ + llvm::orc::JITDylib& GetJITDylib(); + + /*! + * \brief Get the name of this library + * \return The library name + */ + String GetName() const { return name_; } + + /*! \brief Parent execution session (for lifetime management) */ + ObjectPtr session_; + + /*! \brief The LLVM JITDylib */ + llvm::orc::JITDylib* dylib_; + + /*! \brief The LLJIT instance (for addObjectFile API) */ + llvm::orc::LLJIT* jit_; + + /*! \brief Library name */ + String name_; + + /*! \brief Link order tracking (to support incremental linking) */ + llvm::orc::JITDylibSearchOrder link_order_; +}; + +/*! + * \brief DynamicLibrary wrapper for LLVM ORC JIT v2 JITDylib + * + * This class wraps an LLVM JITDylib and provides functionality to: + * - Load object files + * - Link against other dynamic libraries + * - Look up symbols + */ +class ORCJITDynamicLibrary : public Module { + public: + explicit ORCJITDynamicLibrary(const ObjectPtr& ptr) : Module(ptr) {}; + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ORCJITDynamicLibrary, Module, + ORCJITDynamicLibraryObj); +}; + +} // namespace orcjit +} // namespace ffi +} // namespace tvm + +#endif // TVM_FFI_ORCJIT_ORCJIT_DYLIB_H_ diff --git a/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.cc b/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.cc new file mode 100644 index 00000000..5d7de046 --- /dev/null +++ b/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.cc @@ -0,0 +1,160 @@ +/* + * 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. + */ + +/*! + * \file orcjit_session.cc + * \brief LLVM ORC JIT ExecutionSession implementation + */ + +#include "orcjit_session.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "orcjit_dylib.h" +#include "tvm/ffi/object.h" + +namespace tvm { +namespace ffi { +namespace orcjit { + +// Initialize LLVM native target (only once) +struct LLVMInitializer { + LLVMInitializer() { + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + llvm::InitializeNativeTargetAsmParser(); + } +}; + +static LLVMInitializer llvm_initializer; + +ORCJITExecutionSessionObj::ORCJITExecutionSessionObj() : jit_(nullptr), dylib_counter_(0) {} + +void ORCJITExecutionSessionObj::Initialize() { + // Create LLJIT instance + auto jit_or_err = llvm::orc::LLJITBuilder() + .setPlatformSetUp(llvm::orc::ExecutorNativePlatform( + "/usr/lib/llvm-20/lib/clang/20/lib/linux/liborc_rt-x86_64.a")) + .create(); + if (!jit_or_err) { + auto err = jit_or_err.takeError(); + std::string err_msg; + llvm::handleAllErrors(std::move(err), + [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); + TVM_FFI_THROW(InternalError) << "Failed to create LLJIT: " << err_msg; + } + jit_ = std::move(*jit_or_err); +} + +ORCJITExecutionSession::ORCJITExecutionSession() { + ObjectPtr obj = make_object(); + obj->Initialize(); + data_ = std::move(obj); +} + +ORCJITDynamicLibrary ORCJITExecutionSessionObj::CreateDynamicLibrary(const String& name) { + TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; + + // Generate name if not provided + String lib_name = name; + if (lib_name.empty()) { + std::ostringstream oss; + oss << "dylib_" << dylib_counter_++; + lib_name = oss.str(); + } + + // Check if library with this name already exists + TVM_FFI_CHECK(dylibs_.find(lib_name) == dylibs_.end(), ValueError) + << "DynamicLibrary with name '" << lib_name << "' already exists"; + + // Create a new JITDylib + // auto& jd = jit_->getExecutionSession().createBareJITDylib(lib_name.c_str()); + auto& jd = jit_->getMainJITDylib(); + + // Add process symbol resolver to make C/C++ stdlib available + auto dlsg = llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( + jit_->getDataLayout().getGlobalPrefix()); + if (!dlsg) { + TVM_FFI_THROW(InternalError) << "Failed to create process symbol resolver"; + } + jd.addGenerator(std::move(*dlsg)); + + // Add __dso_handle by compiling a minimal LLVM IR module containing it. + // This ensures __dso_handle is allocated in JIT memory (within 2GB of code), + // avoiding "relocation out of range" errors with optimized code. + // + // We create an IR module with a global variable for __dso_handle, then compile + // it through the normal IR compilation path. JITLink will allocate it properly. + auto Ctx = std::make_unique(); + auto M = std::make_unique("__dso_handle_module", *Ctx); + M->setDataLayout(jit_->getDataLayout()); + M->setTargetTriple(jit_->getTargetTriple().str()); + + // Create a global variable: i8 __dso_handle = 0 + auto* Int8Ty = llvm::Type::getInt8Ty(*Ctx); + auto* DsoHandle = new llvm::GlobalVariable( + *M, Int8Ty, + false, // not constant + llvm::GlobalValue::WeakAnyLinkage, // Use weak linkage so multiple dylibs can define it + llvm::ConstantInt::get(Int8Ty, 0), "__dso_handle"); + DsoHandle->setVisibility(llvm::GlobalValue::DefaultVisibility); + + // Add the module to THIS specific JITDylib using the IR layer + auto& CompileLayer = jit_->getIRCompileLayer(); + if (auto Err = CompileLayer.add(jd, llvm::orc::ThreadSafeModule(std::move(M), std::move(Ctx)))) { + std::string err_msg; + llvm::handleAllErrors(std::move(Err), + [&](const llvm::ErrorInfoBase& eib) { err_msg = eib.message(); }); + TVM_FFI_THROW(InternalError) << "Failed to add __dso_handle module: " << err_msg; + } + + // Create the wrapper object + auto dylib = ORCJITDynamicLibrary(make_object( + GetObjectPtr(this), &jd, jit_.get(), lib_name)); + + // Store for lifetime management + dylibs_.insert({lib_name, dylib}); + + return dylib; +} + +llvm::orc::ExecutionSession& ORCJITExecutionSessionObj::GetLLVMExecutionSession() { + TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; + return jit_->getExecutionSession(); +} + +llvm::orc::LLJIT& ORCJITExecutionSessionObj::GetLLJIT() { + TVM_FFI_CHECK(jit_ != nullptr, InternalError) << "ExecutionSession not initialized"; + return *jit_; +} + +} // namespace orcjit +} // namespace ffi +} // namespace tvm diff --git a/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.h b/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.h new file mode 100644 index 00000000..016d0b18 --- /dev/null +++ b/addons/tvm-ffi-orcjit/src/ffi/orcjit_session.h @@ -0,0 +1,115 @@ +/* + * 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. + */ + +/*! + * \file orcjit_session.h + * \brief LLVM ORC JIT ExecutionSession wrapper + */ +#ifndef TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ +#define TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ + +#include +#include +#include + +#include +#include +#include + +namespace tvm { +namespace ffi { +namespace orcjit { + +// Forward declaration +class ORCJITDynamicLibrary; + +/*! + * \brief ExecutionSession object for LLVM ORC JIT v2 + * + * This class manages the lifetime of an LLVM ExecutionSession and provides + * functionality to create and manage multiple JITDylibs (DynamicLibraries). + */ +class ORCJITExecutionSessionObj : public Object { + public: + /*! + * \brief Default constructor (for make_object) + */ + ORCJITExecutionSessionObj(); + + /*! + * \brief Initialize the LLJIT instance + */ + void Initialize(); + + /*! + * \brief Create a new DynamicLibrary (JITDylib) in this session + * \param name Optional name for the library (for debugging) + * \return The created dynamic library instance + */ + ORCJITDynamicLibrary CreateDynamicLibrary(const String& name); + + /*! + * \brief Get the underlying LLVM ExecutionSession + * \return Reference to the LLVM ExecutionSession + */ + llvm::orc::ExecutionSession& GetLLVMExecutionSession(); + + /*! + * \brief Get the underlying LLJIT instance + * \return Reference to the LLJIT instance + */ + llvm::orc::LLJIT& GetLLJIT(); + + static constexpr bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("orcjit.ExecutionSession", ORCJITExecutionSessionObj, Object); + + private: + /*! \brief The LLVM ORC JIT instance */ + std::unique_ptr jit_; + + /*! \brief Counter for auto-generating library names */ + int dylib_counter_ = 0; + + /*! \brief Map of created dynamic libraries for lifetime management */ + std::unordered_map dylibs_; +}; + +/*! + * \brief Reference wrapper for ORCJITExecutionSessionObj + * + * A reference wrapper serves as a reference-counted pointer to the session object. + */ +class ORCJITExecutionSession : public ObjectRef { + public: + /*! + * \brief Create a new ExecutionSession + * \return The created execution session instance + */ + ORCJITExecutionSession(); + + // Required: define object reference methods + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ORCJITExecutionSession, ObjectRef, + ORCJITExecutionSessionObj); +}; + +} // namespace orcjit +} // namespace ffi +} // namespace tvm + +#endif // TVM_FFI_ORCJIT_ORCJIT_SESSION_H_ diff --git a/addons/tvm-ffi-orcjit/tests/CMakeLists.txt b/addons/tvm-ffi-orcjit/tests/CMakeLists.txt new file mode 100644 index 00000000..9d3f5826 --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/CMakeLists.txt @@ -0,0 +1,81 @@ +# 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. + +cmake_minimum_required(VERSION 3.18) +project(tvm_ffi_orcjit_tests) + +# Set C++ standard +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# Run `python -m tvm_ffi.config --cmakedir` to find tvm-ffi package +find_package( + Python + COMPONENTS Interpreter + REQUIRED +) +execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE tvm_ffi_ROOT +) +find_package(tvm_ffi CONFIG REQUIRED) + +# Create object library for test functions +add_library(test_funcs_obj OBJECT sources/test_funcs.cc) +target_link_libraries(test_funcs_obj PRIVATE tvm_ffi_header) +target_compile_options(test_funcs_obj PRIVATE -fPIC -O2) +install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR} + RENAME test_funcs.o +) + +# Create object library for second set of test functions +add_library(test_funcs2_obj OBJECT sources/test_funcs2.cc) +target_link_libraries(test_funcs2_obj PRIVATE tvm_ffi_header) +target_compile_options(test_funcs2_obj PRIVATE -fPIC -O2) +install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR} + RENAME test_funcs2.o +) + +# Create object library for conflicting test functions +add_library(test_funcs_conflict_obj OBJECT sources/test_funcs_conflict.cc) +target_link_libraries(test_funcs_conflict_obj PRIVATE tvm_ffi_header) +target_compile_options(test_funcs_conflict_obj PRIVATE -fPIC -O2) +install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR} + RENAME test_funcs_conflict.o +) + +find_package(CUDAToolkit) + +if (CUDAToolkit_FOUND) + enable_language(CUDA) + message(STATUS "CUDA found: ${CUDAToolkit_VERSION}") + add_library(test_funcs_cuda_obj OBJECT sources/test_funcs_cuda.cu) + target_link_libraries(test_funcs_cuda_obj PRIVATE tvm_ffi_header) + target_compile_options(test_funcs_cuda_obj PRIVATE -fPIC -O2) + install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR} + RENAME test_funcs_cuda.o + ) +endif () diff --git a/addons/tvm-ffi-orcjit/tests/README.md b/addons/tvm-ffi-orcjit/tests/README.md new file mode 100644 index 00000000..852ff6ff --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/README.md @@ -0,0 +1,58 @@ + + + + + + + + + + + + + + + + + +# TVM-FFI-OrcJIT Tests + +This directory contains tests for the tvm-ffi-orcjit package. + +## Building Test Objects + +The tests require pre-built object files. To build them: + +```bash +cd tests +cmake -B build +cmake --build build --target install +``` + +This will compile `sources/test_funcs.cc` and generate `test_funcs.o` in the tests directory. + +## Running Tests + +After building the test objects, run the tests with: + +```bash +pytest tests/ -v +``` + +Or from the repository root: + +```bash +cd addons/tvm-ffi-orcjit +pytest tests/ -v +``` + +## Test Structure + +- `sources/` - C++ source files for test functions +- `test_basic.py` - Python test cases +- `CMakeLists.txt` - Build configuration for test objects +- `test_funcs.o` - Generated object file (after building) + +## CI/CD + +The CI workflow automatically builds the test objects before running tests. See `.github/workflows/tvm-ffi-orcjit/ci_test.yml` for the full workflow. diff --git a/addons/tvm-ffi-orcjit/tests/__init__.py b/addons/tvm-ffi-orcjit/tests/__init__.py new file mode 100644 index 00000000..4d4bdfa8 --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/__init__.py @@ -0,0 +1,17 @@ +# 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. +"""Tests for tvm-ffi-orcjit package.""" diff --git a/addons/tvm-ffi-orcjit/tests/sources/test_funcs.cc b/addons/tvm-ffi-orcjit/tests/sources/test_funcs.cc new file mode 100644 index 00000000..94aecfbc --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/sources/test_funcs.cc @@ -0,0 +1,26 @@ +// 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 + +// Simple addition function +int test_add_impl(int a, int b) { return a + b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_impl); + +// Multiplication function +int test_multiply_impl(int a, int b) { return a * b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_impl); diff --git a/addons/tvm-ffi-orcjit/tests/sources/test_funcs2.cc b/addons/tvm-ffi-orcjit/tests/sources/test_funcs2.cc new file mode 100644 index 00000000..d786bc36 --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/sources/test_funcs2.cc @@ -0,0 +1,26 @@ +// 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 + +// Subtraction function +int test_subtract_impl(int a, int b) { return a - b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_subtract, test_subtract_impl); + +// Division function +int test_divide_impl(int a, int b) { return a / b; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_divide, test_divide_impl); diff --git a/addons/tvm-ffi-orcjit/tests/sources/test_funcs_conflict.cc b/addons/tvm-ffi-orcjit/tests/sources/test_funcs_conflict.cc new file mode 100644 index 00000000..09d35a84 --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/sources/test_funcs_conflict.cc @@ -0,0 +1,26 @@ +// 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 + +// Conflicting test_add function - different implementation +int test_add_conflict_impl(int a, int b) { return a + b + 1000; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_conflict_impl); + +// Conflicting test_multiply function - different implementation +int test_multiply_conflict_impl(int a, int b) { return a * b * 2; } +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_conflict_impl); diff --git a/addons/tvm-ffi-orcjit/tests/sources/test_funcs_cuda.cu b/addons/tvm-ffi-orcjit/tests/sources/test_funcs_cuda.cu new file mode 100644 index 00000000..fda0d7c3 --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/sources/test_funcs_cuda.cu @@ -0,0 +1,61 @@ +// 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 + +#include + +// Simple addition function +__global__ void test_add_kernel(int* a, int* b, int* c) { *c = *a + *b; } +int test_add_impl(int a, int b) { + int c; + int *d_a, *d_b, *d_c; + cudaMalloc(&d_a, sizeof(int)); + cudaMalloc(&d_b, sizeof(int)); + cudaMalloc(&d_c, sizeof(int)); + cudaMemcpy(d_a, &a, sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_b, &b, sizeof(int), cudaMemcpyHostToDevice); + test_add_kernel<<<1, 1>>>(d_a, d_b, d_c); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) printf("Kernel launch error: %s\n", cudaGetErrorString(err)); + cudaMemcpy(&c, d_c, sizeof(int), cudaMemcpyDeviceToHost); + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_c); + return c; +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_add, test_add_impl); + +// Multiplication function + +__global__ void test_multiply_kernel(int* a, int* b, int* c) { *c = *a * *b; } +int test_multiply_impl(int a, int b) { + int c; + int *d_a, *d_b, *d_c; + cudaMalloc(&d_a, sizeof(int)); + cudaMalloc(&d_b, sizeof(int)); + cudaMalloc(&d_c, sizeof(int)); + cudaMemcpy(d_a, &a, sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_b, &b, sizeof(int), cudaMemcpyHostToDevice); + test_multiply_kernel<<<1, 1>>>(d_a, d_b, d_c); + cudaMemcpy(&c, d_c, sizeof(int), cudaMemcpyDeviceToHost); + cudaFree(d_a); + cudaFree(d_b); + cudaFree(d_c); + return c; +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_multiply, test_multiply_impl); diff --git a/addons/tvm-ffi-orcjit/tests/test_basic.py b/addons/tvm-ffi-orcjit/tests/test_basic.py new file mode 100644 index 00000000..7c21b48a --- /dev/null +++ b/addons/tvm-ffi-orcjit/tests/test_basic.py @@ -0,0 +1,261 @@ +# 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. +"""Basic tests for tvm-ffi-orcjit functionality.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from tvm_ffi_orcjit import ExecutionSession + + +def get_test_obj_file(object_file: str) -> Path: + """Get the path to the pre-built test object file. + + Returns + ------- + Path + Path to the test_funcs.o object file. + + """ + # The object file should be built by CMake and located in the tests directory + test_dir = Path(__file__).parent + obj_file = test_dir / object_file + + if not obj_file.exists(): + raise FileNotFoundError( + f"Test object file not found: {obj_file}\n" + "Please build the test object file first:\n" + " cd tests && cmake -B build && cmake --build build" + ) + + return obj_file + + +def test_create_session() -> None: + """Test creating an execution session.""" + session = ExecutionSession() + assert session is not None + + +def test_create_library() -> None: + """Test creating a dynamic library.""" + session = ExecutionSession() + lib = session.create_library() + assert lib is not None + + +def test_load_and_execute_function() -> None: + """Test loading an object file and executing a function.""" + # Get pre-built test object file + obj_file = get_test_obj_file("test_funcs.o") + + # Create session and library + session = ExecutionSession() + lib = session.create_library() + + # Load object file + lib.add(str(obj_file)) + + # Get and call test_add function + add_func = lib.get_function("test_add") + result = add_func(10, 20) + assert result == 30 + + # Get and call test_multiply function + mul_func = lib.get_function("test_multiply") + result = mul_func(7, 6) + assert result == 42 + + +def test_multiple_libraries() -> None: + """Test creating and using multiple libraries.""" + session = ExecutionSession() + + lib1 = session.create_library("lib1") + lib2 = session.create_library("lib2") + + assert lib1 is not None + assert lib2 is not None + + +def test_function_not_found() -> None: + """Test that getting a non-existent function raises an error.""" + # Get pre-built test object file + obj_file = get_test_obj_file("test_funcs.o") + + session = ExecutionSession() + lib = session.create_library() + lib.add(str(obj_file)) + + with pytest.raises(AttributeError, match="Module has no function"): + lib.get_function("nonexistent_function") + + +def test_gradually_add_objects_to_same_library() -> None: + """Test gradually adding multiple object files to the same library.""" + obj_file1 = get_test_obj_file("test_funcs.o") + obj_file2 = get_test_obj_file("test_funcs2.o") + + session = ExecutionSession() + lib = session.create_library() + + # Add first object file + lib.add(str(obj_file1)) + + # Test functions from first object + add_func = lib.get_function("test_add") + assert add_func(5, 3) == 8 + + mul_func = lib.get_function("test_multiply") + assert mul_func(4, 5) == 20 + + # Add second object file to the same library + lib.add(str(obj_file2)) + + # Test functions from second object + sub_func = lib.get_function("test_subtract") + assert sub_func(10, 3) == 7 + + div_func = lib.get_function("test_divide") + assert div_func(20, 4) == 5 + + # Verify first object's functions still work + assert add_func(10, 20) == 30 + assert mul_func(7, 6) == 42 + + +def test_two_separate_libraries() -> None: + """Test creating two separate libraries each with its own object file.""" + obj_file1 = get_test_obj_file("test_funcs.o") + obj_file2 = get_test_obj_file("test_funcs2.o") + + session = ExecutionSession() + + # Create first library with first object + lib1 = session.create_library("lib1") + lib1.add(str(obj_file1)) + + # Create second library with second object + lib2 = session.create_library("lib2") + lib2.add(str(obj_file2)) + + # Test functions from lib1 + add_func = lib1.get_function("test_add") + assert add_func(5, 3) == 8 + + mul_func = lib1.get_function("test_multiply") + assert mul_func(4, 5) == 20 + + # Test functions from lib2 + sub_func = lib2.get_function("test_subtract") + assert sub_func(10, 3) == 7 + + div_func = lib2.get_function("test_divide") + assert div_func(20, 4) == 5 + + # Verify lib1 doesn't have lib2's functions + with pytest.raises(AttributeError, match="Module has no function"): + lib1.get_function("test_subtract") + + # Verify lib2 doesn't have lib1's functions + with pytest.raises(AttributeError, match="Module has no function"): + lib2.get_function("test_add") + + +def test_symbol_conflict_same_library() -> None: + """Test that adding objects with conflicting symbols to same library fails.""" + obj_file1 = get_test_obj_file("test_funcs.o") + obj_file_conflict = get_test_obj_file("test_funcs_conflict.o") + + session = ExecutionSession() + lib = session.create_library() + + # Add first object file + lib.add(str(obj_file1)) + + # Verify first object's function works + add_func = lib.get_function("test_add") + assert add_func(10, 20) == 30 + + # Try to add conflicting object - should raise an error + with pytest.raises(Exception): # LLVM will throw an error for duplicate symbols + lib.add(str(obj_file_conflict)) + + +def test_symbol_conflict_different_libraries() -> None: + """Test that adding objects with conflicting symbols to different libraries works.""" + obj_file1 = get_test_obj_file("test_funcs.o") + obj_file_conflict = get_test_obj_file("test_funcs_conflict.o") + + session = ExecutionSession() + + # Create first library with first object + lib1 = session.create_library("lib1") + lib1.add(str(obj_file1)) + + # Create second library with conflicting object + lib2 = session.create_library("lib2") + lib2.add(str(obj_file_conflict)) + + # Test that both libraries work with their own versions + add_func1 = lib1.get_function("test_add") + result1 = add_func1(10, 20) + assert result1 == 30 # Original implementation + + add_func2 = lib2.get_function("test_add") + result2 = add_func2(10, 20) + assert result2 == 1030 # Conflicting implementation adds 1000 + + # Test multiply functions + mul_func1 = lib1.get_function("test_multiply") + assert mul_func1(5, 6) == 30 # Original: 5 * 6 + + mul_func2 = lib2.get_function("test_multiply") + assert mul_func2(5, 6) == 60 # Conflict: (5 * 6) * 2 + + +def test_load_and_execute_cuda_function() -> None: + """Test loading an object file and executing a function.""" + # Get pre-built test object file + try: + obj_file = get_test_obj_file("test_funcs_cuda.o") + except FileNotFoundError: + return + + # Create session and library + session = ExecutionSession() + lib = session.create_library() + + # Load object file + lib.add(str(obj_file)) + + # Get and call test_add function + add_func = lib.get_function("test_add") + result = add_func(10, 20) + assert result == 30 + + # Get and call test_multiply function + mul_func = lib.get_function("test_multiply") + result = mul_func(7, 6) + assert result == 42 + + +if __name__ == "__main__": + # pytest.main([__file__, "-v"]) + test_load_and_execute_cuda_function() diff --git a/docs/conf.py b/docs/conf.py index bb7f1202..664e8827 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -465,3 +465,6 @@ def footer_html() -> str: html_css_files = ["custom.css"] + + +show_warning_types = True