CPU conv2d direct path writes past the output when padding exceeds a stride's worth of output #4665

The CPU direct conv2d kernel (used for grouped or dilated 2D convs) splits the output into a border region and an interior. The border bound is computed from padding and stride alone and is never clamped to the output size. When the padding is larger than a stride's worth of output on an axis, the border loop runs past the output extent and writes outside the output buffer. With a wide input this is a reliable out-of-bounds write and segfaults.

To reproduce

# repro2.py
# groups=2 forces the direct (non-GEMM) 2D path. iH=1, kernel height 5,
# padding 2, stride 1 -> output height 1, but the height border is ceil(2/1)=2,
# so the first border loop writes a full row past the output. The input is made
# wide so that one stray row lands past the end of the buffer.
import mlx.core as mx
mx.random.seed(0)
W = 500000
x = mx.random.normal((1, 1, W, 2))         # [N, H, W, C]
w = mx.random.normal((2, 5, 1, 1))         # [O, kH, kW, C/groups]
y = mx.conv2d(x, w, stride=(1, 1), padding=(2, 0), groups=2, stream=mx.cpu)
mx.eval(y)
print("done", y.shape)

$ python repro2.py
Segmentation fault: 11

The output shape (1, 1, 500000, 2) is accepted; the fault happens during eval on the CPU. Total allocation is about 12 MB.

A second way in is a very large stride: padding_lo[0] + wt_strides[0] - 1 is an int sum, so a stride near INT_MAX overflows it and the border goes negative, after which the loops run with negative indices. (This one is signed-overflow UB; I did not run it.)

Expected

The op should either compute the correct (small) output or raise, not write outside the output buffer.

Cause

The border bounds are not clamped to the output size and overflow for large strides/kernels:

• height border 1: (padding_lo[0] + wt_strides[0] - 1) / wt_strides[0], no clamp to oH

  mlx/mlx/backend/cpu/conv.cpp

  Line 299 in 3c40e8f

  int oH_border_1 = is_idil_one

• height border 2: max(border1, (iH + padding_lo[0] - wH * wt_dilation[0]) / wt_strides[0]), where wH * wt_dilation[0] can also overflow int

  mlx/mlx/backend/cpu/conv.cpp

  Line 302 in 3c40e8f

  int oH_border_2 = std::max(

• the Case-1 border loop runs oh over [0, oH_border_1) with no check against oH

  mlx/mlx/backend/cpu/conv.cpp

  Line 318 in 3c40e8f

  for (int oh = oH_border_0; oh < oH_border_1; ++oh) {

  (the width borders and Case-2/3 loops have the same shape)

The path is taken for 2D convs with groups > 1 or dilation > 1 (otherwise the op uses the explicit GEMM path):

conv1d's direct path is not affected: it loops only over the real output and checks every tap (ih >= 0 && ih < iH):

Suggested fix

Clamp each border to the output extent, oH_border_k = min(oH, ...) and oW_border_k = min(oW, ...), and compute the border expressions in a wider type (or validate the parameters) so the stride/kernel products cannot overflow int.

Versions

mlx 0.32.3 and mlx-metal 0.32.3 (pip)
main @ 3c40e8f92d0cce0866fc2f436f9b5096cf3b0ea7
macOS 27.0 (build 26A428), Xcode 27.0 (27A266a)
Apple M4 Max, 128 GB