2D Subarray Sum

原始题目：LeetGPU - 2D Subarray Sum

题目描述

计算 $N \times M$ 二维 32 位整数数组的子矩形和。给定行范围 $[S_{ROW}, E_{ROW}]$ 和列范围 $[S_{COL}, E_{COL}]$（0-based，闭区间），求该子矩形内所有元素之和。

实现要求

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

示例

Input:  input = [[1,2,3],[4,5,6],[7,8,9]], S_ROW=0, E_ROW=1, S_COL=1, E_COL=2
Output: 2+3+5+6 = 16

约束条件

• $1 \le N, M \le 10{,}000$，$0 \le S\_ROW \le E\_ROW < N$，$0 \le S\_COL \le E\_COL < M$。

解题思路

与一维子数组类似，可使用二维前缀和 $prefix[i][j] = \sum_{r<i, c<j} input[r][c]$，子矩形和为 $prefix[E_R+1][E_C+1] - prefix[S_R][E_C+1] - prefix[E_R+1][S_C] + prefix[S_R][S_C]$。或直接对子矩形做并行规约。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void sub2d(const int* input, int* output, int M, int SR, int ER, int SC, int EC) {
    __shared__ int bc; if(threadIdx.x==0)bc=0; __syncthreads();
    int rows=ER-SR+1, cols=EC-SC+1, t=rows*cols;
    for(int i=blockIdx.x*blockDim.x+threadIdx.x; i<t; i+=gridDim.x*blockDim.x)
        atomicAdd(&bc,input[(SR+i/cols)*M + (SC+i%cols)]);
    __syncthreads(); if(threadIdx.x==0)atomicAdd(output,bc);
}
extern "C" void solve(const int* input, int* output, int N, int M, int SR, int ER, int SC, int EC) {
    int t=(ER-SR+1)*(EC-SC+1); sub2d<<<min((t+255)/256,1024),256>>>(input,output,M,SR,ER,SC,EC);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def sub2d(input_ptr,output_ptr,M,SR,ER,SC,EC,BLOCK:tl.constexpr):
    rows=ER-SR+1; cols=EC-SC+1; idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
    mask=idx<rows*cols; r=SR+idx//cols; c=SC+idx%cols
    x=tl.load(input_ptr+r*M+c,mask=mask,other=0)
    tl.atomic_add(output_ptr, tl.sum(x,axis=0))