代码

#include <string.h>
#include <immintrin.h>
typedef unsigned int u32;
typedef unsigned char u8;

static u32 tmp[(1u<<27) + (1u<<21)] __attribute__((aligned(64)));
static u32 cnt[65536];
static u32 rbase[65536];
static u32 pbase[65536];
static u32 off[65536];
static u8  pos[65536];
static u32 staging[65536 * 32] __attribute__((aligned(64)));
static u32 scratch[8192] __attribute__((aligned(64)));
static u32 bc[256];

__attribute__((target("avx2")))
static inline void flush_32(u32 *dst, const u32 *src) {
    __m256i v0 = _mm256_loadu_si256((const __m256i*)(src+0));
    __m256i v1 = _mm256_loadu_si256((const __m256i*)(src+8));
    __m256i v2 = _mm256_loadu_si256((const __m256i*)(src+16));
    __m256i v3 = _mm256_loadu_si256((const __m256i*)(src+24));
    _mm256_storeu_si256((__m256i*)(dst+0), v0);
    _mm256_storeu_si256((__m256i*)(dst+8), v1);
    _mm256_storeu_si256((__m256i*)(dst+16), v2);
    _mm256_storeu_si256((__m256i*)(dst+24), v3);
}


static void radix16_top(u32 *src, u32 *dst, int n) {
    int i, j;
    for (j = 0; j < 65536; j++) cnt[j] = 0;
    for (i = 0; i < n; i++) { __builtin_prefetch(&src[i+512],0,0); cnt[src[i] >> 16]++; }
    { u32 t = 0, p = 0; for (j = 0; j < 65536; j++) { u32 c = cnt[j]; rbase[j] = t; pbase[j] = p; t += c; p = (p + c + 31) & ~31u; } }
    for (j = 0; j < 65536; j++) { off[j] = pbase[j]; pos[j] = 0; }
    for (i = 0; i < n; i++) {
        __builtin_prefetch(&src[i+512],0,0);
        u32 x = src[i];
        u32 b = x >> 16;
        u32 p = pos[b];
        staging[b * 32 + p] = x;
        p++;
        if (p == 32) { flush_32(dst + off[b], staging + b * 32); off[b] += 32; p = 0; }
        pos[b] = (u8)p;
    }
    for (j = 0; j < 65536; j++) { u32 p = pos[j]; if (p > 0) memcpy(dst + off[j], staging + j*32, p*4); }
}


static void sort_bucket(u32 *src, u32 *dst, int m) {
    int i, j;
    for (j = 0; j < 256; j++) bc[j] = 0;
    for (i = 0; i < m; i++) bc[src[i] & 255]++;
    for (j = 1; j < 256; j++) bc[j] += bc[j-1];
    for (i = m - 1; i >= 0; i--) { u32 x = src[i]; scratch[--bc[x & 255]] = x; }
    for (j = 0; j < 256; j++) bc[j] = 0;
    for (i = 0; i < m; i++) bc[(scratch[i] >> 8) & 255]++;
    for (j = 1; j < 256; j++) bc[j] += bc[j-1];
    for (i = m - 1; i >= 0; i--) { u32 x = scratch[i]; dst[--bc[(x >> 8) & 255]] = x; }
}

void sort(unsigned *aa, int n) {
    u32 *a = (u32*)aa;
    int j;
    radix16_top(a, tmp, n);
    for (j = 0; j < 65536; j++) {
        int m = (int)cnt[j];
        if (m > 0) sort_bucket(tmp + pbase[j], a + rbase[j], m);
    }
}

评测结果