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

[via llvm-branch-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-pgo

Author: llvmbot

<details>
<summary>Changes</summary>

Backport 48102564a14886db6ba759325446907e2e1cfe71

Requested by: @<!-- -->nikic

---
Full diff: https://github.com/llvm/llvm-project/pull/230041.diff


3 Files Affected:

- (modified) llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp (+7-1) 
- (modified) llvm/test/Analysis/CtxProfAnalysis/handle-select.ll (+3-2) 
- (modified) llvm/test/Transforms/PGOProfile/select1.ll (+9) 


``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index f6fcb416360005..86479f23f995af 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 1880672580eb83..53f8aee99f9a0b 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 51626e63912183..decd58cddbef07 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
+}

``````````

</details>


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