Mojo function
st_matrix_m_layout
st_matrix_m_layout[c_type: DType, WG_BM: Int, num_m_mmas: Int, num_consumer: Int]() -> Layout
Creates a layout for M-major st_matrix in the context of WGMMA C matrix. This meant to be used with swapAB, since the C matrix must be transposed during the write phase. This must also be used in conjuction with st_matrix transposed modifier.
The M-dimension tiling size WG_BM // 16, the number of MMA tiles num_m_mmas
in the N-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_BM (
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!