LayerNorm tutorial: what is the purpose of atomic_xchg for Count?

[artbataev]

I'm trying to understand LayerNorm tutorial.
I figured out the main algorithm, but I still have a question about atomic operations in grouped computation of dw/db in backward.

Specifically, I'm asking about the operation tl.atomic_xchg(Count, 1) in the code:

def _layer_norm_bwd_dx_fused(...):
    ...
    # Accumulate partial sums for dw/db
    partial_dw = (dy * xhat).to(w.dtype)
    partial_db = (dy).to(w.dtype)
    while tl.atomic_cas(Lock, 0, 1) == 1:
        pass
    count = tl.load(Count)
    # First store doesn't accumulate
    if count == 0:
        tl.atomic_xchg(Count, 1)
    else:
        partial_dw += tl.load(DW, mask=mask)
        partial_db += tl.load(DB, mask=mask)
    tl.store(DW, partial_dw, mask=mask)
    tl.store(DB, partial_db, mask=mask)
    # Release the lock
    tl.atomic_xchg(Lock, 0)

Why do we need tl.atomic_xchg here?
Why do we not use just the tl.store(Count, 1)?

If I correctly understand, the program is under the lock while tl.atomic_cas(Lock, 0, 1) == 1:..., and using tl.store should be correct in this case (like the code tl.store(DW, partial_dw, mask=mask) does correct work).