Mojo struct
RegisterAccumulatorLayout
@register_passable(trivial)
struct RegisterAccumulatorLayout[MMA_M: Int, MMA_N: Int, num_m_mmas: Int, num_n_mmas: Int, consumer_group_size: Int, *, frag_simdwidth: Int = 2]
Implemented traits
AnyType,
Copyable,
Movable,
UnknownDestructibility
Aliases
element_layout
alias element_layout = row_major(1, frag_simdwidth)
frag_size
alias frag_size = 0 if (consumer_group_size == 0) else (div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) + -1) if (((rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (consumer_group_size < 0))) else div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">)
num_row_blocks_per_mma
alias num_row_blocks_per_mma = 2
rows_of_frags_layout
alias rows_of_frags_layout = row_major((num_m_mmas * num_n_mmas), 0 if (consumer_group_size == 0) else (div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) + -1) if (((rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (consumer_group_size < 0))) else div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">))
vec_output_layout
alias vec_output_layout = __init__[::Origin[::Bool(__init__[::Origin[::Bool(IntTuple(2, num_m_mmas), IntTuple(0 if ((frag_simdwidth * 2) == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {0}, cond(and(ne(rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), 0), xor(lt(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), 0), lt(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0))), {value = add(div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), -1)}, {value = div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">)})), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2), 0), {1}, {value = mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2)}), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {0}, cond(and(ne(rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), 0), xor(lt(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), 0), lt(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0))), {value = add(div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), -1)}, {value = div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">)})), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2), 0), {1}, {value = mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2)}), "value">) == 0) ^ True) & (((frag_simdwidth * 2) < 0) ^ (0 if (consumer_group_size == 0) else (div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) + -1) if (((rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (consumer_group_size < 0))) else div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {0}, cond(and(ne(rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), 0), xor(lt(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), 0), lt(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0))), {value = add(div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">), -1)}, {value = div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">)})), "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2), 0), {1}, {value = mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int frag_simdwidth, "value">, 2)}), "value">), num_n_mmas), Tuple()), __init__[::Origin[::Bool(IntTuple(frag_simdwidth, 0 if (consumer_group_size == 0) else (div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) + -1) if (((rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (consumer_group_size < 0))) else div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">)), IntTuple((frag_simdwidth * 2), (0 if (consumer_group_size == 0) else (div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) + -1) if (((rem_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (consumer_group_size < 0))) else div_s(mul(#lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_M, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int MMA_N, "value">), #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int consumer_group_size, "value">, 0), {1}, consumer_group_size), "value">) * num_m_mmas)), Tuple()))
Methods
description
static description() -> RegisterAccumulatorDescription
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!