mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
rabitq neon kernels (#5532)
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
This commit is contained in:
lyang24
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
6c51305e64
commit
0db2f30886
@@ -9,15 +9,78 @@
|
||||
|
||||
#ifdef COMPILE_SIMD_ARM_NEON
|
||||
|
||||
#include <arm_neon.h>
|
||||
|
||||
#include <array>
|
||||
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
|
||||
namespace faiss::rabitq {
|
||||
|
||||
namespace {
|
||||
|
||||
// One accumulator per query bit plane. RaBitQuantizer caps qb at 8, so the
|
||||
// accumulator arrays below are fixed size rather than heap allocated.
|
||||
constexpr size_t kMaxQueryBits = 8;
|
||||
|
||||
inline uint32x4_t accumulate_popcount(uint32x4_t acc, uint8x16_t value) {
|
||||
const uint16x8_t pair_counts = vpaddlq_u8(vcntq_u8(value));
|
||||
return vpadalq_u16(acc, pair_counts);
|
||||
}
|
||||
|
||||
inline uint64_t reduce_popcount(uint32x4_t acc) {
|
||||
return vaddlvq_u32(acc);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <>
|
||||
uint64_t bitwise_and_dot_product<SIMDLevel::ARM_NEON>(
|
||||
const uint8_t* query,
|
||||
const uint8_t* data,
|
||||
size_t size,
|
||||
size_t qb) {
|
||||
return bitwise_and_dot_product<SIMDLevel::NONE>(query, data, size, qb);
|
||||
FAISS_THROW_IF_NOT(qb <= kMaxQueryBits);
|
||||
std::array<uint32x4_t, kMaxQueryBits> accumulators;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
accumulators[j] = vdupq_n_u32(0);
|
||||
}
|
||||
|
||||
size_t offset = 0;
|
||||
for (; offset + 16 <= size; offset += 16) {
|
||||
const uint8x16_t value = vld1q_u8(data + offset);
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
const uint8x16_t query_value = vld1q_u8(query + j * size + offset);
|
||||
accumulators[j] = accumulate_popcount(
|
||||
accumulators[j], vandq_u8(query_value, value));
|
||||
}
|
||||
}
|
||||
|
||||
uint64_t result = 0;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
result += reduce_popcount(accumulators[j]) << j;
|
||||
}
|
||||
for (; offset + 8 <= size; offset += 8) {
|
||||
uint64_t value;
|
||||
memcpy(&value, data + offset, sizeof(value));
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
uint64_t query_value;
|
||||
memcpy(&query_value,
|
||||
query + j * size + offset,
|
||||
sizeof(query_value));
|
||||
result += static_cast<uint64_t>(popcount64(query_value & value))
|
||||
<< j;
|
||||
}
|
||||
}
|
||||
for (; offset < size; offset++) {
|
||||
const uint8_t value = data[offset];
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
result += static_cast<uint64_t>(
|
||||
popcount32(query[j * size + offset] & value))
|
||||
<< j;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <>
|
||||
@@ -27,8 +90,52 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
|
||||
const uint8_t* data,
|
||||
size_t size,
|
||||
size_t qb) {
|
||||
return bitwise_and_dot_product_with_popcount<SIMDLevel::NONE>(
|
||||
query, data, size, qb);
|
||||
FAISS_THROW_IF_NOT(qb <= kMaxQueryBits);
|
||||
std::array<uint32x4_t, kMaxQueryBits> dot_accumulators;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
dot_accumulators[j] = vdupq_n_u32(0);
|
||||
}
|
||||
uint32x4_t popcount_accumulator = vdupq_n_u32(0);
|
||||
|
||||
size_t offset = 0;
|
||||
for (; offset + 16 <= size; offset += 16) {
|
||||
const uint8x16_t value = vld1q_u8(data + offset);
|
||||
popcount_accumulator = accumulate_popcount(popcount_accumulator, value);
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
const uint8x16_t query_value = vld1q_u8(query + j * size + offset);
|
||||
dot_accumulators[j] = accumulate_popcount(
|
||||
dot_accumulators[j], vandq_u8(query_value, value));
|
||||
}
|
||||
}
|
||||
|
||||
uint64_t dot_product = 0;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
dot_product += reduce_popcount(dot_accumulators[j]) << j;
|
||||
}
|
||||
uint64_t popcount_sum = reduce_popcount(popcount_accumulator);
|
||||
for (; offset + 8 <= size; offset += 8) {
|
||||
uint64_t value;
|
||||
memcpy(&value, data + offset, sizeof(value));
|
||||
popcount_sum += popcount64(value);
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
uint64_t query_value;
|
||||
memcpy(&query_value,
|
||||
query + j * size + offset,
|
||||
sizeof(query_value));
|
||||
dot_product +=
|
||||
static_cast<uint64_t>(popcount64(query_value & value)) << j;
|
||||
}
|
||||
}
|
||||
for (; offset < size; offset++) {
|
||||
const uint8_t value = data[offset];
|
||||
popcount_sum += popcount32(value);
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
dot_product += static_cast<uint64_t>(
|
||||
popcount32(query[j * size + offset] & value))
|
||||
<< j;
|
||||
}
|
||||
}
|
||||
return {dot_product, popcount_sum};
|
||||
}
|
||||
|
||||
template <>
|
||||
@@ -37,18 +144,185 @@ uint64_t bitwise_xor_dot_product<SIMDLevel::ARM_NEON>(
|
||||
const uint8_t* data,
|
||||
size_t size,
|
||||
size_t qb) {
|
||||
return bitwise_xor_dot_product<SIMDLevel::NONE>(query, data, size, qb);
|
||||
FAISS_THROW_IF_NOT(qb <= kMaxQueryBits);
|
||||
std::array<uint32x4_t, kMaxQueryBits> accumulators;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
accumulators[j] = vdupq_n_u32(0);
|
||||
}
|
||||
|
||||
size_t offset = 0;
|
||||
for (; offset + 16 <= size; offset += 16) {
|
||||
const uint8x16_t value = vld1q_u8(data + offset);
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
const uint8x16_t query_value = vld1q_u8(query + j * size + offset);
|
||||
accumulators[j] = accumulate_popcount(
|
||||
accumulators[j], veorq_u8(query_value, value));
|
||||
}
|
||||
}
|
||||
|
||||
uint64_t result = 0;
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
result += reduce_popcount(accumulators[j]) << j;
|
||||
}
|
||||
for (; offset + 8 <= size; offset += 8) {
|
||||
uint64_t value;
|
||||
memcpy(&value, data + offset, sizeof(value));
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
uint64_t query_value;
|
||||
memcpy(&query_value,
|
||||
query + j * size + offset,
|
||||
sizeof(query_value));
|
||||
result += static_cast<uint64_t>(popcount64(query_value ^ value))
|
||||
<< j;
|
||||
}
|
||||
}
|
||||
for (; offset < size; offset++) {
|
||||
const uint8_t value = data[offset];
|
||||
for (size_t j = 0; j < qb; j++) {
|
||||
result += static_cast<uint64_t>(
|
||||
popcount32(query[j * size + offset] ^ value))
|
||||
<< j;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <>
|
||||
uint64_t popcount<SIMDLevel::ARM_NEON>(const uint8_t* data, size_t size) {
|
||||
return popcount<SIMDLevel::NONE>(data, size);
|
||||
uint32x4_t accumulator = vdupq_n_u32(0);
|
||||
size_t offset = 0;
|
||||
for (; offset + 16 <= size; offset += 16) {
|
||||
accumulator = accumulate_popcount(accumulator, vld1q_u8(data + offset));
|
||||
}
|
||||
|
||||
uint64_t result = reduce_popcount(accumulator);
|
||||
for (; offset + 8 <= size; offset += 8) {
|
||||
uint64_t value;
|
||||
memcpy(&value, data + offset, sizeof(value));
|
||||
result += popcount64(value);
|
||||
}
|
||||
for (; offset < size; offset++) {
|
||||
result += popcount32(data[offset]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace faiss::rabitq
|
||||
|
||||
namespace faiss::rabitq::multibit {
|
||||
|
||||
namespace {
|
||||
|
||||
inline float horizontal_sum(float32x4_t low, float32x4_t high) {
|
||||
return vaddvq_f32(vaddq_f32(low, high));
|
||||
}
|
||||
|
||||
inline float ip_1exbit_neon(
|
||||
const uint8_t* __restrict sign_bits,
|
||||
const uint8_t* __restrict ex_code,
|
||||
const float* __restrict rotated_q,
|
||||
size_t d,
|
||||
float cb) {
|
||||
const uint32x4_t low_bit_positions = {1, 2, 4, 8};
|
||||
const uint32x4_t high_bit_positions = {16, 32, 64, 128};
|
||||
const uint32x4_t sign_weight = vdupq_n_u32(2);
|
||||
const uint32x4_t ex_weight = vdupq_n_u32(1);
|
||||
const float32x4_t bias = vdupq_n_f32(cb);
|
||||
float32x4_t low_accumulator = vdupq_n_f32(0.0f);
|
||||
float32x4_t high_accumulator = vdupq_n_f32(0.0f);
|
||||
|
||||
size_t i = 0;
|
||||
for (; i + 8 <= d; i += 8) {
|
||||
const uint32x4_t sign = vdupq_n_u32(sign_bits[i / 8]);
|
||||
const uint32x4_t extra = vdupq_n_u32(ex_code[i / 8]);
|
||||
|
||||
const uint32x4_t low_values = vaddq_u32(
|
||||
vandq_u32(vtstq_u32(sign, low_bit_positions), sign_weight),
|
||||
vandq_u32(vtstq_u32(extra, low_bit_positions), ex_weight));
|
||||
const uint32x4_t high_values = vaddq_u32(
|
||||
vandq_u32(vtstq_u32(sign, high_bit_positions), sign_weight),
|
||||
vandq_u32(vtstq_u32(extra, high_bit_positions), ex_weight));
|
||||
|
||||
const float32x4_t low_reconstruction =
|
||||
vaddq_f32(vcvtq_f32_u32(low_values), bias);
|
||||
const float32x4_t high_reconstruction =
|
||||
vaddq_f32(vcvtq_f32_u32(high_values), bias);
|
||||
low_accumulator = vfmaq_f32(
|
||||
low_accumulator, vld1q_f32(rotated_q + i), low_reconstruction);
|
||||
high_accumulator = vfmaq_f32(
|
||||
high_accumulator,
|
||||
vld1q_f32(rotated_q + i + 4),
|
||||
high_reconstruction);
|
||||
}
|
||||
|
||||
float result = horizontal_sum(low_accumulator, high_accumulator);
|
||||
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
|
||||
return result;
|
||||
}
|
||||
|
||||
template <size_t ExBits>
|
||||
inline float ip_multibit_neon(
|
||||
const uint8_t* __restrict sign_bits,
|
||||
const uint8_t* __restrict ex_code,
|
||||
const float* __restrict rotated_q,
|
||||
size_t d,
|
||||
float cb) {
|
||||
static_assert(ExBits >= 2 && ExBits <= 7);
|
||||
const uint32x4_t low_bit_positions = {1, 2, 4, 8};
|
||||
const uint32x4_t high_bit_positions = {16, 32, 64, 128};
|
||||
const uint32x4_t sign_weight =
|
||||
vdupq_n_u32(static_cast<uint32_t>(1u << ExBits));
|
||||
const uint64x2_t code_mask =
|
||||
vdupq_n_u64(static_cast<uint64_t>((1u << ExBits) - 1));
|
||||
constexpr int64_t bit_width = static_cast<int64_t>(ExBits);
|
||||
const int64x2_t shifts_01 = {0, -bit_width};
|
||||
const int64x2_t shifts_23 = {-2 * bit_width, -3 * bit_width};
|
||||
const int64x2_t shifts_45 = {-4 * bit_width, -5 * bit_width};
|
||||
const int64x2_t shifts_67 = {-6 * bit_width, -7 * bit_width};
|
||||
const float32x4_t bias = vdupq_n_f32(cb);
|
||||
float32x4_t low_accumulator = vdupq_n_f32(0.0f);
|
||||
float32x4_t high_accumulator = vdupq_n_f32(0.0f);
|
||||
|
||||
size_t i = 0;
|
||||
for (; i + 8 <= d; i += 8) {
|
||||
uint64_t packed_codes = 0;
|
||||
memcpy(&packed_codes, ex_code + (i / 8) * ExBits, ExBits);
|
||||
const uint64x2_t packed = vdupq_n_u64(packed_codes);
|
||||
|
||||
const uint32x4_t low_codes = vcombine_u32(
|
||||
vmovn_u64(vandq_u64(vshlq_u64(packed, shifts_01), code_mask)),
|
||||
vmovn_u64(vandq_u64(vshlq_u64(packed, shifts_23), code_mask)));
|
||||
const uint32x4_t high_codes = vcombine_u32(
|
||||
vmovn_u64(vandq_u64(vshlq_u64(packed, shifts_45), code_mask)),
|
||||
vmovn_u64(vandq_u64(vshlq_u64(packed, shifts_67), code_mask)));
|
||||
|
||||
const uint32x4_t sign = vdupq_n_u32(sign_bits[i / 8]);
|
||||
const uint32x4_t low_values = vaddq_u32(
|
||||
low_codes,
|
||||
vandq_u32(vtstq_u32(sign, low_bit_positions), sign_weight));
|
||||
const uint32x4_t high_values = vaddq_u32(
|
||||
high_codes,
|
||||
vandq_u32(vtstq_u32(sign, high_bit_positions), sign_weight));
|
||||
|
||||
const float32x4_t low_reconstruction =
|
||||
vaddq_f32(vcvtq_f32_u32(low_values), bias);
|
||||
const float32x4_t high_reconstruction =
|
||||
vaddq_f32(vcvtq_f32_u32(high_values), bias);
|
||||
low_accumulator = vfmaq_f32(
|
||||
low_accumulator, vld1q_f32(rotated_q + i), low_reconstruction);
|
||||
high_accumulator = vfmaq_f32(
|
||||
high_accumulator,
|
||||
vld1q_f32(rotated_q + i + 4),
|
||||
high_reconstruction);
|
||||
}
|
||||
|
||||
float result = horizontal_sum(low_accumulator, high_accumulator);
|
||||
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ExBits, cb);
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <>
|
||||
float compute_inner_product<SIMDLevel::ARM_NEON>(
|
||||
const uint8_t* __restrict sign_bits,
|
||||
@@ -57,8 +331,25 @@ float compute_inner_product<SIMDLevel::ARM_NEON>(
|
||||
size_t d,
|
||||
size_t ex_bits,
|
||||
float cb) {
|
||||
return compute_inner_product<SIMDLevel::NONE>(
|
||||
sign_bits, ex_code, rotated_q, d, ex_bits, cb);
|
||||
if (ex_bits == 1) {
|
||||
return ip_1exbit_neon(sign_bits, ex_code, rotated_q, d, cb);
|
||||
}
|
||||
switch (ex_bits) {
|
||||
case 2:
|
||||
return ip_multibit_neon<2>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
case 3:
|
||||
return ip_multibit_neon<3>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
case 4:
|
||||
return ip_multibit_neon<4>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
case 5:
|
||||
return ip_multibit_neon<5>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
case 6:
|
||||
return ip_multibit_neon<6>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
case 7:
|
||||
return ip_multibit_neon<7>(sign_bits, ex_code, rotated_q, d, cb);
|
||||
default:
|
||||
return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace faiss::rabitq::multibit
|
||||
|
||||
@@ -94,6 +94,10 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
|
||||
PROPERTIES COMPILE_OPTIONS "/arch:AVX512"
|
||||
)
|
||||
endif()
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
|
||||
list(APPEND FAISS_TEST_SRC
|
||||
test_rabitq_simd_neon.cpp
|
||||
)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
|
||||
@@ -31,33 +31,8 @@ static std::vector<uint8_t> random_codes(size_t d, size_t qb, uint32_t seed) {
|
||||
return q;
|
||||
}
|
||||
|
||||
static 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;
|
||||
}
|
||||
|
||||
// 32-d chunks and chunk boundaries.
|
||||
static const std::vector<size_t> kDims = {
|
||||
1,
|
||||
8,
|
||||
16,
|
||||
31,
|
||||
32,
|
||||
33,
|
||||
100,
|
||||
128,
|
||||
255,
|
||||
256,
|
||||
257,
|
||||
384,
|
||||
512,
|
||||
768,
|
||||
1024,
|
||||
2048};
|
||||
using faiss_test::kDims;
|
||||
using faiss_test::random_bytes;
|
||||
|
||||
// Note: scalar kernel's own correctness is covered end-to-end by
|
||||
// tests/test_rabitq.py. This target is x86-only (see BUCK).
|
||||
|
||||
@@ -29,37 +29,8 @@
|
||||
|
||||
using faiss::SIMDLevel;
|
||||
|
||||
namespace {
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
// Widths on either side of every register boundary the kernels step over.
|
||||
const std::vector<size_t> kDims = {
|
||||
1,
|
||||
8,
|
||||
16,
|
||||
31,
|
||||
32,
|
||||
33,
|
||||
100,
|
||||
128,
|
||||
255,
|
||||
256,
|
||||
257,
|
||||
384,
|
||||
512,
|
||||
768,
|
||||
1024,
|
||||
2048};
|
||||
|
||||
} // namespace
|
||||
using faiss_test::kDims;
|
||||
using faiss_test::random_bytes;
|
||||
|
||||
TEST(RaBitQVpopcnt, BitwiseKernelsMatchScalarAcrossTails) {
|
||||
if (!faiss::SIMDConfig::is_simd_level_available(
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/utils/rabitq_simd.h>
|
||||
#include <faiss/utils/simd_levels.h>
|
||||
|
||||
#include "test_rabitq_simd_util.h"
|
||||
|
||||
using faiss::SIMDLevel;
|
||||
using faiss_test::kDims;
|
||||
using faiss_test::random_bytes;
|
||||
|
||||
// This target is aarch64-only (see BUCK): it calls the ARM_NEON
|
||||
// specializations directly, and no other architecture compiles them.
|
||||
TEST(RaBitQNeon, BitwiseKernelsMatchScalar) {
|
||||
if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::ARM_NEON)) {
|
||||
GTEST_SKIP() << "ARM_NEON is not available on this CPU";
|
||||
}
|
||||
|
||||
for (size_t d : kDims) {
|
||||
const size_t size = (d + 7) / 8;
|
||||
const auto data = random_bytes(size, 19717 + d);
|
||||
for (size_t qb = 1; qb <= 8; qb++) {
|
||||
const auto query = random_bytes(size * qb, 51691 + d + qb);
|
||||
|
||||
const uint64_t expected_and =
|
||||
faiss::rabitq::bitwise_and_dot_product<SIMDLevel::NONE>(
|
||||
query.data(), data.data(), size, qb);
|
||||
const uint64_t expected_xor =
|
||||
faiss::rabitq::bitwise_xor_dot_product<SIMDLevel::NONE>(
|
||||
query.data(), data.data(), size, qb);
|
||||
const uint64_t expected_popcount =
|
||||
faiss::rabitq::popcount<SIMDLevel::NONE>(data.data(), size);
|
||||
|
||||
const uint64_t actual_and =
|
||||
faiss::rabitq::bitwise_and_dot_product<SIMDLevel::ARM_NEON>(
|
||||
query.data(), data.data(), size, qb);
|
||||
const uint64_t actual_xor =
|
||||
faiss::rabitq::bitwise_xor_dot_product<SIMDLevel::ARM_NEON>(
|
||||
query.data(), data.data(), size, qb);
|
||||
const uint64_t actual_popcount =
|
||||
faiss::rabitq::popcount<SIMDLevel::ARM_NEON>(
|
||||
data.data(), size);
|
||||
const auto actual_fused =
|
||||
faiss::rabitq::bitwise_and_dot_product_with_popcount<
|
||||
SIMDLevel::ARM_NEON>(
|
||||
query.data(), data.data(), size, qb);
|
||||
|
||||
EXPECT_EQ(actual_and, expected_and) << "d=" << d << " qb=" << qb;
|
||||
EXPECT_EQ(actual_xor, expected_xor) << "d=" << d << " qb=" << qb;
|
||||
EXPECT_EQ(actual_popcount, expected_popcount)
|
||||
<< "d=" << d << " qb=" << qb;
|
||||
EXPECT_EQ(actual_fused.dot_product, expected_and)
|
||||
<< "d=" << d << " qb=" << qb;
|
||||
EXPECT_EQ(actual_fused.popcount, expected_popcount)
|
||||
<< "d=" << d << " qb=" << qb;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(RaBitQNeon, MultiBitInnerProductMatchesScalar) {
|
||||
if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::ARM_NEON)) {
|
||||
GTEST_SKIP() << "ARM_NEON is not available on this CPU";
|
||||
}
|
||||
|
||||
std::mt19937 rng(918273);
|
||||
std::uniform_real_distribution<float> values(-2.0f, 2.0f);
|
||||
|
||||
for (size_t d : kDims) {
|
||||
const auto sign_bits = random_bytes((d + 7) / 8, 13579 + d);
|
||||
std::vector<float> query(d);
|
||||
for (float& value : query) {
|
||||
value = values(rng);
|
||||
}
|
||||
for (size_t ex_bits = 1; ex_bits <= 7; ex_bits++) {
|
||||
// ip_scalar reads an 8-byte extraction window, so pad the code to
|
||||
// match the production layout.
|
||||
const auto extra_bits = random_bytes(
|
||||
(d * ex_bits + 7) / 8 + 7, 24680 + d + ex_bits);
|
||||
const float cb = -static_cast<float>((1u << ex_bits) - 0.5f);
|
||||
const float expected =
|
||||
faiss::rabitq::multibit::compute_inner_product<
|
||||
SIMDLevel::NONE>(
|
||||
sign_bits.data(),
|
||||
extra_bits.data(),
|
||||
query.data(),
|
||||
d,
|
||||
ex_bits,
|
||||
cb);
|
||||
const float actual = faiss::rabitq::multibit::compute_inner_product<
|
||||
SIMDLevel::ARM_NEON>(
|
||||
sign_bits.data(),
|
||||
extra_bits.data(),
|
||||
query.data(),
|
||||
d,
|
||||
ex_bits,
|
||||
cb);
|
||||
|
||||
// The NEON kernel accumulates in eight lanes, so it does not
|
||||
// reproduce the scalar sum bit for bit.
|
||||
const float tolerance = 2e-5f * std::max(1.0f, std::abs(expected));
|
||||
EXPECT_NEAR(actual, expected, tolerance)
|
||||
<< "d=" << d << " ex_bits=" << ex_bits;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -5,9 +5,11 @@
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
// Shared by test_rabitq_simd.cpp and test_rabitq_simd_dd.cpp. The AVX512
|
||||
// kernels live in a dynamic dispatch build only, so the tests that call them
|
||||
// sit in the second file, but both compare against the same scalar reference.
|
||||
// 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
|
||||
|
||||
@@ -15,6 +17,7 @@
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/utils/rabitq_simd.h>
|
||||
@@ -24,6 +27,34 @@ 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;
|
||||
|
||||
Reference in New Issue
Block a user