Skip to content

Commit fe597c1

Browse files
authored
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).
1 parent d68befd commit fe597c1

File tree

5 files changed

+94
-82
lines changed

5 files changed

+94
-82
lines changed

src/debug-registry.h

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -99,18 +99,26 @@ class JITDebugInfoRegistry
9999
};
100100
private:
101101

102-
struct ObjectInfo {
103-
const llvm::object::ObjectFile *object = nullptr;
104-
size_t SectionSize = 0;
105-
ptrdiff_t slide = 0;
106-
llvm::object::SectionRef Section{};
107-
llvm::DIContext *context = nullptr;
102+
struct LazyObjectInfo {
103+
SmallVector<uint8_t, 0> data;
104+
size_t uncompressedsize;
105+
std::unique_ptr<const llvm::object::ObjectFile> object;
106+
std::unique_ptr<llvm::DIContext> context;
107+
LazyObjectInfo() = delete;
108+
};
109+
110+
struct SectionInfo {
111+
LazyObjectInfo *object;
112+
size_t SectionSize;
113+
ptrdiff_t slide;
114+
uint64_t SectionIndex;
115+
SectionInfo() = delete;
108116
};
109117

110118
template<typename KeyT, typename ValT>
111119
using rev_map = std::map<KeyT, ValT, std::greater<KeyT>>;
112120

113-
typedef rev_map<size_t, ObjectInfo> objectmap_t;
121+
typedef rev_map<size_t, SectionInfo> objectmap_t;
114122
typedef rev_map<uint64_t, objfileentry_t> objfilemap_t;
115123

116124
objectmap_t objectmap{};

src/debuginfo.cpp

Lines changed: 42 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include <llvm/DebugInfo/DWARF/DWARFContext.h>
88
#include <llvm/Object/SymbolSize.h>
99
#include <llvm/Support/MemoryBuffer.h>
10+
#include <llvm/Support/MemoryBufferRef.h>
1011
#include <llvm/IR/Function.h>
1112
#include <llvm/ADT/StringRef.h>
1213
#include <llvm/ADT/StringMap.h>
@@ -335,8 +336,12 @@ void JITDebugInfoRegistry::registerJITObject(const object::ObjectFile &Object,
335336
#endif // defined(_OS_X86_64_)
336337
#endif // defined(_OS_WINDOWS_)
337338

339+
SmallVector<uint8_t, 0> packed;
340+
compression::zlib::compress(ArrayRef<uint8_t>((uint8_t*)Object.getData().data(), Object.getData().size()), packed, compression::zlib::DefaultCompression);
341+
jl_jit_add_bytes(packed.size());
342+
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)
338343
auto symbols = object::computeSymbolSizes(Object);
339-
bool first = true;
344+
bool hassection = false;
340345
for (const auto &sym_size : symbols) {
341346
const object::SymbolRef &sym_iter = sym_size.first;
342347
object::SymbolRef::Type SymbolType = cantFail(sym_iter.getType());
@@ -385,17 +390,17 @@ void JITDebugInfoRegistry::registerJITObject(const object::ObjectFile &Object,
385390
jl_profile_atomic([&]() JL_NOTSAFEPOINT {
386391
if (mi)
387392
linfomap[Addr] = std::make_pair(Size, mi);
388-
if (first) {
389-
objectmap[SectionLoadAddr] = {&Object,
390-
(size_t)SectionSize,
391-
(ptrdiff_t)(SectionAddr - SectionLoadAddr),
392-
*Section,
393-
nullptr,
394-
};
395-
first = false;
396-
}
393+
hassection = true;
394+
objectmap.insert(std::pair{SectionLoadAddr, SectionInfo{
395+
ObjectCopy,
396+
(size_t)SectionSize,
397+
(ptrdiff_t)(SectionAddr - SectionLoadAddr),
398+
Section->getIndex()
399+
}});
397400
});
398401
}
402+
if (!hassection) // clang-sa demands that we do this to fool cplusplus.NewDeleteLeaks
403+
delete ObjectCopy;
399404
}
400405

401406
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,
12131218
auto fit = objmap.lower_bound(fptr);
12141219
if (fit != objmap.end() && fptr < fit->first + fit->second.SectionSize) {
12151220
*slide = fit->second.slide;
1216-
*Section = fit->second.Section;
1217-
if (context) {
1218-
if (fit->second.context == nullptr)
1219-
fit->second.context = DWARFContext::create(*fit->second.object).release();
1220-
*context = fit->second.context;
1221+
auto lazyobject = fit->second.object;
1222+
if (!lazyobject->object && !lazyobject->data.empty()) {
1223+
if (lazyobject->uncompressedsize) {
1224+
SmallVector<uint8_t, 0> unpacked;
1225+
Error E = compression::zlib::decompress(lazyobject->data, unpacked, lazyobject->uncompressedsize);
1226+
if (E)
1227+
lazyobject->data.clear();
1228+
else
1229+
lazyobject->data = std::move(unpacked);
1230+
jl_jit_add_bytes(lazyobject->data.size() - lazyobject->uncompressedsize);
1231+
lazyobject->uncompressedsize = 0;
1232+
}
1233+
if (!lazyobject->data.empty()) {
1234+
auto obj = object::ObjectFile::createObjectFile(MemoryBufferRef(StringRef((const char*)lazyobject->data.data(), lazyobject->data.size()), "jit.o"));
1235+
if (obj)
1236+
lazyobject->object = std::move(*obj);
1237+
else
1238+
lazyobject->data.clear();
1239+
}
1240+
}
1241+
if (lazyobject->object) {
1242+
*Section = *std::next(lazyobject->object->section_begin(), fit->second.SectionIndex);
1243+
if (context) {
1244+
if (lazyobject->context == nullptr)
1245+
lazyobject->context = DWARFContext::create(*lazyobject->object);
1246+
*context = lazyobject->context.get();
1247+
}
12211248
}
12221249
found = 1;
12231250
}

src/debuginfo.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// This file is a part of Julia. License is MIT: https://julialang.org/license
22

33
// Declarations for debuginfo.cpp
4+
void jl_jit_add_bytes(size_t bytes) JL_NOTSAFEPOINT;
45

56
int jl_DI_for_fptr(uint64_t fptr, uint64_t *symsize, int64_t *slide,
67
llvm::object::SectionRef *Section, llvm::DIContext **context) JL_NOTSAFEPOINT;

src/jitlayers.cpp

Lines changed: 34 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -684,22 +684,19 @@ struct JITObjectInfo {
684684
class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
685685
std::mutex PluginMutex;
686686
std::map<MaterializationResponsibility *, std::unique_ptr<JITObjectInfo>> PendingObjs;
687-
// Resources from distinct `MaterializationResponsibility`s can get merged
688-
// after emission, so we can have multiple debug objects per resource key.
689-
std::map<ResourceKey, SmallVector<std::unique_ptr<JITObjectInfo>, 0>> RegisteredObjs;
690687

691688
public:
692689
void notifyMaterializing(MaterializationResponsibility &MR, jitlink::LinkGraph &G,
693690
jitlink::JITLinkContext &Ctx,
694691
MemoryBufferRef InputObject) override
695692
{
696-
// Keeping around a full copy of the input object file (and re-parsing it) is
697-
// wasteful, but for now, this lets us reuse the existing debuginfo.cpp code.
698-
// Should look into just directly pulling out all the information required in
699-
// a JITLink pass and just keeping the required tables/DWARF sections around
700-
// (perhaps using the LLVM DebuggerSupportPlugin as a reference).
701693
auto NewBuffer =
702694
MemoryBuffer::getMemBufferCopy(InputObject.getBuffer(), G.getName());
695+
// Re-parsing the InputObject is wasteful, but for now, this lets us
696+
// reuse the existing debuginfo.cpp code. Should look into just
697+
// directly pulling out all the information required in a JITLink pass
698+
// and just keeping the required tables/DWARF sections around (perhaps
699+
// using the LLVM DebuggerSupportPlugin as a reference).
703700
auto NewObj =
704701
cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef()));
705702

@@ -733,13 +730,8 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
733730
};
734731

735732
jl_register_jit_object(*NewInfo->Object, getLoadAddress, nullptr);
736-
}
737-
738-
cantFail(MR.withResourceKeyDo([&](ResourceKey K) {
739-
std::lock_guard<std::mutex> lock(PluginMutex);
740-
RegisteredObjs[K].push_back(std::move(PendingObjs[&MR]));
741733
PendingObjs.erase(&MR);
742-
}));
734+
}
743735

744736
return Error::success();
745737
}
@@ -750,32 +742,23 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
750742
PendingObjs.erase(&MR);
751743
return Error::success();
752744
}
745+
753746
#if JL_LLVM_VERSION >= 160000
754747
Error notifyRemovingResources(JITDylib &JD, orc::ResourceKey K) override
755748
#else
756-
Error notifyRemovingResources(ResourceKey K) override
749+
Error notifyRemovingResources(orc::ResourceKey K) override
757750
#endif
758751
{
759-
std::lock_guard<std::mutex> lock(PluginMutex);
760-
RegisteredObjs.erase(K);
761-
// TODO: If we ever unload code, need to notify debuginfo registry.
762752
return Error::success();
763753
}
764754

765755
#if JL_LLVM_VERSION >= 160000
766-
void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) override
756+
void notifyTransferringResources(JITDylib &JD, orc::ResourceKey DstKey,
757+
orc::ResourceKey SrcKey) override {}
767758
#else
768-
void notifyTransferringResources(ResourceKey DstKey, ResourceKey SrcKey) override
759+
void notifyTransferringResources(orc::ResourceKey DstKey,
760+
orc::ResourceKey SrcKey) override {}
769761
#endif
770-
{
771-
std::lock_guard<std::mutex> lock(PluginMutex);
772-
auto SrcIt = RegisteredObjs.find(SrcKey);
773-
if (SrcIt != RegisteredObjs.end()) {
774-
for (std::unique_ptr<JITObjectInfo> &Info : SrcIt->second)
775-
RegisteredObjs[DstKey].push_back(std::move(Info));
776-
RegisteredObjs.erase(SrcIt);
777-
}
778-
}
779762

780763
void modifyPassConfig(MaterializationResponsibility &MR, jitlink::LinkGraph &,
781764
jitlink::PassConfiguration &PassConfig) override
@@ -815,12 +798,12 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
815798

816799
class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin {
817800
private:
818-
std::atomic<size_t> &total_size;
801+
std::atomic<size_t> &jit_bytes_size;
819802

820803
public:
821804

822-
JLMemoryUsagePlugin(std::atomic<size_t> &total_size)
823-
: total_size(total_size) {}
805+
JLMemoryUsagePlugin(std::atomic<size_t> &jit_bytes_size)
806+
: jit_bytes_size(jit_bytes_size) {}
824807

825808
Error notifyFailed(orc::MaterializationResponsibility &MR) override {
826809
return Error::success();
@@ -869,7 +852,7 @@ class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin {
869852
}
870853
(void) code_size;
871854
(void) data_size;
872-
this->total_size.fetch_add(graph_size, std::memory_order_relaxed);
855+
this->jit_bytes_size.fetch_add(graph_size, std::memory_order_relaxed);
873856
jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, graph_size);
874857
jl_timing_counter_inc(JL_TIMING_COUNTER_JITCodeSize, code_size);
875858
jl_timing_counter_inc(JL_TIMING_COUNTER_JITDataSize, data_size);
@@ -985,24 +968,7 @@ void registerRTDyldJITObject(const object::ObjectFile &Object,
985968
const RuntimeDyld::LoadedObjectInfo &L,
986969
const std::shared_ptr<RTDyldMemoryManager> &MemMgr)
987970
{
988-
auto SavedObject = L.getObjectForDebug(Object).takeBinary();
989-
// If the debug object is unavailable, save (a copy of) the original object
990-
// for our backtraces.
991-
// This copy seems unfortunate, but there doesn't seem to be a way to take
992-
// ownership of the original buffer.
993-
if (!SavedObject.first) {
994-
auto NewBuffer =
995-
MemoryBuffer::getMemBufferCopy(Object.getData(), Object.getFileName());
996-
auto NewObj =
997-
cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef()));
998-
SavedObject = std::make_pair(std::move(NewObj), std::move(NewBuffer));
999-
}
1000-
const object::ObjectFile *DebugObj = SavedObject.first.release();
1001-
SavedObject.second.release();
1002-
1003971
StringMap<object::SectionRef> loadedSections;
1004-
// Use the original Object, not the DebugObject, as this is used for the
1005-
// RuntimeDyld::LoadedObjectInfo lookup.
1006972
for (const object::SectionRef &lSection : Object.sections()) {
1007973
auto sName = lSection.getName();
1008974
if (sName) {
@@ -1019,7 +985,9 @@ void registerRTDyldJITObject(const object::ObjectFile &Object,
1019985
return L.getSectionLoadAddress(search->second);
1020986
};
1021987

1022-
jl_register_jit_object(*DebugObj, getLoadAddress,
988+
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.
989+
jl_register_jit_object(DebugObject.getBinary() ? *DebugObject.getBinary() : Object,
990+
getLoadAddress,
1023991
#if defined(_OS_WINDOWS_) && defined(_CPU_X86_64_)
1024992
[MemMgr](void *p) { return lookupWriteAddressFor(MemMgr.get(), p); }
1025993
#else
@@ -1630,7 +1598,7 @@ JuliaOJIT::JuliaOJIT()
16301598
ES, std::move(ehRegistrar)));
16311599

16321600
ObjectLayer.addPlugin(std::make_unique<JLDebuginfoPlugin>());
1633-
ObjectLayer.addPlugin(std::make_unique<JLMemoryUsagePlugin>(total_size));
1601+
ObjectLayer.addPlugin(std::make_unique<JLMemoryUsagePlugin>(jit_bytes_size));
16341602
#else
16351603
ObjectLayer.setNotifyLoaded(
16361604
[this](orc::MaterializationResponsibility &MR,
@@ -2042,19 +2010,20 @@ std::string JuliaOJIT::getMangledName(const GlobalValue *GV)
20422010
return getMangledName(GV->getName());
20432011
}
20442012

2045-
#ifdef JL_USE_JITLINK
20462013
size_t JuliaOJIT::getTotalBytes() const
20472014
{
2048-
return total_size.load(std::memory_order_relaxed);
2015+
auto bytes = jit_bytes_size.load(std::memory_order_relaxed);
2016+
#ifndef JL_USE_JITLINK
2017+
size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT;
2018+
bytes += getRTDyldMemoryManagerTotalBytes(MemMgr.get());
2019+
#endif
2020+
return bytes;
20492021
}
2050-
#else
2051-
size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT;
20522022

2053-
size_t JuliaOJIT::getTotalBytes() const
2023+
void JuliaOJIT::addBytes(size_t bytes)
20542024
{
2055-
return getRTDyldMemoryManagerTotalBytes(MemMgr.get());
2025+
jit_bytes_size.fetch_add(bytes, std::memory_order_relaxed);
20562026
}
2057-
#endif
20582027

20592028
void JuliaOJIT::printTimers()
20602029
{
@@ -2339,3 +2308,9 @@ size_t jl_jit_total_bytes_impl(void)
23392308
{
23402309
return jl_ExecutionEngine->getTotalBytes();
23412310
}
2311+
2312+
// API for adding bytes to record being owned by the JIT
2313+
void jl_jit_add_bytes(size_t bytes)
2314+
{
2315+
jl_ExecutionEngine->addBytes(bytes);
2316+
}

src/jitlayers.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -559,6 +559,7 @@ class JuliaOJIT {
559559
TargetIRAnalysis getTargetIRAnalysis() const JL_NOTSAFEPOINT;
560560

561561
size_t getTotalBytes() const JL_NOTSAFEPOINT;
562+
void addBytes(size_t bytes) JL_NOTSAFEPOINT;
562563
void printTimers() JL_NOTSAFEPOINT;
563564

564565
jl_locked_stream &get_dump_emitted_mi_name_stream() JL_NOTSAFEPOINT {
@@ -605,10 +606,10 @@ class JuliaOJIT {
605606

606607
ResourcePool<orc::ThreadSafeContext, 0, std::queue<orc::ThreadSafeContext>> ContextPool;
607608

609+
std::atomic<size_t> jit_bytes_size{0};
608610
#ifndef JL_USE_JITLINK
609611
const std::shared_ptr<RTDyldMemoryManager> MemMgr;
610612
#else
611-
std::atomic<size_t> total_size{0};
612613
const std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr;
613614
#endif
614615
ObjLayerT ObjectLayer;

0 commit comments

Comments
 (0)