代码

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

typedef unsigned u32;
#define TOP_BUCKETS 65536
#ifndef TOP_CAPACITY
#define TOP_CAPACITY 4096
#endif
#ifndef TOP_BUFFER
#define TOP_BUFFER 32
#endif

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)));
#ifdef HIGH_STAGE_BUFFER
static u32 high_stage[256][HIGH_STAGE_BUFFER] __attribute__((aligned(64)));
static unsigned high_stage_fill[256];
#endif

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)));
}

static __attribute__((always_inline)) inline void flush_top(
        unsigned short *destination, const unsigned short *source) {
#if TOP_BUFFER == 8
    _mm_stream_si128((__m128i *)destination,
                     _mm_load_si128((const __m128i *)source));
#elif TOP_BUFFER == 16
    _mm256_stream_si256((__m256i *)destination,
                        _mm256_load_si256((const __m256i *)source));
#else
    for (unsigned block = 0; block < TOP_BUFFER; block += 32)
        stream32x16(destination + block, source + block);
#endif
}

static __attribute__((always_inline)) inline void emit_top(u32 value) {
    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];
        flush_top(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);
    }
}

#ifdef HIGH_STAGE_BUFFER
static __attribute__((always_inline)) inline void drain_high(
        unsigned high, unsigned amount) {
    const u32 *source = high_stage[high];
    for (unsigned i = 0; i < amount; ++i)
        emit_top(source[i]);
}
#endif

void sort(u32 *a, int n) {
#ifdef HIGH_STAGE_BUFFER
#endif
    for (int i = 0; i < n; ++i) {
#ifdef TOP_PREFETCH
        if (i + TOP_PREFETCH < n) {
            unsigned future = a[i + TOP_PREFETCH] >> 16;
            __builtin_prefetch(top_buffer[future], 1, 1);
        }
#endif
        u32 value = a[i];
#ifdef HIGH_STAGE_BUFFER
        unsigned high = value >> 24;
        unsigned slot = high_stage_fill[high];
        high_stage[high][slot] = value;
        if (slot == HIGH_STAGE_BUFFER - 1) {
            drain_high(high, HIGH_STAGE_BUFFER);
            high_stage_fill[high] = 0;
        } else high_stage_fill[high] = slot + 1;
#else
        emit_top(value);
#endif
    }
#ifdef HIGH_STAGE_BUFFER
    for (unsigned high = 0; high < 256; ++high)
        drain_high(high, high_stage_fill[high]);
#endif
    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();
}

评测结果