UCT/ZE/COPY: Run copies on a copy engine

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.
 #
 
@@ -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 <level_zero/ze_api.h>]])])
+
          CPPFLAGS="$save_CPPFLAGS"
          LDFLAGS="$save_LDFLAGS"
          LIBS="$save_LIBS"

src/uct/ze/base/ze_base.c

@@ -12,6 +12,7 @@
 #include <ucs/sys/module.h>
 #include <ucs/sys/string.h>
 #include <ucs/debug/log.h>
+#include <ucs/debug/memtrack_int.h>
 #include <ucs/sys/math.h>
 
 
@@ -397,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)
 {
@@ -456,6 +448,127 @@ ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device)
     return UCS_SYS_DEVICE_ID_UNKNOWN;
 }
 
+/* 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;
+    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 */
+    status = UCT_ZE_FUNC(zeDeviceGetCommandQueueGroupProperties(
+                                 device, &num_queue_groups, NULL),
+                         log_level);
+    if (status != UCS_OK) {
+        return status;
+    }
+
+    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) {
+        ucs_log(log_level, "device %p: failed to allocate %u queue group "
+                "properties", device, num_queue_groups);
+        return UCS_ERR_NO_MEMORY;
+    }
+
+    for (i = 0; i < num_queue_groups; i++) {
+        queue_props[i].stype = ZE_STRUCTURE_TYPE_COMMAND_QUEUE_GROUP_PROPERTIES;
+    }
+
+    status = UCT_ZE_FUNC(zeDeviceGetCommandQueueGroupProperties(
+                                 device, &num_queue_groups, queue_props),
+                         log_level);
+    if (status != UCS_OK) {
+        goto out;
+    }
+
+    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);
+    }
+
+    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_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);
+
+    *ordinal_p = ordinal;
+    if (num_queues_p != NULL) {
+        *num_queues_p = queue_props[ordinal].numQueues;
+    }
+    status = UCS_OK;
+out:
+    ucs_free(queue_props);
+    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

src/uct/ze/base/ze_base.h

@@ -65,15 +65,28 @@ 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.
+ * 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);
+
+/* 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);
 
 
 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,

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)
 {
@@ -250,27 +256,48 @@ uct_ze_copy_iface_queue_create(uct_ze_copy_iface_t *iface,
     ze_command_queue_desc_t cq_desc = {};
     ze_command_list_desc_t cl_desc  = {};
     uct_ze_copy_queue_t *queue;
+    uint32_t ordinal;
+    ucs_status_t status;
     ze_result_t ret;
 
+    if (engine == UCT_ZE_COPY_ENGINE_COPY) {
+        status = uct_ze_base_get_copy_ordinal(device, log_level, &ordinal,
+                                              NULL);
+    } else {
+        status = uct_ze_base_get_compute_ordinal(device, log_level, &ordinal);
+    }
+
+    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,
                                &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;
     }
 
     ret = zeCommandListCreate(ze_md->ze_context, device, &cl_desc,
                               &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;
 
@@ -282,10 +309,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);
@@ -297,61 +325,71 @@ 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. 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;
 }
 
@@ -379,6 +417,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);
@@ -388,12 +427,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;
 }