From cbd302410e9f27013223a96edcd78dfb597979e1 Mon Sep 17 00:00:00 2001
From: ShatianWang <38512325+ShatianWang@users.noreply.github.com>
Date: Wed, 21 Aug 2024 00:35:07 -0400
Subject: [PATCH] [BOLT] Improve BinaryFunction::inferFallThroughCounts()
 (#105450)

This PR improves how basic block execution count is updated when using
the BOLT option `-infer-fall-throughs`. Previously, if a 0-count
fall-through edge is assigned a positive inferred count N, then the
successor block's execution count will be incremented by N. Since the
successor's execution count is calculated using information besides
inflow sum (such as outflow sum), it likely is already correct, and
incrementing it by an additional N would be wrong. This PR improves how
the successor's execution count is updated by using the max over its
current count and N.
---
 bolt/lib/Core/BinaryFunctionProfile.cpp |  3 +-
 bolt/test/X86/infer-fall-throughs.s     | 45 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 bolt/test/X86/infer-fall-throughs.s

diff --git a/bolt/lib/Core/BinaryFunctionProfile.cpp b/bolt/lib/Core/BinaryFunctionProfile.cpp
index 55ebe5fc900e..726da6a9d082 100644
--- a/bolt/lib/Core/BinaryFunctionProfile.cpp
+++ b/bolt/lib/Core/BinaryFunctionProfile.cpp
@@ -336,7 +336,8 @@ void BinaryFunction::inferFallThroughCounts() {
       if (SuccBI.Count == 0) {
         SuccBI.Count = Inferred;
         SuccBI.MispredictedCount = BinaryBasicBlock::COUNT_INFERRED;
-        Succ->ExecutionCount += Inferred;
+        Succ->ExecutionCount =
+            std::max(Succ->getKnownExecutionCount(), Inferred);
       }
     }
   }
diff --git a/bolt/test/X86/infer-fall-throughs.s b/bolt/test/X86/infer-fall-throughs.s
new file mode 100644
index 000000000000..3289a8ea4aec
--- /dev/null
+++ b/bolt/test/X86/infer-fall-throughs.s
@@ -0,0 +1,45 @@
+## Test that infer-fall-throughs would correctly infer the wrong fall-through
+## edge count in the example
+
+# RUN: llvm-mc --filetype=obj --triple x86_64-unknown-unknown %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# 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 \
+# RUN:         --print-estimate-edge-counts --data=%t.fdata \
+# RUN:     2>&1 | FileCheck --check-prefix=WITHOUTINFERENCE %s
+# RUN: llvm-bolt %t.exe -o %t.bolt --infer-fall-throughs \
+# RUN:         --print-estimate-edge-counts --data=%t.fdata \
+# RUN:     2>&1 | FileCheck --check-prefix=CORRECTINFERENCE %s
+
+
+# WITHOUTINFERENCE: Binary Function "main" after estimate-edge-counts
+# WITHOUTINFERENCE: {{^\.Ltmp0}}
+# WITHOUTINFERENCE: Successors: .Ltmp1 (mispreds: 0, count: 10), .LFT0 (mispreds: 0, count: 0)
+# WITHOUTINFERENCE: {{^\.LFT0}}
+# WITHOUTINFERENCE: Exec Count : 490
+
+# CORRECTINFERENCE: Binary Function "main" after estimate-edge-counts
+# CORRECTINFERENCE: {{^\.Ltmp0}}
+# CORRECTINFERENCE: Successors: .Ltmp1 (mispreds: 0, count: 10), .LFT0 (inferred count: 490)
+# CORRECTINFERENCE: {{^\.LFT0}}
+# CORRECTINFERENCE: Exec Count : 490
+
+
+        .globl  main
+        .type   main, @function
+main:
+LLmain_LLstart:
+        jmp     LLstart
+# FDATA: 1 main #LLmain_LLstart# 1 main #LLstart# 0 500
+LLstart:
+        jge     LLexit
+# FDATA: 1 main #LLstart# 1 main #LLexit# 0 10
+# FDATA: 1 main #LLstart# 1 main #LLmore# 0 0
+LLmore:
+        movl    $5, %eax
+# FDATA: 1 main #LLmore# 1 main #LLexit# 0 490
+LLexit:
+        ret
+.LLmain_end:
+        .size   main, .LLmain_end-main