Improved allreduce segmentation for small sizes

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(

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) \

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;