libc/src/__support/math/exp2bf16.h

//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Implementation of the exp2bf16 function, which computes 2^x for bfloat16.
///
//===----------------------------------------------------------------------===//

#ifndef LLVM_LIBC_SRC___SUPPORT_MATH_EXP2BF16_H
#define LLVM_LIBC_SRC___SUPPORT_MATH_EXP2BF16_H

#include "hdr/fenv_macros.h"
#include "src/__support/FPUtil/FEnvImpl.h"
#include "src/__support/FPUtil/FPBits.h"
#include "src/__support/FPUtil/PolyEval.h"
#include "src/__support/FPUtil/bfloat16.h"
#include "src/__support/FPUtil/cast.h"
#include "src/__support/FPUtil/except_value_utils.h"
#include "src/__support/FPUtil/multiply_add.h"
#include "src/__support/FPUtil/nearest_integer.h"
#include "src/__support/macros/attributes.h"
#include "src/__support/macros/config.h"
#include "src/__support/macros/optimization.h"

namespace LIBC_NAMESPACE_DECL {
namespace math {

namespace exp2bf16_internal {
// Generated by Sollya with the following commands:
//   > display = hexadecimal;
//   > for i in [|-126.5, -127, -127.5, -128, -129, -130, -131, -132, -133|]
//     do print(round(2^i, SG, RN) @ "f,");
LIBC_INLINE_VAR constexpr float EXP2_DENORMAL[9] = {
    0x1.6a09e8p-127f, 0x1p-127f, 0x1.6a09e8p-128f, 0x1p-128f, 0x1p-129f,
    0x1p-130f,        0x1p-131f, 0x1p-132f,        0x1p-133f,
};

// Generated by Sollya with the following commands:
//   > display = hexadecimal;
//   > for i from 0 to 15 do print(round(2^(i/16), SG, RN) @ "f,");
LIBC_INLINE_VAR constexpr float EXP2_MID[16] = {
    0x1p0f,        0x1.0b5586p0f, 0x1.172b84p0f, 0x1.2387a6p0f,
    0x1.306fep0f,  0x1.3dea64p0f, 0x1.4bfdaep0f, 0x1.5ab07ep0f,
    0x1.6a09e6p0f, 0x1.7a1148p0f, 0x1.8ace54p0f, 0x1.9c4918p0f,
    0x1.ae89fap0f, 0x1.c199bep0f, 0x1.d5818ep0f, 0x1.ea4afap0f,
};
} // namespace exp2bf16_internal

LIBC_INLINE bfloat16 exp2bf16(bfloat16 x) {
  using FPBits = fputil::FPBits<bfloat16>;
  FPBits x_bits(x);
  uint16_t x_uint = x_bits.uintval();

  // We first check for NaN
  if (LIBC_UNLIKELY(x_bits.is_nan())) {
    // If it is a signaling NaN, then we raise an FE_INVALID exception
    // and return a quiet NaN
    if (x_bits.is_signaling_nan()) {
      fputil::raise_except_if_required(FE_INVALID);
      return FPBits::quiet_nan().get_val();
    }
    // If not, just return the NaN
    return x;
  }

  // exp2bf16 may only give numeric output for input x in (-134.0, 128.0).
  // All the remaining exceptional cases (overflow, underflow and denormal
  // outputs) have |x| >= 126.5, so we filter them with a single check.
  constexpr uint16_t EXCEPTIONAL_ABS_BASE = 0x42FD; // 0x1.fap6 = 126.5
  if (LIBC_UNLIKELY(x_bits.abs().uintval() >= EXCEPTIONAL_ABS_BASE)) {
    if (x_bits.is_pos()) {
      // Overflow: x >= 128.0 (0x4300 = 0x1p7)
      if (x_uint >= 0x4300) {
        // exp2(Inf) = Inf
        if (x_bits.is_inf())
          return x;

        return fputil::set_overflow<bfloat16>();
      }
    } else {
      // Underflow: x <= -134.0 (0xC306 = -0x1.0cp7)
      if (x_uint >= 0xC306) {
        // exp2(-Inf) = 0
        if (x_bits.is_inf())
          return FPBits::zero().get_val();

        return fputil::set_underflow<bfloat16>();
      }

      // Handle denormal output separately, by using a LUT
      // We are handling these separately because the trick of adding hi
      // directly to the exponent of the result will not work if the
      // output is denormal. There are only 9 such inputs and they are
      // all back-to-back when expressed as an uint16_t, with
      // x_uint >= 0xC2FD (-0x1.fap6 = -126.5) and
      // x_uint <= 0xC305 (-0x1.0ap7 = -133.0).
      constexpr uint16_t DENORMAL_LUT_BASE = 0xC2FD;
      return fputil::cast<bfloat16>(
          exp2bf16_internal::EXP2_DENORMAL[x_uint - DENORMAL_LUT_BASE]);
    }
  }

  // For inputs not handled above in -134.0 < x < 128.0,
  // we perform the following range reduction:
  // x = hi + mid + lo
  // where,
  //      hi is an integer
  //      mid is a grid point of the form j / 2^n in [0, 1)
  //      lo is what's leftover ie x - k / 2^n
  //         and |lo| <= 2^(-n-1)
  // with,
  //      n = 4
  //      k = round(2^n * x)
  //      hi = k >> n
  //      j = k & (2^n - 1)
  // then,
  //      exp2(x) = exp2(hi + mid + lo)
  //              = exp2(hi) * exp2(mid) * exp2(lo)
  //
  // To get exp2(hi), we'll just add hi to the exponent bits of the
  // result of exp2(mid) * exp2(lo), which is exact.
  // exp2(mid) we shall obtain by consulting our LUT,
  // exp2(lo) by the polynomial.

  float x_float = fputil::cast<float>(x);
  float k_float = fputil::nearest_integer(x_float * 16.0f);

  int k_int = static_cast<int>(k_float);
  int hi = k_int >> 4;
  int j = k_int & (16 - 1);

  // lo = x - (hi + mid) = x - k / 2^n
  //    = k * (-2^(-n)) + x
  float lo = fputil::multiply_add(k_float, -0.0625f, x_float);

  float exp2_mid = exp2bf16_internal::EXP2_MID[j];

  // Degree-3 minimax polynomial generated by Sollya with the following
  // commands:
  //   > display = hexadecimal;
  //   > P = fpminimax(2^x, [|1, 2, 3|], [|SG...|], [-2^-5, 2^-5], relative,
  //                   floating, 1);
  //   0x1p0 + x * (0x1.62e43p-1 + x * (0x1.ebfff4p-3 + x * 0x1.c6aceep-5))
  float exp2_lo = fputil::polyeval(lo, 0x1p0f, 0x1.62e43p-1f, 0x1.ebfff4p-3f,
                                   0x1.c6aceep-5f);

  // We tried using Estrin's scheme as follows,
  //    float lo2 = lo * lo;
  //    float p0 = fputil::multiply_add(lo, 0x1.62e43p-1f, 0x1p0f);
  //    float p1 = fputil::multiply_add(lo, 0x1.c6aceep-5f, 0x1.ebfff4p-3f);
  //    float exp2_lo = fputil::multiply_add(lo2, p1, p0);
  // however that led to a 1-ULP error in all RU, negative, small input cases.
  // The latency increase wasn't significant (~1%) so we reverted back
  // to Horner's scheme.

  float res = exp2_mid * exp2_lo;

  fputil::FPBits<float> res_bits(res);
  uint32_t res_int = res_bits.uintval();

  uint32_t exp_hi = hi << fputil::FPBits<float>::FRACTION_LEN;
  res_int += exp_hi;

  res = fputil::FPBits<float>(res_int).get_val();
  return fputil::cast<bfloat16>(res);
}

} // namespace math
} // namespace LIBC_NAMESPACE_DECL

#endif // LLVM_LIBC_SRC___SUPPORT_MATH_EXP2BF16_H