[PGO] MST min edge selection heuristic to ensure non-zero entry count

[davidxl]

Summary

In many programs, event handling loop (not statically infinite) runs forever and the normal exit path is not executed unless there is error condition. Profile dumping can happen asynchronously. For such functions, if the entry block is not instrumented, we will get zero entry count which can be bad -- PGO relies on entry count to get BB count via scaling. See test case in D41058

This patch implements a new heuristic to select min edges -- if entry edges and exit edges have similar weight, prefer to select exit edges to be min edges.

This patch also fixes the problem with the real infinite loop case where entry count is zero.

[davidxl]

Ping ..

[xur]

This is a good improvement.

lib/Transforms/Instrumentation/CFGMST.h

183 ↗

(On Diff #126303)

Will this ever be true? EntryInWeight should be always greater than or equal to weight of Exit BB.

198 ↗

(On Diff #126303)

Do we need to make sure EntryInWeight less than MaxEntryOutWeight?

If Entry BB has more than one successors, and Exit BB has more than one predecessors,
even after adjustment, EntryIncoming->Weight might be greater than ExitOutgoing->Weight.

EntryIncoming ExitOutgoing and ExitIncoming will be in MST (thus not instrumented).
EntryOutgoing will be instrumented. But since it has two successors, we will instrumented the target BB.
(i.e. entryBB is not instrumented).

But I'm sure how this relates to the problem this patch wants to solve. It seems to me that choice b/w instrumenting entry and exit only applies the case of single successor/predecessor for entry and exit nodes.

[davidxl]

lib/Transforms/Instrumentation/CFGMST.h

183 ↗

(On Diff #126303)

Probably not -- will remove.

198 ↗

(On Diff #126303)

EntryIncomingWeight is adjusted to be smaller than MaxExitOutgoingWeight.

In your case, all EntryOutgoingEdges will be selected for instrumentation, and their sum will be the entry count. Unless One of entry outgoing edges is a critical edge which will be selected as Min(ST) edge, in which case, the entryIncomingEdge won't need to be in MST and therefore will be instrumented. That is why we need to make sure entryIncoming weight is smaller than exitoutoging weight.

[davidxl]

address comment: remove unnecessary early return check.

[davidxl]

The instrumentation point will be before the no-return call, so there is probably no issue with instrumenting direct successor of the entry.

[xur]

Look good to me.

lib/Transforms/Instrumentation/CFGMST.h

184 ↗

(On Diff #127198)

NIT: Should we remove this also. This is the same as the removed condition.

llvm/trunk/lib/Transforms/Instrumentation/CFGMST.h

public:
 
   // Store all the edges in CFG. It may contain some stale edges
   // when Removed is set.
   std::vector<std::unique_ptr<Edge>> AllEdges;
 
   // 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)
       G->Group = findAndCompressGroup(static_cast<BBInfo *>(G->Group));
     return static_cast<BBInfo *>(G->Group);
   }
 
   // Union BB1 and BB2 into the same group and return true.
public:
   }
 
   // Traverse the CFG using a stack. Find all the edges and assign the weight.
   // Edges with large weight will be put into MST first so they are less likely
   // to be instrumented.
   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;
     }
 
     static const uint32_t CriticalEdgeMultiplier = 1000;
 
     for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
       TerminatorInst *TI = BB->getTerminator();
       uint64_t BBWeight =
           (BFI != nullptr ? BFI->getBlockFreq(&*BB).getFrequency() : 2);
       uint64_t Weight = 2;
       if (int successors = TI->getNumSuccessors()) {
         for (int i = 0; i != successors; ++i) {
           BasicBlock *TargetBB = TI->getSuccessor(i);
           bool Critical = isCriticalEdge(TI, i);
           uint64_t scaleFactor = BBWeight;
           if (Critical) {
             if (scaleFactor < UINT64_MAX / CriticalEdgeMultiplier)
               scaleFactor *= CriticalEdgeMultiplier;
             else
               scaleFactor = UINT64_MAX;
           }
           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.
   void sortEdgesByWeight() {
     std::stable_sort(AllEdges.begin(), AllEdges.end(),
                      [](const std::unique_ptr<Edge> &Edge1,
                         const std::unique_ptr<Edge> &Edge2) {
                        return Edge1->Weight > Edge2->Weight;
for (auto &Ei : AllEdges) {
             Ei->InMST = true;
         }
       }
     }
 
     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;
     }
   }
 
   // Dump the Debug information about the instrumentation.
   void dumpEdges(raw_ostream &OS, const Twine &Message) const {
     if (!Message.str().empty())

llvm/trunk/test/Transforms/PGOProfile/Inputs/landingpad.proftext

 # :ir is the flag to indicate this is IR level profile.
 :ir
 foo
 59130013419
 4
 3
 1
 2
 0
 
 bar
 24868915205
 2
-1
+3
 2
 

llvm/trunk/test/Transforms/PGOProfile/Inputs/noreturncall.proftext

 # IR level Instrumentation Flag
 :ir
 foo
 # Func Hash:
 36496524737
 # Num Counters:
 3
 # Counter Values:
-20
+21
 21
 0

llvm/trunk/test/Transforms/PGOProfile/branch1.ll

 ; GEN-LINKONCE: @__llvm_profile_raw_version = linkonce constant i64 {{[0-9]+}}
 ; GEN: @__profn_test_br_1 = private constant [9 x i8] c"test_br_1"
 
 define i32 @test_br_1(i32 %i) {
 ; USE-LABEL: @test_br_1
 ; USE-SAME: !prof ![[FUNC_ENTRY_COUNT:[0-9]+]]
 entry:
 ; GEN: entry:
-; GEN-NOT: llvm.instrprof.increment
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([9 x i8], [9 x i8]* @__profn_test_br_1, i32 0, i32 0), i64 25571299074, i32 2, i32 0)
   %cmp = icmp sgt i32 %i, 0
   br i1 %cmp, label %if.then, label %if.end
 ; USE: br i1 %cmp, label %if.then, label %if.end
 ; USE-SAME: !prof ![[BW_ENTRY:[0-9]+]]
 ; USE-DAG: ![[BW_ENTRY]] = !{!"branch_weights", i32 2, i32 1}
 ; USE-LARGE: br i1 %cmp, label %if.then, label %if.end
 ; USE-LARGE-SAME: !prof ![[BW_L_ENTRY:[0-9]+]]
 ; USE-LARGE-DAG: ![[BW_L_ENTRY]] = !{!"branch_weights", i32 -1431655766, i32 1431655765}
 
 if.then:
 ; GEN: if.then:
 ; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([9 x i8], [9 x i8]* @__profn_test_br_1, i32 0, i32 0), i64 25571299074, i32 2, i32 1)
   %add = add nsw i32 %i, 2
   br label %if.end
 
 if.end:
 ; GEN: if.end:
-; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([9 x i8], [9 x i8]* @__profn_test_br_1, i32 0, i32 0), i64 25571299074, i32 2, i32 0)
+; GEN-NOT: llvm.instrprof.increment
+; GEN: ret i32
   %retv = phi i32 [ %add, %if.then ], [ %i, %entry ]
   ret i32 %retv
 }
 ; USE-DAG: {{![0-9]+}} = !{i32 1, !"ProfileSummary", {{![0-9]+}}}
 ; USE-DAG: {{![0-9]+}} = !{!"DetailedSummary", {{![0-9]+}}}
 ; USE-DAG: ![[FUNC_ENTRY_COUNT]] = !{!"function_entry_count", i64 3}
 
 ; ANALYSIS:remark: <unknown>:0:0: sgt_i32_Zero {{.*}}66.67% (total count : 3)

llvm/trunk/test/Transforms/PGOProfile/counter_promo.ll

 ; RUN: opt < %s -pgo-instr-gen -instrprof -do-counter-promotion=true -S | FileCheck --check-prefix=PROMO --check-prefix=NONATOMIC_PROMO %s
 ; RUN: opt < %s --passes=pgo-instr-gen,instrprof -do-counter-promotion=true -S | FileCheck --check-prefix=PROMO --check-prefix=NONATOMIC_PROMO %s 
 ; RUN: opt < %s -pgo-instr-gen -instrprof -do-counter-promotion=true -atomic-counter-update-promoted -S | FileCheck --check-prefix=PROMO --check-prefix=ATOMIC_PROMO %s
 ; RUN: opt < %s --passes=pgo-instr-gen,instrprof -do-counter-promotion=true -atomic-counter-update-promoted -S | FileCheck --check-prefix=PROMO --check-prefix=ATOMIC_PROMO %s 
 
 define void @foo(i32 %n, i32 %N) {
 ; PROMO-LABEL: @foo
+; PROMO: {{.*}} = load {{.*}} @__profc_foo{{.*}} 3)
+; PROMO-NEXT: add 
+; PROMO-NEXT: store {{.*}}@__profc_foo{{.*}}3)
 bb:
   %tmp = add nsw i32 %n, 1
   %tmp1 = add nsw i32 %n, -1
   br label %bb2
 
 bb2:                                              ; preds = %bb9, %bb
 ; PROMO: phi {{.*}}
 ; PROMO-NEXT: phi {{.*}}

 ; NONATOMIC_PROMO-NEXT: add {{.*}} %[[PROMO2]], %[[LIVEOUT2]]
 ; NONATOMIC_PROMO-NEXT: store {{.*}}@__profc_foo{{.*}}1)
 ; NONATOMIC_PROMO-NEXT: %[[PROMO3:[a-z0-9.]+]] = load {{.*}} @__profc_foo{{.*}} 2)
 ; NONATOMIC_PROMO-NEXT: add {{.*}} %[[PROMO3]], %[[LIVEOUT3]]
 ; NONATOMIC_PROMO-NEXT: store {{.*}}@__profc_foo{{.*}}2)
 ; ATOMIC_PROMO: atomicrmw add {{.*}} @__profc_foo{{.*}}0), i64 %[[LIVEOUT1]] seq_cst
 ; ATOMIC_PROMO-NEXT: atomicrmw add {{.*}} @__profc_foo{{.*}}1), i64 %[[LIVEOUT2]] seq_cst
 ; ATOMIC_PROMO-NEXT: atomicrmw add {{.*}} @__profc_foo{{.*}}2), i64 %[[LIVEOUT3]] seq_cst
-; PROMO: {{.*}} = load {{.*}} @__profc_foo{{.*}} 3)
-; PROMO-NEXT: add 
-; PROMO-NEXT: store {{.*}}@__profc_foo{{.*}}3)
 ; PROMO-NOT: @__profc_foo
 
 
 }
 
 declare void @bar(i32)

llvm/trunk/test/Transforms/PGOProfile/infinite_loop_gen.ll

+; RUN: opt < %s -pgo-instr-gen -S -o -   | FileCheck %s
+
+define void @foo() {
+entry:
+  br label %while.body
+  ; CHECK: llvm.instrprof.increment
+
+    while.body:                                       ; preds = %entry, %while.body
+    ; CHECK: llvm.instrprof.increment
+        call void (...) @bar() #2
+    br label %while.body
+}
+
+declare void @bar(...)
+
+attributes #0 = { nounwind }
+

llvm/trunk/test/Transforms/PGOProfile/landingpad.ll

 ; GEN: $__llvm_profile_raw_version = comdat any
 ; GEN: @__llvm_profile_raw_version = constant i64 {{[0-9]+}}, comdat
 ; GEN: @__profn_bar = private constant [3 x i8] c"bar"
 ; GEN: @__profn_foo = private constant [3 x i8] c"foo"
 
 define i32 @bar(i32 %i) {
 entry:
 ; GEN: entry:
-; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @__profn_bar, i32 0, i32 0), i64 24868915205, i32 2, i32 0)
   %rem = srem i32 %i, 3
   %tobool = icmp ne i32 %rem, 0
   br i1 %tobool, label %if.then, label %if.end
 ; USE: br i1 %tobool, label %if.then, label %if.end
 ; USE-SAME: !prof ![[BW_BAR_ENTRY:[0-9]+]]
 
 if.then:
 ; GEN: if.then:
 ; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @__profn_bar, i32 0, i32 0), i64 24868915205, i32 2, i32 1)
   %exception = call i8* @__cxa_allocate_exception(i64 4)
   %tmp = bitcast i8* %exception to i32*
   store i32 %i, i32* %tmp, align 16
   call void @__cxa_throw(i8* %exception, i8* bitcast (i8** @_ZTIi to i8*), i8* null)
   unreachable
 
 if.end:
 ; GEN: if.end:
-; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @__profn_bar, i32 0, i32 0), i64 24868915205, i32 2, i32 0)
+; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: ret i32
   ret i32 0
 }
 
 declare i8* @__cxa_allocate_exception(i64)
 
 declare void @__cxa_throw(i8*, i8*, i8*)
 
 define i32 @foo(i32 %i) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {

llvm/trunk/test/Transforms/PGOProfile/loop1.ll

 ; RUN: opt < %s -pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
 ; RUN: opt < %s -passes=pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
 ; RUN: llvm-profdata merge %S/Inputs/loop1.proftext -o %t.profdata
 ; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
 ; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; GEN: $__llvm_profile_raw_version = comdat any
 ; GEN: @__llvm_profile_raw_version = constant i64 {{[0-9]+}}, comdat
 ; GEN: @__profn_test_simple_for = private constant [15 x i8] c"test_simple_for"
 
 define i32 @test_simple_for(i32 %n) {
 entry:
 ; GEN: entry:
-; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_simple_for, i32 0, i32 0), i64 34137660316, i32 2, i32 1)
   br label %for.cond
 
 for.cond:
 ; GEN: for.cond:
 ; GEN-NOT: call void @llvm.instrprof.increment
   %i = phi i32 [ 0, %entry ], [ %inc1, %for.inc ]
   %sum = phi i32 [ 1, %entry ], [ %inc, %for.inc ]
   %cmp = icmp slt i32 %i, %n

 for.inc:
 ; GEN: for.inc:
 ; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_simple_for, i32 0, i32 0), i64 34137660316, i32 2, i32 0)
   %inc1 = add nsw i32 %i, 1
   br label %for.cond
 
 for.end:
 ; GEN: for.end:
-; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_simple_for, i32 0, i32 0), i64 34137660316, i32 2, i32 1)
+; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: ret i32
   ret i32 %sum
 }

llvm/trunk/test/Transforms/PGOProfile/loop2.ll

 ; RUN: opt < %s -pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
 ; RUN: opt < %s -passes=pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
 ; RUN: llvm-profdata merge %S/Inputs/loop2.proftext -o %t.profdata
 ; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
 ; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; GEN: $__llvm_profile_raw_version = comdat any
 ; GEN: @__llvm_profile_raw_version = constant i64 {{[0-9]+}}, comdat
 ; GEN: @__profn_test_nested_for = private constant [15 x i8] c"test_nested_for"
 
 define i32 @test_nested_for(i32 %r, i32 %s) {
 entry:
 ; GEN: entry:
-; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_nested_for, i32 0, i32 0), i64 53929068288, i32 3, i32 2)
   br label %for.cond.outer
 
 for.cond.outer:
 ; GEN: for.cond.outer:
 ; GEN-NOT: call void @llvm.instrprof.increment
   %i.0 = phi i32 [ 0, %entry ], [ %inc.2, %for.inc.outer ]
   %sum.0 = phi i32 [ 1, %entry ], [ %sum.1, %for.inc.outer ]
   %cmp = icmp slt i32 %i.0, %r

 for.inc.outer:
 ; GEN: for.inc.outer:
 ; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_nested_for, i32 0, i32 0), i64 53929068288, i32 3, i32 1)
   %inc.2 = add nsw i32 %i.0, 1
   br label %for.cond.outer
 
 for.end.outer:
 ; GEN: for.end.outer:
-; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @__profn_test_nested_for, i32 0, i32 0), i64 53929068288, i32 3, i32 2)
+; GEN-NOT: call void @llvm.instrprof.increment
+; GEN: ret i32
   ret i32 %sum.0
 }
 
 ; USE-DAG: ![[BW_FOR_COND_OUTER]] = !{!"branch_weights", i32 10, i32 6}
 ; USE-DAG: ![[BW_FOR_COND_INNER]] = !{!"branch_weights", i32 33, i32 10}
 

llvm/trunk/test/Transforms/PGOProfile/noreturncall.ll

if.then4:
   %call5 = call i32 @bar2(i32 undef)
   br label %if.end7
 
 if.end7:
   %mul = mul nsw i32 %ret.0, %ret.0
   ret i32 %mul
 }
 ; USE: ![[BW_ENTRY]] = !{!"branch_weights", i32 21, i32 0}
-; USE: ![[BW_IF]] = !{!"branch_weights", i32 0, i32 20}
+; USE: ![[BW_IF]] = !{!"branch_weights", i32 0, i32 21}