SIMD 加速：AVX2 指令集实现大小端转换

在应用 thrift 进行 RPC 通信的时候，由于 Thrift 采用了大端序，而 x86_64 等常见的处理器架构均为小端序，因此对于 list 这一类的数据类型需要一个循环来实现小端到大端的转换。而这个过程如果能够利用 SIMD 指令的话，可以极大的提高性能。这篇文章是在探索实现 Thrift 编译后端 Auto-vectorization Pass 的时候进行的一个尝试和学习，使用 avx2 指令集实现了一个简单的大小端转换的功能，并且对比了在不同条件下的加速性能。

原理

大小端转换

计算机数据存储有两种字节优先顺序：高位字节优先（称为大端模式）和低位字节优先（称为小端模式）。

• 大端模式，是指数据的高字节保存在内存的低地址中，而数据的低字节保存在内存的高地址中，这样的存储模式有点儿类似于把数据当作字符串顺序处理：地址由小向大增加，而数据从高位往低位放；这和我们的阅读习惯一致。
• 小端模式，是指数据的高字节保存在内存的高地址中，而数据的低字节保存在内存的低地址中，这种存储模式将地址的高低和数据位权有效地结合起来，高地址部分权值高，低地址部分权值低。

例如，对于内存中存放的数0x12345678来说

• 如果是采用大端模式存放的，则其真实的数是：0x12345678
• 如果是采用小端模式存放的，则其真实的数是：0x78563412

可以使用如下 API 进行转换：

#include <arpa/inet.h>
 
uint32_t htonl(uint32_t hostlong);
uint16_t htons(uint16_t hostshort);
uint32_t ntohl(uint32_t netlong);
uint16_t ntohs(uint16_t netshort);
1
2
3
4
5
6

也可以直接使用移位进行实现

inline uint32_t Swap32(uint32_t x)
{
	return (
		((x & 0x000000FF) << 24) |
		((x & 0x0000FF00) << 8) |
		((x & 0x00FF0000) >> 8) |
		((x & 0xFF000000) >> 24));
}
1
2
3
4
5
6
7
8

bswap

大部分编译器同时提供了 bswap 指令，来帮助实现这一转换过程，例如在 gcc 中，我们可以使用 __builtin_bswap{16,32,64}：

inline uint32_t Swap32(uint32_t x)
{
	return __builtin_bswap32(x);
}
1
2
3
4

这是一个编译器的内置函数。在 x86_64 机器上，它会被编译为这样的指令序列：

Swap32(unsigned int):
  mov eax, edi
  bswap eax
  ret
1
2
3
4

在 arm 机器上，它会被编译为这样的指令序列：

Swap32(unsigned int):
  rev w0, w0
  ret
1
2
3

通常来说，我们自己使用的移位函数实现的大小端转换，在编译器优化 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 位的字节序列完全反转：

const __m256i ShuffleRev = _mm256_set_epi8(
	0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,15, // first 128-bit lane
	0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,15  // second 128-bit lane
);
// Load 32 elements at once into one 32-byte register
__m256i x = _mm256_loadu_si256(
	reinterpret_cast<__m256i*>(&Array8[i])
);
// Reverse each the bytes in each 128-bit lane
x = _mm256_shuffle_epi8(x,ShuffleRev);
1
2
3
4
5
6
7
8
9
10

它接受一个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 指令集进行加速的简单示例：

void reverse64_avx2(void *Array, size_t Count)
{
	uint64_t *Array64 = (uint64_t *)(Array);
	size_t i = 0;
	const __m256i ShuffleRev = _mm256_set_epi8(
		8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7,
		8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7);

	for (; i < (Count & ~0b11); i += 4)
	{
		// Load 4 elements
		__m256i bytes = _mm256_loadu_si256(
			(__m256i *)(&Array64[i]));

		// Reverse the byte order of our 64-byte vectors
		bytes = _mm256_shuffle_epi8(bytes, ShuffleRev);

		// Place them back into the array
		_mm256_storeu_si256(
			(__m256i *)(&Array64[i]),
			bytes);
	}

	// Naive swaps for leftover elements
	for (; i < Count; ++i)
	{
		Array64[i] = Swap64(Array64[i]);
	}
}
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

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

编译指令：

gcc main.c  -mavx2 -fno-tree-vectorize -O3 -o avx 
1

basic 对照函数（这里 Swap64 会被 gcc 自动编译为 bswap 指令）：

void reverse64_basic(void *Array, size_t Count)
{
	uint64_t *Array64 = (uint64_t *)(Array);
	size_t i = 0;

	// Naive swaps
	for (; i < Count; ++i)
	{
		Array64[i] = Swap64(Array64[i]);
	}
}
1
2
3
4
5
6
7
8
9
10
11

我们分别对 64/32/16 位的整数进行大小端转换，并测试 bswap 和 avx2 的加速比：

可以注意到，对于宽度更小的整数的数组，并行度更高，avx2 加速比更加明显。在 64 位时，加速比约为 2.5，在 16 位时，加速比已经达到了 10 倍。

生成的汇编：

objdump -d ./avx > dump.s

我们可以再检查一下生成的汇编指令：

• 使用 bswap 的大小端转换

  0000000000001c30 <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)
  1
  2
  3
  4
  5
  6
  7
  8
  9
  10
  11
  12
  13
  14

• avx2：vpshufb

  0000000000001bb0 <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)
  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

完整源代码，包含性能测试：

• https://github.com/yunwei37/ZJU-CS-GIS-ClassNotes/blob/master/%E4%B8%80%E4%BA%9B%E6%9D%82%E9%A1%B9/demos/avx2/main.c

参考资料

• https://github.com/Wunkolo/qreverse
• https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html