From 90b7a02944a962ccac2df6ab58741f75c85cff86 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 26 Jun 2019 15:09:13 -0700
Subject: [PATCH] Match primitives unroll counts with latest RCCL (#91)

[ROCm/rccl commit: e6a0da444ff98062259888cdb9d73c30df84ef1c]
---
 .../rccl/tools/rccl-prim-test/rccl_prim_test.cpp | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp b/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
index 17b8df3245..5462c1031c 100644
--- a/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
+++ b/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
@@ -35,7 +35,11 @@ THE SOFTWARE.
 
 #define MAX_WORKGROUPS 8
 #define THREADS 256
-#define UNROLL 8
+
+#define COPY_UNROLL       4
+#define REDUCE_UNROLL     2
+#define DOUBLECOPY_UNROLL 2
+#define REDUCECOPY_UNROLL 2
 
 struct transfer_data_t {
   float *dest0; //remote fine grain
@@ -75,11 +79,11 @@ __global__ void flag_sync_kernel(struct transfer_data_t* transfer_data, struct p
 
   int offset = transfer_data->N * blockIdx.x / gridDim.x;
   int n = transfer_data->N / gridDim.x;
-  if (op == OP_COPY) Copy<UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->src0 + offset, n);
-  if (op == OP_LOCALCOPY) Copy<UNROLL, THREADS, float>(transfer_data->dest1 + offset, transfer_data->src0 + offset, n);
-  if (op == OP_DOUBLECOPY) DoubleCopy<UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->dest1 + offset, transfer_data->src0 + offset, n);
-  if (op == OP_REDUCE) Reduce<UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->src0 + offset, transfer_data->src1 + offset, n);
-  if (op == OP_REDUCECOPY) ReduceCopy<UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->dest1 + offset, transfer_data->src0 + offset, transfer_data->src1 + offset, n);
+  if (op == OP_COPY) Copy<COPY_UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->src0 + offset, n);
+  if (op == OP_LOCALCOPY) Copy<COPY_UNROLL, THREADS, float>(transfer_data->dest1 + offset, transfer_data->src0 + offset, n);
+  if (op == OP_DOUBLECOPY) DoubleCopy<DOUBLECOPY_UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->dest1 + offset, transfer_data->src0 + offset, n);
+  if (op == OP_REDUCE) Reduce<REDUCE_UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->src0 + offset, transfer_data->src1 + offset, n);
+  if (op == OP_REDUCECOPY) ReduceCopy<REDUCECOPY_UNROLL, THREADS, float>(transfer_data->dest0 + offset, transfer_data->dest1 + offset, transfer_data->src0 + offset, transfer_data->src1 + offset, n);
 
   __syncthreads();
   if (idx == 0) {