TL/MLX5: addressing sam's comments on PR 989

src/components/tl/mlx5/mcast/tl_mlx5_mcast.h

@@ -17,7 +17,7 @@
 #include "components/tl/ucc_tl_log.h"
 #include "utils/ucc_rcache.h"
 #include "core/ucc_service_coll.h"
-#include "utils/arch/cuda_def.h"
+#include "components/mc/ucc_mc.h"
 
 #define POLL_PACKED       16
 #define REL_DONE          ((void*)-1)
@@ -91,7 +91,7 @@ typedef struct mcast_coll_comm_init_spec {
     int                               scq_moderation;
     int                               wsize;
     int                               max_eager;
-    int                               device_mem_enabled;
+    int                               cuda_mem_enabled;
     void                             *oob;
 } ucc_tl_mlx5_mcast_coll_comm_init_spec_t;
 
@@ -196,17 +196,19 @@ typedef struct ucc_tl_mlx5_mcast_coll_comm {
     ucc_rank_t                              rank;
     ucc_rank_t                              commsize;
     char                                   *grh_buf;
+    ucc_mc_buffer_header_t                 *grh_cuda_header;
     struct ibv_mr                          *grh_mr;
     uint16_t                                mcast_lid;
     union ibv_gid                           mgid;
     unsigned                                max_inline;
     size_t                                  max_eager;
-    int                                     device_mem_enabled;
+    int                                     cuda_mem_enabled;
     int                                     max_per_packet;
     int                                     pending_send;
     int                                     pending_recv;
     struct ibv_mr                          *pp_mr;
     char                                   *pp_buf;
+    ucc_mc_buffer_header_t                 *pp_cuda_header;
     struct pp_packet                       *pp;
     uint32_t                                psn;
     uint32_t                                last_psn;

src/components/tl/mlx5/mcast/tl_mlx5_mcast_coll.c

@@ -283,6 +283,23 @@ void ucc_tl_mlx5_mcast_collective_progress(ucc_coll_task_t *coll_task)
     }
 }
 
+ucc_status_t ucc_tl_mlx5_mcast_check_memory_type_cap(ucc_base_coll_args_t *coll_args,
+                                                     ucc_base_team_t *team)
+{
+    ucc_tl_mlx5_team_t            *mlx5_team = ucc_derived_of(team, ucc_tl_mlx5_team_t);
+    ucc_tl_mlx5_mcast_coll_comm_t *comm      = mlx5_team->mcast->mcast_comm;
+    ucc_coll_args_t               *args      = &coll_args->args;
+
+    if ((comm->cuda_mem_enabled &&
+            args->src.info.mem_type == UCC_MEMORY_TYPE_CUDA) ||
+        (!comm->cuda_mem_enabled &&
+                args->src.info.mem_type == UCC_MEMORY_TYPE_HOST)) {
+        return UCC_OK;
+    }
+
+    return UCC_ERR_NO_RESOURCE;
+}
+
 ucc_status_t ucc_tl_mlx5_mcast_bcast_init(ucc_tl_mlx5_task_t *task)
 {
     task->super.post     = ucc_tl_mlx5_mcast_bcast_start;

src/components/tl/mlx5/mcast/tl_mlx5_mcast_coll.h

@@ -14,4 +14,6 @@ ucc_status_t ucc_tl_mlx5_mcast_bcast_init(ucc_tl_mlx5_task_t *task);
 
 ucc_status_t ucc_tl_mlx5_mcast_test(ucc_tl_mlx5_mcast_coll_req_t* _req);
 
+ucc_status_t ucc_tl_mlx5_mcast_check_memory_type_cap(ucc_base_coll_args_t *coll_args,
+                                                     ucc_base_team_t *team);
 #endif

src/components/tl/mlx5/mcast/tl_mlx5_mcast_helper.c

@@ -300,7 +300,7 @@ ucc_status_t ucc_tl_mlx5_mcast_init_qps(ucc_tl_mlx5_mcast_coll_context_t *ctx,
         return UCC_ERR_NO_RESOURCE;
     }
 
-    if (comm->device_mem_enabled) {
+    if (comm->cuda_mem_enabled) {
         /* max inline send otherwise it segfault during ibv send */
         comm->max_inline = 0;
     } else {
@@ -482,8 +482,8 @@ ucc_status_t ucc_tl_mlx5_clean_mcast_comm(ucc_tl_mlx5_mcast_coll_comm_t *comm)
     }
 
     if (comm->grh_buf) {
-        if (comm->device_mem_enabled) {
-            cudaFree(comm->grh_buf);
+        if (comm->cuda_mem_enabled) {
+            ucc_mc_free(comm->grh_cuda_header);
         } else {
             ucc_free(comm->grh_buf);
         }
@@ -502,8 +502,8 @@ ucc_status_t ucc_tl_mlx5_clean_mcast_comm(ucc_tl_mlx5_mcast_coll_comm_t *comm)
     }
 
     if (comm->pp_buf) {
-        if (comm->device_mem_enabled) {
-            cudaFree(comm->pp_buf);
+        if (comm->cuda_mem_enabled) {
+            ucc_mc_free(comm->pp_cuda_header);
         } else {
             ucc_free(comm->pp_buf);
         }

src/components/tl/mlx5/mcast/tl_mlx5_mcast_helper.h

@@ -75,9 +75,13 @@ static inline ucc_status_t ucc_tl_mlx5_mcast_send(ucc_tl_mlx5_mcast_coll_comm_t
         if (zcopy) {
             pp->context = (uintptr_t) PTR_OFFSET(req->ptr, offset);
         } else {
-            if (comm->device_mem_enabled) {
-                CUDA_FUNC(cudaMemcpy((void*) pp->buf, PTR_OFFSET(req->ptr, offset),
-                                     length, cudaMemcpyDeviceToDevice));
+            if (comm->cuda_mem_enabled) {
+                status = ucc_mc_memcpy((void*) pp->buf, PTR_OFFSET(req->ptr, offset), length,
+                                       UCC_MEMORY_TYPE_CUDA, UCC_MEMORY_TYPE_CUDA);
+                if (ucc_unlikely(status != UCC_OK)) {
+                    tl_error(comm->lib, "failed to copy cuda buffer");
+                    return status;
+                }
             } else {
                 memcpy((void*) pp->buf, PTR_OFFSET(req->ptr, offset), length);
             }

src/components/tl/mlx5/mcast/tl_mlx5_mcast_progress.c

@@ -379,8 +379,13 @@ ucc_status_t ucc_tl_mlx5_mcast_process_packet(ucc_tl_mlx5_mcast_coll_comm_t *com
 
     if (pp->length > 0 ) {
         dest = req->ptr + PSN_TO_RECV_OFFSET(pp->psn, req, comm);
-        if (comm->device_mem_enabled) {
-            cudaMemcpy(dest, (void*) pp->buf, pp->length, cudaMemcpyDeviceToDevice);
+        if (comm->cuda_mem_enabled) {
+            status = ucc_mc_memcpy(dest, (void*) pp->buf, pp->length,
+                                   UCC_MEMORY_TYPE_CUDA, UCC_MEMORY_TYPE_CUDA);
+            if (ucc_unlikely(status != UCC_OK)) {
+                tl_error(comm->lib, "failed to copy cuda buffer");
+                return status;
+            }
         } else {
             memcpy(dest, (void*) pp->buf, pp->length);
         }

src/components/tl/mlx5/mcast/tl_mlx5_mcast_team.c

@@ -12,6 +12,18 @@
 #include "p2p/ucc_tl_mlx5_mcast_p2p.h"
 #include "mcast/tl_mlx5_mcast_helper.h"
 
+static ucc_status_t ucc_tl_mlx5_check_gpudirect_driver()
+{
+    ucc_status_t status = UCC_ERR_NO_RESOURCE;
+    const char  *file   = "/sys/kernel/mm/memory_peers/nv_mem/version";
+
+    if (!access(file, F_OK)) {
+        status = UCC_OK;
+    }
+
+    return status;
+}
+
 static ucc_status_t ucc_tl_mlx5_mcast_service_bcast_post(void *arg, void *buf, size_t size, ucc_rank_t root,
                                                          ucc_service_coll_req_t **bcast_req)
 {
@@ -120,10 +132,16 @@ ucc_status_t ucc_tl_mlx5_mcast_team_init(ucc_base_context_t *base_context,
 
     comm->wsize              = conf_params->wsize;
     comm->max_eager          = conf_params->max_eager;
-    comm->device_mem_enabled = conf_params->device_mem_enabled;
+    comm->cuda_mem_enabled   = conf_params->cuda_mem_enabled;
     comm->comm_id            = team_params->id;
     comm->ctx                = mcast_context;
 
+    if (comm->cuda_mem_enabled && (UCC_OK != ucc_tl_mlx5_check_gpudirect_driver())) {
+        tl_warn(mcast_context->lib, "cuda-aware mcast not available as gpu direct is not ready");
+        status = UCC_ERR_NO_RESOURCE;
+        goto cleanup;
+    }
+
     comm->rcq = ibv_create_cq(mcast_context->ctx, comm->params.rx_depth, NULL, NULL, 0);
     if (!comm->rcq) {
         ibv_dereg_mr(comm->grh_mr);
@@ -212,35 +230,41 @@ ucc_status_t ucc_tl_mlx5_mcast_coll_setup_comm_resources(ucc_tl_mlx5_mcast_coll_
     comm->pending_recv = 0;
     comm->buf_n        = comm->params.rx_depth * 2;
 
-    if (comm->device_mem_enabled) {
-        /* TODO add check to make sure GPUDirect is enabled
-         * lsmod | grep nv_peer */
-        CUDA_FUNC(cudaMalloc((void **)&comm->grh_buf, GRH_LENGTH * sizeof(char)));
-        if (!comm->grh_buf) {
-            tl_error(comm->ctx->lib, "cuda memcpy failed");
-            status = UCC_ERR_NO_MEMORY;
-            goto error;
+    if (comm->cuda_mem_enabled) {
+        status = ucc_mc_alloc(&comm->grh_cuda_header, GRH_LENGTH *
+                              sizeof(char), UCC_MEMORY_TYPE_CUDA);
+        comm->grh_buf = comm->grh_cuda_header->addr;
+        if (ucc_unlikely(status != UCC_OK)) {
+            tl_error(comm->ctx->lib, "failed to allocate cuda memory");
+            return status;
         }
 
-        CUDA_FUNC(cudaMemset(comm->grh_buf, 0, GRH_LENGTH));
+        status = ucc_mc_memset(comm->grh_buf, 0, GRH_LENGTH, UCC_MEMORY_TYPE_CUDA);
+        if (status != UCC_OK) {
+            tl_error(comm->ctx->lib, "could not cuda memset");
+            goto error;
+        }
 
         comm->grh_mr = ibv_reg_mr(comm->ctx->pd, comm->grh_buf, GRH_LENGTH,
                                   IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE);
         if (!comm->grh_mr) {
-            tl_error(comm->ctx->lib, "Could not register device memory for GRH, errno %d", errno);
+            tl_error(comm->ctx->lib, "could not register device memory for GRH, errno %d", errno);
             status = UCC_ERR_NO_RESOURCE;
             goto error;
         }
 
-        // assuming the device page size is same as host page size
-        CUDA_FUNC(cudaMalloc((void**) &comm->pp_buf, buf_size * comm->buf_n));
-        if (!comm->pp_buf) {
-            tl_error(comm->ctx->lib, "cuda memcpy failed");
-            status = UCC_ERR_NO_MEMORY;
+        status = ucc_mc_alloc(&comm->pp_cuda_header, buf_size * comm->buf_n, UCC_MEMORY_TYPE_CUDA);
+        comm->pp_buf = comm->pp_cuda_header->addr;
+        if (ucc_unlikely(status != UCC_OK)) {
+            tl_error(comm->ctx->lib, "failed to allocate cuda memory");
             goto error;
         }
 
-        CUDA_FUNC(cudaMemset(comm->pp_buf, 0, buf_size * comm->buf_n));
+        status = ucc_mc_memset(comm->pp_buf, 0, buf_size * comm->buf_n, UCC_MEMORY_TYPE_CUDA);
+        if (status != UCC_OK) {
+            tl_error(comm->ctx->lib, "could not cuda memset");
+            goto error;
+        }
 
         comm->pp_mr = ibv_reg_mr(comm->ctx->pd, comm->pp_buf, buf_size * comm->buf_n,
                                  IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE |

src/components/tl/mlx5/tl_mlx5.c

@@ -92,8 +92,8 @@ static ucc_config_field_t ucc_tl_mlx5_lib_config_table[] = {
      ucc_offsetof(ucc_tl_mlx5_lib_config_t, mcast_conf.max_eager),
      UCC_CONFIG_TYPE_MEMUNITS},
 
-    {"MCAST_DEVICE_MEM_ENABLE", "0", "Enable GPU memory support for Mcast",
-     ucc_offsetof(ucc_tl_mlx5_lib_config_t, mcast_conf.device_mem_enabled),
+    {"MCAST_CUDA_MEM_ENABLE", "0", "Enable GPU CUDA memory support for Mcast. GPUDirect RDMA must be enabled",
+     ucc_offsetof(ucc_tl_mlx5_lib_config_t, mcast_conf.cuda_mem_enabled),
      UCC_CONFIG_TYPE_INT},
 
     {NULL}};

src/components/tl/mlx5/tl_mlx5_coll.c

@@ -19,6 +19,11 @@ ucc_status_t ucc_tl_mlx5_bcast_mcast_init(ucc_base_coll_args_t *coll_args,
         tl_trace(team->context->lib, "mcast bcast not supported for active sets");
         return UCC_ERR_NOT_SUPPORTED;
     }
+
+    if (UCC_OK != ucc_tl_mlx5_mcast_check_memory_type_cap(coll_args, team)) {
+        tl_trace(team->context->lib, "mcast bcast not compatible with this memory type");
+        return UCC_ERR_NOT_SUPPORTED;
+    }
 
     task = ucc_tl_mlx5_get_task(coll_args, team);