LWN：5.16中支持Intel AMX指令！

關(guān)注了就能看到更多這么棒的文章哦～

Intel AMX support in 5.16

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

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

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

雖然看起來(lái) AMX 可能是一個(gè)用于數(shù)值分析場(chǎng)合的功能，但其真正的目標(biāo)用例似乎是機(jī)器學(xué)習(xí)這種應(yīng)用。這可以解釋為什么是支持 16 位浮點(diǎn)運(yùn)算，而不是 64 位。

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

Development process

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

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

Gleixner 對(duì)這部分代碼的 pull request 中也承認(rèn)，這些代碼還是相對(duì)不成熟的：

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

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

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

Using AMX

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

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

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

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

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

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

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

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

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