UCP/CORE: Add shared topology assignment policy

src/ucp/core/ucp_context.c

@@ -108,6 +108,14 @@ static const char *ucp_fence_modes[] = {
     [UCP_FENCE_MODE_LAST]     = NULL
 };
 
+static const char *ucp_gpu_net_device_modes[] = {
+    [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED]             = "not_shared",
+    [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP]        = "not_shared_flip",
+    [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN] = "not_shared_round_robin",
+    [UCP_GPU_NET_DEVICE_MODE_SHARED]                 = "shared",
+    [UCP_GPU_NET_DEVICE_MODE_LAST]                   = NULL
+};
+
 static const char *ucp_rndv_modes[] = {
     [UCP_RNDV_MODE_AUTO]         = "auto",
     [UCP_RNDV_MODE_GET_ZCOPY]    = "get_zcopy",
@@ -615,6 +623,16 @@ static ucs_config_field_t ucp_context_config_table[] = {
    ucs_offsetof(ucp_context_config_t, connect_all_to_all),
    UCS_CONFIG_TYPE_BOOL},
 
+  {"GPU_NET_DEVICE_MODE", "not_shared",
+   "Select how network devices are assigned to GPUs.\n"
+   " - not_shared             : use the default FLIP assignment.\n"
+   " - not_shared_flip        : use the FLIP assignment.\n"
+   " - not_shared_round_robin : use the round-robin assignment.\n"
+   " - shared                 : assign every network device in a topology\n"
+   "                            group to every GPU in that group.",
+   ucs_offsetof(ucp_context_config_t, gpu_net_device_mode),
+   UCS_CONFIG_TYPE_ENUM(ucp_gpu_net_device_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"

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 network device assignment mode */
+    ucp_gpu_net_device_mode_t              gpu_net_device_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

@@ -33,6 +33,22 @@ typedef struct {
 } ucp_nics_string_buffers_t;
 
 
+ucp_gpu_nic_assignment_policy_t
+ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode)
+{
+    switch (mode) {
+    case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED:
+    case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP:
+        return UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP;
+    case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN:
+        return UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN;
+    case UCP_GPU_NET_DEVICE_MODE_SHARED:
+        return UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED;
+    default:
+        ucs_fatal("invalid gpu network device mode %d", (int)mode);
+    }
+}
+
 static size_t
 ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_policy_t policy,
                                    size_t num_gpus, size_t nic_idx)
@@ -327,16 +343,39 @@ 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_group_shared(ucp_gpu_nic_assignment_t *assignment,
+                                        const ucs_topo_group_t *group)
+{
+    ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
+    const ucs_topo_group_element_t *gpu, *nic;
+
+    ucs_array_for_each(gpu, &group->gpus) {
+        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);
+
+        /* Assign all NICs in the group to the GPU. */
+        ucs_array_for_each(nic, &group->nics) {
+            ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic);
+        }
+    }
+}
+
+static void ucp_gpu_nic_assignment_add_group_not_shared(
+        ucp_gpu_nic_assignment_t *assignment, const ucs_topo_group_t *group,
+        ucp_gpu_nic_assignment_policy_t policy)
 {
     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;
+    const ucs_topo_group_element_t *gpu, *nic;
+    size_t nic_idx, gpu_idx;
+
+    ucs_assertv(policy != UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED,
+                "invalid policy: %d", (int)policy);
 
     ucs_array_for_each_index(nic, nic_idx, &group->nics) {
         gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, nic_idx);
@@ -354,6 +393,18 @@ ucp_gpu_nic_assignment_add_group(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)
+{
+    if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) {
+        ucp_gpu_nic_assignment_add_group_shared(assignment, group);
+    } else {
+        ucp_gpu_nic_assignment_add_group_not_shared(assignment, group, policy);
+    }
+}
+
 ucs_status_t
 ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
                              ucp_gpu_nic_assignment_policy_t policy,

src/ucp/core/ucp_gpu_nic_assignment.h

@@ -7,6 +7,8 @@
 #ifndef UCP_GPU_NIC_ASSIGNMENT_H_
 #define UCP_GPU_NIC_ASSIGNMENT_H_
 
+#include "ucp_types.h"
+
 #include <ucs/datastruct/static_bitmap.h>
 #include <ucs/sys/topo/base/topo_groups.h>
 
@@ -52,10 +54,26 @@ typedef enum {
      */
     UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN,
 
+    /**
+     * Assign every NIC, with all of its ports, to every GPU in the group.
+     */
+    UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED,
+
     UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST
 } ucp_gpu_nic_assignment_policy_t;
 
 
+/**
+ * Get the assignment policy of a GPU network device mode.
+ *
+ * @param [in] mode  GPU network device mode.
+ *
+ * @return Assignment policy corresponding to @a mode.
+ */
+ucp_gpu_nic_assignment_policy_t
+ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode);
+
+
 /**
  * Build GPU-to-NIC assignments from topology groups.
  *

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.
  */
@@ -233,6 +233,18 @@ typedef enum {
 } ucp_fence_mode_t;
 
 
+/**
+ * GPU network device assignment mode.
+ */
+typedef enum {
+    UCP_GPU_NET_DEVICE_MODE_NOT_SHARED,
+    UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP,
+    UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN,
+    UCP_GPU_NET_DEVICE_MODE_SHARED,
+    UCP_GPU_NET_DEVICE_MODE_LAST
+} ucp_gpu_net_device_mode_t;
+
+
 /**
  * Communication scheme in RNDV protocol.
  */