From ce3f490c00051e6587495c5aed679b8264b127ea Mon Sep 17 00:00:00 2001
From: Ilia Yastrebov <iyastrebov@nvidia.com>
Date: Mon, 28 Oct 2024 15:08:30 +0000
Subject: [PATCH] UCP/PROTO/MULTI: Allow header be greater than max_frag, sort
 lanes by BW

---
 src/ucp/proto/proto_multi.c   | 27 +++++++++++++++++++++++++--
 src/ucp/proto/proto_multi.inl |  7 +++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/ucp/proto/proto_multi.c b/src/ucp/proto/proto_multi.c
index 3aa0f4937b5..ce5cd0cbd93 100644
--- a/src/ucp/proto/proto_multi.c
+++ b/src/ucp/proto/proto_multi.c
@@ -16,10 +16,24 @@
 #include "proto_debug.h"
 #include "proto_multi.inl"
 
+#include <ucs/algorithm/qsort_r.h>
 #include <ucs/debug/assert.h>
 #include <ucs/debug/log.h>
 
 
+static int ucp_proto_multi_compare_tl_perf(const void *elem1, const void *elem2,
+                                           void *arg)
+{
+    const ucp_lane_index_t *lane1                = elem1;
+    const ucp_lane_index_t *lane2                = elem2;
+    const ucp_proto_common_tl_perf_t *lanes_perf = arg;
+    int diff;
+
+    diff = lanes_perf[*lane2].bandwidth - lanes_perf[*lane1].bandwidth;
+    /* If bandwidths are equal, prefer to maintain the original ordering */
+    return (diff != 0) ? diff : (elem1 - elem2);
+}
+
 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,
@@ -59,8 +73,7 @@ ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
             &params->super, params->first.lane_type, params->first.tl_cap_flags,
             1, 0, lanes);
     if (num_lanes == 0) {
-        ucs_trace("no lanes for %s",
-                  ucp_proto_id_field(params->super.super.proto_id, name));
+        ucs_trace("no lanes for %s", perf_name);
         return UCS_ERR_NO_ELEM;
     }
 
@@ -86,6 +99,16 @@ ucs_status_t ucp_proto_multi_init(const ucp_proto_multi_init_params_t *params,
         max_bandwidth = ucs_max(max_bandwidth, lane_perf->bandwidth);
     }
 
+    /* Sort lanes by bandwidth */
+    if (num_lanes > 1) {
+        /* If type of the first lane is not the same as in the middle, then
+         * exclude this lane from sorting */
+        lane = (params->first.lane_type == params->middle.lane_type) ? 0 : 1;
+        /* Sort lanes array by lanes_perf bandwidth */
+        ucs_qsort_r(lanes + lane, num_lanes - lane, sizeof(ucp_lane_index_t),
+                    ucp_proto_multi_compare_tl_perf, lanes_perf);
+    }
+
     /* Select the lanes to use, and calculate their aggregate performance */
     perf.bandwidth          = 0;
     perf.send_pre_overhead  = 0;
diff --git a/src/ucp/proto/proto_multi.inl b/src/ucp/proto/proto_multi.inl
index 6ae2371d1e4..3e65fd84aaf 100644
--- a/src/ucp/proto/proto_multi.inl
+++ b/src/ucp/proto/proto_multi.inl
@@ -66,8 +66,11 @@ ucp_proto_multi_max_payload(ucp_request_t *req,
     size_t max_frag = lpriv->max_frag - hdr_size;
     size_t max_payload;
 
-    ucs_assertv(lpriv->max_frag > hdr_size, "max_frag=%zu hdr_size=%zu",
-                lpriv->max_frag, hdr_size);
+    /* If header is greater than max_frag - that's ok, we allow it, but we send
+     * only the header data and keep the payload empty */
+    if (lpriv->max_frag <= hdr_size) {
+        return 0;
+    }
 
     /* Do not split very small sends to chunks, it's not worth it, and
        generic datatype may not be able to pack to a smaller buffer */