關注了就能看到更多這么棒的文章哦～

Intel AMX support in 5.16

By Jonathan Corbet
November 8, 2021
DeepL assisted translation
https://lwn.net/Articles/874846/

x86 指令集中有非常多的指令，但這并沒有阻礙它繼續(xù)擴充。即將推出的英特爾處理器中將推出以 "高級矩陣擴展"（AMX，Advanced Matrix Extensions）為名的一套新指令，用于對矩陣數(shù)據(jù)進行操作。經(jīng)過一個有些波折的開發(fā)過程之后，對 AMX 的支持可以進入即將到來的 5.16 內核了。要使用這套指令集的話，應用程序開發(fā)人員就必須改變一些做法。

AMX（在?https://software.intel.com/content/dam/develop/external/us/en/documents-tps/architecture-instruction-set-extensions-programming-reference.pdf?中有介紹）是一個通用架構，用來加速 x86 處理器上的矩陣操作。其最初形式中實現(xiàn)了一組多達 8 個 "tiles" （也就是 16 行 64 字節(jié)的陣列）。程序員可以在這些 tiles 中存儲任意尺寸的矩陣。16x16 的 32 位浮點值的矩陣可以支持，此外也支持其他規(guī)格的矩陣。目前支持的一個操作是將存儲在兩個 tile 中的矩陣相乘，然后將結果寫入到第三個 tile 中。通過串聯(lián)這類操作，就可以實現(xiàn)任意大小的矩陣乘法。顯然，未來還會支持其他種類的矩陣操作。

雖然看起來 AMX 可能是一個用于數(shù)值分析場合的功能，但其真正的目標用例似乎是機器學習這種應用。這可以解釋為什么是支持 16 位浮點運算，而不是 64 位。

AMX 在設計上就允許內核來控制這些功能是否可以被某個特定進程來使用。設計成這樣是有幾個原因，其中之一是正如人們能預料到的，AMX 指令使用了大量的處理器資源。一個進程如果在共享的計算機（shared computer）上做大量 AMX 計算的進程可能會對其他進程產生很大影響。但是，除非內核和用戶空間進程都為此準備好了，否則 AMX 也無法得到正確支持。

Development process

對 AMX 的支持是由 Chang Bae 在 2020 年 10 月首次發(fā)布出來的，但得到的 review 很少。到了 2 月份第 4 版出來的時候，更多的開發(fā)者開始關注它，他們對這個功能當前被整合到內核現(xiàn)有的浮點單元（FPU）代碼中的這種做法不是很滿意。接下來各種版本就接踵而至，作者的挫折感似乎也在不斷累積。9月底的時候，Len Brown 發(fā)布了一份對話記錄，其中似乎列出了接下來前進的方向。

不幸的是，第二天發(fā)布的第 11 版似乎忽略了許多已經(jīng)做出的決定。這一版就引起了 Thomas Gleixner 的強烈指責，他認為這個功能是在強行要塞進內核，而完全沒有聽取人們的抱怨。AMX 的情況并不樂觀，但是幕后人們仍在繼續(xù)坐著工作。10 月中旬的時候，Gleixner 發(fā)表了他對 FPU 代碼的大規(guī)模 rework，希望能減輕在 kernel 里支持 AMX 的代價。不久之后，一個新的 AMX patchset 就出現(xiàn)了，這就是 5.16 中得到的最終成果。

Gleixner 對這部分代碼的 pull request 中也承認，這些代碼還是相對不成熟的：

注意，這部分代碼還是比較新的，盡管對 AMX 的支持已經(jīng)進行了一年多了。

對 FPU 代碼的重大重構，使其能夠進行適當?shù)恼?，這個工作是在 3 周前開始的。重構現(xiàn)有的 FPU 代碼和原始的 AMX patch 花了一周時間，并進行了廣泛的 review 和 test。唯一沒有經(jīng)歷 review 和 test 的內容也僅限于支持 AMX 的系統(tǒng)，這對那些不使用英特爾和他們的 early access program （新功能的早期使用項目）的人員來說完全沒有影響。像往常一樣，可能會有潛伏的龍（也就是指潛在的 bug），但到目前為止，細化之后的重構已經(jīng)經(jīng)受住了考驗，最終尚未發(fā)現(xiàn)的那些影響是可以使用 bisect 方式來確認出來源的，在 5.16 版本發(fā)布之前應該可以輕松解決。著名的最后補充…

FPU 代碼是相對比較棘手、底層性的工作，所以在新的工作中發(fā)現(xiàn)一兩個潛伏的 bug 確實是不奇怪的。

Using AMX

如上所述，內核能夠控制哪些進程可以使用 AMX 指令。用戶空間進程要想使用它，第一步就是利用一個新的 arch_prctl()命令（名為 ARCH_GET_XCOMP_SUPP）來獲得系統(tǒng)能支持的 feature 列表。如果返回結果中相應 bit 被置位了，那么就可以直接使用 AMX 了。此外還有另一個 arch_prctl()命令（名為 ARCH_REQ_XCOMP_PERM）可以用來請求獲取權限來使用 AMX。這里將會進行一些檢查（其中一個檢查會在本文下面介紹），security modules 也有機會表達意見。不過，通常情況下這些 request 都會被批準。所獲取的權限適用于此進程中的所有線程，并在 fork 后也會繼續(xù)擁有這個權限，不過，調用 execve() 的話就會喪失原進程的特有權限。

AMX 帶來的一個挑戰(zhàn)是處理器在運行 AMX 指令時可能創(chuàng)建大量的內部狀態(tài)（internal state）。如果 CPU 在運行這部分指令過程中被中斷打斷，這些狀態(tài)就必須保存在某個地方，否則會丟失許多進展。因此，如果一個進程使用了 AMX，內核就必須在做其他事情之前先要在其中斷處理程序中保存 10KB 左右的數(shù)據(jù)。這種保存工作是通過 XSAVE 指令完成的。

內核針對這個目的，為每個進程分配了相應的內存空間。不過，為系統(tǒng)中的每個進程都分配 10KB 的內存空間會浪費很多內存。而且大多數(shù)進程不會使用 AMX 指令。不過好消息是，處理器可以配置成在進程第一次執(zhí)行 AMX 指令時 trap 進入 kernel，然后內核可以檢查是否被允許使用這些指令，允許的話就分配一個適當大小的 buffer 來保存 FPU 狀態(tài)并允許此次操作繼續(xù)進行。

這里有一個與 sigaltstack() 系統(tǒng)調用有關的潛在問題，這個調用允許線程來建立一個新的堆棧用來處理信號。如果此進程使用了 AMX，那么這個堆棧也必須要足夠大，才夠用于保存 FPU 的狀態(tài)。多年來，開發(fā)者一直被告知要用 MINSIGSTKSZ 來作為這個堆棧的最小 size，也就是 2KB，但對于使用 AMX 的進程來說這就遠遠不夠了。事實上，這個 size 甚至都不夠用來使用 AVX-512 extension，過去曾經(jīng)因為這個原因而引發(fā)了一些 stack 被破壞的問題。

為了在 AMX 中避免這個問題，內核將仔細檢查，確保所有的 signal stack 都足夠大。這個檢查是在每次調用 sigaltstack()時進行的，但是當一個進程首次請求獲取 AMX 權限時也會對當前堆棧進行檢查。不使用 AMX 的進程將不需要那么大的堆棧，因此，不會在這些檢查的時候被拒絕。而那些想要使用 AMX 的進程則只有擁有足夠大的 signal stack 時才允許繼續(xù)執(zhí)行。

等這些檢查相關的基礎設施都到位了之后，內核就可以確保那些使用 AVX-512 的進程有足夠大小的 signal stack 了。但是，如果強制確保這個條件的話，有可能破壞現(xiàn)在看起來還能正常工作的應用程序，比如也許它們的 signal handler 處理程序從未被實際調用過。為了避免這種問題，準備了一個內核配置選項（STRICT_SIGALTSTACK_SIZE）和一個命令行選項（strict_sas_size=），它們都可以被用來控制在使用 AVX-512 時是否進行嚴格檢查。

假設所有這些都能配合好，那么這也就是 AMX 在 5.16 中得以支持所采取的形式了。想了解更多信息的人可以看看包含 AMX test case 的 commit 以及一些關于 arch_prctl()命令的文檔。同時，在接下來的九周左右的時間里，請留意是有有 bug 的出現(xiàn)。

全文完
LWN 文章遵循 CC BY-SA 4.0 許可協(xié)議。

長按下面二維碼關注，關注 LWN 深度文章以及開源社區(qū)的各種新近言論～