Heavy code refactoring to remove a lot of code in collectives (~1000 lines).

Have all collectives use the same args, the same ring, and the same primitives for synchronization between threads with the same pattern.

src/enqueue.h

@@ -1,31 +1,90 @@
 /*************************************************************************
- * Copyright (c) 2015, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
  *
- * See LICENCE.txt for license information
+ * See LICENSE.txt for license information
  ************************************************************************/
 
 #ifndef enqueue_h_
 #define enqueue_h_
 
 #include "core.h"
+#include "reduce_kernel.h"
 
-int getRingIndex(const ncclComm_t comm, int device);
-void lockEventQueue(EventQueue* eq);
-void releaseEventQueue(EventQueue* eq);
-void CUDART_CB freeEvent(cudaStream_t stream, cudaError_t status, void* eq_void);
+/* Syncronize previous collective (if in different stream) and enqueue
+ * collective. Work is performed asynchronously with the host thread.
+ * The ColFunc class should be templated on the datatype and reduction
+ * operator (if applicable) and define a static entry() method as
+ * follows.
+ *   template <typename T, template <typename> class RedOp>
+ *   class CollectiveFunctor {
+ *     public:
+ *     static ncclResult_t entry(const void* sendbuff, void* recvbuff, int count,
+ *         int root, ncclComm* comm, cudaStream_t stream);
+ *   };
+ * The entry() method can assume that the appropriate cuda device has been set. */
+template< template<typename, template<typename> class> class ColFunc,
+          typename T,
+          template<typename> class Op >
+ncclResult_t enqueue(const void* sendbuff,
+                     void* recvbuff,
+                     int count,
+                     int root,
+                     ncclComm_t comm,
+                     cudaStream_t stream)
+{
+  if (stream != comm->prevStream) { // sync required for calls in different streams
+    comm->prevStream = stream;
+    CUDACHECK( cudaStreamWaitEvent(stream, comm->doneEvent, 0) );
+  }
 
-/* Syncronize with user stream and launch the collective.
- * All work is performed asynchronously with the host thread.
- * The actual collective should be a functor with the
- * folloaing signature.
- * ncclResult_t collective(void* sendbuff, void* recvbuff,
- *                         int count, ncclDataType_t type, ncclRedOp_t op,
- *                         int root, ncclComm_t comm);
- * Unneeded arguments should be ignored. The collective may
- * assume that the appropriate cuda device has been set. */
-template<typename ColFunc>
-ncclResult_t enqueue(ColFunc colfunc,
-                     const void* sendbuff,
+  ncclResult_t ret;
+  ret = ColFunc<T, Op>::entry(sendbuff, recvbuff, count, root, comm, stream);
+
+  // Always have to record done event because we don't know what stream next
+  // collective will be in.
+  CUDACHECK( cudaEventRecord(comm->doneEvent, stream) );
+  comm->opSched += 1;
+  return ret;
+}
+
+
+// This version decodes type
+template< template<typename, template<typename> class> class ColFunc,
+          template<typename> class Op >
+ncclResult_t enqueue(const void* sendbuff,
+                     void* recvbuff,
+                     int count,
+                     ncclDataType_t type,
+                     int root,
+                     ncclComm_t comm,
+                     cudaStream_t stream)
+{
+  switch(type) {
+  case ncclChar:
+    return enqueue<ColFunc, char, Op>(sendbuff, recvbuff, count, root, comm, stream);
+  case ncclInt:
+    return enqueue<ColFunc, int, Op>(sendbuff, recvbuff, count, root, comm, stream);
+#ifdef CUDA_HAS_HALF
+  case ncclHalf:
+    return enqueue<ColFunc, half, Op>(sendbuff, recvbuff, count, root, comm, stream);
+#endif
+  case ncclFloat:
+    return enqueue<ColFunc, float, Op>(sendbuff, recvbuff, count, root, comm, stream);
+  case ncclDouble:
+    return enqueue<ColFunc, double, Op>(sendbuff, recvbuff, count, root, comm, stream);
+  case ncclInt64:
+    return enqueue<ColFunc, long long, Op>(sendbuff, recvbuff, count, root, comm, stream);
+  case ncclUint64:
+    return enqueue<ColFunc, unsigned long long, Op>(sendbuff, recvbuff, count, root, comm, stream);
+  default:
+    WARN("Invalid ncclType %d", type);
+    return ncclInvalidType;
+  }
+}
+
+// This version decodes both type and reduction op
+template< template<typename, template<typename> class> class ColFunc>
+ncclResult_t enqueue(const void* sendbuff,
                      void* recvbuff,
                      int count,
                      ncclDataType_t type,
@@ -34,24 +93,19 @@ ncclResult_t enqueue(ColFunc colfunc,
                      ncclComm_t comm,
                      cudaStream_t stream)
 {
-  int curDevice;
-  CUDACHECK( cudaGetDevice(&curDevice) );
-
-  // No need for a mutex here because we assume that all enqueue operations happen in a fixed
-  // order on all devices. Thus, thread race conditions SHOULD be impossible.
-  EventQueue* eq = &comm->events;
-
-  // Ensure that previous collective is complete
-  cudaError_t flag = cudaEventQuery(eq->isDone[eq->back]);
-  if( flag == cudaErrorNotReady )
-    CUDACHECK( cudaStreamWaitEvent(stream, eq->isDone[eq->back], 0) );
-
-  // Launch the collective here
-  ncclResult_t ret = colfunc(sendbuff, recvbuff, count, type, op, root, comm, stream);
-
-  eq->back = (eq->back + 1) % MAXQUEUE;
-  CUDACHECK( cudaEventRecord(eq->isDone[eq->back], stream) );
-  return ret;
+  switch(op) {
+  case ncclSum:
+    return enqueue<ColFunc, FuncSum>(sendbuff, recvbuff, count, type, root, comm, stream);
+  case ncclProd:
+    return enqueue<ColFunc, FuncProd>(sendbuff, recvbuff, count, type, root, comm, stream);
+  case ncclMax:
+    return enqueue<ColFunc, FuncMax>(sendbuff, recvbuff, count, type, root, comm, stream);
+  case ncclMin:
+    return enqueue<ColFunc, FuncMin>(sendbuff, recvbuff, count, type, root, comm, stream);
+  default:
+    WARN("Invalid ncclRedOp: %d", op);
+    return ncclInvalidOperation;
+  }
 }
 
 #endif // End include guard