1D Convolution

原始题目：LeetGPU - 1D Convolution

题目描述

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

输入包含两个数组：

• input：32 位浮点数的一维数组。
• kernel：32 位浮点数的一维数组，表示卷积核。

输出数组的大小为 $input\_size - kernel\_size + 1$。

卷积操作的数学定义为：

$$output[i] = \sum_{j=0}^{kernel\_size - 1} input[i + j] \cdot kernel[j]$$

其中 $i$ 的取值范围为 $0$ 到 $input\_size - kernel\_size$。

实现要求

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

示例

示例 1

Input:  input  = [1, 2, 3, 4, 5]
        kernel = [1, 0, -1]
Output: [-2, -2, -2]

示例 2

Input:  input  = [2, 4, 6, 8]
        kernel = [0.5, 0.2]
Output: [1.8, 3.2, 4.6]

约束条件

• $1 \le input\_size \le 1{,}500{,}000$。
• $1 \le kernel\_size \le 2{,}047$。
• $kernel\_size \le input\_size$。
• 性能测试在 $input\_size = 1{,}500{,}000,\ kernel\_size = 2{,}047$ 的规模下进行。

解题思路

一维卷积的每个输出元素需要读取 $kernel\_size$ 个输入元素。当 $kernel\_size$ 较大时，每个输入元素会被多个输出共享。可以使用共享内存来缓存输入元素，减少全局内存访问。对于较大的 $kernel\_size$，可以使用常量内存（Constant Memory）存储卷积核以利用广播机制。

代码实现

CUDA

#include <cuda_runtime.h>

__global__ void conv1d_kernel(const float* input, const float* kernel, float* output,
                               int input_size, int kernel_size) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int out_size = input_size - kernel_size + 1;
    if (i < out_size) {
        float sum = 0.0f;
        for (int k = 0; k < kernel_size; k++) {
            sum += input[i + k] * kernel[k];
        }
        output[i] = sum;
    }
}

extern "C" void solve(const float* input, const float* kernel, float* output,
                       int input_size, int kernel_size) {
    int out_size = input_size - kernel_size + 1;
    int threadsPerBlock = 256;
    int blocksPerGrid = (out_size + threadsPerBlock - 1) / threadsPerBlock;
    conv1d_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, kernel, output, input_size, kernel_size);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl

@triton.jit
def conv1d_kernel(input_ptr, kernel_ptr, output_ptr,
                  input_size: tl.constexpr, kernel_size: tl.constexpr,
                  BLOCK_SIZE: tl.constexpr):
    i = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    out_size = input_size - kernel_size + 1
    mask = i < out_size
    acc = tl.zeros((BLOCK_SIZE,), tl.float32)
    for k in range(kernel_size):
        acc += tl.load(input_ptr + i + k, mask=mask) * tl.load(kernel_ptr + k)
    tl.store(output_ptr + i, acc, mask=mask)