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
This commit is contained in:
Avinash Kethineedi authored and GitHub committed 2025-04-08 11:25:31 -05:00
1 parent c652f58cef
commit dc61bca066
16 files changed
+347 -67

No files matched your search

+8
View File
@@ -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>
+52
View File
@@ -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();
}