UCP/CORE: Add shared topology assignment policy

src/ucp/core/ucp_context.c

@@ -108,6 +108,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",
@@ -162,6 +171,7 @@ const char *ucp_extra_op_attr_flags_names[] = {
 static UCS_CONFIG_DEFINE_ARRAY(memunit_sizes, sizeof(size_t),
                                UCS_CONFIG_TYPE_MEMUNITS);
 
+/* clang-format off */
 static ucs_config_field_t ucp_context_config_table[] = {
   {"SELECT_DISTANCE_MD", "cuda_cpy",
    "MD whose distance is queried when evaluating transport selection score",
@@ -615,6 +625,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_NIC_ASSIGNMENT_MODE", "auto",
+   "Select how NICs are assigned to GPUs.\n"
+   " - auto        : enable automatically on supported hardware, otherwise disable.\n"
+   " - off         : disable GPU-to-NIC assignment, fallback to free lane selection.\n"
+   " - flip        : require the not shared flip assignment policy.\n"
+   " - round_robin : require the not shared round-robin assignment policy.\n"
+   " - shared      : require the shared assignment policy.",
+   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"
@@ -650,6 +670,7 @@ static ucs_config_field_t ucp_context_config_table[] = {
 
   {NULL}
 };
+/* clang-format on */
 
 static ucs_config_field_t ucp_config_table[] = {
   {"NET_DEVICES", UCP_RSC_CONFIG_ALL,

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

@@ -33,6 +33,35 @@ typedef struct {
 } ucp_nics_string_buffers_t;
 
 
+int ucp_gpu_nic_assignment_policy_resolve(
+        ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model,
+        ucp_gpu_nic_assignment_policy_t *policy_p)
+{
+    ucs_assert(policy_p != NULL);
+
+    switch (mode) {
+    case UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO:
+        if (cpu_model != UCS_CPU_MODEL_NVIDIA_VERA) {
+            return 0;
+        }
+        *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP;
+        return 1;
+    case UCP_GPU_NIC_ASSIGNMENT_MODE_OFF:
+        return 0;
+    case UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP:
+        *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP;
+        return 1;
+    case UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN:
+        *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN;
+        return 1;
+    case UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED:
+        *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED;
+        return 1;
+    default:
+        ucs_fatal("invalid gpu-nic assignment 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 +356,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 +406,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,9 @@
 #ifndef UCP_GPU_NIC_ASSIGNMENT_H_
 #define UCP_GPU_NIC_ASSIGNMENT_H_
 
+#include "ucp_types.h"
+
+#include <ucs/arch/cpu.h>
 #include <ucs/datastruct/static_bitmap.h>
 #include <ucs/sys/topo/base/topo_groups.h>
 
@@ -52,10 +55,30 @@ 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;
 
 
+/**
+ * Resolve the GPU-to-NIC assignment policy.
+ *
+ * @param [in]  mode       GPU-to-NIC assignment mode.
+ * @param [in]  cpu_model  CPU model used to resolve automatic mode.
+ * @param [out] policy_p   Resolved assignment policy,
+ *                         populated only when assignment is enabled.
+ *
+ * @return Nonzero if assignment is enabled, or zero if it is disabled.
+ */
+int ucp_gpu_nic_assignment_policy_resolve(
+        ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model,
+        ucp_gpu_nic_assignment_policy_t *policy_p);
+
+
 /**
  * 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,19 @@ typedef enum {
 } ucp_fence_mode_t;
 
 
+/**
+ * GPU-to-NIC assignment mode.
+ */
+typedef enum {
+    UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO,
+    UCP_GPU_NIC_ASSIGNMENT_MODE_OFF,
+    UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP,
+    UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN,
+    UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED,
+    UCP_GPU_NIC_ASSIGNMENT_MODE_LAST
+} ucp_gpu_nic_assignment_mode_t;
+
+
 /**
  * Communication scheme in RNDV protocol.
  */