Skip to content

UCT/D2P: Multi-channel support with per-EP DPA address and RING_DB flag #11701

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 1 commit into from
Aug 2, 2026
Merged

UCT/D2P: Multi-channel support with per-EP DPA address and RING_DB flag #11701

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
8 changes: 5 additions & 3 deletions src/uct/api/device/uct_device_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,9 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put(
#endif
#if UCT_D2P_SUPPORTED
if (device_ep->uct_tl_id == UCT_DEVICE_TL_D2P) {
return uct_ib_d2p_ep_put<level>(device_ep, src_uct_elem, mem_elem, address,
remote_address, length, flags, comp);
return uct_ib_d2p_ep_put<level>(device_ep, src_uct_elem, mem_elem,
address, remote_address, length,
channel_id, flags, comp);
}
#endif

Expand Down Expand Up @@ -156,7 +157,8 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_atomic_add(
#if UCT_D2P_SUPPORTED
if (device_ep->uct_tl_id == UCT_DEVICE_TL_D2P) {
return uct_ib_d2p_ep_atomic_add<level>(device_ep, mem_elem, inc_value,
remote_address, flags, comp);
remote_address, channel_id,
flags, comp);
}
#endif

Expand Down
83 changes: 47 additions & 36 deletions src/uct/ib/mlx5/gdaki/d2p.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

template<ucs_device_level_t level>
UCS_F_DEVICE ucs_status_t uct_ib_d2p_post_desc(uct_ib_d2p_gpu_ep_t *ep,
unsigned channel_id,
uint8_t opcode, uint32_t length,
uint32_t lkey, uint64_t laddr,
uint32_t rkey, uint64_t raddr,
Expand All @@ -28,39 +29,41 @@ UCS_F_DEVICE ucs_status_t uct_ib_d2p_post_desc(uct_ib_d2p_gpu_ep_t *ep,
uct_dev_exec_init<level>(lane_id, num_lanes);

if (lane_id == 0) {
const long long depth = UCS_BIT(ep->log_depth);
unsigned long long pi = READ_ONCE(*ep->pi);
auto iface = ep->iface;
auto ch = iface->channels + (channel_id & iface->channel_mask);
const long long depth = UCS_BIT(iface->log_depth);
unsigned long long pi = READ_ONCE(*ch->pi);

for (;;) {
unsigned long long ci = READ_ONCE(*ep->ci);
unsigned long long ci = READ_ONCE(*ch->ci);
if (static_cast<long long>(pi - ci) >= depth) {
status = UCS_ERR_NO_RESOURCE;
break;
}

unsigned long long prev = atomicCAS(ep->pi, pi, pi + 1);
unsigned long long prev = atomicCAS(ch->pi, pi, pi + 1);
if (prev == pi) {
break;
}
pi = prev;
}

if (status == UCS_INPROGRESS) {
const uint32_t slot = pi & UCS_MASK(ep->log_depth);
auto desc = reinterpret_cast<volatile uct_ib_d2p_desc_t*>(
ep->queue_base) +
slot;

desc->opcode = opcode;
desc->length = length;
desc->qp_idx = ep->qp_idx;
desc->lkey = lkey;
desc->laddr = laddr;
desc->rkey = rkey;
desc->raddr = raddr;
desc->add = add;
desc->flags = flags;
const uint32_t owner = (pi >> ep->log_depth) & 0x1;
const uint32_t slot = pi & UCS_MASK(iface->log_depth);
auto desc = reinterpret_cast<volatile uct_ib_d2p_desc_t*>(
ch->queue_base) +
slot;

desc->opcode = opcode;
desc->length = length;
desc->ep_idx = ep->ep_idx;

@ofirfarjun7 ofirfarjun7 Aug 2, 2026 •
edited
Loading

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: ep_idx confused me, maybe we can come up with better name?

desc->lkey = lkey;
desc->laddr = laddr;
desc->rkey = rkey;
desc->raddr = raddr;
desc->add = add;
desc->flags = flags;
const uint32_t owner = (pi >> iface->log_depth) & 0x1;
asm volatile("st.release.sys.global.u8 [%0], %1;" ::
"l"(&desc->owner), "r"(owner) : "memory");
}
Expand All @@ -74,38 +77,46 @@ template<ucs_device_level_t level>
UCS_F_DEVICE ucs_status_t uct_ib_d2p_ep_put(
uct_device_ep_h tl_ep, const uct_device_mem_elem_t *src_uct_elem,
const uct_device_mem_elem_t *tl_mem_elem, const void *address,
uint64_t remote_address, size_t length, uint64_t flags,
uct_device_completion_t *comp)
uint64_t remote_address, size_t length, unsigned channel_id,
uint64_t flags, uct_device_completion_t *comp)
{
auto ep = reinterpret_cast<uct_ib_d2p_gpu_ep_t*>(tl_ep);
auto src_ib = reinterpret_cast<const uct_ib_md_device_mem_element_t*>(
src_uct_elem);
auto rem_ib = reinterpret_cast<const uct_ib_md_device_mem_element_t*>(
tl_mem_elem);

return uct_ib_d2p_post_desc<level>(
ep, UCT_IB_D2P_OP_RDMA_WRITE, length, src_ib->lkey,
reinterpret_cast<uint64_t>(address), rem_ib->rkey, remote_address,
0, comp == nullptr ? 0 : UCT_IB_D2P_FLAG_CQ_UPDATE);
uint16_t desc_flags = (comp != nullptr ? UCT_IB_D2P_FLAG_CQ_UPDATE : 0) |
(flags & UCT_DEVICE_FLAG_NODELAY ?
UCT_IB_D2P_FLAG_RING_DB :

@ofirfarjun7 ofirfarjun7 Aug 2, 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: Maybe we can remove this translation and just use UCT_DEVICE_FLAG_NODELAY?

0);

return uct_ib_d2p_post_desc<level>(ep, channel_id, UCT_IB_D2P_OP_RDMA_WRITE,
length, src_ib->lkey,
reinterpret_cast<uint64_t>(address),
rem_ib->rkey, remote_address, 0,
desc_flags);
}

template<ucs_device_level_t level>
UCS_F_DEVICE ucs_status_t uct_ib_d2p_ep_atomic_add(
uct_device_ep_h tl_ep, const uct_device_mem_elem_t *tl_mem_elem,
uint64_t inc_value, uint64_t remote_address, uint64_t flags,
uct_device_completion_t *comp)
uint64_t inc_value, uint64_t remote_address, unsigned channel_id,
uint64_t flags, uct_device_completion_t *comp)
{
auto ep = reinterpret_cast<uct_ib_d2p_gpu_ep_t*>(tl_ep);
auto rem_ib = reinterpret_cast<const uct_ib_md_device_mem_element_t*>(
tl_mem_elem);

return uct_ib_d2p_post_desc<level>(ep, UCT_IB_D2P_OP_ATOMIC_ADD,
sizeof(uint64_t), ep->atomic_result_lkey,
ep->atomic_result_va, rem_ib->rkey,
remote_address, inc_value,
comp == nullptr ?
0 :
UCT_IB_D2P_FLAG_CQ_UPDATE);
uint16_t desc_flags = (comp != nullptr ? UCT_IB_D2P_FLAG_CQ_UPDATE : 0) |
(flags & UCT_DEVICE_FLAG_NODELAY ?
UCT_IB_D2P_FLAG_RING_DB :
0);

return uct_ib_d2p_post_desc<level>(ep, channel_id, UCT_IB_D2P_OP_ATOMIC_ADD,
sizeof(uint64_t),
ep->iface->atomic_result_lkey,
ep->iface->atomic_result_va,
rem_ib->rkey, remote_address, inc_value,
desc_flags);
}

template<ucs_device_level_t level>
Expand Down
22 changes: 16 additions & 6 deletions src/uct/ib/mlx5/gdaki/d2p.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,25 @@
#include <uct/api/device/uct_device_types.h>

typedef struct {
uct_device_ep_t super;
unsigned long long *pi;
unsigned long long *ci;
void *queue_base;
uint64_t qp_idx;
uint64_t atomic_result_va;
uint32_t atomic_result_lkey;
uint8_t log_depth;
uint8_t pad[3];
} uct_ib_d2p_gpu_channel_t;

typedef struct {
uint8_t channel_mask;
uint8_t log_depth;
uint8_t pad[2];
uint32_t atomic_result_lkey;
uint64_t atomic_result_va;
uct_ib_d2p_gpu_channel_t channels[];
} uct_ib_d2p_gpu_iface_t;

typedef struct {
uct_device_ep_t super;
uct_ib_d2p_gpu_iface_t *iface;
uint64_t ep_idx;
uint8_t pad[8];
} uct_ib_d2p_gpu_ep_t;

#endif /* UCT_D2P_H_ */
3 changes: 2 additions & 1 deletion src/uct/ib/mlx5/gdaki/d2p_proto.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ enum {

enum {
UCT_IB_D2P_FLAG_CQ_UPDATE = UCS_BIT(0),
UCT_IB_D2P_FLAG_RING_DB = UCS_BIT(1),
};


Expand All @@ -25,7 +26,7 @@ typedef struct {
uint8_t opcode;
uint16_t flags;
uint32_t length;
uint64_t qp_idx;
uint64_t ep_idx;
uint32_t lkey;
uint32_t rkey;
uint64_t laddr;
Expand Down