[llvm] 82266d3 - [nfc][ctx_prof] Factor the callsite instrumentation exclusion criteria (#108471)

[via llvm-commits]

Author: Mircea Trofin
Date: 2024-09-13T21:25:47-07:00
New Revision: 82266d3a2b33f49165c0f24d3db5ea9875cc706c

URL: https://github.com/llvm/llvm-project/commit/82266d3a2b33f49165c0f24d3db5ea9875cc706c
DIFF: https://github.com/llvm/llvm-project/commit/82266d3a2b33f49165c0f24d3db5ea9875cc706c.diff

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

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

Added: 
    

Modified: 
    llvm/include/llvm/IR/IntrinsicInst.h
    llvm/lib/Analysis/CtxProfAnalysis.cpp
    llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
    llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
    llvm/unittests/Analysis/CtxProfAnalysisTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/IntrinsicInst.h b/llvm/include/llvm/IR/IntrinsicInst.h
index fc8d1b3d1947e3..4458126ffa759d 100644
--- a/llvm/include/llvm/IR/IntrinsicInst.h
+++ b/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);
 };

diff  --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/llvm/lib/Analysis/CtxProfAnalysis.cpp
index 560d9c742d5e7d..c29709b613410e 100644
--- a/llvm/lib/Analysis/CtxProfAnalysis.cpp
+++ b/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;

diff  --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 9dd4a561edfddc..e5a522dba28f79 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/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);
           }

diff  --git a/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll b/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
index c94c2b4da57a98..1927060de868ea 100644
--- a/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
+++ b/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}
 ;.

diff  --git a/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp b/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
index fbe3a6e45109cc..3fba07ddd0f248 100644
--- a/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
+++ b/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);