From 3535ce8c0aae9e396f99ffa05b2f7b9a953d8316 Mon Sep 17 00:00:00 2001 From: Yiltan Date: Fri, 31 Oct 2025 10:26:44 -0400 Subject: [PATCH] Alltoall linear parallel Optimization (#303) [ROCm/rocshmem commit: 8dd2112ec8e90d29e154990279651fdc630c1e6a] --- .../rocshmem/src/gda/context_gda_tmpl_device.hpp | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp b/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp index d119dd5f6c..c21bef8942 100644 --- a/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp +++ b/projects/rocshmem/src/gda/context_gda_tmpl_device.hpp @@ -618,14 +618,21 @@ __device__ void GDAContext::alltoall_linear(rocshmem_team_t team, T *dst, long *pSync = team_obj->alltoall_pSync; int my_pe_in_team = team_obj->my_pe; + int wf_id = get_flat_block_id() / WF_SIZE; + int wf_count = (int) ceil((double)get_flat_block_size() / (double)WF_SIZE); + bool wf_leader = 0 == get_active_lane_num(); + // Have each PE put their designated data to the other PEs - for (int j = 0; j < pe_size; j++) { + for (int j = wf_id; j < pe_size; j+= wf_count) { int dest_pe = team_obj->get_pe_in_world(j); - put_nbi_wg(&dst[my_pe_in_team * nelems], &src[j * nelems], nelems, dest_pe); + put_nbi_wave(&dst[my_pe_in_team * nelems], &src[j * nelems], nelems, dest_pe); } - if (is_thread_zero_in_block()) { - quiet(); + + for (int j = wf_id; j < pe_size; j+= wf_count) { + int dest_pe = team_obj->get_pe_in_world(j); + pe_quiet(dest_pe); } + // wait until everyone has obtained their designated data internal_sync_wg(my_pe, pe_start, stride, pe_size, pSync); }