[llvm-branch-commits] [llvm] [ctxprof] Extend the notion of "cannot return" (PR #135651)

[Mircea Trofin via llvm-branch-commits]

https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/135651

None

>From 41540073aaa8adc96e6f7889df66e1791cb4dbc9 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 14 Apr 2025 10:03:55 -0700
Subject: [PATCH] [ctxprof] Extend the notion of "cannot return"

---
 .../Instrumentation/PGOCtxProfLowering.cpp    | 19 ++++++++------
 .../ctx-instrumentation-invalid-roots.ll      | 25 +++++++++++--------
 .../PGOProfile/ctx-instrumentation.ll         | 13 ++++++++++
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp b/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
index a314457423819..603022d94838c 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp
@@ -9,6 +9,7 @@
 
 #include "llvm/Transforms/Instrumentation/PGOCtxProfLowering.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CtxProfAnalysis.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/IR/Analysis.h"
@@ -102,6 +103,12 @@ std::pair<uint32_t, uint32_t> getNumCountersAndCallsites(const Function &F) {
   }
   return {NumCounters, NumCallsites};
 }
+
+void emitUnsupportedRoot(const Function &F, StringRef Reason) {
+  F.getContext().emitError("[ctxprof] The function " + F.getName() +
+                           " was indicated as context root but " + Reason +
+                           ", which is not supported.");
+}
 } // namespace
 
 // set up tie-in with compiler-rt.
@@ -144,12 +151,8 @@ CtxInstrumentationLowerer::CtxInstrumentationLowerer(Module &M,
       for (const auto &BB : *F)
         for (const auto &I : BB)
           if (const auto *CB = dyn_cast<CallBase>(&I))
-            if (CB->isMustTailCall()) {
-              M.getContext().emitError(
-                  "The function " + Fname +
-                  " was indicated as a context root, but it features musttail "
-                  "calls, which is not supported.");
-            }
+            if (CB->isMustTailCall())
+              emitUnsupportedRoot(*F, "it features musttail calls");
     }
   }
 
@@ -210,11 +213,13 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
 
   // Probably pointless to try to do anything here, unlikely to be
   // performance-affecting.
-  if (F.doesNotReturn()) {
+  if (!llvm::canReturn(F)) {
     for (auto &BB : F)
       for (auto &I : make_early_inc_range(BB))
         if (isa<InstrProfCntrInstBase>(&I))
           I.eraseFromParent();
+    if (ContextRootSet.contains(&F))
+      emitUnsupportedRoot(F, "it does not return");
     return true;
   }
 
diff --git a/llvm/test/Transforms/PGOProfile/ctx-instrumentation-invalid-roots.ll b/llvm/test/Transforms/PGOProfile/ctx-instrumentation-invalid-roots.ll
index 454780153b823..b5ceb4602c60b 100644
--- a/llvm/test/Transforms/PGOProfile/ctx-instrumentation-invalid-roots.ll
+++ b/llvm/test/Transforms/PGOProfile/ctx-instrumentation-invalid-roots.ll
@@ -1,17 +1,22 @@
-; RUN: not opt -passes=ctx-instr-gen,ctx-instr-lower -profile-context-root=good \
-; RUN:   -profile-context-root=bad \
-; RUN:   -S < %s 2>&1 | FileCheck %s
+; RUN: split-file %s %t
+; RUN: not opt -passes=ctx-instr-gen,ctx-instr-lower -profile-context-root=the_func -S %t/musttail.ll -o - 2>&1 | FileCheck %s
+; RUN: not opt -passes=ctx-instr-gen,ctx-instr-lower -profile-context-root=the_func -S %t/unreachable.ll -o - 2>&1 | FileCheck %s
+; RUN: not opt -passes=ctx-instr-gen,ctx-instr-lower -profile-context-root=the_func -S %t/noreturn.ll -o - 2>&1 | FileCheck %s
 
+;--- musttail.ll
 declare void @foo()
 
-define void @good() {
-  call void @foo()
-  ret void
-}
-
-define void @bad() {
+define void @the_func() {
   musttail call void @foo()
   ret void
 }
+;--- unreachable.ll
+define void @the_func() {
+  unreachable
+}
+;--- noreturn.ll
+define void @the_func() noreturn {
+  unreachable
+}
 
-; CHECK: error: The function bad was indicated as a context root, but it features musttail calls, which is not supported.
+; CHECK: error: [ctxprof] The function the_func was indicated as context root
diff --git a/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll b/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
index 8f72711a9c8b1..6afa37ef286f5 100644
--- a/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
+++ b/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
@@ -323,6 +323,18 @@ define void @does_not_return() noreturn {
 ;
   unreachable
 }
+
+define void @unreachable() {
+; INSTRUMENT-LABEL: define void @unreachable() {
+; INSTRUMENT-NEXT:    call void @llvm.instrprof.increment(ptr @unreachable, i64 742261418966908927, i32 1, i32 0)
+; INSTRUMENT-NEXT:    unreachable
+;
+; LOWERING-LABEL: define void @unreachable(
+; LOWERING-SAME: ) !guid [[META9:![0-9]+]] {
+; LOWERING-NEXT:    unreachable
+;
+  unreachable
+}
 ;.
 ; LOWERING: attributes #[[ATTR0]] = { noreturn }
 ; LOWERING: attributes #[[ATTR1:[0-9]+]] = { nounwind }
@@ -340,4 +352,5 @@ define void @does_not_return() noreturn {
 ; LOWERING: [[META6]] = !{i64 -3771893999295659109}
 ; LOWERING: [[META7]] = !{i64 -4680624981836544329}
 ; LOWERING: [[META8]] = !{i64 5519225910966780583}
+; LOWERING: [[META9]] = !{i64 -565652589829076809}
 ;.