From abec29bd6a0d85838fc52c495aec98e484017395 Mon Sep 17 00:00:00 2001 From: avinashkethineedi Date: Mon, 28 Oct 2024 22:10:18 +0000 Subject: [PATCH] Update all_reduce algorithm to use internal put/get functions for updating pWrk and pSync arrays * Change log_stride calcualtions to stride calculations * Update all_reduce example code to use team based interface --- examples/rocshmem_allreduce_test.cc | 22 ++++++--------- src/ipc/context_ipc_device.hpp | 12 ++++---- src/ipc/context_ipc_tmpl_device.hpp | 31 ++++++++++----------- tests/functional_tests/tester_arguments.cpp | 10 +++---- 4 files changed, 35 insertions(+), 40 deletions(-) diff --git a/examples/rocshmem_allreduce_test.cc b/examples/rocshmem_allreduce_test.cc index 57e8571aa5..d9a2a581ff 100644 --- a/examples/rocshmem_allreduce_test.cc +++ b/examples/rocshmem_allreduce_test.cc @@ -25,8 +25,8 @@ using namespace rocshmem; -__global__ void allreduce_test(int *source, int *dest, int* pWork, long *pSync, size_t nelem) -{ +__global__ void allreduce_test(int *source, int *dest, size_t nelem, + roc_shmem_team_t team) { __shared__ roc_shmem_ctx_t ctx; int64_t ctx_type = 0; @@ -34,7 +34,7 @@ __global__ void allreduce_test(int *source, int *dest, int* pWork, long *pSync, roc_shmem_wg_ctx_create(ctx_type, &ctx); int num_pes = roc_shmem_ctx_n_pes(ctx); - roc_shmem_ctx_int_sum_wg_to_all(ctx, dest, source, nelem, 0, 0, num_pes, pWork, pSync); + roc_shmem_ctx_int_sum_wg_to_all(ctx, team, dest, source, nelem); roc_shmem_ctx_quiet(ctx); __syncthreads(); @@ -97,18 +97,16 @@ int main (int argc, char **argv) result[i] = -1; } - size_t p_wrk_size = ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE; - int *pWrk = (int *)roc_shmem_malloc(p_wrk_size * sizeof(int)); + roc_shmem_team_t team_reduce_world_dup; + team_reduce_world_dup = ROC_SHMEM_TEAM_INVALID; + roc_shmem_team_split_strided(ROC_SHMEM_TEAM_WORLD, 0, 1, npes, nullptr, 0, + &team_reduce_world_dup); - size_t p_sync_size = ROC_SHMEM_REDUCE_SYNC_SIZE; - long *pSync = (long *)roc_shmem_malloc(p_sync_size * sizeof(long)); - for (int i = 0; i < p_sync_size; i++) { - pSync[i] = ROC_SHMEM_SYNC_VALUE; - } CHECK_HIP(hipDeviceSynchronize()); int threadsPerBlock=256; - allreduce_test<<>>(source, result, pWrk, pSync, nelem); + allreduce_test<<>>(source, result, + nelem, team_reduce_world_dup); CHECK_HIP(hipDeviceSynchronize()); bool pass = check_recvbuf(result, npes, rank, nelem); @@ -117,8 +115,6 @@ int main (int argc, char **argv) roc_shmem_free(source); roc_shmem_free(result); - roc_shmem_free(pWrk); - roc_shmem_free(pSync); roc_shmem_finalize(); return 0; diff --git a/src/ipc/context_ipc_device.hpp b/src/ipc/context_ipc_device.hpp index dcb9571833..6983fb0d3b 100644 --- a/src/ipc/context_ipc_device.hpp +++ b/src/ipc/context_ipc_device.hpp @@ -128,11 +128,6 @@ class IPCContext : public Context { __device__ T amo_fetch_cas(void *dst, T value, T cond, int pe); // Collectives - template - __device__ void to_all(T *dest, const T *source, int nreduce, int PE_start, - int logPE_stride, int PE_size, T *pWrk, - long *pSync); // NOLINT(runtime/int) - template __device__ void to_all(roc_shmem_team_t team, T *dest, const T *source, int nreduce); @@ -213,7 +208,12 @@ class IPCContext : public Context { char* g_ret; //internal functions used by collective operations - template + template + __device__ void internal_to_all(T *dest, const T *source, int nreduce, int PE_start, + int stride, int PE_size, T *pWrk, + long *pSync); // NOLINT(runtime/int) + + template __device__ void internal_put_broadcast(T *dst, const T *src, int nelems, int pe_root, int PE_start, int logPE_stride, int PE_size); // NOLINT(runtime/int) diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp index 406367c72e..f55abf0451 100644 --- a/src/ipc/context_ipc_tmpl_device.hpp +++ b/src/ipc/context_ipc_tmpl_device.hpp @@ -164,11 +164,10 @@ __device__ void compute_reduce(T *src, T *dst, int size, int wg_id, template __device__ void IPCContext::internal_direct_allreduce( - T *dst, const T *src, int nelems, int PE_start, int logPE_stride, + T *dst, const T *src, int nelems, int PE_start, int stride, int PE_size, T *pWrk, long *pSync) { // NOLINT(runtime/int) - int stride = 1 << logPE_stride; int finish = PE_start + stride * PE_size; int pe = my_pe; @@ -183,12 +182,12 @@ __device__ void IPCContext::internal_direct_allreduce( for (int i = PE_start; i < finish; i += stride) { if (i != pe) { - putmem_nbi_wg(&pWrk[pe * nelems], reinterpret_cast(src), + internal_putmem_wg(&pWrk[pe * nelems], reinterpret_cast(src), nelems * sizeof(T), i); if (is_thread_zero_in_block()) { fence(); - put_nbi(&pSync[pe], &flag_val, 1, i); + internal_putmem(&pSync[pe], &flag_val, sizeof(*pSync), i); } } } @@ -278,7 +277,7 @@ __device__ void IPCContext::internal_direct_allreduce( template __device__ void IPCContext::internal_ring_allreduce( T *dst, const T *src, int nelems, [[maybe_unused]] int PE_start, - [[maybe_unused]] int logPE_stride, [[maybe_unused]] int PE_size, T *pWrk, + [[maybe_unused]] int stride, [[maybe_unused]] int PE_size, T *pWrk, long *pSync, // NOLINT(runtime/int) int n_seg, int seg_size, int chunk_size) { int off_seg, off_send, off_recv; @@ -300,7 +299,7 @@ __device__ void IPCContext::internal_ring_allreduce( off_send = (((my_pe + 1 - iter + 2 * num_pes) % num_pes) * chunk_size); off_recv = (((my_pe - iter + 2 * num_pes) % num_pes) * chunk_size); - putmem_nbi_wg(reinterpret_cast(&pWrk[off_send]), + internal_putmem_wg(reinterpret_cast(&pWrk[off_send]), reinterpret_cast(&dst[off_send + off_seg]), chunk_size * sizeof(T), send_pe); @@ -308,7 +307,7 @@ __device__ void IPCContext::internal_ring_allreduce( fence(); wait_val = seg + 100; - put_nbi(&pSync[iter], &wait_val, 1, send_pe); + internal_putmem(&pSync[iter], &wait_val, sizeof(*pSync), send_pe); #if defined(__gfx90a__) __threadfence_system(); #endif /* __gfx90a__ */ @@ -329,7 +328,7 @@ __device__ void IPCContext::internal_ring_allreduce( if (is_thread_zero_in_block()) { fence(); wait_val = seg + 100; - put_nbi(&pSync[iter], &wait_val, 1, send_pe); + internal_putmem(&pSync[iter], &wait_val, sizeof(*pSync), send_pe); #if defined(__gfx90a__) __threadfence_system(); #endif /* __gfx90a__ */ @@ -354,19 +353,19 @@ __device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest, /** * Ensure that the stride is a multiple of 2 for GPU_IB. */ - int log_pe_stride = static_cast(team_obj->tinfo_wrt_world->log_stride); + int stride = team_obj->tinfo_wrt_world->stride; int pe_start = team_obj->tinfo_wrt_world->pe_start; int pe_size = team_obj->tinfo_wrt_world->size; long *p_sync = team_obj->barrier_pSync; T *pWrk = reinterpret_cast(team_obj->pWrk); - to_all(dest, source, nreduce, pe_start, log_pe_stride, pe_size, pWrk, + internal_to_all(dest, source, nreduce, pe_start, stride, pe_size, pWrk, p_sync); } template -__device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, - int PE_start, int logPE_stride, +__device__ void IPCContext::internal_to_all(T *dest, const T *source, int nreduce, + int PE_start, int stride, int PE_size, T *pWrk, long *pSync) { // NOLINT(runtime/int) size_t direct_pWrk = num_pes * nreduce; @@ -376,7 +375,7 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, size_t provided_pSync = ROC_SHMEM_REDUCE_SYNC_SIZE; if (provided_pWrk >= direct_pWrk && provided_pSync >= direct_pSync) { - internal_direct_allreduce(dest, source, nreduce, PE_start, logPE_stride, + internal_direct_allreduce(dest, source, nreduce, PE_start, stride, PE_size, pWrk, pSync); } else { if (ring_pSync <= ROC_SHMEM_REDUCE_SYNC_SIZE) { @@ -395,7 +394,7 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, n_seg = 1; } internal_ring_allreduce(dest, source, nreduce, PE_start, - logPE_stride, PE_size, pWrk, pSync, n_seg, + stride, PE_size, pWrk, pSync, n_seg, seg_size, chunk_size); if (n_seg_up > n_seg) { T *p_dst = (dest + (n_seg * seg_size)); @@ -403,7 +402,7 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, int p_count = nreduce - (n_seg * seg_size); int p_chunk = p_count / num_pes; - internal_ring_allreduce(p_dst, p_src, p_count, PE_start, logPE_stride, + internal_ring_allreduce(p_dst, p_src, p_count, PE_start, stride, PE_size, pWrk, pSync, 1, (p_chunk * num_pes), p_chunk); if ((p_chunk * num_pes) < p_count) { @@ -412,7 +411,7 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, p_dst += (p_chunk * num_pes); const T *p_src2 = p_src + (p_chunk * num_pes); - internal_direct_allreduce(p_dst, p_src2, p_count, PE_start, logPE_stride, + internal_direct_allreduce(p_dst, p_src2, p_count, PE_start, stride, PE_size, pWrk, pSync); } } diff --git a/tests/functional_tests/tester_arguments.cpp b/tests/functional_tests/tester_arguments.cpp index 1167fafaeb..87f73dc21b 100644 --- a/tests/functional_tests/tester_arguments.cpp +++ b/tests/functional_tests/tester_arguments.cpp @@ -136,11 +136,11 @@ void TesterArguments::get_rocshmem_arguments() { myid = roc_shmem_my_pe(); TestType type = (TestType)algorithm; - if ((type != ReductionTestType) && (type != BarrierAllTestType) && - (type != SyncAllTestType) && (type != SyncTestType) && - (type != BroadcastTestType) && (type != AllToAllTestType) && - (type != FCollectTestType) && (type != TeamReductionTestType) && - (type != TeamBroadcastTestType) && (type != PingAllTestType)) { + if ((type != BarrierAllTestType) && (type != SyncAllTestType) && + (type != SyncTestType) && (type != BroadcastTestType) && + (type != AllToAllTestType) && (type != FCollectTestType) && + (type != TeamReductionTestType) && (type != TeamBroadcastTestType) && + (type != PingAllTestType)) { if (numprocs != 2) { if (myid == 0) { std::cerr << "This test requires exactly two processes, we have "