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
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
AnyType,
ImplicitlyDestructible
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[*?, *?].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[*?, *?].static_product
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!