LLVM 中對 AArch64 可擴展矩陣擴充功能的支援

1. 簡介

AArch64 SME ACLE 為使用者提供了許多屬性來控制 PSTATE.SM 和 PSTATE.ZA。AArch64 SME ABI 描述了當至少一個函數使用 PSTATE.SM 或 PSTATE.ZA 時,函數之間呼叫的需求。

本文件描述了 SME ACLE 屬性如何對應到 LLVM IR 屬性,以及 LLVM 如何降低這些屬性以實現 ABI 的規則和需求。

以下我們描述 LLVM IR 屬性及其與 C/C++ 層級 ACLE 屬性的關係。

aarch64_pstate_sm_enabled

用於具有 __arm_streaming 的函數

aarch64_pstate_sm_compatible

用於具有 __arm_streaming_compatible 的函數

aarch64_pstate_sm_body

用於具有 __arm_locally_streaming 的函數,且僅在函數定義 (而非宣告) 上有效

aarch64_new_za

用於具有 __arm_new("za") 的函數

aarch64_in_za

用於具有 __arm_in("za") 的函數

aarch64_out_za

用於具有 __arm_out("za") 的函數

aarch64_inout_za

用於具有 __arm_inout("za") 的函數

aarch64_preserves_za

用於具有 __arm_preserves("za") 的函數

aarch64_expanded_pstate_za

用於具有 __arm_new_za 的函數

Clang 必須確保上述屬性被添加到函數的宣告/定義以及其呼叫點。這對於呼叫帶屬性的函數指標非常重要,因為在這種情況下沒有可用的定義或宣告。

2. 處理 PSTATE.SM

當變更 PSTATE.SM 時,FP/向量運算的執行可能會轉移到另一個處理單元。這有三個重要的影響:

  • 執行時 SVE 向量長度可能會改變。

  • FP/AdvSIMD/SVE 暫存器的內容會被歸零。

  • 允許的指令集會改變。

這導致 IR 和優化方面的一些限制。例如,在可能以不同的 PSTATE.SM 值運作的函數之間共享向量長度相關的狀態是未定義行為。前端在產生 LLVM IR 時必須遵守這些限制。

即使執行時 SVE 向量長度可能會改變,但為了 LLVM IR 和幾乎所有 CodeGen 部分的目的,我們可以假設 vscale 的執行時值不會改變。如果我們讓編譯器在呼叫邊界周圍插入適當的 smstartsmstop 指令,那麼可以減輕對 SVE 狀態的影響。透過將狀態變更限制在呼叫周圍非常短暫的時間窗口內,我們可以控制操作的排程方式以及在狀態轉換之間如何保留存活值。

為了在這個粒度層級控制 PSTATE.SM,我們使用函數和呼叫點屬性,而不是內建函數。

屬性限制

  • 將可擴展向量物件 (或其指標) 傳遞或從可能使用不同 SVE 向量長度的函數返回是未定義行為。這包括具有非串流介面但標記為 aarch64_pstate_sm_body 的函數。

  • 不允許函數同時使用 aarch64_pstate_sm_compatibleaarch64_pstate_sm_enabled 進行裝飾。

  • 不允許函數使用以下屬性中的多個進行裝飾:aarch64_new_zaaarch64_in_zaaarch64_out_zaaarch64_inout_zaaarch64_preserves_za

這些限制也適用於更高等級的 SME ACLE,這表示我們可以在 Clang 中發出診斷訊息,以向使用者指出不正確的行為。

編譯器插入的串流模式變更

下表描述了當在具有不同屬性的函數之間進行呼叫時,編譯器必須考慮的 PSTATE.SM 轉換。在本表中,我們使用以下縮寫:

N

具有正常介面的函數 (進入時 PSTATE.SM=0,返回時 PSTATE.SM=0)

S

具有串流介面的函數 (進入時 PSTATE.SM=1,返回時 PSTATE.SM=1)

SC

具有串流相容介面的函數 (進入時 PSTATE.SM 可以是 0 或 1,返回時保持不變)。

具有 __attribute__((arm_locally_streaming)) 的函數不包含在此表中,因為對於呼叫者而言,該屬性與 'streaming' 同義,而對於被呼叫者而言,這僅僅是一個實作細節,明確地不向呼叫者公開。

表 4 具有不同屬性的函數呼叫組合

呼叫前

呼叫後

例外後

N

N

N

S

SMSTART

SMSTOP

N

SC

S

N

SMSTOP

SMSTART

SMSTART

S

S

SMSTART

S

SC

SMSTART

SC

N

如果呼叫前 PSTATE.SM 為 1,則 SMSTOP

如果呼叫前 PSTATE.SM 為 1,則 SMSTART

如果呼叫前 PSTATE.SM 為 1,則 SMSTART

SC

S

如果呼叫前 PSTATE.SM 為 0,則 SMSTART

如果呼叫前 PSTATE.SM 為 0,則 SMSTOP

如果呼叫前 PSTATE.SM 為 1,則 SMSTART

SC

SC

如果呼叫前 PSTATE.SM 為 1,則 SMSTART

由於變更 PSTATE.SM 會將 FP/向量暫存器歸零,因此最好在暫存器分配之前發出 smstartsmstop 指令,以便暫存器分配器可以在模式變更前後溢出/重新載入暫存器。

編譯器也應該有足夠的資訊來判斷哪些操作是呼叫/函數的引數/結果的一部分,以及哪些操作是函數主體的一部分,以便它可以將模式變更放置在完全正確的位置。執行此操作的合適位置似乎是 SelectionDAG,它在其中降低呼叫的引數/傳回值以實作指定的呼叫慣例。SelectionDAG 提供 Chains 和 Glue 來指定操作順序,並初步控制指令的排程。

狀態保存範例

當從具有正常介面的函數傳遞和返回 float 值到/從具有串流介面的函數時,呼叫點將需要確保引數/結果暫存器被保留,並且在 smstart/smstop 和呼叫之間沒有排程其他程式碼。

define float @foo(float %f) nounwind {
  %res = call float @bar(float %f) "aarch64_pstate_sm_enabled"
  ret float %res
}

declare float @bar(float) "aarch64_pstate_sm_enabled"

程式需要保留浮點引數和傳回值在暫存器 s0 中的值。

foo:                                    // @foo
// %bb.0:
        stp     d15, d14, [sp, #-80]!           // 16-byte Folded Spill
        stp     d13, d12, [sp, #16]             // 16-byte Folded Spill
        stp     d11, d10, [sp, #32]             // 16-byte Folded Spill
        stp     d9, d8, [sp, #48]               // 16-byte Folded Spill
        str     x30, [sp, #64]                  // 8-byte Folded Spill
        str     s0, [sp, #76]                   // 4-byte Folded Spill
        smstart sm
        ldr     s0, [sp, #76]                   // 4-byte Folded Reload
        bl      bar
        str     s0, [sp, #76]                   // 4-byte Folded Spill
        smstop  sm
        ldp     d9, d8, [sp, #48]               // 16-byte Folded Reload
        ldp     d11, d10, [sp, #32]             // 16-byte Folded Reload
        ldp     d13, d12, [sp, #16]             // 16-byte Folded Reload
        ldr     s0, [sp, #76]                   // 4-byte Folded Reload
        ldr     x30, [sp, #64]                  // 8-byte Folded Reload
        ldp     d15, d14, [sp], #80             // 16-byte Folded Reload
        ret

在 ISD 節點上設定正確的暫存器遮罩,並在正確的位置插入 smstart/smstop 應該能確保正確完成此操作。

指令選擇節點

AArch64ISD::SMSTART Chain, [SM|ZA|Both], CurrentState, ExpectedState[, RegMask]
AArch64ISD::SMSTOP  Chain, [SM|ZA|Both], CurrentState, ExpectedState[, RegMask]

SMSTART/SMSTOP 節點採用 CurrentStateExpectedState 運算元,用於條件式 SMSTART/SMSTOP 的情況。僅當 CurrentState != ExpectedState 時才會執行指令。

CurrentStateExpectedState 可以在編譯時評估 (即它們都是常數) 時,則會發出無條件的 smstart/smstop 指令。否則,節點會匹配到一個 Pseudo 指令,該指令會擴展為比較/分支和一個 smstart/smstop。這對於實作從 SC -> NSC -> S 的轉換是必要的。

非鏈式函數呼叫

當具有 “aarch64_pstate_sm_enabled” 的函數呼叫不相容串流的函數時,編譯器必須在呼叫前插入 SMSTOP,並在呼叫後插入 SMSTOP。

如果被呼叫的函數是一個沒有副作用的內建函數,而該內建函數又被降低為函數呼叫 (例如 @llvm.cos()),則對 @llvm.cos() 的呼叫不屬於任何 Chain;它可以自由排程。

降低 Callsite 會建立一個小的節點鏈,其中:

  • 開始一個呼叫序列

  • 將輸入值從虛擬暫存器複製到 ABI 指定的實體暫存器

  • 執行分支和連結

  • 停止呼叫序列

  • 將輸出值從其物理暫存器複製到虛擬暫存器

當未使用呼叫點的 Chain 時,僅使用鏈式序列中的結果值,但 Chain 本身會被丟棄。

SMSTARTSMSTOP ISD 節點返回一個 Chain,但沒有實際值,因此當 SMSTART/SMSTOP 節點是未使用 Chain 的一部分時,這些節點不被考慮進行排程,並從 DAG 中移除。為了防止這些節點被移除,我們需要一種方法來確保 CopyFromReg 的結果只能在 執行 SMSTART/SMSTOP 之後 使用。

我們可以使用 CopyToReg -> CopyFromReg 序列來做到這一點,它將值移動到/從虛擬暫存器,並將這些節點與 SMSTART/SMSTOP 鏈接起來,使它們成為計算結果值的表達式的一部分。產生的 COPY 節點會被暫存器分配器移除。

下面的範例展示了如何在 DAG 中使用此方法,該 DAG 不透過 Chain 將結果鏈接在一起,而是透過值鏈接。

            t0: ch,glue = AArch64ISD::SMSTOP ...
          t1: ch,glue = ISD::CALL ....
        t2: res,ch,glue = CopyFromReg t1, ...
      t3: ch,glue = AArch64ISD::SMSTART t2:1, ....   <- this is now part of the expression that returns the result value.
    t4: ch = CopyToReg t3, Register:f64 %vreg, t2
  t5: res,ch = CopyFromReg t4, Register:f64 %vreg
t6: res = FADD t5, t9

我們也需要在本機串流函數中使用此方法,其中需要在函數開始時將 SMSTART 插入到 DAG 中。

具有 __attribute__((arm_locally_streaming)) 的函數

如果函數被標記為 arm_locally_streaming,則序言/後記中的執行時 SVE 向量長度可能與函數主體中的向量長度不同。發生這種情況是因為我們在設定堆疊框架後調用 smstart,並在解除分配堆疊框架之前調用 smstop。

為了確保我們使用正確的 SVE 向量長度來分配本機變數,即使 CPU 尚未處於串流模式,我們也可以使用串流向量長度透過 ADDSVL 指令來分配堆疊槽。

這僅適用於本機變數,而不適用於被呼叫者儲存槽,因為 LLVM 不支援在一個堆疊框架中混合兩種不同的可擴展向量長度。這表示目前不支援將函數標記為 arm_locally_streaming 且需要在序言中溢出 SVE 被呼叫者儲存的情況。然而,在沒有使用者介入的情況下,這種情況不太可能發生,因為 arm_locally_streaming 函數不能接受或返回與向量長度相關的值。否則,這將需要強制使用 SVE PCS (使用 'aarch64_sve_pcs') 並結合使用 arm_locally_streaming 才能遇到這個問題。這種組合可以在 Clang 中透過發出診斷訊息來防止。

一個範例說明了對於使用 arm_locally_streaming 屬性的函數,序言/後記會是什麼樣子:

#define N 64

void __attribute__((arm_streaming_compatible)) some_use(svfloat32_t *);

// Use a float argument type, to check the value isn't clobbered by smstart.
// Use a float return type to check the value isn't clobbered by smstop.
float __attribute__((noinline, arm_locally_streaming)) foo(float arg) {
  // Create local for SVE vector to check local is created with correct
  // size when not yet in streaming mode (ADDSVL).
  float array[N];
  svfloat32_t vector;

  some_use(&vector);
  svst1_f32(svptrue_b32(), &array[0], vector);
  return array[N - 1] + arg;
}

應使用 ADDSVL 分配堆疊空間,並應避免覆蓋傳回/引數值。

_Z3foof:                                // @_Z3foof
// %bb.0:                               // %entry
        stp     d15, d14, [sp, #-96]!           // 16-byte Folded Spill
        stp     d13, d12, [sp, #16]             // 16-byte Folded Spill
        stp     d11, d10, [sp, #32]             // 16-byte Folded Spill
        stp     d9, d8, [sp, #48]               // 16-byte Folded Spill
        stp     x29, x30, [sp, #64]             // 16-byte Folded Spill
        add     x29, sp, #64
        str     x28, [sp, #80]                  // 8-byte Folded Spill
        addsvl  sp, sp, #-1
        sub     sp, sp, #256
        str     s0, [x29, #28]                  // 4-byte Folded Spill
        smstart sm
        sub     x0, x29, #64
        addsvl  x0, x0, #-1
        bl      _Z10some_usePu13__SVFloat32_t
        sub     x8, x29, #64
        ptrue   p0.s
        ld1w    { z0.s }, p0/z, [x8, #-1, mul vl]
        ldr     s1, [x29, #28]                  // 4-byte Folded Reload
        st1w    { z0.s }, p0, [sp]
        ldr     s0, [sp, #252]
        fadd    s0, s0, s1
        str     s0, [x29, #28]                  // 4-byte Folded Spill
        smstop  sm
        ldr     s0, [x29, #28]                  // 4-byte Folded Reload
        addsvl  sp, sp, #1
        add     sp, sp, #256
        ldp     x29, x30, [sp, #64]             // 16-byte Folded Reload
        ldp     d9, d8, [sp, #48]               // 16-byte Folded Reload
        ldp     d11, d10, [sp, #32]             // 16-byte Folded Reload
        ldp     d13, d12, [sp, #16]             // 16-byte Folded Reload
        ldr     x28, [sp, #80]                  // 8-byte Folded Reload
        ldp     d15, d14, [sp], #96             // 16-byte Folded Reload
        ret

防止在串流模式下使用非法指令

  • 當在串流模式 (PSTATE.SM=1) 下執行程式時,SVE/SVE2 指令的子集以及大多數 AdvSIMD/NEON 指令都是無效的。

  • 當在正常模式 (PSTATE.SM=0) 下執行程式時,SME 指令的子集是無效的。

  • 串流相容函數必須僅使用在 PSTATE.SM=0 或 PSTATE.SM=1 時都有效的指令。

PSTATE.SM 的值不是由功能標誌控制,而是由函數屬性控制。這表示我們可以針對 '+sme' 進行編譯,並且編譯器將產生任何指令的程式碼,即使它們在請求的串流模式下不合法。編譯器需要使用函數屬性來確保編譯器不會在假設某些操作在執行時可用的情況下進行轉換。

我們有意識地選擇不使用功能標誌來建模,因為我們仍然希望在任一模式下支援內嵌組合語言 (由使用者手動放置 smstart/smstop),並且由於 TableGen 中的限制,這在個別指令層級變得相當複雜 (請參閱 D120261D121208)。

作為第一步,這表示當函數具有 aarch64_pstate_sm_enabledaarch64_pstate_sm_bodyaarch64_pstate_sm_compatible 屬性中的任一項時,我們將完全停用向量化 (LoopVectorize/SLP),以避免使用向量指令。

稍後我們將致力於放寬這些限制,以使用串流相容指令的子集來啟用可擴展的自動向量化,但這需要變更 CostModel、Legalization 和 SelectionDAG 降低。

我們也將在 Clang 中發出診斷訊息,以防止在函數使用串流模式屬性裝飾時,使用非串流 (或不相容) 操作,例如透過 ACLE 內建函數。

其他考量事項

  • 當呼叫點需要切換 PSTATE.SM 或當被呼叫者的函數主體以與呼叫者不同的串流模式執行時,必須停用內聯。這是必要的,因為函數呼叫是串流模式變更的邊界。

  • 當呼叫點需要切換 PSTATE.SM 時,必須停用尾部呼叫最佳化,以便呼叫者可以還原 PSTATE.SM 的原始值。

3. 處理 PSTATE.ZA

與 PSTATE.SM 相反,啟用 PSTATE.ZA 不會影響 SVE 向量長度,也不會覆蓋 FP/AdvSIMD/SVE 暫存器。這表示使用內建函數切換 PSTATE.ZA 是安全的。這也使得為呼叫 private-ZA 函數 (即可能直接或間接覆蓋 ZA 狀態的函數) 設定延遲儲存機制更簡單。

為了處理標記為 aarch64_new_za 的函數,我們引入了一個新的 LLVM IR pass (SMEABIPass),該 pass 在 SelectionDAG 之前執行。由此 pass 處理的任何此類函數都標記為 aarch64_expanded_pstate_za

設定延遲儲存

提交延遲儲存

例外處理和 ZA

4. 類型

AArch64 謂詞即計數器類型

概述:

謂詞即計數器類型代表 AArch64 SVE 謂詞暫存器中保存的謂詞即計數器值的類型。此類值包含有關活動通道數量、元素寬度以及一個位元,該位元指示是否應反轉產生的遮罩的資訊。應使用 ACLE 內建函數將謂詞即計數器值在謂詞向量之間移動。

該類型存在某些限制:

  • 該類型可用於函數參數和傳回值。

  • 在此類型上支援的 LLVM 操作僅限於 loadstorephiselectalloca 指令。

謂詞即計數器類型是一種可擴展類型。

語法:

target("aarch64.svcount")

5. 參考文獻

  1. SME ACLE Pull-request

  2. SME ABI Pull-request