UCP/PERF: Fixed races in progress (#10935)

* UCP/PERF: Fixed race in progress

* UCP/GTEST: Added concurrent litmus test

* UCP/GDAKI: Fixed race in FC calculation

* UCP/GDAKI: CQ size = 2xTX_QP

* UCP/GDAKI: Improved FC perf

* UCP/GDAKI: Fixed race in progress/reservation with 2 counters

* UCP/GTEST: Fixed test infra

* UCP/GDA: Fixed uninitialized bug

* UCP/GDA: Cleanup & minor changes

* UCP/GTEST: Added stress test for multi

* UCP/GDA: Fixed CQ overflow by validating the CQE after read

* UCP/GTEST: Cosmetic changes

* UCP: Addressed PR comments

* UCP/GTEST: Reduce MAX_THREADS to 128, skip on valgrind

* UCP: Single atomic read in progress

* UCP/GDA: Sync 64 bits of wqe_base

* UCP/GDAKI: Fix for uninit wqe_base

* UCP/GTEST: Disable stress tests temporarily

Author: iyastreb

src/uct/ib/mlx5/gdaki/gdaki.c

@@ -176,8 +176,16 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_ep_t, const uct_ep_params_t *params)
     dev_ep.sq_num       = self->qp.super.qp_num;
     dev_ep.sq_wqe_daddr = UCS_PTR_BYTE_OFFSET(self->ep_gpu,
                                               qp_attr.umem_offset);
-    dev_ep.sq_wqe_num   = qp_attr.max_tx;
-    dev_ep.sq_dbrec     = &self->ep_gpu->qp_dbrec[MLX5_SND_DBR];
+    dev_ep.sq_rsvd_index  = 0;
+    dev_ep.sq_ready_index = 0;
+    dev_ep.sq_wqe_pi      = 0;
+    dev_ep.sq_wqe_num     = qp_attr.max_tx;
+    /* FC mask is used to determine if WQE should be posted with completion.
+     * qp_attr.max_tx must be a power of 2. */
+    dev_ep.sq_fc_mask     = (qp_attr.max_tx >> 1) - 1;
+    dev_ep.avail_count    = qp_attr.max_tx;
+    dev_ep.sq_dbrec       = &self->ep_gpu->qp_dbrec[MLX5_SND_DBR];
+
     dev_ep.cqe_daddr = UCS_PTR_BYTE_OFFSET(self->ep_gpu, cq_attr.umem_offset);
     dev_ep.cqe_num   = cq_attr.cq_size;
     dev_ep.sq_db     = self->sq_db;

src/uct/ib/mlx5/gdaki/gdaki.cuh

@@ -93,71 +93,33 @@ template<ucs_device_level_t level> UCS_F_DEVICE void uct_rc_mlx5_gda_sync(void)
     }
 }
 
-UCS_F_DEVICE uint64_t uct_rc_mlx5_gda_max_alloc_wqe_base(
-    uct_rc_gdaki_dev_ep_t *ep, unsigned count)
+UCS_F_DEVICE uint64_t
+uct_rc_mlx5_gda_reserv_wqe_thread(uct_rc_gdaki_dev_ep_t *ep, unsigned count)
 {
-    /* TODO optimize by including sq_wqe_num in qp->sq_wqe_pi and updating it
-       when processing a new completion */
-    uint64_t pi = doca_gpu_dev_verbs_atomic_read<uint64_t,
-            DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_GPU>(&ep->sq_wqe_pi);
-    return pi + ep->sq_wqe_num - count;
-}
-
-UCS_F_DEVICE uint64_t uct_rc_mlx5_gda_reserv_wqe_thread(
-    uct_rc_gdaki_dev_ep_t *ep, unsigned count)
-{
-    /* Do not attempt to reserve if the available space is less than the
-     * requested count, to avoid starvation of threads trying to rollback the
-     * reservation with atomicCAS. */
-    uint64_t max_wqe_base = uct_rc_mlx5_gda_max_alloc_wqe_base(ep, count);
-    if (ep->sq_rsvd_index > max_wqe_base) {
+    /* Try to reserve optimistically */
+    int32_t prev = atomicAdd(&ep->avail_count, -(int32_t)count);
+    if (prev < (int32_t)count) {
+        /* Rollback */
+        atomicAdd(&ep->avail_count, count);
         return UCT_RC_GDA_RESV_WQE_NO_RESOURCE;
     }
 
-    uint64_t wqe_base = atomicAdd(reinterpret_cast<unsigned long long*>(
-                                          &ep->sq_rsvd_index),
-                                  static_cast<unsigned long long>(count));
-
-    /*
-     *  Attempt to reserve 'count' WQEs by atomically incrementing the reserved
-     *  index. If the reservation exceeds the available space in the work queue,
-     *  enter a rollback loop.
-     *
-     *  Rollback Logic:
-     *  - Calculate the next potential index (wqe_next) after attempting the
-     *    reservation.
-     *  - Use atomic CAS to check if the current reserved index matches wqe_next.
-     *    If it does, revert the reservation by resetting the reserved index to
-     *    wqe_base.
-     *  - A successful CAS indicates no other thread has modified the reserved
-     *    index, allowing the rollback to complete, and the function returns
-     *    UCT_RC_GDA_RESV_WQE_NO_RESOURCE to signal insufficient resources.
-     *  - If CAS fails, it means another thread has modified the reserved index.
-     *    The loop continues to reevaluate resource availability to determine if
-     *    the reservation can now be satisfied, possibly due to other operations
-     *    freeing up resources.
-     */
-    while (wqe_base > max_wqe_base) {
-        uint64_t wqe_next = wqe_base + count;
-        if (atomicCAS(reinterpret_cast<unsigned long long*>(&ep->sq_rsvd_index),
-                      wqe_next, wqe_base) == wqe_next) {
-            return UCT_RC_GDA_RESV_WQE_NO_RESOURCE;
-        }
-
-        max_wqe_base = uct_rc_mlx5_gda_max_alloc_wqe_base(ep, count);
-    }
-
-    return wqe_base;
+    /* We own count elements, now can safely increment the reserved index */
+    return atomicAdd(reinterpret_cast<unsigned long long*>(&ep->sq_rsvd_index),
+                     count);
 }
 
 template<ucs_device_level_t level>
 UCS_F_DEVICE void
 uct_rc_mlx5_gda_reserv_wqe(uct_rc_gdaki_dev_ep_t *ep, unsigned count,
                            unsigned lane_id, uint64_t &wqe_base)
 {
+    wqe_base = 0;
+
     if (lane_id == 0) {
         wqe_base = uct_rc_mlx5_gda_reserv_wqe_thread(ep, count);
     }
+
     if (level == UCS_DEVICE_LEVEL_WARP) {
         wqe_base = __shfl_sync(0xffffffff, wqe_base, 0);
     } else if (level == UCS_DEVICE_LEVEL_BLOCK) {
@@ -231,6 +193,12 @@ UCS_F_DEVICE void uct_rc_mlx5_gda_db(uct_rc_gdaki_dev_ep_t *ep,
             &ep->sq_lock);
 }
 
+UCS_F_DEVICE bool
+uct_rc_mlx5_gda_fc(const uct_rc_gdaki_dev_ep_t *ep, uint16_t wqe_idx)
+{
+    return (wqe_idx & ep->sq_fc_mask) == 1;
+}
+
 template<ucs_device_level_t level>
 UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_single(
         uct_rc_gdaki_dev_ep_t *ep, const uct_device_mem_element_t *tl_mem_elem,
@@ -240,23 +208,20 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_single(
         uint64_t add)
 {
     uct_rc_gda_completion_t *comp = &tl_comp->rc_gda;
-    unsigned cflag = 0;
+    unsigned cflag                = 0;
     uint64_t wqe_base;
-    uint64_t wqe_idx;
     unsigned lane_id;
     unsigned num_lanes;
-    uint32_t fc;
 
     uct_rc_mlx5_gda_exec_init<level>(lane_id, num_lanes);
     uct_rc_mlx5_gda_reserv_wqe<level>(ep, 1, lane_id, wqe_base);
     if (wqe_base == UCT_RC_GDA_RESV_WQE_NO_RESOURCE) {
         return UCS_ERR_NO_RESOURCE;
     }
 
-    fc = doca_gpu_dev_verbs_wqe_idx_inc_mask(ep->sq_wqe_pi, ep->sq_wqe_num / 2);
-    wqe_idx = wqe_base & 0xffff;
     if (lane_id == 0) {
-        if ((comp != nullptr) || (wqe_idx == fc)) {
+        uint16_t wqe_idx = (uint16_t)wqe_base;
+        if ((comp != nullptr) || uct_rc_mlx5_gda_fc(ep, wqe_idx)) {
             cflag = DOCA_GPUNETIO_MLX5_WQE_CTRL_CQ_UPDATE;
             if (comp != nullptr) {
                 comp->wqe_idx = wqe_base;
@@ -325,15 +290,13 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi(
     auto mem_list = reinterpret_cast<const uct_rc_gdaki_device_mem_element_t*>(
             tl_mem_list);
     uct_rc_gda_completion_t *comp = &tl_comp->rc_gda;
-
-    int count                 = mem_list_count;
-    int counter_index         = count - 1;
-    bool atomic               = false;
+    int count                     = mem_list_count;
+    int counter_index             = count - 1;
+    bool atomic                   = false;
     uint64_t wqe_idx;
     unsigned cflag;
     unsigned lane_id;
     unsigned num_lanes;
-    uint32_t fc;
     uint64_t wqe_base;
     size_t length;
     void *address;
@@ -357,7 +320,6 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi(
         return UCS_ERR_NO_RESOURCE;
     }
 
-    fc = doca_gpu_dev_verbs_wqe_idx_inc_mask(ep->sq_wqe_pi, ep->sq_wqe_num / 2);
     wqe_idx = doca_gpu_dev_verbs_wqe_idx_inc_mask(wqe_base, lane_id);
     for (uint32_t i = lane_id; i < count; i += num_lanes) {
         if (i == counter_index) {
@@ -379,7 +341,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi(
 
         cflag = 0;
         if (((comp != nullptr) && (i == count - 1)) ||
-            ((comp == nullptr) && (wqe_idx == fc))) {
+            ((comp == nullptr) && uct_rc_mlx5_gda_fc(ep, wqe_idx))) {
             cflag = DOCA_GPUNETIO_MLX5_WQE_CTRL_CQ_UPDATE;
             if (comp != nullptr) {
                 comp->wqe_idx = wqe_base;
@@ -420,13 +382,12 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi_partial(
     auto mem_list = reinterpret_cast<const uct_rc_gdaki_device_mem_element_t*>(
             tl_mem_list);
     uct_rc_gda_completion_t *comp = &tl_comp->rc_gda;
-    unsigned count            = mem_list_count;
-    bool atomic               = false;
+    unsigned count                = mem_list_count;
+    bool atomic                   = false;
     uint64_t wqe_idx;
     unsigned lane_id;
     unsigned num_lanes;
     unsigned cflag;
-    uint32_t fc;
     uint64_t wqe_base;
     size_t length;
     void *address;
@@ -451,7 +412,6 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi_partial(
         return UCS_ERR_NO_RESOURCE;
     }
 
-    fc = doca_gpu_dev_verbs_wqe_idx_inc_mask(ep->sq_wqe_pi, ep->sq_wqe_num / 2);
     wqe_idx = doca_gpu_dev_verbs_wqe_idx_inc_mask(wqe_base, lane_id);
     for (uint32_t i = lane_id; i < count; i += num_lanes) {
         if (i == mem_list_count) {
@@ -475,7 +435,7 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi_partial(
 
         cflag = 0;
         if (((comp != nullptr) && (i == count - 1)) ||
-            ((comp == nullptr) && (wqe_idx == fc))) {
+            ((comp == nullptr) && uct_rc_mlx5_gda_fc(ep, wqe_idx))) {
             cflag = DOCA_GPUNETIO_MLX5_WQE_CTRL_CQ_UPDATE;
             if (comp != nullptr) {
                 comp->wqe_idx = wqe_base;
@@ -502,20 +462,6 @@ UCS_F_DEVICE ucs_status_t uct_rc_mlx5_gda_ep_put_multi_partial(
     return UCS_INPROGRESS;
 }
 
-UCS_F_DEVICE uint16_t uct_rc_mlx5_gda_bswap16(uint16_t x)
-{
-    uint32_t ret;
-    asm volatile("{\n\t"
-                 ".reg .b32 mask;\n\t"
-                 ".reg .b32 ign;\n\t"
-                 "mov.b32 mask, 0x1;\n\t"
-                 "prmt.b32 %0, %1, ign, mask;\n\t"
-                 "}"
-                 : "=r"(ret)
-                 : "r"((uint32_t)x));
-    return ret;
-}
-
 UCS_F_DEVICE void
 uct_rc_mlx5_gda_qedump(const char *pfx, void *buff, ssize_t len)
 {
@@ -535,16 +481,17 @@ uct_rc_mlx5_gda_qedump(const char *pfx, void *buff, ssize_t len)
 
 UCS_F_DEVICE void uct_rc_mlx5_gda_progress_thread(uct_rc_gdaki_dev_ep_t *ep)
 {
-    void *cqe                = ep->cqe_daddr;
-    size_t cqe_num           = ep->cqe_num;
-    uint64_t cqe_idx         = ep->cqe_ci;
-    const size_t cqe_sz      = DOCA_GPUNETIO_VERBS_CQE_SIZE;
-    uint32_t idx             = cqe_idx & (cqe_num - 1);
-    void *curr_cqe           = (uint8_t*)cqe + idx * cqe_sz;
-    auto *cqe64              = reinterpret_cast<mlx5_cqe64*>(curr_cqe);
-    uint8_t op_owner;
-
-    op_owner = READ_ONCE(cqe64->op_own);
+    void *cqe        = ep->cqe_daddr;
+    size_t cqe_num   = ep->cqe_num;
+    uint64_t cqe_idx = ep->cqe_ci;
+    uint32_t idx     = cqe_idx & (cqe_num - 1);
+    void *curr_cqe   = (uint8_t*)cqe + (idx * DOCA_GPUNETIO_VERBS_CQE_SIZE);
+    auto *cqe64      = reinterpret_cast<mlx5_cqe64*>(curr_cqe);
+
+    /* Read last 3 fields with a single atomic operation */
+    uint32_t *data_ptr = (uint32_t *)&cqe64->wqe_counter;
+    uint32_t data      = READ_ONCE(*data_ptr);
+    uint8_t op_owner   = data >> 24;
     if ((op_owner & MLX5_CQE_OWNER_MASK) ^ !!(cqe_idx & cqe_num)) {
         return;
     }
@@ -555,16 +502,30 @@ UCS_F_DEVICE void uct_rc_mlx5_gda_progress_thread(uct_rc_gdaki_dev_ep_t *ep)
         return;
     }
 
-    uint8_t opcode   = op_owner >> DOCA_GPUNETIO_VERBS_MLX5_CQE_OPCODE_SHIFT;
-    uint16_t wqe_cnt = uct_rc_mlx5_gda_bswap16(cqe64->wqe_counter);
-    uint16_t wqe_idx = wqe_cnt & (ep->sq_wqe_num - 1);
+    uint8_t opcode    = op_owner >> DOCA_GPUNETIO_VERBS_MLX5_CQE_OPCODE_SHIFT;
+    uint32_t data_cpu = doca_gpu_dev_verbs_bswap32(data);
+    uint16_t wqe_cnt  = (data_cpu >> 16) & 0xffff;
+    uint16_t wqe_idx  = wqe_cnt & (ep->sq_wqe_num - 1);
 
     cuda::atomic_ref<uint64_t, cuda::thread_scope_device> pi_ref(ep->sq_wqe_pi);
-    uint64_t sq_wqe_pi = ep->sq_wqe_pi;
-    sq_wqe_pi          = ((wqe_cnt - sq_wqe_pi) & 0xffff) + sq_wqe_pi + 1;
+    uint64_t sq_wqe_pi = pi_ref.load(cuda::std::memory_order_relaxed);
+    uint64_t new_wqe_pi;
+
+    do {
+        /* Skip CQE if it's older than current producer index, could be already
+         * processed by another thread. This handles CQE wrap-around. */
+        if ((int16_t)(wqe_cnt - (uint16_t)sq_wqe_pi) < 0) {
+            return;
+        }
+
+        uint16_t completed_delta = wqe_cnt - (uint16_t)sq_wqe_pi;
+        new_wqe_pi               = sq_wqe_pi + completed_delta + 1;
+    } while (!pi_ref.compare_exchange_weak(sq_wqe_pi, new_wqe_pi,
+                                           cuda::std::memory_order_release,
+                                           cuda::std::memory_order_relaxed));
 
     if (opcode == MLX5_CQE_REQ) {
-        pi_ref.fetch_max(sq_wqe_pi);
+        atomicAdd(&ep->avail_count, (int32_t)(new_wqe_pi - sq_wqe_pi));
         return;
     }
 

src/uct/ib/mlx5/gdaki/gdaki_dev.h

@@ -30,6 +30,8 @@ typedef struct {
     uint32_t                     cqe_num;
     uint16_t                     sq_wqe_num;
     uint32_t                     sq_num;
+    uint16_t                     sq_fc_mask;
+    int32_t                      avail_count;
 } uct_rc_gdaki_dev_ep_t;
 
 

test/gtest/common/cuda.h

@@ -36,6 +36,11 @@ template<typename T> class device_result_ptr {
         return *m_ptr;
     }
 
+    T *operator->() const
+    {
+        return m_ptr.get();
+    }
+
     T *device_ptr()
     {
         T *device_ptr;