From 8bb06c94bec4b20ba91d85dcd344ed4d53c68770 Mon Sep 17 00:00:00 2001
From: Sylvain Jeaugey <sjeaugey@nvidia.com>
Date: Fri, 7 Oct 2016 12:40:05 -0700
Subject: [PATCH] Improved allreduce segmentation for small sizes

---
 src/all_reduce.cu     | 28 +++++++++++++++-------------
 src/common_kernel.h   |  2 ++
 src/reduce_scatter.cu |  2 +-
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/all_reduce.cu b/src/all_reduce.cu
index 86a297f830..a81ee62d60 100644
--- a/src/all_reduce.cu
+++ b/src/all_reduce.cu
@@ -74,11 +74,13 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
     int offset;
     int maxOffset;
     int slice;
+    int chunkSize = min(sliceSize, DIVUP(size-chunkOffset,nranks));
+    ALIGN_SIZE(chunkSize, THREADS*UNROLL);
 
     // step 0: push data to next GPU
     slice = ring.userRank[nranks-1];
-    offset = chunkOffset + slice * sliceSize;
-    maxOffset = size-offset;
+    offset = chunkOffset + slice * chunkSize;
+    maxOffset = min(chunkSize, size-offset);
 
     Prims::Copy(
         thisInput  + offset,
@@ -93,8 +95,8 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
     // k-2 steps: reduce and copy to next GPU
     for (int j=2; j<nranks; ++j) {
       slice = ring.userRank[nranks-j];
-      offset = chunkOffset + slice * sliceSize;
-      maxOffset = size-offset;
+      offset = chunkOffset + slice * chunkSize;
+      maxOffset = min(chunkSize, size-offset);
 
       Prims::Reduce(
           prevInput  + poffset,
@@ -108,11 +110,11 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
       NEXT_STEP;
     }
 
-    // step k - 1: reduce this buffer and data, which will produce the final
+    // step k-1: reduce this buffer and data, which will produce the final
     // result that we store in this data and push to the next GPU
     slice = ring.userRank[0];
-    offset = chunkOffset + slice * sliceSize;
-    maxOffset = size-offset;
+    offset = chunkOffset + slice * chunkSize;
+    maxOffset = min(chunkSize, size-offset);
 
     Prims::ReduceCopy(
         prevInput  + poffset,
@@ -130,8 +132,8 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
       // k-2 steps: copy result to next GPU
       for (int j=1; j<nranks-1; ++j) {
         slice = ring.userRank[nranks - j];
-        offset = chunkOffset + slice * sliceSize;
-	maxOffset = size-offset;
+        offset = chunkOffset + slice * chunkSize;
+        maxOffset = min(chunkSize, size-offset);
 
         Prims::Copy(
             thisOutput + offset,
@@ -147,8 +149,8 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
       // k-2 steps: copy result to next GPU
       for (int j=1; j<nranks-1; ++j) {
         slice = ring.userRank[nranks - j];
-        offset = chunkOffset + slice * sliceSize;
-	maxOffset = size-offset;
+        offset = chunkOffset + slice * chunkSize;
+        maxOffset = min(chunkSize, size-offset);
 
         Prims::DoubleCopy(
             prevInput + poffset,
@@ -164,8 +166,8 @@ __global__ void AllReduceKernel(const KernelArgs<T> args) {
 
       // Make final copy from buffer to dest.
       slice = ring.userRank[1];
-      offset = chunkOffset + slice * sliceSize;
-      maxOffset = size-offset;
+      offset = chunkOffset + slice * chunkSize;
+      maxOffset = min(chunkSize, size-offset);
 
       // Here we need to copy from buffer to this output.
       Prims::Copy(
diff --git a/src/common_kernel.h b/src/common_kernel.h
index 95c9eb4e27..7a17e3682d 100644
--- a/src/common_kernel.h
+++ b/src/common_kernel.h
@@ -17,6 +17,8 @@
 #define WARP_SIZE 32
 #define ROUNDUP(x, y)                                                           \
     (((((x) + (y) - 1) / (y))) * (y))
+#define DIVUP(x, y) \
+    (((x)+(y)-1)/(y))
 #define BAR_EXEC(type, barid, nthreads) \
     asm("bar." #type " " #barid ", " #nthreads ";\n\t")
 #define BAR_EXPAND(type, barid, nthreads) \
diff --git a/src/reduce_scatter.cu b/src/reduce_scatter.cu
index 75f203b61b..f13cbfb06b 100644
--- a/src/reduce_scatter.cu
+++ b/src/reduce_scatter.cu
@@ -96,7 +96,7 @@ __global__ void ReduceScatterKernel(const KernelArgs<T> args) {
       NEXT_STEP;
     }
 
-    // step k - 1: reduce this buffer and data, which will produce the final
+    // step k-1: reduce this buffer and data, which will produce the final
     // result that we store in this data and push to the next GPU
     rankDest = ring.userRank[0];
     offset = chunkOffset + rankDest * size;