simd 排序

最近,阅读了 SIMD- and Cache-Friendly Algorithm for Sorting an Array of Structures 一文后,动手编写了一个 基于 SIMD 的排序算法,具体代码参考:SIMD Sort Test

// file poc/vector_example1/src/sort2.rs

// simd based sort u32x8 
fn sort_u32x8(nums: &mut u32x8) {
...
}

// simd based sort u32x16
fn sort_u32x16(nums: &mut u32x16){
...
}

// simd based sort u32x32
fn sort_u32x32(nums: &mut u32x32)
...
}

// simd based meger sort of 2 u32x16
fn merge_sort_u32x16x2(p1: &mut u32x16, p2: &mut u32x16) {
}

// 基于 merge sort(sort_u32x32 + merge_sort_u32x16x2) 的排序实现,
pub fn sort(nums: &mut [u32]) {
...
}

作为对比,测试了 该算法与 Rust 标准库中的 Slice::sortSlice::sort_unstable 的性能差异:(测试环境: M1-Max 64G)

数据量Slice::sortSlice::sort_unstableMy SIMD sort
32237 ns156 ns56 ns
64642 ms348 ns176 ns
1281.156 us785 ns511 ns
2562.454 us1.635 us1.255 us
5125.192 us3.488 us3.054 us
1K10.857 us7.381 us7.382 us
2K22.746 us15.638 us17.443 us
4K47.517 us32.933 us40.407 us
8K98.424 us69.258 us91.900 us
16K203.153 us145.538 us205.371 us
32K420.905 us305.407 us455.912 us
64K870.353 us641.122 us999.676 us
128K1.809017 ms1.346237 ms2.180088 ms
256K3.821495 ms2.797711 ms4.712874 ms
512K8.038367 ms6.007526 ms10.383386 ms
1M16.589313 ms12.527798 ms22.416753 ms
2M38.806603 ms25.994612 ms47.795221 ms
4M78.304628 ms52.795861 ms100.198899 ms
8M169.534526 ms109.631584 ms211.116117 ms

由于3个算法都是基于 n*log(2, n) 的,因此,我对 t / (n * log(2,n) 进行了归一化,对比如图: img.png

结论:

  1. 在 1K 以下时, SIMD sort 性能比 unstable sort 要更好。(数据量越小的场景,SIMD 版本提升越明显)。
  2. 在 16K 以下时, SIMD sort 性能比 stable sort 要更好,(同样,越小的场景,SIMD 版本提升越明显)
  3. unstable sort 性能比 stable sort 要好20-30%。(而且不需要额外的空间)

说老实话,得出这个结果有一点沮丧,也因此而翻阅了 rust sort 的文档和源码,了解了 driftsort 的设计思路,大概得出如下的结论:

  • driftsoft 内部有大量基于实际经验的优化(比如对少于20个元素的场景进行插入排序),部份函数也支持 auto vectorization。作为 rust 核心库精心 挑选的算法,要想打败它,还是有相当挑战的。
  • 二分排序法也是 stable 的,因此,和 Slice::sort 对比更合适一些。在16K 以下场景时,SIMD 版本还是有领先优势的。所以输的不算太难看。
  • 仔细分析了为什么 SIMD 版本没有预期的好,主要的原因是在 merge sort 部分,SIMD 对新增的每轮 merge 的开销基本上在 1.3-1.4ns/元素,而 quick sort(driftsoft的核心算法)在 0.5-0.8ns。 我推测的原因是, quick sort 在 进行 pivot partition 时,可以利用到当前数据的局部有序性, 而仅 move 更少的元素,而 merge sort 则每次都需要 move 全部的元素,未能利用到数据的局部有序性。
  • 这也颠覆了我之前的认知:我一直以为二分排序法 与 快速排序法,具有相同的时间复杂度,因此,他们的性能应该是一致的。但实际上,quick sort 还是更快一些。 (当然,在不追求 stable 的情况下,也可以减少额外的内存需求。stable sort 和 二分排序法,都需要 N/2 的辅助内存用于交换。)

其它:

  1. 在进行 ns 级的性能优化时,很容易陷入到过早优化的细节中,在进行这个实验过程中,我刚开始也陷入到过早优化的困惑中,后面果断放弃,先完成代码,再利用 samply 等 profile 工具进行性能对比,最后分析出性能瓶颈,最后发现,你可以充分的信赖 rust 的编译器优化,而最后的性能问题,只会时数据结构上的问题, 譬如,在这里,就是 merge sort 阶段的数据局部性的问题。
  2. 后续可以对比 intel x86-simd-sort 这个项目, 对比一下 intel 的解决方案(quick sort版本), 在这个基础上,可以评估编写一版 simd 的 quick sort。