在之前关于排序算法的文章中,我探讨了一些基本的排序方法。这次,我想更进一步,使用CUDA来实现这些算法,并探索通过并行计算提高性能的可能性。目标是了解如何利用并行计算的能力加速我们的排序算法。
几天前我参加了NVIDIA的一次招聘活动,那是一次很棒的经历,它激励我去尝试用CUDA重写这些排序算法。在这里,我将以归并排序(Merge Sort)为例,因为它可以很好地将问题划分为两个等长的子问题,非常适合并行计算。
基本递归归并排序
首先,我们来看传统的自顶向下的递归归并排序逻辑。这个过程通过不断将数组分成两半,直到达到单个元素为止,然后将已排序的子数组重新合并。
要合并两个已排序的数组,我们需要比较它们的起始元素,选择较小的那个放入输出数组中,并将对应的指针向前移动。
MERGE_SORT(arr, left, right)
IF left < right THEN
mid ← left + (right - left) / 2
// 递归地对前半部分进行排序
MERGE_SORT(arr, left, mid)
// 递归地对后半部分进行排序
MERGE_SORT(arr, mid + 1, right)
// 合并已排序的两部分
MERGE(arr, left, mid, right)
ENDIF
END MERGE_SORT
注意事项:
-
函数签名说明:
void merge(uint8_t* arr, uint8_t* temp, long long left, long long mid, long long right)
使用uint8_t
代替int
是为了让数组值保持在较小范围(0-255)。long long
用于索引,支持非常大的数组(例如10^18)。uint8_t* temp
作为临时存储空间,能带来性能提升。void mergeSort(uint8_t* arr, uint8_t* temp, long long left, long long right)
遵循伪代码,将数组划分为两半并分别调用自身。到达单个元素时,调用merge
函数完成合并。
-
GPU与CPU排序的区别:
- 数组根据特定种子生成(例如种子为1)。
- 所有实现方式大致完成相同的计算量。
- 归并排序的结果需要从GPU传回CPU,这会造成开销。
- 更好的对比方式是在GPU上直接对随机数组排序并比较结果。
- 排序方式和排序地点的选择会因应用场景而异。
-
运行时间指的是整个程序的时间,而非仅限于排序部分的时间。
-
正确性检查通过
std::sort
对原始数组进行排序并与结果比较。 -
时间复杂度:O(n log n),空间复杂度:O(n)。
CUDA上的基本递归归并排序
现在让我们看看如何在CUDA中实现归并排序。这个实现与CPU版本的模式类似,是我的第一个CUDA版本。内核在每个合并操作时被启动,递归仍然由CPU执行。
注意事项:
- 包括CUDA Runtime API在内的头文件提供了访问
cudaMalloc()
、cudaMemcpy()
、cudaFree()
、kernel<<<numBlocks, threadsPerBlock>>>(args)
等功能的方法。 - 内核函数
__global__ void mergeSort(uint8_t* arr, uint8_t* temp, long long left, long long right)
目前的操作与CPU实现完全相同。 - 在
mergeSort
内部:- 每次合并操作时调用
merge<<<1, 1>>>(...)
,目前只启动一个线程来处理整个合并过程,这种方式效率较低。 - 线程块数量和每块线程数由
<<<1,1>>>
指定。总线程数为numBlocks * blockSize
,它们可以以一维、二维或三维网格形式排列。 - 使用
cudaDeviceSynchronize()
确保当前合并操作完成后再进入下一步,以避免结果出错。 - 使用
cudaMalloc(...)
分配GPU内存,cudaMemcpy(..., cudaMemcpyHostToDevice)
和cudaMemcpy(...., cudaMemcpyDeviceToHost)
可将数据在CPU和GPU之间复制。 - 使用
cudaFree(cu_arr)
释放GPU上的内存。
- 每次合并操作时调用
CPU与GPU实现的基本递归归并排序对比
可以看到,在图1中,由于每次合并操作都启动了内核,且递归由CPU执行,这种方法效率并不高。CUDA不擅长处理递归,因此我们需要将递归展开为循环。
一些重要问题:
-
为什么CUDA不擅长递归?
- 当前的合并操作只启动了一个GPU线程,而递归则在CPU上完成。深度递归容易导致栈溢出,因为GPU的线程数有限。此外,每次启动内核都会产生一定的开销,递归限制了并行度,同步也是一个问题。
-
如何改进?
- 将递归改写为迭代,并实现自底向上的归并排序。
自底向上的迭代归并排序
由于CUDA难以高效处理递归,因此我们改为实现一种迭代版的归并排序。该方法的核心是从最小的子数组开始,自底向上地逐步合并。
MERGE_SORT(arr, temp, start, end)
FOR sub_size ← 1 TO end STEP 2 × sub_size DO
FOR left ← 0 TO end STEP 2 × sub_size DO
mid ← MIN(left + sub_size - 1, end)
right ← MIN(left + 2 × sub_size - 1, end)
MERGE(arr, temp, left, mid, right)
ENDFOR
ENDFOR
END MERGE_SORT
注意事项:
void mergeSort(uint8_t* arr, uint8_t* temp, long long n) {
long long left, mid, right, size;
for (size = 1; size < n; size *= 2) {
for (left = 0; left < n - size; left += 2 * size) {
mid = left + size - 1;
right = std::min(left + 2 * size - 1, n - 1);
mergeKernel(arr, temp, left, mid, right);
}
}
}
我们将递归转换为循环:
- 外层循环使数组大小从
1
到n
呈2的幂增长,即1, 2, 4, 8...
。即便数组大小不是2的幂,右边界会被夹到数组末尾,这样也能正常处理。 - 内层循环遍历数组,按步长
2 * size
依次合并大小为size
的子数组。 mergeKernel
功能与递归版本中的merge
相同,但现在通过循环调用。
CUDA上的自底向上的迭代归并排序
以下是我个人在这个实现中学到的主要内容。上面提到的实现有两个循环,我最初的想法是将第二个循环在GPU上并行化,以同时处理整个数组的合并操作。
void mergeSort(uint8_t* arr, uint8_t* temp, long long n) {
bool flipflop = tru