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

#pragma once
#include <stdint.h>
#include "ck/utility/amd_ck_fp8.hpp"
#include "ck/utility/e8m0.hpp"
#include "ck/utility/statically_indexed_array.hpp"

/// Definitions from <cstdint>, <cmath> conflict with
/// /opt/rocm/include/hip/amd_detail/amd_hip_vector_types.h.

#if defined(__HIPCC_RTC__) || defined(CK_CODE_GEN_RTC)
#define CHAR_BIT 8
using int8_t   = signed char;
using uint8_t  = unsigned char;
using int16_t  = signed short;
using uint16_t = unsigned short;
using float_t  = float;
#endif // __HIPCC_RTC__

namespace ck {
#if defined(__HIPCC_RTC__) || defined(CK_CODE_GEN_RTC)
using byte = unsigned char;
#else
using std::byte;
#endif

using tf32_t  = _BitInt(19); // 1 sign bit, 8 exponent bits, 10 mantissa bits
using bhalf_t = ushort;
using half_t  = _Float16;
using int4_t  = _BitInt(4);
using f4_t    = unsigned _BitInt(4);
using f6_t    = _BitInt(6);          // e2m3 format
using bf6_t   = unsigned _BitInt(6); // e3m2 format

// native types: double, float, _Float16, ushort, int32_t, int8_t, uint8_t, f8_fnuz_t, bf8_fnuz_t,
// native types: bool
template <typename T>
inline constexpr bool is_native_type()
{
    return is_same_v<T, double> || is_same_v<T, float> || is_same_v<T, half_t> ||
           is_same_v<T, bhalf_t> || is_same_v<T, int32_t> || is_same_v<T, uint32_t> ||
           is_same_v<T, int8_t> || is_same_v<T, uint8_t> || is_same_v<T, _BitInt(8)> ||
           is_same_v<T, unsigned _BitInt(8)> || is_same_v<T, bool>;
}

/**
 * @brief Wrapper for native vector type
 * @tparam T The element type of the vector
 * @tparam Rank The number of elements in the vector
 */
template <typename T, index_t Rank>
using NativeVectorT = T __attribute__((ext_vector_type(Rank)));

/**
 * @brief Mapping of incoming type to local native vector storage type and vector size
 * @tparam T Incoming data type
 */
template <typename T>
struct scalar_type
{
    // Basic data type mapping to unsigned _BitInt of appropriate size
    using type                           = unsigned _BitInt(8 * sizeof(T));
    static constexpr index_t vector_size = 1;
};

/**
 * @brief scalar_type trait override for NativeVectorT
 * @tparam T The vector type
 * @tparam Rank The number of elements in the vector
 */
template <typename T, index_t Rank>
struct scalar_type<NativeVectorT<T, Rank>>
{
    using type                           = T;
    static constexpr index_t vector_size = Rank;
};

struct f4x2_pk_t
{
    static constexpr int packed_size = 2;

    using type = uint8_t;
    type data;
    __host__ __device__ constexpr f4x2_pk_t() : data{type{}} {}
    __host__ __device__ constexpr f4x2_pk_t(const type init) : data{init} {}

    template <index_t I>
    __host__ __device__ inline type unpack(Number<I>) const
    {
        static_assert(I < 2, "Index is out of range.");
        if constexpr(I == 1)
            return (data >> 4);
        else
            return data & 0b00001111;
    }

    __host__ __device__ inline type pack(const type x0, const type x1)
    {
        return (x1 << 4) | (x0 & 0b00001111);
    }

    // Compare operator
    __host__ __device__ friend bool operator==(const f4x2_pk_t& lhs, const f4x2_pk_t& rhs)
    {
        return lhs.data == rhs.data;
    }

    __host__ __device__ friend bool operator!=(const f4x2_pk_t& lhs, const f4x2_pk_t& rhs)
    {
        return !(lhs == rhs);
    }
};

// TODO: Unfortunately, we cannot partially specialize scalar_type for vectors written
// in the following way:
// template<typename T, index_t Rank>
// struct scalar_type<T __attribute__((__vector_size__(sizeof(T) * Rank)))>
// {
//     using type                           = T;
//     static constexpr index_t vector_size = Rank;
// };
// The compiler errors out with "partial specialization is not allowed for this type",
// claiming that the Rank is not a deducible parameter. This might be a compiler bug.
// Note the above type is classified differently from the NativeVectorT<T, Rank> alias,
// even though they are functionally equivalent and are trivially constructibe from each other.
// This is unfortunate, but we have to work around it because some LLVM builtins for some
// operations (e.g., mma) may return the former type.
// For now we have to explicitly specialize for each vector size we need. These are used
// in f6_pk_t below.

/// @brief scalar_type trait override for uint32_t vector of size 3
template <>
struct scalar_type<uint32_t __attribute__((__vector_size__(sizeof(uint32_t) * 3)))>
{
    using type                           = uint32_t;
    static constexpr index_t vector_size = 3;
};

/// @brief scalar_type trait override for uint32_t vector of size 6
template <>
struct scalar_type<uint32_t __attribute__((__vector_size__(sizeof(uint32_t) * 6)))>
{
    using type                           = uint32_t;
    static constexpr index_t vector_size = 6;
};

template <typename BitType, index_t pk_size>
struct f6_pk_t
{
    using element_type = uint32_t; // element storage fundamental type

    static constexpr index_t packed_size   = pk_size; // 16 or 32 for now
    static constexpr index_t num_bits_elem = 6;       // specialized for 6-bit data
    // XXX: CHAR_BIT is not defined in HIPRTC, so we must use 8
    static constexpr index_t num_bits_vec_elem =
        sizeof(element_type) * 8; // 32-bit uint for storage
    static_assert((packed_size * num_bits_elem) % num_bits_vec_elem == 0,
                  "Packed elements must fit exactly into the element storage.");
    static constexpr index_t vector_size =
        (packed_size * num_bits_elem) / num_bits_vec_elem; // 3 or 6 element_type units

    using storage_type = NativeVectorT<element_type, vector_size>;
    storage_type data_{storage_type(0)}; // packed data

    using type = f6_pk_t<BitType, packed_size>;

    /** This class may trivially constructed by the following vector type alias
     * for example from a result of an mma operation. This is primarily for internal use.
     * @note f6x16_pk_t and f6x32_pk_t storage types, may be trivially constructed from
     *       uint32_t vectors of size 3 and 6 respectively for example from mma operation results.
     *       Unfortunately, unsigned int __attribute__((ext_vector_type(6))) a.k.a
     * NativeVectorT<uint32_t, 6> is NOT the same as __attribute__((__vector_size__(6 *
     * sizeof(unsigned int)))) unsigned int which is returned from the mma ops despite being
     * functionally equivalent. This class may be trivially constructed from both, so we can steer
     * the templated ctor below to only consider incoming vectors types other than our two storage
     * types of interest.
     */
    using storage_type_alias =
        element_type __attribute__((__vector_size__(sizeof(element_type) * vector_size)));

    __host__ __device__ constexpr f6_pk_t() {}
    __host__ __device__ constexpr f6_pk_t(const storage_type& init) : data_{init}
    {
        // TODO: consider removing initialization similar to vector_type<T, 256>
    }

    // Initialize from a vector type with the same size as packed_size.
    // Exclude storage_type and storage_type_alias because these are trivially constructible.
    template <
        typename T,
        typename = enable_if_t<!is_same_v<T, storage_type> && !is_same_v<T, storage_type_alias> &&
                               scalar_type<T>::vector_size == packed_size>>
    __host__ __device__ f6_pk_t(const T& v)
    {
        static_assert(scalar_type<T>::vector_size == packed_size,
                      "Input vector size must match packed_size.");
        static_for<0, packed_size, 1>{}(
            [&](auto i) { pack(v[static_cast<index_t>(i)], static_cast<index_t>(i)); });
    }

    // Broadcast single initialization value to all packed elements
    __host__ __device__ f6_pk_t(const int8_t v)
        : f6_pk_t(static_cast<NativeVectorT<int8_t, packed_size>>(v))
    {
        // TODO: consider removing initialization similar to vector_type<T, 256>
    }

    template <typename T>
    __host__ __device__ void pack(const T x, const index_t i)
    {
        static_assert(is_integral<T>::value || is_same_v<T, BitType>,
                      "T must be an integral type.");

        uint32_t bits        = static_cast<uint32_t>(x) & 0x3F;
        const int bit_pos    = i * num_bits_elem;
        const int arr_index  = bit_pos / num_bits_vec_elem;
        const int bit_offset = bit_pos % num_bits_vec_elem;
        const int overhang   = bit_offset + num_bits_elem - num_bits_vec_elem;
        uint32_t old_value   = data_[arr_index];

        // insert bits into the current 32-bit block
        old_value |= (bits << bit_offset);
        data_[arr_index] = old_value;

        // if it crosses into the next block, shift the remainder
        if(overhang > 0 && (arr_index + 1) < vector_size)
        {
            uint32_t next_value = data_[arr_index + 1];
            next_value |= (bits >> (num_bits_elem - overhang));
            data_[arr_index + 1] = next_value;
        }
    }

    __host__ __device__ static inline BitType unpack(const type& pk, const index_t i)
    {
        const int bit_pos    = i * num_bits_elem;
        const int arr_idx    = bit_pos / num_bits_vec_elem;
        const int bit_offset = bit_pos % num_bits_vec_elem;
        const int overhang   = bit_offset + num_bits_elem - num_bits_vec_elem;

        uint32_t bits = pk.data_[arr_idx] >> bit_offset;
        if(overhang > 0 && (arr_idx + 1) < vector_size)
        {
            bits |= (pk.data_[arr_idx + 1] & ((1u << overhang) - 1)) << (num_bits_elem - overhang);
        }

        return static_cast<BitType>(bits & 0x3F);
    }

    __host__ __device__ inline BitType unpack(const index_t i) const { return unpack(*this, i); }

    // Compare operator
    __host__ __device__ friend bool operator==(const f6_pk_t& lhs, const f6_pk_t& rhs)
    {
#pragma unroll
        for(index_t i = 0; i < vector_size; ++i)
        {
            if(lhs.data_[i] != rhs.data_[i])
                return false;
        }
        return true;
    }

    __host__ __device__ friend bool operator!=(const f6_pk_t& lhs, const f6_pk_t& rhs)
    {
        return !(lhs == rhs);
    }
};

using f6x16_pk_t  = f6_pk_t<f6_t, 16>;
using f6x32_pk_t  = f6_pk_t<f6_t, 32>;
using bf6x16_pk_t = f6_pk_t<bf6_t, 16>;
using bf6x32_pk_t = f6_pk_t<bf6_t, 32>;

// custom data type - pack int4 data
struct pk_i4_t
{
    using type = int8_t;
    type data;
    __host__ __device__ constexpr pk_i4_t() : data{type{}} {}
    __host__ __device__ constexpr pk_i4_t(type init) : data{init} {}
};

// is_scalar_type
template <typename TV>
struct is_scalar_type
{
    static constexpr bool value = (scalar_type<remove_cvref_t<TV>>::vector_size == 1);
};

// has_same_scalar_type
template <typename X, typename Y>
using has_same_scalar_type = is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                     typename scalar_type<remove_cvref_t<Y>>::type>;

template <>
struct scalar_type<bool>
{
    using type                           = bool;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<double>
{
    using type                           = double;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<float>
{
    using type                           = float;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<half_t>
{
    using type                           = half_t;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<bhalf_t>
{
    using type                           = bhalf_t;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<int32_t>
{
    using type                           = int32_t;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<int8_t>
{
    using type                           = int8_t;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<uint8_t>
{
    using type                           = uint8_t;
    static constexpr index_t vector_size = 1;
};

#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
template <>
struct scalar_type<int4_t>
{
    using type                           = int4_t;
    static constexpr index_t vector_size = 1;
};
#endif

template <>
struct scalar_type<pk_i4_t>
{
    using type                           = typename pk_i4_t::type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<f8_fnuz_t>
{
    using type                           = typename f8_fnuz_t::data_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<bf8_fnuz_t>
{
    using type                           = typename bf8_fnuz_t::data_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<f8_ocp_t>
{
    using type                           = typename f8_ocp_t::data_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<bf8_ocp_t>
{
    using type                           = typename bf8_ocp_t::data_type;
    static constexpr index_t vector_size = 1;
};

#ifndef CK_CODE_GEN_RTC
template <>
struct scalar_type<e8m0_bexp_t>
{
    using type                           = typename e8m0_bexp_t::type;
    static constexpr index_t vector_size = 1;
};
#endif

template <>
struct scalar_type<f4x2_pk_t>
{
    using type                           = typename f4x2_pk_t::type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<f6x32_pk_t>
{
    using type                           = typename f6x32_pk_t::storage_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<bf6x32_pk_t>
{
    using type                           = typename bf6x32_pk_t::storage_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<f6x16_pk_t>
{
    using type                           = typename f6x16_pk_t::storage_type;
    static constexpr index_t vector_size = 1;
};

template <>
struct scalar_type<bf6x16_pk_t>
{
    using type                           = typename bf6x16_pk_t::storage_type;
    static constexpr index_t vector_size = 1;
};

template <typename T>
struct packed_type_info
{
    private:
    static constexpr auto get_packed_type_info()
    {
        using U = remove_cvref_t<T>;
        if constexpr(is_same_v<U, pk_i4_t>)
            return ck::Tuple<ck::Number<2>, int4_t>{};
        else if constexpr(is_same_v<U, f4x2_pk_t>)
            return ck::Tuple<ck::Number<2>, f4_t>{};
        else if constexpr(is_same_v<U, f6x16_pk_t>)
            return ck::Tuple<ck::Number<16>, f6_t>{};
        else if constexpr(is_same_v<U, bf6x16_pk_t>)
            return ck::Tuple<ck::Number<16>, bf6_t>{};
        else if constexpr(is_same_v<U, f6x32_pk_t>)
            return ck::Tuple<ck::Number<32>, f6_t>{};
        else if constexpr(is_same_v<U, bf6x32_pk_t>)
            return ck::Tuple<ck::Number<32>, bf6_t>{};
        else
            return ck::Tuple<ck::Number<1>, T>{};
    }

    public:
    using element_type = remove_cvref_t<decltype(get_packed_type_info().At(ck::Number<1>{}))>;
    static constexpr auto packed_size =
        static_cast<index_t>(get_packed_type_info().At(ck::Number<0>{}));
};
template <typename T>
using element_type_t = typename packed_type_info<T>::element_type;

template <typename T>
inline constexpr index_t packed_size_v = packed_type_info<T>::packed_size;

template <typename T>
inline constexpr bool is_packed_type_v = packed_size_v<T> > 1;

template <typename T, index_t N = 0>
struct packed_type_maker
{
    private:
    static constexpr auto get_packed_type()
    {
        using U = remove_cvref_t<T>;
        if constexpr(is_same_v<U, int4_t>)
        {
            static_assert(N == 0 || N == 2, "Packed size N for int4_t must be 2.");
            return pk_i4_t{};
        }
        else if constexpr(is_same_v<U, f4_t>)
        {
            static_assert(N == 0 || N == 2, "Packed size N for f4_t must be 2.");
            return f4x2_pk_t{};
        }
        else if constexpr(is_same_v<U, f6_t>)
        {
            static_assert(N == 0 || N == 16 || N == 32, "Packed size N for f6_t must be 16 or 32.");
            if constexpr(N == 16)
                return f6x16_pk_t{};
            else if constexpr(N == 0 || N == 32)
                return f6x32_pk_t{};
        }
        else if constexpr(is_same_v<U, bf6_t>)
        {
            static_assert(N == 0 || N == 16 || N == 32,
                          "Packed size N for bf6_t must be 16 or 32.");
            if constexpr(N == 16)
                return bf6x16_pk_t{};
            else if constexpr(N == 0 || N == 32)
                return bf6x32_pk_t{};
        }
        else
            return T{};
    }

    public:
    using packed_type = remove_cvref_t<decltype(get_packed_type())>;
};

template <typename T, index_t N = 0>
using packed_type_t = typename packed_type_maker<T, N>::packed_type;

#if defined(_WIN32)
using int64_t = long long;
#else
using int64_t = long;
#endif

template <typename T>
inline const char* get_type_name()
{
    if constexpr(is_same_v<T, half_t>)
        return "fp16";
    else if constexpr(is_same_v<T, bhalf_t>)
        return "bf16";
    else if constexpr(is_same_v<T, tf32_t>)
        return "tf32";
    else if constexpr(is_same_v<T, int4_t>)
        return "int4";
    else if constexpr(is_same_v<T, f4_t>)
        return "f4";
    else if constexpr(is_same_v<T, f6_t>)
        return "f6";
    else if constexpr(is_same_v<T, bf6_t>)
        return "bf6";
    else if constexpr(is_same_v<T, f8_t>)
        return "f8";
    else if constexpr(is_same_v<T, bf8_t>)
        return "bf8";
#ifndef CK_CODE_GEN_RTC
    else if constexpr(is_same_v<T, e8m0_bexp_t>)
        return "e8m0";
#endif
    else if constexpr(is_same_v<T, float>)
        return "fp32";
#if defined(__HIPCC_RTC__) || defined(CK_CODE_GEN_RTC)
    else
        return "unknown";
#else
    else
        return typeid(T).name();
#endif
}

} // namespace ck
