diff --git a/mlx/backend/common/utils.cpp b/mlx/backend/common/utils.cpp index 66764173e2..a13c022071 100644 --- a/mlx/backend/common/utils.cpp +++ b/mlx/backend/common/utils.cpp @@ -89,13 +89,16 @@ std::pair 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(shape[i]) > size_cap) { collapsed_shape.push_back(shape[i]); collapsed_strides.push_back(strides[i]); @@ -104,6 +107,10 @@ std::pair 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); diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu index 5d73681a63..980cab3e30 100644 --- a/mlx/backend/cuda/sort.cu +++ b/mlx/backend/cuda/sort.cu @@ -4,6 +4,7 @@ #include #include +#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, diff --git a/mlx/backend/metal/sort.cpp b/mlx/backend/metal/sort.cpp index 15d6bdfb4e..b38d3ac2bd 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" @@ -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(in_nc_str[i])); - out_stride_segment_axis = - std::min(out_stride_segment_axis, static_cast(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(in_stride_segment_axis), 5); + compute_encoder.set_bytes(static_cast(out_stride_segment_axis), 6); } else { compute_encoder.set_bytes(nc_dim, 5); if (nc_shape.empty()) { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 836362e1b1..bd4a8e80dd 100644 --- a/python/tests/test_ops.py +++ b/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):