UCT/CUDA/CUDA_IPC: Release unused cache entry at rkey release.

src/ucp/core/ucp_ep.c

@@ -371,7 +371,7 @@ ucp_ep_peer_mem_get(ucp_context_h context, ucp_ep_h ep, uint64_t address,
 
     data->size = size;
     ucp_ep_rkey_unpack_internal(ep, rkey_buf, 0, UCS_BIT(rkey_ptr_md_index), 0,
-                                UCS_SYS_DEVICE_ID_UNKNOWN, &data->rkey);
+                                UCS_SYS_DEVICE_ID_UNKNOWN, 0, &data->rkey);
     rkey_index = ucs_bitmap2idx(data->rkey->md_map, rkey_ptr_md_index);
     status     = uct_rkey_ptr(data->rkey->tl_rkey[rkey_index].cmpt,
                               &data->rkey->tl_rkey[rkey_index].rkey, address,

src/ucp/core/ucp_rkey.c

@@ -914,10 +914,11 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rkey_proto_resolve,
 
 UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal,
                  (ep, buffer, length, unpack_md_map, skip_md_map, sys_dev,
-                  rkey_p),
+                  uct_flags, rkey_p),
                  ucp_ep_h ep, const void *buffer, size_t length,
                  ucp_md_map_t unpack_md_map, ucp_md_map_t skip_md_map,
-                 ucs_sys_device_t sys_dev, ucp_rkey_h *rkey_p)
+                 ucs_sys_device_t sys_dev, uint64_t uct_flags,
+                 ucp_rkey_h *rkey_p)
 {
     ucp_worker_h worker              = ep->worker;
     const ucp_ep_config_t *ep_config = ucp_ep_config(ep);
@@ -974,8 +975,10 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal,
     rkey->ep       = ep;
 #endif
 
-    unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE;
+    unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE |
+                               UCT_RKEY_UNPACK_FIELD_FLAGS;
     unpack_params.sys_device = sys_dev;
+    unpack_params.flags      = uct_flags;
 
     /* Go over remote MD indices and unpack rkey of each UCT MD */
     rkey_index = 0; /* Index of the rkey in the array */
@@ -1056,18 +1059,34 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal,
     return status;
 }
 
-ucs_status_t ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer,
-                                ucp_rkey_h *rkey_p)
+static ucs_status_t
+ucp_ep_rkey_unpack_flags(ucp_ep_h ep, const void *rkey_buffer,
+                         uint64_t uct_flags, ucp_rkey_h *rkey_p)
 {
     ucs_status_t status;
 
     UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(ep->worker);
-    status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, rkey_p);
+    status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, uct_flags,
+                                          rkey_p);
     UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(ep->worker);
 
     return status;
 }
 
+ucs_status_t
+ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer, ucp_rkey_h *rkey_p)
+{
+    return ucp_ep_rkey_unpack_flags(ep, rkey_buffer,
+                                    UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING,
+                                    rkey_p);
+}
+
+ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer,
+                                         ucp_rkey_h *rkey_p)
+{
+    return ucp_ep_rkey_unpack_flags(ep, rkey_buffer, 0, rkey_p);
+}
+
 void ucp_rkey_dump_packed(const void *buffer, size_t length,
                           ucs_string_buffer_t *strb)
 {

src/ucp/core/ucp_rkey.h

@@ -231,11 +231,15 @@ ucp_memh_exported_unpack(ucp_context_h context, const void *export_mkey_buffer,
 int ucp_memh_buffer_is_dummy(const void *exported_memh_buffer);
 
 
-ucs_status_t ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer,
-                                         size_t length,
-                                         ucp_md_map_t unpack_md_map,
-                                         ucp_md_map_t skip_md_map,
-                                         ucs_sys_device_t sys_device,
+ucs_status_t
+ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer, size_t length,
+                            ucp_md_map_t unpack_md_map,
+                            ucp_md_map_t skip_md_map,
+                            ucs_sys_device_t sys_device, uint64_t uct_flags,
+                            ucp_rkey_h *rkey_p);
+
+/* Unpack a remote key passing zero flags to uct_rkey_unpack */
+ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer,
                                          ucp_rkey_h *rkey_p);
 
 

src/ucp/core/ucp_rkey.inl

@@ -67,13 +67,14 @@ ucp_rkey_get_tl_rkey(ucp_rkey_h rkey, ucp_md_index_t rkey_index)
 
 static UCS_F_ALWAYS_INLINE ucs_status_t
 ucp_ep_rkey_unpack_reachable(ucp_ep_h ep, const void *buffer, size_t length,
-                             ucp_rkey_h *rkey_p)
+                             uint64_t uct_flags, ucp_rkey_h *rkey_p)
 {
     ucp_ep_config_t *config = &ucs_array_elem(&ep->worker->ep_config,
                                               ep->cfg_index);
     return ucp_ep_rkey_unpack_internal(ep, buffer, length,
                                        config->key.reachable_md_map, 0,
-                                       UCS_SYS_DEVICE_ID_UNKNOWN, rkey_p);
+                                       UCS_SYS_DEVICE_ID_UNKNOWN, uct_flags,
+                                       rkey_p);
 }
 
 static UCS_F_ALWAYS_INLINE int

src/ucp/rndv/proto_rndv.c

@@ -857,7 +857,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_proto_rndv_send_reply,
          */
         status = ucp_ep_rkey_unpack_internal(
                 ep, rkey_buffer, rkey_length, ep_config->key.reachable_md_map,
-                ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, &rkey);
+                ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, 0, &rkey);
         if (status != UCS_OK) {
             goto err;
         }

src/ucp/rndv/rndv.c

@@ -952,7 +952,8 @@ static ucs_status_t ucp_rndv_req_send_rma_get(ucp_request_t *rndv_req,
 
     ucp_request_set_super(rndv_req, rreq);
 
-    status = ucp_ep_rkey_unpack(ep, rkey_buf, &rndv_req->send.rndv.rkey);
+    status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf,
+                                         &rndv_req->send.rndv.rkey);
     if (status != UCS_OK) {
         ucs_fatal("failed to unpack rendezvous remote key received from %s: %s",
                   ucp_ep_peer_name(ep), ucs_status_string(status));
@@ -1124,7 +1125,7 @@ ucp_rndv_send_frag_update_get_rkey(ucp_worker_h worker, ucp_request_t *freq,
                            &rkey_buffer, &rkey_size);
     ucs_assert_always(status == UCS_OK);
 
-    status = ucp_ep_rkey_unpack(mem_type_ep, rkey_buffer, rkey_p);
+    status = ucp_ep_rkey_unpack_wo_flags(mem_type_ep, rkey_buffer, rkey_p);
     ucs_assert_always(status == UCS_OK);
     ucp_rkey_buffer_release(rkey_buffer);
 }
@@ -1303,8 +1304,8 @@ ucp_rndv_recv_start_get_pipeline(ucp_worker_h worker, ucp_request_t *rndv_req,
      * Step 3: Send ATS for RNDV request
      */
 
-    status = ucp_ep_rkey_unpack(rndv_req->send.ep, rkey_buffer,
-                                &rndv_req->send.rndv.rkey);
+    status = ucp_ep_rkey_unpack_wo_flags(rndv_req->send.ep, rkey_buffer,
+                                         &rndv_req->send.rndv.rkey);
     if (ucs_unlikely(status != UCS_OK)) {
         ucs_fatal("failed to unpack rendezvous remote key received from %s: %s",
                   ucp_ep_peer_name(rndv_req->send.ep), ucs_status_string(status));
@@ -1500,7 +1501,7 @@ static void ucp_rndv_do_rkey_ptr(ucp_request_t *rndv_req, ucp_request_t *rreq,
 
     ucp_trace_req(rndv_req, "start rkey_ptr rndv rreq %p", rreq);
 
-    status = ucp_ep_rkey_unpack(ep, rkey_buf, &rkey);
+    status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf, &rkey);
     if (status != UCS_OK) {
         ucs_fatal("failed to unpack rendezvous remote key received from %s: %s",
                   ucp_ep_peer_name(ep), ucs_status_string(status));
@@ -2123,7 +2124,8 @@ static ucs_status_t ucp_rndv_send_start_put_pipeline(ucp_request_t *sreq,
     ucp_trace_req(sreq, "using put rndv pipeline protocol: va %p size %zu",
                   (void*)rndv_rtr_hdr->address, rndv_rtr_hdr->size);
 
-    status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1, &sreq->send.rndv.rkey);
+    status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1,
+                                         &sreq->send.rndv.rkey);
     if (status != UCS_OK) {
         ucs_fatal("failed to unpack rendezvous remote key received from %s: %s",
                   ucp_ep_peer_name(ep), ucs_status_string(status));
@@ -2325,8 +2327,8 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler,
              * and we have to use PUT_ZCOPY RNDV scheme instead */
         }
 
-        status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1,
-                                    &sreq->send.rndv.rkey);
+        status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1,
+                                             &sreq->send.rndv.rkey);
         if (status != UCS_OK) {
             ucs_fatal("failed to unpack rendezvous remote key received from %s: %s",
                       ucp_ep_peer_name(ep), ucs_status_string(status));

src/uct/api/v2/uct_v2.h

@@ -764,10 +764,25 @@ typedef struct uct_rkey_compare_params {
  * @brief Rkey unpack parameters field mask.
  */
 typedef enum {
-    UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0)  /**< sys_device field */
+    UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0), /**< sys_device field */
+    UCT_RKEY_UNPACK_FIELD_FLAGS      = UCS_BIT(1) /**< flags field */
 } uct_rkey_unpack_field_mask_t;
 
 
+/**
+ * @ingroup UCT_MD
+ * @brief Rkey unpack flags.
+ */
+typedef enum {
+    /**
+     * When the remote key is released (see @ref uct_rkey_release), free the
+     * local mapping of the remote memory which was created for this remote key
+     * and is not used by any other remote key or operation.
+     */
+    UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING = UCS_BIT(0)
+} uct_rkey_unpack_flags_t;
+
+
 /**
  * @ingroup UCT_MD
  * @brief Parameters for unpacking remote key using @ref uct_rkey_unpack_v2.
@@ -785,6 +800,12 @@ typedef struct uct_rkey_unpack_params {
      * (default behavior).
      */
     ucs_sys_device_t     sys_device;
+
+    /**
+     * Flags to unpack the rkey with, using bits from
+     * @ref uct_rkey_unpack_flags_t.
+     */
+    uint64_t             flags;
 } uct_rkey_unpack_params_t;
 
 

src/uct/cuda/cuda_ipc/cuda_ipc_cache.c

@@ -647,14 +647,23 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr,
     /* use write lock because cache maybe modified */
     pthread_rwlock_wrlock(&cache->lock);
     pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr);
-    ucs_assert(pgt_region != NULL);
-    region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);
+    if (pgt_region == NULL) {
+        ucs_assert(mapped_addr == NULL);
+        pthread_rwlock_unlock(&cache->lock);
+        return;
+    }
 
-    ucs_assert(region->refcount >= 1);
-    region->refcount--;
+    region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);
+    if (mapped_addr != NULL) {
+        ucs_assert(region->refcount >= 1);
+        region->refcount--;
+    }
 
     if (!region->refcount && !cache_enabled) {
-        ucs_assert(region->mapped_addr == mapped_addr);
+        if (mapped_addr != NULL) {
+            ucs_assert(region->mapped_addr == mapped_addr);
+        }
+
         uct_cuda_ipc_cache_region_destroy(cache, region);
     }
 

src/uct/cuda/cuda_ipc/cuda_ipc_md.c

@@ -39,6 +39,7 @@
 typedef struct {
     const void *mapped_addr;
     CUdevice   cu_dev;
+    int        free_unused_mapping;
 } uct_cuda_ipc_rkey_handle_t;
 
 typedef struct {
@@ -474,7 +475,8 @@ uct_cuda_ipc_mkey_pack(uct_md_h md, uct_mem_h tl_memh, void *address,
                                                     memh->dev_num));
 }
 
-static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle()
+static uct_cuda_ipc_rkey_handle_t *
+uct_cuda_ipc_create_rkey_handle(uint64_t flags)
 {
     uct_cuda_ipc_rkey_handle_t *rkey_handle;
 
@@ -484,8 +486,10 @@ static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle()
         return NULL;
     }
 
-    rkey_handle->mapped_addr = NULL;
-    rkey_handle->cu_dev      = CU_DEVICE_INVALID;
+    rkey_handle->mapped_addr         = NULL;
+    rkey_handle->cu_dev              = CU_DEVICE_INVALID;
+    rkey_handle->free_unused_mapping = !!(
+            flags & UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING);
     return rkey_handle;
 }
 
@@ -562,13 +566,15 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack,
     uct_cuda_ipc_unpacked_rkey_t *unpacked;
     uct_cuda_ipc_rkey_handle_t *rkey_handle;
     ucs_sys_device_t sys_dev;
+    uint64_t flags;
     CUdevice cuda_device;
     CUdevice avail_cuda_device;
     ucs_status_t status;
     const uct_cuda_ipc_extended_rkey_t *ext_rkey;
 
     sys_dev = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, sys_device,
                               SYS_DEVICE, UCS_SYS_DEVICE_ID_UNKNOWN);
+    flags   = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, flags, FLAGS, 0);
 
     unpacked = ucs_malloc(sizeof(*unpacked), "uct_cuda_ipc_unpacked_rkey_t");
     if (NULL == unpacked) {
@@ -580,7 +586,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack,
     unpacked->super.super      = *packed;
     unpacked->super.super.pid &= ~UCT_CUDA_IPC_RKEY_FLAG_PID_NS;
 
-    rkey_handle = uct_cuda_ipc_create_rkey_handle();
+    rkey_handle = uct_cuda_ipc_create_rkey_handle(flags);
     if (rkey_handle == NULL) {
         status = UCS_ERR_NO_MEMORY;
         goto err_free_key;
@@ -633,8 +639,12 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey,
     uct_cuda_ipc_extended_rkey_t *extended_rkey;
     uct_cuda_ipc_rkey_t *rkey;
 
-    if ((uct_rkey == 0) || (rkey_handle == NULL) ||
-        (rkey_handle->mapped_addr == NULL)) {
+    if ((uct_rkey == 0) || (rkey_handle == NULL)) {
+        return;
+    }
+
+    if ((rkey_handle->mapped_addr == NULL) &&
+        !rkey_handle->free_unused_mapping) {
         return;
     }
 
@@ -644,7 +654,8 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey,
 
     uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, rkey->d_bptr,
                                  rkey_handle->mapped_addr, rkey_handle->cu_dev,
-                                 uct_cuda_ipc_component.enable_remote_cache);
+                                 uct_cuda_ipc_component.enable_remote_cache &&
+                                         !rkey_handle->free_unused_mapping);
 }
 
 static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component,

test/gtest/ucp/test_ucp_mmap.cc

@@ -245,7 +245,7 @@ ucp_rkey_h test_ucp_mmap::mem_chunk::unpack(ucp_ep_h ep, ucp_md_map_t md_map)
         // Different MD map means different config index on proto v2
         ASSERT_UCS_OK(ucp_ep_rkey_unpack_internal(ep, rkey_buffer, rkey_size,
                                                   md_map, 0,
-                                                  UCS_SYS_DEVICE_ID_UNKNOWN,
+                                                  UCS_SYS_DEVICE_ID_UNKNOWN, 0,
                                                   &rkey));
     }
 
@@ -497,7 +497,7 @@ void test_ucp_mmap::test_rkey_proto(ucp_mem_h memh)
     /* Unpack remote key buffer */
     ucp_rkey_h rkey;
     status = ucp_ep_rkey_unpack_reachable(receiver().ep(), &rkey_buffer[0],
-                                          rkey_size, &rkey);
+                                          rkey_size, 0, &rkey);
     ASSERT_UCS_OK(status);
 
     /* Check rkey configuration and lanes distances */

test/mpi/test_rma_cuda.c

@@ -167,7 +167,7 @@ int main(int argc, char **argv)
         return -1;
     }
 
-    xfer_args.ucp = create_ucp();
+    xfer_args.ucp = create_ucp(UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP);
 
     for (operation_idx = 0; operation_idx < ucs_static_array_size(operations);
          ++operation_idx) {

test/mpi/test_ucp.c

@@ -11,14 +11,14 @@
 
 #include <mpi.h>
 
-static ucp_context_h create_ucp_context()
+static ucp_context_h create_ucp_context(uint64_t features)
 {
     ucp_params_t params;
     ucp_config_t *config;
     ucp_context_h context;
 
     params.field_mask = UCP_PARAM_FIELD_FEATURES | UCP_PARAM_FIELD_NAME;
-    params.features   = UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP;
+    params.features   = features;
     params.name       = "test_context";
     UCX_CHECK(ucp_config_read(NULL, NULL, &config));
     UCX_CHECK(ucp_init(&params, config, &context));
@@ -124,10 +124,10 @@ create_ucp_mem(void *buffer, size_t size, ucp_context_h context)
     return mem;
 }
 
-ucp_t create_ucp()
+ucp_t create_ucp(uint64_t features)
 {
     ucp_t ucp;
-    ucp.context = create_ucp_context();
+    ucp.context = create_ucp_context(features);
     ucp.worker  = create_ucp_worker(ucp.context);
     ucp.ep      = create_ucp_ep(ucp.worker);
     return ucp;

test/mpi/test_ucp.h

@@ -22,7 +22,7 @@ typedef struct {
 } rkey_t;
 
 
-ucp_t create_ucp();
+ucp_t create_ucp(uint64_t features);
 
 
 void destroy_ucp(ucp_t);