?1.5 TFlops 到底有何魅力？首先，這是在電池供電的單核 MacBook Air 2020 上運(yùn)行；
其次，這會(huì)以每條指令約 0.5 納秒的延遲運(yùn)行。
那些強(qiáng)大的加速器或 GPU 張量核不在我們的考慮范疇。我們這里討論的是與 CPU 寄存器相隔一個(gè)周期的實(shí)際線性代數(shù)性能。
奇怪的是，蘋果一直在向我們隱瞞這一點(diǎn)。在本文中，我們將通過(guò)一些代碼來(lái)揭開迷霧。
什么是 AMX 協(xié)處理器？
它可以說(shuō)是 SIMD 的典范。一個(gè)重要的區(qū)別是 AMX:CPU 比率不是 1:1；并非每個(gè)內(nèi)核都有自己的 AMX 協(xié)處理器。
以下是用于加載或存儲(chǔ)值的規(guī)格：
最小值與完整的 AVX512 寄存器一樣寬
但這些值是從哪里加載或存儲(chǔ)的？顯然，這樣的大小會(huì)很快用完整個(gè) NEON 寄存器文件。不過(guò) AMX 有一個(gè)單獨(dú)的寄存器文件，這有些奇怪。
寄存器分為三組：X、Y 和 Z。對(duì)每個(gè)指令，X 和 Y 組保存輸入，Z 組保存輸出。
如我們所見(jiàn)，X 和 Y 相當(dāng)大。二者之間有一個(gè)完整的 KB。Z 則令人稱奇，然后是一些：
（劇透：一條 AMX 指令可以填充 1024 字節(jié)（Z 寄存器的 1/4）。）
那么如何從 X 和 Y 到 Z？方法很多，以至于它不那么適合 ISA 編碼。所以蘋果決定將大部分指令信息編碼在通用寄存器中。事實(shí)證明，這個(gè)決定很贊，因?yàn)榭梢栽?AMX 上執(zhí)行代碼的運(yùn)行時(shí)（動(dòng)態(tài)）配置。
這篇文章旨在提高協(xié)處理器利用效率。有一些 vector-vector 指令可以輸出長(zhǎng)度相同的向量，但不會(huì)使芯片的計(jì)算能力飽和。反而必須借助外積來(lái)進(jìn)行。?
何為外積？假設(shè)有兩個(gè)輸入向量 u 和 v：
外積是一個(gè)矩陣，包含各元素可能組合對(duì)的乘積。（這里給出一些提示，說(shuō)明為什么 Z 寄存器組比 X 和 Y 大得多。）
在 AMX 芯片上，可歸結(jié)為一個(gè)非常簡(jiǎn)單的指令，就像這樣：
可以設(shè)置一個(gè)標(biāo)志，使其從上一個(gè)結(jié)果中累加：
這樣，我們就完全具備了編寫矩陣乘法所需：從輸入矩陣中重復(fù)加載 16 個(gè)浮點(diǎn)數(shù)，并將它們的外積累加成 16x16 輸出?？s小 K 尺寸甚至無(wú)關(guān)緊要！
我們簡(jiǎn)化一下這個(gè)問(wèn)題，并隱式轉(zhuǎn)置矩陣乘法。A 和 B（輸入）都將 K（縮減維度）作為主導(dǎo)維度。這在實(shí)踐中并不重要，但它大大簡(jiǎn)化了我們的代碼。
這里有一個(gè)參考，可用來(lái)檢查我們提出的解決方案：
void reference_16x16xK (float *A, float *B, float *C, uint64_t K) {
for (uint32_t m = 0; m < 16; ++m) {
for (uint32_t n = 0; n < 16; ++n) {
C [n * 16 + m] = 0;
for (uint32_t k = 0; k < K; ++k) {
C [n * 16 + m] += A [k * 16 + m] * B [k * 16 + n];
}
}
}
}下面是我們?cè)?AMX 中的實(shí)現(xiàn)方法：
/only set for k == 0uint64_t reset_z = 1ull << 27;
for (uint32_t k = 0; k < K; ++k) {
uint64_t idx = k % 4;
// 64 bytes = 16 floats
AMX_LDX ((uint64_t) A + k * 64);
AMX_LDY ((uint64_t) B + k * 64);
//now we do 4 indepedent outer products (avoiding pipeline hazards)
AMX_FMA32 (reset_z);
reset_z = 0;
}
神奇的是，我們沒(méi)有處理任何寄存器，但卻悄悄做了些處理。以同樣的方式將 reset_z 編碼為位掩碼，寄存器地址也編碼在傳遞給 AMX_* 的參數(shù)中。指向 A 和 B 的指針最多只能使用 56 位，因此蘋果工程師將信息存儲(chǔ)在其他 8 位中。我們只是意外將其全部設(shè)置為 0。因此，在本例中，對(duì) X 和 Y 我們將寄存器置 “0”。
將 Z 寄存器存儲(chǔ)到內(nèi)存的代碼有點(diǎn)復(fù)雜，因?yàn)槲覀冎惶畛淞说谝涣小Ｋ孕枰@取寄存器 0、4、8 等：
for (uint64_t i = 0; i < 16; ++i) {
const uint64_t z_register = (i * 4ull) << 56;
AMX_STZ (z_register | (uint64_t) C + i * 64);
}但你會(huì)發(fā)現(xiàn)，運(yùn)行上面的代碼非常慢。只有區(qū)區(qū)幾百 GFlops。
為什么會(huì)這樣？有兩個(gè)原因。
開始的減速是流水線冒險(xiǎn)。每個(gè) AMX_FMA32 都依賴于前一個(gè)，因?yàn)槿祭鄯e到寄存器文件的一個(gè)子集中。我們最終只達(dá)到了寄存器文件全節(jié)流的 25%，剩余部分閑置，未能實(shí)現(xiàn)指令級(jí)并行。
接下來(lái)的問(wèn)題是從內(nèi)存中加載的效率很低。我們其實(shí)可以一次加載 128 個(gè)字節(jié)，但上面的代碼只能加載 64 個(gè)字節(jié)。類似地，可以加載到其他寄存器，不必每次都加載到相同的寄存器。也可以實(shí)現(xiàn)一定程度的指令級(jí)并行。
那么計(jì)劃是什么？
我們將向 X 和 Y 加載 128 個(gè)字節(jié)，然后計(jì)算一個(gè) 32x32 塊。這將涉及 16x16 塊的 4 次獨(dú)立計(jì)算，形成指令級(jí)并行，可以更高效利用加載的內(nèi)存（每個(gè) 64 字節(jié)寄存器使用兩次）。
void mm32x32xK (float* A, float* B, float* C, uint64_t K) {
//flag to load/store 128 bytes
const uint64_t load_store_2 = 1ull << 62;
const uint64_t load_store_width = 128; //in bytes
//only set for k == 0
uint64_t reset_z = 1ull << 27;
for (uint32_t k = 0; k < K; ++k) {
uint64_t idx = k % 4;
//load to X, Y (skipping every other index because we're loading 128 bytes)
AMX_LDX (load_store_2 | (idx * 2) << 56 | (uint64_t) A + k * load_store_width);
AMX_LDY (load_store_2 | (idx * 2) << 56 | (uint64_t) B + k * load_store_width);
//offset into X and Y registers is byte-wise
const uint64_t offset = idx * load_store_width;
//now we do 4 indepedent outer products (avoiding pipeline hazards)
AMX_FMA32 (reset_z | (0ull << 20) | ((offset + 0ull) << 10) | ((offset + 0ull) << 0));
AMX_FMA32 (reset_z | (1ull << 20) | ((offset + 64ull) << 10) | ((offset + 0ull) << 0));
AMX_FMA32 (reset_z | (2ull << 20) | ((offset + 0ull) << 10) | ((offset + 64ull) << 0));
AMX_FMA32 (reset_z | (3ull << 20) | ((offset + 64ull) << 10) | ((offset + 64ull) << 0));
reset_z = 0;
}
for (uint64_t i = 0; i < 16; ++i) {
//store interleaved
AMX_STZ (load_store_2 | ((i * 4ull + 0) << 56) | (uint64_t) C + i * load_store_width);
AMX_STZ (load_store_2 | ((i * 4ull + 2) << 56) | (uint64_t) C + (16 + i) * load_store_width);
}
}void mm32x32xK (float* A, float* B, float* C, uint64_t K){
//flag to load/store 128 bytes const uint64_t load_store_2 = 1ull << 62; const uint64_t load_store_width = 128; //in bytes
//only set for k == 0 uint64_t reset_z = 1ull << 27;
for (uint32_t k = 0; k < K; ++k) { uint64_t idx = k % 4; //load to X, Y (skipping every other index because we're loading 128 bytes) AMX_LDX (load_store_2 | (idx * 2) << 56 | (uint64_t) A + k * load_store_width); AMX_LDY (load_store_2 | (idx * 2) << 56 | (uint64_t) B + k * load_store_width);
//offset into X and Y registers is byte-wise const uint64_t offset = idx * load_store_width;
//now we do 4 indepedent outer products (avoiding pipeline hazards) AMX_FMA32 (reset_z | (0ull << 20) | ((offset + 0ull) << 10) | ((offset + 0ull) << 0)); AMX_FMA32 (reset_z | (1ull << 20) | ((offset + 64ull) << 10) | ((offset + 0ull) << 0)); AMX_FMA32 (reset_z | (2ull << 20) | ((offset + 0ull) << 10) | ((offset + 64ull) << 0)); AMX_FMA32 (reset_z | (3ull << 20) | ((offset + 64ull) << 10) | ((offset + 64ull) << 0)); reset_z = 0; }
for (uint64_t i = 0; i < 16; ++i) { //store interleaved AMX_STZ (load_store_2 | ((i * 4ull + 0) << 56) | (uint64_t) C + i * load_store_width); AMX_STZ (load_store_2 | ((i * 4ull + 2) << 56) | (uint64_t) C + (16 + i) * load_store_width); }}我在上面添加了注釋，關(guān)于說(shuō)明性標(biāo)志有些有趣的細(xì)節(jié)。Corsix 在解釋這一點(diǎn)上做得很好，所以我要留下鏈接：
- 加載和存儲(chǔ)標(biāo)志：https://github.com/corsix/amx/blob/main/ldst.md
- FMA 標(biāo)志：https://github.com/corsix/amx/blob/main/fma.md
那么我們到底能有多快？這一定程度上取決于 K，但我們達(dá)到了 1.5TFlops 處理的問(wèn)題更大相對(duì)來(lái)說(shuō)會(huì)獲得更好的性能，這也不足為奇，因?yàn)榫彺婵梢愿玫仡A(yù)熱，CPU 有更多時(shí)間交錯(cuò)指令。
總的來(lái)說(shuō)，在當(dāng)今大型神經(jīng)網(wǎng)絡(luò)競(jìng)相追逐通用 AI 的背景下，這類問(wèn)題顯得微不足道，然而卻為小型神經(jīng)網(wǎng)絡(luò)在現(xiàn)實(shí)計(jì)算中找到一席之地。如果一個(gè)預(yù)測(cè)模型可于幾十納秒內(nèi)在電池供電的筆記本上運(yùn)行，或?qū)樵究赡苁褂锰皆囁惴ǖ牡胤綆?lái)更多價(jià)值。你怎么看？
原文鏈接：https://jott.live/markdown/1.5tflop_m1?
