CUDA算子优化(1)：Reduce

Reduce（规约）是GPU编程中最基础，也是最能体现并行思维的算子之一。本文将一步一步带你优化Reduce sum算子，让你理解怎么写出更快的kernel。

1. Reduce算子基础

Reduce是一类"多输入->单输出"的操作。比如下列操作：

• Sum Reduce：$\sum_{i=0}^{N-1} x_{i}$ （输入为长度N的数组，输出为所有元素的和）
• Max/Min Reduce：$\max \left(x_{0}, x_{1}, \ldots, x_{N-1}\right)$ （输入是长度为N的数组，输出为数组的最大/最小值）
• Dot Product：$\sum_{i=0}^{N-1} a_i \cdot b_i$ （输入是长度都为N的数组a和b，输出为它们的点积值）

1.1 性能瓶颈分析
GPU Kernel的性能上限由计算量和访存量共同决定的。对于sum Reduce来说，每个元素读取一次（1次load）就做一次加法（1次FLOP），算术强度为$\frac{1 \text{FLOP}}{4 \text{Byte}} = 0.25 \text{FLOP/Byte}$，说明Reduce是典型的访存密集型操作，优化的核心在于提高内存带宽的利用率，而不是减少计算次数。

2. 版本V0朴素并行规约

下图所示，以N=8个元素为例，我们仅需要N/2=4个线程即可；首先，第一轮，每个线程计算两个相邻元素的和，其中步长stride=1。

经过上述一轮后，就剩下data[0]、data[2]、data[4]和data[6]，求这几个元素的和即可，第二轮和上一轮一样，线程0和线程2分别计算data[0]+data[2],data[4]+data[6]，其中步长stride=2。

第三轮让线程0计算data[0]+data[4]即可，其中步长stride=4。这样代码就很好写了。

#include <stdio.h>
#include <stdlib.h>


__global__ void kernel(float *d_out, float *d_data, int N) {
    int tid = threadIdx.x;
    int i = 2 * tid;

    for (int stride = 1; stride <= N / 2; stride *= 2) {
        if (i % (2 * stride) == 0) {
            d_data[i] += d_data[i + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        d_out[0] = d_data[0];
    }
}

float sum_cpu(float *h_data, int N) {
    float sum = 0.0f;
    for (int i = 0; i < N; ++i) {
        sum += h_data[i];
    }
    return sum;
}

int main() {
    int N = 2048;
    size_t bytes = N * sizeof(float);
    float *h_data = (float *)malloc(bytes);
    float *h_out = (float *)malloc(sizeof(float));
    for (int i = 0; i < N; ++i) {
        h_data[i] = (float)rand() / RAND_MAX;
    }

    float sum = sum_cpu(h_data, N);

    float *d_data, *d_out;
    cudaMalloc(&d_data, bytes);
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    kernel<<<1, 1024>>>(d_out, d_data, N);
    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    printf("sum:%f, h_out:%f", sum, h_out[0]);

    free(h_data);
    free(h_out);
    cudaFree(d_data);
    cudaFree(d_out);

    return 0;
}

这个方法有 1 个问题，Warp Divergence。
Warp Divergence：GPU 最小的调度单位是Warp（32 个线程）。这 32 个线程都执行同一条指令，如果部分线程满足if 条件，部分线程不满足，就会发生分支分歧，GPU 就会分 2 次执行，先执行满足if 条件的指令，再执行else 条件的指令，这样就会导致性能下降。

3. 版本V1消除Warp

改进思路：改变步长的方向，让步长从大到小，而不是从小到大。整体思路如下图，第一轮，以N=8个元素为例，让data[0]+data[4]，data[1]+data[5]，data[2]+data[6]，data[3]+data[7]，步长为stride = 4.

第二轮，让data[0]+data[2]，data[1]+data[3]，步长为stride=2。

第三轮，让data[0]+data[1]，步长为stride=1

#include <stdio.h>
#include <stdlib.h>


#define BLOCK_SIZE 1024

__global__ void reduce_kernel(float *d_out, float *d_data, int N) {
    int tid = threadIdx.x;

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

    if (tid == 0) {
        d_out[0] = d_data[0];
    }
}

float sum_cpu(float *data, int N) {
    float sum = 0.0;
    for (int i = 0; i < N; ++i) {
        sum += data[i];
    }
    return sum;
}

int main() {
    int N = 2048;
    size_t bytes = N * sizeof(float);
    float *h_data = (float *)malloc(bytes);
    float *h_out = (float *)malloc(sizeof(float));
    for (int i = 0; i < N; ++i) {
        h_data[i] = (float) rand() / RAND_MAX;
    }
    float sum = sum_cpu(h_data, N);

    float *d_data, *d_out;
    cudaMalloc(&d_data, bytes);
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    reduce_kernel<<<1, BLOCK_SIZE>>>(d_out, d_data, N);
    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    printf("sum = %f, h_out = %f\n", sum, h_out[0]);

    free(h_data);
    free(h_out);
    cudaFree(d_data);
    cudaFree(d_out);

    return 0;
}

这样，以N=256为例，我们对比下第一种方法，这种方法效率是如何提升的。
第一轮：

	活跃线程	Warp Divergence
版本V0	每个Warp 只有偶数线程工作	50% 利用率
版本V1	tid 0-127 活跃，恰好 4 个完整的Warp	无分化

第二轮

	活跃线程	Warp Divergence
版本V0	每 4 个线程只有一个工作	25% 利用率
版本V1	tid 0-63 活跃，恰好 2 个完整的Warp	无分化

第三轮

	活跃线程	Warp Divergence
版本V0	每 8 个线程只有一个工作	12.5% 利用率
版本V1	tid 0-31 活跃，恰好 1 个完整的Warp	无分化

依次类推，版本V0越来越差，版本V1相比于V0有较大的提升

3. 版本V2共享内存

上述两种方法均是对全局内存的访问，共享内存的读写速度仅次于寄存器，对整个线程块可见。它的主要作用是减少对全局内存的访问，我们可以利用共享内存加快核函数的运行时间。

共享内存一般使用_shared_关键字，先将全局内存中的数据搬运到共享内存中，然后求和操作在共享内存中进行。

#include <stdio.h>
#include <stdlib.h>


#define BLOCK_SIZE 1024

__global__ void reduce_kernel(float *d_out, float *d_data, int N) {
    int tid = threadIdx.x;

    __shared__ float smem[BLOCK_SIZE];

    smem[tid] = d_data[tid] + d_data[tid + BLOCK_SIZE];
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride >= 1; stride /= 2) {
        if (tid < stride) {
            smem[tid] += smem[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        d_out[0] = smem[0];
    }
}

float sum_cpu(float *data, int N) {
    float sum = 0.0;
    for (int i = 0; i < N; ++i) {
        sum += data[i];
    }
    return sum;
}

int main() {
    int N = 2048;
    size_t bytes = N * sizeof(float);
    float *h_data = (float *)malloc(bytes);
    float *h_out = (float *)malloc(sizeof(float));
    for (int i = 0; i < N; ++i) {
        h_data[i] = (float) rand() / RAND_MAX;
    }
    float sum = sum_cpu(h_data, N);

    float *d_data, *d_out;
    cudaMalloc(&d_data, bytes);
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    reduce_kernel<<<1, BLOCK_SIZE>>>(d_out, d_data, N);
    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    printf("sum = %f, h_out = %f\n", sum, h_out[0]);

    free(h_data);
    free(h_out);
    cudaFree(d_data);
    cudaFree(d_out);

    return 0;
}

4. 版本V3多block情况

当数据量比较大的时候，一个block显然是不够用的，每种GPU的block都有最大的上限，这个时候就需要使用多block。这里很明显的思考是，为每个block创建一个共享内存，分别负责某一块的数据的求和，最后将得到的所有求和值，再求一次和。

#include <stdio.h>
#include <stdlib.h>


#define BLOCK_SIZE 1024

__global__ void reduce_kernel(float *d_out, float *d_data, int N) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    int i = threadIdx.x;

    __shared__ float smem[BLOCK_SIZE];
    if (tid < N) {
        smem[i] = d_data[tid];
    } else {
        smem[i] = 0.0f;
    }
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride >= 1; stride /= 2) {
        if (i < stride && tid < N) {
            smem[i] += smem[i + stride];
        }
        __syncthreads();
    }

    if (i == 0) {
        atomicAdd(&d_out[0], smem[0]);
    }
}

float sum_cpu(float *data, int N) {
    float sum = 0.0;
    for (int i = 0; i < N; ++i) {
        sum += data[i];
    }
    return sum;
}

int main() {
    int N = 18299;
    size_t bytes = N * sizeof(float);
    float *h_data = (float *)malloc(bytes);
    float *h_out = (float *)malloc(sizeof(float));
    for (int i = 0; i < N; ++i) {
        h_data[i] = (float) rand() / RAND_MAX;
    }
    float sum = sum_cpu(h_data, N);

    float *d_data, *d_out;
    cudaMalloc(&d_data, bytes);
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    dim3 grid_size = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    reduce_kernel<<<grid_size, BLOCK_SIZE>>>(d_out, d_data, N);

    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    printf("sum = %f, h_out = %f\n", sum, h_out[0]);

    free(h_data);
    free(h_out);
    cudaFree(d_data);
    cudaFree(d_out);

    return 0;
}

5. 版本V4线程粗化

回顾之前的方法，每一个线程只负责处理一个元素，以规约的第一轮为例（blockDim=256），只有前128个线程在工作，后128个线程彻底闲置了（他们仅仅只是把数据从全局内存搬运到了共享内存）。
线程粗化的意思是让一个线程处理多个元素，如果每个线程在加载阶段就负责 2 个元素并预先求和，那么用同样 256 个线程可以处理 512 个数据，block的数量却可以减半，grid 更小，调度开销降低。

#include <stdio.h>
#include <stdlib.h>


#define BLOCK_SIZE 512
#define COARSEN 4

__global__ void kernel(float *d_out, float *d_data, int N) {
    int tid = threadIdx.x;
    int block_idx = blockDim.x * blockIdx.x * COARSEN;

    __shared__ float smem[BLOCK_SIZE];
    float local_sum = 0.0f;
    for (int i = 0; i < COARSEN; ++i) {
        int idx = block_idx + tid * COARSEN + i;
        if (idx < N) {
            local_sum += d_data[idx];
        }
    }
    smem[tid] = local_sum;
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride >= 1; stride /= 2) {
        if (tid < stride) {
            smem[tid] += smem[tid +stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        atomicAdd(&d_out[0], smem[0]);
    }
}

float sum_cpu(float *h_data, int N) {
    float sum = 0.0f;
    for (int i = 0; i < N; ++i) {
        sum += h_data[i];
    }
    return sum;
}

void init_float_rand(float *arr, int n, float low, float high) {
  for (int i = 0; i < n; i++) {
    float r = (float)rand() / (float)RAND_MAX;
    arr[i] = low + (high - low) * r;
  }
}

int main() {
    int N = 123341;
    size_t bytes = N * sizeof(float);
    float *h_data = (float *)malloc(bytes);
    float *h_out = (float *)malloc(sizeof(float));
    init_float_rand(h_data, N, -1.0, 1.0);

    float sum = sum_cpu(h_data, N);
    float *d_data, *d_out;
    cudaMalloc(&d_data, bytes);
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    dim3 grid = (N + BLOCK_SIZE * COARSEN - 1) / (BLOCK_SIZE * COARSEN);
    kernel<<<grid, BLOCK_SIZE>>>(d_out, d_data, N);
    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    printf("sum = %f, h_out = %f\n", sum, h_out[0]);

    free(h_data);
    free(h_out);
    cudaFree(d_data);
    cudaFree(d_out);
    return 0;
}