diff --git a/src/uct/ib/mlx5/gga/gga_mlx5.c b/src/uct/ib/mlx5/gga/gga_mlx5.c index 0528ffe5d95..a804fbab21c 100644 --- a/src/uct/ib/mlx5/gga/gga_mlx5.c +++ b/src/uct/ib/mlx5/gga/gga_mlx5.c @@ -532,6 +532,7 @@ uct_gga_mlx5_ep_connect_to_ep_v2(uct_ep_h tl_ep, return status; } + uct_rc_mlx5_txwq_set_path_mtu(&ep->super.tx.wq, path_mtu); ep->super.super.atomic_mr_offset = 0; ep->super.super.flags |= UCT_RC_EP_FLAG_CONNECTED; ep->super.super.flush_rkey = diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index 40cbba23578..07bf2544307 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -693,12 +693,15 @@ void uct_ib_mlx5_devx_uar_cleanup(uct_ib_mlx5_devx_uar_t *uar) void uct_ib_mlx5_txwq_reset(uct_ib_mlx5_txwq_t *txwq) { - txwq->curr = txwq->qstart; - txwq->sw_pi = 0; - txwq->prev_sw_pi = UINT16_MAX; + txwq->curr = txwq->qstart; + txwq->sw_pi = 0; + txwq->prev_sw_pi = UINT16_MAX; + txwq->next_wqe_psn = 0; + txwq->path_mtu_mask = 0; + txwq->path_mtu_shift = 0; #if UCS_ENABLE_ASSERT - txwq->hw_ci = 0xFFFF; - txwq->flags = 0; + txwq->hw_ci = 0xFFFF; + txwq->flags = 0; #endif uct_ib_fence_info_init(&txwq->fi); } @@ -712,6 +715,16 @@ void uct_ib_mlx5_init_wq_buf(uct_ib_mlx5_txwq_t *txwq) uct_ib_mlx5_set_ctrl_qpn_ds(uct_ib_mlx5_txwq_get_wqe(txwq, 0xffff), 0, 1); } +static void +uct_ib_mlx5_txwq_vfs_show_next_wqe_psn(void *obj, ucs_string_buffer_t *strb, + void *arg_ptr, uint64_t arg_u64) +{ + uct_ib_mlx5_txwq_t *txwq = arg_ptr; + + ucs_string_buffer_appendf( + strb, "%u\n", uct_ib_mlx5_txwq_get_next_wqe_psn(txwq)); +} + void uct_ib_mlx5_txwq_vfs_populate(uct_ib_mlx5_txwq_t *txwq, void *parent_obj) { ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, @@ -721,6 +734,9 @@ void uct_ib_mlx5_txwq_vfs_populate(uct_ib_mlx5_txwq_t *txwq, void *parent_obj) UCS_VFS_TYPE_U16, "sw_pi"); ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->prev_sw_pi, UCS_VFS_TYPE_U16, "prev_sw_pi"); + ucs_vfs_obj_add_ro_file(parent_obj, + uct_ib_mlx5_txwq_vfs_show_next_wqe_psn, txwq, 0, + "next_wqe_psn"); ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->qstart, UCS_VFS_TYPE_POINTER, "qstart"); ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->qend, diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index 0e013b798ab..aa245ef7f34 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -686,6 +686,8 @@ typedef struct uct_ib_mlx5_txwq { uct_ib_mlx5_qp_t super; uint16_t sw_pi; /* PI for next WQE */ uint16_t prev_sw_pi; /* PI where last WQE *started* */ + uint32_t next_wqe_psn; /* 1st PSN of the next WQE to be + posted */ uct_ib_mlx5_mmio_reg_t *reg; void *curr; volatile uint32_t *dbrec; @@ -693,6 +695,8 @@ typedef struct uct_ib_mlx5_txwq { void *qend; uint16_t bb_max; uint16_t sig_pi; /* PI for last signaled WQE */ + uint16_t path_mtu_mask; /* Path MTU in bytes - 1 */ + uint8_t path_mtu_shift; /* log2(path MTU in bytes) */ #if UCS_ENABLE_ASSERT uint16_t hw_ci; /* First BB index of last completed WQE */ uint8_t flags; /* Debug flags */ @@ -701,6 +705,13 @@ typedef struct uct_ib_mlx5_txwq { } uct_ib_mlx5_txwq_t; +static UCS_F_ALWAYS_INLINE uint32_t +uct_ib_mlx5_txwq_get_next_wqe_psn(const uct_ib_mlx5_txwq_t *txwq) +{ + return txwq->next_wqe_psn & UCS_MASK(24); +} + + /* Receive work-queue */ typedef struct uct_ib_mlx5_rxwq { /* producer index. It updated when new receive wqe is posted */ diff --git a/src/uct/ib/mlx5/ib_mlx5.inl b/src/uct/ib/mlx5/ib_mlx5.inl index 193ccb80bd6..7b5916e6ec1 100644 --- a/src/uct/ib/mlx5/ib_mlx5.inl +++ b/src/uct/ib/mlx5/ib_mlx5.inl @@ -509,10 +509,13 @@ uct_ib_mlx5_set_data_seg(struct mlx5_wqe_data_seg *dptr, static UCS_F_ALWAYS_INLINE size_t uct_ib_mlx5_set_data_seg_iov(uct_ib_mlx5_txwq_t *txwq, struct mlx5_wqe_data_seg *dptr, - const uct_iov_t *iov, size_t iovcnt) + const uct_iov_t *iov, size_t iovcnt, + size_t *iov_length_p) { - size_t wqe_size = 0; + size_t iov_length = 0; + size_t wqe_size = 0; size_t iov_it; + size_t length; for (iov_it = 0; iov_it < iovcnt; ++iov_it) { if (!iov[iov_it].length) { /* Skip zero length WQE*/ @@ -522,13 +525,18 @@ size_t uct_ib_mlx5_set_data_seg_iov(uct_ib_mlx5_txwq_t *txwq, /* place data into the buffer */ dptr = uct_ib_mlx5_txwq_wrap_any(txwq, dptr); - uct_ib_mlx5_set_data_seg(dptr, iov[iov_it].buffer, - uct_iov_get_length(iov + iov_it), + length = uct_iov_get_length(iov + iov_it); + uct_ib_mlx5_set_data_seg(dptr, iov[iov_it].buffer, length, uct_ib_memh_get_lkey(iov[iov_it].memh)); - wqe_size += sizeof(*dptr); + iov_length += length; + wqe_size += sizeof(*dptr); ++dptr; } + if (iov_length_p != NULL) { + *iov_length_p = iov_length; + } + return wqe_size; } diff --git a/src/uct/ib/mlx5/rc/rc_mlx5.h b/src/uct/ib/mlx5/rc/rc_mlx5.h index f40be42c0c5..04f2fc39038 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5.h +++ b/src/uct/ib/mlx5/rc/rc_mlx5.h @@ -213,6 +213,9 @@ ucs_status_t uct_rc_mlx5_iface_create_qp(uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_txwq_t *txwq, uct_ib_mlx5_qp_attr_t *attr); +void uct_rc_mlx5_txwq_set_path_mtu(uct_ib_mlx5_txwq_t *txwq, + enum ibv_mtu path_mtu); + ucs_status_t uct_rc_mlx5_ep_connect_qp(uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_qp_t *qp, uint32_t qp_num, diff --git a/src/uct/ib/mlx5/rc/rc_mlx5.inl b/src/uct/ib/mlx5/rc/rc_mlx5.inl index 50ae3e943f8..302b9197086 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5.inl +++ b/src/uct/ib/mlx5/rc/rc_mlx5.inl @@ -452,6 +452,35 @@ uct_rc_mlx5_ep_fm_cq_update(uct_rc_mlx5_iface_common_t *iface, return fm_ce_se; } +static UCS_F_ALWAYS_INLINE uint32_t +uct_rc_mlx5_num_packets(const uct_ib_mlx5_txwq_t *txwq, + size_t message_length) +{ + ucs_assert(txwq->path_mtu_shift > 0); + + return (message_length + txwq->path_mtu_mask) >> txwq->path_mtu_shift; +} + +static UCS_F_ALWAYS_INLINE void +uct_rc_mlx5_txwq_add_psn(uct_ib_mlx5_txwq_t *txwq, int qp_type, + uint32_t num_packets) +{ + if (qp_type == IBV_QPT_RC) { + txwq->next_wqe_psn += num_packets; + } +} + +static UCS_F_ALWAYS_INLINE void +uct_rc_mlx5_txwq_update_psn(uct_ib_mlx5_txwq_t *txwq, int qp_type, + size_t message_length) +{ + if (qp_type == IBV_QPT_RC) { + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, + uct_rc_mlx5_num_packets(txwq, + message_length)); + } +} + 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, @@ -532,6 +561,7 @@ static UCS_F_ALWAYS_INLINE void uct_rc_mlx5_txqp_inline_iov_post( uct_rc_mlx5_common_post_send(iface, qp_type, txqp, txwq, MLX5_OPCODE_SEND, 0, fm_ce_se, dci_channel, wqe_size, 0, INT_MAX, NULL); + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, 1); } /* @@ -627,6 +657,10 @@ uct_rc_mlx5_txqp_inline_post(uct_rc_mlx5_iface_common_t *iface, int qp_type, uct_rc_mlx5_common_post_send(iface, qp_type, txqp, txwq, opcode, 0, fm_ce_se, dci_channel, wqe_size, imm_val_be, max_log_sge, NULL); + if (opcode != MLX5_OPCODE_NOP) { + /* Inline short operations always fit in one RC packet. */ + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, 1); + } } /* @@ -707,6 +741,7 @@ uct_rc_mlx5_txqp_dptr_post(uct_rc_mlx5_iface_common_t *iface, int qp_type, /* Data segment */ if (length == 0) { wqe_size = ctrl_av_size + sizeof(*raddr); + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, 1); } else { /* dptr cannot wrap, because ctrl+av could be either 2 or 4 segs */ dptr = uct_ib_mlx5_txwq_wrap_none(txwq, raddr + 1); @@ -810,6 +845,7 @@ uct_rc_mlx5_txqp_dptr_post(uct_rc_mlx5_iface_common_t *iface, int qp_type, (opcode_flags & UCT_RC_MLX5_OPCODE_MASK), opmod, fm_ce_se, dci_channel, wqe_size, imm_val_be, max_log_sge, log_sge); + uct_rc_mlx5_txwq_update_psn(txwq, qp_type, length); } static UCS_F_ALWAYS_INLINE @@ -830,6 +866,7 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty struct mlx5_wqe_inl_data_seg *inl; uct_rc_mlx5_hdr_t *rch; unsigned wqe_size, inl_seg_size, ctrl_av_size; + size_t iov_length, message_length; void *next_seg; uint8_t opmod; #if HAVE_MLX5_MMO @@ -852,9 +889,6 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty inl_seg_size = ucs_align_up_pow2(sizeof(*inl) + sizeof(*rch) + am_hdr_len, UCT_IB_MLX5_WQE_SEG_SIZE); - ucs_assert(uct_iov_total_length(iov, iovcnt) + sizeof(*rch) + am_hdr_len <= - iface->super.super.config.seg_size); - /* Inline segment with AM ID and header */ inl = next_seg; inl->byte_count = htonl((sizeof(*rch) + am_hdr_len) | MLX5_INLINE_SEG); @@ -866,9 +900,12 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty /* Data segment with payload */ dptr = (struct mlx5_wqe_data_seg *)((char *)inl + inl_seg_size); wqe_size = ctrl_av_size + inl_seg_size + - uct_ib_mlx5_set_data_seg_iov(txwq, dptr, iov, iovcnt); + uct_ib_mlx5_set_data_seg_iov( + txwq, dptr, iov, iovcnt, &iov_length); opmod = 0; + message_length = iov_length + sizeof(*rch) + am_hdr_len; + ucs_assert(message_length <= iface->super.super.config.seg_size); ucs_assert(wqe_size <= UCT_IB_MLX5_MAX_SEND_WQE_SIZE); break; @@ -881,8 +918,10 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty inl->byte_count = htonl(sizeof(struct ibv_tmh) | MLX5_INLINE_SEG); dptr = uct_ib_mlx5_txwq_wrap_exact(txwq, (char *)inl + inl_seg_size); wqe_size = ctrl_av_size + inl_seg_size + - uct_ib_mlx5_set_data_seg_iov(txwq, dptr, iov, iovcnt); + uct_ib_mlx5_set_data_seg_iov( + txwq, dptr, iov, iovcnt, &iov_length); opmod = 0; + message_length = iov_length + sizeof(struct ibv_tmh); uct_rc_mlx5_fill_tmh((struct ibv_tmh*)(inl + 1), tag, app_ctx, IBV_TMH_EAGER); @@ -895,16 +934,16 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty /* Fall through */ case MLX5_OPCODE_RDMA_WRITE: /* Set RDMA segment */ - ucs_assert(uct_iov_total_length(iov, iovcnt) <= UCT_IB_MAX_MESSAGE_SIZE); - raddr = next_seg; uct_ib_mlx5_ep_set_rdma_seg(raddr, remote_addr, rkey); /* Data segment */ - wqe_size = ctrl_av_size + sizeof(*raddr) + - uct_ib_mlx5_set_data_seg_iov(txwq, (void*)(raddr + 1), - iov, iovcnt); - opmod = 0; + wqe_size = ctrl_av_size + sizeof(*raddr) + + uct_ib_mlx5_set_data_seg_iov( + txwq, (void*)(raddr + 1), iov, iovcnt, + &message_length); + opmod = 0; + ucs_assert(message_length <= UCT_IB_MAX_MESSAGE_SIZE); break; #if HAVE_MLX5_MMO @@ -938,8 +977,10 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty uct_ib_mlx5_set_data_seg(dptr /* scatter segment */, dma_dst_buf, iov[0].length, dma_dst_key); - wqe_size = sizeof(*ctrl) + sizeof(*dma_seg) + (2 * sizeof(*dptr)); - opmod = UCT_IB_MLX5_OPMOD_MMO_DMA; + wqe_size = sizeof(*ctrl) + sizeof(*dma_seg) + + (2 * sizeof(*dptr)); + opmod = UCT_IB_MLX5_OPMOD_MMO_DMA; + message_length = 0; break; #endif default: @@ -950,6 +991,18 @@ void uct_rc_mlx5_txqp_dptr_post_iov(uct_rc_mlx5_iface_common_t *iface, int qp_ty opcode_flags & UCT_RC_MLX5_OPCODE_MASK, opmod, fm_ce_se, dci_channel, wqe_size, ib_imm_be, max_log_sge, NULL); +#if HAVE_MLX5_MMO + if ((opcode_flags & UCT_RC_MLX5_OPCODE_MASK) == MLX5_OPCODE_MMO) { + return; + } +#endif + + uct_rc_mlx5_txwq_update_psn(txwq, qp_type, message_length); + if (opcode_flags == MLX5_OPCODE_RDMA_WRITE) { + /* opcode_flags is constant after inlining, so only zero-length + * PUT_ZCOPY gets the branchless one-PSN adjustment. */ + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, !message_length); + } } /* @@ -1099,6 +1152,8 @@ uct_rc_mlx5_txqp_tag_inline_post(uct_rc_mlx5_iface_common_t *iface, int qp_type, uct_rc_mlx5_common_post_send(iface, qp_type, txqp, txwq, opcode, 0, fm_ce_se, dci_channel, wqe_size, imm_val_be, INT_MAX, NULL); + /* Inline tag operations always fit in one RC packet. */ + uct_rc_mlx5_txwq_add_psn(txwq, qp_type, 1); } static UCS_F_ALWAYS_INLINE void diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index 95e48a643ba..e1cded70c09 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -212,6 +212,7 @@ uct_rc_mlx5_base_ep_put_sgl_zcopy(uct_ep_h tl_ep, void * const *buffers, 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; struct mlx5_wqe_raddr_seg *raddr; struct mlx5_wqe_data_seg *dptr; size_t wqe_size, i; @@ -288,7 +289,8 @@ uct_rc_mlx5_base_ep_put_sgl_zcopy(uct_ep_h tl_ep, void * const *buffers, curr = UCS_PTR_BYTE_OFFSET(ctrl, MLX5_SEND_WQE_BB); curr = uct_ib_mlx5_txwq_wrap_exact(txwq, curr); pi++; - total += lengths[i]; + total += lengths[i]; + num_packets += uct_rc_mlx5_num_packets(txwq, lengths[i]); } res_count = pi - 1 - txwq->prev_sw_pi; @@ -299,6 +301,7 @@ uct_rc_mlx5_base_ep_put_sgl_zcopy(uct_ep_h tl_ep, void * const *buffers, 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(&iface->super, &ep->super.txqp, uct_rc_ep_send_op_completion_handler, comp, sn, @@ -877,6 +880,18 @@ void uct_rc_mlx5_common_packet_dump(uct_base_iface_t *iface, uct_am_trace_type_t valid_length, buffer, max); } +void uct_rc_mlx5_txwq_set_path_mtu(uct_ib_mlx5_txwq_t *txwq, + enum ibv_mtu path_mtu) +{ + size_t mtu = uct_ib_mtu_value(path_mtu); + + ucs_assert(mtu <= UINT16_MAX); + ucs_assert(ucs_is_pow2(mtu)); + + txwq->path_mtu_mask = mtu - 1; + txwq->path_mtu_shift = ucs_ilog2(mtu); +} + ucs_status_t uct_rc_mlx5_ep_connect_qp(uct_rc_mlx5_iface_common_t *iface, uct_ib_mlx5_qp_t *qp, uint32_t qp_num, @@ -974,6 +989,8 @@ uct_rc_mlx5_ep_connect_to_ep_v2(uct_ep_h tl_ep, return status; } + uct_rc_mlx5_txwq_set_path_mtu(&ep->super.tx.wq, path_mtu); + ep->super.super.atomic_mr_offset = uct_ib_md_atomic_offset( rc_addr->atomic_mr_id); ep->super.super.flags |= UCT_RC_EP_FLAG_CONNECTED; diff --git a/src/uct/ib/mlx5/ud/ud_mlx5.c b/src/uct/ib/mlx5/ud/ud_mlx5.c index e4df3a927e9..89b0453877e 100644 --- a/src/uct/ib/mlx5/ud/ud_mlx5.c +++ b/src/uct/ib/mlx5/ud/ud_mlx5.c @@ -320,9 +320,9 @@ static UCS_F_ALWAYS_INLINE ucs_status_t uct_ud_mlx5_ep_inline_iov_post( /* set iov to dptr */ if (iovcnt > 0) { wqe_size = ucs_align_up_pow2(wqe_size, UCT_IB_MLX5_WQE_SEG_SIZE); - wqe_size += uct_ib_mlx5_set_data_seg_iov(&iface->tx.wq, - UCS_PTR_BYTE_OFFSET(ctrl, wqe_size), - iov, iovcnt); + wqe_size += uct_ib_mlx5_set_data_seg_iov( + &iface->tx.wq, UCS_PTR_BYTE_OFFSET(ctrl, wqe_size), iov, + iovcnt, NULL); } uct_ud_mlx5_post_send(iface, ep, 0, ctrl, wqe_size, neth, diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index e9946df00c3..56935f9efaa 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -144,6 +144,108 @@ UCS_TEST_P(test_rc, fence_am_short_consumed, "RC_FENCE=weak") UCT_INSTANTIATE_RC_TEST_CASE(test_rc) +#ifdef HAVE_MLX5_DV +class test_rc_mlx5_psn : public test_rc { +public: + void init() override + { + uct_test::init(); + + modify_config("IB_PATH_MTU", "4096"); + m_e1 = uct_test::create_entity(0); + m_entities.push_back(m_e1); + + check_skip_test(); + + modify_config("IB_PATH_MTU", "512"); + m_e2 = uct_test::create_entity(0); + m_entities.push_back(m_e2); + + connect(); + } + +protected: + uct_ib_mlx5_txwq_t *txwq() + { + return &ucs_derived_of(m_e1->ep(0), + uct_rc_mlx5_ep_t)->super.tx.wq; + } +}; + +UCS_TEST_P(test_rc_mlx5_psn, path_mtu_from_rc_qp) +{ + uct_ib_iface_t *local_iface = + ucs_derived_of(m_e1->iface(), uct_ib_iface_t); + uct_ib_iface_t *peer_iface = + ucs_derived_of(m_e2->iface(), uct_ib_iface_t); + + EXPECT_EQ(IBV_MTU_4096, local_iface->config.path_mtu); + EXPECT_EQ(IBV_MTU_512, peer_iface->config.path_mtu); + EXPECT_EQ(511, txwq()->path_mtu_mask); + EXPECT_EQ(9, txwq()->path_mtu_shift); +} + +UCS_TEST_SKIP_COND_P(test_rc_mlx5_psn, next_wqe_psn, + !check_caps(UCT_IFACE_FLAG_PUT_ZCOPY)) +{ + static const size_t length = 513; + uct_ib_mlx5_txwq_t *wq = txwq(); + mapped_buffer sendbuf(length, 0ul, *m_e1); + mapped_buffer recvbuf(length, 0ul, *m_e2); + uct_iov_t zero_iov = {}; + uct_completion_t comp; + ucs_status_t status; + + EXPECT_EQ(0, uct_ib_mlx5_txwq_get_next_wqe_psn(wq)); + + ASSERT_UCS_OK(uct_ep_am_short(m_e1->ep(0), 0, 0, NULL, 0)); + EXPECT_EQ(1, uct_ib_mlx5_txwq_get_next_wqe_psn(wq)); + + comp.func = [](uct_completion_t*) {}; + comp.count = 1; + comp.status = UCS_OK; + + UCS_TEST_GET_BUFFER_IOV(iov, iovcnt, sendbuf.ptr(), sendbuf.length(), + sendbuf.memh(), 1); + status = uct_ep_put_zcopy(m_e1->ep(0), iov, iovcnt, recvbuf.addr(), + recvbuf.rkey(), &comp); + ASSERT_UCS_OK_OR_INPROGRESS(status); + + /* The RC QP uses the peer's 512-byte path MTU, so 513 bytes use 2 PSNs. */ + EXPECT_EQ(3, uct_ib_mlx5_txwq_get_next_wqe_psn(wq)); + + if (status == UCS_INPROGRESS) { + wait_for_value(&comp.count, 0, true); + } + + zero_iov.buffer = sendbuf.ptr(); + zero_iov.length = 0; + zero_iov.memh = sendbuf.memh(); + zero_iov.stride = 0; + zero_iov.count = 1; + comp.count = 1; + comp.status = UCS_OK; + + status = uct_ep_put_zcopy(m_e1->ep(0), &zero_iov, 1, recvbuf.addr(), + recvbuf.rkey(), &comp); + ASSERT_UCS_OK_OR_INPROGRESS(status); + EXPECT_EQ(4, uct_ib_mlx5_txwq_get_next_wqe_psn(wq)); + + if (status == UCS_INPROGRESS) { + wait_for_value(&comp.count, 0, true); + } + + wq->next_wqe_psn = UCS_MASK(24); + ASSERT_UCS_OK(uct_ep_am_short(m_e1->ep(0), 0, 0, NULL, 0)); + EXPECT_EQ(UCS_BIT(24), wq->next_wqe_psn); + EXPECT_EQ(0, uct_ib_mlx5_txwq_get_next_wqe_psn(wq)); + + flush(); +} + +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_psn, rc_mlx5) +#endif + class test_rc_max_wr : public test_rc { protected: