zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix GPU sort for large arrays (#1285)

Browse Source 
* Fix GPU sort for large arrays
This commit is contained in:
Jagrit Digani
2024-07-24 14:37:10 -07:00
committed by GitHub
parent ebd7135b50
commit 7f914365fd
4 changed files with 34 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/ops.cpp
View File

@@ -1785,15 +1785,6 @@ array sort(const array& a, int axis, StreamOrDevice s /* = {} */) {
throw std::invalid_argument(msg.str());
}
// TODO: Fix GPU kernel
if (a.shape(axis) >= (1u << 21) && to_stream(s).device.type == Device::gpu) {
std::ostringstream msg;
msg << "[sort] GPU sort cannot handle sort axis of >= 2M elements,"
<< " got array with sort axis size " << a.shape(axis) << "."
<< " Please place this operation on the CPU instead.";
throw std::runtime_error(msg.str());
}
return array(
a.shape(), a.dtype(), std::make_shared<Sort>(to_stream(s), axis), {a});
}