Armv9 技術(shù)講堂 | SME 指令介紹

操作 ZA 存儲(chǔ)的 SME 指令主要包括：計(jì)算兩個(gè)向量外積，并累加或累減，將結(jié)果放入一個(gè) ZA tile 的指令；內(nèi)存與 ZA tile 行或列之間的存取操作指令，以及 SVE Z 寄存器和 ZA tile 行或列之間的移動(dòng)指令；水平或垂直方向上，一個(gè) Z 向量與 ZA tile 的加法指令；將一個(gè)標(biāo)量寄存器加上 Streaming SVE 模式的向量長(zhǎng)度倍數(shù)的指令。

為了幫助理解外積并累加或累減指令，讓我們看看如何使用外積操作來(lái)做矩陣乘法。下圖展示了外積運(yùn)算：

計(jì)算 a 和 b 兩個(gè)向量的外積，可以得到他們外積的結(jié)果，即矩陣 C。

現(xiàn)在我們看看如果進(jìn)行矩陣 a 和 b 的矩陣乘：

這個(gè)矩陣乘可以通過(guò)計(jì)算兩次外積操作再累加兩個(gè)外積結(jié)果來(lái)實(shí)現(xiàn)，如下圖所示：

SME 為 8 位、16 位整數(shù)，以及 FP16、BF16、FP32 和 FP64 浮點(diǎn)數(shù)引入了高效的外積并累加或累減指令。這些指令計(jì)算放在兩個(gè) Z 向量寄存器（Zn 和 Zm）里的向量的外積，并將這個(gè)外積結(jié)果累加或累減一個(gè) ZA tile (ZAda) 中已有數(shù)據(jù)，將結(jié)果存入同一 ZA tile (ZAda)。每個(gè)源向量都可以被其相應(yīng)的控制 predicate 寄存器（Pn 和 Pm）獨(dú)立地 predicate。

那些輸入向量和輸出數(shù)組有同樣數(shù)據(jù)類型（FP32 和 FP64）的指令相對(duì)直觀。下面這一例子展示了 FP32 類型的外積并累加或累減指令。

這個(gè)例子中，假設(shè) SVL 向量長(zhǎng)度為 128 位，Zn.S 和 Zm.S 中存放了四個(gè) FP32 數(shù)組成的向量，此指令計(jì)算 Zn.S 和 Zm.S 的外積，外積結(jié)果為圖中灰色的矩陣，然后將此外積結(jié)果累加或累減 ZAda.S 這個(gè) ZA tile（矩陣）中原有的值，將結(jié)果存入同一 ZAda.S tile 中。

為了保持這些數(shù)據(jù)類型結(jié)果的精度，以及充分利用計(jì)算硬件資源和 ZA 存儲(chǔ)，這些指令會(huì)擴(kuò)大計(jì)算結(jié)果數(shù)據(jù)類型，因此這些操作不像前面 FP32 和 FP64 類型指令那么直接。

以下例子展示了 SVL 向量長(zhǎng)度為 128 位的 INT8 UMOPA 指令進(jìn)行的操作：

在這個(gè)例子中，因?yàn)?SVL 向量長(zhǎng)度為 128 位，第一源向量 Zn.B 包含一個(gè)無(wú)符號(hào) 8 位整數(shù)的 4x4 子矩陣，第二源向量 Zm.B 包含一個(gè)無(wú)符號(hào) 8 位整數(shù)的 4x4 子矩陣，UMOPA 指令計(jì)算出 4x4 擴(kuò)大了的 32 位整數(shù)外積的和，然后解構(gòu)性地加上目標(biāo) tile (ZAda) 中的整數(shù)。

更通用地說(shuō)，這條無(wú)符號(hào)整型外積和并累加指令 (UMOPA) 將第一源向量中的子矩陣乘以第二源向量中的子矩陣。每個(gè)源向量包含一個(gè) (SVL/32) x 4 的無(wú)符號(hào) 8 位整數(shù)的子矩陣。然后將得到的 (SVL/32) x (SVL/32) 擴(kuò)大了的 32 位整數(shù)外積和解構(gòu)性地加上一個(gè) 32 位整數(shù)目標(biāo) tile。

下面的例子展示了 SVL 為 128 位的 BF16 BFMOPA 進(jìn)行的操作：

更通用地說(shuō)，這條指令 (BFMOPA) 擴(kuò)大了存放在第一源里的 (SVL/32) x2 BF16 子矩陣的類型為單精度，擴(kuò)大了存放在第二源里的 2x (SVL/32) BF16 子矩陣的類型為單精度，將這兩個(gè)子矩陣相乘。然后將得到的 (SVL/32) x (SVL/32) 單精度外積和解構(gòu)性地加上一個(gè)單精度目標(biāo) tile。

下表顯示了幾種數(shù)據(jù)類型和 SVL長(zhǎng)度的一條外積并累加或累減指令所做的對(duì)應(yīng)數(shù)據(jù)類型的 MAC（乘累加）數(shù)量：

每個(gè)源向量都可以被其相應(yīng)的控制 predicate 寄存器獨(dú)立地 predicate：

Predication 讓矩陣的維數(shù)不是 SVL 的倍數(shù)的情況更容易處理。如下面這一例子所示：

輸入向量 Z0 被 P0 predicate，Z1 被 P1 predicate。在這個(gè)例子里，SVL 為 512 位，Z 寄存器包含 16 個(gè) FP32 數(shù)組成的向量，P0 中最后兩個(gè)元素是 inactive 的，而 P1 中最后一個(gè)元素是 inactive 的。這條指令更新 ZA0.S 中 (16-2) x (16-1) 個(gè) FP32 元素，因?yàn)槭褂昧?Pn/M，ZA0.S 中剩下的元素保持不變。

下圖展示了更多的 predicated 外積并累加或累減的例子。圖中被劃線的文字表示被 inactive predicate 元素影響的計(jì)算部分。

SME 包括 ZA tile 的行或列都加上一個(gè)向量的指令

例如：

它進(jìn)行以下操作：

這個(gè) ADDHA 指令將源向量 Z1 中的每個(gè)元素加上 ZA0.S tile 每一水平 slice 的相應(yīng) active 元素。Tile 中元素被一對(duì)控制 predicate 進(jìn)行 predicate。一個(gè)水平 slice 中的一個(gè)元素在下面情況下可以認(rèn)為是 active：它在第二控制 predicate 對(duì)應(yīng)的元素是 true，并且它在第一 控制 predicate 對(duì)應(yīng)的水平 slice 行號(hào)也為 true，目標(biāo) tile 中 inactive 元素保持不變。

SME tile 存取和移動(dòng) (load, store, move) 指令可以：

此 LD1B 指令執(zhí)行 predicated 的連續(xù)字節(jié)讀取，它從地址為 (X1+X2) 的內(nèi)存讀取數(shù)據(jù)到 ZA0 中行號(hào)為 (W0+imm) 的這個(gè)水平 tile slice 中。目標(biāo) tile slice 中 inactive 的元素被設(shè)置為 0。

此 ST1H 指令執(zhí)行 predicated 連續(xù)半字的存操作，它將 ZA1 中列號(hào)為 (W0+imm) 的垂直 tile slice 存到地址為 (X1+X2*2) 的內(nèi)存，tile slice 中 inactive 的元素不寫入內(nèi)存。

例如：

或

SME LDR 指令從內(nèi)存讀取數(shù)據(jù)到一個(gè) ZA array 向量，SME STR 指令將一個(gè) ZA array 向量中的值存入內(nèi)存。這些指令是不帶 predication 功能的。它們主要是為了軟件上下文切換時(shí)，對(duì) ZA  存儲(chǔ)進(jìn)行保存或恢復(fù)。SME LDR/STR 指令也可以在 Non-streaming SVE 模式下，當(dāng) PSTATE.ZA 使能的情況下使用。例如，下面的 STR 指令的 ZA array 向量是由一個(gè)向量選擇寄存器 Wv（標(biāo)量寄存器 W）加上可選的立即數(shù) (Wv+Imm) 指定。訪問(wèn)內(nèi)存的地址為：一個(gè)標(biāo)量寄存器作為基礎(chǔ)，加上相同的可選立即數(shù)偏移乘以當(dāng)前向量長(zhǎng)度字節(jié)數(shù)。

SME 清零 (ZERO) 指令可以清零一組 64 位 ZA tile：

清零指令可以清零多到八個(gè)名為 ZA0.D 到 ZA8.D 的 ZA tile，哪些 tile 要清零由指令中的 mask 指定，剩下的其他 tile 保持不變。這條指令也可以在 Non-streaming SVE 模式，當(dāng) PSTATE.ZA 使能的情況下使用。如果要清零整個(gè) ZA array，可以使用一個(gè)指令別名，即 ZERO {ZA}。

SME 構(gòu)架擴(kuò)展加入了一些新的 SVE2 指令，這些指令也可以在 PE 實(shí)現(xiàn)了 SVE2，且處于 Non-streaming SVE 模式時(shí)使用。這些指令包括：

PSEL 指令選擇一個(gè) 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] 個(gè)元素為 false，因此目標(biāo)寄存器 P0 被設(shè)置為全 0 (all-false)，如下圖所示：

現(xiàn)在來(lái)看如下指令，仍然假設(shè) W12 的值為 0，但這次立即數(shù)偏移為 1：