UCP: Split an SGL put between the lanes by bytes

src/ucp/dt/datatype_iter.c

@@ -296,10 +296,12 @@ ucs_status_t ucp_datatype_iter_sgl_init(ucp_context_h context,
     ucs_assert(remote != NULL);
 
     dt_iter->dt_class             = UCP_DATATYPE_SGL;
-    dt_iter->length               = count;
+    dt_iter->length               = ucp_dt_sgl_length(local->lengths, count);
     dt_iter->offset               = 0;
     dt_iter->type.sgl.buffers     = local->buffers;
     dt_iter->type.sgl.lengths     = local->lengths;
+    dt_iter->type.sgl.elem_count  = count;
+    dt_iter->type.sgl.elem_index  = 0;
     dt_iter->type.sgl.frag_offset = 0;
 
     if (ucs_unlikely(count == 0)) {
@@ -352,7 +354,7 @@ ucs_status_t ucp_datatype_iter_sgl_mem_reg(ucp_context_h context,
                                            ucp_md_map_t md_map,
                                            unsigned uct_flags)
 {
-    size_t count = dt_iter->length;
+    size_t count = dt_iter->type.sgl.elem_count;
     ucs_status_t status;
     ucp_mem_h *memhs;
     size_t i;
@@ -393,7 +395,7 @@ ucs_status_t ucp_datatype_iter_sgl_mem_reg(ucp_context_h context,
 
 void ucp_datatype_iter_sgl_mem_dereg(ucp_datatype_iter_t *dt_iter)
 {
-    size_t count = dt_iter->length;
+    size_t count = dt_iter->type.sgl.elem_count;
     size_t i;
 
     ucs_assert(ucp_datatype_iter_sgl_owns_memhs(dt_iter));
@@ -419,7 +421,7 @@ void ucp_datatype_iter_sgl_cleanup(ucp_datatype_iter_t *dt_iter, int dereg)
     }
 
     if (UCS_ENABLE_ASSERT) {
-        for (i = 0; i < dt_iter->length; ++i) {
+        for (i = 0; i < dt_iter->type.sgl.elem_count; ++i) {
             ucp_datatatype_iter_memh_cleanup_check(dt_iter->type.sgl.memhs[i]);
         }
     }
@@ -517,7 +519,7 @@ int ucp_datatype_iter_is_user_memh_valid(const ucp_datatype_iter_t *dt_iter,
         break;
     case UCP_DATATYPE_SGL:
         ref = ucp_memory_info_from_memh(memh);
-        for (i = 0; i < dt_iter->length; ++i) {
+        for (i = 0; i < dt_iter->type.sgl.elem_count; ++i) {
             sgl_memh = dt_iter->type.sgl.memhs[i];
             if (sgl_memh == NULL) {
                 ucs_error("sgl[%zu]: got NULL memory handle", i);
@@ -537,7 +539,8 @@ int ucp_datatype_iter_is_user_memh_valid(const ucp_datatype_iter_t *dt_iter,
 
             cur = ucp_memory_info_from_memh(sgl_memh);
             if (ucp_dt_mem_info_verify("sgl", i, &cur, &ref,
-                                       dt_iter->length) != UCS_OK) {
+                                       dt_iter->type.sgl.elem_count) !=
+                UCS_OK) {
                 return 0;
             }
         }

src/ucp/dt/datatype_iter.h

@@ -77,8 +77,9 @@ typedef struct {
             void * const     *buffers;
             const size_t     *lengths;
             ucp_mem_h        *memhs;
+            size_t           elem_count;  /* Number of elements */
+            size_t           elem_index;  /* Index of current element */
             size_t           frag_offset; /* Offset in the current element */
-            /* length = element count, offset = current element index */
         } sgl;
     } type;
 } ucp_datatype_iter_t;

src/ucp/dt/datatype_iter.inl

@@ -384,14 +384,17 @@ ucp_datatype_iter_sgl_check(const ucp_datatype_iter_t *dt_iter)
 
     ucs_assertv(dt_iter->offset <= dt_iter->length,
                 "offset=%zu length=%zu", dt_iter->offset, dt_iter->length);
-    if (dt_iter->offset < dt_iter->length) {
-        elem_length = dt_iter->type.sgl.lengths[dt_iter->offset];
+    ucs_assertv(dt_iter->type.sgl.elem_index <= dt_iter->type.sgl.elem_count,
+                "elem_index=%zu elem_count=%zu", dt_iter->type.sgl.elem_index,
+                dt_iter->type.sgl.elem_count);
+    if (dt_iter->type.sgl.elem_index < dt_iter->type.sgl.elem_count) {
+        elem_length = dt_iter->type.sgl.lengths[dt_iter->type.sgl.elem_index];
         ucs_assertv((elem_length > 0) ?
                     (dt_iter->type.sgl.frag_offset < elem_length) :
                     (dt_iter->type.sgl.frag_offset == 0),
-                    "offset=%zu frag_offset=%zu elem_length=%zu",
-                    dt_iter->offset, dt_iter->type.sgl.frag_offset,
-                    elem_length);
+                    "elem_index=%zu frag_offset=%zu elem_length=%zu",
+                    dt_iter->type.sgl.elem_index,
+                    dt_iter->type.sgl.frag_offset, elem_length);
     }
 }
 
@@ -646,17 +649,19 @@ ucp_datatype_iter_next_sgl_frags(const ucp_datatype_iter_t *dt_iter,
                                  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->type.sgl.elem_index;
+    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);
     ucp_datatype_iter_sgl_check(dt_iter);
 
-    while ((desc_count < max_frag_count) && (elem_index < dt_iter->length)) {
+    while ((desc_count < max_frag_count) &&
+           (elem_index < dt_iter->type.sgl.elem_count)) {
         elem_length = dt_iter->type.sgl.lengths[elem_index];
         if (elem_length > 0) {
             ucs_assert(elem_length > frag_offset);
@@ -670,7 +675,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;
             }
@@ -680,14 +686,18 @@ ucp_datatype_iter_next_sgl_frags(const ucp_datatype_iter_t *dt_iter,
         frag_offset = 0;
     }
 
-    ucs_assertv((elem_index == dt_iter->length) || (desc_count > 0),
-                "dt_iter=%p offset=%zu length=%zu frag_offset=%zu "
+    ucs_assertv((elem_index == dt_iter->type.sgl.elem_count) ||
+                (desc_count > 0),
+                "dt_iter=%p elem_index=%zu elem_count=%zu frag_offset=%zu "
                 "desc_count=%zu",
-                dt_iter, dt_iter->offset, dt_iter->length,
-                dt_iter->type.sgl.frag_offset, desc_count);
+                dt_iter, dt_iter->type.sgl.elem_index,
+                dt_iter->type.sgl.elem_count, dt_iter->type.sgl.frag_offset,
+                desc_count);
 
-    next_iter->offset               = elem_index;
+    next_iter->offset               = dt_iter->offset + total_length;
+    next_iter->type.sgl.elem_index  = elem_index;
     next_iter->type.sgl.frag_offset = frag_offset;
+    ucs_assert(next_iter->offset <= dt_iter->length);
 
     return desc_count;
 }
@@ -709,6 +719,7 @@ ucp_datatype_iter_copy_position(ucp_datatype_iter_t *dt_iter,
         dt_iter->type.iov.iov_index  = src_dt_iter->type.iov.iov_index;
         dt_iter->type.iov.iov_offset = src_dt_iter->type.iov.iov_offset;
     } else if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_SGL, dt_mask)) {
+        dt_iter->type.sgl.elem_index  = src_dt_iter->type.sgl.elem_index;
         dt_iter->type.sgl.frag_offset = src_dt_iter->type.sgl.frag_offset;
     }
 }
@@ -734,10 +745,33 @@ ucp_datatype_iter_rewind(ucp_datatype_iter_t *dt_iter, unsigned dt_mask)
         dt_iter->type.iov.iov_index  = 0;
         dt_iter->type.iov.iov_offset = 0;
     } else if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_SGL, dt_mask)) {
+        dt_iter->type.sgl.elem_index  = 0;
         dt_iter->type.sgl.frag_offset = 0;
     }
 }
 
+static UCS_F_ALWAYS_INLINE void
+ucp_datatype_iter_sgl_seek(ucp_datatype_iter_t *dt_iter, size_t offset)
+{
+    size_t elem_index  = 0;
+    size_t elem_offset = offset;
+    size_t elem_length;
+
+    while (elem_index < dt_iter->type.sgl.elem_count) {
+        elem_length = dt_iter->type.sgl.lengths[elem_index];
+        if (elem_offset < elem_length) {
+            break;
+        }
+
+        elem_offset -= elem_length;
+        ++elem_index;
+    }
+
+    dt_iter->offset               = offset;
+    dt_iter->type.sgl.elem_index  = elem_index;
+    dt_iter->type.sgl.frag_offset = elem_offset;
+}
+
 static UCS_F_ALWAYS_INLINE void
 ucp_datatype_iter_seek(ucp_datatype_iter_t *dt_iter, size_t offset,
                        unsigned dt_mask)
@@ -748,8 +782,7 @@ ucp_datatype_iter_seek(ucp_datatype_iter_t *dt_iter, size_t offset,
     if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_IOV, dt_mask)) {
         ucp_datatype_iter_iov_seek(dt_iter, offset);
     } else if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_SGL, dt_mask)) {
-        dt_iter->type.sgl.frag_offset = 0;
-        dt_iter->offset               = offset;
+        ucp_datatype_iter_sgl_seek(dt_iter, offset);
     } else {
         dt_iter->offset = offset;
     }

src/ucp/dt/dt_sgl.h

@@ -16,6 +16,18 @@
     (((_datatype) & UCP_DATATYPE_CLASS_MASK) == UCP_DATATYPE_SGL)
 
 
+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/rma/put_offload.c

@@ -447,16 +447,17 @@ 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_frag_length       = ucp_proto_multi_max_payload(req, lpriv, 0);
     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 elem_count_total      = dt_iter->type.sgl.elem_count;
+    size_t start_index           = dt_iter->type.sgl.elem_index;
     size_t max_elem_count        = ucs_min(
-            ucp_proto_put_sgl_max_elem_count(lpriv, dt_iter->length),
-            dt_iter->length - start_index);
+            ucp_proto_put_sgl_max_elem_count(lpriv, elem_count_total),
+            elem_count_total - 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);
     uct_rkey_t *uct_rkeys;
@@ -468,7 +469,8 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
     ucs_assert(max_elem_count > 0);
 
     /* Silence compiler warning, in case of an early return below */
-    next_iter->offset               = start_index;
+    next_iter->offset               = dt_iter->offset;
+    next_iter->type.sgl.elem_index  = start_index;
     next_iter->type.sgl.frag_offset = dt_iter->type.sgl.frag_offset;
 
     if (ucs_unlikely((dt_iter->type.sgl.frag_offset != 0) ||
@@ -495,6 +497,7 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
             break;
         }
 
+        next_iter->offset    += lengths[idx];
         uct_memhs[elem_count] = (sgl_memhs != NULL) ?
                                 ucp_datatype_iter_uct_memh(sgl_memhs[idx],
                                                            md_index) :
@@ -503,8 +506,9 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
                                                      rkey_index);
     }
 
-    next_iter->offset               = start_index + elem_count;
+    next_iter->type.sgl.elem_index  = start_index + elem_count;
     next_iter->type.sgl.frag_offset = 0;
+    ucs_assert(next_iter->offset <= dt_iter->length);
 
     status = ucp_proto_put_sgl_offload_post(req, lpriv, &buffers[start_index],
                                             &lengths[start_index], uct_memhs,
@@ -561,7 +565,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_multi_max_payload(req, lpriv, 0);
     ucp_mem_h *sgl_memhs         = dt_iter->type.sgl.memhs;
     ucp_rkey_h const *sgl_rkeys  = req->send.rma.sgl.rkeys;
     void *buffer                 = NULL;

test/gtest/common/test_obj_size.cc

@@ -57,7 +57,7 @@ UCS_TEST_F(test_obj_size, size) {
     EXPECTED_SIZE(ucp_rkey_t, 24);
 #endif
     /* TODO reduce request size to 240 or less after removing old protocols state */
-    EXPECTED_SIZE(ucp_request_t, 272);
+    EXPECTED_SIZE(ucp_request_t, 288);
     EXPECTED_SIZE(ucp_recv_desc_t, 48);
     EXPECTED_SIZE(ucp_mem_t, 160);
     EXPECTED_SIZE(uct_ep_t, 8);

test/gtest/ucp/test_ucp_dt.cc

@@ -382,16 +382,30 @@ class test_ucp_dt_sgl : public ucs::test {
                           m_desc_elem_indices.data());
     }
 
-    void check_next_iter(size_t next_offset, size_t next_frag_offset)
+    size_t flat_offset(size_t elem_index, size_t frag_offset) const
     {
-        EXPECT_EQ(next_offset, m_next_iter.offset);
+        size_t offset = frag_offset;
+
+        for (size_t i = 0; i < elem_index; ++i) {
+            offset += m_lengths[i];
+        }
+
+        return offset;
+    }
+
+    void check_next_iter(size_t next_elem_index, size_t next_frag_offset)
+    {
+        EXPECT_EQ(next_elem_index, m_next_iter.type.sgl.elem_index);
         EXPECT_EQ(next_frag_offset, m_next_iter.type.sgl.frag_offset);
+        EXPECT_EQ(flat_offset(next_elem_index, next_frag_offset),
+                  m_next_iter.offset);
     }
 
-    void check_position(size_t offset, size_t frag_offset)
+    void check_position(size_t elem_index, size_t frag_offset)
     {
-        EXPECT_EQ(offset, m_dt_iter.offset);
+        EXPECT_EQ(elem_index, m_dt_iter.type.sgl.elem_index);
         EXPECT_EQ(frag_offset, m_dt_iter.type.sgl.frag_offset);
+        EXPECT_EQ(flat_offset(elem_index, frag_offset), m_dt_iter.offset);
     }
 
     void start_partial_elem(size_t max_frag_length)
@@ -481,11 +495,13 @@ UCS_TEST_F(test_ucp_dt_sgl, mem_reg_provided_memh_missing_md) {
 
     ucp_datatype_iter_t dt_iter  = {};
     dt_iter.dt_class             = UCP_DATATYPE_SGL;
-    dt_iter.length               = NUM_ELEMS;
+    dt_iter.length               = NUM_ELEMS * LENGTH;
     dt_iter.offset               = 0;
     dt_iter.type.sgl.buffers     = buffers.data();
     dt_iter.type.sgl.lengths     = lengths.data();
     dt_iter.type.sgl.memhs       = memhs.data();
+    dt_iter.type.sgl.elem_count  = NUM_ELEMS;
+    dt_iter.type.sgl.elem_index  = 0;
     dt_iter.type.sgl.frag_offset = 0;
     ucp_memory_info_set_host(&dt_iter.mem_info);
 
@@ -645,12 +661,10 @@ UCS_TEST_F(test_ucp_dt_sgl, iter_next_frag_zero_length) {
     EXPECT_EQ(1u, next_frag(MAX_FRAG));
     check_frag(3, 0, MAX_FRAG, 4, 0);
     advance();
-    EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
 
-    /* Trailing zero-length element produces no descriptor */
-    EXPECT_EQ(0u, next_frag(MAX_FRAG));
-    check_next_iter(5, 0);
-    advance();
+    /* All the bytes are consumed, so the trailing zero-length element does not
+       need another iteration */
+    check_position(4, 0);
     EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
 }
 
@@ -687,14 +701,32 @@ UCS_TEST_F(test_ucp_dt_sgl, iter_seek) {
     init_sgl_iter(2, {96, 64});
     start_partial_elem(MAX_FRAG);
 
-    ucp_datatype_iter_seek(&m_dt_iter, 1, UCS_BIT(UCP_DATATYPE_SGL));
+    /* Seek to an element boundary */
+    ucp_datatype_iter_seek(&m_dt_iter, flat_offset(1, 0),
+                           UCS_BIT(UCP_DATATYPE_SGL));
     check_position(1, 0);
     EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
 
     EXPECT_EQ(1u, next_frag(MAX_FRAG));
     check_frag(1, 0, MAX_FRAG, 1, MAX_FRAG);
 
-    ucp_datatype_iter_seek(&m_dt_iter, 2, UCS_BIT(UCP_DATATYPE_SGL));
+    /* Seek backwards into the middle of an element */
+    ucp_datatype_iter_seek(&m_dt_iter, flat_offset(0, MAX_FRAG),
+                           UCS_BIT(UCP_DATATYPE_SGL));
+    check_position(0, MAX_FRAG);
+
+    EXPECT_EQ(1u, next_frag(MAX_FRAG));
+    check_frag(0, MAX_FRAG, MAX_FRAG, 0, 2 * MAX_FRAG);
+
+    /* Seek forward into the middle of the last element */
+    ucp_datatype_iter_seek(&m_dt_iter, flat_offset(1, MAX_FRAG),
+                           UCS_BIT(UCP_DATATYPE_SGL));
+    check_position(1, MAX_FRAG);
+    EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
+
+    /* Seek to the end */
+    ucp_datatype_iter_seek(&m_dt_iter, flat_offset(2, 0),
+                           UCS_BIT(UCP_DATATYPE_SGL));
     check_position(2, 0);
     EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
 }