cuda随笔1.基数排序

闲聊

其实笔者会写下这篇文章来源与一个问题：TOPK。
在写cpu端的串行代码时，往往会用k个大小的堆来解决这类问题（时间复杂度nlog(k)），当k较大时（k >> n）一般会采用快排解决（相对内存跨度太大的堆来说会有更好的cache命中率）应该会有很多这方面的优质讲解资源，这里简单提一下，如有错误，多谢指出。
当然上面聊的只是为了勾起你的回忆（如果有的话）。写下来介绍一下并行计算中解决topk的方式。
浏览很多深度学习框架会发现topk问题在并行计算中常用基数排序来解决，当然是因为它具有良好的并行性。

基数排序

1.cpu串行实现。
闲聊结束，来看看今天的主角。如何在cpu端实现这个10大排序之一的基数排序。
先讲一下思路：一个数的大小其实是由它的位数和每位上的数字决定的，我们可以从低位(高位也是一样的)按位比较这些数（用两个桶，一个装该位为1的数，一个装该位为0的数，有点像map的感觉）：（假设排序 1111，101，110，1，0这5个数）图来：

第一位排序（最低位），先将所有数据按顺序放入桶内，再将代表1的桶接到0桶后（注意每一步的输出为下一步的输入）。

第二位：

第三位：

第四位（结果）：

实现代码(时间复杂度显然为bitsize*N)：
在这里先提议下一个处理各种类型正负值的小技巧(这里以float为演示):
由于float是由符号位、指数位、尾数组成。对于正浮点而言按位比较就是比较结果（高位越大，该数越大），但对于负浮点而言刚好相反。又考虑到正负浮点的按位比较，所以将正浮点第一位置为1，负浮点全部取反，这样就能统一正负浮点的按位比较。当然最后要转回来。详细代码：

__device__ __host__ unsigned int convert(float v) {
  // if v >= 0 v |= 2^31 else v = ~(v) - 1,ensure anyif v >=0 or v < 0 by bitwise compare is vaild
  unsigned int cmp = *reinterpret_cast<unsigned int*>(&v);
  unsigned int ret = (cmp & (1<<31)) ?  ~(cmp): (cmp | 0x80000000);
  return ret;
}

__device__ __host__ float deconvert(unsigned int v) {
  unsigned int tmp = (v & (1 << 31)) ?  (v ^ 0x80000000) : ~(v);
  return *reinterpret_cast<float*>(&tmp);
}

cpu基数排序实现代码：

void RadixSortHost(float *v,unsigned int size) {
  unsigned int sort_tmp0[size],sort_tmp1[size]; 
  for(int bit = 0;bit < 32;bit++) {
    unsigned int mask = 1 << bit;
    unsigned int cnt0 = 0,cnt1 = 0;
    for (int i = 0;i < size;i++) {
      unsigned int elem = ((bit == 0) ?convert(v[i]):sort_tmp0[i]);
      if((elem & mask) != 0) {
        sort_tmp1[cnt1++] = elem;
      }
      else {
        sort_tmp0[cnt0++] = elem;
      }
    }
    for(int i = 0;i < cnt1;i++) {
      sort_tmp0[cnt0 + i] = sort_tmp1[i];
    }
  }
  for(int i = 0;i < size;i++) {
    v[i] = deconvert(sort_tmp0[i]);
  }
}

2.并行基数排序的实现
为啥要详细讲解基数排序的串行实现，相信很容易猜到：将基数排序分组在并归就是一个并行的基数排序。
2.1分组排序
本人在Cuda C编程权威指南1.并行规约分化+循环展开章节讲过内存合并（交叉读取），这里就不加赘述了。
举一个最简单的例子：

thread0使用基数排序排序data0和data2，thread1使用基数排序排序data1和data3。
对应代码(桶的存储也统一按照交叉的形式，这样线程间互不干扰、方便)：

for (unsigned int bit = 0;bit < 32;bit++) {
      unsigned int mask = 1<<bit;
      unsigned int cnt0 = 0,cnt1 = 0;
      for(unsigned int i = tid;i < n;i += blockDim.x) {
        unsigned int elem = (bit == 0 ? convert(data[i]):sort_tmp0[i]);
        if((elem&mask) != 0) {
          sort_tmp1[cnt1 + tid] = elem;
          cnt1 += blockDim.x;
        }
        else {
          sort_tmp0[cnt0 + tid] = elem;
          cnt0 += blockDim.x;
        }
      }
      for (unsigned int i = 0;i < cnt1;i += blockDim.x) {
        sort_tmp0[cnt0 + i + tid] = sort_tmp1[i + tid];
      }
  }

上面的代码和串行代码几乎一模一样（读取存储的方式为交叉，串行为顺序）。

优化:if((elem&mask) != 0)(可能会出现线程束分化，不过这里修改为下面的代码并未显著提升性能,应该是增加了指令，不过这个代码的性能瓶颈在并归，源码中未作修改)。

for (unsigned int bit = 0;bit < 32;bit++) {
      unsigned int mask = 1<<bit;
      unsigned int cnt[2] = {0,0};
      for(unsigned int i = tid;i < n;i += blockDim.x) {
        unsigned int elem = (bit == 0 ? convert(data[i]):sort[i]);
        unsigned int index_type = (mask&elem)>>bit;
        sort[cnt[index_type] + tid + index_type * n] = elem;
        cnt[index_type] += blockDim.x;
      }
      for (unsigned int i = 0;i < cnt[1];i += blockDim.x) {
        sort[cnt[0] + i + tid] = sort[i + tid + n];
      }
  }

还有一点，每个线程的排序可以放到local来做，可能可以提升性能（先从global复制到本地寄存器再做排序），不过有导致occupancy降低的风险(要注意寄存器的使用，若是线程内使用过多的寄存器同时活跃的线程束数量将受到限制,在最后的源码中也有这个版本在对N=1024*256时耗时减少10%,还是要说的是性能瓶颈在并归)
2.2 归并得出结果
按照上面那张图，归并操作其实就是合并两组有序序列了。是一个很简单的问题，用m个（这里其实就是两个）指针就能解决。(假设这两个有序序列为{1,3 },{2,4})

代码：

__shared__ unsigned int min_value, min_tid;
  __shared__ unsigned int list_idx[512];//512为blockDim.x，即block内的线程数
  unsigned int elem = 0xffffffff;//初始化全为unsigned int max
  list_idx[tid] = 0;//初始化所有指针指向序列头部
  __syncthreads();
  for(unsigned int i = 0;i < n;i++) {
    unsigned int x = (list_idx[tid] * blockDim.x) + tid;//tid 线程的第 list_idx[tid]个数据（经过基数排序从小到大的数据）
    if(x < n) {
      elem = sort_tmp0[x];//合法数据
    }
    else {
      elem = 0xffffffff;//非法数据，置为unsigned int max。意未不参与最小值比较（或最小值就为该值）
    }
    if(tid == 0) {
      min_value = min_tid = 0xffffffff;//block内第一个线程初始化min_value
    }
    __syncthreads();
    atomicMin(&min_value,elem);//原子操作，每个线程的序列头都与min_value比较决出最小值
    __syncthreads();


    if(min_value == elem) {//如果有多个最小值，取thread号最小的那个
      atomicMin(&min_tid, tid);
    }
    __syncthreads();
    if(min_tid == tid) {
      list_idx[tid]++;//序列头右移
      data[i] = deconvert(min_value);//得出第i个结果
    }
    __syncthreads();
  }

宏观来理解上面的代码：在一个block内并归（可以使用共享内存，并且同步方便，因为并归依赖之前的结果，只能按序并归），每个线程都管理着一个有m个数据的有序序列，我们的任务就是合并这些序列（回到上面那张图）。
每个线程都判断自己序列头的值是否为所有序列头中的最小值，是的话得出第i个结果（如果有重复选取最小thread号的那个）。

由于为了最后的并归方便这里只使用了一个block，可以使用更多的block去处理上面的分组排序，再并归。(说实话速度有点慢,有很大的改进空间)

完整代码(看到这不给个星星吗，求求了)：
代码

参考