UCT: add token query support in ucx

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;
 

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);
 

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

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

@@ -25,24 +25,54 @@
 #include "rc_mlx5.inl"
 
 
-ucs_status_t uct_rc_mlx5_base_ep_query(uct_ep_h tl_ep, uct_ep_attr_t *ep_attr)
+#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, uct_ep_attr_t *ep_attr)
 {
-    uct_ib_mlx5_ext_ep_query_attr_t attr = {};
+    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;
     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;
+}
+#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);
+#endif
 
     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) {
-        attr.field_mask |= UCT_IB_MLX5_EXT_EP_QUERY_ATTR_FIELD_TX_TOKEN;
-        attr.tx_token    = ep_attr->tx_token;
-
-        status = uct_ib_mlx5_ext_ep_query(tl_ep, &attr);
-        if (status != UCS_OK) {
-            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;

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

@@ -1040,27 +1040,57 @@ 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;
+#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)
+{
+    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_TX_TOKEN)) {
+        ucs_error("rc mlx5: tx token is required to query rx token");
+        return UCS_ERR_INVALID_PARAM;
     }
 
-    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;
+    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;
     }
 
-    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;
+    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_RX_TOKEN) {
-        ext_attr->field_mask |= UCT_IB_MLX5_EXT_IFACE_QUERY_ATTR_FIELD_RX_TOKEN;
-        ext_attr->rx_token    = iface_attr->rx_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: 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;
 }
+#endif
 
 static ucs_status_t
 uct_rc_mlx5_iface_query_v2(uct_iface_h tl_iface,
@@ -1073,6 +1103,34 @@ 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) {
+        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_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 +1166,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;
 }
 

test/gtest/uct/ib/test_ib_mlx5_ext.cc

@@ -138,55 +138,6 @@ 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_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;
-        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;
-
-    ASSERT_UCS_OK(uct_iface_query_v2(m_e1->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);
-}
-
-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);
-
-    {
-        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(tx_token(), tx_token_value);
-}
-
 UCS_TEST_P(test_uct_ib_mlx5_ext_rc, ep_outstanding_purge)
 {
     uint64_t rx_token_value                  = rx_token();

test/gtest/uct/ib/test_rc.cc

@@ -1323,6 +1323,71 @@ UCT_INSTANTIATE_RC_TEST_CASE(test_rc_keepalive)
 
 #ifdef HAVE_MLX5_DV
 
+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<uint32_t*>(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));
+    }
+};
+
+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 = {};
+    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,
+                                           &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);
+    }
+
+    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);
+
+    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)
+
 class test_rc_srq : public test_rc {
 public:
     test_rc_srq() : m_buf8b(NULL), m_buf8k(NULL)

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 <common/test.h>
 #include <gtest/uct/uct_p2p_test.h>
 
@@ -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)