[llvm] [BOLT] Set InitialDynoStats after EstimateEdgeCounts (PR #93218)

[Amir Ayupov via llvm-commits]

https://github.com/aaupov created https://github.com/llvm/llvm-project/pull/93218

InitialDynoStats used to be assigned inside `runAllPasses`, but the
assignment executed before any of the passes. As we've moved
`EstimateEdgeCounts` into a pass out of ProfileReader, it needs to
execute before initial dyno stats are set.

Thus move `InitialDynoStats` into BinaryContext and assignment into
`DynoStatsSetPass`.


>From 9693efab21ffdf5bceb60d6cd5933bd19e343f6f Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 23 May 2024 10:11:38 -0700
Subject: [PATCH] [BOLT] Set InitialDynoStats after EstimateEdgeCounts

InitialDynoStats used to be assigned inside `runAllPasses`, but the
assignment executed before any of the passes. As we've moved
`EstimateEdgeCounts` into a pass out of ProfileReader, it needs to
execute before initial dyno stats are set.

Thus move `InitialDynoStats` into BinaryContext and assignment into
`DynoStatsSetPass`.
---
 bolt/include/bolt/Core/BinaryContext.h  |  4 ++++
 bolt/include/bolt/Passes/BinaryPasses.h | 23 ++++++++++++++++++++---
 bolt/lib/Core/BinaryContext.cpp         |  2 +-
 bolt/lib/Rewrite/BinaryPassManager.cpp  | 10 ++++------
 4 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 633d83a4368cb..a3f1e6ca08053 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -17,6 +17,7 @@
 #include "bolt/Core/BinaryData.h"
 #include "bolt/Core/BinarySection.h"
 #include "bolt/Core/DebugData.h"
+#include "bolt/Core/DynoStats.h"
 #include "bolt/Core/JumpTable.h"
 #include "bolt/Core/MCPlusBuilder.h"
 #include "bolt/RuntimeLibs/RuntimeLibrary.h"
@@ -717,6 +718,9 @@ class BinaryContext {
     uint64_t NumStaleBlocksWithEqualIcount{0};
   } Stats;
 
+  // Original binary execution count stats.
+  DynoStats InitialDynoStats;
+
   // Address of the first allocated segment.
   uint64_t FirstAllocAddress{std::numeric_limits<uint64_t>::max()};
 
diff --git a/bolt/include/bolt/Passes/BinaryPasses.h b/bolt/include/bolt/Passes/BinaryPasses.h
index a07c9130041fd..ad8473c4aae02 100644
--- a/bolt/include/bolt/Passes/BinaryPasses.h
+++ b/bolt/include/bolt/Passes/BinaryPasses.h
@@ -53,15 +53,31 @@ class BinaryFunctionPass {
   virtual Error runOnFunctions(BinaryContext &BC) = 0;
 };
 
+/// A pass to set initial program-wide dynostats.
+class DynoStatsSetPass : public BinaryFunctionPass {
+public:
+  DynoStatsSetPass() : BinaryFunctionPass(false) {}
+
+  const char *getName() const override {
+    return "set dyno-stats before optimizations";
+  }
+
+  bool shouldPrint(const BinaryFunction &BF) const override { return false; }
+
+  Error runOnFunctions(BinaryContext &BC) override {
+    BC.InitialDynoStats = getDynoStats(BC.getBinaryFunctions(), BC.isAArch64());
+    return Error::success();
+  }
+};
+
 /// A pass to print program-wide dynostats.
 class DynoStatsPrintPass : public BinaryFunctionPass {
 protected:
-  DynoStats PrevDynoStats;
   std::string Title;
 
 public:
-  DynoStatsPrintPass(const DynoStats &PrevDynoStats, const char *Title)
-      : BinaryFunctionPass(false), PrevDynoStats(PrevDynoStats), Title(Title) {}
+  DynoStatsPrintPass(const char *Title)
+      : BinaryFunctionPass(false), Title(Title) {}
 
   const char *getName() const override {
     return "print dyno-stats after optimizations";
@@ -70,6 +86,7 @@ class DynoStatsPrintPass : public BinaryFunctionPass {
   bool shouldPrint(const BinaryFunction &BF) const override { return false; }
 
   Error runOnFunctions(BinaryContext &BC) override {
+    const DynoStats PrevDynoStats = BC.InitialDynoStats;
     const DynoStats NewDynoStats =
         getDynoStats(BC.getBinaryFunctions(), BC.isAArch64());
     const bool Changed = (NewDynoStats != PrevDynoStats);
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 876934b4ddf10..c18f188fa77cf 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -142,7 +142,7 @@ BinaryContext::BinaryContext(std::unique_ptr<MCContext> Ctx,
       AsmInfo(std::move(AsmInfo)), MII(std::move(MII)), STI(std::move(STI)),
       InstPrinter(std::move(InstPrinter)), MIA(std::move(MIA)),
       MIB(std::move(MIB)), MRI(std::move(MRI)), DisAsm(std::move(DisAsm)),
-      Logger(Logger) {
+      Logger(Logger), InitialDynoStats(isAArch64()) {
   Relocation::Arch = this->TheTriple->getArch();
   RegularPageSize = isAArch64() ? RegularPageSizeAArch64 : RegularPageSizeX86;
   PageAlign = opts::NoHugePages ? RegularPageSize : HugePageSize;
diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp
index 0712330a524b0..aaa0e1ff4d46f 100644
--- a/bolt/lib/Rewrite/BinaryPassManager.cpp
+++ b/bolt/lib/Rewrite/BinaryPassManager.cpp
@@ -343,8 +343,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(
       std::make_unique<EstimateEdgeCounts>(PrintEstimateEdgeCounts));
 
-  const DynoStats InitialDynoStats =
-      getDynoStats(BC.getBinaryFunctions(), BC.isAArch64());
+  Manager.registerPass(std::make_unique<DynoStatsSetPass>());
 
   Manager.registerPass(std::make_unique<AsmDumpPass>(),
                        opts::AsmDump.getNumOccurrences());
@@ -456,10 +455,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(std::make_unique<SplitFunctions>(PrintSplit));
 
   // Print final dyno stats right while CFG and instruction analysis are intact.
-  Manager.registerPass(
-      std::make_unique<DynoStatsPrintPass>(
-          InitialDynoStats, "after all optimizations before SCTC and FOP"),
-      opts::PrintDynoStats || opts::DynoStatsAll);
+  Manager.registerPass(std::make_unique<DynoStatsPrintPass>(
+                           "after all optimizations before SCTC and FOP"),
+                       opts::PrintDynoStats || opts::DynoStatsAll);
 
   // Add the StokeInfo pass, which extract functions for stoke optimization and
   // get the liveness information for them