UCP/RMA: Fix PUT-SGL protocol selection length

AUTHORS

@@ -76,6 +76,7 @@ Manjunath Gorentla Venkata <manjugv@gmail.com>
 Manu Shantharam <manu.shantharam@amd.com>
 Marek Schimara <Marek.Schimara@bull.net>
 Mark Allen <markalle@us.ibm.com>
+Mason Pitt <bradley.b.pitt@gmail.com>
 Matthew Baker <bakermb@ornl.gov>
 Michael Braverman <michaelbr@nvidia.com>
 Michal Shalev <mshalev@nvidia.com>

src/ucp/dt/datatype_iter.c

@@ -287,7 +287,7 @@ ucs_status_t ucp_datatype_iter_sgl_init(ucp_context_h context,
                                         ucp_datatype_iter_t *dt_iter,
                                         const ucp_dt_local_sgl_t *local,
                                         const ucp_dt_remote_sgl_t *remote,
-                                        size_t count,
+                                        size_t count, size_t length,
                                         const ucp_request_param_t *param)
 {
     ucs_status_t status;
@@ -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               = length;
     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;
@@ -386,7 +388,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));
@@ -412,7 +414,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]);
         }
     }
@@ -510,7 +512,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);
@@ -530,7 +532,7 @@ 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 SGL elements */
+            size_t           elem_index;  /* Current SGL element */
             size_t           frag_offset; /* Offset in the current element */
-            /* length = element count, offset = current element index */
         } sgl;
     } type;
 } ucp_datatype_iter_t;
@@ -118,7 +119,7 @@ ucs_status_t ucp_datatype_iter_sgl_init(ucp_context_h context,
                                         ucp_datatype_iter_t *dt_iter,
                                         const ucp_dt_local_sgl_t *local,
                                         const ucp_dt_remote_sgl_t *remote,
-                                        size_t count,
+                                        size_t count, size_t length,
                                         const ucp_request_param_t *param);
 
 ucs_status_t ucp_datatype_iter_sgl_mem_reg(ucp_context_h context,

src/ucp/dt/datatype_iter.inl

@@ -10,6 +10,7 @@
 
 #include "datatype_iter.h"
 #include "dt.inl"
+#include "dt_sgl.h"
 
 #include <ucp/core/ucp_context.h>
 #include <ucp/core/ucp_worker.h>
@@ -156,10 +157,12 @@ ucp_datatype_iter_init(ucp_context_h context, void *buffer, size_t count,
                                           dt_iter, param);
     } else if (dt_iter->dt_class == UCP_DATATYPE_SGL) {
         *sg_count = 0;
+        length    = ucp_dt_sgl_length(
+                ((const ucp_dt_local_sgl_t*)buffer)->lengths, count);
         return ucp_datatype_iter_sgl_init(context, dt_iter,
                                           (const ucp_dt_local_sgl_t*)buffer,
                                           (const ucp_dt_remote_sgl_t*)param->remote,
-                                          count, param);
+                                          count, length, param);
     } else if (!ENABLE_PARAMS_CHECK ||
                (dt_iter->dt_class == UCP_DATATYPE_GENERIC)) {
         *sg_count = 0;
@@ -384,14 +387,18 @@ 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 +653,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 +679,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 +690,17 @@ 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 "
-                "desc_count=%zu",
-                dt_iter, dt_iter->offset, dt_iter->length,
+    ucs_assertv((elem_index == dt_iter->type.sgl.elem_count) || (desc_count > 0),
+                "dt_iter=%p offset=%zu length=%zu elem_index=%zu "
+                "elem_count=%zu frag_offset=%zu desc_count=%zu",
+                dt_iter, dt_iter->offset, dt_iter->length, 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 +722,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 +748,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 +785,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,27 @@
     (((_datatype) & UCP_DATATYPE_CLASS_MASK) == UCP_DATATYPE_SGL)
 
 
+/**
+ * Calculate the total byte length of an SGL
+ *
+ * @param [in]     lengths        Array of buffer lengths
+ * @param [in]     count          Number of entries in @a lengths
+ *
+ * @return Total byte length
+ */
+static inline size_t ucp_dt_sgl_length(const size_t *lengths, size_t count)
+{
+    size_t total_length = 0;
+    size_t i;
+
+    for (i = 0; i < count; ++i) {
+        total_length += lengths[i];
+    }
+
+    return total_length;
+}
+
+
 /**
  * Check that all SGL entries match the given memory info
  *

src/ucp/proto/proto_common.inl

@@ -323,11 +323,13 @@ ucp_proto_request_send_op(ucp_ep_h ep, ucp_proto_select_t *proto_select,
         return UCS_STATUS_PTR(status);
     }
 
+    req->send.length = req->send.state.dt_iter.length;
+
     ucp_proto_select_param_init(&sel_param, op_id, param->op_attr_mask,
                                 op_flags, req->send.state.dt_iter.dt_class,
                                 &req->send.state.dt_iter.mem_info, sg_count);
 
-    msg_length = req->send.state.dt_iter.length + header_length;
+    msg_length = req->send.length + header_length;
     return ucp_proto_request_send_op_common(worker, ep, proto_select,
                                             rkey_cfg_index, req, param,
                                             &sel_param, msg_length);
@@ -391,6 +393,8 @@ static UCS_F_ALWAYS_INLINE ucs_status_ptr_t ucp_proto_request_send_op_reply(
         return UCS_STATUS_PTR(status);
     }
 
+    req->send.length = req->send.state.dt_iter.length;
+
     UCS_PROFILE_CALL_VOID(ucp_memory_detect, context, param->reply_buffer,
                           contig_length, &reply_mem_info);
 
@@ -401,8 +405,7 @@ static UCS_F_ALWAYS_INLINE ucs_status_ptr_t ucp_proto_request_send_op_reply(
 
     return ucp_proto_request_send_op_common(worker, ep, proto_select,
                                             rkey_cfg_index, req, param,
-                                            &sel_param,
-                                            req->send.state.dt_iter.length);
+                                            &sel_param, req->send.length);
 }
 
 static UCS_F_ALWAYS_INLINE size_t

src/ucp/proto/proto_reconfig.c

@@ -72,7 +72,7 @@ ucp_proto_reconfig_report_no_rma_emulation_no_proto(ucp_request_t *req,
               ucs_memory_type_names[local_mem_type],
               (op_id == UCP_OP_ID_PUT) ? "to" : "from",
               ucs_memory_type_names[remote_mem_type],
-              req->send.state.dt_iter.length);
+              req->send.length);
     return 1;
 }
 
@@ -101,8 +101,6 @@ static ucs_status_t ucp_proto_reconfig_progress(uct_pending_req_t *self)
                                   req->send.proto_config->rkey_cfg_index,
                                   &req->send.proto_config->select_param,
                                   ucp_operation_names, &strb);
-        ucs_error("cannot find remote protocol for: %s",
-                  ucs_string_buffer_cstr(&strb));
 
         /* No protocol can serve this op on the current lane set - fail
          * the EP so the user error callback fires. */

src/ucp/rma/put_offload.c

@@ -445,9 +445,9 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
     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 start_index           = dt_iter->type.sgl.elem_index;
     size_t max_elem_count        = ucs_min(lpriv->max_sgl_zcopy_count,
-                                           dt_iter->length - start_index);
+                                           dt_iter->type.sgl.elem_count - 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;
@@ -459,7 +459,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) ||
@@ -486,15 +487,17 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
             break;
         }
 
+        next_iter->offset += lengths[idx];
         uct_memhs[elem_count] = (sgl_memhs != NULL) ?
                                 sgl_memhs[idx]->uct[md_index] :
                                 UCT_MEM_HANDLE_NULL;
         uct_rkeys[elem_count] = ucp_rkey_get_tl_rkey(sgl_rkeys[idx],
                                                      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,

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);