2D Convolution

原始题目：LeetGPU - 2D Convolution

题目描述

编写一个 GPU 程序，执行二维卷积操作。给定输入矩阵和卷积核（滤波器），使用 "valid" 边界条件（无填充，卷积核仅在其与输入完全重叠的位置上应用）计算卷积输出。

输入包括：

• input：32 位浮点数的二维矩阵，按行优先顺序以一维数组表示。
• kernel：32 位浮点数的二维卷积核，同样按行优先顺序以一维数组表示。

输出矩阵维度为：

$$\begin{aligned} output\_rows &= input\_rows - kernel\_rows + 1 \\ output\_cols &= input\_cols - kernel\_cols + 1 \end{aligned}$$

卷积操作定义为：

$$output[i][j] = \sum_{m=0}^{kernel\_rows - 1}\sum_{n=0}^{kernel\_cols - 1} input[i + m][j + n] \cdot kernel[m][n]$$

实现要求

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

示例

示例 1

Input:  input (3×3): [[1,2,3], [4,5,6], [7,8,9]]
        kernel (2×2): [[0,1], [1,0]]
Output: output (2×2): [[6,8], [12,14]]

示例 2

Input:  input (4×4): [[1,1,1,1], [1,2,3,1], [1,4,5,1], [1,1,1,1]]
        kernel (1×3): [[1,0,1]]
Output: output (4×2): [[2,2], [4,3], [6,5], [2,2]]

约束条件

• $1 \le input\_rows, input\_cols \le 3{,}072$。
• $1 \le kernel\_rows, kernel\_cols \le 31$。
• $kernel\_rows \le input\_rows$，$kernel\_cols \le input\_cols$。
• 性能测试在 $input = 3{,}072 \times 3{,}072$，$kernel = 15 \times 15$ 的规模下进行。

解题思路

二维卷积的每个输出元素需要读取 $kernel\_rows \times kernel\_cols$ 个输入元素，相邻输出之间存在大量数据重用。核心优化是使用共享内存分块：将输入的子区域加载到共享内存中，线程块内的线程协作计算输出子区域。此外，常量内存适合存储小卷积核（$\le 31 \times 31$），利用广播机制减少寄存器压力。对于更大的输入，还可使用 Winograd 或 FFT 加速。欢迎在 GitHub Discussions 分享你的解法。