代码

#define THREADS 2
#include <pthread.h>
#define HIGH_STAGE_BUFFER 65536
#define TOP_CAPACITY 3296
#define TOP_BUFFER 64
#define sort direct16_unused
#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
    unsigned block = 0;
    for (; block + 32 <= TOP_BUFFER; block += 32)
        stream32x16(destination + block, source + block);
    if (block + 16 <= TOP_BUFFER) {
        _mm256_stream_si256((__m256i *)(destination + block),
            _mm256_load_si256((const __m256i *)(source + block)));
        block += 16;
    }
    if (block + 8 <= TOP_BUFFER) {
        _mm_stream_si128((__m128i *)(destination + block),
            _mm_load_si128((const __m128i *)(source + block)));
        block += 8;
    }
    for (; block < TOP_BUFFER; ++block)
        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

static __attribute__((always_inline)) inline void prefix_u16(
        unsigned short *values) {
    unsigned carry = 0;
    for (unsigned digit = 0; digit < 256; digit += 8) {
        __m128i x = _mm_load_si128((const __m128i *)(values + digit));
        x = _mm_add_epi16(x, _mm_slli_si128(x, 2));
        x = _mm_add_epi16(x, _mm_slli_si128(x, 4));
        x = _mm_add_epi16(x, _mm_slli_si128(x, 8));
        __m128i exclusive = _mm_add_epi16(_mm_slli_si128(x, 2),
                                          _mm_set1_epi16((short)carry));
        _mm_store_si128((__m128i *)(values + digit), exclusive);
        carry += (unsigned short)_mm_extract_epi16(x, 7);
    }
}

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];
        }
#ifdef VECTOR_PREFIX
        prefix_u16(byte_count[0]);
        prefix_u16(byte_count[1]);
#else
        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;
        }
#endif
        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();
}
#undef sort

#ifndef THREADS
#define THREADS 4
#endif

static unsigned top_start[TOP_BUCKETS + 1];
typedef struct __attribute__((aligned(64))) {
    unsigned *a;
    unsigned begin, end;
    unsigned short count[2][256];
    unsigned short scratch[TOP_CAPACITY];
} Worker;
static Worker workers[THREADS];

static void *sort_buckets(void *argument) {
    Worker *worker = argument;
    for (unsigned bucket = worker->begin; bucket < worker->end; ++bucket) {
        unsigned amount = top_count[bucket];
        const unsigned short *source = top_work[bucket];
        unsigned short (*counts)[256] = worker->count;
        __builtin_memset(counts, 0, sizeof(worker->count));
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = source[i];
            ++counts[0][(unsigned char)value];
            ++counts[1][value >> 8];
        }
        unsigned sum0 = 0, sum1 = 0;
        for (unsigned digit = 0; digit < 256; ++digit) {
            unsigned n0 = counts[0][digit], n1 = counts[1][digit];
            counts[0][digit] = (unsigned short)sum0;
            counts[1][digit] = (unsigned short)sum1;
            sum0 += n0; sum1 += n1;
        }
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = source[i];
            worker->scratch[counts[0][(unsigned char)value]++] =
                (unsigned short)value;
        }
        unsigned fixed = bucket << 16, output = top_start[bucket];
        for (unsigned i = 0; i < amount; ++i) {
            unsigned value = worker->scratch[i];
            worker->a[output + counts[1][value >> 8]++] = fixed | value;
        }
    }
    return 0;
}

void sort(u32 *a, int n) {
    for (int i = 0; i < n; ++i) {
        u32 value = a[i];
        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;
    }
    for (unsigned high = 0; high < 256; ++high)
        drain_high(high, high_stage_fill[high]);
    unsigned sum = 0;
    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;
        top_start[bucket] = sum;
        sum += top_count[bucket];
    }
    top_start[TOP_BUCKETS] = sum;
    _mm_sfence();

    pthread_t thread[THREADS - 1];
    unsigned made = 0;
    for (unsigned i = 0; i < THREADS; ++i) {
        workers[i].a = a;
        workers[i].begin = TOP_BUCKETS * i / THREADS;
        workers[i].end = TOP_BUCKETS * (i + 1) / THREADS;
        if (i + 1 < THREADS &&
            pthread_create(&thread[i], 0, sort_buckets, workers + i) == 0)
            ++made;
        else if (i + 1 == THREADS) sort_buckets(workers + i);
        else sort_buckets(workers + i);
    }
    for (unsigned i = 0; i < made; ++i) pthread_join(thread[i], 0);
    _mm256_zeroupper();
}

评测结果