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

[Device] WarpSpeed enablement and single node CU and perf opt for MI350 (#2073)

Browse Source 
[ROCm/rccl commit: d009ab144e]
This commit is contained in:
Mustafa Abduljabbar
2025-12-11 19:04:35 -05:00
committed by GitHub
parent f17357d0d4
commit 2621e0254e
22 changed files with 424 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/device/all_gather.h
+11 -2
View File
@@ -20,11 +20,20 @@ namespace {
const int bid = ncclShmem.channelId - work->channelLo;
int npKitCtxIdx = bid; // unused variable - compiler warning
#endif
#ifdef ENABLE_WARP_SPEED
int warp = threadIdx.x / WARP_SIZE;
ncclRing *ring = &ncclShmem.warpChannel[warp].ring;
#else
ncclRing *ring = &ncclShmem.channel.ring;
#endif
const int *ringRanks = ring->userRanks;
const int nranks = ncclShmem.comm.nRanks;
ssize_t count, partOffset, partCount, chunkCount;
#ifdef ENABLE_WARP_SPEED
ncclCollCbdPart(work, ncclShmem.warpChannelId[warp], Proto::Id, sizeof(T), &count, &partOffset, &partCount, &chunkCount);
#else
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &count, &partOffset, &partCount, &chunkCount);
#endif
ssize_t offset;
ssize_t dataOffset;
int nelem;
@@ -142,7 +151,7 @@ namespace {
#endif
// Final wait/copy.
prims.directRecv(offset, nelem);
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT)
if (tid == 0) {
NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
@@ -671,4 +680,4 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_P
return;
}
}
};
};
projects/rccl/src/device/all_reduce.h
+10
View File
@@ -20,8 +20,14 @@ namespace {
#else
__device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
#endif
#ifdef ENABLE_WARP_SPEED
int warp = threadIdx.x / WARP_SIZE;
ncclRing *ring = &ncclShmem.warpChannel[warp].ring;
#else
ncclRing *ring = &ncclShmem.channel.ring;
#endif
int ringIx = ring->index;
const int nranks = ncclShmem.comm.nRanks;
#if defined(ENABLE_NPKIT)
const int bid = ncclShmem.channelId - work->channelLo;
@@ -31,7 +37,11 @@ namespace {
ssize_t gridOffset;
ssize_t channelCount;
ssize_t chunkCount;
#ifdef ENABLE_WARP_SPEED
ncclCollCbdPart(work, ncclShmem.warpChannelId[warp], Proto::Id, sizeof(T), &size, &gridOffset, &channelCount, &chunkCount);
#else
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &size, &gridOffset, &channelCount, &chunkCount);
#endif
const ssize_t loopCount = nranks * chunkCount;
ssize_t offset;
int nelem;
projects/rccl/src/device/broadcast.h
+9
View File
@@ -19,7 +19,12 @@ namespace {
const int bid = ncclShmem.channelId - work->channelLo;
int npKitCtxIdx = bid; // unused variable - compiler warning
#endif
#ifdef ENABLE_WARP_SPEED
int warp = threadIdx.x / WARP_SIZE;
ncclRing *ring = &ncclShmem.warpChannel[warp].ring;
#else
ncclRing *ring = &ncclShmem.channel.ring;
#endif
const int rank = ring->userRanks[0];
const int nextRank = ring->userRanks[1];
const int root = work->root;
@@ -27,7 +32,11 @@ namespace {
ssize_t chunkCount;
ssize_t channelCount;
ssize_t gridOffset;
#ifdef ENABLE_WARP_SPEED
ncclCollCbdPart(work, ncclShmem.warpChannelId[warp], Proto::Id, sizeof(T), &size, &gridOffset, &channelCount, &chunkCount);
#else
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &size, &gridOffset, &channelCount, &chunkCount);
#endif
size_t offset;
int nelem;
int workNthreads;
projects/rccl/src/device/common.cu
+12 -12
View File
@@ -17,24 +17,24 @@ struct RunWorkNop {
__device__ void run() {}
};
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/1>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_1(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/1>(&argsStorage.args);
}
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/2>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_2(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/2>(&argsStorage.args);
}
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/4>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_4(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/4>(&argsStorage.args);
}
#ifdef ENABLE_COLLTRACE
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/1>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/1>(&argsStorage.args);
}
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/2>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/2>(&argsStorage.args);
}
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/4>(&args4K.args);
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {
ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/4>(&argsStorage.args);
}
#endif
projects/rccl/src/device/common.h
+68 -9
View File
@@ -138,7 +138,11 @@ struct ncclShmemData {
int aborted;
alignas(16) struct ncclDevComm comm;
alignas(16) struct ncclDevChannel channel;
#ifdef ENABLE_WARP_SPEED
int warpComm;
alignas(16) struct ncclDevChannel warpChannel[NCCL_MAX_GROUPS];
int warpChannelId[NCCL_MAX_GROUPS];
#endif
int batchIx, nextBatchIx;
enum ncclDevWorkType workType;
uint8_t directMode;
@@ -442,10 +446,17 @@ struct RunWorkBatch {
if (work->nWarps != workPrev->nWarps) __syncthreads();
}
int subtn = work->nWarps*WARP_SIZE;
#ifdef ENABLE_WARP_SPEED
if (tid < subtn) {
if(ncclShmem.warpComm == 0 || Algo != NCCL_ALGO_RING) RunWorkColl<Fn, T, RedOp, Algo, Proto>().run(tid, subtn, work);
else if (ncclShmem.warpChannelId[tid / WARP_SIZE] >= 0) RunWorkColl<Fn, T, RedOp, Algo, Proto>().run(tid % WARP_SIZE, WARP_SIZE, work);
}
#else
// Coverity reports a possible thread divergence due to not all threads participating in the collective.
// However, the code ensures that the participation is on a per-warp basis.
// coverity[device_thread_diverged:FALSE]
if (tid < subtn) RunWorkColl<Fn, T, RedOp, Algo, Proto>().run(tid, subtn, work);
#endif
}
}
};
@@ -489,7 +500,12 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
int x = tid;
int total = 0, y;
int num = MAXCHANNELS/64 > 0 ? MAXCHANNELS/64 : 1;
#ifdef ENABLE_WARP_SPEED
int warpCount = tn / WARP_SIZE;
int localWarpId = tid / WARP_SIZE;
int globalWarpId = (warpCount * blockIdx.x) + localWarpId;
int laneId = tid % WARP_SIZE;
#endif
// Copy kernel args to shmem and then only read those. Otherwise the compiler
// will end up putting the args into thread local stack which is very wasteful.
if (tid < sizeof(ncclDevKernelArgs)/sizeof(uint32_t)) {
@@ -583,9 +599,52 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
ncclShmem.collTrace = args->comm->collTrace + COLLTRACE_NUM_ITEMS*ncclShmem.channelId;
ncclShmem.collTraceTail = args->comm->collTraceTail + ncclShmem.channelId;
}
#endif
#ifdef ENABLE_WARP_SPEED
if(tid == 0) {
ncclShmem.warpComm = args->comm->warpLevelComm;
}
#endif
__syncthreads(); // publish shmem
#ifdef ENABLE_WARP_SPEED
// Determine per-warp channel assignment for WarpSpeed enablement
total = 0;
if(ncclShmem.warpComm == 1) { // If warpComm is enabled, assign warps to channels that have the corresponding channel mask enabled
ncclShmem.warpChannelId[localWarpId] = -1;
__syncthreads();
for (int i = 0; i < num; i++) {
if (args->channelMask.masks[i] & (1ull<<laneId)) {
y = __popcll(args->channelMask.masks[i] & ((1ull<<laneId)-1));
y = total + y;
if (globalWarpId == y) {
ncclShmem.warpChannelId[localWarpId] = laneId + total;
break;
}
}
total = total + __popcll(args->channelMask.masks[i]);
}
__syncthreads();
if(ncclShmem.warpChannelId[localWarpId] >= 0) {
void* dst = &ncclShmem.warpChannel[localWarpId];
void* src = &((ncclDevCommAndChannels*)ncclShmem.args.comm)->channels[ncclShmem.warpChannelId[localWarpId]];
int bytes = sizeof(ncclDevChannel);
static_assert(sizeof(ncclDevChannel) <= 16*WARP_SIZE, "ncclDevChannel cannot be loaded by a single warp in one insn.");
// assert((tid-localWarpId*WARP_SIZE) >= 0 && (tid-localWarpId*WARP_SIZE) < WARP_SIZE);
copyToShmem16(tid-localWarpId*WARP_SIZE, dst, src, bytes);
}
} else { // If warpComm is disabled, all warps use the same channel as the block
if(laneId == 0) {
ncclShmem.warpChannelId[localWarpId] = ncclShmem.channelId;
}
// Use all threads in the warp to copy the channel data in parallel
void* dst = &ncclShmem.warpChannel[localWarpId];
void* src = &ncclShmem.channel;
int bytes = sizeof(ncclDevChannel);
copyToShmem16(laneId, dst, src, bytes);
}
__syncthreads();
#endif
#ifdef ENABLE_PROFILING
if (tid == 0) {
ncclShmem.prof.count = 0;
@@ -648,17 +707,17 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
#endif
}
__global__ void ncclDevKernel_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernel_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernel_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernel_Generic_1(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
__global__ void ncclDevKernel_Generic_2(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
__global__ void ncclDevKernel_Generic_4(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
#ifdef ENABLE_COLLTRACE
__global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
__global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
__global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
__global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage);
#endif
#define DEFINE_ncclDevKernel_nop(suffix, coll, redop, ty, algo, proto, specializedFnId) \
__global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {}
__global__ void ncclDevKernel_##suffix(ncclDevKernelArgsDefaultStorage NCCL_GRID_CONSTANT const argsStorage) {}
#ifdef USE_INDIRECT_FUNCTION_CALL
#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, acc, pipeline, unroll) \
projects/rccl/src/device/msccl_kernel_impl.h
+3
View File
@@ -146,6 +146,9 @@ __device__ __forceinline__ void mscclRunInterpreter(
}
if (bytes) copyToShmem8(tid%WARP_SIZE, dst, src, bytes);
}
#ifdef ENABLE_WARP_SPEED
ncclShmem.warpComm = 0;
#endif
__syncthreads(); // publish shmem
#if defined(ENABLE_NPKIT)
projects/rccl/src/device/prims_ll.h
+4
View File
@@ -654,7 +654,11 @@ public:
redOp(redOpArg),
tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), group(group), threadsPerBlock(blockDim.x),
stepLines(ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/sizeof(ncclLLFifoLine)) {
#ifdef ENABLE_WARP_SPEED
auto *channel = isMsccl(Metadata) ? &ncclShmem.channel : &ncclShmem.warpChannel[threadIdx.x / WARP_SIZE];
#else
auto *channel = &ncclShmem.channel;
#endif
barriers = &ncclShmem.groups[group].barrier;
// If we are going to support oneshot collNet + LL, then we would need to add connector index here
int nrecv=0, nsend=0;
projects/rccl/src/device/prims_ll128.h
+4
View File
@@ -579,7 +579,11 @@ public:
tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), /*compiler warnings*/
stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)),
warp(tid/WARP_SIZE), warpInBlock(threadIdx.x/WARP_SIZE), flagThread((tid%4)==3), group(group), threadsPerBlock(blockDim.x){
#ifdef ENABLE_WARP_SPEED
auto *channel = isMsccl(Metadata) ? &ncclShmem.channel : &ncclShmem.warpChannel[warpInBlock];
#else
auto *channel = &ncclShmem.channel;
#endif
barriers = &ncclShmem.groups[group].barrier;
int nrecv=0, nsend=0;
while (nrecv < MaxRecv && recvPeers[nrecv] >= 0) {
projects/rccl/src/device/prims_simple.h
+21 -6
View File
@@ -502,14 +502,22 @@ private:
public:
static inline __device__ void sendPeerNotify(int peer, int connIndex, int steps) {
#ifdef ENABLE_WARP_SPEED
ncclDevChannelPeer* peerPtr = ncclShmem.warpChannel[threadIdx.x/WARP_SIZE].peers[peer];
#else
ncclDevChannelPeer* peerPtr = ncclShmem.channel.peers[peer];
#endif
peerPtr->send[connIndex].step += steps;
st_relaxed_sys_global(peerPtr->send[connIndex].tail, peerPtr->send[connIndex].step);
}
static inline __device__ void recvPeerNotify(int peer, int connIndex, int steps) {
int spins = 0;
#ifdef ENABLE_WARP_SPEED
ncclDevChannelPeer* peerPtr = ncclShmem.warpChannel[threadIdx.x/WARP_SIZE].peers[peer];
#else
ncclDevChannelPeer* peerPtr = ncclShmem.channel.peers[peer];
#endif
peerPtr->recv[connIndex].step += steps;
st_relaxed_sys_global(peerPtr->recv[connIndex].head, peerPtr->recv[connIndex].step);
while (ld_volatile_global(peerPtr->recv[connIndex].tail) < peerPtr->recv[connIndex].step) {
@@ -770,13 +778,20 @@ public:
struct ncclDevWorkP2p* p2pWork = nullptr, int stepSize_ = 0, int mode = primsModeDefault
):
tid(tid), tidInBlock(threadIdx.x), nthreads(nthreads), /*compiler warnings*/
#ifdef ENABLE_WARP_SPEED
stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_), group(ncclShmem.warpComm? tidInBlock / WARP_SIZE : group), threadsPerBlock(blockDim.x){
#else
stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_), group(group), threadsPerBlock(blockDim.x){
#endif
barriers = &ncclShmem.groups[group].barrier;
// PAT uses the same barrier for each group
barriers_pat = &ncclShmem.barrier_pat;
this->nworkers = nthreads;
#ifdef ENABLE_WARP_SPEED
auto *channel = isMsccl(Metadata) ? &ncclShmem.channel : &ncclShmem.warpChannel[tidInBlock/WARP_SIZE];
#else
auto *channel = &ncclShmem.channel;
#endif
int peer = -1;
flags = 0;
index = -1;
@@ -831,9 +846,9 @@ public:
}
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
if (flags & (RoleWaitRecv|RolePostRecv)) loadRecvConn(ncclShmem.channel.peers[peer], connIndexRecv, collWork ? collWork->direct : 0, recvIpcReg, recvNetReg);
if (flags & (RoleWaitRecv|RolePostRecv)) loadRecvConn(channel->peers[peer], connIndexRecv, collWork ? collWork->direct : 0, recvIpcReg, recvNetReg);
// coverity[overrun-call] => Coverity think prims.index can be greater than 1
if (flags & (RoleWaitSend|RolePostSend)) loadSendConn(ncclShmem.channel.peers[peer], connIndexSend, collWork ? collWork->direct : 0, sendIpcReg, sendNetReg);
if (flags & (RoleWaitSend|RolePostSend)) loadSendConn(channel->peers[peer], connIndexSend, collWork ? collWork->direct : 0, sendIpcReg, sendNetReg);
// if (barrierAny(flags & NetDeviceUnpack)) {
// flags |= AnyNetDeviceUnpack;
@@ -861,7 +876,7 @@ public:
// Load recv peer
int recvPeer = mode == primsModePatRs ? (rank - delta + nranks) % nranks : (rank + delta) % nranks;
struct ncclPatPeer* peer = ((struct ncclPatPeer*)recvPeers)+tid;
struct ncclConnInfo* conn = peer->conn = ncclShmem.channel.peers[recvPeer]->recv+connIndexRecv;
struct ncclConnInfo* conn = peer->conn = channel->peers[recvPeer]->recv+connIndexRecv;
peer->step = conn->step;
peer->buff = conn->buffs[NCCL_PROTO_SIMPLE];
peer->stepCache = loadStepValue(peer->tailPtr = conn->tail);
@@ -871,7 +886,7 @@ public:
// Load send peer
int sendPeer = mode == primsModePatAg ? (rank - delta + nranks) % nranks : (rank + delta) % nranks;
peer = ((struct ncclPatPeer*)sendPeers)+tid;
conn = peer->conn = ncclShmem.channel.peers[sendPeer]->send+connIndexSend;
conn = peer->conn = channel->peers[sendPeer]->send+connIndexSend;
peer->step = conn->step;
peer->connFifo = conn->connFifo;
peer->buff = conn->buffs[NCCL_PROTO_SIMPLE];
projects/rccl/src/device/reduce.h
+9
View File
@@ -16,7 +16,12 @@ namespace {
#else
__device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
#endif
#ifdef ENABLE_WARP_SPEED
int warp = threadIdx.x / WARP_SIZE;
ncclRing *ring = &ncclShmem.warpChannel[warp].ring;
#else
ncclRing *ring = &ncclShmem.channel.ring;
#endif
const int nranks = ncclShmem.comm.nRanks;
const int rank = ncclShmem.comm.rank;
const int prevRank = ring->userRanks[nranks-1];
@@ -24,7 +29,11 @@ namespace {
size_t chunkCount;
size_t channelCount;
size_t gridOffset;
#ifdef ENABLE_WARP_SPEED
ncclCollCbdPart(work, ncclShmem.warpChannelId[warp], Proto::Id, sizeof(T), (size_t*)nullptr, &gridOffset, &channelCount, &chunkCount);
#else
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), (size_t*)nullptr, &gridOffset, &channelCount, &chunkCount);
#endif
size_t offset;
int nelem;
projects/rccl/src/device/reduce_scatter.h
+9
View File
@@ -16,14 +16,23 @@ namespace {
#else
__device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
#endif
#ifdef ENABLE_WARP_SPEED
int warp = threadIdx.x / WARP_SIZE;
ncclRing *ring = &ncclShmem.warpChannel[warp].ring;
#else
ncclRing *ring = &ncclShmem.channel.ring;
#endif
int const *ringRanks = ring->userRanks;
const int nranks = ncclShmem.comm.nRanks;
size_t count;
size_t gridOffset;
size_t channelCount;
size_t chunkCount;
#ifdef ENABLE_WARP_SPEED
ncclCollCbdPart(work, ncclShmem.warpChannelId[warp], Proto::Id, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
#else
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
#endif
size_t offset;
size_t dataOffset;
uint32_t nelem;