[llvm] [orc] Reduce memory usage from empty materialization info DenseMaps (PR #88167)

[Ben Langmuir via llvm-commits]

https://github.com/benlangmuir created https://github.com/llvm/llvm-project/pull/88167

Saves several MB of memory in larger applications after linking finishes by clearing DenseMap storage that is empty. This does not attempt to shrink partially full materialization infos. The assumption is that adding more after linking finishes is rare.

rdar://126145336

>From 7acc3e11e4009b495ce221818ce256c32611aa09 Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir at apple.com>
Date: Tue, 9 Apr 2024 10:37:30 -0700
Subject: [PATCH] [orc] Reduce memory usage from empty materialization info
 DenseMaps

Saves several MB of memory in larger applications after linking
finishes by clearing DenseMap storage that is empty. This does not
attempt to shrink partially full materialization infos. The assumption
is that adding more after linking finishes is rare.

rdar://126145336
---
 llvm/include/llvm/ExecutionEngine/Orc/Core.h |  4 ++++
 llvm/lib/ExecutionEngine/Orc/Core.cpp        | 23 ++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 09b2f6a10e35e0..7121b3fe762748 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -1333,6 +1333,10 @@ class JITDylib : public ThreadSafeRefCountedBase<JITDylib>,
 
   void unlinkMaterializationResponsibility(MaterializationResponsibility &MR);
 
+  /// Attempt to reduce memory usage from empty \c UnmaterializedInfos and
+  /// \c MaterializingInfos tables.
+  void shrinkMaterializationInfoMemory();
+
   ExecutionSession &ES;
   enum { Open, Closing, Closed } State = Open;
   std::mutex GeneratorsMutex;
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index a9d1998e093054..4841a2d8c4fdb3 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -1002,6 +1002,17 @@ void JITDylib::unlinkMaterializationResponsibility(
   });
 }
 
+void JITDylib::shrinkMaterializationInfoMemory() {
+  // DenseMap::erase never shrinks its storage; use clear to heuristically free
+  // memory since we may have long-lived JDs after linking is done.
+
+  if (UnmaterializedInfos.empty())
+    UnmaterializedInfos.clear();
+
+  if (MaterializingInfos.empty())
+    MaterializingInfos.clear();
+}
+
 void JITDylib::setLinkOrder(JITDylibSearchOrder NewLinkOrder,
                             bool LinkAgainstThisJITDylibFirst) {
   ES.runSessionLocked([&]() {
@@ -1112,6 +1123,8 @@ Error JITDylib::remove(const SymbolNameSet &Names) {
       Symbols.erase(SymI);
     }
 
+    shrinkMaterializationInfoMemory();
+
     return Error::success();
   });
 }
@@ -1313,6 +1326,8 @@ JITDylib::removeTracker(ResourceTracker &RT) {
     Symbols.erase(I);
   }
 
+  shrinkMaterializationInfoMemory();
+
   return Result;
 }
 
@@ -2675,6 +2690,8 @@ void ExecutionSession::OL_completeLookup(
             return true;
           });
 
+      JD.shrinkMaterializationInfoMemory();
+
       // Handle failure.
       if (Err) {
 
@@ -3118,6 +3135,8 @@ void ExecutionSession::IL_makeEDUReady(
 
     JD.MaterializingInfos.erase(MII);
   }
+
+  JD.shrinkMaterializationInfoMemory();
 }
 
 void ExecutionSession::IL_makeEDUEmitted(
@@ -3632,6 +3651,8 @@ ExecutionSession::IL_failSymbols(JITDylib &JD,
           ExtractFailedQueries(DepMI);
           DepJD.MaterializingInfos.erase(SymbolStringPtr(DepName));
         }
+
+        DepJD.shrinkMaterializationInfoMemory();
       }
 
       MI.DependantEDUs.clear();
@@ -3645,6 +3666,8 @@ ExecutionSession::IL_failSymbols(JITDylib &JD,
     JD.MaterializingInfos.erase(Name);
   }
 
+  JD.shrinkMaterializationInfoMemory();
+
 #ifdef EXPENSIVE_CHECKS
   verifySessionState("exiting ExecutionSession::IL_failSymbols");
 #endif