Skip to content

Reduce overhead and optimise partial sorting algorithm #53

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Reduce overhead and optimise partial sorting algorithm #53

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 47 additions & 4 deletions src/avx512-16bit-common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,10 +316,53 @@ static void qselect_16bit_(type_t *arr,
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512<vtype>(
arr, left, right + 1, pivot, &smallest, &biggest);
if ((pivot != smallest) && (pos < pivot_index))
qselect_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
else if ((pivot != biggest) && (pos >= pivot_index))
qselect_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);

if (pos < pivot_index) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would prefer to keep the if conditions the way they are, just more readable I think.

if (pivot != smallest)
qselect_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != biggest)
qselect_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

template <typename vtype, typename type_t>
static void qselsort_16bit_(type_t *arr,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Think we need a better name. May be qpartialsort_*?

int64_t pos,
int64_t left,
int64_t right,
int64_t max_iters)
{
/*
* Resort to std::sort if quicksort isnt making any progress
*/
if (max_iters <= 0) {
std::sort(arr + left, arr + right + 1, comparison_func<vtype>);
return;
}
/*
* Base case: use bitonic networks to sort arrays <= 128
*/
if (right + 1 - left <= 128) {
sort_128_16bit<vtype>(arr + left, (int32_t)(right + 1 - left));
return;
}

type_t pivot = get_pivot_16bit<vtype>(arr, left, right);
type_t smallest = vtype::type_max();
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512<vtype>(
arr, left, right + 1, pivot, &smallest, &biggest);

if (pos < pivot_index) {
if (pivot != smallest)
qselsort_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != smallest)
qsort_16bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
if (pivot != biggest)
qselsort_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

#endif // AVX512_16BIT_COMMON
32 changes: 31 additions & 1 deletion src/avx512-16bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,14 +433,44 @@ void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
if (indx_last_elem >= k) {
qselect_16bit_<zmm_vector<float16>, uint16_t>(
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_partial_qsort(int16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_16bit_<zmm_vector<int16_t>, int16_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_partial_qsort(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_16bit_<zmm_vector<uint16_t>, uint16_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (LIKELY(k > 0)) {
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
qselsort_16bit_<zmm_vector<float16>, uint16_t>(
arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_qsort(int16_t *arr, int64_t arrsize)
{
Expand Down
83 changes: 78 additions & 5 deletions src/avx512-32bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -683,10 +683,52 @@ static void qselect_32bit_(type_t *arr,
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512_unrolled<vtype, 2>(
arr, left, right + 1, pivot, &smallest, &biggest);
if ((pivot != smallest) && (pos < pivot_index))
qselect_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
else if ((pivot != biggest) && (pos >= pivot_index))
qselect_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
if (pos < pivot_index) {
if (pivot != smallest)
qselect_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != biggest)
qselect_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

template <typename vtype, typename type_t>
static void qselsort_32bit_(type_t *arr,
int64_t pos,
int64_t left,
int64_t right,
int64_t max_iters)
{
/*
* Resort to std::sort if quicksort isnt making any progress
*/
if (max_iters <= 0) {
std::sort(arr + left, arr + right + 1);
return;
}
/*
* Base case: use bitonic networks to sort arrays <= 128
*/
if (right + 1 - left <= 128) {
sort_128_32bit<vtype>(arr + left, (int32_t)(right + 1 - left));
return;
}

type_t pivot = get_pivot_32bit<vtype>(arr, left, right);
type_t smallest = vtype::type_max();
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512_unrolled<vtype, 2>(
arr, left, right + 1, pivot, &smallest, &biggest);

if (pos < pivot_index) {
if (pivot != smallest)
qselsort_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != smallest)
qsort_32bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
if (pivot != biggest)
qselsort_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(float *arr, int64_t arrsize)
Expand Down Expand Up @@ -737,14 +779,45 @@ void avx512_qselect<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
{
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
if (indx_last_elem >= k) {
qselect_32bit_<zmm_vector<float>, float>(
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_partial_qsort<int32_t>(int32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_32bit_<zmm_vector<int32_t>, int32_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_partial_qsort<uint32_t>(uint32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_32bit_<zmm_vector<uint32_t>, uint32_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_partial_qsort<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (LIKELY(k > 0)) {
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
qselsort_32bit_<zmm_vector<float>, float>(
arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_qsort<int32_t>(int32_t *arr, int64_t arrsize)
{
Expand Down
86 changes: 80 additions & 6 deletions src/avx512-64bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,7 +732,7 @@ qsort_64bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
return;
}
/*
* Base case: use bitonic networks to sort arrays <= 128
* Base case: use bitonic networks to sort arrays <= 256
*/
if (right + 1 - left <= 256) {
sort_256_64bit<vtype>(arr + left, (int32_t)(right + 1 - left));
Expand Down Expand Up @@ -777,10 +777,53 @@ static void qselect_64bit_(type_t *arr,
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512_unrolled<vtype, 8>(
arr, left, right + 1, pivot, &smallest, &biggest);
if ((pivot != smallest) && (pos < pivot_index))
qselect_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
else if ((pivot != biggest) && (pos >= pivot_index))
qselect_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);

if (pos < pivot_index) {
if (pivot != smallest)
qselect_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != biggest)
qselect_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

template <typename vtype, typename type_t>
static void qselsort_64bit_(type_t *arr,
int64_t pos,
int64_t left,
int64_t right,
int64_t max_iters)
{
/*
* Resort to std::sort if quicksort isnt making any progress
*/
if (max_iters <= 0) {
std::sort(arr + left, arr + right + 1);
return;
}
/*
* Base case: use bitonic networks to sort arrays <= 128
*/
if (right + 1 - left <= 128) {
sort_128_64bit<vtype>(arr + left, (int32_t)(right + 1 - left));
return;
}

type_t pivot = get_pivot_64bit<vtype>(arr, left, right);
type_t smallest = vtype::type_max();
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512_unrolled<vtype, 8>(
arr, left, right + 1, pivot, &smallest, &biggest);

if (pos < pivot_index) {
if (pivot != smallest)
qselsort_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
} else {
if (pivot != smallest)
qsort_64bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
if (pivot != biggest)
qselsort_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
}
}

template <>
Expand All @@ -806,14 +849,45 @@ void avx512_qselect<double>(double *arr, int64_t k, int64_t arrsize, bool hasnan
{
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
if (indx_last_elem >= k) {
qselect_64bit_<zmm_vector<double>, double>(
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_partial_qsort<int64_t>(int64_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_64bit_<zmm_vector<int64_t>, int64_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_partial_qsort<uint64_t>(uint64_t *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (arrsize > 1) {
qselsort_64bit_<zmm_vector<uint64_t>, uint64_t>(
arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_partial_qsort<double>(double *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (LIKELY(k > 0)) {
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
qselsort_64bit_<zmm_vector<double>, double>(
arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_qsort<int64_t>(int64_t *arr, int64_t arrsize)
{
Expand Down
12 changes: 2 additions & 10 deletions src/avx512-common-qsort.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,16 +104,8 @@ void avx512_qselect(T *arr, int64_t k, int64_t arrsize, bool hasnan = false);
void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false);

template <typename T>
inline void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize, bool hasnan = false)
{
avx512_qselect<T>(arr, k - 1, arrsize, hasnan);
avx512_qsort<T>(arr, k - 1);
}
inline void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false)
{
avx512_qselect_fp16(arr, k - 1, arrsize, hasnan);
avx512_qsort_fp16(arr, k - 1);
}
void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize, bool hasnan = false);
void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false);

// key-value sort routines
template <typename T>
Expand Down
15 changes: 14 additions & 1 deletion src/avx512fp16-16bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,14 +157,27 @@ void avx512_qselect(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)
{
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
if (indx_last_elem >= k) {
qselect_16bit_<zmm_vector<_Float16>, _Float16>(
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_partial_qsort(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)
{
if (LIKELY(k > 0)) {
int64_t indx_last_elem = arrsize - 1;
if (UNLIKELY(hasnan)) {
indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
}
qselsort_16bit_<zmm_vector<_Float16>, _Float16>(
arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

template <>
void avx512_qsort(_Float16 *arr, int64_t arrsize)
{
Expand Down