[llvm] [PseudoProbe] Avoid inserting probes between musttail/deoptimize calls and returns (PR #211226)

[chandan singh via llvm-commits]

https://github.com/chandankds updated https://github.com/llvm/llvm-project/pull/211226

>From 8fea69e567924968c2802fd8085c742b234aa92c Mon Sep 17 00:00:00 2001
From: chandankds <chandankds at gmail.com>
Date: Fri, 24 Jul 2026 06:03:33 +0000
Subject: [PATCH] [PseudoProbe] Avoid inserting probes between
 musttail/deoptimize calls and returns

PseudoProbe insertion can generate invalid IR by inserting a probe between a
musttail/llvm.experimental.deoptimize call and its following ret. Insert the
probe before these calls to preserve the required instruction ordering.
---
 .../lib/Transforms/IPO/SampleProfileProbe.cpp | 12 +++++++++
 .../pseudo-probe-musttail-deopt.ll            | 26 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 llvm/test/Transforms/SampleProfile/pseudo-probe-musttail-deopt.ll

diff --git a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
index b6de00de07756..c84808557cb51 100644
--- a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
@@ -393,6 +393,18 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
       J = J->getNextNode();
     }
 
+    // A pseudo probe must not be inserted between a `musttail` or
+    // `llvm.experimental.deoptimize` call and its following `ret`, as this
+    // produces invalid IR. Such a call is required to immediately precede the
+    // block's `ret`, so only that position needs to be checked. Insert the
+    // probe before the call instead.
+    if (auto *Ret = dyn_cast<ReturnInst>(BB->getTerminator()))
+      if (auto *CI = dyn_cast_or_null<CallInst>(Ret->getPrevNode()))
+        if ((CI->isMustTailCall() ||
+             CI->getIntrinsicID() == Intrinsic::experimental_deoptimize) &&
+            !J->comesBefore(CI))
+          J = CI;
+
     IRBuilder<> Builder(J);
     assert(Builder.GetInsertPoint() != BB->end() &&
            "Cannot get the probing point");
diff --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-musttail-deopt.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-musttail-deopt.ll
new file mode 100644
index 0000000000000..1f70070427ec4
--- /dev/null
+++ b/llvm/test/Transforms/SampleProfile/pseudo-probe-musttail-deopt.ll
@@ -0,0 +1,26 @@
+; A pseudo-probe must not be inserted between a musttail or
+; llvm.experimental.deoptimize call and its following return, as this produces
+; invalid IR. Verify that probes are inserted before these calls instead.
+;
+; RUN: opt < %s -passes=pseudo-probe -S | FileCheck %s
+
+; CHECK-LABEL: define i32 @mt(ptr %p)
+; CHECK:      call void @llvm.pseudoprobe
+; CHECK-NEXT: %v = musttail call i32 @callee(ptr %p)
+; CHECK-NEXT: ret i32 %v
+define i32 @mt(ptr %p) {
+  %v = musttail call i32 @callee(ptr %p)
+  ret i32 %v
+}
+
+; CHECK-LABEL: define i32 @deopt()
+; CHECK:      call void @llvm.pseudoprobe
+; CHECK-NEXT: %v = call i32 (...) @llvm.experimental.deoptimize.i32() [ "deopt"() ]
+; CHECK-NEXT: ret i32 %v
+define i32 @deopt() {
+  %v = call i32 (...) @llvm.experimental.deoptimize.i32() [ "deopt"() ]
+  ret i32 %v
+}
+
+declare i32 @callee(ptr)
+declare i32 @llvm.experimental.deoptimize.i32(...)