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 module

smem_types

Shared memory type aliases for GPU kernels.

This module defines the core SMEM type aliases used across SM90, SM100, and other GPU kernel implementations. They depend only on layout and stdlib, making them safe to import without pulling in higher-level kernel packages.

Types:

• SMemTile: Shared memory tile
• RegTile: Register tile
• SMemBarrier: Pointer to shared memory barrier
• PipelineBarrier: Array of pipeline barriers
• SMemTileArray: Array of shared memory tiles
• SMemArray: Generic shared memory array
• SMemPtr: Typed pointer into shared memory
• eval: Compile-time expression evaluator

Retained legacy: SMemTile and RegTile are still backed by the legacy tile type. Their consumers across linalg/, nn/attention/, and test/ construct them from raw pointers and pass them to legacy-only APIs, and many of those files are owned by other migration waves, so the aliases cannot be flipped to TileTensor until those consumers migrate. New code should prefer the native TileTensor-based SMemTile in structured_kernels.tile_types.

comptime values

eval

comptime eval[T: AnyType, //, val: T] = val

Helper alias to force evaluation of expressions at compile time.

Parameters

• ​T (AnyType):
• ​val (T):

PipelineBarrier

comptime PipelineBarrier[num_pipeline_stages: Int] = SMemArray[SharedMemBarrier, num_pipeline_stages]

Type alias for shared memory pipeline barrier array.

Parameters

• ​num_pipeline_stages (Int):

RegTile

comptime RegTile[_dtype: DType, layout: Layout, /, *, element_layout: Layout = Layout(IntTuple(Int(1)), IntTuple(Int(1))), layout_int_type: DType = _get_layout_type(layout, AddressSpace.LOCAL), linear_idx_type: DType = _get_index_type(layout, AddressSpace.LOCAL), masked: Bool = False, alignment: Int = align_of[_dtype]()] = LayoutTensor[_dtype, layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]

Type alias for register (local memory) tile tensors.

Parameters

• ​_dtype (DType):
• ​layout (Layout):
• ​element_layout (Layout):
• ​layout_int_type (DType):
• ​linear_idx_type (DType):
• ​masked (Bool):
• ​alignment (Int):

SMemBarrier

comptime SMemBarrier[origin: MutOrigin] = Pointer[SharedMemBarrier, origin, address_space=AddressSpace.SHARED]

Type alias for shared memory barrier pointer.

Parameters

• ​origin (MutOrigin):

SMemPtr

comptime SMemPtr[type: AnyType] = Pointer[type, MutUntrackedOrigin, address_space=AddressSpace.SHARED]

Parameters

• ​type (AnyType):

SMemTile

comptime SMemTile[_dtype: DType, layout: Layout, /, *, element_layout: Layout = Layout(IntTuple(Int(1)), IntTuple(Int(1))), layout_int_type: DType = _get_layout_type(layout, AddressSpace.SHARED), linear_idx_type: DType = _get_index_type(layout, AddressSpace.SHARED), masked: Bool = False, alignment: Int = align_of[_dtype]()] = LayoutTensor[_dtype, layout, MutAnyOrigin, address_space=AddressSpace.SHARED, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]

Type alias for shared memory tile tensors.

Parameters

• ​_dtype (DType):
• ​layout (Layout):
• ​element_layout (Layout):
• ​layout_int_type (DType):
• ​linear_idx_type (DType):
• ​masked (Bool):
• ​alignment (Int):

Structs

• ​SMemArray: Shared memory array of fixed size.
• ​SMemTileArray: Array of tiles in shared memory.

Functions

• ​reg_tile_to_tile_tensor: Return a TileTensor view of a register tile.