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 加速。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void conv2d_kernel(const float* input, const float* kernel, float* output,
    int H, int W, int KH, int KW) {
    int oc = blockIdx.x * blockDim.x + threadIdx.x;
    int or_ = blockIdx.y * blockDim.y + threadIdx.y;
    int OH = H - KH + 1, OW = W - KW + 1;
    if (or_ < OH && oc < OW) {
        float sum = 0.0f;
        for (int ky = 0; ky < KH; ky++)
            for (int kx = 0; kx < KW; kx++)
                sum += input[(or_+ky)*W + (oc+kx)] * kernel[ky*KW + kx];
        output[or_*OW + oc] = sum;
    }
}
extern "C" void solve(const float* input, const float* kernel, float* output,
    int H, int W, int KH, int KW) {
    int OH = H-KH+1, OW = W-KW+1;
    dim3 threads(16,16), blocks((OW+15)/16, (OH+15)/16);
    conv2d_kernel<<<blocks, threads>>>(input, kernel, output, H, W, KH, KW);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl

@triton.jit
def conv2d_kernel(input_ptr, kernel_ptr, output_ptr, H, W, KH, KW, BLOCK: tl.constexpr):
    ow = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    oh = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
    OW = W - KW + 1
    OH = H - KH + 1
    mask = (oh[:, None] < OH) & (ow[None, :] < OW)
    acc = tl.zeros((BLOCK, BLOCK), tl.float32)
    for ky in range(KH):
        for kx in range(KW):
            val = tl.load(input_ptr + (oh[:, None] + ky) * W + ow[None, :] + kx, mask=mask)
            kw = tl.load(kernel_ptr + ky * KW + kx)
            acc += val * kw
    tl.store(output_ptr + oh[:, None] * OW + ow[None, :], acc, mask=mask)