Mojo module

utils

Aliases

elementwise_compute_lambda_type

alias elementwise_compute_lambda_type = fn[DType, Int, Int](IndexList[2], SIMD[$0, $1]) capturing -> SIMD[$0, $1]

elementwise_epilogue_type

alias elementwise_epilogue_type = fn[DType, Int, Int](IndexList[2], SIMD[$0, $1]) capturing -> None

Structs

• ​GemmShape: Helper class to unpack gemm dimension and layout.
• ​InnerKernelID:
• ​KernelConfig: Static configuration of the matmul inner kernel.
• ​MicroKernelShape: Record describing the inner kernel shape.
• ​SubMatmulConfig: Static configuration of sub-matrices in parallel matmul.

Functions

• ​apply_epilogue:
• ​calculate_tile_n_k: Helper heuristic function to decide on tile size to partition the matmul given the cache size and desired data layout.
• ​dispatch_get_kernel_type:
• ​get_kernel_config: Utility function to extract matmul configuration parameters for exported Functions. TODO: Add target dependent configuration parameters.
• ​get_kernel_type:
• ​get_matmul_arch_factor:
• ​get_matmul_kernel_shape:
• ​get_matmul_kernel_shape_ARM:
• ​get_matmul_kernel_shape_x86:
• ​get_matmul_num_tasks: Compute the number of tasks for parallel matmul. The max number of tasks is typically the number of threads/cores.
• ​get_matmul_prefetch_b_distance_k:
• ​get_min_task_size:
• ​get_pack_data_size: Utility to compute the number of elements to pack in each tile. Returns: The number of elements to pack.
• ​get_packB_unroll_factor:
• ​get_partitioned_matmul:
• ​get_partitioned_matmul_mojo:
• ​get_partitioned_matmul_mojo_shape:
• ​packA_i8mm:
• ​partition_work:
• ​select_inner_kernel:
• ​use_i8mm_fn:
• ​use_vnni_fn:

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