[llvm] [ctx_prof] Fix checks in `PGOCtxprofFlattening` (PR #108467)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

<details>
<summary>Changes</summary>

The assertion that all out-edges of a BB can't be 0 is incorrect: they can be, if that branch is on a cold subgraph.

Added validators and asserts about the expected proprerties of the propagated counters.

---
Full diff: https://github.com/llvm/llvm-project/pull/108467.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp (+56-9) 
- (added) llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll (+55) 


``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
index fc78d8c60ec050..77cfa2cdcff7e8 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
@@ -34,6 +34,7 @@
 #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
 #include "llvm/Transforms/Scalar/DCE.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include <deque>
 
 using namespace llvm;
 
@@ -110,7 +111,7 @@ class ProfileAnnotator final {
     }
 
     bool tryTakeCountFromKnownOutEdges(const BasicBlock &BB) {
-      if (!succ_empty(&BB) && !UnknownCountOutEdges) {
+      if (!UnknownCountOutEdges) {
         computeCountFrom(OutEdges);
         return true;
       }
@@ -118,7 +119,7 @@ class ProfileAnnotator final {
     }
 
     bool tryTakeCountFromKnownInEdges(const BasicBlock &BB) {
-      if (!BB.isEntryBlock() && !UnknownCountInEdges) {
+      if (!UnknownCountInEdges) {
         computeCountFrom(InEdges);
         return true;
       }
@@ -198,6 +199,50 @@ class ProfileAnnotator final {
 
   BBInfo &getBBInfo(const BasicBlock &BB) { return BBInfos.find(&BB)->second; }
 
+  const BBInfo &getBBInfo(const BasicBlock &BB) const {
+    return BBInfos.find(&BB)->second;
+  }
+
+  // validation function after we propagate the counters: all BBs and edges'
+  // counters must have a value.
+  bool allCountersAreAssinged() const {
+    for (const auto &BBInfo : BBInfos)
+      if (!BBInfo.second.hasCount())
+        return false;
+    for (const auto &EdgeInfo : EdgeInfos)
+      if (!EdgeInfo.Count.has_value())
+        return false;
+    return true;
+  }
+
+  bool allTakenPathsExit() const {
+    std::deque<const BasicBlock *> Worklist;
+    DenseSet<const BasicBlock *> Visited;
+    Worklist.push_back(&F.getEntryBlock());
+    Visited.insert(&F.getEntryBlock());
+    while (!Worklist.empty()) {
+      const auto *BB = Worklist.front();
+      Worklist.pop_front();
+      if (succ_size(BB) <= 1)
+        continue;
+      const auto &BBInfo = getBBInfo(*BB);
+      bool Inserted = false;
+      for (auto I = 0U; I < BB->getTerminator()->getNumSuccessors(); ++I) {
+        const auto *Succ = BB->getTerminator()->getSuccessor(I);
+        if (!shouldExcludeEdge(*BB, *Succ)) {
+          if (BBInfo.getEdgeCount(I) > 0)
+            if (Visited.insert(Succ).second) {
+              Worklist.push_back(Succ);
+              Inserted = true;
+            }
+        }
+      }
+      if (!Inserted)
+        return false;
+    }
+    return true;
+  }
+
 public:
   ProfileAnnotator(Function &F, const SmallVectorImpl<uint64_t> &Counters,
                    InstrProfSummaryBuilder &PB)
@@ -268,14 +313,16 @@ class ProfileAnnotator final {
         PB.addInternalCount(EdgeCount);
       }
 
-      if (MaxCount == 0)
-        F.getContext().emitError(
-            "[ctx-prof] Encountered a BB with more than one successor, where "
-            "all outgoing edges have a 0 count. This occurs in non-exiting "
-            "functions (message pumps, usually) which are not supported in the "
-            "contextual profiling case");
-      setProfMetadata(F.getParent(), Term, EdgeCounts, MaxCount);
+      if (MaxCount != 0)
+        setProfMetadata(F.getParent(), Term, EdgeCounts, MaxCount);
     }
+    assert(allCountersAreAssinged() &&
+           "Expected all counters have been assigned.");
+    assert(allTakenPathsExit() &&
+           "[ctx-prof] Encountered a BB with more than one successor, where "
+           "all outgoing edges have a 0 count. This occurs in non-exiting "
+           "functions (message pumps, usually) which are not supported in the "
+           "contextual profiling case");
   }
 };
 
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
new file mode 100644
index 00000000000000..d4ed92438ce5a8
--- /dev/null
+++ b/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
@@ -0,0 +1,55 @@
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: opt -passes=ctx-prof-flatten %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata -S -o - | FileCheck %s
+
+; CHECK-LABEL: entry:
+; CHECK:          br i1 %t, label %yes, label %no, !prof ![[C1:[0-9]+]]
+; CHECK-LABEL: no:
+; CHECK-NOT:   !prof
+; CHECK-LABEL: no1:
+; CHECK-NOT:   !prof
+; CHECK-LABEL: no2:
+; CHECK-NOT:   !prof
+; CHECK-LABEL: yes:
+; CHECK:          br i1 %t3, label %yes1, label %yes2, !prof ![[C1]]
+; CHECK-NOT:   !prof
+; CHECK: ![[C1]] = !{!"branch_weights", i32 6, i32 0}
+
+;--- example.ll
+define void @f1(i32 %cond) !guid !0 {
+entry:
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 0)
+  %t = icmp eq i32 %cond, 1
+  br i1 %t, label %yes, label %no
+
+no:
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 1)
+  %t2 = icmp eq i32 %cond, 2
+  br i1 %t2, label %no1, label %no2
+no1:
+  unreachable
+no2:
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 2)
+  unreachable
+yes:
+  %t3 = icmp eq i32 %cond, 3
+  br i1 %t3, label %yes1, label %yes2
+yes1:
+  br label %exit
+yes2:
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 3)
+  %t4 = icmp eq i32 %cond, 4
+  br i1 %t4, label %yes3, label %yes4
+yes3:
+  br label %exit
+yes4:
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 4)
+  unreachable
+exit:
+  ret void
+}
+
+!0 = !{i64 1234}
+
+;--- profile.json
+[{"Guid":1234, "Counters":[6,0,0,0,0]}]

``````````

</details>


https://github.com/llvm/llvm-project/pull/108467