diff --git a/src/alltoall.cu b/src/alltoall.cu index 4d86bdaae5..41bbc780a8 100644 --- a/src/alltoall.cu +++ b/src/alltoall.cu @@ -8,6 +8,8 @@ #include #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; rreportErrors = 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; rexpected[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