Merge pull request #10084 from tvegas1/odp_cuda_managed

UCT/IB: Support cuda-managed memory when ODP is available

config/ucx.conf

@@ -1,6 +1,7 @@
 [Grace Hopper]
 CPU model=Grace
-UCX_REG_NONBLOCK_MEM_TYPES=host
+UCX_REG_NONBLOCK_MEM_TYPES=host,cuda-managed
+UCX_IB_ODP_MEM_TYPES=host,cuda-managed
 UCX_IB_MLX5_DEVX_OBJECTS=
 UCX_DISTANCE_BW=auto,sys:16500MBs
 

src/uct/ib/base/ib_md.c

@@ -170,6 +170,12 @@ ucs_config_field_t uct_ib_md_config_table[] = {
      "Enable UMR optimization for XGVMI mkeys export/import.",
      ucs_offsetof(uct_ib_md_config_t, xgvmi_umr_enable), UCS_CONFIG_TYPE_BOOL},
 
+    {"ODP_MEM_TYPES", "host",
+     "Advertise non-blocking registration for these memory types, when ODP is "
+     "enabled.\n",
+     ucs_offsetof(uct_ib_md_config_t, ext.odp.mem_types),
+     UCS_CONFIG_TYPE_BITMAP(ucs_memory_type_names)},
+
     {NULL}
 };
 
@@ -581,7 +587,8 @@ ucs_status_t uct_ib_memh_alloc(uct_ib_md_t *md, size_t length,
     memh->rkey = UCT_IB_INVALID_MKEY;
 
     if ((mem_flags & UCT_MD_MEM_FLAG_NONBLOCK) && (length > 0) &&
-        (md->reg_nonblock_mem_types & UCS_BIT(UCS_MEMORY_TYPE_HOST))) {
+        (md->reg_nonblock_mem_types != 0)) {
+        /* Registration will fail if memory does not actually support it */
         memh->flags |= UCT_IB_MEM_FLAG_ODP;
     }
 
@@ -1236,8 +1243,10 @@ 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;
+
     /* Check for GPU-direct support */
-    md->reg_mem_types = UCS_BIT(UCS_MEMORY_TYPE_HOST);
     if (md_config->enable_gpudirect_rdma != UCS_NO) {
         /* Check peer memory driver is loaded, different driver versions use 
          * different paths */

src/uct/ib/base/ib_md.h

@@ -83,6 +83,7 @@ typedef struct uct_ib_md_ext_config {
     struct {
         int                  prefetch;     /**< Auto-prefetch non-blocking memory
                                                 registrations / allocations */
+        uint64_t             mem_types;    /**< Supported mem types for ODP */
     } odp;
 
     unsigned long            gid_index;    /**< IB GID index to use */

src/uct/ib/mlx5/dv/ib_mlx5dv_md.c

@@ -1602,6 +1602,7 @@ static void uct_ib_mlx5_devx_check_odp(uct_ib_mlx5_md_t *md,
 {
     char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_hca_cap_out)] = {};
     char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_hca_cap_in)]   = {};
+    ucs_string_buffer_t strb = UCS_STRING_BUFFER_INITIALIZER;
     ucs_status_t status;
     const void *odp_cap;
     const char *reason;
@@ -1673,10 +1674,17 @@ static void uct_ib_mlx5_devx_check_odp(uct_ib_mlx5_md_t *md,
         goto no_odp;
     }
 
-    ucs_debug("%s: ODP is supported, version %d",
-              uct_ib_device_name(&md->super.dev), version);
-    md->super.reg_nonblock_mem_types = UCS_BIT(UCS_MEMORY_TYPE_HOST);
-    md->super.gva_mem_types          = UCS_BIT(UCS_MEMORY_TYPE_HOST);
+    md->super.gva_mem_types          = md_config->ext.odp.mem_types;
+    md->super.reg_nonblock_mem_types = md_config->ext.odp.mem_types;
+
+    if (ucs_log_is_enabled(UCS_LOG_LEVEL_DEBUG)) {
+        ucs_string_buffer_append_flags(&strb, md->super.reg_nonblock_mem_types,
+                                       ucs_memory_type_names);
+        ucs_debug("%s: ODP is supported, version %d: memory=%s",
+                  uct_ib_device_name(&md->super.dev), version,
+                  ucs_string_buffer_cstr(&strb));
+        ucs_string_buffer_cleanup(&strb);
+    }
 
     if (!md->super.relaxed_order) {
         return;

test/gtest/ucp/test_ucp_mmap.cc

@@ -526,9 +526,9 @@ UCS_TEST_P(test_ucp_mmap, alloc_mem_type) {
             ASSERT_UCS_OK(status);
 
             is_dummy           = (size == 0);
-            expect_rma_offload = !UCP_MEM_IS_CUDA_MANAGED(mem_type) &&
-                                 (is_tl_rdma() || is_tl_shm()) &&
+            expect_rma_offload = (is_tl_rdma() || is_tl_shm()) &&
                                  check_reg_mem_types(sender(), mem_type);
+
             test_rkey_management(memh, is_dummy, expect_rma_offload);
 
             status = ucp_mem_unmap(sender().ucph(), memh);
@@ -571,8 +571,7 @@ UCS_TEST_P(test_ucp_mmap, reg_mem_type) {
             EXPECT_EQ(alloc_mem_type, memh->mem_type);
         }
 
-        expect_rma_offload = !UCP_MEM_IS_CUDA_MANAGED(alloc_mem_type) &&
-                             !UCP_MEM_IS_ROCM_MANAGED(alloc_mem_type) &&
+        expect_rma_offload = !UCP_MEM_IS_ROCM_MANAGED(alloc_mem_type) &&
                              is_tl_rdma() &&
                              check_reg_mem_types(sender(), alloc_mem_type);
         test_rkey_management(memh, is_dummy, expect_rma_offload);
@@ -944,6 +943,32 @@ UCS_TEST_P(test_ucp_mmap, fixed) {
     }
 }
 
+UCS_TEST_P(test_ucp_mmap, gva_allocate, "GVA_ENABLE=y",
+           "IB_MLX5_DEVX_OBJECTS?=")
+{
+    for (auto mem_type : mem_buffer::supported_mem_types()) {
+        ucp_md_map_t md_map = sender().ucph()->gva_md_map[mem_type];
+        if (md_map == 0) {
+            continue;
+        }
+
+        ucp_mem_h memh;
+        ucp_mem_map_params_t params;
+        params.field_mask  = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
+                             UCP_MEM_MAP_PARAM_FIELD_LENGTH |
+                             UCP_MEM_MAP_PARAM_FIELD_FLAGS |
+                             UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE;
+        params.address     = NULL;
+        params.memory_type = mem_type;
+        params.length      = 1 * UCS_MBYTE;
+        params.flags       = UCP_MEM_MAP_ALLOCATE;
+
+        ASSERT_UCS_OK(ucp_mem_map(sender().ucph(), &params, &memh));
+        EXPECT_TRUE(ucs_test_all_flags(memh->md_map, md_map));
+        ASSERT_UCS_OK(ucp_mem_unmap(sender().ucph(), memh));
+    }
+}
+
 UCS_TEST_P(test_ucp_mmap, gva, "GVA_ENABLE=y", "IB_MLX5_DEVX_OBJECTS?=")
 {
     std::list<void*> bufs;