[InstSimplify] Fix a scalable-vector crash

[frasercrmck]

Summary

D143505 fixed/simplified folding of operations with SNaN operands. In
doing so it introduced a crash when handling scalable vector types,
wherein the scalable-vector ConstantVector was cast to a ConstantFP.

Since we know by that point in the code that if we've found a NaN, we're
dealing with a scalable-vector splat (as there are no other kinds of
scalable-vector constant for which that holds), we can grab the splatted
value and re-use the existing code, which will automatically splat the
new NaN back to a scalable vector for us.

[arsenm]

llvm/lib/Analysis/InstructionSimplify.cpp

5471

Why is this down here and not another dyn_cast<ScalableVector> after the FixedVectorType case?

llvm/test/Transforms/InstSimplify/fp-nan.ll

59

Add some scalable vector tests with other FP classes?

[frasercrmck]

llvm/lib/Analysis/InstructionSimplify.cpp

5471

I did it here because the lines above:

if (!In->isNaN())
    return ConstantFP::getNaN(Ty);

Work on scalable vectors, so we don't have to re-implement that part. note how the fixed-vector case does do that.

llvm/test/Transforms/InstSimplify/fp-nan.ll

59

Ah good point. I initially left those out because I thought those fixed-vector tests were there to check non-splat behaviour (which we can't do in scalable vectors) but yeah now I see it's also checking different FP classes.

I'll add the appropriate scalable tests (one for each class below, I suppose)

[frasercrmck]

• extend scalable vector testing

llvm/lib/Analysis/InstructionSimplify.cpp

if (auto *VecTy = dyn_cast<FixedVectorType>(Ty)) {
     return ConstantVector::get(NewC);
   }
 
   // If it is not a fixed vector, but not a simple NaN either, return a
   // canonical NaN.
   if (!In->isNaN())
     return ConstantFP::getNaN(Ty);
 
+  // If we known this is a NaN, and it's scalable vector, we must have a splat
+  // on our hands. Grab that before splatting a QNaN constant.
+  if (isa<ScalableVectorType>(Ty)) {
+    auto *Splat = In->getSplatValue();
+    assert(Splat && Splat->isNaN() &&
+           "Found a scalable-vector NaN but not a splat");
+    In = Splat;
+  }
+
   // Propagate an existing QNaN constant. If it is an SNaN, make it quiet, but
   // preserve the sign/payload.
   return ConstantFP::get(Ty, cast<ConstantFP>(In)->getValue().makeQuiet());
 }
 
 /// Perform folds that are common to any floating-point operation. This implies
 /// transforms based on poison/undef/NaN because the operation itself makes no
 /// difference to the result.

llvm/test/Transforms/InstSimplify/fp-nan.ll

 define <2 x float> @fsub_nan_op1_vec(<2 x float> %x) {
 ; CHECK-LABEL: @fsub_nan_op1_vec(
 ; CHECK-NEXT:    ret <2 x float> <float 0x7FF9000000000000, float 0xFFF9000000000000>
 ;
   %r = fsub <2 x float> %x, <float 0x7FF1000000000000, float 0xFFF1000000000000>
   ret <2 x float> %r
 }
 
+define <vscale x 1 x float> @fsub_nan_op1_scalable_vec_0(<vscale x 1 x float> %x) {
+; CHECK-LABEL: @fsub_nan_op1_scalable_vec_0(
+; CHECK-NEXT:    ret <vscale x 1 x float> shufflevector (<vscale x 1 x float> insertelement (<vscale x 1 x float> poison, float 0x7FF9000000000000, i64 0), <vscale x 1 x float> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fsub <vscale x 1 x float> %x, shufflevector (<vscale x 1 x float> insertelement (<vscale x 1 x float> poison, float 0x7FF1000000000000, i64 0), <vscale x 1 x float> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x float> %r
+}
+
+define <vscale x 1 x float> @fsub_nan_op1_scalable_vec_1(<vscale x 1 x float> %x) {
+; CHECK-LABEL: @fsub_nan_op1_scalable_vec_1(
+; CHECK-NEXT:    ret <vscale x 1 x float> shufflevector (<vscale x 1 x float> insertelement (<vscale x 1 x float> poison, float 0xFFF9000000000000, i64 0), <vscale x 1 x float> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fsub <vscale x 1 x float> %x, shufflevector (<vscale x 1 x float> insertelement (<vscale x 1 x float> poison, float 0xFFF1000000000000, i64 0), <vscale x 1 x float> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x float> %r
+}
+
 ; Signaling and signed - make quiet and preserve the payload and signbit
 
 define double @fmul_nan_op0(double %x) {
 ; CHECK-LABEL: @fmul_nan_op0(
 ; CHECK-NEXT:    ret double 0xFFF8000000000001
 ;
   %r = fmul double 0xFFF0000000000001, %x
   ret double %r
 }
 
 define <2 x double> @fmul_nan_op0_vec(<2 x double> %x) {
 ; CHECK-LABEL: @fmul_nan_op0_vec(
 ; CHECK-NEXT:    ret <2 x double> <double 0xFFF8000000000001, double 0xFFF8DEADDEADDEAD>
 ;
   %r = fmul <2 x double> <double 0xFFF0000000000001, double 0xFFF0DEADDEADDEAD>, %x
   ret <2 x double> %r
 }
 
+define <vscale x 1 x double> @fmul_nan_op0_scalable_vec_0(<vscale x 1 x double> %x) {
+; CHECK-LABEL: @fmul_nan_op0_scalable_vec_0(
+; CHECK-NEXT:    ret <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF8000000000001, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fmul <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF0000000000001, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer), %x
+  ret <vscale x 1 x double> %r
+}
+
+define <vscale x 1 x double> @fmul_nan_op0_scalable_vec_1(<vscale x 1 x double> %x) {
+; CHECK-LABEL: @fmul_nan_op0_scalable_vec_1(
+; CHECK-NEXT:    ret <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF8DEADDEADDEAD, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fmul <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF0DEADDEADDEAD, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer), %x
+  ret <vscale x 1 x double> %r
+}
+
 ; Vector type
 
 define <2 x float> @fmul_nan_op1(<2 x float> %x) {
 ; CHECK-LABEL: @fmul_nan_op1(
 ; CHECK-NEXT:    ret <2 x float> <float 0x7FF8000000000000, float 0x7FF8000000000000>
 ;
   %r = fmul <2 x float> %x, <float 0x7FF8000000000000, float 0x7FF8000000000000>
   ret <2 x float> %r
 }
 
+define <vscale x 1 x double> @fmul_nan_op1_scalable_vec(<vscale x 1 x double> %x) {
+; CHECK-LABEL: @fmul_nan_op1_scalable_vec(
+; CHECK-NEXT:    ret <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0x7FF8000000000000, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fmul <vscale x 1 x double> %x, shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0x7FF8000000000000, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x double> %r
+}
+
 ; Vector signed and non-zero payload
 
 define <2 x double> @fdiv_nan_op0(<2 x double> %x) {
 ; CHECK-LABEL: @fdiv_nan_op0(
 ; CHECK-NEXT:    ret <2 x double> <double 0xFFF800000000000F, double 0xFFF800000000000F>
 ;
   %r = fdiv <2 x double> <double 0xFFF800000000000F, double 0xFFF800000000000F>, %x
   ret <2 x double>  %r
 }
 
+define <vscale x 1 x double> @fdivl_nan_op0_scalable_vec(<vscale x 1 x double> %x) {
+; CHECK-LABEL: @fdivl_nan_op0_scalable_vec(
+; CHECK-NEXT:    ret <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF800000000000F, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fdiv <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF800000000000F, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer), %x
+  ret <vscale x 1 x double> %r
+}
+
 ; Vector with different NaN constant elements
 
 define <2 x half> @fdiv_nan_op1(<2 x half> %x) {
 ; CHECK-LABEL: @fdiv_nan_op1(
 ; CHECK-NEXT:    ret <2 x half> <half 0xH7FFF, half 0xHFF00>
 ;
   %r = fdiv <2 x half> %x, <half 0xH7FFF, half 0xHFF00>
   ret <2 x half> %r
 }
 
+define <vscale x 1 x half> @fdiv_nan_op1_scalable_vec(<vscale x 1 x half> %x) {
+; CHECK-LABEL: @fdiv_nan_op1_scalable_vec(
+; CHECK-NEXT:    ret <vscale x 1 x half> shufflevector (<vscale x 1 x half> insertelement (<vscale x 1 x half> poison, half 0xH7FFF, i64 0), <vscale x 1 x half> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fdiv <vscale x 1 x half> %x, shufflevector (<vscale x 1 x half> insertelement (<vscale x 1 x half> poison, half 0xH7FFF, i64 0), <vscale x 1 x half> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x half> %r
+}
+
 ; Vector with poison element
 
 define <2 x double> @fsub_nan_poison_op1(<2 x double> %x) {
 ; CHECK-LABEL: @fsub_nan_poison_op1(
 ; CHECK-NEXT:    ret <2 x double> <double 0xFFFF00000000DEAD, double poison>
 ;
   %r = fsub <2 x double> %x, <double 0xFFFF00000000DEAD, double poison>
   ret <2 x double> %r

 define <2 x double> @fneg_nan_2(<2 x double> %x) {
 ; CHECK-LABEL: @fneg_nan_2(
 ; CHECK-NEXT:    ret <2 x double> <double 0xFFF9234567890ABC, double 0x7FF8000000000001>
 ;
   %r = fsub <2 x double> <double -0.0, double -0.0>, <double 0xFFF1234567890ABC, double 0x7FF0000000000001>
   ret <2 x double> %r
 }
 
+define <vscale x 1 x double> @fneg_nan_2_scalable_vec() {
+; CHECK-LABEL: @fneg_nan_2_scalable_vec(
+; CHECK-NEXT:    ret <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF9234567890ABC, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+;
+  %r = fsub <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double -0.0, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer), shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF1234567890ABC, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x double> %r
+}
+
 define <2 x double> @unary_fneg_nan_2(<2 x double> %x) {
 ; CHECK-LABEL: @unary_fneg_nan_2(
 ; CHECK-NEXT:    ret <2 x double> <double 0x7FF1234567890ABC, double 0xFFF0000000000001>
 ;
   %r = fneg <2 x double> <double 0xFFF1234567890ABC, double 0x7FF0000000000001>
   ret <2 x double> %r
 }
 
+; FIXME: This doesn't behave the same way as the fixed-length vectors above
+define <vscale x 1 x double> @unary_fneg_nan_2_scalable_vec_0() {
+; CHECK-LABEL: @unary_fneg_nan_2_scalable_vec_0(
+; CHECK-NEXT:    [[R:%.*]] = fneg <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF1234567890ABC, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 1 x double> [[R]]
+;
+  %r = fneg <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0xFFF1234567890ABC, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x double> %r
+}
+
+; FIXME: This doesn't behave the same way as the fixed-length vectors above
+define <vscale x 1 x double> @unary_fneg_nan_2_scalable_vec_1() {
+; CHECK-LABEL: @unary_fneg_nan_2_scalable_vec_1(
+; CHECK-NEXT:    [[R:%.*]] = fneg <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0x7FF0000000000001, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 1 x double> [[R]]
+;
+  %r = fneg <vscale x 1 x double> shufflevector (<vscale x 1 x double> insertelement (<vscale x 1 x double> poison, double 0x7FF0000000000001, i64 0), <vscale x 1 x double> poison, <vscale x 1 x i32> zeroinitializer)
+  ret <vscale x 1 x double> %r
+}
+
 ; Repeat all tests with fast-math-flags. Alternate 'nnan' and 'fast' for more coverage.
 
 define float @fadd_nan_op0_nnan(float %x) {
 ; CHECK-LABEL: @fadd_nan_op0_nnan(
 ; CHECK-NEXT:    ret float poison
 ;
   %r = fadd nnan float 0x7FF8000000000000, %x
   ret float %r