代码

#include <immintrin.h>

typedef unsigned u32;
typedef unsigned long U;
enum { RAD=256, BUF=16 };
static u32 temp[200000000] __attribute__((aligned(4096)));
static u32 counts[4][RAD] __attribute__((aligned(64)));
static u32 pos[RAD], fill[RAD], limit[RAD];
static u32 buffer[RAD][BUF] __attribute__((aligned(64)));

__attribute__((target("avx2"), always_inline))
static inline void flush_full(u32 *d, const u32 *s) {
    _mm256_stream_si256((__m256i *)d,
                        _mm256_load_si256((const __m256i *)s));
    _mm256_stream_si256((__m256i *)(d+8),
                        _mm256_load_si256((const __m256i *)(s+8)));
}

__attribute__((target("avx2")))
static void scatter(const u32 *src,u32 *dst,int n,int shift,const u32 *cnt){
    u32 sum=0;
    for(int b=0;b<RAD;++b){
        pos[b]=sum;sum+=cnt[b];fill[b]=0;
        u32 head=(u32)(-((U)(dst+pos[b])>>2))&7u;
        limit[b]=head?head:BUF;
    }
    for(int i=0;i<n;++i){
        u32 v=src[i],b=(v>>shift)&255u,f=fill[b];
        buffer[b][f++]=v;
        if(f==limit[b]){
            u32*d=dst+pos[b];
            if(f==BUF)flush_full(d,buffer[b]);
            else for(u32 k=0;k<f;++k)d[k]=buffer[b][k];
            pos[b]+=f;fill[b]=0;limit[b]=BUF;
        }else fill[b]=f;
    }
    for(int b=0;b<RAD;++b){
        u32 f=fill[b],*d=dst+pos[b];
        for(u32 k=0;k<f;++k)d[k]=buffer[b][k];
    }
    _mm_sfence();
}

__attribute__((target("avx2")))
void sort(u32*a,int n){
    for(int b=0;b<RAD;++b)
        counts[0][b]=counts[1][b]=counts[2][b]=counts[3][b]=0;
    for(int i=0;i<n;++i){
        u32 v=a[i];++counts[0][v&255];++counts[1][(v>>8)&255];
        ++counts[2][(v>>16)&255];++counts[3][v>>24];
    }
    scatter(a,temp,n,0,counts[0]);scatter(temp,a,n,8,counts[1]);
    scatter(a,temp,n,16,counts[2]);scatter(temp,a,n,24,counts[3]);
    _mm256_zeroupper();
}

评测结果