L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Fast sorting (#3060)

Browse Source
This commit is contained in:
Awni Hannun
2026-01-25 15:10:18 -08:00
committed by GitHub
parent 3ac892b008
commit 9e2d2a5957
2 changed files with 1012 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/sort.cu
+1005 -128
View File
File diff suppressed because it is too large Load Diff
mlx/backend/metal/kernels/sort.h
+7 -3
View File
@@ -125,13 +125,17 @@ struct BlockMergeSort {
short b_idx = 0;
for (int i = 0; i < N_PER_THREAD; ++i) {
auto a = As[a_idx];
auto b = Bs[b_idx];
auto a = (a_idx < A_sz) ? As[a_idx] : ValT(CompareOp::init);
auto b = (b_idx < B_sz) ? Bs[b_idx] : ValT(CompareOp::init);
bool pred = (b_idx < B_sz) && (a_idx >= A_sz || op(b, a));
vals[i] = pred ? b : a;
if (ARG_SORT) {
idxs[i] = pred ? Bs_idx[b_idx] : As_idx[a_idx];
if (pred) {
idxs[i] = Bs_idx[b_idx];
} else {
idxs[i] = (a_idx < A_sz) ? As_idx[a_idx] : IdxT(0);
}
}
b_idx += short(pred);