For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

TmemAddress

struct TmemAddress

Simple TMEM address wrapper for load/store operations.

Encapsulates TMEM address encoding for accumulator fragment access. SM100 MMA accumulators are organized as 32 rows, split into:

• Upper fragment (rows 0-15): accessed via upper_addr()
• Lower fragment (rows 16-31): accessed via lower_addr()

The lower fragment address adds TMEM_LOWER_ROW_OFFSET (16 << 16) to encode the row offset in the upper 16 bits of the address.

Usage: var tmem = TmemAddress(base_offset)

# Load operations
var upper = tmem.load_upper[dtype, size]()
var lower = tmem.load_lower[dtype, size]()
TmemAddress.wait_load()

# Store operations
tmem.store_upper[dtype, size](upper_frag)
tmem.store_lower[dtype, size](lower_frag)
TmemAddress.wait_store()

# Low-level address access for custom operations
raw_upper = tmem.upper_addr()
raw_lower = tmem.lower_addr()

Fields

• ​addr (UInt32):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

Methods

__init__

def __init__(addr: Int) -> Self

Create TmemAddress from integer column address.

Args:

• ​addr (Int): TMEM column address as a plain integer.

def __init__(addr: UInt32) -> Self

Create TmemAddress from hardware address (UInt32).

Args:

• ​addr (UInt32): Raw TMEM hardware address as produced by tcgen05_alloc.

__add__

def __add__(self, offset: Int) -> Self

Create new TmemAddress with column offset added.

Args:

• ​offset (Int): Number of TMEM columns to advance the address by.

upper_addr

def upper_addr(self) -> UInt32

Raw address for upper fragment (rows 0-15).

Returns:

UInt32

lower_addr

def lower_addr(self) -> UInt32

Raw address for lower fragment (rows 16-31).

Returns:

UInt32

load_upper

def load_upper[dtype: DType, width: Int, data_paths: Int = Int(16), bits: Int = Int(256), repeat: Int = Int(1)](self) -> Array[Scalar[dtype], width]

Load upper accumulator fragment (rows 0-15).

Parameters:

• ​dtype (DType): Element type to load from TMEM.
• ​width (Int): Total number of elements per row in the returned Array.
• ​data_paths (Int): Number of SM100 TMEM data paths (defaults to 16).
• ​bits (Int): Bits loaded per data path per repeat (defaults to 256).
• ​repeat (Int): Number of times to repeat the load pattern (defaults to 1).

Returns:

Array[Scalar[dtype], width]

load_lower

def load_lower[dtype: DType, width: Int, data_paths: Int = Int(16), bits: Int = Int(256), repeat: Int = Int(1)](self) -> Array[Scalar[dtype], width]

Load lower accumulator fragment (rows 16-31).

Parameters:

• ​dtype (DType): Element type to load from TMEM.
• ​width (Int): Total number of elements per row in the returned Array.
• ​data_paths (Int): Number of SM100 TMEM data paths (defaults to 16).
• ​bits (Int): Bits loaded per data path per repeat (defaults to 256).
• ​repeat (Int): Number of times to repeat the load pattern (defaults to 1).

Returns:

Array[Scalar[dtype], width]

store_upper

def store_upper[dtype: DType, width: Int, data_paths: Int = Int(16), bits: Int = Int(256), repeat: Int = Int(1)](self, data: Array[Scalar[dtype], width])

Store upper accumulator fragment (rows 0-15).

Parameters:

• ​dtype (DType): Element type to store to TMEM.
• ​width (Int): Total number of elements per row in the data Array.
• ​data_paths (Int): Number of SM100 TMEM data paths (defaults to 16).
• ​bits (Int): Bits stored per data path per repeat (defaults to 256).
• ​repeat (Int): Number of times to repeat the store pattern (defaults to 1).

Args:

• ​data (Array[Scalar[dtype], width]): Array of accumulator data to store.

store_lower

def store_lower[dtype: DType, width: Int, data_paths: Int = Int(16), bits: Int = Int(256), repeat: Int = Int(1)](self, data: Array[Scalar[dtype], width])

Store lower accumulator fragment (rows 16-31).

Parameters:

• ​dtype (DType): Element type to store to TMEM.
• ​width (Int): Total number of elements per row in the data Array.
• ​data_paths (Int): Number of SM100 TMEM data paths (defaults to 16).
• ​bits (Int): Bits stored per data path per repeat (defaults to 256).
• ​repeat (Int): Number of times to repeat the store pattern (defaults to 1).

Args:

• ​data (Array[Scalar[dtype], width]): Array of accumulator data to store.

wait_store

static def wait_store()

Wait for TMEM store operations to complete.

wait_load

static def wait_load()

Wait for TMEM load operations to complete.