[Probe] Grid/system barrier gaps in distributed/sync.h: missing exit fence, single-use grid barrier, unordered in-kernel init

[Rachmanino]

Three related correctness gaps in the device-side synchronization helpers in src/tl_templates/cuda/distributed/sync.h. They are grouped because they all concern the same barrier/fence layer.

Confidence differs per item and is stated explicitly. (1) and (2) are code-level findings that I could not reproduce at runtime; (3) is reproduced only as "the ordering is not enforced", not as an observed failure. None of them is a demonstrated miscompile, so please treat this as an audit report rather than a bug with a repro attached.

1. barrier_blocks has no acquire fence on exit

sync.h:303

template <bool need_fence = true>
TL_DEVICE void barrier_blocks(int offset, int rank, int num_ranks) {
  if constexpr (need_fence) {
    memory_fence_sys();
    __syncthreads();
  }
  // ... arrive: atomicAdd_system / atomicSub_system ...
  while (true) {
    int value = tid < num_ranks ? ld_volatile<int32_t>(BARRIER_PTR(rank) + tid) : 0;
    if (__all_sync(0xffffffff, value <= 0)) break;
  }
  __syncthreads();          // <-- no fence here
}

The entry fence gives the arriving rank release ordering for its own prior writes. The waiting rank, after observing the flags, has no acquire edge before its subsequent loads of peer data.

ld_volatile's "memory" clobber prevents the compiler from sinking later loads above the spin, but it does not force hardware visibility: a later load may still be satisfied from a stale cached line. The counterpart barriers in NVSHMEM/DeepEP place a fence after the wait.

Suggested: memory_fence_sys() between the spin loop and the trailing __syncthreads(), or a comment explaining why the volatile spin is sufficient on the supported hardware.

The arrival protocol itself is sound and reusable — each rank's slot returns to 0 after every barrier — so this is specifically about the missing exit fence.

2. sync_barrier_gpu is single-use; a second call is a no-op

sync.h:247, with init_barrier_gpu at sync.h:219

template <const uint32_t kExpected>
TL_DEVICE void init_barrier_gpu(uint32_t *barrier) {
  if (IS_MASTER_BLOCK() && IS_MASTER_THREAD()) {
    *barrier = BARRIER_MAGIC - kExpected;      // 0x80000000 - expected
  }
  memory_fence_gpu();
}

TL_DEVICE void sync_barrier_gpu(uint32_t *barrier) {
  __syncthreads();
  if (IS_MASTER_THREAD()) {
    atomic_add_release_gpu_u32(barrier, 1);
    uint32_t arrive = ld_acquire_gpu_u32(barrier);
    while (arrive < BARRIER_MAGIC) { arrive = ld_acquire_gpu_u32(barrier); }
  }
  __syncthreads();
}

The counter starts at MAGIC - expected and each block contributes +1, so after expected arrivals it is exactly MAGIC. It is never reset. A second sync_barrier_gpu on the same barrier therefore starts at MAGIC, and arrive < MAGIC is already false on the first load — every block leaves immediately without waiting for anyone.

So a kernel with two grid barriers gets one real barrier and one no-op. Since both init_barrier_gpu and sync_barrier_gpu are exposed through tilelang.language.distributed.sync, a multi-phase kernel is a natural way to hit this.

wait_barrier_gpu (sync.h:236) tests !(arrive & BARRIER_MAGIC) instead of arrive < BARRIER_MAGIC; the two spellings behave the same on first use but are worth unifying.

On reproduction: I tried three ways to observe this and failed each time, so the arithmetic above is the only evidence.

• reusing one barrier buffer across three kernel launches — correct every time, because init_barrier_gpu re-runs per launch and resets the counter
• two sequential sync_barrier_gpu calls in one kernel with 64 blocks — correct, because the blocks are all resident and stay tightly aligned after the first barrier
• the same, with block 0 deliberately slowed by a 300k-iteration loop before its phase-2 contribution — still correct, and this attempt was invalid: the loop is present in the emitted CUDA but reduces to b[tid] += 300000, which NVCC collapses, so block 0 was never actually slow

A valid repro needs a delay NVCC cannot fold (e.g. a dependent chain through global memory, or %%globaltimer). I did not pursue it further.

Suggested: reset the counter at the end of a completed barrier, or use the phase/parity scheme that sync_grids_arrive/sync_grids_wait (sync.h:261) already implements correctly via (oldArrive ^ current_arrive) & 0x80000000.

3. init_barrier_gpu cannot safely initialize a barrier from inside the kernel that uses it

sync.h:219

The initializing store is done by block 0 thread 0 only, followed by memory_fence_gpu(). Nothing orders that store against the atomicAdd that other blocks perform in sync_barrier_gpu. If block 0 is scheduled late, an arrival that already happened is overwritten by the init store, the counter can no longer reach MAGIC, and every block spins forever.

This is inherent: a grid-wide barrier cannot be initialized from within the same kernel without an already-working grid-wide sync. The TODO on sync.h:224 ("Is fence or sync needed here?") points at the same doubt.

testing/python/distributed/primitives/sync/test_barrier_gpu.py uses exactly this pattern:

with T.Kernel(num_blocks, threads=threads) as bid:
    T.init_barrier_gpu(bar, num_blocks)
    ...
    T.atomic_add(A[tid], val[0])
    T.sync_barrier_gpu(bar)

It passes consistently — block 0 is normally scheduled first and there is work before the barrier — but the ordering is not enforced anywhere.

Suggested: initialize the barrier on the host (the buffer is already a user-provided tensor, so a zero_()-style fill of MAGIC - expected costs nothing), or in a separate kernel, and document that init_barrier_gpu is only valid when no block can reach the barrier before the initializing block. If in-kernel init must stay, it needs a comment stating the assumption.

Minor, same file

• atomic_add_release_gpu_u32 (sync.h:44) and every helper in distributed/atomic.h take const uint32_t * while atomically modifying the target.
• sync.h:47 uses atom.add.release.gpu.global.s32 for a uint32_t counter, whereas atomic.h consistently uses .u32. Two's-complement makes this benign, but the inconsistency is confusing next to the signed/unsigned care taken elsewhere.
• barrier_blocks passes a literal 0xffffffff to __all_sync. It works for num_ranks <= 32 with the current call pattern, but it asserts full-warp participation; a note on that assumption would help.

Environment

• branch release/v0.0.2-tilelang-550e25d (PR [Release] TileScale 0.0.2 on upstream TileLang 550e25d #61), read at 42374143
• runtime attempts on 8×B200, CUDA 13.2

[Rachmanino]

Follow-up on item 1 (the missing exit fence): the Python surface makes it easier to hit than it looks, because need_fence is exposed as two differently-named wrappers.

tilelang/language/distributed/sync.py:

def barrier_blocks(barrier):   # need_fence=1
def sync_blocks(barrier):      # need_fence=0

need_fence gates only the entry fence in barrier_blocks (sync.h:310), so combined with the missing exit fence:

	before arrival	after wait
barrier_blocks	fence.acq_rel.sys	none
sync_blocks	none	none

So sync_blocks provides no memory ordering in either direction -- it is purely a counter rendezvous -- while its name and its old docstring ("Synchronize all blocks at a system-level barrier without fence") both read like an ordinary barrier with a minor optimisation.

I have reworded both docstrings on the release branch to state exactly what ordering each one does and does not provide, and to point sync_blocks users at barrier_blocks unless the surrounding code already establishes the ordering. That is a documentation change only -- the fence behaviour itself is unchanged and item 1 above still stands.