Async memory access

Load a chunk of data from global memory into a local memory buffer using TMA. The global address, strides, and buffer size are defined by the tensor descriptor. A barrier object is provided and signaled upon completion of the operation.

Parameters:

Gather independently selected rows from a 2D global tensor into consecutive rows of a shared-memory buffer. The operation is asynchronous and signals barrier after the full TMA transaction completes.

Parameters:

NUM_ROWS must be at least 8 and a multiple of 4. TMA gather supports element types up to 32 bits, and the descriptor block must contain at least 32 / element_bit_width * 8 columns. The x_offsets register layout must provide eight contiguous tensor elements to each issuing thread, arranged as two groups of four, and broadcast them across its warp. The row-index values themselves are independent: they do not need to be consecutive, aligned, or multiples of 4.

Conceptually, one gather4 message consumes four arbitrary row offsets and writes four row tiles into consecutive SMEM rows:

   x_offsets = [3, 17, 42, 5]

   Global:                         SMEM:
   row r0, column tile ---------> result row 0
   row r1, column tile ---------> result row 1
   row r2, column tile ---------> result row 2
   row r3, column tile ---------> result row 3

A row tile is the portion covered by one TMA message. When BLOCK_N is wider than one message, the compiler emits multiple gather4 instructions for the same four-row group, one per column tile.

Example:

   NUM_ROWS: tl.constexpr = 32
   BLOCK_N: tl.constexpr = 128

   desc = tl.make_tensor_descriptor(
       input_ptr,
       shape=[M, N],
       strides=[N, 1],
       block_shape=[1, BLOCK_N],
   )
   buffers = tlx.local_alloc((NUM_ROWS, BLOCK_N), tl.float16, 1)
   result = tlx.local_view(buffers, 0)
   barriers = tlx.alloc_barriers(1)
   barrier = tlx.local_view(barriers, 0)

   # Each warp owns 8 consecutive offsets, broadcast across all its lanes.
   # Each gather4 instruction consumes 4 offsets, so each warp issues 2.
   # This layout assumes 4 warps and NUM_ROWS == 32.
   offset_layout: tl.constexpr = tlx.layout(
       shape=((32, 4), (8,)),
       stride=((0, 8), (1,)),
   )
   offset_ids = tlx.require_layout(tl.arange(0, NUM_ROWS), offset_layout)
   x_offsets = tl.load(row_indices_ptr + offset_ids)

   tlx.barrier_expect_bytes(
       barrier,
       NUM_ROWS * BLOCK_N * tlx.size_of(tl.float16),
   )
   tlx.async_descriptor_gather(
       desc,
       result,
       x_offsets,
       y_offset=0,
       barrier=barrier,
   )
   tlx.barrier_wait(barrier, phase=0)

Hint hardware to load a chunk of data from global memory into a L2 cache to prepare for upcoming async_descriptor_load operations.

Store a chunk of data from shared memory into global memory using TMA. The global address, strides, and buffer size are defined by the tensor descriptor.

Supports optional atomic reduction (store_reduce) and L2 cache eviction hints (eviction_policy). Both regular stores and atomic reduce stores support cache eviction policies.

Parameters:

Example:

   # Regular TMA store with L2 evict_first hint
   tlx.async_descriptor_store(desc_c, c_buf[0], [offs_m, offs_n], eviction_policy="evict_first")

   # TMA atomic reduce-add with L2 evict_first hint
   tlx.async_descriptor_store(desc_c, c_buf[0], [offs_m, offs_n],
                              eviction_policy="evict_first", store_reduce="add")

Asynchronously store a register-source distributed tensor into cluster shared memory, or atomically reduce it with store_reduce="add", "min", "max", "and", "or", "xor", "inc", or "dec". The default store_reduce="" preserves the regular store behavior. ADD, MIN, MAX, AND, OR, and XOR support int32 and uint32; INC and DEC support only uint32. Source/destination shapes and element types must match, and the destination must be mutable.

MIN/MAX use signed comparisons for int32 and unsigned comparisons for uint32. For INC/DEC, each source value is a bound b, not an amount to add or subtract: inc(old, b) returns zero when old >= b, otherwise old + 1; dec(old, b) returns b when old == 0 or old > b, otherwise old - 1.

Reductions require a cluster containing at least two physical CTAs, configured through either num_ctas or ctas_per_cga. A physical two-CTA launch with logical num_ctas=1 is supported. remote_cta_rank may name a peer CTA or the executing CTA itself; destination and completion barrier are mapped to the same rank.

Regular stores use st.async; reductions use scalar red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.<op>.<type>. ADD uses .add.u32 for both signed and unsigned inputs because their 32-bit sums have the same bit-level result. MIN/MAX use .s32 or .u32; bitwise operations use .b32; INC/DEC use .u32. These instructions use the generic proxy, not the async proxy: asynchronous completion alone does not require fence("async_shared_cluster").

Each executed reduction completes 4 bytes per unique logical element on the receiving mbarrier, without an additional arrival. Allocate barriers with scope="cluster" for compiler-owned initialization visibility; this does not replace application-owned destination initialization, expected-byte registration, data-ready synchronization, completion waits, or storage lifetime protection. Expected-byte registration can follow transaction completion, but barrier initialization must precede remote use. Wait on the receiving barrier before consuming the result. Keep every participating CTA's destination and barrier storage alive until all remote operations finish, including a final cluster rendezvous before exit or storage reuse.

Parameters:

Example:

   # Allocate shared memory buffer and barrier
   buffer = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
   barrier = tlx.alloc_barriers(num_barriers=1, arrive_count=1, scope="cluster")

   # Store to remote CTA's shared memory
   tlx.async_remote_shmem_store(buffer[0], src_tensor, remote_cta_rank=1, barrier=barrier[0])

   # Integer register-source reduction after initialization, expected-byte
   # registration, and a cluster rendezvous (int_buffer matches src_i32).
   tlx.async_remote_shmem_store(int_buffer[0], src_i32, remote_cta_rank=1,
                                barrier=barrier[0], store_reduce="add")

Store a local shared memory buffer into a buffer in the remote shared memory of a cluster asynchronously.

Parameters:

Example:

  # Allocate shared memory buffer
  buffer0 = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
  buffer1 = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
  barrier = tlx.alloc_barriers(num_barriers=1, arrive_count=1)

  # Copy to remote CTA's shared memory
  tlx.remote_shmem_store(buffer0[0], buffer1[0], remote_cta_rank=1, barrier=barrier[0])

Allocates global memory for tensor descriptor storage with built-in parameters (nbytes=128, alignment=128 per descriptor). Returns a tensor_descriptor_ptr with 128-byte stride semantics that supports indexing.

Parameters:

Returns:

Example:

   # Allocate storage for 4 tensor descriptors
   desc_ptrs = tlx.allocate_tensor_descriptor(num=4)

   # Access individual descriptors using indexing
   desc_ptr_0 = desc_ptrs[0]  # First descriptor
   desc_ptr_1 = desc_ptrs[1]  # Second descriptor (128 bytes offset)

Create a TMA (Tensor Memory Accelerator) descriptor for efficient asynchronous data movement on Hopper and Blackwell GPUs.

Parameters:

Example:

   # Create a 2D tensor descriptor with automatic scratch allocation
   desc = tlx.make_tensor_descriptor(
       desc_ptr=None,  # Compiler allocates scratch memory automatically
       base=tensor_ptr,
       shape=[M, N],
       strides=[N, tl.constexpr(1)],
       block_shape=[64, 64],
   )

   # Or with explicit descriptor allocation for advanced use cases (e.g., pipelining)
   desc_ptrs = tlx.allocate_tensor_descriptor(num=2)

   # Create descriptor at index 0
   tlx.make_tensor_descriptor(
       desc_ptr=desc_ptrs[0],
       base=tensor_ptr,
       shape=[M, N],
       strides=[N, tl.constexpr(1)],
       block_shape=[64, 64],
   )

   # Reinterpret the descriptor for TMA operations
   desc = tlx.reinterpret_tensor_descriptor(
       desc_ptr=desc_ptrs[0],
       block_shape=[64, 64],
       dtype=tl.float16,
   )

   # Use with async TMA operations
   tlx.async_descriptor_load(desc, buffer, offsets=[m_offset, n_offset], barrier=mbar)

Reinterpret a tensor descriptor pointer as a TMA-backed tensor descriptor object.

Parameters:

Example:

   # Allocate and create descriptor
   desc_ptrs = tlx.allocate_tensor_descriptor(num=2)
   tlx.make_tensor_descriptor(desc_ptr=desc_ptrs[0], base=a_ptr, shape=[M, K], strides=[K, 1], block_shape=[128, 64])

   # Reinterpret for use with TMA
   a_desc = tlx.reinterpret_tensor_descriptor(desc_ptr=desc_ptrs[0], block_shape=[128, 64], dtype=tl.float16)
   tlx.async_descriptor_load(a_desc, buffer, offsets=[offs_m, offs_k], barrier=mbar)

Load a chunk of data from global memory into a local memory buffer asynchronously.

The operation returns a token object which can be used to track the completion of the operation.

On AMD, the non-bulk path uses target-specific direct-to-LDS lowering or a synchronous fallback when the copy cannot be issued directly. Direct transfers require 32 bits per lane on gfx942, 32 or 128 bits on gfx950, or 32, 64, or 128 bits on gfx1250. A supported API does not guarantee a hardware-asynchronous transfer. The bulk=True path is NVIDIA-specific; use TDM descriptor loads on gfx1250 for descriptor-based transfers. See the AMD copy lowering.

Masked copies: The TLX frontend defaults other to 0.0 when mask is provided and other is omitted. Pass another other value explicitly if the masked-out elements need a different fill. This applies to the non-bulk copy path; bulk mode does not accept mask or other. See the frontend implementation.

Commits all prior initiated but uncommitted async_load ops an async group. Optionally, each token represents a tracked async load operation.

Wait for completion of prior asynchronous copy operations. The pendings argument indicates the number of in-flight operations not completed. Optionally, each token represents a tracked async commit group operation.

tlx.buffer_load_to_local

[gfx942, gfx950]

Async load from global memory directly into shared (local) memory, bypassing registers. This is useful for producer warps that prefetch data into shared memory for other warps to consume. Unlike pointer-based async_load, this explicit buffer operation is rejected on gfx1250 by the target verifier; use the TDM path there.

token = tlx.buffer_load_to_local(
    dest,
    ptr,
    offsets,
    mask=None,
    other=None,
    cache_modifier="",
    contiguity=1,
)

Returns: A tlx.async_token that can be used with tlx.async_load_wait_group() to synchronize on the completion of the transfer.

Lowers to amdg.buffer_load_to_local, which is eventually lowered to rocdl.raw.ptr.buffer.load.async.lds — a single hardware instruction that moves data from global memory to LDS without going through VGPRs.

Requirements. The direct-to-LDS copy has the following hardware constraints:

contiguity=N is an opt-in correctness contract, not a speculative hint. It promises that each per-thread group of N offsets is numerically consecutive and suitably aligned, and that mask and other, when present, are uniform within each group. N must be a positive power of two and divide the number of elements owned by each thread. A false promise can change the values loaded.

For N > 1, dest must be the allocation itself or a direct local_view of a user-pinned padded_shared layout. The compiler derives an offset layout from that destination, verifies that it supports an exact direct-to-LDS write, and only then pins the offsets, mask, and other to the derived layout. An incompatible layout is a compilation error. contiguity=1 retains the default best-effort inference behavior.

TDM descriptor loads

tlx.update_tensor_descriptor(desc, add_offsets=None, set_bounds=None, pred=None, clamp_bounds=False) produces a positioned descriptor SSA value. add_offsets advances the tile position without changing bounds; set_bounds rewrites absolute bounds; and pred replaces the inherited predicate. Use clamp_bounds=True with add_offsets to derive the remaining OOB extent of an advanced tile.

tlx.async_amd_descriptor_load(desc, result, offsets=None, pred=None, clamp_bounds=True) issues an AMD TDM descriptor load from global memory to a TLX local buffer. If offsets is omitted, desc is used as already positioned and its predicate is preserved. The analogous store accepts offsets=None for the same reason. Both operations are available on TDM-capable AMD targets (gfx1250+) and should be synchronized with tlx.async_amd_descriptor_wait.

tlx.async_amd_descriptor_load_fused(members, cache_modifier="") fuses two to four AMD TDM descriptor loads behind one static hardware instruction. Each member is a (positioned_desc, destination, warp_used_hint) tuple. Position descriptors with tlx.update_tensor_descriptor before issuing the fused load. The hints must be non-empty, pairwise disjoint, and legal axis-aligned subsets of the CTA's waves; they do not need to cover every wave. Members must have the same descriptor rank but may have different shapes and element types. All members share one cache modifier.

Example:

a_load_desc = tlx.update_tensor_descriptor(
    a_desc, add_offsets=[off_m, k * BLOCK_K], pred=True, clamp_bounds=True)
b_load_desc = tlx.update_tensor_descriptor(
    b_desc, add_offsets=[k * BLOCK_K, off_n], pred=True, clamp_bounds=True)
a_tok = tlx.async_amd_descriptor_load_fused([
    (a_load_desc, tlx.local_view(a_buf, slot), 0b0011),
    (b_load_desc, tlx.local_view(b_buf, slot), 0b1100),
])
tlx.async_amd_descriptor_wait(0, [a_tok])