1. 什么是 Reduce

Reduce 是一个高阶函数,它接收一个函数作为参数,这个函数接收两个参数,然后返回一个值。Reduce 会从左到右依次对数组中的元素进行处理,最终得到一个值。以加法为例,就相当于给定一个数组,对数组求和。用 cpu 实现起来非常简单,一层循环解决问题。本文主要讨论如何用纯函数式的方式实现 Reduce

2. 使用 GPU 实现 Reduce

2.1 CPU 实现

// 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

2.2 GPU 层次结构

首先我们先来回归一下 GPU 的层次结构。也就是代码中的 blockgridblock 是一个线程块,grid 是一个线程网格。block 中的线程可以通过 threadIdx.x 来获取自己的线程 id,grid 中的线程可以通过 blockIdx.x 来获取自己的线程 id。block 中的线程可以通过 blockDim.x 来获取 block 的大小,grid 中的线程可以通过 gridDim.x 来获取 grid 的大小。

对于 Reduce 来说,我们可以按照下面这个图设计:

img1

img1

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

GPU 的计算过程如下图所示:

img2

img2

以上图为例,我们来看一下 Reduce 的计算过程。首先,我们把数组分成了 3block,每个 block 中有 8 个线程。在第一轮计算中,奇数线程会把自己的值累加到偶数线程中。在第二轮计算中,block 中的第 0 个线程会把 4 号线程的值累加到自己的值中。 每个 block 的值都计算完之后还需要对 block 的值进行累加,下面我们来看一下代码要如何实现。

2.3 GPU 实现

我们首先看 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 中。