diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 6ff75d1aea68..c888a6c3e094 100644 --- a/.github/workflows/main.yml +++ b/.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 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: >- diff --git a/docker/install/ubuntu_install_python_package.sh b/docker/install/ubuntu_install_python_package.sh index 8be4a6a8d7ab..7da9ffffbe00 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.1 diff --git a/docs/arch/codegen.rst b/docs/arch/codegen.rst index 1a44f2b37a72..2572c997117c 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.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``. diff --git a/pyproject.toml b/pyproject.toml index 702a525da343..61a207003a05 100644 --- a/pyproject.toml +++ b/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", diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py index 16cef942ad68..bf753e3810f8 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 @@ 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 diff --git a/src/target/llvm/llvm_instance.cc b/src/target/llvm/llvm_instance.cc index 0b1275498cf9..a17e92ce8c82 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::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>()) { @@ -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 0cb141de99bd..5d797d014124 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 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 a2226c5dec37..536db1731f32 100644 --- a/src/target/llvm/llvm_module.cc +++ b/src/target/llvm/llvm_module.cc @@ -26,16 +26,6 @@ #include #include #include -#include -#include -#include -#include -#include -#include -#if _WIN32 -#include -#include -#endif #include #include #include @@ -46,6 +36,8 @@ #include #include #include +#include +#include #if TVM_LLVM_VERSION >= 180 #include #else @@ -54,10 +46,10 @@ #include #include #include -#include #include #include #include +#include #include #include #include @@ -67,14 +59,10 @@ #include #include -#include #include -#include -#include #include #include #include -#include #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 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 module, std::unique_ptr 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 llvm_instance_; - // JIT lock - std::mutex mutex_; - // jit execution engines - llvm::ExecutionEngine* mcjit_ee_{nullptr}; - std::unique_ptr 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 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 tvm_ffi_orcjit_module_; + // The retained module IR used for inspection and object emission. + std::unique_ptr module_; /* \brief names of the external functions declared in this module */ ffi::Array 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(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 LLVMModuleNode::GetFunction(const ffi::String& name) { ffi::ObjectPtr sptr_to_self = ffi::GetObjectPtr(this); if (name == "__tvm_is_system_module") { @@ -181,29 +141,23 @@ ffi::Optional 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 lock(mutex_); + EnsureOrcJITModule(); + return tvm_ffi_orcjit_module_.value()->GetFunction(name); +} - TVMFFISafeCallType faddr; - With llvm_target(*llvm_instance_, LLVMTarget::GetTargetMetadata(*module_)); - ffi::String name_with_prefix = ffi::symbol::tvm_ffi_symbol_prefix + name; - faddr = reinterpret_cast(GetFunctionAddr(name_with_prefix, *llvm_target)); - if (faddr == nullptr) return std::nullopt; - ffi::Module self_strong_ref = ffi::GetRef(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(args.data()), - args.size(), reinterpret_cast(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 CloneLLVMModule(llvm::Module* mod) { return llvm::CloneModule(*mod); } +std::unique_ptr 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 module, std::unique_ptr 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,50 @@ 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 lock(mutex_); - if (mcjit_ee_) { +void LLVMModuleNode::EnsureOrcJITModule() { + if (tvm_ffi_orcjit_module_.has_value()) { return; } - // MCJIT builder - With llvm_target(*llvm_instance_, LLVMTarget::GetTargetMetadata(*module_)); - llvm::EngineBuilder builder(std::move(module_owning_ptr_)); + ffi::Optional get_default_session = + ffi::Function::GetGlobal("tvm_ffi_orcjit.GlobalDefaultSession"); + ffi::Optional 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.1`."; - // 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(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(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(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 lock(mutex_); - if (orcjit_ee_) { - return; - } - // ORCJIT builder With 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 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> { - return std::make_unique(std::move(tm)); - }; - - // linker - const auto linkerBuilder = -#if TVM_LLVM_VERSION >= 230 - [&](llvm::orc::ExecutionSession& session, llvm::jitlink::JITLinkMemoryManager& mem_mgr) - -> llvm::Expected> { -#elif TVM_LLVM_VERSION >= 210 - [&](llvm::orc::ExecutionSession& session) - -> llvm::Expected> { -#else - [&](llvm::orc::ExecutionSession& session, - const llvm::Triple& triple) -> std::unique_ptr { -#endif -#if _WIN32 -#if TVM_LLVM_VERSION >= 210 - auto GetMemMgr = [](const llvm::MemoryBuffer&) { - return std::make_unique(); - }; -#else - auto GetMemMgr = []() { return std::make_unique(); }; -#endif - auto ObjLinkingLayer = - std::make_unique(session, std::move(GetMemMgr)); -#else -#if TVM_LLVM_VERSION >= 230 - auto ObjLinkingLayer = std::make_unique(session, mem_mgr); -#else - auto ObjLinkingLayer = std::make_unique(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::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(); - 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(GetGlobalAddr(ffi::symbol::tvm_ffi_library_ctx, *llvm_target))) { - *ctx_addr = this; + << " and JIT target machine (" << layout.getStringRepresentation() << ")"; + + std::unique_ptr 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::Array> objects = { + ffi::Bytes(object.data(), object.size())}; + ffi::Module dylib = load_module.value()(session, objects, ffi::String("")).cast(); + for (const ffi::Any& imported_module : imports()) { + dylib->ImportModule(imported_module.cast()); } - ffi::Module::VisitContextSymbols([this, &llvm_target](const ffi::String& name, void* symbol) { - if (void** ctx_addr = reinterpret_cast(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(mcjit_ee_->getGlobalValueAddress(name)); - } else if (jit_engine_ == "orcjit") { - auto addr = llvm::cantFail(orcjit_ee_->lookup(name)).getValue(); - return reinterpret_cast(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(mcjit_ee_->getFunctionAddress(name)); - } else if (jit_engine_ == "orcjit") { - auto addr = llvm::cantFail(orcjit_ee_->lookup(name)).getValue(); - return reinterpret_cast(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 +383,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 +499,6 @@ static void LLVMReflectionRegister() { .def("ffi.Module.load_from_file.ll", [](std::string filename, std::string fmt) -> ffi::Module { auto n = ffi::make_object(); - n->SetJITEngine("orcjit"); n->LoadIR(filename); return ffi::Module(n); }) @@ -783,7 +518,6 @@ static void LLVMReflectionRegister() { std::unique_ptr 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 68f1f2e3b7c8..f21be6aaf788 100644 --- a/src/target/target_kind.cc +++ b/src/target/target_kind.cc @@ -131,8 +131,6 @@ TVM_REGISTER_TARGET_KIND("llvm", kDLCPU) .add_attr_option("opt-level") // LLVM command line flags, see below .add_attr_option>("cl-opt") - // LLVM JIT engine mcjit/orcjit - .add_attr_option("jit") // TVM & LLVM custom vector bit width .add_attr_option("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 ee5500fd1b12..91479f72b3a6 100644 --- a/tests/python/runtime/test_runtime_module_load.py +++ b/tests/python/runtime/test_runtime_module_load.py @@ -45,8 +45,7 @@ @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -@pytest.mark.parametrize("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 save_object(names): 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 ed0b4d014221..e18af8c3829f 100644 --- a/tests/python/target/test_target_target.py +++ b/tests/python/target/test_target_target.py @@ -97,13 +97,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 11b2501b7bd8..bd72b5806977 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.1 + 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 524a9f2a2e51..12fb825ec0e3 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.1 + python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER:-not gpu}" tests/python