diff --git a/projects/rccl/src/collectives/device/msccl_kernel_impl.h b/projects/rccl/src/collectives/device/msccl_kernel_impl.h index 44505d1a13..b439212439 100644 --- a/projects/rccl/src/collectives/device/msccl_kernel_impl.h +++ b/projects/rccl/src/collectives/device/msccl_kernel_impl.h @@ -307,7 +307,11 @@ __device__ __forceinline__ void mscclRunInterpreter( int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid]; uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep); while (true){ - uint64_t curFlag = (mscclFlags + dependentBid)->flag; +#if defined(__gfx942__) + uint64_t curFlag = __atomic_load_n(&(mscclFlags + dependentBid)->flag, __ATOMIC_ACQUIRE); +#else + uint64_t curFlag = __atomic_load_n(&(mscclFlags + dependentBid)->flag, __ATOMIC_RELAXED); +#endif if (curFlag >= goalFlag && GET_WORKINDEX_FROM_FLAG(curFlag) == workIndex) break; } } @@ -451,7 +455,11 @@ __device__ __forceinline__ void mscclRunInterpreter( return; } if (t->hasDependence && tid == nthreads-1){ - mscclFlags[bid].flag = (uint64_t) COMPUTE_FLAG(workIndex, iter, step); +#if defined(__gfx942__) + __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELEASE); +#else + __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), __ATOMIC_RELAXED); +#endif } step++; }