Collapse the non sorted axes to pick the contiguous sort kernel

mlx/backend/cuda/sort.cu

@@ -7,6 +7,7 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/device/fp16_math.cuh"
 #include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/common/utils.h"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/dtype_utils.h"
 #include "mlx/primitives.h"
@@ -791,6 +792,20 @@ void single_block_sort(
   };
   contiguous &= check_strides(in, in_stride_sorted_axis);
   contiguous &= check_strides(out, out_stride_sorted_axis);
+  // The contiguous kernel walks the rows with a single stride, so the axes
+  // that are not sorted have to collapse to a single dimension.
+  auto single_run = [](const Shape& shape, const Strides& strides,
+                       int64_t& stride) {
+    auto [cshape, cstrides] = collapse_contiguous_dims(shape, strides);
+    stride = cstrides.empty() ? 0 : cstrides.back();
+    return std::count_if(cshape.begin(), cshape.end(), [](auto d) {
+             return d != 1;
+           }) <= 1;
+  };
+  int64_t in_seg = 0;
+  int64_t out_seg = 0;
+  contiguous &= single_run(nc_shape, in_nc_str, in_seg);
+  contiguous &= single_run(nc_shape, out_nc_str, out_seg);
 
   auto& encoder = cu::get_command_encoder(s);
   out.set_data(cu::malloc_async(out.nbytes(), encoder));
@@ -817,20 +832,11 @@ void single_block_sort(
                 ARG_SORT,
                 BLOCK_THREADS,
                 N_PER_THREAD>;
-            int64_t in_stride_segment_axis = INT64_MAX;
-            int64_t out_stride_segment_axis = INT64_MAX;
-            for (int i = 0; i < nc_shape.size(); i++) {
-              if (nc_shape[i] == 1) {
-                continue;
-              }
-              if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) {
-                throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
-              }
-              in_stride_segment_axis =
-                  std::min(in_stride_segment_axis, in_nc_str[i]);
-              out_stride_segment_axis =
-                  std::min(out_stride_segment_axis, out_nc_str[i]);
+            if (in_seg > INT32_MAX || out_seg > INT32_MAX) {
+              throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
             }
+            int64_t in_stride_segment_axis = in_seg;
+            int64_t out_stride_segment_axis = out_seg;
             encoder.add_kernel_node(
                 kernel,
                 grid,

mlx/backend/metal/sort.cpp

@@ -2,6 +2,7 @@
 
 #include <algorithm>
 
+#include "mlx/backend/common/utils.h"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/kernels.h"
@@ -50,6 +51,20 @@ void single_block_sort(
   };
   contiguous &= check_strides(in, in_stride_sorted_axis);
   contiguous &= check_strides(out, out_stride_sorted_axis);
+  // The contiguous kernel walks the rows with a single stride, so the axes
+  // that are not sorted have to collapse to a single dimension.
+  auto single_run =
+      [](const Shape& shape, const Strides& strides, int64_t& stride) {
+        auto [cshape, cstrides] = collapse_contiguous_dims(shape, strides);
+        stride = cstrides.empty() ? 0 : cstrides.back();
+        return std::count_if(cshape.begin(), cshape.end(), [](auto d) {
+                 return d != 1;
+               }) <= 1;
+      };
+  int64_t in_seg = 0;
+  int64_t out_seg = 0;
+  contiguous &= single_run(nc_shape, in_nc_str, in_seg);
+  contiguous &= single_run(nc_shape, out_nc_str, out_seg);
 
   // Prepare kernel name
   std::ostringstream kname;
@@ -74,20 +89,11 @@ void single_block_sort(
   compute_encoder.set_bytes(out_stride_sorted_axis, 4);
 
   if (contiguous) {
-    int in_stride_segment_axis = INT32_MAX;
-    int out_stride_segment_axis = INT32_MAX;
-    for (int i = 0; i < in_nc_str.size(); i++) {
-      if (nc_shape[i] == 1) {
-        continue;
-      }
-      if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) {
-        throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
-      }
-      in_stride_segment_axis =
-          std::min(in_stride_segment_axis, static_cast<int>(in_nc_str[i]));
-      out_stride_segment_axis =
-          std::min(out_stride_segment_axis, static_cast<int>(out_nc_str[i]));
+    if (in_seg > INT32_MAX || out_seg > INT32_MAX) {
+      throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
     }
+    int in_stride_segment_axis = static_cast<int>(in_seg);
+    int out_stride_segment_axis = static_cast<int>(out_seg);
     compute_encoder.set_bytes(in_stride_segment_axis, 5);
     compute_encoder.set_bytes(out_stride_segment_axis, 6);
   } else {

python/tests/test_ops.py

@@ -4189,6 +4189,28 @@ def test_broadcast_shapes(self):
         with self.assertRaises(ValueError):
             mx.broadcast_shapes()
 
+    def test_sort_transposed(self):
+        # The sorted axis can keep the smallest or largest stride while the
+        # axes that are not sorted are no longer a row major block, which is
+        # what the contiguous kernel's row enumeration assumes.
+        np.random.seed(0)
+        for shape in [(3, 4, 8), (2, 1, 6), (2, 3, 4, 2), (2, 1, 3, 4)]:
+            a_np = np.random.uniform(0, 100, size=shape).astype(np.float32)
+            a_mx = mx.array(a_np)
+            for perm in permutations(range(len(shape))):
+                b_np = np.transpose(a_np, perm)
+                b_mx = mx.transpose(a_mx, perm)
+                for axis in range(len(shape)):
+                    with self.subTest(shape=shape, perm=perm, axis=axis):
+                        s_np = np.sort(b_np, axis=axis)
+                        self.assertTrue(np.array_equal(s_np, mx.sort(b_mx, axis=axis)))
+                        idx = np.array(mx.argsort(b_mx, axis=axis))
+                        self.assertTrue(
+                            np.array_equal(
+                                s_np, np.take_along_axis(b_np, idx, axis=axis)
+                            )
+                        )
+
     def test_sort_nan(self):
         for dtype in [mx.float32, mx.float16, mx.bfloat16]:
             with self.subTest(dtype=dtype):