UCP: Split an SGL put between the lanes by bytes - v1.23.x

src/ucp/core/ucp_request.h

@@ -248,6 +248,7 @@ struct ucp_request {
                     struct {
                         const uint64_t   *remote_addrs;
                         ucp_rkey_h const *rkeys;
+                        size_t           length; /* Total length in bytes */
                     } sgl;
                 } rma;
 

src/ucp/dt/datatype_iter.inl

@@ -640,27 +640,30 @@ static UCS_F_ALWAYS_INLINE size_t
 ucp_datatype_iter_next_sgl_frags(const ucp_datatype_iter_t *dt_iter,
                                  const uint64_t *remote_addrs,
                                  size_t max_frag_count,
-                                 size_t max_frag_length,
+                                 size_t max_total_length,
                                  ucp_datatype_iter_t *next_iter,
                                  void **out_buffers, size_t *out_lengths,
                                  uint64_t *out_remote_addrs,
                                  size_t *out_elem_indices)
 {
-    size_t elem_index  = dt_iter->offset;
-    size_t frag_offset = dt_iter->type.sgl.frag_offset;
-    size_t desc_count  = 0;
+    size_t elem_index   = dt_iter->offset;
+    size_t frag_offset  = dt_iter->type.sgl.frag_offset;
+    size_t desc_count   = 0;
+    size_t total_length = 0;
     size_t elem_length, frag_length;
 
     ucs_assert(dt_iter->dt_class == UCP_DATATYPE_SGL);
     ucs_assert(max_frag_count >= 1);
-    ucs_assert(max_frag_length > 0);
+    ucs_assert(max_total_length > 0);
     ucp_datatype_iter_sgl_check(dt_iter);
 
-    while ((desc_count < max_frag_count) && (elem_index < dt_iter->length)) {
+    while ((desc_count < max_frag_count) && (total_length < max_total_length) &&
+           (elem_index < dt_iter->length)) {
         elem_length = dt_iter->type.sgl.lengths[elem_index];
         if (elem_length > 0) {
             ucs_assert(elem_length > frag_offset);
-            frag_length = ucs_min(elem_length - frag_offset, max_frag_length);
+            frag_length = ucs_min(elem_length - frag_offset,
+                                  max_total_length - total_length);
 
             out_buffers[desc_count]      = UCS_PTR_BYTE_OFFSET(
                     dt_iter->type.sgl.buffers[elem_index], frag_offset);
@@ -670,7 +673,8 @@ ucp_datatype_iter_next_sgl_frags(const ucp_datatype_iter_t *dt_iter,
             out_elem_indices[desc_count] = elem_index;
             ++desc_count;
 
-            frag_offset += frag_length;
+            frag_offset  += frag_length;
+            total_length += frag_length;
             if (frag_offset < elem_length) {
                 continue;
             }

src/ucp/dt/dt_sgl.h

@@ -16,6 +16,26 @@
     (((_datatype) & UCP_DATATYPE_CLASS_MASK) == UCP_DATATYPE_SGL)
 
 
+/**
+ * Get the total length of all SGL elements
+ *
+ * @param [in]     lengths        Array of element lengths
+ * @param [in]     count          Number of entries in the @a lengths array
+ *
+ * @return Total length of all the SGL elements
+ */
+static inline size_t ucp_dt_sgl_length(const size_t *lengths, size_t count)
+{
+    size_t elem_it, total_length = 0;
+
+    for (elem_it = 0; elem_it < count; ++elem_it) {
+        total_length += lengths[elem_it];
+    }
+
+    return total_length;
+}
+
+
 /**
  * Check that all SGL entries match the given memory info
  *

src/ucp/proto/proto_multi.inl

@@ -58,14 +58,26 @@ ucp_proto_multi_scaled_length(uint32_t weight, size_t length)
            UCP_PROTO_MULTI_WEIGHT_SHIFT;
 }
 
+static UCS_F_ALWAYS_INLINE size_t
+ucp_proto_multi_lane_payload(const ucp_proto_multi_lane_priv_t *lpriv,
+                             size_t length, size_t remaining, size_t max_frag)
+{
+    /* 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 */
+    if (remaining < lpriv->min_end_offset) {
+        return max_frag;
+    }
+
+    return ucs_min(ucp_proto_multi_scaled_length(lpriv->weight, length),
+                   max_frag);
+}
+
 static UCS_F_ALWAYS_INLINE size_t
 ucp_proto_multi_max_payload(ucp_request_t *req,
                             const ucp_proto_multi_lane_priv_t *lpriv,
                             size_t hdr_size)
 {
     size_t length = req->send.state.dt_iter.length;
-    size_t offset = length - req->send.state.dt_iter.offset;
-    size_t max_frag;
     size_t max_payload;
 
     /* If header is greater than max_frag - that's ok, we allow it, but we send
@@ -74,16 +86,9 @@ ucp_proto_multi_max_payload(ucp_request_t *req,
         return 0;
     }
 
-    max_frag = lpriv->max_frag - hdr_size;
-
-    /* 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 */
-    if (offset < lpriv->min_end_offset) {
-        return max_frag;
-    }
-
-    max_payload = ucs_min(ucp_proto_multi_scaled_length(lpriv->weight, length),
-                          max_frag);
+    max_payload = ucp_proto_multi_lane_payload(
+            lpriv, length, length - req->send.state.dt_iter.offset,
+            lpriv->max_frag - hdr_size);
     ucs_assertv(max_payload > 0,
                 "length=%zu weight=%zu%% lpriv->max_frag=%zu hdr_size=%zu",
                 length, ucp_proto_multi_scaled_length(lpriv->weight, 100),

src/ucp/rma/put_offload.c

@@ -389,54 +389,59 @@ ucp_proto_put_sgl_offload_post(ucp_request_t *req,
     return status;
 }
 
-static UCS_F_ALWAYS_INLINE int
-ucp_proto_put_sgl_elem_fits(size_t length, size_t max_frag_length)
-{
-    return (length > 0) && (length <= max_frag_length);
-}
-
+/* Split the SGL by its total length, like a multi-lane contiguous message */
 static UCS_F_ALWAYS_INLINE size_t
-ucp_proto_put_sgl_max_elem_count(const ucp_proto_multi_lane_priv_t *lpriv,
-                                 size_t count)
+ucp_proto_put_sgl_max_payload(const ucp_request_t *req,
+                              const ucp_proto_multi_lane_priv_t *lpriv)
 {
-    return ucs_min(lpriv->max_sgl_zcopy_count,
-                   ucp_proto_multi_scaled_length(lpriv->weight, count));
+    size_t length = req->send.rma.sgl.length;
+
+    return ucp_proto_multi_lane_payload(lpriv, length, length,
+                                        lpriv->max_frag);
 }
 
-static UCS_F_ALWAYS_INLINE ucs_status_t
-ucp_proto_put_sgl_offload_send_frag(ucp_request_t *req,
-                                    const ucp_proto_multi_lane_priv_t *lpriv,
-                                    size_t max_frag_length,
-                                    ucp_datatype_iter_t *next_iter)
+static UCS_F_INLINE_OPTIMIZED ucs_status_t
+ucp_proto_put_sgl_offload_send_frags(ucp_request_t *req,
+                                     const ucp_proto_multi_lane_priv_t *lpriv,
+                                     size_t max_frag_count,
+                                     size_t max_payload,
+                                     ucp_datatype_iter_t *next_iter)
 {
     ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter;
     ucp_rsc_index_t md_index     = lpriv->super.md_index;
     ucp_mem_h *sgl_memhs         = dt_iter->type.sgl.memhs;
-    void *buffer                 = NULL;
-    size_t length                = 0;
-    size_t elem_index            = 0;
-    uint64_t remote_addr         = 0;
-    uct_mem_h uct_memh;
-    uct_rkey_t uct_rkey;
-
-    if (ucp_datatype_iter_next_sgl_frags(dt_iter,
-                                         req->send.rma.sgl.remote_addrs, 1,
-                                         max_frag_length, next_iter, &buffer,
-                                         &length, &remote_addr,
-                                         &elem_index) == 0) {
+    ucp_rkey_h const *sgl_rkeys  = req->send.rma.sgl.rkeys;
+    void *buffers[UCP_MAX_IOV];
+    size_t lengths[UCP_MAX_IOV];
+    uint64_t remote_addrs[UCP_MAX_IOV];
+    size_t elem_indices[UCP_MAX_IOV];
+    uct_mem_h uct_memhs[UCP_MAX_IOV];
+    uct_rkey_t uct_rkeys[UCP_MAX_IOV];
+    size_t frag_count, frag_index, elem_index;
+
+    ucs_assert(max_frag_count <= UCP_MAX_IOV);
+
+    frag_count = ucp_datatype_iter_next_sgl_frags(
+            dt_iter, req->send.rma.sgl.remote_addrs, max_frag_count,
+            max_payload, next_iter, buffers, lengths, remote_addrs,
+            elem_indices);
+    if (frag_count == 0) {
         return UCS_OK;
     }
 
-    uct_memh = (sgl_memhs != NULL) ?
-                       ucp_datatype_iter_uct_memh(sgl_memhs[elem_index],
-                                                  md_index) :
-                       UCT_MEM_HANDLE_NULL;
-    uct_rkey = ucp_rkey_get_tl_rkey(req->send.rma.sgl.rkeys[elem_index],
-                                    lpriv->super.rkey_index);
+    for (frag_index = 0; frag_index < frag_count; frag_index++) {
+        elem_index            = elem_indices[frag_index];
+        uct_memhs[frag_index] = (sgl_memhs != NULL) ?
+                                ucp_datatype_iter_uct_memh(sgl_memhs[elem_index],
+                                                           md_index) :
+                                UCT_MEM_HANDLE_NULL;
+        uct_rkeys[frag_index] = ucp_rkey_get_tl_rkey(sgl_rkeys[elem_index],
+                                                     lpriv->super.rkey_index);
+    }
 
-    return ucp_proto_put_sgl_offload_post(req, lpriv, &buffer, &length,
-                                          &uct_memh, &remote_addr, &uct_rkey,
-                                          1);
+    return ucp_proto_put_sgl_offload_post(req, lpriv, buffers, lengths,
+                                          uct_memhs, remote_addrs, uct_rkeys,
+                                          frag_count);
 }
 
 static UCS_F_ALWAYS_INLINE ucs_status_t
@@ -448,37 +453,63 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
     ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter;
     ucp_rsc_index_t md_index     = lpriv->super.md_index;
     ucp_rsc_index_t rkey_index   = lpriv->super.rkey_index;
-    size_t max_frag_length       = lpriv->max_frag;
+    size_t max_payload           = ucp_proto_put_sgl_max_payload(req, lpriv);
+    size_t max_frag_count        = ucs_min(lpriv->max_sgl_zcopy_count,
+                                           UCP_MAX_IOV);
     ucp_mem_h *sgl_memhs         = dt_iter->type.sgl.memhs;
     ucp_rkey_h const *sgl_rkeys  = req->send.rma.sgl.rkeys;
     void *const *buffers         = dt_iter->type.sgl.buffers;
     const size_t *lengths        = dt_iter->type.sgl.lengths;
     const uint64_t *remote_addrs = req->send.rma.sgl.remote_addrs;
     size_t start_index           = dt_iter->offset;
-    size_t max_elem_count        = ucs_min(
-            ucp_proto_put_sgl_max_elem_count(lpriv, dt_iter->length),
-            dt_iter->length - start_index);
-    size_t uct_rkeys_size        = max_elem_count * sizeof(uct_rkey_t);
-    size_t uct_memhs_size        = max_elem_count * sizeof(uct_mem_h);
+    size_t max_elem_count        = ucs_min(lpriv->max_sgl_zcopy_count,
+                                           dt_iter->length - start_index);
+    size_t uct_rkeys_size;
+    size_t uct_memhs_size;
     uct_rkey_t *uct_rkeys;
     uct_mem_h *uct_memhs;
     ucs_status_t status;
-    size_t elem_count, idx;
+    size_t elem_count, idx, length, total_length;
 
-    ucs_assert(max_frag_length > 0);
+    ucs_assert(max_payload > 0);
     ucs_assert(max_elem_count > 0);
 
     /* Silence compiler warning, in case of an early return below */
     next_iter->offset               = start_index;
     next_iter->type.sgl.frag_offset = dt_iter->type.sgl.frag_offset;
 
+    /* The first element must be posted as a whole, otherwise it is fragmented */
     if (ucs_unlikely((dt_iter->type.sgl.frag_offset != 0) ||
-                     !ucp_proto_put_sgl_elem_fits(lengths[start_index],
-                                                  max_frag_length))) {
-        return ucp_proto_put_sgl_offload_send_frag(req, lpriv, max_frag_length,
-                                                   next_iter);
+                     (lengths[start_index] == 0) ||
+                     (lengths[start_index] > max_payload))) {
+        return ucp_proto_put_sgl_offload_send_frags(req, lpriv, max_frag_count,
+                                                    max_payload, next_iter);
     }
 
+    /* Stop at the lane payload, so that each lane sends its share of the
+     * message */
+    for (elem_count = 0, total_length = 0; elem_count < max_elem_count;
+         elem_count++) {
+        length = lengths[start_index + elem_count];
+        if ((length == 0) || ((total_length + length) > max_payload)) {
+            break;
+        }
+
+        total_length += length;
+    }
+
+    /* A long batch is not redirected, to keep posting the user arrays without
+     * copying them */
+    if (ucs_unlikely((total_length < max_payload) &&
+                     (elem_count < max_frag_count) &&
+                     ((start_index + elem_count) < dt_iter->length))) {
+        return ucp_proto_put_sgl_offload_send_frags(req, lpriv, max_frag_count,
+                                                    max_payload, next_iter);
+    }
+
+    uct_rkeys_size = elem_count * sizeof(uct_rkey_t);
+    uct_memhs_size = elem_count * sizeof(uct_mem_h);
+
     uct_rkeys = ucs_alloc_on_stack(uct_rkeys_size, "uct_sgl_rkeys");
     if (uct_rkeys == NULL) {
         return UCS_ERR_NO_MEMORY;
@@ -490,18 +521,13 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
         return UCS_ERR_NO_MEMORY;
     }
 
-    for (elem_count = 0; elem_count < max_elem_count; elem_count++) {
-        idx = start_index + elem_count;
-        if (!ucp_proto_put_sgl_elem_fits(lengths[idx], max_frag_length)) {
-            break;
-        }
-
-        uct_memhs[elem_count] = (sgl_memhs != NULL) ?
-                                ucp_datatype_iter_uct_memh(sgl_memhs[idx],
-                                                           md_index) :
-                                UCT_MEM_HANDLE_NULL;
-        uct_rkeys[elem_count] = ucp_rkey_get_tl_rkey(sgl_rkeys[idx],
-                                                     rkey_index);
+    for (idx = 0; idx < elem_count; idx++) {
+        uct_memhs[idx] = (sgl_memhs != NULL) ?
+                         ucp_datatype_iter_uct_memh(sgl_memhs[start_index + idx],
+                                                    md_index) :
+                         UCT_MEM_HANDLE_NULL;
+        uct_rkeys[idx] = ucp_rkey_get_tl_rkey(sgl_rkeys[start_index + idx],
+                                              rkey_index);
     }
 
     next_iter->offset               = start_index + elem_count;
@@ -512,7 +538,9 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
                                             &remote_addrs[start_index],
                                             uct_rkeys, elem_count);
 
+    /* coverity[incorrect_free] */
     ucs_free_on_stack(uct_memhs, uct_memhs_size);
+    /* coverity[incorrect_free] */
     ucs_free_on_stack(uct_rkeys, uct_rkeys_size);
 
     return status;
@@ -562,7 +590,7 @@ ucp_proto_put_sgl_offload_sw_send_func(ucp_request_t *req,
     uct_ep_h uct_ep              = ucp_ep_get_lane(ep, lane);
     ucp_rsc_index_t md_index     = lpriv->super.md_index;
     ucp_rsc_index_t rkey_index   = lpriv->super.rkey_index;
-    size_t max_frag_length       = lpriv->max_frag;
+    size_t max_frag_length       = ucp_proto_put_sgl_max_payload(req, lpriv);
     ucp_mem_h *sgl_memhs         = dt_iter->type.sgl.memhs;
     ucp_rkey_h const *sgl_rkeys  = req->send.rma.sgl.rkeys;
     void *buffer                 = NULL;

src/ucp/rma/rma_send.c

@@ -195,6 +195,7 @@ ucs_status_ptr_t ucp_put_nbx(ucp_ep_h ep, const void *buffer, size_t count,
 {
     ucp_worker_h worker               = ep->worker;
     size_t contig_length              = 0;
+    size_t sgl_length                 = 0;
     ucp_datatype_t datatype           = ucp_dt_make_contig(1);
     const ucp_dt_remote_sgl_t *remote = NULL;
     ucs_status_ptr_t ret;
@@ -227,8 +228,10 @@ ucs_status_ptr_t ucp_put_nbx(ucp_ep_h ep, const void *buffer, size_t count,
         datatype = param->datatype;
         if (UCP_DT_IS_SGL(datatype)) {
             UCP_PUT_SGL_CHECK_PARAMS(buffer, count, remote_addr, rkey, param);
-            remote = param->remote;
-            rkey   = remote->rkeys[0];
+            remote     = param->remote;
+            rkey       = remote->rkeys[0];
+            sgl_length = ucp_dt_sgl_length(
+                    ((const ucp_dt_local_sgl_t*)buffer)->lengths, count);
         } else if (UCP_DT_IS_CONTIG(datatype)) {
             contig_length = ucp_contig_dt_length(datatype, count);
         }
@@ -244,6 +247,7 @@ ucs_status_ptr_t ucp_put_nbx(ucp_ep_h ep, const void *buffer, size_t count,
     if (remote != NULL) {
         req->send.rma.sgl.remote_addrs = remote->remote_addrs;
         req->send.rma.sgl.rkeys        = remote->rkeys;
+        req->send.rma.sgl.length       = sgl_length;
     }
 
     ret = ucp_proto_request_send_op_rma(