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

.github/workflows/main.yml

@@ -56,13 +56,20 @@ jobs:
           CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
         run: |
           pip install scikit-build-core
-          export CMAKE_ARGS="-DUSE_LLVM=ON -DBUILD_TESTING=OFF"
+          # Link LLVM statically, as the published wheels do. Plain USE_LLVM=ON picks conda's
+          # libLLVM-23.dylib, which would sit in the process alongside the LLVM the orcjit
+          # package links statically; two LLVM copies in one process crash the JIT.
+          export CMAKE_ARGS="-DUSE_LLVM='llvm-config --link-static' -DZLIB_USE_STATIC_LIBS=ON -DBUILD_TESTING=OFF"
           pip wheel --no-deps -w dist . -v
       - name: Install TVM from wheel
         shell: bash -l {0}
         run: |
           export TVM_FFI_VERSION="$(python -c 'from importlib.metadata import version; print(version("apache-tvm-ffi"))')"
           pip install ml_dtypes
+          # --no-deps: this package requires apache-tvm-ffi>=0.1.0, and the source-built core
+          # versions as 0.1.dev1+g<sha> from the shallow checkout, which is < 0.1.0 under PEP 440.
+          # Without it pip replaces the source-matched core with a PyPI build.
+          pip install --no-deps apache-tvm-ffi-orcjit==0.1.1
           # Keep the source-matched tvm-ffi installed by the setup action.
           pip install --no-deps dist/*.whl
           python - <<'PY'
@@ -71,6 +78,7 @@ jobs:
 
           assert version("apache-tvm-ffi") == os.environ["TVM_FFI_VERSION"]
           assert "apache-tvm-ffi>=0.1.13.post2" in requires("apache-tvm")
+          assert "apache-tvm-ffi-orcjit==0.1.1" in requires("apache-tvm")
           PY
 
       - name: Test
@@ -113,6 +121,9 @@ jobs:
         shell: cmd /C call {0}
         run: |
           pip install psutil cloudpickle ml_dtypes numpy packaging scipy tornado typing_extensions
+          rem --no-deps keeps the source-built apache-tvm-ffi (versioned below 0.1.0 from the
+          rem shallow checkout) from being replaced by a PyPI build.
+          pip install --no-deps apache-tvm-ffi-orcjit==0.1.1
       - name: Test
         shell: cmd /C call {0}
         run: >-

docker/install/ubuntu_install_python_package.sh

@@ -41,3 +41,7 @@ uv pip install --upgrade \
     "tornado~=6.4" \
     "ml_dtypes~=0.5" \
     mlc-z3-static==4.16.0
+
+# Provides LLVMModule JIT execution. --no-deps keeps its apache-tvm-ffi requirement from
+# pulling a PyPI core into site-packages; CI installs the submodule-matched core into ./python.
+uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.1

docs/arch/codegen.rst

@@ -245,7 +245,12 @@ exposes it as callable ``PackedFunc``\ s.
      - How Code Is Executed
    * - ``LLVMModule``
      - LLVM IR (in-memory ``llvm::Module``)
-     - JIT-compiled on first call (MCJIT or ORC). Function pointers cached for subsequent calls.
+     - JIT-compiled on first call by the separately installed ``apache-tvm-ffi-orcjit`` package.
+       TVM emits an object in memory and transfers it to the package-backed JITDylib through FFI;
+       it has no local execution engine or fallback. Install it with
+       ``pip install apache-tvm-ffi-orcjit==0.1.1``. Because this boundary transfers an object
+       file, TVM and the package do not need to use the same LLVM version. MCJIT is no longer
+       supported.
    * - ``CUDAModule``
      - PTX or cubin binary
      - Loaded via CUDA driver API (``cuModuleLoad``). Kernels launched via ``cuLaunchKernel``.

pyproject.toml

@@ -51,6 +51,7 @@ classifiers = [
 ]
 dependencies = [
   "apache-tvm-ffi>=0.1.13.post2",
+  "apache-tvm-ffi-orcjit==0.1.1",
   "ml_dtypes",
   "numpy",
   "typing_extensions",

python/tvm/__init__.py

@@ -19,8 +19,8 @@
 """TVM: Open Deep Learning Compiler Stack."""
 
 import multiprocessing
-import sys
 import os
+import sys
 
 # ffi module must load first
 from tvm_ffi import register_object, register_global_func, get_global_func
@@ -125,3 +125,11 @@ def wrapper(exctype, value, trbk):
 from .backend._autoload_backends import _autoload_backends
 
 _autoload_backends()
+
+# Load the optional ORC JIT extension for its global FFI registrations.  LLVM
+# function lookup reports a targeted installation error if this import fails;
+# importing TVM itself remains silent for AOT-only workflows.
+try:
+    import tvm_ffi_orcjit as _tvm_ffi_orcjit
+except Exception:  # pylint: disable=broad-exception-caught
+    pass

src/target/llvm/llvm_instance.cc

@@ -262,18 +262,6 @@ LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance,
     }
   }
 
-  // LLVM JIT engine options
-  if (const auto& v =
-          target.Get("jit").value_or(nullptr).as_or_throw<ffi::Optional<ffi::String>>()) {
-    ffi::String value = v.value();
-    if ((value == "mcjit") || (value == "orcjit")) {
-      jit_engine_ = value;
-    } else {
-      TVM_FFI_THROW(InternalError)
-          << "invalid jit option " << value << " (can be `orcjit` or `mcjit`).";
-    }
-  }
-
   // TVM & LLVM vector width options
   if (const auto& w =
           target.Get("vector-width").value_or(nullptr).as_or_throw<ffi::Optional<int64_t>>()) {
@@ -585,10 +573,6 @@ std::string LLVMTargetInfo::str() const {
     obj.Set(ffi::String("cl-opt"), arr);
   }
 
-  if (jit_engine_ != "orcjit") {
-    obj.Set(ffi::String("jit"), ffi::String(jit_engine_));
-  }
-
   return std::string(ffi::json::Stringify(obj));
 }
 

src/target/llvm/llvm_instance.h

@@ -236,11 +236,6 @@ class LLVMTargetInfo {
    * \return `llvm::FastMathFlags` for this target
    */
   llvm::FastMathFlags GetFastMathFlags() const { return fast_math_flags_; }
-  /*!
-   * \brief Get the LLVM JIT engine type
-   * \return the type name of the JIT engine (default "orcjit" or "mcjit")
-   */
-  const std::string GetJITEngine() const { return jit_engine_; }
   /*!
    * \brief Get the TVM & LLVM vector_width
    * \return number of bits for vector width
@@ -366,7 +361,6 @@ class LLVMTargetInfo {
   llvm::Reloc::Model reloc_model_ = llvm::Reloc::PIC_;
   llvm::CodeModel::Model code_model_ = llvm::CodeModel::Small;
   std::shared_ptr<llvm::TargetMachine> target_machine_;
-  std::string jit_engine_ = "orcjit";
   int vector_width_{0};
 };