Add back __syncthreads() in barrier and adjust stack size (#688)

src/collectives/device/common.h

@@ -21,7 +21,7 @@
 
 #ifdef __GFX9__
 #define STORE(DST, SRC) \
-  { atomicExch((unsigned long long *)(DST), (SRC)); }
+  { __threadfence(); atomicExch((unsigned long long *)(DST), (SRC)); }
 #else
 #define STORE(DST, SRC) \
   { __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST); }

src/collectives/device/primitives.h

@@ -17,12 +17,16 @@
 #define barrier_by_group() do { \
   const int w = threadIdx.x/WARP_SIZE; \
   const int wid = threadIdx.x%WARP_SIZE; \
-  if (wid == 0) { \
-    __asm__ __volatile__("s_waitcnt vmcnt(0) lgkmcnt(0)"); \
-    barrier_next[w] += nthreads/WARP_SIZE; \
-    atomicAdd((unsigned long long *)barriers, 1); \
-    while (atomicAdd((unsigned long long *)barriers, 0) < barrier_next[w]) __builtin_amdgcn_s_sleep(1); \
-    __asm__ __volatile__("s_wakeup"); \
+  __threadfence(); \
+  if (nthreads == NCCL_MAX_NTHREADS) { \
+    __syncthreads(); \
+  } else { \
+    if (wid == 0) { \
+      barrier_next[w] += nthreads/WARP_SIZE; \
+      atomicAdd((unsigned long long *)barriers, 1); \
+      while (atomicAdd((unsigned long long *)barriers, 0) < barrier_next[w]) __builtin_amdgcn_s_sleep(1); \
+      __asm__ __volatile__("s_wakeup"); \
+    } \
   } \
 } while (0)
 

src/init.cc

@@ -1415,7 +1415,7 @@ fail:
 
 #ifdef USE_INDIRECT_FUNCTION_CALL
 NCCL_PARAM(SetStackSize, "SET_STACK_SIZE", 1);
-RCCL_PARAM(StackSizeOverride, "STACK_SIZE_OVERRIDE", 8);
+RCCL_PARAM(StackSizeOverride, "STACK_SIZE_OVERRIDE", 256);
 #else
 NCCL_PARAM(SetStackSize, "SET_STACK_SIZE", 0);
 RCCL_PARAM(StackSizeOverride, "STACK_SIZE_OVERRIDE", 0);