Mojo struct
SMemTileArrayWithLayout
@register_passable(trivial)
struct SMemTileArrayWithLayout[shape_types: Variadic[CoordLike], stride_types: Variadic[CoordLike], //, dtype: DType, tile_layout: Layout[shape_types, stride_types], num_tiles: Int, alignment: 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): 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 (
LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
num_elements
comptime num_elements = (SMemTileArrayWithLayout[dtype, tile_layout, num_tiles, alignment].tile_size * num_tiles)
Storage
comptime Storage = InlineArray[Scalar[dtype], SMemTileArrayWithLayout[dtype, tile_layout, num_tiles, alignment].num_elements]
storage_size
comptime storage_size = (SMemTileArrayWithLayout[dtype, tile_layout, num_tiles, alignment].num_elements * 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__
__init__(ref[AddressSpace._value._mlir_value] storage: InlineArray[Scalar[dtype], SMemTileArrayWithLayout[dtype, tile_layout, num_tiles, alignment].num_elements]) -> Self
Initialize from inline storage.
Args:
- storage (
InlineArray): The inline storage array.
Returns:
Self: A new SMemTileArrayWithLayout pointing to the storage.
__init__[mut: Bool, //, origin: Origin[mut=mut]](unsafe_ptr: LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED, _mlir_origin=origin._mlir_origin, origin=origin]) -> Self
Initialize with a shared memory pointer.
Args:
- unsafe_ptr (
LegacyUnsafePointer): Pointer to shared memory storage.
__getitem__
__getitem__[T: Intable](self, index: T) -> SMemTileArrayWithLayout[dtype, tile_layout, num_tiles, alignment].Tile
Get tile at the given index.
Args:
- index (
T): The tile index.
Returns:
SMemTileArrayWithLayout: A TileTensor with correct swizzled layout at the given index.
slice
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: A new SMemTileArrayWithLayout representing the slice.
stack_allocation
static stack_allocation() -> Self
Allocate the array on the stack (in shared memory).
Returns:
Self: A new SMemTileArrayWithLayout backed by stack-allocated shared memory.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!