[BOLT] Use heuristic for matching split local functions

[aaupov]

Use known order of BOLT split function symbols: fragment symbols
immediately precede the parent fragment symbol.

Depends On: #89648

Test Plan: Added register-fragments-bolt-symbols.s

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Use known order of BOLT split function symbols: fragment symbols
immediately precede the parent fragment symbol.

Depends On: #89648

Test Plan: updated cdsplit-symbol-names.s

---

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

2 Files Affected:

• (modified) bolt/lib/Rewrite/RewriteInstance.cpp (+21)
• (modified) bolt/test/X86/cdsplit-symbol-names.s (+13-2)

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index d08c760f4da18d..8eb2e5a9d9120a 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -1510,6 +1510,26 @@ void RewriteInstance::registerFragments() {
       StopSymbol = *FSI;
 
     uint64_t ParentAddress{0};
+
+    // BOLT split fragment symbols are emitted just before the main function
+    // symbol.
+    for (ELFSymbolRef NextSymbol = Symbol; NextSymbol < StopSymbol;
+         NextSymbol.moveNext()) {
+      Expected<StringRef> NameOrError = Symbol.getName();
+      if (!NameOrError)
+        break;
+      StringRef Name = *NameOrError;
+      if (Name == ParentName) {
+        ParentAddress = cantFail(NextSymbol.getValue());
+        goto registerParent;
+      }
+      if (Name.starts_with(ParentName))
+        // With multi-way splitting, there are multiple fragments with different
+        // suffixes. Parent follows the last fragment.
+        continue;
+      break;
+    }
+
     // Iterate over local file symbols and check symbol names to match parent.
     for (ELFSymbolRef Symbol(FSI[-1]); Symbol < StopSymbol; Symbol.moveNext()) {
       if (cantFail(Symbol.getName()) == ParentName) {
@@ -1518,6 +1538,7 @@ void RewriteInstance::registerFragments() {
       }
     }
 
+registerParent:
     // No local parent is found, use global parent function.
     if (!ParentAddress)
       if (BinaryData *ParentBD = BC->getBinaryDataByName(ParentName))
diff --git a/bolt/test/X86/cdsplit-symbol-names.s b/bolt/test/X86/cdsplit-symbol-names.s
index e53863e22246d6..1d3fa91936af04 100644
--- a/bolt/test/X86/cdsplit-symbol-names.s
+++ b/bolt/test/X86/cdsplit-symbol-names.s
@@ -7,7 +7,7 @@
 # RUN: llvm-strip --strip-unneeded %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
 # RUN: llvm-bolt %t.exe -o %t.bolt --split-functions --split-strategy=cdsplit \
-# RUN:         --call-scale=2 --data=%t.fdata --reorder-blocks=ext-tsp
+# RUN:   --call-scale=2 --data=%t.fdata --reorder-blocks=ext-tsp --enable-bat
 # RUN: llvm-objdump --syms %t.bolt | FileCheck %s --check-prefix=CHECK-SYMS-WARM
 
 # CHECK-SYMS-WARM: 0000000000000000 l df *ABS* 0000000000000000 bolt-pseudo.o
@@ -16,8 +16,19 @@
 # CHECK-SYMS-WARM: .text.cold
 # CHECK-SYMS-WARM-SAME: dummy.cold
 
+# RUN: link_fdata %s %t.bolt %t.preagg PREAGG
+# PREAGG: B X:0 #chain.warm# 1 0
+# RUN: perf2bolt %t.bolt -p %t.preagg --pa -o %t.bat.fdata -w %t.bat.yaml -v=1 \
+# RUN:   | FileCheck %s --check-prefix=CHECK-REGISTER
+
+# CHECK-REGISTER: BOLT-INFO: marking chain.warm/1(*2) as a fragment of chain/2(*2)
+
         .text
-        .globl  chain
+        .type   chain, @function
+chain:
+        ret
+        .size   chain, .-chain
+
         .type   chain, @function
 chain:
         pushq   %rbp