UCT/CUDA_COPY: Fix no context mem flags by tomerg-nvidia · Pull Request #11635 · openucx/ucx

tomerg-nvidia · 2026-07-08T16:18:23Z

What?

Push a cuda context before flag detection if the flag detection is going to export dmabuf fd.

Why?

If there's currently no context and the memory flags detection is called the dmabuf export can fail, making it non-registerable

How?

Try to reuse context from earlier uct_cuda_copy_md_query_attributes, if there's none then do uct_cuda_ctx_primary_push

tvegas1 · 2026-07-09T11:37:30Z

        mem_attr->alloc_length = addr_mem_info.alloc_length;
    }

+    /* mem_flags detection requires a current context; push one if needed */


maybe: DMABUF fd creation/export, also used to detect mem_flags, needs a cuda context

removed it, now the push is much closer to the use

tvegas1 · 2026-07-09T12:19:56Z

    }

+    /* mem_flags detection requires a current context; push one if needed */
+    if ((cuda_device != CU_DEVICE_INVALID) &&


since we add a new function for push/pop, and the need for context really is dmabuf fd export function, and syncmemops already push/pop a CUDA ctx, would it make sense to implement the pushes closer to where they are needed, and finally pop before mem_query return, like below. also this if condition for instance misses the dmabuf supported check also. this could also avoid repetitive push/pop.

diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index 8735bc067..36ce90ec0 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -88,6 +88,14 @@ static ucs_config_field_t uct_cuda_copy_md_config_table[] = { static struct {} uct_cuda_dummy_memh; +struct uct_cuda_copy_md_query_ctx { + CUcontext cuda_ctx; + CUdevice cuda_device; + CUcontext pushed_cuda_ctx; + CUdevice pushed_cuda_device; +}; + + int uct_cuda_copy_md_is_dmabuf_supported() { int dmabuf_supported = 0; @@ -568,12 +576,57 @@ err: return 1; /* return 1 byte to avoid division by zero */ } +static ucs_status_t +uct_cuda_copy_md_ensure_ctx_pushed(uct_cuda_copy_md_query_ctx_t *query_ctx) +{ + ucs_status_t status; + + if ((query_ctx == NULL) || + ((query_ctx->cuda_ctx == NULL) && + (query_ctx->cuda_device == CU_DEVICE_INVALID))) { + return UCS_OK; + } + + if ((query_ctx->pushed_cuda_ctx != NULL) || + (query_ctx->pushed_cuda_device != CU_DEVICE_INVALID)) { + return UCS_OK; + } + + if (query_ctx->cuda_ctx == NULL) { + status = uct_cuda_ctx_primary_push(query_ctx->cuda_device, 0, + UCS_LOG_LEVEL_ERROR); + if (status == UCS_OK) { + query_ctx->pushed_cuda_device = query_ctx->cuda_device; + } + } else { + status = UCT_CUDADRV_FUNC_LOG_ERR( + cuCtxPushCurrent(query_ctx->cuda_ctx)); + if (status == UCS_OK) { + query_ctx->pushed_cuda_ctx = query_ctx->cuda_ctx; + } + } + + return status; +} + +static void +uct_cuda_copy_md_pop_ctx(uct_cuda_copy_md_query_ctx_t *query_ctx) +{ + if (query_ctx == NULL) { + return; + } + + uct_cuda_ctx_pop_and_release(query_ctx->pushed_cuda_device, + query_ctx->pushed_cuda_ctx); + query_ctx->pushed_cuda_device = CU_DEVICE_INVALID; + query_ctx->pushed_cuda_ctx = NULL; +} + static void uct_cuda_copy_md_sync_memops_get_address_range( const uct_cuda_copy_md_t *md, CUdeviceptr address, size_t length, CUcontext cuda_ctx, CUdevice cuda_device, int is_vmm, - ucs_memory_info_t *mem_info) + ucs_memory_info_t *mem_info, uct_cuda_copy_md_query_ctx_t *query_ctx) { - CUcontext tmp_ctx; CUdeviceptr base_address; size_t alloc_length; size_t total_bytes; @@ -582,28 +635,26 @@ static void uct_cuda_copy_md_sync_memops_get_address_range( mem_info->base_address = (void*)address; mem_info->alloc_length = length; - if (cuda_ctx == NULL) { - status = uct_cuda_ctx_primary_push(cuda_device, 0, UCS_LOG_LEVEL_ERROR); - } else { - status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(cuda_ctx)); - } + query_ctx->cuda_ctx = cuda_ctx; + query_ctx->cuda_device = cuda_device; + + status = uct_cuda_copy_md_ensure_ctx_pushed(query_ctx); if (status != UCS_OK) { return; } - /* Wrapped the method by push/pop CUDA context since it sets - CU_CTX_SYNC_MEMOPS flag for the current context */ + /* The current context is required because sync memops is set on it */ uct_cuda_copy_sync_memops(address, is_vmm); if (md->config.alloc_whole_reg == UCS_CONFIG_OFF) { /* Extending the registration range is disable by configuration */ - goto out_ctx_pop; + goto out; } if (UCT_CUDADRV_FUNC_LOG_DEBUG( cuMemGetAddressRange(&base_address, &alloc_length, address)) != UCS_OK) { - goto out_ctx_pop; + goto out; } ucs_trace("query address 0x%llx: 0x%llx..0x%llx length %zu", address, @@ -612,7 +663,7 @@ static void uct_cuda_copy_md_sync_memops_get_address_range( if (md->config.alloc_whole_reg == UCS_CONFIG_AUTO) { total_bytes = uct_cuda_copy_md_get_total_device_mem(cuda_device); if (alloc_length > (total_bytes * md->config.max_reg_ratio)) { - goto out_ctx_pop; + goto out; } } else { ucs_assert(md->config.alloc_whole_reg == UCS_CONFIG_ON); @@ -621,17 +672,15 @@ static void uct_cuda_copy_md_sync_memops_get_address_range( mem_info->base_address = (void*)base_address; mem_info->alloc_length = alloc_length; -out_ctx_pop: - UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(&tmp_ctx)); - if (cuda_ctx == NULL) { - UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device)); - } +out: + return; } static ucs_status_t uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, const void *address, size_t length, - ucs_memory_info_t *mem_info) + ucs_memory_info_t *mem_info, + uct_cuda_copy_md_query_ctx_t *query_ctx) { #define UCT_CUDA_MEM_QUERY_NUM_ATTRS 4 CUmemorytype cuda_mem_type = CU_MEMORYTYPE_HOST; @@ -647,6 +696,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, is_vmm = uct_cuda_copy_detect_vmm(address, &mem_info->type, &cuda_device); if (is_vmm) { + query_ctx->cuda_device = cuda_device; if (mem_info->type == UCS_MEMORY_TYPE_UNKNOWN) { return UCS_ERR_INVALID_ADDR; } @@ -674,6 +724,9 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, return UCS_ERR_INVALID_ADDR; } + query_ctx->cuda_ctx = cuda_mem_ctx; + query_ctx->cuda_device = cuda_device; + if (is_managed) { /* cuMemGetAddress range does not support managed memory so use * provided address and length as base address and alloc length @@ -730,7 +783,7 @@ uct_cuda_copy_md_query_attributes(const uct_cuda_copy_md_t *md, uct_cuda_copy_md_sync_memops_get_address_range(md, (CUdeviceptr)address, length, cuda_mem_ctx, cuda_device, is_vmm, - mem_info); + mem_info, query_ctx); return UCS_OK; out_default_range: @@ -739,13 +792,15 @@ out_default_range: return UCS_OK; } -static int uct_cuda_copy_md_get_dmabuf_fd(uintptr_t address, size_t length, - ucs_sys_device_t sys_dev) +static int uct_cuda_copy_md_get_dmabuf_fd( + uintptr_t address, size_t length, ucs_sys_device_t sys_dev, + uct_cuda_copy_md_query_ctx_t *query_ctx) { #if CUDA_VERSION >= 11070 unsigned long long flags = 0; PFN_cuMemGetHandleForAddressRange_v11070 get_handle_func; CUresult cu_err; + ucs_status_t status; int fd; /* Get fxn ptr for cuMemGetHandleForAddressRange in case installed libcuda @@ -785,6 +840,11 @@ static int uct_cuda_copy_md_get_dmabuf_fd(uintptr_t address, size_t length, } #endif + status = uct_cuda_copy_md_ensure_ctx_pushed(query_ctx); + if (status != UCS_OK) { + return UCT_DMABUF_FD_INVALID; + } + cu_err = get_handle_func((void*)&fd, address, length, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, flags); if (cu_err == CUDA_SUCCESS) { @@ -801,9 +861,9 @@ static int uct_cuda_copy_md_get_dmabuf_fd(uintptr_t address, size_t length, return UCT_DMABUF_FD_INVALID; } -uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address, - size_t length, - ucs_sys_device_t sys_dev) +uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf( + const void *address, size_t length, ucs_sys_device_t sys_dev, + uct_cuda_copy_md_query_ctx_t *query_ctx) { uct_cuda_copy_md_dmabuf_t dmabuf; uintptr_t base_address, aligned_start, aligned_end; @@ -813,7 +873,7 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address, aligned_end = ucs_align_up_pow2(base_address + length, ucs_get_page_size()); dmabuf.fd = uct_cuda_copy_md_get_dmabuf_fd(aligned_start, aligned_end - aligned_start, - sys_dev); + sys_dev, query_ctx); dmabuf.offset = base_address - aligned_start; return dmabuf; } @@ -830,7 +890,13 @@ uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, size_t length, }; uct_cuda_copy_md_t *md = ucs_derived_of(tl_md, uct_cuda_copy_md_t); ucs_memory_info_t addr_mem_info; - ucs_status_t status; + uct_cuda_copy_md_query_ctx_t query_ctx = { + .cuda_ctx = NULL, + .cuda_device = CU_DEVICE_INVALID, + .pushed_cuda_ctx = NULL, + .pushed_cuda_device = CU_DEVICE_INVALID + }; + ucs_status_t status = UCS_OK; uct_cuda_copy_md_dmabuf_t dmabuf; if (!(mem_attr->field_mask & @@ -839,14 +905,15 @@ uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, size_t length, UCT_MD_MEM_ATTR_FIELD_ALLOC_LENGTH | UCT_MD_MEM_ATTR_FIELD_DMABUF_FD | UCT_MD_MEM_ATTR_FIELD_DMABUF_OFFSET))) { - return UCS_OK; + goto out; } if (address != NULL) { status = uct_cuda_copy_md_query_attributes(md, address, length, - &addr_mem_info); + &addr_mem_info, + &query_ctx); if (status != UCS_OK) { - return status; + goto out; } ucs_memtype_cache_update(addr_mem_info.base_address, @@ -874,20 +941,22 @@ uct_cuda_copy_md_mem_query(uct_md_h tl_md, const void *address, size_t length, if ((mem_attr->field_mask & UCT_MD_MEM_ATTR_FIELD_DMABUF_FD) || (mem_attr->field_mask & UCT_MD_MEM_ATTR_FIELD_DMABUF_OFFSET)) { - dmabuf = uct_cuda_copy_md_get_dmabuf(addr_mem_info.base_address, - addr_mem_info.alloc_length, - addr_mem_info.sys_dev); + dmabuf = uct_cuda_copy_md_get_dmabuf( + addr_mem_info.base_address, addr_mem_info.alloc_length, + addr_mem_info.sys_dev, &query_ctx); if (mem_attr->field_mask & UCT_MD_MEM_ATTR_FIELD_DMABUF_FD) { mem_attr->dmabuf_fd = dmabuf.fd; } if (mem_attr->field_mask & UCT_MD_MEM_ATTR_FIELD_DMABUF_OFFSET) { mem_attr->dmabuf_offset = - dmabuf.offset + UCS_PTR_BYTE_DIFF(addr_mem_info.base_address, - address); + dmabuf.offset + + UCS_PTR_BYTE_DIFF(addr_mem_info.base_address, address); } } - return UCS_OK; +out: + uct_cuda_copy_md_pop_ctx(&query_ctx); + return status; } UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_copy_md_detect_memory_type, diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.h b/src/uct/cuda/cuda_copy/cuda_copy_md.h index f081fff02..8fb31b1c7 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.h +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.h @@ -72,6 +72,9 @@ typedef struct { } uct_cuda_copy_md_dmabuf_t; +typedef struct uct_cuda_copy_md_query_ctx uct_cuda_copy_md_query_ctx_t; + + ucs_status_t uct_cuda_copy_md_detect_memory_type(uct_md_h md, const void *address, size_t length, @@ -98,10 +101,13 @@ int uct_cuda_copy_md_is_dmabuf_supported(); file descriptor can be used by a Direct NIC. If sys_dev is UCS_SYS_DEVICE_ID_UNKNOWN, the file descriptor can be used by any device. + * @param query_ctx [inout] Optional context state to use for lazy CUDA context + * push. * @return The dmabuf file descriptor and offset */ -uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address, - size_t length, - ucs_sys_device_t sys_dev); +uct_cuda_copy_md_dmabuf_t +uct_cuda_copy_md_get_dmabuf(const void *address, size_t length, + ucs_sys_device_t sys_dev, + uct_cuda_copy_md_query_ctx_t *query_ctx); #endif

Yes, that's a good idea, I'll make this change

rakhmets · 2026-07-09T11:53:11Z

+    CUcontext tmp_ctx;
+
+    UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(&tmp_ctx));


Suggested change

CUcontext tmp_ctx;

UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(&tmp_ctx));

(void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));

Not relevant anymore

rakhmets · 2026-07-09T11:55:28Z

+
+    UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(&tmp_ctx));
+    if (cuda_ctx == NULL) {
+        UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));


Suggested change

UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));

(void)UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));

More often than not, it is explicitly indicated that the return value is ignodred.

not relevant anymore

rakhmets · 2026-07-09T12:18:17Z

+    *cuda_ctx_p       = NULL;
+    *cuda_device_p    = CU_DEVICE_INVALID;


I think it's better to set values only in case of returning UCS_OK.

rakhmets · 2026-07-09T12:24:28Z

+        ctx_pushed = (uct_cuda_copy_md_push_ctx(cuda_ctx, cuda_device) ==
+                      UCS_OK);


Shouldn't we return from the function if the method returns an error?

Not relevant anymore. Also, for the dmabuf export the context push is best-effort to not change the current behavior which uses current context

rakhmets · 2026-07-09T12:27:13Z

+        ctx_status = cuCtxGetCurrent(&cuda_ctx);
+        if ((ctx_status != CUDA_SUCCESS) || (cuda_ctx != nullptr)) {
+            return;
+        }


Does the test really need to check this?

rakhmets · 2026-07-09T12:27:26Z

+UCS_TEST_P(test_mem_alloc_device, no_current_context_user_mem_registrable,
+           "CUDA_COPY_ASYNC_MEM_TYPE=cuda")
+{
+    const size_t size             = 4 * UCS_MBYTE;


rakhmets · 2026-07-09T12:27:39Z

+
+        mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
+                              UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS;
+        query_status = uct_md_mem_query_v2(md(), (void*)dptr, size, &mem_attr);


static_cast

static_cast can't do it because it's an integer to pointer conversion, used reinterpret_cast

And fix CR comments

tomerg-nvidia changed the title ~~Fix no context mem flags~~ UCT/CUDA_COPY: Fix no context mem flags Jul 8, 2026

tvegas1 self-requested a review July 8, 2026 16:57

tomerg-nvidia added 3 commits July 8, 2026 19:05

TEST/GTEST/UCT: Test cuda_copy mem_query without context

a5d2ba6

UCT/CUDA_COPY: Push context before detecting mem_flags

4614172

TEST/GTEST: format fix

65d1073

tomerg-nvidia force-pushed the fix_no_context_mem_flags branch from f31ea76 to 65d1073 Compare July 8, 2026 19:09

tomerg-nvidia requested review from brminich, gleon99, rakhmets and yosefe July 9, 2026 06:35

tvegas1 reviewed Jul 9, 2026

View reviewed changes

rakhmets reviewed Jul 9, 2026

View reviewed changes

UCT/CUDA_COPY: Track pushed context

be6bdd2

And fix CR comments

tomerg-nvidia force-pushed the fix_no_context_mem_flags branch from f5d17e5 to be6bdd2 Compare July 9, 2026 14:26

UCT/CUDA_COPY: Use existing context if push failed for dmabuf export

ef7c5ac

tomerg-nvidia requested review from rakhmets and tvegas1 July 9, 2026 14:49

		CUcontext tmp_ctx;

		UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(&tmp_ctx));

	UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));
	(void)UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));

		ctx_pushed = (uct_cuda_copy_md_push_ctx(cuda_ctx, cuda_device) ==
		UCS_OK);

Uh oh!

Conversation

tomerg-nvidia commented Jul 8, 2026

What?

Why?

How?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants