CUDA: reduce算子的实现

说在前面：本篇博客主要介绍reduce算子在cpu和gpu上的多种代码实现。

一. cpu实现

        reduce在cpu上的实现思路相对简单，就是遍历所有元素逐个相加就可以获得结果，但这种实现明显效率比较低。

float reduce(const float *x, const int N)
{
    real sum = 0.0;
    for(int i=0; i<N; i++)
    {
        sum += x[i];
    }
    return sum;
}

二. gpu实现

1. native版本

        思路

        1. 折半归约法，在每个block内进行归约求和
        2. 每次将矩阵的后半段加给前半段
        3. 相加的过程相当于是两个向量相加，很容易用GPU实现
        4. 用cpu进行收尾，将长度为N / 1024的d_y数组进行求和得到sum值

void __global__ reduce_global(float *d_x, float *d_y)
{
    const int tid = threadIdx.x;
    float *x = d_x + blockDim.x * blockIdx.x;
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1)
    {
        if (tid < offset)
        {
            x[tid] += x[tid + offset];
        }
        __syncthreads();
    }
    if (tid == 0)
    {
        d_y[blockIdx.x] = x[0];
    }
}

        这种方法的缺点在于全局内存访问比较耗时，并且最后的结果并不是最终值，还需要在cpu上做额外的计算。

2. 优化版本（共享内存+warp洗牌函数+atomicAdd）

        思路

        1. 数组长度 block_size->32 的范围， 使用共享内存来减少全局内存的访问次数
        2. 数组长度 32 ->1 的范围，使用洗牌函数来减少共享内存的访问次数
        3. 使用原子加来替代cpu的收尾

        (个人认为还可以将下面的for循环展开，这样可以避免部分运算和逻辑判断，因为每个线程都要去做这些操作，所以还是存在一定的性能影响的，这里我们没有实现，大家可以尝试下)

void __global__ reduce_shfl(const float *d_x, float *d_y, const int N)
{
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  const int n = bid * blockDim.x + tid;

  extern __shared__ float s_y[];
  s_y[tid] = (n < N) ? d_x[n] : 0.0;
  __syncthreads();

  for (int offset = blockDim.x >> 1; offset >= 32; offset >>= 1)
  {
      if (tid < offset)
      {
          s_y[tid] += s_y[tid + offset];
      }
      __syncthreads();
  }

  float y = s_y[tid];
  for (int offset = 16; offset > 0; offset >>= 1)
  {
      y += __shfl_down_sync(FULL_MASK, y, offset);
  }

  if (tid == 0)
  {
      atomicAdd(d_y, y);
  }
}

        解释

void __global__ reduce_shfl(const float *d_x, float *d_y, const int N)

          d_x是输入的需要进行规约求和的数据，d_y是用于保存求和后的结果的数据。N是需要进行规约的数据的长度。

  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  const int n = bid * blockDim.x + tid;

          tid是block中的线程的局部id，bid是当前block的id，n是线程的全局id。

  extern __shared__ float s_y[];
  s_y[tid] = (n < N) ? d_x[n] : 0.0;
  __syncthreads();

        extern是指动态分配共享内存的大小，s_y[tid] = (n < N) ? d_x[n] : 0.0是将数据转移到共享内存中，超出范围赋值为0.0。__syncthreads()是同步block中的所有线程。

for (int offset = blockDim.x >> 1; offset >= 32; offset >>= 1)
  {
      if (tid < offset)
      {
          s_y[tid] += s_y[tid + offset];
      }
      __syncthreads();
  }

        将每个block中的数据进行折半规约直到每个blcok的元素只有32个。 

  float y = s_y[tid];
  for (int offset = 16; offset > 0; offset >>= 1)
  {
      y += __shfl_down_sync(FULL_MASK, y, offset);
  }

        剩余的block中的32个元素，我们使用warp的洗牌函数进行规约提升效率。思路依然是折半规约。

  if (tid == 0)
  {
      atomicAdd(d_y, y);
  }

        我们之前的计算了每个block的所有元素和，最后使用atomicAdd将每个block的元素和进行累加。这里使用原子加可以避免线程安全问题。