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

2.24.3-1

Browse Source 
Network user buffer support for collectives
 * Leverage user buffer registration to achieve zero-copy
   inter-node communications for Ring, NVLS and Collnet

Add RAS subsystem
 * Create a RAS thread keeping track of all NCCL communicators.
 * Add a ncclras tool contacting the RAS thread and getting a
   report.

Add fp8 support
 * Add support for e5m2 and e4m3 8-bit floating point operations.
 * Use Tree/PAT algorithms when possible for better numerical
   stability.

Add NIC fusion
 * Add a NET API to ask the network plugin to fuse a set of
   interfaces together.
 * Fuse multiple NICs under the same PCI switch as a single,
   larger NIC.

Socket connection failure retry
 * Retry in case of socket connection failure (unreachable host)
 * Avoid "Software caused connection abort" errors on retries

QP connection failure retry
 * Retry in case of IB QP connection failure during ibv_modify_qp.

NET API improvements
 * Allow plugins to force a flush in case data and completion
   ordering is not guaranteed.
 * Indicate when completion is not needed (e.g. for the LL128
   protocol), allowing plugins to skip generating a completion.
 * Allow for full offload of allgather operations when using one
   GPU per node.

NCCL_ALGO/NCCL_PROTO strict enforcement
 * Extend NCCL_ALGO/NCCL_PROTO syntax to be able to specify
   ALGO/PROTO filters for each collective operation.
 * Strictly enforce the ALGO/PROTO filters, no longer fall back
   on the ring algorithm when the filtering leaves no option and
   error out instead.

Enable CUMEM host allocations
 * Use cumem functions for host memory allocation by default.

Improved profiler plugin API
 * Avoid dependencies with NCCL includes.
 * Add information on whether the buffer is registered or not

Adjust PAT tuning
 * Improve transition between PAT and ring at scale.

Fix hangs when running with different CPU architectures
 * Detect when we use a mix of GPU architectures
 * Ensure Algo/Proto decisions are made based on that unified
   state.

Fix FD leak in UDS
 * Fix a leak when mapping buffers intra-node with cumem IPCs.

Fix crash when mixing buffer registration and graph buffer registration.
 * Separate local and graph registration to avoid crashes when we free
   buffers.

Fix user buffer registration with dmabuf
 * Make ncclSend/ncclRecv communication with buffer registration functional
   on network plugins relying on dmabuf for buffer registration.

Fix crash in IB code caused by uninitialized fields.

Fix non-blocking ncclSend/ncclRecv
 * Fix case where ncclSend/ncclRecv would return ncclSuccess in non-blocking
   mode even though the operation was not enqueued onto the stream.
 * Issue #1495

Various compiler tweaks and fixes
 * PR #758

Fix typo in ncclTopoPrintGraph
 * Issue #1468


[ROCm/rccl commit: 6aae379278]
This commit is contained in:
Sylvain Jeaugey
2024-12-18 08:26:06 -08:00
parent 758f28b359
commit db3bfd118f
97 changed files with 12588 additions and 3127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/device/all_gather.h
+88 -59
View File
@@ -9,64 +9,88 @@
#include "primitives.h"
namespace {
template<typename T, typename RedOp, typename Proto>
template<typename T, typename RedOp, typename Proto, bool isNetOffload = false>
__device__ __forceinline__ void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
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;
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>, 1, Proto, 0> prims
(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, 0, 0, work);
for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
/////////////// begin AllGather steps ///////////////
nelem = min(chunkCount, partCount - elemOffset);
dataOffset = partOffset + elemOffset;
// 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;
}
// step 0: push data to next GPU
rankDest = ringRanks[0];
offset = dataOffset + rankDest * count;
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>, 1, Proto, 0, isNetOffload> prims
(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, 0, 0, work, NULL, isNetOffload ? NCCL_MAX_NET_SIZE : 0);
for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
/////////////// begin AllGather steps ///////////////
nelem = min(chunkCount, partCount - elemOffset);
dataOffset = partOffset + elemOffset;
if (inputBuf + dataOffset == outputBuf + offset) { // In place
prims.directSend(dataOffset, offset, nelem);
} else {
prims.directCopySend(dataOffset, offset, nelem);
}
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ringRanks[nranks-j];
// step 0: push data to next GPU
rankDest = ringRanks[0];
offset = dataOffset + rankDest * count;
prims.directRecvCopyDirectSend(offset, nelem);
if ((inputBuf + dataOffset == outputBuf + offset) || isNetOffload) { // In place or onePPN
prims.directSend(dataOffset, offset, nelem);
} else {
prims.directCopySend(dataOffset, 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);
}
// Make final copy from buffer to dest.
rankDest = ringRanks[1];
offset = dataOffset + rankDest * count;
// Final wait/copy.
prims.directRecv(offset, offset, nelem);
}
// Make final copy from buffer to dest.
rankDest = ringRanks[1];
offset = dataOffset + rankDest * count;
// Final wait/copy.
prims.directRecv(offset, offset, nelem);
} 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);
}
// 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);
}
}
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);
bool isNetOffload = work->isOneRPN && work->netRegUsed;
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);
}
};
@@ -96,7 +120,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;
@@ -137,6 +161,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>;
@@ -148,6 +173,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 */
@@ -204,11 +230,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;
@@ -221,15 +247,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,
@@ -238,11 +264,11 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
/*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;
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_DIRECT_WRITE) ? (char*)dstPtrs[d-outIsDst] + userOneBeg
: work->regUsed && (sendDirectFlag & NCCL_P2P_WRITE) ? (char*)dstPtrs[d-outIsDst] + userOneBeg
: (char*)dstPtrs[d-outIsDst] + railAllOffset;
},
delta);
@@ -262,8 +288,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);
@@ -277,9 +304,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();
@@ -288,11 +318,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));
@@ -304,10 +334,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 {
@@ -315,7 +344,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;
@@ -333,7 +362,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
+85 -82
View File
@@ -69,7 +69,7 @@ namespace {
chunkOffset = chunk * chunkCount;
offset = gridOffset + elemOffset + chunkOffset;
nelem = (int)min(chunkCount, remCount - chunkOffset);
prims.directRecvCopyDirectSend(offset, nelem);
prims.directRecvCopyDirectSend(offset, offset, nelem);
}
// Make final copy from buffer to dest.
@@ -139,7 +139,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);
}
}
}
@@ -222,7 +222,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);
}
}
}
@@ -268,22 +268,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.
@@ -291,24 +299,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 {
@@ -316,8 +320,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);
}
}
@@ -327,10 +331,21 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/1, Proto, 0>
prims(tid, nThreadsGather, direct->up, NULL, work->sendbuff, work->recvbuff,
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) {
@@ -342,15 +357,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 {
@@ -394,8 +409,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));
@@ -407,8 +420,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) {
@@ -419,8 +432,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) {
@@ -430,7 +443,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;
@@ -456,6 +470,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>;
@@ -464,38 +479,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
@@ -504,10 +504,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);
}
@@ -660,10 +661,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 {
@@ -673,8 +673,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>
@@ -683,18 +685,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 {
@@ -720,7 +723,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 {
@@ -740,7 +743,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
+32 -20
View File
@@ -15,37 +15,49 @@ namespace {
const int rank = ring->userRanks[0];
const int nextRank = ring->userRanks[1];
const int root = work->root;
size_t chunkCount;
size_t channelCount;
size_t gridOffset;
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), (size_t*)nullptr, &gridOffset, &channelCount, &chunkCount);
ssize_t chunkCount;
ssize_t channelCount;
ssize_t gridOffset;
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), (ssize_t*)nullptr, &gridOffset, &channelCount, &chunkCount);
size_t offset;
int nelem;
int workNthreads;
bool isNetOffload = work->isOneRPN && work->netRegUsed;
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>, 1, Proto, 0>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, 0, 0, work);
workNthreads = isNetOffload ? WARP_SIZE : nthreads;
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
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>, 1, Proto, 0>
prims(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, 0, 0, work);
if (rank == root) {
if (inputBuf == outputBuf) {
prims.directSend(offset, offset, nelem);
for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
offset = gridOffset + elemOffset;
nelem = min(chunkCount, channelCount - elemOffset);
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 (isNetOffload) barrier_sync(14, nthreads);
}
}
projects/rccl/src/device/common.h
+3
View File
@@ -396,6 +396,9 @@ __device__ void ncclDevFunc_Nop();
ncclKernelMain<specializedFnId, RunWorkBatch<coll, ty, redop<ty>, algo, proto>>(&args4K.args); \
}
#define DEFINE_ncclDevKernel_nop(suffix, coll, redop, ty, algo, proto, specializedFnId) \
__global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {}
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
__device__ void ncclDevFunc_##suffix() { \
RunWorkBatch<coll, ty, redop<ty>, algo, proto>().run(); \
projects/rccl/src/device/common_kernel.h
+14 -5
View File
@@ -65,19 +65,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++) {
@@ -163,7 +167,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]);
@@ -173,11 +178,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
+24 -11
View File
@@ -5,7 +5,7 @@ import sys
# Order of redops, tys, protos, algos must match src/include/device.h
all_colls = ["Broadcast","Reduce","AllGather","ReduceScatter","AllReduce","SendRecv"]
all_redops = ["Sum","Prod","MinMax","PreMulSum","SumPostDiv"]
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16"]
all_tys = ["i8","u8","i32","u32","i64","u64","f16","f32","f64","bf16","f8e4m3","f8e5m2"]
all_protos = ["LL","LL128","SIMPLE"]
all_algos = ["TREE","RING","COLLNET_DIRECT","COLLNET_CHAIN","NVLS","NVLS_TREE","PAT"]
@@ -107,6 +107,9 @@ def required_cuda(coll, redop, ty, algo, proto):
if coll in ("AllReduce","Reduce","ReduceScatter"):
if redop=="SumPostDiv" and ty[0] not in ("i","u"): return None
if ty=="bf16": cudart = max(cudart, 11000)
if ty.startswith("f8"):
cudart = max(cudart, 11080)
arch = max(arch, 900)
if "NVLS" in algo:
if coll in ("AllReduce","Reduce","ReduceScatter"):
@@ -125,7 +128,7 @@ def required_cuda(coll, redop, ty, algo, proto):
def equivalent_primary(coll, redop, ty, algo, proto):
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":
return (coll, redop, "u"+ty[1:], algo, proto)
# map signed integer min/max to unsigned for non-NVLS
if redop=="MinMax" and ty[0]=="i" and ("NVLS" not in algo):
@@ -365,7 +368,9 @@ ty_to_cxx = {
"f16": "half",
"f32": "float",
"f64": "double",
"bf16": "__nv_bfloat16"
"bf16": "__nv_bfloat16",
"f8e4m3": "__nv_fp8_e4m3",
"f8e5m2": "__nv_fp8_e5m2"
}
# Generate each <gensrc>/<impl>.cu:
@@ -385,15 +390,23 @@ for name in name_to_funcs.keys():
sym = paste("_", coll, redop, ty, algo, proto)
fn_id = primary_to_index[kfn]
cudart, arch = required_cuda(*kfn)
s = "DEFINE_ncclDevKernel({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {fn_id})\n"
if (cudart, arch) != (0, 0):
out("#if CUDART_VERSION >= %d && __CUDA_ARCH__ >= %d\n" % (cudart, arch))
out(
"DEFINE_ncclDevKernel({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {fn_id})\n"
.format(sym=sym, coll=coll, redop_cxx=redop_to_cxx[redop], ty_cxx=ty_to_cxx[ty],
algo=(algo or "RING"), proto=(proto or "SIMPLE"), fn_id=fn_id)
)
if (cudart, arch) != (0, 0):
out("#endif\n")
# Add conditional compilation logic around s. If CUDART_VERSION is satisfactory
# we must compile a kernel regardless of __CUDA_ARCH__ since the host code has
# to link against some stub.
s = "#if CUDART_VERSION >= {cudart}\n" \
" #if __CUDA_ARCH__ < {arch}\n" \
" DEFINE_ncclDevKernel_nop({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {fn_id})\n" \
" #else\n" \
" " + s + \
" #endif\n" \
"#endif\n"
out(s.format(
cudart=cudart, arch=arch, sym=sym, coll=coll,
redop_cxx=redop_to_cxx[redop], ty_cxx=ty_to_cxx[ty],
algo=(algo or "RING"), proto=(proto or "SIMPLE"), fn_id=fn_id
))
for fn in fns:
(coll, redop, ty, algo, proto) = fn
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
View File
@@ -62,6 +62,10 @@ 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(__CUDA_FP8_TYPES_EXIST__) && __CUDA_ARCH__ >= 900
case ncclFloat8e4m3: kernel = (void const*)&oneRankReduce<FuncPreMulSum<__nv_fp8_e4m3>>; break;
case ncclFloat8e5m2: kernel = (void const*)&oneRankReduce<FuncPreMulSum<__nv_fp8_e5m2>>; break;
#endif
case ncclFloat16: kernel = (void const*)&oneRankReduce<FuncPreMulSum<half>>; break;
#if defined(__CUDA_BF16_TYPES_EXIST__)
case ncclBfloat16: kernel = (void const*)&oneRankReduce<FuncPreMulSum<__nv_bfloat16>>; break;
projects/rccl/src/device/primitives.h
+5 -2
View File
@@ -103,7 +103,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.
@@ -121,9 +121,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
@@ -4,9 +4,9 @@
* See LICENSE.txt for license information
************************************************************************/
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
@@ -8,9 +8,9 @@
#define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS-1)
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
+145 -93
View File
@@ -14,9 +14,9 @@ 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;
@@ -34,11 +34,7 @@ class Primitives<
PatMode = 0x800,
NvlsMinPolling = 0x1000,
NetDeviceUnpack = 0x2000,
AnyNetDeviceUnpack = 0x4000,
NvlsDirectRead = 0x8000,
NvlsDirectWrite = 0x10000,
IpcWrite = 0x20000,
IpcRead = 0x40000;
AnyNetDeviceUnpack = 0x4000;
const int tid, tidInBlock;
const int nthreads;
int nworkers;
@@ -119,12 +115,9 @@ class Primitives<
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;
while (connStepCache + (isSendNotRecv ? NCCL_STEPS : 0) < step + StepPerSlice) {
connStepCache = loadStepValue(connStepPtr);
@@ -134,27 +127,38 @@ class Primitives<
}
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;
@@ -198,7 +202,7 @@ class Primitives<
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.
@@ -252,7 +256,7 @@ class Primitives<
* 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]) {
reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, MaxSend, /*PreOpSrcs*/0>
(tid, nworkers, /*redArg*/0, /*preOpArgs*/nullptr, /*postOp*/false,
1, ncclShmem.groups[group].srcs,
@@ -269,16 +273,32 @@ class Primitives<
} else if (ncclShmem.groups[group].srcs[0] && ncclShmem.groups[group].dsts[0]) {
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);
}
} 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.
@@ -295,10 +315,11 @@ class Primitives<
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;
}
@@ -347,17 +368,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
@@ -440,7 +461,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) {
@@ -452,7 +473,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();
@@ -460,7 +481,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;
@@ -474,7 +495,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
@@ -499,33 +520,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);
@@ -544,27 +566,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;
}
}
}
@@ -574,8 +595,8 @@ private:
__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_) {
@@ -643,11 +664,23 @@ private:
// 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;
@@ -659,8 +692,10 @@ private:
}
}
// 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
}
__device__ ~Primitives() {
@@ -683,6 +718,16 @@ private:
// 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) {
@@ -693,10 +738,10 @@ private:
}
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;
@@ -709,6 +754,7 @@ private:
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);
@@ -727,6 +773,7 @@ private:
directBuff = reinterpret_cast<T*>(ptr);
*slot = nullptr;
} else {
// coverity[var_deref_op]
directBuff = (T*)work->dnOutputs[index];
}
}
@@ -747,8 +794,10 @@ private:
} 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]);
}
@@ -837,11 +886,11 @@ private:
__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);
@@ -860,6 +909,9 @@ private:
__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
+139 -32
View File
@@ -20,6 +20,12 @@ struct IsFloatingPoint<half>: std::true_type {};
template<>
struct IsFloatingPoint<__nv_bfloat16>: std::true_type {};
#endif
#if defined(__CUDA_FP8_TYPES_EXIST__)
template<>
struct IsFloatingPoint<__nv_fp8_e4m3>: std::true_type {};
template<>
struct IsFloatingPoint<__nv_fp8_e5m2>: std::true_type {};
#endif
template<>
struct IsFloatingPoint<float>: std::true_type {};
template<>
@@ -298,6 +304,24 @@ SPECIALIZE_REDUCE(FuncMinMax, double, 1, double, fn.isMinNotMax ? fmin(x, y) : f
#endif
#endif
#if defined(__CUDA_FP8_TYPES_EXIST__)
#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))))
#endif
#endif
#undef SPECIALIZE_REDUCE
////////////////////////////////////////////////////////////////////////////////
@@ -416,9 +440,9 @@ template<>
struct FuncPreMulSum<half> {
using EltType = half;
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
half2 scalar;
__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;
@@ -426,9 +450,9 @@ struct FuncPreMulSum<half> {
#else
float scalar;
__device__ FuncPreMulSum(uint64_t opArg=0) {
union { uint64_t u64; half val; };
union { uint64_t u64; __half val; };
u64 = opArg;
scalar = __half2float(val);
scalar = (float)val;
}
#endif
};
@@ -459,11 +483,39 @@ struct FuncPreMulSum<half> {
};
#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) {
#if defined(__CUDA_FP8_TYPES_EXIST__)
#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;
}
};
#endif
#endif
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);
}
};
@@ -530,6 +582,51 @@ struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
#endif
#endif
////////////////////////////////////////////////////////////////////////////////
// Apply_PreOp of FuncPreMulSum for fp8.
#if defined(__CUDA_FP8_TYPES_EXIST__)
#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)));
}
};
#endif
#endif
////////////////////////////////////////////////////////////////////////////////
// FuncSumPostDiv
@@ -541,34 +638,44 @@ 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 FuncSumPostDiv {
static_assert(T(0) < T(-1), "FuncSumPostDiv is only for implementing ncclAvg on uint types.");
using EltType = T;
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 = 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=*/false>: FuncSum<T> {
using EltType = T;
int divisor;
__device__ FuncSumPostDiv_IntOnly(uint64_t opArg=0): divisor(opArg) {}
};
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);
}
};
@@ -576,7 +683,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
@@ -89,7 +89,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;
@@ -137,6 +137,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>;
@@ -206,10 +207,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;
@@ -222,15 +223,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,
@@ -239,7 +240,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;
},
@@ -264,7 +265,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) __trap();
bool isMultiRail = (direct->nHeads > 1);
@@ -281,15 +283,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;
}
@@ -297,23 +299,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;
@@ -322,9 +323,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();
@@ -333,11 +334,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
@@ -15,33 +15,35 @@ 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;
Primitives<T, RedOp, FanAsymmetric<0, 1>, 1, Proto, 1>
prims(tid, tn, nullptr, &work->sendRank, work->sendAddr, nullptr,
/*redOpArg(ignored)=*/0, group, 1, 1, nullptr,
/*ipcReg=*/work->sendIpcReg, /*netReg=*/work->sendRegistered, ncclShmem.comm.p2pChunkSize);
/*redOpArg(ignored)=*/0, group, 1, 1, nullptr, work, stepSize);
size_t cursor = 0;
do {
int n = min(size_t(chunkSize), bytes-cursor);
prims.directSend(cursor, cursor, n);
cursor += n;
} while (cursor < bytes && work->sendRegistered == 0);
} while (cursor < bytes);
}
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;
Primitives<T, RedOp, FanAsymmetric<1, 0>, 1, Proto, 1>
prims(tid, tn, &work->recvRank, nullptr, nullptr, work->recvAddr,
/*redOpArg(ignored)=*/0, group, 1, 1, nullptr,
/*ipcReg=*/work->recvIpcReg, /*netReg=*/work->recvRegistered, ncclShmem.comm.p2pChunkSize);
/*redOpArg(ignored)=*/0, group, 1, 1, nullptr, work, stepSize);
size_t cursor = 0;
do {
int n = min(size_t(chunkSize), bytes-cursor);
prims.directRecv(cursor, cursor, n);
cursor += n;
} while (cursor < bytes && work->recvRegistered == 0);
} while (cursor < bytes);
}
__device__ __forceinline__ void run() {