Reduce 算子优化笔记

陈王卜2026-08-12 11:49

CUDA Reduce 算子优化笔记（V0 → V7）

从最朴素的规约出发，逐步解决 Warp Divergence、Bank Conflict、Idle 线程、多余同步、循环开销、Shared Memory 依赖、访存效率低等问题。

0. Reduce 基础

Reduce（规约）：多输入 → 单输出。本文以 Sum Reduce 为例，输入长度 N 的 float 数组，输出所有元素之和。

算术强度 = 1 FLOP / 4 Byte = 0.25 FLOP/Byte → 访存密集型（Memory-Bound） ，优化核心是提升内存带宽利用率。

规约本质：树形规约，每轮人数减半，N 个元素约 log₂N 轮收敛。

总体结构 ：每个 block 输出一个 partial sum（写 output[blockIdx.x]），最后由 CPU 加总所有 block 的 partial sum。

V0：朴素并行规约

思路

树形规约，步长从 1 开始每次翻倍。tid % (2*step) == 0 的线程负责加法。

cpp 复制代码 
__global__ void reduce_v0(float* input, float* output, int n) {
    extern __shared__ float smem[];
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    smem[tid] = (gid < n) ? input[gid] : 0.0f;
    __syncthreads();

    for (int step = 1; step < blockDim.x; step *= 2) {
        if (tid % (2 * step) == 0) {
            smem[tid] += smem[tid + step];
        }
        __syncthreads();
    }
    if (tid == 0) output[blockIdx.x] = smem[0];
}

问题：Warp Divergence

GPU 采用 SIMT 模型，一个 Warp（32 线程）共享同一条指令流。当 Warp 内线程走不同分支时，GPU 分两次执行（先满足条件的，再不满足的），吞吐减半。

  • step=1：每 Warp 只有偶数线程工作 → 50% 利用率
  • step=2：每 4 个线程 1 个工作 → 25% 利用率
  • 越来越差

V1：消除 Warp Divergence

思路

用 strided index 替换 tid % (2*step) == 0，让整段 Warp 一起进入/跳过分支。

cpp 复制代码 
for (unsigned int s = 1; s < blockDim.x; s *= 2) {
    int index = threadIdx.x * 2 * s;
    if (index < blockDim.x) {
        smem[index] += smem[index + s];
    }
    __syncthreads();
}

以 blockDim=256 为例：

  • s=1：活跃线程 tid 0~127（4 个完整 Warp）→ 无分化
  • s=2：tid 0~63（2 个完整 Warp）→ 无分化
  • s=4：tid 0~31（1 个完整 Warp）→ 无分化
  • s≥8：活跃线程 <32，仅 Warp0 内部少量分化

前几轮完全消除分化，仅最后几轮有微小分化，开销远小于 V0。

遗留问题：Bank Conflict

V2：解决 Bank Conflict

什么是 Bank Conflict

Shared Memory 划分成 32 个 Bank （每 Bank 4 字节）。同一 Warp 内多个线程访问同一个 Bank 的不同地址 → 访问串行化 → Bank Conflict。

注意：同一 Warp 访问同一地址 是广播，不冲突；只有同一 Bank 的不同地址才冲突。

V1 的冲突分析

V1 的 strided index 让访问地址间隔开，多个线程命中同一 Bank。以 Warp0 第一轮为例：tid0 访问 smem0,1，tid16 访问 smem32,33，而 smem0 和 smem32 都在 Bank0 → 2 路冲突，逐轮恶化。

思路：反转步长方向

从 blockDim.x/2 开始每次减半，且低编号线程始终活跃：

cpp 复制代码 
for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
        smem[tid] += smem[tid + s];
    }
    __syncthreads();
}

以 step=128 为例：tid 0~31 访问 smem0,128、smem1,129...smem31,159，正好覆盖 Bank 0~31 → 无冲突。且活跃线程连续 → 同时消除了 Warp Divergence。

遗留问题：Idle 线程

第一轮后只有一半线程继续工作，后一半线程纯粹当"搬运工"。

V3：解决 Idle 线程

思路：每线程处理 2 个元素

加载阶段每个线程读 2 个元素并提前求和 ，再进规约。用同样 256 线程处理 512 数据，Block 数减半。

cpp 复制代码 
int gid = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
float val = 0.0f;
if (gid < n)              val += input[gid];
if (gid + blockDim.x < n) val += input[gid + blockDim.x];
smem[tid] = val;
__syncthreads();

for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) { ... }

Grid 配置改变：Block 数 = (n + blockDim.x*2 - 1) / (blockDim.x*2)。

收益：Block 数减半 → 调度开销降低、全局访问更规整、每线程做有用功。

遗留问题：多余的 __syncthreads()

当 step <= 32 时只剩 1 个 Warp 工作，此时 __syncthreads() 完全多余（Warp 内天然同步）。

V4：展开最后一个 Warp

思路

规约循环只执行到 step > 32，最后 32 个线程直接展开，省去多余的 __syncthreads()。

cpp 复制代码 
__device__ void warpReduce(volatile float* smem, int tid) {
    smem[tid] += smem[tid + 32];
    smem[tid] += smem[tid + 16];
    smem[tid] += smem[tid +  8];
    smem[tid] += smem[tid +  4];
    smem[tid] += smem[tid +  2];
    smem[tid] += smem[tid +  1];
}

__global__ void reduce_v4(float* input, float* output, int n) {
    // ... 每线程处理 2 元素（继承 V3）
    for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {
        if (tid < s) { smem[tid] += smem[tid + s]; }
        __syncthreads();
    }
    if (tid < 32) warpReduce(smem, tid);   // Warp 内无需同步
    if (tid == 0) output[blockIdx.x] = smem[0];
}

⚠️ warpReduce 的 smem 必须 volatile，防止编译器把中间结果缓存到寄存器。

以 blockDim=256，V3 有 8 轮（step=128,64,32,16,8,4,2,1）每轮同步；V4 把后 5 轮展开，省去 5 次同步。

遗留问题：前几轮仍是运行时循环

V5：完全循环展开

思路

用模板参数把 BLOCK_SIZE 变编译期常量，编译器删除无效分支，生成直线代码。

cpp 复制代码 
template <int BLOCK_SIZE>
__global__ void reduce_v5(float* input, float* output, int n) {
    extern __shared__ float smem[];
    int tid = threadIdx.x;
    int gid = blockIdx.x * (BLOCK_SIZE * 2) + threadIdx.x;

    float val = 0.0f;
    if (gid < n)              val += input[gid];
    if (gid + BLOCK_SIZE < n) val += input[gid + BLOCK_SIZE];
    smem[tid] = val;
    __syncthreads();

    if (BLOCK_SIZE >= 512) { if (tid < 256) smem[tid] += smem[tid+256]; __syncthreads(); }
    if (BLOCK_SIZE >= 256) { if (tid < 128) smem[tid] += smem[tid+128]; __syncthreads(); }
    if (BLOCK_SIZE >= 128) { if (tid <  64) smem[tid] += smem[tid+ 64]; __syncthreads(); }

    if (tid < 32) {
        volatile float* vsmem = smem;
        if (BLOCK_SIZE >= 64) vsmem[tid] += vsmem[tid+32];
        vsmem[tid] += vsmem[tid+16]; vsmem[tid] += vsmem[tid+8];
        vsmem[tid] += vsmem[tid+4];  vsmem[tid] += vsmem[tid+2];
        vsmem[tid] += vsmem[tid+1];
    }
    if (tid == 0) output[blockIdx.x] = smem[0];
}

调用时按 block 大小实例化模板：reduce_v5<256><<<grid, 256, ...>>>。

架构差异（循环展开收益）

差异点 A100 (Ampere) P4 (Pascal)
指令调度器 激进，循环分支易 stall 简单，循环开销占比不高
寄存器 64K/SM，容得下展开代码 32K/SM，展开被寄存器压力抵消
L2 Cache 40 MB，代码缓存命中率高 2 MB，展开代码更大反增 I-Cache miss
SM 数 108，循环开销被放大 40，占比小

结论：循环展开在 Ampere 收益更大，Pascal 上收益有限。

遗留问题：Shared Memory 访问仍是瓶颈

V6：Warp Shuffle 替代 Shared Memory

思路

用 __shfl_down_sync 在寄存器间直接交换数据，省去 Shared Memory 访问。

cpp 复制代码 
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;   // lane 0 持有最终结果
}

__shfl_down_sync(mask, var, delta)：每个线程读取 lane_id + delta 处的寄存器值。比 Shared Memory 快：无地址计算、无 Bank Conflict、延迟更低。

两级规约

  1. Warp 内规约 ：每个 Warp 用 __shfl_down_sync 把 32 值合成 1 值（步长 16→8→4→2→1）
  2. Warp 间规约：各 Warp 结果写入 Shared Memory，再用 Warp0 规约
cpp 复制代码 
__global__ void reduce_v6(float* input, float* output, int n) {
    int tid  = threadIdx.x;
    int gid  = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float val = 0.0f;
    if (gid < n)              val += input[gid];
    if (gid + blockDim.x < n) val += input[gid + blockDim.x];

    val = warpReduceSum(val);                    // 第一级：Warp 内

    __shared__ float warp_results[32];
    if (lane == 0) warp_results[wid] = val;      // 每个 Warp 的 lane0 存结果
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {                              // 第二级：Warp0 汇总
        val = (lane < num_warps) ? warp_results[lane] : 0.0f;
        val = warpReduceSum(val);
    }
    if (tid == 0) output[blockIdx.x] = val;
}

以 blockDim=256：8 个 Warp → 规约得 8 个结果 → 放入 Warp0 再规约。

V7：向量化加载 + Grid Stride Loop

思路

规约流程（两层 Warp 折纸）完全不变，只改数据加载方式：

  1. float4 向量化：一次读 4 个 float，减少访存指令数
  2. Grid Stride Loop：固定 grid 大小，循环扫过全部数据

float4 向量化加载

逐元素 ld.global.f32 每条只读 4 字节，但指令调度开销固定 。用 float4（16 字节）一条指令读 4 个 float，指令数减到 1/4，提升 ILP。

cpp 复制代码 
float4* input4 = reinterpret_cast<float4*>(input);  // 重新解释成 float4 数组
int n4 = n / 4;
...
float4 data = input4[idx];
val += data.x + data.y + data.z + data.w;

前提：指针 16 字节对齐（cudaMalloc 天然满足）。

Grid Stride Loop

之前每个 block 处理固定一小段，N 越大 block 越多，多余 block 排队，且尾部 block 浪费。Grid Stride 固定 grid，每线程循环扫多段。

cpp 复制代码 
int stride = gridDim.x * blockDim.x;               // 步长 = 整个 grid 宽度
for (int idx = blockIdx.x * blockDim.x + tid;
     idx < n4;
     idx += stride) {
    float4 data = input4[idx];
    val += data.x + data.y + data.z + data.w;
}

步长 = 整个 grid 宽度，像扫地毯一样一遍遍扫过去，直到覆盖全部数据。

host 端固定 grid

cpp 复制代码 
int num_sms;
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);
int grid_size = num_sms * 4;      // 恰好填满 GPU
reduce_v7<<<grid_size, BLOCK_SIZE>>>(d_in, d_partial, n);

配套改动：d_partial 大小从 num_blocks 改为 grid_size。

关键心法：grid 大小不再由数据量决定，而是由 SM 数决定。数据再大，人也还是那批，多跑几趟而已。

尾部处理（n 不是 4 的倍数）

cpp 复制代码 
int tail_start = n4 * 4;   // 主循环没碰到的第一个下标
for (int idx = tail_start + blockIdx.x * blockDim.x + tid;
     idx < n;
     idx += stride) {
    val += input[idx];     // 逐元素扫尾巴
}

尾部元素个数永远 < 4（n4 = n/4 整除，零头最多 3）。只有 block0 的 tid 0~2 真正参与（blockIdx.x*blockDim.x + tid < 4），其余越界退出。写成完整公式是为了模式统一 + 防御边界（若尾部更长 / blockDim 更小仍正确）。

完整 kernel

cpp 复制代码 
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

__global__ void reduce_v7(float* input, float* output, int n) {
    int tid  = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float4* input4 = reinterpret_cast<float4*>(input);
    int n4 = n / 4;
    float val = 0.0f;

    int stride = gridDim.x * blockDim.x;
    for (int idx = blockIdx.x * blockDim.x + tid; idx < n4; idx += stride) {
        float4 data = input4[idx];
        val += data.x + data.y + data.z + data.w;
    }

    int tail_start = n4 * 4;
    for (int idx = tail_start + blockIdx.x * blockDim.x + tid; idx < n; idx += stride) {
        val += input[idx];
    }

    val = warpReduceSum(val);

    __shared__ float warp_results[32];
    if (lane == 0) warp_results[wid] = val;
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        val = (lane < num_warps) ? warp_results[lane] : 0.0f;
        val = warpReduceSum(val);
    }
    if (tid == 0) output[blockIdx.x] = val;
}

4 层加法流程（所有版本通用）

层 在哪 谁和谁加 结果存在哪
第1层 每线程自己的 val 线程读到的多个元素自加 每线程寄存器口袋
第2层 warpReduceSum 32 线程 → 1（折纸） lane0 口袋
第3层 再调 warpReduceSum 多个 Warp → 1 block 0 号线程
第4层 host CPU 循环 所有 block → 1 *output

V0~V5 的第1层用 Shared Memory 折纸（V0 小步长大、V1 strided、V2 大步长、V3 双元素、V4 展开最后 warp、V5 全展开）；V6 起第2/3层改用 Warp Shuffle。V7 只在第1层数据加载方式不同（float4 + Grid Stride）。

性能汇总

版本 核心优化点 解决瓶颈
V0 朴素树形 基线
V1 交错寻址 Warp Divergence
V2 步长反转 Bank Conflict + Divergence
V3 每线程双元素 Idle 线程，Grid 减半
V4 展开最后 Warp 多余 syncthreads
V5 模板完全展开 循环开销
V6 Warp Shuffle Shared Memory 依赖
V7 float4 + Grid Stride 访存效率 + GPU 利用率

3080Ti 实测（N = 2^23）

复制代码 
cpu                    diff=0.0000e+00  PASS  0.000 ms
v0_naive               diff=4.6250e+01  PASS  kernel:0.210 ms  160.57 GB/s
v1_divergence          diff=4.6250e+01  PASS  kernel:0.139 ms  241.56 GB/s
v2_bank_conflict       diff=4.6250e+01  PASS  kernel:0.115 ms  291.96 GB/s
v3_idle_threads        diff=4.6250e+01  PASS  kernel:0.069 ms  490.54 GB/s
v4_unroll_last_warp    diff=4.6250e+01  PASS  kernel:0.051 ms  653.05 GB/s
v5_fully_unroll        diff=4.6250e+01  PASS  kernel:0.054 ms  625.62 GB/s
v6_warp_shuffle        diff=4.6250e+01  PASS  kernel:0.053 ms  638.63 GB/s
v7_vectorize           diff=4.6250e+01  PASS  kernel:0.052 ms  644.36 GB/s

注意：所有版本 diff 相同（4.6250e+01），这是浮点累加顺序不同 导致的小误差，但相对误差 diff/expected ≈ 1.1e-5，在容差 1e-4 内，PASS 正确。

工程要点

  1. warpReduceSum 多 .cu 冲突 ：要么各文件 static __device__ 私有化，要么提进头文件用 inline __device__ 共享
  2. 集成 PyTorch ：torch.Tensor.data_ptr() 就是 device 指针，直接喂 (float*, float*, int) 接口
  3. 带宽公式 ：bytes = N * sizeof(float) + num_blocks * sizeof(float)（读 N 元素 + 写 num_blocks partial sum），不要乘 2（那是 copy 的算法）
  4. kernel 计时 ：用 cudaEvent 只测 kernel，host 端 alloc/memcpy 开销会稀释真实性能

代码库：https://gitee.com/chen-buhao/tinyops.git

相关推荐
wabs6661 小时前
关于图论【最短路径之Bellman_ford 算法（判断负权回路）|卡码网95.城市间货物运输II的思考】
数据结构·算法·图论·bellman_ford算法·卡码网·判断负权回路
YouonlyliveonceC1 小时前
《大话数据结构》第4章实战：中缀表达式转后缀 + 后缀表达式求值（完整可运行实现）
算法
phoenix@Capricornus1 小时前
从统计决策到贝叶斯估计
人工智能·算法·机器学习
AICDragon2 小时前
1.8%就够了：K3的896个MoE专家，为什么激活率这么低反而是好事？
人工智能·算法
VALENIAN瓦伦尼安教学设备2 小时前
ASHOOTER激光对中仪如何通过颜色确定调整是否合适
数据库·嵌入式硬件·算法
货拉拉技术2 小时前
重塑 Agent 度量衡：基于 LLM-as-a-Judge 的离线评估体系与实践
算法·设计模式
车压2 小时前
从“为什么”理解注意力机制
算法
wabs6662 小时前
关于哈希表【力扣15.三数之和的思考】
数据结构·算法·leetcode·散列表·哈希表·三数之和
退休倒计时2 小时前
【每日一题】LeetCode 88. 合并两个有序数组 TypeScript
算法·leetcode·职场和发展·typescript
热门推荐
01VSCode+CodeX扩展，无法启动02GitHub 镜像站点03国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）042026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？05微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新06如何新建文件夹？ 电脑新建文件夹的4种方法07AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路08历年考研数学一、数学二、数学三真题试卷及答案PDF092026年8月4日科技热点深度解析：AI大模型群雄逐鹿、卫星互联网组网提速、半导体封装材料革命10Agnes AI 免费 API 接入指南：文本、生图、生视频，一套接口全免费