Tune the weight propagation algorithm for sample profile.

[danielcdh]

Summary

Handle the case when there is only one incoming/outgoing edge for a visited basic block: use the block weight to adjust edge weight even when the edge has been visited before. This can help reduce inaccuracies introduced by incorrect basic block profile, as shown in the updated unittest.

[dnovillo]

LGTM. Thanks.

lib/Transforms/IPO/SampleProfile.cpp

for (const auto &BI : F) {
 
     // Visit all the predecessor and successor edges to determine
     // which ones have a weight assigned already. Note that it doesn't
     // matter that we only keep track of a single unknown edge. The
     // only case we are interested in handling is when only a single
     // edge is unknown (see setEdgeOrBlockWeight).
     for (unsigned i = 0; i < 2; i++) {
       uint64_t TotalWeight = 0;
-      unsigned NumUnknownEdges = 0;
-      Edge UnknownEdge, SelfReferentialEdge;
+      unsigned NumUnknownEdges = 0, NumTotalEdges = 0;
+      Edge UnknownEdge, SelfReferentialEdge, SingleEdge;
 
       if (i == 0) {
         // First, visit all predecessor edges.
+        NumTotalEdges = Predecessors[BB].size();
         for (auto *Pred : Predecessors[BB]) {
           Edge E = std::make_pair(Pred, BB);
           TotalWeight += visitEdge(E, &NumUnknownEdges, &UnknownEdge);
           if (E.first == E.second)
             SelfReferentialEdge = E;
         }
+        if (NumTotalEdges == 1) {
+          SingleEdge = std::make_pair(Predecessors[BB][0], BB);
+        }
       } else {
         // On the second round, visit all successor edges.
+        NumTotalEdges = Successors[BB].size();
         for (auto *Succ : Successors[BB]) {
           Edge E = std::make_pair(BB, Succ);
           TotalWeight += visitEdge(E, &NumUnknownEdges, &UnknownEdge);
         }
+        if (NumTotalEdges == 1) {
+          SingleEdge = std::make_pair(BB, Successors[BB][0]);
+        }
       }
 
       // After visiting all the edges, there are three cases that we
       // can handle immediately:
       //
       // - All the edge weights are known (i.e., NumUnknownEdges == 0).
       //   In this case, we simply check that the sum of all the edges
       //   is the same as BB's weight. If not, we change BB's weight
for (unsigned i = 0; i < 2; i++) {
       //   the self-referential edge as we did in the previous case.
       //
       // In any other case, we must continue iterating. Eventually,
       // all edges will get a weight, or iteration will stop when
       // it reaches SampleProfileMaxPropagateIterations.
       if (NumUnknownEdges <= 1) {
         uint64_t &BBWeight = BlockWeights[EC];
         if (NumUnknownEdges == 0) {
+          if (!VisitedBlocks.count(EC)) {
             // If we already know the weight of all edges, the weight of the
             // basic block can be computed. It should be no larger than the sum
             // of all edge weights.
             if (TotalWeight > BBWeight) {
               BBWeight = TotalWeight;
               Changed = true;
               DEBUG(dbgs() << "All edge weights for " << BB->getName()
                            << " known. Set weight for block: ";
                     printBlockWeight(dbgs(), BB););
             }
-          if (VisitedBlocks.insert(EC).second)
+          } else if (NumTotalEdges == 1 &&
+                     EdgeWeights[SingleEdge] < BlockWeights[EC]) {
+            // If there is only one edge for the visited basic block, use the
+            // block weight to adjust edge weight if edge weight is smaller.
+            EdgeWeights[SingleEdge] = BlockWeights[EC];
             Changed = true;
+          }
         } else if (NumUnknownEdges == 1 && VisitedBlocks.count(EC)) {
           // If there is a single unknown edge and the block has been
           // visited, then we can compute E's weight.
           if (BBWeight >= TotalWeight)
             EdgeWeights[UnknownEdge] = BBWeight - TotalWeight;
           else
             EdgeWeights[UnknownEdge] = 0;
           VisitedEdges.insert(UnknownEdge);

test/Transforms/SampleProfile/fnptr.ll

 ; The two profiles used in this test are the same but encoded in different
 ; formats. This checks that we produce the same profile annotations regardless
 ; of the profile format.
 ;
 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/fnptr.prof | opt -analyze -branch-prob | FileCheck %s
 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/fnptr.binprof | opt -analyze -branch-prob | FileCheck %s
 
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/fnptr.prof | opt -analyze -branch-prob | FileCheck %s
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/fnptr.binprof | opt -analyze -branch-prob | FileCheck %s
 
 ; CHECK:   edge for.body3 -> if.then probability is 0x1a4f3959 / 0x80000000 = 20.55%
 ; CHECK:   edge for.body3 -> if.else probability is 0x65b0c6a7 / 0x80000000 = 79.45%
-; CHECK:   edge for.inc -> for.inc12 probability is 0x33d4a4c1 / 0x80000000 = 40.49%
-; CHECK:   edge for.inc -> for.body3 probability is 0x4c2b5b3f / 0x80000000 = 59.51%
-; CHECK:   edge for.inc12 -> for.end14 probability is 0x3f06d04e / 0x80000000 = 49.24%
-; CHECK:   edge for.inc12 -> for.cond1.preheader probability is 0x40f92fb2 / 0x80000000 = 50.76%
+; CHECK:   edge for.inc -> for.inc12 probability is 0x20dc8dc9 / 0x80000000 = 25.67%
+; CHECK:   edge for.inc -> for.body3 probability is 0x5f237237 / 0x80000000 = 74.33%
+; CHECK:   edge for.inc12 -> for.end14 probability is 0x00000000 / 0x80000000 = 0.00%
+; CHECK:   edge for.inc12 -> for.cond1.preheader probability is 0x80000000 / 0x80000000 = 100.00%
 
 ; Original C++ test case.
 ;
 ; #include <stdlib.h>
 ; #include <math.h>
 ; #include <stdio.h>
 ;
 ; #define N 10000