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
- tensor_layout (
Layout
):
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:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!