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

Support fused all reduce and elementwise operations (#1729)

Browse Source 
* Support fused all reduce and elementwise operations

Add additional "acc" parameter to RCCL Replayer logs

Add flag which indicates availability of new API

* Fix Recorder json parsing

* Remove unreachable code

* Remove extra acc pointer check

* .

* Revert "[DEVICE] Adding ability to choose unroll factor at runtime (#1734)"

This reverts commit 4cadf3597c.

* Use noinline to reduce kernels linking time

* Don't use noinline for gfx942 and gfx950 to avoid perf regression

---------

Co-authored-by: AtlantaPepsi <timhu102@amd.com>
Co-authored-by: BertanDogancay <bertan.dogancay@gmail.com>

[ROCm/rccl commit: 9a4213356d]
This commit is contained in:
Wenkai DuAtlantaPepsiBertanDogancay
 authored and GitHub committed 2025-07-23 09:04:17 -07:00
1 parent cbb648505a
commit caff9764d3
24 files changed
+656 -231
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

projects/rccl/src/enqueue.cc
+29 -36
View File
@@ -28,31 +28,30 @@
using namespace rccl;
/* [RCCL] Determine which GPU kernel to execute */
void* rcclGetKernelIndex(int unroll, bool useCollTrace, struct ncclTaskColl* task)
{
// At this time, unroll factor is controlled only by passed in unroll argument
// After more investigation, this may be further tuned by the actual task being processed
struct ncclKernelMatch {
void* kernelFn;
bool specialized;
};
#ifdef ENABLE_COLLTRACE
int numKernels = sizeof(rcclKernelTable) / sizeof(rcclKernelTable[0]) / 2;
int firstKernel = useCollTrace ? numKernels : 0;
#define ncclGetKernelIndex(p_comm) ((p_comm)->unroll + ((p_comm)->collTraceEnabled ? 3 : 0))
static ncclKernelMatch const ncclKerns[6] = {
{(void *)ncclDevKernel_Generic_1, true},
{(void *)ncclDevKernel_Generic_2, true},
{(void *)ncclDevKernel_Generic_4, true},
{(void *)ncclDevKernelDebug_Generic_1, true},
{(void *)ncclDevKernelDebug_Generic_2, true},
{(void *)ncclDevKernelDebug_Generic_4, true}
};
#else
int numKernels = sizeof(rcclKernelTable) / sizeof(rcclKernelTable[0]);
int firstKernel = 0;
#define ncclGetKernelIndex(p_comm) ((p_comm)->unroll)
static ncclKernelMatch const ncclKerns[3] = {
{(void*)ncclDevKernel_Generic_1, true},
{(void*)ncclDevKernel_Generic_2, true},
{(void*)ncclDevKernel_Generic_4, true}
};
#endif
// Check if the requested unroll exists
for (int kernelIdx = 0; kernelIdx < numKernels; kernelIdx++) {
if (rcclKernelTable[firstKernel + kernelIdx].unroll == unroll) {
return rcclKernelTable[firstKernel + kernelIdx].funcPtr;
}
}
// If does not match, return null
return nullptr;
}
static int rcclProtoGrainSize(int proto, ncclComm *comm){
switch (proto) {
case NCCL_PROTO_LL: return 16;
@@ -82,7 +81,7 @@ NCCL_PARAM(L1SharedMemoryCarveout, "L1_SHARED_MEMORY_CARVEOUT", 0);
// Returns maximum kernel stack size of all CUDA kernels
ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* maxStackSize) {
constexpr int KernelCount = sizeof(rcclKernelTable)/sizeof(rcclKernelTable[0]);
constexpr int KernelCount = sizeof(ncclKerns)/sizeof(ncclKerns[0]);
ncclResult_t result = ncclSuccess;
if (maxStackSize) *maxStackSize = 0;
@@ -95,7 +94,7 @@ ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* ma
int ncclMaxSharedMem = rcclShmemDynamicSize(cudaArch, WarpSize);
for (int k=0; k < KernelCount; k++) {
void* fn = rcclKernelTable[k].funcPtr;
void* fn = ncclKerns[k].kernelFn;
cudaFuncAttributes attr = {0};
if (fn == nullptr) continue;
@@ -356,6 +355,7 @@ ncclResult_t ncclTasksRegAndEnqueue(struct ncclComm* comm) {
devWork.sendbuff = (void*)task->sendbuff;
devWork.recvbuff = (void*)task->recvbuff;
devWork.acc = (void*)task->acc;
devWork.sendbuffOffset = task->sendbuffOffset;
devWork.recvbuffOffset = task->recvbuffOffset;
devWork.sendbuffRmtAddrs = task->sendbuffRmtAddrs;
@@ -806,12 +806,8 @@ static ncclResult_t scheduleCollTasksToPlan(
//plan->channelMask.masks[channelId/64] |= (2ull<<devWork->channelHi) - (1ull<<devWork->channelLo);
plan->threadPerBlock = std::max(plan->threadPerBlock, 192 /* 3*WARP_SIZE */);
if (!plan->kernelSpecialized) {
#ifdef ENABLE_COLLTRACE
plan->kernelFn = rcclGetKernelIndex(comm->unroll, comm->collTraceEnabled);
#else
plan->kernelFn = rcclGetKernelIndex(comm->unroll, false);
#endif
plan->kernelSpecialized = true;
plan->kernelFn = ncclKerns[ncclGetKernelIndex(comm)].kernelFn;
plan->kernelSpecialized = ncclKerns[ncclGetKernelIndex(comm)].specialized;
}
if (comm->rank == 0) {
@@ -1127,12 +1123,8 @@ static ncclResult_t scheduleP2pTasksToPlan(
plan->threadPerBlock = std::max(plan->threadPerBlock, NCCL_MAX_NTHREADS);
if (!plan->kernelSpecialized) {
#ifdef ENABLE_COLLTRACE
plan->kernelFn = rcclGetKernelIndex(comm->unroll, comm->collTraceEnabled);
#else
plan->kernelFn = rcclGetKernelIndex(comm->unroll, false);
#endif
plan->kernelSpecialized = true;
plan->kernelFn = ncclKerns[ncclGetKernelIndex(comm)].kernelFn;
plan->kernelSpecialized = ncclKerns[ncclGetKernelIndex(comm)].specialized;
}
// Compute how much to split operations
@@ -2505,6 +2497,7 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
t->sliceSteps = info->sliceSteps;
t->eActivationMask = __atomic_load_n(&ncclProfilerEventMask, __ATOMIC_RELAXED);
t->opCount = comm->opCount;
t->acc = info->acc;
planner->nTasksColl += 1;
ncclTaskCollSorterInsert(&planner->collSorter, t, t->trafficBytes);
@@ -2554,8 +2547,8 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
}
NCCLCHECKGOTO(ArgsCheck(info), ret, fail);
INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p count %zu datatype %d op %d root %d comm %p [nranks=%d] stream %p task %d globalrank %d",
info->opName, info->comm->opCount, info->sendbuff, info->recvbuff, info->count,
INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p acc %p count %zu datatype %d op %d root %d comm %p [nranks=%d] stream %p task %d globalrank %d",
info->opName, info->comm->opCount, info->sendbuff, info->recvbuff, info->acc, info->count,
info->datatype, info->op, info->root, info->comm, info->comm->nRanks, info->stream,
info->comm->planner.nTasksP2p + info->comm->planner.nTasksColl,
info->comm->localRankToRank[info->comm->localRank]);