Collapse the non sorted axes to pick the contiguous sort kernel

mlx/backend/common/utils.cpp

@@ -89,13 +89,16 @@ std::pair<Shape, Strides> collapse_contiguous_dims(
   Strides collapsed_strides;
 
   if (shape.size() > 0) {
-    collapsed_shape.push_back(shape[0]);
-    collapsed_strides.push_back(strides[0]);
-    for (int i = 1; i < shape.size(); i++) {
+    for (int i = 0; i < shape.size(); i++) {
       if (shape[i] == 1) {
         continue;
-      } else if (
-          strides[i] * shape[i] != collapsed_strides.back() ||
+      }
+      if (collapsed_shape.empty()) {
+        collapsed_shape.push_back(shape[i]);
+        collapsed_strides.push_back(strides[i]);
+        continue;
+      }
+      if (strides[i] * shape[i] != collapsed_strides.back() ||
           collapsed_shape.back() * static_cast<int64_t>(shape[i]) > size_cap) {
         collapsed_shape.push_back(shape[i]);
         collapsed_strides.push_back(strides[i]);
@@ -104,6 +107,10 @@ std::pair<Shape, Strides> collapse_contiguous_dims(
         collapsed_strides.back() = strides[i];
       }
     }
+    if (collapsed_shape.empty()) {
+      collapsed_shape.push_back(1);
+      collapsed_strides.push_back(1);
+    }
   }
 
   return std::make_pair(collapsed_shape, collapsed_strides);

mlx/backend/cuda/sort.cu

@@ -4,6 +4,7 @@
 #include <cassert>
 #include <cstdint>
 
+#include "mlx/backend/common/utils.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/device/fp16_math.cuh"
 #include "mlx/backend/cuda/kernel_utils.cuh"
@@ -781,16 +782,24 @@ void single_block_sort(
   int64_t in_stride_sorted_axis = in.strides()[axis];
   int64_t out_stride_sorted_axis = out.strides()[axis];
 
+  // The contiguous kernel walks the rows with a single stride, so the axes
+  // that are not sorted have to collapse to a single dim.
+  auto single_dim =
+      [](const Shape& nc_shape, const Strides& nc_str, int64_t& stride) {
+        auto [shape, strides] = collapse_contiguous_dims(nc_shape, nc_str);
+        if (shape.size() > 1) {
+          return false;
+        }
+        if (shape.size() > 0) {
+          stride = strides[0];
+        }
+        return true;
+      };
+  int64_t in_stride_segment_axis = 0;
+  int64_t out_stride_segment_axis = 0;
   bool contiguous = in.flags().contiguous;
-  auto check_strides = [](const array& x, int64_t sort_stride) {
-    int64_t min_stride =
-        *std::min_element(x.strides().begin(), x.strides().end());
-    int64_t max_stride =
-        *std::max_element(x.strides().begin(), x.strides().end());
-    return sort_stride == min_stride || sort_stride == max_stride;
-  };
-  contiguous &= check_strides(in, in_stride_sorted_axis);
-  contiguous &= check_strides(out, out_stride_sorted_axis);
+  contiguous &= single_dim(nc_shape, in_nc_str, in_stride_segment_axis);
+  contiguous &= single_dim(nc_shape, out_nc_str, out_stride_segment_axis);
 
   auto& encoder = cu::get_command_encoder(s);
   out.set_data(cu::malloc_async(out.nbytes(), encoder));
@@ -817,19 +826,9 @@ 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_stride_segment_axis > INT32_MAX ||
+                out_stride_segment_axis > INT32_MAX) {
+              throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
             }
             encoder.add_kernel_node(
                 kernel,

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"
@@ -39,17 +40,24 @@ void single_block_sort(
   int in_stride_sorted_axis = in.strides()[axis];
   int out_stride_sorted_axis = out.strides()[axis];
 
-  // We can only use the contiguous kernel if the sorted axis
-  // has the largest or smallest stride.
-  // We also need the input to be contiguous
+  // The contiguous kernel walks the rows with a single stride, so the axes
+  // that are not sorted have to collapse to a single dim.
+  auto single_dim =
+      [](const Shape& nc_shape, const Strides& nc_str, int64_t& stride) {
+        auto [shape, strides] = collapse_contiguous_dims(nc_shape, nc_str);
+        if (shape.size() > 1) {
+          return false;
+        }
+        if (shape.size() > 0) {
+          stride = strides[0];
+        }
+        return true;
+      };
+  int64_t in_stride_segment_axis = 0;
+  int64_t out_stride_segment_axis = 0;
   bool contiguous = in.flags().contiguous;
-  auto check_strides = [](array x, int sort_stride) {
-    int min_stride = *std::min_element(x.strides().begin(), x.strides().end());
-    int max_stride = *std::max_element(x.strides().begin(), x.strides().end());
-    return sort_stride == min_stride || sort_stride == max_stride;
-  };
-  contiguous &= check_strides(in, in_stride_sorted_axis);
-  contiguous &= check_strides(out, out_stride_sorted_axis);
+  contiguous &= single_dim(nc_shape, in_nc_str, in_stride_segment_axis);
+  contiguous &= single_dim(nc_shape, out_nc_str, out_stride_segment_axis);
 
   // Prepare kernel name
   std::ostringstream kname;
@@ -74,22 +82,12 @@ 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_stride_segment_axis > INT32_MAX ||
+        out_stride_segment_axis > INT32_MAX) {
+      throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
     }
-    compute_encoder.set_bytes(in_stride_segment_axis, 5);
-    compute_encoder.set_bytes(out_stride_segment_axis, 6);
+    compute_encoder.set_bytes(static_cast<int>(in_stride_segment_axis), 5);
+    compute_encoder.set_bytes(static_cast<int>(out_stride_segment_axis), 6);
   } else {
     compute_encoder.set_bytes(nc_dim, 5);
     if (nc_shape.empty()) {

python/tests/test_ops.py

@@ -4660,6 +4660,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):