UCP/PROTO: Added option to use single network device. (#10870)

rakhmets · web-flow · commit c73ee1a087e4 · 2025-09-24T09:05:46.000+02:00
diff --git a/src/ucp/api/ucp.h b/src/ucp/api/ucp.h
@@ -125,7 +125,8 @@ enum ucp_params_field {
     UCP_PARAM_FIELD_MT_WORKERS_SHARED = UCS_BIT(5), /**< mt_workers_shared */
     UCP_PARAM_FIELD_ESTIMATED_NUM_EPS = UCS_BIT(6), /**< estimated_num_eps */
     UCP_PARAM_FIELD_ESTIMATED_NUM_PPN = UCS_BIT(7), /**< estimated_num_ppn */
-    UCP_PARAM_FIELD_NAME              = UCS_BIT(8)  /**< name */
+    UCP_PARAM_FIELD_NAME              = UCS_BIT(8), /**< name */
+    UCP_PARAM_FIELD_NODE_LOCAL_ID     = UCS_BIT(9)
 };
 
 
@@ -1144,6 +1145,16 @@ typedef struct ucp_params {
      * unique name will be created for you.
      */
     const char                         *name;
+
+    /**
+     * An optimization hint for a single node. For example, when used from MPI or
+     * OpenSHMEM libraries, this number will specify the local identificator on
+     * a single node in the job. Does not affect semantics, only transport
+     * selection criteria and the resulting performance.
+     * The value can be also set by the UCX_LOCAL_NODE_ID environment variable,
+     * which will override the id set by @e node_local_id
+     */
+    size_t                             node_local_id;
 } ucp_params_t;
 
 
diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c
@@ -577,6 +577,16 @@ static ucs_config_field_t ucp_context_config_table[] = {
    ucs_offsetof(ucp_context_config_t, connect_all_to_all),
    UCS_CONFIG_TYPE_BOOL},
 
+  {"SINGLE_NET_DEVICE", "n", "Use only one network device for all protocols.",
+   ucs_offsetof(ucp_context_config_t, proto_use_single_net_device),
+   UCS_CONFIG_TYPE_BOOL},
+
+  {"NODE_LOCAL_ID", "auto",
+   "An optimization hint for the local identificator on a single node. Does \n"
+   "not affect semantics, only transport selection criteria and the \n"
+   "resulting performance.",
+   ucs_offsetof(ucp_context_config_t, node_local_id), UCS_CONFIG_TYPE_ULUNITS},
+
   {NULL}
 };
 
@@ -1999,6 +2009,9 @@ static void ucp_apply_params(ucp_context_h context, const ucp_params_t *params,
                                                         estimated_num_ppn,
                                                         ESTIMATED_NUM_PPN, 1);
 
+    context->config.node_local_id = UCP_PARAM_FIELD_VALUE(params, node_local_id,
+                                                          NODE_LOCAL_ID, 0);
+
     if ((params->field_mask & UCP_PARAM_FIELD_MT_WORKERS_SHARED) &&
         params->mt_workers_shared) {
         context->mt_lock.mt_type = mt_type;
@@ -2108,6 +2121,12 @@ static ucs_status_t ucp_fill_config(ucp_context_h context,
     ucs_debug("estimated number of endpoints per node is %d",
               context->config.est_num_ppn);
 
+    if (context->config.ext.node_local_id != UCS_ULUNITS_AUTO) {
+        /* node_local_id was set via the env variable. Override current value */
+        context->config.node_local_id = context->config.ext.node_local_id;
+    }
+    ucs_debug("node local id is %lu", context->config.node_local_id);
+
     if (UCS_CONFIG_DBL_IS_AUTO(context->config.ext.bcopy_bw)) {
         /* bcopy_bw wasn't set via the env variable. Calculate the value */
         if (context->config.ext.proto_enable) {
diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h
@@ -215,6 +215,10 @@ typedef struct ucp_context_config {
     /** Extend endpoint lanes connections of each local device to all remote
      *  devices */
     int                                    connect_all_to_all;
+    /** Use only one network device for all protocols */
+    int                                    proto_use_single_net_device;
+    /** Local identificator on a single node */
+    unsigned long                          node_local_id;
 } ucp_context_config_t;
 
 
@@ -415,6 +419,9 @@ typedef struct ucp_context {
         /* How many endpoints are expected to be created on single node */
         int                       est_num_ppn;
 
+        /* Local identificator on a single node */
+        unsigned long             node_local_id;
+
         struct {
             size_t                         size;    /* Request size for user */
             ucp_request_init_callback_t    init;    /* Initialization user callback */
diff --git a/src/ucp/proto/proto_common.c b/src/ucp/proto/proto_common.c
@@ -152,6 +152,37 @@ ucp_proto_common_get_sys_dev(const ucp_proto_init_params_t *params,
     return params->worker->context->tl_rscs[rsc_index].tl_rsc.sys_device;
 }
 
+int ucp_proto_common_add_unique_sys_dev(ucs_sys_device_t sys_dev,
+                                        ucs_sys_device_t *sys_devs,
+                                        ucp_lane_index_t *num_sys_devs,
+                                        ucp_lane_index_t max_sys_devs)
+{
+    ucp_lane_index_t i;
+
+    for (i = 0; i < *num_sys_devs; ++i) {
+        if (sys_dev == sys_devs[i]) {
+            return 0; /* Already exists */
+        }
+    }
+
+    if (*num_sys_devs < max_sys_devs) {
+        sys_devs[(*num_sys_devs)++] = sys_dev;
+        return 1; /* Added */
+    }
+
+    return 0; /* No space */
+}
+
+ucp_lane_index_t
+ucp_proto_common_select_sys_dev_by_node_id(const ucp_proto_init_params_t *params,
+                                           ucp_lane_index_t num_sys_devs)
+{
+    if (num_sys_devs == 0) {
+        return 0;
+    }
+    return params->worker->context->config.node_local_id % num_sys_devs;
+}
+
 /* Pack/unpack local distance to make it equal to the remote one */
 static void
 ucp_proto_common_fp8_pack_unpack_distance(ucs_sys_dev_distance_t *distance)
diff --git a/src/ucp/proto/proto_common.h b/src/ucp/proto/proto_common.h
@@ -255,6 +255,15 @@ ucs_sys_device_t
 ucp_proto_common_get_sys_dev(const ucp_proto_init_params_t *params,
                              ucp_lane_index_t lane);
 
+int ucp_proto_common_add_unique_sys_dev(ucs_sys_device_t sys_dev,
+                                        ucs_sys_device_t *sys_devs,
+                                        ucp_lane_index_t *num_sys_devs,
+                                        ucp_lane_index_t max_sys_devs);
+
+ucp_lane_index_t
+ucp_proto_common_select_sys_dev_by_node_id(const ucp_proto_init_params_t *params,
+                                           ucp_lane_index_t num_sys_devs);
+
 
 void ucp_proto_common_get_lane_distance(const ucp_proto_init_params_t *params,
                                         ucp_lane_index_t lane,
diff --git a/src/ucp/proto/proto_common.inl b/src/ucp/proto/proto_common.inl
@@ -407,4 +407,26 @@ ucp_proto_common_get_dev_index(const ucp_proto_init_params_t *params,
     return params->worker->context->tl_rscs[rsc_index].dev_index;
 }
 
+static UCS_F_ALWAYS_INLINE const uct_tl_resource_desc_t *
+ucp_proto_common_get_tl_rsc(const ucp_proto_init_params_t *params,
+                            ucp_lane_index_t lane)
+{
+    ucp_rsc_index_t rsc_index = ucp_proto_common_get_rsc_index(params, lane);
+    return &params->worker->context->tl_rscs[rsc_index].tl_rsc;
+}
+
+static UCS_F_ALWAYS_INLINE int
+ucp_proto_common_is_net_dev(const ucp_proto_init_params_t *params,
+                            ucp_lane_index_t lane)
+{
+    return ucp_proto_common_get_tl_rsc(params, lane)->dev_type ==
+           UCT_DEVICE_TYPE_NET;
+}
+
+static UCS_F_ALWAYS_INLINE int
+ucp_proto_common_bandwidth_equal(double bw1, double bw2)
+{
+    return fabs(bw1 - bw2) <= UCP_PROTO_PERF_EPSILON;
+}
+
 #endif
diff --git a/src/ucp/proto/proto_multi.c b/src/ucp/proto/proto_multi.c
@@ -151,6 +151,62 @@ ucp_proto_multi_init_flush_sys_dev_mask(const ucp_rkey_config_key_t *key)
     return UCS_BIT(key->sys_dev & ~UCP_SYS_DEVICE_FLUSH_BIT);
 }
 
+static ucp_lane_index_t ucp_proto_multi_filter_net_devices(
+        ucp_lane_index_t num_lanes, const ucp_proto_init_params_t *params,
+        const ucp_proto_common_tl_perf_t *tl_perfs, int fixed_first_lane,
+        ucp_lane_index_t *lanes, ucp_proto_perf_node_t **perf_nodes)
+{
+    ucp_lane_index_t num_max_bw_devs = 0;
+    double max_bandwidth;
+    ucp_lane_index_t i, lane, seed, num_filtered_lanes;
+    ucp_lane_map_t lane_map;
+    ucs_sys_device_t sys_dev;
+    ucs_sys_device_t sys_devs[UCP_PROTO_MAX_LANES];
+    const uct_tl_resource_desc_t *tl_rsc;
+
+    for (lane_map = 0, max_bandwidth = 0, i = 0; i < num_lanes; ++i) {
+        lane = lanes[i];
+        if (!ucp_proto_common_is_net_dev(params, lane)) {
+            continue;
+        }
+
+        lane_map     |= UCS_BIT(lane);
+        max_bandwidth = ucs_max(max_bandwidth, tl_perfs[lane].bandwidth);
+    }
+
+    ucs_for_each_bit(lane, lane_map) {
+        if (!ucp_proto_common_bandwidth_equal(tl_perfs[lane].bandwidth,
+                                              max_bandwidth)) {
+            continue;
+        }
+
+        sys_dev = ucp_proto_common_get_sys_dev(params, lane);
+        ucp_proto_common_add_unique_sys_dev(sys_dev, sys_devs, &num_max_bw_devs,
+                                            UCP_PROTO_MAX_LANES);
+    }
+
+    if (num_max_bw_devs == 0) {
+        return num_lanes;
+    }
+
+    seed = ucp_proto_common_select_sys_dev_by_node_id(params, num_max_bw_devs);
+
+    for (i = !!fixed_first_lane, num_filtered_lanes = i; i < num_lanes; ++i) {
+        lane   = lanes[i];
+        tl_rsc = ucp_proto_common_get_tl_rsc(params, lane);
+        if ((tl_rsc->dev_type == UCT_DEVICE_TYPE_NET) &&
+            (tl_rsc->sys_device != sys_devs[seed])) {
+            ucp_proto_perf_node_deref(&perf_nodes[lane]);
+            ucs_trace("filtered out " UCP_PROTO_LANE_FMT,
+                      UCP_PROTO_LANE_ARG(params, lane, &tl_perfs[lane]));
+        } else {
+            lanes[num_filtered_lanes++] = lane;
+        }
+    }
+
+    return num_filtered_lanes;
+}
+
 ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
                                   const char *perf_name,
                                   ucp_proto_perf_t **perf_p,
@@ -252,6 +308,14 @@ ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
     }
 
     num_lanes = num_fast_lanes;
+    if (context->config.ext.proto_use_single_net_device) {
+        num_lanes = ucp_proto_multi_filter_net_devices(num_lanes,
+                                                       &params->super.super,
+                                                       lanes_perf,
+                                                       fixed_first_lane, lanes,
+                                                       lanes_perf_nodes);
+    }
+
     ucp_proto_multi_select_bw_lanes(&params->super.super, lanes, num_lanes,
                                     params->max_lanes, lanes_perf,
                                     fixed_first_lane, &selection);
diff --git a/src/ucp/proto/proto_single.c b/src/ucp/proto/proto_single.c
@@ -10,13 +10,87 @@
 
 #include "proto_single.h"
 #include "proto_common.h"
+#include "proto_common.inl"
 #include "proto_init.h"
 #include "proto_debug.h"
 
 #include <ucs/debug/assert.h>
 #include <ucs/debug/log.h>
 #include <ucs/sys/math.h>
 
+static double
+ucp_proto_single_get_bandwidth(const ucp_proto_common_init_params_t *params,
+                               ucp_lane_index_t lane)
+{
+    ucp_proto_common_tl_perf_t tl_perf;
+    ucp_proto_perf_node_t *perf_node;
+    ucs_status_t status;
+
+    status = ucp_proto_common_get_lane_perf(params, lane, &tl_perf, &perf_node);
+    if (status != UCS_OK) {
+        return 0;
+    }
+
+    ucp_proto_perf_node_deref(&perf_node);
+    return tl_perf.bandwidth;
+}
+
+static void
+ucp_proto_single_update_lane(const ucp_proto_single_init_params_t *params,
+                             ucp_lane_index_t *lane_p)
+{
+    const ucp_proto_common_init_params_t *common_params = &params->super;
+    const ucp_proto_init_params_t *init_params = &common_params->super;
+    const ucp_context_h context                = init_params->worker->context;
+    double bandwidth;
+    ucp_lane_index_t lanes[UCP_PROTO_MAX_LANES];
+    ucs_sys_device_t sys_devs[UCP_PROTO_MAX_LANES];
+    ucp_lane_index_t num_lanes, num_same_bw_devs, i, lane;
+    ucs_sys_device_t sys_dev;
+
+    if (!context->config.ext.proto_use_single_net_device ||
+        /* skip lane update for node_local_id 0 since the original lane would be
+         * selected anyway */
+        (context->config.node_local_id == 0)) {
+        return;
+    }
+
+    if (!ucp_proto_common_is_net_dev(init_params, *lane_p)) {
+        return;
+    }
+
+    bandwidth   = ucp_proto_single_get_bandwidth(common_params, *lane_p);
+    lanes[0]    = *lane_p;
+    sys_devs[0] = ucp_proto_common_get_sys_dev(init_params, lanes[0]);
+
+    num_lanes = ucp_proto_common_find_lanes(
+            init_params, common_params->flags, params->lane_type,
+            params->tl_cap_flags, UCP_PROTO_MAX_LANES - 1,
+            (common_params->exclude_map | UCS_BIT(lanes[0])),
+            ucp_proto_common_filter_min_frag, lanes + 1);
+
+    for (num_same_bw_devs = 1, i = 1; i < num_lanes; ++i) {
+        lane = lanes[i];
+        if (!ucp_proto_common_is_net_dev(init_params, lane)) {
+            continue;
+        }
+
+        if (!ucp_proto_common_bandwidth_equal(
+                ucp_proto_single_get_bandwidth(common_params, lane), bandwidth)) {
+            continue;
+        }
+
+        sys_dev = ucp_proto_common_get_sys_dev(init_params, lane);
+        if (ucp_proto_common_add_unique_sys_dev(sys_dev, sys_devs,
+                                                &num_same_bw_devs,
+                                                UCP_PROTO_MAX_LANES)) {
+            lanes[num_same_bw_devs - 1] = lane;
+        }
+    }
+
+    *lane_p = lanes[ucp_proto_common_select_sys_dev_by_node_id(init_params,
+                                                               num_same_bw_devs)];
+}
 
 ucs_status_t ucp_proto_single_init(const ucp_proto_single_init_params_t *params,
                                    ucp_proto_perf_t **perf_p,
@@ -47,6 +121,8 @@ ucs_status_t ucp_proto_single_init(const ucp_proto_single_init_params_t *params,
 
     ucs_assert(num_lanes == 1);
 
+    ucp_proto_single_update_lane(params, &lane);
+
     reg_md_map = ucp_proto_common_reg_md_map(&params->super, UCS_BIT(lane));
     if (reg_md_map == 0) {
         spriv->reg_md = UCP_NULL_RESOURCE;
diff --git a/test/gtest/ucp/test_ucp_proto_mock.cc b/test/gtest/ucp/test_ucp_proto_mock.cc
