diff --git a/CMakeLists.txt b/CMakeLists.txt index 87c59ad..203cdcf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,6 +15,13 @@ if(DISABLE_AVX512) message(STATUS "DISABLE_AVX512 is ON") endif() +option(PIXIE_DISABLE_BMI2 "Disable BMI2 instructions" OFF) +if(PIXIE_DISABLE_BMI2) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mno-bmi2") + add_compile_definitions(PIXIE_DISABLE_BMI2) + message(STATUS "PIXIE_DISABLE_BMI2 is ON") +endif() + option(ENABLE_ADDRESS_SANITIZER "Enable AddressSanitizer" OFF) if(ENABLE_ADDRESS_SANITIZER) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=address -fno-omit-frame-pointer") @@ -175,6 +182,15 @@ if(PIXIE_TESTS) gtest_main ${PIXIE_DIAGNOSTICS_LIBS}) + add_executable(select512_experimental_tests + src/tests/select512_experimental_tests.cpp) + target_include_directories(select512_experimental_tests + PUBLIC include) + target_link_libraries(select512_experimental_tests + gtest + gtest_main + ${PIXIE_DIAGNOSTICS_LIBS}) + add_executable(dfuds_tree_tests src/tests/dfuds_tree_tests.cpp) target_include_directories(dfuds_tree_tests @@ -265,6 +281,15 @@ if(PIXIE_BENCHMARKS) benchmark benchmark_main ${PIXIE_DIAGNOSTICS_LIBS}) + + add_executable(select512_benchmarks + src/benchmarks/select512_benchmarks.cpp) + target_include_directories(select512_benchmarks + PUBLIC include) + target_link_libraries(select512_benchmarks + benchmark + benchmark_main + ${PIXIE_DIAGNOSTICS_LIBS}) endif() # --------------------------------------------------------------------------- diff --git a/include/pixie/bits.h b/include/pixie/bits.h index 7d26597..210070c 100644 --- a/include/pixie/bits.h +++ b/include/pixie/bits.h @@ -45,6 +45,10 @@ static inline const __m256i mask_first_half = _mm256_setr_epi8( // clang-format on #endif +#if defined(__BMI2__) && !defined(PIXIE_DISABLE_BMI2) +#define PIXIE_BMI2_SUPPORT +#endif + /** * @brief Test 16 int16 RmM btree child ranges for a node-local target. * @details Each lane represents one child summary. The function checks whether @@ -235,37 +239,144 @@ static inline uint64_t rank_512(const uint64_t* x, uint64_t count) { #endif } +#ifndef PIXIE_BMI2_SUPPORT +struct PixieSelectByteLut { + uint8_t popcounts[256]; + uint8_t select[256][8]; + + constexpr PixieSelectByteLut() : popcounts{}, select{} { + for (int byte = 0; byte < 256; ++byte) { + for (int rank = 0; rank < 8; ++rank) { + select[byte][rank] = 8; + } + + int count = 0; + for (int bit = 0; bit < 8; ++bit) { + if (((byte >> bit) & 1) != 0) { + select[byte][count++] = static_cast(bit); + } + } + popcounts[byte] = static_cast(count); + } + } +}; + +static inline constexpr PixieSelectByteLut pixie_select_byte_lut; + +static inline uint64_t select_64_no_bmi2(uint64_t x, uint64_t rank) { + uint64_t offset = 0; + + uint64_t count = std::popcount(static_cast(x)); + if (rank >= count) { + rank -= count; + x >>= 32; + offset += 32; + } + + count = std::popcount(static_cast(x)); + if (rank >= count) { + rank -= count; + x >>= 16; + offset += 16; + } + + const auto low_byte = static_cast(x); + count = pixie_select_byte_lut.popcounts[low_byte]; + if (rank >= count) { + rank -= count; + x >>= 8; + offset += 8; + } + + return offset + pixie_select_byte_lut.select[static_cast(x)][rank]; +} +#endif + /** * @brief Return position of @p rank 1 bit in @p x */ static inline uint64_t select_64(uint64_t x, uint64_t rank) { - return _tzcnt_u64(_pdep_u64(1ull << rank, x)); +#ifdef PIXIE_BMI2_SUPPORT + return std::countr_zero(_pdep_u64(1ull << rank, x)); +#else + return select_64_no_bmi2(x, rank); +#endif } -/** - * @brief Return position of @p rank 1 bit in @p x - * @details Selecting within 64-bit word can be done - * using combination of _tzcnt_u64(_pdep_u64(1 << rank, x)) - * See Pandey P., Bender M. A., Johnson R. A fast x86 implementation of select - * https://arxiv.org/abs/1706.00990 - * - * To find a 64-bit word inside a 512-bit region we use - * We first popcounts of all 8 64-bit words with _mm512_popcnt_epi64 - * and then perform a linear scan. - * - * Notably a SWAR algorithm for parallel binary search - * http://www-graphics.stanford.edu/~seander/bithacks.html#SelectPosFromMSBRank - * might be used as a backoff algorithm for selecting in a 64-bit word. - * It can also be used as an alternative for linear search but i don't - * see a proper SIMD algorithm to make it faster. - */ -static inline uint64_t select_512(const uint64_t* x, uint64_t rank) { -#ifdef PIXIE_AVX512_SUPPORT +template +static inline uint64_t select_512_word_count(uint64_t word) { + if constexpr (Invert) { + return std::popcount(~word); + } else { + return std::popcount(word); + } +} - __m512i res = _mm512_loadu_epi64(x); - __m512i counts = _mm512_popcnt_epi64(res); - __m512i prefix = counts; +template +static inline uint64_t select_512_selected_word(uint64_t word) { + if constexpr (Invert) { + return ~word; + } else { + return word; + } +} + +template +static inline uint64_t select_512_scalar_impl(const uint64_t* x, + uint64_t rank) { + for (size_t i = 0; i < 8; ++i) { + const uint64_t count = select_512_word_count(x[i]); + if (rank < count) { + return i * 64 + select_64(select_512_selected_word(x[i]), rank); + } + rank -= count; + } + return 512; +} +#ifdef PIXIE_AVX2_SUPPORT +template +static inline void select_512_avx2_counts(const uint64_t* x, uint64_t* counts) { + const __m256i low_mask = _mm256_set1_epi8(0x0F); + const __m256i zero = _mm256_setzero_si256(); + const __m256i sixty_four = _mm256_set1_epi64x(64); + + for (int half = 0; half < 2; ++half) { + const __m256i words = + _mm256_loadu_si256(reinterpret_cast(x + 4 * half)); + + const __m256i low_nibbles = _mm256_and_si256(words, low_mask); + const __m256i high_nibbles = + _mm256_and_si256(_mm256_srli_epi16(words, 4), low_mask); + const __m256i byte_counts = + _mm256_add_epi8(_mm256_shuffle_epi8(lookup_popcount_4, low_nibbles), + _mm256_shuffle_epi8(lookup_popcount_4, high_nibbles)); + __m256i word_counts = _mm256_sad_epu8(byte_counts, zero); + if constexpr (Invert) { + word_counts = _mm256_sub_epi64(sixty_four, word_counts); + } + _mm256_store_si256(reinterpret_cast<__m256i*>(counts + 4 * half), + word_counts); + } +} + +template +static inline uint64_t select_512_avx2_impl(const uint64_t* x, uint64_t rank) { + alignas(32) uint64_t counts[8]; + select_512_avx2_counts(x, counts); + + for (size_t i = 0; i < 8; ++i) { + if (rank < counts[i]) { + return i * 64 + select_64(select_512_selected_word(x[i]), rank); + } + rank -= counts[i]; + } + return 512; +} +#endif + +#ifdef PIXIE_AVX512_SUPPORT +static inline __m512i select_512_avx512_prefix_sum_u64(__m512i prefix) { 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); @@ -275,76 +386,67 @@ static inline uint64_t select_512(const uint64_t* x, uint64_t rank) { 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); + return _mm512_add_epi64(prefix, tmp); +} - __mmask8 mask = _mm512_cmpgt_epu64_mask(prefix, _mm512_set1_epi64(rank)); - uint32_t i = _tzcnt_u32(static_cast(mask)); - uint64_t prev = 0; - if (i != 0) { - __m512i idx_prev = _mm512_set1_epi64(static_cast(i - 1)); - __m512i prev_vec = _mm512_permutexvar_epi64(idx_prev, prefix); - prev = static_cast( - _mm_cvtsi128_si64(_mm512_castsi512_si128(prev_vec))); +static inline uint64_t select_512_avx512_previous_prefix(__m512i prefix, + uint32_t lane) { + if (lane == 0) { + return 0; } - return i * 64 + select_64(x[i], rank - prev); - -#else + const __m512i idx_previous = + _mm512_set1_epi64(static_cast(lane - 1)); + const __m512i previous_vec = _mm512_permutexvar_epi64(idx_previous, prefix); + return static_cast( + _mm_cvtsi128_si64(_mm512_castsi512_si128(previous_vec))); +} - size_t i = 0; - int popcount = std::popcount(x[0]); - while (i < 7 && popcount <= rank) { - rank -= popcount; - popcount = std::popcount(x[++i]); +template +static inline uint64_t select_512_avx512_impl(const uint64_t* x, + uint64_t rank) { + const __m512i words = _mm512_loadu_epi64(x); + __m512i prefix = _mm512_popcnt_epi64(words); + if constexpr (Invert) { + prefix = _mm512_sub_epi64(_mm512_set1_epi64(64), prefix); } - return i * 64 + select_64(x[i], rank); + prefix = select_512_avx512_prefix_sum_u64(prefix); + + const __mmask8 mask = _mm512_cmpgt_epu64_mask( + prefix, _mm512_set1_epi64(static_cast(rank))); + const uint32_t lane = std::countr_zero(static_cast(mask)); + const uint64_t previous = select_512_avx512_previous_prefix(prefix, lane); + return lane * 64 + + select_64(select_512_selected_word(x[lane]), rank - previous); +} +#endif +/** + * @brief Return position of @p rank 1 bit in @p x. + * @details Uses AVX-512, then AVX2, then scalar fallback. The 64-bit in-word + * select uses BMI2 PDEP unless PIXIE_DISABLE_BMI2 is defined or BMI2 is not + * available. + */ +static inline uint64_t select_512(const uint64_t* x, uint64_t rank) { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#elif defined(PIXIE_AVX2_SUPPORT) + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_impl(x, rank); #endif } /** - * @brief Return position of @p rank0 0 bit in @p x + * @brief Return position of @p rank0 0 bit in @p x. * @details select_512 with bit inversion. */ 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(mask)); - uint64_t prev = 0; - if (i != 0) { - __m512i idx_prev = _mm512_set1_epi64(static_cast(i - 1)); - __m512i prev_vec = _mm512_permutexvar_epi64(idx_prev, prefix); - prev = static_cast( - _mm_cvtsi128_si64(_mm512_castsi512_si128(prev_vec))); - } - return i * 64 + select_64(~x[i], rank0 - prev); - + return select_512_avx512_impl(x, rank0); +#elif defined(PIXIE_AVX2_SUPPORT) + return select_512_avx2_impl(x, rank0); #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); - + return select_512_scalar_impl(x, rank0); #endif } diff --git a/include/pixie/experimental/select.h b/include/pixie/experimental/select.h new file mode 100644 index 0000000..155c38b --- /dev/null +++ b/include/pixie/experimental/select.h @@ -0,0 +1,737 @@ +#pragma once + +#include + +#include +#include +#include + +namespace pixie::experimental { + +struct SelectByteLut { + uint8_t popcounts[256]; + uint8_t select[256][8]; + + constexpr SelectByteLut() : popcounts{}, select{} { + for (int b = 0; b < 256; ++b) { + for (int r = 0; r < 8; ++r) { + select[b][r] = 8; + } + + int rank = 0; + for (int i = 0; i < 8; ++i) { + if (((b >> i) & 1) != 0) { + select[b][rank++] = static_cast(i); + } + } + popcounts[b] = static_cast(rank); + } + } +}; + +inline constexpr SelectByteLut kSelectByteLut; + +static inline uint64_t select_64_pdep(uint64_t x, uint64_t rank) noexcept { + return ::select_64(x, rank); +} + +static inline uint64_t select_64_vigna_broadword(uint64_t x, + uint64_t rank) noexcept { + constexpr uint64_t kOnesStep4 = 0x1111111111111111ull; + constexpr uint64_t kOnesStep8 = 0x0101010101010101ull; + constexpr uint64_t kHighBitsStep8 = 0x8080808080808080ull; + + uint64_t sums = x; + sums = sums - ((sums & (0xAull * kOnesStep4)) >> 1); + sums = (sums & (0x3ull * kOnesStep4)) + ((sums >> 2) & (0x3ull * kOnesStep4)); + sums = (sums + (sums >> 4)) & (0xFull * kOnesStep8); + + const uint64_t byte_sums = sums * kOnesStep8; + const uint64_t rank_bytes = rank * kOnesStep8; + const uint64_t ge_rank = + ((rank_bytes | kHighBitsStep8) - byte_sums) & kHighBitsStep8; + const uint64_t byte_index = std::popcount(ge_rank); + const uint64_t byte_rank = + rank - (((byte_sums << 8) >> (8 * byte_index)) & 0xFFull); + const auto byte = static_cast(x >> (8 * byte_index)); + return 8 * byte_index + kSelectByteLut.select[byte][byte_rank]; +} + +static inline uint64_t select_64_byte_lut(uint64_t x, uint64_t rank) noexcept { + for (uint64_t byte_index = 0; byte_index < 8; ++byte_index) { + const auto byte = static_cast(x >> (8 * byte_index)); + const uint8_t count = kSelectByteLut.popcounts[byte]; + if (rank < count) { + return 8 * byte_index + kSelectByteLut.select[byte][rank]; + } + rank -= count; + } + return 64; +} + +static inline uint64_t select_64_broadword_lut(uint64_t x, + uint64_t rank) noexcept { + constexpr uint64_t kOnes = 0x0101010101010101ull; + constexpr uint64_t kHighBits = 0x8080808080808080ull; + + uint64_t byte_counts = x - ((x >> 1) & 0x5555555555555555ull); + byte_counts = (byte_counts & 0x3333333333333333ull) + + ((byte_counts >> 2) & 0x3333333333333333ull); + byte_counts = (byte_counts + (byte_counts >> 4)) & 0x0F0F0F0F0F0F0F0Full; + + const uint64_t prefix = byte_counts * kOnes; + const uint64_t threshold = rank + 1; + const uint64_t prefix_less_than_threshold = + (kOnes * (127 + threshold) - (prefix & 0x7F7F7F7F7F7F7F7Full)) & ~prefix & + kHighBits; + const uint64_t byte_index = std::popcount(prefix_less_than_threshold); + const uint64_t previous = + byte_index == 0 ? 0 : ((prefix >> (8 * (byte_index - 1))) & 0xFFull); + const auto byte = static_cast(x >> (8 * byte_index)); + return 8 * byte_index + kSelectByteLut.select[byte][rank - previous]; +} + +static inline uint64_t select_64_binary_lut(uint64_t x, + uint64_t rank) noexcept { + uint64_t offset = 0; + + uint64_t count = std::popcount(static_cast(x)); + if (rank >= count) { + rank -= count; + x >>= 32; + offset += 32; + } + + count = std::popcount(static_cast(x)); + if (rank >= count) { + rank -= count; + x >>= 16; + offset += 16; + } + + const auto low_byte = static_cast(x); + count = kSelectByteLut.popcounts[low_byte]; + if (rank >= count) { + rank -= count; + x >>= 8; + offset += 8; + } + + return offset + kSelectByteLut.select[static_cast(x)][rank]; +} + +static inline uint64_t select_64_no_pdep(uint64_t x, uint64_t rank) noexcept { + return select_64_binary_lut(x, rank); +} + +static inline uint64_t select_64_experimental_default(uint64_t x, + uint64_t rank) noexcept { +#ifdef PIXIE_EXPERIMENTAL_SELECT_NO_PDEP + return select_64_no_pdep(x, rank); +#else + return select_64_pdep(x, rank); +#endif +} + +template +static inline uint64_t select_512_word_count(uint64_t word) noexcept { + if constexpr (Invert) { + return std::popcount(~word); + } else { + return std::popcount(word); + } +} + +template +static inline uint64_t select_512_select_word(uint64_t word) noexcept { + if constexpr (Invert) { + return ~word; + } else { + return word; + } +} + +template +static inline uint64_t select_512_scalar_impl(const uint64_t* x, + uint64_t rank) noexcept { + for (size_t i = 0; i < 8; ++i) { + const uint64_t count = select_512_word_count(x[i]); + if (rank < count) { + const uint64_t word = select_512_select_word(x[i]); + return i * 64 + SelectWord(word, rank); + } + rank -= count; + } + return 512; +} + +static inline uint64_t select_512_scalar_pdep(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_pdep(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_broadword_lut(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_broadword_lut( + const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_vigna_broadword( + const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_vigna_broadword( + const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select_512_scalar_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +static inline uint64_t select0_512_scalar_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { + return select_512_scalar_impl(x, rank); +} + +#ifdef PIXIE_AVX2_SUPPORT +template +static inline void select_512_avx2_counts(const uint64_t* x, + uint64_t* counts) noexcept { + const __m256i low_mask = _mm256_set1_epi8(0x0F); + const __m256i zero = _mm256_setzero_si256(); + const __m256i sixty_four = _mm256_set1_epi64x(64); + + for (int half = 0; half < 2; ++half) { + const __m256i words = + _mm256_loadu_si256(reinterpret_cast(x + 4 * half)); + + const __m256i low_nibbles = _mm256_and_si256(words, low_mask); + const __m256i high_nibbles = + _mm256_and_si256(_mm256_srli_epi16(words, 4), low_mask); + const __m256i byte_counts = + _mm256_add_epi8(_mm256_shuffle_epi8(lookup_popcount_4, low_nibbles), + _mm256_shuffle_epi8(lookup_popcount_4, high_nibbles)); + __m256i word_counts = _mm256_sad_epu8(byte_counts, zero); + if constexpr (Invert) { + word_counts = _mm256_sub_epi64(sixty_four, word_counts); + } + _mm256_store_si256(reinterpret_cast<__m256i*>(counts + 4 * half), + word_counts); + } +} + +template +static inline uint64_t select_512_avx2_impl(const uint64_t* x, + uint64_t rank) noexcept { + alignas(32) uint64_t counts[8]; + select_512_avx2_counts(x, counts); + + for (size_t i = 0; i < 8; ++i) { + if (rank < counts[i]) { + const uint64_t word = select_512_select_word(x[i]); + return i * 64 + SelectWord(word, rank); + } + rank -= counts[i]; + } + return 512; +} +#endif + +static inline uint64_t select_512_avx2_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_byte_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_byte_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_broadword_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_broadword_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_broadword_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_broadword_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_binary_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_binary_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_vigna_broadword(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_vigna_broadword(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_vigna_broadword( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_vigna_broadword(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_no_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_no_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx2_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select_512_scalar_experimental_default(x, rank); +#endif +} + +static inline uint64_t select0_512_avx2_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX2_SUPPORT + return select_512_avx2_impl(x, rank); +#else + return select0_512_scalar_experimental_default(x, rank); +#endif +} + +#ifdef PIXIE_AVX512_SUPPORT +static inline __m512i select_512_avx512_prefix_sum_u64( + __m512i counts) noexcept { + 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, counts); + counts = _mm512_add_epi64(counts, tmp); + tmp = _mm512_maskz_permutexvar_epi64(0xFC, idx_shift2, counts); + counts = _mm512_add_epi64(counts, tmp); + tmp = _mm512_maskz_permutexvar_epi64(0xF0, idx_shift4, counts); + return _mm512_add_epi64(counts, tmp); +} + +static inline uint64_t select_512_avx512_previous_prefix( + __m512i prefix, + uint32_t lane) noexcept { + if (lane == 0) { + return 0; + } + const __m512i idx_previous = + _mm512_set1_epi64(static_cast(lane - 1)); + const __m512i previous_vec = _mm512_permutexvar_epi64(idx_previous, prefix); + return static_cast( + _mm_cvtsi128_si64(_mm512_castsi512_si128(previous_vec))); +} + +template +static inline uint64_t select_512_avx512_impl(const uint64_t* x, + uint64_t rank) noexcept { + __m512i words = _mm512_loadu_epi64(x); + __m512i prefix = _mm512_popcnt_epi64(words); + if constexpr (Invert) { + prefix = _mm512_sub_epi64(_mm512_set1_epi64(64), prefix); + } + + prefix = select_512_avx512_prefix_sum_u64(prefix); + + const __mmask8 mask = _mm512_cmpgt_epu64_mask( + prefix, _mm512_set1_epi64(static_cast(rank))); + const uint32_t lane = _tzcnt_u32(static_cast(mask)); + + const uint64_t previous = select_512_avx512_previous_prefix(prefix, lane); + const uint64_t word = select_512_select_word(x[lane]); + return lane * 64 + SelectWord(word, rank - previous); +} + +template +static inline uint64_t select_512_avx512_tail_impl( + const uint64_t* x, + uint64_t rank, + uint32_t base_word, + uint32_t word_count) noexcept { + const __mmask8 active_mask = + static_cast<__mmask8>((uint32_t{1} << word_count) - 1); + const __m512i words = _mm512_maskz_loadu_epi64(active_mask, x + base_word); + __m512i counts = _mm512_popcnt_epi64(words); + if constexpr (Invert) { + counts = _mm512_maskz_sub_epi64(active_mask, _mm512_set1_epi64(64), counts); + } + + const __m512i prefix = select_512_avx512_prefix_sum_u64(counts); + const __mmask8 mask = _mm512_mask_cmpgt_epu64_mask( + active_mask, prefix, _mm512_set1_epi64(static_cast(rank))); + if (mask == 0) { + return 512; + } + + const uint32_t lane = _tzcnt_u32(static_cast(mask)); + const uint64_t previous = select_512_avx512_previous_prefix(prefix, lane); + const uint64_t word = select_512_select_word(x[base_word + lane]); + return (base_word + lane) * 64 + SelectWord(word, rank - previous); +} + +template +static inline uint64_t select_512_avx512_hybrid_impl(const uint64_t* x, + uint64_t rank) noexcept { + for (uint32_t i = 0; i < ProbeWords; ++i) { + const uint64_t count = select_512_word_count(x[i]); + if (rank < count) { + const uint64_t word = select_512_select_word(x[i]); + return i * 64 + SelectWord(word, rank); + } + rank -= count; + } + return select_512_avx512_tail_impl(x, rank, ProbeWords, + 8 - ProbeWords); +} +#endif + +static inline uint64_t select_512_avx512_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_byte_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_byte_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_byte_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_broadword_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_broadword_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_broadword_lut( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_broadword_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_binary_lut(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_binary_lut(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_binary_lut(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_vigna_broadword( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_vigna_broadword(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_vigna_broadword( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_vigna_broadword(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_no_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_no_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_no_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_hybrid1_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_hybrid1_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select0_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_hybrid2_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_hybrid2_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select0_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_hybrid4_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_hybrid4_pdep(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select0_512_scalar_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_hybrid2_no_pdep( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select_512_scalar_no_pdep(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_hybrid2_no_pdep( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_hybrid_impl(x, rank); +#else + return select0_512_scalar_no_pdep(x, rank); +#endif +} + +static inline uint64_t select_512_avx512_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select_512_avx2_experimental_default(x, rank); +#endif +} + +static inline uint64_t select0_512_avx512_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_AVX512_SUPPORT + return select_512_avx512_impl(x, rank); +#else + return select0_512_avx2_experimental_default(x, rank); +#endif +} + +static inline uint64_t select_512_experimental_default(const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_EXPERIMENTAL_SELECT_NO_PDEP + return select_512_scalar_no_pdep(x, rank); +#elif defined(PIXIE_AVX512_SUPPORT) + return select_512_avx512_experimental_default(x, rank); +#else + return select_512_scalar_experimental_default(x, rank); +#endif +} + +static inline uint64_t select0_512_experimental_default( + const uint64_t* x, + uint64_t rank) noexcept { +#ifdef PIXIE_EXPERIMENTAL_SELECT_NO_PDEP + return select0_512_scalar_no_pdep(x, rank); +#elif defined(PIXIE_AVX512_SUPPORT) + return select0_512_avx512_experimental_default(x, rank); +#else + return select0_512_scalar_experimental_default(x, rank); +#endif +} + +} // namespace pixie::experimental diff --git a/src/benchmarks/select512_benchmarks.cpp b/src/benchmarks/select512_benchmarks.cpp new file mode 100644 index 0000000..16774a4 --- /dev/null +++ b/src/benchmarks/select512_benchmarks.cpp @@ -0,0 +1,314 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +using pixie::experimental::select0_512_avx2_binary_lut; +using pixie::experimental::select0_512_avx2_broadword_lut; +using pixie::experimental::select0_512_avx2_byte_lut; +using pixie::experimental::select0_512_avx2_no_pdep; +using pixie::experimental::select0_512_avx2_pdep; +using pixie::experimental::select0_512_avx2_vigna_broadword; +using pixie::experimental::select0_512_avx512_binary_lut; +using pixie::experimental::select0_512_avx512_broadword_lut; +using pixie::experimental::select0_512_avx512_byte_lut; +using pixie::experimental::select0_512_avx512_experimental_default; +using pixie::experimental::select0_512_avx512_hybrid1_pdep; +using pixie::experimental::select0_512_avx512_hybrid2_no_pdep; +using pixie::experimental::select0_512_avx512_hybrid2_pdep; +using pixie::experimental::select0_512_avx512_hybrid4_pdep; +using pixie::experimental::select0_512_avx512_no_pdep; +using pixie::experimental::select0_512_avx512_pdep; +using pixie::experimental::select0_512_avx512_vigna_broadword; +using pixie::experimental::select0_512_experimental_default; +using pixie::experimental::select0_512_scalar_binary_lut; +using pixie::experimental::select0_512_scalar_broadword_lut; +using pixie::experimental::select0_512_scalar_byte_lut; +using pixie::experimental::select0_512_scalar_experimental_default; +using pixie::experimental::select0_512_scalar_no_pdep; +using pixie::experimental::select0_512_scalar_pdep; +using pixie::experimental::select0_512_scalar_vigna_broadword; +using pixie::experimental::select_512_avx2_binary_lut; +using pixie::experimental::select_512_avx2_broadword_lut; +using pixie::experimental::select_512_avx2_byte_lut; +using pixie::experimental::select_512_avx2_no_pdep; +using pixie::experimental::select_512_avx2_pdep; +using pixie::experimental::select_512_avx2_vigna_broadword; +using pixie::experimental::select_512_avx512_binary_lut; +using pixie::experimental::select_512_avx512_broadword_lut; +using pixie::experimental::select_512_avx512_byte_lut; +using pixie::experimental::select_512_avx512_experimental_default; +using pixie::experimental::select_512_avx512_hybrid1_pdep; +using pixie::experimental::select_512_avx512_hybrid2_no_pdep; +using pixie::experimental::select_512_avx512_hybrid2_pdep; +using pixie::experimental::select_512_avx512_hybrid4_pdep; +using pixie::experimental::select_512_avx512_no_pdep; +using pixie::experimental::select_512_avx512_pdep; +using pixie::experimental::select_512_avx512_vigna_broadword; +using pixie::experimental::select_512_experimental_default; +using pixie::experimental::select_512_scalar_binary_lut; +using pixie::experimental::select_512_scalar_broadword_lut; +using pixie::experimental::select_512_scalar_byte_lut; +using pixie::experimental::select_512_scalar_experimental_default; +using pixie::experimental::select_512_scalar_no_pdep; +using pixie::experimental::select_512_scalar_pdep; +using pixie::experimental::select_512_scalar_vigna_broadword; + +namespace { + +using Block = std::array; + +struct SelectCase { + Block block; + uint64_t rank1; + uint64_t rank0; +}; + +uint64_t make_word(std::mt19937_64& rng, int fill_permille) { + switch (fill_permille) { + case 125: + return rng() & rng() & rng(); + case 875: + return rng() | rng() | rng(); + case 500: + default: + return rng(); + } +} + +uint64_t count_ones(const Block& block) { + uint64_t count = 0; + for (uint64_t word : block) { + count += std::popcount(word); + } + return count; +} + +void ensure_has_ones_and_zeros(Block& block) { + const uint64_t ones = count_ones(block); + if (ones == 0) { + block[0] = 1; + } else if (ones == 512) { + block[0] = ~uint64_t{1}; + } +} + +uint64_t rank_for_count(uint64_t count, int rank_mode, std::mt19937_64& rng) { + switch (rank_mode) { + case 0: + return 0; + case 1: + return count / 2; + case 2: + return count - 1; + case 3: + default: + return rng() % count; + } +} + +std::vector make_cases(int fill_permille, + int rank_mode, + size_t num_blocks = 4096) { + std::mt19937_64 rng(42 + static_cast(fill_permille) * 17 + + static_cast(rank_mode) * 131); + std::vector cases(num_blocks); + for (SelectCase& select_case : cases) { + for (uint64_t& word : select_case.block) { + word = make_word(rng, fill_permille); + } + ensure_has_ones_and_zeros(select_case.block); + + const uint64_t ones = count_ones(select_case.block); + const uint64_t zeros = 512 - ones; + select_case.rank1 = rank_for_count(ones, rank_mode, rng); + select_case.rank0 = rank_for_count(zeros, rank_mode, rng); + } + return cases; +} + +} // namespace + +#define PIXIE_SELECT512_ARGS \ + ->ArgNames({"FillPermille", "RankMode"}) \ + ->ArgsProduct({{125, 500, 875}, {0, 1, 2, 3}}) + +#define PIXIE_DEFINE_SELECT512_BENCHMARK(Name, Function, RankMember) \ + static void BM_##Name(benchmark::State& state) { \ + const auto cases = make_cases(static_cast(state.range(0)), \ + static_cast(state.range(1))); \ + const size_t num_cases = cases.size(); \ + size_t idx = 0; \ + uint64_t result = 0; \ + for (auto _ : state) { \ + const SelectCase& select_case = cases[idx % num_cases]; \ + result ^= Function(select_case.block.data(), select_case.RankMember); \ + benchmark::DoNotOptimize(result); \ + ++idx; \ + } \ + state.SetItemsProcessed(state.iterations()); \ + } \ + BENCHMARK(BM_##Name) PIXIE_SELECT512_ARGS + +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_Current, ::select_512, rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarPDEP, + select_512_scalar_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarByteLUT, + select_512_scalar_byte_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarBinaryLUT, + select_512_scalar_binary_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarBroadwordLUT, + select_512_scalar_broadword_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarVignaBroadword, + select_512_scalar_vigna_broadword, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarNoPDEP, + select_512_scalar_no_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ScalarExperimentalDefault, + select_512_scalar_experimental_default, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2PDEP, + select_512_avx2_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2ByteLUT, + select_512_avx2_byte_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2BinaryLUT, + select_512_avx2_binary_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2BroadwordLUT, + select_512_avx2_broadword_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2VignaBroadword, + select_512_avx2_vigna_broadword, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX2NoPDEP, + select_512_avx2_no_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512PDEP, + select_512_avx512_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512ByteLUT, + select_512_avx512_byte_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512BinaryLUT, + select_512_avx512_binary_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512BroadwordLUT, + select_512_avx512_broadword_lut, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512VignaBroadword, + select_512_avx512_vigna_broadword, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512NoPDEP, + select_512_avx512_no_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512Hybrid1PDEP, + select_512_avx512_hybrid1_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512Hybrid2PDEP, + select_512_avx512_hybrid2_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512Hybrid4PDEP, + select_512_avx512_hybrid4_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512Hybrid2NoPDEP, + select_512_avx512_hybrid2_no_pdep, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_AVX512ExperimentalDefault, + select_512_avx512_experimental_default, + rank1); +PIXIE_DEFINE_SELECT512_BENCHMARK(Select512_ExperimentalDefault, + select_512_experimental_default, + rank1); + +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_Current, ::select0_512, rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarPDEP, + select0_512_scalar_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarByteLUT, + select0_512_scalar_byte_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarBinaryLUT, + select0_512_scalar_binary_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarBroadwordLUT, + select0_512_scalar_broadword_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarVignaBroadword, + select0_512_scalar_vigna_broadword, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarNoPDEP, + select0_512_scalar_no_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ScalarExperimentalDefault, + select0_512_scalar_experimental_default, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2PDEP, + select0_512_avx2_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2ByteLUT, + select0_512_avx2_byte_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2BinaryLUT, + select0_512_avx2_binary_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2BroadwordLUT, + select0_512_avx2_broadword_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2VignaBroadword, + select0_512_avx2_vigna_broadword, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX2NoPDEP, + select0_512_avx2_no_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512PDEP, + select0_512_avx512_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512ByteLUT, + select0_512_avx512_byte_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512BinaryLUT, + select0_512_avx512_binary_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512BroadwordLUT, + select0_512_avx512_broadword_lut, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512VignaBroadword, + select0_512_avx512_vigna_broadword, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512NoPDEP, + select0_512_avx512_no_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512Hybrid1PDEP, + select0_512_avx512_hybrid1_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512Hybrid2PDEP, + select0_512_avx512_hybrid2_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512Hybrid4PDEP, + select0_512_avx512_hybrid4_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512Hybrid2NoPDEP, + select0_512_avx512_hybrid2_no_pdep, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_AVX512ExperimentalDefault, + select0_512_avx512_experimental_default, + rank0); +PIXIE_DEFINE_SELECT512_BENCHMARK(SelectZero512_ExperimentalDefault, + select0_512_experimental_default, + rank0); + +#undef PIXIE_DEFINE_SELECT512_BENCHMARK +#undef PIXIE_SELECT512_ARGS diff --git a/src/tests/select512_experimental_tests.cpp b/src/tests/select512_experimental_tests.cpp new file mode 100644 index 0000000..507a7c2 --- /dev/null +++ b/src/tests/select512_experimental_tests.cpp @@ -0,0 +1,330 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using pixie::experimental::select0_512_avx2_binary_lut; +using pixie::experimental::select0_512_avx2_broadword_lut; +using pixie::experimental::select0_512_avx2_byte_lut; +using pixie::experimental::select0_512_avx2_no_pdep; +using pixie::experimental::select0_512_avx2_pdep; +using pixie::experimental::select0_512_avx2_vigna_broadword; +using pixie::experimental::select0_512_avx512_binary_lut; +using pixie::experimental::select0_512_avx512_broadword_lut; +using pixie::experimental::select0_512_avx512_byte_lut; +using pixie::experimental::select0_512_avx512_experimental_default; +using pixie::experimental::select0_512_avx512_hybrid1_pdep; +using pixie::experimental::select0_512_avx512_hybrid2_no_pdep; +using pixie::experimental::select0_512_avx512_hybrid2_pdep; +using pixie::experimental::select0_512_avx512_hybrid4_pdep; +using pixie::experimental::select0_512_avx512_no_pdep; +using pixie::experimental::select0_512_avx512_pdep; +using pixie::experimental::select0_512_avx512_vigna_broadword; +using pixie::experimental::select0_512_experimental_default; +using pixie::experimental::select0_512_scalar_binary_lut; +using pixie::experimental::select0_512_scalar_broadword_lut; +using pixie::experimental::select0_512_scalar_byte_lut; +using pixie::experimental::select0_512_scalar_experimental_default; +using pixie::experimental::select0_512_scalar_no_pdep; +using pixie::experimental::select0_512_scalar_pdep; +using pixie::experimental::select0_512_scalar_vigna_broadword; +using pixie::experimental::select_512_avx2_binary_lut; +using pixie::experimental::select_512_avx2_broadword_lut; +using pixie::experimental::select_512_avx2_byte_lut; +using pixie::experimental::select_512_avx2_no_pdep; +using pixie::experimental::select_512_avx2_pdep; +using pixie::experimental::select_512_avx2_vigna_broadword; +using pixie::experimental::select_512_avx512_binary_lut; +using pixie::experimental::select_512_avx512_broadword_lut; +using pixie::experimental::select_512_avx512_byte_lut; +using pixie::experimental::select_512_avx512_experimental_default; +using pixie::experimental::select_512_avx512_hybrid1_pdep; +using pixie::experimental::select_512_avx512_hybrid2_no_pdep; +using pixie::experimental::select_512_avx512_hybrid2_pdep; +using pixie::experimental::select_512_avx512_hybrid4_pdep; +using pixie::experimental::select_512_avx512_no_pdep; +using pixie::experimental::select_512_avx512_pdep; +using pixie::experimental::select_512_avx512_vigna_broadword; +using pixie::experimental::select_512_experimental_default; +using pixie::experimental::select_512_scalar_binary_lut; +using pixie::experimental::select_512_scalar_broadword_lut; +using pixie::experimental::select_512_scalar_byte_lut; +using pixie::experimental::select_512_scalar_experimental_default; +using pixie::experimental::select_512_scalar_no_pdep; +using pixie::experimental::select_512_scalar_pdep; +using pixie::experimental::select_512_scalar_vigna_broadword; +using pixie::experimental::select_64_binary_lut; +using pixie::experimental::select_64_broadword_lut; +using pixie::experimental::select_64_byte_lut; +using pixie::experimental::select_64_experimental_default; +using pixie::experimental::select_64_no_pdep; +using pixie::experimental::select_64_pdep; +using pixie::experimental::select_64_vigna_broadword; + +namespace { + +using Block = std::array; +using SelectFn = uint64_t (*)(const uint64_t*, uint64_t); + +struct SelectVariant { + const char* name; + SelectFn select1; + SelectFn select0; +}; + +static const SelectVariant kVariants[] = { + {"scalar_pdep", select_512_scalar_pdep, select0_512_scalar_pdep}, + {"scalar_byte_lut", select_512_scalar_byte_lut, + select0_512_scalar_byte_lut}, + {"scalar_binary_lut", select_512_scalar_binary_lut, + select0_512_scalar_binary_lut}, + {"scalar_broadword_lut", select_512_scalar_broadword_lut, + select0_512_scalar_broadword_lut}, + {"scalar_vigna_broadword", select_512_scalar_vigna_broadword, + select0_512_scalar_vigna_broadword}, + {"scalar_no_pdep", select_512_scalar_no_pdep, select0_512_scalar_no_pdep}, + {"scalar_experimental_default", select_512_scalar_experimental_default, + select0_512_scalar_experimental_default}, + {"avx2_pdep", select_512_avx2_pdep, select0_512_avx2_pdep}, + {"avx2_byte_lut", select_512_avx2_byte_lut, select0_512_avx2_byte_lut}, + {"avx2_binary_lut", select_512_avx2_binary_lut, + select0_512_avx2_binary_lut}, + {"avx2_broadword_lut", select_512_avx2_broadword_lut, + select0_512_avx2_broadword_lut}, + {"avx2_vigna_broadword", select_512_avx2_vigna_broadword, + select0_512_avx2_vigna_broadword}, + {"avx2_no_pdep", select_512_avx2_no_pdep, select0_512_avx2_no_pdep}, + {"avx512_pdep", select_512_avx512_pdep, select0_512_avx512_pdep}, + {"avx512_byte_lut", select_512_avx512_byte_lut, + select0_512_avx512_byte_lut}, + {"avx512_binary_lut", select_512_avx512_binary_lut, + select0_512_avx512_binary_lut}, + {"avx512_broadword_lut", select_512_avx512_broadword_lut, + select0_512_avx512_broadword_lut}, + {"avx512_vigna_broadword", select_512_avx512_vigna_broadword, + select0_512_avx512_vigna_broadword}, + {"avx512_no_pdep", select_512_avx512_no_pdep, select0_512_avx512_no_pdep}, + {"avx512_hybrid1_pdep", select_512_avx512_hybrid1_pdep, + select0_512_avx512_hybrid1_pdep}, + {"avx512_hybrid2_pdep", select_512_avx512_hybrid2_pdep, + select0_512_avx512_hybrid2_pdep}, + {"avx512_hybrid4_pdep", select_512_avx512_hybrid4_pdep, + select0_512_avx512_hybrid4_pdep}, + {"avx512_hybrid2_no_pdep", select_512_avx512_hybrid2_no_pdep, + select0_512_avx512_hybrid2_no_pdep}, + {"avx512_experimental_default", select_512_avx512_experimental_default, + select0_512_avx512_experimental_default}, + {"experimental_default", select_512_experimental_default, + select0_512_experimental_default}, +}; + +uint64_t naive_select_512(const uint64_t* s, uint64_t rank, bool value) { + uint64_t seen = 0; + for (uint64_t i = 0; i < 512; ++i) { + const bool bit = ((s[i >> 6] >> (i & 63)) & 1ull) != 0; + if (bit == value) { + if (seen == rank) { + return i; + } + ++seen; + } + } + return 512; +} + +uint64_t count_ones(const Block& block) { + uint64_t count = 0; + for (uint64_t word : block) { + count += std::popcount(word); + } + return count; +} + +std::vector sample_ranks(uint64_t count) { + std::vector ranks; + if (count == 0) { + return ranks; + } + ranks.push_back(0); + ranks.push_back(count / 2); + ranks.push_back(count - 1); + std::sort(ranks.begin(), ranks.end()); + ranks.erase(std::unique(ranks.begin(), ranks.end()), ranks.end()); + return ranks; +} + +void expect_select1(const Block& block, + uint64_t rank, + const char* label, + int case_id) { + const uint64_t expected = naive_select_512(block.data(), rank, true); + for (const SelectVariant& variant : kVariants) { + ASSERT_EQ(variant.select1(block.data(), rank), expected) + << variant.name << " select1 rank=" << rank << " case=" << case_id + << " label=" << label; + } +} + +void expect_select0(const Block& block, + uint64_t rank, + const char* label, + int case_id) { + const uint64_t expected = naive_select_512(block.data(), rank, false); + for (const SelectVariant& variant : kVariants) { + ASSERT_EQ(variant.select0(block.data(), rank), expected) + << variant.name << " select0 rank=" << rank << " case=" << case_id + << " label=" << label; + } +} + +void check_sampled_ranks(const Block& block, const char* label, int case_id) { + const uint64_t ones = count_ones(block); + const uint64_t zeros = 512 - ones; + for (uint64_t rank : sample_ranks(ones)) { + expect_select1(block, rank, label, case_id); + } + for (uint64_t rank : sample_ranks(zeros)) { + expect_select0(block, rank, label, case_id); + } +} + +} // namespace + +TEST(Select64Experimental, ByteLutMatchesPdep) { + std::mt19937_64 rng(42); + for (int t = 0; t < 1000; ++t) { + const uint64_t word = rng(); + const uint64_t ones = std::popcount(word); + for (uint64_t rank = 0; rank < ones; ++rank) { + ASSERT_EQ(select_64_byte_lut(word, rank), select_64_pdep(word, rank)) + << "byte_lut case=" << t << " rank=" << rank; + ASSERT_EQ(select_64_binary_lut(word, rank), select_64_pdep(word, rank)) + << "binary_lut case=" << t << " rank=" << rank; + ASSERT_EQ(select_64_broadword_lut(word, rank), select_64_pdep(word, rank)) + << "broadword_lut case=" << t << " rank=" << rank; + ASSERT_EQ(select_64_vigna_broadword(word, rank), + select_64_pdep(word, rank)) + << "vigna_broadword case=" << t << " rank=" << rank; + ASSERT_EQ(select_64_no_pdep(word, rank), select_64_pdep(word, rank)) + << "no_pdep case=" << t << " rank=" << rank; + ASSERT_EQ(select_64_experimental_default(word, rank), + select_64_pdep(word, rank)) + << "experimental_default case=" << t << " rank=" << rank; + } + } +} + +TEST(Select512Experimental, AllOnes) { + const Block block = {std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max()}; + for (uint64_t rank = 0; rank < 512; ++rank) { + expect_select1(block, rank, "all_ones", 0); + } +} + +TEST(Select512Experimental, AllZeros) { + const Block block = {0, 0, 0, 0, 0, 0, 0, 0}; + for (uint64_t rank = 0; rank < 512; ++rank) { + expect_select0(block, rank, "all_zeros", 0); + } +} + +TEST(Select512Experimental, Alternating) { + const Block block = {0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull}; + for (uint64_t rank = 0; rank < 256; ++rank) { + expect_select1(block, rank, "alternating", 0); + expect_select0(block, rank, "alternating", 0); + } +} + +TEST(Select512Experimental, SingleBitPositions) { + for (int position = 0; position < 512; ++position) { + Block block = {0, 0, 0, 0, 0, 0, 0, 0}; + block[position >> 6] = uint64_t{1} << (position & 63); + expect_select1(block, 0, "single_bit", position); + check_sampled_ranks(block, "single_bit", position); + } +} + +TEST(Select512Experimental, EmptyWordGaps) { + const Block block = {0, 0x8000000000000000ull, + 0, 0x0123456789ABCDEFull, + 0, std::numeric_limits::max(), + 0, 0x0000000000000001ull}; + const uint64_t ones = count_ones(block); + const uint64_t zeros = 512 - ones; + for (uint64_t rank = 0; rank < ones; ++rank) { + expect_select1(block, rank, "empty_word_gaps", 0); + } + for (uint64_t rank : sample_ranks(zeros)) { + expect_select0(block, rank, "empty_word_gaps", 0); + } +} + +TEST(Select512Experimental, ExhaustiveLow16) { + for (uint64_t pattern = 0; pattern < (1ull << 16); ++pattern) { + Block block = {pattern, 0, 0, 0, 0, 0, 0, 0}; + const uint64_t ones = std::popcount(pattern); + const uint64_t low_zeros = 16 - ones; + const uint64_t zeros = 512 - ones; + + for (uint64_t rank = 0; rank < ones; ++rank) { + expect_select1(block, rank, "exhaustive_low16", + static_cast(pattern)); + } + for (uint64_t rank = 0; rank < low_zeros; ++rank) { + expect_select0(block, rank, "exhaustive_low16", + static_cast(pattern)); + } + expect_select0(block, zeros - 1, "exhaustive_low16", + static_cast(pattern)); + } +} + +TEST(Select512Experimental, Random) { + const int cases = [] { + const char* env = std::getenv("SELECT512_CASES"); + return env ? std::atoi(env) : 1000; + }(); + const uint64_t seed = [] { + const char* env = std::getenv("SELECT512_SEED"); + return env ? std::strtoull(env, nullptr, 10) : 42ull; + }(); + + std::mt19937_64 rng(seed); + for (int t = 0; t < cases; ++t) { + Block block; + for (uint64_t& word : block) { + word = rng(); + } + + check_sampled_ranks(block, "random", t); + + const uint64_t ones = count_ones(block); + const uint64_t zeros = 512 - ones; + for (int i = 0; i < 8; ++i) { + if (ones != 0) { + expect_select1(block, rng() % ones, "random", t); + } + if (zeros != 0) { + expect_select0(block, rng() % zeros, "random", t); + } + } + } +} diff --git a/src/tests/unittests.cpp b/src/tests/unittests.cpp index 1c435d3..9521b00 100644 --- a/src/tests/unittests.cpp +++ b/src/tests/unittests.cpp @@ -2,12 +2,78 @@ #include #include +#include +#include +#include +#include +#include #include #include +#include using pixie::BitVector; using pixie::BitVectorInterleaved; +namespace { + +using SelectBlock = std::array; + +uint64_t naive_select_512(const uint64_t* bits, uint64_t rank, bool value) { + uint64_t seen = 0; + for (uint64_t i = 0; i < 512; ++i) { + const bool bit = ((bits[i >> 6] >> (i & 63)) & 1ull) != 0; + if (bit == value) { + if (seen == rank) { + return i; + } + ++seen; + } + } + return 512; +} + +uint64_t count_ones(const SelectBlock& block) { + uint64_t count = 0; + for (uint64_t word : block) { + count += std::popcount(word); + } + return count; +} + +std::vector sample_ranks(uint64_t count) { + std::vector ranks; + if (count == 0) { + return ranks; + } + + ranks.push_back(0); + ranks.push_back(count / 2); + ranks.push_back(count - 1); + std::sort(ranks.begin(), ranks.end()); + ranks.erase(std::unique(ranks.begin(), ranks.end()), ranks.end()); + return ranks; +} + +void expect_select_samples(const SelectBlock& block, + const char* label, + int case_id) { + const uint64_t ones = count_ones(block); + const uint64_t zeros = 512 - ones; + + for (uint64_t rank : sample_ranks(ones)) { + ASSERT_EQ(select_512(block.data(), rank), + naive_select_512(block.data(), rank, true)) + << "select1 rank=" << rank << " case=" << case_id << " label=" << label; + } + for (uint64_t rank : sample_ranks(zeros)) { + ASSERT_EQ(select0_512(block.data(), rank), + naive_select_512(block.data(), rank, false)) + << "select0 rank=" << rank << " case=" << case_id << " label=" << label; + } +} + +} // namespace + TEST(Rank512, Ones) { std::array a{std::numeric_limits::max(), std::numeric_limits::max(), @@ -96,6 +162,60 @@ TEST(SelectZero512, Zeros) { } } +TEST(Select512, Alternating) { + const SelectBlock block = {0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull, + 0xAAAAAAAAAAAAAAAAull, 0x5555555555555555ull}; + for (uint64_t rank = 0; rank < 256; ++rank) { + ASSERT_EQ(select_512(block.data(), rank), + naive_select_512(block.data(), rank, true)); + ASSERT_EQ(select0_512(block.data(), rank), + naive_select_512(block.data(), rank, false)); + } +} + +TEST(Select512, SingleBitPositions) { + for (int position = 0; position < 512; ++position) { + SelectBlock block = {0, 0, 0, 0, 0, 0, 0, 0}; + block[position >> 6] = uint64_t{1} << (position & 63); + + ASSERT_EQ(select_512(block.data(), 0), position); + expect_select_samples(block, "single_bit", position); + } +} + +TEST(SelectZero512, SingleZeroPositions) { + for (int position = 0; position < 512; ++position) { + SelectBlock block = {std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max()}; + block[position >> 6] &= ~(uint64_t{1} << (position & 63)); + + ASSERT_EQ(select0_512(block.data(), 0), position); + expect_select_samples(block, "single_zero", position); + } +} + +TEST(Select512, EmptyWordGaps) { + const SelectBlock block = {0, 0x8000000000000000ull, + 0, 0x0123456789ABCDEFull, + 0, std::numeric_limits::max(), + 0, 0x0000000000000001ull}; + const uint64_t ones = count_ones(block); + for (uint64_t rank = 0; rank < ones; ++rank) { + ASSERT_EQ(select_512(block.data(), rank), + naive_select_512(block.data(), rank, true)) + << "rank=" << rank; + } + expect_select_samples(block, "empty_word_gaps", 0); +} + TEST(Select512, Random) { std::array a{std::numeric_limits::max(), std::numeric_limits::max(),