From 19b0e1227ca6653405e4a34627d04a14f2287f26 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke@igalia.com>
Date: Wed, 11 Jun 2025 13:27:14 +0200
Subject: [PATCH] [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?
---
 llvm/lib/Analysis/ConstantFolding.cpp         |  7 ++-
 .../InstSimplify/fp-undef-poison.ll           | 50 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 23ea6966fbf6..1ef0badd2375 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/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])) {
diff --git a/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll b/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll
index cb2026df962c..ffab9c94ddf4 100644
--- a/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll
+++ b/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
+}