bits.h

#pragma once
 
#include <immintrin.h>
 
#include <array>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <numeric>
 
#if defined(__AVX512VPOPCNTDQ__) && defined(__AVX512F__) && \
    defined(__AVX512BW__)
#define PIXIE_AVX512_SUPPORT
#endif
 
#ifdef __AVX2__
#define PIXIE_AVX2_SUPPORT
// Lookup table for 4-bit popcount
// This table maps each 4-bit value (0-15) to its population count
// clang-format off
static inline const __m256i lookup_popcount_4 = _mm256_setr_epi8(
    0, 1, 1, 2,  // 0000, 0001, 0010, 0011
    1, 2, 2, 3,  // 0100, 0101, 0110, 0111
    1, 2, 2, 3,  // 1000, 1001, 1010, 1011
    2, 3, 3, 4,  // 1100, 1101, 1110, 1111
    
    // Same table repeated for high 128 bits
    0, 1, 1, 2,  // 0000, 0001, 0010, 0011
    1, 2, 2, 3,  // 0100, 0101, 0110, 0111
    1, 2, 2, 3,  // 1000, 1001, 1010, 1011
    2, 3, 3, 4   // 1100, 1101, 1110, 1111
);
 
static inline const __m256i mask_first_half = _mm256_setr_epi8(
    0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF,
    0, 0, 0, 0,
    0, 0, 0, 0,
    0, 0, 0, 0,
    0, 0, 0, 0
);
 
// clang-format on
#endif
 
static inline uint64_t first_bits_mask(size_t num) {
  return num >= 64 ? UINT64_MAX : ((1llu << num) - 1);
}

static inline uint64_t rank_512(const uint64_t* x, uint64_t count) {
#ifdef PIXIE_AVX512_SUPPORT
 
  __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_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 < 512 ? count >> 6 : 8;
 
  uint64_t pop_val = 0;
 
  for (int i = 0; i < last_uint; i++) {
    pop_val += std::popcount(x[i]);
  }
 
  pop_val += count < 512
                 ? std::popcount(x[last_uint] & first_bits_mask(count & 63))
                 : 0;
  return pop_val;
 
#endif
}

static inline uint64_t select_64(uint64_t x, uint64_t rank) {
  return _tzcnt_u64(_pdep_u64(1ull << rank, x));
}

static inline uint64_t select_512(const uint64_t* x, uint64_t rank) {
#ifdef PIXIE_AVX512_SUPPORT
 
  __m512i res = _mm512_loadu_epi64(x);
  __m512i counts = _mm512_popcnt_epi64(res);
  __m512i prefix = counts;
 
  const __m512i idx_shift1 = _mm512_set_epi64(6, 5, 4, 3, 2, 1, 0, 0);
  const __m512i idx_shift2 = _mm512_set_epi64(5, 4, 3, 2, 1, 0, 0, 0);
  const __m512i idx_shift4 = _mm512_set_epi64(3, 2, 1, 0, 0, 0, 0, 0);
 
  __m512i tmp = _mm512_maskz_permutexvar_epi64(0xFE, idx_shift1, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
  tmp = _mm512_maskz_permutexvar_epi64(0xFC, idx_shift2, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
  tmp = _mm512_maskz_permutexvar_epi64(0xF0, idx_shift4, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
 
  __mmask8 mask = _mm512_cmpgt_epu64_mask(prefix, _mm512_set1_epi64(rank));
  uint32_t i = _tzcnt_u32(static_cast<uint32_t>(mask));
  uint64_t prev = 0;
  if (i != 0) {
    __m512i idx_prev = _mm512_set1_epi64(static_cast<int64_t>(i - 1));
    __m512i prev_vec = _mm512_permutexvar_epi64(idx_prev, prefix);
    prev = static_cast<uint64_t>(
        _mm_cvtsi128_si64(_mm512_castsi512_si128(prev_vec)));
  }
  return i * 64 + select_64(x[i], rank - prev);
 
#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
}

static inline uint64_t select0_512(const uint64_t* x, uint64_t rank0) {
#ifdef PIXIE_AVX512_SUPPORT
 
  __m512i res = _mm512_loadu_epi64(x);
  res = _mm512_xor_epi64(res, _mm512_set1_epi64(-1));
  __m512i counts = _mm512_popcnt_epi64(res);
  __m512i prefix = counts;
 
  const __m512i idx_shift1 = _mm512_set_epi64(6, 5, 4, 3, 2, 1, 0, 0);
  const __m512i idx_shift2 = _mm512_set_epi64(5, 4, 3, 2, 1, 0, 0, 0);
  const __m512i idx_shift4 = _mm512_set_epi64(3, 2, 1, 0, 0, 0, 0, 0);
 
  __m512i tmp = _mm512_maskz_permutexvar_epi64(0xFE, idx_shift1, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
  tmp = _mm512_maskz_permutexvar_epi64(0xFC, idx_shift2, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
  tmp = _mm512_maskz_permutexvar_epi64(0xF0, idx_shift4, prefix);
  prefix = _mm512_add_epi64(prefix, tmp);
 
  __mmask8 mask = _mm512_cmpgt_epu64_mask(prefix, _mm512_set1_epi64(rank0));
  uint32_t i = _tzcnt_u32(static_cast<uint32_t>(mask));
  uint64_t prev = 0;
  if (i != 0) {
    __m512i idx_prev = _mm512_set1_epi64(static_cast<int64_t>(i - 1));
    __m512i prev_vec = _mm512_permutexvar_epi64(idx_prev, prefix);
    prev = static_cast<uint64_t>(
        _mm_cvtsi128_si64(_mm512_castsi512_si128(prev_vec)));
  }
  return i * 64 + select_64(~x[i], rank0 - prev);
 
#else
 
  size_t i = 0;
  int popcount = std::popcount(~x[0]);
  while (i < 7 && popcount <= rank0) {
    rank0 -= popcount;
    popcount = std::popcount(~x[++i]);
  }
  return i * 64 + select_64(~x[i], rank0);
 
#endif
}

static inline uint16_t lower_bound_4x64(const uint64_t* x, uint64_t y) {
#ifdef PIXIE_AVX512_SUPPORT
 
  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 PIXIE_AVX2_SUPPORT
 
  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
}

static inline uint16_t lower_bound_delta_4x64(const uint64_t* x,
                                              uint64_t y,
                                              const uint64_t* delta_array,
                                              uint64_t delta_scalar) {
#ifdef PIXIE_AVX512_SUPPORT
 
  const __m256i dlt_256 = _mm256_loadu_epi64(delta_array);
  auto x_256 = _mm256_loadu_epi64(x);
  auto dlt_4 = _mm256_set1_epi64x(delta_scalar);
  auto y_4 = _mm256_set1_epi64x(y);
 
  auto tmp = _mm256_add_epi64(dlt_4, dlt_256);
  auto reg_256 = _mm256_sub_epi64(tmp, x_256);
  auto cmp = _mm256_cmpge_epu64_mask(reg_256, y_4);
 
  return _tzcnt_u16(cmp);
 
#else
#ifdef PIXIE_AVX2_SUPPORT
 
  const __m256i dlt_256 =
      _mm256_loadu_si256(reinterpret_cast<const __m256i*>(delta_array));
  auto x_256 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x));
  auto dlt_4 = _mm256_set1_epi64x(delta_scalar);
  auto y_4 = _mm256_set1_epi64x(y);
 
  auto tmp = _mm256_add_epi64(dlt_4, dlt_256);
  auto reg_256 = _mm256_sub_epi64(tmp, x_256);
 
  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 (delta_array[i] + delta_scalar - x[i] >= y) {
      return i;
    }
  }
  return 4;
 
#endif
#endif
}

static inline uint16_t lower_bound_8x64(const uint64_t* x, uint64_t y) {
#ifdef PIXIE_AVX512_SUPPORT
 
  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 PIXIE_AVX2_SUPPORT
 
  uint16_t len = lower_bound_4x64(x, y);
 
  if (len < 4) {
    return len;
  }
 
  return len + lower_bound_4x64(x + 4, y);
 
#else
 
  for (uint16_t i = 0; i < 8; ++i) {
    if (x[i] >= y) {
      return i;
    }
  }
  return 8;
 
#endif
#endif
}

static inline uint16_t lower_bound_delta_8x64(const uint64_t* x,
                                              uint64_t y,
                                              const uint64_t* delta_array,
                                              uint64_t delta_scalar) {
#ifdef PIXIE_AVX512_SUPPORT
 
  const __m512i dlt_512 = _mm512_loadu_epi64(delta_array);
  auto x_512 = _mm512_loadu_epi64(x);
  auto dlt_8 = _mm512_set1_epi64(delta_scalar);
  auto y_8 = _mm512_set1_epi64(y);
 
  auto tmp = _mm512_add_epi64(dlt_8, dlt_512);
  auto reg_512 = _mm512_sub_epi64(tmp, x_512);
  auto cmp = _mm512_cmpge_epu64_mask(reg_512, y_8);
 
  return _tzcnt_u16(cmp);
 
#else
#ifdef PIXIE_AVX2_SUPPORT
 
  uint16_t len = lower_bound_delta_4x64(x, y, delta_array, delta_scalar);
 
  if (len < 4) {
    return len;
  }
 
  return len + lower_bound_delta_4x64(x + 4, y, delta_array + 4, delta_scalar);
 
#else
 
  for (uint16_t i = 0; i < 8; ++i) {
    if (delta_array[i] + delta_scalar - x[i] >= y) {
      return i;
    }
  }
  return 8;
 
#endif
#endif
}

uint16_t lower_bound_32x16(const uint16_t* x, uint16_t y) {
#ifdef PIXIE_AVX512_SUPPORT
 
  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 PIXIE_AVX2_SUPPORT
 
  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
}

uint16_t lower_bound_delta_32x16(const uint16_t* x,
                                 uint16_t y,
                                 const uint16_t* delta_array,
                                 uint16_t delta_scalar) {
#ifdef PIXIE_AVX512_SUPPORT
 
  const __m512i dlt_512 = _mm512_loadu_epi64(delta_array);
  auto x_512 = _mm512_loadu_epi64(x);
  auto dlt_32 = _mm512_set1_epi16(delta_scalar);
  auto y_32 = _mm512_set1_epi16(y);
 
  auto tmp = _mm512_add_epi16(dlt_32, dlt_512);
  auto reg_512 = _mm512_sub_epi16(tmp, x_512);
  auto cmp = _mm512_cmplt_epu16_mask(reg_512, y_32);
  return std::popcount(cmp);
 
#else
#ifdef PIXIE_AVX2_SUPPORT
 
  auto dlt_256 =
      _mm256_loadu_si256(reinterpret_cast<const __m256i*>(delta_array));
  auto x_256 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x));
  auto dlt_16 = _mm256_set1_epi16(delta_scalar);
  auto y_16 = _mm256_set1_epi16(y);
 
  auto tmp = _mm256_add_epi16(dlt_16, dlt_256);
  auto reg_256 = _mm256_sub_epi16(tmp, x_256);
 
  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;
 
  dlt_256 =
      _mm256_loadu_si256(reinterpret_cast<const __m256i*>(delta_array + 16));
  x_256 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x + 16));
 
  tmp = _mm256_add_epi16(dlt_16, dlt_256);
  reg_256 = _mm256_sub_epi16(tmp, x_256);
 
  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 (delta_array[i] + delta_scalar - x[i] < y) {
      cnt++;
    }
  }
  return cnt;
 
#endif
#endif
}

void popcount_64x4(const uint8_t* x, uint8_t* result) {
#ifdef PIXIE_AVX512_SUPPORT
  __m256i data = _mm256_loadu_si256((__m256i const*)x);
 
  // Masks for extracting the lower and upper nibbles
  const __m256i low_bits_mask = _mm256_set1_epi8(0x0F);
 
  // Count bits in the lower half
  __m256i low_bits = _mm256_and_si256(data, low_bits_mask);
  __m256i low_count = _mm256_shuffle_epi8(lookup_popcount_4, low_bits);
 
  // Count bits in the upper half
  __m256i high_bits = _mm256_srli_epi16(data, 4);
  high_bits = _mm256_and_si256(high_bits, low_bits_mask);
  __m256i high_count = _mm256_shuffle_epi8(lookup_popcount_4, high_bits);
 
  // Pack the results into a single output vector
  __m256i result_vec =
      _mm256_or_si256(low_count, _mm256_slli_epi16(high_count, 4));
  _mm256_storeu_epi8(result, result_vec);
#else
  // Fallback implementation for non-AVX2 platforms
  for (size_t i = 0; i < 32; i++) {
    // Count bits in the lower half
    uint8_t a = x[i] & 0x0F;
    uint8_t low_count = std::popcount(a);
    // Count bits in the upper half
    a = (x[i] >> 4) & 0x0F;
    uint8_t high_count = std::popcount(a);
 
    // Pack the counts into the output byte
    result[i] = low_count | (high_count << 4);
  }
#endif
}

void popcount_32x8(const uint8_t* x, uint8_t* result) {
#ifdef PIXIE_AVX512_SUPPORT
  // Load 64 4-bit integers (256 bits total)
  __m256i data = _mm256_loadu_si256((__m256i const*)x);
  auto popcount_8 = _mm256_popcnt_epi8(data);
  _mm256_storeu_si256((__m256i*)result, popcount_8);
#else
#ifdef PIXIE_AVX2_SUPPORT
  // Load 64 4-bit integers (256 bits total)
  __m256i data = _mm256_loadu_si256((__m256i const*)x);
 
  // Masks for extracting the lower and upper nibbles
  const __m256i low_bits_mask = _mm256_set1_epi8(0x0F);
 
  // Count bits in lower half
  __m256i low_bits = _mm256_and_si256(data, low_bits_mask);
  __m256i low_count = _mm256_shuffle_epi8(lookup_popcount_4, low_bits);
 
  // Count bits upper half
  __m256i high_bits = _mm256_srli_epi16(data, 4);
  high_bits = _mm256_and_si256(high_bits, low_bits_mask);
  __m256i high_count = _mm256_shuffle_epi8(lookup_popcount_4, high_bits);
 
  __m256i result_vec = _mm256_add_epi8(low_count, high_count);
  _mm256_storeu_si256((__m256i*)result, result_vec);
#else
  // Fallback implementation for non-AVX2 platforms
  for (size_t i = 0; i < 32; i++) {
    result[i] = std::popcount(x[i]);
  }
#endif
#endif
}

void rank_32x8(const uint8_t* x, uint8_t* result) {
#ifdef PIXIE_AVX512_SUPPORT
  // Step 1: Calculate popcount of each byte
  popcount_32x8(x, result);
  __m256i prefix_sums = _mm256_loadu_si256((__m256i const*)result);
  const __m256i zero = _mm256_setzero_si256();
 
  prefix_sums = _mm256_add_epi8(prefix_sums,
                                _mm256_alignr_epi8(prefix_sums, zero, 16 - 1));
  prefix_sums = _mm256_add_epi8(prefix_sums,
                                _mm256_alignr_epi8(prefix_sums, zero, 16 - 2));
  prefix_sums = _mm256_add_epi8(prefix_sums,
                                _mm256_alignr_epi8(prefix_sums, zero, 16 - 4));
  prefix_sums = _mm256_add_epi8(prefix_sums,
                                _mm256_alignr_epi8(prefix_sums, zero, 16 - 8));
 
  // At this point we have prefix sums for two halfs, the last step is to
  // extract 16-th value and add it to the whole second half
  __m128i low_lane = _mm256_extracti128_si256(prefix_sums, 0);
  __m128i high_lane = _mm256_extracti128_si256(prefix_sums, 1);
  auto last_val_low = _mm_extract_epi8(low_lane, 15);
  __m128i add_to_high = _mm_set1_epi8(last_val_low);
  high_lane = _mm_add_epi8(high_lane, add_to_high);
  prefix_sums = _mm256_set_m128i(high_lane, low_lane);
  _mm256_storeu_epi8(result, prefix_sums);
#else
  // Scalar fallback implementation
  result[0] = std::popcount(x[0]);
  for (size_t i = 1; i < 32; ++i) {
    result[i] = std::popcount(x[i]) + result[i - 1];
  }
#endif
}