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,12 +296,14 @@ 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.remote_addrs = remote->remote_addrs;
     dt_iter->type.sgl.rkeys        = remote->rkeys;
+    dt_iter->type.sgl.elem_count   = count;
+    dt_iter->type.sgl.elem_index   = 0;
 
     if (ucs_unlikely(count == 0)) {
         dt_iter->type.sgl.memhs = NULL;
@@ -347,7 +349,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;
@@ -381,7 +383,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));
@@ -407,7 +409,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]);
         }
     }
@@ -505,7 +507,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);
@@ -525,7 +527,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

@@ -79,7 +79,8 @@ typedef struct {
             ucp_mem_h        *memhs;
             const uint64_t   *remote_addrs;
             ucp_rkey_h const *rkeys;
-            /* length = element count, offset = current element index */
+            size_t           elem_count;
+            size_t           elem_index;
         } sgl;
     } type;
 } ucp_datatype_iter_t;
@@ -119,7 +120,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>
@@ -140,6 +141,8 @@ ucp_datatype_iter_init(ucp_context_h context, void *buffer, size_t count,
                        int is_pack, ucp_datatype_iter_t *dt_iter,
                        uint8_t *sg_count, const ucp_request_param_t *param)
 {
+    const ucp_dt_local_sgl_t *local_sgl;
+    const ucp_dt_remote_sgl_t *remote_sgl;
     size_t length;
 
     dt_iter->dt_class = ucp_datatype_class(datatype);
@@ -155,11 +158,12 @@ ucp_datatype_iter_init(ucp_context_h context, void *buffer, size_t count,
         return ucp_datatype_iov_iter_init(context, buffer, count, length,
                                           dt_iter, param);
     } else if (dt_iter->dt_class == UCP_DATATYPE_SGL) {
-        *sg_count = 0;
-        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);
+        local_sgl  = (const ucp_dt_local_sgl_t*)buffer;
+        remote_sgl = (const ucp_dt_remote_sgl_t*)param->remote;
+        *sg_count  = 0;
+        length     = ucp_dt_sgl_length(local_sgl->lengths, count);
+        return ucp_datatype_iter_sgl_init(context, dt_iter, local_sgl,
+                                          remote_sgl, count, length, param);
     } else if (!ENABLE_PARAMS_CHECK ||
                (dt_iter->dt_class == UCP_DATATYPE_GENERIC)) {
         *sg_count = 0;
@@ -623,8 +627,21 @@ ucp_datatype_iter_next_sgl(const ucp_datatype_iter_t *dt_iter,
                            size_t max_elem_count,
                            ucp_datatype_iter_t *next_iter)
 {
+    size_t elem_index, elem_count;
+
     ucs_assert(dt_iter->dt_class == UCP_DATATYPE_SGL);
-    return ucp_datatype_iter_next(dt_iter, max_elem_count, next_iter);
+    elem_index = dt_iter->type.sgl.elem_index;
+    ucs_assert(elem_index <= dt_iter->type.sgl.elem_count);
+    elem_count = ucs_min(max_elem_count,
+                         dt_iter->type.sgl.elem_count - elem_index);
+
+    next_iter->offset = dt_iter->offset +
+                        ucp_dt_sgl_length(dt_iter->type.sgl.lengths +
+                                          elem_index, elem_count);
+    next_iter->type.sgl.elem_index = elem_index + elem_count;
+    ucs_assert(next_iter->offset <= dt_iter->length);
+
+    return elem_count;
 }
 
 /*
@@ -643,6 +660,9 @@ ucp_datatype_iter_copy_position(ucp_datatype_iter_t *dt_iter,
     if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_IOV, dt_mask)) {
         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;
     }
 }
 
@@ -665,6 +685,9 @@ ucp_datatype_iter_rewind(ucp_datatype_iter_t *dt_iter, unsigned dt_mask)
     if (ucp_datatype_iter_is_class(dt_iter, UCP_DATATYPE_IOV, 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;
     }
 }
 

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_reconfig.c

@@ -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

@@ -376,7 +376,7 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
     uct_ep_h uct_ep              = ucp_ep_get_lane(ep, lane);
     ucp_md_index_t md_index      = ucp_ep_md_index(ep, lane);
     ucp_rsc_index_t rkey_index   = lpriv->super.rkey_index;
-    size_t start_index           = dt_iter->offset;
+    size_t start_elem_index      = dt_iter->type.sgl.elem_index;
     size_t max_sgl_count         = lpriv->max_sgl_zcopy_count;
     size_t elem_count            = ucp_datatype_iter_next_sgl(dt_iter,
                                                               max_sgl_count,
@@ -402,18 +402,19 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
 
     for (i = 0; i < elem_count; i++) {
         uct_memhs[i] = (sgl_memhs != NULL) ?
-                       sgl_memhs[start_index + i]->uct[md_index] :
+                       sgl_memhs[start_elem_index + i]->uct[md_index] :
                        UCT_MEM_HANDLE_NULL;
         uct_rkeys[i] = ucp_rkey_get_tl_rkey(
-                           dt_iter->type.sgl.rkeys[start_index + i], rkey_index);
+                           dt_iter->type.sgl.rkeys[start_elem_index + i],
+                           rkey_index);
     }
 
     status = uct_ep_put_sgl_zcopy(
                  uct_ep,
-                 &dt_iter->type.sgl.buffers[start_index],
-                 &dt_iter->type.sgl.lengths[start_index],
+                 &dt_iter->type.sgl.buffers[start_elem_index],
+                 &dt_iter->type.sgl.lengths[start_elem_index],
                  uct_memhs,
-                 &dt_iter->type.sgl.remote_addrs[start_index],
+                 &dt_iter->type.sgl.remote_addrs[start_elem_index],
                  uct_rkeys,
                  NULL, NULL,
                  elem_count, &req->send.state.uct_comp);
@@ -424,10 +425,10 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req,
     if (!UCS_STATUS_IS_ERR(status)) {
         ucp_proto_put_offload_update_remote_flush(
                 ep, lpriv->flush_sys_dev_mask,
-                ucp_rkey_get_tl_rkey(dt_iter->type.sgl.rkeys[start_index],
+                ucp_rkey_get_tl_rkey(dt_iter->type.sgl.rkeys[start_elem_index],
                                      rkey_index),
                 uct_ep,
-                dt_iter->type.sgl.remote_addrs[start_index]);
+                dt_iter->type.sgl.remote_addrs[start_elem_index]);
     }
 
     return status;
@@ -477,21 +478,21 @@ ucp_proto_put_sgl_offload_sw_send_func(ucp_request_t *req,
     uct_ep_h uct_ep              = ucp_ep_get_lane(ep, lane);
     ucp_md_index_t md_index      = ucp_ep_md_index(ep, lane);
     ucp_rsc_index_t rkey_index   = lpriv->super.rkey_index;
-    size_t index                 = dt_iter->offset;
+    size_t elem_index            = dt_iter->type.sgl.elem_index;
     ucp_mem_h *sgl_memhs         = dt_iter->type.sgl.memhs;
-    uint64_t remote_addr         = dt_iter->type.sgl.remote_addrs[index];
+    uint64_t remote_addr         = dt_iter->type.sgl.remote_addrs[elem_index];
     uct_rkey_t tl_rkey;
     uct_iov_t iov;
     ucs_status_t status;
 
     /* TODO: fragment elements larger than cap.put.max_zcopy */
     ucp_datatype_iter_next_sgl(dt_iter, 1, next_iter);
 
-    tl_rkey     = ucp_rkey_get_tl_rkey(dt_iter->type.sgl.rkeys[index],
+    tl_rkey     = ucp_rkey_get_tl_rkey(dt_iter->type.sgl.rkeys[elem_index],
                                        rkey_index);
-    iov.buffer  = dt_iter->type.sgl.buffers[index];
-    iov.length  = dt_iter->type.sgl.lengths[index];
-    iov.memh    = (sgl_memhs != NULL) ? sgl_memhs[index]->uct[md_index] :
+    iov.buffer  = dt_iter->type.sgl.buffers[elem_index];
+    iov.length  = dt_iter->type.sgl.lengths[elem_index];
+    iov.memh    = (sgl_memhs != NULL) ? sgl_memhs[elem_index]->uct[md_index] :
                                         UCT_MEM_HANDLE_NULL;
     iov.stride  = 0;
     iov.count   = 1;

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, 280);
+    EXPECTED_SIZE(ucp_request_t, 296);
     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

@@ -11,6 +11,7 @@
 
 extern "C" {
 #include <ucp/dt/dt.h>
+#include <ucp/dt/dt_sgl.h>
 #include <ucp/dt/datatype_iter.inl>
 #include <ucp/core/ucp_rkey.h>
 }
@@ -333,11 +334,13 @@ class test_ucp_dt_sgl : public ucs::test {
         param.datatype     = ucp_dt_make_sgl();
         param.remote       = &m_remote;
 
-        ucs_status_t status = ucp_datatype_iter_sgl_init(m_ucph.get(),
-                                                         &m_dt_iter, &m_local,
-                                                         &m_remote, count,
-                                                         &param);
+        size_t length = ucp_dt_sgl_length(m_lengths.data(), count);
+        ucs_status_t status = ucp_datatype_iter_sgl_init(
+                m_ucph.get(), &m_dt_iter, &m_local, &m_remote, count, length,
+                &param);
         ASSERT_UCS_OK(status);
+        ASSERT_EQ(length, m_dt_iter.length);
+        ASSERT_EQ(count, m_dt_iter.type.sgl.elem_count);
     }
 
 private:
@@ -366,13 +369,14 @@ UCS_TEST_F(test_ucp_dt_sgl, iter_next_chunked) {
         size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter,
                                                      MAX_PER_STEP,
                                                      &next_iter);
-        EXPECT_LE(advanced, MAX_PER_STEP);
-        EXPECT_GT(advanced, 0u);
+        ASSERT_LE(advanced, MAX_PER_STEP);
+        ASSERT_GT(advanced, 0u);
         total_advanced += advanced;
         ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
     }
 
     EXPECT_EQ(NUM_ELEMS, total_advanced);
+    EXPECT_EQ(NUM_ELEMS, m_dt_iter.type.sgl.elem_index);
     EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
 }
 
@@ -384,7 +388,7 @@ UCS_TEST_F(test_ucp_dt_sgl, iter_next_single_step) {
         ucp_datatype_iter_t next_iter = {};
         size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter, count,
                                                      &next_iter);
-        EXPECT_EQ(count, advanced);
+        ASSERT_EQ(count, advanced);
         ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
         EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
     }
@@ -396,18 +400,35 @@ UCS_TEST_F(test_ucp_dt_sgl, iter_next_one_by_one) {
     init_sgl_iter(NUM_ELEMS);
 
     for (size_t i = 0; i < NUM_ELEMS; i++) {
-        EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
+        ASSERT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
         ucp_datatype_iter_t next_iter = {};
         size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter, 1,
                                                      &next_iter);
-        EXPECT_EQ(1u, advanced);
-        EXPECT_EQ(i + 1, next_iter.offset);
+        ASSERT_EQ(1u, advanced);
+        ASSERT_EQ(i + 1, next_iter.type.sgl.elem_index);
+        ASSERT_EQ(ucp_dt_sgl_length(m_dt_iter.type.sgl.lengths, i + 1),
+                  next_iter.offset);
         ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
     }
 
     EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
 }
 
+UCS_TEST_F(test_ucp_dt_sgl, iter_rewind) {
+    ucp_datatype_iter_t next_iter = {};
+
+    init_sgl_iter(3);
+    ASSERT_EQ(2u, ucp_datatype_iter_next_sgl(&m_dt_iter, 2, &next_iter));
+    ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
+    ASSERT_EQ(ucp_dt_sgl_length(m_dt_iter.type.sgl.lengths, 2),
+              m_dt_iter.offset);
+    ASSERT_EQ(2u, m_dt_iter.type.sgl.elem_index);
+
+    ucp_datatype_iter_rewind(&m_dt_iter, UINT_MAX);
+    ASSERT_EQ(0u, m_dt_iter.offset);
+    ASSERT_EQ(0u, m_dt_iter.type.sgl.elem_index);
+}
+
 UCS_TEST_F(test_ucp_dt_sgl, init_zero_count) {
     init_sgl_iter(0);