Skip to content

Commit b74f06b

Browse files
tomerg-nvidiagleon99
tomerg-nvidia
and
gleon99
authored
UCT/CUDA_COPY: Fix incorrect mem_flags detection - v1.22.x (#11645)
* UCT/CUDA: Cache allocation flags with active context - v1.22.x * UCT/CUDA_COPY: Fix incorrect mem_flags detection --------- Co-authored-by: Leonid Genkin <lgenkin@nvidia.com>
1 parent d73acf1 commit b74f06b

2 files changed

Lines changed: 189 additions & 25 deletions

File tree

‎src/uct/cuda/cuda_copy/cuda_copy_md.c‎

Lines changed: 57 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,7 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
352352
uct_mem_h *memh_p)
353353
{
354354
uct_cuda_copy_md_t *md = ucs_derived_of(uct_md, uct_cuda_copy_md_t);
355+
uct_md_mem_attr_v2_t mem_attr;
355356
ucs_status_t status;
356357
uct_cuda_copy_alloc_handle_t *alloc_handle;
357358
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,
426427
allocated:
427428
uct_cuda_copy_sync_memops(alloc_handle->ptr, alloc_handle->is_vmm);
428429

430+
/* Cache memory flags as part of uct_cuda_copy_md_mem_query() before
431+
* restoring the CUDA context.
432+
*/
433+
mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
434+
UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV |
435+
UCT_MD_MEM_ATTR_V2_FIELD_BASE_ADDRESS |
436+
UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH |
437+
UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
438+
439+
status = uct_cuda_copy_md_mem_query(uct_md, (void*)alloc_handle->ptr,
440+
alloc_handle->length, &mem_attr);
441+
if (status != UCS_OK) {
442+
(void)uct_md_mem_free(uct_md, alloc_handle);
443+
goto out;
444+
}
445+
429446
*memh_p = alloc_handle;
430447
*address_p = (void*)alloc_handle->ptr;
431448
*length_p = alloc_handle->length;
@@ -465,7 +482,8 @@ uct_cuda_copy_mem_release_fabric(uct_cuda_copy_alloc_handle_t *alloc_handle)
465482

466483
static int uct_cuda_copy_detect_vmm(const void *address,
467484
ucs_memory_type_t *vmm_mem_type,
468-
CUdevice *cuda_device)
485+
CUdevice *cuda_device,
486+
int *is_host_located)
469487
{
470488
#ifdef HAVE_CUMEMRETAINALLOCATIONHANDLE
471489
CUmemGenericAllocationHandle alloc_handle;
@@ -480,8 +498,9 @@ static int uct_cuda_copy_detect_vmm(const void *address,
480498
return 0;
481499
}
482500

483-
*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
484-
*cuda_device = CU_DEVICE_INVALID;
501+
*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
502+
*cuda_device = CU_DEVICE_INVALID;
503+
*is_host_located = 0;
485504

486505
status = UCT_CUDADRV_FUNC_LOG_DEBUG(
487506
cuMemGetAllocationPropertiesFromHandle(&prop, alloc_handle));
@@ -495,7 +514,8 @@ static int uct_cuda_copy_detect_vmm(const void *address,
495514
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA) ||
496515
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT)) {
497516
/* TODO: Marking as CUDA to allow cuda_ipc access vmm for now */
498-
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
517+
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
518+
*is_host_located = 1;
499519
} else
500520
#endif
501521
if (prop.location.type == CU_MEM_LOCATION_TYPE_DEVICE) {
@@ -633,7 +653,7 @@ static ucs_status_t
633653
uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
634654
const void *address, size_t length,
635655
ucs_memory_info_t *mem_info,
636-
int *is_async_managed)
656+
int *is_async_managed, int *is_host_located)
637657
{
638658
#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4
639659
CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST;
@@ -648,8 +668,10 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
648668
ucs_status_t status;
649669

650670
*is_async_managed = 0;
671+
*is_host_located = 0;
651672

652-
is_vmm = uct_cuda_copy_detect_vmm(address, &mem_info->type, &cuda_device);
673+
is_vmm = uct_cuda_copy_detect_vmm(address, &mem_info->type, &cuda_device,
674+
is_host_located);
653675
if (is_vmm) {
654676
if (mem_info->type == UCS_MEMORY_TYPE_UNKNOWN) {
655677
return UCS_ERR_INVALID_ADDR;
@@ -826,7 +848,7 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address,
826848
static uint8_t
827849
uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
828850
const ucs_memory_info_t *mem_info,
829-
int is_async_managed,
851+
int is_async_managed, int is_host_located,
830852
const uct_cuda_copy_md_dmabuf_t *dmabuf)
831853
{
832854
int close_dmabuf = 0;
@@ -836,6 +858,11 @@ uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
836858
return 0;
837859
}
838860

861+
/* Host-located CUDA VMM is registerable even if dmabuf export fails. */
862+
if (is_host_located) {
863+
return UCS_MEM_FLAG_REGISTRABLE;
864+
}
865+
839866
if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
840867
return UCS_MEM_FLAG_REGISTRABLE;
841868
}
@@ -878,8 +905,11 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
878905
.fd = UCT_DMABUF_FD_INVALID,
879906
.offset = 0
880907
};
881-
int dmabuf_queried = 0;
882-
int is_async_managed = 0;
908+
int dmabuf_queried = 0;
909+
int is_async_managed = 0;
910+
int is_host_located = 0;
911+
CUdevice cur_cuda_device = CU_DEVICE_INVALID;
912+
CUdevice avail_cuda_device = CU_DEVICE_INVALID;
883913
ucs_memory_info_t addr_mem_info;
884914
ucs_status_t status;
885915

@@ -897,7 +927,8 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
897927
if (address != NULL) {
898928
status = uct_cuda_copy_md_query_attributes(md, address, length,
899929
&addr_mem_info,
900-
&is_async_managed);
930+
&is_async_managed,
931+
&is_host_located);
901932
if (status != UCS_OK) {
902933
return status;
903934
}
@@ -921,6 +952,17 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
921952
mem_attr->alloc_length = addr_mem_info.alloc_length;
922953
}
923954

955+
/* dmabuf export and mem_flags detection require a current context on the
956+
* memory's device; reuse the current one if suitable, otherwise push it */
957+
if (addr_mem_info.sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
958+
if (uct_cuda_ctx_primary_push_avail(0, addr_mem_info.sys_dev,
959+
&cur_cuda_device, &avail_cuda_device,
960+
UCS_LOG_LEVEL_DEBUG) != UCS_OK) {
961+
/* best-effort: continue with whatever context is current */
962+
avail_cuda_device = cur_cuda_device;
963+
}
964+
}
965+
924966
if ((mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) ||
925967
(mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_OFFSET)) {
926968
dmabuf = uct_cuda_copy_md_get_dmabuf(addr_mem_info.base_address,
@@ -939,7 +981,7 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
939981

940982
if (address != NULL) {
941983
addr_mem_info.mem_flags = uct_cuda_copy_md_detect_mem_flags(
942-
md, &addr_mem_info, is_async_managed,
984+
md, &addr_mem_info, is_async_managed, is_host_located,
943985
dmabuf_queried ? &dmabuf : NULL);
944986
ucs_memtype_cache_update(addr_mem_info.base_address,
945987
addr_mem_info.alloc_length, addr_mem_info.type,
@@ -961,6 +1003,10 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
9611003
ucs_close_fd(&dmabuf.fd);
9621004
}
9631005

1006+
if (cur_cuda_device != avail_cuda_device) {
1007+
uct_cuda_ctx_primary_pop_and_release(avail_cuda_device);
1008+
}
1009+
9641010
return UCS_OK;
9651011
}
9661012

‎test/gtest/uct/cuda/test_switch_cuda_device.cc‎

Lines changed: 132 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/**
2-
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
2+
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
33
*
44
* See file LICENSE for terms.
55
*/
@@ -8,6 +8,7 @@
88
#include <uct/test_p2p_rma.h>
99

1010
extern "C" {
11+
#include <ucs/memory/memtype_cache.h>
1112
#include <ucs/sys/ptr_arith.h>
1213
#include <uct/base/uct_md.h>
1314
}
@@ -60,7 +61,8 @@ class cuda_vmm_mem_buffer {
6061
void *ptr() const;
6162

6263
protected:
63-
void init(size_t size, unsigned handle_type);
64+
void init(size_t size, unsigned handle_type,
65+
CUmemLocationType location_type = CU_MEM_LOCATION_TYPE_DEVICE);
6466

6567
private:
6668
size_t m_size = 0;
@@ -86,20 +88,23 @@ void *cuda_vmm_mem_buffer::ptr() const
8688
return (void*)m_ptr;
8789
}
8890

89-
void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
91+
void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type,
92+
CUmemLocationType location_type)
9093
{
91-
size_t granularity = 0;
92-
CUmemAllocationProp prop = {};
93-
CUmemAccessDesc access_desc = {};
94+
size_t granularity = 0;
95+
CUmemAllocationProp prop = {};
96+
CUmemAccessDesc access_desc[2] = {};
97+
unsigned num_access = 1;
98+
bool host_located = (location_type != CU_MEM_LOCATION_TYPE_DEVICE);
9499
CUdevice device;
95100
if (cuCtxGetDevice(&device) != CUDA_SUCCESS) {
96101
UCS_TEST_ABORT("failed to get the device handle for the current "
97102
"context");
98103
}
99104

100-
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
101-
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
102-
prop.location.id = device;
105+
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
106+
prop.location.type = location_type;
107+
prop.location.id = host_located ? 0 : device;
103108
if (handle_type != 0) {
104109
prop.requestedHandleTypes = (CUmemAllocationHandleType)handle_type;
105110
}
@@ -122,10 +127,16 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
122127
goto err_address_free;
123128
}
124129

125-
access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
126-
access_desc.location.id = device;
127-
access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
128-
if (cuMemSetAccess(m_ptr, m_size, &access_desc, 1) != CUDA_SUCCESS) {
130+
access_desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
131+
access_desc[0].location.id = device;
132+
access_desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
133+
if (host_located) {
134+
access_desc[1].location.type = location_type;
135+
access_desc[1].location.id = 0;
136+
access_desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
137+
num_access = 2;
138+
}
139+
if (cuMemSetAccess(m_ptr, m_size, access_desc, num_access) != CUDA_SUCCESS) {
129140
goto err_mem_unmap;
130141
}
131142

@@ -138,7 +149,7 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
138149
err_mem_release:
139150
cuMemRelease(m_alloc_handle);
140151
err:
141-
UCS_TEST_SKIP_R("failed to allocate CUDA fabric memory");
152+
UCS_TEST_SKIP_R("failed to allocate CUDA VMM memory");
142153
}
143154

144155
#if HAVE_CUDA_FABRIC
@@ -154,6 +165,16 @@ cuda_fabric_mem_buffer::cuda_fabric_mem_buffer(size_t size,
154165
}
155166
#endif
156167

168+
#if CUDA_VERSION >= 12020
169+
class cuda_host_vmm_mem_buffer : public cuda_vmm_mem_buffer {
170+
public:
171+
cuda_host_vmm_mem_buffer(size_t size, ucs_memory_type_t mem_type)
172+
{
173+
init(size, 0, CU_MEM_LOCATION_TYPE_HOST_NUMA);
174+
}
175+
};
176+
#endif
177+
157178
UCS_TEST_P(test_switch_cuda_device, detect_mem_type_cuda)
158179
{
159180
detect_mem_type<mem_buffer>(UCS_MEMORY_TYPE_CUDA);
@@ -431,6 +452,22 @@ class test_mem_alloc_device : public test_switch_cuda_device {
431452
#endif
432453
}
433454

455+
void query_registrable_no_current_context(void *address, size_t size)
456+
{
457+
uct_md_mem_attr_v2_t mem_attr = {};
458+
ucs_status_t query_status = UCS_ERR_NO_ELEM;
459+
460+
std::thread([&]() {
461+
mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
462+
UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
463+
query_status = uct_md_mem_query_v2(md(), address, size, &mem_attr);
464+
}).join();
465+
466+
ASSERT_UCS_OK(query_status);
467+
EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
468+
EXPECT_TRUE(mem_attr.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
469+
}
470+
434471
private:
435472
std::vector<ucs_sys_device_t> m_sys_dev;
436473

@@ -474,6 +511,87 @@ UCS_TEST_P(test_mem_alloc_device, same_device_cuda_fabric_implicit,
474511
test_same_device_alloc(UCS_MEMORY_TYPE_CUDA, false);
475512
}
476513

514+
UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
515+
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
516+
{
517+
ucs_memory_info_t mem_info = {};
518+
ucs_status_t lookup_status = UCS_ERR_NO_ELEM;
519+
ucs_status_t alloc_status = UCS_ERR_NO_MEMORY;
520+
CUresult ctx_status = CUDA_ERROR_UNKNOWN;
521+
CUcontext cuda_ctx = nullptr;
522+
523+
/* UCP initializes the cache before allocating rendezvous fragments. */
524+
ucs_memory_info_t init_mem_info = {};
525+
(void)ucs_memtype_cache_lookup(&init_mem_info, sizeof(init_mem_info),
526+
&init_mem_info);
527+
528+
std::thread([&]() {
529+
ctx_status = cuCtxGetCurrent(&cuda_ctx);
530+
if ((ctx_status != CUDA_SUCCESS) || (cuda_ctx != nullptr)) {
531+
return;
532+
}
533+
534+
alloc_status = allocate(UCS_MEMORY_TYPE_CUDA);
535+
if (alloc_status != UCS_OK) {
536+
return;
537+
}
538+
539+
lookup_status = ucs_memtype_cache_lookup(mem.address, mem.length,
540+
&mem_info);
541+
}).join();
542+
543+
EXPECT_EQ(CUDA_SUCCESS, ctx_status);
544+
EXPECT_EQ(nullptr, cuda_ctx);
545+
ASSERT_UCS_OK(alloc_status);
546+
547+
ucs_status_t free_status = uct_mem_free(&mem);
548+
549+
EXPECT_UCS_OK(lookup_status);
550+
if (lookup_status == UCS_OK) {
551+
EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_info.type);
552+
EXPECT_TRUE(mem_info.mem_flags & UCS_MEM_FLAG_REGISTRABLE);
553+
}
554+
EXPECT_UCS_OK(free_status);
555+
}
556+
557+
UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable,
558+
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
559+
{
560+
constexpr size_t size = 4 * UCS_MBYTE;
561+
CUdeviceptr dptr = 0;
562+
563+
ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc(&dptr, size));
564+
565+
query_registrable_no_current_context(reinterpret_cast<void*>(dptr), size);
566+
567+
EXPECT_EQ(CUDA_SUCCESS, cuMemFree(dptr));
568+
}
569+
570+
UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable,
571+
"CUDA_COPY_ASYNC_MEM_TYPE=cuda")
572+
{
573+
constexpr size_t size = 4 * UCS_MBYTE;
574+
cuda_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
575+
576+
query_registrable_no_current_context(buffer.ptr(), size);
577+
}
578+
579+
#if CUDA_VERSION >= 12020
580+
/* Host-located VMM is not dmabuf-exportable but is registrable by IB.
581+
*
582+
* TODO: REG_WHOLE_ALLOC=off is needed because of an existing issue with
583+
* base_address being set to 0, causing cache pollution with this type of memory
584+
*/
585+
UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable,
586+
"CUDA_COPY_REG_WHOLE_ALLOC=off")
587+
{
588+
constexpr size_t size = 4 * UCS_MBYTE;
589+
cuda_host_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
590+
591+
query_registrable_no_current_context(buffer.ptr(), size);
592+
}
593+
#endif
594+
477595
_UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy);
478596

479597
class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test {

0 commit comments

Comments
 (0)