代码

// 1001 v12: [8][8][16] MSD + vpshufb 3-byte transpose compress + u16 sidecar lower passes.
// Pass 1: scalar top-8 histogram.
// Pass 2: compress a[] -> s01 (u16 low-16) + s2 (u8 byte2), bucketed by top-8, via 16-elem
//         vpshufb transpose (staging in L1, flush to 64-aligned plane regions).
// Pass 3: per bucket, hist byte2 (read s2), scatter low-16 (u16) into scratch by byte2.
// Pass 4: sort16 (2x8-bit on u16) + reconstruct NT to a[].
#include <immintrin.h>
#include <string.h>
#include <emmintrin.h>
typedef unsigned int u32;
typedef unsigned short u16;
typedef unsigned char u8;

#ifndef N
#define N 100000000
#endif

#define NB 256

static u16 s01[N + NB*64] __attribute__((aligned(64)));
static u8  s2[N + NB*64] __attribute__((aligned(64)));

static u32 cnt[NB];
static u32 B[257];       // dense element offsets
static u32 aln[NB];      // 64-aligned element offset per bucket (for planes)
static u32 flcnt[NB];    // flushed element count per bucket
static u8  pos[NB];      // staging count (0..15)
static u32 stg[NB * 16] __attribute__((aligned(64)));  // staging 16 u32 per bucket

static u16 scratch[1048576] __attribute__((aligned(64)));
static u16 scratch2[65536] __attribute__((aligned(64)));
static u16 scratch3[65536] __attribute__((aligned(64)));
static u32 bc[256];
static u32 bc2[256];
static u32 sstart[256];
static u32 spos[256];
static u32 scount[256];

// transpose 4 xmm (16 u32) -> p0 (byte0), p1 (byte1), p2 (byte2), same element order
__attribute__((target("ssse3")))
static inline void transpose3(__m128i r0, __m128i r1, __m128i r2, __m128i r3,
                              __m128i *p0, __m128i *p1, __m128i *p2) {
    const __m128i m0 = _mm_setr_epi8(0,4,8,12, -1,-1,-1,-1, -1,-1,-1,-1, -1,-1,-1,-1);
    const __m128i m1 = _mm_setr_epi8(1,5,9,13, -1,-1,-1,-1, -1,-1,-1,-1, -1,-1,-1,-1);
    const __m128i m2 = _mm_setr_epi8(2,6,10,14, -1,-1,-1,-1, -1,-1,-1,-1, -1,-1,-1,-1);
    __m128i a0 = _mm_shuffle_epi8(r0, m0), a1 = _mm_shuffle_epi8(r1, m0);
    __m128i a2 = _mm_shuffle_epi8(r2, m0), a3 = _mm_shuffle_epi8(r3, m0);
    __m128i t0 = _mm_unpacklo_epi32(a0, a1), t1 = _mm_unpacklo_epi32(a2, a3);
    *p0 = _mm_unpacklo_epi64(t0, t1);
    __m128i b0 = _mm_shuffle_epi8(r0, m1), b1 = _mm_shuffle_epi8(r1, m1);
    __m128i b2 = _mm_shuffle_epi8(r2, m1), b3 = _mm_shuffle_epi8(r3, m1);
    __m128i u0 = _mm_unpacklo_epi32(b0, b1), u1 = _mm_unpacklo_epi32(b2, b3);
    *p1 = _mm_unpacklo_epi64(u0, u1);
    __m128i c0 = _mm_shuffle_epi8(r0, m2), c1 = _mm_shuffle_epi8(r1, m2);
    __m128i c2 = _mm_shuffle_epi8(r2, m2), c3 = _mm_shuffle_epi8(r3, m2);
    __m128i v0 = _mm_unpacklo_epi32(c0, c1), v1 = _mm_unpacklo_epi32(c2, c3);
    *p2 = _mm_unpacklo_epi64(v0, v1);
}

// flush 16 staged elements of bucket k to the planes (full 16, no partial)
__attribute__((target("ssse3")))
static inline void flush_full(int k) {
    __m128i r0 = _mm_load_si128((__m128i*)(stg + k*16 + 0));
    __m128i r1 = _mm_load_si128((__m128i*)(stg + k*16 + 4));
    __m128i r2 = _mm_load_si128((__m128i*)(stg + k*16 + 8));
    __m128i r3 = _mm_load_si128((__m128i*)(stg + k*16 + 12));
    __m128i p0, p1, p2;
    transpose3(r0, r1, r2, r3, &p0, &p1, &p2);
    u32 base = aln[k] + flcnt[k];
    __m128i q0 = _mm_unpacklo_epi8(p0, p1);  // low-16 of elements 0..7
    __m128i q1 = _mm_unpackhi_epi8(p0, p1);  // low-16 of elements 8..15
    _mm_store_si128((__m128i*)(s01 + base + 0), q0);
    _mm_store_si128((__m128i*)(s01 + base + 8), q1);
    _mm_store_si128((__m128i*)(s2 + base), p2);
    flcnt[k] += 16;
}

// flush partial (<16) staged elements of bucket k
static inline void flush_partial(int k, int m) {
    u32 base = aln[k] + flcnt[k];
    for (int j = 0; j < m; j++) {
        u32 x = stg[k*16 + j];
        s01[base + j] = (u16)x;
        s2[base + j] = (u8)(x >> 16);
    }
}

__attribute__((target("avx2")))
static inline void flush16_nt(u32 *dst, const u16 *src, int m, u32 top16) {
    int mi = 0;
    while ((((unsigned long)(dst + mi)) & 31) && mi < m) { dst[mi] = ((u32)src[mi]) | top16; mi++; }
    __m256i t = _mm256_set1_epi32((int)top16);
    for (; mi + 16 <= m; mi += 16) {
        __m128i lo = _mm_loadu_si128((const __m128i*)(src + mi));
        __m128i hi = _mm_loadu_si128((const __m128i*)(src + mi + 8));
        __m256i a0 = _mm256_cvtepu16_epi32(lo);
        __m256i a1 = _mm256_cvtepu16_epi32(hi);
        _mm256_stream_si256((__m256i*)(dst + mi + 0), _mm256_or_si256(a0, t));
        _mm256_stream_si256((__m256i*)(dst + mi + 8), _mm256_or_si256(a1, t));
    }
    for (; mi < m; mi++) dst[mi] = ((u32)src[mi]) | top16;
}

static void sort16(u16 *src, u32 *dst, int m, u32 top16) {
    int i, j;
    for (j = 0; j < 256; j++) bc[j] = 0;
    int m8 = m & ~7;
    for (i = 0; i < m8; i += 8) {
        bc[src[i]&255]++; bc[src[i+1]&255]++; bc[src[i+2]&255]++; bc[src[i+3]&255]++;
        bc[src[i+4]&255]++; bc[src[i+5]&255]++; bc[src[i+6]&255]++; bc[src[i+7]&255]++;
    }
    for (; 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;
    int rm = m & ~7;
    for (i = m - 1; i >= rm; i--) { u16 x = src[i]; scratch2[--bc[x & 255]] = x; bc2[(x >> 8) & 255]++; }
    for (i = rm - 1; i >= 0; i -= 8) {
        u16 x0=src[i],x1=src[i-1],x2=src[i-2],x3=src[i-3],x4=src[i-4],x5=src[i-5],x6=src[i-6],x7=src[i-7];
        scratch2[--bc[x0&255]]=x0; bc2[(x0>>8)&255]++;
        scratch2[--bc[x1&255]]=x1; bc2[(x1>>8)&255]++;
        scratch2[--bc[x2&255]]=x2; bc2[(x2>>8)&255]++;
        scratch2[--bc[x3&255]]=x3; bc2[(x3>>8)&255]++;
        scratch2[--bc[x4&255]]=x4; bc2[(x4>>8)&255]++;
        scratch2[--bc[x5&255]]=x5; bc2[(x5>>8)&255]++;
        scratch2[--bc[x6&255]]=x6; bc2[(x6>>8)&255]++;
        scratch2[--bc[x7&255]]=x7; bc2[(x7>>8)&255]++;
    }
    for (j = 1; j < 256; j++) bc2[j] += bc2[j-1];
    for (i = m - 1; i >= rm; i--) { u16 x = scratch2[i]; scratch3[--bc2[(x >> 8) & 255]] = x; }
    for (i = rm - 1; i >= 0; i -= 8) {
        u16 x0=scratch2[i],x1=scratch2[i-1],x2=scratch2[i-2],x3=scratch2[i-3],x4=scratch2[i-4],x5=scratch2[i-5],x6=scratch2[i-6],x7=scratch2[i-7];
        scratch3[--bc2[(x0>>8)&255]]=x0;
        scratch3[--bc2[(x1>>8)&255]]=x1;
        scratch3[--bc2[(x2>>8)&255]]=x2;
        scratch3[--bc2[(x3>>8)&255]]=x3;
        scratch3[--bc2[(x4>>8)&255]]=x4;
        scratch3[--bc2[(x5>>8)&255]]=x5;
        scratch3[--bc2[(x6>>8)&255]]=x6;
        scratch3[--bc2[(x7>>8)&255]]=x7;
    }
    flush16_nt(dst, scratch3, m, top16);
}

void sort(unsigned *aa, int n) {
    u32 *a = (u32*)aa;
    int i, j;
    (void)n;

    // Pass 1: histogram top-8
    for (j = 0; j < NB; j++) cnt[j] = 0;
    int n8 = N & ~7;
    for (i = 0; i < n8; i += 8) {
        cnt[a[i]>>24]++; cnt[a[i+1]>>24]++; cnt[a[i+2]>>24]++; cnt[a[i+3]>>24]++;
        cnt[a[i+4]>>24]++; cnt[a[i+5]>>24]++; cnt[a[i+6]>>24]++; cnt[a[i+7]>>24]++;
    }
    for (; i < N; i++) cnt[a[i] >> 24]++;
    { u32 acc = 0, acc2 = 0; B[0] = 0; for (j = 0; j < NB; j++) { u32 c = cnt[j]; B[j+1] = acc + c; aln[j] = acc2; acc += c; acc2 = (acc2 + c + 63) & ~63u; } }

    // Pass 2: compress via vpshufb transpose
    for (j = 0; j < NB; j++) { pos[j] = 0; flcnt[j] = 0; }
    for (i = 0; i < n8; i += 8) {
        u32 x0=a[i+0],k0=x0>>24,p0=pos[k0]; stg[k0*16+p0]=x0; p0++; if(p0==16){flush_full(k0); p0=0;} pos[k0]=(u8)p0;
        u32 x1=a[i+1],k1=x1>>24,p1=pos[k1]; stg[k1*16+p1]=x1; p1++; if(p1==16){flush_full(k1); p1=0;} pos[k1]=(u8)p1;
        u32 x2=a[i+2],k2=x2>>24,p2=pos[k2]; stg[k2*16+p2]=x2; p2++; if(p2==16){flush_full(k2); p2=0;} pos[k2]=(u8)p2;
        u32 x3=a[i+3],k3=x3>>24,p3=pos[k3]; stg[k3*16+p3]=x3; p3++; if(p3==16){flush_full(k3); p3=0;} pos[k3]=(u8)p3;
        u32 x4=a[i+4],k4=x4>>24,p4=pos[k4]; stg[k4*16+p4]=x4; p4++; if(p4==16){flush_full(k4); p4=0;} pos[k4]=(u8)p4;
        u32 x5=a[i+5],k5=x5>>24,p5=pos[k5]; stg[k5*16+p5]=x5; p5++; if(p5==16){flush_full(k5); p5=0;} pos[k5]=(u8)p5;
        u32 x6=a[i+6],k6=x6>>24,p6=pos[k6]; stg[k6*16+p6]=x6; p6++; if(p6==16){flush_full(k6); p6=0;} pos[k6]=(u8)p6;
        u32 x7=a[i+7],k7=x7>>24,p7=pos[k7]; stg[k7*16+p7]=x7; p7++; if(p7==16){flush_full(k7); p7=0;} pos[k7]=(u8)p7;
    }
    for (; i < N; i++) {
        u32 x = a[i]; u32 k = x >> 24; u32 p = pos[k];
        stg[k*16 + p] = x; p++;
        if (p == 16) { flush_full(k); p = 0; }
        pos[k] = (u8)p;
    }
    for (j = 0; j < NB; j++) if (pos[j]) flush_partial(j, pos[j]);

    // Pass 3+4: per top-8 bucket
    for (j = 0; j < NB; j++) {
        u32 m = cnt[j];
        if (m == 0) continue;
        u32 base = aln[j];
        u32 baseA = B[j];
        int mb8 = (int)m & ~7;

        for (int q = 0; q < 256; q++) bc[q] = 0;
        for (i = 0; i < mb8; i += 8) {
            bc[s2[base+i]]++; bc[s2[base+i+1]]++; bc[s2[base+i+2]]++; bc[s2[base+i+3]]++;
            bc[s2[base+i+4]]++; bc[s2[base+i+5]]++; bc[s2[base+i+6]]++; bc[s2[base+i+7]]++;
        }
        for (; i < (int)m; i++) bc[s2[base + i]]++;
        { u32 t = 0; for (int q = 0; q < 256; q++) { u32 c = bc[q]; scount[q] = c; sstart[q] = t; spos[q] = t; t += c; } }

        for (i = 0; i < mb8; i += 8) {
            u32 g0=s2[base+i+0]; scratch[spos[g0]++]=s01[base+i+0];
            u32 g1=s2[base+i+1]; scratch[spos[g1]++]=s01[base+i+1];
            u32 g2=s2[base+i+2]; scratch[spos[g2]++]=s01[base+i+2];
            u32 g3=s2[base+i+3]; scratch[spos[g3]++]=s01[base+i+3];
            u32 g4=s2[base+i+4]; scratch[spos[g4]++]=s01[base+i+4];
            u32 g5=s2[base+i+5]; scratch[spos[g5]++]=s01[base+i+5];
            u32 g6=s2[base+i+6]; scratch[spos[g6]++]=s01[base+i+6];
            u32 g7=s2[base+i+7]; scratch[spos[g7]++]=s01[base+i+7];
        }
        for (; i < (int)m; i++) { u32 g = s2[base+i]; scratch[spos[g]++] = s01[base+i]; }
        for (int q = 0; q < 256; q++) { int c = scount[q]; if (c > 0) sort16(scratch + sstart[q], a + baseA + sstart[q], c, ((u32)((j<<8)|q))<<16); }
    }
}

评测结果