[InstCombine] Fix for folding select-like shufflevector into floating point binary operators.

[michele-scandale]

Folding a select-like shufflevector into a floating point binary operators can only be done if the result is preserved for both case. In particular, if the common operand of the shufflevector and the floating point binary operator can be a NaN, then the transformation won't preserve the result value.

[michele-scandale]

This should address the issue reported in #74326

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Michele Scandale (michele-scandale)

Changes

Folding a select-like shufflevector into a floating point binary operators can only be done if the result is preserved for both case. In particular, if the common operand of the shufflevector and the floating point binary operator can be a NaN, then the transformation won't preserve the result value.

---

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

3 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+16-3)
• (modified) llvm/test/Transforms/InstCombine/shuffle_select-inseltpoison.ll (+3-3)
• (modified) llvm/test/Transforms/InstCombine/shuffle_select.ll (+3-3)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 3c4c0f35eb6d48..e3252633153930 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2135,7 +2135,8 @@ static Instruction *foldSelectShuffleOfSelectShuffle(ShuffleVectorInst &Shuf) {
   return new ShuffleVectorInst(X, Y, NewMask);
 }
 
-static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
+static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf,
+                                                const SimplifyQuery &SQ) {
   assert(Shuf.isSelect() && "Must have select-equivalent shuffle");
 
   // Are we shuffling together some value and that same value after it has been
@@ -2159,6 +2160,18 @@ static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   if (!IdC)
     return nullptr;
 
+  Value *X = Op0IsBinop ? Op1 : Op0;
+
+  // Prevent folding in the case the non-binop operand might have NaN values.
+  // If X can have NaN elements then we have that the floating point math
+  // operation in the transformed code may not preserve the exact NaN
+  // bit-pattern -- e.g. `fadd sNaN, 0.0 -> qNaN`.
+  // This makes the transformation incorrect since the original program would
+  // have preserved the exact NaN bit-pattern.
+  // Avoid the folding if X can have NaN elements.
+  if (Shuf.getType()->isFPOrFPVectorTy() && !isKnownNeverNaN(X, 0, SQ))
+    return nullptr;
+
   // Shuffle identity constants into the lanes that return the original value.
   // Example: shuf (mul X, {-1,-2,-3,-4}), X, {0,5,6,3} --> mul X, {-1,1,1,-4}
   // Example: shuf X, (add X, {-1,-2,-3,-4}), {0,1,6,7} --> add X, {0,0,-3,-4}
@@ -2175,7 +2188,6 @@ static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
 
   // shuf (bop X, C), X, M --> bop X, C'
   // shuf X, (bop X, C), M --> bop X, C'
-  Value *X = Op0IsBinop ? Op1 : Op0;
   Instruction *NewBO = BinaryOperator::Create(BOpcode, X, NewC);
   NewBO->copyIRFlags(BO);
 
@@ -2241,7 +2253,8 @@ Instruction *InstCombinerImpl::foldSelectShuffle(ShuffleVectorInst &Shuf) {
   if (Instruction *I = foldSelectShuffleOfSelectShuffle(Shuf))
     return I;
 
-  if (Instruction *I = foldSelectShuffleWith1Binop(Shuf))
+  if (Instruction *I = foldSelectShuffleWith1Binop(
+          Shuf, getSimplifyQuery().getWithInstruction(&Shuf)))
     return I;
 
   BinaryOperator *B0, *B1;
diff --git a/llvm/test/Transforms/InstCombine/shuffle_select-inseltpoison.ll b/llvm/test/Transforms/InstCombine/shuffle_select-inseltpoison.ll
index 44ec77e471bb41..9b4f80d3150758 100644
--- a/llvm/test/Transforms/InstCombine/shuffle_select-inseltpoison.ll
+++ b/llvm/test/Transforms/InstCombine/shuffle_select-inseltpoison.ll
@@ -336,7 +336,7 @@ define <4 x i32> @srem(<4 x i32> %v) {
 
 ; Try FP ops/types.
 
-define <4 x float> @fadd(<4 x float> %v) {
+define <4 x float> @fadd(<4 x float> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fadd(
 ; CHECK-NEXT:    [[S:%.*]] = fadd <4 x float> [[V:%.*]], <float 4.100000e+01, float 4.200000e+01, float -0.000000e+00, float -0.000000e+00>
 ; CHECK-NEXT:    ret <4 x float> [[S]]
@@ -359,7 +359,7 @@ define <4 x double> @fsub(<4 x double> %v) {
 
 ; Propagate any FMF.
 
-define <4 x float> @fmul(<4 x float> %v) {
+define <4 x float> @fmul(<4 x float> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fmul(
 ; CHECK-NEXT:    [[S:%.*]] = fmul nnan ninf <4 x float> [[V:%.*]], <float 4.100000e+01, float 1.000000e+00, float 1.000000e+00, float 1.000000e+00>
 ; CHECK-NEXT:    ret <4 x float> [[S]]
@@ -380,7 +380,7 @@ define <4 x double> @fdiv_constant_op0(<4 x double> %v) {
   ret <4 x double> %s
 }
 
-define <4 x double> @fdiv_constant_op1(<4 x double> %v) {
+define <4 x double> @fdiv_constant_op1(<4 x double> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fdiv_constant_op1(
 ; CHECK-NEXT:    [[S:%.*]] = fdiv reassoc <4 x double> [[V:%.*]], <double undef, double 1.000000e+00, double 4.300000e+01, double 4.400000e+01>
 ; CHECK-NEXT:    ret <4 x double> [[S]]
diff --git a/llvm/test/Transforms/InstCombine/shuffle_select.ll b/llvm/test/Transforms/InstCombine/shuffle_select.ll
index a1b0d782b554f9..aae02a13dc8cd6 100644
--- a/llvm/test/Transforms/InstCombine/shuffle_select.ll
+++ b/llvm/test/Transforms/InstCombine/shuffle_select.ll
@@ -336,7 +336,7 @@ define <4 x i32> @srem(<4 x i32> %v) {
 
 ; Try FP ops/types.
 
-define <4 x float> @fadd(<4 x float> %v) {
+define <4 x float> @fadd(<4 x float> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fadd(
 ; CHECK-NEXT:    [[S:%.*]] = fadd <4 x float> [[V:%.*]], <float 4.100000e+01, float 4.200000e+01, float -0.000000e+00, float -0.000000e+00>
 ; CHECK-NEXT:    ret <4 x float> [[S]]
@@ -359,7 +359,7 @@ define <4 x double> @fsub(<4 x double> %v) {
 
 ; Propagate any FMF.
 
-define <4 x float> @fmul(<4 x float> %v) {
+define <4 x float> @fmul(<4 x float> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fmul(
 ; CHECK-NEXT:    [[S:%.*]] = fmul nnan ninf <4 x float> [[V:%.*]], <float 4.100000e+01, float 1.000000e+00, float 1.000000e+00, float 1.000000e+00>
 ; CHECK-NEXT:    ret <4 x float> [[S]]
@@ -380,7 +380,7 @@ define <4 x double> @fdiv_constant_op0(<4 x double> %v) {
   ret <4 x double> %s
 }
 
-define <4 x double> @fdiv_constant_op1(<4 x double> %v) {
+define <4 x double> @fdiv_constant_op1(<4 x double> nofpclass(nan) %v) {
 ; CHECK-LABEL: @fdiv_constant_op1(
 ; CHECK-NEXT:    [[S:%.*]] = fdiv reassoc <4 x double> [[V:%.*]], <double undef, double 1.000000e+00, double 4.300000e+01, double 4.400000e+01>
 ; CHECK-NEXT:    ret <4 x double> [[S]]

[goldsteinn]

Please keep the old tests so we can see the fix actually take affect.
You can add new tests with nofpclass(nan).

[michele-scandale]

I've added the test case fadd_maybe_nan to check that the transformation doesn't occur if the value might be a NaN. The other tests cases where I added nofpclass(nan) seems to cover the propagation of the fast-math flags in the transformation. Not sure if it is worth duplicating those test cases.

[arsenm]

Shut can't be a scalar FP type, so could just query the element type