ValueTracking: Handle non-splat vectors in computeKnownFPClass

[arsenm]

Summary

Avoids some regressions when the implementation of isKnownNeverNaN is
replaced with computeKnownFPClass.

[arsenm]

656b52a6c61e15dfbd5969417d66fcaf0e730ce1

llvm/lib/Analysis/ValueTracking.cpp

void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
   if (!DemandedElts) {
     // No demanded elts, better to assume we don't know anything.
     Known.resetAll();
     return;
   }
 
   assert(Depth <= MaxAnalysisRecursionDepth && "Limit Search Depth");
 
-  const APFloat *C;
-  if (match(V, m_APFloat(C))) {
-    // We know all of the classes for a scalar constant or a splat vector
-    // constant!
-    Known.KnownFPClasses = C->classify();
-    Known.SignBit = C->isNegative();
+  if (auto *CFP = dyn_cast_or_null<ConstantFP>(V)) {
+    Known.KnownFPClasses = CFP->getValueAPF().classify();
+    Known.SignBit = CFP->isNegative();
+    return;
+  }
+
+  // Try to handle fixed width vector constants
+  auto *VFVTy = dyn_cast<FixedVectorType>(V->getType());
+  const Constant *CV = dyn_cast<Constant>(V);
+  if (VFVTy && CV) {
+    Known.KnownFPClasses = fcNone;
+
+    // For vectors, verify that each element is not NaN.
+    unsigned NumElts = VFVTy->getNumElements();
+    for (unsigned i = 0; i != NumElts; ++i) {
+      Constant *Elt = CV->getAggregateElement(i);
+      if (!Elt) {
+        Known = KnownFPClass();
+        return;
+      }
+      if (isa<UndefValue>(Elt))
+        continue;
+      auto *CElt = dyn_cast<ConstantFP>(Elt);
+      if (!CElt) {
+        Known = KnownFPClass();
+        return;
+      }
+
+      KnownFPClass KnownElt{CElt->getValueAPF().classify(), CElt->isNegative()};
+      Known |= KnownElt;
+    }
+
     return;
   }
 
   FPClassTest KnownNotFromFlags = fcNone;
   if (const auto *CB = dyn_cast<CallBase>(V))
     KnownNotFromFlags |= CB->getRetNoFPClass();
   else if (const auto *Arg = dyn_cast<Argument>(V))
     KnownNotFromFlags |= Arg->getNoFPClass();

llvm/test/Transforms/Attributor/nofpclass.ll

 ; CHECK-LABEL: define noundef nofpclass(nan inf pzero sub norm) <2 x double> @returned_negzero_vector() {
 ; CHECK-NEXT:    call void @unknown()
 ; CHECK-NEXT:    ret <2 x double> <double -0.000000e+00, double -0.000000e+00>
 ;
   call void @unknown()
   ret <2 x double> <double -0.0, double -0.0>
 }
 
+; Test a vector element that's a constant but not ConstantFP.
+define <2 x double> @returned_strange_constant_vector_elt() {
+; CHECK-LABEL: define <2 x double> @returned_strange_constant_vector_elt() {
+; CHECK-NEXT:    call void @unknown()
+; CHECK-NEXT:    ret <2 x double> <double -0.000000e+00, double bitcast (i64 ptrtoint (ptr @unknown to i64) to double)>
+;
+  call void @unknown()
+  ret <2 x double> <double -0.0, double bitcast (i64 ptrtoint (ptr @unknown to i64) to double)>
+}
+
+; Test a vector element that's an undef/poison
+define <3 x double> @returned_undef_constant_vector_elt() {
+; CHECK-LABEL: define nofpclass(nan inf pzero sub norm) <3 x double> @returned_undef_constant_vector_elt() {
+; CHECK-NEXT:    call void @unknown()
+; CHECK-NEXT:    ret <3 x double> <double -0.000000e+00, double poison, double undef>
+;
+  call void @unknown()
+  ret <3 x double> <double -0.0, double poison, double undef>
+}
+
 define <2 x double> @returned_qnan_zero_vector() {
-; CHECK-LABEL: define noundef <2 x double> @returned_qnan_zero_vector() {
+; CHECK-LABEL: define noundef nofpclass(snan inf nzero sub norm) <2 x double> @returned_qnan_zero_vector() {
 ; CHECK-NEXT:    call void @unknown()
 ; CHECK-NEXT:    ret <2 x double> <double 0x7FF8000000000000, double 0.000000e+00>
 ;
   call void @unknown()
   ret <2 x double> <double 0x7FF8000000000000, double 0.0>
 }
 
 ; Return a float trivially nofpclass(nan) (call return attribute)