diff --git a/mlx/backend/cuda/scan.cu b/mlx/backend/cuda/scan.cu index a7dd2eea25..00e484246e 100644 --- a/mlx/backend/cuda/scan.cu +++ b/mlx/backend/cuda/scan.cu @@ -458,7 +458,10 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); - if (in.flags().contiguous && in.strides()[axis_] != 0) { + // The strided kernel writes out[i * stride + j] for i < shape[axis] and + // j < stride, so the scanned axis has to fit inside the allocation. + if (in.flags().contiguous && in.strides()[axis_] != 0 && + in.shape(axis_) * in.strides()[axis_] <= in.data_size()) { if (in.is_donatable() && in.itemsize() == out.itemsize()) { out.copy_shared_buffer(in); } else { diff --git a/mlx/backend/metal/scan.cpp b/mlx/backend/metal/scan.cpp index 18935191ff..3e4017a888 100644 --- a/mlx/backend/metal/scan.cpp +++ b/mlx/backend/metal/scan.cpp @@ -117,7 +117,10 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); auto in = inputs[0]; - if (in.flags().contiguous && in.strides()[axis_] != 0) { + // The strided kernel writes out[i * stride + j] for i < shape[axis] and + // j < stride, so the scanned axis has to fit inside the allocation. + if (in.flags().contiguous && in.strides()[axis_] != 0 && + in.shape(axis_) * in.strides()[axis_] <= in.data_size()) { if (in.is_donatable() && in.itemsize() == out.itemsize()) { out.copy_shared_buffer(in); } else { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 18556e7d2c..3751039cea 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2500,6 +2500,20 @@ def fn(its): mem4 = mx.get_peak_memory() self.assertEqual(mem2, mem4) + def test_scan_does_not_write_past_the_output(self): + # A size one axis can carry any stride, and the strided kernel bounds + # its writes by that stride rather than by the size of the output. + for n, off in ((64, 63), (512, 400)): + with self.subTest(n=n, off=off): + base = mx.arange(1, n + 1, dtype=mx.float32).reshape(1, n) + canaries = [mx.zeros((16,)) + i for i in range(12)] + mx.eval(base, *canaries) + before = [np.array(x) for x in [base] + canaries] + mx.eval(mx.cumsum(base[:, off:], axis=0)) + after = [np.array(x) for x in [base] + canaries] + for b, a in zip(before, after): + self.assertTrue(np.array_equal(b, a)) + def test_cummax_cummin_nan(self): nan = float("nan") cases = [ diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index 8bef07a616..e33e537032 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -712,3 +712,15 @@ TEST_CASE("test layer norm vjp bias grad race") { } CHECK(worst <= 1e-5); } + +TEST_CASE("test scan output holds the strided kernel writes") { + // A size one axis can carry a padded stride. The strided scan addresses + // out[i * stride + j] for i < shape[axis] and j < stride, so the output has + // to hold shape[axis] * stride elements. This slice overruns by one element, + // the tightest case. + auto a = slice(reshape(arange(64.0, float32), {1, 64}), {0, 1}, {1, 64}); + eval(a); + auto out = cumsum(a, 0); + eval(out); + CHECK_LE(out.shape(0) * out.strides()[0], out.data_size()); +}