cuDNN 的初始设计

cuDNN V1.0在2014年的发布,并集成到 Caffe、Paddle 等深度学习框架中。论文 cuDNN: Efficient Primitives for Deep Learning 介绍了 NVIDIA 对于该库的设计和实现。近十年间，NVIDIA 从 Kepler 到 Ada Lovelace 迭代推出了8代架构，cuDNN 也更新到 8.9。硬件上引入了 Tensor Core，软件方面 cuDNN V8 中的 Graph API 相比之前变化较大。然而，作为深度学习领域影响广泛的软件基础设施，一窥 cuDNN 的初始设计仍是有意义的。

cuDNN 类似于 cuBLAS，提供了深度学习原语的高效实现。其卷积实现可以在各种输入大小范围内提供可靠的性能，并利用高度优化的矩阵乘法程序来提供高性能，而无需任何辅助内存。

Library

cuDNN 的主要目标之一是使神经网络框架社区能够平等地从其 API 中受益。相应地，cuDNN 的用户不需要采用任何特定的软件框架，甚至不需要进行数据布局：

• cuDNN 不提供层抽象，而是提供较低级别的计算原语，以简化与现有深度学习框架的集成。因为每个框架都有自己的抽象。
• 大部分 API 专用于对存储在用户控制的缓冲区中的数据执行原语操作的函数。通过维持低级别的 API，该库可以简单地集成到其他框架中。
• cuDNN 支持其所有程序在单精度和双精度浮点运算中的前向和反向传播变体。这些包括卷积、池化和激活函数。该库允许可变数据布局和步幅，以及对输入图像的子部分进行索引。它还包括一组辅助张量变换程序，可以轻松操作4d 张量。

Overview and Handles

cuDNN 公开了一个主机可调用的 C 语言 API，但要求输入和输出数据驻留在 GPU 上，类似于 cuBLAS。

cuDNN 库是线程安全的，其程序可以从不同的主机线程调用。库中提供了一个基于上下文的 API，可以轻松实现多线程和（可选）与 CUDA 流的互操作性。

前向和后向传递的卷积程序使用一个通用的描述符来封装层的属性。Tensor 和 Filter 用不透明的描述符表示，可以灵活地使用张量的每个维度上的任意步长来指定张量布局。 cuDNN句柄、描述符和函数调用的自包含设计以及框架的模块化使集成变得简单。

Spatial Convolutions

卷积神经网络中最重要的计算原语是批处理卷积的一种特殊形式。下表列出了控制此卷积的参数。

Parameter	Meaning
N	Number of images in mini-batch
C	Number of input feature maps
H	Height of input image
W	Width of input image
K	Number of output feature maps
R	Height of filter kernel
S	Width of filter kernel
u	Vertical stride
v	Horizontal stride
pad_h	Height of zero-padding
pad_w	Width of zero-padding

本节描述了这种卷积的前向形式——反向传播所需的其他形式密切相关。卷积有两个输入：

• 输入数据 $D\in {\mathbb{R}}^{NCHW}$；
• 卷积滤波器 $F\in {\mathbb{R}}^{KCRS}$。

输入数据范围为小批量 $N$ 幅图像，$C$ 个输入特征图，每幅图像 $H$ 行和 $W$ 列。滤波器的范围包括 $K$ 个输出特征图、$C$ 个输入特征图、每个滤波器的 $R$ 行和 $S$ 列。

输出也是一个四维张量 $O\in {\mathbb{R}}^{NKPQ}$，其中 $N$ 和 $K$ 如前所定义，$P=f(H,R,u,pad\_h)$, $Q=f(W,S,v,pad\_w)$，这意味着输出图像的高度和宽度取决于图像和滤波器的高度和宽度，以及填充和步幅选择。

步幅参数 $u$ 和 $v$ 允许用户通过仅计算输出像素的子集来减少计算量。填充参数允许用户指定每张图片附加多少行或列的 $0$ 条目。更具体地说，

$$f(H,R,u,pad\_h)=\lceil \frac{H-R+1+2pad\_h}{u}\rceil$$

定义一个访问函数来考虑卷积的步幅、填充和反转：
$$g(p,u,R,r,pad\_h)=p\cdot u+R-r-1-pad\_h$$

然后，前向卷积计算 $O[n,k,p,q]\forall n\in [0,N),\forall k\in [0,K),\forall p\in [0,P),\forall q\in [0,Q)$。为方便起见，将 $D_0$ 定义为 $D$ 的零扩展版本。

$$O[n,k,p,q]=\sum _{c=0}^{C-1}\sum _{r=0}^{R-1}\sum _{s=0}^{S-1}F[k,c,r,s]\cdot D_0[n,c,g(p,u,R,r,pad\_h),g(q,v,S,s,pad\_w)]$$

从上式可以看出，计算卷积涉及一个七层嵌套循环，具有四个独立循环和三个累加循环。有多种实现此计算的方法，本文将在下一节中讨论其中的一些方法。

cuDNN 的卷积程序包含了这些函数的卷积以及互相关变体的实现。这些函数支持用户在输入和输出张量的每个维度上自定义步幅。这很重要，因为不同的框架使用不同的内存布局存储张量。例如，一些框架将特征图交错放置，而另一些则将它们分开。cuDNN 允许用户指定内存布局，这使得集成到现有框架中变得更加简单。对于具有共享参数或有向无环图结构的模型，cuDNN 的程序还有一种模式可以返回原始梯度或将它们累积在缓冲区中。

Implementation

cuDNN 提供的大多数函数都有直接的实现。卷积的实现并不那么明显，因此本文将概述 cuDNN 设计选择背后的动机和推论。有几种方法可以有效地实现卷积：

• Im2col convolution，即 im2col + GEMM；
• FFT convolution；
• Direct convolution；
• Implicit GEMM。

cuDNN 的目标是在不使用辅助内存的情况下，提供尽可能接近矩阵乘法的性能。GPU 内存带宽高，但容量低，因此是一种稀缺资源。在训练深度网络时，理想情况下，GPU 内存中应该装满数据、参数和神经元响应，而不是卷积算法所需的辅助数据结构。几种计算卷积的方法需要大型辅助数据结构，因此我们不考虑将这些方法用于 cuDNN。

Explicit GEMM

一种方法是按照 High Performance Convolutional Neural Networks for Document Processing 将卷积降级为矩阵乘法。这可以通过以下操作来完成：

• 将滤波器张量 $F$ 重塑为维度为 $K\times CRS$ 的矩阵 $F_m$；
• 通过复制原始输入数据收集维度为 $CRS\times NPQ$ 的数据矩阵矩阵 $D_m$。

然后可以使用单个矩阵乘法执行计算，以形成维度为 $K\times NPQ$ 的输出矩阵 $O_m$。下图说明了如何将一个简单的卷积降级为矩阵乘法。

此图中的颜色代表输入特征图，$D$ 和 $F$ 的元素在图中被唯一标记，以显示每个元素如何参与形成 $D_m$ 和 $F_m$：

• 滤波器矩阵 $F_m$ 的维度为 $K\times CRS=2\times 12$。
• 而数据矩阵 $D_m$ 的维度为 $CRS\times NPQ=12\times 4$。请注意，$D$ 的每个元素在 $D_m$ 中最多重复 $RS=4$ 次。
• 输出矩阵 $O_m$ 的维度为 $K\times NPQ=2\times 4$。

将卷积降级到矩阵乘法是有效的，因为矩阵乘法是高度优化的。矩阵乘法速度很快，因为它对每个传输数据字节的浮点运算比率很高。这个比率随着矩阵变大而增加，这意味着矩阵乘法在小矩阵上的效率较低。因此，这种卷积方法在创建大矩阵进行乘法运算时最为有效。如前所述，$D_m$ 和 $F_m$ 的大小取决于卷积参数的乘积，而不是参数本身。这意味着使用这种方法的性能可以非常一致，因为算法不关心其中一个参数是否很小，只要其乘积足够大即可。例如，通常在卷积网络的前几层中，$C$ 很小，但 $R$ 和 $S$ 很大，而在网络的末尾，$C$ 很大，但 $R$ 和 $S$ 很小。然而，对于所有层来说，乘积 $CRS$ 通常都相当大，因此性能可以一直很好。

这种方法的缺点是组成 $D_m$ 需要将输入数据复制多达 $RS$ 次，这可能需要非常大的临时分配。为了解决这个问题，实现有时会逐个具体化 $D_m$，例如，为小批量的每个元素迭代调用矩阵乘法。然而，这限制了实现中的并行性，并可能导致矩阵乘法太小而无法有效利用 GPU 的情况。这种方法还降低了卷积的计算强度，因为除了读取 $D$ 本身之外，还必须写入和读取 $D_m$，作为一种更直接的方法，需要明显更多的内存流量。

因此，我们选择不直接使用此实现，尽管正如我们将解释的那样，cuDNN 的实现是相关的。

Tensor Core 的卷积采用 NHWC 的布局，输出矩阵 $O_m$ 的维度为 $NPQ\times K$。 不需要进行重排。详细可参考 CUTLASS Convolution。

FFT Convolution

另一种方法是使用快速傅立叶变换来计算卷积。FFT 可以显著降低卷积的工作复杂度，巧妙的工程化可以有效地用于深度神经网络 [Fast Training of Convolutional Networks through FFTs]。

然而，基于 FFT 的方法使用了大量的临时内存，因为必须将滤波器填充为与输入相同的大小。这在滤波器相对于图像较小的情况下尤其昂贵，通常发生在卷积网络的前几层。

此外，当步幅参数 $u$ 和 $v$ 大于 $1$ 时，基于 FFT 的方法无法高效执行。这在许多最先进的网络中很常见，例如 [OverFeat] 和 InceptionV1 中的初期层。步幅通过仅计算输出的稀疏子集，将卷积的计算量降低到原来的 $1/uv$。然而，FFT 算法的性质使得计算修整后的 FFT 是一项不寻常的任务，并且通常比计算密集 FFT 后跟一个额外的下采样步骤慢。

由于这些缺点，我们选择放弃 FFT 方法，尽管我们同意它在某些情况下很有用。

Direct Convolution

另一种常用的方法是直接计算卷积。这可能非常有效，但需要大量专门的实现来处理隐含在卷积的 11 维参数空间中的许多极端情况。采用这种方法的实现往往对参数空间中某些部分的卷积进行了很好的优化，但对其他部分则表现不佳。例如，cuda-convnet2 在批量较大时表现良好，但一旦批量降至64或以下时表现不佳。优化和维护所有这些特化是一项艰巨的任务。

由于我们设想该库将被维护一段时间，并被移植到尚未构思的未来架构中，因此我们寻找更简单的东西，使其在参数空间中表现得更稳健，更容易移植到新的架构中。

Implicit GEMM

NVIDIA 提供了一个矩阵乘法程序，可在 GPU 上实现了相当高比重的浮点吞吐量。该程序的算法类似于 Fast implementation of DGEMM on Fermi GPU 中描述的算法。

cudnn 使用 double buffering 技术：

• 依次读取输入矩阵 $A$ 和 $B$ 的固定大小子矩阵到片上存储器，用于计算输出矩阵 $C$ 的子矩阵；
• 在计算 $A$ 和 $B$ 分块的同时，将 $A$ 和 $B$ 的下一个分块从片外内存中提取到片上缓存和其他内存中。

这种技术隐藏了与数据传输相关的内存延迟，使得矩阵乘法计算仅受执行算术所需时间的限制。

如前所述，卷积可以降级到矩阵乘法上。这种方法提供了实现的简单性以及在整个参数空间中性能的一致性，尽管在内存中实物化降级的矩阵可能代价高昂。

cuDNN 的解决方案遵循这种方法，但我们通过仅在片上内存中延迟地实物化 $D_m$，而不是在调用矩阵乘法程序之前在片外内存中实现它，来避免在内存中实现降级矩阵的问题。

由于矩阵乘法程序所需的分块与卷积的任何参数无关，因此 $D_m$ 的分块边界与卷积问题之间的映射是非平凡的。因此，我们的方法需要计算此映射并使用它来将 $A$ 和 $B$ 的正确元素加载到片上内存中。随着计算的进行，这种情况会动态发生，这使得我们的卷积算法能够利用优化的基础设施进行矩阵乘法。

与矩阵乘法相比，我们需要额外的索引算法，但充分利用矩阵乘法的计算引擎来执行工作。计算完成后，我们执行所需的张量转置，将结果存储在用户期望的数据布局中。计算额外的索引需要通过启动时间常数除数来重复计算整数除法和模运算。我们利用 Hacker’s delight 中提出的整数除法和模数算法，将这些代价高昂的运算转换为整数乘法和移位，从而减少 Implicit GEMM 方法所需的索引开销。

参考资料：

• cuDNN: Efficient Primitives for Deep Learning 论文阅读
• Optimizing Convolutional Layers
• 2.5 机器码生成
• Go 编译器介绍
• PyTorch 2.0 发布：除了编译，还是编译！
• 编译原理相关
• MLIR：摩尔定律终结的编译器基础结构 论文解读
• 指令集架构、机器码与 Go 语言
• 编译原理：中间代码IR
• 深入浅出GPU优化系列：GEMM优化（一）
• CUDNN LIBRARY User Guide
• 从AI系统角度回顾GPU架构变迁–从Fermi到Ampere(V1.2)
• Accelerate Machine Learning with the cuDNN Deep Neural Network Library
• CUTLASS Convolution
• OpenPPL 中的卷积优化技巧