使用 Intel AVX-512 更快的 bitset 解码

我将“位集解码”称为在位流中查找 1 的位置的操作。例如，给定整数值 0b11011（或十进制的 27），我想找到 0、1、3、4。

在我之前的帖子“使用英特尔 AVX-512 快速位集解码”中，我解释了如何使用英特尔 AVX-512 系列的新指令更快地解码位集。 AVX-512 指令，顾名思义，通常可以处理 512 位（或 64 字节）的寄存器。

至少有两位读者（Kim Walisch 和 Jatin Bhateja）指出，如果您使用带有 Ice Lake 或 Tiger Lake 微架构的英特尔处理器上可用的最新 AVX-512 指令，您可以做得更好。这些处理器支持 VBMI2 指令，包括vpcompressb指令及其相应的内在函数（例如_mm512_maskz_compress_epi8 ）。该指令的作用是获取一个 64 位字和一个 64 字节寄存器，并且它仅输出（以打包方式）与 64 位字中的设置位相对应的字节。因此，如果您使用值 0b11011 作为 64 位字，并提供一个值为 0、1、2、3、4 的 64 字节寄存器……您将得到结果 0、1、3、4。也就是说，该指令已经有效地进行了解码，但需要注意的是它只会写入字节。在实践中，您通常希望索引为 32 位整数。幸运的是，您可以轻松地从压缩字节转换为压缩 32 位整数。一种可能性是提取连续的 128 位子字（使用vextracti32x4指令或其内在_mm512_extracti32x4_epi32 ），并扩展它们（使用vpmovzxbd指令或其内在_mm512_cvtepu8_epi32 ）。您会得到以下结果：

无效vbmi2_decoder_cvtepu8 ( uint32_t * base_ptr , uint32_t & base ,   
                                           uint32_t idx ， uint64_t 位） {   
  __m512i 索引= _mm512_maskz_compress_epi8 (位, _mm512_set_epi32 (   
    0x3f3e3d3c , 0x3b3a3938 , 0x37363534 , 0x33323130 ,   
    0x2f2e2d2c , 0x2b2a2928 , 0x27262524 , 0x23222120 ,   
    0x1f1e1d1c , 0x1b1a1918 , 0x17161514 , 0x13121110 ,   
    0x0f0e0d0c , 0x0b0a0908 , 0x07060504 , 0x03020100   
  ) ) ;   
  __m512i t0 = _mm512_cvtepu8_epi32 （ _mm512_castsi512_si128 （索引） ） ；   
  __m512i t1 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 1 ） ） ；   
  __m512i t2 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 2 ） ） ；   
  __m512i t3 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 3 ） ） ；   
  __m512i start_index = _mm512_set1_epi32 ( idx ) ;   
     
  _mm512_storeu_si512 （ base_ptr + base ， _mm512_add_epi32 （ t0 ， start_index ） ） ；   
  _mm512_storeu_si512 ( base_ptr + base + 16 , _mm512_add_epi32 ( t1 , start_index ) ) ;   
  _mm512_storeu_si512 ( base_ptr + base + 32 , _mm512_add_epi32 ( t2 , start_index ) ) ;   
  _mm512_storeu_si512 ( base_ptr + base + 48 , _mm512_add_epi32 ( t3 , start_index ) ) ;   
     
  基数+ = _popcnt64 (位) ;   
}   

如果您尝试无条件地使用这种方法，您将为解码的每个 64 位字写入 256 个字节的数据。在实践中，如果你的单词大部分只包含零，那么你会写很多零。

分支对性能不利，但只有在难以预测时才会如此。然而，处理器应该很容易预测我们是否在提供的字中设置了少于 16 位、少于 32 位等等。所以一定程度的分支就足够了。以下功能应该做：

 void vbmi2_decoder_cvtepu8_branchy ( uint32_t * base_ptr , uint32_t & base ,   
                                           uint32_t idx ， uint64_t 位） {   
  if (位= = 0 ) {返回; }   
   
  __m512i 索引= _mm512_maskz_compress_epi8 (位, _mm512_set_epi32 (   
    0x3f3e3d3c , 0x3b3a3938 , 0x37363534 , 0x33323130 ,   
    0x2f2e2d2c , 0x2b2a2928 , 0x27262524 , 0x23222120 ,   
    0x1f1e1d1c , 0x1b1a1918 , 0x17161514 , 0x13121110 ,   
    0x0f0e0d0c , 0x0b0a0908 , 0x07060504 , 0x03020100   
  ) ) ;   
  __m512i start_index = _mm512_set1_epi32 ( idx ) ;   
     
  int count = _popcnt64 （位） ；   
  __m512i t0 = _mm512_cvtepu8_epi32 （ _mm512_castsi512_si128 （索引） ） ；   
  _mm512_storeu_si512 （ base_ptr + base ， _mm512_add_epi32 （ t0 ， start_index ） ） ；   
     
  如果（计数> 16 ） {      
    __m512i t1 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 1 ） ） ；   
    _mm512_storeu_si512 ( base_ptr + base + 16 , _mm512_add_epi32 ( t1 , start_index ) ) ;   
    如果（计数> 32 ） {      
      __m512i t2 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 2 ） ） ；   
      _mm512_storeu_si512 ( base_ptr + base + 32 , _mm512_add_epi32 ( t2 , start_index ) ) ;   
      如果（计数> 48 ） {      
        __m512i t3 = _mm512_cvtepu8_epi32 （ _mm512_extracti32x4_epi32 （索引， 3 ） ） ；   
        _mm512_storeu_si512 ( base_ptr + base + 48 , _mm512_add_epi32 ( t3 , start_index ) ) ;   
      }   
    }   
  }   
  基数+ =计数；   
}   

结果会因输入数据而异，但我已经有一个我正在重用的中等密度（大约 10% 的位已设置）的实际案例。使用 Tiger-Lake 处理器和 GCC 9，当使用相当大的输入时，我得到每个设定值的以下时序：

我的代码可用。

这是一个相当了不起的性能，特别是考虑到我们不需要任何大表或复杂的算法。我们所需要的只是花哨的 AVX-512 指令。

原文： https://lemire.me/blog/2022/05/10/faster-bitset-decoding-using-intel-avx-512/