UCT/IB/MLX5/RC: Implement put_sgl_zcopy

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

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */
@@ -123,6 +123,14 @@ ucs_status_t uct_rc_mlx5_base_ep_put_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_put_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_get_bcopy(uct_ep_h tl_ep, uct_unpack_callback_t unpack_cb,
                               void *arg, size_t length, uint64_t remote_addr,

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

@@ -24,6 +24,16 @@
                                                         uct_rc_mlx5_iface_common_t)
 
 
+static UCS_F_ALWAYS_INLINE size_t
+uct_rc_mlx5_base_put_sgl_zcopy_max_count(uct_rc_mlx5_iface_common_t *iface)
+{
+    size_t max_count = ucs_min((size_t)UCT_RC_MLX5_PUT_SGL_ZCOPY_MAX_COUNT,
+                               iface->super.config.tx_qp_len);
+
+    return ucs_min(max_count, iface->tx.bb_max);
+}
+
+
 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,
                          uct_rkey_t *rkey, uint64_t *addr, uint16_t offset,

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

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */
@@ -134,6 +134,9 @@ enum {
     UCT_RC_MLX5_CQE_APP_OP_TM_CONSUMED_MSG   = 0xA
 };
 
+#define UCT_RC_MLX5_PUT_SGL_ZCOPY_MAX_COUNT 128
+
+
 #define UCT_RC_MLX5_RMA_MAX_IOV(_av_size) \
     ((UCT_IB_MLX5_MAX_SEND_WQE_SIZE - ((_av_size) + \
      sizeof(struct mlx5_wqe_raddr_seg) + sizeof(struct mlx5_wqe_ctrl_seg))) / \

src/uct/ib/mlx5/rc/rc_mlx5_ep.c

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */
@@ -175,6 +175,119 @@ ucs_status_t uct_rc_mlx5_base_ep_put_zcopy(uct_ep_h tl_ep, const uct_iov_t *iov,
     return status;
 }
 
+ucs_status_t
+uct_rc_mlx5_base_ep_put_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)
+{
+    UCT_RC_MLX5_BASE_EP_DECL(tl_ep, iface, ep);
+    uct_ib_mlx5_txwq_t *txwq       = &ep->tx.wq;
+    size_t total                   = 0;
+    struct mlx5_wqe_ctrl_seg *ctrl = NULL;
+    struct mlx5_wqe_raddr_seg *raddr;
+    struct mlx5_wqe_data_seg *dptr;
+    size_t wqe_size, i;
+    uint8_t fm_ce_se, fence_flag;
+    uint16_t sn, pi, res_count;
+    uint64_t addr;
+    uct_rkey_t rkey;
+    void *curr;
+    size_t UCS_V_UNUSED max_count;
+
+    max_count = uct_rc_mlx5_base_put_sgl_zcopy_max_count(iface);
+    UCT_CHECK_PARAM(count <= max_count,
+                    "put_sgl_zcopy count(%zu) should be limited by %zu", count,
+                    max_count);
+    ucs_assert(ep->super.flags & UCT_RC_EP_FLAG_CONNECTED);
+    ucs_assert(count > 0);
+
+    /* TODO: add strided elements support */
+    if (ucs_unlikely((counts != NULL) || (strides != NULL))) {
+        ucs_error("put_sgl_zcopy does not support strided elements");
+        return UCS_ERR_UNSUPPORTED;
+    }
+
+    if (iface->super.tx.cq_available < (int)count) {
+        UCS_STATS_UPDATE_COUNTER(iface->super.stats,
+                                 UCT_RC_IFACE_STAT_NO_CQE, 1);
+        UCS_STATS_UPDATE_COUNTER(ep->super.super.stats,
+                                 UCT_EP_STAT_NO_RES, 1);
+        return UCS_ERR_NO_RESOURCE;
+    }
+
+    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);
+
+    /* Validate all lengths before consuming the fence state */
+    for (i = 0; i < count; i++) {
+        UCT_CHECK_LENGTH(lengths[i], 0, UCT_IB_MAX_MESSAGE_SIZE,
+                         "put_sgl_zcopy");
+    }
+
+    wqe_size = sizeof(*ctrl) + sizeof(*raddr) + sizeof(*dptr);
+    sn       = txwq->sw_pi;
+
+    ucs_assert(!(txwq->flags & UCT_IB_MLX5_TXWQ_FLAG_FAILED));
+    ucs_assert(ucs_div_round_up(wqe_size, MLX5_SEND_WQE_BB) == 1);
+
+    pi   = sn;
+    curr = txwq->curr;
+
+    fence_flag = uct_rc_ep_fm(&iface->super, &txwq->fi, 1) ?
+                 iface->config.put_fence_flag : 0;
+
+    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;
+
+        uct_ib_mlx5_set_ctrl_seg(ctrl, pi, MLX5_OPCODE_RDMA_WRITE, 0,
+                                 txwq->super.qp_num, fm_ce_se, 0, wqe_size);
+
+        addr = remote_addrs[i];
+        if (fence_flag != 0) {
+            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);
+
+        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]));
+
+        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;
+
+    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_txqp_add_send_comp(&iface->super, &ep->super.txqp,
+                              uct_rc_ep_send_op_completion_handler, comp, sn,
+                              UCT_RC_IFACE_SEND_OP_FLAG_ZCOPY, NULL, 0, total);
+
+    UCT_TL_EP_STAT_OP(&ep->super.super, PUT, ZCOPY, total);
+    uct_rc_ep_enable_flush_remote(&ep->super);
+
+    return UCS_INPROGRESS;
+}
+
 ucs_status_t
 uct_rc_mlx5_base_ep_get_bcopy(uct_ep_h tl_ep, uct_unpack_callback_t unpack_cb,
                               void *arg, size_t length, uint64_t remote_addr,

src/uct/ib/mlx5/rc/rc_mlx5_iface.c

@@ -1012,13 +1012,41 @@ static UCS_CLASS_DEFINE_NEW_FUNC(uct_rc_mlx5_iface_t, uct_iface_t, uct_md_h,
 static UCS_CLASS_DEFINE_DELETE_FUNC(uct_rc_mlx5_iface_t, uct_iface_t);
 
 static ucs_status_t
-uct_rc_mlx5_iface_query_v2(uct_iface_h UCS_V_UNUSED iface,
+uct_rc_mlx5_ep_put_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)
+{
+    ucs_status_t status;
+
+    status = uct_ib_mlx5_ext_ep_put_sgl_zcopy(ep, buffers, lengths, memhs,
+                                              remote_addrs, rkeys, counts,
+                                              strides, count, comp);
+    if (status == UCS_ERR_UNSUPPORTED) {
+        status = uct_rc_mlx5_base_ep_put_sgl_zcopy(ep, buffers, lengths, memhs,
+                                                   remote_addrs, rkeys, counts,
+                                                   strides, count, comp);
+    }
+
+    return status;
+}
+
+static ucs_status_t
+uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface,
                            uct_iface_attr_v2_t *iface_attr)
 {
-    size_t max_sgl = uct_ib_mlx5_ext_max_put_sgl_zcopy_count();
+    uct_rc_mlx5_iface_common_t *iface = ucs_derived_of(tl_iface,
+                                                       uct_rc_mlx5_iface_common_t);
+    size_t max_sgl                    = uct_ib_mlx5_ext_max_put_sgl_zcopy_count();
+
+    if (max_sgl == 0) {
+        max_sgl = uct_rc_mlx5_base_put_sgl_zcopy_max_count(iface);
+    }
 
     if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) {
-        iface_attr->cap.flags = (max_sgl > 0) ? UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY : 0;
+        iface_attr->cap.flags = UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY;
     }
 
     if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_MAX_PUT_SGL_ZCOPY_COUNT) {
@@ -1040,7 +1068,7 @@ static uct_rc_iface_ops_t uct_rc_mlx5_iface_ops = {
             .iface_is_reachable_v2  = uct_rc_mlx5_iface_is_reachable_v2,
             .ep_is_connected        = uct_rc_mlx5_base_ep_is_connected,
             .ep_get_device_ep       = (uct_ep_get_device_ep_func_t)ucs_empty_function_return_unsupported,
-            .ep_put_sgl_zcopy       = uct_ib_mlx5_ext_ep_put_sgl_zcopy
+            .ep_put_sgl_zcopy       = uct_rc_mlx5_ep_put_sgl_zcopy
         },
         .create_cq      = uct_rc_mlx5_iface_common_create_cq,
         .destroy_cq     = uct_rc_mlx5_iface_common_destroy_cq,