使用 SIMD 优化 Thrift 编码

前情提要

可以先看下我之前在 JTalk 上分享的实践:https://www.bilibili.com/video/BV1UZ4y1g7ju

这篇文章是对于其中我最后说的“使用 SIMD 优化”部分的详细说明。

TL;DR

List<i64> 场景下提升六倍,List<i32> 提升十二倍。

背景

基于 FastRead/Write 接口,由于我们已经拿到了所有的内存,所以我们可以尝试采用 SIMD 来进一步的优化。

思路

最容易想到的优化点也是公司内最常见的用法 list<i64/i32>,这个比较容易想到使用 SIMD 进行优化。

在 thrift binary 里面,int 类型在复制到 buffer 之前需要先转成大端,也就是 binary.BigEndian.PutInt 一次,这个操作原本需要比较多语句,通过软件来模拟,但是在 amd64 下有一个 BSWAP 指令可以直接完成,这个优化 Go 编译器已经做了,所以现在的伪代码如下:

1
2
3
4
var src, dst
for i := 0; i < len; i++ {
dst[i] = bswap(src[i])
}

可以看出来,这个操作实际上是很有规律的,并且全都是相邻的操作,符合 SIMD 指令的模式。

POC

先使用了 C++ 做了一个 POC(只贴了关键代码,完整代码见 https://gist.github.com/PureWhiteWu/e88f241fc8b62df06ae1eb04923a88ae):

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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
const long long int MASK = 0x0001020304050607;
const __mmask16 bit16mask[17] = {0x0000, 0x0001, 0x0003, 0x0007, 0x000f, 0x001f, 0x003f, 0x007f, 0x00ff, 0x01ff, 0x03ff, 0x07ff, 0x0fff, 0x1fff, 0x3fff, 0x7fff, 0xffff};

void avx512_little_2_big(const long long int *src, long long int *dst, int n)
{
int loop_count = n / 8;
int remainder = n % 8;
__m512i mask = _mm512_set1_epi64(MASK);
for (int i = 0; i < loop_count; i++)
{
int index = i * 8;
__m512i input_data = _mm512_loadu_si512(&src[index]);
__m512i output_data = _mm512_shuffle_epi8(input_data, mask);
_mm512_storeu_si512(&avx512_data[index], output_data);
}
if (remainder != 0)
{
int index = loop_count * 8;
__m512i padding = _mm512_set1_epi64(0);
__m512i input_data = _mm512_mask_loadu_epi64(padding, bit16mask[remainder], &src[index]);
__m512i output_data = _mm512_shuffle_epi8(input_data, mask);
_mm512_mask_storeu_epi64(&avx512_data[index], bit16mask[remainder], output_data);
}
return;
}
void avx2_little_2_big(const long long int *src, long long int *dst, int n)
{
int loop_count = n / 4;
int remainder = n % 4;
__m256i mask = _mm256_set1_epi64x(MASK);
for (int i = 0; i < loop_count; i++)
{
int index = i * 4;
__m256i input_data = _mm256_loadu_si256((__m256i *)&src[index]);
__m256i output_data = _mm256_shuffle_epi8(input_data, mask);
_mm256_storeu_si256((__m256i *)&avx2_data[index], output_data);
}
if (remainder != 0)
{
int index = loop_count * 4;
for (int i = index; i < index + remainder; i++)
{
avx2_data[i] = bswap_64(src[i]);
}
}
return;
}

测试结果

编译命令如下:

1
$ g++ little_2_big_gcc.cpp -o ll2 -mavx512f -mavx512bw -mavx2 -mavx -O3

在 linux 物理机上进行测试,结果如下:

1
2
3
avx512 time: 27009 us
avx2 time: 21920 us
bswap time: 49967 us

可以得出结论:

  1. avx512 的性能很不稳定,有些情况下还不如 avx2;
  2. avx2 相比 bswap 方案基本可以提升一倍以上的性能;
  3. Linus 诚不欺我。

详细解释

bswap 做的事情是将整个字节序进行倒序,以 int32 为例,包含 4 字节,假设原来数据如下:

00000000 00000001 00000010 00000011

那么 bswap 之后,数据为:

00000011 00000010 00000001 00000000

在 avx2 中,也有一个指令 vpshufb 能够达到类似的效果,不过不是纯粹的 bswap,详见:https://software.intel.com/content/www/us/en/develop/documentation/cpp-compiler-developer-guide-and-reference/top/compiler-reference/intrinsics/intrinsics-for-intel-advanced-vector-extensions-2/intrinsics-for-shuffle-operations-1/mm256-shuffle-epi8.html

shuffle 的意思是“洗牌”,作用是可以根据一个传入的 mask 来重排对应 byte 的位置。所以这里最关键的就是代码示例中最上面那行:

1
const long long int MASK = 0x0001020304050607;

为什么用这个 mask 就行了呢?我们得复习一下大小端的知识。

大端字节序是符合人类阅读习惯的顺序,高位在前,还是以刚才的 int32 作为例子,假如大端序表示如下:

00000011(高位在这里) 00000010 00000001 00000000

那么在我们电脑上,小端字节序就是这么存的:

内存地址0123(高位在这里)
00000000000000010000001000000011

这时候对应的 MASK 是 0x00010203,在内存中以小端序表示为:

内存地址0123(高位在这里)
3210

我们的机器都是小端序的,所以,在做 shuffle 的时候,内存地址 0 对应的是 内存地址 3 处的值,内存地址 1 对应的是 内存地址 2 处的值,以此类推。

这样,shuffle 计算下来之后,内存中的值就变成了:

内存地址0123
00000011000000100000000100000000

shuffle

这时候,也就相当于成功完成了一次 bswap 的操作了。

由于 int64 有 8 位,所以 MASK 为 0x00 01 02 03 04 05 06 07 就可以完成一次 int64 的 bswap。

(注:没有 0 键在编写此节时遭到虐待)

Go 中测试结果

最后附上 Go 中的测试结果,我们测试了 List 中有 12345 个元素的 benchmark:

1
2
3
4
5
6
7
8
BenchmarkWriteListI64
BenchmarkWriteListI64-16 703928 1753 ns/op
BenchmarkWriteI64
BenchmarkWriteI64-16 98204 11875 ns/op
BenchmarkWriteListI32
BenchmarkWriteListI32-16 1300507 907 ns/op
BenchmarkWriteI32
BenchmarkWriteI32-16 98522 12580 ns/op

可以看出,在 Go 上的性能提升非常巨大,List<i64> 场景下提升六倍,List<i32> 更是提升了十几倍。

究其原因,应该是 Go 做的优化太少太差,远远比不上 gcc。