AIGC标识 Parallel Reduction 优化

并行归约：把一条加法链折成归约树

归约，SIMD，OpenMP，CUDA，浮点误差。

把 \(n\) 个数加成一个数，最直接的写法是一条首尾相接的链：第 \(i\) 次加法要等第 \(i - 1\) 次的结果。链长是 \(n - 1\)，累加器的读写延迟全串在中间，乱序窗口再宽也藏不住。并行归约要动的就是这条链——把数据切开交给多个执行单元各自累加，再用一棵树把部分结果折起来。

折树改变不了算术强度

先算这个算子本身能吃多少算力。以 float 为例，每个元素读 \(4\) byte、贡献 \(1\) 次加法，也就是每字节 \(\frac{1}{4}\) 次浮点运算。这个比值放在现在的处理器上属于典型的 memory bound，想算得快只能让数据流得快，多出来的执行单元帮不上忙。

这件事对后面的取舍影响很大。折树降低的是依赖深度，换来的是并行度，但每个元素还是要被读一遍，带宽是硬上限。GPU 上卡在显存带宽、CPU 上卡在内存带宽，都是这一类问题。

每个线程一份私有累加器

数据切开之后，每个线程都需要一块自己的累加变量，最后再合并。OpenMP 的 reduction 子句把这件事写成了语法：它为每个线程生成一份私有副本，初值取算子的单位元（加法是 \(0\)，乘法是 \(1\)），线程内先各自归约，出并行区之后再用同一个算子合并回共享变量。要注意私有副本的初值不是循环前面写的那个用户初值，用户初值是在最后合并阶段才参与的[1] 。

手写的时候容易踩到伪共享。多个线程的私有槽如果落在同一条 cacheline 上，写的是不同变量，cacheline 却要在核之间来回搬：

struct alignas(64) PartialSum {
    double value;
    char padding[64 - sizeof(double)];   // 填满一条 cacheline，让相邻线程的私有槽互不干扰
};

\(64\) 是一般 x86 的 cacheline 大小。padding 不是浪费，是让每个槽独占一行。

内层循环还要再拆一次

线程私有解决了线程之间的竞争，单个线程内部却还是一条长度约为 \(\frac{n}{p}\) 的加法链。线程数不够大的时候这条链就是瓶颈，而且它是串行的：一次加法没回来，下一次加法发不出去。

办法是在线程内部再摆几个互不相关的累加器，把一条链拆成几条。拆到 \(4\) 条之后深度降到 \(\frac{1}{4}\)，编译器也才有机会把它们装进向量寄存器：

__m256d acc0 = _mm256_setzero_pd();
__m256d acc1 = _mm256_setzero_pd();
__m256d acc2 = _mm256_setzero_pd();
__m256d acc3 = _mm256_setzero_pd();

std::size_t i = begin;
const std::size_t step = 16;             // 每轮 16 个 double，四条链各吃 4 个
for (; i + step <= end; i += step) {
    acc0 = _mm256_add_pd(acc0, _mm256_loadu_pd(data + i));
    acc1 = _mm256_add_pd(acc1, _mm256_loadu_pd(data + i + 4));
    acc2 = _mm256_add_pd(acc2, _mm256_loadu_pd(data + i + 8));
    acc3 = _mm256_add_pd(acc3, _mm256_loadu_pd(data + i + 12));
}

这四条加法之间没有依赖，流水线可以一直填满。约束出在累加器条数和硬件参数的匹配上：条数太少，前一条加法没回来，下一条就发不出去；条数太多，向量寄存器又不够分。\(4\) 是个常见起点，具体几条合适还是得看一眼编译出来的汇编再定。

收尾再把四条链折到一起。水平归约（_mm256_hadd_pd 这一类）在不少微架构上要拆成若干条 shuffle 加跨 \(128\) 位通道的搬运，比一次纵向加法贵得多，所以只能出现一次，不能放进循环里，编译器那边也一直在想办法少生成几条[2] ：

// 四条链先两两合并，再落到标量，水平归约只做一次
acc0 = _mm256_add_pd(acc0, acc1);
acc2 = _mm256_add_pd(acc2, acc3);
acc0 = _mm256_add_pd(acc0, acc2);

double lanes[4];
_mm256_storeu_pd(lanes, acc0);
double sum = (lanes[0] + lanes[1]) + (lanes[2] + lanes[3]);

for (; i < end; ++i) {          // 剩下不到 16 个的尾巴
    sum += data[i];
}

块内树形归约

线程数上去之后，部分结果本身也有 \(p\) 个，合并阶段又变回一条链。GPU 上的做法把这个阶段也折成树：一个 block 的 \(b\) 个部分和放在 shared memory 里，步长从 \(\frac{b}{2}\) 开始逐次折半。

for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (threadIdx.x < s) {
        sdata[threadIdx.x] += sdata[threadIdx.x + s];
    }
    __syncthreads();
}

每一轮活跃的线程减半，\(b - 1\) 次加法摊在 \(\lceil \log_2 b \rceil\) 轮里，深度从 \(b - 1\) 降到 \(\lceil \log_2 b \rceil\)。__syncthreads() 不能省，它保证这一轮读到的 sdata[threadIdx.x + s] 已经是上一轮的最终值，否则会读到别人还没写完的中间结果。

共享内存是分 bank 的，同一个 warp 内的线程按 \(\frac{b}{2}, \frac{b}{4}, \ldots\) 这样的步长访问，前几轮会大量撞在同一个 bank 上，串行化之后等待时间就白花了。缓解方式有两种，一是往上加一个 offset 把访问错开，二是换成顺序寻址——先在全局内存里按连续的 threadIdx.x 把数据读进来（顺便拿到合并访问），再在 shared memory 里按 \(1, 2, 4\) 的步长向上归约。后一种更干净，因为它顺带解决了全局内存访问的合并问题。

warp 内不用共享内存

warp 里的 \(32\) 个线程本来就是同一批指令一起走，交换数据不必绕共享内存。CUDA 的 shuffle 让线程直接读同一个 warp 里另一个线程的寄存器：

__inline__ __device__ float warpReduceSum(float value) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffff, value, offset);
    }
    return value;
}

\(5\) 轮之后 lane \(0\) 拿到整个 warp 的和，路径上没有共享内存访问，也没有 __syncthreads()。一个 block 里还剩若干 warp 的结果，把它们写进 shared memory，再让第一个 warp 重复一遍上面的操作，一个 block 就收敛到一个值。到这里为止的分层顺序（全局内存合并读、warp 内 shuffle、warp 间共享内存、块间第二轮 kernel）在 CUDA 的 reduction 讲义里是一步步压下来的[3] ，每一步都对应一个具体的开销；同一份内容的课程笔记版见[4] 。

每个线程多领几个元素

再往前是让一个线程处理多个元素，也就是 thread coarsening 或者 grid-stride：

std::size_t idx = 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];   // 同一个线程再吃一个，两个元素先在寄存器里加
}

每线程多读一个元素，block 数就少一半，第二轮要归约的部分和也跟着少一半；元素在寄存器里先加一次，连 shared memory 都不用进。代价是并行度下降，块数少到填不满 SM 的时候反而变慢，尾块不均衡也会更明显。这个平衡点跟 block 大小、每线程元素数、数据规模都有关系，调参的意义大于理论。

多轮 kernel 启动是这套结构里另一块固定开销。第一轮之后剩下的是部分和，还要再来一轮甚至几轮才能收敛到一个标量，中间要在两个 buffer 之间来回切，免得在同一块内存上边读边写。

余数、重结合与不确定性

按块长切分的数据量很少正好是整数倍，余数分支得写清楚。常见做法是靠 idx < n 这类卫语句在边界外直接跳过，而不是单独开一个只处理尾巴的循环，那样等于多一遍循环入口和边界判断。

顺序问题更要紧。浮点加法不满足结合律，\((a + b) + c\) 和 \(a + (b + c)\) 的结果可能差一个舍入。编译器默认不允许重排浮点表达式，只有 -ffast-math 或 -fassociative-math 打开之后才允许向量化这类重排；前面手写的多累加器，等于在源码层面主动放弃了顺序。OpenMP 规范把合并发生的位置和顺序都留成 unspecified，同样的线程数、同样的调度，两次运行也不保证逐位一致[5] 。要可复现性得依赖具体运行时的开关。

整数归约没有这个问题，模 \(2^{k}\) 的加法满足结合律，怎么折树结果都一样。

总操作量、深度和误差

折树不减少加法的总数。一段 \(m\) 个元素的成对归约仍然是 \(m - 1\) 次加法，省下来的是深度：

\[D_{\text{serial}} = m - 1, \qquad D_{\text{tree}} = \lceil \log_2 m \rceil \]

两级结构多出来的是部分和的写出和读回，规模是 \(O(p)\)，相对 \(O(n)\) 的数据量可以忽略；但如果每个线程只领 \(1 \sim 2\) 个元素，\(p\) 就接近 \(n\)，这部分开销会重新变成主要成本，coarsening 正是冲着它去的。

粗化之后深度不再是纯净的对数。每个线程先串行加 \(c\) 个元素得到深度 \(c - 1\)，剩下 \(\frac{n}{c}\) 个部分和再上树：

\[D = (c - 1) + \lceil \log_2 \frac{n}{c} \rceil \]

这个式子也说明收益是有边界的：\(c\) 从 \(1\) 涨到 \(4\)，深度掉一大截；继续往上加，串行那一段就追回来了。

误差用常规模型，\(\text{fl}(a + b) = (a + b)(1 + \delta)\)，\(|\delta| \le \varepsilon\)，并记 \(\gamma_{k} = \frac{k\varepsilon}{1 - k\varepsilon}\)。以同样的 \(\sum_{i} |x_{i}|\) 作基准，串行累加的误差界正比于链长：

\[|S - \hat{S}_{\text{serial}}| \le \gamma_{n - 1} \sum_{i} |x_{i}| \]

成对归约只跟树高有关：

\[|S - \hat{S}_{\text{tree}}| \le \gamma_{\lceil \log_2 n \rceil} \sum_{i} |x_{i}| \]

从 \(O(\varepsilon n)\) 退到 \(O(\varepsilon \log n)\)，这是归约树顺带拿到的好处。两个都是上界，实际误差跟数据分布关系很大，正负相消的序列比这个界乐观得多。另外 \(\gamma_{k}\) 在 \(k\varepsilon\) 接近 \(1\) 的时候就失效了，数据量大到那个程度时这个模型本身不再成立——不过实际算例离这条线很远，暂时不用担心。

误差界这一项在真实归约里通常不是主要矛盾。每个元素只读一遍，算术强度低，时间花在等内存上，所以值得做的是把同一个数组上的多个归约量合并到一次遍历里，比如同时求总和、计数、最小值、最大值，省下来的是一整趟内存访问，而不是几个 \(\log\) 层。

适用场景

这套结构的前提是数据量足够大，而且归约本身确实是瓶颈。算术强度摆在那里，每个元素只参与一次运算，时间大多花在等内存上，所以更值钱的是减少遍历次数——把同一个数组上的和、计数、最值合并到一趟里算完，比继续把树折得更深划算得多。折树解决的是延迟和并行度，只有在数据小到能放进 cache、或者部分和本身变成瓶颈的时候，深度才会重新成为主要矛盾。

输入规模小的时候这套分层是负收益。一个 block 都填不满、或者元素数还不到一条向量的宽度时，部分和的读写、padding 留下的空洞、块间那一轮额外的 kernel 启动，这些固定开销加起来会超过折树省下来的时间，这时候串行累加或者每个线程各算各的更直接。GPU 上尤其明显，两级结构意味着至少两次 kernel 往返，成本要摊到足够多的数据上才划得来。

算子本身也有要求。折树默认加法满足结合律和交换律，浮点加法在数学上并不满足，只能容忍那点舍入误差，前面才会反复强调分块方式会改变结果。如果归约算子真的依赖顺序，比如带顺序语义的合并，那它更接近前缀和，得换 scan 那一套。要求逐位可复现的场合也一样，要么把分块和合并顺序固定下来，要么用补偿求和把舍入压到不影响结论的量级。

线程数这边 CPU 和 GPU 的取舍不一样。CPU 核数有限，线程内累加加线程间合并这两级就够用，再往下细分只会多出同步开销；GPU 上 SM 多、单核弱，才需要 warp 内 shuffle、warp 间共享内存、块间多轮 kernel 这一串分层。coarsening 该切几个元素同样没有通用答案，它跟 block 大小、数据规模、尾块的不均衡程度都相关，只能拿自己的数据实测。

完整实现

完整的 CPU 版并行归约实现：

点击查看代码 
// 线程内四条独立累加链 + 线程间 cacheline 隔离 + 成对合并部分和
#include <algorithm>
#include <cstddef>
#include <vector>

#include <immintrin.h>

namespace {

constexpr std::size_t kCacheLineBytes = 64;       // 一般 x86 的 cacheline 大小
constexpr std::size_t kScalarAccumulators = 4;    // 线程内互相独立的累加链条数

/**
 * @brief 线程私有的部分和槽，独占一条 cacheline 以避免伪共享
 */
struct alignas(kCacheLineBytes) PartialSum {
    double value = 0.0;
    char padding[kCacheLineBytes - sizeof(double)] = {};
};

/**
 * @brief 对一段数据做归约，线程内保留四条独立累加链，收尾只做一次水平归约
 * @param data 数据首地址
 * @param begin 区段起点
 * @param end 区段终点（不含）
 * @return 该区段的和
 * @note 结果依赖分块方式，浮点加法不满足结合律
 */
double reduceChunk(const double* data, std::size_t begin, std::size_t end) {
#if defined(__AVX2__)
    __m256d acc0 = _mm256_setzero_pd();
    __m256d acc1 = _mm256_setzero_pd();
    __m256d acc2 = _mm256_setzero_pd();
    __m256d acc3 = _mm256_setzero_pd();

    std::size_t i = begin;
    const std::size_t step = 4 * kScalarAccumulators;   // 每轮 16 个 double，四条链各吃 4 个
    for (; i + step <= end; i += step) {
        acc0 = _mm256_add_pd(acc0, _mm256_loadu_pd(data + i));
        acc1 = _mm256_add_pd(acc1, _mm256_loadu_pd(data + i + 4));
        acc2 = _mm256_add_pd(acc2, _mm256_loadu_pd(data + i + 8));
        acc3 = _mm256_add_pd(acc3, _mm256_loadu_pd(data + i + 12));
    }

    // 四条链先两两合并，再落到标量，水平归约只在这里出现一次
    acc0 = _mm256_add_pd(acc0, acc1);
    acc2 = _mm256_add_pd(acc2, acc3);
    acc0 = _mm256_add_pd(acc0, acc2);

    double lanes[4];
    _mm256_storeu_pd(lanes, acc0);
    double sum = (lanes[0] + lanes[1]) + (lanes[2] + lanes[3]);

    for (; i < end; ++i) {
        sum += data[i];
    }
    return sum;
#else
    // 没有 AVX2 时保留同样条数的标量累加链，指令级并行的结构不变
    double acc[kScalarAccumulators] = {0.0, 0.0, 0.0, 0.0};

    std::size_t i = begin;
    for (; i + kScalarAccumulators <= end; i += kScalarAccumulators) {
        acc[0] += data[i];
        acc[1] += data[i + 1];
        acc[2] += data[i + 2];
        acc[3] += data[i + 3];
    }

    double sum = (acc[0] + acc[1]) + (acc[2] + acc[3]);
    for (; i < end; ++i) {
        sum += data[i];
    }
    return sum;
#endif
}

/**
 * @brief 把部分和按成对方式合并，合并链深度是 log2(部分和个数)
 * @param partials 部分和数组，返回时只有首元素有意义
 * @return 合并后的结果
 * @note 原地覆盖；个数为奇数时把落单的那个并入前半段的最后一个
 */
double mergePartialsTree(std::vector<PartialSum>& partials) {
    std::size_t count = partials.size();

    while (count > 1) {
        const std::size_t half = (count + 1) / 2;
        for (std::size_t i = 0; i < count - half; ++i) {
            partials[i].value += partials[i + half].value;
        }
        count = half;
    }

    return partials.front().value;
}

}   // namespace

/**
 * @brief 串行参照实现，只用来对照误差
 * @param data 数据首地址
 * @param n 元素个数
 * @return 串行累加的结果
 */
double reduceReference(const double* data, std::size_t n) {
    double sum = 0.0;
    for (std::size_t i = 0; i < n; ++i) {
        sum += data[i];
    }
    return sum;
}

/**
 * @brief 多线程并行归约
 * @param data 数据首地址
 * @param n 元素个数
 * @param threadCount 线程数，小于等于 1 时退化为单线程
 * @return 归约结果
 * @note 每个线程按连续区段取数，尾部区段靠 begin/end 截断处理余数
 * @warning 浮点加法不满足结合律，结果依赖 threadCount 和分块方式
 */
double reduceParallel(const double* data, std::size_t n, int threadCount) {
    if (data == nullptr || n == 0) {
        return 0.0;
    }

    if (threadCount <= 1) {
        return reduceChunk(data, 0, n);
    }

    const std::size_t workers = static_cast<std::size_t>(threadCount);
    std::vector<PartialSum> partials(workers);

#pragma omp parallel for num_threads(threadCount) schedule(static)
    for (int t = 0; t < threadCount; ++t) {
        const std::size_t chunk = (n + workers - 1) / workers;
        const std::size_t begin = std::min(n, chunk * static_cast<std::size_t>(t));
        const std::size_t end = std::min(n, begin + chunk);

        // 相邻线程写的是不同 cacheline，合并阶段之前不会互相干扰
        partials[static_cast<std::size_t>(t)].value = reduceChunk(data, begin, end);
    }

    return mergePartialsTree(partials);
}

  1. https://theartofhpc.com/pcse/omp-reduction.html ↩︎

  2. https://gcc.gnu.org/pipermail/gcc-bugs/2018-January/608916.html ↩︎

  3. https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf ↩︎

  4. https://vuduc.org/teaching/cse6230-hpcta-fa12/slides/cse6230-fa12--05b-reduction-notes.pdf ↩︎

  5. https://www.openmp.org/spec-html/5.0/openmpsu107.html ↩︎

posted @ 2026-09-22 10:58  Ke_scholar  阅读(6)  评论(0)    收藏  举报