Skip to content

Commit 16a7ad9

Browse files
vtjnashKristofferC
authored and
KristofferC
committed
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). (cherry picked from commit fe597c1)
1 parent bcb33ed commit 16a7ad9

File tree

5 files changed

+94
-82
lines changed

5 files changed

+94
-82
lines changed

src/debug-registry.h

+15-7
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

+42-15
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

+1
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

+34-59
Original file line numberDiff line numberDiff line change
@@ -800,22 +800,19 @@ struct JITObjectInfo {
800800
class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
801801
std::mutex PluginMutex;
802802
std::map<MaterializationResponsibility *, std::unique_ptr<JITObjectInfo>> PendingObjs;
803-
// Resources from distinct `MaterializationResponsibility`s can get merged
804-
// after emission, so we can have multiple debug objects per resource key.
805-
std::map<ResourceKey, SmallVector<std::unique_ptr<JITObjectInfo>, 0>> RegisteredObjs;
806803

807804
public:
808805
void notifyMaterializing(MaterializationResponsibility &MR, jitlink::LinkGraph &G,
809806
jitlink::JITLinkContext &Ctx,
810807
MemoryBufferRef InputObject) override
811808
{
812-
// Keeping around a full copy of the input object file (and re-parsing it) is
813-
// wasteful, but for now, this lets us reuse the existing debuginfo.cpp code.
814-
// Should look into just directly pulling out all the information required in
815-
// a JITLink pass and just keeping the required tables/DWARF sections around
816-
// (perhaps using the LLVM DebuggerSupportPlugin as a reference).
817809
auto NewBuffer =
818810
MemoryBuffer::getMemBufferCopy(InputObject.getBuffer(), G.getName());
811+
// Re-parsing the InputObject is wasteful, but for now, this lets us
812+
// reuse the existing debuginfo.cpp code. Should look into just
813+
// directly pulling out all the information required in a JITLink pass
814+
// and just keeping the required tables/DWARF sections around (perhaps
815+
// using the LLVM DebuggerSupportPlugin as a reference).
819816
auto NewObj =
820817
cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef()));
821818

@@ -849,13 +846,8 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
849846
};
850847

851848
jl_register_jit_object(*NewInfo->Object, getLoadAddress, nullptr);
852-
}
853-
854-
cantFail(MR.withResourceKeyDo([&](ResourceKey K) {
855-
std::lock_guard<std::mutex> lock(PluginMutex);
856-
RegisteredObjs[K].push_back(std::move(PendingObjs[&MR]));
857849
PendingObjs.erase(&MR);
858-
}));
850+
}
859851

860852
return Error::success();
861853
}
@@ -866,32 +858,23 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
866858
PendingObjs.erase(&MR);
867859
return Error::success();
868860
}
861+
869862
#if JL_LLVM_VERSION >= 160000
870863
Error notifyRemovingResources(JITDylib &JD, orc::ResourceKey K) override
871864
#else
872-
Error notifyRemovingResources(ResourceKey K) override
865+
Error notifyRemovingResources(orc::ResourceKey K) override
873866
#endif
874867
{
875-
std::lock_guard<std::mutex> lock(PluginMutex);
876-
RegisteredObjs.erase(K);
877-
// TODO: If we ever unload code, need to notify debuginfo registry.
878868
return Error::success();
879869
}
880870

881871
#if JL_LLVM_VERSION >= 160000
882-
void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) override
872+
void notifyTransferringResources(JITDylib &JD, orc::ResourceKey DstKey,
873+
orc::ResourceKey SrcKey) override {}
883874
#else
884-
void notifyTransferringResources(ResourceKey DstKey, ResourceKey SrcKey) override
875+
void notifyTransferringResources(orc::ResourceKey DstKey,
876+
orc::ResourceKey SrcKey) override {}
885877
#endif
886-
{
887-
std::lock_guard<std::mutex> lock(PluginMutex);
888-
auto SrcIt = RegisteredObjs.find(SrcKey);
889-
if (SrcIt != RegisteredObjs.end()) {
890-
for (std::unique_ptr<JITObjectInfo> &Info : SrcIt->second)
891-
RegisteredObjs[DstKey].push_back(std::move(Info));
892-
RegisteredObjs.erase(SrcIt);
893-
}
894-
}
895878

896879
void modifyPassConfig(MaterializationResponsibility &MR, jitlink::LinkGraph &,
897880
jitlink::PassConfiguration &PassConfig) override
@@ -931,12 +914,12 @@ class JLDebuginfoPlugin : public ObjectLinkingLayer::Plugin {
931914

932915
class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin {
933916
private:
934-
std::atomic<size_t> &total_size;
917+
std::atomic<size_t> &jit_bytes_size;
935918

936919
public:
937920

938-
JLMemoryUsagePlugin(std::atomic<size_t> &total_size)
939-
: total_size(total_size) {}
921+
JLMemoryUsagePlugin(std::atomic<size_t> &jit_bytes_size)
922+
: jit_bytes_size(jit_bytes_size) {}
940923

941924
Error notifyFailed(orc::MaterializationResponsibility &MR) override {
942925
return Error::success();
@@ -985,7 +968,7 @@ class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin {
985968
}
986969
(void) code_size;
987970
(void) data_size;
988-
this->total_size.fetch_add(graph_size, std::memory_order_relaxed);
971+
this->jit_bytes_size.fetch_add(graph_size, std::memory_order_relaxed);
989972
jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, graph_size);
990973
jl_timing_counter_inc(JL_TIMING_COUNTER_JITCodeSize, code_size);
991974
jl_timing_counter_inc(JL_TIMING_COUNTER_JITDataSize, data_size);
@@ -1101,24 +1084,7 @@ void registerRTDyldJITObject(const object::ObjectFile &Object,
11011084
const RuntimeDyld::LoadedObjectInfo &L,
11021085
const std::shared_ptr<RTDyldMemoryManager> &MemMgr)
11031086
{
1104-
auto SavedObject = L.getObjectForDebug(Object).takeBinary();
1105-
// If the debug object is unavailable, save (a copy of) the original object
1106-
// for our backtraces.
1107-
// This copy seems unfortunate, but there doesn't seem to be a way to take
1108-
// ownership of the original buffer.
1109-
if (!SavedObject.first) {
1110-
auto NewBuffer =
1111-
MemoryBuffer::getMemBufferCopy(Object.getData(), Object.getFileName());
1112-
auto NewObj =
1113-
cantFail(object::ObjectFile::createObjectFile(NewBuffer->getMemBufferRef()));
1114-
SavedObject = std::make_pair(std::move(NewObj), std::move(NewBuffer));
1115-
}
1116-
const object::ObjectFile *DebugObj = SavedObject.first.release();
1117-
SavedObject.second.release();
1118-
11191087
StringMap<object::SectionRef> loadedSections;
1120-
// Use the original Object, not the DebugObject, as this is used for the
1121-
// RuntimeDyld::LoadedObjectInfo lookup.
11221088
for (const object::SectionRef &lSection : Object.sections()) {
11231089
auto sName = lSection.getName();
11241090
if (sName) {
@@ -1135,7 +1101,9 @@ void registerRTDyldJITObject(const object::ObjectFile &Object,
11351101
return L.getSectionLoadAddress(search->second);
11361102
};
11371103

1138-
jl_register_jit_object(*DebugObj, getLoadAddress,
1104+
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.
1105+
jl_register_jit_object(DebugObject.getBinary() ? *DebugObject.getBinary() : Object,
1106+
getLoadAddress,
11391107
#if defined(_OS_WINDOWS_) && defined(_CPU_X86_64_)
11401108
[MemMgr](void *p) { return lookupWriteAddressFor(MemMgr.get(), p); }
11411109
#else
@@ -1737,7 +1705,7 @@ JuliaOJIT::JuliaOJIT()
17371705
ES, std::move(ehRegistrar)));
17381706

17391707
ObjectLayer.addPlugin(std::make_unique<JLDebuginfoPlugin>());
1740-
ObjectLayer.addPlugin(std::make_unique<JLMemoryUsagePlugin>(total_size));
1708+
ObjectLayer.addPlugin(std::make_unique<JLMemoryUsagePlugin>(jit_bytes_size));
17411709
#else
17421710
ObjectLayer.setNotifyLoaded(
17431711
[this](orc::MaterializationResponsibility &MR,
@@ -2058,19 +2026,20 @@ std::string JuliaOJIT::getMangledName(const GlobalValue *GV)
20582026
return getMangledName(GV->getName());
20592027
}
20602028

2061-
#ifdef JL_USE_JITLINK
20622029
size_t JuliaOJIT::getTotalBytes() const
20632030
{
2064-
return total_size.load(std::memory_order_relaxed);
2031+
auto bytes = jit_bytes_size.load(std::memory_order_relaxed);
2032+
#ifndef JL_USE_JITLINK
2033+
size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT;
2034+
bytes += getRTDyldMemoryManagerTotalBytes(MemMgr.get());
2035+
#endif
2036+
return bytes;
20652037
}
2066-
#else
2067-
size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm) JL_NOTSAFEPOINT;
20682038

2069-
size_t JuliaOJIT::getTotalBytes() const
2039+
void JuliaOJIT::addBytes(size_t bytes)
20702040
{
2071-
return getRTDyldMemoryManagerTotalBytes(MemMgr.get());
2041+
jit_bytes_size.fetch_add(bytes, std::memory_order_relaxed);
20722042
}
2073-
#endif
20742043

20752044
void JuliaOJIT::printTimers()
20762045
{
@@ -2348,3 +2317,9 @@ size_t jl_jit_total_bytes_impl(void)
23482317
{
23492318
return jl_ExecutionEngine->getTotalBytes();
23502319
}
2320+
2321+
// API for adding bytes to record being owned by the JIT
2322+
void jl_jit_add_bytes(size_t bytes)
2323+
{
2324+
jl_ExecutionEngine->addBytes(bytes);
2325+
}

src/jitlayers.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -553,6 +553,7 @@ class JuliaOJIT {
553553
TargetIRAnalysis getTargetIRAnalysis() const JL_NOTSAFEPOINT;
554554

555555
size_t getTotalBytes() const JL_NOTSAFEPOINT;
556+
void addBytes(size_t bytes) JL_NOTSAFEPOINT;
556557
void printTimers() JL_NOTSAFEPOINT;
557558

558559
jl_locked_stream &get_dump_emitted_mi_name_stream() JL_NOTSAFEPOINT {
@@ -597,10 +598,10 @@ class JuliaOJIT {
597598

598599
ResourcePool<orc::ThreadSafeContext, 0, std::queue<orc::ThreadSafeContext>> ContextPool;
599600

601+
std::atomic<size_t> jit_bytes_size{0};
600602
#ifndef JL_USE_JITLINK
601603
const std::shared_ptr<RTDyldMemoryManager> MemMgr;
602604
#else
603-
std::atomic<size_t> total_size{0};
604605
const std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr;
605606
#endif
606607
ObjLayerT ObjectLayer;

0 commit comments

Comments
 (0)