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/vector 操作的执行可能会转移到另一个处理单元。这有三个重要的含义

运行时 SVE 向量长度可能会更改。
FP/AdvSIMD/SVE 寄存器的内容被清零。
允许的指令集发生变化。

这导致了对 IR 和优化的某些限制。例如，在可能使用不同 PSTATE.SM 值的函数之间共享依赖于向量长度的状态是未定义的行为。前端在生成 LLVM IR 时必须遵守这些限制。

即使运行时 SVE 向量长度可能会更改，但为了 LLVM IR 和几乎所有 CodeGen 部分的目的，我们可以假设运行时的 vscale 值不会改变。如果我们让编译器在调用边界周围插入适当的 smstart 和 smstop 指令，那么可以减轻对 SVE 状态的影响。通过将状态更改限制在调用周围非常短暂的窗口内，我们可以控制操作的调度方式以及如何在状态转换之间保持实时值。

为了在这个粒度级别控制 PSTATE.SM，我们使用函数和调用站点属性，而不是内在函数。

属性限制 ¶

将（指向）可伸缩向量对象的指针传递或返回给/从可能使用不同 SVE 向量长度的函数是未定义的行为。这包括具有非流式接口但标记为 aarch64_pstate_sm_body 的函数。
不允许函数同时使用 aarch64_pstate_sm_compatible 和 aarch64_pstate_sm_enabled 进行修饰。
不允许函数使用以下属性中的多个进行修饰：aarch64_new_za、aarch64_in_za、aarch64_out_za、aarch64_inout_za、aarch64_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)) 的函数被排除在此表之外，因为对于调用者来说，该属性与“流式”同义，而对于被调用者来说，它仅仅是一个实现细节，明确地不向调用者公开。

表 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/vector 寄存器，因此最好在寄存器分配之前发出 smstart 和 smstop 指令，以便寄存器分配器可以在模式更改前后溢出/重载寄存器。

编译器还应该有足够的信息了解哪些操作是调用/函数的参数/结果的一部分，哪些操作是函数体的一部分，以便它可以将模式更改放置在完全正确的位置。执行此操作的合适位置似乎是 SelectionDAG，它在其中降低调用的参数/返回值以实现指定的调用约定。SelectionDAG 提供链和 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 节点采用 CurrentState 和 ExpectedState 操作数，用于条件 SMSTART/SMSTOP 的情况。仅当 CurrentState != ExpectedState 时才会执行该指令。

当 CurrentState 和 ExpectedState 可以在编译时评估时（即它们都是常量），则会发出无条件 smstart/smstop 指令。否则，该节点将匹配到一个伪指令，该指令扩展为比较/分支和 smstart/smstop。这对于实现从 SC -> N 和 SC -> S 的转换是必要的。

非链式函数调用 ¶

当具有 “aarch64_pstate_sm_enabled” 的函数调用不兼容流式的函数时，编译器必须在调用之前插入 SMSTOP，并在调用之后插入 SMSTART。

如果被调用的函数是没有副作用的内在函数，而该内在函数又被降低为函数调用（例如 @llvm.cos()），那么对 @llvm.cos() 的调用不属于任何链；它可以自由调度。

降低调用站点会创建一个小的节点链，该链

启动调用序列
将输入值从虚拟寄存器复制到 ABI 指定的物理寄存器
执行分支和链接
停止调用序列
将输出值从其物理寄存器复制到虚拟寄存器

当调用站点的链未被使用时，仅使用链式序列的结果值，但链本身将被丢弃。

SMSTART 和 SMSTOP ISD 节点返回一个链，但没有实际值，因此当 SMSTART/SMSTOP 节点是不使用的链的一部分时，这些节点不被考虑用于调度，并从 DAG 中删除。为了防止这些节点被删除，我们需要一种方法来确保 CopyFromReg 的结果只能在 SMSTART/SMSTOP 执行之后 使用。

我们可以使用 CopyToReg -> CopyFromReg 序列来实现这一点，该序列将值移动到/从虚拟寄存器，并将这些节点与 SMSTART/SMSTOP 链接起来，使其成为计算结果值的表达式的一部分。生成的 COPY 节点由寄存器分配器删除。

下面的示例显示了如何在 DAG 中使用它，该 DAG 不通过链而是通过值将结果链接在一起

            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 向量长度来分配局部变量，我们可以使用流式向量长度通过 ADDSVL 指令来分配堆栈槽，即使 CPU 尚未处于流模式。

这仅适用于局部变量，而不适用于被调用者保存槽，因为 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 的限制，这在单个指令级别上实现起来变得相当复杂（请参阅 D120261 和 D121208）。

作为第一步，这意味着当函数具有 aarch64_pstate_sm_enabled、aarch64_pstate_sm_body 或 aarch64_pstate_sm_compatible 属性中的任何一个时，我们将完全禁用向量化（LoopVectorize/SLP），以避免使用向量指令。

稍后，我们的目标是放宽这些限制，以使用流式兼容指令的子集启用可伸缩自动向量化，但这需要更改 CostModel、Legalization 和 SelectionDAG 降低。

我们还将在 Clang 中发出诊断信息，以防止在使用流模式属性修饰函数时使用非流式（兼容）操作，例如通过 ACLE 内在函数。