代码

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

typedef unsigned u32;
enum { TOP_BUCKETS = 65536, TOP_CAPACITY = 4096, TOP_BUFFER = 32 };

static unsigned short top_work[TOP_BUCKETS][TOP_CAPACITY]
    __attribute__((aligned(2097152)));
static unsigned short top_buffer[TOP_BUCKETS][TOP_BUFFER]
    __attribute__((aligned(64)));
static unsigned top_count[TOP_BUCKETS] __attribute__((aligned(64)));
static unsigned char top_fill[TOP_BUCKETS] __attribute__((aligned(64)));
static unsigned short byte_count[2][256] __attribute__((aligned(64)));
static unsigned short local_work[8192] __attribute__((aligned(64)));

static __attribute__((always_inline)) inline void stream32x16(
        unsigned short *destination, const unsigned short *source) {
    _mm256_stream_si256((__m256i *)destination,
                        _mm256_load_si256((const __m256i *)source));
    _mm256_stream_si256((__m256i *)(destination + 16),
                        _mm256_load_si256((const __m256i *)(source + 16)));
}

void sort(u32 *a, int n) {
    for (int i = 0; i < n; ++i) {
        u32 value = a[i];
        unsigned bucket = value >> 16;
        unsigned slot = top_fill[bucket];
        top_buffer[bucket][slot] = (unsigned short)value;
        if (slot == TOP_BUFFER - 1) {
            unsigned output = top_count[bucket];
            stream32x16(top_work[bucket] + output, top_buffer[bucket]);
            top_count[bucket] = output + TOP_BUFFER;
            top_fill[bucket] = 0;
        } else {
            top_fill[bucket] = (unsigned char)(slot + 1);
        }
    }
    for (unsigned bucket = 0; bucket < TOP_BUCKETS; ++bucket) {
        unsigned output = top_count[bucket];
        unsigned amount = top_fill[bucket];
        for (unsigned i = 0; i < amount; ++i)
            top_work[bucket][output + i] = top_buffer[bucket][i];
        top_count[bucket] = output + amount;
    }
    _mm_sfence();

    unsigned output_begin = 0;
    for (unsigned bucket = 0; bucket < TOP_BUCKETS; ++bucket) {
        unsigned amount = top_count[bucket];
        const unsigned short *source = top_work[bucket];
        __builtin_memset(byte_count, 0, sizeof(byte_count));
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = source[i];
            ++byte_count[0][(unsigned char)value];
            ++byte_count[1][value >> 8];
        }
        unsigned sum0 = 0, sum1 = 0;
        for (unsigned digit = 0; digit < 256; ++digit) {
            unsigned n0 = byte_count[0][digit];
            unsigned n1 = byte_count[1][digit];
            byte_count[0][digit] = (unsigned short)sum0;
            byte_count[1][digit] = (unsigned short)sum1;
            sum0 += n0;
            sum1 += n1;
        }
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = source[i];
            local_work[byte_count[0][(unsigned char)value]++] =
                (unsigned short)value;
        }
        unsigned fixed = bucket << 16;
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = local_work[i];
            a[output_begin + byte_count[1][value >> 8]++] = fixed | value;
        }
        output_begin += amount;
    }
    _mm256_zeroupper();
}

评测结果