Armv9 技術(shù)講堂 | SME 詳解

Arm 技術(shù)專(zhuān)家將分上下兩期帶大家深入了解 SME 及其指令，為你詳細解析 SME 如何確保 Arm CPU 能夠高效地處理新的高能耗生成式 AI 工作負載，從而在數十億臺 Arm 技術(shù)驅動(dòng)的設備上提供更好的 AI 體驗！

Armv9-A 架構引入了在 Arm CPU 上加速和保護如大語(yǔ)言模型 (LLM) 這樣的高級生成式 AI 應用的功能。Arm SME 是為了滿(mǎn)足當前日益復雜和高能耗的 AI 和 ML 應用需求，創(chuàng )新性地設計的 CPU 功能。除了加速現今的 AI，SME 也提供了在 Arm 架構上處理不斷更新的生成式 AI 應用的靈活性。

SME 顯著(zhù)提升了在 Arm CPU 上處理現有 AI 和 ML 應用的能力，從而讓各類(lèi) AI 加持的設備和應用有更快、更及時(shí)的用戶(hù)體驗。它也可以加速很多利用矩陣運算的應用，例如 AR、VR 和圖像處理，它們在 AI 和 ML 扮演著(zhù)越來(lái)越重要的角色。

與 Arm CPU 可以處理各類(lèi)神經(jīng)網(wǎng)絡(luò )中諸多數據類(lèi)型類(lèi)似，SME 提供了不斷進(jìn)化和日益復雜的 AI 和 ML 應用需要的靈活性。這保證了 Arm 架構在 AI 快速演化的年代，仍然承擔 AI 計算重要職責。

Arm SME 是一個(gè)增強矩陣操作的架構擴展。SME 建立在 SVE2 的基礎之上，新增了高效處理矩陣的能力。其關(guān)鍵功能包括：

計算兩個(gè) SVE 向量的外積 (outer product)

矩陣塊 (tile) 的存儲

存取矩陣塊中的向量，向矩陣塊中插入向量和提取矩陣塊里的向量，包括 on-the-fly 矩陣轉置

Streaming SVE 模式

下表總結了 SME、SVE 和 SVE2 的主要功能：

SME 主要支持了以下新功能：

新的架構狀態(tài) (architectural state)，可以用來(lái)存儲二維矩陣塊。

新的 Streaming SVE 模式支持運行 SVE2 指令，Streaming SVE 模式下，SVE2  向量具有矩陣塊一樣寬度，通常比 Non-streaming SVE 模式下的向量要更寬。

新的外積指令用于計算兩個(gè)向量的外積，且可以進(jìn)行外積的累加或累減，將結果放入一個(gè)矩陣塊。

新的存取和移動(dòng) (load, store, move) 指令，用于從矩陣塊的一行或是一列讀取一個(gè)向量到向量寄存器，或是將一個(gè)向量寫(xiě)入矩陣塊的一行或是一列。

與 SVE2 一樣，SME 是一個(gè)可伸縮向量擴展，它同樣支持 Vector Length Agnostic (VLA), per-lane predication, predicate-driven loop control and management 等功能。

Streaming SVE 模式

為了支持 SME，Armv9 架構新增了一個(gè)操作模式 —— Streaming SVE 模式。在此模式下，包含了現有 SVE2 指令的一個(gè)子集（支持大多 SVE2 指令），并增加了新的 SME 特有的指令。

Streaming SVE 模式支持對大數據集進(jìn)行高吞吐 streaming 數據處理，被 stream 的數據通常有較簡(jiǎn)單的循環(huán)控制和有限的條件處理。

在 Non-streaming SVE 模式下，支持全部的 SVE2 指令集，處理一般代碼，通常處理復雜的數據結構和復雜的判斷。

大多數新加入的 SME 指令只能在 Streaming SVE 模式下使用。Streaming SVE 模式下的硬件向量長(cháng)度 (Streaming vector length, SVL) 可以和 Non-streaming SVE  模式的向量長(cháng)度 (Non-streaming vector length, NSVL) 不一樣。預期是，SVL 要比 NSVL 更長(cháng)或是相同，也就是 SVL >= NSVL。例如，NSVL 的長(cháng)度可以為 128 位，而 SVL 的長(cháng)度可以為 512 位。SME 的 SVL 可以是 128 位、256 位、512 位、1024 位 或是 2048 位。SVL 需要是 2 的次冪，而 NSVL 需要是 128 的整數倍。

與 SVE2 類(lèi)似，軟件可以控制 SMCR_ELx.LEN 寄存器位來(lái)設置 EL1、EL2、EL3 想用的生效 SVL 長(cháng)度（可以設置為比硬件支持的 SVL 更短）。

切換 SVE 模式

如果 CPU 硬件實(shí)現既支持 Streaming SVE 模式的 SME，又支持 Non-streaming SVE 模式的 SVE2，應用程序可以根據其需求動(dòng)態(tài)切換這兩個(gè)操作模式。

讓 SME 有一個(gè)獨立的操作模式，使 CPU 硬件實(shí)現可以為同一應用提供不同的向量長(cháng)度。比如 CPU 硬件實(shí)現可以選擇支持一個(gè)更長(cháng)的 Streaming SVE 模式向量長(cháng)度，它可以為 streaming 做硬件優(yōu)化，適用于高吞吐量的數據處理。

在同一應用中很容易動(dòng)態(tài)切換 Streaming SVE 和 Non-streaming SVE 模式。SME 引入了一些新的 PSTATE 位，PSTATE.{SM, ZA} 用于啟用和禁用 Streaming SVE 模式和 SME ZA 存儲：

SM：?jiǎn)⒂煤徒?Streaming SVE 模式

ZA：?jiǎn)⒂煤徒?ZA 存儲訪(fǎng)問(wèn)

可以通過(guò) MSR/MRS 指令操作 Streaming Vector Control Register (SVCR) 來(lái)設置和讀取 PSTATE.{SM, ZA}：

MSR SVCRSM, #

MSR SVCRZA, #

MSR SVCRSMZA, #

SMSTART 指令是設置 PSTATE.SM 和 PSTATE.ZA 的 MSR 指令的別名：

SMSTART：?jiǎn)⒂?Streaming SVE 模式和 ZA 存儲訪(fǎng)問(wèn)

SMSTAT SM：?jiǎn)⒂?Streaming SVE 模式

SMSTART ZA：?jiǎn)⒂?ZA 存儲訪(fǎng)問(wèn)

SMSTOP 指令是清除 PSTATE.SM 和 PSTATE.ZA 的 MSR 指令的別名：

SMSTOP：禁用 Streaming SVE 模式和 ZA 存儲訪(fǎng)問(wèn)

SMSTOP SM：禁用 Streaming SVE 模式

SMSTOP ZA：禁用 ZA 存儲訪(fǎng)問(wèn)

下圖展示了如何在同一應用內切換 Streaming SVE 和 Non-streaming SVE 模式：

SME 架構狀態(tài)

與 SVE2 類(lèi)似，在 Streaming SVE 模式，它有 Z0-Z31 向量寄存器，和 P0-P15 Predicate 寄存器。

物理上，SIMD (Neon)、浮點(diǎn) Vn、Qn、Dn、Sn、Hn、Bn 寄存器和與之對應的 Zn 向量寄存器重疊。也就是改寫(xiě) Vn、Qn、Dn、Sn、Hn、Bn 寄存器就是改寫(xiě) Zn 寄存器的低 bits。

架構上要求，當進(jìn)入 Streaming SVE 模式時(shí)（即 PSTATE.SM 由 0 變 1）和退出 Streaming SVE 模式時(shí)（即 PSTATE.SM 由 1 變 0），Zn 和 Pn 寄存器被硬件清零。

大多數原有的 SVE2 指令還能在 Streaming SVE 模式下使用，但使用的是 SVL 長(cháng)度的向量，與 NSVL 可能不一樣。生效的 SVL 長(cháng)度可以通過(guò) RDSVL 指令讀出來(lái)：

注意：因為 SME 支持 VLA，在 Streaming SVE 模式下，軟件很少需要讀 SVL 向量長(cháng)度，RDSVL 指令通常是在 CPU 核在 Non-streaming SVE 模式時(shí)，軟件讀 SVL 長(cháng)度。

ZA array

SME 引入了一個(gè)新的 ZA（Z Array, ZA 存儲），它是一個(gè)二維正方形數組，大小是 SVL x SVL。之所以叫 Z Array，也是因為它的行和列的長(cháng)度與 Streaming SVE 模式的 Zn 寄存器一致。

例如，如果 SVL 的長(cháng)度為 256 位（32 字節），那么 Zn 寄存器的長(cháng)度為 32 字節（256 位），ZA 的大小為 32 字節 x 32 字節（或 256 位 x 256 位）。

ZA array 可以通過(guò)以下方式訪(fǎng)問(wèn)：

ZA array 向量

ZA tiles

ZA tile slices

ZA array 向量訪(fǎng)問(wèn)

ZA array 的一行可以當成一個(gè) SVL 長(cháng)度的向量（ZA array 向量）來(lái)訪(fǎng)問(wèn)，這個(gè)向量可以放數據類(lèi)型長(cháng)度位 8 位、16 位、32 位、64 位、128 位的元素：

其中 B、H、S、D、Q 分別表示 8 位、16 位、32 位、64 位、128 位。

ZA array 向量的數量與 SVL 中的字節數相同，例如，如果 SVL 是 256 位（32 字節），那么 ZA array 向量的數量就是 32 個(gè) (0<= N <32)。

為了支持上下文切換，SME 引入了新的 LDR/STR 指令用于從內存中存取一個(gè) ZA array 向量：

LDR ZA[,], [{, #, MUL VL}]

STR ZA[, ], [{, #, MUL VL}]

ZA tiles

SME 支持對 ZA 進(jìn)行分塊，分成多個(gè) ZA tiles 進(jìn)行訪(fǎng)問(wèn)和使用。一個(gè) ZA tile 是在 ZA 中的正方形的二維子矩陣。一個(gè) ZA tile 的寬度總是 SVL，這與 ZA array 的寬度一致。

ZA 可以分成多少個(gè) ZA tiles 使用是由元素的數據類(lèi)型大小決定的：

當元素數據類(lèi)型為 8 位時(shí)，ZA 只能是 1 個(gè) ZA tile，ZA0.B

當元素數據類(lèi)型為 16 位時(shí)，ZA 可分為 2 個(gè) ZA tiles，ZA0.H-ZA1.H

當元素數據類(lèi)型為 32 位時(shí)，ZA 可分為 4 個(gè) ZA tiles，ZA0.S-ZA3.S

當元素數據類(lèi)型為 64 位時(shí)，ZA 可分為 8 個(gè) ZA tiles，ZA0.D-ZA7.D

當元素數據類(lèi)型為 128 位時(shí)，ZA 可分為 16 個(gè) ZA tiles，ZA0.Q-ZA15.Q

這樣做可以充分利用 ZA 存儲，例如，如果元素數據類(lèi)型為 32 位浮點(diǎn)數且 SVL 為 256 位（SVL 長(cháng)度可以放 8 個(gè) 32 位浮點(diǎn)數），那么 ZA 的大小為 32 x 32 字節 （可以放 32 行 x 8 列的二維浮點(diǎn)數數值），一個(gè) Z 寄存器可以放 8 個(gè)浮點(diǎn)數元素的向量。那么兩個(gè) Z 寄存器里的向量做外積運算，產(chǎn)生的外積結果是 8 行x 8 列的二維浮點(diǎn)數數組，這個(gè)外積只需要 1/4 的 ZA 存儲。為了充分利用，這樣 ZA 可以分成 4 個(gè) ZA tiles，分別從 ZA0.S 到 ZA3.S。

再例如，如果 SVL 為 256 位（32 字節），元素的數據類(lèi)型大小為 8 位，那么 ZA 可以看成為一個(gè)有 32 行，每行為 32 個(gè) 8 位的數組。

如果 SVL 同為 256 位（32 字節），但元素的數據類(lèi)型大小為 16 位，那么 ZA 可以看出是兩個(gè) ZA tiles， ZA0.H 和 ZA1.H，每個(gè) tile 由 16 行，每行有 16 個(gè) 16 位的數組組成。

ZA tile 的訪(fǎng)問(wèn)

一個(gè) ZA tile 可以作為一個(gè)整體來(lái)訪(fǎng)問(wèn)，也可以以一個(gè)個(gè) ZA tile slice 的方式訪(fǎng)問(wèn)。當作為一個(gè)整體訪(fǎng)問(wèn)時(shí)，指令可以使用 tile 的名字訪(fǎng)問(wèn)：

一個(gè) ZA tile slice 是由其 ZA tile 中水平方向或是垂直方向的連續元素組成的一維數組，即在 ZA tile 中的一行或是一列。

對一個(gè) ZA tile 的向量訪(fǎng)問(wèn)即是讀寫(xiě)一個(gè) ZA tile slice，

是水平或是垂直方向的 ZA tile slice 訪(fǎng)問(wèn)，由加在 ZA tile 名字后的 'H' 或是 'V' 后綴來(lái)表示

是哪一個(gè) ZA tile slice 由 ZA tile 名字后面的 slice 索引 '[N]' 來(lái)表示

例如，如果 SVL 為 128 位 且元素類(lèi)型的大小為 8 位，那么其水平和垂直的 ZA tile slices 可以由下圖展示：

ZA0V.B[0] 和 ZA0V.B[13] 訪(fǎng)問(wèn) ZA0.B ZA tile 垂直方向向量（即一列），ZA0H.B[0] 和 ZA0H.B[15] 訪(fǎng)問(wèn)ZA0.B ZA tile 水平方向向量（即一行）。如果 SVL 為 128 位且元素類(lèi)型的大小為 16 位，那么其水平和垂直的 ZA tile slices 可以由下圖展示：

事實(shí)上，為了實(shí)現高效的硬件對 ZA tile 和 tile slice 的訪(fǎng)問(wèn)，ZA tile 的 tile slice 在 ZA 中以交織方式存儲。例如，如果 SVL 是 256 位，元素數據類(lèi)型大小是 16 位，那么其水平 ZA tile slice 在 ZA 中的存儲可以展示為：

如果 SVL 同樣是 256 位，元素數據類(lèi)型大小是 32 位，那么其水平 ZA tile slice 在 ZA 中的存儲可以展示為：

下圖展示了不同的元素數據類(lèi)型大小的水平和垂直方向 ZA tile slice 的混合視圖:

圖中左邊的那些列顯示了 ZA 存儲的每行可以被各種數據類(lèi)型大小的 tile 水平 slice 訪(fǎng)問(wèn)的方式。 設 'SIZE' 為向量元素的類(lèi)型大小，B、H、S、D 和 Q 類(lèi)型的 SIZE 分別為 1、2、4、8、16。

設 'NUM_OF_ELEMENTS' 為可以放在向量里的元素個(gè)數，即（SVL 的字節數）/SIZE。那么，水平 tile slice，ZAnH.[m] 訪(fǎng)問(wèn) ZA 存儲的第 (m*SIZE+n) 行（完整一行）。例如，ZA2H.S[5] 訪(fǎng)問(wèn) ZA 存儲的第 (5*4+2) 行，即第 22 行。ZA6H.D[0] 訪(fǎng)問(wèn) ZA 存儲的第 (0*8+6) 行，即第 6 行。

垂直 tile slice，ZAnV.[m] 訪(fǎng)問(wèn) ZA 存儲中 (i*SIZE+n) 這些行里第 m 個(gè)元素（元素大小為 SIZE），由這些元素組成的向量，其中 i 為 0 到 (NUM_OF_ELEMNTS-1)。例如，ZA3V.S[4] 訪(fǎng)問(wèn) ZA 存儲中，由第 3、7、11、15、19、23、27、31 這些行中的第 4 個(gè)元素（元素大小為 4）組成的向量。

如果軟件中使用混合數據類(lèi)型大小或是水平/垂直方向 tile slice 時(shí)，需要小心處理 tile slice 間的重疊。

Steaming SVE 模式下支持的指令

有些 Neon/SVE2 指令在 Streaming SVE 模式受到影響：

有些 SVE/SVE2 指令變?yōu)榉欠▓绦?/span>

Gather/Scatter load/store SVE2 指令

使用 First Fault 寄存器的 SVE2 指令

大多的 Neon 指令變?yōu)?UNDEFINED

SME 主要增加以下指令：

矩陣外積并且累加或累減指令，包括 FMOPA、UMOPA、BFMOPA。這些指令利用，

SVE2 向量寄存器 (Z0-Z31) 作為外積運算的行和列輸入

ZA 存儲用來(lái)保存二維的矩陣塊輸出

將 SVE2 Z 向量與 ZA 的行或列做加法運算的指令

對 ZA tiles 的清零操作指令

增加了一些在 Streaming SVE 和 Non-streaming SVE 模式下都能使用的指令