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

Mircea Trofin via llvm-commits llvm-commits at lists.llvm.org
Fri Sep 13 21:26:42 PDT 2024


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

>From e70420ec1894c31cd6d429cacd1220abc3ff58a0 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 12 Sep 2024 15:47:04 -0700
Subject: [PATCH 1/4] [ctx_prof] Fix checks in `PGOCtxprofFlattening`

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.
---
 .../Instrumentation/PGOCtxProfFlattening.cpp  | 65 ++++++++++++++++---
 .../CtxProfAnalysis/flatten-zero-path.ll      | 55 ++++++++++++++++
 2 files changed, 111 insertions(+), 9 deletions(-)
 create mode 100644 llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll

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]}]

>From 02735cf614dee1f642fc0e60084836cd068fe7e4 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 13 Sep 2024 14:58:49 -0700
Subject: [PATCH 2/4] fixes

---
 .../Instrumentation/PGOCtxProfFlattening.cpp  | 30 ++++++++++++-------
 .../CtxProfAnalysis/flatten-zero-path.ll      |  9 +++---
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
index 77cfa2cdcff7e8..cbb02bf5002be5 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
@@ -26,6 +26,7 @@
 #include "llvm/IR/Analysis.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
@@ -59,22 +60,28 @@ class ProfileAnnotator final {
 
     // Pass AssumeAllKnown when we try to propagate counts from edges to BBs -
     // because all the edge counters must be known.
-    uint64_t getEdgeSum(const SmallVector<EdgeInfo *> &Edges,
+    // Return std::nullopt if there were no edges to sum. The user can decide
+    // how to interpret that.
+    std::optional<uint64_t> getEdgeSum(const SmallVector<EdgeInfo *> &Edges,
                         bool AssumeAllKnown) const {
-      uint64_t Sum = 0;
+      std::optional<uint64_t> Sum;
       for (const auto *E : Edges)
-        if (E)
-          Sum += AssumeAllKnown ? *E->Count : E->Count.value_or(0U);
+        if (E) {
+          if (!Sum.has_value())
+            Sum = 0;
+          *Sum += (AssumeAllKnown ? *E->Count : E->Count.value_or(0U));
+        }
       return Sum;
     }
 
-    void computeCountFrom(const SmallVector<EdgeInfo *> &Edges) {
+    bool computeCountFrom(const SmallVector<EdgeInfo *> &Edges) {
       assert(!Count.has_value());
       Count = getEdgeSum(Edges, true);
+      return Count.has_value();
     }
 
     void setSingleUnknownEdgeCount(SmallVector<EdgeInfo *> &Edges) {
-      uint64_t KnownSum = getEdgeSum(Edges, false);
+      uint64_t KnownSum = getEdgeSum(Edges, false).value_or(0U);
       uint64_t EdgeVal = *Count > KnownSum ? *Count - KnownSum : 0U;
       EdgeInfo *E = nullptr;
       for (auto *I : Edges)
@@ -112,16 +119,14 @@ class ProfileAnnotator final {
 
     bool tryTakeCountFromKnownOutEdges(const BasicBlock &BB) {
       if (!UnknownCountOutEdges) {
-        computeCountFrom(OutEdges);
-        return true;
+        return computeCountFrom(OutEdges);
       }
       return false;
     }
 
     bool tryTakeCountFromKnownInEdges(const BasicBlock &BB) {
       if (!UnknownCountInEdges) {
-        computeCountFrom(InEdges);
-        return true;
+        return computeCountFrom(InEdges);
       }
       return false;
     }
@@ -179,7 +184,7 @@ class ProfileAnnotator final {
     bool KeepGoing = true;
     while (KeepGoing) {
       KeepGoing = false;
-      for (const auto &BB : reverse(F)) {
+      for (const auto &BB : F) {
         auto &Info = getBBInfo(BB);
         if (!Info.hasCount())
           KeepGoing |= Info.tryTakeCountFromKnownOutEdges(BB) ||
@@ -261,6 +266,9 @@ class ProfileAnnotator final {
                "profile is managed by IPO transforms");
         (void)Index;
         Count = Counters[Ins->getIndex()->getZExtValue()];
+      } else if (isa<UnreachableInst>(BB.getTerminator())) {
+        // The program presumably didn't crash.
+        Count = 0;
       }
       auto [It, Ins] =
           BBInfos.insert({&BB, {pred_size(&BB), succ_size(&BB), Count}});
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
index d4ed92438ce5a8..889a681a717705 100644
--- a/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
+++ b/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
@@ -23,13 +23,12 @@ entry:
   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)
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 1)
   unreachable
 yes:
   %t3 = icmp eq i32 %cond, 3
@@ -37,13 +36,13 @@ yes:
 yes1:
   br label %exit
 yes2:
-  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 3)
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 2)
   %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)
+  call void @llvm.instrprof.increment(ptr @f1, i64 42, i32 42, i32 3)
   unreachable
 exit:
   ret void
@@ -52,4 +51,4 @@ exit:
 !0 = !{i64 1234}
 
 ;--- profile.json
-[{"Guid":1234, "Counters":[6,0,0,0,0]}]
+[{"Guid":1234, "Counters":[6,0,0,0]}]

>From b44c43d76308a93d96378977ecdbfdf0b9f1740a Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 13 Sep 2024 15:12:44 -0700
Subject: [PATCH 3/4] clang format

---
 llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
index cbb02bf5002be5..389b7617df2e82 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp
@@ -63,7 +63,7 @@ class ProfileAnnotator final {
     // Return std::nullopt if there were no edges to sum. The user can decide
     // how to interpret that.
     std::optional<uint64_t> getEdgeSum(const SmallVector<EdgeInfo *> &Edges,
-                        bool AssumeAllKnown) const {
+                                       bool AssumeAllKnown) const {
       std::optional<uint64_t> Sum;
       for (const auto *E : Edges)
         if (E) {

>From a2884f68697e6ec9d8c2eadc51dbd2f62d43bbe1 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 13 Sep 2024 19:45:53 -0700
Subject: [PATCH 4/4] fixed test that now needs to deal with flattening

---
 llvm/test/ThinLTO/X86/ctxprof.ll | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/llvm/test/ThinLTO/X86/ctxprof.ll b/llvm/test/ThinLTO/X86/ctxprof.ll
index 1e30b90ec23d58..4baea3b25890eb 100644
--- a/llvm/test/ThinLTO/X86/ctxprof.ll
+++ b/llvm/test/ThinLTO/X86/ctxprof.ll
@@ -47,18 +47,21 @@
 ; NOPROF-1-NOT: m2_f1()
 ; NOPROF-2-NOT: m1_f1()
 ;
-; The run with workload definitions - same other options.
+; The run with workload definitions - same other options. We do need to re-generate the .bc
+; files, to include instrumentation.
+; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/m1.ll -o %t/m1-instr.bc
+; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/m2.ll -o %t/m2-instr.bc
 ;
 ; RUN: echo '[ \
 ; RUN:        {"Guid": 6019442868614718803, "Counters": [1], "Callsites": [[{"Guid": 15593096274670919754, "Counters": [1]}]]}, \
 ; RUN:        {"Guid": 15593096274670919754, "Counters": [1], "Callsites": [[{"Guid": 6019442868614718803, "Counters": [1]}]]} \
 ; RUN:  ]' > %t_exp/ctxprof.json
 ; RUN: llvm-ctxprof-util fromJSON --input %t_exp/ctxprof.json --output %t_exp/ctxprof.bitstream
-; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc \
+; RUN: llvm-lto2 run %t/m1-instr.bc %t/m2-instr.bc \
 ; RUN:  -o %t_exp/result.o -save-temps \
 ; RUN:  -use-ctx-profile=%t_exp/ctxprof.bitstream \
-; RUN:  -r %t/m1.bc,m1_f1,plx \
-; RUN:  -r %t/m2.bc,m2_f1,plx
+; RUN:  -r %t/m1-instr.bc,m1_f1,plx \
+; RUN:  -r %t/m2-instr.bc,m2_f1,plx
 ; RUN: llvm-dis %t_exp/result.o.1.3.import.bc -o - | FileCheck %s --check-prefix=FIRST
 ; RUN: llvm-dis %t_exp/result.o.2.3.import.bc -o - | FileCheck %s --check-prefix=SECOND
 ;



More information about the llvm-commits mailing list