Merge pull request #10581 from yosefe/topic/uct-dc-fix-segfault-in-uct-dc

UCT/DC: Fix segfault in uct_dc_mlx5_ep_fence

Author: yosefe

src/uct/ib/mlx5/dc/dc_mlx5.c

@@ -368,7 +368,8 @@ static void uct_ib_mlx5dv_dci_qp_init_attr(uct_ib_qp_init_attr_t *qp_attr,
 
 ucs_status_t uct_dc_mlx5_iface_create_dci(uct_dc_mlx5_iface_t *iface,
                                           uct_dci_index_t dci_index,
-                                          uint8_t num_dci_channels)
+                                          uint8_t num_dci_channels,
+                                          uct_dc_dci_t **dci_p)
 {
     uct_ib_iface_t *ib_iface   = &iface->super.super.super;
     uct_ib_mlx5_qp_attr_t attr = {};
@@ -455,11 +456,6 @@ ucs_status_t uct_dc_mlx5_iface_create_dci(uct_dc_mlx5_iface_t *iface,
         goto err_qp;
     }
 
-    ucs_assertv_always(!uct_dc_mlx5_is_dci_valid(
-                               uct_dc_mlx5_iface_dci(iface, dci_index)),
-                       "iface=%p dci_index=%d", iface, dci_index);
-    ucs_array_elem(&iface->tx.dcis, dci_index) = dci;
-
     if (uct_dc_mlx5_iface_is_policy_shared(iface)) {
         ucs_arbiter_group_init(&dci->arb_group);
     } else {
@@ -477,6 +473,7 @@ ucs_status_t uct_dc_mlx5_iface_create_dci(uct_dc_mlx5_iface_t *iface,
 
     uct_rc_txqp_available_set(&dci->txqp, dci->txwq.bb_max);
 
+    *dci_p = dci;
     return UCS_OK;
 
 err:
@@ -759,11 +756,10 @@ static void uct_dc_mlx5_iface_dci_pool_destroy(uct_dc_mlx5_dci_pool_t *dci_pool)
     ucs_array_cleanup_dynamic(&dci_pool->stack);
 }
 
-void uct_dc_mlx5_destroy_dci(uct_dc_mlx5_iface_t *iface, uint16_t dci_index)
+void uct_dc_mlx5_destroy_dci(uct_dc_mlx5_iface_t *iface, uct_dc_dci_t *dci)
 {
     uct_ib_mlx5_md_t *md = ucs_derived_of(iface->super.super.super.super.md,
                                           uct_ib_mlx5_md_t);
-    uct_dc_dci_t *dci    = uct_dc_mlx5_iface_dci(iface, dci_index);
 
     uct_rc_txqp_cleanup(&iface->super.super, &dci->txqp);
     uct_ib_mlx5_destroy_qp(md, &dci->txwq.super);
@@ -774,17 +770,16 @@ void uct_dc_mlx5_destroy_dci(uct_dc_mlx5_iface_t *iface, uint16_t dci_index)
     uct_ib_mlx5_qp_mmio_cleanup(&dci->txwq.super, dci->txwq.reg);
 
     ucs_free(dci);
-    ucs_array_elem(&iface->tx.dcis, dci_index) = NULL;
 }
 
 static void uct_dc_mlx5_iface_dcis_destroy(uct_dc_mlx5_iface_t *iface)
 {
-    uct_dc_dci_t *dci UCS_V_UNUSED;
     uct_dci_index_t dci_index;
     uint8_t pool_index;
+    uct_dc_dci_t *dci;
 
     uct_dc_iface_for_each_valid_dci(dci_index, dci, iface) {
-        uct_dc_mlx5_destroy_dci(iface, dci_index);
+        uct_dc_mlx5_destroy_dci(iface, dci);
     }
 
     for (pool_index = 0; pool_index < iface->tx.num_dci_pools; pool_index++) {
@@ -885,19 +880,13 @@ uct_dc_mlx5_iface_init_dcis_array(uct_dc_mlx5_iface_t *iface,
         return status;
     }
 
-    ucs_array_length(&iface->tx.dcis) = 0;
-
-    status = uct_dc_mlx5_iface_create_dci(iface, 0, 1);
+    status = uct_dc_mlx5_iface_create_dci(iface, 0, 1, &dci);
     if (status != UCS_OK) {
         return status;
     }
 
-    dci = uct_dc_mlx5_iface_dci(iface, 0);
-
     iface->tx.bb_max = dci->txwq.bb_max;
-    uct_dc_mlx5_destroy_dci(iface, 0);
-
-    ucs_array_length(&iface->tx.dcis) = 0;
+    uct_dc_mlx5_destroy_dci(iface, dci);
 
     return UCS_OK;
 }
@@ -1603,7 +1592,7 @@ static UCS_CLASS_INIT_FUNC(uct_dc_mlx5_iface_t, uct_md_h tl_md, uct_worker_h wor
     self->tx.ndci   = uct_dc_mlx5_iface_is_hw_dcs(self) ? 1 : config->ndci;
 
     if (config->dcis_initial_capacity == UCS_ULUNITS_AUTO) {
-        max_capacity = uct_dc_mlx5_iface_is_dci_rand(self) ? SIZE_MAX : 1;
+        max_capacity = uct_dc_mlx5_iface_is_dci_rand(self) ? SIZE_MAX : 0;
     } else {
         max_capacity = config->dcis_initial_capacity;
     }

src/uct/ib/mlx5/dc/dc_mlx5.h

@@ -412,7 +412,8 @@ void uct_dc_mlx5_iface_reset_dci(uct_dc_mlx5_iface_t *iface,
 
 ucs_status_t uct_dc_mlx5_iface_create_dci(uct_dc_mlx5_iface_t *iface,
                                           uct_dci_index_t dci_index,
-                                          uint8_t num_dci_channels);
+                                          uint8_t num_dci_channels,
+                                          uct_dc_dci_t **dci_p);
 
 ucs_status_t uct_dc_mlx5_iface_resize_and_fill_dcis(uct_dc_mlx5_iface_t *iface,
                                                     uint16_t size);
@@ -429,7 +430,7 @@ uct_dc_mlx5_dci_pool_get_or_create(uct_dc_mlx5_iface_t *iface,
 uint32_t
 uct_dc_mlx5_dci_config_hash(const uct_dc_mlx5_dci_config_t *dci_config);
 
-void uct_dc_mlx5_destroy_dci(uct_dc_mlx5_iface_t *iface, uint16_t dci_index);
+void uct_dc_mlx5_destroy_dci(uct_dc_mlx5_iface_t *iface, uct_dc_dci_t *dci);
 
 static UCS_F_ALWAYS_INLINE uint8_t uct_dc_mlx5_is_dci_valid(const uct_dc_dci_t *dci)
 {
@@ -439,6 +440,9 @@ static UCS_F_ALWAYS_INLINE uint8_t uct_dc_mlx5_is_dci_valid(const uct_dc_dci_t *
 static UCS_F_ALWAYS_INLINE uct_dc_dci_t *
 uct_dc_mlx5_iface_dci(uct_dc_mlx5_iface_t *iface, uct_dci_index_t dci_index)
 {
+    ucs_assertv(dci_index < ucs_array_length(&iface->tx.dcis),
+                "dci_index=%d dcis_length=%d", dci_index,
+                ucs_array_length(&iface->tx.dcis));
     return ucs_array_elem(&iface->tx.dcis, dci_index);
 }
 

src/uct/ib/mlx5/dc/dc_mlx5_ep.c

@@ -273,10 +273,17 @@ ucs_status_t uct_dc_mlx5_ep_fence(uct_ep_h tl_ep, unsigned flags)
 {
     uct_dc_mlx5_iface_t *iface = ucs_derived_of(tl_ep->iface, uct_dc_mlx5_iface_t);
     uct_dc_mlx5_ep_t *ep       = ucs_derived_of(tl_ep, uct_dc_mlx5_ep_t);
+    uct_dc_dci_t *dci;
 
-    return uct_rc_ep_fence(tl_ep,
-                           &uct_dc_mlx5_iface_dci(iface, ep->dci)->txwq.fi,
-                           ep->dci != UCT_DC_MLX5_EP_NO_DCI);
+    if (ep->dci != UCT_DC_MLX5_EP_NO_DCI) {
+        dci = uct_dc_mlx5_iface_dci(iface, ep->dci);
+        ucs_assertv(uct_dc_mlx5_is_dci_valid(dci), "iface=%p ep=%p dci=%d",
+                    iface, ep, ep->dci);
+        return uct_rc_ep_fence(tl_ep, &dci->txwq.fi);
+    }
+
+    UCT_TL_EP_STAT_FENCE(&ep->super);
+    return UCS_OK;
 }
 
 static ucs_status_t UCS_F_ALWAYS_INLINE

src/uct/ib/mlx5/dc/dc_mlx5_ep.h

@@ -365,17 +365,21 @@ uct_dc_mlx5_dci_pool_init_dci(uct_dc_mlx5_iface_t *iface, uint8_t pool_index,
         num_channels = iface->tx.num_dci_channels;
     }
 
-    status = uct_dc_mlx5_iface_create_dci(iface, dci_index, num_channels);
+    status = uct_dc_mlx5_iface_create_dci(iface, dci_index, num_channels, &dci);
     if (status != UCS_OK) {
         ucs_error("iface %p: failed to create dci %u at pool %u", iface,
                   dci_index, pool_index);
         goto err;
     }
 
-    dci             = uct_dc_mlx5_iface_dci(iface, dci_index);
     dci->path_index = pool->config.path_index;
     dci->pool_index = pool_index;
 
+    ucs_assertv_always(!uct_dc_mlx5_is_dci_valid(
+                               uct_dc_mlx5_iface_dci(iface, dci_index)),
+                       "iface=%p dci_index=%d", iface, dci_index);
+    ucs_array_elem(&iface->tx.dcis, dci_index) = dci;
+
     status = uct_dc_mlx5_iface_dci_connect(iface, dci);
     if (status != UCS_OK) {
         ucs_error("iface %p: failed to connect dci %u at pool %u", iface,
@@ -393,7 +397,8 @@ uct_dc_mlx5_dci_pool_init_dci(uct_dc_mlx5_iface_t *iface, uint8_t pool_index,
     return UCS_OK;
 
 err_destroy:
-    uct_dc_mlx5_destroy_dci(iface, dci_index);
+    ucs_array_elem(&iface->tx.dcis, dci_index) = NULL;
+    uct_dc_mlx5_destroy_dci(iface, dci);
 err:
     return status;
 }

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

@@ -520,7 +520,7 @@ ucs_status_t uct_rc_mlx5_base_ep_fence(uct_ep_h tl_ep, unsigned flags)
 {
     uct_rc_mlx5_base_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_mlx5_base_ep_t);
 
-    return uct_rc_ep_fence(tl_ep, &ep->tx.wq.fi, 1);
+    return uct_rc_ep_fence(tl_ep, &ep->tx.wq.fi);
 }
 
 void uct_rc_mlx5_base_ep_post_check(uct_ep_h tl_ep)

src/uct/ib/rc/base/rc_ep.h

@@ -519,14 +519,14 @@ uct_rc_ep_fm(uct_rc_iface_t *iface, uct_ib_fence_info_t* fi, int flag)
     return fence;
 }
 
-static UCS_F_ALWAYS_INLINE ucs_status_t
-uct_rc_ep_fence(uct_ep_h tl_ep, uct_ib_fence_info_t* fi, int fence)
+static UCS_F_ALWAYS_INLINE ucs_status_t uct_rc_ep_fence(uct_ep_h tl_ep,
+                                                        uct_ib_fence_info_t *fi)
 {
     uct_rc_iface_t *iface = ucs_derived_of(tl_ep->iface, uct_rc_iface_t);
 
     /* in case if fence is requested and enabled by configuration
      * we need to schedule fence for next RDMA operation */
-    if (fence && (iface->config.fence_mode != UCT_RC_FENCE_MODE_NONE)) {
+    if (iface->config.fence_mode != UCT_RC_FENCE_MODE_NONE) {
         fi->fence_beat = iface->tx.fi.fence_beat - 1;
     }
 

src/uct/ib/rc/verbs/rc_verbs_ep.c

@@ -516,7 +516,7 @@ ucs_status_t uct_rc_verbs_ep_fence(uct_ep_h tl_ep, unsigned flags)
 {
     uct_rc_verbs_ep_t *ep = ucs_derived_of(tl_ep, uct_rc_verbs_ep_t);
 
-    return uct_rc_ep_fence(tl_ep, &ep->fi, 1);
+    return uct_rc_ep_fence(tl_ep, &ep->fi);
 }
 
 void uct_rc_verbs_ep_post_check(uct_ep_h tl_ep)

test/gtest/uct/ib/test_dc.cc

@@ -186,6 +186,8 @@ UCS_TEST_P(test_dc, dcs_single) {
                            0;
 
     EXPECT_EQ(UCT_DC_MLX5_EP_NO_DCI, ep->dci);
+    status = uct_ep_fence(m_e1->ep(0), 0);
+    EXPECT_UCS_OK(status);
     status = uct_ep_am_short(m_e1->ep(0), 0, 0, NULL, 0);
     EXPECT_UCS_OK(status);
     /* dci 0 must be assigned to the ep */