Go中的平台无关SIMD

# Go 中的平台无关 SIMD - Go 编程语言 来源：https://go.dev/blog/simd-experiment Go 1.26 和 1.27 包含了用于**单指令多数据（SIMD）**操作的实验性 API。SIMD 是许多现代 CPU 的原生特性，允许软件快速地对数据向量执行统一操作，例如在单条指令中将 8 对 `float64` 值相加。它可以显著加速许多计算密集型任务，范围涵盖密码学、数据处理到人工智能。事实上，Go 的 Green Tea 垃圾回收器（https://go.dev/blog/greenteagc）甚至利用 SIMD 来加速对存活对象的内存扫描。在这些新的实验性 API 出现之前，从 Go 访问此功能的唯一方式是编写 Go 汇编语言。这对于真正性能关键的计算内核来说才值得，这意味着许多可以从 SIMD 中受益的软件没有充分利用 CPU。Go 1.26 为 amd64 引入了 SIMD API，Go 1.27 则增加了对 arm64（具体是 NEON）和 wasm 的支持。然而，SIMD API 的一个基本挑战是平台之间巨大的差异，不仅在于支持的操作不同，甚至向量的表示方式也不同。一些平台提供固定大小的向量，通常在 128 位到 512 位之间，而在其他平台上，向量大小在构建时未知，必须在程序启动时查询。为了充分访问这些平台的广度，这些 API 位于架构相关的 `archsimd` 包中。但 Go 1.27 不仅提供了这些架构相关的 API，还引入了一个实验性的、完全可移植的、与平台和大小无关的 SIMD 接口，大致基于 C++ 的 Highway（https://google.github.io/highway/en/master/index.html）。其目标是在支持 SIMD 的平台上编写一次、接近汇编性能的 "simd" 代码，并在那些（尚未）不支持 SIMD 的平台上提供有竞争力的模拟。`simd` 包目前支持 amd64 上的 AVX、AVX2 和 AVX512，arm64 上的 NEON，以及 wasm 的 SIMD 指令。 ## 动机：SIMD 架构之间的差异 SIMD 架构在多个维度上存在差异。 - **向量大小**：一些提供单一的固定向量大小（wasm、PowerPC、s390x，128 位）。一些提供多种固定向量大小（amd64，有 128、256 和 512；loong64 有 128 和 256）。Riscv64 支持未指定大小（128 到 65536 位之间）的向量，但长度限制为 2 的幂。Arm64 支持一种固定大小（128 位，NEON）和一种可变大小（128-2048 位，仅限 2 的幂，SVE）。在特定架构的给定实例上，确定该实例支持哪些大小需要特性检查：amd64，但它是 AVX、AVX2 还是 AVX512？Arm64，但它是 NEON 还是 SVE？如果是 SVE，有多大？是 SVE、SVE2 还是 SVE2.1 的哪个变体？ - **向量掩码**：不同的 SIMD 架构在处理向量掩码方面各不相同。对于向量，如果-则-跨向量可以通过掩码实现；执行操作，但仅在掩码为 "真" 的位置赋值（或加载、或存储）。一些 SIMD 变体不提供掩码；所有操作都作用于所有元素，"掩码" 通过向量位掩码和向量布尔操作完成（wasm、AVX、AVX2、NEON）。一些提供特殊的掩码寄存器，一个位控制一个向量元素的操作（AVX512 和 RVV）。其他（SVE）为每个向量字节分配一个位，但每个元素的掩码位的最低有效位控制掩码操作。AVX2 也支持掩码加载和存储，但使用普通向量作为掩码，且最高有效位控制操作。 - **操作本身**：第三个差异来源是操作本身。每种架构提供自己的基本原语来重新排列向量元素；有些需要常量输入，有些支持变量输入。不同的 SIMD 架构支持不同的密码学相关操作。即使是基本算术也可能有不同的支持；例如 wasm 缺乏对 64 位整数向量的比较。即使对于特定架构上的给定向量长度，指令支持也取决于必须检查的 "特性"。尽管 Go 的架构相关 `archsimd` 包旨在尽可能跨架构统一，但许多这样的怪癖仍然存在，使得设计、编写和测试多平台 SIMD 代码变得繁重。我们可以在 `archsimd` 包中做更多工作，使不同的架构看起来更相似，但如果不牺牲效率，我们只能走到这一步。 ## 概述 新的 `simd` 包通过从类型系统中移除固定大小的向量，并仅支持那些所有不同平台交集中的操作来隐藏这些差异，并用其他 SIMD 指令的高效模拟填补交集中的空白。目标是一组操作： 1. 足以支持许多受益于向量化实现的数据处理算法（但不绑定于特定的向量大小）， 2. 当源代码操作与底层硬件匹配时，与汇编语言一样高效， 3. 否则，尽可能好地进行模拟， 4. 并且易于阅读和理解（即使/尤其是如果 LLM 最终编写了代码）。 在缺乏 SIMD 指令或 `archsimd` 中缺乏支持的平台上，所有操作都被模拟，因此使用 `simd` 编写的代码始终可以运行。要使用此实验性包，请设置 `GOEXPERIMENT=simd`，就像使用实验性 `archsimd` 包一样。`simd` 向量类型只是首字母大写的、复数的、原始类型，例如 `simd.Uint8s` 或 `simd.Float32s`。向量从切片加载和存储到切片，例如： ```go // innerProduct 返回 x 和 y 的内积。 func innerProduct(x, y []float32) float32 { var a simd.Float32s var i int for i = 0; i < len(x)-a.Len()+1; i += a.Len() { u := simd.LoadFloat32s(x[i : i+a.Len()]) v := simd.LoadFloat32s(y[i : i+a.Len()]) a = u.MulAdd(v, a) } if i < len(x) { u, _ := simd.LoadFloat32sPart(x[i:]) v, _ := simd.LoadFloat32sPart(y[i:]) a = u.MulAdd(v, a) } return sum(a) } // sum 返回 x 元素的标量和。 func sum(x simd.Float32s) float32 { s := make([]float32, x.Len()) x.Store(s) var r float32 for _, e := range s { r += e } return r } ``` 此示例也显示了此包第一个实验性版本的局限性之一；因为没有通用的求和方式来对向量的所有元素求和，Go 1.27 中的 `simd` 不支持此操作，尽管 `ReduceSum` 将出现在下一个版本中，因此 `sum` 可以被替换为简单的 `simd.ReduceSum`。 SIMD 比较产生掩码值，这些值对应于相应的向量元素宽度，因此 `Int8s` 的比较产生 `Mask8s`，依此类推，掩码值可用于选择和过滤向量。 ## Go 1.27 中支持的 `simd` 包操作 在此表中，`V` 和 `U` 是向量类型，`M` 是掩码类型，`E` 是标量类型，`W` 是宽度。 ### 包级别加载/广播函数 | 函数 | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `LoadV([]E) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `LoadVPart([]E) (V, int)` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `BroadcastV(E) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | ### 存储/字符串操作 | `(x V).Method(...)` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |---------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `Store(s []E)` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `StorePart(s []E) int` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `String() string` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | ### 算术操作 | `(x V).Method(...) V` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `Abs() V` | Y | Y | Y | | Y | | | | Y | | | `Add(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `AddSaturated(y V) V` | Y | Y | | | Y | | | | | | | `Average(y V) V` | | Y | | | | Y | | | | | | `Div(y V) V` | | | | | | | | | Y | Y | | `IfElse(mask MaskWs, y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Len() int` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Masked(mask MaskWs) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Max(y V) V` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `Min(y V) V` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `Mul(y V) V` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `MulAdd(y V, z V) V` | | | | | | | | | Y | Y | | `Neg() V` | Y | Y | Y | | Y | Y | | | | | | `Not() V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Or(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Sqrt() V` | | | | | | | | | Y | Y | | `Sub(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `SubSaturated(y V) V` | Y | Y | | | Y | | | | | | | `Xor(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | ### 布尔和向量掩码操作 | `(x V).Method(...) V` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `And(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | | | | `AndNot(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | | | | `CarrylessMultiplyEven(y V) V` | Y | | | | | | | | | | | `CarrylessMultiplyOdd(y V) V` | Y | | | | | | | | | | | `IfElse(mask MaskWs, y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Masked(mask MaskWs) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Not() V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Or(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Xor(y V) V` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | ### 比较操作 | `(x V).Method(...) M` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `Equal(y V) MaskWs` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | | `Greater(y V) MaskWs` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `GreaterEqual(y V) MaskWs` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `Less(y V) MaskWs` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `LessEqual(y V) MaskWs` | Y | Y | Y | | Y | Y | Y | | Y | Y | | `NotEqual(y V) MaskWs` | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | ### 转换操作 | `(x V).Method(...) U` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `ConvertToFloatW() FloatWs` | Y | | | | | | | | | | | `ConvertToIntW() IntWs` | Y | Y | Y | Y | Y | | | | | | | `ConvertToUintW() UintWs` | | | Y | Y | | | | | | | | `ToMask() (to MaskWs)` | Y | Y | Y | Y | | | | | | | ### 掩码方法 | `(m M).Method(...) M)` | `Mask8s` | `Mask16s` | `Mask32s` | `Mask64s` | |-------------------------|----------|-----------|-----------|-----------| | `And(y M) M` | Y | Y | Y | Y | | `Or(y V) V` | Y | Y | Y | Y | | `String() string` | Y | Y | Y | Y | | `ToIntWs() (to IntWs)` | Y | Y | Y | Y | ### 移位和循环操作 | `(x V).Method() V` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |----------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `RotateAllLeft(dist uint64) V` | Y | Y | Y | | Y | Y | | | | | | `RotateAllRight(dist uint64) V` | Y | Y | Y | | Y | Y | | | | | | `ShiftAllLeft(dist uint64) V` | Y | Y | Y | | Y | Y | | | | | | `ShiftAllRight(dist uint64) V` | Y | Y | Y | | Y | Y | | | Y | | ### 零成本重塑操作 | `(x V).Method(...) U` | `Int8s` | `Int16s` | `Int32s` | `Int64s` | `Uint8s` | `Uint16s` | `Uint32s` | `Uint64s` | `Float32s` | `Float64s` | |------------------------|---------|----------|----------|----------|----------|-----------|-----------|-----------|------------|------------| | `ToBits() UintWs` | Y | Y | Y | | Y | Y | | | | | | `ReshapeToUint8s() Uint8s` | | Y | Y | | | Y | Y | | | | | `ReshapeToUint16s() Uint16s` | | | Y | | | | Y | | | | | `ReshapeToUint32s() Uint32s` | | | | | | | | Y | | | | `ReshapeToUint64s() Uint64s` | | | | | | | | Y | | | | `BitsToFloatW() FloatWs` | | | | | | | | | Y | Y | | `BitsToIntW() IntWs` | | Y | Y | Y | | | | | | | ## 与平台特定代码的转换 `simd` 包对于特定应用程序的所有部分可能过于有限，或者我们可能尚未为某些必要特性提供足够的模拟。对于这种情况，`simd` 包支持与架构特定的 SIMD 之间的转换。`simd` 包中的每种向量类型都有一个转换方法 `ToArch()`，返回一个 `any`。该 `any` 可以被类型断言为平台的一种架构特定类型。要转换回来，请使用其中一个 `simd.FromArch` 函数。对于可移植代码，这给为每个平台（包括模拟）编写特定架构代码带来了义务。这是一个完整的方法/函数示例，该方法目前缺失，但应在 Go 1.28 中添加。假设你的算法需要 `Int8s.OnesCount()`（Go 1.27 中的 `simd` 缺乏此功能）。无需为每个平台重写整个算法，可以只实现缺失的操作。首先，对于 amd64，AVX 和 AVX2 缺乏该指令，但 AVX512 有： ```go //go:build goexperiment.simd && amd64 package simd_test import ( "simd" "simd/archsimd" ) var popcnt4x16 = [16]int8{0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4} var popcnt4x32 = [32]int8{ 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, } // OnesCount 返回每个元素的 1 位数量。 func OnesCount(v simd.Int8s) simd.Int8s { switch x := v.ToArch().(type) { case archsimd.Int8x16: lut := archsimd.LoadInt8x16Array(&popcnt4x16) mask0f := archsimd.BroadcastInt8x16(0x0f) lo := x.And(mask0f) hi := x.ToBits().ReshapeToUint16s().ShiftAllRight(4). ReshapeToUint8s().BitsToInt8().And(mask0f) return simd.Int8sFromArch(lut.PermuteOrZero(lo). Add(lut.PermuteOrZero(hi))) case archsimd.Int8x32: lut := archsimd.LoadInt8x32Array(&popcnt4x32) mask0f := archsimd.BroadcastInt8x32(0x0f) lo := x.And(mask0f) hi := x.ToBits().ReshapeToUint16s().ShiftAllRight(4). ReshapeToUint8s().BitsToInt8().And(mask0f) return simd.Int8sFromArch(lut.PermuteOrZeroGrouped(lo). Add(lut.PermuteOrZeroGrouped(hi))) case archsimd.Int8x64: return simd.Int8sFromArch(x.OnesCount()) default: // GODEBUG=simd=0 模拟 return OnesCountEmulated(v) } } ``` 接口转换和类型切换看起来效率低下，但 `simd` 在编译器端的实现会特化代码并优化掉类型切换。NEON 和 Wasm 都支持 `Int8s.OnesCount()`，因此它们的实现要简单得多，尽管仍然使用了 `Int8s.ToArch` 和 `Int8sFromArch`。 ```go //go:build goexperiment.simd && (wasm || arm64) package simd_test import ( "simd" "simd/archsimd" ) // OnesCount 返回每个元素的 1 位数量。 func OnesCount(v simd.Int8s) simd.Int8s { // TODO 当添加 SVE 时，这将不起作用 switch x := v.ToArch().(type) { case archsimd.Int8x16: return simd.Int8sFromArch(x.OnesCount()) default: // GODEBUG=simd=0 模拟 return OnesCountEmulated(v) } } ``` 不要忘记有些人没有硬件 SIMD 支持： ```go //go:build goexperiment.simd && !(amd64 || wasm || arm64) package simd_test import ( "simd" ) // OnesCount 返回每个元素的 1 位数量。 func OnesCount(v simd.Int8s) simd.Int8s { return OnesCountEmulated(v) } ``` 为了完成这个练习，一个单独的模拟函数作为所有实现的后备方案共享： ```go //go:build goexperiment.simd package simd_test import ( "simd" ) // OnesCountEmulated 返回每个元素的 1 位数量。 func OnesCountEmulated(v simd.Int8s) simd.Int8s { a := [2]uint64{} v.ToBits().ReshapeToUint64s().Store(a[:]) a0, a1 := a[0], a[1] m1 := uint64(0x5555555555555555) m2 := uint64(0x3333333333333333) m4 := uint64(0x0f0f0f0f0f0f0f0f) a0 = (a0 & m1) + ((a0 >> 1) & m1) a1 = (a1 & m1) + ((a1 >> 1) & m1) a0 = (a0 & m2) + ((a0 >> 2) & m2) a1 = (a1 & m2) + ((a1 >> 2) & m2) a0 = (a0 & m4) + ((a0 >> 4) & m4) a1 = (a1 & m4) + ((a1 >> 4) & m4) a[0], a[1] = a0, a1 return simd.LoadUint64s(a[:]).ReshapeToUint8s().BitsToInt8() } ``` ## API 交集和方法模拟 `simd` 包提供的任何操作都需要在大多数架构上运行良好。第一步，任何普遍支持的操作都可以很容易地在 `simd` 上支持。这往往包括加载、存储、算术和比较（但并非所有比较！）。天真地对不同架构的 SIMD 方法取交集仍然留下许多空缺。这些空缺通过在各种架构特定的 `archsimd` API 中添加模拟来填补。这些 API 已经包含了许多简单的模拟，以简化 Go 程序员的生活；有符号和无符号整数加法使用相同的指令，但就像 Go 为 `int` 和 `uint` 都支持 `+` 运算符一样，`archsimd` 包提供了 `Int8x16.Add(Int8x16)` 和 `Uint8x16.Add(Uint8x16)`，即使它们编译成相同的指令。现代编程语言也不期望程序员知道如何用位操作来实现浮点数的取反和绝对值，因此 `archsimd` 在必要时实现这一点，或者如果你从另一个角度看，它就是在 "模拟"。有许多模拟只需要 2 或 3 条指令；例如，一些架构只支持跨向量元素的相同值移位距离，而另一些则支持每个向量元素不同的移位距离。为了在 `simd` 中支持标量移位，我们用向量移位来模拟标量移位。一些架构缺乏某些无符号比较——这些只是有符号比较加上与常量的两次异或。并非所有缺失的指令都这么简单。"无进位乘法" 指令对密码学和 CRC 校验和很重要，但它并不总是被支持。将其排除在 `simd` API 之外将阻止其用于一些重要的算法。因此，我们提供了一种模拟，而且由于一个重要的用途是在密码学中，它的运行时间不会因...（原文截断）