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