Try to improve radix sort with memory prefetching

src/Common/RadixSort.h

@@ -1,6 +1,5 @@
 #pragma once
 
-
 #include <string.h>
 #if !defined(OS_DARWIN) && !defined(OS_FREEBSD)
 #include <malloc.h>
@@ -13,11 +12,11 @@
 #include <type_traits>
 #include <memory>
 
+#include <Core/Defines.h>
 #include <base/bit_cast.h>
 #include <base/extended_types.h>
 #include <base/sort.h>
-#include <Core/Defines.h>
-
+#include <Common/PODArray.h>
 
 /** Radix sort, has the following functionality:
   *
@@ -225,6 +224,7 @@ struct RadixSort
     static constexpr size_t PART_BITMASK = HISTOGRAM_SIZE - 1;
     static constexpr size_t KEY_BITS = sizeof(Key) * 8;
     static constexpr size_t NUM_PASSES = (KEY_BITS + (Traits::PART_SIZE_BITS - 1)) / Traits::PART_SIZE_BITS;
+    static constexpr size_t PREFETCH_DISTANCE = std::max(1UL, 64 / sizeof(Element) / 2);
 
 
     static KeyBits keyToBits(Key x) { return bit_cast<KeyBits>(x); }
@@ -276,8 +276,18 @@ struct RadixSort
         }
     }
 
-
     template <bool DIRECT_WRITE_TO_DESTINATION>
+    static NO_INLINE void radixSortLSDInternalHelper(Element * arr, size_t size, bool reverse, Result * destination)
+    {
+        if (size <= std::numeric_limits<UInt8>::max())
+            radixSortLSDInternal<DIRECT_WRITE_TO_DESTINATION, UInt8>(arr, size, reverse, destination);
+        else if (size <= std::numeric_limits<UInt16>::max())
+            radixSortLSDInternal<DIRECT_WRITE_TO_DESTINATION, UInt16>(arr, size, reverse, destination);
+        else
+            radixSortLSDInternal<DIRECT_WRITE_TO_DESTINATION, UInt32>(arr, size, reverse, destination);
+    }
+
+    template <bool DIRECT_WRITE_TO_DESTINATION, typename CountType>
     static NO_INLINE void radixSortLSDInternal(Element * arr, size_t size, bool reverse, Result * destination)
     {
         /// If the array is smaller than 256, then it is better to use another algorithm.
@@ -304,15 +314,15 @@ struct RadixSort
         }
 
         {
-            /// Replace the histograms with the accumulated sums: the value in position i is the sum of the previous positions minus one.
+            /// Replace the histograms with the accumulated sums: the value in position i is the sum of the previous positions.
             CountType sums[NUM_PASSES] = {0};
 
             for (size_t i = 0; i < HISTOGRAM_SIZE; ++i)
             {
                 for (size_t pass = 0; pass < NUM_PASSES; ++pass)
                 {
                     CountType tmp = histograms[pass * HISTOGRAM_SIZE + i] + sums[pass];
-                    histograms[pass * HISTOGRAM_SIZE + i] = sums[pass] - 1;
+                    histograms[pass * HISTOGRAM_SIZE + i] = sums[pass];
                     sums[pass] = tmp;
                 }
             }
@@ -326,15 +336,23 @@ struct RadixSort
 
             for (size_t i = 0; i < size; ++i)
             {
-                size_t pos = extractPart(pass, reader[i]);
+                auto element = reader[i];
+                size_t pos = extractPart(pass, element);
+                if (i + PREFETCH_DISTANCE < size) [[likely]]
+                {
+                    size_t next_pos = extractPart(pass, reader[i + PREFETCH_DISTANCE]);
+                    __builtin_prefetch(&writer[histograms[pass * HISTOGRAM_SIZE + next_pos]], 1);
+                }
 
                 /// Place the element on the next free position.
-                auto & dest = writer[++histograms[pass * HISTOGRAM_SIZE + pos]];
-                dest = reader[i];
+                auto & dest = writer[histograms[pass * HISTOGRAM_SIZE + pos]];
+                dest = element;
 
                 /// On the last pass, we do the reverse transformation.
                 if (!Traits::Transform::transform_is_simple && pass == NUM_PASSES - 1)
-                    Traits::extractKey(dest) = bitsToKey(Traits::Transform::backward(keyToBits(Traits::extractKey(reader[i]))));
+                    Traits::extractKey(dest) = bitsToKey(Traits::Transform::backward(keyToBits(Traits::extractKey(element))));
+
+                histograms[pass * HISTOGRAM_SIZE + pos]++;
             }
         }
 
@@ -348,16 +366,30 @@ struct RadixSort
             {
                 for (size_t i = 0; i < size; ++i)
                 {
-                    size_t pos = extractPart(pass, reader[i]);
-                    writer[size - 1 - (++histograms[pass * HISTOGRAM_SIZE + pos])] = Traits::extractResult(reader[i]);
+                    auto element = reader[i];
+                    size_t pos = extractPart(pass, element);
+                    if (i + PREFETCH_DISTANCE < size) [[likely]]
+                    {
+                        size_t next_pos = extractPart(pass, reader[i + PREFETCH_DISTANCE]);
+                        __builtin_prefetch(&writer[size - 1 - histograms[pass * HISTOGRAM_SIZE + next_pos]], 1);
+                    }
+
+                    writer[size - 1 - (histograms[pass * HISTOGRAM_SIZE + pos]++)] = Traits::extractResult(element);
                 }
             }
             else
             {
                 for (size_t i = 0; i < size; ++i)
                 {
-                    size_t pos = extractPart(pass, reader[i]);
-                    writer[++histograms[pass * HISTOGRAM_SIZE + pos]] = Traits::extractResult(reader[i]);
+                    auto element = reader[i];
+                    size_t pos = extractPart(pass, element);
+                    if (i + PREFETCH_DISTANCE < size)
+                    {
+                        size_t next_pos = extractPart(pass, reader[i + PREFETCH_DISTANCE]);
+                        __builtin_prefetch(&writer[histograms[pass * HISTOGRAM_SIZE + next_pos]], 1);
+                    }
+
+                    writer[histograms[pass * HISTOGRAM_SIZE + pos]++] = Traits::extractResult(element);
                 }
             }
         }
@@ -551,7 +583,7 @@ struct RadixSort
             return;
         }
 
-        radixSortLSDInternal<DIRECT_WRITE_TO_DESTINATION>(arr, size, reverse, destination);
+        radixSortLSDInternalHelper<DIRECT_WRITE_TO_DESTINATION>(arr, size, reverse, destination);
     }
 
 public:
@@ -560,12 +592,12 @@ struct RadixSort
       */
     static void executeLSD(Element * arr, size_t size)
     {
-        radixSortLSDInternal<false>(arr, size, false, nullptr);
+        radixSortLSDInternalHelper<false>(arr, size, false, nullptr);
     }
 
     static void executeLSD(Element * arr, size_t size, bool reverse)
     {
-        radixSortLSDInternal<false>(arr, size, reverse, nullptr);
+        radixSortLSDInternalHelper<false>(arr, size, reverse, nullptr);
     }
 
     /** This function will start to sort inplace (modify 'arr')
@@ -576,7 +608,7 @@ struct RadixSort
       */
     static void executeLSD(Element * arr, size_t size, bool reverse, Result * destination)
     {
-        radixSortLSDInternal<true>(arr, size, reverse, destination);
+        radixSortLSDInternalHelper<true>(arr, size, reverse, destination);
     }
 
     /** Tries to fast sort elements for common sorting patterns (unstable).

src/Common/benchmarks/CMakeLists.txt

@@ -13,6 +13,11 @@ target_link_libraries (orc_string_dictionary PRIVATE
     ch_contrib::gbenchmark_all
     dbms)
 
+clickhouse_add_executable(benchmark_radix_sort radix_sort.cpp)
+target_link_libraries (benchmark_radix_sort PRIVATE
+    ch_contrib::gbenchmark_all
+    dbms)
+
 clickhouse_add_executable(wrap_in_nullable wrap_in_nullable.cpp)
 target_link_libraries (wrap_in_nullable PRIVATE
     ch_contrib::gbenchmark_all

src/Common/benchmarks/radix_sort.cpp

@@ -0,0 +1,72 @@
+#include <DataTypes/DataTypesNumber.h>
+#include <Functions/FunctionHelpers.h>
+#include <Storages/StorageGenerateRandom.h>
+#include <benchmark/benchmark.h>
+#include "pcg_random.hpp"
+
+using namespace DB;
+
+static void BM_RadixSort_UInt8(benchmark::State & state)
+{
+    pcg64 rng;
+    UInt64 limit = DEFAULT_BLOCK_SIZE;
+    auto type = std::make_shared<DataTypeUInt8>();
+    auto column = fillColumnWithRandomData(type, limit, 0, 0, rng, nullptr);
+
+    for (auto _ : state)
+    {
+        IColumn::Permutation res;
+        column->getPermutation(IColumn::PermutationSortDirection::Ascending, IColumn::PermutationSortStability::Unstable, 0, 0, res);
+        benchmark::DoNotOptimize(res);
+    }
+}
+
+static void BM_RadixSort_Int16(benchmark::State & state)
+{
+    pcg64 rng;
+    UInt64 limit = DEFAULT_BLOCK_SIZE;
+    auto type = std::make_shared<DataTypeInt16>();
+    auto column = fillColumnWithRandomData(type, limit, 0, 0, rng, nullptr);
+
+    for (auto _ : state)
+    {
+        IColumn::Permutation res;
+        column->getPermutation(IColumn::PermutationSortDirection::Ascending, IColumn::PermutationSortStability::Unstable, 0, 0, res);
+        benchmark::DoNotOptimize(res);
+    }
+}
+
+static void BM_RadixSort_Int32(benchmark::State & state)
+{
+    pcg64 rng;
+    UInt64 limit = DEFAULT_BLOCK_SIZE;
+    auto type = std::make_shared<DataTypeInt32>();
+    auto column = fillColumnWithRandomData(type, limit, 0, 0, rng, nullptr);
+
+    for (auto _ : state)
+    {
+        IColumn::Permutation res;
+        column->getPermutation(IColumn::PermutationSortDirection::Ascending, IColumn::PermutationSortStability::Unstable, 0, 0, res);
+        benchmark::DoNotOptimize(res);
+    }
+}
+
+static void BM_RadixSort_UInt64(benchmark::State & state)
+{
+    pcg64 rng;
+    UInt64 limit = DEFAULT_BLOCK_SIZE;
+    auto type = std::make_shared<DataTypeUInt64>();
+    auto column = fillColumnWithRandomData(type, limit, 0, 0, rng, nullptr);
+
+    for (auto _ : state)
+    {
+        IColumn::Permutation res;
+        column->getPermutation(IColumn::PermutationSortDirection::Ascending, IColumn::PermutationSortStability::Unstable, 0, 0, res);
+        benchmark::DoNotOptimize(res);
+    }
+}
+
+BENCHMARK(BM_RadixSort_UInt8);
+BENCHMARK(BM_RadixSort_Int16);
+BENCHMARK(BM_RadixSort_Int32);
+BENCHMARK(BM_RadixSort_UInt64);