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

Add option to use alltoall, gather and scatter API

Browse Source 
These APIs launche RCCL kernel implementation by default. If environmental
variable RCCL_ALLTOALL_KERNEL_DISABLE=1, then the APIs use wrapper around
ncclSend and ncclRecv.
This commit is contained in:
Wenkai Du
2020-04-13 15:51:57 -07:00
parent 0d7c4db33e
commit 2813c96826
3 changed files with 29 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/alltoall.cu
+8
View File
@@ -8,6 +8,8 @@
#include <hip/hip_runtime.h>
#include "common.h"
#define USE_RCCL_GATHER_SCATTER
void print_header() {
PRINT("# %10s %12s %6s %6s out-of-place in-place \n", "", "", "", "");
PRINT("# %10s %12s %6s %6s %7s %6s %6s %5s %7s %6s %6s %5s\n", "size", "count", "type", "redop",
@@ -65,12 +67,18 @@ testResult_t AlltoAllRunColl(void* sendbuff, void* recvbuff, size_t count, ncclD
size_t rankOffset = count * wordSize(type);
if (count == 0) return testSuccess;
#if NCCL_MAJOR >= 2 && NCCL_MINOR >= 7
#if defined(RCCL_GATHER_SCATTER) && defined(USE_RCCL_GATHER_SCATTER)
NCCLCHECK(ncclAllToAll(sendbuff, recvbuff, count, type, comm, stream));
#else
NCCLCHECK(ncclGroupStart());
for (int r=0; r<nRanks; r++) {
NCCLCHECK(ncclSend(((char*)sendbuff)+r*rankOffset, count, type, r, comm, stream));
NCCLCHECK(ncclRecv(((char*)recvbuff)+r*rankOffset, count, type, r, comm, stream));
}
NCCLCHECK(ncclGroupEnd());
#endif
#endif
return testSuccess;
}
src/gather.cu
+12
View File
@@ -9,6 +9,7 @@
#include "common.h"
//#define DEBUG_PRINT
#define USE_RCCL_GATHER_SCATTER
void print_header() {
PRINT("# %10s %12s %6s out-of-place in-place \n", "", "", "");
@@ -57,6 +58,8 @@ testResult_t GatherInitData(struct threadArgs* args, ncclDataType_t type, ncclRe
}
HIPCHECK(hipDeviceSynchronize());
}
// We don't support in-place gather
args->reportErrors = in_place ? 0 : 1;
return testSuccess;
}
@@ -76,6 +79,13 @@ testResult_t GatherRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDat
int rank;
NCCLCHECK(ncclCommUserRank(comm, &rank));
#if NCCL_MAJOR >= 2 && NCCL_MINOR >= 7
#if defined(RCCL_GATHER_SCATTER) && defined(USE_RCCL_GATHER_SCATTER)
if (rank == root)
NCCLCHECK(ncclGather(sendbuff, recvbuff, count, type, root, comm, stream));
else
NCCLCHECK(ncclGather(sendbuff, 0, count, type, root, comm, stream));
#else
NCCLCHECK(ncclGroupStart());
if (rank == root) {
for (int r=0; r<nRanks; r++)
@@ -83,6 +93,8 @@ testResult_t GatherRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDat
}
NCCLCHECK(ncclSend(sendbuff, count, type, root, comm, stream));
NCCLCHECK(ncclGroupEnd());
#endif
#endif
return testSuccess;
}
src/scatter.cu
+9
View File
@@ -9,6 +9,7 @@
#include "common.h"
//#define DEBUG_PRINT
#define USE_RCCL_GATHER_SCATTER
void print_header() {
PRINT("# %10s %12s %6s %6s out-of-place in-place \n", "", "", "", "");
@@ -59,6 +60,8 @@ testResult_t ScatterInitData(struct threadArgs* args, ncclDataType_t type, ncclR
TESTCHECK(InitData(args->expected[i], recvcount, type, rep, rank));
HIPCHECK(hipDeviceSynchronize());
}
// We don't support in-place scatter
args->reportErrors = in_place ? 0 : 1;
return testSuccess;
}
@@ -78,6 +81,10 @@ testResult_t ScatterRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDa
int rank;
NCCLCHECK(ncclCommUserRank(comm, &rank));
#if NCCL_MAJOR >= 2 && NCCL_MINOR >= 7
#if defined(RCCL_GATHER_SCATTER) && defined(USE_RCCL_GATHER_SCATTER)
NCCLCHECK(ncclScatter(sendbuff, recvbuff, count, type, root, comm, stream));
#else
NCCLCHECK(ncclGroupStart());
if (rank == root) {
for (int r=0; r<nRanks; r++)
@@ -85,6 +92,8 @@ testResult_t ScatterRunColl(void* sendbuff, void* recvbuff, size_t count, ncclDa
}
NCCLCHECK(ncclRecv(recvbuff, count, type, root, comm, stream));
NCCLCHECK(ncclGroupEnd());
#endif
#endif
return testSuccess;
}