math.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_MATH_HPP_
#define GKO_PUBLIC_CORE_BASE_MATH_HPP_
 
 
#include <cmath>
#include <complex>
#include <cstdlib>
#include <limits>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/half.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
 
 
// type manipulations
 
 
namespace detail {
 
 
template <typename T>
struct remove_complex_impl {
    using type = T;
};
 
template <typename T>
struct remove_complex_impl<std::complex<T>> {
    using type = T;
};
 
 
template <typename T>
struct to_complex_impl {
    using type = std::complex<T>;
};
 
template <typename T>
struct to_complex_impl<std::complex<T>> {
    using type = std::complex<T>;
};
 
 
template <typename T>
struct is_complex_impl : public std::integral_constant<bool, false> {};
 
template <typename T>
struct is_complex_impl<std::complex<T>>
    : public std::integral_constant<bool, true> {};
 
 
template <typename T>
struct is_complex_or_scalar_impl : std::is_scalar<T> {};
 
template <>
struct is_complex_or_scalar_impl<half> : std::true_type {};
 
template <>
struct is_complex_or_scalar_impl<bfloat16> : std::true_type {};
 
template <typename T>
struct is_complex_or_scalar_impl<std::complex<T>>
    : is_complex_or_scalar_impl<T> {};
 
 
template <template <typename> class converter, typename T>
struct template_converter {};
 
template <template <typename> class converter, template <typename...> class T,
          typename... Rest>
struct template_converter<converter, T<Rest...>> {
    using type = T<typename converter<Rest>::type...>;
};
 
 
template <typename T, typename = void>
struct remove_complex_s {};
 
template <typename T>
struct remove_complex_s<T,
                        std::enable_if_t<is_complex_or_scalar_impl<T>::value>> {
    using type = typename detail::remove_complex_impl<T>::type;
};
 
template <typename T>
struct remove_complex_s<
    T, std::enable_if_t<!is_complex_or_scalar_impl<T>::value>> {
    using type =
        typename detail::template_converter<detail::remove_complex_impl,
                                            T>::type;
};
 
 
template <typename T, typename = void>
struct to_complex_s {};
 
template <typename T>
struct to_complex_s<T, std::enable_if_t<is_complex_or_scalar_impl<T>::value>> {
    using type = typename detail::to_complex_impl<T>::type;
};
 
template <typename T>
struct to_complex_s<T, std::enable_if_t<!is_complex_or_scalar_impl<T>::value>> {
    using type =
        typename detail::template_converter<detail::to_complex_impl, T>::type;
};
 
 
}  // namespace detail
 
 
template <typename T>
struct cpx_real_type {
    using type = T;
};
 
template <typename T>
struct cpx_real_type<std::complex<T>> {
    using type = typename std::complex<T>::value_type;
};
 
 
template <typename T>
using is_complex_s = detail::is_complex_impl<T>;
 
template <typename T>
GKO_INLINE constexpr bool is_complex()
{
    return detail::is_complex_impl<T>::value;
}
 
 
template <typename T>
using is_complex_or_scalar_s = detail::is_complex_or_scalar_impl<T>;
 
template <typename T>
GKO_INLINE constexpr bool is_complex_or_scalar()
{
    return detail::is_complex_or_scalar_impl<T>::value;
}
 
 
template <typename T>
using remove_complex = typename detail::remove_complex_s<T>::type;
 
 
template <typename T>
using to_complex = typename detail::to_complex_s<T>::type;
 
 
template <typename T>
using to_real = remove_complex<T>;
 
 
namespace detail {
 
 
// singly linked list of all our supported precisions
template <typename T>
struct next_precision_base_impl {};
 
template <>
struct next_precision_base_impl<float> {
    using type = double;
};
 
template <>
struct next_precision_base_impl<double> {
    using type = float;
};
 
template <typename T>
struct next_precision_base_impl<std::complex<T>> {
    using type = std::complex<typename next_precision_base_impl<T>::type>;
};
 
 
template <typename T, int step, typename Visited, typename... Rest>
struct find_precision_list_impl;
 
template <typename T, int step, typename... Visited, typename U,
          typename... Rest>
struct find_precision_list_impl<T, step, std::tuple<Visited...>, U, Rest...> {
    using type =
        typename find_precision_list_impl<T, step, std::tuple<Visited..., U>,
                                          Rest...>::type;
};
 
template <typename T, int step, typename... Visited, typename... Rest>
struct find_precision_list_impl<T, step, std::tuple<Visited...>, T, Rest...> {
    using tuple = std::tuple<T, Rest..., Visited...>;
    constexpr static auto tuple_size =
        static_cast<int>(std::tuple_size_v<tuple>);
    // It turns the first part into positive when step is negative
    constexpr static int index = (tuple_size + step % tuple_size) % tuple_size;
    using type = std::tuple_element_t<index, tuple>;
};
 
 
template <typename T, int step = 1>
struct find_precision_impl {
    using type = typename find_precision_list_impl<T, step, std::tuple<>,
#if GINKGO_ENABLE_HALF
                                                   half,
#endif
#if GINKGO_ENABLE_BFLOAT16
                                                   bfloat16,
#endif
                                                   float, double>::type;
};
 
 
template <typename T, int step>
struct find_precision_impl<std::complex<T>, step> {
    using type = std::complex<typename find_precision_impl<T, step>::type>;
};
 
 
template <typename T>
struct reduce_precision_impl {
    using type = T;
};
 
template <typename T>
struct reduce_precision_impl<std::complex<T>> {
    using type = std::complex<typename reduce_precision_impl<T>::type>;
};
 
template <>
struct reduce_precision_impl<double> {
    using type = float;
};
 
// for block jacobi
template <>
struct reduce_precision_impl<float> {
    using type = half;
};
 
 
template <typename T>
struct increase_precision_impl {
    using type = T;
};
 
template <typename T>
struct increase_precision_impl<std::complex<T>> {
    using type = std::complex<typename increase_precision_impl<T>::type>;
};
 
template <>
struct increase_precision_impl<float> {
    using type = double;
};
 
// for block jacobi
template <>
struct increase_precision_impl<half> {
    using type = float;
};
 
 
template <typename T>
struct infinity_impl {
    // CUDA doesn't allow us to call std::numeric_limits functions
    // so we need to store the value instead.
    static constexpr auto value = std::numeric_limits<T>::infinity();
};
 
 
template <typename T1, typename T2>
struct highest_precision_impl {
    using type = decltype(T1{} + T2{});
};
 
template <typename T1, typename T2>
struct highest_precision_impl<std::complex<T1>, std::complex<T2>> {
    using type = std::complex<typename highest_precision_impl<T1, T2>::type>;
};
 
template <typename Head, typename... Tail>
struct highest_precision_variadic {
    using type = typename highest_precision_impl<
        Head, typename highest_precision_variadic<Tail...>::type>::type;
};
 
template <typename Head>
struct highest_precision_variadic<Head> {
    using type = Head;
};
 
 
}  // namespace detail
 
 
template <typename T>
using next_precision_base = typename detail::next_precision_base_impl<T>::type;
 
 
template <typename T>
using previous_precision_base = next_precision_base<T>;
 
 
template <typename T, int step = 1>
using next_precision = typename detail::find_precision_impl<T, step>::type;
 
template <typename T, int step = 1>
using previous_precision = typename detail::find_precision_impl<T, -step>::type;
 
 
template <typename T>
using reduce_precision = typename detail::reduce_precision_impl<T>::type;
 
 
template <typename T>
using increase_precision = typename detail::increase_precision_impl<T>::type;
 
 
template <typename... Ts>
using highest_precision =
    typename detail::highest_precision_variadic<Ts...>::type;
 
 
template <typename T>
GKO_INLINE constexpr reduce_precision<T> round_down(T val)
{
    return static_cast<reduce_precision<T>>(val);
}
 
 
template <typename T>
GKO_INLINE constexpr increase_precision<T> round_up(T val)
{
    return static_cast<increase_precision<T>>(val);
}
 
 
template <typename FloatType, size_type NumComponents, size_type ComponentId>
class truncated;
 
 
namespace detail {
 
 
template <typename T>
struct truncate_type_impl {
    using type = truncated<T, 2, 0>;
};
 
template <typename T, size_type Components>
struct truncate_type_impl<truncated<T, Components, 0>> {
    using type = truncated<T, 2 * Components, 0>;
};
 
template <typename T>
struct truncate_type_impl<std::complex<T>> {
    using type = std::complex<typename truncate_type_impl<T>::type>;
};
 
 
template <typename T>
struct type_size_impl {
    static constexpr auto value = sizeof(T) * byte_size;
};
 
template <typename T>
struct type_size_impl<std::complex<T>> {
    static constexpr auto value = sizeof(T) * byte_size;
};
 
 
}  // namespace detail
 
 
template <typename T, size_type Limit = sizeof(uint16) * byte_size>
using truncate_type =
    std::conditional_t<detail::type_size_impl<T>::value >= 2 * Limit,
                       typename detail::truncate_type_impl<T>::type, T>;
 
 
template <typename S, typename R>
struct default_converter {
    GKO_ATTRIBUTES R operator()(S val) { return static_cast<R>(val); }
};
 
 
// mathematical functions
 
 
GKO_INLINE constexpr int64 ceildiv(int64 num, int64 den)
{
    return (num + den - 1) / den;
}
 
 
template <typename T>
GKO_INLINE constexpr T zero()
{
    return T{};
}
 
 
template <typename T>
GKO_INLINE constexpr T zero(const T&)
{
    return zero<T>();
}
 
 
template <typename T>
GKO_INLINE constexpr T one()
{
    return T(1);
}
 
template <>
GKO_INLINE constexpr half one<half>()
{
    constexpr auto bits = static_cast<uint16>(0b0'01111'0000000000u);
    return half::create_from_bits(bits);
}
 
template <>
GKO_INLINE constexpr bfloat16 one<bfloat16>()
{
    constexpr auto bits = static_cast<uint16>(0b0'01111111'0000000u);
    return bfloat16::create_from_bits(bits);
}
 
 
template <typename T>
GKO_INLINE constexpr T one(const T&)
{
    return one<T>();
}
 
 
template <typename T>
GKO_INLINE constexpr bool is_zero(T value)
{
    return value == zero<T>();
}
 
 
template <typename T>
GKO_INLINE constexpr bool is_nonzero(T value)
{
    return value != zero<T>();
}
 
 
template <typename T>
GKO_INLINE constexpr T max(const T& x, const T& y)
{
    return x >= y ? x : y;
}
 
 
template <typename T>
GKO_INLINE constexpr T min(const T& x, const T& y)
{
    return x <= y ? x : y;
}
 
 
namespace detail {
 
 
template <typename Ref, typename Dummy = std::void_t<>>
struct has_to_arithmetic_type : std::false_type {
    static_assert(std::is_same<Dummy, void>::value,
                  "Do not modify the Dummy value!");
    using type = Ref;
};
 
template <typename Ref>
struct has_to_arithmetic_type<
    Ref, std::void_t<decltype(std::declval<Ref>().to_arithmetic_type())>>
    : std::true_type {
    using type = decltype(std::declval<Ref>().to_arithmetic_type());
};
 
 
template <typename Ref, typename Dummy = std::void_t<>>
struct has_arithmetic_type : std::false_type {
    static_assert(std::is_same<Dummy, void>::value,
                  "Do not modify the Dummy value!");
};
 
template <typename Ref>
struct has_arithmetic_type<Ref, std::void_t<typename Ref::arithmetic_type>>
    : std::true_type {};
 
 
template <typename Ref>
constexpr GKO_ATTRIBUTES
    std::enable_if_t<has_to_arithmetic_type<Ref>::value,
                     typename has_to_arithmetic_type<Ref>::type>
    to_arithmetic_type(const Ref& ref)
{
    return ref.to_arithmetic_type();
}
 
template <typename Ref>
constexpr GKO_ATTRIBUTES std::enable_if_t<!has_to_arithmetic_type<Ref>::value &&
                                              has_arithmetic_type<Ref>::value,
                                          typename Ref::arithmetic_type>
to_arithmetic_type(const Ref& ref)
{
    return ref;
}
 
template <typename Ref>
constexpr GKO_ATTRIBUTES std::enable_if_t<!has_to_arithmetic_type<Ref>::value &&
                                              !has_arithmetic_type<Ref>::value,
                                          Ref>
to_arithmetic_type(const Ref& ref)
{
    return ref;
}
 
 
// Note: All functions have postfix `impl` so they are not considered for
// overload resolution (in case a class / function also is in the namespace
// `detail`)
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<!is_complex_s<T>::value, T>
real_impl(const T& x)
{
    return x;
}
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<is_complex_s<T>::value,
                                                     remove_complex<T>>
real_impl(const T& x)
{
    return x.real();
}
 
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<!is_complex_s<T>::value, T>
imag_impl(const T&)
{
    return T{};
}
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<is_complex_s<T>::value,
                                                     remove_complex<T>>
imag_impl(const T& x)
{
    return x.imag();
}
 
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<!is_complex_s<T>::value, T>
conj_impl(const T& x)
{
    return x;
}
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr std::enable_if_t<is_complex_s<T>::value, T>
conj_impl(const T& x)
{
    return T{real_impl(x), -imag_impl(x)};
}
 
 
}  // namespace detail
 
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr auto real(const T& x)
{
    return detail::real_impl(detail::to_arithmetic_type(x));
}
 
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr auto imag(const T& x)
{
    return detail::imag_impl(detail::to_arithmetic_type(x));
}
 
 
template <typename T>
GKO_ATTRIBUTES GKO_INLINE constexpr auto conj(const T& x)
{
    return detail::conj_impl(detail::to_arithmetic_type(x));
}
 
 
template <typename T>
GKO_INLINE constexpr auto squared_norm(const T& x)
    -> decltype(real(conj(x) * x))
{
    return real(conj(x) * x);
}
 
using std::abs;
 
template <typename T>
GKO_INLINE constexpr std::enable_if_t<!is_complex_s<T>::value, T> abs(
    const T& x)
{
    return x >= zero<T>() ? x : -x;
}
 
 
template <typename T>
GKO_INLINE constexpr std::enable_if_t<is_complex_s<T>::value, remove_complex<T>>
abs(const T& x)
{
    return sqrt(squared_norm(x));
}
 
// increase the priority in function lookup
GKO_INLINE gko::half abs(const std::complex<gko::half>& x)
{
    // Using float abs not sqrt on norm to avoid overflow
    return static_cast<gko::half>(abs(std::complex<float>(x)));
}
 
GKO_INLINE gko::bfloat16 abs(const std::complex<gko::bfloat16>& x)
{
    // Using float abs not sqrt on norm to avoid overflow
    return static_cast<gko::bfloat16>(abs(std::complex<float>(x)));
}
 
 
using std::sqrt;
 
GKO_INLINE gko::half sqrt(gko::half a)
{
    return gko::half(std::sqrt(float(a)));
}
 
GKO_INLINE std::complex<gko::half> sqrt(std::complex<gko::half> a)
{
    return std::complex<gko::half>(sqrt(std::complex<float>(
        static_cast<float>(a.real()), static_cast<float>(a.imag()))));
}
 
GKO_INLINE gko::bfloat16 sqrt(gko::bfloat16 a)
{
    return gko::bfloat16(std::sqrt(float(a)));
}
 
GKO_INLINE std::complex<gko::bfloat16> sqrt(std::complex<gko::bfloat16> a)
{
    return std::complex<gko::bfloat16>(sqrt(std::complex<float>(
        static_cast<float>(a.real()), static_cast<float>(a.imag()))));
}
 
 
template <typename T>
GKO_INLINE constexpr T pi()
{
    return static_cast<T>(3.1415926535897932384626433);
}
 
 
template <typename T>
GKO_INLINE constexpr std::complex<remove_complex<T>> unit_root(int64 n,
                                                               int64 k = 1)
{
    return std::polar(one<remove_complex<T>>(),
                      remove_complex<T>{2} * pi<remove_complex<T>>() * k / n);
}
 
 
template <typename T>
constexpr uint32 get_significant_bit(const T& n, uint32 hint = 0u) noexcept
{
    return (T{1} << (hint + 1)) > n ? hint : get_significant_bit(n, hint + 1u);
}
 
 
template <typename T>
constexpr T get_superior_power(const T& base, const T& limit,
                               const T& hint = T{1}) noexcept
{
    return hint >= limit ? hint : get_superior_power(base, limit, hint * base);
}
 
 
template <typename T>
GKO_INLINE GKO_ATTRIBUTES std::enable_if_t<!is_complex_s<T>::value, bool>
is_finite(const T& value)
{
    constexpr T infinity{detail::infinity_impl<T>::value};
    return abs(value) < infinity;
}
 
 
template <typename T>
GKO_INLINE GKO_ATTRIBUTES std::enable_if_t<is_complex_s<T>::value, bool>
is_finite(const T& value)
{
    return is_finite(value.real()) && is_finite(value.imag());
}
 
 
template <typename T>
GKO_INLINE GKO_ATTRIBUTES T safe_divide(T a, T b)
{
    return b == zero<T>() ? zero<T>() : a / b;
}
 
 
template <typename T>
GKO_DEPRECATED(
    "is_nan can't be used safely on the device (MSVC+CUDA), and will thus be "
    "removed in a future release, without replacement")
GKO_INLINE GKO_ATTRIBUTES
    std::enable_if_t<!is_complex_s<T>::value, bool> is_nan(const T& value)
{
    using std::isnan;
    return isnan(value);
}
 
 
template <typename T>
GKO_DEPRECATED(
    "is_nan can't be used safely on the device (MSVC+CUDA), and will thus be "
    "removed in a future release, without replacement")
GKO_INLINE GKO_ATTRIBUTES std::enable_if_t<is_complex_s<T>::value, bool> is_nan(
    const T& value)
{
    return is_nan(value.real()) || is_nan(value.imag());
}
 
 
template <typename T>
GKO_INLINE constexpr std::enable_if_t<!is_complex_s<T>::value, T> nan()
{
    return std::numeric_limits<T>::quiet_NaN();
}
 
 
template <typename T>
GKO_INLINE constexpr std::enable_if_t<is_complex_s<T>::value, T> nan()
{
    return T{nan<remove_complex<T>>(), nan<remove_complex<T>>()};
}
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_MATH_HPP_