使用 SIMD 优化 Thrift 编码

发表于 更新于 分类于 阅读次数： Disqus： 本文字数： 1.2k 阅读时长 ≈ 4 分钟

前情提要

可以先看下我之前在 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。