Merge pull request #36 from Yiltan/LWPRHMEM-71-add-coop-groups

Add cooperative groups for sync collective

.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 .

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 */
@@ -85,7 +85,6 @@ IPCBackend::IPCBackend(MPI_Comm comm)
   setup_team_world();
 
   TeamInfo *tinfo = team_tracker.get_team_world()->tinfo_wrt_world;
-  default_context_proxy_ = IPCDefaultContextProxyT(this, tinfo);
 
   roc_shmem_collective_init();
 
@@ -93,8 +92,9 @@ IPCBackend::IPCBackend(MPI_Comm comm)
 
   teams_init();
 
+  default_context_proxy_ = IPCDefaultContextProxyT(this, tinfo);
+
   setup_ctxs();
-  
 }
 
 IPCBackend::~IPCBackend() {
@@ -108,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;
@@ -402,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)
    */
@@ -161,6 +165,10 @@ class IPCBackend : public Backend {
   */
   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()
@@ -259,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

@@ -50,6 +50,10 @@ __host__ IPCContext::IPCContext(Backend *b)
   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() {

src/ipc/context_ipc_device.hpp

@@ -27,6 +27,13 @@
 #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 {
 
 class IPCContext : public Context {
@@ -235,6 +242,10 @@ 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_{};
 

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) {

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;