计算 __m256i 字中的前导零

Question

我正在修改 AVX-2 指令，并且正在寻找一种快速方法来计算 __m256i 中前导零的数量字(有 256 位)。

到目前为止，我已经想出了以下方法:

// Computes the number of leading zero bits.
// Here, avx_word is of type _m256i.

if (!_mm256_testz_si256(avx_word, avx_word)) {
  uint64_t word = _mm256_extract_epi64(avx_word, 0);
  if (word > 0)
    return (__builtin_clzll(word));

  word = _mm256_extract_epi64(avx_word, 1);
  if (word > 0)
    return (__builtin_clzll(word) + 64);

  word = _mm256_extract_epi64(avx_word, 2);
  if (word > 0)
    return (__builtin_clzll(word) + 128);

  word = _mm256_extract_epi64(avx_word, 3);
  return (__builtin_clzll(word) + 192);
} else
  return 256; // word is entirely zero

但是，我发现在 256 位寄存器中找出确切的非零字相当笨拙。

有人知道是否有更优雅(或更快)的方法来做到这一点？

就像附加信息一样:
我实际上想计算由逻辑与创建的任意长 vector 的第一个设置位的索引，并且我正在比较标准 64 位操作与 SSE 和 AVX-2 代码的性能。
这是我的整个测试代码:

#include <stdio.h>
#include <stdlib.h>
#include <immintrin.h>
#include <stdint.h>
#include <assert.h>
#include <time.h>
#include <sys/time.h>
#include <stdalign.h>

#define ALL  0xFFFFFFFF
#define NONE 0x0


#define BV_SHIFTBITS ((size_t)    6)
#define BV_MOD_WORD  ((size_t)   63)
#define BV_ONE       ((uint64_t)  1)
#define BV_ZERO      ((uint64_t)  0)
#define BV_WORDSIZE  ((uint64_t) 64)


uint64_t*
Vector_new(
    size_t num_bits) {

  assert ((num_bits % 256) == 0);
  size_t num_words = num_bits >> BV_SHIFTBITS;
  size_t mod = num_bits & BV_MOD_WORD;
  if (mod > 0)
    assert (0);
  uint64_t* words;
  posix_memalign((void**) &(words), 32, sizeof(uint64_t) * num_words);
  for (size_t i = 0; i < num_words; ++i)
    words[i] = 0;
  return words;
}


void
Vector_set(
    uint64_t* vector,
    size_t pos) {

  const size_t word_index = pos >> BV_SHIFTBITS;
  const size_t offset     = pos & BV_MOD_WORD;
  vector[word_index] |= (BV_ONE << (BV_MOD_WORD - offset));
}


size_t
Vector_and_first_bit(
    uint64_t** vectors,
    const size_t num_vectors,
    const size_t num_words) {

  for (size_t i = 0; i < num_words; ++i) {
    uint64_t word = vectors[0][i];
    for (size_t j = 1; j < num_vectors; ++j)
      word &= vectors[j][i];
    if (word > 0)
      return (1 + i * BV_WORDSIZE + __builtin_clzll(word));
  }
  return 0;
}


size_t
Vector_and_first_bit_256(
    uint64_t** vectors,
    const size_t num_vectors,
    const size_t num_avx_words) {

  for (size_t i = 0; i < num_avx_words; ++i) {
    const size_t addr_offset = i << 2;
    __m256i avx_word = _mm256_load_si256(
        (__m256i const*) (vectors[0] + addr_offset));

    // AND the AVX words
    for (size_t j = 1; j < num_vectors; ++j) {
      avx_word = _mm256_and_si256(
        avx_word,
        _mm256_load_si256((__m256i const*) (vectors[j] + addr_offset))
      );
    }

    // test whether resulting AVX word is not zero
    if (!_mm256_testz_si256(avx_word, avx_word)) {
      uint64_t word = _mm256_extract_epi64(avx_word, 0);
      const size_t shift = i << 8;
      if (word > 0)
        return (1 + shift + __builtin_clzll(word));

      word = _mm256_extract_epi64(avx_word, 1);
      if (word > 0)
        return (1 + shift + __builtin_clzll(word) + 64);

      word = _mm256_extract_epi64(avx_word, 2);
      if (word > 0)
        return (1 + shift + __builtin_clzll(word) + 128);

      word = _mm256_extract_epi64(avx_word, 3);
      return (1 + shift + __builtin_clzll(word) + 192);
    }
  }
  return 0;
}


size_t
Vector_and_first_bit_128(
    uint64_t** vectors,
    const size_t num_vectors,
    const size_t num_avx_words) {

  for (size_t i = 0; i < num_avx_words; ++i) {
    const size_t addr_offset = i << 1;
    __m128i avx_word = _mm_load_si128(
        (__m128i const*) (vectors[0] + addr_offset));

    // AND the AVX words
    for (size_t j = 1; j < num_vectors; ++j) {
      avx_word = _mm_and_si128(
        avx_word,
        _mm_load_si128((__m128i const*) (vectors[j] + addr_offset))
      );
    }

    // test whether resulting AVX word is not zero
    if (!_mm_test_all_zeros(avx_word, avx_word)) {
      uint64_t word = _mm_extract_epi64(avx_word, 0);
      if (word > 0)
        return (1 + (i << 7) + __builtin_clzll(word));

      word = _mm_extract_epi64(avx_word, 1);
      return (1 + (i << 7) + __builtin_clzll(word) + 64);
    }
  }
  return 0;
}


uint64_t*
make_random_vector(
    const size_t num_bits,
    const size_t propability) {

  uint64_t* vector = Vector_new(num_bits);
  for (size_t i = 0; i < num_bits; ++i) {
    const int x = rand() % 10;
    if (x >= (int) propability)
      Vector_set(vector, i);
  }
  return vector;
}


size_t
millis(
    const struct timeval* end,
    const struct timeval* start) {

  struct timeval e = *end;
  struct timeval s = *start;
  return (1000 * (e.tv_sec - s.tv_sec) + (e.tv_usec - s.tv_usec) / 1000);
}


int
main(
    int argc,
    char** argv) {

  if (argc != 6)
    printf("fuck %s\n", argv[0]);

  srand(time(NULL));

  const size_t num_vectors = atoi(argv[1]);
  const size_t size = atoi(argv[2]);
  const size_t num_iterations = atoi(argv[3]);
  const size_t num_dimensions = atoi(argv[4]);
  const size_t propability = atoi(argv[5]);
  const size_t num_words = size / 64;
  const size_t num_sse_words = num_words / 2;
  const size_t num_avx_words = num_words / 4;

  assert(num_vectors > 0);
  assert(size > 0);
  assert(num_iterations > 0);
  assert(num_dimensions > 0);

  struct timeval t1;
  gettimeofday(&t1, NULL);

  uint64_t*** vectors = (uint64_t***) malloc(sizeof(uint64_t**) * num_vectors);
  for (size_t j = 0; j < num_vectors; ++j) {
    vectors[j] = (uint64_t**) malloc(sizeof(uint64_t*) * num_dimensions);
    for (size_t i = 0; i < num_dimensions; ++i)
      vectors[j][i] = make_random_vector(size, propability);
  }

  struct timeval t2;
  gettimeofday(&t2, NULL);
  printf("Creation: %zu ms\n", millis(&t2, &t1));



  size_t* results_64    = (size_t*) malloc(sizeof(size_t) * num_vectors);
  size_t* results_128   = (size_t*) malloc(sizeof(size_t) * num_vectors);
  size_t* results_256   = (size_t*) malloc(sizeof(size_t) * num_vectors);


  gettimeofday(&t1, NULL);
  for (size_t j = 0; j < num_iterations; ++j)
    for (size_t i = 0; i < num_vectors; ++i)
      results_64[i] = Vector_and_first_bit(vectors[i], num_dimensions,
          num_words);
  gettimeofday(&t2, NULL);
  const size_t millis_64 = millis(&t2, &t1);
  printf("64            : %zu ms\n", millis_64);


  gettimeofday(&t1, NULL);
  for (size_t j = 0; j < num_iterations; ++j)
    for (size_t i = 0; i < num_vectors; ++i)
      results_128[i] = Vector_and_first_bit_128(vectors[i],
          num_dimensions, num_sse_words);
  gettimeofday(&t2, NULL);
  const size_t millis_128 = millis(&t2, &t1);
  const double factor_128 = (double) millis_64 / (double) millis_128;
  printf("128           : %zu ms (factor: %.2f)\n", millis_128, factor_128);

  gettimeofday(&t1, NULL);
  for (size_t j = 0; j < num_iterations; ++j)
    for (size_t i = 0; i < num_vectors; ++i)
      results_256[i] = Vector_and_first_bit_256(vectors[i],
          num_dimensions, num_avx_words);
  gettimeofday(&t2, NULL);
  const size_t millis_256 = millis(&t2, &t1);
  const double factor_256 = (double) millis_64 / (double) millis_256;
  printf("256           : %zu ms (factor: %.2f)\n", millis_256, factor_256);


  for (size_t i = 0; i < num_vectors; ++i) {
    if (results_64[i] != results_256[i])
      printf("ERROR: %zu (64) != %zu (256) with i = %zu\n", results_64[i],
          results_256[i], i);
    if (results_64[i] != results_128[i])
      printf("ERROR: %zu (64) != %zu (128) with i = %zu\n", results_64[i],
          results_128[i], i);
  }


  free(results_64);
  free(results_128);
  free(results_256);

  for (size_t j = 0; j < num_vectors; ++j) {
    for (size_t i = 0; i < num_dimensions; ++i)
      free(vectors[j][i]);
    free(vectors[j]);
  }
  free(vectors);
  return 0;
}

编译:

gcc -o main main.c -O3 -Wall -Wextra -pedantic-errors -Werror -march=native -std=c99 -fno-tree-vectorize

执行:

./main 1000 8192 50000 5 9

参数的意思是:1000 个测试用例，长度为 8192 位的 vector ，50000，测试重复(最后两个参数是小调整)。

在我的机器上执行上述调用的示例输出:

Creation: 363 ms
64            : 15000 ms
128           : 10070 ms (factor: 1.49)
256           : 6784 ms (factor: 2.21)

Answer 1

如果您的输入值均匀分布，则几乎所有时间最高设置位都将位于 vector 的前 64 位(2^64 中的 1 个)。在这种情况下的分支将预测得很好。 @Nejc's answer is good for that case 。

但是 lzcnt 是解决方案一部分的许多问题都有均匀分布的输出(或类似的)，因此无分支版本具有优势。不是严格统一的，而是最高设置位通常位于最高 64 位以外的任何地方。

Wim 在比较位图上使用 lzcnt 来查找正确元素的想法是一种非常好的方法。

但是，带有存储/重新加载的 vector 的 运行时变量索引可能比随机播放 更好。存储转发延迟很低(在 Skylake 上可能是 5 到 7 个周期)，并且该延迟与索引生成并行(比较/movemask/lzcnt)。 movd/vpermd/movd 车道交叉洗牌策略在索引已知后需要 5 个周期，才能将正确的元素放入整数寄存器中。 (见 http://agner.org/optimize/ )

我认为这个版本在 Haswell/Skylake(和 Ryzen)上应该有更好的延迟，以及更好的吞吐量 。 ( vpermd 在 Ryzen 上很慢，所以它在那里应该非常好)加载的地址计算应该与存储转发具有相似的延迟，所以这是一个折腾，哪个实际上是关键路径。

将堆栈按 32 对齐以避免 32 字节存储上的缓存行拆分需要额外的指令，因此最好是它可以内联到多次使用它的函数中，或者已经需要对其他一些 __m256i 进行大量对齐。

#include <stdint.h>
#include <immintrin.h>

#ifndef _MSC_VER
#include <stdalign.h>  //MSVC is missing this?
#else
#include <intrin.h>
#pragma intrinsic(_BitScanReverse)  // https://msdn.microsoft.com/en-us/library/fbxyd7zd.aspx suggests this
#endif

// undefined result for mask=0, like BSR
uint32_t bsr_nonzero(uint32_t mask)
{
// on Intel, bsr has a minor advantage for the first step
// for AMD, BSR is slow so you should use 31-LZCNT.

   //return 31 - _lzcnt_u32(mask);
 // Intel's docs say there should be a _bit_scan_reverse(x), maybe try that with ICC

   #ifdef _MSC_VER
     unsigned long tmp;
     _BitScanReverse(&tmp, mask);
     return tmp;
   #else
     return 31 - __builtin_clz(mask);
   #endif
}

和有趣的部分 :

int mm256_lzcnt_si256(__m256i vec)
{
    __m256i   nonzero_elem = _mm256_cmpeq_epi8(vec, _mm256_setzero_si256());
    unsigned  mask = ~_mm256_movemask_epi8(nonzero_elem);

    if (mask == 0)
        return 256;  // if this is rare, branching is probably good.

    alignas(32)  // gcc chooses to align elems anyway, with its clunky code
    uint8_t elems[32];
    _mm256_storeu_si256((__m256i*)elems, vec);

//    unsigned   lz_msk   = _lzcnt_u32(mask);
//    unsigned   idx = 31 - lz_msk;          // can use bsr to get the 31-x, because mask is known to be non-zero.
//  This takes the 31-x latency off the critical path, in parallel with final lzcnt
    unsigned   idx = bsr_nonzero(mask);
    unsigned   lz_msk = 31 - idx;
    unsigned   highest_nonzero_byte = elems[idx];
    return     lz_msk * 8 + _lzcnt_u32(highest_nonzero_byte) - 24;
               // lzcnt(byte)-24, because we don't want to count the leading 24 bits of padding.
}    

On Godbolt with gcc7.3 -O3 -march=haswell ，我们得到这样的 asm 将 ymm1 计入 esi 。

        vpxor   xmm0, xmm0, xmm0
        mov     esi, 256
        vpcmpeqd        ymm0, ymm1, ymm0
        vpmovmskb       eax, ymm0
        xor     eax, -1                      # ~mask and set flags, unlike NOT
        je      .L35
        bsr     eax, eax
        vmovdqa YMMWORD PTR [rbp-48], ymm1   # note no dependency on anything earlier; OoO exec can run it early
        mov     ecx, 31
        mov     edx, eax                     # this is redundant, gcc should just use rax later.  But it's zero-latency on HSW/SKL and Ryzen.
        sub     ecx, eax
        movzx   edx, BYTE PTR [rbp-48+rdx]   # has to wait for the index in edx
        lzcnt   edx, edx
        lea     esi, [rdx-24+rcx*8]          # lzcnt(byte) + lzcnt(vectormask) * 8
.L35:

为了找到最高的非零元素( 31 - lzcnt(~movemask) )， 我们使用 bsr 直接获取位(以及字节)索引，并从关键路径 中减去。只要我们在掩码为零上进行分支，这就是安全的。 (无分支版本需要初始化寄存器以避免越界索引)。

在 AMD CPU 上， bsr 明显慢于 lzcnt 。在 Intel CPU 上，它们的性能相同，除了 output-dependency details 的细微变化。

输入为零的 bsr 使目标寄存器保持不变，但 GCC 不提供利用这一点的方法。 (Intel 仅将其记录为未定义的输出，但 AMD 将 Intel/AMD CPU 的实际行为记录为在目标寄存器中生成旧值)。

如果输入为零，则 bsr 设置 ZF，而不是像大多数指令一样基于输出。 (这和输出依赖性可能是它在 AMD 上缓慢的原因。)在 BSR 标志上进行分支并不比在 xor eax,-1 设置的 ZF 上分支以反转掩码更好，这正是 gcc 所做的。无论如何，英特尔确实会返回 _BitScanReverse(&idx, mask) 的 document a bool intrinsic ，但 gcc 不支持它(甚至不支持 x86intrin.h )。 GNU C 内置函数不会返回一个 bool 值来让您使用标志结果，但如果您检查输入的 C 变量是否为非零，则 gcc 可能会使用 bsr 的标志输出来制作智能汇编。

使用 dword ( uint32_t ) 数组和 vmovmskps 将使第二个 lzcnt 使用内存源操作数，而不需要 movzx 对单个字节进行零扩展。但是 lzcnt 在 Skylake 之前对 Intel CPU 有错误的依赖性，因此编译器可能倾向于单独加载并使用 lzcnt same,same 作为解决方法。 (我没查。)

Wim 的版本需要 lz_msk-24，因为高 24 位始终为零并带有 8 位掩码。但是 32 位掩码填充 32 位寄存器。

这个具有 8 位元素和 32 位掩码的版本是相反的:我们需要对所选字节进行 lzcnt，不包括寄存器中的 24 个前导零位。所以我们的 -24 移动到不同的位置，而不是索引数组的关键路径的一部分。

gcc 选择将其作为单个 3 组件 LEA ( reg + reg*scale - const ) 的一部分来执行，这对吞吐量非常有用，但将其放在最终 lzcnt 之后的关键路径上。 (它不是免费的，因为 3 组件 LEA 与 Intel CPU 上的 reg + reg*scale 相比具有额外的延迟。请参阅 Agner Fog's instruction tables )。

乘以 8 可以作为 lea 的一部分完成，但乘以 32 需要移位(或折叠成两个单独的 LEA)。

Intel's optimization manual 说(表 2-24)即使 Sandybridge 也可以从 256 位存储转发到单字节加载没有问题，所以我认为它在 AVX2 CPU 上很好，与转发到 32 位加载的 4 字节加载相同- 对齐的存储块。