Skip to content

UCP/DEVICE: Add multi-lane support #11268

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
ofirfarjun7 merged 14 commits into from
May 6, 2026
Merged

UCP/DEVICE: Add multi-lane support #11268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
fa7e9fb
UCP/DEVICE: Add multi-lane support
Artemy-Mellanox Mar 9, 2026
60e5be9
UCP/DEVICE: Add multi-lane support - 2
Artemy-Mellanox Mar 17, 2026
82faeb9
UCP/DEVICE: Add multi-lane support - 3
Artemy-Mellanox Mar 23, 2026
1f3445e
UCP/DEVICE: Add multi-lane support - 4
Artemy-Mellanox Mar 24, 2026
71d7b9c
UCP/DEVICE: Add multi-lane support - 5
Artemy-Mellanox Mar 25, 2026
c863ad2
UCP/DEVICE: Add multi-lane support - 6
Artemy-Mellanox Apr 15, 2026
e2b4ec0
UCP/DEVICE: Add multi-lane support - 7
Artemy-Mellanox Apr 24, 2026
ed5762b
UCP/DEVICE: Add multi-lane support - 8
Artemy-Mellanox Apr 24, 2026
3e4eeea
Merge remote-tracking branch 'origin/master' into multilane
Artemy-Mellanox Apr 25, 2026
7986ea3
UCP/DEVICE: Add multi-lane support - 9
Artemy-Mellanox Apr 25, 2026
72d8beb
UCP/DEVICE: Add multi-lane support - 10
Artemy-Mellanox Apr 25, 2026
f0228c0
UCP/DEVICE: Add multi-lane support - 11
Artemy-Mellanox Apr 27, 2026
508fa22
UCP/DEVICE: Add multi-lane support - 12
Artemy-Mellanox Apr 28, 2026
eb078cc
Merge remote-tracking branch 'origin/master' into multilane
Artemy-Mellanox May 5, 2026
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
97 changes: 54 additions & 43 deletions src/ucp/api/device/ucp_device_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
*/
Expand Down Expand Up @@ -105,45 +105,52 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep,
})


#define UCP_DEVICE_GET_LANE(_handle, _channel_id, _lane, _uct_channel_id) \
_lane = _channel_id % _handle->num_lanes; \
_uct_channel_id = _channel_id / _handle->num_lanes;
Comment on lines +108 to +110

@roiedanino roiedanino Apr 27, 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.

ucs_assertv(_handle != NULL, "dest memlist can't be null")
before deref, clearer failure for users

@Artemy-Mellanox Artemy-Mellanox Apr 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

in other APIs we not validate objects for NULL



#define UCP_DEVICE_GET_ELEM(_handle, _index) \
static_cast<ucs_typeof(_handle->mem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \
Comment on lines +113 to +114

@yosefe yosefe May 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

can it be inline function?

_handle->mem_elements, \
(sizeof(_handle->mem_elements[0]) + \
(sizeof(_handle->mem_elements[0].tl[0]) * _handle->num_lanes)) * \

@yosefe yosefe May 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

how much it affects performance?

_index))


UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote(
const ucp_device_remote_mem_list_h dst_mem_list_h,
unsigned dst_mem_list_index, uint64_t &remote_address,
unsigned dst_mem_list_index, uint64_t &remote_address, unsigned lane,
ucp_device_request_t *req, uct_device_ep_t *&device_ep,
const uct_device_mem_element_t *&uct_elem,
uct_device_completion_t *&comp)
const uct_device_mem_elem_t *&uct_elem, uct_device_completion_t *&comp)
{
const size_t elem_size = sizeof(uct_device_remote_mem_list_elem_t);
ucs_status_t status;

status = ucp_device_check_params(dst_mem_list_h, dst_mem_list_index);
if (status != UCS_OK) {
return status;
}

const auto dst_mem_element = static_cast<uct_device_remote_mem_list_elem_t*>(
UCS_PTR_BYTE_OFFSET(dst_mem_list_h->mem_elements,
dst_mem_list_index * elem_size));
const auto dst_mem_element = UCP_DEVICE_GET_ELEM(dst_mem_list_h,
dst_mem_list_index);
remote_address = dst_mem_element->addr;
device_ep = dst_mem_element->device_ep;
uct_elem = &dst_mem_element->uct_mem_element;
device_ep = dst_mem_element->tl[lane].ep;
uct_elem = &dst_mem_element->tl[lane].uct;
ucp_device_request_init(device_ep, req, comp);

return UCS_OK;
}


UCS_F_DEVICE ucs_status_t
ucp_device_prepare_send(const ucp_device_local_mem_list_h src_mem_list_h,
unsigned src_mem_list_index,
const ucp_device_remote_mem_list_h dst_mem_list_h,
unsigned dst_mem_list_index, const void *&address,
uint64_t &remote_address, ucp_device_request_t *req,
uct_device_ep_t *&device_ep,
const uct_device_local_mem_list_elem_t *&src_uct_elem,
const uct_device_mem_element_t *&uct_elem,
uct_device_completion_t *&comp)
UCS_F_DEVICE ucs_status_t ucp_device_prepare_send(
const ucp_device_local_mem_list_h src_mem_list_h,
unsigned src_mem_list_index,
const ucp_device_remote_mem_list_h dst_mem_list_h,
unsigned dst_mem_list_index, const void *&address,
uint64_t &remote_address, unsigned lane, ucp_device_request_t *req,
uct_device_ep_t *&device_ep, const uct_device_mem_elem_t *&src_uct_elem,
const uct_device_mem_elem_t *&uct_elem, uct_device_completion_t *&comp)
{
const size_t elem_size = sizeof(uct_device_local_mem_list_elem_t);
ucs_status_t status;

status = ucp_device_check_params(src_mem_list_h, src_mem_list_index);
Expand All @@ -152,16 +159,16 @@ ucp_device_prepare_send(const ucp_device_local_mem_list_h src_mem_list_h,
}

status = ucp_device_prepare_send_remote(dst_mem_list_h, dst_mem_list_index,
remote_address, req, device_ep,
uct_elem, comp);
remote_address, lane, req,
device_ep, uct_elem, comp);
if (status != UCS_OK) {
return status;
}

src_uct_elem = static_cast<uct_device_local_mem_list_elem_t*>(
UCS_PTR_BYTE_OFFSET(src_mem_list_h->mem_elements,
src_mem_list_index * elem_size));
address = src_uct_elem->addr;
const auto src_mem_elem = UCP_DEVICE_GET_ELEM(src_mem_list_h,
src_mem_list_index);
src_uct_elem = src_mem_elem->tl + lane;
address = src_mem_elem->addr;

return UCS_OK;
}
Expand Down Expand Up @@ -212,18 +219,22 @@ ucp_device_put(const ucp_device_local_mem_list_h src_mem_list_h,
unsigned dst_mem_list_index, size_t dst_offset, size_t length,
unsigned channel_id, uint64_t flags, ucp_device_request_t *req)
{
unsigned uct_channel_id;
unsigned lane;
const void *address;
const uct_device_mem_element_t *uct_elem;
const uct_device_local_mem_list_elem_t *src_uct_elem;
const uct_device_mem_elem_t *uct_elem;
const uct_device_mem_elem_t *src_uct_elem;
uint64_t remote_address;
uct_device_completion_t *comp;
uct_device_ep_t *device_ep;
ucs_status_t status;

UCP_DEVICE_GET_LANE(dst_mem_list_h, channel_id, lane, uct_channel_id);

status = ucp_device_prepare_send(src_mem_list_h, src_mem_list_index,
dst_mem_list_h, dst_mem_list_index,
address, remote_address, req, device_ep,
src_uct_elem, uct_elem, comp);
address, remote_address, lane, req,
device_ep, src_uct_elem, uct_elem, comp);
if (status != UCS_OK) {
return status;
}
Expand All @@ -232,7 +243,7 @@ ucp_device_put(const ucp_device_local_mem_list_h src_mem_list_h,
src_uct_elem, uct_elem,
UCS_PTR_BYTE_OFFSET(address, src_offset),
remote_address + dst_offset, length,
channel_id, flags, comp);
uct_channel_id, flags, comp);
}


Expand Down Expand Up @@ -274,23 +285,27 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc(
unsigned mem_list_index, size_t offset, unsigned channel_id,
uint64_t flags, ucp_device_request_t *req)
{
unsigned uct_channel_id;
unsigned lane;
uint64_t remote_address;
const uct_device_mem_element_t *uct_elem;
const uct_device_mem_elem_t *uct_elem;
uct_device_completion_t *comp;
uct_device_ep_t *device_ep;
ucs_status_t status;

UCP_DEVICE_GET_LANE(mem_list_h, channel_id, lane, uct_channel_id);

status = ucp_device_prepare_send_remote(mem_list_h, mem_list_index,
remote_address, req, device_ep,
uct_elem, comp);
remote_address, lane, req,
device_ep, uct_elem, comp);
if (status != UCS_OK) {
return status;
}

return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_atomic_add, device_ep,
req, uct_elem, inc_value,
remote_address + offset, channel_id, flags,
comp);
remote_address + offset, uct_channel_id,
flags, comp);
}


Expand All @@ -311,20 +326,16 @@ UCS_F_DEVICE ucs_status_t
ucp_device_get_ptr(const ucp_device_remote_mem_list_h mem_list_h,
unsigned mem_list_index, void **addr_p)
{
const size_t elem_size = sizeof(uct_device_remote_mem_list_elem_t);
ucs_status_t status;

status = ucp_device_check_params(mem_list_h, mem_list_index);
if (status != UCS_OK) {
return status;
}

const auto mem_element = static_cast<uct_device_remote_mem_list_elem_t*>(
UCS_PTR_BYTE_OFFSET(mem_list_h->mem_elements,
mem_list_index * elem_size));
const auto mem_element = UCP_DEVICE_GET_ELEM(mem_list_h, mem_list_index);

return uct_device_ep_get_ptr(mem_element->device_ep,
&mem_element->uct_mem_element,
return uct_device_ep_get_ptr(mem_element->tl[0].ep, &mem_element->tl[0].uct,
mem_element->addr, addr_p);
}

Expand Down
24 changes: 17 additions & 7 deletions src/ucp/api/device/ucp_device_types.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
*/
Expand Down Expand Up @@ -30,17 +30,22 @@ typedef struct ucp_device_remote_mem_list {
* Structure version. Allow runtime ABI compatibility checks between host
* and device code.
*/
uint16_t version;
uint16_t version;

/**
* Number of lanes in each memory descriptor.
*/
uint16_t num_lanes;

@yosefe yosefe May 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

maybe uint8_t?


/**
* Number of entries in the memory descriptors array @a elems.
*/
uint32_t length;
uint32_t length;

/**
* UCT memory element objects are allocated contiguously.
*/
uct_device_remote_mem_list_elem_t mem_elements[0];
uct_device_remote_mem_elem_t mem_elements[0];
} ucp_device_remote_mem_list_t;


Expand All @@ -59,14 +64,19 @@ typedef struct ucp_device_local_mem_list {
* Structure version. Allow runtime ABI compatibility checks between host
* and device code.
*/
uint16_t version;
uint16_t version;

/**
* Number of lanes in each memory descriptor.
*/
uint16_t num_lanes;

@yosefe yosefe May 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

consider optimizing this to be something like MAX_LANES or MAX_MDS


/**
* Number of entries in the memory descriptors array @a elems.
*/
uint32_t length;
uint32_t length;

uct_device_local_mem_list_elem_t mem_elements[0];
uct_device_local_mem_elem_t mem_elements[0];
} ucp_device_local_mem_list_t;

#endif /* UCP_DEVICE_TYPES_H */
Loading
Loading