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]
This commit is contained in:
Avinash Kethineedi authored and GitHub committed 2025-04-02 11:58:55 -05:00
1 parent 0cde5f53dc
commit 426bbf525b
22 files changed
+508 -53

No files matched your search

+11 -6
View File
@@ -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;
}