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

Merge remote-tracking branch 'remotes/nccl/master' into rccl_2.5.6

Browse Source
This commit is contained in:
Wenkai Du
2019-11-21 13:41:10 -08:00
parent b1ed4b7fa8 299c554dcc
commit 6648c81dc6
76 changed files with 5119 additions and 3071 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/all_reduce.cc
-1
View File
@@ -6,7 +6,6 @@
************************************************************************/
#include "enqueue.h"
#include "collectives.h"
NCCL_API(ncclResult_t, ncclAllReduce, const void* sendbuff, void* recvbuff, size_t count,
ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream);
src/collectives/collectives.h
-78
View File
@@ -1,78 +0,0 @@
#include "hip/hip_runtime.h"
/*************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#ifndef NCCL_COLLECTIVES_H_
#define NCCL_COLLECTIVES_H_
#define FUNC_INDEX(coll, redop, dtype, ll, al) ((((((coll)*ncclNumOps + (redop))*ncclNumTypes) + (dtype))*2+(al))*2+(ll))
#define NCCL_COLL_NAME(coll, op, dtype) \
coll##_##op##_##dtype
#define NCCL_KERN_NAME(coll, op, dtype) \
coll##Kernel_##op##_##dtype
/* Declare all collective operations */
#define DECL_COLL5(coll, op, dtype) \
extern __device__ __attribute__((noinline)) void NCCL_COLL_NAME(coll, op, dtype)(struct CollectiveArgs* args); \
extern __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl c); \
#define DECL_COLL4(coll, op, dtype) \
DECL_COLL5(coll, op, dtype) \
DECL_COLL5(coll##LL, op, dtype)
#define DECL_COLL3(coll, op, dtype) \
DECL_COLL4(coll##Ring, op, dtype) \
DECL_COLL4(coll##Tree, op, dtype)
#define DECL_COLL2(coll, op) \
DECL_COLL3(coll, op, i8) \
DECL_COLL3(coll, op, u8) \
DECL_COLL3(coll, op, i32) \
DECL_COLL3(coll, op, u32) \
DECL_COLL3(coll, op, i64) \
DECL_COLL3(coll, op, u64) \
DECL_COLL3(coll, op, f16) \
DECL_COLL3(coll, op, f32) \
DECL_COLL3(coll, op, f64) \
DECL_COLL3(coll, op, b16)
#define DECL_COLL(coll) \
DECL_COLL2(coll, sum) \
DECL_COLL2(coll, prod) \
DECL_COLL2(coll, min) \
DECL_COLL2(coll, max)
#define DECL_ALL_COLLS \
DECL_COLL2(ncclBroadcast, copy) \
DECL_COLL(ncclReduce) \
DECL_COLL2(ncclAllGather, copy) \
DECL_COLL(ncclReduceScatter) \
DECL_COLL(ncclAllReduce) \
DECL_ALL_COLLS
// CHUNKSIZE must be a multiple of SLICESIZE
//#define ALLREDUCE_SLICESTEPS (NCCL_STEPS/4)
//#define ALLREDUCE_CHUNKSTEPS (NCCL_STEPS/2)
//#define ALLGATHER_SLICESTEPS (NCCL_STEPS/4)
//#define ALLGATHER_CHUNKSTEPS (NCCL_STEPS/2)
//#define REDUCESCATTER_SLICESTEPS (NCCL_STEPS/4)
//#define REDUCESCATTER_CHUNKSTEPS (NCCL_STEPS/2)
#define ALLREDUCE_SLICESTEPS 4
#define ALLREDUCE_CHUNKSTEPS 4
#define ALLGATHER_SLICESTEPS 4
#define ALLGATHER_CHUNKSTEPS 4
#define REDUCESCATTER_SLICESTEPS 4
#define REDUCESCATTER_CHUNKSTEPS 4
#define BROADCAST_SLICESTEPS 1
#define BROADCAST_CHUNKSTEPS 1
#define REDUCE_SLICESTEPS 1
#define REDUCE_CHUNKSTEPS 1
#endif
src/collectives/device/Makefile
+1 -1
View File
@@ -68,4 +68,4 @@ $(DEVOBJ) : $(LIBOBJ)
$(NVCC) $(NVCUFLAGS) -dlink $^ -o $@
clean:
rm -f $(LIBOBJ) $(DEVOBJ) $(DEPFILES) $(DEPENDFILES) $(STATICLIB) test
rm -f $(LIBOBJ) $(DEVOBJ) $(DEPFILES) $(DEPENDFILES) $(RULESFILE) $(STATICLIB)
src/collectives/device/all_gather.cu
+3 -4
View File
@@ -1,12 +1,11 @@
/*************************************************************************
* Copyright (c) 2015-2018, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "common.h"
#include "all_gather.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL3(ncclAllGather, copy, FuncSum, i8, int8_t, ncclCollAllGather, ncclSum, ncclInt8);
IMPL_COLL_C(ncclAllGather, ncclCollAllGather);
src/collectives/device/all_gather.h
+70 -4
View File
@@ -13,7 +13,7 @@ template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllGatherRingKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
@@ -21,15 +21,15 @@ __device__ void ncclAllGatherRingKernel(struct CollectiveArgs* args) {
const ssize_t size = args->N;
const int nranks = comm->nRanks;
const int stepSize = channel->buffSize / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLREDUCE_CHUNKSTEPS;
const int chunkSize = stepSize * ALLGATHER_CHUNKSTEPS;
const ssize_t loopSize = args->nChannels*(ssize_t)chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
ncclPrimitives<UNROLL, ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLREDUCE_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, args->opCount);
ncclPrimitives<UNROLL, ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, args->nThreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,args->nChannels));
@@ -134,3 +134,69 @@ __device__ void ncclAllGatherRingLLKernel(struct CollectiveArgs* args) {
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllGatherTreeLLKernel(struct CollectiveArgs* args) { }
#include "prims_ll128.h"
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllGatherRingLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int bid = args->bid;
const int nthreads = args->nThreads;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, channel, comm, args->opCount);
const ssize_t size = args->N;
//const int rank = comm->rank;
const int nranks = comm->nRanks;
ssize_t chunkSize = (NCCL_LL128_ELEMS_PER_THREAD*nthreads*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 ssize_t loopSize = args->nChannels*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, args->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];
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);
}
}
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllGatherTreeLL128Kernel(struct CollectiveArgs* args) { }
src/collectives/device/all_reduce.cu
+3 -7
View File
@@ -1,15 +1,11 @@
/*************************************************************************
* Copyright (c) 2015-2018, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "common.h"
#include "all_reduce.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL2(ncclAllReduce, sum, FuncSum, ncclCollAllReduce, ncclSum);
IMPL_COLL2(ncclAllReduce, prod, FuncProd, ncclCollAllReduce, ncclProd);
IMPL_COLL2(ncclAllReduce, min, FuncMin, ncclCollAllReduce, ncclMin);
IMPL_COLL2(ncclAllReduce, max, FuncMax, ncclCollAllReduce, ncclMax);
IMPL_COLL_R(ncclAllReduce, ncclCollAllReduce);
src/collectives/device/all_reduce.h
+169 -17
View File
@@ -13,7 +13,7 @@ template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllReduceRingKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
@@ -34,7 +34,7 @@ __device__ void ncclAllReduceRingKernel(struct CollectiveArgs* args) {
T * __restrict__ thisOutput = (T*)args->ThisOutput;
ncclPrimitives<UNROLL, ALLREDUCE_CHUNKSTEPS/ALLREDUCE_SLICESTEPS, ALLREDUCE_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, args->opCount);
prims(tid, args->nThreads, &ring->prev, &ring->next, thisOutput, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += nranks*loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,nranks*args->nChannels));
@@ -103,25 +103,31 @@ __device__ void ncclAllReduceRingKernel(struct CollectiveArgs* args) {
}
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* tree = &channel->tree;
const ssize_t size = args->N;
const int stepSize = channel->buffSize / (sizeof(T)*NCCL_STEPS);
const int chunkSize = args->lastChunkSize;
int chunkSize = args->lastChunkSize;
const ssize_t minChunkSize = nthreads*8*sizeof(uint64_t) / sizeof(T);
const ssize_t loopSize = args->nChannels*chunkSize;
if (loopSize > size) {
chunkSize = DIVUP(size, args->nChannels*minChunkSize)*minChunkSize;
}
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
do {
struct ncclTree* tree = &channel->treeUp;
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclPrimitives<1, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, nthreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount);
ncclPrimitives<1, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, args->nThreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
@@ -137,8 +143,9 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
} while(0);
do {
struct ncclTree* tree = &channel->treeDn;
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclPrimitives<1, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, nthreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount);
ncclPrimitives<1, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, args->nThreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
@@ -170,6 +177,8 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
//const int rank = comm->rank;
const int nranks = comm->nRanks;
ssize_t chunkSize = NCCL_LL_SLICE_LINES * sizeof(uint64_t) / sizeof(T);
const ssize_t minChunkSize = nthreads * (sizeof(uint64_t)) / sizeof(T);
const ssize_t loopSize = args->nChannels*nranks*chunkSize;
// Compute pointers
@@ -177,10 +186,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
if (size-gridOffset < loopSize) {
chunkSize = args->lastChunkSize;
}
ssize_t chunkOffset = gridOffset + bid*nranks*chunkSize;
chunkSize = min(DIVUP(size-gridOffset, args->nChannels*nranks*minChunkSize)*minChunkSize, chunkSize);
/////////////// begin AllReduce steps ///////////////
ssize_t offset;
@@ -189,7 +195,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
// step 0: push data to next GPU
slice = ring->devUserRanks[nranks-1];
offset = chunkOffset + slice * chunkSize;
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.send(thisInput+offset, nelem);
@@ -197,7 +203,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
// k-2 steps: reduce and copy to next GPU
for (int j=2; j<nranks; ++j) {
slice = ring->devUserRanks[nranks-j];
offset = chunkOffset + slice * chunkSize;
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceSend(thisInput+offset, nelem);
@@ -206,7 +212,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
// 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
slice = ring->devUserRanks[0];
offset = chunkOffset + slice * chunkSize;
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvReduceCopySend(thisInput+offset, thisOutput+offset, nelem);
@@ -214,7 +220,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
// k-2 steps: copy to next GPU
for (int j=1; j<nranks-1; ++j) {
slice = ring->devUserRanks[nranks-j];
offset = chunkOffset + slice * chunkSize;
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvCopySend(thisOutput+offset, nelem);
@@ -222,7 +228,7 @@ __device__ void ncclAllReduceRingLLKernel(struct CollectiveArgs* args) {
// Make final copy from buffer to dest.
slice = ring->devUserRanks[1];
offset = chunkOffset + slice * chunkSize;
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
// Here we need to copy from buffer to this output.
@@ -238,16 +244,21 @@ __device__ void ncclAllReduceTreeLLKernel(struct CollectiveArgs* args) {
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* tree = &channel->tree;
const ssize_t size = args->N;
ssize_t chunkSize = NCCL_LL_SLICE_LINES * sizeof(uint64_t) / sizeof(T);
const ssize_t minChunkSize = nthreads*sizeof(uint64_t) / sizeof(T);
const ssize_t loopSize = args->nChannels*chunkSize;
if (loopSize > size) {
chunkSize = DIVUP(size, args->nChannels*minChunkSize)*minChunkSize;
}
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
do {
struct ncclTree* tree = &channel->treeUp;
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclLLPrimitives<T, FUNC, NCCL_MAX_TREE_ARITY, 1> LLprims(tid, nthreads, tree->down, &tree->up, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
@@ -265,6 +276,7 @@ __device__ void ncclAllReduceTreeLLKernel(struct CollectiveArgs* args) {
} while(0);
do {
struct ncclTree* tree = &channel->treeDn;
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclLLPrimitives<T, FUNC, 1, NCCL_MAX_TREE_ARITY> LLprims(tid, nthreads, &tree->up, tree->down, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
@@ -281,3 +293,143 @@ __device__ void ncclAllReduceTreeLLKernel(struct CollectiveArgs* args) {
}
} while(0);
}
#include "prims_ll128.h"
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllReduceRingLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int bid = args->bid;
const int nthreads = args->nThreads;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, channel, comm, args->opCount);
const ssize_t size = args->N;
//const int rank = comm->rank;
const int nranks = comm->nRanks;
ssize_t chunkSize = (NCCL_LL128_ELEMS_PER_THREAD*nthreads*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 ssize_t loopSize = args->nChannels*nranks*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, args->nChannels*nranks*minChunkSize)*minChunkSize, chunkSize);
/////////////// begin AllReduce steps ///////////////
ssize_t offset;
int nelem;
int slice;
// step 0: push data to next GPU
slice = ring->devUserRanks[nranks-1];
offset = gridOffset + (slice*args->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) {
slice = ring->devUserRanks[nranks-j];
offset = gridOffset + (slice*args->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
slice = ring->devUserRanks[0];
offset = gridOffset + (slice*args->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) {
slice = ring->devUserRanks[nranks-j];
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
LLprims.recvCopySend(thisOutput+offset, nelem);
}
// Make final copy from buffer to dest.
slice = ring->devUserRanks[1];
offset = gridOffset + (slice*args->nChannels+bid) * chunkSize;
nelem = min(chunkSize, size-offset);
// Here we need to copy from buffer to this output.
LLprims.recv(thisOutput+offset, nelem);
}
}
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllReduceTreeLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclTree* treeUp = &channel->treeUp;
struct ncclTree* treeDn = &channel->treeDn;
const ssize_t size = args->N;
ssize_t chunkSize = args->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 = args->nChannels*chunkSize;
int nthreadsSplit = NCCL_LL128_SPLIT(nthreads);
if (loopSize > size) {
chunkSize = DIVUP(size, args->nChannels*minChunkSize)*minChunkSize;
}
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
if (treeUp->up == -1) {
// ReduceAndBroadcast : max number of recv is 3, max number of send is 3
ncclLL128Primitives<T, FUNC, NCCL_MAX_TREE_ARITY, NCCL_MAX_TREE_ARITY> LLprims(tid, nthreads, treeUp->down, treeDn->down, channel, comm, args->opCount);
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_TREE_ARITY, 1> LLprims(tid, nthreadsSplit, treeUp->down, &treeUp->up, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (treeUp->down[0] == -1) {
LLprims.send(thisInput+offset, nelem);
} else {
LLprims.recvReduceSend(thisInput+offset, nelem);
}
}
} else {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclLL128Primitives<T, FUNC, 1, NCCL_MAX_TREE_ARITY> LLprims(tid-nthreadsSplit, nthreads-nthreadsSplit, &treeDn->up, treeDn->down, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down
ssize_t offset = gridOffset + bid*chunkSize;
int nelem = min(chunkSize, size-offset);
if (treeDn->down[0] == -1) {
LLprims.recv(thisOutput+offset, nelem);
} else {
LLprims.recvCopySend(thisOutput+offset, nelem);
}
}
}
}
}
src/collectives/device/broadcast.cu
+3 -4
View File
@@ -1,12 +1,11 @@
/*************************************************************************
* Copyright (c) 2015-2018, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "common.h"
#include "broadcast.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL3(ncclBroadcast, copy, FuncSum, i8, int8_t, ncclCollBroadcast, ncclSum, ncclInt8);
IMPL_COLL_C(ncclBroadcast, ncclCollBroadcast);
src/collectives/device/broadcast.h
+52 -2
View File
@@ -13,7 +13,7 @@ template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclBroadcastRingKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
@@ -36,7 +36,7 @@ __device__ void ncclBroadcastRingKernel(struct CollectiveArgs* args) {
T * __restrict__ thisOutput = (T*)args->ThisOutput;
ncclPrimitives<UNROLL, BROADCAST_CHUNKSTEPS/BROADCAST_SLICESTEPS, BROADCAST_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
prims(tid, args->nThreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,args->nChannels));
@@ -121,3 +121,53 @@ __device__ void ncclBroadcastRingLLKernel(struct CollectiveArgs* args) {
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclBroadcastTreeLLKernel(struct CollectiveArgs* args) { }
#include "prims_ll128.h"
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclBroadcastRingLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int bid = args->bid;
const int nthreads = args->nThreads;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, channel, comm, args->opCount);
const ssize_t size = args->N;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->root;
ssize_t chunkSize = (NCCL_LL128_ELEMS_PER_THREAD*nthreads*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 = args->nChannels*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, args->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 UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclBroadcastTreeLL128Kernel(struct CollectiveArgs* args) { }
src/collectives/device/common.h
+56 -34
View File
@@ -9,10 +9,8 @@
#ifndef NCCL_DEVICE_COMMON_H_
#define NCCL_DEVICE_COMMON_H_
#include "../collectives.h"
#include "collectives.h"
#include "devcomm.h"
#include "nccl.h"
#include <type_traits>
// Exit If Abort Barrier across CTA: make sure all threads exit consistently
// Each thread sets a predicate to true if abort == 1
@@ -23,6 +21,7 @@
if (abort) __atomic_fetch_add(abortCount, 1, __ATOMIC_SEQ_CST); \
__syncthreads(); \
if (LOAD(abortCount)) { asm volatile ("s_endpgm"); return; }
#define __syncwarp()
#else
static inline __device__ void exitIfAbortBarrier(int abort) {
uint32_t popc;
@@ -36,12 +35,13 @@ static inline __device__ void exitIfAbortBarrier(int abort) {
#endif
#define NCCL_FUNC5(coll, op, dtype) \
NCCL_COLL_NAME(coll, op, dtype), \
NCCL_COLL_NAME(coll##LL, op, dtype)
NCCL_COLL_NAME(coll##LL, op, dtype), \
NCCL_COLL_NAME(coll##LL128, op, dtype), \
NCCL_COLL_NAME(coll, op, dtype)
#define NCCL_FUNC4(coll, op, dtype) \
NCCL_FUNC5(coll##Ring, op, dtype), \
NCCL_FUNC5(coll##Tree, op, dtype)
NCCL_FUNC5(coll##Tree, op, dtype), \
NCCL_FUNC5(coll##Ring, op, dtype)
// Must be consistent with ncclDataType_t
#define NCCL_FUNCS3A(coll, op) \
@@ -79,7 +79,7 @@ static inline __device__ void exitIfAbortBarrier(int abort) {
NCCL_FUNCS3B(coll, copy), \
NCCL_FUNCS3B(coll, copy)
// Must be consistent with ncclColl_t
// Must be consistent with ncclFunc_t
#define NCCL_FUNCS() { \
NCCL_FUNCS2B(ncclBroadcast), \
NCCL_FUNCS2A(ncclReduce), \
@@ -88,9 +88,9 @@ static inline __device__ void exitIfAbortBarrier(int abort) {
NCCL_FUNCS2A(ncclAllReduce) }
// Must be consistent with the ncclFuncSet enum
using ncclFunc_t = void (*)(struct CollectiveArgs*);
using ncclKernelFunc_t = void (*)(struct CollectiveArgs*);
static const __device__ constexpr ncclFunc_t ncclFuncs[]{
static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
// Don't try to initialize the host shadow copy of this device-side global
// variable. There is no host pointer to a device-side function, which
// confuses clang. This will be fixed in the next clang release.
@@ -123,36 +123,43 @@ struct Caller<f, f + 1>{
inline
__device__
void NCCL_CALL_FUNCTIONS(struct ncclColl* const c) noexcept {
if (c->funcIndex < 160) {
if (c->funcIndex % 4 == 0) ncclBroadcastRing_copy_i8(&c->args);
else if (c->funcIndex % 4 == 1) ncclBroadcastRingLL_copy_i8(&c->args);
else if (c->funcIndex % 4 == 2) ncclBroadcastTree_copy_i8(&c->args);
else ncclBroadcastTreeLL_copy_i8(&c->args);
if (c->funcIndex < 240) {
if (c->funcIndex % 6 == 0) ncclBroadcastTreeLL_copy_i8(&c->args);
else if (c->funcIndex % 6 == 1) ncclBroadcastTreeLL128_copy_i8(&c->args);
else if (c->funcIndex % 6 == 2) ncclBroadcastTree_copy_i8(&c->args);
else if (c->funcIndex % 6 == 3) ncclBroadcastRingLL_copy_i8(&c->args);
else if (c->funcIndex % 6 == 4) ncclBroadcastRingLL128_copy_i8(&c->args);
else ncclBroadcastRing_copy_i8(&c->args);
}
else if (c->funcIndex < 320) Caller<160, 320>::call(c);
else if (c->funcIndex < 480) {
if (c->funcIndex % 4 == 0) ncclAllGatherRing_copy_i8(&c->args);
else if (c->funcIndex % 4 == 1) ncclAllGatherRingLL_copy_i8(&c->args);
else if (c->funcIndex % 4 == 2) ncclAllGatherTree_copy_i8(&c->args);
else ncclAllGatherTreeLL_copy_i8(&c->args);
else if (c->funcIndex < 480) Caller<240, 480>::call(c);
else if (c->funcIndex < 720) {
if (c->funcIndex % 6 == 0) ncclAllGatherTreeLL_copy_i8(&c->args);
else if (c->funcIndex % 6 == 1) ncclAllGatherTreeLL128_copy_i8(&c->args);
else if (c->funcIndex % 6 == 2) ncclAllGatherTree_copy_i8(&c->args);
else if (c->funcIndex % 6 == 3) ncclAllGatherRingLL_copy_i8(&c->args);
else if (c->funcIndex % 6 == 4) ncclAllGatherRingLL128_copy_i8(&c->args);
else ncclAllGatherRing_copy_i8(&c->args);
}
else Caller<480, 800>::call(c);
else Caller<720, 1200>::call(c);
}
static __device__ void load_parallel(void* dst, void* src, size_t size, int tid, uint32_t* abortCount) {
int* d = (int*)dst;
int* s = (int*)src;
// When aggregation is effective, if some threads have aborted inside the LL kernel,
// make sure the rest of the threads abort as well
exitIfAbortBarrier(0, abortCount);
for (int o = tid; o < (size/sizeof(int)); o += blockDim.x) d[o] = s[o];
__syncthreads();
}
static __device__ void load_coll(struct ncclColl* localColl, struct ncclColl* hostColl, int tid, uint32_t* abortCount) {
static __device__ void load_coll(struct ncclColl* localColl, struct ncclColl* hostColl, int tid, struct ncclDevComm* comm, uint32_t* abortCount) {
// Check whether the last operation was aborted and make sure all threads exit
int abort = tid == 0 ? *(comm->abortFlag) : 0;
exitIfAbortBarrier(abort, abortCount);
load_parallel(localColl, hostColl, sizeof(struct ncclColl), tid, abortCount);
__syncthreads();
if (tid == 0) hostColl->active = 0;
}
extern __device__ volatile uint64_t* ncclShmem;
/* Functions for aggregation case */
#define IMPL_COLL_FUNC(coll, op, ncclFunc, dtype, ctype) \
__device__ void NCCL_COLL_NAME(coll, op, dtype)(struct CollectiveArgs* args) { \
@@ -161,12 +168,14 @@ __device__ void NCCL_COLL_NAME(coll, op, dtype)(struct CollectiveArgs* args) { \
/* Kernels with the first operation inlined */
#define IMPL_COLL_KERN(coll, op, ncclFunc, dtype, ctype, fIndex) \
__launch_bounds__(MAXTHREADS+WARP_SIZE, 1) \
__global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
int tid = threadIdx.x; \
int bid = blockIdx.x; \
__shared__ volatile uint64_t shmem[NCCL_LL128_SHMEM_SIZE]; \
ncclShmem = shmem; \
__shared__ struct ncclColl localColl; \
__shared__ uint32_t abortCount; \
__shared__ uint32_t sync[NCCL_LL128_MAX_NTHREADS/WARP_SIZE]; \
if (tid == 0) abortCount = 0; \
__syncthreads(); \
\
@@ -174,12 +183,13 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
struct ncclChannel* channel = comm->channels+bid; \
struct ncclColl* c; \
channel->abortCount = &abortCount; \
channel->sync = sync; \
if (bid == 0) { \
/* To optimize for latency, (only) the first operation is passed as argument.*/ \
c = &firstColl; \
} else { \
c = &localColl; \
load_coll(c, channel->devCollectives+channel->collFifoHead, tid, &abortCount); \
load_coll(c, channel->devCollectives+channel->collFifoHead, tid, comm, &abortCount); \
} \
while (1) { \
if (tid < c->args.nThreads) { \
@@ -198,7 +208,7 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
\
/* Load next collective operation*/ \
c = &localColl; /* for bid 0 */ \
load_coll(c, channel->devCollectives+nextIndex, tid, &abortCount); \
load_coll(c, channel->devCollectives+nextIndex, tid, comm, &abortCount); \
} \
}
@@ -212,13 +222,14 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
// Only generate inline kernels for LL
#define IMPL_COLL4(coll, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType, al) \
IMPL_COLL_FUNC(coll, op, ncclFunc, dtype, ctype) \
IMPL_COLL_FUNC(coll##LL, op, ncclFunc, dtype, ctype) \
IMPL_COLL_KERN_##op(coll##LL, op, ncclFunc, dtype, ctype, FUNC_INDEX(ncclColl, ncclOp, ncclType, 1, al)) \
IMPL_COLL_FUNC(coll##LL128, op, ncclFunc, dtype, ctype) \
IMPL_COLL_FUNC(coll, op, ncclFunc, dtype, ctype) \
IMPL_COLL_KERN_##op(coll##LL, op, ncclFunc, dtype, ctype, FUNC_INDEX(ncclColl, ncclOp, ncclType, al, NCCL_PROTO_LL)) \
#define IMPL_COLL3(coll, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType) \
IMPL_COLL4(coll##Ring, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType, 0) \
IMPL_COLL4(coll##Tree, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType, 1)
IMPL_COLL4(coll##Tree, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType, NCCL_ALGO_TREE) \
IMPL_COLL4(coll##Ring, op, ncclFunc, dtype, ctype, ncclColl, ncclOp, ncclType, NCCL_ALGO_RING)
#define IMPL_COLL2(coll, op, ncclFunc, ncclColl, ncclOp) \
IMPL_COLL3(coll, op, ncclFunc, i8, int8_t, ncclColl, ncclOp, ncclInt8) \
@@ -232,6 +243,17 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclColl firstColl) { \
IMPL_COLL3(coll, op, ncclFunc, f64, double, ncclColl, ncclOp, ncclFloat64) \
IMPL_COLL3(coll, op, ncclFunc, b16, rccl_bfloat16, ncclColl, ncclOp, ncclBfloat16)
// Reduction define all functions
#define IMPL_COLL_R(collf, colln) \
IMPL_COLL2(collf, sum, FuncSum, colln, ncclSum); \
IMPL_COLL2(collf, prod, FuncProd, colln, ncclProd); \
IMPL_COLL2(collf, min, FuncMin, colln, ncclMin); \
IMPL_COLL2(collf, max, FuncMax, colln, ncclMax);
// Copy primitives only define one
#define IMPL_COLL_C(collf, colln) \
IMPL_COLL3(collf, copy, FuncSum, i8, int8_t, colln, ncclSum, ncclInt8);
#define COLL_UNROLL 2
#endif
src/collectives/device/common_kernel.h
-2
View File
@@ -300,8 +300,6 @@ __device__ void ReduceCopyMulti(const int tid, const int nthreads,
}
}
#define WARP_SIZE 64
template<class FUNC, typename T, int UNROLL, int MINSRCS, int MAXSRCS, int MINDSTS, int MAXDSTS>
__device__ void ReduceCopy128bMulti( const int w, const int nw, const int t,
int nsrcs, const T* s[MAXSRCS], int ndsts, T* d[MAXDSTS],
src/collectives/device/functions.cu
+70
View File
@@ -9,5 +9,75 @@
#include "collectives.h"
#include "common.h"
__device__ volatile uint64_t* ncclShmem;
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
NCCL_FUNC5(coll, op, dtype) \
NCCL_COLL_NAME(coll##LL, op, dtype), \
NCCL_COLL_NAME(coll##LL128, op, dtype), \
NCCL_COLL_NAME(coll, op, dtype)
#define NCCL_FUNC4(coll, op, dtype) \
NCCL_FUNC5(coll##Tree, op, dtype), \
NCCL_FUNC5(coll##Ring, op, dtype)
// Must be consistent with ncclDataType_t
#define NCCL_FUNCS3A(coll, op) \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, u8), \
NCCL_FUNC4(coll, op, i32), \
NCCL_FUNC4(coll, op, u32), \
NCCL_FUNC4(coll, op, i64), \
NCCL_FUNC4(coll, op, u64), \
NCCL_FUNC4(coll, op, f16), \
NCCL_FUNC4(coll, op, f32), \
NCCL_FUNC4(coll, op, f64)
#define NCCL_FUNCS3B(coll, op) \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8), \
NCCL_FUNC4(coll, op, i8)
// Must be consistent with ncclRedOp_t
#define NCCL_FUNCS2A(coll) \
NCCL_FUNCS3A(coll, sum ), \
NCCL_FUNCS3A(coll, prod), \
NCCL_FUNCS3A(coll, max ), \
NCCL_FUNCS3A(coll, min )
#define NCCL_FUNCS2B(coll) \
NCCL_FUNCS3B(coll, copy), \
NCCL_FUNCS3B(coll, copy), \
NCCL_FUNCS3B(coll, copy), \
NCCL_FUNCS3B(coll, copy)
// Must be consistent with ncclFunc_t
#define NCCL_FUNCS() { \
NCCL_FUNCS2B(ncclBroadcast), \
NCCL_FUNCS2A(ncclReduce), \
NCCL_FUNCS2B(ncclAllGather), \
NCCL_FUNCS2A(ncclReduceScatter), \
NCCL_FUNCS2A(ncclAllReduce) }
// Must be consistent with the ncclFuncSet enum
__device__ ncclKern_t ncclFuncs[NCCL_NUM_FUNCTIONS*ncclNumOps*ncclNumTypes*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS] = {
// Don't try to initialize the host shadow copy of this device-side global
// variable. There is no host pointer to a device-side function, which
// confuses clang. This will be fixed in the next clang release.
#if __CUDA_ARCH__
NCCL_FUNCS2B(ncclBroadcast),
NCCL_FUNCS2A(ncclReduce),
NCCL_FUNCS2B(ncclAllGather),
NCCL_FUNCS2A(ncclReduceScatter),
NCCL_FUNCS2A(ncclAllReduce)
#endif
};
#endif
// Workaround for https://reviews.llvm.org/D55580
__device__ void ncclWorkaroundClangD55580() {}
src/collectives/device/op128.h
+63
View File
@@ -0,0 +1,63 @@
/*************************************************************************
* Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#ifndef OP128_H_
#define OP128_H_
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
v0=LOAD(ptr);
v1=LOAD(ptr+1);
}
inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) {
STORE(ptr, v0);
STORE(ptr+1, v1);
}
inline __device__ uint64_t* shmemCvtPtr(volatile uint64_t* shmemGenericPtr) {
return (uint64_t*)shmemGenericPtr;
}
inline __device__ void loadShmem128(uint64_t* shmemAsmPtr, uint64_t &v0, uint64_t &v1) {
v0=LOAD(shmemAsmPtr);
v1=LOAD(shmemAsmPtr+1);
}
inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_t v1) {
STORE(shmemAsmPtr, v0);
STORE(shmemAsmPtr+1, v1);
}
#else
inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];"
: "=l"(v0), "=l"(v1) : "l"(ptr));
}
inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) {
asm volatile("st.volatile.global.v2.u64 [%2], {%0,%1};"
:: "l"(v0), "l"(v1), "l"(ptr));
}
inline __device__ uint64_t* shmemCvtPtr(volatile uint64_t* shmemGenericPtr) {
uint64_t* shmemAsmPtr;
asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(shmemAsmPtr) : "l"(shmemGenericPtr));
return shmemAsmPtr;
}
inline __device__ void loadShmem128(uint64_t* shmemAsmPtr, uint64_t &v0, uint64_t &v1) {
asm volatile("ld.volatile.shared.v2.u64 {%0,%1}, [%2];"
: "=l"(v0), "=l"(v1) : "l"(shmemAsmPtr));
}
inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_t v1) {
asm volatile("st.volatile.shared.v2.u64 [%2], {%0,%1};"
:: "l"(v0), "l"(v1), "l"(shmemAsmPtr));
}
#endif
#endif
src/collectives/device/primitives.h
+187 -419
View File
@@ -38,15 +38,27 @@ class ncclPrimitives {
private:
const int tid;
const int nthreads;
const int wid;
const int stepSize;
int nrecv = 0;
int nsend = 0;
const int stepSize;
struct ncclConnInfo* recvConn[NRECV];
struct ncclConnInfo* sendConn[NSEND];
volatile uint64_t* waitPtr;
struct ncclConnInfo* recvConn = NULL;
volatile uint64_t* recvConnHeadPtr = NULL;
uint64_t recvConnHead;
volatile uint64_t* recvConnTailPtr = NULL;
uint64_t recvConnTail;
uint64_t recvConnTailCache; // Cache last seen value
struct ncclConnInfo* sendConn = NULL;
volatile int* sendConnFifoPtr = NULL;
volatile uint64_t* sendConnTailPtr = NULL;
uint64_t sendConnTail;
volatile uint64_t* sendConnHeadPtr = NULL;
uint64_t sendConnHead;
uint64_t sendConnHeadCache; // Cache last seen value
uint64_t recvStep[NRECV];
uint64_t sendStep[NSEND];
uint64_t sendConnHead[NSEND];
#if defined(RCCL_USE_DIRECT_BUFFER)
const T* recvDirectBuff[NRECV];
T* sendDirectBuff[NSEND];
@@ -56,12 +68,12 @@ class ncclPrimitives {
struct ncclDevComm* comm;
uint32_t* abortCount;
__device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepSize; }
__device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepSize; }
__device__ const T* recvPtr(int i) { return ((const T*)recvBuff[i])+recvOffset(i); }
__device__ T* sendPtr(int i) { return ((T*)sendBuff[i])+sendOffset(i); }
inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepSize; }
inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepSize; }
inline __device__ const T* recvPtr(int i) { return ((const T*)recvBuff[i])+recvOffset(i); }
inline __device__ T* sendPtr(int i) { return ((T*)sendBuff[i])+sendOffset(i); }
__device__ void barrier() {
inline __device__ void barrier() {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
@@ -72,11 +84,11 @@ class ncclPrimitives {
uint32_t mismatch = 0;
const uint64_t opCount;
__device__ void checkMismatch(volatile uint64_t* remoteOpCount) {
inline __device__ void checkMismatch(struct ncclConnInfo* conn) {
if (mismatch) {
// In non-LL, we use _threadfence_system before incrementing opCount, yet we are still waiting for credits here, so there must be a size mismatch
STORE(comm->fatalDevError, ncclDevAssertedMismatch);
} else if (remoteOpCount && LOAD(remoteOpCount) > opCount) {
} else if (conn && LOAD(conn->opCountRem) > opCount) {
mismatch += 1;
}
}
@@ -84,68 +96,76 @@ class ncclPrimitives {
uint32_t spins = 0;
uint32_t abort = 0;
__device__ int checkAbort(volatile uint64_t* remoteOpCount) {
inline __device__ int checkAbort(int i, int send) {
spins++;
if (spins == SPINS_BEFORE_CHECK_ABORT) {
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = LOAD(comm->abortFlag);
checkMismatch(remoteOpCount);
if (wid == i) checkMismatch(send ? sendConn : recvConn);
spins = 0;
}
return abort;
}
__device__ void waitRecv(int i) {
inline __device__ void waitSend(int nbytes) {
spins = 0;
mismatch = 0;
recvStep[i] += SLICESTEPS;
if (tid == i) {
if (sendConnHeadPtr) {
#ifdef ENABLE_PROFILING
auto devProf = comm->devProf;
uint64_t t0 = clock64();
#endif
while (LOAD(waitPtr) < recvStep[i]) {
if (checkAbort(recvConn[i]->opCountRem)) break;
}
#ifdef ENABLE_PROFILING
__atomic_fetch_add(&devProf->wait_recv_cycle[blockIdx.x], clock64() - t0, __ATOMIC_SEQ_CST);
#endif
}
}
__device__ void waitSend(int i) {
spins = 0;
mismatch = 0;
sendStep[i] += SLICESTEPS;
if (tid == WARP_SIZE+i) {
#ifdef ENABLE_PROFILING
auto devProf = comm->devProf;
uint64_t t0 = clock64();
#endif
while (sendConnHead[i] + NCCL_STEPS < sendStep[i]) {
sendConnHead[i] = LOAD(waitPtr);
if (checkAbort(sendConn[i]->opCountRem)) break;
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + SLICESTEPS) {
sendConnHeadCache = LOAD(sendConnHeadPtr);
if (checkAbort(wid, 1)) break;
}
#ifdef ENABLE_PROFILING
__atomic_fetch_add(&devProf->wait_send_cycle[blockIdx.x], clock64() - t0, __ATOMIC_SEQ_CST);
#endif
if (sendConnFifoPtr) {
STORE(sendConnFifoPtr+sendConnHead%NCCL_STEPS, nbytes);
}
sendConnHead += SLICESTEPS;
}
}
inline __device__ void postRecv(int i) {
STORE(recvConn[i]->head, recvStep[i]);
inline __device__ void waitRecv() {
spins = 0;
mismatch = 0;
if (recvConnTailPtr) {
#ifdef ENABLE_PROFILING
auto devProf = comm->devProf;
uint64_t t0 = clock64();
#endif
while (recvConnTailCache < recvConnTail + SLICESTEPS) {
recvConnTailCache = LOAD(recvConnTailPtr);
if (checkAbort(wid, 0)) break;
}
#ifdef ENABLE_PROFILING
__atomic_fetch_add(&devProf->wait_recv_cycle[blockIdx.x], clock64() - t0, __ATOMIC_SEQ_CST);
#endif
recvConnTail += SLICESTEPS;
}
}
inline __device__ void postSend(int i) {
if (sendConn[i]->next_hdp_reg) STORE(sendConn[i]->next_hdp_reg, 0x1);
STORE(sendConn[i]->tail, sendStep[i]);
inline __device__ void incRecv(int i) {
recvStep[i] += SLICESTEPS;
}
inline __device__ void postRecv() {
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += SLICESTEPS);
}
__device__ void postSendSize(int i, int size) {
if (sendConn[i]->fifo) STORE(sendConn[i]->fifo+((sendStep[i]-SLICESTEPS)%NCCL_STEPS), size);
inline __device__ void incSend(int i) {
sendStep[i] += SLICESTEPS;
}
inline __device__ void postSend() {
if (sendConnTailPtr) {
if (sendConn->next_hdp_reg) STORE(sendConn->next_hdp_reg, 0x1);
STORE(sendConnTailPtr, sendConnTail += SLICESTEPS);
}
}
template <int DIRECTRECV>
__device__ const T* directRecvPtr(int i, int directOffset) {
inline __device__ const T* directRecvPtr(int i, int directOffset) {
#if defined(RCCL_USE_DIRECT_BUFFER)
return DIRECTRECV && recvDirectBuff[i] ? recvDirectBuff[i]+directOffset : recvPtr(i);
#else
@@ -154,19 +174,38 @@ class ncclPrimitives {
}
template <int DIRECTSEND>
__device__ T* directSendPtr(int i, int directOffset) {
#if defined(RCCL_USE_DIRECT_BUFFER)
inline __device__ T* directSendPtr(int i, int directOffset) {
#if defined(RCCL_USE_DIRECT_BUFFER)
return DIRECTSEND && sendDirectBuff[i] ? sendDirectBuff[i]+directOffset : sendPtr(i);
#else
#else
return sendPtr(i);
#endif
#endif
}
template <int DIRECTRECV>
inline __device__ int directRecvInc(int i, int directInc, int sliceInc) {
#if defined(RCCL_USE_DIRECT_BUFFER)
return DIRECTRECV && recvDirectBuff[i] ? directInc : sliceInc;
#else
return sliceInc;
#endif
}
template <int DIRECTSEND>
inline __device__ int directSendInc(int i, int directInc, int sliceInc) {
#if defined(RCCL_USE_DIRECT_BUFFER)
return DIRECTSEND && sendDirectBuff[i] ? directInc : sliceInc;
#else
return sliceInc;
#endif
}
template <int DIRECTRECV, int DIRECTSEND, int RECV, int SEND, int SRC, int DST>
__device__ void
inline __device__ void
GenericOp(const T* srcPtr, T* dstPtr, int nelem, int directOffset) {
int offset = 0;
int sliceSize = stepSize * SLICESTEPS;
int sliceSize = stepSize*SLICESTEPS;
int dataSize = max(DIVUP(nelem, 16*SLICESPERCHUNK)*16, sliceSize/32);
const T* srcs[RECV*NRECV+SRC];
srcs[0] = SRC ? srcPtr : directRecvPtr<DIRECTRECV>(0, directOffset);
@@ -182,11 +221,11 @@ class ncclPrimitives {
for (int i=1; i<NSEND && i<nsend; i++) dsts[DST+i] = directSendPtr<DIRECTSEND>(i, directOffset);
}
#pragma unroll 1
#pragma unroll
for (int slice=0; slice<SLICESPERCHUNK; ++slice) {
int realSize = max(0, min(sliceSize, nelem-offset));
FOR_SEND(waitSend);
FOR_RECV(waitRecv);
int realSize = max(0, min(dataSize, nelem-offset));
if (SEND) waitSend(realSize*sizeof(T));
if (RECV) waitRecv();
if (realSize > 0) {
barrier();
#if defined(RCCL_USE_DIRECT_BUFFER)
@@ -202,458 +241,187 @@ class ncclPrimitives {
ReduceOrCopyMulti<UNROLL, FUNC, T, RECV+SRC, RECV*NRECV+SRC, SEND+DST, SEND*NSEND+DST>(tid, nthreads, RECV*nrecv+SRC, srcs, SEND*nsend+DST, dsts, realSize);
#endif
}
exitIfAbortBarrier(abort, abortCount);
if (tid == 0)
{
FOR_SEND(postSendSize, realSize*sizeof(T));
__threadfence_system();
FOR_SEND(postSend);
FOR_RECV(postRecv);
barrier();
FOR_SEND(incSend);
FOR_RECV(incRecv);
if (tid >= nthreads-WARP_SIZE) {
if (SEND) {
if (realSize > 0 && wid == 0) __threadfence_system();
__syncwarp();
postSend();
}
if (RECV) postRecv();
}
for (int i=0; i<RECV*NRECV+SRC; i++) srcs[i] += sliceSize;
for (int i=0; i<SEND*NSEND+DST; i++) dsts[i] += sliceSize;
offset += sliceSize;
srcs[0] += SRC ? realSize : directRecvInc<DIRECTRECV>(0, realSize, sliceSize);
for (int i=1-SRC; i<RECV*NRECV; i++) srcs[SRC+i] += sliceSize;
dsts[0] += DST ? realSize : directSendInc<DIRECTSEND>(0, realSize, sliceSize);
for (int i=1-DST; i<SEND*NSEND; i++) dsts[DST+i] += directSendInc<DIRECTSEND>(i, realSize, sliceSize);
offset += realSize;
}
}
__device__ void loadRecvConn(struct ncclConnInfo* conn, int i, T* directBuff) {
recvConn[i] = conn;
recvBuff[i] = (const T*)LOAD(&recvConn[i]->buff);
recvStep[i] = LOAD(&recvConn[i]->step);
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i, T* directBuff) {
recvBuff[i] = (const T*)LOAD(&conn->buff);
recvStep[i] = LOAD(&conn->step);
recvStep[i] = ROUNDUP(recvStep[i], SLICESPERCHUNK*SLICESTEPS);
// Return credits in case we rounded up.
if (tid == 0) STORE(recvConn[i]->head, recvStep[i]);
if (tid == i) {
waitPtr = LOAD(&recvConn[i]->tail);
STORE(recvConn[i]->opCountLoc, opCount);
}
#if defined(RCCL_USE_DIRECT_BUFFER)
recvDirectBuff[i] = NULL;
if (directBuff && recvConn[i]->direct) {
if (directBuff && LOAD(&conn->direct)) {
recvDirectBuff[i] = directBuff;
if (tid == 0) STORE(recvConn[i]->ptrExchange, directBuff);
if (tid == 0) STORE(conn->ptrExchange, directBuff);
}
#endif
if (wid == i) recvConn = conn;
if (wid == i) recvConnTail = recvConnHead = recvStep[i]; // Make sure we set this after rounding up
nrecv++;
}
__device__ void loadSendConn(struct ncclConnInfo* conn, int i, T* directBuff) {
sendConn[i] = conn;
sendBuff[i] = (T*)LOAD(&sendConn[i]->buff);
sendStep[i] = LOAD(&sendConn[i]->step);
sendStep[i] = ROUNDUP(sendStep[i], SLICESPERCHUNK*SLICESTEPS);
if (tid == WARP_SIZE+i) {
waitPtr = LOAD(&sendConn[i]->head);
sendConnHead[i] = LOAD(waitPtr);
STORE(sendConn[i]->opCountLoc, opCount);
__device__ __forceinline__ void loadRecvSync() {
if (tid >= WARP_SIZE && tid < 2*WARP_SIZE && wid<nrecv) {
recvConnTailPtr = LOAD(&recvConn->tail);
recvConnTailCache = LOAD(recvConnTailPtr);
}
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
recvConnHeadPtr = LOAD(&recvConn->head);
// Return credits in case we rounded up.
STORE(recvConnHeadPtr, recvConnHead);
// Update opCount in case we skipped some operations
STORE(recvConn->opCountLoc, opCount);
}
}
__device__ void loadSendConn(struct ncclConnInfo* conn, int i, T* directBuff) {
sendBuff[i] = (T*)LOAD(&conn->buff);
sendStep[i] = LOAD(&conn->step);
sendStep[i] = ROUNDUP(sendStep[i], SLICESPERCHUNK*SLICESTEPS);
#if defined(RCCL_USE_DIRECT_BUFFER)
sendDirectBuff[i] = NULL;
if (directBuff && sendConn[i]->direct) {
void* volatile* ptr = sendConn[i]->ptrExchange;
if (directBuff && LOAD(&conn->direct)) {
void* volatile* ptr = LOAD(&conn->ptrExchange);
while ((sendDirectBuff[i] = (T*)(LOAD(ptr))) == NULL);
__syncthreads();
barrier();
if (tid == 0) STORE(ptr, NULL);
}
#endif
if (wid == i) sendConn = conn;
if (wid == i) sendConnTail = sendConnHead = sendStep[i]; // Make sure we set this after rounding up
nsend++;
}
__device__ void saveRecvConn(int i) {
if (tid == i) {
STORE(&recvConn[i]->step, recvStep[i]);
__threadfence_system();
__atomic_fetch_add(recvConn[i]->opCountLoc, 1, __ATOMIC_SEQ_CST);
__device__ void loadSendSync() {
if (tid < nsend) {
sendConnHeadPtr = LOAD(&sendConn->head);
sendConnHeadCache = LOAD(sendConnHeadPtr);
sendConnFifoPtr = LOAD(&sendConn->fifo);
STORE(sendConn->opCountLoc, opCount);
}
if (tid >= nthreads-WARP_SIZE && wid < nsend) {
sendConnTailPtr = LOAD(&sendConn->tail);
}
}
__device__ void saveSendConn(int i) {
if (tid == WARP_SIZE+i) {
STORE(&sendConn[i]->step, sendStep[i]);
__device__ void saveRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
STORE(&recvConn->step, recvConnHead);
STORE(recvConn->opCountLoc, opCount+1);
__threadfence_system();
}
}
__device__ void saveSendSync() {
if (tid < nsend) {
STORE(&sendConn->step, sendConnHead);
STORE(sendConn->opCountLoc, opCount+1);
__threadfence_system();
__atomic_fetch_add(sendConn[i]->opCountLoc, 1, __ATOMIC_SEQ_CST);
}
}
public:
__device__
__device__ __forceinline__
ncclPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, T* directBuff, int stepSize, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount)
: comm(comm), tid(tid), nthreads(nthreads), stepSize(stepSize), opCount(opCount) {
// Make sure step is updated before we read it
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), stepSize(stepSize), opCount(opCount) {
// Make sure step is updated before we read it.
abortCount = channel->abortCount;
__syncthreads();
barrier();
// disable directBuff
for (int i=0; i<NRECV && recvPeers[i] >= 0; i++) loadRecvConn(&channel->devPeers[recvPeers[i]].recv.conn, i, 0);
for (int i=0; i<NSEND && sendPeers[i] >= 0; i++) loadSendConn(&channel->devPeers[sendPeers[i]].send.conn, i, 0);
loadRecvSync();
loadSendSync();
}
__device__ void
__device__ __forceinline__ void
send(const T* src, int nelem) {
GenericOp<0, 0, 0, 1, 1, 0>(src, NULL, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
directSend(const T* src, int directOffset, int nelem) {
GenericOp<0, 1, 0, 1, 1, 0>(src, NULL, nelem, directOffset);
}
__device__ void
__device__ __forceinline__ void
recv(T* dst, int nelem) {
GenericOp<0, 0, 1, 0, 0, 1>(NULL, dst, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
directRecv(T* dst, int directOffset, int nelem) {
GenericOp<1, 0, 1, 0, 0, 1>(NULL, dst, nelem, directOffset);
}
__device__ void
__device__ __forceinline__ void
copySend(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 0, 1, 1, 1>(src, dst, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
directCopySend(const T* src, T* dst, int directOffset, int nelem) {
GenericOp<0, 1, 0, 1, 1, 1>(src, dst, nelem, directOffset);
}
__device__ void
__device__ __forceinline__ void
recvCopySend(T* dst, int nelem) {
GenericOp<0, 0, 1, 1, 0, 1>(NULL, dst, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
directRecvCopySend(T* dst, int directOffset, int nelem) {
GenericOp<1, 1, 1, 1, 0, 1>(NULL, dst, nelem, directOffset);
}
__device__ void
__device__ __forceinline__ void
recvReduceCopy(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 1, 0, 1, 1>(src, dst, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
recvReduceSend(const T* src, int nelem) {
GenericOp<0, 0, 1, 1, 1, 0>(src, NULL, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
recvReduceCopySend(const T* src, T* dst, int nelem) {
GenericOp<0, 0, 1, 1, 1, 1>(src, dst, nelem, 0);
}
__device__ void
__device__ __forceinline__ void
directRecvReduceCopySend(const T* src, T* dst, int directOffset, int nelem) {
// Direct is only for the send part
GenericOp<0, 1, 1, 1, 1, 1>(src, dst, nelem, directOffset);
}
__device__ ~ncclPrimitives() {
// Save steps for next collective. Have thread 0 do it to be compatible
// with the way LL works.
for (int i=0; i<NRECV && i<nrecv; i++) saveRecvConn(i);
for (int i=0; i<NSEND && i<nsend; i++) saveSendConn(i);
}
};
template <typename T, class FUNC, int NRECV, int NSEND>
class ncclLLPrimitives {
private:
const int tid;
const int nthreads;
int nrecv = 0;
int nsend = 0;
struct ncclConnInfo* recvConn[NRECV];
struct ncclConnInfo* sendConn[NSEND];
volatile uint64_t* waitPtr;
volatile uint64_t* postPtr;
volatile int* fifoPtr;
uint64_t recvStep[NRECV];
uint64_t sendStep[NSEND];
uint64_t sendConnHead;
union ncclLLFifoLine* recvBuff[NRECV];
union ncclLLFifoLine* sendBuff[NSEND];
struct ncclDevComm* comm;
uint32_t* abortCount;
__device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
__device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
__device__ union ncclLLFifoLine* recvPtr(int i) { return recvBuff[i]+recvOffset(i); }
__device__ union ncclLLFifoLine* sendPtr(int i) { return sendBuff[i]+sendOffset(i); }
__device__ uint32_t recvFlag(int i) { return NCCL_LL_FLAG(recvStep[i]+1); }
__device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(sendStep[i]+1); }
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
// Exit If Abort Barrier : make sure all threads exit consistently
// Each thread sets a predicate to true if val == 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()
__device__ void exitIfAbortLocalBarrier() {
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, 14, %1, barr_pred;" : "=r"(popc) : "r"(nthreads));
asm ("}");
if (popc) {
// Make sure threads not participating in the operation get the abort and all threads exit
exitIfAbortBarrier(1);
}
}
#endif
__device__ void barrier() {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
asm volatile ("bar.sync 1, %0;" :: "r"(nthreads));
#endif
}
uint32_t mismatch = 0;
const uint64_t opCount;
__device__ void checkMismatch(volatile uint64_t* remoteOpCount) {
if (mismatch > 20) {
// We have seen that the peer advanced opcount so many times yet we are still waiting for credit of current op, so it is _most likely_ a mismatch
// Note that we are not using _threadfence_system in LL so the error cannot be asserted
STORE(comm->fatalDevError, ncclDevSuspectedMismatch);
} else if (remoteOpCount && LOAD(remoteOpCount) > opCount) {
mismatch += 1;
}
}
uint32_t spins = 0;
uint32_t abort = 0;
__device__ int checkAbort(volatile uint64_t* remoteOpCount) {
spins++;
if (spins == SPINS_BEFORE_CHECK_ABORT) {
abort = LOAD(comm->abortFlag);
checkMismatch(remoteOpCount);
spins = 0;
}
return abort;
}
__device__ void waitSend(int i, int nbytes) {
spins = 0;
mismatch = 0;
if (tid == WARP_SIZE+i) {
while (sendConnHead + NCCL_STEPS < sendStep[i] + 1) {
sendConnHead = LOAD(waitPtr);
if (checkAbort(sendConn[i]->opCountRem)) break;
}
if (fifoPtr) {
int size = ((sendStep[i] & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? NCCL_LL_SLICE_LINES*sizeof(union ncclLLFifoLine) : nbytes;
STORE(fifoPtr+sendStep[i]%NCCL_STEPS, size);
}
}
}
__device__ void postRecv(int i) {
recvStep[i]++;
if (tid == i) STORE(postPtr, recvStep[i]);
}
__device__ void postSend(int i, int offset) {
// LL Cleanup : write all flags in the slice to make sure we don't have
// data corruption when flag loops over.
if ((sendStep[i] & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) {
for (int o = offset; o<NCCL_LL_SLICE_LINES; o+=nthreads) storeLL(sendPtr(i)+o, 0, sendFlag(i));
}
sendStep[i]++;
}
__device__ __attribute__((noinline)) uint64_t readLL(int i, int offset) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
uint32_t flag = recvFlag(i);
uint32_t data1, flag1, data2, flag2;
spins = 0;
mismatch = 0;
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
do {
asm volatile ("flat_load_dwordx4 %0, %1, glc\n"
"s_waitcnt vmcnt(0)\n"
"buffer_wbinvl1_vol\n" : "=v"(i4) : "v"(src));
if (i4[1] == flag && i4[3] == flag) break;
} while (!checkAbort(recvConn[i]->opCountRem));
uint64_t val64 = (uint64_t)(i4[0]) + (((uint64_t)i4[2]) << 32);
#else
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(recvConn[i]->opCountRem)) break;
} while ((flag1 != flag) || (flag2 != flag));
uint64_t val64 = data1 + (((uint64_t)data2) << 32);
#endif
return val64;
}
__device__ __attribute__((noinline)) void storeLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = val & 0xffffffff;
i4[1] = flag;
i4[2] = (val >> 32);
i4[3] = flag;
asm volatile ("flat_store_dwordx4 %0, %1, glc\n"
"s_waitcnt vmcnt(0)\n"
"buffer_wbinvl1_vol\n" : : "v"(dst), "v"(i4));
#else
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));
#endif
}
// 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;
}
__device__ void storeAL(uint64_t* dst, uint64_t val, uint32_t nbytes) {
memcpy((char*)dst, (char*)&val, nbytes);
}
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);
FOR_SEND(waitSend, nbytes*2);
barrier();
uint32_t npack = DIVUP(nbytes, sizeof(uint64_t));
uint64_t* srcPack = (uint64_t*)srcPtr;
uint64_t* dstPack = (uint64_t*)dstPtr;
int offset = tid;
// Do multiples of 64 bits
#pragma unroll 1
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);
}
}
// 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));
}
}
}
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
exitIfAbortBarrier(abort, abortCount);
#else
exitIfAbortLocalBarrier();
#endif
FOR_RECV(postRecv);
FOR_SEND(postSend, offset);
}
__device__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
recvConn[i] = conn;
recvBuff[i] = recvConn[i]->llBuff;
recvStep[i] = recvConn[i]->step;
if (tid == i) {
postPtr = recvConn[i]->head;
STORE(recvConn[i]->opCountLoc, opCount);
}
nrecv++;
}
__device__ void loadSendConn(struct ncclConnInfo* conn, int i) {
sendConn[i] = conn;
sendBuff[i] = sendConn[i]->llBuff;
sendStep[i] = sendConn[i]->step;
if (tid == WARP_SIZE+i) {
waitPtr = sendConn[i]->head;
fifoPtr = sendConn[i]->fifo;
sendConnHead = LOAD(waitPtr);
STORE(sendConn[i]->opCountLoc, opCount);
}
nsend++;
}
__device__ void saveRecvConn(int i) {
if (tid == i) {
recvConn[i]->step = recvStep[i];
__atomic_fetch_add(recvConn[i]->opCountLoc, 1, __ATOMIC_SEQ_CST);
__threadfence_block();
}
}
__device__ void saveSendConn(int i) {
if (tid == WARP_SIZE+i) {
sendConn[i]->step = sendStep[i];
__atomic_fetch_add(sendConn[i]->opCountLoc, 1, __ATOMIC_SEQ_CST);
__threadfence_block();
}
}
public:
__device__
ncclLLPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount)
: comm(comm), tid(tid), nthreads(nthreads), opCount(opCount) {
// Make sure step is updated before we read it.
abortCount = channel->abortCount;
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);
}
__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__ ~ncclLLPrimitives() {
__device__ __forceinline__ ~ncclPrimitives() {
// Save steps for the next operation
for (int i=0; i<NRECV && i<nrecv; i++) saveRecvConn(i);
for (int i=0; i<NSEND && i<nsend; i++) saveSendConn(i);
saveRecvSync();
saveSendSync();
}
};
#include "prims_ll.h"
//#include "prims_ll128.h"
#ifdef ENABLE_PROFILING
#define INIT_COUNTER \
if (tid==0) { t0 = clock64(); ws = LOAD(&(devProf->wait_send_cycle[blockIdx.x])); \
if (tid == 0) { t0 = clock64(); ws = LOAD(&(devProf->wait_send_cycle[blockIdx.x])); \
wr = LOAD(&(devProf->wait_recv_cycle[blockIdx.x])); }
#define ACCUMULATE_COUNTER(prim) \
if (tid==0) { __atomic_fetch_add(&(devProf->prim##_cycle), clock64() - t0 \
if (tid == 0) { __atomic_fetch_add(&(devProf->prim##_cycle), clock64() - t0 \
+ ws - LOAD(&(devProf->wait_send_cycle[blockIdx.x])) \
+ wr - LOAD(&(devProf->wait_recv_cycle[blockIdx.x])), \
__ATOMIC_SEQ_CST); \
src/collectives/device/prims_ll.h
+293
View File
@@ -0,0 +1,293 @@
/*************************************************************************
* Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
template <typename T, class FUNC, int NRECV, int NSEND>
class ncclLLPrimitives {
private:
const int tid;
const int nthreads;
const int wid;
int nrecv = 0;
int nsend = 0;
struct ncclConnInfo* recvConn = NULL;
volatile uint64_t* recvConnHeadPtr = NULL;
uint64_t recvConnHead;
struct ncclConnInfo* sendConn = NULL;
volatile int* sendConnFifoPtr = NULL;
volatile uint64_t* sendConnHeadPtr = NULL;
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;
inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
inline __device__ union ncclLLFifoLine* recvPtr(int i) { return recvBuff[i]+recvOffset(i); }
inline __device__ union ncclLLFifoLine* sendPtr(int i) { return sendBuff[i]+sendOffset(i); }
inline __device__ uint32_t recvFlag(int i) { return NCCL_LL_FLAG(recvStep[i]+1); }
inline __device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(sendStep[i]+1); }
inline __device__ void barrier() {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
asm volatile ("bar.sync 1, %0;" :: "r"(nthreads));
#endif
}
uint32_t mismatch = 0;
const uint64_t opCount;
inline __device__ void checkMismatch(struct ncclConnInfo* conn) {
if (mismatch > 20) {
// We have seen that the peer advanced opcount so many times yet we are still waiting for credit of current op, so it is _most likely_ a mismatch
// Note that we are not using _threadfence_system in LL so the error cannot be asserted
STORE(comm->fatalDevError, ncclDevSuspectedMismatch);
} else if (conn && LOAD(conn->opCountRem) > opCount) {
mismatch += 1;
}
}
uint32_t spins = 0;
uint32_t abort = 0;
inline __device__ int checkAbort(int i, int send) {
spins++;
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = LOAD(comm->abortFlag);
if (wid == i) checkMismatch(send ? sendConn : recvConn);
spins = 0;
}
return abort;
}
inline __device__ void waitSend(int nbytes) {
spins = 0;
mismatch = 0;
if (sendConnHeadPtr) {
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
sendConnHeadCache = LOAD(sendConnHeadPtr);
if (checkAbort(wid, 1)) break;
}
if (sendConnFifoPtr) {
int size = ((sendConnHead & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? NCCL_LL_SLICE_LINES*sizeof(union ncclLLFifoLine) : nbytes;
STORE(sendConnFifoPtr+sendConnHead%NCCL_STEPS, size);
}
sendConnHead += 1;
}
barrier();
}
inline __device__ void incRecv(int i) {
recvStep[i] += 1;
}
inline __device__ void postRecv() {
barrier();
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
}
inline __device__ void incSend(int i, int offset) {
// LL Cleanup : write all flags in the slice to make sure we don't have
// data corruption when flag loops over.
if ((sendStep[i] & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) {
for (int o = offset; o<NCCL_LL_SLICE_LINES; o+=nthreads) storeLL(sendPtr(i)+o, 0, sendFlag(i));
}
sendStep[i]++;
}
__device__ uint64_t readLL(int i, int offset) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
uint32_t flag = recvFlag(i);
uint32_t data1, flag1, data2, flag2;
spins = 0;
mismatch = 0;
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
do {
asm volatile ("flat_load_dwordx4 %0, %1, glc\n"
"s_waitcnt vmcnt(0)\n"
"buffer_wbinvl1_vol\n" : "=v"(i4) : "v"(src));
} while ((i4[1] != flag) || (i4[3] != flag));
uint64_t val64 = (uint64_t)(i4[0]) + (((uint64_t)i4[2]) << 32);
#else
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;
} while ((flag1 != flag) || (flag2 != flag));
uint64_t val64 = data1 + (((uint64_t)data2) << 32);
#endif
return val64;
}
__device__ void storeLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
using Vec = uint32_t __attribute__((ext_vector_type(4)));
Vec i4;
i4[0] = val & 0xffffffff;
i4[1] = flag;
i4[2] = (val >> 32);
i4[3] = flag;
asm volatile ("flat_store_dwordx4 %0, %1, glc\n"
"s_waitcnt vmcnt(0)\n"
"buffer_wbinvl1_vol\n" : : "v"(dst), "v"(i4));
#else
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));
#endif
}
// 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;
}
__device__ void storeAL(uint64_t* dst, uint64_t val, uint32_t nbytes) {
memcpy((char*)dst, (char*)&val, nbytes);
}
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;
// Always waitSend in case of cleanup
if (SEND) waitSend(npack*sizeof(union ncclLLFifoLine));
// Do multiples of 64 bits
#pragma unroll 1
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);
}
}
// 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));
}
}
}
FOR_RECV(incRecv); if (RECV) postRecv();
FOR_SEND(incSend, offset);
}
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
recvBuff[i] = LOAD(&conn->llBuff);
recvStep[i] = LOAD(&conn->step);
if (wid == i) recvConn = conn;
nrecv++;
}
__device__ __forceinline__ void loadRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
recvConnHeadPtr = LOAD(&recvConn->head);
recvConnHead = LOAD(&recvConn->step);
// Update opCount in case we skipped some operations
STORE(recvConn->opCountLoc, opCount);
}
}
__device__ __forceinline__ void loadSendConn(struct ncclConnInfo* conn, int i) {
sendBuff[i] = LOAD(&conn->llBuff);
sendStep[i] = LOAD(&conn->step);
if (wid == i) sendConn = conn;
nsend++;
}
__device__ __forceinline__ void loadSendSync() {
if (tid < nsend) {
sendConnHeadPtr = LOAD(&sendConn->head);
sendConnHeadCache = LOAD(sendConnHeadPtr);
sendConnHead = LOAD(&sendConn->step);
sendConnFifoPtr = LOAD(&sendConn->fifo);
STORE(sendConn->opCountLoc, opCount);
}
}
__device__ __forceinline__ void saveRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
STORE(&recvConn->step, recvConnHead);
STORE(recvConn->opCountLoc, opCount+1);
__threadfence_block();
}
}
__device__ __forceinline__ void saveSendSync() {
if (tid < nsend) {
STORE(&sendConn->step, sendConnHead);
STORE(sendConn->opCountLoc, opCount+1);
__threadfence_block();
}
}
public:
__device__ __forceinline__
ncclLLPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount)
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), opCount(opCount) {
// 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);
loadRecvSync();
loadSendSync();
}
__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() {
// Save steps for the next operation
saveRecvSync();
saveSendSync();
}
};
src/collectives/device/prims_ll128.h
+427
View File
@@ -0,0 +1,427 @@
/*************************************************************************
* Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "op128.h"
#define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS-1)
template <typename T, class FUNC, int NRECV, int NSEND>
class ncclLL128Primitives {
private:
const int tid;
const int nthreads;
const int wid;
const int warp;
const bool flagThread;
int nrecv = 0;
int nsend = 0;
struct ncclConnInfo* recvConn = NULL;
volatile uint64_t* recvConnHeadPtr = NULL;
uint64_t recvConnHead;
struct ncclConnInfo* sendConn = NULL;
volatile int* sendConnFifoPtr = NULL;
volatile uint64_t* sendConnTailPtr = NULL;
uint64_t sendConnTail;
volatile uint64_t* sendConnHeadPtr = NULL;
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;
uint32_t* sync;
inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*NCCL_LL128_SLICE_ELEMS; }
inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*NCCL_LL128_SLICE_ELEMS; }
inline __device__ uint64_t* recvPtr(int i) { return recvBuff[i]+recvOffset(i); }
inline __device__ uint64_t* sendPtr(int i) { return sendBuff[i]+sendOffset(i); }
inline __device__ uint64_t recvFlag(int i) { return recvStep[i]+1; }
inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; }
inline __device__ void barrier() {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
if (NSEND>NRECV) {
asm volatile ("bar.sync 2, %0;" :: "r"(nthreads));
} else {
asm volatile ("bar.sync 3, %0;" :: "r"(nthreads));
}
#endif
}
uint32_t mismatch = 0;
const uint64_t opCount;
inline __device__ void checkMismatch(struct ncclConnInfo* conn) {
if (mismatch > 20) {
// We have seen that the peer advanced opcount so many times yet we are still waiting for credit of current op, so it is _most likely_ a mismatch
// Note that we are not using _threadfence_system in LL so the error cannot be asserted
STORE(comm->fatalDevError, ncclDevSuspectedMismatch);
} else if (conn && LOAD(conn->opCountRem) > opCount) {
mismatch += 1;
}
}
uint32_t spins = 0;
uint32_t abort = 0;
inline __device__ int checkAbort(int i, int send) {
spins++;
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
abort = LOAD(comm->abortFlag);
if (wid == i) checkMismatch(send ? sendConn : recvConn);
spins = 0;
}
return abort;
}
inline __device__ void waitSend(int nbytes) {
spins = 0;
mismatch = 0;
if (sendConnHeadPtr) {
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
sendConnHeadCache = LOAD(sendConnHeadPtr);
if (checkAbort(wid, 1)) break;
}
if (sendConnFifoPtr) {
STORE(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes);
}
sendConnHead += 1;
}
}
inline __device__ void incRecv(int i) {
recvStep[i] += 1;
}
inline __device__ void postRecv() {
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
}
inline __device__ void incSend(int i) {
sendStep[i] += 1;
}
inline __device__ void postSend() {
if (sendConnTailPtr) { __threadfence(); STORE(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);
}
}
#endif
}
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];
}
}
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);
#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]);
}
}
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];
}
}
#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];
/************* 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 **************************/
if (RECV) {
uint64_t flag = recvFlag(0);
uint64_t* ptr = recvPtr(0)+ll128Offset;
bool needReload;
uint64_t v0, v1;
do {
if (wid == 0) STORE(sync, 0);
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);
}
if (needReload) __atomic_fetch_add(sync, 1, __ATOMIC_SEQ_CST);
if (LOAD(sync) == 0) break;
} while (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++) {
uint64_t flag = recvFlag(i);
uint64_t* ptr = recvPtr(i)+ll128Offset;
uint64_t v0, v1;
do {
if (wid == 0) STORE(sync, 0);
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);
}
if (needReload) __atomic_fetch_add(sync, 1, __ATOMIC_SEQ_CST);
if (LOAD(sync) == 0) break;
} while (checkAbort(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]);
}
}
}
/********************** End Recv ************************/
/************************ Send **************************/
if (SEND) {
for (int i=1; i<NSEND && i<nsend; i++) {
int flag = sendFlag(i);
uint64_t* ptr = sendPtr(i)+ll128Offset;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
store128(ptr+u*WARP_SIZE, v[u], flagThread ? flag : v[u+1]);
}
}
int flag = sendFlag(0);
uint64_t* ptr = sendPtr(0)+ll128Offset;
#pragma unroll
for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
store128(ptr+u*WARP_SIZE, v[u], flagThread ? flag : v[u+1]);
}
}
/********************** 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))
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);
int ll128Offset = LL128INC*warp+2*wid;
int elemOffset = ELEMINC*warp;
const int nwarps = nthreads/WARP_SIZE;
if (SEND) waitSend(DIVUP(nelem*sizeof(T), ELEMINC*sizeof(uint64_t))*LL128INC*sizeof(uint64_t));
barrier();
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;
}
barrier();
FOR_SEND(incSend); if (SEND) postSend();
FOR_RECV(incRecv); if (RECV) postRecv();
}
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
recvBuff[i] = LOAD(&conn->ll128Buff);
recvStep[i] = LOAD(&conn->step);
if (wid == i) recvConn = conn;
nrecv++;
}
__device__ __forceinline__ void loadRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
recvConnHeadPtr = LOAD(&recvConn->head);
recvConnHead = LOAD(&recvConn->step);
// Update opCount in case we skipped some operations
STORE(recvConn->opCountLoc, opCount);
}
}
__device__ __forceinline__ void loadSendConn(struct ncclConnInfo* conn, int i) {
sendBuff[i] = LOAD(&conn->ll128Buff);
sendStep[i] = LOAD(&conn->step);
if (wid == i) sendConn = conn;
nsend++;
}
__device__ __forceinline__ void loadSendSync() {
if (tid < nsend) {
sendConnHeadPtr = LOAD(&sendConn->head);
sendConnHeadCache = LOAD(sendConnHeadPtr);
sendConnHead = LOAD(&sendConn->step);
sendConnFifoPtr = LOAD(&sendConn->fifo);
STORE(sendConn->opCountLoc, opCount);
}
if (tid >= nthreads-WARP_SIZE && wid<nsend) {
if (sendConn->fifo) {
sendConnTailPtr = LOAD(&sendConn->tail);
sendConnTail = LOAD(&sendConn->step);
}
}
}
__device__ __forceinline__ void saveRecvSync() {
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
STORE(&recvConn->step, recvConnHead);
STORE(recvConn->opCountLoc, opCount+1);
__threadfence_block();
}
}
__device__ __forceinline__ void saveSendSync() {
if (tid < nsend) {
STORE(&sendConn->step, sendConnHead);
STORE(sendConn->opCountLoc, opCount+1);
__threadfence_block();
}
}
public:
__device__ __forceinline__
ncclLL128Primitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount)
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), warp(tid/WARP_SIZE), flagThread((tid%8)==7), opCount(opCount), shmem(ncclShmem+(threadIdx.x/WARP_SIZE)*NCCL_LL128_SHMEM_ELEMS_PER_THREAD*WARP_SIZE+2*wid) {
// for __any_sync
if (NSEND > NRECV)
sync = channel->sync + 2 + tid/WARP_SIZE;
else
sync = channel->sync + tid/WARP_SIZE;
// 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);
loadRecvSync();
loadSendSync();
}
__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() {
// Save steps for the next operation
saveRecvSync();
saveSendSync();
}
};
src/collectives/device/reduce.cu
+3 -7
View File
@@ -1,15 +1,11 @@
/*************************************************************************
* Copyright (c) 2015-2018, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "common.h"
#include "reduce.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL2(ncclReduce, sum, FuncSum, ncclCollReduce, ncclSum);
IMPL_COLL2(ncclReduce, prod, FuncProd, ncclCollReduce, ncclProd);
IMPL_COLL2(ncclReduce, min, FuncMin, ncclCollReduce, ncclMin);
IMPL_COLL2(ncclReduce, max, FuncMax, ncclCollReduce, ncclMax);
IMPL_COLL_R(ncclReduce, ncclCollReduce);
src/collectives/device/reduce.h
+49 -2
View File
@@ -13,7 +13,7 @@ template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceRingKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
@@ -32,7 +32,7 @@ __device__ void ncclReduceRingKernel(struct CollectiveArgs* args) {
T * __restrict__ thisOutput = (T*)args->ThisOutput;
ncclPrimitives<UNROLL, REDUCE_CHUNKSTEPS/REDUCE_SLICESTEPS, REDUCE_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
prims(tid, args->nThreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,args->nChannels));
@@ -98,3 +98,50 @@ __device__ void ncclReduceRingLLKernel(struct CollectiveArgs* args) {
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceTreeLLKernel(struct CollectiveArgs* args) { }
#include "prims_ll128.h"
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceRingLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int bid = args->bid;
const int nthreads = args->nThreads;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, channel, comm, args->opCount);
const ssize_t size = args->N;
const int rank = comm->rank;
const int nranks = comm->nRanks;
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->root;
ssize_t chunkSize = (NCCL_LL128_ELEMS_PER_THREAD*nthreads*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 = args->nChannels*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, args->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);
}
}
}
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceTreeLL128Kernel(struct CollectiveArgs* args) { }
src/collectives/device/reduce_kernel.h
-18
View File
@@ -260,15 +260,12 @@ static __device__ uint32_t addChar4(const uint32_t x, const uint32_t y) {
template<>
struct FuncSum<int8_t> {
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vadd4.s32.s32.s32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
return rv;
#else
return addChar4(x, y);
#endif
#endif
}
__device__ int8_t operator()(const int8_t x, const int8_t y) const {
@@ -278,15 +275,12 @@ struct FuncSum<int8_t> {
template<>
struct FuncSum<uint8_t> {
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vadd4.u32.u32.u32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
return rv;
#else
return addChar4(x, y);
#endif
#endif
}
__device__ uint8_t operator()(const uint8_t x, const uint8_t y) const {
@@ -330,8 +324,6 @@ template<>
struct FuncMax<int8_t> {
union converter { uint32_t storage; char4 a; };
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vmax4.s32.s32.s32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
@@ -345,7 +337,6 @@ struct FuncMax<int8_t> {
cr.a.z = max(cx.a.z, cy.a.z);
cr.a.w = max(cx.a.w, cy.a.w);
return cr.storage;
#endif
#endif
}
__device__ int8_t operator()(const int8_t x, const int8_t y) const {
@@ -356,8 +347,6 @@ template<>
struct FuncMax<uint8_t> {
union converter { uint32_t storage; uchar4 a; };
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vmax4.u32.u32.u32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
@@ -371,7 +360,6 @@ struct FuncMax<uint8_t> {
cr.a.z = max(cx.a.z, cy.a.z);
cr.a.w = max(cx.a.w, cy.a.w);
return cr.storage;
#endif
#endif
}
__device__ uint8_t operator()(const uint8_t x, const uint8_t y) const {
@@ -383,8 +371,6 @@ template<>
struct FuncMin<int8_t> {
union converter { uint32_t storage; char4 a; };
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vmin4.s32.s32.s32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
@@ -398,7 +384,6 @@ struct FuncMin<int8_t> {
cr.a.z = min(cx.a.z, cy.a.z);
cr.a.w = min(cx.a.w, cy.a.w);
return cr.storage;
#endif
#endif
}
__device__ int8_t operator()(const int8_t x, const int8_t y) const {
@@ -409,8 +394,6 @@ template<>
struct FuncMin<uint8_t> {
union converter { uint32_t storage; uchar4 a; };
__device__ uint32_t operator()(const uint32_t x, const uint32_t y) const {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#else
#if (__CUDA_ARCH__ >= 300) && (__CUDA_ARCH__ < 500)
int32_t rv, z=0;
asm("vmin4.u32.u32.u32 %0, %1, %2, %3;" : "=r"(rv) : "r"(x), "r"(y), "r"(z));
@@ -424,7 +407,6 @@ struct FuncMin<uint8_t> {
cr.a.z = min(cx.a.z, cy.a.z);
cr.a.w = min(cx.a.w, cy.a.w);
return cr.storage;
#endif
#endif
}
__device__ uint8_t operator()(const uint8_t x, const uint8_t y) const {
src/collectives/device/reduce_scatter.cu
+3 -7
View File
@@ -1,15 +1,11 @@
/*************************************************************************
* Copyright (c) 2015-2018, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "common.h"
#include "reduce_scatter.h"
#include "common.h"
#include "collectives.h"
IMPL_COLL2(ncclReduceScatter, sum, FuncSum, ncclCollReduceScatter, ncclSum);
IMPL_COLL2(ncclReduceScatter, prod, FuncProd, ncclCollReduceScatter, ncclProd);
IMPL_COLL2(ncclReduceScatter, min, FuncMin, ncclCollReduceScatter, ncclMin);
IMPL_COLL2(ncclReduceScatter, max, FuncMax, ncclCollReduceScatter, ncclMax);
IMPL_COLL_R(ncclReduceScatter, ncclCollReduceScatter);
src/collectives/device/reduce_scatter.h
+66 -3
View File
@@ -13,7 +13,7 @@ template<int UNROLL, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceScatterRingKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
const int nthreads = args->nThreads;
const int bid = args->bid;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
@@ -21,7 +21,7 @@ __device__ void ncclReduceScatterRingKernel(struct CollectiveArgs* args) {
const ssize_t size = args->N;
const int nranks = comm->nRanks;
const int stepSize = channel->buffSize / (sizeof(T)*NCCL_STEPS);
const int chunkSize = stepSize * ALLREDUCE_CHUNKSTEPS;
const int chunkSize = stepSize * REDUCESCATTER_CHUNKSTEPS;
const ssize_t loopSize = args->nChannels*(ssize_t)chunkSize;
// Compute pointers
@@ -29,7 +29,7 @@ __device__ void ncclReduceScatterRingKernel(struct CollectiveArgs* args) {
T * __restrict__ thisOutput = (T*)args->ThisOutput;
ncclPrimitives<UNROLL, REDUCESCATTER_CHUNKSTEPS/REDUCESCATTER_SLICESTEPS, REDUCESCATTER_SLICESTEPS, T, 1, 1, FUNC>
prims(tid, nthreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
prims(tid, args->nThreads, &ring->prev, &ring->next, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
int realChunkSize = min(chunkSize, DIVUP(size-gridOffset,args->nChannels));
@@ -126,3 +126,66 @@ __device__ void ncclReduceScatterRingLLKernel(struct CollectiveArgs* args) {
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceScatterTreeLLKernel(struct CollectiveArgs* args) { }
#include "prims_ll128.h"
template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceScatterRingLL128Kernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x;
const int bid = args->bid;
const int nthreads = args->nThreads;
struct ncclDevComm* comm = args->comm;
struct ncclChannel* channel = comm->channels+blockIdx.x;
struct ncclRing* ring = &channel->ring;
ncclLL128Primitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &ring->prev, &ring->next, channel, comm, args->opCount);
const ssize_t size = args->N;
//const int rank = comm->rank;
const int nranks = comm->nRanks;
ssize_t chunkSize = (NCCL_LL128_ELEMS_PER_THREAD*nthreads*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 ssize_t loopSize = args->nChannels*chunkSize;
// Compute pointers
const T * __restrict__ thisInput = (const T*)args->ThisInput;
T * __restrict__ thisOutput = (T*)args->ThisOutput;
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
chunkSize = min(DIVUP(size-gridOffset, args->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<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclReduceScatterTreeLL128Kernel(struct CollectiveArgs* args) { }