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 和几乎所有代码生成部分的目的,我们可以假设 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)) 的函数被排除在此表之外,因为对于调用方而言,该属性与“流”同义,而对于被调用方而言,它仅仅是一个实现细节,明确地不对调用方公开。

表 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 指令。否则,该节点将与扩展为比较/分支和 smstart/smstop 的伪指令匹配。这对于实现从 SC -> NSC -> S 的转换是必要的。

未链接的函数调用

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

如果被调用的函数是一个没有副作用的内联函数,该内联函数又降低为函数调用(例如 @llvm.cos()),则对 @llvm.cos() 的调用不属于任何 Chain;它可以自由调度。

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

  • 启动调用序列

  • 将输入值从虚拟寄存器复制到 ABI 指定的物理寄存器

  • 执行跳转并链接

  • 停止调用序列

  • 将输出值从其物理寄存器复制到虚拟寄存器

当调用站点的 Chain 未被使用时,仅使用来自链式序列的结果值,但 Chain 本身会被丢弃。

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

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

下面的示例展示了如何在不通过 Chain 而是通过值链接结果的 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 向量长度来分配局部变量,即使 CPU 尚未处于流模式,我们也可以通过 ADDSVL 指令使用流向量长度来分配栈槽。

这仅适用于局部变量,而不适用于被调用者保存的槽,因为 LLVM 不支持在一个栈帧中混合两种不同的可扩展向量长度。这意味着当一个函数被标记为 arm_locally_streaming 并且需要在序言中溢出 SVE 被调用者保存时,当前不支持这种情况。但是,在没有用户干预的情况下,这种情况不太可能发生,因为 arm_locally_streaming 函数不能获取或返回依赖于向量长度的值。否则,这将需要强制使用 ' aarch64_sve_pcs' 的 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),并且在单个指令级别实现这一点变得相当复杂(参见 D120261D121208),因为 TableGen 的限制。

作为第一步,这意味着当函数具有 aarch64_pstate_sm_enabledaarch64_pstate_sm_bodyaarch64_pstate_sm_compatible 属性中的任何一个时,我们将完全禁用向量化(LoopVectorize/SLP),以避免使用向量指令。

稍后,我们将旨在放宽这些限制,以使用与流兼容的指令子集启用可扩展的自动向量化,但这需要对成本模型、合法化和 SelectionDAG 降级进行更改。

我们还将在 Clang 中发出诊断,以防止在函数使用流模式属性装饰时使用非流(兼容)操作,例如通过 ACLE 内联函数。

其他需要考虑的事项

  • 当调用方需要切换 PSTATE.SM 或被调用方的函数体在与其调用方不同的流模式下执行时,必须禁用内联。这是因为函数调用是流模式更改的边界。

  • 当调用方需要切换 PSTATE.SM 时,必须禁用尾调用优化,以便调用方可以恢复 PSTATE.SM 的原始值。

3. 处理 PSTATE.ZA

与 PSTATE.SM 相反,启用 PSTATE.ZA 不会影响 SVE 向量长度,也不会修改 FP/AdvSIMD/SVE 寄存器。这意味着可以使用内联函数安全地切换 PSTATE.ZA。这也使为对私有 ZA 函数(即可能直接或间接修改 ZA 状态的函数)的调用设置延迟保存机制变得更加简单。

为了处理用 aarch64_new_za 标记的函数,我们在 SelectionDAG 之前引入了一个新的 LLVM IR 传递 (SMEABIPass)。此传递处理的任何此类函数都将用 aarch64_expanded_pstate_za 标记。

设置延迟保存

提交延迟保存

异常处理和 ZA

4. 类型

AArch64 谓词作为计数器类型

概述:

谓词作为计数器类型表示保存在 AArch64 SVE 谓词寄存器中的谓词作为计数器值的类型。此类值包含有关活动通道数、元素宽度以及指示生成的掩码是否应反转的位的的信息。应使用 ACLE 内联函数将谓词作为计数器值移动到/从谓词向量。

该类型有一些限制

  • 该类型可用于函数参数和返回值。

  • 此类型支持的 LLVM 操作仅限于 loadstorephiselectalloca 指令。

谓词作为计数器类型是可扩展类型。

语法:

target("aarch64.svcount")

5. 参考

  1. SME ACLE 拉取请求

  2. SME ABI 拉取请求