Mojo function

wgmma_c_layout

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

Generates three layouts based on the WGMMA instruction dimensions and the C matrix layout.

The function takes mma_m and mma_n (the dimensions of a single WGMMA instruction) and C: Layout (the layout of the C matrix within a thread block) to generate three layouts:

1. idx -> i
2. idx -> j
3. `(TV_to_idx, (num_m_mma, num_n_mma)) -> idx

Here, idx represents the linearized coordinate of C, while idx -> {i, j} are the projection functions that map the linearized coordinate into the Cartesian components i and j. This function is particularly useful for mapping register values into the coordinate system of C, such as in attention masking and the matmul epilogue.

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