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_MEMPOOL_HANDLE],
+                        [], [], [[#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;
 }
 
@@ -901,7 +902,7 @@ ucs_status_t ucp_device_counter_init(ucp_worker_h worker,
     mem_type = ucp_device_counter_mem_type(worker->context, counter_ptr,
                                            params);
     ucp_dt_contig_unpack(worker, counter_ptr, &counter_value,
-                         sizeof(counter_value), mem_type,
+                         sizeof(counter_value), mem_type, NULL,
                          sizeof(counter_value));
     return UCS_OK;
 }
@@ -916,6 +917,7 @@ uint64_t ucp_device_counter_read(ucp_worker_h worker,
     mem_type = ucp_device_counter_mem_type(worker->context, counter_ptr,
                                            params);
     ucp_dt_contig_pack(worker, &counter_value, counter_ptr,
-                       sizeof(counter_value), mem_type, sizeof(counter_value));
+                       sizeof(counter_value), mem_type, NULL,
+                       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,92 @@ 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_lane(ucp_ep_h ep, 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 = &context->tl_mds[ucp_ep_md_index(ep, lane)].attr;
+        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;
+        }
+    }
+
+    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_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 +162,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 +204,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

@@ -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);

src/ucp/dt/dt_iov.h

@@ -55,7 +55,8 @@ static inline size_t ucp_dt_iov_length(const ucp_dt_iov_t *iov, size_t iovcnt)
  */
 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);
 
 
 /**
@@ -81,7 +82,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);
 
 
 /**

src/ucp/rma/amo_offload.c

@@ -23,7 +23,7 @@ ucp_amo_memtype_unpack_reply_buffer(ucp_request_t *req)
 {
     ucp_dt_contig_unpack(req->send.ep->worker, req->send.amo.reply_buffer,
                          &req->send.amo.result, req->send.state.dt_iter.length,
-                         ucp_amo_request_reply_mem_type(req),
+                         ucp_amo_request_reply_mem_type(req), NULL,
                          req->send.state.dt_iter.length);
 }
 
@@ -73,7 +73,10 @@ ucp_proto_amo_progress(uct_pending_req_t *self, ucp_operation_id_t op_id,
                                     UCS_MEMORY_TYPE_HOST;
             ucp_dt_contig_pack(ep->worker, &req->send.amo.value,
                                req->send.state.dt_iter.type.contig.buffer,
-                               op_size, mem_type, op_size);
+                               op_size, mem_type,
+                               is_memtype ? &req->send.state.dt_iter.mem_info :
+                                            NULL,
+                               op_size);
             req->flags |= UCP_REQUEST_FLAG_PROTO_AMO_PACKED;
         }
 
@@ -86,7 +89,8 @@ ucp_proto_amo_progress(uct_pending_req_t *self, ucp_operation_id_t op_id,
         if (op_id == UCP_OP_ID_AMO_CSWAP) {
             ucp_dt_contig_pack(ep->worker, &req->send.amo.result,
                                req->send.amo.reply_buffer, op_size,
-                               ucp_amo_request_reply_mem_type(req), op_size);
+                               ucp_amo_request_reply_mem_type(req), NULL,
+                               op_size);
         }
 
         status = ucp_ep_rma_handle_fence(ep, req, UCS_BIT(spriv->super.lane));

src/ucp/rma/amo_sw.c

@@ -39,10 +39,11 @@ static size_t ucp_amo_sw_pack(void *dest, ucp_request_t *req, int fetch,
 
     if (worker->context->config.ext.proto_enable) {
         ucp_dt_contig_pack(worker, atomich + 1, &req->send.amo.value, size,
-                           UCS_MEMORY_TYPE_HOST, size);
+                           UCS_MEMORY_TYPE_HOST, NULL, size);
         if (req->send.amo.uct_op == UCT_ATOMIC_OP_CSWAP) {
-            ucp_dt_contig_pack(worker, cswaph, req->send.amo.reply_buffer, size,
-                               ucp_amo_request_reply_mem_type(req), size);
+            ucp_dt_contig_pack(worker, cswaph, req->send.amo.reply_buffer,
+                               size, ucp_amo_request_reply_mem_type(req), NULL,
+                               size);
             length += size;
         }
     } else {
@@ -305,7 +306,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_atomic_rep_handler, (arg, data, length, am_fl
     if (worker->context->config.ext.proto_enable) {
         ucp_dt_contig_unpack(worker, req->send.amo.reply_buffer, hdr + 1,
                              frag_length, ucp_amo_request_reply_mem_type(req),
-                             frag_length);
+                             NULL, frag_length);
     } else {
         memcpy(req->send.buffer, hdr + 1, frag_length);
     }