Skip to main content

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 = float32, cta_group: Int = 1, a_swizzle: TensorMapSwizzle = TensorMapSwizzle(3), b_swizzle: TensorMapSwizzle = TensorMapSwizzle(3), transpose_b: Bool = False]

Fields

  • idesc (UMMAInsDescriptor[UMMAKind(2)]):

Implemented traits

AnyType, Copyable, Defaultable, Movable, UnknownDestructibility

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: SIMD[uint32, 1], 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?