卷积是现代机器学习中的基础操作之一,但它也是一个更广泛工程理念的有用示例:许多看似专用的张量操作可以被重新组织为人们熟知的计算模式,从而更容易实现、优化和扩展。
本文介绍了一种实现二维卷积的实用方法,即把它视为矩阵乘法的结构化形式。目标不是推广某个特定产品或基准测试结果,而是为构建 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; // 结果累加器,长度为 _N_STRIDE
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 推理和张量内核的工程师来说,这才是更重要的启示。良好的性能并不是从繁重的底层调优开始的,而是从选择正确的问题理解方式开始的。
来源: 海外媒体
登录后即可参与技术讨论与同行交流
免费注册