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

add 256 channels support

Browse Source
This commit is contained in:
Nusrat Islam
2024-05-15 16:58:28 -05:00
parent 6475da2ed9
commit 506f16c506
11 changed files with 95 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/device/common.cu
+2 -2
View File
@@ -17,11 +17,11 @@ struct RunWorkNop {
__device__ void run(ncclWork *w) {}
};
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
ncclKernelMain<-1, RunWorkNop, false>(comm, channelMask, workHead);
}
#ifdef ENABLE_COLLTRACE
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
ncclKernelMain<-1, RunWorkNop, true>(comm, channelMask, workHead);
}
#endif
src/device/common.h
+22 -12
View File
@@ -13,6 +13,7 @@
#include "op128.h"
#include "device_table.h"
#include "network/unpack/unpack_defs.h"
#include "comm.h"
#if defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
#define COLL_UNROLL 2
@@ -226,22 +227,31 @@ static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we
}
template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE>
__forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
__forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
const int tid = threadIdx.x;
int x = tid;
int total = 0, y;
int num = MAXCHANNELS/64 > 0 ? MAXCHANNELS/64 : 1;
switch (tid/WARP_SIZE) {
case 0:
if (channelMask & (1ull<<x)) {
int y = __popcll(channelMask & ((1ull<<x)-1));
ncclShmem.channelId = blockIdx.x;
/*for (int i = 0; i < num; i++) {
if (channelMask.masks[i] & (1ull<<x)) {
y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
y = total + y;
if (blockIdx.x == y) ncclShmem.channelId = x;
}
if (WARP_SIZE < MAXCHANNELS) {
if (WARP_SIZE < 64) {
x = WARP_SIZE + tid;
if (channelMask & (1ull<<x)) {
int y = __popcll(channelMask & ((1ull<<x)-1));
if (channelMask.masks[i] & (1ull<<x)) {
y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
y = y + total;
if (blockIdx.x == y) ncclShmem.channelId = x;
}
}
total = __popcll(channelMask.masks[i]);
}*/
break;
case 1:
if (tid < WARP_SIZE + NCCL_MAX_GROUPS)
@@ -361,23 +371,23 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, uint64_
#endif
}
__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
#ifdef ENABLE_COLLTRACE
__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
#endif
#ifdef ENABLE_COLLTRACE
#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
__global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
__global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
} \
\
__global__ void ncclDevKernelDebug_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
__global__ void ncclDevKernelDebug_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, true>(comm, channelMask, workHead); \
}
#else
#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
__global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
__global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
}
#endif
@@ -394,4 +404,4 @@ __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t ch
}
#endif
#endif
#endif