Mojo function

wgmma_c_layout

wgmma_c_layout[mma_m: Int, mma_n: Int, C: Layout]() -> List[Layout]

Generates three layouts for mapping WGMMA C matrix coordinates.

This function creates three layout mappings that are essential for working with WGMMA (Warp Group Matrix Multiply-Accumulate) operations:

1. A projection layout that maps linearized indices to row coordinates (i)
2. A projection layout that maps linearized indices to column coordinates (j)
3. A composite layout that maps thread and vector coordinates to linearized indices across multiple MMA tiles

These layouts are particularly useful for operations like attention masking and matrix multiplication epilogues, where register values need to be mapped to the coordinate system of the C matrix.

Note: This function enforces constraints on the WGMMA dimensions and ensures the C matrix dimensions are compatible with the WGMMA instruction size.

Parameters:

• ​mma_m (Int): The M dimension (rows) of a single WGMMA instruction, must be 64.
• ​mma_n (Int): The N dimension (columns) of a single WGMMA instruction, must be multiple of 8.
• ​C (Layout): The layout of the C matrix within a thread block.

Returns:

List[Layout] - A list containing three layouts:

1. proj_i: Maps linearized indices to row coordinates
2. proj_j: Maps linearized indices to column coordinates
3. TV_tile_to_idx: Maps thread/vector/tile coordinates to linearized indices

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit