UCT/CUDA_COPY: Skip dmabuf probe for async managed memory

src/uct/cuda/cuda_copy/cuda_copy_md.c

@@ -927,6 +927,11 @@ uct_cuda_copy_md_is_registrable(uct_cuda_copy_md_t *md,
         return 0;
     }
 
+    /* Managed memory is registered through ODP and is not dmabuf-exportable. */
+    if (mem_info->type == UCS_MEMORY_TYPE_CUDA_MANAGED) {
+        return 1;
+    }
+
     /* Host-located CUDA VMM is registerable even if dmabuf export fails. */
     if (is_host_located || (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ||
         !md->config.dmabuf_supported) {
@@ -982,8 +987,9 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
     int is_host_located        = 0;
     CUdevice cur_cuda_device   = CU_DEVICE_INVALID;
     CUdevice avail_cuda_device = CU_DEVICE_INVALID;
+    ucs_memory_info_t detected_mem_info = {};
+    ucs_memory_info_t addr_mem_info     = {};
     ucs_memory_info_t cached_mem_info;
-    ucs_memory_info_t addr_mem_info;
     ucs_status_t cache_status;
     ucs_status_t status;
 
@@ -1009,8 +1015,10 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
             return status;
         }
 
-        /* CUDA reports device symbols as device memory, so preserve the type
-         * explicitly provided by the UCM allocation event. */
+        /* Preserve the driver-reported type for flags. CUDA reports device
+         * symbols as device memory, so use the UCM event type for callers. */
+        detected_mem_info = addr_mem_info;
+
         if ((cache_status == UCS_OK) &&
             (cached_mem_info.type == UCS_MEMORY_TYPE_CUDA_MANAGED)) {
             addr_mem_info.type = cached_mem_info.type;
@@ -1064,7 +1072,7 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
 
     if (address != NULL) {
         addr_mem_info.mem_flags = uct_cuda_copy_md_detect_mem_flags(
-                md, &addr_mem_info, is_async_managed, is_host_located,
+                md, &detected_mem_info, is_async_managed, is_host_located,
                 dmabuf_queried ? &dmabuf : NULL);
         ucs_memtype_cache_update(addr_mem_info.base_address,
                                  addr_mem_info.alloc_length, addr_mem_info.type,

test/gtest/common/mem_buffer.cc

@@ -52,6 +52,31 @@
 
 #endif
 
+#if HAVE_CUDA && CUDART_VERSION >= 13000
+static cudaError_t
+mem_buffer_cuda_get_default_managed_pool(cudaMemPool_t *pool_p)
+{
+    cudaMemLocation location = {};
+    int device;
+    cudaError_t cuda_status;
+
+    cuda_status = cudaGetDevice(&device);
+    if (cuda_status != cudaSuccess) {
+        return cuda_status;
+    }
+
+    location.type = cudaMemLocationTypeDevice;
+    location.id   = device;
+    cuda_status   = cudaMemGetDefaultMemPool(pool_p, &location,
+                                             cudaMemAllocationTypeManaged);
+    if (cuda_status != cudaSuccess) {
+        (void)cudaGetLastError(); /* do not leak the error to other tests */
+    }
+
+    return cuda_status;
+}
+#endif
+
 #if HAVE_ROCM
 #include <hip_runtime.h>
 #include <hip_version.h>
@@ -473,7 +498,22 @@ void *mem_buffer::allocate(size_t size, ucs_memory_type_t mem_type, bool async)
         }
         return ptr;
     case UCS_MEMORY_TYPE_CUDA_MANAGED:
-        CUDA_CALL(cudaMallocManaged(&ptr, size), ": size=" << size);
+        if (async) {
+#if CUDART_VERSION >= 13000
+            cudaMemPool_t pool;
+
+            CUDA_CALL(mem_buffer_cuda_get_default_managed_pool(&pool), "");
+            CUDA_CALL(cudaMallocFromPoolAsync(&ptr, size, pool, 0),
+                      ": size=" << size);
+            CUDA_CALL(cudaStreamSynchronize(0), "");
+#else
+            UCS_TEST_ABORT("asynchronous allocation for " +
+                           std::string(ucs_memory_type_names[mem_type]) +
+                           " memory type is not supported");
+#endif
+        } else {
+            CUDA_CALL(cudaMallocManaged(&ptr, size), ": size=" << size);
+        }
         return ptr;
 #endif
 #if HAVE_ROCM
@@ -549,6 +589,7 @@ void mem_buffer::release(void *ptr, ucs_memory_type_t mem_type, bool async)
             break;
 #if HAVE_CUDA
         case UCS_MEMORY_TYPE_CUDA:
+        case UCS_MEMORY_TYPE_CUDA_MANAGED:
             if (async) {
 #if CUDART_VERSION >= 11020
                 cudaStreamSynchronize(0);
@@ -562,9 +603,6 @@ void mem_buffer::release(void *ptr, ucs_memory_type_t mem_type, bool async)
                 CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr);
             }
             break;
-        case UCS_MEMORY_TYPE_CUDA_MANAGED:
-            CUDA_CALL(cudaFree(ptr), ": ptr=" << ptr);
-            break;
 #endif
 #if HAVE_ROCM
         case UCS_MEMORY_TYPE_ROCM:
@@ -851,10 +889,21 @@ bool mem_buffer::is_async_supported(ucs_memory_type_t mem_type)
     if (!is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) {
         return false;
     }
-    return mem_type == UCS_MEMORY_TYPE_CUDA;
-#else
-    return false;
+
+    if (mem_type == UCS_MEMORY_TYPE_CUDA) {
+        return true;
+    }
+
+#if CUDART_VERSION >= 13000
+    if (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) {
+        cudaMemPool_t pool;
+
+        return mem_buffer_cuda_get_default_managed_pool(&pool) == cudaSuccess;
+    }
 #endif
+#endif
+
+    return false;
 }
 
 mem_buffer::mem_buffer(size_t size, ucs_memory_type_t mem_type) :

test/gtest/common/mem_buffer.h

@@ -124,7 +124,7 @@ class mem_buffer {
     /* Allocation mode. */
     enum class alloc_mode {
         DEFAULT, /* Default allocation mode, using cudaMalloc */
-        ASYNC    /* Asynchronous allocation mode, using cudaMallocAsync */
+        ASYNC    /* Asynchronous CUDA allocation mode */
     };
 
     mem_buffer(size_t size, ucs_memory_type_t mem_type);

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -342,6 +342,44 @@ class test_mem_alloc_device : public test_switch_cuda_device {
         EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
     }
 
+    void check_managed_memory_context(void *address)
+    {
+        CUpointer_attribute attr_type = CU_POINTER_ATTRIBUTE_CONTEXT;
+        CUcontext cuda_mem_ctx        = nullptr;
+
+        ASSERT_EQ(CUDA_SUCCESS,
+                  cuPointerGetAttribute(&cuda_mem_ctx, attr_type,
+                                        (CUdeviceptr)address));
+        EXPECT_NE(nullptr, cuda_mem_ctx);
+    }
+
+    void query_managed_registrable(void *address, size_t size)
+    {
+        uct_md_mem_attr_v2_t mem_attr = {};
+
+        mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
+                              UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
+        EXPECT_UCS_OK(uct_md_mem_query_v2(md(), address, size, &mem_attr));
+        EXPECT_EQ(UCS_MEMORY_TYPE_CUDA_MANAGED, mem_attr.mem_type);
+        EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
+    }
+
+    void test_async_managed_mem_pool_registrable()
+    {
+        constexpr size_t size = 192;
+
+        if (!mem_buffer::is_async_supported(UCS_MEMORY_TYPE_CUDA_MANAGED)) {
+            UCS_TEST_SKIP_R("asynchronous CUDA managed memory is not "
+                            "supported");
+        }
+
+        mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED,
+                          mem_buffer::alloc_mode::ASYNC);
+        buffer.memset(0);
+
+        query_managed_registrable(buffer.ptr(), size);
+    }
+
 private:
     std::vector<ucs_sys_device_t> m_sys_dev;
 
@@ -441,6 +479,46 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable,
     EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr));
 }
 
+UCS_TEST_P(test_mem_alloc_device, user_managed_mem_registrable,
+           "CUDA_COPY_PREF_LOC=gpu")
+{
+    constexpr size_t size = 4 * UCS_MBYTE;
+    mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA_MANAGED);
+
+    check_managed_memory_context(buffer.ptr());
+    query_managed_registrable(buffer.ptr(), size);
+}
+
+UCS_TEST_P(test_mem_alloc_device, uct_alloc_managed_mem_registrable,
+           "CUDA_COPY_PREF_LOC=gpu")
+{
+    ASSERT_UCS_OK(allocate(UCS_MEMORY_TYPE_CUDA_MANAGED));
+
+    check_managed_memory_context(mem.address);
+    query_managed_registrable(mem.address, mem.length);
+
+    EXPECT_UCS_OK(uct_mem_free(&mem));
+}
+
+UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_registrable,
+           "CUDA_COPY_DMABUF=try")
+{
+    test_async_managed_mem_pool_registrable();
+}
+
+UCS_TEST_P(test_mem_alloc_device, async_managed_mem_pool_gpu_pref_loc,
+           "CUDA_COPY_DMABUF=try", "CUDA_COPY_PREF_LOC=gpu")
+{
+    test_async_managed_mem_pool_registrable();
+}
+
+UCS_TEST_P(test_mem_alloc_device,
+           async_managed_mem_pool_registrable_cuda_type,
+           "CUDA_COPY_DMABUF=try", "CUDA_COPY_ASYNC_MEM_TYPE=cuda")
+{
+    test_async_managed_mem_pool_registrable();
+}
+
 UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable,
            "CUDA_COPY_ASYNC_MEM_TYPE=cuda")
 {