Detect int32 shape-product overflow at MLX compute-shape boundaries

[qflen]

Summary

Fixes #3327.

Shapes whose per-dim values fit in int32 but whose product exceeds 2^31 silently wrapped at five sites: reshape(-1), flatten, take, the Metal conv dispatcher's out.size() / O cast, and per-thread output offsets inside the Metal conv kernels. All five now either succeed correctly or raise std::overflow_error with a clear [op] Shape dimension ... message, extending the diagnostic style #3425 established at the Python binding to the internal C++ boundaries.

ShapeElem stays int32_t. No public API change. The flatten case from #2681 is also resolved.

Root cause

• Flatten::output_shape accumulated flat_size *= ... as int32_t via auto.
• Reshape::output_shape and unflatten truncated input.size() / size from size_t to ShapeElem.
• indices_or_default used std::multiplies<int>.
• Metal conv.cpp: int implicit_M = out.size() / O truncated; int implicit_M = N * oS[0] * oS[1] wrapped. Wrapped neg dim fed nbytes() which sign-extends to ~18 EB.
• Metal steel_conv{,_3d,_general}.h: per-thread output offset c_row * (N * groups) + c_col wrapped int32 when M < 2^31 but M * O > 2^31.
• CUDA conv/{gemm_conv,gemm_grouped_conv}.cu: same int mat_M = out.size() / O truncation as the Metal sites.

Change

• New check_shape_dim(int64_t, std::string_view op) in mlx/utils.h.
• Flatten, Reshape, unflatten, indices_or_default widen accumulators to int64_t and narrow through the helper. Backend-agnostic.
• Four Metal and two CUDA conv dispatchers guard implicit_M / mat_M. Widen Metal inp_large / out_large to int64_t.
• Three Metal steel conv kernel headers promote per-thread output offset to size_t. This is the substance of the prior attempts Fix int32 overflow in Metal conv_general output offset for large tensors #3294 and Fix conv_general output offset overflow in Metal writeback #3320, now exercised by an end-to-end test.

Tests (tests/gpu_tests.cpp)

• test gpu int32 shape overflow errors covers flatten, reshape(-1), take, eval. Eval branch guarded by max_buffer_length.
• test gpu conv2d large output offset: output 2.15 G fp16 elements, varying batch values, allclose vs CPU. Fails on y[-1] without the steel_conv_general.h patch (verified by stash/restore).

Validation

Apple M5 32GB, Release + MLX_METAL_JIT=OFF.

• CPU: test_ops.py (139), test_array.py (73), test_conv.py (18), test_compile.py (54), and the reproducer on CPU stream all pass.
• GPU / Metal: 248/248. Metal API validation layer: 0 errors.
• GPU / CUDA: deferred to CI; the change is a two-file mirror of the Metal dispatcher fix with no kernel-side component.

Performance

One host-side int64 compare per kernel dispatch. Trimmed-mean of 500 iters, per-call µs:

Workload	Baseline	Patched
conv2d (32,64,64,64) -> (..,128), Metal	10080	9204
conv2d_general (16,64,64,1) -> (..,17), Metal	1164	717
take((8192,64), 5 idx), Metal	326	205
reshape((4096,1024) -> -1), Metal	66	77
flatten((4096,1024)), Metal	83	74
conv2d (4,64,64,32) -> (..,64), CPU	3623	3348
reshape(-1), CPU	27	22
flatten, CPU	28	23

Patched is faster in 7/8 cells; deltas are within sub-ms jitter.

Scope

I noticed that same-class fixes exist at adjacent sites the reproducer does not reach:
mlx/backend/metal/matmul.cpp gather-MM dispatchers tile / repeat / kron / concatenate / pad, CPU mlx/backend/cpu/conv.cpp, and dilate_size / conv_out_axis_size.

Left for follow-up to keep this PR focused, but I can also work on them here.

[zcbenz]

Nice fix, thanks!

[zcbenz]

Allocating such a large array would drag down the test speed a lot, I prefer just removing this test since it is not a serious issue.

[zcbenz]

Sorry I meant deleting the tests that actually allocates the 4GB array (including the "test gpu conv2d large output offset"), the tests that check overflow are nice to keep.

[qflen]

Oh my bad, I'll rewrite.

[ltctech]

Standalone mx.conv3d reproducer for the steel_conv_3d.h output-offset overflow your PR fixes.

Confirmed on MLX 0.31.2, macOS Tahoe 26.4.1, Apple M1 Max 64 GB.

import time
import mlx.core as mx

mx.set_default_device(mx.gpu)
mx.set_cache_limit(1024 * 1024 * 1024)
print("MLX:", mx.__version__)

# Output shape: (1, 457, 72, 128, 512)
# M = 1 * 457 * 72 * 128 = 4,211,712
# N = 512
# M * N = 2,156,396,544  >  INT32_MAX (2,147,483,647)
#
# C=64 satisfies C_per_group % 16 == 0 (implicit path),
# and keeps input/weight memory reasonable.

x = mx.ones((1, 457, 72, 128, 64), dtype=mx.float16)
w = mx.ones((512, 3, 3, 3, 64), dtype=mx.float16)
t0 = time.perf_counter()
out = mx.conv3d(x, w, padding=1)
mx.eval(out)
print(f"conv eval: {time.perf_counter() - t0:.2f}s")

for label, d, expected in [
    ("before boundary", 454, 1728),
    ("after boundary",  455, 1728),
    ("last depth",      456, 1152),
]:
    vals = out[0, d, 36, 64, :8]
    mx.eval(vals)
    print(f"{label:16s} expected={expected:4d} actual={vals}")

print(f"peak GB: {mx.get_peak_memory() / 1e9:.2f}")

On 0.31.2 (before your fix):

before boundary  expected=1728 actual=array([1728, 1728, 1728, ..., 1728, 1728, 1728], dtype=float16)
after boundary   expected=1728 actual=array([0, 0, 0, ..., 0, 0, 0], dtype=float16)
last depth       expected=1152 actual=array([0, 0, 0, ..., 0, 0, 0], dtype=float16)

The overflow boundary is at 2^31 / (72 * 128 * 512) = depth index 455.11, which is exactly where the output goes to zero. The affected region is unwritten (no garbage values, just zeros), consistent with the stores targeting an incorrect address and missing the intended output region entirely.

This may be useful as a regression test for the steel_conv_3d.h path specifically, since it exercises mx.conv3d directly rather than conv_general.

[qflen]

Thanks for the repro and pinpointing the 2^31 / (72 * 128 * 512) boundary.

The size_t cast on steel_conv_3d.h:97 covers your case since per-thread offsets store_result_safe adds on top are tile-bounded, so they don't re-introduce overflow.

With @zcbenz's earlier feedback, the >=4 GB-allocating tests were dropped, and this one would also need ~4.3 GB to trigger (M * N > INT32_MAX). The host-side overflow checks in gpu_tests.cpp stay.