UCP/UCT/CUDA: Avoid gdr_copy for non gdr capable memory

config/m4/cuda.m4

@@ -255,6 +255,9 @@ AS_IF([test "x$cuda_checked" != "xyes"],
          AC_CHECK_DECLS([CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN],
                         [], [], [[#include <cuda.h>]])
 
+         AC_CHECK_DECLS([CU_POINTER_ATTRIBUTE_IS_GPU_DIRECT_RDMA_CAPABLE],
+                        [], [], [[#include <cuda.h>]])
+
          CPPFLAGS="$save_CPPFLAGS"
          LDFLAGS="$save_LDFLAGS"
          LIBS="$save_LIBS"

src/ucp/core/ucp_device.c

@@ -285,7 +285,8 @@ static ucs_status_t ucp_device_mem_list_export_handle(
         return status;
     }
 
-    ucp_mem_type_unpack(worker, mem->address, handle, handle_size, mem_type);
+    ucp_mem_type_unpack(worker, mem->address, handle, handle_size, mem_type,
+                        NULL);
     return UCS_OK;
 }
 
@@ -872,36 +873,44 @@ void ucp_device_mem_list_release(void *mem_list_h)
     uct_mem_free(&info.mem);
 }
 
-static ucs_memory_type_t
-ucp_device_counter_mem_type(ucp_context_h context, const void *counter_ptr,
-                            const ucp_device_counter_params_t *params)
+/*
+ * Resolve the memory info of the counter area. Returns NULL when the memory
+ * flags cannot be resolved, in which case only @a mem_info->type is set and
+ * the memory type endpoint looks the flags up by itself.
+ */
+static const ucp_memory_info_t *
+ucp_device_counter_mem_info(ucp_context_h context, const void *counter_ptr,
+                            const ucp_device_counter_params_t *params,
+                            ucp_memory_info_t *mem_info)
 {
-    ucs_memory_info_t mem_info;
-
     if (params->field_mask & UCP_DEVICE_COUNTER_PARAMS_FIELD_MEMH) {
-        return params->memh->mem_type;
+        *mem_info = ucp_memory_info_from_memh(params->memh);
+        return mem_info;
     }
 
     if (params->field_mask & UCP_DEVICE_COUNTER_PARAMS_FIELD_MEM_TYPE) {
-        return params->mem_type;
+        /* The memory flags are not part of the caller's hint */
+        mem_info->type = params->mem_type;
+        return NULL;
     }
 
-    ucp_memory_detect_internal(context, counter_ptr, sizeof(uint64_t),
-                               &mem_info);
-    return mem_info.type;
+    ucp_memory_detect(context, counter_ptr, sizeof(uint64_t), mem_info);
+    return mem_info;
 }
 
 ucs_status_t ucp_device_counter_init(ucp_worker_h worker,
                                      const ucp_device_counter_params_t *params,
                                      void *counter_ptr)
 {
     uint64_t counter_value = 0;
-    ucs_memory_type_t mem_type;
+    ucp_memory_info_t mem_info;
+    const ucp_memory_info_t *resolved;
 
-    mem_type = ucp_device_counter_mem_type(worker->context, counter_ptr,
-                                           params);
+    resolved = ucp_device_counter_mem_info(worker->context, counter_ptr, params,
+                                           &mem_info);
     ucp_dt_contig_unpack(worker, counter_ptr, &counter_value,
-                         sizeof(counter_value), mem_type,
+                         sizeof(counter_value),
+                         (ucs_memory_type_t)mem_info.type, resolved,
                          sizeof(counter_value));
     return UCS_OK;
 }
@@ -911,11 +920,13 @@ uint64_t ucp_device_counter_read(ucp_worker_h worker,
                                  void *counter_ptr)
 {
     uint64_t counter_value = 0;
-    ucs_memory_type_t mem_type;
+    ucp_memory_info_t mem_info;
+    const ucp_memory_info_t *resolved;
 
-    mem_type = ucp_device_counter_mem_type(worker->context, counter_ptr,
-                                           params);
+    resolved = ucp_device_counter_mem_info(worker->context, counter_ptr, params,
+                                           &mem_info);
     ucp_dt_contig_pack(worker, &counter_value, counter_ptr,
-                       sizeof(counter_value), mem_type, sizeof(counter_value));
+                       sizeof(counter_value), (ucs_memory_type_t)mem_info.type,
+                       resolved, sizeof(counter_value));
     return counter_value;
 }

src/ucp/dt/datatype_iter.inl

@@ -415,7 +415,7 @@ ucp_datatype_iter_next_pack(const ucp_datatype_iter_t *dt_iter,
                                      dt_iter->offset);
         ucp_dt_contig_pack(worker, dest, src, length,
                            (ucs_memory_type_t)dt_iter->mem_info.type,
-                           dt_iter->length);
+                           &dt_iter->mem_info, dt_iter->length);
         break;
     case UCP_DATATYPE_IOV:
         ucp_datatype_iter_iov_check(dt_iter);
@@ -427,7 +427,7 @@ ucp_datatype_iter_next_pack(const ucp_datatype_iter_t *dt_iter,
                               &next_iter->type.iov.iov_offset,
                               &next_iter->type.iov.iov_index,
                               (ucs_memory_type_t)dt_iter->mem_info.type,
-                              dt_iter->length);
+                              &dt_iter->mem_info, dt_iter->length);
         break;
     case UCP_DATATYPE_GENERIC:
         if (max_length != 0) {
@@ -482,7 +482,7 @@ ucp_datatype_iter_unpack(ucp_datatype_iter_t *dt_iter, ucp_worker_h worker,
         dest = UCS_PTR_BYTE_OFFSET(dt_iter->type.contig.buffer, offset);
         ucp_dt_contig_unpack(worker, dest, src, length,
                              (ucs_memory_type_t)dt_iter->mem_info.type,
-                             dt_iter->length);
+                             &dt_iter->mem_info, dt_iter->length);
         status = UCS_OK;
         break;
     case UCP_DATATYPE_IOV:
@@ -493,7 +493,7 @@ ucp_datatype_iter_unpack(ucp_datatype_iter_t *dt_iter, ucp_worker_h worker,
                                            &dt_iter->type.iov.iov_offset,
                                            &dt_iter->type.iov.iov_index,
                                            (ucs_memory_type_t)dt_iter->mem_info.type,
-                                           dt_iter->length);
+                                           &dt_iter->mem_info, dt_iter->length);
         ucs_assert(unpacked_length <= length);
         dt_iter->offset += unpacked_length;
         status           = UCS_OK;

src/ucp/dt/dt.c

@@ -60,26 +60,97 @@ ucs_status_t ucp_dt_mem_info_verify(const char *dt_name, size_t index,
 }
 
 
+/*
+ * Select the first RMA lane of the mem type endpoint whose memory domain
+ * accepts the buffer. @a mem_info holds the memory flags resolved by the
+ * caller, or is NULL when they are unknown, in which case they are resolved on
+ * the first memory domain which actually requires them. Returns UCP_NULL_LANE
+ * if no lane is compatible.
+ */
+static ucp_lane_index_t
+ucp_mem_type_ep_find_lane(ucp_ep_h ep, const void *address, size_t length,
+                          const ucp_memory_info_t *mem_info)
+{
+    ucp_context_h context          = ep->worker->context;
+    const ucp_ep_config_key_t *key = &ucp_ep_config(ep)->key;
+    const uct_md_attr_v2_t *md_attr;
+    ucp_memory_info_t detected;
+    ucp_lane_index_t lane;
+    unsigned i;
+
+    for (i = 0; key->rma_lanes[i] != UCP_NULL_LANE; ++i) {
+        lane    = key->rma_lanes[i];
+        md_attr = ucp_ep_md_attr(ep, lane);
+        if (md_attr->required_mem_flags == 0) {
+            return lane;
+        }
+
+        if (mem_info == NULL) {
+            /* Coverity wrongly resolves the memory domain deregister function
+             * pointer to 'uct_cuda_ipc_mem_dereg' and considers 'address' as
+             * freed */
+            /* coverity[pass_freed_arg] */
+            ucp_memory_detect(context, address, length, &detected);
+            mem_info = &detected;
+        }
+
+        if (ucs_test_all_flags(mem_info->flags,
+                               md_attr->required_mem_flags)) {
+            return lane;
+        }
+
+        /* TODO: Also take the memory flags into account in the protocol
+         * performance estimation, as they are already part of the selection
+         * key. Performance model here diverges from the lane actually used. */
+    }
+
+    return UCP_NULL_LANE;
+}
+
+/*
+ * Select the mem type endpoint lane which can register the buffer, and register
+ * it. @a mem_info holds the memory flags resolved by the caller, or is NULL
+ * when they are unknown.
+ */
+static ucs_status_t
+ucp_mem_type_lane_reg(ucp_worker_h worker, ucp_ep_h ep, void *address,
+                      size_t length, ucs_memory_type_t mem_type,
+                      const ucp_memory_info_t *mem_info,
+                      ucp_lane_index_t *lane_p,
+                      ucp_mtype_pack_context_t *pack_context)
+{
+    ucp_lane_index_t lane = ucp_mem_type_ep_find_lane(ep, address, length,
+                                                      mem_info);
+
+    if (lane == UCP_NULL_LANE) {
+        ucs_error("no mem type lane can register %s buffer %p length %zu",
+                  ucs_memory_type_names[mem_type], address, length);
+        return UCS_ERR_UNSUPPORTED;
+    }
+
+    *lane_p = lane;
+    return ucp_mem_type_reg_buffers(worker, address, length, mem_type,
+                                    ucp_ep_md_index(ep, lane), pack_context);
+}
+
 UCS_PROFILE_FUNC_VOID(ucp_mem_type_unpack,
-                      (worker, buffer, recv_data, recv_length, mem_type),
+                      (worker, buffer, recv_data, recv_length, mem_type,
+                       mem_info),
                       ucp_worker_h worker, void *buffer, const void *recv_data,
-                      size_t recv_length, ucs_memory_type_t mem_type)
+                      size_t recv_length, ucs_memory_type_t mem_type,
+                      const ucp_memory_info_t *mem_info)
 {
     ucp_ep_h ep = worker->mem_type_ep[mem_type];
     ucp_lane_index_t lane;
-    unsigned md_index;
     ucs_status_t status;
     ucp_mtype_pack_context_t pack_context;
 
     if (recv_length == 0) {
         return;
     }
 
-    lane     = ucp_ep_config(ep)->key.rma_lanes[0];
-    md_index = ucp_ep_md_index(ep, lane);
-
-    status = ucp_mem_type_reg_buffers(worker, buffer, recv_length, mem_type,
-                                      md_index, &pack_context);
+    status = ucp_mem_type_lane_reg(worker, ep, buffer, recv_length, mem_type,
+                                   mem_info, &lane, &pack_context);
     if (status != UCS_OK) {
         ucs_fatal("failed to register buffer with mem type domain %s",
                   ucs_memory_type_names[mem_type]);
@@ -96,25 +167,22 @@ UCS_PROFILE_FUNC_VOID(ucp_mem_type_unpack,
 }
 
 UCS_PROFILE_FUNC_VOID(ucp_mem_type_pack,
-                      (worker, dest, src, length, mem_type),
+                      (worker, dest, src, length, mem_type, mem_info),
                       ucp_worker_h worker, void *dest, const void *src,
-                      size_t length, ucs_memory_type_t mem_type)
+                      size_t length, ucs_memory_type_t mem_type,
+                      const ucp_memory_info_t *mem_info)
 {
     ucp_ep_h ep = worker->mem_type_ep[mem_type];
     ucp_lane_index_t lane;
-    ucp_md_index_t md_index;
     ucs_status_t status;
     ucp_mtype_pack_context_t pack_context;
 
     if (length == 0) {
         return;
     }
 
-    lane     = ucp_ep_config(ep)->key.rma_lanes[0];
-    md_index = ucp_ep_md_index(ep, lane);
-
-    status = ucp_mem_type_reg_buffers(worker, (void *)src, length, mem_type,
-                                      md_index, &pack_context);
+    status = ucp_mem_type_lane_reg(worker, ep, (void*)src, length, mem_type,
+                                   mem_info, &lane, &pack_context);
     if (status != UCS_OK) {
         ucs_fatal("failed to register buffer with mem type domain %s",
                   ucs_memory_type_names[mem_type]);
@@ -141,18 +209,21 @@ size_t ucp_dt_pack(ucp_worker_h worker, ucp_datatype_t datatype,
         return length;
     }
 
+    /* This path is driven by 'ucp_dt_state_t', which unlike the datatype
+     * iterator does not carry the resolved memory flags of the buffer */
     switch (datatype & UCP_DATATYPE_CLASS_MASK) {
     case UCP_DATATYPE_CONTIG:
         ucp_dt_contig_pack(worker, dest,
                            UCS_PTR_BYTE_OFFSET(src, state->offset),
-                           length, mem_type, length);
+                           length, mem_type, NULL, length);
         result_len = length;
         break;
 
     case UCP_DATATYPE_IOV:
         UCS_PROFILE_CALL_VOID(ucp_dt_iov_gather, worker, dest, src, length,
                               &state->dt.iov.iov_offset,
-                              &state->dt.iov.iovcnt_offset, mem_type, length);
+                              &state->dt.iov.iovcnt_offset, mem_type, NULL,
+                              length);
         result_len = length;
         break;
 

src/ucp/dt/dt.h

@@ -78,12 +78,19 @@ ucp_memory_info_equal(const ucp_memory_info_t *a, const ucp_memory_info_t *b)
            (a->flags == b->flags);
 }
 
+/*
+ * Copy to/from memory which is not accessible from the CPU, through the memory
+ * type endpoint. @a mem_info holds the memory flags already resolved for the
+ * buffer, or is NULL when they are unknown and have to be looked up.
+ */
 void ucp_mem_type_pack(ucp_worker_h worker, void *dest, const void *src,
-                       size_t length, ucs_memory_type_t mem_type);
+                       size_t length, ucs_memory_type_t mem_type,
+                       const ucp_memory_info_t *mem_info);
 
 void ucp_mem_type_unpack(ucp_worker_h worker, void *buffer,
                          const void *recv_data, size_t recv_length,
-                         ucs_memory_type_t mem_type);
+                         ucs_memory_type_t mem_type,
+                         const ucp_memory_info_t *mem_info);
 
 
 static UCS_F_ALWAYS_INLINE void

src/ucp/dt/dt_contig.h

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2015. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
  * Copyright (C) Advanced Micro Devices, Inc. 2024. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
@@ -38,24 +38,26 @@ static UCS_F_ALWAYS_INLINE size_t ucp_contig_dt_length(ucp_datatype_t datatype,
 
 static UCS_F_ALWAYS_INLINE void
 ucp_dt_contig_pack(ucp_worker_h worker, void *dest, const void *src,
-                   size_t length, ucs_memory_type_t mem_type, size_t total_len)
+                   size_t length, ucs_memory_type_t mem_type,
+                   const ucp_memory_info_t *mem_info, size_t total_len)
 {
     if (ucs_likely(UCP_MEM_IS_ACCESSIBLE_FROM_CPU(mem_type))) {
         ucp_memcpy_pack(dest, src, length, total_len, "memcpy_pack");
     } else {
-        ucp_mem_type_pack(worker, dest, src, length, mem_type);
+        ucp_mem_type_pack(worker, dest, src, length, mem_type, mem_info);
     }
 }
 
 
 static UCS_F_ALWAYS_INLINE void
 ucp_dt_contig_unpack(ucp_worker_h worker, void *dest, const void *src,
-                     size_t length, ucs_memory_type_t mem_type, size_t total_len)
+                     size_t length, ucs_memory_type_t mem_type,
+                     const ucp_memory_info_t *mem_info, size_t total_len)
 {
     if (ucs_likely(UCP_MEM_IS_ACCESSIBLE_FROM_CPU(mem_type))) {
         ucp_memcpy_unpack(dest, src, length, total_len, "memcpy_unpack");
     } else {
-        ucp_mem_type_unpack(worker, dest, src, length, mem_type);
+        ucp_mem_type_unpack(worker, dest, src, length, mem_type, mem_info);
     }
 }
 

src/ucp/dt/dt_iov.c

@@ -25,7 +25,8 @@
 
 void ucp_dt_iov_gather(ucp_worker_h worker, void *dest, const ucp_dt_iov_t *iov,
                        size_t length, size_t *iov_offset, size_t *iovcnt_offset,
-                       ucs_memory_type_t mem_type, size_t total_len)
+                       ucs_memory_type_t mem_type,
+                       const ucp_memory_info_t *mem_info, size_t total_len)
 {
     size_t length_it = 0;
     size_t item_len, item_reminder, item_len_to_copy;
@@ -40,7 +41,7 @@ void ucp_dt_iov_gather(ucp_worker_h worker, void *dest, const ucp_dt_iov_t *iov,
         ucp_dt_contig_pack(worker, UCS_PTR_BYTE_OFFSET(dest, length_it),
                            UCS_PTR_BYTE_OFFSET(iov[*iovcnt_offset].buffer,
                                                *iov_offset),
-                           item_len_to_copy, mem_type, total_len);
+                           item_len_to_copy, mem_type, mem_info, total_len);
         length_it += item_len_to_copy;
 
         ucs_assert(length_it <= length);
@@ -56,7 +57,8 @@ void ucp_dt_iov_gather(ucp_worker_h worker, void *dest, const ucp_dt_iov_t *iov,
 size_t ucp_dt_iov_scatter(ucp_worker_h worker, const ucp_dt_iov_t *iov,
                           size_t iovcnt, const void *src, size_t length,
                           size_t *iov_offset, size_t *iovcnt_offset,
-                          ucs_memory_type_t mem_type, size_t total_len)
+                          ucs_memory_type_t mem_type,
+                          const ucp_memory_info_t *mem_info, size_t total_len)
 {
     size_t length_it = 0;
     size_t item_len, item_len_to_copy;
@@ -71,7 +73,7 @@ size_t ucp_dt_iov_scatter(ucp_worker_h worker, const ucp_dt_iov_t *iov,
                              UCS_PTR_BYTE_OFFSET(iov[*iovcnt_offset].buffer,
                                                  *iov_offset),
                              UCS_PTR_BYTE_OFFSET(src, length_it),
-                             item_len_to_copy, mem_type, total_len);
+                             item_len_to_copy, mem_type, mem_info, total_len);
         length_it += item_len_to_copy;
 
         ucs_assert(length_it <= length);