代码

#pragma GCC target("avx2")
#pragma GCC optimize("O3", "unroll-loops")
#include <immintrin.h>

typedef unsigned u32;
typedef unsigned long U;

enum { RAD = 65536, BUF = 16 };

static u32 temp[200000000] __attribute__((aligned(2097152)));
static u32 low_count[RAD], high_count[RAD];
static u32 position[RAD];
static unsigned char fill[RAD], limit[RAD];
static u32 buffer[RAD][BUF] __attribute__((aligned(64)));

static inline __attribute__((always_inline))
void stream16(u32 *destination, const u32 *source) {
    _mm256_stream_si256((__m256i *)destination,
                        _mm256_load_si256((const __m256i *)source));
    _mm256_stream_si256((__m256i *)(destination + 8),
                        _mm256_load_si256((const __m256i *)(source + 8)));
}

#define SCATTER(NAME, DIGIT)                                                \
static void NAME(const u32 *source, u32 *destination, int n,               \
                 const u32 *counts) {                                      \
    u32 sum = 0;                                                           \
    for (unsigned bucket = 0; bucket < RAD; ++bucket) {                    \
        position[bucket] = sum;                                            \
        sum += counts[bucket];                                             \
        fill[bucket] = 0;                                                  \
        unsigned head = (unsigned)(-(U)(destination + position[bucket])    \
                                    >> 2) & 7u;                            \
        limit[bucket] = (unsigned char)(head ? head : BUF);                \
    }                                                                      \
    for (int i = 0; i < n; ++i) {                                         \
        u32 value = source[i];                                             \
        unsigned bucket = (DIGIT);                                         \
        unsigned amount = fill[bucket];                                    \
        buffer[bucket][amount++] = value;                                  \
        if (amount == limit[bucket]) {                                     \
            u32 *output = destination + position[bucket];                  \
            if (amount == BUF) stream16(output, buffer[bucket]);           \
            else for (unsigned j = 0; j < amount; ++j)                    \
                output[j] = buffer[bucket][j];                             \
            position[bucket] += amount;                                    \
            fill[bucket] = 0;                                              \
            limit[bucket] = BUF;                                           \
        } else fill[bucket] = (unsigned char)amount;                       \
    }                                                                      \
    for (unsigned bucket = 0; bucket < RAD; ++bucket) {                    \
        unsigned amount = fill[bucket];                                    \
        u32 *output = destination + position[bucket];                      \
        for (unsigned j = 0; j < amount; ++j) output[j] = buffer[bucket][j];\
    }                                                                      \
    _mm_sfence();                                                          \
}

SCATTER(scatter_low, value & 65535u)
SCATTER(scatter_high, value >> 16)

void sort(u32 *a, int n) {
    __builtin_memset(low_count, 0, sizeof(low_count));
    __builtin_memset(high_count, 0, sizeof(high_count));
    for (int i = 0; i < n; ++i) {
        u32 value = a[i];
        ++low_count[value & 65535u];
        ++high_count[value >> 16];
    }
    scatter_low(a, temp, n, low_count);
    scatter_high(temp, a, n, high_count);
    _mm256_zeroupper();
}

评测结果