Armv9 技術(shù)講堂 | SME 指令介紹
加入技術(shù)交流群
Arm 可伸縮矩陣擴展 (SME) 作為 Armv9 架構(gòu)中的一項創(chuàng)新特性，旨在滿足當前日益復(fù)雜和高能耗的人工智能 (AI) 和機器視覺 (ML) 應(yīng)用需求。除了加速現(xiàn)今的 AI，SME 也提供了在 Arm 架構(gòu)上處理不斷更新的生成式 AI 應(yīng)用的靈活性。在 上一篇內(nèi)容 中，Arm 技術(shù)專家為大家簡要介紹了 SME，本周我們將帶各位更為詳細地來了解 SME 的指令，助力你在應(yīng)用中高效使用 SME！
本文引用地址：http://butianyuan.cn/article/202408/461972.htm操作 ZA 存儲的 SME 指令主要包括：計算兩個向量外積，并累加或累減，將結(jié)果放入一個 ZA tile 的指令；內(nèi)存與 ZA tile 行或列之間的存取操作指令，以及 SVE Z 寄存器和 ZA tile 行或列之間的移動指令；水平或垂直方向上，一個 Z 向量與 ZA tile 的加法指令；將一個標量寄存器加上 Streaming SVE 模式的向量長度倍數(shù)的指令。
外積并累加或累減指令
為了幫助理解外積并累加或累減指令，讓我們看看如何使用外積操作來做矩陣乘法。下圖展示了外積運算：
計算 a 和 b 兩個向量的外積，可以得到他們外積的結(jié)果，即矩陣 C。
現(xiàn)在我們看看如果進行矩陣 a 和 b 的矩陣乘：
這個矩陣乘可以通過計算兩次外積操作再累加兩個外積結(jié)果來實現(xiàn)，如下圖所示：
SME 為 8 位、16 位整數(shù)，以及 FP16、BF16、FP32 和 FP64 浮點數(shù)引入了高效的外積并累加或累減指令。這些指令計算放在兩個 Z 向量寄存器（Zn 和 Zm）里的向量的外積，并將這個外積結(jié)果累加或累減一個 ZA tile (ZAda) 中已有數(shù)據(jù)，將結(jié)果存入同一 ZA tile (ZAda)。每個源向量都可以被其相應(yīng)的控制 predicate 寄存器（Pn 和 Pm）獨立地 predicate。
FP32、FP64 外積并累加或累減指令
那些輸入向量和輸出數(shù)組有同樣數(shù)據(jù)類型（FP32 和 FP64）的指令相對直觀。下面這一例子展示了 FP32 類型的外積并累加或累減指令。
這個例子中，假設(shè) SVL 向量長度為 128 位，Zn.S 和 Zm.S 中存放了四個 FP32 數(shù)組成的向量，此指令計算 Zn.S 和 Zm.S 的外積，外積結(jié)果為圖中灰色的矩陣，然后將此外積結(jié)果累加或累減 ZAda.S 這個 ZA tile（矩陣）中原有的值，將結(jié)果存入同一 ZAda.S tile 中。
FP16、BF16、INT16、INT8、 I16I64 類型
的 外積并累加或累減指令
為了保持這些數(shù)據(jù)類型結(jié)果的精度，以及充分利用計算硬件資源和 ZA 存儲，這些指令會擴大計算結(jié)果數(shù)據(jù)類型，因此這些操作不像前面 FP32 和 FP64 類型指令那么直接。
BF16 指令計算兩個 BF16 的外積的和，擴大結(jié)果類型為 FP32，然后這些結(jié)構(gòu)被解構(gòu)性地從目標 tile 中加或減。
INT8 指令計算四個 INT8 的外積的和，擴大結(jié)果類型為 INT32，然后這些結(jié)構(gòu)被解構(gòu)性地從目標 tile 中加或減。
INT16 指令計算兩個 INT16 的外積的和，擴大結(jié)果類型為 INT32，然后這些結(jié)構(gòu)被解構(gòu)性地從目標 tile 中加或減。
FP16 指令計算兩個 FP16 的外積的和，擴大結(jié)果類型為 INT32，然后這些結(jié)構(gòu)被解構(gòu)性地從目標tile中加或減。
如果實現(xiàn)了 FEAT_SME_I16I64，I16I64 指令計算四個 INT16 的外積的和，擴大結(jié)果類型為 INT64，然后這些結(jié)構(gòu)被解構(gòu)性地從目標 tile 中加或減。
以下例子展示了 SVL 向量長度為 128 位的 INT8 UMOPA 指令進行的操作：
對于 UMPOA 指令，每個輸入向量 (Zn.B, Zm.B) 被當成一個有 4x4 元素的矩陣，可以看作是四個連續(xù)的元素組成的塊（如圖中紅線標出）被轉(zhuǎn)置了。
在這個例子中，因為 SVL 向量長度為 128 位，第一源向量 Zn.B 包含一個無符號 8 位整數(shù)的 4x4 子矩陣，第二源向量 Zm.B 包含一個無符號 8 位整數(shù)的 4x4 子矩陣，UMOPA 指令計算出 4x4 擴大了的 32 位整數(shù)外積的和，然后解構(gòu)性地加上目標 tile (ZAda) 中的整數(shù)。
更通用地說，這條無符號整型外積和并累加指令 (UMOPA) 將第一源向量中的子矩陣乘以第二源向量中的子矩陣。每個源向量包含一個 (SVL/32) x 4 的無符號 8 位整數(shù)的子矩陣。然后將得到的 (SVL/32) x (SVL/32) 擴大了的 32 位整數(shù)外積和解構(gòu)性地加上一個 32 位整數(shù)目標 tile。
下面的例子展示了 SVL 為 128 位的 BF16 BFMOPA 進行的操作：
這個例子中，因為 SVL 為 128 位，第一源向量 Zn.H 包含一個 BF16 浮點數(shù)的 4x2 子矩陣，它被擴大為單精度浮點數(shù)；第二源向量 Zm.H 包含一個 BF16 浮點數(shù)的 2x4 子矩陣，它被擴大為單精度浮點數(shù)；BFMOPA 指令計算出 4x4 單精度外積的和，然后解構(gòu)性地加上目標tile (ZAda) 中的單精度數(shù)。
更通用地說，這條指令 (BFMOPA) 擴大了存放在第一源里的 (SVL/32) x2 BF16 子矩陣的類型為單精度，擴大了存放在第二源里的 2x (SVL/32) BF16 子矩陣的類型為單精度，將這兩個子矩陣相乘。然后將得到的 (SVL/32) x (SVL/32) 單精度外積和解構(gòu)性地加上一個單精度目標 tile。
下表顯示了幾種數(shù)據(jù)類型和 SVL長度的一條外積并累加或累減指令所做的對應(yīng)數(shù)據(jù)類型的 MAC（乘累加）數(shù)量：
帶 Predication 的 SME 指令
每個源向量都可以被其相應(yīng)的控制 predicate 寄存器獨立地 predicate：
外積并累加或累減指令使用 Pn/M 和 Pn/M（沒有 /Z 形式）：inactive 的源元素被當成具有 0 值。
Slice 移動 (move) 指令使用 Pg/M：目標 slice 中 inactive 的元素保持不變。
Tile slice 加載 (load) 指令使用 Pg/Z：目標 tile slice 中的 inactive 元素被設(shè)置為 0。
Tile slice 存儲 (store) 指令使用 Pg：inactive 的元素不會寫入內(nèi)存。
Predication 讓矩陣的維數(shù)不是 SVL 的倍數(shù)的情況更容易處理。如下面這一例子所示：
輸入向量 Z0 被 P0 predicate，Z1 被 P1 predicate。在這個例子里，SVL 為 512 位，Z 寄存器包含 16 個 FP32 數(shù)組成的向量，P0 中最后兩個元素是 inactive 的，而 P1 中最后一個元素是 inactive 的。這條指令更新 ZA0.S 中 (16-2) x (16-1) 個 FP32 元素，因為使用了 Pn/M，ZA0.S 中剩下的元素保持不變。
下圖展示了更多的 predicated 外積并累加或累減的例子。圖中被劃線的文字表示被 inactive predicate 元素影響的計算部分。
ZA tile 與一個 Z 向量的加運算
SME 包括 ZA tile 的行或列都加上一個向量的指令
例如：
它進行以下操作：
這個 ADDHA 指令將源向量 Z1 中的每個元素加上 ZA0.S tile 每一水平 slice 的相應(yīng) active 元素。Tile 中元素被一對控制 predicate 進行 predicate。一個水平 slice 中的一個元素在下面情況下可以認為是 active：它在第二控制 predicate 對應(yīng)的元素是 true，并且它在第一 控制 predicate 對應(yīng)的水平 slice 行號也為 true，目標 tile 中 inactive 元素保持不變。
Tile 存取和移動指令
SME tile 存取和移動 (load, store, move) 指令可以：
從內(nèi)存讀取數(shù)據(jù)，放入 ZA tile 的行或列
將 ZA tile 的行或列寫入內(nèi)存
將 ZA tile 的行移動到 SVE Z 向量寄存器
將 SVE Z 向量寄存器移動到 ZA tile 行或列
Tile slice 存取指令
LD1B、LD1H、LD1S、LD1D 和 LD1Q 指令從內(nèi)存中連續(xù)讀取一些值到 ZA tile slice，數(shù)據(jù)類型分別是 8 位、16 位、32 位、64 位和 128 位元素。ST1B、ST1H、ST1S、ST1D 和 ST1Q 指令將 ZA tile slice 存入連續(xù)內(nèi)存中，數(shù)據(jù)類型分別是 8 位、16 位、32 位、64 位和 128 位元素。這些指令也用 predication 的支持，例如：
此 LD1B 指令執(zhí)行 predicated 的連續(xù)字節(jié)讀取，它從地址為 (X1+X2) 的內(nèi)存讀取數(shù)據(jù)到 ZA0 中行號為 (W0+imm) 的這個水平 tile slice 中。目標 tile slice 中 inactive 的元素被設(shè)置為 0。
此 ST1H 指令執(zhí)行 predicated 連續(xù)半字的存操作，它將 ZA1 中列號為 (W0+imm) 的垂直 tile slice 存到地址為 (X1+X2*2) 的內(nèi)存，tile slice 中 inactive 的元素不寫入內(nèi)存。
Tile slice 移動指令
MOV 指令（MOVA 指令的別名）將一個 Z 向量寄存器的值移動到一個 ZA tile slice，或?qū)⒁粋€ ZA tile slice 中的值移動到一個 Z 向量寄存器。這條指令操作帶指定元素大小的 ZA tile 的單個水平或垂直 tile slice。Slice 的行號/列號由 slice 的檢索寄存器加上立即數(shù)偏移指定。目標 slice 中 inactive 的元素保持不變。
例如：
或
此指令將向量寄存器 Z0.B 中的值移動到 ZA0H.B[W0, #imm] 這個水平 ZA tile slice 中，使用 P0 作為 predication 寄存器。目標 tile slice 中 inactive 的元素保持不變。
ZA array 向量存取指令
SME LDR 指令從內(nèi)存讀取數(shù)據(jù)到一個 ZA array 向量，SME STR 指令將一個 ZA array 向量中的值存入內(nèi)存。這些指令是不帶 predication 功能的。它們主要是為了軟件上下文切換時，對 ZA 存儲進行保存或恢復(fù)。SME LDR/STR 指令也可以在 Non-streaming SVE 模式下，當 PSTATE.ZA 使能的情況下使用。例如，下面的 STR 指令的 ZA array 向量是由一個向量選擇寄存器 Wv（標量寄存器 W）加上可選的立即數(shù) (Wv+Imm) 指定。訪問內(nèi)存的地址為：一個標量寄存器作為基礎(chǔ)，加上相同的可選立即數(shù)偏移乘以當前向量長度字節(jié)數(shù)。
ZA tile 清零指令
SME 清零 (ZERO) 指令可以清零一組 64 位 ZA tile：
清零指令可以清零多到八個名為 ZA0.D 到 ZA8.D 的 ZA tile，哪些 tile 要清零由指令中的 mask 指定，剩下的其他 tile 保持不變。這條指令也可以在 Non-streaming SVE 模式，當 PSTATE.ZA 使能的情況下使用。如果要清零整個 ZA array，可以使用一個指令別名，即 ZERO {ZA}。
新的 SVE2 指令
SME 構(gòu)架擴展加入了一些新的 SVE2 指令，這些指令也可以在 PE 實現(xiàn)了 SVE2，且處于 Non-streaming SVE 模式時使用。這些指令包括：
選擇一個 predicate 寄存器或是 all-false 的 predicate 選擇指令
翻轉(zhuǎn) (reverse) 64 位雙字元素的指令
有符號/無符號鉗位為更小/更大值向量的指令
PSEL 指令
PSEL 指令選擇一個 predicate 寄存器或是 all-false 到目標 predicate 寄存器，如下所示：
如果指令中第二源 predicate 寄存器 (Pm) 中指定的元素為 true，這條指令將第一源 predicate 寄存器 (Pn) 的內(nèi)容放到目標 predicate 寄存器 (Pd)，否則設(shè)置目標 predicate 寄存器的值全部為 false。例如以下指令，假設(shè) W12 的值為 0：
第二源 predicate 寄存器的 [W12+0] 即 [0] 個元素為 false，因此目標寄存器 P0 被設(shè)置為全 0 (all-false)，如下圖所示：
現(xiàn)在來看如下指令，仍然假設(shè) W12 的值為 0，但這次立即數(shù)偏移為 1：
第二源 predicate 寄存器的 [W12+1] 即 [1] 個元素為 true，因此選擇第一源 predicate 寄存器的值到目標寄存器 P0，如下圖所示：
評論