/**
 * Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See file LICENSE for terms.
 */

#ifdef __cplusplus
extern "C" {
#endif

#include "utils/arch/cuda_def.h"
#include "../tl_cuda.h"

#ifdef __cplusplus
}
#endif

#include "nvls.cuh"

// vectorized allreduce kernel for 32-bit lanes
template <typename NvlsOps>
__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
    allreduce_kernel_vec32(
        ucc_tl_cuda_nvls_control_t *mc_bar, ucc_tl_cuda_nvls_control_t *uc_bar,
        const uint32_t total_blocks, uint64_t launch_counter,
        uint32_t *base_u32, size_t count_u32, uint32_t rank, uint32_t tsize)
{
    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 4;
    size_t stride        = blockDim.x * gridDim.x * 4;
    size_t chunk_start;
    size_t chunk_end;
    uint4  val;

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 1));

    if (count_u32 >= 4 * tsize) {
        chunk_start = ((int64_t)count_u32 * (int64_t)rank) / (int64_t)tsize;
        chunk_end = ((int64_t)count_u32 * (int64_t)(rank + 1)) / (int64_t)tsize;
    } else {
        chunk_start = 0;
        chunk_end   = (rank == 0) ? count_u32 : 0;
    }

    for (size_t idx = chunk_start + thread_offset; idx < chunk_end;
         idx += stride) {
        NvlsOps::ld(val, base_u32 + idx);
        NvlsOps::st(val, base_u32 + idx);
    }

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 2));
}

template <typename NvlsOps>
__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
    allreduce_kernel_scalar32(
        ucc_tl_cuda_nvls_control_t *mc_bar, ucc_tl_cuda_nvls_control_t *uc_bar,
        const uint32_t total_blocks, uint64_t launch_counter,
        uint32_t *base_u32, size_t count_u32, uint32_t rank, uint32_t tsize)
{
    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 4;
    size_t stride        = blockDim.x * gridDim.x * 4;
    size_t chunk_start;
    size_t chunk_end;
    typename NvlsOps::value_type v0, v1, v2, v3;

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 1));

    if (count_u32 >= 4 * tsize) {
        chunk_start = ((int64_t)count_u32 * (int64_t)rank) / (int64_t)tsize;
        chunk_end = ((int64_t)count_u32 * (int64_t)(rank + 1)) / (int64_t)tsize;
    } else {
        chunk_start = 0;
        chunk_end   = (rank == 0) ? count_u32 : 0;
    }

    for (size_t idx = chunk_start + thread_offset; idx < chunk_end;
         idx += stride) {
        NvlsOps::ld(v0, base_u32 + idx + 0);
        NvlsOps::ld(v1, base_u32 + idx + 1);
        NvlsOps::ld(v2, base_u32 + idx + 2);
        NvlsOps::ld(v3, base_u32 + idx + 3);
        NvlsOps::st(v0, base_u32 + idx + 0);
        NvlsOps::st(v1, base_u32 + idx + 1);
        NvlsOps::st(v2, base_u32 + idx + 2);
        NvlsOps::st(v3, base_u32 + idx + 3);
    }

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 2));
}

template <typename NvlsOps>
__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
    allreduce_kernel_scalar64(
        ucc_tl_cuda_nvls_control_t *mc_bar, ucc_tl_cuda_nvls_control_t *uc_bar,
        const uint32_t total_blocks, uint64_t launch_counter,
        uint64_t *base_u64, size_t count_u64, uint32_t rank, uint32_t tsize)
{
    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 2;
    size_t stride        = blockDim.x * gridDim.x * 2;
    size_t chunk_start;
    size_t chunk_end;
    typename NvlsOps::value_type v0, v1;

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 1));

    if (count_u64 >= 2 * tsize) {
        chunk_start = ((int64_t)count_u64 * (int64_t)rank) / (int64_t)tsize;
        chunk_end = ((int64_t)count_u64 * (int64_t)(rank + 1)) / (int64_t)tsize;
    } else {
        chunk_start = 0;
        chunk_end   = (rank == 0) ? count_u64 : 0;
    }

    for (size_t idx = chunk_start + thread_offset; idx < chunk_end;
         idx += stride) {
        NvlsOps::ld(v0, base_u64 + idx + 0);
        NvlsOps::ld(v1, base_u64 + idx + 1);
        NvlsOps::st(v0, base_u64 + idx + 0);
        NvlsOps::st(v1, base_u64 + idx + 1);
    }

    nvls_bar(
        &(mc_bar->arrival_counter),
        &(uc_bar->arrival_counter),
        total_blocks * (launch_counter * 2 + 2));
}

#ifdef __cplusplus
extern "C" {
#endif

ucc_status_t post_allreduce_kernel(cudaStream_t stream, uint32_t sm_count,
                                   uint32_t threads, CUdeviceptr mc_base_addr,
                                   size_t src_size_bytes,
                                   CUdeviceptr mc_control_addr,
                                   CUdeviceptr uc_control_addr,
                                   uint64_t launch_counter,
                                   uint32_t rank,
                                   uint32_t tsize, ucc_datatype_t datatype)
{
    assert(sm_count > 0 && sm_count <= UCC_TL_CUDA_MAX_NVLS_SM_COUNT);
    assert(threads > 0 && threads <= UCC_TL_CUDA_MAX_NVLS_THREADS);
    uint32_t *base_u32   = reinterpret_cast<uint32_t *>(mc_base_addr);
    size_t    count_u32  = src_size_bytes / sizeof(uint32_t);
    ucc_tl_cuda_nvls_control_t *mc_bar = reinterpret_cast<ucc_tl_cuda_nvls_control_t *>(mc_control_addr);
    ucc_tl_cuda_nvls_control_t *uc_bar = reinterpret_cast<ucc_tl_cuda_nvls_control_t *>(uc_control_addr);
    uint32_t expected_blocks = sm_count * tsize; // total num of blocks in the multicast group, num gpus * num blocks per gpu, used for barrier synchronization

    assert(((uintptr_t)(mc_base_addr) % 16) == 0);
    if (datatype == UCC_DT_FLOAT32 || datatype == UCC_DT_BFLOAT16) {
        assert(count_u32 % (4 * tsize) == 0);
    }
    switch (datatype) {
    case UCC_DT_FLOAT32:
        allreduce_kernel_vec32<NvlsFp32Ops><<<sm_count, threads, 0, stream>>>(
            mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
        break;
    case UCC_DT_BFLOAT16:
        allreduce_kernel_vec32<NvlsBf16Ops><<<sm_count, threads, 0, stream>>>(
            mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
        break;
    case UCC_DT_INT32:
        allreduce_kernel_scalar32<NvlsInt32Ops>
            <<<sm_count, threads, 0, stream>>>(
                mc_bar,
                uc_bar,
                expected_blocks,
                launch_counter,
                base_u32,
                count_u32,
                rank,
                tsize);
        break;
    case UCC_DT_UINT32:
        allreduce_kernel_scalar32<NvlsUint32Ops>
            <<<sm_count, threads, 0, stream>>>(
                mc_bar,
                uc_bar,
                expected_blocks,
                launch_counter,
                base_u32,
                count_u32,
                rank,
                tsize);
        break;
    case UCC_DT_INT64:
    {
        uint64_t *base_u64  = reinterpret_cast<uint64_t *>(mc_base_addr);
        size_t    count_u64 = src_size_bytes / sizeof(uint64_t);
        allreduce_kernel_scalar64<NvlsInt64Ops>
            <<<sm_count, threads, 0, stream>>>(
                mc_bar,
                uc_bar,
                expected_blocks,
                launch_counter,
                base_u64,
                count_u64,
                rank,
                tsize);
    } break;
    case UCC_DT_UINT64:
    {
        uint64_t *base_u64  = reinterpret_cast<uint64_t *>(mc_base_addr);
        size_t    count_u64 = src_size_bytes / sizeof(uint64_t);
        allreduce_kernel_scalar64<NvlsUint64Ops>
            <<<sm_count, threads, 0, stream>>>(
                mc_bar,
                uc_bar,
                expected_blocks,
                launch_counter,
                base_u64,
                count_u64,
                rank,
                tsize);
    } break;
    default:
        return UCC_ERR_NOT_SUPPORTED;
    }
    CUDA_CHECK(cudaGetLastError());

    return UCC_OK;
}

#ifdef __cplusplus
}
#endif
