Go 语言字符串相等函数的 SIMD 实现
背景说明
Go 语言是 Google 开发的一种静态强类型、编译型、并发型,并具有垃圾回收功能的编程语言。它用批判吸收的眼光,融合 C 语言、Java 等众家之长,将简洁、高效演绎得淋漓尽致。
SIMD 全称 Single Instruction Multiple Data,即单指令多数据流,用一条指令对多个数据进行操作。一般用向量寄存器来实现。常用于加速数据密集型运算,如数列求和、矩阵乘法。
本实验对 Go 语言自带的字符串相等函数的源码进行分析,用自己实现的函数进行替换,并比较性能。
探索过程
不同 CPU 对 SIMD 的支持不同。用 CPU-Z 查看当前 CPU 支持的 SIMD 指令集。最新的指令集为 AVX2。
Go 语言汇编器基于 Plan 9 汇编器的输入风格,与 GNU 汇编器不同。阅读和编写代码时需要注意。
Go 语言字符串相等函数的代码在 GOROOT/src/internal/bytealg/equal_[arch].s
文件中。当前 CPU 架构为 AMD64,对应文件为 equal_amd64.s
。
文件中定义了三个函数:runtime·memequal
,runtime·memequal_varlen
,memeqbody
。前两个函数为 ABI,供应用程序调用。它们将设置相应寄存器并跳转到第三个函数。部分定义如下。
1 |
|
我们只需要关注真正进行相等判断的 memeqbody
函数。它接收两个字符串的地址(SI
、DI
)以及字符串的长度(BX
),返回 0 或 1(AX
)。Go 语言的编译器保证这两个字符串的长度相等(否则可以直接判断不相等)。
原代码巧妙地利用了 SIMD。思路如下:
- 如果字符串的长度不小于 64,则一轮循环比较 64 个字符(512 位),直到剩余长度小于 64;
- 如果字符串的长度不小于 8,则一轮循环比较 8 个字符(64 位),直到剩余长度小于 8;
- 比较剩余字符(不用循环)。
结合代码分析。在函数开头,根据字符串的长度进入不同的循环。
1 |
|
第 6 行的代码判断 CPU 是否支持 AVX2 指令集,如果支持则用 Y
系列寄存器比较 64 个字符,否则用 X
系列寄存器比较 64 个字符。64 个字符的循环如下。
1 |
|
如果发现不同,则直接返回 0,否则继续循环,直到剩余长度小于 64,进入 8 个字符的循环。一个通用寄存器刚好可以装下 8 个字符,所以不需要 SIMD。
1 |
|
最后比较剩余字符。不过这里并没有用循环,而是直接加载字符串末尾的 8 个字符(可能与之前判断过的字符重叠)。
1 |
|
如果字符串的长度本来就小于 8,这么做会加载一些不属于字符串的字符。代码中对这种情况也做了处理。
可以看出,这个函数尽可能地使用了 SIMD,用一条指令判断多个字符是否相等,加快了处理速度,同时也保证了边界情况下的正确性,对较短的字符串进行特殊处理。
尝试用自己编写的函数进行替换。首先用最简单的方法,直接一个一个字符进行比较。
1 |
|
代码很短,但显然效率不够高。添加 8 个字符的循环,剩余字符还是一个一个比较。
1 |
|
当前 CPU 支持 MMX、SSE4.2、AVX2 指令集。MMX 的寄存器为 64 位寄存器,一个寄存器可以装下 8 个字符。将 8 个字符的循环改为使用 MMX 的寄存器。
1 |
|
第 4、5 行将从地址 SI
、DI
开始的 8 个字符分别存入 M0
、M1
寄存器,即进行了打包。第 6 行用 PCMPEQB
指令(compare packed bytes for equal)比较 M0
、M1
打包字节整数值的相等性,并将比较结果存入 M1
。如果 M0
、M1
中的某个字节相等,则 M1
中的这个字节会变成全 1(即有符号数的 -1),否则变成全 0。因为 MMX 指令不会修改状态寄存器,所以需要将 M1
的值存入 CX
,再与 -1 比较。
SSE4.2 的寄存器为 128 位寄存器,一个寄存器可以装下 16 个字符。添加 16 个字符的循环。
1 |
|
第 4、5 行用 MOVUPD
指令(move two unaligned packed double-precision floating-point values between XMM registers and memory)将从地址 SI
、DI
开始的 16 个字符分别存入 X0
、X1
寄存器。第 6 行将比较结果存入 X1
。第 7 行用 PMOVMSKB
指令(move byte mask)将 X1
中每个字节的最高位提取出来存入 CX
的低位。这是因为 X1
寄存器有 128 位,无法直接存入通用寄存器进行判断。而比较某个字节时,如果相等则这个字节会变成全 1,所以只要提取每个字节的最高位即可判断是否全部相等。
AVX2 的寄存器为 256 位寄存器,一个寄存器可以装下 32 个字符。添加 32 个字符的循环。
1 |
|
AVX2 指令集与 SSE4.2 指令集类似,只需在指令前加字母 V。唯一不同的是第 6 行的 VPCMPEQB
指令,它需要三个操作数,将前两个操作数的比较结果存入第三个操作数。
由于 CPU 不支持 AVX512 指令集,因此无法使用 512 位寄存器。不过,还可以用循环展开来优化代码。AVX2 的寄存器有 16 个,所以可以展开 2 次、4 次、8 次循环。以展开 4 次循环为例,添加 128 个字符的循环。
1 |
|
第 4 行到第 7 行将连续 128 个字符存入 Y
系列寄存器,第 12 行到第 15 行分别比较四对 Y
系列寄存器,第 16 行到第 18 行将比较结果作按位与(因为相等的比较结果为全 1),第 19 行将结果每个字节的最高位提取出来存入 CX
的低位。展开 2 次、8 次循环的代码类似。
至此,已基本实现了用 SIMD 优化的字符串相等函数。
效果分析
编写测试函数,用于测试字符串相等函数的正确性。
1 |
|
用 go test -run ^TestEqual$
命令运行测试函数。实验中编写的每一种循环都可以通过测试。
编写基准函数,用于测试字符串相等函数的性能。
1 |
|
用 go test -run ^$ -bench ^BenchmarkEqual$ -count 10
命令运行基准函数,可以得到字符串相等函数的运行次数和平均运行时间。字符串大小均为 4KB。据此可以算出函数处理数据的速率。
在表格的说明中,R1 表示一轮比较 1 个字符,R8 表示用通用寄存器一轮比较 8 个字符,M8 表示用 MMX 寄存器一轮比较 8 个字符,X16 表示用 SSE4.2 寄存器一轮比较 16 个字符,Y32 表示用 AVX2 寄存器一轮比较 32 个字符,Y64、Y128、Y256 分别表示用 AVX2 寄存器和 2 次、4 次、8 次循环展开一轮比较 64 个、128 个、256 个字符。
测试程序 | 说明 | 运行次数 | 运行时间(ns/op) | 处理速率(GB/s) |
---|---|---|---|---|
original.asm |
Go 语言自带 | 13499010 | 87.92 | 43.39 |
loop_1.asm |
R1 | 424518 | 2721 | 1.402 |
loop_8.asm |
R8 + R1 | 3993097 | 298.0 | 12.80 |
loop_8_mmx.asm |
M8 + R1 | 2980017 | 395.4 | 9.648 |
loop_16_sse.asm |
X16 + R8 + R1 | 5688907 | 206.8 | 18.45 |
loop_32_avx.asm |
Y32 + R8 + R1 | 10291144 | 115.2 | 33.11 |
loop_64_unroll.asm |
Y64 + R8 + R1 | 13146477 | 88.00 | 43.35 |
loop_128_unroll.asm |
Y128 + Y32 + R8 + R1 | 17871183 | 64.31 | 59.32 |
loop_256_unroll.asm |
Y256 + Y32 + R8 + R1 | 21847575 | 54.56 | 69.92 |
可以看出:
- 随着 SIMD 寄存器位数的增加,函数的运行时间会减少,处理数据的速率也会变快;
- 用通用寄存器比较 8 个字符比用 MMX 寄存器比较 8 个字符要快,是因为 MMX 寄存器比较相等之后还需要移回通用寄存器进行跳转判断;
- Go 语言自带的函数用 AVX2 的寄存器时只展开了 2 次循环,实验中展开了 2 次循环的函数与 Go 语言自带的函数速度相近,而展开了 4 次、8 次循环的函数速度更快;
- 对于大小为 4KB 的字符串,实验中最快的函数相比最慢的函数速度提升 4887%,相比 Go 语言自带的函数速度提升 61.14%。
以上结果符合预期。
参考文献
- A Quick Guide to Go’s Assembler, https://go.dev/doc/asm.
- A Manual for the Plan 9 assembler, https://9p.io/sys/doc/asm.html.
- x64 Cheat Sheet, https://cs.brown.edu/courses/cs033/docs/guides/x64_cheatsheet.pdf.
- x86 and amd64 instruction reference, https://www.felixcloutier.com/x86/index.html.
- x86 Assembly Language Reference Manual, https://docs.oracle.com/cd/E37838_01/html/E61064/index.html.
- Intel® Instruction Set Extensions Technology, https://www.intel.com/content/www/us/en/support/articles/000005779/processors.html.
- equal_amd64.s, https://github.com/golang/go/blob/master/src/internal/bytealg/equal_amd64.s.
- testing package, https://pkg.go.dev/testing.
本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!