UCT/RC: Reserve TX credits for multi-credit sends

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

@@ -2109,7 +2109,7 @@ uct_rc_mlx5_iface_poll_tx(uct_rc_mlx5_iface_common_t *iface, int poll_flags)
                    ep, qp_num, hw_ci);
 
     uct_rc_mlx5_txqp_process_tx_cqe(&ep->super.txqp, cqe, hw_ci);
-    ucs_arbiter_group_schedule(&iface->super.tx.arbiter, &ep->super.arb_group);
+    uct_rc_ep_schedule_pending(&ep->super);
     uct_rc_mlx5_ep_update_tx_qp_res(ep, hw_ci);
     uct_rc_mlx5_iface_update_tx_cq_res(&iface->super, ep, hw_ci);
     uct_rc_iface_arbiter_dispatch(&iface->super);

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

@@ -209,6 +209,7 @@ void uct_rc_mlx5_iface_handle_failure(uct_ib_iface_t *ib_iface, void *arg,
 
     ucs_arbiter_group_purge(&iface->tx.arbiter, &ep->super.arb_group,
                             uct_rc_ep_arbiter_purge_internal_cb, NULL);
+    uct_rc_ep_schedule_pending(&ep->super);
     uct_ib_mlx5_txwq_update_flags(&ep->tx.wq, UCT_IB_MLX5_TXWQ_FLAG_FAILED, 0);
 
     if (ep->super.flags & (UCT_RC_EP_FLAG_ERR_HANDLER_INVOKED |

src/uct/ib/rc/base/rc_ep.c

@@ -146,10 +146,9 @@ UCS_CLASS_INIT_FUNC(uct_rc_ep_t, uct_rc_iface_t *iface, uint32_t qp_num,
         return status;
     }
 
-    self->path_index   = UCT_EP_PARAMS_GET_PATH_INDEX(params);
-    self->flags        = 0;
-    self->txqp_reserve = 0;
-    self->cq_reserve   = 0;
+    self->path_index = UCT_EP_PARAMS_GET_PATH_INDEX(params);
+    self->flags      = 0;
+    uct_rc_ep_reset_reserve(self);
 
     status = uct_rc_fc_init(&self->fc, iface UCS_STATS_ARG(self->super.stats));
     if (status != UCS_OK) {
@@ -425,8 +424,9 @@ ucs_arbiter_cb_result_t uct_rc_ep_process_pending(ucs_arbiter_t *arbiter,
 
     status = uct_rc_iface_invoke_pending_cb(iface, req);
     if (status == UCS_OK) {
-        ep->txqp_reserve = 0;
-        ep->cq_reserve   = 0;
+        /* Keep the reserve because other requests may still be queued, and new
+         * sends must not pass them. It is reset when the group is empty, on TX
+         * completion or purge. */
         return UCS_ARBITER_CB_RESULT_REMOVE_ELEM;
     } else if (status == UCS_INPROGRESS) {
         return UCS_ARBITER_CB_RESULT_NEXT_GROUP;
@@ -512,6 +512,8 @@ void uct_rc_ep_pending_purge(uct_ep_h tl_ep, uct_pending_purge_callback_t cb,
 
     ucs_arbiter_group_purge(&iface->tx.arbiter, &ep->arb_group,
                             uct_rc_ep_arbiter_purge_cb, &args);
+
+    uct_rc_ep_reset_reserve(ep);
 }
 
 ucs_status_t uct_rc_ep_fc_grant(uct_pending_req_t *self)

src/uct/ib/rc/base/rc_ep.h

@@ -96,15 +96,14 @@ enum {
 #define UCT_RC_CHECK_CQE_RET(_iface, _ep, _ret) \
     /* tx_moderation == 0 for TLs which don't support it */ \
     if (ucs_unlikely((_iface)->tx.cq_available <= \
-        (signed)(_iface)->config.tx_moderation)) { \
+                     (signed)ucs_max((_iface)->config.tx_moderation, \
+                                     (_ep)->cq_reserve))) { \
         UCT_RC_CHECK_CQE_VALUE_RET(_iface, _ep, _ret, 0); \
         (_ep)->txqp.unsignaled = RC_UNSIGNALED_INF; \
     }
 
 #define UCT_RC_CHECK_CQE_VALUE_RET(_iface, _ep, _ret, _value) \
-    if ((_iface)->tx.cq_available <= (signed)(_value)) { \
-        (_ep)->cq_reserve = ucs_max((_ep)->cq_reserve, \
-                                    (uint16_t)(_value)); \
+    if (uct_rc_ep_reserve_cqe(_iface, _ep, (uint16_t)(_value))) { \
         UCS_STATS_UPDATE_COUNTER((_iface)->stats, UCT_RC_IFACE_STAT_NO_CQE, 1); \
         UCS_STATS_UPDATE_COUNTER((_ep)->super.stats, UCT_EP_STAT_NO_RES, 1); \
         return _ret; \
@@ -114,9 +113,7 @@ enum {
     UCT_RC_CHECK_TXQP_VALUE_RET(_iface, _ep, _ret, 0)
 
 #define UCT_RC_CHECK_TXQP_VALUE_RET(_iface, _ep, _ret, _value) \
-    if (uct_rc_txqp_available(&(_ep)->txqp) <= (_value)) { \
-        (_ep)->txqp_reserve = ucs_max((_ep)->txqp_reserve, \
-                                      (uint16_t)(_value)); \
+    if (uct_rc_ep_reserve_txqp(_ep, (uint16_t)(_value))) { \
         UCS_STATS_UPDATE_COUNTER((_ep)->txqp.stats, UCT_RC_TXQP_STAT_QP_FULL, 1); \
         UCS_STATS_UPDATE_COUNTER((_ep)->super.stats, UCT_EP_STAT_NO_RES, 1); \
         return _ret; \
@@ -375,6 +372,51 @@ static UCS_F_ALWAYS_INLINE void uct_rc_fc_restore_wnd(uct_rc_iface_t *iface,
     UCS_STATS_SET_COUNTER(fc->stats, UCT_RC_FC_STAT_FC_WND, fc->fc_wnd);
 }
 
+static UCS_F_ALWAYS_INLINE void uct_rc_ep_reset_reserve(uct_rc_ep_t *ep)
+{
+    ep->txqp_reserve = 0;
+    ep->cq_reserve   = 0;
+}
+
+static UCS_F_ALWAYS_INLINE void
+uct_rc_ep_schedule_pending(uct_rc_ep_t *ep)
+{
+    uct_rc_iface_t *iface = ucs_derived_of(ep->super.super.iface, uct_rc_iface_t);
+
+    if (ucs_arbiter_group_is_empty(&ep->arb_group)) {
+        uct_rc_ep_reset_reserve(ep);
+    } else {
+        ucs_arbiter_group_schedule_nonempty(&iface->tx.arbiter, &ep->arb_group);
+    }
+}
+
+static UCS_F_ALWAYS_INLINE int
+uct_rc_ep_reserve_txqp(uct_rc_ep_t *ep, uint16_t value)
+{
+    uint16_t reserve = ucs_max(ep->txqp_reserve, value);
+
+    if (uct_rc_txqp_available(&ep->txqp) > reserve) {
+        return 0;
+    }
+
+    ep->txqp_reserve = reserve;
+    return 1;
+}
+
+static UCS_F_ALWAYS_INLINE int
+uct_rc_ep_reserve_cqe(const uct_rc_iface_t *iface, uct_rc_ep_t *ep,
+                      uint16_t value)
+{
+    uint16_t reserve = ucs_max(ep->cq_reserve, value);
+
+    if (iface->tx.cq_available > (signed)reserve) {
+        return 0;
+    }
+
+    ep->cq_reserve = reserve;
+    return 1;
+}
+
 static UCS_F_ALWAYS_INLINE int uct_rc_ep_has_tx_resources(uct_rc_ep_t *ep)
 {
     uct_rc_iface_t *iface = ucs_derived_of(ep->super.super.iface, uct_rc_iface_t);

src/uct/ib/rc/verbs/rc_verbs_iface.c

@@ -101,6 +101,7 @@ static void uct_rc_verbs_handle_failure(uct_ib_iface_t *ib_iface, void *arg,
                                   ep->txcnt.ci + count, 0);
     ucs_arbiter_group_purge(&iface->tx.arbiter, &ep->super.arb_group,
                             uct_rc_ep_arbiter_purge_internal_cb, NULL);
+    uct_rc_ep_schedule_pending(&ep->super);
     uct_rc_verbs_update_tx_res(iface, ep, count);
 
     if (ep->super.flags & (UCT_RC_EP_FLAG_ERR_HANDLER_INVOKED |
@@ -156,8 +157,7 @@ uct_rc_verbs_iface_poll_tx(uct_rc_verbs_iface_t *iface)
                        iface, wc[i].wr_id, ep, wc[i].qp_num, count);
 
         uct_rc_txqp_completion_desc(&ep->super.txqp, ep->txcnt.ci + count);
-        ucs_arbiter_group_schedule(&iface->super.tx.arbiter,
-                                   &ep->super.arb_group);
+        uct_rc_ep_schedule_pending(&ep->super);
         uct_rc_verbs_update_tx_res(&iface->super, ep, count);
         ucs_arbiter_dispatch(&iface->super.tx.arbiter, 1, uct_rc_ep_process_pending,
                              NULL);

test/gtest/uct/ib/test_dc.cc

@@ -806,8 +806,14 @@ class test_dc_flow_control : public test_rc_flow_control {
         ucs_array_length(&iface->tx.dci_pool[0].stack) = 0;
         iface->tx.dci_pool[0].stack_top                = 0;
 
-        uint8_t pool_index;
-        for (pool_index = 0; pool_index < iface->tx.num_dci_pools;
+        dispatch_pending(e);
+    }
+
+    virtual void dispatch_pending(entity *e) {
+        uct_dc_mlx5_iface_t *iface = ucs_derived_of(e->iface(),
+                                                    uct_dc_mlx5_iface_t);
+
+        for (uint8_t pool_index = 0; pool_index < iface->tx.num_dci_pools;
              ++pool_index) {
             uct_dc_mlx5_iface_progress_pending(iface, pool_index);
         }

test/gtest/uct/ib/test_rc.cc

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 * Copyright (C) UT-Battelle, LLC. 2016. ALL RIGHTS RESERVED.
 * Copyright (C) ARM Ltd. 2016.All rights reserved.
 * See file LICENSE for terms.
@@ -150,6 +150,87 @@ UCS_TEST_P(test_rc, fence_am_short_consumed, "RC_FENCE=weak")
     EXPECT_EQ(rc_iface(m_e1)->tx.fi.fence_beat, fence_info->fence_beat);
 }
 
+void test_rc::test_sgl_reserve(bool add_pending)
+{
+    uct_iface_attr_v2_t attr = {};
+    const size_t length      = 64;
+    const size_t count       = 2;
+    ucs_status_t status;
+    uct_pending_req_t pend_req;
+
+    attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS |
+                      UCT_IFACE_ATTR_FIELD_MAX_PUT_SGL_ZCOPY_COUNT;
+    ASSERT_UCS_OK(uct_iface_query_v2(m_e1->iface(), &attr));
+    if (!(attr.cap.flags & UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY) ||
+        (attr.max_put_sgl_zcopy_count < count)) {
+        UCS_TEST_SKIP_R("put_sgl_zcopy is not supported");
+    }
+
+    mapped_buffer s0(length, 0ul, *m_e1), s1(length, 0ul, *m_e1);
+    mapped_buffer r0(length, 0ul, *m_e2), r1(length, 0ul, *m_e2);
+    void *buffers[]    = {s0.ptr(), s1.ptr()};
+    size_t lengths[]   = {length, length};
+    uct_mem_h memhs[]  = {s0.memh(), s1.memh()};
+    uint64_t raddrs[]  = {r0.addr(), r1.addr()};
+    uct_rkey_t rkeys[] = {r0.rkey(), r1.rkey()};
+
+    /* A single PUT SHORT consumes exactly one TXQP credit */
+    const int16_t initial = uct_rc_txqp_available(&rc_ep(m_e1)->txqp);
+    for (int16_t i = initial; i > (int16_t)(count - 1); --i) {
+        ASSERT_UCS_OK(uct_ep_put_short(m_e1->ep(0), s0.ptr(), 8, r0.addr(),
+                                       r0.rkey()));
+    }
+
+    ASSERT_EQ(uct_rc_txqp_available(&rc_ep(m_e1)->txqp), (int16_t)(count - 1));
+    ASSERT_GT(rc_iface(m_e1)->tx.cq_available, (int16_t)(count - 1));
+
+    status = uct_ep_put_sgl_zcopy(m_e1->ep(0), buffers, lengths, memhs, raddrs,
+                                  rkeys, NULL, NULL, count, NULL);
+    ASSERT_EQ(UCS_ERR_NO_RESOURCE, status);
+    EXPECT_GT(rc_ep(m_e1)->txqp_reserve, 0);
+
+    if (add_pending) {
+        pend_req.func = [](uct_pending_req_t*) -> ucs_status_t {
+            return UCS_ERR_NO_RESOURCE;
+        };
+
+        ASSERT_UCS_OK(uct_ep_pending_add(m_e1->ep(0), &pend_req, 0));
+    }
+
+    EXPECT_EQ(UCS_ERR_NO_RESOURCE,
+              uct_ep_put_sgl_zcopy(m_e1->ep(0), buffers, lengths, memhs,
+                                   raddrs, rkeys, NULL, NULL, count, NULL));
+    EXPECT_EQ(UCS_ERR_NO_RESOURCE,
+              uct_ep_am_short(m_e1->ep(0), 0, 0, NULL, 0));
+
+    if (add_pending) {
+        /* The pending request never completes, so purge it to drop the
+         * reserve */
+        uct_ep_pending_purge(m_e1->ep(0), NULL, NULL);
+    }
+
+    /* Without pending requests, the reserve is dropped on TX completion */
+    flush();
+
+    EXPECT_EQ(0, rc_ep(m_e1)->txqp_reserve);
+    EXPECT_EQ(0, rc_ep(m_e1)->cq_reserve);
+}
+
+UCS_TEST_SKIP_COND_P(test_rc, sgl_pending_am_order,
+                     !check_caps(UCT_IFACE_FLAG_AM_SHORT |
+                                 UCT_IFACE_FLAG_PENDING  |
+                                 UCT_IFACE_FLAG_PUT_SHORT))
+{
+    test_sgl_reserve(true);
+}
+
+UCS_TEST_SKIP_COND_P(test_rc, sgl_reserve_released_on_completion,
+                     !check_caps(UCT_IFACE_FLAG_AM_SHORT |
+                                 UCT_IFACE_FLAG_PUT_SHORT))
+{
+    test_sgl_reserve(false);
+}
+
 UCS_TEST_SKIP_COND_P(test_rc, relaxed_order_required_strong_fence,
                      GetParam()->tl_name != "rc_mlx5",
                      "IB_PCI_RELAXED_ORDERING=yes")
@@ -1353,12 +1434,10 @@ void test_rc_flow_control::test_pending_grant(int16_t wnd)
 
     wait_fc_hard_resend(m_e1);
 
-    /* Enable send capabilities of m_e2 and send short put message to force
-     * pending queue dispatch. Can't send AM message for that, because it may
-     * trigger reordering assert due to disable/enable entity hack. */
+    /* Enable send capabilities of m_e2 and dispatch the pending queue directly,
+     * so it does not depend on a send to generate a TX completion. */
     enable_entity(m_e2);
-    set_tx_moderation(m_e2, 0);
-    EXPECT_EQ(UCS_OK, uct_ep_put_short(m_e2->ep(0), NULL, 0, 0, 0));
+    dispatch_pending(m_e2);
 
     /* Check that m_e1 got grant */
     validate_grant(m_e1);

test/gtest/uct/ib/test_rc.h

@@ -50,6 +50,8 @@ class test_rc : public uct_test {
 
     void test_iface_ops(int cq_len);
 
+    void test_sgl_reserve(bool add_pending);
+
     static ucs_status_t am_dummy_handler(void *arg, void *data, size_t length,
                                          unsigned flags) {
         return UCS_OK;
@@ -83,6 +85,10 @@ class test_rc_flow_control : public test_rc {
         rc_iface(e)->tx.cq_available = cq_num;
     }
 
+    virtual void dispatch_pending(entity *e) {
+        uct_rc_iface_arbiter_dispatch(rc_iface(e));
+    }
+
     virtual void set_tx_moderation(entity *e, int val) {
         rc_iface(e)->config.tx_moderation = val;
     }