代码

// This code is AI-generated. (AI 生成的代码)
// 1001: sort 100,000,000 unsigned ints.
//
// Strategy (multi-pass MSD radix sort):
//   1. MSD on the top 8 bits, scattering into tmp with per-bucket staging so
//      writes are sequential and flushed with non-temporal stores.
//   2. For each top-byte bucket, scatter the low 16 bits into a u16 scratch
//      array grouped by bits 16..23 (the sub-bucket).  Working on u16 halves
//      the memory width of the two remaining passes.
//   3. sort16() does two LSD passes (byte 0 then byte 1) on each sub-bucket and
//      writes the final u32 (sub-bucket top16 OR'ed in) back to the output with
//      non-temporal stores.
#include <immintrin.h>
#include <string.h>
#include <pthread.h>

typedef unsigned int u32;
typedef unsigned short u16;
typedef unsigned char u8;

#define NB 256
#define LINE 128

#define NMAX 100000000

static u32 tmp[NMAX + NB * LINE] __attribute__((aligned(64)));
static u32 staging[NB * LINE] __attribute__((aligned(64)));
static u32 cnt[NB], rbase[NB], pbase[NB], off[NB];
static u8  pos[NB];
#define SORT_THREADS 4
struct SortWorker {
    u16 scratch[1048576] __attribute__((aligned(64)));
    u16 scratch2[65536] __attribute__((aligned(64)));
    u16 scratch3[65536] __attribute__((aligned(64)));
    u32 bc[256], bc2[256];
    u32 sstart[256], spos[256], scount[256];
} __attribute__((aligned(64)));
static struct SortWorker workers[SORT_THREADS];
static __thread struct SortWorker *W;

__attribute__((target("avx2")))
static inline void flush_nt(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));
    __m256i v4  = _mm256_loadu_si256((const __m256i *)(src + 32));
    __m256i v5  = _mm256_loadu_si256((const __m256i *)(src + 40));
    __m256i v6  = _mm256_loadu_si256((const __m256i *)(src + 48));
    __m256i v7  = _mm256_loadu_si256((const __m256i *)(src + 56));
    __m256i v8  = _mm256_loadu_si256((const __m256i *)(src + 64));
    __m256i v9  = _mm256_loadu_si256((const __m256i *)(src + 72));
    __m256i v10 = _mm256_loadu_si256((const __m256i *)(src + 80));
    __m256i v11 = _mm256_loadu_si256((const __m256i *)(src + 88));
    __m256i v12 = _mm256_loadu_si256((const __m256i *)(src + 96));
    __m256i v13 = _mm256_loadu_si256((const __m256i *)(src + 104));
    __m256i v14 = _mm256_loadu_si256((const __m256i *)(src + 112));
    __m256i v15 = _mm256_loadu_si256((const __m256i *)(src + 120));
    _mm256_stream_si256((__m256i *)(dst + 0), v0);
    _mm256_stream_si256((__m256i *)(dst + 8), v1);
    _mm256_stream_si256((__m256i *)(dst + 16), v2);
    _mm256_stream_si256((__m256i *)(dst + 24), v3);
    _mm256_stream_si256((__m256i *)(dst + 32), v4);
    _mm256_stream_si256((__m256i *)(dst + 40), v5);
    _mm256_stream_si256((__m256i *)(dst + 48), v6);
    _mm256_stream_si256((__m256i *)(dst + 56), v7);
    _mm256_stream_si256((__m256i *)(dst + 64), v8);
    _mm256_stream_si256((__m256i *)(dst + 72), v9);
    _mm256_stream_si256((__m256i *)(dst + 80), v10);
    _mm256_stream_si256((__m256i *)(dst + 88), v11);
    _mm256_stream_si256((__m256i *)(dst + 96), v12);
    _mm256_stream_si256((__m256i *)(dst + 104), v13);
    _mm256_stream_si256((__m256i *)(dst + 112), v14);
    _mm256_stream_si256((__m256i *)(dst + 120), v15);
}

// Scatter src into dst, grouped by the top byte.  Writes go through one LINE
// sized staging buffer per bucket and are flushed with NT stores.
static void scatter_top(const u32 *src, u32 *dst, int n) {
    int i, j;
    for (j = 0; j < NB; j++) cnt[j] = 0;
    int n8 = n & ~7;
    for (i = 0; i < n8; i += 8) {
        cnt[src[i    ] >> 24]++; cnt[src[i + 1] >> 24]++;
        cnt[src[i + 2] >> 24]++; cnt[src[i + 3] >> 24]++;
        cnt[src[i + 4] >> 24]++; cnt[src[i + 5] >> 24]++;
        cnt[src[i + 6] >> 24]++; cnt[src[i + 7] >> 24]++;
    }
    for (; i < n; i++) cnt[src[i] >> 24]++;

    u32 t = 0, p = 0;
    for (j = 0; j < NB; j++) {
        u32 c = cnt[j];
        rbase[j] = t;
        pbase[j] = p;
        t += c;
        p = (p + c + LINE - 1) & ~(LINE - 1u);
    }
    for (j = 0; j < NB; j++) { off[j] = pbase[j]; pos[j] = 0; }

#define S1(X) do {                                  \
        u32 x_ = (X);                               \
        u32 b_ = x_ >> 24;                          \
        u32 p_ = pos[b_];                           \
        staging[b_ * LINE + p_] = x_;               \
        if (++p_ == LINE) {                         \
            u32 go_ = off[b_];                      \
            flush_nt(dst + go_, staging + b_ * LINE); \
            off[b_] = go_ + LINE;                   \
            p_ = 0;                                 \
        }                                           \
        pos[b_] = (u8)p_;                           \
    } while (0)

    int n8s = n & ~7;
    for (i = 0; i < n8s; i += 8) {
        S1(src[i]); S1(src[i + 1]); S1(src[i + 2]); S1(src[i + 3]);
        S1(src[i + 4]); S1(src[i + 5]); S1(src[i + 6]); S1(src[i + 7]);
    }
    for (; i < n; i++) S1(src[i]);
#undef S1

    for (j = 0; j < NB; j++) {
        u32 p_ = pos[j];
        if (p_) __builtin_memcpy(dst + off[j], staging + j * LINE, p_ * 4);
    }
}

__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;
}

// Two LSD byte passes on a u16 sub-bucket, then expand to u32 and NT-store.
static void sort16(const u16 *src, u32 *dst, int m, u32 top16) {
    int i, j;
    for (j = 0; j < 256; j++) W->bc[j] = 0;
    int m8 = m & ~7;
    for (i = 0; i < m8; i += 8) {
        W->bc[src[i] & 255]++;     W->bc[src[i + 1] & 255]++;
        W->bc[src[i + 2] & 255]++; W->bc[src[i + 3] & 255]++;
        W->bc[src[i + 4] & 255]++; W->bc[src[i + 5] & 255]++;
        W->bc[src[i + 6] & 255]++; W->bc[src[i + 7] & 255]++;
    }
    for (; i < m; i++) W->bc[src[i] & 255]++;
    for (j = 1; j < 256; j++) W->bc[j] += W->bc[j - 1];

    for (j = 0; j < 256; j++) W->bc2[j] = 0;
    int rm = m & ~7;
    for (i = m - 1; i >= rm; i--) {
        u16 x = src[i];
        W->scratch2[--W->bc[x & 255]] = x;
        W->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];
        u16 x4 = src[i - 4], x5 = src[i - 5], x6 = src[i - 6], x7 = src[i - 7];
        W->scratch2[--W->bc[x0 & 255]] = x0; W->bc2[(x0 >> 8) & 255]++;
        W->scratch2[--W->bc[x1 & 255]] = x1; W->bc2[(x1 >> 8) & 255]++;
        W->scratch2[--W->bc[x2 & 255]] = x2; W->bc2[(x2 >> 8) & 255]++;
        W->scratch2[--W->bc[x3 & 255]] = x3; W->bc2[(x3 >> 8) & 255]++;
        W->scratch2[--W->bc[x4 & 255]] = x4; W->bc2[(x4 >> 8) & 255]++;
        W->scratch2[--W->bc[x5 & 255]] = x5; W->bc2[(x5 >> 8) & 255]++;
        W->scratch2[--W->bc[x6 & 255]] = x6; W->bc2[(x6 >> 8) & 255]++;
        W->scratch2[--W->bc[x7 & 255]] = x7; W->bc2[(x7 >> 8) & 255]++;
    }
    for (j = 1; j < 256; j++) W->bc2[j] += W->bc2[j - 1];

    for (i = m - 1; i >= rm; i--)
        W->scratch3[--W->bc2[(W->scratch2[i] >> 8) & 255]] = W->scratch2[i];
    for (i = rm - 1; i >= 0; i -= 8) {
        u16 x0 = W->scratch2[i],     x1 = W->scratch2[i - 1], x2 = W->scratch2[i - 2], x3 = W->scratch2[i - 3];
        u16 x4 = W->scratch2[i - 4], x5 = W->scratch2[i - 5], x6 = W->scratch2[i - 6], x7 = W->scratch2[i - 7];
        W->scratch3[--W->bc2[(x0 >> 8) & 255]] = x0;
        W->scratch3[--W->bc2[(x1 >> 8) & 255]] = x1;
        W->scratch3[--W->bc2[(x2 >> 8) & 255]] = x2;
        W->scratch3[--W->bc2[(x3 >> 8) & 255]] = x3;
        W->scratch3[--W->bc2[(x4 >> 8) & 255]] = x4;
        W->scratch3[--W->bc2[(x5 >> 8) & 255]] = x5;
        W->scratch3[--W->bc2[(x6 >> 8) & 255]] = x6;
        W->scratch3[--W->bc2[(x7 >> 8) & 255]] = x7;
    }
    flush16_nt(dst, W->scratch3, m, top16);
}

// Sort one top-byte bucket by bits 0..23 using a u16 sidecar.
static void sort_bucket(const u32 *src, u32 *dst, int m, int bj) {
    int i, j;
    for (j = 0; j < 256; j++) W->bc[j] = 0;
    int m8 = m & ~7;
    for (i = 0; i < m8; i += 8) {
        W->bc[(src[i    ] >> 16) & 255]++; W->bc[(src[i + 1] >> 16) & 255]++;
        W->bc[(src[i + 2] >> 16) & 255]++; W->bc[(src[i + 3] >> 16) & 255]++;
        W->bc[(src[i + 4] >> 16) & 255]++; W->bc[(src[i + 5] >> 16) & 255]++;
        W->bc[(src[i + 6] >> 16) & 255]++; W->bc[(src[i + 7] >> 16) & 255]++;
    }
    for (; i < m; i++) W->bc[(src[i] >> 16) & 255]++;

    u32 t = 0;
    for (j = 0; j < 256; j++) { W->scount[j] = W->bc[j]; W->sstart[j] = t; W->spos[j] = t; t += W->bc[j]; }

    for (i = 0; i < m8; i += 8) {
        u32 x0 = src[i],     b0 = (x0 >> 16) & 255; W->scratch[W->spos[b0]++] = (u16)x0;
        u32 x1 = src[i + 1], b1 = (x1 >> 16) & 255; W->scratch[W->spos[b1]++] = (u16)x1;
        u32 x2 = src[i + 2], b2 = (x2 >> 16) & 255; W->scratch[W->spos[b2]++] = (u16)x2;
        u32 x3 = src[i + 3], b3 = (x3 >> 16) & 255; W->scratch[W->spos[b3]++] = (u16)x3;
        u32 x4 = src[i + 4], b4 = (x4 >> 16) & 255; W->scratch[W->spos[b4]++] = (u16)x4;
        u32 x5 = src[i + 5], b5 = (x5 >> 16) & 255; W->scratch[W->spos[b5]++] = (u16)x5;
        u32 x6 = src[i + 6], b6 = (x6 >> 16) & 255; W->scratch[W->spos[b6]++] = (u16)x6;
        u32 x7 = src[i + 7], b7 = (x7 >> 16) & 255; W->scratch[W->spos[b7]++] = (u16)x7;
    }
    for (; i < m; i++) {
        u32 x = src[i];
        W->scratch[W->spos[(x >> 16) & 255]++] = (u16)x;
    }

    for (j = 0; j < 256; j++) {
        int c = W->scount[j];
        if (c > 0)
            sort16(W->scratch + W->sstart[j], dst + W->sstart[j], c,
                   ((u32)((bj << 8) | j)) << 16);
    }
}

static u32 *g_a;
struct SortJob { int lo, hi, id; };
static void *sort_worker(void *arg) {
    struct SortJob *job = (struct SortJob *)arg;
    W = &workers[job->id];
    for (int j = job->lo; j < job->hi; j++) {
        int m = (int)cnt[j];
        if (m > 0)
            sort_bucket(tmp + pbase[j], g_a + rbase[j], m, j);
    }
    return 0;
}
void sort(unsigned *aa, int n) {
    u32 *a = (u32 *)aa;
    g_a = a;
    scatter_top(a, tmp, n);
    struct SortJob job[SORT_THREADS];
    pthread_t th[SORT_THREADS - 1];
    unsigned char created[SORT_THREADS - 1] = {0};
    for (int t = 0; t < SORT_THREADS; t++) {
        job[t].lo = (int)((long long)t * NB / SORT_THREADS);
        job[t].hi = (int)((long long)(t + 1) * NB / SORT_THREADS);
        job[t].id = t;
    }
    for (int t = 1; t < SORT_THREADS; t++)
        if (pthread_create(&th[t - 1], 0, sort_worker, &job[t]) == 0)
            created[t - 1] = 1;
        else
            sort_worker(&job[t]);
    sort_worker(&job[0]);
    for (int t = 1; t < SORT_THREADS; t++)
        if (created[t - 1]) pthread_join(th[t - 1], 0);
}

评测结果