ACCELERATING CONVOLUTION WITH TENSOR CORES IN CUTLASS

Manish Gupta, April 13, 2021

致谢 (ACKNOWLEDGEMENTS)

CUTLASS GitHub 社区:
- 每月35k次克隆，1.1k星标，以及许多活跃用户。
CUTLASS 团队:
- Andrew Kerr, Haicheng Wu, Manish Gupta, Dustyn Blasig, Duane Merrill, Pradeep Ramani, Vijay Thakkar
贡献者:
- Cris Ceka, Timothy Costa, Maila Farooqui, Markus Hohnerbach, Alan Kaatz, Wei Liu, Piotr Majcher, Dhiraj Reddy Nallapa, Mathew Nicely, Kyrylo Perelygin, Aniket Shivam, Paul Springer, Pawel Tabaszewski, Chinmay Talegaonkar, John Tran, Jin Wang, Yang Xu, Scott Yokin
鸣谢:
- Olivier Giroux, Mostafa Hagog, Bryce Lelbach, Julien Demouth, Joel McCormack, Aartem Belevich, Peter Han, Timmy Liu, Yang Wang, Nich Zhao, Jack Yang, Vicki Wang, Junkai Wu, Ivan Yin, Aditya Alturi, Shang Zhang, Takuma Yamaguchi, Stephen Jones, Luke Durant, Harun Bayraktar

议程 (AGENDA)

概述 (Overview)
- CUTLASS 2.4-2.6 版本及卷积定义
深入探讨隐式 GEMM 卷积 (Deep dive implicit GEMM convolutions)
- 构建连贯且完整的抽象
卷积抽象的优化实现 (Optimized implementation of convolution abstractions)
- 预计算不变量
结尾融合 (Epilogue Fusion)
- 支持的结尾融合模式

概述 (OVERVIEW)

CUTLASS 发展历程

CUTLASS 是一个用于深度学习和线性代数的 CUDA C++ 模板库。下图展示了其从 CUTLASS 1.3 到 2.6 的发展时间线，以及与 CUDA 版本的对应关系。GTC 2021 的重点是 2.4 至 2.6 版本，主要涵盖了隐式 GEMM 卷积、张量规约和结尾融合等功能。

近期更新回顾

CUTLASS 2.4 - 2020年11月
- 隐式 GEMM 2D 卷积:
  - 支持前向传播 (Fprop2D)、反向数据梯度 (Dgrad2D) 和反向权重梯度 (Wgrad2D)。
  - 支持 NHWC 和 NCHWx 布局的 4D 张量。
  - 支持 S4, S8, S32, F16, BF16, TF32, F32, complex<F32> 等数据类型。
  - 支持 Tensor Cores 和 CUDA Cores。
  - 支持 Ampere, Turing, Volta, Pascal, Maxwell 架构。
CUTLASS 2.5 - 2021年2月
- 隐式 GEMM 3D 卷积:
  - 支持前向传播 (Fprop3D)、反向数据梯度 (Dgrad3D) 和反向权重梯度 (Wgrad3D)。
  - 支持 NDHWC 布局的 5D 张量。
- 张量规约 (Tensor Reductions):
  - 支持 m-to-n 的仿射布局张量规约。
  - 自定义规约函数。
  - 对 2^63 个元素的大规模支持。
- 融合卷积 + 卷积示例。
CUTLASS 2.6 - 即将发布
- 结尾融合模式 (Epilogue fusion pattern):
  - 列上广播向量 (例如：加偏置 Bias add)。
  - 列上部分规约 (例如：批量归一化 Batch normalization)。

CUTLASS 卷积性能 (相对于 CUDNN)

下图展示了在 NVIDIA A100 GPU 和 CUDA 11.3 环境下，使用混合精度训练 (F16<->F16+F32) 时，CUTLASS 2.5 在 ResNet50 各层上的性能与 cuDNN 的对比。几何平均 (Geomean) 结果显示，CUTLASS 的性能达到了 cuDNN 的 90%。

4D 张量上的 2D 卷积：前向传播

卷积操作涉及三个张量：激活张量 (x)、滤波器张量 (w) 和输出张量 (y)。

输出张量 (y): 维度为 NPQK = {1, 3, 3, 4}
- N: 批量大小 (Batch size)
- P: 输出张量的高度
- Q: 输出张量的宽度
- K: 输出通道数
激活张量 (x): 维度为 NHWC = {1, 4, 4, 3}
- N: 批量大小
- H: 输入张量的高度
- W: 输入张量的宽度
- C: 输入通道数
滤波器张量 (w): 维度为 KRSC = {4, 2, 2, 3}
- K: 滤波器数量
- R: 滤波器高度
- S: 滤波器宽度
- C: 滤波器通道数

2D 卷积定义

前向传播的数学定义如下：
2D 卷积的数学定义和图示

其中：
* y[n, p, q, k] 是输出张量在特定位置的值。
* x[...] 是输入激活张量的值。
* w[...] 是滤波器张量的值。
* h_bar 和 w_bar 是根据步长 (stride)、填充 (pad) 和扩张 (dilation) 计算出的输入坐标。

卷积到 GEMM 的映射

卷积操作可以映射为通用的矩阵乘法 (GEMM) 操作，这使得可以利用高度优化的 GEMM 核函数来加速卷积。

卷积: y = CONV(x, w)
GEMM: C = GEMM(A, B)

映射关系如下：
* 4D 激活张量 x[N,H,W,C] -> 2D 卷积矩阵 A[NPQ, RSC]
* 4D 滤波器张量 w[K,R,S,C] -> 2D 滤波器矩阵 B[RSC, K]
* 4D 输出张量 y[N,P,Q,K] -> 2D 输出矩阵 C[NPQ, K]

滤波器矩阵映射

从 GEMM 的滤波器矩阵坐标 (gemm_k, gemm_n) 到 4D 滤波器张量坐标 (k, r, s, c) 的映射关系。

输出矩阵映射

从 GEMM 的输出矩阵坐标 (gemm_m, gemm_n) 到 4D 输出张量坐标 (n, p, q, k) 的映射关系。

卷积矩阵映射

从 GEMM 的卷积矩阵坐标 (gemm_m, gemm_k) 到 4D 激活张量坐标 (n, h_bar, w_bar, c) 的映射关系。这是隐式 GEMM 的核心，即在计算过程中动态计算激活张量的坐标，而无需事先构造一个巨大的矩阵（im2col）。

GEMM 三重嵌套循环

以下伪代码展示了如何通过一个三重循环来实现映射到 GEMM 的卷积操作。循环遍历 GEMM 输出矩阵的维度 (M, N) 和内积维度 (K)，并在循环体内将 GEMM 坐标实时转换回卷积张量的坐标，以访问激活张量和滤波器张量中的相应元素。

显式GEMM卷积（前向传播）

前向传播（Fprop）的卷积运算 y = CONV(x,w) 可以通过通用矩阵乘法（GEMM）实现。其中：
- x[N,H,W,C]：4D激活张量
- w[K,R,S,C]：4D滤波器张量
- y[N,P,Q,K]：4D输出张量

卷积运算映射到GEMM的维度如下：
- GEMM-M = NPQ
- GEMM-N = K
- GEMM-K = RSC

朴素显式GEMM卷积实现会在全局内存（Global Memory）中创建卷积矩阵（即im2col操作）。这种方法存在一个主要缺点：显式GEMM卷积会使全局内存占用和流量增加RS倍，其中R和S是滤波器的高度和宽度。

隐式GEMM卷积（前向传播）

隐式GEMM卷积是一种更高效的实现方式。它在共享内存（Shared Memory）中动态（on the fly）构建卷积矩阵，而不是在全局内存中物化整个矩阵。

与显式方法相比，隐式GEMM卷积不会增加全局内存的占用和流量，从而显著提高了效率。

分块GEMM回顾

分块GEMM是一种分层、切片的计算模型，旨在通过在共享内存和寄存器中重用数据来优化计算。其数据流如下：
1. 全局内存：存储原始的GEMM A和GEMM B矩阵。
2. 共享内存：加载一个线程块（Thread Block）大小的数据块（Tile）。
3. 寄存器文件：从共享内存加载一个线程束（Warp）大小的数据块。
4. CUDA/Tensor Cores：在这些计算单元上执行计算。
5. SMEM/CUDA Cores：通过Epilogue Tile和Epilogue Functor处理计算结果。
6. 全局内存：将最终结果写回。

更多关于此模型的详细信息，可以参考CUTLASS GTC 2018和2020的演讲。

隐式GEMM卷积与CUTLASS组件

实现隐式GEMM卷积需要新的CUTLASS组件。下图展示了隐式GEMM卷积如何映射到CUTLASS的分层结构中：

cutlass::conv::threadblock: 负责处理隐式GEMM卷积的数据加载部分，将数据从全局内存加载到共享内存。这是本次演讲（GTC talk 2021）的重点。
cutlass::gemm::warp: 负责线程束级别的计算，这部分重用了现有的GEMM组件（在GTC talks 2018-2020中介绍）。
cutlass::epilogue: 负责处理计算的收尾工作，如Epilogue Fusion（将在CUTLASS 2.6版本中发布）。

深入探讨：隐式GEMM卷积

隐式GEMM卷积算法

隐式GEMM卷积算法主要包括以下三个步骤：
1. 将卷积矩阵和滤波器矩阵的一个数据块（tile）加载到共享内存中。
2. 对共享内存中的操作数执行矩阵乘累加（mma）计算。
3. 沿RSC维度进行迭代。

步骤1：加载数据块

以一个具体的例子说明数据块的加载过程：
- 示例数据块大小:
- Tile_M = 128
- Tile_N = 128
- Tile_K = 32
- 输入类型: F16

卷积矩阵块:
- 加载128x32个元素。
- 每一行映射到唯一的(n, p, q)坐标。
- 每一列映射到唯一的通道c坐标。
滤波器矩阵块:
- 加载32x128个元素。
- 每一行映射到唯一的通道c坐标。
- 每一列映射到唯一的滤波器k坐标。
加载机制: 每个线程从每个数据块中发出多个向量加载指令。例如，对于F16操作数，线程T0的第一次加载会获取c=0..7的数据到共享内存。

步骤2：计算矩阵乘累加（mma）

此步骤直接使用cutlass::gemm::warp中的线程束级（warp-level）mma操作符，以充分利用NVIDIA Tensor Cores的计算能力。数据从共享内存经由寄存器文件流向Tensor Cores进行计算。

Page 25: 使用CUTLASS warp级组件在Tensor Cores上进行计算。

步骤3：沿RSC维度迭代

迭代过程包括：
a) 前进以加载下一个数据块到共享内存。
b) 确保对所有滤波器位置(r, s)和通道c进行累加。

为了覆盖整个GEMM-K（RSC）维度，需要在滤波器位置s、r以及通道c上进行迭代。需要启动足够多的平铺迭代来覆盖所有通道元素（C）和滤波器位置（R-by-S）。

平铺迭代次数的计算公式为：
num_tiled_iterations = R * S * ((C + Tile_K - 1) / Tile_K)

CUTLASS：构建连贯且完整的抽象

cutlass::conv::threadblock::Iterators是为实现隐式GEMM卷积而设计的关键抽象。这些迭代器封装了复杂的地址计算和数据加载逻辑。

迭代器实现的抽象接口:
- advance(): 移动到GEMM-K维度上的下一个平铺迭代。
- operator++(): 移动到线程的下一个加载位置。
- at(): 应用函数将输出坐标(p, q)映射到激活张量中的坐标(h, w)。
- valid(): 检查对全局内存中张量的访问是否越界。
- get(): 根据张量坐标获取全局内存中的指针。

卷积矩阵的分析式瓦片迭代器

以下是迭代器中部分核心功能的伪代码实现，展示了如何从逻辑上的卷积矩阵坐标分析计算出物理内存地址。

at(): 根据输出坐标(p, q)、滤波器偏移(r, s)、步长和扩张率计算出在激活张量中的实际坐标(h, w)。
valid(): 检查at()函数计算出的坐标是否在激活张量的有效范围内。
get(): 根据多维坐标和张量的步长（Strides）计算一维内存偏移量，并返回指针。

然而，这种朴素的分析式实现方式存在性能问题。

一个朴素的分析式实现会产生过多的非Tensor Core数学指令，这会影响整体性能。

卷积抽象的优化实现

预计算不变量 (Precompute Invariants)

为了优化卷积抽象的实现，采用了以下策略：

增量表 (Delta tables)
- 用于访问激活张量的偏移量。
- 在整个内核执行期间保持不变（Invariant）。
- 减少了主循环中的指针算术运算。
掩码谓词 (Mask predicates)
- 使用32位谓词向量进行越界（OOB）检查。
- 对于整个线程块（thread block）是不变的。
- 减少了主循环中的逻辑算术运算。

下图展示了卷积运算中涉及的矩阵：滤波器矩阵 (B, Filter matrix)，卷积矩阵 (A, Convolution matrix)，以及输出矩阵 (C, Output matrix)。

预计算增量表 (Precomputed Delta Tables)

在标准的实现中，访问激活张量的坐标需要复杂的计算，如 valid() 和 get() 函数所示。这些计算涉及多次乘法和加法，以确定内存偏移量，并在主循环中重复执行，效率低下。

为了覆盖卷积矩阵中单个 npq（批次、输出高、输出宽）对应的 RSC（滤波器高、滤波器宽、输入通道）维度，需要在 r, s, c 维度上移动并索引到每个 rsc 位置。这导致了过多的非张量核心（non-tensor core）运算。

通过预计算增量，可以简化地址计算。这些增量值对于固定的npq和问题规模是恒定的。

Δs: 当滤波器位置从 s=0 移动到 s=1 时，在激活张量（NHW-by-C格式）中垂直移动的元素数量。Δs = C 个元素。

Δr: 当滤波器位置从 r=0 移动到 r=1 时，在激活张量中垂直移动的元素数量。Δr = W * C 个元素。

Δc: 移动到下一个 Tile_K 通道元素时，在激活张量中水平移动的元素数量。Δc = Tile_K 个元素。

增量表的不变性 (Invariance)

关键在于，Δs, Δr, 和 Δc 这三个增量值在整个内核执行过程中都是不变的。因此，可以预先计算它们，并在循环中通过简单的加法来更新激活张量的指针，从而极大地减少了计算开销。下图展示了如何使用这些预计算的增量来遍历激活张量。

掩码谓词 (Mask Predicates)

掩码谓词策略旨在优化越界（Out-of-Bounds, OOB）检查。原始的 valid() 函数需要在循环内进行多次逻辑比较。

该策略利用了访问模式的规律性。对于一个给定的线程块，其访问的激活张量区域是固定的。因此，哪些访问是有效的（在边界内），哪些是无效的（越界），是可以预先知道的。

下图展示了n=0时，线程T0在不同Tile_K迭代中的首次访存情况。

第一次 Tile_K 迭代 (r=0, s=0, c=0-7): 计算出的 h_bar 为 -1，这是一个越界访问（OOB）。

第二次 Tile_K 迭代 (r=0, s=1, c=0-7): 计算出的 h_bar 仍为 -1，越界。

第三次 Tile_K 迭代 (r=0, s=2, c=0-7): 计算出的 h_bar 仍为 -1，越界。

掩码谓词的不变性 (Invariance)

随着迭代的进行，最终会出现有效的访问。

掩码谓词对于一个线程块是不变的。
这些谓词可以被预先计算并存储在一个位向量（bit-vector）中。
每个线程、每个空间滤波器维度（spatial filter dimension）使用一个32位的位向量来存储这些预计算的有效性信息。

通过这种方式，循环内的复杂逻辑判断被替换为简单的位运算和查表，从而提高了效率。

隐式 GEMM 卷积 (Implicit GEMM Convolution) - 反向数据梯度 (Dgrad)

反向数据梯度（Dgrad）的计算可以表示为 dx = CONV(dy, w)，其中 dy 是输出梯度张量，w 是滤波器张量，dx 是激活梯度张量。这个卷积操作同样可以映射为一个通用的矩阵乘法（GEMM）操作。

GEMM-M = NHW
GEMM-N = C
GEMM-K = KRS

下图展示了如何将 dy（输出梯度）和 w（滤波器）重排成矩阵，通过 GEMM 运算得到 dx（激活梯度）。

隐式 GEMM 卷积 (Wgrad)

后向权重梯度 (Wgrad) 的计算 dw = CONV(dy, x) 也可以映射为 GEMM。
- dy[N,P,Q,K]: 4D 输出梯度张量
- x[N,H,W,C]: 4D 激活张量
- dw[K,R,S,C]: 4D 滤波器梯度张量

GEMM 的维度映射如下：
- GEMM M = K
- GEMM N = RSC
- GEMM K = NPQ

在此映射中：
- GEMM-A 是输出梯度矩阵 dy，维度为 K x NPQ。
- GEMM-B 是卷积矩阵 x（激活张量），维度为 NPQ x RSC。
- GEMM-C 是滤波器梯度矩阵 dw，维度为 K x RSC。

CUTLASS 隐式 GEMM 卷积在 NVIDIA GPU 上的应用

CUTLASS 在 Tensor Cores 上相对于 CUDA Cores 的性能（训练）

基于 Resnet50 层的 Tensor Cores 性能 (F16 <= F16*F16 + F32)

峰值理论数学吞吐量:
CUDA Cores (F32): 1x
Tensor Core (F16): 16x
在 NVIDIA A100 上，使用 F16 输入的 Tensor Cores 是训练最快的数据类型。
对于 F16 输入，CUTLASS 在 NVIDIA A100 上的性能达到了 cuDNN 的 93% 以内 (10.8x / 11.6x)。

Page 48 - CUTLASS 2.5 在 NVIDIA A100 上使用 F16 混合精度进行训练时，相对于 CUDA Cores 的加速比

图表显示了在 NVIDIA A100 上，使用 Tensor Cores 进行混合精度训练 (F16 <= F16*F16 + F32) 时，相对于 CUDA Cores (F32) 的加速比。cuDNN 的几何平均加速比为 11.6x，而 CUTLASS 为 10.8x。

基于 Resnet50 层的 Tensor Cores 性能 (F32 <= TF32*TF32 + F32)

峰值理论数学吞吐量:
CUDA Cores (F32): 1x
Tensor Core (TF32): 8x
TensorFloat32 (TF32) 使 Tensor Cores 能够直接处理 F32 输入。
对于 F32 输入，CUTLASS 在 Tensor Cores 上的性能相对于 CUDA Cores 实现了 5.3x 的加速。

Page 49 - CUTLASS 2.5 在 NVIDIA A100 上使用 TensorFloat32 进行训练时，相对于 CUDA Cores 的加速比

图表显示了 CUTLASS 在 Tensor Cores 上使用 TensorOp(TF32) 进行训练时的性能。前向传播、后向数据梯度和后向权重梯度的几何平均加速比为 5.3x。

CUTLASS 卷积性能相对于 CUDNN (F16[NHWC]) (推理)

基于 Resnet50 层的 Tensor Cores 性能

以下图表展示了 CUTLASS 2.5 在 NVIDIA A100 和 2080Ti 上，使用 F16、S8 和 S4 数据类型进行推理时，相对于 cuDNN F16[NHWC] 的性能。

F16 数据类型 (F16 <= F16*F16 + F16)

对于 F16 数据类型，CUTLASS 在 Tensor Cores 上的推理性能与 cuDNN 相当。
在 A100 上达到 1.02x 加速比。
在 2080Ti 上达到 0.97x 加速比。

Page 50 - CUTLASS 使用 F16 数据类型在 A100 和 2080Ti 上的推理性能对比

S8 数据类型 (S8 <= S8*S8 + S32)

相对于 F16，CUTLASS 使用 S8 数据类型在 Tensor Cores 上的推理性能获得了显著提升：
在 A100 上，对于 NHWC 和 NC32HW32 交错布局，均获得 1.23x 的增益。
在 2080Ti 上，对于 NHWC 和 NC32HW32 交错布局，分别获得 1.62x 和 1.74x 的增益。

Page 51 - CUTLASS 使用 S8 数据类型在 A100 和 2080Ti 上的推理性能对比

S4 数据类型 (S4 <= S4*S4 + S32)

相对于 F16，CUTLASS 使用 S4 数据类型在 Tensor Cores 上的推理性能获得了更大提升：
在 A100 上，对于 S4[NHWC] 和 S4[NC64HW64] 交错布局，分别获得 1.54x 和 1.55x 的增益。
在 2080Ti 上，对于 S4[NHWC] 和 S4[NC64HW64] 交错布局，分别获得 2.28x 和 2.43x 的增益。

Page 52 - CUTLASS 使用 S4 数据类型在 A100 和 2080Ti 上的推理性能对比

Epilogue Fusion (结尾融合)

Epilogue Fusion

Epilogue Fusion 是指在 GEMM 计算（位于寄存器中）完成之后、将结果写回全局内存之前，将一个或多个附加操作融合进来的技术。

CUTLASS 支持的 Epilogue Fusion 模式及其应用：

Epilogue Fusion 模式	应用
Element-wise operators (按元素操作)	Scale (缩放), bias (偏置), activations (激活函数)
Data type conversion (数据类型转换)	F32->F16, Int32->Int8
Broadcast vector over columns (列向量广播)	Bias add (偏置加法)
Partial reduction over columns (列部分规约)	Sum or sum-of-squares for batch norm (批量归一化的求和或平方和)

Epilogue Fusion 示例

下图展示了两种 Epilogue Fusion 模式：列向量广播（用于偏置加法）和列部分规约（用于批量归一化）。

Page 55 - CUTLASS 中 Epilogue Fusion 模式的详细图解

可用性:
CUTLASS 2.5: 已支持按元素操作和数据类型转换。
CUTLASS 2.6 (展望): 计划支持列向量广播和列部分规约。

结论

结论：面向 NVIDIA GPU 的隐式 GEMM 卷积

CUTLASS 2.4 和 2.5: Nov 2020 和 Feb 2021
- 开源的 CUDA C++ 模板库，用于 CUDA 开发。
- 提供可重用的构建模块，用于利用 Tensor Cores 进行 GEMM 和卷积运算。
- CUTLASS 卷积性能与 cuDNN 相当（> 90%）。

CUTLASS 2.6: 即将发布的版本展望
- 支持新的 epilogue fusion 模式：
- 列向量广播（偏置加法）
- 列部分规约（批量归一化）

立即开始! https://github.com/NVIDIA/cutlass

参考文献

CUTLASS: https://github.com/NVIDIA/cutlass
GTC 2018 talk (S8854): CUTLASS: Software primitives for dense linear algebra at all levels and scales within CUDA
GTC 2019 talk (S9593): cuTENSOR:High-performance Tensor Operations in CUDA (joint talk with cuTENSOR)
GTC 2020 talk (S21745): Developing CUDA kernels to push Tensor Cores to the Absolute Limit on NVIDIA A100
CUTLASS Parallel For All blog post