For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
get_partitioned_matmul
def get_partitioned_matmul[a_type: DType, b_type: DType, c_type: DType, kernel_rows: Int, kernel_cols: Int](m: Int, n: Int, k: Int, task_id: Int, num_tasks: Int) -> SubMatmulConfig
Returns the sub-matmul config for a given task in a parallel matmul.
When the i8mm instruction is available the partition is forced to have an even number of rows, except possibly for the last range.
Parameters:
- a_type (
DType): The dtype of the A tensor. - b_type (
DType): The dtype of the B tensor. - c_type (
DType): The dtype of the C tensor. - kernel_rows (
Int): The static number of rows of the micro kernel. - kernel_cols (
Int): The static number of columns of the micro kernel.
Args:
- m (
Int): The M dimension of the matmul. - n (
Int): The N dimension of the matmul. - k (
Int): The K dimension of the matmul. - task_id (
Int): The index of the task to compute the partition for. - num_tasks (
Int): The total number of tasks sharing the matmul.
Returns:
SubMatmulConfig: The sub-matmul offset and shape for this task.