[llvm] [ctx_prof] Fix `ProfileAnnotator::allTakenPathsExit` (PR #109183)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/109183

None

>From 6d23eef034b0b8dece993a8bd2294594f1c7992e Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 17 Sep 2024 21:46:47 -0700
Subject: [PATCH] [ctx_prof] Fix `ProfileAnnotator::allTakenPathsExit`

---
 .../Instrumentation/PGOCtxProfFlattening.cpp  | 29 ++++---
 .../CtxProfAnalysis/flatten-check-path.ll     | 85 +++++++++++++++++++
 2 files changed, 104 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll

diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
index e76689e2f5f0a5..91f950e2ba4c3e 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
@@ -233,28 +233,37 @@ class ProfileAnnotator final {
     std::deque<const BasicBlock *> Worklist;
     DenseSet<const BasicBlock *> Visited;
     Worklist.push_back(&F.getEntryBlock());
-    Visited.insert(&F.getEntryBlock());
+    bool HitExit = false;
     while (!Worklist.empty()) {
       const auto *BB = Worklist.front();
       Worklist.pop_front();
-      if (succ_size(BB) <= 1)
+      if (!Visited.insert(BB).second)
         continue;
+      if (succ_size(BB) == 0) {
+        if (isa<UnreachableInst>(BB->getTerminator()))
+          return false;
+        HitExit = true;
+        continue;
+      }
+      if (succ_size(BB) == 1) {
+        llvm::append_range(Worklist, successors(BB));
+        continue;
+      }
       const auto &BBInfo = getBBInfo(*BB);
-      bool Inserted = false;
+      bool HasAWayOut = 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 (BBInfo.getEdgeCount(I) > 0) {
+            HasAWayOut = true;
+            Worklist.push_back(Succ);
+          }
         }
       }
-      if (!Inserted)
+      if (!HasAWayOut)
         return false;
     }
-    return true;
+    return HitExit;
   }
 
 public:
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll
new file mode 100644
index 00000000000000..dbe14c3522a8a9
--- /dev/null
+++ b/llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll
@@ -0,0 +1,85 @@
+; REQUIRES: asserts
+; Check that the profile annotator works: we hit an exit and non-zero paths to
+; already visited blocks are not disconsidered.
+;
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_ok.json --output=%t/profile_ok.ctxprofdata
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_pump.json --output=%t/profile_pump.ctxprofdata
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_unreachable.json --output=%t/profile_unreachable.ctxprofdata
+;
+; RUN: opt -passes=ctx-prof-flatten %t/example_ok.ll -use-ctx-profile=%t/profile_ok.ctxprofdata -S -o - | FileCheck %s
+; RUN: not --crash opt -passes=ctx-prof-flatten %t/message_pump.ll -use-ctx-profile=%t/profile_pump.ctxprofdata -S 2>&1 | FileCheck %s --check-prefix=ASSERTION
+; RUN: not --crash opt -passes=ctx-prof-flatten %t/unreachable.ll -use-ctx-profile=%t/profile_unreachable.ctxprofdata -S 2>&1 | FileCheck %s --check-prefix=ASSERTION
+
+; CHECK: br i1 %x, label %b1, label %exit, !prof ![[PROF1:[0-9]+]]
+; CHECK: br i1 %y, label %blk, label %exit, !prof ![[PROF2:[0-9]+]]
+; CHECK: ![[PROF1]] = !{!"branch_weights", i32 1, i32 1}
+; CHECK: ![[PROF2]] = !{!"branch_weights", i32 0, i32 1}
+; ASSERTION: Assertion `allTakenPathsExit()
+
+; b1->exit is the only way out from b1, but the exit block would have been
+; already visited from blk. That should not result in an assertion, though.
+;--- example_ok.ll
+define void @foo(i32 %t) !guid !0 {
+entry:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 0)
+  br label %blk
+blk:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 1)
+  %x = icmp eq i32 %t, 0
+  br i1 %x, label %b1, label %exit
+b1:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 2)
+  %y = icmp eq i32 %t, 0
+  br i1 %y, label %blk, label %exit
+exit:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 3)
+  ret void
+}
+!0 = !{i64 1234}
+
+;--- profile_ok.json
+[{"Guid":1234, "Counters":[2, 2, 1, 2]}]
+
+;--- message_pump.ll
+; This is a message pump: the loop never exits. This should result in an
+; assertion because we can't reach an exit BB
+
+define void @foo(i32 %t) !guid !0 {
+entry:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 0)  
+  br label %blk
+blk:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 1)
+  %x = icmp eq i32 %t, 0
+  br i1 %x, label %blk, label %exit
+exit:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 2)
+  ret void
+}
+!0 = !{i64 1234}
+
+;--- profile_pump.json
+[{"Guid":1234, "Counters":[2, 10, 0]}]
+
+;--- unreachable.ll
+; An unreachable block is reached, that's an error
+define void @foo(i32 %t) !guid !0 {
+entry:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 0)
+  br label %blk
+blk:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 1)
+  %x = icmp eq i32 %t, 0
+  br i1 %x, label %b1, label %exit
+b1:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 2)
+  unreachable
+exit:
+  call void @llvm.instrprof.increment(ptr @foo, i64 42, i32 42, i32 3)
+  ret void
+}
+!0 = !{i64 1234}
+
+;--- profile_unreachable.json
+[{"Guid":1234, "Counters":[2, 1, 1, 2]}]
\ No newline at end of file