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
RegisterAccumulatorLayout
struct RegisterAccumulatorLayout[MMA_M: Int, MMA_N: Int, num_m_mmas: Int, num_n_mmas: Int, consumer_group_size: Int, *, frag_simdwidth: Int = Int(2)]
Describes how UMMA accumulator fragments are distributed across the threads of a consumer warp group.
Parameters
- MMA_M (
Int): The M dimension of a single UMMA instruction tile. - MMA_N (
Int): The N dimension of a single UMMA instruction tile. - num_m_mmas (
Int): Number of UMMA tiles along the M dimension. - num_n_mmas (
Int): Number of UMMA tiles along the N dimension. - consumer_group_size (
Int): Number of threads in the consumer warp group. - frag_simdwidth (
Int): SIMD width of each accumulator fragment (defaults to 2).
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
element_layout
comptime element_layout = Layout.row_major(Int(1), frag_simdwidth)
frag_size
comptime frag_size = ((MMA_M * MMA_N) // consumer_group_size)
num_row_blocks_per_mma
comptime num_row_blocks_per_mma = 2
rows_of_frags_layout
comptime rows_of_frags_layout = Layout.row_major((num_m_mmas * num_n_mmas), (Int((mul MMA_M, MMA_N)) // consumer_group_size))
vec_output_layout
comptime vec_output_layout = Layout(IntTuple(IntTuple(Int(2), num_m_mmas), IntTuple(((Int((mul MMA_M, MMA_N)) // consumer_group_size) // Int((mul frag_simdwidth, 2))), num_n_mmas), __list_literal__=NoneType(None)), IntTuple(IntTuple(frag_simdwidth, (Int((mul MMA_M, MMA_N)) // consumer_group_size)), IntTuple((Int(2) * frag_simdwidth), Int((mul (Int((mul MMA_M, MMA_N)) // consumer_group_size), num_m_mmas))), __list_literal__=NoneType(None)))
Methods
description
static def description() -> RegisterAccumulatorDescription
Returns:
RegisterAccumulatorDescription