[ConstantFolding] Fold sqrt poison -> poison

[lukel97]

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?

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Luke Lau (lukel97)

Changes

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

This patch calls simplifyFPOp to simplify sqrt, similar to fma/fmuladd, and folds poison as well as nnan/ninf with nan/inf.

There's a lot of other FP intrinsics that probably need to go through this. And 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?

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/InstructionSimplify.cpp (+5)
• (modified) llvm/test/Transforms/InstSimplify/fp-undef-poison.ll (+40)

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 23e147ba8c6a1..f4d1bfb1e9554 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -6306,6 +6306,11 @@ static Value *simplifyUnaryIntrinsic(Function *F, Value *Op0,
     if (computeKnownFPSignBit(Op0, /*Depth=*/0, Q) == false)
       return Op0;
     break;
+  case Intrinsic::sqrt:
+    if (Value *V = simplifyFPOp(Op0, Call->getFastMathFlags(), Q, fp::ebIgnore,
+                                RoundingMode::NearestTiesToEven))
+      return V;
+    break;
   case Intrinsic::bswap:
     // bswap(bswap(x)) -> x
     if (match(Op0, m_BSwap(m_Value(X))))
diff --git a/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll b/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll
index cb2026df962c8..537473873486e 100644
--- a/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll
+++ b/llvm/test/Transforms/InstSimplify/fp-undef-poison.ll
@@ -293,3 +293,43 @@ 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 <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:    ret float poison
+;
+  %sqrt = call nnan float @llvm.sqrt(float 0x7ff8000000000000)
+  ret float %sqrt
+}
+
+define float @sqrt_ninf_inf() {
+; CHECK-LABEL: @sqrt_ninf_inf(
+; CHECK-NEXT:    ret float poison
+;
+  %sqrt = call ninf float @llvm.sqrt(float 0xfff0000000000000)
+  ret float %sqrt
+}

[nikic]

I think the correct starting point here would be to support poison propagation in ConstantFolding rather than InstSimplify. I think

llvm-project/llvm/lib/Analysis/ConstantFolding.cpp

Line 2209 in b035580

if (IntrinsicID == Intrinsic::canonicalize)

is the relevant code for this case.

To generalize things, something to consider would be to handle this centrally based on the propagatesPoison() VT query. Which, incidentally, does not list the FP intrinsics.

[lukel97]

I think the correct starting point here would be to support poison propagation in ConstantFolding rather than InstSimplify. I think

llvm-project/llvm/lib/Analysis/ConstantFolding.cpp

Line 2209 in b035580

if (IntrinsicID == Intrinsic::canonicalize)

is the relevant code for this case.
To generalize things, something to consider would be to handle this centrally based on the propagatesPoison() VT query. Which, incidentally, does not list the FP intrinsics.

I saw that case in ConstantFolding but it looks like we currently only fold scalar intrinsics, and for my use case I need to fold scalable vectors.

I think fixed length vectors would get their poison propagated if we extended the case in ConstantFoldScalarCall1. Maybe we should look into constant scalable vector intrinsics with splatted operands instead then?

[nikic]

I think fixed length vectors would get their poison propagated if we extended the case in ConstantFoldScalarCall1.

Yes, fixed vectors are folded element-wise.

Maybe we should look into constant scalable vector intrinsics with splatted operands instead then?

That would make sense to me. Surprised this isn't already the case...

[lukel97]

I've updated this to just handle this in ConstantFoldScalarCall1, and I've left around the test cases for scalable vectors + nnan/ninf for later (let me know if you want me to take them out though)

I'll work on a patch to constant fold scalable intrinsics where the operands are splatted next.

But I guess we'll still need to figure out how to handle the case whenever not all arguments are constant, e.g only one operand is poison. Should that be handled by ValueTracking?

[dtcxzyw]

But I guess we'll still need to figure out how to handle the case whenever not all arguments are constant, e.g only one operand is poison. Should that be handled by ValueTracking?

Yeah. Generally we should handle these cases in propagatesPoison and simplifyIntrinsic. Unfortunately, we have to repeat this in ConstantFolding for non-splat vectors :(

[dtcxzyw]

LG

[dtcxzyw]

Can you add a case with <float 1.0, float poison>?