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 操作等效于不存在）。欢迎在 GitHub Discussions 分享你的解法。