[llvm] [BOLT] Allow pass-through blocks in YAMLProfileReader (PR #91828)

[Amir Ayupov via llvm-commits]

https://github.com/aaupov updated https://github.com/llvm/llvm-project/pull/91828

>From 4450f9de369e685c3b7a960edcd2dff8e1e3f1c8 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Fri, 10 May 2024 16:37:31 -0700
Subject: [PATCH 1/2] [BOLT] Allow pass-through blocks in YAMLProfileReader

Align YAMLProfileReader with DataReader in handling non-matching edges.

ccb99dd1261a9b40a3141fc7a63437dab5770cb2 introduced a compatibility
feature to allow pass-through blocks in non-matching edges:
match the profile edge A -> C to the CFG with blocks A -> B -> C.

A similar case arises when profiling BOLTed binaries with block B
removed.

Test Plan: Added profile-passthrough-block.test
---
 bolt/lib/Profile/YAMLProfileReader.cpp       | 30 ++++++---
 bolt/test/X86/profile-passthrough-block.test | 66 ++++++++++++++++++++
 2 files changed, 87 insertions(+), 9 deletions(-)
 create mode 100644 bolt/test/X86/profile-passthrough-block.test

diff --git a/bolt/lib/Profile/YAMLProfileReader.cpp b/bolt/lib/Profile/YAMLProfileReader.cpp
index e4673f6e3c301..156a0efd89cd9 100644
--- a/bolt/lib/Profile/YAMLProfileReader.cpp
+++ b/bolt/lib/Profile/YAMLProfileReader.cpp
@@ -218,17 +218,29 @@ bool YAMLProfileReader::parseFunctionProfile(
         continue;
       }
 
-      BinaryBasicBlock &SuccessorBB = *Order[YamlSI.Index];
-      if (!BB.getSuccessor(SuccessorBB.getLabel())) {
-        if (opts::Verbosity >= 1)
-          errs() << "BOLT-WARNING: no successor for block " << BB.getName()
-                 << " that matches index " << YamlSI.Index << " or block "
-                 << SuccessorBB.getName() << '\n';
-        ++MismatchedEdges;
-        continue;
+      BinaryBasicBlock *ToBB = Order[YamlSI.Index];
+      if (!BB.getSuccessor(ToBB->getLabel())) {
+        // Allow for BOLT-removed passthrough blocks to align with DataReader
+        // behavior.
+        BinaryBasicBlock *FTSuccessor = BB.getConditionalSuccessor(false);
+        if (FTSuccessor && FTSuccessor->succ_size() == 1 &&
+            FTSuccessor->getSuccessor(ToBB->getLabel())) {
+          BinaryBasicBlock::BinaryBranchInfo &FTBI =
+              FTSuccessor->getBranchInfo(*ToBB);
+          FTBI.Count += YamlSI.Count;
+          FTBI.MispredictedCount += YamlSI.Mispreds;
+          ToBB = FTSuccessor;
+        } else {
+          if (opts::Verbosity >= 1)
+            errs() << "BOLT-WARNING: no successor for block " << BB.getName()
+                   << " that matches index " << YamlSI.Index << " or block "
+                   << ToBB->getName() << '\n';
+          ++MismatchedEdges;
+          continue;
+        }
       }
 
-      BinaryBasicBlock::BinaryBranchInfo &BI = BB.getBranchInfo(SuccessorBB);
+      BinaryBasicBlock::BinaryBranchInfo &BI = BB.getBranchInfo(*ToBB);
       BI.Count += YamlSI.Count;
       BI.MispredictedCount += YamlSI.Mispreds;
     }
diff --git a/bolt/test/X86/profile-passthrough-block.test b/bolt/test/X86/profile-passthrough-block.test
new file mode 100644
index 0000000000000..54920e05d5520
--- /dev/null
+++ b/bolt/test/X86/profile-passthrough-block.test
@@ -0,0 +1,66 @@
+## This reproduces a bug with BOLT setting incorrect discriminator for
+## secondary entry points in YAML profile.
+
+# REQUIRES: system-linux
+# RUN: split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %t/main.s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt %t.exe -o %t.out --data %t/yaml --profile-ignore-hash -v=1 \
+# RUN:   2>&1 | FileCheck %s
+
+# CHECK: BOLT-INFO: 1 out of 1 functions in the binary (100.0%) have non-empty execution profile
+# CHECK-NOT: BOLT-WARNING: no successor for block .LFT0 that matches index 3 or block .Ltmp0
+
+#--- main.s
+.globl main
+.type	main, @function
+main:
+  .cfi_startproc
+.LBB00:
+  pushq   %rbp
+  movq    %rsp, %rbp
+  subq    $16, %rsp
+  testq   %rax, %rax
+  js      .LBB03
+.LBB01:
+  jne     .LBB04
+.LBB02:
+  nop
+.LBB03:
+  xorl    %eax, %eax
+  addq    $16, %rsp
+  popq    %rbp
+  retq
+.LBB04:
+  xorl    %eax, %eax
+  addq    $16, %rsp
+  popq    %rbp
+  retq
+## For relocations against .text
+.LBB05:
+  call exit
+  .cfi_endproc
+  .size	main, .-main
+
+#--- yaml
+---
+header:
+  profile-version: 1
+  binary-name:     'profile-passthrough-block.s.tmp.exe'
+  binary-build-id: '<unknown>'
+  profile-flags:   [ lbr ]
+  profile-origin:  branch profile reader
+  profile-events:  ''
+  dfs-order:       false
+  hash-func:       xxh3
+functions:
+  - name:            main
+    fid:             0
+    hash:            0x0000000000000000
+    exec:            1
+    nblocks:         6
+    blocks:
+      - bid:             1
+        insns:           1
+        succ:            [ { bid: 3, cnt: 1} ]
+...

>From addda4f3b7216c104cacec849368d31358f92fc1 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Sat, 11 May 2024 19:09:40 -0700
Subject: [PATCH 2/2] Fix test

---
 bolt/test/X86/profile-passthrough-block.test | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/bolt/test/X86/profile-passthrough-block.test b/bolt/test/X86/profile-passthrough-block.test
index 54920e05d5520..8cc07c5a71362 100644
--- a/bolt/test/X86/profile-passthrough-block.test
+++ b/bolt/test/X86/profile-passthrough-block.test
@@ -6,9 +6,10 @@
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %t/main.s -o %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
 # RUN: llvm-bolt %t.exe -o %t.out --data %t/yaml --profile-ignore-hash -v=1 \
-# RUN:   2>&1 | FileCheck %s
+# RUN:   --print-cfg 2>&1 | FileCheck %s
 
-# CHECK: BOLT-INFO: 1 out of 1 functions in the binary (100.0%) have non-empty execution profile
+# CHECK: Binary Function "main" after building cfg
+# CHECK: Profile Acc : 100.0%
 # CHECK-NOT: BOLT-WARNING: no successor for block .LFT0 that matches index 3 or block .Ltmp0
 
 #--- main.s