Subarray Sum

原始题目：LeetGPU - Subarray Sum

题目描述

计算 32 位整数数组的子数组和。给定长度为 $N$ 的输入数组和两个索引 $S$、$E$（0-based，闭区间），计算 $input[S..E]$ 的和。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

示例

Input:  input = [1, 2, 1, 3, 4], S = 1, E = 3
Output: 6 (2 + 1 + 3)

约束条件

• $1 \le N \le 100{,}000{,}000$，$0 \le S \le E < N$。

解题思路

子数组和 = 对 $[S, E]$ 范围内元素的规约求和。可以先计算前缀和 $prefix[E+1] - prefix[S]$（预处理），也可以直接对子范围做并行规约。当子范围远小于 $N$ 时，直接规约更高效。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void subsum(const int* input, int* output, int S, int E) {
    __shared__ int bc; if(threadIdx.x==0)bc=0; __syncthreads();
    int n=E-S+1;
    for(int i=blockIdx.x*blockDim.x+threadIdx.x; i<n; i+=gridDim.x*blockDim.x)
        atomicAdd(&bc,input[S+i]);
    __syncthreads(); if(threadIdx.x==0)atomicAdd(output,bc);
}
extern "C" void solve(const int* input, int* output, int N, int S, int E) {
    int n=E-S+1; subsum<<<min((n+255)/256,1024),256>>>(input,output,S,E);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def subsum(input_ptr,output_ptr,S,E,BLOCK:tl.constexpr):
    n=E-S+1; idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=idx<n
    x=tl.load(input_ptr+S+idx,mask=mask,other=0)
    tl.atomic_add(output_ptr, tl.sum(x,axis=0))