Fix compiled negative-strided slice updates

mlx/backend/metal/compiled.cpp

@@ -12,6 +12,20 @@
 #include "mlx/utils.h"
 
 namespace mlx::core {
+namespace {
+
+bool has_negative_strides(const std::vector<Strides>& strides) {
+  for (const auto& stride : strides) {
+    for (auto value : stride) {
+      if (value < 0) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+} // namespace
 
 inline void build_kernel(
     std::string& os,
@@ -139,8 +153,10 @@ inline void build_kernel(
     os += fmt::format("  {0} index_{1} = ", idx_type, xname);
     if (ndim == 1) {
       int offset = i * ndim;
-      os +=
-          fmt::format("elem_to_loc_1<uint>(pos.x, in_strides[{0}]);\n", offset);
+      os += fmt::format(
+          "elem_to_loc_1<{0}>(pos.x, in_strides[{1}]);\n",
+          idx_type,
+          offset);
     } else if (ndim == 2) {
       int offset = i * ndim;
       os += fmt::format(
@@ -316,20 +332,18 @@ void Compiled::eval_gpu(
           /* dynamic_dims = */ false,
           /* use_big_index = */ false,
           /* work_per_thread = */ i > 3 ? 2 : 1);
-      if (i > 1) {
-        build_kernel(
-            kernel,
-            kernel_lib_ + "_strided_" + std::to_string(i) + "_large",
-            inputs_,
-            outputs_,
-            tape_,
-            is_constant_,
-            /* contiguous = */ false,
-            /* ndim = */ i,
-            /* dynamic_dims = */ false,
-            /* use_big_index = */ true,
-            /* work_per_thread = */ i > 3 ? 4 : 1);
-      }
+      build_kernel(
+          kernel,
+          kernel_lib_ + "_strided_" + std::to_string(i) + "_large",
+          inputs_,
+          outputs_,
+          tape_,
+          is_constant_,
+          /* contiguous = */ false,
+          /* ndim = */ i,
+          /* dynamic_dims = */ false,
+          /* use_big_index = */ true,
+          /* work_per_thread = */ i > 3 ? 4 : 1);
     }
     build_kernel(
         kernel,
@@ -364,7 +378,8 @@ void Compiled::eval_gpu(
       compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
 
   // Whether to use large index.
-  bool large = compiled_use_large_index(inputs, outputs, contiguous);
+  bool large = compiled_use_large_index(inputs, outputs, contiguous) ||
+      (!contiguous && has_negative_strides(strides));
 
   // Get the kernel from the lib
   int ndim = shape.size();

python/tests/test_compile.py

@@ -5,6 +5,7 @@
 import io
 import math
 import threading
+import unittest
 from functools import partial, wraps
 from io import StringIO
 
@@ -71,6 +72,25 @@ def test_compile_nonfinite_constants(self):
             self.assertEqual(out[0].item(), 1.0)
             self.assertEqual(out[1].item(), float("-inf"))
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
+    def test_compile_negative_strided_slice_update_expr(self):
+        def add_update(x):
+            out = mx.zeros_like(x)
+            out[::-1] += 2.0 * x[::-1]
+            return out
+
+        def set_update(x):
+            out = mx.zeros_like(x)
+            out[::-1] = 2.0 * x[::-1]
+            return out
+
+        mx.set_default_device(mx.gpu)
+        x = mx.arange(6, dtype=mx.float32)
+        for fn in (add_update, set_update):
+            expected = fn(x)
+            actual = mx.compile(fn)(x)
+            self.assertEqualArray(actual, expected)
+
     def test_compile_tuple_output_in_thread(self):
         @mx.compile
         def fun(x):