Mojo struct
AMD_MMA
struct AMD_MMA[out_type: DType, in_type: DType, shape: IndexList[3], transpose_b: Bool, k_group_size: Int, num_k_tiles: Int, num_m_mmas: Int, num_n_mmas: Int, simd_width: Int, swizzle: Swizzle, BK: Int, WK: Int]
Implemented traits
AnyType,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = True
MMARegTileType
alias MMARegTileType[num_mmas: Int] = LayoutTensor[in_type, Layout.row_major((num_k_tiles * num_mmas), simd_width), MutableAnyOrigin, address_space=AddressSpace(5), alignment=align_of[SIMD[in_type, simd_width]]()]
Parameters
- num_mmas (
Int):
SharedMemTileType
alias SharedMemTileType[smem_layout: Layout] = LayoutTensor[in_type, smem_layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=align_of[SIMD[in_type, simd_width]]()]
Parameters
- smem_layout (
Layout):
SharedMemWarpTileType
alias SharedMemWarpTileType[warp_rows: Int, smem_layout: Layout] = LayoutTensor[in_type, LayoutTensor._compute_tile_layout[True, in_type, smem_layout, MutableAnyOrigin, AddressSpace(3), Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(smem_layout, AddressSpace(3)), _get_index_type(smem_layout, AddressSpace(3)), False, align_of[SIMD[in_type, simd_width]](), warp_rows, WK]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(smem_layout, AddressSpace(3)), linear_idx_type=_get_index_type(smem_layout, AddressSpace(3)), masked=_tile_is_masked[smem_layout, warp_rows, WK](), alignment=align_of[SIMD[in_type, simd_width]]()]
Parameters
tensor_core_mma
alias tensor_core_mma = TensorCoreKGroup[out_type, in_type, shape, k_group_size, transpose_b]()
type_alignment
alias type_alignment = align_of[SIMD[in_type, simd_width]]()
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!