For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

conv_transpose

Implements transposed convolution kernels for CPU and GPU backends.

Includes direct register-tiled computation, filter packing, and a cuDNN-backed path for NVIDIA GPUs.

Structs

ConvTransposedPacked: Holds the packed tensors and partition state for one task of a direct transposed convolution.

accumulate_wo_tile: Accumulates one width tile of the transposed convolution into the output buffer.
conv_transpose_naive: Implements the ConvTranspose operator from the MO spec.
conv_transpose_shape: Compute the output shape of a conv-transpose operation, and assert the inputs are compatible.
conv_transposed_cpu: Runs a transposed convolution on the CPU using the direct register-tiled kernel.
conv_transposed_cudnn: Runs a 2D transposed convolution on the GPU using cuDNN backward-data.
conv_transposed_gpu: Runs a transposed convolution on the GPU via the cuDNN backward-data path.
get_num_partitions: Partition the workload in (batch&group, C, F, H) dimensions. HOWO is the combination of HO and WO dimensions. The actual number of tasks are the product of return num_partitions.
get_partition: Computes the partition of the transposed convolution workload assigned to a single task.
pack_filter: Packs the filter from RSFC layout into the FRSCf layout expected by the direct transposed convolution kernel.
pack_filter_shape: Compute the output shape of transposed convolution filter packing.
update_w_tile_2d: Updates one output tile of a 2D transposed convolution using register tiling.
update_w_tile_3d: Updates one output tile of a 3D transposed convolution using register tiling.