Condense AscendingComparator and DescendingComparator into one class

Raghuveer Devulapalli · Raghuveer Devulapalli · commit b0d09297f6bf · 2024-03-28T09:54:39.000-07:00
diff --git a/src/avx512-16bit-qsort.hpp b/src/avx512-16bit-qsort.hpp
@@ -574,11 +574,11 @@ X86_SIMD_SORT_INLINE void avx512_qsort_fp16(uint16_t *arr,
                     arr, arrsize);
         }
         if (descending) {
-            qsort_<vtype, DescendingComparator<vtype>, uint16_t>(
+            qsort_<vtype, Comparator<vtype, true>, uint16_t>(
                     arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
         }
         else {
-            qsort_<vtype, AscendingComparator<vtype>, uint16_t>(
+            qsort_<vtype, Comparator<vtype, false>, uint16_t>(
                     arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
         }
         replace_inf_with_nan(arr, arrsize, nan_count, descending);
@@ -599,15 +599,15 @@ X86_SIMD_SORT_INLINE void avx512_qselect_fp16(uint16_t *arr,
     }
     if (indx_last_elem >= k) {
         if (descending) {
-            qselect_<vtype, DescendingComparator<vtype>, uint16_t>(
+            qselect_<vtype, Comparator<vtype, true>, uint16_t>(
                     arr,
                     k,
                     0,
                     indx_last_elem,
                     2 * (arrsize_t)log2(indx_last_elem));
         }
         else {
-            qselect_<vtype, AscendingComparator<vtype>, uint16_t>(
+            qselect_<vtype, Comparator<vtype, false>, uint16_t>(
                     arr,
                     k,
                     0,
diff --git a/src/avx512fp16-16bit-qsort.hpp b/src/avx512fp16-16bit-qsort.hpp
@@ -208,8 +208,8 @@ avx512_qsort(_Float16 *arr, arrsize_t arrsize, bool hasnan)
     using vtype = zmm_vector<_Float16>;
     using comparator =
             typename std::conditional<descending,
-                                      DescendingComparator<vtype>,
-                                      AscendingComparator<vtype>>::type;
+                                      Comparator<vtype, true>,
+                                      Comparator<vtype, false>>::type;
 
     if (arrsize > 1) {
         arrsize_t nan_count = 0;
@@ -231,8 +231,8 @@ avx512_qselect(_Float16 *arr, arrsize_t k, arrsize_t arrsize, bool hasnan)
     using vtype = zmm_vector<_Float16>;
     using comparator =
             typename std::conditional<descending,
-                                      DescendingComparator<vtype>,
-                                      AscendingComparator<vtype>>::type;
+                                      Comparator<vtype, true>,
+                                      Comparator<vtype, false>>::type;
 
     arrsize_t index_first_elem = 0;
     arrsize_t index_last_elem = arrsize - 1;
diff --git a/src/xss-common-comparators.hpp b/src/xss-common-comparators.hpp
@@ -34,124 +34,94 @@ type_t next_value(type_t value)
 template <typename vtype, typename mm_t>
 X86_SIMD_SORT_INLINE void COEX(mm_t &a, mm_t &b);
 
-template <typename vtype>
-struct AscendingComparator {
+template <typename vtype, bool descend>
+struct Comparator {
     using reg_t = typename vtype::reg_t;
     using opmask_t = typename vtype::opmask_t;
     using type_t = typename vtype::type_t;
 
     X86_SIMD_SORT_FINLINE bool STDSortComparator(const type_t &a,
                                                  const type_t &b)
     {
-        return comparison_func<vtype>(a, b);
+        if constexpr (descend) { return comparison_func<vtype>(b, a); }
+        else {
+            return comparison_func<vtype>(a, b);
+        }
     }
 
     X86_SIMD_SORT_FINLINE opmask_t PartitionComparator(reg_t a, reg_t b)
     {
-        return vtype::ge(a, b);
+        if constexpr (descend) { return vtype::ge(b, a); }
+        else {
+            return vtype::ge(a, b);
+        }
     }
 
     X86_SIMD_SORT_FINLINE void COEX(reg_t &a, reg_t &b)
     {
-        ::COEX<vtype, reg_t>(a, b);
+        if constexpr (descend) { ::COEX<vtype, reg_t>(b, a); }
+        else {
+            ::COEX<vtype, reg_t>(a, b);
+        }
     }
 
     // Returns a vector of values that would be sorted as far right as possible
     // For ascending order, this is the maximum possible value
     X86_SIMD_SORT_FINLINE reg_t rightmostPossibleVec()
     {
-        return vtype::zmm_max();
+        if constexpr (descend) { return vtype::zmm_min(); }
+        else {
+            return vtype::zmm_max();
+        }
     }
 
     // Returns the value that would be leftmost of the two when sorted
     // For ascending order, that is the smaller value
     X86_SIMD_SORT_FINLINE type_t leftmost(type_t smaller, type_t larger)
     {
-        UNUSED(larger);
-        return smaller;
+        if constexpr (descend) {
+            UNUSED(smaller);
+            return larger;
+        }
+        else {
+            UNUSED(larger);
+            return smaller;
+        }
     }
 
     // Returns the value that would be rightmost of the two when sorted
     // For ascending order, that is the larger value
     X86_SIMD_SORT_FINLINE type_t rightmost(type_t smaller, type_t larger)
     {
-        UNUSED(smaller);
-        return larger;
+        if constexpr (descend) {
+            UNUSED(larger);
+            return smaller;
+        }
+        else {
+            UNUSED(smaller);
+            return larger;
+        }
     }
 
     // If median == smallest, that implies approximately half the array is equal to smallest, unless we were very unlucky with our sample
     // Try just doing the next largest value greater than this seemingly very common value to seperate them out
     X86_SIMD_SORT_FINLINE type_t choosePivotMedianIsSmallest(type_t median)
     {
-        return next_value<type_t>(median);
+        if constexpr (descend) { return median; }
+        else {
+            return next_value<type_t>(median);
+        }
     }
 
     // If median == largest, that implies approximately half the array is equal to largest, unless we were very unlucky with our sample
     // Thus, median probably is a fine pivot, since it will move all of this common value into its own partition
     X86_SIMD_SORT_FINLINE type_t choosePivotMedianIsLargest(type_t median)
     {
-        return median;
-    }
-};
-
-template <typename vtype>
-struct DescendingComparator {
-    using reg_t = typename vtype::reg_t;
-    using opmask_t = typename vtype::opmask_t;
-    using type_t = typename vtype::type_t;
-
-    X86_SIMD_SORT_FINLINE bool STDSortComparator(const type_t &a,
-                                                 const type_t &b)
-    {
-        return comparison_func<vtype>(b, a);
-    }
-
-    X86_SIMD_SORT_FINLINE opmask_t PartitionComparator(reg_t a, reg_t b)
-    {
-        return vtype::ge(b, a);
-    }
-
-    X86_SIMD_SORT_FINLINE void COEX(reg_t &a, reg_t &b)
-    {
-        ::COEX<vtype, reg_t>(b, a);
-    }
-
-    // Returns a vector of values that would be sorted as far right as possible
-    // For descending order, this is the minimum possible value
-    X86_SIMD_SORT_FINLINE reg_t rightmostPossibleVec()
-    {
-        return vtype::zmm_min();
-    }
-
-    // Returns the value that would be leftmost of the two when sorted
-    // For descending order, that is the larger value
-    X86_SIMD_SORT_FINLINE type_t leftmost(type_t smaller, type_t bigger)
-    {
-        UNUSED(smaller);
-        return bigger;
-    }
-
-    // Returns the value that would be rightmost of the two when sorted
-    // For descending order, that is the smaller value
-    X86_SIMD_SORT_FINLINE type_t rightmost(type_t smaller, type_t bigger)
-    {
-        UNUSED(bigger);
-        return smaller;
-    }
-
-    // If median == smallest, that implies approximately half the array is equal to smallest, unless we were very unlucky with our sample
-    // Thus, median probably is a fine pivot, since it will move all of this common value into its own partition
-    X86_SIMD_SORT_FINLINE type_t choosePivotMedianIsSmallest(type_t median)
-    {
-        return median;
-    }
-
-    // If median == largest, that implies approximately half the array is equal to largest, unless we were very unlucky with our sample
-    // Try just doing the next smallest value less than this seemingly very common value to seperate them out
-    X86_SIMD_SORT_FINLINE type_t choosePivotMedianIsLargest(type_t median)
-    {
-        return prev_value<type_t>(median);
+        if constexpr (descend) { return prev_value<type_t>(median); }
+        else {
+            return median;
+        }
     }
 };
 
-#endif // XSS_COMMON_COMPARATORS
+#endif // XSS_COMMON_COMPARATORS
diff --git a/src/xss-common-qsort.h b/src/xss-common-qsort.h
@@ -615,8 +615,8 @@ X86_SIMD_SORT_INLINE void xss_qsort(T *arr, arrsize_t arrsize, bool hasnan)
 {
     using comparator =
             typename std::conditional<descending,
-                                      DescendingComparator<vtype>,
-                                      AscendingComparator<vtype>>::type;
+                                      Comparator<vtype, true>,
+                                      Comparator<vtype, false>>::type;
 
     if (arrsize > 1) {
         arrsize_t nan_count = 0;
@@ -641,8 +641,8 @@ xss_qselect(T *arr, arrsize_t k, arrsize_t arrsize, bool hasnan)
 {
     using comparator =
             typename std::conditional<descending,
-                                      DescendingComparator<vtype>,
-                                      AscendingComparator<vtype>>::type;
+                                      Comparator<vtype, true>,
+                                      Comparator<vtype, false>>::type;
 
     arrsize_t index_first_elem = 0;
     arrsize_t index_last_elem = arrsize - 1;
diff --git a/src/xss-pivot-selection.hpp b/src/xss-pivot-selection.hpp
@@ -115,7 +115,7 @@ get_pivot_smart(type_t *arr, const arrsize_t left, const arrsize_t right)
     // Sort the samples
     // Note that this intentionally uses the AscendingComparator
     // instead of the provided comparator
-    sort_vectors<vtype, AscendingComparator<vtype>, numVecs>(vecs);
+    sort_vectors<vtype, Comparator<vtype, false>, numVecs>(vecs);
 
     type_t samples[N];
     for (int i = 0; i < numVecs; i++) {
