radix_sort.hpp

/*
	WORK IN PROGRESS: C++ implementation of a 8xW-bit radix sort
	WARNING: DO NOT USE IN PRODUCTION CODE. #lol

	See https://github.com/eloj/radix-sorting

	TODO:
		Do away with shift-table (support any number of passes?)
		Support C-arrays for histograms.
		Add special case for bool.
*/
#pragma once

#include <array>
#include <cinttypes>
#include <cstring> // for std::memcpy

#include "radix_sort_basic_kdf.hpp"

#ifndef RESTRICT
#define RESTRICT __restrict__
#endif

// 8xW-bit Radix Sort
//
// T is the type being sorted.
// KeyFunc is the KDF.
// Hist is storage for the histograms, sized to 256*passes*sizeof(counter-type)
// KeyType is derived from the return value of the KeyFunc (an unsigned integer)
//
template<typename T, typename KeyFunc = decltype(basic_kdfs::kdf<T>), typename Hist, typename KeyType=typename std::result_of_t<KeyFunc&&(T)>>
T* rs_sort_main(T* RESTRICT src, T* RESTRICT aux, size_t n, Hist& histogram, KeyFunc && kf = basic_kdfs::kdf) {
	typedef typename Hist::value_type HVT;
	static_assert(sizeof(KeyType) <= 8, "KeyType must be 64-bits or less");
	static_assert(std::is_unsigned<KeyType>(), "KeyType must be unsigned");

	if (n < 2)
		return src;

	constexpr size_t wc = sizeof(KeyType);
	constexpr std::array<uint8_t, 8> shift_table = { 0, 8, 16, 24, 32, 40, 48, 56 };
	constexpr unsigned int hist_len = 256;
	unsigned int cols[wc];
	unsigned int ncols = 0;
	KeyType key0;

	// Histograms
	size_t n_unsorted = n;
	for (size_t i = 0 ; i < n ; ++i) {
		// pre-sorted detection
		key0 = kf(src[i]);
		if ((i < n - 1) && (key0 <= kf(src[i+1]))) {
			--n_unsorted;
		}
		for (unsigned int j = 0 ; j < wc ; ++j) {
			++histogram[(hist_len*j) + ((key0 >> shift_table[j]) & 0xFF)];
		}
	}

	if (n_unsorted < 2) {
		return src;
	}

	// Sample first key to determine if any columns can be skipped
	key0 = kf(*src);
	for (unsigned int i = 0 ; i < wc ; ++i) {
		if (histogram[(hist_len*i) + ((key0 >> shift_table[i]) & 0xFF)] != n) {
			cols[ncols++] = i;
		}
	}

	// Calculate offsets (exclusive scan)
	for (unsigned int i = 0 ; i < ncols ; ++i) {
		HVT a = 0;
		for (unsigned int j = 0 ; j < hist_len ; ++j) {
			HVT b = histogram[(hist_len*cols[i]) + j];
			histogram[(hist_len*cols[i]) + j] = a;
			a += b;
		}
	}

	// Sort
	for (unsigned int i = 0 ; i < ncols ; ++i) {
		for (size_t j = 0 ; j < n ; ++j) {
			auto k = src[j];
			size_t dst = histogram[(hist_len*cols[i]) + ((kf(k) >> shift_table[cols[i]]) & 0xFF)]++;
			aux[dst] = std::move(k);
		}
		std::swap(src, aux);
	}

	return src;
}

// This version is for automatically selecting the smallest
// possible counter data-type for the histograms.
// Histograms stored on stack (2KiB-16KiB).
template<typename T, typename KeyFunc = decltype(basic_kdfs::kdf<T>), int passes = sizeof(typename std::result_of_t<KeyFunc&&(T)>)>
T* radix_sort(T* RESTRICT src, T* RESTRICT aux, size_t n, KeyFunc && kf = basic_kdfs::kdf) {
	if (n < 2) {
		return src;
	} else if (n < 256) {
		std::array<uint8_t,256*passes> histogram{0};
		return rs_sort_main(src, aux, n, histogram, kf);
	} else if (n < (1ULL << 16ULL)) {
		std::array<uint16_t,256*passes> histogram{0};
		return rs_sort_main(src, aux, n, histogram, kf);
	} else if (n < (1ULL << 32ULL)) {
		std::array<uint32_t,256*passes> histogram{0};
		return rs_sort_main(src, aux, n, histogram, kf);
	} else {
		std::array<uint64_t,256*passes> histogram{0};
		return rs_sort_main(src, aux, n, histogram, kf);
	}
}