CPU AMX 詳解

@(System)[CPU]

概述

2016 年開始，隨著 NV GPU AI 能力的不斷加強，隱隱感覺到威脅的 Intel 也不斷在面向數(shù)據(jù)中心的至強系列 CPU 上堆砌計算能力饮戳，增加 core count 然想、提高 frequency 梯嗽、增強向量協(xié)處理器計算能力三管其下匕得。幾乎每一代 CPU 都在 AI 計算能力上有所增強或拓展，從這個方面來講栈虚，如果我們說它沒認識到勢袖外，沒有采取行動，也是不公平的魂务。

從上圖不難看到曼验，2015年的 Sky Lake 首次引入了 AVX-512 (Advanced Vector eXtensions)向量協(xié)處理器，與上一代 Broadwell 的 AVX2 相比粘姜， 每個向量處理器單元的單精度浮點乘加吞吐翻倍鬓照。接著的Cascade Lake 和 Cooper Lake又拓展了 AVX-512 ，增加了對 INT8 和 BF16 精度的支持孤紧，奮力想守住 inference 的基本盤豺裆。一直到 Sapphire Rapids，被市場和客戶用腳投票号显，前有狼（NVIDIA）后有虎（AMD）臭猜，都把自己的食盆都快拱翻了，終于意識到在AI的計算能力上不能在按摩爾定律線性發(fā)育了押蚤，最終也步Google和NVIDIA的后塵获讳，把AVX升一維成了AMX(Advanced Matrix eXtension)，即矩陣協(xié)處理器了活喊。充分說明一句老話，你永遠叫不醒一個裝睡的人，要用火燒他钾菊。不管怎么樣帅矗，這下總算是賽道對齊了，終于不是拿長茅對火槍了煞烫。

算力如何

AI 工作負載 Top-2 的算子：

• Convolution
• MatMul/Fully Connected

這倆本質(zhì)上都是矩陣乘浑此。怎么計算矩陣乘，有兩種化歸方法：

• 化歸成向量點積的組合滞详，這在CPU中就對應(yīng)AVX
• 化過程分塊矩陣乘的組合凛俱，這在CPU就對應(yīng)AMX

我們展開講講。

問題定義

假設(shè)有如下矩陣乘問題：

AVX如何解決矩陣乘問題

AVX把向量作為一等公民料饥，每次計算一個輸出元素蒲犬，而該元素等于的第行與的第列的點積，即有：

不就化歸成向量點積了嘛岸啡。那向量的長度是可以任意指定的原叮，但硬件是有固定長度的，怎么辦巡蘸？很簡單奋隶，就把每個向量切成每個長度為的塊，多做幾次就好了悦荒。這個就是區(qū)分AVX各代的主要因素唯欣。下面以AVX2為例淺釋一下。

AVX2 FP32 (k=8)
AVX2使用的寄存器長度為256 bit搬味，也就是8個FP32數(shù)境氢，此時。AVX的乘加> 指令操作示意如下：

一個時鐘周期可以完成兩個8維向量的點積操作身腻，也叫FMA(Fused Multiply > Add）操作产还。因此每個AVX單元的FLOPS為：16 FLOPS/cycle。

以FP32/BF16為例嘀趟，AVX算力的代際演進如下脐区，可以看出相鄰代際增長是平平無奇的2倍。

AMX如何解決矩陣乘問題

以BF16為例她按，AMX把矩陣乘操作化歸為若干個的分塊矩陣乘的組合牛隅，如下所示。

需要注意的是整個操作需要16個cycle完成酌泰，因此不難計算每個AMX單元的FLOPS為：1024 OPS/cycle媒佣。這下單AMX單元與單AVX單元的每時鐘周期的算力提高了16倍，有點像樣了陵刹。目前Sapphire Rapids每個核有一個AMX單元默伍，而有兩個AVX單元，因此每核的每時鐘周期算力提高倍數(shù)為8倍。

如何計算含有AMX CPU的peak TFLOPS
公式：

假設(shè)你有一個56核也糊，每核有1個AMX單元炼蹦，且AMX頻率為1.9 GHz的CPU。其BF16 peak TFLOPS應(yīng)為：

如何實現(xiàn)的

AMX圍繞矩陣這一一等公民的支持分為計算和數(shù)據(jù)兩個部分狸剃。

• 計算部分：目前僅有矩陣乘支持掐隐，由稱為TMUL(Tile Matrix mULtiply Unit)的模塊來實現(xiàn)。但也為后面支持其他的矩陣運算留了想像钞馁。
• 數(shù)據(jù)部分：由一組稱為TILES的二維寄存器來承擔虑省。

其系統(tǒng)框圖如下：

計算部分

TMUL 硬件層面的實現(xiàn)也比較直觀，是一個典型的systolic array設(shè)計僧凰。比較好的是array的每一行都復(fù)用了原來的AVX-512 BF16的設(shè)計探颈，堆疊了16個AVX-512 BF16單元，在一個cycle內(nèi)完成了一個的運算允悦，因此完成整個的計算需要16個cycle膝擂。

Systolic形式的邏輯圖，如下隙弛〖懿觯可以看出每個cycle輸出 的一行結(jié)果。

數(shù)據(jù)部分

每個AMX單元共有8組TILES寄存器全闷，TILE寄存器可以存放一個二維矩陣的子矩陣叉寂，有專門的load/store指令。

每個TILE寄存器容量為：16行 64 bytes总珠，因此可用于存放：

• 的 FP32 矩陣
• 的 BF16 矩陣
• 的 INT8 矩陣

路才開始

邁出腳只是路的開始屏鳍，而不是結(jié)束。后面有的是路(問題)：

• 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開銷
  • ......

讓我們邊走邊看钓瞭！