mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
lyang24
0db2f30886
Summary: This diff adds native ARM NEON kernels for the RaBitQ bitwise popcount operations and for the multi-bit inner product. The kernels replace the scalar fallbacks on ARM. They handle every supported bit width and every unaligned dimension. - `tests/test_rabitq_simd_neon.cpp` holds the correctness test, and `tests/BUCK` holds its `cpp_unittest` target. The target carries `target_compatible_with` on `ovr_config//cpu:arm64`, which mirrors the x86_64 gate on `test_rabitq_simd`. - `tests/test_rabitq_simd_util.h` holds the shared `kDims` list and the shared `random_bytes` helper. `test_rabitq_simd.cpp` and `test_rabitq_simd_dd.cpp` now use them too. - The three kernels that index a fixed array of eight accumulators assert `qb <= kMaxQueryBits`. `RaBitQuantizer::set_query` caps `qb` at 8, but the kernels are public templates and carry no other bound. The scalar kernel and the AVX-512 kernel need no such assertion, because they shift inside the loop and hold one accumulator. The multi-bit test compares with a relative tolerance of `2e-5`. The NEON kernel accumulates in eight lanes, so it does not reproduce the scalar sum bit for bit. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5532 Test Plan: `buck2 test fbcode//faiss/tests:test_ivf_index fbcode//faiss/tests:test_custom_result_handler fbcode//faiss/tests:test_hnsw fbcode//faiss/tests:test_rabitq_simd_neon fbcode//faiss/tests:test_params_override fbcode//faiss/tests:test_hnsw_rabitq fbcode//faiss/tests:test_rabitq fbcode//faiss/tests:test_search_params fbcode//faiss/tests:test_io_corrupted fbcode//faiss/tests:test_rabitq_utils`, run at the top of the stack on aarch64: 173 pass, 0 fail, 3 skip. `buck2 build fbcode//faiss:faiss` succeeds in the static aarch64 build and with `-c faiss.dynamic_dispatch=true`. The host holds Neoverse V2 cores with ASIMDDP, so `test_rabitq_simd_neon` runs its kernels natively rather than skipping. Reviewed By: alibeklfc Differential Revision: D121846717 Pulled By: mnorris11 fbshipit-source-id: de079613a52ced326045830fb799dfc9b169fe2e
125 lines
3.5 KiB
C++
125 lines
3.5 KiB
C++
/*
|
|
* 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.
|
|
*/
|
|
|
|
// Shared by test_rabitq_simd.cpp, test_rabitq_simd_dd.cpp and
|
|
// test_rabitq_simd_neon.cpp. The AVX512 kernels live in a dynamic dispatch
|
|
// build only, and the NEON kernels live in an aarch64 build only, so each set
|
|
// of tests sits in its own file. They all compare against the same scalar
|
|
// reference.
|
|
|
|
#pragma once
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <random>
|
|
#include <vector>
|
|
|
|
#include <faiss/utils/rabitq_simd.h>
|
|
#include <faiss/utils/simd_levels.h>
|
|
|
|
namespace faiss_test {
|
|
|
|
using faiss::SIMDLevel;
|
|
|
|
// Widths on either side of every register boundary the kernels step over.
|
|
inline const std::vector<size_t> kDims = {
|
|
1,
|
|
8,
|
|
16,
|
|
31,
|
|
32,
|
|
33,
|
|
100,
|
|
128,
|
|
255,
|
|
256,
|
|
257,
|
|
384,
|
|
512,
|
|
768,
|
|
1024,
|
|
2048};
|
|
|
|
inline std::vector<uint8_t> random_bytes(size_t n, uint32_t seed) {
|
|
std::mt19937 rng(seed);
|
|
std::vector<uint8_t> v(n);
|
|
for (size_t i = 0; i < n; i++) {
|
|
v[i] = static_cast<uint8_t>(rng());
|
|
}
|
|
return v;
|
|
}
|
|
|
|
template <SIMDLevel SL>
|
|
inline void check_quantization_matches_scalar() {
|
|
constexpr size_t d = 257;
|
|
constexpr float v_min = -2.0f;
|
|
constexpr float inv_delta = 16.0f;
|
|
|
|
for (uint8_t max_code : {uint8_t(1), uint8_t(15), uint8_t(255)}) {
|
|
std::vector<float> values(d);
|
|
for (size_t i = 0; i < d; i++) {
|
|
const uint8_t code = i % (size_t(max_code) + 1);
|
|
values[i] = v_min + code / inv_delta;
|
|
}
|
|
|
|
for (bool centered : {false, true}) {
|
|
std::vector<uint8_t> scalar_codes(d);
|
|
size_t scalar_sum = 0;
|
|
int64_t scalar_sum2 = 0;
|
|
faiss::rabitq::quantize_query_values<SIMDLevel::NONE>(
|
|
values.data(),
|
|
d,
|
|
v_min,
|
|
inv_delta,
|
|
max_code,
|
|
centered,
|
|
scalar_codes.data(),
|
|
scalar_sum,
|
|
scalar_sum2);
|
|
|
|
std::vector<uint8_t> simd_codes(d);
|
|
size_t simd_sum = 0;
|
|
int64_t simd_sum2 = 0;
|
|
faiss::rabitq::quantize_query_values<SL>(
|
|
values.data(),
|
|
d,
|
|
v_min,
|
|
inv_delta,
|
|
max_code,
|
|
centered,
|
|
simd_codes.data(),
|
|
simd_sum,
|
|
simd_sum2);
|
|
|
|
EXPECT_EQ(simd_codes, scalar_codes)
|
|
<< "max_code=" << int(max_code) << " centered=" << centered;
|
|
EXPECT_EQ(simd_sum, scalar_sum);
|
|
EXPECT_EQ(simd_sum2, scalar_sum2);
|
|
}
|
|
}
|
|
|
|
constexpr float lut_min = -2.0f;
|
|
constexpr float lut_scale = 4.0f;
|
|
float lut[16];
|
|
for (size_t i = 0; i < 16; i++) {
|
|
lut[i] = lut_min + i / lut_scale;
|
|
}
|
|
uint8_t scalar_lut[16];
|
|
uint8_t simd_lut[16];
|
|
faiss::rabitq::lut_quantize_16_to_uint8<SIMDLevel::NONE>(
|
|
lut, lut_min, lut_scale, scalar_lut);
|
|
faiss::rabitq::lut_quantize_16_to_uint8<SL>(
|
|
lut, lut_min, lut_scale, simd_lut);
|
|
EXPECT_EQ(
|
|
std::vector<uint8_t>(simd_lut, simd_lut + 16),
|
|
std::vector<uint8_t>(scalar_lut, scalar_lut + 16));
|
|
}
|
|
|
|
} // namespace faiss_test
|