src/__support/math/atanbf16.h

//===-- Implementation for atanbf16(x) function ---------------------------===//
//
// 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.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_LIBC_SRC___SUPPORT_MATH_ATANBF16_H
#define LLVM_LIBC_SRC___SUPPORT_MATH_ATANBF16_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/multiply_add.h"
#include "src/__support/FPUtil/sqrt.h"
#include "src/__support/macros/optimization.h"

namespace LIBC_NAMESPACE_DECL {

namespace math {

LIBC_INLINE bfloat16 atanbf16(bfloat16 x) {
  // Generated by Sollya using the following command:
  // > display = hexadecimal;
  // > round(pi/2, SG, RN);
  constexpr float PI_2 = 0x1.921fb6p0f;

  using FPBits = fputil::FPBits<bfloat16>;
  FPBits xbits(x);

  uint16_t x_u = xbits.uintval();
  uint16_t x_abs = x_u & 0x7fff;
  bool x_sign = x_u >> 15;
  float sign = (x_sign ? -1.0f : 1.0f);

  // Taylor series -> [x - x^3/3 + x^5/5 - x^7/7 ...]
  // x * [1 - x^2/3 + x^4/5 - x^6/7...] -> x * P(x)
  // atan(x) = x * poly(x^2)
  // atan(x)/x = poly(x^2)

  // Degree 14 polynomial of atan(x) generated using Sollya with command :
  // > display = hexadecimal ;
  // > P = fpminimax(atan(x)/x, [|0, 2, 4, 6, 8, 10, 12, 14|], [|1, SG,
  // SG..SG|], [0, 1]);
  //
  // relative error for the polynomial given by:
  // > dirtyinfnorm(atan(x)/x - P(x), [0, 1]);
  // gives error ~ 0x1.ee44001p-24
  // worst case error for it being ~ 0x1.dcf750p-23
  // satisfying -> error < worst_case
  auto atan_eval = [](float x0) {
    return fputil::polyeval(x0, -0x1.5552c4p-2f, 0x1.990f2p-3f, -0x1.1f7dccp-3f,
                            0x1.97e49p-4f, -0x1.ebff34p-5f, 0x1.938e46p-6f,
                            -0x1.3a28bcp-8f);
  };

  float xf = x;
  float x_sq = xf * xf;

  // |x| <= 1
  if (x_abs <= 0x3f80) {
    // atanbf16(+/-0) = +/-0
    if (LIBC_UNLIKELY(x_abs == 0))
      return x;

    // For smaller x
    if (LIBC_UNLIKELY(x_abs <= 0x3db8))
      return fputil::cast<bfloat16>(fputil::multiply_add(xf, -0x1p-9f, xf));

    float result = atan_eval(x_sq);
    return fputil::cast<bfloat16>(fputil::multiply_add(xf * x_sq, result, xf));
  }

  // |x| is +/-inf or NaN
  if (LIBC_UNLIKELY(x_abs >= 0x7F80)) {
    // NaN
    if (xbits.is_nan()) {
      if (xbits.is_signaling_nan()) {
        fputil::raise_except_if_required(FE_INVALID);
        return FPBits::quiet_nan().get_val();
      }
      return x;
    }
    // atanbf16(+/-inf) = +/-pi/2
    return fputil::cast<bfloat16>(sign * PI_2);
  }

  // If |x| > 1:
  // atan(x) = sign(x) * (pi/2 - atan(1/|x|))
  // Since 1/|x| < 1, we can use the same polynomial.
  float x_inv_sq = 1.0f / x_sq;
  float x_inv = fputil::sqrt<float>(x_inv_sq);

  float result = atan_eval(x_inv_sq);
  float atan_inv = fputil::multiply_add(x_inv * x_inv_sq, result, x_inv);
  return fputil::cast<bfloat16>(sign * (PI_2 - atan_inv));
}

} // namespace math
} // namespace LIBC_NAMESPACE_DECL

#endif // LLVM_LIBC_SRC___SUPPORT_MATH_ATANBF16_H