Reduction

原始题目：LeetGPU - Reduction

题目描述

编写一个 GPU 程序，对 32 位浮点数数组执行并行规约以计算其总和。程序应接受一个输入数组并输出所有元素之和的单个值。

实现要求

• 只允许使用 GPU 原生功能（不允许使用外部库）。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 变量中。

示例

示例 1

Input:  [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
Output: 36.0

示例 2

Input:  [-2.5, 1.5, -1.0, 2.0]
Output: 0.0

约束条件

• $1 \le N \le 100{,}000{,}000$。
• $-1000.0 \le input[i] \le 1000.0$。
• 最终总和始终适合 32 位浮点数。
• 性能测试在 $N = 4{,}194{,}304$ 的规模下进行。

解题思路

规约是 GPU 并行编程中最基础也最重要的模式之一。朴素做法中每个线程累加一段子数组，最后汇总。更高效的做法是树形规约（Tree Reduction）：在共享内存中通过多轮折半累加，将复杂度从 $O(N)$ 降到 $O(\log N)$ 轮。关键优化包括：避免 bank conflict 的步长选择、使用 warp shuffle 指令（__shfl_down_sync）在 warp 内部完成最后几轮规约、以及使用 atomicAdd 合并多个 block 的部分结果。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void reduce_kernel(const float* input, float* output, int N) {
    __shared__ float sdata[256];
    int tid = threadIdx.x;
    float sum = 0.0f;
    for (int i = blockIdx.x * blockDim.x + tid; i < N; i += gridDim.x * blockDim.x)
        sum += input[i];
    sdata[tid] = sum; __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) atomicAdd(output, sdata[0]);
}
extern "C" void solve(const float* input, float* output, int N) {
    int threads = 256;
    reduce_kernel<<<min((N+255)/256, 1024), threads>>>(input, output, N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl, torch
@triton.jit
def reduce_kernel(input_ptr, output_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    idx = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask, other=0.0)
    tl.atomic_add(output_ptr, tl.sum(x, axis=0))