Fix _approx_ftz_d implementations for rcp + rsqrt

Fix the rcp_approx_ftz_d and rsqrt_approx_ftz_d implementations
to better match the PTX spec, which states that the inputs and outputs
should zero the lower 32-bits of the mantissa.

Author: LewisCrawford

llvm/lib/Analysis/ConstantFolding.cpp

@@ -2006,6 +2006,15 @@ static const APFloat FTZPreserveSign(const APFloat &V) {
   return V;
 }
 
+// Get only the upper word of the input double in 1.11.20 format
+// by making the lower 32-bits of the mantissa all 0.
+static const APFloat ZeroLower32Bits(const APFloat &V) {
+  assert(V.getSizeInBits(V.getSemantics()) == 64);
+  uint64_t DoubleBits = V.bitcastToAPInt().getZExtValue();
+  DoubleBits &= 0xffffffff00000000;
+  return APFloat(V.getSemantics(), APInt(64, DoubleBits, false, false));
+}
+
 Constant *ConstantFoldFP(double (*NativeFP)(double), const APFloat &V, Type *Ty,
                          bool ShouldFTZPreservingSign = false) {
   llvm_fenv_clearexcept();
@@ -2651,6 +2660,8 @@ static Constant *ConstantFoldScalarCall1(StringRef Name,
         bool IsFTZ = nvvm::RCPShouldFTZ(IntrinsicID);
 
         auto Denominator = IsFTZ ? FTZPreserveSign(APF) : APF;
+        if (IntrinsicID == Intrinsic::nvvm_rcp_approx_ftz_d)
+          Denominator = ZeroLower32Bits(Denominator);
         if (IsApprox && Denominator.isZero()) {
           // According to the PTX spec, approximate rcp should return infinity
           // with the same sign as the denominator when dividing by 0.
@@ -2663,6 +2674,8 @@ static Constant *ConstantFoldScalarCall1(StringRef Name,
         if (Status == APFloat::opOK || Status == APFloat::opInexact) {
           if (IsFTZ)
             Res = FTZPreserveSign(Res);
+          if (IntrinsicID == Intrinsic::nvvm_rcp_approx_ftz_d)
+            Res = ZeroLower32Bits(Res);
           return ConstantFP::get(Ty->getContext(), Res);
         }
         return nullptr;
@@ -2680,14 +2693,24 @@ static Constant *ConstantFoldScalarCall1(StringRef Name,
       case Intrinsic::nvvm_rsqrt_approx_f: {
         bool IsFTZ = nvvm::UnaryMathIntrinsicShouldFTZ(IntrinsicID);
         auto V = IsFTZ ? FTZPreserveSign(APF) : APF;
+
+        if (IntrinsicID == Intrinsic::nvvm_rsqrt_approx_ftz_d)
+          V = ZeroLower32Bits(V);
+
         APFloat SqrtV(sqrt(V.convertToDouble()));
 
-        bool lost;
-        SqrtV.convert(APF.getSemantics(), APFloat::rmNearestTiesToEven, &lost);
+        if (Ty->isFloatTy()) {
+          bool lost;
+          SqrtV.convert(APF.getSemantics(), APFloat::rmNearestTiesToEven,
+                        &lost);
+        }
 
         APFloat Res = APFloat::getOne(APF.getSemantics());
         Res.divide(SqrtV, APFloat::rmNearestTiesToEven);
 
+        if (IntrinsicID == Intrinsic::nvvm_rsqrt_approx_ftz_d)
+          Res = ZeroLower32Bits(Res);
+
         // We do not need to flush the output for ftz because it is impossible
         // for 1/sqrt(x) to be a denormal value. If x is the largest fp value,
         // sqrt(x) will be a number with the exponent approximately halved and

llvm/test/Transforms/InstSimplify/const-fold-nvvm-unary-arithmetic.ll

@@ -551,7 +551,7 @@ define float @test_rcp_approx_ftz_f_0_5() {
 
 define double @test_rcp_approx_ftz_d_neg_subnorm() {
 ; CHECK-LABEL: define double @test_rcp_approx_ftz_d_neg_subnorm() {
-; CHECK-NEXT:    ret double 0xC7D0000020000040
+; CHECK-NEXT:    ret double 0xC7D0000000000000
 ;
   %res = call double @llvm.nvvm.rcp.approx.ftz.d(double 0xB80FFFFFC0000000)
   ret double %res
@@ -568,7 +568,7 @@ define float @test_rcp_approx_ftz_f_neg_subnorm() {
 
 define double @test_rcp_approx_ftz_d_pos_subnorm() {
 ; CHECK-LABEL: define double @test_rcp_approx_ftz_d_pos_subnorm() {
-; CHECK-NEXT:    ret double 0x47D0000020000040
+; CHECK-NEXT:    ret double 0x47D0000000000000
 ;
   %res = call double @llvm.nvvm.rcp.approx.ftz.d(double 0x380FFFFFC0000000)
   ret double %res
@@ -658,7 +658,7 @@ define float @test_rsqrt_approx_f_1_25() {
 
 define double @test_rsqrt_approx_ftz_d_1_25() {
 ; CHECK-LABEL: define double @test_rsqrt_approx_ftz_d_1_25() {
-; CHECK-NEXT:    ret double 0x3FEC9F25C5BFEDD9
+; CHECK-NEXT:    ret double 0x3FEC9F2500000000
 ;
   %res = call double @llvm.nvvm.rsqrt.approx.ftz.d(double 1.25)
   ret double %res
@@ -690,7 +690,7 @@ define float @test_rsqrt_approx_f_pos_subnorm() {
 
 define double @test_rsqrt_approx_ftz_d_pos_subnorm() {
 ; CHECK-LABEL: define double @test_rsqrt_approx_ftz_d_pos_subnorm() {
-; CHECK-NEXT:    ret double 0x43E0000010000018
+; CHECK-NEXT:    ret double 0x43E0000000000000
 ;
   %res = call double @llvm.nvvm.rsqrt.approx.ftz.d(double 0x380FFFFFC0000000)
   ret double %res