For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

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.

Parameters

• ​dtype (DType): Element data type stored in the shared memory buffer.
• ​layout (Layout): 2D layout of a single pipeline stage's tile in shared memory.
• ​pipeline_stages (Int): Number of pipeline stages buffered in shared memory.
• ​BM (Int): Block tile height in rows; must equal the row extent of layout.
• ​BN (Int): Block tile width in columns; must equal the column extent of layout.
• ​WM (Int): Warp tile height in rows; must evenly divide BM.
• ​WN (Int): Warp tile width in columns; must evenly divide BN.

Fields

• ​buffer (SMemBuffer[dtype, layout, pipeline_stages, BM, BN, WM, WN].SMemTile):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

BlockTileType

comptime BlockTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[BM, BN]()[Int(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=Int(128)]

SMemTile

comptime SMemTile = LayoutTensor[dtype, pipeline_layout[layout, pipeline_stages](), MutAnyOrigin, address_space=AddressSpace.SHARED, alignment=Int(128)]

WarpTileType

comptime WarpTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[WM, WN]()[Int(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]()[Int(0)], WM, WN](), alignment=Int(128)]

Methods

__init__

def __init__() -> Self

get_tile

def get_tile(self, stage: Int) -> Self.BlockTileType

Returns:

Self.BlockTileType