/*
 * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
 *
 * See COPYRIGHT for license information
 */

#include <cuda.h>
#include <cuda_runtime.h>
#include "nvshmem.h"
#include "nvshmemx.h"
#include "nvshmemi_util.h"
#include "nvshmemi_coll.h"
#include "nvshmemi_transfer.h"
#include "gpu_coll.h"
#include "barrier.h"
#include "barrier_device.cuh"

__global__ void barrier_on_stream_kernel(int start, int stride, int size, long *pSync, long *counter);
__global__ void barrier_on_stream_kernel_warp(int start, int stride, int size, long *pSync, long *counter);
__global__ void barrier_on_stream_kernel_block(int start, int stride, int size, long *pSync, long *counter);

__global__ void sync_on_stream_kernel(int start, int stride, int size, long *pSync, long *counter);
__global__ void sync_on_stream_kernel_warp(int start, int stride, int size, long *pSync, long *counter);
__global__ void sync_on_stream_kernel_block(int start, int stride, int size, long *pSync, long *counter);
__global__ void sync_all_on_stream_kernel();
__global__ void sync_all_on_stream_kernel_warp();
__global__ void sync_all_on_stream_kernel_block();

template <threadgroup_t SCOPE>
__global__ void barrier_on_stream_kernel_threadgroup(nvshmem_team_t team) {
#ifdef __CUDA_ARCH__
    int myidx = nvshmemi_thread_id_in_threadgroup<SCOPE>();

    if (nvshmemi_device_state_d.job_connectivity >= NVSHMEMI_JOB_GPU_PROXY) {
        if (!myidx) nvshmemi_transfer_quiet(false);
        nvshmemi_threadgroup_sync<SCOPE>();
    }

    nvshmemi_sync_algo_threadgroup<SCOPE>(team);

    if (!myidx) {
        if (nvshmemi_device_state_d.job_connectivity > NVSHMEMI_JOB_GPU_PROXY)
            nvshmemi_transfer_enforce_consistency_at_target(false);
    }
#endif
}

template <threadgroup_t SCOPE>
__global__ void sync_on_stream_kernel_threadgroup(nvshmem_team_t team) {
#ifdef __CUDA_ARCH__
    nvshmemi_sync_algo_threadgroup<SCOPE>(team);
#endif
}

int nvshmemi_call_barrier_on_stream_kernel(nvshmem_team_t team, cudaStream_t stream) {
    int num_blocks = 1;
    int num_threads_per_block;
    if (nvshmemi_job_connectivity <= NVSHMEMI_JOB_GPU_LDST) {
        int size = nvshmemi_team_pool[team]->size;
        num_threads_per_block = size - 1; // Have enough threads for alltoall algo
    } else {
        num_threads_per_block = nvshmemi_options.BARRIER_TG_DISSEM_KVAL;
    }

    if (num_threads_per_block <= 32) {
        barrier_on_stream_kernel_threadgroup<NVSHMEMI_THREADGROUP_WARP><<<num_blocks, 32, 0, stream>>>(team);
    }
    else {
        barrier_on_stream_kernel_threadgroup<NVSHMEMI_THREADGROUP_BLOCK>
            <<<num_blocks, num_threads_per_block, 0, stream>>>(team);
    }

    CUDA_RUNTIME_CHECK(cudaGetLastError());
    return 0;
}

int nvshmemi_call_sync_on_stream_kernel(nvshmem_team_t team, cudaStream_t stream) {
    int num_blocks = 1;
    int num_threads_per_block;
    if (nvshmemi_job_connectivity <= NVSHMEMI_JOB_GPU_LDST) {
        int size = nvshmemi_team_pool[team]->size;
        num_threads_per_block = size - 1; // Have enough threads for alltoall algo
    } else {
        num_threads_per_block = nvshmemi_options.BARRIER_TG_DISSEM_KVAL;
    }

    if (num_threads_per_block <= 32) {
        sync_on_stream_kernel_threadgroup<NVSHMEMI_THREADGROUP_WARP><<<num_blocks, 32, 0, stream>>>(team);
    } else {
        sync_on_stream_kernel_threadgroup<NVSHMEMI_THREADGROUP_BLOCK>
            <<<num_blocks, num_threads_per_block, 0, stream>>>(team);
    }
    CUDA_RUNTIME_CHECK(cudaGetLastError());
    return 0;
}
