[flang] Fold NEAREST() and its relatives

[klausler]

Summary

Implement constant folding for the intrinsic function NEAREST()
and the related functions IEEE_NEXT_AFTER(), IEEE_NEXT_UP(), and
IEEE_NEXT_DOWN().

[jeanPerier]

Looks good

flang/include/flang/Evaluate/real.h

public:
   }
   ValueWithRealFlags<Real> Multiply(
       const Real &, Rounding rounding = defaultRounding) const;
   ValueWithRealFlags<Real> Divide(
       const Real &, Rounding rounding = defaultRounding) const;
 
   ValueWithRealFlags<Real> SQRT(Rounding rounding = defaultRounding) const;
 
+  // NEAREST(), IEEE_NEXT_AFTER(), IEEE_NEXT_UP(), and IEEE_NEXT_DOWN()
+  ValueWithRealFlags<Real> NEAREST(bool upward) 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();

flang/lib/Evaluate/fold-real.cpp

return FoldElementalIntrinsic<T, T, T>(context, std::move(funcRef),
             }));
   } 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 == "nearest") {
+    if (const auto *sExpr{UnwrapExpr<Expr<SomeReal>>(args[1])}) {
+      return std::visit(
+          [&](const auto &sVal) {
+            using TS = ResultType<decltype(sVal)>;
+            return FoldElementalIntrinsic<T, T, TS>(context, std::move(funcRef),
+                ScalarFunc<T, T, TS>([&](const Scalar<T> &x,
+                                         const Scalar<TS> &s) -> Scalar<T> {
+                  if (s.IsZero()) {
+                    context.messages().Say(
+                        "NEAREST: S argument is zero"_warn_en_US);
+                  }
+                  auto result{x.NEAREST(!s.IsNegative())};
+                  if (result.flags.test(RealFlag::Overflow)) {
+                    context.messages().Say(
+                        "NEAREST intrinsic folding overflow"_warn_en_US);
+                  } else if (result.flags.test(RealFlag::InvalidArgument)) {
+                    context.messages().Say(
+                        "NEAREST intrinsic folding: bad argument"_warn_en_US);
+                  }
+                  return result.value;
+                }));
+          },
+          sExpr->u);
+    }
   } 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);
#endif
   } 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()};
+  } else if (name == "__builtin_ieee_next_after") {
+    if (const auto *yExpr{UnwrapExpr<Expr<SomeReal>>(args[1])}) {
+      return std::visit(
+          [&](const auto &yVal) {
+            using TY = ResultType<decltype(yVal)>;
+            return FoldElementalIntrinsic<T, T, TY>(context, std::move(funcRef),
+                ScalarFunc<T, T, TY>([&](const Scalar<T> &x,
+                                         const Scalar<TY> &y) -> Scalar<T> {
+                  bool upward{true};
+                  switch (x.Compare(Scalar<T>::Convert(y).value)) {
+                  case Relation::Unordered:
+                    context.messages().Say(
+                        "IEEE_NEXT_AFTER intrinsic folding: bad argument"_warn_en_US);
+                    return x;
+                  case Relation::Equal:
+                    return x;
+                  case Relation::Less:
+                    upward = true;
+                    break;
+                  case Relation::Greater:
+                    upward = false;
+                    break;
+                  }
+                  auto result{x.NEAREST(upward)};
+                  if (result.flags.test(RealFlag::Overflow)) {
+                    context.messages().Say(
+                        "IEEE_NEXT_AFTER intrinsic folding overflow"_warn_en_US);
+                  }
+                  return result.value;
+                }));
+          },
+          yExpr->u);
+    }
+  } else if (name == "__builtin_ieee_next_up" ||
+      name == "__builtin_ieee_next_down") {
+    bool upward{name == "__builtin_ieee_next_up"};
+    const char *iName{upward ? "IEEE_NEXT_UP" : "IEEE_NEXT_DOWN"};
+    return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
+        ScalarFunc<T, T>([&](const Scalar<T> &x) -> Scalar<T> {
+          auto result{x.NEAREST(upward)};
+          if (result.flags.test(RealFlag::Overflow)) {
+            context.messages().Say(
+                "%s intrinsic folding overflow"_warn_en_US, iName);
+          } else if (result.flags.test(RealFlag::InvalidArgument)) {
+            context.messages().Say(
+                "%s intrinsic folding: bad argument"_warn_en_US, iName);
+          }
+          return result.value;
+        }));
   }
   // TODO: dim, dot_product, fraction, matmul,
-  // modulo, nearest, norm2, rrspacing,
-  // __builtin_next_after/down/up,
+  // modulo, norm2, rrspacing,
   // set_exponent, spacing, transfer,
   // bessel_jn (transformational) and bessel_yn (transformational)
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Real, KIND>> FoldOperation(
     FoldingContext &context, ComplexComponent<KIND> &&x) {

flang/lib/Evaluate/real.cpp

if (!nextAfterSq.flags.test(RealFlag::Overflow) &&
         }
       }
     }
     result.flags.set(RealFlag::Inexact);
   }
   return result;
 }
 
+template <typename W, int P>
+ValueWithRealFlags<Real<W, P>> Real<W, P>::NEAREST(bool upward) const {
+  ValueWithRealFlags<Real> result;
+  if (IsFinite()) {
+    Fraction fraction{GetFraction()};
+    int expo{Exponent()};
+    Fraction one{1};
+    Fraction nearest;
+    bool isNegative{IsNegative()};
+    if (upward != isNegative) { // upward in magnitude
+      auto next{fraction.AddUnsigned(one)};
+      if (next.carry) {
+        ++expo;
+        nearest = Fraction::Least(); // MSB only
+      } else {
+        nearest = next.value;
+      }
+    } else { // downward in magnitude
+      if (IsZero()) {
+        nearest = 1; // smallest magnitude negative subnormal
+        isNegative = !isNegative;
+      } else {
+        auto sub1{fraction.SubtractSigned(one)};
+        if (sub1.overflow) {
+          nearest = Fraction{0}.NOT();
+          --expo;
+        } else {
+          nearest = sub1.value;
+        }
+      }
+    }
+    result.flags = result.value.Normalize(isNegative, expo, nearest);
+  } else {
+    result.flags.set(RealFlag::InvalidArgument);
+    result.value = *this;
+  }
+  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 {

flang/test/Evaluate/fold-nearest.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of NEAREST() and its relatives
+module m1
+  real, parameter :: minSubnormal = 1.e-45
+  logical, parameter :: test_1 = nearest(0., 1.) == minSubnormal
+  logical, parameter :: test_2 = nearest(minSubnormal, -1.) == 0
+  logical, parameter :: test_3 = nearest(1., 1.) == 1.0000001
+  logical, parameter :: test_4 = nearest(1.0000001, -1.) == 1
+  !WARN: warning: NEAREST intrinsic folding overflow
+  real, parameter :: inf = nearest(huge(1.), 1.)
+  !WARN: warning: NEAREST intrinsic folding: bad argument
+  logical, parameter :: test_5 = nearest(inf, 1.) == inf
+  !WARN: warning: NEAREST intrinsic folding: bad argument
+  logical, parameter :: test_6 = nearest(-inf, -1.) == -inf
+  logical, parameter :: test_7 = nearest(1.9999999, 1.) == 2.
+  logical, parameter :: test_8 = nearest(2., -1.) == 1.9999999
+  logical, parameter :: test_9 = nearest(1.9999999999999999999_10, 1.) == 2._10
+  logical, parameter :: test_10 = nearest(-1., 1.) == -.99999994
+  logical, parameter :: test_11 = nearest(-1., -2.) == -1.0000001
+  real, parameter :: negZero = sign(0., -1.)
+  logical, parameter :: test_12 = nearest(negZero, 1.) == minSubnormal
+  logical, parameter :: test_13 = nearest(negZero, -1.) == -minSubnormal
+  !WARN: warning: NEAREST: S argument is zero
+  logical, parameter :: test_14 = nearest(0., negZero) == -minSubnormal
+  !WARN: warning: NEAREST: S argument is zero
+  logical, parameter :: test_15 = nearest(negZero, 0.) == minSubnormal
+end module
+
+module m2
+  use ieee_arithmetic, only: ieee_next_after
+  real, parameter :: minSubnormal = 1.e-45
+  logical, parameter :: test_0 = ieee_next_after(0., 0.) == 0.
+  logical, parameter :: test_1 = ieee_next_after(0., 1.) == minSubnormal
+  logical, parameter :: test_2 = ieee_next_after(minSubnormal, -1.) == 0
+  logical, parameter :: test_3 = ieee_next_after(1., 2.) == 1.0000001
+  logical, parameter :: test_4 = ieee_next_after(1.0000001, -1.) == 1
+  !WARN: warning: division by zero
+  real, parameter :: inf = 1. / 0.
+  logical, parameter :: test_5 = ieee_next_after(inf, inf) == inf
+  logical, parameter :: test_6 = ieee_next_after(inf, -inf) == inf
+  logical, parameter :: test_7 = ieee_next_after(-inf, inf) == -inf
+  logical, parameter :: test_8 = ieee_next_after(-inf, -1.) == -inf
+  logical, parameter :: test_9 = ieee_next_after(1.9999999, 3.) == 2.
+  logical, parameter :: test_10 = ieee_next_after(2., 1.) == 1.9999999
+  logical, parameter :: test_11 = ieee_next_after(1.9999999999999999999_10, 3.) == 2._10
+  logical, parameter :: test_12 = ieee_next_after(1., 1.) == 1.
+  !WARN: warning: invalid argument on division
+  real, parameter :: nan = 0. / 0.
+  !WARN: warning: IEEE_NEXT_AFTER intrinsic folding: bad argument
+  real, parameter :: x13 = ieee_next_after(nan, nan)
+  logical, parameter :: test_13 = .not. (x13 == x13)
+  !WARN: warning: IEEE_NEXT_AFTER intrinsic folding: bad argument
+  real, parameter :: x14 = ieee_next_after(nan, 0.)
+  logical, parameter :: test_14 = .not. (x14 == x14)
+end module
+
+module m3
+  use ieee_arithmetic, only: ieee_next_up, ieee_next_down
+  real(kind(0.d0)), parameter :: minSubnormal = 5.d-324
+  logical, parameter :: test_1 = ieee_next_up(0.d0) == minSubnormal
+  logical, parameter :: test_2 = ieee_next_down(0.d0) == -minSubnormal
+  logical, parameter :: test_3 = ieee_next_up(1.d0) == 1.0000000000000002d0
+  logical, parameter :: test_4 = ieee_next_down(1.0000000000000002d0) == 1.d0
+  !WARN: warning: division by zero
+  real(kind(0.d0)), parameter :: inf = 1.d0 / 0.d0
+  !WARN: warning: IEEE_NEXT_UP intrinsic folding overflow
+  logical, parameter :: test_5 = ieee_next_up(huge(0.d0)) == inf
+  !WARN: warning: IEEE_NEXT_DOWN intrinsic folding overflow
+  logical, parameter :: test_6 = ieee_next_down(-huge(0.d0)) == -inf
+  !WARN: warning: IEEE_NEXT_UP intrinsic folding: bad argument
+  logical, parameter :: test_7 = ieee_next_up(inf) == inf
+  !WARN: warning: IEEE_NEXT_DOWN intrinsic folding: bad argument
+  logical, parameter :: test_8 = ieee_next_down(inf) == inf
+  !WARN: warning: IEEE_NEXT_UP intrinsic folding: bad argument
+  logical, parameter :: test_9 = ieee_next_up(-inf) == -inf
+  !WARN: warning: IEEE_NEXT_DOWN intrinsic folding: bad argument
+  logical, parameter :: test_10 = ieee_next_down(-inf) == -inf
+  logical, parameter :: test_11 = ieee_next_up(1.9999999999999997d0) == 2.d0
+  logical, parameter :: test_12 = ieee_next_down(2.d0) == 1.9999999999999997d0
+  !WARN: warning: invalid argument on division
+  real(kind(0.d0)), parameter :: nan = 0.d0 / 0.d0
+  !WARN: warning: IEEE_NEXT_UP intrinsic folding: bad argument
+  real(kind(0.d0)), parameter :: x13 = ieee_next_up(nan)
+  logical, parameter :: test_13 = .not. (x13 == x13)
+  !WARN: warning: IEEE_NEXT_DOWN intrinsic folding: bad argument
+  real(kind(0.d0)), parameter :: x14 = ieee_next_down(nan)
+  logical, parameter :: test_14 = .not. (x14 == x14)
+end module