Mojo module
conv_utils
Aliases
-
elementwise_epilogue_type = fn[Int](coords: Index[$0], f_size: Int) capturing -> None
: -
elementwise_simd_epilogue_type = fn[DType, Int, Int](Index[$1], SIMD[$0, $2]) 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 insertinglast2nd
andlast
at the end. -
extend_shape
: Extend input shape by insertingfirst
andlast
at both ends. -
get_conv2d_shape
: -
get_conv_num_partitions
: Partition the worload 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!