1 files changed, 67 insertions, 7 deletions

diff --git a/lib/Transforms/Instrumentation/CFGMST.h b/lib/Transforms/Instrumentation/CFGMST.h
index 16e2e6b4e730..075e5672cff8 100644
--- a/lib/Transforms/Instrumentation/CFGMST.h
+++ b/lib/Transforms/Instrumentation/CFGMST.h
@@ -46,6 +46,10 @@ public:
   // This map records the auxiliary information for each BB.
   DenseMap<const BasicBlock *, std::unique_ptr<BBInfo>> BBInfos;
 
+  // Whehter the function has an exit block with no successors.
+  // (For function with an infinite loop, this block may be absent)
+  bool ExitBlockFound = false;
+
   // Find the root group of the G and compress the path from G to the root.
   BBInfo *findAndCompressGroup(BBInfo *G) {
     if (G->Group != G)
@@ -95,14 +99,20 @@ public:
   void buildEdges() {
     DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
 
-    const BasicBlock *BB = &(F.getEntryBlock());
+    const BasicBlock *Entry = &(F.getEntryBlock());
     uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
+    Edge *EntryIncoming = nullptr, *EntryOutgoing = nullptr,
+        *ExitOutgoing = nullptr, *ExitIncoming = nullptr;
+    uint64_t MaxEntryOutWeight = 0, MaxExitOutWeight = 0, MaxExitInWeight = 0;
+
     // Add a fake edge to the entry.
-    addEdge(nullptr, BB, EntryWeight);
+    EntryIncoming = &addEdge(nullptr, Entry, EntryWeight);
+    DEBUG(dbgs() << "  Edge: from fake node to " << Entry->getName()
+                     << " w = " << EntryWeight << "\n");
 
     // Special handling for single BB functions.
-    if (succ_empty(BB)) {
-      addEdge(BB, nullptr, EntryWeight);
+    if (succ_empty(Entry)) {
+      addEdge(Entry, nullptr, EntryWeight);
       return;
     }
 
@@ -126,16 +136,62 @@ public:
           }
           if (BPI != nullptr)
             Weight = BPI->getEdgeProbability(&*BB, TargetBB).scale(scaleFactor);
-          addEdge(&*BB, TargetBB, Weight).IsCritical = Critical;
+          auto *E = &addEdge(&*BB, TargetBB, Weight);
+          E->IsCritical = Critical;
           DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to "
                        << TargetBB->getName() << "  w=" << Weight << "\n");
+
+          // Keep track of entry/exit edges:
+          if (&*BB == Entry) {
+            if (Weight > MaxEntryOutWeight) {
+              MaxEntryOutWeight = Weight;
+              EntryOutgoing = E;
+            }
+          }
+
+          auto *TargetTI = TargetBB->getTerminator();
+          if (TargetTI && !TargetTI->getNumSuccessors()) {
+            if (Weight > MaxExitInWeight) {
+              MaxExitInWeight = Weight;
+              ExitIncoming = E;
+            }
+          }
         }
       } else {
-        addEdge(&*BB, nullptr, BBWeight);
-        DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to exit"
+        ExitBlockFound = true;
+        Edge *ExitO = &addEdge(&*BB, nullptr, BBWeight);
+        if (BBWeight > MaxExitOutWeight) {
+          MaxExitOutWeight = BBWeight;
+          ExitOutgoing = ExitO;
+        }
+        DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to fake exit"
                      << " w = " << BBWeight << "\n");
       }
     }
+
+    // Entry/exit edge adjustment heurisitic:
+    // prefer instrumenting entry edge over exit edge
+    // if possible. Those exit edges may never have a chance to be
+    // executed (for instance the program is an event handling loop)
+    // before the profile is asynchronously dumped.
+    //
+    // If EntryIncoming and ExitOutgoing has similar weight, make sure
+    // ExitOutging is selected as the min-edge. Similarly, if EntryOutgoing
+    // and ExitIncoming has similar weight, make sure ExitIncoming becomes
+    // the min-edge.
+    uint64_t EntryInWeight = EntryWeight;
+
+    if (EntryInWeight >= MaxExitOutWeight &&
+        EntryInWeight * 2 < MaxExitOutWeight * 3) {
+      EntryIncoming->Weight = MaxExitOutWeight;
+      ExitOutgoing->Weight = EntryInWeight + 1;
+    }
+
+    if (MaxEntryOutWeight >= MaxExitInWeight &&
+        MaxEntryOutWeight * 2 < MaxExitInWeight * 3) {
+      EntryOutgoing->Weight = MaxExitInWeight;
+      ExitIncoming->Weight = MaxEntryOutWeight + 1;
+    }
   }
 
   // Sort CFG edges based on its weight.
@@ -167,6 +223,10 @@ public:
     for (auto &Ei : AllEdges) {
       if (Ei->Removed)
         continue;
+      // If we detect infinite loops, force
+      // instrumenting the entry edge:
+      if (!ExitBlockFound && Ei->SrcBB == nullptr)
+        continue;
       if (unionGroups(Ei->SrcBB, Ei->DestBB))
         Ei->InMST = true;
     }