UCP/RMA: Add SGL support to ucp_get_nbx

src/ucp/am/eager_multi.c

@@ -46,7 +46,8 @@ ucp_am_eager_multi_bcopy_proto_probe(const ucp_proto_init_params_t *init_params)
         .first.tl_cap_flags  = UCT_IFACE_FLAG_AM_BCOPY,
         .middle.lane_type    = UCP_LANE_TYPE_AM_BW,
         .middle.tl_cap_flags = UCT_IFACE_FLAG_AM_BCOPY,
-        .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID
+        .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID,
+        .min_sgl_frag_offs   = UCP_PROTO_COMMON_OFFSET_INVALID
     };
 
     if (!ucp_am_check_init_params(init_params, UCP_PROTO_AM_OP_ID_MASK,
@@ -213,6 +214,7 @@ static void ucp_am_eager_multi_zcopy_proto_probe_common(
         .min_chunk           = 0,
         .initial_reg_md_map  = 0,
         .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID,
+        .min_sgl_frag_offs   = UCP_PROTO_COMMON_OFFSET_INVALID,
         .first.lane_type     = UCP_LANE_TYPE_AM,
         .first.tl_cap_flags  = UCT_IFACE_FLAG_AM_ZCOPY,
         .middle.lane_type    = UCP_LANE_TYPE_AM_BW,

src/ucp/api/ucp_def.h

@@ -162,8 +162,8 @@ typedef struct ucp_rkey                  *ucp_rkey_h;
  * @brief Invalid remote address sentinel.
  *
  * This value should be passed as the @a remote_addr parameter of
- * @ref ucp_put_nbx when remote addresses are provided through an SGL
- * descriptor (@ref ucp_dt_remote_sgl_t) instead.
+ * @ref ucp_put_nbx or @ref ucp_get_nbx when remote addresses are provided
+ * through an SGL descriptor (@ref ucp_dt_remote_sgl_t) instead.
  */
 #define UCP_REMOTE_ADDR_INVALID  UINT64_MAX
 
@@ -173,8 +173,8 @@ typedef struct ucp_rkey                  *ucp_rkey_h;
  * @brief Invalid remote key sentinel.
  *
  * This value should be passed as the @a rkey parameter of
- * @ref ucp_put_nbx when remote keys are provided through an SGL
- * descriptor (@ref ucp_dt_remote_sgl_t) instead.
+ * @ref ucp_put_nbx or @ref ucp_get_nbx when remote keys are provided through
+ * an SGL descriptor (@ref ucp_dt_remote_sgl_t) instead.
  */
 #define UCP_RKEY_INVALID         NULL
 

src/ucp/core/ucp_request.inl

@@ -218,7 +218,7 @@ UCS_PTR_MAP_IMPL(request, 0);
               (((_param)->datatype & UCP_DATATYPE_CLASS_MASK) == \
                UCP_DATATYPE_SGL)))) { \
             ucs_error("SGL datatype and remote descriptor parameters are " \
-                      "only supported for ucp_put_nbx"); \
+                      "only supported for ucp_put_nbx and ucp_get_nbx"); \
             return UCS_STATUS_PTR(UCS_ERR_INVALID_PARAM); \
         } \
     } while (0)

src/ucp/proto/proto.c

@@ -35,6 +35,8 @@
     _macro(ucp_put_rndv_proto) \
     _macro(ucp_put_sgl_offload_proto) \
     _macro(ucp_put_sgl_offload_sw_proto) \
+    _macro(ucp_get_sgl_offload_proto) \
+    _macro(ucp_get_sgl_offload_sw_proto) \
     _macro(ucp_eager_bcopy_multi_proto) \
     _macro(ucp_eager_sync_bcopy_multi_proto) \
     _macro(ucp_eager_zcopy_multi_proto) \

src/ucp/proto/proto_multi.c

@@ -876,6 +876,10 @@ ucp_proto_multi_init_priv(const ucp_proto_multi_init_params_t *params,
         lpriv->flush_sys_dev_mask =
                 ucp_proto_multi_init_flush_sys_dev_mask(params, lane);
 
+        lpriv->min_sgl_frag   = ucp_proto_common_get_iface_attr_field(
+                ucp_proto_common_get_iface_attr(&params->super.super, lane),
+                params->min_sgl_frag_offs, 0);
+
         if (v2_cap_flags & (UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY |
                             UCT_IFACE_FLAG_V2_GET_SGL_ZCOPY)) {
             rsc_index = ucp_proto_common_get_rsc_index(&params->super.super,

src/ucp/proto/proto_multi.h

@@ -90,6 +90,11 @@ typedef struct {
      * uct_ep_get_sgl_zcopy on this lane, cached from uct_iface_attr_v2 at
      * protocol init when SGL zcopy is selected, otherwise zero */
      size_t                      max_sgl_zcopy_count;
+
+    /* Minimal length of a single SGL element fragment posted on this lane,
+     * cached from uct_iface_attr_t at protocol init according to
+     * ucp_proto_multi_init_params_t::min_sgl_frag_offs */
+    size_t                      min_sgl_frag;
 } ucp_proto_multi_lane_priv_t;
 
 
@@ -128,6 +133,12 @@ typedef struct {
      * messages that do not use all selected lanes */
     int                            use_single_lane_min_length;
 
+    /* Offset of the iface attribute (uct_iface_attr_t) holding the minimal
+     * length of a single SGL element fragment, cached in
+     * ucp_proto_multi_lane_priv_t::min_sgl_frag. Can be
+     * UCP_PROTO_COMMON_OFFSET_INVALID if there is no minimal length */
+    ptrdiff_t                      min_sgl_frag_offs;
+
     /* MDs on which the buffer is expected to be already registered, so no need
        to account for the overhead of registering on them */
     ucp_md_map_t                   initial_reg_md_map;

src/ucp/proto/proto_multi.inl

@@ -381,4 +381,193 @@ ucp_proto_multi_rma_init_func(ucp_request_t *req)
     return ucp_ep_rma_handle_fence(req->send.ep, req, mpriv->lane_map);
 }
 
+typedef ucs_status_t (*ucp_proto_rma_sgl_post_func_t)(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        void *const *buffers, const size_t *lengths, uct_mem_h const *memhs,
+        const uint64_t *remote_addrs, uct_rkey_t const *rkeys, size_t count);
+
+/* Returns UCS_INPROGRESS if the fragment was already submitted */
+typedef ucs_status_t (*ucp_proto_rma_sgl_adjust_func_t)(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        size_t elem_index, void **buffer_p, size_t *length_p,
+        uint64_t *remote_addr_p);
+
+/* Split the SGL by its total length, like a multi-lane contiguous message */
+static UCS_F_ALWAYS_INLINE size_t
+ucp_proto_rma_sgl_max_payload(const ucp_request_t *req,
+                              const ucp_proto_multi_lane_priv_t *lpriv)
+{
+    size_t length = req->send.rma.sgl.length;
+
+    return ucp_proto_multi_lane_payload(lpriv, length, length,
+                                        lpriv->max_frag);
+}
+
+static UCS_F_INLINE_OPTIMIZED ucs_status_t
+ucp_proto_rma_sgl_offload_send_frags(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        size_t max_frag_count, ucp_datatype_iter_t *next_iter,
+        ucp_proto_rma_sgl_post_func_t post_func,
+        ucp_proto_rma_sgl_adjust_func_t adjust_func)
+{
+    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;
+    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_status_t status;
+
+    ucs_assert(max_frag_count <= UCP_MAX_IOV);
+    ucs_assert((adjust_func == NULL) || (max_frag_count == 1));
+
+    if (max_frag_count == 1) {
+        /* Silence compiler warning, a single fragment is read by the adjust
+         * and post functions */
+        buffers[0]      = NULL;
+        lengths[0]      = 0;
+        remote_addrs[0] = 0;
+    }
+
+    frag_count = ucp_datatype_iter_next_sgl_frags(
+            dt_iter, req->send.rma.sgl.remote_addrs, max_frag_count,
+            ucp_proto_rma_sgl_max_payload(req, lpriv), next_iter, buffers,
+            lengths, remote_addrs, elem_indices);
+    if (frag_count == 0) {
+        return UCS_OK;
+    }
+
+    if (adjust_func != NULL) {
+        status = adjust_func(req, lpriv, elem_indices[0], &buffers[0],
+                             &lengths[0], &remote_addrs[0]);
+        if (ucs_unlikely(status != UCS_OK)) {
+            return (status == UCS_INPROGRESS) ? UCS_OK : status;
+        }
+    }
+
+    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 post_func(req, lpriv, buffers, lengths, uct_memhs, remote_addrs,
+                     uct_rkeys, frag_count);
+}
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_proto_rma_sgl_offload_send_func(
+        ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv,
+        ucp_datatype_iter_t *next_iter,
+        ucp_proto_rma_sgl_post_func_t post_func,
+        ucp_proto_rma_sgl_adjust_func_t adjust_func)
+{
+    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_payload           = ucp_proto_rma_sgl_max_payload(req, lpriv);
+    size_t min_length            = ucs_max(lpriv->min_sgl_frag, 1);
+    size_t max_frag_count        = ucs_min(lpriv->max_sgl_zcopy_count,
+                                           UCP_MAX_IOV);
+    size_t frags_count           = (adjust_func == NULL) ? max_frag_count : 1;
+    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(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, length, total_length;
+
+    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) ||
+                     (lengths[start_index] < min_length) ||
+                     (lengths[start_index] > max_payload))) {
+        return ucp_proto_rma_sgl_offload_send_frags(req, lpriv, frags_count,
+                                                    next_iter, post_func,
+                                                    adjust_func);
+    }
+
+    /* 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 < min_length) || ((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((adjust_func == NULL) && (total_length < max_payload) &&
+                     (elem_count < max_frag_count) &&
+                     ((start_index + elem_count) < dt_iter->length))) {
+        return ucp_proto_rma_sgl_offload_send_frags(req, lpriv, frags_count,
+                                                    next_iter, post_func,
+                                                    adjust_func);
+    }
+
+    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;
+    }
+
+    uct_memhs = ucs_alloc_on_stack(uct_memhs_size, "uct_sgl_memhs");
+    if (uct_memhs == NULL) {
+        ucs_free_on_stack(uct_rkeys, uct_rkeys_size);
+        return UCS_ERR_NO_MEMORY;
+    }
+
+    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;
+    next_iter->type.sgl.frag_offset = 0;
+
+    status = post_func(req, lpriv, &buffers[start_index], &lengths[start_index],
+                       uct_memhs, &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;
+}
+
 #endif