之前為了驗證一個技術(shù)上的想法，隨手寫的一段匯編：
https://github.com/pigirons/sgemm_hsw/blob/master/sgemm_kernel_x64_fma.S
然后被朋友轉(zhuǎn)到知乎上：
立交橋跳水冠軍：大佬是怎么優(yōu)雅實現(xiàn)矩陣乘法的？
這段匯編基本描述了這樣一個寄存器分塊的外積矩陣乘法：
圖中長矩形表示SIMD寄存器，方塊表示標量。用一個4 × 3的寄存器分塊來存儲累加結(jié)果，然后從矩陣A里依次按列讀取四個標量廣播到向量寄存器的每個lane里，和矩陣B的三個行向量寄存器做乘法，累加到矩陣C的對應(yīng)向量寄存器中。
根據(jù)我在github上給出的測試數(shù)據(jù)，這段代碼只要控制矩陣三個維度（m, k, n）的大小使A，B，C三個矩陣都能放到L1 cache里，就可以達到十分接近理論峰值的性能。比如m = 24，k = 64，n = 24，在我的Zen2架構(gòu)4750G上可以做到大于99%的浮點峰值性能。大家對于這個算法的點評得基本到位，但這里還是要補充幾點：
關(guān)于使用的寄存器數(shù)量
大家討論的一個點：這個程序恰好用滿了AVX2指令的最多16個寄存器。但寄存器用了多少其實無關(guān)緊要，我這里只是恰好用到最大，可以支持FMA指令更長的延遲，更多的發(fā)射端口。
在我之前的文章里：
提到過FMA的吞吐和延遲的關(guān)系。以Intel Haswell架構(gòu)為例，F(xiàn)MA指令延遲是5個周期，每周期可以吞吐2條。這樣為了掩蓋所有FMA的執(zhí)行周期，至少需要10條無數(shù)據(jù)依賴關(guān)系的FMA指令才能填滿各級流水線。按照文章里面測試FMA峰值的代碼，反映在時空圖上就是這樣的：
縱軸代表haswell架構(gòu)port0和port1上各有一個FMA單元，每個FMA單元有五級流水線（S0->S5）。橫軸代表時鐘周期：第一個周期，最開始的兩條fma指令（紅色）同時分別發(fā)射到port0和port5，并執(zhí)行流水線的第一階段（S0）；隨著時鐘的tick，這兩條指令分別進入之后的流水階段（S2到S5）；從第二周期開始，再發(fā)射兩條FMA指令（藍色）到S0；以此類推，到第五個周期，最后兩條指令開始發(fā)射（灰色），同時最開始的兩條指令執(zhí)行最后一個階段（S5）；第六個周期開始，循環(huán)也回到開頭，繼續(xù)發(fā)射兩條紅色的FMA指令，完美銜接，沒有任何氣泡。可以想象，只要程序足夠長，兩個FMA單元的10個流水線階段（相當于計算資源）絕大部分時間都是有指令在執(zhí)行的，除了開頭進入流水線，和結(jié)尾結(jié)束流水線。
這個原理反應(yīng)在寄存器分塊上，就是存儲累加結(jié)果的分塊使用的寄存器數(shù)量，要大于等于10，才能保證中間的FMA指令沒有停頓。我這里安排4 × 3 = 12這個分塊，就沒有問題。實際上，安排2 × 5或者5 × 2，都沒有問題。這個情況最少用2 × 5 + 2 + 1 = 13個寄存器就可以打滿峰值。對于后續(xù)的skylake等一系列新架構(gòu)，F(xiàn)MA指令的延遲降低到4個周期，同時保持一個周期發(fā)射兩條，那么在這種情況下，最少可以用2 × 4 + 2 + 1 = 11個寄存器就可以打滿峰值了。我這里安排4 × 3=12，就是為了照顧各種不同架構(gòu)FMA發(fā)射端口數(shù)量以及延遲可能有差別。
這同時引出了下一個問題：
向量讀取指令和FMA指令的關(guān)系
既然寄存器分塊只要保證10個以上的寄存器，那能不能用10 × 1的分塊呢？我們回顧一下Intel Haswell架構(gòu)圖：
可以看出，除了Port0和Port1分別有一條FMA單元，Port2和Port3也分別有一條LD單元（Load），這個LD單元可以讀取向量。我們假設(shè)所有的數(shù)據(jù)都在L1 cache里，LD指令就可以流水執(zhí)行，大約是4個周期可以把一條向量讀取到寄存器。這樣，每個周期，Haswell CPU就可以同時發(fā)射兩條FMA以及兩條vmovaps(或者vbroadcastss）指令。這就要求，總的Load指令數(shù)量最多跟FMA指令數(shù)量相等，一旦超過，F(xiàn)MA將不是瓶頸，LD指令將變成瓶頸。對于10 × 1的分塊，我們需要10 + 1 = 11次LD才能喂給10條FMA指令，這樣一定不能打滿FMA的峰值；對于4 × 3的分塊，我們需要4 + 3 = 7次LD就能喂給12條FMA指令，Haswell架構(gòu)應(yīng)對這個比例非常輕松。
對于不同的處理器，這個結(jié)論可以推廣為：LD和FMA指令的比例，要小于等于處理器提供的LD和FMA單元同周期發(fā)射能力的比例。對于某些順序單發(fā)射（比如某已流片的RISC-V）或者亂序發(fā)射能力有限（比如多數(shù)早期arm如a9, a15, a57）的架構(gòu)，甚至要把這個比例在寄存器數(shù)量能支撐的范圍內(nèi)降到最低，這樣就可以保證更多的周期是在發(fā)射FMA。其實對于m × n的寄存器分塊，LD的數(shù)量一般就是m + n，所以盡量構(gòu)造一個足夠大的接近方形的分塊，可以最大限度地利用FMA的峰值性能。
關(guān)于ymm15
可以看到程序里面對于矩陣B的三個向量的讀取，分別用了三個寄存器（ymm12, ymm13和ymm14），但對于矩陣A的四個標量的讀取，只用一個ymm15寄存器，這樣不會有問題嗎？我們拿出程序的一段結(jié)構(gòu)觀察：
vmovaps 0(%rbx), %ymm12
vmovaps 32(%rbx), %ymm13
vmovaps 64(%rbx), %ymm14
vbroadcastss 0(%r12, %rax, 4), %ymm15
vfmadd231ps %ymm12, %ymm15, %ymm0
vfmadd231ps %ymm13, %ymm15, %ymm1
vfmadd231ps %ymm14, %ymm15, %ymm2
vbroadcastss 0(%r13, %rax, 4), %ymm15
vfmadd231ps %ymm12, %ymm15, %ymm3
vfmadd231ps %ymm13, %ymm15, %ymm4
vfmadd231ps %ymm14, %ymm15, %ymm5
vbroadcastss 0(%r14, %rax, 4), %ymm15
vfmadd231ps %ymm12, %ymm15, %ymm6
vfmadd231ps %ymm13, %ymm15, %ymm7
vfmadd231ps %ymm14, %ymm15, %ymm8
vbroadcastss 0(%r15, %rax, 4), %ymm15
vfmadd231ps %ymm12, %ymm15, %ymm9
vfmadd231ps %ymm13, %ymm15, %ymm10
vfmadd231ps %ymm14, %ymm15, %ymm11
可以看到，第二條vbroadcastss指令要寫入第二個標量給ymm15。從這個程序執(zhí)行的性能結(jié)果分析，幾乎每個周期都緊湊地有兩條FMA發(fā)射出去，第二組三條FMA指令又依賴第二條vbroadcastss讀取到的數(shù)據(jù)，所以這條廣播指令一定被CPU亂序調(diào)度到第一組三條FMA指令更前面去了，因為LD指令本身需要4個周期。從程序邏輯來看，第一組三條FMA依賴的ymm15數(shù)據(jù)，是從第一條廣播指令讀取來的，但是第二條廣播指令又會被調(diào)度到第一組三條FMA之前，這樣數(shù)據(jù)不就亂套了嗎？
其實不然，CPU有一個強大的register renamer部件，我們能操縱的ymm寄存器，只是邏輯寄存器，具體映射到實際的物理寄存器，是有一套復(fù)雜的轉(zhuǎn)換規(guī)則。對于這個例子，CPU能分析出指令的數(shù)據(jù)流依賴關(guān)系，自動把第二組ymm15映射到與第一組ymm15不同的物理寄存器上，這樣兩組數(shù)據(jù)就毫無關(guān)系了，可以隨意改變發(fā)射順序，不影響結(jié)果的正確性。
這個技術(shù)部分地解決了邏輯寄存器不夠用的問題。對于前面構(gòu)造寄存器分塊的分析，實際上m + n條LD指令，并不需要m + n個寄存器存儲，只需要min(m, n) + 1條寄存器就夠了。這就增加了寄存器分塊本身可以利用的寄存器的數(shù)量，構(gòu)造更大的分塊，進一步減少LD和FMA的比例。
這就是sgemm_hsw這個程序里面構(gòu)造4 × 3寄存器分塊暗含的三個玄機，一環(huán)套一環(huán)，雖然有其他的分塊方案仍然能做到接近的峰值，但這個方案對不同x86架構(gòu)的魯棒性更好，適用范圍更廣。這些參數(shù)不是恰巧湊出來的，都是一步步精細算出來的，希望這篇文章對大家有所幫助。
最后強調(diào)一點，這還只是在L1 cache里的矩陣乘法分塊，真正想優(yōu)化任意大小的矩陣乘法，還有很多工作要做，比如CPU的prefetch機制，多級cache的性能特性，TLB miss的處理，多核多線程的任務(wù)分配調(diào)度等，這些技術(shù)大家去參考Goto的經(jīng)典論文即可。
