Barriers and pipelines

Verb-style API over the raw mbarrier.* wrappers, plus two storage types: BarrierArray (one homogeneous slot ring) and BarrierSet (a typed, named layout over a whole SMEM mbarrier arena — the shape a warp-specialized kernel's synchronization plan actually has). Mirrors <cuda/barrier> in scope: init / arrive / arriveexpecttx / wait / cluster-arrive, with the arrive variants covering producer (expect_tx) and consumer (plain arrive) sides.

Sister module PTX.Pipelines builds the N-stage producer/consumer ring on top of these — mirrors <cuda/pipeline>.

Contents are plain sm90 (mbarrier, mapa, fence). Kernels that use them only become arch-specific via what other ops they pull in (e.g. wgmma → sm90a, tcgen05.mma → sm_100a).

BarrierArray{N}(base::Core.LLVMPtr{UInt64, AS.Shared})

Typed SMEM mbarrier array. N is the slot count; base points at the first of N consecutive 8-byte mbarriers. Indexing (0-based stage) returns the raw shared-AS pointer expected by all mbarrier ops (and by TMA copies that take a completion mbarrier), so callers can pass full[stage] directly to wrappers without an unwrap step.

barrier_init(mbar, count)

Initialize an SMEM mbarrier with an arrival count — mbarrier.init.shared.b64.

barrier_arrive(mbar)

Consumer-side arrival — mbarrier.arrive.shared.b64. Returns the UInt64 state token.

barrier_arrive_expect_tx(mbar, tx)

Producer-side arrival declaring tx expected transaction bytes — mbarrier.arrive.expect_tx.shared.b64. Pair with TMA copies whose completion mechanism is mbarrier::complete_tx::bytes.

barrier_wait(mbar, phase)

Spin on mbarrier.test_wait.parity until the named phase has been observed. Lowers to a setp ; @!pred bra loop — identical to the hand-rolled while !test_wait_parity ... end callers were writing.

test_wait vs try_wait are NOT interchangeable: testwait is a non-suspending poll; trywait may park the thread until a scheduler wake. The Blackwell tcgen05 commit-drain idiom is specifically try_wait.parity, and casually swapping the two — or the phase literal — is a known hang footgun. Hence two distinct verbs (barrier_try_wait is the other), not one with a flag: when porting, pick the one the source kernel uses, verbatim.

barrier_try_wait(mbar, phase)

Spin on mbarrier.try_wait.parity (the potentially-suspending form) until phase is observed — the tcgen05 commit → mbarrier::arrive::one drain. Same loop shape as barrier_wait; lowers byte-identically to the hand-rolled while !try_wait_parity ... end the blackwell kernels were writing.

barrier_arrive_cluster(mbar, remote_rank)

Cluster-scope arrive: translate the local mbarrier address through mapa.shared::cluster to the remote CTA's view, then arrive on it. The mbarrier itself lives in some CTA's SMEM; mapa rebases the SMEM offset to a cluster-mapped address so the remote CTA's hardware sees the arrival.

BarrierSet{SPEC}(base::Core.LLVMPtr{UInt64, AS.Shared})

Typed, named layout over a contiguous SMEM mbarrier arena. SPEC is a NamedTuple value type parameter mapping each group name to a (shape, count) pair:

• shape — () for a single barrier, (n,) / (n, m) / deeper for indexed groups (row-major, 0-based, matching BarrierArray);
• count — the arrival count barrierset_init! initializes every barrier of the group with.

Groups are laid out in declaration order, 8 bytes per barrier. bars.name returns the raw shared-AS pointer for a ()-shaped group and an indexable BarrierGroup otherwise; every offset is a compile-time constant (the name lookup is a tuple recursion over keys(SPEC) that constant propagation folds to a literal).

const RING = (full = ((4,), 1), free = ((4,), 1),
              stats = ((2, 2), 128), done = ((), 1))
bars = BarrierSet{RING}(pointer(smem_bars))
barrier_wait(bars.full[slot], phase)
barrier_arrive(bars.stats[stage, parity])
barrier_arrive_expect_tx(bars.done, tx)

The arena footprint for the enclosing SMEM map is barrier_bytes(typeof(bars)); per-group offsets are exposed for host-side layout assertions via barrier_offset.

BarrierGroup{SHAPE}

Indexed view over one named group of a BarrierSet. Row-major, 0-based indexing with one index per shape dimension returns the raw shared-AS pointer, like BarrierArray. Obtained via bars.name; not constructed directly.

barrier_bytes(::Type{<:BarrierSet}) -> Int

Total arena footprint in bytes (8 per barrier, groups in declaration order) — for computing the enclosing SMEM map.

barrier_offset(::Type{<:BarrierSet}, name::Symbol) -> Int

Byte offset of group name from the arena base. Host-side companion to the device accessors, for layout assertions in tests.

barrierset_init!(bs::BarrierSet)

Initialize every barrier in the arena with its group's declared arrival count, then publish with fence.proxy.async.shared::cta. Same caller contract as PTX.Pipelines.pipeline_init!: invoke from a single thread and follow with bar.sync. Ring self-crediting pre-arrives remain the caller's job — this initializes counts only.

Body is fully unrolled via @generated; the emitted PTX matches the hand-rolled init sequence byte-for-byte.

N-stage producer/consumer ring built on top of PTX.MBarriers. Mirrors <cuda/pipeline> in scope: a phantom Pipeline{N} type, the (stage, phase) = pipeline_cursor(P, k_iter) arithmetic, and the thread-0 init ritual that every Hopper warp-spec kernel repeats verbatim.

Pure leaf module: depends on PTX.MBarriers for the verbs + BarrierArray storage type, but no method dispatches on a Barrier type — Pipelines just calls verb functions and treats BarrierArray as an indexable pointer.

Pipeline{N_STAGES}

Phantom type carrying the stage count for cursor arithmetic. Use pipeline_cursor(Pipeline{N_STAGES}, k_iter) to get (stage, phase).

pipeline_stage(::Type{Pipeline{N}}, k_iter) -> Int32

Stage cursor: k_iter mod N. Compile-time-constant divisor → LLVM strength-reduces to & (N-1) for power-of-two N and to a multiply- and-shift otherwise.

pipeline_phase(::Type{Pipeline{N}}, k_iter) -> UInt32

Phase parity: which round through this stage's mbarrier we're on — (k_iter ÷ N) & 1. For pow2 N this folds to a single shift; for non-pow2 N (e.g. N=3) it lowers to a divide-by-constant which LLVM strength-reduces.

pipeline_cursor(::Type{Pipeline{N}}, k_iter) -> (stage, phase)

The (pipeline_stage, pipeline_phase) pair for iteration k_iter.

pipeline_init!(full::BarrierArray{N}, empty::BarrierArray{N},
               ::Val{EMPTY_COUNT}, ::Val{CLUSTER})

Initialize the producer/consumer mbarrier pair for an N-stage ring and pre-arrive empty[s] once per consumer so the producer's first wait succeeds without a prior consumer arrival.

• full : mbarriers consumers wait on after the producer finishes a stage. Init count = 1 (the producer's expect_tx arrive flips the phase by itself).
• empty : mbarriers the producer waits on after consumers finish a stage. Init count = EMPTY_COUNT (one arrive per releasing party per stage).
• EMPTY_COUNT is the per-stage release count: NUMCONSUMERS for a single CTA, NUMCONSUMERS × CLUSTERS for a clustered pipeline.
• CLUSTER=true emits fence.mbarrier_init.release.cluster (visible cluster-wide) instead of the CTA-scope fence.proxy.async.shared::cta.

Caller responsibility:

• Invoke from a single thread (typically thread 0).
• Follow with bar.sync(0) to publish init within the CTA.
• If CLUSTER=true, follow with ptx"barrier.cluster.arrive"(); ptx"barrier.cluster.wait"() (PTX tier-2 wrappers — convergent by registry, and unlike CUDACore's clusterarrive/clusterwait they don't trap on Julia ≤ 1.11; see wrappers/barrier_cluster.jl).

Body is fully unrolled across stages and pre-arrives via @generated; the emitted PTX matches the hand-rolled version byte-for-byte.

Warp-level collective idioms on top of the shfl.sync wrappers.

Sister module to PTX.MBarriers/PTX.Pipelines: no new hardware surface — every call lowers to the reviewed shfl.sync.bfly.b32 wrapper — but it names the butterfly-reduce ladder that every row-reduction kernel (softmax, the norm family, attention softmax streams) hand-rolls. Not exported; access as using PTX.Warps: warp_reduce.

warp_reduce(op, v::T[, ::Val{W} = Val(32)]) -> T

Butterfly (XOR-shuffle) all-reduce: after log₂(W) rounds every lane of each W-lane segment holds op folded over the segment's values. W must be a power of two in 2:32 (32 = the full warp); T must be a 4-byte isbits type (Float32, Int32, UInt32 — the shuffle moves raw b32 lanes and op sees T).

op is any two-argument function, applied between shuffle rounds exactly as written — pass the combining op the kernel means. ptx"..." ops are callable singletons, so they pass directly: warp_reduce(ptx"max.f32", v) and warp_reduce(max, v) (Julia max lowers to the NaN-propagating max.NaN.f32) are different instructions, and the choice stays at the call site.

Convergent: all 32 lanes of the warp must reach the call with the full membermask's lanes active, including lanes whose segment result goes unused.

The ladder is pinned: XOR offsets descend W/2, …, 1, and the shuffle c-operand is ((32 - W) << 8) | 0x1F (CUDA's __shfl_*_sync width encoding; segments never mix because XOR offsets below W cannot cross a W-aligned boundary). At W = 32 this is instruction- identical to the hand-rolled descending ladders it replaces.