Added records sorting. (2 columns)

Author: PatwinchIR

CMakeLists.txt

@@ -13,7 +13,7 @@ include_directories(third_party/googletest/googlemock/include)
 include_directories(third_party/ips4o)
 include_directories(third_party/pdqsort)
 
-set(CMAKE_CXX_FLAGS "-g -O3 -flto -Wall -march=native")
+set(CMAKE_CXX_FLAGS "-g -O3 -flto -Wall -march=native -fopenmp")
 
 file(GLOB_RECURSE SOURCE_FILES
         "src/*.cpp" "test/*.cpp")

src/avx256/merge_util.cpp

@@ -59,12 +59,13 @@ template void MaskedMergeRuns4<double, __m256d>(double *&arr, size_t N);
 template<typename InType, typename RegType>
 void MergePass8(InType *&arr, InType *buffer, size_t N, unsigned int run_size) {
   int UNIT_RUN_SIZE = 8;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -119,12 +120,13 @@ template void MergePass8<float, __m256>(float *&arr, float *buffer, size_t N, un
 template<typename InType, typename RegType>
 void MaskedMergePass8(InType *&arr, InType *buffer, size_t N, unsigned int run_size) {
   int UNIT_RUN_SIZE = 8;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -179,12 +181,13 @@ template void MaskedMergePass8<float, __m256>(float *&arr, float *buffer, size_t
 template<typename InType, typename RegType>
 void MergePass4(InType *&arr, InType *buffer, size_t N, unsigned int run_size) {
   int UNIT_RUN_SIZE = 4;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -238,12 +241,13 @@ template void MergePass4<double, __m256d>(double *&arr, double *buffer, size_t N
 template<typename InType, typename RegType>
 void MaskedMergePass4(InType *&arr, InType *buffer, size_t N, unsigned int run_size) {
   int UNIT_RUN_SIZE = 4;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);

src/avx256/simd_sort.cpp

@@ -69,6 +69,21 @@ void SIMDSort(size_t N, std::pair<int, int> *&arr) {
   }
 }
 
+void SIMDOrderBy(std::pair<int, int> *&result_arr, size_t N, std::pair<int, int> *arr, int order_by) {
+  int64_t *kv_arr;
+  aligned_init<int64_t>(kv_arr, N);
+  aligned_init<std::pair<int, int>>(result_arr, N);
+  for (int i = 0; i < N; ++i) {
+    auto value = (int64_t) (order_by == 0 ? arr[i].first : arr[i].second);
+    kv_arr[i] = (((value) << 32) | (0x00000000ffffffff & i));
+  }
+  SIMDSort(N, kv_arr);
+  for (int j = 0; j < N; ++j) {
+    auto index = 0x00000000ffffffff & kv_arr[j];
+    result_arr[j] = arr[index];
+  }
+}
+
 void SIMDSort(size_t N, std::pair<float, float> *&arr) {
   float *kv_arr;
   size_t Nkv = N * 2;

src/avx512/merge_util.cpp

@@ -59,12 +59,13 @@ template void MaskedMergeRuns8<double, __m512d>(double *&arr, size_t N);
 template<typename InType, typename RegType>
 void MergePass16(InType *&arr, InType *buffer, size_t N, int run_size) {
   int UNIT_RUN_SIZE = 16;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -119,12 +120,13 @@ template void MergePass16<float, __m512>(float *&arr, float *buffer, size_t N, i
 template<typename InType, typename RegType>
 void MaskedMergePass16(InType *&arr, InType *buffer, size_t N, int run_size) {
   int UNIT_RUN_SIZE = 16;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -179,12 +181,13 @@ template void MaskedMergePass16<float, __m512>(float *&arr, float *buffer, size_
 template<typename InType, typename RegType>
 void MergePass8(InType *&arr, InType *buffer, size_t N, int run_size) {
   int UNIT_RUN_SIZE = 8;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);
@@ -239,12 +242,13 @@ template void MergePass8<double, __m512d>(double *&arr, double *buffer, size_t N
 template<typename InType, typename RegType>
 void MaskedMergePass8(InType *&arr, InType *buffer, size_t N, int run_size) {
   int UNIT_RUN_SIZE = 8;
-  RegType ra, rb;
-  int buffer_offset = 0;
+#pragma omp parallel for
   for (int i = 0; i < N; i += 2 * run_size) {
     int start = i;
     int mid = i + run_size;
     int end = i + 2 * run_size;
+    int buffer_offset = start;
+    RegType ra, rb;
     int p1_ptr = start;
     int p2_ptr = mid;
     LoadReg(ra, &arr[p1_ptr]);

src/avx512/simd_sort.cpp

@@ -61,6 +61,21 @@ void SIMDSort(size_t N, std::pair<int, int> *&arr) {
   }
 }
 
+void SIMDOrderBy(std::pair<int, int> *&result_arr, size_t N, std::pair<int, int> *arr, int order_by) {
+  int64_t *kv_arr;
+  aligned_init<int64_t>(kv_arr, N);
+  aligned_init<std::pair<int, int>>(result_arr, N);
+  for (int i = 0; i < N; ++i) {
+    auto value = (int64_t) (order_by == 0 ? arr[i].first : arr[i].second);
+    kv_arr[i] = (((value) << 32) | (0x00000000ffffffff & i));
+  }
+  SIMDSort(N, kv_arr);
+  for (int j = 0; j < N; ++j) {
+    auto index = 0x00000000ffffffff & kv_arr[j];
+    result_arr[j] = arr[index];
+  }
+}
+
 void SIMDSort(size_t N, std::pair<float, float> *&arr) {
   float *kv_arr;
   size_t Nkv = N * 2;

src/include/avx256/simd_sort.h

@@ -12,6 +12,7 @@ namespace avx2{
   void SIMDSort(size_t N, float *&arr);
   void SIMDSort(size_t N, double *&arr);
   void SIMDSort(size_t N, std::pair<int,int> *&arr);
+  void SIMDOrderBy(std::pair<int, int> *&result_arr, size_t N, std::pair<int, int> *arr, int order_by=0);
   void SIMDSort(size_t N, std::pair<float, float> *&arr);
   void SIMDSort(size_t N, std::pair<int64_t ,int64_t> *&arr);
   void SIMDSort(size_t N, std::pair<double, double> *&arr);

src/include/avx512/simd_sort.h

@@ -12,6 +12,7 @@ namespace avx512{
   void SIMDSort(size_t N, float *&arr);
   void SIMDSort(size_t N, double *&arr);
   void SIMDSort(size_t N, std::pair<int,int> *&arr);
+  void SIMDOrderBy(std::pair<int, int> *&result_arr, size_t N, std::pair<int, int> *arr, int order_by=0);
   void SIMDSort(size_t N, std::pair<float, float> *&arr);
   void SIMDSort(size_t N, std::pair<int64_t ,int64_t> *&arr);
   void SIMDSort(size_t N, std::pair<double, double> *&arr);

src/include/common.h

@@ -5,6 +5,7 @@
 #include <string>
 #include <cstdint>
 #include <cassert>
+#include <omp.h>
 
 /**
  * Common definitions

test/avx256/simd_sort_test.cpp

@@ -345,6 +345,49 @@ TEST(SIMDSortTests, AVX256SIMDSort32BitKeyValueIntTest) {
   delete soln_arr;
 }
 
+TEST(SIMDSortTests, AVX256SIMDOrderBy32BitIntTest) {
+  using T = int;
+  size_t N = NNUM;
+  T lo = LO;
+  T hi = HI;
+  std::pair<T, T> *rand_arr;
+  std::pair<T, T> *soln_arr1, *soln_arr2, *input_arr1, *input_arr2;
+  double start, end;
+
+  // Initialization
+  TestUtil::RandGenIntEntries(rand_arr, N, lo, hi);
+
+  aligned_init<std::pair<T, T>>(input_arr1, N);
+  aligned_init<std::pair<T, T>>(soln_arr1, N);
+  std::copy(rand_arr, rand_arr + N, input_arr1);
+  std::vector<std::pair<T, T>> check_arr1(rand_arr, rand_arr + N);
+  start = currentSeconds();
+  SIMDOrderBy(soln_arr1, N, input_arr1);
+  end = currentSeconds();
+  std::sort(check_arr1.begin(), check_arr1.end(), [](std::pair<T, T> &left, std::pair<T, T> &right) {
+    return left.first < right.first;
+  });
+  for (int i = 0; i < N; i++) {
+    EXPECT_EQ(check_arr1[i].first, soln_arr1[i].first);
+  }
+
+  aligned_init<std::pair<T, T>>(input_arr2, N);
+  aligned_init<std::pair<T, T>>(soln_arr2, N);
+  std::copy(rand_arr, rand_arr + N, input_arr2);
+  std::vector<std::pair<T, T>> check_arr2(rand_arr, rand_arr + N);
+  SIMDOrderBy(soln_arr2, N, input_arr2, 1);
+  std::sort(check_arr2.begin(), check_arr2.end(), [](std::pair<T, T> &left, std::pair<T, T> &right) {
+    return left.second < right.second;
+  });
+  for (int i = 0; i < N; i++) {
+    EXPECT_EQ(check_arr2[i].second, soln_arr2[i].second);
+  }
+  printf("[avx256::orderby] %lu elements: %.8f seconds\n", N, end - start);
+  delete rand_arr;
+  delete soln_arr1;
+  delete soln_arr2;
+}
+
 TEST(SIMDSortTests, AVX256SIMDSort64BitKeyValueIntTest) {
   using T = int64_t;
   size_t N = NNUM;

test/avx512/simd_sort_test.cpp

@@ -347,6 +347,49 @@ TEST(SIMDSortTests, AVX512SIMDSort32BitKeyValueIntTest) {
   delete soln_arr;
 }
 
+TEST(SIMDSortTests, AVX256SIMDOrderBy32BitIntTest) {
+  using T = int;
+  size_t N = NNUM;
+  T lo = LO;
+  T hi = HI;
+  std::pair<T, T> *rand_arr;
+  std::pair<T, T> *soln_arr1, *soln_arr2, *input_arr1, *input_arr2;
+  double start, end;
+
+  // Initialization
+  TestUtil::RandGenIntEntries(rand_arr, N, lo, hi);
+
+  aligned_init<std::pair<T, T>>(input_arr1, N);
+  aligned_init<std::pair<T, T>>(soln_arr1, N);
+  std::copy(rand_arr, rand_arr + N, input_arr1);
+  std::vector<std::pair<T, T>> check_arr1(rand_arr, rand_arr + N);
+  start = currentSeconds();
+  SIMDOrderBy(soln_arr1, N, input_arr1);
+  end = currentSeconds();
+  std::sort(check_arr1.begin(), check_arr1.end(), [](std::pair<T, T> &left, std::pair<T, T> &right) {
+    return left.first < right.first;
+  });
+  for (int i = 0; i < N; i++) {
+    EXPECT_EQ(check_arr1[i].first, soln_arr1[i].first);
+  }
+
+  aligned_init<std::pair<T, T>>(input_arr2, N);
+  aligned_init<std::pair<T, T>>(soln_arr2, N);
+  std::copy(rand_arr, rand_arr + N, input_arr2);
+  std::vector<std::pair<T, T>> check_arr2(rand_arr, rand_arr + N);
+  SIMDOrderBy(soln_arr2, N, input_arr2, 1);
+  std::sort(check_arr2.begin(), check_arr2.end(), [](std::pair<T, T> &left, std::pair<T, T> &right) {
+    return left.second < right.second;
+  });
+  for (int i = 0; i < N; i++) {
+    EXPECT_EQ(check_arr2[i].second, soln_arr2[i].second);
+  }
+  printf("[avx256::orderby] %lu elements: %.8f seconds\n", N, end - start);
+  delete rand_arr;
+  delete soln_arr1;
+  delete soln_arr2;
+}
+
 TEST(SIMDSortTests, AVX512SIMDSort64BitKeyValueIntTest) {
   using T = int64_t;
   size_t N = NNUM;

test/include/test_util.h

@@ -58,6 +58,18 @@ struct TestUtil{
     }
   }
 
+  template <typename T>
+  static void RandGenIntEntries(std::pair<T, T>* &arr, size_t N, T lo, T hi, unsigned int offset_start=0) {
+    aligned_init<std::pair<T, T>>(arr, N);
+    std::random_device rd;  //Will be used to obtain a seed for the random number engine
+    std::mt19937 gen(rd()); //Standard mersenne_twister_engine seeded with rd()
+    std::uniform_int_distribution<T> dis(lo, hi);
+    for(size_t i = 0; i < N; i++) {
+      arr[i].first = dis(gen);
+      arr[i].second = dis(gen);
+    }
+  }
+
   template <typename T>
   static void RandGenFloatRecords(std::pair<T, T>* &arr, size_t N, T lo, T hi, unsigned int offset_start=0) {
     aligned_init<std::pair<T, T>>(arr, N);