For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
SMemTileArray2DRowMajor
struct SMemTileArray2DRowMajor[dtype: DType, dim0: Int, dim1: Int, num_tiles: Int, alignment: Int = Int(128)]
Array of TileTensor tiles in shared memory with row_major layout.
Unlike SMemTileArray2D which uses swizzled internal_k_major layout, this type uses simple row_major layout. Suitable for 1D vectors (like A-scales) or output tiles where swizzling is not needed.
Example: comptime MyArray = SMemTileArray2DRowMajor[DType.float32, 1, 64, 4]
var array = MyArray.stack_allocation() var tile = array[0] # Returns TileTensor with row_major layout
Parameters
- dtype (
DType): Tile element data type. - dim0 (
Int): First dimension (rows). - dim1 (
Int): Second dimension (columns). - 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, _safe=False]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
num_elements
comptime num_elements = (Int((mul dim0, dim1)) * num_tiles)
Storage
comptime Storage = Array[Scalar[dtype], Int((mul dim0, dim1, num_tiles))]
storage_size
comptime storage_size = (Int((mul dim0, dim1, num_tiles)) * size_of[dtype]())
Tile
comptime Tile = TileTensor[dtype, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]
tile_layout
comptime tile_layout = row_major[dim0, dim1]()
tile_size
comptime tile_size = (dim0 * dim1)
Methods
__init__
def __init__(ref[AddressSpace._value] storage: Array[Scalar[dtype], Int((mul dim0, dim1, num_tiles))]) -> Self
Initialize from inline storage.
Args:
- storage (
Array[Scalar[dtype], Int((mul dim0, dim1, num_tiles))]): The inline storage array.
Returns:
Self: A new SMemTileArray2DRowMajor pointing to the storage.
def __init__(unsafe_ptr: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]) -> Self
Initialize with a shared memory pointer.
Args:
- unsafe_ptr (
Pointer[Scalar[dtype], address_space=AddressSpace.SHARED, _safe=False]): 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 toInt.
Args:
- index (
T): The tile index.
Returns:
Self.Tile: A TileTensor-based tile at the given index with row_major layout.
slice
def slice[length: Int](self, start: Int) -> SMemTileArray2DRowMajor[dtype, dim0, dim1, length, alignment]
Get a slice of the array.
Parameters:
- length (
Int): The length of the slice.
Args:
- start (
Int): The starting index.
Returns:
SMemTileArray2DRowMajor[dtype, dim0, dim1, length, alignment]: A new SMemTileArray2DRowMajor representing the slice.
stack_allocation
static def stack_allocation() -> Self
Allocate the array on the stack (in shared memory).
Returns:
Self: A new SMemTileArray2DRowMajor 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!