[博客翻译]用CUDA实现排序算法


原文地址:https://ashwanirathee.com/blog/2025/sort2/


在之前关于排序算法的文章中,我探讨了一些基本的排序方法。这次,我想更进一步,使用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不擅长处理递归,因此我们需要将递归展开为循环。

merge_sort_comparison-1400.webp

一些重要问题:

  • 为什么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);
    }
  }
}

我们将递归转换为循环:

  • 外层循环使数组大小从1n呈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