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

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

@@ -243,6 +243,7 @@ uct_rc_mlx5_base_ep_put_sgl_zcopy(uct_ep_h tl_ep, void * const *buffers,
                          "put_sgl_zcopy");
     }
 
+    UCT_RC_CHECK_PENDING(&ep->super);
     UCT_RC_CHECK_CQE_VALUE_RET(&iface->super, &ep->super,
                                UCS_ERR_NO_RESOURCE, count - 1);
 

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

@@ -147,10 +147,10 @@ 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;
+    self->in_pending = 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) {
@@ -316,7 +316,8 @@ ucs_status_t uct_rc_ep_pending_add(uct_ep_h tl_ep, uct_pending_req_t *n,
     uct_rc_iface_t *iface = ucs_derived_of(tl_ep->iface, uct_rc_iface_t);
     uct_rc_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_ep_t);
 
-    if (uct_rc_ep_have_tx_cqe_avail(ep) &&
+    if (!UCT_RC_EP_MUST_BLOCK_SEND(ep) &&
+        uct_rc_ep_have_tx_cqe_avail(ep) &&
         uct_rc_ep_has_tx_resources(ep) &&
         uct_rc_iface_has_tx_resources(iface)) {
         return UCS_ERR_BUSY;
@@ -345,10 +346,11 @@ ucs_arbiter_cb_result_t uct_rc_ep_process_pending(ucs_arbiter_t *arbiter,
     uct_rc_iface_t *iface  = ucs_derived_of(ep->super.super.iface, uct_rc_iface_t);
     ucs_status_t status;
 
-    status = uct_rc_iface_invoke_pending_cb(iface, req);
+    ep->in_pending = 1;
+    status         = uct_rc_iface_invoke_pending_cb(iface, req);
+    ep->in_pending = 0;
     if (status == UCS_OK) {
-        ep->txqp_reserve = 0;
-        ep->cq_reserve   = 0;
+        uct_rc_ep_reset_reserve(ep);
         return UCS_ARBITER_CB_RESULT_REMOVE_ELEM;
     } else if (status == UCS_INPROGRESS) {
         return UCS_ARBITER_CB_RESULT_NEXT_GROUP;
@@ -434,6 +436,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

@@ -90,10 +90,25 @@ enum {
 #define UCT_RC_DEFINE_ATOMIC_HANDLER_FUNC_NAME(_num_bits, _is_be) \
     uct_rc_ep_atomic_handler_##_num_bits##_be##_is_be
 
+#define UCT_RC_EP_MUST_BLOCK_SEND(_ep) \
+    (!(_ep)->in_pending && \
+     (!ucs_arbiter_group_is_empty(&(_ep)->arb_group) || \
+      (((_ep)->cq_reserve | (_ep)->txqp_reserve) != 0)))
+
+#define UCT_RC_CHECK_PENDING_RET(_ep, _ret) \
+    if (ucs_unlikely(UCT_RC_EP_MUST_BLOCK_SEND(_ep))) { \
+        UCS_STATS_UPDATE_COUNTER((_ep)->super.stats, UCT_EP_STAT_NO_RES, 1); \
+        return _ret; \
+    }
+
+#define UCT_RC_CHECK_PENDING(_ep) \
+    UCT_RC_CHECK_PENDING_RET(_ep, UCS_ERR_NO_RESOURCE)
+
 /*
  * Check for send resources
  */
 #define UCT_RC_CHECK_CQE_RET(_iface, _ep, _ret) \
+    UCT_RC_CHECK_PENDING_RET(_ep, _ret) \
     /* tx_moderation == 0 for TLs which don't support it */ \
     if (ucs_unlikely((_iface)->tx.cq_available <= \
         (signed)(_iface)->config.tx_moderation)) { \
@@ -201,8 +216,7 @@ enum {
 
 #define UCT_RC_CHECK_RES_AND_FC(_iface, _ep, _am_id) \
     UCT_RC_CHECK_RES(_iface, _ep) \
-    UCT_RC_CHECK_FC(_iface, _ep, _am_id) \
-    uct_rc_iface_check_pending(_iface, &(_ep)->arb_group);
+    UCT_RC_CHECK_FC(_iface, _ep, _am_id)
 
 /* this is a common type for all rc and dc transports */
 struct uct_rc_txqp {
@@ -233,6 +247,7 @@ struct uct_rc_ep {
     uint16_t            cq_reserve;
     uint8_t             path_index;
     uint8_t             flags;
+    uint8_t             in_pending;
 };
 
 
@@ -354,6 +369,12 @@ 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 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);

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.
@@ -9,6 +9,10 @@
 #include <uct/ib/rc/verbs/rc_verbs.h>
 #include <uct/test_peer_failure.h>
 
+extern "C" {
+#include <uct/api/v2/uct_v2.h>
+}
+
 #ifdef HAVE_MLX5_DV
 extern "C" {
 #include <uct/ib/mlx5/rc/rc_mlx5_common.h>
@@ -142,6 +146,62 @@ 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);
 }
 
+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))
+{
+    uct_iface_attr_v2_t attr = {};
+    const size_t length      = 64;
+    const size_t count       = 2;
+    unsigned n;
+    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()};
+
+    n = 0;
+    while (uct_rc_txqp_available(&rc_ep(m_e1)->txqp) > (int16_t)(count - 1)) {
+        ASSERT_LT(n++, 1024u);
+        ASSERT_UCS_OK(uct_ep_put_short(m_e1->ep(0), s0.ptr(), 8, r0.addr(),
+                                       r0.rkey()));
+    }
+
+    ASSERT_GT(uct_rc_txqp_available(&rc_ep(m_e1)->txqp), 0);
+    ASSERT_LT(uct_rc_txqp_available(&rc_ep(m_e1)->txqp), (int16_t)count);
+    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);
+
+    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_am_short(m_e1->ep(0), 0, 0, NULL, 0));
+
+    uct_ep_pending_purge(m_e1->ep(0), NULL, NULL);
+    flush();
+}
+
 UCT_INSTANTIATE_RC_TEST_CASE(test_rc)