Support 8-byte dtypes in assignment scatter on Metal

mlx/backend/metal/indexing.cpp

@@ -230,10 +230,12 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Scatter::eval_gpu(const std::vector<array>& 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<array>& inputs, array& out) {
         nidx ? get_type_string(inputs[1].dtype()) : "bool";
     std::string op_type = make_op<Scatter>(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<array>& 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<array>& inputs, array& out) {
 }
 
 void ScatterAxis::eval_gpu(const std::vector<array>& 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<array>& 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<array>& inputs, array& out) {
           large ? "int64_t" : "int",
           op_type,
           uc ? "true" : "false",
-          ic ? "true" : "false");
+          ic ? "true" : "false",
+          out_ptr_type);
     }
     return kernel_source;
   });

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)]],

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<T>* out,
+    device OutT* out,
     const constant int* upd_shape,
     const constant int64_t* upd_strides,
     const constant size_t& upd_ndim,

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<T>* 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)]],

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 <typename T>
+  void atomic_update(device T* out, T val, size_t offset = 0) thread {
+    out[offset] = val;
+  }
 };
 
 template <typename U = bool>

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());
   }
 

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,))