diff --git a/include/bits.h b/include/bits.h index b619bcb..2d46e5f 100644 --- a/include/bits.h +++ b/include/bits.h @@ -5,9 +5,21 @@ #include #include -/** - * TODO: compile/runtime checks fror SIMD - */ + +#if defined(__AVX512VPOPCNTDQ__) && defined(__AVX512F__) && defined(__AVX512BW__) +#define AVX512SUPP +#endif + + +#ifdef __AVX2__ +#define AVX2SUPP +#endif + + +inline uint64_t first_bits_mask(size_t num) { + return num >= 64 ? UINT64_MAX : ((1llu << num) - 1); +} + /** * @brief Number of 1 bits in positions 0 .. count - 1 @@ -30,18 +42,38 @@ * 64 bits and then reduce_add to sum the result. */ uint64_t rank_512(const uint64_t *x, uint64_t count) { +#ifdef AVX512SUPP + __m512i a = _mm512_maskz_set1_epi64((1ull << ((count >> 6))) - 1, std::numeric_limits::max()); __m512i b = _mm512_maskz_set1_epi64((1ull << ((count >> 6) + 1)) - 1, std::numeric_limits::max()); __m512i mask = _mm512_shldv_epi64(a, b, _mm512_set1_epi64(count % 64)); - __m512i res = _mm512_load_epi64(x); + __m512i res = _mm512_loadu_epi64(x); res = _mm512_and_epi64(res, mask); __m512i cnt = _mm512_popcnt_epi64(res); return _mm512_reduce_add_epi64(cnt); + +#else + + uint64_t last_uint = count >> 6; + + uint64_t pop_val = 0; + + for (int i = 0; i < last_uint; i++) { + pop_val += std::popcount(x[i]); + } + + uint64_t final = x[last_uint] & first_bits_mask(count & 63); + + pop_val += std::popcount(final); + return pop_val; + +#endif } + /** * @brief Return position of @p rank 1 bit in @p x */ @@ -50,6 +82,7 @@ uint64_t select_64(uint64_t x, uint64_t rank) { return _tzcnt_u64(a); } + /** * @brief Return position of @p rank 1 bit in @p x * @details Selecting within 64-bit word can be done @@ -68,13 +101,145 @@ uint64_t select_64(uint64_t x, uint64_t rank) { * see a proper SIMD algorithm to make it faster. */ uint64_t select_512(const uint64_t *x, uint64_t rank) { - __m512i res = _mm512_load_epi64(x); +#ifdef AVX512SUPP + + __m512i res = _mm512_loadu_epi64(x); std::array counts; - _mm512_store_epi64(counts.data(), _mm512_popcnt_epi64(res)); + _mm512_storeu_epi64(counts.data(), _mm512_popcnt_epi64(res)); size_t i = 0; while (i < 8 && counts[i] <= rank) { rank -= counts[i++]; } return i * 64 + select_64(x[i], rank); + +#else + + size_t i = 0; + int popcount = std::popcount(x[0]); + while (i < 7 && popcount <= rank) { + rank -= popcount; + popcount = std::popcount(x[++i]); + } + return i * 64 + select_64(x[i], rank); + +#endif +} + + +/** + * @brief Compare 4 64-bit numbers of @p x with @p y and + * return the length of the prefix where @p y is less then @p x + */ +uint16_t cmpl_pref_len_256(const uint64_t* x, uint64_t y) { +#ifdef AVX512SUPP + + auto y_4 = _mm256_set1_epi64x(y); + auto reg_256 = _mm256_loadu_epi64(x); + auto cmp = _mm256_cmpge_epu64_mask(reg_256, y_4); + + return _tzcnt_u16(cmp); + +#else +#ifdef AVX2SUPP + + auto y_4 = _mm256_set1_epi64x(y); + __m256i reg_256 = _mm256_loadu_si256(reinterpret_cast(x)); + + const __m256i offset = _mm256_set1_epi64x(0x8000000000000000ULL); + __m256i x_offset = _mm256_xor_si256(reg_256, offset); + __m256i y_offset = _mm256_xor_si256(y_4, offset); + auto mask = _mm256_movemask_epi8(_mm256_cmpgt_epi64(x_offset, _mm256_sub_epi64(y_offset, _mm256_set1_epi64x(1)))); + + return _tzcnt_u32(mask) >> 3; + +#else + + for (uint16_t i = 0; i < 4; ++i) + if (x[i] >= y) + return i; + return 4; + +#endif +#endif +} + + +/** + * @brief Compare 8 64-bit numbers of @p x with @p y and + * return the length of the prefix where @p y is less then @p x + */ +uint16_t cmpl_pref_len_512(const uint64_t* x, uint64_t y) { +#ifdef AVX512SUPP + + auto y_8 = _mm512_set1_epi64(y); + auto reg_512 = _mm512_loadu_epi64(x); + auto cmp = _mm512_cmpge_epu64_mask(reg_512, y_8); + + return _tzcnt_u16(cmp); + +#else +#ifdef AVX2SUPP + + uint16_t len = cmpl_pref_len_256(x, y); + + if (len < 4) + return len; + + return len + cmpl_pref_len_256(x + 4, y); + +#else + + for (uint16_t i = 0; i < 8; ++i) + if (x[i] >= y) + return i; + return 8; + +#endif +#endif +} + + +/** + * @brief Compare 32 16-bit numbers of @p x with @p y and + * return the count of numbers where @p x is less then @p y + */ +uint16_t cmpl_count_512(const uint16_t* x, uint64_t y) { +#ifdef AVX512SUPP + + auto y_32 = _mm512_set1_epi16(y); + auto reg_512 = _mm512_loadu_epi16(x); + auto cmp = _mm512_cmplt_epu16_mask(reg_512, y_32); + return std::popcount(cmp); + +#else +#ifdef AVX2SUPP + + auto y_16 = _mm256_set1_epi16(y); + __m256i reg_256 = _mm256_loadu_si256(reinterpret_cast(x)); + + const __m256i offset = _mm256_set1_epi16(0x8000); + __m256i x_offset = _mm256_xor_si256(reg_256, offset); + __m256i y_offset = _mm256_xor_si256(y_16, offset); + uint32_t mask = _mm256_movemask_epi8(_mm256_cmpgt_epi16(y_offset, x_offset)); + + uint16_t count = std::popcount(mask) >> 1; + + reg_256 = _mm256_loadu_si256(reinterpret_cast(x + 16)); + + x_offset = _mm256_xor_si256(reg_256, offset); + mask = _mm256_movemask_epi8(_mm256_cmpgt_epi16(y_offset, x_offset)); + + return count + (std::popcount(mask) >> 1); + +#else + + uint16_t cnt = 0; + for (uint16_t i = 0; i < 32; ++i) + if (x[i] < y) + cnt++; + return cnt; + +#endif +#endif } \ No newline at end of file diff --git a/include/bitvector.h b/include/bitvector.h index c6ce83c..79aa3ca 100644 --- a/include/bitvector.h +++ b/include/bitvector.h @@ -81,10 +81,6 @@ class BitVector { std::span bits; - inline uint64_t first_bits_mask(size_t num) const { - return num >= 64 ? UINT64_MAX : ((1llu << num) - 1); - } - /** * Precompute rank for fast queries */ @@ -144,21 +140,17 @@ class BitVector { uint64_t find_superblock(uint64_t rank) const { uint64_t left = select_samples[rank / kSelectSampleFrequency]; - auto rank_8 = _mm512_set1_epi64(rank); - while (left + 7 < super_block_rank.size()) { - auto reg_512 = _mm512_loadu_epi64(&super_block_rank[left]); - auto cmp = _mm512_cmpge_epu64_mask(reg_512, rank_8); - if (cmp) { - return left + _tzcnt_u16(cmp) - 1; + auto len = cmpl_pref_len_512(&super_block_rank[left], rank); + if (len < 8) { + return left + len - 1; } left += 8; } if (left + 3 < super_block_rank.size()) { - auto reg_256 = _mm256_loadu_epi64(&super_block_rank[left]); - auto cmp = _mm256_cmpge_epu64_mask(reg_256, *(__m256i *)(&rank_8)); - if (cmp) { - return left + _tzcnt_u16(cmp) - 1; + auto len = cmpl_pref_len_256(&super_block_rank[left], rank); + if (len < 4) { + return left + len - 1; } left += 4; } @@ -174,14 +166,10 @@ class BitVector { * 4 iterations. */ uint64_t find_basicblock(uint16_t local_rank, uint64_t s_block) const { - auto rank_32 = _mm512_set1_epi16(local_rank); auto pos = 0; for (size_t i = 0; i < 4; ++i) { - auto batch = - _mm512_loadu_epi16(&basic_block_rank[128 * s_block + 32 * i]); - auto cmp = _mm512_cmplt_epu16_mask(batch, rank_32); - auto count = std::popcount(cmp); + auto count = cmpl_count_512(&basic_block_rank[128 * s_block + 32 * i], local_rank); if (count < 32) { return kBlocksPerSuperBlock * s_block + pos + count - 1; } @@ -206,11 +194,8 @@ class BitVector { uint64_t pos = kBlocksPerSuperBlock * local_rank / (upper - lower); pos = pos + 16 < 32 ? 0 : (pos - 16); pos = pos > 96 ? 96 : pos; - auto rank_32 = _mm512_set1_epi16(local_rank); while (pos < 96) { - auto batch = _mm512_loadu_epi16(&basic_block_rank[128 * s_block + pos]); - auto cmp = _mm512_cmplt_epu16_mask(batch, rank_32); - auto count = std::popcount(cmp); + auto count = cmpl_count_512(&basic_block_rank[128 * s_block + pos], local_rank); if (count == 0) { return find_basicblock(local_rank, s_block); } @@ -220,9 +205,7 @@ class BitVector { pos += 32; } pos = 96; - auto batch = _mm512_loadu_epi16(&basic_block_rank[128 * s_block + pos]); - auto cmp = _mm512_cmplt_epu16_mask(batch, rank_32); - auto count = std::popcount(cmp); + auto count = cmpl_count_512(&basic_block_rank[128 * s_block + pos], local_rank); if (count == 0) { return find_basicblock(local_rank, s_block); } diff --git a/src/unittests.cpp b/src/unittests.cpp index 3dab3e9..ce423f4 100644 --- a/src/unittests.cpp +++ b/src/unittests.cpp @@ -283,6 +283,58 @@ TEST(BitVectorInterleavedTest, MainTest) { } } +TEST(CMPLPREFLEN256, Random) { + std::vector x(4); + std::mt19937_64 rng(42); + for (size_t i = 0; i < 1000; i++) { + uint64_t y = rng(); + uint16_t cnt = 0; + bool fl = 1; + for (size_t j = 0; j < 4;j++) { + x[j] = rng(); + fl &= x[j] < y; + cnt += fl; + } + if (cnt < 4) + EXPECT_EQ(cmpl_pref_len_256(x.data(), y), cnt); + else + EXPECT_GE(cmpl_pref_len_256(x.data(), y), cnt); + } +} + +TEST(CMPLPREFLEN512, Random) { + std::vector x(8); + std::mt19937_64 rng(42); + for (size_t i = 0; i < 1000; i++) { + uint64_t y = rng(); + uint16_t cnt = 0; + bool fl = 1; + for (size_t j = 0; j < 8;j++) { + x[j] = rng(); + fl &= x[j] < y; + cnt += fl; + } + if (cnt < 8) + EXPECT_EQ(cmpl_pref_len_512(x.data(), y), cnt); + else + EXPECT_GE(cmpl_pref_len_512(x.data(), y), cnt); + } +} + +TEST(CMPLCOUNT512, Random) { + std::vector x(32); + std::mt19937 rng(42); + for (size_t i = 0; i < 1000; i++) { + uint16_t y = rng(); + uint16_t cnt = 0; + for (size_t j = 0; j < 32;j++) { + x[j] = rng(); + cnt += x[j] < y; + } + EXPECT_EQ(cmpl_count_512(x.data(), y), cnt); + } +} + // TEST(BitVectorTest, SelectBasic) { // std::vector bits = {0b10110}; // BitVector bv(bits, 5);