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/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" :

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,

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

@@ -54,6 +54,9 @@ typedef struct uct_rc_mlx5_base_ep {
     } tx;
 } uct_rc_mlx5_base_ep_t;
 
+typedef __be32 uct_rc_mlx5_tx_token_t;
+typedef __be32 uct_rc_mlx5_rx_token_t;
+
 
 /**
  * Generic RC remote endpoint with HW TM support