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

Port alltoall[v]

Browse Source
This commit is contained in:
Wenkai Du
2020-12-04 18:52:32 -05:00
parent 105db19a11
commit f4d5d3d620
25 changed files with 332 additions and 383 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/device/all_to_all.cu
+1 -2
View File
@@ -8,5 +8,4 @@
#include "common.h"
#include "collectives.h"
IMPL_COLL_FUNC(ncclAllToAll, copy, FuncSum, i8, int8_t);
IMPL_COLL_KERN(ncclAllToAll, copy, FuncSum, i8, int8_t, 0);
IMPL_COLL_FUNC(AllToAll, RING, SIMPLE, Sum, int8_t);
src/collectives/device/all_to_all.h
+68 -67
View File
@@ -9,80 +9,81 @@
#include "primitives.h"
#include "collectives.h"
template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllToAllKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = args->coll.nThreads;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t size = args->coll.count;
const int nranks = comm->nRanks;
const int bid = args->coll.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLTOALL_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllToAll, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ __attribute__((noinline)) void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t size = args->coll.count;
const int nranks = comm->nRanks;
const int bid = args->coll.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLTOALL_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
// 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) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0) {
if (tid < nthreads && thisInput != thisOutput) {
const T* sendbuff = thisInput+chunkOffset+rank*size;
T* recvbuff = thisOutput+chunkOffset+rank*size;
// local copy
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem);
}
}
}
}
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0) {
if (tid < nthreads && thisInput != thisOutput) {
const T* sendbuff = thisInput+chunkOffset+rank*size;
T* recvbuff = thisOutput+chunkOffset+rank*size;
// local copy
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem);
if ((blockIdx.x*peersPerChan+i)%nranks != 0) {
int nthreadsSplit = nthreads/2;
if (tid < nthreadsSplit ) {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
ncclPrimitives<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 0, 1, 0, FUNC>
prims(tid, nthreadsSplit, NULL, &peerSend, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
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, 0, 0, FUNC>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peerRecv, NULL, NULL, stepSize, channel, comm, ncclShmem->ptrs, 1);
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);
}
}
}
}
}
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);
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);
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);
}
}
}
}
}
};
src/collectives/device/all_to_allv.cu
Executable → Regular
+1 -2
View File
@@ -8,5 +8,4 @@
#include "common.h"
#include "collectives.h"
IMPL_COLL_FUNC(ncclAllToAllv, copy, FuncSum, i8, int8_t);
IMPL_COLL_KERN(ncclAllToAllv, copy, FuncSum, i8, int8_t, 0);
IMPL_COLL_FUNC(AllToAllv, RING, SIMPLE, Sum, int8_t);
src/collectives/device/all_to_allv.h
Executable → Regular
+76 -74
View File
@@ -9,87 +9,89 @@
#include "primitives.h"
#include "collectives.h"
template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllToAllvKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = args->a2av.nThreads;
const int nChannels = args->a2av.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t typesize = args->a2av.count;
const int nranks = comm->nRanks;
const int bid = args->a2av.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLTOALLV_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
template<class FUNC, typename T, int UNROLL>
class ncclFunction<ncclFuncAllToAllv, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE, FUNC, T, UNROLL> {
public:
__device__ __attribute__((noinline)) void run(struct ncclWorkElem* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int nChannels = args->a2av.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t typesize = args->a2av.count;
const int nranks = comm->nRanks;
const int bid = args->a2av.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLTOALLV_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->sendbuff;
T * __restrict__ thisOutput = (T*)args->recvbuff;
size_t* params = channel->a2avParams + nranks*4*args->index;
size_t *sendcounts = params;
size_t *sdispls = params + nranks;
size_t *recvcounts = params + nranks*2;
size_t *rdispls = params + nranks*3;
ssize_t size = sendcounts[rank]*typesize;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0) {
if (tid < nthreads && thisInput != thisOutput) {
const T* sendbuff = thisInput+chunkOffset+sdispls[rank]*typesize;
T* recvbuff = thisOutput+chunkOffset+rdispls[rank]*typesize;
// local copy
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem);
}
}
}
}
size_t *sendcounts = args->a2av.extra;
size_t *sdispls = args->a2av.extra + nranks;
size_t *recvcounts = args->a2av.extra + nranks*2;
size_t *rdispls = args->a2av.extra + nranks*3;
ssize_t size = sendcounts[rank]*typesize;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0) {
if (tid < nthreads && thisInput != thisOutput) {
const T* sendbuff = thisInput+chunkOffset+sdispls[rank]*typesize;
T* recvbuff = thisOutput+chunkOffset+rdispls[rank]*typesize;
// local copy
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem);
if ((blockIdx.x*peersPerChan+i)%nranks != 0) {
int nthreadsSplit = nthreads/2;
if (tid < nthreadsSplit ) {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
ncclPrimitives<UNROLL, ALLTOALLV_CHUNKSTEPS/ALLTOALLV_SLICESTEPS, ALLTOALLV_SLICESTEPS, T, 0, 1, 0, FUNC>
prims(tid, nthreadsSplit, NULL, &peerSend, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
size = sendcounts[peerSend]*typesize;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1)));
ALIGN_SIZE(realChunkSize, nthreadsSplit*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 + sdispls[peerSend]*typesize;
prims.send(thisInput+send_offset, nelem);
}
} else {
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
ncclPrimitives<UNROLL, ALLTOALLV_CHUNKSTEPS/ALLTOALLV_SLICESTEPS, ALLTOALLV_SLICESTEPS, T, 1, 0, 0, FUNC>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peerRecv, NULL, NULL, stepSize, channel, comm, ncclShmem->ptrs, 1);
size = recvcounts[peerRecv]*typesize;
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-nthreadsSplit)*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 + rdispls[peerRecv]*typesize;
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, ALLTOALLV_CHUNKSTEPS/ALLTOALLV_SLICESTEPS, ALLTOALLV_SLICESTEPS, T, 2, 1, 0, FUNC>
prims(tid, nthreadsSplit, peerNone, &peerSend, NULL, stepSize, channel, comm);
size = sendcounts[peerSend]*typesize;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1)));
ALIGN_SIZE(realChunkSize, nthreadsSplit*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 + sdispls[peerSend]*typesize;
prims.send(thisInput+send_offset, nelem);
}
} else {
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
ncclPrimitives<UNROLL, ALLTOALLV_CHUNKSTEPS/ALLTOALLV_SLICESTEPS, ALLTOALLV_SLICESTEPS, T, 1, 2, 0, FUNC>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peerRecv, peerNone, NULL, stepSize, channel, comm);
size = recvcounts[peerRecv]*typesize;
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-nthreadsSplit)*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 + rdispls[peerRecv]*typesize;
prims.recv(thisOutput+recv_offset, nelem);
}
}
}
}
}
};
src/collectives/device/common.h
+8 -2
View File
@@ -77,7 +77,9 @@
NCCL_FUNCS2B(AllGather), \
NCCL_FUNCS2A(ReduceScatter), \
NCCL_FUNCS2A(AllReduce), \
NCCL_FUNC_NAME(SendRecv, RING, SIMPLE, Sum, int8_t) }
NCCL_FUNC_NAME(SendRecv, RING, SIMPLE, Sum, int8_t), \
NCCL_FUNC_NAME(AllToAll, RING, SIMPLE, Sum, int8_t), \
NCCL_FUNC_NAME(AllToAllv, RING, SIMPLE, Sum, int8_t) }
// Must be consistent with the ncclFuncSet enum
using ncclKernelFunc_t = void (*)(struct ncclWorkElem* args);
@@ -93,6 +95,8 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
NCCL_FUNCS2A(ReduceScatter),
NCCL_FUNCS2A(AllReduce),
NCCL_FUNC_NAME(SendRecv, RING, SIMPLE, Sum, int8_t),
NCCL_FUNC_NAME(AllToAll, RING, SIMPLE, Sum, int8_t),
NCCL_FUNC_NAME(AllToAllv, RING, SIMPLE, Sum, int8_t),
#endif
};
@@ -140,7 +144,9 @@ void NCCL_CALL_FUNCTIONS(struct ncclWorkElem* const c) noexcept {
else ncclFunction_AllGather_COLLNET_SIMPLE_Sum_int8_t(c);
}
else if (c->funcIndex < 1800) Caller<1080, 1800>::call(c);
else ncclFunction_SendRecv_RING_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex == 1800) ncclFunction_SendRecv_RING_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex == 1801) ncclFunction_AllToAll_RING_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex == 1802) ncclFunction_AllToAllv_RING_SIMPLE_Sum_int8_t(c);
}
static __device__ void load_parallel(void* dst, void* src, size_t size, int tid) {
src/collectives/device/gather.cu
-12
View File
@@ -1,12 +0,0 @@
/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "gather.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL_FUNC(ncclGather, copy, FuncSum, i8, int8_t);
IMPL_COLL_KERN(ncclGather, copy, FuncSum, i8, int8_t, 0);
src/collectives/device/gather.h
-74
View File
@@ -1,74 +0,0 @@
/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclGatherKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = args->coll.nThreads;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t size = args->coll.count;
const int nranks = comm->nRanks;
const int bid = args->coll.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * GATHER_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
const int root = args->coll.root;
// 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) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0 && rank == root) {
const T* sendbuff = thisInput+chunkOffset;
T* recvbuff = thisOutput+chunkOffset+rank*size;
if (tid < nthreads && sendbuff != recvbuff) {
// local copy
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;
if (rank == root) {
ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
ssize_t recv_offset = chunkOffset + peerRecv*size;
prims.recv(thisOutput+recv_offset, nelem);
}
else {
if (peerSend == root) {
ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
ssize_t send_offset = chunkOffset;
prims.send(thisInput+send_offset, nelem);
}
}
}
}
}
}
src/collectives/device/scatter.cu
-12
View File
@@ -1,12 +0,0 @@
/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "scatter.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL_FUNC(ncclScatter, copy, FuncSum, i8, int8_t);
IMPL_COLL_KERN(ncclScatter, copy, FuncSum, i8, int8_t, 0);
src/collectives/device/scatter.h
-74
View File
@@ -1,74 +0,0 @@
/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "devcomm.h"
#include "primitives.h"
#include "collectives.h"
template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclScatterKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = args->coll.nThreads;
const int nChannels = args->coll.nChannels;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
const ssize_t size = args->coll.count;
const int nranks = comm->nRanks;
const int bid = args->coll.bid;
const int rank = ring->devUserRanks[0];
const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * SCATTER_CHUNKSTEPS;
const int peersPerChan = DIVUP(nranks, nChannels);
const ssize_t loopSize = (peersPerChan == 1 ? (nChannels/nranks)*(ssize_t)chunkSize : (ssize_t)chunkSize);
const int root = args->coll.root;
// 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) {
for (int i = 0; i < peersPerChan; i++) {
if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) ||
(peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks))
continue;
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);
if ((blockIdx.x*peersPerChan+i)%nranks == 0 && rank == root) {
const T* sendbuff = thisInput+chunkOffset+rank*size;
T* recvbuff = thisOutput+chunkOffset;
if (tid < nthreads && sendbuff != recvbuff) {
// local copy
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;
if (rank == root) {
ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
ssize_t send_offset = chunkOffset + peerSend*size;
prims.send(thisInput+send_offset, nelem);
}
else {
if (peerRecv == root) {
ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
ssize_t recv_offset = chunkOffset;
prims.recv(thisOutput+recv_offset, nelem);
}
}
}
}
}
}