如何把 Conv2D 转化为实用的 GPU 计算模式

卷积是现代机器学习中的基础运算之一，但它也体现了一个更广泛的工程理念：许多看似专用的张量运算可以重组为易于理解、更易实现、优化和扩展的计算模式。

本文介绍一种将 2D 卷积视为结构化矩阵乘法来实现的实用方法。本文的目的不是推广某个特定产品或基准测试结果，而是为构建 GPU 计算内核的工程师展示一种可复用的优化思路。

核心思想：1×1 卷积本质上就是矩阵乘法，而更大的 K×K 卷积可以表示为一系列矩阵式的累加，只需注意缓存行为和循环顺序。

为什么这很重要

Conv2D 广泛存在于推理流水线、图像模型和特征提取阶段。然而，实现的挑战往往不在于数学定义，而在于执行策略：数据如何布局、工作如何分组、内存流量如何控制。

一种有效的方法是避免把卷积当作完全独特的原语来处理，而是将其解读为与矩阵乘法密切相关的计算模式。这样可以降低实现复杂度，并为优化开辟更清晰的路径。

• 通过复用熟悉的执行模型，降低了工程投入。
• 降低了维护成本，因为优化可以在相关内核之间迁移。
• 提高了未来软硬件改进同时惠及矩阵乘法和卷积路径的可能性。

从张量形状说起

标准的 Conv2D 操作可以描述为：在输入张量上应用一个学习到的滤波器张量，以产生输出张量。

C = A * B

其中：
A 是大小为 N × H × W × C_in 的输入张量
B 是大小为 H_filter × W_filter × C_in × C_out 的权重张量
C 是大小为 N × H_out × W_out × C_out 的输出张量

在实践中，最重要的化简之一来自于以匹配计算引擎的方式来审视数据，而不是依据原始的数学记号。

最简单的情形：为什么 1×1 卷积很特殊

1×1 卷积是理解矩阵乘法类比的最清晰的切入点。在每个空间位置上，输入贡献一个长度为 C_in 的特征向量。权重将该向量映射为长度为 C_out 的输出向量。

如果输入张量已经以 NHWC 顺序存储，那么每个空间特征向量天然就以矩阵行的形式排列，无需显式的数据重排。

实用要点：对于 NHWC 输入，1×1 卷积可以直接视为从 C_in 到 C_out 的矩阵乘法，以 NHWC 顺序产生输出，无需转置步骤。

这一点很重要，因为不必要的数据搬移往往是最容易损失性能的因素之一。当原始布局已经匹配目标访问模式时，实现就可以保持简单高效。

从 1×1 到完整的 K×K 卷积

一旦滤波器超过 1×1，每个输出元素就依赖于多个空间位置。例如，3×3 卷积会组合九个独立的空间贡献。从概念上讲，这变成了一组需要累加在一起的矩阵式乘积。

这意味着问题不再只是“做矩阵乘法”，而是“在最小化内存浪费的同时进行反复的类矩阵累加”。

对于 3×3 的滤波器，每个输出位置有九个 C_in × C_out 的贡献项。朴素的实现可以获得较好的算术利用率，但可能会通过缓存层级反复清空并重新加载数据。这正是执行顺序变得关键的地方。

循环顺序是一种优化决策

卷积优化中最有用的经验之一是：等价的数学运算可能产生截然不同的内存行为。

一种直接的顺序可能是把每个滤波器位置当作单独的矩阵乘法来处理并累加结果。虽然这很简单，但可能导致同一输入区域被反复取入 L1 缓存，甚至 L2 缓存。

重排循环可以改善局部性、减少重复取数据的流量，并提高缓存的有效复用。换句话说，算术运算保持不变，但内存系统面对的模式友好得多。

优化原则：当一个内核在计算上已经正确有效时，最有价值的改动之一往往是在尝试低级指令调优之前，先重构循环嵌套以改善数据复用。

一个代表性的内核模式

下面的示例展示了卷积循环嵌套的核心结构。细节因实现而异，但模式大体上可以辨认：加载权重、计算输入坐标、读取向量化的输入值、同步本地内存，并累加部分乘积。

1	VEC_TYPE out_vec = 0.0;  // accumulator for results, with _N_STRIDE length
2	while ((idx_x += 2) < K_STEPS * 2) {
3	  for (int fh = 0; fh < F_H; ++fh) {
4	    for (int fw = 0; fw < F_W; ++fw) { 
5	     int idx_in2 = 
6	          IDX_FROM_4D(fh, fw, F_W, idx_x * 4 + load_in2_k_offset, CI, load_in2_n_offset, CO); 
7	     ((__local float *)localBuf)[local_idx] = (load_in2_n_offset >= CO) ? 0 : in2[idx_in2];
8	 
9	      const int hi = h * STRIDE_H + fh * DILATION_H - PAD_TOP;
10	      const int wi = w * STRIDE_W + fw * DILATION_W - PAD_LEFT;
11	 
12	      idx_y = IDX_FROM_3D(b, hi, IN_H, wi, IN_W);
13	      input1_8.lo = (float4)read_imagef(in1, samp, (int2){idx_x, idx_y});
14	      input1_8.hi = (float4)read_imagef(in1, samp, (int2){idx_x + 1, idx_y});
15	      barrier(CLK_LOCAL_MEM_FENCE);
16	      batch8_F32_MM(input1_8, (__local float16 *)localBuf, (float2 *)&out_vec,
17	                    (float2 *)&out_vec);
18	      barrier(CLK_LOCAL_MEM_FENCE);
19	    }
20	  }
21	}

这类内核值得研究，因为它揭示了性能通常所在之处：

• 用向量化读取取代标量流量
• 用本地内存暂存可复用的权重
• 对滤波器位置进行结构化累加
• 围绕共享数据的显式同步
• 跨多个输出的线程级并行

填充、步长与边界处理

真实的卷积内核不仅由滤波器数学决定，还受填充、步长和膨胀等参数的影响。这些参数既影响正确性，也影响性能。

填充策略

零填充是最常见的方法，也是源实现所采用的方法。一种典型的边界安全模式是在发出读取之前，先检查计算出的输入坐标是否落在有效的图像边界内。

1	float8 input1_8 = 0;
2	if ((hi >= 0) && (hi < IN_H) && (wi >= 0) && (wi < IN_W)) 
3	{
4	  input1_8[i] = { ... }
5	}

这是正确的，也易于推理，但热循环内部的分支可能代价高昂。一种替代方案是把边界处理移到最内层循环之外，将越界读取重定向到一条用零填充的哑行或缓冲区。这用分支开销换取了设置逻辑，通常值得一试。

重要：边界检查不仅仅是正确性问题。它们在循环嵌套中的位置可能对吞吐量产生可测量的影响。

步长与滤波器尺寸

滤波器尺寸和步长还会以不太明显的方式影响性能。

• 更大的滤波器 通常会提高输入复用率，从而改善缓存效率。
• 更大的步长 会减少相邻输出之间的复用，可能使内存行为不那么有利。

这些影响往往是间接的而非绝对的。更大的滤波器不一定会让内核更快，更大的步长也不一定会让它变慢。但两者都会影响算术工作量与内存流量之间的平衡，而这正是许多 Conv2D 内核成败的关键所在。

工程师能从这种方法中学到什么

这种 Conv2D 实现中最具可迁移价值的不是某个单一的微优化，而是其背后的方法论。

• 从一个清晰的张量解释出发。
• 判断该操作何时能自然地映射为矩阵乘法。
• 尽可能保留数据布局，以避免不必要的数据移动。
• 把早期优化精力集中在循环顺序和缓存复用上。
• 把填充、步长和边界逻辑视为与性能相关的设计选择。

这种方法往往能产出更易于理解、更易于演进的实现。它还为未来的加速打下更坚实的基础，因为许多硬件平台已经在稠密线性代数模式上投入了大量资源。

常见陷阱

在把 Conv2D 转化为类矩阵的 GPU 内核时，以下错误反复出现：

• 以为算术开销是主要瓶颈，而实际上内存流量才是主导因素
• 在未检查源布局是否已经适用之前，就过早引入显式数据变换
• 把边界分支留在最热的循环里，却不衡量其影响
• 在验证缓存行为和访问顺序之前就优化指令细节
• 忽视步长对数据复用模式的改变

注意：数学上优雅的实现并不一定是高性能的实现。等价的公式在真实硬件上的表现可能截然不同。

一个实用的思维模型

如果你正在设计或评审 GPU 卷积内核，从三个层次来思考会很有帮助：

• 数学层：哪些值需要被相乘和累加？
• 布局层：这些值在内存中如何排列，该布局能否被直接复用？
• 执行层：工作应以什么顺序运行，才能高效利用缓存、本地内存和向量单元？

优秀的内核通常来自三个层次的对齐，而不是只对其中一个过度优化。

最后的要点

Conv2D 常被当作一种专用的深度学习原语来介绍，但从实现的角度看，它往往可以被理解为结构化的矩阵乘法，再加上对空间邻域的细致累加。

这一视角非常有力，因为它简化了设计决策。它鼓励复用经过验证的计算模式，减少不必要的数据移动，并把优化精力放在通常最重要的地方：局部性、循环顺序和带宽效率。

对于从事 GPU 推理和张量内核工作的工程师来说，这是更重要的启示。良好的性能并非始于高深的底层调优，而是始于为问题选择正确的解释。