C语言 计算__m256i字中的前导零

vcudknz3  于 11个月前  发布在  其他
关注(0)|答案(4)|浏览(125)

我正在修改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位寄存器中找出确切的非零字是相当笨拙的。
有没有人知道有没有更好(或更快)的方法来做到这一点?
只是作为一个额外的信息:我实际上想计算由逻辑AND创建的任意长向量的第一个设置位的索引,我正在比较标准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)

zy1mlcev

zy1mlcev1#

如果你的 input 值是均匀分布的,几乎所有的时候最高的设置位都将在向量的前64位(1/2^64)。在这种情况下的分支将预测得很好。@Nejc的答案对这种情况很好。
但是,许多问题,其中lzcnt是解决方案的一部分,具有更均匀的分布 * 输出 *,其中最高设置位通常不是最高64位。
Wim在比较位图上使用lzcnt来找到正确的元素的想法是一种非常好的方法。
但是,使用store/reload对向量进行运行时变量索引可能比shuffle更好。(在Skylake上可能是5到7个周期),其中一些延迟与索引生成并行(compare / movemask / lzcnt). movd/vpermd/movd跨通道 Shuffle 策略在Skylake上需要7个周期才能将正确的元素放入整数寄存器,L1 d缓存命中的加载使用延迟从地址到数据只有5个周期,当数据来自存储缓冲区时可能类似。(4c only in pointer-chasing scenarios)。参见http://agner.org/optimize/,特别是https://uops.info/,它具有movd的延迟测量。如果您查看details,其中一个测量是往返于xmm 0的movd链。在Haswell上,每次往返2个周期,因此总共5个周期,中间有一个vpermd。在Skylake上,一个movd往返有4个周期延迟,大概是3个周期,1个周期,所以用vpermd循环7次。

我认为这个版本在Haswell/Skylake上的延迟应该更好或相等,吞吐量也更好。在Ryzen上要好得多。(vpermd比Zen 4之前的Zen上的英特尔慢。而Zen 4上的movd往返大约是7个周期,其中一个movd指令是2个uop。在Zen 2和3上使用单uop movd进行6周期往返。)

lzcnt结果上需要一些标量数学来获得加载索引,这会消耗一些延迟和吞吐量优势,这取决于编译器的智能程度。vmovmskps结果上的lzcnt直接用作vpermd的shuffle索引。
堆栈按32对齐可以避免在32字节的存储区上进行缓存行拆分,但这需要额外的指令。因此,如果它可以内联到一个多次使用它的函数中,或者已经需要对其他一些__m256i进行如此多的对齐,那么这是最好的。GCC将对齐堆栈,无论您是否要求它,但是MSVC和Clang不会。我认为,即使存储本身在大多数现代CPU上是一个缓存行分割,从相对于存储对齐的dword进行存储转发也可以工作。

static inline
int lzcnt_si256(__m256i vec)
{
    // or just lzcnt_si256_memsrc(&vec) optimizes away store/reload
    __m256i   vnonzero = _mm256_cmpeq_epi32(vec, _mm256_setzero_si256());
    uint32_t  nzmask = ~ _mm256_movemask_epi8(vnonzero);   //  1 for bytes that are part of non-zero dwords
    // nzmask |= 0xf;  // branchless clamp to last elem
    if (nzmask == 0)     // all 32 bits came from vpmovmskb, so NOT didn't introduce any constant 1s
        return 256;      // don't access outside the array
    alignas(32) uint32_t elems[8];
    _mm256_storeu_si256((__m256i*)elems, vec);

    unsigned  lzbytes = _lzcnt_u32(nzmask);   // bytes above the dword containing the nonzero bit.
    unsigned char *end_elem = 28 + (unsigned char*)elems;
    uint32_t *nz_elem = (uint32_t*)(end_elem - lzbytes);  // idx = 31-(lzcnt+3) = 28-lzcnt
    return    8*lzbytes + _lzcnt_u32(*nz_elem);  // this is an aligned load, memcpy not needed
}

// unaligned and strict-aliasing safe
static inline
int lzcnt_si256_memsrc(const void *m256)
{
    __m256i  vec = _mm256_loadu_si256((const __m256i*)m256);
    __m256i  vnonzero = _mm256_cmpeq_epi32(vec, _mm256_setzero_si256());
    uint32_t nzmask = ~ _mm256_movemask_epi8(vnonzero);
    // nzmask |= 0xf;  // branchless clamp to last elem
    if (nzmask == 0)    // all 32 bits came from vpmovmskb, so NOT didn't introduce any constant 1s
        return 256;     // don't access outside the array

    unsigned  lzbytes = _lzcnt_u32(nzmask);
    unsigned char *end_elem = 28 + (unsigned char*)m256;  // can be done as part of the addressing mode after sub
    uint32_t *nz_elem = (uint32_t*)(end_elem - lzbytes);  // idx = MSB_idx-3 = 31-(lzcnt+3) = 28-lzcnt
    uint32_t nz_dword;
    memcpy(&nz_dword, nz_elem, sizeof(nz_dword));  // For strict-aliasing safety,  and/or if m256 wasn't aligned by 4.  __attribute__((aligned(1),may_alias)) on the pointer would work in GNU C
    return    8*lzbytes + _lzcnt_u32(nz_dword);
}

字符串
GCC和clang不会优化storeu的副本,如果vector刚刚加载,所以我做了一个单独的版本,不幸的是。(如果你的数据不是4字节对齐的,并且你在dword加载和vector加载时都得到了缓存行分割,请考虑使用movzx加载的字节版本。)
在Godbolt和GCC 13 -O3 -march=x86-64-v3上,我们得到这样的asm,将ymm0计数到esi中(内联到main的循环中)。

# GCC13.2 -O3 -march=x86-64-v3     (or -march=haswell)
# lzcnt_si256_memsrc inside a loop in main, after printf
   vmovdqa ymm0, YMMWORD PTR [rsp]
   vpxor   xmm1, xmm1, xmm1
   mov     esi, 256
   vpcmpeqd        ymm0, ymm0, ymm1  # could have used a memory source operand
   vpmovmskb       eax, ymm0
   xor     eax, -1                   # ~mask and set FLAGS, unlike the NOT instruction
   je      .L33
   xor     ecx, ecx            # dep-breaking for lzcnt
   mov     edx, 28
   lzcnt   ecx, eax
   sub     rdx, rcx
   lzcnt   edx, DWORD PTR [rsp+32+rdx]  # Un-laminates into 2 uops with an indexed addr mode.  mov rdx, rsp ; sub would have allowed [rdx+32+28] here
   lea     esi, [rdx+rcx*8]          # lzbytes*8 + lzcnt(nz chunk)
.L33:   # jump target for mask==0


Clang喜欢使用vptest ymm0, ymm0/jne来处理早期输出,这比等待测试movemask结果要花费更多的uop。也许[[unlikely]]注解或不可移植的等价物会有所帮助。
GCC的非内联版本在某些方面更好(sub针对指针arg,+28作为lzcnt edx, [rdi+28]的addr模式的一部分,lzcnt edx, [rdi+28]可以stay micro-fused on Intel,因为它只使用一个reg。)但GCC浪费了一个mov reg副本和两个深度中断异或归零指令,即使两个lzcnt指令都可以覆盖它们的输入(或者一个reg,它保存一个指向mem-src版本的指针)。有时候,我们可以重新排列C语言,但这取决于它所内联的代码。

**掩码上的bsr而不是31 - lzcnt**可以减少英特尔上的关键路径延迟:没有nag或NEG,只是添加了一些东西作为标量加载的寻址模式的一部分。GCC 8和更早的版本将为31-__builtin_clz()发出它,但当前的GCC仅使用31-lzcnt31^lzcnt,甚至与-march=haswell一起使用,两者具有相同的性能特性(包括输出依赖性)。

如果你是专门为英特尔优化的,BSR可能仍然是一个好主意。但是对于便携式软件,BSR在AMD上比LZCNT慢得多,LZCNT在除了x86-64 macOS之外的任何地方都是相关的。但是希望MSVC之外的编译器能发出它。

#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
static inline 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.

 // 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);  // GCC 9 and later use lzcnt even with -march=haswell implying tune for Intel, leading to worse code.  Same for clang
 #endif
}


In C++20还有31 - std::countr_zero(mask)-所有AVX 2的CPU都有BMI 1/BMI 2,所以编译器可以使用lzcnt。(在没有BMI 1的CPU上,countr_zero会稍微慢一点,因为它保证了掩码=0时的32,这与bsr指令或内部指令不同,或者GNU __builtin_clz。所以它会在输入为零时进行分支或CMOV。)

使用字节元素的BSR版本

一般情况下不推荐,但我不想放弃它。它在MSVC上使用BSR编译,但在GCC和Clang上更差。将lzcnt#if更改为1,用于未对齐数据的字节源版本,以避免标量加载(但不是vec)上的缓存行和页面分裂。

int lzcnt_si256_byte_memsrc(const void *m256)
{
    __m256i  vec = _mm256_loadu_si256((const __m256i*)m256);
    __m256i   nonzero_elem = _mm256_cmpeq_epi8(vec, _mm256_setzero_si256());
    uint32_t  mask = ~_mm256_movemask_epi8(nonzero_elem);   //  1 where there are non-zero bytes

    // alternative: mask |= 1 to lzcnt the low byte in that case, giving 32-24 = 8
    if (mask == 0)
        return 256;
    unsigned char *elems = (unsigned char*)m256;

#if 0  // GCC13 compiles both the same, clang's __builtin_clz version is worse (xor twice)
// Unless you're specifically tuning for Intel, use this version
    uint8_t *end_elem = elems + 31;
    unsigned   lz_msk   = _lzcnt_u32(mask);
    size_t   idx_from_end = -(size_t)lz_msk;    // idx = 31 - lz_msk;  // relative to start
    unsigned   highest_nonzero_byte = end_elem[idx_from_end];
#else
    unsigned   idx = bsr_nonzero(mask);   // use bsr to get the 31-x for free, because it's known to be non-zero
    unsigned   lz_msk = 31 - idx;    // off the critical path, if compilers actually use BSR
// MSVC notices that (31-idx)*8 - 24 + x  is (28-idx)*8 + x,  allowing a 2-component LEA
// GCC and clang miss that, in older versions that actually use BSR
    unsigned   highest_nonzero_byte = elems[idx];
#endif
    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.
}


在AMD CPU上,bsrlzcnt慢得多。在Intel CPU上,它们的性能相同,除了output-dependency details的微小变化。(lzcnt在Skylake之前有一个false依赖关系。bsr对所有CPU的输出都有一个 true 依赖关系。)
输入=0的bsr使目标寄存器保持不变,但intrinsic并没有提供一种方法来利用这一点免费获得类似于CMOV的行为。(Intel只将其记录为未定义的输出,但AMD将Intel / AMD CPU的实际行为记录为在目标寄存器中产生旧值)。
bsrinput 为零时设置ZF,而不是像大多数指令那样基于输出。(这和输出依赖性可能是它在AMD上慢的原因。)在BSR标志上分支并不比在ZF上分支特别好,如xor eax,-1设置的那样,以反转掩码,这就是gcc所做的。无论如何,Intel确实记录了一个返回bool_BitScanReverse(&idx, mask)内部函数,但gcc不支持它(即使使用x86intrin.h也不行)。GNU C内置程序不返回布尔值来让您使用标志结果,并且通常不使用FLAGS结果,即使在检查输入C变量是否为非零时也是如此。
Wim的版本需要lz_msk-24,因为使用8位掩码时,高24位始终为0,但32位掩码填充32位reg。
这个版本有8位元素和32位掩码,正好相反:我们需要lzcnt选择的字节,* 不 * 包括寄存器中的24个前导零位。所以我们的-24移动到不同的位置,而不是索引数组的关键路径的一部分。
GCC选择将其作为单个3组件莱亚的一部分(reg + reg*scale - const),这对于吞吐量来说是很好的,但是把它放在了最后一个lzcnt之后的关键路径上。(3组件莱亚与Ice Lake之前的Intel上的reg + reg*scale相比具有额外的延迟。有关Ice Lake和桤木Lake上不同莱亚地址模式的测试,请参见Agner Fog's instruction tableshttps://uops.info/;任何缩放的索引现在都是慢LEA,但在桤木和更多端口之前仍有1c延迟。)
带有一些-march选项的Clang会造成混乱,通过使用更多的指令而不是复杂的莱亚,它有利于延迟而不是吞吐量。
乘以8可以作为lea的一部分来完成,但是乘以32需要移位(或者折叠成两个单独的LEA)。
Intel's optimization manual表示(表2-24)即使Sandybridge也可以从256位存储转发到单字节加载而没有问题,所以我认为这在AVX 2 CPU上很好,就像转发到32位加载存储的4字节对齐块一样。

AVX-512版本,可能不会更快,只是一个实验

Store/reload可能仍然是一个更好的策略,特别是对于已经在内存中的数据,以避免store部分。我们可以使低元素特殊,总是真或总是假,用vptestnmb反对除低dword之外的全1向量,免费获得nz |= 0xf
AVX-512有vplzcntq来执行每个元素的64位位扫描。在存储之前执行该操作可以通过将lzcnt与元素索引计算重叠来缩短关键路径延迟。但是它需要一个向量ALU而不是Intel CPU上的port-1 ALU,后者不能运行512位向量uop(或者当有任何512位uop正在运行时的任何向量uop)。
有趣的方法是使用掩码从四个元素中选择一个来进行 Shuffle 和混合,我希望我可以使用vpcompressq来左打包,但这会得到最低而不是最高。水平和类型的 Shuffle +混合模式只有在高元素的掩码位为零时才能抓住较低的元素。 Shuffle 本身可以使用合并掩码,包括在同一个reg中进行单微操作缩减步骤。但是准备掩码的掩码指令并不快。

// See Godbolt for the first 2 attempts at this

// this isn't great either
int lzcnt_si256_avx512_v2(__m256i vec)
{
    __mmask8 nz = _mm256_test_epi64_mask(vec, vec); // 1 where input is non-zero, useful for merging higher into lower elements
    __m256i vclz = _mm256_lzcnt_epi64(vec);
    vclz = _mm256_add_epi64(vclz, _mm256_set_epi64x(0, 64, 128, 192));  // add the number of bits to the left of each element
    int nz_shr = nz>>1;
    vclz = _mm256_mask_unpackhi_epi64(vclz, nz_shr, vclz, vclz); //  even elements = higher elem if it was non-zero (nz_mask=1 merge masking), else keep the lower element
    nz_shr |= nz;  // kand: 1 cycle latency on port 0
    nz_shr >>= 2;  // kshiftr: 4 cycle latency on port 5 (Intel).  Zen 4 is 1c latency on 2 ports
    __m128i low = _mm256_mask_extracti32x4_epi32(_mm256_castsi256_si128(vclz), nz_shr, vclz, 1);
    return _mm_cvtsi128_si32(low);
}

基于非零掩码通过vpcompressd将零计数尾随到左包

// Just for fun, use vpcompressd to grab the first non-zero element
// Worth comparing for throughput and latency on Intel and Zen 4, vs. store/reload
static inline
int tzcnt_si256_avx512(__m256i vec)
{
    __mmask8 mask = _mm256_test_epi32_mask(vec, vec);
    __m256i pack_nz = _mm256_maskz_compress_epi32(mask, vec);
    uint32_t first_nz = _mm256_extract_epi32(pack_nz, 0);  // optimizes to vmovd
    uint32_t mask32 = mask | (-1U<<7); // produce 256 = 7*32 + 32 for an all-zero vector where tzcnt(first_nz) = 32.
      // -1U<<7 fits in signed imm8.  With 1<<7, compilers were going nuts trying to optimize
      // alternative: if (mask==0) return 256;
    return _tzcnt_u32(first_nz) + 32*_tzcnt_u32(mask32);
}


在Godbolt上,我包含了一个使用VBMI 2(Ice Lake)vpcompressb的版本,因此最终的x + 32*y可以是*8,允许莱亚。但这使得处理全零掩码变得更加困难;请参阅Godbolt上的代码注解。
另一个tzcnt策略可能涉及vplzcntq( v & -v ),并从set_epi64x(192+63, 128+63, 64+63, 63)的向量中减去它,得到tzcnt=63-lzcnt(blsi(vec))。然后选择正确的非零元素?这是更多的向量操作,但是与压缩 Shuffle 并行运行lzcnt dep链。vpcompress*是2个uops,第一个只需要掩码作为输入,而不是正在被shuffle的vector。(假设它将掩码处理为vperm* 的shuffle-control)。这将优化延迟,但不是吞吐量。如果使用512位vector,更多的vector uops甚至更不利。

处理全零输入的情况可能需要一个带有compress策略的分支,除非我们想将合并掩码的sub转换为set1(256)的向量。63-lzcnt只适用于[0..63]中的lzcnt;它需要-1而不是64来获得我们想要的输入=0的结果。这不会严重损害并行级:vplzcntq仍然可以与compare-into-mask并行运行,而vpsubq必须等待两者都准备好。这也与compress从compare阅读mask分开。

j0pj023g

j0pj023g2#

(更新:自2019-01-31以来的新答案)
三个备选方案是:

  • Peter Cordes' excellent answer.快速。这个解决方案不是无分支的,这应该不是问题,除非输入经常是零,并且出现不规则的模式。
  • 我之前的答案在这个答案的edit history中,比彼得·科德斯的答案效率低,但没有分支。
  • 这个答案。如果来自2个小查找表的数据在L1缓存中,则速度非常快。L1缓存占用空间为128字节。无分支。当不经常调用时,可能会出现缓存未命中。

在这个答案中,将输入epi64向量与零进行比较,从而产生掩码。通过索引i_mask,从两个查找表读取两个值:1.第一个非零64位元素的索引,以及2.前一个非零元素的非零数。(从左到右)零元素。最后,计算第一个非零64位元素的_lzcnt_u64并将其添加到查找表值中。函数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是同一函数的替代版本,但现在指向查找表和临时数组的指针随函数调用一起传递。这导致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可能会被提升到循环之外。

i2byvkas

i2byvkas3#

既然你也要求更优雅(即更简单)的方法来做到这一点:在我的计算机上,你的代码运行得和下面的代码一样快。在两种情况下,计算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(but he provided a better optimized solution)的方法。一旦我开始做基准测试,我就改变了我的方法,因为我完全忽略了这样一个事实,即实际上我所有的输入都具有位于AVX字的前64位内的最高有效位。

在我意识到我犯的错误之后,我决定尝试更正确地做基准测试。我将在下面展示两个结果。我搜索了我的帖子的编辑历史,并从那里复制粘贴了我提交的函数。(但后来编辑出来)之前,我改变了我的方法,去分支版本。该函数如下所示。我比较了我的“分支”函数的性能,我的“无分支”函数和由@PeterCorbes.His version is superior to mine in terms of performance - see his excellently written post that contains lots of useful details独立开发的无分支函数。

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万次重复,从我的帖子顶部开始的函数需要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毫秒才能完成,而我的“无分支”实现需要55毫秒才能完成同样的工作。
我不认为我敢从中得出任何一般性的结论。在我看来,无分支方法更好,因为它提供了更稳定的计算时间,但你是否需要这种稳定性可能取决于用例。

编辑:随机生成器

这是对@PeterCorbes评论的扩展回复。正如我上面所说的,基准测试代码只是伪代码。如果有人感兴趣,我实际上是如何生成数字的,这里有一个快速描述。
我使用xoroshiro 128+算法,该算法已发布到公共领域,可用于at this website。使用AVX指令重写算法非常简单,以便并行生成四个数字。我编写了一个类,接受所谓的初始种子(128位)作为参数。我通过首先复制初始种子四次来获得四个并行生成器中每一个的种子(状态);之后,我在第i个并行生成器上使用跳转指令i次; i = {0,1,2,3}。每一次跳跃都会使内部状态J=2^64步向前推进。这意味着我可以生成4*J个数字(对于所有日常用途来说,这已经足够了),在任何并行生成器开始重复已经由当前会话中的任何其他生成器产生的数字序列之前,一次四个。我用_mm256_srli_epi64指令控制产生的数字的范围;第一次测试我使用移位63,第二次测试没有移位。

syqv5f0l

syqv5f0l4#

我有一个版本,这是不是真的“优雅”,但更快在这里(苹果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);
    }
}

字符串
它将一个更大的问题分解为更小的问题,并利用了这样一个事实:如果向量分布是均匀的,那么高位比低位更有可能是非零的。
如果期望均匀分布以获得额外的性能,则只需添加#define UNIFORM_DISTRIBUTION

相关问题