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

[DEVICE] Enable PAT algo for RCCL 1ppn (#1756)

Browse Source 
* Enable PAT algo for RCCL 1ppn


[ROCm/rccl commit: e96c8473a1]
This commit is contained in:
Bertan Dogancay
2025-07-04 13:45:18 -04:00
committed by GitHub
parent 82a822b646
commit 471fc6bff2
12 changed files with 108 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/device/all_gather.h
+5 -7
View File
@@ -220,7 +220,6 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128
template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE> {
__device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
#if __CUDA_ARCH__ >= 600
using Proto = ProtoSimple<1, 1>;
const int nranks = ncclShmem.comm.nRanks;
const int rank = ncclShmem.comm.rank;
@@ -242,8 +241,8 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
int step = 0;
while (1) {
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
cuda::atomic_ref<int, cuda::thread_scope_block> poll(ps->flags);
while (poll.load(cuda::memory_order_acquire) != 0) pollCount++; // Wait for workers to be done with step 'step-NCCL_SHMEM_PAT_STEPS'
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) != 0) pollCount++; // Wait for workers to be done with step 'step-NCCL_SHMEM_PAT_STEPS'
patAlgo.getNextOp(ps);
int last = ps->last;
step++;
@@ -267,16 +266,15 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
int step = group;
while(1) {
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
cuda::atomic_ref<int, cuda::thread_scope_block> poll(ps->flags);
while (poll.load(cuda::memory_order_acquire) == 0) pollCount++; // Wait for compute thread
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) == 0) pollCount++; // Wait for compute thread
int last = ps->last;
prims.patCopy(ps, shmem);
if (tidInGroup == 0) poll.store(0, cuda::memory_order_release); // Return element to compute thread
if (tidInGroup == 0) __hip_atomic_store(poll, 0, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_WORKGROUP); // Return element to compute thread
if (last) break;
step += nGroups;
}
}
#endif
}
};
projects/rccl/src/device/common.h
+7 -2
View File
@@ -164,6 +164,7 @@ struct ncclShmemData {
#ifdef ENABLE_FAULT_INJECTION
uint64_t faults;
#endif
uint64_t barrier_pat;
};
extern __shared__ ncclShmemData ncclShmem;
@@ -526,8 +527,10 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
}
break;
case 1:
if (tid < WARP_SIZE + NCCL_MAX_GROUPS)
if (tid < WARP_SIZE + NCCL_MAX_GROUPS) {
if (tid == WARP_SIZE) ncclShmem.barrier_pat = 0;
ncclShmem.groups[tid-WARP_SIZE].barrier = 0;
}
break;
case 2:
#ifdef ENABLE_FAULT_INJECTION
@@ -611,8 +614,10 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
__syncthreads();
switch (tid/WARP_SIZE) {
case 1:
if (tid < WARP_SIZE + NCCL_MAX_GROUPS)
if (tid < WARP_SIZE + NCCL_MAX_GROUPS) {
if (tid == WARP_SIZE) ncclShmem.barrier_pat = 0;
ncclShmem.groups[tid-WARP_SIZE].barrier = 0;
}
break;
default:
break;
projects/rccl/src/device/generate.py
+4 -4
View File
@@ -8,7 +8,7 @@ all_colls = ["AllGather","AllReduce","AllToAllPivot","Broadcast","Reduce","Redu
all_redops = ["Sum","Prod","MinMax","PreMulSum","SumPostDiv"]
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16","f8e4m3","f8e5m2"]
all_protos = ["LL","LL128","SIMPLE"]
all_algos = ["TREE","RING"]
all_algos = ["TREE","RING", "PAT"]
all_unroll = ["1", "2", "4"]
all_params = [all_colls, all_algos, all_protos, all_redops, all_tys, all_unroll]
@@ -85,12 +85,12 @@ else:
################################################################################
algos_of_coll = {
"AllGather": ["RING"],
"AllReduce": all_algos,
"AllGather": ["RING", "PAT"],
"AllReduce": ["RING", "TREE"],
"AllToAllPivot": ["RING"],
"Broadcast": ["RING"],
"Reduce": ["RING"],
"ReduceScatter": ["RING"],
"ReduceScatter": ["RING", "PAT"],
"SendRecv": ["RING"]
}
projects/rccl/src/device/primitives.h
+6 -12
View File
@@ -15,19 +15,19 @@
#define NCCL_SPINS_BEFORE_CHECK_ABORT 10000
#define barrier_by_group_common(__THREAD_FENCE) do { \
#define barrier_generic(__THREAD_FENCE, NWORKERS, BARRIER_NEXT, BARRIERS_PTR) do { \
if (nthreads == NCCL_MAX_NTHREADS) { \
__THREAD_FENCE; __builtin_amdgcn_s_barrier(); \
} else { \
const int w = threadIdx.x/WARP_SIZE; \
const int wid = threadIdx.x%WARP_SIZE; \
if (wid == 0) { \
barrier_next += nthreads/WARP_SIZE; \
(BARRIER_NEXT) += (NWORKERS) / WARP_SIZE; \
__THREAD_FENCE; \
__hip_atomic_fetch_add(barriers, 1, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_WORKGROUP); \
__hip_atomic_fetch_add((BARRIERS_PTR), 1, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_WORKGROUP); \
int spins = 0; \
int rate_limit = 50; \
while (__hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) < barrier_next) { \
while (__hip_atomic_load((BARRIERS_PTR), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) < (BARRIER_NEXT)) { \
spins++; \
if (spins == NCCL_SPINS_BEFORE_CHECK_ABORT) { \
if (__atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST)) { \
@@ -37,8 +37,8 @@
spins = 0; \
} \
if (spins == 0 && rate_limit > 0) { \
rate_limit --; \
traceData(__LINE__, threadIdx.x, __hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP), barrier_next); \
rate_limit--; \
traceData(__LINE__, threadIdx.x, __hip_atomic_load((BARRIERS_PTR), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP), (BARRIER_NEXT)); \
} \
__builtin_amdgcn_s_sleep(1); \
} \
@@ -47,12 +47,6 @@
} \
} while (0)
#define barrier_by_group() barrier_by_group_common(__threadfence())
#if defined(__gfx942__) || defined(__gfx950__)
#define barrier_by_group_block() barrier_by_group_common(__threadfence_block())
#endif
/* Protocol classes: ProtoSimple, ProtoLL, ProtoLL128
* We use these as template args to the Primtiives class instead of integral
* enums (e.g. NCCL_PROTO_LL) because for SIMPLE we need to carry a few extra
projects/rccl/src/device/prims_ll.h
+2 -2
View File
@@ -72,9 +72,9 @@ private:
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
if (nthreads != WARP_SIZE)
#if defined(__gfx942__) || (defined(__gfx950__) && defined(HIP_HOST_UNCACHED_MEMORY))
barrier_by_group_block();
barrier_generic(__threadfence_block(), nthreads, barrier_next, barriers);
#else
barrier_by_group();
barrier_generic(__threadfence(), nthreads, barrier_next, barriers);
#endif
#else
if (nthreads == WARP_SIZE) {
projects/rccl/src/device/prims_ll128.h
+2 -2
View File
@@ -77,9 +77,9 @@ private:
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
if (nthreads != WARP_SIZE)
#if defined(__gfx942__) || defined(__gfx950__)
barrier_by_group_block();
barrier_generic(__threadfence_block(), nthreads, barrier_next, barriers);
#else
barrier_by_group();
barrier_generic(__threadfence(), nthreads, barrier_next, barriers);
#endif
#else
barrier_sync(15-group, nthreads);
projects/rccl/src/device/prims_simple.h
+11 -4
View File
@@ -63,6 +63,8 @@ class Primitives<
uint32_t* next_hdp_reg;
uint64_t* barriers;
uint64_t barrier_next = 0;
uint64_t* barriers_pat;
uint64_t barrier_next_pat = 0;
int repeat;
#if defined(ENABLE_NPKIT)
@@ -80,9 +82,9 @@ private:
__syncwarp();
else
#if defined(__gfx942__) || defined(__gfx950__)
barrier_by_group_block();
barrier_generic(__threadfence_block(), nworkers, barrier_next, barriers);
#else
barrier_by_group();
barrier_generic(__threadfence(), nworkers, barrier_next, barriers);
#endif
}
inline __device__ void subBarrier() {
@@ -92,7 +94,11 @@ private:
}
inline __device__ void patBarrier() {
barrier();
#if defined(__gfx942__) || defined(__gfx950__)
barrier_generic(__threadfence_block(), NCCL_PAT_NWORKERS, barrier_next_pat, barriers_pat);
#else
barrier_generic(__threadfence(), NCCL_PAT_NWORKERS, barrier_next_pat, barriers_pat);
#endif
}
inline __device__ void barrierAny() {
@@ -743,8 +749,9 @@ public:
tid(tid), nthreads(nthreads), tidInBlock(threadIdx.x), group(group),
stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_) {
// For send operations, we need an extra warp to overlap the threadfence and the copy
barriers = &ncclShmem.groups[group].barrier;
// PAT uses the same barrier for each group
barriers_pat = &ncclShmem.barrier_pat;
this->nworkers = nthreads;
int peer = -1;
projects/rccl/src/device/reduce_scatter.h
+5 -7
View File
@@ -170,7 +170,6 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_L
template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE> {
__device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
#if __CUDA_ARCH__ >= 600
using Proto = ProtoSimple<1, 1>;
const int nranks = ncclShmem.comm.nRanks;
const int rank = ncclShmem.comm.rank;
@@ -192,8 +191,8 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SI
int step = 0;
while (1) {
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
cuda::atomic_ref<int, cuda::thread_scope_block> poll(ps->flags);
while (poll.load(cuda::memory_order_acquire) != 0) pollCount++; // Wait for workers to be done with step 'step-NCCL_SHMEM_PAT_STEPS'
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) != 0) pollCount++; // Wait for workers to be done with step 'step-NCCL_SHMEM_PAT_STEPS'
patAlgo.getNextOp(ps);
int last = ps->last;
step++;
@@ -217,16 +216,15 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SI
int step = group;
while(1) {
struct ncclPatStep* ps = shmem->patSteps+(step%NCCL_SHMEM_PAT_STEPS);
cuda::atomic_ref<int, cuda::thread_scope_block> poll(ps->flags);
while (poll.load(cuda::memory_order_acquire) == 0) pollCount++; // Wait for compute thread
int* poll = &ps->flags;
while (__hip_atomic_load(poll, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) == 0) pollCount++; // Wait for compute thread
int last = ps->last;
prims.patReduce(ps, shmem);
if (tidInGroup == 0) poll.store(0, cuda::memory_order_release); // Return element to compute thread
if (tidInGroup == 0) __hip_atomic_store(poll, 0, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_WORKGROUP); // Return element to compute thread
if (last) break;
step += nGroups;
}
}
#endif
}
};