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
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,
Deinitable,
ImplicitlyCopyable,
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