[nfc][ctx_prof] Factor the callsite instrumentation exclusion criteria (#108471)

Reusing this in the logic fetching the instrumentation in `CtxProfAnalysis`.

Author: mtrofin

llvm/include/llvm/IR/IntrinsicInst.h

@@ -1587,6 +1587,12 @@ class InstrProfCallsite : public InstrProfCntrInstBase {
   static bool classof(const Value *V) {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
+  // We instrument direct calls (but not to intrinsics), or indirect calls.
+  static bool canInstrumentCallsite(const CallBase &CB) {
+    return !CB.isInlineAsm() &&
+           (CB.isIndirectCall() ||
+            (CB.getCalledFunction() && !CB.getCalledFunction()->isIntrinsic()));
+  }
   Value *getCallee() const;
   void setCallee(Value *Callee);
 };

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -234,7 +234,7 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
 }
 
 InstrProfCallsite *CtxProfAnalysis::getCallsiteInstrumentation(CallBase &CB) {
-  while (auto *Prev = CB.getPrevNode())
+  for (auto *Prev = CB.getPrevNode(); Prev; Prev = Prev->getPrevNode())
     if (auto *IPC = dyn_cast<InstrProfCallsite>(Prev))
       return IPC;
   return nullptr;

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -945,9 +945,7 @@ void FunctionInstrumenter::instrument() {
       for (auto &BB : F)
         for (auto &Instr : BB)
           if (auto *CS = dyn_cast<CallBase>(&Instr)) {
-            if ((CS->getCalledFunction() &&
-                 CS->getCalledFunction()->isIntrinsic()) ||
-                dyn_cast<InlineAsm>(CS->getCalledOperand()))
+            if (!InstrProfCallsite::canInstrumentCallsite(*CS))
               continue;
             Visitor(CS);
           }

llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll

@@ -250,6 +250,26 @@ define void @no_counters() {
   call void @bar()
   ret void
 }
+
+; Ensure "calls" to inline asm don't get callsite-instrumented.
+define void @inlineasm() {
+; INSTRUMENT-LABEL: define void @inlineasm() {
+; INSTRUMENT-NEXT:    call void @llvm.instrprof.increment(ptr @inlineasm, i64 742261418966908927, i32 1, i32 0)
+; INSTRUMENT-NEXT:    call void asm "nop", ""()
+; INSTRUMENT-NEXT:    ret void
+;
+; LOWERING-LABEL: define void @inlineasm(
+; LOWERING-SAME: ) !guid [[META6:![0-9]+]] {
+; LOWERING-NEXT:    [[TMP1:%.*]] = call ptr @__llvm_ctx_profile_get_context(ptr @inlineasm, i64 -3771893999295659109, i32 1, i32 0)
+; LOWERING-NEXT:    [[TMP2:%.*]] = ptrtoint ptr [[TMP1]] to i64
+; LOWERING-NEXT:    [[TMP3:%.*]] = and i64 [[TMP2]], -2
+; LOWERING-NEXT:    [[TMP4:%.*]] = inttoptr i64 [[TMP3]] to ptr
+; LOWERING-NEXT:    call void asm "nop", ""()
+; LOWERING-NEXT:    ret void
+;
+  call void asm "nop", ""()
+  ret void
+}
 ;.
 ; LOWERING: attributes #[[ATTR0:[0-9]+]] = { nounwind }
 ; LOWERING: attributes #[[ATTR1:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
@@ -262,4 +282,5 @@ define void @no_counters() {
 ; LOWERING: [[META3]] = !{i64 -3006003237940970099}
 ; LOWERING: [[META4]] = !{i64 5679753335911435902}
 ; LOWERING: [[META5]] = !{i64 5458232184388660970}
+; LOWERING: [[META6]] = !{i64 -3771893999295659109}
 ;.

llvm/unittests/Analysis/CtxProfAnalysisTest.cpp

@@ -64,6 +64,11 @@ define void @another_entrypoint_no_callees(i32 %a) {
   ret void
 }
 
+define void @inlineasm() {
+  call void asm "nop", ""()
+  ret void
+}
+
 attributes #0 = { noinline }
 !0 = !{ i64 11872291593386833696 }
 )IR";
@@ -115,6 +120,27 @@ TEST_F(CtxProfAnalysisTest, GetCallsiteIDTest) {
   EXPECT_THAT(InsValues, testing::ElementsAre(0, 1));
 }
 
+TEST_F(CtxProfAnalysisTest, GetCallsiteIDInlineAsmTest) {
+  ModulePassManager MPM;
+  MPM.addPass(PGOInstrumentationGen(PGOInstrumentationType::CTXPROF));
+  EXPECT_FALSE(MPM.run(*M, MAM).areAllPreserved());
+  auto *F = M->getFunction("inlineasm");
+  ASSERT_NE(F, nullptr);
+  std::vector<const Instruction *> InsValues;
+
+  for (auto &BB : *F)
+    for (auto &I : BB)
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        // Skip instrumentation inserted intrinsics.
+        if (CB->getCalledFunction() && CB->getCalledFunction()->isIntrinsic())
+          continue;
+        auto *Ins = CtxProfAnalysis::getCallsiteInstrumentation(*CB);
+        InsValues.push_back(Ins);
+      }
+
+  EXPECT_THAT(InsValues, testing::ElementsAre(nullptr));
+}
+
 TEST_F(CtxProfAnalysisTest, GetCallsiteIDNegativeTest) {
   auto *F = M->getFunction("foo");
   ASSERT_NE(F, nullptr);