Skip to main content

Mojo struct

SharedMemoryBuffer

struct SharedMemoryBuffer[BufferType: DType, layout: Layout, pipeline_stages: Int, warp_rows: Int, warp_cols: Int]

Fields

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

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, UnknownDestructibility

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial = LayoutTensor[BufferType, layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128].__copyinit__is_trivial

__del__is_trivial

alias __del__is_trivial = LayoutTensor[BufferType, layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128].__del__is_trivial

smem_2D_tile_layout

alias smem_2D_tile_layout[tensor_layout: Layout] = Layout.__init__(IntTuple.__init__[__origin_of()](tensor_layout.shape[1], tensor_layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](tensor_layout.stride[1], tensor_layout.stride[2], Tuple[]())) if ((tensor_layout.rank() == 2) ^ True) else tensor_layout

Parameters

SmemTileType

alias SmemTileType = LayoutTensor[BufferType, layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128]

SmemTileType2D

alias SmemTileType2D = LayoutTensor[BufferType, Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128]

WarpTileType

alias WarpTileType = LayoutTensor[BufferType, LayoutTensor._compute_tile_layout[True, BufferType, Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, MutableAnyOrigin, AddressSpace(3), Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, AddressSpace(3)), _get_index_type(Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, AddressSpace(3)), False, 128, warp_rows, warp_cols]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, AddressSpace(3)), linear_idx_type=_get_index_type(Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, AddressSpace(3)), masked=_tile_is_masked[Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, warp_rows, warp_cols](), alignment=128]

Methods

__init__

__init__(out self)

get_tile

get_tile(self, stage: Int) -> LayoutTensor[BufferType, Layout.__init__(IntTuple.__init__[__origin_of()](layout.shape[1], layout.shape[2], Tuple[]()), IntTuple.__init__[__origin_of()](layout.stride[1], layout.stride[2], Tuple[]())) if ((layout.rank() == 2) ^ True) else layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=128]

Returns:

LayoutTensor

Was this page helpful?