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);

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

inline uint32_t Swap32(uint32_t x)
{
 return (
  ((x & 0x000000FF) << 24) |
  ((x & 0x0000FF00) << 8) |
  ((x & 0x00FF0000) >> 8) |
  ((x & 0xFF000000) >> 24));
}

bswap

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

inline uint32_t Swap32(uint32_t x)
{
 return __builtin_bswap32(x);
}

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

Swap32(unsigned int):
  mov eax, edi
  bswap eax
  ret

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

Swap32(unsigned int):
  rev w0, w0
  ret

通常来说，我们自己使用的移位函数实现的大小端转换，在编译器优化 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);

它接受一个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]);
 }
}

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 

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]);
 }
}

我们分别对 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)
  

• 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)
  

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

• main.c

参考资料

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