diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu index d9c77387ce..1d3a2c423c 100644 --- a/mlx/backend/cuda/sort.cu +++ b/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, diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp index 65f144c026..6201e7ae0a 100644 --- a/mlx/backend/metal/sort.cpp +++ b/mlx/backend/metal/sort.cpp @@ -2,6 +2,7 @@ #include +#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(in_nc_str[i])); - out_stride_segment_axis = - std::min(out_stride_segment_axis, static_cast(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(in_seg); + int out_stride_segment_axis = static_cast(out_seg); compute_encoder.set_bytes(in_stride_segment_axis, 5); compute_encoder.set_bytes(out_stride_segment_axis, 6); } else { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 9bcce60fce..c921b4ffa0 100644 --- a/python/tests/test_ops.py +++ b/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):