Reduce 是一个高阶函数,它接收一个函数作为参数,这个函数接收两个参数,然后返回一个值。Reduce 会从左到右依次对数组中的元素进行处理,最终得到一个值。以加法为例,就相当于给定一个数组,对数组求和。用 cpu 实现起来非常简单,一层循环解决问题。本文主要讨论如何用纯函数式的方式实现 Reduce。
// reduce cpu versionint reduce(int *arr, int len) { int sum = 0; for (int i = 0; i < len; i++) { sum += arr[i]; } return sum;}
可以看到,cpu 实现非常简单,一层循环就可以解决问题。但是,这个实现并不是纯函数式的,因为它有副作用,即修改了 sum 的值。我们可以用 GPU 来实现一个纯函数式的 Reduce。
首先我们先来回归一下 GPU 的层次结构。也就是代码中的 block 和 grid。block 是一个线程块,grid 是一个线程网格。block 中的线程可以通过 threadIdx.x 来获取自己的线程 id,grid 中的线程可以通过 blockIdx.x 来获取自己的线程 id。block 中的线程可以通过 blockDim.x 来获取 block 的大小,grid 中的线程可以通过 gridDim.x 来获取 grid 的大小。
对于 Reduce 来说,我们可以按照下面这个图设计:

img1
我们把数组分成了若干个 block,每个 block 中有若干个线程。每个线程负责处理一个元素。每个 block 的线程会把自己处理的元素的值累加到 block 的第一个线程中。最后,每个 block 的第一个线程会把自己 block 中的所有元素的值累加到 grid 的第一个线程中。最后,grid 的第一个线程就是我们要求的结果。一般我们会把 block 的大小设置为 32,但是图中为了方便演示,我们把 block 的大小设置为 8。
GPU 的计算过程如下图所示:

img2
以上图为例,我们来看一下 Reduce 的计算过程。首先,我们把数组分成了 3 个 block,每个 block 中有 8 个线程。在第一轮计算中,奇数线程会把自己的值累加到偶数线程中。在第二轮计算中,block 中的第 0 个线程会把 4 号线程的值累加到自己的值中。 每个 block 的值都计算完之后还需要对 block 的值进行累加,下面我们来看一下代码要如何实现。
我们首先看 Kernel 的实现,代码如下:
template <int BLOCKSIZE>__global__ void reduce_naive_kernel(int *arr, int *out, int len){ __shared__ int sdata[BLOCKSIZE]; int tid = threadIdx.x; // 线程 id (block 内) int bid = blockIdx.x; // block id (grid 内) int bdim = blockDim.x; // block 大小 int i = bid * bdim + tid; // 全局 id // 将数据拷贝到共享内存 if (i < len) { sdata[tid] = arr[i]; } __syncthreads(); // 等待所有线程完成 // 每个线程计算 log2(bdim)-1 个轮回 // 比如 bdim = 8, 则每个线程计算 2 个轮回 for (int s = 1; s < bdim; s *= 2) { if (tid % (2 * s) == 0 && i + s < len) { sdata[tid] += sdata[tid + s]; } // 等待所有线程完成 后再进行下一轮计算 __syncthreads(); } // 每个 block 的第一个线程将结果写入到 out 中 if (tid == 0) { out[bid] = sdata[0]; }}
我们可以对照图来看一下代码。首先,我们把数组拷贝到共享内存中,然后等待所有线程完成。接着,每个线程计算 log2(bdim)-1 个轮回,每个轮回中,线程 tid 会把 tid + s 的值累加到 tid 中。最后,每个 block 的第一个线程会把自己的值写入到 out 中。