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

Update broadcast function to use stride calculations instead of log_stride

Browse Source 
[ROCm/rocshmem commit: 5975b8c621]
This commit is contained in:
avinashkethineedi
2024-10-29 19:10:05 +00:00
parent 058aa3ccf3
commit b1a52b906b
2 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rocshmem/src/ipc/context_ipc_device.hpp
+5 -4
View File
@@ -137,10 +137,6 @@ class IPCContext : public Context {
int nelems, int pe_root);
template <typename T>
__device__ void broadcast(T *dest, const T *source, int nelems, int pe_root,
int pe_start, int log_pe_stride, int pe_size,
long *p_sync); // NOLINT(runtime/int)
template <typename T>
__device__ void alltoall(roc_shmem_team_t team, T *dest, const T *source,
int nelems);
template <typename T>
@@ -213,6 +209,11 @@ class IPCContext : public Context {
int stride, int PE_size, T *pWrk,
long *pSync); // NOLINT(runtime/int)
template <typename T>
__device__ void internal_broadcast(T *dest, const T *source, int nelems, int pe_root,
int pe_start, int stride, int pe_size,
long *p_sync); // NOLINT(runtime/int)
template <typename T>
__device__ void internal_put_broadcast(T *dst, const T *src, int nelems,
int pe_root, int PE_start,
projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp
+7 -8
View File
@@ -424,9 +424,8 @@ __device__ void IPCContext::internal_to_all(T *dest, const T *source, int nreduc
template <typename T>
__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) { // NOLINT(runtime/int)
int stride, int pe_size) { // 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) {
@@ -452,31 +451,31 @@ __device__ void IPCContext::broadcast(roc_shmem_team_t team, T *dst,
/**
* Ensure that the stride is a multiple of 2 .
*/
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->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<T>(dst, src, nelems, pe_root_world, pe_start, log_pe_stride,
internal_broadcast<T>(dst, src, nelems, pe_root_world, pe_start, stride,
pe_size, p_sync);
}
template <typename T>
__device__ void IPCContext::broadcast(T *dst, const T *src, int nelems,
__device__ void IPCContext::internal_broadcast(T *dst, const T *src, int nelems,
int pe_root, int pe_start,
int log_pe_stride, int pe_size,
int 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,
internal_put_broadcast(dst, src, nelems, pe_root, pe_start, stride,
pe_size);
} else {
internal_get_broadcast(dst, src, nelems, pe_root);
}
// Synchronize on completion of broadcast
internal_sync(my_pe, pe_start, (1 << log_pe_stride), pe_size, p_sync);
internal_sync(my_pe, pe_start, stride, pe_size, p_sync);
}
template <typename T>