概述
2016 年开始,随着 NV GPU AI 能力的不断加强,隐隐感觉到威胁的 Intel 也不断在面向数据中心的至强系列 CPU 上堆砌计算能力,增加 core count 、提高 frequency 、增强向量协处理器计算能力三管其下。几乎每一代 CPU 都在 AI 计算能力上有所增强或拓展,从这个方面来讲,如果我们说它没认识到势,没有采取行动,也是不公平的。
![](https://i-blog.csdnimg.cn/blog_migrate/a4db49321ca944fdba1edd538e1e8df4.png)
从上图不难看到,2015年的 Sky Lake 首次引入了 AVX-512 (Advanced Vector eXtensions)向量协处理器,与上一代 Broadwell 的 AVX2 相比, 每个向量处理器单元的单精度浮点乘加吞吐翻倍。接着的Cascade Lake 和 Cooper Lake又拓展了 AVX-512 ,增加了对 INT8 和 BF16 精度的支持,奋力想守住 inference 的基本盘。一直到 Sapphire Rapids,被市场和客户用脚投票,前有狼(NVIDIA)后有虎(AMD),都把自己的食盆都快拱翻了,终于意识到在AI的计算能力上不能在按摩尔定律线性发育了,最终也步Google和NVIDIA的后尘,把AVX升一维成了AMX(Advanced Matrix eXtension),即矩阵协处理器了。充分说明一句老话,你永远叫不醒一个装睡的人,要用火烧他。不管怎么样,这下总算是赛道对齐了,终于不是拿长茅对火枪了。
![](https://i-blog.csdnimg.cn/blog_migrate/722a752c188abf0d284d787055f7aa2f.png)
算力如何
AI 工作负载 Top-2 的算子:
Convolution
MatMul/Fully Connected
这俩本质上都是矩阵乘。怎么计算矩阵乘,有两种化归方法:
化归成向量点积的组合,这在CPU中就对应AVX
化过程分块矩阵乘的组合,这在CPU就对应AMX
我们展开讲讲。
问题定义
假设有如下矩阵乘问题:
![](https://i-blog.csdnimg.cn/blog_migrate/b8033f0349ce2301c22e05e3ecb8a80b.png)
![](https://i-blog.csdnimg.cn/blog_migrate/72dbdd448fc119d2182b7cf5cea7c10b.png)
AVX如何解决矩阵乘问题
![](https://i-blog.csdnimg.cn/blog_migrate/e83d6c48ebe3d712f13fdedee6e913d9.png)
AVX把向量作为一等公民,每次计算一个输出元素
![](https://i-blog.csdnimg.cn/blog_migrate/5014fbc281821ad4395dd4049d423285.png)
,而该元素等于
![](https://i-blog.csdnimg.cn/blog_migrate/38e5e221923431d07a096ba8b60800e7.png)
的第
![](https://i-blog.csdnimg.cn/blog_migrate/4431162d77dd7015a5567cc6aba479b7.png)
行与
![](https://i-blog.csdnimg.cn/blog_migrate/d0e7f664c8d8dff9a3a753a6277d8e87.png)
的第
![](https://i-blog.csdnimg.cn/blog_migrate/878184c1a0cab3f8ae796ef8c2f3f548.png)
列的点积,即有:
![](https://i-blog.csdnimg.cn/blog_migrate/0d7873c56042889d0e90a060482cbde5.png)
不就化归成向量点积了嘛。那向量的长度是可以任意指定的,但硬件是有固定长度的,怎么办?很简单,就把每个向量切成每个长度为
![](https://i-blog.csdnimg.cn/blog_migrate/30e6b5a7a10ed5a39bc165b5543bfa0c.png)
的块,多做几次就好了。这个
![](https://i-blog.csdnimg.cn/blog_migrate/5536466e7966856984d56254198faddc.png)
就是区分AVX各代的主要因素。下面以AVX2为例浅释一下。
AVX2 FP32 (k=8)
AVX2使用的寄存器长度为256 bit,也就是8个FP32数,此时
。AVX的乘加> 指令操作示意如下:
一个时钟周期可以完成两个8维向量的点积操作,也叫FMA(Fused Multiply > Add)操作。因此每个AVX单元的FLOPS为:16 FLOPS/cycle。
以FP32/BF16为例,AVX算力的代际演进如下,可以看出相邻代际增长是平平无奇的2倍。
![](https://i-blog.csdnimg.cn/blog_migrate/2c240ae7b90a8878c3a246c77966bfcb.png)
AMX如何解决矩阵乘问题
以BF16为例,AMX把矩阵乘操作化归为若干个
![](https://i-blog.csdnimg.cn/blog_migrate/994857179550293a1933629ec6fd2d02.png)
的分块矩阵乘的组合,如下所示。
![](https://i-blog.csdnimg.cn/blog_migrate/45e29384a4006dd81f736ac2ac85bd3e.png)
需要注意的是整个操作需要16个cycle完成,因此不难计算每个AMX单元的FLOPS为:1024 OPS/cycle。这下单AMX单元与单AVX单元的每时钟周期的算力提高了16倍,有点像样了。目前Sapphire Rapids每个核有一个AMX单元,而有两个AVX单元,因此每核的每时钟周期算力提高倍数为8倍。
如何计算含有AMX CPU的peak TFLOPS
公式:
假设你有一个56核,每核有1个AMX单元,且AMX频率为1.9 GHz的CPU。其BF16 peak TFLOPS应为:
如何实现的
AMX围绕矩阵这一一等公民的支持分为计算和数据两个部分。
![](https://i-blog.csdnimg.cn/blog_migrate/a6fa1dfe59a30f298f22da4bb8253992.png)
计算部分:目前仅有矩阵乘支持,由称为TMUL(Tile Matrix mULtiply Unit)的模块来实现。但也为后面支持其他的矩阵运算留了想像。
数据部分:由一组称为TILES的二维寄存器来承担。
其系统框图如下:
![](https://i-blog.csdnimg.cn/blog_migrate/b3c598c357b5d383eab8b179f8ce1549.png)
计算部分
TMUL 硬件层面的实现也比较直观,是一个典型的systolic array设计。比较好的是array的每一行都复用了原来的AVX-512 BF16的设计,堆叠了16个AVX-512 BF16单元,在一个cycle内完成了一个
![](https://i-blog.csdnimg.cn/blog_migrate/ba8f3316955df09598461116c5e9c929.png)
的运算,因此完成整个
![](https://i-blog.csdnimg.cn/blog_migrate/d5acd3833a7e189276b7a26a868a0ff0.png)
的计算需要16个cycle。
![](https://i-blog.csdnimg.cn/blog_migrate/f9842c49e4cca3881e50a44e440fed1d.png)
Systolic形式的逻辑图,如下。可以看出每个cycle输出
![](https://i-blog.csdnimg.cn/blog_migrate/808aa065047fee6f62acb9673453e55a.png)
的一行结果。
![](https://i-blog.csdnimg.cn/blog_migrate/7d98198ad355eb858858c2474b243129.png)
数据部分
每个AMX单元共有8组TILES寄存器,TILE寄存器可以存放一个二维矩阵的子矩阵,有专门的load/store指令。
![](https://i-blog.csdnimg.cn/blog_migrate/d08359e4fc15c637360743d1949be2d8.png)
每个TILE寄存器容量为:16行
![](https://i-blog.csdnimg.cn/blog_migrate/d701db7c9d52c61de6a14c15e224bd7d.png)
64 bytes,因此可用于存放:
![](https://i-blog.csdnimg.cn/blog_migrate/521661ca7f761a52802d41bee98d343f.png)
的 FP32 矩阵
![](https://i-blog.csdnimg.cn/blog_migrate/7dba28eb63a6fc8613c16f34bdbe72b8.png)
的 BF16 矩阵
![](https://i-blog.csdnimg.cn/blog_migrate/5cb96bb865eb4e164e43d42ca615d15e.png)
的 INT8 矩阵
路才开始
迈出脚只是路的开始,而不是结束。后面有的是路(问题):
HW
TILE 和 memory 之间的 load 和 save 带宽与TMUL计算能力的匹配度
AI workload 一般都是矩阵操作(matmul, conv等)与向量操作混杂,而向量操作有分为 element-wise 操作和 reduce 类操作
这3类操作算力的匹配度
矩阵寄存器与向量寄存器之间的 data path 通畅度如何
……
SW
如何提高SW efficiency
如何摆平AI框架要求的plain data layout与AMX硬件要求的data layout之间的re-layout开销
……