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

[via llvm-commits]

Author: Ben Langmuir
Date: 2024-04-09T13:01:00-07:00
New Revision: aacb8985f734e7e11fc94947f7bc348d1d39f7af

URL: https://github.com/llvm/llvm-project/commit/aacb8985f734e7e11fc94947f7bc348d1d39f7af
DIFF: https://github.com/llvm/llvm-project/commit/aacb8985f734e7e11fc94947f7bc348d1d39f7af.diff

LOG: [orc] Reduce memory usage from empty materialization info DenseMaps (#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

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Core.h
    llvm/lib/ExecutionEngine/Orc/Core.cpp

Removed: 
    


################################################################################
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