Mojo struct
SMemBuffer
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 (
SMemBuffer[dtype, layout, pipeline_stages, BM, BN, WM, WN].SMemTile):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
BlockTileType
comptime BlockTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[BM, BN]()[0], MutAnyOrigin, 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]
SMemTile
comptime SMemTile = LayoutTensor[dtype, pipeline_layout[layout, pipeline_stages](), MutAnyOrigin, address_space=AddressSpace.SHARED, alignment=128]
WarpTileType
comptime WarpTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[WM, WN]()[0], MutAnyOrigin, 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]() or _tile_is_masked[LayoutTensor._compute_tile_layout[BM, BN]()[0], WM, WN](), alignment=128]
Methods
__init__
__init__() -> Self
get_tile
get_tile(self, stage: Int) -> SMemBuffer[dtype, layout, pipeline_stages, BM, BN, WM, WN].BlockTileType
Returns:
SMemBuffer[dtype, layout, pipeline_stages, BM, BN, WM, WN].BlockTileType
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!