This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch orcjit-windows-arm64-855 in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
commit 811965c0794b9d17329e272eb90c6e0c8a78f758 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 26 03:31:28 2026 +0000 [FIX] Bind ARM64 lifecycle ranges to object memory ownership --- .../ffi/llvm_patches/win_dll_import_generator.cc | 1 + addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 39 +++++++++++++++++----- 2 files changed, 32 insertions(+), 8 deletions(-) diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc index ec484d10..528d462c 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.cc @@ -39,6 +39,7 @@ #include <llvm/ADT/DenseMap.h> #include <llvm/ExecutionEngine/JITLink/JITLink.h> #include <llvm/ExecutionEngine/JITLink/x86_64.h> +#include <llvm/ExecutionEngine/Orc/AbsoluteSymbols.h> #include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h> #include <llvm/Support/DynamicLibrary.h> #include <llvm/Support/Error.h> diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc index d52d04d5..a3db0990 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc @@ -180,18 +180,41 @@ std::unique_ptr<llvm::orc::ObjectLayer> CreateARM64COFFLayer(llvm::orc::Executio bool is_init; }; using Responsibility = llvm::orc::MaterializationResponsibility; - using PendingRanges = std::unordered_map<Responsibility*, std::vector<LifecycleRange>>; + using Ranges = std::vector<LifecycleRange>; + using PendingRanges = std::unordered_map<const char*, std::shared_ptr<Ranges>>; + class ObjectMemoryManager : public llvm::SectionMemoryManager { + public: + ObjectMemoryManager(std::shared_ptr<PendingRanges> pending, const char* object) + : llvm::SectionMemoryManager(nullptr, true), + pending_(std::move(pending)), + object_(object), + ranges_(std::make_shared<Ranges>()) { + (*pending_)[object_] = ranges_; + } + ~ObjectMemoryManager() override { + // NotifyEmitted is skipped on materialization failure. Release retained + // lifecycle ranges together with the memory they describe on every path. + auto it = pending_->find(object_); + // An emitted object's buffer can be reused before its memory is released. + if (it != pending_->end() && it->second == ranges_) pending_->erase(it); + } + + private: + std::shared_ptr<PendingRanges> pending_; + const char* object_; + std::shared_ptr<Ranges> ranges_; + }; auto pending = std::make_shared<PendingRanges>(); - auto layer = - std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(ES, [](const llvm::MemoryBuffer&) { + auto layer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>( + ES, [pending](const llvm::MemoryBuffer& object) { // Keep each object's code and data within ARM64 PC-relative range. - return std::make_unique<llvm::SectionMemoryManager>(nullptr, true); + return std::make_unique<ObjectMemoryManager>(pending, object.getBufferStart()); }); layer->setProcessAllSections(true); layer->setOverrideObjectFlagsWithResponsibilityFlags(true); layer->setNotifyLoaded([pending](Responsibility& R, const llvm::object::ObjectFile& Obj, const llvm::RuntimeDyld::LoadedObjectInfo& Info) { - auto& ranges = (*pending)[&R]; + auto& ranges = *pending->at(Obj.getMemoryBufferRef().getBufferStart()); ranges.clear(); for (const auto& section : Obj.sections()) { auto name = section.getName(); @@ -213,12 +236,12 @@ std::unique_ptr<llvm::orc::ObjectLayer> CreateARM64COFFLayer(llvm::orc::Executio } }); layer->setNotifyEmitted( - [pending, session](Responsibility& R, std::unique_ptr<llvm::MemoryBuffer>) { - auto it = pending->find(&R); + [pending, session](Responsibility& R, std::unique_ptr<llvm::MemoryBuffer> object) { + auto it = pending->find(object->getBufferStart()); if (it == pending->end()) return; auto ranges = std::move(it->second); pending->erase(it); - for (const auto& range : ranges) { + for (const auto& range : *ranges) { if (!range.address) continue; for (uint64_t offset = 0; offset + sizeof(uint64_t) <= range.size; offset += sizeof(uint64_t)) {
