Skip to content

Commit e6f7bb5

Browse files
mdouzefacebook-github-bot
authored andcommitted
move distances implementation into SIMD specific compile units
Summary: The main objective is to move the IVFFlat scanning code inside SIMD compile units, so that the loop over codes runs inside SIMD code. Unfortunately, I could not instant iate all the implementations with with_MetricType, hence the macro to define them. Reviewed By: algoriddle Differential Revision: D95936389
1 parent eddea2a commit e6f7bb5

13 files changed

Lines changed: 344 additions & 314 deletions

faiss/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,7 @@ set(FAISS_HEADERS
327327
utils/hamming_distance/avx512-inl.h
328328
utils/simd_impl/distances_autovec-inl.h
329329
utils/simd_impl/distances_simdlib256.h
330+
utils/simd_impl/IVFFlatScanner-inl.h
330331
utils/simd_impl/distances_sse-inl.h
331332
)
332333

faiss/IndexIVFFlat.cpp

Lines changed: 5 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313

1414
#include <cinttypes>
1515
#include <cstdio>
16-
#include <numeric>
1716

1817
#include <faiss/IndexFlat.h>
1918

@@ -23,6 +22,11 @@
2322
#include <faiss/impl/FaissAssert.h>
2423
#include <faiss/impl/expanded_scanners.h>
2524
#include <faiss/utils/extra_distances.h>
25+
26+
#define THE_SIMD_LEVEL SIMDLevel::NONE
27+
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
28+
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
29+
2630
#include <faiss/utils/utils.h>
2731

2832
namespace faiss {
@@ -145,48 +149,6 @@ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
145149
}
146150
}
147151

148-
namespace {
149-
150-
template <typename VectorDistance>
151-
struct IVFFlatScanner : InvertedListScanner {
152-
VectorDistance vd;
153-
using C = typename VectorDistance::C;
154-
155-
IVFFlatScanner(
156-
const VectorDistance& vd,
157-
bool store_pairs,
158-
const IDSelector* sel)
159-
: InvertedListScanner(store_pairs, sel), vd(vd) {
160-
keep_max = vd.is_similarity;
161-
code_size = vd.d * sizeof(float);
162-
}
163-
164-
const float* xi;
165-
void set_query(const float* query) override {
166-
this->xi = query;
167-
}
168-
169-
void set_list(idx_t list_no, float /* coarse_dis */) override {
170-
this->list_no = list_no;
171-
}
172-
173-
float distance_to_code(const uint8_t* code) const final {
174-
const float* yj = (float*)code;
175-
return vd(xi, yj);
176-
}
177-
178-
// redefining the scan_codes allows to inline the distance_to_code
179-
size_t scan_codes(
180-
size_t list_size,
181-
const uint8_t* codes,
182-
const idx_t* ids,
183-
ResultHandler& handler) const {
184-
return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler);
185-
}
186-
};
187-
188-
} // anonymous namespace
189-
190152
InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
191153
bool store_pairs,
192154
const IDSelector* sel,

faiss/IndexIVFFlat.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,38 @@ struct IndexIVFFlat : IndexIVF {
6262
IndexIVFFlat();
6363
};
6464

65+
template <typename VectorDistance>
66+
struct IVFFlatScanner : InvertedListScanner {
67+
VectorDistance vd;
68+
using C = typename VectorDistance::C;
69+
70+
IVFFlatScanner(
71+
const VectorDistance& vd,
72+
bool store_pairs,
73+
const IDSelector* sel)
74+
: InvertedListScanner(store_pairs, sel), vd(vd) {
75+
keep_max = vd.is_similarity;
76+
code_size = vd.d * sizeof(float);
77+
}
78+
79+
const float* xi;
80+
void set_query(const float* query) override {
81+
this->xi = query;
82+
}
83+
84+
void set_list(idx_t list_no, float /* coarse_dis */) override {
85+
this->list_no = list_no;
86+
}
87+
88+
float distance_to_code(const uint8_t* code) const final;
89+
90+
size_t scan_codes(
91+
size_t list_size,
92+
const uint8_t* codes,
93+
const idx_t* ids,
94+
ResultHandler& handler) const;
95+
};
96+
6597
struct IndexIVFFlatDedup : IndexIVFFlat {
6698
/** Maps ids stored in the index to the ids of vectors that are
6799
* the same. When a vector is unique, it does not appear in the

faiss/utils/distances_simd.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
#include <faiss/impl/platform_macros.h>
1919
#include <faiss/impl/simdlib/simdlib_dispatch.h>
2020

21-
#define AUTOVEC_LEVEL SIMDLevel::NONE
21+
#define THE_SIMD_LEVEL SIMDLevel::NONE
2222
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
2323
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
2424

faiss/utils/extra_distances-inl.h

Lines changed: 0 additions & 235 deletions
This file was deleted.

0 commit comments

Comments
 (0)