UCT/IB/MLX5/RC: Implement get_sgl_zcopy

src/uct/ib/mlx5/ib_mlx5_ext.c

@@ -69,6 +69,46 @@ uct_ib_mlx5_ext_ep_put_sgl_zcopy(uct_ep_h ep, void * const *buffers,
     return UCS_ERR_UNSUPPORTED;
 }
 
+size_t uct_ib_mlx5_ext_max_get_sgl_zcopy_count(void)
+{
+    uct_ib_mlx5_ext_plugin_t *plugin;
+
+    ucs_list_for_each(plugin, &uct_ib_mlx5_ext_plugins, list) {
+        if (ucs_unlikely(uct_ib_mlx5_ext_is_unsupported_op(
+                    (const void*)plugin->ops.max_get_sgl_zcopy_count))) {
+            continue;
+        }
+
+        return plugin->ops.max_get_sgl_zcopy_count();
+    }
+
+    return 0;
+}
+
+ucs_status_t
+uct_ib_mlx5_ext_ep_get_sgl_zcopy(uct_ep_h ep, void * const *buffers,
+                                 const size_t *lengths, uct_mem_h const *memhs,
+                                 const uint64_t *remote_addrs,
+                                 uct_rkey_t const *rkeys, const size_t *counts,
+                                 const size_t *strides, size_t count,
+                                 uct_completion_t *comp)
+{
+    uct_ib_mlx5_ext_plugin_t *plugin;
+
+    ucs_list_for_each(plugin, &uct_ib_mlx5_ext_plugins, list) {
+        if (ucs_unlikely(uct_ib_mlx5_ext_is_unsupported_op(
+                    (const void*)plugin->ops.ep_get_sgl_zcopy))) {
+            continue;
+        }
+
+        return plugin->ops.ep_get_sgl_zcopy(ep, buffers, lengths, memhs,
+                                            remote_addrs, rkeys, counts,
+                                            strides, count, comp);
+    }
+
+    return UCS_ERR_UNSUPPORTED;
+}
+
 void uct_ib_mlx5_ext_cleanup(void)
 {
     uct_ib_mlx5_ext_plugin_t *plugin, *tmp;
@@ -120,12 +160,17 @@ ucs_status_t uct_ib_mlx5_ext_register(const uct_ib_mlx5_ext_ops_t *ops)
     num_plugins = ucs_list_length(&uct_ib_mlx5_ext_plugins);
 
     ucs_debug("ib mlx5 ext: registered plugin name=%s put_sgl_zcopy=%s "
+              "get_sgl_zcopy=%s "
               "(total=%u)",
               plugin->ops.name,
               uct_ib_mlx5_ext_is_unsupported_op(
                       (const void*)plugin->ops.ep_put_sgl_zcopy) ?
                       "unsupported" :
                       "supported",
+              uct_ib_mlx5_ext_is_unsupported_op(
+                      (const void*)plugin->ops.ep_get_sgl_zcopy) ?
+                      "unsupported" :
+                      "supported",
               num_plugins);
     return UCS_OK;
 }

src/uct/ib/mlx5/ib_mlx5_ext.h

@@ -29,13 +29,24 @@ BEGIN_C_DECLS
  */
 typedef size_t (*uct_ib_mlx5_ext_max_put_sgl_zcopy_count_func_t)(void);
 
+/**
+ * @brief External plugin maximum GET SGL zero-copy entry count callback.
+ *
+ * @return Maximum number of SGL entries supported by the plugin's
+ *         @ref uct_ib_mlx5_ext_ep_get_sgl_zcopy implementation, or 0 if
+ *         unsupported.
+ */
+typedef size_t (*uct_ib_mlx5_ext_max_get_sgl_zcopy_count_func_t)(void);
+
 /**
  * @brief External plugin operations.
  */
 typedef struct uct_ib_mlx5_ext_ops {
     char                                           name[UCT_COMPONENT_NAME_MAX]; /**< Plugin name */
     uct_ib_mlx5_ext_max_put_sgl_zcopy_count_func_t max_put_sgl_zcopy_count;      /**< Maximum PUT SGL zero-copy entry count callback */
     uct_ep_put_sgl_zcopy_func_t                    ep_put_sgl_zcopy;             /**< PUT SGL zero-copy callback */
+    uct_ib_mlx5_ext_max_get_sgl_zcopy_count_func_t max_get_sgl_zcopy_count;      /**< Maximum GET SGL zero-copy entry count callback */
+    uct_ep_get_sgl_zcopy_func_t                    ep_get_sgl_zcopy;             /**< GET SGL zero-copy callback */
 } uct_ib_mlx5_ext_ops_t;
 
 /**
@@ -72,6 +83,19 @@ ucs_status_t uct_ib_mlx5_ext_ep_put_sgl_zcopy(uct_ep_h ep,
                                               size_t count,
                                               uct_completion_t *comp);
 
+size_t uct_ib_mlx5_ext_max_get_sgl_zcopy_count(void);
+
+ucs_status_t uct_ib_mlx5_ext_ep_get_sgl_zcopy(uct_ep_h ep,
+                                              void * const *buffers,
+                                              const size_t *lengths,
+                                              uct_mem_h const *memhs,
+                                              const uint64_t *remote_addrs,
+                                              uct_rkey_t const *rkeys,
+                                              const size_t *counts,
+                                              const size_t *strides,
+                                              size_t count,
+                                              uct_completion_t *comp);
+
 END_C_DECLS
 
 #endif /* UCT_IB_MLX5_EXT_H_ */

src/uct/ib/mlx5/rc/rc_mlx5.h

@@ -25,6 +25,14 @@
                                    UCS_STATUS_PTR(UCS_ERR_NO_RESOURCE))
 
 
+#define UCT_RC_MLX5_SGL_ZCOPY_CHECK_STRIDES(_counts, _strides, _name) \
+    /* TODO: add strided elements support */ \
+    if (ucs_unlikely(((_counts) != NULL) || ((_strides) != NULL))) { \
+        ucs_error("%s does not support strided elements", (_name)); \
+        return UCS_ERR_UNSUPPORTED; \
+    }
+
+
 enum {
     UCT_RC_MLX5_IFACE_ADDR_TYPE_BASIC,
 
@@ -145,6 +153,14 @@ ucs_status_t uct_rc_mlx5_base_ep_get_zcopy(uct_ep_h tl_ep, const uct_iov_t *iov,
                                            uct_rkey_t rkey,
                                            uct_completion_t *comp);
 
+ucs_status_t
+uct_rc_mlx5_base_ep_get_sgl_zcopy(uct_ep_h tl_ep, void * const *buffers,
+                                  const size_t *lengths, uct_mem_h const *memhs,
+                                  const uint64_t *remote_addrs,
+                                  uct_rkey_t const *rkeys, const size_t *counts,
+                                  const size_t *strides, size_t count,
+                                  uct_completion_t *comp);
+
 ucs_status_t uct_rc_mlx5_base_ep_am_short(uct_ep_h tl_ep, uint8_t id,
                                           uint64_t header, const void *payload,
                                           unsigned length);

src/uct/ib/mlx5/rc/rc_mlx5.inl

@@ -25,7 +25,7 @@
 
 
 static UCS_F_ALWAYS_INLINE size_t
-uct_rc_mlx5_base_put_sgl_zcopy_max_count(uct_rc_mlx5_iface_common_t *iface)
+uct_rc_mlx5_base_sgl_zcopy_max_count(uct_rc_mlx5_iface_common_t *iface)
 {
     size_t max_count;
 
@@ -37,6 +37,29 @@ uct_rc_mlx5_base_put_sgl_zcopy_max_count(uct_rc_mlx5_iface_common_t *iface)
     return max_count;
 }
 
+static UCS_F_ALWAYS_INLINE ucs_status_t
+uct_rc_mlx5_base_ep_sgl_zcopy_check(uct_rc_mlx5_iface_common_t *iface,
+                                    uct_rc_mlx5_base_ep_t *ep,
+                                    const size_t *lengths, size_t count,
+                                    size_t min_length, size_t max_length,
+                                    const char *name)
+{
+    size_t i;
+
+    ucs_assert(ep->super.flags & UCT_RC_EP_FLAG_CONNECTED);
+    ucs_assert(count > 0);
+
+    for (i = 0; i < count; i++) {
+        UCT_CHECK_LENGTH(lengths[i], min_length, max_length, name);
+    }
+
+    UCT_RC_CHECK_CQE_VALUE_RET(&iface->super, &ep->super, UCS_ERR_NO_RESOURCE,
+                               count - 1);
+    UCT_RC_CHECK_NUM_RDMA_READ_RET(&iface->super, UCS_ERR_NO_RESOURCE);
+    UCT_RC_CHECK_TXQP_VALUE_RET(&iface->super, &ep->super, UCS_ERR_NO_RESOURCE,
+                                count - 1);
+    return UCS_OK;
+}
 
 static UCS_F_ALWAYS_INLINE void
 uct_rc_mlx5_ep_fence_put(uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_txwq_t *txwq,
@@ -481,6 +504,100 @@ uct_rc_mlx5_txwq_update_psn(uct_ib_mlx5_txwq_t *txwq, int qp_type,
     }
 }
 
+static UCS_F_ALWAYS_INLINE size_t
+uct_rc_mlx5_base_ep_sgl_zcopy_post(uct_rc_mlx5_base_ep_t *ep,
+                                   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,
+                                   uint8_t opcode, uint8_t fence_flag,
+                                   int resolve_atomic_rkey,
+                                   uct_rc_send_handler_t handler,
+                                   uint16_t op_flags, uct_completion_t *comp)
+{
+    uct_ib_mlx5_txwq_t *txwq       = &ep->tx.wq;
+    size_t total                   = 0;
+    struct mlx5_wqe_ctrl_seg *ctrl = NULL;
+    uint32_t num_packets           = 0;
+    const size_t non_data_wqe_size = sizeof(struct mlx5_wqe_ctrl_seg) +
+                                     sizeof(struct mlx5_wqe_raddr_seg);
+    const size_t data_wqe_size     = non_data_wqe_size +
+                                     sizeof(struct mlx5_wqe_data_seg);
+    uct_rc_mlx5_iface_common_t *iface;
+    struct mlx5_wqe_raddr_seg *raddr;
+    struct mlx5_wqe_data_seg *dptr;
+    size_t wqe_size, i;
+    uint8_t fm_ce_se;
+    uint16_t pi, res_count;
+    uint64_t addr;
+    uct_rkey_t rkey;
+    void *curr;
+
+    ucs_assert(count > 0);
+
+    pi = txwq->sw_pi;
+
+    ucs_assert(!(txwq->flags & UCT_IB_MLX5_TXWQ_FLAG_FAILED));
+    ucs_assert(ucs_div_round_up(data_wqe_size, MLX5_SEND_WQE_BB) == 1);
+
+    curr = txwq->curr;
+
+    for (i = 0; i < count; i++) {
+        fm_ce_se = ((i == 0) ? fence_flag : 0) |
+                   ((i == count - 1) ? MLX5_WQE_CTRL_CQ_UPDATE : 0);
+        ctrl     = curr;
+
+        addr = remote_addrs[i];
+        if (resolve_atomic_rkey) {
+            rkey = uct_ib_resolve_atomic_rkey(rkeys[i],
+                                              ep->super.atomic_mr_offset,
+                                              &addr);
+        } else {
+            rkey = uct_ib_md_direct_rkey(rkeys[i]);
+        }
+
+        raddr = uct_ib_mlx5_txwq_wrap_none(txwq, ctrl + 1);
+        uct_ib_mlx5_ep_set_rdma_seg(raddr, addr, rkey);
+
+        if (ucs_unlikely(lengths[i] == 0)) {
+            wqe_size     = non_data_wqe_size;
+            num_packets += 1;
+        } else {
+            dptr         = uct_ib_mlx5_txwq_wrap_none(txwq, raddr + 1);
+            uct_ib_mlx5_set_data_seg(dptr, buffers[i], lengths[i],
+                                     uct_ib_memh_get_lkey(memhs[i]));
+            wqe_size     = data_wqe_size;
+            num_packets += uct_rc_mlx5_num_packets(txwq, lengths[i]);
+        }
+
+        uct_ib_mlx5_set_ctrl_seg(ctrl, pi, opcode, 0, txwq->super.qp_num,
+                                 fm_ce_se, 0, wqe_size);
+
+        curr = UCS_PTR_BYTE_OFFSET(ctrl, MLX5_SEND_WQE_BB);
+        curr = uct_ib_mlx5_txwq_wrap_exact(txwq, curr);
+        pi++;
+        total += lengths[i];
+    }
+
+    res_count         = pi - 1 - txwq->prev_sw_pi;
+    txwq->prev_sw_pi += res_count;
+    txwq->sw_pi       = pi;
+    txwq->curr        = curr;
+    txwq->sig_pi      = txwq->prev_sw_pi;
+
+    iface = ucs_derived_of(ep->super.super.super.iface,
+                           uct_rc_mlx5_iface_common_t);
+    uct_rc_txqp_posted(&ep->super.txqp, &iface->super, res_count, 1);
+    uct_ib_mlx5_txwq_ring_doorbell(txwq, ctrl, txwq->sw_pi, 1);
+    uct_rc_mlx5_txwq_add_psn(txwq, IBV_QPT_RC, num_packets);
+
+    uct_rc_txqp_add_send_comp_sgl_always(
+            &iface->super, &ep->super.txqp, handler, comp, txwq->sig_pi,
+            UCT_RC_IFACE_SEND_OP_FLAG_ZCOPY | op_flags, buffers, lengths,
+            count, total);
+    return total;
+}
+
 static UCS_F_ALWAYS_INLINE void
 uct_rc_mlx5_common_post_send(uct_rc_mlx5_iface_common_t *iface, int qp_type,
                              uct_rc_txqp_t *txqp, uct_ib_mlx5_txwq_t *txwq,