[BOLT][NFC] Track fragment relationships using EquivalenceClasses

[aaupov]

Three-way splitting can create references between split fragments (warm
to cold or vice versa) that are not handled by
isChildOf/isParentOf/isChildOrParentOf. Generalize fragment
relationships to allow checking if two functions belong to one group,
potentially in presence of ICF which can join multiple groups.

Test Plan: NFC for existing tests

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Three-way splitting can create references between split fragments (warm
to cold or vice versa) that are not handled by
isChildOf/isParentOf/isChildOrParentOf. Generalize fragment
relationships to allow checking if two functions belong to one group,
potentially in presence of ICF which can join multiple groups.

Test Plan: NFC for existing tests

---

Full diff: https://github.com/llvm/llvm-project/pull/99979.diff

4 Files Affected:

• (modified) bolt/include/bolt/Core/BinaryContext.h (+14-1)
• (modified) bolt/include/bolt/Core/BinaryFunction.h (-10)
• (modified) bolt/lib/Core/BinaryContext.cpp (+5-4)
• (modified) bolt/lib/Core/Exceptions.cpp (+1-1)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index b3cf9f834cc08..5fb32a1ea6784 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -23,6 +23,7 @@
 #include "bolt/RuntimeLibs/RuntimeLibrary.h"
 #include "llvm/ADT/AddressRanges.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/BinaryFormat/Dwarf.h"
@@ -241,6 +242,10 @@ class BinaryContext {
   /// Function fragments to skip.
   std::unordered_set<BinaryFunction *> FragmentsToSkip;
 
+  /// Fragment equivalence classes to query belonging to the same "family" in
+  /// presence of multiple fragments/multiple parents.
+  EquivalenceClasses<const BinaryFunction *> FragmentClasses;
+
   /// The runtime library.
   std::unique_ptr<RuntimeLibrary> RtLibrary;
 
@@ -1032,7 +1037,15 @@ class BinaryContext {
   ///   fragment_name == parent_name.cold(.\d+)?
   /// True if the Function is registered, false if the check failed.
   bool registerFragment(BinaryFunction &TargetFunction,
-                        BinaryFunction &Function) const;
+                        BinaryFunction &Function);
+
+  /// Return true if two functions belong to the same "family": are fragments
+  /// of one another, or fragments of the same parent, or transitively fragment-
+  /// related.
+  bool areRelatedFragments(const BinaryFunction *LHS,
+                           const BinaryFunction *RHS) const {
+    return FragmentClasses.isEquivalent(LHS, RHS);
+  }
 
   /// Add interprocedural reference for \p Function to \p Address
   void addInterproceduralReference(BinaryFunction *Function, uint64_t Address) {
diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index da3fc433b7a3b..24c7db2f5d69c 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -1793,11 +1793,6 @@ class BinaryFunction {
     return ParentFragments.contains(&Other);
   }
 
-  /// Returns if this function is a parent of \p Other function.
-  bool isParentOf(const BinaryFunction &Other) const {
-    return Fragments.contains(&Other);
-  }
-
   /// Return the child fragment form parent function
   iterator_range<FragmentsSetTy::const_iterator> getFragments() const {
     return iterator_range<FragmentsSetTy::const_iterator>(Fragments.begin(),
@@ -1807,11 +1802,6 @@ class BinaryFunction {
   /// Return the parent function for split function fragments.
   FragmentsSetTy *getParentFragments() { return &ParentFragments; }
 
-  /// Returns if this function is a parent or child of \p Other function.
-  bool isParentOrChildOf(const BinaryFunction &Other) const {
-    return isChildOf(Other) || isParentOf(Other);
-  }
-
   /// Set the profile data for the number of times the function was called.
   BinaryFunction &setExecutionCount(uint64_t Count) {
     ExecutionCount = Count;
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 83a5484f097ef..874cdd26ce6ea 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -646,7 +646,7 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
     const BinaryFunction *TargetBF = getBinaryFunctionContainingAddress(Value);
     const bool DoesBelongToFunction =
         BF.containsAddress(Value) ||
-        (TargetBF && TargetBF->isParentOrChildOf(BF));
+        (TargetBF && areRelatedFragments(TargetBF, &BF));
     if (!DoesBelongToFunction) {
       LLVM_DEBUG({
         if (!BF.containsAddress(Value)) {
@@ -841,7 +841,7 @@ BinaryContext::getOrCreateJumpTable(BinaryFunction &Function, uint64_t Address,
     // Prevent associating a jump table to a specific fragment twice.
     // This simple check arises from the assumption: no more than 2 fragments.
     if (JT->Parents.size() == 1 && JT->Parents[0] != &Function) {
-      assert(JT->Parents[0]->isParentOrChildOf(Function) &&
+      assert(areRelatedFragments(JT->Parents[0], &Function) &&
              "cannot re-use jump table of a different function");
       // Duplicate the entry for the parent function for easy access
       JT->Parents.push_back(&Function);
@@ -1209,12 +1209,13 @@ void BinaryContext::generateSymbolHashes() {
 }
 
 bool BinaryContext::registerFragment(BinaryFunction &TargetFunction,
-                                     BinaryFunction &Function) const {
+                                     BinaryFunction &Function) {
   assert(TargetFunction.isFragment() && "TargetFunction must be a fragment");
   if (TargetFunction.isChildOf(Function))
     return true;
   TargetFunction.addParentFragment(Function);
   Function.addFragment(TargetFunction);
+  FragmentClasses.unionSets(&TargetFunction, &Function);
   if (!HasRelocations) {
     TargetFunction.setSimple(false);
     Function.setSimple(false);
@@ -1336,7 +1337,7 @@ void BinaryContext::processInterproceduralReferences() {
 
     if (TargetFunction) {
       if (TargetFunction->isFragment() &&
-          !TargetFunction->isChildOf(Function)) {
+          !areRelatedFragments(TargetFunction, &Function)) {
         this->errs()
             << "BOLT-WARNING: interprocedural reference between unrelated "
                "fragments: "
diff --git a/bolt/lib/Core/Exceptions.cpp b/bolt/lib/Core/Exceptions.cpp
index 82bddf76d5b8c..6a46a49a983d8 100644
--- a/bolt/lib/Core/Exceptions.cpp
+++ b/bolt/lib/Core/Exceptions.cpp
@@ -207,7 +207,7 @@ Error BinaryFunction::parseLSDA(ArrayRef<uint8_t> LSDASectionData,
                "BOLT-ERROR: cannot find landing pad fragment");
         BC.addInterproceduralReference(this, Fragment->getAddress());
         BC.processInterproceduralReferences();
-        assert(isParentOrChildOf(*Fragment) &&
+        assert(BC.areRelatedFragments(this, Fragment) &&
                "BOLT-ERROR: cannot have landing pads in different functions");
         setHasIndirectTargetToSplitFragment(true);
         BC.addFragmentsToSkip(this);

[rafaelauler]

How expensive is this data structure, considering we can work with 1M+ functions?

[aaupov]

It's a set that only keeps a copy of entry (a pointer in our case):

llvm-project/llvm/include/llvm/ADT/EquivalenceClasses.h

Lines 26 to 35 in 73ffeea

	/// EquivalenceClasses - This represents a collection of equivalence classes and
	/// supports three efficient operations: insert an element into a class of its
	/// own, union two classes, and find the class for a given element. In
	/// addition to these modification methods, it is possible to iterate over all
	/// of the equivalence classes and all of the elements in a class.
	///
	/// This implementation is an efficient implementation that only stores one copy
	/// of the element being indexed per entry in the set, and allows any arbitrary
	/// type to be indexed (as long as it can be ordered with operator< or a
	/// comparator is provided).

Moreover, we will only insert fragments into it, so it shouldn't be too expensive even for the largest binaries (we split up to 10s of thousands of functions).

[rafaelauler]

LG