UCP/CORE: Add shared topology assignment policy

src/ucp/core/ucp_context.c

@@ -110,6 +110,15 @@ static const char *ucp_fence_modes[] = {
     [UCP_FENCE_MODE_LAST]     = NULL
 };
 
+static const char *ucp_gpu_nic_assignment_modes[] = {
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO]        = "auto",
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_OFF]         = "off",
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP]        = "flip",
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN] = "round_robin",
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED]      = "shared",
+    [UCP_GPU_NIC_ASSIGNMENT_MODE_LAST]        = NULL
+};
+
 static const char *ucp_rndv_modes[] = {
     [UCP_RNDV_MODE_AUTO]         = "auto",
     [UCP_RNDV_MODE_GET_ZCOPY]    = "get_zcopy",
@@ -617,6 +626,26 @@ static ucs_config_field_t ucp_context_config_table[] = {
    ucs_offsetof(ucp_context_config_t, connect_all_to_all),
    UCS_CONFIG_TYPE_BOOL},
 
+  {"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"
+   "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"
+   " - auto        : use 'flip' on hardware with a known GPU-NIC topology,\n"
+   "                 otherwise 'off'.\n"
+   " - off         : do not assign; select lanes from all NICs.\n"
+   " - flip        : assign each NIC to a single GPU, walking the N GPUs of\n"
+   "                 the group forward then backward:\n"
+   "                 0, 1, .., N-1, N-1, .., 1, 0, 0, 1, ..\n"
+   " - round_robin : assign each NIC to a single GPU, walking the N GPUs of\n"
+   "                 the group in ascending order:\n"
+   "                 0, 1, .., N-1, 0, 1, .., N-1, 0, ..\n"
+   " - shared      : assign all NICs of a group to every GPU of that group.",
+   ucs_offsetof(ucp_context_config_t, gpu_nic_assignment_mode),
+   UCS_CONFIG_TYPE_ENUM(ucp_gpu_nic_assignment_modes)},
+
   {"SINGLE_NET_DEVICE", "n",
    "Restrict each protocol's lanes to one network device.\n"
    "The device is picked among the closest eligible network devices using \n"
@@ -2755,22 +2784,32 @@ ucp_version_check(unsigned api_major_version, unsigned api_minor_version)
 }
 
 static ucs_status_t
-ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t **assignment_p)
+ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
+                                    ucp_gpu_nic_assignment_t **assignment_p)
 {
     ucp_gpu_nic_assignment_t *assignment;
     ucs_topo_groups_t groups;
     ucs_status_t status;
 
     *assignment_p = NULL;
 
-    /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */
-    if (ucs_arch_get_cpu_model() != UCS_CPU_MODEL_NVIDIA_VERA) {
-        ucs_debug("gpu-nic assignment is not supported on %s architecture, "
-                  "skipping",
-                  ucs_cpu_model_name());
+    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) {
+            ucs_debug("gpu-nic assignment is not supported on %s "
+                      "architecture, skipping",
+                      ucs_cpu_model_name());
+            return UCS_OK;
+        }
+
+        mode = UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP;
+    } else if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_OFF) {
+        ucs_debug("gpu-nic assignment is disabled by configuration");
         return UCS_OK;
     }
 
+    ucs_debug("gpu-nic assignment mode %s", ucp_gpu_nic_assignment_modes[mode]);
+
     status = ucs_topo_build_groups(&groups);
     if (status != UCS_OK) {
         return status;
@@ -2788,9 +2827,7 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t **assignment_p)
         goto out_release_groups;
     }
 
-    status = ucp_gpu_nic_assignment_build(&groups,
-                                          UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP,
-                                          assignment);
+    status = ucp_gpu_nic_assignment_build(&groups, mode, assignment);
     if (status != UCS_OK) {
         ucs_free(assignment);
         goto out_release_groups;
@@ -2857,7 +2894,9 @@ 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->gpu_nic_assignment);
+    status = ucp_context_gpu_nic_assignment_init(
+            context->config.ext.gpu_nic_assignment_mode,
+            &context->gpu_nic_assignment);
     if (status != UCS_OK) {
         goto err_free_res;
     }

src/ucp/core/ucp_context.h

@@ -255,6 +255,8 @@ typedef struct ucp_context_config {
     /** Extend endpoint lanes connections of each local device to all remote
      *  devices */
     int                                    connect_all_to_all;
+    /** GPU-to-NIC assignment mode */
+    ucp_gpu_nic_assignment_mode_t          gpu_nic_assignment_mode;
     /** Restrict lanes to one network device per protocol */
     int                                    proto_use_single_net_device;
     /** Max HCAs for GPU memory registration: auto=closest, N=limit, inf=all */

src/ucp/core/ucp_gpu_nic_assignment.c

@@ -34,20 +34,20 @@ typedef struct {
 
 
 static size_t
-ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_policy_t policy,
+ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_mode_t mode,
                                    size_t num_gpus, size_t nic_idx)
 {
     size_t gpu_idx, leg;
 
     ucs_assert(num_gpus > 0);
 
     gpu_idx = nic_idx % num_gpus;
-    if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN) {
+    if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN) {
         return gpu_idx;
     }
 
-    ucs_assertv(policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP,
-                "invalid policy: %d", (int)policy);
+    ucs_assertv(mode == UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, "invalid mode: %d",
+                (int)mode);
     leg = nic_idx / num_gpus;
     if ((leg % 2) != 0) {
         gpu_idx = num_gpus - 1 - gpu_idx;
@@ -327,45 +327,61 @@ ucp_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t *assignment,
 }
 
 static void
-ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment,
-                                 const ucs_topo_group_t *group,
-                                 ucp_gpu_nic_assignment_policy_t policy)
+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)
 {
-    size_t num_gpus = ucs_array_length(&group->gpus);
     ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
-    const ucs_topo_group_element_t *nic;
-    const ucs_topo_group_element_t *gpu;
-    size_t nic_idx;
-    size_t gpu_idx;
 
-    ucs_array_for_each_index(nic, nic_idx, &group->nics) {
-        gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, nic_idx);
-        gpu     = &ucs_array_elem(&group->gpus, gpu_idx);
+    ucs_assert((gpu->num_sys_devs > 0) &&
+               (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN));
 
-        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*,
+            ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0]));
+    ucs_assert(nic_sys_dev_bitmap != NULL);
 
-        nic_sys_dev_bitmap = ucs_const_cast(
-                ucp_gpu_nic_sys_dev_bitmap_t*,
-                ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0]));
-        ucs_assert(nic_sys_dev_bitmap != NULL);
+    ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic);
+}
 
-        ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic);
+static void
+ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment,
+                                 const ucs_topo_group_t *group,
+                                 ucp_gpu_nic_assignment_mode_t mode)
+{
+    size_t num_gpus = ucs_array_length(&group->gpus);
+    const ucs_topo_group_element_t *gpu, *nic;
+    size_t nic_idx, gpu_idx;
+
+    ucs_array_for_each_index(nic, nic_idx, &group->nics) {
+        if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED) {
+            /* Assign the NIC to every GPU in the group. */
+            ucs_array_for_each(gpu, &group->gpus) {
+                ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic);
+            }
+        } else {
+            gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(mode, num_gpus,
+                                                         nic_idx);
+            gpu     = &ucs_array_elem(&group->gpus, gpu_idx);
+            ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic);
+        }
     }
 }
 
 ucs_status_t
 ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
-                             ucp_gpu_nic_assignment_policy_t policy,
+                             ucp_gpu_nic_assignment_mode_t mode,
                              ucp_gpu_nic_assignment_t *assignment_p)
 {
     ucp_gpu_nic_assignment_t assignment;
     const ucs_topo_group_t *group;
     ucs_status_t status;
 
     ucs_assert(groups != NULL);
-    ucs_assertv(policy < UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST,
-                "invalid gpu-nic assignment policy %d", (int)policy);
+    ucs_assertv((mode == UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP) ||
+                        (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN) ||
+                        (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED),
+                "invalid gpu-nic assignment mode %d", (int)mode);
 
     status = ucp_gpu_nic_assignment_init(&assignment, groups);
     if (status != UCS_OK) {
@@ -377,7 +393,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
             continue;
         }
 
-        ucp_gpu_nic_assignment_add_group(&assignment, group, policy);
+        ucp_gpu_nic_assignment_add_group(&assignment, group, mode);
     }
 
     ucp_gpu_nic_assignment_log(&assignment, groups);

src/ucp/core/ucp_gpu_nic_assignment.h

@@ -37,43 +37,22 @@ struct ucp_gpu_nic_assignment {
 };
 
 
-typedef enum {
-    /**
-     * Assign each NIC, with all of its ports, going forward and then backward
-     * across the GPUs.
-     * The group's NICs are assigned to the following GPU indices in order:
-     * 0, 1, ..., num_gpus-1, num_gpus-1, ..., 1, 0, 0, 1, ...
-     */
-    UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP,
-
-    /**
-     * Assign each NIC, with all of its ports, to GPUs repeatedly in ascending
-     * order.
-     * The group's NICs are assigned to the following GPU indices in order:
-     * 0, 1, ..., num_gpus-1, 0, 1, ..., num_gpus-1, 0, 1, ...
-     */
-    UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN,
-
-    UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST
-} ucp_gpu_nic_assignment_policy_t;
-
-
 /**
  * Build GPU-to-NIC assignments from topology groups.
  *
  * @note The groups must be disjoint: each sys_dev may appear in at most one
  *       topology group.
  *
  * @param [in]  groups        Topology groups to build assignments from.
- * @param [in]  policy        Assignment policy.
+ * @param [in]  mode          Assignment mode: flip, round-robin or shared.
  * @param [out] assignment_p  Completed assignment. Updated only on success.
  *
  * @return UCS_OK on success, or an error status if assignment construction
  *         failed.
  */
 ucs_status_t
 ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
-                             ucp_gpu_nic_assignment_policy_t policy,
+                             ucp_gpu_nic_assignment_mode_t mode,
                              ucp_gpu_nic_assignment_t *assignment_p);
 
 

src/ucp/core/ucp_types.h

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2017. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -234,6 +234,39 @@ typedef enum {
 } ucp_fence_mode_t;
 
 
+/**
+ * GPU-to-NIC assignment mode.
+ */
+typedef enum {
+    /** Use flip on hardware with a known GPU-NIC topology, otherwise off */
+    UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO,
+
+    /** Do not assign NICs to GPUs, select lanes from all NICs */
+    UCP_GPU_NIC_ASSIGNMENT_MODE_OFF,
+
+    /**
+     * Assign each NIC, with all of its ports, going forward and then backward
+     * 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, ...
+     */
+    UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP,
+
+    /**
+     * Assign each NIC, with all of its ports, to the N GPUs of the group
+     * repeatedly in ascending order.
+     * The group's NICs are assigned to the following GPU indices in order:
+     * 0, 1, ..., N-1, 0, 1, ..., N-1, 0, 1, ...
+     */
+    UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN,
+
+    /** Assign every NIC, with all of its ports, to every GPU in the group */
+    UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED,
+
+    UCP_GPU_NIC_ASSIGNMENT_MODE_LAST
+} ucp_gpu_nic_assignment_mode_t;
+
+
 /**
  * Communication scheme in RNDV protocol.
  */