From d1342d1655e20ffdadd4232ffdee99ed997b8262 Mon Sep 17 00:00:00 2001
From: Thomas Vegas <tvegas@nvidia.com>
Date: Thu, 26 Oct 2023 16:32:13 +0300
Subject: [PATCH] UCT/IB: Skip multi-thread memory registration for symmetric
 key

---
 src/uct/ib/mlx5/dv/ib_mlx5dv_md.c | 16 +++++++++++++---
 test/gtest/uct/ib/test_ib_md.cc   | 15 ++++++++++++++-
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c
index f86fd719bc5d..e2d2d36f8539 100644
--- a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c
+++ b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c
@@ -556,6 +556,8 @@ static void uct_ib_mlx5_devx_reg_symmetric(uct_ib_mlx5_md_t *md,
     uint32_t symmetric_rkey;
     ucs_status_t status;
 
+    ucs_assert((memh->super.flags & UCT_IB_MEM_MULTITHREADED) == 0);
+
     /* Best effort, only allocate in the range below the atomic keys. */
     while (md->smkey_index < md->super.mkey_by_name_reserve.size) {
         status = uct_ib_mlx5_devx_reg_ksm_data_contig(
@@ -588,6 +590,12 @@ static void uct_ib_mlx5_devx_reg_symmetric(uct_ib_mlx5_md_t *md,
               md->super.mkey_by_name_reserve.base + start);
 }
 
+static UCS_F_ALWAYS_INLINE int
+uct_ib_mlx5_devx_symmetric_rkey(const uct_ib_mlx5_md_t *md, unsigned flags)
+{
+    return (flags & UCT_MD_MEM_SYMMETRIC_RKEY) &&
+           (md->flags & UCT_IB_MLX5_MD_FLAG_MKEY_BY_NAME_RESERVE);
+}
 
 static ucs_status_t
 uct_ib_mlx5_devx_reg_mr(uct_ib_mlx5_md_t *md, uct_ib_mlx5_devx_mem_t *memh,
@@ -598,11 +606,14 @@ uct_ib_mlx5_devx_reg_mr(uct_ib_mlx5_md_t *md, uct_ib_mlx5_devx_mem_t *memh,
 {
     uint64_t access_flags = uct_ib_memh_access_flags(&md->super, &memh->super) &
                             access_mask;
+    unsigned flags        = UCT_MD_MEM_REG_FIELD_VALUE(params, flags,
+                                                       FIELD_FLAGS, 0);
     ucs_status_t status;
     uint32_t mkey;
 
     if ((length >= md->super.config.min_mt_reg) &&
-        !(access_flags & IBV_ACCESS_ON_DEMAND)) {
+        !(access_flags & IBV_ACCESS_ON_DEMAND) &&
+        !uct_ib_mlx5_devx_symmetric_rkey(md, flags)) {
         /* Verbs transports can issue atomic operations to the default key */
         status = uct_ib_mlx5_devx_reg_mt(md, address, length,
                                          (memh->super.flags &
@@ -672,8 +683,7 @@ uct_ib_mlx5_devx_mem_reg(uct_md_h uct_md, void *address, size_t length,
         goto err_memh_free;
     }
 
-    if ((flags & UCT_MD_MEM_SYMMETRIC_RKEY) &&
-        (md->flags & UCT_IB_MLX5_MD_FLAG_MKEY_BY_NAME_RESERVE)) {
+    if (uct_ib_mlx5_devx_symmetric_rkey(md, flags)) {
         uct_ib_mlx5_devx_reg_symmetric(md, memh, address);
     }
 
diff --git a/test/gtest/uct/ib/test_ib_md.cc b/test/gtest/uct/ib/test_ib_md.cc
index e7ab20876df5..0839727a46d4 100644
--- a/test/gtest/uct/ib/test_ib_md.cc
+++ b/test/gtest/uct/ib/test_ib_md.cc
@@ -27,6 +27,8 @@ class test_ib_md : public test_md
                                 uct_rkey_t *rkey_p = NULL);
     void check_smkeys(uct_rkey_t rkey1, uct_rkey_t rkey2);
 
+    void test_smkey_reg_atomic(void);
+
 private:
 #ifdef HAVE_MLX5_DV
     uint32_t m_mlx5_flags = 0;
@@ -197,7 +199,7 @@ void test_ib_md::check_smkeys(uct_rkey_t rkey1, uct_rkey_t rkey2)
               uct_ib_md_atomic_rkey(rkey2) - uct_ib_md_direct_rkey(rkey2));
 }
 
-UCS_TEST_P(test_ib_md, smkey_reg_atomic)
+void test_ib_md::test_smkey_reg_atomic(void)
 {
     static const size_t size = 8192;
     void *buffer;
@@ -222,4 +224,15 @@ UCS_TEST_P(test_ib_md, smkey_reg_atomic)
     ucs_mmap_free(buffer, size);
 }
 
+UCS_TEST_P(test_ib_md, smkey_reg_atomic)
+{
+    test_smkey_reg_atomic();
+}
+
+UCS_TEST_P(test_ib_md, smkey_reg_atomic_mt, "REG_MT_THRESH=1k",
+           "REG_MT_CHUNK=1k")
+{
+    test_smkey_reg_atomic();
+}
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_ib_md, ib)