Skip to main content
Log in

Mojo function

st_matrix_n_layout

st_matrix_n_layout[c_type: DType, WG_BN: Int, num_m_mmas: Int, num_consumer: Int]() -> Layout

Creates a layout for N-major st_matrix in the context of WGMMA C matrix.

The layout modes are: the warp group local thread index, the N-dimension tiling size WG_BN // 16, the number of MMA tiles num_m_mmas in the M-dimension, and the number of consumers num_consumer. The output is an offset for a logical array with the element type c_type.

Parameters:

  • c_type (DType): Data type of the C matrix.
  • WG_BN (Int): Size of the K dimension in the C matrix in shared memory.
  • num_m_mmas (Int): Number of MMA tiles in the M dimension.
  • num_consumer (Int): Number of consumers.

Returns:

Layout - A layout that maps warp group local thread index to an offset for a logical array with the element type c_type.