Armv9 技術(shù)講堂 | Neon、SVE 和 SME 實(shí)現(xiàn)矩陣-矩陣乘法的比較
Arm 始終專注于架構(gòu)演進(jìn),確保生態(tài)系統(tǒng)能夠適應(yīng)未來的技術(shù)趨勢(shì)和不斷變化的計(jì)算需求。Armv9 架構(gòu)上的 可伸縮矩陣擴(kuò)展 (SME) 顯著提高了 Arm CPU 對(duì)現(xiàn)有人工智能 (AI) 和機(jī)器學(xué)習(xí) (ML) 工作負(fù)載的處理能力,從而在各種 AI 驅(qū)動(dòng)的設(shè)備和應(yīng)用中帶來速度更快、響應(yīng)更靈敏的用戶體驗(yàn)。
本文引用地址:http://cafeforensic.com/article/202409/462643.htm在此前的內(nèi)容中,Arm 技術(shù)專家為大家簡要介紹了 SME 和 SME 指令 ,本期將帶你了解如何使用 Neon、可伸縮向量擴(kuò)展 (SVE) 和 SME 這三種不同的 Arm 技術(shù)實(shí)現(xiàn)相同的矩陣-矩陣乘法算法。這三個(gè)示例展現(xiàn)了這些技術(shù)之間的關(guān)鍵差異,為開發(fā)者將代碼從 Neon、SVE/SVE2 移植到 SME/SME2 提供指導(dǎo)。
架構(gòu)演進(jìn)
Armv7 引入了高級(jí) SIMD 擴(kuò)展,為一系列整型和浮點(diǎn)型提供單指令多數(shù)據(jù) (SIMD) 操作。Neon 是高級(jí) SIMD 指令的一種實(shí)現(xiàn)方案,作為部分 Arm Cortex-A 系列處理器的擴(kuò)展提供。2011 年 Armv7-A 中引入了 Neon。Neon 提供固定寬度的 128 位寄存器。這意味著每條 Neon 指令對(duì)固定數(shù)量的數(shù)據(jù)值進(jìn)行操作,例如四個(gè) 32 位數(shù)據(jù)值。
2016 年 Armv8-A 中引入了 SVE,2021 年 Armv9-A 中引入了 SVE2,它們提供可變長度寄存器。寄存器的大小由實(shí)現(xiàn)方案定義,從 128 位到 2048 位寄存器不等。這意味著程序員不知道可用寄存器的大小,因此必須將代碼編寫為與向量長度無關(guān)。因此,每條指令處理的數(shù)據(jù)值的數(shù)量不是固定的,而是可變的。
2021 年 Armv9-A 中引入了 SME 和 SME2,也提供可變長度寄存器。SME 引入了兩個(gè)關(guān)鍵的新架構(gòu)特性:Streaming SVE 模式和 ZA 存儲(chǔ)。Streaming SVE 模式是一種高吞吐量的矩陣數(shù)據(jù)處理模式,ZA 存儲(chǔ)則是一種專用的二維數(shù)組,可以方便地進(jìn)行常見的矩陣操作。這些特性使 SME 和 SME2 能夠高效地處理矩陣和基于向量的工作負(fù)載。
這些 SIMD 架構(gòu)擴(kuò)展提供的指令可加速多種應(yīng)用,包括媒體和信號(hào)處理應(yīng)用、高性能計(jì)算 (HPC) 應(yīng)用,以及 ML 應(yīng)用。
本文中的示例使用了內(nèi)置函數(shù) (intrinsics),即編譯器提供的與特定 Arm 指令相對(duì)應(yīng)的函數(shù)。這使得程序員能夠用 C 語言而不是匯編語言編寫整個(gè)程序。
矩陣-矩陣乘法
本文中的所有三個(gè)示例都能夠?qū)崿F(xiàn)矩陣-矩陣乘法。矩陣乘法需要接收兩個(gè)輸入矩陣,通過將第一個(gè)矩陣一行的每個(gè)元素與第二個(gè)矩陣一列的相應(yīng)元素相乘,然后對(duì)這些乘積求和,從而生成一個(gè)結(jié)果矩陣。結(jié)果矩陣的維度由第一個(gè)矩陣的行數(shù)和第二個(gè)矩陣的列數(shù)決定。例如,3 x 2 矩陣乘以 2 x 3 矩陣將得到 3 x 3 矩陣。
要將矩陣 A 和矩陣 B 相乘,矩陣 A 的列數(shù)必須等于矩陣 B 的行數(shù)。將矩陣 A 和矩陣 B 相乘將得到矩陣 C。
Neon
此示例使用 Neon 內(nèi)置函數(shù)來執(zhí)行矩陣-矩陣乘法。相關(guān)代碼執(zhí)行以下操作:
兩個(gè)輸入矩陣包含以列優(yōu)先格式存儲(chǔ)的 32 位浮點(diǎn)數(shù)據(jù)。
代碼以 4 x 4 的塊形式迭代處理這些矩陣中的所有數(shù)據(jù)。
vld 內(nèi)置函數(shù)將輸入矩陣的行和列中的四個(gè)值加載到 Neon 寄存器中。
每個(gè) fma Neon 內(nèi)置函數(shù)執(zhí)行四次乘加運(yùn)算,計(jì)算正在處理的 4 x 4 塊的結(jié)果。
vst 內(nèi)置函數(shù)將結(jié)果矩陣存儲(chǔ)到內(nèi)存中。
以下是使用 Neon 內(nèi)置函數(shù)的示例代碼:
此示例使用以下 Neon 代碼特性:
此示例使用 4 x 4 的固定塊大小。這意味著輸入矩陣在兩個(gè)維度上都必須是四的倍數(shù)。可以通過用零填充矩陣來處理其他大小的矩陣。
SVE/SVE2
此示例使用 SVE2 內(nèi)置函數(shù)來執(zhí)行矩陣-矩陣乘法。
Neon 示例和 SVE2 示例之間的主要區(qū)別在于 SVE2 使用可變長度向量。Neon 示例可以使用 4 x 4 的固定塊大小來匹配 Neon 寄存器中的四個(gè) 32 位值,但程序員要到運(yùn)行時(shí)才能知道 SVE2 寄存器的大小。這意味著代碼必須與向量長度無關(guān)。此示例使用 predication 來控制 SVE2 內(nèi)置函數(shù)操作的數(shù)據(jù)值的數(shù)量。這意味著無論實(shí)現(xiàn)的大小如何,它們都能夠精準(zhǔn)地適配 SVE2 寄存器。Neon 示例使用 32 位浮點(diǎn)數(shù)據(jù)類型 float32x4_t,其中 的“4”表示每個(gè) Neon 寄存器可以包含四個(gè) 32 位值。SVE2 示例使用 svfloat32_t 數(shù)據(jù)類型,因?yàn)闊o法在運(yùn)行前知道 SVE2 寄存器的大小。
相關(guān)代碼執(zhí)行以下操作:
兩個(gè)輸入矩陣包含以列優(yōu)先格式存儲(chǔ)的 32 位浮點(diǎn)數(shù)據(jù)。
代碼以四行一組的形式迭代處理這些矩陣中的所有數(shù)據(jù)。它使用 svcntw 內(nèi)置函數(shù),返回向量中 32 位元素的數(shù)量,以匹配加載到 SVE2 寄存器大小的列數(shù)。這有助于避免對(duì)外循環(huán)每次迭代中的元素?cái)?shù)量進(jìn)行硬編碼。whileit 內(nèi)置函數(shù)生成一個(gè) predicate,以確保不超過矩陣的界限。
四個(gè) svld 內(nèi)置函數(shù)使用之前生成的 predicate 將矩陣數(shù)據(jù)加載到 SVE2 寄存器中。
svlma 內(nèi)置函數(shù)執(zhí)行乘加運(yùn)算,計(jì)算當(dāng)前迭代的結(jié)果。
svst 內(nèi)置函數(shù)將結(jié)果矩陣存儲(chǔ)到內(nèi)存中。
以下是使用 SVE2 內(nèi)置函數(shù)的示例代碼:
此示例使用以下 SVE2 代碼特性:
SME/SME2
此示例使用 SME2 匯編指令來執(zhí)行矩陣-矩陣乘法。SME2 示例與其他示例的區(qū)別如下:
SME2 示例使用匯編代碼,而不是其他示例所使用的內(nèi)置函數(shù)。
SME2 提供 ZA 存儲(chǔ),這是專為矩陣運(yùn)算設(shè)計(jì)的二維數(shù)據(jù)數(shù)組。此 ZA 存儲(chǔ)內(nèi)的子數(shù)組可作為 tile 進(jìn)行訪問,并且 tile 內(nèi)的元素可以垂直或水平訪問。這為操作矩陣數(shù)據(jù)提供了非常靈活的機(jī)制。
SME2 提供了執(zhí)行矩陣運(yùn)算的新指令。例如,fmopa 指令可計(jì)算外積。
SME2 提供了一種多向量二維 predication 機(jī)制,以確保不超出矩陣邊界。
Streaming SVE 模式(使用 smstart 指令進(jìn)入的)啟用 SME2 指令和 ZA 存儲(chǔ)。
此示例使用 matLeft 和 matRight 作為輸入矩陣。示例中用到了以下運(yùn)算原理:兩個(gè)矩陣相乘等同于依次對(duì) matLeft 的每一列和 matRight 的每一行的外積求和。
初始輸入矩陣作為行優(yōu)先數(shù)組存儲(chǔ)在內(nèi)存中。矩陣乘法是將 matLeft 的一列與 matRight 的一行的外積求和。由于外積需要 matLeft 的列元素,因此代碼重新排列 matLeft 數(shù)據(jù),以便列元素連續(xù)存儲(chǔ)在內(nèi)存中。為了簡潔起見,本文未展示此類數(shù)據(jù)重排,如需參考可查看 SME 程序員指南。
此示例包含三個(gè)嵌套循環(huán):
最外層循環(huán)用于迭代處理結(jié)果矩陣的行。
中間層循環(huán)用于迭代處理結(jié)果矩陣的列。
最內(nèi)層循環(huán)用于迭代處理 K 維度,通過對(duì)乘積求和生成結(jié)果矩陣元素。
使用 ld1w 指令將矩陣數(shù)據(jù)從內(nèi)存加載到 ZA 存儲(chǔ)。外積計(jì)算使用 fmopa 指令。每個(gè) fmopa 指令讀取兩個(gè) SVE Z 輸入向量,并使用結(jié)果更新整個(gè) SME ZA tile。二維 predication 確保不超出矩陣的邊界。最后,st1w 指令將結(jié)果從 ZA 存儲(chǔ)寫入內(nèi)存。
評(píng)論