From 9a380a65ff6f2e117c8efabd69d9fe0fea51e2e8 Mon Sep 17 00:00:00 2001 From: Edgar Gabriel Date: Tue, 13 Aug 2024 14:09:15 -0700 Subject: [PATCH] first cut on collectives and sync code is based on the GPUIB implementations of the routines, which seem however generic enough to work also for the IPC conduit. Some code is in for broadcast, fcollect, and alltoall. [ROCm/rocshmem commit: 0de3b5e6fcdde7413053a55d9ac86b2e4782d2c0] --- projects/rocshmem/src/ipc/CMakeLists.txt | 1 + .../rocshmem/src/ipc/context_ipc_device.cpp | 12 -- .../rocshmem/src/ipc/context_ipc_device.hpp | 39 +++++- .../src/ipc/context_ipc_device_coll.cpp | 120 ++++++++++++++++++ .../src/ipc/context_ipc_tmpl_device.hpp | 114 ++++++++++++++++- 5 files changed, 268 insertions(+), 18 deletions(-) create mode 100644 projects/rocshmem/src/ipc/context_ipc_device_coll.cpp diff --git a/projects/rocshmem/src/ipc/CMakeLists.txt b/projects/rocshmem/src/ipc/CMakeLists.txt index e8d64ff523..07a40538a5 100644 --- a/projects/rocshmem/src/ipc/CMakeLists.txt +++ b/projects/rocshmem/src/ipc/CMakeLists.txt @@ -30,4 +30,5 @@ target_sources( context_ipc_host.cpp backend_ipc.cpp ipc_team.cpp + context_ipc_device_coll.cpp ) diff --git a/projects/rocshmem/src/ipc/context_ipc_device.cpp b/projects/rocshmem/src/ipc/context_ipc_device.cpp index 7b9f3dd469..48d7dd9b4f 100644 --- a/projects/rocshmem/src/ipc/context_ipc_device.cpp +++ b/projects/rocshmem/src/ipc/context_ipc_device.cpp @@ -100,18 +100,6 @@ __device__ void *IPCContext::shmem_ptr(const void *dest, int pe) { return ret; } -__device__ void IPCContext::barrier_all() { - __syncthreads(); -} - -__device__ void IPCContext::sync_all() { - __syncthreads(); -} - -__device__ void IPCContext::sync(roc_shmem_team_t team) { - __syncthreads(); -} - __device__ void IPCContext::putmem_wg(void *dest, const void *source, size_t nelems, int pe) { // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable) diff --git a/projects/rocshmem/src/ipc/context_ipc_device.hpp b/projects/rocshmem/src/ipc/context_ipc_device.hpp index a8cefb446a..01fb1c1fb8 100644 --- a/projects/rocshmem/src/ipc/context_ipc_device.hpp +++ b/projects/rocshmem/src/ipc/context_ipc_device.hpp @@ -196,13 +196,44 @@ class IPCContext : public Context { private: - //context class has IpcImpl object (ipcImpl_) - IpcImpl *ipcImpl{nullptr}; + //context class has IpcImpl object (ipcImpl_) + IpcImpl *ipcImpl{nullptr}; - uint64_t* atomic_base_ptr{nullptr}; + uint64_t* atomic_base_ptr{nullptr}; - char* g_ret; + char* g_ret; + //internal functions used by collective operations + 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, + long *pSync); // NOLINT(runtime/int) + + template + __device__ void internal_get_broadcast(T *dst, const T *src, int nelems, + int pe_root, + long *pSync); // NOLINT(runtime/int) + + template + __device__ void fcollect_linear(roc_shmem_team_t team, T *dest, + const T *source, int nelems); + + template + __device__ void alltoall_linear(roc_shmem_team_t team, T *dest, + const T *source, int nelems); + + __device__ void internal_sync(int pe, int PE_start, int stride, int PE_size, + int64_t *pSync); + + __device__ void internal_direct_barrier(int pe, int PE_start, int stride, + int n_pes, int64_t *pSync); + + __device__ void internal_atomic_barrier(int pe, int PE_start, int stride, + int n_pes, int64_t *pSync); + + //Temporary scratchpad memory used by internal barrier algorithms. + int64_t *barrier_sync{nullptr}; }; } // namespace rocshmem diff --git a/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp b/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp new file mode 100644 index 0000000000..804be2efbf --- /dev/null +++ b/projects/rocshmem/src/ipc/context_ipc_device_coll.cpp @@ -0,0 +1,120 @@ +/****************************************************************************** + * 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. + *****************************************************************************/ + +#include "roc_shmem/roc_shmem.hpp" +#include "../context_incl.hpp" +#include "context_ipc_tmpl_device.hpp" +#include "../util.hpp" +#include "ipc_team.hpp" + +namespace rocshmem { + +__device__ void IPCContext::internal_direct_barrier(int pe, int PE_start, + int stride, int n_pes, + int64_t *pSync) { + int64_t flag_val = 1; + if (pe == PE_start) { + // Go through all PE offsets (except current offset = 0) + // and wait until they all reach + for (size_t i = 1; i < n_pes; i++) { + wait_until(&pSync[i], ROC_SHMEM_CMP_EQ, flag_val); + pSync[i] = ROC_SHMEM_SYNC_VALUE; + } + threadfence_system(); + + // Announce to other PEs that all have reached + for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) { + put_nbi(&pSync[0], &flag_val, 1, j); + } + } else { + // Mark current PE offset as reached + size_t pe_offset = (pe - PE_start) / stride; + put_nbi(&pSync[pe_offset], &flag_val, 1, PE_start); + wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val); + pSync[0] = ROC_SHMEM_SYNC_VALUE; + threadfence_system(); + } +} + +__device__ void IPCContext::internal_atomic_barrier(int pe, int PE_start, + int stride, int n_pes, + int64_t *pSync) { + int64_t flag_val = 1; + if (pe == PE_start) { + wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, (int64_t)(n_pes - 1)); + pSync[0] = ROC_SHMEM_SYNC_VALUE; + threadfence_system(); + + for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) { + put_nbi(&pSync[0], &flag_val, 1, j); + } + } else { + amo_add(&pSync[0], flag_val, PE_start); + wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val); + pSync[0] = ROC_SHMEM_SYNC_VALUE; + threadfence_system(); + } +} + +// Uses PE values that are relative to world +__device__ void IPCContext::internal_sync(int pe, int PE_start, int stride, + int PE_size, int64_t *pSync) { + __syncthreads(); + if (is_thread_zero_in_block()) { + if (PE_size < 64) { + internal_direct_barrier(pe, PE_start, stride, PE_size, pSync); + } else { + internal_atomic_barrier(pe, PE_start, stride, PE_size, pSync); + } + } + __threadfence(); + __syncthreads(); +} + +__device__ void IPCContext::sync(roc_shmem_team_t team) { + IPCTeam *team_obj = reinterpret_cast(team); + + /** + * Ensure that the stride is a multiple of 2. + */ + int log_pe_stride = static_cast(team_obj->tinfo_wrt_world->log_stride); + int pe = team_obj->my_pe_in_world; + int pe_start = team_obj->tinfo_wrt_world->pe_start; + int pe_stride = (1 << log_pe_stride); + int pe_size = team_obj->num_pes; + + internal_sync(pe, pe_start, pe_stride, pe_size, barrier_sync); +} + +__device__ void IPCContext::sync_all() { + internal_sync(my_pe, 0, 1, num_pes, barrier_sync); +} + +__device__ void IPCContext::barrier_all() { + if (is_thread_zero_in_block()) { + quiet(); + } + sync_all(); + __syncthreads(); +} + +} // namespace rocshmem diff --git a/projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp b/projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp index 4bf652781e..af79cd20b1 100644 --- a/projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp +++ b/projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp @@ -27,6 +27,7 @@ #include "roc_shmem/roc_shmem.hpp" #include "context_ipc_device.hpp" #include "../util.hpp" +#include "ipc_team.hpp" namespace rocshmem { @@ -150,11 +151,48 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce, long *pSync) { // NOLINT(runtime/int) } +template +__device__ void IPCContext::internal_put_broadcast( + T *dst, const T *src, int nelems, int pe_root, int pe_start, + int log_pe_stride, int pe_size, + [[maybe_unused]] long *p_sync) { // NOLINT(runtime/int) + if (my_pe == pe_root) { + int stride = 1 << log_pe_stride; + int finish = pe_start + stride * pe_size; + for (int i = pe_start; i < finish; i += stride) { + if (i != my_pe) { + put_nbi_wg(dst, src, nelems, i); + } + } + } +} + +template +__device__ void IPCContext::internal_get_broadcast( + T *dst, const T *src, int nelems, int pe_root, + [[maybe_unused]] long *pSync) { // NOLINT(runtime/int) + if (my_pe != pe_root) { + get_wg(dst, src, nelems, pe_root); + } +} + template __device__ void IPCContext::broadcast(roc_shmem_team_t team, T *dst, const T *src, int nelems, int pe_root) { - //broadcast(dst, src, nelems, pe_root_world, pe_start, log_pe_stride, - // pe_size, p_sync); + IPCTeam *team_obj = reinterpret_cast(team); + + /** + * Ensure that the stride is a multiple of 2 . + */ + 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->bcast_pSync; + + // Passed pe_root is relative to team, convert to world root + int pe_root_world = team_obj->get_pe_in_world(pe_root); + broadcast(dst, src, nelems, pe_root_world, pe_start, log_pe_stride, + pe_size, p_sync); } template @@ -162,16 +200,88 @@ __device__ void IPCContext::broadcast(T *dst, const T *src, int nelems, int pe_root, int pe_start, int log_pe_stride, int pe_size, long *p_sync) { // NOLINT(runtime/int) + if (num_pes < 4) { + internal_put_broadcast(dst, src, nelems, pe_root, pe_start, log_pe_stride, + pe_size, p_sync); + } else { + internal_get_broadcast(dst, src, nelems, pe_root, p_sync); + } + + // Synchronize on completion of broadcast + internal_sync(my_pe, pe_start, (1 << log_pe_stride), pe_size, p_sync); } template __device__ void IPCContext::alltoall(roc_shmem_team_t team, T *dst, const T *src, int nelems) { + alltoall_linear(team, dst, src, nelems); +} + +template +__device__ void IPCContext::alltoall_linear(roc_shmem_team_t team, T *dst, + const T *src, int nelems) { + IPCTeam *team_obj = reinterpret_cast(team); + + /** + * Ensure that the stride is a multiple of 2 + */ + 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->num_pes; + int stride = 1 << log_pe_stride; + long *pSync = team_obj->alltoall_pSync; + int my_pe_in_team = team_obj->my_pe; + + /* TODO (egabriel): shouldn't the loop handle pe_start, pe_size, and stride? + * otherwise what is the point of the few lines above? + */ + // Have each PE put their designated data to the other PEs + for (int j = 0; j < pe_size; j++) { + 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); + } + if (is_thread_zero_in_block()) { + quiet(); + } + // wait until everyone has obtained their designated data + internal_sync(my_pe, pe_start, stride, pe_size, pSync); } template __device__ void IPCContext::fcollect(roc_shmem_team_t team, T *dst, const T *src, int nelems) { + fcollect_linear(team, dst, src, nelems); +} + +template +__device__ void IPCContext::fcollect_linear(roc_shmem_team_t team, T *dst, + const T *src, int nelems) { + IPCTeam *team_obj = reinterpret_cast(team); + + /** + * Ensure that the stride is a multiple of 2. + */ + 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->num_pes; + int stride = 1 << log_pe_stride; + long *pSync = team_obj->alltoall_pSync; + int my_pe_in_team = team_obj->my_pe; + + /* TODO (egabriel): shouldn't the loop handle pe_start, pe_size, and stride? + * otherwise what is the point of the few lines above? + */ + // Have each PE put their designated data to the other PEs + for (int j = 0; j < pe_size; j++) { + int dest_pe = team_obj->get_pe_in_world(j); + put_nbi_wg(&dst[my_pe_in_team * nelems], src, nelems, dest_pe); + } + + if (is_thread_zero_in_block()) { + quiet(); + } + // wait until everyone has obtained their designated data + internal_sync(my_pe, pe_start, stride, pe_size, pSync); } // Block/wave functions