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

[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()) {
+      DWARFDie Die(CU, &Entry);
+      const DWARFAbbreviationDeclaration *Abbrev =
+          Die.getAbbreviationDeclarationPtr();
+      if (!Abbrev || !RefAddrAbbrevs.count(Abbrev))
+        continue;
+      for (const DWARFAttribute &Attr : Die.attributes()) {
+        if (Attr.Value.getForm() != dwarf::DW_FORM_ref_addr)
+          continue;
+        auto OptRef = Attr.Value.getAsDebugInfoReference();
+        if (!OptRef)
+          continue;
+        DWARFUnit *TargetCU = FindCuForOffset(*OptRef);
+        if (!TargetCU)
+          continue;
+        if (CrossRefSet.insert(CU).second)
+          EC.insert(CU);
+        if (CrossRefSet.insert(TargetCU).second)
+          EC.insert(TargetCU);
+        EC.unionSets(CU, TargetCU);
       }
-      if (CrossCURefFound)
-        break;
     }
-    if (CrossCURefFound) {
-      Vec[0].push_back(CU.get());
-    } else {
-      ++Counter;
-      Vec.back().push_back(CU.get());
+  }
+
+  DenseMap<DWARFUnit *, std::vector<DWARFUnit *>> MembersByLeader;
+  for (DWARFUnit *CU : AllCUs) {
+    if (!CrossRefSet.count(CU))
+      continue;
+    MembersByLeader[EC.getLeaderValue(CU)].push_back(CU);
+  }
+
+  std::vector<DWARFUnit *> Leaders;
+  Leaders.reserve(MembersByLeader.size());
+  for (auto &[Leader, Members] : MembersByLeader)
+    Leaders.push_back(Leader);
+  llvm::sort(Leaders, [&](DWARFUnit *A, DWARFUnit *B) {
+    return MembersByLeader[A].front()->getOffset() <
+           MembersByLeader[B].front()->getOffset();
+  });
+
+  // Emit cross-ref buckets, then singleton non-cross-ref CUs.
+  std::vector<std::vector<DWARFUnit *>> Vec;
----------------
Thrrreeee wrote:

`dwarf-preserve-order` is still useful. `cu-processing-batch-size` is now dead. After this PR merged, we will use a new pr to clean it, because many testcases need to be changed.

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