[IR] Disable select ? C : undef -> C fold in ConstantFoldSelectInstruction unless we know C isn't poison.

This matches the recent change to InstSimplify from D83440.

Differential Revision: https://reviews.llvm.org/D83535

Author: topperc

llvm/lib/IR/ConstantFold.cpp

@@ -779,10 +779,30 @@ Constant *llvm::ConstantFoldSelectInstruction(Constant *Cond,
     if (isa<UndefValue>(V1)) return V1;
     return V2;
   }
-  if (isa<UndefValue>(V1)) return V2;
-  if (isa<UndefValue>(V2)) return V1;
+
   if (V1 == V2) return V1;
 
+  // If the true or false value is undef, we can fold to the other value as
+  // long as the other value isn't poison.
+  auto NotPoison = [](Constant *C) {
+    // TODO: We can analyze ConstExpr by opcode to determine if there is any
+    //       possibility of poison.
+    if (isa<ConstantExpr>(C))
+      return false;
+
+    if (isa<ConstantInt>(C) || isa<GlobalVariable>(C) || isa<ConstantFP>(C) ||
+        isa<ConstantPointerNull>(C) || isa<Function>(C))
+      return true;
+
+    if (C->getType()->isVectorTy())
+      return !C->containsUndefElement() && !C->containsConstantExpression();
+
+    // TODO: Recursively analyze aggregates or other constants.
+    return false;
+  };
+  if (isa<UndefValue>(V1) && NotPoison(V2)) return V2;
+  if (isa<UndefValue>(V2) && NotPoison(V1)) return V1;
+
   if (ConstantExpr *TrueVal = dyn_cast<ConstantExpr>(V1)) {
     if (TrueVal->getOpcode() == Instruction::Select)
       if (TrueVal->getOperand(0) == Cond)

llvm/test/Transforms/InferAddressSpaces/AMDGPU/select.ll

@@ -221,7 +221,7 @@ define amdgpu_kernel void @store_select_group_global_mismatch_inttoptr_flat_null
 }
 
 ; CHECK-LABEL: @store_select_group_global_mismatch_undef_undef_constexpr(
-; CHECK: store i32 7, i32 addrspace(3)* null
+; CHECK: store i32 7, i32* select (i1 icmp eq (i32 ptrtoint (i32 addrspace(3)* @lds1 to i32), i32 4), i32* addrspacecast (i32 addrspace(3)* null to i32*), i32* undef), align 4
 define amdgpu_kernel void @store_select_group_global_mismatch_undef_undef_constexpr() #0 {
   store i32 7, i32* select (i1 icmp eq (i32 ptrtoint (i32 addrspace(3)* @lds1 to i32), i32 4), i32* addrspacecast (i32 addrspace(3)* null to i32*), i32* addrspacecast (i32 addrspace(1)* undef to i32*)), align 4
   ret void

llvm/test/Transforms/InstSimplify/select.ll

@@ -858,3 +858,73 @@ define <2 x i32> @false_undef_true_constextpr_vec(i1 %cond) {
   %s = select i1 %cond, <2 x i32> <i32 undef, i32 ptrtoint (i32* @g to i32)>, <2 x i32> <i32 ptrtoint (i32* @g to i32), i32 undef>
   ret <2 x i32> %s
 }
+
+define i32 @all_constant_true_undef() {
+; CHECK-LABEL: @all_constant_true_undef(
+; CHECK-NEXT:    ret i32 1
+;
+  %s = select i1 ptrtoint (i32 ()* @all_constant_true_undef to i1), i32 undef, i32 1
+  ret i32 %s
+}
+
+define float @all_constant_false_undef() {
+; CHECK-LABEL: @all_constant_false_undef(
+; CHECK-NEXT:    ret float 1.000000e+00
+;
+  %s = select i1 ptrtoint (float ()* @all_constant_false_undef to i1), float undef, float 1.0
+  ret float %s
+}
+
+define <2 x i32> @all_constant_true_undef_vec() {
+; CHECK-LABEL: @all_constant_true_undef_vec(
+; CHECK-NEXT:    ret <2 x i32> <i32 1, i32 -1>
+;
+  %s = select i1 ptrtoint (<2 x i32> ()* @all_constant_true_undef_vec to i1), <2 x i32> undef, <2 x i32> <i32 1, i32 -1>
+  ret <2 x i32> %s
+}
+
+define <2 x float> @all_constant_false_undef_vec() {
+; CHECK-LABEL: @all_constant_false_undef_vec(
+; CHECK-NEXT:    ret <2 x float> <float 1.000000e+00, float -1.000000e+00>
+;
+  %s = select i1 ptrtoint (<2 x float> ()* @all_constant_false_undef_vec to i1), <2 x float> undef, <2 x float> <float 1.0, float -1.0>
+  ret <2 x float> %s
+}
+
+; Negative tests. Don't fold if the non-undef operand is a constexpr.
+define i32 @all_constant_false_undef_true_constexpr() {
+; CHECK-LABEL: @all_constant_false_undef_true_constexpr(
+; CHECK-NEXT:    [[S:%.*]] = select i1 ptrtoint (i32 ()* @all_constant_false_undef_true_constexpr to i1), i32 ptrtoint (i32 ()* @all_constant_false_undef_true_constexpr to i32), i32 undef
+; CHECK-NEXT:    ret i32 [[S]]
+;
+  %s = select i1 ptrtoint (i32 ()* @all_constant_false_undef_true_constexpr to i1), i32 ptrtoint (i32 ()* @all_constant_false_undef_true_constexpr to i32), i32 undef
+  ret i32 %s
+}
+
+define i32 @all_constant_true_undef_false_constexpr() {
+; CHECK-LABEL: @all_constant_true_undef_false_constexpr(
+; CHECK-NEXT:    [[S:%.*]] = select i1 ptrtoint (i32 ()* @all_constant_true_undef_false_constexpr to i1), i32 undef, i32 ptrtoint (i32 ()* @all_constant_true_undef_false_constexpr to i32)
+; CHECK-NEXT:    ret i32 [[S]]
+;
+  %s = select i1 ptrtoint (i32 ()* @all_constant_true_undef_false_constexpr to i1), i32 undef, i32 ptrtoint (i32 ()* @all_constant_true_undef_false_constexpr to i32)
+  ret i32 %s
+}
+
+; Negative tests. Don't fold if the non-undef operand is a vector containing a constexpr.
+define <2 x i32> @all_constant_false_undef_true_constexpr_vec() {
+; CHECK-LABEL: @all_constant_false_undef_true_constexpr_vec(
+; CHECK-NEXT:    [[S:%.*]] = select i1 ptrtoint (<2 x i32> ()* @all_constant_false_undef_true_constexpr_vec to i1), <2 x i32> <i32 ptrtoint (<2 x i32> ()* @all_constant_false_undef_true_constexpr_vec to i32), i32 -1>, <2 x i32> undef
+; CHECK-NEXT:    ret <2 x i32> [[S]]
+;
+  %s = select i1 ptrtoint (<2 x i32> ()* @all_constant_false_undef_true_constexpr_vec to i1), <2 x i32> <i32 ptrtoint (<2 x i32> ()* @all_constant_false_undef_true_constexpr_vec to i32), i32 -1>, <2 x i32> undef
+  ret <2 x i32> %s
+}
+
+define <2 x i32> @all_constant_true_undef_false_constexpr_vec() {
+; CHECK-LABEL: @all_constant_true_undef_false_constexpr_vec(
+; CHECK-NEXT:    [[S:%.*]] = select i1 ptrtoint (<2 x i32> ()* @all_constant_true_undef_false_constexpr_vec to i1), <2 x i32> undef, <2 x i32> <i32 -1, i32 ptrtoint (<2 x i32> ()* @all_constant_true_undef_false_constexpr_vec to i32)>
+; CHECK-NEXT:    ret <2 x i32> [[S]]
+;
+  %s = select i1 ptrtoint (<2 x i32> ()* @all_constant_true_undef_false_constexpr_vec to i1), <2 x i32> undef, <2 x i32><i32 -1, i32 ptrtoint (<2 x i32> ()* @all_constant_true_undef_false_constexpr_vec to i32)>
+  ret <2 x i32> %s
+}