代码

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

typedef unsigned long U;
#define N 200000000U
#define TOP 16U
#define TOP_CAP (((N + TOP - 1) / TOP) + 8192U)
#define FULL_BUFFER 256U
#define PACK_BUFFER 64U

static unsigned full_work[TOP_CAP * TOP] __attribute__((aligned(64)));
static unsigned hist8[256], starts8[257];
static unsigned position[TOP], fill[TOP];
static unsigned full_buffer[TOP][FULL_BUFFER] __attribute__((aligned(64)));
static unsigned pack_buffer[TOP][PACK_BUFFER] __attribute__((aligned(64)));

static inline void full_flush(unsigned *destination, const unsigned *source) {
    for (unsigned i = 0; i < FULL_BUFFER; i += 8)
        _mm256_stream_si256((__m256i *)(destination + i),
                            _mm256_load_si256((const __m256i *)(source + i)));
}

static inline void pack16(__m256i first, __m256i second,
                          __m128i *out0, __m128i *out1, __m128i *out2) {
    const __m128i compact = _mm_setr_epi8(
        0,1,2,4,5,6,8,9,10,12,13,14,-1,-1,-1,-1);
    __m128i p0 = _mm_shuffle_epi8(_mm256_castsi256_si128(first), compact);
    __m128i p1 = _mm_shuffle_epi8(_mm256_extracti128_si256(first,1), compact);
    __m128i p2 = _mm_shuffle_epi8(_mm256_castsi256_si128(second), compact);
    __m128i p3 = _mm_shuffle_epi8(_mm256_extracti128_si256(second,1), compact);
    *out0 = _mm_or_si128(p0, _mm_slli_si128(p1,12));
    *out1 = _mm_or_si128(_mm_srli_si128(p1,4), _mm_slli_si128(p2,8));
    *out2 = _mm_or_si128(_mm_srli_si128(p2,8), _mm_slli_si128(p3,4));
}

static inline void packed_flush(unsigned char *destination,
                                const unsigned *source) {
    __m128i x0,x1,x2,x3,x4,x5;
    pack16(_mm256_load_si256((const __m256i *)(source+0)),
           _mm256_load_si256((const __m256i *)(source+8)),&x0,&x1,&x2);
    pack16(_mm256_load_si256((const __m256i *)(source+16)),
           _mm256_load_si256((const __m256i *)(source+24)),&x3,&x4,&x5);
    _mm256_stream_si256((__m256i *)(destination+0),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x0),x1,1));
    _mm256_stream_si256((__m256i *)(destination+32),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x2),x3,1));
    _mm256_stream_si256((__m256i *)(destination+64),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x4),x5,1));
    pack16(_mm256_load_si256((const __m256i *)(source+32)),
           _mm256_load_si256((const __m256i *)(source+40)),&x0,&x1,&x2);
    pack16(_mm256_load_si256((const __m256i *)(source+48)),
           _mm256_load_si256((const __m256i *)(source+56)),&x3,&x4,&x5);
    _mm256_stream_si256((__m256i *)(destination+96),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x0),x1,1));
    _mm256_stream_si256((__m256i *)(destination+128),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x2),x3,1));
    _mm256_stream_si256((__m256i *)(destination+160),
        _mm256_inserti128_si256(_mm256_castsi128_si256(x4),x5,1));
}

static inline void store24(unsigned char *destination, unsigned value) {
    destination[0]=(unsigned char)value;
    destination[1]=(unsigned char)(value>>8);
    destination[2]=(unsigned char)(value>>16);
}

void sort(unsigned *a, int n) {
    unsigned amount=(unsigned)n;
    __builtin_memset(hist8,0,sizeof(hist8));
    for(unsigned d=0;d<TOP;++d) position[d]=d*TOP_CAP;
    for(unsigned i=0;i<amount;++i){
        unsigned v=a[i],d=v>>28,slot=fill[d];
        ++hist8[v>>24];
        full_buffer[d][slot++]=v;
        if(slot==FULL_BUFFER){full_flush(full_work+position[d],full_buffer[d]);position[d]+=FULL_BUFFER;fill[d]=0;}
        else fill[d]=slot;
    }
    for(unsigned d=0;d<TOP;++d){for(unsigned j=0;j<fill[d];++j)full_work[position[d]+j]=full_buffer[d][j];position[d]+=fill[d];fill[d]=0;}
    _mm_sfence();
    unsigned sum=0;
    for(unsigned d=0;d<256;++d){starts8[d]=sum;sum+=hist8[d];}
    starts8[256]=amount;
    unsigned char *packed=(unsigned char *)a;
    for(unsigned top=0;top<TOP;++top){
        unsigned begin=top*TOP_CAP,end=begin;
        for(unsigned d=0;d<TOP;++d){position[d]=starts8[(top<<4)|d];end+=hist8[(top<<4)|d];}
        for(unsigned i=begin;i<end;++i){
            unsigned v=full_work[i],d=(v>>24)&15,output=position[d];
            unsigned char *destination=packed+3U*output;
            if((U)destination&31U){store24(destination,v);position[d]=output+1;}
            else {unsigned slot=fill[d];pack_buffer[d][slot]=v;if(slot==PACK_BUFFER-1){packed_flush(destination,pack_buffer[d]);position[d]=output+PACK_BUFFER;fill[d]=0;}else fill[d]=slot+1;}
        }
        for(unsigned d=0;d<TOP;++d){unsigned output=position[d];for(unsigned j=0;j<fill[d];++j)store24(packed+3U*(output+j),pack_buffer[d][j]);fill[d]=0;}
    }
    _mm_sfence();
}

评测结果