当前位置：首页 > news >正文

ValKey中使用SIMD指令优化bitcount命令

news 来源：原创 2025/7/4 3:21:13

一、AVX/AVX2 的历史演进

随着计算机技术的飞速发展，数据处理需求呈指数级增长，SIMD（单指令多数据）技术应运而生。它通过一条指令同时处理多个数据元素，大幅提升计算效率，从早期的 MMX 技术起步，历经 SSE 系列的逐步完善，不断发展成熟。
2011年，英特尔推出 AVX（Advanced Vector Extensions）指令集，这是 SIMD 技术发展的重要里程碑。AVX 将寄存器宽度从 SSE 的128位扩展至256位，使处理器能够并行处理更多数据，极大地增强了并行计算能力，为高性能计算领域带来了重大突破，也为后续的技术演进奠定了坚实基础。
2013年，AVX2作为 AVX 的增强版本问世。AVX2 进一步扩展了整数指令集，新增了一系列乘法、除法等位操作指令，显著优化了整数运算性能；同时，它强化了数据读写机制，提升了内存访问效率，实现了计算性能的全方位升级。

二、AVX/AVX2 的原理剖析

（一）寄存器架构

AVX/AVX2的核心优势体现在256位YMM 寄存器上。相较于 128 位寄存器，YMM 寄存器的数据承载量翻倍。例如，在处理单精度浮点数时，存储量从 4 个提升至 8 个；对于整数数据的存储能力同样实现倍增，这为高效的数据并行处理提供了强大的硬件支撑。

（二）指令集特性

数据并行指令：AVX/AVX2 支持丰富的数据并行指令，涵盖算术运算、逻辑操作、移位运算等多个方面。单条指令即可同时处理8个单精度浮点数，相比传统指令执行方式，大幅提升了数据处理速度。
融合乘法累加（FMA）指令：AVX2 引入的 FMA 指令是一大亮点，它将乘法和累加运算合二为一，有效减少了指令执行周期。在科学计算、深度学习等对计算性能要求极高的领域，FMA 指令发挥着关键作用。
灵活的数据操作：AVX/AVX2 能够兼容整数、浮点数等多种数据类型，支持灵活的数据转换与对齐操作。这使其能够适配不同的内存访问需求，进一步优化数据读取效率，提升整体计算性能。

（三）数据处理流程

当执行 AVX/AVX2 指令时，数据首先从内存加载至256位 YMM 寄存器，一次可读取多个数据元素。随后，处理器根据指令类型，对寄存器中的数据进行并行运算处理。运算完成后，将结果写回内存或存储至其他寄存器。整个过程通过数据并行处理，极大地提高了计算资源利用率和数据处理效率。

三、AVX/AVX2 的广泛应用

（一）科学计算领域

科学计算常常涉及大规模的矩阵运算、数值模拟等复杂计算任务。AVX/AVX2 利用其并行指令特性，能够并行处理矩阵元素的乘累加运算。与传统的循环计算方式相比，计算效率得到显著提升。在气象预测、物理模拟等科研场景中，AVX/AVX2 加速了复杂计算过程，大幅缩短了科研时间，提高了科研效率。

（二）多媒体处理领域

音频、视频等多媒体数据具有数据量大、处理复杂度高的特点。在视频编解码过程中，AVX/AVX2 可以并行处理多个像素的变换、滤波等操作，实现视频的流畅播放和高效编码；在音频处理方面，AVX/AVX2 加速了滤波、混音等操作，有效提升了音频质量。

（三）深度学习领域

深度学习模型的训练和推理过程依赖大量的矩阵、向量运算。AVX/AVX2 的并行指令能够加速神经网络中的梯度计算和权重更新等关键操作。TensorFlow、PyTorch 等主流深度学习框架对 AVX/AVX2 进行优化后，在图像识别、语音识别等应用领域，模型的训练和推理速度得到大幅提升。

（四）数据压缩与解压缩领域

在数据存储和传输过程中，AVX/AVX2 能够加速 LZMA、Zstandard 等压缩算法的执行。通过并行处理数据，实现了快速的数据压缩和解压缩，显著缩短了数据存储和传输时间，提高了数据处理效率。
AVX/AVX2 作为先进的 SIMD 技术，在不断的发展演进中，凭借其独特的技术原理和强大的性能优势，在众多领域得到了广泛应用。随着技术的持续进步，AVX/AVX2 有望在更多领域发挥更大的作用，为计算机技术的发展和各行业的进步提供强大动力。

四、ValKey中的应用

以下代码通过使用 AVX2 指令集实现了一个高效的位计数（popcount）函数popcountAVX2，用于计算给定内存区域中设置为 1 的位的总数。以下是对代码中如何应用 AVX2 的详细解释：

1. 定义目标架构属性

ATTRIBUTE_TARGET_AVX2

这行代码用于明确指定该函数针对 AVX2 架构进行优化，确保编译器生成适配 AVX2 指令集的机器码，从而充分发挥 AVX2 架构的性能优势。

2. 初始化和数据准备

long long popcountAVX2(void *s, long count) {long i = 0;unsigned char *p = (unsigned char *)s;long long bits = 0;/* clang-format off */const __m256i lookup = _mm256_setr_epi8(/* First Lane [0:127] *//* 0 */ 0, /* 1 */ 1, /* 2 */ 1, /* 3 */ 2,/* 4 */ 1, /* 5 */ 2, /* 6 */ 2, /* 7 */ 3,/* 8 */ 1, /* 9 */ 2, /* a */ 2, /* b */ 3,/* c */ 2, /* d */ 3, /* e */ 3, /* f */ 4,/* Second Lane [128:255] identical to first lane due to lane isolation in _mm256_shuffle_epi8.* For more information, see following URL* https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_shuffle_epi8 *//* 0 */ 0, /* 1 */ 1, /* 2 */ 1, /* 3 */ 2,/* 4 */ 1, /* 5 */ 2, /* 6 */ 2, /* 7 */ 3,/* 8 */ 1, /* 9 */ 2, /* a */ 2, /* b */ 3,/* c */ 2, /* d */ 3, /* e */ 3, /* f */ 4);/* clang-format on */const __m256i low_mask = _mm256_set1_epi8(0x0f);__m256i acc = _mm256_setzero_si256();

lookup变量是一个__m256i类型的向量，存储了一个查找表。通过这个查找表，可以快速获取每个字节中 1 的个数，为后续的位计数操作提供便利。
low_mask是一个全为0x0f的向量，其作用是提取每个字节的低 4 位，以便进行后续的位计数计算。
acc是一个初始化为 0 的累加器向量，用于存储位计数过程中的中间计算结果。

3. 定义数据处理宏

/* Count 32 bytes per iteration. */
#define ITER_32_BYTES                                                             \{                                                                             \const __m256i vec = _mm256_loadu_si256((const __m256i *)(p + i));         \const __m256i lo = _mm256_and_si256(vec, low_mask);                       \const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(vec, 4), low_mask); \const __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);                  \const __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);                  \local = _mm256_add_epi8(local, popcnt1);                                  \local = _mm256_add_epi8(local, popcnt2);                                  \i += 32;                                                                  \}

_mm256_loadu_si256：从内存中加载32字节的数据到__m256i类型的向量vec中，其中u表示非对齐加载，这种方式可以灵活地从内存中读取数据。
_mm256_and_si256：将向量vec与low_mask进行按位与操作，分别提取每个字节的低4位和高4位，为后续的位计数查找做准备。
_mm256_srli_epi16：将向量vec的每个16位元素右移4位，从而提取出每个字节的高4位。
_mm256_shuffle_epi8：根据lookup表以及提取的低4位和高4位，查找每个字节中1的个数，并返回结果向量。
_mm256_add_epi8：将查找得到的位计数结果累加到local向量中，实现位计数的累加操作。

4. 数据分块处理

/* Part A: loop unrolling, processing 8 * 32 bytes per iteration. */while (i + 8 * 32 <= count) {__m256i local = _mm256_setzero_si256();ITER_32_BYTESITER_32_BYTESITER_32_BYTESITER_32_BYTESITER_32_BYTESITER_32_BYTESITER_32_BYTESITER_32_BYTESacc = _mm256_add_epi64(acc, _mm256_sad_epu8(local, _mm256_setzero_si256()));}/* Part B: when the remaining data length is less than 8 * 32 bytes,* process 32 bytes per iteration. */__m256i local = _mm256_setzero_si256();while (i + 32 <= count) {ITER_32_BYTES;}acc = _mm256_add_epi64(acc, _mm256_sad_epu8(local, _mm256_setzero_si256()));#undef ITER_32_BYTES

Part A：通过循环展开技术，每次处理8个32字节的数据块，充分利用 AVX2 的并行处理能力，提高指令执行并行度，加快位计数计算速度。
Part B：当剩余数据长度小于8*32字节时，每次处理32字节的数据，确保对所有数据进行完整的位计数处理。
_mm256_sad_epu8：计算两个向量对应元素的绝对差值之和，在这里用于将local向量中的字节计数结果进行累加。
_mm256_add_epi64：将累加结果累加到acc向量中，实现位计数结果的逐步汇总。

5. 结果提取和剩余数据处理

    bits += _mm256_extract_epi64(acc, 0);bits += _mm256_extract_epi64(acc, 1);bits += _mm256_extract_epi64(acc, 2);bits += _mm256_extract_epi64(acc, 3);/* Part C: count the remaining bytes. */for (; i < count; i++) {bits += bitsinbyte[p[i]];}return bits;

_mm256_extract_epi64：从acc向量中提取64位整数，并累加到bits变量中，将acc向量中的位计数结果进行汇总。
Part C：对于剩余的字节，通过传统的数组查找方式（bitsinbyte数组）计算位计数，并累加到bits中。最终，返回总的位计数结果，完成对给定内存区域中1的位总数的计算。

通过以上步骤，代码充分利用了 AVX2 指令集的向量处理能力，通过批量处理数据和循环展开技术，显著提高了位计数操作的效率。

五、优化效果

Value Size	QPS (After optimization)	QPS (Before optimization)	change
16 B	114925	115924	-0.8%
256 B	112619	112201	+0.3%
4 KB	105523	96251	+9.6%
64 KB	79723	36796	+116%
1MB	21306	3466	+514%

CPU: AMD EPYC 9754 128-Core Processor * 8
OS: Ubuntu Server 22.04 LTS 64bit
Memory: 16GB
VM: Tencent cloud SA5.2XLARGE16