[llvm-branch-commits] [llvm] release/23.x: [PGO] Fix PGO instrumentation for selects with poison condition (#229002) (PR #230041)

[via llvm-branch-commits]

https://github.com/llvmbot created https://github.com/llvm/llvm-project/pull/230041

Backport 48102564a14886db6ba759325446907e2e1cfe71

Requested by: @nikic

>From ab4cd6da4e1b0803eb50c4c4caf694a704b08668 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 8 Oct 2026 10:05:47 +0200
Subject: [PATCH] [PGO] Fix PGO instrumentation for selects with poison
 condition (#229002)

PGO instrumentation for selects currently emits an
instrprof.increment.step, where the step is a zext of the select
condition (0 if false, 1 if true). However, if the condition is poison,
then the zext is poison and ends up taking an arbitray, possibly very
large, value, which then gets detected as a corrupted profile.

To avoid this, freeze the condition before extending. A poison select
condition can then add 0 or 1 to the counter, but no other value. I
believe that is the best we can do for poison conditions.

Fixes https://github.com/llvm/llvm-project/issues/228895.

(cherry picked from commit 48102564a14886db6ba759325446907e2e1cfe71)
---
 .../Transforms/Instrumentation/PGOInstrumentation.cpp    | 8 +++++++-
 llvm/test/Analysis/CtxProfAnalysis/handle-select.ll      | 5 +++--
 llvm/test/Transforms/PGOProfile/select1.ll               | 9 +++++++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index f6fcb41636000..86479f23f995a 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -66,6 +66,7 @@
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
@@ -1824,7 +1825,12 @@ void SelectInstVisitor::instrumentOneSelectInst(SelectInst &SI) {
   Module *M = F.getParent();
   IRBuilder<> Builder(&SI);
   Type *Int64Ty = Builder.getInt64Ty();
-  auto *Step = Builder.CreateZExt(SI.getCondition(), Int64Ty);
+  Value *Cond = SI.getCondition();
+  // Freeze the condition so that a poison condition can only increment by
+  // 0 or 1, but not some other value.
+  if (!isGuaranteedNotToBePoison(Cond))
+    Cond = Builder.CreateFreeze(Cond);
+  auto *Step = Builder.CreateZExt(Cond, Int64Ty);
   auto *NormalizedFuncNameVarPtr =
       ConstantExpr::getPointerBitCastOrAddrSpaceCast(
           FuncNameVar, PointerType::get(M->getContext(), 0));
diff --git a/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll b/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll
index 1880672580eb8..53f8aee99f9a0 100644
--- a/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll
+++ b/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll
@@ -21,8 +21,9 @@
 ; INSTR-NEXT:   call void @llvm.instrprof.increment(ptr @bar, i64 [[#]], i32 2, i32 0)
 ; INSTR-NEXT:   %inc =
 ; INSTR:        %test = icmp eq i32 %t, 0
-; INSTR-NEXT:   %1  = zext i1 %test to i64
-; INSTR-NEXT:   call void @llvm.instrprof.increment.step(ptr @bar, i64 [[#]], i32 2, i32 1, i64 %1)
+; INSTR-NEXT:   [[FREEZE:%.*]] = freeze i1 %test
+; INSTR-NEXT:   [[ZEXT:%.*]] = zext i1 [[FREEZE]] to i64
+; INSTR-NEXT:   call void @llvm.instrprof.increment.step(ptr @bar, i64 [[#]], i32 2, i32 1, i64 [[ZEXT]])
 ; INSTR-NEXT:   %res = select
 
 ; POST-INL-LABEL: yes:
diff --git a/llvm/test/Transforms/PGOProfile/select1.ll b/llvm/test/Transforms/PGOProfile/select1.ll
index 51626e6391218..decd58cddbef0 100644
--- a/llvm/test/Transforms/PGOProfile/select1.ll
+++ b/llvm/test/Transforms/PGOProfile/select1.ll
@@ -29,3 +29,12 @@ if.end:
   %retv = phi i32 [ %add, %if.then ], [ %sub, %if.else ]
   ret i32 %retv
 }
+
+define i32 @test_select_maybe_poison(i1 %cmp) {
+; GEN: %[[FREEZE:[0-9]+]] = freeze i1 %cmp
+; GEN: %[[STEP:[0-9]+]] = zext i1 %[[FREEZE]] to i64
+; GEN: call void @llvm.instrprof.increment.step({{.*}} i32 2, i32 1, i64 %[[STEP]])
+; NOSELECT-NOT: call void @llvm.instrprof.increment.step
+  %s = select i1 %cmp, i32 1, i32 2
+  ret i32 %s
+}