For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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.