Merge pull request #10405 from Akshay-Venkatesh/topic/cuda-ipc-buffer-id-va-recycle

UCT/CUDA_IPC: Use buffer id to detect VA recylcing

Author: yosefe

src/uct/cuda/cuda_ipc/cuda_ipc_cache.c

@@ -124,7 +124,8 @@ static ucs_status_t uct_cuda_ipc_close_memhandle(uct_cuda_ipc_cache_region_t *re
                         (CUdeviceptr)region->mapped_addr, region->key.b_len));
         }
     } else if (region->key.ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL) {
-        return UCT_CUDADRV_FUNC_LOG_WARN(cuMemPoolDestroy(region->key.ph.pool));
+        return UCT_CUDADRV_FUNC_LOG_WARN(
+                cuMemFree((CUdeviceptr)region->mapped_addr));
     } else
 #endif
     {
@@ -335,22 +336,22 @@ static ucs_status_t uct_cuda_ipc_open_memhandle(uct_cuda_ipc_rkey_t *key,
                                                 CUdeviceptr *mapped_addr)
 {
 
-#if HAVE_CUDA_FABRIC
     ucs_trace("key handle type %u", key->ph.handle_type);
 
-    if (key->ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY) {
+    switch(key->ph.handle_type) {
+    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY:
         return uct_cuda_ipc_open_memhandle_legacy(key->ph.handle.legacy,
                                                   mapped_addr);
-    } else if (key->ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM) {
+#if HAVE_CUDA_FABRIC
+    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM:
         return uct_cuda_ipc_open_memhandle_vmm(key, mapped_addr);
-    } else if (key->ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL) {
+    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL:
         return uct_cuda_ipc_open_memhandle_mempool(key, mapped_addr);
-    } else {
+#endif
+    default:
+        ucs_error("unsupported key handle type");
         return UCS_ERR_INVALID_PARAM;
     }
-#else
-    return uct_cuda_ipc_open_memhandle_legacy(key->ph, mapped_addr);
-#endif
 }
 
 static void uct_cuda_ipc_cache_invalidate_regions(uct_cuda_ipc_cache_t *cache,
@@ -480,13 +481,6 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_map_memhandle, (key, mapped_addr),
     ucs_pgt_region_t *pgt_region;
     uct_cuda_ipc_cache_region_t *region;
     int ret;
-    size_t cmp_size;
-
-#if HAVE_CUDA_FABRIC
-    cmp_size = sizeof(key->ph.handle);
-#else
-    cmp_size = sizeof(key->ph);
-#endif
 
     status = uct_cuda_ipc_get_remote_cache(key->pid, &cache);
     if (status != UCS_OK) {
@@ -498,8 +492,8 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_map_memhandle, (key, mapped_addr),
                                   &cache->pgtable, key->d_bptr);
     if (ucs_likely(pgt_region != NULL)) {
         region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);
-        if (memcmp((const void *)&key->ph, (const void *)&region->key.ph,
-                   cmp_size) == 0) {
+
+        if (key->ph.buffer_id == region->key.ph.buffer_id) {
             /*cache hit */
             ucs_trace("%s: cuda_ipc cache hit addr:%p size:%lu region:"
                       UCS_PGT_REGION_FMT, cache->name, (void *)key->d_bptr,

src/uct/cuda/cuda_ipc/cuda_ipc_md.c

@@ -70,11 +70,7 @@ uct_cuda_ipc_get_dev_cache(uct_cuda_ipc_component_t *component,
     int ret;
 
     key.uuid = rkey->uuid;
-#if HAVE_CUDA_FABRIC
     key.type = rkey->ph.handle_type;
-#else
-    key.type = 0;
-#endif
 
     iter = kh_put(cuda_ipc_uuid_hash, hash, key, &ret);
     if (ret == UCS_KH_PUT_KEY_PRESENT) {
@@ -112,11 +108,10 @@ static ucs_status_t
 uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
                          uct_cuda_ipc_lkey_t **key_p)
 {
-    CUipcMemHandle *legacy_handle;
     uct_cuda_ipc_lkey_t *key;
     ucs_status_t status;
 #if HAVE_CUDA_FABRIC
-#define UCT_CUDA_IPC_QUERY_NUM_ATTRS 2
+#define UCT_CUDA_IPC_QUERY_NUM_ATTRS 3
     CUmemGenericAllocationHandle handle;
     CUmemoryPool mempool;
     CUpointer_attribute attr_type[UCT_CUDA_IPC_QUERY_NUM_ATTRS];
@@ -130,10 +125,16 @@ uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
         return UCS_ERR_NO_MEMORY;
     }
 
-    legacy_handle = (CUipcMemHandle*)&key->ph;
     UCT_CUDADRV_FUNC_LOG_ERR(cuMemGetAddressRange(&key->d_bptr, &key->b_len,
                 (CUdeviceptr)addr));
 
+    status = UCT_CUDADRV_FUNC_LOG_ERR(cuPointerGetAttribute(&key->ph.buffer_id,
+                                      CU_POINTER_ATTRIBUTE_BUFFER_ID,
+                                      (CUdeviceptr)addr));
+    if (status != UCS_OK) {
+        goto err;
+    }
+
 #if HAVE_CUDA_FABRIC
     /* cuda_ipc can handle VMM, mallocasync, and legacy pinned device so need to
      * pack appropriate handle */
@@ -142,6 +143,8 @@ uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
     attr_data[0] = &legacy_capable;
     attr_type[1] = CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES;
     attr_data[1] = &allowed_handle_types;
+    attr_type[2] = CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE;
+    attr_data[2] = &mempool;
 
     status = UCT_CUDADRV_FUNC_LOG_ERR(
             cuPointerGetAttributes(ucs_static_array_size(attr_data), attr_type,
@@ -151,8 +154,6 @@ uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
     }
 
     if (legacy_capable) {
-        key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY;
-        legacy_handle       = &key->ph.handle.legacy;
         goto legacy_path;
     }
 
@@ -184,9 +185,7 @@ uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
         goto common_path;
     }
 
-    status = UCT_CUDADRV_FUNC_LOG_ERR(cuPointerGetAttribute(&mempool,
-                CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE, (CUdeviceptr)addr));
-    if ((status != UCS_OK) || (mempool == 0)) {
+    if (mempool == 0) {
         /* cuda_ipc can only handle UCS_MEMORY_TYPE_CUDA, which has to be either
          * legacy type, or VMM type, or mempool type. Return error if memory
          * does not belong to any of the three types */
@@ -216,16 +215,18 @@ uct_cuda_ipc_mem_add_reg(void *addr, uct_cuda_ipc_memh_t *memh,
     goto common_path;
 #endif
 legacy_path:
-    status = UCT_CUDADRV_FUNC(cuIpcGetMemHandle(legacy_handle, (CUdeviceptr)addr),
-                              UCS_LOG_LEVEL_ERROR);
+    key->ph.handle_type = UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY;
+    status              = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuIpcGetMemHandle(&key->ph.handle.legacy, (CUdeviceptr)addr));
     if (status != UCS_OK) {
         goto err;
     }
 
 common_path:
     ucs_list_add_tail(&memh->list, &key->link);
-    ucs_trace("registered addr:%p/%p length:%zd dev_num:%d",
-              addr, (void *)key->d_bptr, key->b_len, (int)memh->dev_num);
+    ucs_trace("registered addr:%p/%p length:%zd dev_num:%d buffer_id:%llu",
+              addr, (void *)key->d_bptr, key->b_len, (int)memh->dev_num,
+              key->ph.buffer_id);
 
     *key_p = key;
     return UCS_OK;

src/uct/cuda/cuda_ipc/cuda_ipc_md.h

@@ -14,28 +14,30 @@
 #include <ucs/config/types.h>
 
 
-#if HAVE_CUDA_FABRIC
 typedef enum uct_cuda_ipc_key_handle {
     UCT_CUDA_IPC_KEY_HANDLE_TYPE_ERROR = 0,
     UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY, /* cudaMalloc memory */
+#if HAVE_CUDA_FABRIC
     UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM, /* cuMemCreate memory */
     UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL /* cudaMallocAsync memory */
+#endif
 } uct_cuda_ipc_key_handle_t;
 
 
 typedef struct uct_cuda_ipc_md_handle {
     uct_cuda_ipc_key_handle_t handle_type;
     union {
         CUipcMemHandle        legacy;        /* Legacy IPC handle */
+#if HAVE_CUDA_FABRIC
         CUmemFabricHandle     fabric_handle; /* VMM/Mallocasync export handle */
+#endif
     } handle;
+#if HAVE_CUDA_FABRIC
     CUmemPoolPtrExportData    ptr;
     CUmemoryPool              pool;
-} uct_cuda_ipc_md_handle_t;
-#else
-typedef CUipcMemHandle uct_cuda_ipc_md_handle_t;
 #endif
-
+    unsigned long long        buffer_id;
+} uct_cuda_ipc_md_handle_t;
 
 /**
  * @brief cuda ipc MD descriptor