From 4e68229c8b365d2e9108f06ae86ab39ebf1f311e Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Thu, 12 Nov 2020 15:47:34 -0800
Subject: [PATCH] Skip unused peer connection in scatter and gather (#301)

---
 src/collectives/device/gather.h  | 5 +++--
 src/collectives/device/scatter.h | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/collectives/device/gather.h b/src/collectives/device/gather.h
index c704e39421..b6c3865875 100644
--- a/src/collectives/device/gather.h
+++ b/src/collectives/device/gather.h
@@ -52,9 +52,10 @@ __device__ void ncclGatherKernel(struct CollectiveArgs* args) {
       else {
         int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
         int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
+        int peerNone = -1;
         if (rank == root) {
           ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
-            prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
+            prims(tid, nthreads, &peerRecv, &peerNone, NULL, stepSize, channel, comm);
 
           ssize_t recv_offset = chunkOffset + peerRecv*size;
           prims.recv(thisOutput+recv_offset, nelem);
@@ -62,7 +63,7 @@ __device__ void ncclGatherKernel(struct CollectiveArgs* args) {
         else {
           if (peerSend == root) {
             ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
-              prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
+              prims(tid, nthreads, &peerNone, &peerSend, NULL, stepSize, channel, comm);
 
             ssize_t send_offset = chunkOffset;
             prims.send(thisInput+send_offset, nelem);
diff --git a/src/collectives/device/scatter.h b/src/collectives/device/scatter.h
index f08544068a..848b47713d 100644
--- a/src/collectives/device/scatter.h
+++ b/src/collectives/device/scatter.h
@@ -52,9 +52,10 @@ __device__ void ncclScatterKernel(struct CollectiveArgs* args) {
       else {
         int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
         int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
+        int peerNone = -1;
         if (rank == root) {
           ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
-            prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
+            prims(tid, nthreads, &peerNone, &peerSend, NULL, stepSize, channel, comm);
 
           ssize_t send_offset = chunkOffset + peerSend*size;
           prims.send(thisInput+send_offset, nelem);
@@ -62,7 +63,7 @@ __device__ void ncclScatterKernel(struct CollectiveArgs* args) {
         else {
           if (peerRecv == root) {
             ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
-              prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
+              prims(tid, nthreads, &peerRecv, &peerNone, NULL, stepSize, channel, comm);
 
             ssize_t recv_offset = chunkOffset;
             prims.recv(thisOutput+recv_offset, nelem);