计算 __m256i 个单词中的前导零

Count leading zeros in __m256i word

我正在研究 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 位寄存器中找出确切的非零字相当笨拙。

有人知道是否有更优雅(或更快)的方法吗?

作为附加信息: 我实际上想计算由逻辑与创建的任意长向量的第一个设置位的索引,并且我正在将标准 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 位的向量,50000,测试重复(最后两个参数是小调整)。

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

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

(更新:自 2019-01-31 以来的新答案)

三个备选方案是:

  • 。快速地。 这个解决方案不是无分支的,这应该不是问题,除非输入 经常为零,出现的模式不规则。

  • 我之前的答案现在在这个答案的edit history中。效率较低 比彼得科德斯的回答,但无分支。

  • 这个回答。如果来自 2 个微型查找 table 的数据位于 L1 缓存中,则速度非常快。 L1 高速缓存占用空间为 128 字节。无枝。它可能会遭受缓存未命中 当不经常调用时。

在此答案中,将输入 epi64 向量与零进行比较,从而生成掩码。 此掩码被转换为 4 位索引 i_mask(通过 _mm256_movemask_pd)。 使用索引 i_mask 从两个查找 table 中读取两个值: 1. 第一个非零 64 位元素的索引,以及 2. 前面(从左到右)零元素的非零数。 最后,计算并添加第一个非零 64 位元素的 _lzcnt_u64 查找 table 值。函数 mm256_lzcnt_si256 实现了这个方法:

#include <stdio.h>
#include <stdint.h>
#include <x86intrin.h>
#include <stdalign.h>
/* gcc -Wall -m64 -O3 -march=haswell clz_avx256_upd.c */


int mm256_lzcnt_si256(__m256i input)
{   
    /* Version with lookup tables and scratch array included in the function                                                                  */

    /* Two tiny lookup tables (64 bytes each, less space is possible with uint8_t or uint16_t arrays instead of uint32_t):                       */
    /* i_mask  (input==0)                 0000 0001 0010 0011 0100 0101 0110 0111 1000 1001 1010 1011 1100 1101 1110 1111                        */
    /* ~i_mask (input!=0)                 1111 1110 1101 1100 1011 1010 1001 1000 0111 0110 0101 0100 0011 0010 0001 0000                        */
    static const uint32_t indx[16]   = {   3,   3,   3,   3,   3,   3,   3,   3,   2,   2,   2,   2,   1,   1,   0,   0};
    static const uint32_t lz_msk[16] = {   0,   0,   0,   0,   0,   0,   0,   0,  64,  64,  64,  64, 128, 128, 192, 192};

    alignas(32)  uint64_t tmp[4]     = {   0,   0,   0,   0};                /* tmp is a scratch array of 32 bytes, preferably 32 byte aligned   */ 

                          _mm256_storeu_si256((__m256i*)&tmp[0], input);     /* Store input in the scratch array                                 */
    __m256i  mask       = _mm256_cmpeq_epi64(input, _mm256_setzero_si256()); /* Check which 64 bits elements are zero                            */
    uint32_t i_mask     = _mm256_movemask_pd(_mm256_castsi256_pd(mask));     /* Move vector mask to integer mask                                 */
    uint64_t input_i    = tmp[indx[i_mask]];                                 /* Load the first (from the left) non-zero 64 bit element input_i   */
    int32_t  lz_input_i = _lzcnt_u64(input_i);                               /* Count the number of leading zeros in input_i                     */
    int32_t  lz         = lz_msk[i_mask] + lz_input_i;                       /* Add the number of leading zeros of the preceding 64 bit elements */
             return lz;
}    


int mm256_lzcnt_si256_v2(__m256i input, uint64_t* restrict tmp, const uint32_t* indx, const uint32_t* lz_msk)
{   
    /* Version that compiles to nice assembly, although, after inlining there won't be any difference between the different versions.            */
                          _mm256_storeu_si256((__m256i*)&tmp[0], input);     /* Store input in the scratch array                                 */
    __m256i  mask       = _mm256_cmpeq_epi64(input, _mm256_setzero_si256()); /* Check which 64 bits elements are zero                            */
    uint32_t i_mask     = _mm256_movemask_pd(_mm256_castsi256_pd(mask));     /* Move vector mask to integer mask                                 */
    uint64_t input_i    = tmp[indx[i_mask]];                                 /* Load the first (from the left) non-zero 64 bit element input_i   */
    int32_t  lz_input_i = _lzcnt_u64(input_i);                               /* Count the number of leading zeros in input_i                     */
    int32_t  lz         = lz_msk[i_mask] + lz_input_i;                       /* Add the number of leading zeros of the preceding 64 bit elements */
             return lz;
}    


__m256i bit_mask_avx2_lsb(unsigned int n)               
{           
    __m256i ones       = _mm256_set1_epi32(-1);
    __m256i cnst32_256 = _mm256_set_epi32(256,224,192,160, 128,96,64,32);
    __m256i shift      = _mm256_set1_epi32(n);   
            shift      = _mm256_subs_epu16(cnst32_256,shift);  
                  return _mm256_srlv_epi32(ones,shift);
}


int print_avx2_hex(__m256i ymm)
{
    long unsigned int x[4];
        _mm256_storeu_si256((__m256i*)x,ymm);
        printf("%016lX %016lX %016lX %016lX  ", x[3],x[2],x[1],x[0]);
    return 0;
}


int main()
{
    unsigned int i;
    __m256i x;

    printf("mm256_lzcnt_si256\n");
    for (i = 0; i < 257; i++){
        printf("x=");
        x = bit_mask_avx2_lsb(i);
        print_avx2_hex(x);
        printf("lzcnt(x)=%i \n", mm256_lzcnt_si256(x));
    }
    printf("\n");

    x = _mm256_set_epi32(0,0,0,0, 0,15,1,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256(x));
    x = _mm256_set_epi32(0,0,0,8, 0,0,0,256);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256(x));
    x = _mm256_set_epi32(0,0x100,0,8, 0,192,0,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256(x));
    x = _mm256_set_epi32(-1,0x100,0,8, 0,0,32,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256(x));

   /* Set arrays for mm256_lzcnt_si256_v2:                          */
    alignas(32) static const uint32_t indx[16]   = {   3,   3,   3,   3,   3,   3,   3,   3,   2,   2,   2,   2,   1,   1,   0,   0};
    alignas(32) static const uint32_t lz_msk[16] = {   0,   0,   0,   0,   0,   0,   0,   0,  64,  64,  64,  64, 128, 128, 192, 192};
    alignas(32)              uint64_t tmp[4]     = {   0,   0,   0,   0};
    printf("\nmm256_lzcnt_si256_v2\n");
    for (i = 0; i < 257; i++){
        printf("x=");
        x = bit_mask_avx2_lsb(i);
        print_avx2_hex(x);
        printf("lzcnt(x)=%i \n", mm256_lzcnt_si256_v2(x, tmp, indx, lz_msk));
    }
    printf("\n");

    x = _mm256_set_epi32(0,0,0,0, 0,15,1,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256_v2(x, tmp, indx, lz_msk));
    x = _mm256_set_epi32(0,0,0,8, 0,0,0,256);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256_v2(x, tmp, indx, lz_msk));
    x = _mm256_set_epi32(0,0x100,0,8, 0,192,0,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256_v2(x, tmp, indx, lz_msk));
    x = _mm256_set_epi32(-1,0x100,0,8, 0,0,32,0);
    printf("x=");print_avx2_hex(x);printf("lzcnt(x)=%i \n", mm256_lzcnt_si256_v2(x, tmp, indx, lz_msk));

    return 0;
}

输出表明代码正确:

$ ./a.out
mm256_lzcnt_si256
x=0000000000000000 0000000000000000 0000000000000000 0000000000000000  lzcnt(x)=256 
x=0000000000000000 0000000000000000 0000000000000000 0000000000000001  lzcnt(x)=255 
...
x=0000000000000000 0000000000000000 7FFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF  lzcnt(x)=129 
x=0000000000000000 0000000000000000 FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF  lzcnt(x)=128 
x=0000000000000000 0000000000000001 FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF  lzcnt(x)=127 
...
x=7FFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF  lzcnt(x)=1 
x=FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF FFFFFFFFFFFFFFFF  lzcnt(x)=0 

x=0000000000000000 0000000000000000 000000000000000F 0000000100000000  lzcnt(x)=188 
x=0000000000000000 0000000000000008 0000000000000000 0000000000000100  lzcnt(x)=124 
x=0000000000000100 0000000000000008 00000000000000C0 0000000000000000  lzcnt(x)=55 
x=FFFFFFFF00000100 0000000000000008 0000000000000000 0000002000000000  lzcnt(x)=0 

函数mm256_lzcnt_si256_v2是同一函数的替代版本, 但现在指向查找 tables 和临时数组的指针是通过 函数调用。这导致 clean assembly code (无堆栈操作),并给出 内联后需要哪些指令的印象mm256_lzcnt_si256 在循环中。

使用 gcc 8.2 和选项 -m64 -O3 -march=skylake:

mm256_lzcnt_si256_v2:
        vpxor   xmm1, xmm1, xmm1
        vmovdqu YMMWORD PTR [rdi], ymm0
        vpcmpeqq        ymm0, ymm0, ymm1
        vmovmskpd       ecx, ymm0
        mov     eax, DWORD PTR [rsi+rcx*4]
        lzcnt   rax, QWORD PTR [rdi+rax*8]
        add     eax, DWORD PTR [rdx+rcx*4]
        vzeroupper
        ret

在循环上下文中,通过内联,vpxor 可能会被提升到循环之外。

我有一个版本不是 "elegant",但这里更快(Apple LLVM 版本 9.0.0 (clang-900.0.39.2)):

#define NOT_ZERO(x) (!!(x))

#ifdef UNIFORM_DISTRIBUTION
#define LIKELY(x)           __builtin_expect(NOT_ZERO(x), 1)
#define UNLIKELY(x)         __builtin_expect(NOT_ZERO(x), 0)
#else
#define LIKELY(x)           (x)
#define UNLIKELY(x)         (x)
#endif


inline unsigned int clz_u128(uint64_t a, uint64_t b, int not_a, int not_b) {
    if(UNLIKELY(not_a)) {
        if(UNLIKELY(not_b)) {
            return 128;
        } else {
            return (__builtin_clzll(b)) + 64;
        }
    } else {
        return (__builtin_clzll(a));
    }
}

unsigned int clz_u256(__m256i packed) {
    const uint64_t a_0 = (uint64_t)_mm256_extract_epi64(packed, 0);
    const uint64_t a_1 = (uint64_t)_mm256_extract_epi64(packed, 1);
    const uint64_t b_0 = (uint64_t)_mm256_extract_epi64(packed, 2);
    const uint64_t b_1 = (uint64_t)_mm256_extract_epi64(packed, 3);

    const int not_a_0 = !a_0;
    const int not_a_1 = !a_1;

    if(UNLIKELY(not_a_0 & not_a_1)) {
        return clz_u128(b_0, b_1, !b_0, !b_1) + 128;
    } else {
        return clz_u128(a_0, a_1, not_a_0, not_a_1);
    }
}

它将一个更大的问题拆分成更小的问题,并利用这样一个事实,即如果向量分布均匀,则高位 non-zero 比低位更有可能。

如果需要均匀分布以获得额外性能,只需添加 #define UNIFORM_DISTRIBUTION

由于您还要求更优雅(即更简单)的方法来执行此操作:在我的计算机上,您的代码运行速度与下面的代码一样快。在这两种情况下,计算 1000 万个 256 位字的结果都需要 45 毫秒。

由于我用(四个)随机生成的均匀分布的 64 位整数(而不是均匀分布的 256 位整数)填充 AVX 寄存器,数组的迭代顺序对我的基准测试结果没有影响。此外,尽管这几乎不用说,但编译器足够聪明,可以展开循环。

uint32_t countLeadZeros(__m256i const& reg)
{
  alignas(32) uint64_t v[4];
  _mm256_store_si256((__m256i*)&v[0], reg);

  for (int i = 3; i >= 0; --i)
    if (v[i]) return _lzcnt_u64(v[i]) + (3 - i)*64;

  return 256;
}

编辑:从我的回答下方的讨论和我的编辑历史中可以看出,我最初采用了类似于@PeterCorbes 的方法().一旦我开始做基准测试,我就改变了我的方法,因为我完全忽略了一个事实,即几乎我所有的输入都有位于 AVX 字的前 64 位中的最高有效位。

在意识到自己犯的错误后,我决定尝试更正确地进行基准测试。下面我将展示两个结果。我搜索了 post 的编辑历史记录,并从那里 copy-pasted 我提交的函数(但后来 edited-out),然后我改变了我的方法并选择了分支版本。该功能如下所示。我比较了我的 "branched" 函数、我的 "branchless" 函数和@PeterCorbes 独立开发的无分支函数的性能。

int countLeadZeros(__m256i const& reg){

  __m256i zero = _mm256_setzero_si256();
  __m256i cmp = _mm256_cmpeq_epi64(reg, zero);

  int mask = _mm256_movemask_epi8(cmp);

  if (mask == 0xffffffff) return 256;

  int first_nonzero_idx = 3 - (_lzcnt_u32(~mask) >> 3);

  alignas(32) uint64_t stored[4]; // edit: added alignas(32)
  _mm256_store_si256((__m256i*)stored, reg);

  int lead_zero_count = _lzcnt_u64(stored[first_nonzero_idx]);

  return (3 - first_nonzero_idx) * 64 + lead_zero_count;
}

基准数1

为了简短起见,我将以伪代码形式呈现测试代码。我实际上使用了随机数生成器的 AVX 实现,它可以非常快速地生成随机数。首先,让我们对使分支预测非常困难的输入进行测试:

tick()
for(int i = 0; i < N; ++i)
{
   // "xoroshiro128+"-based random generator was actually used
   __m256i in = _mm256_set_epi64x(rand()%2, rand()%2, rand()%2, rand()%2);

   res = countLeadZeros(in);  
}
tock();

对于 1000 万次重复,我的 post 顶部的函数需要 200 毫秒。我最初开发的实现只需要 65 毫秒就可以完成同样的工作。但是@PeterCorbes 提供的功能只消耗了 60 毫秒。

基准数2

下面我们来测试一下我原来用的。同样,伪代码:

tick()
for(int i = 0; i < N; ++i)
{
   // "rand()" represents random 64-bit int; xoroshiro128+ waw actually used here
   __m256i in = _mm256_set_epi64x(rand(), rand(), rand(), rand());

   res = countLeadZeros(in);  
}
tock();

在这种情况下,有分支的版本更快;计算 1000 万个结果需要 45 毫秒。 @PeterCorbes 的功能需要 50 毫秒才能完成,而我的 "branchless" 实现需要 55 毫秒才能完成相同的工作。

我认为我不敢由此得出任何普遍的结论。在我看来,无分支方法更好,因为它提供了更稳定的计算时间,但是您是否需要这种稳定性可能取决于用例。

编辑:随机生成器

这是对@PeterCorbes 评论的扩展回复。如上所述,基准测试代码只是伪代码。如果有人对我实际如何生成数字感兴趣,这里有一个简短的描述。

我使用了 xoroshiro128+ 算法,该算法已发布到 public 域并且可用 at this website。用 AVX 指令重写算法非常简单,可以并行生成四个数字。我写了一个 class 接受 so-called 初始种子(128 位)作为参数。 我通过首先复制初始种子四次来获得四个并行生成器中每一个的种子(状态);之后我在 i-th 并行生成器 i-times 上使用跳转指令;我 = {0, 1, 2, 3}。每次跳跃都会将内部状态向前推进 J=2^64 步。这意味着我可以生成 4*J 个数字(对于所有日常用途来说已经足够了),在任何并行生成器开始重复当前会话中任何其他生成器已经生成的数字序列之前一次生成四个。我用 _mm256_srli_epi64 指令控制生成数字的范围;我第一次测试使用 shift 63,第二次测试没有使用 shift。

如果您的 输入 值是均匀分布的,几乎所有时间最高设置位都在向量的前 64 位(2^64 中的 1)。这种情况下的分支将预测得很好。 .

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


Wim 的 lzcnt 在比较位图上找到正确元素的想法是一个很好的方法。

但是,runtime-variable 使用 store/reload 对向量进行索引可能比随机播放 更好。 Store-forwarding 延迟很低(在 Skylake 上可能是 5 到 7 个周期),并且该延迟与索引生成并行(compare / movemask / lzcnt)。 movd/vpermd/movd lane-crossing 洗牌策略在索引已知后需要 5 个周期,才能将正确的元素放入整数寄存器。 (参见 http://agner.org/optimize/

我认为这个版本在 Haswell/Skylake(和 Ryzen)上的延迟应该更好,吞吐量也更好。 (vpermd 在 Ryzen 上很慢,所以它应该很好)负载的地址计算应该有与 store-forwarding 相似的延迟,所以它是一个 toss-up 实际上是关键路径。

将堆栈按 32 对齐以避免 cache-line 在 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:

为了找到最高的 non-zero 元素(31 - lzcnt(~movemask)),我们使用 bsr 直接获取位(以及字节)索引,并取减去关键路径。只要我们将掩码分支为零,这就是安全的。 (无分支版本需要初始化寄存器以避免 out-of-bounds 索引)。

在 AMD CPU 上,bsrlzcnt 慢得多。在 Intel CPU 上,它们具有相同的性能,除了 output-dependency details.

的微小变化

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

如果 输入 为零,

bsr 设置 ZF,而不是像大多数指令那样基于输出。 (这和输出依赖性可能是它在 AMD 上运行缓慢的原因。)BSR 标志上的分支并不比 xor eax,-1 设置的 ZF 上的分支特别好以反转掩码,这正是 gcc 所做的。无论如何,Intel document a _BitScanReverse(&idx, mask) intrinsic return 是 bool,但 gcc 不支持它(甚至 x86intrin.h 也不支持)。 GNU C 内置函数没有 return 布尔值让你使用标志结果,但如果你检查输入 C 变量是否为 non-zero.


使用双字 (uint32_t) 数组和 vmovmskps 将使第二个 lzcnt 使用内存源操作数而不需要 movzx 到 zero-extend 一个字节。但是 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 与英特尔 CPU 上的 reg + reg*scale 相比有额外的延迟。参见 Agner Fog's instruction tables)。

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


Intel's optimization manual 说 (Table 2-24) 即使 Sandybridge 也可以从 256 位存储转发到 single-byte 加载没有问题,所以我认为它在 AVX2 CPU 上很好,与转发到 32 位负载相同,即存储的 4-byte-aligned 块。