// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#ifndef BFLOAT16_H_
#define BFLOAT16_H_

#include <iostream>
#include <miopen/config.h>

class bfloat16
{
public:
    bfloat16() : data_{0} {}
    explicit bfloat16(float rhs)
    {
        union
        {
            float float_st;
            std::uint32_t bf16_st;
        } bits_st = {rhs};

        // BF16 round and NaN preservation code matches
        // https://github.com/ROCm/rocBLAS/blob/develop/library/include/rocblas_bfloat16.h
        if((~bits_st.bf16_st & 0x7f800000) == 0) // Inf or NaN
        {
            // When all of the exponent bits are 1, the value is Inf or NaN.
            // Inf is indicated by a zero mantissa. NaN is indicated by any nonzero
            // mantissa bit. Quiet NaN is indicated by the most significant mantissa
            // bit being 1. Signaling NaN is indicated by the most significant
            // mantissa bit being 0 but some other bit(s) being 1. If any of the
            // lower 16 bits of the mantissa are 1, we set the least significant bit
            // of the bfloat16 mantissa, in order to preserve signaling NaN in case
            // the bloat16's mantissa bits are all 0.
            if((bits_st.bf16_st & 0xffff) != 0)
            {
                bits_st.bf16_st |= 0x10000; // Preserve signaling NaN
            }
        }
        else
        {
#if MIOPEN_USE_RNE_BFLOAT16 == 1
            // When the exponent bits are not all 1s, then the value is zero, normal,
            // or subnormal. We round the bfloat16 mantissa up by adding 0x7FFF, plus
            // 1 if the least significant bit of the bfloat16 mantissa is 1 (odd).
            // This causes the bfloat16's mantissa to be incremented by 1 if the 16
            // least significant bits of the float mantissa are greater than 0x8000,
            // or if they are equal to 0x8000 and the least significant bit of the
            // bfloat16 mantissa is 1 (odd). This causes it to be rounded to even when
            // the lower 16 bits are exactly 0x8000. If the bfloat16 mantissa already
            // has the value 0x7f, then incrementing it causes it to become 0x00 and
            // the exponent is incremented by one, which is the next higher FP value
            // to the unrounded bfloat16 value. When the bfloat16 value is subnormal
            // with an exponent of 0x00 and a mantissa of 0x7F, it may be rounded up
            // to a normal value with an exponent of 0x01 and a mantissa of 0x00.
            // When the bfloat16 value has an exponent of 0xFE and a mantissa of 0x7F,
            // incrementing it causes it to become an exponent of 0xFF and a mantissa
            // of 0x00, which is Inf, the next higher value to the unrounded value.
            bits_st.bf16_st +=
                (0x7fff + ((bits_st.bf16_st >> 16) & 1)); // Round to nearest, round to even
#else                                                     // truncation
// do nothing
#endif
        }
        data_ = bits_st.bf16_st >> 16;
    }
    operator float() const
    {
        union
        {
            std::uint32_t bf16_st;
            float float_st;
        } bits_st = {data_};

        bits_st.bf16_st = bits_st.bf16_st << 16;
        return bits_st.float_st;
    }

    bfloat16 operator-() const { return bfloat16(-static_cast<float>(*this)); }
    bfloat16 operator+() const { return *this; }

    bfloat16& operator=(const float rhs)
    {
        *this = bfloat16(rhs);
        return *this;
    }
    bfloat16& operator+=(bfloat16 rhs)
    {
        *this = bfloat16(static_cast<float>(*this) + static_cast<float>(rhs));
        return *this;
    }

    bfloat16& operator+=(float rhs)
    {
        *this = bfloat16(static_cast<float>(*this) + rhs);
        return *this;
    }

    bfloat16& operator-=(bfloat16 rhs)
    {
        *this += -rhs;
        return *this;
    }
    bfloat16& operator*=(bfloat16 rhs)
    {
        *this = bfloat16(static_cast<float>(*this) * static_cast<float>(rhs));
        return *this;
    }
    bfloat16& operator*=(float rhs)
    {
        *this = bfloat16(static_cast<float>(*this) * rhs);
        return *this;
    }

    bfloat16& operator/=(bfloat16 rhs)
    {
        *this = bfloat16(static_cast<float>(*this) / static_cast<float>(rhs));
        return *this;
    }
    bool operator<(bfloat16 rhs) const
    {
        return static_cast<float>(*this) < static_cast<float>(rhs);
    }
    bool operator==(bfloat16 rhs) const { return std::equal_to<float>()(*this, rhs); }

    static constexpr bfloat16 generate(uint16_t val) { return bfloat16{val, true}; }

private:
    constexpr bfloat16(std::uint16_t val, bool) : data_{val} {}

    std::uint16_t data_;
};

inline bfloat16 operator+(bfloat16 a, const bfloat16& b)
{
    a += b;
    return a;
}

inline bfloat16 operator-(bfloat16 a, const bfloat16& b)
{
    a -= b;
    return a;
}

inline bfloat16 operator*(bfloat16 a, const bfloat16& b)
{
    a *= b;
    return a;
}

inline bfloat16 operator/(bfloat16 a, const bfloat16& b)
{
    a /= b;
    return a;
}

namespace std {
template <>
class numeric_limits<bfloat16>
{
public:
    static constexpr bool is_specialized = true;
    static constexpr bfloat16 min() noexcept { return bfloat16::generate(0x0080); } // 0x1.00p-126
    static constexpr bfloat16 max() noexcept { return bfloat16::generate(0x7F7F); }
    static constexpr bfloat16 lowest() noexcept { return bfloat16::generate(0xFF7F); }
    static constexpr bfloat16 epsilon() noexcept { return bfloat16::generate(0x3C00); }
    static constexpr bfloat16 infinity() noexcept { return bfloat16::generate(0x7F80); }
    static constexpr bfloat16 quiet_NaN() noexcept { return bfloat16::generate(0x7FC0); } // qnan(0)
    static constexpr bfloat16 signaling_NaN() noexcept
    {
        return bfloat16::generate(0x7F81); // snan(1)
    }
    static constexpr bfloat16 denorm_min() noexcept
    {
        return bfloat16::generate(0x0001); // 0x0.02p-126
    }
};
} // namespace std
#endif
