代码

#define N 200000000U
#pragma GCC target("avx2")
#pragma GCC optimize("O3", "unroll-loops")
#include <immintrin.h>

typedef unsigned long U;
#ifndef N
#define N 100000000U
#endif
#define BUCKETS 256U
#ifndef U32_BUFFER_SIZE
#define U32_BUFFER_SIZE 64U
#endif
#ifndef PACKED_BUFFER_SIZE
#define PACKED_BUFFER_SIZE 64U
#endif
#define HIGH_BUFFER_SIZE 64U
#ifndef SIMD_UNPACK
#define SIMD_UNPACK 1
#endif
#define RESERVED_CAPACITY (((N + BUCKETS - 1) / BUCKETS) + 8192U)
#define WORK_N (RESERVED_CAPACITY * BUCKETS)

static unsigned char work[3U * WORK_N] __attribute__((aligned(64)));
static unsigned position[BUCKETS], starts[BUCKETS + 1];
static unsigned count[3][BUCKETS];
static unsigned mid_start[BUCKETS + 1];
static unsigned short low16_count[2][BUCKETS];
static unsigned short low16_work[8192] __attribute__((aligned(64)));
#ifndef MID16_BUFFER_SIZE
#define MID16_BUFFER_SIZE 64U
#endif
#ifndef TINY_LOW_SORT
#define TINY_LOW_SORT 0
#endif
static unsigned short mid16_buffer[BUCKETS][MID16_BUFFER_SIZE]
    __attribute__((aligned(64)));
static unsigned fill[BUCKETS], high_fill[BUCKETS];
static unsigned high_buffers[BUCKETS][HIGH_BUFFER_SIZE]
    __attribute__((aligned(64)));
#define BUFFER_BYTES (U32_BUFFER_SIZE * 4U > 3U * PACKED_BUFFER_SIZE ? \
                      U32_BUFFER_SIZE * 4U : 3U * PACKED_BUFFER_SIZE)
static unsigned char buffers[BUCKETS][BUFFER_BYTES]
    __attribute__((aligned(64)));
static unsigned actual_n;

typedef struct __attribute__((packed, may_alias)) {
    unsigned value;
} PackedU32;
typedef struct __attribute__((packed, may_alias)) {
    unsigned short value;
} PackedU16;

static __attribute__((always_inline)) inline unsigned load24(
        const unsigned char *source) {
    return ((const PackedU32 *)source)->value & 0x00ffffffU;
}

static __attribute__((always_inline)) inline void store24(
        unsigned char *destination, unsigned value) {
    destination[0] = (unsigned char)value;
    destination[1] = (unsigned char)(value >> 8);
    destination[2] = (unsigned char)(value >> 16);
}

static __attribute__((always_inline)) inline void store24_buffer(
        unsigned char *destination, unsigned value) {
    ((PackedU32 *)destination)->value = value;
}

static __attribute__((always_inline)) inline void pack16(
        __m256i first, __m256i second,
        __m128i *out0, __m128i *out1, __m128i *out2) {
    const __m128i compact = _mm_setr_epi8(
        0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, -1, -1, -1, -1);
    __m128i p0 = _mm_shuffle_epi8(_mm256_castsi256_si128(first), compact);
    __m128i p1 = _mm_shuffle_epi8(_mm256_extracti128_si256(first, 1), compact);
    __m128i p2 = _mm_shuffle_epi8(_mm256_castsi256_si128(second), compact);
    __m128i p3 = _mm_shuffle_epi8(_mm256_extracti128_si256(second, 1), compact);
    *out0 = _mm_or_si128(p0, _mm_slli_si128(p1, 12));
    *out1 = _mm_or_si128(_mm_srli_si128(p1, 4), _mm_slli_si128(p2, 8));
    *out2 = _mm_or_si128(_mm_srli_si128(p2, 8), _mm_slli_si128(p3, 4));
}

static __attribute__((always_inline)) inline void flush_packed64(
        unsigned char *destination, const unsigned *source,
        unsigned non_temporal) {
    __m128i x0, x1, x2, x3, x4, x5;
    pack16(_mm256_load_si256((const __m256i *)(source + 0)),
           _mm256_load_si256((const __m256i *)(source + 8)),
           &x0, &x1, &x2);
    pack16(_mm256_load_si256((const __m256i *)(source + 16)),
           _mm256_load_si256((const __m256i *)(source + 24)),
           &x3, &x4, &x5);
    __m256i y0 = _mm256_inserti128_si256(_mm256_castsi128_si256(x0), x1, 1);
    __m256i y1 = _mm256_inserti128_si256(_mm256_castsi128_si256(x2), x3, 1);
    __m256i y2 = _mm256_inserti128_si256(_mm256_castsi128_si256(x4), x5, 1);
    if (non_temporal) {
        _mm256_stream_si256((__m256i *)(destination + 0), y0);
        _mm256_stream_si256((__m256i *)(destination + 32), y1);
        _mm256_stream_si256((__m256i *)(destination + 64), y2);
    } else {
        _mm256_storeu_si256((__m256i *)(destination + 0), y0);
        _mm256_storeu_si256((__m256i *)(destination + 32), y1);
        _mm256_storeu_si256((__m256i *)(destination + 64), y2);
    }
    pack16(_mm256_load_si256((const __m256i *)(source + 32)),
           _mm256_load_si256((const __m256i *)(source + 40)),
           &x0, &x1, &x2);
    pack16(_mm256_load_si256((const __m256i *)(source + 48)),
           _mm256_load_si256((const __m256i *)(source + 56)),
           &x3, &x4, &x5);
    y0 = _mm256_inserti128_si256(_mm256_castsi128_si256(x0), x1, 1);
    y1 = _mm256_inserti128_si256(_mm256_castsi128_si256(x2), x3, 1);
    y2 = _mm256_inserti128_si256(_mm256_castsi128_si256(x4), x5, 1);
    if (non_temporal) {
        _mm256_stream_si256((__m256i *)(destination + 96), y0);
        _mm256_stream_si256((__m256i *)(destination + 128), y1);
        _mm256_stream_si256((__m256i *)(destination + 160), y2);
    } else {
        _mm256_storeu_si256((__m256i *)(destination + 96), y0);
        _mm256_storeu_si256((__m256i *)(destination + 128), y1);
        _mm256_storeu_si256((__m256i *)(destination + 160), y2);
    }
}

static __attribute__((always_inline)) inline void flush_u32_regular(
        unsigned *destination, const unsigned *source) {
    for (unsigned offset = 0; offset < U32_BUFFER_SIZE; offset += 8)
        _mm256_storeu_si256((__m256i *)(destination + offset),
                            _mm256_load_si256(
                                (const __m256i *)(source + offset)));
}

static void scatter_high_packed(const unsigned *source) {
    for (unsigned digit = 0; digit < BUCKETS; ++digit)
        position[digit] = digit * RESERVED_CAPACITY;
    for (unsigned i = 0; i < actual_n; ++i) {
        unsigned value = source[i];
        unsigned digit = value >> 24;
        unsigned output = position[digit];
        unsigned char *destination = work + 3U * output;
        if ((U)destination & 31U) {
            store24(destination, value);
            position[digit] = output + 1;
        } else {
            unsigned slot = high_fill[digit];
            high_buffers[digit][slot] = value;
            if (slot == HIGH_BUFFER_SIZE - 1) {
                flush_packed64(destination,
                               high_buffers[digit], 1);
                position[digit] = output + HIGH_BUFFER_SIZE;
                high_fill[digit] = 0;
            } else high_fill[digit] = slot + 1;
        }
    }
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        unsigned output = position[digit];
        unsigned amount = high_fill[digit];
        const unsigned *source_buffer = high_buffers[digit];
        for (unsigned j = 0; j < amount; ++j)
            store24(work + 3U * (output + j), source_buffer[j]);
        position[digit] = output + amount;
        high_fill[digit] = 0;
    }
    _mm_sfence();
}

static inline void prefix(unsigned *values, unsigned base) {
    unsigned sum = base;
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        unsigned amount = values[digit];
        values[digit] = sum;
        sum += amount;
    }
}

static __attribute__((always_inline)) inline void packed_to_u32(
        const unsigned char *source, unsigned source_begin,
        unsigned source_end, unsigned *destination, const unsigned *start,
        unsigned shift, unsigned high) {
    const __m128i expand = _mm_setr_epi8(
        0, 1, 2, -1, 3, 4, 5, -1, 6, 7, 8, -1, 9, 10, 11, -1);
#define EMIT_U32(VALUE) do {                                                \
        unsigned value = (VALUE) | high;                                    \
        unsigned digit = (value >> shift) & 255U;                           \
        unsigned sequence = fill[digit]++;                                  \
        unsigned slot = sequence & (U32_BUFFER_SIZE - 1);                   \
        ((unsigned *)buffers[digit])[slot] = value;                         \
        if (slot == U32_BUFFER_SIZE - 1)                                    \
            flush_u32_regular(destination + start[digit] + sequence + 1 -   \
                              U32_BUFFER_SIZE,                              \
                              (const unsigned *)buffers[digit]);            \
    } while (0)
    unsigned i = source_begin;
    const unsigned char *packed = source + 3U * source_begin;
#if SIMD_UNPACK
    for (; i + 16 <= source_end; i += 16, packed += 48) {
        __m128i in0 = _mm_loadu_si128((const __m128i *)(packed + 0));
        __m128i in1 = _mm_loadu_si128((const __m128i *)(packed + 16));
        __m128i in2 = _mm_loadu_si128((const __m128i *)(packed + 32));
        __m128i v0 = _mm_shuffle_epi8(in0, expand);
        __m128i v1 = _mm_shuffle_epi8(_mm_alignr_epi8(in1, in0, 12), expand);
        __m128i v2 = _mm_shuffle_epi8(_mm_alignr_epi8(in2, in1, 8), expand);
        __m128i v3 = _mm_shuffle_epi8(_mm_srli_si128(in2, 4), expand);
        EMIT_U32((unsigned)_mm_cvtsi128_si32(v0));
        EMIT_U32((unsigned)_mm_extract_epi32(v0, 1));
        EMIT_U32((unsigned)_mm_extract_epi32(v0, 2));
        EMIT_U32((unsigned)_mm_extract_epi32(v0, 3));
        EMIT_U32((unsigned)_mm_cvtsi128_si32(v1));
        EMIT_U32((unsigned)_mm_extract_epi32(v1, 1));
        EMIT_U32((unsigned)_mm_extract_epi32(v1, 2));
        EMIT_U32((unsigned)_mm_extract_epi32(v1, 3));
        EMIT_U32((unsigned)_mm_cvtsi128_si32(v2));
        EMIT_U32((unsigned)_mm_extract_epi32(v2, 1));
        EMIT_U32((unsigned)_mm_extract_epi32(v2, 2));
        EMIT_U32((unsigned)_mm_extract_epi32(v2, 3));
        EMIT_U32((unsigned)_mm_cvtsi128_si32(v3));
        EMIT_U32((unsigned)_mm_extract_epi32(v3, 1));
        EMIT_U32((unsigned)_mm_extract_epi32(v3, 2));
        EMIT_U32((unsigned)_mm_extract_epi32(v3, 3));
    }
#endif
    for (; i < source_end; ++i, packed += 3)
        EMIT_U32(load24(packed));
#undef EMIT_U32
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        unsigned amount = fill[digit] & (U32_BUFFER_SIZE - 1);
        unsigned output = start[digit] +
                          (fill[digit] & ~(U32_BUFFER_SIZE - 1));
        const unsigned *source_buffer = (const unsigned *)buffers[digit];
        for (unsigned j = 0; j < amount; ++j)
            destination[output + j] = source_buffer[j];
        fill[digit] = 0;
    }
}

static __attribute__((always_inline)) inline void u32_to_packed(
        const unsigned *source, unsigned begin, unsigned end,
        unsigned char *destination, const unsigned *start, unsigned shift) {
    for (unsigned i = begin; i < end; ++i) {
        unsigned value = source[i];
        unsigned digit = (value >> shift) & 255U;
        unsigned sequence = fill[digit]++;
        unsigned slot = sequence & (PACKED_BUFFER_SIZE - 1);
        ((unsigned *)buffers[digit])[slot] = value;
        if (slot == PACKED_BUFFER_SIZE - 1)
            flush_packed64(destination +
                           3U * (start[digit] + sequence + 1 -
                                   PACKED_BUFFER_SIZE),
                           (const unsigned *)buffers[digit], 0);
    }
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        unsigned amount = fill[digit] & (PACKED_BUFFER_SIZE - 1);
        unsigned output = start[digit] +
                          (fill[digit] & ~(PACKED_BUFFER_SIZE - 1));
        const unsigned *source_buffer = (const unsigned *)buffers[digit];
        for (unsigned j = 0; j < amount; ++j)
            store24(destination + 3U * (output + j), source_buffer[j]);
        fill[digit] = 0;
    }
}

static __attribute__((always_inline)) inline void packed_to_u16(
        const unsigned char *source, unsigned source_begin,
        unsigned source_end, unsigned short *destination,
        const unsigned *start) {
    unsigned output_position[BUCKETS], pass_fill[BUCKETS] = {0};
    for (unsigned digit = 0; digit < BUCKETS; ++digit)
        output_position[digit] = start[digit];
    const unsigned char *input = source + 3U * source_begin;
    for (unsigned i = source_begin; i < source_end; ++i, input += 3) {
        unsigned value = ((const PackedU16 *)input)->value;
        unsigned digit = input[2];
        unsigned output = output_position[digit];
        unsigned slot = pass_fill[digit];
        mid16_buffer[digit][slot] = (unsigned short)value;
        if (slot == MID16_BUFFER_SIZE - 1U) {
            for (unsigned j = 0; j < MID16_BUFFER_SIZE; j += 16)
                _mm256_storeu_si256((__m256i *)(destination + output + j),
                    _mm256_load_si256(
                        (const __m256i *)(mid16_buffer[digit] + j)));
            output_position[digit] = output + MID16_BUFFER_SIZE;
            pass_fill[digit] = 0;
        } else pass_fill[digit] = slot + 1;
    }
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        unsigned output = output_position[digit], amount = pass_fill[digit];
        for (unsigned i = 0; i < amount; ++i)
            destination[output + i] = mid16_buffer[digit][i];
    }
}

void sort(unsigned *a, int n) {
    actual_n = (unsigned)n;
    scatter_high_packed(a);
    unsigned sum = 0;
    for (unsigned digit = 0; digit < BUCKETS; ++digit) {
        starts[digit] = sum;
        sum += position[digit] - digit * RESERVED_CAPACITY;
    }
    starts[BUCKETS] = actual_n;

    unsigned short *compressed = (unsigned short *)a;
    for (unsigned bucket = BUCKETS; bucket-- != 0;) {
        unsigned begin = starts[bucket], end = starts[bucket + 1];
        unsigned source_begin = bucket * RESERVED_CAPACITY;
        unsigned source_end = source_begin + end - begin;
        __builtin_memset(count[2], 0, sizeof(count[2]));
        const unsigned char *packed_source = work + 3U * source_begin;
        for (unsigned i = source_begin; i < source_end;
             ++i, packed_source += 3)
            ++count[2][packed_source[2]];
        unsigned middle_sum = begin;
        for (unsigned digit = 0; digit < BUCKETS; ++digit) {
            mid_start[digit] = middle_sum;
            middle_sum += count[2][digit];
            count[2][digit] = mid_start[digit];
        }
        mid_start[BUCKETS] = end;
        unsigned high = bucket << 24;
        packed_to_u16(work, source_begin, source_end, compressed, count[2]);
        for (unsigned mid = BUCKETS; mid-- != 0;) {
            unsigned low_begin = mid_start[mid], low_end = mid_start[mid + 1];
#if TINY_LOW_SORT
            __builtin_memset(low16_count[1], 0, sizeof(low16_count[1]));
            for (unsigned i = low_begin; i < low_end; ++i)
                ++low16_count[1][compressed[i] >> 8];
            unsigned high_sum = 0;
            for (unsigned digit = 0; digit < BUCKETS; ++digit) {
                unsigned amount = low16_count[1][digit];
                low16_count[0][digit] = (unsigned short)high_sum;
                high_sum += amount;
                low16_count[1][digit] = (unsigned short)high_sum;
            }
            for (unsigned i = low_begin; i < low_end; ++i) {
                unsigned value = compressed[i], digit = value >> 8;
                low16_work[low16_count[0][digit]++] = (unsigned short)value;
            }
            unsigned tiny_begin = 0;
            for (unsigned digit = 0; digit < BUCKETS; ++digit) {
                unsigned tiny_end = low16_count[1][digit];
                for (unsigned i = tiny_begin + 1; i < tiny_end; ++i) {
                    unsigned short value = low16_work[i];
                    unsigned j = i;
                    while (j > tiny_begin && low16_work[j - 1] > value) {
                        low16_work[j] = low16_work[j - 1];
                        --j;
                    }
                    low16_work[j] = value;
                }
                tiny_begin = tiny_end;
            }
            unsigned fixed = high | (mid << 16), amount = low_end - low_begin;
            for (unsigned i = 0; i < amount; ++i)
                a[low_begin + i] = fixed | low16_work[i];
#else
            __builtin_memset(low16_count, 0, sizeof(low16_count));
            for (unsigned i = low_begin; i < low_end; ++i) {
                unsigned value = compressed[i];
                ++low16_count[0][(unsigned char)value];
                ++low16_count[1][(unsigned char)(value >> 8)];
            }
            unsigned sum0 = 0, sum1 = 0;
            for (unsigned digit = 0; digit < BUCKETS; ++digit) {
                unsigned amount0 = low16_count[0][digit];
                unsigned amount1 = low16_count[1][digit];
                low16_count[0][digit] = (unsigned short)sum0;
                low16_count[1][digit] = (unsigned short)sum1;
                sum0 += amount0;
                sum1 += amount1;
            }
            for (unsigned i = low_begin; i < low_end; ++i) {
                unsigned value = compressed[i];
                low16_work[low16_count[0][(unsigned char)value]++] =
                    (unsigned short)value;
            }
            unsigned fixed = high | (mid << 16), amount = low_end - low_begin;
            for (unsigned i = 0; i < amount; ++i) {
                unsigned value = low16_work[i];
                a[low_begin + low16_count[1][value >> 8]++] = fixed | value;
            }
#endif
        }
    }
    _mm256_zeroupper();
}

评测结果