Fix compiled kernel correctness for negative-strided inputs

mlx/backend/common/compiled.cpp

@@ -82,14 +82,24 @@ std::string get_type_string(Dtype d) {
   }
 }
 
-bool compiled_check_contiguity(
+bool has_negative_strides(const array& x) {
+  for (auto s : x.strides()) {
+    if (s < 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
+std::pair<bool, bool> compiled_check_contiguity(
     const std::vector<array>& inputs,
     const Shape& shape) {
   bool contiguous = true;
   bool all_contig = true;
   bool all_row_contig = true;
   bool all_col_contig = true;
   int non_scalar_inputs = 0;
+  bool negative_strides = false;
   for (const auto& x : inputs) {
     if (is_scalar(x)) {
       continue;
@@ -99,6 +109,7 @@ bool compiled_check_contiguity(
     all_contig &= (x.flags().contiguous && shape_eq);
     all_row_contig &= (x.flags().row_contiguous && shape_eq);
     all_col_contig &= (x.flags().col_contiguous && shape_eq);
+    negative_strides |= has_negative_strides(x);
   }
   if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) {
     contiguous = false;
@@ -107,7 +118,7 @@ bool compiled_check_contiguity(
   } else if (non_scalar_inputs == 0 && !shape.empty()) {
     contiguous = false;
   }
-  return contiguous;
+  return {contiguous, negative_strides};
 }
 
 void compiled_allocate_outputs(
@@ -170,14 +181,16 @@ void compiled_allocate_outputs(
   }
 }
 
-std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
+std::tuple<bool, bool, Shape, std::vector<Strides>>
+compiled_collapse_contiguous_dims(
     const std::vector<array>& inputs,
     const array& out,
     const std::function<bool(size_t)>& is_constant) {
   const Shape& shape = out.shape();
-  bool contiguous = compiled_check_contiguity(inputs, shape);
-  if (contiguous) {
-    return {true, shape, {}};
+  auto [contiguous, negative_strides] =
+      compiled_check_contiguity(inputs, shape);
+  if (contiguous && !negative_strides) {
+    return {true, false, shape, {}};
   }
 
   std::vector<Strides> strides_vec{out.strides()};
@@ -218,7 +231,11 @@ std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
   }
 
   auto tup = collapse_contiguous_dims(shape, strides_vec, INT32_MAX);
-  return {false, std::move(std::get<0>(tup)), std::move(std::get<1>(tup))};
+  return {
+      false,
+      negative_strides,
+      std::move(std::get<0>(tup)),
+      std::move(std::get<1>(tup))};
 }
 
 bool compiled_use_large_index(

mlx/backend/common/compiled.h

@@ -66,8 +66,9 @@ inline bool is_scalar(const array& x) {
   return x.ndim() == 0;
 }
 
-// Check if we can use a contiguous operation given inputs and the output shape
-bool compiled_check_contiguity(
+// Check if we can use a contiguous operation given inputs and the output shape.
+// Also returns if any input has negative strides.
+std::pair<bool, bool> compiled_check_contiguity(
     const std::vector<array>& inputs,
     const Shape& shape);
 
@@ -81,7 +82,8 @@ void compiled_allocate_outputs(
         allocator::malloc);
 
 // Collapse contiguous dims ignoring scalars and constants.
-std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
+std::tuple<bool, bool, Shape, std::vector<Strides>>
+compiled_collapse_contiguous_dims(
     const std::vector<array>& inputs,
     const array& out,
     const std::function<bool(size_t)>& is_constant);

mlx/backend/cpu/compiled.cpp

@@ -298,7 +298,7 @@ void Compiled::eval_cpu(
 
   // Collapse contiguous dims to route to a faster kernel if possible. Also
   // handle all broadcasting.
-  auto [contiguous, shape, strides] =
+  auto [contiguous, negative_strides, shape, strides] =
       compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
 
   // Collect function input arguments.

mlx/backend/cuda/compiled.cpp

@@ -285,11 +285,12 @@ void Compiled::eval_gpu(
 
   // Collapse contiguous dims to route to a faster kernel if possible. Also
   // handle all broadcasting.
-  auto [contiguous, shape, strides_vec] =
+  auto [contiguous, negative_strides, shape, strides_vec] =
       compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
 
-  // Whether to use large index.
-  bool large = compiled_use_large_index(inputs, outputs, contiguous);
+  // Whether to use large index (also true for negative strides).
+  bool large =
+      negative_strides || compiled_use_large_index(inputs, outputs, contiguous);
 
   cu::KernelArgs args;
   // Put inputs.

mlx/backend/metal/compiled.cpp

@@ -139,8 +139,8 @@ 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 +316,20 @@ 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);
-      }
+      // Generate int64_t index variant for all ndim, including ndim=1.
+      // Negative strides force large mode even for small arrays.
+      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,
@@ -360,11 +360,12 @@ void Compiled::eval_gpu(
 
   // Collapse contiguous dims to route to a faster kernel if possible. Also
   // handle all broadcasting.
-  auto [contiguous, shape, strides] =
+  auto [contiguous, negative_strides, shape, strides] =
       compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
 
-  // Whether to use large index.
-  bool large = compiled_use_large_index(inputs, outputs, contiguous);
+  // Whether to use large index (also true for negative strides).
+  bool large =
+      negative_strides || compiled_use_large_index(inputs, outputs, contiguous);
 
   // Get the kernel from the lib
   int ndim = shape.size();

python/tests/test_compile.py

@@ -1454,6 +1454,53 @@ def fun(x):
             np.asarray(out, copy=False).__array_interface__["data"][0], in_ptr
         )
 
+    def test_compile_negative_strides(self):
+        # 1D negative stride with elementwise expression
+        @mx.compile
+        def f(x):
+            return 2.0 * x[::-1]
+
+        x = mx.arange(8, dtype=mx.float32)
+        expected = 2.0 * x[::-1]
+        self.assertTrue(mx.array_equal(f(x), expected))
+
+        # 1D negative stride with slice update
+        def g_eager(x):
+            base = mx.zeros_like(x)
+            base[::-1] += 2.0 * x[::-1]
+            return base
+
+        g_compiled = mx.compile(g_eager)
+        expected = g_eager(x)
+        self.assertTrue(mx.array_equal(g_compiled(x), expected))
+
+        # 2D negative stride
+        @mx.compile
+        def h(x):
+            return x[::-1] + 1.0
+
+        y = mx.arange(12, dtype=mx.float32).reshape(3, 4)
+        expected = y[::-1] + 1.0
+        self.assertTrue(mx.array_equal(h(y), expected))
+
+        # Mixed positive and negative strides
+        @mx.compile
+        def m(x):
+            return x[::-1, ::2] * 3.0
+
+        z = mx.arange(24, dtype=mx.float32).reshape(4, 6)
+        expected = z[::-1, ::2] * 3.0
+        self.assertTrue(mx.array_equal(m(z), expected))
+
+        # 4D negative stride (exercises work_per_thread > 1 path)
+        @mx.compile
+        def p(x):
+            return x + 1.0
+
+        w = mx.arange(120, dtype=mx.float32).reshape(2, 3, 4, 5)
+        expected = w[::-1, :, ::-1, :] + 1.0
+        self.assertTrue(mx.array_equal(p(w[::-1, :, ::-1, :]), expected))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()