im2col Expansion, FFT Convolution, and Winograd Transforms

A convolution layer is awkward for hardware built around dense matrix multiplication. The inner loop reads two operands on a grid with overlapping windows, not two contiguous rows. Three transformations bridge the gap, and each pays for the bridge differently: im2col trades memory for contiguous arithmetic, FFT trades precision and boundary handling for an asymptotically cheaper operator, and Winograd trades numerical range for fewer multiplications. Choosing between them requires knowing what each actually removes.

The accounting that governs all three is the arithmetic intensity of the layer, not its operation count. A convolution with $C_{in} C_{out} K^2 H W$ multiply-adds and $C_{in} H W + C_{in} C_{out} K^2$ operands has a high ratio of arithmetic to bytes because each input value is reused $C_{out} K^2$ times and each weight is reused $HW$ times. A GEMM with the same arithmetic but no reuse beyond a two-operand tile becomes bandwidth-bound the moment the operands do not fit in registers and cache. The transformations differ in how much of that reuse they preserve while making the data contiguous.

im2col: replicating patches so the GEMM can be naive

The straightforward route materialises every patch. For an output of extent $H_{out} \times W_{out}$, the input is read as a matrix of shape $(C_{in} K_h K_w) \times (H_{out} W_{out})$, where column $(i,j)$ holds the flattened patch that output cell $(i,j)$ sees. The kernel becomes a row of shape $C_{out} \times (C_{in} K_h K_w)$, and the layer is a single matrix product.

import numpy as np

def im2col(x, kh, kw, stride=1, pad=0):
    c, h, w = x.shape
    xp = np.pad(x, ((0, 0), (pad, pad), (pad, pad)))
    ho = (h + 2 * pad - kh) // stride + 1
    wo = (w + 2 * pad - kw) // stride + 1
    cols = np.empty((c * kh * kw, ho * wo), dtype=x.dtype)
    idx = 0
    for ci in range(c):
        for i in range(kh):
            for j in range(kw):
                patch = xp[ci, i:i + stride * ho:stride, j:j + stride * wo:stride]
                cols[idx] = patch.ravel()          # each tap becomes one row
                idx += 1
    return cols

x = np.random.default_rng(0).normal(size=(3, 32, 32))
cols = im2col(x, 3, 3, stride=1, pad=1)
print(cols.shape)     # (27, 1024)

The expanded matrix holds $C_{in} K_h K_w H_{out} W_{out}$ elements. For a 3×3 kernel that is nine times the input, and the expansion is not free: the same input value appears in up to nine columns, so the layer writes nine copies of the input and then reads them back through the GEMM. At $C_{in}=256$, $H=W=56$, $K=3$, the buffer holds about 7.2 million elements, near 29 MB in 32-bit precision for a single sample, far larger than the level-2 cache of a typical server core, so the expansion itself becomes a bandwidth-bound pass that runs at memory speed rather than at arithmetic speed.

The alternative is implicit expansion: an indexing scheme inside the GEMM’s load path that computes the address of the required element without materialising it. The arithmetic stays the same and the expansion writes disappear, at the cost of more complex address generation and less regular access. A middle path tiles the output space so that each tile’s expanded patch fits in cache; only the tile is materialised.

The structural weakness of the im2col formulation is the zero multiplier. When the input extent is not a multiple of the stride, the last output row is a partial window; padding adds further zeros; and the expanded matrix contains the padding explicitly. Every zero multiplies and is added, which is why a layer whose nominal shape implies $2C_{in}C_{out}K^2 HW$ floating-point operations does that many regardless of the padding fraction. On a 3×3 kernel with one block of padding at a 56-pixel boundary the fraction is small; at $224 \times 224$ with a large kernel, or with a dilation that leaves gaps, it is not.

Tensor cores impose a second layout

A GEMM on tensor core hardware requires the operands in the hardware’s fragment layout, which means the matrix product is neither the end nor the natural form of the computation. Kernels split the reduction over $C_{in} K_h K_w$ into chunks matching the fragment depth, accumulate in registers at higher precision than the input dtype, and write the accumulated tile directly to the output’s channel-major layout. At that level the im2col matrix never exists in any physical form; the mapping from convolution coordinates to fragment indices is what the kernel’s code generator emits. The arithmetic count is still that of im2col, and the padding zeros still multiply.

FFT convolution moves the cost from $K^2$ to $\log$

The convolution theorem turns a spatial convolution into a pointwise product of spectra. For a kernel large enough that $K^2$ is significant compared to $\log H$, the transform route wins: the cost becomes $O(HW \log HW)$ per channel pair instead of $O(HW K^2)$.

The details are where implementations diverge from the textbook. The product of two spectra is a circular convolution on the transform length, so an $N$-point transform computes the sum over indices modulo $N$; to get a linear convolution of an $N$-pixel signal with a $K$-tap kernel the transform must be padded to at least $N + K - 1$, which inflates every buffer by a factor approaching two for large $K$. Everything in that region is zero; the zeros cost nothing to multiply, but they must be written and transformed, and for a kernel that overlaps most of the signal they outweigh the content.

Two-dimensional transforms exploit the separability of the kernel: a 2-D transform is applied as transforms along rows and then along columns. That means the intermediate is in a transposed layout, so either the pass gathers strided elements or the buffer is transposed between passes, and the transpose costs bandwidth comparable to the transforms themselves.

Batching is the third consideration. The transform can be amortised across the $C_{out}$ output channels and across the batch: transforming the input patch once and multiplying against $C_{out}$ spectra is the reuse that makes the method competitive. Without that reuse, the transform overhead per output value dominates for any kernel that a small-tap spatial method can compute cheaply.

The practical boundary is that FFT is used for large kernels, roughly 5×5 and above at moderate channel counts, and for filters that are already defined in the frequency domain. Below that, the transform and its padding cost more than the taps they replace. The radial filters used in some video and audio models are naturally expressed in the frequency domain and are the clearest case where the transform is not a trick but the natural coordinate system; the same reasoning underlies the band splitting in which a separable stage handles low frequencies spatially and high frequencies pointwise.

One subtlety with consequences: the convolution theorem applies to true convolution, not cross-correlation. To match a spatial implementation that uses correlation, one of the spectra must be conjugated, or the kernel reversed, before the product. Omission of that step produces a mirror-image result whose effect is invisible while weights are random and reversed once learned, which is itself a reason to state the convention explicitly rather than rely on symmetry.

Winograd: fewer multiplications, more additions

Winograd’s minimal filtering algorithms compute an $m$-output, $r$-tap FIR filter with $m + r - 1$ multiplications instead of $mr$. For the $F(2,3)$ case, two outputs from a 3-tap kernel over a window of four inputs, the direct method needs 6 multiplications and the transform needs 4, a saving of 1.5×. $F(4,3)$ needs 6 multiplications for 4 outputs versus 12, a factor of 2. In two dimensions the savings compound: the 2-D transform of $F(2\times2, 3\times3)$ needs 16 multiplications for 4 outputs, against 36 in the direct method, a factor of 2.25.

The mechanism is a change of basis. The input tile is transformed by a matrix whose entries are small integers, the kernel is transformed by another, the products are formed elementwise, and the result is transformed back. The transforms are additions and shifts rather than multiplications, so they are cheap on hardware where multiplies and adds have similar cost but the multiply throughput is the bottleneck.

The entries of the transform matrices are not all small. The kernel transform of $F(4,3)$ carries factors of $1/6$ and $1/24$, while the input and inverse transforms carry integer coefficients as large as 5 and 8. Elementwise products of transformed operands therefore span a wider dynamic range than the original values, and the accumulated error is larger than for the direct method. The consequence is dtype-specific: in 16-bit floating point the error at full range is large enough to matter, so implementations either restrict themselves to $F(2,3)$ at that precision or accumulate in 32-bit and clip. The trade is the whole story of Winograd: the flop count falls, the numeric margin shrinks, and the transform buffers raise the working-set size.

import numpy as np

# F(2,3): 4 multiplications instead of 6
G = np.array([[1, 0, 0], [0.5, 0.5, 0.5], [0.5, -0.5, 0.5], [0, 0, 1]])
Bt = np.array([[1, 0, -1, 0], [0, 1, 1, 0], [0, -1, 1, 0], [0, 1, 0, -1]])
At = np.array([[1, 1, 1, 0], [0, 1, -1, -1]])

d = np.array([1.0, 2.0, 3.0, 4.0])       # input tile
g = np.array([2.0, 0.0, -1.0])           # filter
direct = np.array([d[:3] @ g, d[1:] @ g])
via = At @ ((Bt @ d) * (G @ g))
assert np.allclose(direct, via)

Where each transform loses

The three methods fail on different axes. im2col fails on memory: the expansion multiplies the input footprint by $K^2$ and turns a compute-bound layer into a copy-bound one whenever the expanded buffer does not fit the cache level that the GEMM’s blocking expects. FFT fails on constants and boundaries: the padding buffer doubles the footprint, the transpose pass costs bandwidth, and small kernels never amortise the transform. Winograd fails on precision and locality: the inverse transform is applied to a tile of products, so the accumulation order is fixed by the tile shape rather than chosen by the compiler, and the transform buffers add to the working set.

There is a fourth consideration that none of the three addresses: sparsity and structure in the kernel. A depthwise or grouped layer has a block-diagonal weight matrix in the im2col view, so the GEMM has large zero blocks and achieves a fraction of its nominal throughput; the formulations in depthwise convolutions and spectral efficiency avoid im2col for exactly this reason, and the output-size arithmetic that decides how much of the expansion is padding zeros comes from the index bounds derived in cross-correlation and convolution in the forward pass. A kernel with dilation $D$ has its taps at spacing $D$, so the im2col rows are strided reads from the same rows of the input; the expansion is the same size but the access pattern changes, and direct tiled kernels are usually faster than the expanded form.

The decision rule that survives contact with hardware is therefore not a formula. It is: measure the achieved throughput of the candidate kernel against the layer’s arithmetic intensity, and choose the method whose bottleneck is not the one the layer is already up against. A layer with many channels and a 3×3 kernel is compute-bound after expansion, so im2col with a well-tuned GEMM is close to optimal. A layer with a large kernel and moderate channels is bandwidth-bound, so the transform route’s lower flop count translates into a real gain. A layer with few channels and small spatial extent is launch-bound, and the choice among the three is dominated by fixed overheads rather than by any of this arithmetic.

What is not negotiable is the numerical contract. Any of the three must produce results indistinguishable from the direct sum within the tolerance the training procedure assumed, and the transform methods change the accumulation order, the intermediate magnitude, and in the FFT case the boundary behaviour. A layer whose output is compared to a reference implementation bit-for-bit will differ under every one of them, and the differences are largest exactly where the padding and the tile boundaries meet.