popcnt也能向量化?

2024-07-30 15:22:43 浏览数 (2)

popcnt指令本身也支持sse向量化了,但如果序列非常大 popcnt只能处理8B,怎么办

直观的方法就是把序列按照8B拆开,分段popcnt,或者,向量化?大块?

回忆一下popcnt实现

代码语言:javascript复制
int count ( uint64_t x) {
    int v = 0;
    while (x != 0) {
        x &= x - 1;
        v   ;
    }
    return v;
}

当然编译器会优化成popcnt

考虑avx512,这个足够大了吧

chatgpt很快就能帮咱们实现一个

代码语言:javascript复制
#include <immintrin.h>

// Function to perform popcnt using AVX-512 for 64-bit integers
void avx512_popcnt_epi64(const uint64_t *input, uint64_t *output, size_t size) {
    for (size_t i = 0; i < size; i  = 8) { // Process 8 x 64-bit integers at a time
        __m512i data = _mm512_loadu_si512(&input[i]); // Load 512 bits (8 x 64-bit integers)
        __m512i result = _mm512_popcnt_epi64(data);   // Perform popcnt on each 64-bit integer
        _mm512_storeu_si512(&output[i], result);      // Store the results
    }
}

但很多硬件是不支持avx512的(比如arm), 怎么办?模拟,只需要avx2就行

但数字大于512呢,怎么拆分呢?

Harley-Seal算法 和 Faster Population Counts Using AVX2 Instructions[1]

如果没有avx512也可以avx2的话类似_mm256_shuffle_epi8也可以利用上

借助 PSHUFB可以多组popcnt 甚至可以自己主动划分组搞流水线 这里引入Harley-Seal算法

核心思想就是 Carry-Save Adder(CSA):

给定三个数 a b c 那他们和可以分成两部分

代码语言:javascript复制
void CSA (uint64_t * h , uint64_t * l, uint64_t a , uint64_t b , uint64_t c) {
    uint64_t u = a ˆ b;
    *h = (a & b) | (u & c);
    *l = u ˆ c;
}

具体实现类似这样

代码语言:javascript复制
uint64_t harley_seal ( uint64_t * d , size_t size ) {
    uint64_t total = 0, ones = 0, twos = 0,
    fours = 0, eights = 0, sixteens = 0;
    uint64_t twosA , twosB , foursA , foursB , eightsA , eightsB ;
    for ( size_t i = 0; i < size - size % 16; i  = 16) {
        CSA (& twosA , & ones , ones , d[i  0] , d[i  1]) ;
        CSA (& twosB , & ones , ones , d[i  2] , d[i  3]) ;
        CSA (& foursA , & twos , twos , twosA , twosB );
        CSA (& twosA , & ones , ones , d[i  4] , d[i  5]) ;
        CSA (& twosB , & ones , ones , d[i  6] , d[i  7]) ;
        CSA (& foursB , & twos , twos , twosA , twosB );
        CSA (& eightsA , & fours , fours , foursA , foursB );
        CSA (& twosA , & ones , ones , d[i  8] , d[i  9]) ;
        CSA (& twosB , & ones , ones , d[i  10] , d[i  11]) ;
        CSA (& foursA , & twos , twos , twosA , twosB );
        CSA (& twosA , & ones , ones , d[i  12] , d[i  13]) ;
        CSA (& twosB , & ones , ones , d[i  14] , d[i  15]) ;
        CSA (& foursB , & twos , twos , twosA , twosB );
        CSA (& eightsB , & fours , fours , foursA , foursB );
        CSA (& sixteens , & eights , eights , eightsA , eightsB );
        total  = count ( sixteens );
    }
    total = 16 * total   8 * count ( eights )   4 * count ( fours )   2 * count ( twos )   count ( ones );
    for ( size_t i = size - size % 16 ; i < size ; i   )
    total  = count (d[i ]) ;
    return total ;
}

这个算法可以三个数驱动,那自然可以pipeline话 另外这个CSA也可以用avx2或者avx512重写

比如 avx2

代码语言:javascript复制
#include <immintrin.h>
#include <stdint.h>
#include <stddef.h>

// Carry-Save Adder (CSA) implementation using AVX2
void CSA(__m256i* h, __m256i* l, __m256i a, __m256i b, __m256i c) {
    __m256i u = _mm256_xor_si256(a, b);
    *h = _mm256_or_si256(_mm256_and_si256(a, b), _mm256_and_si256(u, c));
    *l = _mm256_xor_si256(u, c);
}

__m256i count ( __m256i v) {
    __m256i lookup =
        _mm256_setr_epi8 (0 , 1, 1, 2, 1, 2, 2, 3, 1, 2,
        2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3,
        1, 2, 2, 3, 2, 3, 3, 4) ;
    __m256i low_mask = _mm256_set1_epi8 (0 x0f );
    __m256i lo = = _mm256_and_si256 (v , low_mask );
    __m256i hi = _mm256_and_si256 ( _mm256_srli_epi32
    (v , 4) , low_mask );
    __m256i popcnt1 = _mm256_shuffle_epi8 ( lookup ,
    lo );
    __m256i popcnt2 = _mm256_shuffle_epi8 ( lookup ,
    hi );
    __m256i total = _mm256_add_epi8 ( popcnt1 , popcnt2
    );
    return _mm256_sad_epu8 ( total ,
        _mm256_setzero_si256 () );
}

uint64_t avx_hs ( __m256i * d , uint64_t size ) {
    __m256i total = _mm256_setzero_si256 () ;
    __m256i ones = _mm256_setzero_si256 () ;
    __m256i twos = _mm256_setzero_si256 () ;
    __m256i fours = _mm256_setzero_si256 () ;
    __m256i eights = _mm256_setzero_si256 () ;
    __m256i sixteens = _mm256_setzero_si256 () ;
    __m256i twosA , twosB , foursA , foursB ,
    eightsA , eightsB ;
    for ( uint64_t i = 0; i < size ; i  = 16) {
        CSA (& twosA , & ones , ones , d[i], d[i  1]) ;
        CSA (& twosB , & ones , ones , d[i  2] , d[i  3]) ;
        CSA (& foursA , & twos , twos , twosA , twosB );
        CSA (& twosA , & ones , ones , d[i  4] , d[i  5]) ;
        CSA (& twosB , & ones , ones , d[i  6] , d[i  7]) ;
        CSA (& foursB ,& twos , twos , twosA , twosB );
        CSA (& eightsA ,& fours , fours , foursA , foursB );
        CSA (& twosA , & ones , ones , d[i  8] , d[i  9]) ;
        CSA (& twosB , & ones , ones , d[i  10] , d[i  11]) ;
        CSA (& foursA , & twos , twos , twosA , twosB );
        CSA (& twosA , & ones , ones , d[i  12] , d[i  13]) ;
        CSA (& twosB , & ones , ones , d[i  14] , d[i  15]) ;
        CSA (& foursB , & twos , twos , twosA , twosB );
        CSA (& eightsB , & fours , fours , foursA , foursB );
        CSA (& sixteens , & eights , eights , eightsA , eightsB );
        total = _mm256_add_epi64 ( total , count ( sixteens ));
    }
    total = _mm256_slli_epi64 ( total , 4) ;
    total = _mm256_add_epi64 ( total ,
        _mm256_slli_epi64 ( count ( eights ) , 3) );
    total = _mm256_add_epi64 ( total ,
        _mm256_slli_epi64 ( count ( fours ) , 2) );
    total = _mm256_add_epi64 ( total ,
        _mm256_slli_epi64 ( count ( twos ) , 1) );
    total = _mm256_add_epi64 ( total , count ( ones ));
    return _mm256_extract_epi64 ( total , 0)
          _mm256_extract_epi64 ( total , 1)
          _mm256_extract_epi64 ( total , 2)
          _mm256_extract_epi64 ( total , 3) ;
}

好了,你大概已经明白这个算法了,avx512版本就不贴了,可以看libpopcnt代码

性能数据

libpopcnt[2]

直接贴性能数据吧

Algorithm

32 B

64 B

128 B

256 B

512 B

1024 B

2048 B

4096 B

lookup-8

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

bit-parallel-mul

1.41

1.54

1.63

1.78

1.60

1.62

1.63

1.64

builtin-popcnt

4.75

6.36

8.58

8.55

6.72

7.60

7.88

7.94

avx2-harley-seal

1.15

1.85

3.22

4.17

8.46

10.74

12.52

13.66

avx512-harley-seal

0.35

1.49

2.54

3.83

5.63

15.12

22.18

25.60

显然 avx512-harley-seal 非常快

sse-popcnt[3]的结论差不多,就不贴数据了

算法厉害,但是用的上吗?

bitmap是经典场景了,但是bitmap不方便管理特长数据

roaringbitmap方案就是截断,优化思路和上面的论文思路相同

另外就是压缩了 bitmagic 压缩bit[4]

吸收了harley-seal的思想

引用链接

[1] Faster Population Counts Using AVX2 Instructions: https://arxiv.org/pdf/1611.07612 [2] libpopcnt: https://github.com/kimwalisch/libpopcnt [3] sse-popcnt: https://github.com/WojciechMula/sse-popcount/blob/master/results/cannonlake/cannonlake-i3-8121U-gcc-8.3.1.rst [4] bitmagic 压缩bit: https://github.com/tlk00/BitMagic

0 人点赞