mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Measure 20-byte Hamming codes eight at a time (#5587)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5587 **TL;DR:** Adds a batched 20-byte Hamming kernel that measures eight codes per call, the only thing that helps the 160-bit width, worth 1.21x end to end. A 20-byte code has no vector kernel at any x86 SIMD level. `HammingComputer20` measures one distance with three scalar popcounts at every level, including `AVX512_VPOPCNT`, so a scan of 160-bit codes runs the same code everywhere. Vectorizing a single 20-byte distance does not help. A masked 256-bit load plus `vpopcntq` plus a horizontal reduction measures the same as the scalar body, because the reduction costs more than the three popcounts it replaces. Measuring several codes per call does help, because it amortizes that reduction. This adds `hamming_batch()` to `HammingComputer20` at `AVX512_VPOPCNT`, which measures eight codes at once: - Eight codes span 160 bytes. `_mm512_popcnt_epi8` gives a count per byte. - `_mm512_sad_epu8` sums each aligned 8-byte group. - A 20-byte code covers two whole groups plus half of a group it shares with its neighbour, so each code needs one scalar popcount for its half. `IVFBinaryScannerL2::scan_codes` and `IVFBinaryScannerL2::scan_codes_range` call `hamming_batch()` when the computer declares a `batch_size`, then measure the remaining codes one at a time. A computer without a `batch_size` keeps the existing loop, so no other code size changes. The scanner holds the tiled query, and the constructor rejects a code size that does not match the computer. A caller therefore cannot reach the batch kernel with a stride the kernel does not read. The byte-wise popcount needs `AVX512_BITALG`, which is a separate CPUID bit from `AVX512_VPOPCNTDQ`. This adds the flag to the `AVX512_VPOPCNT` level and requires both bits to select it. Every CPU that has VPOPCNTDQ also has BITALG, so no CPU loses the level. An alternative byte-wise popcount built from `vpshufb`, which needs neither bit and would run at the plain `AVX512` level, measures only 1.08x against 1.54x, so it is not worth the wider reach. Reviewed By: alibeklfc Differential Revision: D118971285 fbshipit-source-id: 5fd4637c972e5c55e2695b32b9b15fba8181a242
This commit is contained in:
Michael Norris
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
3c7042f17a
commit
6bbb068ad8
@@ -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 $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16>)
|
||||
target_compile_options(faiss_c_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mavx512bitalg -mpopcnt -mavx512fp16 -mavx512bf16>)
|
||||
else()
|
||||
target_compile_options(faiss_c_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
|
||||
endif()
|
||||
|
||||
@@ -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 $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16>)
|
||||
target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mavx512bitalg -mpopcnt -mavx512fp16 -mavx512bf16>)
|
||||
else()
|
||||
target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
|
||||
endif()
|
||||
|
||||
@@ -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 $<$<COMPILE_LANGUAGE:CXX>:-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 $<$<COMPILE_LANGUAGE:CXX>:-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 $<$<COMPILE_LANGUAGE:CXX>:/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"
|
||||
|
||||
@@ -39,17 +39,50 @@ namespace faiss {
|
||||
|
||||
namespace {
|
||||
|
||||
// Whether the computer can measure a whole batch of codes in one call.
|
||||
template <class HammingComputer>
|
||||
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 <class HammingComputer>
|
||||
struct BatchQueryTile {};
|
||||
|
||||
template <class HammingComputer>
|
||||
requires has_hamming_batch<HammingComputer>
|
||||
struct BatchQueryTile<HammingComputer> {
|
||||
alignas(64) uint8_t batch_query
|
||||
[HammingComputer::batch_size * HammingComputer::get_code_size()];
|
||||
};
|
||||
|
||||
template <class HammingComputer>
|
||||
struct IVFBinaryScannerL2 : BinaryInvertedListScanner,
|
||||
BatchQueryTile<HammingComputer> {
|
||||
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<HammingComputer>) {
|
||||
// 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>) {
|
||||
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 <class Accept>
|
||||
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<HammingComputer>) {
|
||||
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<uint32_t>(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<uint32_t>(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<C>(k, simi, idxi, dis, id);
|
||||
bound = static_cast<uint32_t>(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<C>(k, simi, idxi, dis, id);
|
||||
bound = static_cast<uint32_t>(simi[0]);
|
||||
nup++;
|
||||
accept(static_cast<int32_t>(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<uint32_t>(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<uint32_t>(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<uint32_t>(radius)) {
|
||||
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
|
||||
result.add(dis, id);
|
||||
if (dis < bound) {
|
||||
accept(static_cast<int32_t>(dis), j);
|
||||
}
|
||||
codes += code_size;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -17,4 +17,8 @@
|
||||
#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
|
||||
// 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<faiss::SIMDLevel::AVX512_VPOPCNT>>);
|
||||
|
||||
#endif // COMPILE_SIMD_AVX512_VPOPCNT
|
||||
|
||||
@@ -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 <cassert>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
#include <faiss/utils/popcount.h>
|
||||
|
||||
#include <faiss/impl/platform_macros.h>
|
||||
#include <faiss/utils/hamming_distance/hamming_computer-avx512.h>
|
||||
@@ -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<SIMDLevel::AVX512_VPOPCNT>
|
||||
: HammingComputer20_tpl<SIMDLevel::AVX512> {
|
||||
using HammingComputer20_tpl<SIMDLevel::AVX512>::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<int32_t>(
|
||||
grp[g] + grp[g + 1] + popcount32(xh ^ qh));
|
||||
} else {
|
||||
memcpy(&xh, codes + s, 4);
|
||||
memcpy(&qh, tile + s, 4);
|
||||
dis[k] = static_cast<int32_t>(
|
||||
popcount32(xh ^ qh) + grp[g + 1] + grp[g + 2]);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct HammingComputer32_tpl<SIMDLevel::AVX512_VPOPCNT> {
|
||||
const uint8_t* a8;
|
||||
|
||||
@@ -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<int>(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<int>(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<int>(SIMDLevel::AVX512_SPR));
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user