From 9c6d3e8510487977791c3779c473e853a2c3d5e5 Mon Sep 17 00:00:00 2001 From: Leonid Genkin Date: Thu, 9 Jul 2026 12:23:14 +0300 Subject: [PATCH] UCT/CUDA: Cache allocation flags with active context - v1.22.x --- src/uct/cuda/cuda_copy/cuda_copy_md.c | 17 +++++++ .../gtest/uct/cuda/test_switch_cuda_device.cc | 44 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 80798e8a15a..de76992664c 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/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; diff --git a/test/gtest/uct/cuda/test_switch_cuda_device.cc b/test/gtest/uct/cuda/test_switch_cuda_device.cc index e5fa35275c1..c5a5ebf0643 100644 --- a/test/gtest/uct/cuda/test_switch_cuda_device.cc +++ b/test/gtest/uct/cuda/test_switch_cuda_device.cc @@ -8,6 +8,7 @@ #include extern "C" { +#include #include #include } @@ -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 {