// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause

#ifndef GKO_HIP_COMPONENTS_COOPERATIVE_GROUPS_HIP_HPP_
#define GKO_HIP_COMPONENTS_COOPERATIVE_GROUPS_HIP_HPP_


#include <type_traits>

#include "common/cuda_hip/base/config.hpp"
#include "common/cuda_hip/base/types.hpp"


namespace gko {
namespace kernels {
namespace hip {


/**
 * Ginkgo uses cooperative groups to handle communication among the threads.
 *
 * However, HIP's implementation of cooperative groups is still quite limited
 * in functionality, and some parts are not supported on all hardware
 * interesting for Ginkgo. For this reason, Ginkgo exposes only a part of the
 * original functionality, and possibly extends it if it is required. Thus,
 * developers should include and use this header and the gko::group namespace
 * instead of the standard cooperative_groups.h header. The interface exposed
 * by Ginkgo's implementation is equivalent to the standard interface, with some
 * useful extensions.
 *
 * A cooperative group (both from standard HIP and from Ginkgo) is not a
 * specific type, but a concept. That is, any type satisfying the interface
 * imposed by the cooperative groups API is considered a cooperative
 * group (a.k.a. "duck typing"). To maximize the generality of components that
 * need cooperative groups, instead of creating the group manually, consider
 * requesting one as an input parameter. Make sure its type is a template
 * parameter to maximize the set of groups for which your algorithm can be
 * invoked. To maximize the amount of contexts in which your algorithm can be
 * called and avoid hidden requirements, do not depend on a specific setup of
 * kernel launch parameters (i.e. grid dimensions and block dimensions).
 * Instead, use the thread_rank() method of the group to distinguish between
 * distinct threads of a group.
 *
 * The original HIP implementation does not provide ways to verify if a certain
 * type represents a cooperative group. Ginkgo's implementation provides
 * metafunctions which do that. Additionally, not all cooperative groups have
 * equivalent functionality, so Ginkgo splits the cooperative group concept into
 * three sub-concepts which describe what functionality is available. Here is a
 * list of concepts and their interfaces:
 *
 * ```c++
 * concept Group {
 *   unsigned size() const;
 *   unsigned thread_rank() const;
 * };
 *
 * concept SynchronizableGroup : Group {
 *   void sync();
 * };
 *
 * concept CommunicatorGroup : SynchronizableGroup {
 *   template <typename T>
 *   T shfl(T var, int srcLane);
 *   T shfl_up(T var, unsigned delta);
 *   T shfl_down(T var, unsigned delta);
 *   T shfl_xor(T var, int laneMask);
 *   int all(int predicate);
 *   int any(int predicate);
 *   unsigned ballot(int predicate);
 * };
 * ```
 *
 * To check if a group T satisfies one of the concepts, one can use the
 * metafunctions is_group<T>::value, is_synchronizable_group<T>::value and
 * is_communicator_group<T>::value.
 *
 * @note Please note that the current implementation of cooperative groups
 *       contains only a subset of functionalities provided by those APIs. If
 *       you need more functionality, please add the appropriate implementations
 *       to existing cooperative groups, or create new groups if the existing
 *       groups do not cover your use-case. For an example, see the
 *       enable_extended_shuffle mixin, which adds extended shuffles support
 *       to built-in HIP cooperative groups.
 */
namespace group {


// metafunctions


namespace detail {


template <typename T>
struct is_group_impl : std::false_type {};


template <typename T>
struct is_synchronizable_group_impl : std::false_type {};


template <typename T>
struct is_communicator_group_impl : std::false_type {};

}  // namespace detail


/**
 * Check if T is a Group.
 */
template <typename T>
using is_group = detail::is_group_impl<std::decay_t<T>>;


/**
 * Check if T is a SynchronizableGroup.
 */
template <typename T>
using is_synchronizable_group =
    detail::is_synchronizable_group_impl<std::decay_t<T>>;


/**
 * Check if T is a CommunicatorGroup.
 */
template <typename T>
using is_communicator_group =
    detail::is_communicator_group_impl<std::decay_t<T>>;


// types


namespace detail {


/**
 * This is a limited implementation of the HIP thread_block_tile.
 * `any` and `all` are only supported when the size is config::warp_size
 *
 */
template <unsigned Size>
class thread_block_tile {
    /**
     * Mask with Size consecutive ones starting at the least significant bit.
     */
    static constexpr auto lane_mask_base = ~config::lane_mask_type{} >>
                                           (config::warp_size - Size);

public:
    __device__ thread_block_tile() : data_{Size, 0, 0, lane_mask_base}
    {
        auto tid =
            unsigned(threadIdx.x +
                     blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z));
        data_.rank = tid % Size;
        data_.lane_offset = (tid % config::warp_size) / Size * Size;
        data_.mask <<= data_.lane_offset;
    }

    __device__ __forceinline__ unsigned thread_rank() const noexcept
    {
        return data_.rank;
    }

    __device__ __forceinline__ unsigned size() const noexcept { return Size; }

    __device__ __forceinline__ void sync() const noexcept
    {
#if GINKGO_HIP_PLATFORM_NVCC
        __syncwarp(data_.mask);
#endif  // GINKGO_HIP_PLATFORM_NVCC
    }

#if GINKGO_HIP_PLATFORM_HCC
#define GKO_BIND_SHFL(ShflOp, ValueType, SelectorType)                       \
    __device__ __forceinline__ ValueType ShflOp(                             \
        ValueType var, SelectorType selector) const noexcept                 \
    {                                                                        \
        return __##ShflOp(var, selector, Size);                              \
    }                                                                        \
    static_assert(true,                                                      \
                  "This assert is used to counter the false positive extra " \
                  "semi-colon warnings")
#else
#define GKO_BIND_SHFL(ShflOp, ValueType, SelectorType)                       \
    __device__ __forceinline__ ValueType ShflOp(                             \
        ValueType var, SelectorType selector) const noexcept                 \
    {                                                                        \
        return __##ShflOp##_sync(data_.mask, var, selector, Size);           \
    }                                                                        \
    static_assert(true,                                                      \
                  "This assert is used to counter the false positive extra " \
                  "semi-colon warnings")
#endif

    GKO_BIND_SHFL(shfl, int32, int32);
    GKO_BIND_SHFL(shfl, float, int32);
    GKO_BIND_SHFL(shfl, uint32, int32);
    GKO_BIND_SHFL(shfl, double, int32);

    GKO_BIND_SHFL(shfl_up, int32, uint32);
    GKO_BIND_SHFL(shfl_up, uint32, uint32);
    GKO_BIND_SHFL(shfl_up, float, uint32);
    GKO_BIND_SHFL(shfl_up, double, uint32);

    GKO_BIND_SHFL(shfl_down, int32, uint32);
    GKO_BIND_SHFL(shfl_down, uint32, uint32);
    GKO_BIND_SHFL(shfl_down, float, uint32);
    GKO_BIND_SHFL(shfl_down, double, uint32);

    GKO_BIND_SHFL(shfl_xor, int32, int32);
    GKO_BIND_SHFL(shfl_xor, float, int32);
    GKO_BIND_SHFL(shfl_xor, uint32, int32);
    GKO_BIND_SHFL(shfl_xor, double, int32);

    /**
     * Returns true iff the predicate is true for at least one threads in the
     * group. Note that the whole group needs to execute the same operation.
     */
    __device__ __forceinline__ int any(int predicate) const noexcept
    {
#if GINKGO_HIP_PLATFORM_HCC
        if (Size == config::warp_size) {
            return __any(predicate);
        } else {
            return (__ballot(predicate) & data_.mask) != 0;
        }
#else
        return __any_sync(data_.mask, predicate);
#endif
    }

    /**
     * Returns true iff the predicate is true for all threads in the group.
     * Note that the whole group needs to execute the same operation.
     */
    __device__ __forceinline__ int all(int predicate) const noexcept
    {
#if GINKGO_HIP_PLATFORM_HCC
        if (Size == config::warp_size) {
            return __all(predicate);
        } else {
            return (__ballot(predicate) & data_.mask) == data_.mask;
        }
#else
        return __all_sync(data_.mask, predicate);
#endif
    }

    /**
     * Returns a bitmask containing the value of the given predicate
     * for all threads in the group.
     * This means that the ith bit is equal to the predicate of the
     * thread with thread_rank() == i in the group.
     * Note that the whole group needs to execute the same operation.
     */
    __device__ __forceinline__ config::lane_mask_type ballot(
        int predicate) const noexcept
    {
#if GINKGO_HIP_PLATFORM_HCC
        if (Size == config::warp_size) {
            return __ballot(predicate);
        } else {
            return (__ballot(predicate) & data_.mask) >> data_.lane_offset;
        }
#else
        if (Size == config::warp_size) {
            return __ballot_sync(data_.mask, predicate);
        } else {
            return __ballot_sync(data_.mask, predicate) >> data_.lane_offset;
        }
#endif
    }

private:
    struct alignas(8) {
        unsigned size;
        unsigned rank;
        unsigned lane_offset;
        config::lane_mask_type mask;
    } data_;
};


}  // namespace detail


namespace detail {


// Adds generalized shuffles that support any type to the group.
template <typename Group>
class enable_extended_shuffle : public Group {
public:
    using Group::Group;
    using Group::shfl;
    using Group::shfl_down;
    using Group::shfl_up;
    using Group::shfl_xor;

#define GKO_ENABLE_SHUFFLE_OPERATION(_name, SelectorType)                   \
    template <typename ValueType>                                           \
    __device__ __forceinline__ ValueType _name(const ValueType& var,        \
                                               SelectorType selector) const \
    {                                                                       \
        return shuffle_impl(                                                \
            [this](uint32 v, SelectorType s) {                              \
                return static_cast<const Group*>(this)->_name(v, s);        \
            },                                                              \
            var, selector);                                                 \
    }

    GKO_ENABLE_SHUFFLE_OPERATION(shfl, int32)
    GKO_ENABLE_SHUFFLE_OPERATION(shfl_up, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION(shfl_down, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION(shfl_xor, int32)

#undef GKO_ENABLE_SHUFFLE_OPERATION

// hip does not support 16bit shuffle directly
#define GKO_ENABLE_SHUFFLE_OPERATION_HALF(_name, SelectorType)        \
    __device__ __forceinline__ device_type<float16> _name(            \
        const device_type<float16>& var, SelectorType selector) const \
    {                                                                 \
        uint32 u;                                                     \
        memcpy(&u, &var, sizeof(device_type<float16>));               \
        u = static_cast<const Group*>(this)->_name(u, selector);      \
        device_type<float16> result;                                  \
        memcpy(&result, &u, sizeof(device_type<float16>));            \
        return result;                                                \
    }

    GKO_ENABLE_SHUFFLE_OPERATION_HALF(shfl, int32)
    GKO_ENABLE_SHUFFLE_OPERATION_HALF(shfl_up, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION_HALF(shfl_down, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION_HALF(shfl_xor, int32)

#undef GKO_ENABLE_SHUFFLE_OPERATION_HALF

// hip does not support 16bit shuffle directly
#define GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16(_name, SelectorType)     \
    __device__ __forceinline__ device_type<bfloat16> _name(            \
        const device_type<bfloat16>& var, SelectorType selector) const \
    {                                                                  \
        uint32 u;                                                      \
        memcpy(&u, &var, sizeof(device_type<bfloat16>));               \
        u = static_cast<const Group*>(this)->_name(u, selector);       \
        device_type<bfloat16> result;                                  \
        memcpy(&result, &u, sizeof(device_type<bfloat16>));            \
        return result;                                                 \
    }

    GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16(shfl, int32)
    GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16(shfl_up, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16(shfl_down, uint32)
    GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16(shfl_xor, int32)

#undef GKO_ENABLE_SHUFFLE_OPERATION_BFLOAT16


private:
    template <typename ShuffleOperator, typename ValueType,
              typename SelectorType>
    static __device__ __forceinline__ ValueType
    shuffle_impl(ShuffleOperator intrinsic_shuffle, const ValueType var,
                 SelectorType selector)
    {
        static_assert(sizeof(ValueType) % sizeof(uint32) == 0,
                      "Unable to shuffle sizes which are not 4-byte multiples");
        constexpr auto value_size = sizeof(ValueType) / sizeof(uint32);
        ValueType result;
        auto var_array = reinterpret_cast<const uint32*>(&var);
        auto result_array = reinterpret_cast<uint32*>(&result);
#pragma unroll
        for (std::size_t i = 0; i < value_size; ++i) {
            result_array[i] = intrinsic_shuffle(var_array[i], selector);
        }
        return result;
    }
};


}  // namespace detail


// Implementing this as a using directive messes up with SFINAE for some reason,
// probably a bug in NVCC. If it is a complete type, everything works fine.
template <unsigned Size>
struct thread_block_tile
    : detail::enable_extended_shuffle<detail::thread_block_tile<Size>> {
    using detail::enable_extended_shuffle<
        detail::thread_block_tile<Size>>::enable_extended_shuffle;
};


// Only support tile_partition with 1, 2, 4, 8, 16, 32, 64 (hip).
template <size_type Size, typename Group>
__device__ __forceinline__
    std::enable_if_t<(Size <= kernels::hip::config::warp_size) && (Size > 0) &&
                         (kernels::hip::config::warp_size % Size == 0),
                     thread_block_tile<Size>>
    tiled_partition(const Group&)
{
    return thread_block_tile<Size>();
}


namespace detail {


template <unsigned Size>
struct is_group_impl<group::thread_block_tile<Size>> : std::true_type {};
template <unsigned Size>
struct is_synchronizable_group_impl<group::thread_block_tile<Size>>
    : std::true_type {};
template <unsigned Size>
struct is_communicator_group_impl<group::thread_block_tile<Size>>
    : std::true_type {};


}  // namespace detail


class thread_block {
    friend __device__ __forceinline__ thread_block this_thread_block();

public:
    __device__ __forceinline__ unsigned thread_rank() const noexcept
    {
        return data_.rank;
    }

    __device__ __forceinline__ unsigned size() const noexcept
    {
        return data_.size;
    }

    __device__ __forceinline__ void sync() const noexcept { __syncthreads(); }

private:
    __device__ thread_block()
        : data_{static_cast<unsigned>(blockDim.x * blockDim.y * blockDim.z),
                static_cast<unsigned>(
                    threadIdx.x +
                    blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z))}
    {}
    struct alignas(8) {
        unsigned size;
        unsigned rank;
    } data_;
};


__device__ __forceinline__ thread_block this_thread_block()
{
    return thread_block();
}


namespace detail {

template <>
struct is_group_impl<thread_block> : std::true_type {};
template <>
struct is_synchronizable_group_impl<thread_block> : std::true_type {};


}  // namespace detail


/**
 * This is a limited implementation of the CUDA grid_group that works even on
 * devices that do not support device-wide synchronization and without special
 * kernel launch syntax.
 *
 * Note that this implementation (as well as the one from CUDA's cooperative
 * groups) does not support large grids, since it uses 32 bits to represent
 * sizes and ranks, while at least 73 bits (63 bit grid + 10 bit block) would
 * have to be used to represent the full space of thread ranks.
 */
class grid_group {
    friend __device__ grid_group this_grid();

public:
    __device__ unsigned size() const noexcept { return data_.size; }

    __device__ unsigned thread_rank() const noexcept { return data_.rank; }

private:
    // clang-format off
    __device__ grid_group()
        : data_{
                blockDim.x * blockDim.y * blockDim.z *
                    gridDim.x * gridDim.y * gridDim.z,
                threadIdx.x + blockDim.x *
                    (threadIdx.y + blockDim.y *
                        (threadIdx.z + blockDim.z *
                            (blockIdx.x + gridDim.x *
                                (blockIdx.y + gridDim.y * blockIdx.z))))}
    {}
    // clang-format on

    struct alignas(8) {
        unsigned size;
        unsigned rank;
    } data_;
};

// Not using this, as grid_group is not universally supported.
// grid_group this_grid()
// using cooperative_groups::this_grid;
// Instead, use our limited implementation:
__device__ inline grid_group this_grid() { return {}; }


/**
 * During cuda12.2 - cuda12.4, compiler might have some optimization issue on
 * ballot on H100. The ballot result on subwarp might keep the result from the
 * other subwarp in the same warp. For example with subwarp size = 8, subwarp 0
 * will get the result from (subwarp 3, subwarp 2, subwarp 1, subwarp 0),
 * subwarp 1 will get the result from (subwarp 3, subwarp 2, subwarp 1), and so
 * on so forth. They are shifted but not masked. cuda 12.1.1 and cuda 12.5.1
 * does not have the issue.
 */
template <unsigned Size>
__device__ __forceinline__ auto ballot(const thread_block_tile<Size>& g,
                                       int predicate)
{
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12020 && CUDA_VERSION < 12051
    constexpr auto subwarp_mask =
        ~config::lane_mask_type{} >> (config::warp_size - Size);
    return subwarp_mask & g.ballot(predicate);
#else
    return g.ballot(predicate);
#endif
}


}  // namespace group
}  // namespace hip
}  // namespace kernels
}  // namespace gko


#endif  // GKO_HIP_COMPONENTS_COOPERATIVE_GROUPS_HIP_HPP_
