mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Let the binary Hamming paths reach the VPOPCNT kernels
Summary: **TL;DR:** Points `IndexBinaryIVF` at the VPOPCNT level and instantiates its scanner in the translation unit compiled with `-mavx512vpopcntdq`, worth 1.58x at 256-bit and 1.41x at 512-bit on Genoa. Two things kept the binary Hamming paths on scalar popcount, whatever the host. `IndexBinaryIVF` builds its scanner and runs its searches through `with_simd_level`, which uses the BASE level mask. That mask holds no `AVX512_VPOPCNT` bit, so on a host that reports that level the dispatch falls through to `AVX512`, where the wide computers use scalar popcount. Two build variants, `faiss_omp_mock` and `faiss_no_multithreading`, hardcoded dynamic dispatch off. They compiled the AVX2 kernels only, so a consumer of either variant could not use runtime dispatch at all, whatever the build setting said. This change: - Moves the three dispatch sites in `IndexBinaryIVF.cpp` to `with_simd_level_with_vpopcnt`. Below that level the dispatch falls through to `AVX512`, as before. - Instantiates `IndexBinaryIVF_impl.h` in `hamming_avx512_vpopcnt.cpp`. That is the one translation unit compiled with `-mavx512vpopcntdq`, which the VPOPCNT computers need. Without it the new dispatch has no symbol to call. The include must follow the computer specializations, so it is placed after them. - Lets both build variants honour the dynamic dispatch setting. `HammingComputer16` and `HammingComputer20` still inherit the scalar body at this level, so a 16-byte or 20-byte code gains nothing here. `IVFBinaryScannerL2::scan_codes` also reads the heap top into a local `bound` instead of loading `simi[0]` on every code. Only `scan_codes` writes `simi`, so it refreshes `bound` after each `heap_replace_top` and the two forms return the same results. The test plan measures it. Reviewed By: alibeklfc Differential Revision: D118874433 fbshipit-source-id: dc74294ecaf178919fdd713abc60c44bbaa80347
This commit is contained in:
Michael Norris
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
e6bd1d252c
commit
3c7042f17a
@@ -348,10 +348,6 @@ void IndexBinaryIVF::replace_invlists(InvertedLists* il, bool own) {
|
||||
own_invlists = own;
|
||||
}
|
||||
|
||||
// IVFBinaryScannerL2, search_knn_hamming_count, BlockSearch,
|
||||
// BlockSearchVariableK, search_knn_hamming_per_invlist are now in
|
||||
// impl/binary_hamming/IndexBinaryIVF_impl.h (compiled per-ISA)
|
||||
|
||||
namespace {
|
||||
|
||||
void search_knn_hamming_heap(
|
||||
@@ -461,13 +457,9 @@ void search_knn_hamming_heap(
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// The remaining template code (search_knn_hamming_count,
|
||||
// search_knn_hamming_per_invlist, etc.) has been moved to
|
||||
// impl/binary_hamming/IndexBinaryIVF_impl.h
|
||||
|
||||
BinaryInvertedListScanner* IndexBinaryIVF::get_InvertedListScanner(
|
||||
bool store_pairs) const {
|
||||
return with_simd_level([&]<SIMDLevel SL>() {
|
||||
return with_simd_level_with_vpopcnt([&]<SIMDLevel SL>() {
|
||||
return make_binary_ivf_scanner_fixSL<SL>(code_size, store_pairs);
|
||||
});
|
||||
}
|
||||
@@ -483,7 +475,7 @@ void IndexBinaryIVF::search_preassigned(
|
||||
bool store_pairs,
|
||||
const IVFSearchParameters* params) const {
|
||||
if (per_invlist_search) {
|
||||
with_simd_level([&]<SIMDLevel SL>() {
|
||||
with_simd_level_with_vpopcnt([&]<SIMDLevel SL>() {
|
||||
search_knn_hamming_per_invlist_fixSL<SL>(
|
||||
code_size,
|
||||
this,
|
||||
@@ -501,7 +493,7 @@ void IndexBinaryIVF::search_preassigned(
|
||||
search_knn_hamming_heap(
|
||||
this, n, x, k, cidx, cdis, dis, idx, store_pairs, params);
|
||||
} else {
|
||||
with_simd_level([&]<SIMDLevel SL>() {
|
||||
with_simd_level_with_vpopcnt([&]<SIMDLevel SL>() {
|
||||
search_knn_hamming_count_fixSL<SL>(
|
||||
code_size,
|
||||
store_pairs,
|
||||
|
||||
@@ -232,7 +232,14 @@ struct BinaryInvertedListScanner {
|
||||
virtual uint32_t distance_to_code(const uint8_t* code) const = 0;
|
||||
|
||||
/** compute the distances to codes. (distances, labels) should be
|
||||
* organized as a min- or max-heap
|
||||
* organized as a max-heap: the scan accepts a code whose distance is
|
||||
* below distances[0] and replaces the top. A min-heap caller reads its
|
||||
* neutral value as a huge unsigned bound and gets wrong results.
|
||||
|
||||
*
|
||||
* The heap top is the only bound. To keep the k nearest codes inside a
|
||||
* radius, seed every heap slot with that radius instead of the neutral
|
||||
* value; an unfilled slot keeps its label of -1.
|
||||
*
|
||||
* @param n number of codes to scan
|
||||
* @param codes codes to scan (n * code_size)
|
||||
|
||||
@@ -70,12 +70,15 @@ struct IVFBinaryScannerL2 : BinaryInvertedListScanner {
|
||||
size_t k) const override {
|
||||
using C = CMax<int32_t, idx_t>;
|
||||
|
||||
uint32_t bound = static_cast<uint32_t>(simi[0]);
|
||||
|
||||
size_t nup = 0;
|
||||
for (size_t j = 0; j < n; j++) {
|
||||
uint32_t dis = hc.hamming(codes);
|
||||
if (dis < static_cast<uint32_t>(simi[0])) {
|
||||
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++;
|
||||
}
|
||||
codes += code_size;
|
||||
|
||||
@@ -8,8 +8,13 @@
|
||||
#ifdef COMPILE_SIMD_AVX512_VPOPCNT
|
||||
|
||||
#define THE_SIMD_LEVEL SIMDLevel::AVX512_VPOPCNT
|
||||
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
|
||||
#include <faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h>
|
||||
#include <faiss/utils/hamming_distance/hamming_impl.h>
|
||||
|
||||
// Must follow the computer specializations above.
|
||||
// clang-format off
|
||||
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
|
||||
#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
|
||||
// clang-format on
|
||||
|
||||
#endif // COMPILE_SIMD_AVX512_VPOPCNT
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""
|
||||
Tests for BinaryInvertedListScanner::scan_codes, the bulk entry point a
|
||||
caller uses to fold an inverted list into a heap it owns.
|
||||
|
||||
Each test compares against a NumPy popcount-xor reference, which is
|
||||
level-independent, so per-level dispatch drift surfaces as a failure. The
|
||||
@for_all_simd_levels decorator runs every case at each available level.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
|
||||
from common_faiss_tests import for_all_simd_levels
|
||||
|
||||
# with_HammingComputer dispatches on the code size: 4, 8, 16, 20, 32 and 64
|
||||
# each get their own computer, and anything else lands on the generic one.
|
||||
# scan_codes runs the same loop for all of them, so the semantic cases below
|
||||
# cover every computer rather than only the 20-byte one.
|
||||
WIDTHS = [4, 8, 16, 20, 24, 32, 64]
|
||||
CODE_SIZE = 20
|
||||
NEUTRAL = np.iinfo(np.int32).max
|
||||
|
||||
|
||||
def _hamming(query, codes):
|
||||
return np.unpackbits(codes ^ query, axis=1).sum(1).astype("int32")
|
||||
|
||||
|
||||
def _expected(query, codes, k, radius=None):
|
||||
dis = _hamming(query, codes)
|
||||
if radius is not None:
|
||||
dis = dis[dis < radius]
|
||||
return sorted(dis.tolist())[:k]
|
||||
|
||||
|
||||
@for_all_simd_levels
|
||||
class TestBinaryIVFScanner(unittest.TestCase):
|
||||
def make_scanner(self, rs, query, store_pairs=False, cs=CODE_SIZE):
|
||||
dim = cs * 8
|
||||
quantizer = faiss.IndexBinaryFlat(dim)
|
||||
quantizer.add(self.random_codes(rs, 1, cs))
|
||||
ivf = faiss.IndexBinaryIVF(quantizer, dim, 1)
|
||||
ivf.is_trained = True
|
||||
scanner = ivf.get_InvertedListScanner(store_pairs)
|
||||
scanner.set_query(faiss.swig_ptr(query))
|
||||
scanner.set_list(0, 0)
|
||||
# The index and quantizer must outlive the scanner.
|
||||
self.keep_alive = (ivf, quantizer)
|
||||
return scanner
|
||||
|
||||
def random_codes(self, rs, n, cs=CODE_SIZE):
|
||||
return rs.randint(0, 256, size=(n, cs)).astype("uint8")
|
||||
|
||||
def sequential_ids(self, n):
|
||||
return (np.arange(n) * 3 + 100).astype("int64")
|
||||
|
||||
def run_scan(self, scanner, codes, ids, slice_sizes, k, radius=None):
|
||||
"""Seeds the heap, scans each slice into it, then drops the slots the
|
||||
scan never filled. Seeding every slot with a radius makes the heap top
|
||||
reject any code at or beyond it."""
|
||||
simi = np.full(k, NEUTRAL if radius is None else radius, dtype="int32")
|
||||
idxi = np.full(k, -1, dtype="int64")
|
||||
offset = 0
|
||||
for count in slice_sizes:
|
||||
scanner.scan_codes(
|
||||
count,
|
||||
faiss.swig_ptr(codes[offset:]),
|
||||
faiss.swig_ptr(ids[offset:]),
|
||||
faiss.swig_ptr(simi),
|
||||
faiss.swig_ptr(idxi),
|
||||
k,
|
||||
)
|
||||
offset += count
|
||||
keep = idxi >= 0
|
||||
order = np.argsort(simi[keep], kind="stable")
|
||||
return simi[keep][order].tolist(), idxi[keep][order].tolist()
|
||||
|
||||
def test_unseeded_heap_keeps_the_k_nearest(self):
|
||||
"""Every code width, so the shared scan loop is covered for each
|
||||
HammingComputer specialization and for the generic one."""
|
||||
rs = np.random.RandomState(1234)
|
||||
n, k = 2000, 10
|
||||
for cs in WIDTHS:
|
||||
query = self.random_codes(rs, 1, cs)
|
||||
codes = self.random_codes(rs, n, cs)
|
||||
ids = self.sequential_ids(n)
|
||||
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query, cs=cs), codes, ids, [n], k
|
||||
)
|
||||
self.assertEqual(got, _expected(query, codes, k), f"code_size={cs}")
|
||||
|
||||
def test_seeded_heap_bounds_every_width(self):
|
||||
"""The radius seed must behave the same for every computer."""
|
||||
rs = np.random.RandomState(1235)
|
||||
n, k = 1500, 12
|
||||
for cs in WIDTHS:
|
||||
query = self.random_codes(rs, 1, cs)
|
||||
codes = self.random_codes(rs, n, cs)
|
||||
ids = self.sequential_ids(n)
|
||||
radius = int(cs * 8 * 0.42) # well below the mean of cs*4
|
||||
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query, cs=cs),
|
||||
codes,
|
||||
ids,
|
||||
[n],
|
||||
k,
|
||||
radius,
|
||||
)
|
||||
self.assertEqual(
|
||||
got, _expected(query, codes, k, radius), f"code_size={cs}"
|
||||
)
|
||||
|
||||
def test_seeded_heap_applies_the_radius_and_k_together(self):
|
||||
rs = np.random.RandomState(21)
|
||||
n = 4000
|
||||
query = self.random_codes(rs, 1)
|
||||
codes, ids = self.random_codes(rs, n), self.sequential_ids(n)
|
||||
|
||||
for k, radius in [(64, 60), (5, 70)]:
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query), codes, ids, [n], k, radius
|
||||
)
|
||||
self.assertEqual(got, _expected(query, codes, k, radius), f"k={k}")
|
||||
|
||||
def test_results_accumulate_across_calls(self):
|
||||
rs = np.random.RandomState(7)
|
||||
n, k, radius = 3000, 16, 75
|
||||
query = self.random_codes(rs, 1)
|
||||
codes, ids = self.random_codes(rs, n), self.sequential_ids(n)
|
||||
|
||||
for layout in [[1, n - 1], [0, n], [1000, 0, 1000, 1, 999]]:
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query), codes, ids, layout, k, radius
|
||||
)
|
||||
self.assertEqual(
|
||||
got, _expected(query, codes, k, radius), str(layout)
|
||||
)
|
||||
|
||||
def test_the_seeded_bound_is_exclusive(self):
|
||||
rs = np.random.RandomState(31)
|
||||
# An all-zero query, so a code with p bits set sits at distance p.
|
||||
query = np.zeros((1, CODE_SIZE), dtype="uint8")
|
||||
codes = np.zeros((2, CODE_SIZE), dtype="uint8")
|
||||
codes[0][0], codes[1][0] = 0x0F, 0x1F # 4 and 5 bits set
|
||||
ids = self.sequential_ids(2)
|
||||
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query), codes, ids, [2], 10, 5
|
||||
)
|
||||
self.assertEqual(got, [4])
|
||||
|
||||
# The same rule at its limit.
|
||||
got, _ = self.run_scan(
|
||||
self.make_scanner(rs, query), codes, ids, [2], 10, 0
|
||||
)
|
||||
self.assertEqual(got, [])
|
||||
Reference in New Issue
Block a user