Mojo struct
TensorCoreAsync
struct TensorCoreAsync[c_type: DType, a_type: DType, b_type: DType, mma_shape: Index[3], /, a_swizzle: TensorMapSwizzle = TensorMapSwizzle(__init__[__mlir_type.!kgen.int_literal](0)), b_swizzle: TensorMapSwizzle = TensorMapSwizzle(__init__[__mlir_type.!kgen.int_literal](0)), transpose_b: Bool = False]
Implemented traits
AnyType,
UnknownDestructibility
Methods
__init__
__init__(out self)
wgmma
static wgmma[num_warp_groups: Int = 1](a_smem_tile: LayoutTensor[a_type, layout, mut=mut, origin=origin, address_space=AddressSpace(3), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], b_smem_tile: LayoutTensor[b_type, layout, mut=mut, origin=origin, address_space=AddressSpace(3), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], c_reg_tile: LayoutTensor[c_type, layout, mut=mut, origin=origin, address_space=AddressSpace(5), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], wg_idx: Int = 0)
static wgmma[num_warp_groups: Int = 1](a_frag: LayoutTensor[a_type, layout, mut=mut, origin=origin, address_space=AddressSpace(5), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], b_smem_tile: LayoutTensor[b_type, layout, mut=mut, origin=origin, address_space=AddressSpace(3), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], c_reg_tile: LayoutTensor[c_type, layout, mut=mut, origin=origin, address_space=AddressSpace(5), element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], wg_idx: Int = 0)
arrive
static arrive()
commit_group
static commit_group()
wait_for_all
static wait_for_all()
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!