Merge remote-tracking branch 'nccl/master' into develop

This commit is contained in:
BertanDogancay
2025-06-20 07:53:59 -05:00
136 changed files with 8510 additions and 5421 deletions
+239 -211
View File
@@ -12,6 +12,7 @@
#include "nccl.h"
#include "nccl_common.h"
#include "device.h"
#define NCCL_MAX_NET_SIZE (1024*1024*1024L) // Rather than send INT_MAX which is 2G-1, send a power of two.
// CHUNKSIZE must be a multiple of SLICESIZE
@@ -396,6 +397,42 @@ public:
~RingBCAlgorithm() {}
};
#if !defined (__CUDA_ARCH__) || __CUDA_ARCH__ >= 600
// #include <cuda/atomic>
#endif
// Need a power of two to ensure it divides by parallelFactor (which is also a power of two)
#define NCCL_PAT_NWORKERS 512
static constexpr int PatUsed = 0x1,
PatSkipped = 0x2;
struct ncclPatStep {
int recvDim, sendDim, recvOffset, sendOffset, stepOffset, postRecv, postSend, nelem, last, flags;
size_t inpIx, outIx;
};
struct ncclPatPeer {
uint64_t step;
struct ncclConnInfo* conn;
struct ncclConnFifo* connFifo;
void* buff;
uint64_t *headPtr;
uint64_t *tailPtr;
uint64_t stepCache;
long long int accSize;
int connStepSize;
};
#define NCCL_SHMEM_PAT_STEPS 32
struct ncclPatShmem {
struct ncclPatStep patSteps[NCCL_SHMEM_PAT_STEPS];
int parallelFactor;
long long int localAccSize;
struct ncclPatPeer sendDims[32]; // Should cover 2^32 ranks
struct ncclPatPeer recvDims[32];
};
template<typename T>
class PatRSAlgorithm{
size_t offset;
@@ -408,18 +445,17 @@ class PatRSAlgorithm{
int nrPow2;
int postFreq;
int lastA;
int parallelFactor;
int aggFactor;
int as; // aggregated steps
int a; // step inside aggregated step
int sendSkipped; // number of skipped steps during aggregation
int recvSkipped; // number of skipped steps during aggregation
int phase2recv; // receive offset for phase 2
int stepOffset;
int aggDelta;
int scale;
int phase;
__device__ __host__ int min(int a, int b) {
__device__ __host__ ssize_t min(ssize_t a, ssize_t b) {
return (a<b)?a:b;
}
@@ -447,16 +483,16 @@ class PatRSAlgorithm{
__device__ __host__ void resetA() {
a = 0;
sendSkipped = recvSkipped = 0;
sendSkipped = stepOffset = 0;
lastA = aggFactor;
if (phase >= 2) lastA /= 2*scale;
if (phase == 4) lastA = 1;
}
__device__ __host__ void reset() {
nelem = getNelem();
phase = 0;
scale = 1;
phase2recv = 0;
as = aggDelta - 1;
resetA();
}
@@ -479,8 +515,9 @@ class PatRSAlgorithm{
}
public:
__device__ __host__ PatRSAlgorithm(int stepSize, int stepDepth, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
__device__ __host__ PatRSAlgorithm(int stepSize, int stepDepth, int maxParallelFactor, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
offset(offset), end(end), count(count), chunkCount(chunkCount), rank(rank), nranks(nranks) {
parallelFactor = maxParallelFactor;
aggDelta = nrPow2 = (1<<log2Up(nranks));
aggFactor = 1;
@@ -490,6 +527,7 @@ public:
aggDelta /= 2;
}
postFreq = aggFactor;
if (postFreq < parallelFactor) parallelFactor = postFreq;
int d = stepDepth;
while (d > 1 && aggFactor < nranks/2) {
d /= 2;
@@ -500,160 +538,151 @@ public:
reset();
}
__device__ __host__ void getNextOp(int &recvDim, int &sendDim, size_t &inpIx, size_t &outIx, int &recvOffset, int &sendOffset, int &sendStepOffset, int &nelemOut, int &postRecv, int &postSend, int &last) {
restart:
last = 0;
nelemOut = nelem;
outIx = offset;
__device__ __host__ int getParallelFactor() {
return parallelFactor;
}
__device__ __host__ void getNextOp(struct ncclPatStep* ps) {
ps->last = 0;
ps->nelem = nelem;
ps->outIx = offset;
ps->stepOffset = stepOffset;
int skip = 0;
//printf("Phase %d as %d/%d a %d/%d scale %d\n", phase, as, aggDelta, a, lastA, scale);
if (phase == 0) {
if (a >= lastA) {
skip = 1;
} else if (phase == 0) {
int s = mirrorInvert(a, lastA)*aggDelta + as;
if (s >= nranks) skip = 1;
int sendDataRank = (rank + s) % nranks;
inpIx = sendDataRank * count + offset;
recvDim = -1;
sendDim = 0;
outIx = 0;
recvOffset = -1;
sendOffset = ((a - sendSkipped)%postFreq) * nelem;
sendStepOffset = 0;
if ((((a - sendSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
postSend = 1;
ps->inpIx = sendDataRank * count + offset;
ps->recvDim = -1;
ps->sendDim = 0;
ps->outIx = 0;
ps->recvOffset = -1;
ps->sendOffset = (a%postFreq) * nelem;
if (((a%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
ps->postSend = 1;
} else {
postSend = 0;
ps->postSend = 0;
}
postRecv = 0;
if (skip) sendSkipped++;
if (++a == lastA) {
phase = as == 1 ? (aggFactor > 1 ? 2 : 4) : 1; // If as == 1, switch to phase 2
resetA();
}
if (skip == 0) return;
ps->postRecv = 0;
} else if (phase == 1) {
int s = mirrorInvert(a, lastA)*aggDelta + as;
if (s >= nranks) skip = 1;
recvDim = firstBitSet(s, nrPow2);
sendOffset = ((a - sendSkipped)%postFreq)*nelem;
recvOffset = ((a - recvSkipped)%postFreq)*nelem;
postSend = 0;
if (recvDim == 0) {
if ((((a - sendSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) postSend = 1;
sendStepOffset = 0;
ps->recvDim = firstBitSet(s, nrPow2);
ps->sendOffset = (a%postFreq)*nelem;
ps->recvOffset = (a%postFreq)*nelem;
ps->postSend = 0;
if (ps->recvDim == 0 && (((a%postFreq) + 1 >= postFreq) || (a == lastA-1))) ps->postSend = 1;
if (((a%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
ps->postRecv = 1;
} else {
sendStepOffset = (a - sendSkipped)/postFreq;
ps->postRecv = 0;
}
if ((((a - recvSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
postRecv = 1;
} else {
postRecv = 0;
}
s -= (1<<recvDim);
s -= (1<<ps->recvDim);
int recvDataRank = (rank + nranks + s) % nranks;
inpIx = recvDataRank * count + offset;
sendDim = s ? firstBitSet(s, nrPow2) : -1;
if (sendDim == -1) {
sendOffset = -1;
sendStepOffset = 0;
} else if (as - (1<<recvDim) == 0) {
if (newPeer(a, aggFactor)) sendSkipped = a;
ps->inpIx = recvDataRank * count + offset;
ps->sendDim = s ? firstBitSet(s, nrPow2) : -1;
if (ps->sendDim == -1) {
ps->sendOffset = -1;
} else if (as - (1<<ps->recvDim) == 0) {
if (newPeer(a, aggFactor)) { sendSkipped = a; ps->stepOffset = stepOffset = 0; }
int foffset = a - sendSkipped;
sendStepOffset = recvDim == 0 ? 0 : foffset/postFreq;
sendOffset = (foffset%postFreq)*nelem;
ps->sendOffset = (foffset%postFreq)*nelem;
}
int recvDim = ps->recvDim;
if (s < nranks && skip) {
recvDim = -1;
recvOffset = -1;
postRecv = 0;
ps->recvDim = -1;
ps->recvOffset = -1;
ps->postRecv = 0;
skip = 0;
}
if (skip || recvDim == -1) recvSkipped++;
if (skip) sendSkipped++;
if (++a == lastA) {
as--;
phase = as % 2 == 1 ? 0 : 1;
resetA();
}
if (skip == 0) return;
if (recvDim > 0 && (((a-sendSkipped)%postFreq) + 1 >= postFreq) && skip == 0) stepOffset++;
} else if (phase == 2) {
int s = (2*mirrorInvert(a, lastA)+1)*scale*aggDelta + 1;
postRecv = 0;
ps->postRecv = 0;
if (s >= nranks) skip = 1;
recvDim = 0;
postSend = a == lastA-1 ? 1 : 0;
ps->recvDim = 0;
ps->postSend = a == lastA-1 ? 1 : 0;
s -= 1;
if (s < nranks && skip) {
recvDim = -1;
recvOffset = -1;
ps->recvDim = -1;
ps->recvOffset = -1;
skip = 0;
} else if (!skip) {
int foffset = phase2recv;
phase2recv++;
postRecv |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
recvOffset = (foffset%postFreq) * nelem;
int foffset = a + aggFactor - aggFactor/scale;
ps->postRecv |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
ps->recvOffset = (foffset%postFreq) * nelem;
}
int recvDataRank = (rank + nranks + s) % nranks;
inpIx = recvDataRank * count + offset;
sendDim = s ? firstBitSet(s, nrPow2) : -1;
int foffset = a - sendSkipped;
postSend |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
sendStepOffset = 0;
sendOffset = (foffset%postFreq) * nelem;
if (skip || sendDim == -1) sendSkipped++;
if (++a == lastA) {
phase = 3;
resetA();
}
if (skip == 0) return;
ps->inpIx = recvDataRank * count + offset;
ps->sendDim = s ? firstBitSet(s, nrPow2) : -1;
int foffset = a;
ps->postSend |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
ps->sendOffset = (foffset%postFreq) * nelem;
} else if (phase == 3) {
int s = (2*mirrorInvert(a, lastA)+1)*scale*aggDelta;
postRecv = a == lastA-1 ? 1 : 0;
ps->postRecv = a == lastA-1 ? 1 : 0;
if (s >= nranks) skip = 1;
recvDim = firstBitSet(s, nrPow2);
postSend = 0;
s -= (1<<recvDim);
int foffset = a - recvSkipped;
postRecv |= (foffset+1)%postFreq == 0 ? 1 : 0;
recvOffset = (foffset%postFreq) * nelem;
ps->recvDim = firstBitSet(s, nrPow2);
ps->postSend = 0;
s -= (1<<ps->recvDim);
int foffset = a;
ps->postRecv |= (foffset+1)%postFreq == 0 ? 1 : 0;
ps->recvOffset = (foffset%postFreq) * nelem;
int recvDataRank = (rank + nranks + s) % nranks;
inpIx = recvDataRank * count + offset;
sendDim = s ? firstBitSet(s, nrPow2) : -1;
ps->inpIx = recvDataRank * count + offset;
ps->sendDim = s ? firstBitSet(s, nrPow2) : -1;
if (s < nranks && skip) {
recvDim = -1;
recvOffset = -1;
postRecv = 0;
ps->recvDim = -1;
ps->recvOffset = -1;
ps->postRecv = 0;
skip = 0;
}
if (newPeer(a, aggFactor/(2*scale))) sendSkipped = a;
if (newPeer(a, aggFactor/(2*scale))) { sendSkipped = a; ps->stepOffset = stepOffset = 0; }
foffset = a - sendSkipped;
sendStepOffset = foffset / postFreq; // Accumulate on next steps
sendOffset = sendDim >= 0 ? (foffset%postFreq) * nelem : -1;
if (skip || recvDim == -1) recvSkipped++;
if (skip) sendSkipped++;
if (++a == lastA) {
scale *= 2;
phase = scale < aggFactor ? 2 : 4;
if ((foffset%postFreq) + 1 >= postFreq && skip == 0) stepOffset++;
ps->sendOffset = ps->sendDim >= 0 ? (foffset%postFreq) * nelem : -1;
} else if (phase == 4) {
ps->recvDim = 0;
ps->sendDim = -1;
ps->inpIx = rank * count + offset;
ps->recvOffset = ((aggFactor-1)%postFreq) * nelem;
ps->sendOffset = -1;
ps->postRecv = 1;
ps->postSend = 0;
offset += chunkCount;
}
a++;
if (a >= lastA && a >= parallelFactor) {
int p = phase;
if (p == 1) as--;
if (p == 3) scale *= 2;
phase =
p == 0 ? as == 1 ? (aggFactor > 1 ? 2 : 4) : 1 :
p == 1 ? as % 2 == 1 ? 0 : 1 :
p == 2 ? 3 :
p == 3 ? scale < aggFactor ? 2 : 4 :
5;
if (p == 4) {
if (offset >= end) {
ps->last = 2;
} else {
reset();
}
} else {
resetA();
}
if (skip == 0) return;
} else if (phase == 4) {
recvDim = 0;
sendDim = -1;
inpIx = rank * count + offset;
recvOffset = (phase2recv%postFreq) * nelem;
sendStepOffset = 0;
sendOffset = -1;
postRecv = 1;
postSend = 0;
offset += chunkCount;
if (offset >= end) {
last = 1;
} else {
reset();
}
return;
} else if (phase == 4 && offset >= end) {
ps->last = 1;
}
goto restart;
int flags = PatUsed | (skip ? PatSkipped : 0);
#if __CUDA_ARCH__ >= 600
cuda::atomic_ref<int, cuda::thread_scope_block> a(ps->flags);
a.store(flags, cuda::memory_order_release);
#else
ps->flags = flags;
#endif
}
};
@@ -669,14 +698,12 @@ class PatAGAlgorithm{
int nrPow2;
int postFreq;
int lastA;
int parallelFactor;
int aggFactor;
int as; // aggregated steps
int a; // step inside aggregated step
int aggDelta;
int scale;
int phase;
// AS computation
@@ -685,7 +712,7 @@ class PatAGAlgorithm{
int bitCount[32];
int bitZeroStep[32];
__device__ __host__ int min(int a, int b) {
__device__ __host__ ssize_t min(ssize_t a, ssize_t b) {
return (a<b)?a:b;
}
@@ -752,8 +779,9 @@ class PatAGAlgorithm{
public:
__device__ __host__ PatAGAlgorithm(int stepSize, int stepDepth, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
__device__ __host__ PatAGAlgorithm(int stepSize, int stepDepth, int maxParallelFactor, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
offset(offset), end(end), count(count), chunkCount(chunkCount), rank(rank), nranks(nranks) {
parallelFactor = maxParallelFactor;
aggDelta = nrPow2 = (1<<log2Up(nranks));
aggFactor = 1;
@@ -763,120 +791,120 @@ public:
aggDelta /= 2;
}
postFreq = aggFactor;
if (postFreq < parallelFactor) parallelFactor = postFreq;
int d = stepDepth;
while (d > 1 && aggFactor < nranks/2) {
d /= 2;
aggFactor *= 2;
aggDelta /= 2;
}
//printf("AggFactor %d PostFreq %d AggDelta %d\n", aggFactor, postFreq, aggDelta);
asDim = log2Up(aggDelta);
reset();
}
__device__ __host__ void getNextOp(int &recvDim, int &sendDim, size_t &inpIx, size_t &outIx, int &recvOffset, int &sendOffset, int &recvStepOffset, int &nelemOut, int &postRecv, int &postSend, int &last) {
restart:
//printf("Phase %d as %d/%d a %d/%d scale %d\n", phase, as, aggDelta, a, lastA, scale);
last = 0;
nelemOut = nelem;
inpIx = offset;
__device__ __host__ int getParallelFactor() {
return parallelFactor;
}
__device__ __host__ void getNextOp(struct ncclPatStep* ps) {
ps->last = 0;
ps->nelem = nelem;
ps->inpIx = offset;
int skip = 0;
if (phase == 0) {
if (a >= lastA) {
skip = 1;
} else if (phase == 0) {
int s = a*aggDelta + as;
if (s >= nranks) skip = 1;
int nextSkip = (a+1)*aggDelta + as >= nranks ? 1 : 0;
int recvDataRank = (rank + s) % nranks;
outIx = recvDataRank * count + offset;
sendDim = -1;
recvDim = 0;
inpIx = 0;
sendOffset = -1;
recvOffset = (a % postFreq) * nelem;
recvStepOffset = 0;
postRecv = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
postSend = 0;
a++;
if (nextSkip) {
as = nextAs();
if (as == aggDelta/2) {
offset += chunkCount;
if (offset >= end) {
last = 1;
} else {
reset();
}
return;
}
phase = 1;
resetA();
}
if (skip == 0) return;
ps->outIx = recvDataRank * count + offset;
ps->sendDim = -1;
ps->recvDim = 0;
ps->inpIx = 0;
ps->sendOffset = -1;
ps->recvOffset = (a % postFreq) * nelem;
ps->stepOffset = 0;
ps->postRecv = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
ps->postSend = 0;
} else if (phase == 1) {
int s = a*aggDelta + as;
if (s >= nranks) skip = 1;
sendDim = firstBitSet(s, nrPow2);
s -= (1<<sendDim);
ps->sendDim = firstBitSet(s, nrPow2);
s -= (1<<ps->sendDim);
int sendDataRank = (rank + nranks + s) % nranks;
outIx = sendDataRank * count + offset;
recvDim = s ? firstBitSet(s, nrPow2) : -1;
sendOffset = recvOffset = (a % postFreq) * nelem;
postSend = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
postRecv = (sendDim == 0) && ((a % postFreq == postFreq-1) || ((a+1)*aggDelta+as-1 >= nranks)) ? 1 : 0;
recvStepOffset = (sendDim == 0) ? 0 : a/postFreq;
if (recvDim == -1) {
recvOffset = -1;
postRecv = 0;
} else if (as - (1<<sendDim) == 0) {
int foffset = (a*aggDelta) >> (recvDim+1);
recvOffset = (foffset%postFreq)*nelem;
postRecv = (sendDim == 0) && ((foffset % postFreq == postFreq-1) || ((((foffset+1)*2)+1)<<recvDim) >= nranks) ? 1 : 0;
recvStepOffset = (sendDim == 0) ? 0 : foffset/postFreq;
ps->outIx = sendDataRank * count + offset;
ps->recvDim = s ? firstBitSet(s, nrPow2) : -1;
ps->sendOffset = ps->recvOffset = (a % postFreq) * nelem;
ps->postSend = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
ps->postRecv = (ps->sendDim == 0) && ((a % postFreq == postFreq-1) || ((a+1)*aggDelta+as-1 >= nranks)) ? 1 : 0;
ps->stepOffset = (ps->sendDim == 0) ? 0 : a/postFreq;
if (ps->recvDim == -1) {
ps->recvOffset = -1;
ps->postRecv = 0;
} else if (as - (1<<ps->sendDim) == 0) {
int foffset = (a*aggDelta) >> (ps->recvDim+1);
ps->recvOffset = (foffset%postFreq)*nelem;
ps->postRecv = (ps->sendDim == 0) && ((foffset % postFreq == postFreq-1) || ((((foffset+1)*2)+1)<<ps->recvDim) >= nranks) ? 1 : 0;
ps->stepOffset = (ps->sendDim == 0) ? 0 : foffset/postFreq;
}
if (s < nranks && sendDim == 0 && skip) {
if (s < nranks && ps->sendDim == 0 && skip) {
// Don't forget to receive at least once even if we don't send afterwards
sendDim = -1;
sendOffset = -1;
postSend = 0;
ps->sendDim = -1;
ps->sendOffset = -1;
ps->postSend = 0;
skip = 0;
}
if (++a == lastA) {
if (as % 2 == 1) {
phase = 0;
} else {
as = nextAs();
}
resetA();
}
if (skip == 0) return;
} else if (phase == 2) {
int s = (2*a+1)*scale*aggDelta;
postSend = (a % postFreq == postFreq-1) || ((2*(a+1)+1)*scale*aggDelta >= nranks) ? 1 : 0;
postRecv = 0;
ps->postSend = (a % postFreq == postFreq-1) || ((2*(a+1)+1)*scale*aggDelta >= nranks) ? 1 : 0;
ps->postRecv = 0;
if (s >= nranks) skip = 1;
sendDim = firstBitSet(s, nrPow2);
s -= (1<<sendDim);
sendOffset = (a%postFreq) * nelem;
recvStepOffset = a / postFreq;
ps->sendDim = firstBitSet(s, nrPow2);
s -= (1<<ps->sendDim);
ps->sendOffset = (a%postFreq) * nelem;
ps->stepOffset = a / postFreq;
int sendDataRank = (rank + nranks + s) % nranks;
outIx = sendDataRank * count + offset;
recvDim = s ? firstBitSet(s, nrPow2) : -1;
if (recvDim == -1) {
recvOffset = -1;
ps->outIx = sendDataRank * count + offset;
ps->recvDim = s ? firstBitSet(s, nrPow2) : -1;
if (ps->recvDim == -1) {
ps->recvOffset = -1;
} else {
s -= (1<<recvDim);
int foffset = (a*2*scale*aggDelta) >> (recvDim+1);
recvOffset = (foffset%postFreq)*nelem;
recvStepOffset = foffset / postFreq;
s -= (1<<ps->recvDim);
int foffset = (a*2*scale*aggDelta) >> (ps->recvDim+1);
ps->recvOffset = (foffset%postFreq)*nelem;
ps->stepOffset = foffset / postFreq;
}
if (++a == lastA) {
scale /= 2;
phase = scale ? 2 : 1;
}
a++;
if (a >= lastA && a >= parallelFactor) {
int p = phase;
if (p == 2) scale /= 2;
phase =
p == 2 ? scale ? 2 : 1 :
p == 1 ? as % 2 == 1 ? 0 : 1 :
1;
if (p == 0 || (p == 1 && as % 2 == 0)) as = nextAs();
if (p == 0 && as == aggDelta/2) {
offset += chunkCount;
if (offset >= end) {
ps->last = 2;
} else {
reset();
}
} else {
resetA();
}
if (skip == 0) return;
} else if (phase == 0 && as == 1 && offset + chunkCount >= end && a-1 >= ((lastA-1) / parallelFactor) * parallelFactor) {
ps->last = 1;
}
goto restart;
int flags = PatUsed | (skip ? PatSkipped : 0);
#if __CUDA_ARCH__ >= 600
cuda::atomic_ref<int, cuda::thread_scope_block> a(ps->flags);
a.store(flags, cuda::memory_order_release);
#else
ps->flags = flags;
#endif
}
};
#endif