Mojo struct
MmaOpSM100_SS
@register_passable(trivial)
struct MmaOpSM100_SS[c_type: DType, a_type: DType, b_type: DType, block_tile_shape: IndexList[3], mma_shape: IndexList[3], /, *, accum_type: DType = DType.float32, cta_group: Int = 1, cluster_shape: IndexList[3] = Index(1, 1, 1), a_swizzle: TensorMapSwizzle = 3, b_swizzle: TensorMapSwizzle = 3, transpose_b: Bool = False]
Fields
- idesc (
UMMAInsDescriptor[MmaOpSM100_SS._get_umma_kind[c_type, a_type, b_type, block_tile_shape, mma_shape, accum_type, cta_group, cluster_shape, a_swizzle, b_swizzle, transpose_b, a_type]()]): - mask (
UInt16):
Implemented traits
AnyType,
Copyable,
Defaultable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
__moveinit__is_trivial
alias __moveinit__is_trivial = True
Methods
__init__
__init__() -> Self
mma
mma(self, a: LayoutTensor[dtype, layout, origin, address_space=AddressSpace(3), element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], b: LayoutTensor[dtype, layout, origin, address_space=AddressSpace(3), element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], c_tmem: UInt32, init_c: Bool)
MMA input tiles.
The layout assumes that coalesce(A) has shape (bm, sw_k, num_sw_k), we currently assumes bm = mma_m. In future, we can tile it to (mma_m, sw_k, num_sw_k, num_mma_m) The same logic applies to matrix B.
commit
commit(self, ptr_mbar: UnsafePointer[type, address_space=AddressSpace(3), alignment=alignment, mut=mut, origin=origin])
wait
wait(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!