Module batchnorm 

BatchNorm1d Triton kernels.

Layout: input x is [N, C] row-major. Element x[n, c] lives at flat offset n * C + c.

Parallelism: one CTA per channel. Each CTA iterates over all N batch elements in BLOCK_N-wide tiles. This avoids cross-CTA synchronisation entirely — C channels execute concurrently across SMs.

Training requires two sequential kernel launches separated by a host sync:

1. batch_norm_stats_forward — computes per-channel mean + rstd, updates running stats.
2. batch_norm_normalize_forward — normalises x using the saved stats.

Inference uses a single kernel that reads the frozen running statistics.

Structs

BatchNorm2dNchwBackward

Computes gradients for NCHW BatchNorm2d inference.

BatchNorm2dNchwForwardInference

Normalises NCHW input x using frozen running statistics.

BatchNorm2dNchwInferenceRuntimeOp

RuntimeOp for NCHW BatchNorm2d inference.

BatchNormBackward

Computes gradients for BatchNorm.

BatchNormForwardInference

Normalises input x using the frozen running_mean / running_var.

BatchNormNormalizeForward

Normalises x using the mean and rstd produced by batch_norm_stats_forward.

BatchNormNormalizeRuntimeOp

RuntimeOp for the normalize kernel node in a training BatchNorm graph.

BatchNormStatsForward

Computes per-channel mean and rstd from the current mini-batch, saves them for the normalisation kernel and the backward pass, and updates the running statistics with exponential moving average.

BatchNormStatsRuntimeOp

RuntimeOp for the stats kernel node in a training BatchNorm graph.

Functions

batch_norm_2d_nchw_backward

Computes gradients for NCHW BatchNorm2d inference.

batch_norm_2d_nchw_forward_inference

Normalises NCHW input x using frozen running statistics.

batch_norm_backward

Computes gradients for BatchNorm.

batch_norm_forward_inference

Normalises input x using the frozen running_mean / running_var.

batch_norm_normalize_forward

Normalises x using the mean and rstd produced by batch_norm_stats_forward.

batch_norm_stats_forward

Computes per-channel mean and rstd from the current mini-batch, saves them for the normalisation kernel and the backward pass, and updates the running statistics with exponential moving average.