Prefix Sum

agicy2026/6/6大约 1 分钟

Prefix Sum

原始题目：LeetGPU - Prefix Sum

题目描述

编写一个 GPU 程序，计算 32 位浮点数数组的前缀和（累积和）。对于输入数组 $[a_0, a_1, a_2, \dots, a_{N-1}]$ ，前缀和为：

output[i] = \sum_{j=0}^{i} input[j]

即 $[a_0, a_0+a_1, a_0+a_1+a_2, \dots]$ 。

实现要求

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

示例

示例 1

Input:  [1.0, 2.0, 3.0, 4.0]
Output: [1.0, 3.0, 6.0, 10.0]

示例 2

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

约束条件

$1 \le N \le 100{,}000{,}000$ 。
$-1000.0 \le input[i] \le 1000.0$ 。
输出数组中的最大值适合 32 位浮点数。
性能测试在 $N = 250{,}000$ 的规模下进行。

前缀和（Scan）是并行计算中的核心原语，广泛应用于排序（基数排序）、流压缩、稀疏矩阵等算法。经典方法有 Hillis-Steele（简单但效率低， $O(n \log n)$ ）和 Blelloch（work-efficient，两趟：up-sweep + down-sweep）。实际实现中每个 block 先做本地 scan，再用跨 block 的全局 scan 连接。关键优化是避免 bank conflict 和使用 warp-level shuffle。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void prefix_sum_kernel(const float* input, float* output, int N) {
    extern __shared__ float temp[];
    int tid = threadIdx.x, i = blockIdx.x * blockDim.x + tid;
    temp[tid] = (i < N) ? input[i] : 0.0f; __syncthreads();
    for (int s = 1; s < blockDim.x; s *= 2) {
        float t = (tid >= s) ? temp[tid - s] : 0.0f;
        __syncthreads(); temp[tid] += t; __syncthreads();
    }
    if (i < N) output[i] = temp[tid];
}
extern "C" void solve(const float* input, float* output, int N) {
    prefix_sum_kernel<<<(N+255)/256, 256, 256*sizeof(float)>>>(input, output, N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def scan_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.store(output_ptr + idx, tl.associative_scan(x, 0, tl.sum), mask=mask)