Skip to content

[LLVM] Use external tvm-ffi ORC JIT for LLVMModule execution - #20292

Open
cyx-6 wants to merge 7 commits into
apache:mainfrom
cyx-6:llvmmodule-jit-contract
Open

cyx-6 wants to merge 7 commits into
apache:mainfrom
cyx-6:llvmmodule-jit-contract

Conversation

@cyx-6

@cyx-6 cyx-6 commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

This makes LLVMModule execution go through the separately installed
apache-tvm-ffi-orcjit package instead of TVM hosting its own execution engine.

Original work by @tqchen (first three commits); I rebased onto main and added the CI
installation commit.

What changes

LLVMModuleNode no longer owns an llvm::ExecutionEngine / llvm::orc::LLJIT. On first
function lookup it emits the retained module as an object file in memory and transfers
the bytes to the package through the global FFI functions
tvm_ffi_orcjit.GlobalDefaultSession / tvm_ffi_orcjit.SessionLoadModule, then forwards
GetFunction to the returned dylib. Because the boundary transfers an object file, TVM and
the package do not need to be built against the same LLVM version.

This removes ~340 lines from llvm_module.cc, including all the
#if TVM_LLVM_VERSION >= 210/230 LLJIT and ObjectLinkingLayer shims. MCJIT is no longer
supported, so the now-dead jit target attribute is removed from target_kind.cc and
LLVMTargetInfo, along with its tests.

There is no local fallback: if the package is missing, function lookup raises a targeted
install error. Importing tvm stays silent so AOT-only workflows are unaffected.

CI

The package is a new runtime dependency, so CI has to provide it:

  • Linux (Jenkins): installed in docker/install/ubuntu_install_python_package.sh so the
    ~20MB wheel is baked into the CI image rather than downloaded on every run. This
    docker/ change trips git_change_docker.sh, so the images rebuild for this PR.
  • macOS / Windows (GitHub Actions): both jobs install TVM with --no-deps and would
    otherwise silently skip the new dependency, so it is installed explicitly.

Pinned to ==0.1.1 in pyproject.toml and everywhere CI installs it, so the version moves
deliberately rather than drifting.

Verification

  • apache-tvm-ffi-orcjit 0.1.1 publishes prebuilt wheels for all four CI platforms
    (manylinux x86_64/aarch64, macOS arm64, win_amd64), so no CI leg has to build LLVM.
  • Confirmed the published wheel registers both global functions the C++ contract looks up.
  • Ran the addon's quick-start example end to end (compile object -> JIT load -> call) against
    a tvm_ffi core built from the submodule pin, matching how CI installs the core. Also
    checked the case where the PyPI core is additionally present in site-packages: the
    PYTHONPATH copy correctly shadows it and JIT execution still passes.
  • Verified the built metadata emits Requires-Dist: apache-tvm-ffi-orcjit==0.1.1 verbatim,
    so the macOS job's requires() assertion matches.

Note: my local box has USE_LLVM=OFF, so the LLVM codegen tests themselves were not run
locally — that coverage comes from CI (test_minimal_target_codegen_llvm.py on macOS/Windows
and the full Linux suite).

cc @tqchen

🤖 Generated with Claude Code

tqchen and others added 6 commits September 9, 2026 00:18
LLVMModule execution now goes through the separately installed
apache-tvm-ffi-orcjit package, so CI has to provide it:

- Bake it into the Linux CI docker image alongside the other python
  packages, so the ~20MB wheel is cached in the image rather than
  downloaded on every run.
- Install it explicitly in the macOS and Windows jobs, which install
  TVM with --no-deps and would otherwise skip the new dependency.

Pin to ==0.1.1 (the current release) in pyproject.toml and everywhere
CI installs it, so the version is bumped deliberately.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The macOS job replaced the source-built apache-tvm-ffi with a PyPI build:
the shallow CI checkout makes the source core version as 0.1.dev1+g<sha>,
which is below 0.1.0 under PEP 440, so it does not satisfy the orcjit
package's apache-tvm-ffi>=0.1.0 and pip resolves it from PyPI instead.

Install the package with --no-deps everywhere so the source-matched core
is always the one that stays installed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A docker/ change rebuilds the CI images but does not feed them back into
the same run: the test stages pull the pinned tag from docker-images.ini,
and the Jenkinsfile deliberately leaves the built image unassigned. So the
image-level install alone leaves LLVM JIT unavailable in the Linux legs,
which failed with the missing-package error from EnsureOrcJITModule.

Install the package in the test scripts as well, next to where they set up
tvm-ffi. The image install stays so the wheel is cached once the rebuilt
image is promoted.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@cyx-6

cyx-6 commented Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Status after taking this over and rebasing onto main.

Working

Windows passes test_llvm_add_pipeline through the new external JIT path — that is the
end-to-end confirmation that the object-file contract works: TVM emits the object, hands it to
the package's JITDylib over FFI, and the compiled function runs correctly.

I also verified locally (aarch64 Linux) that the published 0.1.1 wheel registers both global
functions EnsureOrcJITModule looks up, and ran the addon's quick-start example end to end
against a tvm_ffi core built from this PR's submodule pin.

CI wiring (fixed here)

Two things the CI needed that were not obvious:

  1. macOS/Windows install TVM with --no-deps, so the new dependency was skipped entirely.
    Installing it plainly then broke macOS a different way: the shallow CI checkout makes the
    source-built core version as 0.1.dev1+g<sha>, which is below 0.1.0 under PEP 440, so it
    does not satisfy the package's apache-tvm-ffi>=0.1.0 and pip silently replaced the
    source-matched core with a PyPI build. Installing with --no-deps fixes both.

  2. A docker/ change does not reach the test containers in the same run. The image-level
    install alone left every Linux leg failing with the missing-package error from
    EnsureOrcJITModule, because the test stages pull the pinned tag from docker-images.ini and
    the Jenkinsfile deliberately leaves the newly built image unassigned
    (// ci_cpu = build_image('ci_cpu') is commented out). So the package is now also installed in
    task_python_unittest.sh / task_python_docs.sh, next to where they set up tvm-ffi. The
    image install stays so the wheel is cached once a rebuilt image is promoted.

Pinned to ==0.1.1 everywhere so the version moves deliberately.

Open issue: macOS segfaults in the JIT

macOS gets past install and build, then segfaults (exit 139) inside
test_llvm_add_pipeline
— the same test Windows passes:

collecting ... collected 3 items
Segmentation fault: 11  python -m pytest -v tests/python/all-platform-minimal-test
tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py::test_llvm_add_pipeline

This looks like a genuine platform issue in the JIT path rather than CI wiring, and macOS is the
platform the addon treats specially: it skips ExecutorNativePlatform/MachOPlatform to avoid
the compact-unwind 32-bit-delta bug in JITLink's CompactUnwindSupport, and relies on
InitFiniPlugin plus a ___cxa_atexit shim instead. Note the failing kernel is the parallel one
(sch.parallel(xo)), so it goes through TVMBackendParallelLaunch and the context-symbol
injection that EnsureOrcJITModule no longer does itself (it now depends on the package's
InitContextSymbols).

I have not root-caused it further — that likely needs a debug build with a stack trace on a real
arm64 macOS box, which I do not have (my box is aarch64 Linux and USE_LLVM=OFF).

@tqchen how would you like to handle the macOS leg? Options as I see them: dig into the addon's
MachO path and fix it there (probably a tvm-ffi side change + a 0.1.2), or land the contract
with macOS JIT execution known-broken and tracked. Happy to take either.

macOS was segfaulting inside test_llvm_add_pipeline. It is the only leg
where TVM links LLVM as a shared library: conda's libLLVM-23.dylib, per
'Found LLVM_LIBS=.../libLLVM-23.dylib' in the build log. The orcjit
package statically links its own LLVM, so the process ended up with two
ODR-incompatible LLVM copies. The package hides its LLVM symbols with
--exclude-libs, but that guard is ELF-only, so nothing prevents the
interposition on Mach-O.

Windows already links LLVM statically and passes the same test, and the
published wheels use --link-static on every platform, so this aligns the
macOS job with the policy the shipped artifacts already follow.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@cyx-6

cyx-6 commented Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Root-caused the macOS segfault: it is the only leg where TVM links LLVM as a shared library.

From the macOS build log:

-- Found LLVM_LIBS=/Users/runner/miniconda3/envs/tvm-build/lib/libLLVM-23.dylib
-- Found TVM_LLVM_VERSION=231

versus Windows, which passes the same test:

-- Found LLVM_LIBS=...\lib\LLVMWindowsManifest.lib;...\lib\LLVMXRay.lib;... (static archives)

The orcjit wheel statically links its own LLVM (22.1.0), so on macOS the process ends up with two
ODR-incompatible LLVM copies — conda's libLLVM-23.dylib pulled in by TVM, and the package's
static LLVM. The package does try to prevent exactly this (its CMake comment says the symbols are
hidden so they "cannot interpose with a host process's own copies (e.g. PyTorch's bundled LLVM)"),
but that --exclude-libs guard is Linux|Android|*BSD only, so nothing suppresses the
interposition on Mach-O.

Fixed by linking LLVM statically in the macOS job, which is what Windows already does and what
every published wheel already doesbuild-wheel-for-publish passes
-DUSE_LLVM='...llvm-config --link-static' on all three platforms. So this was a gap between the
macOS CI job and the shipped-artifact policy, not something the shipped macOS wheel suffers from.
The conda build env already provides zlib/zstd-static that static linking needs.

Pushed; macOS is re-running. All other legs are green on the previous commit (Linux CPU/GPU/ARM/
wasm/docker/lint and Windows).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants