Mojo module

conv_utils

Aliases

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

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

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_conv2d_shape:
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: