Armv9 技術(shù)講堂 | SME 指令介紹
Arm 可伸縮矩陣擴展 (SME) 作為 Armv9 架構(gòu)中的一項創(chuàng)新特性,旨在滿足當(dāng)前日益復(fù)雜和高能耗的人工智能 (AI) 和機器視覺 (ML) 應(yīng)用需求。除了加速現(xiàn)今的 AI,SME 也提供了在 Arm 架構(gòu)上處理不斷更新的生成式 AI 應(yīng)用的靈活性。在 上一篇內(nèi)容 中,Arm 技術(shù)專家為大家簡要介紹了 SME,本周我們將帶各位更為詳細(xì)地來了解 SME 的指令,助力你在應(yīng)用中高效使用 SME!
本文引用地址:http://www.ex-cimer.com/article/202408/461972.htm操作 ZA 存儲的 SME 指令主要包括:計算兩個向量外積,并累加或累減,將結(jié)果放入一個 ZA tile 的指令;內(nèi)存與 ZA tile 行或列之間的存取操作指令,以及 SVE Z 寄存器和 ZA tile 行或列之間的移動指令;水平或垂直方向上,一個 Z 向量與 ZA tile 的加法指令;將一個標(biāo)量寄存器加上 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)性地從目標(biāo) tile 中加或減。
INT8 指令計算四個 INT8 的外積的和,擴大結(jié)果類型為 INT32,然后這些結(jié)構(gòu)被解構(gòu)性地從目標(biāo) tile 中加或減。
INT16 指令計算兩個 INT16 的外積的和,擴大結(jié)果類型為 INT32,然后這些結(jié)構(gòu)被解構(gòu)性地從目標(biāo) tile 中加或減。
FP16 指令計算兩個 FP16 的外積的和,擴大結(jié)果類型為 INT32,然后這些結(jié)構(gòu)被解構(gòu)性地從目標(biāo)tile中加或減。
如果實現(xiàn)了 FEAT_SME_I16I64,I16I64 指令計算四個 INT16 的外積的和,擴大結(jié)果類型為 INT64,然后這些結(jié)構(gòu)被解構(gòu)性地從目標(biāo) tile 中加或減。
以下例子展示了 SVL 向量長度為 128 位的 INT8 UMOPA 指令進行的操作:
對于 UMPOA 指令,每個輸入向量 (Zn.B, Zm.B) 被當(dāng)成一個有 4x4 元素的矩陣,可以看作是四個連續(xù)的元素組成的塊(如圖中紅線標(biāo)出)被轉(zhuǎn)置了。
在這個例子中,因為 SVL 向量長度為 128 位,第一源向量 Zn.B 包含一個無符號 8 位整數(shù)的 4x4 子矩陣,第二源向量 Zm.B 包含一個無符號 8 位整數(shù)的 4x4 子矩陣,UMOPA 指令計算出 4x4 擴大了的 32 位整數(shù)外積的和,然后解構(gòu)性地加上目標(biāo) tile (ZAda) 中的整數(shù)。
更通用地說,這條無符號整型外積和并累加指令 (UMOPA) 將第一源向量中的子矩陣乘以第二源向量中的子矩陣。每個源向量包含一個 (SVL/32) x 4 的無符號 8 位整數(shù)的子矩陣。然后將得到的 (SVL/32) x (SVL/32) 擴大了的 32 位整數(shù)外積和解構(gòu)性地加上一個 32 位整數(shù)目標(biāo) tile。
下面的例子展示了 SVL 為 128 位的 BF16 BFMOPA 進行的操作:
這個例子中,因為 SVL 為 128 位,第一源向量 Zn.H 包含一個 BF16 浮點數(shù)的 4x2 子矩陣,它被擴大為單精度浮點數(shù);第二源向量 Zm.H 包含一個 BF16 浮點數(shù)的 2x4 子矩陣,它被擴大為單精度浮點數(shù);BFMOPA 指令計算出 4x4 單精度外積的和,然后解構(gòu)性地加上目標(biāo)tile (ZAda) 中的單精度數(shù)。
更通用地說,這條指令 (BFMOPA) 擴大了存放在第一源里的 (SVL/32) x2 BF16 子矩陣的類型為單精度,擴大了存放在第二源里的 2x (SVL/32) BF16 子矩陣的類型為單精度,將這兩個子矩陣相乘。然后將得到的 (SVL/32) x (SVL/32) 單精度外積和解構(gòu)性地加上一個單精度目標(biāo) tile。
下表顯示了幾種數(shù)據(jù)類型和 SVL長度的一條外積并累加或累減指令所做的對應(yīng)數(shù)據(jù)類型的 MAC(乘累加)數(shù)量:
帶 Predication 的 SME 指令
每個源向量都可以被其相應(yīng)的控制 predicate 寄存器獨立地 predicate:
外積并累加或累減指令使用 Pn/M 和 Pn/M(沒有 /Z 形式):inactive 的源元素被當(dāng)成具有 0 值。
Slice 移動 (move) 指令使用 Pg/M:目標(biāo) slice 中 inactive 的元素保持不變。
Tile slice 加載 (load) 指令使用 Pg/Z:目標(biāo) 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 中的一個元素在下面情況下可以認(rèn)為是 active:它在第二控制 predicate 對應(yīng)的元素是 true,并且它在第一 控制 predicate 對應(yīng)的水平 slice 行號也為 true,目標(biāo) 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 中。目標(biāo) 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ù)偏移指定。目標(biāo) slice 中 inactive 的元素保持不變。
例如:
或
此指令將向量寄存器 Z0.B 中的值移動到 ZA0H.B[W0, #imm] 這個水平 ZA tile slice 中,使用 P0 作為 predication 寄存器。目標(biāo) 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 模式下,當(dāng) PSTATE.ZA 使能的情況下使用。例如,下面的 STR 指令的 ZA array 向量是由一個向量選擇寄存器 Wv(標(biāo)量寄存器 W)加上可選的立即數(shù) (Wv+Imm) 指定。訪問內(nèi)存的地址為:一個標(biāo)量寄存器作為基礎(chǔ),加上相同的可選立即數(shù)偏移乘以當(dāng)前向量長度字節(jié)數(shù)。
ZA tile 清零指令
SME 清零 (ZERO) 指令可以清零一組 64 位 ZA tile:
清零指令可以清零多到八個名為 ZA0.D 到 ZA8.D 的 ZA tile,哪些 tile 要清零由指令中的 mask 指定,剩下的其他 tile 保持不變。這條指令也可以在 Non-streaming SVE 模式,當(dāng) 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 到目標(biāo) predicate 寄存器,如下所示:
如果指令中第二源 predicate 寄存器 (Pm) 中指定的元素為 true,這條指令將第一源 predicate 寄存器 (Pn) 的內(nèi)容放到目標(biāo) predicate 寄存器 (Pd),否則設(shè)置目標(biāo) predicate 寄存器的值全部為 false。例如以下指令,假設(shè) W12 的值為 0:
第二源 predicate 寄存器的 [W12+0] 即 [0] 個元素為 false,因此目標(biāo)寄存器 P0 被設(shè)置為全 0 (all-false),如下圖所示:
現(xiàn)在來看如下指令,仍然假設(shè) W12 的值為 0,但這次立即數(shù)偏移為 1:
第二源 predicate 寄存器的 [W12+1] 即 [1] 個元素為 true,因此選擇第一源 predicate 寄存器的值到目標(biāo)寄存器 P0,如下圖所示:
評論