From 0c3751f280bbc31c1ec133e256f82a86ad116b56 Mon Sep 17 00:00:00 2001 From: Jameson Nash Date: Fri, 19 Jul 2024 21:08:28 +0000 Subject: [PATCH] compress DebugObject to save memory Account for debug info memory as well, since this info is usually about 10x the size of the data that was being tracked, although this compression appears to reduce that by at least half. --- src/debug-registry.h | 3 ++- src/debuginfo.cpp | 39 +++++++++++++++++++++++++++++++-------- src/debuginfo.h | 1 + src/jitlayers.cpp | 31 +++++++++++++++++++------------ src/jitlayers.h | 3 ++- 5 files changed, 55 insertions(+), 22 deletions(-) diff --git a/src/debug-registry.h b/src/debug-registry.h index 91f68c9d2058a7..85a94245ce6aa3 100644 --- a/src/debug-registry.h +++ b/src/debug-registry.h @@ -100,7 +100,8 @@ class JITDebugInfoRegistry private: struct LazyObjectInfo { - std::unique_ptr data; + SmallVector data; + size_t uncompressedsize; std::unique_ptr object; std::unique_ptr context; LazyObjectInfo() = delete; diff --git a/src/debuginfo.cpp b/src/debuginfo.cpp index d40954b8ed8cd4..fd803c9e62ec4d 100644 --- a/src/debuginfo.cpp +++ b/src/debuginfo.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -335,7 +336,10 @@ void JITDebugInfoRegistry::registerJITObject(const object::ObjectFile &Object, #endif // defined(_OS_X86_64_) #endif // defined(_OS_WINDOWS_) - auto ObjectCopy = new LazyObjectInfo{MemoryBuffer::getMemBufferCopy(Object.getData())}; // intentionally leaked so that we don't need to ref-count it + 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); for (const auto &sym_size : symbols) { const object::SymbolRef &sym_iter = sym_size.first; @@ -1211,13 +1215,32 @@ int jl_DI_for_fptr(uint64_t fptr, uint64_t *symsize, int64_t *slide, if (fit != objmap.end() && fptr < fit->first + fit->second.SectionSize) { *slide = fit->second.slide; auto lazyobject = fit->second.object; - if (!lazyobject->object) - lazyobject->object = cantFail(object::ObjectFile::createObjectFile(lazyobject->data->getMemBufferRef())); - *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(); + 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 5b5cdcb82d534f..6cd75289107652 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 ac403fc157e946..af8860d4d1116a 100644 --- a/src/jitlayers.cpp +++ b/src/jitlayers.cpp @@ -781,12 +781,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(); @@ -835,7 +835,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); @@ -1579,7 +1579,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, @@ -1991,19 +1991,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() { @@ -2288,3 +2289,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 1eea3692aec71c..101f5714abd112 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;