[llvm] [GVNSink] Don't sink pseudo probes (PR #224669)

[Timur Baidusenov via llvm-commits]

https://github.com/bai-tim updated https://github.com/llvm/llvm-project/pull/224669

>From 9fa582266ae71d7bc586553d01edfc990313507f Mon Sep 17 00:00:00 2001
From: Timur Baidusenov <timurbaidusenov at gmail.com>
Date: Fri, 18 Sep 2026 17:20:12 +0300
Subject: [PATCH] [GVNSink] Don't sink pseudo probes

llvm.pseudoprobe marks the block it is in, so sinking the probes of two sibling
blocks into their common successor misattributes samples and, when the indices
differ, merges them into a phi node that extractProbe() does not expect. Report
the probe from shouldAvoidSinkingInstruction(), as EarlyCSE and SimplifyCFG
skip pseudo probes for the same reason.

Fixes #166306.
---
 llvm/lib/Transforms/Scalar/GVNSink.cpp       |  3 +-
 llvm/test/Transforms/GVNSink/pseudo-probe.ll | 88 ++++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/GVNSink/pseudo-probe.ll

diff --git a/llvm/lib/Transforms/Scalar/GVNSink.cpp b/llvm/lib/Transforms/Scalar/GVNSink.cpp
index eb7a9c68b5b618..2963bf0bffec20 100644
--- a/llvm/lib/Transforms/Scalar/GVNSink.cpp
+++ b/llvm/lib/Transforms/Scalar/GVNSink.cpp
@@ -50,6 +50,7 @@
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
@@ -536,7 +537,7 @@ class GVNSink {
   bool shouldAvoidSinkingInstruction(Instruction *I) {
     // These instructions may change or break semantics if moved.
     if (isa<PHINode>(I) || I->isEHPad() || isa<AllocaInst>(I) ||
-        I->getType()->isTokenTy())
+        I->getType()->isTokenTy() || isa<PseudoProbeInst>(I))
       return true;
     return false;
   }
diff --git a/llvm/test/Transforms/GVNSink/pseudo-probe.ll b/llvm/test/Transforms/GVNSink/pseudo-probe.ll
new file mode 100644
index 00000000000000..9f53cfa83152b0
--- /dev/null
+++ b/llvm/test/Transforms/GVNSink/pseudo-probe.ll
@@ -0,0 +1,88 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S < %s -passes=gvn-sink | FileCheck %s
+
+declare void @llvm.pseudoprobe(i64, i64, i32, i64)
+
+; A pseudo probe marks the block it is in, so sinking it would attribute the
+; samples of that block to another one. Sinking two probes with different
+; indices would also leave a phi node as the index, which extractProbe() does
+; not expect. The instructions after the probe are still sunk.
+
+define i32 @different_indices(i1 %c, i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @different_indices(
+; CHECK-SAME: i1 [[C:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[C]], label %[[IF_THEN:.*]], label %[[IF_ELSE:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+; CHECK-NEXT:    [[SUB0:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    br label %[[IF_END:.*]]
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    call void @llvm.pseudoprobe(i64 1234, i64 3, i32 0, i64 -1)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    br label %[[IF_END]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB1_SINK:%.*]] = phi i32 [ [[SUB1]], %[[IF_ELSE]] ], [ [[SUB0]], %[[IF_THEN]] ]
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i32 [[SUB1_SINK]], [[A]]
+; CHECK-NEXT:    ret i32 [[MUL1]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:
+  call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+  %sub0 = sub i32 %a, %b
+  %mul0 = mul i32 %sub0, %a
+  br label %if.end
+
+if.else:
+  call void @llvm.pseudoprobe(i64 1234, i64 3, i32 0, i64 -1)
+  %sub1 = sub i32 %a, %b
+  %mul1 = mul i32 %sub1, %a
+  br label %if.end
+
+if.end:
+  %phi = phi i32 [ %mul0, %if.then ], [ %mul1, %if.else ]
+  ret i32 %phi
+}
+
+; Probes with the same index do not need a phi node to be sunk, but merging
+; them would still misattribute the samples of one block to the other.
+
+define i32 @same_index(i1 %c, i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @same_index(
+; CHECK-SAME: i1 [[C:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[C]], label %[[IF_THEN:.*]], label %[[IF_ELSE:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+; CHECK-NEXT:    [[SUB0:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    br label %[[IF_END:.*]]
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+; CHECK-NEXT:    [[SUB1:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    br label %[[IF_END]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB1_SINK:%.*]] = phi i32 [ [[SUB1]], %[[IF_ELSE]] ], [ [[SUB0]], %[[IF_THEN]] ]
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i32 [[SUB1_SINK]], [[A]]
+; CHECK-NEXT:    ret i32 [[MUL1]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:
+  call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+  %sub0 = sub i32 %a, %b
+  %mul0 = mul i32 %sub0, %a
+  br label %if.end
+
+if.else:
+  call void @llvm.pseudoprobe(i64 1234, i64 2, i32 0, i64 -1)
+  %sub1 = sub i32 %a, %b
+  %mul1 = mul i32 %sub1, %a
+  br label %if.end
+
+if.end:
+  %phi = phi i32 [ %mul0, %if.then ], [ %mul1, %if.else ]
+  ret i32 %phi
+}