CUDA 编程：Parallel Prefix Sum / Scan 并行前缀和

Michael阿明

发布于 2026-07-06 14:08:40

170

文章目录

1. 本课目标
2. 什么是 Prefix Sum？
- 2.1 Inclusive Scan
- 2.2 Exclusive Scan
3. 为什么 Prefix Sum 难以直接并行？
4. Blelloch Scan 核心思想
5. 实验内容
6. CUDA C++ 代码
7. 编译与运行
8. 输出

第 10 课我们学习了Reduction 归约：把很多数变成一个数。

另一个非常核心的并行算法：

❝Prefix Sum / Scan：把一个数组变成它的前缀累积结果。

它是很多 GPU 算子的基础，例如：

stream compaction、 radix sort 、histogram、 稀疏矩阵 CSR、 token offset 计算、 batch packing、 并行过滤、 并行分配输出位置

1. 本课目标

本节课重点掌握：

1. 什么是 inclusive scan 和 exclusive scan

2. 为什么 scan 不能简单并行

3. 掌握 Blelloch Scan 的 up-sweep / down-sweep 思想

4. 用 shared memory 实现 block-level scan

5. 用多轮 block scan 实现大数组 scan

6. 对比 CPU scan 和 GPU scan 的性能

2. 什么是 Prefix Sum？

假设输入数组是：

input = [3, 1, 7, 0, 4, 1, 6, 3]

2.1 Inclusive Scan

包含自己：

inclusive = [3, 4, 11, 11, 15, 16, 22, 25]

也就是：

out[i] = input[0] + input[1] + ... + input[i]

2.2 Exclusive Scan

不包含自己：

exclusive = [0, 3, 4, 11, 11, 15, 16, 22]

也就是：

out[i] = input[0] + input[1] + ... + input[i-1]

CUDA 中很多场景更常用exclusive scan。

例如你有一个 flag 数组：

flag = [1, 0, 1, 1, 0, 1]

exclusive scan 后：

pos = [0, 1, 1, 2, 3, 3]

它可以告诉每个有效元素应该写到输出数组的哪个位置。

3. 为什么 Prefix Sum 难以直接并行？

CPU 串行写法很简单：

out[0] = 0;

for (int i = 1; i < n; ++i) {
    out[i] = out[i - 1] + input[i - 1];
}

问题是：

out[i] 依赖 out[i-1] out[i-1] 又依赖 out[i-2]

所以它天然是链式依赖，不能像 vector add 那样每个线程独立计算一个元素。

GPU 上要并行做 scan，需要改变计算结构。

4. Blelloch Scan 核心思想

Blelloch Scan 分两步：

1. Up-sweep：向上规约，构造部分和树

2. Down-sweep：向下传播，生成 exclusive scan

以 8 个元素为例：

输入： [3, 1, 7, 0, 4, 1, 6, 3]

Up-sweep： 先两两求和 再四个一组求和 最后得到总和

Down-sweep： 把根节点置 0 再把前缀和向下分发 最终得到 exclusive scan

最终：

exclusive scan = [0, 3, 4, 11, 11, 15, 16, 22]

它的优势是：

串行复杂度：O(n) 并行 scan：O(log n) 层级同步

在一个 block 内，可以用 shared memory 高效实现。

在这里插入图片描述

5. 实验内容

我们实现：

1. CPU exclusive scan

2. GPU block-level Blelloch scan

3. 多轮 GPU scan 支持大数组

4. 结果校验

5. CPU vs GPU 时间对比

注意：

❝本实验主要测 GPU scan kernel 路径时间，不把 H2D / D2H 作为核心对比对象。

在这里插入图片描述

6. CUDA C++ 代码

代码如下：

#include <cuda_runtime.h>

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

#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)

/*
 * CPU exclusive scan:
 *
 * input : [x0, x1, x2, x3, ...]
 * output: [0,  x0, x0+x1, x0+x1+x2, ...]
 */
void exclusive_scan_cpu(const std::vector<float>& input,
                        std::vector<float>& output) {
    if (input.empty()) return;

    output[0] = 0.0f;

    for (size_t i = 1; i < input.size(); ++i) {
        output[i] = output[i - 1] + input[i - 1];
    }
}

/*
 * block-level exclusive scan.
 *
 * 每个 block 处理 2 * blockDim.x 个元素。
 *
 * 例如 blockDim.x = 256：
 * 一个 block 处理 512 个 float。
 *
 * block_sums[blockIdx.x] 保存当前 block 的总和，
 * 后续要对 block_sums 再做 scan，得到每个 block 的 offset。
 */
__global__ void block_exclusive_scan_kernel(const float* input,
                                            float* output,
                                            float* block_sums,
                                            size_t n) {
    extern __shared__ float temp[];

    unsignedint tid = threadIdx.x;

    size_t block_start = static_cast<size_t>(blockIdx.x) * blockDim.x * 2;
    size_t idx1 = block_start + tid;
    size_t idx2 = block_start + tid + blockDim.x;

    /*
     * 1. 从 global memory 读入 shared memory。
     *
     * 每个线程最多读两个元素。
     * 越界位置补 0。
     */
    temp[tid] = (idx1 < n) ? input[idx1] : 0.0f;
    temp[tid + blockDim.x] = (idx2 < n) ? input[idx2] : 0.0f;

    __syncthreads();

    /*
     * 2. Up-sweep 阶段。
     *
     * 构造一棵求和树。
     *
     * 例如 512 个元素：
     * 512 -> 256 -> 128 -> ... -> 1
     */
    unsignedint offset = 1;
    unsignedint total_elems = blockDim.x * 2;

    for (unsignedint d = total_elems >> 1; d > 0; d >>= 1) {
        if (tid < d) {
            unsignedint ai = offset * (2 * tid + 1) - 1;
            unsignedint bi = offset * (2 * tid + 2) - 1;

            temp[bi] += temp[ai];
        }

        offset <<= 1;
        __syncthreads();
    }

    /*
     * 3. 保存 block 总和，并把根节点置 0。
     *
     * 置 0 是 exclusive scan 的关键。
     */
    if (tid == 0) {
        if (block_sums != nullptr) {
            block_sums[blockIdx.x] = temp[total_elems - 1];
        }

        temp[total_elems - 1] = 0.0f;
    }

    __syncthreads();

    /*
     * 4. Down-sweep 阶段。
     *
     * 把前缀和从树根向下传播。
     */
    for (unsignedint d = 1; d < total_elems; d <<= 1) {
        offset >>= 1;

        if (tid < d) {
            unsignedint ai = offset * (2 * tid + 1) - 1;
            unsignedint bi = offset * (2 * tid + 2) - 1;

            float t = temp[ai];
            temp[ai] = temp[bi];
            temp[bi] += t;
        }

        __syncthreads();
    }

    /*
     * 5. 写回 global memory。
     */
    if (idx1 < n) {
        output[idx1] = temp[tid];
    }

    if (idx2 < n) {
        output[idx2] = temp[tid + blockDim.x];
    }
}

/*
 * 给每个 block 的局部 scan 结果加上 block offset。
 *
 * block_offsets 是对 block_sums 做 exclusive scan 后得到的结果。
 */
__global__ void add_block_offsets_kernel(float* data,
                                         const float* block_offsets,
                                         size_t n,
                                         int elems_per_block) {
    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;

    if (idx < n) {
        size_t block_id = idx / elems_per_block;
        data[idx] += block_offsets[block_id];
    }
}

/*
 * 递归式 GPU exclusive scan。
 *
 * 思路：
 * 1. 每个 block 做局部 scan，得到 output 和 block_sums
 * 2. 对 block_sums 再做 exclusive scan，得到 block_offsets
 * 3. output 每个元素加上自己 block 的 offset
 */
void exclusive_scan_gpu_recursive(const float* d_input,
                                  float* d_output,
                                  size_t n,
                                  int block_size) {
    if (n == 0) return;

    int elems_per_block = block_size * 2;
    size_t num_blocks = (n + elems_per_block - 1) / elems_per_block;

    float* d_block_sums = nullptr;

    CUDA_CHECK(cudaMalloc(&d_block_sums, num_blocks * sizeof(float)));

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

    block_exclusive_scan_kernel<<<static_cast<int>(num_blocks),
                                  block_size,
                                  shared_bytes>>>(
        d_input,
        d_output,
        d_block_sums,
        n
    );
    CUDA_CHECK(cudaGetLastError());

    if (num_blocks > 1) {
        float* d_block_offsets = nullptr;
        CUDA_CHECK(cudaMalloc(&d_block_offsets, num_blocks * sizeof(float)));

        /*
         * 递归地对 block_sums 做 exclusive scan。
         */
        exclusive_scan_gpu_recursive(d_block_sums,
                                     d_block_offsets,
                                     num_blocks,
                                     block_size);

        /*
         * 把 block offset 加回每个元素。
         */
        int add_threads = 256;
        int add_grid = static_cast<int>((n + add_threads - 1) / add_threads);

        add_block_offsets_kernel<<<add_grid, add_threads>>>(
            d_output,
            d_block_offsets,
            n,
            elems_per_block
        );
        CUDA_CHECK(cudaGetLastError());

        CUDA_CHECK(cudaFree(d_block_offsets));
    }

    CUDA_CHECK(cudaFree(d_block_sums));
}

/*
 * 计时 GPU scan。
 *
 * 注意：
 * 这里用 cudaEvent 测 GPU kernel 路径时间。
 * 临时 cudaMalloc/cudaFree 不应作为核心性能结论。
 * 工程中应该预分配临时 buffer。
 */
float time_gpu_scan(const float* d_input,
                    float* d_output,
                    size_t n,
                    int block_size,
                    int repeat) {
    /*
     * warmup
     */
    exclusive_scan_gpu_recursive(d_input, d_output, n, block_size);
    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(cudaEventRecord(start));

        exclusive_scan_gpu_recursive(d_input, d_output, n, block_size);

        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;
}

float max_abs_diff(const std::vector<float>& a,
                   const std::vector<float>& b) {
    float max_diff = 0.0f;

    for (size_t i = 0; i < a.size(); ++i) {
        float diff = std::fabs(a[i] - b[i]);
        if (diff > max_diff) {
            max_diff = diff;
        }
    }

    return max_diff;
}

int main(int argc, char** argv) {
    size_t n = 1ULL << 22;   // 4,194,304 floats，约 16 MB
    int block_size = 256;
    int repeat = 10;

    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";
        return1;
    }

    /*
     * 为了让 float 结果容易校验，默认全部初始化为 1.0f。
     * 当 n <= 2^24 时，前缀和整数部分可以被 float 精确表示。
     */
    std::vector<float> h_input(n, 1.0f);
    std::vector<float> h_cpu(n, 0.0f);
    std::vector<float> h_gpu(n, 0.0f);

    size_t bytes = n * sizeof(float);

    std::cout << "CUDA Lesson 11: Parallel Prefix Sum / Scan\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";

    /*
     * CPU timing。
     */
    auto cpu_start = std::chrono::high_resolution_clock::now();
    exclusive_scan_cpu(h_input, h_cpu);
    auto cpu_end = std::chrono::high_resolution_clock::now();

    double cpu_ms = std::chrono::duration<double, std::milli>(
        cpu_end - cpu_start
    ).count();

    float* d_input = nullptr;
    float* d_output = nullptr;

    CUDA_CHECK(cudaMalloc(&d_input, bytes));
    CUDA_CHECK(cudaMalloc(&d_output, bytes));

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

    float gpu_ms = time_gpu_scan(d_input,
                                 d_output,
                                 n,
                                 block_size,
                                 repeat);

    CUDA_CHECK(cudaMemcpy(h_gpu.data(),
                          d_output,
                          bytes,
                          cudaMemcpyDeviceToHost));

    float diff = max_abs_diff(h_cpu, h_gpu);

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

    std::cout << "\n[Timing]\n";
    std::cout << "CPU exclusive scan time : " << cpu_ms << " ms\n";
    std::cout << "GPU exclusive scan time : " << gpu_ms << " ms\n";
    std::cout << "Speedup                 : " << cpu_ms / gpu_ms << "x\n";

    std::cout << "\n[Check]\n";
    std::cout << "Max abs diff            : " << diff << "\n";
    std::cout << "Result                  : " << (diff < 1e-3f ? "PASS" : "CHECK") << "\n";

    std::cout << "\n[Sample]\n";
    std::cout << "input[0..7] : ";
    for (int i = 0; i < 8 && i < static_cast<int>(n); ++i) {
        std::cout << h_input[i] << " ";
    }
    std::cout << "\n";

    std::cout << "gpu[0..7]   : ";
    for (int i = 0; i < 8 && i < static_cast<int>(n); ++i) {
        std::cout << h_gpu[i] << " ";
    }
    std::cout << "\n";

    CUDA_CHECK(cudaFree(d_input));
    CUDA_CHECK(cudaFree(d_output));

    return0;
}

7. 编译与运行

Tesla T4：

nvcc -O3 -arch=sm_75 lesson11_prefix_scan.cu -o lesson11_scan

运行默认实验：

./lesson11_scan

也可以指定参数：

./lesson11_scan 16777216 256 10

参数含义：

第 1 个参数：元素数量 n

第 2 个参数：block_size

第 3 个参数：repeat

8. 输出

CUDA Lesson 11: Parallel Prefix Sum / Scan

Elements : 16777216

Data size : 64 MB

Block size : 256

Repeat : 10

[Timing]

CPU exclusive scan time : 23.7002 ms

GPU exclusive scan time : 2.9987 ms

Speedup : 7.9035x

[Check] Max abs diff : 0.0000

Result :

PASS [Sample]

input[0..7] : 1.0000 1.0000 1.0000 1.0000 1.0000 1.0000 1.0000 1.0000

gpu[0..7] : 0.0000 1.0000 2.0000 3.0000 4.0000 5.0000 6.0000 7.0000

GPU 执行速度是 CPU 的 8 倍

总结：

❝Prefix Scan 的本质，是把“前一个结果依赖后一个结果”的串行链条，改造成可以在 shared memory 中并行传播的树形结构。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2026-07-04，如有侵权请联系 cloudcommunity@tencent.com 删除

cuda

本文分享自 Michael阿明微信公众号，前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

登录后参与评论

0 条评论

热度