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

Fix small sort with metal validation (#1695)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-12 09:21:45 -08:00
committed by GitHub
parent 6bd28d246e
commit 9111999af3
2 changed files with 17 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
mlx/backend/metal/sort.cpp
View File

@@ -82,9 +82,17 @@ void single_block_sort(
compute_encoder.set_bytes(out_stride_segment_axis, 6);
} else {
compute_encoder.set_bytes(nc_dim, 5);
compute_encoder.set_vector_bytes(nc_shape, 6);
compute_encoder.set_vector_bytes(in_nc_str, 7);
compute_encoder.set_vector_bytes(out_nc_str, 8);
if (nc_shape.empty()) {
int shape = 0;
int64_t stride = 0;
compute_encoder.set_bytes(shape, 6);
compute_encoder.set_bytes(stride, 7);
compute_encoder.set_bytes(stride, 8);
} else {
compute_encoder.set_vector_bytes(nc_shape, 6);
compute_encoder.set_vector_bytes(in_nc_str, 7);
compute_encoder.set_vector_bytes(out_nc_str, 8);
}
}
MTL::Size group_dims = MTL::Size(bn, 1, 1);