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
__del__is_trivial
comptime __del__is_trivial = True
a_smem_layout
comptime a_smem_layout = tile_layout_k_major[a_type, BM, BK, a_swizzle]()
b_smem_layout
comptime b_smem_layout = tile_layout_k_major[b_type, BN, BK, b_swizzle]() if transpose_b else tile_layout_mn_major[b_type, BN, BK, b_swizzle]()
c_smem_layout
comptime c_smem_layout = Layout.row_major(OutputM, OutputN)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!