Mojo module

conv

Structs

ConvDirectNHWC: Implement the outer loops for direct convolution. Collapse N, HO, WO into one dimension n_ho_wo. Tile n_ho_wo, C, and F. The tile factor for C and F are chosen by a heuristic prioritizing C. n_ho_wo is tiled by micro kernel's height.
CuDNNConvMeta:
Naive2dConvolution: Struct wrapper for naive 2d convolution implementation.

accumulate_wo_tile_1d: Update one row in the output for a given (c, f) tile.
accumulate_wo_tile_2d:
accumulate_wo_tile_3d:
check_cudnn_error:
conv1d_update_wo_tile:
conv2d_gpu_naive_nhwc_rscf:
conv2d_update_wo_tile:
conv3d_gpu_naive_ndhwc_qrscf:
conv3d_update_wo_tile:
conv_cudnn:
conv_gpu:
conv_nhwc_direct:
conv_shape: Compute the output shape of a conv operation, and assert the inputs are compatible.
get_cudnn_dtype: Map Mojo DType to cuDNN data type.
pack_conv_filter_shape: Compute the output shape of convolution filter packing.
pack_filter: This packs the filter form RSCF to FRSCf. Use the default micro kernel size for dynamic shapes.
pack_filter_shape: Compute the shape of packed filter. The packed layout is FRSCf. shape_ref should be allocated with size 5 outside this kernel.
pack_filter_shape_impl: Compute the shape of packed filter. The packed layout is FRSCf. shape_ref should be allocated with size 5 outside this kernel.