[译]将stencil计算转化为Tensor Core 上的矩阵乘法

原文地址： dl.acm.org/doi/10.1145…

原文作者： Yuetao Chen，Kun Li

写在前面

本文主要工作为通过将stencil计算转化为矩阵乘法，进一步充分利用GPU计算能力，从而提高性能。

摘要

TCU 越来越多的集成到现代高性能处理器中，以提高矩阵乘法性能。然后，受限于其规格过高，它在改进模版计算等其他关键科学运算方面的提升空间依旧很大。

本文所介绍的 ConvStencil ，是一种新的模版计算系统，将模版计算高效的转换为在TCU上的矩阵乘法。作者首先为 ConvStencil 建立了一个性能模型，以指导TCU上的算法设计和优化。基于该模型，本文提出了三种技术：

内存高效布局转换
使用双重细分和核融合技术来计算计算密集型的任务
使用查找表和脏位填充来消除算法实现与硬件设计之间的冲突

介绍

随着深度学习模型的普及，处理器越来越多的采用专门的TCU来加速矩阵乘法操作。然而，在高性能计算领域中，计算模式更加多样复杂，大部分无法直接用矩阵乘法表达，HPC中比较常见一种计算模式是stencil，他是一个预定义的模式，通过时间维度迭代更新d维空间网格中的每个点，一个点在时间t的值是本身与邻近点在前一个时间t-1的加权和（译者注：也就是常见的生命游戏）。是被广泛用于科学和工程建模领域的一种计算技术。

ConvStencil的灵感来源于其与深度学习中的卷积操作的相似性，两种计算方法都涉及使用模版核（或卷积核）形成滑动窗口，对输入矩阵窗口内的数据进行加权计算。为有效支持张量核上的卷积，基于 GEMM 的卷积计算中使用了 im2row(col) 方法。它包括将输入和滤波器转换为矩阵，从而允许通过矩阵计算卷积。

然而模版计算和卷积计算仍包含一些细微差距，导致利用im2row来实现模版计算仍有以下挑战：

首先，模版计算时的模版核的数量以及通道数量只有一个，因此相对于卷积计算，对Tensor Core的利用率会较低，同时由于卷积计算直接在原始数据上进行卷积，而模版计算需要将输入数据转换为矩阵，这会增加内存的占用率。

此外，可能会遇到算法实现和硬件设计之间的冲突，比如warp分歧和bank冲突，这也会大大影响性能。

由于GPU没有复杂的分支预测，当同一个warp中线程在执行相同的指令时，如果进入不同的分支，同一时刻除了正在执行的分支外，其他的分支都被阻塞了

在布局转换中，我们引入了stencil2row，以减少内存消耗，为MM创建高效的内存布局。与im2row相比，他实现了70.0% 到96.4%的内存占用减少。在计算自适应中，我们提出了双重细分，通过矩阵细分来提高张量核心的利用率，将张量核心的使用率从12.5%提高到87.5%。同时，核融合降低了矩阵稀疏性，进一步提高了模版核的计算密度。在冲突消除中，我们设计了一个查找表，以避免昂贵的操作并减少冗余的寻址计算。此外，脏位填充使用填充区来写入脏数据并避开条件分支，从而实现了无冲突的实现，进一步提高了性能。

背景和挑战

在Tensor Core上基于GEMM的卷积计算

Tensor Core 具有混合精度的矩阵相乘和累加的能力，如下公式

D_{m \times n} = A_{m \times k} \times B_{k \times n} + C_{m \times n}

基于GEMM库的卷积计算将卷积转换为矩阵计算，成为在Tensor Core上计算卷积的有效方法。输入矩阵是通过将图像的每个内核大小的块展开成为一行。卷积核矩阵是通过将过滤器权重展开成一列来创建的，然后在此基础上加进行矩阵运算。

挑战

即使两者很相似，但仍有一些挑战

矩阵转换带来的高内存需求（对于卷积操作，可以接受，因为卷积核也有很多，平衡内存和计算开销，但是stencil只有一个核）
Tensor Core 利用率低， FP64 A100 仅支持884，浪费了7列
算法和硬件之间的冲突：
- 内存访问的大量重复偏移计算引起了与标准stencil计算的冲突，这些冲突消耗了计算资源并导致性能下降
- 布局转换中存在条件分支和bank冲突，导致warp分歧和串行内存访问。

ConvStencil

ConvStencil 包括布局转换、计算适应、冲突消除等基本组件。

布局转换阶段，提出了stencil2row方法，将输入重塑为两个较小的矩阵
在计算适应阶段，通过对双重细分对stencil2row矩阵中选定的切片应用tensor core矩阵乘法累加操作，生成结果。
在冲突消除部分，预先计算指针偏移量以避免耗时的整数除法和模运算。同时提出了脏位填充法，通过利用填充区域来消除负载冲突和消除条件分支。

性能模型

ConvStencil性能模型包括三个公式

T = max(T_{compute}, T_{memory}) \\ T_{compute} = \frac{1}{fN_{tcu}} \sum^{K_{tcu}}_{i=0} (K_{tcu} \times CPI_{tcu_i}) \\ T_{memory} = max(\frac{data_R}{bw_G} + \frac{data_W}{bw_G}, \frac{data_{transW}}{bw_S} + \frac{data_{transR}}{bw_S})

公式1 表示总时间是计算时间和内存访问时间中的较大值
公式2 计算所需的时间是时钟频率的倒数与所需周期数的乘积。所需的周期数是通过将程序中每种类型指令的数量与该指令所花费的周期数的乘积相加来计算的。在 NVIDIA A100 GPU 上，TCU 上的 FP64 MMA 指令的周期数为 16 。
公式3 内存访问所需的时间是跨不同内存层次结构的读/写时间的最大总和。

内存转换

im2row转换会导致内存爆炸是因为有冗余元素，基于三个观测结果提出了stencil2row

im2row矩阵大部分元素都是冗余的
im2row转换中，冗余行的数据已经存储在其他行中
共享内存的访问延迟比全局内存低一个数量级

stencil2row通过两种方式节省了数据传输开销

隐式的在共享内存中构建了stencil2row矩阵的切片，从而减少了全局内存读写操作的开销
相较于im2row，减少了内存占用，从而减少了写入共享内存的数据量

计算适应

本文提出了一种基于 stencil2row转换的新算法，称为双重细分，用于高效计算tensor core 上的stencil结果。双重细分通过迭代调用来逐步计算所有的stencil。每个双重细分首先构建两个半结果矩阵，然后将两个半结果矩阵相加得到stencil计算结果。该算法能够有效利用Tensor Cores进行stencil计算，并且基于stencil2row矩阵构建。

冲突移除

通过引入查找表，预先计算指针偏移量，避免了整数除法和模运算，同时引入了脏位填充，以消除条件分支语句。填充可以改变数据映射到共享内存的方式，从而消除负载冲突，通过将脏位转储到填充空间中可以消除条件分支和相应的计算。