From dba36917f3bee9e96a294528aab29dd8feff56f3 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 23 Oct 2026 10:36:18 -0700 Subject: [PATCH 1/5] UCT/ZE/COPY: Prefer a dedicated copy engine ze_copy created its command queues and lists on queue group 0. On the tested devices that is the compute group, so ze_copy copies ran on the compute engine, next to the application's kernels. Copies to a peer or to host should run on the copy engine, as the driver recommends. Move the queue group lookup from ze_ipc to ze_base and use it wherever ze_copy creates a queue, so the queues of other devices, created on first use, get a dedicated copy group too when the device has one, or else any copy-capable group. Add a test which checks every queue of the iface, including one created for another device. Signed-off-by: Yaser Afshar --- src/uct/ze/base/ze_base.c | 56 +++++++++++++++++++++++++ src/uct/ze/base/ze_base.h | 6 +++ src/uct/ze/copy/ze_copy_iface.c | 15 +++++++ src/uct/ze/ze_ipc/ze_ipc_iface.c | 59 +------------------------- test/gtest/uct/ze/test_ze.cc | 71 ++++++++++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 58 deletions(-) diff --git a/src/uct/ze/base/ze_base.c b/src/uct/ze/base/ze_base.c index 95691fabb0a..25c610e40a9 100644 --- a/src/uct/ze/base/ze_base.c +++ b/src/uct/ze/base/ze_base.c @@ -456,6 +456,62 @@ ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device) return UCS_SYS_DEVICE_ID_UNKNOWN; } +ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, + uint32_t *ordinal_p, + uint32_t *num_queues_p) +{ + ze_command_queue_group_properties_t queue_props[16]; + uint32_t num_queue_groups = 16; + uint32_t i; + ze_result_t ret; + + /* Initialize structure types */ + for (i = 0; i < 16; i++) { + queue_props[i].stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; + queue_props[i].pNext = NULL; + } + + /* Get queue group properties */ + ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, + queue_props); + if (ret != ZE_RESULT_SUCCESS) { + ucs_error("zeDeviceGetCommandQueueGroupProperties(device=%p) failed: " + "0x%x", + device, ret); + return UCS_ERR_IO_ERROR; + } + + /* First pass: find dedicated copy engine (COPY flag but NOT COMPUTE) */ + for (i = 0; i < num_queue_groups; i++) { + if ((queue_props[i].flags & + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) && + !(queue_props[i].flags & + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE)) { + *ordinal_p = i; + *num_queues_p = queue_props[i].numQueues; + ucs_debug("device %p: using dedicated copy engine queue group " + "%u (numQueues=%u)", + device, i, queue_props[i].numQueues); + return UCS_OK; + } + } + + /* Second pass: find any queue group that supports COPY */ + for (i = 0; i < num_queue_groups; i++) { + if (queue_props[i].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) { + *ordinal_p = i; + *num_queues_p = queue_props[i].numQueues; + ucs_debug("device %p: using copy-capable queue group %u " + "(numQueues=%u)", + device, i, queue_props[i].numQueues); + return UCS_OK; + } + } + + ucs_error("device %p: no copy-capable queue group found", device); + return UCS_ERR_NO_DEVICE; +} + /** * Query MD resources - returns one MD per sub-device * This is correct because each sub-device has separate memory diff --git a/src/uct/ze/base/ze_base.h b/src/uct/ze/base/ze_base.h index d8eb8a9517e..441fa11455e 100644 --- a/src/uct/ze/base/ze_base.h +++ b/src/uct/ze/base/ze_base.h @@ -65,6 +65,12 @@ ze_device_handle_t uct_ze_base_get_device_handle_from_subdevice( ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device); +/* Find the queue group to run memory copies on a device: a dedicated copy + * engine group if the device has one, or else any group which can copy */ +ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, + uint32_t *ordinal_p, + uint32_t *num_queues_p); + /* Get root device handle by ordinal, returns NULL if out of range */ ze_device_handle_t uct_ze_base_get_device(int ordinal); diff --git a/src/uct/ze/copy/ze_copy_iface.c b/src/uct/ze/copy/ze_copy_iface.c index 631324a120f..abe6f3d336a 100644 --- a/src/uct/ze/copy/ze_copy_iface.c +++ b/src/uct/ze/copy/ze_copy_iface.c @@ -249,9 +249,24 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, uct_ze_copy_md_t); ze_command_queue_desc_t cq_desc = {}; ze_command_list_desc_t cl_desc = {}; + uint32_t ordinal, num_queues; uct_ze_copy_queue_t *queue; + ucs_status_t status; ze_result_t ret; + /* Run on a copy engine rather than the default compute group: the + * driver recommends it for copies to a peer or to host, and the copies + * then do not compete with the application's kernels */ + status = uct_ze_base_get_copy_ordinal(device, &ordinal, &num_queues); + if (status != UCS_OK) { + return status; + } + + cq_desc.stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC; + cq_desc.ordinal = ordinal; + cl_desc.stype = ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC; + cl_desc.commandQueueGroupOrdinal = ordinal; + queue = ucs_array_append(&iface->queues, return UCS_ERR_NO_MEMORY); ret = zeCommandQueueCreate(ze_md->ze_context, device, &cq_desc, diff --git a/src/uct/ze/ze_ipc/ze_ipc_iface.c b/src/uct/ze/ze_ipc/ze_ipc_iface.c index c25eb04f789..178351a2219 100644 --- a/src/uct/ze/ze_ipc/ze_ipc_iface.c +++ b/src/uct/ze/ze_ipc/ze_ipc_iface.c @@ -527,63 +527,6 @@ static uct_iface_internal_ops_t uct_ze_ipc_iface_internal_ops = { }; -/* Find the copy engine queue group ordinal for a device */ -static ucs_status_t uct_ze_ipc_find_copy_ordinal(ze_device_handle_t device, - uint32_t *ordinal_p, - uint32_t *num_queues_p) -{ - ze_command_queue_group_properties_t queue_props[16]; - uint32_t num_queue_groups = 16; - uint32_t i; - ze_result_t ret; - - /* Initialize structure types */ - for (i = 0; i < 16; i++) { - queue_props[i].stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; - queue_props[i].pNext = NULL; - } - - /* Get queue group properties */ - ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, - queue_props); - if (ret != ZE_RESULT_SUCCESS) { - ucs_error("ze_ipc_iface: zeDeviceGetCommandQueueGroupProperties " - "failed: 0x%x", - ret); - return UCS_ERR_IO_ERROR; - } - - /* First pass: find dedicated copy engine (COPY flag but NOT COMPUTE) */ - for (i = 0; i < num_queue_groups; i++) { - if ((queue_props[i].flags & - ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) && - !(queue_props[i].flags & - ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE)) { - *ordinal_p = i; - *num_queues_p = queue_props[i].numQueues; - ucs_debug("ze_ipc_iface: using dedicated copy engine queue group " - "%u (numQueues=%u)", - i, queue_props[i].numQueues); - return UCS_OK; - } - } - - /* Second pass: find any queue group that supports COPY */ - for (i = 0; i < num_queue_groups; i++) { - if (queue_props[i].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) { - *ordinal_p = i; - *num_queues_p = queue_props[i].numQueues; - ucs_debug("ze_ipc_iface: using copy-capable queue group %u " - "(numQueues=%u)", - i, queue_props[i].numQueues); - return UCS_OK; - } - } - - ucs_error("ze_ipc_iface: no copy-capable queue group found"); - return UCS_ERR_NO_DEVICE; -} - static UCS_CLASS_INIT_FUNC(uct_ze_ipc_iface_t, uct_md_h md, uct_worker_h worker, const uct_iface_params_t *params, const uct_iface_config_t *tl_config) @@ -614,7 +557,7 @@ static UCS_CLASS_INIT_FUNC(uct_ze_ipc_iface_t, uct_md_h md, uct_worker_h worker, self->next_cmd_list = 0; /* Find copy engine queue group ordinal and available queue count */ - status = uct_ze_ipc_find_copy_ordinal(self->ze_device, ©_ordinal, + status = uct_ze_base_get_copy_ordinal(self->ze_device, ©_ordinal, &num_queues); if (status != UCS_OK) { return status; diff --git a/test/gtest/uct/ze/test_ze.cc b/test/gtest/uct/ze/test_ze.cc index 19dc0e6af9b..790514bf07c 100644 --- a/test/gtest/uct/ze/test_ze.cc +++ b/test/gtest/uct/ze/test_ze.cc @@ -278,6 +278,59 @@ class test_ze_copy_rma : public uct_p2p_rma_test { return ucs_array_elem(&ze_iface()->queues, queue_index).device; } + /* Check that the queue runs on a copy group, and on a dedicated copy + * group when its device has one */ + void check_copy_engine(const uct_ze_copy_queue_t &queue) + { + const ze_command_queue_group_property_flags_t copy_only = + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY; + const ze_command_queue_group_property_flags_t copy_compute = + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY | + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE; + std::vector groups; + uint32_t num_groups = 0; + bool has_copy_only = false; + uint32_t cmdq_ordinal = UINT32_MAX; + uint32_t cmdl_ordinal = UINT32_MAX; + + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeDeviceGetCommandQueueGroupProperties(queue.device, + &num_groups, NULL)); + groups.resize(num_groups); + for (auto &group : groups) { + group = {}; + group.stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; + } + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeDeviceGetCommandQueueGroupProperties(queue.device, + &num_groups, + groups.data())); + + for (const auto &group : groups) { + if ((group.flags & copy_compute) == copy_only) { + has_copy_only = true; + } + } + + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeCommandQueueGetOrdinal(queue.cmdq, &cmdq_ordinal)); + ASSERT_EQ(ZE_RESULT_SUCCESS, + zeCommandListGetOrdinal(queue.cmdl, &cmdl_ordinal)); + UCS_TEST_MESSAGE << "device " << queue.device << ": queue group " + << cmdq_ordinal << " of " << num_groups + << (has_copy_only ? ", device has a copy engine" : + ""); + + /* The list must be submitted to the group of its queue */ + EXPECT_EQ(cmdq_ordinal, cmdl_ordinal); + ASSERT_LT(cmdq_ordinal, num_groups); + EXPECT_TRUE(groups[cmdq_ordinal].flags & + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY); + if (has_copy_only) { + EXPECT_EQ(copy_only, groups[cmdq_ordinal].flags & copy_compute); + } + } + void zcopy(void *buffer, void *remote_addr, size_t length, bool is_put) { uct_iov_t iov; @@ -487,4 +540,22 @@ UCS_TEST_SKIP_COND_P(test_ze_copy_rma, ze_caps_and_mem_types, UCS_BIT(UCS_MEMORY_TYPE_ZE_HOST)); } +UCS_TEST_P(test_ze_copy_rma, queue_on_copy_engine) +{ + ze_device_handle_t device = other_device(); + std::vector host_buf(64); + const uct_ze_copy_queue_t *queue; + + /* Also create the queue of another device, if there is one; the queue + * stays in the iface after the buffer is released */ + if (device != NULL) { + external_buffer other_buf(device, host_buf.size()); + EXPECT_NE(0u, get_queue(other_buf.ptr(), host_buf.data())); + } + + ucs_array_for_each(queue, &ze_iface()->queues) { + check_copy_engine(*queue); + } +} + UCT_INSTANTIATE_ZE_TEST_CASE(test_ze_copy_rma) From 2324f6ac3e03a84c395c09fb9402513c266eb33b Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 23 Oct 2026 11:40:56 -0700 Subject: [PATCH 2/5] UCT/ZE/COPY: Address review: query group count, guard 1.9 API in test Signed-off-by: Yaser Afshar --- config/m4/ze.m4 | 5 +++ src/uct/ze/base/ze_base.c | 66 +++++++++++++++++++++++--------- src/uct/ze/base/ze_base.h | 5 ++- src/uct/ze/copy/ze_copy_iface.c | 4 +- src/uct/ze/ze_ipc/ze_ipc_iface.c | 5 ++- test/gtest/uct/ze/test_ze.cc | 38 ++++++++++++++++-- 6 files changed, 95 insertions(+), 28 deletions(-) diff --git a/config/m4/ze.m4 b/config/m4/ze.m4 index 1ba48d302c8..02e766df468 100644 --- a/config/m4/ze.m4 +++ b/config/m4/ze.m4 @@ -47,6 +47,11 @@ AS_IF([test "x$ze_checked" != "xyes"], [AC_CHECK_LIB([ze_loader], [zeInit], [ZE_LIBS="$ZE_LIBS -lze_loader"], [ze_happy="no"])]) + # Queue group ordinal queries, Level Zero API 1.9 (used by gtest) + AS_IF([test "x$ze_happy" = "xyes"], + [AC_CHECK_DECLS([zeCommandQueueGetOrdinal, zeCommandListGetOrdinal], + [], [], [[#include ]])]) + CPPFLAGS="$save_CPPFLAGS" LDFLAGS="$save_LDFLAGS" LIBS="$save_LIBS" diff --git a/src/uct/ze/base/ze_base.c b/src/uct/ze/base/ze_base.c index 25c610e40a9..0a6936c6ed9 100644 --- a/src/uct/ze/base/ze_base.c +++ b/src/uct/ze/base/ze_base.c @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -457,28 +458,48 @@ ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device) } ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, uint32_t *ordinal_p, uint32_t *num_queues_p) { - ze_command_queue_group_properties_t queue_props[16]; - uint32_t num_queue_groups = 16; - uint32_t i; + uint32_t num_queue_groups = 0; + ze_command_queue_group_properties_t *queue_props; + ucs_status_t status; ze_result_t ret; + uint32_t i; - /* Initialize structure types */ - for (i = 0; i < 16; i++) { + /* Query the group count first: the driver fills at most the given count + * and may leave the count itself unchanged */ + ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, + NULL); + if (ret != ZE_RESULT_SUCCESS) { + ucs_log(log_level, "zeDeviceGetCommandQueueGroupProperties failed " + "for device %p with 0x%x", device, ret); + return UCS_ERR_IO_ERROR; + } + + if (num_queue_groups == 0) { + ucs_log(log_level, "device %p: no command queue groups", device); + return UCS_ERR_NO_DEVICE; + } + + queue_props = ucs_calloc(num_queue_groups, sizeof(*queue_props), + "ze_queue_group_props"); + if (queue_props == NULL) { + return UCS_ERR_NO_MEMORY; + } + + for (i = 0; i < num_queue_groups; i++) { queue_props[i].stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; - queue_props[i].pNext = NULL; } - /* Get queue group properties */ ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, queue_props); if (ret != ZE_RESULT_SUCCESS) { - ucs_error("zeDeviceGetCommandQueueGroupProperties(device=%p) failed: " - "0x%x", - device, ret); - return UCS_ERR_IO_ERROR; + ucs_log(log_level, "zeDeviceGetCommandQueueGroupProperties failed " + "for device %p with 0x%x", device, ret); + status = UCS_ERR_IO_ERROR; + goto out; } /* First pass: find dedicated copy engine (COPY flag but NOT COMPUTE) */ @@ -487,29 +508,36 @@ ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) && !(queue_props[i].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE)) { - *ordinal_p = i; - *num_queues_p = queue_props[i].numQueues; ucs_debug("device %p: using dedicated copy engine queue group " "%u (numQueues=%u)", device, i, queue_props[i].numQueues); - return UCS_OK; + goto out_found; } } /* Second pass: find any queue group that supports COPY */ for (i = 0; i < num_queue_groups; i++) { if (queue_props[i].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) { - *ordinal_p = i; - *num_queues_p = queue_props[i].numQueues; ucs_debug("device %p: using copy-capable queue group %u " "(numQueues=%u)", device, i, queue_props[i].numQueues); - return UCS_OK; + goto out_found; } } - ucs_error("device %p: no copy-capable queue group found", device); - return UCS_ERR_NO_DEVICE; + ucs_log(log_level, "device %p: no copy-capable queue group found", device); + status = UCS_ERR_NO_DEVICE; + goto out; + +out_found: + *ordinal_p = i; + if (num_queues_p != NULL) { + *num_queues_p = queue_props[i].numQueues; + } + status = UCS_OK; +out: + ucs_free(queue_props); + return status; } /** diff --git a/src/uct/ze/base/ze_base.h b/src/uct/ze/base/ze_base.h index 441fa11455e..54d4016d479 100644 --- a/src/uct/ze/base/ze_base.h +++ b/src/uct/ze/base/ze_base.h @@ -66,8 +66,11 @@ ze_device_handle_t uct_ze_base_get_device_handle_from_subdevice( ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device); /* Find the queue group to run memory copies on a device: a dedicated copy - * engine group if the device has one, or else any group which can copy */ + * engine group if the device has one, or else any group which can copy. + * A failure is logged at log_level, as the caller may be able to continue + * without the group. num_queues_p may be NULL. */ ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, uint32_t *ordinal_p, uint32_t *num_queues_p); diff --git a/src/uct/ze/copy/ze_copy_iface.c b/src/uct/ze/copy/ze_copy_iface.c index abe6f3d336a..f67305ec8cb 100644 --- a/src/uct/ze/copy/ze_copy_iface.c +++ b/src/uct/ze/copy/ze_copy_iface.c @@ -249,15 +249,15 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, uct_ze_copy_md_t); ze_command_queue_desc_t cq_desc = {}; ze_command_list_desc_t cl_desc = {}; - uint32_t ordinal, num_queues; uct_ze_copy_queue_t *queue; + uint32_t ordinal; ucs_status_t status; ze_result_t ret; /* Run on a copy engine rather than the default compute group: the * driver recommends it for copies to a peer or to host, and the copies * then do not compete with the application's kernels */ - status = uct_ze_base_get_copy_ordinal(device, &ordinal, &num_queues); + status = uct_ze_base_get_copy_ordinal(device, log_level, &ordinal, NULL); if (status != UCS_OK) { return status; } diff --git a/src/uct/ze/ze_ipc/ze_ipc_iface.c b/src/uct/ze/ze_ipc/ze_ipc_iface.c index 178351a2219..72d4dfe725f 100644 --- a/src/uct/ze/ze_ipc/ze_ipc_iface.c +++ b/src/uct/ze/ze_ipc/ze_ipc_iface.c @@ -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. */ @@ -557,7 +557,8 @@ static UCS_CLASS_INIT_FUNC(uct_ze_ipc_iface_t, uct_md_h md, uct_worker_h worker, self->next_cmd_list = 0; /* Find copy engine queue group ordinal and available queue count */ - status = uct_ze_base_get_copy_ordinal(self->ze_device, ©_ordinal, + status = uct_ze_base_get_copy_ordinal(self->ze_device, + UCS_LOG_LEVEL_ERROR, ©_ordinal, &num_queues); if (status != UCS_OK) { return status; diff --git a/test/gtest/uct/ze/test_ze.cc b/test/gtest/uct/ze/test_ze.cc index 790514bf07c..4de6749cce8 100644 --- a/test/gtest/uct/ze/test_ze.cc +++ b/test/gtest/uct/ze/test_ze.cc @@ -4,6 +4,10 @@ * See file LICENSE for terms. */ +#ifdef HAVE_CONFIG_H +# include "config.h" +#endif + #include "uct/test_p2p_rma.h" #include @@ -278,6 +282,7 @@ class test_ze_copy_rma : public uct_p2p_rma_test { return ucs_array_elem(&ze_iface()->queues, queue_index).device; } +#if HAVE_DECL_ZECOMMANDQUEUEGETORDINAL && HAVE_DECL_ZECOMMANDLISTGETORDINAL /* Check that the queue runs on a copy group, and on a dedicated copy * group when its device has one */ void check_copy_engine(const uct_ze_copy_queue_t &queue) @@ -287,11 +292,11 @@ class test_ze_copy_rma : public uct_p2p_rma_test { const ze_command_queue_group_property_flags_t copy_compute = ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY | ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE; - std::vector groups; uint32_t num_groups = 0; bool has_copy_only = false; uint32_t cmdq_ordinal = UINT32_MAX; uint32_t cmdl_ordinal = UINT32_MAX; + std::vector groups; ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGetCommandQueueGroupProperties(queue.device, @@ -330,6 +335,13 @@ class test_ze_copy_rma : public uct_p2p_rma_test { EXPECT_EQ(copy_only, groups[cmdq_ordinal].flags & copy_compute); } } +#else + void check_copy_engine(const uct_ze_copy_queue_t &queue) + { + UCS_TEST_SKIP_R("needs Level Zero API 1.9 " + "(zeCommandQueueGetOrdinal)"); + } +#endif void zcopy(void *buffer, void *remote_addr, size_t length, bool is_put) { @@ -545,15 +557,33 @@ UCS_TEST_P(test_ze_copy_rma, queue_on_copy_engine) ze_device_handle_t device = other_device(); std::vector host_buf(64); const uct_ze_copy_queue_t *queue; + ze_device_handle_t owner; - /* Also create the queue of another device, if there is one; the queue - * stays in the iface after the buffer is released */ + /* Also create the queue of another device, if there is one and the MD + * context resolves its memory; the queue stays in the iface after the + * buffer is released */ if (device != NULL) { external_buffer other_buf(device, host_buf.size()); - EXPECT_NE(0u, get_queue(other_buf.ptr(), host_buf.data())); + + owner = mem_device(other_buf.ptr()); + if (owner == NULL) { + UCS_TEST_MESSAGE << "the MD context does not resolve memory of " + "another context, checking the MD queue only"; + } else { + ASSERT_EQ(device, owner); + EXPECT_NE(0u, get_queue(other_buf.ptr(), host_buf.data())); + } } ucs_array_for_each(queue, &ze_iface()->queues) { + if (queue->cmdq == NULL) { + /* Device whose queue could not be created, copies on it use the + * MD device queue; see md_queue_when_create_fails */ + UCS_TEST_MESSAGE << "device " << queue->device + << ": no queue of its own"; + continue; + } + check_copy_engine(*queue); } } From f81cb38aa5acb55824142dffc53a782cffdfda75 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Sun, 1 Nov 2026 13:03:35 -0800 Subject: [PATCH 3/5] UCT/ZE/COPY: Keep copies within one device on the compute engine Routing all ze_copy transfers through copy engines also moved same-device device-to-device copies off the compute path, slowing some cache-resident transfers. Select a queue group with both COMPUTE and COPY flags when both buffers are device memory of the same device. Keep host, shared, and cross-device copies on the copy-engine path. On the tested devices this restores the earlier engine choice, with the trade-off that same-device copies can contend with application kernels. Key queue entries by device and engine, keeping entry 0 on the MD device's copy engine. Remove the single-device shortcut so buffer properties are checked even when only one device is visible. Share queue-group discovery between the compute and copy paths while preserving the dedicated-copy preference. Log queue-group selection without '=' and skip ordinal checks in tests when the Level Zero driver does not implement the query. Signed-off-by: Yaser Afshar --- src/uct/ze/base/ze_base.c | 98 ++++++++++++----- src/uct/ze/base/ze_base.h | 7 ++ src/uct/ze/copy/ze_copy_iface.c | 88 ++++++++++------ src/uct/ze/copy/ze_copy_iface.h | 24 +++-- test/gtest/uct/ze/test_ze.cc | 181 +++++++++++++++++++++++++++----- 5 files changed, 307 insertions(+), 91 deletions(-) diff --git a/src/uct/ze/base/ze_base.c b/src/uct/ze/base/ze_base.c index 0a6936c6ed9..f0d85e92c09 100644 --- a/src/uct/ze/base/ze_base.c +++ b/src/uct/ze/base/ze_base.c @@ -457,15 +457,42 @@ ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device) return UCS_SYS_DEVICE_ID_UNKNOWN; } -ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, - ucs_log_level_t log_level, - uint32_t *ordinal_p, - uint32_t *num_queues_p) +/* Index of the first queue group with all of the required flags and none of + * the avoided ones, or -1 */ +static int +uct_ze_base_find_queue_group(const ze_command_queue_group_properties_t *props, + uint32_t num_groups, + ze_command_queue_group_property_flags_t required, + ze_command_queue_group_property_flags_t avoided) +{ + uint32_t i; + + for (i = 0; i < num_groups; i++) { + if (((props[i].flags & required) == required) && + !(props[i].flags & avoided)) { + return i; + } + } + + return -1; +} + +/* Find the queue group of a device for the engine named in the log: a group + * with the required flags and without the avoided ones, or else any group + * with the required flags */ +static ucs_status_t +uct_ze_base_get_queue_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, + ze_command_queue_group_property_flags_t required, + ze_command_queue_group_property_flags_t avoided, + const char *engine_name, uint32_t *ordinal_p, + uint32_t *num_queues_p) { uint32_t num_queue_groups = 0; ze_command_queue_group_properties_t *queue_props; ucs_status_t status; ze_result_t ret; + int ordinal; uint32_t i; /* Query the group count first: the driver fills at most the given count @@ -502,37 +529,27 @@ ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, goto out; } - /* First pass: find dedicated copy engine (COPY flag but NOT COMPUTE) */ - for (i = 0; i < num_queue_groups; i++) { - if ((queue_props[i].flags & - ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) && - !(queue_props[i].flags & - ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE)) { - ucs_debug("device %p: using dedicated copy engine queue group " - "%u (numQueues=%u)", - device, i, queue_props[i].numQueues); - goto out_found; - } + ordinal = uct_ze_base_find_queue_group(queue_props, num_queue_groups, + required, avoided); + if ((ordinal < 0) && (avoided != 0)) { + ordinal = uct_ze_base_find_queue_group(queue_props, num_queue_groups, + required, 0); } - /* Second pass: find any queue group that supports COPY */ - for (i = 0; i < num_queue_groups; i++) { - if (queue_props[i].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY) { - ucs_debug("device %p: using copy-capable queue group %u " - "(numQueues=%u)", - device, i, queue_props[i].numQueues); - goto out_found; - } + if (ordinal < 0) { + ucs_log(log_level, "device %p: no %s queue group found", device, + engine_name); + status = UCS_ERR_NO_DEVICE; + goto out; } - ucs_log(log_level, "device %p: no copy-capable queue group found", device); - status = UCS_ERR_NO_DEVICE; - goto out; + ucs_debug("device %p: using queue group %d as %s, flags 0x%x, %u queues", + device, ordinal, engine_name, queue_props[ordinal].flags, + queue_props[ordinal].numQueues); -out_found: - *ordinal_p = i; + *ordinal_p = ordinal; if (num_queues_p != NULL) { - *num_queues_p = queue_props[i].numQueues; + *num_queues_p = queue_props[ordinal].numQueues; } status = UCS_OK; out: @@ -540,6 +557,29 @@ ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, return status; } +ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, + uint32_t *ordinal_p, + uint32_t *num_queues_p) +{ + /* A dedicated copy engine first, or else any group which can copy */ + return uct_ze_base_get_queue_ordinal( + device, log_level, ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY, + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE, "copy engine", + ordinal_p, num_queues_p); +} + +ucs_status_t uct_ze_base_get_compute_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, + uint32_t *ordinal_p) +{ + /* The group has to copy as well, since its queues run memory copies */ + return uct_ze_base_get_queue_ordinal( + device, log_level, ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE | + ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY, 0, "compute engine", + ordinal_p, NULL); +} + /** * Query MD resources - returns one MD per sub-device * This is correct because each sub-device has separate memory diff --git a/src/uct/ze/base/ze_base.h b/src/uct/ze/base/ze_base.h index 54d4016d479..e57b9057ca3 100644 --- a/src/uct/ze/base/ze_base.h +++ b/src/uct/ze/base/ze_base.h @@ -74,6 +74,13 @@ ucs_status_t uct_ze_base_get_copy_ordinal(ze_device_handle_t device, uint32_t *ordinal_p, uint32_t *num_queues_p); +/* Find the queue group to run memory copies on the compute engine of a + * device: a group with both the compute and the copy flag. A failure is + * logged at log_level. */ +ucs_status_t uct_ze_base_get_compute_ordinal(ze_device_handle_t device, + ucs_log_level_t log_level, + uint32_t *ordinal_p); + /* Get root device handle by ordinal, returns NULL if out of range */ ze_device_handle_t uct_ze_base_get_device(int ordinal); diff --git a/src/uct/ze/copy/ze_copy_iface.c b/src/uct/ze/copy/ze_copy_iface.c index f67305ec8cb..729800da270 100644 --- a/src/uct/ze/copy/ze_copy_iface.c +++ b/src/uct/ze/copy/ze_copy_iface.c @@ -237,11 +237,17 @@ static uct_iface_internal_ops_t uct_ze_copy_iface_internal_ops = { .ep_outstanding_purge = (uct_ep_outstanding_purge_func_t)ucs_empty_function_return_unsupported }; -/* Create the queue and list of a device; a driver failure is logged at - * log_level, as the caller may be able to continue without them */ +static const char *uct_ze_copy_engine_names[] = { + [UCT_ZE_COPY_ENGINE_COPY] = "copy", + [UCT_ZE_COPY_ENGINE_COMPUTE] = "compute" +}; + +/* Create the queue and list of an engine of a device; a driver failure is + * logged at log_level, as the caller may be able to continue without them */ static ucs_status_t uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, ze_device_handle_t device, + uct_ze_copy_engine_t engine, ucs_log_level_t log_level, unsigned *queue_index_p) { @@ -254,10 +260,18 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, ucs_status_t status; ze_result_t ret; - /* Run on a copy engine rather than the default compute group: the - * driver recommends it for copies to a peer or to host, and the copies - * then do not compete with the application's kernels */ - status = uct_ze_base_get_copy_ordinal(device, log_level, &ordinal, NULL); + if (engine == UCT_ZE_COPY_ENGINE_COPY) { + /* The driver recommends a copy engine for copies to a peer or to + * host, and there the copies do not compete with the application's + * kernels */ + status = uct_ze_base_get_copy_ordinal(device, log_level, &ordinal, + NULL); + } else { + /* The driver keeps its own copies within one device on the compute + * engine, which can be faster for cache-resident copies */ + status = uct_ze_base_get_compute_ordinal(device, log_level, &ordinal); + } + if (status != UCS_OK) { return status; } @@ -273,7 +287,8 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, &queue->cmdq); if (ret != ZE_RESULT_SUCCESS) { ucs_log(log_level, "ze_copy_iface: zeCommandQueueCreate failed for " - "device %p with 0x%x", device, ret); + "the %s engine of device %p with 0x%x", + uct_ze_copy_engine_names[engine], device, ret); goto err_pop; } @@ -281,11 +296,13 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, &queue->cmdl); if (ret != ZE_RESULT_SUCCESS) { ucs_log(log_level, "ze_copy_iface: zeCommandListCreate failed for " - "device %p with 0x%x", device, ret); + "the %s engine of device %p with 0x%x", + uct_ze_copy_engine_names[engine], device, ret); goto err_destroy_cmdq; } queue->device = device; + queue->engine = engine; *queue_index_p = ucs_array_length(&iface->queues) - 1; return UCS_OK; @@ -297,10 +314,11 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, } /* Return the device which owns the allocation, or NULL for host memory and - * for allocations the context does not recognize */ + * for allocations the context does not recognize; is_device_p tells whether + * the allocation is device memory, which cannot migrate to the host */ static ze_device_handle_t uct_ze_copy_iface_get_mem_device(uct_ze_copy_iface_t *iface, - const void *address) + const void *address, int *is_device_p) { uct_ze_copy_md_t *ze_md = ucs_derived_of(iface->super.md, uct_ze_copy_md_t); @@ -312,61 +330,72 @@ uct_ze_copy_iface_get_mem_device(uct_ze_copy_iface_t *iface, ret = zeMemGetAllocProperties(ze_md->ze_context, address, &props, &device); if ((ret != ZE_RESULT_SUCCESS) || (props.type == ZE_MEMORY_TYPE_UNKNOWN)) { + *is_device_p = 0; return NULL; } + *is_device_p = (props.type == ZE_MEMORY_TYPE_DEVICE); return device; } unsigned uct_ze_copy_iface_get_queue(uct_ze_copy_iface_t *iface, const void *src, const void *dst) { + uct_ze_copy_engine_t engine = UCT_ZE_COPY_ENGINE_COPY; uct_ze_copy_queue_t *queue; - ze_device_handle_t device; + ze_device_handle_t device, dst_device; + int src_is_device, dst_is_device; unsigned queue_index; ucs_status_t status; unsigned i; - /* With one device every device buffer is on the MD device */ - if (!iface->multi_device) { - return 0; - } - /* Run the copy on the device which owns the source buffer, so it reads * locally, or else on the owner of the destination. A queue on an * unrelated device has to reach both buffers over PCIe. */ - device = uct_ze_copy_iface_get_mem_device(iface, src); + device = uct_ze_copy_iface_get_mem_device(iface, src, &src_is_device); if (device == NULL) { - device = uct_ze_copy_iface_get_mem_device(iface, dst); + device = uct_ze_copy_iface_get_mem_device(iface, dst, &dst_is_device); if (device == NULL) { return 0; } + } else if (src_is_device) { + /* A copy within one device runs on the compute engine of that + * device, as the driver does with its own copies; it can be faster + * there for cache-resident copies. Shared memory may reside on the + * host, so it does not count. */ + dst_device = uct_ze_copy_iface_get_mem_device(iface, dst, + &dst_is_device); + if (dst_is_device && (dst_device == device)) { + engine = UCT_ZE_COPY_ENGINE_COMPUTE; + } } for (i = 0; i < ucs_array_length(&iface->queues); ++i) { queue = &ucs_array_elem(&iface->queues, i); - if (queue->device == device) { - /* A device without a queue of its own uses the MD device queue */ + if ((queue->device == device) && (queue->engine == engine)) { + /* An engine without a queue of its own uses the MD device queue */ return (queue->cmdq == NULL) ? 0 : i; } } - status = uct_ze_copy_iface_queue_create(iface, device, UCS_LOG_LEVEL_DEBUG, - &queue_index); + status = uct_ze_copy_iface_queue_create(iface, device, engine, + UCS_LOG_LEVEL_DEBUG, &queue_index); if (status == UCS_OK) { return queue_index; } /* Fall back to the MD device queue, which reaches the buffer if the MD - * device has peer access to its device. Record the device without a + * device has peer access to its device. Record the engine without a * queue, so that the next copy does not try to create one again. Without * the record the copy still runs, and the next one tries again. */ queue = ucs_array_append(&iface->queues, return 0); queue->device = device; + queue->engine = engine; queue->cmdq = NULL; queue->cmdl = NULL; - ucs_warn("ze_copy_iface: could not create a queue on device %p, using " - "the queue of the MD device", device); + ucs_warn("ze_copy_iface: could not create a %s engine queue on device %p, " + "using the copy engine queue of the MD device", + uct_ze_copy_engine_names[engine], device); return 0; } @@ -394,6 +423,7 @@ static UCS_CLASS_INIT_FUNC(uct_ze_copy_iface_t, uct_md_h md, ucs_array_init_dynamic(&self->queues); status = uct_ze_copy_iface_queue_create(self, ze_md->ze_device, + UCT_ZE_COPY_ENGINE_COPY, UCS_LOG_LEVEL_ERROR, &queue_index); if (status != UCS_OK) { ucs_array_cleanup_dynamic(&self->queues); @@ -403,12 +433,10 @@ static UCS_CLASS_INIT_FUNC(uct_ze_copy_iface_t, uct_md_h md, /* The fallback to the MD device queue uses element 0 */ ucs_assert(queue_index == 0); - self->id = ucs_generate_uuid((uintptr_t)self); - self->multi_device = (uct_ze_base_get_num_subdevices() > 1); + self->id = ucs_generate_uuid((uintptr_t)self); - ucs_debug("ze_copy_iface: initialized iface on device %p, %s", - ze_md->ze_device, - self->multi_device ? "queues per device" : "single device"); + ucs_debug("ze_copy_iface: initialized iface on device %p", + ze_md->ze_device); return UCS_OK; } diff --git a/src/uct/ze/copy/ze_copy_iface.h b/src/uct/ze/copy/ze_copy_iface.h index afb4d68f076..8d106d323cd 100644 --- a/src/uct/ze/copy/ze_copy_iface.h +++ b/src/uct/ze/copy/ze_copy_iface.h @@ -17,10 +17,21 @@ typedef uint64_t uct_ze_copy_iface_addr_t; -/* Command queue and list which run copies on one device; both handles are - * NULL for a device whose queue could not be created */ +/* Engine which runs the copies of a queue: the driver recommends a copy + * engine for copies to host or to another device, and keeps copies within + * one device on the compute engine, which can be faster for cache-resident + * same-device copies */ +typedef enum { + UCT_ZE_COPY_ENGINE_COPY, + UCT_ZE_COPY_ENGINE_COMPUTE +} uct_ze_copy_engine_t; + + +/* Command queue and list which run copies on one engine of one device; both + * handles are NULL for an engine whose queue could not be created */ typedef struct { ze_device_handle_t device; + uct_ze_copy_engine_t engine; ze_command_queue_handle_t cmdq; ze_command_list_handle_t cmdl; } uct_ze_copy_queue_t; @@ -33,13 +44,10 @@ UCS_ARRAY_DECLARE_TYPE(uct_ze_copy_queue_array_t, unsigned, typedef struct uct_ze_copy_iface { uct_base_iface_t super; uct_ze_copy_iface_addr_t id; - /* Element 0 is on the MD device, the rest are created on first use. - * Growth may relocate the elements and nothing here locks, so one iface - * must not be used concurrently. */ + /* Element 0 is the copy engine queue of the MD device, the rest are + * created on first use. Growth may relocate the elements and nothing + * here locks, so one iface must not be used concurrently. */ uct_ze_copy_queue_array_t queues; - /* More than one ZE device is visible, so a buffer may belong to a - * device other than the MD device */ - int multi_device; } uct_ze_copy_iface_t; diff --git a/test/gtest/uct/ze/test_ze.cc b/test/gtest/uct/ze/test_ze.cc index 4de6749cce8..c6eb562bead 100644 --- a/test/gtest/uct/ze/test_ze.cc +++ b/test/gtest/uct/ze/test_ze.cc @@ -126,29 +126,40 @@ class test_ze_copy_rma : public uct_p2p_rma_test { } } - /* Device memory allocated outside UCX in its own context, as an - * application such as a framework allocates it. ze_copy has to resolve + /* Device or shared memory allocated outside UCX in its own context, as + * an application such as a framework allocates it. ze_copy has to resolve * such memory through the MD context, as the MD mem_query already does, * so this deliberately does not allocate from the MD context. */ class external_buffer { public: - external_buffer(ze_device_handle_t device, size_t length) + external_buffer(ze_device_handle_t device, size_t length, + ze_memory_type_t type = ZE_MEMORY_TYPE_DEVICE) { ze_context_desc_t ctx_desc = {}; ze_device_mem_alloc_desc_t alloc_desc = {}; + ze_host_mem_alloc_desc_t host_desc = {}; + ze_result_t ret; ctx_desc.stype = ZE_STRUCTURE_TYPE_CONTEXT_DESC; alloc_desc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC; + host_desc.stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC; if (zeContextCreate(uct_ze_base_get_driver(), &ctx_desc, &m_context) != ZE_RESULT_SUCCESS) { UCS_TEST_ABORT("zeContextCreate failed"); } - if (zeMemAllocDevice(m_context, &alloc_desc, length, 64, device, - &m_ptr) != ZE_RESULT_SUCCESS) { + if (type == ZE_MEMORY_TYPE_SHARED) { + ret = zeMemAllocShared(m_context, &alloc_desc, &host_desc, + length, 64, device, &m_ptr); + } else { + ret = zeMemAllocDevice(m_context, &alloc_desc, length, 64, + device, &m_ptr); + } + + if (ret != ZE_RESULT_SUCCESS) { zeContextDestroy(m_context); - UCS_TEST_ABORT("zeMemAllocDevice failed"); + UCS_TEST_ABORT("zeMemAlloc failed"); } } @@ -258,20 +269,44 @@ class test_ze_copy_rma : public uct_p2p_rma_test { ASSERT_EQ(device, owner); } - void fill_and_copy_through(void *remote_addr, size_t length) + static void fill_pattern(std::vector &buf, uint8_t seed) { - std::vector host_buf(length), check_buf(length, 0); size_t i; - for (i = 0; i < length; ++i) { - host_buf[i] = (uint8_t)(i * 7 + 1); + for (i = 0; i < buf.size(); ++i) { + buf[i] = (uint8_t)(i * 7 + seed); } + } + void fill_and_copy_through(void *remote_addr, size_t length) + { + std::vector host_buf(length), check_buf(length, 0); + + fill_pattern(host_buf, 1); zcopy(host_buf.data(), remote_addr, length, true); zcopy(check_buf.data(), remote_addr, length, false); EXPECT_EQ(host_buf, check_buf); } + /* Move a pattern host -> src -> dst -> host, with the device-to-device + * step as a put from src or as a get into dst, and check it arrived */ + void copy_through_devices(void *src, void *dst, size_t length, + bool is_put, uint8_t seed) + { + std::vector host_buf(length), check_buf(length, 0); + + fill_pattern(host_buf, seed); + zcopy(host_buf.data(), src, length, true); + if (is_put) { + zcopy(src, dst, length, true); + } else { + zcopy(dst, src, length, false); + } + + zcopy(check_buf.data(), dst, length, false); + EXPECT_EQ(host_buf, check_buf); + } + unsigned get_queue(const void *src, const void *dst) { return uct_ze_copy_iface_get_queue(ze_iface(), src, dst); @@ -282,10 +317,15 @@ class test_ze_copy_rma : public uct_p2p_rma_test { return ucs_array_elem(&ze_iface()->queues, queue_index).device; } + uct_ze_copy_engine_t queue_engine(unsigned queue_index) + { + return ucs_array_elem(&ze_iface()->queues, queue_index).engine; + } + #if HAVE_DECL_ZECOMMANDQUEUEGETORDINAL && HAVE_DECL_ZECOMMANDLISTGETORDINAL - /* Check that the queue runs on a copy group, and on a dedicated copy - * group when its device has one */ - void check_copy_engine(const uct_ze_copy_queue_t &queue) + /* Check that the queue runs on the engine it was created for: a copy + * group, dedicated when its device has one, or a compute group */ + void check_queue_engine(const uct_ze_copy_queue_t &queue) { const ze_command_queue_group_property_flags_t copy_only = ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY; @@ -297,6 +337,7 @@ class test_ze_copy_rma : public uct_p2p_rma_test { uint32_t cmdq_ordinal = UINT32_MAX; uint32_t cmdl_ordinal = UINT32_MAX; std::vector groups; + ze_result_t ret; ASSERT_EQ(ZE_RESULT_SUCCESS, zeDeviceGetCommandQueueGroupProperties(queue.device, @@ -317,18 +358,38 @@ class test_ze_copy_rma : public uct_p2p_rma_test { } } - ASSERT_EQ(ZE_RESULT_SUCCESS, - zeCommandQueueGetOrdinal(queue.cmdq, &cmdq_ordinal)); - ASSERT_EQ(ZE_RESULT_SUCCESS, - zeCommandListGetOrdinal(queue.cmdl, &cmdl_ordinal)); - UCS_TEST_MESSAGE << "device " << queue.device << ": queue group " - << cmdq_ordinal << " of " << num_groups + /* The loader reports the ordinal only if the driver implements the + * query, which is newer than the rest of the API in use */ + ret = zeCommandQueueGetOrdinal(queue.cmdq, &cmdq_ordinal); + if (ret == ZE_RESULT_SUCCESS) { + ret = zeCommandListGetOrdinal(queue.cmdl, &cmdl_ordinal); + } + + if ((ret == ZE_RESULT_ERROR_UNSUPPORTED_FEATURE) || + (ret == ZE_RESULT_ERROR_UNSUPPORTED_VERSION)) { + UCS_TEST_SKIP_R("the Level Zero driver does not implement the " + "queue ordinal query"); + } + + ASSERT_EQ(ZE_RESULT_SUCCESS, ret); + UCS_TEST_MESSAGE << "device " << queue.device << ": " + << ((queue.engine == UCT_ZE_COPY_ENGINE_COMPUTE) ? + "compute" : "copy") + << " engine queue on group " << cmdq_ordinal + << " of " << num_groups << (has_copy_only ? ", device has a copy engine" : ""); /* The list must be submitted to the group of its queue */ EXPECT_EQ(cmdq_ordinal, cmdl_ordinal); ASSERT_LT(cmdq_ordinal, num_groups); + + if (queue.engine == UCT_ZE_COPY_ENGINE_COMPUTE) { + /* The group runs memory copies, so it has to copy as well */ + EXPECT_EQ(copy_compute, groups[cmdq_ordinal].flags & copy_compute); + return; + } + EXPECT_TRUE(groups[cmdq_ordinal].flags & ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY); if (has_copy_only) { @@ -336,7 +397,7 @@ class test_ze_copy_rma : public uct_p2p_rma_test { } } #else - void check_copy_engine(const uct_ze_copy_queue_t &queue) + void check_queue_engine(const uct_ze_copy_queue_t &queue) { UCS_TEST_SKIP_R("needs Level Zero API 1.9 " "(zeCommandQueueGetOrdinal)"); @@ -370,18 +431,19 @@ UCS_TEST_SKIP_COND_P(test_ze_copy_rma, queue_follows_buffer_device, const size_t length = UCS_MBYTE; ze_device_handle_t device = other_device(); std::vector host_buf(length), check_buf(length); - unsigned queue_index; + unsigned queue_index, compute_index; if (device == NULL) { UCS_TEST_SKIP_R("needs a second ZE device"); } external_buffer other_buf(device, length); + external_buffer other_buf2(device, length); external_buffer md_buf(ze_md()->ze_device, length); check_buffer_resolved(ze_md()->ze_device, md_buf.ptr()); check_buffer_resolved(device, other_buf.ptr()); - EXPECT_TRUE(ze_iface()->multi_device); + check_buffer_resolved(device, other_buf2.ptr()); /* Host-to-host copies and buffers on the MD device use the default * queue */ @@ -390,18 +452,31 @@ UCS_TEST_SKIP_COND_P(test_ze_copy_rma, queue_follows_buffer_device, EXPECT_EQ(0u, get_queue(host_buf.data(), md_buf.ptr())); EXPECT_EQ(ze_md()->ze_device, queue_device(0)); - /* Both directions select one queue on the buffer device */ + /* Both directions select one copy engine queue on the buffer device */ queue_index = get_queue(other_buf.ptr(), host_buf.data()); ASSERT_LT(queue_index, ucs_array_length(&ze_iface()->queues)); EXPECT_NE(0u, queue_index); EXPECT_EQ(device, queue_device(queue_index)); + EXPECT_EQ(UCT_ZE_COPY_ENGINE_COPY, queue_engine(queue_index)); EXPECT_EQ(queue_index, get_queue(host_buf.data(), other_buf.ptr())); /* Between two devices, the source device runs the copy */ EXPECT_EQ(queue_index, get_queue(other_buf.ptr(), md_buf.ptr())); EXPECT_EQ(0u, get_queue(md_buf.ptr(), other_buf.ptr())); + /* Within the other device, the copy runs on its compute engine */ + compute_index = get_queue(other_buf.ptr(), other_buf2.ptr()); + ASSERT_LT(compute_index, ucs_array_length(&ze_iface()->queues)); + EXPECT_NE(0u, compute_index); + EXPECT_NE(queue_index, compute_index); + EXPECT_EQ(device, queue_device(compute_index)); + EXPECT_EQ(UCT_ZE_COPY_ENGINE_COMPUTE, queue_engine(compute_index)); + + /* Data moves through the copy queue and through the compute queue of + * the other device */ fill_and_copy_through(other_buf.ptr(), length); + copy_through_devices(other_buf.ptr(), other_buf2.ptr(), length, true, 2); + check_queue_engine(ucs_array_elem(&ze_iface()->queues, compute_index)); } UCS_TEST_SKIP_COND_P(test_ze_copy_rma, md_queue_when_create_fails, @@ -437,6 +512,7 @@ UCS_TEST_SKIP_COND_P(test_ze_copy_rma, md_queue_when_create_fails, queue = ucs_array_append(&ze_iface()->queues, UCS_TEST_ABORT("cannot append a queue")); queue->device = device; + queue->engine = UCT_ZE_COPY_ENGINE_COPY; queue->cmdq = NULL; queue->cmdl = NULL; @@ -584,8 +660,65 @@ UCS_TEST_P(test_ze_copy_rma, queue_on_copy_engine) continue; } - check_copy_engine(*queue); + check_queue_engine(*queue); } } +UCS_TEST_SKIP_COND_P(test_ze_copy_rma, same_device_copy_on_compute_engine, + !check_caps(UCT_IFACE_FLAG_GET_ZCOPY | + UCT_IFACE_FLAG_PUT_ZCOPY)) +{ + const size_t length = UCS_MBYTE; + ze_device_handle_t device = ze_md()->ze_device; + std::vector host_buf(length); + unsigned queue_index; + + external_buffer src_buf(device, length); + external_buffer dst_buf(device, length); + + check_buffer_resolved(device, src_buf.ptr()); + check_buffer_resolved(device, dst_buf.ptr()); + + /* Copies to and from host memory stay on the copy engine queue */ + EXPECT_EQ(0u, get_queue(src_buf.ptr(), host_buf.data())); + EXPECT_EQ(0u, get_queue(host_buf.data(), src_buf.ptr())); + EXPECT_EQ(UCT_ZE_COPY_ENGINE_COPY, queue_engine(0)); + + /* A copy within the device selects the compute engine queue of the + * device, created on first use, in both directions */ + queue_index = get_queue(src_buf.ptr(), dst_buf.ptr()); + ASSERT_LT(queue_index, ucs_array_length(&ze_iface()->queues)); + EXPECT_NE(0u, queue_index); + EXPECT_EQ(device, queue_device(queue_index)); + EXPECT_EQ(UCT_ZE_COPY_ENGINE_COMPUTE, queue_engine(queue_index)); + EXPECT_EQ(queue_index, get_queue(dst_buf.ptr(), src_buf.ptr())); + + /* Move data between the two device buffers with a put and with a get */ + copy_through_devices(src_buf.ptr(), dst_buf.ptr(), length, true, 1); + copy_through_devices(src_buf.ptr(), dst_buf.ptr(), length, false, 3); + + check_queue_engine(ucs_array_elem(&ze_iface()->queues, queue_index)); +} + +UCS_TEST_SKIP_COND_P(test_ze_copy_rma, shared_memory_copy_on_copy_engine, + !check_caps(UCT_IFACE_FLAG_GET_ZCOPY | + UCT_IFACE_FLAG_PUT_ZCOPY)) +{ + const size_t length = 64 * UCS_KBYTE; + ze_device_handle_t device = ze_md()->ze_device; + + external_buffer device_buf(device, length); + external_buffer shared_buf(device, length, ZE_MEMORY_TYPE_SHARED); + + check_buffer_resolved(device, device_buf.ptr()); + check_buffer_resolved(device, shared_buf.ptr()); + + /* Shared memory may reside on the host, so a copy with it is not a copy + * within the device and stays on the copy engine queue */ + EXPECT_EQ(0u, get_queue(device_buf.ptr(), shared_buf.ptr())); + EXPECT_EQ(0u, get_queue(shared_buf.ptr(), device_buf.ptr())); + + fill_and_copy_through(shared_buf.ptr(), length); +} + UCT_INSTANTIATE_ZE_TEST_CASE(test_ze_copy_rma) From 8e25e4bc572b3c82e9c1fdebfcc7cf135ab29e93 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Sun, 1 Nov 2026 15:29:25 -0800 Subject: [PATCH 4/5] UCT/ZE/COPY: Address review: reuse UCT_ZE_FUNC, remove count getter Use UCT_ZE_FUNC for both queue-group queries instead of duplicating error checking and logging. Remove the test-only sub-device count getter; the test now enumerates sub-devices until the lookup returns NULL. Signed-off-by: Yaser Afshar --- src/uct/ze/base/ze_base.c | 31 +++++++++---------------------- src/uct/ze/base/ze_base.h | 3 --- test/gtest/uct/ze/test_ze.cc | 28 +++++++++++++++++++--------- 3 files changed, 28 insertions(+), 34 deletions(-) diff --git a/src/uct/ze/base/ze_base.c b/src/uct/ze/base/ze_base.c index f0d85e92c09..458e20e85bf 100644 --- a/src/uct/ze/base/ze_base.c +++ b/src/uct/ze/base/ze_base.c @@ -398,15 +398,6 @@ int uct_ze_base_get_num_devices(void) return uct_ze_base.num_devices; } -int uct_ze_base_get_num_subdevices(void) -{ - if (uct_ze_base_init() != ZE_RESULT_SUCCESS) { - return 0; - } - - return ze_num_subdevices; -} - ze_device_handle_t uct_ze_base_get_device_handle_from_subdevice(const uct_ze_subdevice_t *subdevice) { @@ -491,18 +482,16 @@ uct_ze_base_get_queue_ordinal(ze_device_handle_t device, uint32_t num_queue_groups = 0; ze_command_queue_group_properties_t *queue_props; ucs_status_t status; - ze_result_t ret; int ordinal; uint32_t i; /* Query the group count first: the driver fills at most the given count * and may leave the count itself unchanged */ - ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, - NULL); - if (ret != ZE_RESULT_SUCCESS) { - ucs_log(log_level, "zeDeviceGetCommandQueueGroupProperties failed " - "for device %p with 0x%x", device, ret); - return UCS_ERR_IO_ERROR; + status = UCT_ZE_FUNC(zeDeviceGetCommandQueueGroupProperties( + device, &num_queue_groups, NULL), + log_level); + if (status != UCS_OK) { + return status; } if (num_queue_groups == 0) { @@ -520,12 +509,10 @@ uct_ze_base_get_queue_ordinal(ze_device_handle_t device, queue_props[i].stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; } - ret = zeDeviceGetCommandQueueGroupProperties(device, &num_queue_groups, - queue_props); - if (ret != ZE_RESULT_SUCCESS) { - ucs_log(log_level, "zeDeviceGetCommandQueueGroupProperties failed " - "for device %p with 0x%x", device, ret); - status = UCS_ERR_IO_ERROR; + status = UCT_ZE_FUNC(zeDeviceGetCommandQueueGroupProperties( + device, &num_queue_groups, queue_props), + log_level); + if (status != UCS_OK) { goto out; } diff --git a/src/uct/ze/base/ze_base.h b/src/uct/ze/base/ze_base.h index e57b9057ca3..64e8f2d419e 100644 --- a/src/uct/ze/base/ze_base.h +++ b/src/uct/ze/base/ze_base.h @@ -87,9 +87,6 @@ ze_device_handle_t uct_ze_base_get_device(int ordinal); int uct_ze_base_get_num_devices(void); -/* Number of sub-devices over all root devices, one memory domain each */ -int uct_ze_base_get_num_subdevices(void); - ucs_status_t uct_ze_base_query_md_resources(uct_component_h component, diff --git a/test/gtest/uct/ze/test_ze.cc b/test/gtest/uct/ze/test_ze.cc index c6eb562bead..b924a2879b4 100644 --- a/test/gtest/uct/ze/test_ze.cc +++ b/test/gtest/uct/ze/test_ze.cc @@ -193,6 +193,21 @@ class test_ze_copy_rma : public uct_p2p_rma_test { return ucs_derived_of(sender().iface(), uct_ze_copy_iface_t); } + /* All sub-devices: the global IDs are contiguous from 0, and the lookup + * returns NULL past the last one */ + static std::vector subdevices() + { + std::vector result; + const uct_ze_subdevice_t *subdevice; + + while ((subdevice = uct_ze_base_get_subdevice_by_global_id( + (int)result.size())) != NULL) { + result.push_back(subdevice); + } + + return result; + } + /* Return a device other than the MD device, preferring one on another * card over another tile of the same card, or NULL if there is none */ ze_device_handle_t other_device() @@ -200,23 +215,18 @@ class test_ze_copy_rma : public uct_p2p_rma_test { ze_device_handle_t md_device = ze_md()->ze_device; const uct_ze_subdevice_t *md_subdevice = NULL; ze_device_handle_t other_tile = NULL; - int num_subdevices = - uct_ze_base_get_num_subdevices(); - const uct_ze_subdevice_t *subdevice; + const auto all_subdevices = subdevices(); ze_device_handle_t device; - int id; - for (id = 0; id < num_subdevices; ++id) { - subdevice = uct_ze_base_get_subdevice_by_global_id(id); + for (const uct_ze_subdevice_t *subdevice : all_subdevices) { if (uct_ze_base_get_device_handle_from_subdevice(subdevice) == md_device) { md_subdevice = subdevice; } } - for (id = 0; id < num_subdevices; ++id) { - subdevice = uct_ze_base_get_subdevice_by_global_id(id); - device = uct_ze_base_get_device_handle_from_subdevice(subdevice); + for (const uct_ze_subdevice_t *subdevice : all_subdevices) { + device = uct_ze_base_get_device_handle_from_subdevice(subdevice); if ((device == NULL) || (device == md_device)) { continue; } From 906b7c23b3a8daa8836b7f4a6026b105c464bda8 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Mon, 2 Nov 2026 06:30:27 -0800 Subject: [PATCH 5/5] UCT/ZE/COPY: Address review: log allocation failure, document engine once Signed-off-by: Yaser Afshar --- config/m4/ze.m4 | 2 +- src/uct/ze/base/ze_base.c | 2 ++ src/uct/ze/copy/ze_copy_iface.c | 10 ++-------- src/uct/ze/copy/ze_copy_iface.h | 6 +++--- test/gtest/uct/ze/test_ze.cc | 5 ++--- 5 files changed, 10 insertions(+), 15 deletions(-) diff --git a/config/m4/ze.m4 b/config/m4/ze.m4 index 02e766df468..2db57d5528c 100644 --- a/config/m4/ze.m4 +++ b/config/m4/ze.m4 @@ -1,5 +1,5 @@ # -# Copyright (c) Intel Corporation, 2023. ALL RIGHTS RESERVED. +# Copyright (c) Intel Corporation, 2023-2026. ALL RIGHTS RESERVED. # See file LICENSE for terms. # diff --git a/src/uct/ze/base/ze_base.c b/src/uct/ze/base/ze_base.c index 458e20e85bf..7b3e40bc79c 100644 --- a/src/uct/ze/base/ze_base.c +++ b/src/uct/ze/base/ze_base.c @@ -502,6 +502,8 @@ uct_ze_base_get_queue_ordinal(ze_device_handle_t device, queue_props = ucs_calloc(num_queue_groups, sizeof(*queue_props), "ze_queue_group_props"); if (queue_props == NULL) { + ucs_log(log_level, "device %p: failed to allocate %u queue group " + "properties", device, num_queue_groups); return UCS_ERR_NO_MEMORY; } diff --git a/src/uct/ze/copy/ze_copy_iface.c b/src/uct/ze/copy/ze_copy_iface.c index 729800da270..c5df6790058 100644 --- a/src/uct/ze/copy/ze_copy_iface.c +++ b/src/uct/ze/copy/ze_copy_iface.c @@ -261,14 +261,9 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface, ze_result_t ret; if (engine == UCT_ZE_COPY_ENGINE_COPY) { - /* The driver recommends a copy engine for copies to a peer or to - * host, and there the copies do not compete with the application's - * kernels */ status = uct_ze_base_get_copy_ordinal(device, log_level, &ordinal, NULL); } else { - /* The driver keeps its own copies within one device on the compute - * engine, which can be faster for cache-resident copies */ status = uct_ze_base_get_compute_ordinal(device, log_level, &ordinal); } @@ -360,9 +355,8 @@ unsigned uct_ze_copy_iface_get_queue(uct_ze_copy_iface_t *iface, } } else if (src_is_device) { /* A copy within one device runs on the compute engine of that - * device, as the driver does with its own copies; it can be faster - * there for cache-resident copies. Shared memory may reside on the - * host, so it does not count. */ + * device. Shared memory may reside on the host, so it does not + * count. */ dst_device = uct_ze_copy_iface_get_mem_device(iface, dst, &dst_is_device); if (dst_is_device && (dst_device == device)) { diff --git a/src/uct/ze/copy/ze_copy_iface.h b/src/uct/ze/copy/ze_copy_iface.h index 8d106d323cd..fc3b55d17ae 100644 --- a/src/uct/ze/copy/ze_copy_iface.h +++ b/src/uct/ze/copy/ze_copy_iface.h @@ -18,9 +18,9 @@ typedef uint64_t uct_ze_copy_iface_addr_t; /* Engine which runs the copies of a queue: the driver recommends a copy - * engine for copies to host or to another device, and keeps copies within - * one device on the compute engine, which can be faster for cache-resident - * same-device copies */ + * engine for copies to host or to another device, where they do not compete + * with the application's kernels, and keeps copies within one device on the + * compute engine, which can be faster for cache-resident same-device copies */ typedef enum { UCT_ZE_COPY_ENGINE_COPY, UCT_ZE_COPY_ENGINE_COMPUTE diff --git a/test/gtest/uct/ze/test_ze.cc b/test/gtest/uct/ze/test_ze.cc index b924a2879b4..c4c9f0a0c37 100644 --- a/test/gtest/uct/ze/test_ze.cc +++ b/test/gtest/uct/ze/test_ze.cc @@ -337,9 +337,9 @@ class test_ze_copy_rma : public uct_p2p_rma_test { * group, dedicated when its device has one, or a compute group */ void check_queue_engine(const uct_ze_copy_queue_t &queue) { - const ze_command_queue_group_property_flags_t copy_only = + constexpr ze_command_queue_group_property_flags_t copy_only = ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY; - const ze_command_queue_group_property_flags_t copy_compute = + constexpr ze_command_queue_group_property_flags_t copy_compute = ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COPY | ZE_COMMAND_QUEUE_GROUP_PROPERTY_FLAG_COMPUTE; uint32_t num_groups = 0; @@ -354,7 +354,6 @@ class test_ze_copy_rma : public uct_p2p_rma_test { &num_groups, NULL)); groups.resize(num_groups); for (auto &group : groups) { - group = {}; group.stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES; } ASSERT_EQ(ZE_RESULT_SUCCESS,