2D Max Pooling

原始题目：LeetGPU - 2D Max Pooling

题目描述

编写一个 GPU 程序，实现图像/特征图的二维最大池化下采样。给定输入张量（形状 $[N, C, H, W]$）和池化参数（kernel_size、stride、padding），计算最大池化输出。

对于每个池化窗口，输出为该窗口内所有元素的最大值。输出尺寸为：

$$\begin{aligned} H_{out} &= \left\lfloor\frac{H + 2 \cdot padding - kernel\_size}{stride}\right\rfloor + 1 \\ W_{out} &= \left\lfloor\frac{W + 2 \cdot padding - kernel\_size}{stride}\right\rfloor + 1 \end{aligned}$$

实现要求

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

示例

示例 1

Input:  input (1×1×3×3): [[[1,2,3],[4,5,6],[7,8,9]]]
        kernel_size=2, stride=1, padding=0
Output: (1×1×2×2): [[[5,6],[8,9]]]

示例 2

Input:  input (1×1×5×5): [[[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20],[21,22,23,24,25]]]
        kernel_size=3, stride=1, padding=1
Output: (1×1×5×5): [[[7,8,9,10,10],[12,13,14,15,15],[17,18,19,20,20],[22,23,24,25,25],[22,23,24,25,25]]]

约束条件

• $1 \le N \le 100$，$1 \le C \le 512$，$1 \le H, W \le 1{,}024$。
• $1 \le kernel\_size \le 16$，$1 \le stride \le 16$，$0 \le padding \le 16$。
• 输入和输出张量使用 float32。
• 性能测试在 $N = 4,\ kernel\_size = 3,\ stride = 2$ 的规模下进行。

解题思路

最大池化与卷积类似，但用 max 操作替代了乘加操作。每个输出元素对应一个池化窗口，窗口内的所有输入都需要被读取。共享内存分块策略同样适用，但 max 操作比 FMA 更轻量，因此这是一个更纯粹的内存带宽受限内核。注意边界处理——padding 区域应视为 $-\infty$（对 max 操作等效于不存在）。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void maxpool_kernel(const float* input, float* output, int H, int W,
    int K, int stride, int padding) {
    int ow=blockIdx.x*blockDim.x+threadIdx.x, oh=blockIdx.y*blockDim.y+threadIdx.y;
    int OH=(H+2*padding-K)/stride+1, OW=(W+2*padding-K)/stride+1;
    if(oh<OH && ow<OW) {
        float mx=-INFINITY;
        for(int ky=0;ky<K;ky++)
            for(int kx=0;kx<K;kx++) {
                int ih=oh*stride+ky-padding, iw=ow*stride+kx-padding;
                if(ih>=0&&ih<H&&iw>=0&&iw<W) mx=fmaxf(mx,input[ih*W+iw]);
            }
        output[oh*OW+ow]=mx;
    }
}
extern "C" void solve(const float* input, float* output, int H, int W,
    int K, int stride, int padding) {
    int OH=(H+2*padding-K)/stride+1, OW=(W+2*padding-K)/stride+1;
    dim3 t(16,16), b((OW+15)/16, (OH+15)/16);
    maxpool_kernel<<<b,t>>>(input, output, H, W, K, stride, padding);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def maxpool_kernel(input_ptr, output_ptr, H, W, K, stride, padding, BLOCK: tl.constexpr):
    ow=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); oh=tl.program_id(1)*BLOCK+tl.arange(0,BLOCK)
    OH=(H+2*padding-K)//stride+1; OW=(W+2*padding-K)//stride+1
    mask=(oh[:,None]<OH)&(ow[None,:]<OW)
    mx=tl.full((BLOCK,BLOCK), -float('inf'), tl.float32)
    for ky in range(K):
        for kx in range(K):
            ih=oh[:,None]*stride+ky-padding; iw=ow[None,:]*stride+kx-padding
            vm=(ih>=0)&(ih<H)&(iw>=0)&(iw<W)
            mx=tl.maximum(mx, tl.load(input_ptr+ih*W+iw, mask=vm, other=-float('inf')))
    tl.store(output_ptr+oh[:,None]*OW+ow[None,:], mx, mask=mask)