Skip to content

UCT/CUDA/CUDA_IPC: Unmap memory handle at ucp_device_mem_list_release. #11730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rakhmets merged 1 commit into from
Aug 14, 2026
Merged

UCT/CUDA/CUDA_IPC: Unmap memory handle at ucp_device_mem_list_release. #11730

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 124 additions & 40 deletions src/ucp/core/ucp_device.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,18 @@
#include "ucp_ep.inl"
#include "ucp_mm.inl"

typedef struct {
uct_md_h md;
void *release_handle;
} ucp_device_mem_elem_release_handle_t;

UCS_ARRAY_DECLARE_TYPE(ucp_device_mem_elem_release_handles_t, unsigned,
ucp_device_mem_elem_release_handle_t);

typedef struct {
uct_allocated_memory_t mem;
uint32_t mem_list_length;
uct_allocated_memory_t mem;
uint32_t mem_list_length;
ucp_device_mem_elem_release_handles_t release_handles;
} ucp_device_handle_info_t;

#define ucp_device_handle_hash_key(_handle) \
Expand Down Expand Up @@ -61,9 +69,40 @@ void ucp_device_cleanup(void)
ucs_spinlock_destroy(&ucp_device_handle_hash_lock);
}

static ucs_status_t
ucp_device_mem_handle_hash_insert(const uct_allocated_memory_t *mem_handle,
uint32_t mem_list_length)
static ucs_status_t ucp_device_mem_elem_release_handles_append(
ucp_device_mem_elem_release_handles_t *release_handles, uct_md_h md,
void *release_handle)
{
ucp_device_mem_elem_release_handle_t *it;

it = ucs_array_append(release_handles, return UCS_ERR_NO_MEMORY);
it->md = md;
it->release_handle = release_handle;
return UCS_OK;
}

static void ucp_device_mem_elem_release_handles_cleanup(
ucp_device_mem_elem_release_handles_t *release_handles)
{
ucp_device_mem_elem_release_handle_t *it;

/* Check buffer explicitly: Coverity issues FORWARD_NULL error;
ucs_array_is_empty is not enough as Coverity does not correlate
non-zero length with a non-NULL buffer after create_handle paths. */
if (ucs_array_begin(release_handles) == NULL) {
return;
}

ucs_array_for_each(it, release_handles) {
uct_md_mem_elem_release(it->md, it->release_handle);
}

ucs_array_cleanup_dynamic(release_handles);
}

static ucs_status_t ucp_device_mem_handle_hash_insert(
const uct_allocated_memory_t *mem_handle, uint32_t mem_list_length,
ucp_device_mem_elem_release_handles_t *release_handles)
{
ucs_status_t status;
khiter_t iter;
Expand All @@ -72,6 +111,7 @@ ucp_device_mem_handle_hash_insert(const uct_allocated_memory_t *mem_handle,

info.mem = *mem_handle;
info.mem_list_length = mem_list_length;
info.release_handles = *release_handles;

ucs_spin_lock(&ucp_device_handle_hash_lock);
iter = kh_put(ucp_device_handle_allocs, &ucp_device_handle_hash,
Expand All @@ -85,25 +125,27 @@ ucp_device_mem_handle_hash_insert(const uct_allocated_memory_t *mem_handle,
} else {
kh_value(&ucp_device_handle_hash, iter) = info;
status = UCS_OK;
/* Ownership of the array moved to the hash */
ucs_array_init_dynamic(release_handles);
}

ucs_spin_unlock(&ucp_device_handle_hash_lock);
return status;
}

static uct_allocated_memory_t ucp_device_mem_handle_hash_remove(void *handle)
static ucp_device_handle_info_t ucp_device_mem_handle_hash_remove(void *handle)
{
khiter_t iter;
uct_allocated_memory_t mem;
ucp_device_handle_info_t info;

ucs_spin_lock(&ucp_device_handle_hash_lock);
iter = kh_get(ucp_device_handle_allocs, &ucp_device_handle_hash, handle);
ucs_assertv_always((iter != kh_end(&ucp_device_handle_hash)), "handle=%p",
handle);
mem = kh_value(&ucp_device_handle_hash, iter).mem;
info = kh_value(&ucp_device_handle_hash, iter);
kh_del(ucp_device_handle_allocs, &ucp_device_handle_hash, iter);
ucs_spin_unlock(&ucp_device_handle_hash_lock);
return mem;
return info;
}

static ucs_status_t
Expand Down Expand Up @@ -182,13 +224,15 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker,
static ucs_status_t ucp_device_local_mem_list_element_pack(
const ucp_worker_h worker, const ucp_worker_iface_t *wiface,
const ucp_device_mem_list_elem_t *element,
const ucs_memory_type_t mem_type, uct_device_mem_elem_t *mem_element)
const ucs_memory_type_t mem_type, uct_device_mem_elem_t *mem_element,
ucp_device_mem_elem_release_handles_t *release_handles)
{
ucp_tl_resource_desc_t *resource;
ucp_md_index_t md_index;
ucp_mem_h memh;
uct_mem_h uct_memh;
ucp_tl_md_t *ucp_md;
void *release_handle;
ucs_status_t status;

if (wiface == NULL) {
Expand All @@ -207,11 +251,18 @@ static ucs_status_t ucp_device_local_mem_list_element_pack(
}

status = uct_md_mem_elem_pack(ucp_md->md, uct_memh, UCT_INVALID_RKEY,
mem_element);
mem_element, &release_handle);
if (status != UCS_OK) {
ucs_error("failed to pack local mem element for memh=%p", memh);
return status;
}

status = ucp_device_mem_elem_release_handles_append(release_handles,
ucp_md->md,
release_handle);
if (status != UCS_OK) {
uct_md_mem_elem_release(ucp_md->md, release_handle);
}
return status;
}

Expand Down Expand Up @@ -240,7 +291,8 @@ static ucs_status_t ucp_device_mem_list_export_handle(

static ucs_status_t ucp_device_local_mem_list_create_handle(
const ucp_device_mem_list_params_t *params, ucs_memory_type_t mem_type,
uct_allocated_memory_t *mem, ucs_sys_device_t local_sys_dev)
ucs_sys_device_t local_sys_dev, uct_allocated_memory_t *mem,
ucp_device_mem_elem_release_handles_t *release_handles)
{
const ucp_worker_h worker = UCS_PARAM_VALUE(
UCP_DEVICE_MEM_LIST_PARAMS_FIELD, params, worker, WORKER, NULL);
Expand Down Expand Up @@ -283,7 +335,8 @@ static ucs_status_t ucp_device_local_mem_list_create_handle(
status = ucp_device_local_mem_list_element_pack(worker, wiface,
ucp_element,
mem_type,
tl_element);
tl_element,
release_handles);
if (status != UCS_OK) {
ucs_error("failed to pack local mem list element for "
"element=%zu",
Expand Down Expand Up @@ -377,6 +430,8 @@ ucs_status_t
ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params,
ucp_device_local_mem_list_h *mem_list_h)
{
ucp_device_mem_elem_release_handles_t release_handles =
UCS_ARRAY_DYNAMIC_INITIALIZER;

@ofirfarjun7 ofirfarjun7 Aug 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Next PR: We know the mem list length in advance, can we move to fixed size array?

const ucp_worker_h worker = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_PARAMS_FIELD,
params, worker, WORKER, NULL);
ucs_memory_type_t export_mem_type;
Expand Down Expand Up @@ -405,21 +460,27 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params,
}

status = ucp_device_local_mem_list_create_handle(params, export_mem_type,
&mem, local_sys_dev);
local_sys_dev, &mem,
&release_handles);
if (status != UCS_OK) {
ucs_error("failed to create local mem list handle: %s",
ucs_status_string(status));
return status;
goto err;
}

/* Track memory allocator for later release */
status = ucp_device_mem_handle_hash_insert(&mem, params->num_elements);
status = ucp_device_mem_handle_hash_insert(&mem, params->num_elements,
&release_handles);
if (status != UCS_OK) {
uct_mem_free(&mem);
} else {
*mem_list_h = mem.address;
goto err;
}

*mem_list_h = mem.address;
return UCS_OK;

err:
ucp_device_mem_elem_release_handles_cleanup(&release_handles);
return status;
}

Expand Down Expand Up @@ -458,15 +519,18 @@ ucp_device_ep_check_lanes(const ucp_ep_h ep, ucp_tl_bitmap_t *tl_bitmap)

static ucs_status_t ucp_device_remote_mem_list_element_pack(
const ucp_device_mem_list_elem_t *element, ucp_rsc_index_t tl_id,
uct_device_remote_tl_elem_t *mem_element)
uct_device_remote_tl_elem_t *mem_element,
ucp_device_mem_elem_release_handles_t *release_handles)
{
const ucp_ep_h ep = element->ep;
const ucp_rkey_h rkey = element->rkey;
ucp_ep_config_t *ep_config = ucp_ep_config(ep);
uct_md_h md;
uint8_t rkey_index;
uct_rkey_t uct_rkey;
uct_ep_h uct_ep;
uct_device_ep_h device_ep;
void *release_handle;
ucs_status_t status;
ucp_lane_index_t lane;

Expand All @@ -488,11 +552,19 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack(
uct_rkey = ucp_rkey_get_tl_rkey(rkey, rkey_index);
ucs_assert(uct_rkey != UCT_INVALID_RKEY);

md = ucp_ep_md(ep, lane);
mem_element->ep = device_ep;
status = uct_md_mem_elem_pack(ucp_ep_md(ep, lane), NULL, uct_rkey,
&mem_element->uct);
status = uct_md_mem_elem_pack(md, NULL, uct_rkey, &mem_element->uct,
&release_handle);
if (status != UCS_OK) {
ucs_error("failed to pack uct memory element for lane=%u", lane);
return status;
}

status = ucp_device_mem_elem_release_handles_append(release_handles, md,
release_handle);
if (status != UCS_OK) {
uct_md_mem_elem_release(md, release_handle);
}

return status;
Expand Down Expand Up @@ -532,10 +604,11 @@ static ucp_ep_h ucp_device_remote_mem_list_get_first_ep(
return NULL;
}

static ucs_status_t
ucp_device_remote_mem_list_fill(const ucp_device_mem_list_elem_t *ucp_element,
ucp_tl_bitmap_t *tl_bitmap, size_t num_lanes,
uct_device_remote_mem_elem_t *uct_element)
static ucs_status_t ucp_device_remote_mem_list_fill(
const ucp_device_mem_list_elem_t *ucp_element,
const ucp_tl_bitmap_t *tl_bitmap, size_t num_lanes,
uct_device_remote_mem_elem_t *uct_element,
ucp_device_mem_elem_release_handles_t *release_handles)
{
uct_device_remote_tl_elem_t *tl_element;
ucp_rsc_index_t tl_id;
Expand All @@ -547,7 +620,8 @@ ucp_device_remote_mem_list_fill(const ucp_device_mem_list_elem_t *ucp_element,
for (i = 0; i < num_lanes;) {
UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) {
status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id,
tl_element);
tl_element,
release_handles);
if (status != UCS_OK) {
return status;
}
Expand All @@ -561,9 +635,10 @@ ucp_device_remote_mem_list_fill(const ucp_device_mem_list_elem_t *ucp_element,
}

static ucs_status_t ucp_device_remote_mem_list_create_handle(
const ucp_device_mem_list_params_t *params, const ucp_ep_h ep,
const ucs_memory_type_t mem_type, uct_allocated_memory_t *mem,
const ucs_sys_device_t local_sys_dev)
const ucp_device_mem_list_params_t *params, ucp_ep_h ep,
ucs_memory_type_t mem_type, ucs_sys_device_t local_sys_dev,
uct_allocated_memory_t *mem,
ucp_device_mem_elem_release_handles_t *release_handles)
{
size_t uct_elem_size;
size_t handle_size = 0;
Expand Down Expand Up @@ -622,7 +697,8 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle(

status = ucp_device_remote_mem_list_fill(ucp_element,
&tl_bitmap[tl_type],
num_lanes, uct_element);
num_lanes, uct_element,
release_handles);
if (status != UCS_OK) {
goto out;
}
Expand Down Expand Up @@ -698,7 +774,9 @@ ucs_status_t
ucp_device_remote_mem_list_create(const ucp_device_mem_list_params_t *params,
ucp_device_remote_mem_list_h *mem_list_h)
{
const ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params);
ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params);
ucp_device_mem_elem_release_handles_t release_handles =
UCS_ARRAY_DYNAMIC_INITIALIZER;
ucs_memory_type_t export_mem_type;
ucs_sys_device_t sys_dev;
uct_allocated_memory_t mem;
Expand All @@ -721,8 +799,8 @@ ucp_device_remote_mem_list_create(const ucp_device_mem_list_params_t *params,
}

status = ucp_device_remote_mem_list_create_handle(params, ep,
export_mem_type,
&mem, sys_dev);
export_mem_type, sys_dev,
&mem, &release_handles);
if (status != UCS_OK) {
/*
* Do not log error for UCS_ERR_NOT_CONNECTED because it is expected
Expand All @@ -732,17 +810,22 @@ ucp_device_remote_mem_list_create(const ucp_device_mem_list_params_t *params,
if (status != UCS_ERR_NOT_CONNECTED) {
ucs_error("failed to create handle: %s", ucs_status_string(status));
}
return status;
goto err;
}

/* Track memory allocator for later release */
status = ucp_device_mem_handle_hash_insert(&mem, params->num_elements);
status = ucp_device_mem_handle_hash_insert(&mem, params->num_elements,
&release_handles);
if (status != UCS_OK) {
uct_mem_free(&mem);
} else {
*mem_list_h = mem.address;
goto err;
}

*mem_list_h = mem.address;
return UCS_OK;

err:
ucp_device_mem_elem_release_handles_cleanup(&release_handles);
return status;
}

Expand All @@ -766,10 +849,11 @@ uint32_t ucp_device_get_mem_list_length(const void *mem_list_h)

void ucp_device_mem_list_release(void *mem_list_h)
{
uct_allocated_memory_t mem;
ucp_device_handle_info_t info;

mem = ucp_device_mem_handle_hash_remove(mem_list_h);
uct_mem_free(&mem);
info = ucp_device_mem_handle_hash_remove(mem_list_h);
ucp_device_mem_elem_release_handles_cleanup(&info.release_handles);
uct_mem_free(&info.mem);
}

static ucs_memory_type_t
Expand Down
23 changes: 18 additions & 5 deletions src/uct/api/v2/uct_v2.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1535,15 +1535,28 @@ ucs_status_t uct_rkey_unpack_v2(uct_component_h component,
* @ingroup UCT_MD
* @brief Pack a memh and rkey into a device memory element structure.
*
* @param [in] md Memory domain.
* @param [in] memh Memory handle to pack (can be NULL).
* @param [in] rkey Remote key to pack (can be UCT_INVALID_RKEY).
* @param [out] mem_elem Filled with the packed memh and rkey.
* @param [in] md Memory domain.
* @param [in] memh Memory handle to pack (can be NULL).
* @param [in] rkey Remote key to pack (can be UCT_INVALID_RKEY).
* @param [out] mem_elem Filled with the packed memh and rkey.
* @param [out] release_handle_p Handle for releasing resources allocated
* during packing.
*
* @return UCS_OK on success or error code in case of failure.
*/
ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
uct_device_mem_elem_t *mem_elem);
uct_device_mem_elem_t *mem_elem,
void **release_handle_p);


/**
* @ingroup UCT_MD
* @brief Release resources allocated by @ref uct_md_mem_elem_pack.
*
* @param [in] md Memory domain.
* @param [in] release_handle Handle for releasing resources.
*/
void uct_md_mem_elem_release(uct_md_h md, void *release_handle);


/**
Expand Down
Loading