ValueTracking: Implement computeKnownFPClass for sqrt

[arsenm]

Summary

Could be slightly smarter in cases that are probably uninteresting.

[arsenm]

Update unit test that was now too smart

[jcranmer-intel]

llvm/lib/Analysis/ValueTracking.cpp

4449

This also includes -subnormal in DAZ=preserve-sign mode.

[arsenm]

llvm/lib/Analysis/ValueTracking.cpp

4449

The comment could elaborate on the preserve sign handling, but I think this is correct as is? There is a pre-existing bug in cannotBeOrderedLessThanZeroImpl for denormal handling

[arsenm]

Fix DAZ handling

[arsenm]

ping

[arsenm]

ping

[arsenm]

ping

[foad]

llvm/lib/Analysis/ValueTracking.cpp

4135

There's a better version of this in D148230 that uses inputDenormalIsIEEEOrPosZero.

4445

This should reuse info from the first call to computeKnownFPClass, instead of making a second call to cannotBeOrderedLessThanZeroImpl.

[arsenm]

Rebase to use new stuff

llvm/lib/Analysis/ValueTracking.cpp

4445

Oh right, this predates when I added cannotBeOrderedLessThanZero to KnownFPClass

[foad]

llvm/include/llvm/Analysis/ValueTracking.h

259

"known"

llvm/lib/Analysis/ValueTracking.cpp

4442–4443

Is this valid? You haven't implemented snan logic for most of the other intrinsics.

4447

Should really add OrderedLessThanZeroMask to the InterestedClasses for the recursive call for this case.

[arsenm]

llvm/lib/Analysis/ValueTracking.cpp

4442–4443

Yes, I need to go back to refine this. You are never allowed to introduce signaling nans. Non-constrained ops just don't guarantee quieting. The propagateNaN helper introduced later helps with this (although it doesn't quite apply here since there are additional nan conditions

[foad]

LGTM with or without the InterestedClasses improvement.

[arsenm]

Interested classes

[arsenm]

fe5786d41644ec6c739b6ace49d654a21a07cca0

llvm/include/llvm/Analysis/ValueTracking.h

bool isKnownNeverNaN() const {
     return (KnownFPClasses & fcNan) == fcNone;
   }
 
   /// Return true if it's known this can never be an infinity.
   bool isKnownNeverInfinity() const {
     return (KnownFPClasses & fcInf) == fcNone;
   }
 
+  /// Return true if it's known this can never be +infinity.
+  bool isKnownNeverPosInfinity() const {
+    return (KnownFPClasses & fcPosInf) == fcNone;
+  }
+
   /// Return true if it's known this can never be a zero.
   bool isKnownNeverZero() const {
     return (KnownFPClasses & fcZero) == fcNone;
   }
 
   KnownFPClass &operator|=(const KnownFPClass &RHS) {
     KnownFPClasses = KnownFPClasses | RHS.KnownFPClasses;
 
     if (SignBit != RHS.SignBit)
       SignBit = std::nullopt;
     return *this;

llvm/lib/Analysis/ValueTracking.cpp

 /// \p F for \p Val.
 static bool inputDenormalIsIEEE(const Function &F, const Value *Val) {
   Type *Ty = Val->getType()->getScalarType();
   return F.getDenormalMode(Ty->getFltSemantics()).Input == DenormalMode::IEEE;
 }
 
 /// Returns a pair of values, which if passed to llvm.is.fpclass, returns the
 /// same result as an fcmp with the given operands.
 std::pair<Value *, FPClassTest> llvm::fcmpToClassTest(FCmpInst::Predicate Pred,
                                                       const Function &F,
                                                       Value *LHS, Value *RHS,
                                                       bool LookThroughSrc) {
   const APFloat *ConstRHS;
   if (!match(RHS, m_APFloat(ConstRHS)))
     return {nullptr, fcNone};
 
   if (ConstRHS->isZero()) {
if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(Op)) {
 
         computeKnownFPClass(II->getArgOperand(0), DemandedElts,
                             InterestedClasses, Known, Depth + 1, Q, TLI);
         computeKnownFPClass(II->getArgOperand(1), DemandedElts,
                             InterestedClasses, KnownSign, Depth + 1, Q, TLI);
         Known.copysign(KnownSign);
         break;
       }
+      case Intrinsic::sqrt: {
+        KnownFPClass KnownSrc;
+        computeKnownFPClass(II->getArgOperand(0), DemandedElts,
+                            InterestedClasses, KnownSrc, Depth + 1, Q, TLI);
+
+        if (KnownSrc.isKnownNeverPosInfinity())
+          Known.knownNot(fcPosInf);
+
+        if (KnownSrc.isKnownNeverNaN() &&
+            (InterestedClasses & fcNan) != fcNone &&
+            cannotBeOrderedLessThanZeroImpl(II->getArgOperand(0), TLI, false,
+                                            Depth + 1))
+          Known.knownNot(fcNan);
+
+        // Any negative value besides -0 returns a nan.
+        if (KnownSrc.isKnownNeverNaN() && KnownSrc.SignBit &&
+            !*KnownSrc.SignBit) {
+          Known.knownNot(fcNan);
+          Known.SignBit = false;
+        }
+
+        // The only negative value that can be returned is -0 for -0 inputs.
+        Known.knownNot(fcNegInf | fcNegSubnormal | fcNegNormal);
+        if ((KnownSrc.KnownFPClasses & fcNegZero) == fcNone) {
+          Known.knownNot(fcNegZero);
+          Known.SignBit = false;
+        }
+
+        break;
+      }
       case Intrinsic::sin:
       case Intrinsic::cos: {
         // Return NaN on infinite inputs.
         KnownFPClass KnownSrc;
         computeKnownFPClass(II->getArgOperand(0), DemandedElts,
                             InterestedClasses, KnownSrc, Depth + 1, Q, TLI);
         Known.knownNot(fcInf);
         if (KnownSrc.isKnownNeverInfinity())

llvm/test/Transforms/Attributor/nofpclass-sqrt.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -S < %s | FileCheck %s --check-prefixes=CHECK,TUNIT
 
 declare float @llvm.sqrt.f32(float)
 
 define float @ret_sqrt(float %arg0) {
-; CHECK-LABEL: define float @ret_sqrt
+; CHECK-LABEL: define nofpclass(ninf nsub nnorm) float @ret_sqrt
 ; CHECK-SAME: (float [[ARG0:%.*]]) #[[ATTR1:[0-9]+]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2:[0-9]+]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(ninf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2:[0-9]+]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_noinf(float nofpclass(inf) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_noinf
+; CHECK-LABEL: define nofpclass(inf nsub nnorm) float @ret_sqrt_noinf
 ; CHECK-SAME: (float nofpclass(inf) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_nopinf(float nofpclass(pinf) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_nopinf
+; CHECK-LABEL: define nofpclass(inf nsub nnorm) float @ret_sqrt_nopinf
 ; CHECK-SAME: (float nofpclass(pinf) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_noninf(float nofpclass(ninf) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_noninf
+; CHECK-LABEL: define nofpclass(ninf nsub nnorm) float @ret_sqrt_noninf
 ; CHECK-SAME: (float nofpclass(ninf) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(ninf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_nonan(float nofpclass(nan) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_nonan
+; CHECK-LABEL: define nofpclass(ninf nsub nnorm) float @ret_sqrt_nonan
 ; CHECK-SAME: (float nofpclass(nan) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(ninf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_nonan_noinf(float nofpclass(nan inf) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_nonan_noinf
+; CHECK-LABEL: define nofpclass(inf nsub nnorm) float @ret_sqrt_nonan_noinf
 ; CHECK-SAME: (float nofpclass(nan inf) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_nonan_noinf_nozero(float nofpclass(nan inf zero) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_nonan_noinf_nozero
+; CHECK-LABEL: define nofpclass(inf nzero nsub nnorm) float @ret_sqrt_nonan_noinf_nozero
 ; CHECK-SAME: (float nofpclass(nan inf zero) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nzero nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_noinf_nozero(float nofpclass(inf zero) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_noinf_nozero
+; CHECK-LABEL: define nofpclass(inf nzero nsub nnorm) float @ret_sqrt_noinf_nozero
 ; CHECK-SAME: (float nofpclass(inf zero) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nzero nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_noinf_nonegzero(float nofpclass(inf nzero) %arg0) {
-; CHECK-LABEL: define float @ret_sqrt_noinf_nonegzero
+; CHECK-LABEL: define nofpclass(inf nzero nsub nnorm) float @ret_sqrt_noinf_nonegzero
 ; CHECK-SAME: (float nofpclass(inf nzero) [[ARG0:%.*]]) #[[ATTR1]] {
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(inf nzero nsub nnorm) float @llvm.sqrt.f32(float [[ARG0]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %call = call float @llvm.sqrt.f32(float %arg0)
   ret float %call
 }
 
 define float @ret_sqrt_positive_source(i32 %arg) {
-; CHECK-LABEL: define float @ret_sqrt_positive_source
+; CHECK-LABEL: define nofpclass(nan inf nzero nsub nnorm) float @ret_sqrt_positive_source
 ; CHECK-SAME: (i32 [[ARG:%.*]]) #[[ATTR1]] {
 ; CHECK-NEXT:    [[UITOFP:%.*]] = uitofp i32 [[ARG]] to float
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[UITOFP]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(nan inf nzero nsub nnorm) float @llvm.sqrt.f32(float [[UITOFP]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %uitofp = uitofp i32 %arg to float
   %call = call float @llvm.sqrt.f32(float %uitofp)
   ret float %call
 }
 
 ; Could produce a nan because we don't know if the multiply is negative.
 define float @ret_sqrt_unknown_sign(float nofpclass(nan) %arg0, float nofpclass(nan) %arg1) {
-; CHECK-LABEL: define float @ret_sqrt_unknown_sign
+; CHECK-LABEL: define nofpclass(ninf nsub nnorm) float @ret_sqrt_unknown_sign
 ; CHECK-SAME: (float nofpclass(nan) [[ARG0:%.*]], float nofpclass(nan) [[ARG1:%.*]]) #[[ATTR1]] {
 ; CHECK-NEXT:    [[UNKNOWN_SIGN_NOT_NAN:%.*]] = fmul nnan float [[ARG0]], [[ARG1]]
-; CHECK-NEXT:    [[CALL:%.*]] = call float @llvm.sqrt.f32(float [[UNKNOWN_SIGN_NOT_NAN]]) #[[ATTR2]]
+; CHECK-NEXT:    [[CALL:%.*]] = call nofpclass(ninf nsub nnorm) float @llvm.sqrt.f32(float [[UNKNOWN_SIGN_NOT_NAN]]) #[[ATTR2]]
 ; CHECK-NEXT:    ret float [[CALL]]
 ;
   %unknown.sign.not.nan = fmul nnan float %arg0, %arg1
   %call = call float @llvm.sqrt.f32(float %unknown.sign.not.nan)
   ret float %call
 }
 
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; TUNIT: {{.*}}

llvm/unittests/Analysis/ValueTrackingTest.cpp

parseAssembly(
       "  %A = call float @llvm.copysign.f32(float %fabs, float %arg1)"
       "  ret float %A\n"
       "}\n");
   expectKnownFPClass(~fcNan, std::nullopt);
 }
 
 TEST_F(ComputeKnownFPClassTest, CopySignNInfSrc0_NegSign) {
   parseAssembly(
-      "declare float @llvm.sqrt.f32(float)\n"
+      "declare float @llvm.log.f32(float)\n"
       "declare float @llvm.copysign.f32(float, float)\n"
       "define float @test(float %arg0, float %arg1) {\n"
-      "  %ninf = call ninf float @llvm.sqrt.f32(float %arg0)"
+      "  %ninf = call ninf float @llvm.log.f32(float %arg0)"
       "  %A = call float @llvm.copysign.f32(float %ninf, float -1.0)"
       "  ret float %A\n"
       "}\n");
   expectKnownFPClass(fcNegFinite, true);
 }
 
 TEST_F(ComputeKnownFPClassTest, CopySignNInfSrc0_PosSign) {
   parseAssembly(