Fix to https://github.com/NVIDIA/nccl/issues/560

ncclGroup's containing operations of mixed datatype, element, or collective
would induce crash.


[ROCm/rccl commit: 5f2f2f670f]

projects/rccl/src/collectives/device/common.h

@@ -71,9 +71,30 @@ template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
 struct RunWork {
   __device__ void run(ncclWork *w) {
     int tid = threadIdx.x;
-    #pragma unroll 1
-    for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].active != 0; e++) {
-      if (tid < w->elems[e].nThreads)
+    /* Some invariants that must hold:
+     * 1. All elems[] have same funcIndex.
+     * 2. All elems[] have same nThreads.
+     * 3. The thread-to-group relation (as in prims group numbers) is the same
+     *    for all elems[].
+     *
+     * If (1) isn't true then we might be in the wrong function since dispatch
+     * on ncclFuncs[w->elems[0].funcIndex] is how we got here.
+     *
+     * If (2) or (3) aren't true, then threads from different work elements
+     * could race for barrier resources (barrier numbers 0...15) which is fatal.
+     *
+     * Important, to ensure (3), implementations of
+     * `RunWorkElement<Fn,T,RedOp,Algo,Proto>::run()` may only use values which
+     * are the same for all elems[] when deciding how to map threads to groups,
+     * such as  the following:
+     *    Fn, T, RedOp, Algo, Proto, nThreads
+     *
+     * This last one is difficult to enforce and diagnosing it is a headeache.
+     * Device-side developers, consider yourselves warned.
+     */
+    if (tid < w->elems[0].nThreads) {
+      #pragma unroll 1
+      for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].active != 0; e++)
         RunWorkElement<Fn, T, RedOp, Algo, Proto>().run(&w->elems[e]);
     }
   }