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

2.10.3-1

Browse Source 
Add support for bfloat16.
Add ncclAvg reduction operation.
Improve performance for aggregated operations.
Improve performance for tree.
Improve network error reporting.
Add NCCL_NET parameter to force a specific network.
Add NCCL_IB_QPS_PER_CONNECTION parameter to split IB traffic onto multiple queue pairs.
Fix topology detection error in WSL2.
Fix proxy memory elements affinity (improve alltoall performance).
Fix graph search on cubemesh topologies.
Fix hang in cubemesh during NVB connections.
This commit is contained in:
Ke Wen
2021-07-08 14:12:04 -07:00
parent 3fec2fa5ee
commit 7e51592129
52 changed files with 3496 additions and 2469 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/device/Makefile
+2 -2
View File
@@ -1,5 +1,5 @@
#
# Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
#
# See LICENSE.txt for license information
#
@@ -32,7 +32,7 @@ all_deps: $(DEPENDFILES)
$(RULESFILE) :
@printf "Generating %-35s > %s\n" rules $@
@mkdir -p $(OBJDIR)
@./gen_rules.sh $(OBJDIR) > $@
@CUDA_MAJOR=${CUDA_MAJOR} CUDA_MINOR=${CUDA_MINOR} ./gen_rules.sh $(OBJDIR) > $@
-include $(RULESFILE)
src/collectives/device/all_gather.h
+72 -181
View File
@@ -5,204 +5,95 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllGather, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-WARP_SIZE;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLGATHER_CHUNKSTEPS;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
const ssize_t size = args->coll.count;
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
ncclRing *ring = &ncclShmem.channel.ring;
const int *ringRanks = ring->devUserRanks;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLGATHER_CHUNKSTEPS : 1));
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2);
const int nranks = ncclShmem.comm.nRanks;
const ssize_t loopSize = nChannels*int(chunkSize);
const ssize_t size = args->coll.count;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
T *inputBuf = (T*)args->sendbuff;
T *outputBuf = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 1, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf);
ncclPrimitives<UNROLL, ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS, T, 1, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nChannels));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t chunkOffset = gridOffset + bid*realChunkSize;
/////////////// begin AllGather steps ///////////////
ssize_t offset;
int nelem = min(realChunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[0];
offset = chunkOffset + rankDest * size;
if (thisInput + chunkOffset == thisOutput + offset) { // In place
prims.directSend(thisInput+chunkOffset, offset, nelem);
} else {
prims.directCopySend(thisInput+chunkOffset, thisOutput+offset, offset, nelem);
}
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
prims.directRecvCopySend(thisOutput+offset, offset, nelem);
}
// Make final copy from buffer to dest.
rankDest = ring->devUserRanks[1];
offset = chunkOffset + rankDest * size;
// Final wait/copy.
prims.directRecv(thisOutput+offset, offset, nelem);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
realChunkSize = min(chunkSize, divUp(size-gridOffset,nChannels));
realChunkSize = roundUp(realChunkSize, (nthreads-WARP_SIZE)*sizeof(uint64_t)/sizeof(T));
}
}
};
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
realChunkSize = int(realChunkSize);
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllGather, NCCL_ALGO_RING, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
ssize_t chunkOffset = gridOffset + int(bid*realChunkSize);
ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepLines, channel, comm);
/////////////// begin AllGather steps ///////////////
ssize_t offset;
int nelem = min(realChunkSize, size-chunkOffset);
int rankDest;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
// step 0: push data to next GPU
rankDest = ringRanks[0];
offset = chunkOffset + rankDest * size;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
if (size-gridOffset < loopSize) {
chunkSize = args->coll.lastChunkSize;
}
ssize_t chunkOffset = gridOffset + bid*chunkSize;
/////////////// begin AllGather steps ///////////////
ssize_t offset;
int nelem = min(chunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[0];
offset = chunkOffset + rankDest * size;
if (thisInput + chunkOffset == thisOutput + offset) { // In place
LLprims.send(thisInput+chunkOffset, nelem);
} else {
LLprims.copySend(thisInput+chunkOffset, thisOutput+offset, nelem);
}
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
LLprims.recvCopySend(thisOutput+offset, nelem);
}
// step k-1: final store
rankDest = ring->devUserRanks[1];
offset = chunkOffset + rankDest * size;
LLprims.recv(thisOutput+offset, nelem);
if (inputBuf + chunkOffset == outputBuf + offset) { // In place
prims.directSend(chunkOffset, offset, nelem);
} else {
prims.directCopySend(chunkOffset, offset, offset, nelem);
}
}
};
#include "prims_ll128.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllGather, NCCL_ALGO_RING, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = stepSize*NCCL_LL128_DATAELEMS*sizeof(uint64_t) / (NCCL_LL128_LINEELEMS*sizeof(T));
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T))/2;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepSize, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*minChunkSize)*minChunkSize, chunkSize);
ssize_t chunkOffset = gridOffset + bid*chunkSize;
/////////////// begin AllGather steps ///////////////
ssize_t offset;
int nelem = min(chunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[0];
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ringRanks[nranks-j];
offset = chunkOffset + rankDest * size;
if (thisInput + chunkOffset == thisOutput + offset) { // In place
LLprims.send(thisInput+chunkOffset, nelem);
} else {
LLprims.copySend(thisInput+chunkOffset, thisOutput+offset, nelem);
}
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
LLprims.recvCopySend(thisOutput+offset, nelem);
}
// step k-1: final store
rankDest = ring->devUserRanks[1];
offset = chunkOffset + rankDest * size;
LLprims.recv(thisOutput+offset, nelem);
prims.directRecvCopySend(offset, offset, nelem);
}
// Make final copy from buffer to dest.
rankDest = ringRanks[1];
offset = chunkOffset + rankDest * size;
// Final wait/copy.
prims.directRecv(offset, nelem);
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
using Proto = ProtoSimple<ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS>;
runRing<T, RedOp, Proto>(args);
}
};
template<int PROTO, class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllGather, NCCL_ALGO_TREE, PROTO, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL>(args);
}
};
template<int PROTO, class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllGather, NCCL_ALGO_COLLNET, PROTO, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL128>(args);
}
};
src/collectives/device/all_reduce.h
+280 -462
View File
@@ -5,566 +5,384 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-WARP_SIZE;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLREDUCE_CHUNKSTEPS;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
ncclRing *ring = &ncclShmem.channel.ring;
int ringIx = ring->index;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLREDUCE_CHUNKSTEPS : 1));
const int nranks = ncclShmem.comm.nRanks;
const ssize_t loopSize = nChannels*nranks*chunkSize;
const ssize_t size = args->coll.count;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
int minChunkSize;
if (Proto::Id == NCCL_PROTO_LL)
minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T));
if (Proto::Id == NCCL_PROTO_LL128) {
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2;
}
ncclPrimitives<UNROLL, ALLREDUCE_CHUNKSTEPS/ALLREDUCE_SLICESTEPS, ALLREDUCE_SLICESTEPS, T, 1, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
Primitives<T, RedOp, FanSymmetric<1>, 1, Proto> prims
(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += nranks*loopSize) {
ssize_t realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nranks*nChannels));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t chunkOffset = gridOffset + bid*nranks*realChunkSize;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*nranks));
realChunkSize = roundUp(realChunkSize, (nthreads-WARP_SIZE)*sizeof(uint64_t)/sizeof(T));
}
else
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*nranks*minChunkSize)*minChunkSize);
realChunkSize = int(realChunkSize);
auto calcOffset = [&]__device__(int chunk)->ssize_t {
if (Proto::Id == NCCL_PROTO_SIMPLE)
return gridOffset + bid*nranks*realChunkSize + chunk*realChunkSize;
else
return gridOffset + (chunk*nChannels + bid)*realChunkSize;
};
auto modRanks = [&]__device__(int r)->int {
return r - (r >= nranks ? nranks : 0);
};
/////////////// begin AllReduce steps ///////////////
ssize_t offset;
int nelem;
int chunk;
// step 0: push data to next GPU
chunk = ring->devUserRanks[nranks-1];
offset = chunkOffset + chunk * realChunkSize;
chunk = modRanks(ringIx + nranks-1);
offset = calcOffset(chunk);
nelem = min(realChunkSize, size-offset);
prims.send(thisInput+offset, nelem);
prims.send(offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = chunkOffset + chunk * realChunkSize;
chunk = modRanks(ringIx + nranks-j);
offset = calcOffset(chunk);
nelem = min(realChunkSize, size-offset);
prims.recvReduceSend(thisInput+offset, nelem);
prims.recvReduceSend(offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final
// result that we store in this data and push to the next GPU
chunk = ring->devUserRanks[0];
offset = chunkOffset + chunk * realChunkSize;
chunk = ringIx + 0;
offset = calcOffset(chunk);
nelem = min(realChunkSize, size-offset);
prims.directRecvReduceCopySend(thisInput+offset, thisOutput+offset, offset, nelem);
prims.directRecvReduceCopySend(offset, offset, offset, nelem, /*postOp=*/true);
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = chunkOffset + chunk * realChunkSize;
chunk = modRanks(ringIx + nranks-j);
offset = calcOffset(chunk);
nelem = min(realChunkSize, size-offset);
prims.directRecvCopySend(thisOutput+offset, offset, nelem);
prims.directRecvCopySend(offset, offset, nelem);
}
// Make final copy from buffer to dest.
chunk = ring->devUserRanks[1];
offset = chunkOffset + chunk * realChunkSize;
chunk = modRanks(ringIx + 1);
offset = calcOffset(chunk);
nelem = min(realChunkSize, size-offset);
// Final wait/copy.
prims.directRecv(thisOutput+offset, offset, nelem);
prims.directRecv(offset, nelem);
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
template<typename T, typename RedOp, typename Proto>
__device__ void runTreeUpDown(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-2*WARP_SIZE;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* tree = &channel->tree;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
int chunkSize = args->coll.lastChunkSize;
const ssize_t minChunkSize = nthreads*8*sizeof(uint64_t) / sizeof(T);
const ssize_t loopSize = nChannels*chunkSize;
ncclTree *tree = &ncclShmem.channel.tree;
ssize_t chunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? args->coll.lastChunkSize
/* LL & LL128 */ : Proto::calcBytePerStep()/sizeof(T));
const ssize_t minChunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? (nthreads-2*WARP_SIZE)*8*(sizeof(uint64_t)/sizeof(T))
/* LL & LL128 */ : nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const ssize_t loopSize = int(nChannels*chunkSize);
const ssize_t size = args->coll.count;
if (loopSize > size) {
chunkSize = DIVUP(size, nChannels*minChunkSize)*minChunkSize;
}
if (loopSize > size)
chunkSize = divUp((int)size, int(nChannels*minChunkSize))*int(minChunkSize);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
#if 1
if (tid < nthreads+WARP_SIZE) {
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_DEV_ARITY, 1, 0, FUNC>
prims(tid, nthreads, tree->down, &tree->up, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->up == -1) {
prims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);
} else if (tree->down[0] == -1) {
prims.send(thisInput+offset, nelem);
} else {
prims.recvReduceSend(thisInput+offset, nelem);
{ // Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto> prims
(tid, nthreads, tree->down, &tree->up, args->sendbuff, args->recvbuff);
if (tree->up == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.recvReduceCopy(offset, offset, nelem, /*postOp=*/true);
}
}
else if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.send(offset, nelem);
}
}
else {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.recvReduceSend(offset, nelem);
}
}
}
if (tid < nthreads+WARP_SIZE) {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, 1, NCCL_MAX_DEV_ARITY, 1, FUNC>
prims(tid, nthreads, &tree->up, tree->down, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->up == -1) {
prims.directSend(thisOutput+offset, offset, nelem);
} else if (tree->down[0] == -1) {
prims.directRecv(thisOutput+offset, offset, nelem);
} else {
prims.directRecvCopySend(thisOutput+offset, offset, nelem);
{ // Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto> prims
(tid, nthreads, &tree->up, tree->down, args->sendbuff, args->recvbuff);
if (tree->up == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directSendFromOutput(offset, offset, nelem);
}
}
else if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecv(offset, nelem);
}
}
else {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecvCopySend(offset, offset, nelem);
}
}
}
#else
int nthreadsSplit = nthreads/2;
if (nthreadsSplit >= 256) nthreadsSplit += 64;
}
template<typename T, typename RedOp, typename Proto>
__device__ void runTreeSplit(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
ncclTree *tree = &ncclShmem.channel.tree;
ssize_t chunkSize = int(
Proto::Id != NCCL_PROTO_LL ? args->coll.lastChunkSize
: Proto::calcBytePerStep()/sizeof(T));
const ssize_t minChunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? (nthreads - 2*WARP_SIZE)*8*(sizeof(uint64_t)/sizeof(T)) :
Proto::Id == NCCL_PROTO_LL ? nthreads*(Proto::calcBytePerGrain()/sizeof(T))
/* LL128 */ : nthreads*(Proto::calcBytePerGrain()/sizeof(T))/8);
const ssize_t loopSize = int(nChannels*chunkSize);
const ssize_t size = args->coll.count;
int nthreadsSplit;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
nthreadsSplit = nthreads/2;
if (nthreadsSplit >= 256) nthreadsSplit += 64;
} else { // LL & LL128
// Receiving from up to 3 sources is more compute intensive than sending
// to 3 dests. Use 70% for reduce and 30% for bcast.
nthreadsSplit = (nthreads*7/(10*WARP_SIZE))*WARP_SIZE;
}
if (loopSize > size)
chunkSize = divUp((int)size, nChannels*int(minChunkSize))*int(minChunkSize);
if (tree->up == -1) {
if (tid < nthreads+WARP_SIZE) {
// ReduceAndBroadcast : max number of recv is 3, max number of send is 3
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_DEV_ARITY, NCCL_MAX_DEV_ARITY, 1, FUNC>
prims(tid, nthreads, tree->down, tree->down, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
// Reduce and broadcast. Max number of recv is 3, max number of send is 3
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto>
prims(tid, nthreads, tree->down, tree->down, args->sendbuff, args->recvbuff);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecvReduceCopySend(offset, offset, offset, nelem, /*doPost=*/true);
}
}
else if (tid < nthreadsSplit) {
/* Reduce up. Max number of recv is 3, max number of send is 1 (binary tree + local).
* Why Direct=1????
* Answer: Because despite not performing any direct operations, the ctor
* must assume Direct so that it can exchange direct pointers with remote ctors
* that are Direct, otherwise it hangs. A cleaner solution would be to seperate
* into DirectRecv and DirectSend capabilities, this ctor would have both=0,
* but the ctor above for tree roots would be DirectRecv=0 DirectSend=1.
*/
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/1, Proto>
prims(tid, nthreadsSplit, tree->down, &tree->up, args->sendbuff, args->recvbuff, 0*Proto::MaxGroupWidth);
if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*chunkSize;
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecvReduceCopySend(thisInput+offset, thisOutput+offset, offset, nelem);
prims.send(offset, nelem);
}
}
} else {
if (tid < nthreadsSplit + WARP_SIZE) {
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_DEV_ARITY, 1, 0, FUNC>
prims(tid, nthreadsSplit, tree->down, &tree->up, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
else {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
if (tree->down[0] == -1) {
prims.send(thisInput+offset, nelem);
} else {
prims.recvReduceSend(thisInput+offset, nelem);
}
}
} else {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, 1, NCCL_MAX_DEV_ARITY, 1, FUNC>
prims(tid-nthreadsSplit-WARP_SIZE, nthreads-nthreadsSplit, &tree->up, tree->down, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 2);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->down[0] == -1) {
prims.directRecv(thisOutput+offset, offset, nelem);
} else {
prims.directRecvCopySend(thisOutput+offset, offset, nelem);
}
prims.recvReduceSend(offset, nelem);
}
}
}
#endif
else {
// Broadcast down. Max number of recv is 1, max number of send is 3 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &tree->up, tree->down, args->sendbuff, args->recvbuff, 1*Proto::MaxGroupWidth);
if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecv(offset, nelem);
}
}
else {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
prims.directRecvCopySend(offset, offset, nelem);
}
}
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
using Proto = ProtoSimple<ALLREDUCE_CHUNKSTEPS/ALLREDUCE_SLICESTEPS, ALLREDUCE_SLICESTEPS>;
runRing<T, RedOp, Proto>(args);
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
#define COLLNET_COPY_THREADS 96
public:
__device__ void run(struct ncclWorkElem* args) {
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
#if CUDART_VERSION >= 11020 && CUDART_VERSION < 11040 && __CUDA_ARCH__ >= 800
runTreeUpDown<T, RedOp, ProtoSimple<1, 1>>(args);
#else
runTreeSplit<T, RedOp, ProtoSimple<1, 1>>(args);
#endif
}
};
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
static constexpr int COLLNET_COPY_THREADS = 96;
const int tid = threadIdx.x;
//const int nthreads = args->nThreads-3*WARP_SIZE;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclDirect* tree = &channel->collTree;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
int chunkSize = args->coll.lastChunkSize;
struct ncclDirect* tree = &ncclShmem.channel.collTree;
const ssize_t chunkSize = int(args->coll.lastChunkSize);
const ssize_t size = args->coll.count;
const ssize_t loopSize = nChannels*tree->nHeads*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
const int hasUp = (tree->up[0] >= 0) ? 1 : 0;
const int hasDn = (tree->down[0] >= 0) ? 1 : 0;
const int nThreadsScatter = (hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 3*COLLNET_COPY_THREADS : 0;
const int nThreadsGather = (hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 2*COLLNET_COPY_THREADS : 0;
const int nThreadsBcast = (hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 0 : 2*COLLNET_COPY_THREADS;
// Gather does not need sync threads, sparing one more warp for reduce
const int nThreadsReduce = NCCL_SIMPLE_MAX_NTHREADS + WARP_SIZE - nThreadsScatter - nThreadsGather - nThreadsBcast;
const int nThreadsScatter = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 3*COLLNET_COPY_THREADS : 0);
const int nThreadsGather = ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 2*COLLNET_COPY_THREADS : 0);
const int nThreadsBcast = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 0 : 2*COLLNET_COPY_THREADS);
const int nThreadsReduce = args->nThreads - nThreadsScatter - nThreadsGather - nThreadsBcast;
const int tidStartBcast = nThreadsGather;
const int tidStartScatter = tidStartBcast + nThreadsBcast + WARP_SIZE;
const int tidStartReduce = tidStartScatter + nThreadsScatter + WARP_SIZE;
const int tidStartScatter = tidStartBcast + nThreadsBcast;
const int tidStartReduce = tidStartScatter + nThreadsScatter;
using Proto = ProtoSimple<1, 1>;
if (tid >= tidStartScatter && tid < tidStartReduce && hasUp) {
// Scatter
ncclPrimitives<UNROLL, 1, 1, T, 0, NCCL_MAX_DIRECT_ARITY, 0, FUNC>
prims(tid-tidStartScatter, nThreadsScatter, NULL, tree->up, NULL, stepSize, channel, comm, ncclShmem->ptrs, 4);
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto>
prims(tid-tidStartScatter, nThreadsScatter, NULL, tree->up, args->sendbuff, args->recvbuff, 2*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*tree->nHeads*chunkSize;
int nelem = min(tree->nHeads*chunkSize, size-offset);
prims.scatter(thisInput+offset, nelem, chunkSize, tree->headRank, tree->shift);
prims.scatter(offset, nelem, chunkSize, tree->headRank, tree->shift);
}
} else if (tid >= tidStartReduce && tree->out != -1) {
// Reduce, send to network
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_DIRECT_ARITY, 1, 0, FUNC>
prims(tid-tidStartReduce, nThreadsReduce, tree->down, &tree->out, NULL, stepSize, channel, comm, ncclShmem->ptrs, 6);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
if (hasDn) {
prims.recvReduceSend(thisInput+offset, nelem);
} else {
prims.send(thisInput+offset, nelem);
if (hasDn) {
// Reduce, send to network
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto>
prims(tid-tidStartReduce, nThreadsReduce, tree->down, &tree->out, args->sendbuff, args->recvbuff, 3*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
prims.recvReduceSend(offset, nelem);
}
} else {
// Directly send to network
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto>
prims(tid-tidStartReduce, nThreadsReduce, nullptr, &tree->out, args->sendbuff, args->recvbuff, 3*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
prims.send(offset, nelem);
}
}
} else if (tid < tidStartBcast && hasUp) {
// Gather
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_DIRECT_ARITY, 0, 0, FUNC>
prims(tid, nThreadsGather, tree->up, NULL, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/0, Proto>
prims(tid, nThreadsGather, tree->up, NULL, args->sendbuff, args->recvbuff, 0*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*tree->nHeads*chunkSize;
int nelem = min(tree->nHeads*chunkSize, size-offset);
prims.gather(thisOutput+offset, nelem, chunkSize, tree->headRank, tree->shift);
prims.gather(offset, nelem, chunkSize, tree->headRank, tree->shift);
}
} else if (tid >= tidStartBcast && tid < tidStartScatter && tree->out != -1) {
// Recv from network, broadcast
ncclPrimitives<UNROLL, 1, 1, T, 1, NCCL_MAX_DIRECT_ARITY, 0, FUNC>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, tree->down, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 2);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
if (hasDn) {
prims.recvCopySend(thisOutput+offset, nelem);
} else {
prims.recv(thisOutput+offset, nelem);
}
}
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_RING, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const ssize_t minChunkSize = nthreads * (sizeof(uint64_t)) / sizeof(T);
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*nranks*chunkSize;
const ssize_t size = args->coll.count;
ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepLines, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*nranks*minChunkSize)*minChunkSize, chunkSize);
/////////////// begin AllReduce steps ///////////////
ssize_t offset;
int nelem;
int chunk;
// step 0: push data to next GPU
chunk = ring->devUserRanks[nranks-1];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.send(thisInput+offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceSend(thisInput+offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final
// result that we store in this data and push to the next GPU
chunk = ring->devUserRanks[0];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceCopySend(thisInput+offset, thisOutput+offset, nelem);
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvCopySend(thisOutput+offset, nelem);
}
// Make final copy from buffer to dest.
chunk = ring->devUserRanks[1];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
// Here we need to copy from buffer to this output.
LLprims.recv(thisOutput+offset, nelem);
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_TREE, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* tree = &channel->tree;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const ssize_t minChunkSize = nthreads*sizeof(uint64_t) / sizeof(T);
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
if (loopSize > size) {
chunkSize = DIVUP(size, nChannels*minChunkSize)*minChunkSize;
}
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
do {
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclLLPrimitives<T, FUNC, NCCL_MAX_DEV_ARITY, 1> LLprims(tid, nthreads, tree->down, &tree->up, stepLines, channel, comm);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->up == -1) {
LLprims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);
} else if (tree->down[0] == -1) {
LLprims.send(thisInput+offset, nelem);
} else {
LLprims.recvReduceSend(thisInput+offset, nelem);
}
}
} while(0);
do {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclLLPrimitives<T, FUNC, 1, NCCL_MAX_DEV_ARITY> LLprims(tid, nthreads, &tree->up, tree->down, stepLines, channel, comm);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->up == -1) {
LLprims.send(thisOutput+offset, nelem);
} else if (tree->down[0] == -1) {
LLprims.recv(thisOutput+offset, nelem);
} else {
LLprims.recvCopySend(thisOutput+offset, nelem);
}
}
} while(0);
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_COLLNET, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) { }
};
#include "prims_ll128.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_RING, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = stepSize*NCCL_LL128_DATAELEMS*sizeof(uint64_t) / (NCCL_LL128_LINEELEMS*sizeof(T));
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T))/2;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*nranks*chunkSize;
const ssize_t size = args->coll.count;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepSize, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*nranks*minChunkSize)*minChunkSize, chunkSize);
/////////////// begin AllReduce steps ///////////////
ssize_t offset;
int nelem;
int chunk;
// step 0: push data to next GPU
chunk = ring->devUserRanks[nranks-1];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.send(thisInput+offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceSend(thisInput+offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final
// result that we store in this data and push to the next GPU
chunk = ring->devUserRanks[0];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceCopySend(thisInput+offset, thisOutput+offset, nelem);
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
chunk = ring->devUserRanks[nranks-j];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvCopySend(thisOutput+offset, nelem);
}
// Make final copy from buffer to dest.
chunk = ring->devUserRanks[1];
offset = gridOffset + (chunk*nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
// Here we need to copy from buffer to this output.
LLprims.recv(thisOutput+offset, nelem);
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_TREE, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* tree = &channel->tree;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = args->coll.lastChunkSize;
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T))/8;
const ssize_t loopSize = nChannels*chunkSize;
int nthreadsSplit = NCCL_LL128_SPLIT(nthreads);
const ssize_t size = args->coll.count;
if (loopSize > size) {
chunkSize = DIVUP(size, nChannels*minChunkSize)*minChunkSize;
}
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
if (tree->up == -1) {
// ReduceAndBroadcast : max number of recv is 3, max number of send is 3
ncclLL128Primitives<T, FUNC, NCCL_MAX_DEV_ARITY, NCCL_MAX_DEV_ARITY> LLprims(tid, nthreads, tree->down, tree->down, stepSize, channel, comm);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
LLprims.recvReduceCopySend(thisInput+offset, thisOutput+offset, nelem);
}
} else {
if (tid < nthreadsSplit) {
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclLL128Primitives<T, FUNC, NCCL_MAX_DEV_ARITY, 1> LLprims(tid, nthreadsSplit, tree->down, &tree->up, stepSize, channel, comm);
if (hasDn) {
// Recv from network, broadcast
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, tree->down, args->sendbuff, args->recvbuff, 1*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->down[0] == -1) {
LLprims.send(thisInput+offset, nelem);
} else {
LLprims.recvReduceSend(thisInput+offset, nelem);
}
prims.recvCopySend(offset, nelem, /*postOp=*/true);
}
} else {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclLL128Primitives<T, FUNC, 1, NCCL_MAX_DEV_ARITY> LLprims(tid-nthreadsSplit, nthreads-nthreadsSplit, &tree->up, tree->down, stepSize, channel, comm);
// Recv from network (no post thread needed)
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, nullptr, args->sendbuff, args->recvbuff, 1*Proto::MaxGroupWidth);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
if (tree->down[0] == -1) {
LLprims.recv(thisOutput+offset, nelem);
} else {
LLprims.recvCopySend(thisOutput+offset, nelem);
}
prims.recv(offset, nelem, /*postOp=*/true);
}
}
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_COLLNET, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) { }
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL>(args);
}
};
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runTreeSplit<T, RedOp, ProtoLL>(args);
}
};
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL128>(args);
}
};
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runTreeSplit<T, RedOp, ProtoLL128>(args);
}
};
src/collectives/device/broadcast.h
+60 -140
View File
@@ -5,158 +5,78 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncBroadcast, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-WARP_SIZE;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * BROADCAST_CHUNKSTEPS;
const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
const ssize_t size = args->coll.count;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->coll.root;
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
ncclRing *ring = &ncclShmem.channel.ring;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? BROADCAST_CHUNKSTEPS : 1));
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->coll.root;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
T *inputBuf = (T*)args->sendbuff;
T *outputBuf = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf);
ncclPrimitives<UNROLL, BROADCAST_CHUNKSTEPS/BROADCAST_SLICESTEPS, BROADCAST_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels));
realChunkSize = roundUp(realChunkSize, (nthreads-WARP_SIZE)*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
realChunkSize = int(realChunkSize);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nChannels));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t offset = gridOffset + bid*realChunkSize;
int nelem = min(realChunkSize, size-offset);
ssize_t offset = gridOffset + int(bid*realChunkSize);
int nelem = min(realChunkSize, size-offset);
if (rank == root) {
if (thisInput == thisOutput) {
prims.send(thisInput+offset, nelem);
} else {
prims.copySend(thisInput+offset, thisOutput+offset, nelem);
}
} else if (nextRank == root) {
prims.recv(thisOutput+offset, nelem);
if (rank == root) {
if (inputBuf == outputBuf) {
prims.send(offset, nelem);
} else {
prims.recvCopySend(thisOutput+offset, nelem);
prims.copySend(offset, offset, nelem);
}
} else if (nextRank == root) {
prims.recv(offset, nelem);
} else {
prims.recvCopySend(offset, nelem);
}
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
using Proto = ProtoSimple<BROADCAST_CHUNKSTEPS/BROADCAST_SLICESTEPS, BROADCAST_SLICESTEPS>;
runRing<T, RedOp, Proto>(args);
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncBroadcast, NCCL_ALGO_RING, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->coll.root;
ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepLines, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
if (size-gridOffset < loopSize) {
chunkSize = args->coll.lastChunkSize;
}
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (rank == root) {
if (thisInput == thisOutput) {
LLprims.send(thisInput+offset, nelem);
} else {
LLprims.copySend(thisInput + offset, thisOutput + offset, nelem);
}
} else if (nextRank == root) {
LLprims.recv(thisOutput + offset, nelem);
} else {
LLprims.recvCopySend(thisOutput + offset, nelem);
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL>(args);
}
};
#include "prims_ll128.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncBroadcast, NCCL_ALGO_RING, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = stepSize*NCCL_LL128_DATAELEMS*sizeof(uint64_t) / (NCCL_LL128_LINEELEMS*sizeof(T));
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T));
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->coll.root;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepSize, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*minChunkSize)*minChunkSize, chunkSize);
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (rank == root) {
if (thisInput == thisOutput) {
LLprims.send(thisInput+offset, nelem);
} else {
LLprims.copySend(thisInput + offset, thisOutput + offset, nelem);
}
} else if (nextRank == root) {
LLprims.recv(thisOutput + offset, nelem);
} else {
LLprims.recvCopySend(thisOutput + offset, nelem);
}
}
}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncBroadcast, NCCL_ALGO_TREE, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncBroadcast, NCCL_ALGO_COLLNET, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL128>(args);
}
};
src/collectives/device/common.h
+116 -69
View File
@@ -10,108 +10,152 @@
#include "collectives.h"
#include "devcomm.h"
#if __CUDA_ARCH__ >= 800
#define COLL_UNROLL 8
#define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1) // Using balanced tree instead of split tree
#else
#define COLL_UNROLL 4
#define NCCL_MAX_DEV_ARITY NCCL_MAX_TREE_ARITY
#endif
// Exit If Abort Barrier across CTA: make sure all threads exit consistently
// Each thread sets a predicate to true if abort == 1
// all CTA's threads enter the barrier and do a popc on their predicates being True
// If any of the thread's predicate was True, all the threads call exit()
static inline __device__ void exitIfAbortBarrier(int abort) {
#define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1) // Using balanced tree instead of split tree
__device__ inline bool barrierReduceAny(int bit) {
uint32_t popc;
asm ("{");
asm volatile (" .reg .pred barr_pred;");
asm volatile (" setp.eq.u32 barr_pred,%0,1;" :: "r"(abort));
asm volatile (" bar.red.popc.u32 %0, 0, barr_pred;" : "=r"(popc));
asm ("}");
if (popc) { asm volatile ("exit;"); }
asm ("{"
".reg .pred barr_pred;"
"setp.eq.u32 barr_pred, %1, 1;"
"bar.red.popc.u32 %0, 0, barr_pred;"
"}" : "=r"(popc) : "r"(bit));
return popc != 0;
}
typedef void(*ncclKern_t)(struct ncclWorkElem* args);
extern __device__ ncclKern_t ncclFuncs[];
template<typename T>
__device__ int copyToShmem(T *dst, T const *src, int turn=0) {
static_assert(sizeof(uint64_t) <= alignof(T), "Uhoh");
uint64_t *d = reinterpret_cast<uint64_t*>(dst);
uint64_t const *s = reinterpret_cast<uint64_t const*>(src);
int t = threadIdx.x - turn;
if (t < 0) t += blockDim.x;
int n = sizeof(T)/sizeof(uint64_t);
static __device__ void load_parallel(void* dst, void* src, size_t size, int tid) {
int* d = (int*)dst;
int* s = (int*)src;
for (int o = tid; o < (size/sizeof(int)); o += blockDim.x) d[o] = s[o];
int delta = (n + WARP_SIZE-1) & -WARP_SIZE; // round up to warp lane 0
if (delta < blockDim.x) {
turn += delta;
if (turn >= blockDim.x) turn -= blockDim.x;
}
else
turn = 0;
n -= t;
d += t;
s += t;
#pragma unroll
for (int i=0; i < divUp(sizeof(T), WARP_SIZE*sizeof(uint64_t)); i++) {
if (n > 0) {
*d = *s;
d += blockDim.x;
s += blockDim.x;
n -= blockDim.x;
}
}
return turn;
}
static __device__ void load_coll(struct ncclWork* localWork, struct ncclWork *hostWork, struct ncclWork* workFifo, int tid, struct ncclDevComm* comm) {
load_parallel(localWork, workFifo, sizeof(struct ncclWork), tid);
// Check whether the last operation was aborted and make sure all threads exit
int abort = tid == 0 ? *(comm->abortFlag) : 0;
exitIfAbortBarrier(abort);
if (tid == 0) hostWork->elems[0].active = 0;
}
template <ncclFunc_t FUNCTION, int ALGO, int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
struct RunWorkElement {
__device__ void run(ncclWorkElem*) {
// Put NOT IMPLEMENTED behavior here.
}
};
struct ncclShmemPtrs {
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
struct RunWork {
__device__ void run(ncclWork *w) {
int tid = threadIdx.x;
#pragma unroll 1
for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].active != 0; e++) {
if (tid < w->elems[e].nThreads)
RunWorkElement<Fn, T, RedOp, Algo, Proto>().run(&w->elems[e]);
}
}
};
typedef void(*ncclKern_t)();
extern __device__ ncclKern_t ncclFuncs[];
struct ncclShmemGroup {
ncclConnInfo *recvConns[NCCL_MAX_DIRECT_ARITY];
ncclConnInfo *sendConns[NCCL_MAX_DIRECT_ARITY];
void* srcs[NCCL_MAX_DIRECT_ARITY+1];
void* dsts[NCCL_MAX_DIRECT_ARITY+1];
};
struct ncclShmemData {
union {
volatile uint64_t data[NCCL_LL128_SHMEM_SIZE];
struct ncclShmemPtrs ptrs[NCCL_MAX_GROUPS];
uint64_t ll128warp[NCCL_LL128_MAX_NTHREADS/WARP_SIZE][NCCL_LL128_SHMEM_ELEMS_PER_THREAD*WARP_SIZE];
struct ncclShmemGroup groups[NCCL_MAX_GROUPS];
};
struct ncclWork localWork;
ncclDevComm comm;
ncclChannel channel;
ncclWork work;
};
extern __device__ struct ncclShmemData *ncclShmem;
template <ncclFunc_t FUNCTION, int ALGO, int PROTO, class REDOP, typename T, int UNROLL, int FINDEX>
__device__ void ncclKernel(struct ncclWorkElem first) {
extern __shared__ ncclShmemData ncclShmem;
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int FnIndex>
__device__ void ncclKernel(ncclWorkElem first) {
int tid = threadIdx.x;
int bid = blockIdx.x;
__shared__ struct ncclShmemData shmem;
ncclShmem = &shmem;
auto f = ncclFunction<FUNCTION, ALGO, PROTO, REDOP, T, UNROLL>();
int turn = copyToShmem(&ncclShmem.comm, first.comm);
// get address of channel without incurring indirect load from ncclDevCom::channels
ncclChannel *channel = &((ncclDevCommAndChannels*)first.comm)->channels[bid];
turn = copyToShmem(&ncclShmem.channel, channel, turn);
struct ncclDevComm* comm = first.comm;
struct ncclChannel* channel = comm->channels+bid;
struct ncclWorkElem* w = NULL;
// To optimize for latency, (only) the first operation is passed as argument.
if (bid == 0 && first.active != 0) {
turn = copyToShmem(&ncclShmem.work.elems[0], &first, turn);
if (tid == 0) ncclShmem.work.elems[1].active = 0;
}
__syncthreads(); // publish ncclShmem
/* To optimize for latency, (only) the first operation is passed as argument.*/
if (bid == 0 && first.funcIndex != FUNC_INDEX_P2P) w = &first;
ncclWork *workFifoHost = ncclShmem.channel.workFifo;
ncclWork *workFifoDev = ncclShmem.channel.workFifoDev;
int workFifoIx = ncclShmem.channel.index;
while (1) {
if (w == NULL) {
w = shmem.localWork.elems;
__syncthreads();
load_coll(&shmem.localWork, channel->workFifo+channel->index, channel->workFifoDev+channel->index, tid, comm);
if (bid == 0 && first.active != 0)
goto SkipLoadWork;
while (true) {
copyToShmem(&ncclShmem.work, &workFifoDev[workFifoIx]); // turn no longer helps
{ // Check whether the last operation was aborted and make sure all threads exit
int aborted = tid == 0 ? *ncclShmem.comm.abortFlag : 0;
if (barrierReduceAny(aborted)) // publish ncclShmem.work
break;
if (tid == 0)
workFifoHost[workFifoIx].elems[0].active = 0;
}
if (tid < w->nThreads) {
if (w->funcIndex == FINDEX) {
f.run(w);
} else {
ncclFuncs[w->funcIndex](w);
}
}
if (tid == 0) channel->index = (channel->index+1) % NCCL_MAX_OPS;
if (w->active == 2) {
return;
}
w = NULL;
SkipLoadWork:
workFifoIx = (workFifoIx + 1)%NCCL_MAX_OPS;
if (tid == 0)
channel->index = workFifoIx; // write back to real channel, not shmem shadow
if (ncclShmem.work.elems[0].funcIndex == FnIndex)
RunWork<Fn, T, RedOp, Algo, Proto>().run(&ncclShmem.work);
else
ncclFuncs[ncclShmem.work.elems[0].funcIndex]();
if (ncclShmem.work.elems[0].active == 2)
break;
__syncthreads();
}
}
// Only generate kernels for SUM
#if NCCL_OP == 0
#define IMPL_COLL_KERN(func, algo, proto, redop, type, fIndex) \
__global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(struct ncclWorkElem first) { \
ncclKernel<ncclFunc##func, NCCL_ALGO_##algo, NCCL_PROTO_##proto, Func##redop<type>, type, COLL_UNROLL, fIndex>(first); \
__global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(ncclWorkElem first) { \
ncclKernel<ncclFunc##func, type, Func##redop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex>(first); \
}
#else
#define IMPL_COLL_KERN(func, algo, proto, redop, type, fInded)
@@ -119,9 +163,8 @@ __global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(struct ncclWorkEl
// Examples : AllReduce, RING, LL, Sum, uint8
#define IMPL_COLL_FUNC(func, algo, proto, redop, type) \
__device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)(struct ncclWorkElem* args) { \
auto f = ncclFunction<ncclFunc##func, NCCL_ALGO_##algo, NCCL_PROTO_##proto, Func##redop<type>, type, COLL_UNROLL>(); \
f.run(args); \
__device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)() { \
RunWork<ncclFunc##func, type, Func##redop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto>().run(&ncclShmem.work); \
}
// Only generate inline kernels for LL
@@ -154,6 +197,8 @@ __device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)(struct ncclWorkEl
#define IMPL_COLL2(func, redop) IMPL_COLL3(func, redop, float, ncclFloat32)
#elif NCCL_TYPE == 8
#define IMPL_COLL2(func, redop) IMPL_COLL3(func, redop, double, ncclFloat64)
#elif NCCL_TYPE == 9 && defined(__CUDA_BF16_TYPES_EXIST__)
#define IMPL_COLL2(func, redop) IMPL_COLL3(func, redop, __nv_bfloat16, ncclBfloat16)
#endif
// Reduction define all functions
@@ -165,6 +210,8 @@ __device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)(struct ncclWorkEl
#define IMPL_COLL_R(func) IMPL_COLL2(func, Min);
#elif NCCL_OP == 3
#define IMPL_COLL_R(func) IMPL_COLL2(func, Max);
#elif NCCL_OP == 4
#define IMPL_COLL_R(func) IMPL_COLL2(func, Avg);
#endif
#if NCCL_OP == 0 && NCCL_TYPE == 0
src/collectives/device/common_kernel.h
+309 -51
View File
@@ -1,5 +1,5 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -24,10 +24,19 @@ inline __device__ void loadPtr(void** ptr, T* &v) {
typedef uint64_t PackType;
template<typename Fn>
struct FuncTraits /*{
__device__ static Fn make();
__device__ static T preOp(Fn, T);
__device__ static T postOp(Fn, T);
}*/;
// unpack x and y to elements of type T and apply FUNC to each element
template<class FUNC, typename T>
struct MULTI {
__device__ PackType operator()(const PackType x, const PackType y) const;
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const;
__device__ PackType preOp(FUNC fn, PackType x) const;
__device__ PackType postOp(FUNC fn, PackType x) const;
};
template<class FUNC>
@@ -41,17 +50,39 @@ struct MULTI<FUNC, int8_t> {
};
};
__device__ PackType operator()(const PackType x, const PackType y) const {
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
converter cx, cy, cr;
cx.storage = x;
cy.storage = y;
// for char, we do these as vector ops
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
cr.a = fn(cx.a, cy.a);
cr.b = fn(cx.b, cy.b);
return cr.storage;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
int8_t elt[8];
} u;
u.pack = x;
#pragma unroll
for (int i=0; i < 8; i++)
u.elt[i] = FuncTraits<FUNC>().preOp(fn, u.elt[i]);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
int8_t elt[8];
} u;
u.pack = x;
#pragma unroll
for (int i=0; i < 8; i++)
u.elt[i] = FuncTraits<FUNC>().postOp(fn, u.elt[i]);
return u.pack;
}
};
template<class FUNC>
@@ -65,17 +96,39 @@ struct MULTI<FUNC, uint8_t> {
};
};
__device__ PackType operator()(const PackType x, const PackType y) const {
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
converter cx, cy, cr;
cx.storage = x;
cy.storage = y;
// for char, we do these as vector ops
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
cr.a = fn(cx.a, cy.a);
cr.b = fn(cx.b, cy.b);
return cr.storage;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint8_t elt[8];
} u;
u.pack = x;
#pragma unroll
for (int i=0; i < 8; i++)
u.elt[i] = FuncTraits<FUNC>().preOp(fn, u.elt[i]);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint8_t elt[8];
} u;
u.pack = x;
#pragma unroll
for (int i=0; i < 8; i++)
u.elt[i] = FuncTraits<FUNC>().postOp(fn, u.elt[i]);
return u.pack;
}
};
template<class FUNC>
@@ -89,16 +142,36 @@ struct MULTI<FUNC, int32_t> {
};
};
__device__ PackType operator()(const PackType x, const PackType y) const {
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
converter cx, cy, cr;
cx.storage = x;
cy.storage = y;
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
cr.a = fn(cx.a, cy.a);
cr.b = fn(cx.b, cy.b);
return cr.storage;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
int32_t elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().preOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().preOp(fn, u.elt[1]);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
int32_t elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().postOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().postOp(fn, u.elt[1]);
return u.pack;
}
};
template<class FUNC>
@@ -112,16 +185,36 @@ struct MULTI<FUNC, uint32_t> {
};
};
__device__ PackType operator()(const PackType x, const PackType y) const {
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
converter cx, cy, cr;
cx.storage = x;
cy.storage = y;
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
cr.a = fn(cx.a, cy.a);
cr.b = fn(cx.b, cy.b);
return cr.storage;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint32_t elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().preOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().preOp(fn, u.elt[1]);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint32_t elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().postOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().postOp(fn, u.elt[1]);
return u.pack;
}
};
template<class FUNC>
@@ -129,22 +222,69 @@ struct MULTI<FUNC, half> {
static_assert(sizeof(PackType) == 4 * sizeof(half),
"PackType must be four times the size of half.");
struct PackHalf2 {
half2 a, b;
union Converter {
PackType pack;
half2 h2[2];
};
__device__ PackType operator()(const PackType x, const PackType y) const {
struct PackHalf2 cx, cy, cr;
cx = *(reinterpret_cast<const struct PackHalf2*>(&x));
cy = *(reinterpret_cast<const struct PackHalf2*>(&y));
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
return *(reinterpret_cast<PackType*>(&cr));
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
Converter cx, cy, cr;
cx.pack = x;
cy.pack = y;
cr.h2[0] = fn(cx.h2[0], cy.h2[0]);
cr.h2[1] = fn(cx.h2[1], cy.h2[1]);
return cr.pack;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
Converter c;
c.pack = x;
c.h2[0] = FuncTraits<FUNC>().preOp(fn, c.h2[0]);
c.h2[1] = FuncTraits<FUNC>().preOp(fn, c.h2[1]);
return c.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
Converter c;
c.pack = x;
c.h2[0] = FuncTraits<FUNC>().postOp(fn, c.h2[0]);
c.h2[1] = FuncTraits<FUNC>().postOp(fn, c.h2[1]);
return c.pack;
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<class FUNC>
struct MULTI<FUNC, __nv_bfloat16> {
static_assert(sizeof(PackType) == 4 * sizeof(__nv_bfloat16),
"PackType must be four times the size of __nv_bfloat16.");
union Converter {
PackType pack;
__nv_bfloat162 h2[2];
};
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
Converter cx, cy, cr;
cx.pack = x;
cy.pack = y;
cr.h2[0] = fn(cx.h2[0], cy.h2[0]);
cr.h2[1] = fn(cx.h2[1], cy.h2[1]);
return cr.pack;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
Converter c;
c.pack = x;
c.h2[0] = FuncTraits<FUNC>().preOp(fn, c.h2[0]);
c.h2[1] = FuncTraits<FUNC>().preOp(fn, c.h2[1]);
return c.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
Converter c;
c.pack = x;
c.h2[0] = FuncTraits<FUNC>().postOp(fn, c.h2[0]);
c.h2[1] = FuncTraits<FUNC>().postOp(fn, c.h2[1]);
return c.pack;
}
};
#endif
template<class FUNC>
struct MULTI<FUNC, float> {
static_assert(sizeof(PackType) == 2 * sizeof(float),
@@ -156,46 +296,120 @@ struct MULTI<FUNC, float> {
};
};
__device__ PackType operator()(const PackType x, const PackType y) const {
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
converter cx, cy, cr;
cx.storage = x;
cy.storage = y;
cr.a = FUNC()(cx.a, cy.a);
cr.b = FUNC()(cx.b, cy.b);
cr.a = fn(cx.a, cy.a);
cr.b = fn(cx.b, cy.b);
return cr.storage;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
float elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().preOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().preOp(fn, u.elt[1]);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
float elt[2];
} u;
u.pack = x;
u.elt[0] = FuncTraits<FUNC>().postOp(fn, u.elt[0]);
u.elt[1] = FuncTraits<FUNC>().postOp(fn, u.elt[1]);
return u.pack;
}
};
template<class FUNC>
struct MULTI<FUNC, double> {
static_assert(sizeof(PackType) == sizeof(double),
"PackType must be the same size as double.");
__device__ PackType operator()(const PackType x, const PackType y) const {
double rv = FUNC()(__longlong_as_double(x), __longlong_as_double(y));
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
double rv = fn(__longlong_as_double(x), __longlong_as_double(y));
return __double_as_longlong(rv);
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
double elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().preOp(fn, u.elt);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
double elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().postOp(fn, u.elt);
return u.pack;
}
};
template<class FUNC>
struct MULTI<FUNC, uint64_t> {
static_assert(sizeof(PackType) == sizeof(uint64_t),
"PackType must be the same size as uint64_t.");
__device__ PackType operator()(const PackType x, const PackType y) const {
uint64_t rv = FUNC()(x, y);
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
uint64_t rv = fn(x, y);
return rv;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint64_t elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().preOp(fn, u.elt);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
uint64_t elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().postOp(fn, u.elt);
return u.pack;
}
};
template<class FUNC>
struct MULTI<FUNC, int64_t> {
static_assert(sizeof(PackType) == sizeof(int64_t),
"PackType must be the same size as int64_t.");
__device__ PackType operator()(const PackType x, const PackType y) const {
int64_t rv = FUNC()((int64_t)x, (int64_t)y);
__device__ PackType operator()(FUNC fn, const PackType x, const PackType y) const {
int64_t rv = fn((int64_t)x, (int64_t)y);
return rv;
}
__device__ PackType preOp(FUNC fn, PackType x) const {
union {
PackType pack;
int64_t elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().preOp(fn, u.elt);
return u.pack;
}
__device__ PackType postOp(FUNC fn, PackType x) const {
union {
PackType pack;
int64_t elt;
} u;
u.pack = x;
u.elt = FuncTraits<FUNC>().postOp(fn, u.elt);
return u.pack;
}
};
template<typename T> inline __device__
@@ -234,13 +448,35 @@ void vStore<half>(volatile half* ptr, const half val) {
}
#endif
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<> inline __device__
__nv_bfloat16 vFetch<__nv_bfloat16>(const volatile __nv_bfloat16* ptr) {
__nv_bfloat16 r;
r = ((__nv_bfloat16*)ptr)[0];
return r;
}
template<> inline __device__
void vStore<__nv_bfloat16>(volatile __nv_bfloat16* ptr, const __nv_bfloat16 val) {
((__nv_bfloat16*)ptr)[0] = val;
}
#endif
typedef ulong2 Pack128;
template<class FUNC, typename T>
struct MULTI128 {
__device__ void operator()(Pack128& x, Pack128& y) {
x.x = MULTI<FUNC, T>()(x.x, y.x);
x.y = MULTI<FUNC, T>()(x.y, y.y);
__device__ void operator()(FUNC fn, Pack128& x, Pack128 const& y) const {
x.x = MULTI<FUNC, T>()(fn, x.x, y.x);
x.y = MULTI<FUNC, T>()(fn, x.y, y.y);
}
__device__ void preOp(FUNC fn, Pack128 &x) const {
x.x = MULTI<FUNC, T>().preOp(fn, x.x);
x.y = MULTI<FUNC, T>().preOp(fn, x.y);
}
__device__ void postOp(FUNC fn, Pack128 &x) const {
x.x = MULTI<FUNC, T>().postOp(fn, x.x);
x.y = MULTI<FUNC, T>().postOp(fn, x.y);
}
};
@@ -253,7 +489,8 @@ inline __device__ void Store128(Pack128* p, Pack128& v) {
template<class FUNC, typename T, int UNROLL, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const int t,
int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Nelem) {
FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Nelem
) {
const int inc = nw * UNROLL * WARP_SIZE;
int offset = w * UNROLL * WARP_SIZE + t;
@@ -266,22 +503,30 @@ __device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const
T vals[UNROLL];
// Load and reduce
for (int u = 0; u < UNROLL; ++u) vals[u] = vFetch(srcs[0]+u*WARP_SIZE);
if (preOpSrc0) {
for (int u = 0; u < UNROLL; ++u) vals[u] = FuncTraits<FUNC>().preOp(fn, vals[u]);
}
#pragma unroll
for (int i=1; i<MINSRCS; i++) {
T vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) vals2[u] = vFetch(srcs[i]+u*WARP_SIZE);
for (int u = 0; u < UNROLL; ++u) vals[u] = FUNC()(vals[u], vals2[u]);
for (int u = 0; u < UNROLL; ++u) vals[u] = fn(vals[u], vals2[u]);
}
#pragma unroll
for (int i=MINSRCS; i<MAXSRCS; i++) {
if (i<nsrcs) {
T vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) vals2[u] = vFetch(srcs[i]+u*WARP_SIZE);
for (int u = 0; u < UNROLL; ++u) vals[u] = FUNC()(vals[u], vals2[u]);
for (int u = 0; u < UNROLL; ++u) vals[u] = fn(vals[u], vals2[u]);
}
}
if (postOp) {
#pragma unroll
for (int u = 0; u < UNROLL; ++u) vals[u] = FuncTraits<FUNC>().postOp(fn, vals[u]);
}
// Store
#pragma unroll
for (int i = 0; i < MINDSTS; i++) {
@@ -301,7 +546,8 @@ __device__ __forceinline__ void ReduceCopyMulti(const int w, const int nw, const
template<class FUNC, typename T, int UNROLL, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceCopy128bMulti(const int w, const int nw, const int t,
int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Npack) {
FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** s, int ndsts, T** d, const int elemOffset, const int Npack
) {
const int inc = nw * UNROLL * WARP_SIZE;
int offset = w * UNROLL * WARP_SIZE + t;
@@ -314,22 +560,30 @@ __device__ __forceinline__ void ReduceCopy128bMulti(const int w, const int nw, c
Pack128 vals[UNROLL];
// Load and reduce
for (int u = 0; u < UNROLL; ++u) Fetch128(vals[u], srcs[0]+u*WARP_SIZE);
if (preOpSrc0) {
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>().preOp(fn, vals[u]);
}
#pragma unroll
for (int i=1; i<MINSRCS; i++) {
Pack128 vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) Fetch128(vals2[u], srcs[i]+u*WARP_SIZE);
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(vals[u], vals2[u]);
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(fn, vals[u], vals2[u]);
}
#pragma unroll
for (int i=MINSRCS; i<MAXSRCS; i++) {
if (i<nsrcs) {
Pack128 vals2[UNROLL];
for (int u = 0; u < UNROLL; ++u) Fetch128(vals2[u], srcs[i]+u*WARP_SIZE);
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(vals[u], vals2[u]);
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(fn, vals[u], vals2[u]);
}
}
if (postOp) {
#pragma unroll
for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>().postOp(fn, vals[u]);
}
// Store
#pragma unroll
for (int i = 0; i < MINDSTS; i++) {
@@ -353,9 +607,9 @@ __device__ int ptrAlign128(T* ptr) { return (uint64_t)ptr % alignof(Pack128); }
#define PACKELEMS (sizeof(Pack128) / sizeof(T))
template<int UNROLL, class FUNC, typename T, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ __forceinline__ void ReduceOrCopyMulti(const int tid, const int nthreads,
int nsrcs, const T** srcs, int ndsts, T** dsts,
int N) {
__device__ __forceinline__ void ReduceOrCopyMulti(
const int tid, const int nthreads, FUNC fn, bool preOpSrc0, bool postOp, int nsrcs, const T** srcs, int ndsts, T** dsts, int N
) {
int Nrem = N;
if (Nrem <= 0) return;
@@ -381,7 +635,8 @@ __device__ __forceinline__ void ReduceOrCopyMulti(const int tid, const int nthre
int Npack = (Nrem / (PACKELEMS*UNROLL*WARP_SIZE)) * (UNROLL*WARP_SIZE); // round down
int Nelem = Npack * PACKELEMS;
ReduceCopy128bMulti<FUNC, T, UNROLL, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>(w, nw, t, nsrcs, srcs, ndsts, dsts, offset, Npack);
ReduceCopy128bMulti<FUNC, T, UNROLL, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
Nrem -= Nelem;
if (Nrem == 0) return;
@@ -391,7 +646,8 @@ __device__ __forceinline__ void ReduceOrCopyMulti(const int tid, const int nthre
Npack = Nrem / PACKELEMS;
Nelem = Npack * PACKELEMS;
ReduceCopy128bMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>(w, nw, t, nsrcs, srcs, ndsts, dsts, offset, Npack);
ReduceCopy128bMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Npack);
Nrem -= Nelem;
if (Nrem == 0) return;
@@ -401,14 +657,16 @@ __device__ __forceinline__ void ReduceOrCopyMulti(const int tid, const int nthre
// unrolled, by-type (mostly for unaligned buffers)
int Nelem = (Nrem / (UNROLL*PACKELEMS/2*WARP_SIZE)) * (UNROLL*PACKELEMS/2*WARP_SIZE); // round down
ReduceCopyMulti<FUNC, T, UNROLL*PACKELEMS/2, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>(w, nw, t, nsrcs, srcs, ndsts, dsts, offset, Nelem);
ReduceCopyMulti<FUNC, T, UNROLL*PACKELEMS/2, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Nelem);
Nrem -= Nelem;
if (Nrem == 0) return;
offset += Nelem;
// no unroll, by type. Should finish what's remaining.
ReduceCopyMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>(w, nw, t, nsrcs, srcs, ndsts, dsts, offset, Nrem);
ReduceCopyMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>
(w, nw, t, fn, preOpSrc0, postOp, nsrcs, srcs, ndsts, dsts, offset, Nrem);
}
#endif // COMMON_KERNEL_H_
src/collectives/device/functions.cu
+32 -3
View File
@@ -1,5 +1,5 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -8,7 +8,7 @@
#include "collectives.h"
#include "common.h"
__device__ struct ncclShmemData* ncclShmem;
__shared__ ncclShmemData ncclShmem;
#define NCCL_FUNC5(func, algo, redop, type) \
NCCL_FUNC_NAME(func, algo, LL, redop, type), \
@@ -20,6 +20,31 @@ __device__ struct ncclShmemData* ncclShmem;
NCCL_FUNC5(func, RING, redop, type), \
NCCL_FUNC5(func, COLLNET, redop, type)
#if defined(__CUDA_BF16_TYPES_EXIST__)
// Must be consistent with ncclDataType_t
#define NCCL_FUNCS3A(func, redop) \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, uint8_t), \
NCCL_FUNC4(func, redop, int32_t), \
NCCL_FUNC4(func, redop, uint32_t), \
NCCL_FUNC4(func, redop, int64_t), \
NCCL_FUNC4(func, redop, uint64_t), \
NCCL_FUNC4(func, redop, half), \
NCCL_FUNC4(func, redop, float), \
NCCL_FUNC4(func, redop, double), \
NCCL_FUNC4(func, redop, __nv_bfloat16)
#define NCCL_FUNCS3B(func, redop) \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t)
#else
// Must be consistent with ncclDataType_t
#define NCCL_FUNCS3A(func, redop) \
NCCL_FUNC4(func, redop, int8_t), \
@@ -41,17 +66,21 @@ __device__ struct ncclShmemData* ncclShmem;
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t), \
NCCL_FUNC4(func, redop, int8_t)
#endif
// Must be consistent with ncclRedOp_t
#define NCCL_FUNCS2A(func) \
NCCL_FUNCS3A(func, Sum ), \
NCCL_FUNCS3A(func, Prod), \
NCCL_FUNCS3A(func, Max ), \
NCCL_FUNCS3A(func, Min )
NCCL_FUNCS3A(func, Min ), \
NCCL_FUNCS3A(func, Avg)
#define NCCL_FUNCS2B(func) \
NCCL_FUNCS3B(func, Sum), \
NCCL_FUNCS3B(func, Sum), \
NCCL_FUNCS3B(func, Sum), \
NCCL_FUNCS3B(func, Sum), \
NCCL_FUNCS3B(func, Sum)
// Must be consistent with ncclFunc_t
src/collectives/device/gen_rules.sh
+10 -3
View File
@@ -1,19 +1,26 @@
#!/bin/bash
#
# Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.
#
# See LICENSE.txt for license information
#
dir=$1
datatypes="i8 u8 i32 u32 i64 u64 f16 f32 f64"
if [ "$CUDA_MAJOR" -ge 11 ]
then
datatypes+=" bf16"
fi
targets="GENOBJS := \\\\\n"
for base in sendrecv all_reduce all_gather broadcast reduce reduce_scatter; do
opn=0
for op in sum prod min max; do
for op in sum prod min max avg; do
dtn=0
for dt in i8 u8 i32 u32 i64 u64 f16 f32 f64; do
# Order must match that of the ncclDataType_t enum
for dt in ${datatypes}; do
echo "${dir}/${base}_${op}_${dt}.o : ${base}.cu ${dir}/${base}.dep"
echo " @printf \"Compiling %-35s > %s\\\\n\" ${base}.cu ${dir}/${base}_${op}_${dt}.o"
echo " mkdir -p ${dir}"
src/collectives/device/op128.h
+32
View File
@@ -33,4 +33,36 @@ inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_
:: "l"(v0), "l"(v1), "l"(shmemAsmPtr));
}
template<typename T>
inline __device__ void loadShmemMisaligned128(T *ptr, uint64_t &v0, uint64_t &v1) {
union {
uint32_t tmp4[4];
uint64_t tmp8[2];
};
if(sizeof(T) < 4) {
uint32_t *ptr4 = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(ptr) & -uintptr_t(4));
#pragma unroll
for(int e=0; e < 4; e++) {
// Produce 4 bytes of sub-register type by reading 2 4-byte
// aligned values and shifting.
uint32_t lo, hi;
asm("ld.shared.b32 %0,[%1];" : "=r"(lo) : "l"(ptr4+e+0));
asm("ld.shared.b32 %0,[%1];" : "=r"(hi) : "l"(ptr4+e+1));
tmp4[e] = __funnelshift_r(lo, hi, 8*(int(reinterpret_cast<uintptr_t>(ptr))%4));
}
}
else if(sizeof(T) == 4) {
#pragma unroll
for(int e=0; e < 4; e++)
asm("ld.shared.b32 %0,[%1];" : "=r"(tmp4[e]) : "l"(ptr+e));
}
else /*sizeof(T)==8*/ {
#pragma unroll
for(int e=0; e < 2; e++)
asm("ld.shared.b64 %0,[%1];" : "=l"(tmp8[e]) : "l"(ptr+e));
}
v0 = tmp8[0];
v1 = tmp8[1];
}
#endif
src/collectives/device/primitives.h
+120 -369
View File
@@ -11,381 +11,132 @@
#include "reduce_kernel.h" // for reduction funcs
#include "common.h"
#define SPINS_BEFORE_CHECK_ABORT 1000000
#define NCCL_SPINS_BEFORE_CHECK_ABORT 1000000
// Unroll unconditionally the first send/recv since nsend/nrecv should be at
// least 1 if SEND/RECV is set.
#define FOR_SEND(func, ...) do { \
if (SEND) { \
/* Send to far first, then close */ \
for (int i=1; i<NSEND && i<nsend; i++) func(i, ##__VA_ARGS__); \
func(0, ##__VA_ARGS__); \
} \
} while (0)
/* Protocol classes: ProtoSimple, ProtoLL, ProtoLL128
* We use these as template args to the Primtiives class instead of integral
* enums (e.g. NCCL_PROTO_LL) because for SIMPLE we need to carry a few extra
* numbers. Also these types hold methods which let us compute numbers important
* to how that protocol operates with a consistent interface so that our
* algorithm code can operate protocol parametrically.
*/
template<int SlicePerChunk_1, int StepPerSlice_1, int Unroll_1 = COLL_UNROLL>
struct ProtoSimple {
static constexpr int Id = NCCL_PROTO_SIMPLE;
static constexpr int SlicePerChunk = SlicePerChunk_1;
static constexpr int StepPerSlice = StepPerSlice_1;
static constexpr int Unroll = Unroll_1;
#define FOR_RECV(func, ...) do { \
if (RECV) { \
/* Recv from close first, then far */ \
func(0, ##__VA_ARGS__); \
for (int i=1; i<NRECV && i<nrecv; i++) func(i, ##__VA_ARGS__); \
} \
} while (0)
#define ROLE_SRC 0x01
#define ROLE_DST 0x02
#define ROLE_WAIT_RECV 0x04
#define ROLE_WAIT_SEND 0x08
#define ROLE_POST_SEND 0x10
#define ROLE_POST_RECV 0x20
// Implementation of primitive types
template <int UNROLL, int SLICESPERCHUNK, int SLICESTEPS, typename T, int NRECV, int NSEND, int DIRECT, class FUNC>
class ncclPrimitives {
private:
const int tid;
int nthreads;
int nworkers;
const int stepSize;
int nrecv = 0;
int nsend = 0;
struct ncclConnInfo* conn = NULL;
volatile int* connSizesFifoPtr = NULL;
void** connPtrsFifoPtr = NULL;
volatile uint64_t* connHeadPtr = NULL;
volatile uint64_t* connTailPtr = NULL;
uint64_t connTailCache; // Cache last seen value
uint64_t connHeadCache; // Cache last seen value
int index; // Peer index I'm responsible for
int peer = -1;
int role = 0;
int group;
uint64_t step;
T* direct = NULL;
T* buff;
struct ncclDevComm* comm;
const T** srcs;
T** dsts;
// Don't use barrier 0 as it's used by the final sync
inline __device__ void barrier() {
if (nthreads == WARP_SIZE) __syncwarp();
else asm volatile ("bar.sync %0, %1;" :: "r"(group+1), "r"(nthreads));
// Data bytes (no flags etc) in one step of the fifo queue.
__device__ static int calcBytePerStep() {
return ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS;
}
inline __device__ void subBarrier() {
if (nworkers == nthreads) barrier();
else asm volatile ("bar.sync %0, %1;" :: "r"(group+2), "r"(nworkers));
// Granularity of data bytes transferred per thread.
__device__ static int calcBytePerGrain() {
return sizeof(uint64_t); // Bogus value? Nobody queries this metric for simple.
}
uint32_t spins = 0;
uint32_t abort = 0;
inline __device__ int checkAbort() {
spins++;
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = *(comm->abortFlag);
spins = 0;
}
return abort;
}
template <int DIRECTPTR>
inline __device__ T* directPtr(ssize_t directOffset) {
return DIRECTPTR && direct ? direct+directOffset : buff+(step%NCCL_STEPS)*stepSize;
}
template <int DST, int DIRECTSEND>
inline __device__ void waitSend(ssize_t directOffset, int nbytes) {
spins = 0;
while (connHeadCache + NCCL_STEPS < step + SLICESTEPS) {
connHeadCache = *connHeadPtr;
if (checkAbort()) break;
}
if (connSizesFifoPtr) {
connSizesFifoPtr[step%NCCL_STEPS] = nbytes;
}
if (connPtrsFifoPtr) loadPtr(connPtrsFifoPtr+step%NCCL_STEPS, dsts[DST+index]);
else dsts[DST+index] = directPtr<DIRECTSEND>(directOffset);
step += SLICESTEPS;
}
template <int SRC, int DIRECTRECV>
inline __device__ void waitRecv(ssize_t directOffset) {
spins = 0;
while (connTailCache < step + SLICESTEPS) {
connTailCache = *connTailPtr;
if (checkAbort()) break;
}
if (connPtrsFifoPtr) loadPtr(connPtrsFifoPtr+step%NCCL_STEPS, srcs[SRC+index]);
else srcs[SRC+index] = directPtr<DIRECTRECV>(directOffset);
step += SLICESTEPS;
}
inline __device__ void postRecv() {
*connHeadPtr = step += SLICESTEPS;
}
inline __device__ void postSend() {
*connTailPtr = step += SLICESTEPS;
}
template <int DIRECTRECV, int DIRECTSEND, int RECV, int SEND, int SRC, int DST>
inline __device__ void
GenericOp(const T* srcPtr, T* dstPtr, int nelem, ssize_t directOffset) {
int offset = 0;
int sliceSize = stepSize*SLICESTEPS;
int dataSize = max(DIVUP(nelem, 16*SLICESPERCHUNK)*16, sliceSize/32);
#pragma unroll
for (int slice=0; slice<SLICESPERCHUNK; ++slice) {
int realSize = max(0, min(dataSize, nelem-offset));
if (tid < nworkers) {
if (SRC && (role & ROLE_SRC)) srcs[0] = srcPtr+offset;
if (RECV && (role & ROLE_WAIT_RECV)) waitRecv<SRC, DIRECTRECV>(directOffset+offset);
if (DST && (role & ROLE_DST)) dsts[0] = dstPtr+offset;
if (SEND && (role & ROLE_WAIT_SEND)) waitSend<DST, DIRECTSEND>(directOffset+offset, realSize*sizeof(T));
if (realSize > 0) {
subBarrier();
if (DIRECTRECV && srcs[0] == dsts[0]) {
// We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
if (SEND) {
// (1-SEND) is only there to avoid compilation errors in case NSEND=0 (and SEND=0).
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, (1-SEND)+NSEND>(tid, nworkers, 1, srcs, nsend, dsts+1, realSize);
}
} else {
ReduceOrCopyMulti<UNROLL, FUNC, T, RECV+SRC, RECV*NRECV+SRC, SEND+DST, SEND*NSEND+DST>(tid, nworkers, RECV*nrecv+SRC, srcs, SEND*nsend+DST, dsts, realSize);
}
}
}
barrier();
if (SEND && (role & ROLE_POST_SEND) && realSize > 0 && index == 0) __threadfence_system();
__syncwarp();
if (SEND && (role & ROLE_POST_SEND)) postSend();
if (RECV && (role & ROLE_POST_RECV)) postRecv();
offset += realSize;
}
}
// Scatter and gather do not support DIRECT
template <int RECV, int SEND>
inline __device__ void
ScatterGatherOp(const T* srcPtr, T* dstPtr, int totalElem, int peerElem, int skip, int shift) {
int offset = 0; // slice offset
int sliceSize = stepSize*SLICESTEPS;
int dataSize = max(DIVUP(peerElem, 16*SLICESPERCHUNK)*16, sliceSize/32); // per-peer slice size
#pragma unroll
for (int slice=0; slice<SLICESPERCHUNK; ++slice) {
int realSize = max(0, min(dataSize, peerElem-offset));
if (tid < nworkers) {
if (RECV && (role & ROLE_WAIT_RECV)) waitRecv<0, 0>(0);
// realSize is not accurate here; but intra-node does not rely on sizes FIFO
if (SEND && (role & ROLE_WAIT_SEND)) waitSend<0, 0>(0, realSize*sizeof(T));
subBarrier();
if (SEND) {
#pragma unroll
for (int j=0; j<nsend; j++) {
int i = (j+shift)%nsend;
int peerOffset = i*peerElem + offset;
if (skip >=0 && i >= skip) peerOffset += peerElem;
const T* src0 = srcPtr + peerOffset;
int realPeerSize = min(realSize, totalElem-peerOffset);
if (realPeerSize > 0) ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nworkers, 1, &src0, 1, dsts+i, realPeerSize);
}
} else if (RECV) {
#pragma unroll
for (int j=0; j<nrecv; j++) {
int i = (j+shift)%nrecv;
int peerOffset = i*peerElem + offset;
if (skip >= 0 && i >= skip) peerOffset += peerElem;
T* dst0 = dstPtr + peerOffset;
int realPeerSize = min(realSize, totalElem-peerOffset);
if (realPeerSize > 0) ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nworkers, 1, srcs+i, 1, &dst0, realPeerSize);
}
}
}
barrier();
if (SEND && (role & ROLE_POST_SEND) && realSize > 0 && index == 0) __threadfence_system();
__syncwarp();
if (SEND && (role & ROLE_POST_SEND)) postSend();
if (RECV && (role & ROLE_POST_RECV)) postRecv();
offset += realSize;
}
}
__device__ __forceinline__ void loadRecvConn(struct ncclChannel* channel, T* directBuff) {
if (role & (ROLE_WAIT_RECV|ROLE_POST_RECV)) {
// For oneshot: groups 0,2 use conn 0, groups 4,6 use conn 1
const int connIndex = (NSEND == NCCL_MAX_DIRECT_ARITY || NRECV == NCCL_MAX_DIRECT_ARITY) ? group/4 : 0;
conn = &channel->devPeers[peer].recv[connIndex].conn;
step = conn->step;
step = ROUNDUP(step, SLICESPERCHUNK*SLICESTEPS);
if (role & ROLE_POST_RECV) {
connHeadPtr = conn->head;
// Return credits in case we rounded up.
*connHeadPtr = step;
}
if (role & ROLE_WAIT_RECV) {
buff = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (DIRECT && (conn->direct & NCCL_DIRECT_GPU)) {
direct = directBuff;
*conn->ptrExchange = directBuff;
}
connTailPtr = conn->tail;
connTailCache = *connTailPtr;
connPtrsFifoPtr = conn->ptrsFifo;
}
}
}
__device__ __forceinline__ void loadSendConn(struct ncclChannel* channel) {
if (role & (ROLE_WAIT_SEND|ROLE_POST_SEND)) {
// For oneshot: groups 0,2 use conn 0, groups 4,6 use conn 1
const int connIndex = (NSEND == NCCL_MAX_DIRECT_ARITY || NRECV == NCCL_MAX_DIRECT_ARITY) ? group/4 : 0;
conn = &channel->devPeers[peer].send[connIndex].conn;
step = conn->step;
step = ROUNDUP(step, SLICESPERCHUNK*SLICESTEPS);
if (role & ROLE_POST_SEND) {
connTailPtr = conn->tail;
}
if (role & ROLE_WAIT_SEND) {
buff = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (DIRECT && (conn->direct & NCCL_DIRECT_GPU)) {
void* volatile* ptr = conn->ptrExchange;
while ((direct = (T*)(*ptr)) == NULL) { if (checkAbort()) break; }
*ptr = NULL;
}
connHeadPtr = conn->head;
connHeadCache = *connHeadPtr;
connSizesFifoPtr = conn->sizesFifo;
connPtrsFifoPtr = conn->ptrsFifo;
}
}
}
__device__ __forceinline__ void saveSync() {
if (role & (ROLE_POST_SEND|ROLE_POST_RECV)) {
conn->step = step;
__threadfence_system();
}
}
public:
__device__ __forceinline__
ncclPrimitives(const int tid, const int nworkers, int* recvPeers, int* sendPeers, T* directBuff, int stepSize, struct ncclChannel* channel, struct ncclDevComm* comm, struct ncclShmemPtrs* ptrs, int group)
: comm(comm), tid(tid), nworkers(nworkers), stepSize(stepSize), srcs((const T**)ptrs[group].srcs), dsts((T**)ptrs[group].dsts), group(group) {
nthreads = nworkers;
// For send operations, we need an extra warp to overlap the threadfence and the copy
int postThreads = NSEND && nworkers >= 64 ? WARP_SIZE : 0;
nthreads += postThreads;
// Make sure step is updated before we read it.
barrier();
for (int i=0; i<NRECV; i++) if (recvPeers[i] != -1) nrecv++;
for (int i=0; i<NSEND; i++) if (sendPeers[i] != -1) nsend++;
#define SYNC_GROUP 8
static_assert(NSEND < SYNC_GROUP && NRECV < SYNC_GROUP, "Not enough threads to cover all peers");
int g = tid / SYNC_GROUP;
int ng = nthreads / SYNC_GROUP;
index = tid % SYNC_GROUP;
if (g == 0) {
if (index < nrecv) role |= ROLE_WAIT_RECV;
if (index == nrecv) role |= ROLE_SRC;
} else if (g == 1) {
if (index < nsend) role |= ROLE_WAIT_SEND;
if (index == nsend) role |= ROLE_DST;
} else if (g == ng - 2) {
if (index < nrecv) role |= ROLE_POST_RECV;
} else if (g == ng - 1) {
if (index < nsend) role |= ROLE_POST_SEND;
}
if (role & (ROLE_WAIT_RECV|ROLE_POST_RECV)) peer = recvPeers[index];
if (role & (ROLE_WAIT_SEND|ROLE_POST_SEND)) peer = sendPeers[index];
loadRecvConn(channel, directBuff);
loadSendConn(channel);
}
__device__ __forceinline__ void
send(const T* src, int nelem) {
GenericOp<0, 0, 0, 1, 1, 0>(src, NULL, nelem, 0);
}
__device__ __forceinline__ void
directSend(const T* src, ssize_t directOffset, int nelem) {
GenericOp<0, 1, 0, 1, 1, 0>(src, NULL, nelem, directOffset);
}
__device__ __forceinline__ void
recv(T* dst, int nelem) {
GenericOp<0, 0, 1, 0, 0, 1>(NULL, dst, nelem, 0);
}
__device__ __forceinline__ void
directRecv(T* dst, ssize_t directOffset, int nelem) {
GenericOp<1, 0, 1, 0, 0, 1>(NULL, dst, nelem, directOffset);
}
__device__ __forceinline__ void
copySend(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 0, 1, 1, 1>(src, dst, nelem, 0);
}
__device__ __forceinline__ void
directCopySend(const T* src, T* dst, ssize_t directOffset, int nelem) {
GenericOp<0, 1, 0, 1, 1, 1>(src, dst, nelem, directOffset);
}
__device__ __forceinline__ void
recvCopySend(T* dst, int nelem) {
GenericOp<0, 0, 1, 1, 0, 1>(NULL, dst, nelem, 0);
}
__device__ __forceinline__ void
directRecvCopySend(T* dst, ssize_t directOffset, int nelem) {
GenericOp<1, 1, 1, 1, 0, 1>(NULL, dst, nelem, directOffset);
}
__device__ __forceinline__ void
recvReduceCopy(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 1, 0, 1, 1>(src, dst, nelem, 0);
}
__device__ __forceinline__ void
recvReduceSend(const T* src, int nelem) {
GenericOp<0, 0, 1, 1, 1, 0>(src, NULL, nelem, 0);
}
__device__ __forceinline__ void
recvReduceCopySend(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 1, 1, 1, 1>(src, dst, nelem, 0);
}
__device__ __forceinline__ void
directRecvReduceCopySend(const T* src, T* dst, ssize_t directOffset, int nelem) {
// Direct is only for the send part
GenericOp<0, 1, 1, 1, 1, 1>(src, dst, nelem, directOffset);
}
__device__ __forceinline__ void
scatter(const T* src, int totalElem, int peerElem, int skip, int shift) {
ScatterGatherOp<0, 1>(src, NULL, totalElem, peerElem, skip, shift);
}
__device__ __forceinline__ void
gather(T* dst, int totalElem, int peerElem, int skip, int shift) {
ScatterGatherOp<1, 0>(NULL, dst, totalElem, peerElem, skip, shift);
}
__device__ __forceinline__ ~ncclPrimitives() {
// Save steps for the next operation
saveSync();
// Group width is how many consecutive group values a subchannel occupies.
static constexpr int MaxGroupWidth = 2;
__device__ static int calcGroupWidth(bool send, int nthreads) {
return send && nthreads-WARP_SIZE >= 64 ? 2 : 1;
}
};
#include "prims_ll.h"
//#include "prims_ll128.h"
struct ProtoLL {
static constexpr int Id = NCCL_PROTO_LL;
// Data bytes (no flags etc) in one step of the fifo queue.
__device__ static int calcBytePerStep() {
return ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/2; // Half is data
}
// Granularity of data bytes transferred per thread.
__device__ static int calcBytePerGrain() {
return sizeof(uint64_t); // One 16-byte line has 8-bytes of data
}
// Group width is how many consecutive group values a subchannel occupies.
static constexpr int MaxGroupWidth = 1;
__device__ static int calcGroupWidth(bool send, int nthreads) {
return 1;
}
};
struct ProtoLL128 {
static constexpr int Id = NCCL_PROTO_LL128;
// Data bytes (no flags etc) in one step of the fifo queue.
__device__ static int calcBytePerStep() {
return (ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS)*NCCL_LL128_DATAELEMS/NCCL_LL128_LINEELEMS;
}
// Granularity of data bytes transferred per thread.
__device__ static int calcBytePerGrain() {
return NCCL_LL128_SHMEM_ELEMS_PER_THREAD*NCCL_LL128_DATAELEMS*sizeof(uint64_t)/NCCL_LL128_LINEELEMS;
}
// Group width is how many consecutive group values a subchannel occupies.
static constexpr int MaxGroupWidth = 1;
__device__ static int calcGroupWidth(bool send, int nthreads) {
return 1;
}
};
/* Fan (as in fan-in & fan-out) classes hold recv and send counts. The template
* arguments are static bounds on the maximum values. Asymmetric counts are
* independent. Symmetric is a static guarantee that nrecv==nsend, so it only
* stores one value at runtime. This optimization save 32-bit register, but more
* importantly uses fewer predicate registers when unrolling loops.
*/
template<int MaxRecv_, int MaxSend_>
struct FanAsymmetric {
static constexpr int MaxRecv = MaxRecv_, MaxSend = MaxSend_;
int nr, ns;
FanAsymmetric() = default;
__device__ FanAsymmetric(int nrecv, int nsend): nr(nrecv), ns(nsend) {
// assert(nrecv <= MaxRecv && nsend <= MaxSend);
}
__device__ int nrecv() const { return MaxRecv ? nr : 0; }
__device__ int nsend() const { return MaxSend ? ns : 0; }
};
template<int MaxArity>
struct FanSymmetric {
static constexpr int MaxRecv = MaxArity, MaxSend = MaxArity;
int n;
FanSymmetric() = default;
__device__ FanSymmetric(int nrecv, int nsend): n(nrecv) {
// assert(nrecv == nsend && nrecv <= MaxArity);
}
__device__ int nrecv() const { return n; }
__device__ int nsend() const { return n; }
};
// The primitives class. Specialized per protocol in the other headers.
template<typename T, typename RedOp, typename Fan, int Direct, typename Proto>
class Primitives;
// Used by LL & LL128 to implement direct members in the naive way.
template<typename RealPrimitives>
struct PrimitivesWithoutDirect {
__device__ void directSend(intptr_t inpIx, intptr_t remoteOutIx, int eltN) {
static_cast<RealPrimitives*>(this)->send(inpIx, eltN);
}
__device__ void directSendFromOutput(intptr_t outIx, intptr_t remoteOutIx, int eltN) {
static_cast<RealPrimitives*>(this)->sendFromOutput(outIx, eltN);
}
__device__ void directRecv(intptr_t outIx, int eltN) {
static_cast<RealPrimitives*>(this)->recv(outIx, eltN, /*postOp=*/false);
}
__device__ void directCopySend(intptr_t inpIx, intptr_t outIx, intptr_t remoteOutIx, int eltN, bool postOp=false) {
static_cast<RealPrimitives*>(this)->copySend(inpIx, outIx, eltN, postOp);
}
__device__ void directRecvCopySend(intptr_t outIx, intptr_t remoteOutIx, int eltN) {
static_cast<RealPrimitives*>(this)->recvCopySend(outIx, eltN, /*postOp=*/false);
}
__device__ void directRecvReduceCopySend(intptr_t inpIx, intptr_t outIx, intptr_t remoteOutIx, int eltN, bool postOp=false) {
// Direct is only for the send part
static_cast<RealPrimitives*>(this)->recvReduceCopySend(inpIx, outIx, eltN, postOp);
}
};
#include "prims_simple.h"
#include "prims_ll.h"
#include "prims_ll128.h"
#endif
src/collectives/device/prims_ll.h
+258 -118
View File
@@ -4,15 +4,20 @@
* See LICENSE.txt for license information
************************************************************************/
template <typename T, class FUNC, int NRECV, int NSEND>
class ncclLLPrimitives {
private:
template<typename T, typename RedOp, typename Fan, int Direct>
class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL>> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
RedOp redOp;
const int tid;
const int nthreads;
const int wid;
const int group;
const int stepLines;
int nrecv = 0;
int nsend = 0;
Fan fan;
T *userBufs[2];
struct ncclConnInfo* recvConn = NULL;
volatile uint64_t* recvConnHeadPtr = NULL;
uint64_t recvConnHead;
@@ -23,11 +28,10 @@ class ncclLLPrimitives {
uint64_t sendConnHead;
uint64_t sendConnHeadCache; // Cache last seen value
uint64_t recvStep[NRECV];
uint64_t sendStep[NSEND];
union ncclLLFifoLine* recvBuff[NRECV];
union ncclLLFifoLine* sendBuff[NSEND];
struct ncclDevComm* comm;
uint64_t recvStep[MaxRecv];
uint64_t sendStep[MaxSend];
union ncclLLFifoLine* recvBuff[MaxRecv];
union ncclLLFifoLine* sendBuff[MaxSend];
inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepLines; }
inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepLines; }
@@ -37,27 +41,26 @@ class ncclLLPrimitives {
inline __device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(sendStep[i]+1); }
inline __device__ void barrier() {
asm volatile ("bar.sync 1, %0;" :: "r"(nthreads));
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(1+group));
}
uint32_t spins = 0;
uint32_t abort = 0;
inline __device__ int checkAbort(int i, int send) {
inline __device__ int checkAbort(int &spins, int send) {
spins++;
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = *(comm->abortFlag);
if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
abort = *ncclShmem.comm.abortFlag;
spins = 0;
}
return abort;
}
inline __device__ void waitSend(int nbytes) {
spins = 0;
if (sendConnHeadPtr) {
int spins = 0;
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
sendConnHeadCache = *sendConnHeadPtr;
if (checkAbort(wid, 1)) break;
if (checkAbort(spins, 1)) break;
}
if (sendConnFifoPtr) {
int size = ((sendConnHead & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? stepLines*sizeof(union ncclLLFifoLine) : nbytes;
@@ -85,83 +88,212 @@ class ncclLLPrimitives {
sendStep[i]++;
}
__device__ uint64_t readLL(int i, int offset) {
__device__ uint64_t readLL(int offset, int i) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
uint32_t flag = recvFlag(i);
uint32_t data1, flag1, data2, flag2;
spins = 0;
int spins = 0;
do {
asm volatile("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(data1), "=r"(flag1), "=r"(data2), "=r"(flag2) : "l"(&src->i4));
if (checkAbort(i, 0)) break;
asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(data1), "=r"(flag1), "=r"(data2), "=r"(flag2) : "l"(&src->i4));
if (checkAbort(spins, 0)) break;
} while ((flag1 != flag) || (flag2 != flag));
uint64_t val64 = data1 + (((uint64_t)data2) << 32);
return val64;
}
template<int BeginIx>
__device__ void readLLBeginAll(int offset, ncclLLFifoLine(&line)[MaxRecv]) {
#pragma unroll
for (int i=BeginIx; i < MaxRecv; i++) {
if (i < fan.nrecv()) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
}
}
}
__device__ uint64_t readLLFinish(int offset, ncclLLFifoLine(&line)[MaxRecv], int i) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
uint32_t flag = recvFlag(i);
int spins = 0;
while (line[i].flag1 != flag || line[i].flag2 != flag) {
asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
if (checkAbort(spins, 0)) break;
}
uint64_t val64 = line[i].data1 + (((uint64_t)line[i].data2) << 32);
return val64;
}
__device__ void storeLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
asm volatile("st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};" :: "l"(&dst->i4), "r"((uint32_t)val), "r"(flag), "r"((uint32_t)(val >> 32)), "r"(flag));
}
// Using memcpy handles misaligned pointers.
__device__ uint64_t readAL(uint64_t* src) {
uint64_t val;
memcpy((char*)&val, (char*)src, sizeof(uint64_t));
return val;
static constexpr int EltPerLine = sizeof(uint64_t)/sizeof(T);
template<typename U>
__device__ static U load(U *src) {
union {
U elt;
uint16_t u2;
uint32_t u4;
uint64_t u8;
};
if(sizeof(U) == 1)
asm("ld.volatile.global.b8 %0,[%1];" : "=r"(u4) : "l"(src));
else if(sizeof(U) == 2)
asm("ld.volatile.global.b16 %0,[%1];" : "=h"(u2) : "l"(src));
else if(sizeof(U) == 4)
asm("ld.volatile.global.b32 %0,[%1];" : "=r"(u4) : "l"(src));
else
asm("ld.volatile.global.b64 %0,[%1];" : "=l"(u8) : "l"(src));
return elt;
}
__device__ void storeAL(uint64_t* dst, uint64_t val, uint32_t nbytes) {
memcpy((char*)dst, (char*)&val, nbytes);
template<typename U>
__device__ static void store(U *dst, U val) {
union {
U elt;
uint16_t u2;
uint32_t u4;
uint64_t u8;
};
elt = val;
if(sizeof(U) == 1)
asm("st.volatile.global.b8 [%0],%1;" :: "l"(dst), "r"(u4));
else if(sizeof(U) == 2)
asm("st.volatile.global.b16 [%0],%1;" :: "l"(dst), "h"(u2));
else if(sizeof(U) == 4)
asm("st.volatile.global.b32 [%0],%1;" :: "l"(dst), "r"(u4));
else
asm("st.volatile.global.b64 [%0],%1;" :: "l"(dst), "l"(u8));
}
template <int RECV, int SEND, int SRC, int DST>
__device__ void LLGenericOp(const T* srcPtr, T* dstPtr, int nelem) {
uint32_t nbytes = nelem < 0 ? 0 : nelem*sizeof(T);
uint32_t npack = DIVUP(nbytes, sizeof(uint64_t));
uint64_t* srcPack = (uint64_t*)srcPtr;
uint64_t* dstPack = (uint64_t*)dstPtr;
int offset = tid;
struct DataLoader {
int misalign;
union {
uint32_t u4[sizeof(T) <= 2 ? 3 : 2];
uint64_t u8;
T elt[EltPerLine];
};
// Always waitSend in case of cleanup
if (SEND) waitSend(npack*sizeof(union ncclLLFifoLine));
// Do multiples of 64 bits
#pragma unroll 2
for (; offset<npack; offset+=nthreads) {
// Recv : local, then intra-node, then inter-node
uint64_t val = SRC ? readAL(srcPack+offset) : readLL(0, offset);
if (RECV) {
if (SRC) val = MULTI<FUNC, T>()(readLL(0, offset), val);
for (int i=1; i<NRECV && i<nrecv; i++) {
val = MULTI<FUNC, T>()(readLL(i, offset), val);
}
__device__ void loadBegin(T *src, int eltN) {
if (sizeof(T) <= 2) {
misalign = reinterpret_cast<uintptr_t>(src)%4;
uint32_t *p = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(src) & -uintptr_t(4));
u4[0] = load(p+0);
u4[1] = misalign + eltN*sizeof(T) > 4 ? load(p+1) : 0;
// u4[2] would be simpler, but that throws warnings on some compilers
u4[sizeof(T) <= 2 ? 2 : 0] = misalign + eltN*sizeof(T) > 8 ? load(p+2) : 0;
}
// Send : inter-node, then intra-node, then local
if (SEND) {
for (int i=1; i<NSEND && i<nsend; i++) storeLL(sendPtr(i)+offset, val, sendFlag(i));
storeLL(sendPtr(0)+offset, val, sendFlag(0));
}
if (DST) {
if (((offset*sizeof(uint64_t)) ^ nbytes) < sizeof(uint64_t)) {
// Last incomplete word
storeAL(dstPack+offset, val, nbytes & 0x7);
} else {
storeAL(dstPack+offset, val, sizeof(uint64_t));
else {
#pragma unroll
for(int i=0; i < EltPerLine; i++) {
if(i==0 || i < eltN)
elt[i] = load(src + i);
}
}
}
FOR_RECV(incRecv); if (RECV) postRecv();
FOR_SEND(incSend, offset);
__device__ uint64_t loadFinish() {
if (sizeof(T) <= 2) {
u4[0] = __funnelshift_r(u4[0], u4[1], 8*misalign);
// u4[2] would be simpler, but that throws warnings on some compilers
u4[1] = __funnelshift_r(u4[1], u4[sizeof(T) <= 2 ? 2 : 0], 8*misalign);
}
return u8;
}
};
__device__ void storeData(T *dst, uint64_t val, int eltN) {
union {
uint64_t u8;
T elt[EltPerLine];
};
u8 = val;
#pragma unroll
for(int i=0; i < EltPerLine; i++) {
if (i==0 || i < eltN)
//store(dst+i, elt[i]);
dst[i] = elt[i];
}
}
template <int RECV, int SEND, int SrcBuf, int DstBuf>
__device__ void LLGenericOp(intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp) {
constexpr int SRC = SrcBuf != -1 ? 1 : 0;
constexpr int DST = DstBuf != -1 ? 1 : 0;
T *srcElts = SrcBuf == -1 ? nullptr : userBufs[SrcBuf] + srcIx;
T *dstElts = DstBuf == -1 ? nullptr : userBufs[DstBuf] + dstIx;
// Always waitSend in case of cleanup
nelem = nelem < 0 ? 0 : nelem;
if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine));
nelem -= tid*EltPerLine;
srcElts += tid*EltPerLine;
dstElts += tid*EltPerLine;
int offset = tid;
int eltPerTrip = nthreads*EltPerLine;
while (nelem > 0) {
int eltInLine = EltPerLine < nelem ? EltPerLine : nelem;
DataLoader dl;
ncclLLFifoLine line[MaxRecv];
uint64_t data, peerData;
if (SRC) {
dl.loadBegin(srcElts, eltInLine);
srcElts += eltPerTrip;
}
if (RECV) {
readLLBeginAll<1>(offset, line);
peerData = readLL(offset, 0);
}
if (SRC) {
data = dl.loadFinish();
if (SrcBuf == Input) data = MULTI<RedOp, T>().preOp(redOp, data);
}
if (RECV) {
data = !SRC ? peerData : MULTI<RedOp,T>()(redOp, peerData, data);
#pragma unroll MaxRecv
for (int i=1; i < MaxRecv && i < fan.nrecv(); i++) {
peerData = readLLFinish(offset, line, i);
data = MULTI<RedOp,T>()(redOp, peerData, data);
}
}
if (postOp) data = MULTI<RedOp, T>().postOp(redOp, data);
// Send : inter-node, then intra-node, then local
if (SEND) {
for (int i=1; i < MaxSend && i < fan.nsend(); i++)
storeLL(sendPtr(i)+offset, data, sendFlag(i));
storeLL(sendPtr(0)+offset, data, sendFlag(0));
}
if (DST) {
storeData(dstElts, data, eltInLine);
dstElts += eltPerTrip;
}
nelem -= eltPerTrip;
offset += nthreads;
}
if (RECV) {
for (int i=0; i < MaxRecv; i++) incRecv(i);
postRecv();
}
if (SEND) {
for (int i=1; i < MaxSend && i < fan.nsend(); i++)
incSend(i, offset);
incSend(0, offset);
}
}
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
recvBuff[i] = (union ncclLLFifoLine*)conn->buffs[NCCL_PROTO_LL];
recvStep[i] = conn->step;
if (wid == i) recvConn = conn;
nrecv++;
}
__device__ __forceinline__ void loadRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv()) {
recvConnHeadPtr = recvConn->head;
recvConnHead = recvConn->step;
}
@@ -171,10 +303,9 @@ class ncclLLPrimitives {
sendBuff[i] = (union ncclLLFifoLine*)conn->buffs[NCCL_PROTO_LL];
sendStep[i] = conn->step;
if (wid == i) sendConn = conn;
nsend++;
}
__device__ __forceinline__ void loadSendSync() {
if (tid < nsend) {
if (tid < fan.nsend()) {
sendConnHeadPtr = sendConn->head;
sendConnHeadCache = *sendConnHeadPtr;
sendConnHead = sendConn->step;
@@ -182,65 +313,74 @@ class ncclLLPrimitives {
}
}
__device__ __forceinline__ void saveRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
recvConn->step = recvConnHead;
__threadfence_block();
}
}
__device__ __forceinline__ void saveSendSync() {
if (tid < nsend) {
sendConn->step = sendConnHead;
__threadfence_block();
}
}
public:
__device__ __forceinline__
ncclLLPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, int stepLines, struct ncclChannel* channel, struct ncclDevComm* comm)
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), stepLines(stepLines) {
// Make sure step is updated before we read it.
barrier();
__device__ Primitives(
const int tid, const int nthreads, int const *recvPeers, int const *sendPeers,
void const *inputBuf, void *outputBuf, int group=0
):
redOp(FuncTraits<RedOp>().make(ncclShmem.comm.nRanks)),
tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), group(group),
stepLines(ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/sizeof(ncclLLFifoLine)) {
auto *channel = &ncclShmem.channel;
// If we are going to support oneshot collNet + LL, then we would need to add connector index here
for (int i=0; i<NRECV && recvPeers[i] >= 0; i++) loadRecvConn(&channel->devPeers[recvPeers[i]].recv->conn, i);
for (int i=0; i<NSEND && sendPeers[i] >= 0; i++) loadSendConn(&channel->devPeers[sendPeers[i]].send->conn, i);
int nrecv=0, nsend=0;
while (nrecv < MaxRecv && recvPeers[nrecv] >= 0) {
loadRecvConn(&channel->devPeers[recvPeers[nrecv]].recv->conn, nrecv);
nrecv++;
}
while (nsend < MaxSend && sendPeers[nsend] >= 0) {
loadSendConn(&channel->devPeers[sendPeers[nsend]].send->conn, nsend);
nsend++;
}
this->fan = Fan(nrecv, nsend);
loadRecvSync();
loadSendSync();
setDataPtrs(inputBuf, outputBuf);
}
__device__ void send(const T* src, int nelem) {
return LLGenericOp<0, 1, 1, 0>(src, NULL, nelem);
}
__device__ void recv(T* dst, int nelem) {
return LLGenericOp<1, 0, 0, 1>(NULL, dst, nelem);
}
__device__ void recvReduceSend(const T* src, int nelem) {
return LLGenericOp<1, 1, 1, 0>(src, NULL, nelem);
}
__device__ void recvReduceCopy(const T* src, T* dst, int nelem) {
return LLGenericOp<1, 0, 1, 1>(src, dst, nelem);
}
__device__ void copySend(const T* src, T* dst, int nelem) {
return LLGenericOp<0, 1, 1, 1>(src, dst, nelem);
}
__device__ void recvCopySend(T* dst, int nelem) {
return LLGenericOp<1, 1, 0, 1>(NULL, dst, nelem);
}
__device__ void recvReduceCopySend(const T* src, T* dst, int nelem) {
return LLGenericOp<1, 1, 1, 1>(src, dst, nelem);
}
__device__ __forceinline__ ~ncclLLPrimitives() {
__device__ ~Primitives() {
// Save steps for the next operation
saveRecvSync();
saveSendSync();
if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv())
recvConn->step = recvConnHead;
if (tid < fan.nsend())
sendConn->step = sendConnHead;
// Ensure all steps written back
barrier();
}
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
userBufs[Input] = (T*)inputBuf;
userBufs[Output] = (T*)outputBuf;
}
__device__ void moveDataPtrs(intptr_t delta) {
userBufs[Input] += delta;
userBufs[Output] += delta;
}
__device__ void send(intptr_t inpIx, int eltN) {
return LLGenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
}
__device__ void sendFromOutput(intptr_t outIx, int eltN) {
return LLGenericOp<0, 1, Output, -1>(outIx, -1, eltN, false);
}
__device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
return LLGenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
}
__device__ void recvReduceSend(intptr_t inpIx, int eltN) {
return LLGenericOp<1, 1, Input, -1>(inpIx, -1, eltN, false);
}
__device__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return LLGenericOp<1, 0, Input, Output>(inpIx, outIx, eltN, postOp);
}
__device__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return LLGenericOp<0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
}
__device__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
return LLGenericOp<1, 1, -1, Output>(-1, outIx, eltN, postOp);
}
__device__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return LLGenericOp<1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
}
};
src/collectives/device/prims_ll128.h
+250 -217
View File
@@ -8,17 +8,22 @@
#define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS-1)
template <typename T, class FUNC, int NRECV, int NSEND>
class ncclLL128Primitives {
private:
template<typename T, typename RedOp, typename Fan, int Direct>
class Primitives<T, RedOp, Fan, Direct, ProtoLL128>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128>> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
RedOp redOp;
const int tid;
const int nthreads;
const int wid;
const int stepSize;
const int warp;
const bool flagThread;
int nrecv = 0;
int nsend = 0;
const int group;
Fan fan;
T *userBufs[2];
struct ncclConnInfo* recvConn = NULL;
volatile uint64_t* recvConnHeadPtr = NULL;
uint64_t recvConnHead;
@@ -31,13 +36,10 @@ class ncclLL128Primitives {
uint64_t sendConnHead;
uint64_t sendConnHeadCache; // Cache last seen value
uint64_t recvStep[NRECV];
uint64_t sendStep[NSEND];
uint64_t* recvBuff[NRECV];
uint64_t* sendBuff[NSEND];
struct ncclDevComm* comm;
volatile uint64_t* shmem;
uint64_t recvStep[MaxRecv];
uint64_t sendStep[MaxSend];
uint64_t* recvBuff[MaxRecv];
uint64_t* sendBuff[MaxSend];
inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepSize; }
inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepSize; }
@@ -47,31 +49,26 @@ class ncclLL128Primitives {
inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; }
inline __device__ void barrier() {
if (NSEND>NRECV) {
asm volatile ("bar.sync 1, %0;" :: "r"(nthreads));
} else {
asm volatile ("bar.sync 2, %0;" :: "r"(nthreads));
}
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(1+group));
}
uint32_t spins = 0;
uint32_t abort = 0;
inline __device__ int checkAbort(int i, int send) {
inline __device__ int checkAbort(int &spins, int i, int send) {
spins++;
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = *(comm->abortFlag);
if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
abort = *ncclShmem.comm.abortFlag;
spins = 0;
}
return abort;
}
inline __device__ void waitSend(int nbytes) {
spins = 0;
if (sendConnHeadPtr) {
int spins = 0;
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
sendConnHeadCache = *sendConnHeadPtr;
if (checkAbort(wid, 1)) break;
if (checkAbort(spins, wid, 1)) break;
}
if (sendConnFifoPtr) {
sendConnFifoPtr[sendStep[wid]%NCCL_STEPS] = nbytes;
@@ -80,137 +77,185 @@ class ncclLL128Primitives {
}
}
inline __device__ void incRecv(int i) {
recvStep[i] += 1;
}
inline __device__ void postRecv() {
if (recvConnHeadPtr) *recvConnHeadPtr = recvConnHead += 1;
}
inline __device__ void incSend(int i) {
sendStep[i] += 1;
}
inline __device__ void postSend() {
if (sendConnTailPtr) { __threadfence(); *sendConnTailPtr = sendConnTail += 1; }
}
template <int ELEMS_PER_THREAD>
inline __device__ void loadSrcToShmem128(int maxOffset, const uint64_t* src64Ptr) {
#if 0
uint64_t v[ELEMS_PER_THREAD];
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
if (u*WARP_SIZE < maxOffset) load128(src64Ptr+u*WARP_SIZE, v[u], v[u+1]);
}
uint64_t* shmemAsmPtr = shmemCvtPtr(shmem);
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
storeShmem128(shmemAsmPtr+u*WARP_SIZE, v[u], v[u+1]);
}
#else
uint64_t* shmemAsmPtr = shmemCvtPtr(shmem);
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
if (u*WARP_SIZE < maxOffset) {
uint64_t v0, v1;
load128(src64Ptr+u*WARP_SIZE, v0, v1);
storeShmem128(shmemAsmPtr+u*WARP_SIZE, v0, v1);
template<int WordPerThread>
__device__ __forceinline__ void loadRegsBegin(uint64_t(&regs)[WordPerThread], T const *src, int eltN) {
constexpr int EltPer16B = 16/sizeof(T);
if(reinterpret_cast<uintptr_t>(src)%16 == 0) {
/* We are aligned to 16 bytes, so load directly to registers no shmem.
* Flag threads load half as much data which gets shuffled to the even
* registers during Finish. The point of splitting into two phases is to
* defer that shuffle, which incurs a dependency stall, until after other
* memops are launched by the caller.
*/
#pragma unroll
for(int g=0; g < WordPerThread/2; g++) {
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
if(!flagThread || g%2==0) {
if(ix*EltPer16B < eltN)
load128((uint64_t*)(src + ix*EltPer16B), regs[2*g+0], regs[2*g+1]);
}
}
}
#endif
}
else {
// Not aligned. Stage the smallest 16 byte aligned region subsuming the
// buffer into shmem.
int misalignment = reinterpret_cast<uintptr_t>(src) % 16;
uint64_t *src8 = reinterpret_cast<uint64_t*>(reinterpret_cast<uintptr_t>(src) & -uintptr_t(16));
uint64_t *shm8 = shmemCvtPtr(ncclShmem.ll128warp[warp]);
#pragma unroll
for(int g=0; g < WordPerThread/2; g++)
if((g*WARP_SIZE + wid)*16 < misalignment + eltN*sizeof(T))
load128(src8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]);
#pragma unroll
for(int g=0; g < WordPerThread/2; g++)
storeShmem128(shm8 + 2*(g*WARP_SIZE + wid), regs[2*g+0], regs[2*g+1]);
inline __device__ void loadSrcToShmem(int start, int end, const T* srcPtr) {
T* shmemPtr = (T*)(shmem-2*wid);
for (int offset = start+wid; offset < end; offset += WARP_SIZE) {
shmemPtr[offset] = srcPtr[offset];
__syncwarp();
// Now load from shmem stage to regs. Preserve the same pre-shuffled layout
// as the aligned case since Finish() will be applied regardless.
T *shm = (T*)shm8 + misalignment/sizeof(T);
#pragma unroll
for(int g=0; g < WordPerThread/2; g++) {
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
if(!flagThread || g%2==0) {
if(ix*EltPer16B < eltN)
loadShmemMisaligned128(shm + ix*EltPer16B, regs[2*g+0], regs[2*g+1]);
}
}
}
}
template <int ELEMS_PER_THREAD>
inline __device__ void storeShmemToDst128(int maxOffset, uint64_t* dst64Ptr) {
uint64_t v[ELEMS_PER_THREAD];
uint64_t* shmemAsmPtr = shmemCvtPtr(shmem);
template<int WordPerThread>
__device__ __forceinline__ void loadRegsFinish(uint64_t(&regs)[WordPerThread]) {
// Move data out of flag registers into the vacant registers.
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
loadShmem128(shmemAsmPtr+u*WARP_SIZE, v[u], v[u+1]);
}
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
if (u*WARP_SIZE < maxOffset) store128(dst64Ptr+u*WARP_SIZE, v[u], v[u+1]);
for (int g=1; g < WordPerThread/2; g+=2) {
if (flagThread) regs[2*g] = regs[2*g-1];
}
}
inline __device__ void storeShmemToDst(int start, int end, T* dstPtr) {
T* shmemPtr = (T*)(shmem-2*wid);
for (int offset = start+wid; offset < end; offset += WARP_SIZE) {
dstPtr[offset] = shmemPtr[offset];
template<int WordPerThread>
__device__ __forceinline__ void storeRegs(T *dst, uint64_t(&regs)[WordPerThread], int eltN) {
constexpr int EltPer16B = 16/sizeof(T);
// Reverse Finish() register permuatation.
#pragma unroll
for (int g=1; g < WordPerThread/2; g+=2) {
if (flagThread) regs[2*g-1] = regs[2*g];
}
// Write to dst if 16-byte aligned, shmem otherwise.
int misalignment = reinterpret_cast<uintptr_t>(dst)%16;
uint64_t *shm8 = shmemCvtPtr(ncclShmem.ll128warp[warp]);
#pragma unroll
for(int g=0; g < WordPerThread/2; g++) {
int ix = g*WARP_SIZE - 4*(g/2) + wid - (g%2)*(wid/8);
if (!flagThread || g%2==0) {
if(misalignment == 0 && (ix+1)*EltPer16B <= eltN)
store128((uint64_t*)(dst + ix*EltPer16B), regs[2*g+0], regs[2*g+1]);
else
storeShmem128(shm8+2*ix, regs[2*g+0], regs[2*g+1]);
}
}
__syncwarp();
// Write rest from shmem to dst. No need to coalesce stores to 16-bytes,
// the hardware keeps up fine.
T *shm = (T*)ncclShmem.ll128warp[warp];
int skip = misalignment == 0 ? eltN & -EltPer16B : 0;
for(int i=skip+wid; i < eltN; i += WARP_SIZE)
dst[i] = shm[i];
}
#define WARP_MASK 0xffffffff
template <int ELEMS_PER_THREAD, int RECV, int SEND, int SRC, int DST>
__device__ __forceinline__ void recvReduceSendCopy(int ll128Offset) {
uint64_t v[ELEMS_PER_THREAD];
template <int ELEMS_PER_THREAD, int RECV, int SEND, int SrcBuf, int DstBuf>
__device__ __forceinline__ void recvReduceSendCopy(uint64_t(&v)[ELEMS_PER_THREAD], int ll128Offset, bool postOp) {
constexpr int SRC = SrcBuf != -1 ? 1 : 0;
uint64_t vr[ELEMS_PER_THREAD];
/************* Data Loading : SHMEM -> REG **************/
if (SRC) {
volatile uint64_t* shmem64Ptr = shmem - (2*wid)/NCCL_LL128_LINEELEMS;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
v[u] = shmem64Ptr[u*(WARP_SIZE-2)];
if (!flagThread) v[u+1] = shmem64Ptr[u*(WARP_SIZE-2)+1];
}
}
/*********** End Data Loading : SHMEM -> REG ************/
/************************ Recv **************************/
__syncwarp();
/************************ Wait first recv ********************/
if (RECV) {
uint64_t flag = recvFlag(0);
uint64_t* ptr = recvPtr(0)+ll128Offset;
uint64_t flag = recvFlag(0);
bool needReload;
uint64_t v0, v1;
int spins = 0;
do {
needReload = false;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
load128(ptr+u*WARP_SIZE, v0, v1);
needReload |= flagThread && (v1 != flag);
load128(ptr+u*WARP_SIZE, vr[u], vr[u+1]);
needReload |= flagThread && (vr[u+1] != flag);
}
} while (__any_sync(WARP_MASK, needReload) && checkAbort(spins, 0, 0) == 0);
}
/************* Finish register load **************/
if (SRC) {
// By deferring register shuffle here we've overlapped spinning on first
// peer's data with memory loads of src data.
loadRegsFinish(v);
if (SrcBuf == Input) {
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
v[u] = MULTI<RedOp, T>().preOp(redOp, v[u]);
if (!flagThread)
v[u+1] = MULTI<RedOp, T>().preOp(redOp, v[u+1]);
}
}
}
/************************ Recv rest *********************/
if (RECV) {
{ // Consume data from first recv
uint64_t* ptr = recvPtr(0)+ll128Offset;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
v[u] = SRC ? MULTI<RedOp, T>()(redOp, vr[u], v[u]) : vr[u];
v[u+1] = SRC ? MULTI<RedOp, T>()(redOp, vr[u+1], v[u+1]) : vr[u+1];
}
} while (__any_sync(WARP_MASK, needReload) && checkAbort(0, 0) == 0);
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
load128(ptr+u*WARP_SIZE, v0, v1);
v[u] = SRC ? MULTI<FUNC, T>()(v0, v[u]) : v0;
v[u+1] = SRC ? MULTI<FUNC, T>()(v1, v[u+1]) : v1;
}
for (int i=1; i<NRECV && i<nrecv; i++) {
for (int i=1; i<MaxRecv && i<fan.nrecv(); i++) {
uint64_t flag = recvFlag(i);
uint64_t* ptr = recvPtr(i)+ll128Offset;
uint64_t v0, v1;
bool needReload;
int spins = 0;
do {
needReload = false;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
load128(ptr+u*WARP_SIZE, v0, v1);
needReload |= flagThread && (v1 != flag);
load128(ptr+u*WARP_SIZE, vr[u], vr[u+1]);
needReload |= flagThread && (vr[u+1] != flag);
}
} while (__any_sync(WARP_MASK, needReload) && checkAbort(i, 0) == 0);
} while (__any_sync(WARP_MASK, needReload) && checkAbort(spins, i, 0) == 0);
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
load128(ptr+u*WARP_SIZE, v0, v1);
v[u] = MULTI<FUNC, T>()(v0, v[u]);
v[u+1] = MULTI<FUNC, T>()(v1, v[u+1]);
v[u] = MULTI<RedOp, T>()(redOp, vr[u], v[u]);
v[u+1] = MULTI<RedOp, T>()(redOp, vr[u+1], v[u+1]);
}
}
}
/********************** End Recv ************************/
if (postOp && !FuncTraits<RedOp>::IsPostOpIdentity) {
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
v[u] = MULTI<RedOp, T>().postOp(redOp, v[u]);
v[u+1] = MULTI<RedOp, T>().postOp(redOp, v[u+1]);
}
}
/************************ Send **************************/
if (SEND) {
for (int i=1; i<NSEND && i<nsend; i++) {
for (int i=1; i<MaxSend && i<fan.nsend(); i++) {
uint64_t flag = sendFlag(i);
uint64_t* ptr = sendPtr(i)+ll128Offset;
#pragma unroll
@@ -226,82 +271,61 @@ class ncclLL128Primitives {
}
}
/********************** End Send ************************/
/************* Data Storing : REG -> SHMEM **************/
if (DST) {
volatile uint64_t* shmem64Ptr = shmem - (2*wid)/NCCL_LL128_LINEELEMS;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
shmem64Ptr[u*(WARP_SIZE-2)] = v[u];
if (!flagThread) shmem64Ptr[u*(WARP_SIZE-2)+1] = v[u+1];
}
}
/*********** End data Storing : REG -> SHMEM ************/
}
#define LL128INC (WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD)
#define ELEMINC (LL128INC-(LL128INC/NCCL_LL128_LINEELEMS))
static constexpr int WireWordPerSlice = WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD;
static constexpr int DataEltPerSlice = (WireWordPerSlice - WireWordPerSlice/NCCL_LL128_LINEELEMS)*(sizeof(uint64_t)/sizeof(T));
template <int RECV, int SEND, int SRC, int DST>
__device__ void GenericOp(const T* srcPtr, T* dstPtr, int nelem) {
if (nelem <= 0) {
// Don't move any data but still increase steps and sync with prev/next
if (SEND) waitSend(0);
FOR_SEND(incSend); if (SEND) postSend();
FOR_RECV(incRecv); if (RECV) postRecv();
return;
}
const int nelem64 = ((nelem*sizeof(T))/(2*sizeof(uint64_t)))*2;
const uint64_t* src64Ptr = ((uint64_t*)srcPtr);
uint64_t* dst64Ptr = ((uint64_t*)dstPtr);
template <int RECV, int SEND, int SrcBuf, int DstBuf>
__device__ void GenericOp(intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp) {
constexpr int SRC = SrcBuf != -1 ? 1 : 0;
constexpr int DST = DstBuf != -1 ? 1 : 0;
static_assert(-1<=SrcBuf && SrcBuf < 2, "Uhoh");
static_assert(-1<=DstBuf && DstBuf < 2, "Uhoh");
static_assert(DstBuf!=Input, "Mistake?");
#if 0
assert((SrcBuf==-1) == (srcIx==-1));
assert((DstBuf==-1) == (dstIx==-1));
#endif
int ll128Offset = LL128INC*warp+2*wid;
int elemOffset = ELEMINC*warp;
T const *srcPtr = SrcBuf == -1 ? nullptr : userBufs[SrcBuf] + srcIx;
T *dstPtr = DstBuf == -1 ? nullptr : userBufs[DstBuf] + dstIx;
int wireOffset = WireWordPerSlice*warp + 2*wid;
const int nwarps = nthreads/WARP_SIZE;
nelem = nelem < 0 ? 0 : nelem;
if (SEND) waitSend(DIVUP(nelem*sizeof(T), ELEMINC*sizeof(uint64_t))*LL128INC*sizeof(uint64_t));
if (SEND) waitSend(divUp(nelem, DataEltPerSlice)*WireWordPerSlice*sizeof(uint64_t));
barrier();
nelem -= DataEltPerSlice*warp;
srcPtr += DataEltPerSlice*warp;
dstPtr += DataEltPerSlice*warp;
while (nelem > 0) {
const int eltInSlice = min(nelem, DataEltPerSlice);
uint64_t regs[NCCL_LL128_SHMEM_ELEMS_PER_THREAD];
if (SRC) loadRegsBegin(regs, srcPtr, eltInSlice);
recvReduceSendCopy<NCCL_LL128_SHMEM_ELEMS_PER_THREAD, RECV, SEND, SrcBuf, DstBuf>(regs, wireOffset, postOp);
if (DST) storeRegs(dstPtr, regs, eltInSlice);
while (elemOffset*(sizeof(uint64_t)/sizeof(T)) < nelem) {
const int maxOffset128 = min(nelem64-elemOffset, (int)ELEMINC);
const int maxOffset = min(nelem-(elemOffset*((int)(sizeof(uint64_t)/sizeof(T)))), (int)(ELEMINC*(sizeof(uint64_t)/sizeof(T))));
if (SRC) {
int done = 0;
if ((((uint64_t)srcPtr)&0xf) == 0) {
loadSrcToShmem128<NCCL_LL128_SHMEM_ELEMS_PER_THREAD>(maxOffset128-2*wid, src64Ptr+elemOffset+2*wid);
done = maxOffset128*(sizeof(uint64_t)/sizeof(T));
}
loadSrcToShmem(done, maxOffset, (T*)(src64Ptr+elemOffset));
}
__syncwarp();
recvReduceSendCopy<NCCL_LL128_SHMEM_ELEMS_PER_THREAD, RECV, SEND, SRC, DST>(ll128Offset);
__syncwarp();
if (DST) {
int done = 0;
if ((((uint64_t)dstPtr)&0xf) == 0) {
storeShmemToDst128<NCCL_LL128_SHMEM_ELEMS_PER_THREAD>(maxOffset128-2*wid, dst64Ptr+elemOffset+2*wid);
done = maxOffset128*(sizeof(uint64_t)/sizeof(T));
}
storeShmemToDst(done, maxOffset, (T*)(dst64Ptr+elemOffset));
}
__syncwarp();
ll128Offset += LL128INC*nwarps;
elemOffset += ELEMINC*nwarps;
wireOffset += WireWordPerSlice*nwarps;
srcPtr += DataEltPerSlice*nwarps;
dstPtr += DataEltPerSlice*nwarps;
nelem -= DataEltPerSlice*nwarps;
}
barrier();
FOR_SEND(incSend); if (SEND) postSend();
FOR_RECV(incRecv); if (RECV) postRecv();
if (SEND) for (int i=0; i < MaxSend; i++) sendStep[i] += 1;
if (SEND) postSend();
if (RECV) for (int i=0; i < MaxRecv; i++) recvStep[i] += 1;
if (RECV) postRecv();
}
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
recvBuff[i] = (uint64_t*)conn->buffs[NCCL_PROTO_LL128];
recvStep[i] = conn->step;
if (wid == i) recvConn = conn;
nrecv++;
}
__device__ __forceinline__ void loadRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv()) {
recvConnHeadPtr = recvConn->head;
recvConnHead = recvConn->step;
}
@@ -311,16 +335,15 @@ class ncclLL128Primitives {
sendBuff[i] = (uint64_t*)conn->buffs[NCCL_PROTO_LL128];
sendStep[i] = conn->step;
if (wid == i) sendConn = conn;
nsend++;
}
__device__ __forceinline__ void loadSendSync() {
if (tid < nsend) {
if (tid < fan.nsend()) {
sendConnHeadPtr = sendConn->head;
sendConnHeadCache = *sendConnHeadPtr;
sendConnHead = sendConn->step;
sendConnFifoPtr = sendConn->sizesFifo;
}
if (tid >= nthreads-WARP_SIZE && wid<nsend) {
if (tid >= nthreads-WARP_SIZE && wid<fan.nsend()) {
if (sendConn->sizesFifo) {
sendConnTailPtr = sendConn->tail;
sendConnTail = sendConn->step;
@@ -328,64 +351,74 @@ class ncclLL128Primitives {
}
}
__device__ __forceinline__ void saveRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
recvConn->step = recvConnHead;
__threadfence_block();
public:
__device__ Primitives(
const int tid, const int nthreads, int const *recvPeers, int const *sendPeers,
void const *inputBuf, void *outputBuf, int group=0
):
redOp(FuncTraits<RedOp>().make(ncclShmem.comm.nRanks)),
tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), warp(tid/WARP_SIZE),
flagThread((tid%8)==7), group(group),
stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)) {
auto *channel = &ncclShmem.channel;
int nrecv=0, nsend=0;
while (nrecv < MaxRecv && recvPeers[nrecv] >= 0) {
loadRecvConn(&channel->devPeers[recvPeers[nrecv]].recv->conn, nrecv);
nrecv++;
}
}
__device__ __forceinline__ void saveSendSync() {
if (tid < nsend) {
sendConn->step = sendConnHead;
__threadfence_block();
while (nsend < MaxSend && sendPeers[nsend] >= 0) {
loadSendConn(&channel->devPeers[sendPeers[nsend]].send->conn, nsend);
nsend++;
}
}
public:
__device__ __forceinline__
ncclLL128Primitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, int stepSize, struct ncclChannel* channel, struct ncclDevComm* comm)
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), warp(tid/WARP_SIZE), flagThread((tid%8)==7), stepSize(stepSize), shmem(ncclShmem->data+(threadIdx.x/WARP_SIZE)*NCCL_LL128_SHMEM_ELEMS_PER_THREAD*WARP_SIZE+2*wid) {
// Make sure step is updated before we read it.
barrier();
for (int i=0; i<NRECV && recvPeers[i] >= 0; i++) loadRecvConn(&channel->devPeers[recvPeers[i]].recv->conn, i);
for (int i=0; i<NSEND && sendPeers[i] >= 0; i++) loadSendConn(&channel->devPeers[sendPeers[i]].send->conn, i);
this->fan = Fan(nrecv, nsend);
loadRecvSync();
loadSendSync();
setDataPtrs(inputBuf, outputBuf);
}
__device__ void send(const T* src, int nelem) {
return GenericOp<0, 1, 1, 0>(src, NULL, nelem);
}
__device__ void recv(T* dst, int nelem) {
return GenericOp<1, 0, 0, 1>(NULL, dst, nelem);
}
__device__ void recvReduceSend(const T* src, int nelem) {
return GenericOp<1, 1, 1, 0>(src, NULL, nelem);
}
__device__ void recvReduceCopy(const T* src, T* dst, int nelem) {
return GenericOp<1, 0, 1, 1>(src, dst, nelem);
}
__device__ void copySend(const T* src, T* dst, int nelem) {
return GenericOp<0, 1, 1, 1>(src, dst, nelem);
}
__device__ void recvCopySend(T* dst, int nelem) {
return GenericOp<1, 1, 0, 1>(NULL, dst, nelem);
}
__device__ void recvReduceCopySend(const T* src, T* dst, int nelem) {
return GenericOp<1, 1, 1, 1>(src, dst, nelem);
}
__device__ __forceinline__ ~ncclLL128Primitives() {
__device__ ~Primitives() {
// Save steps for the next operation
saveRecvSync();
saveSendSync();
if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv())
recvConn->step = recvConnHead;
if (tid < fan.nsend())
sendConn->step = sendConnHead;
// Ensure all steps written back
barrier();
}
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
userBufs[Input] = (T*)inputBuf;
userBufs[Output] = (T*)outputBuf;
}
__device__ void moveDataPtrs(intptr_t delta) {
userBufs[Input] += delta;
userBufs[Output] += delta;
}
__device__ void send(intptr_t inpIx, int eltN) {
return GenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
}
__device__ void sendFromOutput(intptr_t outIx, int eltN) {
return GenericOp<0, 1, Output, -1>(outIx, -1, eltN, false);
}
__device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
return GenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
}
__device__ void recvReduceSend(intptr_t inpIx, int eltN) {
return GenericOp<1, 1, Input, -1>(inpIx, -1, eltN, false);
}
__device__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return GenericOp<1, 0, Input, Output>(inpIx, outIx, eltN, postOp);
}
__device__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return GenericOp<0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
}
__device__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
return GenericOp<1, 1, -1, Output>(-1, outIx, eltN, postOp);
}
__device__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
return GenericOp<1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
}
};
src/collectives/device/prims_simple.h
+463
View File
@@ -0,0 +1,463 @@
/*************************************************************************
* Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
template<typename T, typename RedOp, typename Fan, int Direct,
int SlicePerChunk, int StepPerSlice, int Unroll>
class Primitives<
T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll>
> {
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,
PtrsFifoEnabled = 0x80,
SizesFifoEnabled = 0x100,
DirectEnabled = 0x200,
ThreadsSynced = 0x400;
const int tid;
int nthreads;
int nworkers;
const int stepSize;
Fan fan;
RedOp const redOp;
int index; // Peer index I'm responsible for
int flags;
int group;
uint64_t step;
union {
void **connPtrsFifoPtr; // (flags & PtrsFifoEnabled)
T *userBuff; // (flags & (RoleInput|RoleOutput))
T *connEltsFifo; // !(flags & (PtrsFifoEnabled|RoleInput|RoleOutput))
};
union {
int volatile *connSizesFifoPtr; // (flags & SizesFifoEnabled)
T *directBuff; // !(flags & SizesFifoEnabled)
};
uint64_t volatile *connStepPtr;
uint64_t connStepCache; // Cache last seen value of (*connStepPtr)
// Don't use barrier 0 as it's used by the final sync
inline __device__ void barrier() {
if (nthreads == WARP_SIZE)
__syncwarp();
else
asm volatile("bar.sync %0, %1;" :: "r"(group+1), "r"(nthreads));
flags |= ThreadsSynced;
}
inline __device__ void subBarrier() {
if (nworkers == nthreads)
barrier();
else
asm volatile("bar.sync %0, %1;" :: "r"(group+2), "r"(nworkers));
}
inline __device__ bool checkAbort(int &spins) {
spins++;
if (!(flags & Aborted) && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
flags |= *ncclShmem.comm.abortFlag ? Aborted : 0;
spins = 0;
}
return flags & Aborted;
}
template <int DirectRecv, int DirectSend, int Recv, int Send, int Src, int Dst>
inline __device__ void waitPeer(intptr_t dstIx, intptr_t remoteOutIx, int offset, int nelts) {
if (flags & (Recv*RoleWaitRecv | Send*RoleWaitSend)) {
bool const isSendNotRecv = (Send && Recv) ? (flags & RoleWaitSend) : Send;
int spins = 0;
while (connStepCache + (isSendNotRecv ? NCCL_STEPS : 0) < step + StepPerSlice) {
connStepCache = *connStepPtr;
if (checkAbort(spins)) break;
//if (spins == 0) printf("r=%d b=%d t=%d SPUN OUT got=%d want=%d\n", ncclShmem.comm.rank, blockIdx.x, threadIdx.x, int(connStepCache + (isSendNotRecv ? NCCL_STEPS : 0)), int(step+StepPerSlice));
}
if (isSendNotRecv && (flags & SizesFifoEnabled))
connSizesFifoPtr[step%NCCL_STEPS] = nelts*sizeof(T);
void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst)
: (ncclShmem.groups[group].srcs + Src);
if (flags & PtrsFifoEnabled)
loadPtr(connPtrsFifoPtr + step%NCCL_STEPS, ptrs[index]);
else if ((isSendNotRecv ? DirectSend : DirectRecv) && (flags & DirectEnabled))
ptrs[index] = directBuff + (isSendNotRecv ? remoteOutIx : dstIx) + offset;
else
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
step += StepPerSlice;
}
}
template<int Recv, int Send>
inline __device__ void postPeer() {
if (flags & (Recv*RolePostRecv | Send*RolePostSend)) {
step += StepPerSlice;
*connStepPtr = step;
}
}
template <int DirectRecv1, int DirectSend1, int Recv, int Send, int SrcBuf, int DstBuf>
inline __device__ void genericOp(
intptr_t srcIx, intptr_t dstIx, intptr_t remoteOutIx, int nelem, bool postOp
) {
constexpr int DirectRecv = 1 && Direct && DirectRecv1;
constexpr int DirectSend = 1 && Direct && DirectSend1;
constexpr int Src = SrcBuf != -1;
constexpr int Dst = DstBuf != -1;
nelem = nelem < 0 ? 0 : nelem;
int sliceSize = stepSize*StepPerSlice;
sliceSize = max(divUp(nelem, 16*SlicePerChunk)*16, sliceSize/32);
int slice = 0;
int offset = 0;
if (tid < nworkers && offset < nelem) {
// 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.
// Using "number of branch insns (taken or not) encountered dynamically"
// as the performance metric, then:
// perf_orig = 2*numslices
// perf_new = 2+numslices
// So the new code and old code behave the same for numslices=2, and for
// numslices>2 the new code is superior. And note that in the case
// numslices=1, the loop is trivially unrollable (single iteration) so we
// don't incur that that tail branch and we still have perf_new=2.
//
// ORIGINAL CODE:
// unrolled for(slices) {
// if(worker) { // This branch removed
// wait();
// subBarrier();
// if(slice not empty) // This branch removed
// ReduceCopyMulti();
// }
// barrier();
// post();
// } // Since we no longer unroll, new branch added here
#if __CUDA_ARCH__ < 700
// Yeah, so all that above don't matter a lick on older hardware.
#pragma unroll SlicePerChunk
#else
#pragma unroll 1
#endif
do {
sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
if (Src && (flags & (SrcBuf==Input ? RoleInput : RoleOutput)))
ncclShmem.groups[group].srcs[0] = userBuff + srcIx + offset;
if (Dst && (flags & (DstBuf==Input ? RoleInput : RoleOutput)))
ncclShmem.groups[group].dsts[0] = userBuff + dstIx + offset;
waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(dstIx, remoteOutIx, offset, sliceSize);
subBarrier();
if (DirectRecv && ncclShmem.groups[group].srcs[0] == ncclShmem.groups[group].dsts[0]) {
// We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
if (Send) {
// (1-Send) is only there to avoid compilation errors in case MaxSend=0 (and Send=0).
ReduceOrCopyMulti<Unroll, RedOp, T, 1, 1, 1, (1-Send)+MaxSend>
(tid, nworkers, redOp, false, false,
1, (T const**)ncclShmem.groups[group].srcs,
fan.nsend(), (T**)ncclShmem.groups[group].dsts+1,
sliceSize);
}
} else {
ReduceOrCopyMulti<Unroll, RedOp, T, Recv+Src, Recv*MaxRecv+Src, Send+Dst, Send*MaxSend+Dst>
(tid, nworkers, redOp, SrcBuf==Input, postOp,
Recv*fan.nrecv()+Src, (T const**)ncclShmem.groups[group].srcs,
Send*fan.nsend()+Dst, (T**)ncclShmem.groups[group].dsts,
sliceSize);
}
barrier(); // This barrier has a counterpart in following loop
if (Send && (flags & RolePostSend) && index == 0) __threadfence_system();
__syncwarp();
postPeer<Recv, Send>();
offset += sliceSize;
slice += 1;
} while (slice < SlicePerChunk && offset < nelem);
}
// Non-workers come straight here. Workers too but only once the remaining
// slices are all empty. Since empty slices are the uncommon case, and
// worker perf is the limiter, perf-wise this loop is effectively unentered,
// hence just a single branch insn.
#pragma unroll 1
while (slice < SlicePerChunk) {
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);
}
barrier(); // Has couterpart in preceding worker-only loop.
if (Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) __threadfence_system();
__syncwarp();
postPeer<Recv, Send>();
offset += sliceSize;
slice += 1;
}
}
// Scatter and gather do not support Direct
template <int Recv, int Send>
inline __device__ void
ScatterGatherOp(intptr_t inpIx, intptr_t outIx, int totalElem, int peerElem, int skip, int shift, bool postOp) {
int offset = 0; // slice offset
int sliceSize = stepSize*StepPerSlice;
int dataSize = max(DIVUP(peerElem, 16*SlicePerChunk)*16, sliceSize/32); // per-peer slice size
#pragma unroll
for (int slice=0; slice<SlicePerChunk; ++slice) {
int realSize = max(0, min(dataSize, peerElem-offset));
if (tid < nworkers) {
if (Send && (flags & RoleInput)) ncclShmem.groups[group].srcs[0] = userBuff + inpIx + offset;
if (Recv && (flags & RoleOutput)) ncclShmem.groups[group].dsts[0] = userBuff + outIx + offset;
// realSize is not accurate here; but intra-node does not rely on sizes FIFO
waitPeer<0, 0, Recv, Send, 0, 0>(0, 0, 0, realSize);
subBarrier();
if (Send) {
#pragma unroll
for (int j=0; j<fan.nsend(); j++) {
int i = (j+shift)%fan.nsend();
int peerOffset = i*peerElem;
if (skip >= 0 && i >= skip) peerOffset += peerElem;
const T* src0 = (T*)ncclShmem.groups[group].srcs[0] + peerOffset;
int realPeerSize = min(realSize, totalElem-peerOffset);
if (realPeerSize > 0) ReduceOrCopyMulti<Unroll, RedOp, T, 1, 1, 1, 1>(tid, nworkers, redOp, true, false, 1, &src0, 1, (T**)ncclShmem.groups[group].dsts+i, realPeerSize);
}
} else if (Recv) {
#pragma unroll
for (int j=0; j<fan.nrecv(); j++) {
int i = (j+shift)%fan.nrecv();
int peerOffset = i*peerElem;
if (skip >= 0 && i >= skip) peerOffset += peerElem;
T* dst0 = (T*)ncclShmem.groups[group].dsts[0] + peerOffset;
int realPeerSize = min(realSize, totalElem-peerOffset);
if (realPeerSize > 0) ReduceOrCopyMulti<Unroll, RedOp, T, 1, 1, 1, 1>(tid, nworkers, redOp, false, postOp, 1, (T const**)ncclShmem.groups[group].srcs+i, 1, &dst0, realPeerSize);
}
}
}
barrier();
if (Send && (flags & RolePostSend) && realSize > 0 && index == 0) __threadfence_system();
__syncwarp();
postPeer<Recv, Send>();
offset += realSize;
}
}
__device__ __forceinline__ void loadRecvConn(ncclPeer *peer) {
if (flags & (RoleWaitRecv|RolePostRecv)) {
// For other colls: group <= 2, hence always use conn 0
// For P2P: Direct is set to 1, hence always use conn 0
// Ideally we should be accepting connIndex from the constructor!
const int connIndex = Direct ? 0 : group/4;
auto *conn = &peer->recv[connIndex].conn;
step = conn->step;
step = roundUp(step, SlicePerChunk*StepPerSlice);
if (flags & RolePostRecv) {
connStepPtr = conn->head;
*connStepPtr = step; // Return credits in case we rounded up.
}
if (flags & RoleWaitRecv) {
ncclShmem.groups[group].recvConns[index] = conn; // WaitRecv role saves since that's who needs it in setDataPtrs()
connStepPtr = conn->tail;
connStepCache = *connStepPtr;
flags |= (conn->ptrsFifo != nullptr) ? PtrsFifoEnabled : 0;
flags |= (Direct && (conn->direct & NCCL_DIRECT_GPU)) ? DirectEnabled : 0;
if (flags & PtrsFifoEnabled)
connPtrsFifoPtr = conn->ptrsFifo;
else
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
}
}
}
__device__ __forceinline__ void loadSendConn(ncclPeer *peer) {
if (flags & (RoleWaitSend|RolePostSend)) {
// For other colls: group <= 2, hence always use conn 0
// For P2P: Direct is set to 1, hence always use conn 0
// Ideally we should be accepting connIndex from the constructor!
const int connIndex = Direct ? 0 : group/4;
auto *conn = &peer->send[connIndex].conn;
step = conn->step;
step = roundUp(step, SlicePerChunk*StepPerSlice);
if (flags & RolePostSend) {
connStepPtr = conn->tail;
}
if (flags & RoleWaitSend) {
ncclShmem.groups[group].sendConns[index] = conn; // WaitSend role saves since that's who needs it in setDataPtrs()
connStepPtr = conn->head;
connStepCache = *connStepPtr;
flags |= (conn->ptrsFifo != nullptr) ? PtrsFifoEnabled : 0;
if (flags & PtrsFifoEnabled)
connPtrsFifoPtr = conn->ptrsFifo;
else
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (conn->sizesFifo != nullptr) {
flags |= SizesFifoEnabled;
connSizesFifoPtr = conn->sizesFifo;
}
else if (Direct && (conn->direct & NCCL_DIRECT_GPU))
flags |= DirectEnabled;
}
}
}
public:
__device__ Primitives(
int tid, int nthreads, int const *recvPeers, int const *sendPeers,
void const *inputBuf, void *outputBuf, int group=0
):
tid(tid),
stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T)),
redOp(FuncTraits<RedOp>::make(ncclShmem.comm.nRanks)) {
// For send operations, we need an extra warp to overlap the threadfence and the copy
this->nthreads = nthreads;
this->nworkers = nthreads - (MaxSend > 0 && nthreads-WARP_SIZE >= 64 ? WARP_SIZE : 0);
this->group = group;
int nrecv=0, nsend=0;
while (nrecv < MaxRecv && recvPeers[nrecv] != -1) nrecv++;
while (nsend < MaxSend && sendPeers[nsend] != -1) nsend++;
this->fan = Fan(nrecv, nsend);
constexpr int ThreadPerSync = 8;
static_assert(MaxSend < ThreadPerSync && MaxRecv < ThreadPerSync, "Not enough threads to cover all peers");
int g = tid / ThreadPerSync;
int ng = nthreads / ThreadPerSync;
index = tid % ThreadPerSync;
flags = 0;
if (g == 0) {
if (index < nrecv) flags |= RoleWaitRecv;
if (index == nrecv) flags |= RoleInput;
} else if (g == 1) {
if (index < nsend) flags |= RoleWaitSend;
if (index == nsend) flags |= RoleOutput;
} else if (g == ng - 2) {
if (index < nrecv) flags |= RolePostRecv;
} else if (g == ng - 1) {
if (index < nsend) flags |= RolePostSend;
}
int peer = 0;
if (flags & (RoleWaitRecv|RolePostRecv)) peer = recvPeers[index];
if (flags & (RoleWaitSend|RolePostSend)) peer = sendPeers[index];
loadRecvConn(&ncclShmem.channel.devPeers[peer]);
loadSendConn(&ncclShmem.channel.devPeers[peer]);
setDataPtrs(inputBuf, outputBuf);
}
__device__ ~Primitives() {
// Ensure ncclShmem.groups[].send/recvConns are available
if (!(flags & ThreadsSynced))
barrier();
// Save steps for the next operation
if (flags & (RolePostSend|RolePostRecv)) {
auto *conns = (flags & RolePostSend) ? ncclShmem.groups[group].sendConns : ncclShmem.groups[group].recvConns;
conns[index]->step = step;
}
// Make sure all threads are done writing back conn->step and done using
// ncclShmem.groups[group]
barrier();
}
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
if (flags & RoleInput) userBuff = (T*)inputBuf;
if (flags & RoleOutput) userBuff = (T*)outputBuf;
if (Direct && flags == (flags|RoleWaitRecv|DirectEnabled)) {
int spins = 0;
void *volatile *slot = ncclShmem.groups[group].recvConns[index]->ptrExchange;
// Wait for consumer to consume previous value before trampling it.
while (*slot != nullptr && !checkAbort(spins));
directBuff = (T*)outputBuf;
// Encode pointer by XOR'ing against some address they definitely wouldn't send
// since we want to allow them sending us nullptr while not colliding with
// the empty slot value.
*slot = reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(outputBuf) ^ reinterpret_cast<uintptr_t>(slot));
}
if (Direct && flags == (flags|RoleWaitSend|DirectEnabled)) {
int spins = 0;
void *volatile *slot = ncclShmem.groups[group].sendConns[index]->ptrExchange;
void *ptr;
while (true) {
ptr = *slot;
if (ptr != nullptr || checkAbort(spins)) break;
}
directBuff = reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(ptr) ^ reinterpret_cast<uintptr_t>(slot));
*slot = nullptr;
}
}
__device__ void moveDataPtrs(intptr_t delta) {
if (flags & (RoleInput|RoleOutput))
userBuff += delta;
}
__device__ __forceinline__ void send(intptr_t inpIx, int eltN) {
genericOp<0, 0, 0, 1, Input, -1>(inpIx, -1, -1, eltN, false);
}
__device__ __forceinline__ void sendFromOutput(intptr_t outIx, int eltN) {
genericOp<0, 0, 0, 1, Output, -1>(outIx, -1, -1, eltN, false);
}
__device__ __forceinline__ void directSend(intptr_t inpIx, intptr_t remoteOutIx, int eltN) {
genericOp<0, 1, 0, 1, Input, -1>(inpIx, -1, remoteOutIx, eltN, false);
}
__device__ __forceinline__ void directSendFromOutput(intptr_t outIx, intptr_t remoteOutIx, int eltN) {
genericOp<0, 1, 0, 1, Output, -1>(outIx, -1, remoteOutIx, eltN, false);
}
__device__ __forceinline__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 0, -1, Output>(-1, outIx, -1, eltN, postOp);
}
__device__ __forceinline__ void directRecv(intptr_t outIx, int eltN) {
genericOp<1, 0, 1, 0, -1, Output>(-1, outIx, -1, eltN, /*postOp=*/false);
}
__device__ __forceinline__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 0, 1, Input, Output>(inpIx, outIx, -1, eltN, postOp);
}
__device__ __forceinline__ void directCopySend(intptr_t inpIx, intptr_t outIx, intptr_t remoteOutIx, int eltN, bool postOp=false) {
genericOp<0, 1, 0, 1, Input, Output>(inpIx, outIx, remoteOutIx, eltN, postOp);
}
__device__ __forceinline__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 1, -1, Output>(-1, outIx, -1, eltN, postOp);
}
__device__ __forceinline__ void directRecvCopySend(intptr_t outIx, intptr_t remoteOutIx, int eltN) {
genericOp<1, 1, 1, 1, -1, Output>(-1, outIx, remoteOutIx, eltN, false);
}
__device__ __forceinline__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 0, Input, Output>(inpIx, outIx, -1, eltN, postOp);
}
__device__ __forceinline__ void recvReduceSend(intptr_t inpIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 1, Input, -1>(inpIx, -1, -1, eltN, postOp);
}
__device__ __forceinline__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
genericOp<0, 0, 1, 1, Input, Output>(inpIx, outIx, -1, eltN, postOp);
}
__device__ __forceinline__ void directRecvReduceCopySend(intptr_t inpIx, intptr_t outIx, intptr_t remoteOutIx, int eltN, bool postOp=false) {
// Direct is only for the send part
genericOp<0, 1, 1, 1, Input, Output>(inpIx, outIx, remoteOutIx, eltN, postOp);
}
__device__ __forceinline__ void
scatter(intptr_t inpIx, int totalElem, int peerElem, int skip, int shift) {
ScatterGatherOp<0, 1>(inpIx, -1, totalElem, peerElem, skip, shift, /*postOp=*/false);
}
__device__ __forceinline__ void
gather(intptr_t outIx, int totalElem, int peerElem, int skip, int shift, bool postOp=false) {
ScatterGatherOp<1, 0>(-1, outIx, totalElem, peerElem, skip, shift, postOp);
}
};
src/collectives/device/reduce.h
+63 -124
View File
@@ -5,148 +5,87 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduce, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-WARP_SIZE;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * REDUCE_CHUNKSTEPS;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
const ssize_t size = args->coll.count;
const int rank = ring->devUserRanks[0];
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->coll.root;
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
ncclRing *ring = &ncclShmem.channel.ring;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCE_CHUNKSTEPS : 1));
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const int nranks = ncclShmem.comm.nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = ncclShmem.comm.rank;
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->coll.root;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff);
ncclPrimitives<UNROLL, REDUCE_CHUNKSTEPS/REDUCE_SLICESTEPS, REDUCE_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
auto calcChunkSize = [&]__device__(ssize_t gridOffset)->int {
int realChunkSize;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels));
realChunkSize = roundUp(realChunkSize, (nthreads-WARP_SIZE)*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128, chunkSize);
return realChunkSize;
};
if (prevRank == root) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nChannels));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
int realChunkSize = calcChunkSize(gridOffset);
ssize_t offset = gridOffset + bid*realChunkSize;
int nelem = min(realChunkSize, size-offset);
if (prevRank == root) {
prims.send(thisInput+offset, nelem);
} else if (rank == root) {
prims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);
} else {
prims.recvReduceSend(thisInput+offset, nelem);
}
prims.send(offset, nelem);
}
}
};
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduce, NCCL_ALGO_RING, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = comm->rank;
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->coll.root;
ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepLines, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
else if (rank == root) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
if (size-gridOffset < loopSize) {
chunkSize = args->coll.lastChunkSize;
}
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (prevRank == root) {
LLprims.send(thisInput+offset, nelem);
} else if (rank == root) {
LLprims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);
} else {
LLprims.recvReduceSend(thisInput+offset, nelem);
}
int realChunkSize = calcChunkSize(gridOffset);
ssize_t offset = gridOffset + bid*realChunkSize;
int nelem = min(realChunkSize, size-offset);
prims.recvReduceCopy(offset, offset, nelem, /*postOp=*/true);
}
}
};
#include "prims_ll128.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduce, NCCL_ALGO_RING, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = stepSize*NCCL_LL128_DATAELEMS*sizeof(uint64_t) / (NCCL_LL128_LINEELEMS*sizeof(T));
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T));
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const int rank = comm->rank;
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->coll.root;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepSize, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
else {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*minChunkSize)*minChunkSize, chunkSize);
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (prevRank == root) {
LLprims.send(thisInput+offset, nelem);
} else if (rank == root) {
LLprims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);
} else {
LLprims.recvReduceSend(thisInput+offset, nelem);
}
int realChunkSize = calcChunkSize(gridOffset);
ssize_t offset = gridOffset + bid*realChunkSize;
int nelem = min(realChunkSize, size-offset);
prims.recvReduceSend(offset, nelem);
}
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
using Proto = ProtoSimple<REDUCE_CHUNKSTEPS/REDUCE_SLICESTEPS, REDUCE_SLICESTEPS>;
runRing<T, RedOp, Proto>(args);
}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncReduce, NCCL_ALGO_TREE, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL>(args);
}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncReduce, NCCL_ALGO_COLLNET, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL128>(args);
}
};
src/collectives/device/reduce_kernel.h
+357 -1
View File
@@ -1,5 +1,5 @@
/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -10,6 +10,7 @@
#include "common_kernel.h"
#include <limits>
#include <type_traits>
template<typename T>
struct FuncNull {
@@ -46,6 +47,18 @@ struct FuncMin {
}
};
template<typename Fn>
struct FuncTraits { // generic implementation for FuncSum,Prod,Min,Max
static constexpr bool IsPreOpIdentity = true;
static constexpr bool IsPostOpIdentity = true;
__device__ static Fn make(int rankN) { return Fn(); }
template<typename T>
__device__ static T preOp(Fn, T x) { return x; }
template<typename T>
__device__ static T postOp(Fn, T x) { return x; }
};
#define MASK0 0x00ff00ff
#define MASK1 0xff00ff00
static __device__ uint32_t addChar4(const uint32_t x, const uint32_t y) {
@@ -239,6 +252,31 @@ struct FuncSum<half> {
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<>
struct FuncSum<__nv_bfloat16> {
__device__ __nv_bfloat162 operator()(const __nv_bfloat162 x, const __nv_bfloat162 y) const {
#if __CUDA_ARCH__ >= 800
return __hadd2(x, y);
#else
float fxl, fxh, fyl, fyh;
fxl = __low2float(x);
fxh = __high2float(x);
fyl = __low2float(y);
fyh = __high2float(y);
return __floats2bfloat162_rn(fxl + fyl, fxh + fyh);
#endif
}
__device__ __nv_bfloat16 operator()(const __nv_bfloat16 x, const __nv_bfloat16 y) const {
#if __CUDA_ARCH__ >= 800
return __hadd(x, y);
#else
return __float2bfloat16( __bfloat162float(x) + __bfloat162float(y) );
#endif
}
};
#endif
template<>
struct FuncProd<half> {
__device__ half2 operator()(const half2 x, const half2 y) const {
@@ -262,6 +300,31 @@ struct FuncProd<half> {
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<>
struct FuncProd<__nv_bfloat16> {
__device__ __nv_bfloat162 operator()(const __nv_bfloat162 x, const __nv_bfloat162 y) const {
#if __CUDA_ARCH__ >= 800
return __hmul2(x, y);
#else
float fxl, fxh, fyl, fyh;
fxl = __low2float(x);
fxh = __high2float(x);
fyl = __low2float(y);
fyh = __high2float(y);
return __floats2bfloat162_rn(fxl * fyl, fxh * fyh);
#endif
}
__device__ __nv_bfloat16 operator()(const __nv_bfloat16 x, const __nv_bfloat16 y) const {
#if __CUDA_ARCH__ >= 800
return __hmul(x, y);
#else
return __float2bfloat16( __bfloat162float(x) * __bfloat162float(y) );
#endif
}
};
#endif
template<>
struct FuncMax<half> {
__device__ half2 operator()(const half2 x, const half2 y) const {
@@ -281,6 +344,34 @@ struct FuncMax<half> {
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<>
struct FuncMax<__nv_bfloat16> {
__device__ __nv_bfloat162 operator()(const __nv_bfloat162 x, const __nv_bfloat162 y) const {
#if __CUDA_ARCH__ >= 800
return __hmax2(x, y);
#else
float fxl, fxh, fyl, fyh;
fxl = __low2float(x);
fxh = __high2float(x);
fyl = __low2float(y);
fyh = __high2float(y);
return __floats2bfloat162_rn(fmaxf(fxl, fyl), fmaxf(fxh, fyh));
#endif
}
__device__ __nv_bfloat16 operator()(const __nv_bfloat16 x, const __nv_bfloat16 y) const {
#if __CUDA_ARCH__ >= 800
return __hmax(x, y);
#else
float fx, fy;
fx = __bfloat162float(x);
fy = __bfloat162float(y);
return __float2bfloat16(fmaxf(fx, fy));
#endif
}
};
#endif
template<>
struct FuncMin<half> {
__device__ half2 operator()(const half2 x, const half2 y) const {
@@ -299,4 +390,269 @@ struct FuncMin<half> {
return __float2half(fm);
}
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<>
struct FuncMin<__nv_bfloat16> {
__device__ __nv_bfloat162 operator()(const __nv_bfloat162 x, const __nv_bfloat162 y) const {
#if __CUDA_ARCH__ >= 800
return __hmin2(x, y);
#else
float fxl, fxh, fyl, fyh;
fxl = __low2float(x);
fxh = __high2float(x);
fyl = __low2float(y);
fyh = __high2float(y);
return __floats2bfloat162_rn(fminf(fxl, fyl), fminf(fxh, fyh));
#endif
}
__device__ __nv_bfloat16 operator()(const __nv_bfloat16 x, const __nv_bfloat16 y) const {
#if __CUDA_ARCH__ >= 800
return __hmin(x, y);
#else
float fx, fy;
fx = __bfloat162float(x);
fy = __bfloat162float(y);
return __float2bfloat16(fminf(fx, fy));
#endif
}
};
#endif
template<>
struct FuncMax<float> {
__device__ float operator()(float x, float y) const {
return fmaxf(x, y);
}
};
template<>
struct FuncMin<float> {
__device__ float operator()(float x, float y) const {
return fminf(x, y);
}
};
template<>
struct FuncMax<double> {
__device__ double operator()(double x, double y) const {
return fmax(x, y);
}
};
template<>
struct FuncMin<double> {
__device__ double operator()(double x, double y) const {
return fmin(x, y);
}
};
template<typename T>
struct FuncAvg: FuncSum<T> {
static_assert(!std::is_floating_point<T>::value, "Uhoh");
static constexpr bool IsPreOpIdentity = true;
static constexpr bool IsPostOpIdentity = false;
int n;
template<typename ...Arg>
__device__ FuncAvg(int n): n(n) {}
__device__ T preOp(T x) const {
return x;
}
__device__ T postOp(T x) const {
return T(x/n);
}
};
template<>
struct FuncAvg<double>: FuncSum<double> {
static constexpr bool IsPreOpIdentity = false;
static constexpr bool IsPostOpIdentity = true;
double rcp;
__device__ FuncAvg(int n) {
rcp = __drcp_rn(double(n));
}
// inherits FuncSum::operator()
__device__ double preOp(double x) const {
return IsPreOpIdentity ? x : x*rcp;
}
__device__ double postOp(double x) const {
return IsPostOpIdentity ? x : x*rcp;
}
};
template<>
struct FuncAvg<float>: FuncSum<float> {
static constexpr bool IsPreOpIdentity = false;
static constexpr bool IsPostOpIdentity = true;
float rcp;
__device__ FuncAvg(int n) {
rcp = __frcp_rn(float(n));
}
// inherits FuncSum::operator()
__device__ float preOp(float x) const {
return IsPreOpIdentity ? x : x*rcp;
}
__device__ float postOp(float x) const {
return IsPostOpIdentity ? x : x*rcp;
}
};
template<>
struct FuncAvg<half>: FuncSum<half> {
// Change these to switch between all prescale, all postscale, or both by sqrt(N).
// Obviously, the only invalid combination is both true. An improvement would be
// make this parameterized as a build time setting and passed here through
// preprocessor definitions.
static constexpr bool IsPreOpIdentity = false;
static constexpr bool IsPostOpIdentity = true;
#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
half2 scale;
__device__ FuncAvg(int n) {
if (!IsPreOpIdentity && !IsPostOpIdentity)
scale.x = __float2half(__frsqrt_rn(float(n)));
else
scale.x = __float2half(__frcp_rn(float(n)));
scale.y = scale.x;
}
// inherits FuncSum::operator()
__device__ half preOp(half x) const {
return IsPreOpIdentity ? x : __hmul(x, scale.x);
}
__device__ half2 preOp(half2 x) const {
return IsPreOpIdentity ? x : __hmul2(x, scale);
}
__device__ half postOp(half x) const {
return IsPostOpIdentity ? x : __hmul(x, scale.x);
}
__device__ half2 postOp(half2 x) const {
return IsPostOpIdentity ? x : __hmul2(x, scale);
}
#else
float scale;
__device__ FuncAvg(int n) {
if (!IsPreOpIdentity && !IsPostOpIdentity)
scale = __frsqrt_rn(float(n));
else
scale = __frcp_rn(float(n));
}
// inherits FuncSum::operator()
__device__ half preOp(half x) const {
return IsPreOpIdentity ? x : __float2half(__half2float(x)*scale);
}
__device__ half2 preOp(half2 x) const {
if (IsPreOpIdentity)
return x;
else {
float2 a = __half22float2(x);
a.x *= scale;
a.y *= scale;
return __float22half2_rn(a);
}
}
__device__ half postOp(half x) const {
return IsPostOpIdentity ? x : __float2half(__half2float(x)*scale);
}
__device__ half2 postOp(half2 x) const {
if (IsPostOpIdentity)
return x;
else {
float2 a = __half22float2(x);
a.x *= scale;
a.y *= scale;
return __float22half2_rn(a);
}
}
#endif
};
#if defined(__CUDA_BF16_TYPES_EXIST__)
template<>
struct FuncAvg<__nv_bfloat16>: FuncSum<__nv_bfloat16> {
// Change these to switch between all prescale, all postscale, or both by sqrt(N).
// Obviously, the only invalid combination is both true. An improvement would be
// make this parameterized as a build time setting and passed here through
// preprocessor definitions.
static constexpr bool IsPreOpIdentity = false;
static constexpr bool IsPostOpIdentity = true;
#if __CUDA_ARCH__ >= 800
__nv_bfloat162 scale;
__device__ FuncAvg(int n) {
if (!IsPreOpIdentity && !IsPostOpIdentity)
scale.x = __float2bfloat16(__frsqrt_rn(float(n)));
else
scale.x = __float2bfloat16(__frcp_rn(float(n)));
scale.y = scale.x;
}
// inherits FuncSum::operator()
__device__ __nv_bfloat16 preOp(__nv_bfloat16 x) const {
return IsPreOpIdentity ? x : __hmul(x, scale.x);
}
__device__ __nv_bfloat162 preOp(__nv_bfloat162 x) const {
return IsPreOpIdentity ? x : __hmul2(x, scale);
}
__device__ __nv_bfloat16 postOp(__nv_bfloat16 x) const {
return IsPostOpIdentity ? x : __hmul(x, scale.x);
}
__device__ __nv_bfloat162 postOp(__nv_bfloat162 x) const {
return IsPostOpIdentity ? x : __hmul2(x, scale);
}
#else
float scale;
__device__ FuncAvg(int n) {
if (!IsPreOpIdentity && !IsPostOpIdentity)
scale = __frsqrt_rn(float(n));
else
scale = __frcp_rn(float(n));
}
// inherits FuncSum::operator()
__device__ __nv_bfloat16 preOp(__nv_bfloat16 x) const {
return IsPreOpIdentity ? x : __float2bfloat16(__bfloat162float(x)*scale);
}
__device__ __nv_bfloat162 preOp(__nv_bfloat162 x) const {
if (IsPreOpIdentity)
return x;
else {
float fxl, fxh;
fxl = __low2float(x);
fxh = __high2float(x);
return __floats2bfloat162_rn(fxl * scale, fxh * scale);
}
}
__device__ __nv_bfloat16 postOp(__nv_bfloat16 x) const {
return IsPostOpIdentity ? x : __float2bfloat16(__bfloat162float(x)*scale);
}
__device__ __nv_bfloat162 postOp(__nv_bfloat162 x) const {
if (IsPostOpIdentity)
return x;
else {
float fxl, fxh;
fxl = __low2float(x);
fxh = __high2float(x);
return __floats2bfloat162_rn(fxl * scale, fxh * scale);
}
}
#endif
};
#endif
template<typename T>
struct FuncTraits<FuncAvg<T>> {
static constexpr bool IsPreOpIdentity = FuncAvg<T>::IsPreOpIdentity;
static constexpr bool IsPostOpIdentity = FuncAvg<T>::IsPostOpIdentity;
__device__ static FuncAvg<T> make(int rankN) {
return FuncAvg<T>(rankN);
}
template<typename U>
__device__ static U preOp(FuncAvg<T> fn, U x) {
return fn.preOp(x);
}
template<typename U>
__device__ static U postOp(FuncAvg<T> fn, U x) {
return fn.postOp(x);
}
};
#endif // REDUCE_KERNEL_H_
src/collectives/device/reduce_scatter.h
+65 -172
View File
@@ -5,192 +5,85 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduceScatter, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads-WARP_SIZE;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * REDUCESCATTER_CHUNKSTEPS;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
const ssize_t size = args->coll.count;
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
ncclRing *ring = &ncclShmem.channel.ring;
int const *ringRanks = ring->devUserRanks;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCESCATTER_CHUNKSTEPS : 1));
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2);
const int nranks = ncclShmem.comm.nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff);
ncclPrimitives<UNROLL, REDUCESCATTER_CHUNKSTEPS/REDUCESCATTER_SLICESTEPS, REDUCESCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nChannels));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t chunkOffset = gridOffset + bid*realChunkSize;
/////////////// begin ReduceScatter steps ///////////////
ssize_t offset;
int nelem = min(realChunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[nranks-1];
offset = chunkOffset + rankDest * size;
prims.send(thisInput+offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
prims.recvReduceSend(thisInput+offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final result
rankDest = ring->devUserRanks[0];
offset = chunkOffset + rankDest * size;
prims.recvReduceCopy(thisInput+offset, thisOutput+chunkOffset, nelem);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels));
realChunkSize = roundUp(realChunkSize, (nthreads-WARP_SIZE)*sizeof(uint64_t)/sizeof(T));
}
}
};
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128, chunkSize);
realChunkSize = int(realChunkSize);
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduceScatter, NCCL_ALGO_RING, NCCL_PROTO_LL, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepLines = comm->buffSizes[NCCL_PROTO_LL] / (sizeof(union ncclLLFifoLine)*NCCL_STEPS);
ssize_t chunkSize = stepLines * sizeof(uint64_t) / sizeof(T);
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
ssize_t chunkOffset = gridOffset + bid*int(realChunkSize);
ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepLines, channel, comm);
/////////////// begin ReduceScatter steps ///////////////
ssize_t offset;
int nelem = min(realChunkSize, size-chunkOffset);
int rankDest;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
// step 0: push data to next GPU
rankDest = ringRanks[nranks-1];
offset = chunkOffset + rankDest * size;
prims.send(offset, nelem);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
if (size-gridOffset < loopSize) {
chunkSize = args->coll.lastChunkSize;
}
ssize_t chunkOffset = gridOffset + bid*chunkSize;
/////////////// begin ReduceScatter steps ///////////////
ssize_t offset;
int nelem = min(chunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[nranks-1];
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
rankDest = ringRanks[nranks-j];
offset = chunkOffset + rankDest * size;
LLprims.send(thisInput+offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
LLprims.recvReduceSend(thisInput+offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final
// result that we store in this data
rankDest = ring->devUserRanks[0];
offset = chunkOffset + rankDest * size;
LLprims.recvReduceCopy(thisInput+offset, thisOutput+chunkOffset, nelem);
prims.recvReduceSend(offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final result
rankDest = ringRanks[0];
offset = chunkOffset + rankDest * size;
prims.recvReduceCopy(offset, chunkOffset, nelem, /*postOp=*/true);
}
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWorkElem *args) {
using Proto = ProtoSimple<REDUCESCATTER_CHUNKSTEPS/REDUCESCATTER_SLICESTEPS, REDUCESCATTER_SLICESTEPS>;
runRing<T, RedOp, Proto>(args);
}
};
#include "prims_ll128.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncReduceScatter, NCCL_ALGO_RING, NCCL_PROTO_LL128, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const int stepSize = comm->buffSizes[NCCL_PROTO_LL128] / (sizeof(uint64_t)*NCCL_STEPS);
ssize_t chunkSize = stepSize*NCCL_LL128_DATAELEMS*sizeof(uint64_t) / (NCCL_LL128_LINEELEMS*sizeof(T));
// We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
const ssize_t minChunkSize = (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*nthreads*NCCL_LL128_DATAELEMS*sizeof(uint64_t))/(NCCL_LL128_LINEELEMS*sizeof(T))/2;
const int nranks = comm->nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, stepSize, channel, comm);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, nChannels*minChunkSize)*minChunkSize, chunkSize);
ssize_t chunkOffset = gridOffset + bid*chunkSize;
/////////////// begin ReduceScatter steps ///////////////
ssize_t offset;
int nelem = min(chunkSize, size-chunkOffset);
int rankDest;
// step 0: push data to next GPU
rankDest = ring->devUserRanks[nranks-1];
offset = chunkOffset + rankDest * size;
LLprims.send(thisInput+offset, nelem);
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
rankDest = ring->devUserRanks[nranks-j];
offset = chunkOffset + rankDest * size;
LLprims.recvReduceSend(thisInput+offset, nelem);
}
// step k-1: reduce this buffer and data, which will produce the final
// result that we store in this data
rankDest = ring->devUserRanks[0];
offset = chunkOffset + rankDest * size;
LLprims.recvReduceCopy(thisInput+offset, thisOutput+chunkOffset, nelem);
}
}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL>(args);
}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncReduceScatter, NCCL_ALGO_TREE, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
};
template<int PROTO, class REDOP, typename T, int UNROLL>
class ncclFunction<ncclFuncReduceScatter, NCCL_ALGO_COLLNET, PROTO, REDOP, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* args) {}
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ void run(ncclWorkElem *args) {
runRing<T, RedOp, ProtoLL128>(args);
}
};
src/collectives/device/sendrecv.h
+72 -74
View File
@@ -5,89 +5,87 @@
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
#include "primitives.h"
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncSendRecv, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ void run(struct ncclWorkElem* firstArgs) {
struct ncclWorkElem* args = firstArgs;
int tid = threadIdx.x;
int group = 0;
for (int s=0; s<NCCL_MAX_WORK_ELEMENTS; s++) {
int nThreadsSegment = args->p2p.nThreads;
if (nThreadsSegment == 0) return; // Nothing else to do
int groupRecv = group;
group += 1;
int groupSend = group;
group += nThreadsSegment > 128 ? 2 : 1;
if (tid < nThreadsSegment) {
const int nThreads = nThreadsSegment > 128 ? nThreadsSegment-WARP_SIZE : nThreadsSegment;
template<typename T, typename RedOp>
struct RunWork<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ void run(ncclWork *work) {
int tid = threadIdx.x;
int group = 0;
const int rank = ncclShmem.comm.rank;
const int nRanks = ncclShmem.comm.nRanks;
using Proto = ProtoSimple<1, 1>;
// Compute pointers
const T* sendbuff = (const T*)args->sendbuff;
T* recvbuff = (T*)args->recvbuff;
const ssize_t sendCount = args->p2p.sendCount;
const ssize_t recvCount = args->p2p.recvCount;
for (int s=0; s<NCCL_MAX_WORK_ELEMENTS; s++) {
ncclWorkElem *args = &work->elems[s];
int nThreadsSegment = args->p2p.nThreads;
if (args->active == 0 || nThreadsSegment == 0) break;
const int delta = args->p2p.delta;
if (delta == 0) {
if (tid < nThreads && sendbuff != recvbuff) {
// local copy : ReduceOrCopyMulti takes an int as number of elements,
// so we split it in blocks of 1G elements.
int blockSize = 1<<30;
for (size_t offset=0; offset<sendCount; offset += blockSize) {
size_t remaining = sendCount - offset;
if (remaining < blockSize) blockSize = remaining;
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nThreads, 1, &sendbuff, 1, &recvbuff, blockSize);
sendbuff += blockSize; recvbuff += blockSize;
}
}
} else {
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
int nThreadsSplit = (nThreadsSegment - (nThreadsSegment > 128 ? WARP_SIZE : 0))/2;
int groupRecv = group;
group += Proto::calcGroupWidth(/*send=*/false, nThreadsSplit);
int groupSend = group;
group += Proto::calcGroupWidth(/*send=*/true, nThreadsSegment - nThreadsSplit);
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE]/(sizeof(T)*NCCL_STEPS);
if (tid < nThreadsSegment) {
// Compute pointers
T const* sendbuff = (const T*)args->sendbuff;
T* recvbuff = (T*)args->recvbuff;
ssize_t const sendCount = args->p2p.sendCount;
ssize_t const recvCount = args->p2p.recvCount;
int const delta = args->p2p.delta;
int nThreadsSplit = nThreads/2;
if ((tid < nThreadsSplit) && recvCount >= 0) {
const int chunkSize = args->p2p.recvChunkSize/sizeof(T);
int peer = (comm->rank-delta+comm->nRanks)%comm->nRanks;
int nt = nThreadsSplit;
ncclPrimitives<UNROLL, 1, 1, T, 1, 0, 1, FUNC>
prims(tid, nt, &peer, NULL, recvbuff, stepSize, channel, comm, ncclShmem->ptrs, groupRecv);
if (recvCount == 0) {
prims.recv(recvbuff, 0);
} else for (ssize_t offset = 0; offset < recvCount; offset += chunkSize) {
int realChunkSize = min(chunkSize, recvCount-offset);
ALIGN_SIZE(realChunkSize, nt*sizeof(uint64_t)/sizeof(T));
int nelem = min(realChunkSize, recvCount-offset);
prims.directRecv(recvbuff+offset, offset, nelem);
}
}
if ((tid >= nThreadsSplit) && sendCount >= 0) {
const int chunkSize = args->p2p.sendChunkSize/sizeof(T);
int peer = (comm->rank+delta)%comm->nRanks;
int nt = nThreads-nThreadsSplit;
ncclPrimitives<UNROLL, 1, 1, T, 0, 1, 1, FUNC>
prims(tid-nThreadsSplit, nt, NULL, &peer, recvbuff, stepSize, channel, comm, ncclShmem->ptrs, groupSend);
if (sendCount == 0) {
prims.send(sendbuff, 0);
} else for (ssize_t offset = 0; offset < sendCount; offset += chunkSize) {
int realChunkSize = min(chunkSize, sendCount-offset);
ALIGN_SIZE(realChunkSize, nt*sizeof(uint64_t)/sizeof(T));
int nelem = min(realChunkSize, sendCount-offset);
prims.directSend(sendbuff+offset, offset, nelem);
}
if (delta == 0) {
if (sendbuff != recvbuff) {
// local copy : ReduceOrCopyMulti takes an int as number of elements,
// so we split it in blocks of 1G elements.
int blockSize = 1<<30;
for (size_t offset=0; offset<sendCount; offset += blockSize) {
size_t remaining = sendCount - offset;
if (remaining < blockSize) blockSize = remaining;
ReduceOrCopyMulti<COLL_UNROLL, RedOp, T, 1, 1, 1, 1>(tid, nThreadsSegment, RedOp(), false, false, 1, &sendbuff, 1, &recvbuff, blockSize);
sendbuff += blockSize;
recvbuff += blockSize;
}
}
}
tid -= nThreadsSegment;
if (tid < 0) return;
args++;
else {
if ((tid < nThreadsSplit) && recvCount >= 0) {
int const peer = (rank - delta + nRanks)%nRanks;
int const t0 = 0;
int const nt = nThreadsSplit;
int const chunkSize = args->p2p.recvChunkSize/sizeof(T);
Primitives<T, RedOp, FanAsymmetric<1, 0>, 1, Proto> prims
(tid-t0, nt, &peer, nullptr, nullptr, recvbuff, groupRecv);
ssize_t offset = 0;
do {
int nelem = roundUp(chunkSize, nt*(sizeof(uint64_t)/sizeof(T)));
nelem = min(chunkSize, recvCount-offset);
prims.directRecv(offset, nelem);
offset += nelem;
} while(offset < recvCount);
}
if ((tid >= nThreadsSplit) && sendCount >= 0) {
int const peer = (rank + delta)%nRanks;
int const t0 = nThreadsSplit;
int const nt = nThreadsSegment - nThreadsSplit;
int const chunkSize = args->p2p.sendChunkSize/sizeof(T);
Primitives<T, RedOp, FanAsymmetric<0, 1>, 1, Proto> prims
(tid-t0, nt, nullptr, &peer, sendbuff, nullptr, groupSend);
ssize_t offset = 0;
do {
int nelem = roundUp(chunkSize, nt*(sizeof(uint64_t)/sizeof(T)));
nelem = min(chunkSize, sendCount-offset);
prims.directSend(offset, offset, nelem);
offset += nelem;
} while(offset < sendCount);
}
}
break;
}
tid -= nThreadsSegment;
}
}
};