msccl: use atomic to set dependency flags (#941)

[ROCm/rccl commit: a7400218a2]
Šī revīzija ir iekļauta:
Wenkai Du
2023-10-31 14:46:57 -07:00
revīziju iesūtīja GitHub
vecāks 691df735a3
revīzija 3eeaea3f00
@@ -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++;
}