[llvm] [RFC][BOLT] Add a new parallel DWARF processing(2/2) (PR #197859)

[Alexander Yermolovich via llvm-commits]

================
@@ -539,43 +607,92 @@ static void emitDWOBuilder(const std::string &DWOName,
                          StrOffstsWriter, StrWriter, TempRangesSectionWriter);
 }
 
-using DWARFUnitVec = std::vector<DWARFUnit *>;
-using CUPartitionVector = std::vector<DWARFUnitVec>;
-/// Partitions CUs in to buckets. Bucket size is controlled by
-/// cu-processing-batch-size. All the CUs that have cross CU reference reference
-/// as a source are put in to the same initial bucket.
-static CUPartitionVector partitionCUs(DWARFContext &DwCtx) {
-  CUPartitionVector Vec(2);
-  unsigned Counter = 0;
-  const DWARFDebugAbbrev *Abbr = DwCtx.getDebugAbbrev();
-  for (std::unique_ptr<DWARFUnit> &CU : DwCtx.compile_units()) {
-    Expected<const DWARFAbbreviationDeclarationSet *> AbbrDeclSet =
-        Abbr->getAbbreviationDeclarationSet(CU->getAbbreviationsOffset());
-    if (!AbbrDeclSet) {
-      consumeError(AbbrDeclSet.takeError());
-      return Vec;
-    }
-    bool CrossCURefFound = false;
-    for (const DWARFAbbreviationDeclaration &Decl : *AbbrDeclSet.get()) {
-      for (const DWARFAbbreviationDeclaration::AttributeSpec &Attr :
-           Decl.attributes()) {
-        if (Attr.Form == dwarf::DW_FORM_ref_addr) {
-          CrossCURefFound = true;
+static std::vector<std::vector<DWARFUnit *>> partitionCUs(DWARFContext &DwCtx) {
+  SmallVector<DWARFUnit *, 0> AllCUs;
+  for (auto &CU : DwCtx.compile_units())
+    AllCUs.push_back(CU.get());
+  if (AllCUs.empty())
+    return {};
+  auto FindCuForOffset = [&](uint64_t Offset) -> DWARFUnit * {
+    auto *It =
+        llvm::upper_bound(AllCUs, Offset, [](uint64_t Off, DWARFUnit *U) {
+          return Off < U->getOffset();
+        });
+    if (It == AllCUs.begin())
+      return nullptr;
+    DWARFUnit *TargetCU = *--It;
+    // Ensure offset falls within TargetCU's range.
+    if (Offset >= TargetCU->getNextUnitOffset())
+      return nullptr;
+    return TargetCU;
+  };
+
+  DenseSet<DWARFUnit *> CrossRefSet;
+  EquivalenceClasses<DWARFUnit *> EC;
+  for (DWARFUnit *CU : AllCUs) {
+    const DWARFAbbreviationDeclarationSet *AbbrevSet = CU->getAbbreviations();
+    if (!AbbrevSet)
+      continue;
+    SmallDenseSet<const DWARFAbbreviationDeclaration *, 4> RefAddrAbbrevs;
+    for (const auto &Decl : *AbbrevSet)
+      for (const auto &Spec : Decl.attributes())
+        if (Spec.Form == dwarf::DW_FORM_ref_addr) {
+          RefAddrAbbrevs.insert(&Decl);
           break;
         }
+    if (RefAddrAbbrevs.empty())
+      continue;
+    // Track CUs involved in cross-CU references via DW_FORM_ref_addr.
+    for (const DWARFDebugInfoEntry &Entry : CU->dies()) {
----------------
ayermolo wrote:

This will extract all the dies into a vector. This whole thing is already a memory monster, probably best not to add to it.
Can you change this to be similar to DIEBuilder::constructFromUnit.

https://github.com/llvm/llvm-project/pull/197859