Skip to main content

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]

Implemented traits

AnyType, UnknownDestructibility

Aliases

MMARegTileType

alias MMARegTileType = LayoutTensor[in_type, row_major((num_k_tiles * num_mmas), simd_width), MutableAnyOrigin, address_space=AddressSpace(5), alignment=alignof[::AnyType,__mlir_type.!kgen.target]()]

SharedMemTileType

alias SharedMemTileType = LayoutTensor[in_type, smem_layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=alignof[::AnyType,__mlir_type.!kgen.target]()]

SharedMemWarpTileType

alias SharedMemWarpTileType = LayoutTensor[in_type, _compute_tile_layout[*::Int]()[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[::Layout,*::Int](), alignment=alignof[::AnyType,__mlir_type.!kgen.target]()]

tensor_core_mma

alias tensor_core_mma = TensorCoreKGroup()

type_alignment

alias type_alignment = alignof[::AnyType,__mlir_type.!kgen.target]()

Was this page helpful?