Module ml::gpu_cudnn

#include <std/ml/gpu_cudnn.h> source

cuDNN convolution, pooling and U-Net on CUDA.

The GPU counterpart of cnn.h / unet.h, layered on gpu_cuda.h: tensors are CudaBufs in NCHW order (cnn.h's CHW FeatureMap with a leading batch axis), and every op is queued on gpu_cuda's stream, so it orders with the CudaBuf ops there and can be recorded into a CudaGraph.

Convolutions run through cuDNN. The first call for each shape (and for plain vs. +relu) builds the fused graph conv -> + bias (-> relu) with cuDNN's graph API, times the engine configs its heuristics suggest and the legacy API's best algorithm (cudnnFind…) on the call's buffers, and keeps the fastest, with its workspace, in a plan cache — so make that first call outside graph capture. Later calls cost a few microseconds of host time. fp32 convolutions run in full fp32 unless cuda_set_tf32(1) is in effect, which also lets cuDNN use TF32 tensor cores (same switch as the cuBLAS GEMMs). Link with: -lcudnn plus gpu_cuda's -lcuda -lcublas -lcublasLt -lnvrtc (cuDNN 9).

Structs

CudaConv2D	A Conv2D (cnn.h) whose weight [outC, inC, kH, kW] and bias [outC] live on the GPU. weight.dev == 0 if the upload failed.
CudaHalfConv2D	The same convolution on 16-bit NCHW tensors (gpu_cuda.h's CudaHalfBuf): cuDNN's half / bfloat16 convolutions (fp32 compute, tensor cores where cuDNN picks them), rounding once per output; the bias add and ReLU are one fused 16-bit kernel (cuda_half_bias). Weights are rounded from a Conv2D's fp32 ones on upload. x, y and the layer must share one dtype.
CudaUNet	Weights plus every intermediate activation, sized once for a fixed batch x height x width, so a forward pass allocates nothing.

Functions

cuda_conv_out_dim	Output side length of a convolution or pooling window: (dim + 2*padding - kernel)/stride + 1.
cuda_maxpool2d	kernel x kernel window, given stride, no padding (FeatureMap::maxpool2d / FeatureMap::avgpool2d). x: [batch, channels, height, width].
cuda_avgpool2d
cuda_concat_channels	out[b] = (a[b], c[b]) along channels, for each of 'batch' images whose planes are 'plane' = height*width elements (concat_channels).
conv2d_forward_cuda	Same signatures and results as the CPU functions; each falls back to the CPU implementation if there is no usable GPU. Every call uploads the weights — inside a loop, upload once and use the queued ops above.
unet_forward_cuda