From fa7e9fbefd722a3f1c82c0f9f164a4959d26afd8 Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Mon, 9 Mar 2026 19:02:05 +0200 Subject: [PATCH 01/12] UCP/DEVICE: Add multi-lane support --- src/ucp/api/device/ucp_device_impl.h | 47 ++++-- src/ucp/api/device/ucp_device_types.h | 3 + src/ucp/core/ucp_device.c | 219 ++++++++------------------ src/ucp/wireup/address.c | 3 +- src/ucp/wireup/select.c | 8 +- src/uct/api/uct.h | 3 + src/uct/ib/mlx5/gdaki/gdaki.c | 6 +- test/gtest/ucp/cuda/test_kernels.cu | 62 -------- test/gtest/ucp/cuda/test_kernels.h | 2 - test/gtest/ucp/test_ucp_device.cc | 18 +-- test/gtest/uct/test_device.cc | 24 +-- 11 files changed, 125 insertions(+), 270 deletions(-) diff --git a/src/ucp/api/device/ucp_device_impl.h b/src/ucp/api/device/ucp_device_impl.h index 86388082b84..07ca5729a44 100644 --- a/src/ucp/api/device/ucp_device_impl.h +++ b/src/ucp/api/device/ucp_device_impl.h @@ -105,14 +105,28 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep, }) +#define UCP_DEVICE_GET_LANE(_handle, _channel_id) \ + ({ \ + unsigned _lane = _channel_id % _handle->num_lanes; \ + _channel_id /= _handle->num_lanes; \ + _lane; \ + }) + + +#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \ + static_castmem_elements[0])*>( \ + UCS_PTR_BYTE_OFFSET(_handle->mem_elements, \ + ((_index * _handle->num_lanes) + _lane) * \ + sizeof(_handle->mem_elements[0]))) + + 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 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); @@ -120,9 +134,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote( return status; } - const auto dst_mem_element = static_cast( - 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, lane); remote_address = dst_mem_element->addr; device_ep = dst_mem_element->device_ep; uct_elem = &dst_mem_element->uct_mem_element; @@ -137,13 +150,12 @@ 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, + uint64_t &remote_address, unsigned lane, + 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) { - 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); @@ -152,15 +164,14 @@ 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( - UCS_PTR_BYTE_OFFSET(src_mem_list_h->mem_elements, - src_mem_list_index * elem_size)); + src_uct_elem = UCP_DEVICE_GET_ELEM(src_mem_list_h, src_mem_list_index, + lane); address = src_uct_elem->addr; return UCS_OK; @@ -212,6 +223,7 @@ 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) { + const unsigned lane = UCP_DEVICE_GET_LANE(dst_mem_list_h, channel_id); const void *address; const uct_device_mem_element_t *uct_elem; const uct_device_local_mem_list_elem_t *src_uct_elem; @@ -222,8 +234,8 @@ ucp_device_put(const ucp_device_local_mem_list_h src_mem_list_h, 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; } @@ -274,6 +286,7 @@ 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) { + const unsigned lane = UCP_DEVICE_GET_LANE(mem_list_h, channel_id); uint64_t remote_address; const uct_device_mem_element_t *uct_elem; uct_device_completion_t *comp; @@ -281,8 +294,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc( ucs_status_t status; 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; } diff --git a/src/ucp/api/device/ucp_device_types.h b/src/ucp/api/device/ucp_device_types.h index 7847f314fdd..71933280eac 100644 --- a/src/ucp/api/device/ucp_device_types.h +++ b/src/ucp/api/device/ucp_device_types.h @@ -32,6 +32,8 @@ typedef struct ucp_device_remote_mem_list { */ uint16_t version; + uint16_t num_lanes; + /** * Number of entries in the memory descriptors array @a elems. */ @@ -61,6 +63,7 @@ typedef struct ucp_device_local_mem_list { */ uint16_t version; + uint16_t num_lanes; /** * Number of entries in the memory descriptors array @a elems. */ diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 0fedc394803..2bcd885a3ef 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -138,127 +138,28 @@ ucp_device_detect_local_sys_dev(ucp_context_h context, return UCS_OK; } -static ucp_md_map_t -ucp_device_detect_local_md_map(const ucp_context_h context, - ucs_sys_device_t local_sys_dev) +static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, + ucp_tl_bitmap_t *tl_bitmap, + ucs_sys_device_t local_sys_dev) { - ucp_md_map_t local_md_map = 0; - ucp_md_index_t md_index; - - /* Build MD map from MDs that can access the local_sys_dev */ - for (md_index = 0; md_index < context->num_mds; md_index++) { - ucp_sys_dev_map_t sys_dev_map = context->tl_mds[md_index].sys_dev_map; - - if (sys_dev_map & UCS_BIT(local_sys_dev)) { - local_md_map |= UCS_BIT(md_index); - } - } - - ucs_trace("detected local_md_map=0x%" PRIx64 " for local_sys_dev=%u", - local_md_map, local_sys_dev); - return local_md_map; -} - -static void ucp_device_mem_list_lane_lookup( - ucp_ep_h ep, ucp_ep_config_t *ep_config, ucs_sys_device_t local_sys_dev, - ucp_md_map_t local_md_map, ucs_sys_device_t remote_sys_dev, - ucp_md_map_t remote_md_map, - ucp_lane_index_t lanes[UCP_DEVICE_MEM_LIST_MAX_EPS]) -{ - double best_bw[UCP_DEVICE_MEM_LIST_MAX_EPS] = {-1., -1.}; - ucp_lane_index_t lane; - double bandwidth; - ucp_ep_config_key_lane_t *lane_key; - ucs_sys_device_t src_sys_dev; - ucp_md_index_t src_md_index; - - lanes[0] = UCP_NULL_LANE; - lanes[1] = UCP_NULL_LANE; - - for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { - if (!(ep_config->key.lanes[lane].lane_types & - UCS_BIT(UCP_LANE_TYPE_DEVICE))) { - continue; - } - - lane_key = &ep_config->key.lanes[lane]; - /* Check lane remote sys dev only when remote memory is not host */ - if ((remote_sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) && - (remote_sys_dev != lane_key->dst_sys_dev)) { - ucs_trace("lane[%u] wrong destination sys_dev: dst_sys_dev=%u", - lane, lane_key->dst_sys_dev); - continue; - } - - if (!(remote_md_map & UCS_BIT(lane_key->dst_md_index))) { - ucs_trace("lane[%u] missing remote md: dst_md_index=%u", lane, - lane_key->dst_md_index); - continue; - } - - src_sys_dev = ucp_ep_get_tl_rsc(ep, lane)->sys_device; - if (src_sys_dev != local_sys_dev) { - ucs_trace("lane[%u] wrong source sys_dev: src_sys_dev=%u", lane, - src_sys_dev); - continue; - } + const ucp_worker_iface_t *wiface; + ucp_rsc_index_t tl_id; - src_md_index = ucp_ep_md_index(ep, lane); - if (!(local_md_map & UCS_BIT(src_md_index))) { - ucs_trace("lane[%u] missing local md: src_md_index=%u", lane, - src_md_index); + /** TODO: Maybe cache results */ + UCS_STATIC_BITMAP_RESET_ALL(tl_bitmap); + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &worker->context->tl_bitmap) { + wiface = ucp_worker_iface(worker, tl_id); + if (!(wiface->attr.cap.flags & UCT_IFACE_FLAG_DEVICE_EP)) { continue; } - - bandwidth = ucp_worker_iface_bandwidth(ep->worker, - ucp_ep_get_rsc_index(ep, lane)); - ucs_trace("checking lane[%u] src_md_index=%u dst_md_index=%u " - "src_sys_dev=%u dst_sys_dev=%u bandwidth=%lfMB/s", - lane, src_md_index, lane_key->dst_md_index, src_sys_dev, - lane_key->dst_sys_dev, bandwidth / UCS_MBYTE); - - UCS_STATIC_ASSERT(UCP_DEVICE_MEM_LIST_MAX_EPS == 2); - if (bandwidth > best_bw[0]) { - best_bw[1] = best_bw[0]; - lanes[1] = lanes[0]; - best_bw[0] = bandwidth; - lanes[0] = lane; - } else if (bandwidth > best_bw[1]) { - best_bw[1] = bandwidth; - lanes[1] = lane; - } else { + if (wiface->attr.ctl_device != local_sys_dev) { continue; } - ucs_trace("best lanes: lane[%u]=%lfMB/s lane[%u]=%lfMB/s", lanes[0], - best_bw[0] / UCS_MBYTE, lanes[1], best_bw[1] / UCS_MBYTE); + UCS_STATIC_BITMAP_SET(tl_bitmap, tl_id); } } - -static ucp_worker_iface_t * -ucp_device_get_worker_iface_by_device_id(ucp_worker_h worker, - ucs_sys_device_t device_mem_id) -{ - const ucp_tl_resource_desc_t *resource; - const uct_md_attr_v2_t *md_attr; - ucp_md_index_t md_index; - unsigned i; - - /** TODO: Maybe cache results */ - for (i = 0; i < worker->num_ifaces; i++) { - resource = &worker->context->tl_rscs[worker->ifaces[i]->rsc_index]; - md_index = resource->md_index; - md_attr = &worker->context->tl_mds[md_index].attr; - if ((md_attr->flags & UCT_MD_FLAG_NEED_MEMH) && - (resource->tl_rsc.sys_device == device_mem_id)) { - return worker->ifaces[i]; - } - } - - return NULL; -} - 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, @@ -334,41 +235,46 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( const ucp_worker_iface_t *wiface; ucp_device_local_mem_list_t *handle; uct_device_local_mem_list_elem_t *uct_element; - size_t i; + size_t i, num_lanes; ucs_status_t status; + ucp_tl_bitmap_t tl_bitmap; + ucp_rsc_index_t tl_id; - handle_size = (uct_elem_size * params->num_elements) + sizeof(*handle); + ucp_device_get_tl_bitmap(worker, &tl_bitmap, local_sys_dev); + num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap); + handle_size = (uct_elem_size * params->num_elements * num_lanes) + + sizeof(*handle); handle = ucs_calloc(1, handle_size, "ucp_device_local_mem_list_t"); if (handle == NULL) { ucs_error("failed to allocate ucp_device_local_mem_list_t"); return UCS_ERR_NO_MEMORY; } - /* TODO: To support multi lanes we need to pack all memhs of ifaces that require memhs */ - wiface = ucp_device_get_worker_iface_by_device_id(worker, local_sys_dev); - if (wiface == NULL) { - ucs_debug("no worker iface found for device_id=%u", local_sys_dev); - } - /* Populate element specific parameters */ ucp_element = params->elements; uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); for (i = 0; i < params->num_elements; i++) { - status = ucp_device_local_mem_list_element_pack(worker, wiface, - ucp_element, mem_type, - uct_element); - if (status != UCS_OK) { - ucs_error("failed to pack local mem list element for element=%zu", - i); - goto out; - } + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { + wiface = ucp_worker_iface(worker, tl_id); + status = ucp_device_local_mem_list_element_pack(worker, wiface, + ucp_element, + mem_type, + uct_element); + if (status != UCS_OK) { + ucs_error( + "failed to pack local mem list element for element=%zu", + i); + goto out; + } + uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); + } ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); } handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1; handle->length = params->num_elements; + handle->num_lanes = num_lanes; status = ucp_device_mem_list_export_handle( worker, handle, handle_size, mem_type, local_sys_dev, mem, "ucp_device_local_mem_list_handle_t"); @@ -467,21 +373,12 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params, } static ucs_status_t ucp_device_remote_mem_list_element_pack( - const ucp_device_mem_list_elem_t *element, - const ucs_sys_device_t local_sys_dev, const ucs_memory_type_t mem_type, + const ucp_device_mem_list_elem_t *element, ucp_rsc_index_t tl_id, uct_device_remote_mem_list_elem_t *mem_element) { const ucp_ep_h ep = element->ep; const ucp_rkey_h rkey = element->rkey; - const ucp_md_map_t local_md_map = - ucp_device_detect_local_md_map(ep->worker->context, local_sys_dev); - const ucp_worker_cfg_index_t rkey_cfg_index = element->rkey->cfg_index; - const ucp_rkey_config_t rkey_config = - ucs_array_elem(&ep->worker->rkey_config, rkey_cfg_index); - const ucs_sys_device_t remote_sys_dev = rkey_config.key.sys_dev; - const ucp_md_map_t remote_md_map = rkey_config.key.md_map; ucp_ep_config_t *ep_config = ucp_ep_config(ep); - ucp_lane_index_t lanes[UCP_DEVICE_MEM_LIST_MAX_EPS]; uint8_t rkey_index; uct_rkey_t uct_rkey; uct_ep_h uct_ep; @@ -489,11 +386,13 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack( ucs_status_t status; ucp_lane_index_t lane; - ucp_device_mem_list_lane_lookup(ep, ep_config, local_sys_dev, local_md_map, - remote_sys_dev, remote_md_map, lanes); - lane = lanes[0]; + for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { + if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { + break; + } + } - if (lane == UCP_NULL_LANE) { + if (lane == ep_config->key.num_lanes) { ucs_error("no lane found for ep=%p", ep); return UCS_ERR_NO_DEVICE; } @@ -566,7 +465,9 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucp_device_remote_mem_list_t *handle; uct_device_remote_mem_list_elem_t *uct_element; ucs_sys_device_t local_sys_dev; - size_t i; + ucp_tl_bitmap_t tl_bitmap; + ucp_rsc_index_t tl_id; + size_t i, num_lanes; ucs_status_t status; if (ep == NULL) { @@ -580,7 +481,10 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( return status; } - handle_size = sizeof(*handle) + (uct_elem_size * params->num_elements); + ucp_device_get_tl_bitmap(ep->worker, &tl_bitmap, local_sys_dev); + num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap); + handle_size = sizeof(*handle) + + (uct_elem_size * params->num_elements * num_lanes); handle = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t"); if (handle == NULL) { ucs_error("failed to allocate ucp_device_remote_mem_list_t"); @@ -590,24 +494,30 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucp_element = params->elements; uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); for (i = 0; i < params->num_elements; i++) { - if (!UCP_DEVICE_MEM_ELEMENT_IS_GAP(ucp_element)) { - status = ucp_device_remote_mem_list_element_pack(ucp_element, - local_sys_dev, - mem_type, - uct_element); - if (status != UCS_OK) { - ucs_error("failed to pack uct memory element for element=%zu", - i); - goto out; + if (UCP_DEVICE_MEM_ELEMENT_IS_GAP(ucp_element)) { + uct_element = UCS_PTR_BYTE_OFFSET(uct_element, + uct_elem_size * num_lanes); + } else { + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { + status = ucp_device_remote_mem_list_element_pack(ucp_element, + tl_id, + uct_element); + if (status != UCS_OK) { + ucs_error( + "failed to pack uct memory element for element=%zu", + i); + goto out; + } + uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); } } ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); } handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1; handle->length = params->num_elements; + handle->num_lanes = num_lanes; status = ucp_device_mem_list_export_handle( ep->worker, handle, handle_size, mem_type, local_sys_dev, mem, "ucp_device_remote_mem_list_handle_t"); @@ -716,7 +626,6 @@ ucp_device_remote_mem_list_create(const ucp_device_mem_list_params_t *params, return status; } - uint32_t ucp_device_get_mem_list_length(const void *mem_list_h) { khiter_t iter; diff --git a/src/ucp/wireup/address.c b/src/ucp/wireup/address.c index 71b0ea1d0fc..7e70a407852 100644 --- a/src/ucp/wireup/address.c +++ b/src/ucp/wireup/address.c @@ -455,7 +455,8 @@ ucp_address_gather_devices(ucp_worker_h worker, const ucp_ep_config_key_t *key, dev->rsc_index = rsc_index; UCS_STATIC_BITMAP_SET(&dev->tl_bitmap, rsc_index); - dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); + + dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); } *devices_p = devices; diff --git a/src/ucp/wireup/select.c b/src/ucp/wireup/select.c index 1111ff8317b..4d6d72b5bb5 100644 --- a/src/ucp/wireup/select.c +++ b/src/ucp/wireup/select.c @@ -1821,7 +1821,13 @@ static int ucp_wireup_add_bw_lanes_pairwise( /* Account for possible path override */ local_num_paths = iface_attr->dev_num_paths; - remote_num_paths = ae->dev_num_paths; + + if (bw_info->criteria.lane_type == UCP_LANE_TYPE_DEVICE) { + remote_num_paths = 1; + } else { + remote_num_paths = ae->dev_num_paths; + } + if (allow_extra_path && ((skip_dev_index != UCP_NULL_RESOURCE) && /* clang sanitizer */ (skip_dev_index == dev_index))) { diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index 02ffab3ced3..b06f9fcbf2e 100644 --- a/src/uct/api/uct.h +++ b/src/uct/api/uct.h @@ -1183,6 +1183,9 @@ struct uct_iface_attr { achieve higher total bandwidth compared to using only a single endpoint. */ + + ucs_sys_device_t ctl_device; /**< System device controlling this iface, + if it's not CPU. */ }; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index 414b5c9c856..f0b72e2688f 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -511,6 +511,7 @@ uct_rc_gdaki_iface_query(uct_iface_h tl_iface, uct_iface_attr_t *iface_attr) iface_attr->cap.put.min_zcopy = 0; iface_attr->cap.put.max_zcopy = uct_ib_iface_port_attr(&iface->super.super.super)->max_msg_sz; + iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev); return UCS_OK; } @@ -1053,7 +1054,6 @@ uct_gdaki_query_tl_devices(uct_md_h tl_md, uct_tl_device_resource_t *tl_devices; ucs_status_t status; CUdevice device; - ucs_sys_device_t dev; int i; uct_gdaki_dev_matrix_elem_t *ibdesc; char dmabuf_str[8]; @@ -1120,14 +1120,12 @@ uct_gdaki_query_tl_devices(uct_md_h tl_md, goto err; } - dev = uct_cuda_get_sys_dev(device); - snprintf(tl_devices[num_tl_devices].name, sizeof(tl_devices[num_tl_devices].name), "%s%d-%s:%d", UCT_DEVICE_CUDA_NAME, device, uct_ib_device_name(&ib_md->dev), ib_md->dev.first_port); tl_devices[num_tl_devices].type = UCT_DEVICE_TYPE_NET; - tl_devices[num_tl_devices].sys_device = dev; + tl_devices[num_tl_devices].sys_device = ib_md->dev.sys_dev; num_tl_devices++; } diff --git a/test/gtest/ucp/cuda/test_kernels.cu b/test/gtest/ucp/cuda/test_kernels.cu index 309bb2fa265..3894c576190 100644 --- a/test/gtest/ucp/cuda/test_kernels.cu +++ b/test/gtest/ucp/cuda/test_kernels.cu @@ -108,56 +108,6 @@ private: ucp_device_request_t *m_ptr; }; -UCS_F_DEVICE ucs_status_t -ucp_test_kernel_get_state(const test_ucp_device_kernel_params_t ¶ms, - test_ucp_device_kernel_result_t &result) -{ - ucp_device_request_t *req_ptr = nullptr; - const uct_device_mem_element_t *uct_elem; - uct_device_ep_t *device_ep; - uint64_t remote_address; - uct_device_completion_t *comp; - ucs_status_t status = UCS_OK; - - if (nullptr == params.remote_mem_list) { - return UCS_OK; - } - - __syncthreads(); - if (threadIdx.x == 0) { - for (unsigned i = 0; i < params.remote_mem_list->length; ++i) { - status = ucp_device_prepare_send_remote(params.remote_mem_list, i, - remote_address, req_ptr, - device_ep, uct_elem, comp); - if ((status == UCS_OK) && (device_ep != nullptr)) { - break; - } - } - - result.producer_index = 0; - result.ready_index = 0; -#if HAVE_MLX5_DV - if ((status == UCS_OK) && - (device_ep != nullptr) && - (device_ep->uct_tl_id == UCT_DEVICE_TL_RC_MLX5_GDA)) { - uint16_t wqe_cnt; - uct_rc_gdaki_dev_ep_t *ep = - reinterpret_cast(device_ep); - unsigned i; - - for (i = 0; i < params.num_channels; i++) { - result.producer_index += - uct_rc_mlx5_gda_parse_cqe(ep, i, &wqe_cnt, nullptr) + 1; - result.ready_index += ep->qps[i].sq_ready_index; - } - } -#endif - } - - __syncthreads(); - return status; -} - template UCS_F_DEVICE void ucp_test_kernel_job(const test_ucp_device_kernel_params_t ¶ms, @@ -175,11 +125,6 @@ ucp_test_kernel_job(const test_ucp_device_kernel_params_t ¶ms, shared_reqs[device_request::num_shared_reqs()]; device_request req(shared_reqs); - status = ucp_test_kernel_get_state(params, *result_ptr); - if (status != UCS_OK) { - return; - } - ucp_device_request_t *req_ptr = params.with_request ? req.ptr() : nullptr; uint64_t flags = params.with_no_delay ? UCT_DEVICE_FLAG_NODELAY : 0; @@ -196,11 +141,6 @@ ucp_test_kernel_job(const test_ucp_device_kernel_params_t ¶ms, // function to the API. status = ucp_test_kernel_do_operation(params, UCT_DEVICE_FLAG_NODELAY, req.ptr()); - if (status != UCS_OK) { - return; - } - - status = ucp_test_kernel_get_state(params, *result_ptr); } template @@ -256,8 +196,6 @@ launch_test_ucp_device_kernel(const test_ucp_device_kernel_params_t ¶ms) ucx_cuda::device_result_ptr result; result->status = UCS_ERR_NOT_IMPLEMENTED; - result->producer_index = 0; - result->ready_index = 0; switch (params.level) { case UCS_DEVICE_LEVEL_THREAD: diff --git a/test/gtest/ucp/cuda/test_kernels.h b/test/gtest/ucp/cuda/test_kernels.h index fcac4d39312..34ec4a972e0 100644 --- a/test/gtest/ucp/cuda/test_kernels.h +++ b/test/gtest/ucp/cuda/test_kernels.h @@ -51,8 +51,6 @@ typedef struct { struct test_ucp_device_kernel_result_t { ucs_status_t status; - uint64_t producer_index; - uint64_t ready_index; }; test_ucp_device_kernel_result_t diff --git a/test/gtest/ucp/test_ucp_device.cc b/test/gtest/ucp/test_ucp_device.cc index fe2f192d235..586a9d277ef 100644 --- a/test/gtest/ucp/test_ucp_device.cc +++ b/test/gtest/ucp/test_ucp_device.cc @@ -538,21 +538,6 @@ class test_ucp_device_kernel : public test_ucp_device { ASSERT_UCS_OK(result.status); return result; } - - void check_result(const test_ucp_device_kernel_params_t ¶ms, - const test_ucp_device_kernel_result_t &result, - unsigned count) - { - unsigned num_threads = params.num_threads; - if (params.level == UCS_DEVICE_LEVEL_WARP) { - num_threads /= UCS_DEVICE_NUM_THREADS_IN_WARP; - } - - uint64_t expected = params.num_iters * num_threads * count; - EXPECT_UCS_OK(result.status); - EXPECT_EQ(expected, result.producer_index); - EXPECT_EQ(expected, result.ready_index); - } }; UCS_TEST_P(test_ucp_device_kernel, local_counter) @@ -781,11 +766,10 @@ UCS_TEST_SKIP_COND_P(test_ucp_device_xfer, put_stress_test, params.num_iters = 1000; params.num_blocks = 1; params.num_threads = MAX_THREADS; - auto result = launch_kernel(params); + launch_kernel(params); // Check proper index received data list.dst_pattern_check(mem_list_index, mem_list::SEED_SRC); - check_result(params, result, 1); } UCS_TEST_P(test_ucp_device_xfer, counter) diff --git a/test/gtest/uct/test_device.cc b/test/gtest/uct/test_device.cc index ffb320d96ef..c565fa966d9 100644 --- a/test/gtest/uct/test_device.cc +++ b/test/gtest/uct/test_device.cc @@ -101,17 +101,6 @@ class test_device : public uct_test { ucs_status_t status; uct_test::init(); - m_cuda_dev = uct_cuda_get_cuda_device(GetParam()->sys_device); - ASSERT_NE(m_cuda_dev, CU_DEVICE_INVALID) << " sys_device " - << static_cast(GetParam()->sys_device); - - status = UCT_CUDADRV_FUNC_LOG_ERR( - cuDevicePrimaryCtxRetain(&ctx, m_cuda_dev)); - ASSERT_UCS_OK(status); - - status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(ctx)); - ASSERT_UCS_OK(status); - m_receiver = uct_test::create_entity(0); m_entities.push_back(m_receiver); @@ -119,6 +108,19 @@ class test_device : public uct_test { m_entities.push_back(m_sender); m_sender->connect(0, *m_receiver, 0); + + m_cuda_dev = uct_cuda_get_cuda_device( + m_sender->iface_attr().ctl_device); + ASSERT_NE(m_cuda_dev, CU_DEVICE_INVALID) + << " sys_device " + << static_cast(m_sender->iface_attr().ctl_device); + + status = UCT_CUDADRV_FUNC_LOG_ERR( + cuDevicePrimaryCtxRetain(&ctx, m_cuda_dev)); + ASSERT_UCS_OK(status); + + status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(ctx)); + ASSERT_UCS_OK(status); } void cleanup() From 60e5be916fe168467b38f932d05af05fc907bb7f Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Wed, 18 Mar 2026 00:43:47 +0200 Subject: [PATCH 02/12] UCP/DEVICE: Add multi-lane support - 2 --- src/ucp/core/ucp_device.c | 109 +++++++++++++++++++++++++++++----- src/ucp/wireup/address.c | 3 +- src/uct/api/uct.h | 1 + src/uct/ib/mlx5/gdaki/gdaki.c | 1 + 4 files changed, 97 insertions(+), 17 deletions(-) diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 2bcd885a3ef..3deae80e040 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -140,6 +140,7 @@ ucp_device_detect_local_sys_dev(ucp_context_h context, static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, ucp_tl_bitmap_t *tl_bitmap, + ucp_tl_bitmap_t *tl_bitmap_nolkey, ucs_sys_device_t local_sys_dev) { const ucp_worker_iface_t *wiface; @@ -147,16 +148,23 @@ static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, /** TODO: Maybe cache results */ UCS_STATIC_BITMAP_RESET_ALL(tl_bitmap); + UCS_STATIC_BITMAP_RESET_ALL(tl_bitmap_nolkey); UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &worker->context->tl_bitmap) { wiface = ucp_worker_iface(worker, tl_id); + if (!(wiface->attr.cap.flags & UCT_IFACE_FLAG_DEVICE_EP)) { continue; } + if (wiface->attr.ctl_device != local_sys_dev) { continue; } - UCS_STATIC_BITMAP_SET(tl_bitmap, tl_id); + if (wiface->attr.cap.flags & UCT_IFACE_FLAG_DEVICE_LKEY) { + UCS_STATIC_BITMAP_SET(tl_bitmap, tl_id); + } else { + UCS_STATIC_BITMAP_SET(tl_bitmap_nolkey, tl_id); + } } } @@ -238,9 +246,11 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( size_t i, num_lanes; ucs_status_t status; ucp_tl_bitmap_t tl_bitmap; + ucp_tl_bitmap_t tl_bitmap_nolkey; ucp_rsc_index_t tl_id; - ucp_device_get_tl_bitmap(worker, &tl_bitmap, local_sys_dev); + ucp_device_get_tl_bitmap(worker, &tl_bitmap, &tl_bitmap_nolkey, + local_sys_dev); num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap); handle_size = (uct_elem_size * params->num_elements * num_lanes) + sizeof(*handle); @@ -372,6 +382,28 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params, return status; } +static int ucp_device_remote_mem_list_check_lanes(const ucp_ep_h ep, + ucp_tl_bitmap_t *tl_bitmap) + +{ + ucp_ep_config_t *ep_config = ucp_ep_config(ep); + ucp_lane_index_t lane; + ucp_rsc_index_t tl_id; + + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) { + for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { + if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { + break; + } + } + if (lane == ep_config->key.num_lanes) { + return 0; + } + } + + return 1; +} + 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_mem_list_elem_t *mem_element) @@ -454,6 +486,29 @@ 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, + uct_device_remote_mem_list_elem_t **uct_element_p) +{ + const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); + ucp_rsc_index_t tl_id; + ucs_status_t status; + + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) { + status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id, + *uct_element_p); + if (status != UCS_OK) { + return status; + } + + *uct_element_p = UCS_PTR_BYTE_OFFSET(*uct_element_p, uct_elem_size); + } + + return UCS_OK; +} + + static ucs_status_t ucp_device_remote_mem_list_create_handle( const ucp_device_mem_list_params_t *params, ucs_memory_type_t mem_type, uct_allocated_memory_t *mem) @@ -466,9 +521,10 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( uct_device_remote_mem_list_elem_t *uct_element; ucs_sys_device_t local_sys_dev; ucp_tl_bitmap_t tl_bitmap; - ucp_rsc_index_t tl_id; + ucp_tl_bitmap_t tl_bitmap_nolkey; size_t i, num_lanes; ucs_status_t status; + int has_lkey, has_nolkey; if (ep == NULL) { ucs_error("no ep found in remote mem list"); @@ -481,8 +537,20 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( return status; } - ucp_device_get_tl_bitmap(ep->worker, &tl_bitmap, local_sys_dev); - num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap); + ucp_device_get_tl_bitmap(ep->worker, &tl_bitmap, &tl_bitmap_nolkey, + local_sys_dev); + ucp_element = params->elements; + + has_lkey = ucp_device_remote_mem_list_check_lanes(ep, &tl_bitmap); + has_nolkey = ucp_device_remote_mem_list_check_lanes(ep, &tl_bitmap_nolkey); + + num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap) * has_lkey + + UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap_nolkey) * has_nolkey; + + if (!num_lanes) { + ucs_error("failed to pack uct memory element for first element"); + } + handle_size = sizeof(*handle) + (uct_elem_size * params->num_elements * num_lanes); handle = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t"); @@ -491,27 +559,38 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( return UCS_ERR_NO_MEMORY; } - ucp_element = params->elements; uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); for (i = 0; i < params->num_elements; i++) { if (UCP_DEVICE_MEM_ELEMENT_IS_GAP(ucp_element)) { uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size * num_lanes); } else { - UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { - status = ucp_device_remote_mem_list_element_pack(ucp_element, - tl_id, - uct_element); + ucs_assert(has_lkey == + ucp_device_remote_mem_list_check_lanes(ucp_element->ep, + &tl_bitmap)); + ucs_assert( + has_nolkey == + ucp_device_remote_mem_list_check_lanes(ucp_element->ep, + &tl_bitmap_nolkey)); + + if (has_lkey) { + status = ucp_device_remote_mem_list_fill(ucp_element, + &tl_bitmap, + &uct_element); if (status != UCS_OK) { - ucs_error( - "failed to pack uct memory element for element=%zu", - i); goto out; } - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); } - } + if (has_nolkey) { + status = ucp_device_remote_mem_list_fill(ucp_element, + &tl_bitmap_nolkey, + &uct_element); + if (status != UCS_OK) { + goto out; + } + } + } ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); } diff --git a/src/ucp/wireup/address.c b/src/ucp/wireup/address.c index 7e70a407852..71b0ea1d0fc 100644 --- a/src/ucp/wireup/address.c +++ b/src/ucp/wireup/address.c @@ -455,8 +455,7 @@ ucp_address_gather_devices(ucp_worker_h worker, const ucp_ep_config_key_t *key, dev->rsc_index = rsc_index; UCS_STATIC_BITMAP_SET(&dev->tl_bitmap, rsc_index); - - dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); + dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); } *devices_p = devices; diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index b06f9fcbf2e..d26e976a2a4 100644 --- a/src/uct/api/uct.h +++ b/src/uct/api/uct.h @@ -438,6 +438,7 @@ typedef enum uct_atomic_op { /* Interface capability */ #define UCT_IFACE_FLAG_INTER_NODE UCS_BIT(54) /**< Interface is inter-node capable */ #define UCT_IFACE_FLAG_DEVICE_EP UCS_BIT(55) /**< Interface supports device endpoint */ +#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /**< Interface require lkey for device operations */ /** * @} */ diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index f0b72e2688f..8423e301dd0 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -502,6 +502,7 @@ uct_rc_gdaki_iface_query(uct_iface_h tl_iface, uct_iface_attr_t *iface_attr) iface_attr->cap.flags = UCT_IFACE_FLAG_CONNECT_TO_EP | UCT_IFACE_FLAG_INTER_NODE | UCT_IFACE_FLAG_DEVICE_EP | + UCT_IFACE_FLAG_DEVICE_LKEY | UCT_IFACE_FLAG_ERRHANDLE_PEER_FAILURE; iface_attr->ep_addr_len = sizeof(uct_ib_uint24_t) * iface->num_channels; From 82faeb9955796ede3abc5a07e016d171c76f224b Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Mon, 23 Mar 2026 12:21:03 +0200 Subject: [PATCH 03/12] UCP/DEVICE: Add multi-lane support - 3 --- src/ucp/core/ucp_device.c | 201 ++++++++++++++++++++---------------- src/ucs/sys/math.c | 25 +++++ src/ucs/sys/math.h | 12 +++ test/gtest/ucs/test_math.cc | 15 +++ 4 files changed, 163 insertions(+), 90 deletions(-) diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 3deae80e040..2a96b4e384d 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -43,6 +43,12 @@ static ucs_spinlock_t ucp_device_handle_hash_lock; /* Size of temporary allocation for local sys_dev detection */ #define UCP_DEVICE_LOCAL_SYS_DEV_DETECT_SIZE 64 +enum { + UCP_DEVICE_TL_TYPE_LKEY, + UCP_DEVICE_TL_TYPE_NOLKEY, + UCP_DEVICE_TL_TYPE_NUM, +}; + void ucp_device_init(void) { @@ -138,17 +144,21 @@ ucp_device_detect_local_sys_dev(ucp_context_h context, return UCS_OK; } -static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, - ucp_tl_bitmap_t *tl_bitmap, - ucp_tl_bitmap_t *tl_bitmap_nolkey, - ucs_sys_device_t local_sys_dev) +static size_t +ucp_device_get_tl_bitmap(const ucp_worker_h worker, + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM], + ucs_sys_device_t local_sys_dev) { const ucp_worker_iface_t *wiface; ucp_rsc_index_t tl_id; + int tl_type; + size_t num, num_lanes = 0; /** TODO: Maybe cache results */ - UCS_STATIC_BITMAP_RESET_ALL(tl_bitmap); - UCS_STATIC_BITMAP_RESET_ALL(tl_bitmap_nolkey); + for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + UCS_STATIC_BITMAP_RESET_ALL(&tl_bitmap[tl_type]); + } + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &worker->context->tl_bitmap) { wiface = ucp_worker_iface(worker, tl_id); @@ -161,11 +171,23 @@ static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, } if (wiface->attr.cap.flags & UCT_IFACE_FLAG_DEVICE_LKEY) { - UCS_STATIC_BITMAP_SET(tl_bitmap, tl_id); + tl_type = UCP_DEVICE_TL_TYPE_LKEY; } else { - UCS_STATIC_BITMAP_SET(tl_bitmap_nolkey, tl_id); + tl_type = UCP_DEVICE_TL_TYPE_NOLKEY; + } + UCS_STATIC_BITMAP_SET(&tl_bitmap[tl_type], tl_id); + } + + for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + num = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]); + if (num_lanes == 0) { + num_lanes = num; + } else if (num > 0) { + num_lanes = ucs_lcm(num_lanes, num); } } + + return num_lanes; } static ucs_status_t ucp_device_local_mem_list_element_pack( @@ -243,15 +265,12 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( const ucp_worker_iface_t *wiface; ucp_device_local_mem_list_t *handle; uct_device_local_mem_list_elem_t *uct_element; - size_t i, num_lanes; + size_t i, j, num_lanes; ucs_status_t status; - ucp_tl_bitmap_t tl_bitmap; - ucp_tl_bitmap_t tl_bitmap_nolkey; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM]; ucp_rsc_index_t tl_id; - ucp_device_get_tl_bitmap(worker, &tl_bitmap, &tl_bitmap_nolkey, - local_sys_dev); - num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap); + num_lanes = ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); handle_size = (uct_elem_size * params->num_elements * num_lanes) + sizeof(*handle); handle = ucs_calloc(1, handle_size, "ucp_device_local_mem_list_t"); @@ -260,26 +279,32 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( return UCS_ERR_NO_MEMORY; } - /* Populate element specific parameters */ - ucp_element = params->elements; - uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); - for (i = 0; i < params->num_elements; i++) { - UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { - wiface = ucp_worker_iface(worker, tl_id); - status = ucp_device_local_mem_list_element_pack(worker, wiface, - ucp_element, - mem_type, - uct_element); - if (status != UCS_OK) { - ucs_error( - "failed to pack local mem list element for element=%zu", - i); - goto out; + if (UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY])) { + /* Populate element specific parameters */ + ucp_element = params->elements; + uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); + for (i = 0; i < params->num_elements; i++) { + for (j = 0; j < num_lanes;) { + UCS_STATIC_BITMAP_FOR_EACH_BIT( + tl_id, &tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY]) { + wiface = ucp_worker_iface(worker, tl_id); + status = ucp_device_local_mem_list_element_pack( + worker, wiface, ucp_element, mem_type, uct_element); + if (status != UCS_OK) { + ucs_error("failed to pack local mem list element for " + "element=%zu", + i); + goto out; + } + + uct_element = UCS_PTR_BYTE_OFFSET(uct_element, + uct_elem_size); + j++; + } } - - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); + ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, + params->element_size); } - ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); } handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1; @@ -382,21 +407,38 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params, return status; } -static int ucp_device_remote_mem_list_check_lanes(const ucp_ep_h ep, - ucp_tl_bitmap_t *tl_bitmap) - +static int ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id, + ucp_lane_index_t *lane_p) { ucp_ep_config_t *ep_config = ucp_ep_config(ep); ucp_lane_index_t lane; + + for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { + if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { + break; + } + } + + if (lane == ep_config->key.num_lanes) { + return 0; + } + + *lane_p = lane; + return 1; +} + +static int +ucp_device_ep_check_lanes(const ucp_ep_h ep, ucp_tl_bitmap_t *tl_bitmap) +{ + ucp_lane_index_t lane; ucp_rsc_index_t tl_id; + if (UCS_STATIC_BITMAP_POPCOUNT(*tl_bitmap) == 0) { + return 0; + } + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) { - for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { - if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { - break; - } - } - if (lane == ep_config->key.num_lanes) { + if (!ucp_device_ep_find_lane(ep, tl_id, &lane)) { return 0; } } @@ -410,7 +452,7 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack( { const ucp_ep_h ep = element->ep; const ucp_rkey_h rkey = element->rkey; - ucp_ep_config_t *ep_config = ucp_ep_config(ep); + ucp_ep_config_t *ep_config = ucp_ep_config(ep); uint8_t rkey_index; uct_rkey_t uct_rkey; uct_ep_h uct_ep; @@ -418,13 +460,7 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack( ucs_status_t status; ucp_lane_index_t lane; - for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { - if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { - break; - } - } - - if (lane == ep_config->key.num_lanes) { + if (!ucp_device_ep_find_lane(ep, tl_id, &lane)) { ucs_error("no lane found for ep=%p", ep); return UCS_ERR_NO_DEVICE; } @@ -489,26 +525,29 @@ static ucp_ep_h ucp_device_remote_mem_list_get_first_ep( 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, - uct_device_remote_mem_list_elem_t **uct_element_p) + uct_device_remote_mem_list_elem_t **uct_element_p, size_t num_lanes) { const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); ucp_rsc_index_t tl_id; ucs_status_t status; + size_t i; - UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) { - status = ucp_device_remote_mem_list_element_pack(ucp_element, tl_id, - *uct_element_p); - if (status != UCS_OK) { - return status; - } + 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, + *uct_element_p); + if (status != UCS_OK) { + return status; + } - *uct_element_p = UCS_PTR_BYTE_OFFSET(*uct_element_p, uct_elem_size); + *uct_element_p = UCS_PTR_BYTE_OFFSET(*uct_element_p, uct_elem_size); + i++; + } } return UCS_OK; } - static ucs_status_t ucp_device_remote_mem_list_create_handle( const ucp_device_mem_list_params_t *params, ucs_memory_type_t mem_type, uct_allocated_memory_t *mem) @@ -520,11 +559,10 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucp_device_remote_mem_list_t *handle; uct_device_remote_mem_list_elem_t *uct_element; ucs_sys_device_t local_sys_dev; - ucp_tl_bitmap_t tl_bitmap; - ucp_tl_bitmap_t tl_bitmap_nolkey; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM]; size_t i, num_lanes; ucs_status_t status; - int has_lkey, has_nolkey; + int tl_type; if (ep == NULL) { ucs_error("no ep found in remote mem list"); @@ -537,18 +575,12 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( return status; } - ucp_device_get_tl_bitmap(ep->worker, &tl_bitmap, &tl_bitmap_nolkey, - local_sys_dev); + num_lanes = ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev); ucp_element = params->elements; - has_lkey = ucp_device_remote_mem_list_check_lanes(ep, &tl_bitmap); - has_nolkey = ucp_device_remote_mem_list_check_lanes(ep, &tl_bitmap_nolkey); - - num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap) * has_lkey + - UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap_nolkey) * has_nolkey; - if (!num_lanes) { ucs_error("failed to pack uct memory element for first element"); + return UCS_ERR_INVALID_PARAM; } handle_size = sizeof(*handle) + @@ -565,30 +597,19 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size * num_lanes); } else { - ucs_assert(has_lkey == - ucp_device_remote_mem_list_check_lanes(ucp_element->ep, - &tl_bitmap)); - ucs_assert( - has_nolkey == - ucp_device_remote_mem_list_check_lanes(ucp_element->ep, - &tl_bitmap_nolkey)); - - if (has_lkey) { - status = ucp_device_remote_mem_list_fill(ucp_element, - &tl_bitmap, - &uct_element); - if (status != UCS_OK) { - goto out; + for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + if (ucp_device_ep_check_lanes(ucp_element->ep, + &tl_bitmap[tl_type])) { + break; } } - if (has_nolkey) { - status = ucp_device_remote_mem_list_fill(ucp_element, - &tl_bitmap_nolkey, - &uct_element); - if (status != UCS_OK) { - goto out; - } + ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_NUM); + status = ucp_device_remote_mem_list_fill(ucp_element, + &tl_bitmap[tl_type], + &uct_element, num_lanes); + if (status != UCS_OK) { + goto out; } } ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); diff --git a/src/ucs/sys/math.c b/src/ucs/sys/math.c index 421eaa239be..4a2f429f285 100644 --- a/src/ucs/sys/math.c +++ b/src/ucs/sys/math.c @@ -28,6 +28,31 @@ static uint64_t ucs_large_primes[] = { 9929050207ull, 9929050217ull, 9929050249ull, 9929050253ull }; +uint64_t ucs_gcd(uint64_t a, uint64_t b) +{ + uint64_t t; + + while (b != 0) { + t = a % b; + a = b; + b = t; + } + + return a; +} + +uint64_t ucs_lcm(uint64_t a, uint64_t b) +{ + uint64_t g; + + if ((a == 0) || (b == 0)) { + return 0; + } + + g = ucs_gcd(a, b); + return a / g * b; +} + uint64_t ucs_get_prime(unsigned index_val) { static const unsigned num_primes = sizeof(ucs_large_primes) / sizeof(ucs_large_primes[0]); diff --git a/src/ucs/sys/math.h b/src/ucs/sys/math.h index f23ec234063..9cf7aad89fb 100644 --- a/src/ucs/sys/math.h +++ b/src/ucs/sys/math.h @@ -177,6 +177,18 @@ static UCS_F_ALWAYS_INLINE size_t ucs_double_to_sizet(double value, size_t max) ((_submask) = ((_submask )+ ~(_mask)) & (_mask)) : 0) +/* + * Greatest common divisor of @a and @b. + */ +uint64_t ucs_gcd(uint64_t a, uint64_t b); + + +/* + * Least common multiple of @a and @b. + */ +uint64_t ucs_lcm(uint64_t a, uint64_t b); + + /* * Generate a large prime number */ diff --git a/test/gtest/ucs/test_math.cc b/test/gtest/ucs/test_math.cc index 45ed41656ab..01e6ca93326 100644 --- a/test/gtest/ucs/test_math.cc +++ b/test/gtest/ucs/test_math.cc @@ -306,3 +306,18 @@ UCS_TEST_F(test_math, double_to_sizet) { EXPECT_EQ(10, ucs_double_to_sizet(10.0, SIZE_MAX)); EXPECT_EQ(UCS_MBYTE, ucs_double_to_sizet(UCS_MBYTE, SIZE_MAX)); } + +UCS_TEST_F(test_math, gcd_lcm) { + EXPECT_EQ(0u, ucs_gcd(0, 0)); + EXPECT_EQ(7u, ucs_gcd(0, 7)); + EXPECT_EQ(7u, ucs_gcd(7, 0)); + EXPECT_EQ(6u, ucs_gcd(54, 24)); + EXPECT_EQ(6u, ucs_gcd(24, 54)); + EXPECT_EQ(1u, ucs_gcd(17, 13)); + EXPECT_EQ(17u, ucs_gcd(17, 17)); + + EXPECT_EQ(0u, ucs_lcm(0, 5)); + EXPECT_EQ(0u, ucs_lcm(7, 0)); + EXPECT_EQ(216u, ucs_lcm(54, 24)); + EXPECT_EQ(221u, ucs_lcm(17, 13)); +} From 1f3445e84c2711cdfda542b34e98aaa48edc9ffa Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Tue, 24 Mar 2026 05:23:04 +0200 Subject: [PATCH 04/12] UCP/DEVICE: Add multi-lane support - 4 --- src/ucp/core/ucp_device.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 2a96b4e384d..37fc244aea4 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -267,7 +267,7 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( uct_device_local_mem_list_elem_t *uct_element; size_t i, j, num_lanes; ucs_status_t status; - ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM]; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM] = {}; ucp_rsc_index_t tl_id; num_lanes = ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); @@ -559,7 +559,7 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucp_device_remote_mem_list_t *handle; uct_device_remote_mem_list_elem_t *uct_element; ucs_sys_device_t local_sys_dev; - ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM]; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM] = {}; size_t i, num_lanes; ucs_status_t status; int tl_type; From 71d7b9cf83eb798a5747259d253667bd3dddf02d Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Wed, 25 Mar 2026 12:27:11 +0200 Subject: [PATCH 05/12] UCP/DEVICE: Add multi-lane support - 5 --- src/ucp/wireup/address.c | 7 +++++-- src/uct/ib/mlx5/gdaki/gdaki.c | 1 + 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/ucp/wireup/address.c b/src/ucp/wireup/address.c index 71b0ea1d0fc..5f4b71977f6 100644 --- a/src/ucp/wireup/address.c +++ b/src/ucp/wireup/address.c @@ -315,6 +315,7 @@ ucp_address_get_device(ucp_context_h context, ucp_rsc_index_t rsc_index, dev = &devices[(*num_devices_p)++]; memset(dev, 0, sizeof(*dev)); + dev->num_paths = 1; out: return dev; } @@ -453,9 +454,11 @@ ucp_address_gather_devices(ucp_worker_h worker, const ucp_ep_config_key_t *key, return UCS_ERR_UNSUPPORTED; } - dev->rsc_index = rsc_index; + dev->rsc_index = rsc_index; UCS_STATIC_BITMAP_SET(&dev->tl_bitmap, rsc_index); - dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); + if (!(iface_attr->cap.flags & UCT_IFACE_FLAG_DEVICE_EP)) { + dev->num_paths = ucs_min(max_num_paths, iface_attr->dev_num_paths); + } } *devices_p = devices; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index 8423e301dd0..49b4dedce31 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -513,6 +513,7 @@ uct_rc_gdaki_iface_query(uct_iface_h tl_iface, uct_iface_attr_t *iface_attr) iface_attr->cap.put.max_zcopy = uct_ib_iface_port_attr(&iface->super.super.super)->max_msg_sz; iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev); + iface_attr->dev_num_paths = 1; return UCS_OK; } From c863ad2de4127c99d89fcedad047ff2605530c6a Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Wed, 15 Apr 2026 11:19:34 +0300 Subject: [PATCH 06/12] UCP/DEVICE: Add multi-lane support - 6 --- src/ucp/api/device/ucp_device_impl.h | 2 +- src/ucp/api/device/ucp_device_types.h | 2 +- src/ucp/core/ucp_device.c | 48 +++++++++++++-------------- src/ucp/wireup/address.c | 2 +- src/ucp/wireup/select.c | 2 +- src/ucs/sys/math.c | 2 +- src/ucs/sys/math.h | 2 +- test/gtest/ucp/cuda/test_kernels.cu | 2 +- test/gtest/ucp/cuda/test_kernels.h | 2 +- test/gtest/ucp/test_ucp_device.cc | 2 +- test/gtest/ucs/test_math.cc | 2 +- test/gtest/uct/test_device.cc | 2 +- 12 files changed, 35 insertions(+), 35 deletions(-) diff --git a/src/ucp/api/device/ucp_device_impl.h b/src/ucp/api/device/ucp_device_impl.h index 07ca5729a44..75a3c233ace 100644 --- a/src/ucp/api/device/ucp_device_impl.h +++ b/src/ucp/api/device/ucp_device_impl.h @@ -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. */ diff --git a/src/ucp/api/device/ucp_device_types.h b/src/ucp/api/device/ucp_device_types.h index 71933280eac..10c056018c1 100644 --- a/src/ucp/api/device/ucp_device_types.h +++ b/src/ucp/api/device/ucp_device_types.h @@ -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. */ diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 37fc244aea4..a9df0cfc648 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -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. */ @@ -44,9 +44,10 @@ static ucs_spinlock_t ucp_device_handle_hash_lock; #define UCP_DEVICE_LOCAL_SYS_DEV_DETECT_SIZE 64 enum { - UCP_DEVICE_TL_TYPE_LKEY, + UCP_DEVICE_TL_TYPE_FIRST, + UCP_DEVICE_TL_TYPE_LKEY = UCP_DEVICE_TL_TYPE_FIRST, UCP_DEVICE_TL_TYPE_NOLKEY, - UCP_DEVICE_TL_TYPE_NUM, + UCP_DEVICE_TL_TYPE_LAST, }; @@ -146,7 +147,7 @@ ucp_device_detect_local_sys_dev(ucp_context_h context, static size_t ucp_device_get_tl_bitmap(const ucp_worker_h worker, - ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM], + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST], ucs_sys_device_t local_sys_dev) { const ucp_worker_iface_t *wiface; @@ -155,7 +156,8 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker, size_t num, num_lanes = 0; /** TODO: Maybe cache results */ - for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; + tl_type < UCP_DEVICE_TL_TYPE_LAST; tl_type++) { UCS_STATIC_BITMAP_RESET_ALL(&tl_bitmap[tl_type]); } @@ -178,7 +180,8 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker, UCS_STATIC_BITMAP_SET(&tl_bitmap[tl_type], tl_id); } - for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; + tl_type < UCP_DEVICE_TL_TYPE_LAST; tl_type++) { num = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]); if (num_lanes == 0) { num_lanes = num; @@ -261,13 +264,13 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( UCP_DEVICE_MEM_LIST_PARAMS_FIELD, params, worker, WORKER, NULL); const size_t uct_elem_size = sizeof(uct_device_local_mem_list_elem_t); size_t handle_size = 0; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {}; const ucp_device_mem_list_elem_t *ucp_element; const ucp_worker_iface_t *wiface; ucp_device_local_mem_list_t *handle; uct_device_local_mem_list_elem_t *uct_element; size_t i, j, num_lanes; ucs_status_t status; - ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM] = {}; ucp_rsc_index_t tl_id; num_lanes = ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); @@ -407,24 +410,18 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params, return status; } -static int ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id, - ucp_lane_index_t *lane_p) +static ucp_lane_index_t ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id) { ucp_ep_config_t *ep_config = ucp_ep_config(ep); ucp_lane_index_t lane; for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { - break; + return lane; } } - if (lane == ep_config->key.num_lanes) { - return 0; - } - - *lane_p = lane; - return 1; + return UCP_NULL_LANE; } static int @@ -438,7 +435,8 @@ ucp_device_ep_check_lanes(const ucp_ep_h ep, ucp_tl_bitmap_t *tl_bitmap) } UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, tl_bitmap) { - if (!ucp_device_ep_find_lane(ep, tl_id, &lane)) { + lane = ucp_device_ep_find_lane(ep, tl_id); + if (lane == UCP_NULL_LANE) { return 0; } } @@ -460,7 +458,8 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack( ucs_status_t status; ucp_lane_index_t lane; - if (!ucp_device_ep_find_lane(ep, tl_id, &lane)) { + lane = ucp_device_ep_find_lane(ep, tl_id); + if (lane == UCP_NULL_LANE) { ucs_error("no lane found for ep=%p", ep); return UCS_ERR_NO_DEVICE; } @@ -524,8 +523,8 @@ static ucp_ep_h ucp_device_remote_mem_list_get_first_ep( 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, - uct_device_remote_mem_list_elem_t **uct_element_p, size_t num_lanes) + ucp_tl_bitmap_t *tl_bitmap, size_t num_lanes, + uct_device_remote_mem_list_elem_t **uct_element_p) { const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); ucp_rsc_index_t tl_id; @@ -555,11 +554,11 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( const ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params); const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); size_t handle_size = 0; + ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {}; const ucp_device_mem_list_elem_t *ucp_element; ucp_device_remote_mem_list_t *handle; uct_device_remote_mem_list_elem_t *uct_element; ucs_sys_device_t local_sys_dev; - ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_NUM] = {}; size_t i, num_lanes; ucs_status_t status; int tl_type; @@ -597,17 +596,18 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size * num_lanes); } else { - for (tl_type = 0; tl_type < UCP_DEVICE_TL_TYPE_NUM; tl_type++) { + for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; + tl_type < UCP_DEVICE_TL_TYPE_LAST; tl_type++) { if (ucp_device_ep_check_lanes(ucp_element->ep, &tl_bitmap[tl_type])) { break; } } - ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_NUM); + ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST); status = ucp_device_remote_mem_list_fill(ucp_element, &tl_bitmap[tl_type], - &uct_element, num_lanes); + num_lanes, &uct_element); if (status != UCS_OK) { goto out; } diff --git a/src/ucp/wireup/address.c b/src/ucp/wireup/address.c index 5f4b71977f6..70406b0766d 100644 --- a/src/ucp/wireup/address.c +++ b/src/ucp/wireup/address.c @@ -1,5 +1,5 @@ /** - * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2016. ALL RIGHTS RESERVED. + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ diff --git a/src/ucp/wireup/select.c b/src/ucp/wireup/select.c index 4d6d72b5bb5..126101f53f0 100644 --- a/src/ucp/wireup/select.c +++ b/src/ucp/wireup/select.c @@ -1,5 +1,5 @@ /** - * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2016. ALL RIGHTS RESERVED. + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * Copyright (C) Los Alamos National Security, LLC. 2019 ALL RIGHTS RESERVED. * * See file LICENSE for terms. diff --git a/src/ucs/sys/math.c b/src/ucs/sys/math.c index 4a2f429f285..4d61168abe1 100644 --- a/src/ucs/sys/math.c +++ b/src/ucs/sys/math.c @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2012. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ diff --git a/src/ucs/sys/math.h b/src/ucs/sys/math.h index 9cf7aad89fb..c135b36b92d 100644 --- a/src/ucs/sys/math.h +++ b/src/ucs/sys/math.h @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2014. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED. * * See file LICENSE for terms. diff --git a/test/gtest/ucp/cuda/test_kernels.cu b/test/gtest/ucp/cuda/test_kernels.cu index 3894c576190..18a5dc7ff9a 100644 --- a/test/gtest/ucp/cuda/test_kernels.cu +++ b/test/gtest/ucp/cuda/test_kernels.cu @@ -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. */ diff --git a/test/gtest/ucp/cuda/test_kernels.h b/test/gtest/ucp/cuda/test_kernels.h index 34ec4a972e0..b9d10f9a142 100644 --- a/test/gtest/ucp/cuda/test_kernels.h +++ b/test/gtest/ucp/cuda/test_kernels.h @@ -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. */ diff --git a/test/gtest/ucp/test_ucp_device.cc b/test/gtest/ucp/test_ucp_device.cc index 586a9d277ef..0e281e714f7 100644 --- a/test/gtest/ucp/test_ucp_device.cc +++ b/test/gtest/ucp/test_ucp_device.cc @@ -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. */ diff --git a/test/gtest/ucs/test_math.cc b/test/gtest/ucs/test_math.cc index 01e6ca93326..cf3097438e2 100644 --- a/test/gtest/ucs/test_math.cc +++ b/test/gtest/ucs/test_math.cc @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2012. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * Copyright (C) UT-Battelle, LLC. 2014. ALL RIGHTS RESERVED. * See file LICENSE for terms. */ diff --git a/test/gtest/uct/test_device.cc b/test/gtest/uct/test_device.cc index c565fa966d9..5a6024cfdd2 100644 --- a/test/gtest/uct/test_device.cc +++ b/test/gtest/uct/test_device.cc @@ -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. */ From e2b4ec0f49ed00b624db2f349435b890e712f5fc Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Fri, 24 Apr 2026 10:34:13 +0300 Subject: [PATCH 07/12] UCP/DEVICE: Add multi-lane support - 7 --- src/ucp/api/device/ucp_device_impl.h | 52 +++++---- src/ucp/core/ucp_device.c | 154 ++++++++++++------------- src/ucs/sys/math.c | 27 +---- src/ucs/sys/math.h | 14 +-- src/uct/api/device/uct_device_impl.h | 11 +- src/uct/api/device/uct_device_types.h | 12 +- src/uct/api/uct_def.h | 1 + src/uct/cuda/cuda_ipc/cuda_ipc_iface.c | 3 + src/uct/ib/mlx5/gdaki/gdaki.c | 2 +- src/uct/ib/mlx5/gdaki/gdaki.cuh | 5 +- test/gtest/ucs/test_math.cc | 17 +-- 11 files changed, 125 insertions(+), 173 deletions(-) diff --git a/src/ucp/api/device/ucp_device_impl.h b/src/ucp/api/device/ucp_device_impl.h index 75a3c233ace..325d31c40e4 100644 --- a/src/ucp/api/device/ucp_device_impl.h +++ b/src/ucp/api/device/ucp_device_impl.h @@ -105,19 +105,17 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep, }) -#define UCP_DEVICE_GET_LANE(_handle, _channel_id) \ - ({ \ - unsigned _lane = _channel_id % _handle->num_lanes; \ - _channel_id /= _handle->num_lanes; \ - _lane; \ - }) +#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; #define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \ - static_castmem_elements[0])*>( \ - UCS_PTR_BYTE_OFFSET(_handle->mem_elements, \ - ((_index * _handle->num_lanes) + _lane) * \ - sizeof(_handle->mem_elements[0]))) + static_castmem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \ + _handle->mem_elements, \ + (sizeof(_handle->mem_elements[0]) + \ + (sizeof(_handle->mem_elements[0].tl[0]) * _handle->num_lanes)) * \ + _index)) UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote( @@ -137,8 +135,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote( const auto dst_mem_element = UCP_DEVICE_GET_ELEM(dst_mem_list_h, dst_mem_list_index, lane); 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; @@ -152,7 +150,7 @@ ucp_device_prepare_send(const ucp_device_local_mem_list_h src_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_local_mem_list_elem_t *&src_uct_elem, + const uct_device_mem_element_t *&src_uct_elem, const uct_device_mem_element_t *&uct_elem, uct_device_completion_t *&comp) { @@ -170,9 +168,10 @@ ucp_device_prepare_send(const ucp_device_local_mem_list_h src_mem_list_h, return status; } - src_uct_elem = UCP_DEVICE_GET_ELEM(src_mem_list_h, src_mem_list_index, - lane); - address = src_uct_elem->addr; + const auto src_mem_elem = UCP_DEVICE_GET_ELEM(src_mem_list_h, + src_mem_list_index, lane); + src_uct_elem = src_mem_elem->tl + lane; + address = src_mem_elem->addr; return UCS_OK; } @@ -223,15 +222,18 @@ 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) { - const unsigned lane = UCP_DEVICE_GET_LANE(dst_mem_list_h, channel_id); + 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_element_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, lane, req, @@ -244,7 +246,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); } @@ -286,13 +288,16 @@ 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) { - const unsigned lane = UCP_DEVICE_GET_LANE(mem_list_h, channel_id); + unsigned uct_channel_id; + unsigned lane; uint64_t remote_address; const uct_device_mem_element_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, lane, req, device_ep, uct_elem, comp); @@ -302,8 +307,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc( 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); } @@ -336,8 +341,7 @@ ucp_device_get_ptr(const ucp_device_remote_mem_list_h mem_list_h, UCS_PTR_BYTE_OFFSET(mem_list_h->mem_elements, mem_list_index * elem_size)); - 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); } diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index a9df0cfc648..39e942325bd 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -145,7 +145,7 @@ ucp_device_detect_local_sys_dev(ucp_context_h context, return UCS_OK; } -static size_t +static void ucp_device_get_tl_bitmap(const ucp_worker_h worker, ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST], ucs_sys_device_t local_sys_dev) @@ -153,7 +153,6 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker, const ucp_worker_iface_t *wiface; ucp_rsc_index_t tl_id; int tl_type; - size_t num, num_lanes = 0; /** TODO: Maybe cache results */ for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; @@ -168,7 +167,8 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker, continue; } - if (wiface->attr.ctl_device != local_sys_dev) { + if ((wiface->attr.ctl_device != UCS_SYS_DEVICE_ID_UNKNOWN) && + (wiface->attr.ctl_device != local_sys_dev)) { continue; } @@ -179,28 +179,13 @@ ucp_device_get_tl_bitmap(const ucp_worker_h worker, } UCS_STATIC_BITMAP_SET(&tl_bitmap[tl_type], tl_id); } - - for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; - tl_type < UCP_DEVICE_TL_TYPE_LAST; tl_type++) { - num = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]); - if (num_lanes == 0) { - num_lanes = num; - } else if (num > 0) { - num_lanes = ucs_lcm(num_lanes, num); - } - } - - return num_lanes; } 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_local_mem_list_elem_t *mem_element) + const ucs_memory_type_t mem_type, uct_device_mem_element_t *mem_element) { - void *local_addr = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, element, - local_addr, LOCAL_ADDR, NULL); ucp_tl_resource_desc_t *resource; ucp_md_index_t md_index; ucp_mem_h memh; @@ -208,7 +193,6 @@ static ucs_status_t ucp_device_local_mem_list_element_pack( ucp_tl_md_t *ucp_md; ucs_status_t status; - mem_element->addr = local_addr; if (wiface == NULL) { return UCS_OK; } @@ -225,7 +209,7 @@ 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->uct_mem_element); + mem_element); if (status != UCS_OK) { ucs_error("failed to pack local mem element for memh=%p", memh); } @@ -262,52 +246,57 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( { const ucp_worker_h worker = UCS_PARAM_VALUE( UCP_DEVICE_MEM_LIST_PARAMS_FIELD, params, worker, WORKER, NULL); - const size_t uct_elem_size = sizeof(uct_device_local_mem_list_elem_t); - size_t handle_size = 0; + size_t uct_elem_size; + size_t handle_size; + int tl_type = UCP_DEVICE_TL_TYPE_LKEY; ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {}; const ucp_device_mem_list_elem_t *ucp_element; const ucp_worker_iface_t *wiface; ucp_device_local_mem_list_t *handle; uct_device_local_mem_list_elem_t *uct_element; - size_t i, j, num_lanes; + uct_device_mem_element_t *tl_element; + size_t i, num_lanes; ucs_status_t status; ucp_rsc_index_t tl_id; + void *local_addr; - num_lanes = ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); - handle_size = (uct_elem_size * params->num_elements * num_lanes) + - sizeof(*handle); - handle = ucs_calloc(1, handle_size, "ucp_device_local_mem_list_t"); + ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); + num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]); + + uct_elem_size = sizeof(uct_device_local_mem_list_elem_t) + + (sizeof(uct_device_mem_element_t) * num_lanes); + handle_size = (uct_elem_size * params->num_elements) + sizeof(*handle); + handle = ucs_calloc(1, handle_size, "ucp_device_local_mem_list_t"); if (handle == NULL) { ucs_error("failed to allocate ucp_device_local_mem_list_t"); return UCS_ERR_NO_MEMORY; } - if (UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY])) { - /* Populate element specific parameters */ - ucp_element = params->elements; - uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); - for (i = 0; i < params->num_elements; i++) { - for (j = 0; j < num_lanes;) { - UCS_STATIC_BITMAP_FOR_EACH_BIT( - tl_id, &tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY]) { - wiface = ucp_worker_iface(worker, tl_id); - status = ucp_device_local_mem_list_element_pack( - worker, wiface, ucp_element, mem_type, uct_element); - if (status != UCS_OK) { - ucs_error("failed to pack local mem list element for " - "element=%zu", - i); - goto out; - } - - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, - uct_elem_size); - j++; - } + /* Populate element specific parameters */ + ucp_element = params->elements; + uct_element = UCS_PTR_TYPE_OFFSET(handle, *handle); + for (i = 0; i < params->num_elements; i++) { + local_addr = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, + ucp_element, local_addr, LOCAL_ADDR, NULL); + uct_element->addr = local_addr; + tl_element = uct_element->tl; + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap[tl_type]) { + wiface = ucp_worker_iface(worker, tl_id); + status = ucp_device_local_mem_list_element_pack(worker, wiface, + ucp_element, + mem_type, + tl_element); + if (status != UCS_OK) { + ucs_error("failed to pack local mem list element for " + "element=%zu", + i); + goto out; } - ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, - params->element_size); + + tl_element = UCS_PTR_TYPE_OFFSET(tl_element, *tl_element); } + uct_element = (void*)tl_element; + ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); } handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1; @@ -412,10 +401,9 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params, static ucp_lane_index_t ucp_device_ep_find_lane(const ucp_ep_h ep, ucp_rsc_index_t tl_id) { - ucp_ep_config_t *ep_config = ucp_ep_config(ep); ucp_lane_index_t lane; - for (lane = 0; lane < ep_config->key.num_lanes; ++lane) { + for (lane = 0; lane < ucp_ep_num_lanes(ep); ++lane) { if (ucp_ep_get_rsc_index(ep, lane) == tl_id) { return lane; } @@ -446,7 +434,7 @@ 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_mem_list_elem_t *mem_element) + uct_device_remote_tl_list_elem_t *mem_element) { const ucp_ep_h ep = element->ep; const ucp_rkey_h rkey = element->rkey; @@ -476,10 +464,9 @@ 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); - mem_element->device_ep = device_ep; - mem_element->addr = element->remote_addr; - status = uct_md_mem_elem_pack(ucp_ep_md(ep, lane), NULL, uct_rkey, - &mem_element->uct_mem_element); + mem_element->ep = device_ep; + status = uct_md_mem_elem_pack(ucp_ep_md(ep, lane), NULL, uct_rkey, + &mem_element->uct); if (status != UCS_OK) { ucs_error("failed to pack uct memory element for lane=%u", lane); } @@ -521,25 +508,27 @@ 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_list_elem_t **uct_element_p) +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_list_elem_t *uct_element) { - const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); + uct_device_remote_tl_list_elem_t *tl_element; ucp_rsc_index_t tl_id; ucs_status_t status; size_t i; + uct_element->addr = ucp_element->remote_addr; + tl_element = uct_element->tl; 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, - *uct_element_p); + tl_element); if (status != UCS_OK) { return status; } - *uct_element_p = UCS_PTR_BYTE_OFFSET(*uct_element_p, uct_elem_size); + tl_element = UCS_PTR_TYPE_OFFSET(tl_element, *tl_element); i++; } } @@ -552,7 +541,7 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( uct_allocated_memory_t *mem) { const ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params); - const size_t uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t); + size_t uct_elem_size; size_t handle_size = 0; ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {}; const ucp_device_mem_list_elem_t *ucp_element; @@ -574,28 +563,32 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( return status; } - num_lanes = ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev); - ucp_element = params->elements; - + ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev); + num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY]); if (!num_lanes) { - ucs_error("failed to pack uct memory element for first element"); - return UCS_ERR_INVALID_PARAM; + if (!UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_NOLKEY])) { + ucs_error("failed to pack uct memory element for first element"); + return UCS_ERR_INVALID_PARAM; + } + + ucs_assert(UCS_STATIC_BITMAP_POPCOUNT( + tl_bitmap[UCP_DEVICE_TL_TYPE_NOLKEY]) == 1); + num_lanes = 1; } - handle_size = sizeof(*handle) + - (uct_elem_size * params->num_elements * num_lanes); + ucp_element = params->elements; + uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t) + + (sizeof(uct_device_remote_tl_list_elem_t) * num_lanes); + handle_size = sizeof(*handle) + (params->num_elements * uct_elem_size); handle = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t"); if (handle == NULL) { ucs_error("failed to allocate ucp_device_remote_mem_list_t"); return UCS_ERR_NO_MEMORY; } - uct_element = UCS_PTR_BYTE_OFFSET(handle, sizeof(*handle)); + uct_element = UCS_PTR_TYPE_OFFSET(handle, *handle); for (i = 0; i < params->num_elements; i++) { - if (UCP_DEVICE_MEM_ELEMENT_IS_GAP(ucp_element)) { - uct_element = UCS_PTR_BYTE_OFFSET(uct_element, - uct_elem_size * num_lanes); - } else { + if (!UCP_DEVICE_MEM_ELEMENT_IS_GAP(ucp_element)) { for (tl_type = UCP_DEVICE_TL_TYPE_FIRST; tl_type < UCP_DEVICE_TL_TYPE_LAST; tl_type++) { if (ucp_device_ep_check_lanes(ucp_element->ep, @@ -607,11 +600,12 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST); status = ucp_device_remote_mem_list_fill(ucp_element, &tl_bitmap[tl_type], - num_lanes, &uct_element); + num_lanes, uct_element); if (status != UCS_OK) { goto out; } } + uct_element = UCS_PTR_BYTE_OFFSET(uct_element, uct_elem_size); ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size); } diff --git a/src/ucs/sys/math.c b/src/ucs/sys/math.c index 4d61168abe1..421eaa239be 100644 --- a/src/ucs/sys/math.c +++ b/src/ucs/sys/math.c @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2012. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -28,31 +28,6 @@ static uint64_t ucs_large_primes[] = { 9929050207ull, 9929050217ull, 9929050249ull, 9929050253ull }; -uint64_t ucs_gcd(uint64_t a, uint64_t b) -{ - uint64_t t; - - while (b != 0) { - t = a % b; - a = b; - b = t; - } - - return a; -} - -uint64_t ucs_lcm(uint64_t a, uint64_t b) -{ - uint64_t g; - - if ((a == 0) || (b == 0)) { - return 0; - } - - g = ucs_gcd(a, b); - return a / g * b; -} - uint64_t ucs_get_prime(unsigned index_val) { static const unsigned num_primes = sizeof(ucs_large_primes) / sizeof(ucs_large_primes[0]); diff --git a/src/ucs/sys/math.h b/src/ucs/sys/math.h index c135b36b92d..f23ec234063 100644 --- a/src/ucs/sys/math.h +++ b/src/ucs/sys/math.h @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2014. ALL RIGHTS RESERVED. * Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED. * * See file LICENSE for terms. @@ -177,18 +177,6 @@ static UCS_F_ALWAYS_INLINE size_t ucs_double_to_sizet(double value, size_t max) ((_submask) = ((_submask )+ ~(_mask)) & (_mask)) : 0) -/* - * Greatest common divisor of @a and @b. - */ -uint64_t ucs_gcd(uint64_t a, uint64_t b); - - -/* - * Least common multiple of @a and @b. - */ -uint64_t ucs_lcm(uint64_t a, uint64_t b); - - /* * Generate a large prime number */ diff --git a/src/uct/api/device/uct_device_impl.h b/src/uct/api/device/uct_device_impl.h index 0a9ec3d0447..710152f145e 100644 --- a/src/uct/api/device/uct_device_impl.h +++ b/src/uct/api/device/uct_device_impl.h @@ -112,12 +112,11 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put_single( * @return Error code as defined by @ref ucs_status_t */ template -UCS_F_DEVICE ucs_status_t -uct_device_ep_put(uct_device_ep_h device_ep, - const uct_device_local_mem_list_elem_t *src_uct_elem, - const uct_device_mem_element_t *mem_elem, const void *address, - uint64_t remote_address, size_t length, unsigned channel_id, - uint64_t flags, uct_device_completion_t *comp) +UCS_F_DEVICE ucs_status_t uct_device_ep_put( + uct_device_ep_h device_ep, const uct_device_mem_element_t *src_uct_elem, + const uct_device_mem_element_t *mem_elem, const void *address, + uint64_t remote_address, size_t length, unsigned channel_id, + uint64_t flags, uct_device_completion_t *comp) { #if UCT_RC_MLX5_GDA_SUPPORTED if (device_ep->uct_tl_id == UCT_DEVICE_TL_RC_MLX5_GDA) { diff --git a/src/uct/api/device/uct_device_types.h b/src/uct/api/device/uct_device_types.h index c2104bd07fc..3d863f227a3 100644 --- a/src/uct/api/device/uct_device_types.h +++ b/src/uct/api/device/uct_device_types.h @@ -80,14 +80,18 @@ union uct_device_mem_element { struct uct_device_local_mem_list_elem { void *addr; - uct_device_mem_element_t uct_mem_element; + uct_device_mem_element_t tl[0]; }; +struct uct_device_remote_tl_list_elem { + uct_device_ep_h ep; + uct_device_mem_element_t uct; +}; + struct uct_device_remote_mem_list_elem { - uct_device_ep_h device_ep; - uint64_t addr; - uct_device_mem_element_t uct_mem_element; + uint64_t addr; + uct_device_remote_tl_list_elem_t tl[0]; }; #endif diff --git a/src/uct/api/uct_def.h b/src/uct/api/uct_def.h index 688d08689fc..758ba08d214 100644 --- a/src/uct/api/uct_def.h +++ b/src/uct/api/uct_def.h @@ -113,6 +113,7 @@ typedef void* uct_conn_request_h; typedef struct uct_device_ep *uct_device_ep_h; typedef union uct_device_mem_element uct_device_mem_element_t; typedef struct uct_device_local_mem_list_elem uct_device_local_mem_list_elem_t; +typedef struct uct_device_remote_tl_list_elem uct_device_remote_tl_list_elem_t; typedef struct uct_device_remote_mem_list_elem uct_device_remote_mem_list_elem_t; /** diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c index 2a960409925..9b5e909cd1c 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c @@ -268,6 +268,9 @@ static ucs_status_t uct_cuda_ipc_iface_query(uct_iface_h tl_iface, UCT_IFACE_FLAG_GET_ZCOPY | UCT_IFACE_FLAG_PUT_ZCOPY | UCT_IFACE_FLAG_DEVICE_EP; + + iface_attr->ctl_device = UCS_SYS_DEVICE_ID_UNKNOWN; + if (md->enable_mnnvl) { iface_attr->cap.flags |= UCT_IFACE_FLAG_INTER_NODE; } diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index 49b4dedce31..da15eeb9ee8 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -512,7 +512,7 @@ uct_rc_gdaki_iface_query(uct_iface_h tl_iface, uct_iface_attr_t *iface_attr) iface_attr->cap.put.min_zcopy = 0; iface_attr->cap.put.max_zcopy = uct_ib_iface_port_attr(&iface->super.super.super)->max_msg_sz; - iface_attr->ctl_device = uct_cuda_get_cuda_device(iface->cuda_dev); + iface_attr->ctl_device = uct_cuda_get_sys_dev(iface->cuda_dev); iface_attr->dev_num_paths = 1; return UCS_OK; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.cuh b/src/uct/ib/mlx5/gdaki/gdaki.cuh index 949e49a9392..f6991012b58 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.cuh +++ b/src/uct/ib/mlx5/gdaki/gdaki.cuh @@ -310,8 +310,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_single( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( - uct_device_ep_h tl_ep, - const uct_device_local_mem_list_elem_t *src_uct_elem, + uct_device_ep_h tl_ep, const uct_device_mem_element_t *src_uct_elem, const uct_device_mem_element_t *tl_mem_elem, const void *address, uint64_t remote_address, size_t length, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) @@ -321,7 +320,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( tl_mem_elem); auto local_mem_elem = reinterpret_cast( - &src_uct_elem->uct_mem_element); + src_uct_elem); auto cid = channel_id & ep->channel_mask; return uct_rc_mlx5_gda_ep_single(ep, tl_mem_elem, address, diff --git a/test/gtest/ucs/test_math.cc b/test/gtest/ucs/test_math.cc index cf3097438e2..45ed41656ab 100644 --- a/test/gtest/ucs/test_math.cc +++ b/test/gtest/ucs/test_math.cc @@ -1,5 +1,5 @@ /** -* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. +* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2012. ALL RIGHTS RESERVED. * Copyright (C) UT-Battelle, LLC. 2014. ALL RIGHTS RESERVED. * See file LICENSE for terms. */ @@ -306,18 +306,3 @@ UCS_TEST_F(test_math, double_to_sizet) { EXPECT_EQ(10, ucs_double_to_sizet(10.0, SIZE_MAX)); EXPECT_EQ(UCS_MBYTE, ucs_double_to_sizet(UCS_MBYTE, SIZE_MAX)); } - -UCS_TEST_F(test_math, gcd_lcm) { - EXPECT_EQ(0u, ucs_gcd(0, 0)); - EXPECT_EQ(7u, ucs_gcd(0, 7)); - EXPECT_EQ(7u, ucs_gcd(7, 0)); - EXPECT_EQ(6u, ucs_gcd(54, 24)); - EXPECT_EQ(6u, ucs_gcd(24, 54)); - EXPECT_EQ(1u, ucs_gcd(17, 13)); - EXPECT_EQ(17u, ucs_gcd(17, 17)); - - EXPECT_EQ(0u, ucs_lcm(0, 5)); - EXPECT_EQ(0u, ucs_lcm(7, 0)); - EXPECT_EQ(216u, ucs_lcm(54, 24)); - EXPECT_EQ(221u, ucs_lcm(17, 13)); -} From ed5762bfb8caa0ab044544eec663d946fe4294cf Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Fri, 24 Apr 2026 17:08:21 +0300 Subject: [PATCH 08/12] UCP/DEVICE: Add multi-lane support - 8 --- src/ucp/api/device/ucp_device_impl.h | 31 ++++++------ src/ucp/api/device/ucp_device_types.h | 23 ++++++--- src/ucp/core/ucp_device.c | 22 ++++----- src/uct/api/device/uct_device_impl.h | 14 +++--- src/uct/api/device/uct_device_types.h | 20 ++++---- src/uct/api/uct_def.h | 8 +-- src/uct/api/v2/uct_v2.h | 2 +- src/uct/base/uct_md.c | 2 +- src/uct/base/uct_md.h | 2 +- src/uct/cuda/cuda_ipc/cuda_ipc.cuh | 17 +++---- src/uct/cuda/cuda_ipc/cuda_ipc_md.c | 2 +- src/uct/ib/mlx5/dv/ib_mlx5dv_md.c | 2 +- src/uct/ib/mlx5/gdaki/gdaki.cuh | 14 +++--- test/gtest/uct/cuda/test_cuda_ipc_device.cc | 28 +++++------ test/gtest/uct/cuda/test_kernels.cu | 20 ++++---- test/gtest/uct/cuda/test_kernels.h | 8 +-- test/gtest/uct/cuda/test_kernels_uct.cu | 44 ++++++++--------- test/gtest/uct/cuda/test_kernels_uct.h | 14 +++--- test/gtest/uct/test_device.cc | 54 ++++++++++----------- 19 files changed, 160 insertions(+), 167 deletions(-) diff --git a/src/ucp/api/device/ucp_device_impl.h b/src/ucp/api/device/ucp_device_impl.h index 325d31c40e4..32ea314a376 100644 --- a/src/ucp/api/device/ucp_device_impl.h +++ b/src/ucp/api/device/ucp_device_impl.h @@ -122,8 +122,7 @@ 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 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) { ucs_status_t status; @@ -143,16 +142,14 @@ UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote( } -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_element_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) { ucs_status_t status; @@ -225,8 +222,8 @@ ucp_device_put(const ucp_device_local_mem_list_h src_mem_list_h, unsigned uct_channel_id; unsigned lane; const void *address; - const uct_device_mem_element_t *uct_elem; - const uct_device_mem_element_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; @@ -291,7 +288,7 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc( 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; @@ -329,7 +326,7 @@ 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); + const size_t elem_size = sizeof(uct_device_remote_mem_elem_t); ucs_status_t status; status = ucp_device_check_params(mem_list_h, mem_list_index); @@ -337,7 +334,7 @@ ucp_device_get_ptr(const ucp_device_remote_mem_list_h mem_list_h, return status; } - const auto mem_element = static_cast( + const auto mem_element = static_cast( UCS_PTR_BYTE_OFFSET(mem_list_h->mem_elements, mem_list_index * elem_size)); diff --git a/src/ucp/api/device/ucp_device_types.h b/src/ucp/api/device/ucp_device_types.h index 10c056018c1..121607345cb 100644 --- a/src/ucp/api/device/ucp_device_types.h +++ b/src/ucp/api/device/ucp_device_types.h @@ -30,19 +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; - uint16_t num_lanes; + /** + * Number of lanes in each memory descriptor. + */ + uint16_t num_lanes; /** * 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; @@ -61,15 +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; - uint16_t num_lanes; /** * 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 */ diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 39e942325bd..839b5a7a4b9 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -184,7 +184,7 @@ 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_element_t *mem_element) + const ucs_memory_type_t mem_type, uct_device_mem_elem_t *mem_element) { ucp_tl_resource_desc_t *resource; ucp_md_index_t md_index; @@ -253,8 +253,8 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( const ucp_device_mem_list_elem_t *ucp_element; const ucp_worker_iface_t *wiface; ucp_device_local_mem_list_t *handle; - uct_device_local_mem_list_elem_t *uct_element; - uct_device_mem_element_t *tl_element; + uct_device_local_mem_elem_t *uct_element; + uct_device_mem_elem_t *tl_element; size_t i, num_lanes; ucs_status_t status; ucp_rsc_index_t tl_id; @@ -263,8 +263,8 @@ static ucs_status_t ucp_device_local_mem_list_create_handle( ucp_device_get_tl_bitmap(worker, tl_bitmap, local_sys_dev); num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[tl_type]); - uct_elem_size = sizeof(uct_device_local_mem_list_elem_t) + - (sizeof(uct_device_mem_element_t) * num_lanes); + uct_elem_size = sizeof(uct_device_local_mem_elem_t) + + (sizeof(uct_device_mem_elem_t) * num_lanes); handle_size = (uct_elem_size * params->num_elements) + sizeof(*handle); handle = ucs_calloc(1, handle_size, "ucp_device_local_mem_list_t"); if (handle == NULL) { @@ -434,7 +434,7 @@ 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_list_elem_t *mem_element) + uct_device_remote_tl_elem_t *mem_element) { const ucp_ep_h ep = element->ep; const ucp_rkey_h rkey = element->rkey; @@ -511,9 +511,9 @@ static ucp_ep_h ucp_device_remote_mem_list_get_first_ep( 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_list_elem_t *uct_element) + uct_device_remote_mem_elem_t *uct_element) { - uct_device_remote_tl_list_elem_t *tl_element; + uct_device_remote_tl_elem_t *tl_element; ucp_rsc_index_t tl_id; ucs_status_t status; size_t i; @@ -546,7 +546,7 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {}; const ucp_device_mem_list_elem_t *ucp_element; ucp_device_remote_mem_list_t *handle; - uct_device_remote_mem_list_elem_t *uct_element; + uct_device_remote_mem_elem_t *uct_element; ucs_sys_device_t local_sys_dev; size_t i, num_lanes; ucs_status_t status; @@ -577,8 +577,8 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( } ucp_element = params->elements; - uct_elem_size = sizeof(uct_device_remote_mem_list_elem_t) + - (sizeof(uct_device_remote_tl_list_elem_t) * num_lanes); + uct_elem_size = sizeof(uct_device_remote_mem_elem_t) + + (sizeof(uct_device_remote_tl_elem_t) * num_lanes); handle_size = sizeof(*handle) + (params->num_elements * uct_elem_size); handle = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t"); if (handle == NULL) { diff --git a/src/uct/api/device/uct_device_impl.h b/src/uct/api/device/uct_device_impl.h index 710152f145e..2949015a2a1 100644 --- a/src/uct/api/device/uct_device_impl.h +++ b/src/uct/api/device/uct_device_impl.h @@ -59,7 +59,7 @@ union uct_device_completion { */ template UCS_F_DEVICE ucs_status_t uct_device_ep_put_single( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, const void *address, uint64_t remote_address, size_t length, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -113,8 +113,8 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put_single( */ template UCS_F_DEVICE ucs_status_t uct_device_ep_put( - uct_device_ep_h device_ep, const uct_device_mem_element_t *src_uct_elem, - const uct_device_mem_element_t *mem_elem, const void *address, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *src_uct_elem, + const uct_device_mem_elem_t *mem_elem, const void *address, uint64_t remote_address, size_t length, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -165,7 +165,7 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put( */ template UCS_F_DEVICE ucs_status_t uct_device_ep_atomic_add( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, uint64_t inc_value, uint64_t remote_address, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -200,7 +200,7 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_atomic_add( * @return Error code as defined by @ref ucs_status_t */ UCS_F_DEVICE ucs_status_t uct_device_ep_get_ptr( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, uint64_t address, void **addr_p) { if (device_ep->uct_tl_id != UCT_DEVICE_TL_CUDA_IPC) { @@ -256,7 +256,7 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_get_ptr( */ template UCS_F_DEVICE ucs_status_t uct_device_ep_put_multi( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, uint64_t counter_remote_address, @@ -341,7 +341,7 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put_multi( */ template UCS_F_DEVICE ucs_status_t uct_device_ep_put_multi_partial( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *local_offsets, const size_t *remote_offsets, diff --git a/src/uct/api/device/uct_device_types.h b/src/uct/api/device/uct_device_types.h index 3d863f227a3..ec10a2370b4 100644 --- a/src/uct/api/device/uct_device_types.h +++ b/src/uct/api/device/uct_device_types.h @@ -72,26 +72,26 @@ typedef union uct_device_completion uct_device_completion_t; /* Union structure of all device memory elements types */ -union uct_device_mem_element { +union uct_device_mem_elem { uct_ib_md_device_mem_element_t ib_md_mem_element; uct_cuda_ipc_md_device_mem_element_t cuda_ipc_md_mem_element; }; -struct uct_device_local_mem_list_elem { - void *addr; - uct_device_mem_element_t tl[0]; +struct uct_device_local_mem_elem { + void *addr; + uct_device_mem_elem_t tl[0]; }; -struct uct_device_remote_tl_list_elem { - uct_device_ep_h ep; - uct_device_mem_element_t uct; +struct uct_device_remote_tl_elem { + uct_device_ep_h ep; + uct_device_mem_elem_t uct; }; -struct uct_device_remote_mem_list_elem { - uint64_t addr; - uct_device_remote_tl_list_elem_t tl[0]; +struct uct_device_remote_mem_elem { + uint64_t addr; + uct_device_remote_tl_elem_t tl[0]; }; #endif diff --git a/src/uct/api/uct_def.h b/src/uct/api/uct_def.h index 758ba08d214..0895ce7469d 100644 --- a/src/uct/api/uct_def.h +++ b/src/uct/api/uct_def.h @@ -111,10 +111,10 @@ typedef uint64_t uct_tag_t; /* tag type - 64 bit */ typedef int uct_worker_cb_id_t; typedef void* uct_conn_request_h; typedef struct uct_device_ep *uct_device_ep_h; -typedef union uct_device_mem_element uct_device_mem_element_t; -typedef struct uct_device_local_mem_list_elem uct_device_local_mem_list_elem_t; -typedef struct uct_device_remote_tl_list_elem uct_device_remote_tl_list_elem_t; -typedef struct uct_device_remote_mem_list_elem uct_device_remote_mem_list_elem_t; +typedef union uct_device_mem_elem uct_device_mem_elem_t; +typedef struct uct_device_local_mem_elem uct_device_local_mem_elem_t; +typedef struct uct_device_remote_tl_elem uct_device_remote_tl_elem_t; +typedef struct uct_device_remote_mem_elem uct_device_remote_mem_elem_t; /** * @} diff --git a/src/uct/api/v2/uct_v2.h b/src/uct/api/v2/uct_v2.h index 254e825315c..c42a7cd63ff 100644 --- a/src/uct/api/v2/uct_v2.h +++ b/src/uct/api/v2/uct_v2.h @@ -1260,7 +1260,7 @@ ucs_status_t uct_rkey_unpack_v2(uct_component_h component, * @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_element_t *mem_elem); + uct_device_mem_elem_t *mem_elem); END_C_DECLS diff --git a/src/uct/base/uct_md.c b/src/uct/base/uct_md.c index 62231f2d52a..28db7881908 100644 --- a/src/uct/base/uct_md.c +++ b/src/uct/base/uct_md.c @@ -658,7 +658,7 @@ ucs_status_t uct_md_dummy_mem_dereg(uct_md_h uct_md, } ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey, - uct_device_mem_element_t *mem_elem) + uct_device_mem_elem_t *mem_elem) { return md->ops->mem_elem_pack(md, memh, rkey, mem_elem); } diff --git a/src/uct/base/uct_md.h b/src/uct/base/uct_md.h index 8219ac652d7..31909de3c22 100644 --- a/src/uct/base/uct_md.h +++ b/src/uct/base/uct_md.h @@ -128,7 +128,7 @@ typedef ucs_status_t (*uct_md_detect_memory_type_func_t)(uct_md_h md, typedef ucs_status_t (*uct_md_mem_elem_pack_func_t)( uct_md_h md, uct_mem_h memh, uct_rkey_t rkey, - uct_device_mem_element_t *mem_elem_p); + uct_device_mem_elem_t *mem_elem_p); /** * Memory domain operations diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc.cuh b/src/uct/cuda/cuda_ipc/cuda_ipc.cuh index 5fd29725897..60abf07f48f 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc.cuh +++ b/src/uct/cuda/cuda_ipc/cuda_ipc.cuh @@ -292,7 +292,7 @@ void uct_cuda_ipc_copy_level(void *dst, const void *src, template UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_single( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, const void *address, uint64_t remote_address, size_t length, uint64_t flags, uct_device_completion_t *comp) { @@ -309,7 +309,7 @@ UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_single( template UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_multi( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, uint64_t counter_remote_address, @@ -351,7 +351,7 @@ UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_multi( template UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_multi_partial( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *local_offsets, const size_t *remote_offsets, @@ -396,11 +396,10 @@ UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_put_multi_partial( } template -UCS_F_DEVICE ucs_status_t -uct_cuda_ipc_ep_atomic_add(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, - uint64_t inc_value, uint64_t remote_address, - uint64_t flags, uct_device_completion_t *comp) +UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_atomic_add( + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, + uint64_t inc_value, uint64_t remote_address, uint64_t flags, + uct_device_completion_t *comp) { auto cuda_ipc_mem_element = reinterpret_cast( @@ -420,7 +419,7 @@ uct_cuda_ipc_ep_atomic_add(uct_device_ep_h device_ep, } UCS_F_DEVICE ucs_status_t uct_cuda_ipc_ep_get_ptr( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, uint64_t remote_address, void **addr_p) { auto cuda_ipc_mem_element = diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index 540793b7591..369be410dc3 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -559,7 +559,7 @@ static void uct_cuda_ipc_md_close(uct_md_h md) static ucs_status_t uct_cuda_ipc_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey, - uct_device_mem_element_t *mem_elem_p) + uct_device_mem_elem_t *mem_elem_p) { uct_cuda_ipc_unpacked_rkey_t *key = (uct_cuda_ipc_unpacked_rkey_t*)rkey; uct_cuda_ipc_md_device_mem_element_t *cuda_ipc_md_mem_element = diff --git a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c index 7433e4ca22f..b2240d8c505 100644 --- a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c +++ b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c @@ -3165,7 +3165,7 @@ uct_ib_mlx5_devx_md_open(struct ibv_device *ibv_device, static ucs_status_t uct_ib_md_mlx5_devx_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey, - uct_device_mem_element_t *mem_elem_p) + uct_device_mem_elem_t *mem_elem_p) { uct_ib_md_device_mem_element_t *mem_elem = (uct_ib_md_device_mem_element_t*) mem_elem_p; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.cuh b/src/uct/ib/mlx5/gdaki/gdaki.cuh index f6991012b58..c5fa8b44300 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.cuh +++ b/src/uct/ib/mlx5/gdaki/gdaki.cuh @@ -263,7 +263,7 @@ uct_rc_mlx5_gda_fc(const uct_rc_gdaki_dev_ep_t *ep, uint16_t wqe_idx) template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_single( - uct_rc_gdaki_dev_ep_t *ep, const uct_device_mem_element_t *tl_mem_elem, + uct_rc_gdaki_dev_ep_t *ep, const uct_device_mem_elem_t *tl_mem_elem, const void *address, uint32_t lkey, uint64_t remote_address, uint32_t rkey, size_t length, unsigned cid, uint64_t flags, uct_device_completion_t *tl_comp, uint32_t opcode, bool is_atomic, @@ -310,8 +310,8 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_single( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( - uct_device_ep_h tl_ep, const uct_device_mem_element_t *src_uct_elem, - const uct_device_mem_element_t *tl_mem_elem, const void *address, + 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, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -332,7 +332,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( - uct_device_ep_h tl_ep, const uct_device_mem_element_t *tl_mem_elem, + uct_device_ep_h tl_ep, const uct_device_mem_elem_t *tl_mem_elem, const void *address, uint64_t remote_address, size_t length, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -350,7 +350,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_single( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_atomic_add( - uct_device_ep_h tl_ep, const uct_device_mem_element_t *tl_mem_elem, + uct_device_ep_h tl_ep, const uct_device_mem_elem_t *tl_mem_elem, uint64_t value, uint64_t remote_address, unsigned channel_id, uint64_t flags, uct_device_completion_t *comp) { @@ -368,7 +368,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_atomic_add( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi( - uct_device_ep_h tl_ep, const uct_device_mem_element_t *tl_mem_list, + uct_device_ep_h tl_ep, const uct_device_mem_elem_t *tl_mem_list, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, uint64_t counter_remote_address, @@ -460,7 +460,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi( template UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi_partial( - uct_device_ep_h tl_ep, const uct_device_mem_element_t *tl_mem_list, + uct_device_ep_h tl_ep, const uct_device_mem_elem_t *tl_mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *local_offsets, const size_t *remote_offsets, diff --git a/test/gtest/uct/cuda/test_cuda_ipc_device.cc b/test/gtest/uct/cuda/test_cuda_ipc_device.cc index e09a0e9a006..74f75d064a9 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_device.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_device.cc @@ -189,7 +189,7 @@ UCS_TEST_P(test_cuda_ipc_rma, get_mem_elem_pack) mapped_buffer sendbuf(length, SEED1, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); mapped_buffer recvbuf(length, SEED2, *m_receiver, 0, UCS_MEMORY_TYPE_CUDA); - uct_device_mem_element_t mem_elem_host; + uct_device_mem_elem_t mem_elem_host; EXPECT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(), recvbuf.rkey(), &mem_elem_host)); } @@ -212,7 +212,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_zcopy_device) size_t length = base_length + offset; unsigned num_blocks = get_num_blocks(); uct_device_ep_h device_ep; - uct_device_mem_element_t *mem_elem; + uct_device_mem_elem_t *mem_elem; void *send_buf; if (device_level == UCS_DEVICE_LEVEL_GRID) { @@ -228,7 +228,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_zcopy_device) ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &device_ep)); - uct_device_mem_element_t mem_elem_host; + uct_device_mem_elem_t mem_elem_host; ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&mem_elem, mem_elem_size)); ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(), recvbuf.rkey(), &mem_elem_host)); @@ -257,7 +257,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_device) size_t length = iovcnt * (base_length + offset); uint64_t signal_val = 4; uct_device_ep_h device_ep; - uct_device_mem_element_t *mem_elem; + uct_device_mem_elem_t *mem_elem; uint64_t *remote_addresses_dev, remote_addresses[iovcnt]; size_t *lengths_dev, lengths[iovcnt]; void **addresses_dev, *addresses[iovcnt]; @@ -278,8 +278,8 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_device) size_t total_mem_elem_size = mem_elem_size * (iovcnt + 1); uct_cuda_ipc_md_device_mem_element_t mem_elem_host_arr[iovcnt + 1]; - uct_device_mem_element_t *mem_elem_host = - (uct_device_mem_element_t*)mem_elem_host_arr; + uct_device_mem_elem_t *mem_elem_host = (uct_device_mem_elem_t*) + mem_elem_host_arr; ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&mem_elem, total_mem_elem_size)); ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&remote_addresses_dev, @@ -297,13 +297,13 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_device) lengths[i] = base_length; ASSERT_UCS_OK(uct_md_mem_elem_pack( m_sender->md(), sendbuf.memh(), recvbuf.rkey(), - (uct_device_mem_element_t*) + (uct_device_mem_elem_t*) UCS_PTR_BYTE_OFFSET(mem_elem_host, mem_elem_size * i))); } ASSERT_UCS_OK( uct_md_mem_elem_pack(m_sender->md(), nullptr, signal.rkey(), - (uct_device_mem_element_t*)UCS_PTR_BYTE_OFFSET( + (uct_device_mem_elem_t*)UCS_PTR_BYTE_OFFSET( mem_elem_host, mem_elem_size * iovcnt))); @@ -354,7 +354,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_partial_device) int counter_index = 1; std::vector offsets(iovcnt, 0); uct_device_ep_h device_ep; - uct_device_mem_element_t *mem_elements; + uct_device_mem_elem_t *mem_elements; uint64_t *remote_addresses_dev, remote_addresses[iovcnt + 1]; size_t *lengths_dev, lengths[iovcnt]; void **addresses_dev, *addresses[iovcnt + 1]; @@ -377,8 +377,8 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_partial_device) size_t total_mem_elem_size = mem_elem_size * (iovcnt + 1); uct_cuda_ipc_md_device_mem_element_t mem_elements_host_arr[iovcnt + 1]; - uct_device_mem_element_t *mem_elements_host = - (uct_device_mem_element_t*)mem_elements_host_arr; + uct_device_mem_elem_t *mem_elements_host = (uct_device_mem_elem_t*) + mem_elements_host_arr; ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&mem_elements, total_mem_elem_size)); ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&remote_addresses_dev, @@ -393,7 +393,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_multi_partial_device) /* Fill indices and pack PUT entries */ int idx = 0; for (int i = 0; i < iovcnt + 1; i++) { - uct_device_mem_element_t *mem_elem = (uct_device_mem_element_t*) + uct_device_mem_elem_t *mem_elem = (uct_device_mem_elem_t*) UCS_PTR_BYTE_OFFSET(mem_elements_host, mem_elem_size * i); if (i == counter_index) { ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), nullptr, @@ -465,7 +465,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, atomic_add_device) unsigned num_threads = get_num_threads(); unsigned num_blocks = get_num_blocks(); uct_device_ep_h device_ep; - uct_device_mem_element_t *mem_elem; + uct_device_mem_elem_t *mem_elem; if (device_level == UCS_DEVICE_LEVEL_GRID) { GTEST_SKIP() << "Grid level is not supported"; @@ -478,7 +478,7 @@ UCS_TEST_P(test_cuda_ipc_rma_device, atomic_add_device) mapped_buffer signal(sizeof(uint64_t), 0, *m_receiver, 0, UCS_MEMORY_TYPE_CUDA); ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &device_ep)); - uct_device_mem_element_t mem_elem_host; + uct_device_mem_elem_t mem_elem_host; ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&mem_elem, mem_elem_size)); ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), nullptr, signal.rkey(), &mem_elem_host)); diff --git a/test/gtest/uct/cuda/test_kernels.cu b/test/gtest/uct/cuda/test_kernels.cu index 008efdc91f7..046f578b75f 100644 --- a/test/gtest/uct/cuda/test_kernels.cu +++ b/test/gtest/uct/cuda/test_kernels.cu @@ -13,7 +13,7 @@ namespace ucx_cuda { static __global__ void -uct_put_single_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_elem, +uct_put_single_kernel(uct_device_ep_h ep, uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length, ucs_status_t *status_p) { @@ -37,7 +37,7 @@ uct_put_single_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_elem, * Basic single element put operation. */ ucs_status_t launch_uct_put_single(uct_device_ep_h ep, - uct_device_mem_element_t *mem_elem, + uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length) { device_result_ptr status = UCS_ERR_NOT_IMPLEMENTED; @@ -49,7 +49,7 @@ ucs_status_t launch_uct_put_single(uct_device_ep_h ep, } static __global__ void -uct_atomic_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_elem, +uct_atomic_kernel(uct_device_ep_h ep, uct_device_mem_elem_t *mem_elem, uint64_t rva, uint64_t add, ucs_status_t *status_p) { uct_device_completion_t comp; @@ -72,7 +72,7 @@ uct_atomic_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_elem, * Atomic operation. */ ucs_status_t launch_uct_atomic(uct_device_ep_h ep, - uct_device_mem_element_t *mem_elem, uint64_t rva, + uct_device_mem_elem_t *mem_elem, uint64_t rva, uint64_t add) { device_result_ptr status = UCS_ERR_NOT_IMPLEMENTED; @@ -84,7 +84,7 @@ ucs_status_t launch_uct_atomic(uct_device_ep_h ep, template static __global__ void -uct_put_multi_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +uct_put_multi_kernel(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length, ucs_status_t *status_p) { @@ -123,7 +123,7 @@ uct_put_multi_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, */ template ucs_status_t -launch_uct_put_multi(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +launch_uct_put_multi(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length) { @@ -142,13 +142,13 @@ launch_uct_put_multi(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, } template ucs_status_t launch_uct_put_multi( - uct_device_ep_h ep, uct_device_mem_element_t *mem_list, const void *va, + uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length); template static __global__ void -uct_put_partial_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +uct_put_partial_kernel(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length, ucs_status_t *status_p) { @@ -191,7 +191,7 @@ uct_put_partial_kernel(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, */ template ucs_status_t -launch_uct_put_partial(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +launch_uct_put_partial(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length) { @@ -211,7 +211,7 @@ launch_uct_put_partial(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, template ucs_status_t launch_uct_put_partial( - uct_device_ep_h ep, uct_device_mem_element_t *mem_list, const void *va, + uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length); } // namespace ucx_cuda diff --git a/test/gtest/uct/cuda/test_kernels.h b/test/gtest/uct/cuda/test_kernels.h index 9b0a683fd48..33c718f70e6 100644 --- a/test/gtest/uct/cuda/test_kernels.h +++ b/test/gtest/uct/cuda/test_kernels.h @@ -14,22 +14,22 @@ namespace ucx_cuda { ucs_status_t launch_uct_put_single(uct_device_ep_h ep, - uct_device_mem_element_t *mem_elem, + uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length); ucs_status_t launch_uct_atomic(uct_device_ep_h ep, - uct_device_mem_element_t *mem_elem, uint64_t rva, + uct_device_mem_elem_t *mem_elem, uint64_t rva, uint64_t add); template ucs_status_t -launch_uct_put_multi(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +launch_uct_put_multi(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length); template ucs_status_t -launch_uct_put_partial(uct_device_ep_h ep, uct_device_mem_element_t *mem_list, +launch_uct_put_partial(uct_device_ep_h ep, uct_device_mem_elem_t *mem_list, const void *va, uint64_t rva, uint64_t atomic_rva, size_t length); diff --git a/test/gtest/uct/cuda/test_kernels_uct.cu b/test/gtest/uct/cuda/test_kernels_uct.cu index a6ad8a17881..799a6138a6c 100644 --- a/test/gtest/uct/cuda/test_kernels_uct.cu +++ b/test/gtest/uct/cuda/test_kernels_uct.cu @@ -97,13 +97,13 @@ template class device_result_ptr { return false; } - template - static __global__ void - uct_put_single_kernel(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, - const void *address, uint64_t remote_address, - size_t length, ucs_status_t *status) - { +template +static __global__ void +uct_put_single_kernel(uct_device_ep_h device_ep, + const uct_device_mem_elem_t *mem_elem, + const void *address, uint64_t remote_address, + size_t length, ucs_status_t *status) +{ uct_device_completion_t comp; if (is_op_enabled(level)) { @@ -117,13 +117,11 @@ template class device_result_ptr { * Basic single element put operation. */ ucs_status_t launch_uct_put_single(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, + const uct_device_mem_elem_t *mem_elem, const void *address, uint64_t remote_address, - size_t length, - ucs_device_level_t level, - unsigned num_threads, - unsigned num_blocks) - { + size_t length, ucs_device_level_t level, + unsigned num_threads, unsigned num_blocks) +{ device_result_ptr status = UCS_ERR_NOT_IMPLEMENTED; cudaError_t st; @@ -168,8 +166,7 @@ ucs_status_t launch_uct_put_single(uct_device_ep_h device_ep, template static __global__ void -uct_atomic_kernel(uct_device_ep_h ep, - const uct_device_mem_element_t *mem_elem, +uct_atomic_kernel(uct_device_ep_h ep, const uct_device_mem_elem_t *mem_elem, uint64_t rva, uint64_t add, ucs_status_t *status_p) { uct_device_completion_t comp; @@ -182,11 +179,9 @@ uct_atomic_kernel(uct_device_ep_h ep, } ucs_status_t launch_uct_atomic(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, - uint64_t rva, - uint64_t add, - ucs_device_level_t level, - unsigned num_threads, + const uct_device_mem_elem_t *mem_elem, + uint64_t rva, uint64_t add, + ucs_device_level_t level, unsigned num_threads, unsigned num_blocks) { device_result_ptr status = UCS_ERR_NOT_IMPLEMENTED; @@ -224,8 +219,7 @@ ucs_status_t launch_uct_atomic(uct_device_ep_h device_ep, template static __global__ void -uct_put_multi_kernel(uct_device_ep_h ep, - const uct_device_mem_element_t *mem_list, +uct_put_multi_kernel(uct_device_ep_h ep, const uct_device_mem_elem_t *mem_list, size_t mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, @@ -245,7 +239,7 @@ uct_put_multi_kernel(uct_device_ep_h ep, ucs_status_t launch_uct_put_multi(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_list, + const uct_device_mem_elem_t *mem_list, size_t mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, @@ -307,7 +301,7 @@ launch_uct_put_multi(uct_device_ep_h device_ep, template static __global__ void uct_put_multi_partial_kernel( - uct_device_ep_h ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h ep, const uct_device_mem_elem_t *mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *offsets, const size_t *lengths, unsigned counter_index, @@ -326,7 +320,7 @@ static __global__ void uct_put_multi_partial_kernel( } ucs_status_t launch_uct_put_multi_partial( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *offsets, const size_t *lengths, unsigned counter_index, diff --git a/test/gtest/uct/cuda/test_kernels_uct.h b/test/gtest/uct/cuda/test_kernels_uct.h index beaba9b03e7..3d92f7cdb47 100644 --- a/test/gtest/uct/cuda/test_kernels_uct.h +++ b/test/gtest/uct/cuda/test_kernels_uct.h @@ -16,22 +16,20 @@ namespace cuda_uct { int launch_memcmp(const void *s1, const void *s2, size_t size); ucs_status_t launch_uct_put_single(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, + const uct_device_mem_elem_t *mem_elem, const void *address, uint64_t remote_address, size_t length, ucs_device_level_t level, unsigned num_threads, unsigned num_blocks); ucs_status_t launch_uct_atomic(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_elem, - uint64_t rva, - uint64_t add, - ucs_device_level_t level, - unsigned num_threads, + const uct_device_mem_elem_t *mem_elem, + uint64_t rva, uint64_t add, + ucs_device_level_t level, unsigned num_threads, unsigned num_blocks); ucs_status_t launch_uct_put_multi(uct_device_ep_h device_ep, - const uct_device_mem_element_t *mem_list, + const uct_device_mem_elem_t *mem_list, size_t mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *lengths, uint64_t counter_inc_value, @@ -39,7 +37,7 @@ launch_uct_put_multi(uct_device_ep_h device_ep, unsigned num_threads, unsigned num_blocks); ucs_status_t launch_uct_put_multi_partial( - uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_list, + uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_list, const unsigned *mem_list_indices, unsigned mem_list_count, void *const *addresses, const uint64_t *remote_addresses, const size_t *offsets, const size_t *lengths, unsigned counter_index, diff --git a/test/gtest/uct/test_device.cc b/test/gtest/uct/test_device.cc index 5a6024cfdd2..7eba027cc9b 100644 --- a/test/gtest/uct/test_device.cc +++ b/test/gtest/uct/test_device.cc @@ -145,19 +145,18 @@ UCS_TEST_P(test_device, single) mapped_buffer sendbuf(length, SEED1, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); mapped_buffer recvbuf(length, SEED2, *m_receiver, 0, UCS_MEMORY_TYPE_CUDA); - mapped_buffer elembuf_host(sizeof(uct_device_mem_element_t), 0, *m_sender, - 0, UCS_MEMORY_TYPE_HOST); - mapped_buffer elembuf(sizeof(uct_device_mem_element_t), 0, *m_sender, 0, + mapped_buffer elembuf_host(sizeof(uct_device_mem_elem_t), 0, *m_sender, 0, + UCS_MEMORY_TYPE_HOST); + mapped_buffer elembuf(sizeof(uct_device_mem_elem_t), 0, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); - uct_device_mem_element_t *mem_elem_host = (uct_device_mem_element_t*) - elembuf_host.ptr(); - uct_device_mem_element_t *mem_elem = (uct_device_mem_element_t*) - elembuf.ptr(); + uct_device_mem_elem_t *mem_elem_host = (uct_device_mem_elem_t*) + elembuf_host.ptr(); + uct_device_mem_elem_t *mem_elem = (uct_device_mem_elem_t*)elembuf.ptr(); ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(), recvbuf.rkey(), mem_elem_host)); /* Copy packed element from host to GPU */ ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)mem_elem, mem_elem_host, - sizeof(uct_device_mem_element_t))); + sizeof(uct_device_mem_elem_t))); uct_device_ep_h dev_ep; ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &dev_ep)); @@ -176,18 +175,17 @@ UCS_TEST_P(test_device, atomic) uint64_t signal_val = 0; size_t i; - mapped_buffer elembuf_host(sizeof(uct_device_mem_element_t), 0, *m_sender, - 0, UCS_MEMORY_TYPE_HOST); - mapped_buffer elembuf(sizeof(uct_device_mem_element_t), 0, *m_sender, 0, + mapped_buffer elembuf_host(sizeof(uct_device_mem_elem_t), 0, *m_sender, 0, + UCS_MEMORY_TYPE_HOST); + mapped_buffer elembuf(sizeof(uct_device_mem_elem_t), 0, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); - uct_device_mem_element_t *mem_elem_host = (uct_device_mem_element_t*) - elembuf_host.ptr(); - uct_device_mem_element_t *mem_elem = (uct_device_mem_element_t*) - elembuf.ptr(); + uct_device_mem_elem_t *mem_elem_host = (uct_device_mem_elem_t*) + elembuf_host.ptr(); + uct_device_mem_elem_t *mem_elem = (uct_device_mem_elem_t*)elembuf.ptr(); ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), nullptr, signal.rkey(), mem_elem_host)); ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)mem_elem, mem_elem_host, - sizeof(uct_device_mem_element_t))); + sizeof(uct_device_mem_elem_t))); uct_device_ep_h dev_ep; ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &dev_ep)); @@ -215,7 +213,7 @@ UCS_TEST_P(test_device, multi) uint64_t signal_val = 0; size_t i; - size_t total_elem_size = sizeof(uct_device_mem_element_t) * (iovcnt + 1); + size_t total_elem_size = sizeof(uct_device_mem_elem_t) * (iovcnt + 1); mapped_buffer elembuf_host(total_elem_size, 0, *m_sender, 0, UCS_MEMORY_TYPE_HOST); mapped_buffer elembuf(total_elem_size, 0, *m_sender, 0, @@ -223,16 +221,16 @@ UCS_TEST_P(test_device, multi) for (i = 0; i < iovcnt; i++) { ASSERT_UCS_OK(uct_md_mem_elem_pack( m_sender->md(), sendbuf.memh(), recvbuf.rkey(), - (uct_device_mem_element_t*)UCS_PTR_BYTE_OFFSET( + (uct_device_mem_elem_t*)UCS_PTR_BYTE_OFFSET( elembuf_host.ptr(), - sizeof(uct_device_mem_element_t) * i))); + sizeof(uct_device_mem_elem_t) * i))); } ASSERT_UCS_OK(uct_md_mem_elem_pack( m_sender->md(), NULL, signal.rkey(), - (uct_device_mem_element_t*)UCS_PTR_BYTE_OFFSET( + (uct_device_mem_elem_t*)UCS_PTR_BYTE_OFFSET( elembuf_host.ptr(), - sizeof(uct_device_mem_element_t) * iovcnt))); + sizeof(uct_device_mem_elem_t) * iovcnt))); /* Copy all packed elements from host to GPU in one operation */ ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)elembuf.ptr(), @@ -242,7 +240,7 @@ UCS_TEST_P(test_device, multi) ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &dev_ep)); for (i = 0; i < 100; i++) { ASSERT_UCS_OK(ucx_cuda::launch_uct_put_multi( - dev_ep, (uct_device_mem_element_t*)elembuf.ptr(), sendbuf.ptr(), + dev_ep, (uct_device_mem_elem_t*)elembuf.ptr(), sendbuf.ptr(), (uintptr_t)recvbuf.ptr(), (uintptr_t)signal.ptr(), length)); signal_val += 4; @@ -267,7 +265,7 @@ UCS_TEST_P(test_device, partial) uint64_t signal_val = 0; size_t i; - size_t total_elem_size = sizeof(uct_device_mem_element_t) * (iovcnt + 1); + size_t total_elem_size = sizeof(uct_device_mem_elem_t) * (iovcnt + 1); mapped_buffer elembuf_host(total_elem_size, 0, *m_sender, 0, UCS_MEMORY_TYPE_HOST); mapped_buffer elembuf(total_elem_size, 0, *m_sender, 0, @@ -275,16 +273,16 @@ UCS_TEST_P(test_device, partial) for (i = 0; i < iovcnt; i++) { ASSERT_UCS_OK(uct_md_mem_elem_pack( m_sender->md(), sendbuf.memh(), recvbuf.rkey(), - (uct_device_mem_element_t*)UCS_PTR_BYTE_OFFSET( + (uct_device_mem_elem_t*)UCS_PTR_BYTE_OFFSET( elembuf_host.ptr(), - sizeof(uct_device_mem_element_t) * i))); + sizeof(uct_device_mem_elem_t) * i))); } ASSERT_UCS_OK(uct_md_mem_elem_pack( m_sender->md(), NULL, signal.rkey(), - (uct_device_mem_element_t*)UCS_PTR_BYTE_OFFSET( + (uct_device_mem_elem_t*)UCS_PTR_BYTE_OFFSET( elembuf_host.ptr(), - sizeof(uct_device_mem_element_t) * iovcnt))); + sizeof(uct_device_mem_elem_t) * iovcnt))); /* Copy all packed elements from host to GPU in one operation */ ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)elembuf.ptr(), @@ -294,7 +292,7 @@ UCS_TEST_P(test_device, partial) ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &dev_ep)); for (i = 0; i < 100; i++) { ASSERT_UCS_OK(ucx_cuda::launch_uct_put_partial( - dev_ep, (uct_device_mem_element_t*)elembuf.ptr(), sendbuf.ptr(), + dev_ep, (uct_device_mem_elem_t*)elembuf.ptr(), sendbuf.ptr(), (uintptr_t)recvbuf.ptr(), (uintptr_t)signal.ptr(), length)); signal_val += 4; From 7986ea39c7c53a304a1fbb55c5a8db16b7867a27 Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Sat, 25 Apr 2026 11:51:01 +0300 Subject: [PATCH 09/12] UCP/DEVICE: Add multi-lane support - 9 --- test/gtest/uct/cuda/test_cuda_ipc_device.cc | 26 +++++++++------------ test/gtest/uct/cuda/test_kernels.cu | 5 ++-- test/gtest/uct/cuda/test_kernels.h | 2 +- test/gtest/uct/cuda/test_kernels_uct.cu | 14 +++++------ test/gtest/uct/cuda/test_kernels_uct.h | 11 ++++----- test/gtest/uct/test_device.cc | 24 +++++++++---------- 6 files changed, 36 insertions(+), 46 deletions(-) diff --git a/test/gtest/uct/cuda/test_cuda_ipc_device.cc b/test/gtest/uct/cuda/test_cuda_ipc_device.cc index 266610152f6..1f64979875e 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_device.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_device.cc @@ -221,26 +221,22 @@ UCS_TEST_P(test_cuda_ipc_rma_device, put_device) mapped_buffer sendbuf(length, SEED1, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); mapped_buffer recvbuf(length, SEED2, *m_receiver, 0, UCS_MEMORY_TYPE_CUDA); - uct_device_local_mem_list_elem_t src_elem_host; + uct_device_mem_elem_t src_elem_host; ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(), - recvbuf.rkey(), - &src_elem_host.uct_mem_element)); + recvbuf.rkey(), &src_elem_host)); - uct_device_local_mem_list_elem_t *src_elem; - ASSERT_EQ(CUDA_SUCCESS, - cuMemAlloc((CUdeviceptr*)&src_elem, - sizeof(uct_device_local_mem_list_elem_t))); - ASSERT_EQ(CUDA_SUCCESS, - cuMemcpyHtoD((CUdeviceptr)src_elem, &src_elem_host, - sizeof(uct_device_local_mem_list_elem_t))); + uct_device_mem_elem_t *src_elem; + ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&src_elem, + sizeof(uct_device_mem_elem_t))); + ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)src_elem, &src_elem_host, + sizeof(uct_device_mem_elem_t))); - uct_device_mem_element_t *mem_elem; + uct_device_mem_elem_t *mem_elem; ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc((CUdeviceptr*)&mem_elem, - sizeof(uct_device_mem_element_t))); - ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)mem_elem, - &src_elem_host.uct_mem_element, - sizeof(uct_device_mem_element_t))); + sizeof(uct_device_mem_elem_t))); + ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)mem_elem, &src_elem_host, + sizeof(uct_device_mem_elem_t))); uct_device_ep_h device_ep; ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &device_ep)); diff --git a/test/gtest/uct/cuda/test_kernels.cu b/test/gtest/uct/cuda/test_kernels.cu index 09ff7ec1b57..ca28709880a 100644 --- a/test/gtest/uct/cuda/test_kernels.cu +++ b/test/gtest/uct/cuda/test_kernels.cu @@ -13,8 +13,7 @@ namespace ucx_cuda { static __global__ void -uct_put_kernel(uct_device_ep_h ep, - const uct_device_local_mem_elem_t *src_elem, +uct_put_kernel(uct_device_ep_h ep, const uct_device_mem_elem_t *src_elem, const uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length, ucs_status_t *status_p) { @@ -39,7 +38,7 @@ uct_put_kernel(uct_device_ep_h ep, * Basic single element put operation. */ ucs_status_t launch_uct_put(uct_device_ep_h ep, - const uct_device_local_mem_elem_t *src_elem, + const uct_device_mem_elem_t *src_elem, const uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length) { diff --git a/test/gtest/uct/cuda/test_kernels.h b/test/gtest/uct/cuda/test_kernels.h index 23192af51df..ff23a4f2d28 100644 --- a/test/gtest/uct/cuda/test_kernels.h +++ b/test/gtest/uct/cuda/test_kernels.h @@ -14,7 +14,7 @@ namespace ucx_cuda { ucs_status_t launch_uct_put(uct_device_ep_h ep, - const uct_device_local_mem_elem_t *src_elem, + const uct_device_mem_elem_t *src_elem, const uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length); diff --git a/test/gtest/uct/cuda/test_kernels_uct.cu b/test/gtest/uct/cuda/test_kernels_uct.cu index fe73f3c7b73..4665f5cf5db 100644 --- a/test/gtest/uct/cuda/test_kernels_uct.cu +++ b/test/gtest/uct/cuda/test_kernels_uct.cu @@ -99,8 +99,7 @@ template class device_result_ptr { template static __global__ void -uct_put_kernel(uct_device_ep_h ep, - const uct_device_local_mem_elem_t *src_elem, +uct_put_kernel(uct_device_ep_h ep, const uct_device_mem_elem_t *src_elem, const uct_device_mem_elem_t *mem_elem, const void *va, uint64_t rva, size_t length, ucs_status_t *status_p) { @@ -119,12 +118,11 @@ uct_put_kernel(uct_device_ep_h ep, } } -ucs_status_t launch_uct_put(uct_device_ep_h device_ep, - const uct_device_local_mem_elem_t *src_elem, - const uct_device_mem_elem_t *mem_elem, - const void *va, uint64_t rva, size_t length, - ucs_device_level_t level, unsigned num_threads, - unsigned num_blocks) +ucs_status_t +launch_uct_put(uct_device_ep_h device_ep, const uct_device_mem_elem_t *src_elem, + const uct_device_mem_elem_t *mem_elem, const void *va, + uint64_t rva, size_t length, ucs_device_level_t level, + unsigned num_threads, unsigned num_blocks) { device_result_ptr status = UCS_ERR_NOT_IMPLEMENTED; cudaError_t st; diff --git a/test/gtest/uct/cuda/test_kernels_uct.h b/test/gtest/uct/cuda/test_kernels_uct.h index add641b78a7..b62ecda67a6 100644 --- a/test/gtest/uct/cuda/test_kernels_uct.h +++ b/test/gtest/uct/cuda/test_kernels_uct.h @@ -15,12 +15,11 @@ namespace cuda_uct { int launch_memcmp(const void *s1, const void *s2, size_t size); -ucs_status_t launch_uct_put(uct_device_ep_h device_ep, - const uct_device_local_mem_elem_t *src_elem, - const uct_device_mem_elem_t *mem_elem, - const void *va, uint64_t rva, size_t length, - ucs_device_level_t level, unsigned num_threads, - unsigned num_blocks); +ucs_status_t +launch_uct_put(uct_device_ep_h device_ep, const uct_device_mem_elem_t *src_elem, + const uct_device_mem_elem_t *mem_elem, const void *va, + uint64_t rva, size_t length, ucs_device_level_t level, + unsigned num_threads, unsigned num_blocks); ucs_status_t launch_uct_atomic(uct_device_ep_h device_ep, const uct_device_mem_elem_t *mem_elem, diff --git a/test/gtest/uct/test_device.cc b/test/gtest/uct/test_device.cc index e453c1dede1..24e00098245 100644 --- a/test/gtest/uct/test_device.cc +++ b/test/gtest/uct/test_device.cc @@ -145,29 +145,27 @@ UCS_TEST_P(test_device, put) mapped_buffer sendbuf(length, SEED1, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); mapped_buffer recvbuf(length, SEED2, *m_receiver, 0, UCS_MEMORY_TYPE_CUDA); - uct_device_local_mem_list_elem_t src_elem_host; - src_elem_host.addr = nullptr; + uct_device_mem_elem_t src_elem_host; ASSERT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(), - recvbuf.rkey(), - &src_elem_host.uct_mem_element)); + recvbuf.rkey(), &src_elem_host)); - mapped_buffer src_elembuf(sizeof(uct_device_local_mem_list_elem_t), 0, - *m_sender, 0, UCS_MEMORY_TYPE_CUDA); + mapped_buffer src_elembuf(sizeof(src_elem_host), 0, *m_sender, 0, + UCS_MEMORY_TYPE_CUDA); ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)src_elembuf.ptr(), &src_elem_host, - sizeof(uct_device_local_mem_list_elem_t))); + sizeof(src_elem_host))); - mapped_buffer rem_elembuf(sizeof(uct_device_mem_element_t), 0, *m_sender, 0, + mapped_buffer rem_elembuf(sizeof(src_elem_host), 0, *m_sender, 0, UCS_MEMORY_TYPE_CUDA); - ASSERT_EQ(CUDA_SUCCESS, cuMemcpyHtoD((CUdeviceptr)rem_elembuf.ptr(), - &src_elem_host.uct_mem_element, - sizeof(uct_device_mem_element_t))); + ASSERT_EQ(CUDA_SUCCESS, + cuMemcpyHtoD((CUdeviceptr)rem_elembuf.ptr(), &src_elem_host, + sizeof(src_elem_host))); uct_device_ep_h dev_ep; ASSERT_UCS_OK(uct_ep_get_device_ep(m_sender->ep(0), &dev_ep)); ASSERT_UCS_OK(ucx_cuda::launch_uct_put( - dev_ep, (const uct_device_local_mem_list_elem_t*)src_elembuf.ptr(), - (const uct_device_mem_element_t*)rem_elembuf.ptr(), sendbuf.ptr(), + dev_ep, (const uct_device_mem_elem_t*)src_elembuf.ptr(), + (const uct_device_mem_elem_t*)rem_elembuf.ptr(), sendbuf.ptr(), (uintptr_t)recvbuf.ptr(), length)); recvbuf.pattern_check(SEED1); From 72d8beb72940705e9aa9388e3e92fba634378327 Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Sat, 25 Apr 2026 14:53:44 +0300 Subject: [PATCH 10/12] UCP/DEVICE: Add multi-lane support - 10 --- src/ucp/core/ucp_device.c | 3 --- test/gtest/uct/test_device.cc | 11 ++++++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 839b5a7a4b9..26d907eb5cc 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -23,9 +23,6 @@ #include "ucp_mm.inl" -#define UCP_DEVICE_MEM_LIST_MAX_EPS 2 - - typedef struct { uct_allocated_memory_t mem; uint32_t mem_list_length; diff --git a/test/gtest/uct/test_device.cc b/test/gtest/uct/test_device.cc index 24e00098245..c2b26551049 100644 --- a/test/gtest/uct/test_device.cc +++ b/test/gtest/uct/test_device.cc @@ -111,9 +111,9 @@ class test_device : public uct_test { m_cuda_dev = uct_cuda_get_cuda_device( m_sender->iface_attr().ctl_device); - ASSERT_NE(m_cuda_dev, CU_DEVICE_INVALID) - << " sys_device " - << static_cast(m_sender->iface_attr().ctl_device); + if (m_cuda_dev == CU_DEVICE_INVALID) { + return; + } status = UCT_CUDADRV_FUNC_LOG_ERR( cuDevicePrimaryCtxRetain(&ctx, m_cuda_dev)); @@ -126,6 +126,11 @@ class test_device : public uct_test { void cleanup() { uct_test::cleanup(); + + if (m_cuda_dev == CU_DEVICE_INVALID) { + return; + } + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL)); (void)UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(m_cuda_dev)); } From f0228c0a71dd17aaacc6ede1d21ded093402248e Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Mon, 27 Apr 2026 19:05:52 +0300 Subject: [PATCH 11/12] UCP/DEVICE: Add multi-lane support - 11 --- src/uct/ib/base/ib_md.c | 4 ++-- src/uct/ib/base/ib_md.h | 2 +- src/uct/ib/mlx5/gdaki/gdaki.c | 6 +++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c index 793bf723457..639cc8f6c3d 100644 --- a/src/uct/ib/base/ib_md.c +++ b/src/uct/ib/base/ib_md.c @@ -120,10 +120,10 @@ ucs_config_field_t uct_ib_md_config_table[] = { "Use GPU Direct RDMA for HCA to access GPU pages directly\n", ucs_offsetof(uct_ib_md_config_t, enable_gpudirect_rdma), UCS_CONFIG_TYPE_TERNARY}, - {"GDA_MAX_HCA_PER_GPU", "1", + {"GDA_MAX_HCA_PER_GPU", "auto", "Max number of HCA devices to use for GDA per one GPU device.", ucs_offsetof(uct_ib_md_config_t, ext.gda_max_hca_per_gpu), - UCS_CONFIG_TYPE_UINT}, + UCS_CONFIG_TYPE_ULUNITS}, {"GDA_DMABUF_ENABLE", "try", "Enable DMA-BUF in GDA.", diff --git a/src/uct/ib/base/ib_md.h b/src/uct/ib/base/ib_md.h index b7ea5d0eb66..ba60cfa58da 100644 --- a/src/uct/ib/base/ib_md.h +++ b/src/uct/ib/base/ib_md.h @@ -112,7 +112,7 @@ typedef struct uct_ib_md_ext_config { unsigned long reg_retry_cnt; /**< Memory registration retry count */ unsigned smkey_block_size; /**< Mkey indexes in a symmetric block */ int direct_nic; /**< Direct NIC with GPU functionality */ - unsigned gda_max_hca_per_gpu; /**< Threshold of IB per GPU */ + unsigned long gda_max_hca_per_gpu; /**< Threshold of IB per GPU */ int gda_dmabuf_enable; /**< Enable DMA-BUF in GDA */ } uct_ib_md_ext_config_t; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index 53b13021e44..218c894160b 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -1246,7 +1246,7 @@ static int uct_gdaki_dev_matrix_score(const void *pa, const void *pb, void *arg) uct_gdaki_dev_matrix_elem_t * uct_gdaki_dev_matrix_init(const uct_ib_md_t *ib_md, size_t *dmat_length_p) { - unsigned ib_per_cuda = ib_md->config.gda_max_hca_per_gpu; + unsigned long ib_per_cuda = ib_md->config.gda_max_hca_per_gpu; uct_gdaki_dev_matrix_elem_t *dmat = NULL; ucs_status_t status; int ibdev_index, cudadev_index, ibdev_count, cudadev_count; @@ -1325,6 +1325,10 @@ uct_gdaki_dev_matrix_init(const uct_ib_md_t *ib_md, size_t *dmat_length_p) ucs_assert(cudadev_count < UCT_GDAKI_MAX_CUDA_PER_IB); + if (ib_per_cuda == UCS_ULUNITS_AUTO) { + ib_per_cuda = ibdev_count / cudadev_count; + } + /* Map each CUDA device to the best suited IB devices */ for (cudadev_index = 0; cudadev_index < cudadev_count; cudadev_index++) { status = UCT_CUDADRV_FUNC_LOG_ERR( From 508fa22a802f2a2cda1faba2aef68a34c32037dc Mon Sep 17 00:00:00 2001 From: Artemy Kovalyov Date: Tue, 28 Apr 2026 11:17:39 +0300 Subject: [PATCH 12/12] UCP/DEVICE: Add multi-lane support - 12 --- src/ucp/api/device/ucp_device_impl.h | 11 ++++------- src/ucp/core/ucp_device.c | 12 +++++++++++- src/uct/api/uct.h | 2 +- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/ucp/api/device/ucp_device_impl.h b/src/ucp/api/device/ucp_device_impl.h index 32ea314a376..ce004d7677a 100644 --- a/src/ucp/api/device/ucp_device_impl.h +++ b/src/ucp/api/device/ucp_device_impl.h @@ -110,7 +110,7 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep, _uct_channel_id = _channel_id / _handle->num_lanes; -#define UCP_DEVICE_GET_ELEM(_handle, _index, _lane) \ +#define UCP_DEVICE_GET_ELEM(_handle, _index) \ static_castmem_elements[0])*>(UCS_PTR_BYTE_OFFSET( \ _handle->mem_elements, \ (sizeof(_handle->mem_elements[0]) + \ @@ -132,7 +132,7 @@ UCS_F_DEVICE ucs_status_t ucp_device_prepare_send_remote( } const auto dst_mem_element = UCP_DEVICE_GET_ELEM(dst_mem_list_h, - dst_mem_list_index, lane); + dst_mem_list_index); remote_address = dst_mem_element->addr; device_ep = dst_mem_element->tl[lane].ep; uct_elem = &dst_mem_element->tl[lane].uct; @@ -166,7 +166,7 @@ UCS_F_DEVICE ucs_status_t ucp_device_prepare_send( } const auto src_mem_elem = UCP_DEVICE_GET_ELEM(src_mem_list_h, - src_mem_list_index, lane); + src_mem_list_index); src_uct_elem = src_mem_elem->tl + lane; address = src_mem_elem->addr; @@ -326,7 +326,6 @@ 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_elem_t); ucs_status_t status; status = ucp_device_check_params(mem_list_h, mem_list_index); @@ -334,9 +333,7 @@ ucp_device_get_ptr(const ucp_device_remote_mem_list_h mem_list_h, return status; } - const auto mem_element = static_cast( - 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->tl[0].ep, &mem_element->tl[0].uct, mem_element->addr, addr_p); diff --git a/src/ucp/core/ucp_device.c b/src/ucp/core/ucp_device.c index 26d907eb5cc..8efc601fe15 100644 --- a/src/ucp/core/ucp_device.c +++ b/src/ucp/core/ucp_device.c @@ -561,6 +561,11 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( } ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev); + + /* handle->num_lanes is the least common multiple of both lane types, so: + * - each lane is replicated num_lanes / popcount(tl_bitmap) times + * - channel_id % num_lanes maps to the correct lane, regardless of lane type + */ num_lanes = UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_LKEY]); if (!num_lanes) { if (!UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap[UCP_DEVICE_TL_TYPE_NOLKEY])) { @@ -594,7 +599,12 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle( } } - ucs_assert(tl_type < UCP_DEVICE_TL_TYPE_LAST); + if (tl_type == UCP_DEVICE_TL_TYPE_LAST) { + ucs_error("lane not found for element %zd", i); + status = UCS_ERR_INVALID_PARAM; + goto out; + } + status = ucp_device_remote_mem_list_fill(ucp_element, &tl_bitmap[tl_type], num_lanes, uct_element); diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index 590aa9de3e6..082f30d72bb 100644 --- a/src/uct/api/uct.h +++ b/src/uct/api/uct.h @@ -438,7 +438,7 @@ typedef enum uct_atomic_op { /* Interface capability */ #define UCT_IFACE_FLAG_INTER_NODE UCS_BIT(54) /**< Interface is inter-node capable */ #define UCT_IFACE_FLAG_DEVICE_EP UCS_BIT(55) /**< Interface supports device endpoint */ -#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /**< Interface require lkey for device operations */ +#define UCT_IFACE_FLAG_DEVICE_LKEY UCS_BIT(56) /**< Interface requires lkey for device operations */ /** * @} */