卷积的实现与优化

常量缓存

设输入矩阵为，滤波器为。这份代码采用步长为 1、不填充边界的 valid 模式，因此输出尺寸为

每个输出元素的计算过程为

严格来说，代码没有将滤波器沿两个维度翻转，计算的是二维互相关；深度学习框架通常沿用“卷积”的名称。这里也没有 Batch、Channel、Padding、Stride 或 Dilation 维度，只处理一个二维输入和一个二维滤波器。

最直接的并行方式是让每个线程负责一个输出元素。二维 Grid 中，threadIdx.y 对应输出行，threadIdx.x 对应输出列；线程在两个内层循环中遍历完整滤波器，在寄存器变量 acc 中累加点积结果。相邻线程计算相邻输出列，在固定的 kr 和 kc 下读取连续的输入地址，因此输入访问可以形成合并访问。

滤波器具有只读、尺寸较小，并且被所有线程反复访问的特点，适合存放在 Constant Memory。一个 Warp 中的线程以相同顺序执行 kr、kc 循环，在同一时刻读取相同的 c_kernel[kr * kernel_cols + kc]，Constant Cache 可以将一次读取广播给整个 Warp。若同一 Warp 的线程读取不同常量地址，访问可能被拆分甚至串行化；当前访问模式恰好满足广播条件。

#include <cuda_runtime.h>

#define BLOCK_SIZE 16
#define MAX_KERNEL_SIZE 4096

__constant__ float c_kernel[MAX_KERNEL_SIZE];

__global__ void conv2d_kernel(const float *input, const float *kernel,
                              float *output, int input_rows, int input_cols,
                              int kernel_rows, int kernel_cols, int output_rows,
                              int output_cols) {
    int row = blockDim.y * blockIdx.y + threadIdx.y;
    int col = blockDim.x * blockIdx.x + threadIdx.x;
    if (row >= output_rows || col >= output_cols) {
        return;
    }

    float acc = 0.0f;
    for (int kr = 0; kr < kernel_rows; ++kr) {
        for (int kc = 0; kc < kernel_cols; ++kc) {
            float x = input[(row + kr) * input_cols + (col + kc)];
            float k = c_kernel[kr * kernel_cols + kc];
            acc += x * k;
        }
    }

    output[row * output_cols + col] = acc;
}

// input, kernel, output are device pointers
extern "C" void solve(const float *input, const float *kernel, float *output,
                      int input_rows, int input_cols, int kernel_rows,
                      int kernel_cols) {
    int output_rows = input_rows - kernel_rows + 1;
    int output_cols = input_cols - kernel_cols + 1;

    int kernel_size = kernel_rows * kernel_cols;
    cudaMemcpyToSymbol(c_kernel, kernel, kernel_size * sizeof(float), 0,
                       cudaMemcpyDeviceToDevice);

    dim3 threadsPerBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 blocksPerGrid((output_cols + BLOCK_SIZE - 1) / BLOCK_SIZE,
                       (output_rows + BLOCK_SIZE - 1) / BLOCK_SIZE);

    conv2d_kernel<<<blocksPerGrid, threadsPerBlock>>>(
        input, kernel, output, input_rows, input_cols, kernel_rows, kernel_cols,
        output_rows, output_cols);
}

solve 先根据输入和滤波器尺寸计算输出范围，再通过 cudaMemcpyToSymbol 将滤波器从 Device Memory 复制到常量数组 c_kernel。Kernel 的 kernel 形参没有参与计算，实际读取的滤波器数据全部来自 c_kernel。MAX_KERNEL_SIZE=4096 对应 4096 个 float，共 16 KiB；这一版本默认输出尺寸为正，并且滤波器元素数量不超过该上限。

每个输出元素需要执行次乘加，约为 FLOPs。忽略缓存复用和 Constant Memory 的滤波器流量，只计算输入读取与最终输出写回时，单线程需要读取个 float 并写入一个 float，算术强度近似为

当滤波器增大时，该值趋近于，因此仅把滤波器放入 Constant Memory 后，主要带宽开销仍来自输入。相邻输出的感受野高度重叠，同一个输入元素会被多个线程重复读取；L1、L2 Cache 可能提供隐式复用，但代码本身还没有显式保留这些输入数据。

平铺并行卷积

平铺版本进一步利用相邻输出窗口之间的输入重叠。一个边长为的线程 Block 计算个输出元素。为了覆盖这些输出的全部感受野，需要加载的输入 Tile 尺寸不是，而是

其中超出输出 Tile 的部分称为 halo。代码令，将包含 halo 的完整输入 Tile 放入 Shared Memory，然后由 Block 内所有线程反复复用。

#include <cuda_runtime.h>

#define BLOCK_SIZE 16
#define MAX_KERNEL_SIZE 4096

__constant__ float c_kernel[MAX_KERNEL_SIZE];

__global__ void conv2d_kernel(const float *__restrict__ input,
                              float *__restrict__ output, int input_rows,
                              int input_cols, int kernel_rows, int kernel_cols,
                              int output_rows, int output_cols) {
    extern __shared__ float tile[];
    const int out_row = blockIdx.y * blockDim.y + threadIdx.y;
    const int out_col = blockIdx.x * blockDim.x + threadIdx.x;
    const int ty = threadIdx.y;
    const int tx = threadIdx.x;

    const int tile_rows = BLOCK_SIZE + kernel_rows - 1;
    const int tile_cols = BLOCK_SIZE + kernel_cols - 1;
    const int tile_size = tile_rows * tile_cols;

    // Cooperative loading of input tile + halo.
    const int num_threads = BLOCK_SIZE * BLOCK_SIZE;
    const int tid = ty * BLOCK_SIZE + tx;
    for (int i = tid; i < tile_size; i += num_threads) {
        const int tile_row = i / tile_cols;
        const int tile_col = i % tile_cols;

        const int in_row = blockIdx.y * BLOCK_SIZE + tile_row;
        const int in_col = blockIdx.x * BLOCK_SIZE + tile_col;

        tile[i] = (in_row < input_rows && in_col < input_cols)
                      ? input[in_row * input_cols + in_col]
                      : 0.0f;
    }
    __syncthreads();

    if (out_row >= output_rows || out_col >= output_cols) {
        return;
    }

    float acc = 0.0f;
    for (int kr = 0; kr < kernel_rows; ++kr) {
        const int tile_offset = (ty + kr) * tile_cols + tx;
        const int kernel_offset = kr * kernel_cols;

        for (int kc = 0; kc < kernel_cols; ++kc) {
            acc += tile[tile_offset + kc] * c_kernel[kernel_offset + kc];
        }
    }

    output[out_row * output_cols + out_col] = acc;
}

// input, kernel, output are device pointers.
extern "C" void solve(const float *input, const float *kernel, float *output,
                      int input_rows, int input_cols, int kernel_rows,
                      int kernel_cols) {
    const int output_rows = input_rows - kernel_rows + 1;
    const int output_cols = input_cols - kernel_cols + 1;
    const int kernel_size = kernel_rows * kernel_cols;

    if (output_rows <= 0 || output_cols <= 0 || kernel_size > MAX_KERNEL_SIZE)
        return;

    cudaMemcpyToSymbol(c_kernel, kernel, kernel_size * sizeof(float), 0,
                       cudaMemcpyDeviceToDevice);

    const dim3 block(BLOCK_SIZE, BLOCK_SIZE);
    const dim3 grid((output_cols + BLOCK_SIZE - 1) / BLOCK_SIZE,
                    (output_rows + BLOCK_SIZE - 1) / BLOCK_SIZE);

    const size_t shared_bytes =
        static_cast<size_t>(BLOCK_SIZE + kernel_rows - 1) *
        (BLOCK_SIZE + kernel_cols - 1) * sizeof(float);

    conv2d_kernel<<<grid, block, shared_bytes>>>(
        input, output, input_rows, input_cols, kernel_rows, kernel_cols,
        output_rows, output_cols);
}

Shared Memory 的大小取决于运行时传入的滤波器尺寸，因此 Kernel 使用 extern __shared__ float tile[] 声明动态共享内存，Host 端通过 shared_bytes 在启动 Kernel 时传入所需字节数。二维线程坐标被展平为 tid，所有线程以 num_threads 为步长遍历 tile_size：当 Tile 元素数大于 256 时，一个线程可能加载多个元素；当元素数较少时，部分线程不执行加载。线性下标再被还原成 tile_row 和 tile_col，用于计算对应的全局输入位置。

最后一个 Block 可能覆盖到输入边界之外，越界位置会被填为 0。__syncthreads() 必须位于输出边界判断之前，因为即使某个线程不对应有效输出，它仍可能负责加载其他有效线程所需的 halo；如果这些线程提前返回，其余线程既可能缺少数据，也可能在同步屏障处永久等待。同步完成后，越界线程才退出，有效线程则从 Shared Memory 中读取自己的局部窗口。

计算阶段中，线程 (ty, tx) 对应 Tile 内左上角为 (ty, tx) 的窗口。tile_offset 定位当前滤波器行在 Shared Memory 中的起点，kernel_offset 定位 Constant Memory 中的对应滤波器行，内层 kc 循环沿连续列完成乘加。__restrict__ 表示 input 与 output 不会指向重叠的存储区域，使编译器可以更积极地优化内存访问。

若不考虑硬件缓存，一个输出块在直接实现中需要读取个输入元素；平铺后只需从 Global Memory 读取个输入元素。Block 内理论输入复用倍数为

例如、时，直接实现需要读取个输入元素，而平铺版本只需加载个，Block 内的 Global Memory 输入读取量约减少到原来的。这个估算忽略了缓存命中、边缘 Tile 和相邻 Block 之间重复加载的 halo，但能够反映 Shared Memory 显式复用的主要收益。

平铺的代价是 Shared Memory 占用随滤波器尺寸增长。每个 Block 需要

当、时仅需 1296 Bytes；滤波器较大时，动态共享内存可能限制单个 SM 上同时驻留的 Block 数量，甚至超过设备允许的每 Block 上限。不同 Block 的 halo 仍会重复从 Global Memory 加载，而且当前实现仍是一线程计算一个输出元素，尚未使用寄存器分块、向量化加载或异步拷贝等进一步优化。

卷积的实现与优化

常量缓存

平铺并行卷积

AI Infra

评论

搜索文章

常量缓存

平铺并行卷积

AI Infra

相关文章

Softmax 算子的实现与优化

GEMM 算子的实现与优化

Triton Attention Kernel

评论