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")

Store a distributed tensor into a buffer in the remote shared memory of a cluster asynchronously. Signals the provided mbarrier when the store completes.

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)

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

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.

MI350X caveat: When mask is provided, callers should also provide other—typically 0.0 for numerical kernels. Elements for which mask=True are copied from global memory into the destination local buffer. For elements where mask=False, other is written only if it is provided. Otherwise, the corresponding LDS locations retain unspecified, potentially stale contents from an earlier use of the buffer.

For example, in FlashAttention, when the sequence length is not a multiple of BLOCK_N, the final V tile contains masked-out rows. If other is omitted, those rows may contain unspecified values, including bit patterns representing NaN or infinity. These values can propagate through the subsequent matrix multiplication and produce incorrect output. Use other=0.0 for such padded tiles.

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

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.

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])