【CUDA笔记】04 CUDA 归约，原子操作，Warp 交换

引言

前几节处理的问题，最终一个线程都会对应输出一个结果。今天将讨论的问题，多个线程将会对应输出一个，或者少于启动线程数量的结果，称之为归约（Reduction）。

几种归约方法的实现

以实现求和方法为例，下面列出几种归约的实现

1.使用 Cuda API 的原始操作

关键API

复制代码

float atomicAdd(float* address, float val);

这里 Cuda 也有提供其他数据类型的加法的原子操作，可以详见

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html?highlight=atomicAdd#atomicadd

示例核函数代码

复制代码

// 1. 直接使用原子操作进行归约
__global__ void reduction_atmoic(const float *gdata, float *out, int inputDataSize)
{
  // 获取当前线程 id
  size_t index = threadIdx.x + blockDim.x * blockIdx.x;
  if (index < inputDataSize)
  {
    atomicAdd(out, gdata[index]);
  }
}

补充

任何 原子操作 都可以借助 atomicCAS 来实现。例如"atomicAdd() for double-precision floating-point numbers is not available on devices with compute capability lower than 6.0 " ，可以借助单精度或这双精度浮点型的加法原子操作借助 atomicCAS 来实现。

复制代码

#if __CUDA_ARCH__ < 600
__device__ double atomicAdd(double* address, double val)
{
    unsigned long long int* address_as_ull =
                              (unsigned long long int*)address;
    unsigned long long int old = *address_as_ull, assumed;

    do {
        assumed = old;
        old = atomicCAS(address_as_ull, assumed,
                        __double_as_longlong(val +
                               __longlong_as_double(assumed)));

    // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
    } while (assumed != old);

    return __longlong_as_double(old);
}
#endif

2.借助 shared memory, 使用经典并行归约方法，最后依赖原子操作进行归约

复制代码

// 与前一个方法的区别仅在于最后使用原子操作来进行归约, 直接将所有的和加到 out 指向的内存空间
__global__ void reduce_atomicEnding(const float *gdata, float *out, int inputDataSize)
{
  // 申请 BLOCK_SIZE 个 float 大小的共享内存空间
  __shared__ float sdata[BLOCK_SIZE];
  // 获取当前全局线程 index
  size_t index = threadIdx.x + blockDim.x * blockIdx.x;
  // 获取当前 block 中的线程ID  tid, 并初始化对应地址 shared memory 的初始值
  size_t tid = threadIdx.x;
  sdata[tid] = 0;
  // 使用 grid stride load(这里因为启动的总的线程数量为 gridDim.x * blockDim.x, 所以以这个数为步长)
  while (index < inputDataSize)
  {
    sdata[tid] += gdata[index];
    index += blockDim.x * gridDim.x;
  }

  // 对 sdata 中的数据进行归约计算, for 循环下标以 blockDim.x 为起点
  for (size_t limit = blockDim.x / 2; limit > 0; limit >>= 1)
  {
    // 先进行数据同步
    __syncthreads();

    // tid 小于 limit 的数据才进行进一步的计算
    if (tid < limit)
    {
      sdata[tid] += sdata[tid + limit];
    }
  }

  // 将合并的结果协会 out 数组当中， 仅当tid 为 0 的线程做这一步
  if (tid == 0)
  {
    atomicAdd(out, sdata[0]);
  }
}

3. 借助 warp sharp 机制，进行归约

关键 API

复制代码

warpSize  
// This variable is of type int and contains the warp size in threads (see SIMT Architecture for the definition of a warp). Cuda 定义的内置变量，通常是 32, 由具体 GPU 架构有决定

T __shfl_down_sync(unsigned mask, T var, unsigned int delta, int width=warpSize);
// Copy from a lane with higher ID relative to caller.Calculates a source lane ID by adding delta to the caller's lane ID. The value of var held by the resulting lane ID is returned: this has the effect of shifting var down the warp by delta lanes.

示例核函数代码

复制代码

__global__ void reduction_warpSharp(const float *gdata, float *out, int inputDataSize)
{
  // 申请 32 个 float 大小的共享内存空间（相当于在 block 下又细分了一个层级）
  // 这里设置为 32 的原因是, block 中的线程数不会超过 1024，又 warpSize 的大小为 32， 所以一个 block 中的 warp 数量最多为 32
  //
  __shared__ float sdata[32];
  // 获取当前全局线程 index
  size_t index = threadIdx.x + blockDim.x * blockIdx.x;
  // 获取当前 block 中的线程ID  tid
  size_t tid = threadIdx.x;
  // warp sharp 方法要用到的一些参数的初始化, warp 的 val， mask, lane, warpID
  float val = 0.0f;
  unsigned mask = 0xFFFFFFFFU;         // 指定 warp 中的哪些线程要参与后面的__shfl_down_sync, 这里暂时先一值用这个值就行
  int lane = threadIdx.x % warpSize;   // 这里这个值的取值范围 就是 0 到 31， 指代在同一个 warp 中的线程的唯一标识符
  int warpID = threadIdx.x / warpSize; // block 中 warp 的唯一标识符
  // 也是先使用 gride stride loop 加载数据，对 val 进行初始化（初步加载数据顺便先求一次和）
  while (index < inputDataSize)
  {
    val += gdata[index];
    index += blockDim.x * gridDim.x;
  }
  // 第一次先 使用__shfl_down_sync 函数来求一次和, 将同一个 warp 中的线程进行求和， 都累加到一个 warp 里 lane 为 1 的线程当中。
  // for 循环 offset 以 warpSize / 2 为起点
  for (size_t offset = warpSize / 2; offset > 0; offset >>= 1)
  {
    val += __shfl_down_sync(mask, val, offset);
  }

  // 当 lane == 0, 将 val 的值赋值到 sdata 当中， 完了之后 同步一次数据
  if (lane == 0)
  {
    sdata[warpID] = val;
  }
  __syncthreads();
  // 对于此时 warpID == 0 的线程， 再通过 _shfl_down_sync 进行一次归约
  if (warpID == 0)
  {
    // 在条件当中，如果 tid < blockDim.x / warpSize, 则 将 val 赋值为 sdata[lane], 否则赋值为 0
    if (tid < blockDim.x / warpSize)
    {
      val = sdata[lane];
    }
    else
    {
      val = 0;
    }

    // 然后使用 __shfl_down_sync 进行最后一次 归约
    for (size_t offset = warpSize / 2; offset > 0; offset >>= 1)
    {
      val += __shfl_down_sync(mask, val, offset);
    }
    // 当 tid == 0， 使用 原子操作 将 块中计算的 和 val 累加到 out 指向的值当中
    if (tid == 0)
    {
      atomicAdd(out, val);
    }
  }
}

方法试验比较

实验中设置输入数组的大小为

复制代码

 const size_t N = 16ULL * 1024ULL * 1024ULL; // data size 16M
 const int BLOCK_SIZE = 256;

试验几次结果如下

复制代码

reduction_atmoic took 0.063000 seconds.
atomic sum reduction correct!
reduce_atomicEnding took 0.002000 seconds.
reduction w/atomic sum correct!
reduction_warpSharp took 0.001000 seconds.
reduction warp shuffle sum correct!

reduction_atmoic took 0.061000 seconds.
atomic sum reduction correct!
reduce_atomicEnding took 0.002000 seconds.
reduction w/atomic sum correct!
reduction_warpSharp took 0.002000 seconds.
reduction warp shuffle sum correct!

reduction_atmoic took 0.063000 seconds.
atomic sum reduction correct!
reduce_atomicEnding took 0.003000 seconds.
reduction w/atomic sum correct!
reduction_warpSharp took 0.001000 seconds.
reduction warp shuffle sum correct!

4.另外一种通过多次启动核函数的归约方式，以求最大值为例

复制代码

const size_t N = 8ULL * 1024ULL * 1024ULL; // data size
const int BLOCK_SIZE = 256;                // CUDA maximum is 1024

__global__ void reduce(float *gdata, float *out, size_t n)
{
  __shared__ float sdata[BLOCK_SIZE];
  int tid = threadIdx.x;
  sdata[tid] = 0.0f;
  size_t idx = threadIdx.x + blockDim.x * blockIdx.x;

  while (idx < n)
  { // grid stride loop to load data
    sdata[tid] = max(sdata[tid], gdata[idx]);
    idx += gridDim.x * blockDim.x;
  }

  for (size_t s = blockDim.x / 2; s > 0; s >>= 1)
  {
    __syncthreads();
    if (tid < s)
    {
      sdata[tid] = max(sdata[tid], sdata[tid + s]);
    }
  }
  if (tid == 0)
  {
    out[blockIdx.x] = sdata[0];
  }
}

基本和前面归约方法求和的类似，除了使用 max 方法来进行归约。

然后在 CPU 侧的代码逻辑，将分两次启动这个核函数，来实现最终的求解最大值。

假设总的比较的数量是 M * N，第一次先分配 M 个Block 启动核函数，每个 block 归出一个最大值，然后再配置 N 个 Block 再启动一次核函数，对剩下的 N 个数再归一次，取得最终的最大值。类似

复制代码

reduce<<<M, BLOCK_SIZE>>>(d_A, d_sums, M * N); // reduce stage 1
reduce<<<N, BLOCK_SIZE>>>(d_A, d_sums, N); // reduce stage 2

因为当前版本的 Cuda 提供了只整型取max 的原子操作，没有还没有直接的float 型的 max 浮点操作。否则就是借助文章前面提到的 atomicCAS 来实现浮点数的取最大值操作的原子操作。

小结

本节主要通过结合代码。介绍几种基本的归约方法的实现。

【CUDA笔记】04 CUDA 归约， 原子操作，Warp 交换

引言

几种归约方法的实现

1.使用 Cuda API 的原始操作

关键API

示例核函数代码

补充

2.借助 shared memory, 使用经典并行归约方法，最后依赖原子操作进行归约

3. 借助 warp sharp 机制， 进行归约

关键 API

示例核函数代码

方法试验比较

4.另外一种通过多次启动核函数的归约方式， 以求最大值为例

小结

【CUDA笔记】04 CUDA 归约，原子操作，Warp 交换

3. 借助 warp sharp 机制，进行归约

4.另外一种通过多次启动核函数的归约方式，以求最大值为例