Skip to content

Commit 75a31bc

Browse files
tomerg-nvidiatomerg-nvidia
authored
UCT/CUDA_COPY: Fix incorrect mem_flags detection (#11635)
1 parent 1bb34d6 commit 75a31bc

2 files changed

Lines changed: 128 additions & 25 deletions

File tree

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

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -482,7 +482,8 @@ uct_cuda_copy_mem_release_fabric(uct_cuda_copy_alloc_handle_t *alloc_handle)
482482

483483
static int uct_cuda_copy_detect_vmm(const void *address,
484484
ucs_memory_type_t *vmm_mem_type,
485-
CUdevice *cuda_device)
485+
CUdevice *cuda_device,
486+
int *is_host_located)
486487
{
487488
#ifdef HAVE_CUMEMRETAINALLOCATIONHANDLE
488489
CUmemGenericAllocationHandle alloc_handle;
@@ -497,8 +498,9 @@ static int uct_cuda_copy_detect_vmm(const void *address,
497498
return 0;
498499
}
499500

500-
*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
501-
*cuda_device = CU_DEVICE_INVALID;
501+
*vmm_mem_type = UCS_MEMORY_TYPE_UNKNOWN;
502+
*cuda_device = CU_DEVICE_INVALID;
503+
*is_host_located = 0;
502504

503505
status = UCT_CUDADRV_FUNC_LOG_DEBUG(
504506
cuMemGetAllocationPropertiesFromHandle(&prop, alloc_handle));
@@ -512,7 +514,8 @@ static int uct_cuda_copy_detect_vmm(const void *address,
512514
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA) ||
513515
(prop.location.type == CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT)) {
514516
/* TODO: Marking as CUDA to allow cuda_ipc access vmm for now */
515-
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
517+
*vmm_mem_type = UCS_MEMORY_TYPE_CUDA;
518+
*is_host_located = 1;
516519
} else
517520
#endif
518521
if (prop.location.type == CU_MEM_LOCATION_TYPE_DEVICE) {
@@ -650,7 +653,7 @@ static ucs_status_t
650653
uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
651654
const void *address, size_t length,
652655
ucs_memory_info_t *mem_info,
653-
int *is_async_managed)
656+
int *is_async_managed, int *is_host_located)
654657
{
655658
#define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4
656659
CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST;
@@ -665,8 +668,10 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md,
665668
ucs_status_t status;
666669

667670
*is_async_managed = 0;
671+
*is_host_located = 0;
668672

669-
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);
670675
if (is_vmm) {
671676
if (mem_info->type == UCS_MEMORY_TYPE_UNKNOWN) {
672677
return UCS_ERR_INVALID_ADDR;
@@ -843,7 +848,7 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address,
843848
static uint8_t
844849
uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
845850
const ucs_memory_info_t *mem_info,
846-
int is_async_managed,
851+
int is_async_managed, int is_host_located,
847852
const uct_cuda_copy_md_dmabuf_t *dmabuf)
848853
{
849854
int close_dmabuf = 0;
@@ -853,6 +858,11 @@ uct_cuda_copy_md_detect_mem_flags(uct_cuda_copy_md_t *md,
853858
return 0;
854859
}
855860

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+
856866
if (mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
857867
return UCS_MEM_FLAG_REGISTRABLE;
858868
}
@@ -895,8 +905,11 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
895905
.fd = UCT_DMABUF_FD_INVALID,
896906
.offset = 0
897907
};
898-
int dmabuf_queried = 0;
899-
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;
900913
ucs_memory_info_t cached_mem_info;
901914
ucs_memory_info_t addr_mem_info;
902915
ucs_status_t cache_status;
@@ -918,7 +931,8 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
918931
&cached_mem_info);
919932
status = uct_cuda_copy_md_query_attributes(md, address, length,
920933
&addr_mem_info,
921-
&is_async_managed);
934+
&is_async_managed,
935+
&is_host_located);
922936
if (status != UCS_OK) {
923937
return status;
924938
}
@@ -949,6 +963,17 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
949963
mem_attr->alloc_length = addr_mem_info.alloc_length;
950964
}
951965

966+
/* dmabuf export and mem_flags detection require a current context on the
967+
* memory's device; reuse the current one if suitable, otherwise push it */
968+
if (addr_mem_info.sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
969+
if (uct_cuda_ctx_primary_push_avail(0, addr_mem_info.sys_dev,
970+
&cur_cuda_device, &avail_cuda_device,
971+
UCS_LOG_LEVEL_DEBUG) != UCS_OK) {
972+
/* best-effort: continue with whatever context is current */
973+
avail_cuda_device = cur_cuda_device;
974+
}
975+
}
976+
952977
if ((mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) ||
953978
(mem_attr->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_OFFSET)) {
954979
dmabuf = uct_cuda_copy_md_get_dmabuf(addr_mem_info.base_address,
@@ -967,7 +992,7 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
967992

968993
if (address != NULL) {
969994
addr_mem_info.mem_flags = uct_cuda_copy_md_detect_mem_flags(
970-
md, &addr_mem_info, is_async_managed,
995+
md, &addr_mem_info, is_async_managed, is_host_located,
971996
dmabuf_queried ? &dmabuf : NULL);
972997
ucs_memtype_cache_update(addr_mem_info.base_address,
973998
addr_mem_info.alloc_length, addr_mem_info.type,
@@ -989,6 +1014,10 @@ ucs_status_t uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address,
9891014
ucs_close_fd(&dmabuf.fd);
9901015
}
9911016

1017+
if (cur_cuda_device != avail_cuda_device) {
1018+
uct_cuda_ctx_primary_pop_and_release(avail_cuda_device);
1019+
}
1020+
9921021
return UCS_OK;
9931022
}
9941023

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

Lines changed: 88 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
*/
@@ -61,7 +61,8 @@ class cuda_vmm_mem_buffer {
6161
void *ptr() const;
6262

6363
protected:
64-
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);
6566

6667
private:
6768
size_t m_size = 0;
@@ -87,20 +88,23 @@ void *cuda_vmm_mem_buffer::ptr() const
8788
return (void*)m_ptr;
8889
}
8990

90-
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)
9193
{
92-
size_t granularity = 0;
93-
CUmemAllocationProp prop = {};
94-
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);
9599
CUdevice device;
96100
if (cuCtxGetDevice(&device) != CUDA_SUCCESS) {
97101
UCS_TEST_ABORT("failed to get the device handle for the current "
98102
"context");
99103
}
100104

101-
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
102-
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
103-
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;
104108
if (handle_type != 0) {
105109
prop.requestedHandleTypes = (CUmemAllocationHandleType)handle_type;
106110
}
@@ -123,10 +127,16 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
123127
goto err_address_free;
124128
}
125129

126-
access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
127-
access_desc.location.id = device;
128-
access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
129-
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) {
130140
goto err_mem_unmap;
131141
}
132142

@@ -139,7 +149,7 @@ void cuda_vmm_mem_buffer::init(size_t size, unsigned handle_type)
139149
err_mem_release:
140150
cuMemRelease(m_alloc_handle);
141151
err:
142-
UCS_TEST_SKIP_R("failed to allocate CUDA fabric memory");
152+
UCS_TEST_SKIP_R("failed to allocate CUDA VMM memory");
143153
}
144154

145155
#if HAVE_CUDA_FABRIC
@@ -155,6 +165,16 @@ cuda_fabric_mem_buffer::cuda_fabric_mem_buffer(size_t size,
155165
}
156166
#endif
157167

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+
158178
UCS_TEST_P(test_switch_cuda_device, detect_mem_type_cuda)
159179
{
160180
detect_mem_type<mem_buffer>(UCS_MEMORY_TYPE_CUDA);
@@ -432,6 +452,22 @@ class test_mem_alloc_device : public test_switch_cuda_device {
432452
#endif
433453
}
434454

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+
435471
private:
436472
std::vector<ucs_sys_device_t> m_sys_dev;
437473

@@ -518,6 +554,44 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_cuda_registrable,
518554
EXPECT_UCS_OK(free_status);
519555
}
520556

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+
521595
_UCT_MD_INSTANTIATE_TEST_CASE(test_mem_alloc_device, cuda_cpy);
522596

523597
class test_p2p_no_current_cuda_ctx : public uct_p2p_rma_test {

0 commit comments

Comments
 (0)