[BOLT][NFC] Make estimateEdgeCounts a BinaryFunctionPass

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

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

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 {

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;
 

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());