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

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

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

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

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

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

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

矩陣塊 (tile) 的存儲(chǔ)

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

Streaming SVE 模式

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

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

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

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

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

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

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

Streaming SVE 模式

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

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

在 Non-streaming SVE 模式下，支持全部的 SVE2 指令集，處理一般代碼，通常處理復(fù)雜的數(shù)據(jù)結(jié)構(gòu)和復(fù)雜的判斷。

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

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

切換 SVE 模式

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

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

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

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

ZA：?jiǎn)⒂煤徒?ZA 存儲(chǔ)訪問

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

MSR SVCRSM, #

MSR SVCRZA, #

MSR SVCRSMZA, #

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

SMSTART：?jiǎn)⒂?Streaming SVE 模式和 ZA 存儲(chǔ)訪問

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

SMSTART ZA：?jiǎn)⒂?ZA 存儲(chǔ)訪問

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

SMSTOP：禁用 Streaming SVE 模式和 ZA 存儲(chǔ)訪問

SMSTOP SM：禁用 Streaming SVE 模式

SMSTOP ZA：禁用 ZA 存儲(chǔ)訪問

下圖展示了如何在同一應(yīng)用內(nèi)切換 Streaming SVE 和 Non-streaming SVE 模式：

SME 架構(gòu)狀態(tài)

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

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

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

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

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

ZA array

SME 引入了一個(gè)新的 ZA（Z Array, ZA 存儲(chǔ)），它是一個(gè)二維正方形數(shù)組，大小是 SVL x SVL。之所以叫 Z Array，也是因?yàn)樗男泻土械拈L(zhǎng)度與 Streaming SVE 模式的 Zn 寄存器一致。

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

ZA array 可以通過以下方式訪問：

ZA array 向量

ZA tiles

ZA tile slices

ZA array 向量訪問

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

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

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

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

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

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

ZA tiles

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

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

當(dāng)元素?cái)?shù)據(jù)類型為 8 位時(shí)，ZA 只能是 1 個(gè) ZA tile，ZA0.B

當(dāng)元素?cái)?shù)據(jù)類型為 16 位時(shí)，ZA 可分為 2 個(gè) ZA tiles，ZA0.H-ZA1.H

當(dāng)元素?cái)?shù)據(jù)類型為 32 位時(shí)，ZA 可分為 4 個(gè) ZA tiles，ZA0.S-ZA3.S

當(dāng)元素?cái)?shù)據(jù)類型為 64 位時(shí)，ZA 可分為 8 個(gè) ZA tiles，ZA0.D-ZA7.D

當(dāng)元素?cái)?shù)據(jù)類型為 128 位時(shí)，ZA 可分為 16 個(gè) ZA tiles，ZA0.Q-ZA15.Q

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

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

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

ZA tile 的訪問

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

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

對(duì)一個(gè) ZA tile 的向量訪問即是讀寫一個(gè) ZA tile slice，

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

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

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

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

事實(shí)上，為了實(shí)現(xiàn)高效的硬件對(duì) ZA tile 和 tile slice 的訪問，ZA tile 的 tile slice 在 ZA 中以交織方式存儲(chǔ)。例如，如果 SVL 是 256 位，元素?cái)?shù)據(jù)類型大小是 16 位，那么其水平 ZA tile slice 在 ZA 中的存儲(chǔ)可以展示為：

如果 SVL 同樣是 256 位，元素?cái)?shù)據(jù)類型大小是 32 位，那么其水平 ZA tile slice 在 ZA 中的存儲(chǔ)可以展示為：

下圖展示了不同的元素?cái)?shù)據(jù)類型大小的水平和垂直方向 ZA tile slice 的混合視圖:

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

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

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

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

Steaming SVE 模式下支持的指令

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

有些 SVE/SVE2 指令變?yōu)榉欠▓?zhí)行

Gather/Scatter load/store SVE2 指令

使用 First Fault 寄存器的 SVE2 指令

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

SME 主要增加以下指令：

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

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

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

將 SVE2 Z 向量與 ZA 的行或列做加法運(yùn)算的指令

對(duì) ZA tiles 的清零操作指令

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