UCT/GDAKI: Export EP to GPU (#10833)

Artemy-Mellanox · web-flow · commit 34317a76c8b3 · 2025-09-02T07:17:52.000Z
diff --git a/src/uct/ib/mlx5/dv/ib_mlx5_dv.c b/src/uct/ib/mlx5/dv/ib_mlx5_dv.c
@@ -197,12 +197,14 @@ ucs_status_t uct_ib_mlx5_devx_create_qp_common(uct_ib_iface_t *iface,
     UCT_IB_MLX5DV_SET(qpc, qpc, user_index, attr->uidx);
     UCT_IB_MLX5DV_SET(qpc, qpc, ts_format, UCT_IB_MLX5_QPC_TS_FORMAT_DEFAULT);
 
-    if (qp->devx.wq_buf == NULL) {
+    if ((qp->devx.wq_buf == NULL) && (attr->umem_offset == 0)) {
         UCT_IB_MLX5DV_SET(qpc, qpc, no_sq, true);
         UCT_IB_MLX5DV_SET(qpc, qpc, offload_type, true);
         UCT_IB_MLX5DV_SET(create_qp_in, in, wq_umem_id, md->zero_mem.mem->umem_id);
     } else {
         UCT_IB_MLX5DV_SET(create_qp_in, in, wq_umem_id, qp->devx.mem.mem->umem_id);
+        UCT_IB_MLX5DV_SET64(create_qp_in, in, wq_umem_offset,
+                            attr->umem_offset);
     }
 
     if (md->super.ece_enable) {
@@ -248,13 +250,11 @@ ucs_status_t uct_ib_mlx5_devx_create_qp_common(uct_ib_iface_t *iface,
     attr->super.cap.max_recv_wr = 0;
 
     if (tx != NULL) {
-        ucs_assert(qp->devx.wq_buf != NULL);
         tx->reg    = &uar->super;
         tx->qstart = qp->devx.wq_buf;
         tx->qend   = UCS_PTR_BYTE_OFFSET(qp->devx.wq_buf, attr->len);
         tx->dbrec  = &qp->devx.dbrec->db[MLX5_SND_DBR];
         tx->bb_max = attr->max_tx - 2 * UCT_IB_MLX5_MAX_BB;
-        ucs_assert(*tx->dbrec == 0);
         uct_ib_mlx5_txwq_reset(tx);
     } else {
         ucs_assert(qp->devx.wq_buf == NULL);
@@ -602,7 +602,7 @@ uct_ib_mlx5_devx_create_cq_common(uct_ib_iface_t *iface, uct_ib_dir_t dir,
 
     /* Set CQ umem related bits */
     UCT_IB_MLX5DV_SET(create_cq_in, in, cq_umem_id, cq->devx.mem.mem->umem_id);
-    UCT_IB_MLX5DV_SET64(create_cq_in, in, cq_umem_offset, 0);
+    UCT_IB_MLX5DV_SET64(create_cq_in, in, cq_umem_offset, attr->umem_offset);
 
     UCT_IB_MLX5DV_SET(cqc, cqctx, log_cq_size, log_cq_size);
     UCT_IB_MLX5DV_SET(cqc, cqctx, cqe_sz, (attr->cqe_size == 128) ? 1 : 0);
diff --git a/src/uct/ib/mlx5/dv/ib_mlx5_ifc.h b/src/uct/ib/mlx5/dv/ib_mlx5_ifc.h
@@ -1563,7 +1563,7 @@ struct uct_ib_mlx5_create_qp_in_bits {
 
     struct uct_ib_mlx5_qpc_bits qpc;
 
-    uint8_t         reserved_at_800[0x40];
+    uint8_t         wq_umem_offset[0x40];
 
     uint8_t         wq_umem_id[0x20];
 
diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c
@@ -16,6 +16,11 @@
 #include <uct/cuda/base/cuda_iface.h>
 
 #include <doca_log.h>
+#include <cuda.h>
+
+
+#define UCT_GDAKI_DOCA_NOTUSE    1
+#define UCT_GDAKI_DOCA_NOTUSEPTR (void*)1
 
 typedef struct {
     uct_rc_iface_common_config_t      super;
@@ -42,48 +47,181 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_ep_t, const uct_ep_params_t *params)
     uct_ib_mlx5_md_t *md = ucs_derived_of(iface->super.super.super.super.md,
                                           uct_ib_mlx5_md_t);
     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_rc_gdaki_dev_ep_t dev_ep       = {};
     ucs_status_t status;
+    doca_error_t derr;
+    size_t dev_ep_size;
+    uct_ib_mlx5_dbrec_t dbrec;
+    CUcontext primary_ctx;
 
     UCS_CLASS_CALL_SUPER_INIT(uct_base_ep_t, &iface->super.super.super.super);
 
-    init_attr.cq_len[UCT_IB_DIR_TX] = iface->super.super.config.tx_qp_len *
-                                      UCT_IB_MLX5_MAX_BB;
-    init_attr.flags                 = UCT_IB_CQ_IGNORE_OVERRUN;
-    status = uct_ib_mlx5_devx_create_cq(&iface->super.super.super,
-                                        UCT_IB_DIR_TX, &init_attr, &self->cq, 0,
-                                        0);
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuDevicePrimaryCtxRetain(&primary_ctx, iface->cuda_dev));
     if (status != UCS_OK) {
         return status;
     }
 
+    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(primary_ctx));
+    if (status != UCS_OK) {
+        goto err_ctx;
+    }
+
+    init_attr.cq_len[UCT_IB_DIR_TX] = iface->super.super.config.tx_qp_len *
+                                      UCT_IB_MLX5_MAX_BB;
+    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);
+    uct_ib_mlx5_wq_calc_sizes(&qp_attr);
+
+    cq_attr.flags      |= UCT_IB_MLX5_CQ_IGNORE_OVERRUN;
+    cq_attr.umem_offset = ucs_align_up_pow2(
+            sizeof(uct_rc_gdaki_dev_ep_t) +
+                    qp_attr.max_tx * sizeof(uct_rc_gdaki_op_t),
+            ucs_get_page_size());
+
     qp_attr.mmio_mode     = UCT_IB_MLX5_MMIO_MODE_DB;
     qp_attr.super.srq_num = 0;
-    status = uct_ib_mlx5_devx_create_qp(&iface->super.super.super, &self->cq,
-                                        &self->cq, &self->qp.super, &self->qp,
-                                        &qp_attr);
+    qp_attr.umem_offset   = ucs_align_up_pow2(cq_attr.umem_offset +
+                                                      cq_attr.umem_len,
+                                              ucs_get_page_size());
+
+    dev_ep_size = qp_attr.umem_offset + qp_attr.len;
+    /*
+     * dev_ep layout:
+     * +---------------------+-------+---------+---------+
+     * | counters, dbr       | ops   | cq buff | wq buff |
+     * +---------------------+-------+---------+---------+
+     */
+    derr = doca_gpu_mem_alloc(iface->gpu_dev, dev_ep_size, ucs_get_page_size(),
+                              DOCA_GPU_MEM_TYPE_GPU, (void**)&self->ep_gpu,
+                              NULL);
+    if (derr != DOCA_SUCCESS) {
+        ucs_error("doca_gpu_mem_alloc failed: %s", doca_error_get_descr(derr));
+        status = UCS_ERR_IO_ERROR;
+        goto err_ctx;
+    }
+
+    /* TODO add dmabuf_fd support */
+    self->umem = mlx5dv_devx_umem_reg(md->super.dev.ibv_context, self->ep_gpu,
+                                      dev_ep_size, IBV_ACCESS_LOCAL_WRITE);
+    if (self->umem == NULL) {
+        uct_ib_check_memlock_limit_msg(md->super.dev.ibv_context,
+                                       UCS_LOG_LEVEL_ERROR,
+                                       "mlx5dv_devx_umem_reg(size=%zu)",
+                                       dev_ep_size);
+        status = UCS_ERR_NO_MEMORY;
+        goto err_mem;
+    }
+
+    self->cq.devx.mem.mem       = self->umem;
+    self->qp.super.devx.mem.mem = self->umem;
+
+    dbrec.mem_id        = self->umem->umem_id;
+    dbrec.offset        = ucs_offsetof(uct_rc_gdaki_dev_ep_t, cq_dbrec);
+    self->cq.devx.dbrec = &dbrec;
+    status = uct_ib_mlx5_devx_create_cq_common(&iface->super.super.super,
+                                               UCT_IB_DIR_TX, &cq_attr,
+                                               &self->cq, 0, 0);
+    if (status != UCS_OK) {
+        goto err_umem;
+    }
+
+    dbrec.offset              = ucs_offsetof(uct_rc_gdaki_dev_ep_t, qp_dbrec);
+    self->qp.super.devx.dbrec = &dbrec;
+    status = uct_ib_mlx5_devx_create_qp_common(&iface->super.super.super,
+                                               &self->cq, &self->cq,
+                                               &self->qp.super, &self->qp,
+                                               &qp_attr);
     if (status != UCS_OK) {
+        goto err_cq;
+    }
+
+    derr = doca_gpu_verbs_bridge_export_qp(
+            iface->gpu_dev, self->qp.super.qp_num,
+            UCS_PTR_BYTE_OFFSET(self->ep_gpu, qp_attr.umem_offset),
+            qp_attr.max_tx, self->ep_gpu->qp_dbrec, self->qp.reg->addr.ptr,
+            UCT_IB_MLX5_BF_REG_SIZE * 2, self->cq.cq_num,
+            UCS_PTR_BYTE_OFFSET(self->ep_gpu, cq_attr.umem_offset),
+            cq_attr.cq_size, self->ep_gpu->cq_dbrec, UCT_GDAKI_DOCA_NOTUSE,
+            UCT_GDAKI_DOCA_NOTUSEPTR, UCT_GDAKI_DOCA_NOTUSE,
+            UCT_GDAKI_DOCA_NOTUSEPTR, UCT_GDAKI_DOCA_NOTUSE,
+            UCT_GDAKI_DOCA_NOTUSE, UCT_GDAKI_DOCA_NOTUSEPTR,
+            UCT_GDAKI_DOCA_NOTUSE, UCT_GDAKI_DOCA_NOTUSEPTR, 0, &self->qp_cpu);
+    if (derr != DOCA_SUCCESS) {
+        ucs_error("doca_gpu_verbs_bridge_export_qp failed: %s",
+                  doca_error_get_descr(derr));
+        status = UCS_ERR_INVALID_PARAM;
         goto err_qp;
     }
 
+    derr = doca_gpu_verbs_get_qp_dev(self->qp_cpu, &self->qp_gpu);
+    if (derr != DOCA_SUCCESS) {
+        status = UCS_ERR_INVALID_PARAM;
+        goto err_dev_ep;
+    }
+
+    dev_ep.qp = self->qp_gpu;
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuMemsetD8((CUdeviceptr)self->ep_gpu, 0, dev_ep_size));
+    if (status != UCS_OK) {
+        goto err_dev_ep;
+    }
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuMemsetD8((CUdeviceptr)UCS_PTR_BYTE_OFFSET(self->ep_gpu,
+                                                        cq_attr.umem_offset),
+                       0xff, cq_attr.umem_len));
+    if (status != UCS_OK) {
+        goto err_dev_ep;
+    }
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuMemcpyHtoD((CUdeviceptr)self->ep_gpu, &dev_ep, sizeof(dev_ep)));
+    if (status != UCS_OK) {
+        goto err_dev_ep;
+    }
+
+    (void)UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPopCurrent(NULL));
     return UCS_OK;
 
+err_dev_ep:
+    doca_gpu_verbs_unexport_qp(iface->gpu_dev, self->qp_cpu);
 err_qp:
-    uct_ib_mlx5_devx_destroy_cq(md, &self->cq);
+    uct_ib_mlx5_devx_destroy_qp_common(&self->qp.super);
+err_cq:
+    uct_ib_mlx5_devx_destroy_cq_common(&self->cq);
+err_umem:
+    mlx5dv_devx_umem_dereg(self->umem);
+err_mem:
+    doca_gpu_mem_free(iface->gpu_dev, self->ep_gpu);
+err_ctx:
+    (void)UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPopCurrent(NULL));
+    (void)UCT_CUDADRV_FUNC_LOG_ERR(cuDevicePrimaryCtxRelease(iface->cuda_dev));
     return status;
 }
 
 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);
-    uct_ib_mlx5_md_t *md = ucs_derived_of(iface->super.super.super.super.md,
-                                          uct_ib_mlx5_md_t);
+    doca_error_t derr;
+
+    derr = doca_gpu_verbs_unexport_qp(iface->gpu_dev, self->qp_cpu);
+    if (derr != DOCA_SUCCESS) {
+        ucs_error("doca_gpu_rdma_verbs_unexport_qp failed: %s",
+                  doca_error_get_descr(derr));
+    }
 
-    uct_ib_mlx5_devx_destroy_qp(md, &self->qp.super);
-    uct_ib_mlx5_devx_destroy_cq(md, &self->cq);
+    uct_ib_mlx5_devx_destroy_qp_common(&self->qp.super);
+    uct_ib_mlx5_devx_destroy_cq_common(&self->cq);
+    mlx5dv_devx_umem_dereg(self->umem);
+    doca_gpu_mem_free(iface->gpu_dev, self->ep_gpu);
+    (void)UCT_CUDADRV_FUNC_LOG_ERR(cuDevicePrimaryCtxRelease(iface->cuda_dev));
 }
 
 UCS_CLASS_DEFINE(uct_rc_gdaki_ep_t, uct_base_ep_t);
diff --git a/src/uct/ib/mlx5/gdaki/gdaki.h b/src/uct/ib/mlx5/gdaki/gdaki.h
@@ -3,15 +3,16 @@
  * See file LICENSE for terms.
  */
 
-#ifndef UCT_GDAKI_IFACE_H
-#define UCT_GDAKI_IFACE_H
+#ifndef UCT_GDAKI_H
+#define UCT_GDAKI_H
 
 #include <uct/ib/mlx5/rc/rc_mlx5_common.h>
 #include <uct/base/uct_iface.h>
 
 #include <cuda.h>
 #include <doca_gpunetio.h>
 
+#include "gdaki_dev.h"
 
 #define UCT_DEV_CUDA_NAME "cuda"
 #define UCT_DEV_CUDA_NAME_LEN 4
@@ -25,9 +26,13 @@ typedef struct uct_rc_gdaki_iface {
 
 
 typedef struct uct_rc_gdaki_ep {
-    uct_base_ep_t      super;
-    uct_ib_mlx5_cq_t   cq;
-    uct_ib_mlx5_txwq_t qp;
+    uct_base_ep_t                super;
+    uct_ib_mlx5_cq_t             cq;
+    uct_ib_mlx5_txwq_t           qp;
+    struct mlx5dv_devx_umem      *umem;
+    struct doca_gpu_verbs_qp     *qp_cpu;
+    struct doca_gpu_dev_verbs_qp *qp_gpu;
+    uct_rc_gdaki_dev_ep_t        *ep_gpu;
 } uct_rc_gdaki_ep_t;
 
 #endif /* UCT_GDAKI_IFACE_H */
diff --git a/src/uct/ib/mlx5/gdaki/gdaki_dev.h b/src/uct/ib/mlx5/gdaki/gdaki_dev.h
@@ -0,0 +1,21 @@
+/**
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
+ * See file LICENSE for terms.
+ */
+
+#ifndef UCT_GDAKI_DEV_H
+#define UCT_GDAKI_DEV_H
+
+typedef struct {
+    /* TODO add uct completion */
+} uct_rc_gdaki_op_t;
+
+
+typedef struct {
+    struct doca_gpu_dev_verbs_qp *qp;
+    uint32_t                     cq_dbrec[2];
+    uint32_t                     qp_dbrec[2];
+    uct_rc_gdaki_op_t            ops[0];
+} uct_rc_gdaki_dev_ep_t;
+
+#endif /* UCT_GDAKI_DEV_H */
diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h
@@ -585,6 +585,7 @@ typedef struct uct_ib_mlx5_cq_attr {
     unsigned cq_size;
     unsigned cqe_size;
     size_t   umem_len;
+    size_t   umem_offset;
     unsigned flags;
 } uct_ib_mlx5_cq_attr_t;
 
@@ -643,6 +644,7 @@ typedef struct uct_ib_mlx5_qp_attr {
     uint8_t                     log_num_dci_stream_channels;
     unsigned                    max_tx;
     unsigned                    len;
+    size_t                      umem_offset;
 } uct_ib_mlx5_qp_attr_t;
 
 
