From 844096ccf92fc02c93cfe32b603ccffb6dcbf3ae Mon Sep 17 00:00:00 2001 From: Daniel Lemire Date: Wed, 26 Aug 2026 23:03:53 -0400 Subject: [PATCH] Speed up the AVX-512 VBMI2 bitset decoders Both bitset_extract_setbits_avx512 and bitset_extract_setbits_avx512_uint16 kept the VPCOMPRESSB kernel but stored full 64-byte blocks unconditionally, relying on the next word to overwrite the slack. That cost two things: the caller had to leave 64 values of padding (so the last values always fell back to the scalar tail), and every word paid for all four blocks even when it only produced a handful of values. Each block is now stored under a mask derived from BZHI(-1, popcount), so a store writes exactly the values it produced, and only the ceil(popcount/16) blocks that carry a value are computed at all. Empty words are skipped and the base is carried in a ZMM incremented by a constant instead of being rebuilt per word. Measured on a Xeon Gold 6548N (Emerald Rapids), ns per value, kernel only: cardinality scalar before after 256 2.91 6.93 2.95 1024 1.29 1.95 1.02 2048 0.76 1.07 0.65 4096 0.54 0.54 0.37 8192 0.42 0.27 0.18 The uint32 decoder now beats the scalar loop at every cardinality measured from 16 to 65536, so bitset_container_to_uint32_array no longer needs its cardinality heuristic. array_container_from_bitset is the opposite case: every caller checks that the result fits an array container first, so its input is always under 6.25% density, and it had no heuristic at all. The old kernel was a net regression there -- 2.4x slower than the scalar loop at 256 values. The new one is faster from about 1024 values up, so that is where the gate goes. End to end, with the repository's own benchmark: microbench/ToArray/census-income 6671 ns -> 4705 ns (1.42x) microbench/ToArray/weather_sept_85 25895 ns -> 11611 ns (2.23x) bitset_container/to_array_convert 2.08 ns -> 1.47 ns (1.42x) The other datasets hold no bitset containers and are unchanged. Benchmarks outside these paths move by up to 1.5x in either direction, but rebuilding the unmodified source reproduces the same swings, so they are code layout noise. The masked-store idea is from RoaringBitmap/roaring#557, which applies it to the Go port's ToArray. A new test compares both decoders against the scalar ones across 64 densities, at exactly the cardinality and at truncated capacities. Claude-Session: https://claude.ai/code/session_015iZSVG9KJ595ktzGqaK7qR --- include/roaring/bitset_util.h | 15 ++++ src/bitset_util.c | 135 +++++++++++++++++++++------------- src/containers/bitset.c | 6 +- src/containers/convert.c | 9 ++- tests/util_unit.c | 82 +++++++++++++++++++++ 5 files changed, 195 insertions(+), 52 deletions(-) diff --git a/include/roaring/bitset_util.h b/include/roaring/bitset_util.h index 9705e4d7b..53e4a9fa2 100644 --- a/include/roaring/bitset_util.h +++ b/include/roaring/bitset_util.h @@ -179,6 +179,15 @@ size_t bitset_extract_setbits_avx2(const uint64_t *words, size_t length, uint32_t *out, size_t outcapacity, uint32_t base); +/* + * Same as bitset_extract_setbits_avx2, using AVX-512 (VBMI2) decoding. + * + * It writes at most "outcapacity" values and requires no padding beyond them: + * "out" only needs room for the values that are actually produced. + * + * Requires AVX-512 VBMI2 (VPCOMPRESSB): only call it when + * croaring_hardware_support() reports ROARING_SUPPORTS_AVX512. + */ size_t bitset_extract_setbits_avx512(const uint64_t *words, size_t length, uint32_t *out, size_t outcapacity, uint32_t base); @@ -214,6 +223,12 @@ size_t bitset_extract_setbits_sse_uint16(const uint64_t *words, size_t length, uint16_t *out, size_t outcapacity, uint16_t base); +/* + * Same as bitset_extract_setbits_avx512, writing 16-bit values. + * + * Requires AVX-512 VBMI2 (VPCOMPRESSB): only call it when + * croaring_hardware_support() reports ROARING_SUPPORTS_AVX512. + */ size_t bitset_extract_setbits_avx512_uint16(const uint64_t *words, size_t length, uint16_t *out, size_t outcapacity, uint16_t base); diff --git a/src/bitset_util.c b/src/bitset_util.c index c75bfdf15..8ce7c6284 100644 --- a/src/bitset_util.c +++ b/src/bitset_util.c @@ -573,43 +573,73 @@ const uint8_t vbmi2_table[64] = { 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63}; +// Reference: +// https://lemire.me/blog/2022/05/10/faster-bitset-decoding-using-intel-avx-512/ +// +// One VPCOMPRESSB (AVX-512 VBMI2) turns a whole 64-bit word into the (up to +// 64) byte-sized positions of its set bits in a single shot; the positions are +// then widened 16 at a time with VPMOVZXBD and added to a running base vector. +// +// Two details matter for speed. Each block is stored under a mask derived from +// BZHI(-1, popcount), so a store writes exactly the values it produced: the +// caller needs no padding past the values it asked for, and no store bandwidth +// goes to values nobody wants. And only the ceil(popcount/16) blocks that +// carry a value are computed at all -- bitset containers are usually well +// under half full, so the later blocks are almost always skipped. size_t bitset_extract_setbits_avx512(const uint64_t *words, size_t length, uint32_t *vout, size_t outcapacity, uint32_t base) { if (outcapacity == 0) return 0; - uint32_t *out = (uint32_t *)vout; + uint32_t *out = vout; uint32_t *initout = out; uint32_t *safeout = out + outcapacity; - __m512i base_v = _mm512_set1_epi32(base); - __m512i index_table = _mm512_loadu_si512(vbmi2_table); + // base_v holds the value of bit 0 of the word being decoded. + __m512i base_v = _mm512_set1_epi32((int)base); + const __m512i inc_v = _mm512_set1_epi32(64); + const __m512i index_table = _mm512_loadu_si512(vbmi2_table); size_t i = 0; - for (; (i < length) && ((out + 64) < safeout); i += 1) { + for (; i < length; i++) { uint64_t v = words[i]; - __m512i vec = _mm512_maskz_compress_epi8(v, index_table); - - uint8_t advance = (uint8_t)roaring_hamming(v); - - __m512i vbase = - _mm512_add_epi32(base_v, _mm512_set1_epi32((int)(i * 64))); - __m512i r1 = _mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(vec, 0)); - __m512i r2 = _mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(vec, 1)); - __m512i r3 = _mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(vec, 2)); - __m512i r4 = _mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(vec, 3)); - - r1 = _mm512_add_epi32(r1, vbase); - r2 = _mm512_add_epi32(r2, vbase); - r3 = _mm512_add_epi32(r3, vbase); - r4 = _mm512_add_epi32(r4, vbase); - _mm512_storeu_si512((__m512i *)out, r1); - _mm512_storeu_si512((__m512i *)(out + 16), r2); - _mm512_storeu_si512((__m512i *)(out + 32), r3); - _mm512_storeu_si512((__m512i *)(out + 48), r4); - - out += advance; + if (v != 0) { + size_t advance = (size_t)roaring_hamming(v); + if (out + advance > safeout) break; + __m512i vec = _mm512_maskz_compress_epi8((__mmask64)v, index_table); + // BZHI(-1, advance): advance is in [1, 64] here, so the shift is + // well defined. + uint64_t bits = ~UINT64_C(0) >> (64 - advance); + _mm512_mask_storeu_epi32( + out, (__mmask16)bits, + _mm512_add_epi32( + base_v, _mm512_cvtepu8_epi32(_mm512_castsi512_si128(vec)))); + if (advance > 16) { + _mm512_mask_storeu_epi32( + out + 16, (__mmask16)(bits >> 16), + _mm512_add_epi32(base_v, + _mm512_cvtepu8_epi32( + _mm512_extracti32x4_epi32(vec, 1)))); + if (advance > 32) { + _mm512_mask_storeu_epi32( + out + 32, (__mmask16)(bits >> 32), + _mm512_add_epi32( + base_v, _mm512_cvtepu8_epi32( + _mm512_extracti32x4_epi32(vec, 2)))); + if (advance > 48) { + _mm512_mask_storeu_epi32( + out + 48, (__mmask16)(bits >> 48), + _mm512_add_epi32( + base_v, + _mm512_cvtepu8_epi32( + _mm512_extracti32x4_epi32(vec, 3)))); + } + } + } + out += advance; + } + base_v = _mm512_add_epi32(base_v, inc_v); } - base += i * 64; + base += (uint32_t)(i * 64); for (; (i < length) && (out < safeout); ++i) { uint64_t w = words[i]; @@ -628,48 +658,53 @@ size_t bitset_extract_setbits_avx512(const uint64_t *words, size_t length, return out - initout; } -// Reference: -// https://lemire.me/blog/2022/05/10/faster-bitset-decoding-using-intel-avx-512/ +// Same kernel as bitset_extract_setbits_avx512, writing 16-bit values, so the +// blocks hold 32 values instead of 16. This one is only worth calling on +// reasonably dense input: see array_container_from_bitset. size_t bitset_extract_setbits_avx512_uint16(const uint64_t *array, size_t length, uint16_t *vout, size_t capacity, uint16_t base) { if (capacity == 0) return 0; - uint16_t *out = (uint16_t *)vout; + uint16_t *out = vout; uint16_t *initout = out; uint16_t *safeout = vout + capacity; - __m512i base_v = _mm512_set1_epi16(base); - __m512i index_table = _mm512_loadu_si512(vbmi2_table); + __m512i base_v = _mm512_set1_epi16((short)base); + const __m512i inc_v = _mm512_set1_epi16(64); + const __m512i index_table = _mm512_loadu_si512(vbmi2_table); size_t i = 0; - for (; (i < length) && ((out + 64) < safeout); i++) { + for (; i < length; i++) { uint64_t v = array[i]; - __m512i vec = _mm512_maskz_compress_epi8(v, index_table); - - uint8_t advance = (uint8_t)roaring_hamming(v); - - __m512i vbase = - _mm512_add_epi16(base_v, _mm512_set1_epi16((short)(i * 64))); - __m512i r1 = _mm512_cvtepi8_epi16(_mm512_extracti32x8_epi32(vec, 0)); - __m512i r2 = _mm512_cvtepi8_epi16(_mm512_extracti32x8_epi32(vec, 1)); - - r1 = _mm512_add_epi16(r1, vbase); - r2 = _mm512_add_epi16(r2, vbase); - - _mm512_storeu_si512((__m512i *)out, r1); - _mm512_storeu_si512((__m512i *)(out + 32), r2); - out += advance; + if (v != 0) { + size_t advance = (size_t)roaring_hamming(v); + if (out + advance > safeout) break; + __m512i vec = _mm512_maskz_compress_epi8((__mmask64)v, index_table); + uint64_t bits = ~UINT64_C(0) >> (64 - advance); + _mm512_mask_storeu_epi16( + out, (__mmask32)bits, + _mm512_add_epi16( + base_v, _mm512_cvtepu8_epi16(_mm512_castsi512_si256(vec)))); + if (advance > 32) { + _mm512_mask_storeu_epi16( + out + 32, (__mmask32)(bits >> 32), + _mm512_add_epi16(base_v, + _mm512_cvtepu8_epi16( + _mm512_extracti64x4_epi64(vec, 1)))); + } + out += advance; + } + base_v = _mm512_add_epi16(base_v, inc_v); } - base += i * 64; + base = (uint16_t)(base + i * 64); for (; (i < length) && (out < safeout); ++i) { uint64_t w = array[i]; while ((w != 0) && (out < safeout)) { int r = roaring_trailing_zeroes(w); // on x64, should compile to TZCNT - uint32_t val = r + base; - memcpy(out, &val, sizeof(uint16_t)); + *out = (uint16_t)(r + base); out++; w &= (w - 1); } diff --git a/src/containers/bitset.c b/src/containers/bitset.c index 769ccee12..7a452a524 100644 --- a/src/containers/bitset.c +++ b/src/containers/bitset.c @@ -952,7 +952,11 @@ int bitset_container_to_uint32_array( #if CROARING_IS_X64 int support = croaring_hardware_support(); #if CROARING_COMPILER_SUPPORTS_AVX512 - if(( support & ROARING_SUPPORTS_AVX512 ) && (bc->cardinality >= 8192)) // heuristic + // Unlike the AVX2 kernel, the AVX-512 one needs no cardinality heuristic: + // it skips empty words and only stores the blocks that carry a value, so it + // beat the scalar loop at every cardinality measured from 16 to 65536 (on + // an Emerald Rapids Xeon, 1.3x at 512 values, 6x at 4096, 7.5x at 8192). + if( support & ROARING_SUPPORTS_AVX512 ) return (int) bitset_extract_setbits_avx512(bc->words, BITSET_CONTAINER_SIZE_IN_WORDS, out, bc->cardinality, base); else diff --git a/src/containers/convert.c b/src/containers/convert.c index 6cd31d34c..796cd6c47 100644 --- a/src/containers/convert.c +++ b/src/containers/convert.c @@ -58,7 +58,14 @@ array_container_t *array_container_from_bitset(const bitset_container_t *bits) { result->cardinality = bits->cardinality; #if CROARING_IS_X64 #if CROARING_COMPILER_SUPPORTS_AVX512 - if (croaring_hardware_support() & ROARING_SUPPORTS_AVX512) { + // Every caller first checks that the cardinality fits an array container + // (<= DEFAULT_MAX_SIZE), so the input here is sparse: under 6.25% density, + // a couple of set bits per word. + // Below ~1024 values the vector decoder does not make back its per-word + // cost and the scalar loop is as fast or faster; from 1024 up it wins + // (1.2x at 1024, 1.4x at 4096, measured on an Emerald Rapids Xeon). + if ((bits->cardinality >= 1024) && + (croaring_hardware_support() & ROARING_SUPPORTS_AVX512)) { bitset_extract_setbits_avx512_uint16( bits->words, BITSET_CONTAINER_SIZE_IN_WORDS, result->array, bits->cardinality, 0); diff --git a/tests/util_unit.c b/tests/util_unit.c index 3fffe0665..93e6677e4 100644 --- a/tests/util_unit.c +++ b/tests/util_unit.c @@ -9,6 +9,7 @@ #include #include +#include #include #ifdef __cplusplus // stronger type checking errors if C built in C++ mode @@ -74,12 +75,93 @@ DEFINE_TEST(setandextract_uint32) { } } +#if CROARING_IS_X64 && CROARING_COMPILER_SUPPORTS_AVX512 +// Checks the AVX-512 (VBMI2) decoders against the scalar ones. The kernels +// write with masks derived from the popcount of each word, so the interesting +// cases are the block boundaries (16/32/48/64 values) and an output buffer +// sized exactly to the cardinality, with no padding. +DEFINE_TEST(avx512_extract_matches_scalar) { + if ((croaring_hardware_support() & ROARING_SUPPORTS_AVX512) == 0) { + return; // no VBMI2 here + } + const size_t length = (1 << 16) / 64; + uint64_t* words = (uint64_t*)malloc(length * sizeof(uint64_t)); + assert_non_null(words); + + // A deterministic set of densities, from very sparse to almost full. + for (unsigned int step = 1; step <= 64; step++) { + size_t cardinality = 0; + for (size_t i = 0; i < length; ++i) { + uint64_t w = 0; + for (unsigned int b = (unsigned int)(i % step); b < 64; b += step) { + w |= UINT64_C(1) << b; + } + words[i] = w; + cardinality += (size_t)roaring_hamming(w); + } + + const uint32_t base32 = 1u << 20; + const uint16_t base16 = 3; + // Exactly the cardinality, plus a few truncated capacities. + size_t capacities[] = {cardinality, cardinality / 2, 64, 63, 16, 1, 0}; + for (size_t k = 0; k < sizeof(capacities) / sizeof(capacities[0]); + k++) { + size_t capacity = capacities[k]; + + uint32_t* got32 = (uint32_t*)malloc((capacity + 1) * 4); + uint32_t* want32 = (uint32_t*)malloc((capacity + 1) * 4); + uint16_t* got16 = (uint16_t*)malloc((capacity + 1) * 2); + uint16_t* want16 = (uint16_t*)malloc((capacity + 1) * 2); + assert_non_null(got32); + assert_non_null(want32); + assert_non_null(got16); + assert_non_null(want16); + + // Scalar reference, truncated to the same capacity. + size_t n = 0; + uint32_t base = base32; + for (size_t i = 0; (i < length) && (n < capacity); ++i) { + uint64_t w = words[i]; + while ((w != 0) && (n < capacity)) { + int r = roaring_trailing_zeroes(w); + want32[n] = (uint32_t)(r + base); + want16[n] = (uint16_t)(r + base16 + i * 64); + n++; + w &= (w - 1); + } + base += 64; + } + + size_t n32 = bitset_extract_setbits_avx512(words, length, got32, + capacity, base32); + size_t n16 = bitset_extract_setbits_avx512_uint16( + words, length, got16, capacity, base16); + assert_int_equal(n32, n); + assert_int_equal(n16, n); + for (size_t j = 0; j < n; ++j) { + assert_int_equal(got32[j], want32[j]); + assert_int_equal(got16[j], want16[j]); + } + + free(got32); + free(want32); + free(got16); + free(want16); + } + } + free(words); +} +#endif // CROARING_IS_X64 && CROARING_COMPILER_SUPPORTS_AVX512 + int main() { tellmeall(); const struct CMUnitTest tests[] = { cmocka_unit_test(setandextract_uint16), cmocka_unit_test(setandextract_uint32), +#if CROARING_IS_X64 && CROARING_COMPILER_SUPPORTS_AVX512 + cmocka_unit_test(avx512_extract_matches_scalar), +#endif }; return cmocka_run_group_tests(tests, NULL, NULL);