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 module

conv_utils

Utilities for direct convolution kernels on CPU.

Provides convolution shape description, L2 tiling heuristics, micro-kernel shape selection, and work partitioning helpers used by the direct and im2col convolution implementations.

comptime values

elementwise_epilogue_type

comptime elementwise_epilogue_type = def[rank: Int](coords: IndexList[rank], f_size: Int) capturing thin -> None

elementwise_simd_epilogue_type

comptime elementwise_simd_epilogue_type = def[dtype: DType, rank: Int, width: SIMDLength, alignment: Int = Int(1)](IndexList[rank], SIMD[dtype, width]) capturing thin -> None

Structs

• ​ConvAlgorithm: Tags the convolution algorithm selected for a given operation.
• ​ConvInfoStatic: Holds statically known convolution attributes (padding, stride, dilation, groups).
• ​ConvPartition: Work range for a partition.
• ​ConvShape: A shape struct describing the convolution dimensions.

Functions

• ​align_down_residual: Returns the remainder after aligning down value to alignment.
• ​append_shape: Append input shape by inserting last2nd and last at the end.
• ​extend_shape: Extend input shape by inserting first and last at both ends.
• ​get_conv_num_partitions: Partition the workload in (batch, C, F, HOWO) dimensions. HOWO is the combination of HO and WO dimensions. The actual number of tasks are the product of return num_partitions.
• ​get_conv_num_tasks: Returns the number of tasks to partition the convolution into.
• ​get_conv_shape: Builds a ConvShape from the output, input, and filter tile tensors.
• ​get_conv_tile_shape: Compute the (c, f) tile shape in L2.
• ​get_conv_tile_size: Returns the convolution L2 cache tile size in elements for the target.
• ​get_direct_conv_micro_kernel_height: Returns the micro-kernel height for direct convolution on the target CPU.
• ​get_direct_conv_micro_kernel_width: Returns the micro-kernel width for direct convolution on the target CPU.
• ​get_micro_kernel_shape: Selects the (height, width) micro-kernel tile shape for the target.
• ​get_partition: Computes the ConvPartition work range for a given task id.
• ​reorder_padding: Reorders padding entries into (lower, upper) pairs per dimension.

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