AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update all_reduce algorithm to use internal put/get functions for updating pWrk and pSync arrays

Browse Source 
* Change log_stride calcualtions to stride calculations
* Update all_reduce example code to use team based interface
This commit is contained in:
avinashkethineedi
2024-10-28 22:10:18 +00:00
parent c22048112e
commit abec29bd6a
4 changed files with 35 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/rocshmem_allreduce_test.cc
+9 -13
View File
@@ -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<<<dim3(1), dim3(threadsPerBlock), 0, 0>>>(source, result, pWrk, pSync, nelem);
allreduce_test<<<dim3(1), dim3(threadsPerBlock), 0, 0>>>(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;
src/ipc/context_ipc_device.hpp
+6 -6
View File
@@ -128,11 +128,6 @@ class IPCContext : public Context {
__device__ T amo_fetch_cas(void *dst, T value, T cond, int pe);
// Collectives
template <typename T, ROC_SHMEM_OP Op>
__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 <typename T, ROC_SHMEM_OP Op>
__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 <typename T>
template <typename T, ROC_SHMEM_OP Op>
__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 <typename T>
__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)
src/ipc/context_ipc_tmpl_device.hpp
+15 -16
View File
@@ -164,11 +164,10 @@ __device__ void compute_reduce(T *src, T *dst, int size, int wg_id,
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,
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<const void *>(src),
internal_putmem_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);
internal_putmem(&pSync[pe], &flag_val, sizeof(*pSync), i);
}
}
}
@@ -278,7 +277,7 @@ __device__ void IPCContext::internal_direct_allreduce(
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,
[[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<void *>(&pWrk[off_send]),
internal_putmem_wg(reinterpret_cast<void *>(&pWrk[off_send]),
reinterpret_cast<void *>(&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<int>(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<T *>(team_obj->pWrk);
to_all<T, Op>(dest, source, nreduce, pe_start, log_pe_stride, pe_size, pWrk,
internal_to_all<T, Op>(dest, source, nreduce, pe_start, 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,
__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<T, Op>(dest, source, nreduce, PE_start, logPE_stride,
internal_direct_allreduce<T, Op>(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<T, Op>(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<T, Op>(p_dst, p_src, p_count, PE_start, logPE_stride,
internal_ring_allreduce<T, Op>(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<T, Op>(p_dst, p_src2, p_count, PE_start, logPE_stride,
internal_direct_allreduce<T, Op>(p_dst, p_src2, p_count, PE_start, stride,
PE_size, pWrk, pSync);
}
}
tests/functional_tests/tester_arguments.cpp
+5 -5
View File
@@ -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 "