Tlx API Reference

TLX API Quick Reference

Warp Specialization

Warp specialization skeleton

with tlx.async_tasks():
    with tlx.async_task([0]):       # Producer
        # TMA loads
    with tlx.async_task([1, 2]):    # Consumers
        # MMA compute

Memory Barriers

mbarrier (shared-memory allocated)

Ordinary barrier count rules:

• alloc_barriers(..., arrive_count=N) initializes the number of logical arrivals required to complete a phase. Transaction bytes registered with barrier_expect_bytes are an additional completion condition; do not infer the software arrival topology from the transaction byte count.
• For TMA full/data-ready barriers where one task registers the transaction, use arrive_count=1 unless the surrounding protocol explicitly requires additional software arrivals.
• For local software barriers shared by replicated tlx.async_task consumers, the ordinary count is typically the number of consumer replicas that arrive once per phase.
• For cross-CTA software arrivals using remote_cta_rank, derive the count from the exact CTA protocol. Do not assume the local-task rule applies.

Warp-barrier semantics and safe conversion

alloc_warp_barrier changes the physical arrival protocol, not merely the spelling of the allocation. With an ordinary barrier, TLX selects a leader to perform the arrival and may synchronize the participating threads first. With a warp barrier, every participating thread performs its own arrival, avoiding that leader-path synchronization but issuing more mbarrier arrival operations.

The initialized count is:

expected arrivals per phase = num_warps * 32 * num_arrivals

Here num_warps is the number of warps in each participating task replica, and num_arrivals is the number of such all-thread arrival events targeting the same barrier in one phase. For example, two replicated four-warp consumers that each release one shared buffer slot use num_warps=4, num_arrivals=2, for 256 arrivals. A consumer-owned slot released by one four-warp replica uses num_warps=4, num_arrivals=1, for 128 arrivals. Keep the corresponding barrier_arrive calls at their original final-use points and normally use the default unit arrival count.

Before converting an ordinary barrier, classify its role and arrival source:

Safe-conversion checklist:

1. Identify the protected buffer and confirm the barrier is signaled by explicit software barrier_arrive calls rather than TMA, MMA, multicast, or remote completion.
2. Enumerate every task replica, warp, lane, predicate, and arrival call that targets the barrier during one phase.
3. Verify num_warps * 32 * num_arrivals equals the exact number of unit arrivals.
4. Prove every counted lane reaches the arrival after its final access to the buffer, including prologue, tail, and persistent-loop iterations.
5. Preserve buffer indices, phase calculations, wait sites, and arrival sites; change only the allocator for the first A/B experiment.
6. Keep full/TMA barriers ordinary even when the paired empty/reuse barriers are converted.
7. Run correctness and benchmark both forms. Warp barriers trade additional per-thread arrivals for removal of leader-path synchronization, so conversion is an optimization candidate, not a universal rule.

Named barriers (hardware-allocated, indices 0–15)

num_threads is the total number of threads required to flip the barrier phase: num_waiting_threads + num_arriving_threads. Wait and Arrive calls for the same barrier phase must use the same value. The count must be a multiple of 32 (warp size); it is typically num_warp_groups * warps_per_group * 32.

Used for PingPong scheduling to prevent tensor core contention between consumer warp groups.

Memory Operations

SMEM / TMEM allocation

Storage kinds: tlx.storage_kind.smem, tlx.storage_kind.tmem (Blackwell), tlx.storage_kind.smemCluster