Mojo function
get_mma_shape
get_mma_shape[input_type: DType, accum_type: DType, shape_id: Int = 0]() -> Index[3]
Returns the appropriate matrix multiply-accumulate (MMA) shape for tensor core operations.
Selects the optimal MMA shape based on the GPU architecture, input data type, accumulation data type, and optional shape identifier. This function handles different configurations for both NVIDIA and AMD GPUs.
Parameters:
- input_type (
DType): The data type of the input matrices (A and B). - accum_type (
DType): The data type used for accumulation (C and D). - shape_id (
Int): Optional identifier to select between multiple valid shapes (default: 0).
Returns:
An IndexList[3] containing the MMA dimensions in the format [M, N, K],
where M×N is the output matrix size and K is the reduction dimension.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!