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
OutputRegisterBufferRDNA
struct OutputRegisterBufferRDNA[dtype: DType, num_m_mmas: Int, num_n_mmas: Int]
Output accumulator register buffer. Layout is (num_n_mmas * num_m_mmas, RDNA_CD_FRAG_SIZE) row_major: one row per MMA tile, one column per per-lane C/D register.
Parameters
- dtype (
DType): Element type of the output accumulator registers. - num_m_mmas (
Int): Number of MMA tiles along the M (query) axis. - num_n_mmas (
Int): Number of MMA tiles along the N (key) axis.
Fields
- reg_tile (
OutputRegisterBufferRDNA[dtype, num_m_mmas, num_n_mmas].RegisterTileType):
Implemented traits
comptime members
output_frag_size
comptime output_frag_size = RDNA_CD_FRAG_SIZE
reg_dtype
comptime reg_dtype = dtype
reg_tile_layout
comptime reg_tile_layout = row_major[(num_n_mmas * num_m_mmas), Int(8)]()
RegisterTileType
comptime RegisterTileType = TileTensor[dtype, Layout[*(), *()], MutAnyOrigin, address_space=AddressSpace.LOCAL]
Methods
__init__
def __init__(out self)
get_dtype
zero
def zero(self)
get_reg_tile
def get_reg_tile[stage: Int = Int(0)](self) -> Self.RegisterTileType
Returns:
Self.RegisterTileType