Skip to main content

Mojo struct

TMemOperand

@register_passable(trivial) struct TMemOperand[dtype: DType, num_m_mmas: Int, num_n_mmas: Int, MMA_M: Int, MMA_N: Int, MMA_K: Int, num_consumer_threads: Int]

Fields

  • tmem_addr (SIMD[uint32, 1]):

Implemented traits

AnyType, Copyable, Movable, UnknownDestructibility, WriteableMMAOperandDescriptor

Aliases

frag_size

alias frag_size = 0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)

reg_layout

alias reg_layout = RegisterAccumulatorLayout[MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_consumer_threads]

reg_tile_t

alias reg_tile_t = LayoutTensor[dtype, __init__[::Origin[::Bool(__init__[::Origin[::Bool(IntTuple(2, num_m_mmas), IntTuple((div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4) + -1) if ((0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) < 0) & ((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4) == 0) ^ True)) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4), num_n_mmas), Tuple()), __init__[::Origin[::Bool(IntTuple(2, 0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)), IntTuple(4, (0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) * num_m_mmas)), Tuple())), MutableAnyOrigin, address_space=AddressSpace(5), element_layout=row_major(1, 2)]

vec_output_layout

alias vec_output_layout = __init__[::Origin[::Bool(__init__[::Origin[::Bool(IntTuple(2, num_m_mmas), IntTuple((div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4) + -1) if ((0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) < 0) & ((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4) == 0) ^ True)) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)})), "value">, 4), num_n_mmas), Tuple()), __init__[::Origin[::Bool(IntTuple(2, 0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">)), IntTuple(4, (0 if (num_consumer_threads == 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) == 0) ^ True) & (((MMA_M * MMA_N) < 0) ^ (num_consumer_threads < 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 num_consumer_threads, "value">, 0), {1}, num_consumer_threads), "value">) * num_m_mmas)), Tuple()))

Methods

__init__

__init__(tmem_addr: SIMD[uint32, 1]) -> Self

offset

offset[m_mma: Int, k_mma: Int](self) -> SIMD[uint32, 1]

copy_from

copy_from[src_type: DType, src_layout: Layout, src_element_layout: Layout, //](self, src: LayoutTensor[src_type, src_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=src_element_layout])

copy_to

copy_to[dst_type: DType, dst_layout: Layout, dst_element_layout: Layout, //](self, dst: LayoutTensor[dst_type, dst_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=dst_element_layout])

Was this page helpful?