CUDA编程:Parallel Reduction 并行归约与 Warp Shuffle 优化

0 阅读10分钟

@[toc]

前面几课我们已经学习了:

第 6 课:Shared Memory 分块复用
第 7 课:Global Memory 合并访问
第 8 课:Shared Memory Bank Conflict
第 9 课:Occupancy 与 block size

第 10 课进入 CUDA 算子开发里非常核心的一类模式:

Reduction,并行归约。

典型任务包括:

数组求和 sum
最大值 max
最小值 min
均值 mean
方差 variance
softmax 中的 max/sum
layer norm 中的 mean/variance

这些都是深度学习算子、科学计算、统计计算中的基础操作。


一、本节课目标

本节课重点掌握:

1. 理解 reduction 为什么不能简单每个线程写一个结果
2. 理解 atomicAdd 为什么简单但慢
3. 掌握 shared memory block reduction
4. 掌握 warp shuffle reduction 的基本思想
5. 对比 atomic / shared / shuffle 三种实现的性能

二、核心原理

1. 什么是 Reduction?

Reduction 就是把很多数据归约成一个或少量结果。

例如数组求和:

input:  x0, x1, x2, x3, ..., xN-1

output: sum = x0 + x1 + x2 + ... + xN-1

CPU 上最直接写法是:

float sum = 0.0f;

for (int i = 0; i < n; ++i) {
    sum += data[i];
}

但是在 GPU 上,如果很多线程同时执行:

sum += data[i];

就会发生数据竞争。

因为多个线程会同时读写同一个 sum


2. 方法一:Atomic Add

最简单的 GPU 写法是:

atomicAdd(result, data[idx]);

这样可以保证多个线程不会同时破坏结果。

但是问题是:

所有线程都往同一个地址 atomicAdd
竞争非常严重
大量线程排队
性能很差

所以 atomic 方法适合做基线,不适合作为高性能方案。


3. 方法二:Shared Memory Block Reduction

更好的方式是:

每个 block 先在 shared memory 内部求一个局部和
每个 block 只写出一个 partial sum
最后再对 partial sum 继续归约

流程:

Global Memory 输入
        ↓
每个线程读取 1~2 个元素
        ↓
写入 shared memory
        ↓
block 内部并行树形归约
        ↓
每个 block 输出一个 partial sum
        ↓
继续归约 partial sum
        ↓
最终得到总和

这种方式大幅减少了全局原子竞争。


4. 方法三:Warp Shuffle Reduction

Shared Memory reduction 需要:

shared memory
__syncthreads()

但一个 warp 内的 32 个线程本来就是同步执行的。

CUDA 提供了 warp-level primitive,例如:

__shfl_down_sync(...)

它可以让同一个 warp 内线程直接交换寄存器数据。

优点是:

不需要 shared memory 做 warp 内归约
减少 __syncthreads()
减少 shared memory 访问
通常更快

典型 warp 内求和:

for (int offset = 16; offset > 0; offset >>= 1) {
    value += __shfl_down_sync(0xffffffff, value, offset);
}

可以理解为:

32 个线程先两两相加
再 16 个
再 8 个
再 4 个
再 2 个
最后得到一个 warp sum

就是不断的把后半部分的原地加到前半部分上,长度减半 在这里插入图片描述

在这里插入图片描述


三、本节实验设计

实现三种 GPU 求和方式:

1. atomicAdd 全局原子求和
2. shared memory block reduction
3. warp shuffle block reduction

比较指标:

kernel time
speedup vs atomic
结果误差

注意:本实验计时主要是 GPU kernel 计算时间,不包含完整 H2D/D2H 端到端时间。


四、完整可运行 CUDA C++ 代码

保存为:

lesson10_reduction.cu

代码如下:

#include <cuda_runtime.h>

#include <cmath>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <vector>

#define CUDA_CHECK(call)                                                        \
    do {                                                                        \
        cudaError_t err = call;                                                 \
        if (err != cudaSuccess) {                                               \
            std::cerr << "CUDA Error: " << cudaGetErrorString(err)              \
                      << " at " << __FILE__ << ":" << __LINE__ << std::endl;    \
            std::exit(EXIT_FAILURE);                                            \
        }                                                                       \
    } while (0)

/*
 * 方法 1:atomicAdd 版本。
 *
 * 每个线程读取一个元素,然后 atomicAdd 到同一个 result。
 * 写法最简单,但所有线程竞争同一个地址,性能通常很差。
 */
__global__ void reduce_atomic_kernel(const float* input,
                                     float* result,
                                     size_t n) {
    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;

    if (idx < n) {
        atomicAdd(result, input[idx]);
    }
}

/*
 * 方法 2:shared memory block reduction。
 *
 * 每个 block 输出一个 partial sum。
 * 每个线程最多读取两个元素,提高每个线程的工作量,减少 block 数量。
 */
__global__ void reduce_shared_kernel(const float* input,
                                     float* partial,
                                     size_t n) {
    extern __shared__ float sdata[];

    unsigned int tid = threadIdx.x;

    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x * 2 + threadIdx.x;

    float sum = 0.0f;

    if (idx < n) {
        sum += input[idx];
    }

    if (idx + blockDim.x < n) {
        sum += input[idx + blockDim.x];
    }

    sdata[tid] = sum;

    __syncthreads();

    /*
     * 树形归约:
     *
     * blockDim.x = 256 时:
     * 256 -> 128 -> 64 -> 32 -> 16 -> 8 -> 4 -> 2 -> 1
     */
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }

        __syncthreads();
    }

    if (tid == 0) {
        partial[blockIdx.x] = sdata[0];
    }
}

/*
 * warp 内归约。
 *
 * __shfl_down_sync 可以让同一个 warp 内的线程直接交换寄存器中的值。
 */
__inline__ __device__ float warp_reduce_sum(float value) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffff, value, offset);
    }

    return value;
}

/*
 * 方法 3:warp shuffle block reduction。
 *
 * 每个线程先读取数据到寄存器。
 * 每个 warp 内用 shuffle 做归约。
 * 每个 warp 的结果写入 shared memory。
 * 第一个 warp 再把各 warp 的结果归约成 block sum。
 */
__global__ void reduce_shuffle_kernel(const float* input,
                                      float* partial,
                                      size_t n) {
    unsigned int tid = threadIdx.x;

    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x * 2 + threadIdx.x;

    float sum = 0.0f;

    if (idx < n) {
        sum += input[idx];
    }

    if (idx + blockDim.x < n) {
        sum += input[idx + blockDim.x];
    }

    /*
     * 第一步:每个 warp 内部求和。
     */
    sum = warp_reduce_sum(sum);

    /*
     * 一个 block 最多 1024 个线程,也就是最多 32 个 warp。
     */
    __shared__ float warp_sums[32];

    int lane = tid % warpSize;
    int warp_id = tid / warpSize;

    /*
     * 每个 warp 的 lane 0 保存该 warp 的 sum。
     */
    if (lane == 0) {
        warp_sums[warp_id] = sum;
    }

    __syncthreads();

    /*
     * 第二步:第一个 warp 负责把所有 warp 的 sum 再归约一次。
     */
    int num_warps = (blockDim.x + warpSize - 1) / warpSize;

    float block_sum = 0.0f;

    if (tid < num_warps) {
        block_sum = warp_sums[tid];
    }

    if (warp_id == 0) {
        block_sum = warp_reduce_sum(block_sum);

        if (tid == 0) {
            partial[blockIdx.x] = block_sum;
        }
    }
}

/*
 * CPU reference。
 */
double reduce_cpu(const std::vector<float>& input) {
    double sum = 0.0;

    for (float v : input) {
        sum += static_cast<double>(v);
    }

    return sum;
}

/*
 * 计时 atomic kernel。
 *
 * 注意:
 * cudaMemset 不计入 kernel 时间。
 */
float time_atomic(const float* d_input,
                  float* d_result,
                  size_t n,
                  int block_size,
                  int repeat) {
    int grid_size = static_cast<int>((n + block_size - 1) / block_size);

    /*
     * warmup
     */
    CUDA_CHECK(cudaMemset(d_result, 0, sizeof(float)));
    reduce_atomic_kernel<<<grid_size, block_size>>>(d_input, d_result, n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());

    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    float total_ms = 0.0f;

    for (int r = 0; r < repeat; ++r) {
        CUDA_CHECK(cudaMemset(d_result, 0, sizeof(float)));

        CUDA_CHECK(cudaEventRecord(start));

        reduce_atomic_kernel<<<grid_size, block_size>>>(d_input, d_result, n);
        CUDA_CHECK(cudaGetLastError());

        CUDA_CHECK(cudaEventRecord(stop));
        CUDA_CHECK(cudaEventSynchronize(stop));

        float ms = 0.0f;
        CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
        total_ms += ms;
    }

    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));

    return total_ms / repeat;
}

/*
 * 多轮 reduction。
 *
 * reduce_kernel 每一轮把 input 归约成 partial。
 * 如果 partial 数量仍然大于 1,就继续归约。
 */
template <typename KernelFunc>
float time_multi_pass_reduction(KernelFunc kernel,
                                const float* d_input,
                                float* d_buf1,
                                float* d_buf2,
                                size_t n,
                                int block_size,
                                int repeat) {
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    float total_ms = 0.0f;

    for (int r = 0; r < repeat; ++r) {
        const float* current_input = d_input;
        float* current_output = d_buf1;
        float* next_output = d_buf2;

        size_t current_n = n;

        CUDA_CHECK(cudaEventRecord(start));

        while (current_n > 1) {
            int grid_size = static_cast<int>(
                (current_n + block_size * 2 - 1) / (block_size * 2)
            );

            size_t shared_bytes = static_cast<size_t>(block_size) * sizeof(float);

            kernel<<<grid_size, block_size, shared_bytes>>>(
                current_input,
                current_output,
                current_n
            );
            CUDA_CHECK(cudaGetLastError());

            current_n = static_cast<size_t>(grid_size);

            /*
             * 下一轮输入变成这一轮输出。
             */
            current_input = current_output;

            /*
             * 在两个 buffer 之间来回切换。
             */
            float* tmp = current_output;
            current_output = next_output;
            next_output = tmp;
        }

        CUDA_CHECK(cudaEventRecord(stop));
        CUDA_CHECK(cudaEventSynchronize(stop));

        float ms = 0.0f;
        CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
        total_ms += ms;
    }

    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));

    return total_ms / repeat;
}

/*
 * 执行一次 multi-pass reduction,并把最终结果拷回 CPU。
 */
template <typename KernelFunc>
float run_multi_pass_reduction_once(KernelFunc kernel,
                                    const float* d_input,
                                    float* d_buf1,
                                    float* d_buf2,
                                    size_t n,
                                    int block_size) {
    const float* current_input = d_input;
    float* current_output = d_buf1;
    float* next_output = d_buf2;

    size_t current_n = n;

    while (current_n > 1) {
        int grid_size = static_cast<int>(
            (current_n + block_size * 2 - 1) / (block_size * 2)
        );

        size_t shared_bytes = static_cast<size_t>(block_size) * sizeof(float);

        kernel<<<grid_size, block_size, shared_bytes>>>(
            current_input,
            current_output,
            current_n
        );
        CUDA_CHECK(cudaGetLastError());
        CUDA_CHECK(cudaDeviceSynchronize());

        current_n = static_cast<size_t>(grid_size);

        current_input = current_output;

        float* tmp = current_output;
        current_output = next_output;
        next_output = tmp;
    }

    float result = 0.0f;
    CUDA_CHECK(cudaMemcpy(&result,
                          current_input,
                          sizeof(float),
                          cudaMemcpyDeviceToHost));

    return result;
}

int main(int argc, char** argv) {
    size_t n = 1ULL << 24;   // 16,777,216 floats,约 64 MB
    int block_size = 256;
    int repeat = 5;

    if (argc >= 2) {
        n = static_cast<size_t>(std::atoll(argv[1]));
    }
    if (argc >= 3) {
        block_size = std::atoi(argv[2]);
    }
    if (argc >= 4) {
        repeat = std::atoi(argv[3]);
    }

    if (block_size <= 0 || block_size > 1024) {
        std::cerr << "block_size must be in (0, 1024].\n";
        return 1;
    }

    size_t bytes = n * sizeof(float);

    std::cout << "CUDA Lesson 10: Parallel Reduction\n";
    std::cout << "Elements   : " << n << "\n";
    std::cout << "Data size  : " << bytes / 1024.0 / 1024.0 << " MB\n";
    std::cout << "Block size : " << block_size << "\n";
    std::cout << "Repeat     : " << repeat << "\n";

    /*
     * 为了避免浮点累加误差过大,这里全部初始化为 1.0f。
     *
     * n 默认是 2^24,float 可以较好表示这个结果。
     */
    std::vector<float> h_input(n, 1.0f);

    double cpu_ref = reduce_cpu(h_input);

    float* d_input = nullptr;
    float* d_atomic_result = nullptr;
    float* d_buf1 = nullptr;
    float* d_buf2 = nullptr;

    CUDA_CHECK(cudaMalloc(&d_input, bytes));
    CUDA_CHECK(cudaMemcpy(d_input, h_input.data(), bytes, cudaMemcpyHostToDevice));

    CUDA_CHECK(cudaMalloc(&d_atomic_result, sizeof(float)));

    /*
     * partial buffer 的最大规模是第一轮 block 数。
     */
    size_t max_partials = (n + block_size * 2 - 1) / (block_size * 2);

    CUDA_CHECK(cudaMalloc(&d_buf1, max_partials * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_buf2, max_partials * sizeof(float)));

    /*
     * 计时 atomic。
     */
    float atomic_ms = time_atomic(d_input, d_atomic_result, n, block_size, repeat);

    CUDA_CHECK(cudaMemset(d_atomic_result, 0, sizeof(float)));
    reduce_atomic_kernel<<<static_cast<int>((n + block_size - 1) / block_size),
                           block_size>>>(d_input, d_atomic_result, n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());

    float atomic_result = 0.0f;
    CUDA_CHECK(cudaMemcpy(&atomic_result,
                          d_atomic_result,
                          sizeof(float),
                          cudaMemcpyDeviceToHost));

    /*
     * 计时 shared memory reduction。
     */
    float shared_ms = time_multi_pass_reduction(
        reduce_shared_kernel,
        d_input,
        d_buf1,
        d_buf2,
        n,
        block_size,
        repeat
    );

    float shared_result = run_multi_pass_reduction_once(
        reduce_shared_kernel,
        d_input,
        d_buf1,
        d_buf2,
        n,
        block_size
    );

    /*
     * 计时 warp shuffle reduction。
     */
    float shuffle_ms = time_multi_pass_reduction(
        reduce_shuffle_kernel,
        d_input,
        d_buf1,
        d_buf2,
        n,
        block_size,
        repeat
    );

    float shuffle_result = run_multi_pass_reduction_once(
        reduce_shuffle_kernel,
        d_input,
        d_buf1,
        d_buf2,
        n,
        block_size
    );

    std::cout << std::fixed << std::setprecision(4);

    std::cout << "\n[Reference]\n";
    std::cout << "CPU reference       : " << cpu_ref << "\n";

    std::cout << "\n[Timing: kernel only]\n";
    std::cout << "Atomic time         : " << atomic_ms << " ms\n";
    std::cout << "Shared reduce time  : " << shared_ms << " ms\n";
    std::cout << "Shuffle reduce time : " << shuffle_ms << " ms\n";

    std::cout << "\n[Speedup]\n";
    std::cout << "Shared vs Atomic    : " << atomic_ms / shared_ms << "x\n";
    std::cout << "Shuffle vs Atomic   : " << atomic_ms / shuffle_ms << "x\n";
    std::cout << "Shuffle vs Shared   : " << shared_ms / shuffle_ms << "x\n";

    std::cout << "\n[Result]\n";
    std::cout << "Atomic result       : " << atomic_result << "\n";
    std::cout << "Shared result       : " << shared_result << "\n";
    std::cout << "Shuffle result      : " << shuffle_result << "\n";

    std::cout << "\n[Abs error]\n";
    std::cout << "Atomic abs error    : " << std::fabs(atomic_result - cpu_ref) << "\n";
    std::cout << "Shared abs error    : " << std::fabs(shared_result - cpu_ref) << "\n";
    std::cout << "Shuffle abs error   : " << std::fabs(shuffle_result - cpu_ref) << "\n";

    CUDA_CHECK(cudaFree(d_input));
    CUDA_CHECK(cudaFree(d_atomic_result));
    CUDA_CHECK(cudaFree(d_buf1));
    CUDA_CHECK(cudaFree(d_buf2));

    return 0;
}

五、编译与运行

Tesla T4:

nvcc -O3 -arch=sm_75 lesson10_reduction.cu -o lesson10_reduction

运行实验:

./lesson10_reduction

也可以指定参数:

./lesson10_reduction 16777216 256 5

参数含义:

第 1 个参数:元素数量 n
第 2 个参数:block_size
第 3 个参数:repeat

实验结果:

方法Kernel time (ms)相对说明
Atomic Reduction133.7433最慢,全局 atomic 竞争严重
Shared Memory Reduction4.2167明显快于 Atomic
Warp Shuffle Reduction3.8721本次最快
对比项Speedup
Shared vs Atomic31.7174x
Shuffle vs Atomic34.5399x
Shuffle vs Shared1.0890x
Atomic 最慢
Shared reduction 明显更快
Warp shuffle reduction 通常最快或接近最快

原因:

atomicAdd:所有线程竞争一个全局地址
shared:每个 block 内部先局部归约,减少全局写入
shuffle:warp 内直接用寄存器交换,减少 shared memory 和同步开销

六、分析实验结果

1. 为什么 atomic 慢?

因为所有线程都执行:

atomicAdd(result, input[idx]);

这意味着所有线程都抢同一个地址。

虽然 atomic 保证正确性,但它会产生严重串行化。


2. 为什么 shared reduction 快?

因为它把问题拆成了两层:

block 内部:shared memory 快速归约
block 之间:只输出少量 partial sum

如果输入有 16M 个元素,atomic 是 16M 次全局原子加。

Shared reduction 第一轮只输出:

n / (block_size * 2)

个 partial sum。

block_size=256 时:

16,777,216 / 512 = 32,768 个 partial sum

全局写入数量大幅减少。


3. 为什么 shuffle 可能更快?

Shared reduction 在每一层都需要:

写 shared memory
读 shared memory
__syncthreads()

而 warp shuffle 在 warp 内部直接交换寄存器数据:

__shfl_down_sync(...)

它避免了很多 shared memory 访问。

所以它通常更快。

七、总结

核心结论:

1. Reduction 是 CUDA 算子开发中的基础模式
2. atomicAdd 写法简单,但全局竞争严重,通常很慢
3. Shared Memory Reduction 通过 block 内局部归约减少全局写入
4. Warp Shuffle Reduction 用寄存器交换减少 shared memory 和同步开销
5. 高性能 reduction 通常采用多级归约结构
6. 归约类算子要注意浮点误差和加法顺序

一句话总结:

Reduction 的高性能实现,本质是把“所有线程抢一个结果”改成“先分层局部求和,再逐级合并”。