Implementation on AVX2

include/bits.h

@@ -5,9 +5,21 @@
 #include <immintrin.h>
 #include <numeric>
 
-/**
- * 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<uint64_t>::max());
   __m512i b = _mm512_maskz_set1_epi64((1ull << ((count >> 6) + 1)) - 1,
                                       std::numeric_limits<uint64_t>::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<uint64_t, 8> 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<const __m256i*>(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<const __m256i*>(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<const __m256i*>(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
 }

include/bitvector.h

@@ -81,10 +81,6 @@ class BitVector {
 
   std::span<const uint64_t> 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);
     }

src/unittests.cpp

@@ -283,6 +283,58 @@ TEST(BitVectorInterleavedTest, MainTest) {
   }
 }
 
+TEST(CMPLPREFLEN256, Random) {
+  std::vector<uint64_t> 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<uint64_t> 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<uint16_t> 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<uint64_t> bits = {0b10110};
 //     BitVector bv(bits, 5);