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

Merge remote-tracking branch 'nccl/master' into develop

Browse Source 
[ROCm/rccl commit: a6bf9bfc9e]
This commit is contained in:
BertanDogancay
2025-04-23 20:46:36 -07:00
parent 07620c7efd 8f0723e904
commit d045d0ca23
118 changed files with 12789 additions and 3358 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/device/all_gather.h
+137 -102
View File
@@ -10,7 +10,7 @@
#include "primitives.h"
namespace {
template<typename T, typename RedOp, typename Proto>
template<typename T, typename RedOp, typename Proto, bool isNetOffload = false>
#if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx942__) && !defined(__gfx950__)
__device__ void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
#else
@@ -20,10 +20,10 @@ namespace {
ncclRing *ring = &ncclShmem.channel.ring;
const int *ringRanks = ring->userRanks;
const int nranks = ncclShmem.comm.nRanks;
size_t count, partOffset, partCount, chunkCount;
ssize_t count, partOffset, partCount, chunkCount;
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &count, &partOffset, &partCount, &chunkCount);
size_t offset;
size_t dataOffset;
ssize_t offset;
ssize_t dataOffset;
int nelem;
int rankDest;
@@ -51,103 +51,126 @@ namespace {
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
int workNthreads;
T *inputBuf = (T*)work->sendbuff;
T *outputBuf = (T*)work->recvbuff;
// Coverity reports that the callee treats &ring->next as an array. However, due to the use of
// FanSymmetric<1>, only the first element is ever accessed, so it's fine.
// coverity[callee_ptr_arith:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex, work->connIndex, work);
// If isNetOffload == true, we only use 1 warp to drive Ring algo/network communication
// and the rest of warps proceed to copy src data into dst buffer in parallel when AG
// is not in-place.
if (isNetOffload) {
workNthreads = WARP_SIZE;
chunkCount = NCCL_MAX_NET_SIZE;
} else {
workNthreads = nthreads;
}
if (tid < workNthreads) {
// Coverity reports that the callee treats &ring->next as an array. However, due to the use of
// FanSymmetric<1>, only the first element is ever accessed, so it's fine.
// coverity[callee_ptr_arith:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0, isNetOffload> prims
(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex, work->connIndex, work, NULL, isNetOffload ? NCCL_MAX_NET_SIZE : 0);
#if defined(ENABLE_NPKIT)
if (tid == 0) {
prims.npKitCtxIdx = npKitCtxIdx;
}
if (tid == 0) {
prims.npKitCtxIdx = npKitCtxIdx;
}
#endif
for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
/////////////// begin AllGather steps ///////////////
nelem = min(chunkCount, partCount - elemOffset);
dataOffset = partOffset + elemOffset;
for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
/////////////// begin AllGather steps ///////////////
nelem = min(chunkCount, partCount - elemOffset);
dataOffset = partOffset + elemOffset;
// step 0: push data to next GPU
rankDest = ringRanks[0];
offset = dataOffset + rankDest * count;
// step 0: push data to next GPU
rankDest = ringRanks[0];
offset = dataOffset + rankDest * count;
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_SEND_ENTRY)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
#endif
if (inputBuf + dataOffset == outputBuf + offset) { // In place
prims.directSend(dataOffset, offset, nelem);
} else {
prims.directCopySend(dataOffset, offset, nelem);
}
if ((inputBuf + dataOffset == outputBuf + offset) || isNetOffload) { // In place or onePPN
prims.directSend(dataOffset, offset, nelem);
} else {
prims.directCopySend(dataOffset, offset, nelem);
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_SEND_EXIT)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_ENTRY)
if (tid == 0 && nranks > 2) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_ENTRY, nelem*(nranks-2)*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
if (tid == 0 && nranks > 2) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_ENTRY, nelem*(nranks-2)*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
#endif
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ringRanks[nranks-j];
offset = dataOffset + rankDest * count;
prims.directRecvCopyDirectSend(offset, nelem);
}
// k-2 steps: copy to next GPU
for (int j = 1; j < nranks - 1; ++j) {
rankDest = ringRanks[nranks - j];
offset = dataOffset + rankDest * count;
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_EXIT)
if (tid == 0 && nranks > 2) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_EXIT, nelem*(nranks-2)*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
if (tid == 0 && nranks > 2) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_EXIT, nelem*(nranks-2)*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
// Make final copy from buffer to dest.
rankDest = ringRanks[1];
offset = dataOffset + rankDest * count;
// Make final copy from buffer to dest.
rankDest = ringRanks[1];
offset = dataOffset + rankDest * count;
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_ENTRY)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
prims.npKitDataProcessTotalTime = 0;
}
#endif
// Final wait/copy.
prims.directRecv(offset, offset, nelem);
// Final wait/copy.
prims.directRecv(offset, offset, nelem);
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
}
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_EXIT)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_EXIT, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_EXIT, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
} else if (inputBuf != outputBuf + ringRanks[0] * count) {
inputBuf = inputBuf + partOffset;
outputBuf = outputBuf + partOffset + ringRanks[0] * count;
reduceCopy<COLL_UNROLL, RedOp, T, 0, 1, 1, 0, 1, 1, /*PreOpSrcs=*/0>
(tid - workNthreads, nthreads - workNthreads, work->redOpArg, &work->redOpArg, false, 1, (void**)&inputBuf, 1, (void**)&outputBuf, partCount);
}
#if !defined(__HIP_PLATFORM_AMD__) && !defined(__HIPCC__)
// we have to wait for all warps before we can proceed to the next work;
// otherwise, we can have contention if next work will use the outputBuf
// in this work. We use bar 14 to avoid conflicts with prims barrier and
// __syncthread().
if (isNetOffload) barrier_sync(14, nThreads);
#endif
}
}
@@ -155,8 +178,15 @@ namespace {
template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
using Proto = ProtoSimple<ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS>;
runRing<T, RedOp, Proto>(tid, nthreads, work);
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
bool isNetOffload = false;
#else
bool isNetOffload = work->isOneRPN && work->netRegUsed;
#endif
if (isNetOffload)
runRing<T, RedOp, ProtoSimple<1, 1>, true>(tid, nthreads, work);
else
runRing<T, RedOp, ProtoSimple<ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS>, false>(tid, nthreads, work);
}
};
@@ -186,7 +216,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE
T *inputBuf = (T*)work->sendbuff;
T *outputBuf = (T*)work->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, NULL, NULL, inputBuf, outputBuf, work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, nullptr, false, false, 0, primsModePatAg);
(tid, nthreads, NULL, NULL, inputBuf, outputBuf, work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, nullptr, nullptr, 0, primsModePatAg);
PatAGAlgorithm<T> patAlgo(chunkCount*sizeof(T), NCCL_STEPS, channelOffset, channelOffset + channelCount, count, chunkCount, rank, nranks);
int last = 0;
@@ -227,6 +257,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
nelem = min(chunkCount, channelCount - elemOffset);
prims.gather(offset, nvls->nHeads * count, nelem, count, -1, 0);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndBcast) {
// Bcast through NVLS
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
@@ -238,6 +269,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
nelem = min(chunkCount, channelCount - elemOffset);
prims.send(offset, nelem);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
}
} else {
/* direct allgather */
@@ -294,11 +326,11 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
int part = ncclShmem.channelId - work->channelLo;
char* inbuf = (char*)work->sendbuff;
char* outbuf = (char*)work->recvbuff;
ssize_t sizePerRank = work->collnet.count*sizeof(T);
bool inPlace = (inbuf == outbuf + ncclShmem.comm.rank*sizePerRank);
ssize_t countPerRank = work->collnet.count*sizeof(T);
bool inPlace = (inbuf == outbuf + ncclShmem.comm.rank*countPerRank);
ssize_t railAllBeg = min(railGridOffset + part*chunkSize, nNodes*sizePerRank);
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes*sizePerRank);
ssize_t railAllBeg = min(railGridOffset + part*chunkSize, nNodes*countPerRank);
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes*countPerRank);
int railAllSize = railAllEnd - railAllBeg;
if (tid < nDsts) dstSizes[tid] = railAllSize;
@@ -311,15 +343,15 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
if (rail == nRails) rail = 0;
}
do {
int node = railAllBeg/sizePerRank;
int node = railAllBeg/countPerRank;
int railAllOffset = 0;
while (railAllOffset < railAllSize) {
ssize_t railOneBeg = node*sizePerRank;
ssize_t railOneEnd = railOneBeg + sizePerRank;
ssize_t railOneBeg = node*countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t railOneOffset = (railAllBeg+railAllOffset) - railOneBeg;
int delta = min(railAllEnd, railOneEnd) - (railAllBeg+railAllOffset);
int rank = ncclShmem.comm.collNetDenseToUserRank[node*nRails + rail];
ssize_t userOneBeg = rank*sizePerRank + railOneOffset;
ssize_t userOneBeg = rank*countPerRank + railOneOffset;
int outIsDst = (inPlace && rank == ncclShmem.comm.rank) ? 0 : 1;
if (nSrcs != 0 && outIsDst+nDsts != 0) {
reduceCopy<ncclCollUnroll(), RedOp, T,
@@ -327,15 +359,15 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
/*MultimemDsts=*/0, 0+MinDsts, 1+MaxDsts,
/*PreOpSrcs=*/0>
(tid, tn, 0, nullptr, false,
/*nSrcs=*/1, [=]__device__(int s/*==0*/) -> void* {
return work->regUsed && (recvDirectFlag & NCCL_DIRECT_READ) ? (char*)srcPtrs[src] + userOneBeg : (char*)srcPtrs[src] + railAllOffset;
},
/*nDsts=*/outIsDst+nDsts, [=]__device__(int d) -> void* {
return d < outIsDst ? outbuf + userOneBeg
: work->regUsed && (sendDirectFlag & NCCL_DIRECT_WRITE) ? (char*)dstPtrs[d-outIsDst] + userOneBeg
: (char*)dstPtrs[d-outIsDst] + railAllOffset;
},
delta);
/*nSrcs=*/1, [=]__device__(int s/*==0*/) -> void* {
return work->regUsed && (recvDirectFlag & NCCL_P2P_READ) ? (char*)srcPtrs[src] + userOneBeg : (char*)srcPtrs[src] + railAllOffset;
},
/*nDsts=*/outIsDst+nDsts, [=]__device__(int d) -> void* {
return d < outIsDst ? outbuf + userOneBeg
: work->regUsed && (sendDirectFlag & NCCL_P2P_WRITE) ? (char*)dstPtrs[d-outIsDst] + userOneBeg
: (char*)dstPtrs[d-outIsDst] + railAllOffset;
},
delta);
}
railAllOffset += delta;
node += 1;
@@ -352,8 +384,9 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
const int nChannels = work->channelHi - work->channelLo + 1;
struct ncclDirect* direct = &ncclShmem.channel.collnetDirect;
int const &nNodes = ncclShmem.comm.nNodes;
ssize_t sizePerRank = work->collnet.count*sizeof(T);
ssize_t countPerRank = work->collnet.count;
size_t chunkSize = work->collnet.chunkCount;
const int hasDn = (direct->down[0] >= 0) ? 1 : 0;
bool isMultiRail = (direct->nHeads > 1);
int nWarps1 = 1;
int nWarps2 = (isMultiRail ? 2 : 1);
@@ -367,9 +400,12 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
int tn = nWarps1*WARP_SIZE;
if (tid < tn) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (tid == 0) {
int steps = (int)divUp(nNodes * sizePerRank * sizeof(T), NCCL_MAX_COLLNET_SIZE);
// If this rank has local peers (i.e, hasDn == true), we cannot offload all data to network.
// In this case, steps should be computed based on chunkSize and so on; otherwise, we just
// bump the step by 1 to kick off collnet progress.
int steps = hasDn ? (int)divUp(nNodes * countPerRank, nChannels * chunkSize) : 1;
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>::sendPeerNotify(direct->out, 1, steps);
}
__syncwarp();
@@ -378,11 +414,11 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
prims(tid, tn, nullptr, &direct->out, work->sendbuff, nullptr,
/*redOpArg=*/0, 0 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * sizePerRank; railGridOffset += nChannels * chunkSize) {
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkSize) {
ssize_t railAllBeg = railGridOffset + part * chunkSize;
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes * sizePerRank);
ssize_t railOneBeg = ncclShmem.comm.node * sizePerRank;
ssize_t railOneEnd = railOneBeg + sizePerRank;
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes * countPerRank);
ssize_t railOneBeg = ncclShmem.comm.node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t beg = max(railAllBeg, railOneBeg);
ssize_t end = min(railAllEnd, railOneEnd);
prims.send(beg - railOneBeg, max(ssize_t(0), end - beg));
@@ -394,10 +430,9 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
tn = nWarps2*WARP_SIZE;
if (tid < tn) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed && !hasDn) {
if (tid == 0) {
int steps = (int)divUp(nNodes * sizePerRank * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>::recvPeerNotify(direct->out, 0, steps);
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>::recvPeerNotify(direct->out, 0, 1);
}
__syncwarp();
} else {
@@ -405,7 +440,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/1, Proto, 0>
prims(tid, tn, &direct->out, direct->heads + 1, nullptr, work->recvbuff,
/*redOpArg=*/0, 1 * Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * sizePerRank; railGridOffset += nChannels * chunkSize) {
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkSize) {
Scatterer</*BcastSendNotRecv=*/true> scat;
scat.work = work;
scat.chunkSize = chunkSize;
@@ -423,7 +458,7 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/1, Proto, 0>
prims(tid, tn, direct->heads+1, nullptr, nullptr, work->recvbuff,
/*redOpArg=*/0, 2*Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t railGridOffset=0; railGridOffset < nNodes*sizePerRank; railGridOffset += nChannels*chunkSize) {
for (ssize_t railGridOffset=0; railGridOffset < nNodes*countPerRank; railGridOffset += nChannels*chunkSize) {
Scatterer</*BcastSendNotRecv=*/false> scat;
scat.work = work;
scat.chunkSize = chunkSize;
projects/rccl/src/device/all_reduce.h
+86 -83
View File
@@ -166,7 +166,7 @@ namespace {
chunkOffset = chunk * chunkCount;
offset = gridOffset + elemOffset + chunkOffset;
nelem = (int)min(chunkCount, remCount - chunkOffset);
prims.directRecvCopyDirectSend(offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_EXIT)
@@ -336,7 +336,7 @@ namespace {
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.directRecvCopyDirectSend(offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
}
@@ -535,7 +535,7 @@ namespace {
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
prims.directRecvCopyDirectSend(offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
}
@@ -599,22 +599,30 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
const int tidStartBcast = nThreadsGather;
const int tidStartScatter = tidStartBcast + nThreadsBcast;
const int tidStartReduce = tidStartScatter + nThreadsScatter;
using Proto = ProtoSimple<1, 1>;
if (tid >= tidStartScatter && tid < tidStartReduce && hasUp) {
// Scatter
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/1, Proto, 0>
prims(tid-tidStartScatter, nThreadsScatter, NULL, direct->up, work->sendbuff, work->recvbuff,
work->redOpArg, 2*Proto::MaxGroupWidth, 1, 1);
work->redOpArg, 2*Proto::MaxGroupWidth, 1, 1, work);
ssize_t offsetBase, peerOffset;
ssize_t maxNelems;
if (work->netRegUsed) {
offsetBase = bid * chunkSize;
maxNelems = size; // never be the min
peerOffset = nChannels * chunkSize;
} else {
offsetBase = bid * direct->nHeads * chunkSize;
maxNelems = direct->nHeads * chunkSize;
peerOffset = chunkSize;
}
// For collnet UB case, we need to organize buffers differently for contiguous buffer access
// across channels. This access pattern should be consistent with code in coll_net.cc
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*direct->nHeads*chunkSize;
int nelem = min(direct->nHeads*chunkSize, size-offset);
if (work->regUsed) {
prims.directScatter(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
} else {
prims.scatter(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
}
ssize_t offset = gridOffset + offsetBase;
ssize_t nelem = min(maxNelems, size - offset);
prims.scatter(offset, nelem, chunkSize, peerOffset, direct->headRank, direct->shift);
}
// Coverity complains about a possible overrun inside the destructor of "prims", but that's actually
// a false positive.
@@ -622,24 +630,20 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
} else if (tid >= tidStartReduce && direct->out != -1) {
if (hasDn) {
// Reduce, send to network
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/1, Proto, 0>
prims(tid-tidStartReduce, nThreadsReduce, direct->down, &direct->out, work->sendbuff, work->recvbuff,
work->redOpArg, 3*Proto::MaxGroupWidth, 1, 1);
work->redOpArg, 3*Proto::MaxGroupWidth, 1, 1, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
if (work->regUsed) {
prims.directRecvReduceSend(offset, nelem);
} else {
prims.recvReduceSend(offset, nelem);
}
ssize_t offset = work->netRegUsed ? gridOffset + (bid + direct->headRank * nChannels) * chunkSize
: gridOffset + (bid * direct->nHeads + direct->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.recvReduceDirectSend(offset, offset, nelem);
}
} else {
// Directly send to network
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (tid == tidStartReduce) {
int steps = (int)divUp(size * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>::sendPeerNotify(direct->out, 1, steps);
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>::sendPeerNotify(direct->out, 1, 1);
}
__syncwarp();
} else {
@@ -647,8 +651,8 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
prims(tid-tidStartReduce, nThreadsReduce, nullptr, &direct->out, work->sendbuff, work->recvbuff,
work->redOpArg, 3*Proto::MaxGroupWidth, 1, 1);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
ssize_t offset = gridOffset + (bid * direct->nHeads + direct->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.send(offset, nelem);
}
}
@@ -657,11 +661,22 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
// Gather
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, direct->up, NULL, work->sendbuff, work->recvbuff,
work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, work);
work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, work);
ssize_t offsetBase, peerOffset;
ssize_t maxNelems;
if (work->netRegUsed) {
offsetBase = bid * chunkSize;
maxNelems = size; // never be the min
peerOffset = nChannels * chunkSize;
} else {
offsetBase = bid * direct->nHeads * chunkSize;
maxNelems = direct->nHeads * chunkSize;
peerOffset = chunkSize;
}
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*direct->nHeads*chunkSize;
int nelem = min(direct->nHeads*chunkSize, size-offset);
prims.directGather(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
ssize_t offset = gridOffset + offsetBase;
ssize_t nelem = min(maxNelems, size - offset);
prims.directGather(offset, nelem, chunkSize, peerOffset, direct->headRank, direct->shift);
}
} else if (tid >= tidStartBcast && tid < tidStartScatter && direct->out != -1) {
if (hasDn) {
@@ -673,15 +688,15 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
prims(tid-tidStartBcast, nThreadsBcast, &direct->out, direct->down, work->sendbuff, work->recvbuff,
work->redOpArg, 1*Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
prims.recvCopyDirectSend(offset, nelem, /*postOp=*/true);
ssize_t offset = work->netRegUsed ? gridOffset + (bid + direct->headRank * nChannels) * chunkSize
: gridOffset + (bid * direct->nHeads + direct->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.directRecvCopyDirectSend(offset, offset, nelem, /*postOp=*/true);
}
} else {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (tid == tidStartBcast) {
int steps = (int)divUp(size * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>::recvPeerNotify(direct->out, 0, steps);
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>::recvPeerNotify(direct->out, 0, 1);
}
__syncwarp();
} else {
@@ -725,8 +740,6 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
ssize_t gridOffset, channelCount, chunkSize;
ncclCollCbdPart(work, ncclShmem.channelId, NCCL_PROTO_SIMPLE, sizeof(T), (ssize_t*)nullptr, &gridOffset, &channelCount, &chunkSize);
const ssize_t loopCount = nvls->nHeads * chunkSize;
ssize_t offset;
int nelem;
int remCount = channelCount%(nvls->nHeads*chunkSize);
int lastChunkSize = alignUp(divUp(remCount, nvls->nHeads), 16384/sizeof(T));
@@ -738,8 +751,8 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
work->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount) {
if (channelCount - elemOffset < loopCount) chunkSize = lastChunkSize;
offset = gridOffset + elemOffset;
nelem = work->regUsed ? 0 : min(loopCount, channelCount - elemOffset);
ssize_t offset = gridOffset + elemOffset;
int nelem = work->regUsed ? 0 : min(loopCount, channelCount - elemOffset);
prims.scatter(offset, nelem, chunkSize, chunkSize, -1, 0);
}
} else if (tid < tidEndGather) {
@@ -750,8 +763,8 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
work->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount) {
if (channelCount - elemOffset < loopCount) chunkSize = lastChunkSize;
offset = gridOffset + elemOffset;
nelem = work->regUsed ? 0 : min(loopCount, channelCount - elemOffset);
ssize_t offset = gridOffset + elemOffset;
int nelem = work->regUsed ? 0 : min(loopCount, channelCount - elemOffset);
prims.gather(offset, nelem, chunkSize, chunkSize, -1, 0);
}
} else if (tid < tidEndReduce) {
@@ -761,7 +774,8 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->down, NULL, NULL,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount) {
ssize_t chunkOffset;
ssize_t chunkOffset, offset;
int nelem;
if (channelCount - elemOffset < loopCount) chunkSize = lastChunkSize;
chunkOffset = elemOffset + nvls->headRank * chunkSize;
offset = gridOffset + chunkOffset;
@@ -787,6 +801,7 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
int nelem = work->regUsed ? 0 : min(nvls->nHeads * chunkSize, size - offset);
prims.scatter(offset, nelem, chunkSize, chunkSize, -1, 0);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndGather) {
// Gather
using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
@@ -795,38 +810,23 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
work->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid * nvls->nHeads * chunkSize;
int nelem = work->regUsed ? 0 :min(nvls->nHeads * chunkSize, size - offset);
int nelem = work->regUsed ? 0 : min(nvls->nHeads * chunkSize, size - offset);
prims.gather(offset, nelem, chunkSize, chunkSize, -1, 0);
}
} else if (tid < tidEndReduce && nvls->headRank != -1) {
if (!hasOut) {
// Reduce, broadcast through NVLS
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 1>;
// Coverity complains about a possible overrun inside the class below, but that's actually
// a false positive.
// coverity[identity_transfer:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->down, NULL, NULL,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.directRecvDirectSend(offset, offset, nelem);
}
} else {
// Reduce, send to network
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
// Coverity complains about a possible overrun inside the class below, but that's actually
// a false positive.
// coverity[identity_transfer:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->out, NULL, NULL,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 1, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.directRecvDirectSend(offset, offset, nelem);
}
// Reduce, send to network
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
// Coverity complains about a possible overrun inside the class below, but that's actually
// a false positive.
// coverity[identity_transfer:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->out, NULL, work->recvbuff,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 1, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = work->regUsed && work->netRegUsed ? gridOffset + (nvls->headRank * nChannels + bid) * chunkSize
: gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.directRecvDirectSend(offset, offset, nelem);
}
} else if (tid < tidEndBcast && nvls->headRank != -1) {
// Recv from network, broadcast
@@ -835,10 +835,11 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPL
// a false positive.
// coverity[identity_transfer:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
prims(tid - tidEndReduce, nThreadsBcast, &nvls->out, &nvls->down, NULL, NULL,
prims(tid - tidEndReduce, nThreadsBcast, &nvls->out, &nvls->down, NULL, work->recvbuff,
work->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, work);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
ssize_t offset = work->regUsed && work->netRegUsed ? gridOffset + (nvls->headRank * nChannels + bid) * chunkSize
: gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
int nelem = min(chunkSize, size - offset);
prims.directRecvDirectSend(offset, offset, nelem);
}
@@ -991,10 +992,9 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PR
if (tid < nthreadsSplit) {
if (recv == -1) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (groupTid == 0) {
int steps = (int)divUp(size * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>::sendPeerNotify(send, connIndex, steps);
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>::sendPeerNotify(send, connIndex, 1);
}
__syncwarp();
} else {
@@ -1004,8 +1004,10 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PR
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid * int(chunkSize);
int nelem = min(chunkSize, size - offset);
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
prims.directSend(offset, offset, nelem);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
}
} else {
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
@@ -1014,18 +1016,19 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PR
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid * int(chunkSize);
int nelem = min(chunkSize, size - offset);
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
prims.directRecvReduceDirectSend(offset, offset, nelem);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
}
}
else {
if (recv == nranks) {
// I'm the first in the broadcast chain, I need to perform the division (postOp)
if (send == -1) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (groupTid == 0) {
int steps = (int)divUp(size * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>::recvPeerNotify(recv, connIndex, steps);
Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>::recvPeerNotify(recv, connIndex, 1);
}
__syncwarp();
} else {
@@ -1051,7 +1054,7 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PR
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid * int(chunkSize);
int nelem = min(chunkSize, size - offset);
prims.directRecvCopyDirectSend(offset, nelem, /*postOp*/true);
prims.directRecvCopyDirectSend(offset, offset, nelem, /*postOp*/true);
}
}
} else {
@@ -1071,7 +1074,7 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PR
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecvCopyDirectSend(offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
}
}
projects/rccl/src/device/broadcast.h
+37 -23
View File
@@ -20,13 +20,15 @@ namespace {
const int rank = ring->userRanks[0];
const int nextRank = ring->userRanks[1];
const int root = work->root;
size_t size;
size_t chunkCount;
size_t channelCount;
size_t gridOffset;
ssize_t size;
ssize_t chunkCount;
ssize_t channelCount;
ssize_t gridOffset;
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &size, &gridOffset, &channelCount, &chunkCount);
size_t offset;
int nelem;
int workNthreads;
bool isNetOffload = work->isOneRPN && work->netRegUsed;
#if defined(ENABLE_NPKIT)
int npKitCtxIdx = bid;
@@ -55,39 +57,51 @@ namespace {
T *inputBuf = (T*)work->sendbuff;
T *outputBuf = (T*)work->recvbuff;
// Coverity reports that the callee treats &ring->next as an array. However, due to the use of
// FanSymmetric<1>, only the first element is ever accessed, so it's fine.
// coverity[callee_ptr_arith:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex, work->connIndex, work);
workNthreads = isNetOffload ? WARP_SIZE : nthreads;
if (tid < workNthreads) {
// Coverity reports that the callee treats &ring->next as an array. However, due to the use of
// FanSymmetric<1>, only the first element is ever accessed, so it's fine.
// coverity[callee_ptr_arith:FALSE]
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
prims(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex, work->connIndex, work);
#if defined(ENABLE_NPKIT)
if (tid == 0) {
prims.npKitCtxIdx = npKitCtxIdx;
}
if (tid == 0) {
prims.npKitCtxIdx = npKitCtxIdx;
}
#endif
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
if (rank == root) {
if (inputBuf == outputBuf) {
prims.directSend(offset, offset, nelem);
if (rank == root) {
if (inputBuf == outputBuf || isNetOffload) {
prims.directSend(offset, offset, nelem);
} else {
prims.directCopySend(offset, offset, nelem);
}
} else if (nextRank == root) {
prims.directRecv(offset, offset, nelem);
} else {
prims.directCopySend(offset, offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
} else if (nextRank == root) {
prims.directRecv(offset, offset, nelem);
} else {
prims.directRecvCopyDirectSend(offset, nelem);
}
} else if (inputBuf != outputBuf && rank == root) {
inputBuf = inputBuf + gridOffset;
outputBuf = outputBuf + gridOffset;
reduceCopy<COLL_UNROLL, RedOp, T, 0, 1, 1, 0, 1, 1, /*PreOpSrcs=*/0>
(tid - workNthreads, nthreads - workNthreads, work->redOpArg, &work->redOpArg, false, 1, (void**)&inputBuf, 1, (void**)&outputBuf, channelCount);
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_BROADCAST_RING_EXIT)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_BROADCAST_RING_EXIT, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
#if !defined(__HIP_PLATFORM_AMD__) && !defined(__HIPCC__)
if (isNetOffload) barrier_sync(14, nThreads);
#endif
}
}
projects/rccl/src/device/common.h
+5 -2
View File
@@ -59,8 +59,8 @@
collTrace->p2p.recvConnIndex = p2pWork->recvConnIndex; \
collTrace->p2p.sendProtoLL = p2pWork->sendProtoLL; \
collTrace->p2p.recvProtoLL = p2pWork->recvProtoLL; \
collTrace->p2p.sendRegistered = p2pWork->sendRegistered; \
collTrace->p2p.recvRegistered = p2pWork->recvRegistered; \
collTrace->p2p.sendRegistered = p2pWork->sendNetReg; \
collTrace->p2p.recvRegistered = p2pWork->recvNetReg; \
collTrace->p2pOpCount[0] = p2pWork->sendOpCount; \
collTrace->p2pOpCount[1] = p2pWork->recvOpCount; \
collTrace->type = (launch_type) | ncclCollTraceP2pElemType; \
@@ -628,6 +628,9 @@ __global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONST
__global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
#endif
#define DEFINE_ncclDevKernel_nop(suffix, coll, redop, ty, algo, proto, specializedFnId) \
__global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {}
#ifdef USE_INDIRECT_FUNCTION_CALL
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, unroll) \
__device__ void ncclDevFunc_##suffix() { \
projects/rccl/src/device/common_kernel.h
+14 -5
View File
@@ -67,19 +67,23 @@ __device__ __forceinline__ void reduceCopyPacks(
uintptr_t minSrcs[MinSrcs + !MinSrcs];
uintptr_t minDsts[MinDsts + !MinDsts];
#pragma unroll
for (int s=0; s < MinSrcs; s++)
for (int s=0; s < MinSrcs; s++) {
minSrcs[s] = cvta_to_global(srcPtrFn(s)) + threadBytesBehind;
}
#pragma unroll
for (int d=0; d < MinDsts; d++)
for (int d=0; d < MinDsts; d++) {
// Yes, for some template arguments this code will be unreachable. That's fine.
// coverity[dead_error_line]
minDsts[d] = cvta_to_global(dstPtrFn(d)) + threadBytesBehind;
}
// We dictate loop termination condition according to whether partial hunks
// can be handled or not.
while (Unroll==1 ? (BytePerPack <= threadBytesAhead) : (0 < nHunksAhead)) {
BytePack<BytePerPack> acc[Unroll];
// minSrcs[0] cannot be nullptr so we always process it
{ RedFn preFn(0 < PreOpSrcs ? preOpArgs[0] : 0);
#pragma unroll Unroll
for (int u=0; u < Unroll; u++) {
@@ -165,7 +169,8 @@ __device__ __forceinline__ void reduceCopyPacks(
}
}
for (int d=MinDsts; (MinDsts < MaxDsts) && (d < MaxDsts) && (d < nDsts); d++) {
uintptr_t dst = cvta_to_global(dstPtrFn(d)) + threadBytesBehind;
uintptr_t dstPtr = cvta_to_global(dstPtrFn(d));
uintptr_t dst = dstPtr + threadBytesBehind;
#pragma unroll Unroll
for (int u=0; u < Unroll; u++) {
st_global<BytePerPack>(dst, acc[u]);
@@ -175,11 +180,15 @@ __device__ __forceinline__ void reduceCopyPacks(
nWarps = nThreads/WARP_SIZE;
#pragma unroll
for (int s=0; s < MinSrcs; s++) minSrcs[s] += (nWarps-1)*BytePerHunk;
for (int s=0; s < MinSrcs; s++) {
minSrcs[s] += (nWarps-1)*BytePerHunk;
}
#pragma unroll
// Yes, for some template arguments this code will be unreachable. That's fine.
// coverity[dead_error_line]
for (int d=0; d < MinDsts; d++) minDsts[d] += (nWarps-1)*BytePerHunk;
for (int d=0; d < MinDsts; d++) {
minDsts[d] += (nWarps-1)*BytePerHunk;
}
threadBytesBehind += nWarps*BytePerHunk;
threadBytesAhead -= nWarps*BytePerHunk;
nHunksAhead -= nWarps;
projects/rccl/src/device/generate.py
+4 -4
View File
@@ -6,7 +6,7 @@ import subprocess
# Order of redops, tys, protos, algos must match src/include/device.h
all_colls = ["AllGather","AllReduce","AllToAllPivot","Broadcast","Reduce","ReduceScatter","SendRecv"]
all_redops = ["Sum","Prod","MinMax","PreMulSum","SumPostDiv"]
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16", "f8", "bf8"]
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16","f8e4m3","f8e5m2"]
all_protos = ["LL","LL128","SIMPLE"]
all_algos = ["TREE","RING"]
all_unroll = ["1", "2", "4"]
@@ -253,7 +253,7 @@ def equivalent_primary(coll, algo, proto, redop, ty, unroll):
unroll = str(coll_unroll)
if coll in ("AllReduce", "Reduce", "ReduceScatter"):
# map signed integer sum/prod to unsigned
if redop in ("Sum","Prod","PreMulSum") and ty[0]=="i":
if redop in ("Sum","Prod","PreMulSum","SumPostDiv") and ty[0]=="i":
ty = "u"+ty[1:]
# map signed integer min/max to unsigned for non-NVLS
elif redop=="MinMax" and ty[0]=="i" and ("NVLS" not in algo):
@@ -510,8 +510,8 @@ ty_to_cxx = {
"f32": "float",
"f64": "double",
"bf16": "hip_bfloat16",
"f8": "rccl_float8",
"bf8": "rccl_bfloat8",
"f8e4m3": "rccl_float8",
"f8e5m2": "rccl_bfloat8"
}
# Generate each <gensrc>/<impl>.cpp:
projects/rccl/src/device/network/unpack/unpack.h
+4
View File
@@ -33,17 +33,21 @@ inline __device__ void load64gpu(const uint64_t* ptr, uint64_t &v) {
// Map internal association of handle with group and peer index (called once at init time)
inline __device__ void ncclNetDeviceUnpackSetup(void* ohandle, const int group, const int index) {
struct unpackNetDeviceHandle* handle = (struct unpackNetDeviceHandle*) ohandle;
// coverity[index_parm:FALSE]
ncclShmem.groups[group].devicePlugin.unpack.g_meta[index] = handle->meta;
ncclShmem.devicePlugin.unpack.bounce_buf = handle->bounce_buf;
// coverity[index_parm:FALSE]
ncclShmem.groups[group].devicePlugin.unpack.head[index] = handle->head;
}
inline __device__ void ncclNetDeviceIncrementHead(const int group, const int index) {
// coverity[index_parm:FALSE]
ncclShmem.groups[group].devicePlugin.unpack.head[index]++;
}
inline __device__ void ncclNetDeviceSaveHead(void* ohandle, const int group, const int index) {
struct unpackNetDeviceHandle* handle = (struct unpackNetDeviceHandle*) ohandle;
// coverity[index_parm:FALSE]
handle->head = ncclShmem.groups[group].devicePlugin.unpack.head[index];
}
projects/rccl/src/device/onerank.cu
+4 -4
View File
@@ -71,13 +71,13 @@ ncclResult_t ncclLaunchOneRank(void* dst, void const* src, size_t nElts, struct
case ncclUint32: kernel = (void const*)&oneRankReduce<FuncPreMulSum<uint32_t>>; break;
case ncclInt64: kernel = (void const*)&oneRankReduce<FuncPreMulSum<int64_t>>; break;
case ncclUint64: kernel = (void const*)&oneRankReduce<FuncPreMulSum<uint64_t>>; break;
#if defined(RCCL_FLOAT8)
case ncclFloat8e4m3: kernel = (void const*)&oneRankReduce<FuncPreMulSum<rccl_float8>>; break;
case ncclFloat8e5m2: kernel = (void const*)&oneRankReduce<FuncPreMulSum<rccl_bfloat8>>; break;
#endif
case ncclFloat16: kernel = (void const*)&oneRankReduce<FuncPreMulSum<half>>; break;
#if defined(RCCL_BFLOAT16)
case ncclBfloat16: kernel = (void const*)&oneRankReduce<FuncPreMulSum<hip_bfloat16>>; break;
#endif
#if defined(RCCL_FLOAT8)
case ncclFp8E4M3: kernel = (void const*)&oneRankReduce<FuncPreMulSum<rccl_float8>>; break;
case ncclFp8E5M2: kernel = (void const*)&oneRankReduce<FuncPreMulSum<rccl_bfloat8>>; break;
#endif
case ncclFloat32: kernel = (void const*)&oneRankReduce<FuncPreMulSum<float>>; break;
case ncclFloat64: kernel = (void const*)&oneRankReduce<FuncPreMulSum<double>>; break;
projects/rccl/src/device/primitives.h
+5 -2
View File
@@ -142,7 +142,7 @@ struct FanSymmetric {
};
// The primitives class. Specialized per protocol in the other headers.
template<typename T, typename RedOp, typename Fan, int Direct, typename Proto, int P2p>
template<typename T, typename RedOp, typename Fan, int Direct, typename Proto, int P2p, bool isNetOffload = false>
class Primitives;
// Used by LL & LL128 to implement direct members in the naive way.
@@ -160,9 +160,12 @@ struct PrimitivesWithoutDirect {
__device__ void directCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
static_cast<RealPrimitives*>(this)->copySend(inpIx, outIx, eltN, postOp);
}
__device__ void directRecvCopyDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
__device__ void directRecvCopyDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
static_cast<RealPrimitives*>(this)->recvCopySend(outIx, eltN, /*postOp=*/false);
}
__device__ void directRecvDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return;
}
__device__ void recvReduceCopyDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
// Direct is only for the send part
static_cast<RealPrimitives*>(this)->recvReduceCopySend(inpIx, outIx, eltN, postOp);
projects/rccl/src/device/prims_ll.h
+3 -3
View File
@@ -10,9 +10,9 @@
#include "npkit/npkit.h"
#endif
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>> {
template<typename T, typename RedOp, typename Fan, int Direct, int P2p, bool isNetOffload>
class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p, isNetOffload>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p, isNetOffload>> {
// In the case of Fan::MaxRecv == 0, we need to force MaxRecv to 1 for this to compile
// This is because of a recv buffer which is allocated to MaxRecv length in send-only cases
projects/rccl/src/device/prims_ll128.h
+3 -3
View File
@@ -20,9 +20,9 @@
#endif
#endif
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>> {
template<typename T, typename RedOp, typename Fan, int Direct, int P2p, bool isNetOffload>
class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload>> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
projects/rccl/src/device/prims_simple.h
+158 -106
View File
@@ -21,31 +21,27 @@ enum primsMode {
};
template<typename T, typename RedOp, typename Fan, int Direct,
int SlicePerChunk, int StepPerSlice, int Unroll, int P2p, int MultimemSrcs, int MultimemDsts>
int SlicePerChunk, int StepPerSlice, int Unroll, int P2p, int MultimemSrcs, int MultimemDsts, bool isNetOffload>
class Primitives<
T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll, MultimemSrcs, MultimemDsts>, P2p
T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll, MultimemSrcs, MultimemDsts>, P2p, isNetOffload
> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
static constexpr int RoleInput = 0x01,
RoleOutput = 0x02,
RoleWaitRecv = 0x04,
RoleWaitSend = 0x08,
RolePostSend = 0x10,
RolePostRecv = 0x20,
Aborted = 0x40,
NetRegMode = 0x80,
ConnFifoEnabled = 0x100,
DirectWrite = 0x200,
DirectRead = 0x400,
PatMode = 0x800,
NvlsMinPolling = 0x1000,
NetDeviceUnpack = 0x2000,
AnyNetDeviceUnpack = 0x4000,
NvlsDirectRead = 0x8000,
NvlsDirectWrite = 0x10000,
IpcWrite = 0x20000,
IpcRead = 0x40000;
RoleOutput = 0x02,
RoleWaitRecv = 0x04,
RoleWaitSend = 0x08,
RolePostSend = 0x10,
RolePostRecv = 0x20,
Aborted = 0x40,
NetRegMode = 0x80,
ConnFifoEnabled = 0x100,
DirectWrite = 0x200,
DirectRead = 0x400,
PatMode = 0x800,
NvlsMinPolling = 0x1000,
NetDeviceUnpack = 0x2000,
AnyNetDeviceUnpack = 0x4000;
const int tid, tidInBlock;
const int nthreads;
int nworkers;
@@ -121,12 +117,9 @@ private:
template <int DirectRecv, int DirectSend, int Recv, int Send, int Src, int Dst>
__device__ __forceinline__ void waitPeer(intptr_t srcIx, intptr_t dstIx, int offset, int nelts) {
const bool isSendNotRecv = (Send && Recv) ? (flags & RoleWaitSend) : Send;
const bool noRecvWait = DirectRecv && Src && (flags & (DirectRead | IpcRead)); // no wait when directly reading from remote input
const bool noSendWait = DirectSend && (flags & (DirectRead|DirectWrite)); // no wait in empty send (e.g. directScatter) or direct remote write
// Yes, for some template arguments this code will be unreachable. That's fine.
// coverity[dead_error_line]
if (((flags & (Recv*RoleWaitRecv)) && !noRecvWait) ||
((flags & (Send*RoleWaitSend)) && !noSendWait)) {
if ((flags & (Recv * RoleWaitRecv)) || (flags & (Send * RoleWaitSend))) {
int spins = 0;
repeat = 50;
while (connStepCache + (isSendNotRecv ? NCCL_STEPS : 0) < step + StepPerSlice) {
@@ -143,27 +136,38 @@ private:
}
if (flags & (Recv*RoleWaitRecv | Send*RoleWaitSend)) {
if (flags & ConnFifoEnabled)
if ((flags & ConnFifoEnabled) && (flags & (Send * RoleWaitSend)))
connFifo[step%NCCL_STEPS].size = nelts*sizeof(T);
void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst)
: (ncclShmem.groups[group].srcs + Src);
if (flags & NetRegMode) {
// Do nothing
if (P2p) {
ptrs[index] = NULL;
} else {
if (isSendNotRecv) {
if (!Recv)
ptrs[index] = NULL;
else
ptrs[index] = (T*)ncclShmem.groups[group].userOutput + dstIx + offset;
} else {
ptrs[index] = (T*)ncclShmem.groups[group].userOutput + srcIx + offset;
}
}
} else if ((flags & ConnFifoEnabled) && connFifo[step%NCCL_STEPS].mode == NCCL_MODE_OFFSET) {
ptrs[index] = connEltsFifo + loadInt(&connFifo[step%NCCL_STEPS].offset)/sizeof(T);
} else if (isSendNotRecv && DirectSend) {
if (flags & (DirectWrite | NvlsDirectWrite | IpcWrite)) {
if (flags & DirectWrite) {
ptrs[index] = directBuff + dstIx + offset;
} else if ((flags & DirectRead) || (flags & IpcRead)) { // empty send
} else if (flags & DirectRead) { // empty send
ptrs[index] = nullptr;
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*connStepSize;
}
} else if (!isSendNotRecv && DirectRecv) {
if (flags & (DirectRead | NvlsDirectRead | IpcRead)) {
if (flags & DirectRead) {
ptrs[index] = directBuff + srcIx + offset;
} else if ((flags & DirectWrite) || (flags & IpcWrite)) {
} else if (flags & DirectWrite) {
ptrs[index] = directBuff + dstIx + offset; // send to next from my output buffer
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*connStepSize;
@@ -214,7 +218,7 @@ private:
int slice = 0;
int offset = 0;
if (tid < nworkers && offset < nelem && ((flags & NetRegMode) == 0)) {
if (tid < nworkers && offset < nelem && !isNetOffload) {
// Worker-only loop for non-empty slices. Non-workers and empty slices are
// processed in the loop following this if block. The benefit of splitting
// the loop like this is we pull two branches out of the critical path.
@@ -263,7 +267,7 @@ private:
* so we need to check whether MultimemSrcs and MultimemDsts are 0. */
&& MultimemSrcs == 0 && MultimemDsts == 0 && !Src) {
// We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
if (Send) {
if (Send && Dst && ncclShmem.groups[group].srcs[0] != ncclShmem.groups[group].dsts[1]) {
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY)
if (tid == 0) {
@@ -350,13 +354,24 @@ private:
constexpr int PreOpSrcs = SrcBuf != Input ? 0 :
DirectRecv*MaxRecv == NCCL_MAX_DIRECT_ARITY ? (1+NCCL_MAX_DIRECT_ARITY) : 1;
reduceCopy<Unroll, RedOp, T,
MultimemSrcs, Recv+Src, Recv*MaxRecv+Src,
MultimemDsts, Send+Dst, Send*MaxSend+Dst, PreOpSrcs>
(tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, postOp,
Recv*fan.nrecv()+Src, ncclShmem.groups[group].srcs,
Send*fan.nsend()+Dst, ncclShmem.groups[group].dsts,
workSize);
if (Send && Dst && ncclShmem.groups[group].dsts[1] == nullptr) {
// this case should only be directCopySend() with registered buffers and send to net peer
reduceCopy<Unroll, RedOp, T,
0, Recv + Src, Recv * MaxRecv + Src,
0, 1, 1, PreOpSrcs>
(tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, postOp,
Recv * fan.nrecv() + Src, ncclShmem.groups[group].srcs,
1, ncclShmem.groups[group].dsts,
workSize);
} else {
reduceCopy<Unroll, RedOp, T,
MultimemSrcs, Recv + Src, Recv * MaxRecv + Src,
MultimemDsts, Send + Dst, Send * MaxSend + Dst, PreOpSrcs>
(tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, postOp,
Recv * fan.nrecv() + Src, ncclShmem.groups[group].srcs,
Send * fan.nsend() + Dst, ncclShmem.groups[group].dsts,
workSize);
}
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
if (tid == 0) {
@@ -372,9 +387,14 @@ private:
}
#endif
} else {
// we will come here when calling prims.directSend with net peer,
// in this case, ncclShmem.groups[group].dsts[0] == NULL, so we
// skip data flush.
workSize = 0;
}
barrier(); // This barrier has a counterpart in following loop
postPeer<Recv, Send>(0 < sliceSize);
postPeer<Recv, Send>(0 < workSize);
offset += sliceSize;
slice += 1;
// Yes, for some template arguments this code will be unreachable. That's fine.
@@ -391,10 +411,11 @@ private:
sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
{ // Only workers could have Wait roles so we know the slice must be empty
// since we've exited the loop above.
waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(0, 0, 0, 0);
waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(0, 0, 0, sliceSize);
}
barrier(); // Has couterpart in preceding worker-only loop.
postPeer<Recv, Send>(0 < sliceSize);
int workSize = ncclShmem.aborted ? 0 : sliceSize;
postPeer<Recv, Send>(0 < workSize);
offset += sliceSize;
slice += 1;
}
@@ -487,17 +508,17 @@ public:
ptrs[index] = connEltsFifo + offset/sizeof(T);
} else if (Direct && fn.work->regUsed) {
if (isSendNotRecv) {
if (flags & (DirectWrite | IpcWrite)) {
if (flags & DirectWrite) {
ptrs[index] = directBuff;
} else if (flags & (DirectRead | IpcRead)) { // empty send
} else if (flags & DirectRead) { // empty send
ptrs[index] = nullptr;
} else {
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
}
} else {
if (flags & (DirectRead | IpcRead)) {
if (flags & DirectRead) {
ptrs[index] = directBuff;
} else if (flags & (DirectWrite | IpcWrite)) {
} else if (flags & DirectWrite) {
if (Send)
ptrs[index] = directBuff; // send to next from my output buffer
else
@@ -580,7 +601,7 @@ private:
int i = (j+shift)%fan.nsend();
ssize_t pOffset = i*peerOffset;
// Skip the data I am responsible of reducing myself
if (skip >= 0 && i >= skip) pOffset += peerElem;
if (skip >= 0 && i >= skip) pOffset += peerOffset;
void* src0 = (T*)ncclShmem.groups[group].srcs[0] + pOffset;
ssize_t realPeerSize = min(realSize, totalElem-pOffset);
if (realPeerSize > 0 && ncclShmem.groups[group].dsts[i] != nullptr) {
@@ -592,7 +613,7 @@ private:
} else if (Recv) {
if (tid==0) ncclShmem.groups[group].dsts[0] = (T*)ncclShmem.groups[group].userOutput + outIx + offset;
ssize_t pOffset = index*peerOffset;
if (skip >= 0 && index >= skip) pOffset += peerElem;
if (skip >= 0 && index >= skip) pOffset += peerOffset;
// Adjust remote index with peer offset in case we are directly pulling from peer's output buffer
waitPeer<DirectRecv, 0, 1, 0, 0, 1>(outIx+pOffset, outIx+pOffset, offset, realSize);
subBarrier();
@@ -600,7 +621,7 @@ private:
for (int j=0; j<fan.nrecv(); j++) {
int i = (j+shift)%fan.nrecv();
pOffset = i*peerOffset;
if (skip >= 0 && i >= skip) pOffset += peerElem;
if (skip >= 0 && i >= skip) pOffset += peerOffset;
void* dst0 = (T*)ncclShmem.groups[group].dsts[0] + pOffset;
ssize_t realPeerSize = min(realSize, totalElem-pOffset);
if (DirectRecv && ncclShmem.groups[group].srcs[i] == dst0) realPeerSize = 0;
@@ -614,7 +635,7 @@ private:
}
}
__device__ __forceinline__ void loadRecvConn(ncclDevChannelPeer *peer, int connIndex, uint32_t direct, int regFlag) {
__device__ __forceinline__ void loadRecvConn(ncclDevChannelPeer *peer, int connIndex, uint32_t direct, int ipcRegFlag, int netRegFlag) {
conn = &peer->recv[connIndex];
if (conn->netDeviceHandle.netDeviceType == NCCL_NET_DEVICE_UNPACK) {
// handle must be a device ptr
@@ -639,33 +660,34 @@ private:
if (conn->connFifo != nullptr) {
flags |= ConnFifoEnabled;
connFifo = conn->connFifo;
} else if (Direct && regFlag) {
// User buffers have been registered
if (conn->flags & (NCCL_IPC_READ | NCCL_IPC_WRITE)) {
if (P2p) {
flags |= conn->flags & NCCL_IPC_WRITE ? IpcWrite : IpcRead;
} else if (connIndex == 1 && direct) {
flags |= IpcRead;
} else {
flags |= direct & NCCL_DIRECT_READ ? IpcRead : IpcWrite;
}
if (Direct) {
if (ipcRegFlag) {
// User buffers have been registered
if (conn->flags & (NCCL_P2P_READ | NCCL_P2P_WRITE)) {
if (P2p) {
flags |= conn->flags & NCCL_P2P_WRITE ? DirectWrite : DirectRead;
} else if (connIndex == 1 && direct) {
flags |= DirectRead;
} else {
flags |= direct & NCCL_P2P_READ ? DirectRead : DirectWrite;
}
} else if ((conn->flags & NCCL_NVLS_MIN_POLL)) {
/* NVLS direct */
flags |= DirectRead;
}
} else if (conn->flags & (NCCL_DIRECT_WRITE | NCCL_DIRECT_READ)) {
if (P2p) {
flags |= conn->flags & NCCL_DIRECT_WRITE ? DirectWrite : DirectRead;
} else if (connIndex == 1 && direct) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
flags |= direct & NCCL_DIRECT_READ ? DirectRead : DirectWrite;
}
if (netRegFlag) {
if (conn->flags & NCCL_DIRECT_NIC) {
flags |= NetRegMode;
connFifo[step % NCCL_STEPS].size = 0;
}
} else if ((conn->flags & NCCL_NVLS_MIN_POLL)) {
/* NVLS direct */
flags |= NvlsDirectRead;
}
}
}
}
__device__ __forceinline__ void loadSendConn(ncclDevChannelPeer *peer, int connIndex, uint32_t direct, int regFlag) {
__device__ __forceinline__ void loadSendConn(ncclDevChannelPeer *peer, int connIndex, uint32_t direct, int ipcRegFlag, int netRegFlag) {
conn = &peer->send[connIndex];
step = conn->step;
step = roundUp(step, SlicePerChunk*StepPerSlice);
@@ -685,27 +707,26 @@ private:
connStepCache = loadStepValue(connStepPtr);
connStepSize = conn->stepSize/sizeof(T);
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (connFifo == nullptr && Direct && regFlag) {
// User buffers have been registered
if (conn->flags & (NCCL_IPC_READ | NCCL_IPC_WRITE)) {
if (P2p) {
flags |= conn->flags & NCCL_IPC_WRITE ? IpcWrite : IpcRead;
} else if (connIndex == 1 && direct) {
flags |= IpcRead;
} else {
flags |= direct & NCCL_DIRECT_READ ? IpcRead : IpcWrite;
if (Direct) {
if (ipcRegFlag) {
// User buffers have been registered
if (conn->flags & (NCCL_P2P_WRITE | NCCL_P2P_READ)) {
if (P2p) {
flags |= conn->flags & NCCL_P2P_WRITE ? DirectWrite : DirectRead;
} else if (connIndex == 1 && direct) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
flags |= direct & NCCL_P2P_READ ? DirectRead : DirectWrite;
}
} else if ((conn->flags & NCCL_NVLS_MIN_POLL)) {
/* NVLS direct */
flags |= DirectWrite;
}
} else if (conn->flags & (NCCL_DIRECT_WRITE | NCCL_DIRECT_READ)) {
if (P2p) {
flags |= conn->flags & NCCL_DIRECT_WRITE ? DirectWrite : DirectRead;
} else if (connIndex == 1 && direct) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
flags |= direct & NCCL_DIRECT_READ ? DirectRead : DirectWrite;
}
if (netRegFlag) {
if (conn->flags & NCCL_DIRECT_NIC) {
flags |= NetRegMode;
}
} else if ((conn->flags & NCCL_NVLS_MIN_POLL)) {
/* NVLS direct */
flags |= NvlsDirectWrite;
}
}
}
@@ -715,8 +736,8 @@ public:
__forceinline__ __device__ Primitives(
int tid, int nthreads, int const *recvPeers, int const *sendPeers,
void const *inputBuf, void *outputBuf, uint64_t redOpArg, uint8_t group=0,
uint8_t connIndexRecv = 0, uint8_t connIndexSend = 0, struct ncclDevWorkColl* e = nullptr,
bool ipcReg = false, bool netReg = false, int stepSize_ = 0, int mode = primsModeDefault
uint8_t connIndexRecv = 0, uint8_t connIndexSend = 0, struct ncclDevWorkColl* collWork = nullptr,
struct ncclDevWorkP2p* p2pWork = nullptr, int stepSize_ = 0, int mode = primsModeDefault
):
tid(tid), nthreads(nthreads), tidInBlock(threadIdx.x), group(group),
stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_) {
@@ -785,11 +806,23 @@ public:
// Coverity thinks that index could be -1 here but that's not actually the case.
// coverity[negative_returns:FALSE]
if (flags & (RoleWaitRecv|RolePostRecv)) loadRecvConn(ncclShmem.channel.peers[peer], connIndexRecv, e ? e->direct : 0, e ? e->regUsed : ipcReg);
// coverity[negative_returns:FALSE]
if (flags & (RoleWaitSend|RolePostSend)) loadSendConn(ncclShmem.channel.peers[peer], connIndexSend, e ? e->direct : 0, e ? e->regUsed : ipcReg);
if (netReg) flags |= NetRegMode;
int sendIpcReg;
int recvIpcReg;
int sendNetReg;
int recvNetReg;
if (P2p) {
sendIpcReg = p2pWork ? p2pWork->sendIpcReg : 0;
recvIpcReg = p2pWork ? p2pWork->recvIpcReg : 0;
sendNetReg = p2pWork ? p2pWork->sendNetReg : 0;
recvNetReg = p2pWork ? p2pWork->recvNetReg : 0;
} else {
recvIpcReg = sendIpcReg = collWork ? collWork->regUsed : 0;
recvNetReg = sendNetReg = collWork ? collWork->netRegUsed : 0;
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
if (flags & (RoleWaitRecv|RolePostRecv)) loadRecvConn(ncclShmem.channel.peers[peer], connIndexRecv, collWork ? collWork->direct : 0, recvIpcReg, recvNetReg);
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
if (flags & (RoleWaitSend|RolePostSend)) loadSendConn(ncclShmem.channel.peers[peer], connIndexSend, collWork ? collWork->direct : 0, sendIpcReg, sendNetReg);
// if (barrierAny(flags & NetDeviceUnpack)) {
// flags |= AnyNetDeviceUnpack;
@@ -801,8 +834,10 @@ public:
// }
// }
// coverity[negative_returns:FALSE]
setDataPtrs(inputBuf, outputBuf, redOpArg, (struct ncclDevWorkCollReg*)e, (uint8_t)(e ? e->regUsed : ipcReg), peer);
// coverity[negative_returns:FALSE] => coverity thinks that index could be -1 but that's not actually the case
// coverity[var_deref_model] => coverity thinks work can dereferenced if NULL but this is not the case
setDataPtrs(inputBuf, outputBuf, redOpArg, (struct ncclDevWorkCollReg*)collWork, sendIpcReg || recvIpcReg, peer);
// coverity[uninit_member] => coverity thinks fan.n is not initialized
}
__forceinline__ __device__ ~Primitives() {
@@ -825,6 +860,16 @@ public:
// Make sure all threads are done writing back conn->step and done using
// ncclShmem.groups[group]
barrier();
if ((flags & DirectRead) && (flags & RoleWaitSend) && P2p) {
// For sendrecv DirectRead, sender needs to wait for receiver reading data from src.
// This has to be done after barrier() since post thread might have contention with
// this check.
int spins = 0;
volatile uint64_t* tail = conn->tail;
volatile uint64_t* head = conn->head;
while (*tail > *head) if (checkAbort(spins)) break;
}
}
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf, uint64_t redOpArg, struct ncclDevWorkCollReg* work, uint8_t ipcReg, int peer) {
@@ -835,10 +880,10 @@ public:
}
if (Direct && ipcReg) {
bool recvProvider = (flags & RoleWaitRecv) && (flags & DirectWrite || flags & IpcWrite);
bool sendAcceptor = (flags & RoleWaitSend) && (flags & DirectWrite || flags & IpcWrite || flags & NvlsDirectWrite);
bool sendProvider = (flags & RoleWaitSend) && (flags & DirectRead || flags & IpcRead); // sender provides direct buffer (to be fetched)
bool recvAcceptor = (flags & RoleWaitRecv) && (flags & DirectRead || flags & IpcRead || flags & NvlsDirectRead); // receiver accepts direct buffer
bool recvProvider = (flags & RoleWaitRecv) && (flags & DirectWrite);
bool sendAcceptor = (flags & RoleWaitSend) && (flags & DirectWrite);
bool sendProvider = (flags & RoleWaitSend) && (flags & DirectRead); // sender provides direct buffer (to be fetched)
bool recvAcceptor = (flags & RoleWaitRecv) && (flags & DirectRead); // receiver accepts direct buffer
if (recvProvider) {
int spins = 0;
void* volatile* slot = ncclShmem.groups[group].recvConns[index]->ptrExchange;
@@ -851,6 +896,7 @@ public:
exchgPtr = (T*)outputBuf;
} else {
int localPeer = ncclShmem.comm.rankToLocalRank[peer];
// coverity[deref_parm:FALSE] => work cannot be NULL if ipcReg != NULL
exchgPtr = (T*)(work->coll.recvbuffOffset + work->coll.recvbuffRmtAddrs[localPeer]);
}
*slot = reinterpret_cast<void*>(exchgPtr);
@@ -869,6 +915,7 @@ public:
directBuff = reinterpret_cast<T*>(ptr);
*slot = nullptr;
} else {
// coverity[var_deref_op]
directBuff = (T*)work->dnOutputs[index];
}
}
@@ -889,8 +936,10 @@ public:
} else {
int localPeer = ncclShmem.comm.rankToLocalRank[peer];
if (MaxRecv == 0)
// coverity[var_deref_op]
exchgPtr = (T*)(work->coll.sendbuffOffset + work->coll.sendbuffRmtAddrs[localPeer]);
else
// coverity[var_deref_op]
exchgPtr = (T*)(work->coll.recvbuffOffset + work->coll.recvbuffRmtAddrs[localPeer]);
}
@@ -987,11 +1036,11 @@ public:
__device__ __forceinline__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
}
__device__ __forceinline__ void directRecvCopyDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<1, 1, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
__device__ __forceinline__ void directRecvCopyDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<1, 1, 1, 1, -1, Output>(inpIx, outIx, eltN, postOp);
}
__device__ __forceinline__ void directRecvDirectSend(intptr_t inpIx, intptr_t outIx, int eltN) {
genericOp<1, 1, 1, 1, -1, -1>(inpIx, outIx, eltN, false);
__device__ __forceinline__ void directRecvDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<1, 1, 1, 1, -1, -1>(inpIx, outIx, eltN, postOp);
}
__device__ __forceinline__ void recvCopyDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 1, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
@@ -1010,6 +1059,9 @@ public:
__device__ __forceinline__ void directRecvReduceSend(intptr_t inpIx, int eltN, bool postOp=false) {
genericOp<1, 0, 1, 1, Input, -1>(inpIx, -1, eltN, postOp);
}
__device__ __forceinline__ void recvReduceDirectSend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 1, 1, 1, Input, -1>(inpIx, outIx, eltN, postOp);
}
__device__ __forceinline__ void directRecvReduceDirectSend(intptr_t inpIx, intptr_t outIx, ssize_t eltN, bool postOp=false) {
genericOp<1, 1, 1, 1, Input, -1>(inpIx, outIx, eltN, postOp);
}
projects/rccl/src/device/reduce_kernel.h
+151 -25
View File
@@ -289,13 +289,30 @@ SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, fn.isMinNotMax ? __hmin(x, y) : __h
#endif
#if defined(RCCL_FLOAT8)
#if __CUDA_ARCH__ >= 900
SPECIALIZE_REDUCE(FuncSum, __nv_fp8_e4m3, 1, __nv_fp8_e4m3, __nv_fp8_e4m3(__hadd(__half(x),__half(y))))
SPECIALIZE_REDUCE(FuncSum, __nv_fp8_e4m3, 2, __nv_fp8x2_e4m3, __nv_fp8x2_e4m3(__hadd2(__half2(x),__half2(y))))
SPECIALIZE_REDUCE(FuncProd, __nv_fp8_e4m3, 1, __nv_fp8_e4m3, __nv_fp8_e4m3(__hmul(__half(x),__half(y))))
SPECIALIZE_REDUCE(FuncProd, __nv_fp8_e4m3, 2, __nv_fp8x2_e4m3, __nv_fp8x2_e4m3(__hmul2(__half2(x),__half2(y))))
SPECIALIZE_REDUCE(FuncMinMax, __nv_fp8_e4m3, 1, __nv_fp8_e4m3, __nv_fp8_e4m3(fn.isMinNotMax ? __hmin(__half(x),__half(y)) : __hmax(__half(x),__half(y))))
SPECIALIZE_REDUCE(FuncMinMax, __nv_fp8_e4m3, 2, __nv_fp8x2_e4m3, __nv_fp8x2_e4m3(fn.isMinNotMax ? __hmin2(__half2(x),__half2(y)) : __hmax2(__half2(x),__half2(y))))
SPECIALIZE_REDUCE(FuncSum, __nv_fp8_e5m2, 1, __nv_fp8_e5m2, __nv_fp8_e5m2(__hadd(__half(x),__half(y))))
SPECIALIZE_REDUCE(FuncSum, __nv_fp8_e5m2, 2, __nv_fp8x2_e5m2, __nv_fp8x2_e5m2(__hadd2(__half2(x),__half2(y))))
SPECIALIZE_REDUCE(FuncProd, __nv_fp8_e5m2, 1, __nv_fp8_e5m2, __nv_fp8_e5m2(__hmul(__half(x),__half(y))))
SPECIALIZE_REDUCE(FuncProd, __nv_fp8_e5m2, 2, __nv_fp8x2_e5m2, __nv_fp8x2_e5m2(__hmul2(__half2(x),__half2(y))))
SPECIALIZE_REDUCE(FuncMinMax, __nv_fp8_e5m2, 1, __nv_fp8_e5m2, __nv_fp8_e5m2(fn.isMinNotMax ? __hmin(__half(x), __half(y)) : __hmax(__half(x), __half(y))))
SPECIALIZE_REDUCE(FuncMinMax, __nv_fp8_e5m2, 2, __nv_fp8x2_e5m2, __nv_fp8x2_e5m2(fn.isMinNotMax ? __hmin2(__half2(x), __half2(y)) : __hmax2(__half2(x), __half2(y))))
#else
SPECIALIZE_REDUCE(FuncSum, rccl_float8, 1, rccl_float8, rccl_float8(float(x) + float(y)))
SPECIALIZE_REDUCE(FuncProd, rccl_float8, 1, rccl_float8, rccl_float8(float(x) * float(y)))
SPECIALIZE_REDUCE(FuncMinMax, rccl_float8, 1, rccl_float8, rccl_float8(fn.isMinNotMax ? fminf(float(x), float(y)) : fmaxf(float(x), float(y))))
SPECIALIZE_REDUCE(FuncSum, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(float(x) + float(y)))
SPECIALIZE_REDUCE(FuncProd, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(float(x) * float(y)))
SPECIALIZE_REDUCE(FuncMinMax, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(fn.isMinNotMax ? fminf(float(x), float(y)) : fmaxf(float(x), float(y))))
#endif
#endif
#undef SPECIALIZE_REDUCE
@@ -416,7 +433,7 @@ struct FuncPreMulSum<half> {
using EltType = half;
half2 scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; half val; };
union { uint64_t u64; __half val; };
u64 = opArg;
scalar.x = val;
scalar.y = val;
@@ -450,6 +467,31 @@ struct FuncPreMulSum<half> {
#endif
#if defined(RCCL_FLOAT8)
#if __CUDA_ARCH__ >= 900
template<>
struct FuncPreMulSum<__nv_fp8_e4m3> {
using EltType = __nv_fp8_e4m3;
__half2 scalar2;
__device__ FuncPreMulSum(uint64_t opArg) {
union { uint64_t u64; __nv_fp8_storage_t val; };
u64 = opArg;
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E4M3));
scalar2.y = scalar2.x;
}
};
template<>
struct FuncPreMulSum<__nv_fp8_e5m2> {
using EltType = __nv_fp8_e5m2;
__half2 scalar2;
__device__ FuncPreMulSum(uint64_t opArg) {
union { uint64_t u64; __nv_fp8_storage_t val; };
u64 = opArg;
scalar2.x = __half(__nv_cvt_fp8_to_halfraw(val, __NV_E5M2));
scalar2.y = scalar2.x;
}
};
#else
template<>
struct FuncPreMulSum<rccl_float8> {
// Change these to switch between all prescale, all postscale, or both by sqrt(N).
@@ -480,12 +522,13 @@ struct FuncPreMulSum<half> {
}
};
#endif
#endif
template<typename T>
struct Apply_Reduce<FuncPreMulSum<T>, /*EltPerPack=*/1> {
__device__ static BytePack<sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
template<typename T, int EltPerPack>
struct Apply_Reduce<FuncPreMulSum<T>, EltPerPack> {
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
// FuncPreMulSum reduce dispatches to FuncSum.
return Apply_Reduce<FuncSum<T>, 1>::reduce(FuncSum<T>(), a, b);
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
}
};
@@ -548,7 +591,49 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
#endif
#endif
////////////////////////////////////////////////////////////////////////////////
// Apply_PreOp of FuncPreMulSum for fp8.
#if defined(RCCL_FLOAT8)
#if __CUDA_ARCH__ >= 900
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8_e4m3)> preOp(
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8_e4m3)> a
) {
return toPack<__nv_fp8_e4m3>(__nv_fp8_e4m3(__hmul(__half(fromPack<__nv_fp8_e4m3>(a)), fn.scalar2.x)));
}
};
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e4m3>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8x2_e4m3)> preOp(
FuncPreMulSum<__nv_fp8_e4m3> fn, BytePack<sizeof(__nv_fp8x2_e4m3)> a
) {
return toPack<__nv_fp8x2_e4m3>(__nv_fp8x2_e4m3(__hmul2(__half2(fromPack<__nv_fp8x2_e4m3>(a)), fn.scalar2)));
}
};
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8_e5m2)> preOp(
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8_e5m2)> a
) {
return toPack<__nv_fp8_e5m2>(__nv_fp8_e5m2(__hmul(__half(fromPack<__nv_fp8_e5m2>(a)), fn.scalar2.x)));
}
};
template<>
struct Apply_PreOp<FuncPreMulSum<__nv_fp8_e5m2>, /*EltPerPack=*/2> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(__nv_fp8x2_e5m2)> preOp(
FuncPreMulSum<__nv_fp8_e5m2> fn, BytePack<sizeof(__nv_fp8x2_e5m2)> a
) {
return toPack<__nv_fp8x2_e5m2>(__nv_fp8x2_e5m2(__hmul2(__half2(fromPack<__nv_fp8x2_e5m2>(a)), fn.scalar2)));
}
};
#else
template<>
struct Apply_PreOp<FuncPreMulSum<rccl_float8>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
@@ -571,6 +656,7 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
}
};
#endif
#endif
////////////////////////////////////////////////////////////////////////////////
// FuncSumPostDiv
@@ -583,34 +669,74 @@ struct RedOpArg<FuncSumPostDiv<T>> {
}
};
template<typename T, bool IsFloating=IsFloatingPoint<T>::value>
struct FuncSumPostDiv_IntOnly;
template<typename T>
struct FuncSumPostDiv: FuncSumPostDiv_IntOnly<T> {
__device__ FuncSumPostDiv(uint64_t opArg=0):
FuncSumPostDiv_IntOnly<T>(opArg) {
struct Divider {
__device__ __forceinline__ static T divide(T dividend, T divisor) {
return dividend / divisor;
}
};
template<>
struct Divider<uint64_t> {
__device__ __forceinline__ static uint64_t divide(uint64_t dividend, uint64_t divisor) {
if (divisor == 0) {
return UINT64_MAX;
}
uint64_t quotient = 0;
uint64_t remainder = 0;
#pragma unroll 64
for (int i = 63; i >= 0; --i) {
remainder = (remainder << 1) | ((dividend >> i) & 1);
if (remainder >= divisor) {
remainder -= divisor;
quotient |= (1ULL << i);
}
}
return quotient;
}
};
template<typename T>
struct FuncSumPostDiv_IntOnly<T, /*IsFloating=*/false>: FuncSum<T> {
struct FuncSumPostDiv {
static_assert(T(0) < T(-1), "FuncSumPostDiv is only for implementing ncclAvg on uint types.");
using EltType = T;
int divisor;
__device__ FuncSumPostDiv_IntOnly(uint64_t opArg=0): divisor(opArg) {}
using UintType = typename std::conditional<sizeof(T)==8, uint64_t, uint32_t>::type;
uint32_t divisor:31, isSigned:1;
UintType recip;
__device__ FuncSumPostDiv(uint64_t opArg=0) {
isSigned = opArg & 1;
divisor = opArg >> 1;
recip = Divider<UintType>::divide(UintType(-1), divisor);
}
__device__ T divide(T x) {
// x is negative iff we are in signed mode and the top bit is set
bool xneg = isSigned && (x & ~(T(-1)>>1));
// Compute abs(x):
// T(-x) vs -T(x) is critical. We have to negate then truncate the bits. Consider
// if we are doing signed 8-bit types, thus T=uint8_t. The value -1 is encoded
// as 0xff. -T(0xff) when promoted to 32-bit (which is implicit by compiler)
// gives 0xffffff01, but T(-0xff) is 0x1, and that is the abs value we want.
UintType xabs = xneg ? T(-x) : x;
// Compute quotient by multiplying by reciprical.
UintType q = sizeof(T)==8 ? __umul64hi(xabs, recip) : __umulhi(xabs, recip);
// Quotient may be off by one so do a fixup.
if (xabs - q*divisor >= divisor) q += 1;
// If original x was negative then we have to negate it back since we were
// working with its abs val.
return xneg ? -T(q) : T(q);
}
};
template<typename T>
struct FuncSumPostDiv_IntOnly<T, /*IsFloating=*/true> {
static_assert(sizeof(T)!=sizeof(T), "FuncSumPostDiv is only for implementing ncclAvg on integral types.");
};
template<typename T>
struct Apply_Reduce<FuncSumPostDiv<T>, /*EltPerPack=*/1>:
Apply_Reduce<FuncSum<T>, 1> {
__device__ static BytePack<sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
template<typename T, int EltPerPack>
struct Apply_Reduce<FuncSumPostDiv<T>, EltPerPack>:
Apply_Reduce<FuncSum<T>, EltPerPack> {
__device__ static BytePack<EltPerPack*sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b) {
// FuncSumPostDiv reduce dispatches to FuncSum.
return Apply_Reduce<FuncSum<T>, 1>::reduce(FuncSum<T>(), a, b);
return Apply_Reduce<FuncSum<T>, EltPerPack>::reduce(FuncSum<T>(), a, b);
}
};
@@ -618,7 +744,7 @@ template<typename T>
struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
static constexpr bool IsIdentity = false;
__device__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
return toPack<T>(fromPack<T>(a) / fn.divisor);
return toPack<T>(fn.divide(fromPack<T>(a)));
}
};
projects/rccl/src/device/reduce_scatter.h
+28 -27
View File
@@ -165,7 +165,7 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SI
T *inputBuf = (T*)work->sendbuff;
T *outputBuf = (T*)work->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, NULL, NULL, inputBuf, outputBuf, work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, nullptr, false, false, 0, primsModePatRs);
(tid, nthreads, NULL, NULL, inputBuf, outputBuf, work->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, nullptr, nullptr, 0, primsModePatRs);
PatRSAlgorithm<T> patAlgo(chunkCount*sizeof(T), NCCL_STEPS, channelOffset, channelOffset + channelCount, count, chunkCount, rank, nranks);
int last = 0;
@@ -213,6 +213,7 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_S
nelem = min(chunkCount, channelCount - elemOffset);
prims.scatter(offset, nvls->nHeads * count, nelem, count, -1, 0);
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
} else if (tid < tidEndReduce) {
// Reduce through NVLS
using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
@@ -282,10 +283,10 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
int nRails = direct->nHeads;
int part = ncclShmem.channelId - work->channelLo;
void* inbuf = (void*)work->sendbuff;
ssize_t sizePerRank = work->collnet.count;
ssize_t countPerRank = work->collnet.count;
ssize_t railAllBeg = min(railGridOffset + part*chunkSize, nNodes*sizePerRank);
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes*sizePerRank);
ssize_t railAllBeg = min(railGridOffset + part*chunkSize, nNodes*countPerRank);
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes*countPerRank);
int railAllSize = railAllEnd - railAllBeg;
if (tid < nDsts) dstSizes[tid] = railAllSize;
@@ -298,15 +299,15 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
if (rail == nRails) rail = 0;
}
do {
int node = railAllBeg/sizePerRank;
int node = railAllBeg/countPerRank;
int railAllOffset = 0;
while (railAllOffset < railAllSize) {
ssize_t railOneBeg = node*sizePerRank;
ssize_t railOneEnd = railOneBeg + sizePerRank;
ssize_t railOneBeg = node*countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t railOneOffset = (railAllBeg+railAllOffset) - railOneBeg;
int delta = min(railAllEnd, railOneEnd) - (railAllBeg+railAllOffset);
int rank = ncclShmem.comm.collNetDenseToUserRank[node*nRails + rail];
ssize_t userOneBeg = rank*sizePerRank + railOneOffset;
ssize_t userOneBeg = rank*countPerRank + railOneOffset;
if (nDsts != 0) {
reduceCopy<ncclCollUnroll(), RedOp, T,
/*MultimemSrcs=*/0, 1+MinSrcs, 1+MaxSrcs,
@@ -315,7 +316,7 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
(tid, tn, work->redOpArg, &work->redOpArg, false,
/*nSrcs=*/1+nSrcs, [=]__device__(int s) {
return s==0 ? (T*)inbuf + userOneBeg
: work->regUsed && (recvDirectFlag & NCCL_DIRECT_READ)
: work->regUsed && (recvDirectFlag & NCCL_P2P_READ)
? (T*)srcPtrs[s-1] + userOneBeg
: (T*)srcPtrs[s-1] + railAllOffset;
},
@@ -340,7 +341,8 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
struct ncclDirect* direct = &ncclShmem.channel.collnetDirect;
int const &nNodes = ncclShmem.comm.nNodes;
ssize_t chunkSize = int(work->collnet.chunkCount);
ssize_t sizePerRank = work->collnet.count;
ssize_t countPerRank = work->collnet.count;
const int hasDn = (direct->down[0] >= 0) ? 1 : 0;
if (direct->out == -1) __builtin_trap();
bool isMultiRail = (direct->nHeads > 1);
@@ -357,15 +359,15 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
int tn = nWarps1*WARP_SIZE;
if (tid < tn) {
// Phase 1: Scatter inputs to peers
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/1, Proto, 0>
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, tn, nullptr, direct->heads+1, work->sendbuff, nullptr,
work->redOpArg, 0*Proto::MaxGroupWidth, 1, 1, work);
for (ssize_t railGridOffset=0; railGridOffset < nNodes*sizePerRank; railGridOffset += nChannels*chunkSize) {
work->redOpArg, 0*Proto::MaxGroupWidth, 1, 1);
for (ssize_t railGridOffset=0; railGridOffset < nNodes*countPerRank; railGridOffset += nChannels*chunkSize) {
Scatterer</*ReduceSendNotRecv=*/true> scat;
scat.work = work;
scat.chunkSize = chunkSize;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/0, /*Send=*/1>(scat, NCCL_DIRECT_READ, 0);
prims.template process</*Recv=*/0, /*Send=*/1>(scat, 0, 0);
}
return;
}
@@ -373,23 +375,22 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
tn = nWarps2*WARP_SIZE;
if (tid < tn) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed && !hasDn) {
if (tid == 0) {
int steps = (int)divUp(nNodes * sizePerRank * sizeof(T), NCCL_MAX_COLLNET_SIZE);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>::sendPeerNotify(direct->out, 1, steps);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>::sendPeerNotify(direct->out, 1, 1);
}
__syncwarp();
} else {
// Phase 2: Reduce from peers + local input -> send to network
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/1, Proto, 0>
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>
prims(tid, tn, direct->heads + 1, &direct->out, nullptr, nullptr,
work->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1, work);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * sizePerRank; railGridOffset += nChannels * chunkSize) {
work->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkSize) {
Scatterer</*ReduceSendNotRecv=*/false> scat;
scat.work = work;
scat.chunkSize = chunkSize;
scat.railGridOffset = railGridOffset;
prims.template process</*Recv=*/1, /*Send=*/1>(scat, 0, NCCL_DIRECT_READ);
prims.template process</*Recv=*/1, /*Send=*/1>(scat, 0, 0);
}
}
return;
@@ -398,9 +399,9 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
tn = nWarps3*WARP_SIZE;
if (tid < tn) {
if (work->regUsed == NCCL_COLLNET_REG_BUFFER) {
if (work->netRegUsed) {
if (tid == 0) {
int steps = (int)divUp(nNodes * sizePerRank * sizeof(T), NCCL_MAX_COLLNET_SIZE);
int steps = hasDn ? (int)divUp(nNodes * countPerRank, nChannels * chunkSize) : 1;
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>::recvPeerNotify(direct->out, 0, steps);
}
__syncwarp();
@@ -409,11 +410,11 @@ struct RunWorkColl<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NC
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
prims(tid, tn, &direct->out, nullptr, nullptr, work->recvbuff,
work->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0);
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * sizePerRank; railGridOffset += nChannels * chunkSize) {
for (ssize_t railGridOffset = 0; railGridOffset < nNodes * countPerRank; railGridOffset += nChannels * chunkSize) {
ssize_t railAllBeg = railGridOffset + part * chunkSize;
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes * sizePerRank);
ssize_t railOneBeg = ncclShmem.comm.node * sizePerRank;
ssize_t railOneEnd = railOneBeg + sizePerRank;
ssize_t railAllEnd = min(railAllBeg + chunkSize, nNodes * countPerRank);
ssize_t railOneBeg = ncclShmem.comm.node * countPerRank;
ssize_t railOneEnd = railOneBeg + countPerRank;
ssize_t beg = max(railAllBeg, railOneBeg);
ssize_t end = min(railAllEnd, railOneEnd);
prims.recv(beg - railOneBeg, max(ssize_t(0), end - beg), /*postOp=*/true);
projects/rccl/src/device/sendrecv.h
+10 -8
View File
@@ -19,7 +19,9 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
template<typename Proto>
__device__ void runSend(int tid, int tn, int group, struct ncclDevWorkP2p* work) {
size_t bytes = work->sendBytes;
int chunkSize = work->sendIpcReg && ncclShmem.comm.isNvlink ? (1 << 30) : u32fp8Decode(work->sendChunkSize_u32fp8);
bool useLargeChunk = (work->sendIpcReg && ncclShmem.comm.isAllNvlink) || work->sendNetReg;
int chunkSize = useLargeChunk ? NCCL_MAX_NET_SIZE : u32fp8Decode(work->sendChunkSize_u32fp8);
int stepSize = useLargeChunk ? NCCL_MAX_NET_SIZE : ncclShmem.comm.p2pChunkSize;
#if defined(ENABLE_NPKIT)
bool isNpKitThread = (tid == 0);
@@ -42,8 +44,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
Primitives<T, RedOp, FanAsymmetric<0, 1>, 0, Proto, 1>
prims(tid, tn, nullptr, &work->sendRank, work->sendAddr, nullptr,
/*redOpArg(ignored)=*/0, group, work->sendConnIndex, work->sendConnIndex, nullptr,
/*ipcReg=*/work->sendIpcReg, /*netReg=*/work->sendRegistered, ncclShmem.comm.p2pChunkSize);
/*redOpArg(ignored)=*/0, group, work->sendConnIndex, work->sendConnIndex, nullptr, work, stepSize);
#if defined(ENABLE_NPKIT)
if (isNpKitThread) {
@@ -64,7 +65,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
int n = min(size_t(chunkSize), bytes-cursor);
prims.directSend(cursor, cursor, n);
cursor += n;
} while (cursor < bytes && work->sendRegistered == 0);
} while (cursor < bytes);
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_SEND_EXIT)
if (isNpKitThread) {
@@ -77,7 +78,9 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
template<typename Proto>
__device__ void runRecv(int tid, int tn, int group, struct ncclDevWorkP2p* work) {
size_t bytes = work->recvBytes;
int chunkSize = work->recvIpcReg && ncclShmem.comm.isNvlink ? (1 << 30) : u32fp8Decode(work->recvChunkSize_u32fp8);
bool useLargeChunk = (work->recvIpcReg && ncclShmem.comm.isAllNvlink) || work->recvNetReg;
int chunkSize = useLargeChunk ? NCCL_MAX_NET_SIZE : u32fp8Decode(work->recvChunkSize_u32fp8);
int stepSize = useLargeChunk ? NCCL_MAX_NET_SIZE : ncclShmem.comm.p2pChunkSize;
#if defined(ENABLE_NPKIT)
bool isNpKitThread = (tid == 0);
@@ -100,8 +103,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
Primitives<T, RedOp, FanAsymmetric<1, 0>, 0, Proto, 1>
prims(tid, tn, &work->recvRank, nullptr, nullptr, work->recvAddr,
/*redOpArg(ignored)=*/0, group, work->recvConnIndex, work->recvConnIndex, nullptr,
/*ipcReg=*/work->recvIpcReg, /*netReg=*/work->recvRegistered, ncclShmem.comm.p2pChunkSize);
/*redOpArg(ignored)=*/0, group, work->recvConnIndex, work->recvConnIndex, nullptr, work, stepSize);
#if defined(ENABLE_NPKIT)
if (isNpKitThread) {
@@ -122,7 +124,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
int n = min(size_t(chunkSize), bytes-cursor);
prims.directRecv(cursor, cursor, n);
cursor += n;
} while (cursor < bytes && work->recvRegistered == 0);
} while (cursor < bytes);
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_RECV_EXIT)
if (isNpKitThread) {