UCP/CORE: Ignore DPU in topology assignments

src/ucp/core/ucp_context.c

@@ -629,7 +629,10 @@ static ucs_config_field_t ucp_context_config_table[] = {
   {"GPU_NIC_ASSIGNMENT_MODE", "auto",
    "Assign NICs to GPUs within each topology group, and restrict the lanes\n"
    "for a GPU's memory to the NICs assigned to that GPU.\n"
-   "All ports of a NIC are assigned together.\n"
+   "All ports of a NIC are assigned together, and all GPUs with the same PCI\n"
+   "address are assigned together (e.g. MLOPart partitions of a GPU).\n"
+   "NICs without memory registration and DPUs are skipped during assignment.\n"
+   "NICs excluded by UCX_NET_DEVICES or UCX_TLS are ignored.\n"
    "The 'flip', 'round_robin' and 'shared' modes apply on any hardware.\n"
    "With 'flip' and 'round_robin', a group with fewer NICs than GPUs leaves\n"
    "some GPUs without any NICs.\n"
@@ -2783,16 +2786,67 @@ ucp_version_check(unsigned api_major_version, unsigned api_minor_version)
     ucs_debug("Configured with: %s", UCX_CONFIGURE_FLAGS);
 }
 
-static ucs_status_t
-ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
-                                    ucp_gpu_nic_assignment_t **assignment_p)
+static int
+ucp_context_is_net_device_assignable(ucp_context_h context,
+                                     const ucp_tl_resource_desc_t *rsc)
+{
+    const uct_tl_resource_desc_t *resource = &rsc->tl_rsc;
+    const ucp_tl_md_t *md                  = &context->tl_mds[rsc->md_index];
+
+    if ((resource->dev_type != UCT_DEVICE_TYPE_NET) ||
+        (resource->sys_device == UCS_SYS_DEVICE_ID_UNKNOWN)) {
+        return 0;
+    }
+
+    if (rsc->flags & UCP_TL_RSC_FLAG_AUX) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is ignored for gpu-nic assignment: auxiliary transport",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    if (!(md->attr.flags & UCT_MD_FLAG_REG)) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is ignored for gpu-nic assignment: no memory registration",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    if (ucs_topo_sys_device_get_flags(resource->sys_device) &
+        UCS_TOPO_DEVICE_FLAG_DPU) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is ignored for gpu-nic assignment: dpu device",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    ucs_trace(UCT_TL_RESOURCE_DESC_FMT " is used for gpu-nic assignment",
+              UCT_TL_RESOURCE_DESC_ARG(resource));
+    return 1;
+}
+
+static void ucp_context_gpu_nic_assignment_net_device_filter_init(
+        ucp_context_h context, ucs_sys_device_bitmap_t *net_device_filter)
+{
+    const ucp_tl_resource_desc_t *rsc;
+
+    ucs_carray_for_each(rsc, context->tl_rscs, context->num_tls) {
+        if (ucp_context_is_net_device_assignable(context, rsc)) {
+            UCS_STATIC_BITMAP_SET(net_device_filter, rsc->tl_rsc.sys_device);
+        }
+    }
+}
+
+static ucs_status_t ucp_context_gpu_nic_assignment_init(ucp_context_h context)
 {
+    ucp_gpu_nic_assignment_mode_t mode =
+            context->config.ext.gpu_nic_assignment_mode;
+    ucs_sys_device_bitmap_t net_device_filter =
+            UCS_STATIC_BITMAP_ZERO_INITIALIZER;
     ucp_gpu_nic_assignment_t *assignment;
     ucs_topo_groups_t groups;
     ucs_status_t status;
 
-    *assignment_p = NULL;
-
     if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO) {
         /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */
         if (ucs_arch_get_cpu_model() != UCS_CPU_MODEL_NVIDIA_VERA) {
@@ -2810,7 +2864,23 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
 
     ucs_debug("gpu-nic assignment mode %s", ucp_gpu_nic_assignment_modes[mode]);
 
-    status = ucs_topo_build_groups(&groups);
+    ucp_context_gpu_nic_assignment_net_device_filter_init(context,
+                                                          &net_device_filter);
+    if (UCS_STATIC_BITMAP_IS_ZERO(net_device_filter)) {
+        if (context->config.ext.gpu_nic_assignment_mode ==
+            UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO) {
+            ucs_diag("gpu-nic assignment is disabled: no assignable network "
+                     "devices");
+            return UCS_OK;
+        }
+
+        ucs_error("no assignable network devices for gpu-nic assignment mode "
+                  "%s, set UCX_GPU_NIC_ASSIGNMENT_MODE=off to disable it",
+                  ucp_gpu_nic_assignment_modes[mode]);
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    status = ucs_topo_build_groups(&net_device_filter, &groups);
     if (status != UCS_OK) {
         return status;
     }
@@ -2833,7 +2903,7 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
         goto out_release_groups;
     }
 
-    *assignment_p = assignment;
+    context->gpu_nic_assignment = assignment;
 
 out_release_groups:
     ucs_topo_release_groups(&groups);
@@ -2894,9 +2964,7 @@ ucs_status_t ucp_init_version(unsigned api_major_version, unsigned api_minor_ver
         goto err_thread_lock_finalize;
     }
 
-    status = ucp_context_gpu_nic_assignment_init(
-            context->config.ext.gpu_nic_assignment_mode,
-            &context->gpu_nic_assignment);
+    status = ucp_context_gpu_nic_assignment_init(context);
     if (status != UCS_OK) {
         goto err_free_res;
     }

src/ucp/core/ucp_gpu_nic_assignment.c

@@ -57,7 +57,7 @@ ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_mode_t mode,
 }
 
 static void
-ucp_gpu_nic_bitmap_add_nic(ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap,
+ucp_gpu_nic_bitmap_add_nic(ucs_sys_device_bitmap_t *nic_sys_dev_bitmap,
                            const ucs_topo_group_element_t *nic)
 {
     const ucs_sys_device_t *nic_sys_dev;
@@ -72,7 +72,7 @@ ucp_gpu_nic_bitmap_add_nic(ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap,
     }
 }
 
-const ucp_gpu_nic_sys_dev_bitmap_t *
+const ucs_sys_device_bitmap_t *
 ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
                               ucs_sys_device_t gpu_sys_dev)
 {
@@ -135,12 +135,12 @@ static void
 ucp_gpu_nic_assignment_log_gpu(const ucp_gpu_nic_assignment_t *assignment,
                                const ucs_topo_group_t *group,
                                const ucs_topo_group_element_t *gpu,
-                               ucp_gpu_nic_sys_dev_bitmap_t *assigned_nics)
+                               ucs_sys_device_bitmap_t *assigned_nics)
 {
     UCS_STRING_BUFFER_ONSTACK(gpu_strb, 128);
     UCP_GPU_NIC_STRING_BUFFERS_ONSTACK(nic_strbs, 256);
     char bdf_name[UCS_SYS_BDF_NAME_MAX];
-    const ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
+    const ucs_sys_device_bitmap_t *nic_sys_dev_bitmap;
     const ucs_topo_group_element_t *nic;
 
     ucs_assert((gpu->num_sys_devs > 0) &&
@@ -182,7 +182,7 @@ ucp_gpu_nic_assignment_log_gpu(const ucp_gpu_nic_assignment_t *assignment,
 
 static void ucp_gpu_nic_assignment_log_unassigned_nics(
         const ucs_topo_groups_t *groups,
-        const ucp_gpu_nic_sys_dev_bitmap_t *assigned_nics)
+        const ucs_sys_device_bitmap_t *assigned_nics)
 {
     UCP_GPU_NIC_STRING_BUFFERS_ONSTACK(nic_strbs, 256);
     const ucs_topo_group_element_t *nic;
@@ -211,8 +211,7 @@ static void
 ucp_gpu_nic_assignment_log(const ucp_gpu_nic_assignment_t *assignment,
                            const ucs_topo_groups_t *groups)
 {
-    ucp_gpu_nic_sys_dev_bitmap_t assigned_nics =
-            UCS_STATIC_BITMAP_ZERO_INITIALIZER;
+    ucs_sys_device_bitmap_t assigned_nics = UCS_STATIC_BITMAP_ZERO_INITIALIZER;
     const ucs_topo_group_element_t *gpu;
     const ucs_topo_group_t *group;
 
@@ -331,13 +330,13 @@ ucp_gpu_nic_assignment_add_nic_to_gpu(ucp_gpu_nic_assignment_t *assignment,
                                       const ucs_topo_group_element_t *gpu,
                                       const ucs_topo_group_element_t *nic)
 {
-    ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
+    ucs_sys_device_bitmap_t *nic_sys_dev_bitmap;
 
     ucs_assert((gpu->num_sys_devs > 0) &&
                (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN));
 
     nic_sys_dev_bitmap = ucs_const_cast(
-            ucp_gpu_nic_sys_dev_bitmap_t*,
+            ucs_sys_device_bitmap_t*,
             ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0]));
     ucs_assert(nic_sys_dev_bitmap != NULL);
 
@@ -404,7 +403,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
     return UCS_OK;
 }
 
-int ucp_gpu_nic_bitmap_get(const ucp_gpu_nic_sys_dev_bitmap_t *bitmap,
+int ucp_gpu_nic_bitmap_get(const ucs_sys_device_bitmap_t *bitmap,
                            ucs_sys_device_t net_sys_dev)
 {
     ucs_assert(bitmap != NULL);

src/ucp/core/ucp_gpu_nic_assignment.h

@@ -21,13 +21,9 @@ BEGIN_C_DECLS
 #define UCP_GPU_NIC_BITMAP_INDEX_INVALID UCS_SYS_DEVICE_ID_COUNT
 
 
-typedef ucs_static_bitmap_s(UCS_SYS_DEVICE_ID_COUNT)
-        ucp_gpu_nic_sys_dev_bitmap_t;
-
-
 struct ucp_gpu_nic_assignment {
     /* NIC sys_dev bitmaps referenced by bitmap_idx_by_gpu_sys_dev. */
-    ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmaps;
+    ucs_sys_device_bitmap_t      *nic_sys_dev_bitmaps;
 
     /* Number of entries in nic_sys_dev_bitmaps. */
     size_t                       num_bitmaps;
@@ -66,7 +62,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
  *         has no assigned NICs. Returns NULL if @a gpu_sys_dev is unknown or
  *         is not represented in the assignment.
  */
-const ucp_gpu_nic_sys_dev_bitmap_t *
+const ucs_sys_device_bitmap_t *
 ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
                               ucs_sys_device_t gpu_sys_dev);
 
@@ -80,7 +76,7 @@ ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
  * @return Nonzero if @a net_sys_dev is present in @a bitmap, or zero if the
  *         system device is unknown or the bit is not set.
  */
-int ucp_gpu_nic_bitmap_get(const ucp_gpu_nic_sys_dev_bitmap_t *bitmap,
+int ucp_gpu_nic_bitmap_get(const ucs_sys_device_bitmap_t *bitmap,
                            ucs_sys_device_t net_sys_dev);
 
 

src/ucp/core/ucp_types.h

@@ -249,6 +249,9 @@ typedef enum {
      * across the N GPUs of the group.
      * The group's NICs are assigned to the following GPU indices in order:
      * 0, 1, ..., N-1, N-1, ..., 1, 0, 0, 1, ...
+     * This targets architectures where only some of the NICs share a PCIe
+     * switch with a local SSD, so that each GPU gets NICs both with and without
+     * direct SSD access.
      */
     UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP,
 

src/ucs/sys/topo/base/topo.c

@@ -517,6 +517,7 @@ ucs_topo_find_device_by_bus_id_value(const ucs_sys_bus_id_t *bus_id,
         device->pci_id          = pci_id;
         device->user_value      = user_value;
         device->device_class    = UCS_TOPO_DEVICE_CLASS_UNKNOWN;
+        device->flags           = 0;
         device->class_ordinal   = UCS_SYS_DEVICE_ORDINAL_INVALID;
         device->sibling_role    = UCS_TOPO_SIBLING_ROLE_NONE;
         device->sibling_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
@@ -1119,6 +1120,51 @@ ucs_status_t ucs_topo_sys_device_set_class(ucs_sys_device_t sys_dev,
     return status;
 }
 
+ucs_status_t
+ucs_topo_sys_device_add_flags(ucs_sys_device_t sys_dev, unsigned flags)
+{
+    ucs_status_t status = UCS_OK;
+
+    if (sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
+        ucs_error("system device %d is unknown", sys_dev);
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    ucs_spin_lock(&ucs_topo_global_ctx.lock);
+
+    if (sys_dev >= ucs_topo_global_ctx.num_devices) {
+        ucs_error("system device %d is invalid (max: %u)", sys_dev,
+                  ucs_topo_global_ctx.num_devices);
+        status = UCS_ERR_INVALID_PARAM;
+        goto out_unlock;
+    }
+
+    ucs_topo_global_ctx.devices[sys_dev].flags |= flags;
+
+out_unlock:
+    ucs_spin_unlock(&ucs_topo_global_ctx.lock);
+    return status;
+}
+
+unsigned ucs_topo_sys_device_get_flags(ucs_sys_device_t sys_dev)
+{
+    unsigned flags;
+
+    if (sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
+        return 0;
+    }
+
+    ucs_spin_lock(&ucs_topo_global_ctx.lock);
+    if (sys_dev < ucs_topo_global_ctx.num_devices) {
+        flags = ucs_topo_global_ctx.devices[sys_dev].flags;
+    } else {
+        flags = 0;
+    }
+    ucs_spin_unlock(&ucs_topo_global_ctx.lock);
+
+    return flags;
+}
+
 ucs_status_t
 ucs_topo_device_class_mark_incomplete(ucs_topo_device_class_t device_class)
 {
@@ -1446,14 +1492,16 @@ static void ucs_topo_release_devices()
     }
 }
 
-ucs_status_t ucs_topo_build_groups(ucs_topo_groups_t *groups_p)
+ucs_status_t
+ucs_topo_build_groups(const ucs_sys_device_bitmap_t *net_device_filter,
+                      ucs_topo_groups_t *groups_p)
 {
     ucs_status_t status;
 
     ucs_spin_lock(&ucs_topo_global_ctx.lock);
     status = ucs_topo_build_groups_inner(ucs_topo_global_ctx.devices,
                                          ucs_topo_global_ctx.num_devices,
-                                         groups_p);
+                                         net_device_filter, groups_p);
     ucs_spin_unlock(&ucs_topo_global_ctx.lock);
 
     return status;

src/ucs/sys/topo/base/topo.h

@@ -106,6 +106,15 @@ typedef enum {
 } ucs_topo_device_class_t;
 
 
+/**
+ * @ingroup UCS_RESOURCE
+ * Properties of a system device, set by the owning transport.
+ */
+typedef enum {
+    UCS_TOPO_DEVICE_FLAG_DPU = UCS_BIT(0) /**< DPU device (e.g. BlueField) */
+} ucs_topo_device_flags_t;
+
+
 /**
  * @ingroup UCS_RESOURCE
  * Global state of the topology subsystem.
@@ -473,6 +482,27 @@ ucs_status_t
 ucs_topo_sys_device_set_class(ucs_sys_device_t sys_dev,
                               ucs_topo_device_class_t device_class);
 
+/**
+ * Add flags to a given system device.
+ *
+ * @param [in]  sys_dev  System device index.
+ * @param [in]  flags    Flags to add, from @ref ucs_topo_device_flags_t.
+ *
+ * @return UCS_OK on success, error otherwise.
+ */
+ucs_status_t
+ucs_topo_sys_device_add_flags(ucs_sys_device_t sys_dev, unsigned flags);
+
+/**
+ * Get the flags of a given system device.
+ *
+ * @param [in]  sys_dev  System device index.
+ *
+ * @return Flags from @ref ucs_topo_device_flags_t, or 0 if the system device
+ *         is unknown or invalid.
+ */
+unsigned ucs_topo_sys_device_get_flags(ucs_sys_device_t sys_dev);
+
 /**
  * Mark a device class as incomplete.
  *