代码

#include <immintrin.h>
#include <string.h>
#include <sys/mman.h>

typedef unsigned u32;
typedef unsigned long U;
enum { RAD=2048, BUF=16 };
static u32 temp[200000000] __attribute__((aligned(2097152)));
static u32 counts[3][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 stream16(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,int rad){
    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)&(u32)(rad-1),f=fill[b];
        buffer[b][f++]=v;
        if(f==limit[b]){
            u32*d=dst+pos[b];
            if(f==BUF)stream16(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){
    U bytes=(U)(u32)n*4u;
    madvise((void*)((U)a&~4095ul),bytes,MADV_HUGEPAGE);
    madvise(temp,bytes,MADV_HUGEPAGE);
    for(int b=0;b<RAD;++b)
        counts[0][b]=counts[1][b]=counts[2][b]=0;
    for(int i=0;i<n;++i){
        u32 v=a[i];++counts[0][v&2047];++counts[1][(v>>11)&2047];
        ++counts[2][v>>22];
    }
    scatter(a,temp,n,0,counts[0],2048);
    scatter(temp,a,n,11,counts[1],2048);
    scatter(a,temp,n,22,counts[2],1024);
    memcpy(a,temp,bytes);
    _mm256_zeroupper();
}

评测结果