From 53c855b83421150958ce546888d26dcc61392340 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 4 Sep 2026 05:50:19 -0700 Subject: [PATCH 01/13] UCT/ZE/COPY: Allocate DMA-BUF-exportable memory for RDMA A Level Zero device-memory allocation may need to be created for external DMA-BUF export before its handle can be used for RDMA. The ze_copy memory domain advertised ze-device in 'dmabuf_mem_types' and obtained a file descriptor through zeMemGetAllocProperties(), but allocated without an export descriptor. On the tested system, registering a DMA-BUF obtained that way completed successfully, but RDMA reads returned incorrect data while the transfer reported success. Request DMA-BUF exportability in uct_ze_copy_mem_alloc(), and only for the memory types advertised in 'dmabuf_mem_types', so that allocating a type without a DMA-BUF path cannot fail because of the descriptor. Add UCX_ZE_COPY_DMABUF, following the cuda_copy and rocm_copy pattern: query DMA-BUF export support once at md_open() using zeDeviceGetExternalMemoryProperties(), fail md_open() when 'y' is requested but unsupported, and advertise 'dmabuf_mem_types' and return a DMA-BUF descriptor only when it is effectively enabled. Memory allocated by an application without requesting external export cannot be repaired by UCX, so 'n' is available to disable DMA-BUF registration for it. The presence of the Intel Xe module alone does not prove that an IB memory domain can register Level Zero device memory directly. UCP routes DMA-BUF registration by combining a provider's 'dmabuf_mem_types' with an RDMA memory domain's UCT_MD_FLAG_REG_DMABUF. Remove the Intel Xe module heuristic from the IB memory domain. With UCX_ZE_COPY_DMABUF=n on the tested system, that heuristic made UCP attempt ordinary registration, which failed with EFAULT instead of staging through ze_copy. Removing ZE device memory from the IB MD's 'reg_mem_types' also disables the packet-tracer bound, so apply the bound whenever the IB memory domain supports DMA-BUF registration; otherwise tracing may dereference device memory from the CPU. Add test_ucp_mem_type_rndv_zcopy, which sends from memory allocated by ucp_mem_map(UCP_MEM_MAP_ALLOCATE) to a host buffer over a rendezvous zero-copy lane and verifies a non-zero pattern. It asserts that the source is registered on the memory domain used by the operation, so a staged transfer cannot make the test pass silently, and skips when that lane cannot register the memory type. Reverting only the export descriptor makes the ze-device variants fail the pattern check with zeros, in both the get_zcopy and put_zcopy directions. Tested on Intel Arc Pro B60 Graphics with rc_mlx5 over RoCE: 16 gtest variants pass, and ucx_perftest tag_lat with -V passes for ze-device, ze-host and ze-managed, including the DMABUF=n and DMABUF=y configurations. Signed-off-by: Yaser Afshar --- src/uct/ib/base/ib_md.c | 18 ++- src/uct/ze/copy/ze_copy_md.c | 80 ++++++++++- src/uct/ze/copy/ze_copy_md.h | 8 +- test/gtest/ucp/test_ucp_mem_type.cc | 202 ++++++++++++++++++++++++++++ 4 files changed, 292 insertions(+), 16 deletions(-) diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c index 912ac7b8c20..43b7a92acdf 100644 --- a/src/uct/ib/base/ib_md.c +++ b/src/uct/ib/base/ib_md.c @@ -2,6 +2,7 @@ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * Copyright (C) The University of Tennessee and The University * of Tennessee Research Foundation. 2016. ALL RIGHTS RESERVED. + * Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -1446,15 +1447,14 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md, /* check if ROCM KFD driver is loaded */ uct_ib_check_gpudirect_driver(md, "/dev/kfd", UCS_MEMORY_TYPE_ROCM); - /* Check if Intel Xe driver is loaded */ - uct_ib_check_gpudirect_driver(md, "/sys/module/xe/srcversion", - UCS_MEMORY_TYPE_ZE_DEVICE); - /* Check for HabanaLabs Gaudi DMABuf support */ uct_ib_check_gpudirect_driver(md, "/dev/accel/accel0", UCS_MEMORY_TYPE_GAUDI); uct_ib_check_gpudirect_driver(md, "/dev/hl0", UCS_MEMORY_TYPE_GAUDI); + /* Do not infer Level Zero device-memory registration support from the + * presence of the Xe module. DMA-BUF support is detected below. */ + /* Check for dma-buf support */ uct_ib_md_check_dmabuf(md); } @@ -1463,15 +1463,19 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md, !(md->cap_flags & UCT_MD_FLAG_REG_DMABUF) && (md_config->enable_gpudirect_rdma == UCS_YES)) { ucs_error("%s: Couldn't enable GPUDirect RDMA. Please make sure " - "nv_peer_mem, amdgpu plugin, or Intel Xe driver is " - "installed correctly, or dmabuf is supported.", + "nv_peer_mem or amdgpu plugin is installed correctly, or " + "dmabuf is supported.", uct_ib_device_name(&md->dev)); status = UCS_ERR_UNSUPPORTED; goto err_cleanup_device; } + /* Limit the number of SGEs dumped by the packet tracer, so that it does not + * dereference non-host memory from the CPU. DMA-BUF registration makes such + * memory reachable without it being listed in 'reg_mem_types'. */ md->dev.max_zcopy_log_sge = INT_MAX; - if (md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) { + if ((md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) || + (md->cap_flags & UCT_MD_FLAG_REG_DMABUF)) { md->dev.max_zcopy_log_sge = 1; } diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index 99730ffcf74..81ff3f92126 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -22,6 +22,11 @@ #include +/* Memory types which can be exported as a DMA-BUF file descriptor */ +#define UCT_ZE_COPY_MD_DMABUF_MEM_TYPES \ + (UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE)) + + static ucs_config_field_t uct_ze_copy_md_config_table[] = { {"", "", NULL, ucs_offsetof(uct_ze_copy_md_config_t, super), UCS_CONFIG_TYPE_TABLE(uct_md_config_table)}, @@ -31,11 +36,43 @@ static ucs_config_field_t uct_ze_copy_md_config_table[] = { ucs_offsetof(uct_ze_copy_md_config_t, device_ordinal), UCS_CONFIG_TYPE_INT}, + {"DMABUF", "try", + "Enable using cross-device dmabuf file descriptor.\n" + "Level Zero device memory may have to be allocated with external\n" + "DMA-BUF export enabled. Without it, registration can succeed but an\n" + "RDMA device may read incorrect data, and no error is reported. Set to\n" + "'n' if the application allocates Level Zero memory without\n" + "requesting external export.", + ucs_offsetof(uct_ze_copy_md_config_t, enable_dmabuf), + UCS_CONFIG_TYPE_TERNARY}, + {NULL} }; +static int uct_ze_copy_md_is_dmabuf_supported(ze_device_handle_t ze_device) +{ + ze_device_external_memory_properties_t props = { + .stype = ZE_STRUCTURE_TYPE_DEVICE_EXTERNAL_MEMORY_PROPERTIES + }; + int supported; + + if (UCT_ZE_FUNC_LOG_DEBUG( + zeDeviceGetExternalMemoryProperties(ze_device, &props)) != + UCS_OK) { + return 0; + } + + supported = !!(props.memoryAllocationExportTypes & + ZE_EXTERNAL_MEMORY_TYPE_FLAG_DMA_BUF); + ucs_debug("ze device %p: DMA-BUF export is%s supported", ze_device, + supported ? "" : " not"); + return supported; +} + static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr) { + uct_ze_copy_md_t *ze_md = ucs_derived_of(md, uct_ze_copy_md_t); + uct_md_base_md_query(md_attr); md_attr->flags = UCT_MD_FLAG_REG | UCT_MD_FLAG_ALLOC; md_attr->reg_mem_types = UCS_BIT(UCS_MEMORY_TYPE_HOST) | @@ -55,8 +92,8 @@ static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr) md_attr->detect_mem_types = UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) | UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED); - md_attr->dmabuf_mem_types = UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | - UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE); + md_attr->dmabuf_mem_types = ze_md->enable_dmabuf ? + UCT_ZE_COPY_MD_DMABUF_MEM_TYPES : 0; md_attr->max_alloc = SIZE_MAX; return UCS_OK; } @@ -66,7 +103,18 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p, ucs_memory_type_t mem_type, ucs_sys_device_t sys_dev, unsigned flags, const char *alloc_name, uct_mem_h *memh_p) { - uct_ze_copy_md_t *md = ucs_derived_of(tl_md, uct_ze_copy_md_t); + uct_ze_copy_md_t *md = ucs_derived_of(tl_md, uct_ze_copy_md_t); + /* Request DMA-BUF exportability for the memory types advertised in + * 'dmabuf_mem_types', and only for those, so that allocating a type without + * a DMA-BUF path cannot fail because of the descriptor. Level Zero may + * return a DMA-BUF for an allocation which was not created for external + * export, but such a handle does not necessarily describe the requested + * range, and an RDMA device may read incorrect data through it. */ + ze_external_memory_export_desc_t export_desc = { + .stype = ZE_STRUCTURE_TYPE_EXTERNAL_MEMORY_EXPORT_DESC, + .flags = ZE_EXTERNAL_MEMORY_TYPE_FLAG_DMA_BUF + }; + void *export_desc_p = NULL; ze_host_mem_alloc_desc_t host_desc = { .stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC }; @@ -76,13 +124,20 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p, size_t alignment = ucs_get_page_size(); ucs_status_t status; + if (md->enable_dmabuf && + (UCS_BIT(mem_type) & UCT_ZE_COPY_MD_DMABUF_MEM_TYPES)) { + export_desc_p = &export_desc; + } + switch (mem_type) { case UCS_MEMORY_TYPE_ZE_HOST: + host_desc.pNext = export_desc_p; status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocHost(md->ze_context, &host_desc, *length_p, alignment, address_p)); break; case UCS_MEMORY_TYPE_ZE_DEVICE: + dev_desc.pNext = export_desc_p; status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocDevice(md->ze_context, &dev_desc, *length_p, alignment, md->ze_device, @@ -198,12 +253,14 @@ static ucs_status_t uct_ze_copy_md_mem_query(uct_md_h md, const void *addr, const size_t length, uct_md_mem_attr_v2_t *mem_attr_p) { - int dmabuf_fd = UCT_DMABUF_FD_INVALID; - int *dmabuf_fd_p = NULL; + uct_ze_copy_md_t *ze_md = ucs_derived_of(md, uct_ze_copy_md_t); + int dmabuf_fd = UCT_DMABUF_FD_INVALID; + int *dmabuf_fd_p = NULL; ucs_memory_info_t mem_info; ucs_status_t status; - if (mem_attr_p->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) { + if ((mem_attr_p->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) && + ze_md->enable_dmabuf) { mem_attr_p->dmabuf_fd = UCT_DMABUF_FD_INVALID; dmabuf_fd_p = &dmabuf_fd; } @@ -302,6 +359,7 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, ze_context_desc_t context_desc = {}; ze_result_t ret; const uct_ze_subdevice_t *subdevice; + int dmabuf_supported; ze_driver = uct_ze_base_get_driver(); if (ze_driver == NULL) { @@ -337,6 +395,16 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, return UCS_ERR_NO_DEVICE; } + dmabuf_supported = uct_ze_copy_md_is_dmabuf_supported(md->ze_device); + if ((config->enable_dmabuf == UCS_YES) && !dmabuf_supported) { + ucs_error("%s: DMA-BUF is not supported by the Level Zero device", + md_name); + zeContextDestroy(md->ze_context); + ucs_free(md); + return UCS_ERR_UNSUPPORTED; + } + + md->enable_dmabuf = (config->enable_dmabuf != UCS_NO) && dmabuf_supported; md->super.ops = &md_ops; md->super.component = &uct_ze_copy_component; diff --git a/src/uct/ze/copy/ze_copy_md.h b/src/uct/ze/copy/ze_copy_md.h index f47554b6c39..46bf6d47755 100644 --- a/src/uct/ze/copy/ze_copy_md.h +++ b/src/uct/ze/copy/ze_copy_md.h @@ -1,5 +1,5 @@ /* - * Copyright (C) Intel Corporation, 2023-2024. ALL RIGHTS RESERVED. + * Copyright (C) Intel Corporation, 2023-2026. ALL RIGHTS RESERVED. * See file LICENSE for terms. */ @@ -21,6 +21,7 @@ typedef struct uct_ze_copy_md { uct_md_t super; /**< Domain info */ ze_context_handle_t ze_context; ze_device_handle_t ze_device; + int enable_dmabuf; /**< Expose memory as DMA-BUF */ } uct_ze_copy_md_t; @@ -28,8 +29,9 @@ typedef struct uct_ze_copy_md { * ze copy domain configuration. */ typedef struct uct_ze_copy_md_config { - uct_md_config_t super; - int device_ordinal; + uct_md_config_t super; + int device_ordinal; + ucs_ternary_auto_value_t enable_dmabuf; } uct_ze_copy_md_config_t; diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index 9181eb48bf4..ce6b3e140d4 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -1,5 +1,6 @@ /** * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2017. ALL RIGHTS RESERVED. +* Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -11,6 +12,7 @@ extern "C" { #include #include #include +#include } @@ -98,6 +100,206 @@ UCS_TEST_P(test_ucp_mem_type_alloc_before_init, xfer) { UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_alloc_before_init, all, "all") + +/* + * Moves a payload out of memory allocated by UCP, using the transport which + * performs the rendezvous zero-copy operation. This covers the registration + * path a network transport uses to reach non-host memory directly, including + * registration through a DMA-BUF handle: memory which is registered but not + * readable by the transport is detected by the payload check rather than + * silently reported as a successful transfer. + */ +class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { +public: + static const uint64_t SEED = 0x1234567890abcdeflu; + +protected: + typedef ucs::handle mem_handle_t; + + static void unmap_memh(ucp_mem_h memh, ucp_context_h context) + { + ucs_status_t status = ucp_mem_unmap(context, memh); + if (status != UCS_OK) { + ucs_warn("failed to unmap memory: %s", ucs_status_string(status)); + } + } + + void init() override + { + test_ucp_mem_type::init(); + sender().connect(&receiver(), get_ep_params()); + receiver().connect(&sender(), get_ep_params()); + } + + /* Allocate through UCP so that a memory domain performs the allocation */ + void *alloc_mem(entity &e, size_t length, ucs_memory_type_t mem_type, + mem_handle_t &memh_handle) + { + ucp_mem_map_params_t params; + ucp_mem_attr_t attr; + ucp_mem_h memh; + ucs_status_t status; + + params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS | + UCP_MEM_MAP_PARAM_FIELD_LENGTH | + UCP_MEM_MAP_PARAM_FIELD_FLAGS | + UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; + params.address = NULL; + params.length = length; + params.flags = UCP_MEM_MAP_ALLOCATE; + params.memory_type = mem_type; + + status = ucp_mem_map(e.ucph(), ¶ms, &memh); + if (status != UCS_OK) { + UCS_TEST_SKIP_R(std::string("cannot allocate ") + + ucs_memory_type_names[mem_type] + " memory: " + + ucs_status_string(status)); + } + + memh_handle.reset(memh, unmap_memh, e.ucph()); + + attr.field_mask = UCP_MEM_ATTR_FIELD_ADDRESS; + status = ucp_mem_query(memh, &attr); + if (status != UCS_OK) { + UCS_TEST_ABORT("ucp_mem_query() failed: " + + std::string(ucs_status_string(status))); + } + + return attr.address; + } + + /* + * Memory domains used by the endpoint's rendezvous zero-copy lanes. + * + * 'remote' returns the destination domains, which is what a GET issued by + * this endpoint addresses on its peer, so the indices are in the peer's + * domain space. Otherwise the endpoint's own domains are returned, which is + * what a PUT issued by this endpoint registers locally. + */ + ucp_md_map_t rndv_zcopy_md_map(entity &e, bool remote) + { + const ucp_ep_config_t *ep_config = ucp_ep_config(e.ep()); + ucp_context_h context = e.ucph(); + ucp_md_map_t md_map = 0; + ucp_lane_index_t i, lane; + ucp_rsc_index_t rsc_index; + + for (i = 0; i < UCP_MAX_LANES; ++i) { + lane = ep_config->key.rma_bw_lanes[i]; + if (lane == UCP_NULL_LANE) { + break; + } + + if (remote) { + md_map |= UCS_BIT(ep_config->key.lanes[lane].dst_md_index); + } else { + rsc_index = ep_config->key.lanes[lane].rsc_index; + md_map |= UCS_BIT(context->tl_rscs[rsc_index].md_index); + } + } + + return md_map; + } + + std::string md_map_str(entity &e, ucp_md_map_t md_map) + { + ucp_context_h context = e.ucph(); + std::string names; + ucp_md_index_t md_index; + + ucs_for_each_bit(md_index, md_map) { + if (!names.empty()) { + names += ","; + } + names += (md_index < context->num_mds) ? + context->tl_mds[md_index].rsc.md_name : + ""; + } + + return names.empty() ? "" : names; + } + + void test_xfer_from_mem_type(size_t length, bool is_get); +}; + +void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, + bool is_get) +{ + const ucs_memory_type_t send_mem_type = mem_type(); + ucp_request_param_t param; + mem_handle_t send_memh, recv_memh; + void *send_buf, *recv_buf; + void *sreq, *rreq; + + if (!mem_buffer::is_mem_type_supported(send_mem_type)) { + UCS_TEST_SKIP_R(std::string(ucs_memory_type_names[send_mem_type]) + + " memory is not supported"); + } + + /* + * Resolve the lanes and the registration capability before allocating, so + * that skipping does not leak the memory mappings. + */ + entity &op_entity = is_get ? receiver() : sender(); + ucp_md_map_t lane_md_map = rndv_zcopy_md_map(op_entity, is_get); + if (lane_md_map == 0) { + UCS_TEST_SKIP_R("no rendezvous zero-copy lane is available"); + } + + if (!(sender().ucph()->reg_md_map[send_mem_type] & lane_md_map)) { + UCS_TEST_SKIP_R( + std::string("the rendezvous zero-copy lane cannot register ") + + ucs_memory_type_names[send_mem_type] + " memory"); + } + + send_buf = alloc_mem(sender(), length, send_mem_type, send_memh); + recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); + + UCS_TEST_MESSAGE << (is_get ? "get" : "put") << " lane mds: " + << md_map_str(sender(), lane_md_map) + << " | source memh mds: " + << md_map_str(sender(), send_memh->md_map); + + EXPECT_NE(0u, send_memh->md_map & lane_md_map) + << ucs_memory_type_names[send_mem_type] + << " memory is not registered on the memory domain used by the " + "rendezvous zero-copy lane, the payload would be staged instead " + "of accessed by the transport"; + + mem_buffer::pattern_fill(send_buf, length, SEED, send_mem_type); + mem_buffer::pattern_fill(recv_buf, length, 0, UCS_MEMORY_TYPE_HOST); + + param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; + param.memh = recv_memh; + rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf, length, 1, 1, + ¶m); + + param.memh = send_memh; + sreq = ucp_tag_send_nbx(sender().ep(), send_buf, length, 1, ¶m); + + EXPECT_UCS_OK(request_wait(sreq)); + EXPECT_UCS_OK(request_wait(rreq)); + + mem_buffer::pattern_check(recv_buf, length, SEED, UCS_MEMORY_TYPE_HOST); +} + +UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, get_zcopy, "RNDV_THRESH=0", + "RNDV_SCHEME=get_zcopy") +{ + test_xfer_from_mem_type(4 * UCS_MBYTE, true); +} + +UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, put_zcopy, "RNDV_THRESH=0", + "RNDV_SCHEME=put_zcopy") +{ + test_xfer_from_mem_type(4 * UCS_MBYTE, false); +} + +UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_rndv_zcopy, rcx, + "rc_x") +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcx_ze, + "rc_x,ze_copy") + class test_ucp_cuda : public ucp_test { public: static void get_test_variants(std::vector& variants) { From 6c01e4a8ab582a554643eadeba857a31b247070f Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 9 Oct 2026 13:23:17 -0700 Subject: [PATCH 02/13] UCT/ZE/COPY: Test transfers with DMA-BUF disabled Document the behavior of UCX_ZE_COPY_DMABUF=n and test transfers when the network rendezvous lane cannot register Level Zero memory: device memory must be staged through host fragments and CPU-accessible memory copied by the CPU. Check the selected protocol as well as the payload, since a forced zero-copy scheme can fall back to eager or to AM rendezvous when it cannot access the source. Cover memory allocated by UCP and by the application without an export descriptor, and instantiate the tests on rc_v as well, so that they run in builds without mlx5. Signed-off-by: Yaser Afshar --- src/uct/ze/copy/ze_copy_md.c | 3 +- test/gtest/ucp/test_ucp_mem_type.cc | 204 ++++++++++++++++++++++++---- 2 files changed, 178 insertions(+), 29 deletions(-) diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index 81ff3f92126..5f2df629bf1 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -42,7 +42,8 @@ static ucs_config_field_t uct_ze_copy_md_config_table[] = { "DMA-BUF export enabled. Without it, registration can succeed but an\n" "RDMA device may read incorrect data, and no error is reported. Set to\n" "'n' if the application allocates Level Zero memory without\n" - "requesting external export.", + "requesting external export; device memory is then staged through\n" + "host memory for RDMA.", ucs_offsetof(uct_ze_copy_md_config_t, enable_dmabuf), UCS_CONFIG_TYPE_TERNARY}, diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index ce6b3e140d4..8259a432585 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -8,10 +8,15 @@ #include "ucp_test.h" #include +#include + extern "C" { #include #include #include +#include +#include +#include #include } @@ -102,17 +107,23 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_alloc_before_init, all, "all") /* - * Moves a payload out of memory allocated by UCP, using the transport which - * performs the rendezvous zero-copy operation. This covers the registration - * path a network transport uses to reach non-host memory directly, including - * registration through a DMA-BUF handle: memory which is registered but not - * readable by the transport is detected by the payload check rather than - * silently reported as a successful transfer. + * Moves a payload out of memory of the tested type, allocated by UCP or by the + * application, over a rendezvous protocol. The zero-copy mode covers the + * registration path a network transport uses to reach the memory directly, + * including registration through a DMA-BUF handle: memory which is registered + * but not readable by the transport is detected by the payload check rather + * than silently reported as a successful transfer. The staged mode covers the + * transfer when the lane cannot register the memory. */ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { public: static const uint64_t SEED = 0x1234567890abcdeflu; + enum xfer_mode { + XFER_ZCOPY, /* the transport accesses the source memory directly */ + XFER_STAGED /* the source is not registered on the rendezvous lane */ + }; + protected: typedef ucs::handle mem_handle_t; @@ -169,7 +180,9 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { } /* - * Memory domains used by the endpoint's rendezvous zero-copy lanes. + * Memory domains used by the endpoint's rendezvous zero-copy lanes, without + * the domains used for memory type detection: those register their own + * memory type regardless of DMA-BUF support. * * 'remote' returns the destination domains, which is what a GET issued by * this endpoint addresses on its peer, so the indices are in the peer's @@ -180,9 +193,16 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { { const ucp_ep_config_t *ep_config = ucp_ep_config(e.ep()); ucp_context_h context = e.ucph(); + ucp_md_map_t copy_md_map = 0; ucp_md_map_t md_map = 0; ucp_lane_index_t i, lane; ucp_rsc_index_t rsc_index; + ucp_md_index_t md_index; + + for (md_index = 0; md_index < context->num_mem_type_detect_mds; + ++md_index) { + copy_md_map |= UCS_BIT(context->mem_type_detect_mds[md_index]); + } for (i = 0; i < UCP_MAX_LANES; ++i) { lane = ep_config->key.rma_bw_lanes[i]; @@ -198,7 +218,7 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { } } - return md_map; + return md_map & ~copy_md_map; } std::string md_map_str(entity &e, ucp_md_map_t md_map) @@ -219,20 +239,70 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { return names.empty() ? "" : names; } - void test_xfer_from_mem_type(size_t length, bool is_get); + /* + * Protocol selected for a tagged send of 'length' bytes from 'mem_type' + * memory, with its description as reported by UCX_PROTO_INFO. + */ + std::string tag_send_proto(ucs_memory_type_t mem_type, size_t length, + std::string &desc) + { + ucp_worker_h worker = sender().worker(); + ucp_worker_cfg_index_t ep_cfg_index = sender().ep()->cfg_index; + ucp_memory_info_t mem_info = { + .type = static_cast(mem_type), + .sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN, + .flags = UCS_MEM_FLAG_REGISTRABLE + }; + const ucp_proto_threshold_elem_t *thresh; + ucp_proto_select_param_t select_param; + ucp_proto_select_elem_t *select_elem; + ucp_proto_query_attr_t attr; + + ucp_proto_select_param_init(&select_param, UCP_OP_ID_TAG_SEND, 0, 0, + UCP_DATATYPE_CONTIG, &mem_info, 1); + select_elem = ucp_proto_select_lookup_slow( + worker, + &ucs_array_elem(&worker->ep_config, ep_cfg_index).proto_select, + 0, ep_cfg_index, UCP_WORKER_CFG_INDEX_NULL, &select_param); + if (select_elem == NULL) { + desc.clear(); + return ""; + } + + thresh = ucp_proto_thresholds_search_slow(select_elem->thresholds, + length); + ucp_proto_config_query(worker, &thresh->proto_config, length, &attr); + desc = attr.desc; + return thresh->proto_config.proto->name; + } + + /* + * Send 'length' bytes of the tested memory type to a host buffer over a + * rendezvous protocol. 'app_owned' allocates the source without UCP and + * sends it without a memory handle, like an application would. + */ + void test_xfer_from_mem_type(size_t length, bool is_get, xfer_mode mode, + bool app_owned); }; void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, - bool is_get) + bool is_get, + xfer_mode mode, + bool app_owned) { const ucs_memory_type_t send_mem_type = mem_type(); + const char *mem_type_name = + ucs_memory_type_names[send_mem_type]; ucp_request_param_t param; mem_handle_t send_memh, recv_memh; + std::unique_ptr app_buf; + std::string proto_name, proto_desc; void *send_buf, *recv_buf; void *sreq, *rreq; + bool can_reg; if (!mem_buffer::is_mem_type_supported(send_mem_type)) { - UCS_TEST_SKIP_R(std::string(ucs_memory_type_names[send_mem_type]) + + UCS_TEST_SKIP_R(std::string(mem_type_name) + " memory is not supported"); } @@ -246,25 +316,64 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, UCS_TEST_SKIP_R("no rendezvous zero-copy lane is available"); } - if (!(sender().ucph()->reg_md_map[send_mem_type] & lane_md_map)) { + can_reg = sender().ucph()->reg_md_map[send_mem_type] & lane_md_map; + if ((mode == XFER_ZCOPY) && !can_reg) { UCS_TEST_SKIP_R( std::string("the rendezvous zero-copy lane cannot register ") + - ucs_memory_type_names[send_mem_type] + " memory"); + mem_type_name + " memory"); + } else if ((mode == XFER_STAGED) && can_reg) { + UCS_TEST_SKIP_R( + std::string("the rendezvous zero-copy lane can register ") + + mem_type_name + " memory, nothing to stage"); } - send_buf = alloc_mem(sender(), length, send_mem_type, send_memh); - recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); + /* A forced zero-copy scheme can fall back to eager or to AM rendezvous when + * it cannot access the source, so check the selected protocol as well as + * the payload */ + proto_name = tag_send_proto(send_mem_type, length, proto_desc); + UCS_TEST_MESSAGE << "protocol: " << proto_name << ", " << proto_desc; + EXPECT_EQ("tag/rndv", proto_name) << "no rendezvous protocol is selected"; + if (mode == XFER_ZCOPY) { + EXPECT_NE(std::string::npos, proto_desc.find("zero-copy")) + << "the source memory is not accessed by the transport"; + } else { + EXPECT_EQ(std::string::npos, proto_desc.find("zero-copy")) + << "the source memory is accessed by the transport"; + /* CPU-accessible memory is copied by the CPU instead of being staged */ + if (!(UCS_BIT(send_mem_type) & UCS_MEMORY_TYPES_CPU_ACCESSIBLE)) { + EXPECT_NE(std::string::npos, proto_desc.find("frag host")) + << "device memory is not staged through host fragments"; + } + } - UCS_TEST_MESSAGE << (is_get ? "get" : "put") << " lane mds: " - << md_map_str(sender(), lane_md_map) - << " | source memh mds: " - << md_map_str(sender(), send_memh->md_map); + if (app_owned) { + app_buf.reset(new mem_buffer(length, send_mem_type)); + send_buf = app_buf->ptr(); + } else { + send_buf = alloc_mem(sender(), length, send_mem_type, send_memh); + } + recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); - EXPECT_NE(0u, send_memh->md_map & lane_md_map) - << ucs_memory_type_names[send_mem_type] - << " memory is not registered on the memory domain used by the " - "rendezvous zero-copy lane, the payload would be staged instead " - "of accessed by the transport"; + if (!app_owned) { + UCS_TEST_MESSAGE << (is_get ? "get" : "put") << " lane mds: " + << md_map_str(sender(), lane_md_map) + << " | source memh mds: " + << md_map_str(sender(), send_memh->md_map); + + if (mode == XFER_ZCOPY) { + EXPECT_NE(0u, send_memh->md_map & lane_md_map) + << mem_type_name + << " memory is not registered on the memory domain used by " + "the rendezvous zero-copy lane, the payload would be " + "staged instead of accessed by the transport"; + } else { + EXPECT_EQ(0u, send_memh->md_map & lane_md_map) + << mem_type_name + << " memory is registered on the memory domain used by the " + "rendezvous zero-copy lane, the payload would be " + "accessed by the transport instead of staged"; + } + } mem_buffer::pattern_fill(send_buf, length, SEED, send_mem_type); mem_buffer::pattern_fill(recv_buf, length, 0, UCS_MEMORY_TYPE_HOST); @@ -274,8 +383,12 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf, length, 1, 1, ¶m); - param.memh = send_memh; - sreq = ucp_tag_send_nbx(sender().ep(), send_buf, length, 1, ¶m); + if (app_owned) { + param.op_attr_mask = 0; + } else { + param.memh = send_memh; + } + sreq = ucp_tag_send_nbx(sender().ep(), send_buf, length, 1, ¶m); EXPECT_UCS_OK(request_wait(sreq)); EXPECT_UCS_OK(request_wait(rreq)); @@ -286,19 +399,54 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, get_zcopy, "RNDV_THRESH=0", "RNDV_SCHEME=get_zcopy") { - test_xfer_from_mem_type(4 * UCS_MBYTE, true); + test_xfer_from_mem_type(4 * UCS_MBYTE, true, XFER_ZCOPY, false); } UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, put_zcopy, "RNDV_THRESH=0", "RNDV_SCHEME=put_zcopy") { - test_xfer_from_mem_type(4 * UCS_MBYTE, false); + test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_ZCOPY, false); } UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_rndv_zcopy, rcx, "rc_x") UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcx_ze, "rc_x,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcv_ze, + "rc_v,ze_copy") + + +/* + * With DMA-BUF export disabled, the rendezvous zero-copy lane cannot register + * Level Zero memory and the payload must be copied through host memory: device + * memory is staged through fragments, CPU-accessible memory is copied in and + * out. The rendezvous scheme is left to UCP: forcing a zero-copy scheme would + * exclude these protocols. + */ +class test_ucp_mem_type_rndv_staged : public test_ucp_mem_type_rndv_zcopy { +public: + test_ucp_mem_type_rndv_staged() + { + m_env.push_back(new ucs::scoped_setenv("UCX_ZE_COPY_DMABUF", "n")); + } +}; + +UCS_TEST_P(test_ucp_mem_type_rndv_staged, staged, "RNDV_THRESH=0", + "RNDV_FRAG_MEM_TYPES=host") +{ + test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_STAGED, false); +} + +UCS_TEST_P(test_ucp_mem_type_rndv_staged, staged_app_mem, "RNDV_THRESH=0", + "RNDV_FRAG_MEM_TYPES=host") +{ + test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_STAGED, true); +} + +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_staged, rcx_ze, + "rc_x,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_staged, rcv_ze, + "rc_v,ze_copy") class test_ucp_cuda : public ucp_test { public: From ac618d10118842ee41f7cc83fb7c6bc43dc255ec Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 9 Oct 2026 14:38:50 -0700 Subject: [PATCH 03/13] UCT/ZE/COPY: Unwind MD open through cleanup labels Also drop the IB comment which described removed code. Signed-off-by: Yaser Afshar --- src/uct/ib/base/ib_md.c | 3 --- src/uct/ze/copy/ze_copy_md.c | 24 +++++++++++++++--------- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c index 150bda0cdc9..a1e1e2841f6 100644 --- a/src/uct/ib/base/ib_md.c +++ b/src/uct/ib/base/ib_md.c @@ -1455,9 +1455,6 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md, UCS_MEMORY_TYPE_GAUDI); uct_ib_check_gpudirect_driver(md, "/dev/hl0", UCS_MEMORY_TYPE_GAUDI); - /* Do not infer Level Zero device-memory registration support from the - * presence of the Xe module. DMA-BUF support is detected below. */ - /* Check for dma-buf support */ uct_ib_md_check_dmabuf(md); } diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index 5f2df629bf1..24ee4145a76 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -361,6 +361,7 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, ze_result_t ret; const uct_ze_subdevice_t *subdevice; int dmabuf_supported; + ucs_status_t status; ze_driver = uct_ze_base_get_driver(); if (ze_driver == NULL) { @@ -377,32 +378,31 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, subdevice = uct_ze_base_get_subdevice_by_global_id(config->device_ordinal); if (subdevice == NULL) { ucs_error("Failed to get sub-device at ordinal %d", config->device_ordinal); - ucs_free(md); - return UCS_ERR_NO_DEVICE; + status = UCS_ERR_NO_DEVICE; + goto err_free_md; } /* Get the actual device handle from the sub-device */ md->ze_device = uct_ze_base_get_device_handle_from_subdevice(subdevice); if (md->ze_device == NULL) { ucs_error("Failed to get device handle for sub-device %d", config->device_ordinal); - ucs_free(md); - return UCS_ERR_NO_DEVICE; + status = UCS_ERR_NO_DEVICE; + goto err_free_md; } ret = zeContextCreate(ze_driver, &context_desc, &md->ze_context); if (ret != ZE_RESULT_SUCCESS) { ucs_error("zeContextCreate failed with error %x", ret); - ucs_free(md); - return UCS_ERR_NO_DEVICE; + status = UCS_ERR_NO_DEVICE; + goto err_free_md; } dmabuf_supported = uct_ze_copy_md_is_dmabuf_supported(md->ze_device); if ((config->enable_dmabuf == UCS_YES) && !dmabuf_supported) { ucs_error("%s: DMA-BUF is not supported by the Level Zero device", md_name); - zeContextDestroy(md->ze_context); - ucs_free(md); - return UCS_ERR_UNSUPPORTED; + status = UCS_ERR_UNSUPPORTED; + goto err_destroy_context; } md->enable_dmabuf = (config->enable_dmabuf != UCS_NO) && dmabuf_supported; @@ -411,6 +411,12 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, *md_p = (uct_md_h)md; return UCS_OK; + +err_destroy_context: + zeContextDestroy(md->ze_context); +err_free_md: + ucs_free(md); + return status; } uct_component_t uct_ze_copy_component = { From 4139ae87ce3c1e60d4c4206c9ac1fabb945ca6d9 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Sat, 10 Oct 2026 06:23:21 -0700 Subject: [PATCH 04/13] UCT/ZE/COPY: Limit the staged rendezvous test to Level Zero memory Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index ca4e9ecad99..a0ca59b1d4e 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -118,7 +118,7 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_alloc_before_init, all, "all") */ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { public: - static const uint64_t SEED = 0x1234567890abcdeflu; + static constexpr uint64_t SEED = 0x1234567890abcdeflu; enum xfer_mode { XFER_ZCOPY, /* the transport accesses the source memory directly */ @@ -126,7 +126,7 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { }; protected: - typedef ucs::handle mem_handle_t; + using mem_handle_t = ucs::handle; static void unmap_memh(ucp_mem_h memh, ucp_context_h context) { @@ -430,6 +430,17 @@ class test_ucp_mem_type_rndv_staged : public test_ucp_mem_type_rndv_zcopy { { m_env.push_back(new ucs::scoped_setenv("UCX_ZE_COPY_DMABUF", "n")); } + + /* Disabling DMA-BUF export changes how only Level Zero memory is reached */ + static void get_test_variants(std::vector &variants) + { + static constexpr uint64_t ze_mem_types = + UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | + UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) | + UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED); + + add_variant_memtypes(variants, get_test_variants_base, ze_mem_types); + } }; UCS_TEST_P(test_ucp_mem_type_rndv_staged, staged, "RNDV_THRESH=0", From 9b5c28255ca5d094ff9591c20e25e1c4668a05ad Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Mon, 19 Oct 2026 12:45:02 -0700 Subject: [PATCH 05/13] UCT/ZE/COPY: Test zero-copy from a pool sub-allocated buffer Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index a0ca59b1d4e..fdfa3ed2015 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -295,7 +295,7 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, const char *mem_type_name = ucs_memory_type_names[send_mem_type]; ucp_request_param_t param; - mem_handle_t send_memh, recv_memh; + mem_handle_t pad_memh, send_memh, recv_memh; std::unique_ptr app_buf; std::string proto_name, proto_desc; void *send_buf, *recv_buf; @@ -351,6 +351,11 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, app_buf.reset(new mem_buffer(length, send_mem_type)); send_buf = app_buf->ptr(); } else { + /* Keep another buffer of the same size allocated, so that a driver + * which sub-allocates both from one memory pool places the source at a + * nonzero offset in the pool, where a DMA-BUF exported for an + * allocation not created for export does not describe the source */ + alloc_mem(sender(), length, send_mem_type, pad_memh); send_buf = alloc_mem(sender(), length, send_mem_type, send_memh); } recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); @@ -397,16 +402,19 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, mem_buffer::pattern_check(recv_buf, length, SEED, UCS_MEMORY_TYPE_HOST); } +/* Small enough for a driver to sub-allocate from a memory pool */ +static constexpr size_t RNDV_ZCOPY_LENGTH = 64 * UCS_KBYTE; + UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, get_zcopy, "RNDV_THRESH=0", "RNDV_SCHEME=get_zcopy") { - test_xfer_from_mem_type(4 * UCS_MBYTE, true, XFER_ZCOPY, false); + test_xfer_from_mem_type(RNDV_ZCOPY_LENGTH, true, XFER_ZCOPY, false); } UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, put_zcopy, "RNDV_THRESH=0", "RNDV_SCHEME=put_zcopy") { - test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_ZCOPY, false); + test_xfer_from_mem_type(RNDV_ZCOPY_LENGTH, false, XFER_ZCOPY, false); } UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_rndv_zcopy, rcx, From 3ed6aa8150511edb78c89a6b3bf2dd978366b794 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 30 Oct 2026 11:23:10 -0700 Subject: [PATCH 06/13] UCT/ZE/COPY: Verify rendezvous protocols used by tests Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 206 ++++++++++++++++++++-------- 1 file changed, 152 insertions(+), 54 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index fdfa3ed2015..cec2cdfc5a2 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -8,15 +8,16 @@ #include "ucp_test.h" #include +#include +#include #include extern "C" { #include #include #include -#include -#include -#include +#include +#include #include #include } @@ -138,6 +139,8 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { void init() override { + /* Count protocol selections, which used_protos() reports */ + modify_config("PROTO_INFO", "used"); test_ucp_mem_type::init(); sender().connect(&receiver(), get_ep_params()); receiver().connect(&sender(), get_ep_params()); @@ -240,41 +243,102 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { return names.empty() ? "" : names; } + struct proto_selections_t { + ucp_worker_h worker; + ucs_memory_type_t remote_mem_type; /* UNKNOWN without a remote key */ + unsigned selections; + }; + + using proto_counts_t = std::map; + + struct used_proto_t { + std::string name; + std::string desc; + bool is_estimation; /* desc predicts the peer's part */ + ucs_memory_type_t mem_type; /* Local buffer */ + ucs_memory_type_t remote_mem_type; /* Remote key, or UNKNOWN */ + }; + + static void add_selections(ucp_worker_h worker, + const ucp_proto_select_t *proto_select, + ucs_memory_type_t remote_mem_type, + proto_counts_t &counts) + { + const ucp_proto_threshold_elem_t *thresh; + khiter_t khiter; + + for (khiter = kh_begin(proto_select->hash); + khiter != kh_end(proto_select->hash); ++khiter) { + if (!kh_exist(proto_select->hash, khiter)) { + continue; + } + + thresh = kh_val(proto_select->hash, khiter).thresholds; + do { + counts[&thresh->proto_config] = { + worker, remote_mem_type, thresh->proto_config.selections + }; + } while ((thresh++)->max_msg_length < SIZE_MAX); + } + } + /* - * Protocol selected for a tagged send of 'length' bytes from 'mem_type' - * memory, with its description as reported by UCX_PROTO_INFO. + * Selection counters of every protocol configured on the sender and the + * receiver, which count requests only with PROTO_INFO=used. */ - std::string tag_send_proto(ucs_memory_type_t mem_type, size_t length, - std::string &desc) + proto_counts_t proto_selections() { - ucp_worker_h worker = sender().worker(); - ucp_worker_cfg_index_t ep_cfg_index = sender().ep()->cfg_index; - ucp_memory_info_t mem_info = { - .type = static_cast(mem_type), - .sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN, - .flags = UCS_MEM_FLAG_REGISTRABLE - }; - const ucp_proto_threshold_elem_t *thresh; - ucp_proto_select_param_t select_param; - ucp_proto_select_elem_t *select_elem; + proto_counts_t counts; + ucp_rkey_config_t **rkey_config_p; + ucp_ep_config_t *ep_config; + + for (entity *e : {&sender(), &receiver()}) { + ucp_worker_h worker = e->worker(); + + ucs_array_for_each(ep_config, &worker->ep_config) { + add_selections(worker, &ep_config->proto_select, + UCS_MEMORY_TYPE_UNKNOWN, counts); + } + + ucs_array_for_each(rkey_config_p, &worker->rkey_config) { + add_selections(worker, &(*rkey_config_p)->proto_select, + (*rkey_config_p)->key.mem_type, counts); + } + } + + return counts; + } + + /* + * Protocols which started a request since 'before' was taken, including + * the rendezvous protocols which move the data after the tag/rndv request + * is matched. + */ + std::vector + used_protos(const proto_counts_t &before, size_t length) + { + std::vector used; ucp_proto_query_attr_t attr; - ucp_proto_select_param_init(&select_param, UCP_OP_ID_TAG_SEND, 0, 0, - UCP_DATATYPE_CONTIG, &mem_info, 1); - select_elem = ucp_proto_select_lookup_slow( - worker, - &ucs_array_elem(&worker->ep_config, ep_cfg_index).proto_select, - 0, ep_cfg_index, UCP_WORKER_CFG_INDEX_NULL, &select_param); - if (select_elem == NULL) { - desc.clear(); - return ""; + for (const auto &count : proto_selections()) { + auto prev = before.find(count.first); + unsigned prev_sels = (prev == before.end()) ? + 0 : prev->second.selections; + if (count.second.selections <= prev_sels) { + continue; + } + + ucp_proto_config_query(count.second.worker, count.first, length, + &attr); + used.push_back({count.first->proto->name, attr.desc, + static_cast(attr.is_estimation), + static_cast( + count.first->select_param.mem_type), + count.second.remote_mem_type}); } - thresh = ucp_proto_thresholds_search_slow(select_elem->thresholds, - length); - ucp_proto_config_query(worker, &thresh->proto_config, length, &attr); - desc = attr.desc; - return thresh->proto_config.proto->name; + return used; } /* @@ -297,7 +361,7 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, ucp_request_param_t param; mem_handle_t pad_memh, send_memh, recv_memh; std::unique_ptr app_buf; - std::string proto_name, proto_desc; + proto_counts_t sels_before; void *send_buf, *recv_buf; void *sreq, *rreq; bool can_reg; @@ -328,25 +392,6 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, mem_type_name + " memory, nothing to stage"); } - /* A forced zero-copy scheme can fall back to eager or to AM rendezvous when - * it cannot access the source, so check the selected protocol as well as - * the payload */ - proto_name = tag_send_proto(send_mem_type, length, proto_desc); - UCS_TEST_MESSAGE << "protocol: " << proto_name << ", " << proto_desc; - EXPECT_EQ("tag/rndv", proto_name) << "no rendezvous protocol is selected"; - if (mode == XFER_ZCOPY) { - EXPECT_NE(std::string::npos, proto_desc.find("zero-copy")) - << "the source memory is not accessed by the transport"; - } else { - EXPECT_EQ(std::string::npos, proto_desc.find("zero-copy")) - << "the source memory is accessed by the transport"; - /* CPU-accessible memory is copied by the CPU instead of being staged */ - if (!(UCS_BIT(send_mem_type) & UCS_MEMORY_TYPES_CPU_ACCESSIBLE)) { - EXPECT_NE(std::string::npos, proto_desc.find("frag host")) - << "device memory is not staged through host fragments"; - } - } - if (app_owned) { app_buf.reset(new mem_buffer(length, send_mem_type)); send_buf = app_buf->ptr(); @@ -384,6 +429,8 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, mem_buffer::pattern_fill(send_buf, length, SEED, send_mem_type); mem_buffer::pattern_fill(recv_buf, length, 0, UCS_MEMORY_TYPE_HOST); + sels_before = proto_selections(); + param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; param.memh = recv_memh; rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf, length, 1, 1, @@ -400,6 +447,57 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, EXPECT_UCS_OK(request_wait(rreq)); mem_buffer::pattern_check(recv_buf, length, SEED, UCS_MEMORY_TYPE_HOST); + + /* A forced zero-copy scheme can fall back to eager or to AM rendezvous when + * it cannot access the source, so check the protocols which moved the data + * as well as the payload */ + std::vector used = used_protos(sels_before, length); + std::string used_str; + for (const used_proto_t &proto : used) { + used_str += std::string(" [") + proto.name + ": " + + (proto.is_estimation ? "(?) " : "") + proto.desc + + ", local " + ucs_memory_type_names[proto.mem_type] + + ", remote " + ucs_memory_type_names[proto.remote_mem_type] + + "]"; + } + UCS_TEST_MESSAGE << "used protocols:" << used_str; + + /* A zero-copy protocol on the source memory: a PUT or an AM send reads + * its local buffer, a GET issued by the receiver reads the remote one */ + auto is_used_on_source = [&](const std::string &name) { + return std::any_of(used.begin(), used.end(), + [&](const used_proto_t &proto) { + return (proto.name == name) && + ((proto.mem_type == send_mem_type) || + (proto.remote_mem_type == + send_mem_type)); + }); + }; + + if (mode == XFER_ZCOPY) { + EXPECT_TRUE(is_used_on_source(is_get ? "rndv/get/zcopy" : + "rndv/put/zcopy")) + << "the source memory is not accessed by the transport"; + } else { + EXPECT_TRUE(is_used_on_source("tag/rndv")) + << "the source memory is not sent by a rendezvous protocol"; + EXPECT_FALSE(is_used_on_source("rndv/get/zcopy") || + is_used_on_source("rndv/put/zcopy") || + is_used_on_source("rndv/am/zcopy")) + << "the source memory is accessed by the transport"; + /* CPU-accessible memory may be copied by the CPU instead of being + * staged. The rendezvous request and reply only estimate the peer's + * protocol, so skip estimated descriptions. */ + if (!(UCS_BIT(send_mem_type) & UCS_MEMORY_TYPES_CPU_ACCESSIBLE)) { + EXPECT_TRUE(std::any_of(used.begin(), used.end(), + [](const used_proto_t &proto) { + return !proto.is_estimation && + (proto.desc.find("frag host") != + std::string::npos); + })) + << "device memory is not staged through host fragments"; + } + } } /* Small enough for a driver to sub-allocate from a memory pool */ @@ -428,9 +526,9 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcv_ze, /* * With DMA-BUF export disabled, the rendezvous zero-copy lane cannot register * Level Zero memory and the payload must be copied through host memory: device - * memory is staged through fragments, CPU-accessible memory is copied in and - * out. The rendezvous scheme is left to UCP: forcing a zero-copy scheme would - * exclude these protocols. + * memory is staged through fragments, CPU-accessible memory is staged or copied + * in and out. The rendezvous scheme is left to UCP: forcing a zero-copy scheme + * would exclude these protocols. */ class test_ucp_mem_type_rndv_staged : public test_ucp_mem_type_rndv_zcopy { public: From 2cf346f483cfa3f2844a1f48c36d2f3a3aefc001 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 30 Oct 2026 11:58:45 -0700 Subject: [PATCH 07/13] UCT/ZE/COPY: Trim rendezvous test helpers and fix MD index space Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 208 ++++++++++------------------ 1 file changed, 71 insertions(+), 137 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index cec2cdfc5a2..a7b91d6385c 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -129,14 +129,6 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { protected: using mem_handle_t = ucs::handle; - static void unmap_memh(ucp_mem_h memh, ucp_context_h context) - { - ucs_status_t status = ucp_mem_unmap(context, memh); - if (status != UCS_OK) { - ucs_warn("failed to unmap memory: %s", ucs_status_string(status)); - } - } - void init() override { /* Count protocol selections, which used_protos() reports */ @@ -151,7 +143,6 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { mem_handle_t &memh_handle) { ucp_mem_map_params_t params; - ucp_mem_attr_t attr; ucp_mem_h memh; ucs_status_t status; @@ -171,16 +162,10 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { ucs_status_string(status)); } - memh_handle.reset(memh, unmap_memh, e.ucph()); - - attr.field_mask = UCP_MEM_ATTR_FIELD_ADDRESS; - status = ucp_mem_query(memh, &attr); - if (status != UCS_OK) { - UCS_TEST_ABORT("ucp_mem_query() failed: " + - std::string(ucs_status_string(status))); - } - - return attr.address; + memh_handle.reset(memh, [](ucp_mem_h handle, ucp_context_h context) { + EXPECT_UCS_OK(ucp_mem_unmap(context, handle)); + }, e.ucph()); + return ucp_memh_address(memh); } /* @@ -188,15 +173,16 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { * the domains used for memory type detection: those register their own * memory type regardless of DMA-BUF support. * - * 'remote' returns the destination domains, which is what a GET issued by - * this endpoint addresses on its peer, so the indices are in the peer's - * domain space. Otherwise the endpoint's own domains are returned, which is - * what a PUT issued by this endpoint registers locally. + * The indices are in the domain space of 'md_entity'. When it is the peer + * of 'e', the destination domains are returned, which is what a GET issued + * by 'e' addresses on its peer. Otherwise the endpoint's own domains are + * returned, which is what a PUT issued by 'e' registers locally. */ - ucp_md_map_t rndv_zcopy_md_map(entity &e, bool remote) + ucp_md_map_t rndv_zcopy_md_map(entity &e, entity &md_entity) { const ucp_ep_config_t *ep_config = ucp_ep_config(e.ep()); - ucp_context_h context = e.ucph(); + const bool remote = &md_entity != &e; + ucp_context_h context = md_entity.ucph(); ucp_md_map_t copy_md_map = 0; ucp_md_map_t md_map = 0; ucp_lane_index_t i, lane; @@ -225,117 +211,79 @@ class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { return md_map & ~copy_md_map; } - std::string md_map_str(entity &e, ucp_md_map_t md_map) - { - ucp_context_h context = e.ucph(); - std::string names; - ucp_md_index_t md_index; - - ucs_for_each_bit(md_index, md_map) { - if (!names.empty()) { - names += ","; - } - names += (md_index < context->num_mds) ? - context->tl_mds[md_index].rsc.md_name : - ""; - } - - return names.empty() ? "" : names; - } - - struct proto_selections_t { + /* Selection count of a protocol configuration, which counts requests only + * with PROTO_INFO=used */ + struct proto_count_t { ucp_worker_h worker; ucs_memory_type_t remote_mem_type; /* UNKNOWN without a remote key */ unsigned selections; }; - using proto_counts_t = std::map; + using proto_counts_t = std::map; struct used_proto_t { - std::string name; - std::string desc; - bool is_estimation; /* desc predicts the peer's part */ - ucs_memory_type_t mem_type; /* Local buffer */ - ucs_memory_type_t remote_mem_type; /* Remote key, or UNKNOWN */ + const ucp_proto_config_t *config; + ucs_memory_type_t remote_mem_type; + ucp_proto_query_attr_t attr; }; - static void add_selections(ucp_worker_h worker, - const ucp_proto_select_t *proto_select, - ucs_memory_type_t remote_mem_type, - proto_counts_t &counts) + static void add_counts(ucp_worker_h worker, + const ucp_proto_select_t *select, + ucs_memory_type_t remote_mem_type, + proto_counts_t &counts) { const ucp_proto_threshold_elem_t *thresh; - khiter_t khiter; - - for (khiter = kh_begin(proto_select->hash); - khiter != kh_end(proto_select->hash); ++khiter) { - if (!kh_exist(proto_select->hash, khiter)) { - continue; + khiter_t iter; + + for (iter = kh_begin(select->hash); iter != kh_end(select->hash); + ++iter) { + if (kh_exist(select->hash, iter)) { + thresh = kh_val(select->hash, iter).thresholds; + do { + counts[&thresh->proto_config] = { + worker, remote_mem_type, + thresh->proto_config.selections}; + } while ((thresh++)->max_msg_length < SIZE_MAX); } - - thresh = kh_val(proto_select->hash, khiter).thresholds; - do { - counts[&thresh->proto_config] = { - worker, remote_mem_type, thresh->proto_config.selections - }; - } while ((thresh++)->max_msg_length < SIZE_MAX); } } - /* - * Selection counters of every protocol configured on the sender and the - * receiver, which count requests only with PROTO_INFO=used. - */ - proto_counts_t proto_selections() + /* Selection counts of the sender's and the receiver's protocols */ + proto_counts_t proto_counts() { proto_counts_t counts; ucp_rkey_config_t **rkey_config_p; ucp_ep_config_t *ep_config; for (entity *e : {&sender(), &receiver()}) { - ucp_worker_h worker = e->worker(); - - ucs_array_for_each(ep_config, &worker->ep_config) { - add_selections(worker, &ep_config->proto_select, - UCS_MEMORY_TYPE_UNKNOWN, counts); + ucs_array_for_each(ep_config, &e->worker()->ep_config) { + add_counts(e->worker(), &ep_config->proto_select, + UCS_MEMORY_TYPE_UNKNOWN, counts); } - - ucs_array_for_each(rkey_config_p, &worker->rkey_config) { - add_selections(worker, &(*rkey_config_p)->proto_select, - (*rkey_config_p)->key.mem_type, counts); + ucs_array_for_each(rkey_config_p, &e->worker()->rkey_config) { + add_counts(e->worker(), &(*rkey_config_p)->proto_select, + (*rkey_config_p)->key.mem_type, counts); } } return counts; } - /* - * Protocols which started a request since 'before' was taken, including - * the rendezvous protocols which move the data after the tag/rndv request - * is matched. - */ + /* Protocols which started a request since 'before', including the + * rendezvous protocols which the peers select after tag/rndv */ std::vector used_protos(const proto_counts_t &before, size_t length) { std::vector used; - ucp_proto_query_attr_t attr; - - for (const auto &count : proto_selections()) { - auto prev = before.find(count.first); - unsigned prev_sels = (prev == before.end()) ? - 0 : prev->second.selections; - if (count.second.selections <= prev_sels) { - continue; - } - ucp_proto_config_query(count.second.worker, count.first, length, - &attr); - used.push_back({count.first->proto->name, attr.desc, - static_cast(attr.is_estimation), - static_cast( - count.first->select_param.mem_type), - count.second.remote_mem_type}); + for (const auto &count : proto_counts()) { + auto prev = before.find(count.first); + if (count.second.selections > + ((prev == before.end()) ? 0 : prev->second.selections)) { + used.push_back({count.first, count.second.remote_mem_type}); + ucp_proto_config_query(count.second.worker, count.first, + length, &used.back().attr); + } } return used; @@ -361,7 +309,7 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, ucp_request_param_t param; mem_handle_t pad_memh, send_memh, recv_memh; std::unique_ptr app_buf; - proto_counts_t sels_before; + proto_counts_t counts_before; void *send_buf, *recv_buf; void *sreq, *rreq; bool can_reg; @@ -376,20 +324,16 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, * that skipping does not leak the memory mappings. */ entity &op_entity = is_get ? receiver() : sender(); - ucp_md_map_t lane_md_map = rndv_zcopy_md_map(op_entity, is_get); + ucp_md_map_t lane_md_map = rndv_zcopy_md_map(op_entity, sender()); if (lane_md_map == 0) { UCS_TEST_SKIP_R("no rendezvous zero-copy lane is available"); } can_reg = sender().ucph()->reg_md_map[send_mem_type] & lane_md_map; - if ((mode == XFER_ZCOPY) && !can_reg) { - UCS_TEST_SKIP_R( - std::string("the rendezvous zero-copy lane cannot register ") + - mem_type_name + " memory"); - } else if ((mode == XFER_STAGED) && can_reg) { - UCS_TEST_SKIP_R( - std::string("the rendezvous zero-copy lane can register ") + - mem_type_name + " memory, nothing to stage"); + if (can_reg != (mode == XFER_ZCOPY)) { + UCS_TEST_SKIP_R(std::string("the rendezvous zero-copy lane ") + + (can_reg ? "can" : "cannot") + " register " + + mem_type_name + " memory"); } if (app_owned) { @@ -406,30 +350,21 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); if (!app_owned) { - UCS_TEST_MESSAGE << (is_get ? "get" : "put") << " lane mds: " - << md_map_str(sender(), lane_md_map) - << " | source memh mds: " - << md_map_str(sender(), send_memh->md_map); - if (mode == XFER_ZCOPY) { EXPECT_NE(0u, send_memh->md_map & lane_md_map) - << mem_type_name - << " memory is not registered on the memory domain used by " - "the rendezvous zero-copy lane, the payload would be " - "staged instead of accessed by the transport"; + << mem_type_name << " memory is not registered on the " + << "rendezvous zero-copy lane's memory domain"; } else { EXPECT_EQ(0u, send_memh->md_map & lane_md_map) - << mem_type_name - << " memory is registered on the memory domain used by the " - "rendezvous zero-copy lane, the payload would be " - "accessed by the transport instead of staged"; + << mem_type_name << " memory is registered on the " + << "rendezvous zero-copy lane's memory domain"; } } mem_buffer::pattern_fill(send_buf, length, SEED, send_mem_type); mem_buffer::pattern_fill(recv_buf, length, 0, UCS_MEMORY_TYPE_HOST); - sels_before = proto_selections(); + counts_before = proto_counts(); param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; param.memh = recv_memh; @@ -451,13 +386,11 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, /* A forced zero-copy scheme can fall back to eager or to AM rendezvous when * it cannot access the source, so check the protocols which moved the data * as well as the payload */ - std::vector used = used_protos(sels_before, length); + std::vector used = used_protos(counts_before, length); std::string used_str; for (const used_proto_t &proto : used) { - used_str += std::string(" [") + proto.name + ": " + - (proto.is_estimation ? "(?) " : "") + proto.desc + - ", local " + ucs_memory_type_names[proto.mem_type] + - ", remote " + ucs_memory_type_names[proto.remote_mem_type] + + used_str += std::string(" [") + proto.config->proto->name + ": " + + (proto.attr.is_estimation ? "(?) " : "") + proto.attr.desc + "]"; } UCS_TEST_MESSAGE << "used protocols:" << used_str; @@ -467,8 +400,9 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, auto is_used_on_source = [&](const std::string &name) { return std::any_of(used.begin(), used.end(), [&](const used_proto_t &proto) { - return (proto.name == name) && - ((proto.mem_type == send_mem_type) || + return (name == proto.config->proto->name) && + ((proto.config->select_param.mem_type == + send_mem_type) || (proto.remote_mem_type == send_mem_type)); }); @@ -491,9 +425,9 @@ void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, if (!(UCS_BIT(send_mem_type) & UCS_MEMORY_TYPES_CPU_ACCESSIBLE)) { EXPECT_TRUE(std::any_of(used.begin(), used.end(), [](const used_proto_t &proto) { - return !proto.is_estimation && - (proto.desc.find("frag host") != - std::string::npos); + return !proto.attr.is_estimation && + strstr(proto.attr.desc, + "frag host"); })) << "device memory is not staged through host fragments"; } From 65561c7bb4b975144a01e773ac6f99b085298783 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 30 Oct 2026 13:58:03 -0700 Subject: [PATCH 08/13] UCT/ZE/COPY: Set ZE_COPY_DMABUF through the test config helpers Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index a7b91d6385c..8eccd816929 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -466,9 +466,10 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcv_ze, */ class test_ucp_mem_type_rndv_staged : public test_ucp_mem_type_rndv_zcopy { public: - test_ucp_mem_type_rndv_staged() + void init() override { - m_env.push_back(new ucs::scoped_setenv("UCX_ZE_COPY_DMABUF", "n")); + modify_config("ZE_COPY_DMABUF", "n", SETENV_IF_NOT_EXIST); + test_ucp_mem_type_rndv_zcopy::init(); } /* Disabling DMA-BUF export changes how only Level Zero memory is reached */ From b402de777b268ad458d1f2b944ab414cccc614ba Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Wed, 4 Nov 2026 09:09:54 -0800 Subject: [PATCH 09/13] UCT/ZE/COPY: Address review: check DMA-BUF before context, rename flag Signed-off-by: Yaser Afshar --- src/uct/ze/copy/ze_copy_md.c | 30 ++++++++++++++---------------- src/uct/ze/copy/ze_copy_md.h | 2 +- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index af2e27c1001..aa8fec9d2dc 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -93,7 +93,7 @@ static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr) md_attr->detect_mem_types = UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) | UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED); - md_attr->dmabuf_mem_types = ze_md->enable_dmabuf ? + md_attr->dmabuf_mem_types = ze_md->dmabuf_supported ? UCT_ZE_COPY_MD_DMABUF_MEM_TYPES : 0; md_attr->max_alloc = SIZE_MAX; return UCS_OK; @@ -125,7 +125,7 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p, size_t alignment = ucs_get_page_size(); ucs_status_t status; - if (md->enable_dmabuf && + if (md->dmabuf_supported && (UCS_BIT(mem_type) & UCT_ZE_COPY_MD_DMABUF_MEM_TYPES)) { export_desc_p = &export_desc; } @@ -261,9 +261,8 @@ static ucs_status_t uct_ze_copy_md_mem_query(uct_md_h md, const void *addr, ucs_status_t status; if ((mem_attr_p->field_mask & UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD) && - ze_md->enable_dmabuf) { - mem_attr_p->dmabuf_fd = UCT_DMABUF_FD_INVALID; - dmabuf_fd_p = &dmabuf_fd; + ze_md->dmabuf_supported) { + dmabuf_fd_p = &dmabuf_fd; } status = uct_ze_copy_md_query_attributes(md, addr, length, &mem_info, @@ -390,22 +389,23 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, goto err_free_md; } - ret = zeContextCreate(ze_driver, &context_desc, &md->ze_context); - if (ret != ZE_RESULT_SUCCESS) { - ucs_error("zeContextCreate failed with error %x", ret); - status = UCS_ERR_NO_DEVICE; - goto err_free_md; - } - dmabuf_supported = uct_ze_copy_md_is_dmabuf_supported(md->ze_device); if ((config->enable_dmabuf == UCS_YES) && !dmabuf_supported) { ucs_error("%s: DMA-BUF is not supported by the Level Zero device", md_name); status = UCS_ERR_UNSUPPORTED; - goto err_destroy_context; + goto err_free_md; } - md->enable_dmabuf = (config->enable_dmabuf != UCS_NO) && dmabuf_supported; + ret = zeContextCreate(ze_driver, &context_desc, &md->ze_context); + if (ret != ZE_RESULT_SUCCESS) { + ucs_error("zeContextCreate failed with error %x", ret); + status = UCS_ERR_NO_DEVICE; + goto err_free_md; + } + + md->dmabuf_supported = (config->enable_dmabuf != UCS_NO) && + dmabuf_supported; md->super.ops = &md_ops; md->super.component = &uct_ze_copy_component; @@ -415,8 +415,6 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, ucs_memtype_cache_global_create(); return UCS_OK; -err_destroy_context: - zeContextDestroy(md->ze_context); err_free_md: ucs_free(md); return status; diff --git a/src/uct/ze/copy/ze_copy_md.h b/src/uct/ze/copy/ze_copy_md.h index 46bf6d47755..5187ce4a304 100644 --- a/src/uct/ze/copy/ze_copy_md.h +++ b/src/uct/ze/copy/ze_copy_md.h @@ -21,7 +21,7 @@ typedef struct uct_ze_copy_md { uct_md_t super; /**< Domain info */ ze_context_handle_t ze_context; ze_device_handle_t ze_device; - int enable_dmabuf; /**< Expose memory as DMA-BUF */ + int dmabuf_supported; /**< Expose memory as DMA-BUF */ } uct_ze_copy_md_t; From 5582a585cc498c334eb6ea53481b8e6f6948640a Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Wed, 4 Nov 2026 09:10:14 -0800 Subject: [PATCH 10/13] UCT/ZE/COPY: Address review: check DMA-BUF transfers at UCT Signed-off-by: Yaser Afshar --- test/gtest/ucp/test_ucp_mem_type.cc | 423 +++++----------------------- test/gtest/uct/ib/test_ib_xfer.cc | 167 +++++++++++ 2 files changed, 245 insertions(+), 345 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index 8eccd816929..07fe1fe4ca0 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -8,16 +8,11 @@ #include "ucp_test.h" #include -#include -#include -#include - extern "C" { #include #include #include #include -#include #include #include } @@ -109,396 +104,134 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_alloc_before_init, all, "all") /* - * Moves a payload out of memory of the tested type, allocated by UCP or by the - * application, over a rendezvous protocol. The zero-copy mode covers the - * registration path a network transport uses to reach the memory directly, - * including registration through a DMA-BUF handle: memory which is registered - * but not readable by the transport is detected by the payload check rather - * than silently reported as a successful transfer. The staged mode covers the - * transfer when the lane cannot register the memory. + * ZE_COPY_DMABUF decides whether the network memory domains register Level Zero + * memory through a DMA-BUF. uct_p2p_rma_test_dmabuf checks that the network + * reads the memory correctly through it. */ -class test_ucp_mem_type_rndv_zcopy : public test_ucp_mem_type { +class test_ucp_mem_type_dmabuf : public test_ucp_mem_type { public: static constexpr uint64_t SEED = 0x1234567890abcdeflu; - enum xfer_mode { - XFER_ZCOPY, /* the transport accesses the source memory directly */ - XFER_STAGED /* the source is not registered on the rendezvous lane */ - }; + static void get_test_variants(std::vector &variants) + { + add_variant_memtypes(variants, get_test_variants_base, + UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | + UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE)); + } protected: - using mem_handle_t = ucs::handle; - void init() override { - /* Count protocol selections, which used_protos() reports */ - modify_config("PROTO_INFO", "used"); test_ucp_mem_type::init(); sender().connect(&receiver(), get_ep_params()); - receiver().connect(&sender(), get_ep_params()); - } - - /* Allocate through UCP so that a memory domain performs the allocation */ - void *alloc_mem(entity &e, size_t length, ucs_memory_type_t mem_type, - mem_handle_t &memh_handle) - { - ucp_mem_map_params_t params; - ucp_mem_h memh; - ucs_status_t status; - - params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS | - UCP_MEM_MAP_PARAM_FIELD_LENGTH | - UCP_MEM_MAP_PARAM_FIELD_FLAGS | - UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; - params.address = NULL; - params.length = length; - params.flags = UCP_MEM_MAP_ALLOCATE; - params.memory_type = mem_type; - - status = ucp_mem_map(e.ucph(), ¶ms, &memh); - if (status != UCS_OK) { - UCS_TEST_SKIP_R(std::string("cannot allocate ") + - ucs_memory_type_names[mem_type] + " memory: " + - ucs_status_string(status)); - } - - memh_handle.reset(memh, [](ucp_mem_h handle, ucp_context_h context) { - EXPECT_UCS_OK(ucp_mem_unmap(context, handle)); - }, e.ucph()); - return ucp_memh_address(memh); - } - - /* - * Memory domains used by the endpoint's rendezvous zero-copy lanes, without - * the domains used for memory type detection: those register their own - * memory type regardless of DMA-BUF support. - * - * The indices are in the domain space of 'md_entity'. When it is the peer - * of 'e', the destination domains are returned, which is what a GET issued - * by 'e' addresses on its peer. Otherwise the endpoint's own domains are - * returned, which is what a PUT issued by 'e' registers locally. - */ - ucp_md_map_t rndv_zcopy_md_map(entity &e, entity &md_entity) - { - const ucp_ep_config_t *ep_config = ucp_ep_config(e.ep()); - const bool remote = &md_entity != &e; - ucp_context_h context = md_entity.ucph(); - ucp_md_map_t copy_md_map = 0; - ucp_md_map_t md_map = 0; - ucp_lane_index_t i, lane; - ucp_rsc_index_t rsc_index; - ucp_md_index_t md_index; - - for (md_index = 0; md_index < context->num_mem_type_detect_mds; - ++md_index) { - copy_md_map |= UCS_BIT(context->mem_type_detect_mds[md_index]); - } - - for (i = 0; i < UCP_MAX_LANES; ++i) { - lane = ep_config->key.rma_bw_lanes[i]; - if (lane == UCP_NULL_LANE) { - break; - } - - if (remote) { - md_map |= UCS_BIT(ep_config->key.lanes[lane].dst_md_index); - } else { - rsc_index = ep_config->key.lanes[lane].rsc_index; - md_map |= UCS_BIT(context->tl_rscs[rsc_index].md_index); - } - } - - return md_map & ~copy_md_map; - } - - /* Selection count of a protocol configuration, which counts requests only - * with PROTO_INFO=used */ - struct proto_count_t { - ucp_worker_h worker; - ucs_memory_type_t remote_mem_type; /* UNKNOWN without a remote key */ - unsigned selections; - }; - - using proto_counts_t = std::map; - - struct used_proto_t { - const ucp_proto_config_t *config; - ucs_memory_type_t remote_mem_type; - ucp_proto_query_attr_t attr; - }; - - static void add_counts(ucp_worker_h worker, - const ucp_proto_select_t *select, - ucs_memory_type_t remote_mem_type, - proto_counts_t &counts) - { - const ucp_proto_threshold_elem_t *thresh; - khiter_t iter; - - for (iter = kh_begin(select->hash); iter != kh_end(select->hash); - ++iter) { - if (kh_exist(select->hash, iter)) { - thresh = kh_val(select->hash, iter).thresholds; - do { - counts[&thresh->proto_config] = { - worker, remote_mem_type, - thresh->proto_config.selections}; - } while ((thresh++)->max_msg_length < SIZE_MAX); - } - } - } - - /* Selection counts of the sender's and the receiver's protocols */ - proto_counts_t proto_counts() - { - proto_counts_t counts; - ucp_rkey_config_t **rkey_config_p; - ucp_ep_config_t *ep_config; - - for (entity *e : {&sender(), &receiver()}) { - ucs_array_for_each(ep_config, &e->worker()->ep_config) { - add_counts(e->worker(), &ep_config->proto_select, - UCS_MEMORY_TYPE_UNKNOWN, counts); - } - ucs_array_for_each(rkey_config_p, &e->worker()->rkey_config) { - add_counts(e->worker(), &(*rkey_config_p)->proto_select, - (*rkey_config_p)->key.mem_type, counts); - } - } - - return counts; - } - - /* Protocols which started a request since 'before', including the - * rendezvous protocols which the peers select after tag/rndv */ - std::vector - used_protos(const proto_counts_t &before, size_t length) - { - std::vector used; - - for (const auto &count : proto_counts()) { - auto prev = before.find(count.first); - if (count.second.selections > - ((prev == before.end()) ? 0 : prev->second.selections)) { - used.push_back({count.first, count.second.remote_mem_type}); - ucp_proto_config_query(count.second.worker, count.first, - length, &used.back().attr); - } - } - - return used; } /* - * Send 'length' bytes of the tested memory type to a host buffer over a - * rendezvous protocol. 'app_owned' allocates the source without UCP and - * sends it without a memory handle, like an application would. + * Send memory of the tested type to a host buffer over a rendezvous + * protocol. 'app_owned' allocates the source without UCP and sends it + * without a memory handle, like an application would. */ - void test_xfer_from_mem_type(size_t length, bool is_get, xfer_mode mode, - bool app_owned); + void test_xfer(bool app_owned, bool exp_dmabuf); }; -void test_ucp_mem_type_rndv_zcopy::test_xfer_from_mem_type(size_t length, - bool is_get, - xfer_mode mode, - bool app_owned) +void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) { - const ucs_memory_type_t send_mem_type = mem_type(); - const char *mem_type_name = - ucs_memory_type_names[send_mem_type]; - ucp_request_param_t param; - mem_handle_t pad_memh, send_memh, recv_memh; - std::unique_ptr app_buf; - proto_counts_t counts_before; - void *send_buf, *recv_buf; - void *sreq, *rreq; - bool can_reg; - - if (!mem_buffer::is_mem_type_supported(send_mem_type)) { - UCS_TEST_SKIP_R(std::string(mem_type_name) + - " memory is not supported"); - } - - /* - * Resolve the lanes and the registration capability before allocating, so - * that skipping does not leak the memory mappings. - */ - entity &op_entity = is_get ? receiver() : sender(); - ucp_md_map_t lane_md_map = rndv_zcopy_md_map(op_entity, sender()); - if (lane_md_map == 0) { - UCS_TEST_SKIP_R("no rendezvous zero-copy lane is available"); - } - - can_reg = sender().ucph()->reg_md_map[send_mem_type] & lane_md_map; - if (can_reg != (mode == XFER_ZCOPY)) { - UCS_TEST_SKIP_R(std::string("the rendezvous zero-copy lane ") + - (can_reg ? "can" : "cannot") + " register " + - mem_type_name + " memory"); - } - - if (app_owned) { - app_buf.reset(new mem_buffer(length, send_mem_type)); - send_buf = app_buf->ptr(); - } else { - /* Keep another buffer of the same size allocated, so that a driver - * which sub-allocates both from one memory pool places the source at a - * nonzero offset in the pool, where a DMA-BUF exported for an - * allocation not created for export does not describe the source */ - alloc_mem(sender(), length, send_mem_type, pad_memh); - send_buf = alloc_mem(sender(), length, send_mem_type, send_memh); - } - recv_buf = alloc_mem(receiver(), length, UCS_MEMORY_TYPE_HOST, recv_memh); - - if (!app_owned) { - if (mode == XFER_ZCOPY) { - EXPECT_NE(0u, send_memh->md_map & lane_md_map) - << mem_type_name << " memory is not registered on the " - << "rendezvous zero-copy lane's memory domain"; - } else { - EXPECT_EQ(0u, send_memh->md_map & lane_md_map) - << mem_type_name << " memory is registered on the " - << "rendezvous zero-copy lane's memory domain"; - } - } - - mem_buffer::pattern_fill(send_buf, length, SEED, send_mem_type); - mem_buffer::pattern_fill(recv_buf, length, 0, UCS_MEMORY_TYPE_HOST); - - counts_before = proto_counts(); - - param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; - param.memh = recv_memh; - rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf, length, 1, 1, - ¶m); + static constexpr size_t length = 4 * UCS_MBYTE; + ucp_context_h context = sender().ucph(); + const char *mem_type_name = ucs_memory_type_names[mem_type()]; + mem_buffer recv_buf(length, UCS_MEMORY_TYPE_HOST); + ucs::handle memh_handle; + ucs::auto_ptr app_buf; + ucp_request_param_t param, recv_param; + ucp_mem_map_params_t params; + void *send_buf, *sreq, *rreq; + ucp_mem_h memh; + + EXPECT_EQ(exp_dmabuf, (context->reg_md_map[mem_type()] & + context->dmabuf_reg_md_map) != 0) + << mem_type_name << " memory registration through a DMA-BUF"; if (app_owned) { + app_buf.reset(new mem_buffer(length, mem_type())); + send_buf = app_buf->ptr(); param.op_attr_mask = 0; } else { - param.memh = send_memh; - } + params.field_mask = UCP_MEM_MAP_PARAM_FIELD_LENGTH | + UCP_MEM_MAP_PARAM_FIELD_FLAGS | + UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; + params.length = length; + params.flags = UCP_MEM_MAP_ALLOCATE; + params.memory_type = mem_type(); + ASSERT_UCS_OK(ucp_mem_map(context, ¶ms, &memh)); + memh_handle.reset(memh, [](ucp_mem_h handle, ucp_context_h ctx) { + EXPECT_UCS_OK(ucp_mem_unmap(ctx, handle)); + }, context); + + EXPECT_EQ(exp_dmabuf, (memh->md_map & context->dmabuf_reg_md_map) != 0) + << mem_type_name << " memory handle registration through a " + << "DMA-BUF"; + send_buf = ucp_memh_address(memh); + param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; + param.memh = memh; + } + + mem_buffer::pattern_fill(send_buf, length, SEED, mem_type()); + mem_buffer::pattern_fill(recv_buf.ptr(), length, 0, UCS_MEMORY_TYPE_HOST); + + recv_param.op_attr_mask = 0; + rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf.ptr(), length, 1, 1, + &recv_param); sreq = ucp_tag_send_nbx(sender().ep(), send_buf, length, 1, ¶m); - EXPECT_UCS_OK(request_wait(sreq)); EXPECT_UCS_OK(request_wait(rreq)); - mem_buffer::pattern_check(recv_buf, length, SEED, UCS_MEMORY_TYPE_HOST); - - /* A forced zero-copy scheme can fall back to eager or to AM rendezvous when - * it cannot access the source, so check the protocols which moved the data - * as well as the payload */ - std::vector used = used_protos(counts_before, length); - std::string used_str; - for (const used_proto_t &proto : used) { - used_str += std::string(" [") + proto.config->proto->name + ": " + - (proto.attr.is_estimation ? "(?) " : "") + proto.attr.desc + - "]"; - } - UCS_TEST_MESSAGE << "used protocols:" << used_str; - - /* A zero-copy protocol on the source memory: a PUT or an AM send reads - * its local buffer, a GET issued by the receiver reads the remote one */ - auto is_used_on_source = [&](const std::string &name) { - return std::any_of(used.begin(), used.end(), - [&](const used_proto_t &proto) { - return (name == proto.config->proto->name) && - ((proto.config->select_param.mem_type == - send_mem_type) || - (proto.remote_mem_type == - send_mem_type)); - }); - }; - - if (mode == XFER_ZCOPY) { - EXPECT_TRUE(is_used_on_source(is_get ? "rndv/get/zcopy" : - "rndv/put/zcopy")) - << "the source memory is not accessed by the transport"; - } else { - EXPECT_TRUE(is_used_on_source("tag/rndv")) - << "the source memory is not sent by a rendezvous protocol"; - EXPECT_FALSE(is_used_on_source("rndv/get/zcopy") || - is_used_on_source("rndv/put/zcopy") || - is_used_on_source("rndv/am/zcopy")) - << "the source memory is accessed by the transport"; - /* CPU-accessible memory may be copied by the CPU instead of being - * staged. The rendezvous request and reply only estimate the peer's - * protocol, so skip estimated descriptions. */ - if (!(UCS_BIT(send_mem_type) & UCS_MEMORY_TYPES_CPU_ACCESSIBLE)) { - EXPECT_TRUE(std::any_of(used.begin(), used.end(), - [](const used_proto_t &proto) { - return !proto.attr.is_estimation && - strstr(proto.attr.desc, - "frag host"); - })) - << "device memory is not staged through host fragments"; - } - } + mem_buffer::pattern_check(recv_buf.ptr(), length, SEED, + UCS_MEMORY_TYPE_HOST); } -/* Small enough for a driver to sub-allocate from a memory pool */ -static constexpr size_t RNDV_ZCOPY_LENGTH = 64 * UCS_KBYTE; - -UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, get_zcopy, "RNDV_THRESH=0", - "RNDV_SCHEME=get_zcopy") +UCS_TEST_P(test_ucp_mem_type_dmabuf, alloc, "RNDV_THRESH=0") { - test_xfer_from_mem_type(RNDV_ZCOPY_LENGTH, true, XFER_ZCOPY, false); -} + if ((sender().ucph()->dmabuf_mds[mem_type()] == UCP_NULL_RESOURCE) || + (sender().ucph()->dmabuf_reg_md_map == 0)) { + UCS_TEST_SKIP_R("DMA-BUF registration is not supported"); + } -UCS_TEST_P(test_ucp_mem_type_rndv_zcopy, put_zcopy, "RNDV_THRESH=0", - "RNDV_SCHEME=put_zcopy") -{ - test_xfer_from_mem_type(RNDV_ZCOPY_LENGTH, false, XFER_ZCOPY, false); + test_xfer(false, true); } -UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_rndv_zcopy, rcx, - "rc_x") -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcx_ze, - "rc_x,ze_copy") -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_zcopy, rcv_ze, - "rc_v,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf, rcx_ze, "rc_x,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf, rcv_ze, "rc_v,ze_copy") /* - * With DMA-BUF export disabled, the rendezvous zero-copy lane cannot register - * Level Zero memory and the payload must be copied through host memory: device - * memory is staged through fragments, CPU-accessible memory is staged or copied - * in and out. The rendezvous scheme is left to UCP: forcing a zero-copy scheme - * would exclude these protocols. + * With DMA-BUF export disabled, the network memory domains cannot register + * Level Zero memory: device memory is staged through host memory, and + * CPU-accessible memory is copied. */ -class test_ucp_mem_type_rndv_staged : public test_ucp_mem_type_rndv_zcopy { -public: +class test_ucp_mem_type_dmabuf_off : public test_ucp_mem_type_dmabuf { +protected: void init() override { modify_config("ZE_COPY_DMABUF", "n", SETENV_IF_NOT_EXIST); - test_ucp_mem_type_rndv_zcopy::init(); - } - - /* Disabling DMA-BUF export changes how only Level Zero memory is reached */ - static void get_test_variants(std::vector &variants) - { - static constexpr uint64_t ze_mem_types = - UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | - UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) | - UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED); - - add_variant_memtypes(variants, get_test_variants_base, ze_mem_types); + test_ucp_mem_type_dmabuf::init(); } }; -UCS_TEST_P(test_ucp_mem_type_rndv_staged, staged, "RNDV_THRESH=0", +UCS_TEST_P(test_ucp_mem_type_dmabuf_off, alloc, "RNDV_THRESH=0", "RNDV_FRAG_MEM_TYPES=host") { - test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_STAGED, false); + test_xfer(false, false); } -UCS_TEST_P(test_ucp_mem_type_rndv_staged, staged_app_mem, "RNDV_THRESH=0", +UCS_TEST_P(test_ucp_mem_type_dmabuf_off, app_mem, "RNDV_THRESH=0", "RNDV_FRAG_MEM_TYPES=host") { - test_xfer_from_mem_type(4 * UCS_MBYTE, false, XFER_STAGED, true); + test_xfer(true, false); } -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_staged, rcx_ze, +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf_off, rcx_ze, "rc_x,ze_copy") -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_rndv_staged, rcv_ze, +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf_off, rcv_ze, "rc_v,ze_copy") class test_ucp_cuda : public ucp_test { diff --git a/test/gtest/uct/ib/test_ib_xfer.cc b/test/gtest/uct/ib/test_ib_xfer.cc index 280f9594ba3..9fb7e0a1a4a 100644 --- a/test/gtest/uct/ib/test_ib_xfer.cc +++ b/test/gtest/uct/ib/test_ib_xfer.cc @@ -1,5 +1,6 @@ /** * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2014. ALL RIGHTS RESERVED. +* Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -225,3 +226,169 @@ UCS_TEST_P(uct_p2p_mix_test_indirect_atomic, mix1000_indirect_atomic, UCT_INSTANTIATE_IB_TEST_CASE(uct_p2p_mix_test_indirect_atomic) + + +/* + * Moves memory allocated by a DMA-BUF provider, such as a GPU copy memory + * domain, which the IB memory domain registers through the DMA-BUF. Nothing + * falls back at UCT, so a correct payload shows that the NIC accessed the + * memory through the DMA-BUF. + */ +class uct_p2p_rma_test_dmabuf : public uct_p2p_rma_test { +protected: + /* Small enough for a driver to sub-allocate from a memory pool */ + static constexpr size_t LENGTH = 64 * UCS_KBYTE; + + struct provider_mem : uct_allocated_memory_t { + provider_mem(uct_md_h md, ucs_memory_type_t mem_type) + { + uct_alloc_method_t method = UCT_ALLOC_METHOD_MD; + uct_mem_alloc_params_t params = {}; + + params.field_mask = UCT_MEM_ALLOC_PARAM_FIELD_FLAGS | + UCT_MEM_ALLOC_PARAM_FIELD_MEM_TYPE | + UCT_MEM_ALLOC_PARAM_FIELD_MDS | + UCT_MEM_ALLOC_PARAM_FIELD_NAME; + params.flags = UCT_MD_MEM_ACCESS_ALL; + params.mem_type = mem_type; + params.mds.mds = &md; + params.mds.count = 1; + params.name = "dmabuf_test"; + ASSERT_UCS_OK(uct_mem_alloc(LENGTH, &method, 1, ¶ms, this)); + } + + ~provider_mem() + { + uct_mem_free(this); + } + }; + + void test_dmabuf_xfer(uct_md_h provider_md, ucs_memory_type_t mem_type, + bool is_get); + + void test_dmabuf_xfer(bool is_get); +}; + +void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, + ucs_memory_type_t mem_type, + bool is_get) +{ + /* The NIC accesses the sender's memory for PUT, the receiver's for GET */ + entity &mem_entity = is_get ? receiver() : sender(); + uct_md_mem_attr_v2_t mem_attr = {}; + uct_md_mem_reg_params_t reg_params = {}; + uct_allocated_memory_t reg_mem; + uct_rkey_bundle_t rkey = {}; + uct_mem_h uct_memh; + ucs_status_t status; + + UCS_TEST_MESSAGE << ucs_memory_type_names[mem_type] << " memory"; + + /* Keep another buffer of the same size allocated, so that a driver which + * sub-allocates both from one memory pool places the transferred buffer at + * a nonzero offset in the pool */ + provider_mem pad(provider_md, mem_type); + provider_mem mem(provider_md, mem_type); + mapped_buffer host_buf(LENGTH, 0, is_get ? sender() : receiver()); + /* Declared after the memory, so it is deregistered before being freed */ + ucs::handle memh; + + mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_FD | + UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_OFFSET; + ASSERT_UCS_OK(uct_md_mem_query_v2(provider_md, mem.address, LENGTH, + &mem_attr)); + ASSERT_NE(UCT_DMABUF_FD_INVALID, mem_attr.dmabuf_fd); + + reg_params.field_mask = UCT_MD_MEM_REG_FIELD_FLAGS | + UCT_MD_MEM_REG_FIELD_DMABUF_FD | + UCT_MD_MEM_REG_FIELD_DMABUF_OFFSET; + reg_params.flags = UCT_MD_MEM_ACCESS_ALL; + reg_params.dmabuf_fd = mem_attr.dmabuf_fd; + reg_params.dmabuf_offset = mem_attr.dmabuf_offset; + status = uct_md_mem_reg_v2(mem_entity.md(), mem.address, LENGTH, + ®_params, &uct_memh); + close(mem_attr.dmabuf_fd); + ASSERT_UCS_OK(status); + memh.reset(uct_memh, [](uct_mem_h handle, uct_md_h md) { + EXPECT_UCS_OK(uct_md_mem_dereg(md, handle)); + }, mem_entity.md()); + + mem_buffer::pattern_fill(mem.address, LENGTH, SEED1, mem_type); + + uct_iov_t iov = {mem.address, LENGTH, memh, 0, 1}; + if (is_get) { + reg_mem = mem; + reg_mem.memh = memh; + mem_entity.rkey_unpack(®_mem, &rkey); + } + + do { + status = is_get ? + uct_ep_get_zcopy(sender_ep(), host_buf.iov(), 1, + (uintptr_t)mem.address, rkey.rkey, NULL) : + uct_ep_put_zcopy(sender_ep(), &iov, 1, host_buf.addr(), + host_buf.rkey(), NULL); + progress(); + } while (status == UCS_ERR_NO_RESOURCE); + EXPECT_FALSE(UCS_STATUS_IS_ERR(status)) << ucs_status_string(status); + flush(); + + if (is_get) { + mem_entity.rkey_release(&rkey); + } + + host_buf.pattern_check(SEED1); +} + +void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(bool is_get) +{ + unsigned num_mem_types = 0; + uct_md_attr_v2_t md_attr; + unsigned mem_type; + + if (!(sender().md_attr().flags & UCT_MD_FLAG_REG_DMABUF)) { + UCS_TEST_SKIP_R("the memory domain cannot register a DMA-BUF"); + } + + for (const md_resource &rsc : enum_md_resources()) { + ucs::handle md_config; + ucs::handle md; + + UCS_TEST_CREATE_HANDLE(uct_md_config_t*, md_config, + (void (*)(uct_md_config_t*))uct_config_release, + uct_md_config_read, rsc.cmpt, NULL, NULL); + UCS_TEST_CREATE_HANDLE(uct_md_h, md, uct_md_close, uct_md_open, + rsc.cmpt, rsc.rsc_desc.md_name, md_config); + + md_attr.field_mask = UCT_MD_ATTR_FIELD_ALLOC_MEM_TYPES | + UCT_MD_ATTR_FIELD_DMABUF_MEM_TYPES; + ASSERT_UCS_OK(uct_md_query_v2(md, &md_attr)); + ucs_for_each_bit(mem_type, md_attr.alloc_mem_types & + md_attr.dmabuf_mem_types) { + if (mem_buffer::is_mem_type_supported( + (ucs_memory_type_t)mem_type)) { + test_dmabuf_xfer(md, (ucs_memory_type_t)mem_type, is_get); + ++num_mem_types; + } + } + } + + if (num_mem_types == 0) { + UCS_TEST_SKIP_R("no memory domain exports allocated memory as a " + "DMA-BUF"); + } +} + +UCS_TEST_SKIP_COND_P(uct_p2p_rma_test_dmabuf, put_zcopy, + !check_caps(UCT_IFACE_FLAG_PUT_ZCOPY)) +{ + test_dmabuf_xfer(false); +} + +UCS_TEST_SKIP_COND_P(uct_p2p_rma_test_dmabuf, get_zcopy, + !check_caps(UCT_IFACE_FLAG_GET_ZCOPY)) +{ + test_dmabuf_xfer(true); +} + +UCT_INSTANTIATE_RC_TEST_CASE(uct_p2p_rma_test_dmabuf) From 75f974d2d602242d3b1fb9ad4c9992ff64284e0a Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Thu, 5 Nov 2026 08:19:35 -0800 Subject: [PATCH 11/13] UCT/ZE/COPY: Address review: test DMA-BUF settings on CUDA and ROCm Signed-off-by: Yaser Afshar --- src/uct/ze/copy/ze_copy_md.c | 6 ++-- test/gtest/ucp/test_ucp_mem_type.cc | 55 ++++++++++++++++++++--------- test/gtest/uct/ib/test_ib_xfer.cc | 2 +- 3 files changed, 43 insertions(+), 20 deletions(-) diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index aa8fec9d2dc..fbaa37469bf 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -133,12 +133,14 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p, switch (mem_type) { case UCS_MEMORY_TYPE_ZE_HOST: host_desc.pNext = export_desc_p; + status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocHost(md->ze_context, &host_desc, *length_p, alignment, address_p)); break; case UCS_MEMORY_TYPE_ZE_DEVICE: dev_desc.pNext = export_desc_p; + status = UCT_ZE_FUNC_LOG_ERR(zeMemAllocDevice(md->ze_context, &dev_desc, *length_p, alignment, md->ze_device, @@ -406,8 +408,8 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name, md->dmabuf_supported = (config->enable_dmabuf != UCS_NO) && dmabuf_supported; - md->super.ops = &md_ops; - md->super.component = &uct_ze_copy_component; + md->super.ops = &md_ops; + md->super.component = &uct_ze_copy_component; *md_p = (uct_md_h)md; diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index 07fe1fe4ca0..c3bee823af1 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -104,9 +104,10 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_alloc_before_init, all, "all") /* - * ZE_COPY_DMABUF decides whether the network memory domains register Level Zero - * memory through a DMA-BUF. uct_p2p_rma_test_dmabuf checks that the network - * reads the memory correctly through it. + * The CUDA_COPY_DMABUF, ROCM_COPY_DMABUF and ZE_COPY_DMABUF settings decide + * whether the GPU memory domains provide a DMA-BUF that the network memory + * domains register. uct_p2p_rma_test_dmabuf checks that the network reads the + * memory correctly through it. */ class test_ucp_mem_type_dmabuf : public test_ucp_mem_type { public: @@ -115,6 +116,8 @@ class test_ucp_mem_type_dmabuf : public test_ucp_mem_type { static void get_test_variants(std::vector &variants) { add_variant_memtypes(variants, get_test_variants_base, + UCS_BIT(UCS_MEMORY_TYPE_CUDA) | + UCS_BIT(UCS_MEMORY_TYPE_ROCM) | UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST) | UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE)); } @@ -146,10 +149,18 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) ucp_mem_map_params_t params; void *send_buf, *sreq, *rreq; ucp_mem_h memh; - - EXPECT_EQ(exp_dmabuf, (context->reg_md_map[mem_type()] & - context->dmabuf_reg_md_map) != 0) - << mem_type_name << " memory registration through a DMA-BUF"; + /* The network memory domains register Level Zero memory only through a + * DMA-BUF, and CUDA or ROCm memory also through peer memory */ + bool dmabuf_only = (mem_type() == UCS_MEMORY_TYPE_ZE_HOST) || + (mem_type() == UCS_MEMORY_TYPE_ZE_DEVICE); + + EXPECT_EQ(exp_dmabuf, context->dmabuf_mds[mem_type()] != UCP_NULL_RESOURCE) + << mem_type_name << " memory DMA-BUF provider"; + if (dmabuf_only) { + EXPECT_EQ(exp_dmabuf, (context->reg_md_map[mem_type()] & + context->dmabuf_reg_md_map) != 0) + << mem_type_name << " memory registration through a DMA-BUF"; + } if (app_owned) { app_buf.reset(new mem_buffer(length, mem_type())); @@ -167,9 +178,13 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) EXPECT_UCS_OK(ucp_mem_unmap(ctx, handle)); }, context); - EXPECT_EQ(exp_dmabuf, (memh->md_map & context->dmabuf_reg_md_map) != 0) - << mem_type_name << " memory handle registration through a " - << "DMA-BUF"; + if (exp_dmabuf || dmabuf_only) { + EXPECT_EQ(exp_dmabuf, + (memh->md_map & context->dmabuf_reg_md_map) != 0) + << mem_type_name << " memory handle registration through " + << "a DMA-BUF"; + } + send_buf = ucp_memh_address(memh); param.op_attr_mask = UCP_OP_ATTR_FIELD_MEMH; param.memh = memh; @@ -179,6 +194,7 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) mem_buffer::pattern_fill(recv_buf.ptr(), length, 0, UCS_MEMORY_TYPE_HOST); recv_param.op_attr_mask = 0; + rreq = ucp_tag_recv_nbx(receiver().worker(), recv_buf.ptr(), length, 1, 1, &recv_param); sreq = ucp_tag_send_nbx(sender().ep(), send_buf, length, 1, ¶m); @@ -199,19 +215,24 @@ UCS_TEST_P(test_ucp_mem_type_dmabuf, alloc, "RNDV_THRESH=0") test_xfer(false, true); } -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf, rcx_ze, "rc_x,ze_copy") -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf, rcv_ze, "rc_v,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_dmabuf, rcx, + "rc_x,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_dmabuf, rcv, + "rc_v,ze_copy") /* * With DMA-BUF export disabled, the network memory domains cannot register * Level Zero memory: device memory is staged through host memory, and - * CPU-accessible memory is copied. + * CPU-accessible memory is copied. CUDA and ROCm memory may still be + * registered through peer memory. */ class test_ucp_mem_type_dmabuf_off : public test_ucp_mem_type_dmabuf { protected: void init() override { + modify_config("CUDA_COPY_DMABUF", "n", SETENV_IF_NOT_EXIST); + modify_config("ROCM_COPY_DMABUF", "n", SETENV_IF_NOT_EXIST); modify_config("ZE_COPY_DMABUF", "n", SETENV_IF_NOT_EXIST); test_ucp_mem_type_dmabuf::init(); } @@ -229,10 +250,10 @@ UCS_TEST_P(test_ucp_mem_type_dmabuf_off, app_mem, "RNDV_THRESH=0", test_xfer(true, false); } -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf_off, rcx_ze, - "rc_x,ze_copy") -UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_mem_type_dmabuf_off, rcv_ze, - "rc_v,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_dmabuf_off, rcx, + "rc_x,ze_copy") +UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_mem_type_dmabuf_off, rcv, + "rc_v,ze_copy") class test_ucp_cuda : public ucp_test { public: diff --git a/test/gtest/uct/ib/test_ib_xfer.cc b/test/gtest/uct/ib/test_ib_xfer.cc index 9fb7e0a1a4a..cafffcfff5e 100644 --- a/test/gtest/uct/ib/test_ib_xfer.cc +++ b/test/gtest/uct/ib/test_ib_xfer.cc @@ -227,7 +227,6 @@ UCS_TEST_P(uct_p2p_mix_test_indirect_atomic, mix1000_indirect_atomic, UCT_INSTANTIATE_IB_TEST_CASE(uct_p2p_mix_test_indirect_atomic) - /* * Moves memory allocated by a DMA-BUF provider, such as a GPU copy memory * domain, which the IB memory domain registers through the DMA-BUF. Nothing @@ -305,6 +304,7 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, reg_params.flags = UCT_MD_MEM_ACCESS_ALL; reg_params.dmabuf_fd = mem_attr.dmabuf_fd; reg_params.dmabuf_offset = mem_attr.dmabuf_offset; + status = uct_md_mem_reg_v2(mem_entity.md(), mem.address, LENGTH, ®_params, &uct_memh); close(mem_attr.dmabuf_fd); From 4a2d4f443076ea2dcebe04cc7586dee0be667f62 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Thu, 5 Nov 2026 09:47:52 -0800 Subject: [PATCH 12/13] UCT/ZE/COPY: Address review: skip unopenable MDs, share DMA-BUF export probe Signed-off-by: Yaser Afshar --- src/uct/ze/copy/ze_copy_md.c | 18 +++++----- test/gtest/common/mem_buffer.cc | 20 +++++++++++ test/gtest/common/mem_buffer.h | 8 +++++ test/gtest/ucp/test_ucp_mem_type.cc | 26 ++++++++++---- test/gtest/uct/ib/test_ib_xfer.cc | 56 ++++++++++++++++++++--------- test/gtest/uct/test_md.cc | 18 ++-------- 6 files changed, 99 insertions(+), 47 deletions(-) diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index fbaa37469bf..9d5f67be79e 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -105,24 +105,22 @@ uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p, unsigned flags, const char *alloc_name, uct_mem_h *memh_p) { uct_ze_copy_md_t *md = ucs_derived_of(tl_md, uct_ze_copy_md_t); - /* Request DMA-BUF exportability for the memory types advertised in - * 'dmabuf_mem_types', and only for those, so that allocating a type without - * a DMA-BUF path cannot fail because of the descriptor. Level Zero may - * return a DMA-BUF for an allocation which was not created for external - * export, but such a handle does not necessarily describe the requested - * range, and an RDMA device may read incorrect data through it. */ + /* Request export only for 'dmabuf_mem_types', so that other types cannot + * fail because of it. A DMA-BUF of memory not allocated for export may not + * describe the requested range, and an RDMA device may read incorrect data + * through it. */ ze_external_memory_export_desc_t export_desc = { .stype = ZE_STRUCTURE_TYPE_EXTERNAL_MEMORY_EXPORT_DESC, .flags = ZE_EXTERNAL_MEMORY_TYPE_FLAG_DMA_BUF }; - void *export_desc_p = NULL; - ze_host_mem_alloc_desc_t host_desc = { + void *export_desc_p = NULL; + ze_host_mem_alloc_desc_t host_desc = { .stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC }; - ze_device_mem_alloc_desc_t dev_desc = { + ze_device_mem_alloc_desc_t dev_desc = { .stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC }; - size_t alignment = ucs_get_page_size(); + size_t alignment = ucs_get_page_size(); ucs_status_t status; if (md->dmabuf_supported && diff --git a/test/gtest/common/mem_buffer.cc b/test/gtest/common/mem_buffer.cc index c9acdbf2778..238eb591c42 100644 --- a/test/gtest/common/mem_buffer.cc +++ b/test/gtest/common/mem_buffer.cc @@ -416,6 +416,26 @@ bool mem_buffer::cuda_mem_is_gdr_capable() #endif } +ucs_status_t mem_buffer::query_dmabuf_export(uct_md_h md, + ucs_memory_type_t mem_type, + bool *exported_p) +{ + uct_md_mem_attr_t mem_attr = {}; + mem_buffer buffer(ucs_get_page_size(), mem_type); + ucs_status_t status; + + mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_DMABUF_FD; + status = uct_md_mem_query(md, buffer.ptr(), buffer.size(), &mem_attr); + if (status == UCS_OK) { + *exported_p = (mem_attr.dmabuf_fd != UCT_DMABUF_FD_INVALID); + if (*exported_p) { + close(mem_attr.dmabuf_fd); + } + } + + return status; +} + #if HAVE_ZE namespace { diff --git a/test/gtest/common/mem_buffer.h b/test/gtest/common/mem_buffer.h index 5eb02bf55d5..198a0014dae 100644 --- a/test/gtest/common/mem_buffer.h +++ b/test/gtest/common/mem_buffer.h @@ -1,5 +1,6 @@ /** * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. + * Copyright (C) Intel Corporation, 2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -7,6 +8,7 @@ #ifndef GTEST_MEM_BUFFER_H_ #define GTEST_MEM_BUFFER_H_ +#include #include #include #include @@ -117,6 +119,12 @@ class mem_buffer { * which it does not under an MPS MLOPart device */ static bool cuda_mem_is_gdr_capable(); + /* Query 'md' for a DMA-BUF of a page of 'mem_type' memory, and set + * 'exported_p' to whether it returned one */ + static ucs_status_t query_dmabuf_export(uct_md_h md, + ucs_memory_type_t mem_type, + bool *exported_p); + /* Return free memory on the BAR1 / GPU. If GPU is not used * SIZE_MAX is returned */ static size_t get_bar1_free_size() diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index c3bee823af1..30b00e66383 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -142,6 +142,10 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) static constexpr size_t length = 4 * UCS_MBYTE; ucp_context_h context = sender().ucph(); const char *mem_type_name = ucs_memory_type_names[mem_type()]; + /* The network memory domains register Level Zero memory only through a + * DMA-BUF, and CUDA or ROCm memory also through peer memory */ + bool dmabuf_only = (mem_type() == UCS_MEMORY_TYPE_ZE_HOST) || + (mem_type() == UCS_MEMORY_TYPE_ZE_DEVICE); mem_buffer recv_buf(length, UCS_MEMORY_TYPE_HOST); ucs::handle memh_handle; ucs::auto_ptr app_buf; @@ -149,10 +153,6 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) ucp_mem_map_params_t params; void *send_buf, *sreq, *rreq; ucp_mem_h memh; - /* The network memory domains register Level Zero memory only through a - * DMA-BUF, and CUDA or ROCm memory also through peer memory */ - bool dmabuf_only = (mem_type() == UCS_MEMORY_TYPE_ZE_HOST) || - (mem_type() == UCS_MEMORY_TYPE_ZE_DEVICE); EXPECT_EQ(exp_dmabuf, context->dmabuf_mds[mem_type()] != UCP_NULL_RESOURCE) << mem_type_name << " memory DMA-BUF provider"; @@ -207,11 +207,25 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) UCS_TEST_P(test_ucp_mem_type_dmabuf, alloc, "RNDV_THRESH=0") { - if ((sender().ucph()->dmabuf_mds[mem_type()] == UCP_NULL_RESOURCE) || - (sender().ucph()->dmabuf_reg_md_map == 0)) { + ucp_context_h context = sender().ucph(); + ucp_md_index_t md_index = context->dmabuf_mds[mem_type()]; + bool exported; + + if ((md_index == UCP_NULL_RESOURCE) || (context->dmabuf_reg_md_map == 0)) { UCS_TEST_SKIP_R("DMA-BUF registration is not supported"); } + /* Only an MPS MLOPart device may refuse an export the memory domain + * advertises, and the network cannot register the memory then */ + if ((mem_type() == UCS_MEMORY_TYPE_CUDA) && + !mem_buffer::cuda_mem_is_gdr_capable()) { + ASSERT_UCS_OK(mem_buffer::query_dmabuf_export( + context->tl_mds[md_index].md, mem_type(), &exported)); + if (!exported) { + UCS_TEST_SKIP_R("the device refuses DMA-BUF export"); + } + } + test_xfer(false, true); } diff --git a/test/gtest/uct/ib/test_ib_xfer.cc b/test/gtest/uct/ib/test_ib_xfer.cc index cafffcfff5e..f19df9333ae 100644 --- a/test/gtest/uct/ib/test_ib_xfer.cc +++ b/test/gtest/uct/ib/test_ib_xfer.cc @@ -227,12 +227,9 @@ UCS_TEST_P(uct_p2p_mix_test_indirect_atomic, mix1000_indirect_atomic, UCT_INSTANTIATE_IB_TEST_CASE(uct_p2p_mix_test_indirect_atomic) -/* - * Moves memory allocated by a DMA-BUF provider, such as a GPU copy memory - * domain, which the IB memory domain registers through the DMA-BUF. Nothing - * falls back at UCT, so a correct payload shows that the NIC accessed the - * memory through the DMA-BUF. - */ +/* Moves memory of a DMA-BUF provider, which the IB memory domain registers + * through the DMA-BUF. UCT has no fallback, so a correct payload shows that + * the NIC accessed the memory through the DMA-BUF. */ class uct_p2p_rma_test_dmabuf : public uct_p2p_rma_test { protected: /* Small enough for a driver to sub-allocate from a memory pool */ @@ -262,13 +259,14 @@ class uct_p2p_rma_test_dmabuf : public uct_p2p_rma_test { } }; - void test_dmabuf_xfer(uct_md_h provider_md, ucs_memory_type_t mem_type, + /* Return false if the device refuses to export the memory */ + bool test_dmabuf_xfer(uct_md_h provider_md, ucs_memory_type_t mem_type, bool is_get); void test_dmabuf_xfer(bool is_get); }; -void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, +bool uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, ucs_memory_type_t mem_type, bool is_get) { @@ -276,8 +274,8 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, entity &mem_entity = is_get ? receiver() : sender(); uct_md_mem_attr_v2_t mem_attr = {}; uct_md_mem_reg_params_t reg_params = {}; - uct_allocated_memory_t reg_mem; uct_rkey_bundle_t rkey = {}; + uct_allocated_memory_t reg_mem; uct_mem_h uct_memh; ucs_status_t status; @@ -296,14 +294,27 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, UCT_MD_MEM_ATTR_V2_FIELD_DMABUF_OFFSET; ASSERT_UCS_OK(uct_md_mem_query_v2(provider_md, mem.address, LENGTH, &mem_attr)); - ASSERT_NE(UCT_DMABUF_FD_INVALID, mem_attr.dmabuf_fd); + if (mem_attr.dmabuf_fd == UCT_DMABUF_FD_INVALID) { + /* Only an MPS MLOPart device may refuse an export the memory domain + * advertises; anywhere else a refusal is a failure */ + if ((mem_type == UCS_MEMORY_TYPE_CUDA) && + !mem_buffer::cuda_mem_is_gdr_capable()) { + UCS_TEST_MESSAGE << "skipping cuda (device refuses DMA-BUF export)"; + return false; + } + + UCS_TEST_ABORT(ucs_memory_type_names[mem_type] + << " memory is not exported as a DMA-BUF"); + } reg_params.field_mask = UCT_MD_MEM_REG_FIELD_FLAGS | UCT_MD_MEM_REG_FIELD_DMABUF_FD | - UCT_MD_MEM_REG_FIELD_DMABUF_OFFSET; + UCT_MD_MEM_REG_FIELD_DMABUF_OFFSET | + UCT_MD_MEM_REG_FIELD_MEM_TYPE; reg_params.flags = UCT_MD_MEM_ACCESS_ALL; reg_params.dmabuf_fd = mem_attr.dmabuf_fd; reg_params.dmabuf_offset = mem_attr.dmabuf_offset; + reg_params.mem_type = mem_type; status = uct_md_mem_reg_v2(mem_entity.md(), mem.address, LENGTH, ®_params, &uct_memh); @@ -330,7 +341,7 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, host_buf.rkey(), NULL); progress(); } while (status == UCS_ERR_NO_RESOURCE); - EXPECT_FALSE(UCS_STATUS_IS_ERR(status)) << ucs_status_string(status); + ASSERT_UCS_OK_OR_INPROGRESS(status); flush(); if (is_get) { @@ -338,6 +349,7 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(uct_md_h provider_md, } host_buf.pattern_check(SEED1); + return true; } void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(bool is_get) @@ -353,12 +365,24 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(bool is_get) for (const md_resource &rsc : enum_md_resources()) { ucs::handle md_config; ucs::handle md; + ucs_status_t status; + uct_md_h uct_md; UCS_TEST_CREATE_HANDLE(uct_md_config_t*, md_config, (void (*)(uct_md_config_t*))uct_config_release, uct_md_config_read, rsc.cmpt, NULL, NULL); - UCS_TEST_CREATE_HANDLE(uct_md_h, md, uct_md_close, uct_md_open, - rsc.cmpt, rsc.rsc_desc.md_name, md_config); + + /* A listed memory domain may fail to open, as in enum_resources() */ + { + scoped_log_handler slh(hide_errors_logger); + status = uct_md_open(rsc.cmpt, rsc.rsc_desc.md_name, md_config, + &uct_md); + } + if (status != UCS_OK) { + continue; + } + + md.reset(uct_md, uct_md_close); md_attr.field_mask = UCT_MD_ATTR_FIELD_ALLOC_MEM_TYPES | UCT_MD_ATTR_FIELD_DMABUF_MEM_TYPES; @@ -366,8 +390,8 @@ void uct_p2p_rma_test_dmabuf::test_dmabuf_xfer(bool is_get) ucs_for_each_bit(mem_type, md_attr.alloc_mem_types & md_attr.dmabuf_mem_types) { if (mem_buffer::is_mem_type_supported( - (ucs_memory_type_t)mem_type)) { - test_dmabuf_xfer(md, (ucs_memory_type_t)mem_type, is_get); + (ucs_memory_type_t)mem_type) && + test_dmabuf_xfer(md, (ucs_memory_type_t)mem_type, is_get)) { ++num_mem_types; } } diff --git a/test/gtest/uct/test_md.cc b/test/gtest/uct/test_md.cc index 5750fb2e3e8..dc658c2f63c 100644 --- a/test/gtest/uct/test_md.cc +++ b/test/gtest/uct/test_md.cc @@ -263,27 +263,15 @@ bool test_md::check_alloc_mem_type(ucs_memory_type_t mem_type) * refuse it, as it does under an MPS MLOPart device */ bool test_md::check_dmabuf_export(ucs_memory_type_t mem_type) { - uct_md_mem_attr_t mem_attr = {}; + bool exported; if (!(m_md_attr.dmabuf_mem_types & UCS_BIT(mem_type)) || !mem_buffer::is_mem_type_supported(mem_type)) { return false; } - mem_buffer buffer(ucs_get_page_size(), mem_type); - - mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_DMABUF_FD; - if (uct_md_mem_query(md(), buffer.ptr(), buffer.size(), &mem_attr) != - UCS_OK) { - return false; - } - - if (mem_attr.dmabuf_fd == UCT_DMABUF_FD_INVALID) { - return false; - } - - close(mem_attr.dmabuf_fd); - return true; + return (mem_buffer::query_dmabuf_export(md(), mem_type, &exported) == + UCS_OK) && exported; } /* Under an MPS MLOPart device the driver allocates no GDR capable memory and From 346c1231e8f84350b18f0c8e88f5e8a06e5e7f46 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Thu, 5 Nov 2026 12:24:07 -0800 Subject: [PATCH 13/13] UCT/ZE/COPY: Address review: clarify DMA-BUF comments, help text and test message Signed-off-by: Yaser Afshar --- src/uct/ib/base/ib_md.c | 3 ++- src/uct/ze/copy/ze_copy_md.c | 6 +++--- test/gtest/common/mem_buffer.h | 4 ++-- test/gtest/ucp/test_ucp_mem_type.cc | 4 ++-- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c index a1e1e2841f6..d5345c19b16 100644 --- a/src/uct/ib/base/ib_md.c +++ b/src/uct/ib/base/ib_md.c @@ -1472,7 +1472,8 @@ ucs_status_t uct_ib_md_open_common(uct_ib_md_t *md, /* Limit the number of SGEs dumped by the packet tracer, so that it does not * dereference non-host memory from the CPU. DMA-BUF registration makes such - * memory reachable without it being listed in 'reg_mem_types'. */ + * memory registrable without it being listed in 'reg_mem_types', so the + * limit also applies to host-only runs on a DMA-BUF capable device. */ md->dev.max_zcopy_log_sge = INT_MAX; if ((md->reg_mem_types & ~UCS_BIT(UCS_MEMORY_TYPE_HOST)) || (md->cap_flags & UCT_MD_FLAG_REG_DMABUF)) { diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index 9d5f67be79e..a071aa1ea37 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/src/uct/ze/copy/ze_copy_md.c @@ -38,9 +38,9 @@ static ucs_config_field_t uct_ze_copy_md_config_table[] = { {"DMABUF", "try", "Enable using cross-device dmabuf file descriptor.\n" - "Level Zero device memory may have to be allocated with external\n" - "DMA-BUF export enabled. Without it, registration can succeed but an\n" - "RDMA device may read incorrect data, and no error is reported. Set to\n" + "Level Zero memory may have to be allocated with external DMA-BUF\n" + "export enabled. Without it, registration can succeed but an RDMA\n" + "device may access incorrect data, and no error is reported. Set to\n" "'n' if the application allocates Level Zero memory without\n" "requesting external export; device memory is then staged through\n" "host memory for RDMA.", diff --git a/test/gtest/common/mem_buffer.h b/test/gtest/common/mem_buffer.h index 198a0014dae..658c09d1c17 100644 --- a/test/gtest/common/mem_buffer.h +++ b/test/gtest/common/mem_buffer.h @@ -119,8 +119,8 @@ class mem_buffer { * which it does not under an MPS MLOPart device */ static bool cuda_mem_is_gdr_capable(); - /* Query 'md' for a DMA-BUF of a page of 'mem_type' memory, and set - * 'exported_p' to whether it returned one */ + /* Query 'md' for a DMA-BUF of a page of 'mem_type' memory allocated by + * mem_buffer, and set 'exported_p' to whether it returned one */ static ucs_status_t query_dmabuf_export(uct_md_h md, ucs_memory_type_t mem_type, bool *exported_p); diff --git a/test/gtest/ucp/test_ucp_mem_type.cc b/test/gtest/ucp/test_ucp_mem_type.cc index 30b00e66383..10336966be9 100644 --- a/test/gtest/ucp/test_ucp_mem_type.cc +++ b/test/gtest/ucp/test_ucp_mem_type.cc @@ -181,8 +181,8 @@ void test_ucp_mem_type_dmabuf::test_xfer(bool app_owned, bool exp_dmabuf) if (exp_dmabuf || dmabuf_only) { EXPECT_EQ(exp_dmabuf, (memh->md_map & context->dmabuf_reg_md_map) != 0) - << mem_type_name << " memory handle registration through " - << "a DMA-BUF"; + << mem_type_name << " memory handle registration on a " + << "DMA-BUF capable memory domain"; } send_buf = ucp_memh_address(memh);