Update Barrier and Sync APIs (#73)

* Add thread, wavefront, and workgroup-level `barrier` APIs in IPC and RO conduits; remove collectives on default context
 - Implemented `barrier` APIs for thread, wavefront, and workgroup scopes
 - Added support into both IPC and RO conduits
 - Added functional tests to cover all `barrier` APIs
 - Removed collective operations on default context

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

* update naming convention for context-based `barrier` APIs

src/backend_bc.cpp

@@ -132,6 +132,9 @@ void Backend::dump_stats() {
   printf("BarrierAll %llu\n", device_stats.getStat(NUM_BARRIER_ALL));
   printf("WAVE_BarrierAll %llu\n", device_stats.getStat(NUM_BARRIER_ALL_WAVE));
   printf("WG_BarrierAll %llu\n", device_stats.getStat(NUM_BARRIER_ALL_WG));
+  printf("Barrier %llu\n", device_stats.getStat(NUM_BARRIER));
+  printf("WAVE_Barrier %llu\n", device_stats.getStat(NUM_BARRIER_WAVE));
+  printf("WG_Barrier %llu\n", device_stats.getStat(NUM_BARRIER_WG));
   printf("Wait Until %llu\n", device_stats.getStat(NUM_WAIT_UNTIL));
   printf("Wait Until Any %llu\n", device_stats.getStat(NUM_WAIT_UNTIL_ANY));
   printf("Wait Until All %llu\n", device_stats.getStat(NUM_WAIT_UNTIL_ALL));
@@ -157,6 +160,9 @@ void Backend::dump_stats() {
   printf("SyncAll %llu\n", device_stats.getStat(NUM_SYNC_ALL));
   printf("WAVE_SyncAll %llu\n", device_stats.getStat(NUM_SYNC_ALL_WAVE));
   printf("WG_SyncAll %llu\n", device_stats.getStat(NUM_SYNC_ALL_WG));
+  printf("Sync %llu\n", device_stats.getStat(NUM_SYNC));
+  printf("WAVE_Sync %llu\n", device_stats.getStat(NUM_SYNC_WAVE));
+  printf("WG_Sync %llu\n", device_stats.getStat(NUM_SYNC_WG));
 
   const auto& host_stats{globalHostStats};
   printf("HOST STATS\n");

src/context.hpp

@@ -143,12 +143,20 @@ class Context {
 
   __device__ void barrier(rocshmem_team_t team);
 
+  __device__ void barrier_wave(rocshmem_team_t team);
+
+  __device__ void barrier_wg(rocshmem_team_t team);
+
   __device__ void sync_all();
 
   __device__ void sync_all_wave();
 
   __device__ void sync_all_wg();
 
+  __device__ void sync(rocshmem_team_t team);
+
+  __device__ void sync_wave(rocshmem_team_t team);
+
   __device__ void sync_wg(rocshmem_team_t team);
 
   template <typename T>

src/context_device.cpp

@@ -161,11 +161,23 @@ __device__ void Context::barrier_all_wg() {
 }
 
 __device__ void Context::barrier(rocshmem_team_t team) {
-  ctxStats.incStat(NUM_BARRIER_ALL);
+  ctxStats.incStat(NUM_BARRIER);
 
   DISPATCH(barrier(team));
 }
 
+__device__ void Context::barrier_wave(rocshmem_team_t team) {
+  ctxStats.incStat(NUM_BARRIER_WAVE);
+
+  DISPATCH(barrier_wave(team));
+}
+
+__device__ void Context::barrier_wg(rocshmem_team_t team) {
+  ctxStats.incStat(NUM_BARRIER_WG);
+
+  DISPATCH(barrier_wg(team));
+}
+
 __device__ void Context::sync_all() {
   ctxStats.incStat(NUM_SYNC_ALL);
 
@@ -184,8 +196,20 @@ __device__ void Context::sync_all_wg() {
   DISPATCH(sync_all_wg());
 }
 
+__device__ void Context::sync(rocshmem_team_t team) {
+  ctxStats.incStat(NUM_SYNC);
+
+  DISPATCH(sync(team));
+}
+
+__device__ void Context::sync_wave(rocshmem_team_t team) {
+  ctxStats.incStat(NUM_SYNC_WAVE);
+
+  DISPATCH(sync_wave(team));
+}
+
 __device__ void Context::sync_wg(rocshmem_team_t team) {
-  ctxStats.incStat(NUM_SYNC_ALL_WG);
+  ctxStats.incStat(NUM_SYNC_WG);
 
   DISPATCH(sync_wg(team));
 }

src/ipc/context_ipc_device.hpp

@@ -67,12 +67,20 @@ class IPCContext : public Context {
 
   __device__ void barrier(rocshmem_team_t team);
 
+  __device__ void barrier_wave(rocshmem_team_t team);
+
+  __device__ void barrier_wg(rocshmem_team_t team);
+
   __device__ void sync_all();
 
   __device__ void sync_all_wave();
 
   __device__ void sync_all_wg();
 
+  __device__ void sync(rocshmem_team_t team);
+
+  __device__ void sync_wave(rocshmem_team_t team);
+
   __device__ void sync_wg(rocshmem_team_t team);
 
   template <typename T>

src/ipc/context_ipc_device_coll.cpp

@@ -118,6 +118,30 @@ __device__ void IPCContext::internal_sync_wg(int pe, int PE_start, int stride,
   __syncthreads();
 }
 
+__device__ void IPCContext::sync(rocshmem_team_t team) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = team_obj->tinfo_wrt_world->stride;
+  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);
+}
+
+__device__ void IPCContext::sync_wave(rocshmem_team_t team) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = team_obj->tinfo_wrt_world->stride;
+  int pe_size = team_obj->num_pes;
+  long *p_sync = team_obj->barrier_pSync;
+
+  internal_sync_wave(pe, pe_start, pe_stride, pe_size, p_sync);
+}
+
 __device__ void IPCContext::sync_wg(rocshmem_team_t team) {
   IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
 
@@ -171,6 +195,34 @@ __device__ void IPCContext::barrier(rocshmem_team_t team) {
   int pe_size = team_obj->num_pes;
   long *p_sync = team_obj->barrier_pSync;
 
+  quiet();
+  internal_sync(pe, pe_start, pe_stride, pe_size, p_sync);
+}
+
+__device__ void IPCContext::barrier_wave(rocshmem_team_t team) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = team_obj->tinfo_wrt_world->stride;
+  int pe_size = team_obj->num_pes;
+  long *p_sync = team_obj->barrier_pSync;
+
+  if (is_thread_zero_in_wave()) {
+    quiet();
+  }
+  internal_sync_wave(pe, pe_start, pe_stride, pe_size, p_sync);
+}
+
+__device__ void IPCContext::barrier_wg(rocshmem_team_t team) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = team_obj->tinfo_wrt_world->stride;
+  int pe_size = team_obj->num_pes;
+  long *p_sync = team_obj->barrier_pSync;
+
   if (is_thread_zero_in_block()) {
     quiet();
   }

src/reverse_offload/context_ro_device.cpp

@@ -193,6 +193,22 @@ __device__ void ROContext::barrier_all_wg() {
 }
 
 __device__ void ROContext::barrier(rocshmem_team_t team) {
+  ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
+  build_queue_element(RO_NET_BARRIER, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
+                      nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
+                      true, get_status_flag(), is_default_ctx);
+}
+
+__device__ void ROContext::barrier_wave(rocshmem_team_t team) {
+  ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
+  if (is_thread_zero_in_wave()) {
+    build_queue_element(RO_NET_BARRIER, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
+                        nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
+                        true, get_status_flag(), is_default_ctx);
+  }
+}
+
+__device__ void ROContext::barrier_wg(rocshmem_team_t team) {
   ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
   if (is_thread_zero_in_block()) {
     build_queue_element(RO_NET_BARRIER, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
@@ -225,6 +241,22 @@ __device__ void ROContext::sync_all_wg() {
   __syncthreads();
 }
 
+__device__ void ROContext::sync(rocshmem_team_t team) {
+  ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
+  build_queue_element(RO_NET_SYNC, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
+                      nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
+                      true, get_status_flag(), is_default_ctx);
+}
+
+__device__ void ROContext::sync_wave(rocshmem_team_t team) {
+  ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
+  if (is_thread_zero_in_wave()) {
+    build_queue_element(RO_NET_SYNC, nullptr, nullptr, 0, 0, 0, 0, 0, nullptr,
+                        nullptr, team_obj->mpi_comm, ro_net_win_id, block_handle,
+                        true, get_status_flag(), is_default_ctx);
+  }
+}
+
 __device__ void ROContext::sync_wg(rocshmem_team_t team) {
   ROTeam *team_obj = reinterpret_cast<ROTeam *>(team);
   if (is_thread_zero_in_block()) {

src/reverse_offload/context_ro_device.hpp

@@ -71,12 +71,20 @@ class ROContext : public Context {
 
   __device__ void barrier(rocshmem_team_t team);
 
+  __device__ void barrier_wave(rocshmem_team_t team);
+
+  __device__ void barrier_wg(rocshmem_team_t team);
+
   __device__ void sync_all();
 
   __device__ void sync_all_wave();
 
   __device__ void sync_all_wg();
 
+  __device__ void sync(rocshmem_team_t team);
+
+  __device__ void sync_wave(rocshmem_team_t team);
+
   __device__ void sync_wg(rocshmem_team_t team);
 
   template <typename T>

src/rocshmem_gpu.cpp

@@ -588,14 +588,22 @@ __device__ void rocshmem_ctx_wg_barrier_all(rocshmem_ctx_t ctx) {
   get_internal_ctx(ctx)->barrier_all_wg();
 }
 
-__device__ void rocshmem_wg_barrier_all() {
-  rocshmem_ctx_wg_barrier_all(ROCSHMEM_CTX_DEFAULT);
-}
-
-__device__ void rocshmem_barrier(rocshmem_team_t team) {
+__device__ void rocshmem_ctx_barrier(rocshmem_ctx_t ctx, rocshmem_team_t team) {
   GPU_DPRINTF("Function: rocshmem_barrier\n");
 
-  get_internal_ctx(ROCSHMEM_CTX_DEFAULT)->barrier(team);
+  get_internal_ctx(ctx)->barrier(team);
+}
+
+__device__ void rocshmem_ctx_wave_barrier(rocshmem_ctx_t ctx, rocshmem_team_t team) {
+  GPU_DPRINTF("Function: rocshmem_wave_barrier\n");
+
+  get_internal_ctx(ctx)->barrier_wave(team);
+}
+
+__device__ void rocshmem_ctx_wg_barrier(rocshmem_ctx_t ctx, rocshmem_team_t team) {
+  GPU_DPRINTF("Function: rocshmem_wg_barrier\n");
+
+  get_internal_ctx(ctx)->barrier_wg(team);
 }
 
 __device__ void rocshmem_ctx_sync_all(rocshmem_ctx_t ctx) {
@@ -604,39 +612,37 @@ __device__ void rocshmem_ctx_sync_all(rocshmem_ctx_t ctx) {
   get_internal_ctx(ctx)->sync_all();
 }
 
-__device__ void rocshmem_sync_all() {
-  rocshmem_ctx_sync_all(ROCSHMEM_CTX_DEFAULT);
-}
-
 __device__ void rocshmem_ctx_wave_sync_all(rocshmem_ctx_t ctx) {
   GPU_DPRINTF("Function: rocshmem_ctx_wave_sync_all\n");
 
   get_internal_ctx(ctx)->sync_all_wave();
 }
 
-__device__ void rocshmem_wave_sync_all() {
-  rocshmem_ctx_wave_sync_all(ROCSHMEM_CTX_DEFAULT);
-}
-
 __device__ void rocshmem_ctx_wg_sync_all(rocshmem_ctx_t ctx) {
   GPU_DPRINTF("Function: rocshmem_ctx_wg_sync_all\n");
 
   get_internal_ctx(ctx)->sync_all_wg();
 }
 
-__device__ void rocshmem_wg_sync_all() {
-  rocshmem_ctx_wg_sync_all(ROCSHMEM_CTX_DEFAULT);
-}
-
-__device__ void rocshmem_ctx_wg_team_sync(rocshmem_ctx_t ctx,
+__device__ void rocshmem_ctx_team_sync(rocshmem_ctx_t ctx,
                                            rocshmem_team_t team) {
   GPU_DPRINTF("Function: rocshmem_ctx_sync_all\n");
 
   get_internal_ctx(ctx)->sync_wg(team);
 }
 
-__device__ void rocshmem_wg_team_sync(rocshmem_team_t team) {
-  rocshmem_ctx_wg_team_sync(ROCSHMEM_CTX_DEFAULT, team);
+__device__ void rocshmem_ctx_wave_team_sync(rocshmem_ctx_t ctx,
+  rocshmem_team_t team) {
+GPU_DPRINTF("Function: rocshmem_ctx_wave_sync_all\n");
+
+get_internal_ctx(ctx)->sync_wg(team);
+}
+
+__device__ void rocshmem_ctx_wg_team_sync(rocshmem_ctx_t ctx,
+  rocshmem_team_t team) {
+GPU_DPRINTF("Function: rocshmem_ctx_wg_sync_all\n");
+
+get_internal_ctx(ctx)->sync_wg(team);
 }
 
 __device__ int rocshmem_ctx_n_pes(rocshmem_ctx_t ctx) {

src/stats.hpp

@@ -45,6 +45,9 @@ enum rocshmem_stats {
   NUM_BARRIER_ALL,
   NUM_BARRIER_ALL_WAVE,
   NUM_BARRIER_ALL_WG,
+  NUM_BARRIER,
+  NUM_BARRIER_WAVE,
+  NUM_BARRIER_WG,
   NUM_WAIT_UNTIL,
   NUM_WAIT_UNTIL_ANY,
   NUM_WAIT_UNTIL_ALL,
@@ -74,6 +77,9 @@ enum rocshmem_stats {
   NUM_SYNC_ALL,
   NUM_SYNC_ALL_WAVE,
   NUM_SYNC_ALL_WG,
+  NUM_SYNC,
+  NUM_SYNC_WAVE,
+  NUM_SYNC_WG,
   NUM_BROADCAST,
   NUM_PUT_WG,
   NUM_PUT_NBI_WG,