GPU Microarch 学习笔记【3】Tensor Core-编程知识

Nvidia自从Volta/Turing（2018）架构开始，在stream multi processor中加入了tensor core，用于加速矩阵计算。如下图所示，其中每个SM有两个tensor core。相信大家也看了这个图很多次，那么一个tensor core里面的64个绿色小格子代表的是什么级别的计算呢？

图自[3]

指令与架构

Tensor core支持mixed precision和FP16矩阵计算：D = A x B + C. mixed precision指的是 A和B为FP16，但是C和D为FP32或FP16.

CUDA9.0支持的最小矩阵乘法为16x16x16（不是一个周期就能计算完毕的），下图所示为Tensor Core load A, B, C，进行mma（matrix multiply-accumulate），而后store结果D的PTX指令，其中每个指令都有sync，用于进行warp-wide的同步。

Tensor Core PTX指令

Volta的每个SM core内部有两个tensor core，每个Tensor Core每个周期可以完成一个 4x4x4的MACC操作，如下图所示:

这里的4x4x4指的是A*B矩阵中A的维度为4x4， B的维度为4x4，A*B的计算量即为4x4x4。

这里写的two 4x4x4是因为有两个tensor core。

下图右侧即为Turing架构新引入的tensor core，在A和B矩阵的精度为FP16时，每次A的4x4，B的矩阵4x4，每次进行乘法累加，得到下方绿色的4x4。INT8和INT4复用乘法单元，计算能力分别x2，x4。

传统的CUDA教学里面，每个线程的寄存器是private，仅自己可见。

但是对于Tensor Core来说，一个warp内的线程，互相之间的寄存器都是可见的。因此一个warp的32个线程，可以各自load输入矩阵A,B,C的一部分，加载到register file，读取时也可以复用其他寄存器加载的寄存器。

与传统CUDA相同，tensor core的一个warp仍然是32个thread，但是不同的是，32个thread被切成了8个thread group，每个thread group 4个线程。

Load

首先将是将矩阵加载到寄存器中：

如左图所示，一共16行，每4行为一个segment，每个segment由两个thread group一同加载。一行16个元素，每个元素16bit，共计256bit，需要两条load.E.128指令完成。

这里个人理解，因为一共四行，每行可以两个thread加载，所以可以每个thread执行一个load.E.128，一共32个thread。

32(thread 个数)*128(load.E.128) = 16(矩阵行数)*16(矩阵列数)*16(每个元素大小)

文章未具体介绍。

计算MMA

加载之后就是计算了，一条mma 16*16*16的指令如下图所示，

在Nvidia的SASS汇编级别会被展开成16条指令：

这里只展示了mixed precision的汇编。FP16的汇编指令与此相同，但是概念相通，这里就不赘述。

上图右侧显示了计算完毕一条16x16x16需要的时间，54个cycle，除了开始需要10个cycle，后面每条指令基本上只需要2个cycle。

需要注意的是，这里的16条指令会在一个warp的32个thread上同时执行。

上图左侧，一共有4个set，每个set内部又有4个step。

Set, Step与Thread Group

Set,step都是汇编级别的指令，都是逐一计算的，但是thread，thread group这些thread的概念，是同时进行的。具体step，set和thread group的切分如下图所示。

从set的角度看矩阵的计算：

如上图所示的是thread group 0的计算每个set可以计算完毕一个4*4*8的计算，4次set计算完毕之后，就可以得到最终的C矩阵中的一个4*8小矩阵。与此同时，其他的thread group 1-7也在进行计算，最终我们就可以得到完整的C矩阵。

更细致的计算解释如下图所示：

如上图(b)所示，每个step完成了2*4*4的计算，得到了2*4的矩阵输出（上图b的ADBC标错了，应该是ABCD）。

那么是如何完成2*4*4的计算的呢？

OCTET

这里又要引入一个新的概念OCTET，具体的OCTET负责的计算一个8*8的小矩阵具体如下图所示:

之所以引入这个概念，是因为OCTET内部的threadgroup存在加载数据的依赖关系，比如上图b）中，每个小方块为4x4的矩阵。

d,c,b,a,D,C,B,A由thread group 0加载
h,g,f,e,H,G,F,E由thread group4加载

因此当thread group0计算a*A和a*E时，需要等待thread group4加载完毕，也就是Octet内部的计算需要另一个thread group load完毕数据，但是Octet之间不存在数据依赖关系。

而这个依赖关系，个人理解后面会体现在thread group 0和4公用的matrix buffer上。

如下图所示，step0/1，threadgroup0和4公用thread0加载的A，step2/3，threadgroup0和4公用thread0加载的E.

这里的a[0:1]代表a的0行和1行

而再回忆一下刚才介绍的输出C的矩阵：

可以看出OCTET 0的这个部分是由Thread group 0和 Thread group 4共同计算完成的。

Tensor Core 微架构

如上图所示，右上角的为一个FEDP unit （four-element dop-product），这个灰色的单元，每个可以完成一个行矩阵和一个列矩阵的点乘计算，即一个1*4的行矩阵和一个4*1的列矩阵的点乘计算。这个4次的乘法累加，每周期可以计算完毕一次。

每个thread group内部有4个，那么就是4*4，可以理解成一个1*4的A矩阵和一个4*4的B矩阵，每个周期可以计算完毕一次A和B的成累加，那么两个周期，就可以计算完毕一个两行四列的A矩阵和一个4行4列的B矩阵的计算，得到一个两行四列的C矩阵。对应了上面的单个step的计算图。

我们再回忆一下，tensor core支持每周期4*4*4的矩阵计算，这里一个thread group每个周期可以计算完毕1*4*4的计算，每个Octet内部有两个thread group，即2*4*4，而每个tensor core内部又有两个Octet，这样每个周期就可以有4*4*4的计算量了。

Octet0内部的Thread group0和Thread group共用了一个matrix B buffer，利用了计算时复用矩阵的B的特点，减轻了加载操作数到计算单元的负担。

Final

到这里我们已经从架构，指令，指令分解到微架构已经明白了Volta的tensor core是如何计算的。

那么具体的每个tensor core内部有2个Octet，每个Octet有2个thread group，每个thread group有4个thread，每个thread执行在一个FEDP上，每个FEDP内部有四个乘法单元，2*2*4*4 = 64.

也就是说最开始的那个问题，tensor core里面的绿色小格子就是一个乘法单元：

也就是微架构图里的这个：

本文是基于文章[1]和[2]的理解。

[1] Modeling Deep Learning Accelerator Enabled GPUs

[2] Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking

[3] https://learnopencv.com/demystifying-gpu-architectures-for-deep-learning-part-2/