UCT/CUDA_IPC: Support VMM with multiple memory allocations

Handle CUDA VMM allocations spanning multiple cuMemCreate chunks by
discovering all chunks, exporting their fabric handles into a GPU
metadata buffer, and sharing that buffer's fabric handle via the rkey.

On the receiver, the metadata is fetched and a persistent contiguous
VA mapping is created by importing each chunk individually. Address
translation for put/get uses this mapping directly.

Author: tomerg-nvidia

src/uct/cuda/cuda_ipc/cuda_ipc.inl

@@ -83,22 +83,36 @@ uct_cuda_ipc_check_and_pop_ctx(int is_ctx_pushed)
 }
 
 static UCS_F_ALWAYS_INLINE ucs_status_t uct_cuda_ipc_get_remote_address(
-        uct_cuda_ipc_extended_rkey_t *rkey, uint64_t raddr, CUdevice cu_dev,
+        uct_cuda_ipc_unpacked_rkey_t *rkey, uint64_t raddr, CUdevice cu_dev,
         void **laddr_p, void **base_addr_p)
 {
     ucs_status_t status;
     ptrdiff_t offset;
     void *mapped_addr;
 
-    status = uct_cuda_ipc_map_memhandle(rkey, cu_dev, &mapped_addr,
+#if HAVE_CUDA_FABRIC
+    if (rkey->super.super.ph.handle_type ==
+            UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM_MULTI) {
+        offset = raddr - rkey->super.super.d_bptr;
+        ucs_assertv(offset <= rkey->super.super.b_len,
+                    "offset:%ld b_len:%lu", offset, rkey->super.super.b_len);
+        *laddr_p = (void*)(rkey->contig_va + offset);
+        if (base_addr_p != NULL) {
+            *base_addr_p = NULL;
+        }
+        return UCS_OK;
+    }
+#endif
+
+    status = uct_cuda_ipc_map_memhandle(&rkey->super, cu_dev, &mapped_addr,
                                         UCS_LOG_LEVEL_ERROR);
     if (ucs_unlikely(status != UCS_OK)) {
         return status;
     }
 
-    offset = UCS_PTR_BYTE_DIFF(rkey->super.d_bptr, raddr);
-    ucs_assertv(offset <= rkey->super.b_len, "offset:%ld b_len:%lu", offset,
-                rkey->super.b_len);
+    offset = UCS_PTR_BYTE_DIFF(rkey->super.super.d_bptr, raddr);
+    ucs_assertv(offset <= rkey->super.super.b_len, "offset:%ld b_len:%lu",
+                offset, rkey->super.super.b_len);
     *laddr_p = UCS_PTR_BYTE_OFFSET(mapped_addr, offset);
     if (base_addr_p != NULL) {
         *base_addr_p = mapped_addr;

src/uct/cuda/cuda_ipc/cuda_ipc_cache.c

@@ -297,14 +297,6 @@ uct_cuda_ipc_open_memhandle_legacy(CUipcMemHandle memh, CUdevice cu_dev,
 }
 
 #if HAVE_CUDA_FABRIC
-static void
-uct_cuda_ipc_init_access_desc(CUmemAccessDesc *access_desc, CUdevice cu_dev)
-{
-    access_desc->location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-    access_desc->flags         = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
-    access_desc->location.id   = cu_dev;
-}
-
 static ucs_status_t
 uct_cuda_ipc_open_memhandle_vmm(const uct_cuda_ipc_rkey_t *key, CUdevice cu_dev,
                                 CUdeviceptr *mapped_addr,

src/uct/cuda/cuda_ipc/cuda_ipc_ep.c

@@ -120,9 +120,8 @@ uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
         return status;
     }
 
-    status = uct_cuda_ipc_get_remote_address(&key->super, remote_addr,
-                                             cuda_device, &mapped_rem_addr,
-                                             &mapped_addr);
+    status = uct_cuda_ipc_get_remote_address(key, remote_addr, cuda_device,
+                                             &mapped_rem_addr, &mapped_addr);
     if (ucs_unlikely(status != UCS_OK)) {
         goto out;
     }

src/uct/cuda/cuda_ipc/cuda_ipc_iface.c

@@ -311,6 +311,10 @@ static void uct_cuda_ipc_complete_event(uct_iface_h tl_iface,
                                                                uct_cuda_ipc_event_desc_t);
     ucs_status_t status;
 
+    if (cuda_ipc_event->mapped_addr == NULL) {
+        return; /* VMM_MULTI persistent mapping: cleanup at rkey_release */
+    }
+
     status = uct_cuda_ipc_unmap_memhandle(cuda_ipc_event->pid,
                                           cuda_ipc_event->pid_ns,
                                           cuda_ipc_event->d_bptr,