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

Convolution kernels for CPU and GPU targets.

Provides direct (register-tiled) convolution, cuDNN-backed convolution (NVIDIA), MIOpen-backed convolution (AMD), and naive GPU reference kernels for 1D, 2D, and 3D convolutions in NHWC/NDHWC layouts.

Structs

• ​CachedCuDNNMetaNHWCFull: Caches cuDNN descriptors, selected forward algorithm, and workspace size for a full NHWC 2D convolution, keyed by input/filter/output shapes and convolution parameters.
• ​CachedMIOpenMeta: Caches MIOpen handle, tensor/filter/convolution descriptors, selected forward algorithm, and workspace size for a convolution of the given rank, keyed by input/filter/output shapes and convolution parameters.
• ​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: Holds a cuDNN handle and the associated input, filter, convolution, and output descriptors for a single device.
• ​Naive2dConvolution: Struct wrapper for naive 2d convolution implementation.

Functions

• ​accumulate_wo_tile_1d: Update one row in the output for a given (c, f) tile.
• ​accumulate_wo_tile_2d: Accumulates one output row tile for a 2D convolution by iterating over the R and S filter-window dimensions and delegating each row to accumulate_wo_tile_1d.
• ​accumulate_wo_tile_3d: Accumulates one output row tile for a 3D convolution by iterating over the Q filter-window depth dimension and delegating each depth slice to accumulate_wo_tile_2d.
• ​check_cudnn_error: Raises an error if a cuDNN call returns a non-success status.
• ​conv1d_update_wo_tile: Updates one micro tile of the 1D convolution output for a given (c, f) tile, accumulating over the S filter window and optionally applying an elementwise epilogue on the last C tile.
• ​conv2d_gpu_naive_nhwc_rscf: Naive GPU kernel for 2D NHWC convolution with RSCF filter layout.
• ​conv2d_update_wo_tile: Updates one micro tile of the 2D convolution output for a given (c, f) tile, accumulating over the R x S filter window and optionally applying an elementwise epilogue on the last C tile.
• ​conv3d_cudnn: Runs a 3D convolution via cuDNN using Nd descriptor APIs, activating the device context before dispatching.
• ​conv3d_gpu_naive_ndhwc_qrscf: Naive GPU kernel for 3D NDHWC convolution with QRSCF filter layout.
• ​conv3d_update_wo_tile: Updates one micro tile of the 3D convolution output for a given (c, f) tile, accumulating over the Q x R x S filter window and optionally applying an elementwise epilogue on the last C tile.
• ​conv_cudnn: Runs a 2D convolution via cuDNN with NHWC input/output and FCRS filter layout, activating the device context before dispatching.
• ​conv_gpu: Dispatches a GPU convolution to the best available backend for the current device and shape, including SM100 structured conv, im2col+matmul, AMD 4-wave, Apple M5 fused, cuDNN, MIOpen, and naive reference kernels, with optional asymmetric padding pre-processing and elementwise epilogue fusion.
• ​conv_miopen: Runs a convolution via MIOpen on AMD GPUs, transposing the filter to FRSC physical layout and dispatching to the cached MIOpen forward convolution path.
• ​conv_nhwc_direct: Runs a direct (register-tiled) NHWC convolution on CPU, bridging TileTensor inputs to LayoutTensors and dispatching to ConvDirectNHWC.run with optional elementwise epilogue fusion.
• ​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_from_fcrs: This packs the filter from FCRS to FRSCf (2D) or FCQRS to FQRSCf (3D).
• ​pack_filter_lt: This packs the filter form RSCF to FRSCf.
• ​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.