CUDA 并行算法：从 Reduce 到矩阵转置

Q1：如何在 GPU 上实现 Reduce（规约）操作？有哪些优化方法？

什么是 Reduce？

Reduce（规约）是并行计算中最基础的操作之一：将 N 个元素通过某种二元运算（求和、求最大值、求最小值等）合并为一个标量结果。CPU 上的串行实现是 O(N)；GPU 上的并行实现可以达到 O(log N) 步，但需要精心设计。

朴素实现与问题

最直观的想法是让每个线程将相邻元素两两相加，迭代 log₂(blockDim.x) 轮：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
__global__ void reduce_naive(float* input, float* output, int n) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + threadIdx.x;

    sdata[tid] = (gid < n) ? input[gid] : 0.0f;
    __syncthreads();

    // 树形规约
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }

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

这个版本正确，但存在 Warp 发散：随着 stride 减小，参与计算的线程越来越少，大量线程空转。

优化一：消除 Warp 发散（连续线程活跃）

将参与规约的线程改为连续的低编号线程，减少 Warp 内空闲：

1
2
3
4
5
6
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
        sdata[tid] += sdata[tid + stride];
    }
    __syncthreads();
}

这种写法（高地址向低地址归并）已经消除了大部分发散。

优化二：每线程处理多个元素（减少 Block 数量）

让每个线程在载入共享内存前先做一次全局内存的规约，减少所需 Block 数，同时提升内存访问效率：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
__global__ void reduce_multi(float* input, float* output, int n) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    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];
    sdata[tid] = val;
    __syncthreads();

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

优化三：Warp 内无需同步（最后 32 线程特殊处理）

当 stride ≤ 32 时，所有参与规约的线程都在同一个 Warp 内，Warp 内天然同步，无需 __syncthreads()。可以使用 __shfl_down_sync 做 Warp 内规约（见 Q2），大幅减少同步开销。

优化四：使用 Cooperative Groups（推荐方式）

CUDA 提供了 Cooperative Groups API，内置硬件加速的规约操作（计算能力 ≥ 8.0 时启用硬件路径）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
namespace cg = cooperative_groups;

__global__ void reduce_cg(float* input, float* output, int n) {
    cg::thread_block block = cg::this_thread_block();
    extern __shared__ float sdata[];

    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[threadIdx.x] = (gid < n) ? input[gid] : 0.0f;
    block.sync();

    // CG 提供的 reduce，内部自动选择最优实现
    float sum = cg::reduce(block, sdata[threadIdx.x], cg::plus<float>());

    if (block.thread_rank() == 0) output[blockIdx.x] = sum;
}

cg::reduce 支持的操作符：

操作符 含义
cg::plus<T>() 求和
cg::less<T>() 求最小值
cg::greater<T>() 求最大值
cg::bit_and<T>() 按位与
cg::bit_or<T>() 按位或
cg::bit_xor<T>() 按位异或

计算能力 ≥ 8.0 时，4 字节整型的规约由硬件直接加速。

两阶段 Reduce

单次 kernel 只能规约一个 Block 内的数据。对于超大数组，需要两阶段规约：

1
2
第一阶段：每个 Block 规约出一个局部结果，写到 output[blockIdx.x]
第二阶段：用 1 个 Block 规约 output[] 数组得到最终结果

也可以用原子加将各 Block 结果直接累加到全局变量，但原子操作在大量 Block 并发时会有争用（适合中等规模）：

1
if (tid == 0) atomicAdd(global_result, sdata[0]);

性能总结

优化策略 效果
消除 Warp 发散 减少空闲线程
每线程处理多元素 提升内存带宽利用率
Warp 内 shfl 规约 避免最后几轮 __syncthreads
Cooperative Groups 自动最优实现，代码简洁
两阶段规约 处理任意规模数组

Q2：如何实现 Warp-level Reduce（使用 __shfl_down_sync）？

Warp Shuffle 原语

Warp Shuffle 函数允许同一 Warp 内的线程直接交换寄存器数据，无需经过共享内存，延迟更低、带宽更高。关键函数签名：

1
2
3
4
T __shfl_sync    (unsigned mask, T value, int srcLane, int width=32);
T __shfl_up_sync (unsigned mask, T value, unsigned delta, int width=32);
T __shfl_down_sync(unsigned mask, T value, unsigned delta, int width=32);
T __shfl_xor_sync (unsigned mask, T value, int laneMask, int width=32);
  • mask：参与操作的线程掩码，通常为 0xFFFFFFFF（全部 32 线程）
  • value：当前线程贡献的值
  • delta：读取来源 Lane 的偏移量
  • width：操作的子 Warp 宽度，必须是 2 的幂（1/2/4/8/16/32）

__shfl_down_sync 的语义：每个线程读取 laneId + delta 号线程的 value。高编号线程（lane >= width - delta）保持自身原值不变。

Warp-level 求和规约

利用 __shfl_down_sync 实现标准的蝶形规约，5 次 shuffle 完成 32 个元素的求和：

1
2
3
4
5
6
7
8
9
10
__device__ float warp_reduce_sum(float val) {
    // 蝶形规约：每次将高半部分的值加到低半部分
    val += __shfl_down_sync(0xFFFFFFFF, val, 16);
    val += __shfl_down_sync(0xFFFFFFFF, val, 8);
    val += __shfl_down_sync(0xFFFFFFFF, val, 4);
    val += __shfl_down_sync(0xFFFFFFFF, val, 2);
    val += __shfl_down_sync(0xFFFFFFFF, val, 1);
    // lane 0 持有最终结果
    return val;
}

规约过程（以 8 线程简化举例）：

1
2
3
4
5
6
7
初始:  [v0, v1, v2, v3, v4, v5, v6, v7]
步骤1 (delta=4): lane[i] += lane[i+4]
       [v0+v4, v1+v5, v2+v6, v3+v7, v4, v5, v6, v7]
步骤2 (delta=2): lane[i] += lane[i+2]
       [v0+v2+v4+v6, v1+v3+v5+v7, ...]
步骤3 (delta=1): lane[i] += lane[i+1]
       [v0+v1+v2+...+v7, ...]   ← lane 0 是最终结果

Block-level 规约（组合 Warp 结果）

一个 Block 有多个 Warp，需要将各 Warp 的结果汇集：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
__device__ float block_reduce_sum(float val) {
    static __shared__ float warp_sums[32]; // 最多 32 个 Warp
    int lane   = threadIdx.x % 32;
    int warp_id = threadIdx.x / 32;

    // 第一步：Warp 内规约
    val = warp_reduce_sum(val);

    // 第二步：每个 Warp 的 lane 0 写入共享内存
    if (lane == 0) warp_sums[warp_id] = val;
    __syncthreads();

    // 第三步：第一个 Warp 规约所有 Warp 的结果
    int num_warps = blockDim.x / 32;
    val = (lane < num_warps) ? warp_sums[lane] : 0.0f;
    if (warp_id == 0) val = warp_reduce_sum(val);

    return val; // 所有线程持有结果（或只在 thread 0 使用）
}

计算能力 8.0+ 的硬件规约

对于 ≥ Ampere 的 GPU，CUDA 提供了更简洁的硬件规约 intrinsics：

1
2
3
4
5
6
7
// 对 mask 指定的线程求和，直接返回结果
unsigned __reduce_add_sync(unsigned mask, unsigned value);
unsigned __reduce_min_sync(unsigned mask, unsigned value);
unsigned __reduce_max_sync(unsigned mask, unsigned value);
unsigned __reduce_and_sync(unsigned mask, unsigned value);
unsigned __reduce_or_sync (unsigned mask, unsigned value);
unsigned __reduce_xor_sync(unsigned mask, unsigned value);

注意：这些 intrinsics 仅支持 32 位整型，且不提供内存屏障。

使用注意事项

  1. 掩码必须精确：参与 shuffle 的所有线程必须都包含在 mask 中，否则行为未定义
  2. 不得在分支内不对称地调用：若 lane 0 不参与调用，则行为未定义
  3. width 必须是 2 的幂：1、2、4、8、16、32，其他值结果未定义
  4. 不保证内存顺序：shuffle 操作本身不是内存屏障
1
2
3
4
5
// 错误示例：lane 0 不参与，行为未定义
int result = (laneId > 0) ? __shfl_sync(0xFFFFFFFF, value, 0) : 0;

// 正确示例：所有线程都参与
int result = __shfl_sync(0xFFFFFFFF, value, 0);

Q3：如何在 CUDA 中实现 Prefix Sum（前缀和/Scan）？

什么是 Prefix Sum？

Prefix Sum（前缀和）分两种：

  • Exclusive Scan（排他前缀和）：output[i] = sum(input[0..i-1])，第 i 个输出不含第 i 个输入
  • Inclusive Scan（包含前缀和）：output[i] = sum(input[0..i])，第 i 个输出含第 i 个输入
1
2
3
input:          [3, 1, 7, 2, 5]
inclusive scan: [3, 4, 11, 13, 18]
exclusive scan: [0, 3, 4, 11, 13]

Prefix Sum 是并行计算的基础原语，用于流压缩、稀疏矩阵、排序等众多算法。

方法一：Hillis-Steele 算法（Work-Inefficient）

Hillis-Steele 是最简单的并行 scan，步骤数为 O(log N)，但总工作量是 O(N log N)（不节省工作）：

1
2
3
4
步骤 1 (stride=1): 每个位置加左边 1 步的值
步骤 2 (stride=2): 每个位置加左边 2 步的值
步骤 3 (stride=4): 每个位置加左边 4 步的值
...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
__global__ void scan_hillis_steele(float* data, float* output, int n) {
    extern __shared__ float temp[];
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;

    temp[tid] = (gid < n) ? data[gid] : 0.0f;
    __syncthreads();

    for (int stride = 1; stride < blockDim.x; stride *= 2) {
        float val = (tid >= stride) ? temp[tid - stride] : 0.0f;
        __syncthreads();
        temp[tid] += val;
        __syncthreads();
    }

    if (gid < n) output[gid] = temp[tid];
}

方法二：Blelloch 算法（Work-Efficient，推荐）

Blelloch 算法分两个阶段，总工作量 O(N)，步骤数 O(log N)：

上扫（Up-sweep / Reduce）阶段：从叶到根，建立树状求和结构
下扫（Down-sweep）阶段：从根到叶，传播前缀和

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
__global__ void scan_blelloch(float* data, float* output, int n) {
    extern __shared__ float temp[];
    int tid = threadIdx.x;
    int offset = 1;

    // 载入数据
    int ai = tid, bi = tid + (n / 2);
    temp[ai] = data[ai];
    temp[bi] = data[bi];

    // ===== 上扫（Reduce）=====
    for (int d = n >> 1; d > 0; d >>= 1) {
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            temp[bi] += temp[ai];
        }
        offset *= 2;
    }

    // 将根节点清零（转为 exclusive scan）
    if (tid == 0) temp[n - 1] = 0;

    // ===== 下扫（Down-sweep）=====
    for (int d = 1; d < n; d *= 2) {
        offset >>= 1;
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            float t = temp[ai];
            temp[ai] = temp[bi];
            temp[bi] += t;
        }
    }

    __syncthreads();
    output[ai] = temp[ai];
    output[bi] = temp[bi];
}

方法三：使用 Cooperative Groups（最简洁）

CUDA 的 Cooperative Groups 提供了开箱即用的 scan 函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#include <cooperative_groups.h>
#include <cooperative_groups/scan.h>
namespace cg = cooperative_groups;

__global__ void scan_cg(int* data, int* output, int n) {
    cg::thread_block block = cg::this_thread_block();
    int tid = block.thread_rank();
    int val = data[tid];

    // Exclusive scan（每个线程的结果是其之前所有线程的和）
    int prefix = cg::exclusive_scan(block, val, cg::plus<int>());

    // Inclusive scan（含自身）
    // int prefix = cg::inclusive_scan(block, val, cg::plus<int>());

    output[tid] = prefix;
}

exclusive_scan 和 inclusive_scan 支持与 reduce 相同的算子（plus、less、greater、bit_and 等），也支持 lambda 和自定义函数对象。

大数组的分块处理策略

单个 Block 只能处理 blockDim.x 个元素。处理大数组需要分三步：

1
2
3
1. 每个 Block 做局部 scan，同时记录每个 Block 的总和到 block_sums[]
2. 对 block_sums[] 做一次 scan，得到各 Block 的偏移量 offsets[]
3. 每个 Block 将 offsets[blockIdx.x] 加到本 Block 所有局部结果上

这也是 CUB 库中 cub::DeviceScan 的实现思路。实际工程中推荐直接使用 CUB：

1
2
3
4
5
6
7
8
9
10
#include <cub/cub.cuh>

// 获取临时存储大小
void* d_temp = nullptr;
size_t temp_bytes = 0;
cub::DeviceScan::ExclusiveSum(d_temp, temp_bytes, d_input, d_output, n);

// 分配并执行
cudaMalloc(&d_temp, temp_bytes);
cub::DeviceScan::ExclusiveSum(d_temp, temp_bytes, d_input, d_output, n);

Q4：如何实现 Softmax 的 GPU 并行化？

Softmax 的定义与挑战

Softmax 是深度学习中最常用的归一化操作：

softmax(xi)=exi∑jexj\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}

直接计算存在数值溢出问题：当 x 较大时，exp(x) 会超出 float 范围。标准做法是减去最大值进行稳定化：

softmax(xi)=exi−max⁡(x)∑jexj−max⁡(x)\text{softmax}(x_i) = \frac{e^{x_i - \max(x)}}{\sum_j e^{x_j - \max(x)}}

朴素三遍实现

将 softmax 分解为三个 Kernel：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
// Kernel 1: 求每行最大值
__global__ void find_max(float* input, float* max_vals, int cols) {
    int row = blockIdx.x;
    float max_val = -FLT_MAX;
    for (int col = threadIdx.x; col < cols; col += blockDim.x)
        max_val = max(max_val, input[row * cols + col]);
    // Warp reduce 求 max
    max_val = warp_reduce_max(max_val);
    if (threadIdx.x == 0) max_vals[row] = max_val;
}

// Kernel 2: exp(x - max) 并求和
__global__ void exp_and_sum(float* input, float* output, float* max_vals,
                             float* sum_vals, int cols) {
    int row = blockIdx.x;
    float sum = 0.0f;
    float m = max_vals[row];
    for (int col = threadIdx.x; col < cols; col += blockDim.x) {
        float e = expf(input[row * cols + col] - m);
        output[row * cols + col] = e;
        sum += e;
    }
    sum = warp_reduce_sum(sum);
    if (threadIdx.x == 0) sum_vals[row] = sum;
}

// Kernel 3: 归一化
__global__ void normalize(float* output, float* sum_vals, int cols) {
    int row = blockIdx.x;
    float inv_sum = 1.0f / sum_vals[row];
    for (int col = threadIdx.x; col < cols; col += blockDim.x)
        output[row * cols + col] *= inv_sum;
}

三遍实现需要三次全局内存读写，带宽利用率低。

两遍融合实现（Online Softmax）

利用 Online Softmax 算法，在一遍扫描中同时维护最大值和归一化因子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
__global__ void softmax_fused(float* input, float* output, int rows, int cols) {
    int row = blockIdx.x;
    float* x = input + row * cols;
    float* y = output + row * cols;

    // 第一遍：求最大值和 sum（Online 方式可合并，这里分开以清晰展示）
    float max_val = -FLT_MAX;
    float sum = 0.0f;

    // 求 max（每线程处理多列）
    for (int col = threadIdx.x; col < cols; col += blockDim.x)
        max_val = max(max_val, x[col]);
    // Block 内规约 max
    max_val = block_reduce_max(max_val);  // 使用共享内存规约

    // 求 exp(x-max) 的和
    for (int col = threadIdx.x; col < cols; col += blockDim.x)
        sum += expf(x[col] - max_val);
    sum = block_reduce_sum(sum);

    // 归一化
    float inv_sum = 1.0f / sum;
    for (int col = threadIdx.x; col < cols; col += blockDim.x)
        y[col] = expf(x[col] - max_val) * inv_sum;
}

Online Softmax 的单遍算法

Flash Attention 引入的 Online Softmax 通过维护一对状态变量 (m, l) 实现真正的单遍扫描：

  • m：当前见过的最大值
  • l：当前归一化因子（累积 exp 之和）

每当新的最大值 m_new 出现时，旧的 l 需要用 exp(m_old - m_new) 重新缩放：

1
2
3
4
5
6
7
8
9
// 单遍 Online Softmax 核心状态更新
float m = -FLT_MAX, l = 0.0f;
for (int j = 0; j < cols; j++) {
    float x_j = x[j];
    float m_new = max(m, x_j);
    l = l * expf(m - m_new) + expf(x_j - m_new);
    m = m_new;
}
// 最终：output[j] = exp(x[j] - m) / l

这在 GPU 上通过 Warp Shuffle 完成并行规约，是 Flash Attention 的核心技巧。

性能建议

场景 推荐方案
列数 ≤ 1024 单 Block 处理一行，用 shfl 规约
列数很大 多 Block 协作，两阶段规约
Transformer 注意力 使用 Flash Attention（IO-aware）
训练框架 torch.nn.functional.softmax（已高度优化）

Q5：如何实现 LayerNorm / BatchNorm 的 CUDA Kernel？

数学定义

BatchNorm：在 batch 维度和空间维度上归一化，适合 CNN

x^i=xi−μBσB2+ϵ⋅γ+β\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}} \cdot \gamma + \beta

其中 μB\mu_B、σB2\sigma_B^2 在每个 batch 内的对应通道上统计。

LayerNorm：在每个样本的特征维度上归一化，适合 Transformer

x^i=xi−μσ2+ϵ⋅γ+β\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta

其中 μ\mu、σ2\sigma^2 在单个样本的所有特征上统计。

LayerNorm CUDA 实现

LayerNorm 每行独立计算，非常适合 GPU 并行——每个 Block 负责一行：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
__global__ void layer_norm(float* input, float* output,
                            float* gamma, float* beta,
                            int N, float eps) {
    // 每个 Block 处理一行（一个样本的特征向量）
    int row = blockIdx.x;
    float* x = input + row * N;
    float* y = output + row * N;

    // ===== 第一步：计算均值 =====
    float mean = 0.0f;
    for (int i = threadIdx.x; i < N; i += blockDim.x)
        mean += x[i];
    mean = block_reduce_sum(mean) / N;

    // ===== 第二步：计算方差 =====
    float var = 0.0f;
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        float diff = x[i] - mean;
        var += diff * diff;
    }
    var = block_reduce_sum(var) / N;
    float inv_std = rsqrtf(var + eps);

    // ===== 第三步：归一化 + 仿射变换 =====
    for (int i = threadIdx.x; i < N; i += blockDim.x)
        y[i] = (x[i] - mean) * inv_std * gamma[i] + beta[i];
}

关键细节：

  • block_reduce_sum 使用共享内存 + Warp shuffle 实现（见 Q1/Q2）
  • rsqrtf 是 GPU 内建的快速倒数平方根，比 1.0f / sqrtf() 更快
  • 均值需要广播到所有线程，通过共享内存传递

Welford 在线算法（一遍计算均值和方差）

上面的实现需要两遍遍历数据（一次求均值，一次求方差）。Welford 算法可以一遍同时计算均值和方差，减少全局内存访问：

1
2
3
4
5
6
7
8
9
10
11
// Welford 在线更新
float mean = 0.0f, m2 = 0.0f;
int count = 0;
for (int i = threadIdx.x; i < N; i += blockDim.x) {
    count++;
    float delta = x[i] - mean;
    mean += delta / count;
    float delta2 = x[i] - mean;
    m2 += delta * delta2;
}
// m2 / count 即为方差（需要在 Block 间合并 Welford 状态）

Welford 的状态合并需要特殊公式，适合封装为函数。PyTorch 的 LayerNorm CUDA 实现就使用了这一技术。

BatchNorm 与 LayerNorm 的区别

对比项 BatchNorm LayerNorm
统计维度 Batch + 空间 特征维度（单样本）
适合场景 CNN（大 batch） Transformer、RNN（小 batch）
推理行为 使用训练时统计的移动均值 实时计算，无移动均值
GPU 并行策略 需要跨样本规约（复杂） 每样本独立（简单）
并行化难度 较高（需要跨 Block 同步） 低（Block 独立）

BatchNorm 需要在 Batch 维度上规约，跨 Block 协调更复杂，通常需要两 Kernel 或 Cooperative Groups 全局同步。

RMSNorm（简化版 LayerNorm）

大语言模型（如 LLaMA）常用 RMSNorm，去掉了均值归一化，仅保留方差归一化，速度更快：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
__global__ void rms_norm(float* input, float* output, float* gamma,
                          int N, float eps) {
    int row = blockIdx.x;
    float* x = input + row * N;
    float* y = output + row * N;

    // 只计算均方根（无需均值）
    float rms = 0.0f;
    for (int i = threadIdx.x; i < N; i += blockDim.x)
        rms += x[i] * x[i];
    rms = block_reduce_sum(rms) / N;
    float inv_rms = rsqrtf(rms + eps);

    for (int i = threadIdx.x; i < N; i += blockDim.x)
        y[i] = x[i] * inv_rms * gamma[i];
}

Q6：如何实现 Histogram（直方图）统计？

问题与挑战

直方图统计需要对每个输入元素原子地更新对应的 bin 计数，是典型的随机写、高竞争场景。直接对全局内存做原子加会导致大量竞争，性能很差。

方法一：共享内存局部直方图

标准优化策略：每个 Block 先在共享内存上维护局部直方图，Block 结束时再原子加到全局内存：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
__global__ void histogram_smem(int* input, int* output,
                                 int n, int nbins) {
    extern __shared__ int local_hist[];  // 大小 = nbins

    // 初始化共享内存直方图
    for (int i = threadIdx.x; i < nbins; i += blockDim.x)
        local_hist[i] = 0;
    __syncthreads();

    // 统计：对共享内存原子加（竞争范围仅在 Block 内）
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid < n) {
        int bin = input[gid];  // 假设 bin 范围已合法
        atomicAdd(&local_hist[bin], 1);
    }
    __syncthreads();

    // 将局部结果合并到全局内存
    for (int i = threadIdx.x; i < nbins; i += blockDim.x)
        atomicAdd(&output[i], local_hist[i]);
}

// 启动：动态共享内存 = nbins * sizeof(int)
histogram_smem<<<gridDim, blockDim, nbins * sizeof(int)>>>(input, output, n, nbins);

适用条件：nbins 须能放入共享内存。一般不超过 4096 个 bin（16 KB = 4096 个 int）。

方法二：分布式共享内存直方图（Thread Block Cluster）

当 bin 数量超过单个 Block 的共享内存时，可以利用 CUDA 的 Thread Block Cluster 功能，将多个 Block 的共享内存合并为分布式共享内存，扩大可用容量。Cluster 内的线程块可以直接读写其他线程块的共享内存：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
#include <cooperative_groups.h>

__global__ void clusterHist_kernel(int* bins, const int nbins,
                                    const int bins_per_block,
                                    const int* __restrict__ input,
                                    size_t array_size) {
    extern __shared__ int smem[];
    namespace cg = cooperative_groups;

    int tid = cg::this_grid().thread_rank();
    cg::cluster_group cluster = cg::this_cluster();

    // 初始化本 Block 的共享内存段
    for (int i = threadIdx.x; i < bins_per_block; i += blockDim.x)
        smem[i] = 0;
    cluster.sync();  // 确保所有 Block 初始化完毕

    // 处理数据：根据 bin id 找到目标 Block，写入其共享内存
    for (int i = tid; i < array_size; i += blockDim.x * gridDim.x) {
        int binid = max(0, min(input[i], nbins - 1));
        int dst_block = binid / bins_per_block;
        int dst_offset = binid % bins_per_block;

        // 通过分布式共享内存指针跨 Block 原子加
        int* dst_smem = cluster.map_shared_rank(smem, dst_block);
        atomicAdd(dst_smem + dst_offset, 1);
    }
    cluster.sync();

    // 将各 Block 的局部直方图写入全局内存
    int* lbins = bins + cluster.block_rank() * bins_per_block;
    for (int i = threadIdx.x; i < bins_per_block; i += blockDim.x)
        atomicAdd(&lbins[i], smem[i]);
}

Cluster 大小根据 nbins 动态决定，cluster_size = nbins / bins_per_block。这样整个 Cluster 的分布式共享内存恰好能放下所有 bin。

方法三：多副本直方图（减少竞争）

为每个 Block 或每个 Warp 维护独立的直方图副本，最后归并。这用内存换竞争减少：

1
2
3
每线程维护独立直方图 → 极少竞争，但内存消耗大
每 Warp 维护共享副本 → 平衡点较好
每 Block 一份（标准方法）→ 最常用

方法选择

nbins 大小 推荐方案
≤ 1K bins 共享内存直方图（方法一）
1K ~ 64K bins Cluster 分布式共享内存（方法二，需 CC ≥ 9.0）
> 64K bins 或旧 GPU 直接全局内存原子（接受竞争开销）
特殊：float/大范围 先量化到固定 bin 数，再统计

Q7：如何将数组按奇偶位置分组（类似 Partition）？

问题定义

Partition 操作将数组元素按某一条件（谓词）分为两组，稳定地输出满足条件的元素（前半部分）和不满足条件的元素（后半部分）。奇偶分组是最简单的例子。

1
2
3
输入: [1, 2, 3, 4, 5, 6, 7, 8]
谓词: 奇数
输出: [1, 3, 5, 7, 2, 4, 6, 8]  // 奇数在前，偶数在后

核心：Stream Compaction（流压缩）

Partition 本质是 Stream Compaction，其标准并行实现分三步：

  1. 生成谓词掩码：对每个元素计算 0/1 谓词
  2. 前缀和：对谓词数组做 exclusive scan，得到每个满足条件元素的目标位置
  3. 散射（Scatter）：将满足条件的元素写到前缀和指示的位置
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// Step 1: 计算谓词（是否为奇数）
__global__ void compute_predicate(int* input, int* predicate, int n) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid < n) predicate[gid] = input[gid] % 2;  // 1=奇数, 0=偶数
}

// Step 2: 对 predicate[] 做 exclusive scan，得到奇数写入位置
// （使用 CUB: cub::DeviceScan::ExclusiveSum）

// Step 3: 散射奇数
__global__ void scatter_odd(int* input, int* predicate, int* scan_result,
                              int* output, int n, int num_odd) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid < n && predicate[gid] == 1)
        output[scan_result[gid]] = input[gid];  // 写到前半部分
}

// Step 4: 类似地，对 (1-predicate) 做 exclusive scan，散射偶数到后半部分

Warp 级别的快速 Partition

在 Warp 内，利用 __ballot_sync 获得谓词位掩码，再用 __popc（popcount）计算前缀偏移，实现 O(1) 的 Warp 内 partition：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
__device__ void warp_partition(int* val, int predicate, int* out_true, int* out_false,
                                 int* n_true, int* n_false) {
    unsigned mask = 0xFFFFFFFF;
    int lane = threadIdx.x % 32;

    // 获取本 Warp 中满足谓词的线程掩码
    unsigned ballot = __ballot_sync(mask, predicate);
    int num_true = __popc(ballot);  // 满足条件的线程数

    // 计算本线程的输出位置
    unsigned lower_mask = (1u << lane) - 1;
    int true_idx  = __popc(ballot & lower_mask);       // 前面有多少个满足条件的
    int false_idx = lane - true_idx;                    // 前面有多少个不满足的

    // 散射到输出
    if (predicate) out_true[true_idx]   = *val;
    else           out_false[false_idx] = *val;

    *n_true  = num_true;
    *n_false = 32 - num_true;
}

使用 CUB（实际工程推荐）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#include <cub/cub.cuh>

// CUB 提供完整的 DevicePartition
void* d_temp = nullptr;
size_t temp_bytes = 0;

// 获取所需临时空间大小
cub::DevicePartition::If(d_temp, temp_bytes,
    d_input, d_output, d_num_selected, n,
    IsOdd{});  // IsOdd 是实现了 operator() 的函数对象

// 分配并执行
cudaMalloc(&d_temp, temp_bytes);
cub::DevicePartition::If(d_temp, temp_bytes,
    d_input, d_output, d_num_selected, n, IsOdd{});

使用场景

  • Sparse 计算：压缩掉零元素，只处理非零值
  • 光线追踪：过滤命中/未命中的光线
  • 图神经网络：筛选活跃节点
  • 排序：Radix Sort 内部大量使用 partition

Q8：如何在 CUDA 上高效实现 Transpose（矩阵转置）？

朴素实现的问题

转置操作：C[col][row] = A[row][col]

1
2
3
4
5
6
__global__ void naive_transpose(float* A, float* C, int M) {
    int col = blockDim.x * blockIdx.x + threadIdx.x;
    int row = blockDim.y * blockIdx.y + threadIdx.y;
    if (row < M && col < M)
        C[col * M + row] = A[row * M + col];
}

看似简单，但存在严重性能问题：

  • 读 A：A[row * M + col]，相邻线程读相邻列，内存访问合并，带宽利用率 100%
  • 写 C：C[col * M + row]，相邻线程（不同 col）写的地址相差 M 个元素，完全不合并

当 M > 32 时，写操作的带宽利用率只有 12.5%，成为瓶颈。

共享内存优化转置

核心思想：先将一个 Tile 以合并方式读入共享内存，再以合并方式写出到转置后的位置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
#define TILE 32

__global__ void smem_transpose(float* A, float* C, int M) {
    // 注意：+1 避免 Bank Conflict（列方向相邻元素步长 = 33，不整除 32）
    __shared__ float tile[TILE][TILE + 1];

    // 确定当前 Tile 的左上角位置
    int tileCol = blockIdx.x * TILE;
    int tileRow = blockIdx.y * TILE;

    // ===== 合并读 A =====
    // threadIdx.x 是列偏移，threadIdx.y 是行偏移
    // 行主序：相邻 threadIdx.x 读相邻内存地址 → 合并
    int srcCol = tileCol + threadIdx.x;
    int srcRow = tileRow + threadIdx.y;
    if (srcRow < M && srcCol < M)
        tile[threadIdx.y][threadIdx.x] = A[srcRow * M + srcCol];
    __syncthreads();

    // ===== 合并写 C（转置后）=====
    // 目标位置：原来的 (row, col) 块 → 变为 (col, row) 块
    int dstCol = tileRow + threadIdx.x;   // 目标列 = 源行块起始 + 列内偏移
    int dstRow = tileCol + threadIdx.y;   // 目标行 = 源列块起始 + 行内偏移
    // 从共享内存读：tile[threadIdx.x][threadIdx.y]（转置了 xy）
    // 相邻 threadIdx.x 读 tile 的相邻行首元素 → 不同行，但步长=TILE+1 → 无 Bank Conflict
    if (dstRow < M && dstCol < M)
        C[dstRow * M + dstCol] = tile[threadIdx.x][threadIdx.y];
}

为什么 tile[TILE][TILE+1] 能避免 Bank Conflict？

共享内存有 32 个 Bank，每个 Bank 宽 4 字节。二维数组 tile[r][c] 的地址为 r * (TILE+1) + c。

写入时（tile[threadIdx.y][threadIdx.x]），同一行的 threadIdx.x 从 0 到 31 访问 bank 0 到 31，无冲突。

读取时（tile[threadIdx.x][threadIdx.y]），同一组线程的 threadIdx.x 从 0 到 31：

  • 若使用 tile[TILE][TILE]，地址为 threadIdx.x * 32 + threadIdx.y，步长 32，全部落在同一 Bank → 32 路冲突
  • 若使用 tile[TILE][TILE+1]，地址为 threadIdx.x * 33 + threadIdx.y，步长 33 与 32 互质 → 无冲突

性能对比

实现 读带宽 写带宽 总效率
朴素转置 100% 合并 12.5% 未合并 ~50%
共享内存转置（无padding） 100% 合并 100% 合并 ~80%（Bank Conflict）
共享内存转置（+1 padding） 100% 合并 100% 合并 ~95%

优化后的转置性能接近内存带宽的理论上限（copy 操作的两倍带宽需求）。

Q9：如何实现 NCHW 到 NHWC 格式的转换？

格式定义

深度学习框架中张量的内存布局有两种主要格式：

  • NCHW（Channels-First）：[Batch, Channel, Height, Width]，内存顺序为 W 变化最快
    • 地址：n * C*H*W + c * H*W + h * W + w
    • PyTorch 默认格式，适合 CUDA GEMM 计算
  • NHWC（Channels-Last）：[Batch, Height, Width, Channel]，内存顺序为 C 变化最快
    • 地址：n * H*W*C + h * W*C + w * C + c
    • cuDNN 的最优格式，Tensor Core 对此格式效率更高

格式转换的内存访问分析

NCHW → NHWC 转换：output[n][h][w][c] = input[n][c][h][w]

如果简单实现：

1
2
3
4
5
6
7
8
9
10
11
12
__global__ void nchw_to_nhwc_naive(float* input, float* output,
                                     int N, int C, int H, int W) {
    int n = blockIdx.z;
    int c = blockIdx.y * blockDim.y + threadIdx.y;
    int hw = blockIdx.x * blockDim.x + threadIdx.x;
    int h = hw / W, w = hw % W;

    if (n < N && c < C && h < H && w < W) {
        float val = input[n * C*H*W + c * H*W + h * W + w];  // 读不合并（c 方向跳步）
        output[n * H*W*C + h * W*C + w * C + c] = val;       // 写不合并（c 方向跳步）
    }
}

两个方向的访问都不合并，性能很差。

优化一：共享内存 Tile 转换

类似矩阵转置，先合并读入共享内存，再合并写出：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
#define TILE_C 32
#define TILE_HW 32

__global__ void nchw_to_nhwc_smem(float* input, float* output,
                                    int N, int C, int H, int W) {
    // 每个 Block 处理一个 [TILE_C × TILE_HW] 的 tile
    __shared__ float tile[TILE_HW][TILE_C + 1];  // +1 避免 Bank Conflict

    int n  = blockIdx.z;
    int c0 = blockIdx.y * TILE_C;   // 当前 tile 的起始 channel
    int hw0 = blockIdx.x * TILE_HW; // 当前 tile 的起始 hw 索引

    int tc = threadIdx.x;   // tile 内的 channel 偏移
    int thw = threadIdx.y;  // tile 内的 hw 偏移

    int c  = c0 + tc;
    int hw = hw0 + thw;
    int h = hw / W, w = hw % W;

    // ===== 合并读 NCHW =====
    // 固定 c，相邻 thw 对应相邻 hw，地址连续 → 合并
    if (n < N && c < C && hw < H * W)
        tile[thw][tc] = input[n * C*H*W + c * H*W + hw];
    __syncthreads();

    // ===== 合并写 NHWC =====
    // 固定 hw，相邻 tc 对应相邻 channel，地址连续 → 合并
    if (n < N && c < C && hw < H * W)
        output[n * H*W*C + hw * C + c] = tile[thw][tc];
}

优化二：float4 向量化读写

当 C 是 4 的倍数时，使用 float4 每次读写 16 字节，将内存事务数减少到 1/4：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
__global__ void nchw_to_nhwc_vec4(float4* input, float4* output,
                                    int N, int C4, int H, int W) {
    // C4 = C / 4，每次处理 4 个 channel
    int n  = blockIdx.z;
    int c4 = blockIdx.y * blockDim.y + threadIdx.y;  // 处理第 c4*4 ~ c4*4+3 个 channel
    int hw = blockIdx.x * blockDim.x + threadIdx.x;

    if (n < N && c4 < C4 && hw < H * W) {
        float4 val = input[n * C4 * H * W + c4 * H * W + hw];
        // NHWC 写：output[n][hw][c4] - 需要 AOS to SOA 转换（更复杂）
        // 简化版（连续 channel 组合并写）
        output[n * H * W * C4 + hw * C4 + c4] = val;
    }
}

注意：向量化要求地址 16 字节对齐，即 (float*)ptr % 16 == 0。

推荐实践

实际工程中不建议手动实现这些格式转换，而是：

  1. 使用 cuDNN：cudnnTransformTensor 支持任意格式转换，有高度优化
  2. 使用 PyTorch：tensor.contiguous(memory_format=torch.channels_last) 内部调用优化的 CUDA 实现
  3. 设计时统一格式：在模型设计阶段决定统一使用 NHWC 或 NCHW，避免运行时转换

NCHW vs NHWC 性能对比

操作 NCHW 优势 NHWC 优势
一般 GEMM ✓（PyTorch 默认）
cuDNN 卷积 ✓（最优路径）
Tensor Core 计算 ✓（需要特定内存布局）
逐通道操作（BatchNorm） ✓（channel 连续，合并访问）
通道混洗（PixelShuffle） ✓

Ampere（A100）及之后的架构，NHWC + Tensor Core 的组合是大多数深度学习推理任务的最优选择。

总结：并行算法的 CUDA 实现要点

算法 核心技术 关键注意事项
Reduce 树形归约 + Warp shuffle 两阶段归约处理大数组
Warp Reduce __shfl_down_sync mask 必须精确，不能在分支内不对称调用
Prefix Sum Blelloch 算法 / CG scan 大数组需分块 + 偏移修正
Softmax Max + sum + normalize 数值稳定需减最大值；Online 算法可单遍完成
LayerNorm Block 内 reduce 每行独立，Welford 算法省一遍扫描
Histogram 共享内存局部累积 bin 数决定是否用 Cluster 分布式共享内存
Partition Scan + Scatter 可用 __ballot_sync 加速 Warp 内 partition
Transpose 共享内存 Tile + padding [TILE][TILE+1] 消除 Bank Conflict
NCHW↔NHWC 共享内存 Tile / 向量化 实际工程优先用 cuDNN/PyTorch