diff --git a/src/uct/ib/base/ib_md.c b/src/uct/ib/base/ib_md.c index 37af9045a7e..d09872d2539 100644 --- 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 index e9b679c883a..dbb524a9374 100644 --- 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 index fcbf1772201..fa62a2820fc 100644 --- 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(&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(&test_ucp_rma::put_nbi)); }