Bound the strided scan to the size of its output

[kapellirohith]

Problem

This is not #206 and it is not bounds checking. Nothing here validates a caller
supplied index. Scan::eval_gpu sizes the output from in.data_size()
(metal/scan.cpp:129, cuda/scan.cu:469) and then hands the kernel
in.strides(), while the kernel bounds its writes by the stride
(kernels/scan.h:445, cuda/scan.cu:236) and so addresses
shape[axis] * stride elements. Those agree for a densely packed array. They
stop agreeing when the scanned axis has size one: check_contiguity requires
strides[i] == prod(shape[i+1:]) only when shape[i] != 1, so a size one axis
may carry any stride and the array is still flagged contiguous and is not
copied. A sliced view gives it exactly that, and MLX's own kernel then writes
past MLX's own allocation.

Minimal reproducer (M3 Pro, macOS 26.x, main @ 3d23f7d):

import mlx.core as mx
import numpy as np

base = mx.arange(1, 65, dtype=mx.float32).reshape(1, 64)
canaries = [mx.zeros((16,)) + float(i) for i in range(12)]
mx.eval(base, *canaries)
snapshot = [np.array(x) for x in [base] + canaries]

mx.eval(mx.cumsum(base[:, 63:], axis=0))   # a (1, 1) view with strides (64, 1)

after = [np.array(x) for x in [base] + canaries]
print(all(np.array_equal(s, a) for s, a in zip(snapshot, after)))
# before: False, in 60 of 60 fresh processes
# after:  True,  in 60 of 60

The scan is asked for one element and writes 64. Nothing raises, and the
tensors it lands on are unrelated to the operation.

Instrumenting the dispatch to print the allocation against the kernel's write
extent, over 120 consecutive dispatches of cumsum(base[:, 200:], axis=0) on a
(1, 256) array:

alloc=56 elems   kernel_writes=256 elems      writes_exceeding_alloc = 120/120

Those element counts are identical for float16, float32, int32, int64 and
complex64. Only the byte overrun scales with itemsize, 224 B allocated against
1024 B written at 4 bytes and 448 B against 2048 B at 8. Whether a given
overrun lands on a tensor you can observe depends on allocator placement and
moves between releases, so the invariant above, not any corruption rate, is the
load bearing evidence here.

What makes it fire and what makes it stop, read off the arrays themselves:

input	axis	allocated	kernel writes	overruns
(1,256) unsliced	0	256	256	no
(1,256)[:, 0:] zero offset	0	256	256	no
(1,256)[:, 200:]	0	56	256	yes
(1,256)[:, 200:] on the size 56 axis	1	56	56	no
(1,64)[:, 63:]	0	1	64	yes
(2,256)[:, 200:] scanned axis size 2	0	312	512	not eligible

It needs a scanned axis of size one carrying a stride larger than the packed
extent. The last row isolates that: with a size two axis the same slice is no
longer contiguous, so the fast path is not taken at all.

Fix

Take the no copy path only when the scanned axis fits inside the allocation,
which is exactly the kernel's maximum write index:

in.shape(axis_) * in.strides()[axis_] <= in.data_size()

Inputs that fail it fall to contiguous_copy_gpu, already the else branch
directly below. For a row major array the condition holds by construction,
since scanning axis i of shape (d0 ... dn) gives
di * prod(d[i+1:]) <= prod(d), so nothing that works today changes path. The
same dispatch, allocation and kernel bound are present in cuda/scan.cu, so it
changes there too.

Introduced by c423074 "redesign for faster cpu/gpu synch (#1869)"
(2025-03-06), which is where the dispatch began sizing the output from
data_size() while passing strides(). No existing scan test scans a size one
axis, and none asserts that an operation left memory it does not own alone.

Testing

Two tests. The Python one allocates canaries and asserts none of them, nor the
parent, moved. The C++ one asserts
out.shape(axis) * out.strides()[axis] <= out.data_size() on a slice that
overruns by exactly one element; it is there because the invariant is not
observable from Python, where strides are not exposed, and because it is the
only deterministic detector of a one element overrun.

All figures on main @ 3d23f7d, M3 Pro, macOS 26.x.

check	main behaviour	with this change
Python test, standalone, DEVICE=gpu	100/100 red	200/200 pass
Python test, in full test_ops, DEVICE=gpu	30/30 red	10/10 pass
Python test, standalone, DEVICE=cpu	n/a, CPU path	200/200 pass
Python test, in full test_ops, DEVICE=cpu	n/a, CPU path	10/10 pass
C++ invariant test, DEVICE=gpu	FAILURE	SUCCESS
C++ invariant test, DEVICE=cpu	n/a, CPU path	SUCCESS
parent clobber, (1,10)[:, 3:], 120 dispatches	60/120	0/120
dtype and op matrix, 16 rows including control	8 byte dtypes 40/40	0/40 every row
scan outputs, 5 shapes x every axis x 4 ops x reverse x inclusive	n/a	224/224 bit identical
C++ suite, both devices	n/a	277/277
Python suite, DEVICE=gpu	n/a	840 tests, exit 0
MLX_METAL_JIT=ON, new test	n/a	20/20
pre-commit run --all-files	n/a	5 hooks clean

Mutation testing of the new condition, each mutant built and run against both
tests:

mutant	C++ invariant	Python canary	classification
drop the guard, that is main	caught	caught 10/10	the defect
compare against nbytes()	caught	survives	permits an overrun
allow one element of slack	caught	survives	permits an overrun
< instead of <=	survives	survives	safe, not equivalent
size() instead of data_size()	survives	survives	equivalent
drop the shape[axis] factor	survives	survives	equivalent

Both mutants that still permit an overrun are caught only by the C++ test. The
survivors were classified by enumerating the inputs reachable through reshape,
slice, transpose, broadcast_to and as_strided: < never permits an overrun
but forces needless copies, and the other two never diverge from the shipped
condition on any reachable input.

tests/CMakeLists.txt gates gpu_tests.cpp on
if(MLX_BUILD_METAL OR MLX_BUILD_CUDA), so the invariant test compiles and runs
on the CUDA CI legs as well as the Metal ones. Metal is covered by both tests,
CUDA by the invariant test.

Perf is unchanged for inputs that already worked: they satisfy the new
condition, so they take the same path and the same kernel. Toggling the new term
at runtime, so that both configurations run alternately in one process, gives
paired median differences of +2.8, +0.9 and -0.6 percent over 12 pairs on three
contiguous scan shapes, mixed in sign and inside this machine's run to run
spread. (1, 4096)[:, 1:] goes from 53.2 ms to 0.14 ms, since on main the
strided kernel is dispatched with stride = 4096 for a scan of length one.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)