Skip to content
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

Jump to bottom

Add avx512_argselect for 32-bit and 64-bit dtypes #56

Merged
merged 6 commits into from
Jul 21, 2023
Merged

Add avx512_argselect for 32-bit and 64-bit dtypes #56

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
243fbe4
Add avx512_argselect
r-devulap Jul 18, 2023
608538d
More test condition for argselect
r-devulap Jul 18, 2023
8bd9c42
Use std_argselect_withnan
r-devulap Jul 20, 2023
437fd4a
Add tests for arrays with NAN
r-devulap Jul 20, 2023
e1b7fa7
clang format
r-devulap Jul 20, 2023
0ae431a
Fix build issues
r-devulap Jul 21, 2023
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
4 changes: 2 additions & 2 deletions src/avx512-16bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,11 +433,11 @@ 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));
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

Expand Down
14 changes: 10 additions & 4 deletions src/avx512-32bit-qsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -715,7 +715,10 @@ replace_inf_with_nan(float *arr, int64_t arrsize, int64_t nan_count)
}

template <>
void avx512_qselect<int32_t>(int32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
void avx512_qselect<int32_t>(int32_t *arr,
int64_t k,
int64_t arrsize,
bool hasnan)
{
if (arrsize > 1) {
qselect_32bit_<zmm_vector<int32_t>, int32_t>(
Expand All @@ -724,7 +727,10 @@ void avx512_qselect<int32_t>(int32_t *arr, int64_t k, int64_t arrsize, bool hasn
}

template <>
void avx512_qselect<uint32_t>(uint32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
void avx512_qselect<uint32_t>(uint32_t *arr,
int64_t k,
int64_t arrsize,
bool hasnan)
{
if (arrsize > 1) {
qselect_32bit_<zmm_vector<uint32_t>, uint32_t>(
Expand All @@ -737,11 +743,11 @@ 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));
arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
}
}

Expand Down
152 changes: 139 additions & 13 deletions src/avx512-64bit-argsort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,28 @@
#define AVX512_ARGSORT_64BIT

#include "avx512-64bit-common.h"
#include "avx512-common-argsort.h"
#include "avx512-64bit-keyvalue-networks.hpp"
#include "avx512-common-argsort.h"

template <typename T>
void std_argselect_withnan(
T *arr, int64_t *arg, int64_t k, int64_t left, int64_t right)
{
std::nth_element(arg + left,
arg + k,
arg + right,
[arr](int64_t a, int64_t b) -> bool {
if ((!std::isnan(arr[a])) && (!std::isnan(arr[b]))) {
return arr[a] < arr[b];
}
else if (std::isnan(arr[a])) {
return false;
}
else {
return true;
}
});
}

/* argsort using std::sort */
template <typename T>
Expand All @@ -18,9 +38,15 @@ void std_argsort_withnan(T *arr, int64_t *arg, int64_t left, int64_t right)
std::sort(arg + left,
arg + right,
[arr](int64_t left, int64_t right) -> bool {
if ((!std::isnan(arr[left])) && (!std::isnan(arr[right]))) {return arr[left] < arr[right];}
else if (std::isnan(arr[left])) {return false;}
else {return true;}
if ((!std::isnan(arr[left])) && (!std::isnan(arr[right]))) {
return arr[left] < arr[right];
}
else if (std::isnan(arr[left])) {
return false;
}
else {
return true;
}
});
}

Expand Down Expand Up @@ -284,7 +310,42 @@ inline void argsort_64bit_(type_t *arr,
}

template <typename vtype, typename type_t>
bool has_nan(type_t* arr, int64_t arrsize)
static void argselect_64bit_(type_t *arr,
int64_t *arg,
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_argsort(arr, arg, left, right + 1);
return;
}
/*
* Base case: use bitonic networks to sort arrays <= 64
*/
if (right + 1 - left <= 64) {
argsort_64_64bit<vtype>(arr, arg + left, (int32_t)(right + 1 - left));
return;
}
type_t pivot = get_pivot_64bit<vtype>(arr, arg, left, right);
type_t smallest = vtype::type_max();
type_t biggest = vtype::type_min();
int64_t pivot_index = partition_avx512_unrolled<vtype, 4>(
arr, arg, left, right + 1, pivot, &smallest, &biggest);
if ((pivot != smallest) && (pos < pivot_index))
argselect_64bit_<vtype>(
arr, arg, pos, left, pivot_index - 1, max_iters - 1);
else if ((pivot != biggest) && (pos >= pivot_index))
argselect_64bit_<vtype>(
arr, arg, pos, pivot_index, right, max_iters - 1);
}

template <typename vtype, typename type_t>
bool has_nan(type_t *arr, int64_t arrsize)
{
using opmask_t = typename vtype::opmask_t;
using zmm_t = typename vtype::zmm_t;
Expand All @@ -299,7 +360,7 @@ bool has_nan(type_t* arr, int64_t arrsize)
else {
in = vtype::loadu(arr);
}
opmask_t nanmask = vtype::template fpclass<0x01|0x80>(in);
opmask_t nanmask = vtype::template fpclass<0x01 | 0x80>(in);
arr += vtype::numlanes;
arrsize -= vtype::numlanes;
if (nanmask != 0x00) {
Expand All @@ -310,8 +371,9 @@ bool has_nan(type_t* arr, int64_t arrsize)
return found_nan;
}

/* argsort methods for 32-bit and 64-bit dtypes */
template <typename T>
void avx512_argsort(T* arr, int64_t *arg, int64_t arrsize)
void avx512_argsort(T *arr, int64_t *arg, int64_t arrsize)
{
if (arrsize > 1) {
argsort_64bit_<zmm_vector<T>>(
Expand All @@ -320,7 +382,7 @@ void avx512_argsort(T* arr, int64_t *arg, int64_t arrsize)
}

template <>
void avx512_argsort(double* arr, int64_t *arg, int64_t arrsize)
void avx512_argsort(double *arr, int64_t *arg, int64_t arrsize)
{
if (arrsize > 1) {
if (has_nan<zmm_vector<double>>(arr, arrsize)) {
Expand All @@ -333,9 +395,8 @@ void avx512_argsort(double* arr, int64_t *arg, int64_t arrsize)
}
}


template <>
void avx512_argsort(int32_t* arr, int64_t *arg, int64_t arrsize)
void avx512_argsort(int32_t *arr, int64_t *arg, int64_t arrsize)
{
if (arrsize > 1) {
argsort_64bit_<ymm_vector<int32_t>>(
Expand All @@ -344,7 +405,7 @@ void avx512_argsort(int32_t* arr, int64_t *arg, int64_t arrsize)
}

template <>
void avx512_argsort(uint32_t* arr, int64_t *arg, int64_t arrsize)
void avx512_argsort(uint32_t *arr, int64_t *arg, int64_t arrsize)
{
if (arrsize > 1) {
argsort_64bit_<ymm_vector<uint32_t>>(
Expand All @@ -353,7 +414,7 @@ void avx512_argsort(uint32_t* arr, int64_t *arg, int64_t arrsize)
}

template <>
void avx512_argsort(float* arr, int64_t *arg, int64_t arrsize)
void avx512_argsort(float *arr, int64_t *arg, int64_t arrsize)
{
if (arrsize > 1) {
if (has_nan<ymm_vector<float>>(arr, arrsize)) {
Expand All @@ -367,12 +428,77 @@ void avx512_argsort(float* arr, int64_t *arg, int64_t arrsize)
}

template <typename T>
std::vector<int64_t> avx512_argsort(T* arr, int64_t arrsize)
std::vector<int64_t> avx512_argsort(T *arr, int64_t arrsize)
{
std::vector<int64_t> indices(arrsize);
std::iota(indices.begin(), indices.end(), 0);
avx512_argsort<T>(arr, indices.data(), arrsize);
return indices;
}

/* argselect methods for 32-bit and 64-bit dtypes */
template <typename T>
void avx512_argselect(T *arr, int64_t *arg, int64_t k, int64_t arrsize)
{
if (arrsize > 1) {
argselect_64bit_<zmm_vector<T>>(
arr, arg, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_argselect(double *arr, int64_t *arg, int64_t k, int64_t arrsize)
{
if (arrsize > 1) {
if (has_nan<zmm_vector<double>>(arr, arrsize)) {
std_argselect_withnan(arr, arg, k, 0, arrsize);
}
else {
argselect_64bit_<zmm_vector<double>>(
arr, arg, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}
}

template <>
void avx512_argselect(int32_t *arr, int64_t *arg, int64_t k, int64_t arrsize)
{
if (arrsize > 1) {
argselect_64bit_<ymm_vector<int32_t>>(
arr, arg, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_argselect(uint32_t *arr, int64_t *arg, int64_t k, int64_t arrsize)
{
if (arrsize > 1) {
argselect_64bit_<ymm_vector<uint32_t>>(
arr, arg, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}

template <>
void avx512_argselect(float *arr, int64_t *arg, int64_t k, int64_t arrsize)
{
if (arrsize > 1) {
if (has_nan<ymm_vector<float>>(arr, arrsize)) {
std_argselect_withnan(arr, arg, k, 0, arrsize);
}
else {
argselect_64bit_<ymm_vector<float>>(
arr, arg, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
}
}
}

template <typename T>
std::vector<int64_t> avx512_argselect(T *arr, int64_t k, int64_t arrsize)
{
std::vector<int64_t> indices(arrsize);
std::iota(indices.begin(), indices.end(), 0);
avx512_argselect<T>(arr, indices.data(), k, arrsize);
return indices;
}

#endif // AVX512_ARGSORT_64BIT
64 changes: 32 additions & 32 deletions src/avx512-64bit-keyvalue-networks.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,14 +136,14 @@ X86_SIMD_SORT_INLINE void bitonic_merge_four_zmm_64bit(zmm_t *key_zmm,
typename vtype1::opmask_t movmask1 = vtype1::eq(key_zmm_t1, key_zmm[0]);
typename vtype1::opmask_t movmask2 = vtype1::eq(key_zmm_t2, key_zmm[1]);

index_type index_zmm_t1 = vtype2::mask_mov(
index_zmm3r, movmask1, index_zmm[0]);
index_type index_zmm_m1 = vtype2::mask_mov(
index_zmm[0], movmask1, index_zmm3r);
index_type index_zmm_t2 = vtype2::mask_mov(
index_zmm2r, movmask2, index_zmm[1]);
index_type index_zmm_m2 = vtype2::mask_mov(
index_zmm[1], movmask2, index_zmm2r);
index_type index_zmm_t1
= vtype2::mask_mov(index_zmm3r, movmask1, index_zmm[0]);
index_type index_zmm_m1
= vtype2::mask_mov(index_zmm[0], movmask1, index_zmm3r);
index_type index_zmm_t2
= vtype2::mask_mov(index_zmm2r, movmask2, index_zmm[1]);
index_type index_zmm_m2
= vtype2::mask_mov(index_zmm[1], movmask2, index_zmm2r);

// 2) Recursive half clearer: 16
zmm_t key_zmm_t3 = vtype1::permutexvar(rev_index1, key_zmm_m2);
Expand All @@ -159,14 +159,14 @@ X86_SIMD_SORT_INLINE void bitonic_merge_four_zmm_64bit(zmm_t *key_zmm,
movmask1 = vtype1::eq(key_zmm0, key_zmm_t1);
movmask2 = vtype1::eq(key_zmm2, key_zmm_t3);

index_type index_zmm0 = vtype2::mask_mov(
index_zmm_t2, movmask1, index_zmm_t1);
index_type index_zmm1 = vtype2::mask_mov(
index_zmm_t1, movmask1, index_zmm_t2);
index_type index_zmm2 = vtype2::mask_mov(
index_zmm_t4, movmask2, index_zmm_t3);
index_type index_zmm3 = vtype2::mask_mov(
index_zmm_t3, movmask2, index_zmm_t4);
index_type index_zmm0
= vtype2::mask_mov(index_zmm_t2, movmask1, index_zmm_t1);
index_type index_zmm1
= vtype2::mask_mov(index_zmm_t1, movmask1, index_zmm_t2);
index_type index_zmm2
= vtype2::mask_mov(index_zmm_t4, movmask2, index_zmm_t3);
index_type index_zmm3
= vtype2::mask_mov(index_zmm_t3, movmask2, index_zmm_t4);

key_zmm[0] = bitonic_merge_zmm_64bit<vtype1, vtype2>(key_zmm0, index_zmm0);
key_zmm[1] = bitonic_merge_zmm_64bit<vtype1, vtype2>(key_zmm1, index_zmm1);
Expand Down Expand Up @@ -212,22 +212,22 @@ X86_SIMD_SORT_INLINE void bitonic_merge_eight_zmm_64bit(zmm_t *key_zmm,
typename vtype1::opmask_t movmask3 = vtype1::eq(key_zmm_t3, key_zmm[2]);
typename vtype1::opmask_t movmask4 = vtype1::eq(key_zmm_t4, key_zmm[3]);

index_type index_zmm_t1 = vtype2::mask_mov(
index_zmm7r, movmask1, index_zmm[0]);
index_type index_zmm_m1 = vtype2::mask_mov(
index_zmm[0], movmask1, index_zmm7r);
index_type index_zmm_t2 = vtype2::mask_mov(
index_zmm6r, movmask2, index_zmm[1]);
index_type index_zmm_m2 = vtype2::mask_mov(
index_zmm[1], movmask2, index_zmm6r);
index_type index_zmm_t3 = vtype2::mask_mov(
index_zmm5r, movmask3, index_zmm[2]);
index_type index_zmm_m3 = vtype2::mask_mov(
index_zmm[2], movmask3, index_zmm5r);
index_type index_zmm_t4 = vtype2::mask_mov(
index_zmm4r, movmask4, index_zmm[3]);
index_type index_zmm_m4 = vtype2::mask_mov(
index_zmm[3], movmask4, index_zmm4r);
index_type index_zmm_t1
= vtype2::mask_mov(index_zmm7r, movmask1, index_zmm[0]);
index_type index_zmm_m1
= vtype2::mask_mov(index_zmm[0], movmask1, index_zmm7r);
index_type index_zmm_t2
= vtype2::mask_mov(index_zmm6r, movmask2, index_zmm[1]);
index_type index_zmm_m2
= vtype2::mask_mov(index_zmm[1], movmask2, index_zmm6r);
index_type index_zmm_t3
= vtype2::mask_mov(index_zmm5r, movmask3, index_zmm[2]);
index_type index_zmm_m3
= vtype2::mask_mov(index_zmm[2], movmask3, index_zmm5r);
index_type index_zmm_t4
= vtype2::mask_mov(index_zmm4r, movmask4, index_zmm[3]);
index_type index_zmm_m4
= vtype2::mask_mov(index_zmm[3], movmask4, index_zmm4r);

zmm_t key_zmm_t5 = vtype1::permutexvar(rev_index1, key_zmm_m4);
zmm_t key_zmm_t6 = vtype1::permutexvar(rev_index1, key_zmm_m3);
Expand Down
Loading