Split primitive class to smaller structures
This commit is contained in:
Wenkai Du
@@ -127,7 +127,8 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
|
||||
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, args->nThreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount);
|
||||
ncclPrimitivesRecvData<T, NCCL_MAX_TREE_ARITY> recvData;
|
||||
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, recvData);
|
||||
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
|
||||
// Up
|
||||
ssize_t offset = gridOffset + bid*chunkSize;
|
||||
@@ -145,7 +146,8 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
|
||||
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, args->nThreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount);
|
||||
ncclPrimitivesSendData<T, NCCL_MAX_TREE_ARITY> sendData;
|
||||
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, sendData);
|
||||
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
|
||||
// Down
|
||||
ssize_t offset = gridOffset + bid*chunkSize;
|
||||
@@ -260,7 +262,8 @@ __device__ void ncclAllReduceTreeLLKernel(struct CollectiveArgs* args) {
|
||||
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);
|
||||
ncclLLPrimitivesRecvData<T, NCCL_MAX_TREE_ARITY> recvData;
|
||||
ncclLLPrimitives<T, FUNC, NCCL_MAX_TREE_ARITY, 1> LLprims(tid, nthreads, tree->down, &tree->up, channel, comm, args->opCount, recvData);
|
||||
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
|
||||
// Up
|
||||
ssize_t offset = gridOffset + bid*chunkSize;
|
||||
@@ -278,7 +281,8 @@ __device__ void ncclAllReduceTreeLLKernel(struct CollectiveArgs* args) {
|
||||
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);
|
||||
ncclLLPrimitivesSendData<T, NCCL_MAX_TREE_ARITY> sendData;
|
||||
ncclLLPrimitives<T, FUNC, 1, NCCL_MAX_TREE_ARITY> LLprims(tid, nthreads, &tree->up, tree->down, channel, comm, args->opCount, sendData);
|
||||
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
|
||||
// Down
|
||||
ssize_t offset = gridOffset + bid*chunkSize;
|
||||
|
||||
@@ -32,6 +32,41 @@
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
template <typename T, int NRECV>
|
||||
class ncclPrimitivesRecvData {
|
||||
public:
|
||||
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
|
||||
|
||||
uint64_t recvStep[NRECV];
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
const T* recvDirectBuff[NRECV];
|
||||
#endif
|
||||
const T* recvBuff[NRECV];
|
||||
};
|
||||
|
||||
template <typename T, int NSEND>
|
||||
class ncclPrimitivesSendData {
|
||||
public:
|
||||
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 sendStep[NSEND];
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
const T* sendDirectBuff[NRECV];
|
||||
#endif
|
||||
T* sendBuff[NSEND];
|
||||
};
|
||||
|
||||
// Implementation of primitive types
|
||||
template <int UNROLL, int SLICESPERCHUNK, int SLICESTEPS, typename T, int NRECV, int NSEND, class FUNC>
|
||||
class ncclPrimitives {
|
||||
@@ -42,35 +77,18 @@ class ncclPrimitives {
|
||||
const int stepSize;
|
||||
int nrecv = 0;
|
||||
int nsend = 0;
|
||||
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
|
||||
typename std::conditional<NRECV == NCCL_MAX_TREE_ARITY,
|
||||
ncclPrimitivesRecvData<T, NRECV>&, ncclPrimitivesRecvData<T, NRECV>>::type r;
|
||||
typename std::conditional<NSEND == NCCL_MAX_TREE_ARITY,
|
||||
ncclPrimitivesSendData<T, NSEND>&, ncclPrimitivesSendData<T, NSEND>>::type s;
|
||||
|
||||
uint64_t recvStep[NRECV];
|
||||
uint64_t sendStep[NSEND];
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
const T* recvDirectBuff[NRECV];
|
||||
T* sendDirectBuff[NSEND];
|
||||
#endif
|
||||
const T* recvBuff[NRECV];
|
||||
T* sendBuff[NSEND];
|
||||
const struct ncclDevComm* comm;
|
||||
|
||||
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); }
|
||||
inline __device__ int recvOffset(int i) { return (r.recvStep[i]%NCCL_STEPS)*stepSize; }
|
||||
inline __device__ int sendOffset(int i) { return (s.sendStep[i]%NCCL_STEPS)*stepSize; }
|
||||
inline __device__ const T* recvPtr(int i) { return ((const T*)r.recvBuff[i])+recvOffset(i); }
|
||||
inline __device__ T* sendPtr(int i) { return ((T*)s.sendBuff[i])+sendOffset(i); }
|
||||
|
||||
inline __device__ void barrier() {
|
||||
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
|
||||
@@ -99,7 +117,7 @@ class ncclPrimitives {
|
||||
spins++;
|
||||
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
|
||||
abort = LOAD(comm->abortFlag);
|
||||
if (wid == i) checkMismatch(send ? sendConn : recvConn);
|
||||
if (wid == i) checkMismatch(send ? s.sendConn : r.recvConn);
|
||||
spins = 0;
|
||||
}
|
||||
return abort;
|
||||
@@ -108,65 +126,65 @@ class ncclPrimitives {
|
||||
inline __device__ void waitSend(int nbytes) {
|
||||
spins = 0;
|
||||
mismatch = 0;
|
||||
if (sendConnHeadPtr) {
|
||||
if (s.sendConnHeadPtr) {
|
||||
#ifdef ENABLE_PROFILING
|
||||
auto devProf = comm->devProf;
|
||||
uint64_t t0 = __rtc64();
|
||||
#endif
|
||||
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + SLICESTEPS) {
|
||||
sendConnHeadCache = LOAD(sendConnHeadPtr);
|
||||
while (s.sendConnHeadCache + NCCL_STEPS < s.sendConnHead + SLICESTEPS) {
|
||||
s.sendConnHeadCache = LOAD(s.sendConnHeadPtr);
|
||||
if (checkAbort(wid, 1)) break;
|
||||
}
|
||||
#ifdef ENABLE_PROFILING
|
||||
__atomic_fetch_add(&devProf->wait_send_cycle[blockIdx.x], __rtc64() - t0, __ATOMIC_SEQ_CST);
|
||||
#endif
|
||||
if (sendConnFifoPtr) {
|
||||
STORE(sendConnFifoPtr+sendConnHead%NCCL_STEPS, nbytes);
|
||||
if (s.sendConnFifoPtr) {
|
||||
STORE(s.sendConnFifoPtr+s.sendConnHead%NCCL_STEPS, nbytes);
|
||||
}
|
||||
sendConnHead += SLICESTEPS;
|
||||
s.sendConnHead += SLICESTEPS;
|
||||
}
|
||||
}
|
||||
|
||||
inline __device__ void waitRecv() {
|
||||
spins = 0;
|
||||
mismatch = 0;
|
||||
if (recvConnTailPtr) {
|
||||
if (r.recvConnTailPtr) {
|
||||
#ifdef ENABLE_PROFILING
|
||||
auto devProf = comm->devProf;
|
||||
uint64_t t0 = __rtc64();
|
||||
#endif
|
||||
while (recvConnTailCache < recvConnTail + SLICESTEPS) {
|
||||
recvConnTailCache = LOAD(recvConnTailPtr);
|
||||
while (r.recvConnTailCache < r.recvConnTail + SLICESTEPS) {
|
||||
r.recvConnTailCache = LOAD(r.recvConnTailPtr);
|
||||
if (checkAbort(wid, 0)) break;
|
||||
}
|
||||
#ifdef ENABLE_PROFILING
|
||||
__atomic_fetch_add(&devProf->wait_recv_cycle[blockIdx.x], __rtc64() - t0, __ATOMIC_SEQ_CST);
|
||||
#endif
|
||||
recvConnTail += SLICESTEPS;
|
||||
r.recvConnTail += SLICESTEPS;
|
||||
}
|
||||
}
|
||||
|
||||
inline __device__ void incRecv(int i) {
|
||||
recvStep[i] += SLICESTEPS;
|
||||
r.recvStep[i] += SLICESTEPS;
|
||||
}
|
||||
inline __device__ void postRecv() {
|
||||
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += SLICESTEPS);
|
||||
if (r.recvConnHeadPtr) STORE(r.recvConnHeadPtr, r.recvConnHead += SLICESTEPS);
|
||||
}
|
||||
|
||||
inline __device__ void incSend(int i) {
|
||||
sendStep[i] += SLICESTEPS;
|
||||
s.sendStep[i] += SLICESTEPS;
|
||||
}
|
||||
inline __device__ void postSend() {
|
||||
if (sendConnTailPtr) {
|
||||
if (sendConn->next_hdp_reg) STORE(sendConn->next_hdp_reg, 0x1);
|
||||
STORE(sendConnTailPtr, sendConnTail += SLICESTEPS);
|
||||
if (s.sendConnTailPtr) {
|
||||
if (s.sendConn->next_hdp_reg) STORE(s.sendConn->next_hdp_reg, 0x1);
|
||||
STORE(s.sendConnTailPtr, s.sendConnTail += SLICESTEPS);
|
||||
}
|
||||
}
|
||||
|
||||
template <int DIRECTRECV>
|
||||
inline __device__ const T* directRecvPtr(int i, int directOffset) {
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
return DIRECTRECV && recvDirectBuff[i] ? recvDirectBuff[i]+directOffset : recvPtr(i);
|
||||
return DIRECTRECV && r.recvDirectBuff[i] ? r.recvDirectBuff[i]+directOffset : recvPtr(i);
|
||||
#else
|
||||
return recvPtr(i);
|
||||
#endif
|
||||
@@ -175,7 +193,7 @@ class ncclPrimitives {
|
||||
template <int DIRECTSEND>
|
||||
inline __device__ T* directSendPtr(int i, int directOffset) {
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
return DIRECTSEND && sendDirectBuff[i] ? sendDirectBuff[i]+directOffset : sendPtr(i);
|
||||
return DIRECTSEND && s.sendDirectBuff[i] ? s.sendDirectBuff[i]+directOffset : sendPtr(i);
|
||||
#else
|
||||
return sendPtr(i);
|
||||
#endif
|
||||
@@ -184,7 +202,7 @@ class ncclPrimitives {
|
||||
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;
|
||||
return DIRECTRECV && r.recvDirectBuff[i] ? directInc : sliceInc;
|
||||
#else
|
||||
return sliceInc;
|
||||
#endif
|
||||
@@ -193,7 +211,7 @@ inline __device__ int directRecvInc(int i, int directInc, int sliceInc) {
|
||||
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;
|
||||
return DIRECTSEND && s.sendDirectBuff[i] ? directInc : sliceInc;
|
||||
#else
|
||||
return sliceInc;
|
||||
#endif
|
||||
@@ -228,7 +246,7 @@ inline __device__ int directSendInc(int i, int directInc, int sliceInc) {
|
||||
if (realSize > 0) {
|
||||
barrier();
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
if (DIRECTRECV && recvDirectBuff[0]) {
|
||||
if (DIRECTRECV && r.recvDirectBuff[0]) {
|
||||
// We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
|
||||
if (SEND) {
|
||||
ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, NSEND>(tid, nthreads, 1, srcs, nsend, dsts+1, realSize);
|
||||
@@ -260,84 +278,81 @@ inline __device__ int directSendInc(int i, int directInc, int sliceInc) {
|
||||
}
|
||||
|
||||
__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);
|
||||
r.recvBuff[i] = (const T*)LOAD(&conn->buff);
|
||||
r.recvStep[i] = LOAD(&conn->step);
|
||||
r.recvStep[i] = ROUNDUP(r.recvStep[i], SLICESPERCHUNK*SLICESTEPS);
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
recvDirectBuff[i] = NULL;
|
||||
r.recvDirectBuff[i] = NULL;
|
||||
if (directBuff && LOAD(&conn->direct)) {
|
||||
recvDirectBuff[i] = directBuff;
|
||||
r.recvDirectBuff[i] = 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
|
||||
if (wid == i) r.recvConn = conn;
|
||||
if (wid == i) r.recvConnTail = r.recvConnHead = r.recvStep[i]; // Make sure we set this after rounding up
|
||||
nrecv++;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void loadRecvSync() {
|
||||
if (tid >= WARP_SIZE && tid < 2*WARP_SIZE && wid<nrecv) {
|
||||
recvConnTailPtr = LOAD(&recvConn->tail);
|
||||
recvConnTailCache = LOAD(recvConnTailPtr);
|
||||
r.recvConnTailPtr = LOAD(&r.recvConn->tail);
|
||||
r.recvConnTailCache = LOAD(r.recvConnTailPtr);
|
||||
}
|
||||
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
|
||||
recvConnHeadPtr = LOAD(&recvConn->head);
|
||||
r.recvConnHeadPtr = LOAD(&r.recvConn->head);
|
||||
// Return credits in case we rounded up.
|
||||
STORE(recvConnHeadPtr, recvConnHead);
|
||||
STORE(r.recvConnHeadPtr, r.recvConnHead);
|
||||
// Update opCount in case we skipped some operations
|
||||
STORE(recvConn->opCountLoc, opCount);
|
||||
STORE(r.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);
|
||||
s.sendBuff[i] = (T*)LOAD(&conn->buff);
|
||||
s.sendStep[i] = LOAD(&conn->step);
|
||||
s.sendStep[i] = ROUNDUP(s.sendStep[i], SLICESPERCHUNK*SLICESTEPS);
|
||||
#if defined(RCCL_USE_DIRECT_BUFFER)
|
||||
sendDirectBuff[i] = NULL;
|
||||
s.sendDirectBuff[i] = NULL;
|
||||
if (directBuff && LOAD(&conn->direct)) {
|
||||
void* volatile* ptr = LOAD(&conn->ptrExchange);
|
||||
while ((sendDirectBuff[i] = (T*)(LOAD(ptr))) == NULL);
|
||||
while ((s.sendDirectBuff[i] = (T*)(LOAD(ptr))) == NULL);
|
||||
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
|
||||
if (wid == i) s.sendConn = conn;
|
||||
if (wid == i) s.sendConnTail = s.sendConnHead = s.sendStep[i]; // Make sure we set this after rounding up
|
||||
nsend++;
|
||||
}
|
||||
__device__ void loadSendSync() {
|
||||
if (tid < nsend) {
|
||||
sendConnHeadPtr = LOAD(&sendConn->head);
|
||||
sendConnHeadCache = LOAD(sendConnHeadPtr);
|
||||
sendConnFifoPtr = LOAD(&sendConn->fifo);
|
||||
STORE(sendConn->opCountLoc, opCount);
|
||||
s.sendConnHeadPtr = LOAD(&s.sendConn->head);
|
||||
s.sendConnHeadCache = LOAD(s.sendConnHeadPtr);
|
||||
s.sendConnFifoPtr = LOAD(&s.sendConn->fifo);
|
||||
STORE(s.sendConn->opCountLoc, opCount);
|
||||
}
|
||||
if (tid >= nthreads-WARP_SIZE && wid < nsend) {
|
||||
sendConnTailPtr = LOAD(&sendConn->tail);
|
||||
s.sendConnTailPtr = LOAD(&s.sendConn->tail);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ void saveRecvSync() {
|
||||
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
|
||||
STORE(&recvConn->step, recvConnHead);
|
||||
STORE(recvConn->opCountLoc, opCount+1);
|
||||
STORE(&r.recvConn->step, r.recvConnHead);
|
||||
STORE(r.recvConn->opCountLoc, opCount+1);
|
||||
__threadfence_system();
|
||||
}
|
||||
}
|
||||
|
||||
__device__ void saveSendSync() {
|
||||
if (tid < nsend) {
|
||||
STORE(&sendConn->step, sendConnHead);
|
||||
STORE(sendConn->opCountLoc, opCount+1);
|
||||
STORE(&s.sendConn->step, s.sendConnHead);
|
||||
STORE(s.sendConn->opCountLoc, opCount+1);
|
||||
__threadfence_system();
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
__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), wid(tid%WARP_SIZE), stepSize(stepSize), opCount(opCount) {
|
||||
inline __device__ void init(int* recvPeers, int* sendPeers, struct ncclChannel* channel) {
|
||||
// Make sure step is updated before we read it.
|
||||
barrier();
|
||||
|
||||
@@ -347,6 +362,25 @@ inline __device__ int directSendInc(int i, int directInc, int sliceInc) {
|
||||
loadSendSync();
|
||||
}
|
||||
|
||||
public:
|
||||
__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), wid(tid%WARP_SIZE), stepSize(stepSize), opCount(opCount) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__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, ncclPrimitivesRecvData<T, NRECV>& r)
|
||||
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), stepSize(stepSize), opCount(opCount), r(r) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__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, ncclPrimitivesSendData<T, NSEND>& s)
|
||||
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), stepSize(stepSize), opCount(opCount), s(s) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void
|
||||
send(const T* src, int nelem) {
|
||||
GenericOp<0, 0, 0, 1, 1, 0>(src, NULL, nelem, 0);
|
||||
|
||||
@@ -5,6 +5,28 @@
|
||||
* See LICENSE.txt for license information
|
||||
************************************************************************/
|
||||
|
||||
template <typename T, int NRECV>
|
||||
class ncclLLPrimitivesRecvData {
|
||||
public:
|
||||
struct ncclConnInfo* recvConn = NULL;
|
||||
volatile uint64_t* recvConnHeadPtr = NULL;
|
||||
uint64_t recvConnHead;
|
||||
uint64_t recvStep[NRECV];
|
||||
union ncclLLFifoLine* recvBuff[NRECV];
|
||||
};
|
||||
|
||||
template <typename T, int NSEND>
|
||||
class ncclLLPrimitivesSendData {
|
||||
public:
|
||||
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 sendStep[NSEND];
|
||||
union ncclLLFifoLine* sendBuff[NSEND];
|
||||
};
|
||||
|
||||
template <typename T, class FUNC, int NRECV, int NSEND>
|
||||
class ncclLLPrimitives {
|
||||
private:
|
||||
@@ -13,28 +35,19 @@ class ncclLLPrimitives {
|
||||
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); }
|
||||
typename std::conditional<NRECV == NCCL_MAX_TREE_ARITY,
|
||||
ncclLLPrimitivesRecvData<T, NRECV>&, ncclLLPrimitivesRecvData<T, NRECV>>::type r;
|
||||
typename std::conditional<NSEND == NCCL_MAX_TREE_ARITY,
|
||||
ncclLLPrimitivesSendData<T, NSEND>&, ncclLLPrimitivesSendData<T, NSEND>>::type s;
|
||||
|
||||
inline __device__ int recvOffset(int i) { return (r.recvStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
|
||||
inline __device__ int sendOffset(int i) { return (s.sendStep[i]%NCCL_STEPS)*NCCL_LL_SLICE_LINES; }
|
||||
inline __device__ union ncclLLFifoLine* recvPtr(int i) { return r.recvBuff[i]+recvOffset(i); }
|
||||
inline __device__ union ncclLLFifoLine* sendPtr(int i) { return s.sendBuff[i]+sendOffset(i); }
|
||||
inline __device__ uint32_t recvFlag(int i) { return NCCL_LL_FLAG(r.recvStep[i]+1); }
|
||||
inline __device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(s.sendStep[i]+1); }
|
||||
|
||||
inline __device__ void barrier() {
|
||||
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
|
||||
@@ -64,7 +77,7 @@ class ncclLLPrimitives {
|
||||
spins++;
|
||||
if (abort == 0 && spins == SPINS_BEFORE_CHECK_ABORT) {
|
||||
abort = LOAD(comm->abortFlag);
|
||||
if (wid == i) checkMismatch(send ? sendConn : recvConn);
|
||||
if (wid == i) checkMismatch(send ? s.sendConn : r.recvConn);
|
||||
spins = 0;
|
||||
}
|
||||
return abort;
|
||||
@@ -73,35 +86,35 @@ class ncclLLPrimitives {
|
||||
inline __device__ void waitSend(int nbytes) {
|
||||
spins = 0;
|
||||
mismatch = 0;
|
||||
if (sendConnHeadPtr) {
|
||||
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
|
||||
sendConnHeadCache = LOAD(sendConnHeadPtr);
|
||||
if (s.sendConnHeadPtr) {
|
||||
while (s.sendConnHeadCache + NCCL_STEPS < s.sendConnHead + 1) {
|
||||
s.sendConnHeadCache = LOAD(s.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);
|
||||
if (s.sendConnFifoPtr) {
|
||||
int size = ((s.sendConnHead & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? NCCL_LL_SLICE_LINES*sizeof(union ncclLLFifoLine) : nbytes;
|
||||
STORE(s.sendConnFifoPtr+s.sendConnHead%NCCL_STEPS, size);
|
||||
}
|
||||
sendConnHead += 1;
|
||||
s.sendConnHead += 1;
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
inline __device__ void incRecv(int i) {
|
||||
recvStep[i] += 1;
|
||||
r.recvStep[i] += 1;
|
||||
}
|
||||
inline __device__ void postRecv() {
|
||||
barrier();
|
||||
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
|
||||
if (r.recvConnHeadPtr) STORE(r.recvConnHeadPtr, r.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) {
|
||||
if ((s.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]++;
|
||||
s.sendStep[i]++;
|
||||
}
|
||||
|
||||
__device__ uint64_t readLL(int i, int offset) {
|
||||
@@ -198,56 +211,53 @@ class ncclLLPrimitives {
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
|
||||
recvBuff[i] = LOAD(&conn->llBuff);
|
||||
recvStep[i] = LOAD(&conn->step);
|
||||
if (wid == i) recvConn = conn;
|
||||
r.recvBuff[i] = LOAD(&conn->llBuff);
|
||||
r.recvStep[i] = LOAD(&conn->step);
|
||||
if (wid == i) r.recvConn = conn;
|
||||
nrecv++;
|
||||
}
|
||||
__device__ __forceinline__ void loadRecvSync() {
|
||||
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
|
||||
recvConnHeadPtr = LOAD(&recvConn->head);
|
||||
recvConnHead = LOAD(&recvConn->step);
|
||||
r.recvConnHeadPtr = LOAD(&r.recvConn->head);
|
||||
r.recvConnHead = LOAD(&r.recvConn->step);
|
||||
// Update opCount in case we skipped some operations
|
||||
STORE(recvConn->opCountLoc, opCount);
|
||||
STORE(r.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;
|
||||
s.sendBuff[i] = LOAD(&conn->llBuff);
|
||||
s.sendStep[i] = LOAD(&conn->step);
|
||||
if (wid == i) s.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);
|
||||
s.sendConnHeadPtr = LOAD(&s.sendConn->head);
|
||||
s.sendConnHeadCache = LOAD(s.sendConnHeadPtr);
|
||||
s.sendConnHead = LOAD(&s.sendConn->step);
|
||||
s.sendConnFifoPtr = LOAD(&s.sendConn->fifo);
|
||||
STORE(s.sendConn->opCountLoc, opCount);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void saveRecvSync() {
|
||||
if (tid >= nthreads-WARP_SIZE && wid < nrecv) {
|
||||
STORE(&recvConn->step, recvConnHead);
|
||||
STORE(recvConn->opCountLoc, opCount+1);
|
||||
STORE(&r.recvConn->step, r.recvConnHead);
|
||||
STORE(r.recvConn->opCountLoc, opCount+1);
|
||||
__threadfence_block();
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void saveSendSync() {
|
||||
if (tid < nsend) {
|
||||
STORE(&sendConn->step, sendConnHead);
|
||||
STORE(sendConn->opCountLoc, opCount+1);
|
||||
STORE(&s.sendConn->step, s.sendConnHead);
|
||||
STORE(s.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) {
|
||||
inline __device__ void init(int* recvPeers, int* sendPeers, struct ncclChannel* channel) {
|
||||
// Make sure step is updated before we read it.
|
||||
barrier();
|
||||
|
||||
@@ -257,6 +267,25 @@ class ncclLLPrimitives {
|
||||
loadSendSync();
|
||||
}
|
||||
|
||||
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) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
ncclLLPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount, ncclLLPrimitivesRecvData<T, NRECV>& r)
|
||||
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), opCount(opCount), r(r) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
ncclLLPrimitives(const int tid, const int nthreads, int* recvPeers, int* sendPeers, struct ncclChannel* channel, struct ncclDevComm* comm, const uint64_t opCount, ncclLLPrimitivesSendData<T, NSEND>& s)
|
||||
: comm(comm), tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), opCount(opCount), s(s) {
|
||||
init(recvPeers, sendPeers, channel);
|
||||
}
|
||||
|
||||
__device__ void send(const T* src, int nelem) {
|
||||
return LLGenericOp<0, 1, 1, 0>(src, NULL, nelem);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user