[ctx_prof] API to get the instrumentation of a BB

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -92,7 +92,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

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) {

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