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 d8b428677a51c1f1dc4f1cdd8ccd95223fb33fe8 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 26 02:18:04 2026 +0000 Use RuntimeDyld for Windows ARM64 COFF objects --- .../ffi/llvm_patches/win_dll_import_generator.cc | 18 +++-- .../ffi/llvm_patches/win_dll_import_generator.h | 7 +- addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc | 90 +++++++++++++++++++++- 3 files changed, 105 insertions(+), 10 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 69fa3fdb..ec484d10 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 @@ -38,7 +38,6 @@ #include <llvm/ADT/DenseMap.h> #include <llvm/ExecutionEngine/JITLink/JITLink.h> -#include <llvm/ExecutionEngine/JITLink/aarch64.h> #include <llvm/ExecutionEngine/JITLink/x86_64.h> #include <llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h> #include <llvm/Support/DynamicLibrary.h> @@ -127,6 +126,14 @@ llvm::Error DLLImportDefinitionGenerator::tryToGenerate( } if (Resolved.empty()) return llvm::Error::success(); + if (!L_) { + llvm::orc::SymbolMap Symbols; + for (auto& [Name, Addr] : Resolved) { + Symbols[Name] = llvm::orc::ExecutorSymbolDef(Addr, llvm::JITSymbolFlags::Exported); + } + return JD.define(llvm::orc::absoluteSymbols(std::move(Symbols))); + } + // Step 2: Build a LinkGraph with __imp_ pointers and PLT jump stubs. auto G = std::make_unique<llvm::jitlink::LinkGraph>( "<DLL_IMPORT_STUBS>", ES_.getSymbolStringPool(), ES_.getTargetTriple(), @@ -141,19 +148,16 @@ llvm::Error DLLImportDefinitionGenerator::tryToGenerate( G->getPointerSize(), llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Local, false); // __imp_XXX pointer (GOT-like entry) - bool IsAArch64 = ES_.getTargetTriple().getArch() == llvm::Triple::aarch64; - auto& Ptr = IsAArch64 ? llvm::jitlink::aarch64::createAnonymousPointer(*G, Sec, &Target) - : llvm::jitlink::x86_64::createAnonymousPointer(*G, Sec, &Target); + auto& Ptr = llvm::jitlink::x86_64::createAnonymousPointer(*G, Sec, &Target); Ptr.setName(G->intern(("__imp_" + *InternedName).str())); Ptr.setLinkage(llvm::jitlink::Linkage::Strong); Ptr.setScope(llvm::jitlink::Scope::Default); // XXX jump stub (PLT-like entry) for direct calls - auto& StubBlock = IsAArch64 ? llvm::jitlink::aarch64::createPointerJumpStubBlock(*G, Sec, Ptr) - : llvm::jitlink::x86_64::createPointerJumpStubBlock(*G, Sec, Ptr); + auto& StubBlock = llvm::jitlink::x86_64::createPointerJumpStubBlock(*G, Sec, Ptr); G->addDefinedSymbol(StubBlock, 0, *InternedName, StubBlock.getSize(), llvm::jitlink::Linkage::Strong, llvm::jitlink::Scope::Default, true, false); } - return L_.add(JD, std::move(G)); + return L_->add(JD, std::move(G)); } } // namespace orcjit diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h index 3b0d34c4..f37161d4 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h +++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_dll_import_generator.h @@ -71,7 +71,10 @@ namespace orcjit { class DLLImportDefinitionGenerator : public llvm::orc::DefinitionGenerator { public: DLLImportDefinitionGenerator(llvm::orc::ExecutionSession& ES, llvm::orc::ObjectLinkingLayer& L) - : ES_(ES), L_(L) {} + : ES_(ES), L_(&L) {} + + // RuntimeDyld creates ARM64 COFF import slots and branch trampolines itself. + explicit DLLImportDefinitionGenerator(llvm::orc::ExecutionSession& ES) : ES_(ES), L_(nullptr) {} llvm::Error tryToGenerate(llvm::orc::LookupState& LS, llvm::orc::LookupKind K, llvm::orc::JITDylib& JD, llvm::orc::JITDylibLookupFlags JDLookupFlags, @@ -81,7 +84,7 @@ class DLLImportDefinitionGenerator : public llvm::orc::DefinitionGenerator { static void* FindInProcessModules(const std::string& Name); llvm::orc::ExecutionSession& ES_; - llvm::orc::ObjectLinkingLayer& L_; + llvm::orc::ObjectLinkingLayer* L_; }; } // namespace orcjit diff --git a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc index d525367d..d52d04d5 100644 --- a/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc +++ b/addons/tvm_ffi_orcjit/src/ffi/orcjit_session.cc @@ -29,6 +29,10 @@ #include <llvm/ExecutionEngine/Orc/ExecutionUtils.h> #include <llvm/ExecutionEngine/Orc/LLJIT.h> #include <llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h> +#if defined(_WIN32) && (defined(_M_ARM64) || defined(__aarch64__)) +#include <llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h> +#include <llvm/ExecutionEngine/SectionMemoryManager.h> +#endif #include <llvm/Support/Error.h> #include <llvm/Support/MemoryBuffer.h> #include <llvm/Support/Process.h> @@ -40,6 +44,7 @@ #include <cstddef> #include <cstdint> +#include <cstring> #include <mutex> #include <sstream> @@ -162,6 +167,79 @@ class LibStdCxxNonsharedGenerator final : public llvm::orc::DefinitionGenerator #endif #endif +#if defined(_WIN32) && (defined(_M_ARM64) || defined(__aarch64__)) +// LLVM's COFF JITLink supports x86-64 only. RuntimeDyld supplies ARM64 COFF +// relocations, import slots and long-branch stubs. Retain our local CRT +// lifecycle policy, collecting pointers only after relocation has finished. +std::unique_ptr<llvm::orc::ObjectLayer> CreateARM64COFFLayer(llvm::orc::ExecutionSession& ES, + ORCJITExecutionSessionObj* session) { + struct LifecycleRange { + uint64_t address; + uint64_t size; + int priority; + bool is_init; + }; + using Responsibility = llvm::orc::MaterializationResponsibility; + using PendingRanges = std::unordered_map<Responsibility*, std::vector<LifecycleRange>>; + auto pending = std::make_shared<PendingRanges>(); + auto layer = + std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(ES, [](const llvm::MemoryBuffer&) { + // Keep each object's code and data within ARM64 PC-relative range. + return std::make_unique<llvm::SectionMemoryManager>(nullptr, true); + }); + 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]; + ranges.clear(); + for (const auto& section : Obj.sections()) { + auto name = section.getName(); + if (!name) { + llvm::consumeError(name.takeError()); + continue; + } + bool is_init = name->starts_with(".CRT$XC"); + if (!is_init && !name->starts_with(".CRT$XT")) continue; + auto suffix = name->substr(7); + int priority = 0; + if (!suffix.empty()) { + priority = static_cast<int>(suffix[0]) * 100000; + int number = 0; + suffix.substr(1).getAsInteger(10, number); + priority += number; + } + ranges.push_back({Info.getSectionLoadAddress(section), section.getSize(), priority, is_init}); + } + }); + layer->setNotifyEmitted( + [pending, session](Responsibility& R, std::unique_ptr<llvm::MemoryBuffer>) { + auto it = pending->find(&R); + if (it == pending->end()) return; + auto ranges = std::move(it->second); + pending->erase(it); + for (const auto& range : ranges) { + if (!range.address) continue; + for (uint64_t offset = 0; offset + sizeof(uint64_t) <= range.size; + offset += sizeof(uint64_t)) { + uint64_t address = 0; + std::memcpy(&address, reinterpret_cast<const void*>(range.address + offset), + sizeof(address)); + if (!address) continue; + ORCJITExecutionSessionObj::InitFiniEntry entry{llvm::orc::ExecutorAddr(address), + range.priority}; + if (range.is_init) { + session->AddPendingInitializer(&R.getTargetJITDylib(), entry); + } else { + session->AddPendingDeinitializer(&R.getTargetJITDylib(), entry); + } + } + } + }); + return layer; +} +#endif + // Install ExecutorNativePlatform per the `orc_rt` selector (see the ctor doc). // A no-op off Linux/ELF: those targets retain LLJIT's generic platform support // and use the local lifecycle adapter below. @@ -264,6 +342,12 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const Optional<Variant<Stri return std::make_unique<llvm::orc::ObjectLinkingLayer>(ES, memory_manager); }); } // if (memory_manager_) +#elif defined(_WIN32) && (defined(_M_ARM64) || defined(__aarch64__)) + builder.setObjectLinkingLayerCreator( + [this](llvm::orc::ExecutionSession& ES, llvm::jitlink::JITLinkMemoryManager&) + -> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> { + return CreateARM64COFFLayer(ES, this); + }); #elif defined(__APPLE__) || defined(_WIN32) // Force ObjectLinkingLayer (JITLink) so we can attach InitFiniPlugin. // macOS: LLJIT already defaults to JITLink for Darwin, but the explicit @@ -342,7 +426,7 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const Optional<Variant<Stri // building. See llvm_patches/win_coff_pdata_strip.h for the rationale. jit_->getObjTransformLayer().setTransform(&StripCoffPdataXdata); #endif -#if defined(__APPLE__) || defined(_WIN32) +#if defined(__APPLE__) || (defined(_WIN32) && !defined(_M_ARM64) && !defined(__aarch64__)) // Mach-O and COFF still need the local lifecycle adapter. Linux uses LLVM // 23's ELFNixPlatform initialize/deinitialize path instead. auto& objlayer = jit_->getObjLinkingLayer(); @@ -355,9 +439,13 @@ ORCJITExecutionSessionObj::ORCJITExecutionSessionObj(const Optional<Variant<Stri // comprehensive generator that searches all loaded DLLs (vcruntime140, // ucrtbase, tvm_ffi, etc.) and creates __imp_* pointer stubs. if (auto PSG = jit_->getProcessSymbolsJITDylib()) { +#if defined(_M_ARM64) || defined(__aarch64__) + PSG->addGenerator(std::make_unique<DLLImportDefinitionGenerator>(jit_->getExecutionSession())); +#else auto& ObjLayer = static_cast<llvm::orc::ObjectLinkingLayer&>(jit_->getObjLinkingLayer()); PSG->addGenerator( std::make_unique<DLLImportDefinitionGenerator>(jit_->getExecutionSession(), ObjLayer)); +#endif } #endif }
