[flang] evaluate: Fold SQRT, HYPOT, & CABS

[klausler]

Summary

Implement IEEE Real::SQRT() operation, then use it to
also implement Real::HYPOT(), which can then be used directly
to implement Complex::ABS().

flang/include/flang/Evaluate/complex.h

ValueWithRealFlags<Complex> Add(
       const Complex &, Rounding rounding = defaultRounding) const;
   ValueWithRealFlags<Complex> Subtract(
       const Complex &, Rounding rounding = defaultRounding) const;
   ValueWithRealFlags<Complex> Multiply(
       const Complex &, Rounding rounding = defaultRounding) const;
   ValueWithRealFlags<Complex> Divide(
       const Complex &, Rounding rounding = defaultRounding) const;
 
+  // ABS/CABS = HYPOT(re_, imag_) = SQRT(re_**2 + im_**2)
+  ValueWithRealFlags<Part> ABS(Rounding rounding = defaultRounding) const {
+    return re_.HYPOT(im_, rounding);
+  }
+
   constexpr Complex FlushSubnormalToZero() const {
     return {re_.FlushSubnormalToZero(), im_.FlushSubnormalToZero()};
   }
 
   static constexpr Complex NotANumber() {
     return {Part::NotANumber(), Part::NotANumber()};
   }
 
   std::string DumpHexadecimal() const;
   llvm::raw_ostream &AsFortran(llvm::raw_ostream &, int kind) const;
 
-  // TODO: (C)ABS once Real::HYPOT is done
   // TODO: unit testing
 
 private:
   Part re_, im_;
 };
 
 extern template class Complex<Real<Integer<16>, 11>>;
 extern template class Complex<Real<Integer<16>, 8>>;
 extern template class Complex<Real<Integer<32>, 24>>;
 extern template class Complex<Real<Integer<64>, 53>>;
 extern template class Complex<Real<Integer<80>, 64>>;
 extern template class Complex<Real<Integer<128>, 113>>;
 } // namespace Fortran::evaluate::value
 #endif // FORTRAN_EVALUATE_COMPLEX_H_

flang/include/flang/Evaluate/real.h

ValueWithRealFlags<Real> Subtract(
       const Real &y, Rounding rounding = defaultRounding) const {
     return Add(y.Negate(), rounding);
   }
   ValueWithRealFlags<Real> Multiply(
       const Real &, Rounding rounding = defaultRounding) const;
   ValueWithRealFlags<Real> Divide(
       const Real &, Rounding rounding = defaultRounding) const;
 
-  // SQRT(x**2 + y**2) but computed so as to avoid spurious overflow
-  // TODO: not yet implemented; needed for CABS
+  ValueWithRealFlags<Real> SQRT(Rounding rounding = defaultRounding) const;
+
+  // HYPOT(x,y)=SQRT(x**2 + y**2) computed so as to avoid spurious
+  // intermediate overflows.
   ValueWithRealFlags<Real> HYPOT(
       const Real &, Rounding rounding = defaultRounding) const;
 
   template <typename INT> constexpr INT EXPONENT() const {
     if (Exponent() == maxExponent) {
       return INT::HUGE();
     } else {
       return {UnbiasedExponent()};

flang/lib/Evaluate/fold-real.cpp

Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
   CHECK(intrinsic);
   std::string name{intrinsic->name};
   if (name == "acos" || name == "acosh" || name == "asin" || name == "asinh" ||
       (name == "atan" && args.size() == 1) || name == "atanh" ||
       name == "bessel_j0" || name == "bessel_j1" || name == "bessel_y0" ||
       name == "bessel_y1" || name == "cos" || name == "cosh" || name == "erf" ||
       name == "erfc" || name == "erfc_scaled" || name == "exp" ||
       name == "gamma" || name == "log" || name == "log10" ||
-      name == "log_gamma" || name == "sin" || name == "sinh" ||
-      name == "sqrt" || name == "tan" || name == "tanh") {
+      name == "log_gamma" || name == "sin" || name == "sinh" || name == "tan" ||
+      name == "tanh") {
     CHECK(args.size() == 1);
     if (auto callable{GetHostRuntimeWrapper<T, T>(name)}) {
       return FoldElementalIntrinsic<T, T>(
           context, std::move(funcRef), *callable);
     } else {
       context.messages().Say(
           "%s(real(kind=%d)) cannot be folded on host"_en_US, name, KIND);
     }
   } else if (name == "amax0" || name == "amin0" || name == "amin1" ||
       name == "amax1" || name == "dmin1" || name == "dmax1") {
     return RewriteSpecificMINorMAX(context, std::move(funcRef));
-  } else if (name == "atan" || name == "atan2" || name == "hypot" ||
-      name == "mod") {
+  } else if (name == "atan" || name == "atan2" || name == "mod") {
     std::string localName{name == "atan" ? "atan2" : name};
     CHECK(args.size() == 2);
     if (auto callable{GetHostRuntimeWrapper<T, T, T>(localName)}) {
       return FoldElementalIntrinsic<T, T, T>(
           context, std::move(funcRef), *callable);
     } else {
       context.messages().Say(
           "%s(real(kind=%d), real(kind%d)) cannot be folded on host"_en_US,
if (args.size() == 2) { // elemental
       }
     }
   } else if (name == "abs") {
     // Argument can be complex or real
     if (auto *x{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
       return FoldElementalIntrinsic<T, T>(
           context, std::move(funcRef), &Scalar<T>::ABS);
     } else if (auto *z{UnwrapExpr<Expr<SomeComplex>>(args[0])}) {
-      if (auto callable{GetHostRuntimeWrapper<T, ComplexT>("abs")}) {
-        return FoldElementalIntrinsic<T, ComplexT>(
-            context, std::move(funcRef), *callable);
-      } else {
-        context.messages().Say(
-            "abs(complex(kind=%d)) cannot be folded on host"_en_US, KIND);
-      }
+      return FoldElementalIntrinsic<T, ComplexT>(context, std::move(funcRef),
+          ScalarFunc<T, ComplexT>([](const Scalar<ComplexT> &z) -> Scalar<T> {
+            return z.ABS().value;
+          }));
     } else {
       common::die(" unexpected argument type inside abs");
     }
   } else if (name == "aimag") {
     return FoldElementalIntrinsic<T, ComplexT>(
         context, std::move(funcRef), &Scalar<ComplexT>::AIMAG);
   } else if (name == "aint" || name == "anint") {
     // ANINT rounds ties away from zero, not to even
if (auto scalars{GetScalarConstantArguments<T, T>(context, args)}) {
       return Fold(context,
           Expr<T>{Multiply<T>{
               Expr<T>{std::get<0>(*scalars)}, Expr<T>{std::get<1>(*scalars)}}});
     }
   } else if (name == "epsilon") {
     return Expr<T>{Scalar<T>::EPSILON()};
   } else if (name == "huge") {
     return Expr<T>{Scalar<T>::HUGE()};
+  } else if (name == "hypot") {
+    CHECK(args.size() == 2);
+    return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),
+        ScalarFunc<T, T, T>(
+            [](const Scalar<T> &x, const Scalar<T> &y) -> Scalar<T> {
+              return x.HYPOT(y).value;
+            }));
   } else if (name == "max") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Greater);
   } else if (name == "maxval") {
     return FoldMaxvalMinval<T>(context, std::move(funcRef),
         RelationalOperator::GT, T::Scalar::HUGE().Negate());
   } else if (name == "merge") {
     return FoldMerge<T>(context, std::move(funcRef));
   } else if (name == "min") {
     return FoldMINorMAX(context, std::move(funcRef), Ordering::Less);
   } else if (name == "minval") {
     return FoldMaxvalMinval<T>(
         context, std::move(funcRef), RelationalOperator::LT, T::Scalar::HUGE());
   } else if (name == "product") {
     auto one{Scalar<T>::FromInteger(value::Integer<8>{1}).value};
     return FoldProduct<T>(context, std::move(funcRef), one);
   } else if (name == "real") {
     if (auto *expr{args[0].value().UnwrapExpr()}) {
       return ToReal<KIND>(context, std::move(*expr));
     }
   } else if (name == "sign") {
     return FoldElementalIntrinsic<T, T, T>(
         context, std::move(funcRef), &Scalar<T>::SIGN);
+  } else if (name == "sqrt") {
+    return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
+        ScalarFunc<T, T>(
+            [](const Scalar<T> &x) -> Scalar<T> { return x.SQRT().value; }));
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "tiny") {
     return Expr<T>{Scalar<T>::TINY()};
   }
   // TODO: dim, dot_product, fraction, matmul,
   // maxloc, minloc, modulo, nearest, norm2, rrspacing, scale,
   // set_exponent, spacing, transfer,

flang/lib/Evaluate/intrinsics-library.cpp

static constexpr HostRuntimeFunction table[]{
       FolderFactory<F2, F2{std::atan2}>::Create("atan2"),
       FolderFactory<F, F{std::atanh}>::Create("atanh"),
       FolderFactory<F, F{std::cos}>::Create("cos"),
       FolderFactory<F, F{std::cosh}>::Create("cosh"),
       FolderFactory<F, F{std::erf}>::Create("erf"),
       FolderFactory<F, F{std::erfc}>::Create("erfc"),
       FolderFactory<F, F{std::exp}>::Create("exp"),
       FolderFactory<F, F{std::tgamma}>::Create("gamma"),
-      FolderFactory<F2, F2{std::hypot}>::Create("hypot"),
       FolderFactory<F, F{std::log}>::Create("log"),
       FolderFactory<F, F{std::log10}>::Create("log10"),
       FolderFactory<F, F{std::lgamma}>::Create("log_gamma"),
       FolderFactory<F2, F2{std::fmod}>::Create("mod"),
       FolderFactory<F2, F2{std::pow}>::Create("pow"),
       FolderFactory<F, F{std::sin}>::Create("sin"),
       FolderFactory<F, F{std::sinh}>::Create("sinh"),
-      FolderFactory<F, F{std::sqrt}>::Create("sqrt"),
       FolderFactory<F, F{std::tan}>::Create("tan"),
       FolderFactory<F, F{std::tanh}>::Create("tanh"),
   };
   // Note: cmath does not have modulo and erfc_scaled equivalent
 
   // Note regarding  lack of bessel function support:
   // C++17 defined standard Bessel math functions std::cyl_bessel_j
   // and std::cyl_neumann that can be used for Fortran j and y

flang/lib/Evaluate/real.cpp

if (IsInfinite()) {
       NormalizeAndRound(
           result, isNegative, exponent, quotient, rounding, roundingBits);
     }
   }
   return result;
 }
 
 template <typename W, int P>
+ValueWithRealFlags<Real<W, P>> Real<W, P>::SQRT(Rounding rounding) const {
+  ValueWithRealFlags<Real> result;
+  if (IsNotANumber()) {
+    result.value = NotANumber();
+    if (IsSignalingNaN()) {
+      result.flags.set(RealFlag::InvalidArgument);
+    }
+  } else if (IsNegative()) {
+    if (IsZero()) {
+      // SQRT(-0) == -0 in IEEE-754.
+      result.value.word_ = result.value.word_.IBSET(bits - 1);
+    } else {
+      result.value = NotANumber();
+    }
+  } else if (IsInfinite()) {
+    // SQRT(+Inf) == +Inf
+    result.value = Infinity(false);
+  } else {
+    // Slow but reliable bit-at-a-time method.  Start with a clear significand
+    // and half the unbiased exponent, and then try to set significand bits
+    // in descending order of magnitude without exceeding the exact result.
+    int expo{UnbiasedExponent()};
+    if (IsSubnormal()) {
+      expo -= GetFraction().LEADZ();
+    }
+    expo = expo / 2 + exponentBias;
+    result.value.Normalize(false, expo, Fraction::MASKL(1));
+    for (int bit{significandBits - 1}; bit >= 0; --bit) {
+      Word word{result.value.word_};
+      result.value.word_ = word.IBSET(bit);
+      auto squared{result.value.Multiply(result.value, rounding)};
+      if (squared.flags.test(RealFlag::Overflow) ||
+          squared.flags.test(RealFlag::Underflow) ||
+          Compare(squared.value) == Relation::Less) {
+        result.value.word_ = word;
+      }
+    }
+    // The computed square root, when squared, has a square that's not greater
+    // than the original argument.  Check this square against the square of the
+    // next Real value, and return that one if its square is closer in magnitude
+    // to the original argument.
+    Real resultSq{result.value.Multiply(result.value).value};
+    Real diff{Subtract(resultSq).value.ABS()};
+    if (diff.IsZero()) {
+      return result; // exact
+    }
+    Real ulp;
+    ulp.Normalize(false, expo, Fraction::MASKR(1));
+    Real nextAfter{result.value.Add(ulp).value};
+    auto nextAfterSq{nextAfter.Multiply(nextAfter)};
+    if (!nextAfterSq.flags.test(RealFlag::Overflow) &&
+        !nextAfterSq.flags.test(RealFlag::Underflow)) {
+      Real nextAfterDiff{Subtract(nextAfterSq.value).value.ABS()};
+      if (nextAfterDiff.Compare(diff) == Relation::Less) {
+        result.value = nextAfter;
+        if (nextAfterDiff.IsZero()) {
+          return result; // exact
+        }
+      }
+    }
+    result.flags.set(RealFlag::Inexact);
+  }
+  return result;
+}
+
+// HYPOT(x,y) = SQRT(x**2 + y**2) by definition, but those squared intermediate
+// values are susceptible to over/underflow when computed naively.
+// Assuming that x>=y, calculate instead:
+//   HYPOT(x,y) = SQRT(x**2 * (1+(y/x)**2))
+//              = ABS(x) * SQRT(1+(y/x)**2)
+template <typename W, int P>
+ValueWithRealFlags<Real<W, P>> Real<W, P>::HYPOT(
+    const Real &y, Rounding rounding) const {
+  ValueWithRealFlags<Real> result;
+  if (IsNotANumber() || y.IsNotANumber()) {
+    result.flags.set(RealFlag::InvalidArgument);
+    result.value = NotANumber();
+  } else if (ABS().Compare(y.ABS()) == Relation::Less) {
+    return y.HYPOT(*this);
+  } else if (IsZero()) {
+    return result; // x==y==0
+  } else {
+    auto yOverX{y.Divide(*this, rounding)}; // y/x
+    bool inexact{yOverX.flags.test(RealFlag::Inexact)};
+    auto squared{yOverX.value.Multiply(yOverX.value, rounding)}; // (y/x)**2
+    inexact |= squared.flags.test(RealFlag::Inexact);
+    Real one;
+    one.Normalize(false, exponentBias, Fraction::MASKL(1)); // 1.0
+    auto sum{squared.value.Add(one, rounding)}; // 1.0 + (y/x)**2
+    inexact |= sum.flags.test(RealFlag::Inexact);
+    auto sqrt{sum.value.SQRT()};
+    inexact |= sqrt.flags.test(RealFlag::Inexact);
+    result = sqrt.value.Multiply(ABS(), rounding);
+    if (inexact) {
+      result.flags.set(RealFlag::Inexact);
+    }
+  }
+  return result;
+}
+
+template <typename W, int P>
 ValueWithRealFlags<Real<W, P>> Real<W, P>::ToWholeNumber(
     common::RoundingMode mode) const {
   ValueWithRealFlags<Real> result{*this};
   if (IsNotANumber()) {
     result.flags.set(RealFlag::InvalidArgument);
     result.value = NotANumber();
   } else if (IsInfinite()) {
     result.flags.set(RealFlag::Overflow);

flang/test/Evaluate/folding28.f90

+! RUN: %S/test_folding.sh %s %t %flang_fc1
+! REQUIRES: shell
+! Tests folding of SQRT()
+module m
+  implicit none
+  ! +Inf
+  real(8), parameter :: inf8 = z'7ff0000000000000'
+  logical, parameter :: test_inf8 = sqrt(inf8) == inf8
+  ! max finite
+  real(8), parameter :: h8 = huge(1.0_8), h8z = z'7fefffffffffffff'
+  logical, parameter :: test_h8 = h8 == h8z
+  real(8), parameter :: sqrt_h8 = sqrt(h8), sqrt_h8z = z'5fefffffffffffff'
+  logical, parameter :: test_sqrt_h8 = sqrt_h8 == sqrt_h8z
+  real(8), parameter :: sqr_sqrt_h8 = sqrt_h8 * sqrt_h8, sqr_sqrt_h8z = z'7feffffffffffffe'
+  logical, parameter :: test_sqr_sqrt_h8 = sqr_sqrt_h8 == sqr_sqrt_h8z
+  ! -0 (sqrt is -0)
+  real(8), parameter :: n08 = z'8000000000000000'
+  real(8), parameter :: sqrt_n08 = sqrt(n08)
+!WARN: division by zero
+  real(8), parameter :: inf_n08 = 1.0_8 / sqrt_n08, inf_n08z = z'fff0000000000000'
+  logical, parameter :: test_n08 = inf_n08 == inf_n08z
+  ! min normal
+  real(8), parameter :: t8 = tiny(1.0_8), t8z = z'0010000000000000'
+  logical, parameter :: test_t8 = t8 == t8z
+  real(8), parameter :: sqrt_t8 = sqrt(t8), sqrt_t8z = z'2000000000000000'
+  logical, parameter :: test_sqrt_t8 = sqrt_t8 == sqrt_t8z
+  real(8), parameter :: sqr_sqrt_t8 = sqrt_t8 * sqrt_t8
+  logical, parameter :: test_sqr_sqrt_t8 = sqr_sqrt_t8 == t8
+  ! max subnormal
+  real(8), parameter :: maxs8 = z'000fffffffffffff'
+  real(8), parameter :: sqrt_maxs8 = sqrt(maxs8), sqrt_maxs8z = z'2000000000000000'
+  logical, parameter :: test_sqrt_maxs8 = sqrt_maxs8 == sqrt_maxs8z
+  ! min subnormal
+  real(8), parameter :: mins8 = z'1'
+  real(8), parameter :: sqrt_mins8 = sqrt(mins8), sqrt_mins8z = z'1e60000000000000'
+  logical, parameter :: test_sqrt_mins8 = sqrt_mins8 == sqrt_mins8z
+  real(8), parameter :: sqr_sqrt_mins8 = sqrt_mins8 * sqrt_mins8
+  logical, parameter :: test_sqr_sqrt_mins8 = sqr_sqrt_mins8 == mins8
+end module
+