Mojo struct
SMemBuffer
@register_passable(trivial)
struct SMemBuffer[dtype: DType, layout: Layout, pipeline_stages: Int, BM: Int, BN: Int, WM: Int, WN: Int]
Manages shared memory and returns 2D tile slices of the buffer.
Fields
- buffer (
LayoutTensor[dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, address_space=AddressSpace.SHARED, alignment=128]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
BlockTileType
alias BlockTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, BM, BN]()[0], MutableAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), linear_idx_type=_get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), masked=_tile_is_masked[pipeline_layout[layout, pipeline_stages](), BM, BN](), alignment=128]
SMemTileType
alias SMemTileType = LayoutTensor[dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, address_space=AddressSpace.SHARED, alignment=128]
WarpTileType
alias WarpTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, BM, BN]()[0], MutableAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _tile_is_masked[pipeline_layout[layout, pipeline_stages](), BM, BN](), 128, WM, WN]()[0], MutableAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), linear_idx_type=_get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), masked=_tile_is_masked[pipeline_layout[layout, pipeline_stages](), BM, BN]() if _tile_is_masked[pipeline_layout[layout, pipeline_stages](), BM, BN]() else _tile_is_masked[LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, BM, BN]()[0], WM, WN](), alignment=128]
Methods
__init__
__init__() -> Self
get_tile
get_tile(self, stage: Int) -> LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, pipeline_layout[layout, pipeline_stages](), MutableAnyOrigin, AddressSpace.SHARED, Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), _get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), False, 128, BM, BN]()[0], MutableAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), linear_idx_type=_get_index_type(pipeline_layout[layout, pipeline_stages](), AddressSpace.SHARED), masked=_tile_is_masked[pipeline_layout[layout, pipeline_stages](), BM, BN](), alignment=128]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!