[llvm] [InstCombine] Enable select freeze poison folding when storing value (PR #129776)

[John McIver via llvm-commits]

https://github.com/jmciver updated https://github.com/llvm/llvm-project/pull/129776

>From 892c18ea5269e7e1055fa163e56ee24c0056503f Mon Sep 17 00:00:00 2001
From: John McIver <john.mciver.iii at gmail.com>
Date: Tue, 4 Mar 2025 14:10:38 -0700
Subject: [PATCH] [InstCombine] Enable select freeze poison folding when
 storing value

The non-freeze poison argument to select can be one of the following: global,
constant, and noundef arguments.

Alive2 test validation:
- https://alive2.llvm.org/ce/z/jbtCS6
- https://alive2.llvm.org/ce/z/YFA-5S
---
 .../InstCombine/InstructionCombining.cpp      | 16 ++++++--
 llvm/test/Transforms/InstCombine/select.ll    | 38 ++++++++++++++++---
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 4c14dcfb4d75f..c4829d8600b03 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -4812,7 +4812,14 @@ Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   //
   // TODO: This could use getBinopAbsorber() / getBinopIdentity() to avoid
   //       duplicating logic for binops at least.
-  auto getUndefReplacement = [&I](Type *Ty) {
+  auto getUndefReplacement = [&I, &AC = this->AC,
+                              &DT = this->DT](Type *Ty) -> Value * {
+    Value *SelectArgument = nullptr;
+    if (match(&I, m_OneUse(m_Freeze(m_Poison()))) &&
+        match(I.user_back(),
+              m_c_Select(m_Specific(&I), m_Value(SelectArgument))) &&
+        isGuaranteedNotToBeUndefOrPoison(SelectArgument, &AC, &I, &DT))
+      return SelectArgument;
     Constant *BestValue = nullptr;
     Constant *NullValue = Constant::getNullValue(Ty);
     for (const auto *U : I.users()) {
@@ -4842,8 +4849,11 @@ Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
 
   Constant *C;
   if (match(Op0, m_Constant(C)) && C->containsUndefOrPoisonElement()) {
-    Constant *ReplaceC = getUndefReplacement(I.getType()->getScalarType());
-    return replaceInstUsesWith(I, Constant::replaceUndefsWith(C, ReplaceC));
+    Value *Replace = getUndefReplacement(I.getType()->getScalarType());
+    if (Constant *ReplaceC = dyn_cast<Constant>(Replace))
+      return replaceInstUsesWith(I, Constant::replaceUndefsWith(C, ReplaceC));
+    else
+      return replaceInstUsesWith(I, Replace);
   }
 
   // Replace uses of Op with freeze(Op).
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index cc25f4ce24d9a..02e32d9a5ddcc 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -4904,8 +4904,7 @@ define i32 @src_simplify_2x_at_once_and(i32 %x, i32 %y) {
 
 define void @select_freeze_poison_parameter(ptr noundef %addr.src, ptr %addr.tgt, i1 %cond) {
 ; CHECK-LABEL: @select_freeze_poison_parameter(
-; CHECK-NEXT:    [[ADDR_SRC:%.*]] = select i1 [[COND:%.*]], ptr [[ADDR_SRC1:%.*]], ptr null
-; CHECK-NEXT:    store ptr [[ADDR_SRC]], ptr [[ADDR_TGT:%.*]], align 8
+; CHECK-NEXT:    store ptr [[ADDR_SRC:%.*]], ptr [[ADDR_TGT:%.*]], align 8
 ; CHECK-NEXT:    ret void
 ;
   %freeze = freeze ptr poison
@@ -4918,8 +4917,7 @@ define void @select_freeze_poison_parameter(ptr noundef %addr.src, ptr %addr.tgt
 
 define void @select_freeze_poison_global(ptr %addr.tgt, i1 %cond) {
 ; CHECK-LABEL: @select_freeze_poison_global(
-; CHECK-NEXT:    [[SELECT_ADDR:%.*]] = select i1 [[COND:%.*]], ptr @glb, ptr null
-; CHECK-NEXT:    store ptr [[SELECT_ADDR]], ptr [[ADDR_TGT:%.*]], align 8
+; CHECK-NEXT:    store ptr @glb, ptr [[ADDR_TGT:%.*]], align 8
 ; CHECK-NEXT:    ret void
 ;
   %freeze = freeze ptr poison
@@ -4930,8 +4928,7 @@ define void @select_freeze_poison_global(ptr %addr.tgt, i1 %cond) {
 
 define void @select_freeze_poison_constant(ptr %addr.tgt, i1 %cond) {
 ; CHECK-LABEL: @select_freeze_poison_constant(
-; CHECK-NEXT:    [[SELECT_ADDR:%.*]] = select i1 [[COND:%.*]], i32 72, i32 0
-; CHECK-NEXT:    store i32 [[SELECT_ADDR]], ptr [[ADDR_TGT:%.*]], align 4
+; CHECK-NEXT:    store i32 72, ptr [[ADDR_TGT:%.*]], align 4
 ; CHECK-NEXT:    ret void
 ;
   %freeze = freeze i32 poison
@@ -4939,3 +4936,32 @@ define void @select_freeze_poison_constant(ptr %addr.tgt, i1 %cond) {
   store i32 %select.addr, ptr %addr.tgt
   ret void
 }
+
+define <2 x i8> @select_freeze_poison_mask_vector(i1 %cond, <2 x i8> noundef %y) {
+; CHECK-LABEL: @select_freeze_poison_mask_vector(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> [[Y:%.*]], <2 x i8> zeroinitializer
+; CHECK-NEXT:    ret <2 x i8> [[SEL]]
+;
+  %freeze = freeze <2 x i8> <i8 0, i8 poison>
+  %sel = select i1 %cond, <2 x i8> %y, <2 x i8> %freeze
+  ret <2 x i8> %sel
+}
+
+define <2 x i8> @select_freeze_poison_splat_vector(i1 %cond, <2 x i8> noundef %y) {
+; CHECK-LABEL: @select_freeze_poison_splat_vector(
+; CHECK-NEXT:    ret <2 x i8> [[Y:%.*]]
+;
+  %freeze = freeze <2 x i8> <i8 poison, i8 poison>
+  %sel = select i1 %cond, <2 x i8> %y, <2 x i8> %freeze
+  ret <2 x i8> %sel
+}
+
+define <2 x i8> @select_freeze_constant_vector(i1 %cond, <2 x i8> noundef %y) {
+; CHECK-LABEL: @select_freeze_constant_vector(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], <2 x i8> [[Y:%.*]], <2 x i8> zeroinitializer
+; CHECK-NEXT:    ret <2 x i8> [[SEL]]
+;
+  %freeze = freeze <2 x i8> <i8 0, i8 0>
+  %sel = select i1 %cond, <2 x i8> %y, <2 x i8> %freeze
+  ret <2 x i8> %sel
+}