CUDA与大数组的双调排序

CUDA 与大数组的双调排序

双调排序

双调排序，又称作 Bitonic Sort，整体复杂度为 $n{\log }^2(n)$。

归并排序劣势

相比于归并排序的 $n\log (n)$ 的复杂度，看起来优势不大，但是胜在能够比归并排序更加有效的运用多核多线程，因此在 CUDA 加持下，效率往往快于归并排序。

归并排序的缺点： 当归并排序到后期时，待归并的段主要有少数的“大段”构成，而这几个大段的归并比较难运用到多核多线程的优势。尽管可以通过一定手段优化几个大段的优势，但是实现起来比较复杂。

既然归并这么复杂，为什么不用实现起来更简单的双调排序呢？

双调排序原理

双调排序的原理，许多博主都有详细的说明，其中关键的性质就是 Batcher 定理。

双调序列

双调序列(Bitonic Sequence)是指由一个非严格增序列 $X$ 和非严格减序列 $Y$ 构成的序列，比如序列（23,10,8,3,5,7,11,78）。

定义:一个序列 $a_1,a_2,\dots ,a_n$ 是双调序列(Bitonic Sequence)，如果：

1. 存在一个$a_k(1\le k\le n)$, 使得 $a_1\ge \dots \ge a_k\le \dots \le a_n$ 成立；或者
2. 序列能够循环移位满足条件 1

Batcher 定理

将任意一个长为 $2n$ 的双调序列 $A$ 分为等长的两半 $X$ 和 $Y$，将 $X$ 中的元素与 $Y$ 中的元素一一按原序比较，即 $a[i]$ 与 $a[i+n],(i<n)$比较，将较大者放入 MAX 序列，较小者放入 MIN 序列。则得到的 MAX 和 MIN 序列仍然是双调序列，并且 MAX 序列中的任意一个元素不小于 MIN 序列中的任意一个元素。

双调排序详解

不过多赘述了，附上两篇挺好的博客：

• 三十分钟理解：双调排序Bitonic Sort，适合并行计算的排序算法

• 【并行计算】Bitonic Sort（双调排序）基础

CUDA 代码实现

实际代码

#include <stdio.h>
const char* version_name = "Naive, sort.";

__device__ void swap_float(float* f1, float* f2) {
    float tmp = *f1;
    *f1 = *f2;
    *f2 = tmp;
}

__global__ void _bitonic_sort(float* d_arr, unsigned stride, unsigned inner_stride) {
    unsigned flipper = inner_stride >> 1;
    unsigned tid = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned tid_other = tid ^ flipper;

    if (tid < tid_other) {
        // 操纵左侧的半部分
        if ((tid & stride) == 0) {
            // 此处将留升序
            if (d_arr[tid] > d_arr[tid_other]) {
                swap_float(&d_arr[tid], &d_arr[tid_other]);
            }
        } else {
            // 此处将留降序
            if (d_arr[tid] < d_arr[tid_other]) {
                swap_float(&d_arr[tid], &d_arr[tid_other]);
            }
        }
    }
}

/// entry point for gpu float sorting
/// \param arr memory on gpu device
/// \param len the length of the array
void float_sort(float arr[], int len) {
    // 首先检查长度是否为 2 的幂
    unsigned twoUpper = 1;
    for (;twoUpper < len; twoUpper <<= 1) {
        if (twoUpper == len) {
            break;
        }
    }

    // 如果是 host 指针，返回
    cudaPointerAttributes attrs;
    cudaPointerGetAttributes(&attrs, arr);
    if (attrs.type != cudaMemoryTypeDevice) {
        return;
    }

    float* d_input_arr;
    unsigned input_arr_len;
    if (twoUpper == len) {
        input_arr_len = len;
        d_input_arr = arr;
    } else {
        // 需要 padding
        input_arr_len = twoUpper;
        cudaMalloc(&d_input_arr, sizeof(float) * input_arr_len);
        // 然后初始化
        cudaMemcpy(d_input_arr, arr, sizeof(float) * len, cudaMemcpyHostToDevice);
        cudaMemset(d_input_arr + len, 0x7f, sizeof(float) * (input_arr_len - len));
    }

    dim3 grid_dim((input_arr_len / 256 == 0)? 1 : input_arr_len / 256);
    dim3 block_dim((input_arr_len / 256 == 0)? input_arr_len : 256);
    
    // 排序过程(重点)
    for (unsigned stride = 2; stride <= input_arr_len; stride <<= 1) {
        for (unsigned inner_stride = stride; inner_stride >= 2; inner_stride >>= 1) {
            _bitonic_sort<<<grid_dim, block_dim>>>(d_input_arr, stride, inner_stride);
        }
    }

    // 如果 padding 过，则此处还原
    if (twoUpper != len) {
        cudaMemcpy(arr, d_input_arr, sizeof(float) * len, cudaMemcpyDeviceToDevice);
        cudaFree(d_input_arr);
    }
}

代码解析与例子分析

函数 float_sort(float arr[], int len) 是我们的排序函数主函数体。

在这里，float arr[] 必须是由 CUDA 分配的显存空间，而且已经通过 cudaMemcpy 装载好了待排序的数据。如果传入了内存空间，会直接 return。

由于原始的双调排序只能对 $2^n$ 的数据排序，因此对于长度不为 $2^n$ 的数组，先进行 padding。

在排序过程中，会调用 _bitonic_sort(float* d_arr, unsigned stride, unsigned inner_stride)，而这是双调排序的主体环节。
float* d_arr 为 padding 后的数组。

unsigned stride 表示当前双调排序的步长，下面举一个详细的例子：

当双调序列为 $1,2,3,4,5,3,2,1,3,4,5,6,7,8,7,3$ 且步长为 8 时，表示以步长为 8 进行分组，每一个组中都为一个双调序列，在这里就是两组：$1,2,3,4,5,3,2,1$ 和 $3,4,5,6,7,8,7,3$。

之后，我们想把这两个序列处理为前者递增，后者递减，使整个长度为 16 的数组变成一个双调序列，因此之后我们对这两组序列分别进行双调排序，只不过排出来的结果中，前者为升序，后者后降序。

unsignd inner_stride 表示对一个已经通过 stride 分组的双调序列进行排序。此处继续沿用上述的列子，并将后半段 $3,4,5,6,7,8,7,3$ 排序为降序序列。在 float_sort() 函数体中，可见 inner_stride 从 stride 依次指数减小到 2。

1. inner_stride = 8，对长度为 8 的后半段分为两节，使分出来的两节中，左半节大于右半节，则依靠 Batcher 定理得到两节 $7,8,7,6$ 和 $3,4,5,3$（注意这里 $7,8,7,6$ 移位后是 $7,6,7,8$，仍为双调序列，满足 Batcher 定理）
2. inner_stride = 4，这里为简便，只处理 $7,8,7,6$，另一节 $3,4,5,3$ 类似。将长度为 4 的 $7,8,7,6$ 根据 Batcher 定理分为两节，得到 $7,8$ 和 $7,6$
3. inner_stride = 2，之后对长度为 2 的 $7,8$ 和 $7,6$ 分别分为两节并排序，得到 $8$，$7$，$7$ 和 $6$，此时后半段的前半节就已经排好序为 $8,7,7,6$ 的降序序列了。同理，后半段的后半节排序为 $5,4,3,3$。
4. 此时，后半段就排序为 $8,7,7,6,5,4,3,3$，成为降序序列。同理，对前半段排序，得到升序序列 $1,1,2,2,3,3,4,5$，之后得到新的数组 $1,1,2,2,3,3,4,5,8,7,7,6,5,4,3,3$，可见又是新的双调序列。

当 stride = 8 时，从 $1,2,3,4,5,3,2,1,3,4,5,6,7,8,7,3$ 转为 $1,1,2,2,3,3,4,5,8,7,7,6,5,4,3,3$。之后，只需要对 stride = 16 重新执行上述过程，就能得到升序的最终序列了。

总结

上面我们通过一个例子，详细理解了代码以及思想。

这个代码没有用到 device 中的 share_memory，是因为数据大小可能非常巨大，share_memory 无法存下整个数组，因此只能通过 global 共享显存地址，并通过 float_sort() 中的循环进行 threadBlock 之间的同步。