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: Layout - A layout that maps warp group local thread index to an offset
for a logical array with the element type c_type.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!