1 files changed, 15 insertions, 1 deletions
diff --git a/llvm/lib/Transforms/Instrumentation/CFGMST.h b/llvm/lib/Transforms/Instrumentation/CFGMST.h
index 8bb6f47c4846f..9addb5d1ba938 100644
--- a/llvm/lib/Transforms/Instrumentation/CFGMST.h
+++ b/llvm/lib/Transforms/Instrumentation/CFGMST.h
@@ -20,6 +20,7 @@
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
@@ -28,6 +29,11 @@
 
 #define DEBUG_TYPE "cfgmst"
 
+using namespace llvm;
+static cl::opt<bool> PGOInstrumentEntry(
+    "pgo-instrument-entry", cl::init(false), cl::Hidden,
+    cl::desc("Force to instrument function entry basicblock."));
+
 namespace llvm {
 
 /// An union-find based Minimum Spanning Tree for CFG
@@ -100,8 +106,11 @@ public:
 
     const BasicBlock *Entry = &(F.getEntryBlock());
     uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
+    // If we want to instrument the entry count, lower the weight to 0.
+    if (PGOInstrumentEntry)
+      EntryWeight = 0;
     Edge *EntryIncoming = nullptr, *EntryOutgoing = nullptr,
-        *ExitOutgoing = nullptr, *ExitIncoming = nullptr;
+         *ExitOutgoing = nullptr, *ExitIncoming = nullptr;
     uint64_t MaxEntryOutWeight = 0, MaxExitOutWeight = 0, MaxExitInWeight = 0;
 
     // Add a fake edge to the entry.
@@ -135,6 +144,8 @@ public:
           }
           if (BPI != nullptr)
             Weight = BPI->getEdgeProbability(&*BB, TargetBB).scale(scaleFactor);
+          if (Weight == 0)
+            Weight++;
           auto *E = &addEdge(&*BB, TargetBB, Weight);
           E->IsCritical = Critical;
           LLVM_DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to "
@@ -278,6 +289,9 @@ public:
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();
+    if (PGOInstrumentEntry && (AllEdges.size() > 1))
+      std::iter_swap(std::move(AllEdges.begin()),
+                     std::move(AllEdges.begin() + AllEdges.size() - 1));
   }
 };