[flang] Fold SET_EXPONENT() and FRACTION()

Fold the intrinsic function SET_EXPONENT() and its special case,
FRACTION().

Differential Revision: https://reviews.llvm.org/D130378

4 files changed, 57 insertions, 2 deletions

diff --git a/flang/include/flang/Evaluate/real.h b/flang/include/flang/Evaluate/real.h
index 37c4010..916b8be 100644
--- a/flang/include/flang/Evaluate/real.h
+++ b/flang/include/flang/Evaluate/real.h
@@ -170,6 +170,8 @@ public:
   static constexpr int MINEXPONENT{2 - exponentBias};
   Real RRSPACING() const;
   Real SPACING() const;
+  Real SET_EXPONENT(int) const;
+  Real FRACTION() const;
 
   // SCALE(); also known as IEEE_SCALB and (in IEEE-754 '08) ScaleB.
   template <typename INT>
diff --git a/flang/lib/Evaluate/fold-real.cpp b/flang/lib/Evaluate/fold-real.cpp
index 159b2ed..6bdc922 100644
--- a/flang/lib/Evaluate/fold-real.cpp
+++ b/flang/lib/Evaluate/fold-real.cpp
@@ -49,6 +49,7 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
     FunctionRef<Type<TypeCategory::Real, KIND>> &&funcRef) {
   using T = Type<TypeCategory::Real, KIND>;
   using ComplexT = Type<TypeCategory::Complex, KIND>;
+  using Int4 = Type<TypeCategory::Integer, 4>;
   ActualArguments &args{funcRef.arguments()};
   auto *intrinsic{std::get_if<SpecificIntrinsic>(&funcRef.proc().u)};
   CHECK(intrinsic);
@@ -86,7 +87,6 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
   } else if (name == "bessel_jn" || name == "bessel_yn") {
     if (args.size() == 2) { // elemental
       // runtime functions use int arg
-      using Int4 = Type<TypeCategory::Integer, 4>;
       if (auto callable{GetHostRuntimeWrapper<T, Int4, T>(name)}) {
         return FoldElementalIntrinsic<T, Int4, T>(
             context, std::move(funcRef), *callable);
@@ -144,6 +144,10 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
     }
   } else if (name == "epsilon") {
     return Expr<T>{Scalar<T>::EPSILON()};
+  } else if (name == "fraction") {
+    return FoldElementalIntrinsic<T, T>(context, std::move(funcRef),
+        ScalarFunc<T, T>(
+            [](const Scalar<T> &x) -> Scalar<T> { return x.FRACTION(); }));
   } else if (name == "huge") {
     return Expr<T>{Scalar<T>::HUGE()};
   } else if (name == "hypot") {
@@ -250,6 +254,12 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
           },
           byExpr->u);
     }
+  } else if (name == "set_exponent") {
+    return FoldElementalIntrinsic<T, T, Int4>(context, std::move(funcRef),
+        ScalarFunc<T, T, Int4>(
+            [&](const Scalar<T> &x, const Scalar<Int4> &i) -> Scalar<T> {
+              return x.SET_EXPONENT(i.ToInt64());
+            }));
   } else if (name == "sign") {
     return FoldElementalIntrinsic<T, T, T>(
         context, std::move(funcRef), &Scalar<T>::SIGN);
@@ -315,7 +325,7 @@ Expr<Type<TypeCategory::Real, KIND>> FoldIntrinsicFunction(
           return result.value;
         }));
   }
-  // TODO: dot_product, fraction, matmul, norm2, set_exponent
+  // TODO: dot_product, matmul, norm2
   return Expr<T>{std::move(funcRef)};
 }
 
diff --git a/flang/lib/Evaluate/real.cpp b/flang/lib/Evaluate/real.cpp
index c22d3c8..5ffbf2f 100644
--- a/flang/lib/Evaluate/real.cpp
+++ b/flang/lib/Evaluate/real.cpp
@@ -755,6 +755,27 @@ template <typename W, int P> Real<W, P> Real<W, P>::SPACING() const {
   }
 }
 
+// 16.9.171
+template <typename W, int P>
+Real<W, P> Real<W, P>::SET_EXPONENT(int expo) const {
+  if (IsNotANumber()) {
+    return *this;
+  } else if (IsInfinite()) {
+    return NotANumber();
+  } else if (IsZero()) {
+    return *this;
+  } else {
+    Real result;
+    result.Normalize(IsNegative(), expo + exponentBias - 1, GetFraction());
+    return result;
+  }
+}
+
+// 16.9.171
+template <typename W, int P> Real<W, P> Real<W, P>::FRACTION() const {
+  return SET_EXPONENT(0);
+}
+
 template class Real<Integer<16>, 11>;
 template class Real<Integer<16>, 8>;
 template class Real<Integer<32>, 24>;
diff --git a/flang/test/Evaluate/folding07.f90 b/flang/test/Evaluate/folding07.f90
index b8eca95..d5ce9ac 100644
--- a/flang/test/Evaluate/folding07.f90
+++ b/flang/test/Evaluate/folding07.f90
@@ -197,13 +197,25 @@ module m
   logical, parameter :: test_tiny10 = tiny10 == ztiny10
   logical, parameter :: test_tiny16 = tiny16 == ztiny16
 
+  real,    parameter :: nan = real(z'7fc12345')
+  integer, parameter :: nanInt = int(z'7fc12345')
+  real,    parameter :: inf = real(z'7f800000')
   logical, parameter :: test_exponent_0 = exponent(0.0) == 0
+  logical, parameter :: test_fraction_0 = fraction(0.) == 0.
   logical, parameter :: test_exponent_r8 = exponent(0.125) == -2
+  logical, parameter :: test_fraction_r8 = fraction(0.125) == 0.5
   logical, parameter :: test_exponent_r4 = exponent(0.25) == -1
+  logical, parameter :: test_fraction_mr4 = fraction(-0.25) == -0.5
   logical, parameter :: test_exponent_r2 = exponent(0.5) == 0
+  logical, parameter :: test_fraction_r2 = fraction(0.5) == 0.5
   logical, parameter :: test_exponent_1 = exponent(1.0) == 1
+  logical, parameter :: test_fraction_1 = fraction(1.) == 0.5
   logical, parameter :: test_exponent_4 = exponent(4.1) == 3
+  logical, parameter :: test_fraction_m4 = fraction(-4.5) == -0.5625
   logical, parameter :: test_exponent_12 = exponent(12.9) == 4
+  real,    parameter :: fraction_inf = fraction(inf)
+  logical, parameter :: test_fraction_inf = fraction_inf /= fraction_inf ! must be NaN
+  logical, parameter :: test_fraction_nan = transfer(fraction(nan),0) == nanInt
 
   integer, parameter :: &
     max2 = maxexponent(0._2), &
@@ -259,4 +271,14 @@ module m
   logical, parameter :: test_zrange10 = arange10 == 4931 .and. zrange10 == 4931
   logical, parameter :: test_zrange16 = arange16 == 4931 .and. zrange16 == 4931
 
+  logical, parameter :: test_set_exponent_z = set_exponent(0., 999) == 0.
+  logical, parameter :: test_set_exponent_nan = transfer(set_exponent(nan, 0), 0) == nanInt
+  real,    parameter :: set_expo_inf = set_exponent(inf, 0)
+  integer, parameter :: set_expo_inf_int = transfer(set_expo_inf, 0)
+  logical, parameter :: test_set_exponent_inf = shiftr(set_expo_inf_int, 23) == 255 .and. &
+                                                shiftl(set_expo_inf_int, 9) /= 0 ! NaN
+  logical, parameter :: test_set_exponent_0 = set_exponent(1., 0) == 0.5
+  logical, parameter :: test_set_exponent_1 = set_exponent(1., 1) == 1.
+  logical, parameter :: test_set_exponent_2 = set_exponent(1., 2) == 2.
+
 end module