Fix Metal sort for arrays with negative strides

mlx/backend/metal/sort.cpp

@@ -278,9 +278,27 @@ void gpu_merge_sort(
     array& out,
     int axis_,
     bool argsort) {
+  // Negative strides cause MSL elem_to_loc to compute negative relative pointer
+  // offsets, attempting to read memory prior to Metal's bound buffer window
+  // (a_buf + in.offset()). To prevent GPU memory protection traps while
+  // matching NumPy behavior, we create a temporary contiguous copy for negative
+  // strides.
+  bool has_negative_stride = false;
+  for (auto st : in.strides()) {
+    if (st < 0) {
+      has_negative_stride = true;
+      break;
+    }
+  }
+
+  array in_dense = in;
+  if (has_negative_stride) {
+    in_dense = contiguous_copy_gpu(in, s);
+  }
+
   // Get size info
-  int axis = axis_ < 0 ? axis_ + in.ndim() : axis_;
-  int size_sorted_axis = in.shape(axis);
+  int axis = axis_ < 0 ? axis_ + in_dense.ndim() : axis_;
+  int size_sorted_axis = in_dense.shape(axis);
 
   // Get kernel size
   int tn = 4;
@@ -299,17 +317,18 @@ void gpu_merge_sort(
     bn = 32;
   }
 
-  if (bn == 512 && size_of(in.dtype()) > 4) {
+  if (bn == 512 && size_of(in_dense.dtype()) > 4) {
     bn = 256;
   }
 
   int n_per_block = bn * tn;
   int n_blocks = (size_sorted_axis + n_per_block - 1) / n_per_block;
 
   if (n_blocks > 1) {
-    return multi_block_sort(s, d, in, out, axis, bn, tn, n_blocks, argsort);
+    return multi_block_sort(
+        s, d, in_dense, out, axis, bn, tn, n_blocks, argsort);
   } else {
-    return single_block_sort(s, d, in, out, axis, bn, tn, argsort);
+    return single_block_sort(s, d, in_dense, out, axis, bn, tn, argsort);
   }
 }
 

python/tests/test_ops.py

@@ -3939,6 +3939,20 @@ def test_zeros_ones_empty_like_dtype(self):
         e = mx.empty_like(x, dtype=mx.float32)
         self.assertEqual(e.dtype, mx.float32)
 
+    def test_strided_negative_stride_sort(self):
+        a_np = np.arange(12, dtype=np.float32).reshape(3, 4)
+        a_mx = mx.array(a_np)[::-1, :]
+        a_np_strided = a_np[::-1, :]
+
+        expected_sort = np.sort(a_np_strided, axis=-1)
+        expected_topk = np.sort(a_np_strided, axis=-1)[:, -2:]
+
+        self.assertTrue(np.allclose(np.array(mx.sort(a_mx, axis=-1)), expected_sort))
+        self.assertTrue(np.allclose(np.array(mx.topk(a_mx, 2, axis=-1)), expected_topk))
+        self.assertTrue(
+            np.allclose(np.array(mx.partition(a_mx, 2, axis=-1)), expected_sort)
+        )
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()