Skip to main content

Mojo struct

MmaOpSM100_SS

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[Int, Int, Int](1, 1, 1), a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, b_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, transpose_b: Bool = False]

Fields

  • idesc (UMMAInsDescriptor[MmaOpSM100_SS._get_umma_kind[a_type]()]):
  • mask (UInt16):

Implemented traits

AnyType, Copyable, Defaultable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

Methods

__init__

__init__() -> Self

mma

mma(self, a: TileTensor[a.dtype, a.LayoutType, a.origin, address_space=AddressSpace.SHARED, linear_idx_type=a.linear_idx_type, element_size=a.element_size], b: TileTensor[b.dtype, b.LayoutType, b.origin, address_space=AddressSpace.SHARED, linear_idx_type=b.linear_idx_type, element_size=b.element_size], c_tmem: UInt32, init_c: Bool)

Issue MMA operations over K tiles from shared memory to TMEM.

Args:

  • a (TileTensor): A operand tile in shared memory.
  • b (TileTensor): B operand tile in shared memory.
  • c_tmem (UInt32): TMEM address for the accumulator.
  • init_c (Bool): When True, zero-initialize the accumulator on the first K slice instead of accumulating.

commit

commit(self, ptr_mbar: UnsafePointer[ptr_mbar.type, ptr_mbar.origin, address_space=AddressSpace.SHARED])

wait

wait(self)

Was this page helpful?