Initial commit
This commit is contained in:
yichuan520030910320
334
packages/leann-backend-hnsw/third_party/faiss/tests/test_distances_simd.cpp
vendored
Normal file
334
packages/leann-backend-hnsw/third_party/faiss/tests/test_distances_simd.cpp
vendored
Normal file
@@ -0,0 +1,334 @@
|
||||
/*
|
||||
* 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 <cstddef>
|
||||
#include <cstdint>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/utils/distances.h>
|
||||
|
||||
// reference implementations
|
||||
void fvec_inner_products_ny_ref(
|
||||
float* ip,
|
||||
const float* x,
|
||||
const float* y,
|
||||
size_t d,
|
||||
size_t ny) {
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
ip[i] = faiss::fvec_inner_product(x, y, d);
|
||||
y += d;
|
||||
}
|
||||
}
|
||||
|
||||
void fvec_L2sqr_ny_ref(
|
||||
float* dis,
|
||||
const float* x,
|
||||
const float* y,
|
||||
size_t d,
|
||||
size_t ny) {
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
dis[i] = faiss::fvec_L2sqr(x, y, d);
|
||||
y += d;
|
||||
}
|
||||
}
|
||||
|
||||
// test templated versions of fvec_L2sqr_ny
|
||||
TEST(TestFvecL2sqrNy, D2) {
|
||||
// we're using int values in order to get 100% accurate
|
||||
// results with floats.
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> u(0, 32);
|
||||
|
||||
for (const auto dim : {2, 4, 8, 12}) {
|
||||
std::vector<float> x(dim, 0);
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = u(rng);
|
||||
}
|
||||
|
||||
for (const auto nrows : {1, 2, 5, 10, 15, 20, 25}) {
|
||||
std::vector<float> y(nrows * dim);
|
||||
for (size_t i = 0; i < y.size(); i++) {
|
||||
y[i] = u(rng);
|
||||
}
|
||||
|
||||
std::vector<float> distances(nrows, 0);
|
||||
faiss::fvec_L2sqr_ny(
|
||||
distances.data(), x.data(), y.data(), dim, nrows);
|
||||
|
||||
std::vector<float> distances_ref(nrows, 0);
|
||||
fvec_L2sqr_ny_ref(
|
||||
distances_ref.data(), x.data(), y.data(), dim, nrows);
|
||||
|
||||
ASSERT_EQ(distances, distances_ref)
|
||||
<< "Mismatching results for dim = " << dim
|
||||
<< ", nrows = " << nrows;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fvec_inner_products_ny
|
||||
TEST(TestFvecInnerProductsNy, D2) {
|
||||
// we're using int values in order to get 100% accurate
|
||||
// results with floats.
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> u(0, 32);
|
||||
|
||||
for (const auto dim : {2, 4, 8, 12}) {
|
||||
std::vector<float> x(dim, 0);
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = u(rng);
|
||||
}
|
||||
|
||||
for (const auto nrows : {1, 2, 5, 10, 15, 20, 25}) {
|
||||
std::vector<float> y(nrows * dim);
|
||||
for (size_t i = 0; i < y.size(); i++) {
|
||||
y[i] = u(rng);
|
||||
}
|
||||
|
||||
std::vector<float> distances(nrows, 0);
|
||||
faiss::fvec_inner_products_ny(
|
||||
distances.data(), x.data(), y.data(), dim, nrows);
|
||||
|
||||
std::vector<float> distances_ref(nrows, 0);
|
||||
fvec_inner_products_ny_ref(
|
||||
distances_ref.data(), x.data(), y.data(), dim, nrows);
|
||||
|
||||
ASSERT_EQ(distances, distances_ref)
|
||||
<< "Mismatching results for dim = " << dim
|
||||
<< ", nrows = " << nrows;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecL2sqr, distances_L2_squared_y_transposed) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
// modulo 8 results - 16 is to repeat the loop in the function
|
||||
int ny = 11; // this value will hit all the codepaths
|
||||
for (const auto d : {1, 2, 3, 4, 5, 6, 7, 8, 16}) {
|
||||
// initialize inputs
|
||||
std::vector<float> x(d);
|
||||
float x_sqlen = 0;
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = uniform(rng);
|
||||
x_sqlen += x[i] * x[i];
|
||||
}
|
||||
std::vector<float> y(d * ny);
|
||||
std::vector<float> y_sqlens(ny, 0);
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
for (size_t j = 0; j < y.size(); j++) {
|
||||
y[j] = uniform(rng);
|
||||
y_sqlens[i] += y[j] * y[j];
|
||||
}
|
||||
}
|
||||
|
||||
// perform function
|
||||
std::vector<float> true_distances(ny, 0);
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (size_t j = 0; j < d; j++) {
|
||||
dp += x[j] * y[i + j * ny];
|
||||
}
|
||||
true_distances[i] = x_sqlen + y_sqlens[i] - 2 * dp;
|
||||
}
|
||||
|
||||
std::vector<float> distances(ny);
|
||||
faiss::fvec_L2sqr_ny_transposed(
|
||||
distances.data(),
|
||||
x.data(),
|
||||
y.data(),
|
||||
y_sqlens.data(),
|
||||
d,
|
||||
ny, // no need for special offset to test all lines of code
|
||||
ny);
|
||||
|
||||
ASSERT_EQ(distances, true_distances)
|
||||
<< "Mismatching fvec_L2sqr_ny_transposed results for d = " << d;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecL2sqr, nearest_L2_squared_y_transposed) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
// modulo 8 results - 16 is to repeat the loop in the function
|
||||
int ny = 11; // this value will hit all the codepaths
|
||||
for (const auto d : {1, 2, 3, 4, 5, 6, 7, 8, 16}) {
|
||||
// initialize inputs
|
||||
std::vector<float> x(d);
|
||||
float x_sqlen = 0;
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = uniform(rng);
|
||||
x_sqlen += x[i] * x[i];
|
||||
}
|
||||
std::vector<float> y(d * ny);
|
||||
std::vector<float> y_sqlens(ny, 0);
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
for (size_t j = 0; j < y.size(); j++) {
|
||||
y[j] = uniform(rng);
|
||||
y_sqlens[i] += y[j] * y[j];
|
||||
}
|
||||
}
|
||||
|
||||
// get distances
|
||||
std::vector<float> distances(ny, 0);
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (size_t j = 0; j < d; j++) {
|
||||
dp += x[j] * y[i + j * ny];
|
||||
}
|
||||
distances[i] = x_sqlen + y_sqlens[i] - 2 * dp;
|
||||
}
|
||||
// find nearest
|
||||
size_t true_nearest_idx = 0;
|
||||
float min_dis = HUGE_VALF;
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
if (distances[i] < min_dis) {
|
||||
min_dis = distances[i];
|
||||
true_nearest_idx = i;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<float> buffer(ny);
|
||||
size_t nearest_idx = faiss::fvec_L2sqr_ny_nearest_y_transposed(
|
||||
buffer.data(),
|
||||
x.data(),
|
||||
y.data(),
|
||||
y_sqlens.data(),
|
||||
d,
|
||||
ny, // no need for special offset to test all lines of code
|
||||
ny);
|
||||
|
||||
ASSERT_EQ(nearest_idx, true_nearest_idx)
|
||||
<< "Mismatching fvec_L2sqr_ny_nearest_y_transposed results for d = "
|
||||
<< d;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecL1, manhattan_distance) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
// modulo 8 results - 16 is to repeat the while loop in the function
|
||||
for (const auto nrows : {8, 9, 10, 11, 12, 13, 14, 15, 16}) {
|
||||
std::vector<float> x(nrows);
|
||||
std::vector<float> y(nrows);
|
||||
float true_distance = 0;
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = uniform(rng);
|
||||
y[i] = uniform(rng);
|
||||
true_distance += std::abs(x[i] - y[i]);
|
||||
}
|
||||
|
||||
auto distance = faiss::fvec_L1(x.data(), y.data(), x.size());
|
||||
|
||||
ASSERT_EQ(distance, true_distance)
|
||||
<< "Mismatching fvec_Linf results for nrows = " << nrows;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecLinf, chebyshev_distance) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
// modulo 8 results - 16 is to repeat the while loop in the function
|
||||
for (const auto nrows : {8, 9, 10, 11, 12, 13, 14, 15, 16}) {
|
||||
std::vector<float> x(nrows);
|
||||
std::vector<float> y(nrows);
|
||||
float true_distance = 0;
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = uniform(rng);
|
||||
y[i] = uniform(rng);
|
||||
true_distance = std::max(true_distance, std::abs(x[i] - y[i]));
|
||||
}
|
||||
|
||||
auto distance = faiss::fvec_Linf(x.data(), y.data(), x.size());
|
||||
|
||||
ASSERT_EQ(distance, true_distance)
|
||||
<< "Mismatching fvec_Linf results for nrows = " << nrows;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecMadd, multiple_add) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
// modulo 8 results - 16 is to repeat the while loop in the function
|
||||
for (const auto nrows : {8, 9, 10, 11, 12, 13, 14, 15, 16}) {
|
||||
std::vector<float> a(nrows);
|
||||
std::vector<float> b(nrows);
|
||||
const float bf = uniform(rng);
|
||||
std::vector<float> true_distances(nrows);
|
||||
for (size_t i = 0; i < a.size(); i++) {
|
||||
a[i] = uniform(rng);
|
||||
b[i] = uniform(rng);
|
||||
true_distances[i] = a[i] + bf * b[i];
|
||||
}
|
||||
|
||||
std::vector<float> distances(nrows);
|
||||
faiss::fvec_madd(a.size(), a.data(), bf, b.data(), distances.data());
|
||||
|
||||
ASSERT_EQ(distances, true_distances)
|
||||
<< "Mismatching fvec_madd results for nrows = " << nrows;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecAdd, add_array) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
for (const auto nrows : {1, 2, 5, 10, 15, 20, 25}) {
|
||||
std::vector<float> a(nrows);
|
||||
std::vector<float> b(nrows);
|
||||
std::vector<float> true_distances(nrows);
|
||||
for (size_t i = 0; i < a.size(); i++) {
|
||||
a[i] = uniform(rng);
|
||||
b[i] = uniform(rng);
|
||||
true_distances[i] = a[i] + b[i];
|
||||
}
|
||||
|
||||
std::vector<float> distances(nrows);
|
||||
faiss::fvec_add(a.size(), a.data(), b.data(), distances.data());
|
||||
|
||||
ASSERT_EQ(distances, true_distances)
|
||||
<< "Mismatching array-array fvec_add results for nrows = "
|
||||
<< nrows;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestFvecAdd, add_value) {
|
||||
// ints instead of floats for 100% accuracy
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> uniform(0, 32);
|
||||
|
||||
for (const auto nrows : {1, 2, 5, 10, 15, 20, 25}) {
|
||||
std::vector<float> a(nrows);
|
||||
const float b = uniform(rng); // value to add
|
||||
std::vector<float> true_distances(nrows);
|
||||
for (size_t i = 0; i < a.size(); i++) {
|
||||
a[i] = uniform(rng);
|
||||
true_distances[i] = a[i] + b;
|
||||
}
|
||||
|
||||
std::vector<float> distances(nrows);
|
||||
faiss::fvec_add(a.size(), a.data(), b, distances.data());
|
||||
|
||||
ASSERT_EQ(distances, true_distances)
|
||||
<< "Mismatching array-value fvec_add results for nrows = "
|
||||
<< nrows;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user