From f3dc732b3623c86164f9d95f7563f982cecc5558 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Wed, 22 May 2024 11:59:00 -0700
Subject: [PATCH] [BOLT][NFC] Make estimateEdgeCounts a BinaryFunctionPass
 (#93074)

---
 bolt/include/bolt/Passes/MCF.h         | 17 +++++++++++++++--
 bolt/lib/Passes/MCF.cpp                | 24 +++++++++++++++++++++++-
 bolt/lib/Profile/DataReader.cpp        |  2 --
 bolt/lib/Profile/YAMLProfileReader.cpp |  4 +---
 bolt/lib/Rewrite/BinaryPassManager.cpp |  9 +++++++++
 5 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/bolt/include/bolt/Passes/MCF.h b/bolt/include/bolt/Passes/MCF.h
index 4b87401498fa5a..3fe674463bf136 100644
--- a/bolt/include/bolt/Passes/MCF.h
+++ b/bolt/include/bolt/Passes/MCF.h
@@ -9,10 +9,12 @@
 #ifndef BOLT_PASSES_MCF_H
 #define BOLT_PASSES_MCF_H
 
+#include "bolt/Passes/BinaryPasses.h"
+#include "llvm/Support/CommandLine.h"
+
 namespace llvm {
 namespace bolt {
 
-class BinaryFunction;
 class DataflowInfoManager;
 
 /// Implement the idea in "SamplePGO - The Power of Profile Guided Optimizations
@@ -23,7 +25,18 @@ void equalizeBBCounts(DataflowInfoManager &Info, BinaryFunction &BF);
 
 /// Fill edge counts based on the basic block count. Used in nonLBR mode when
 /// we only have bb count.
-void estimateEdgeCounts(BinaryFunction &BF);
+class EstimateEdgeCounts : public BinaryFunctionPass {
+  void runOnFunction(BinaryFunction &BF);
+
+public:
+  explicit EstimateEdgeCounts(const cl::opt<bool> &PrintPass)
+      : BinaryFunctionPass(PrintPass) {}
+
+  const char *getName() const override { return "estimate-edge-counts"; }
+
+  /// Pass entry point
+  Error runOnFunctions(BinaryContext &BC) override;
+};
 
 } // end namespace bolt
 } // end namespace llvm
diff --git a/bolt/lib/Passes/MCF.cpp b/bolt/lib/Passes/MCF.cpp
index b2723cd8dcb8b6..77dea7369140e7 100644
--- a/bolt/lib/Passes/MCF.cpp
+++ b/bolt/lib/Passes/MCF.cpp
@@ -12,9 +12,11 @@
 
 #include "bolt/Passes/MCF.h"
 #include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/ParallelUtilities.h"
 #include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include <algorithm>
 #include <vector>
@@ -432,7 +434,7 @@ void equalizeBBCounts(DataflowInfoManager &Info, BinaryFunction &BF) {
   }
 }
 
-void estimateEdgeCounts(BinaryFunction &BF) {
+void EstimateEdgeCounts::runOnFunction(BinaryFunction &BF) {
   EdgeWeightMap PredEdgeWeights;
   EdgeWeightMap SuccEdgeWeights;
   if (!opts::IterativeGuess) {
@@ -453,5 +455,25 @@ void estimateEdgeCounts(BinaryFunction &BF) {
   recalculateBBCounts(BF, /*AllEdges=*/false);
 }
 
+Error EstimateEdgeCounts::runOnFunctions(BinaryContext &BC) {
+  if (llvm::none_of(llvm::make_second_range(BC.getBinaryFunctions()),
+                    [](const BinaryFunction &BF) {
+                      return BF.getProfileFlags() == BinaryFunction::PF_SAMPLE;
+                    }))
+    return Error::success();
+
+  ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+    runOnFunction(BF);
+  };
+  ParallelUtilities::PredicateTy SkipFunc = [&](const BinaryFunction &BF) {
+    return BF.getProfileFlags() != BinaryFunction::PF_SAMPLE;
+  };
+
+  ParallelUtilities::runOnEachFunction(
+      BC, ParallelUtilities::SchedulingPolicy::SP_BB_QUADRATIC, WorkFun,
+      SkipFunc, "EstimateEdgeCounts");
+  return Error::success();
+}
+
 } // namespace bolt
 } // namespace llvm
diff --git a/bolt/lib/Profile/DataReader.cpp b/bolt/lib/Profile/DataReader.cpp
index 06c5e96b780643..f2e999bbfdc6dc 100644
--- a/bolt/lib/Profile/DataReader.cpp
+++ b/bolt/lib/Profile/DataReader.cpp
@@ -598,8 +598,6 @@ void DataReader::readSampleData(BinaryFunction &BF) {
   }
 
   BF.ExecutionCount = TotalEntryCount;
-
-  estimateEdgeCounts(BF);
 }
 
 void DataReader::convertBranchData(BinaryFunction &BF) const {
diff --git a/bolt/lib/Profile/YAMLProfileReader.cpp b/bolt/lib/Profile/YAMLProfileReader.cpp
index 29d94067f459fc..aa38dda47eb564 100644
--- a/bolt/lib/Profile/YAMLProfileReader.cpp
+++ b/bolt/lib/Profile/YAMLProfileReader.cpp
@@ -253,10 +253,8 @@ bool YAMLProfileReader::parseFunctionProfile(
     if (BB.getExecutionCount() == BinaryBasicBlock::COUNT_NO_PROFILE)
       BB.setExecutionCount(0);
 
-  if (YamlBP.Header.Flags & BinaryFunction::PF_SAMPLE) {
+  if (YamlBP.Header.Flags & BinaryFunction::PF_SAMPLE)
     BF.setExecutionCount(FunctionExecutionCount);
-    estimateEdgeCounts(BF);
-  }
 
   ProfileMatched &= !MismatchedBlocks && !MismatchedCalls && !MismatchedEdges;
 
diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp
index cbb7199a53ddd1..0712330a524b01 100644
--- a/bolt/lib/Rewrite/BinaryPassManager.cpp
+++ b/bolt/lib/Rewrite/BinaryPassManager.cpp
@@ -23,6 +23,7 @@
 #include "bolt/Passes/JTFootprintReduction.h"
 #include "bolt/Passes/LongJmp.h"
 #include "bolt/Passes/LoopInversionPass.h"
+#include "bolt/Passes/MCF.h"
 #include "bolt/Passes/PLTCall.h"
 #include "bolt/Passes/PatchEntries.h"
 #include "bolt/Passes/RegReAssign.h"
@@ -90,6 +91,11 @@ PrintAfterLowering("print-after-lowering",
   cl::desc("print function after instruction lowering"),
   cl::Hidden, cl::cat(BoltOptCategory));
 
+static cl::opt<bool> PrintEstimateEdgeCounts(
+    "print-estimate-edge-counts",
+    cl::desc("print function after edge counts are set for no-LBR profile"),
+    cl::Hidden, cl::cat(BoltOptCategory));
+
 cl::opt<bool>
 PrintFinalized("print-finalized",
   cl::desc("print function after CFG is finalized"),
@@ -334,6 +340,9 @@ Error BinaryFunctionPassManager::runPasses() {
 Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   BinaryFunctionPassManager Manager(BC);
 
+  Manager.registerPass(
+      std::make_unique<EstimateEdgeCounts>(PrintEstimateEdgeCounts));
+
   const DynoStats InitialDynoStats =
       getDynoStats(BC.getBinaryFunctions(), BC.isAArch64());