sequence_tree_kernels.h

#pragma once
 
#include <bit>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <span>
 
#if SIZE_MAX == UINT64_MAX && (defined(__AVX512F__) || defined(__AVX2__))
#include <immintrin.h>
#endif

namespace pixie::permutations::detail {

inline std::size_t node_select_scalar(std::span<const std::size_t> ends,
                                      std::size_t index) noexcept {
  std::size_t i = 0;
  while (i < ends.size() && ends[i] <= index) {
    ++i;
  }
  return i;
}

inline std::size_t node_select(std::span<const std::size_t> ends,
                               std::size_t index) noexcept {
  if (ends.empty()) {
    return 0;
  }
#if SIZE_MAX == UINT64_MAX && defined(__AVX512F__)
  const auto query =
      _mm512_set1_epi64(std::bit_cast<long long>(std::uint64_t{index}));
  for (std::size_t i = 0; i < ends.size();) {
    const auto remaining = ends.size() - i;
    const auto count = remaining < 8 ? remaining : 8;
    const auto active = static_cast<__mmask8>((1u << count) - 1);
    const auto values = _mm512_maskz_loadu_epi64(active, ends.data() + i);
    const auto greater = static_cast<unsigned>(
        _mm512_mask_cmp_epu64_mask(active, values, query, _MM_CMPINT_GT));
    if (greater != 0) {
      return i + std::countr_zero(greater);
    }
    i += count;
  }
  return ends.size();
#elif SIZE_MAX == UINT64_MAX && defined(__AVX2__)
  const auto sign = _mm256_set1_epi64x(std::numeric_limits<long long>::min());
  const auto query = _mm256_xor_si256(
      _mm256_set1_epi64x(std::bit_cast<long long>(std::uint64_t{index})), sign);
  std::size_t i = 0;
  for (; ends.size() - i >= 4; i += 4) {
    const auto values =
        _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ends.data() + i));
    const auto greater =
        static_cast<unsigned>(_mm256_movemask_pd(_mm256_castsi256_pd(
            _mm256_cmpgt_epi64(_mm256_xor_si256(values, sign), query))));
    if (greater != 0) {
      return i + std::countr_zero(greater);
    }
  }
  return i + node_select_scalar(ends.subspan(i), index);
#else
  return node_select_scalar(ends, index);
#endif
}
 
}  // namespace pixie::permutations::detail