CPU AMX 詳解

概述

2016 年開始，隨著 NV GPU AI 能力的不斷加強，隱隱感覺到威脅的 Intel 也不斷在面向資料中心的至強系列 CPU 上堆砌計算能力，增加 core count 、提高 frequency 、增強向量協處理器計算能力三管其下。幾乎每一代 CPU 都在 AI 計算能力上有所增強或拓展，從這個方面來講，如果我們說它沒認識到勢，沒有采取行動，也是不公平的。

從上圖不難看到，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)，即矩陣協處理器了。充分說明一句老話，你永遠叫不醒一個裝睡的人，要用火燒他。不管怎麼樣，這下總算是賽道對齊了，終於不是拿長茅對火槍了。

算力如何

AI 工作負載 Top-2 的運算元：

• Convolution

• MatMul/Fully Connected

這倆本質上都是矩陣乘。怎麼計算矩陣乘，有兩種化歸方法：

• 化歸成向量點積的組合，這在CPU中就對應AVX

• 化過程分塊矩陣乘的組合，這在CPU就對應AMX

我們展開講講。

問題定義

假設有如下矩陣乘問題：

 

AVX如何解決矩陣乘問題

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

不就化歸成向量點積了嘛。那向量的長度是可以任意指定的，但硬體是有固定長度的，怎麼辦？很簡單，就把每個向量切成每個長度為的塊，多做幾次就好了。這個就是區分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倍。

AMX如何解決矩陣乘問題

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

需要注意的是整個操作需要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圍繞矩陣這一一等公民的支援分為計算和資料兩個部分。

• 計算部分：目前僅有矩陣乘支援，由稱為TMUL(Tile Matrix mULtiply Unit)的模組來實現。但也為後面支援其他的矩陣運算留了想像。

• 資料部分：由一組稱為TILES的二維暫存器來承擔。

其系統框圖如下：

計算部分

TMUL 硬體層面的實現也比較直觀，是一個典型的systolic array設計。比較好的是array的每一行都複用了原來的AVX-512 BF16的設計，堆疊了16個AVX-512 BF16單元，在一個cycle內完成了一個的運算，因此完成整個的計算需要16個cycle。

Systolic形式的邏輯圖，如下。可以看出每個cycle輸出 的一行結果。

資料部分

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

每個TILE暫存器容量為：16行 64 bytes，因此可用於存放：

• 的 FP32 矩陣

• 的 BF16 矩陣

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

  • ……

讓我們邊走邊看！