From 0589762e4e42a21796ca74eeb356cdfc50eaa232 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n@gmail.com>
Date: Tue, 16 Jul 2024 21:30:36 +0800
Subject: [PATCH] [ValueTracking] Consistently propagate `DemandedElts` is
 `computeKnownFPClass`

Closes #99080
---
 llvm/include/llvm/Analysis/ValueTracking.h    | 22 ++++++++++++++-----
 llvm/lib/Analysis/ValueTracking.cpp           |  5 +++--
 .../Analysis/ValueTracking/known-fpclass.ll   |  7 +-----
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/Analysis/ValueTracking.h b/llvm/include/llvm/Analysis/ValueTracking.h
index 354ad5bc953175..2c2f965a3cd6f6 100644
--- a/llvm/include/llvm/Analysis/ValueTracking.h
+++ b/llvm/include/llvm/Analysis/ValueTracking.h
@@ -526,16 +526,17 @@ inline KnownFPClass computeKnownFPClass(
 }
 
 /// Wrapper to account for known fast math flags at the use instruction.
-inline KnownFPClass computeKnownFPClass(const Value *V, FastMathFlags FMF,
-                                        FPClassTest InterestedClasses,
-                                        unsigned Depth,
-                                        const SimplifyQuery &SQ) {
+inline KnownFPClass
+computeKnownFPClass(const Value *V, const APInt &DemandedElts,
+                    FastMathFlags FMF, FPClassTest InterestedClasses,
+                    unsigned Depth, const SimplifyQuery &SQ) {
   if (FMF.noNaNs())
     InterestedClasses &= ~fcNan;
   if (FMF.noInfs())
     InterestedClasses &= ~fcInf;
 
-  KnownFPClass Result = computeKnownFPClass(V, InterestedClasses, Depth, SQ);
+  KnownFPClass Result =
+      computeKnownFPClass(V, DemandedElts, InterestedClasses, Depth, SQ);
 
   if (FMF.noNaNs())
     Result.KnownFPClasses &= ~fcNan;
@@ -544,6 +545,17 @@ inline KnownFPClass computeKnownFPClass(const Value *V, FastMathFlags FMF,
   return Result;
 }
 
+inline KnownFPClass computeKnownFPClass(const Value *V, FastMathFlags FMF,
+                                        FPClassTest InterestedClasses,
+                                        unsigned Depth,
+                                        const SimplifyQuery &SQ) {
+  auto *FVTy = dyn_cast<FixedVectorType>(V->getType());
+  APInt DemandedElts =
+      FVTy ? APInt::getAllOnes(FVTy->getNumElements()) : APInt(1, 1);
+  return computeKnownFPClass(V, DemandedElts, FMF, InterestedClasses, Depth,
+                             SQ);
+}
+
 /// Return true if we can prove that the specified FP value is never equal to
 /// -0.0. Users should use caution when considering PreserveSign
 /// denormal-fp-math.
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index f54de030d33445..6e039ad2deadb8 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -5274,8 +5274,9 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
     }
       // reverse preserves all characteristics of the input vec's element.
     case Intrinsic::vector_reverse:
-      Known = computeKnownFPClass(II->getArgOperand(0), II->getFastMathFlags(),
-                                  InterestedClasses, Depth + 1, Q);
+      Known = computeKnownFPClass(
+          II->getArgOperand(0), DemandedElts.reverseBits(),
+          II->getFastMathFlags(), InterestedClasses, Depth + 1, Q);
       break;
     case Intrinsic::trunc:
     case Intrinsic::floor:
diff --git a/llvm/test/Analysis/ValueTracking/known-fpclass.ll b/llvm/test/Analysis/ValueTracking/known-fpclass.ll
index 225120584edef9..2b8e6298d746ae 100644
--- a/llvm/test/Analysis/ValueTracking/known-fpclass.ll
+++ b/llvm/test/Analysis/ValueTracking/known-fpclass.ll
@@ -26,12 +26,7 @@ define <4 x i1> @vector_reverse_fpclass2(<4 x double> nofpclass(nzero) %x) {
 
 define i1 @vector_reverse_fpclass_demanded(<4 x double> %vec, double nofpclass(nzero nan) %x) {
 ; CHECK-LABEL: @vector_reverse_fpclass_demanded(
-; CHECK-NEXT:    [[X_ABS:%.*]] = call double @llvm.fabs.f64(double [[X:%.*]])
-; CHECK-NEXT:    [[VEC_X:%.*]] = insertelement <4 x double> [[VEC:%.*]], double [[X_ABS]], i64 1
-; CHECK-NEXT:    [[REV:%.*]] = call <4 x double> @llvm.vector.reverse.v4f64(<4 x double> [[VEC_X]])
-; CHECK-NEXT:    [[ELE:%.*]] = extractelement <4 x double> [[REV]], i64 2
-; CHECK-NEXT:    [[CMP:%.*]] = fcmp oge double [[ELE]], 0.000000e+00
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
 
   %x.abs = call double @llvm.fabs.f64(double %x)