Max Subarray Sum

原始题目：LeetGPU - Max Subarray Sum

题目描述

编写一个 GPU 程序，计算长度为 $window\_size$ 的任意连续子数组的最大和。给定长度为 $N$ 的 32 位有符号整数数组 input 和整数 window_size。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 变量中。

示例

示例 1

Input:  input = [1, 2, 4, 2, 3], window_size = 2
Output: 6 (子数组 [4, 2] 或 [2, 4]? 不，[4, 2] 和 = 6)

示例 2

Input:  input = [-1, -4, -2, 1], window_size = 3
Output: -5 (子数组 [-1, -4, -2] 的和... 等等, [-4, -2, 1] = -5)

约束条件

• $1 \le N \le 50{,}000$。
• $-10 \le input[i] \le 10$。
• $1 \le window\_size \le N$。
• 性能测试在 $N = 50{,}000$ 的规模下进行。

解题思路

固定窗口大小的最大子数组和本质上是滑动窗口前缀和问题。可以先计算前缀和 $prefix[i] = \sum_{j=0}^{i-1} input[j]$，则窗口 $[i, i+window\_size)$ 的和为 $prefix[i+window\_size] - prefix[i]$。对每个 $i$ 计算并取最大值。这可以并行化：每个线程处理一个起始位置，独立计算窗口和并比较。也可以使用分块策略，每个 block 处理一段，内部做 max-reduce。由于 $N$ 不大（$\le 50k$），单 block 处理就够了。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void max_subsum(const int* input, int* output, int N, int W) {
    __shared__ int smax[256]; int tid=threadIdx.x; int mx=-2147483648;
    for(int i=blockIdx.x*blockDim.x+tid; i<=N-W; i+=gridDim.x*blockDim.x) {
        int sum=0; for(int j=0;j<W;j++)sum+=input[i+j];
        if(sum>mx)mx=sum;
    }
    smax[tid]=mx; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s&&smax[tid+s]>smax[tid])smax[tid]=smax[tid+s];__syncthreads();}
    if(tid==0)atomicMax(output,smax[0]);
}
extern "C" void solve(const int* input, int* output, int N, int W) {
    max_subsum<<<min(N-W+1,1024),256>>>(input,output,N,W);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def max_subsum(input_ptr,output_ptr,N,W,BLOCK:tl.constexpr):
    i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=i<=N-W
    acc=tl.zeros((BLOCK,),tl.int32)
    for j in range(W): acc+=tl.load(input_ptr+i+j,mask=mask,other=0)
    tl.atomic_max(output_ptr, tl.max(acc,axis=0))