Minor cleanup changes

Author: sterrettm2

src/avx2-32bit-half.hpp

@@ -29,37 +29,26 @@
 template <typename vtype, typename reg_t = typename vtype::reg_t>
 X86_SIMD_SORT_INLINE reg_t sort_ymm_32bit_half(reg_t ymm)
 {
-    //static_assert(vtype::numlanes == 0, "This function is not implemented");
-    typename vtype::type_t buffer[vtype::numlanes];
-    vtype::storeu(buffer, ymm);
-    std::sort(&buffer[0], &buffer[vtype::numlanes], comparison_func<vtype>);
-    return vtype::loadu(buffer);
-    /*
-    const typename vtype::opmask_t oxAA = _mm256_set_epi32(
-            0xFFFFFFFF, 0, 0xFFFFFFFF, 0, 0xFFFFFFFF, 0, 0xFFFFFFFF, 0);
-    const typename vtype::opmask_t oxCC = _mm256_set_epi32(
-            0xFFFFFFFF, 0xFFFFFFFF, 0, 0, 0xFFFFFFFF, 0xFFFFFFFF, 0, 0);
-    const typename vtype::opmask_t oxF0 = _mm256_set_epi32(
-            0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0, 0, 0, 0);
-
-    const typename vtype::ymmi_t rev_index = vtype::seti(NETWORK_32BIT_AVX2_2);
-    ymm = cmp_merge<vtype>(
-            ymm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(ymm), oxAA);
+    using swizzle = typename vtype::swizzle_ops;
+    
+    const typename vtype::opmask_t oxAA
+            = vtype::seti(-1, 0, -1, 0);
+    const typename vtype::opmask_t oxCC
+            = vtype::seti(-1, -1, 0, 0);
+            
     ymm = cmp_merge<vtype>(
             ymm,
-            vtype::permutexvar(vtype::seti(NETWORK_32BIT_AVX2_1), ymm),
-            oxCC);
-    ymm = cmp_merge<vtype>(
-            ymm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(ymm), oxAA);
-    ymm = cmp_merge<vtype>(ymm, vtype::permutexvar(rev_index, ymm), oxF0);
+            swizzle::template swap_n<vtype, 2>(ymm),
+            oxAA);
     ymm = cmp_merge<vtype>(
             ymm,
-            vtype::permutexvar(vtype::seti(NETWORK_32BIT_AVX2_3), ymm),
+            vtype::reverse(ymm),
             oxCC);
     ymm = cmp_merge<vtype>(
-            ymm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(ymm), oxAA);
+            ymm,
+            swizzle::template swap_n<vtype, 2>(ymm),
+            oxAA);
     return ymm;
-    */
 }
 
 struct avx2_32bit_half_swizzle_ops;

src/avx512-64bit-argsort.hpp

@@ -8,14 +8,9 @@
 #define AVX512_ARGSORT_64BIT
 
 #include "xss-common-qsort.h"
-//#include "avx512-64bit-common.h"
-//#include "avx2-32bit-half.hpp"
 #include "xss-network-keyvaluesort.hpp"
 #include <numeric>
 
-template <typename T>
-struct avx2_half_vector;
-
 template <typename T>
 X86_SIMD_SORT_INLINE void std_argselect_withnan(
         T *arr, arrsize_t *arg, arrsize_t k, arrsize_t left, arrsize_t right)
@@ -146,9 +141,9 @@ X86_SIMD_SORT_INLINE int32_t partition_vec(type_t *arg,
                                            reg_t *smallest_vec,
                                            reg_t *biggest_vec)
 {
-    if constexpr (sizeof (argreg_t) == 64){
+    if constexpr (vtype::vec_type == simd_type::AVX512){
         return partition_vec_avx512<vtype, argtype, type_t>(arg, left, right, arg_vec, curr_vec, pivot_vec, smallest_vec, biggest_vec);
-    }else if constexpr (sizeof (argreg_t) == 32){
+    }else if constexpr (vtype::vec_type == simd_type::AVX2){
         return partition_vec_avx2<vtype, argtype, type_t>(arg, left, right, arg_vec, curr_vec, pivot_vec, smallest_vec, biggest_vec);
     }else{
         static_assert(sizeof(argreg_t) == 0, "Should not get here");

src/xss-network-keyvaluesort.hpp

@@ -362,6 +362,16 @@ X86_SIMD_SORT_INLINE void argsort_n_vec(typename keyType::type_t *keys,
 
     kreg_t keyVecs[numVecs];
     ireg_t indexVecs[numVecs];
+    
+    // Generate masks for loading and storing
+    typename keyType::opmask_t ioMasks[numVecs - numVecs / 2];
+    X86_SIMD_SORT_UNROLL_LOOP(64)
+    for (int i = numVecs / 2, j = 0; i < numVecs; i++, j++) {
+        uint64_t num_to_read
+                = std::min((uint64_t)std::max(0, N - i * keyType::numlanes),
+                           (uint64_t)keyType::numlanes);
+        ioMasks[j] = keyType::get_partial_loadmask(num_to_read);
+    }
 
     // Unmasked part of the load
     X86_SIMD_SORT_UNROLL_LOOP(64)
@@ -373,20 +383,13 @@ X86_SIMD_SORT_INLINE void argsort_n_vec(typename keyType::type_t *keys,
     // Masked part of the load
     X86_SIMD_SORT_UNROLL_LOOP(64)
     for (int i = numVecs / 2; i < numVecs; i++) {
-        uint64_t num_to_read
-                = std::min((uint64_t)std::max(0, N - i * keyType::numlanes),
-                           (uint64_t)keyType::numlanes);
-       
-        auto indexMask = indexType::get_partial_loadmask(num_to_read);
-        auto keyMask = keyType::get_partial_loadmask(num_to_read);
-        
         indexVecs[i] = indexType::mask_loadu(indexType::zmm_max(),
-                                             indexMask,
+                                             extend_mask<keyType, indexType>(ioMasks[i - numVecs/2]),
                                              indices + i * indexType::numlanes);
 
         keyVecs[i] = keyType::template mask_i64gather<sizeof(
                 typename keyType::type_t)>(
-                keyType::zmm_max(), keyMask, indexVecs[i], keys);
+                keyType::zmm_max(), ioMasks[i - numVecs / 2], indexVecs[i], keys);
     }
 
     // Sort each loaded vector
@@ -406,13 +409,8 @@ X86_SIMD_SORT_INLINE void argsort_n_vec(typename keyType::type_t *keys,
     // Masked part of the store
     X86_SIMD_SORT_UNROLL_LOOP(64)
     for (int i = numVecs / 2, j = 0; i < numVecs; i++, j++) {
-        uint64_t num_to_read
-                = std::min((uint64_t)std::max(0, N - i * keyType::numlanes),
-                           (uint64_t)keyType::numlanes);
-        
-        auto indexMask = indexType::get_partial_loadmask(num_to_read);
         indexType::mask_storeu(
-                indices + i * indexType::numlanes, indexMask, indexVecs[i]);
+                indices + i * indexType::numlanes, extend_mask<keyType, indexType>(ioMasks[i - numVecs/2]), indexVecs[i]);
     }
 }