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
SmemLayouts
struct SmemLayouts[a_type: DType, b_type: DType, BM: Int, BN: Int, BK: Int, OutputM: Int, OutputN: Int, a_swizzle: TensorMapSwizzle, b_swizzle: TensorMapSwizzle, transpose_b: Bool]
Common SMEM layout definitions for matmul-family kernels.
Centralizes the A/B/C tile layout computation including the transpose-conditional B layout logic, eliminating ~10 lines of duplicated layout definitions from each SMEM struct.
Parameters
- a_type (
DType): Data type for A matrix tiles. - b_type (
DType): Data type for B matrix tiles. - BM (
Int): Block tile M dimension. - BN (
Int): Block tile N dimension. - BK (
Int): Block tile K dimension. - OutputM (
Int): Output tile M dimension. - OutputN (
Int): Output tile N dimension. - a_swizzle (
TensorMapSwizzle): Swizzle mode for A tiles. - b_swizzle (
TensorMapSwizzle): Swizzle mode for B tiles. - transpose_b (
Bool): Whether B is transposed (K-major).
Implemented traits
comptime members
a_smem_layout_typed
comptime a_smem_layout_typed = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())))
a_tile_elems
comptime a_tile_elems = Layout[TypeList[Coord[ComptimeInt[Int(8)], ComptimeInt[(BM // Int(8))]], Coord[ComptimeInt[(a_swizzle.bytes() // size_of[a_type]())], ComptimeInt[(BK // (a_swizzle.bytes() // size_of[a_type]()))]]](), TypeList[Coord[ComptimeInt[(a_swizzle.bytes() // size_of[a_type]())], ComptimeInt[Int(0) if identical(BM, 8) else Int((mul (a_swizzle.bytes() // size_of[a_type]()), 8))]], Coord[ComptimeInt[Int(1)], ComptimeInt[Int(0) if identical((a_swizzle.bytes() // size_of[a_type]()), BK) else Int((mul (a_swizzle.bytes() // size_of[a_type]()), BM))]]]()].static_product
b_smem_layout_typed
comptime b_smem_layout_typed = Layout(Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())), Coord(Coord(ComptimeInt(), ComptimeInt()), Coord(ComptimeInt(), ComptimeInt())))
b_tile_elems
comptime b_tile_elems = Layout[TypeList[Coord[ComptimeInt[Int(8)], ComptimeInt[(BN // Int(8))]], Coord[ComptimeInt[(b_swizzle.bytes() // size_of[b_type]())], ComptimeInt[(BK // (b_swizzle.bytes() // size_of[b_type]()))]]](), TypeList[Coord[ComptimeInt[(b_swizzle.bytes() // size_of[b_type]())], ComptimeInt[Int(0) if identical(BN, 8) else Int((mul (b_swizzle.bytes() // size_of[b_type]()), 8))]], Coord[ComptimeInt[Int(1)], ComptimeInt[Int(0) if identical((b_swizzle.bytes() // size_of[b_type]()), BK) else Int((mul (b_swizzle.bytes() // size_of[b_type]()), BN))]]]()].static_product