From fe597c1e60f4984bff2f76ac7ae6702e2e151a18 Mon Sep 17 00:00:00 2001 From: Jameson Nash Date: Tue, 23 Jul 2024 11:52:17 -0400 Subject: [PATCH] compress jit debuginfo for easy memory savings (#55180) In some ad-hoc testing, I had JIT about 19 MB of code and data, which generated about 170 MB of debuginfo alongside it, and that debuginfo then compressed to about 50 MB with this change, which simply compresses the ObjectFile until it is actually required (which it very rarely is needed). --- src/debug-registry.h | 22 +++++++---- src/debuginfo.cpp | 57 ++++++++++++++++++++------- src/debuginfo.h | 1 + src/jitlayers.cpp | 93 ++++++++++++++++---------------------------- src/jitlayers.h | 3 +- 5 files changed, 94 insertions(+), 82 deletions(-) diff --git a/src/debug-registry.h b/src/debug-registry.h index f30049eb5b210..85a94245ce6aa 100644 --- a/src/debug-registry.h +++ b/src/debug-registry.h @@ -99,18 +99,26 @@ class JITDebugInfoRegistry }; private: - struct ObjectInfo { - const llvm::object::ObjectFile *object = nullptr; - size_t SectionSize = 0; - ptrdiff_t slide = 0; - llvm::object::SectionRef Section{}; - llvm::DIContext *context = nullptr; + struct LazyObjectInfo { + SmallVector data; + size_t uncompressedsize; + std::unique_ptr object; + std::unique_ptr context; + LazyObjectInfo() = delete; + }; + + struct SectionInfo { + LazyObjectInfo *object; + size_t SectionSize; + ptrdiff_t slide; + uint64_t SectionIndex; + SectionInfo() = delete; }; template using rev_map = std::map>; - typedef rev_map objectmap_t; + typedef rev_map objectmap_t; typedef rev_map objfilemap_t; objectmap_t objectmap{}; diff --git a/src/debuginfo.cpp b/src/debuginfo.cpp index 5cb7401a036d4..84550811072fe 100644 --- a/src/debuginfo.cpp +++ b/src/debuginfo.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -335,8 +336,12 @@ void JITDebugInfoRegistry::registerJITObject(const object::ObjectFile &Object, #endif // defined(_OS_X86_64_) #endif // defined(_OS_WINDOWS_) + SmallVector packed; + compression::zlib::compress(ArrayRef((uint8_t*)Object.getData().data(), Object.getData().size()), packed, compression::zlib::DefaultCompression); + jl_jit_add_bytes(packed.size()); + auto ObjectCopy = new LazyObjectInfo{packed, Object.getData().size()}; // intentionally leaked so that we don't need to ref-count it, intentionally copied so that we exact-size the allocation (since no shrink_to_fit function) auto symbols = object::computeSymbolSizes(Object); - bool first = true; + bool hassection = false; for (const auto &sym_size : symbols) { const object::SymbolRef &sym_iter = sym_size.first; object::SymbolRef::Type SymbolType = cantFail(sym_iter.getType()); @@ -385,17 +390,17 @@ void JITDebugInfoRegistry::registerJITObject(const object::ObjectFile &Object, jl_profile_atomic([&]() JL_NOTSAFEPOINT { if (mi) linfomap[Addr] = std::make_pair(Size, mi); - if (first) { - objectmap[SectionLoadAddr] = {&Object, - (size_t)SectionSize, - (ptrdiff_t)(SectionAddr - SectionLoadAddr), - *Section, - nullptr, - }; - first = false; - } + hassection = true; + objectmap.insert(std::pair{SectionLoadAddr, SectionInfo{ + ObjectCopy, + (size_t)SectionSize, + (ptrdiff_t)(SectionAddr - SectionLoadAddr), + Section->getIndex() + }}); }); } + if (!hassection) // clang-sa demands that we do this to fool cplusplus.NewDeleteLeaks + delete ObjectCopy; } void jl_register_jit_object(const object::ObjectFile &Object, @@ -1213,11 +1218,33 @@ int jl_DI_for_fptr(uint64_t fptr, uint64_t *symsize, int64_t *slide, auto fit = objmap.lower_bound(fptr); if (fit != objmap.end() && fptr < fit->first + fit->second.SectionSize) { *slide = fit->second.slide; - *Section = fit->second.Section; - if (context) { - if (fit->second.context == nullptr) - fit->second.context = DWARFContext::create(*fit->second.object).release(); - *context = fit->second.context; + auto lazyobject = fit->second.object; + if (!lazyobject->object && !lazyobject->data.empty()) { + if (lazyobject->uncompressedsize) { + SmallVector unpacked; + Error E = compression::zlib::decompress(lazyobject->data, unpacked, lazyobject->uncompressedsize); + if (E) + lazyobject->data.clear(); + else + lazyobject->data = std::move(unpacked); + jl_jit_add_bytes(lazyobject->data.size() - lazyobject->uncompressedsize); + lazyobject->uncompressedsize = 0; + } + if (!lazyobject->data.empty()) { + auto obj = object::ObjectFile::createObjectFile(MemoryBufferRef(StringRef((const char*)lazyobject->data.data(), lazyobject->data.size()), "jit.o")); + if (obj) + lazyobject->object = std::move(*obj); + else + lazyobject->data.clear(); + } + } + if (lazyobject->object) { + *Section = *std::next(lazyobject->object->section_begin(), fit->second.SectionIndex); + if (context) { + if (lazyobject->context == nullptr) + lazyobject->context = DWARFContext::create(*lazyobject->object); + *context = lazyobject->context.get(); + } } found = 1; } diff --git a/src/debuginfo.h b/src/debuginfo.h index 5b5cdcb82d534..6cd7528910765 100644 --- a/src/debuginfo.h +++ b/src/debuginfo.h @@ -1,6 +1,7 @@ // This file is a part of Julia. License is MIT: https://julialang.org/license // Declarations for debuginfo.cpp +void jl_jit_add_bytes(size_t bytes) JL_NOTSAFEPOINT; int jl_DI_for_fptr(uint64_t fptr, uint64_t *symsize, int64_t *slide, llvm::object::SectionRef *Section, llvm::DIContext **context) JL_NOTSAFEPOINT; diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp index 6c40f4230ab69..487e6a0462500 100644 --- a/src/jitlayers.cpp +++ b/src/jitlayers.cpp @@ -684,22 +684,19 @@ struct JITObjectInfo { class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin { std::mutex PluginMutex; std::map> PendingObjs; - // Resources from distinct `MaterializationResponsibility`s can get merged - // after emission, so we can have multiple debug objects per resource key. - std::map, 0>> RegisteredObjs; public: void notifyMaterializing(MaterializationResponsibility &MR, jitlink::LinkGraph &G, jitlink::JITLinkContext &Ctx, MemoryBufferRef InputObject) override { - // Keeping around a full copy of the input object file (and re-parsing it) is - // wasteful, but for now, this lets us reuse the existing debuginfo.cpp code. - // Should look into just directly pulling out all the information required in - // a JITLink pass and just keeping the required tables/DWARF sections around - // (perhaps using the LLVM DebuggerSupportPlugin as a reference). auto NewBuffer = MemoryBuffer::getMemBufferCopy(InputObject.getBuffer(), G.getName()); + // Re-parsing the InputObject is wasteful, but for now, this lets us + // reuse the existing debuginfo.cpp code. Should look into just + // directly pulling out all the information required in a JITLink pass + // and just keeping the required tables/DWARF sections around (perhaps + // using the LLVM DebuggerSupportPlugin as a reference). auto NewObj = cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef())); @@ -733,13 +730,8 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin { }; jl_register_jit_object(*NewInfo->Object, getLoadAddress, nullptr); - } - - cantFail(MR.withResourceKeyDo([&](ResourceKey K) { - std::lock_guard lock(PluginMutex); - RegisteredObjs[K].push_back(std::move(PendingObjs[&MR])); PendingObjs.erase(&MR); - })); + } return Error::success(); } @@ -750,32 +742,23 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin { PendingObjs.erase(&MR); return Error::success(); } + #if JL_LLVM_VERSION >= 160000 Error notifyRemovingResources(JITDylib &JD, orc::ResourceKey K) override #else - Error notifyRemovingResources(ResourceKey K) override + Error notifyRemovingResources(orc::ResourceKey K) override #endif { - std::lock_guard lock(PluginMutex); - RegisteredObjs.erase(K); - // TODO: If we ever unload code, need to notify debuginfo registry. return Error::success(); } #if JL_LLVM_VERSION >= 160000 - void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) override + void notifyTransferringResources(JITDylib &JD, orc::ResourceKey DstKey, + orc::ResourceKey SrcKey) override {} #else - void notifyTransferringResources(ResourceKey DstKey, ResourceKey SrcKey) override + void notifyTransferringResources(orc::ResourceKey DstKey, + orc::ResourceKey SrcKey) override {} #endif - { - std::lock_guard lock(PluginMutex); - auto SrcIt = RegisteredObjs.find(SrcKey); - if (SrcIt != RegisteredObjs.end()) { - for (std::unique_ptr &Info : SrcIt->second) - RegisteredObjs[DstKey].push_back(std::move(Info)); - RegisteredObjs.erase(SrcIt); - } - } void modifyPassConfig(MaterializationResponsibility &MR, jitlink::LinkGraph &, jitlink::PassConfiguration &PassConfig) override @@ -815,12 +798,12 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin { class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin { private: - std::atomic &total_size; + std::atomic &jit_bytes_size; public: - JLMemoryUsagePlugin(std::atomic &total_size) - : total_size(total_size) {} + JLMemoryUsagePlugin(std::atomic &jit_bytes_size) + : jit_bytes_size(jit_bytes_size) {} Error notifyFailed(orc::MaterializationResponsibility &MR) override { return Error::success(); @@ -869,7 +852,7 @@ class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin { } (void) code_size; (void) data_size; - this->total_size.fetch_add(graph_size, std::memory_order_relaxed); + this->jit_bytes_size.fetch_add(graph_size, std::memory_order_relaxed); jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, graph_size); jl_timing_counter_inc(JL_TIMING_COUNTER_JITCodeSize, code_size); jl_timing_counter_inc(JL_TIMING_COUNTER_JITDataSize, data_size); @@ -985,24 +968,7 @@ void registerRTDyldJITObject(const object::ObjectFile &Object, const RuntimeDyld::LoadedObjectInfo &L, const std::shared_ptr &MemMgr) { - auto SavedObject = L.getObjectForDebug(Object).takeBinary(); - // If the debug object is unavailable, save (a copy of) the original object - // for our backtraces. - // This copy seems unfortunate, but there doesn't seem to be a way to take - // ownership of the original buffer. - if (!SavedObject.first) { - auto NewBuffer = - MemoryBuffer::getMemBufferCopy(Object.getData(), Object.getFileName()); - auto NewObj = - cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef())); - SavedObject = std::make_pair(std::move(NewObj), std::move(NewBuffer)); - } - const object::ObjectFile *DebugObj = SavedObject.first.release(); - SavedObject.second.release(); - StringMap loadedSections; - // Use the original Object, not the DebugObject, as this is used for the - // RuntimeDyld::LoadedObjectInfo lookup. for (const object::SectionRef &lSection : Object.sections()) { auto sName = lSection.getName(); if (sName) { @@ -1019,7 +985,9 @@ void registerRTDyldJITObject(const object::ObjectFile &Object, return L.getSectionLoadAddress(search->second); }; - jl_register_jit_object(*DebugObj, getLoadAddress, + auto DebugObject = L.getObjectForDebug(Object); // ELF requires us to make a copy to mutate the header with the section load addresses. On other platforms this is a no-op. + jl_register_jit_object(DebugObject.getBinary() ? *DebugObject.getBinary() : Object, + getLoadAddress, #if defined(_OS_WINDOWS_) && defined(_CPU_X86_64_) [MemMgr](void *p) { return lookupWriteAddressFor(MemMgr.get(), p); } #else @@ -1630,7 +1598,7 @@ JuliaOJIT::JuliaOJIT() ES, std::move(ehRegistrar))); ObjectLayer.addPlugin(std::make_unique()); - ObjectLayer.addPlugin(std::make_unique(total_size)); + ObjectLayer.addPlugin(std::make_unique(jit_bytes_size)); #else ObjectLayer.setNotifyLoaded( [this](orc::MaterializationResponsibility &MR, @@ -2042,19 +2010,20 @@ std::string JuliaOJIT::getMangledName(const GlobalValue *GV) return getMangledName(GV->getName()); } -#ifdef JL_USE_JITLINK size_t JuliaOJIT::getTotalBytes() const { - return total_size.load(std::memory_order_relaxed); + auto bytes = jit_bytes_size.load(std::memory_order_relaxed); +#ifndef JL_USE_JITLINK + size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT; + bytes += getRTDyldMemoryManagerTotalBytes(MemMgr.get()); +#endif + return bytes; } -#else -size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT; -size_t JuliaOJIT::getTotalBytes() const +void JuliaOJIT::addBytes(size_t bytes) { - return getRTDyldMemoryManagerTotalBytes(MemMgr.get()); + jit_bytes_size.fetch_add(bytes, std::memory_order_relaxed); } -#endif void JuliaOJIT::printTimers() { @@ -2339,3 +2308,9 @@ size_t jl_jit_total_bytes_impl(void) { return jl_ExecutionEngine->getTotalBytes(); } + +// API for adding bytes to record being owned by the JIT +void jl_jit_add_bytes(size_t bytes) +{ + jl_ExecutionEngine->addBytes(bytes); +} diff --git a/src/jitlayers.h b/src/jitlayers.h index 1eea3692aec71..101f5714abd11 100644 --- a/src/jitlayers.h +++ b/src/jitlayers.h @@ -559,6 +559,7 @@ class JuliaOJIT { TargetIRAnalysis getTargetIRAnalysis() const JL_NOTSAFEPOINT; size_t getTotalBytes() const JL_NOTSAFEPOINT; + void addBytes(size_t bytes) JL_NOTSAFEPOINT; void printTimers() JL_NOTSAFEPOINT; jl_locked_stream &get_dump_emitted_mi_name_stream() JL_NOTSAFEPOINT { @@ -605,10 +606,10 @@ class JuliaOJIT { ResourcePool> ContextPool; + std::atomic jit_bytes_size{0}; #ifndef JL_USE_JITLINK const std::shared_ptr MemMgr; #else - std::atomic total_size{0}; const std::unique_ptr MemMgr; #endif ObjLayerT ObjectLayer;