Armv9 技術(shù)講堂 | Neon、SVE 和 SME 實(shí)現(xiàn)矩陣-矩陣乘法的比較

Arm 始終專注于架構(gòu)演進(jìn)，確保生態(tài)系統(tǒng)能夠適應(yīng)未來(lái)的技術(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)用中帶來(lái)速度更快、響應(yīng)更靈敏的用戶體驗(yàn)。

在此前的內(nèi)容中，Arm 技術(shù)專家為大家簡(jiǎn)要介紹了 SME 和 SME 指令 ，本期將帶你了解如何使用 Neon、可伸縮向量擴(kuò)展 (SVE) 和 SME 這三種不同的 Arm 技術(shù)實(shí)現(xiàn)相同的矩陣-矩陣乘法算法。這三個(gè)示例展現(xiàn)了這些技術(shù)之間的關(guān)鍵差異，為開(kāi)發(fā)者將代碼從 Neon、SVE/SVE2 移植到 SME/SME2 提供指導(dǎo)。

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，它們提供可變長(zhǎng)度寄存器。寄存器的大小由實(shí)現(xiàn)方案定義，從 128 位到 2048 位寄存器不等。這意味著程序員不知道可用寄存器的大小，因此必須將代碼編寫為與向量長(zhǎng)度無(wú)關(guān)。因此，每條指令處理的數(shù)據(jù)值的數(shù)量不是固定的，而是可變的。

2021 年 Armv9-A 中引入了 SME 和 SME2，也提供可變長(zhǎng)度寄存器。SME 引入了兩個(gè)關(guān)鍵的新架構(gòu)特性：Streaming SVE 模式和 ZA 存儲(chǔ)。Streaming SVE 模式是一種高吞吐量的矩陣數(shù)據(jù)處理模式，ZA 存儲(chǔ)則是一種專用的二維數(shù)組，可以方便地進(jìn)行常見(jià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 語(yǔ)言而不是匯編語(yǔ)言編寫整個(gè)程序。

本文中的所有三個(gè)示例都能夠?qū)崿F(xiàn)矩陣-矩陣乘法。矩陣乘法需要接收兩個(gè)輸入矩陣，通過(guò)將第一個(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 內(nèi)置函數(shù)來(lái)執(zhí)行矩陣-矩陣乘法。相關(guān)代碼執(zhí)行以下操作：

以下是使用 Neon 內(nèi)置函數(shù)的示例代碼：

此示例使用 4 x 4 的固定塊大小。這意味著輸入矩陣在兩個(gè)維度上都必須是四的倍數(shù)。可以通過(guò)用零填充矩陣來(lái)處理其他大小的矩陣。

此示例使用 SVE2 內(nèi)置函數(shù)來(lái)執(zhí)行矩陣-矩陣乘法。

Neon 示例和 SVE2 示例之間的主要區(qū)別在于 SVE2 使用可變長(zhǎng)度向量。Neon 示例可以使用 4 x 4 的固定塊大小來(lái)匹配 Neon 寄存器中的四個(gè) 32 位值，但程序員要到運(yùn)行時(shí)才能知道 SVE2 寄存器的大小。這意味著代碼必須與向量長(zhǎng)度無(wú)關(guān)。此示例使用 predication 來(lái)控制 SVE2 內(nèi)置函數(shù)操作的數(shù)據(jù)值的數(shù)量。這意味著無(wú)論實(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(wú)法在運(yùn)行前知道 SVE2 寄存器的大小。

相關(guān)代碼執(zhí)行以下操作：

以下是使用 SVE2 內(nèi)置函數(shù)的示例代碼：

此示例使用以下 SVE2 代碼特性：

此示例使用 SME2 匯編指令來(lái)執(zhí)行矩陣-矩陣乘法。SME2 示例與其他示例的區(qū)別如下：

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)存中。為了簡(jiǎn)潔起見(jiàn)，本文未展示此類數(shù)據(jù)重排，如需參考可查看 SME 程序員指南。

此示例包含三個(gè)嵌套循環(huán)：