From f6010c77bdbd841ed30c00856c4ba257c2c27a8e Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Thu, 20 Aug 2026 12:03:54 +0300 Subject: [PATCH 01/13] UCT/IB/MLX5: impl token query --- src/uct/ib/mlx5/dv/ib_mlx5_dv.c | 5 +- src/uct/ib/mlx5/ib_mlx5.h | 3 + src/uct/ib/mlx5/rc/rc_mlx5.h | 8 +++ src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 36 ++++++++-- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 96 ++++++++++++++++++++------ test/gtest/uct/ib/test_ib_mlx5_ext.cc | 98 +++++++++++++++++++++++---- 6 files changed, 204 insertions(+), 42 deletions(-) diff --git a/src/uct/ib/mlx5/dv/ib_mlx5_dv.c b/src/uct/ib/mlx5/dv/ib_mlx5_dv.c index 6e90f34986c..19072bb1015 100644 --- a/src/uct/ib/mlx5/dv/ib_mlx5_dv.c +++ b/src/uct/ib/mlx5/dv/ib_mlx5_dv.c @@ -364,9 +364,8 @@ ucs_status_t uct_ib_mlx5_devx_modify_qp(uct_ib_mlx5_qp_t *qp, return UCS_OK; } -static ucs_status_t -uct_ib_mlx5_devx_query_qp(uct_ib_mlx5_qp_t *qp, void *in, size_t inlen, - void *out, size_t outlen) +ucs_status_t uct_ib_mlx5_devx_query_qp(uct_ib_mlx5_qp_t *qp, void *in, + size_t inlen, void *out, size_t outlen) { int ret; diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index 0e013b798ab..22aed800ba0 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -996,6 +996,9 @@ ucs_status_t uct_ib_mlx5_devx_modify_qp(uct_ib_mlx5_qp_t *qp, const void *in, size_t inlen, void *out, size_t outlen); +ucs_status_t uct_ib_mlx5_devx_query_qp(uct_ib_mlx5_qp_t *qp, void *in, + size_t inlen, void *out, size_t outlen); + ucs_status_t uct_ib_mlx5_devx_modify_qp_state(uct_ib_mlx5_qp_t *qp, enum ibv_qp_state state); diff --git a/src/uct/ib/mlx5/rc/rc_mlx5.h b/src/uct/ib/mlx5/rc/rc_mlx5.h index f40be42c0c5..1b01ddee663 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5.h +++ b/src/uct/ib/mlx5/rc/rc_mlx5.h @@ -54,6 +54,14 @@ typedef struct uct_rc_mlx5_base_ep { } tx; } uct_rc_mlx5_base_ep_t; +typedef struct { + uint32_t remote_qpn; +} uct_rc_mlx5_tx_token_t; + +typedef struct { + uint32_t receiver_next_psn; +} uct_rc_mlx5_rx_token_t; + /** * Generic RC remote endpoint with HW TM support diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index 95e48a643ba..5bca1f29c83 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -25,9 +25,33 @@ #include "rc_mlx5.inl" +static ucs_status_t UCS_F_ALWAYS_INLINE +uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, void *token) +{ +#if HAVE_DEVX + char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; + char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; + uct_rc_mlx5_tx_token_t *tx_token = token; + ucs_status_t status; + void *qpc; + + status = uct_ib_mlx5_devx_query_qp(&ep->tx.wq.super, in, sizeof(in), out, + sizeof(out)); + if (status != UCS_OK) { + return status; + } + + qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); + tx_token->remote_qpn = UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn); + return UCS_OK; +#else + return UCS_ERR_UNSUPPORTED; +#endif +} + ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) { - uct_ib_mlx5_ext_ep_query_attr_t attr = {}; + uct_rc_mlx5_base_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t); ucs_status_t status; if (ep_attr->field_mask & (UCT_EP_ATTR_FIELD_LOCAL_SOCKADDR | @@ -36,11 +60,15 @@ ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) } if (ep_attr->field_mask & UCT_EP_ATTR_FIELD_TX_TOKEN) { - attr.field_mask |= UCT_IB_MLX5_EXT_EP_QUERY_ATTR_FIELD_TX_TOKEN; - attr.tx_token = ep_attr->tx_token; + if (ep_attr->tx_token == NULL) { + ucs_error("rc mlx5: tx token is NULL"); + return UCS_ERR_INVALID_PARAM; + } - status = uct_ib_mlx5_ext_ep_query(tl_ep, &attr); + status = uct_rc_mlx5_ep_query_tx_token(ep, ep_attr->tx_token); if (status != UCS_OK) { + ucs_error("rc mlx5: failed to query tx token: %s", + ucs_status_string(status)); return status; } } diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 596b1fca239..0c7346e24e3 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1040,26 +1040,46 @@ static void uct_rc_mlx5_iface_fill_ext_query_attr( if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) { ext_attr->field_mask |= UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS; } +} - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH) { - ext_attr->field_mask |= - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN_LEN; - } +static ucs_status_t +uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, + uct_iface_attr_v2_t *iface_attr) +{ +#if HAVE_DEVX + char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; + char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; + const uct_rc_mlx5_tx_token_t *tx_token; + uct_rc_mlx5_rx_token_t *rx_token; + uct_rc_iface_t *iface; + uct_rc_mlx5_base_ep_t *ep; + uct_rc_ep_t *rc_ep; + ucs_status_t status; + void *qpc; - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH) { - ext_attr->field_mask |= - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN_LEN; + tx_token = iface_attr->tx_token; + iface = ucs_derived_of(tl_iface, uct_rc_iface_t); + rc_ep = uct_rc_iface_lookup_ep(iface, tx_token->remote_qpn); + if (rc_ep == NULL) { + ucs_error("rc mlx5: no ep for QPN %u", tx_token->remote_qpn); + return UCS_ERR_INVALID_PARAM; } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN) { - ext_attr->field_mask |= UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN; - ext_attr->tx_token = iface_attr->tx_token; + ep = ucs_derived_of(rc_ep, uct_rc_mlx5_base_ep_t); + status = uct_ib_mlx5_devx_query_qp(&ep->tx.wq.super, in, sizeof(in), out, + sizeof(out)); + if (status != UCS_OK) { + ucs_error("rc mlx5: failed to query qp: %s", ucs_status_string(status)); + return status; } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { - ext_attr->field_mask |= UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN; - ext_attr->rx_token = iface_attr->rx_token; - } + qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); + rx_token = iface_attr->rx_token; + rx_token->receiver_next_psn = UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn); + return UCS_OK; +#else + return UCS_ERR_UNSUPPORTED; +#endif } static ucs_status_t @@ -1073,6 +1093,46 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, uct_iface_query_v2_init(tl_iface, iface_attr); +#if HAVE_DEVX + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { + if (!(iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN)) { + ucs_error("rc mlx5: tx token is required to query rx token"); + return UCS_ERR_INVALID_PARAM; + } + + if ((iface_attr->rx_token == NULL) || (iface_attr->tx_token == NULL)) { + ucs_error("rc mlx5: tx token or rx token is NULL"); + return UCS_ERR_INVALID_PARAM; + } + + status = uct_rc_mlx5_iface_query_rx_token(tl_iface, iface_attr); + if (status != UCS_OK) { + ucs_error("rc mlx5: failed to query rx token: %s", + ucs_status_string(status)); + return status; + } + } + + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) { + iface_attr->cap.flags |= UCT_IFACE_FLAG_V2_QUERY_TOKEN; + } + + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH) { + iface_attr->tx_token_length = sizeof(uct_rc_mlx5_tx_token_t); + } + + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH) { + iface_attr->rx_token_length = sizeof(uct_rc_mlx5_rx_token_t); + } +#else + if (iface_attr->field_mask & + (UCT_IFACE_ATTR_FIELD_TX_TOKEN | UCT_IFACE_ATTR_FIELD_RX_TOKEN | + UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | + UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH)) { + return UCS_ERR_UNSUPPORTED; + } +#endif + if (iface_attr->field_mask & (UCT_IFACE_ATTR_FIELD_CAP_FLAGS | UCT_IFACE_ATTR_FIELD_MAX_PUT_SGL_ZCOPY_COUNT)) { @@ -1108,14 +1168,6 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, iface_attr->cap.flags |= ext_attr.cap.flags; } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH) { - iface_attr->tx_token_length = ext_attr.tx_token_len; - } - - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH) { - iface_attr->rx_token_length = ext_attr.rx_token_len; - } - return UCS_OK; } diff --git a/test/gtest/uct/ib/test_ib_mlx5_ext.cc b/test/gtest/uct/ib/test_ib_mlx5_ext.cc index 80c12a69318..8ec48890199 100644 --- a/test/gtest/uct/ib/test_ib_mlx5_ext.cc +++ b/test/gtest/uct/ib/test_ib_mlx5_ext.cc @@ -9,6 +9,7 @@ extern "C" { #include #include +#include extern ucs_list_link_t uct_ib_mlx5_ext_plugins; } @@ -30,6 +31,34 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { } protected: + static ucs_status_t am_counter(void *arg, void*, size_t, unsigned) + { + ++*static_cast(arg); + return UCS_OK; + } + + static void query_tx_token(uct_ep_h ep, uct_rc_mlx5_ft_tx_token_t *tx_token) + { + uct_ep_attr_t attr = {}; + + attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN; + attr.tx_token = tx_token; + ASSERT_UCS_OK(uct_ep_query(ep, &attr)); + } + + static void query_rx_token(uct_iface_h iface, + const uct_rc_mlx5_ft_tx_token_t *tx_token, + uct_rc_mlx5_ft_rx_token_t *rx_token) + { + uct_iface_attr_v2_t attr = {}; + + attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN | + UCT_IFACE_ATTR_FIELD_RX_TOKEN; + attr.tx_token = tx_token; + attr.rx_token = rx_token; + ASSERT_UCS_OK(uct_iface_query_v2(iface, &attr)); + } + static uint64_t tx_token() { return 0x1234; @@ -140,12 +169,10 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { UCS_TEST_P(test_uct_ib_mlx5_ext_rc, iface_query) { - uint64_t tx_token_value = tx_token(); - uint64_t rx_token_value = 0; + uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; + uct_rc_mlx5_ft_rx_token_t rx_token_value = {}; uct_iface_attr_v2_t attr = {}; - register_plugin("token", iface_query, ep_query, purge); - { scoped_log_handler wrap_err(wrap_errors_logger); attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN; @@ -162,19 +189,19 @@ UCS_TEST_P(test_uct_ib_mlx5_ext_rc, iface_query) attr.tx_token = &tx_token_value; attr.rx_token = &rx_token_value; - ASSERT_UCS_OK(uct_iface_query_v2(m_e1->iface(), &attr)); + query_tx_token(m_e1->ep(0), &tx_token_value); + ASSERT_UCS_OK(uct_iface_query_v2(m_e2->iface(), &attr)); EXPECT_TRUE(attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); - EXPECT_EQ(sizeof(uint64_t), attr.tx_token_length); - EXPECT_EQ(sizeof(uint64_t), attr.rx_token_length); - EXPECT_EQ(rx_token(), rx_token_value); + EXPECT_EQ(sizeof(tx_token_value), attr.tx_token_length); + EXPECT_EQ(sizeof(rx_token_value), attr.rx_token_length); } UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_query) { - uint64_t tx_token_value = 0; - uct_ep_attr_t attr = {}; - - register_plugin("token", iface_query, ep_query, purge); + uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; + uct_rc_mlx5_base_ep_t *remote_ep = ucs_derived_of(m_e2->ep(0), + uct_rc_mlx5_base_ep_t); + uct_ep_attr_t attr = {}; { scoped_log_handler wrap_err(wrap_errors_logger); @@ -184,7 +211,52 @@ UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_query) attr.tx_token = &tx_token_value; ASSERT_UCS_OK(uct_ep_query(m_e1->ep(0), &attr)); - EXPECT_EQ(tx_token(), tx_token_value); + EXPECT_EQ(remote_ep->tx.wq.super.qp_num, tx_token_value.remote_qpn); +} + +UCS_TEST_P(test_uct_ib_mlx5_ext_rc, token_query_psn) +{ + uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; + uct_rc_mlx5_ft_rx_token_t rx_token_pre = {}; + uct_rc_mlx5_ft_rx_token_t rx_token_post = {}; + uct_rc_mlx5_base_ep_t *sender_ep = ucs_derived_of(m_e1->ep(0), + uct_rc_mlx5_base_ep_t); + const uint32_t num_messages = UCS_PTR_BYTE_DIFF(sender_ep->tx.wq.qstart, + sender_ep->tx.wq.qend) / + MLX5_SEND_WQE_BB + + 1; + uint32_t rx_count = 0; + + ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_counter, + &rx_count, 0)); + query_tx_token(m_e1->ep(0), &tx_token_value); + query_rx_token(m_e2->iface(), &tx_token_value, &rx_token_pre); + + EXPECT_EQ(sender_ep->tx.wq.next_first_psn, rx_token_pre.receiver_next_psn) + << "initial HW receive PSN must match SW send PSN"; + + for (uint32_t i = 0; i < num_messages; ++i) { + ucs_status_t status; + + do { + status = uct_ep_am_short(m_e1->ep(0), 0, i, NULL, 0); + if (status == UCS_ERR_NO_RESOURCE) { + progress(); + } + } while (status == UCS_ERR_NO_RESOURCE); + ASSERT_UCS_OK(status); + } + + while (rx_count < num_messages) { + progress(); + } + + query_rx_token(m_e2->iface(), &tx_token_value, &rx_token_post); + EXPECT_EQ((rx_token_pre.receiver_next_psn + num_messages) & + UCS_MASK(UCT_IB_MLX5_PSN_BITS), + rx_token_post.receiver_next_psn); + EXPECT_EQ(sender_ep->tx.wq.next_first_psn, rx_token_post.receiver_next_psn) + << "updated HW receive PSN must match SW send PSN"; } UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_outstanding_purge) From a44fd1fc043d69021e72483c9098165da787c107 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Thu, 20 Aug 2026 12:25:33 +0300 Subject: [PATCH 02/13] GTEST/UCT/IB: add gtest --- test/gtest/uct/ib/test_ib_mlx5_ext.cc | 121 -------------------------- test/gtest/uct/ib/test_rc.cc | 66 ++++++++++++++ 2 files changed, 66 insertions(+), 121 deletions(-) diff --git a/test/gtest/uct/ib/test_ib_mlx5_ext.cc b/test/gtest/uct/ib/test_ib_mlx5_ext.cc index 8ec48890199..67528ff7b9c 100644 --- a/test/gtest/uct/ib/test_ib_mlx5_ext.cc +++ b/test/gtest/uct/ib/test_ib_mlx5_ext.cc @@ -9,7 +9,6 @@ extern "C" { #include #include -#include extern ucs_list_link_t uct_ib_mlx5_ext_plugins; } @@ -31,34 +30,6 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { } protected: - static ucs_status_t am_counter(void *arg, void*, size_t, unsigned) - { - ++*static_cast(arg); - return UCS_OK; - } - - static void query_tx_token(uct_ep_h ep, uct_rc_mlx5_ft_tx_token_t *tx_token) - { - uct_ep_attr_t attr = {}; - - attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN; - attr.tx_token = tx_token; - ASSERT_UCS_OK(uct_ep_query(ep, &attr)); - } - - static void query_rx_token(uct_iface_h iface, - const uct_rc_mlx5_ft_tx_token_t *tx_token, - uct_rc_mlx5_ft_rx_token_t *rx_token) - { - uct_iface_attr_v2_t attr = {}; - - attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN | - UCT_IFACE_ATTR_FIELD_RX_TOKEN; - attr.tx_token = tx_token; - attr.rx_token = rx_token; - ASSERT_UCS_OK(uct_iface_query_v2(iface, &attr)); - } - static uint64_t tx_token() { return 0x1234; @@ -167,98 +138,6 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { }; -UCS_TEST_P(test_uct_ib_mlx5_ext_rc, iface_query) -{ - uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; - uct_rc_mlx5_ft_rx_token_t rx_token_value = {}; - uct_iface_attr_v2_t attr = {}; - - { - scoped_log_handler wrap_err(wrap_errors_logger); - attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN; - attr.tx_token = &tx_token_value; - EXPECT_EQ(UCS_ERR_INVALID_PARAM, - uct_iface_query_v2(m_e1->iface(), &attr)); - } - - attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS | - UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | - UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH | - UCT_IFACE_ATTR_FIELD_TX_TOKEN | - UCT_IFACE_ATTR_FIELD_RX_TOKEN; - attr.tx_token = &tx_token_value; - attr.rx_token = &rx_token_value; - - query_tx_token(m_e1->ep(0), &tx_token_value); - ASSERT_UCS_OK(uct_iface_query_v2(m_e2->iface(), &attr)); - EXPECT_TRUE(attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); - EXPECT_EQ(sizeof(tx_token_value), attr.tx_token_length); - EXPECT_EQ(sizeof(rx_token_value), attr.rx_token_length); -} - -UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_query) -{ - uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; - uct_rc_mlx5_base_ep_t *remote_ep = ucs_derived_of(m_e2->ep(0), - uct_rc_mlx5_base_ep_t); - uct_ep_attr_t attr = {}; - - { - scoped_log_handler wrap_err(wrap_errors_logger); - attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN; - EXPECT_EQ(UCS_ERR_INVALID_PARAM, uct_ep_query(m_e1->ep(0), &attr)); - } - - attr.tx_token = &tx_token_value; - ASSERT_UCS_OK(uct_ep_query(m_e1->ep(0), &attr)); - EXPECT_EQ(remote_ep->tx.wq.super.qp_num, tx_token_value.remote_qpn); -} - -UCS_TEST_P(test_uct_ib_mlx5_ext_rc, token_query_psn) -{ - uct_rc_mlx5_ft_tx_token_t tx_token_value = {}; - uct_rc_mlx5_ft_rx_token_t rx_token_pre = {}; - uct_rc_mlx5_ft_rx_token_t rx_token_post = {}; - uct_rc_mlx5_base_ep_t *sender_ep = ucs_derived_of(m_e1->ep(0), - uct_rc_mlx5_base_ep_t); - const uint32_t num_messages = UCS_PTR_BYTE_DIFF(sender_ep->tx.wq.qstart, - sender_ep->tx.wq.qend) / - MLX5_SEND_WQE_BB + - 1; - uint32_t rx_count = 0; - - ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_counter, - &rx_count, 0)); - query_tx_token(m_e1->ep(0), &tx_token_value); - query_rx_token(m_e2->iface(), &tx_token_value, &rx_token_pre); - - EXPECT_EQ(sender_ep->tx.wq.next_first_psn, rx_token_pre.receiver_next_psn) - << "initial HW receive PSN must match SW send PSN"; - - for (uint32_t i = 0; i < num_messages; ++i) { - ucs_status_t status; - - do { - status = uct_ep_am_short(m_e1->ep(0), 0, i, NULL, 0); - if (status == UCS_ERR_NO_RESOURCE) { - progress(); - } - } while (status == UCS_ERR_NO_RESOURCE); - ASSERT_UCS_OK(status); - } - - while (rx_count < num_messages) { - progress(); - } - - query_rx_token(m_e2->iface(), &tx_token_value, &rx_token_post); - EXPECT_EQ((rx_token_pre.receiver_next_psn + num_messages) & - UCS_MASK(UCT_IB_MLX5_PSN_BITS), - rx_token_post.receiver_next_psn); - EXPECT_EQ(sender_ep->tx.wq.next_first_psn, rx_token_post.receiver_next_psn) - << "updated HW receive PSN must match SW send PSN"; -} - UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_outstanding_purge) { uint64_t rx_token_value = rx_token(); diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index e9946df00c3..69ff4857bf8 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1323,6 +1323,72 @@ UCT_INSTANTIATE_RC_TEST_CASE(test_rc_keepalive) #ifdef HAVE_MLX5_DV +class test_rc_mlx5_token_query : public test_rc { +protected: + static ucs_status_t am_handler(void *arg, void*, size_t, unsigned) + { + ++*static_cast(arg); + return UCS_OK; + } + + static void query_tx_token(uct_ep_h ep, uct_rc_mlx5_tx_token_t *tx_token) + { + uct_ep_attr_t attr = {}; + + attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN; + attr.tx_token = tx_token; + ASSERT_UCS_OK(uct_ep_query(ep, &attr)); + } + + static void query_rx_token(uct_iface_h iface, + const uct_rc_mlx5_tx_token_t *tx_token, + uct_rc_mlx5_rx_token_t *rx_token) + { + uct_iface_attr_v2_t attr = {}; + + attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN | + UCT_IFACE_ATTR_FIELD_RX_TOKEN; + attr.tx_token = tx_token; + attr.rx_token = rx_token; + ASSERT_UCS_OK(uct_iface_query_v2(iface, &attr)); + } +}; + +UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, + !check_caps(UCT_IFACE_FLAG_AM_SHORT)) +{ + uct_rc_mlx5_tx_token_t tx_token = {}; + uct_rc_mlx5_rx_token_t rx_token = {}; + const uint32_t num_message = 10; + uint32_t ep1_rx_count = 0; + uint32_t ep2_rx_count = 0; + + ASSERT_UCS_OK(uct_iface_set_am_handler(m_e1->iface(), 0, am_handler, + &ep1_rx_count, 0)); + ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_handler, + &ep2_rx_count, 0)); + + for (uint32_t i = 0; i < num_message; ++i) { + ucs_status_t status; + + UCT_TEST_CALL_AND_TRY_AGAIN(uct_ep_am_short(m_e1->ep(0), 0, i, NULL, 0), + status); + ASSERT_UCS_OK(status); + UCT_TEST_CALL_AND_TRY_AGAIN(uct_ep_am_short(m_e2->ep(0), 0, i, NULL, 0), + status); + ASSERT_UCS_OK(status); + } + + wait_for_value(&ep1_rx_count, num_message, true); + wait_for_value(&ep2_rx_count, num_message, true); + + query_tx_token(m_e1->ep(0), &tx_token); + query_rx_token(m_e2->iface(), &tx_token, &rx_token); + EXPECT_EQ(num_message, rx_token.receiver_next_psn - 1); +} + +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_token_query, rc_mlx5) + class test_rc_srq : public test_rc { public: test_rc_srq() : m_buf8b(NULL), m_buf8k(NULL) From d8356a8fa480bb989aedb1f5bc9f204d2f20d8db Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Thu, 20 Aug 2026 14:46:26 +0300 Subject: [PATCH 03/13] GTEST/UCT: update token test --- test/gtest/uct/v2/test_uct_query.cc | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/test/gtest/uct/v2/test_uct_query.cc b/test/gtest/uct/v2/test_uct_query.cc index 2c81f924876..46efb78574c 100644 --- a/test/gtest/uct/v2/test_uct_query.cc +++ b/test/gtest/uct/v2/test_uct_query.cc @@ -4,6 +4,10 @@ * See file LICENSE for terms. */ +#ifdef HAVE_CONFIG_H +# include "config.h" +#endif + #include #include @@ -113,12 +117,19 @@ UCS_TEST_P(test_uct_query, query_perf) UCS_TEST_P(test_uct_query, query_token_support) { uct_iface_attr_v2_t attr = {}; + uint64_t iface_cap_flags = 0; attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS; attr.cap.flags = UINT64_MAX; ASSERT_UCS_OK(uct_iface_query_v2(get_iface(), &attr)); - EXPECT_EQ(0ul, attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); + +#if HAVE_DEVX + if (has_transport("rc_mlx5")) { + iface_cap_flags = UCT_IFACE_FLAG_V2_QUERY_TOKEN; + } +#endif + EXPECT_EQ(iface_cap_flags, attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); } UCT_INSTANTIATE_TEST_CASE(test_uct_query) From a895bef454d795e8f905f85b7585870d2461239f Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Fri, 21 Aug 2026 04:40:28 +0300 Subject: [PATCH 04/13] GTEST/UCT: fix psn match --- test/gtest/uct/ib/test_rc.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 69ff4857bf8..d62f7561803 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1384,7 +1384,7 @@ UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, query_tx_token(m_e1->ep(0), &tx_token); query_rx_token(m_e2->iface(), &tx_token, &rx_token); - EXPECT_EQ(num_message, rx_token.receiver_next_psn - 1); + EXPECT_EQ(num_message, rx_token.receiver_next_psn); } _UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_token_query, rc_mlx5) From 9b1ba9197607a63a7aae43c7d9271029e0db52c8 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Fri, 21 Aug 2026 07:48:43 +0300 Subject: [PATCH 05/13] UCT/IB/MLX5: fix build and update test --- src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 38 ++++++++++++++++-------------- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 32 ++++++++++++------------- test/gtest/uct/ib/test_rc.cc | 29 +++++++++++------------ 3 files changed, 49 insertions(+), 50 deletions(-) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index 5bca1f29c83..c3de49a0574 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -25,34 +25,42 @@ #include "rc_mlx5.inl" +#if HAVE_DEVX static ucs_status_t UCS_F_ALWAYS_INLINE -uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, void *token) +uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) { -#if HAVE_DEVX char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; - uct_rc_mlx5_tx_token_t *tx_token = token; + uct_rc_mlx5_tx_token_t *tx_token; ucs_status_t status; void *qpc; + if (ep_attr->tx_token == NULL) { + ucs_error("rc mlx5: tx token is NULL"); + return UCS_ERR_INVALID_PARAM; + } + + tx_token = ep_attr->tx_token; status = uct_ib_mlx5_devx_query_qp(&ep->tx.wq.super, in, sizeof(in), out, sizeof(out)); if (status != UCS_OK) { + ucs_error("rc mlx5: ep %p failed to query tx token: %s", ep, + ucs_status_string(status)); return status; } qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); tx_token->remote_qpn = UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn); + return UCS_OK; -#else - return UCS_ERR_UNSUPPORTED; -#endif } +#endif ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) { +#if HAVE_DEVX uct_rc_mlx5_base_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t); - ucs_status_t status; +#endif if (ep_attr->field_mask & (UCT_EP_ATTR_FIELD_LOCAL_SOCKADDR | UCT_EP_ATTR_FIELD_REMOTE_SOCKADDR)) { @@ -60,17 +68,11 @@ ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) } if (ep_attr->field_mask & UCT_EP_ATTR_FIELD_TX_TOKEN) { - if (ep_attr->tx_token == NULL) { - ucs_error("rc mlx5: tx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - status = uct_rc_mlx5_ep_query_tx_token(ep, ep_attr->tx_token); - if (status != UCS_OK) { - ucs_error("rc mlx5: failed to query tx token: %s", - ucs_status_string(status)); - return status; - } +#if HAVE_DEVX + return uct_rc_mlx5_ep_query_tx_token(ep, ep_attr); +#else + return UCS_ERR_UNSUPPORTED; +#endif } return UCS_OK; diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 0c7346e24e3..786bfe57ae1 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1042,11 +1042,11 @@ static void uct_rc_mlx5_iface_fill_ext_query_attr( } } +#if HAVE_DEVX static ucs_status_t uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, uct_iface_attr_v2_t *iface_attr) { -#if HAVE_DEVX char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; const uct_rc_mlx5_tx_token_t *tx_token; @@ -1057,6 +1057,16 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, ucs_status_t status; void *qpc; + if (!(iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN)) { + ucs_error("rc mlx5: tx token is required to query rx token"); + return UCS_ERR_INVALID_PARAM; + } + + if ((iface_attr->tx_token == NULL) || (iface_attr->rx_token == NULL)) { + ucs_error("rc mlx5: tx token or rx token is NULL"); + return UCS_ERR_INVALID_PARAM; + } + tx_token = iface_attr->tx_token; iface = ucs_derived_of(tl_iface, uct_rc_iface_t); rc_ep = uct_rc_iface_lookup_ep(iface, tx_token->remote_qpn); @@ -1069,18 +1079,18 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, status = uct_ib_mlx5_devx_query_qp(&ep->tx.wq.super, in, sizeof(in), out, sizeof(out)); if (status != UCS_OK) { - ucs_error("rc mlx5: failed to query qp: %s", ucs_status_string(status)); + ucs_error("rc mlx5: iface %p failed to query rx token: %s", tl_iface, + ucs_status_string(status)); return status; } qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); rx_token = iface_attr->rx_token; rx_token->receiver_next_psn = UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn); + return UCS_OK; -#else - return UCS_ERR_UNSUPPORTED; -#endif } +#endif static ucs_status_t uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, @@ -1095,20 +1105,8 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, #if HAVE_DEVX if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { - if (!(iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN)) { - ucs_error("rc mlx5: tx token is required to query rx token"); - return UCS_ERR_INVALID_PARAM; - } - - if ((iface_attr->rx_token == NULL) || (iface_attr->tx_token == NULL)) { - ucs_error("rc mlx5: tx token or rx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - status = uct_rc_mlx5_iface_query_rx_token(tl_iface, iface_attr); if (status != UCS_OK) { - ucs_error("rc mlx5: failed to query rx token: %s", - ucs_status_string(status)); return status; } } diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index d62f7561803..703afec86fa 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1325,6 +1325,8 @@ UCT_INSTANTIATE_RC_TEST_CASE(test_rc_keepalive) class test_rc_mlx5_token_query : public test_rc { protected: + static constexpr uint32_t NUM_MESSAGES = 10; + static ucs_status_t am_handler(void *arg, void*, size_t, unsigned) { ++*static_cast(arg); @@ -1354,37 +1356,34 @@ class test_rc_mlx5_token_query : public test_rc { } }; +constexpr uint32_t test_rc_mlx5_token_query::NUM_MESSAGES; + UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, !check_caps(UCT_IFACE_FLAG_AM_SHORT)) { uct_rc_mlx5_tx_token_t tx_token = {}; uct_rc_mlx5_rx_token_t rx_token = {}; - const uint32_t num_message = 10; - uint32_t ep1_rx_count = 0; - uint32_t ep2_rx_count = 0; + uct_rc_mlx5_base_ep_t *e2_ep; + uint32_t rx_count = 0; + ucs_status_t status; - ASSERT_UCS_OK(uct_iface_set_am_handler(m_e1->iface(), 0, am_handler, - &ep1_rx_count, 0)); ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_handler, - &ep2_rx_count, 0)); - - for (uint32_t i = 0; i < num_message; ++i) { - ucs_status_t status; + &rx_count, 0)); + for (uint32_t i = 0; i < NUM_MESSAGES; ++i) { UCT_TEST_CALL_AND_TRY_AGAIN(uct_ep_am_short(m_e1->ep(0), 0, i, NULL, 0), status); ASSERT_UCS_OK(status); - UCT_TEST_CALL_AND_TRY_AGAIN(uct_ep_am_short(m_e2->ep(0), 0, i, NULL, 0), - status); - ASSERT_UCS_OK(status); } - wait_for_value(&ep1_rx_count, num_message, true); - wait_for_value(&ep2_rx_count, num_message, true); + wait_for_value(&rx_count, NUM_MESSAGES, true); query_tx_token(m_e1->ep(0), &tx_token); query_rx_token(m_e2->iface(), &tx_token, &rx_token); - EXPECT_EQ(num_message, rx_token.receiver_next_psn); + + e2_ep = ucs_derived_of(m_e2->ep(0), uct_rc_mlx5_base_ep_t); + EXPECT_EQ(e2_ep->tx.wq.super.qp_num, tx_token.remote_qpn); + EXPECT_EQ(NUM_MESSAGES, rx_token.receiver_next_psn); } _UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_token_query, rc_mlx5) From ea9be34c3202c4485de637573c1190f82f8c8191 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Fri, 21 Aug 2026 10:18:17 +0300 Subject: [PATCH 06/13] GTEST/UCT: code style --- test/gtest/uct/ib/test_rc.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 703afec86fa..2a23854af0f 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1363,8 +1363,8 @@ UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, { uct_rc_mlx5_tx_token_t tx_token = {}; uct_rc_mlx5_rx_token_t rx_token = {}; - uct_rc_mlx5_base_ep_t *e2_ep; uint32_t rx_count = 0; + uct_rc_mlx5_base_ep_t *e2_ep; ucs_status_t status; ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_handler, From 88929cb304795bbcc47bccdbb9cf48a396409f50 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Mon, 24 Aug 2026 13:46:45 +0300 Subject: [PATCH 07/13] UCT/IB/MLX5: address comments --- src/uct/ib/mlx5/rc/rc_mlx5.h | 9 ++------- src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 4 ++-- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 12 +++++++----- test/gtest/uct/ib/test_rc.cc | 10 +++++----- 4 files changed, 16 insertions(+), 19 deletions(-) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5.h b/src/uct/ib/mlx5/rc/rc_mlx5.h index 1b01ddee663..8bacf4985e6 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5.h +++ b/src/uct/ib/mlx5/rc/rc_mlx5.h @@ -54,13 +54,8 @@ typedef struct uct_rc_mlx5_base_ep { } tx; } uct_rc_mlx5_base_ep_t; -typedef struct { - uint32_t remote_qpn; -} uct_rc_mlx5_tx_token_t; - -typedef struct { - uint32_t receiver_next_psn; -} uct_rc_mlx5_rx_token_t; +typedef __be32 uct_rc_mlx5_tx_token_t; +typedef __be32 uct_rc_mlx5_rx_token_t; /** diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index c3de49a0574..f413acd3a6c 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -49,8 +49,8 @@ uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) return status; } - qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); - tx_token->remote_qpn = UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn); + qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); + *tx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn)); return UCS_OK; } diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 786bfe57ae1..5f826c4f75d 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1055,6 +1055,7 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, uct_rc_mlx5_base_ep_t *ep; uct_rc_ep_t *rc_ep; ucs_status_t status; + uint32_t remote_qpn; void *qpc; if (!(iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN)) { @@ -1067,11 +1068,12 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, return UCS_ERR_INVALID_PARAM; } - tx_token = iface_attr->tx_token; - iface = ucs_derived_of(tl_iface, uct_rc_iface_t); - rc_ep = uct_rc_iface_lookup_ep(iface, tx_token->remote_qpn); + tx_token = iface_attr->tx_token; + iface = ucs_derived_of(tl_iface, uct_rc_iface_t); + remote_qpn = be32toh(*tx_token); + rc_ep = uct_rc_iface_lookup_ep(iface, remote_qpn); if (rc_ep == NULL) { - ucs_error("rc mlx5: no ep for QPN %u", tx_token->remote_qpn); + ucs_error("rc mlx5: no ep for QPN %u", remote_qpn); return UCS_ERR_INVALID_PARAM; } @@ -1086,7 +1088,7 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); rx_token = iface_attr->rx_token; - rx_token->receiver_next_psn = UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn); + *rx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn)); return UCS_OK; } diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 2a23854af0f..7f7046bc661 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1327,9 +1327,9 @@ class test_rc_mlx5_token_query : public test_rc { protected: static constexpr uint32_t NUM_MESSAGES = 10; - static ucs_status_t am_handler(void *arg, void*, size_t, unsigned) + static ucs_status_t am_handler(void *rx_count, void*, size_t, unsigned) { - ++*static_cast(arg); + ++*static_cast(rx_count); return UCS_OK; } @@ -1363,7 +1363,7 @@ UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, { uct_rc_mlx5_tx_token_t tx_token = {}; uct_rc_mlx5_rx_token_t rx_token = {}; - uint32_t rx_count = 0; + uint32_t rx_count = 0; uct_rc_mlx5_base_ep_t *e2_ep; ucs_status_t status; @@ -1382,8 +1382,8 @@ UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, query_rx_token(m_e2->iface(), &tx_token, &rx_token); e2_ep = ucs_derived_of(m_e2->ep(0), uct_rc_mlx5_base_ep_t); - EXPECT_EQ(e2_ep->tx.wq.super.qp_num, tx_token.remote_qpn); - EXPECT_EQ(NUM_MESSAGES, rx_token.receiver_next_psn); + EXPECT_EQ(e2_ep->tx.wq.super.qp_num, be32toh(tx_token)); + EXPECT_EQ(NUM_MESSAGES, be32toh(rx_token)); } _UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_token_query, rc_mlx5) From 0bb751f34a65427f1c81e3d933a36d214c0a23f1 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Tue, 25 Aug 2026 09:59:16 +0300 Subject: [PATCH 08/13] UCT/IB/MLX5: address comments --- src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 11 ++++------- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 9 ++------- test/gtest/uct/ib/test_rc.cc | 6 ++++-- 3 files changed, 10 insertions(+), 16 deletions(-) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index f413acd3a6c..322cd0b7bc9 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -26,7 +26,7 @@ #if HAVE_DEVX -static ucs_status_t UCS_F_ALWAYS_INLINE +static ucs_status_t uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) { char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; @@ -40,7 +40,6 @@ uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) return UCS_ERR_INVALID_PARAM; } - tx_token = ep_attr->tx_token; status = uct_ib_mlx5_devx_query_qp(&ep->tx.wq.super, in, sizeof(in), out, sizeof(out)); if (status != UCS_OK) { @@ -50,6 +49,7 @@ uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) } qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); + tx_token = ep_attr->tx_token; *tx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn)); return UCS_OK; @@ -58,10 +58,6 @@ uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) { -#if HAVE_DEVX - uct_rc_mlx5_base_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t); -#endif - if (ep_attr->field_mask & (UCT_EP_ATTR_FIELD_LOCAL_SOCKADDR | UCT_EP_ATTR_FIELD_REMOTE_SOCKADDR)) { return UCS_ERR_UNSUPPORTED; @@ -69,7 +65,8 @@ ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) if (ep_attr->field_mask & UCT_EP_ATTR_FIELD_TX_TOKEN) { #if HAVE_DEVX - return uct_rc_mlx5_ep_query_tx_token(ep, ep_attr); + return uct_rc_mlx5_ep_query_tx_token( + ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t), ep_attr); #else return UCS_ERR_UNSUPPORTED; #endif diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 5f826c4f75d..8357eab94d5 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1063,11 +1063,6 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, return UCS_ERR_INVALID_PARAM; } - if ((iface_attr->tx_token == NULL) || (iface_attr->rx_token == NULL)) { - ucs_error("rc mlx5: tx token or rx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - tx_token = iface_attr->tx_token; iface = ucs_derived_of(tl_iface, uct_rc_iface_t); remote_qpn = be32toh(*tx_token); @@ -1086,8 +1081,8 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, return status; } - qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); - rx_token = iface_attr->rx_token; + qpc = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc); + rx_token = iface_attr->rx_token; *rx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn)); return UCS_OK; diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 7f7046bc661..d6e3e419ecd 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1327,9 +1327,11 @@ class test_rc_mlx5_token_query : public test_rc { protected: static constexpr uint32_t NUM_MESSAGES = 10; - static ucs_status_t am_handler(void *rx_count, void*, size_t, unsigned) + static ucs_status_t am_handler(void *arg, void*, size_t, unsigned) { - ++*static_cast(rx_count); + uint32_t *rx_count = static_cast(arg); + + ++*rx_count; return UCS_OK; } From c82cfe40bdedda107f08d9e15c5bf787efb59703 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Tue, 25 Aug 2026 12:57:32 +0300 Subject: [PATCH 09/13] UCT/IB/MLX5: fix coverity --- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 8357eab94d5..f52007335a2 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1058,6 +1058,10 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, uint32_t remote_qpn; void *qpc; + /* suppress coverity false-positive */ + ucs_assert((iface_attr->tx_token != NULL) && + (iface_attr->rx_token != NULL)); + if (!(iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN)) { ucs_error("rc mlx5: tx token is required to query rx token"); return UCS_ERR_INVALID_PARAM; From 483f0e21ce29ce90dcbfe170774198ba2090ad4c Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Tue, 25 Aug 2026 15:42:45 +0300 Subject: [PATCH 10/13] UCT/IB/MLX5: address comments --- src/uct/ib/mlx5/ib_mlx5_ext.c | 146 +------------------------- src/uct/ib/mlx5/ib_mlx5_ext.h | 84 +-------------- src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 1 - test/gtest/uct/ib/test_ib_mlx5_ext.cc | 52 +-------- 4 files changed, 8 insertions(+), 275 deletions(-) diff --git a/src/uct/ib/mlx5/ib_mlx5_ext.c b/src/uct/ib/mlx5/ib_mlx5_ext.c index 8160e98cc9f..643e0dbc21c 100644 --- a/src/uct/ib/mlx5/ib_mlx5_ext.c +++ b/src/uct/ib/mlx5/ib_mlx5_ext.c @@ -29,54 +29,6 @@ static int uct_ib_mlx5_ext_is_unsupported_op(const void *op) (op == (const void*)ucs_empty_function_return_unsupported); } -static ucs_status_t -uct_ib_mlx5_ext_iface_query_check_param(uct_ib_mlx5_ext_iface_query_attr_t *attr) -{ - if (attr == NULL) { - ucs_error("ib mlx5 ext: iface query attribute is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - if ((attr->field_mask & UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN) && - (attr->tx_token == NULL)) { - ucs_error("ib mlx5 ext: tx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - if ((attr->field_mask & UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN) && - (attr->rx_token == NULL)) { - ucs_error("ib mlx5 ext: rx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - if (!!(attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN) != - !!(attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN)) { - ucs_error("ib mlx5 ext: tx token and rx token must be set together"); - return UCS_ERR_INVALID_PARAM; - } - - return UCS_OK; -} - -static ucs_status_t -uct_ib_mlx5_ext_ep_query_check_param(uct_ib_mlx5_ext_ep_query_attr_t *attr) -{ - if (attr == NULL) { - ucs_error("ib mlx5 ext: ep query attribute is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - if ((attr->field_mask & UCT_IB_MLX5_EXT_EP_QUERY_ATTR_FIELD_TX_TOKEN) && - (attr->tx_token == NULL)) { - ucs_error("ib mlx5 ext: tx token is NULL"); - return UCS_ERR_INVALID_PARAM; - } - - return UCS_OK; -} - static uint64_t uct_ib_mlx5_ext_iface_query_cap_flags(uct_iface_h iface) { uint64_t cap_flags = 0; @@ -103,66 +55,13 @@ static uint64_t uct_ib_mlx5_ext_iface_query_cap_flags(uct_iface_h iface) return cap_flags; } -static uct_ib_mlx5_ext_plugin_t * -uct_ib_mlx5_ext_find_plugin(uct_iface_h iface, uint64_t cap_flags) -{ - uct_ib_mlx5_ext_iface_query_attr_t attr; - uct_ib_mlx5_ext_plugin_t *plugin; - ucs_status_t status; - - ucs_list_for_each(plugin, &uct_ib_mlx5_ext_plugins, list) { - /* Each plugin must support iface query. */ - ucs_assert(!(uct_ib_mlx5_ext_is_unsupported_op( - (const void*)plugin->ops.iface_query))); - - attr.field_mask = UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS; - status = plugin->ops.iface_query(iface, &attr); - if (status != UCS_OK) { - continue; - } - - if (!ucs_test_all_flags(attr.cap.flags, cap_flags)) { - continue; - } - - return plugin; - } - - ucs_error("ib mlx5 ext: no plugin found for cap flags: %" PRIu64, - cap_flags); - return NULL; -} - ucs_status_t uct_ib_mlx5_ext_iface_query(uct_iface_h iface, uct_ib_mlx5_ext_iface_query_attr_t *attr) { - const uint64_t token_mask = - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN_LEN | - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN_LEN | - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN | - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN; - uct_ib_mlx5_ext_plugin_t *plugin; - ucs_status_t status; - - status = uct_ib_mlx5_ext_iface_query_check_param(attr); - if (status != UCS_OK) { - return status; - } - - if (ucs_test_flags(attr->field_mask, token_mask)) { - plugin = uct_ib_mlx5_ext_find_plugin(iface, - UCT_IFACE_FLAG_V2_QUERY_TOKEN); - if (plugin == NULL) { - return UCS_ERR_UNSUPPORTED; - } - - status = plugin->ops.iface_query(iface, attr); - if (status != UCS_OK) { - ucs_error("ib mlx5 ext: iface query failed: %s", - ucs_status_string(status)); - return status; - } + if (attr == NULL) { + ucs_error("ib mlx5 ext: iface query attribute is NULL"); + return UCS_ERR_INVALID_PARAM; } if (attr->field_mask & UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS) { @@ -172,38 +71,6 @@ uct_ib_mlx5_ext_iface_query(uct_iface_h iface, return UCS_OK; } -ucs_status_t -uct_ib_mlx5_ext_ep_query(uct_ep_h ep, uct_ib_mlx5_ext_ep_query_attr_t *attr) -{ - uct_ib_mlx5_ext_plugin_t *plugin; - ucs_status_t status; - - status = uct_ib_mlx5_ext_ep_query_check_param(attr); - if (status != UCS_OK) { - return status; - } - - if (attr->field_mask & UCT_IB_MLX5_EXT_EP_QUERY_ATTR_FIELD_TX_TOKEN) { - plugin = uct_ib_mlx5_ext_find_plugin(ep->iface, - UCT_IFACE_FLAG_V2_QUERY_TOKEN); - if (plugin == NULL) { - return UCS_ERR_UNSUPPORTED; - } - - ucs_assert(!uct_ib_mlx5_ext_is_unsupported_op( - (const void*)plugin->ops.ep_query)); - - status = plugin->ops.ep_query(ep, attr); - if (status != UCS_OK) { - ucs_error("ib mlx5 ext: ep query failed: %s", - ucs_status_string(status)); - return status; - } - } - - return UCS_OK; -} - size_t uct_ib_mlx5_ext_max_put_sgl_zcopy_count(void) { uct_ib_mlx5_ext_plugin_t *plugin; @@ -344,17 +211,12 @@ 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 iface_query=%s " - "ep_query=%s put_sgl_zcopy=%s outstanding_purge=%s " - "(total=%u)", + "put_sgl_zcopy=%s outstanding_purge=%s (total=%u)", plugin->ops.name, uct_ib_mlx5_ext_is_unsupported_op( (const void*)plugin->ops.iface_query) ? "unsupported" : "supported", - uct_ib_mlx5_ext_is_unsupported_op( - (const void*)plugin->ops.ep_query) ? - "unsupported" : - "supported", uct_ib_mlx5_ext_is_unsupported_op( (const void*)plugin->ops.ep_put_sgl_zcopy) ? "unsupported" : diff --git a/src/uct/ib/mlx5/ib_mlx5_ext.h b/src/uct/ib/mlx5/ib_mlx5_ext.h index 8d6331934be..04a9d2a44b1 100644 --- a/src/uct/ib/mlx5/ib_mlx5_ext.h +++ b/src/uct/ib/mlx5/ib_mlx5_ext.h @@ -28,19 +28,7 @@ BEGIN_C_DECLS */ enum uct_ib_mlx5_ext_iface_query_attr_field { /** Enables @ref uct_ib_mlx5_ext_iface_query_attr_t::cap */ - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS = UCS_BIT(0), - - /** Enables @ref uct_ib_mlx5_ext_iface_query_attr_t::tx_token_len */ - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN_LEN = UCS_BIT(1), - - /** Enables @ref uct_ib_mlx5_ext_iface_query_attr_t::rx_token_len */ - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN_LEN = UCS_BIT(2), - - /** Enables @ref uct_ib_mlx5_ext_iface_query_attr_t::tx_token */ - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN = UCS_BIT(3), - - /** Enables @ref uct_ib_mlx5_ext_iface_query_attr_t::rx_token */ - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN = UCS_BIT(4) + UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS = UCS_BIT(0) }; /** @@ -58,62 +46,8 @@ typedef struct uct_ib_mlx5_ext_iface_query_attr { struct { uint64_t flags; /**< Flags from @ref UCT_RESOURCE_IFACE_CAP_V2 */ } cap; - - /** TX token length in bytes. */ - size_t tx_token_len; - - /** RX token length in bytes. */ - size_t rx_token_len; - - /** - * TX token input buffer. - * Valid when @ref UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN is set. - * Caller sets this to a buffer of @ref tx_token_len bytes containing - * the TX token received from the sender. - * @ref UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN must be set together. - */ - const void *tx_token; - - /** - * RX token output buffer. - * Valid when @ref UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN is set. - * Caller sets this to a pre-allocated buffer of @ref rx_token_len - * bytes; callee fills it with RX token. - * @ref UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN must be set together. - */ - void *rx_token; } uct_ib_mlx5_ext_iface_query_attr_t; -/** - * @brief EP query attributes field mask. - * - * The enumeration allows specifying which fields in - * @ref uct_ib_mlx5_ext_ep_query_attr_t are present. - */ -enum uct_ib_mlx5_ext_ep_query_attr_field { - /** Enables @ref uct_ib_mlx5_ext_ep_query_attr_t::tx_token */ - UCT_IB_MLX5_EXT_EP_QUERY_ATTR_FIELD_TX_TOKEN = UCS_BIT(0) -}; - -/** - * @brief EP query parameters. - */ -typedef struct uct_ib_mlx5_ext_ep_query_attr { - /** - * Mask of valid fields in this structure, using bits from - * @ref uct_ib_mlx5_ext_ep_query_attr_field. Fields not specified in this - * mask will be ignored. - */ - uint64_t field_mask; - - /** - * Pointer to a caller-allocated buffer for TX token data. The buffer size - * must be at least the TX token length returned by - * @ref uct_ib_mlx5_ext_iface_query. - */ - void *tx_token; -} uct_ib_mlx5_ext_ep_query_attr_t; - /** * @brief External plugin iface query callback. * @@ -126,18 +60,6 @@ typedef struct uct_ib_mlx5_ext_ep_query_attr { typedef ucs_status_t (*uct_ib_mlx5_ext_iface_query_func_t)( uct_iface_h iface, uct_ib_mlx5_ext_iface_query_attr_t *attr); -/** - * @brief External plugin EP query callback. - * - * @param [in] ep Endpoint to query. - * @param [in,out] attr Query parameters. Only fields selected by - * @a attr->field_mask should be accessed. - * - * @return UCS_OK on success, or an error if the operation failed. - */ -typedef ucs_status_t (*uct_ib_mlx5_ext_ep_query_func_t)( - uct_ep_h ep, uct_ib_mlx5_ext_ep_query_attr_t *attr); - /** * @brief External plugin maximum PUT SGL zero-copy entry count callback. * @@ -153,7 +75,6 @@ typedef size_t (*uct_ib_mlx5_ext_max_put_sgl_zcopy_count_func_t)(void); typedef struct uct_ib_mlx5_ext_ops { char name[UCT_COMPONENT_NAME_MAX]; /**< Plugin name */ uct_ib_mlx5_ext_iface_query_func_t iface_query; /**< Iface query callback */ - uct_ib_mlx5_ext_ep_query_func_t ep_query; /**< EP query callback */ 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_ep_outstanding_purge_func_t ep_outstanding_purge; /**< Outstanding operation purge callback */ @@ -184,9 +105,6 @@ ucs_status_t uct_ib_mlx5_ext_iface_query(uct_iface_h iface, uct_ib_mlx5_ext_iface_query_attr_t *attr); -ucs_status_t -uct_ib_mlx5_ext_ep_query(uct_ep_h ep, uct_ib_mlx5_ext_ep_query_attr_t *attr); - size_t uct_ib_mlx5_ext_max_put_sgl_zcopy_count(void); ucs_status_t uct_ib_mlx5_ext_ep_put_sgl_zcopy(uct_ep_h ep, diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index 322cd0b7bc9..026137eddfd 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -14,7 +14,6 @@ #endif #include -#include #include #include #include diff --git a/test/gtest/uct/ib/test_ib_mlx5_ext.cc b/test/gtest/uct/ib/test_ib_mlx5_ext.cc index 67528ff7b9c..0f747b40b45 100644 --- a/test/gtest/uct/ib/test_ib_mlx5_ext.cc +++ b/test/gtest/uct/ib/test_ib_mlx5_ext.cc @@ -7,7 +7,6 @@ #include "test_rc.h" extern "C" { -#include #include extern ucs_list_link_t uct_ib_mlx5_ext_plugins; @@ -58,44 +57,6 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { } } - static ucs_status_t - iface_query(uct_iface_h, uct_ib_mlx5_ext_iface_query_attr_t *attr) - { - if (attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_CAP_FLAGS) { - attr->cap.flags = UCT_IFACE_FLAG_V2_QUERY_TOKEN; - } - - if (attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_TX_TOKEN_LEN) { - attr->tx_token_len = sizeof(uint64_t); - } - - if (attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN_LEN) { - attr->rx_token_len = sizeof(uint64_t); - } - - if (attr->field_mask & - UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN) { - if ((attr->tx_token == NULL) || (attr->rx_token == NULL)) { - ADD_FAILURE() << "token buffers must be non-NULL"; - return UCS_ERR_INVALID_PARAM; - } - *static_cast(attr->rx_token) = - *static_cast(attr->tx_token) + 1; - } - - return UCS_OK; - } - - static ucs_status_t - ep_query(uct_ep_h, uct_ib_mlx5_ext_ep_query_attr_t *attr) - { - *static_cast(attr->tx_token) = tx_token(); - return UCS_OK; - } - static ucs_status_t purge_fail(uct_ep_h, const uct_ep_outstanding_purge_params_t*) { @@ -119,15 +80,11 @@ class test_uct_ib_mlx5_ext_rc : public test_rc { static void register_plugin(const char *name, - uct_ib_mlx5_ext_iface_query_func_t iface_query_cb = NULL, - uct_ib_mlx5_ext_ep_query_func_t ep_query_cb = NULL, uct_ep_outstanding_purge_func_t purge_cb = NULL) { uct_ib_mlx5_ext_ops_t ops = {}; ucs_strncpy_zero(ops.name, name, sizeof(ops.name)); - ops.iface_query = iface_query_cb; - ops.ep_query = ep_query_cb; ops.ep_outstanding_purge = purge_cb; ASSERT_UCS_OK(uct_ib_mlx5_ext_register(&ops)); } @@ -150,12 +107,9 @@ UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_outstanding_purge) uct_ep_outstanding_purge(m_e1->ep(0), ¶ms)); } - register_plugin("stub", - (uct_ib_mlx5_ext_iface_query_func_t) - ucs_empty_function_return_unsupported, - NULL, NULL); - register_plugin("fail", NULL, NULL, purge_fail); - register_plugin("token", iface_query, ep_query, purge); + register_plugin("stub"); + register_plugin("fail", purge_fail); + register_plugin("token", purge); params.field_mask = UCT_EP_OUTSTANDING_FIELD_RX_TOKEN | UCT_EP_OUTSTANDING_FIELD_CB | From 231a24474ec33f8358187d2155aa0cc5b14201ff Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Tue, 25 Aug 2026 15:42:45 +0300 Subject: [PATCH 11/13] UCT/IB/MLX5: address comments --- src/uct/ib/mlx5/rc/rc_mlx5_ep.c | 17 ++++++--- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 59 +++++++++++++++++------------ test/gtest/uct/ib/test_rc.cc | 6 +++ test/gtest/uct/v2/test_uct_query.cc | 13 +++++-- 4 files changed, 61 insertions(+), 34 deletions(-) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c index 026137eddfd..72a4bd8063f 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_ep.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_ep.c @@ -24,10 +24,10 @@ #include "rc_mlx5.inl" -#if HAVE_DEVX static ucs_status_t uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) { +#if HAVE_DEVX char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; uct_rc_mlx5_tx_token_t *tx_token; @@ -52,23 +52,28 @@ uct_rc_mlx5_ep_query_tx_token(uct_rc_mlx5_base_ep_t *ep, uct_ep_attr_t *ep_attr) *tx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, remote_qpn)); return UCS_OK; -} +#else + return UCS_ERR_UNSUPPORTED; #endif +} ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr) { + uct_ib_mlx5_md_t *md; + if (ep_attr->field_mask & (UCT_EP_ATTR_FIELD_LOCAL_SOCKADDR | UCT_EP_ATTR_FIELD_REMOTE_SOCKADDR)) { return UCS_ERR_UNSUPPORTED; } if (ep_attr->field_mask & UCT_EP_ATTR_FIELD_TX_TOKEN) { -#if HAVE_DEVX + md = uct_ib_mlx5_iface_md(ucs_derived_of(tl_ep->iface, uct_ib_iface_t)); + if (!(md->flags & UCT_IB_MLX5_MD_FLAG_DEVX)) { + return UCS_ERR_UNSUPPORTED; + } + return uct_rc_mlx5_ep_query_tx_token( ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t), ep_attr); -#else - return UCS_ERR_UNSUPPORTED; -#endif } return UCS_OK; diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index f52007335a2..3a6f6fb3157 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1042,11 +1042,11 @@ static void uct_rc_mlx5_iface_fill_ext_query_attr( } } -#if HAVE_DEVX static ucs_status_t uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, uct_iface_attr_v2_t *iface_attr) { +#if HAVE_DEVX char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)] = {}; char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {}; const uct_rc_mlx5_tx_token_t *tx_token; @@ -1070,7 +1070,12 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, tx_token = iface_attr->tx_token; iface = ucs_derived_of(tl_iface, uct_rc_iface_t); remote_qpn = be32toh(*tx_token); - rc_ep = uct_rc_iface_lookup_ep(iface, remote_qpn); + if (remote_qpn >= UCS_BIT(UCT_IB_QPN_ORDER)) { + ucs_error("rc mlx5: remote QPN %u is out of range", remote_qpn); + return UCS_ERR_INVALID_PARAM; + } + + rc_ep = uct_rc_iface_lookup_ep(iface, remote_qpn); if (rc_ep == NULL) { ucs_error("rc mlx5: no ep for QPN %u", remote_qpn); return UCS_ERR_INVALID_PARAM; @@ -1090,8 +1095,10 @@ uct_rc_mlx5_iface_query_rx_token(uct_iface_h tl_iface, *rx_token = htobe32(UCT_IB_MLX5DV_GET(qpc, qpc, next_rcv_psn)); return UCS_OK; -} +#else + return UCS_ERR_UNSUPPORTED; #endif +} static ucs_status_t uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, @@ -1099,38 +1106,40 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, { uct_ib_mlx5_ext_iface_query_attr_t ext_attr = {0}; uct_rc_mlx5_iface_common_t *mlx5_iface; + uct_ib_mlx5_md_t *md; size_t max_sgl; ucs_status_t status; uct_iface_query_v2_init(tl_iface, iface_attr); -#if HAVE_DEVX - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { - status = uct_rc_mlx5_iface_query_rx_token(tl_iface, iface_attr); - if (status != UCS_OK) { - return status; + md = uct_ib_mlx5_iface_md(ucs_derived_of(tl_iface, uct_ib_iface_t)); + if (md->flags & UCT_IB_MLX5_MD_FLAG_DEVX) { + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { + status = uct_rc_mlx5_iface_query_rx_token(tl_iface, iface_attr); + if (status != UCS_OK) { + return status; + } } - } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) { - iface_attr->cap.flags |= UCT_IFACE_FLAG_V2_QUERY_TOKEN; - } + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) { + iface_attr->cap.flags |= UCT_IFACE_FLAG_V2_QUERY_TOKEN; + } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH) { - iface_attr->tx_token_length = sizeof(uct_rc_mlx5_tx_token_t); - } + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH) { + iface_attr->tx_token_length = sizeof(uct_rc_mlx5_tx_token_t); + } - if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH) { - iface_attr->rx_token_length = sizeof(uct_rc_mlx5_rx_token_t); - } -#else - if (iface_attr->field_mask & - (UCT_IFACE_ATTR_FIELD_TX_TOKEN | UCT_IFACE_ATTR_FIELD_RX_TOKEN | - UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | - UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH)) { - return UCS_ERR_UNSUPPORTED; + if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH) { + iface_attr->rx_token_length = sizeof(uct_rc_mlx5_rx_token_t); + } + } else { + if (iface_attr->field_mask & + (UCT_IFACE_ATTR_FIELD_TX_TOKEN | UCT_IFACE_ATTR_FIELD_RX_TOKEN | + UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | + UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH)) { + return UCS_ERR_UNSUPPORTED; + } } -#endif if (iface_attr->field_mask & (UCT_IFACE_ATTR_FIELD_CAP_FLAGS | diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index d6e3e419ecd..dd83d846a00 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -1366,9 +1366,15 @@ UCS_TEST_SKIP_COND_P(test_rc_mlx5_token_query, am_short, uct_rc_mlx5_tx_token_t tx_token = {}; uct_rc_mlx5_rx_token_t rx_token = {}; uint32_t rx_count = 0; + uct_ib_mlx5_md_t *md = uct_ib_mlx5_iface_md( + ucs_derived_of(m_e1->iface(), uct_ib_iface_t)); uct_rc_mlx5_base_ep_t *e2_ep; ucs_status_t status; + if (!(md->flags & UCT_IB_MLX5_MD_FLAG_DEVX)) { + UCS_TEST_SKIP_R("DEVX is not supported"); + } + ASSERT_UCS_OK(uct_iface_set_am_handler(m_e2->iface(), 0, am_handler, &rx_count, 0)); diff --git a/test/gtest/uct/v2/test_uct_query.cc b/test/gtest/uct/v2/test_uct_query.cc index 46efb78574c..fefcd50969a 100644 --- a/test/gtest/uct/v2/test_uct_query.cc +++ b/test/gtest/uct/v2/test_uct_query.cc @@ -16,6 +16,9 @@ extern "C" { #include #include #include +#if HAVE_DEVX +#include +#endif } @@ -118,17 +121,21 @@ UCS_TEST_P(test_uct_query, query_token_support) { uct_iface_attr_v2_t attr = {}; uint64_t iface_cap_flags = 0; + uct_ib_mlx5_md_t *md; attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS; attr.cap.flags = UINT64_MAX; ASSERT_UCS_OK(uct_iface_query_v2(get_iface(), &attr)); -#if HAVE_DEVX if (has_transport("rc_mlx5")) { - iface_cap_flags = UCT_IFACE_FLAG_V2_QUERY_TOKEN; + md = uct_ib_mlx5_iface_md(ucs_derived_of(get_iface(), uct_ib_iface_t)); + + if (md->flags & UCT_IB_MLX5_MD_FLAG_DEVX) { + iface_cap_flags = UCT_IFACE_FLAG_V2_QUERY_TOKEN; + } } -#endif + EXPECT_EQ(iface_cap_flags, attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); } From c0df5345e25098a6bbe15d45ff4b7b32e85e16f1 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Wed, 26 Aug 2026 13:55:48 +0300 Subject: [PATCH 12/13] UCT/IB/MLX5: return invalid if rx/tx not set togather --- src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index 3a6f6fb3157..2710ea16a35 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -1112,6 +1112,17 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, uct_iface_query_v2_init(tl_iface, iface_attr); + if (ucs_test_flags(iface_attr->field_mask, + UCT_IFACE_ATTR_FIELD_TX_TOKEN | + UCT_IFACE_ATTR_FIELD_RX_TOKEN) && + !ucs_test_all_flags(iface_attr->field_mask, + UCT_IFACE_ATTR_FIELD_TX_TOKEN | + UCT_IFACE_ATTR_FIELD_RX_TOKEN)) { + ucs_error("rc mlx5: tx and rx token fields must be set together"); + + return UCS_ERR_INVALID_PARAM; + } + md = uct_ib_mlx5_iface_md(ucs_derived_of(tl_iface, uct_ib_iface_t)); if (md->flags & UCT_IB_MLX5_MD_FLAG_DEVX) { if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_RX_TOKEN) { @@ -1133,10 +1144,11 @@ uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface, iface_attr->rx_token_length = sizeof(uct_rc_mlx5_rx_token_t); } } else { - if (iface_attr->field_mask & - (UCT_IFACE_ATTR_FIELD_TX_TOKEN | UCT_IFACE_ATTR_FIELD_RX_TOKEN | - UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | - UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH)) { + if (ucs_test_flags(iface_attr->field_mask, + UCT_IFACE_ATTR_FIELD_TX_TOKEN | + UCT_IFACE_ATTR_FIELD_RX_TOKEN | + UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH | + UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH)) { return UCS_ERR_UNSUPPORTED; } } From b4cd8f9fca2fb51d444232e0433ebe944117fea8 Mon Sep 17 00:00:00 2001 From: Zihao Zhao Date: Wed, 26 Aug 2026 13:49:29 +0300 Subject: [PATCH 13/13] UCT/IB/MLX5: address comments --- test/gtest/uct/v2/test_uct_query.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/test/gtest/uct/v2/test_uct_query.cc b/test/gtest/uct/v2/test_uct_query.cc index fefcd50969a..9977220ad0a 100644 --- a/test/gtest/uct/v2/test_uct_query.cc +++ b/test/gtest/uct/v2/test_uct_query.cc @@ -16,7 +16,7 @@ extern "C" { #include #include #include -#if HAVE_DEVX +#ifdef HAVE_MLX5_DV #include #endif } @@ -121,20 +121,21 @@ UCS_TEST_P(test_uct_query, query_token_support) { uct_iface_attr_v2_t attr = {}; uint64_t iface_cap_flags = 0; - uct_ib_mlx5_md_t *md; attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS; attr.cap.flags = UINT64_MAX; ASSERT_UCS_OK(uct_iface_query_v2(get_iface(), &attr)); +#ifdef HAVE_MLX5_DV if (has_transport("rc_mlx5")) { - md = uct_ib_mlx5_iface_md(ucs_derived_of(get_iface(), uct_ib_iface_t)); - + uct_ib_mlx5_md_t *md = uct_ib_mlx5_iface_md( + ucs_derived_of(get_iface(), uct_ib_iface_t)); if (md->flags & UCT_IB_MLX5_MD_FLAG_DEVX) { iface_cap_flags = UCT_IFACE_FLAG_V2_QUERY_TOKEN; } } +#endif EXPECT_EQ(iface_cap_flags, attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN); }