[BOLT][DWARF] Change to process and write out TUs first then CUs in batches

To reduce memory footprint changed so that we process and write out TUs first,
reset DIEBuilder and process CUs. CUs are processed in buckets. First bucket
contains all the CUs with cross CU references. Rest processd one at a time.

clang-17 build in debug mode, by clang-17.
before
8:25.81 real, 834.37 user, 86.03 sys, 0 amem, 79525064 mmem
8:02.20 real, 820.46 user, 81.81 sys, 0 amem, 79501616 mmem
7:52.69 real, 802.01 user, 83.99 sys, 0 amem, 79534392 mmem

after
7:49.35 real, 822.04 user, 66.19 sys, 0 amem, 34934260 mmem
7:42.16 real, 825.46 user, 63.52 sys, 0 amem, 34951660 mmem
7:46.71 real, 821.11 user, 63.14 sys, 0 amem, 34981164 mmem

Reviewed By: maksfb

Differential Revision: https://reviews.llvm.org/D151909

Author: ayermolo

bolt/include/bolt/Core/DIEBuilder.h

@@ -60,6 +60,8 @@ class DIEBuilder {
     std::unordered_map<uint64_t, uint32_t> DIEIDMap;
   };
 
+  enum class ProcessingType { DWARF4TUs, DWARF5TUs, CUs };
+
 private:
   /// Contains information so that we we can update references in locexpr after
   /// we calculated all the final DIE offsets.
@@ -83,41 +85,54 @@ class DIEBuilder {
     DWARFAbbreviationDeclaration::AttributeSpec AttrSpec;
   };
 
-  /// A map of Units to Unit Index.
-  std::unordered_map<uint64_t, uint32_t> UnitIDMap;
-  /// A map of Type Units to Type DIEs.
-  std::unordered_map<DWARFUnit *, DIE *> TypeDIEMap;
-  std::vector<DWARFUnit *> DUList;
-  std::vector<DWARFUnitInfo> CloneUnitCtxMap;
-  std::vector<std::pair<DIEInfo *, AddrReferenceInfo>> AddrReferences;
+  struct State {
+    /// A map of Units to Unit Index.
+    std::unordered_map<uint64_t, uint32_t> UnitIDMap;
+    /// A map of Type Units to Type DIEs.
+    std::unordered_map<DWARFUnit *, DIE *> TypeDIEMap;
+    std::list<DWARFUnit *> DUList;
+    std::vector<DWARFUnitInfo> CloneUnitCtxMap;
+    std::vector<std::pair<DIEInfo *, AddrReferenceInfo>> AddrReferences;
+    std::vector<DWARFUnit *> DWARF4TUVector;
+    std::vector<DWARFUnit *> DWARF5TUVector;
+    std::vector<DWARFUnit *> DWARFCUVector;
+    std::vector<LocWithReference> LocWithReferencesToProcess;
+    BumpPtrAllocator DIEAlloc;
+    ProcessingType Type;
+  };
+
+  std::unique_ptr<State> BuilderState;
   FoldingSet<DIEAbbrev> AbbreviationsSet;
   std::vector<std::unique_ptr<DIEAbbrev>> Abbreviations;
-  std::vector<DWARFUnit *> DWARF4TUVector;
-  std::vector<LocWithReference> LocWithReferencesToProcess;
-  BumpPtrAllocator DIEAlloc;
+  DWARFContext *DwarfContext{nullptr};
+  bool IsDWO{false};
+  uint64_t UnitSize{0};
+  llvm::DenseSet<uint64_t> AllProcessed;
 
+  /// Returns current state of the DIEBuilder
+  State &getState() { return *BuilderState.get(); }
   /// Resolve the reference in DIE, if target is not loaded into IR,
   /// pre-allocate it. \p RefCU will be updated to the Unit specific by \p
   /// RefValue.
-  DWARFDie resolveDIEReference(const DWARFFormValue &RefValue,
-                               DWARFUnit *&RefCU,
-                               DWARFDebugInfoEntry &DwarfDebugInfoEntry,
-                               const std::vector<DWARFUnit *> &DUOffsetList);
+  DWARFDie resolveDIEReference(
+      const DWARFFormValue &RefValue,
+      const DWARFAbbreviationDeclaration::AttributeSpec AttrSpec,
+      DWARFUnit *&RefCU, DWARFDebugInfoEntry &DwarfDebugInfoEntry);
 
   /// Resolve the reference in DIE, if target is not loaded into IR,
   /// pre-allocate it. \p RefCU will be updated to the Unit specific by \p
   /// RefValue.
-  DWARFDie resolveDIEReference(const uint64_t ReffOffset, DWARFUnit *&RefCU,
-                               DWARFDebugInfoEntry &DwarfDebugInfoEntry,
-                               const std::vector<DWARFUnit *> &DUOffsetList);
+  DWARFDie resolveDIEReference(
+      const DWARFAbbreviationDeclaration::AttributeSpec AttrSpec,
+      const uint64_t ReffOffset, DWARFUnit *&RefCU,
+      DWARFDebugInfoEntry &DwarfDebugInfoEntry);
 
   /// Clone one attribute according to the format. \return the size of this
   /// attribute.
   void
   cloneAttribute(DIE &Die, const DWARFDie &InputDIE, DWARFUnit &U,
                  const DWARFFormValue &Val,
-                 const DWARFAbbreviationDeclaration::AttributeSpec AttrSpec,
-                 const std::vector<DWARFUnit *> &DUOffsetList);
+                 const DWARFAbbreviationDeclaration::AttributeSpec AttrSpec);
 
   /// Clone an attribute in string format.
   void cloneStringAttribute(
@@ -129,7 +144,7 @@ class DIEBuilder {
   void cloneDieReferenceAttribute(
       DIE &Die, const DWARFUnit &U, const DWARFDie &InputDIE,
       const DWARFAbbreviationDeclaration::AttributeSpec AttrSpec,
-      const DWARFFormValue &Val, const std::vector<DWARFUnit *> &DUOffsetList);
+      const DWARFFormValue &Val);
 
   /// Clone an attribute in block format.
   void cloneBlockAttribute(
@@ -175,23 +190,23 @@ class DIEBuilder {
   /// Update the Offset and Size of DIE.
   uint32_t computeDIEOffset(const DWARFUnit &CU, DIE &Die, uint32_t &CurOffset);
 
-  void registerUnit(DWARFUnit &DU);
+  void registerUnit(DWARFUnit &DU, bool NeedSort);
 
   /// \return the unique ID of \p U if it exists.
   std::optional<uint32_t> getUnitId(const DWARFUnit &DU);
 
   DWARFUnitInfo &getUnitInfo(uint32_t UnitId) {
-    return CloneUnitCtxMap[UnitId];
+    return getState().CloneUnitCtxMap[UnitId];
   }
 
   DIEInfo &getDIEInfo(uint32_t UnitId, uint32_t DIEId) {
-    if (CloneUnitCtxMap[UnitId].DieInfoVector.size() > DIEId)
-      return *CloneUnitCtxMap[UnitId].DieInfoVector[DIEId].get();
+    if (getState().CloneUnitCtxMap[UnitId].DieInfoVector.size() > DIEId)
+      return *getState().CloneUnitCtxMap[UnitId].DieInfoVector[DIEId].get();
 
     errs() << "BOLT-WARNING: [internal-dwarf-error]: The DIE is not allocated "
               "before looking up, some"
            << "unexpected corner cases happened.\n";
-    return *CloneUnitCtxMap[UnitId].DieInfoVector.front().get();
+    return *getState().CloneUnitCtxMap[UnitId].DieInfoVector.front().get();
   }
 
   std::optional<uint32_t> getAllocDIEId(const DWARFUnit &DU,
@@ -223,16 +238,28 @@ class DIEBuilder {
 
   /// Construct IR for \p DU. \p DUOffsetList specific the Unit in current
   /// Section.
-  void constructFromUnit(DWARFUnit &DU, std::vector<DWARFUnit *> &DUOffsetList);
+  void constructFromUnit(DWARFUnit &DU);
 
   /// Construct a DIE for \p DDie in \p U. \p DUOffsetList specific the Unit in
   /// current Section.
-  DIE *constructDIEFast(DWARFDie &DDie, DWARFUnit &U, uint32_t UnitId,
-                        std::vector<DWARFUnit *> &DUOffsetList);
+  DIE *constructDIEFast(DWARFDie &DDie, DWARFUnit &U, uint32_t UnitId);
 
 public:
   DIEBuilder(DWARFContext *DwarfContext, bool IsDWO = false);
 
+  /// Returns enum to what we are currently processing.
+  ProcessingType getCurrentProcessingState() { return getState().Type; }
+
+  /// Constructs IR for Type Units.
+  void buildTypeUnits(const bool Init = true);
+  /// Constructs IR for all the CUs.
+  void buildCompileUnits(const bool Init = true);
+  /// Constructs IR for CUs in a vector.
+  void buildCompileUnits(const std::vector<DWARFUnit *> &CUs);
+  /// Preventing implicit conversions.
+  template <class T> void buildCompileUnits(T) = delete;
+  void buildBoth();
+
   /// Returns DWARFUnitInfo for DWARFUnit
   DWARFUnitInfo &getUnitInfoByDwarfUnit(const DWARFUnit &DwarfUnit) {
     std::optional<uint32_t> UnitId = getUnitId(DwarfUnit);
@@ -247,16 +274,26 @@ class DIEBuilder {
     return Abbreviations;
   }
   DIE *getTypeDIE(DWARFUnit &DU) {
-    if (TypeDIEMap.count(&DU))
-      return TypeDIEMap[&DU];
+    if (getState().TypeDIEMap.count(&DU))
+      return getState().TypeDIEMap[&DU];
 
     errs() << "BOLT-ERROR: unable to find TypeUnit for Type Unit at offset 0x"
            << DU.getOffset() << "\n";
     return nullptr;
   }
 
-  std::vector<DWARFUnit *> getDWARF4TUVector() { return DWARF4TUVector; }
-  bool isEmpty() { return CloneUnitCtxMap.empty(); }
+  std::vector<DWARFUnit *> &getDWARF4TUVector() {
+    return getState().DWARF4TUVector;
+  }
+  std::vector<DWARFUnit *> &getDWARF5TUVector() {
+    return getState().DWARF5TUVector;
+  }
+  std::vector<DWARFUnit *> &getDWARFCUVector() {
+    return getState().DWARFCUVector;
+  }
+  /// Returns list of CUs for which IR was build.
+  std::list<DWARFUnit *> &getProcessedCUs() { return getState().DUList; }
+  bool isEmpty() { return getState().CloneUnitCtxMap.empty(); }
 
   DIE *getUnitDIEbyUnit(const DWARFUnit &DU) {
     const DWARFUnitInfo &U = getUnitInfoByDwarfUnit(DU);
@@ -272,37 +309,40 @@ class DIEBuilder {
   void finish();
 
   // Interface to edit DIE
-  template <class T> T *allocateDIEValue() { return new (DIEAlloc) T; }
+  template <class T> T *allocateDIEValue() {
+    return new (getState().DIEAlloc) T;
+  }
 
   DIEValueList::value_iterator addValue(DIEValueList *Die, const DIEValue &V) {
-    return Die->addValue(DIEAlloc, V);
+    return Die->addValue(getState().DIEAlloc, V);
   }
 
   template <class T>
   DIEValueList::value_iterator addValue(DIEValueList *Die,
                                         dwarf::Attribute Attribute,
                                         dwarf::Form Form, T &&Value) {
-    return Die->addValue(DIEAlloc, Attribute, Form, std::forward<T>(Value));
+    return Die->addValue(getState().DIEAlloc, Attribute, Form,
+                         std::forward<T>(Value));
   }
 
   template <class T>
   bool replaceValue(DIEValueList *Die, dwarf::Attribute Attribute,
                     dwarf::Form Form, T &&NewValue) {
-    return Die->replaceValue(DIEAlloc, Attribute, Form,
+    return Die->replaceValue(getState().DIEAlloc, Attribute, Form,
                              std::forward<T>(NewValue));
   }
 
   template <class T>
   bool replaceValue(DIEValueList *Die, dwarf::Attribute Attribute,
                     dwarf::Attribute NewAttribute, dwarf::Form Form,
                     T &&NewValue) {
-    return Die->replaceValue(DIEAlloc, Attribute, NewAttribute, Form,
+    return Die->replaceValue(getState().DIEAlloc, Attribute, NewAttribute, Form,
                              std::forward<T>(NewValue));
   }
 
   bool replaceValue(DIEValueList *Die, dwarf::Attribute Attribute,
                     dwarf::Form Form, DIEValue &NewValue) {
-    return Die->replaceValue(DIEAlloc, Attribute, Form, NewValue);
+    return Die->replaceValue(getState().DIEAlloc, Attribute, Form, NewValue);
   }
 
   template <class T>

bolt/include/bolt/Core/DebugData.h

@@ -22,6 +22,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 #include <map>
+#include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
@@ -315,17 +316,12 @@ class DebugAddrWriter {
   /// Adds {\p Address, \p Index} to \p CU.
   void addIndexAddress(uint64_t Address, uint32_t Index, DWARFUnit &CU);
 
-  /// Creates consolidated .debug_addr section, and builds DWOID to offset map.
-  virtual AddressSectionBuffer finalize();
+  /// Write out entries in to .debug_addr section for CUs.
+  virtual void update(DIEBuilder &DIEBlder, DWARFUnit &CUs);
 
-  /// Given DWARFUnit \p Unit returns offset of this CU in to .debug_addr
-  /// section.
-  virtual uint64_t getOffset(DWARFUnit &Unit);
-
-  /// Returns True if CU exists in the DebugAddrWriter.
-  bool doesCUExist(DWARFUnit &Unit) {
-    return DWOIdToOffsetMap.count(getCUID(Unit)) > 0;
-  }
+  /// Return buffer with all the entries in .debug_addr already writen out using
+  /// update(...).
+  virtual AddressSectionBuffer &finalize() { return *Buffer; }
 
   /// Returns False if .debug_addr section was created..
   bool isInitialized() const { return !AddressMaps.empty(); }
@@ -387,22 +383,21 @@ class DebugAddrWriter {
   BinaryContext *BC;
   /// Maps DWOID to AddressForDWOCU.
   std::unordered_map<uint64_t, AddressForDWOCU> AddressMaps;
-  /// Maps DWOID to offset within .debug_addr section.
-  std::unordered_map<uint64_t, uint64_t> DWOIdToOffsetMap;
   /// Mutex used for parallel processing of debug info.
   std::mutex WriterMutex;
+  std::unique_ptr<AddressSectionBuffer> Buffer;
+  std::unique_ptr<raw_svector_ostream> AddressStream;
+  /// Used to track sections that were not modified so that they can be re-used.
+  DenseMap<uint64_t, uint64_t> UnmodifiedAddressOffsets;
 };
 
 class DebugAddrWriterDwarf5 : public DebugAddrWriter {
 public:
   DebugAddrWriterDwarf5() = delete;
   DebugAddrWriterDwarf5(BinaryContext *BC) : DebugAddrWriter(BC) {}
 
-  /// Creates consolidated .debug_addr section, and builds DWOID to offset map.
-  AddressSectionBuffer finalize() override;
-  /// Given DWARFUnit \p Unit returns offset of this CU in to .debug_addr
-  /// section.
-  uint64_t getOffset(DWARFUnit &Unit) override;
+  /// Write out entries in to .debug_addr section for CUs.
+  virtual void update(DIEBuilder &DIEBlder, DWARFUnit &CUs) override;
 
 protected:
   /// Given DWARFUnit \p Unit returns either DWO ID or it's offset within
@@ -433,7 +428,7 @@ class DebugStrOffsetsWriter {
   void updateAddressMap(uint32_t Index, uint32_t Address);
 
   /// Writes out current sections entry into .debug_str_offsets.
-  void finalizeSection(DWARFUnit &Unit);
+  void finalizeSection(DWARFUnit &Unit, DIEBuilder &DIEBldr);
 
   /// Returns False if no strings were added to .debug_str.
   bool isFinalized() const { return !StrOffsetsBuffer->empty(); }
@@ -447,7 +442,7 @@ class DebugStrOffsetsWriter {
   std::unique_ptr<DebugStrOffsetsBufferVector> StrOffsetsBuffer;
   std::unique_ptr<raw_svector_ostream> StrOffsetsStream;
   std::map<uint32_t, uint32_t> IndexToAddressMap;
-  DenseSet<uint64_t> ProcessedBaseOffsets;
+  std::unordered_map<uint64_t, uint64_t> ProcessedBaseOffsets;
   // Section size not including header.
   uint32_t CurrentSectionSize{0};
   bool StrOffsetSectionWasModified = false;

bolt/include/bolt/Rewrite/DWARFRewriter.h

@@ -105,6 +105,9 @@ class DWARFRewriter {
   /// DWARFLegacy is all DWARF versions before DWARF 5.
   enum class DWARFVersion { DWARFLegacy, DWARF5 };
 
+  /// Used to track last CU offset for GDB Index.
+  uint32_t CUOffset{0};
+
   /// Update debug info for all DIEs in \p Unit.
   void updateUnitDebugInfo(DWARFUnit &Unit, DIEBuilder &DIEBldr,
                            DebugLocWriter &DebugLocWriter,
@@ -128,8 +131,17 @@ class DWARFRewriter {
   std::unique_ptr<DebugBufferVector>
   makeFinalLocListsSection(DWARFVersion Version);
 
+  /// Finalize type sections in the main binary.
+  CUOffsetMap finalizeTypeSections(DIEBuilder &DIEBlder, DIEStreamer &Streamer);
+
+  /// Process and write out CUs that are passsed in.
+  void finalizeCompileUnits(DIEBuilder &DIEBlder, DIEStreamer &Streamer,
+                            CUOffsetMap &CUMap,
+                            const std::list<DWARFUnit *> &CUs);
+
   /// Finalize debug sections in the main binary.
-  CUOffsetMap finalizeDebugSections(DIEBuilder &DIEBlder);
+  void finalizeDebugSections(DIEBuilder &DIEBlder, DIEStreamer &Streamer,
+                             raw_svector_ostream &ObjOS, CUOffsetMap &CUMap);
 
   /// Patches the binary for DWARF address ranges (e.g. in functions and lexical
   /// blocks) to be updated.