UCT/CUDA: Cache allocation flags with active context - v1.22.x

src/uct/cuda/cuda_copy/cuda_copy_md.c

@@ -352,6 +352,7 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
                         uct_mem_h *memh_p)
 {
     uct_cuda_copy_md_t *md = ucs_derived_of(uct_md, uct_cuda_copy_md_t);
+    uct_md_mem_attr_v2_t mem_attr;
     ucs_status_t status;
     uct_cuda_copy_alloc_handle_t *alloc_handle;
     ucs_log_level_t log_level;
@@ -426,6 +427,22 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
 allocated:
     uct_cuda_copy_sync_memops(alloc_handle->ptr, alloc_handle->is_vmm);
 
+    /* Cache memory flags as part of uct_cuda_copy_md_mem_query() before
+     * restoring the CUDA context.
+     */
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE     |
+                          UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV      |
+                          UCT_MD_MEM_ATTR_V2_FIELD_BASE_ADDRESS |
+                          UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH |
+                          UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
+
+    status = uct_cuda_copy_md_mem_query(uct_md, (void*)alloc_handle->ptr,
+                                        alloc_handle->length, &mem_attr);
+    if (status != UCS_OK) {
+        (void)uct_md_mem_free(uct_md, alloc_handle);
+        goto out;
+    }
+
     *memh_p    = alloc_handle;
     *address_p = (void*)alloc_handle->ptr;
     *length_p  = alloc_handle->length;

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -8,6 +8,7 @@
 #include <uct/test_p2p_rma.h>
 
 extern "C" {
+#include <ucs/memory/memtype_cache.h>
 #include <ucs/sys/ptr_arith.h>
 #include <uct/base/uct_md.h>
 }
@@ -474,6 +475,49 @@ UCS_TEST_P(test_mem_alloc_device, same_device_cuda_fabric_implicit,
     test_same_device_alloc(UCS_MEMORY_TYPE_CUDA, false);
 }
 
+UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
+           "CUDA_COPY_ASYNC_MEM_TYPE=cuda")
+{
+    ucs_memory_info_t mem_info = {};
+    ucs_status_t lookup_status = UCS_ERR_NO_ELEM;
+    ucs_status_t alloc_status  = UCS_ERR_NO_MEMORY;
+    CUresult ctx_status        = CUDA_ERROR_UNKNOWN;
+    CUcontext cuda_ctx         = nullptr;
+
+    /* UCP initializes the cache before allocating rendezvous fragments. */
+    ucs_memory_info_t init_mem_info = {};
+    (void)ucs_memtype_cache_lookup(&init_mem_info, sizeof(init_mem_info),
+                                   &init_mem_info);
+
+    std::thread([&]() {
+        ctx_status = cuCtxGetCurrent(&cuda_ctx);
+        if ((ctx_status != CUDA_SUCCESS) || (cuda_ctx != nullptr)) {
+            return;
+        }
+
+        alloc_status = allocate(UCS_MEMORY_TYPE_CUDA);
+        if (alloc_status != UCS_OK) {
+            return;
+        }
+
+        lookup_status = ucs_memtype_cache_lookup(mem.address, mem.length,
+                                                 &mem_info);
+    }).join();
+
+    EXPECT_EQ(CUDA_SUCCESS, ctx_status);
+    EXPECT_EQ(nullptr, cuda_ctx);
+    ASSERT_UCS_OK(alloc_status);
+
+    ucs_status_t free_status = uct_mem_free(&mem);
+
+    EXPECT_UCS_OK(lookup_status);
+    if (lookup_status == UCS_OK) {
+        EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_info.type);
+        EXPECT_TRUE(mem_info.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
+    }
+    EXPECT_UCS_OK(free_status);
+}
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy);
 
 class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test {