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

SMemTileArrayWithLayout

struct SMemTileArrayWithLayout[dtype: DType, tile_layout: Layout[tile_layout.shape_types, tile_layout.stride_types], num_tiles: Int, alignment: Int = Int(128)]

Array of TileTensor tiles with explicit Layout (preserves swizzle info).

Unlike SMemTileArray2D which uses row_major internally, this type preserves the full layout information from TMA swizzling, enabling .to_layout_tensor() to produce correctly swizzled LayoutTensors.

Example: comptime swizzled_layout = tile_layout_k_majordtype, BM, BK, ... comptime MyArray = SMemTileArrayWithLayout[dtype, swizzled_layout, 4]

var array = MyArray.stack_allocation() var tile = array[0] # Returns TileTensor with swizzled layout var lt = tile.to_layout_tensor() # Correctly swizzled!

Parameters

• ​dtype (DType): Tile element data type.
• ​tile_layout (Layout[tile_layout.shape_types, tile_layout.stride_types]): The full layout including swizzle information.
• ​num_tiles (Int): Number of tiles in the array.
• ​alignment (Int): Memory alignment (default 128 for shared memory).

Fields

• ​ptr (Pointer[Scalar[dtype], MutUntrackedOrigin, address_space=AddressSpace.SHARED]):

Implemented traits

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

comptime members

num_elements

comptime num_elements = (tile_layout.product() * num_tiles)

Storage

comptime Storage = Array[Scalar[dtype], Int((mul tile_layout.product(), num_tiles))]

storage_size

comptime storage_size = (Int((mul tile_layout.product(), num_tiles)) * size_of[dtype]())

Tile

comptime Tile = TileTensor[dtype, Layout[shape_types, stride_types], MutAnyOrigin, address_space=AddressSpace.SHARED]

tile_size

comptime tile_size = tile_layout.product()

Methods

__init__

def __init__(ref[AddressSpace._value] storage: Array[Scalar[dtype], Int((mul tile_layout.product(), num_tiles))]) -> Self

Initialize from inline storage.

Args:

• ​storage (Array[Scalar[dtype], Int((mul tile_layout.product(), num_tiles))]): The inline storage array.

Returns:

Self: A new SMemTileArrayWithLayout pointing to the storage.

def __init__(unsafe_ptr: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> Self

Initialize with a shared memory pointer.

Args:

• ​unsafe_ptr (Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]): Pointer to shared memory storage.

__getitem__

def __getitem__[T: Intable](self, index: T) -> Self.Tile

Get tile at the given index.

Parameters:

• ​T (Intable): Index value type, must be convertible to Int.

Args:

• ​index (T): The tile index.

Returns:

Self.Tile: A TileTensor with correct swizzled layout at the given index.

slice

def slice[length: Int](self, start: Int) -> SMemTileArrayWithLayout[dtype, tile_layout, length, alignment]

Get a slice of the array.

Parameters:

• ​length (Int): The length of the slice.

Args:

• ​start (Int): The starting index.

Returns:

SMemTileArrayWithLayout[dtype, tile_layout, length, alignment]: A new SMemTileArrayWithLayout representing the slice.

stack_allocation

static def stack_allocation() -> Self

Allocate the array on the stack (in shared memory).

Returns:

Self: A new SMemTileArrayWithLayout backed by stack-allocated shared memory.