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

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

#include "gdaki.h"

#include <ucs/sys/sock.h>
#include <ucs/sys/topo/base/topo.h>
#include <ucs/time/time.h>
#include <ucs/datastruct/string_buffer.h>
#include <ucs/algorithm/qsort_r.h>
#include <ucs/type/init_once.h>
#include <ucs/type/serialize.h>
#include <uct/ib/base/ib_verbs.h>
#include <uct/ib/mlx5/rc/rc_mlx5.h>
#include <uct/cuda/cuda_copy/cuda_copy_md.h>
#include <uct/cuda/base/cuda_util.h>
#include <uct/cuda/base/cuda_ctx.h>
#include <uct/cuda/base/cuda_nvml.h>

#include <string.h>

#include "gpunetio/common/doca_gpunetio_verbs_def.h"

#define UCT_GDAKI_MAX_CUDA_DEVICES 64

#define UCT_GDAKI_CQ_PAGE_OFFSET_SHIFT 6
#define UCT_GDAKI_DEV_EP_SIZE          64
#define UCT_GDAKI_DEV_QP_SIZE          128

#define UCT_GDAKI_CUDA_REG_FLAGS \
    (CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP | \
     CU_MEMHOSTREGISTER_IOMEMORY)

typedef enum {
    UCT_RC_GDAKI_EP_ALLOC_MODE_POOL,
    UCT_RC_GDAKI_EP_ALLOC_MODE_DIRECT,
    UCT_RC_GDAKI_EP_ALLOC_MODE_LAST
} uct_rc_gdaki_ep_alloc_mode_t;

static const char *uct_rc_gdaki_ep_alloc_mode_names[] = {
    [UCT_RC_GDAKI_EP_ALLOC_MODE_POOL]   = "pool",
    [UCT_RC_GDAKI_EP_ALLOC_MODE_DIRECT] = "direct",
    [UCT_RC_GDAKI_EP_ALLOC_MODE_LAST]   = NULL
};

enum {
    UCT_GDAKI_SUPPORTED   = UCS_BIT(0),
    UCT_GDAKI_DMABUF_UMEM = UCS_BIT(1)
};

typedef struct {
    uct_rc_iface_common_config_t      super;
    uct_rc_mlx5_iface_common_config_t mlx5;
    uct_rc_gdaki_ep_alloc_mode_t      ep_alloc_mode;
    unsigned                          num_channels;
    unsigned                          num_eps;
} uct_rc_gdaki_iface_config_t;

ucs_config_field_t uct_rc_gdaki_iface_config_table[] = {
    {UCT_IB_CONFIG_PREFIX, "", NULL,
     ucs_offsetof(uct_rc_gdaki_iface_config_t, super),
     UCS_CONFIG_TYPE_TABLE(uct_rc_iface_common_config_table)},

    {UCT_IB_CONFIG_PREFIX, "", NULL,
     ucs_offsetof(uct_rc_gdaki_iface_config_t, mlx5),
     UCS_CONFIG_TYPE_TABLE(uct_rc_mlx5_common_config_table)},

    {"NUM_CHANNELS", "1",
     "Number of channels. \nRounded up to the next power-of-2 value.\n"
     "Maximum value is 256.",
     ucs_offsetof(uct_rc_gdaki_iface_config_t, num_channels),
     UCS_CONFIG_TYPE_UINT},

    {"EP_ALLOC_MODE", "pool",
     "Endpoint channel allocation mode.\n"
     " - pool   : allocate channels from pooled chunks\n"
     " - direct : allocate channels per endpoint directly",
     ucs_offsetof(uct_rc_gdaki_iface_config_t, ep_alloc_mode),
     UCS_CONFIG_TYPE_ENUM(uct_rc_gdaki_ep_alloc_mode_names)},

    {"NUM_PRE_ALLOC_EPS", "32",
     "Expected number of pre-allocated endpoints. The channel pool grows if\n"
     "more are needed.",
     ucs_offsetof(uct_rc_gdaki_iface_config_t, num_eps), UCS_CONFIG_TYPE_UINT},

    {NULL}
};


ucs_status_t
uct_rc_gdaki_alloc(size_t size, size_t align, void **buf_p, CUdeviceptr *orig_p)
{
    unsigned int flag = 1;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuMemAlloc(orig_p, size + align - 1));
    if (status != UCS_OK) {
        return status;
    }

    *buf_p = (void*)ucs_align_up_pow2_ptr(*orig_p, align);
    status = UCT_CUDADRV_FUNC_LOG_ERR(
            cuPointerSetAttribute(&flag, CU_POINTER_ATTRIBUTE_SYNC_MEMOPS,
                                  (CUdeviceptr)*buf_p));
    if (status != UCS_OK) {
        goto err;
    }

    return UCS_OK;

err:
    cuMemFree(*orig_p);
    return status;
}

static void uct_rc_gdaki_calc_dev_ep_layout(size_t num_channels, size_t wq_len,
                                            size_t *dev_ep_size_p,
                                            uint64_t *pgsz_bitmap_p)
{
    uint64_t max_page_size;
    size_t wq_umem_offset;

    UCS_STATIC_ASSERT(sizeof(uct_rc_gdaki_dev_ep_t) == UCT_GDAKI_DEV_EP_SIZE);
    UCS_STATIC_ASSERT(sizeof(uct_rc_gdaki_dev_qp_t) == UCT_GDAKI_DEV_QP_SIZE);

    wq_umem_offset = UCT_GDAKI_DEV_EP_SIZE +
                     (UCT_GDAKI_DEV_QP_SIZE * num_channels);
    *dev_ep_size_p = wq_umem_offset + wq_len * num_channels;
    max_page_size  = ucs_min(UCT_GDAKI_DEV_EP_SIZE, UCT_GDAKI_DEV_QP_SIZE)
                     << UCT_GDAKI_CQ_PAGE_OFFSET_SHIFT;
    *pgsz_bitmap_p = (max_page_size << 1) - 1;
}

static int uct_gdaki_check_umem_dmabuf(const uct_ib_md_t *md)
{
    int ret = 0;
#if HAVE_DECL_MLX5DV_UMEM_MASK_DMABUF
    struct mlx5dv_devx_umem_in umem_in = {};
    struct mlx5dv_devx_umem *umem;
    uct_cuda_copy_md_dmabuf_t dmabuf;
    CUdeviceptr buff;

    if (UCT_CUDADRV_FUNC_LOG_ERR(cuMemAlloc(&buff, 1)) != UCS_OK) {
        goto out;
    }

    dmabuf = uct_cuda_copy_md_get_dmabuf((void*)buff, 1,
                                         UCS_SYS_DEVICE_ID_UNKNOWN);
    if (dmabuf.fd == UCT_DMABUF_FD_INVALID) {
        goto out_free;
    }

    umem_in.addr        = (void *)(uintptr_t)dmabuf.offset;
    umem_in.size        = 1;
    umem_in.access      = IBV_ACCESS_LOCAL_WRITE;
    umem_in.pgsz_bitmap = UINT64_MAX & ~(ucs_get_page_size() - 1);
    umem_in.comp_mask   = MLX5DV_UMEM_MASK_DMABUF;
    umem_in.dmabuf_fd   = dmabuf.fd;

    umem = mlx5dv_devx_umem_reg_ex(md->dev.ibv_context, &umem_in);
    if (umem == NULL) {
        goto out_close;
    }

    ret = 1;
out_dereg:
    mlx5dv_devx_umem_dereg(umem);
out_close:
    ucs_close_fd(&dmabuf.fd);
out_free:
    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemFree(buff));
out:
#endif
    return ret;
}

static int uct_gdaki_is_dmabuf_supported(const uct_ib_md_t *md)
{
    static int dmabuf_supported = -1;
    int loglevel                = (md->config.gda_dmabuf_enable == UCS_YES) ?
                                                         UCS_LOG_LEVEL_DIAG :
                                                         UCS_LOG_LEVEL_DEBUG;

    if (dmabuf_supported != -1) {
        return dmabuf_supported;
    }

    if (!(md->cap_flags & UCT_MD_FLAG_REG_DMABUF)) {
        dmabuf_supported = 0;
        ucs_log(loglevel, "IB doesn't support DMA-BUF");
    } else if (!uct_cuda_copy_md_is_dmabuf_supported()) {
        dmabuf_supported = 0;
        ucs_log(loglevel, "CUDA doesn't support DMA-BUF");
    } else if (!uct_gdaki_check_umem_dmabuf(md)) {
        dmabuf_supported = 0;
        ucs_log(loglevel, "DEVX UMEM doesn't support DMA-BUF");
    } else {
        dmabuf_supported = 1;
    }

    return dmabuf_supported;
}

static int uct_gdaki_md_check_uar(const uct_ib_md_t *md)
{
    int ret = 0;
    struct mlx5dv_devx_uar *uar;
    ucs_status_t status;

    uar = mlx5dv_devx_alloc_uar(md->dev.ibv_context, 0);
    if (uar == NULL) {
        return 0;
    }

    status = UCT_CUDADRV_FUNC_LOG_DEBUG(
            cuMemHostRegister(uar->reg_addr, UCT_IB_MLX5_BF_REG_SIZE,
                              UCT_GDAKI_CUDA_REG_FLAGS));
    if (status != UCS_OK) {
        goto out;
    }

    UCT_CUDADRV_FUNC_LOG_WARN(cuMemHostUnregister(uar->reg_addr));
    ret = 1;

out:
    mlx5dv_devx_free_uar(uar);
    return ret;
}

static CUdevice uct_gdaki_push_primary_ctx(int retain_inactive_ctx)
{
    CUdevice cuda_dev;
    ucs_status_t status;

    status = uct_cuda_ctx_primary_push_first_active(&cuda_dev);
    if (status == UCS_OK) {
        return cuda_dev;
    }

    if ((status != UCS_ERR_NO_DEVICE) || !retain_inactive_ctx) {
        if (status == UCS_ERR_NO_DEVICE) {
            ucs_diag("no active primary CUDA context on any device. Please set "
                     "UCX_IB_GDA_RETAIN_INACTIVE_CTX=yes to retain inactive "
                     "context.");
        }
        return CU_DEVICE_INVALID;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&cuda_dev, 0));
    if (status != UCS_OK) {
        return CU_DEVICE_INVALID;
    }

    status = uct_cuda_ctx_primary_push(cuda_dev, 1, UCS_LOG_LEVEL_ERROR);
    if (status != UCS_OK) {
        return CU_DEVICE_INVALID;
    }

    return cuda_dev;
}

static uint32_t
uct_gdaki_check_cuda_ctx_dependent_features(const uct_ib_md_t *md)
{
    uint32_t features = 0;
    CUdevice cuda_dev;

    cuda_dev = uct_gdaki_push_primary_ctx(md->config.gda_retain_inactive_ctx);
    if (cuda_dev == CU_DEVICE_INVALID) {
        return features;
    }

    if (!uct_gdaki_md_check_uar(md)) {
        ucs_diag("GDAKI not supported, please add NVreg_RegistryDwords="
                 "\"PeerMappingOverride=1;\" option for nvidia kernel driver");
        goto out;
    }

    if ((md->config.gda_dmabuf_enable != UCS_NO) &&
        uct_gdaki_is_dmabuf_supported(md)) {
        features |= UCT_GDAKI_DMABUF_UMEM;
        ucs_debug("%s: using dmabuf for gda transport",
                  uct_ib_device_name(&md->dev));
    } else if ((md->config.gda_dmabuf_enable != UCS_YES) &&
               (md->reg_mem_types & UCS_BIT(UCS_MEMORY_TYPE_CUDA))) {
        ucs_debug("%s: using peermem for gda transport",
                  uct_ib_device_name(&md->dev));
    } else {
        if (md->config.gda_dmabuf_enable != UCS_YES) {
            ucs_diag("GDAKI not supported, please load Nvidia peermem driver "
                     "by running \"modprobe nvidia_peermem\"");
        }
        goto out;
    }

    features |= UCT_GDAKI_SUPPORTED;

out:
    uct_cuda_ctx_primary_pop_and_release(cuda_dev);
    return features;
}

static uint32_t uct_gdaki_get_driver_features(const uct_ib_md_t *md)
{
    static ucs_init_once_t once = UCS_INIT_ONCE_INITIALIZER;
    static uint32_t features;

    UCS_INIT_ONCE(&once) {
        features = uct_gdaki_check_cuda_ctx_dependent_features(md);
    }

    return features;
}

static struct mlx5dv_devx_umem *
uct_rc_gdaki_umem_reg(const uct_ib_md_t *md, struct ibv_context *ibv_context,
                      void *address, size_t length, uint64_t pgsz_bitmap)
{
#if HAVE_DECL_MLX5DV_UMEM_MASK_DMABUF
    struct mlx5dv_devx_umem_in umem_in = {};
    uct_cuda_copy_md_dmabuf_t dmabuf   = {
        .fd     = UCT_DMABUF_FD_INVALID,
        .offset = 0
    };
    struct mlx5dv_devx_umem *umem;

    umem_in.size        = length;
    umem_in.access      = IBV_ACCESS_LOCAL_WRITE;
    umem_in.pgsz_bitmap = pgsz_bitmap;

    if (uct_gdaki_get_driver_features(md) & UCT_GDAKI_DMABUF_UMEM) {
        dmabuf = uct_cuda_copy_md_get_dmabuf(address, length,
                                             UCS_SYS_DEVICE_ID_UNKNOWN);
    }

    if (dmabuf.fd == UCT_DMABUF_FD_INVALID) {
        umem_in.addr      = address;
    } else {
        umem_in.comp_mask = MLX5DV_UMEM_MASK_DMABUF;
        umem_in.dmabuf_fd = dmabuf.fd;
        umem_in.addr      = (void *)(uintptr_t)dmabuf.offset;
    }

    umem = mlx5dv_devx_umem_reg_ex(ibv_context, &umem_in);
    if (dmabuf.fd != UCT_DMABUF_FD_INVALID) {
        ucs_close_fd(&dmabuf.fd);
    }

    return umem;
#else
    return mlx5dv_devx_umem_reg(ibv_context, address, length,
                                IBV_ACCESS_LOCAL_WRITE);
#endif
}

static ucs_status_t
uct_rc_gdaki_init_umem(uct_rc_gdaki_iface_t *iface, uint64_t pgsz_bitmap,
                       size_t mem_size, uct_rc_gdaki_channel_block_mem_t *mem)
{
    uct_ib_md_t *md = ucs_derived_of(iface->super.super.super.super.md,
                                     uct_ib_md_t);
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(iface->cuda_ctx));
    if (status != UCS_OK) {
        return status;
    }

    status = uct_rc_gdaki_alloc(mem_size, ucs_get_page_size(), &mem->gpu_mem,
                                &mem->gpu_raw);
    if (status != UCS_OK) {
        goto out_ctx;
    }

    mem->umem = uct_rc_gdaki_umem_reg(md, md->dev.ibv_context, mem->gpu_mem,
                                      mem_size, pgsz_bitmap);
    if (mem->umem == NULL) {
        uct_ib_check_memlock_limit_msg(md->dev.ibv_context, UCS_LOG_LEVEL_ERROR,
                                       "mlx5dv_devx_umem_reg(ptr=%p size=%zu)",
                                       mem->gpu_mem, mem_size);
        status = UCS_ERR_NO_MEMORY;
        goto err_umem;
    }

    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));
    return UCS_OK;

err_umem:
    cuMemFree(mem->gpu_raw);
    mem->gpu_mem = NULL;
    mem->gpu_raw = 0;
out_ctx:
    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));
    return status;
}

typedef struct {
    uint64_t pgsz_bitmap;
    size_t   dev_ep_size;
} uct_rc_gdaki_pool_priv_t;

static uct_rc_gdaki_channel_block_t *
uct_rc_gdaki_channel_block(uct_rc_gdaki_iface_t *iface, ucs_mpool_t *mp,
                           void *elems, unsigned ep_index)
{
    if (iface->ep_alloc_mode == UCT_RC_GDAKI_EP_ALLOC_MODE_POOL) {
        return ucs_mpool_chunk_obj(mp, elems, ep_index);
    } else {
        return elems;
    }
}

static void uct_rc_gdaki_chunk_channels_destroy(uct_rc_gdaki_iface_t *iface,
                                                ucs_mpool_t *mp, void *elems,
                                                unsigned num_elems,
                                                unsigned end_ep,
                                                unsigned end_ch)
{
    uct_rc_gdaki_channel_block_t *channel_block;
    unsigned ep_index, channel_index, num_channels;

    for (ep_index = 0; ep_index < num_elems; ++ep_index) {
        channel_block = uct_rc_gdaki_channel_block(iface, mp, elems, ep_index);
        num_channels  = (ep_index < end_ep) ? end_ch + 1 : end_ch;
        for (channel_index = 0; channel_index < num_channels; ++channel_index) {
            uct_ib_mlx5_devx_destroy_qp_common(
                    &channel_block->channels[channel_index].qp.super);
            uct_ib_mlx5_devx_destroy_cq_common(
                    &channel_block->channels[channel_index].cq);
        }
    }
}

static ucs_status_t
uct_rc_gdaki_init_channel_chunk(uct_rc_gdaki_iface_t *iface,
                                uct_rc_gdaki_channel_block_mem_t *mem,
                                size_t dev_ep_size, ucs_mpool_t *mp,
                                void *elems, unsigned num_elems)
{
    uct_ib_iface_init_attr_t init_attr = {};
    uct_ib_mlx5_cq_attr_t cq_attr      = {};
    uct_ib_mlx5_qp_attr_t qp_attr      = {};
    uct_ib_mlx5_dbrec_t dbrec;
    size_t ep_offset, ep_qp_offset, ep_wq_offset;
    unsigned ep_index, channel_index;
    uct_rc_gdaki_channel_block_t *channel_block;
    uct_rc_gdaki_channel_t *channel;
    ucs_status_t status;

    init_attr.cq_len[UCT_IB_DIR_TX] = 1;
    uct_ib_mlx5_cq_calc_sizes(&iface->super.super.super, UCT_IB_DIR_TX,
                              &init_attr, 0, &cq_attr);
    uct_rc_iface_fill_attr(&iface->super.super, &qp_attr.super,
                           iface->super.super.config.tx_qp_len, NULL);

    cq_attr.flags                           |= UCT_IB_MLX5_CQ_IGNORE_OVERRUN;
    qp_attr.mmio_mode                        = UCT_IB_MLX5_MMIO_MODE_DB;
    qp_attr.super.srq_num                    = 0;
    qp_attr.super.max_inl_cqe[UCT_IB_DIR_TX] = 0;
    uct_ib_mlx5_wq_calc_sizes(&qp_attr);

    ep_qp_offset = UCT_GDAKI_DEV_EP_SIZE;
    ep_wq_offset = ep_qp_offset + iface->num_channels * UCT_GDAKI_DEV_QP_SIZE;
    dbrec.mem_id = mem->umem->umem_id;

    for (channel_index = 0; channel_index < iface->num_channels;
         ++channel_index) {
        size_t ch_cq_offset = channel_index * UCT_GDAKI_DEV_QP_SIZE;
        size_t ch_wq_offset = channel_index * qp_attr.len;
        for (ep_index = 0; ep_index < num_elems; ++ep_index) {
            ep_offset     = ep_index * dev_ep_size;
            channel_block = uct_rc_gdaki_channel_block(iface, mp, elems,
                                                       ep_index);
            channel       = &channel_block->channels[channel_index];
            if (channel_index == 0) {
                channel_block->gpu_ptr = (uintptr_t)
                        UCS_PTR_BYTE_OFFSET(mem->gpu_mem, ep_offset);
            }

            channel->cq.devx.mem.mem       = mem->umem;
            channel->qp.super.devx.mem.mem = mem->umem;

            cq_attr.umem_offset    = ep_offset + ep_qp_offset + ch_cq_offset;
            dbrec.offset           = cq_attr.umem_offset +
                                     ucs_offsetof(uct_rc_gdaki_dev_qp_t, cq_dbrec);
            channel->cq.devx.dbrec = &dbrec;
            status = uct_ib_mlx5_devx_create_cq_common(&iface->super.super.super,
                                                       UCT_IB_DIR_TX, &cq_attr,
                                                       &channel->cq, 0, 0);
            if (status != UCS_OK) {
                goto err_cleanup;
            }

            qp_attr.umem_offset = ep_offset + ep_wq_offset + ch_wq_offset;
            dbrec.offset        = cq_attr.umem_offset +
                                  ucs_offsetof(uct_rc_gdaki_dev_qp_t, qp_dbrec);
            channel->qp.super.devx.dbrec = &dbrec;
            status                       = uct_ib_mlx5_devx_create_qp_common(
                    &iface->super.super.super, &channel->cq, &channel->cq,
                    &channel->qp.super, &channel->qp, &qp_attr);
            if (status != UCS_OK) {
                uct_ib_mlx5_devx_destroy_cq_common(&channel->cq);
                goto err_cleanup;
            }
        }
    }

    return UCS_OK;

err_cleanup:
    uct_rc_gdaki_chunk_channels_destroy(iface, mp, elems, num_elems, ep_index,
                                        channel_index);
    ucs_error("gdaki channel chunk init: create_cq/qp failed (ch=%u ep=%u)",
              channel_index, ep_index);
    return status;
}

static ucs_status_t
uct_rc_gdaki_pool_chunk_alloc(ucs_mpool_t *mp, size_t *size_p, void **chunk_p)
{
    uct_rc_gdaki_pool_priv_t *priv = ucs_mpool_priv(mp);
    uct_rc_gdaki_iface_t *iface    = ucs_container_of(mp, uct_rc_gdaki_iface_t,
                                                      channel_pool);
    uct_rc_gdaki_channel_block_mem_t *hdr;
    size_t host_alloc_size, gpu_alloc_size;
    unsigned num_elems;
    ucs_status_t status;

    host_alloc_size = sizeof(*hdr) + *size_p;
    hdr             = ucs_malloc(host_alloc_size, "gdaki pool chunk");
    if (hdr == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    *chunk_p  = hdr + 1;
    num_elems = ucs_mpool_num_elems_per_chunk(mp, *chunk_p, *size_p);
    if (num_elems == 0) {
        status = UCS_ERR_NO_MEMORY;
        goto err_free_hdr;
    }

    gpu_alloc_size = num_elems * priv->dev_ep_size;
    status = uct_rc_gdaki_init_umem(iface, priv->pgsz_bitmap, gpu_alloc_size,
                                    hdr);
    if (status != UCS_OK) {
        goto err_free_hdr;
    }

    status = uct_rc_gdaki_init_channel_chunk(
            iface, hdr, priv->dev_ep_size, mp,
            ucs_mpool_chunk_elems(mp, (void*)(hdr + 1)), num_elems);
    if (status != UCS_OK) {
        goto err_umem;
    }

    return UCS_OK;

err_umem:
    mlx5dv_devx_umem_dereg(hdr->umem);
    cuMemFree(hdr->gpu_raw);
err_free_hdr:
    ucs_free(hdr);
    return status;
}


static void uct_rc_gdaki_pool_chunk_release(ucs_mpool_t *mp, void *chunk)
{
    uct_rc_gdaki_iface_t *iface = ucs_container_of(mp, uct_rc_gdaki_iface_t,
                                                   channel_pool);
    uct_rc_gdaki_channel_block_mem_t *hdr =
            (uct_rc_gdaki_channel_block_mem_t*)chunk - 1;
    ucs_mpool_chunk_t *mp_chunk = chunk;

    uct_rc_gdaki_chunk_channels_destroy(iface, mp, mp_chunk->elems,
                                        mp_chunk->num_elems,
                                        mp_chunk->num_elems,
                                        iface->num_channels - 1);
    mlx5dv_devx_umem_dereg(hdr->umem);
    cuMemFree(hdr->gpu_raw);
    ucs_free(hdr);
}

static ucs_mpool_ops_t uct_rc_gdaki_pool_mpool_ops = {
    .chunk_alloc   = uct_rc_gdaki_pool_chunk_alloc,
    .chunk_release = uct_rc_gdaki_pool_chunk_release,
    .obj_init      = NULL,
    .obj_cleanup   = NULL,
    .obj_str       = NULL,
};

static ucs_status_t
uct_rc_gdaki_iface_init_channel_pool(uct_rc_gdaki_iface_t *iface,
                                     const uct_rc_gdaki_iface_config_t *config)
{
    uct_ib_mlx5_qp_attr_t qp_attr = {};
    ucs_mpool_params_t mp_params;
    uct_rc_gdaki_pool_priv_t *priv;
    ucs_status_t status;

    ucs_mpool_params_reset(&mp_params);
    mp_params.priv_size       = sizeof(uct_rc_gdaki_pool_priv_t);
    mp_params.elem_size       = sizeof(uct_rc_gdaki_channel_block_t) +
                                (sizeof(uct_rc_gdaki_channel_t) *
                                 iface->num_channels);
    mp_params.alignment       = UCS_SYS_CACHE_LINE_SIZE;
    mp_params.elems_per_chunk = config->num_eps;
    mp_params.max_elems       = UINT_MAX;
    mp_params.ops             = &uct_rc_gdaki_pool_mpool_ops;
    mp_params.name            = "gdaki_channel_pool";

    status = ucs_mpool_init(&mp_params, &iface->channel_pool);
    if (status != UCS_OK) {
        return status;
    }

    uct_rc_iface_fill_attr(&iface->super.super, &qp_attr.super,
                           iface->super.super.config.tx_qp_len, NULL);
    uct_ib_mlx5_wq_calc_sizes(&qp_attr);

    priv = ucs_mpool_priv(&iface->channel_pool);
    uct_rc_gdaki_calc_dev_ep_layout(iface->num_channels, qp_attr.len,
                                    &priv->dev_ep_size, &priv->pgsz_bitmap);

    return UCS_OK;
}

static void uct_rc_gdaki_iface_cleanup_channel_pool(uct_rc_gdaki_iface_t *iface)
{
    ucs_mpool_t *mp = &iface->channel_pool;

    if (mp->data == NULL) {
        return;
    }
    ucs_mpool_cleanup(mp, 0);
}

static ucs_status_t
uct_rc_gdaki_ep_init_channels_pooled(uct_rc_gdaki_iface_t *iface,
                                     uct_rc_gdaki_ep_t *ep)
{
    ep->channel_block = ucs_mpool_get(&iface->channel_pool);

    if (ep->channel_block == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    return UCS_OK;
}

static UCS_F_ALWAYS_INLINE void
uct_rc_gdaki_ep_reset_channels(uct_rc_gdaki_ep_t *ep)
{
    ep->mem.umem      = NULL;
    ep->mem.gpu_mem   = NULL;
    ep->mem.gpu_raw   = 0;
    ep->channel_block = NULL;
}

static void uct_rc_gdaki_cleanup_channels_pooled(uct_rc_gdaki_ep_t *ep)
{
    if (ep->channel_block == NULL) {
        return;
    }

    ucs_mpool_put(ep->channel_block);
    uct_rc_gdaki_ep_reset_channels(ep);
}

static ucs_status_t
uct_rc_gdaki_ep_init_channels_direct(uct_rc_gdaki_iface_t *iface,
                                     uct_rc_gdaki_ep_t *ep)
{
    uct_ib_mlx5_qp_attr_t qp_attr = {};
    size_t channel_block_size     = sizeof(uct_rc_gdaki_channel_block_t) +
                                    (iface->num_channels *
                                     sizeof(uct_rc_gdaki_channel_t));
    uint64_t pgsz_bitmap;
    size_t dev_ep_size;
    ucs_status_t status;

    uct_rc_iface_fill_attr(&iface->super.super, &qp_attr.super,
                           iface->super.super.config.tx_qp_len, NULL);
    uct_ib_mlx5_wq_calc_sizes(&qp_attr);
    uct_rc_gdaki_calc_dev_ep_layout(iface->num_channels, qp_attr.len,
                                    &dev_ep_size, &pgsz_bitmap);

    ep->channel_block = ucs_calloc(1, channel_block_size, "gdaki channels");
    if (ep->channel_block == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    status = uct_rc_gdaki_init_umem(iface, pgsz_bitmap, dev_ep_size, &ep->mem);
    if (status != UCS_OK) {
        goto err_block;
    }

    ep->channel_block->gpu_ptr = (uintptr_t)ep->mem.gpu_mem;

    status = uct_rc_gdaki_init_channel_chunk(iface, &ep->mem, dev_ep_size, NULL,
                                             ep->channel_block, 1);
    if (status != UCS_OK) {
        goto err_umem;
    }

    return UCS_OK;

err_umem:
    mlx5dv_devx_umem_dereg(ep->mem.umem);
    cuMemFree(ep->mem.gpu_raw);
err_block:
    ucs_free(ep->channel_block);
    uct_rc_gdaki_ep_reset_channels(ep);
    return status;
}

static void uct_rc_gdaki_cleanup_channels_direct(uct_rc_gdaki_iface_t *iface,
                                                 uct_rc_gdaki_ep_t *ep)
{
    if (ep->channel_block == NULL) {
        return;
    }

    uct_rc_gdaki_chunk_channels_destroy(iface, NULL, ep->channel_block, 1, 1,
                                        iface->num_channels - 1);
    mlx5dv_devx_umem_dereg(ep->mem.umem);
    cuMemFree(ep->mem.gpu_raw);
    ucs_free(ep->channel_block);
    uct_rc_gdaki_ep_reset_channels(ep);
}

static ucs_status_t uct_rc_gdaki_ep_init_channels(uct_rc_gdaki_iface_t *iface,
                                                  uct_rc_gdaki_ep_t *ep)
{
    if (iface->ep_alloc_mode == UCT_RC_GDAKI_EP_ALLOC_MODE_POOL) {
        return uct_rc_gdaki_ep_init_channels_pooled(iface, ep);
    }

    return uct_rc_gdaki_ep_init_channels_direct(iface, ep);
}

static void uct_rc_gdaki_ep_cleanup_channels(uct_rc_gdaki_iface_t *iface,
                                             uct_rc_gdaki_ep_t *ep)
{
    if (iface->ep_alloc_mode == UCT_RC_GDAKI_EP_ALLOC_MODE_POOL) {
        uct_rc_gdaki_cleanup_channels_pooled(ep);
        return;
    }

    uct_rc_gdaki_cleanup_channels_direct(iface, ep);
}

static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_ep_t, const uct_ep_params_t *params)
{
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(params->iface,
                                                 uct_rc_gdaki_iface_t);
    ucs_status_t status;

    if (iface->cuda_ctx == NULL) {
        status = UCT_CUDADRV_FUNC_LOG_ERR(
                cuDevicePrimaryCtxRetain(&iface->cuda_ctx, iface->cuda_dev));
        if (status != UCS_OK) {
            iface->cuda_ctx = NULL;
            return status;
        }
    }

    UCS_CLASS_CALL_SUPER_INIT(uct_base_ep_t, &iface->super.super.super.super);

    self->dev_ep_init = 0;
    uct_rc_gdaki_ep_reset_channels(self);

    return uct_rc_gdaki_ep_init_channels(iface, self);
}

static UCS_CLASS_CLEANUP_FUNC(uct_rc_gdaki_ep_t)
{
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(self->super.super.iface,
                                                 uct_rc_gdaki_iface_t);
    unsigned i;

    if (self->dev_ep_init) {
        uct_rc_gdaki_channel_t *channels = self->channel_block->channels;
        for (i = 0; i < iface->num_channels; i++) {
            (void)cuMemHostUnregister(channels[i].qp.reg->addr.ptr);
        }
    }
    uct_rc_gdaki_ep_cleanup_channels(iface, self);
}

UCS_CLASS_DEFINE(uct_rc_gdaki_ep_t, uct_base_ep_t);
UCS_CLASS_DEFINE_NEW_FUNC(uct_rc_gdaki_ep_t, uct_ep_t, const uct_ep_params_t *);
UCS_CLASS_DEFINE_DELETE_FUNC(uct_rc_gdaki_ep_t, uct_ep_t);

static ucs_status_t
uct_rc_gdaki_ep_get_address(uct_ep_h tl_ep, uct_ep_addr_t *addr)
{
    uct_rc_gdaki_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_gdaki_ep_t);
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(tl_ep->iface,
                                                 uct_rc_gdaki_iface_t);
    void *rc_addr               = (void*)addr;
    unsigned i;

    for (i = 0; i < iface->num_channels; i++) {
        uct_ib_pack_uint24(*ucs_serialize_next(&rc_addr, uct_ib_uint24_t),
                           ep->channel_block->channels[i].qp.super.qp_num);
    }
    return UCS_OK;
}

static ucs_status_t uct_rc_gdaki_iface_get_address(uct_iface_h tl_iface,
                                                   uct_iface_addr_t *addr)
{
    *(uint8_t*)addr = UCT_RC_MLX5_IFACE_ADDR_TYPE_BASIC;
    return UCS_OK;
}

static ucs_status_t
uct_rc_gdaki_ep_connect_to_ep_v2(uct_ep_h tl_ep,
                                 const uct_device_addr_t *device_addr,
                                 const uct_ep_addr_t *ep_addr,
                                 const uct_ep_connect_to_ep_params_t *params)
{
    uct_rc_gdaki_ep_t *ep       = ucs_derived_of(tl_ep, uct_rc_gdaki_ep_t);
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(tl_ep->iface,
                                                 uct_rc_gdaki_iface_t);
    uct_rc_gdaki_channel_t *channels = ep->channel_block->channels;
    const uct_ib_address_t *ib_addr  = (void*)device_addr;
    uint8_t path_index               = 0;
    struct ibv_ah_attr ah_attr;
    enum ibv_mtu path_mtu;
    uint32_t dest_qp_num;
    ucs_status_t status;
    unsigned i;

    status = uct_ib_iface_fill_ah_attr_from_addr(&iface->super.super.super,
                                                 ib_addr, path_index, &ah_attr,
                                                 &path_mtu);
    if (status != UCS_OK) {
        return status;
    }

    ucs_assert(path_mtu != UCT_IB_ADDRESS_INVALID_PATH_MTU);

    for (i = 0; i < iface->num_channels; i++) {
        dest_qp_num = uct_ib_unpack_uint24(
                *ucs_serialize_next(&ep_addr, uct_ib_uint24_t));
        status      = uct_rc_mlx5_iface_common_devx_connect_qp(
                &iface->super, &channels[i].qp.super, dest_qp_num, &ah_attr,
                path_mtu, path_index, iface->super.super.config.max_rd_atomic);
        if (status != UCS_OK) {
            return status;
        }
    }

    return UCS_OK;
}

int uct_rc_gdaki_ep_is_connected(uct_ep_h tl_ep,
                                 const uct_ep_is_connected_params_t *params)
{
    uct_rc_gdaki_ep_t *ep       = ucs_derived_of(tl_ep, uct_rc_gdaki_ep_t);
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(ep->super.super.iface,
                                                 uct_rc_gdaki_iface_t);
    uint32_t addr_qp            = 0;
    uct_rc_mlx5_base_ep_address_t *rc_addr;
    ucs_status_t status;
    struct ibv_ah_attr ah_attr;
    uint32_t qp_num;
    union ibv_gid *rgid;
    const uct_ib_address_t *ib_addr;

    status = uct_ib_mlx5_query_qp_peer_info(
            &iface->super.super.super, &ep->channel_block->channels[0].qp.super,
            &ah_attr, &qp_num);
    if (status != UCS_OK) {
        return 0;
    }

    /* TODO unite code with uct_rc_mlx5_base_ep_is_connected */
    if (params->field_mask & UCT_EP_IS_CONNECTED_FIELD_EP_ADDR) {
        rc_addr = (uct_rc_mlx5_base_ep_address_t*)params->ep_addr;
        addr_qp = uct_ib_unpack_uint24(rc_addr->qp_num);
    }

    if ((addr_qp != 0) && (qp_num != addr_qp)) {
        return 0;
    }

    rgid    = (ah_attr.is_global) ? &ah_attr.grh.dgid : NULL;
    ib_addr = (const uct_ib_address_t*)params->device_addr;
    return uct_ib_iface_is_same_device(ib_addr, ah_attr.dlid, rgid);
}

static ucs_status_t
uct_rc_gdaki_iface_query(uct_iface_h tl_iface, uct_iface_attr_t *iface_attr)
{
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(tl_iface,
                                                 uct_rc_gdaki_iface_t);
    ucs_status_t status;

    status = uct_ib_iface_query(&iface->super.super.super, iface_attr);
    if (status != UCS_OK) {
        return status;
    }

    /* TODO:
     *  - add UCT_IFACE_FLAG_PUT_BATCH
     *  - PENDING and PUT_ZCOPY will be needed to establish rma_bw lanes
     *  - As this lane does not really support PUT_ZCOPY and PENDING, this could be
     *    causing issue when trying to send standard PUT. Eventually we must probably
     *    introduce another type of lane (rma_batch#x).
     */
    iface_attr->cap.flags = UCT_IFACE_FLAG_CONNECT_TO_EP |
                            UCT_IFACE_FLAG_INTER_NODE |
                            UCT_IFACE_FLAG_DEVICE_EP |
                            UCT_IFACE_FLAG_DEVICE_LKEY |
                            UCT_IFACE_FLAG_ERRHANDLE_PEER_FAILURE;

    iface_attr->ep_addr_len    = sizeof(uct_ib_uint24_t) * iface->num_channels;
    iface_attr->iface_addr_len = sizeof(uint8_t);
    iface_attr->overhead       = UCT_RC_MLX5_IFACE_OVERHEAD;

    iface_attr->cap.put.min_zcopy = 0;
    iface_attr->cap.put.max_zcopy =
            uct_ib_iface_port_attr(&iface->super.super.super)->max_msg_sz;
    iface_attr->ctl_device    = uct_cuda_get_sys_dev(iface->cuda_dev);
    iface_attr->dev_num_paths = 1;

    return UCS_OK;
}

ucs_status_t
uct_rc_gdaki_create_cq(uct_ib_iface_t *ib_iface, uct_ib_dir_t dir,
                       const uct_ib_iface_init_attr_t *init_attr,
                       int preferred_cpu, size_t inl)
{
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(ib_iface,
                                                 uct_rc_gdaki_iface_t);

    iface->super.cq[dir].type = UCT_IB_MLX5_OBJ_TYPE_NULL;
    return UCS_OK;
}

ucs_status_t
uct_rc_gdaki_ep_get_device_ep(uct_ep_h tl_ep, uct_device_ep_h *device_ep_p)
{
    uct_rc_gdaki_ep_t *ep       = ucs_derived_of(tl_ep, uct_rc_gdaki_ep_t);
    uct_rc_gdaki_iface_t *iface = ucs_derived_of(ep->super.super.iface,
                                                 uct_rc_gdaki_iface_t);
    uct_rc_gdaki_dev_ep_t *dev_ep;
    uct_rc_gdaki_channel_t *channel;
    ucs_status_t status;
    CUdeviceptr sq_db;
    unsigned i;
    size_t dev_ep_host_size;

    pthread_mutex_lock(&iface->ep_init_lock);

    if (!ep->dev_ep_init) {
        status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(iface->cuda_ctx));
        if (status != UCS_OK) {
            goto out_unlock;
        }

        dev_ep_host_size = UCT_GDAKI_DEV_EP_SIZE +
                           iface->num_channels * UCT_GDAKI_DEV_QP_SIZE;
        dev_ep           = ucs_calloc(1, dev_ep_host_size, "dev_ep");
        if (dev_ep == NULL) {
            ucs_error("failed to allocate dev_ep on host buffer size=%zu ep=%p "
                      "iface=%p",
                      dev_ep_host_size, ep, iface);
            status = UCS_ERR_NO_MEMORY;
            goto out_ctx;
        }

        status = UCT_CUDADRV_FUNC_LOG_ERR(
                cuMemsetD8((CUdeviceptr)ep->channel_block->gpu_ptr, 0,
                           dev_ep_host_size));
        if (status != UCS_OK) {
            goto out_free;
        }

        dev_ep->atomic_va    = iface->atomic_buff;
        dev_ep->atomic_lkey  = htonl(iface->atomic_mr->lkey);
        dev_ep->sq_wqe_num   = uct_ib_mlx5_devx_sq_length(
                iface->super.super.config.tx_qp_len);
        dev_ep->sq_fc_mask   = (dev_ep->sq_wqe_num >> 1) - 1;
        dev_ep->channel_mask = iface->num_channels - 1;
        dev_ep->sq_wqe_daddr = (uint8_t*)(ep->channel_block->gpu_ptr +
                                          UCT_GDAKI_DEV_EP_SIZE +
                                          iface->num_channels *
                                                  UCT_GDAKI_DEV_QP_SIZE);

        for (i = 0; i < iface->num_channels; ++i) {
            channel = &ep->channel_block->channels[i];
            (void)cuMemHostRegister(channel->qp.reg->addr.ptr,
                                    UCT_IB_MLX5_BF_REG_SIZE * 2,
                                    UCT_GDAKI_CUDA_REG_FLAGS);

            status = UCT_CUDADRV_FUNC_LOG_ERR(
                    cuMemHostGetDevicePointer(&sq_db, channel->qp.reg->addr.ptr,
                                              0));
            if (status != UCS_OK) {
                goto out_unreg;
            }

            dev_ep->qps[i].sq_db  = (uint64_t *)sq_db;
            dev_ep->qps[i].sq_num = channel->qp.super.qp_num;
            dev_ep->qps[i].qpn_ds = htonl(channel->qp.super.qp_num <<
                                          DOCA_GPUNETIO_VERBS_WQE_IDX_SHIFT);
            memset(&dev_ep->qps[i].cq_buff, 0xff, 64);
        }

        status = UCT_CUDADRV_FUNC_LOG_ERR(
                cuMemcpyHtoD((CUdeviceptr)ep->channel_block->gpu_ptr, dev_ep,
                             dev_ep_host_size));
        if (status != UCS_OK) {
            goto out_unreg;
        }

        ucs_free(dev_ep);
        (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));

        ep->dev_ep_init = 1;
    }

    *device_ep_p = &((uct_rc_gdaki_dev_ep_t*)ep->channel_block->gpu_ptr)->super;
    pthread_mutex_unlock(&iface->ep_init_lock);
    return UCS_OK;

out_unreg:
    do {
        (void)cuMemHostUnregister(
                ep->channel_block->channels[i].qp.reg->addr.ptr);
    } while (i-- > 0);
out_free:
    ucs_free(dev_ep);
out_ctx:
    (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));
out_unlock:
    pthread_mutex_unlock(&iface->ep_init_lock);
    return status;
}

static UCS_CLASS_DECLARE_NEW_FUNC(uct_rc_gdaki_iface_t, uct_iface_t, uct_md_h,
                                  uct_worker_h, const uct_iface_params_t*,
                                  const uct_iface_config_t*);

static UCS_CLASS_DECLARE_DELETE_FUNC(uct_rc_gdaki_iface_t, uct_iface_t);

static uct_rc_iface_ops_t uct_rc_gdaki_internal_ops = {
    .super = {
        .super = {
            .iface_query_v2         = uct_iface_base_query_v2,
            .iface_estimate_perf    = uct_ib_iface_estimate_perf,
            .iface_vfs_refresh      = (uct_iface_vfs_refresh_func_t)ucs_empty_function,
            .ep_query               = (uct_ep_query_func_t)ucs_empty_function_return_unsupported,
            .ep_invalidate          = (uct_ep_invalidate_func_t)ucs_empty_function_return_unsupported,
            .ep_connect_to_ep_v2    = uct_rc_gdaki_ep_connect_to_ep_v2,
            .iface_is_reachable_v2  = (uct_iface_is_reachable_v2_func_t)ucs_empty_function_return_one_int,
            .ep_is_connected        = uct_rc_gdaki_ep_is_connected,
            .ep_get_device_ep       = uct_rc_gdaki_ep_get_device_ep,
        },
        .create_cq  = uct_rc_gdaki_create_cq,
        .destroy_cq = (uct_ib_iface_destroy_cq_func_t)ucs_empty_function_return_success,
    },
    .init_rx    = (uct_rc_iface_init_rx_func_t)ucs_empty_function_return_success,
    .cleanup_rx = (uct_rc_iface_cleanup_rx_func_t)
            ucs_empty_function_return_success,
};

static uct_iface_ops_t uct_rc_gdaki_iface_tl_ops = {
    .ep_flush          = uct_base_ep_flush,
    .ep_fence          = uct_base_ep_fence,
    .ep_create         = UCS_CLASS_NEW_FUNC_NAME(uct_rc_gdaki_ep_t),
    .ep_destroy        = UCS_CLASS_DELETE_FUNC_NAME(uct_rc_gdaki_ep_t),
    .ep_get_address    = uct_rc_gdaki_ep_get_address,
    .ep_connect_to_ep  = uct_base_ep_connect_to_ep,
    .ep_pending_purge  = (uct_ep_pending_purge_func_t)ucs_empty_function,
    .iface_close       = UCS_CLASS_DELETE_FUNC_NAME(uct_rc_gdaki_iface_t),
    .iface_query       = uct_rc_gdaki_iface_query,
    .iface_get_address = uct_rc_gdaki_iface_get_address,
    .iface_get_device_address = uct_ib_iface_get_device_address,
    .iface_is_reachable       = uct_base_iface_is_reachable,
    .iface_flush              = (uct_iface_flush_func_t)
            ucs_empty_function_return_success,
    .iface_fence              = (uct_iface_fence_func_t)
            ucs_empty_function_return_unsupported,
    .iface_progress_enable    = (uct_iface_progress_enable_func_t)
            ucs_empty_function_return_unsupported,
    .iface_progress_disable   = (uct_iface_progress_disable_func_t)
            ucs_empty_function_return_unsupported,
    .iface_progress           = (uct_iface_progress_func_t)
            ucs_empty_function_return_unsupported,
};

static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_iface_t, uct_md_h tl_md,
                           uct_worker_h worker,
                           const uct_iface_params_t *params,
                           const uct_iface_config_t *tl_config)
{
    uct_rc_gdaki_iface_config_t *config =
            ucs_derived_of(tl_config, uct_rc_gdaki_iface_config_t);
    uct_ib_mlx5_md_t *md = ucs_derived_of(tl_md, uct_ib_mlx5_md_t);
    uct_ib_iface_init_attr_t init_attr = {};
    uct_md_mem_reg_params_t reg_params = {};
    UCS_STRING_BUFFER_ONSTACK(strb, 64);
    char *gpu_name, *ib_name;
    char pci_addr[UCS_SYS_BDF_NAME_MAX];
    ucs_status_t status;
    int cuda_id;
    int ret;

    if (config->num_channels > UINT8_MAX + 1) {
         ucs_error("num_channels exceeds maximum value of 256");
         return UCS_ERR_INVALID_PARAM;
    }

    self->num_channels = ucs_roundup_pow2(config->num_channels);
    self->ep_alloc_mode = config->ep_alloc_mode;

    status = uct_rc_mlx5_dp_ordering_ooo_init(md, &self->super,
                                              md->dp_ordering_cap_devx.rc,
                                              md->ddp_support_dv.rc,
                                              &config->mlx5, "rc_gda");
    if (status != UCS_OK) {
        return status;
    }

    ucs_string_buffer_appendf(&strb, "%s", params->mode.device.dev_name);
    gpu_name = ucs_string_buffer_next_token(&strb, NULL, "-");
    ib_name  = ucs_string_buffer_next_token(&strb, gpu_name, "-");

    init_attr.seg_size      = config->super.super.seg_size;
    init_attr.qp_type       = IBV_QPT_RC;
    init_attr.dev_name      = ib_name;
    init_attr.max_rd_atomic = IBV_DEV_ATTR(&md->super.dev, max_qp_rd_atom);

    UCS_CLASS_CALL_SUPER_INIT(uct_rc_mlx5_iface_common_t,
                              &uct_rc_gdaki_iface_tl_ops,
                              &uct_rc_gdaki_internal_ops, tl_md, worker, params,
                              &config->super, &config->mlx5, &init_attr);

    if (memcmp(gpu_name, UCT_DEVICE_CUDA_NAME, UCT_DEVICE_CUDA_NAME_LEN)) {
        ucs_error("wrong device name: %s\n", gpu_name);
        return status;
    }

    cuda_id = atoi(gpu_name + UCT_DEVICE_CUDA_NAME_LEN);
    status  = UCT_CUDADRV_FUNC_LOG_ERR(
            cuDeviceGetPCIBusId(pci_addr, UCS_SYS_BDF_NAME_MAX, cuda_id));
    if (status != UCS_OK) {
        return status;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&self->cuda_dev, cuda_id));
    if (status != UCS_OK) {
        return status;
    }

    self->cuda_ctx = NULL;

    ret = ucs_posix_memalign((void**)&self->atomic_buff,
                             UCS_SYS_CACHE_LINE_SIZE, sizeof(uint64_t),
                             "gdaki_atomic_buf");
    if (ret != 0) {
        ucs_error("failed to allocate AMO buffer");
        return UCS_ERR_NO_MEMORY;
    }

    status = uct_ib_reg_mr(&md->super, self->atomic_buff, sizeof(uint64_t),
                           &reg_params, UCT_IB_MEM_ACCESS_FLAGS, NULL,
                           &self->atomic_mr);
    if (status != UCS_OK) {
        goto err_atomic_buff;
    }

    if (pthread_mutex_init(&self->ep_init_lock, NULL) != 0) {
        status = UCS_ERR_IO_ERROR;
        goto err_lock;
    }

    if (self->ep_alloc_mode == UCT_RC_GDAKI_EP_ALLOC_MODE_POOL) {
        status = uct_rc_gdaki_iface_init_channel_pool(self, config);
        if (status != UCS_OK) {
            goto err_pool;
        }
    }

    return UCS_OK;

err_pool:
    pthread_mutex_destroy(&self->ep_init_lock);
err_lock:
    ibv_dereg_mr(self->atomic_mr);
err_atomic_buff:
    ucs_free(self->atomic_buff);
    return status;
}

static UCS_CLASS_CLEANUP_FUNC(uct_rc_gdaki_iface_t)
{
    pthread_mutex_destroy(&self->ep_init_lock);
    ibv_dereg_mr(self->atomic_mr);
    ucs_free(self->atomic_buff);
    if (self->ep_alloc_mode == UCT_RC_GDAKI_EP_ALLOC_MODE_POOL) {
        uct_rc_gdaki_iface_cleanup_channel_pool(self);
    }

    if (self->cuda_ctx != NULL) {
        (void)UCT_CUDADRV_FUNC_LOG_WARN(
                cuDevicePrimaryCtxRelease(self->cuda_dev));
    }
}

UCS_CLASS_DEFINE(uct_rc_gdaki_iface_t, uct_rc_mlx5_iface_common_t);

static UCS_CLASS_DEFINE_NEW_FUNC(uct_rc_gdaki_iface_t, uct_iface_t, uct_md_h,
                                 uct_worker_h, const uct_iface_params_t*,
                                 const uct_iface_config_t*);

static UCS_CLASS_DEFINE_DELETE_FUNC(uct_rc_gdaki_iface_t, uct_iface_t);

typedef struct {
    unsigned               index;
    ucs_sys_dev_distance_t dist;
    int                    usecount;
} uct_gdaki_dev_score_t;

typedef struct {
    ucs_sys_device_t sys_dev;
    uint64_t         cuda_map;
    int              direct_nic;
} uct_gdaki_dev_matrix_elem_t;

typedef struct {
    ucs_sys_device_t sys_dev;
    ucs_sys_bus_id_t bus_id;
    int              cuda_idx; /* CUDA driver index, -1 if not CUDA-visible */
} uct_gdaki_gpu_info_t;

static const uct_gdaki_gpu_info_t *
uct_gdaki_gpu_bus_id_lookup(const uct_gdaki_gpu_info_t *gpus, unsigned count,
                            const ucs_sys_bus_id_t *bus_id)
{
    ucs_bus_id_bit_rep_t bus_id_key = ucs_topo_get_bus_id_bit_repr(bus_id);
    unsigned gpu_idx;

    for (gpu_idx = 0; gpu_idx < count; gpu_idx++) {
        if (bus_id_key ==
            ucs_topo_get_bus_id_bit_repr(&gpus[gpu_idx].bus_id)) {
            return &gpus[gpu_idx];
        }
    }

    return NULL;
}

static int uct_gdaki_dev_matrix_score(const void *pa, const void *pb, void *arg)
{
    const uct_gdaki_dev_score_t *a = pa;
    const uct_gdaki_dev_score_t *b = pb;

    /* Prefer lower latency device, and if same latency prefer the less utilized */
    return (ucs_fp_compare(a->dist.latency, b->dist.latency) *
            UCT_GDAKI_MAX_CUDA_DEVICES) +
           ucs_signum(a->usecount - b->usecount);
}

static ucs_status_t
uct_gdaki_get_cuda_sys_dev(int cuda_idx, ucs_sys_device_t *sys_dev_p)
{
    ucs_status_t status;
    CUdevice cuda_dev;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&cuda_dev, cuda_idx));
    if (status != UCS_OK) {
        return status;
    }

    *sys_dev_p = uct_cuda_get_sys_dev(cuda_dev);
    return UCS_OK;
}

static ucs_status_t
uct_gdaki_enum_gpus(uct_gdaki_gpu_info_t *gpus, unsigned *count_p)
{
    uct_gdaki_gpu_info_t cuda_gpus[UCT_GDAKI_MAX_CUDA_DEVICES];
    const uct_gdaki_gpu_info_t *gpu_info;
    unsigned nvml_dev_count, nvml_idx, gpu_count, cuda_gpu_count;
    int cuda_dev_count, cuda_idx;
    ucs_sys_bus_id_t bus_id;
    nvmlDevice_t nvml_dev;
    nvmlPciInfo_t nvml_pci;
    ucs_sys_device_t sys_dev;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetCount(&cuda_dev_count));
    if (status != UCS_OK) {
        return status;
    }

    ucs_assert_always(cuda_dev_count <= UCT_GDAKI_MAX_CUDA_DEVICES);

    cuda_gpu_count = 0;
    for (cuda_idx = 0; cuda_idx < cuda_dev_count; cuda_idx++) {
        status = uct_gdaki_get_cuda_sys_dev(cuda_idx, &sys_dev);
        if (status != UCS_OK) {
            return status;
        }

        status = ucs_topo_get_device_bus_id(sys_dev, &bus_id);
        if (status != UCS_OK) {
            return status;
        }

        if (uct_gdaki_gpu_bus_id_lookup(cuda_gpus, cuda_gpu_count,
                                        &bus_id) != NULL) {
            ucs_debug("skip cuda device %d with duplicate bdf "
                      "%04x:%02x:%02x.%u", cuda_idx,
                      (unsigned)bus_id.domain, (unsigned)bus_id.bus,
                      (unsigned)bus_id.slot, (unsigned)bus_id.function);
            /* Same BDF. TODO: support MLOPart. */
            continue;
        }

        cuda_gpus[cuda_gpu_count].sys_dev  = sys_dev;
        cuda_gpus[cuda_gpu_count].bus_id   = bus_id;
        cuda_gpus[cuda_gpu_count].cuda_idx = cuda_idx;
        cuda_gpu_count++;
    }

    status = UCT_CUDA_NVML_WRAP_CALL(nvmlDeviceGetCount_v2, &nvml_dev_count);
    if (status != UCS_OK) {
        ucs_diag("NVML unavailable: using legacy CUDA-only enumeration");
        memcpy(gpus, cuda_gpus, cuda_gpu_count * sizeof(*gpus));
        *count_p = cuda_gpu_count;
        return UCS_OK;
    }

    ucs_assert_always(nvml_dev_count <= UCT_GDAKI_MAX_CUDA_DEVICES);

    /* Add non-CUDA-visible devices from NVML. TODO: support MLOPart. */
    gpu_count = 0;
    for (nvml_idx = 0; nvml_idx < nvml_dev_count; nvml_idx++) {
        status = UCT_CUDA_NVML_WRAP_CALL(nvmlDeviceGetHandleByIndex, nvml_idx,
                                         &nvml_dev);
        if (status != UCS_OK) {
            return status;
        }

        status = UCT_CUDA_NVML_WRAP_CALL(nvmlDeviceGetPciInfo_v3, nvml_dev,
                                         &nvml_pci);
        if (status != UCS_OK) {
            return status;
        }

        bus_id.domain   = nvml_pci.domain;
        bus_id.bus      = nvml_pci.bus;
        bus_id.slot     = nvml_pci.device;
        bus_id.function = 0;

        gpu_info = uct_gdaki_gpu_bus_id_lookup(cuda_gpus, cuda_gpu_count,
                                               &bus_id);
        if (gpu_info != NULL) {
            gpus[gpu_count++] = *gpu_info;
            /* Already added as CUDA-visible. */
            continue;
        }

        status = ucs_topo_find_device_by_bus_id(&bus_id, &sys_dev);
        if (status != UCS_OK) {
            return status;
        }

        if (gpu_count == UCT_GDAKI_MAX_CUDA_DEVICES) {
            ucs_error("exceeded maximal number of GDAKI CUDA devices (%u)",
                      UCT_GDAKI_MAX_CUDA_DEVICES);
            return UCS_ERR_EXCEEDS_LIMIT;
        }

        gpus[gpu_count].sys_dev  = sys_dev;
        gpus[gpu_count].bus_id   = bus_id;
        gpus[gpu_count].cuda_idx = -1;
        gpu_count++;
    }

    *count_p = gpu_count;
    return UCS_OK;
}

uct_gdaki_dev_matrix_elem_t *
uct_gdaki_dev_matrix_init(const uct_ib_md_t *ib_md, size_t *dmat_length_p)
{
    unsigned long ib_per_cuda         = ib_md->config.gda_max_hca_per_gpu;
    uct_gdaki_dev_matrix_elem_t *dmat = NULL;
    uct_gdaki_gpu_info_t gpus[UCT_GDAKI_MAX_CUDA_DEVICES];
    ucs_status_t status;
    int ibdev_index, ibdev_count;
    unsigned gpu_count, gpu_index;
    struct ibv_device **device_list;
    struct ibv_device *ibdev;
    uct_gdaki_dev_score_t *scores;
    uct_gdaki_dev_matrix_elem_t *ibdesc;
    struct ibv_context *context;

    /* Obtain the list of IB devices */
    device_list = ibv_get_device_list(&ibdev_count);
    if (device_list == NULL) {
        return NULL;
    }

    ucs_assert(ibdev_count > 0);
    /* Allocate memory for device matrix */
    dmat = ucs_calloc(ibdev_count, sizeof(*dmat), "dev matrix");
    if (dmat == NULL) {
        status = UCS_ERR_NO_MEMORY;
        goto out_dev;
    }

    /* Allocate memory for device scores */
    scores = ucs_calloc(ibdev_count, sizeof(*scores), "dev scores");
    if (scores == NULL) {
        status = UCS_ERR_NO_MEMORY;
        goto out_dmat;
    }

    /* Query each IB device so its sys_dev is resolved and its port state is
     * known before direct-NIC sibling matching can be triggered. */
    for (ibdev_index = 0; ibdev_index < ibdev_count; ibdev_index++) {
        uct_ib_device_t tmp_dev = {0};

        ibdesc  = &dmat[ibdev_index];
        ibdev   = device_list[ibdev_index];
        context = ibv_open_device(ibdev);
        if (context == NULL) {
            ucs_error("ibv_open_device(%s) failed: %m",
                      ibv_get_device_name(ibdev));
            status = UCS_ERR_IO_ERROR;
            goto out;
        }

        tmp_dev.ibv_context = context;
        status              = uct_ib_device_query(&tmp_dev, ibdev);
        if (status != UCS_OK) {
            ibv_close_device(context);
            goto out;
        }

        ibdesc->sys_dev    = tmp_dev.sys_dev;
        ibdesc->direct_nic = uct_ib_mlx5dv_check_direct_nic(&tmp_dev, 1) !=
                             UCS_SYS_DEVICE_ID_UNKNOWN;
        scores[ibdev_index].index = ibdev_index;
        ibv_close_device(context);
    }

    /* Enumerate physical GPUs via NVML (or CUDA fallback). NVML is used here
     * so that CUDA_VISIBLE_DEVICES does not hide GPUs from the topology
     * calculation */
    status = uct_gdaki_enum_gpus(gpus, &gpu_count);
    if (status != UCS_OK) {
        goto out;
    }

    if (gpu_count == 0) {
        goto out;
    }

    if (ib_per_cuda == UCS_ULUNITS_AUTO) {
        ib_per_cuda = ibdev_count / gpu_count;
    }

    /* Map each GPU to the best suited IB devices */
    for (gpu_index = 0; gpu_index < gpu_count; gpu_index++) {
        /* Update PCI distance in IB device scores */
        for (ibdev_index = 0; ibdev_index < ibdev_count; ibdev_index++) {
            ibdesc = &dmat[scores[ibdev_index].index];
            status = ucs_topo_get_distance(gpus[gpu_index].sys_dev,
                                           ibdesc->sys_dev,
                                           &scores[ibdev_index].dist);
            if (status != UCS_OK) {
                goto out;
            }
        }

        /* Sort and select the best suited IB devices for this GPU */
        ucs_qsort_r(scores, ibdev_count, sizeof(*scores),
                    uct_gdaki_dev_matrix_score, NULL);

        for (ibdev_index = 0; ibdev_index < ib_per_cuda; ibdev_index++) {
            ibdesc = &dmat[scores[ibdev_index].index];
            if (gpus[gpu_index].cuda_idx >= 0) {
                ibdesc->cuda_map |= UCS_BIT(gpus[gpu_index].cuda_idx);
            }
            scores[ibdev_index].usecount++;

            /* direct NIC is closest and can be only one */
            if (ibdesc->direct_nic) {
                ucs_assert_always(ibdev_index == 0);
                break;
            }
        }
    }

    /* Output processed device matrix length */
    *dmat_length_p = ibdev_count;

out:
    ucs_free(scores);
out_dmat:
    if (status != UCS_OK) {
        ucs_free(dmat);
        dmat = NULL;
    }
out_dev:
    ibv_free_device_list(device_list);
    return dmat;
}

static ucs_status_t
uct_gdaki_query_tl_devices(uct_md_h tl_md,
                           uct_tl_device_resource_t **tl_devices_p,
                           unsigned *num_tl_devices_p)
{
    uct_ib_mlx5_md_t *ib_mlx5_md     = ucs_derived_of(tl_md, uct_ib_mlx5_md_t);
    uct_ib_md_t *ib_md               = &ib_mlx5_md->super;
    static ucs_init_once_t dmat_once = UCS_INIT_ONCE_INITIALIZER;
    static uct_gdaki_dev_matrix_elem_t *dmat;
    static size_t dmat_length;
    unsigned num_tl_devices;
    uct_tl_device_resource_t *tl_devices;
    ucs_status_t status;
    CUdevice device;
    int i;
    uct_gdaki_dev_matrix_elem_t *ibdesc;

    if (!(uct_gdaki_get_driver_features(ib_md) & UCT_GDAKI_SUPPORTED)) {
        ucs_debug("%s: GDAKI is not supported",
                  uct_ib_device_name(&ib_md->dev));
        status = UCS_ERR_NO_DEVICE;
        goto out;
    }

    UCS_INIT_ONCE(&dmat_once) {
        dmat = uct_gdaki_dev_matrix_init(ib_md, &dmat_length);
    }

    if (dmat == NULL) {
        ucs_debug("%s: global device matrix initialization failed",
                  uct_ib_device_name(&ib_md->dev));
        status = UCS_ERR_NO_DEVICE;
        goto out;
    }

    for (ibdesc = dmat; ibdesc - dmat < dmat_length; ibdesc++) {
        if (ibdesc->sys_dev == ib_md->dev.sys_dev) {
            break;
        }
    }

    ucs_assertv(ibdesc - dmat < dmat_length, "dev %s",
                uct_ib_device_name(&ib_md->dev));

    if (ibdesc->cuda_map == 0) {
        ucs_debug("%s: no assigned gpu found", uct_ib_device_name(&ib_md->dev));
        status = UCS_ERR_NO_DEVICE;
        goto out;
    }

    tl_devices = ucs_malloc(sizeof(*tl_devices) *
                                    ucs_popcount(ibdesc->cuda_map),
                            "gdaki_tl_devices");
    if (tl_devices == NULL) {
        status = UCS_ERR_NO_MEMORY;
        goto out;
    }

    num_tl_devices = 0;
    ucs_for_each_bit(i, ibdesc->cuda_map) {
        status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&device, i));
        if (status != UCS_OK) {
            goto err;
        }

        snprintf(tl_devices[num_tl_devices].name,
                 sizeof(tl_devices[num_tl_devices].name), "%s%d-%s:%d",
                 UCT_DEVICE_CUDA_NAME, device, uct_ib_device_name(&ib_md->dev),
                 ib_md->dev.first_port);
        tl_devices[num_tl_devices].type       = UCT_DEVICE_TYPE_NET;
        tl_devices[num_tl_devices].sys_device = ib_md->dev.sys_dev;
        num_tl_devices++;
    }

    *num_tl_devices_p = num_tl_devices;
    *tl_devices_p     = tl_devices;
    return UCS_OK;

err:
    ucs_free(tl_devices);
out:
    return status;
}

UCT_TL_DEFINE_ENTRY(&uct_ib_component, rc_gda, uct_gdaki_query_tl_devices,
                    uct_rc_gdaki_iface_t, "RC_GDA_",
                    uct_rc_gdaki_iface_config_table,
                    uct_rc_gdaki_iface_config_t);

UCT_TL_INIT(&uct_ib_component, rc_gda, ctor, , )
