UCT/IB/MLX5: Purge outstanding AM short operations

src/uct/ib/mlx5/ib_mlx5.c

@@ -858,6 +858,23 @@ void *uct_ib_mlx5_txwq_get_wqe(const uct_ib_mlx5_txwq_t *txwq, uint16_t pi)
     return UCS_PTR_BYTE_OFFSET(txwq->qstart, (pi % num_bb) * MLX5_SEND_WQE_BB);
 }
 
+size_t uct_ib_mlx5_wqe_size(const struct mlx5_wqe_ctrl_seg *ctrl)
+{
+    uint8_t ds = ntohl(ctrl->qpn_ds) & UINT8_MAX;
+
+    ucs_assertv_always(
+            (ds > 0) &&
+            ((ds * UCT_IB_MLX5_WQE_SEG_SIZE) <=
+             UCT_IB_MLX5_MAX_SEND_WQE_SIZE),
+            "ds=%u", ds);
+    return ds * UCT_IB_MLX5_WQE_SEG_SIZE;
+}
+
+uint16_t uct_ib_mlx5_txwq_next_wqe_index(uint16_t index, size_t wqe_size)
+{
+    return index + ucs_div_round_up(wqe_size, MLX5_SEND_WQE_BB);
+}
+
 uint16_t uct_ib_mlx5_txwq_num_posted_wqes(const uct_ib_mlx5_txwq_t *txwq,
                                           uint16_t outstanding)
 {
@@ -873,8 +890,8 @@ uint16_t uct_ib_mlx5_txwq_num_posted_wqes(const uct_ib_mlx5_txwq_t *txwq,
     ucs_assert(pi == txwq->hw_ci);
     do {
         ctrl     = uct_ib_mlx5_txwq_get_wqe(txwq, pi);
-        wqe_size = (ctrl->qpn_ds >> 24) * UCT_IB_MLX5_WQE_SEG_SIZE;
-        pi      += (wqe_size + MLX5_SEND_WQE_BB - 1) / MLX5_SEND_WQE_BB;
+        wqe_size = uct_ib_mlx5_wqe_size(ctrl);
+        pi       = uct_ib_mlx5_txwq_next_wqe_index(pi, wqe_size);
         ++count;
     } while (pi != txwq->sw_pi);
 

src/uct/ib/mlx5/ib_mlx5.h

@@ -64,6 +64,8 @@
 #define UCT_IB_MLX5_CQ_SET_CI            0
 #define UCT_IB_MLX5_CQ_ARM_DB            1
 #define UCT_IB_MLX5_LOG_MAX_MSG_SIZE     30
+#define UCT_IB_MLX5_PSN_BITS             24
+#define UCT_IB_MLX5_PSN_MASK             UCS_MASK(UCT_IB_MLX5_PSN_BITS)
 #define UCT_IB_MLX5_ATOMIC_MODE_COMP     1
 #define UCT_IB_MLX5_ATOMIC_MODE_EXT      3
 #define UCT_IB_MLX5_CQE_FLAG_L3_IN_DATA  UCS_BIT(28) /* GRH/IP in the receive buffer */
@@ -709,7 +711,7 @@ typedef struct uct_ib_mlx5_txwq {
 static UCS_F_ALWAYS_INLINE uint32_t
 uct_ib_mlx5_txwq_get_next_wqe_psn(const uct_ib_mlx5_txwq_t *txwq)
 {
-    return txwq->next_wqe_psn & UCS_MASK(24);
+    return txwq->next_wqe_psn & UCT_IB_MLX5_PSN_MASK;
 }
 
 
@@ -917,6 +919,12 @@ ucs_status_t uct_ib_mlx5_txwq_init(uct_priv_worker_t *worker,
 /* Get pointer to a WQE by producer index */
 void *uct_ib_mlx5_txwq_get_wqe(const uct_ib_mlx5_txwq_t *txwq, uint16_t pi);
 
+/* Get the WQE size in bytes from its control segment */
+size_t uct_ib_mlx5_wqe_size(const struct mlx5_wqe_ctrl_seg *ctrl);
+
+/* Get the index of the WQE that follows a WQE of the given size */
+uint16_t uct_ib_mlx5_txwq_next_wqe_index(uint16_t index, size_t wqe_size);
+
 /* Count how many WQEs are currently posted */
 uint16_t uct_ib_mlx5_txwq_num_posted_wqes(const uct_ib_mlx5_txwq_t *txwq,
                                           uint16_t outstanding);

src/uct/ib/mlx5/ib_mlx5_ext.c

@@ -111,55 +111,6 @@ uct_ib_mlx5_ext_ep_put_sgl_zcopy(uct_ep_h ep, void * const *buffers,
     return UCS_ERR_UNSUPPORTED;
 }
 
-static ucs_status_t uct_ib_mlx5_ext_ep_outstanding_purge_check_param(
-        const uct_ep_outstanding_purge_params_t *params)
-{
-    if (params == NULL) {
-        ucs_error("ib mlx5 ext: outstanding purge parameters are NULL");
-        return UCS_ERR_INVALID_PARAM;
-    }
-
-    if (!(params->field_mask & UCT_EP_OUTSTANDING_FIELD_RX_TOKEN) ||
-        !(params->field_mask & UCT_EP_OUTSTANDING_FIELD_CB) ||
-        (params->rx_token == NULL) || (params->cb == NULL)) {
-        ucs_error("ib mlx5 ext: rx token or callback is not set or is NULL");
-        return UCS_ERR_INVALID_PARAM;
-    }
-
-    return UCS_OK;
-}
-
-ucs_status_t uct_ib_mlx5_ext_ep_outstanding_purge(
-        uct_ep_h ep, const uct_ep_outstanding_purge_params_t *params)
-{
-    ucs_status_t status;
-    uct_ib_mlx5_ext_plugin_t *plugin;
-
-    status = uct_ib_mlx5_ext_ep_outstanding_purge_check_param(params);
-    if (status != UCS_OK) {
-        return status;
-    }
-
-    ucs_list_for_each(plugin, &uct_ib_mlx5_ext_plugins, list) {
-        if (ucs_unlikely(uct_ib_mlx5_ext_is_unsupported_op(
-                    (const void*)plugin->ops.ep_outstanding_purge))) {
-            continue;
-        }
-
-        /* The provider validates parameters and token signatures itself, and
-         * returns an error if it should not handle this request. */
-        status = plugin->ops.ep_outstanding_purge(ep, params);
-        if (status != UCS_OK) {
-            continue;
-        }
-
-        return UCS_OK;
-    }
-
-    ucs_error("ib mlx5 ext: ep outstanding purge is not supported");
-    return UCS_ERR_UNSUPPORTED;
-}
-
 void uct_ib_mlx5_ext_cleanup(void)
 {
     uct_ib_mlx5_ext_plugin_t *plugin, *tmp;
@@ -211,7 +162,7 @@ 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 "
-              "put_sgl_zcopy=%s outstanding_purge=%s (total=%u)",
+              "put_sgl_zcopy=%s (total=%u)",
               plugin->ops.name,
               uct_ib_mlx5_ext_is_unsupported_op(
                       (const void*)plugin->ops.iface_query) ?
@@ -221,10 +172,6 @@ ucs_status_t uct_ib_mlx5_ext_register(const uct_ib_mlx5_ext_ops_t *ops)
                       (const void*)plugin->ops.ep_put_sgl_zcopy) ?
                       "unsupported" :
                       "supported",
-              uct_ib_mlx5_ext_is_unsupported_op(
-                      (const void*)plugin->ops.ep_outstanding_purge) ?
-                      "unsupported" :
-                      "supported",
               num_plugins);
     return UCS_OK;
 }

src/uct/ib/mlx5/ib_mlx5_ext.h

@@ -77,7 +77,6 @@ typedef struct uct_ib_mlx5_ext_ops {
     uct_ib_mlx5_ext_iface_query_func_t             iface_query;                  /**< Iface 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 */
 } uct_ib_mlx5_ext_ops_t;
 
 /**
@@ -118,9 +117,6 @@ ucs_status_t uct_ib_mlx5_ext_ep_put_sgl_zcopy(uct_ep_h ep,
                                               size_t count,
                                               uct_completion_t *comp);
 
-ucs_status_t uct_ib_mlx5_ext_ep_outstanding_purge(
-        uct_ep_h ep, const uct_ep_outstanding_purge_params_t *params);
-
 END_C_DECLS
 
 #endif /* UCT_IB_MLX5_EXT_H_ */

src/uct/ib/mlx5/ib_mlx5_log.c

@@ -16,11 +16,6 @@
 #include <string.h>
 
 
-static void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
-                                 void *qend, int max_sge, int dump_qp,
-                                 uct_log_data_dump_func_t packet_dump_cb,
-                                 char *buffer, size_t max, uct_ib_log_sge_t *log_sge);
-
 static void uct_ib_mlx5_resp_error_dump(const uct_ib_mlx5_srq_seg_t *seg,
                                         unsigned max_strides, char *buffer,
                                         size_t max);
@@ -281,10 +276,10 @@ static size_t uct_ib_mlx5_dump_dgram(char *buf, size_t max, void *seg, int is_et
            UCT_IB_MLX5_AV_FULL_SIZE : UCT_IB_MLX5_AV_BASE_SIZE;
 }
 
-static void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
-                                 void *qend, int max_sge, int dump_qp,
-                                 uct_log_data_dump_func_t packet_dump_cb,
-                                 char *buffer, size_t max, uct_ib_log_sge_t *log_sge)
+void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
+                          void *qend, int max_sge, int dump_qp,
+                          uct_log_data_dump_func_t packet_dump_cb,
+                          char *buffer, size_t max, uct_ib_log_sge_t *log_sge)
 {
     static uct_ib_opcode_t opcodes[] = {
         [MLX5_OPCODE_NOP]              = { "NOP",        0 },
@@ -308,18 +303,26 @@ static void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
     uint8_t opcode                 = ctrl->opmod_idx_opcode >> 24;
     uint8_t opmod                  = ctrl->opmod_idx_opcode & 0xff;
     uint32_t qp_num                = ntohl(ctrl->qpn_ds) >> 8;
-    int ds                         = ctrl->qpn_ds >> 24;
-    uct_ib_opcode_t *op            = &opcodes[opcode];
+    int ds                         = ntohl(ctrl->qpn_ds) & UINT8_MAX;
     char *s                        = buffer;
     char *ends                     = buffer + max;
-    const char* sg_prefix_arr      = (op->flags & UCT_IB_OPCODE_FLAG_HAS_DMA) ?
-                                     "GS" : NULL;
+    uct_ib_opcode_t *op;
+    const char *sg_prefix_arr;
     struct ibv_sge sg_list[16];
     uint64_t inline_bitmap;
     int i, is_inline, is_eth;
     size_t dg_size;
     void *seg;
 
+    if ((opcode >= ucs_static_array_size(opcodes)) ||
+        (opcodes[opcode].name == NULL)) {
+        snprintf(buffer, max, "unknown opcode 0x%x", opcode);
+        return;
+    }
+
+    op            = &opcodes[opcode];
+    sg_prefix_arr = (op->flags & UCT_IB_OPCODE_FLAG_HAS_DMA) ? "GS" : NULL;
+
     /* QP and WQE index */
     if (dump_qp) {
         snprintf(s, ends - s, "QP 0x%x [%03ld] ", qp_num,
@@ -387,8 +390,8 @@ static void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
 
     /* Extended atomic segment */
     if (op->flags & UCT_IB_OPCODE_FLAG_HAS_EXT_ATOMIC) {
-        uint64_t add, boundary, compare, swap, compare_mask, swap_mask;
         int size = 1 << ((opmod & 7) + 2);
+        uint64_t add, boundary, compare, swap, compare_mask, swap_mask;
 
         if (opcode == MLX5_OPCODE_ATOMIC_MASKED_FA) {
             add      = network_to_host(seg, size);

src/uct/ib/mlx5/ib_mlx5_log.h

@@ -24,6 +24,12 @@ ucs_status_t uct_ib_mlx5_completion_with_err(uct_ib_iface_t *iface,
                                              ucs_log_level_t log_level);
 
 
+void uct_ib_mlx5_wqe_dump(uct_ib_iface_t *iface, void *wqe, void *qstart,
+                          void *qend, int max_sge, int dump_qp,
+                          uct_log_data_dump_func_t packet_dump_cb,
+                          char *buffer, size_t max, uct_ib_log_sge_t *log_sge);
+
+
 void __uct_ib_mlx5_log_tx(const char *file, int line, const char *function,
                           uct_ib_iface_t *iface, void *wqe, void *qstart,
                           void *qend, int max_sge, uct_ib_log_sge_t *log_sge,

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

@@ -209,6 +209,9 @@ ucs_status_t
 uct_rc_mlx5_base_ep_invalidate(uct_ep_h tl_ep,
                                const uct_ep_invalidate_params_t *params);
 
+ucs_status_t uct_rc_mlx5_ep_outstanding_purge(
+        uct_ep_h tl_ep, const uct_ep_outstanding_purge_params_t *params);
+
 ucs_status_t uct_rc_mlx5_base_ep_fc_ctrl(uct_ep_t *tl_ep, unsigned op,
                                          uct_rc_pending_req_t *req);