Merge pull request #32 from edgargabriel/topic/to_all_first_version
ipc/to_all: add direct allreduce algorithm
This commit is contained in:
Edgar Gabriel
GitHub
@@ -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 <iostream>
|
||||
|
||||
#define __HIP_PLATFORM_AMD__
|
||||
#include <hip/hip_runtime_api.h>
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <roc_shmem.hpp>
|
||||
|
||||
#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<count; i++) {
|
||||
int result = expected + nprocs * (i%9);
|
||||
if (recvbuf[i] != result) {
|
||||
res = false;
|
||||
#ifdef VERBOSE
|
||||
printf("recvbuf[%d] = %d expected %d \n", i, recvbuf[i], result);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
#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 *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<nelem; i++) {
|
||||
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));
|
||||
|
||||
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<<<dim3(1), dim3(threadsPerBlock), 0, 0>>>(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;
|
||||
}
|
||||
@@ -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 <iostream>
|
||||
|
||||
#define __HIP_PLATFORM_AMD__
|
||||
#include <hip/hip_runtime_api.h>
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <roc_shmem.hpp>
|
||||
|
||||
#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<nelem; i++) {
|
||||
src[i] = 0;
|
||||
dst[i] = 1;
|
||||
}
|
||||
CHECK_HIP(hipDeviceSynchronize());
|
||||
|
||||
int threadsPerBlock=256;
|
||||
simple_getmem_test<<<dim3(1), dim3(threadsPerBlock), 0, 0>>>(src, dst, nelem);
|
||||
roc_shmem_barrier_all();
|
||||
CHECK_HIP(hipDeviceSynchronize());
|
||||
|
||||
bool pass = true;
|
||||
for (int i=0; i<nelem; i++) {
|
||||
if (dst[i] != 0) {
|
||||
pass = false;
|
||||
#if VERBOSE
|
||||
printf("[%d] Error in element %d expected 0 got %d\n", rank, i, dst[i]);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
printf("Test %s \t %s\n", argv[0], pass ? "[PASS]" : "[FAIL]");
|
||||
|
||||
roc_shmem_free(src);
|
||||
roc_shmem_free(dst);
|
||||
roc_shmem_finalize();
|
||||
return 0;
|
||||
}
|
||||
@@ -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 <ROC_SHMEM_OP Op>
|
||||
struct OpWrap {
|
||||
template <typename T>
|
||||
__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<ROC_SHMEM_SUM> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] += src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_MAX> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] = max(dst[i], src[i]);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_MIN> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] = min(dst[i], src[i]);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_PROD> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] *= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_AND> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] &= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_OR> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] |= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_XOR> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] ^= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__device__ void compute_reduce(T *src, T *dst, int size, int wg_id,
|
||||
int wg_size) {
|
||||
|
||||
@@ -58,6 +58,7 @@ __host__ IPCContext::IPCContext(Backend *b)
|
||||
}
|
||||
|
||||
__device__ void IPCContext::threadfence_system() {
|
||||
__threadfence_system();
|
||||
}
|
||||
|
||||
__device__ void IPCContext::ctx_create() {
|
||||
|
||||
@@ -239,6 +239,16 @@ class IPCContext : public Context {
|
||||
__device__ void internal_atomic_barrier(int pe, int PE_start, int stride,
|
||||
int n_pes, int64_t *pSync);
|
||||
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__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);
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__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
|
||||
__device__ void internal_putmem(void *dest, const void *source,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -28,6 +28,7 @@
|
||||
#include "context_ipc_device.hpp"
|
||||
#include "../util.hpp"
|
||||
#include "ipc_team.hpp"
|
||||
#include "../roc_shmem_calc.hpp"
|
||||
|
||||
namespace rocshmem {
|
||||
|
||||
@@ -152,18 +153,273 @@ __device__ T IPCContext::amo_fetch_cas(void *dest, T value, T cond, int pe) {
|
||||
}
|
||||
|
||||
// Collectives
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__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<Op>::Calc(src, dst, i);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__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<const void *>(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<T, Op>(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();
|
||||
}
|
||||
|
||||
/*
|
||||
* 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 <typename T, ROC_SHMEM_OP Op>
|
||||
__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;
|
||||
// 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);
|
||||
|
||||
putmem_nbi_wg(reinterpret_cast<void *>(&pWrk[off_send]),
|
||||
reinterpret_cast<void *>(&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<T, Op>(&pWrk[off_recv], &dst[off_seg + off_recv],
|
||||
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<void *>(&dst[off_send + off_seg]),
|
||||
reinterpret_cast<void *>(&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 <typename T, ROC_SHMEM_OP Op>
|
||||
__device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest,
|
||||
const T *source, int nreduce) {
|
||||
//to_all<T, Op>(dest, source, nreduce, pe_start, log_pe_stride, pe_size, pWrk,
|
||||
// p_sync);
|
||||
const T *source, int nreduce) {
|
||||
IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
|
||||
|
||||
/**
|
||||
* Ensure that the stride is a multiple of 2 for GPU_IB.
|
||||
*/
|
||||
int log_pe_stride = static_cast<int>(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<T *>(team_obj->pWrk);
|
||||
|
||||
to_all<T, Op>(dest, source, nreduce, pe_start, log_pe_stride, pe_size, pWrk,
|
||||
p_sync);
|
||||
}
|
||||
|
||||
template <typename T, ROC_SHMEM_OP Op>
|
||||
__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)
|
||||
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<T, Op>(dest, source, nreduce, PE_start, logPE_stride,
|
||||
PE_size, pWrk, pSync);
|
||||
} else {
|
||||
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_up = (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<T, Op>(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<T, Op>(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<T, Op>(p_dst, p_src2, p_count, PE_start, logPE_stride,
|
||||
PE_size, pWrk, pSync);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
GPU_DPRINTF("Unsupported reduction size for IPC conduit.\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
|
||||
@@ -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 <ROC_SHMEM_OP Op>
|
||||
struct OpWrap {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
static_assert(true, "Unimplemented ipc collective.");
|
||||
}
|
||||
};
|
||||
)
|
||||
// clang-format on
|
||||
|
||||
/******************************************************************************
|
||||
************************** TEMPLATE SPECIALIZATIONS **************************
|
||||
*****************************************************************************/
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_SUM> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] += src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_MAX> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] = max(dst[i], src[i]);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_MIN> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] = min(dst[i], src[i]);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_PROD> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] *= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_AND> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] &= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_OR> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] |= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct OpWrap<ROC_SHMEM_XOR> {
|
||||
template <typename T>
|
||||
__device__ static void Calc(T *src, T *dst, int i) {
|
||||
dst[i] ^= src[i];
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
#endif // LIBRARY_SRC_ROC_SHMEM_CALC_HPP_
|
||||
Reference in New Issue
Block a user