UCP/GTEST: Workaround KSM ODP failure - 2

openucx · Nov 27, 2024 · 6c564cc · 6c564cc
1 parent 6239d0d
commit 6c564cc
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c
@@ -104,6 +104,10 @@ ucs_config_field_t uct_ib_md_config_table[] = {
      "Use indirect atomic\n",
      ucs_offsetof(uct_ib_md_config_t, ext.enable_indirect_atomic), UCS_CONFIG_TYPE_BOOL},
 
+    {"ATOMIC_OFFSET", "y",
+     "Use atomic offset\n",
+     ucs_offsetof(uct_ib_md_config_t, ext.enable_atomic_offset), UCS_CONFIG_TYPE_BOOL},
+
     {"GID_INDEX", "auto",
      "Port GID index to use.",
      ucs_offsetof(uct_ib_md_config_t, ext.gid_index), UCS_CONFIG_TYPE_ULUNITS},

diff --git a/src/uct/ib/base/ib_md.h b/src/uct/ib/base/ib_md.h
@@ -81,6 +81,7 @@ typedef struct uct_ib_md_ext_config {
     int                      prefer_nearest_device; /**< Give priority for near
                                                          device */
     int                      enable_indirect_atomic; /** Enable indirect atomic */
+    int                      enable_atomic_offset;
 
     struct {
         int                  prefetch;     /**< Auto-prefetch non-blocking memory
@@ -333,7 +334,10 @@ uct_ib_md_is_flush_rkey_valid(uint32_t flush_rkey) {
 
 static UCS_F_ALWAYS_INLINE uint8_t uct_ib_md_get_atomic_mr_id(uct_ib_md_t *md)
 {
-    return md->flush_rkey >> 8;
+    if (md->config.enable_atomic_offset) {
+        return md->flush_rkey >> 8;
+    }
+    return 0;
 }
 
 void uct_ib_md_parse_relaxed_order(uct_ib_md_t *md,

diff --git a/test/gtest/ucp/test_ucp_rma.cc b/test/gtest/ucp/test_ucp_rma.cc
@@ -408,11 +408,11 @@ class test_ucp_rma_reg : public test_ucp_rma {
 };
 
 /* TODO temp workaround. Pending RM 4170682 fix */
-UCS_TEST_P(test_ucp_rma_reg, put_blocking, "IB_PCI_RELAXED_ORDERING?=n") {
+UCS_TEST_P(test_ucp_rma_reg, put_blocking, "IB_ATOMIC_OFFSET?=n") {
     test_reg(static_cast<send_func_t>(&test_ucp_rma::put_b));
 }
 
-UCS_TEST_P(test_ucp_rma_reg, put_nonblocking, "IB_PCI_RELAXED_ORDERING?=n") {
+UCS_TEST_P(test_ucp_rma_reg, put_nonblocking, "IB_ATOMIC_OFFSET?=n") {
     test_reg(static_cast<send_func_t>(&test_ucp_rma::put_nbi));
 }