From 8c074744faa4964b283bb0ef1dc75a5032bfc82b Mon Sep 17 00:00:00 2001 From: Erwin Zhang Date: Mon, 17 Aug 2026 07:20:51 -0400 Subject: [PATCH] Support 8-byte dtypes in assignment scatter on Metal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Metal scatter rejected every 8-byte output type, so mx.eye and mx.diag both failed on the GPU for int64, uint64 and complex64. The cause is that all scatters write through mlx_atomic. For a type with no native Metal atomic that falls back to packing several values into one atomic, which cannot represent an 8-byte value at all, so the dispatch refused them outright. Only the reducing scatters need read-modify-write. Assignment is a plain store: scattering duplicate indices under Scatter::None is already a race whose result is whichever thread writes last, and a plain store has exactly those semantics. So the output pointer is now chosen by the op — mlx_atomic for the reducing scatters, plain T for assignment — and the two type gates only reject the reducing case. Kernel names already encode the op, so the two forms cannot collide in the library cache. This fixes eye and diag as a consequence rather than as a special case, and it covers put_along_axis and every other 8-byte assignment scatter with it. test_put_along_axis asserted the old limitation and now verifies the result on both streams. test_scatter_8_byte_types covers the reported cases, checks assignment against the CPU with unique indices, and keeps the invariant that a reducing scatter still raises rather than returning something plausible. Indices are unique in the new test on purpose: assignment with a repeated index is a race by definition, so the GPU and the CPU are both correct and need not agree. --- mlx/backend/metal/indexing.cpp | 25 +++++++-- mlx/backend/metal/jit/indexing.h | 2 +- mlx/backend/metal/kernels/indexing/scatter.h | 5 +- .../metal/kernels/indexing/scatter_axis.h | 5 +- mlx/backend/metal/kernels/reduction/ops.h | 8 +++ mlx/ops.cpp | 7 ++- python/tests/test_ops.py | 56 ++++++++++++++++--- 7 files changed, 85 insertions(+), 23 deletions(-) diff --git a/mlx/backend/metal/indexing.cpp b/mlx/backend/metal/indexing.cpp index 117e402f7b..0a6be0c776 100644 --- a/mlx/backend/metal/indexing.cpp +++ b/mlx/backend/metal/indexing.cpp @@ -230,10 +230,12 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { } void Scatter::eval_gpu(const std::vector& inputs, array& out) { - if (size_of(out.dtype()) == 8 && + // Assignment is dispatched non-atomically below and has no size limit. + if (size_of(out.dtype()) == 8 && reduce_type_ != Scatter::None && !(out.dtype() == complex64 && reduce_type_ == Scatter::Sum)) { std::ostringstream msg; - msg << "[Scatter::eval_gpu] Does not support " << out.dtype(); + msg << "[Scatter::eval_gpu] Does not support " << out.dtype() + << " with a reducing scatter"; throw std::invalid_argument(msg.str()); } @@ -328,6 +330,10 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { nidx ? get_type_string(inputs[1].dtype()) : "bool"; std::string op_type = make_op(reduce_type_, out_type_str); auto [idx_args, idx_arr] = make_index_args(idx_type_str, nidx); + // Kernel names carry the op, so the two pointer forms cannot collide. + std::string out_ptr_type = reduce_type_ == Scatter::None + ? out_type_str + : "mlx_atomic<" + out_type_str + ">"; kernel_source += fmt::format( scatter_kernels, @@ -340,7 +346,8 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { idx_arr, upd_contig, nwork, - large ? "int64_t" : "int"); + large ? "int64_t" : "int", + out_ptr_type); return kernel_source; }); @@ -520,10 +527,11 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { } void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { - if (size_of(out.dtype()) == 8 && + if (size_of(out.dtype()) == 8 && reduce_type_ != ScatterAxis::None && !(out.dtype() == complex64 && reduce_type_ == ScatterAxis::Sum)) { std::ostringstream msg; - msg << "[ScatterAxis::eval_gpu] Does not support " << out.dtype(); + msg << "[ScatterAxis::eval_gpu] Does not support " << out.dtype() + << " with a reducing scatter"; throw std::invalid_argument(msg.str()); } @@ -590,6 +598,10 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { break; } + std::string out_ptr_type = reduce_type_ == ScatterAxis::None + ? out_type_str + : "mlx_atomic<" + out_type_str + ">"; + for (int i = 0; i < 4; ++i) { bool uc = i & 1; bool ic = i & 2; @@ -601,7 +613,8 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { large ? "int64_t" : "int", op_type, uc ? "true" : "false", - ic ? "true" : "false"); + ic ? "true" : "false", + out_ptr_type); } return kernel_source; }); diff --git a/mlx/backend/metal/jit/indexing.h b/mlx/backend/metal/jit/indexing.h index b36093b857..638bf7523b 100644 --- a/mlx/backend/metal/jit/indexing.h +++ b/mlx/backend/metal/jit/indexing.h @@ -36,7 +36,7 @@ constexpr std::string_view gather_kernels = R"( constexpr std::string_view scatter_kernels = R"( [[kernel]] void scatter{0}_{4}_updc_{7}_nwork{8}_{9}( const device {1}* updates [[buffer(1)]], - device mlx_atomic<{1}>* out [[buffer(2)]], + device {10}* out [[buffer(2)]], const constant int* upd_shape [[buffer(3)]], const constant int64_t* upd_strides [[buffer(4)]], const constant size_t& upd_ndim [[buffer(5)]], diff --git a/mlx/backend/metal/kernels/indexing/scatter.h b/mlx/backend/metal/kernels/indexing/scatter.h index 3604f20f61..7e4e8bde7f 100644 --- a/mlx/backend/metal/kernels/indexing/scatter.h +++ b/mlx/backend/metal/kernels/indexing/scatter.h @@ -11,10 +11,11 @@ template < int NIDX, bool UPD_ROW_CONTIG, int NWORK, - typename LocT> + typename LocT, + typename OutT> METAL_FUNC void scatter_impl( const device T* updates, - device mlx_atomic* out, + device OutT* out, const constant int* upd_shape, const constant int64_t* upd_strides, const constant size_t& upd_ndim, diff --git a/mlx/backend/metal/kernels/indexing/scatter_axis.h b/mlx/backend/metal/kernels/indexing/scatter_axis.h index 73fd7ab4a3..206ed53130 100644 --- a/mlx/backend/metal/kernels/indexing/scatter_axis.h +++ b/mlx/backend/metal/kernels/indexing/scatter_axis.h @@ -8,11 +8,12 @@ template < typename LocT, typename Op, bool UpdC, - bool IdxC> + bool IdxC, + typename OutT> [[kernel]] void scatter_axis( const device T* upd [[buffer(0)]], const device IdxT* indices [[buffer(1)]], - device mlx_atomic* out [[buffer(2)]], + device OutT* out [[buffer(2)]], const constant int* shape [[buffer(3)]], const constant int64_t* upd_strides [[buffer(4)]], const constant int64_t* idx_strides [[buffer(5)]], diff --git a/mlx/backend/metal/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h index f875e51c92..bbecfd3382 100644 --- a/mlx/backend/metal/kernels/reduction/ops.h +++ b/mlx/backend/metal/kernels/reduction/ops.h @@ -32,6 +32,14 @@ struct None { thread { mlx_atomic_store_explicit(out, val, offset); } + + // Non-atomic overload, used for types with no atomic representation such as + // the 8-byte ones. Duplicate indices already race under assignment and the + // last write wins, which is what a plain store does. + template + void atomic_update(device T* out, T val, size_t offset = 0) thread { + out[offset] = val; + } }; template diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 1c09e728d0..ed911e9ca1 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3909,12 +3909,13 @@ array scatter( idx = astype(idx, dtype, s); } - // TODO, remove when scatter supports 64-bit outputs + // Only the reducing modes need an atomic, which the GPU lacks at 8 bytes. if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8 && + mode != Scatter::None && !(a.dtype() == complex64 && mode == Scatter::Sum)) { std::ostringstream msg; - msg << "[scatter] GPU scatter does not yet support " << a.dtype() - << " for the input or updates."; + msg << "[scatter] GPU scatter does not support " << a.dtype() + << " for a reducing scatter."; throw std::invalid_argument(msg.str()); } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a09b96669e..b74e32a890 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1538,19 +1538,20 @@ def test_put_along_axis(self): self.assertEqual(b.size, 0) self.assertEqual(b.shape, a.shape) - # 64-bit outputs are not supported by the Metal scatter and should - # raise a clean error rather than failing the Metal JIT build. The CPU - # and CUDA backends handle them fine. + # Assignment needs no atomic, so 8-byte outputs work on every backend. idx = mx.array([[0], [1], [2], [3]]) + expected = np.zeros((4, 8)) + expected[np.arange(4), np.arange(4)] = 1 for dt in (mx.int64, mx.uint64): x = mx.zeros((4, 8), dtype=dt) upd = mx.ones((4, 1), dtype=dt) - if mx.metal.is_available(): - with self.assertRaises(ValueError): - mx.eval(mx.put_along_axis(x, idx, upd, axis=1, stream=mx.gpu)) - out = mx.put_along_axis(x, idx, upd, axis=1, stream=mx.cpu) - self.assertEqual(out.dtype, dt) - mx.eval(out) + for stream in (mx.cpu, mx.gpu): + if stream == mx.gpu and not mx.metal.is_available(): + continue + out = mx.put_along_axis(x, idx, upd, axis=1, stream=stream) + self.assertEqual(out.dtype, dt) + mx.eval(out) + self.assertTrue(np.array_equal(np.array(out), expected)) def test_split(self): a = mx.array([1, 2, 3]) @@ -2931,6 +2932,43 @@ def test_eye(self): # Test with negative k parameter self.assertCmpNumpy([5, 6], mx.eye, np.eye, k=-2) + def test_scatter_8_byte_types(self): + # eye and diag build their result with an assignment scatter, which used + # to reject 8-byte output types on the GPU. + for dt in (mx.int64, mx.uint64, mx.complex64): + for stream in (mx.cpu, mx.gpu): + if stream == mx.gpu and not mx.metal.is_available(): + continue + d = mx.diag(mx.arange(5).astype(dt), stream=stream) + mx.eval(d) + self.assertEqual(d.dtype, dt) + self.assertTrue(np.array_equal(np.array(d), np.diag(np.arange(5)))) + + if dt != mx.complex64: + e = mx.eye(4, dtype=dt, stream=stream) + mx.eval(e) + self.assertEqual(e.dtype, dt) + self.assertTrue(np.array_equal(np.array(e), np.eye(4))) + + # Unique indices: a repeated one is a race by definition, so + # the GPU and the CPU need not agree. + base = mx.zeros((10,), dtype=dt) + idx = mx.array([7, 1, 4]) + upd = mx.array([30, 10, 20]).astype(dt) + out = base + out[idx] = upd + mx.eval(out) + want = np.zeros(10) + want[[7, 1, 4]] = [30, 10, 20] + self.assertTrue(np.array_equal(np.array(out), want)) + + # The reducing scatters still have no atomic to use. + if mx.metal.is_available(): + for dt in (mx.int64, mx.uint64): + x = mx.zeros((8,), dtype=dt, stream=mx.gpu) + with self.assertRaises(ValueError): + mx.eval(x.at[mx.array([0, 1, 1])].add(mx.ones((3,), dtype=dt))) + def test_stack(self): a = mx.ones((2,)) np_a = np.ones((2,))