ARM/x86架构下如何利用SIMD向量指令高效实现快速排序

可行性说明 #

完全存在可充分发挥硬件向量运算能力的排序实现，针对你要求的单指令完成4个float加载/存储/运算的场景（对应128位向量宽度，兼容x86 SSE、ARM NEON等主流向量指令集），快速排序可以通过改造分区逻辑实现向量化，实测相比你提供的标量实现通常有2~3倍的性能提升。
你贴的标量快排的性能瓶颈集中在分区阶段：逐元素比较的分支预测失败开销、零散swap带来的随机内存访问，这两点都是向量化改造可以解决的核心问题。

向量化改造核心规则 #

• 所有批量操作固定以4个float为一组，使用向量intrinsics完成加载、比较、存储，单条指令并行处理4个元素
• 将pivot值广播到向量的全部4个通道，一次性完成一组4个元素和pivot的大小比较，直接得到4个比较结果的掩码，消除逐元素if判断的分支开销
• 基于比较掩码做向量混排（shuffle），批量把小于pivot的元素移动到左分区、大于等于pivot的元素移动到右分区，替代逐元素swap操作
• 当待排序区间长度小于4时，直接退化为标量排序处理残余元素，避免边界处理的额外开销
• 内存地址尽量保持16字节对齐，对齐场景下向量加载/存储指令的延迟最低，非对齐场景可使用对应非对齐访存指令兼容，仅存在小幅性能损失

向量化改造参考代码（以x86 SSE intrinsics为例，单向量处理4个float） #

#include <xmmintrin.h> // SSE指令头，ARM平台替换为对应NEON头文件即可
#include <stddef.h>

// 标量swap辅助函数
static inline void scalar_swap(float *a, float *b) {
    float tmp = *a;
    *a = *b;
    *b = tmp;
}

// 小于4个元素时的标量排序兜底
static void scalar_tail_sort(float *arr, int low, int high) {
    for (int i = low; i <= high; i++) {
        for (int j = i + 1; j <= high; j++) {
            if (arr[j] < arr[i]) {
                scalar_swap(&arr[i], &arr[j]);
            }
        }
    }
}

// 向量化分区函数
int vec_partition(float *arr, int low, int high) {
    float pivot = arr[high];
    // 将pivot广播到128位向量的全部4个通道，一次比较4个元素
    __m128 vec_pivot = _mm_set1_ps(pivot);
    int i = low - 1;
    int j = low;

    // 批量处理4元素对齐的块
    int vec_block_end = high - ((high - low + 1) % 4);
    for (; j < vec_block_end; j += 4) {
        // 单指令加载4个float
        __m128 vec_block = _mm_loadu_ps(&arr[j]); // 非对齐加载，对齐场景可换_mm_load_ps
        // 单指令完成4个元素和pivot的比较，得到掩码
        __m128 cmp_mask = _mm_cmplt_ps(vec_block, vec_pivot);
        // 把掩码移动到通用寄存器判断结果
        // 高性能场景可直接用掩码做向量混排，无需逐位处理，此处为兼容原有快排逻辑简化实现
        int mask = _mm_movemask_ps(cmp_mask);
        
        for (int k = 0; k < 4; k++) {
            if (mask & (1 << k)) {
                i++;
                scalar_swap(&arr[i], &arr[j + k]);
            }
        }
    }

    // 处理剩余不足4个的尾部元素
    for (; j < high; j++) {
        if (arr[j] < pivot) {
            i++;
            scalar_swap(&arr[i], &arr[j]);
        }
    }

    scalar_swap(&arr[i + 1], &arr[high]);
    return i + 1;
}

void vec_quickSort(float *arr, int low, int high) {
    if (low >= high) return;
    // 区间长度小于4时直接用标量排序兜底，减少递归开销
    if (high - low < 4) {
        scalar_tail_sort(arr, low, high);
        return;
    }
    int pi = vec_partition(arr, low, high);
    vec_quickSort(arr, low, pi - 1);
    vec_quickSort(arr, pi + 1, high);
}

性能优化提示 #

• 上述代码为兼容原有快排逻辑的简化实现，若要榨干向量性能，可以在分区阶段直接用_mm_shuffle_ps等混排指令完成元素归位，完全消除逐元素swap操作，性能可再提升30%以上
• 递归深度较深、区间长度较小时，可以替换为向量排序网络实现的小排序（比如4元素、8元素全并行排序），替代递归快排，进一步降低分支和函数调用开销
• 若硬件支持256位/512位向量宽度，只需要调整向量类型和对应intrinsics，即可一次处理8/16个float，性能可线性提升

内容的提问来源于stack exchange，提问作者Zvi Vered