From d2580c8cf5d336dedc676ed59a6c5e8687588f4c Mon Sep 17 00:00:00 2001 From: Wenkai Du Date: Wed, 8 Sep 2021 16:14:32 -0500 Subject: [PATCH] Improve barrier implementation [ROCm/rccl commit: adb8d633528861165252ac53b043d5f96f95dd24] --- .../rccl/src/collectives/device/primitives.h | 16 ++++++++++------ projects/rccl/src/include/devcomm.h | 4 ++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/projects/rccl/src/collectives/device/primitives.h b/projects/rccl/src/collectives/device/primitives.h index c64d0bd153..c2c2e2cd10 100644 --- a/projects/rccl/src/collectives/device/primitives.h +++ b/projects/rccl/src/collectives/device/primitives.h @@ -33,12 +33,16 @@ } while (0) #define barrier_by_group() do { \ - const int w = threadIdx.x/WARP_SIZE; \ - const int wid = threadIdx.x%WARP_SIZE; \ - if (wid == 0) { \ - barrier_next[w] += nthreads/WARP_SIZE; \ - __atomic_fetch_add(barriers, 1, __ATOMIC_SEQ_CST); \ - while (LOAD(barriers) < barrier_next[w]) /* spin */; \ + if (nthreads == NCCL_MAX_NTHREADS) \ + __syncthreads(); \ + else { \ + const int w = threadIdx.x/WARP_SIZE; \ + const int wid = threadIdx.x%WARP_SIZE; \ + if (wid == 0) { \ + barrier_next[w] += nthreads/WARP_SIZE; \ + __atomic_fetch_add(barriers, 1, __ATOMIC_SEQ_CST); \ + while (LOAD(barriers) < barrier_next[w]) /* spin */; \ + } \ } \ } while (0) diff --git a/projects/rccl/src/include/devcomm.h b/projects/rccl/src/include/devcomm.h index 42976e36c7..b7923f97eb 100644 --- a/projects/rccl/src/include/devcomm.h +++ b/projects/rccl/src/include/devcomm.h @@ -18,8 +18,8 @@ // Convert volatile access to atomic #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) -#define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_SEQ_CST) -#define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST) +#define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_ACQUIRE) +#define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_RELEASE) #else #define LOAD(VAR) *(VAR) #define STORE(DST, SRC) *(DST) = (SRC)