From 4469d04e5d237d1ec12224a1c311e0236309809d Mon Sep 17 00:00:00 2001
From: Tanish Jain <recklurker@gmail.com>
Date: Thu, 13 Aug 2026 11:38:07 +0530
Subject: [PATCH] fix(metal): handle negative strides in gpu_merge_sort (#4225)

---
 mlx/backend/metal/sort.cpp | 29 ++++++++++++++++++++++++-----
 python/tests/test_ops.py   | 14 ++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp
index 65f144c026..2167a95fd0 100644
--- a/mlx/backend/metal/sort.cpp
+++ b/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,7 +317,7 @@ 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;
   }
 
@@ -307,9 +325,10 @@ void gpu_merge_sort(
   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);
   }
 }
 
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index d569dbd3c7..a6e0de26e3 100644
--- a/python/tests/test_ops.py
+++ b/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()