From 2813c968264b2368264dbfc9f8e864b3d0ef7aa4 Mon Sep 17 00:00:00 2001
From: Wenkai Du <wenkai.du@amd.com>
Date: Mon, 13 Apr 2020 15:51:57 -0700
Subject: [PATCH] Add option to use alltoall, gather and scatter API

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.
---
 src/alltoall.cu |  8 ++++++++
 src/gather.cu   | 12 ++++++++++++
 src/scatter.cu  |  9 +++++++++
 3 files changed, 29 insertions(+)

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 <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;
 }
diff --git a/src/gather.cu b/src/gather.cu
index 5230fdc81d..65dc714893 100644
--- a/src/gather.cu
+++ b/src/gather.cu
@@ -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;
 }
 
diff --git a/src/scatter.cu b/src/scatter.cu
index 18ebb75b07..d741391300 100644
--- a/src/scatter.cu
+++ b/src/scatter.cu
@@ -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;
 }