Skip to content

Commit 1da5637

Browse files
yaoyaodingcyx-6
yaoyaoding
and
cyx-6
authored
[OrcJIT] Add LLVM ORC JIT v2 dynamic object loader addon (#254)
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 ```python 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: ```python 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) --------- Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
1 parent dbf64e8 commit 1da5637

51 files changed

Lines changed: 5375 additions & 92 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/actions/build-orcjit-wheel/action.yml‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
name: Build OrcJIT Wheel
19+
description: >
20+
Build and test OrcJIT wheels for a given OS/architecture combination.
21+
Handles LLVM caching, conda installation, and cibuildwheel execution.
22+
23+
inputs:
24+
arch:
25+
description: "Target architecture (e.g., x86_64, aarch64, arm64, AMD64)"
26+
required: true
27+
build:
28+
description: "cibuildwheel build selector (e.g., cp312-manylinux_x86_64)"
29+
required: true
30+
checkout_ref:
31+
description: "Branch, tag, or SHA to check out before building"
32+
required: true
33+
34+
runs:
35+
using: "composite"
36+
steps:
37+
- name: Check out source
38+
uses: actions/checkout@v5
39+
with:
40+
ref: ${{ inputs.checkout_ref }}
41+
submodules: recursive
42+
43+
- uses: ./.github/actions/detect-env-vars
44+
id: env_vars
45+
46+
# ---- Cache LLVM prefix ----
47+
- name: Cache LLVM
48+
uses: actions/cache@v4
49+
id: llvm-cache
50+
with:
51+
path: ${{ runner.os == 'Windows' && 'C:/opt/llvm' || '/opt/llvm' }}
52+
key: llvm-22.1.0-${{ runner.os }}-${{ inputs.arch }}-v3
53+
54+
# ---- Install LLVM via conda (cache miss only) ----
55+
- name: Setup conda
56+
if: steps.llvm-cache.outputs.cache-hit != 'true'
57+
uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0
58+
continue-on-error: true
59+
id: conda1
60+
with:
61+
miniforge-version: latest
62+
63+
- name: Setup conda (retry with tar.bz2)
64+
if: steps.llvm-cache.outputs.cache-hit != 'true' && steps.conda1.outcome == 'failure'
65+
uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0
66+
with:
67+
miniforge-version: latest
68+
use-only-tar-bz2: true
69+
70+
- name: Create /opt/llvm (macOS)
71+
if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'macOS'
72+
shell: bash
73+
run: sudo mkdir -p /opt/llvm && sudo chown -R $(whoami) /opt/llvm
74+
75+
- name: Install LLVM (Unix)
76+
if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os != 'Windows'
77+
shell: bash -l {0}
78+
run: |
79+
conda create -q -p /opt/llvm -c conda-forge \
80+
llvmdev=22.1.0 clangdev=22.1.0 compiler-rt=22.1.0 zlib zstd-static \
81+
-y
82+
83+
- name: Install LLVM (Windows)
84+
if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'Windows'
85+
shell: cmd /C call {0}
86+
run: |
87+
conda create -q -p C:\opt\llvm -c conda-forge llvmdev=22.1.0 zlib zstd-static -y
88+
89+
# ---- Build and test wheels ----
90+
- name: Build and test wheels
91+
uses: pypa/cibuildwheel@298ed2fb2c105540f5ed055e8a6ad78d82dd3a7e # v3.3.1
92+
with:
93+
package-dir: addons/tvm_ffi_orcjit
94+
output-dir: wheelhouse
95+
env:
96+
CIBW_BUILD: ${{ inputs.build }}
97+
CIBW_ARCHS_LINUX: ${{ inputs.arch }}
98+
CIBW_ARCHS_MACOS: ${{ inputs.arch }}
99+
CIBW_ARCHS_WINDOWS: ${{ inputs.arch }}
100+
CIBW_BUILD_VERBOSITY: 1
101+
CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
102+
CIBW_ENVIRONMENT: LLVM_PREFIX=/opt/llvm
103+
CIBW_ENVIRONMENT_WINDOWS: LLVM_PREFIX="C:/opt/llvm"
104+
CIBW_CONTAINER_ENGINE: "docker; create_args: --volume /opt/llvm:/opt/llvm"
105+
CIBW_TEST_REQUIRES: pytest ninja
106+
CIBW_TEST_COMMAND: >-
107+
pip install --force-reinstall {project} &&
108+
pytest {project}/addons/tvm_ffi_orcjit/tests -v &&
109+
python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang cpp &&
110+
python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang c
111+
CIBW_TEST_COMMAND_WINDOWS: >-
112+
pip install --force-reinstall {project} &&
113+
pytest {project}/addons/tvm_ffi_orcjit/tests -v &&
114+
python {project}/addons/tvm_ffi_orcjit/examples/quick-start/run.py --lang c

‎.github/workflows/ci_test.yml‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ jobs:
3535
should_skip_ci_docs_only: ${{ steps.detect.outputs.should_skip_ci_docs_only }}
3636
cpp_changed: ${{ steps.cpp_files.outputs.changed }}
3737
cpp_files: ${{ steps.cpp_files.outputs.files }}
38+
orcjit_changed: ${{ steps.orcjit_files.outputs.changed }}
3839
steps:
3940
- uses: actions/checkout@v5
4041
with:
@@ -54,6 +55,13 @@ jobs:
5455
src/ tests/ | grep -E '\.(c|cc|cpp|cxx)$' | tr '\n' ' ')
5556
echo "files=$FILES" >> $GITHUB_OUTPUT
5657
[ -n "$FILES" ] && echo "changed=true" >> $GITHUB_OUTPUT || echo "changed=false" >> $GITHUB_OUTPUT
58+
- name: Get changed OrcJIT files
59+
id: orcjit_files
60+
run: |
61+
FILES=$(git diff --name-only --diff-filter=ACMR origin/${{ github.base_ref }}...HEAD -- \
62+
addons/tvm_ffi_orcjit/ .github/actions/build-orcjit-wheel/ | tr '\n' ' ')
63+
echo "files=$FILES" >> $GITHUB_OUTPUT
64+
[ -n "$FILES" ] && echo "changed=true" >> $GITHUB_OUTPUT || echo "changed=false" >> $GITHUB_OUTPUT
5765
5866
lint:
5967
needs: [prepare]
@@ -199,3 +207,37 @@ jobs:
199207
working-directory: rust
200208
run: |
201209
cargo test
210+
211+
orcjit:
212+
needs: [lint, prepare]
213+
if: >
214+
needs.prepare.outputs.orcjit_changed == 'true' &&
215+
needs.prepare.outputs.should_skip_ci_commit != 'true' &&
216+
needs.prepare.outputs.should_skip_ci_docs_only != 'true'
217+
name: orcjit ${{ matrix.os }} (${{ matrix.arch }})
218+
runs-on: ${{ matrix.os }}
219+
strategy:
220+
fail-fast: false
221+
matrix:
222+
include:
223+
- {os: ubuntu-latest, arch: x86_64, build: "cp312-manylinux_x86_64"}
224+
- {os: ubuntu-24.04-arm, arch: aarch64, build: "cp312-manylinux_aarch64"}
225+
- {os: macos-14, arch: arm64, build: "cp312-macosx_arm64"}
226+
- {os: windows-latest, arch: AMD64, build: "cp312-win_amd64"}
227+
228+
steps:
229+
- uses: actions/checkout@v5
230+
with:
231+
fetch-depth: 1
232+
233+
- name: Build OrcJIT wheel
234+
uses: ./.github/actions/build-orcjit-wheel
235+
with:
236+
arch: ${{ matrix.arch }}
237+
build: ${{ matrix.build }}
238+
checkout_ref: ${{ github.sha }}
239+
240+
- uses: actions/upload-artifact@v4
241+
with:
242+
name: cibw-orcjit-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }}
243+
path: wheelhouse/*.whl

‎.github/workflows/publish_orcjit_wheel.yml‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
name: Publish OrcJIT wheel
19+
20+
on:
21+
workflow_dispatch:
22+
inputs:
23+
tag:
24+
description: "Tag to publish (manual run)"
25+
required: true
26+
27+
jobs:
28+
build_wheels:
29+
name: ${{ matrix.os }} (${{ matrix.arch }})
30+
runs-on: ${{ matrix.os }}
31+
strategy:
32+
fail-fast: false
33+
matrix:
34+
include:
35+
- {os: ubuntu-latest, arch: x86_64, build: "cp312-manylinux_x86_64"}
36+
- {os: ubuntu-24.04-arm, arch: aarch64, build: "cp312-manylinux_aarch64"}
37+
- {os: macos-14, arch: arm64, build: "cp312-macosx_arm64"}
38+
- {os: windows-latest, arch: AMD64, build: "cp312-win_amd64"}
39+
40+
steps:
41+
- name: Checkout repository (for local composite action)
42+
uses: actions/checkout@v5
43+
with:
44+
fetch-depth: 1
45+
46+
- name: Build OrcJIT wheel
47+
uses: ./.github/actions/build-orcjit-wheel
48+
with:
49+
arch: ${{ matrix.arch }}
50+
build: ${{ matrix.build }}
51+
checkout_ref: ${{ inputs.tag }}
52+
53+
- name: Upload wheels
54+
uses: actions/upload-artifact@v4
55+
with:
56+
name: cibw-orcjit-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }}
57+
path: ./wheelhouse/*.whl
58+
59+
upload_pypi:
60+
needs: [build_wheels]
61+
runs-on: ubuntu-latest
62+
environment: pypi
63+
permissions:
64+
id-token: write
65+
attestations: write
66+
if: github.event_name == 'workflow_dispatch'
67+
steps:
68+
- uses: actions/download-artifact@v4
69+
with:
70+
pattern: cibw-orcjit-*
71+
path: dist
72+
merge-multiple: true
73+
74+
- name: Generate artifact attestation for wheels
75+
uses: actions/attest-build-provenance@v1
76+
with:
77+
subject-path: dist/*
78+
79+
- name: Publish package distributions to PyPI
80+
uses: pypa/gh-action-pypi-publish@release/v1
81+
with:
82+
attestations: true
83+
verbose: true

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,8 @@ celerybeat-schedule
122122
# Allow composite actions under .github/actions despite build-* ignore pattern
123123
!.github/actions/build-wheel-for-publish/
124124
!.github/actions/build-wheel-for-publish/action.yml
125+
!.github/actions/build-orcjit-wheel/
126+
!.github/actions/build-orcjit-wheel/action.yml
125127
celerybeat.pid
126128

127129
# SageMath parsed files

0 commit comments

Comments
 (0)