Mojo struct
OutputRegisterBuffer
struct OutputRegisterBuffer[dtype: DType, num_m_mmas: Int, num_n_mmas: Int, output_frag_size: Int]
Fields
- reg_tile (
OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].RegType):
Implemented traits
AnyType,
ImplicitlyDestructible
comptime members
reg_dtype
comptime reg_dtype = dtype
reg_layout
comptime reg_layout = row_major[(num_n_mmas * num_m_mmas), output_frag_size]()
RegType
comptime RegType = TileTensor[dtype, Layout[*?, *?], MutExternalOrigin, address_space=AddressSpace.LOCAL]
Methods
__init__
__init__(out self)
apply_softmax_denominator
apply_softmax_denominator[layout_type: TensorLayout, //](self, rowsum: TileTensor[dtype, layout_type, address_space=rowsum.address_space, linear_idx_type=rowsum.linear_idx_type, element_size=rowsum.element_size])
zero
zero(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!