Mojo struct
OutputRegisterBuffer
struct OutputRegisterBuffer[dtype: DType, num_m_mmas: Int, num_n_mmas: Int, output_frag_size: Int]
Fields
- reg_tile (
LayoutTensor[dtype, OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]):
Implemented traits
AnyType,
RegisterBuffer,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = True
reg_dtype
alias reg_dtype = dtype
reg_tile_layout
alias reg_tile_layout = Layout.row_major((num_n_mmas * num_m_mmas), output_frag_size)
RegisterTileType
alias RegisterTileType = LayoutTensor[dtype, OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Methods
__init__
__init__(out self)
get_dtype
vectorize
vectorize(self) -> LayoutTensor[dtype, coalesce(LayoutTensor._compute_tile_layout[True, dtype, OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, MutAnyOrigin, AddressSpace.LOCAL, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL), _get_index_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL), False, align_of[dtype](), 1, output_frag_size]()[1], True), MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=LayoutTensor._divide_tiles[True, dtype, OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, MutAnyOrigin, AddressSpace.LOCAL, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL), _get_index_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL), False, align_of[dtype](), 1, output_frag_size]()[0], layout_int_type=_get_layout_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL), linear_idx_type=_get_index_type(OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, AddressSpace.LOCAL)]
Returns:
apply_softmax_denominator
apply_softmax_denominator(self, rowsum: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment])
zero
zero(self)
get_reg_tile
get_reg_tile(self) -> LayoutTensor[dtype, OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!