CUDA 经典问题:数组归约 | CUDA

本文最后更新于:2022年1月30日

归约(reduction)是指一种解题思路,是指将某个计算问题变换为另一个问题的过程。而数据归约计算则是指使用归约的方法求一个数组总和/最大值/最小值/均值等的操作。

求和

块内归约

下面给出的两个例子都是仅在 block 内完成归约,并没有完整实现求数组元素加和。

使用全局内存

__global__ void ReduceSumKernel(int *in, int *out) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;

  for (int i = blockDim.x / 2; i >= 1; i /= 2) {
    if (threadIdx.x < i) {
      in[tid] += in[tid + i];
    }
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    out[blockIdx.x] = in[offset];
  }
}

这个函数实现要求全局内存数组的长度是线程块大小的整数倍。

循环语句优化

针对在循环语句可以做两个优化:

  • 在循环语句中使用常量
  • 用位运算代替除法操作
for (int i = BLOCK_SIZE >> 1; i >= 1; i >>= 1) {
  // 循环体
}

使用共享内存

__global__ void ReduceSumKernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;

  __shared__ int buffer[BLOCK_SIZE];

  if (tid < n) {
    buffer[threadIdx.x] = in[tid];
  } else {
    buffer[threadIdx.x] = 0;
  }
  __syncthreads();

  for (int i = BLOCK_SIZE >> 1; i >= 1; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    out[blockIdx.x] = buffer[0];
  }
}

使用共享内存相对于仅使用全局内存有两个好处:

  • 不再要求全局内存数组的长度是线程块大小的整数倍
  • 在归约的过程中不会改变全局内存数组中的数据

全局归约

在前面几个版本的实现中,核函数并没有做全部的计算,只是将一个长一些的数组 in 变成了一个短一些的数组 out,后者中的每个元素为前者中若干元素的和。在调用这些核函数之后,将短一些的数组复制到主机,然后在主机中完成了余下的求和。

使用原子函数

__global__ void ReduceSumKernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;

  __shared__ int buffer[BLOCK_SIZE];

  if (tid < n) {
    buffer[threadIdx.x] = in[tid];
  } else {
    buffer[threadIdx.x] = 0;
  }
  __syncthreads();

  for (int i = BLOCK_SIZE >> 1; i >= 1; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncthreads();
  }

  if (threadIdx.x == 0) {
    atomicAdd(out, buffer[0]);
  }
}

使用 __syncwarp() 函数

__global__ void ReduceSumKernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;

  __shared__ int buffer[BLOCK_SIZE];

  if (tid < n) {
    buffer[threadIdx.x] = in[tid];
  } else {
    buffer[threadIdx.x] = 0;
  }
  __syncthreads();

  for (int i = BLOCK_SIZE >> 1; i >= 32; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncthreads();
  }

  for (int i = 16; i >= 1; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncwarp();
  }

  if (threadIdx.x == 0) {
    atomicAdd(out, buffer[0]);
  }
}

在归约问题中,当所涉及的线程都在一个线程束内时,可以将线程块同步函数 __syncthreads() 换成一个开销更小的线程束内同步函数 __syncwarp()

使用 __shfl_down_sync() 函数

__global__ void ReduceSumKernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;

  __shared__ int buffer[BLOCK_SIZE];

  if (tid < n) {
    buffer[threadIdx.x] = in[tid];
  } else {
    buffer[threadIdx.x] = 0;
  }
  __syncthreads();

  for (int i = BLOCK_SIZE >> 1; i >= 32; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncthreads();
  }

  int t = buffer[threadIdx.x];
  for (int i = 16; i >= 1; i >>= 1) {
    t += __shfl_down_sync(0xffffffff, t, i);
  }

  if (threadIdx.x == 0) {
    atomicAdd(out, t);
  }
}

在进行线程束内的循环之前,这里将共享内存中的数据复制到了寄存器。在线程束内使用洗牌函数进行归约时,不再需要使用共享内存。因为寄存器访问延迟比共享内存更低,所以能用寄存器尽量用寄存器。

提高线程利用率

template<int block_size>
__global__ void ReduceSumKernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;
  int stride = blockDim.x * gridDim.x;

  __shared__ int buffer[block_size];

  // 在上一个版本中,共享内存数组中的每一个元素仅仅保存了一个全局内存数组的数据
  // 为了提高归约之前所做计算的比例,可以在归约之前将多个全局内存数组的数据累加到一个共享内存数组的一个元素中
  // 如果一个线程处理相邻的几个数据,会导致全局内存的非合并访问,所以必须让相邻的线程访问相邻的数据
  // 这就意味着同一个线程访问的数据之间有一个跨度,这里使用整个 grid 的大小作为跨度
  int t = 0;
  for (int i = tid; i < n; i += stride) {
    t += in[i];
  }
  buffer[threadIdx.x] = t;
  __syncthreads();

  for (int i = block_size >> 1; i >= 32; i >>= 1) {
    if (threadIdx.x < i) {
      buffer[threadIdx.x] += buffer[threadIdx.x + i];
    }
    __syncthreads();
  }

  t = buffer[threadIdx.x];
  for (int i = 16; i >= 1; i >>= 1) {
    t += __shfl_down_sync(0xffffffff, t, i);
  }

  if (threadIdx.x == 0) {
    // 这里没有使用原子函数的方式完成加和,而是调用核函数两次
    out[blockIdx.x] = t;
  }
}

void ReduceSum() {
  // 在使用两个核函数时,将第一次计算结果归约到最终结果的计算也使用了折半求和,比直接累加要更精确,避免数字太大『吃』掉小一点的数
  ReduceSumKernel<BLOCK_SIZE><<<BLOCK_NUM, BLOCK_SIZE>>>((int *)(gpu_data), (int *)(gpu_data), N);
  ReduceSumKernel<1024><<<1, 1024>>>((int *)(gpu_data), (int *)(gpu_data), BLOCK_NUM);
}

向量 top2

向量 top2 问题即求包含 N 个元素的数组的最大的两个值。

这道题包括两个阶段:

  • 每个线程把本线程读入的数据中最大的两个值按照顺序写入共享内存
  • 与其他线程合作,求出其所在的 block 内共享内存中所有数据的两个最大值,并写入到输出全局内存

Baseline

__global__ void VecTop2Kernel(int *in, int *out, int n) {
  int offset = blockIdx.x * blockDim.x;
  int tid = offset + threadIdx.x;
  int stride = blockDim.x * gridDim.x;

  extern __shared__ int lich[];

  // 每个线程先读取一个数据,并将数据写入共享内存
  // 在该核函数被第一次调用时,每个线程中的 top1 为它读取的数,top2 为 INT_MIN
  int top1 = INT_MIN;
  int top2 = INT_MIN;

  for (int i = tid; i < n; i += stride) {
    if (in[i] > top2) {
      top2 = min(in[i], top1);
      top1 = max(in[i], top1);
    }
  }
  lich[2 * threadIdx.x] = top1;
  lich[2 * threadIdx.x + 1] = top2;
  __syncthreads();

  // 每次调用一半的线程,每个线程读取四个数据,从中选出两个最大的值,把这两个最大的值写入共享内存中的相应位置
  // 每次循环都筛选掉一半数字,i 是每次循环使用的线程数,c、d 与 a、b 错开了 2 * i 个数
  for (int i = blockDim.x / 2; i >= 1; i /= 2) {
    if (threadIdx.x < i) {
      int a = lich[2 * threadIdx.x];
      int b = lich[2 * threadIdx.x + 1];
      int c = lich[2 * threadIdx.x + 2 * i];
      int d = lich[2 * threadIdx.x + 1 + 2 * i];
      top1 = max(a, c);
      top2 = min(max(a, d), max(b, c));

      lich[2 * threadIdx.x] = top1;
      lich[2 * threadIdx.x + 1] = top2;
      __syncthreads();
    }
  }

  // 将每个 block 中最大的两个值写入到全局内存的相应位置
  if (offset < n) {
    if (threadIdx.x == 0) {
      out[2 * blockIdx.x] = lich[0]