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

MMAOperandOffsetFn

struct MMAOperandOffsetFn[dtype: DType, BMN: Int, BK: Int, swizzle: TensorMapSwizzle, is_k_major: Bool, WMMA_MN: Int, WMMA_K: Int]

Computes the shared-memory layout and byte offsets for MMA operand tiles, bridging typed tile layouts to legacy MMA descriptors.

Parameters

dtype (DType): The element type of the operand tile.
BMN (Int): The non-K dimension of the operand tile in elements.
BK (Int): The K dimension of the operand tile in elements.
swizzle (TensorMapSwizzle): The shared-memory swizzle mode applied to the operand layout.
is_k_major (Bool): Whether the operand is stored K-major versus M/N-major.
WMMA_MN (Int): The M or N dimension of a single warp-level MMA tile.
WMMA_K (Int): The K dimension of a single warp-level MMA tile.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

`comptime` members

`canonical_K`

comptime canonical_K = (swizzle.bytes() // size_of[dtype]()) if (swizzle != TensorMapSwizzle.SWIZZLE_NONE) else BK

`canonical_layout`

comptime canonical_layout = tile_to_descriptor[dtype, MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_layout_flat, is_k_major]()

`canonical_layout_flat`

comptime canonical_layout_flat = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt()))).to_layout() if is_k_major else MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].layout

`canonical_layout_size`

comptime canonical_layout_size = MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].canonical_layout.size()

`layout`

comptime layout = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt()))).to_layout() if is_k_major else Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt()))).transpose().to_layout()

`layout_size`

comptime layout_size = MMAOperandOffsetFn[dtype, BMN, BK, swizzle, is_k_major, WMMA_MN, WMMA_K].layout.size()

Methods

`init`

def __init__() -> Self

Parameters​

Implemented traits​

comptime members​

canonical_K​

canonical_layout​

canonical_layout_flat​

canonical_layout_size​

layout​

layout_size​

Methods​

__init__​