[llvm-branch-commits] [llvm] [ctx_prof] API to get the instrumentation of a BB (PR #105468)

[Mircea Trofin via llvm-branch-commits]

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

>From c5ee379ec43215d8268219ec3cfced3f7f730fc8 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 20 Aug 2024 21:09:16 -0700
Subject: [PATCH] [ctx_prof] API to get the instrumentation of a BB

---
 llvm/include/llvm/Analysis/CtxProfAnalysis.h  |  5 +++++
 llvm/lib/Analysis/CtxProfAnalysis.cpp         |  7 ++++++
 .../Analysis/CtxProfAnalysisTest.cpp          | 22 +++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/llvm/include/llvm/Analysis/CtxProfAnalysis.h b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
index 23abcbe2c6e9d2..0b4dd8ae3a0dc7 100644
--- a/llvm/include/llvm/Analysis/CtxProfAnalysis.h
+++ b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
@@ -95,7 +95,12 @@ class CtxProfAnalysis : public AnalysisInfoMixin<CtxProfAnalysis> {
 
   PGOContextualProfile run(Module &M, ModuleAnalysisManager &MAM);
 
+  /// Get the instruction instrumenting a callsite, or nullptr if that cannot be
+  /// found.
   static InstrProfCallsite *getCallsiteInstrumentation(CallBase &CB);
+
+  /// Get the instruction instrumenting a BB, or nullptr if not present.
+  static InstrProfIncrementInst *getBBInstrumentation(BasicBlock &BB);
 };
 
 class CtxProfAnalysisPrinterPass
diff --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/llvm/lib/Analysis/CtxProfAnalysis.cpp
index fffc8de2b36c8e..46daa4a4506189 100644
--- a/llvm/lib/Analysis/CtxProfAnalysis.cpp
+++ b/llvm/lib/Analysis/CtxProfAnalysis.cpp
@@ -202,6 +202,13 @@ InstrProfCallsite *CtxProfAnalysis::getCallsiteInstrumentation(CallBase &CB) {
   return nullptr;
 }
 
+InstrProfIncrementInst *CtxProfAnalysis::getBBInstrumentation(BasicBlock &BB) {
+  for (auto &I : BB)
+    if (auto *Incr = dyn_cast<InstrProfIncrementInst>(&I))
+      return Incr;
+  return nullptr;
+}
+
 static void
 preorderVisit(const PGOCtxProfContext::CallTargetMapTy &Profiles,
               function_ref<void(const PGOCtxProfContext &)> Visitor) {
diff --git a/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp b/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
index 5f9bf3ec540eb3..fbe3a6e45109cc 100644
--- a/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
+++ b/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
@@ -132,4 +132,26 @@ TEST_F(CtxProfAnalysisTest, GetCallsiteIDNegativeTest) {
   EXPECT_EQ(IndIns, nullptr);
 }
 
+TEST_F(CtxProfAnalysisTest, GetBBIDTest) {
+  ModulePassManager MPM;
+  MPM.addPass(PGOInstrumentationGen(PGOInstrumentationType::CTXPROF));
+  EXPECT_FALSE(MPM.run(*M, MAM).areAllPreserved());
+  auto *F = M->getFunction("foo");
+  ASSERT_NE(F, nullptr);
+  std::map<std::string, int> BBNameAndID;
+
+  for (auto &BB : *F) {
+    auto *Ins = CtxProfAnalysis::getBBInstrumentation(BB);
+    if (Ins)
+      BBNameAndID[BB.getName().str()] =
+          static_cast<int>(Ins->getIndex()->getZExtValue());
+    else
+      BBNameAndID[BB.getName().str()] = -1;
+  }
+
+  EXPECT_THAT(BBNameAndID,
+              testing::UnorderedElementsAre(
+                  testing::Pair("", 0), testing::Pair("yes", 1),
+                  testing::Pair("no", -1), testing::Pair("exit", -1)));
+}
 } // namespace