代码

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

static u32 tmp[(1u<<27) + (1u<<21)] __attribute__((aligned(64)));
static u32 cnt[256];
static u32 rbase[256];
static u32 pbase[256];
static u32 off[256];
static u8 pos[256];
static u32 staging[256 * 64] __attribute__((aligned(64)));
static u32 scratch[1u<<20] __attribute__((aligned(64)));
static u32 scratch2[1u<<20] __attribute__((aligned(64)));
static u32 bc[256];
static u32 bc2[256];
static u32 sstart[256];
static u32 spos[256];
static u32 scount[256];

static void scatter_top(u32 *src, u32 *dst, int n) {
    int i, j;
    for (j = 0; j < 256; j++) cnt[j] = 0;
    for (i = 0; i < n; i++) cnt[src[i] >> 24]++;
    { u32 t = 0, p = 0; for (j = 0; j < 256; j++) { u32 c = cnt[j]; rbase[j] = t; pbase[j] = p; t += c; p = (p + c + 63) & ~63u; } }
    for (j = 0; j < 256; j++) { off[j] = pbase[j]; pos[j] = 0; }
    for (i = 0; i < n; i++) {
        u32 x = src[i]; u32 b = x >> 24; u32 p = pos[b];
        staging[b * 64 + p] = x; p++;
        if (p == 64) { u32 go = off[b]; __builtin_memcpy(dst+go, staging+b*64, 256); off[b] = go + 64; p = 0; }
        pos[b] = (u8)p;
    }
    for (j = 0; j < 256; j++) { u32 p = pos[j]; if (p > 0) memcpy(dst+off[j], staging+j*64, p*4); }
}

// sort m elements by 16 bits (2x8-bit LSD, fused hist), src -> dst (8KB sub-bucket, L2-resident)
static void sort16(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 (j = 0; j < 256; j++) bc2[j] = 0;
    for (i = m - 1; i >= 0; i--) { u32 x = src[i]; scratch2[--bc[x & 255]] = x; bc2[(x >> 8) & 255]++; }
    for (j = 1; j < 256; j++) bc2[j] += bc2[j-1];
    for (i = m - 1; i >= 0; i--) { u32 x = scratch2[i]; dst[--bc2[(x >> 8) & 255]] = x; }
}

// sort 2MB bucket by 24 bits: scatter by bits16-23, then sort16 each 8KB sub-bucket
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] >> 16) & 255]++;
    { u32 t = 0; for (j = 0; j < 256; j++) { u32 c = bc[j]; scount[j] = c; sstart[j] = t; spos[j] = t; t += c; } }
    for (i = 0; i < m; i++) { u32 x = src[i]; u32 b = (x >> 16) & 255; scratch[spos[b]++] = x; }
    for (j = 0; j < 256; j++) { int c = scount[j]; if (c > 0) sort16(scratch + sstart[j], dst + sstart[j], c); }
}

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

评测结果