3D Subarray Sum

原始题目：LeetGPU - 3D Subarray Sum

题目描述

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

实现要求

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

约束条件

• $1 \le N, M, K \le 100$。

解题思路

二维子数组和的直接推广。可使用三维前缀和（基于容斥原理的 8 项加减公式），或直接对子体积做并行规约。由于三维数组维度较小（$\le 100$），计算量可控。

代码实现

CUDA

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

Triton

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