Mojo struct
PRegisterBuffer
struct PRegisterBuffer[accum_type_: DType, dtype: DType, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, num_m_mmas: Int, num_n_mmas: Int, output_frag_size: Int, shared_memory_backed: Bool, mma_shape: IndexList[3], k_group_size: Int, tr_load_enabled: Bool = False]
Fields
- reg_tile (
LayoutTensor[accum_type_, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]): - shared_memory_tile (
LayoutTensor[dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, MutAnyOrigin, address_space=AddressSpace.SHARED]):
Implemented traits
AnyType,
RegisterBuffer,
RegisterMMABuffer,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = True
mma_dtype
alias mma_dtype = dtype
mma_tile_layout
alias mma_tile_layout = Layout.row_major(num_m_mmas, simd_width_of[dtype]())
MMATileType
alias MMATileType = LayoutTensor[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
reg_dtype
alias reg_dtype = accum_type_
reg_tile_layout
alias reg_tile_layout = Layout.row_major((num_n_mmas * num_m_mmas), output_frag_size)
RegisterTileType
alias RegisterTileType = LayoutTensor[accum_type_, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
shared_memory_layout
alias shared_memory_layout = blocked_product(Layout.row_major(BM, BK), Layout.row_major(1, (BN // BK)), False)
SharedMemoryTileType
alias SharedMemoryTileType = LayoutTensor[dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, MutAnyOrigin, address_space=AddressSpace.SHARED]
Methods
__init__
__init__(out self, shared_ptr: LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED, mut=mut, origin=origin])
get_mma_tile_reg
get_mma_tile_reg[tile_idx: Int, k_idx: Int](self) -> LayoutTensor[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
get_mma_tile_shared
get_mma_tile_shared[tile_idx: Int, k_idx: Int](self) -> LayoutTensor[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
get_mma_tile
get_mma_tile[tile_idx: Int, k_idx: Int](self) -> LayoutTensor[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].mma_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
get_dtype
vectorize
vectorize(self) -> LayoutTensor[accum_type_, coalesce(LayoutTensor._compute_tile_layout[True, accum_type_, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, MutAnyOrigin, AddressSpace.LOCAL, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL), _get_index_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL), False, align_of[accum_type_](), 1, output_frag_size]()[1], True), MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=LayoutTensor._divide_tiles[True, accum_type_, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, MutAnyOrigin, AddressSpace.LOCAL, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL), _get_index_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL), False, align_of[accum_type_](), 1, output_frag_size]()[0], layout_int_type=_get_layout_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL), linear_idx_type=_get_index_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, AddressSpace.LOCAL)]
Returns:
zero
zero(self)
get_reg_tile
get_reg_tile(self) -> LayoutTensor[accum_type_, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].reg_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
get_shared_memory_tile
get_shared_memory_tile(self, tile_idx: Int) -> LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, MutAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, AddressSpace.SHARED), _get_index_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, AddressSpace.SHARED), False, align_of[dtype](), BM, BK]()[0], MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, AddressSpace.SHARED), linear_idx_type=_get_index_type(PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, AddressSpace.SHARED), masked=_tile_is_masked[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, k_group_size, tr_load_enabled].shared_memory_layout, BM, BK]()]
Returns:
copy_to_shared
copy_to_shared(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!