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

#include "ck_tile/core/config.hpp"
#include "ck_tile/core/utility/bit_cast.hpp"
#include "ck_tile/core/numeric/half.hpp"
#include "ck_tile/core/numeric/integral_constant.hpp"
#include "ck_tile/core/numeric/numeric.hpp"
#include "ck_tile/core/numeric/ext_vector_base.hpp"
#if CK_TILE_USE_LLVM_BUILTIN_BF16
#include <hip/hip_bfloat16.h>
#endif
#include <stdint.h>

#pragma once

namespace ck_tile {

enum class bf16_rounding_mode
{
    standard = 0, // rtn
    truncate_with_nan,
    truncate,
    standard_asm,
    rta_asm,       // round to nearest away
    standard_cnan, // rtn with canonical NaN
};

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE constexpr uint16_t float_to_bf16_raw(float f, constant<rounding> = {});

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE constexpr uint16_t double_to_bf16_raw(double f, constant<rounding> = {});

CK_TILE_HOST_DEVICE
constexpr float bf16_to_float_raw(uint16_t x);

CK_TILE_HOST_DEVICE
constexpr double bf16_to_double_raw(uint16_t x);

#if CK_TILE_USE_CUSTOM_DATA_TYPE
// HIP use __hip_bfloat16 as struct
struct alignas(2) bfloat16_t
{
    using raw_type = uint16_t;
    raw_type data;

    CK_TILE_HOST_DEVICE
    static constexpr bfloat16_t bit_cast(raw_type x)
    {
        bfloat16_t y;
        y.data = x;
        return y;
    }

    // constructor
    constexpr bfloat16_t() : data() {}

    // construct from float
    CK_TILE_HOST_DEVICE
    explicit constexpr bfloat16_t(const float& x) : data(float_to_bf16_raw(x)) {}

    // construct from double
    CK_TILE_HOST_DEVICE
    explicit constexpr bfloat16_t(const double& x) : data(double_to_bf16_raw(x)) {}

    // construct from int
    CK_TILE_HOST_DEVICE
    explicit constexpr bfloat16_t(const int& x) : data(float_to_bf16_raw(static_cast<float>(x))) {}

    // construct from unsigned int
    CK_TILE_HOST_DEVICE
    explicit constexpr bfloat16_t(const unsigned int& x)
        : data(float_to_bf16_raw(static_cast<float>(x)))
    {
    }

    // cast to float
    CK_TILE_HOST_DEVICE
    explicit constexpr operator float() const { return bf16_to_float_raw(data); }

    // cast to float
    CK_TILE_HOST_DEVICE
    explicit constexpr operator double() const { return bf16_to_double_raw(data); }

    // cast to int
    CK_TILE_HOST_DEVICE
    explicit constexpr operator int() const { return static_cast<int>(bf16_to_float_raw(data)); }

    // internal access
    CK_TILE_HOST_DEVICE
    constexpr raw_type& get() { return data; }

    CK_TILE_HOST_DEVICE
    constexpr raw_type get() const { return data; }
};
template <typename>
struct native_t;

template <>
struct native_t<bfloat16_t>
{
    using type = ushort;
};
using bf16_t     = bfloat16_t;
using bf16_raw_t = typename bf16_t::raw_type;
#else
#if CK_TILE_USE_LLVM_BUILTIN_BF16
using bfloat16_t = __bf16;
#else
using bfloat16_t = ushort;
#endif
using bf16_t     = bfloat16_t;
using bf16_raw_t = uint16_t;
#endif
// round to nearest
CK_TILE_HOST_DEVICE
constexpr uint16_t float_to_bf16_rtn_raw(float f)
{
    uint32_t bits = bit_cast<uint32_t>(f);
    if(~bits & 0x7f800000)
    {
        // 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 += 0x7fff + ((bits >> 16) & 1); // Round to nearest, round to even
    }
    else if(bits & 0xffff)
    {
        // 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.
        bits |= 0x10000; // Preserve signaling NaN
    }
    return uint16_t(bits >> 16);
}

CK_TILE_HOST
constexpr uint16_t float_to_bf16_rtn_asm(float f) { return float_to_bf16_rtn_raw(f); }

CK_TILE_DEVICE
uint16_t float_to_bf16_rtn_asm(float f)
{
    union
    {
        float fp32;
        uint32_t int32;
    } u = {f};

    static constexpr uint32_t FP32_NAN            = 0x7fff0000;
    static constexpr uint32_t ROUND_BIAS_FOR_BF16 = 0x7fff;

#if defined(__GFX9__)
    using uint32x2_t = uint32_t __attribute__((ext_vector_type(2)));
    uint32x2_t check_nan;
#else
    uint32_t check_nan;
#endif
    uint32_t tmp;
    asm volatile("\n \
            v_cmp_u_f32 %0, %2, %2 \n \
            v_bfe_u32 %1, %2, 16, 1 \n \
            v_add3_u32 %1, %2, %1, %3 \n \
            v_cndmask_b32 %2, %1, %4, %0 \n \
            v_lshrrev_b32 %2, 16, %2 \n \
            "
                 : "=s"(check_nan), "+v"(tmp), "+v"(u.fp32)
                 : "v"(ROUND_BIAS_FOR_BF16), "v"(FP32_NAN));

    return uint16_t(u.int32);
}

// TODO: do we need this on host?
CK_TILE_HOST
uint16_t float_to_bf16_rta_asm(float f) { return float_to_bf16_rtn_raw(f); }

CK_TILE_DEVICE
uint16_t float_to_bf16_rta_asm(float f)
{
    union
    {
        float fp32;
        struct
        {
            uint16_t lo;
            uint16_t hi;
        };
    } u = {f};

    const uint32_t low_nan = 0x7fff;
    const uint32_t hi_nan  = 0x7fff0000;

#if defined(__GFX9__)
    using uint32x2_t = uint32_t __attribute__((ext_vector_type(2)));
    uint32x2_t check_nan;
#else
    uint32_t check_nan;
#endif

    asm volatile("v_cmp_u_f32 %[s_cnan], %[v_x], %[v_x] \n"
                 "v_add3_u32 %[v_x], %[v_x], %[v_blo], 1 \n"
                 "v_cndmask_b32 %[v_x], %[v_x], %[v_bhi], %[s_cnan]"
                 : [s_cnan] "+s"(check_nan), [v_x] "+v"(u.fp32)
                 : [v_blo] "v"(low_nan), [v_bhi] "v"(hi_nan));

    // Note: in above code snipet, we use hi 16 bit
    return u.hi;
}

CK_TILE_HOST_DEVICE
constexpr bool float_is_nan_raw(float f)
{
#if defined(__has_builtin) && __has_builtin(__builtin_isnan)
    return __builtin_isnan(f);
#else
    uint32_t bits                = bit_cast<uint32_t>(f);
    constexpr uint32_t exp_mask  = 0x7f800000;
    constexpr uint32_t mant_mask = 0x007fffff;

    return (bits & exp_mask) == exp_mask && (bits & mant_mask);
#endif
}

// Round to nearest even, but canonicalize any NaN input to the canonical quiet bf16 NaN
// (`0x7fff`). Unlike `float_to_bf16_rtn_raw`, this does not preserve signaling NaN
// payload/state.
CK_TILE_HOST_DEVICE
constexpr uint16_t float_to_bf16_rtn_cnan_raw(float f)
{
#if defined(__FAST_MATH__) || (defined(__FINITE_MATH_ONLY__) && __FINITE_MATH_ONLY__)
    // Fast/finite-math can fold the NaN predicate away, so fall back to standard RTN.
    return float_to_bf16_rtn_raw(f);
#else
    // `-fgpu-flush-denormals-to-zero` only affects denormals, not NaN handling.
    uint32_t bits = bit_cast<uint32_t>(f);
    uint32_t tmp  = (bits >> 16) & 1;
    uint32_t res  = float_is_nan_raw(f) ? 0x7fff0000 : bits + tmp + 0x7fff;

    return uint16_t(res >> 16);
#endif
}

// Truncate instead of rounding, preserving SNaN
CK_TILE_HOST_DEVICE
constexpr uint16_t float_to_bf16_truc_nan_raw(float f)
{
    uint32_t bits = bit_cast<uint32_t>(f);
    return static_cast<uint16_t>(bits >> 16) | (!(~bits & 0x7f800000) && (bits & 0xffff));
}

// Fast truncate instead of rounding, RTZ
CK_TILE_HOST_DEVICE
constexpr uint16_t float_to_bf16_truc_raw(float f)
{
    uint32_t bits = bit_cast<uint32_t>(f);
    return static_cast<uint16_t>(bits >> 16);
}

template <bf16_rounding_mode rounding>
CK_TILE_HOST_DEVICE constexpr uint16_t float_to_bf16_raw(float f, constant<rounding>)
{
    if constexpr(rounding == bf16_rounding_mode::standard)
        return float_to_bf16_rtn_raw(f);
    else if constexpr(rounding == bf16_rounding_mode::standard_asm)
        return float_to_bf16_rtn_asm(f);
    else if constexpr(rounding == bf16_rounding_mode::standard_cnan)
        return float_to_bf16_rtn_cnan_raw(f);
    else if constexpr(rounding == bf16_rounding_mode::truncate_with_nan)
        return float_to_bf16_truc_nan_raw(f);
    else if constexpr(rounding == bf16_rounding_mode::rta_asm)
        return float_to_bf16_rta_asm(f);
    else
        return float_to_bf16_truc_raw(f);
}

template <bf16_rounding_mode rounding>
CK_TILE_HOST_DEVICE constexpr uint16_t double_to_bf16_raw(double f, constant<rounding>)
{
    return float_to_bf16_raw(static_cast<float>(f), constant<rounding>{});
}

CK_TILE_HOST_DEVICE
constexpr float bf16_to_float_raw(uint16_t x)
{
    union
    {
        uint32_t int32;
        float fp32;
    } u = {uint32_t(x) << 16};
    return u.fp32;
}

CK_TILE_HOST_DEVICE
constexpr double bf16_to_double_raw(uint16_t x)
{
    return static_cast<double>(bf16_to_float_raw(x));
}

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE constexpr bfloat16_t float_to_bf16(float f, constant<rounding> = {})
{
// Use builtin bfloat16 conversion only on gfx950 as its predecessors do not support bf16 cvt
// instructions, resulting in suboptimal performance; Add host side marcro check for consistency
// during accuracy tests.
#if CK_TILE_USE_LLVM_BUILTIN_BF16 && (defined(__gfx950__) || defined(CK_GFX950_SUPPORT))
    return static_cast<bfloat16_t>(f);
#else
    return bit_cast<bfloat16_t>(float_to_bf16_raw(f, constant<rounding>{}));
#endif
}

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE constexpr bfloat16_t double_to_bf16(double f, constant<rounding> = {})
{
    return bit_cast<bfloat16_t>(double_to_bf16_raw(f, constant<rounding>{}));
}

CK_TILE_HOST_DEVICE
constexpr float bf16_to_float(bfloat16_t x) { return bf16_to_float_raw(bit_cast<uint16_t>(x)); }

CK_TILE_HOST_DEVICE
constexpr double bf16_to_double(bfloat16_t x) { return static_cast<double>(bf16_to_float_raw(x)); }

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE bfloat16_t constexpr fp16_to_bf16(half_t f, constant<rounding> = {})
{
    return bit_cast<bfloat16_t>(float_to_bf16_raw(static_cast<float>(f), constant<rounding>{}));
}

CK_TILE_HOST_DEVICE
constexpr half_t bf16_to_fp16(bfloat16_t x) { return static_cast<fp16_t>(static_cast<float>(x)); }

template <class T>
struct numeric;

template <>
struct numeric<bfloat16_t>
{
    // minimum finite value, or minimum positive normalized value for float
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t min()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x0080));
    }

    // minumum finite value
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t lowest()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0xff7f));
    }

    // maximum finite value
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t max()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x7f7f));
    }

    // difference between 1.0 and next value representable by float
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t epsilon()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x1000));
    }

    // maximum rounding error
    // maximum rounding error
    // bin :  f edcba 9876543210
    // bits:  s eeeeeeee mmmmmmm
    //        0 01111110 0000000 (0.5)
    //
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t round_error()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x3f00));
    }

    // positive infinity value
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t infinity()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x7f80));
    }

    // quiet NaN
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t quiet_NaN()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x7FFF));
    }

    // signaling NaN
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t signaling_NaN()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x7FFF));
    }

    // smallest positive subnormal value
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t denorm_min()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0x0001));
    }
    CK_TILE_HOST_DEVICE static constexpr bfloat16_t zero()
    {
        return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(0));
    }
};

template <>
struct numeric_traits<bfloat16_t>
{
    static constexpr int exp        = 8;
    static constexpr int mant       = 7;
    static constexpr int PackedSize = 1;
};

#if CK_TILE_USE_CUSTOM_DATA_TYPE
CK_TILE_ARITHMETIC_USING_FLOAT(CK_TILE_HOST_DEVICE, bfloat16_t)
#endif

// math
CK_TILE_HOST_DEVICE
bfloat16_t abs(const bfloat16_t& x)
{
    return bit_cast<bfloat16_t>(static_cast<bf16_raw_t>(bit_cast<bf16_raw_t>(x) & 0x7fff));
}

CK_TILE_HOST_DEVICE
bool isnan(const bfloat16_t& x)
{
    uint16_t xx = bit_cast<bf16_raw_t>(x);
    return (xx & 0x7FFF) > 0x7C00;
}

CK_TILE_DEVICE
bfloat16_t sqrt(bfloat16_t x)
{
    return static_cast<bfloat16_t>(__builtin_amdgcn_sqrtf(static_cast<float>(x)));
};

CK_TILE_DEVICE
bfloat16_t exp(bfloat16_t x)
{
    return static_cast<bfloat16_t>(__ocml_exp_f32(static_cast<float>(x)));
};

CK_TILE_DEVICE
bfloat16_t exp2(bfloat16_t x) { return static_cast<bfloat16_t>(exp2f(static_cast<float>(x))); };

CK_TILE_DEVICE
bfloat16_t log(bfloat16_t x) { return static_cast<bfloat16_t>(__logf(static_cast<float>(x))); };

using bf16x2_t = bfloat16_t __attribute__((ext_vector_type(2)));
using fp32x2_t = float __attribute__((ext_vector_type(2)));

template <bf16_rounding_mode rounding =
              static_cast<bf16_rounding_mode>(CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT)>
CK_TILE_HOST_DEVICE constexpr bf16x2_t fp32x2_to_bf16x2(const fp32x2_t& x)
{
    return bf16x2_t{float_to_bf16<rounding>(x.x), float_to_bf16<rounding>(x.y)};
}

// Available on gfx94x (gfx942, gfx950) and later
CK_TILE_DEVICE bf16x2_t cvt_pk_bf16_f32(float a, float b)
{
#if defined(__gfx94__) && CK_TILE_USE_LLVM_BUILTIN_BF16
    return __builtin_convertvector(fp32x2_t{a, b}, bf16x2_t);
#else
    return fp32x2_to_bf16x2(fp32x2_t{a, b});
#endif
}

// Packed bf16x2 to fp32x2 conversion
CK_TILE_HOST_DEVICE constexpr fp32x2_t bf16x2_to_fp32x2(bf16x2_t x)
{
#if CK_TILE_USE_LLVM_BUILTIN_BF16
    return __builtin_convertvector(x, fp32x2_t);
#else
    uint32_t packed = bit_cast<uint32_t>(x);
    float f0        = bit_cast<float>(packed << 16);
    float f1        = bit_cast<float>(packed & 0xFFFF0000u);
    return fp32x2_t{f0, f1};
#endif
}

#ifndef CK_TILE_TF32_USE_PACKED_CVT
#define CK_TILE_TF32_USE_PACKED_CVT 1
#endif

template <int VecSize>
CK_TILE_DEVICE void convert_float_to_bf16_pairs(const ext_vector_t<float, VecSize>& reg_f32,
                                                ext_vector_t<bfloat16_t, VecSize>& reg_bf16_big,
                                                ext_vector_t<bfloat16_t, VecSize>& reg_bf16_small)
{
#if defined(__gfx94__) && CK_TILE_TF32_USE_PACKED_CVT && CK_TILE_USE_LLVM_BUILTIN_BF16
    static_assert(VecSize % 2 == 0, "VecSize must be even for packed operations");

#pragma unroll
    for(int i = 0; i < VecSize; i += 2)
    {
        fp32x2_t orig = {reg_f32[i], reg_f32[i + 1]};

        bf16x2_t big_pair   = cvt_pk_bf16_f32(orig[0], orig[1]);
        fp32x2_t big_f32    = bf16x2_to_fp32x2(big_pair);
        fp32x2_t diff       = orig - big_f32;
        bf16x2_t small_pair = cvt_pk_bf16_f32(diff[0], diff[1]);

        reinterpret_cast<bf16x2_t*>(&reg_bf16_big)[i / 2]   = big_pair;
        reinterpret_cast<bf16x2_t*>(&reg_bf16_small)[i / 2] = small_pair;
    }
#else
#pragma unroll
    for(int i = 0; i < VecSize; i++)
    {
        reg_bf16_big[i]   = float_to_bf16(reg_f32[i]);
        reg_bf16_small[i] = float_to_bf16(reg_f32[i] - bf16_to_float(reg_bf16_big[i]));
    }
#endif
}

} // namespace ck_tile
