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?
Thank you! We'll create more content like this.
Thank you for helping us improve!