Update Barrier_All and Sync_All APIs (#72)

* Fix deadlock in `rocshmem_ctx_wg_barrier_all` API in IPC conduit by adding per-context pSync buffers and context IDs
  - Added separate pSync buffers for each device context
  - Resolved deadlock when invoking barrier API (`rocshmem_ctx_wg_barrier_all`) concurrently from multiple contexts

* Update barrier_all functional tests for multi-context support

* Add thread, wavefront, and workgroup-level barrier_all APIs in IPC and RO conduits
  - Implemented barrier_all APIs at thread, wavefront, and workgroup granularity
  - Added support in both IPC and RO conduits
  - Updated functional tests to cover all `barrier_all` APIs

* Add thread, wavefront, and workgroup-level sync_all APIs in IPC and RO conduits
  - Implemented sync_all APIs for thread, wavefront, and workgroup scopes
  - Added support into both IPC and RO conduits
  - Added functional tests to cover all `sync_all` APIs

[ROCm/rocshmem commit: c652f58cef]

projects/rocshmem/src/ipc/backend_ipc.cpp

@@ -114,8 +114,9 @@ IPCBackend::~IPCBackend() {
 
 void IPCBackend::setup_ctxs() {
   CHECK_HIP(hipMalloc(&ctx_array, sizeof(IPCContext) * maximum_num_contexts_));
+  // 0th context is default context
   for (size_t i = 0; i < maximum_num_contexts_; i++) {
-    new (&ctx_array[i]) IPCContext(this);
+    new (&ctx_array[i]) IPCContext(this, i + 1);
     ctx_free_list.get()->push_back(ctx_array + i);
   }
 }
@@ -278,9 +279,10 @@ void IPCBackend::init_wrk_sync_buffer() {
   auto max_num_teams{team_tracker.get_max_num_teams()};
 
   /**
-   * size of barrier sync
+   * size of barrier sync for all the contexts
    */
-  Wrk_Sync_buffer_size_ += sizeof(*barrier_sync) * ROCSHMEM_BARRIER_SYNC_SIZE;
+  Wrk_Sync_buffer_size_ += sizeof(*barrier_sync) * ROCSHMEM_BARRIER_SYNC_SIZE *
+                           (maximum_num_contexts_ + 1);
 
   /**
    * Size of sync arrays for the teams
@@ -378,15 +380,18 @@ void IPCBackend::rocshmem_collective_init() {
   /*
    * Allocate heap space for barrier_sync
    */
-  size_t one_sync_size_bytes{sizeof(*barrier_sync)};
-  size_t sync_size_bytes{one_sync_size_bytes * ROCSHMEM_BARRIER_SYNC_SIZE};
+  size_t one_sync_size_bytes {sizeof(*barrier_sync)};
+  size_t total_sync_elems {
+    ROCSHMEM_BARRIER_SYNC_SIZE * (maximum_num_contexts_ + 1)};
+  size_t sync_size_bytes {one_sync_size_bytes * total_sync_elems};
+
   barrier_sync = reinterpret_cast<int64_t*>(temp_Wrk_Sync_buff_ptr_);
   temp_Wrk_Sync_buff_ptr_ += sync_size_bytes;
 
   /*
    * Initialize the barrier synchronization array with default values.
    */
-  for (int i = 0; i < num_pes; i++) {
+  for (int i = 0; i < total_sync_elems; i++) {
     barrier_sync[i] = ROCSHMEM_SYNC_VALUE;
   }
 

projects/rocshmem/src/ipc/context_ipc_device.cpp

@@ -36,15 +36,18 @@
 
 namespace rocshmem {
 
-__host__ IPCContext::IPCContext(Backend *b)
+__host__ IPCContext::IPCContext(Backend *b, unsigned int ctx_id)
     : Context(b, false) {
   IPCBackend *backend{static_cast<IPCBackend *>(b)};
   ipcImpl_.ipc_bases = b->ipcImpl.ipc_bases;
   ipcImpl_.shm_size = b->ipcImpl.shm_size;
 
-  barrier_sync = backend->barrier_sync;
+  size_t barrier_sync_offset = ctx_id * ROCSHMEM_BARRIER_SYNC_SIZE;
+
+  barrier_sync = backend->barrier_sync + barrier_sync_offset;
   fence_pool = backend->fence_pool;
   Wrk_Sync_buffer_bases_ = backend->get_wrk_sync_bases();
+  ctx_id_ = ctx_id;
 
   orders_.store = detail::atomic::rocshmem_memory_order::memory_order_seq_cst;
 }

projects/rocshmem/src/ipc/context_ipc_device.hpp

@@ -31,9 +31,9 @@ namespace rocshmem {
 
 class IPCContext : public Context {
  public:
-  __host__ IPCContext(Backend *b);
+  __host__ IPCContext(Backend *b, unsigned int ctx_id);
 
-  __device__ IPCContext(Backend *b);
+  __device__ IPCContext(Backend *b, unsigned int ctx_id);
 
   __device__ void threadfence_system();
 
@@ -61,11 +61,19 @@ class IPCContext : public Context {
 
   __device__ void barrier_all();
 
+  __device__ void barrier_all_wave();
+
+  __device__ void barrier_all_wg();
+
   __device__ void barrier(rocshmem_team_t team);
 
   __device__ void sync_all();
 
-  __device__ void sync(rocshmem_team_t team);
+  __device__ void sync_all_wave();
+
+  __device__ void sync_all_wg();
+
+  __device__ void sync_wg(rocshmem_team_t team);
 
   template <typename T>
   __device__ void p(T *dest, T value, int pe);
@@ -240,6 +248,12 @@ class IPCContext : public Context {
   __device__ void internal_sync(int pe, int PE_start, int stride, int PE_size,
                                 int64_t *pSync);
 
+  __device__ void internal_sync_wave(int pe, int PE_start, int stride, int PE_size,
+                                int64_t *pSync);
+
+  __device__ void internal_sync_wg(int pe, int PE_start, int stride, int PE_size,
+                                int64_t *pSync);
+
   __device__ void internal_direct_barrier(int pe, int PE_start, int stride,
                                           int n_pes, int64_t *pSync);
 
@@ -289,6 +303,11 @@ class IPCContext : public Context {
   */
   char **Wrk_Sync_buffer_bases_{nullptr};
 
+  /**
+   * @brief Decive context Id
+   */
+  unsigned int ctx_id_{};
+
  public:
   //TODO(Avinash):
   //Make tinfo private variable, it requires changes to the context

projects/rocshmem/src/ipc/context_ipc_device_coll.cpp

@@ -84,9 +84,29 @@ __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) {
+  if (PE_size < 64) {
+    internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
+  } else {
+    internal_atomic_barrier(pe, PE_start, stride, PE_size, pSync);
+  }
+}
+
+__device__ void IPCContext::internal_sync_wave(int pe, int PE_start, int stride,
+                                          int PE_size, int64_t *pSync) {
+  if (is_thread_zero_in_wave()) {
+    if (PE_size < 64) {
+      internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
+    } else {
+      internal_atomic_barrier(pe, PE_start, stride, PE_size, pSync);
+    }
+  }
+}
+
+// Uses PE values that are relative to world
+__device__ void IPCContext::internal_sync_wg(int pe, int PE_start, int stride,
+                                          int PE_size, int64_t *pSync) {
   __syncthreads();
   if (is_thread_zero_in_block()) {
     if (PE_size < 64) {
@@ -98,7 +118,7 @@ __device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
   __syncthreads();
 }
 
-__device__ void IPCContext::sync(rocshmem_team_t team) {
+__device__ void IPCContext::sync_wg(rocshmem_team_t team) {
   IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
 
   int pe = team_obj->my_pe_in_world;
@@ -107,18 +127,38 @@ __device__ void IPCContext::sync(rocshmem_team_t team) {
   int pe_size = team_obj->num_pes;
   long *p_sync = team_obj->barrier_pSync;
 
-  internal_sync(pe, pe_start, pe_stride, pe_size, p_sync);
+  internal_sync_wg(pe, pe_start, pe_stride, pe_size, p_sync);
 }
 
 __device__ void IPCContext::sync_all() {
   internal_sync(my_pe, 0, 1, num_pes, barrier_sync);
 }
 
+__device__ void IPCContext::sync_all_wave() {
+  internal_sync_wave(my_pe, 0, 1, num_pes, barrier_sync);
+}
+
+__device__ void IPCContext::sync_all_wg() {
+  internal_sync_wg(my_pe, 0, 1, num_pes, barrier_sync);
+}
+
 __device__ void IPCContext::barrier_all() {
+  quiet();
+  sync_all();
+}
+
+__device__ void IPCContext::barrier_all_wave() {
+  if (is_thread_zero_in_wave()) {
+    quiet();
+  }
+  sync_all_wave();
+}
+
+__device__ void IPCContext::barrier_all_wg() {
   if (is_thread_zero_in_block()) {
     quiet();
   }
-  sync_all();
+  sync_all_wg();
   __syncthreads();
 }
 
@@ -134,7 +174,7 @@ __device__ void IPCContext::barrier(rocshmem_team_t team) {
   if (is_thread_zero_in_block()) {
     quiet();
   }
-  internal_sync(pe, pe_start, pe_stride, pe_size, p_sync);
+  internal_sync_wg(pe, pe_start, pe_stride, pe_size, p_sync);
   __syncthreads();
 }
 

projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp

@@ -468,7 +468,7 @@ __device__ void IPCContext::internal_broadcast(T *dst, const T *src, int nelems,
   }
 
   // Synchronize on completion of broadcast
-  internal_sync(my_pe, pe_start, stride, pe_size, p_sync);
+  internal_sync_wg(my_pe, pe_start, stride, pe_size, p_sync);
 }
 
 template <typename T>
@@ -497,7 +497,7 @@ __device__ void IPCContext::alltoall_linear(rocshmem_team_t team, T *dst,
     quiet();
   }
   // wait until everyone has obtained their designated data
-  internal_sync(my_pe, pe_start, stride, pe_size, pSync);
+  internal_sync_wg(my_pe, pe_start, stride, pe_size, pSync);
 }
 
 template <typename T>
@@ -527,7 +527,7 @@ __device__ void IPCContext::fcollect_linear(rocshmem_team_t team, T *dst,
     quiet();
   }
   // wait until everyone has obtained their designated data
-  internal_sync(my_pe, pe_start, stride, pe_size, pSync);
+  internal_sync_wg(my_pe, pe_start, stride, pe_size, pSync);
 }
 
 // Block/wave functions

projects/rocshmem/src/ipc/ipc_context_proxy.hpp

@@ -45,7 +45,7 @@ class IPCDefaultContextProxy {
                                   size_t num_elems = 1)
   : constructed_{true}, proxy_{num_elems} {
     auto ctx{proxy_.get()};
-    new (ctx) IPCContext(reinterpret_cast<Backend*>(backend));
+    new (ctx) IPCContext(reinterpret_cast<Backend*>(backend), 0);
     ctx->tinfo = tinfo;
     rocshmem_ctx_t local{ctx, tinfo};
     set_internal_ctx(&local);