UCT/IB/MLX5: support in progress error handler

src/uct/ib/mlx5/ib_mlx5.c

@@ -698,10 +698,11 @@ void uct_ib_mlx5_txwq_reset(uct_ib_mlx5_txwq_t *txwq)
     txwq->sw_pi          = 0;
     txwq->prev_sw_pi     = UINT16_MAX;
     txwq->next_wqe_psn   = 0;
+    txwq->hw_ci          = UINT16_MAX;
+    txwq->ft_ci          = UINT16_MAX;
     txwq->path_mtu_mask  = 0;
     txwq->path_mtu_shift = 0;
 #if UCS_ENABLE_ASSERT
-    txwq->hw_ci          = 0xFFFF;
     txwq->flags          = 0;
 #endif
     uct_ib_fence_info_init(&txwq->fi);
@@ -746,10 +747,10 @@ void uct_ib_mlx5_txwq_vfs_populate(uct_ib_mlx5_txwq_t *txwq, void *parent_obj)
                             UCS_VFS_TYPE_U16, "bb_max");
     ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->sig_pi,
                             UCS_VFS_TYPE_U16, "sig_pi");
-#if UCS_ENABLE_ASSERT
+    ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->ft_ci,
+                            UCS_VFS_TYPE_U16, "ft_ci");
     ucs_vfs_obj_add_ro_file(parent_obj, ucs_vfs_show_primitive, &txwq->hw_ci,
                             UCS_VFS_TYPE_U16, "hw_ci");
-#endif
 }
 
 ucs_status_t

src/uct/ib/mlx5/ib_mlx5.h

@@ -695,10 +695,11 @@ typedef struct uct_ib_mlx5_txwq {
     void                        *qend;
     uint16_t                    bb_max;
     uint16_t                    sig_pi;     /* PI for last signaled WQE */
+    uint16_t                    hw_ci;      /* First BB index of last completed WQE */
+    uint16_t                    ft_ci;      /* First BB index of last ft completed WQE */
     uint16_t                    path_mtu_mask;  /* Path MTU in bytes - 1 */
     uint8_t                     path_mtu_shift; /* log2(path MTU in bytes) */
 #if UCS_ENABLE_ASSERT
-    uint16_t                    hw_ci; /* First BB index of last completed WQE */
     uint8_t                     flags; /* Debug flags */
 #endif
     uct_ib_fence_info_t         fi;

src/uct/ib/mlx5/ib_mlx5.inl

@@ -191,9 +191,7 @@ uct_ib_mlx5_poll_cq(uct_ib_iface_t *iface, uct_ib_mlx5_cq_t *cq, int poll_flags,
 static UCS_F_ALWAYS_INLINE uint16_t
 uct_ib_mlx5_txwq_update_bb(uct_ib_mlx5_txwq_t *wq, uint16_t hw_ci)
 {
-#if UCS_ENABLE_ASSERT
     wq->hw_ci = hw_ci;
-#endif
     return wq->bb_max - (wq->prev_sw_pi - hw_ci);
 }
 

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

@@ -52,6 +52,7 @@ typedef struct uct_rc_mlx5_base_ep {
     struct {
         uct_ib_mlx5_txwq_t   wq;
     } tx;
+    uint8_t                  err_handler_inprogress;
 } uct_rc_mlx5_base_ep_t;
 
 typedef __be32 uct_rc_mlx5_tx_token_t;

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

@@ -1931,31 +1931,36 @@ uct_rc_mlx5_iface_common_atomic_data(unsigned opcode, unsigned size, uint64_t va
 }
 
 static UCS_F_ALWAYS_INLINE void
-uct_rc_mlx5_iface_update_tx_res(uct_rc_iface_t *rc_iface,
-                                uct_rc_mlx5_base_ep_t *rc_mlx5_base_ep,
-                                uint16_t hw_ci)
+uct_rc_mlx5_iface_update_tx_cq_res(uct_rc_iface_t *rc_iface,
+                                   uct_rc_mlx5_base_ep_t *ep, uint16_t hw_ci)
 {
-    uct_ib_mlx5_txwq_t *txwq = &rc_mlx5_base_ep->tx.wq;
-    uct_rc_txqp_t *txqp      = &rc_mlx5_base_ep->super.txqp;
-    uint16_t bb_num;
-
-    bb_num = uct_ib_mlx5_txwq_update_bb(txwq, hw_ci) -
-             uct_rc_txqp_available(txqp);
-
-    /* Must always have positive number of released resources. The first
-     * completion will report bb_num=1 (because prev_sw_pi is initialized to -1)
-     * and all the rest report the amount of BBs the previous WQE has consumed.
-     */
-    ucs_assertv(bb_num > 0, "hw_ci=%d prev_sw_pi=%d available=%d bb_num=%d",
-                hw_ci, txwq->prev_sw_pi, txqp->available, bb_num);
+    uct_ib_mlx5_txwq_t *txwq = &ep->tx.wq;
+    uint16_t prev_hw_ci      = txwq->hw_ci;
+    uint16_t bb_num          = hw_ci - prev_hw_ci;
 
-    uct_rc_txqp_available_add(txqp, bb_num);
-    ucs_assert(uct_rc_txqp_available(txqp) <= txwq->bb_max);
+    ucs_assertv(bb_num > 0, "hw_ci=%d prev_hw_ci=%d ft_ci=%d bb_num=%d", hw_ci,
+                prev_hw_ci, txwq->ft_ci, bb_num);
 
+    txwq->hw_ci = hw_ci;
     uct_rc_iface_update_reads(rc_iface);
     uct_rc_iface_add_cq_credits(rc_iface, bb_num);
 }
 
+static UCS_F_ALWAYS_INLINE void
+uct_rc_mlx5_ep_update_tx_qp_res(uct_rc_mlx5_base_ep_t *ep, uint16_t sw_ci)
+{
+    uct_ib_mlx5_txwq_t *txwq = &ep->tx.wq;
+    uct_rc_txqp_t *txqp      = &ep->super.txqp;
+    int16_t prev_available   = uct_rc_txqp_available(txqp);
+    uint16_t available       = txwq->bb_max -
+                               (txwq->prev_sw_pi - sw_ci);
+
+    ucs_assert(available >= prev_available);
+    uct_rc_txqp_available_add(txqp, available - prev_available);
+
+    ucs_assert(uct_rc_txqp_available(txqp) <= txwq->bb_max);
+}
+
 static UCS_F_ALWAYS_INLINE unsigned
 uct_rc_mlx5_iface_poll_tx(uct_rc_mlx5_iface_common_t *iface, int poll_flags)
 {
@@ -1986,7 +1991,8 @@ uct_rc_mlx5_iface_poll_tx(uct_rc_mlx5_iface_common_t *iface, int poll_flags)
 
     uct_rc_mlx5_txqp_process_tx_cqe(&ep->super.txqp, cqe, hw_ci);
     ucs_arbiter_group_schedule(&iface->super.tx.arbiter, &ep->super.arb_group);
-    uct_rc_mlx5_iface_update_tx_res(&iface->super, ep, hw_ci);
+    uct_rc_mlx5_ep_update_tx_qp_res(ep, hw_ci);
+    uct_rc_mlx5_iface_update_tx_cq_res(&iface->super, ep, hw_ci);
     uct_rc_iface_arbiter_dispatch(&iface->super);
     uct_ib_mlx5_update_db_cq_ci(&iface->cq[UCT_IB_DIR_TX]);
 

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

@@ -1242,6 +1242,7 @@ UCS_CLASS_INIT_FUNC(uct_rc_mlx5_base_ep_t, const uct_ep_params_t *params)
 
     UCS_CLASS_CALL_SUPER_INIT(uct_rc_ep_t, &iface->super,
                               self->tx.wq.super.qp_num, params);
+    self->err_handler_inprogress = 0;
 
     if (self->tx.wq.super.type == UCT_IB_MLX5_OBJ_TYPE_VERBS) {
         status = uct_rc_iface_qp_init(&iface->super,

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

@@ -194,6 +194,10 @@ void uct_rc_mlx5_iface_handle_failure(uct_ib_iface_t *ib_iface, void *arg,
                                                                       qp_num),
                                                uct_rc_mlx5_base_ep_t);
     uint16_t pi               = ntohs(cqe->wqe_counter);
+#if UCS_ENABLE_ASSERT
+    uct_iface_attr_v2_t iface_attr;
+    ucs_status_t query_status;
+#endif
     ucs_log_level_t log_lvl;
     ucs_status_t status;
 
@@ -203,27 +207,57 @@ void uct_rc_mlx5_iface_handle_failure(uct_ib_iface_t *ib_iface, void *arg,
         goto out;
     }
 
-    uct_rc_txqp_purge_outstanding(iface, &ep->super.txqp, ep_status, pi, 0);
     ucs_arbiter_group_purge(&iface->tx.arbiter, &ep->super.arb_group,
                             uct_rc_ep_arbiter_purge_internal_cb, NULL);
-    uct_rc_mlx5_iface_update_tx_res(iface, ep, pi);
     uct_ib_mlx5_txwq_update_flags(&ep->tx.wq, UCT_IB_MLX5_TXWQ_FLAG_FAILED, 0);
 
     if (ep->super.flags & (UCT_RC_EP_FLAG_ERR_HANDLER_INVOKED |
                            UCT_RC_EP_FLAG_FLUSH_CANCEL)) {
-        goto out;
+        goto out_update_tx_res;
     }
 
     ep->super.flags |= UCT_RC_EP_FLAG_ERR_HANDLER_INVOKED;
     uct_rc_fc_restore_wnd(iface, &ep->super.fc);
 
     status  = uct_iface_handle_ep_err(&iface->super.super.super,
                                       &ep->super.super.super, ep_status);
-    log_lvl = uct_base_iface_failure_log_level(&ib_iface->super, status,
-                                               ep_status);
+
+    if (status == UCS_INPROGRESS) {
+#if UCS_ENABLE_ASSERT
+        iface_attr.field_mask = UCT_IFACE_ATTR_FIELD_CAP_FLAGS;
+
+        query_status = uct_iface_query_v2(&iface->super.super.super,
+                                          &iface_attr);
+        ucs_assert(query_status == UCS_OK);
+        ucs_assert(iface_attr.cap.flags & UCT_IFACE_FLAG_V2_QUERY_TOKEN);
+#endif
+
+        /* Save last completed WQE. TX QP resources are reserved until purge. */
+        ep->tx.wq.ft_ci            = ep->tx.wq.prev_sw_pi -
+                                     (ep->tx.wq.bb_max -
+                                      uct_rc_txqp_available(&ep->super.txqp));
+        ep->err_handler_inprogress = 1;
+
+        ucs_debug("ep %p outstanding WQE range (%u, %u)", ep, ep->tx.wq.ft_ci,
+                  ep->tx.wq.sw_pi);
+
+        log_lvl = uct_base_iface_failure_log_level(&ib_iface->super, UCS_OK,
+                                                   ep_status);
+    } else {
+        log_lvl = uct_base_iface_failure_log_level(&ib_iface->super, status,
+                                                   ep_status);
+    }
 
     uct_ib_mlx5_completion_with_err(ib_iface, arg, &ep->tx.wq, log_lvl);
 
+out_update_tx_res:
+    if (!(ep->err_handler_inprogress)) {
+        uct_rc_txqp_purge_outstanding(iface, &ep->super.txqp, ep_status, pi, 0);
+        uct_rc_mlx5_ep_update_tx_qp_res(ep, pi);
+    }
+
+    uct_rc_mlx5_iface_update_tx_cq_res(iface, ep, pi);
+
 out:
     uct_rc_iface_arbiter_dispatch(iface);
 }