Introduction and context
I am working on a small application which involves sorting 20GB files as one of the operations. This is being done, due to memory constraints, by breaking them into ~1GB chunks, sorting in memory and then writing back to disk. Then, as stage 2, reading the head of each file and picking the smallest/largest and streaming to the final output file.
This falls into the domain of External sorting.
This CR will focus on just one small part of that which involves "picking which chunk to take the next element from" when streaming to the final file. It turns out that this part, and specifically the element comparison, is CPU bound and the bottleneck. The elements being sorted are std::array<std::byte, 20>, so a non-trivial comparison, which compiles down to memcmp.
EDIT
It turns out that we can be significantly faster than memcmp with some very simple "big 64-bits load, byte-swap and compares".
See my self answer below END-EDIT
So we need a data structure which makes it easy to pick the next element. Typically there are around 20 chunks, so my first attempt was just using std::min_element because linear search can be really fast for a small number of elements. When I discovered the comparisons were the CPU bottleneck I moved to using a std::priority_queue. This reduced the comparisons by > 2x, so that was a big gain.
For clarification, that reduction in the number of comparisons, also brought a 2x speed up in the full application, from 6mins for the merge stage of the 20 x 1GB files, down to 3mins. But at 3mins, the CPU is still @ 100% and almost all inside the operator< of the 20 byte array.
This CR is about whether we can do better still, by reducing the number of comparisons further.
You cannot change/update elements in a std::priority_queue, and that is exactly what we we need to do. This is the algorithm of the critical loop
std::priority_queue<head, std::vector<head>, decltype(cmp)> heads(cmp);
while (!heads.empty()) {
const head& t = heads.top();
sorted.write(t.value);
std::size_t chunk_idx = t.idx;
heads.pop();
if (auto& chunk = chunks[chunk_idx]; chunk.current != chunk.end) {
heads.push({*(chunk.current), chunk_idx});
++(chunk.current);
}
}
So, it seems very inefficient to call queue.pop() and then queue.push() both of which need to move elements through the heap. It would seem more efficient to just "change" the .top() (having read it and written it to output), to the new value (essentially calling std::next or ++ on the iterator of that chunk) and then "re-heapifying".
Let get under the hood of std::priority_queue. The canonical way to do this using the c++ STL heap algorithms is:
// Note I am using a max-heap, but the comparator is customisable below
std::pop_heap(heap.begin(), heap.end(), comp);
heap.pop_back();
heap.push_back(newval);
std::push_heap(heap.begin(), heap.end(), comp);
Most heap implementations use primitives underneath called something like "bubble/sift up/down". All we actually need is the bubble_down implementation and call that on the changed, top() element.
The c++ standard does not offer primitives, but implementations are not hard and there are 2 (one iterative and one recursive) in the code below.
Investigation Focus
- Relatively small heaps (3 => 100 elements, but easily adjustable in code below)
- Operation: Update the
heap.top()element ONLY - Optimisation criteria = number of comparisons to re-heapify
This is a very simple case. No question of "knowing where the element to be changed is in the heap" - a typical concern with this approach. So surely we can do better, right...?
The code below studies 4 implementations:
- STL way as above (using libstdc++-11). This is the base case. (Note that I believe libc++ is less sophisticated in this area).
- A recursive "bubble down" algo
- An iterative "bubble down" algo
- A libstd++
__adjust_heapfollowed by__push_heapalgo
I generate random heap sizes, contents and replacement value and iterate 1M times.
Findings
- The STL is Good
- The recursive and iterative bubble down alogs are nearly identical (expected, because the compiler, on
-O3, "tailcall optimises" the recursion away) and, surprisingly, consistently have more comparisons than the STL even for this very specialised case. So just using a heap-primitive doesn't give you any gains. - We can beat the STL using... the STL , by copying out the code of the unpublished functions, "cleaning them up" (of underscores etc) and using them in a specialised way to solve this limited, specialised problem. Gains are in the order of 10-20%. A modest, but significant gain?
Typical results:
method avg_cmp_cnt
std::heap / std::priority_queue 7.568285
bubble up recursively 8.031054
bubble up iteratively 8.047352
libstc++ __adjust_heap 6.327297
The code:
#include "fmt/core.h"
#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <cstdlib>
#include <execution>
#include <iostream>
#include <random>
#include <stdexcept>
template <std::unsigned_integral T>
T log2(T x) {
T log = 0;
while (x >>= 1U) ++log;
return log;
}
template <typename T, typename Comp = std::less<>>
void print_heap(const std::vector<T>& heap, Comp comp = {}) {
std::size_t levels = log2(heap.size()) + 1;
unsigned width = 6 * (1U << (levels - 1U));
std::cout << "\n\n";
for (const auto& e: heap) std::cout << e << " ";
std::cout << "\n";
std::cout << fmt::format("is_heap = {:}\n\n", std::is_heap(heap.begin(), heap.end(), comp));
if (heap.empty()) {
std::cout << "<empty heap>\n";
return;
}
unsigned idx = 0;
bool done = false;
for (unsigned l = 0; l != levels; ++l) {
for (unsigned e = 0; e != 1U << l; ++e) {
std::cout << fmt::format("{:^{}}", heap[idx], width);
++idx;
if (idx == heap.size()) {
done = true;
break;
}
}
width /= 2;
std::cout << "\n\n";
if (done) break;
}
}
template <typename T, typename Comp = std::less<>>
void replace_top_using_stl(std::vector<T>& heap, T newval, Comp comp = {}) {
if (heap.empty()) throw std::domain_error("can't replace_top on an empty heap");
assert(std::is_heap(heap.begin(), heap.end(), comp));
std::pop_heap(heap.begin(), heap.end(), comp);
heap.pop_back();
heap.push_back(newval);
std::push_heap(heap.begin(), heap.end(), comp);
}
template <typename T, typename Comp = std::less<>>
// NOLINTNEXTLINE recursion is tailcall eliminated by compiler
void bubble_down_recursively(std::vector<T>& heap, std::size_t i, Comp comp = {}) {
const auto left = 2 * i + 1;
const auto right = 2 * i + 2;
const auto n = heap.size();
using std::swap; // enable ADL
if (left >= n) { // no children
return;
} else if (right >= n) { // left exists right does not. NOLINT else after return
if (comp(heap[i], heap[left])) {
swap(heap[i], heap[left]);
bubble_down_recursively(heap, left, comp);
}
} else { // both children exist
// 'larger' is only well named if comp = std::less<>{}
auto larger = comp(heap[right], heap[left]) ? left : right;
if (comp(heap[i], heap[larger])) {
swap(heap[i], heap[larger]);
bubble_down_recursively(heap, larger, comp);
}
}
}
template <typename T, typename Comp = std::less<>>
void replace_top_using_bubble_down_recursively(std::vector<T>& heap, T newval, Comp comp = {}) {
if (heap.empty()) throw std::domain_error("can't replace_top on an empty heap");
assert(std::is_heap(heap.begin(), heap.end(), comp));
heap[0] = newval;
bubble_down_recursively(heap, 0, comp);
}
template <typename T, typename Comp = std::less<>>
void bubble_down_iteratively(std::vector<T>& heap, std::size_t i, Comp comp = {}) {
const auto n = heap.size();
while (true) {
const std::size_t left = 2 * i + 1;
const std::size_t right = 2 * i + 2;
std::size_t largest = i;
if ((left < n) && comp(heap[largest], heap[left])) {
largest = left;
}
if ((right < n) && comp(heap[largest], heap[right])) {
largest = right;
}
if (largest == i) {
break;
}
using std::swap; // enable ADL
swap(heap[i], heap[largest]);
i = largest;
}
}
template <typename T, typename Comp = std::less<>>
void replace_top_using_bubble_down_iteratively(std::vector<T>& heap, T newval, Comp comp = {}) {
if (heap.empty()) throw std::domain_error("can't replace_top on an empty heap");
assert(std::is_heap(heap.begin(), heap.end(), comp));
heap[0] = newval; // stick it in anyway
bubble_down_iteratively(heap, 0, comp); // and fix the heap
}
// borrowed from libstdc++ __push_heap
template <typename RandomAccessIterator, typename Distance, typename Tp, typename Compare>
constexpr void push_heap(RandomAccessIterator first, Distance holeIndex, Distance topIndex,
Tp value, Compare& comp) {
Distance parent = (holeIndex - 1) / 2;
while (holeIndex > topIndex && comp(*(first + parent), value)) {
*(first + holeIndex) = *(first + parent);
holeIndex = parent;
parent = (holeIndex - 1) / 2;
}
*(first + holeIndex) = std::move(value);
}
// borrowed from libstdc++ __adjust_heap
template <typename RandomAccessIterator, typename Distance, typename Tp, typename Compare>
constexpr void adjust_heap(RandomAccessIterator first, Distance holeIndex, Distance len, Tp value,
Compare comp) {
const Distance topIndex = holeIndex;
Distance secondChild = holeIndex;
while (secondChild < (len - 1) / 2) {
secondChild = 2 * (secondChild + 1);
if (comp(*(first + secondChild), *(first + (secondChild - 1)))) secondChild--;
*(first + holeIndex) = *(first + secondChild);
holeIndex = secondChild;
}
if ((len & 1) == 0 && secondChild == (len - 2) / 2) {
secondChild = 2 * (secondChild + 1);
*(first + holeIndex) = *(first + (secondChild - 1));
holeIndex = secondChild - 1;
}
push_heap(first, holeIndex, topIndex, value, comp);
}
template <typename T, typename Comp = std::less<>>
void replace_top_using_adjust_heap(std::vector<T>& heap, T newval, Comp comp = {}) {
if (heap.empty()) throw std::domain_error("can't replace_top on an empty heap");
assert(std::is_heap(heap.begin(), heap.end(), comp));
heap[0] = newval;
adjust_heap(heap.begin(), 0L, heap.end() - heap.begin(), newval, comp);
}
template <typename T>
struct cmp_counter {
static std::size_t cmpcount; // NOLINT must be static because STL takes Comp by value
bool operator()(T a, T b) {
++cmpcount;
return a < b; // effectively std::less<>{};
}
static void reset() { cmpcount = 0; }
};
template <typename T>
std::size_t cmp_counter<T>::cmpcount = 0; // NOLINT global static
int main() {
using ValueType = int;
struct method {
using cb_t = void (*)(std::vector<ValueType>&, ValueType, cmp_counter<ValueType>);
std::string label;
cb_t cb;
};
auto methods = std::vector<method>{
{"std::heap / std::priority_queue", &replace_top_using_stl},
{"bubble up recursively", &replace_top_using_bubble_down_recursively},
{"bubble up iteratively", &replace_top_using_bubble_down_iteratively},
{"libstc++ __adjust_heap", &replace_top_using_adjust_heap},
};
std::cout << fmt::format("{:35s} {:s}\n", "method", "avg_cmp_cnt");
for (auto& method: methods) {
auto prng = std::mt19937_64(1); // NOLINT fixed seed for repeatability
auto heap_element_dist = std::uniform_int_distribution<>(1, 100);
auto heap_size_dist = std::uniform_int_distribution<std::size_t>(3, 100);
const std::size_t number_of_trials = 1'000'000;
std::size_t total_cmpcount = 0;
cmp_counter<ValueType> comp;
for (unsigned i = 0; i != number_of_trials; ++i) {
std::vector<int> h(heap_size_dist(prng));
std::generate(h.begin(), h.end(), [&] { return ValueType(heap_element_dist(prng)); });
std::make_heap(h.begin(), h.end(), comp);
auto newval = ValueType(heap_element_dist(prng));
cmp_counter<ValueType>::reset();
method.cb(h, newval, comp);
total_cmpcount += cmp_counter<ValueType>::cmpcount;
if (!std::is_heap(h.begin(), h.end(), comp)) {
std::cerr << method.label << "NOT A HEAP ANYMORE!!\n";
return EXIT_FAILURE;
}
}
std::cout << fmt::format("{:35s} {:f}\n", method.label,
double(total_cmpcount) / number_of_trials);
}
}
std::array<std::byte, 20>. I proved that to be the case. I am trying to "reduce the thing that the CPU is spending all its time on". And that for me is the no of comparisons, because I doubt I can improvememcmp. TBH, no of comparisons is also a very standard metric for sorting algorithms? (Note that this code is using ValueType = int, for ease of debugging and printing, which is perhaps misleading. The actual application is 20 byte arrays, as explained). \$\endgroup\$std::min_elementtostd::priority_queueand the number of comparisons more than halved, the wall clock execution time for the merge stage went from 6mins to 3mins. And it's still all at 100% CPU. All inoperator<for the 20 byte array. Unbelievably for me, the bottleneck is NOT the disk for this stage. \$\endgroup\$assert(std::is_heap(..., comp))calls that were also counting, if I add-DNDEBUGto the compiler options I get the same results as you. \$\endgroup\$memcmpcan not optimize memory access with long word reads and must iterate byte by byte in the worst case. (That, however, should not matter, as hash keys should need just one byte comparison in 255/256 cases on average... provided they are quite well distributed.) \$\endgroup\$