Skip to main content

Mojo struct

SharedMemoryBuffer

struct SharedMemoryBuffer[BufferType: DType, smem_tile_layout: Layout, pipeline_stages: Int, block_rows: Int, block_cols: Int, warp_rows: Int, warp_cols: Int]

Manages shared memory and returns 2D tile slices of the buffer.

Fields

  • buffer (LayoutTensor[BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, address_space=AddressSpace(3), alignment=128]):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, UnknownDestructibility

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial = True

__del__is_trivial

alias __del__is_trivial = True

SmemTensorType

alias SmemTensorType = LayoutTensor[BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, address_space=AddressSpace(3), alignment=128]

SmemTileType

alias SmemTileType = LayoutTensor[BufferType, LayoutTensor._compute_tile_layout[True, BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, AddressSpace(3), Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), False, 128, block_rows, block_cols]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), linear_idx_type=_get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), masked=_tile_is_masked[pipeline_layout[smem_tile_layout, pipeline_stages](), block_rows, block_cols](), alignment=128]

WarpTileType

alias WarpTileType = LayoutTensor[BufferType, LayoutTensor._compute_tile_layout[True, BufferType, LayoutTensor._compute_tile_layout[True, BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, AddressSpace(3), Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), False, 128, block_rows, block_cols]()[0], MutableAnyOrigin, AddressSpace(3), Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _tile_is_masked[pipeline_layout[smem_tile_layout, pipeline_stages](), block_rows, block_cols](), 128, warp_rows, warp_cols]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), linear_idx_type=_get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), masked=_tile_is_masked[pipeline_layout[smem_tile_layout, pipeline_stages](), block_rows, block_cols]() if _tile_is_masked[pipeline_layout[smem_tile_layout, pipeline_stages](), block_rows, block_cols]() else _tile_is_masked[LayoutTensor._compute_tile_layout[True, BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, AddressSpace(3), Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), False, 128, block_rows, block_cols]()[0], warp_rows, warp_cols](), alignment=128]

Methods

__init__

__init__(out self)

get_tile

get_tile(self, stage: Int) -> LayoutTensor[BufferType, LayoutTensor._compute_tile_layout[True, BufferType, pipeline_layout[smem_tile_layout, pipeline_stages](), MutableAnyOrigin, AddressSpace(3), Layout(IntTuple(1), IntTuple(1)), _get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), _get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), False, 128, block_rows, block_cols]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), linear_idx_type=_get_index_type(pipeline_layout[smem_tile_layout, pipeline_stages](), AddressSpace(3)), masked=_tile_is_masked[pipeline_layout[smem_tile_layout, pipeline_stages](), block_rows, block_cols](), alignment=128]

Returns:

LayoutTensor

Was this page helpful?