[flang] Fix test regression from SQRT folding

[klausler]

Summary

The algorithm used to fold SQRT has some holes that
led to test failures; debug and add more tests.

[jeanPerier]

LGTM

flang/lib/Evaluate/real.cpp

if (IsZero()) {
       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();
-    }
+    if (expo < -1 || expo > 1) {
+      // Reduce the range to [0.5 .. 4.0) by dividing by an integral power
+      // of four to avoid trouble with very large and very small values
+      // (esp. truncation of subnormals).
+      // SQRT(2**(2a) * x) = SQRT(2**(2a)) * SQRT(x) = 2**a * SQRT(x)
+      Real scaled;
+      int adjust{expo / 2};
+      scaled.Normalize(false, expo - 2 * adjust + exponentBias, GetFraction());
+      result = scaled.SQRT(rounding);
+      result.value.Normalize(false,
+          result.value.UnbiasedExponent() + adjust + exponentBias,
+          result.value.GetFraction());
+      return result;
+    }
+    // Compute the square root of the reduced value with the slow but
+    // reliable bit-at-a-time method.  Start with a clear significand and
+    // half of the unbiased exponent, and then try to set significand bits
+    // in descending order of magnitude without exceeding the exact result.
     expo = expo / 2 + exponentBias;
     result.value.Normalize(false, expo, Fraction::MASKL(1));
+    Real initialSq{result.value.Multiply(result.value).value};
+    if (Compare(initialSq) == Relation::Less) {
+      // Initial estimate is too large; this can happen for values just
+      // under 1.0.
+      --expo;
+      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.
+    // The computed square root has a square that's not greater than the
+    // original argument.  Check this square against the square of the next
+    // larger Real 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};

flang/test/Evaluate/folding28.f90

!WARN: division by zero
   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'
+  real(8), parameter :: sqrt_maxs8 = sqrt(maxs8), sqrt_maxs8z = z'1fffffffffffffff'
   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
+  ! regression tests: cases near 1.
+  real(4), parameter :: sqrt_under1 = sqrt(.96875)
+  logical, parameter :: test_sqrt_under1 = sqrt_under1 == .984250962734222412109375
+  ! oddball case: the value before 1. is also its own sqrt, but not its own square
+  real(4), parameter :: before_1 = z'3f7fffff' ! .999999940395355224609375
+  real(4), parameter :: sqrt_before_1 = sqrt(before_1)
+  logical, parameter :: test_before_1 = sqrt_before_1 == before_1
+  real(4), parameter :: sq_sqrt_before_1 = sqrt_before_1 * sqrt_before_1
+  logical, parameter :: test_sq_before_1 = sq_sqrt_before_1 < before_1
 end module
-