SIMD 加速:AVX2 指令集实现大小端转换
在应用 thrift 进行 RPC 通信的时候,由于 Thrift 采用了大端序,而 x86_64 等常见的处理器架构均为小端序,因此对于 list 这一类的数据类型需要一个循环来实现小端到大端的转换。而这个过程如果能够利用 SIMD 指令的话,可以极大的提高性能。这篇文章是在探索实现 Thrift 编译后端 Auto-vectorization Pass 的时候进行的一个尝试和学习,使用 avx2 指令集实现了一个简单的大小端转换的功能,并且对比了在不同条件下的加速性能。
原理
大小端转换
计算机数据存储有两种字节优先顺序:高位字节优先(称为大端模式)和低位字节优先(称为小端模式)。
大端模式
,是指数据的高字节保存在内存的低地址中,而数据的低字节保存在内存的高地址中,这样的存储模式有点儿类似于把数据当作字符串顺序处理:地址由小向大增加,而数据从高位往低位放;这和我们的阅读习惯一致。小端模式
,是指数据的高字节保存在内存的高地址中,而数据的低字节保存在内存的低地址中,这种存储模式将地址的高低和数据位权有效地结合起来,高地址部分权值高,低地址部分权值低。
例如,对于内存中存放的数0x12345678来说
- 如果是采用大端模式存放的,则其真实的数是:0x12345678
- 如果是采用小端模式存放的,则其真实的数是:0x78563412
可以使用如下 API 进行转换:
1 |
|
也可以直接使用移位进行实现
1 | inline uint32_t Swap32(uint32_t x) |
bswap
大部分编译器同时提供了 bswap
指令,来帮助实现这一转换过程,例如在 gcc 中,我们可以使用 __builtin_bswap{16,32,64}
:
1 | inline uint32_t Swap32(uint32_t x) |
这是一个编译器的内置函数。在 x86_64 机器上,它会被编译为这样的指令序列:
1 | Swap32(unsigned int): |
在 arm 机器上,它会被编译为这样的指令序列:
1 | Swap32(unsigned int): |
通常来说,我们自己使用的移位函数实现的大小端转换,在编译器优化 O2 时也会被自动识别替换为 bswap 指令。
avx2 指令集
使用 SIMD 对于这样可以高度并行化的计算应该是一个更快的选择。bswap指令可以反转 2, 4, 或 8 字节的字节顺序,但 x86 中的 SIMD 扩展允许仅使用一条指令对多条数据通道进行并行操作。就像原子地反转寄存器中的所有四个字节一样,它提供了一个完整的算术指令集,允许使用一条指令同时并行处理多个数据实例。这些被操作的数据块往往被称为 vectors。一般来说可用的有如下几种 SIMD 指令集:
- MMX (1996)
- SSE (1999)
- SSE2 (2001)
- SSE3 (2004)
- SSSE3 (2006)
- SSE4 a/1/2 (2006)
- AVX (2008)
- AVX2 (2013)
- AVX512 (2015)
目前较为常用的是 avx/avx2 指令集,早先的某些指令集主要是为了兼容性而保留的。具体的指令信息,可以参考 Intel 指令集查询。
我们这里主要使用的是 _mm256_shuffle_epi8
指令,在 C 中它被定义在了 #include <immintrin.h>
头文件中。它实现了一个 vector 中字节的重排序,例如将一个 128 位的字节序列完全反转:
1 | const __m256i ShuffleRev = _mm256_set_epi8( |
它接受一个ShuffleRev,定义具体每个字节应该被放到哪个位置。注意每128位为一个通道,重排范围只能在128位内,不能将前128位的内容重排至后128位。可以参照下图,可能会比较直观:
来源:https://www.officedaytime.com/simd512e/simdimg/si.php?f=pshufb
在 gcc -O3 中,Auto-vectorization Pass 可以帮助我们自动识别可以被向量化的循环,并且使用 avx 指令集进行并行化优化。
avx2 code
这是一个对于 64 位整数的大小端转换 load-swap-store 循环,使用 avx2 指令集进行加速的简单示例:
1 | void reverse64_avx2(void *Array, size_t Count) |
avx2 指令集的向量是 256 位长度,相当于 4 个 64bit 的整数。由于输入的数组并不一定被 4 整除,因此结尾的部分使用一般转换法逐个进行大小端转换。
benchmark
测试环境:
- Linux ubuntu 5.13.0-51-generic #58~20.04.1-Ubuntu SMP Tue Jun 14 11:29:12 UTC 2022 x86_64
- Intel Core i7-10750H
编译指令:
1 | gcc main.c -mavx2 -fno-tree-vectorize -O3 -o avx |
basic 对照函数(这里 Swap64 会被 gcc 自动编译为 bswap 指令):
1 | void reverse64_basic(void *Array, size_t Count) |
我们分别对 64/32/16 位的整数进行大小端转换,并测试 bswap 和 avx2 的加速比:
array size | avx2 64bit | basic 64bit | avx2 32bit | basic 32bit | avx2 16bit | basic 16bit |
---|---|---|---|---|---|---|
4 | 2ns | 3ns | 3ns | 3ns | 4ns | 2ns |
8 | 3ns | 4ns | 2ns | 4ns | 15ns | 9ns |
16 | 5ns | 9ns | 3ns | 9ns | 3ns | 10ns |
32 | 9ns | 37ns | 5ns | 18ns | 4ns | 19ns |
64 | 19ns | 34ns | 9ns | 34ns | 6ns | 59ns |
128 | 35ns | 181ns | 15ns | 76ns | 9ns | 82ns |
256 | 52ns | 120ns | 26ns | 477ns | 11ns | 712ns |
512 | 86ns | 248ns | 44ns | 192ns | 29ns | 254ns |
1024 | 179ns | 510ns | 96ns | 422ns | 47ns | 486ns |
2048 | 383ns | 996ns | 179ns | 812ns | 96ns | 981ns |
4096 | 726ns | 2190ns | 457ns | 2675ns | 384ns | 1878ns |
8192 | 1544ns | 4170ns | 748ns | 3434ns | 401ns | 4511ns |
16384 | 3570ns | 8977ns | 1704ns | 6771ns | 793ns | 7941ns |
可以注意到,对于宽度更小的整数的数组,并行度更高,avx2 加速比更加明显。在 64 位时,加速比约为 2.5,在 16 位时,加速比已经达到了 10 倍。
生成的汇编
objdump -d ./avx > dump.s
我们可以再检查一下生成的汇编指令:
使用 bswap 的大小端转换
1
2
3
4
5
6
7
8
9
10
11
12
13
140000000000001c30 <reverse64_basic>:
1c30: f3 0f 1e fa endbr64
1c34: 48 85 f6 test %rsi,%rsi
1c37: 74 1a je 1c53 <reverse64_basic+0x23>
1c39: 48 8d 14 f7 lea (%rdi,%rsi,8),%rdx
1c3d: 0f 1f 00 nopl (%rax)
1c40: 48 8b 07 mov (%rdi),%rax
1c43: 48 83 c7 08 add $0x8,%rdi
1c47: 48 0f c8 bswap %rax
1c4a: 48 89 47 f8 mov %rax,-0x8(%rdi)
1c4e: 48 39 d7 cmp %rdx,%rdi
1c51: 75 ed jne 1c40 <reverse64_basic+0x10>
1c53: c3 retq
1c54: 66 66 2e 0f 1f 84 00 data16 nopw %cs:0x0(%rax,%rax,1)avx2:vpshufb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
340000000000001bb0 <reverse64_avx2>:
1bb0: f3 0f 1e fa endbr64
1bb4: 48 89 f2 mov %rsi,%rdx
1bb7: 48 83 e2 fc and $0xfffffffffffffffc,%rdx
1bbb: 74 46 je 1c03 <reverse64_avx2+0x53>
1bbd: c5 fd 6f 0d fb 14 00 vmovdqa 0x14fb(%rip),%ymm1 # 30c0 <_IO_stdin_used+0xc0>
1bc4: 00
1bc5: 48 8d 4a ff lea -0x1(%rdx),%rcx
1bc9: 48 89 f8 mov %rdi,%rax
1bcc: 48 8d 14 d7 lea (%rdi,%rdx,8),%rdx
1bd0: c5 fa 6f 10 vmovdqu (%rax),%xmm2
1bd4: c4 e3 6d 38 40 10 01 vinserti128 $0x1,0x10(%rax),%ymm2,%ymm0
1bdb: 48 83 c0 20 add $0x20,%rax
1bdf: c4 e2 7d 00 c1 vpshufb %ymm1,%ymm0,%ymm0
1be4: c5 fa 7f 40 e0 vmovdqu %xmm0,-0x20(%rax)
1be9: c4 e3 7d 39 40 f0 01 vextracti128 $0x1,%ymm0,-0x10(%rax)
1bf0: 48 39 d0 cmp %rdx,%rax
1bf3: 75 db jne 1bd0 <reverse64_avx2+0x20>
1bf5: 48 83 e1 fc and $0xfffffffffffffffc,%rcx
1bf9: 48 89 ca mov %rcx,%rdx
1bfc: 48 83 c2 04 add $0x4,%rdx
1c00: c5 f8 77 vzeroupper
1c03: 48 39 d6 cmp %rdx,%rsi
1c06: 76 1b jbe 1c23 <reverse64_avx2+0x73>
1c08: 48 8d 04 d7 lea (%rdi,%rdx,8),%rax
1c0c: 48 8d 0c f7 lea (%rdi,%rsi,8),%rcx
1c10: 48 8b 10 mov (%rax),%rdx
1c13: 48 83 c0 08 add $0x8,%rax
1c17: 48 0f ca bswap %rdx
1c1a: 48 89 50 f8 mov %rdx,-0x8(%rax)
1c1e: 48 39 c1 cmp %rax,%rcx
1c21: 75 ed jne 1c10 <reverse64_avx2+0x60>
1c23: c3 retq
1c24: 66 66 2e 0f 1f 84 00 data16 nopw %cs:0x0(%rax,%rax,1)
完整源代码,包含性能测试
参考资料
SIMD 加速:AVX2 指令集实现大小端转换