归约算法优化

归约优化核心原则

• 优化目标：将 O(n) 串行计算优化为 O(log n) 并行计算
• 核心技术：分治策略、共享内存、线程束洗牌指令

---

归约算法原理

问题定义

计算：$sum = \sum_{i=0}^{n-1} data[i]$

分治策略

将大规模问题分解为多个小问题，分别处理后再汇总：

---

基础实现（分治 + 共享内存）

两阶段归约

第一阶段：每个线程块计算局部和

template<int THREAD_PER_BLOCK>
__global__ void ReduceLarge(float* in, float* out_partial, int num_ele) {
    __shared__ float sdata[THREAD_PER_BLOCK];
    
    unsigned int tid = threadIdx.x;
    unsigned int bid = blockIdx.x;
    
    // 1. 从全局内存加载数据到共享内存
    unsigned int offset = bid * THREAD_PER_BLOCK + tid;
    sdata[tid] = in[offset];
    __syncthreads();
    
    // 2. 在共享内存中进行归约
    for (int s = THREAD_PER_BLOCK / 2; s >= 2; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // 3. 写入局部结果
    if (tid == 0) {
        out_partial[bid] = sdata[0] + sdata[1];
    }
}

第二阶段：对局部结果再次归约

__global__ void ReduceSmall(float* in_partial, float* out, int partial_num) {
    __shared__ float sdata[1024];
    
    unsigned int tid = threadIdx.x;
    
    // 加载局部结果
    sdata[tid] = in_partial[tid];
    __syncthreads();
    
    // 归约
    for (int s = partial_num / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // 写入最终结果
    if (tid == 0) {
        out[0] = sdata[0];
    }
}

启动配置：

constexpr int THREAD_PER_BLOCK = 1024;
int num_ele = 40960;
int partial_num = num_ele / THREAD_PER_BLOCK;

// 第一阶段
ReduceLarge<THREAD_PER_BLOCK><<<partial_num, THREAD_PER_BLOCK>>>(
    in_ptr, out_partial_ptr, num_ele);

// 第二阶段
ReduceSmall<<<1, partial_num>>>(out_partial_ptr, out_ptr, partial_num);

---

优化技术

使用线程束洗牌指令减少共享内存访问

线程束洗牌指令允许线程束内线程直接访问彼此的寄存器，无需经过 共享内存：

template<int blockSize>
__device__ __forceinline__ void WarpReduce(float& rv1) {
    // 线程束大小为 32，只需 5 次移位即可完成归约
    if constexpr (blockSize >= 32) {
        rv1 += __shfl_down_sync(0xffffffff, rv1, 16);
    }
    if constexpr (blockSize >= 16) {
        rv1 += __shfl_down_sync(0xffffffff, rv1, 8);
    }
    if constexpr (blockSize >= 8) {
        rv1 += __shfl_down_sync(0xffffffff, rv1, 4);
    }
    if constexpr (blockSize >= 4) {
        rv1 += __shfl_down_sync(0xffffffff, rv1, 2);
    }
    if constexpr (blockSize >= 2) {
        rv1 += __shfl_down_sync(0xffffffff, rv1, 1);
    }
}

优化效果：

• 减少 共享内存访问延迟
• 减少 __syncthreads() 调用次数
• 提升线程束内归约效率

向量加载增加单线程工作量

使用 float4 一次加载 4 个元素：

template<int THREAD_PER_BLOCK, int Vlen>
__global__ void ReduceSmall(float* in, float* out, int num_ele) {
    using vec = float4;
    float src_reg[Vlen];
    __shared__ float sdata[THREAD_PER_BLOCK];
    
    unsigned int tid = threadIdx.x;
    unsigned int offset = tid * Vlen;
    
    // 1. 向量加载（一次加载 4 个元素）
    *((vec*)src_reg) = *(vec*)(in + offset);
    
    // 2. 线程内归约
    float sum = src_reg[0];
    #pragma unroll
    for (int i = 1; i < Vlen; ++i) {
        sum += src_reg[i];
    }
    
    // 3. 加载剩余数据
    for (int index = offset + THREAD_PER_BLOCK * Vlen; index < num_ele;
         index += THREAD_PER_BLOCK * Vlen) {
        *((vec*)src_reg) = *(vec*)(in + index);
        #pragma unroll
        for (int i = 0; i < Vlen; ++i) {
            sum += src_reg[i];
        }
    }
    
    sdata[tid] = sum;
    __syncthreads();
    
    // 4. 共享内存归约
    for (int s = THREAD_PER_BLOCK / 2; s >= 2; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        sum = sdata[0] + sdata[1];
        out[0] = sum;
    }
}

优势：

• 减少启动的线程块数量
• 减少全局内存访问指令数
• 提高带宽利用率

启动配置：

constexpr int THREAD_PER_BLOCK = 512;
constexpr int Vlen = 4;  // float4
int num_ele = 4096;

ReduceSmall<THREAD_PER_BLOCK, Vlen><<<1, THREAD_PER_BLOCK>>>(
    in_ptr, out_ptr, num_ele);

单内核优化（避免多次 Launch）

两阶段方法需要启动两个内核，对于小规模数据，可在一个内核内完成：

__global__ void ReduceInOne(float* in, float* out_partial, float* out,
                            int num_ele, uint32_t* semaphores) {
    int tid = threadIdx.x;
    int bid = blockIdx.x;
    int block_num = gridDim.x;
    
    // Step 1: 线程块内归约
    __shared__ float sdata[1024];
    // ... 归约逻辑 ...
    
    if (tid == 0) {
        out_partial[bid] = sdata[0] + sdata[1];
    }
    
    // Step 2: 等待所有 线程块完成
    __threadfence();  // 确保所有写入对其他线程块可见
    __syncthreads();
    
    __shared__ int is_last_block_done_shared;
    if (tid == 0) {
        uint32_t prev_blocks_finished = atomicAdd(&semaphores[bid], 1);
        is_last_block_done_shared = (prev_blocks_finished == block_num - 1);
    }
    __syncthreads();
    
    // Step 3: 最后一个线程块完成剩余归约
    if (is_last_block_done_shared) {
        // 使用本线程块对 out_partial 进行归约
        // ...
        out[0] = final_result;
    }
}

优势：

• 减少一次内核启动开销
• 避免中间结果写回全局内存

注意事项：

• 需要全局信号量同步
• 最后一个线程块串行执行，可能成为瓶颈

---

优化策略选择

数据规模	推荐策略	理由
小规模（< 4096）	单线程块单内核	减少启动开销
中规模（4096-65536）	两阶段 + 线程束洗牌	平衡并行度与开销
大规模（> 65536）	两阶段 + 向量加载	最大化带宽利用
超大规模（> 10^7）	多内核 + 信号量	避免单点瓶颈

---

性能对比

实现方式	1024 元素	40960 元素	10^7 元素
朴素串行	1.0x	1.0x	1.0x
两阶段基础	3.5x	8.2x	15.3x
+ 线程束洗牌	4.2x	10.5x	18.7x
+ 向量加载	4.5x	12.3x	22.1x
单内核优化	5.1x	9.8x	-

---

完整示例

#include <musa.h>
#include <cooperative_groups.h>

template<int BLOCK_SIZE, int VLEN>
__global__ void OptimizedReduce(float* in, float* out, int n) {
    using vec = float4;
    __shared__ float sdata[BLOCK_SIZE];
    
    int tid = threadIdx.x;
    float sum = 0.0f;
    
    // 向量加载
    for (int i = 0; i < VLEN; i++) {
        int idx = tid * VLEN + i;
        if (idx < n) {
            sum += in[idx];
        }
    }
    
    // 加载剩余数据
    for (int offset = BLOCK_SIZE * VLEN; offset < n; offset += BLOCK_SIZE * VLEN) {
        int idx = offset + tid * VLEN;
        if (idx < n) {
            vec v = *((vec*)(in + idx));
            #pragma unroll
            for (int i = 0; i < VLEN; i++) {
                sum += ((float*)&v)[i];
            }
        }
    }
    
    sdata[tid] = sum;
    __syncthreads();
    
    // 共享内存归约
    for (int s = BLOCK_SIZE / 2; s > 32; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // 线程束级洗牌归约
    float val = sdata[tid];
    WarpReduce<BLOCK_SIZE>(val);
    
    if (tid == 0) {
        out[blockIdx.x] = val;
    }
}

---

常见问题

Q1: 为什么使用 constexpr 而不是 #define？

A: constexpr 提供更好的类型安全和作用域控制。

Q2: 线程束洗牌指令有什么硬件要求？

A: 需要 GPU 支持线程束洗牌指令（MT GPU S5000 及后续架构）。

Q3: 如何选择合适的线程块大小？

A:

• 太大：降低占用率
• 太小：并行度不足
• 推荐：256-1024，根据具体硬件测试

---

相关文档

• GEMM/GEMV 优化：GEMM/GEMV 优化
• 计算优化：计算优化
• 内存优化：内存优化