GPU上的排序

排序方法的选择

这个其实很难说,在深入了解这块之前我只知道 Merge sort和Odd-Even sort可以并行。
特别是Odd-Even sort,GPU Gems2有一章单独介绍: Chapter 46. Improved GPU Sorting,但是很因缺思厅的是在 GPU Gems2 出来之前 GPGPU这个概念还没落地,因此整一个算法是run在fragment shader上的,结果就是其需要n多个Pass才能运行成功,听起来性能就很差
Merge sort嘛应该还可以,不过我没尝试这个

再看看其他人咋选的8:在NV官方的CUDA库CUB中有三种Radix sort的实现和一种Merge sort的实现
那肯定得选Radix sort啊UwU

欸那你为什么不用CUDA调库呢 似不似啥
因为我的这些东西要在Game Engine上run的那必须得全平台啊...
所以最后是需要写成Compute shader的捏QwQ

Radix sort思想

基数排序确实很简单...

  1. 按位序把每种数字分别放入不同的桶中
  2. 然后用一趟遍历串起来
  3. 接着遍历其他的位。

局限是不能对浮点数进行排序--不过也有Trick可以避免这个问题

但是这个算法无法并行,因此有了一些改进。


一个简单的演示:
对于数字串 <5 1 2 4 3>
其二进制为 101 001 010 100 011
对于最低位,其为 1 1 0 0 1,那么做一个Histogram:

|0 |1 |
|--|--|
|2 |3 |

对原来的序列做一个Exclusive Prefix Sum (不包含自身的前缀和)
-> 0,1,2,2,2

对Histogram做一个Exclusive Prefix Sum
-> 0, 2

那么由Histogram的前缀和可以知道: 对任意一个1前面必定有两个0(理解为箱子之间的偏移)
再由序列的前缀和可以知道某个1/0前面有几个1/0(理解为箱子内部的偏移)

那么自然就得到了第一趟的位置序列: 2, 3, 0, 1, 4 => 于是第一趟的排序结果就是2, 4, 5, 1, 3


这是一个基数为2的例子,这种情况下对于Uint32需要进行32趟排序,当然也可以利用4为基数,那就只需要16趟了!
为什么说这种算法比 Naive 的Radix sort方便并行呢?
回顾一下,发现其中耗时最多的操作应该是“前缀和”,而并行前缀和其实是一个Solved problem

Parallel Prefix Sum

首先需要介绍并行下的几个基础操作:
Reduce, Scan

Reduce

reduce就是一个有着两个输入的函数:一个数组和一个满足结合律的操作符
$$
R(x,\oplus)
$$

在这里我们定义:数组x为[1,2,3,.....,n],操作符$\oplus$为+
那么这个Reduce就等价于:
1+2+3+....+n

可以简单写出伪代码:

sum = 0
for i in x: 
    sum = sum + i

也就可视作为$(((1+2)+3)+.....)$
如图

由于加法有结合律,所以自然的我们能想出更并行的方案
$((1+2)+(3+4)) + ((...)+(...))$
也就是

简单写出伪代码:

#define TID: current thread ID   
#define Input: input elements, with size equal to total threads number
offset = 1
Parallel for each Thread TID:
    for i  = Input.size / 2 to 1 step i = i / 2:
        SynWithOtherThreads()
        if TID < i:
            temp = Input[TID] + Input[i+TID]
            Input[TID+i] = temp
        offset = offset * 2

但是如图,这种方法实际上是构建了一颗满二叉树
所以当且仅当数组长度为2的n次幂时才能使用。

Scan

Scan就是Reduce的序列,记为
$S(x,\oplus)$
那么Scan操作就是:
$$ Input: [a_0,a_1,a_2]
$$ $$ Output: [a_0,a_0\oplus a_1,a_0\oplus a_1 \oplus a_2, ....] $$

串行的Scan伪代码就是:

acc = identitiy   
for i in x:  
    acc = acc op i
    Output[idx] = acc

Scan的并行一看就很难的样子,因为后一个操作完全依赖于上一个...但也不是没有办法!

Hillis and Steele Scan

直接举例来说明会更好:

其实就是个错位相加....灰常简单,每次的错位step会翻倍...

#define TID: current thread ID   
#define Input: input elements, with size equal to total threads number

Parallel for each Thread TID:
    for i = 0 to Input.size, i*=2 :
        SynWithOtherThreads()
        temp = Identity
        if TID <= i:
            temp = Input[TID]
        else 
            temp = Input[TID - i] op Input[TID]
        SynWithOtherThreads()
        Input[TID] = temp

Blelloch Scan

这种Scan方法又被称作为Reduce-Then-Scan,因为其包含两个阶段,Upsweep(Reduction)和Downsweep 还是用栗子来说明

在上扫阶段,其会执行Reduce时介绍的策略,而下扫则类似于上扫的逆运算,不过其运算为: $$ Output2 = Input1 \oplus Input2 $$ $$ Output1 = Input2 $$

简单写出伪代码:

#define TID: current thread ID   
#define Input: input elements, with size equal to total threads number
Parallel for each Thread TID:
    offset = 1
    for i  = Input.size / 2 to 1 step i = i / 2:
        SynWithOtherThreads()
        if TID < i:
            temp = Input[TID] + Input[i+TID]
            Input[TID+i] = temp
        offset = offset * 2
    
    if TID == 0 :
        Input[Input.size - 1] = 0
    
    for i  = 1 to n step i = i * 2:
        SynWithOtherThreads()
        offset = offset / 2
        if TID < i:
            offset1 = offset * (2 * TID + 1) - 1
            offset2 = offset * (2 * TID + 2) - 1
            Input1 = Input[offset1]
            Input2 = Input[offset2]
            Input[offset1] = Input2 
            Input[offset2] = Input1 + Input2
         

这种方法同时也具有Reduce的限制:数组长度必须小于2的n次幂

两种Scan的对比

把并行算法的复杂度分为两个维度:
Step complexity和Work Complexity
Step可以理解为一个线程需要跑几次循环
Work代表所有线程总共干了多少任务

对于Hillis and Steele Scan ,很明显其Step就是O(logn),而Work是O(nlogn)

对于Blelloch Scan,其出现了两次logn循环,于是Step为O(logn),而Work为O(n)
于是乎:Hillis and Steele是Step Efficient,而 Blelloch Scan是work efficient
这代表着:

  1. 如果存在 处理器数量 >= 任务 的情况,应当选择Step Efficient的Blelloch Scan,因为此时可以一个work对应一个处理器,因此瓶颈不在任务数量上
  2. 如果 任务 > 处理器数 ,那么应当选择work efficient的Blelloch Scan,因为此时任务过多,执行速度受到了处理器数量的限制。

回到主话题上

那么,可以用以上我介绍的两种方法来计算前缀和,这就解决Radix sort在并行排序下最大的问题啦!
剩下的算每个元素的offset那都是简简单单了嘿嘿

Further Reading

实际上Blelloch Scan可以通过一定的技巧处理非2的幂次情况,详见:Guy E. Blelloch. “Prefix Sums and Their Applications”
上面提到的前缀和部分实际上还能进行加速,在2022年NV的研究员提出了一种称作为Onesweep的算法,通过一种叫chain-scan的方法把Blelloch Scan的两次sweep压到了一次,详见:Onesweep: A Faster Least Significant Digit Radix Sort for GPUs
但是Onesweep貌似需要一些硬件上的额外工作,也就是说它的可移植性不高(或者说手机用不了...硬件厂商随便一点幺蛾子就能要我命

加之我现在有点困,所以等有空我再更新Zzzzz

Reference

Parallel Prefix Sum (Scan) with CUDA

AMS 148 Chapter 5: Reduce and Scan

END