/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2019. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif

#include "cuda_iface.h"
#include "cuda_ctx.inl"
#include "cuda_md.h"

#include <ucs/sys/string.h>


#define UCT_CUDA_DEV_NAME "cuda"


ucs_status_t
uct_cuda_base_query_devices_common(
        uct_md_h md, uct_device_type_t dev_type,
        uct_tl_device_resource_t **tl_devices_p, unsigned *num_tl_devices_p)
{
    ucs_sys_device_t sys_device = UCS_SYS_DEVICE_ID_UNKNOWN;
    CUdevice cuda_device;
    ucs_status_t status;

    if (uct_cuda_ctx_is_active()) {
        status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxGetDevice(&cuda_device));
        if (status != UCS_OK) {
            return status;
        }

        sys_device = uct_cuda_get_sys_dev(cuda_device);
    } else {
        ucs_debug("set cuda sys_device to `unknown` as no context is"
                  " currently active");
    }

    return uct_single_device_resource(md, UCT_CUDA_DEV_NAME, dev_type,
                                      sys_device, tl_devices_p,
                                      num_tl_devices_p);
}

ucs_status_t
uct_cuda_base_query_devices(
        uct_md_h md, uct_tl_device_resource_t **tl_devices_p,
        unsigned *num_tl_devices_p)
{
    return uct_cuda_base_query_devices_common(md, UCT_DEVICE_TYPE_ACC,
                                              tl_devices_p, num_tl_devices_p);
}

#if (__CUDACC_VER_MAJOR__ >= 100000)
void CUDA_CB uct_cuda_base_iface_stream_cb_fxn(void *arg)
#else
void CUDA_CB uct_cuda_base_iface_stream_cb_fxn(CUstream hStream, CUresult status,
                                               void *arg)
#endif
{
    uct_cuda_iface_t *cuda_iface = arg;

    ucs_async_eventfd_signal(cuda_iface->eventfd);
}

ucs_status_t uct_cuda_base_iface_event_fd_get(uct_iface_h tl_iface, int *fd_p)
{
    uct_cuda_iface_t *iface = ucs_derived_of(tl_iface, uct_cuda_iface_t);
    ucs_status_t status;

    if (iface->eventfd == UCS_ASYNC_EVENTFD_INVALID_FD) {
        status = ucs_async_eventfd_create(&iface->eventfd);
        if (status != UCS_OK) {
            return status;
        }
    }

    *fd_p = iface->eventfd;
    return UCS_OK;
}

ucs_status_t uct_cuda_base_check_device_name(const uct_iface_params_t *params)
{
    UCT_CHECK_PARAM(params->field_mask & UCT_IFACE_PARAM_FIELD_DEVICE,
                    "UCT_IFACE_PARAM_FIELD_DEVICE is not defined");

    if (strncmp(params->mode.device.dev_name, UCT_CUDA_DEV_NAME,
                strlen(UCT_CUDA_DEV_NAME)) != 0) {
        ucs_error("no device was found: %s", params->mode.device.dev_name);
        return UCS_ERR_NO_DEVICE;
    }

    return UCS_OK;
}

static UCS_F_ALWAYS_INLINE int
uct_cuda_base_event_is_flush(const uct_cuda_event_desc_t *event)
{
    return event->event == NULL;
}

static UCS_F_ALWAYS_INLINE unsigned
uct_cuda_base_queue_head_ready(ucs_queue_head_t *queue_head)
{
    uct_cuda_event_desc_t *cuda_event;

    if (ucs_queue_is_empty(queue_head)) {
        return 0;
    }

    cuda_event = ucs_queue_head_elem_non_empty(queue_head,
                                               uct_cuda_event_desc_t, queue);
    return uct_cuda_base_event_is_flush(cuda_event) ||
           (CUDA_SUCCESS == cuEventQuery(cuda_event->event));
}

ucs_status_t uct_cuda_base_iface_event_fd_arm(uct_iface_h tl_iface,
                                              unsigned events)
{
    uct_cuda_iface_t *iface = ucs_derived_of(tl_iface, uct_cuda_iface_t);
    ucs_status_t status;
    CUstream *stream;
    ucs_queue_head_t *event_q;
    uct_cuda_queue_desc_t *q_desc;
    ucs_queue_iter_t iter;

    ucs_queue_for_each_safe(q_desc, iter, &iface->active_queue, queue) {
        event_q = &q_desc->event_queue;
        if (uct_cuda_base_queue_head_ready(event_q)) {
            return UCS_ERR_BUSY;
        }
    }

    status = ucs_async_eventfd_poll(iface->eventfd);
    if (status == UCS_OK) {
        return UCS_ERR_BUSY;
    } else if (status == UCS_ERR_IO_ERROR) {
        return status;
    }

    ucs_assertv(status == UCS_ERR_NO_PROGRESS, "%s", ucs_status_string(status));

    ucs_queue_for_each_safe(q_desc, iter, &iface->active_queue, queue) {
        event_q = &q_desc->event_queue;
        stream  = &q_desc->stream;
        if (!ucs_queue_is_empty(event_q)) {
            status =
#if (__CUDACC_VER_MAJOR__ >= 100000)
                UCT_CUDADRV_FUNC_LOG_ERR(
                        cuLaunchHostFunc(*stream,
                                         uct_cuda_base_iface_stream_cb_fxn,
                                         iface));
#else
                UCT_CUDADRV_FUNC_LOG_ERR(
                        cuStreamAddCallback(*stream,
                                            uct_cuda_base_iface_stream_cb_fxn,
                                            iface, 0));
#endif
            if (UCS_OK != status) {
                return status;
            }
        }
    }

    return UCS_OK;
}

static void uct_cuda_base_stream_flushed_cb(uct_completion_t *self)
{
    uct_cuda_flush_stream_desc_t *desc =
            ucs_container_of(self, uct_cuda_flush_stream_desc_t, comp);

    if (--desc->flush_desc->stream_counter == 0) {
        uct_invoke_completion(desc->flush_desc->comp, UCS_OK);
        ucs_free(desc->flush_desc);
    }
}

/* Flush is done by enqueueing flush events on all active streams and wait
 * for them to finish. On each flush event completion, we decrement a shared
 * counter and once it hits zero, flush is completed. */
ucs_status_t
uct_cuda_base_ep_flush(uct_ep_h tl_ep, unsigned flags, uct_completion_t *comp)
{
    uct_base_ep_t UCS_V_UNUSED *ep = ucs_derived_of(tl_ep, uct_base_ep_t);
    uct_cuda_iface_t *iface        = ucs_derived_of(tl_ep->iface,
                                                    uct_cuda_iface_t);
    uct_cuda_flush_desc_t *flush_desc;
    uct_cuda_flush_stream_desc_t *flush_stream_desc;
    uct_cuda_queue_desc_t *q_desc;
    uct_cuda_event_desc_t *event_desc;
    ucs_queue_head_t *event_queue;
    ucs_queue_iter_t iter;
    unsigned stream_index;

    if (ucs_queue_is_empty(&iface->active_queue) || (comp == NULL)) {
        UCT_TL_EP_STAT_FLUSH(ep);
        return UCS_OK;
    }

    /* Allocate base flush descriptor */
    flush_desc = ucs_malloc(sizeof(*flush_desc), "cuda_flush_desc");
    if (flush_desc == NULL) {
        ucs_error("failed to allocate cuda_flush_desc");
        return UCS_ERR_NO_MEMORY;
    }

    flush_desc->comp           = comp;
    flush_desc->stream_counter = 0;

    /* For each active stream, init a flush event and enqueue it on the
     * stream */
    ucs_queue_for_each(q_desc, &iface->active_queue, queue) {
        flush_stream_desc = ucs_mpool_get(&iface->flush_mpool);
        if (flush_stream_desc == NULL) {
            goto error;
        }

        flush_stream_desc->flush_desc  = flush_desc;
        flush_stream_desc->comp.func   = uct_cuda_base_stream_flushed_cb;
        flush_stream_desc->comp.count  = 1;
        flush_stream_desc->super.comp  = &flush_stream_desc->comp;
        flush_stream_desc->super.event = NULL;
        ucs_queue_push(&q_desc->event_queue, &flush_stream_desc->super.queue);
        flush_desc->stream_counter++;
    }

    UCT_TL_EP_STAT_FLUSH_WAIT(ep);
    return UCS_INPROGRESS;

error:
    /* Rollback enqueued items in case of error */
    for (iter = ucs_queue_iter_begin(&iface->active_queue), stream_index = 0;
         stream_index < flush_desc->stream_counter;
         iter = ucs_queue_iter_next(iter), ++stream_index) {
        event_queue = &ucs_queue_iter_elem(q_desc, iter, queue)->event_queue;
        event_desc  = ucs_queue_tail_elem_non_empty(event_queue,
                                                    uct_cuda_event_desc_t,
                                                    queue);

        ucs_queue_remove(event_queue, &event_desc->queue);
        ucs_mpool_put(event_desc);
    }

    ucs_free(flush_desc);
    return UCS_ERR_NO_MEMORY;
}

static UCS_F_ALWAYS_INLINE unsigned
uct_cuda_base_progress_event_queue(uct_cuda_iface_t *iface,
                                   ucs_queue_head_t *queue_head,
                                   unsigned max_events)
{
    unsigned count = 0;
    uct_cuda_event_desc_t *cuda_event;

    ucs_queue_for_each_extract(cuda_event, queue_head, queue,
                               (count < max_events) &&
                               ucs_unlikely(uct_cuda_base_event_is_flush(
                                       cuda_event) ||
                               (cuEventQuery(cuda_event->event) == CUDA_SUCCESS))) {
        ucs_trace_data("cuda event %p completed", cuda_event);
        if (cuda_event->comp != NULL) {
            uct_invoke_completion(cuda_event->comp, UCS_OK);
        }

        if (ucs_likely(!uct_cuda_base_event_is_flush(cuda_event))) {
            iface->ops->complete_event(&iface->super.super, cuda_event);
        }

        ucs_mpool_put(cuda_event);
        count++;
    }

    return count;
}

unsigned uct_cuda_base_iface_progress(uct_iface_h tl_iface)
{
    uct_cuda_iface_t *iface = ucs_derived_of(tl_iface, uct_cuda_iface_t);
    unsigned max_events     = iface->config.max_poll;
    unsigned count          = 0;
    ucs_queue_head_t *event_q;
    uct_cuda_queue_desc_t *q_desc;
    ucs_queue_iter_t iter;

    ucs_queue_for_each_safe(q_desc, iter, &iface->active_queue, queue) {
        event_q = &q_desc->event_queue;
        count  += uct_cuda_base_progress_event_queue(iface, event_q,
                                                     max_events - count);
        if (ucs_queue_is_empty(event_q)) {
            ucs_queue_del_iter(&iface->active_queue, iter);
        }
    }

    return count;
}

ucs_status_t uct_cuda_base_iface_flush(uct_iface_h tl_iface, unsigned flags,
                                       uct_completion_t *comp)
{
    uct_cuda_iface_t *iface = ucs_derived_of(tl_iface, uct_cuda_iface_t);

    if (comp != NULL) {
        return UCS_ERR_UNSUPPORTED;
    }

    if (!ucs_queue_is_empty(&iface->active_queue)) {
        UCT_TL_IFACE_STAT_FLUSH_WAIT(ucs_derived_of(tl_iface, uct_base_iface_t));
        return UCS_INPROGRESS;
    }

    UCT_TL_IFACE_STAT_FLUSH(ucs_derived_of(tl_iface, uct_base_iface_t));
    return UCS_OK;
}

static ucs_status_t
uct_cuda_base_ctx_cmp(const uct_cuda_ctx_rsc_t *ctx_rsc,
                      CUcontext cuda_ctx)
{
#if CUDA_VERSION >= 12000
    unsigned long long ctx_id;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_DEBUG(
            uct_cuda_ctx_get_id(cuda_ctx, &ctx_id));
    if (status != UCS_OK) {
        return status;
    }

    return (ctx_id == ctx_rsc->ctx_id) ? UCS_OK : UCS_ERR_NO_DEVICE;
#else
    return (cuda_ctx == ctx_rsc->ctx) ? UCS_OK : UCS_ERR_NO_DEVICE;
#endif
}

static ucs_status_t
uct_cuda_base_ctx_rsc_check(const uct_cuda_ctx_rsc_t *ctx_rsc)
{
    CUcontext primary_ctx;
    ucs_status_t status;

    if (ctx_rsc->cuda_device == CU_DEVICE_INVALID) {
        return UCS_OK;
    }

    status = uct_cuda_ctx_primary_retain(ctx_rsc->cuda_device, 0,
                                         &primary_ctx);
    if (status != UCS_OK) {
        return status;
    }

    status = uct_cuda_base_ctx_cmp(ctx_rsc, primary_ctx);
    UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(ctx_rsc->cuda_device));
    return status;
}

void uct_cuda_base_stream_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
                                  CUstream *stream)
{
    if (*stream == NULL) {
        return;
    }

    if (uct_cuda_base_ctx_rsc_check(ctx_rsc) != UCS_OK) {
        return;
    }

    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuStreamDestroy(*stream));
}

static void
uct_cuda_base_event_desc_init(ucs_mpool_t *mp, void *obj, void *chunk)
{
    uct_cuda_event_desc_t *event_desc = obj;

    UCT_CUDADRV_FUNC_LOG_ERR(cuEventCreate(&event_desc->event,
                                           CU_EVENT_DISABLE_TIMING));
}

static void uct_cuda_base_event_desc_cleanup(ucs_mpool_t *mp, void *obj)
{
    uct_cuda_event_desc_t *event_desc = obj;
    uct_cuda_ctx_rsc_t *ctx_rsc       = ucs_container_of(mp, uct_cuda_ctx_rsc_t,
                                                         event_mp);

    if (uct_cuda_base_ctx_rsc_check(ctx_rsc) != UCS_OK) {
        return;
    }

    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuEventDestroy(event_desc->event));
}

void uct_cuda_base_queue_desc_init(uct_cuda_queue_desc_t *qdesc)
{
    qdesc->stream = NULL;
    ucs_queue_head_init(&qdesc->event_queue);
}

void uct_cuda_base_queue_desc_destroy(const uct_cuda_ctx_rsc_t *ctx_rsc,
                                      uct_cuda_queue_desc_t *qdesc)
{
    if (!ucs_queue_is_empty(&qdesc->event_queue)) {
        ucs_warn("cuda context %llu stream being destroyed with  %zu "
                 "outstanding events", ctx_rsc->ctx_id,
                 ucs_queue_length(&qdesc->event_queue));
    }

    uct_cuda_base_stream_destroy(ctx_rsc, &qdesc->stream);
}

static ucs_mpool_ops_t uct_cuda_event_desc_mpool_ops = {
    .chunk_alloc   = ucs_mpool_chunk_malloc,
    .chunk_release = ucs_mpool_chunk_free,
    .obj_init      = uct_cuda_base_event_desc_init,
    .obj_cleanup   = uct_cuda_base_event_desc_cleanup,
    .obj_str       = NULL
};

static void uct_cuda_base_ctx_rsc_release_primary_ctx(
        uct_cuda_ctx_rsc_t *ctx_rsc)
{
    if (ctx_rsc->cuda_device == CU_DEVICE_INVALID) {
        return;
    }

    UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(ctx_rsc->cuda_device));
    ctx_rsc->cuda_device = CU_DEVICE_INVALID;
}

static ucs_status_t
uct_cuda_base_ctx_rsc_retain_primary_ctx(uct_cuda_ctx_rsc_t *ctx_rsc)
{
    CUcontext primary_ctx;
    CUdevice cuda_device;
    ucs_status_t status;

    ctx_rsc->cuda_device = CU_DEVICE_INVALID;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxGetDevice(&cuda_device));
    if (status != UCS_OK) {
        return status;
    }

    status = uct_cuda_ctx_primary_retain(cuda_device, 0, &primary_ctx);
    if (status == UCS_ERR_NO_DEVICE) {
        return UCS_OK;
    } else if (status != UCS_OK) {
        return status;
    }

    status = uct_cuda_base_ctx_cmp(ctx_rsc, primary_ctx);
    if (status == UCS_ERR_NO_DEVICE) {
        ucs_debug("cuda context %llu is not the primary context on device %d",
                  ctx_rsc->ctx_id, cuda_device);
        status = UCS_OK;
        goto out_release_primary_ctx;
    } else if (status != UCS_OK) {
        goto out_release_primary_ctx;
    }

    ctx_rsc->ctx         = primary_ctx;
    ctx_rsc->cuda_device = cuda_device;
    return UCS_OK;

out_release_primary_ctx:
    UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));
    return status;
}

ucs_status_t uct_cuda_base_ctx_rsc_create(uct_cuda_iface_t *iface,
                                          unsigned long long ctx_id,
                                          uct_cuda_ctx_rsc_t **ctx_rsc_p)
{
    CUcontext ctx;
    ucs_status_t status;
    int ret;
    khiter_t iter;
    uct_cuda_ctx_rsc_t *ctx_rsc;
    ucs_mpool_params_t mp_params;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxGetCurrent(&ctx));
    if (status != UCS_OK) {
        return status;
    } else if (ctx == NULL) {
        ucs_error("no cuda context bound to calling thread");
        return UCS_ERR_IO_ERROR;
    }

    iter = kh_put(cuda_ctx_rscs, &iface->ctx_rscs, ctx_id, &ret);
    if (ret == UCS_KH_PUT_FAILED) {
        ucs_error("failed to allocate cuda context resource hash entry");
        return UCS_ERR_NO_MEMORY;
    }

    ucs_assertv_always(ret != UCS_KH_PUT_KEY_PRESENT,
                       "the key has already been added iface=%p key=%llu",
                       iface, ctx_id);

    ctx_rsc = iface->ops->create_rsc(&iface->super.super);
    if (ctx_rsc == NULL) {
        status = UCS_ERR_NO_MEMORY;
        goto err_del_iter;
    }

    ctx_rsc->ctx    = ctx;
    ctx_rsc->ctx_id = ctx_id;

    status = uct_cuda_base_ctx_rsc_retain_primary_ctx(ctx_rsc);
    if (status != UCS_OK) {
        goto err_free_ctx_rsc;
    }

    ucs_mpool_params_reset(&mp_params);
    mp_params.elem_size       = iface->config.event_desc_size;
    mp_params.elems_per_chunk = 128;
    mp_params.max_elems       = iface->config.max_events;
    mp_params.ops             = &uct_cuda_event_desc_mpool_ops;
    mp_params.name            = "cuda_event_descriptors";

    status = ucs_mpool_init(&mp_params, &ctx_rsc->event_mp);
    if (status != UCS_OK) {
        goto err_release_primary_ctx;
    }

    kh_value(&iface->ctx_rscs, iter) = ctx_rsc;
    *ctx_rsc_p                       = ctx_rsc;
    return UCS_OK;

err_release_primary_ctx:
    uct_cuda_base_ctx_rsc_release_primary_ctx(ctx_rsc);
err_free_ctx_rsc:
    iface->ops->destroy_rsc(&iface->super.super, ctx_rsc);
err_del_iter:
    kh_del(cuda_ctx_rscs, &iface->ctx_rscs, iter);
    return status;
}

static void uct_cuda_base_ctx_rsc_destroy(uct_cuda_iface_t *iface,
                                          uct_cuda_ctx_rsc_t *ctx_rsc)
{
    const CUdevice cuda_device = ctx_rsc->cuda_device;

    ucs_mpool_cleanup(&ctx_rsc->event_mp, 1);
    iface->ops->destroy_rsc(&iface->super.super, ctx_rsc);

    if (cuda_device != CU_DEVICE_INVALID) {
        UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));
    }
}

static ucs_mpool_ops_t uct_cuda_flush_desc_mpool_ops = {
    .chunk_alloc   = ucs_mpool_chunk_malloc,
    .chunk_release = ucs_mpool_chunk_free,
    .obj_init      = NULL,
    .obj_cleanup   = NULL,
    .obj_str       = NULL
};

UCS_CLASS_INIT_FUNC(uct_cuda_iface_t, uct_iface_ops_t *tl_ops,
                    uct_iface_internal_ops_t *ops, uct_md_h md,
                    uct_worker_h worker, const uct_iface_params_t *params,
                    const uct_iface_config_t *tl_config,
                    const char *dev_name)
{
    ucs_mpool_params_t mp_params;
    ucs_status_t status;

    UCS_CLASS_CALL_SUPER_INIT(uct_base_iface_t, tl_ops, ops, md, worker, params,
                              tl_config UCS_STATS_ARG(params->stats_root)
                              UCS_STATS_ARG(dev_name));

    self->eventfd = UCS_ASYNC_EVENTFD_INVALID_FD;
    kh_init_inplace(cuda_ctx_rscs, &self->ctx_rscs);
    ucs_queue_head_init(&self->active_queue);

    ucs_mpool_params_reset(&mp_params);
    mp_params.elem_size = sizeof(uct_cuda_flush_stream_desc_t);
    mp_params.ops       = &uct_cuda_flush_desc_mpool_ops;
    mp_params.name      = "cuda_flush_descriptors";

    status = ucs_mpool_init(&mp_params, &self->flush_mpool);
    if (status != UCS_OK) {
        kh_destroy_inplace(cuda_ctx_rscs, &self->ctx_rscs);
    }

    return status;
}

static UCS_CLASS_CLEANUP_FUNC(uct_cuda_iface_t)
{
    uct_cuda_ctx_rsc_t *ctx_rsc;

    uct_base_iface_progress_disable(&self->super.super,
                                    UCT_PROGRESS_SEND | UCT_PROGRESS_RECV);

    kh_foreach_value(&self->ctx_rscs, ctx_rsc, {
        uct_cuda_base_ctx_rsc_destroy(self, ctx_rsc);
    });

    kh_destroy_inplace(cuda_ctx_rscs, &self->ctx_rscs);
    ucs_async_eventfd_destroy(self->eventfd);
    ucs_mpool_cleanup(&self->flush_mpool, 1);
}

UCS_CLASS_DEFINE(uct_cuda_iface_t, uct_base_iface_t);
