我正在修改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)
型
4条答案
按热度按时间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上使用单uopmovd
进行6周期往返。)在
lzcnt
结果上需要一些标量数学来获得加载索引,这会消耗一些延迟和吞吐量优势,这取决于编译器的智能程度。vmovmskps
结果上的lzcnt
直接用作vpermd
的shuffle索引。堆栈按32对齐可以避免在32字节的存储区上进行缓存行拆分,但这需要额外的指令。因此,如果它可以内联到一个多次使用它的函数中,或者已经需要对其他一些
__m256i
进行如此多的对齐,那么这是最好的。GCC将对齐堆栈,无论您是否要求它,但是MSVC和Clang不会。我认为,即使存储本身在大多数现代CPU上是一个缓存行分割,从相对于存储对齐的dword进行存储转发也可以工作。字符串
GCC和clang不会优化
storeu
的副本,如果vector刚刚加载,所以我做了一个单独的版本,不幸的是。(如果你的数据不是4字节对齐的,并且你在dword加载和vector加载时都得到了缓存行分割,请考虑使用movzx
加载的字节版本。)在Godbolt和GCC 13
-O3 -march=x86-64-v3
上,我们得到这样的asm,将ymm0
计数到esi
中(内联到main
的循环中)。型
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-lzcnt
或31^lzcnt
,甚至与-march=haswell
一起使用,两者具有相同的性能特性(包括输出依赖性)。如果你是专门为英特尔优化的,BSR可能仍然是一个好主意。但是对于便携式软件,BSR在AMD上比LZCNT慢得多,LZCNT在除了x86-64 macOS之外的任何地方都是相关的。但是希望MSVC之外的编译器能发出它。
型
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)上的缓存行和页面分裂。型
在AMD CPU上,
bsr
比lzcnt
慢得多。在Intel CPU上,它们的性能相同,除了output-dependency details的微小变化。(lzcnt
在Skylake之前有一个false依赖关系。bsr
对所有CPU的输出都有一个 true 依赖关系。)输入=0的
bsr
使目标寄存器保持不变,但intrinsic并没有提供一种方法来利用这一点免费获得类似于CMOV的行为。(Intel只将其记录为未定义的输出,但AMD将Intel / AMD CPU的实际行为记录为在目标寄存器中产生旧值)。bsr
在 input 为零时设置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 tables和https://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中进行单微操作缩减步骤。但是准备掩码的掩码指令并不快。型
基于非零掩码通过
vpcompressd
将零计数尾随到左包型
在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分开。j0pj023g2#
(更新:自2019-01-31以来的新答案)
三个备选方案是:
在这个答案中,将输入
epi64
向量与零进行比较,从而产生掩码。通过索引i_mask
,从两个查找表读取两个值:1.第一个非零64位元素的索引,以及2.前一个非零元素的非零数。(从左到右)零元素。最后,计算第一个非零64位元素的_lzcnt_u64
并将其添加到查找表值中。函数mm256_lzcnt_si256
实现了以下方法:字符串
输出表明代码是正确的:
型
函数
mm256_lzcnt_si256_v2
是同一函数的替代版本,但现在指向查找表和临时数组的指针随函数调用一起传递。这导致clean assembly code(无堆栈操作),并给出了在循环中内联mm256_lzcnt_si256
后需要哪些指令的印象。使用gcc 8.2和选项
-m64 -O3 -march=skylake
:型
在循环上下文中,使用内联,
vpxor
可能会被提升到循环之外。i2byvkas3#
既然你也要求更优雅(即更简单)的方法来做到这一点:在我的计算机上,你的代码运行得和下面的代码一样快。在两种情况下,计算1000万个256位字的结果都花了45毫秒。
由于我用随机生成的均匀分布的64位整数(而不是均匀分布的256位整数)填充AVX寄存器,因此通过数组的迭代顺序对我的基准测试结果没有影响。此外,尽管这几乎不用说,编译器足够聪明,可以展开循环。
字符串
编辑:正如在我的答案下面的讨论和我的编辑历史中可以看到的那样,我最初采取了类似于@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独立开发的无分支函数。
型
基准数1
我将在伪代码中展示测试代码,以使其简短。我实际上使用了随机数生成器的AVX实现,它可以非常快速地生成随机数。首先,让我们对使分支预测非常困难的输入进行测试:
型
对于1000万次重复,从我的帖子顶部开始的函数需要200毫秒。我最初开发的实现只需要65毫秒来完成同样的工作。但是@PeterCorbes提供的函数只需要60毫秒。
基准数2
现在让我们转向我最初使用的测试。同样,伪代码:
型
在这种情况下,有分支的版本更快;计算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,第二次测试没有移位。syqv5f0l4#
我有一个版本,这是不是真的“优雅”,但更快在这里(苹果LLVM版本9.0.0(clang-900.0.39.2)):
字符串
它将一个更大的问题分解为更小的问题,并利用了这样一个事实:如果向量分布是均匀的,那么高位比低位更有可能是非零的。
如果期望均匀分布以获得额外的性能,则只需添加
#define UNIFORM_DISTRIBUTION
。