This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 076e83f62f [LLVM] Use external tvm-ffi ORC JIT for LLVMModule
execution (#20292)
076e83f62f is described below
commit 076e83f62f453d5be64d74a296507d68d568dbc0
Author: Yaxing Cai <[email protected]>
AuthorDate: Mon Sep 28 09:39:27 2026 +0800
[LLVM] Use external tvm-ffi ORC JIT for LLVMModule execution (#20292)
This makes `LLVMModule` execution go through the separately installed
`apache-tvm-ffi-orcjit` package instead of TVM hosting its own execution
engine.
### 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.2` in `pyproject.toml` and everywhere CI installs it.
The tvm-ffi
submodule and Python requirement move to v0.1.14-post1, which provides
the process-lifetime
module registry required by the new addon loader. macOS CI returns to
its normal
`USE_LLVM=ON` configuration: 0.1.2 keeps its LLVM symbols local, so the
addon can coexist
with the LLVM dylib used by TVM.
### Verification
- `apache-tvm-ffi-orcjit` 0.1.2 publishes prebuilt wheels for all four
CI platforms
(manylinux x86_64/aarch64, macOS arm64, win_amd64), so no CI leg has to
build LLVM.
- On macOS arm64, installed the published 0.1.2 addon and 0.1.14.post1
core wheels in an
isolated environment. The addon's C and C++ quick-start examples both
passed.
- Loaded the compiled C and C++ object bytes through the same
five-argument global FFI call
used by `LLVMModuleNode`, then called `add(10, 20)` successfully in both
modules.
- Confirmed the macOS wheel exports only `TVMFFIOrcJITInitialize` and
compiled the new C++
FFI call against the pinned tvm-ffi headers. The macOS CI job checks the
built TVM wheel's
`Requires-Dist` pins.
The full TVM LLVM codegen test was not run locally because this checkout
has no TVM LLVM
build. CI runs `test_minimal_target_codegen_llvm.py` on macOS and
Windows and the full Linux
suite.
Co-authored-by: Tianqi Chen <[email protected]>
---
.github/workflows/main.yml | 10 +-
3rdparty/tvm-ffi | 2 +-
docker/install/ubuntu_install_python_package.sh | 4 +
docs/arch/codegen.rst | 7 +-
pyproject.toml | 3 +-
python/tvm/__init__.py | 10 +-
src/target/llvm/llvm_instance.cc | 16 -
src/target/llvm/llvm_instance.h | 6 -
src/target/llvm/llvm_module.cc | 391 ++++-------------------
src/target/target_kind.cc | 2 -
tests/python/runtime/test_runtime_module_load.py | 5 +-
tests/python/target/test_target_target.py | 7 -
tests/scripts/task_python_docs.sh | 3 +
tests/scripts/task_python_unittest.sh | 5 +
14 files changed, 105 insertions(+), 366 deletions(-)
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2827ed917a..7a524732cf 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -63,6 +63,10 @@ jobs:
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.14.post1, 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.2
# Keep the source-matched tvm-ffi installed by the setup action.
pip install --no-deps dist/*.whl
python - <<'PY'
@@ -70,7 +74,8 @@ jobs:
from importlib.metadata import requires, version
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>=0.1.14.post1" in requires("apache-tvm")
+ assert "apache-tvm-ffi-orcjit==0.1.2" in requires("apache-tvm")
PY
- name: Test
@@ -113,6 +118,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.14.post1 from the
+ rem shallow checkout) from being replaced by a PyPI build.
+ pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
- name: Test
shell: cmd /C call {0}
run: >-
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index daf594da8e..e251c6abee 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit daf594da8e6950fdc02ae35fcf4e0ab2e59a4979
+Subproject commit e251c6abee7a111c6467f84d854be512e33e2db0
diff --git a/docker/install/ubuntu_install_python_package.sh
b/docker/install/ubuntu_install_python_package.sh
index 8be4a6a8d7..484ceabf1e 100755
--- a/docker/install/ubuntu_install_python_package.sh
+++ b/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.2
diff --git a/docs/arch/codegen.rst b/docs/arch/codegen.rst
index d14e4b8a07..799568b420 100644
--- a/docs/arch/codegen.rst
+++ b/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.2``. 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``.
diff --git a/pyproject.toml b/pyproject.toml
index 24d3594c37..e2fcc2918b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,7 +50,8 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
- "apache-tvm-ffi>=0.1.13.post2",
+ "apache-tvm-ffi>=0.1.14.post1",
+ "apache-tvm-ffi-orcjit==0.1.2",
"ml_dtypes",
"numpy",
"typing_extensions",
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index 77e30a9e2a..388e2180f3 100644
--- a/python/tvm/__init__.py
+++ b/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 @@ sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
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
diff --git a/src/target/llvm/llvm_instance.cc b/src/target/llvm/llvm_instance.cc
index 0b1275498c..a17e92ce8c 100644
--- a/src/target/llvm/llvm_instance.cc
+++ b/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));
}
diff --git a/src/target/llvm/llvm_instance.h b/src/target/llvm/llvm_instance.h
index 0cb141de99..5d797d0141 100644
--- a/src/target/llvm/llvm_instance.h
+++ b/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};
};
diff --git a/src/target/llvm/llvm_module.cc b/src/target/llvm/llvm_module.cc
index a2226c5dec..952638c38e 100644
--- a/src/target/llvm/llvm_module.cc
+++ b/src/target/llvm/llvm_module.cc
@@ -26,16 +26,6 @@
#include <llvm/ADT/SmallString.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/Bitcode/BitcodeWriter.h>
-#include <llvm/ExecutionEngine/ExecutionEngine.h>
-#include <llvm/ExecutionEngine/MCJIT.h>
-#include <llvm/ExecutionEngine/Orc/LLJIT.h>
-#include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/support/io.h>
-#if _WIN32
-#include <llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h>
-#include <llvm/ExecutionEngine/SectionMemoryManager.h>
-#endif
#include <llvm/IR/DataLayout.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Intrinsics.h>
@@ -46,6 +36,8 @@
#include <llvm/IR/Module.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/Support/FileSystem.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/support/io.h>
#if TVM_LLVM_VERSION >= 180
#include <llvm/TargetParser/Host.h>
#else
@@ -54,10 +46,10 @@
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Target/TargetMachine.h>
-#include <llvm/Target/TargetOptions.h>
#include <llvm/Transforms/Utils/Cloning.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/string.h>
@@ -67,14 +59,10 @@
#include <tvm/target/codegen.h>
#include <tvm/target/target.h>
-#include <algorithm>
#include <memory>
-#include <mutex>
-#include <sstream>
#include <string>
#include <system_error>
#include <utility>
-#include <vector>
#include "../../runtime/file_utils.h"
#include "codegen_blob.h"
@@ -91,12 +79,12 @@ using ffi::PackedArgs;
class LLVMModuleNode final : public ffi::ModuleObj {
public:
- ~LLVMModuleNode();
-
const char* kind() const final { return "llvm"; }
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final;
+ void ImportModule(const ffi::Module& other) final;
+
/*! \brief Get the property of the runtime module .*/
// TODO(tvm-team): Make it serializable
int GetPropertyMask() const override {
@@ -111,50 +99,22 @@ class LLVMModuleNode final : public ffi::ModuleObj {
void Init(std::unique_ptr<llvm::Module> module,
std::unique_ptr<LLVMInstance> llvm_instance);
void LoadIR(const std::string& file_name);
- bool ImplementsFunction(const ffi::String& name) final;
-
- void SetJITEngine(const std::string& jit_engine) { jit_engine_ = jit_engine;
}
-
private:
- void InitMCJIT();
- void InitORCJIT();
- bool IsCompatibleWithHost(const llvm::TargetMachine* tm) const;
- void* GetGlobalAddr(const std::string& name, const LLVMTarget& llvm_target)
const;
- void* GetFunctionAddr(const std::string& name, const LLVMTarget&
llvm_target) const;
+ void EnsureOrcJITModule();
// The LLVM scope object.
std::unique_ptr<LLVMInstance> llvm_instance_;
- // JIT lock
- std::mutex mutex_;
- // jit execution engines
- llvm::ExecutionEngine* mcjit_ee_{nullptr};
- std::unique_ptr<llvm::orc::LLJIT> orcjit_ee_{nullptr};
- // The raw pointer to the module.
- llvm::Module* module_{nullptr};
- // The unique_ptr owning the module. This becomes empty once JIT has been
initialized
- // (EngineBuilder takes ownership of the module).
- std::unique_ptr<llvm::Module> module_owning_ptr_;
+ // JIT execution is provided by the separately installed
apache-tvm-ffi-orcjit package. TVM
+ // only emits the retained LLVM module as object-file bytes and transfers
them to the package.
+ // Dropping this reference releases the package JITDylib; a returned package
Function retains
+ // that dylib independently when the function outlives this LLVMModule.
+ ffi::Optional<ffi::Module> tvm_ffi_orcjit_module_;
+ // The retained module IR used for inspection and object emission.
+ std::unique_ptr<llvm::Module> module_;
/* \brief names of the external functions declared in this module */
ffi::Array<ffi::String> function_names_;
- std::string jit_engine_;
};
-LLVMModuleNode::~LLVMModuleNode() {
- if (mcjit_ee_ != nullptr) {
- mcjit_ee_->runStaticConstructorsDestructors(true);
- delete mcjit_ee_;
- }
- if (orcjit_ee_ != nullptr) {
- auto dtors = llvm::orc::getDestructors(*module_);
- auto dtorRunner =
std::make_unique<llvm::orc::CtorDtorRunner>(orcjit_ee_->getMainJITDylib());
- dtorRunner->add(dtors);
- auto err = dtorRunner->run();
- TVM_FFI_ICHECK(!err) << llvm::toString(std::move(err));
- orcjit_ee_.reset();
- }
- module_owning_ptr_.reset();
-}
-
ffi::Optional<ffi::Function> LLVMModuleNode::GetFunction(const ffi::String&
name) {
ffi::ObjectPtr<ffi::Object> sptr_to_self =
ffi::GetObjectPtr<ffi::Object>(this);
if (name == "__tvm_is_system_module") {
@@ -181,29 +141,23 @@ ffi::Optional<ffi::Function>
LLVMModuleNode::GetFunction(const ffi::String& name
return ffi::Function(
[target_string](ffi::PackedArgs args, ffi::Any* rv) { *rv =
target_string; });
}
- TVM_FFI_ICHECK(jit_engine_.size()) << "JIT engine type is missing";
- if ((jit_engine_ == "mcjit") && (mcjit_ee_ == nullptr)) InitMCJIT();
- if ((jit_engine_ == "orcjit") && (orcjit_ee_ == nullptr)) InitORCJIT();
-
- std::lock_guard<std::mutex> lock(mutex_);
+ EnsureOrcJITModule();
+ return tvm_ffi_orcjit_module_.value()->GetFunction(name);
+}
- TVMFFISafeCallType faddr;
- With<LLVMTarget> llvm_target(*llvm_instance_,
LLVMTarget::GetTargetMetadata(*module_));
- ffi::String name_with_prefix = ffi::symbol::tvm_ffi_symbol_prefix + name;
- faddr =
reinterpret_cast<TVMFFISafeCallType>(GetFunctionAddr(name_with_prefix,
*llvm_target));
- if (faddr == nullptr) return std::nullopt;
- ffi::Module self_strong_ref = ffi::GetRef<ffi::Module>(this);
- return ffi::Function::FromPacked([faddr, self_strong_ref](ffi::PackedArgs
args, ffi::Any* rv) {
- TVM_FFI_ICHECK_LT(rv->type_index(),
ffi::TypeIndex::kTVMFFIStaticObjectBegin);
- TVM_FFI_CHECK_SAFE_CALL((*faddr)(nullptr, reinterpret_cast<const
TVMFFIAny*>(args.data()),
- args.size(),
reinterpret_cast<TVMFFIAny*>(rv)));
- });
+void LLVMModuleNode::ImportModule(const ffi::Module& other) {
+ ffi::ModuleObj::ImportModule(other);
+ if (tvm_ffi_orcjit_module_.has_value()) {
+ tvm_ffi_orcjit_module_.value()->ImportModule(other);
+ }
}
namespace {
constexpr auto llvm_open_output_flag = llvm::sys::fs::OF_None;
-std::unique_ptr<llvm::Module> CloneLLVMModule(llvm::Module* mod) { return
llvm::CloneModule(*mod); }
+std::unique_ptr<llvm::Module> CloneLLVMModule(const llvm::Module* mod) {
+ return llvm::CloneModule(*mod);
+}
#if TVM_LLVM_VERSION <= 170
constexpr auto llvm_object_file_target = llvm::CGFT_ObjectFile;
@@ -214,7 +168,7 @@ constexpr auto llvm_assembly_file_target =
llvm::CodeGenFileType::AssemblyFile;
#endif
bool LLVMAddPassesToEmitFile(llvm::TargetMachine* tm,
llvm::legacy::PassManager* pm,
- llvm::raw_fd_ostream* dest,
+ llvm::raw_pwrite_stream* dest,
decltype(llvm_object_file_target)
llvm_file_target) {
return tm->addPassesToEmitFile(*pm, *dest, nullptr, llvm_file_target);
}
@@ -242,7 +196,7 @@ void LLVMModuleNode::WriteToFile(const ffi::String&
file_name_str,
auto err = LLVMAddPassesToEmitFile(tm, &pass, &dest, llvm_file_target);
TVM_FFI_ICHECK(!err) << "Cannot emit target CGFT_ObjectFile";
- pass.run(*CloneLLVMModule(module_));
+ pass.run(*CloneLLVMModule(module_.get()));
} else if (fmt == "ll") {
module_->print(dest, nullptr);
} else if (fmt == "bc") {
@@ -332,10 +286,8 @@ void LLVMModuleNode::Init(const IRModule& mod, const
Target& target) {
cg->AddMainFunction(entry_func);
}
- module_owning_ptr_ = cg->Finish();
- module_ = module_owning_ptr_.get();
- jit_engine_ = llvm_target->GetJITEngine();
- llvm_target->SetTargetMetadata(module_);
+ module_ = cg->Finish();
+ llvm_target->SetTargetMetadata(module_.get());
module_->addModuleFlag(llvm::Module::Override, "Debug Info Version",
llvm::DEBUG_METADATA_VERSION);
@@ -351,8 +303,7 @@ void LLVMModuleNode::Init(const IRModule& mod, const
Target& target) {
void LLVMModuleNode::Init(std::unique_ptr<llvm::Module> module,
std::unique_ptr<LLVMInstance> llvm_instance) {
- module_owning_ptr_ = std::move(module);
- module_ = module_owning_ptr_.get();
+ module_ = std::move(module);
llvm_instance_ = std::move(llvm_instance);
}
@@ -362,264 +313,53 @@ void LLVMModuleNode::LoadIR(const std::string&
file_name) {
Init(std::move(module), std::move(llvm_instance));
}
-bool LLVMModuleNode::ImplementsFunction(const ffi::String& name) {
- return std::find(function_names_.begin(), function_names_.end(),
- ffi::symbol::tvm_ffi_symbol_prefix + name) !=
function_names_.end();
-}
-
-void LLVMModuleNode::InitMCJIT() {
- std::lock_guard<std::mutex> lock(mutex_);
- if (mcjit_ee_) {
+void LLVMModuleNode::EnsureOrcJITModule() {
+ if (tvm_ffi_orcjit_module_.has_value()) {
return;
}
- // MCJIT builder
- With<LLVMTarget> llvm_target(*llvm_instance_,
LLVMTarget::GetTargetMetadata(*module_));
- llvm::EngineBuilder builder(std::move(module_owning_ptr_));
+ ffi::Optional<ffi::Function> get_default_session =
+ ffi::Function::GetGlobal("tvm_ffi_orcjit.GlobalDefaultSession");
+ ffi::Optional<ffi::Function> load_module =
+ ffi::Function::GetGlobal("tvm_ffi_orcjit.SessionLoadModule");
+ TVM_FFI_CHECK(get_default_session.has_value() && load_module.has_value(),
InternalError)
+ << "LLVMModule execution requires the separately installed
apache-tvm-ffi-orcjit "
+ "package and its global execution-session functions. "
+ "Install it with `pip install apache-tvm-ffi-orcjit==0.1.2`.";
- // set options
- builder.setEngineKind(llvm::EngineKind::JIT);
-#if TVM_LLVM_VERSION <= 170
- builder.setOptLevel(llvm::CodeGenOpt::Aggressive);
-#else
- builder.setOptLevel(llvm::CodeGenOptLevel::Aggressive);
-#endif
- builder.setMCPU(llvm_target->GetCPU());
- builder.setMAttrs(llvm_target->GetTargetFeatures());
- builder.setTargetOptions(llvm_target->GetTargetOptions());
-
- // create the taget machine
- auto tm = std::unique_ptr<llvm::TargetMachine>(builder.selectTarget());
- if (!IsCompatibleWithHost(tm.get())) {
- TVM_FFI_THROW(InternalError) << "Cannot run module, architecture mismatch";
- }
-
- // data layout
- llvm::DataLayout layout(tm->createDataLayout());
- TVM_FFI_ICHECK(layout == module_->getDataLayout())
- << "Data layout mismatch between module("
- << module_->getDataLayout().getStringRepresentation() << ")"
- << " and ExecutionEngine (" << layout.getStringRepresentation() << ")";
-
- // create MCJIT
- mcjit_ee_ = builder.create(tm.release());
- TVM_FFI_ICHECK(mcjit_ee_ != nullptr) << "Failed to initialize LLVM MCJIT
engine for "
-#if TVM_LLVM_VERSION >= 210
- << module_->getTargetTriple().str();
-#else
- << module_->getTargetTriple();
-#endif
-
- VLOG(2) << "LLVM MCJIT execute " << module_->getModuleIdentifier() << " for
triple `"
- << llvm_target->GetTargetTriple() << "`"
- << " on cpu `" << llvm_target->GetCPU() << "`";
-
- // run ctors
- mcjit_ee_->runStaticConstructorsDestructors(false);
-
- if (void** ctx_addr =
-
reinterpret_cast<void**>(GetGlobalAddr(ffi::symbol::tvm_ffi_library_ctx,
*llvm_target))) {
- *ctx_addr = this;
- }
-
- ffi::Module::VisitContextSymbols([this, &llvm_target](const ffi::String&
name, void* symbol) {
- if (void** ctx_addr = reinterpret_cast<void**>(GetGlobalAddr(name,
*llvm_target))) {
- *ctx_addr = symbol;
- }
- });
- // There is a problem when a JITed function contains a call to a runtime
function.
- // The runtime function (e.g. __truncsfhf2) may not be resolved, and calling
it will
- // lead to a runtime crash.
- // Do name lookup on a symbol that doesn't exist. This will force MCJIT to
finalize
- // all loaded objects, which will resolve symbols in JITed code.
- mcjit_ee_->getFunctionAddress(
- "__some_name_that_hopefully_doesnt_exist__b49f8aaade5877eaba7583b91");
-}
-
-void LLVMModuleNode::InitORCJIT() {
- std::lock_guard<std::mutex> lock(mutex_);
- if (orcjit_ee_) {
- return;
- }
- // ORCJIT builder
With<LLVMTarget> llvm_target(*llvm_instance_,
LLVMTarget::GetTargetMetadata(*module_));
- llvm::orc::JITTargetMachineBuilder
tm_builder(llvm::Triple(llvm_target->GetTargetTriple()));
-
- // set options
- tm_builder.setCPU(llvm_target->GetCPU());
- tm_builder.setFeatures(llvm_target->GetTargetFeatureString());
- tm_builder.setOptions(llvm_target->GetTargetOptions());
-#if TVM_LLVM_VERSION <= 170
- tm_builder.setCodeGenOptLevel(llvm::CodeGenOpt::Aggressive);
-#else
- tm_builder.setCodeGenOptLevel(llvm::CodeGenOptLevel::Aggressive);
-#endif
-
- // Default is no explicit JIT code & reloc model
- // Propagate instance code & reloc for RISCV case.
- auto arch = tm_builder.getTargetTriple().getArch();
- if (arch == llvm::Triple::riscv32 || arch == llvm::Triple::riscv64) {
- tm_builder.setRelocationModel(llvm_target->GetTargetRelocModel());
- tm_builder.setCodeModel(llvm_target->GetTargetCodeModel());
- }
-
- // create the taget machine
- std::unique_ptr<llvm::TargetMachine> tm =
llvm::cantFail(tm_builder.createTargetMachine());
- if (!IsCompatibleWithHost(tm.get())) {
- TVM_FFI_THROW(InternalError) << "Cannot run module, architecture mismatch";
- }
+ llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
- // data layout
- ffi::String module_name = module_->getModuleIdentifier();
llvm::DataLayout layout(tm->createDataLayout());
TVM_FFI_ICHECK(layout == module_->getDataLayout())
<< "Data layout mismatch between module("
<< module_->getDataLayout().getStringRepresentation() << ")"
- << " and ExecutionEngine (" << layout.getStringRepresentation() << ")";
-
- // compiler
- const auto compilerBuilder = [&](const llvm::orc::JITTargetMachineBuilder&)
- ->
llvm::Expected<std::unique_ptr<llvm::orc::IRCompileLayer::IRCompiler>> {
- return std::make_unique<llvm::orc::TMOwningSimpleCompiler>(std::move(tm));
- };
-
- // linker
- const auto linkerBuilder =
-#if TVM_LLVM_VERSION >= 230
- [&](llvm::orc::ExecutionSession& session,
llvm::jitlink::JITLinkMemoryManager& mem_mgr)
- -> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
-#elif TVM_LLVM_VERSION >= 210
- [&](llvm::orc::ExecutionSession& session)
- -> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
-#else
- [&](llvm::orc::ExecutionSession& session,
- const llvm::Triple& triple) ->
std::unique_ptr<llvm::orc::ObjectLayer> {
-#endif
-#if _WIN32
-#if TVM_LLVM_VERSION >= 210
- auto GetMemMgr = [](const llvm::MemoryBuffer&) {
- return std::make_unique<llvm::SectionMemoryManager>();
- };
-#else
- auto GetMemMgr = []() { return
std::make_unique<llvm::SectionMemoryManager>(); };
-#endif
- auto ObjLinkingLayer =
- std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(session,
std::move(GetMemMgr));
-#else
-#if TVM_LLVM_VERSION >= 230
- auto ObjLinkingLayer =
std::make_unique<llvm::orc::ObjectLinkingLayer>(session, mem_mgr);
-#else
- auto ObjLinkingLayer =
std::make_unique<llvm::orc::ObjectLinkingLayer>(session);
-#endif
-#endif
-#if TVM_LLVM_VERSION >= 210
- if (tm_builder.getTargetTriple().isOSBinFormatCOFF()) {
-#else
- if (triple.isOSBinFormatCOFF()) {
-#endif
- ObjLinkingLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
- ObjLinkingLayer->setAutoClaimResponsibilityForObjectSymbols(true);
- }
-#if TVM_LLVM_VERSION >= 210
- return
llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>>(std::move(ObjLinkingLayer));
-#else
- return ObjLinkingLayer;
-#endif
- }; // NOLINT(readability/braces)
-
- // create LLJIT
- orcjit_ee_ = llvm::cantFail(llvm::orc::LLJITBuilder()
- .setDataLayout(layout)
- .setCompileFunctionCreator(compilerBuilder)
- .setObjectLinkingLayerCreator(linkerBuilder)
- .create());
-
- TVM_FFI_ICHECK(orcjit_ee_ != nullptr) << "Failed to initialize LLVM ORCJIT
engine for "
-#if TVM_LLVM_VERSION >= 210
- << module_->getTargetTriple().str();
-#else
- << module_->getTargetTriple();
-#endif
-
- // store ctors
- auto ctors = llvm::orc::getConstructors(*module_);
- llvm::orc::CtorDtorRunner ctorRunner(orcjit_ee_->getMainJITDylib());
- ctorRunner.add(ctors);
-
- // resolve system symbols (like pthread, dl, m, etc.)
- auto gen =
-
llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(layout.getGlobalPrefix());
- TVM_FFI_ICHECK(gen) << llvm::toString(gen.takeError()) << "\n";
- orcjit_ee_->getMainJITDylib().addGenerator(std::move(gen.get()));
-
- // transfer module to a clone
- auto uctx = std::make_unique<llvm::LLVMContext>();
- auto umod = llvm::CloneModule(*(std::move(module_owning_ptr_)));
-
- // add the llvm module to run
- llvm::orc::ThreadSafeModule tsm(std::move(umod), std::move(uctx));
- auto err = orcjit_ee_->addIRModule(std::move(tsm));
- TVM_FFI_ICHECK(!err) << llvm::toString(std::move(err));
-
- VLOG(2) << "LLVM ORCJIT execute " << module_->getModuleIdentifier() << " for
triple `"
- << llvm_target->GetTargetTriple() << "`"
- << " on cpu `" << llvm_target->GetCPU() << "`";
-
- // run ctors
- err = ctorRunner.run();
- TVM_FFI_ICHECK(!err) << llvm::toString(std::move(err));
-
- if (void** ctx_addr =
-
reinterpret_cast<void**>(GetGlobalAddr(ffi::symbol::tvm_ffi_library_ctx,
*llvm_target))) {
- *ctx_addr = this;
+ << " and JIT target machine (" << layout.getStringRepresentation() <<
")";
+
+ std::unique_ptr<llvm::Module> object_module = CloneLLVMModule(module_.get());
+
+ llvm::SmallString<0> object;
+ llvm::raw_svector_ostream object_stream(object);
+ llvm::legacy::PassManager pass;
+ TVM_FFI_ICHECK(!LLVMAddPassesToEmitFile(tm, &pass, &object_stream,
llvm_object_file_target))
+ << "Cannot emit LLVM object for apache-tvm-ffi-orcjit";
+ pass.run(*object_module);
+
+ ffi::ObjectRef session =
get_default_session.value()().cast<ffi::ObjectRef>();
+ ffi::Array<ffi::Variant<ffi::String, ffi::Bytes>> objects = {
+ ffi::Bytes(object.data(), object.size())};
+ ffi::Module dylib = load_module
+ .value()(session, objects, ffi::String(""),
ffi::Optional<ffi::String>{},
+ ffi::Optional<ffi::String>{})
+ .cast<ffi::Module>();
+ for (const ffi::Any& imported_module : imports()) {
+ dylib->ImportModule(imported_module.cast<ffi::Module>());
}
- ffi::Module::VisitContextSymbols([this, &llvm_target](const ffi::String&
name, void* symbol) {
- if (void** ctx_addr = reinterpret_cast<void**>(GetGlobalAddr(name,
*llvm_target))) {
- *ctx_addr = symbol;
- }
- });
-}
-
-bool LLVMModuleNode::IsCompatibleWithHost(const llvm::TargetMachine* tm) const
{
- LLVMTargetInfo host_target(*llvm_instance_, "llvm");
- auto tm_host = host_target.GetOrCreateTargetMachine();
- if (tm_host->getTargetTriple().getArch() != tm->getTargetTriple().getArch())
{
- LOG(INFO) << "Architecture mismatch: module=" <<
tm->getTargetTriple().str()
- << " host=" << tm_host->getTargetTriple().str();
- return false;
- }
- return true;
-}
-// Get global address from execution engine.
-void* LLVMModuleNode::GetGlobalAddr(const std::string& name, const LLVMTarget&
llvm_target) const {
- // first verifies if GV exists.
- if (module_->getGlobalVariable(name) != nullptr) {
- if (jit_engine_ == "mcjit") {
- return reinterpret_cast<void*>(mcjit_ee_->getGlobalValueAddress(name));
- } else if (jit_engine_ == "orcjit") {
- auto addr = llvm::cantFail(orcjit_ee_->lookup(name)).getValue();
- return reinterpret_cast<void*>(addr);
- } else {
- TVM_FFI_THROW(InternalError) << "Either `mcjit` or `orcjit` are not
initialized.";
- }
- }
- return nullptr;
-}
+ tvm_ffi_orcjit_module_ = dylib;
-void* LLVMModuleNode::GetFunctionAddr(const std::string& name,
- const LLVMTarget& llvm_target) const {
- // first verifies if GV exists.
- if (module_->getFunction(name) != nullptr) {
- if (jit_engine_ == "mcjit") {
- return reinterpret_cast<void*>(mcjit_ee_->getFunctionAddress(name));
- } else if (jit_engine_ == "orcjit") {
- auto addr = llvm::cantFail(orcjit_ee_->lookup(name)).getValue();
- return reinterpret_cast<void*>(addr);
- } else {
- TVM_FFI_THROW(InternalError) << "Either `mcjit` or `orcjit` are not
initialized.";
- }
- }
- return nullptr;
+ VLOG(2) << "apache-tvm-ffi-orcjit execute " <<
module_->getModuleIdentifier() << " for triple `"
+ << llvm_target->GetTargetTriple() << "` on cpu `" <<
llvm_target->GetCPU()
+ << "` with features `" << llvm_target->GetTargetFeatureString() <<
"`";
}
static void LLVMReflectionRegister() {
@@ -646,7 +386,6 @@ static void LLVMReflectionRegister() {
#endif
module->setDataLayout(llvm_target->GetOrCreateTargetMachine()->createDataLayout());
n->Init(std::move(module), std::move(llvm_instance));
- n->SetJITEngine(llvm_target->GetJITEngine());
return ffi::Module(n);
})
.def("target.llvm_lookup_intrinsic_id",
@@ -763,7 +502,6 @@ static void LLVMReflectionRegister() {
.def("ffi.Module.load_from_file.ll",
[](std::string filename, std::string fmt) -> ffi::Module {
auto n = ffi::make_object<LLVMModuleNode>();
- n->SetJITEngine("orcjit");
n->LoadIR(filename);
return ffi::Module(n);
})
@@ -783,7 +521,6 @@ static void LLVMReflectionRegister() {
std::unique_ptr<llvm::Module> blob =
CodeGenBlob(data, system_lib, llvm_target.get(),
c_symbol_prefix);
n->Init(std::move(blob), std::move(llvm_instance));
- n->SetJITEngine(llvm_target->GetJITEngine());
return ffi::Module(n);
});
}
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 93c66869dc..0e326f6a6f 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -144,8 +144,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_option<int64_t>("opt-level")
// LLVM command line flags, see below
.def_option<ffi::Array<ffi::String>>("cl-opt")
- // LLVM JIT engine mcjit/orcjit
- .def_option<ffi::String>("jit")
// TVM & LLVM custom vector bit width
.def_option<int64_t>("vector-width")
.set_default_keys({"cpu"})
diff --git a/tests/python/runtime/test_runtime_module_load.py
b/tests/python/runtime/test_runtime_module_load.py
index ee5500fd1b..91479f72b3 100644
--- a/tests/python/runtime/test_runtime_module_load.py
+++ b/tests/python/runtime/test_runtime_module_load.py
@@ -45,8 +45,7 @@ print("Finish runtime checking...")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
[email protected]("target", ["llvm", {"kind": "llvm", "jit": "mcjit"}])
-def test_dso_module_load(target):
+def test_dso_module_load():
dtype = "int64"
temp = utils.tempdir()
@@ -65,7 +64,7 @@ def test_dso_module_load(target):
mod = tvm.IRModule.from_expr(
tvm.tirx.PrimFunc([Ab], stmt).with_attr("global_symbol", "main")
)
- m = tvm.tirx.build(mod, target=target)
+ m = tvm.tirx.build(mod, target="llvm")
for name in names:
m.write_to_file(name)
diff --git a/tests/python/target/test_target_target.py
b/tests/python/target/test_target_target.py
index 73074b72ea..aa81d3be8d 100644
--- a/tests/python/target/test_target_target.py
+++ b/tests/python/target/test_target_target.py
@@ -98,13 +98,6 @@ def test_target_llvm_options():
)
-def test_target_llvm_jit_options():
- target = tvm.target.Target({"kind": "llvm", "jit": "mcjit"})
- assert target.attrs["jit"] == "mcjit"
- target = tvm.target.Target({"kind": "llvm", "jit": "orcjit"})
- assert target.attrs["jit"] == "orcjit"
-
-
def test_target_llvm_vector_width():
target = tvm.target.Target({"kind": "llvm", "vector-width": 256})
assert target.attrs["vector-width"] == 256
diff --git a/tests/scripts/task_python_docs.sh
b/tests/scripts/task_python_docs.sh
index 11b2501b7b..e45d439929 100755
--- a/tests/scripts/task_python_docs.sh
+++ b/tests/scripts/task_python_docs.sh
@@ -132,6 +132,9 @@ find . -type f -path "*.pyc" | xargs rm -f
# setup tvm-ffi into python folder
uv pip install -v --target=python ./3rdparty/tvm-ffi/
+# LLVMModule JIT execution, needed by the tutorials this build runs.
+uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
+
cd docs
PYTHONPATH=$(pwd)/../python make htmldepoly SPHINXOPTS='-j auto' |& tee
/tmp/$$.log.txt
diff --git a/tests/scripts/task_python_unittest.sh
b/tests/scripts/task_python_unittest.sh
index 524a9f2a2e..511b3ec2a5 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -24,4 +24,9 @@ export PYTEST_ADDOPTS="${CI_PYTEST_ADD_OPTIONS:-}
${PYTEST_ADDOPTS:-}"
# setup tvm-ffi into python folder
uv pip install -v --target=python ./3rdparty/tvm-ffi/
+# LLVMModule JIT execution. Installed here as well as in the CI image because
a docker/
+# change does not reach the test containers in the same run (they use the
pinned tag).
+# --no-deps keeps the submodule-built tvm-ffi above from being replaced by a
PyPI build.
+uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
+
python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER:-not gpu}" tests/python
