From ba21cb7b85729f4aa839dfde9ae9a101502ee7e7 Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Fri, 27 Sep 2024 21:25:41 +0000 Subject: [PATCH 1/6] ipc/to_all: add direct allreduce algorithm add a simple version of an allreduce algorithm as a starting point. --- src/gpu_ib/context_ib_tmpl_device.hpp | 72 +------------------ src/ipc/context_ipc_device.hpp | 6 ++ src/ipc/context_ipc_tmpl_device.hpp | 84 +++++++++++++++++++++- src/roc_shmem_calc.hpp | 100 ++++++++++++++++++++++++++ 4 files changed, 189 insertions(+), 73 deletions(-) create mode 100644 src/roc_shmem_calc.hpp diff --git a/src/gpu_ib/context_ib_tmpl_device.hpp b/src/gpu_ib/context_ib_tmpl_device.hpp index a3f463a342..c575cb0832 100644 --- a/src/gpu_ib/context_ib_tmpl_device.hpp +++ b/src/gpu_ib/context_ib_tmpl_device.hpp @@ -29,80 +29,10 @@ #include "gpu_ib_team.hpp" #include "queue_pair.hpp" #include "../util.hpp" +#include "../roc_shmem_calc.hpp" namespace rocshmem { -// clang-format off -NOWARN(-Wunused-parameter, -template -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - static_assert(true, "Unimplemented gpu_ib collective."); - } -}; -) -// clang-format on - -/****************************************************************************** - ************************** TEMPLATE SPECIALIZATIONS ************************** - *****************************************************************************/ -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] += src[i]; - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] = max(dst[i], src[i]); - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] = min(dst[i], src[i]); - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] *= src[i]; - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] &= src[i]; - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] |= src[i]; - } -}; - -template <> -struct OpWrap { - template - __device__ static void Calc(T *src, T *dst, int i) { - dst[i] ^= src[i]; - } -}; - template __device__ void compute_reduce(T *src, T *dst, int size, int wg_id, int wg_size) { diff --git a/src/ipc/context_ipc_device.hpp b/src/ipc/context_ipc_device.hpp index 2a7aaebd62..6bf8885d6c 100644 --- a/src/ipc/context_ipc_device.hpp +++ b/src/ipc/context_ipc_device.hpp @@ -239,6 +239,12 @@ class IPCContext : public Context { __device__ void internal_atomic_barrier(int pe, int PE_start, int stride, int n_pes, int64_t *pSync); + template + __device__ void internal_direct_allreduce(T *dst, const T *src, + int nelems, int PE_start, int + logPE_stride, int PE_size, + T *pWrk, long *pSync); + //internal functions used by collectives routines to write/read to //work/sync buffers __device__ void internal_putmem(void *dest, const void *source, diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp index 91bdbd45e7..33bb91a1f1 100644 --- a/src/ipc/context_ipc_tmpl_device.hpp +++ b/src/ipc/context_ipc_tmpl_device.hpp @@ -28,6 +28,7 @@ #include "context_ipc_device.hpp" #include "../util.hpp" #include "ipc_team.hpp" +#include "../roc_shmem_calc.hpp" namespace rocshmem { @@ -152,11 +153,88 @@ __device__ T IPCContext::amo_fetch_cas(void *dest, T value, T cond, int pe) { } // Collectives +template +__device__ void compute_reduce(T *src, T *dst, int size, int wg_id, + int wg_size) { + for (size_t i = wg_id; i < size; i += wg_size) { + OpWrap::Calc(src, dst, i); + } + __syncthreads(); +} + +template +__device__ void IPCContext::internal_direct_allreduce( + T *dst, const T *src, int nelems, int PE_start, int logPE_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; + + int wg_id = get_flat_block_id(); + int wg_size = get_flat_block_size(); + int64_t flag_val = 1; + + for (int i = wg_id; i < nelems; i += wg_size) { + dst[i] = src[i]; + } + __syncthreads(); + + for (int i = PE_start; i < finish; i += stride) { + if (i != pe) { + putmem_nbi_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); + } + } + } + threadfence_system(); + __syncthreads(); + + // Do the compute and pSync reset in parallel. + for (int i = PE_start; i < finish; i += stride) { + if (i != pe) { + // Wait for leader thread to see that the buffer is ready. + if (is_thread_zero_in_block()) { + wait_until(&pSync[i], ROC_SHMEM_CMP_EQ, flag_val); + } + __syncthreads(); + + T *ptr = &pWrk[i * nelems]; + compute_reduce(ptr, dst, nelems, wg_id, wg_size); + threadfence_system(); + } + } + + __syncthreads(); + + for (int i = wg_id; i < num_pes; i += wg_size) { + pSync[i] = ROC_SHMEM_SYNC_VALUE; + } + threadfence_system(); + __syncthreads(); +} + template __device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest, const T *source, int nreduce) { - //to_all(dest, source, nreduce, pe_start, log_pe_stride, pe_size, pWrk, - // p_sync); + IPCTeam *team_obj = reinterpret_cast(team); + + /** + * 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 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, + p_sync); } template @@ -164,6 +242,8 @@ __device__ void IPCContext::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) + internal_direct_allreduce(dest, source, nreduce, PE_start, logPE_stride, + PE_size, pWrk, pSync); } template diff --git a/src/roc_shmem_calc.hpp b/src/roc_shmem_calc.hpp new file mode 100644 index 0000000000..5420cde3ea --- /dev/null +++ b/src/roc_shmem_calc.hpp @@ -0,0 +1,100 @@ +/****************************************************************************** + * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + *****************************************************************************/ + +#ifndef LIBRARY_SRC_ROC_SHMEM_CALC_HPP_ +#define LIBRARY_SRC_ROC_SHMEM_CALC_HPP_ + +namespace rocshmem { + +// clang-format off +NOWARN(-Wunused-parameter, +template +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + static_assert(true, "Unimplemented ipc collective."); + } +}; +) +// clang-format on + +/****************************************************************************** + ************************** TEMPLATE SPECIALIZATIONS ************************** + *****************************************************************************/ +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] += src[i]; + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] = max(dst[i], src[i]); + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] = min(dst[i], src[i]); + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] *= src[i]; + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] &= src[i]; + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] |= src[i]; + } +}; + +template <> +struct OpWrap { + template + __device__ static void Calc(T *src, T *dst, int i) { + dst[i] ^= src[i]; + } +}; + +} +#endif // LIBRARY_SRC_ROC_SHMEM_CALC_HPP_ From 1fbb89bc73a06fe4bc268539c34f36de5cdc3994 Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Mon, 7 Oct 2024 20:49:45 +0000 Subject: [PATCH 2/6] ipc: add ring_allreduce algorithms add the ring allreduce algorithm to the ipc conduit in order to be able to execute slightly largers reductions. --- src/ipc/context_ipc_device.hpp | 10 ++- src/ipc/context_ipc_tmpl_device.hpp | 114 ++++++++++++++++++++++++++-- 2 files changed, 115 insertions(+), 9 deletions(-) diff --git a/src/ipc/context_ipc_device.hpp b/src/ipc/context_ipc_device.hpp index 6bf8885d6c..dcb9571833 100644 --- a/src/ipc/context_ipc_device.hpp +++ b/src/ipc/context_ipc_device.hpp @@ -241,9 +241,13 @@ class IPCContext : public Context { template __device__ void internal_direct_allreduce(T *dst, const T *src, - int nelems, int PE_start, int - logPE_stride, int PE_size, - T *pWrk, long *pSync); + int nelems, int PE_start, int logPE_stride, + int PE_size, T *pWrk, long *pSync); + template + __device__ void internal_ring_allreduce(T *dst, const T *src, + int nelems, int PE_start, int logPE_stride, + int PE_size, T *pWrk, long *pSync, + int n_seg, int seg_size, int chunk_size); //internal functions used by collectives routines to write/read to //work/sync buffers diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp index 33bb91a1f1..81e669d942 100644 --- a/src/ipc/context_ipc_tmpl_device.hpp +++ b/src/ipc/context_ipc_tmpl_device.hpp @@ -219,9 +219,78 @@ __device__ void IPCContext::internal_direct_allreduce( __syncthreads(); } +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, + long *pSync, // NOLINT(runtime/int) + int n_seg, int seg_size, int chunk_size) { + int off_seg, off_send, off_recv; + int send_pe = (my_pe + 1) % num_pes; + long wait_val; // NOLINT(runtime/int) + + int wg_size = get_flat_block_size(); + int wg_id = get_flat_block_id(); + + for (size_t i = wg_id; i < nelems; i += wg_size) { + dst[i] = src[i]; + } + __syncthreads(); + + for (size_t seg = 0; seg < n_seg; seg++) { + off_seg = seg * seg_size; + for (int iter = 0; iter < num_pes - 1; iter++) { + 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]), + reinterpret_cast(&dst[off_send + off_seg]), + chunk_size * sizeof(T), send_pe); + + if (is_thread_zero_in_block()) { + fence(); + + wait_val = seg + 100; + put_nbi(&pSync[iter], &wait_val, 1, send_pe); +#if defined(__gfx90a__) + __threadfence_system(); +#endif /* __gfx90a__ */ + wait_until(&pSync[iter], ROC_SHMEM_CMP_EQ, wait_val); + } + __syncthreads(); + compute_reduce(&pWrk[off_recv], &dst[off_seg + off_recv], + chunk_size, wg_id, wg_size); + } + + for (size_t iter = num_pes - 1; iter < 2 * num_pes - 2; iter++) { + off_send = (((my_pe + 1 - iter + 2 * num_pes) % num_pes) * chunk_size); + putmem_nbi_wg(reinterpret_cast(&dst[off_send + off_seg]), + reinterpret_cast(&dst[off_send + off_seg]), + chunk_size * sizeof(T), send_pe); + + if (is_thread_zero_in_block()) { + fence(); + wait_val = seg + 100; + put_nbi(&pSync[iter], &wait_val, 1, send_pe); +#if defined(__gfx90a__) + __threadfence_system(); +#endif /* __gfx90a__ */ + wait_until(&pSync[iter], ROC_SHMEM_CMP_EQ, wait_val); + } + __syncthreads(); + } + } + __syncthreads(); + + for (size_t i = wg_id; i < 2 * num_pes - 2; i += wg_size) { + pSync[i] = ROC_SHMEM_SYNC_VALUE; + } + __syncthreads(); +} + template __device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest, - const T *source, int nreduce) { + const T *source, int nreduce) { IPCTeam *team_obj = reinterpret_cast(team); /** @@ -239,11 +308,44 @@ __device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest, template __device__ void IPCContext::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) - internal_direct_allreduce(dest, source, nreduce, PE_start, logPE_stride, - PE_size, pWrk, pSync); + int PE_start, int logPE_stride, + int PE_size, T *pWrk, + long *pSync) { // NOLINT(runtime/int) + size_t direct_pWrk = num_pes * nreduce; + size_t direct_pSync = num_pes; + size_t ring_pSync = 2 * num_pes; + size_t provided_pWrk = max(nreduce / 2 + 1, ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE); + 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, + PE_size, pWrk, pSync); + } else { + /* TODO (Edgar): some nreduce values cannot be evenly divided + ** among num_pes and/or segments. THe algorithm + ** currently cannot handle that one segment is of + ** different size than other segments. + */ + if (ring_pSync <= ROC_SHMEM_REDUCE_SYNC_SIZE) { + size_t ring_pWrk = ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE; + // integer division truncating value + int chunk_size = ring_pWrk / num_pes; + int seg_size = chunk_size * num_pes; + + // integer division rounding up + int n_seg = (nreduce + (seg_size -1)) / seg_size; + // recalculate chunk_size + chunk_size = seg_size / num_pes; + if (n_seg == 0) { + n_seg = 1; + } + internal_ring_allreduce(dest, source, nreduce, PE_start, + logPE_stride, PE_size, pWrk, pSync, n_seg, + seg_size, chunk_size); + } else { + GPU_DPRINTF("Unsupported reduction size for IPC conduit.\n"); + } + } } template From a0ac7b2d60d575e8bebe2df53bd2e1d270e6119e Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Thu, 17 Oct 2024 19:56:53 +0000 Subject: [PATCH 3/6] add some example code first examples include a getmem testcase and an allreduce (to_all) example. --- examples/rocshmem_allreduce_test.cc | 125 ++++++++++++++++++++++++++++ examples/rocshmem_getmem_test.cc | 94 +++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 examples/rocshmem_allreduce_test.cc create mode 100644 examples/rocshmem_getmem_test.cc diff --git a/examples/rocshmem_allreduce_test.cc b/examples/rocshmem_allreduce_test.cc new file mode 100644 index 0000000000..57e8571aa5 --- /dev/null +++ b/examples/rocshmem_allreduce_test.cc @@ -0,0 +1,125 @@ +/* +** hipcc -c -fgpu-rdc -x hip rocshmem_allreduce_test.cc -I/opt/rocm/include +** -I$ROCHSMEM_INSTALL_DIR/include -I$OPENMPI_UCX_INSTALL_DIR/include/ +** hipcc -fgpu-rdc --hip-link rocshmem_allreduce_test.o -o rocshmem_allreduce_test +** $ROCHSMEM_INSTALL_DIR/lib/librocshmem.a $OPENMPI_UCX_INSTALL_DIR/lib/libmpi.so +** -L/opt/rocm/lib -lamdhip64 -lhsa-runtime64 +** +** ROC_SHMEM_MAX_NUM_CONTEXTS=2 mpirun -np 8 ./rocshmem_allreduce_test +*/ + +#include + +#define __HIP_PLATFORM_AMD__ +#include +#include +#include + +#define CHECK_HIP(condition) { \ + hipError_t error = condition; \ + if(error != hipSuccess){ \ + fprintf(stderr,"HIP error: %d line: %d\n", error, __LINE__); \ + MPI_Abort(MPI_COMM_WORLD, error); \ + } \ + } + +using namespace rocshmem; + +__global__ void allreduce_test(int *source, int *dest, int* pWork, long *pSync, size_t nelem) +{ + __shared__ roc_shmem_ctx_t ctx; + int64_t ctx_type = 0; + + roc_shmem_wg_init(); + 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_quiet(ctx); + __syncthreads(); + + roc_shmem_wg_ctx_destroy(&ctx); + roc_shmem_wg_finalize(); +} + +static void init_sendbuf (int *sendbuf, int count, int mynode) +{ + for (int i = 0; i < count; i++) { + sendbuf[i] = mynode + i%9; + } +} + +static bool check_recvbuf(int *recvbuf, int nprocs, int rank, int count) +{ + bool res=true; + int expected = nprocs * (nprocs -1) / 2; + + for (int i=0; i 1) { + nelem = atoi(argv[1]); + } + + roc_shmem_init(); + int npes = roc_shmem_n_pes(); + int *source = (int *)roc_shmem_malloc(nelem * sizeof(int)); + int *result = (int *)roc_shmem_malloc(nelem * sizeof(int)); + if (NULL == source || NULL == result) { + std::cout << "Error allocating memory from symmetric heap" << std::endl; + roc_shmem_global_exit(1); + } + + init_sendbuf(source, nelem, rank); + for (int i=0; i>>(source, result, pWrk, pSync, nelem); + CHECK_HIP(hipDeviceSynchronize()); + + bool pass = check_recvbuf(result, npes, rank, nelem); + + printf("Test %s \t nelem %d %s\n", argv[0], nelem, pass ? "[PASS]" : "[FAIL]"); + + roc_shmem_free(source); + roc_shmem_free(result); + roc_shmem_free(pWrk); + roc_shmem_free(pSync); + + roc_shmem_finalize(); + return 0; +} diff --git a/examples/rocshmem_getmem_test.cc b/examples/rocshmem_getmem_test.cc new file mode 100644 index 0000000000..48664aa048 --- /dev/null +++ b/examples/rocshmem_getmem_test.cc @@ -0,0 +1,94 @@ +/* +** hipcc -c -fgpu-rdc -x hip rocshmem_getmem_test.cc -I/opt/rocm/include +** -I$ROCHSMEM_INSTALL_DIR/include -I$OPENMPI_UCX_INSTALL_DIR/include/ +** hipcc -fgpu-rdc --hip-link rocshmem_getmem_test.o -o rocshmem_getmem_test +** $ROCHSMEM_INSTALL_DIR/lib/librocshmem.a $OPENMPI_UCX_INSTALL_DIR/lib/libmpi.so +** -L/opt/rocm/lib -lamdhip64 -lhsa-runtime64 +** +** ROC_SHMEM_MAX_NUM_CONTEXTS=2 mpirun -np 2 ./rocshmem_getmem_test +*/ + +#include + +#define __HIP_PLATFORM_AMD__ +#include +#include +#include + +#define CHECK_HIP(condition) { \ + hipError_t error = condition; \ + if(error != hipSuccess){ \ + fprintf(stderr,"HIP error: %d line: %d\n", error, __LINE__); \ + MPI_Abort(MPI_COMM_WORLD, error); \ + } \ + } + +using namespace rocshmem; + +__global__ void simple_getmem_test(int *src, int *dst, size_t nelem) +{ + roc_shmem_wg_init(); + + int threadId = blockIdx.x * blockDim.x + threadIdx.x; + if (threadId == 0) { + int rank = roc_shmem_my_pe(); + int peer = rank ? 0 : 1; + roc_shmem_getmem(dst, src, nelem * sizeof(int), peer); + roc_shmem_quiet(); + } + + __syncthreads(); + roc_shmem_wg_finalize(); +} + +#define MAX_ELEM 256 + +int main (int argc, char **argv) +{ + int rank = roc_shmem_my_pe(); + int ndevices, my_device = 0; + CHECK_HIP(hipGetDeviceCount(&ndevices)); + my_device = rank % ndevices; + CHECK_HIP(hipSetDevice(my_device)); + int nelem = MAX_ELEM; + + if (argc > 1) { + nelem = atoi(argv[1]); + } + + roc_shmem_init(); + int npes = roc_shmem_n_pes(); + int *src = (int *)roc_shmem_malloc(nelem * sizeof(int)); + int *dst = (int *)roc_shmem_malloc(nelem * sizeof(int)); + if (NULL == src || NULL == dst) { + std::cout << "Error allocating memory from symmetric heap" << std::endl; + roc_shmem_global_exit(1); + } + + for (int i=0; i>>(src, dst, nelem); + roc_shmem_barrier_all(); + CHECK_HIP(hipDeviceSynchronize()); + + bool pass = true; + for (int i=0; i Date: Tue, 22 Oct 2024 17:22:51 +0000 Subject: [PATCH 4/6] fix barrier synchronization on gfx90a --- src/ipc/context_ipc_device.cpp | 1 + src/ipc/context_ipc_device_coll.cpp | 3 +++ 2 files changed, 4 insertions(+) diff --git a/src/ipc/context_ipc_device.cpp b/src/ipc/context_ipc_device.cpp index 275b15fe82..ee80694578 100644 --- a/src/ipc/context_ipc_device.cpp +++ b/src/ipc/context_ipc_device.cpp @@ -58,6 +58,7 @@ __host__ IPCContext::IPCContext(Backend *b) } __device__ void IPCContext::threadfence_system() { + __threadfence_system(); } __device__ void IPCContext::ctx_create() { diff --git a/src/ipc/context_ipc_device_coll.cpp b/src/ipc/context_ipc_device_coll.cpp index 7d7f39919c..aa7e7be9ae 100644 --- a/src/ipc/context_ipc_device_coll.cpp +++ b/src/ipc/context_ipc_device_coll.cpp @@ -47,6 +47,9 @@ __device__ void IPCContext::internal_direct_barrier(int pe, int PE_start, // Announce to other PEs that all have reached for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) { internal_putmem(&pSync[0], &flag_val, sizeof(*pSync), j); +#if defined(__gfx90a__) + __threadfence_system(); +#endif /* __gfx90a__ */ } } else { // Mark current PE offset as reached From a4b4281f504f8b6fd8f28c8641f047102ccc3b55 Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Wed, 23 Oct 2024 15:22:31 +0000 Subject: [PATCH 5/6] fix odd-case allreduce scenarios if the number of elements to be used in the allreduce operation is not exact multiple of the work-array buffer size and number of pe's, we need to adjust the algorithm to: - initially perform a ring_allreduce on n_segments * chunk_size (which is the integer division of the number of elements and the work-buffer size, i.e. will not cover the entire buffer) - perform another ring_allreduce where chunk_size is reduced to match the remaining elements - if the remaining elements from the previous step cannot evenly be divded by the number of pe's, we need to perform a direct_allreduce on the outstanding number of elements. --- src/ipc/context_ipc_tmpl_device.hpp | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp index 81e669d942..57a9177c04 100644 --- a/src/ipc/context_ipc_tmpl_device.hpp +++ b/src/ipc/context_ipc_tmpl_device.hpp @@ -321,19 +321,16 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, internal_direct_allreduce(dest, source, nreduce, PE_start, logPE_stride, PE_size, pWrk, pSync); } else { - /* TODO (Edgar): some nreduce values cannot be evenly divided - ** among num_pes and/or segments. THe algorithm - ** currently cannot handle that one segment is of - ** different size than other segments. - */ if (ring_pSync <= ROC_SHMEM_REDUCE_SYNC_SIZE) { size_t ring_pWrk = ROC_SHMEM_REDUCE_MIN_WRKDATA_SIZE; // integer division truncating value int chunk_size = ring_pWrk / num_pes; int seg_size = chunk_size * num_pes; + // integer division truncating value + int n_seg = nreduce / seg_size; // integer division rounding up - int n_seg = (nreduce + (seg_size -1)) / seg_size; + int n_seg_up = (nreduce + (seg_size -1)) / seg_size; // recalculate chunk_size chunk_size = seg_size / num_pes; if (n_seg == 0) { @@ -342,6 +339,25 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, internal_ring_allreduce(dest, source, nreduce, PE_start, logPE_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)); + const T *p_src = (source + (n_seg * seg_size)); + 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, + PE_size, pWrk, pSync, 1, (p_chunk * num_pes), p_chunk); + + if ((p_chunk * num_pes) < p_count) { + // Final elements need to use direct_allreduce + p_count -= (p_chunk * num_pes); + 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, + PE_size, pWrk, pSync); + } + } } else { GPU_DPRINTF("Unsupported reduction size for IPC conduit.\n"); } From 11df5427a6dfd5648d67226ad3757d0d709dcbf5 Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Wed, 23 Oct 2024 19:23:53 +0000 Subject: [PATCH 6/6] add ascii art for ring allredude --- src/ipc/context_ipc_tmpl_device.hpp | 58 +++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/src/ipc/context_ipc_tmpl_device.hpp b/src/ipc/context_ipc_tmpl_device.hpp index 57a9177c04..406367c72e 100644 --- a/src/ipc/context_ipc_tmpl_device.hpp +++ b/src/ipc/context_ipc_tmpl_device.hpp @@ -219,6 +219,62 @@ __device__ void IPCContext::internal_direct_allreduce( __syncthreads(); } +/* + * Visual representation of the ring_allreduce algorithm below + * assuming 4 PEs and a single segment. + * + * Initial state + * PE# 0 1 2 3 + * [00] [10] [20] [30] + * [01] [11] [21] [31] + * [02] [12] [22] [32] + * [03] [13] [23] [33] + * + * Loop 1: + * iter 0 + * PE# 0 1 2 3 + * [00+30] [10] [20] [30] + * [01] [01+11] [21] [31] + * [02] [12] [12+22] [32] + * [03] [13] [23] [23+33] + * + * iter 1 + * PE# 0 1 2 3 + * [00+30] [00+10+30] [20] [30] + * [01] [01+11] [01+11+21] [31] + * [02] [12] [12+22] [12+22+32] + * [03+23+33] [13] [23] [23+33] + * + * iter 2 + * PE# 0 1 2 3 + * [00+30] [00+10+30] [00+10+20+30] [30] + * [01] [01+11] [01+11+21] [01+11+21+31] + * [02+12+22+32] [12] [12+22] [12+22+32] + * [03+23+33] [03+13+23+33] [23] [23+33] + * + * Loop 2: + * + * iter 3 + * PE# 0 1 2 3 + * [00+30] [00+10+30] [00+10+20+30] [00+10+20+30] + * [01+11+21+31] [01+11] [01+11+21] [01+11+21+31] + * [02+12+22+32] [02+12+22+32] [12+22] [12+22+32] + * [03+23+33] [03+13+23+33] [03+13+23+33] [23+33] + * + * iter 4 + * PE# 0 1 2 3 + * [00+10+20+30] [00+10+30] [00+10+20+30] [00+10+20+30] + * [01+11+21+31] [01+11+21+31] [01+11+21] [01+11+21+31] + * [02+12+22+32] [02+12+22+32] [02+12+22+32] [12+22+32] + * [03+23+33] [03+13+23+33] [03+13+23+33] [03+13+23+33] + * + * iter 5 + * PE# 0 1 2 3 + * [00+10+20+30] [00+10+20+30] [00+10+20+30] [00+10+20+30] + * [01+11+21+31] [01+11+21+31] [01+11+21+31] [01+11+21+31] + * [02+12+22+32] [02+12+22+32] [02+12+22+32] [02+12+22+32] + * [03+13+23+33] [03+13+23+33] [03+13+23+33] [03+13+23+33] + */ template __device__ void IPCContext::internal_ring_allreduce( T *dst, const T *src, int nelems, [[maybe_unused]] int PE_start, @@ -239,6 +295,7 @@ __device__ void IPCContext::internal_ring_allreduce( for (size_t seg = 0; seg < n_seg; seg++) { off_seg = seg * seg_size; + // Loop 2 in the algorithm above for (int iter = 0; iter < num_pes - 1; iter++) { 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); @@ -262,6 +319,7 @@ __device__ void IPCContext::internal_ring_allreduce( chunk_size, wg_id, wg_size); } + // Loop 2 in the example above for (size_t iter = num_pes - 1; iter < 2 * num_pes - 2; iter++) { off_send = (((my_pe + 1 - iter + 2 * num_pes) % num_pes) * chunk_size); putmem_nbi_wg(reinterpret_cast(&dst[off_send + off_seg]),