modified team based to_all -> reduce

[ROCm/rocshmem commit: 997eb69b5a]
This commit is contained in:
Yiltan Hassan Temucin
2024-10-29 06:32:13 -07:00
committed by Yiltan Temucin
parent b43ef2b45b
commit 7f3dd33144
14 changed files with 52 additions and 50 deletions
+2 -4
View File
@@ -192,8 +192,7 @@ class Context {
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);
__device__ int reduce(roc_shmem_team_t team, T* dest, const T* source, int nreduce);
template <typename T>
__device__ void put(T* dest, const T* source, size_t nelems, int pe);
@@ -361,8 +360,7 @@ class Context {
long* pSync); // NOLINT(runtime/int)
template <typename T, ROC_SHMEM_OP Op>
__host__ void to_all(roc_shmem_team_t team, T* dest, const T* source,
int nreduce);
__host__ int reduce(roc_shmem_team_t team, T* dest, const T* source, int nreduce);
template <typename T>
__host__ void wait_until(T *ivars, int cmp, T val);
@@ -80,17 +80,17 @@ __device__ void Context::to_all(T *dest, const T *source, int nreduce,
}
template <typename T, ROC_SHMEM_OP Op>
__device__ void Context::to_all(roc_shmem_team_t team, T *dest, const T *source,
__device__ int Context::reduce(roc_shmem_team_t team, T *dest, const T *source,
int nreduce) {
if (nreduce == 0) {
return;
return ROC_SHMEM_SUCCESS;
}
if (is_thread_zero_in_block()) {
ctxStats.incStat(NUM_TO_ALL);
}
DISPATCH(to_all<PAIR(T, Op)>(team, dest, source, nreduce));
DISPATCH_RET(reduce<PAIR(T, Op)>(team, dest, source, nreduce));
}
template <typename T>
+4 -4
View File
@@ -222,15 +222,15 @@ __host__ void Context::to_all(T *dest, const T *source, int nreduce,
}
template <typename T, ROC_SHMEM_OP Op>
__host__ void Context::to_all(roc_shmem_team_t team, T *dest, const T *source,
int nreduce) { // NOLINT(runtime/int)
__host__ int Context::reduce(roc_shmem_team_t team, T *dest, const T *source,
int nreduce) { // NOLINT(runtime/int)
if (nreduce == 0) {
return;
return ROC_SHMEM_SUCCESS;
}
ctxHostStats.incStat(NUM_HOST_TO_ALL);
HOST_DISPATCH(to_all<PAIR(T, Op)>(team, dest, source, nreduce));
HOST_DISPATCH_RET(reduce<PAIR(T, Op)>(team, dest, source, nreduce));
}
template <typename T>
+1 -2
View File
@@ -207,8 +207,7 @@ class HostInterface {
long* p_sync); // NOLINT(runtime/int)
template <typename T, ROC_SHMEM_OP Op>
__host__ void to_all(roc_shmem_team_t team, T* dest, const T* source,
int nreduce);
__host__ int reduce(roc_shmem_team_t team, T* dest, const T* source, int nreduce);
template <typename T>
__host__ void wait_until(T *ivars, int cmp, T val,
@@ -376,9 +376,9 @@ __host__ void HostInterface::to_all(T* dest, const T* source, int nreduce,
}
template <typename T, ROC_SHMEM_OP Op>
__host__ void HostInterface::to_all(roc_shmem_team_t team, T* dest,
__host__ int HostInterface::reduce(roc_shmem_team_t team, T* dest,
const T* source, int nreduce) {
DPRINTF("Function: Team-based host_to_all\n");
DPRINTF("Function: Team-based host_reduce\n");
/*
* Get the MPI communicator of this team
@@ -388,7 +388,7 @@ __host__ void HostInterface::to_all(roc_shmem_team_t team, T* dest,
to_all_internal<T, Op>(mpi_comm, dest, source, nreduce);
return;
return ROC_SHMEM_SUCCESS;
}
template <typename T>
@@ -122,8 +122,7 @@ class IPCContext : public Context {
// Collectives
template <typename T, ROC_SHMEM_OP Op>
__device__ void to_all(roc_shmem_team_t team, T *dest, const T *source,
int nreduce);
__device__ int reduce(roc_shmem_team_t team, T *dest, const T *source, int nreduce);
template <typename T>
__device__ void broadcast(roc_shmem_team_t team, T *dest, const T *source,
@@ -95,8 +95,7 @@ class IPCHostContext : public Context {
long *p_sync);
template <typename T, ROC_SHMEM_OP Op>
__host__ void to_all(roc_shmem_team_t team, T *dest, const T *source,
int nreduce);
__host__ int reduce(roc_shmem_team_t team, T *dest, const T *source, int nreduce);
template <typename T>
__host__ void wait_until(T *ivars, int cmp, T val);
@@ -151,7 +151,7 @@ __device__ T IPCContext::amo_fetch_cas(void *dest, T value, T cond, int pe) {
reinterpret_cast<T *>(ipcImpl_.ipc_bases[pe] + L_offset), cond,
value);
}
// Collectives
template <typename T, ROC_SHMEM_OP Op>
__device__ void compute_reduce(T *src, T *dst, int size, int wg_id,
@@ -346,8 +346,8 @@ __device__ void IPCContext::internal_ring_allreduce(
}
template <typename T, ROC_SHMEM_OP Op>
__device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest,
const T *source, int nreduce) {
__device__ int IPCContext::reduce(roc_shmem_team_t team, T *dest,
const T *source, int nreduce) {
IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
/**
@@ -361,6 +361,7 @@ __device__ void IPCContext::to_all(roc_shmem_team_t team, T *dest,
internal_to_all<T, Op>(dest, source, nreduce, pe_start, stride, pe_size, pWrk,
p_sync);
return ROC_SHMEM_SUCCESS;
}
template <typename T, ROC_SHMEM_OP Op>
@@ -109,9 +109,9 @@ __host__ void IPCHostContext::to_all(T *dest, const T *source, int nreduce,
}
template <typename T, ROC_SHMEM_OP Op>
__host__ void IPCHostContext::to_all(roc_shmem_team_t team, T *dest,
__host__ int IPCHostContext::reduce(roc_shmem_team_t team, T *dest,
const T *source, int nreduce) {
host_interface->to_all<T, Op>(team, dest, source, nreduce);
return host_interface->reduce<T, Op>(team, dest, source, nreduce);
}
template <typename T>
+7 -7
View File
@@ -734,13 +734,13 @@ __host__ void roc_shmem_to_all([[maybe_unused]] roc_shmem_ctx_t ctx, T *dest,
}
template <typename T, ROC_SHMEM_OP Op>
__host__ void roc_shmem_to_all([[maybe_unused]] roc_shmem_ctx_t ctx,
__host__ int roc_shmem_reduce([[maybe_unused]] roc_shmem_ctx_t ctx,
roc_shmem_team_t team, T *dest, const T *source,
int nreduce) {
DPRINTF("Host function: Team-based roc_shmem_to_all\n");
DPRINTF("Host function: Team-based roc_shmem_reduce\n");
get_internal_ctx(ROC_SHMEM_HOST_CTX_DEFAULT)
->to_all<T, Op>(team, dest, source, nreduce);
return get_internal_ctx(ROC_SHMEM_HOST_CTX_DEFAULT)
->reduce<T, Op>(team, dest, source, nreduce);
}
template <typename T>
@@ -821,7 +821,7 @@ __host__ int roc_shmem_test(T *ivars, int cmp, T val) {
template __host__ void roc_shmem_to_all<T, Op>( \
roc_shmem_ctx_t ctx, T * dest, const T *source, int nreduce, \
int PE_start, int logPE_stride, int PE_size, T *pWrk, long *pSync); \
template __host__ void roc_shmem_to_all<T, Op>( \
template __host__ int roc_shmem_reduce<T, Op>( \
roc_shmem_ctx_t ctx, roc_shmem_team_t team, T * dest, const T *source, \
int nreduce);
@@ -977,10 +977,10 @@ __host__ int roc_shmem_test(T *ivars, int cmp, T val) {
roc_shmem_to_all<T, Op>(ctx, dest, source, nreduce, PE_start, \
logPE_stride, PE_size, pWrk, pSync); \
} \
__host__ void roc_shmem_ctx_##TNAME##_##Op_API##_to_all( \
__host__ int roc_shmem_ctx_##TNAME##_##Op_API##_reduce( \
roc_shmem_ctx_t ctx, roc_shmem_team_t team, T *dest, const T *source, \
int nreduce) { \
roc_shmem_to_all<T, Op>(ctx, team, dest, source, nreduce); \
return roc_shmem_reduce<T, Op>(ctx, team, dest, source, nreduce); \
}
#define ARITH_REDUCTION_DEF_GEN(T, TNAME) \
+7 -7
View File
@@ -430,11 +430,11 @@ __device__ void *roc_shmem_ptr(const void *dest, int pe) {
}
template <typename T, ROC_SHMEM_OP Op>
__device__ void roc_shmem_wg_to_all(roc_shmem_ctx_t ctx, roc_shmem_team_t team,
T *dest, const T *source, int nreduce) {
GPU_DPRINTF("Function: roc_shmem_to_all\n");
__device__ int roc_shmem_wg_reduce(roc_shmem_ctx_t ctx, roc_shmem_team_t team,
T *dest, const T *source, int nreduce) {
GPU_DPRINTF("Function: roc_shmem_reduce\n");
get_internal_ctx(ctx)->to_all<T, Op>(team, dest, source, nreduce);
return get_internal_ctx(ctx)->reduce<T, Op>(team, dest, source, nreduce);
}
template <typename T>
@@ -864,7 +864,7 @@ __device__ int roc_shmem_team_translate_pe(roc_shmem_team_t src_team,
* Template generator for reductions
*/
#define REDUCTION_GEN(T, Op) \
template __device__ void roc_shmem_wg_to_all<T, Op>( \
template __device__ int roc_shmem_wg_reduce<T, Op>( \
roc_shmem_ctx_t ctx, roc_shmem_team_t team, T * dest, const T *source, \
int nreduce);
@@ -1072,10 +1072,10 @@ __device__ int roc_shmem_team_translate_pe(roc_shmem_team_t src_team,
**/
#define REDUCTION_DEF_GEN(T, TNAME, Op_API, Op) \
__device__ void roc_shmem_ctx_##TNAME##_##Op_API##_wg_to_all( \
__device__ int roc_shmem_ctx_##TNAME##_##Op_API##_wg_reduce( \
roc_shmem_ctx_t ctx, roc_shmem_team_t team, T *dest, const T *source, \
int nreduce) { \
roc_shmem_wg_to_all<T, Op>(ctx, team, dest, source, nreduce); \
return roc_shmem_wg_reduce<T, Op>(ctx, team, dest, source, nreduce); \
}
#define ARITH_REDUCTION_DEF_GEN(T, TNAME) \