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

Ensure small sort doesn't use indices if not argsort

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-08-16 21:43:24 -07:00
parent 1ba18ff7d9
commit b33457ae4d
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/backend/metal/kernels/sort.h
View File

@@ -45,7 +45,9 @@ struct ThreadSort {
for (short j = i & 1; j < N_PER_THREAD - 1; j += 2) {
if (op(vals[j + 1], vals[j])) {
thread_swap(vals[j + 1], vals[j]);
thread_swap(idxs[j + 1], idxs[j]);
if (ARG_SORT) {
thread_swap(idxs[j + 1], idxs[j]);
}
}
}
}
@@ -111,7 +113,9 @@ struct BlockMergeSort {
bool pred = (b_idx < B_sz) && (a_idx >= A_sz || op(b, a));
vals[i] = pred ? b : a;
idxs[i] = pred ? Bs_idx[b_idx] : As_idx[a_idx];
if (ARG_SORT) {
idxs[i] = pred ? Bs_idx[b_idx] : As_idx[a_idx];
}
b_idx += short(pred);
a_idx += short(!pred);