Mojo module

conv_utils

comptime values

elementwise_epilogue_type

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

elementwise_simd_epilogue_type

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

Structs

• ​ConvAlgorithm:
• ​ConvInfoStatic:
• ​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:
• ​get_conv_shape:
• ​get_conv_tile_shape: Compute the (c, f) tile shape in L2. Assume NHWC layout, the tile shape is (R, S, c_tile, f_tile). R and S are by default fully covered. The heuristic tried to block in C as much as possible. If C is small, it would start to block F.
• ​get_conv_tile_size:
• ​get_direct_conv_micro_kernel_height:
• ​get_direct_conv_micro_kernel_width:
• ​get_micro_kernel_shape:
• ​get_partition:
• ​reorder_padding:

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