UCP/RNDV: Skip inter-node rkey_ptr without exportable memory

src/ucp/rndv/proto_rndv.c

@@ -88,6 +88,25 @@ ucp_proto_rndv_ctrl_get_md_map(const ucp_proto_rndv_ctrl_init_params_t *params,
             continue;
         }
 
+        /* A remote key of a memory domain which accesses memory by rkey_ptr
+         * (for example cuda_ipc) can be opened by a peer on a different node
+         * only if the memory is exportable to that node. Do not advertise a
+         * remote key which the peer would fail to open. Both the deprecated MD
+         * flag and the component flag are checked, since cuda_ipc reports only
+         * the former.
+         */
+        if (((md_attr->flags & UCT_MD_FLAG_RKEY_PTR) ||
+             (cmpt_attr->flags & UCT_COMPONENT_FLAG_RKEY_PTR)) &&
+            ucp_ep_config_is_inter_node(ep_config_key) &&
+            !(params->super.reg_mem_info.flags &
+              UCS_MEM_FLAG_RKEY_PTR_INTER_NODE)) {
+            ucs_trace_req("lane[%d]: md %s rkey_ptr is not usable inter-node, "
+                          "mem_flags 0x%x",
+                          lane, context->tl_mds[md_index].rsc.md_name,
+                          params->super.reg_mem_info.flags);
+            continue;
+        }
+
         /* Check reachability between mem_sys_dev and current lane's sys_dev */
         if (!ucs_topo_is_reachable(ep_sys_dev, mem_sys_dev)) {
             continue;
@@ -170,6 +189,14 @@ ucp_proto_rndv_rkey_mem_flags_estimate(const ucp_proto_init_params_t *params)
         /*
          * Derive UCS_MEM_FLAG_REGISTRABLE from matching local MDs which
          * require it and whose remote MDs are present in the rkey.
+         *
+         * UCS_MEM_FLAG_RKEY_PTR_INTER_NODE is intentionally not inferred
+         * here: unlike REGISTRABLE it is not an MD required_mem_flags bit,
+         * and is checked only when packing keys in
+         * ucp_proto_rndv_ctrl_get_md_map(). Runtime packing uses the real
+         * local buffer flags; omitting it from this estimate can only make
+         * nested ctrl modeling (e.g. peer rndv/rtr) slightly pessimistic
+         * for fabric-exportable remotes, not incorrect.
          */
         md_index   = context->tl_rscs[lane_cfg->rsc_index].md_index;
         mem_flags |= context->tl_mds[md_index].attr.required_mem_flags;

src/ucs/memory/memory_type.h

@@ -52,7 +52,14 @@ typedef enum ucs_memory_type {
 
 
 typedef enum ucs_mem_flags {
-    UCS_MEM_FLAG_REGISTRABLE = UCS_BIT(0) /**< Memory is registrable by MDs */
+    UCS_MEM_FLAG_REGISTRABLE = UCS_BIT(0), /**< Memory is registrable by MDs */
+
+    /**
+     * Memory can be accessed by @ref uct_rkey_ptr on a peer which resides on a
+     * different node (for example CUDA IPC over MNNVL). Checked only by users
+     * which know the peer is on a different node.
+     */
+    UCS_MEM_FLAG_RKEY_PTR_INTER_NODE = UCS_BIT(1)
 } ucs_mem_flags_t;
 
 

src/uct/cuda/cuda_copy/cuda_copy_md.c

@@ -845,49 +845,86 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address,
     return dmabuf;
 }
 
+/*
+ * Detect whether the allocation can be opened by a peer residing on a different
+ * node, following the same handle type selection as
+ * uct_cuda_ipc_mem_add_reg(): legacy IPC handles are node-local, while fabric
+ * handles can be exported over MNNVL.
+ */
 static uint8_t
-uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
-                                  const ucs_memory_info_t *mem_info,
-                                  int is_async_managed, int is_host_located,
-                                  const uct_cuda_copy_md_dmabuf_t *dmabuf)
+uct_cuda_copy_md_detect_rkey_ptr_flags(const ucs_memory_info_t *mem_info)
 {
-    int close_dmabuf = 0;
-    uct_cuda_copy_md_dmabuf_t local_dmabuf;
+#if HAVE_CUDA_FABRIC
+    CUpointer_attribute attr_type[2];
+    void *attr_data[2];
+    uint64_t allowed_handle_types;
+    int legacy_capable;
+    ucs_status_t status;
 
-    if (is_async_managed) {
+    /* Only memory which cuda_ipc can register is relevant */
+    if (mem_info->type != UCS_MEMORY_TYPE_CUDA) {
         return 0;
     }
 
-    /* Host-located CUDA VMM is registerable even if dmabuf export fails. */
-    if (is_host_located) {
-        return UCS_MEM_FLAG_REGISTRABLE;
-    }
+    attr_type[0] = CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE;
+    attr_data[0] = &legacy_capable;
+    attr_type[1] = CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES;
+    attr_data[1] = &allowed_handle_types;
 
-    if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
-        return UCS_MEM_FLAG_REGISTRABLE;
+    status = UCT_CUDADRV_FUNC_LOG_DEBUG(
+            cuPointerGetAttributes(ucs_static_array_size(attr_data), attr_type,
+                                   attr_data,
+                                   (CUdeviceptr)mem_info->base_address));
+    if (status != UCS_OK) {
+        return 0;
     }
 
-    if (!md->config.dmabuf_supported) {
-        return UCS_MEM_FLAG_REGISTRABLE;
+    if (legacy_capable || !(allowed_handle_types & CU_MEM_HANDLE_TYPE_FABRIC)) {
+        return 0;
     }
 
-    if (dmabuf == NULL) {
-        local_dmabuf = uct_cuda_copy_md_get_dmabuf(mem_info->base_address,
-                                                   mem_info->alloc_length,
-                                                   mem_info->sys_dev);
-        dmabuf       = &local_dmabuf;
-        close_dmabuf = 1;
-    }
+    return UCS_MEM_FLAG_RKEY_PTR_INTER_NODE;
+#else
+    return 0;
+#endif
+}
 
-    if (dmabuf->fd == UCT_DMABUF_FD_INVALID) {
-        return 0;
-    }
+static uint8_t
+uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
+                                  const ucs_memory_info_t *mem_info,
+                                  int is_async_managed, int is_host_located,
+                                  const uct_cuda_copy_md_dmabuf_t *dmabuf)
+{
+    uint8_t mem_flags = UCS_MEM_FLAG_REGISTRABLE |
+                        uct_cuda_copy_md_detect_rkey_ptr_flags(mem_info);
+    int close_dmabuf  = 0;
+    uct_cuda_copy_md_dmabuf_t local_dmabuf;
+
+    if (!is_async_managed) {
+        /* Host-located CUDA VMM is registerable even if dmabuf export fails. */
+        if (is_host_located ||
+            (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ||
+            !md->config.dmabuf_supported) {
+            return mem_flags;
+        }
 
-    if (close_dmabuf) {
-        ucs_close_fd(&local_dmabuf.fd);
+        if (dmabuf == NULL) {
+            local_dmabuf = uct_cuda_copy_md_get_dmabuf(mem_info->base_address,
+                                                       mem_info->alloc_length,
+                                                       mem_info->sys_dev);
+            dmabuf       = &local_dmabuf;
+            close_dmabuf = 1;
+        }
+
+        if (dmabuf->fd != UCT_DMABUF_FD_INVALID) {
+            if (close_dmabuf) {
+                ucs_close_fd(&local_dmabuf.fd);
+            }
+            return mem_flags;
+        }
     }
 
-    return UCS_MEM_FLAG_REGISTRABLE;
+    return mem_flags & ~UCS_MEM_FLAG_REGISTRABLE;
 }
 
 ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,

test/gtest/ucp/test_ucp_proto.cc

@@ -14,6 +14,7 @@
 
 extern "C" {
 #include <ucp/core/ucp_rkey.h>
+#include <ucp/core/ucp_ep.inl>
 #include <ucp/dt/datatype_iter.inl>
 #include <ucp/proto/proto.h>
 #include <ucp/proto/proto_debug.h>
@@ -1075,6 +1076,135 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_cuda_async_non_reg, rcx,
 UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_cuda_async_non_reg, rcv,
                               "rc_v,cuda_copy")
 
+/*
+ * Inter-node rkey_ptr (e.g. cuda_ipc over MNNVL) may only advertise a remote
+ * key when the local buffer is exportable to another node. Force an inter-node
+ * endpoint config and check that put/rndv's md_map follows
+ * UCS_MEM_FLAG_RKEY_PTR_INTER_NODE.
+ */
+class test_ucp_proto_rkey_ptr_inter_node : public test_ucp_proto {
+protected:
+    void init() override
+    {
+        modify_config("PROTOS", "put/rndv,rndv/*");
+        modify_config("RNDV_THRESH", "0");
+        test_ucp_proto::init();
+    }
+
+    static int md_is_rkey_ptr(ucp_context_h ctx, ucp_md_index_t md_index)
+    {
+        return (ctx->tl_mds[md_index].attr.flags & UCT_MD_FLAG_RKEY_PTR) ||
+               (ucp_cmpt_attr_by_md_index(ctx, md_index)->flags &
+                UCT_COMPONENT_FLAG_RKEY_PTR);
+    }
+
+    ucp_md_index_t find_rkey_ptr_md_index()
+    {
+        ucp_context_h ctx = context();
+        ucp_md_index_t md_index;
+
+        for (md_index = 0; md_index < ctx->num_mds; ++md_index) {
+            if (md_is_rkey_ptr(ctx, md_index) &&
+                (ctx->reg_md_map[UCS_MEMORY_TYPE_CUDA] & UCS_BIT(md_index))) {
+                return md_index;
+            }
+        }
+
+        return UCP_NULL_RESOURCE;
+    }
+
+    int has_rkey_ptr_lane(ucp_md_index_t md_index)
+    {
+        const ucp_ep_config_t *ep_config =
+                &ucs_array_elem(&worker()->ep_config, sender().ep()->cfg_index);
+        ucp_lane_index_t lane;
+
+        for (lane = 0; lane < ep_config->key.num_lanes; ++lane) {
+            if ((ep_config->key.lanes[lane].rsc_index != UCP_NULL_RESOURCE) &&
+                (context()->tl_rscs[ep_config->key.lanes[lane].rsc_index]
+                         .md_index == md_index)) {
+                return 1;
+            }
+        }
+
+        return 0;
+    }
+
+    ucp_md_map_t
+    select_put_rndv_md_map(ucp_ep_config_t *ep_config, uint8_t mem_flags)
+    {
+        ucp_worker_cfg_index_t ep_cfg_index = sender().ep()->cfg_index;
+        ucp_memory_info_t mem_info         = {
+            .type    = UCS_MEMORY_TYPE_CUDA,
+            .sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN,
+            .flags   = mem_flags
+        };
+        ucp_proto_select_param_t select_param;
+        const ucp_proto_select_elem_t *select_elem;
+        const ucp_proto_threshold_elem_t *thresh;
+        const ucp_proto_rndv_ctrl_priv_t *rpriv;
+
+        ucp_proto_select_param_init(&select_param, UCP_OP_ID_PUT, 0, 0,
+                                    UCP_DATATYPE_CONTIG, &mem_info, 1);
+        select_elem = ucp_proto_select_lookup_slow(
+                worker(), &ep_config->proto_select, 0, ep_cfg_index,
+                UCP_WORKER_CFG_INDEX_NULL, &select_param);
+        if (select_elem == nullptr) {
+            UCS_TEST_SKIP_R("put protocol was not selected");
+        }
+
+        thresh = ucp_proto_thresholds_search_slow(select_elem->thresholds,
+                                                  UCS_MBYTE);
+        if (strcmp(thresh->proto_config.proto->name, "put/rndv") != 0) {
+            UCS_TEST_SKIP_R("put/rndv was not selected");
+        }
+
+        rpriv = static_cast<const ucp_proto_rndv_ctrl_priv_t*>(
+                thresh->proto_config.priv);
+        return rpriv->md_map;
+    }
+};
+
+UCS_TEST_P(test_ucp_proto_rkey_ptr_inter_node, filter_md_map)
+{
+    ucp_md_index_t rkey_ptr_md;
+    ucp_ep_config_t *ep_config;
+    unsigned orig_flags;
+    ucp_md_map_t md_map_no_flag;
+    ucp_md_map_t md_map_with_flag;
+
+    require_cuda_memory();
+
+    rkey_ptr_md = find_rkey_ptr_md_index();
+    if (rkey_ptr_md == UCP_NULL_RESOURCE) {
+        UCS_TEST_SKIP_R("no CUDA rkey_ptr MD");
+    }
+
+    if (!has_rkey_ptr_lane(rkey_ptr_md)) {
+        UCS_TEST_SKIP_R("no endpoint lane for CUDA rkey_ptr MD");
+    }
+
+    ep_config  = &ucs_array_elem(&worker()->ep_config, sender().ep()->cfg_index);
+    orig_flags = ep_config->key.flags;
+
+    /* Force inter-node so the new flag is required for rkey_ptr MDs. */
+    ep_config->key.flags &= ~(UCP_EP_CONFIG_KEY_FLAG_INTRA_NODE |
+                              UCP_EP_CONFIG_KEY_FLAG_SELF);
+    ASSERT_TRUE(ucp_ep_config_is_inter_node(&ep_config->key));
+
+    md_map_no_flag = select_put_rndv_md_map(ep_config, 0);
+    EXPECT_FALSE(md_map_no_flag & UCS_BIT(rkey_ptr_md));
+
+    md_map_with_flag = select_put_rndv_md_map(ep_config,
+                                              UCS_MEM_FLAG_RKEY_PTR_INTER_NODE);
+    EXPECT_TRUE(md_map_with_flag & UCS_BIT(rkey_ptr_md));
+
+    ep_config->key.flags = orig_flags;
+}
+
+UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_proto_rkey_ptr_inter_node,
+                                        shm_rc_ipc, "rc_x,cuda_ipc,rocm_ipc")
+
 class test_perf_node : public test_ucp_proto {
 };
 

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -339,7 +339,6 @@ class test_mem_alloc_device : public test_switch_cuda_device {
 
         ASSERT_UCS_OK(query_status);
         EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
-        EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
     }
 
 private:
@@ -423,7 +422,6 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
     EXPECT_UCS_OK(lookup_status);
     if (lookup_status == UCS_OK) {
         EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_info.type);
-        EXPECT_TRUE(mem_info.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
     }
     EXPECT_UCS_OK(free_status);
 }
@@ -466,6 +464,40 @@ UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable,
 }
 #endif
 
+UCS_TEST_P(test_mem_alloc_device, legacy_cuda_no_rkey_ptr_inter_node)
+{
+    constexpr size_t size = 4096;
+    uct_md_mem_attr_v2_t mem_attr = {};
+    CUdeviceptr dptr              = 0;
+
+    ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc(&dptr, size));
+
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
+                          UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
+    ASSERT_UCS_OK(uct_md_mem_query_v2(md(), reinterpret_cast<void*>(dptr), size,
+                                      &mem_attr));
+    EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
+    EXPECT_EQ(0, mem_attr.mem_flags & UCS_MEM_FLAG_RKEY_PTR_INTER_NODE);
+
+    EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr));
+}
+
+#if HAVE_CUDA_FABRIC
+UCS_TEST_P(test_mem_alloc_device, fabric_cuda_rkey_ptr_inter_node,
+           "CUDA_COPY_ENABLE_FABRIC=y")
+{
+    constexpr size_t size = 4 * UCS_MBYTE;
+    uct_md_mem_attr_v2_t mem_attr = {};
+    cuda_fabric_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
+
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
+                          UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
+    ASSERT_UCS_OK(uct_md_mem_query_v2(md(), buffer.ptr(), size, &mem_attr));
+    EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
+    EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_RKEY_PTR_INTER_NODE);
+}
+#endif
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy);
 
 class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test {