UCT/IB: Check EFA-specific GPUDirect support

Author: tvegas1

src/uct/ib/base/ib_md.c

@@ -1203,8 +1203,8 @@ void uct_ib_md_parse_relaxed_order(uct_ib_md_t *md,
               uct_ib_device_name(&md->dev), md->relaxed_order ? "en" : "dis");
 }
 
-static void uct_ib_check_gpudirect_driver(uct_ib_md_t *md, const char *file,
-                                          ucs_memory_type_t mem_type)
+void uct_ib_check_gpudirect_driver(uct_ib_md_t *md, const char *file,
+                                   ucs_memory_type_t mem_type)
 {
     if (md->reg_mem_types & UCS_BIT(mem_type)) {
         return;
@@ -1316,8 +1316,8 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
         md->check_subnet_filter = 1;
     }
 
-    md->reg_mem_types = UCS_BIT(UCS_MEMORY_TYPE_HOST) |
-                        md->reg_nonblock_mem_types;
+    md->reg_mem_types |= UCS_BIT(UCS_MEMORY_TYPE_HOST) |
+                         md->reg_nonblock_mem_types;
 
     /* Check for GPU-direct support */
     if (md_config->enable_gpudirect_rdma != UCS_NO) {
@@ -1333,7 +1333,6 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md,
                 md, "/sys/module/nv_peer_mem/version",
                 UCS_MEMORY_TYPE_CUDA);
 
-
         /* check if ROCM KFD driver is loaded */
         uct_ib_check_gpudirect_driver(md, "/dev/kfd", UCS_MEMORY_TYPE_ROCM);
 

src/uct/ib/base/ib_md.h

@@ -431,4 +431,8 @@ ucs_status_t uct_ib_memh_alloc(uct_ib_md_t *md, size_t length,
                                unsigned mem_flags, size_t memh_base_size,
                                size_t mr_size, uct_ib_mem_t **memh_p);
 
+void uct_ib_check_gpudirect_driver(uct_ib_md_t *md,
+                                   const char *file,
+                                   ucs_memory_type_t mem_type);
+
 #endif

src/uct/ib/efa/base/ib_efa_md.c

@@ -75,6 +75,12 @@ static ucs_status_t uct_ib_efa_md_open(struct ibv_device *ibv_device,
      */
     dev->req_notify_cq_support = 0;
 
+    if (md_config->enable_gpudirect_rdma != UCS_NO) {
+        uct_ib_check_gpudirect_driver(
+                &md->super, "/sys/module/efa_nv_peermem/version",
+                UCS_MEMORY_TYPE_CUDA);
+    }
+
     status = uct_ib_md_open_common(&md->super, ibv_device, md_config);
     if (status != UCS_OK) {
         goto err_md_free;