diff --git a/c_api/CMakeLists.txt b/c_api/CMakeLists.txt index d9dbe7960..065c67deb 100644 --- a/c_api/CMakeLists.txt +++ b/c_api/CMakeLists.txt @@ -77,7 +77,7 @@ endif() if(NOT WIN32) # Architecture mode to support AVX512 extensions available since Intel(R) Sapphire Rapids. # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide - target_compile_options(faiss_c_avx512_spr PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16>) + target_compile_options(faiss_c_avx512_spr PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mavx512bitalg -mpopcnt -mavx512fp16 -mavx512bf16>) else() target_compile_options(faiss_c_avx512_spr PRIVATE $<$:/arch:AVX512>) endif() diff --git a/cmake/link_to_faiss_lib.cmake b/cmake/link_to_faiss_lib.cmake index da81e8429..49d62b5fc 100644 --- a/cmake/link_to_faiss_lib.cmake +++ b/cmake/link_to_faiss_lib.cmake @@ -31,7 +31,7 @@ function(link_to_faiss_lib target) if(NOT WIN32) # Architecture mode to support AVX512 extensions available since Intel (R) Sapphire Rapids. # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide - target_compile_options(${target} PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16>) + target_compile_options(${target} PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mavx512bitalg -mpopcnt -mavx512fp16 -mavx512bf16>) else() target_compile_options(${target} PRIVATE $<$:/arch:AVX512>) endif() diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 9d893d886..fcf26ab10 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -39,9 +39,9 @@ set(FAISS_SIMD_AVX512_SRC utils/simd_impl/rabitq_avx512.cpp utils/simd_impl/super_kmeans_kernels_avx512.cpp ) -# AVX-512 sources that only require AVX512_VPOPCNTDQ on top of the baseline -# AVX-512F/CD/VL/DQ/BW feature set. In DD builds these are available on CPUs -# such as Ice Lake and Zen 4 without requiring the full SPR feature set. +# AVX-512 sources that require AVX512_VPOPCNTDQ and AVX512_BITALG on top of +# the baseline AVX-512F/CD/VL/DQ/BW feature set. In DD builds these are +# available on CPUs such as Ice Lake and Zen 4 without the full SPR set. set(FAISS_SIMD_AVX512_VPOPCNT_SRC utils/simd_impl/rabitq_avx512_vpopcnt.cpp utils/hamming_distance/hamming_avx512_vpopcnt.cpp @@ -492,7 +492,7 @@ endif() if(NOT WIN32) # Architecture mode to support AVX512 extensions available since Intel(R) Sapphire Rapids. # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide - target_compile_options(faiss_avx512_spr PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vnni -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16 -mprefer-vector-width=512 ${FAISS_BMI2_FLAGS}>) + target_compile_options(faiss_avx512_spr PRIVATE $<$:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vnni -mavx512vpopcntdq -mavx512bitalg -mpopcnt -mavx512fp16 -mavx512bf16 -mprefer-vector-width=512 ${FAISS_BMI2_FLAGS}>) else() target_compile_options(faiss_avx512_spr PRIVATE $<$:/arch:AVX512>) # we need bigobj for the swig wrapper @@ -567,17 +567,17 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt" ) - # VPOPCNT sources only require VPOPCNTDQ in addition to baseline AVX-512. + # VPOPCNT sources require VPOPCNTDQ and BITALG on top of baseline AVX-512. set_source_files_properties(${FAISS_SIMD_AVX512_VPOPCNT_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS - "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mpopcnt" + "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mavx512bitalg;-mpopcnt" ) # The remaining SPR sources require the full SPR feature set. set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS - "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mpopcnt;-mavx512vnni;-mavx512fp16;-mavx512bf16" + "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mavx512bitalg;-mpopcnt;-mavx512vnni;-mavx512fp16;-mavx512bf16" ) else() # Per-file SIMD flags (MSVC) @@ -590,7 +590,7 @@ if(FAISS_OPT_LEVEL STREQUAL "dd") TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS "/arch:AVX512" ) - # MSVC has no per-feature flag for VPOPCNTDQ. + # MSVC has no per-feature flag for VPOPCNTDQ or BITALG. set_source_files_properties(${FAISS_SIMD_AVX512_VPOPCNT_SRC} TARGET_DIRECTORY faiss PROPERTIES COMPILE_OPTIONS "/arch:AVX512" diff --git a/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h b/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h index 473990158..3f26200fa 100644 --- a/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +++ b/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h @@ -39,17 +39,50 @@ namespace faiss { namespace { +// Whether the computer can measure a whole batch of codes in one call. template -struct IVFBinaryScannerL2 : BinaryInvertedListScanner { +constexpr bool has_hamming_batch = + requires(const uint8_t* tile, const uint8_t* codes, int32_t* dis) { + HammingComputer::batch_size; + HammingComputer::get_code_size(); + HammingComputer::build_batch_query(tile, nullptr); + HammingComputer::hamming_batch(tile, codes, dis); + }; + +// The query repeated batch_size times, which is what hamming_batch() XORs +// against. Empty for a computer that does not batch, so those scanners carry +// neither the buffer nor its alignment. +template +struct BatchQueryTile {}; + +template + requires has_hamming_batch +struct BatchQueryTile { + alignas(64) uint8_t batch_query + [HammingComputer::batch_size * HammingComputer::get_code_size()]; +}; + +template +struct IVFBinaryScannerL2 : BinaryInvertedListScanner, + BatchQueryTile { HammingComputer hc; size_t code_size; bool store_pairs; IVFBinaryScannerL2(size_t code_size_, bool store_pairs_) - : code_size(code_size_), store_pairs(store_pairs_) {} + : code_size(code_size_), store_pairs(store_pairs_) { + if constexpr (has_hamming_batch) { + // The batch kernel reads a fixed stride, so a caller that pairs + // this computer with another code size would read past the codes. + FAISS_THROW_IF_NOT(code_size == HammingComputer::get_code_size()); + } + } void set_query(const uint8_t* query_vector) override { hc.set(query_vector, code_size); + if constexpr (has_hamming_batch) { + HammingComputer::build_batch_query(query_vector, this->batch_query); + } } idx_t list_no = 0; @@ -61,6 +94,31 @@ struct IVFBinaryScannerL2 : BinaryInvertedListScanner { return hc.hamming(code); } + // Measures whole batches while at least batch_size codes remain, then + // leaves codes and j on the first code the caller must measure singly. + // bound is read per lane, so a caller that raises its heap top inside + // accept prunes the rest of the batch. + template + void scan_batch_prefix( + size_t n, + const uint8_t* __restrict& codes, + size_t& j, + const uint32_t& bound, + Accept&& accept) const { + if constexpr (has_hamming_batch) { + constexpr size_t B = HammingComputer::batch_size; + int32_t batch[B]; + for (; j + B <= n; j += B, codes += B * code_size) { + HammingComputer::hamming_batch(this->batch_query, codes, batch); + for (size_t t = 0; t < B; t++) { + if (static_cast(batch[t]) < bound) { + accept(batch[t], j + t); + } + } + } + } + } + size_t scan_codes( size_t n, const uint8_t* __restrict codes, @@ -73,15 +131,22 @@ struct IVFBinaryScannerL2 : BinaryInvertedListScanner { uint32_t bound = static_cast(simi[0]); size_t nup = 0; - for (size_t j = 0; j < n; j++) { + size_t j = 0; + + auto accept = [&](int32_t dis, size_t at) { + idx_t id = store_pairs ? lo_build(list_no, at) : ids[at]; + heap_replace_top(k, simi, idxi, dis, id); + bound = static_cast(simi[0]); + nup++; + }; + + scan_batch_prefix(n, codes, j, bound, accept); + + for (; j < n; j++, codes += code_size) { uint32_t dis = hc.hamming(codes); if (dis < bound) { - idx_t id = store_pairs ? lo_build(list_no, j) : ids[j]; - heap_replace_top(k, simi, idxi, dis, id); - bound = static_cast(simi[0]); - nup++; + accept(static_cast(dis), j); } - codes += code_size; } return nup; } @@ -92,13 +157,21 @@ struct IVFBinaryScannerL2 : BinaryInvertedListScanner { const idx_t* __restrict ids, int radius, RangeQueryResult& result) const override { - for (size_t j = 0; j < n; j++) { + const uint32_t bound = static_cast(radius); + size_t j = 0; + + auto accept = [&](int32_t dis, size_t at) { + int64_t id = store_pairs ? lo_build(list_no, at) : ids[at]; + result.add(static_cast(dis), id); + }; + + scan_batch_prefix(n, codes, j, bound, accept); + + for (; j < n; j++, codes += code_size) { uint32_t dis = hc.hamming(codes); - if (dis < static_cast(radius)) { - int64_t id = store_pairs ? lo_build(list_no, j) : ids[j]; - result.add(dis, id); + if (dis < bound) { + accept(static_cast(dis), j); } - codes += code_size; } } }; diff --git a/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp b/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp index a724fe1d7..6faa86d2a 100644 --- a/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp +++ b/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp @@ -17,4 +17,8 @@ #include // clang-format on +// Fails the build if the batch path stops being selected for this computer. +static_assert(faiss::has_hamming_batch< + faiss::HammingComputer20_tpl>); + #endif // COMPILE_SIMD_AVX512_VPOPCNT diff --git a/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h b/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h index 751eede05..807b389d6 100644 --- a/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h +++ b/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h @@ -8,12 +8,16 @@ #ifndef HAMMING_COMPUTER_AVX512_VPOPCNT_H #define HAMMING_COMPUTER_AVX512_VPOPCNT_H -// AVX512_VPOPCNT HammingComputer specializations using VPOPCNTDQ. This gives +// AVX512_VPOPCNT HammingComputer specializations. The 32/64/Default kernels +// use VPOPCNTDQ; the batched 20-byte kernel uses AVX512_BITALG. This gives // a faster path than the scalar popcount fallback used in the base AVX512 // specializations when compiled without -mavx512vpopcntdq. #include #include +#include + +#include #include #include @@ -34,7 +38,6 @@ namespace faiss { } FAISS_INHERIT_HAMMING_VPOPCNT(HammingComputer16); -FAISS_INHERIT_HAMMING_VPOPCNT(HammingComputer20); FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer8); FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer16); FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer32); @@ -46,6 +49,66 @@ FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputerM8); * Custom AVX512_VPOPCNT specializations using VPOPCNTDQ. ***************************************************************************/ +template <> +struct HammingComputer20_tpl + : HammingComputer20_tpl { + using HammingComputer20_tpl::HammingComputer20_tpl; + + static constexpr size_t batch_size = 8; + static constexpr size_t kStride = get_code_size(); + // 160 bytes is what the three loads in hamming_batch() cover, and the + // 16+4 or 4+16 split it applies per lane is written out for 20 bytes: + // the literal offsets and the two group indices are not derived from + // kStride, so another width needs the body reworked, not just retuned. + static_assert(batch_size * kStride == 160); + static_assert(kStride == 20, "hamming_batch() hardcodes the 16+4 split"); + static constexpr __mmask64 kTailMask = 0xFFFFFFFFull; + + /// Writes the query repeated batch_size times. The caller owns the buffer, + /// so a computer used only through hamming() carries no batch state. + static void build_batch_query(const uint8_t* a8, uint8_t* tile) { + for (size_t k = 0; k < batch_size; k++) { + memcpy(tile + k * kStride, a8, kStride); + } + } + + static void hamming_batch( + const uint8_t* tile, + const uint8_t* codes, + int32_t* dis) { + const __m512i zero = _mm512_setzero_si512(); + const __m512i p0 = _mm512_popcnt_epi8(_mm512_xor_si512( + _mm512_loadu_si512(codes), _mm512_loadu_si512(tile))); + const __m512i p1 = _mm512_popcnt_epi8(_mm512_xor_si512( + _mm512_loadu_si512(codes + 64), _mm512_loadu_si512(tile + 64))); + const __m512i p2 = _mm512_popcnt_epi8(_mm512_xor_si512( + _mm512_maskz_loadu_epi8(kTailMask, codes + 128), + _mm512_maskz_loadu_epi8(kTailMask, tile + 128))); + + alignas(64) uint64_t grp[24]; + _mm512_store_si512(grp, _mm512_sad_epu8(p0, zero)); + _mm512_store_si512(grp + 8, _mm512_sad_epu8(p1, zero)); + _mm512_store_si512(grp + 16, _mm512_sad_epu8(p2, zero)); + + for (size_t k = 0; k < batch_size; k++) { + const size_t s = k * kStride; + const size_t g = s / 8; + uint32_t xh, qh; + if (s % 8 == 0) { + memcpy(&xh, codes + s + 16, 4); + memcpy(&qh, tile + s + 16, 4); + dis[k] = static_cast( + grp[g] + grp[g + 1] + popcount32(xh ^ qh)); + } else { + memcpy(&xh, codes + s, 4); + memcpy(&qh, tile + s, 4); + dis[k] = static_cast( + popcount32(xh ^ qh) + grp[g + 1] + grp[g + 2]); + } + } + } +}; + template <> struct HammingComputer32_tpl { const uint8_t* a8; diff --git a/faiss/utils/simd_levels.cpp b/faiss/utils/simd_levels.cpp index 576d1d150..3190cec2a 100644 --- a/faiss/utils/simd_levels.cpp +++ b/faiss/utils/simd_levels.cpp @@ -267,6 +267,8 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { unsigned int ecx7 = regs[2]; [[maybe_unused]] bool has_avx512_vnni = (ecx7 & (1 << 11)) != 0; [[maybe_unused]] bool has_avx512_vpopcntdq = (ecx7 & (1 << 14)) != 0; + // Bit 12 = AVX512_BITALG, needed for the byte-wise popcount kernels. + [[maybe_unused]] bool has_avx512_bitalg = (ecx7 & (1 << 12)) != 0; uint64_t xcr0 = xgetbv0(); @@ -292,7 +294,7 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { (1 << static_cast(SIMDLevel::AVX512)); #if defined(COMPILE_SIMD_AVX512_VPOPCNT) - if (has_avx512_vpopcntdq) { + if (has_avx512_vpopcntdq && has_avx512_bitalg) { detected_level = SIMDLevel::AVX512_VPOPCNT; supported_simd_levels |= (1 << static_cast(SIMDLevel::AVX512_VPOPCNT)); @@ -302,7 +304,8 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { #if defined(COMPILE_SIMD_AVX512_SPR) // Check for Sapphire Rapids features. // The SPR code path is compiled with AVX512_VNNI, BF16, - // FP16, and VPOPCNTDQ, so all four features are required. + // FP16 and VPOPCNTDQ, and falls back to the VPOPCNT kernels, + // which need BITALG. All five features are required. // AMD Zen 4 has VPOPCNTDQ and BF16 but not FP16, and must // remain on the AVX512_VPOPCNT level. // CPUID EAX=7, ECX=1: EAX bit 5 = AVX512_BF16 @@ -312,7 +315,7 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() { cpuid_count(7, 1, regs); const bool has_avx512_bf16 = (regs[0] & (1 << 5)) != 0; if (has_avx512_vnni && has_avx512_vpopcntdq && - has_avx512_bf16 && has_avx512_fp16) { + has_avx512_bitalg && has_avx512_bf16 && has_avx512_fp16) { detected_level = SIMDLevel::AVX512_SPR; supported_simd_levels |= (1 << static_cast(SIMDLevel::AVX512_SPR)); diff --git a/faiss/utils/simd_levels.h b/faiss/utils/simd_levels.h index 781d0b800..80e15a85f 100644 --- a/faiss/utils/simd_levels.h +++ b/faiss/utils/simd_levels.h @@ -21,7 +21,7 @@ enum class SIMDLevel { // x86 AVX2, AVX512, - AVX512_SPR, // Sapphire Rapids: AVX512 + BF16 + FP16 + VNNI + VPOPCNTDQ + AVX512_SPR, // Sapphire Rapids: AVX512_VPOPCNT + BF16 + FP16 + VNNI // arm & aarch64 ARM_NEON, ARM_SVE, // Scalable Vector Extension (ARMv8.2+) @@ -29,7 +29,8 @@ enum class SIMDLevel { RISCV_RVV, // RISC-V Vector Extension (rv64gcv) // Appended to preserve the numeric values of the existing public enum. - // AVX-512 core features plus AVX512_VPOPCNTDQ (Ice Lake, Zen 4, etc.). + // AVX-512 core features plus AVX512_VPOPCNTDQ and AVX512_BITALG + // (Ice Lake, Zen 4, etc.). AVX512_VPOPCNT, COUNT diff --git a/tests/test_binary_ivf_scanner.py b/tests/test_binary_ivf_scanner.py index ff15c2def..54c3dcf25 100644 --- a/tests/test_binary_ivf_scanner.py +++ b/tests/test_binary_ivf_scanner.py @@ -162,3 +162,69 @@ class TestBinaryIVFScanner(unittest.TestCase): self.make_scanner(rs, query), codes, ids, [2], 10, 0 ) self.assertEqual(got, []) + + def test_every_distance_matches_at_this_level(self): + """Three shapes. (64, 64) has k == n, so every distance is compared + and all eight lanes are covered; a lane returning too large a distance + would never reach a smaller heap. (2003, 12) runs the batch loop and + then the tail. (5, 12) is below the batch size, so the batch loop must + not run at all.""" + rs = np.random.RandomState(61) + query = self.random_codes(rs, 1) + + for n, k in [(2003, 12), (5, 12), (64, 64)]: + codes, ids = self.random_codes(rs, n), self.sequential_ids(n) + got, _ = self.run_scan( + self.make_scanner(rs, query), codes, ids, [n], k + ) + self.assertEqual(got, _expected(query, codes, k), f"n={n}") + + def test_range_search_matches_the_reference(self): + """range_search reaches scan_codes_range. n is not a multiple of the + batch size, so the batch loop and the tail both run.""" + rs = np.random.RandomState(63) + n, nq = 2003, 4 + + for cs in WIDTHS: + dim = cs * 8 + codes = self.random_codes(rs, n, cs) + queries = self.random_codes(rs, nq, cs) + radius = int(dim * 0.42) + + quantizer = faiss.IndexBinaryFlat(dim) + index = faiss.IndexBinaryIVF(quantizer, dim, 4) + index.train(codes) + index.add(codes) + index.nprobe = 4 # every list, so the scan is exhaustive + + lims, dis, _ = index.range_search(queries, radius) + for q in range(nq): + got = sorted(dis[lims[q] : lims[q + 1]].tolist()) + truth = _hamming(queries[q : q + 1], codes) + self.assertEqual( + got, + sorted(truth[truth < radius].tolist()), + f"code_size={cs}, q={q}", + ) + + def test_store_pairs_offsets(self): + """store_pairs makes the batch path build ids with + lo_build(list_no, j + t) instead of reading them.""" + rs = np.random.RandomState(62) + n, k = 100, 8 + query = self.random_codes(rs, 1) + codes, ids = self.random_codes(rs, n), self.sequential_ids(n) + + distances, labels = self.run_scan( + self.make_scanner(rs, query, store_pairs=True), codes, ids, [n], k + ) + self.assertEqual(len(distances), k) + + self.assertEqual([faiss.lo_listno(x) for x in labels], [0] * k) + + # Each distance must be the distance to the code at the offset the + # same label decodes to. + truth = _hamming(query, codes) + offsets = [faiss.lo_offset(x) for x in labels] + self.assertEqual(distances, [int(truth[o]) for o in offsets]) + self.assertEqual(distances, _expected(query, codes, k)) diff --git a/tests/test_simd_dispatch.py b/tests/test_simd_dispatch.py index 16e0e4e4f..4ac6a649b 100644 --- a/tests/test_simd_dispatch.py +++ b/tests/test_simd_dispatch.py @@ -256,7 +256,7 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): """SPR detection must agree with the CPU's real feature flags. The SPR code path is compiled with -mavx512fp16, so AVX512_SPR must be reported available if and only if the CPU actually has the full - AVX512 core feature set, VNNI, VPOPCNTDQ, BF16, and FP16. + AVX512 core feature set, VNNI, VPOPCNTDQ, BITALG, BF16 and FP16. """ import platform @@ -298,6 +298,7 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): avx512_core <= flags and "avx512_vnni" in flags and "avx512_vpopcntdq" in flags + and "avx512_bitalg" in flags and "avx512_bf16" in flags and "avx512_fp16" in flags ) @@ -314,6 +315,7 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): f"(avx512_core={avx512_core <= flags}, " f"vnni={'avx512_vnni' in flags}, " f"vpopcntdq={'avx512_vpopcntdq' in flags}, " + f"bitalg={'avx512_bitalg' in flags}, " f"bf16={'avx512_bf16' in flags}, " f"fp16={'avx512_fp16' in flags}). " "detected=True with fp16=False is the D107684495 regression " @@ -321,7 +323,7 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): ) def test_vpopcnt_detection_matches_cpu_features(self): - """The VPOPCNT level requires baseline AVX-512 and VPOPCNTDQ.""" + """The VPOPCNT level requires baseline AVX-512, VPOPCNTDQ and BITALG.""" import platform try: @@ -358,7 +360,9 @@ for lvl in range(int(faiss.SIMDLevel_COUNT)): "avx512bw", } capable = ( - avx512_core <= flags and "avx512_vpopcntdq" in flags + avx512_core <= flags + and "avx512_vpopcntdq" in flags + and "avx512_bitalg" in flags ) detected = faiss.SIMDConfig.is_simd_level_available( faiss.SIMDLevel_AVX512_VPOPCNT