mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
ARM SVE optimization for L2 distance kernels adapted to new simd_impl structure (#5098)
Summary: This PR introduces ARM SVE (Scalable Vector Extension) optimization for L2 metric functions. The new implementation extends the existing NEON/SVE SIMD logic with SVE vectorized kernels for distance computations, improving both indexing and search performance. **Benchmarking Setup:** To evaluate the performance impact of the new SVE-enabled L2 distance computations, a dedicated benchmarking script was developed. This script systematically compares the earlier FAISS version and the latest SVE-optimized version using the SIFT1M dataset with 128-dimensional vectors. The benchmarks cover all major index types like FLAT, IVF_FLAT, IVF_PQ, IVF_SQ8, and HNSW. **Results:** Index | OSS QPS | SVE QPS | Improvement -- | -- | -- | -- Flat (L2) | 888.53 | 1188.12 | 33.70% IVF1024-Flat | 33,437.76 | 45,871.00 | 37.20% IVF1024-PQ32x8 | 33,597.01 | 40,650.00 | 21.00% HNSW (M=16, ef=64) | 72,837.67 | 77,482.00 | 6.40% IVF1024-SQ8 | 17,002.79 | 20,657.00 | 21.50% The above results clearly demonstrate that enabling SVE vectorization yields consistent performance gains across all tested index types, where both indexing and search times improved significantly. These improvements validate the effectiveness of the SVE implementation in accelerating L2 distance computations on modern ARM platforms. **Co-author:** This work was jointly developed with abhijain1204fujitsu. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5098 Test Plan: Run on an aarch64 host with SVE2, `svcntw() == 4` (128-bit vectors). **The SVE code only exists in dynamic-dispatch builds.** In the default static aarch64 Buck build, `xplat.bzl` filters `utils/simd_impl/distances_arm_sve.cpp` out of `srcs` and `SINGLE_SIMD_LEVEL` resolves to `ARM_NEON`, so a static-mode run exercises nothing here. CI runs both modes. ## Correctness ``` buck2 test -c faiss.dynamic_dispatch=true \ fbcode//faiss/tests:test_distances_simd \ fbcode//faiss/tests:test_distances_dispatch \ fbcode//faiss/tests:test_pq_encoding \ fbcode//faiss/tests:test_ivfpq_indexing \ fbcode//faiss/tests:test_ivfpq_codec \ fbcode//faiss/tests:test_pq_code_distance \ fbcode//faiss/tests:test_disable_pq_sdc_tables # Pass 32. Fail 0. buck2 test -c faiss.dynamic_dispatch=true \ fbcode//faiss/tests:test_build_blocks fbcode//faiss/tests:test_product_quantizer # Pass 105. Fail 0. buck2 test -c faiss.dynamic_dispatch=true fbcode//faiss/tests:test_index # Pass buck2 test fbcode//faiss/tests:test_distances_simd fbcode//faiss/tests:test_distances_dispatch # static mode: Pass 9. Skip 6 (only one SIMD level available). ``` `buck2 build fbcode//mode/opt fbcode//faiss:faiss -c faiss.dynamic_dispatch=true` and the static build both clean. `arc f` / `arc lint` clean. `test_distances_dispatch` forces every available level and compares against `SIMDLevel::NONE`; in DD mode on this host it reports `NONE`/`ARM_NEON`/`ARM_SVE` and no longer skips, so the `ARM_SVE` specializations are genuinely under test. Added coverage in `tests/test_distances_simd.cpp`: - `TestFvecL2sqrNy.D2`: `dim` swept over `{1,2,3,4,5,8,12,16,24,32,48,64}` (was `{2,4,8,12}`) so the `lanes1`..`lanes4` kernels are reached on 128/256/512-bit SVE, not just the `d1/d2/d4/d8` cases. - Both `_y_transposed` tests: `ny` swept over `{5,8,11,16,32}` (was hardcoded `11`, which never produced a full-vector block on any SVE width, so the no-tail path was untested). ## Kernel microbenchmark `mode/opt`, DD, pinned core, alternating runs, min-of-N x median-of-7, run-to-run spread <1%. `ny = 256` (ksub), `d = dsub`. vs this diff's parent: | kernel | speedup range over parent | |---|---| | `fvec_L2sqr_ny` | +17% to +725% (dsub 3 and 6 hit the scalar fallback, unchanged) | | `fvec_L2sqr_ny_nearest` | +48% to +130% | | `fvec_L2sqr_ny_transposed` | +234% to +298% | | `fvec_L2sqr_ny_nearest_y_transposed` | +318% to +658% | ## Index-level `faiss.contrib.datasets.SyntheticDataset`, d=128, nb=400k, nq=2k, k=1, single-threaded, min/best of 3 alternating rounds. "kernels" = SVE kernels with the pre-existing A0 dispatch; "wired" = this diff. PQ encode, `add()` seconds: | Index | parent | kernels only | wired | speedup | |---|---|---|---|---| | IVF1024-PQ32x8 | 16.17 | 16.16 | 6.32 | **2.56x** | | PQ32x8 (flat) | 14.80 | 14.82 | 4.97 | **2.98x** | | Flat / IVF-Flat / HNSW / IVF-SQ8 | - | - | - | 1.00x | Search QPS is unchanged everywhere (-0.9% to +0.5%, inside 0.2-3.0% run-to-run variance). recall@1 is identical to the parent for every index. The "kernels only" column is 1.00x because `ProductQuantizer` dispatched through `with_simd_level()` = `AVAILABLE_SIMD_LEVELS_A0`, which lacks the `ARM_SVE` bit, so the `case SIMDLevel::ARM_SVE:` arm failed its `if constexpr` and fell through to `ARM_NEON`. Verified by putting `abort()` in `fvec_L2sqr_ny_nearest<ARM_SVE>`: PQ encoding completed without hitting it, while `test_distances_dispatch` (A1) aborted immediately. The three `with_simd_level_a1` call sites in `impl/ProductQuantizer.cpp` are what close that gap. `compute_inner_prod_table` still uses A0 and so cannot reach the pre-existing `fvec_inner_products_ny<ARM_SVE>`; left alone as it is outside this diff's L2 scope and unmeasured here. Reviewed By: junjieqi Differential Revision: D106149015 Pulled By: mnorris11 fbshipit-source-id: c1144fc8fd0eefb30c9dcc8084b472152623a88e
This commit is contained in:
committed by
meta-codesync[bot]
parent
2832ac3f67
commit
317a207bca
@@ -280,7 +280,8 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
|
||||
} // namespace
|
||||
|
||||
void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
|
||||
with_simd_level([&]<SIMDLevel SL>() {
|
||||
// a1: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE specializations
|
||||
with_simd_level_a1([&]<SIMDLevel SL>() {
|
||||
switch (nbits) {
|
||||
case 8:
|
||||
compute_1_code<PQEncoder8, SL>(*this, x, code);
|
||||
@@ -294,7 +295,7 @@ void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
|
||||
compute_1_code<PQEncoderGeneric, SL>(*this, x, code);
|
||||
break;
|
||||
}
|
||||
}); // with_simd_level
|
||||
}); // with_simd_level_a1
|
||||
}
|
||||
|
||||
template <class PQDecoder>
|
||||
@@ -442,7 +443,8 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
|
||||
|
||||
void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
|
||||
const {
|
||||
with_simd_level([&]<SIMDLevel SL>() {
|
||||
// a1: fvec_L2sqr_ny / _transposed have ARM_SVE specializations
|
||||
with_simd_level_a1([&]<SIMDLevel SL>() {
|
||||
if (transposed_centroids.empty()) {
|
||||
// use regular version
|
||||
for (size_t m = 0; m < M; m++) {
|
||||
@@ -824,7 +826,8 @@ void ProductQuantizer::compute_sdc_table() {
|
||||
sdc_table.resize(M * ksub * ksub);
|
||||
|
||||
if (dsub < 4) {
|
||||
with_simd_level([&]<SIMDLevel SL>() {
|
||||
// a1: fvec_L2sqr_ny has an ARM_SVE specialization
|
||||
with_simd_level_a1([&]<SIMDLevel SL>() {
|
||||
#pragma omp parallel for
|
||||
for (int64_t mk = 0; mk < static_cast<int64_t>(M * ksub); mk++) {
|
||||
// allow omp to schedule in a more fine-grained way
|
||||
|
||||
@@ -238,6 +238,18 @@ inline auto with_simd_level_256bit(LambdaType&& action) {
|
||||
std::forward<LambdaType>(action));
|
||||
}
|
||||
|
||||
/**
|
||||
* Use for functions that have A0-level implementations plus a dedicated
|
||||
* ARM_SVE specialization. Plain with_simd_level() uses A0, which omits the
|
||||
* ARM_SVE bit, so on an SVE host the ARM_SVE case falls through to ARM_NEON
|
||||
* and the SVE specialization is never instantiated.
|
||||
*/
|
||||
template <typename LambdaType>
|
||||
inline auto with_simd_level_a1(LambdaType&& action) {
|
||||
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
|
||||
std::forward<LambdaType>(action));
|
||||
}
|
||||
|
||||
/**
|
||||
* Use for functions that have A0-level implementations plus an AVX512_SPR
|
||||
* specialization (e.g. using VPOPCNTDQ).
|
||||
|
||||
@@ -129,6 +129,22 @@ struct ElementOpIP {
|
||||
}
|
||||
};
|
||||
|
||||
struct ElementOpL2 {
|
||||
static svfloat32_t op(svbool_t pg, svfloat32_t x, svfloat32_t y) {
|
||||
const svfloat32_t diff = svsub_f32_x(pg, x, y);
|
||||
return svmul_f32_x(pg, diff, diff);
|
||||
}
|
||||
|
||||
static svfloat32_t merge(
|
||||
svbool_t pg,
|
||||
svfloat32_t z,
|
||||
svfloat32_t x,
|
||||
svfloat32_t y) {
|
||||
const svfloat32_t diff = svsub_f32_x(pg, x, y);
|
||||
return svmla_f32_x(pg, z, diff, diff);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename ElementOp>
|
||||
void fvec_op_ny_sve_d1(float* dis, const float* x, const float* y, size_t ny) {
|
||||
const size_t lanes = svcntw();
|
||||
@@ -513,10 +529,42 @@ void fvec_L2sqr_ny<SIMDLevel::ARM_SVE>(
|
||||
const float* y,
|
||||
size_t d,
|
||||
size_t ny) {
|
||||
// Use autovectorized L2sqr in a loop
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
dis[i] = fvec_L2sqr<SIMDLevel::ARM_SVE>(x, y, d);
|
||||
y += d;
|
||||
const size_t lanes = static_cast<size_t>(svcntw());
|
||||
|
||||
switch (d) {
|
||||
case 1:
|
||||
fvec_op_ny_sve_d1<ElementOpL2>(dis, x, y, ny);
|
||||
break;
|
||||
|
||||
case 2:
|
||||
fvec_op_ny_sve_d2<ElementOpL2>(dis, x, y, ny);
|
||||
break;
|
||||
|
||||
case 4:
|
||||
fvec_op_ny_sve_d4<ElementOpL2>(dis, x, y, ny);
|
||||
break;
|
||||
|
||||
case 8:
|
||||
fvec_op_ny_sve_d8<ElementOpL2>(dis, x, y, ny);
|
||||
break;
|
||||
|
||||
default:
|
||||
if (d == lanes)
|
||||
fvec_op_ny_sve_lanes1<ElementOpL2>(dis, x, y, ny);
|
||||
else if (d == lanes * 2)
|
||||
fvec_op_ny_sve_lanes2<ElementOpL2>(dis, x, y, ny);
|
||||
else if (d == lanes * 3)
|
||||
fvec_op_ny_sve_lanes3<ElementOpL2>(dis, x, y, ny);
|
||||
else if (d == lanes * 4)
|
||||
fvec_op_ny_sve_lanes4<ElementOpL2>(dis, x, y, ny);
|
||||
else {
|
||||
// Fallback: use autovectorized L2sqr
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
dis[i] = fvec_L2sqr<SIMDLevel::ARM_SVE>(x, y, d);
|
||||
y += d;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -527,14 +575,40 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_SVE>(
|
||||
const float* y,
|
||||
size_t d,
|
||||
size_t ny) {
|
||||
fvec_L2sqr_ny<SIMDLevel::ARM_SVE>(distances_tmp_buffer, x, y, d, ny);
|
||||
const size_t lanes = static_cast<size_t>(svcntw());
|
||||
|
||||
size_t nearest_idx = 0;
|
||||
float min_dis = HUGE_VALF;
|
||||
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
if (distances_tmp_buffer[i] < min_dis) {
|
||||
min_dis = distances_tmp_buffer[i];
|
||||
for (size_t i = 0; i < ny; ++i) {
|
||||
const float* yi = y + i * d;
|
||||
size_t j = 0;
|
||||
|
||||
svfloat32_t accv = svdup_n_f32(0.0f);
|
||||
|
||||
for (; j + lanes <= d; j += lanes) {
|
||||
const svbool_t pg = svptrue_b32();
|
||||
const svfloat32_t xv = svld1_f32(pg, x + j);
|
||||
const svfloat32_t yv = svld1_f32(pg, yi + j);
|
||||
const svfloat32_t diff = svsub_f32_x(pg, xv, yv);
|
||||
accv = svmla_f32_x(pg, accv, diff, diff);
|
||||
}
|
||||
|
||||
if (j < d) {
|
||||
const svbool_t pg = svwhilelt_b32_u64(j, d);
|
||||
const svfloat32_t xv = svld1_f32(pg, x + j);
|
||||
const svfloat32_t yv = svld1_f32(pg, yi + j);
|
||||
const svfloat32_t diff = svsub_f32_x(pg, xv, yv);
|
||||
// Merging predication: lanes outside the tail keep the partial
|
||||
// sums accumulated above, which the reduction below still adds in.
|
||||
// `_x` would leave them unspecified.
|
||||
accv = svmla_f32_m(pg, accv, diff, diff);
|
||||
}
|
||||
|
||||
const float dist = svaddv_f32(svptrue_b32(), accv);
|
||||
distances_tmp_buffer[i] = dist;
|
||||
if (dist < min_dis) {
|
||||
min_dis = dist;
|
||||
nearest_idx = i;
|
||||
}
|
||||
}
|
||||
@@ -542,7 +616,6 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_SVE>(
|
||||
return nearest_idx;
|
||||
}
|
||||
|
||||
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
|
||||
template <>
|
||||
void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
|
||||
float* dis,
|
||||
@@ -552,22 +625,26 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
|
||||
size_t d,
|
||||
size_t d_offset,
|
||||
size_t ny) {
|
||||
float x_sqlen = 0;
|
||||
FAISS_PRAGMA_IMPRECISE_LOOP
|
||||
for (size_t j = 0; j < d; j++) {
|
||||
x_sqlen += x[j] * x[j];
|
||||
}
|
||||
const size_t lanes = static_cast<size_t>(svcntw());
|
||||
const float x_sq = fvec_norm_L2sqr(x, d);
|
||||
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
FAISS_PRAGMA_IMPRECISE_LOOP
|
||||
for (size_t j = 0; j < d; j++) {
|
||||
dp += x[j] * y[i + j * d_offset];
|
||||
for (size_t k = 0; k < ny; k += lanes) {
|
||||
svbool_t pg = svwhilelt_b32_u64(k, ny);
|
||||
svfloat32_t acc = svdup_n_f32(0.0f);
|
||||
|
||||
for (size_t j = 0; j < d; ++j) {
|
||||
svfloat32_t ychunk = svld1_f32(pg, y + j * d_offset + k);
|
||||
svfloat32_t xj = svdup_n_f32(x[j]);
|
||||
acc = svmla_f32_x(pg, acc, xj, ychunk);
|
||||
}
|
||||
dis[i] = x_sqlen + y_sqlen[i] - 2 * dp;
|
||||
|
||||
svfloat32_t ysq = svld1_f32(pg, y_sqlen + k);
|
||||
svfloat32_t two_acc = svmul_f32_x(pg, acc, svdup_n_f32(2.0f));
|
||||
svfloat32_t sum = svadd_f32_x(pg, svdup_n_f32(x_sq), ysq);
|
||||
svfloat32_t res = svsub_f32_x(pg, sum, two_acc);
|
||||
svst1_f32(pg, dis + k, res);
|
||||
}
|
||||
}
|
||||
FAISS_PRAGMA_IMPRECISE_FUNCTION_END
|
||||
|
||||
template <>
|
||||
size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
|
||||
@@ -578,20 +655,53 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
|
||||
size_t d,
|
||||
size_t d_offset,
|
||||
size_t ny) {
|
||||
fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
|
||||
distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
|
||||
const size_t lanes = svcntw();
|
||||
const float x_sq = fvec_norm_L2sqr(x, d);
|
||||
|
||||
size_t nearest_idx = 0;
|
||||
float min_dis = HUGE_VALF;
|
||||
size_t current_min_idx = 0;
|
||||
svfloat32_t current_min_v = svdup_n_f32(HUGE_VALF);
|
||||
|
||||
for (size_t i = 0; i < ny; i++) {
|
||||
if (distances_tmp_buffer[i] < min_dis) {
|
||||
min_dis = distances_tmp_buffer[i];
|
||||
nearest_idx = i;
|
||||
// Max SVE vector length is 2048 bits → 64 fp32 lanes
|
||||
float tmp_buf[64];
|
||||
|
||||
for (size_t k = 0; k < ny; k += lanes) {
|
||||
svbool_t pg = svwhilelt_b32_u64(k, ny);
|
||||
svfloat32_t acc = svdup_n_f32(0.0f);
|
||||
|
||||
for (size_t j = 0; j < d; ++j) {
|
||||
svfloat32_t ychunk = svld1_f32(pg, y + j * d_offset + k);
|
||||
svfloat32_t xj = svdup_n_f32(x[j]);
|
||||
acc = svmla_f32_x(pg, acc, xj, ychunk);
|
||||
}
|
||||
|
||||
svfloat32_t ysq = svld1_f32(pg, y_sqlen + k);
|
||||
svfloat32_t two_acc = svmul_f32_x(pg, acc, svdup_n_f32(2.0f));
|
||||
svfloat32_t sum = svadd_f32_x(pg, svdup_n_f32(x_sq), ysq);
|
||||
svfloat32_t res = svsub_f32_x(pg, sum, two_acc);
|
||||
|
||||
svst1_f32(pg, distances_tmp_buffer + k, res);
|
||||
|
||||
svbool_t less_mask = svcmplt_f32(pg, res, current_min_v);
|
||||
|
||||
if (svptest_any(pg, less_mask)) {
|
||||
float vec_min = svminv_f32(pg, res);
|
||||
|
||||
current_min_v =
|
||||
svmin_f32_x(pg, current_min_v, svdup_n_f32(vec_min));
|
||||
|
||||
svst1_f32(pg, tmp_buf, res);
|
||||
|
||||
size_t cnt = (size_t)svcntw();
|
||||
for (size_t lane = 0; lane < cnt && (k + lane) < ny; ++lane) {
|
||||
if (tmp_buf[lane] == vec_min) {
|
||||
current_min_idx = k + lane;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nearest_idx;
|
||||
return current_min_idx;
|
||||
}
|
||||
|
||||
template <>
|
||||
|
||||
@@ -46,7 +46,10 @@ TEST(TestFvecL2sqrNy, D2) {
|
||||
std::default_random_engine rng(123);
|
||||
std::uniform_int_distribution<int32_t> u(0, 32);
|
||||
|
||||
for (const auto dim : {2, 4, 8, 12}) {
|
||||
// 1/2/4/8 hit the dedicated per-dimension kernels; the multiples of 4, 8
|
||||
// and 16 reach the `lanes1`..`lanes4` kernels on 128/256/512-bit SVE
|
||||
// respectively, and 3/5 fall through to the scalar tail.
|
||||
for (const auto dim : {1, 2, 3, 4, 5, 8, 12, 16, 24, 32, 48, 64}) {
|
||||
std::vector<float> x(dim, 0);
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
x[i] = u(rng);
|
||||
@@ -112,47 +115,51 @@ TEST(TestFvecL2sqr, distances_L2_squared_y_transposed) {
|
||||
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 (int 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];
|
||||
// ny = 8/16/32 are exact multiples of the SVE vector length on 256/512/1024
|
||||
// bit implementations, so they leave no partial-predicate tail; 5 and 11
|
||||
// do.
|
||||
for (const auto ny : {5, 8, 11, 16, 32}) {
|
||||
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];
|
||||
}
|
||||
}
|
||||
|
||||
// perform function
|
||||
std::vector<float> true_distances(ny, 0);
|
||||
for (int i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (int j = 0; j < d; j++) {
|
||||
dp += x[j] * y[i + j * ny];
|
||||
std::vector<float> y(d * ny);
|
||||
std::vector<float> y_sqlens(ny, 0);
|
||||
for (int 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];
|
||||
}
|
||||
}
|
||||
true_distances[i] = x_sqlen + y_sqlens[i] - 2 * dp;
|
||||
|
||||
// perform function
|
||||
std::vector<float> true_distances(ny, 0);
|
||||
for (int i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (int 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 << ", ny = " << ny;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -161,57 +168,60 @@ TEST(TestFvecL2sqr, nearest_L2_squared_y_transposed) {
|
||||
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 (int 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];
|
||||
// ny = 8/16/32 are exact multiples of the SVE vector length on 256/512/1024
|
||||
// bit implementations, so they leave no partial-predicate tail; 5 and 11
|
||||
// do.
|
||||
for (const auto ny : {5, 8, 11, 16, 32}) {
|
||||
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];
|
||||
}
|
||||
}
|
||||
|
||||
// get distances
|
||||
std::vector<float> distances(ny, 0);
|
||||
for (int i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (int j = 0; j < d; j++) {
|
||||
dp += x[j] * y[i + j * ny];
|
||||
std::vector<float> y(d * ny);
|
||||
std::vector<float> y_sqlens(ny, 0);
|
||||
for (int 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];
|
||||
}
|
||||
}
|
||||
distances[i] = x_sqlen + y_sqlens[i] - 2 * dp;
|
||||
}
|
||||
// find nearest
|
||||
size_t true_nearest_idx = 0;
|
||||
float min_dis = HUGE_VALF;
|
||||
for (int i = 0; i < ny; i++) {
|
||||
if (distances[i] < min_dis) {
|
||||
min_dis = distances[i];
|
||||
true_nearest_idx = i;
|
||||
|
||||
// get distances
|
||||
std::vector<float> distances(ny, 0);
|
||||
for (int i = 0; i < ny; i++) {
|
||||
float dp = 0;
|
||||
for (int 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 (int 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 << ", ny = " << ny;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user