faiss 0.2.0 → 0.2.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/CHANGELOG.md +16 -0
- data/LICENSE.txt +1 -1
- data/README.md +7 -7
- data/ext/faiss/extconf.rb +6 -3
- data/ext/faiss/numo.hpp +4 -4
- data/ext/faiss/utils.cpp +1 -1
- data/ext/faiss/utils.h +1 -1
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +292 -291
- data/vendor/faiss/faiss/AutoTune.h +55 -56
- data/vendor/faiss/faiss/Clustering.cpp +365 -194
- data/vendor/faiss/faiss/Clustering.h +102 -35
- data/vendor/faiss/faiss/IVFlib.cpp +171 -195
- data/vendor/faiss/faiss/IVFlib.h +48 -51
- data/vendor/faiss/faiss/Index.cpp +85 -103
- data/vendor/faiss/faiss/Index.h +54 -48
- data/vendor/faiss/faiss/Index2Layer.cpp +126 -224
- data/vendor/faiss/faiss/Index2Layer.h +22 -36
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +407 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +195 -0
- data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
- data/vendor/faiss/faiss/IndexBinary.h +140 -132
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
- data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
- data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
- data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
- data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
- data/vendor/faiss/faiss/IndexFlat.cpp +115 -176
- data/vendor/faiss/faiss/IndexFlat.h +42 -59
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +67 -0
- data/vendor/faiss/faiss/IndexFlatCodes.h +47 -0
- data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
- data/vendor/faiss/faiss/IndexHNSW.h +57 -41
- data/vendor/faiss/faiss/IndexIVF.cpp +545 -453
- data/vendor/faiss/faiss/IndexIVF.h +169 -118
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +316 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +121 -0
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +247 -252
- data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +459 -517
- data/vendor/faiss/faiss/IndexIVFPQ.h +75 -67
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
- data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +163 -150
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +38 -25
- data/vendor/faiss/faiss/IndexLSH.cpp +66 -113
- data/vendor/faiss/faiss/IndexLSH.h +20 -38
- data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
- data/vendor/faiss/faiss/IndexLattice.h +11 -16
- data/vendor/faiss/faiss/IndexNNDescent.cpp +229 -0
- data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
- data/vendor/faiss/faiss/IndexNSG.cpp +301 -0
- data/vendor/faiss/faiss/IndexNSG.h +85 -0
- data/vendor/faiss/faiss/IndexPQ.cpp +387 -495
- data/vendor/faiss/faiss/IndexPQ.h +64 -82
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
- data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
- data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
- data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
- data/vendor/faiss/faiss/IndexRefine.cpp +139 -127
- data/vendor/faiss/faiss/IndexRefine.h +32 -23
- data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
- data/vendor/faiss/faiss/IndexReplicas.h +62 -56
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +111 -172
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -59
- data/vendor/faiss/faiss/IndexShards.cpp +256 -240
- data/vendor/faiss/faiss/IndexShards.h +85 -73
- data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
- data/vendor/faiss/faiss/MatrixStats.h +7 -10
- data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
- data/vendor/faiss/faiss/MetaIndexes.h +40 -34
- data/vendor/faiss/faiss/MetricType.h +7 -7
- data/vendor/faiss/faiss/VectorTransform.cpp +654 -475
- data/vendor/faiss/faiss/VectorTransform.h +64 -89
- data/vendor/faiss/faiss/clone_index.cpp +78 -73
- data/vendor/faiss/faiss/clone_index.h +4 -9
- data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
- data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +198 -171
- data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
- data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
- data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
- data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
- data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
- data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
- data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
- data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
- data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
- data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
- data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
- data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
- data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
- data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
- data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
- data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
- data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
- data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
- data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
- data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
- data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
- data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +503 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +175 -0
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
- data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
- data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
- data/vendor/faiss/faiss/impl/FaissException.h +41 -29
- data/vendor/faiss/faiss/impl/HNSW.cpp +606 -617
- data/vendor/faiss/faiss/impl/HNSW.h +179 -200
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +855 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +244 -0
- data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
- data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
- data/vendor/faiss/faiss/impl/NSG.cpp +679 -0
- data/vendor/faiss/faiss/impl/NSG.h +199 -0
- data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
- data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
- data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +758 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +188 -0
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +647 -707
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
- data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
- data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
- data/vendor/faiss/faiss/impl/index_read.cpp +631 -480
- data/vendor/faiss/faiss/impl/index_write.cpp +547 -407
- data/vendor/faiss/faiss/impl/io.cpp +76 -95
- data/vendor/faiss/faiss/impl/io.h +31 -41
- data/vendor/faiss/faiss/impl/io_macros.h +60 -29
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +301 -0
- data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
- data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
- data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
- data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
- data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
- data/vendor/faiss/faiss/index_factory.cpp +619 -397
- data/vendor/faiss/faiss/index_factory.h +8 -6
- data/vendor/faiss/faiss/index_io.h +23 -26
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
- data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
- data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
- data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
- data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
- data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
- data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
- data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
- data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
- data/vendor/faiss/faiss/utils/Heap.h +186 -209
- data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
- data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
- data/vendor/faiss/faiss/utils/distances.cpp +305 -312
- data/vendor/faiss/faiss/utils/distances.h +170 -122
- data/vendor/faiss/faiss/utils/distances_simd.cpp +498 -508
- data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
- data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
- data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
- data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
- data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
- data/vendor/faiss/faiss/utils/hamming.h +62 -85
- data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
- data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
- data/vendor/faiss/faiss/utils/partitioning.h +26 -21
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
- data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
- data/vendor/faiss/faiss/utils/random.cpp +39 -63
- data/vendor/faiss/faiss/utils/random.h +13 -16
- data/vendor/faiss/faiss/utils/simdlib.h +4 -2
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
- data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
- data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
- data/vendor/faiss/faiss/utils/utils.cpp +304 -287
- data/vendor/faiss/faiss/utils/utils.h +54 -49
- metadata +29 -4
|
@@ -15,7 +15,6 @@
|
|
|
15
15
|
|
|
16
16
|
namespace faiss {
|
|
17
17
|
|
|
18
|
-
|
|
19
18
|
/** Fast scan version of IVFPQ. Works for 4-bit PQ for now.
|
|
20
19
|
*
|
|
21
20
|
* The codes in the inverted lists are not stored sequentially but
|
|
@@ -32,13 +31,12 @@ namespace faiss {
|
|
|
32
31
|
* 13: idem, collect results in reservoir
|
|
33
32
|
*/
|
|
34
33
|
|
|
35
|
-
struct IndexIVFPQFastScan: IndexIVF {
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
ProductQuantizer pq; ///< produces the codes
|
|
34
|
+
struct IndexIVFPQFastScan : IndexIVF {
|
|
35
|
+
bool by_residual; ///< Encode residual or plain vector?
|
|
36
|
+
ProductQuantizer pq; ///< produces the codes
|
|
39
37
|
|
|
40
38
|
// size of the kernel
|
|
41
|
-
int bbs;
|
|
39
|
+
int bbs; // set at build time
|
|
42
40
|
|
|
43
41
|
// M rounded up to a multiple of 2
|
|
44
42
|
size_t M2;
|
|
@@ -57,88 +55,114 @@ struct IndexIVFPQFastScan: IndexIVF {
|
|
|
57
55
|
int qbs = 0;
|
|
58
56
|
size_t qbs2 = 0;
|
|
59
57
|
|
|
60
|
-
IndexIVFPQFastScan
|
|
61
|
-
Index
|
|
62
|
-
size_t
|
|
63
|
-
|
|
58
|
+
IndexIVFPQFastScan(
|
|
59
|
+
Index* quantizer,
|
|
60
|
+
size_t d,
|
|
61
|
+
size_t nlist,
|
|
62
|
+
size_t M,
|
|
63
|
+
size_t nbits_per_idx,
|
|
64
|
+
MetricType metric = METRIC_L2,
|
|
65
|
+
int bbs = 32);
|
|
64
66
|
|
|
65
|
-
IndexIVFPQFastScan
|
|
67
|
+
IndexIVFPQFastScan();
|
|
66
68
|
|
|
67
69
|
// built from an IndexIVFPQ
|
|
68
|
-
explicit IndexIVFPQFastScan(const IndexIVFPQ
|
|
70
|
+
explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
|
|
69
71
|
|
|
70
72
|
/// orig's inverted lists (for debugging)
|
|
71
|
-
InvertedLists
|
|
73
|
+
InvertedLists* orig_invlists = nullptr;
|
|
72
74
|
|
|
73
|
-
void train_residual
|
|
75
|
+
void train_residual(idx_t n, const float* x) override;
|
|
74
76
|
|
|
75
77
|
/// build precomputed table, possibly updating use_precomputed_table
|
|
76
|
-
void precompute_table
|
|
78
|
+
void precompute_table();
|
|
77
79
|
|
|
78
80
|
/// same as the regular IVFPQ encoder. The codes are not reorganized by
|
|
79
81
|
/// blocks a that point
|
|
80
82
|
void encode_vectors(
|
|
81
|
-
idx_t n,
|
|
82
|
-
const
|
|
83
|
+
idx_t n,
|
|
84
|
+
const float* x,
|
|
85
|
+
const idx_t* list_nos,
|
|
86
|
+
uint8_t* codes,
|
|
83
87
|
bool include_listno = false) const override;
|
|
84
88
|
|
|
85
|
-
void add_with_ids
|
|
86
|
-
idx_t n, const float * x, const idx_t *xids) override;
|
|
89
|
+
void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
|
|
87
90
|
|
|
88
91
|
void search(
|
|
89
|
-
idx_t n,
|
|
90
|
-
float*
|
|
92
|
+
idx_t n,
|
|
93
|
+
const float* x,
|
|
94
|
+
idx_t k,
|
|
95
|
+
float* distances,
|
|
96
|
+
idx_t* labels) const override;
|
|
91
97
|
|
|
92
98
|
// prepare look-up tables
|
|
93
99
|
|
|
94
100
|
void compute_LUT(
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
101
|
+
size_t n,
|
|
102
|
+
const float* x,
|
|
103
|
+
const idx_t* coarse_ids,
|
|
104
|
+
const float* coarse_dis,
|
|
105
|
+
AlignedTable<float>& dis_tables,
|
|
106
|
+
AlignedTable<float>& biases) const;
|
|
100
107
|
|
|
101
108
|
void compute_LUT_uint8(
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
109
|
+
size_t n,
|
|
110
|
+
const float* x,
|
|
111
|
+
const idx_t* coarse_ids,
|
|
112
|
+
const float* coarse_dis,
|
|
113
|
+
AlignedTable<uint8_t>& dis_tables,
|
|
114
|
+
AlignedTable<uint16_t>& biases,
|
|
115
|
+
float* normalizers) const;
|
|
108
116
|
|
|
109
117
|
// internal search funcs
|
|
110
118
|
|
|
111
|
-
template<bool is_max>
|
|
119
|
+
template <bool is_max>
|
|
112
120
|
void search_dispatch_implem(
|
|
113
|
-
idx_t n,
|
|
114
|
-
float*
|
|
121
|
+
idx_t n,
|
|
122
|
+
const float* x,
|
|
123
|
+
idx_t k,
|
|
124
|
+
float* distances,
|
|
125
|
+
idx_t* labels) const;
|
|
115
126
|
|
|
116
|
-
template<class C>
|
|
127
|
+
template <class C>
|
|
117
128
|
void search_implem_1(
|
|
118
|
-
idx_t n,
|
|
119
|
-
float*
|
|
129
|
+
idx_t n,
|
|
130
|
+
const float* x,
|
|
131
|
+
idx_t k,
|
|
132
|
+
float* distances,
|
|
133
|
+
idx_t* labels) const;
|
|
120
134
|
|
|
121
|
-
template<class C>
|
|
135
|
+
template <class C>
|
|
122
136
|
void search_implem_2(
|
|
123
|
-
idx_t n,
|
|
124
|
-
float*
|
|
137
|
+
idx_t n,
|
|
138
|
+
const float* x,
|
|
139
|
+
idx_t k,
|
|
140
|
+
float* distances,
|
|
141
|
+
idx_t* labels) const;
|
|
125
142
|
|
|
126
143
|
// implem 10 and 12 are not multithreaded internally, so
|
|
127
144
|
// export search stats
|
|
128
|
-
template<class C>
|
|
145
|
+
template <class C>
|
|
129
146
|
void search_implem_10(
|
|
130
|
-
idx_t n,
|
|
131
|
-
float*
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
147
|
+
idx_t n,
|
|
148
|
+
const float* x,
|
|
149
|
+
idx_t k,
|
|
150
|
+
float* distances,
|
|
151
|
+
idx_t* labels,
|
|
152
|
+
int impl,
|
|
153
|
+
size_t* ndis_out,
|
|
154
|
+
size_t* nlist_out) const;
|
|
155
|
+
|
|
156
|
+
template <class C>
|
|
135
157
|
void search_implem_12(
|
|
136
|
-
idx_t n,
|
|
137
|
-
float*
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
158
|
+
idx_t n,
|
|
159
|
+
const float* x,
|
|
160
|
+
idx_t k,
|
|
161
|
+
float* distances,
|
|
162
|
+
idx_t* labels,
|
|
163
|
+
int impl,
|
|
164
|
+
size_t* ndis_out,
|
|
165
|
+
size_t* nlist_out) const;
|
|
142
166
|
};
|
|
143
167
|
|
|
144
168
|
struct IVFFastScanStats {
|
|
@@ -148,13 +172,15 @@ struct IVFFastScanStats {
|
|
|
148
172
|
uint64_t reservoir_times[4];
|
|
149
173
|
|
|
150
174
|
double Mcy_at(int i) {
|
|
151
|
-
return times[i] / (1000*1000.0);
|
|
175
|
+
return times[i] / (1000 * 1000.0);
|
|
152
176
|
}
|
|
153
177
|
|
|
154
178
|
double Mcy_reservoir_at(int i) {
|
|
155
|
-
return reservoir_times[i] / (1000*1000.0);
|
|
179
|
+
return reservoir_times[i] / (1000 * 1000.0);
|
|
180
|
+
}
|
|
181
|
+
IVFFastScanStats() {
|
|
182
|
+
reset();
|
|
156
183
|
}
|
|
157
|
-
IVFFastScanStats() {reset();}
|
|
158
184
|
void reset() {
|
|
159
185
|
memset(this, 0, sizeof(*this));
|
|
160
186
|
}
|
|
@@ -162,5 +188,4 @@ struct IVFFastScanStats {
|
|
|
162
188
|
|
|
163
189
|
FAISS_API extern IVFFastScanStats IVFFastScan_stats;
|
|
164
190
|
|
|
165
|
-
|
|
166
191
|
} // namespace faiss
|
|
@@ -12,115 +12,116 @@
|
|
|
12
12
|
#include <cinttypes>
|
|
13
13
|
|
|
14
14
|
#include <faiss/utils/Heap.h>
|
|
15
|
-
#include <faiss/utils/utils.h>
|
|
16
15
|
#include <faiss/utils/distances.h>
|
|
16
|
+
#include <faiss/utils/utils.h>
|
|
17
17
|
|
|
18
18
|
#include <faiss/impl/FaissAssert.h>
|
|
19
19
|
|
|
20
|
-
|
|
21
20
|
namespace faiss {
|
|
22
21
|
|
|
23
22
|
/*****************************************
|
|
24
23
|
* IndexIVFPQR implementation
|
|
25
24
|
******************************************/
|
|
26
25
|
|
|
27
|
-
IndexIVFPQR::IndexIVFPQR
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
26
|
+
IndexIVFPQR::IndexIVFPQR(
|
|
27
|
+
Index* quantizer,
|
|
28
|
+
size_t d,
|
|
29
|
+
size_t nlist,
|
|
30
|
+
size_t M,
|
|
31
|
+
size_t nbits_per_idx,
|
|
32
|
+
size_t M_refine,
|
|
33
|
+
size_t nbits_per_idx_refine)
|
|
34
|
+
: IndexIVFPQ(quantizer, d, nlist, M, nbits_per_idx),
|
|
35
|
+
refine_pq(d, M_refine, nbits_per_idx_refine),
|
|
36
|
+
k_factor(4) {
|
|
35
37
|
by_residual = true;
|
|
36
38
|
}
|
|
37
39
|
|
|
38
|
-
IndexIVFPQR::IndexIVFPQR ()
|
|
39
|
-
k_factor (1)
|
|
40
|
-
{
|
|
40
|
+
IndexIVFPQR::IndexIVFPQR() : k_factor(1) {
|
|
41
41
|
by_residual = true;
|
|
42
42
|
}
|
|
43
43
|
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
void IndexIVFPQR::reset()
|
|
47
|
-
{
|
|
44
|
+
void IndexIVFPQR::reset() {
|
|
48
45
|
IndexIVFPQ::reset();
|
|
49
46
|
refine_codes.clear();
|
|
50
47
|
}
|
|
51
48
|
|
|
49
|
+
void IndexIVFPQR::train_residual(idx_t n, const float* x) {
|
|
50
|
+
float* residual_2 = new float[n * d];
|
|
51
|
+
ScopeDeleter<float> del(residual_2);
|
|
52
52
|
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
void IndexIVFPQR::train_residual (idx_t n, const float *x)
|
|
56
|
-
{
|
|
57
|
-
|
|
58
|
-
float * residual_2 = new float [n * d];
|
|
59
|
-
ScopeDeleter <float> del(residual_2);
|
|
60
|
-
|
|
61
|
-
train_residual_o (n, x, residual_2);
|
|
53
|
+
train_residual_o(n, x, residual_2);
|
|
62
54
|
|
|
63
55
|
if (verbose)
|
|
64
|
-
printf
|
|
65
|
-
|
|
56
|
+
printf("training %zdx%zd 2nd level PQ quantizer on %" PRId64
|
|
57
|
+
" %dD-vectors\n",
|
|
58
|
+
refine_pq.M,
|
|
59
|
+
refine_pq.ksub,
|
|
60
|
+
n,
|
|
61
|
+
d);
|
|
66
62
|
|
|
67
63
|
refine_pq.cp.max_points_per_centroid = 1000;
|
|
68
64
|
refine_pq.cp.verbose = verbose;
|
|
69
65
|
|
|
70
|
-
refine_pq.train
|
|
71
|
-
|
|
66
|
+
refine_pq.train(n, residual_2);
|
|
72
67
|
}
|
|
73
68
|
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
add_core (n, x, xids, nullptr);
|
|
69
|
+
void IndexIVFPQR::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
|
|
70
|
+
add_core(n, x, xids, nullptr);
|
|
77
71
|
}
|
|
78
72
|
|
|
79
|
-
void IndexIVFPQR::add_core
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
73
|
+
void IndexIVFPQR::add_core(
|
|
74
|
+
idx_t n,
|
|
75
|
+
const float* x,
|
|
76
|
+
const idx_t* xids,
|
|
77
|
+
const idx_t* precomputed_idx) {
|
|
78
|
+
float* residual_2 = new float[n * d];
|
|
79
|
+
ScopeDeleter<float> del(residual_2);
|
|
84
80
|
|
|
85
81
|
idx_t n0 = ntotal;
|
|
86
82
|
|
|
87
|
-
add_core_o
|
|
88
|
-
|
|
89
|
-
refine_codes.resize (ntotal * refine_pq.code_size);
|
|
90
|
-
|
|
91
|
-
refine_pq.compute_codes (
|
|
92
|
-
residual_2, &refine_codes[n0 * refine_pq.code_size], n);
|
|
83
|
+
add_core_o(n, x, xids, residual_2, precomputed_idx);
|
|
93
84
|
|
|
85
|
+
refine_codes.resize(ntotal * refine_pq.code_size);
|
|
94
86
|
|
|
87
|
+
refine_pq.compute_codes(
|
|
88
|
+
residual_2, &refine_codes[n0 * refine_pq.code_size], n);
|
|
95
89
|
}
|
|
96
90
|
#define TIC t0 = get_cycles()
|
|
97
|
-
#define TOC get_cycles
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
91
|
+
#define TOC get_cycles() - t0
|
|
92
|
+
|
|
93
|
+
void IndexIVFPQR::search_preassigned(
|
|
94
|
+
idx_t n,
|
|
95
|
+
const float* x,
|
|
96
|
+
idx_t k,
|
|
97
|
+
const idx_t* idx,
|
|
98
|
+
const float* L1_dis,
|
|
99
|
+
float* distances,
|
|
100
|
+
idx_t* labels,
|
|
104
101
|
bool store_pairs,
|
|
105
|
-
const IVFSearchParameters
|
|
106
|
-
) const
|
|
107
|
-
{
|
|
102
|
+
const IVFSearchParameters* params,
|
|
103
|
+
IndexIVFStats* stats) const {
|
|
108
104
|
uint64_t t0;
|
|
109
105
|
TIC;
|
|
110
106
|
size_t k_coarse = long(k * k_factor);
|
|
111
|
-
idx_t
|
|
112
|
-
ScopeDeleter<idx_t> del1
|
|
107
|
+
idx_t* coarse_labels = new idx_t[k_coarse * n];
|
|
108
|
+
ScopeDeleter<idx_t> del1(coarse_labels);
|
|
113
109
|
{ // query with quantizer levels 1 and 2.
|
|
114
|
-
float
|
|
110
|
+
float* coarse_distances = new float[k_coarse * n];
|
|
115
111
|
ScopeDeleter<float> del(coarse_distances);
|
|
116
112
|
|
|
117
|
-
IndexIVFPQ::search_preassigned
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
113
|
+
IndexIVFPQ::search_preassigned(
|
|
114
|
+
n,
|
|
115
|
+
x,
|
|
116
|
+
k_coarse,
|
|
117
|
+
idx,
|
|
118
|
+
L1_dis,
|
|
119
|
+
coarse_distances,
|
|
120
|
+
coarse_labels,
|
|
121
|
+
true,
|
|
122
|
+
params);
|
|
121
123
|
}
|
|
122
124
|
|
|
123
|
-
|
|
124
125
|
indexIVFPQ_stats.search_cycles += TOC;
|
|
125
126
|
|
|
126
127
|
TIC;
|
|
@@ -130,91 +131,92 @@ void IndexIVFPQR::search_preassigned (
|
|
|
130
131
|
#pragma omp parallel reduction(+ : n_refine)
|
|
131
132
|
{
|
|
132
133
|
// tmp buffers
|
|
133
|
-
float
|
|
134
|
-
ScopeDeleter<float> del
|
|
135
|
-
float
|
|
134
|
+
float* residual_1 = new float[2 * d];
|
|
135
|
+
ScopeDeleter<float> del(residual_1);
|
|
136
|
+
float* residual_2 = residual_1 + d;
|
|
136
137
|
#pragma omp for
|
|
137
138
|
for (idx_t i = 0; i < n; i++) {
|
|
138
|
-
const float
|
|
139
|
-
const idx_t
|
|
140
|
-
float
|
|
141
|
-
idx_t
|
|
142
|
-
maxheap_heapify
|
|
139
|
+
const float* xq = x + i * d;
|
|
140
|
+
const idx_t* shortlist = coarse_labels + k_coarse * i;
|
|
141
|
+
float* heap_sim = distances + k * i;
|
|
142
|
+
idx_t* heap_ids = labels + k * i;
|
|
143
|
+
maxheap_heapify(k, heap_sim, heap_ids);
|
|
143
144
|
|
|
144
145
|
for (int j = 0; j < k_coarse; j++) {
|
|
145
146
|
idx_t sl = shortlist[j];
|
|
146
147
|
|
|
147
|
-
if (sl == -1)
|
|
148
|
+
if (sl == -1)
|
|
149
|
+
continue;
|
|
148
150
|
|
|
149
151
|
int list_no = lo_listno(sl);
|
|
150
152
|
int ofs = lo_offset(sl);
|
|
151
153
|
|
|
152
|
-
assert
|
|
153
|
-
assert
|
|
154
|
+
assert(list_no >= 0 && list_no < nlist);
|
|
155
|
+
assert(ofs >= 0 && ofs < invlists->list_size(list_no));
|
|
154
156
|
|
|
155
157
|
// 1st level residual
|
|
156
|
-
quantizer->compute_residual
|
|
158
|
+
quantizer->compute_residual(xq, residual_1, list_no);
|
|
157
159
|
|
|
158
160
|
// 2nd level residual
|
|
159
|
-
const uint8_t
|
|
160
|
-
invlists->get_single_code (list_no, ofs);
|
|
161
|
+
const uint8_t* l2code = invlists->get_single_code(list_no, ofs);
|
|
161
162
|
|
|
162
|
-
pq.decode
|
|
163
|
+
pq.decode(l2code, residual_2);
|
|
163
164
|
for (int l = 0; l < d; l++)
|
|
164
165
|
residual_2[l] = residual_1[l] - residual_2[l];
|
|
165
166
|
|
|
166
167
|
// 3rd level residual's approximation
|
|
167
|
-
idx_t id = invlists->get_single_id
|
|
168
|
-
assert
|
|
169
|
-
refine_pq.decode
|
|
170
|
-
|
|
168
|
+
idx_t id = invlists->get_single_id(list_no, ofs);
|
|
169
|
+
assert(0 <= id && id < ntotal);
|
|
170
|
+
refine_pq.decode(
|
|
171
|
+
&refine_codes[id * refine_pq.code_size], residual_1);
|
|
171
172
|
|
|
172
|
-
float dis = fvec_L2sqr
|
|
173
|
+
float dis = fvec_L2sqr(residual_1, residual_2, d);
|
|
173
174
|
|
|
174
175
|
if (dis < heap_sim[0]) {
|
|
175
176
|
idx_t id_or_pair = store_pairs ? sl : id;
|
|
176
|
-
maxheap_replace_top
|
|
177
|
+
maxheap_replace_top(k, heap_sim, heap_ids, dis, id_or_pair);
|
|
177
178
|
}
|
|
178
|
-
n_refine
|
|
179
|
+
n_refine++;
|
|
179
180
|
}
|
|
180
|
-
maxheap_reorder
|
|
181
|
+
maxheap_reorder(k, heap_sim, heap_ids);
|
|
181
182
|
}
|
|
182
183
|
}
|
|
183
184
|
indexIVFPQ_stats.nrefine += n_refine;
|
|
184
185
|
indexIVFPQ_stats.refine_cycles += TOC;
|
|
185
186
|
}
|
|
186
187
|
|
|
187
|
-
void IndexIVFPQR::reconstruct_from_offset
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
188
|
+
void IndexIVFPQR::reconstruct_from_offset(
|
|
189
|
+
int64_t list_no,
|
|
190
|
+
int64_t offset,
|
|
191
|
+
float* recons) const {
|
|
192
|
+
IndexIVFPQ::reconstruct_from_offset(list_no, offset, recons);
|
|
191
193
|
|
|
192
|
-
idx_t id = invlists->get_single_id
|
|
193
|
-
assert
|
|
194
|
+
idx_t id = invlists->get_single_id(list_no, offset);
|
|
195
|
+
assert(0 <= id && id < ntotal);
|
|
194
196
|
|
|
195
197
|
std::vector<float> r3(d);
|
|
196
|
-
refine_pq.decode
|
|
198
|
+
refine_pq.decode(&refine_codes[id * refine_pq.code_size], r3.data());
|
|
197
199
|
for (int i = 0; i < d; ++i) {
|
|
198
|
-
|
|
200
|
+
recons[i] += r3[i];
|
|
199
201
|
}
|
|
200
202
|
}
|
|
201
203
|
|
|
202
|
-
void IndexIVFPQR::merge_from
|
|
203
|
-
|
|
204
|
-
IndexIVFPQR *other = dynamic_cast<IndexIVFPQR *> (&other_in);
|
|
204
|
+
void IndexIVFPQR::merge_from(IndexIVF& other_in, idx_t add_id) {
|
|
205
|
+
IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&other_in);
|
|
205
206
|
FAISS_THROW_IF_NOT(other);
|
|
206
207
|
|
|
207
|
-
IndexIVF::merge_from
|
|
208
|
+
IndexIVF::merge_from(other_in, add_id);
|
|
208
209
|
|
|
209
|
-
refine_codes.insert
|
|
210
|
-
|
|
211
|
-
|
|
210
|
+
refine_codes.insert(
|
|
211
|
+
refine_codes.end(),
|
|
212
|
+
other->refine_codes.begin(),
|
|
213
|
+
other->refine_codes.end());
|
|
212
214
|
other->refine_codes.clear();
|
|
213
215
|
}
|
|
214
216
|
|
|
215
217
|
size_t IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
|
|
216
|
-
|
|
217
|
-
|
|
218
|
+
FAISS_THROW_MSG("not implemented");
|
|
219
|
+
return 0;
|
|
218
220
|
}
|
|
219
221
|
|
|
220
222
|
} // namespace faiss
|
|
@@ -13,23 +13,24 @@
|
|
|
13
13
|
|
|
14
14
|
#include <faiss/IndexIVFPQ.h>
|
|
15
15
|
|
|
16
|
-
|
|
17
16
|
namespace faiss {
|
|
18
17
|
|
|
19
|
-
|
|
20
|
-
|
|
21
18
|
/** Index with an additional level of PQ refinement */
|
|
22
|
-
struct IndexIVFPQR: IndexIVFPQ {
|
|
23
|
-
ProductQuantizer refine_pq;
|
|
24
|
-
std::vector
|
|
19
|
+
struct IndexIVFPQR : IndexIVFPQ {
|
|
20
|
+
ProductQuantizer refine_pq; ///< 3rd level quantizer
|
|
21
|
+
std::vector<uint8_t> refine_codes; ///< corresponding codes
|
|
25
22
|
|
|
26
23
|
/// factor between k requested in search and the k requested from the IVFPQ
|
|
27
24
|
float k_factor;
|
|
28
25
|
|
|
29
|
-
IndexIVFPQR
|
|
30
|
-
Index
|
|
31
|
-
size_t
|
|
32
|
-
size_t
|
|
26
|
+
IndexIVFPQR(
|
|
27
|
+
Index* quantizer,
|
|
28
|
+
size_t d,
|
|
29
|
+
size_t nlist,
|
|
30
|
+
size_t M,
|
|
31
|
+
size_t nbits_per_idx,
|
|
32
|
+
size_t M_refine,
|
|
33
|
+
size_t nbits_per_idx_refine);
|
|
33
34
|
|
|
34
35
|
void reset() override;
|
|
35
36
|
|
|
@@ -41,26 +42,30 @@ struct IndexIVFPQR: IndexIVFPQ {
|
|
|
41
42
|
void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
|
|
42
43
|
|
|
43
44
|
/// same as add_with_ids, but optionally use the precomputed list ids
|
|
44
|
-
void add_core
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
void
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
void
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
45
|
+
void add_core(
|
|
46
|
+
idx_t n,
|
|
47
|
+
const float* x,
|
|
48
|
+
const idx_t* xids,
|
|
49
|
+
const idx_t* precomputed_idx) override;
|
|
50
|
+
|
|
51
|
+
void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
|
|
52
|
+
const override;
|
|
53
|
+
|
|
54
|
+
void merge_from(IndexIVF& other, idx_t add_id) override;
|
|
55
|
+
|
|
56
|
+
void search_preassigned(
|
|
57
|
+
idx_t n,
|
|
58
|
+
const float* x,
|
|
59
|
+
idx_t k,
|
|
60
|
+
const idx_t* assign,
|
|
61
|
+
const float* centroid_dis,
|
|
62
|
+
float* distances,
|
|
63
|
+
idx_t* labels,
|
|
64
|
+
bool store_pairs,
|
|
65
|
+
const IVFSearchParameters* params = nullptr,
|
|
66
|
+
IndexIVFStats* stats = nullptr) const override;
|
|
61
67
|
|
|
62
68
|
IndexIVFPQR();
|
|
63
69
|
};
|
|
64
70
|
|
|
65
|
-
|
|
66
71
|
} // namespace faiss
|