[ConstantFolding] Fold sqrt poison -> poison (#141821)

I noticed this when a sqrt produced by VectorCombine with a poison
operand wasn't getting folded away to poison.

Most intrinsics in general could probably be folded to poison if one of
their arguments are poison too. Are there any exceptions to this we need
to be aware of?

Author: lukel97

llvm/lib/Analysis/ConstantFolding.cpp

@@ -2223,8 +2223,13 @@ static Constant *ConstantFoldScalarCall1(StringRef Name,
 
   if (isa<PoisonValue>(Operands[0])) {
     // TODO: All of these operations should probably propagate poison.
-    if (IntrinsicID == Intrinsic::canonicalize)
+    switch (IntrinsicID) {
+    case Intrinsic::canonicalize:
+    case Intrinsic::sqrt:
       return PoisonValue::get(Ty);
+    default:
+      break;
+    }
   }
 
   if (isa<UndefValue>(Operands[0])) {

llvm/test/Transforms/InstSimplify/fp-undef-poison.ll

@@ -293,3 +293,53 @@ define double @fmul_nnan_inf_op1(double %x) {
   %r = fmul nnan double %x, 0xfff0000000000000
   ret double %r
 }
+
+define float @sqrt_poison() {
+; CHECK-LABEL: @sqrt_poison(
+; CHECK-NEXT:    ret float poison
+;
+  %sqrt = call float @llvm.sqrt(float poison)
+  ret float %sqrt
+}
+
+define <2 x float> @sqrt_poison_fixed_vec() {
+; CHECK-LABEL: @sqrt_poison_fixed_vec(
+; CHECK-NEXT:    ret <2 x float> poison
+;
+  %sqrt = call <2 x float> @llvm.sqrt(<2 x float> poison)
+  ret <2 x float> %sqrt
+}
+
+define <2 x float> @sqrt_poison_elt_fixed_vec() {
+; CHECK-LABEL: @sqrt_poison_elt_fixed_vec(
+; CHECK-NEXT:    ret <2 x float> <float 1.000000e+00, float poison>
+;
+  %sqrt = call <2 x float> @llvm.sqrt(<2 x float> <float 1.0, float poison>)
+  ret <2 x float> %sqrt
+}
+
+define <vscale x 2 x float> @sqrt_poison_scalable_vec() {
+; CHECK-LABEL: @sqrt_poison_scalable_vec(
+; CHECK-NEXT:    ret <vscale x 2 x float> poison
+;
+  %sqrt = call <vscale x 2 x float> @llvm.sqrt(<vscale x 2 x float> poison)
+  ret <vscale x 2 x float> %sqrt
+}
+
+define float @sqrt_nnan_nan() {
+; CHECK-LABEL: @sqrt_nnan_nan(
+; CHECK-NEXT:    [[SQRT:%.*]] = call nnan float @llvm.sqrt.f32(float 0x7FF8000000000000)
+; CHECK-NEXT:    ret float [[SQRT]]
+;
+  %sqrt = call nnan float @llvm.sqrt(float 0x7ff8000000000000)
+  ret float %sqrt
+}
+
+define float @sqrt_ninf_inf() {
+; CHECK-LABEL: @sqrt_ninf_inf(
+; CHECK-NEXT:    [[SQRT:%.*]] = call ninf float @llvm.sqrt.f32(float 0xFFF0000000000000)
+; CHECK-NEXT:    ret float [[SQRT]]
+;
+  %sqrt = call ninf float @llvm.sqrt(float 0xfff0000000000000)
+  ret float %sqrt
+}