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

Merge pull request #223 from wenkaidu/sendrecv

Browse Source 
Use separate threads for send and receive

[ROCm/rccl commit: a144a85465]
This commit is contained in:
Wenkai Du
2020-06-30 10:50:06 -07:00
committed by GitHub
parent 117f0f6e6e ac87c9db37
commit effafd456b
17 changed files with 130 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/collectives/device/all_to_all.h
+33 -9
View File
@@ -48,16 +48,40 @@ __device__ void ncclAllToAllKernel(struct CollectiveArgs* args) {
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem);
}
}
else {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
ncclPrimitives<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm, args->opCount);
}
}
ssize_t send_offset = chunkOffset + peerSend*size;
ssize_t recv_offset = chunkOffset + peerRecv*size;
prims.send(thisInput+send_offset, nelem);
prims.recv(thisOutput+recv_offset, nelem);
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
if ((blockIdx.x*peersPerChan+i)%nranks != 0) {
int nthreadsSplit = nthreads/2;
int peerNone[2] = {-1,-1};
if (tid < nthreadsSplit ) {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
ncclPrimitives<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 2, 1, 0, FUNC>
prims(tid, nthreadsSplit, peerNone, &peerSend, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1)));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t chunkOffset = gridOffset + (peersPerChan == 1 ? (bid/nranks)*realChunkSize : 0);
int nelem = min(realChunkSize, size-chunkOffset);
ssize_t send_offset = chunkOffset + peerSend*size;
prims.send(thisInput+send_offset, nelem);
}
} else {
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
ncclPrimitives<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 1, 2, 0, FUNC>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peerRecv, peerNone, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1)));
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t chunkOffset = gridOffset + (peersPerChan == 1 ? (bid/nranks)*realChunkSize : 0);
int nelem = min(realChunkSize, size-chunkOffset);
ssize_t recv_offset = chunkOffset + peerRecv*size;
prims.recv(thisOutput+recv_offset, nelem);
}
}
}
}
projects/rccl/src/collectives/device/common.h
+9 -1
View File
@@ -246,11 +246,19 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclDevComm* comm) { \
ALLOCATE_SHMEM; \
__shared__ struct ncclColl localColl; \
__shared__ uint32_t abortCount; \
__shared__ uint64_t barrier[MAXBARRIERS]; \
__shared__ uint64_t barrier_next[MAXBARRIERS*MAXWARPS]; \
if (tid == 0) abortCount = 0; \
__syncthreads(); \
\
struct ncclChannel* channel = comm->channels+bid; \
channel->sync = sync; \
if (tid == 0) { \
channel->sync = sync; \
channel->barrier = barrier; \
channel->barrier_next = barrier_next; \
for (auto i = 0; i < MAXBARRIERS; i++) barrier[i] = 0; \
for (auto i = 0; i < MAXBARRIERS*MAXWARPS; i++) barrier_next[i] = 0; \
} \
if (!load_coll(&localColl, channel->collectives+channel->collFifoHead, tid, comm, &abortCount)) { \
if (tid == 0) traceAbort(-1); \
return; \
projects/rccl/src/collectives/device/primitives.h
+18 -1
View File
@@ -32,6 +32,13 @@
} \
} while (0)
#define barrier_by_id(id) do { \
const int w = threadIdx.x/WARP_SIZE; \
barrier_next[id*MAXWARPS+w] += nthreads/WARP_SIZE; \
__atomic_fetch_add(barriers+id, 1, __ATOMIC_SEQ_CST); \
while (LOAD(barriers+id) < barrier_next[id*MAXWARPS+w]) /* spin */; \
} while (0)
template <typename T, int NRECV>
class ncclPrimitivesRecvData {
public:
@@ -90,9 +97,17 @@ class ncclPrimitives {
inline __device__ const T* recvPtr(int i) { return ((const T*)r.recvBuff[i])+recvOffset(i); }
inline __device__ T* sendPtr(int i) { return ((T*)s.sendBuff[i])+sendOffset(i); }
uint64_t* barriers;
uint64_t* barrier_next;
inline __device__ void barrier() {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
if (nthreads == gridDim.x) {
__syncthreads();
} else if (wid == 0) {
if (NRECV < NSEND) barrier_by_id(0);
else barrier_by_id(1);
}
#else
asm volatile ("bar.sync 1, %0;" :: "r"(nthreads+WARP_SIZE));
#endif
@@ -359,6 +374,8 @@ inline __device__ int directSendInc(int i, int directInc, int sliceInc) {
}
inline __device__ void init(int* recvPeers, int* sendPeers, struct ncclChannel* channel) {
barriers = channel->barrier;
barrier_next = channel->barrier_next;
// Make sure step is updated before we read it.
barrier();
projects/rccl/src/collectives/device/sendrecv.h
+48 -28
View File
@@ -22,8 +22,15 @@ __device__ void ncclSendRecvKernel(struct CollectiveArgs* args) {
if (args->p2p.delta == 0) {
if (tid < nthreads && sendbuff != recvbuff) {
// local copy
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, args->p2p.sendCount);
// 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<args->p2p.sendCount; offset += blockSize) {
size_t remaining = args->p2p.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;
}
}
return;
}
@@ -31,40 +38,53 @@ __device__ void ncclSendRecvKernel(struct CollectiveArgs* args) {
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
const ssize_t sendSize = args->p2p.sendCount;
const ssize_t recvSize = args->p2p.recvCount;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize;
int peerRecv = recvSize >= 0 ? (comm->rank-(int)args->p2p.delta+comm->nRanks)%comm->nRanks : -1;
int peerSend = sendSize >= 0 ? (comm->rank+(int)args->p2p.delta)%comm->nRanks : -1;
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE]/(sizeof(T)*NCCL_STEPS)/SENDRECV_SLICEFACTOR;
ncclPrimitives<UNROLL, 1, 1, T, 1, 1, 1, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm, args->opCount);
int maxSize = sendSize-chunkSize>recvSize ? sendSize-chunkSize : recvSize;
if (sendSize >= 0) {
int realChunkSize = min(chunkSize, sendSize);
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
int nelem = min(realChunkSize, sendSize);
prims.send(sendbuff, nelem);
int nthreadsSplit;
const int64_t sendCount = static_cast<int64_t>(args->p2p.sendCount);
const int64_t recvCount = static_cast<int64_t>(args->p2p.recvCount);
if (sendCount >= 0 && recvCount >= 0)
nthreadsSplit = nthreads/2;
else {
if (sendCount >= 0) nthreadsSplit = nthreads;
else nthreadsSplit = 0;
}
// We set NRECV or NSEND to 2 to use different barriers in primitives for the send threads and
// receive threads, but then we define all peers to -1 since sender threads don't receive and
// receive threads don't send.
int peerNone[2] = {-1,-1};
for (ssize_t gridOffset = 0; gridOffset < maxSize; gridOffset += chunkSize) {
if (gridOffset+chunkSize < sendSize) {
int realChunkSize = min(chunkSize, sendSize-gridOffset-chunkSize);
if (tid < nthreadsSplit ) {
const ssize_t sendSize = args->p2p.sendCount;
if (sendSize < 0) return;
int peer = (comm->rank+(int)args->p2p.delta)%comm->nRanks;
ncclPrimitives<UNROLL, 1, 1, T, 2, 1, 1, FUNC>
prims(tid, nthreadsSplit, peerNone, &peer, recvbuff, stepSize*SENDRECV_SLICEFACTOR, channel, comm, args->opCount);
if (sendSize == 0) {
prims.send(sendbuff, 0);
} else for (ssize_t offset = 0; offset < sendSize; offset += stepSize) {
int realChunkSize = min(stepSize, sendSize-offset);
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t offset = gridOffset + chunkSize;
int nelem = min(realChunkSize, sendSize-offset);
prims.send(sendbuff+offset, nelem);
prims.directSend(sendbuff+offset, offset, nelem);
}
if (gridOffset < recvSize) {
int realChunkSize = min(chunkSize, recvSize-gridOffset);
} else {
const ssize_t recvSize = args->p2p.recvCount;
if (recvSize < 0) return;
int peer = (comm->rank-(int)args->p2p.delta+comm->nRanks)%comm->nRanks;
ncclPrimitives<UNROLL, 1, 1, T, 1, 2, 1, FUNC>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peer, peerNone, recvbuff, stepSize*SENDRECV_SLICEFACTOR, channel, comm, args->opCount);
if (recvSize == 0) {
prims.recv(recvbuff, 0);
} else for (ssize_t offset = 0; offset < recvSize; offset += stepSize) {
int realChunkSize = min(stepSize, recvSize-offset);
ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
ssize_t offset = gridOffset;
int nelem = min(realChunkSize, recvSize-offset);
prims.recv(recvbuff+offset, nelem);
prims.directRecv(recvbuff+offset, offset, nelem);
}
}
if (recvSize == 0) prims.recv(recvbuff,0);
}
projects/rccl/src/include/collectives.h
+1 -1
View File
@@ -82,7 +82,7 @@ DECL_ALL_COLLS
#define BROADCAST_CHUNKSTEPS 1
#define REDUCE_SLICESTEPS 1
#define REDUCE_CHUNKSTEPS 1
#define SENDRECV_SLICEFACTOR 4
#define SENDRECV_SLICEFACTOR 1
#define GATHER_SLICESTEPS 4
#define GATHER_CHUNKSTEPS 4
#define SCATTER_SLICESTEPS 4
projects/rccl/src/include/devcomm.h
+5
View File
@@ -89,6 +89,9 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
#define NCCL_DIRECT_GPU 0x01
#define NCCL_DIRECT_NIC 0x10
#define MAXBARRIERS 2
#define MAXWARPS (NCCL_MAX_NTHREADS/WARP_SIZE)
struct ncclConnInfo {
// Regular comm mechanism
char *buffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send
@@ -220,6 +223,8 @@ struct ncclChannel {
int collFifoTail; // Only used by CPU
uint32_t* sync;
uint64_t* barrier;
uint64_t* barrier_next;
#ifdef ENABLE_PROFILING
struct timeval tvs;
uint64_t sizes;
projects/rccl/src/proxy.cc
+2 -2
View File
@@ -152,13 +152,13 @@ ncclResult_t ncclProxySaveP2p(struct ncclInfo* info, struct ncclChannel* channel
args.dtype = info->datatype;
if (info->delta > 0 && info->sendbytes >= 0) {
int peersend = (info->comm->rank+info->delta)%info->comm->nRanks;
args.nsteps = DIVUP(info->sendbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS);
args.nsteps = DIVUP(info->sendbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/SENDRECV_SLICEFACTOR);
if (args.nsteps == 0) args.nsteps = 1;
NCCLCHECK(SaveProxy<proxySend>(peersend, &args));
}
if (info->delta > 0 && info->recvbytes >= 0) {
int peerrecv = (info->comm->nRanks+info->comm->rank-info->delta)%info->comm->nRanks;
args.nsteps = DIVUP(info->recvbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS);
args.nsteps = DIVUP(info->recvbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/SENDRECV_SLICEFACTOR);
if (args.nsteps == 0) args.nsteps = 1;
NCCLCHECK(SaveProxy<proxyRecv>(peerrecv, &args));
}
projects/rccl/test/test_AllGather.cpp
+2 -2
View File
@@ -105,9 +105,9 @@ namespace CorrectnessTests
ncclFloat64,
ncclBfloat16),
// Number of elements
testing::Values(3072, 3145728),
testing::Values(2520, 3026520),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_AllReduce.cpp
+1 -1
View File
@@ -55,7 +55,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_AllToAll.cpp
+1 -1
View File
@@ -59,7 +59,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false),
testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1")));
projects/rccl/test/test_Broadcast.cpp
+1 -1
View File
@@ -63,7 +63,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_CombinedCalls.cpp
+2 -2
View File
@@ -91,9 +91,9 @@ namespace CorrectnessTests
ncclFloat64,
ncclBfloat16),
// Number of elements
testing::Values(3072, 3145728),
testing::Values(2520, 3026520),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_Gather.cpp
+1 -1
View File
@@ -63,7 +63,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false),
testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1")));
projects/rccl/test/test_GroupCalls.cpp
+2 -2
View File
@@ -111,9 +111,9 @@ namespace CorrectnessTests
ncclFloat64,
ncclBfloat16),
// Number of elements
testing::Values(3072, 3145728),
testing::Values(2520, 3026520),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_Reduce.cpp
+1 -1
View File
@@ -63,7 +63,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_ReduceScatter.cpp
+2 -2
View File
@@ -59,9 +59,9 @@ namespace CorrectnessTests
ncclFloat64,
ncclBfloat16),
// Number of elements
testing::Values(3072, 3145728),
testing::Values(2520, 3026520),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false, true),
testing::Values("")));
projects/rccl/test/test_Scatter.cpp
+1 -1
View File
@@ -63,7 +63,7 @@ namespace CorrectnessTests
// Number of elements
testing::Values(1024, 1048576),
// Number of devices
testing::Values(2,3,4),
testing::Values(2,3,4,5,6,7,8),
// In-place or not
testing::Values(false),
testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1")));