Merge branch 'ROCm:develop' into ipc_parameterized_simple_tests_10-01-24

.gitignore

@@ -0,0 +1 @@
+build/

CMakeLists.txt

@@ -68,6 +68,7 @@ option(USE_FUNC_CALL "Force compiler to use function calls on library API" OFF)
 option(USE_SHARED_CTX "Request support for shared ctx between WG" OFF)
 option(USE_SINGLE_NODE "Enable single node support only." OFF)
 option(USE_HOST_SIDE_HDP_FLUSH "Use a polling thread to flush the HDP cache on the host." OFF)
+option(USE_COOPERATIVE_GROUPS "Use cooperative groups for internal syncronization" OFF)
 option(BUILD_FUNCTIONAL_TESTS "Build the functional tests" ON)
 option(BUILD_SOS_TESTS "Build the host-facing tests" OFF)
 option(BUILD_UNIT_TESTS "Build the unit tests" ON)

cmake/config.h.in

@@ -14,3 +14,4 @@
 #cmakedefine USE_FUNC_CALL
 #cmakedefine USE_SINGLE_NODE
 #cmakedefine USE_HOST_SIDE_HDP_FLUSH
+#cmakedefine USE_COOPERATIVE_GROUPS

scripts/build_configs/ipc_single_cg

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+
+if [ -z $1 ]
+then
+  install_path=~/rocshmem
+else
+  install_path=$1
+fi
+
+src_path=$(dirname "$(realpath $0)")/../../
+
+cmake \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=$install_path \
+    -DCMAKE_VERBOSE_MAKEFILE=OFF \
+    -DDEBUG=OFF \
+    -DPROFILE=OFF \
+    -DUSE_GPU_IB=OFF \
+    -DUSE_RO=OFF \
+    -DUSE_DC=OFF \
+    -DUSE_IPC=ON \
+    -DUSE_COHERENT_HEAP=ON \
+    -DUSE_THREADS=OFF \
+    -DUSE_WF_COAL=OFF \
+    -DUSE_SINGLE_NODE=ON \
+    -DUSE_HOST_SIDE_HDP_FLUSH=OFF \
+    -DUSE_COOPERATIVE_GROUPS=ON \
+    $src_path
+cmake --build . --parallel 8
+cmake --install .

scripts/unit_tests/driver.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+
+# Function to display help information
+function display_help {
+    echo "Usage:"
+    echo "  $0 binary_name all                     # Runs all standard tests"
+    echo "  $0 binary_name custom <ranks> <filter> # Runs custom test configuration"
+    echo
+    echo "Arguments:"
+    echo "  binary_name: Name of the binary to run."
+    echo "  all: Executes predefined test configurations."
+    echo "  custom: Executes a test with custom MPI ranks and GTest filter."
+    echo "  ranks: Number of MPI ranks (required for custom mode)."
+    echo "  filter: GTest filter string (required for custom mode)."
+    echo
+}
+
+# Validate number of arguments for each mode
+if [[ "$#" -lt 2 ]] || 
+   { [[ "$2" == "all" ]] && [[ "$#" -ne 2 ]]; } ||
+   { [[ "$2" == "custom" ]] && [[ "$#" -ne 4 ]]; }; then
+    display_help
+    exit 1
+fi
+
+binary_name=$1
+mode=$2
+timestamp=$(date "+%Y-%m-%d-%H:%M:%S")
+log_file="unit_tests_${timestamp}.log"
+
+# Function to execute mpirun command
+function run_mpirun {
+    local np=$1
+    local gtest_filter=$2
+    echo "mpirun -np $np $binary_name --gtest_filter='$gtest_filter'" | tee -a "$log_file"
+    mpirun -np "$np" "$binary_name" --gtest_filter="$gtest_filter" >> "$log_file" 2>&1
+}
+
+# Processing modes
+case $mode in
+    all)
+        run_mpirun 4 "-IPCImplSimpleCoarseTestFixture.*:IPCImplSimpleFineTestFixture.*"
+        run_mpirun 2 "IPCImplSimpleCoarseTestFixture.*:IPCImplSimpleFineTestFixture.*"
+        ;;
+    custom)
+        # Check if ranks is a positive integer
+        if [[ "$3" -le 1 ]]; then
+            echo "Error: 'ranks' must be a positive integer."
+            display_help
+            exit 1
+        fi
+        run_mpirun $3 $4
+        ;;
+    *)
+        echo "Error: Invalid mode '$mode'." | tee -a "$log_file"
+        display_help
+        exit 1
+        ;;
+esac
+
+echo "Tests Completed"
+echo "log file: '$log_file'"

src/assembly.hpp

@@ -135,12 +135,12 @@ NOWARN(-Wdeprecated-volatile,
             : "v"(src));
 #endif
         break;
+      case 8:
 #if defined(__gfx906__)
 #endif
 #if defined(__gfx908__)
 #endif
 #if defined(__gfx90a__)
-      case 8:
         asm volatile(
             "global_load_dwordx2 %0 %1 off glc slc \n"
             "s_waitcnt vmcnt(0)"

src/gpu_ib/backend_ib.cpp

@@ -133,8 +133,6 @@ __device__ bool GPUIBBackend::create_ctx(int64_t options,
     return false;
   }
   ctx_ = pop_result.value;
-
-  ctx_->dev_mtx_.shareable_ = static_cast<bool>(options >> 3);
   ctx->ctx_opaque = ctx_;
   return true;
 }
@@ -296,7 +294,8 @@ void GPUIBBackend::setup_default_ctx() {
   CHECK_HIP(hipGetSymbolAddress(reinterpret_cast<void **>(&symbol_address),
                                 HIP_SYMBOL(ROC_SHMEM_CTX_DEFAULT)));
 
-  roc_shmem_ctx_t ctx_default_host{default_ctx_, nullptr};
+  TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
+  roc_shmem_ctx_t ctx_default_host{default_ctx_, tinfo};
 
   hipStream_t stream;
   CHECK_HIP(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));

src/gpu_ib/segment_builder.cpp

@@ -22,6 +22,7 @@
 
 #include "segment_builder.hpp"
 
+#include "../util.hpp"
 #include "endian.hpp"
 
 namespace rocshmem {
@@ -35,9 +36,9 @@ __device__ SegmentBuilder::SegmentBuilder(uint64_t wqe_idx, void *base) {
 __device__ void SegmentBuilder::update_cntrl_seg(
     uint8_t opcode, uint16_t wqe_idx, uint32_t ctrl_qp_sq, uint64_t ctrl_sig,
     ConnectionImpl *connection_policy, bool zero_byte_rd) {
-  mlx5_wqe_ctrl_seg *ctrl_seg = &seg_ptr->ctrl_seg;
+  mlx5_wqe_ctrl_seg ctrl_seg;
 
-  ctrl_seg->opmod_idx_opcode = (opcode << 24) | (wqe_idx << 8);
+  ctrl_seg.opmod_idx_opcode = (opcode << 24) | (wqe_idx << 8);
 
   uint32_t DS = 2;
   if (zero_byte_rd == false) {
@@ -48,39 +49,43 @@ __device__ void SegmentBuilder::update_cntrl_seg(
 
   DS += connection_policy->wqeCntrlOffset();
 
-  ctrl_seg->qpn_ds = (DS << 24) | ctrl_qp_sq;
+  ctrl_seg.qpn_ds = (DS << 24) | ctrl_qp_sq;
 
-  ctrl_seg->signature = ctrl_sig;
+  ctrl_seg.signature = ctrl_sig;
 
-  ctrl_seg->fm_ce_se = ctrl_sig >> 24;
+  ctrl_seg.fm_ce_se = ctrl_sig >> 24;
 
-  ctrl_seg->imm = ctrl_sig >> 32;
+  ctrl_seg.imm = ctrl_sig >> 32;
+
+  memcpy(&seg_ptr->ctrl_seg, &ctrl_seg, sizeof(mlx5_wqe_ctrl_seg));
 
   seg_ptr++;
 }
 
 __device__ void SegmentBuilder::update_atomic_data_seg(uint64_t atomic_data,
                                                        uint64_t atomic_cmp) {
-  mlx5_wqe_atomic_seg *atomic_seg = &seg_ptr->atomic_seg;
+  mlx5_wqe_atomic_seg atomic_seg;
 
-  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg->swap_add),
+  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg.swap_add),
                     atomic_data);
 
-  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg->compare),
+  swap_endian_store(reinterpret_cast<uint64_t *>(&atomic_seg.compare),
                     atomic_cmp);
 
+  memcpy(&seg_ptr->atomic_seg, &atomic_seg, sizeof(mlx5_wqe_atomic_seg));
   seg_ptr++;
 }
 
 __device__ void SegmentBuilder::update_rdma_seg(uintptr_t *raddr,
                                                 uint32_t rkey) {
-  mlx5_wqe_raddr_seg *raddr_seg = &seg_ptr->raddr_seg;
+  mlx5_wqe_raddr_seg raddr_seg;
 
-  raddr_seg->rkey = rkey;
+  raddr_seg.rkey = rkey;
 
-  swap_endian_store(reinterpret_cast<uint64_t *>(&raddr_seg->raddr),
+  swap_endian_store(reinterpret_cast<uint64_t *>(&raddr_seg.raddr),
                     reinterpret_cast<uint64_t>(raddr));
 
+  memcpy(&seg_ptr->raddr_seg, &raddr_seg, sizeof(mlx5_wqe_raddr_seg));
   seg_ptr++;
 }
 
@@ -90,33 +95,36 @@ __device__ void SegmentBuilder::update_data_seg(uintptr_t *laddr, int32_t size,
     return;
   }
 
-  mlx5_wqe_data_seg *data_seg = &seg_ptr->data_seg;
+  mlx5_wqe_data_seg data_seg;
+  data_seg.lkey = lkey;
 
-  data_seg->lkey = lkey;
-
-  swap_endian_store(&data_seg->byte_count, size & 0x7FFFFFFFU);
-
-  swap_endian_store(reinterpret_cast<uint64_t *>(&data_seg->addr),
+  swap_endian_store(&data_seg.byte_count, size & 0x7FFFFFFFU);
+  swap_endian_store(reinterpret_cast<uint64_t *>(&data_seg.addr),
                     reinterpret_cast<uint64_t>(laddr));
 
+  memcpy(&seg_ptr->data_seg, &data_seg, sizeof(mlx5_wqe_data_seg));
   seg_ptr++;
 }
 
 __device__ void SegmentBuilder::update_inl_data_seg(uintptr_t *laddr,
                                                     int32_t size) {
-  mlx5_wqe_inl_data_seg *inl_data_seg = &seg_ptr->inl_data_seg;
+  mlx5_wqe_inl_data_seg inl_data_seg;
 
-  swap_endian_store(&inl_data_seg->byte_count, (size & 0x3FF) | 0x80000000);
+  swap_endian_store(&inl_data_seg.byte_count, (size & 0x3FF) | 0x80000000);
 
   // Assume fence HDP flush
   // TODO(khamidou): Rework fence interface to avoid this
+  size_t field_size{sizeof(mlx5_wqe_inl_data_seg)};
   if (!laddr) {
     uint8_t flush_val = 1;
-    memcpy(inl_data_seg + 1, &flush_val, sizeof(flush_val));
+    memcpy(&inl_data_seg + 1, &flush_val, sizeof(flush_val));
+    field_size += sizeof(flush_val);
   } else {
-    memcpy(inl_data_seg + 1, laddr, size);
+    memcpy(&inl_data_seg + 1, laddr, size);
+    field_size += size;
   }
 
+  memcpy(&seg_ptr->inl_data_seg, &inl_data_seg, field_size);
   seg_ptr++;
 }
 

src/ipc/backend_ipc.cpp

@@ -56,18 +56,18 @@ int get_ls_non_zero_bit(char *bitmask, int mask_length) {
 IPCBackend::IPCBackend(MPI_Comm comm)
     :  Backend() {
   type = BackendType::IPC_BACKEND;
-    
+
   if (auto maximum_num_contexts_str = getenv("ROC_SHMEM_MAX_NUM_CONTEXTS")) {
     std::stringstream sstream(maximum_num_contexts_str);
     sstream >> maximum_num_contexts_;
   }
 
   init_mpi_once(comm);
-    
+
   initIPC();
-    
+
   auto *bp{ipc_backend_proxy.get()};
-    
+
   bp->heap_ptr = &heap;
 
   /* Initialize the host interface */
@@ -82,16 +82,19 @@ IPCBackend::IPCBackend(MPI_Comm comm)
 
   allocate_atomic_region(&bp->atomic_ret, MAX_NUM_BLOCKS);
 
-  default_context_proxy_ = IPCDefaultContextProxyT(this);
-
   setup_team_world();
 
+  TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
+
   roc_shmem_collective_init();
 
+  setup_fence_buffer();
+
   teams_init();
 
+  default_context_proxy_ = IPCDefaultContextProxyT(this, tinfo);
+
   setup_ctxs();
-  
 }
 
 IPCBackend::~IPCBackend() {
@@ -105,7 +108,7 @@ IPCBackend::~IPCBackend() {
    * Free the atomic_ret array.
    */
   CHECK_HIP(hipFree(bp->atomic_ret->atomic_base_ptr));
-  
+
   // TODO(Avinash) Free g_ret
 
   // delete host_interface;
@@ -141,6 +144,8 @@ __device__ bool IPCBackend::create_ctx(int64_t options, roc_shmem_ctx_t *ctx) {
   ctx_ = pop_result.value;
 
   ctx->ctx_opaque = ctx_;
+
+  ctx_->tinfo = reinterpret_cast<TeamInfo *>(ctx->team_opaque);
   return true;
 }
 
@@ -287,6 +292,14 @@ void IPCBackend::teams_destroy() {
   free(reduced_bitmask_);
 }
 
+void IPCBackend::setup_fence_buffer() {
+  /*
+  * Allocate heap space for fence
+  */
+  fence_pool = reinterpret_cast<int *>(roc_shmem_malloc(
+      sizeof(int) * num_pes));
+}
+
 void IPCBackend::roc_shmem_collective_init() {
   /*
    * Allocate heap space for barrier_sync
@@ -389,4 +402,4 @@ void IPCBackend::teams_init() {
   NET_CHECK(MPI_Barrier(thread_comm));
 }
 
-}  // namespace rocshmem
+}  // namespace rocshmem

src/ipc/backend_ipc.hpp

@@ -32,6 +32,10 @@
 #include "ipc_context_proxy.hpp"
 #include "../ipc_policy.hpp"
 
+#ifndef USE_COOPERATIVE_GROUPS
+#include "../memory/notifier.hpp"
+#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */
+
 namespace rocshmem {
 
 class IPCBackend : public Backend {
@@ -104,7 +108,7 @@ class IPCBackend : public Backend {
                        TeamInfo *team_info_wrt_world, int num_pes,
                        int my_pe_in_new_team, MPI_Comm team_comm,
                        roc_shmem_team_t *new_team) override;
-  
+
   /**
    * @copydoc Backend::team_destroy(roc_shmem_team_t)
    */
@@ -156,6 +160,15 @@ class IPCBackend : public Backend {
    */
   void *pAta_pool{nullptr};
 
+  /**
+   * @brief Handle for raw memory for fence/quiet
+  */
+  int *fence_pool{nullptr};
+
+#ifndef USE_COOPERATIVE_GROUPS
+  NotifierProxy<HIPAllocator, detail::atomic::memory_scope_agent> notifier_{};
+#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */
+
  protected:
    /**
    * @copydoc Backend::dump_backend_stats()
@@ -203,6 +216,11 @@ class IPCBackend : public Backend {
    */
   void roc_shmem_collective_init();
 
+  /**
+   * @brief Allocate buffer for fence/quiet operation
+   */
+  void setup_fence_buffer();
+
  private:
   /**
    * @brief Proxy for the default context
@@ -249,4 +267,4 @@ class IPCBackend : public Backend {
 
 }  // namespace rocshmem
 
-#endif  // LIBRARY_SRC_IPC_BACKEND_HPP_
+#endif  // LIBRARY_SRC_IPC_BACKEND_HPP_

src/ipc/context_ipc_device.cpp

@@ -47,6 +47,13 @@ __host__ IPCContext::IPCContext(Backend *b)
   barrier_sync = backend->barrier_sync;
   g_ret = bp->g_ret;
   atomic_base_ptr = bp->atomic_ret->atomic_base_ptr;
+  fence_pool = backend->fence_pool;
+
+  orders_.store = detail::atomic::rocshmem_memory_order::memory_order_seq_cst;
+
+#ifndef USE_COOPERATIVE_GROUPS
+  notifier_ = backend->notifier_.get();
+#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */
 }
 
 __device__ void IPCContext::threadfence_system() {
@@ -85,12 +92,17 @@ __device__ void IPCContext::getmem_nbi(void *dest, const void *source,
 }
 
 __device__ void IPCContext::fence() {
+  for (int i{0}, j{tinfo->pe_start}; i < tinfo->size; i++, j += tinfo->stride) {
+    detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[j], 1, orders_);
+  }
 }
 
 __device__ void IPCContext::fence(int pe) {
+  detail::atomic::store<int, detail::atomic::memory_scope_system>(&fence_pool[pe], 1, orders_);
 }
 
 __device__ void IPCContext::quiet() {
+  fence();
 }
 
 __device__ void *IPCContext::shmem_ptr(const void *dest, int pe) {

src/ipc/context_ipc_device.hpp

@@ -24,6 +24,15 @@
 #define LIBRARY_SRC_IPC_CONTEXT_DEVICE_HPP_
 
 #include "../context.hpp"
+#include "../atomic.hpp"
+#include "../team.hpp"
+
+#ifdef USE_COOPERATIVE_GROUPS
+#include <hip/hip_cooperative_groups.h>
+namespace cg = cooperative_groups;
+#else
+#include "../memory/notifier.hpp"
+#endif /* USE_COOPERATIVE_GROUPS */
 
 namespace rocshmem {
 
@@ -232,6 +241,24 @@ class IPCContext : public Context {
 
   //Temporary scratchpad memory used by internal barrier algorithms.
   int64_t *barrier_sync{nullptr};
+
+#ifndef USE_COOPERATIVE_GROUPS
+  Notifier<detail::atomic::memory_scope_agent> *notifier_{nullptr};
+#endif /* NOT DEFINED: USE_COOPERATIVE_GROUPS */
+
+  //Struct defining memory ordering for atomic operations.
+  detail::atomic::rocshmem_memory_orders orders_{};
+
+  //Buffer to perform Atomic store to enforce memory ordering
+  int *fence_pool{nullptr};
+
+ public:
+  //TODO(Avinash):
+  //Make tinfo private variable, it requires changes to the context
+  //creation API in backend
+
+  //Team information for the team associated with the context
+  TeamInfo *tinfo{nullptr};
 };
 
 }  // namespace rocshmem

src/ipc/context_ipc_device_coll.cpp

@@ -35,6 +35,9 @@ __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start,
   if (pe == PE_start) {
     // Go through all PE offsets (except current offset = 0)
     // and wait until they all reach
+#if defined(__gfx90a__)
+    __threadfence_system();
+#endif /* __gfx90a__ */
     for (size_t i = 1; i < n_pes; i++) {
       wait_until(&pSync[i], ROC_SHMEM_CMP_EQ, flag_val);
       pSync[i] = ROC_SHMEM_SYNC_VALUE;
@@ -49,6 +52,9 @@ __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start,
     // Mark current PE offset as reached
     size_t pe_offset = (pe - PE_start) / stride;
     put_nbi(&pSync[pe_offset], &flag_val, 1, PE_start);
+#if defined(__gfx90a__)
+    __threadfence_system();
+#endif /* __gfx90a__ */
     wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val);
     pSync[0] = ROC_SHMEM_SYNC_VALUE;
     threadfence_system();
@@ -78,8 +84,15 @@ __device__ void IPCContext::internal_atomic_barrier(int pe, int PE_start,
 // Uses PE values that are relative to world
 __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
                                           int PE_size, int64_t *pSync) {
-  __syncthreads();
-  if (is_thread_zero_in_block()) {
+#ifdef USE_COOPERATIVE_GROUPS
+  cg::grid_group grid = cg::this_grid();
+  grid.sync();
+  if (0 == grid.thread_rank())
+#else
+  notifier_->sync();
+  if (0 == get_flat_id())
+#endif /* USE_COOPERATIVE_GROUPS */
+  {
     if (PE_size < 64) {
       internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
     } else {
@@ -87,7 +100,11 @@ __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
     }
   }
   __threadfence();
-  __syncthreads();
+#ifdef USE_COOPERATIVE_GROUPS
+  grid.sync();
+#else
+  notifier_->sync();
+#endif /* USE_COOPERATIVE_GROUPS */
 }
 
 __device__ void IPCContext::sync(roc_shmem_team_t team) {

src/ipc/ipc_context_proxy.hpp

@@ -41,10 +41,11 @@ class IPCDefaultContextProxy {
   /*
    * Placement new the memory which is allocated by proxy_
    */
-  explicit IPCDefaultContextProxy(IPCBackend* backend) : constructed_{true} {
+  explicit IPCDefaultContextProxy(IPCBackend* backend, TeamInfo *tinfo)
+  : constructed_{true} {
     auto ctx{proxy_.get()};
     new (ctx) IPCContext(reinterpret_cast<Backend*>(backend));
-    roc_shmem_ctx_t local{ctx, nullptr};
+    roc_shmem_ctx_t local{ctx, tinfo};
     set_internal_ctx(&local);
   }
 

src/reverse_offload/backend_ro.cpp

@@ -94,7 +94,9 @@ ROBackend::ROBackend(MPI_Comm comm)
 
   default_block_handle_proxy_ = DefaultBlockHandleProxyT(
       bp->g_ret, bp->atomic_ret, &queue_, &ipcImpl, hdp_proxy_.get());
-  default_context_proxy_ = DefaultContextProxyT(this);
+
+  TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
+  default_context_proxy_ = DefaultContextProxyT(this, tinfo);
 
   block_handle_proxy_ = BlockHandleProxyT(bp->g_ret, bp->atomic_ret, &queue_,
                                           &ipcImpl, hdp_proxy_.get());

src/reverse_offload/context_proxy.hpp

@@ -42,10 +42,11 @@ class DefaultContextProxy {
   /*
    * Placement new the memory which is allocated by proxy_
    */
-  explicit DefaultContextProxy(ROBackend* backend) : constructed_{true} {
+  explicit DefaultContextProxy(ROBackend* backend, TeamInfo *tinfo)
+  : constructed_{true} {
     auto ctx{proxy_.get()};
     new (ctx) ROContext(reinterpret_cast<Backend*>(backend), -1);
-    roc_shmem_ctx_t local{ctx, nullptr};
+    roc_shmem_ctx_t local{ctx, tinfo};
     set_internal_ctx(&local);
   }
 

src/roc_shmem_gpu.cpp

@@ -262,9 +262,9 @@ __device__ int roc_shmem_wg_ctx_create(long option, roc_shmem_ctx_t *ctx) {
   GPU_DPRINTF("Function: roc_shmem_ctx_create\n");
   bool result{true};
   if (get_flat_block_id() == 0) {
+    ctx->team_opaque = reinterpret_cast<TeamInfo *>(ROC_SHMEM_CTX_DEFAULT.team_opaque);
     device_backend_proxy->create_ctx(option, ctx);
     reinterpret_cast<Context *>(ctx->ctx_opaque)->setFence(option);
-    ctx->team_opaque = nullptr;
   }
   __syncthreads();
   return result == true ? 0 : -1;
@@ -279,11 +279,11 @@ __device__ int roc_shmem_wg_team_create_ctx(roc_shmem_team_t team, long options,
 
   bool result{true};
   if (get_flat_block_id() == 0) {
-    result = device_backend_proxy->create_ctx(options, ctx);
-    reinterpret_cast<Context *>(ctx->ctx_opaque)->setFence(options);
     Team *team_obj{get_internal_team(team)};
     TeamInfo *info_wrt_world = team_obj->tinfo_wrt_world;
     ctx->team_opaque = info_wrt_world;
+    result = device_backend_proxy->create_ctx(options, ctx);
+    reinterpret_cast<Context *>(ctx->ctx_opaque)->setFence(options);
   }
   __syncthreads();
 
@@ -412,7 +412,9 @@ __device__ void roc_shmem_ctx_fence(roc_shmem_ctx_t ctx) {
 __device__ void roc_shmem_ctx_fence(roc_shmem_ctx_t ctx, int pe) {
   GPU_DPRINTF("Function: roc_shmem_ctx_fence\n");
 
-  get_internal_ctx(ctx)->fence(pe);
+  int pe_in_world = translate_pe(ctx, pe);
+
+  get_internal_ctx(ctx)->fence(pe_in_world);
 }
 
 __device__ void roc_shmem_ctx_quiet(roc_shmem_ctx_t ctx) {

tests/functional_tests/sync_tester.cpp

@@ -83,9 +83,21 @@ void SyncTester::launchKernel(dim3 gridSize, dim3 blockSize, int loop,
   roc_shmem_team_split_strided(ROC_SHMEM_TEAM_WORLD, 0, 1, n_pes, nullptr, 0,
                                &team_sync_world_dup);
 
+#ifdef USE_COOPERATIVE_GROUPS
+  void* kernelParams[] = {(void*)&loop,
+                          (void*)&args.skip,
+                          (void*)&timer,
+                          (void*)&_type,
+                          (void*)&_shmem_context,
+                          (void*)&team_sync_world_dup};
+
+  CHECK_HIP(hipLaunchCooperativeKernel(SyncTest, gridSize, blockSize,
+                                       kernelParams, shared_bytes, stream));
+#else
   hipLaunchKernelGGL(SyncTest, gridSize, blockSize, shared_bytes, stream, loop,
                      args.skip, timer, _type, _shmem_context,
                      team_sync_world_dup);
+#endif /* USE_COOPERATIVE_GROUPS */
 
   num_msgs = (loop + args.skip) * gridSize.x;
   num_timed_msgs = loop;

tests/unit_tests/shmem_gtest.cpp

@@ -21,8 +21,27 @@
  *****************************************************************************/
 
 #include "gtest/gtest.h"
+#include <mpi.h>
 
 int main(int argc, char **argv) {
   ::testing::InitGoogleTest(&argc, argv);
-  return RUN_ALL_TESTS();
+  
+  int initialized;
+  MPI_Initialized(&initialized);
+  if (!initialized) {
+    int provided;
+    MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
+    if (provided != MPI_THREAD_MULTIPLE) {
+      std::cerr << "MPI_THREAD_MULTIPLE support disabled.\n";
+    }
+  }
+
+  int ret_val = RUN_ALL_TESTS();
+  
+  int finalized{0};
+  MPI_Finalized(&finalized);
+  if (!finalized) {
+    MPI_Finalize();
+  }
+  return ret_val;
 }