faiss 0.2.0 → 0.2.1
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/CHANGELOG.md +4 -0
- 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 +334 -195
- data/vendor/faiss/faiss/Clustering.h +88 -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 +139 -164
- data/vendor/faiss/faiss/Index2Layer.h +22 -22
- 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 +116 -147
- data/vendor/faiss/faiss/IndexFlat.h +35 -46
- data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
- data/vendor/faiss/faiss/IndexHNSW.h +57 -41
- data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
- data/vendor/faiss/faiss/IndexIVF.h +146 -113
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
- data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
- data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
- 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 +125 -133
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
- data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
- data/vendor/faiss/faiss/IndexLSH.h +21 -26
- data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
- data/vendor/faiss/faiss/IndexLattice.h +11 -16
- data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
- data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
- data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
- data/vendor/faiss/faiss/IndexNSG.h +85 -0
- data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
- data/vendor/faiss/faiss/IndexPQ.h +64 -67
- 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 +115 -131
- data/vendor/faiss/faiss/IndexRefine.h +22 -23
- data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
- data/vendor/faiss/faiss/IndexReplicas.h +62 -56
- data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
- data/vendor/faiss/faiss/IndexResidual.h +152 -0
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
- 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 +652 -474
- data/vendor/faiss/faiss/VectorTransform.h +61 -89
- data/vendor/faiss/faiss/clone_index.cpp +77 -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 +197 -170
- 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/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 +270 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -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 +595 -611
- data/vendor/faiss/faiss/impl/HNSW.h +179 -200
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -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 +682 -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 +448 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
- 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 +547 -479
- data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
- data/vendor/faiss/faiss/impl/io.cpp +75 -94
- data/vendor/faiss/faiss/impl/io.h +31 -41
- data/vendor/faiss/faiss/impl/io_macros.h +40 -29
- 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 +269 -218
- data/vendor/faiss/faiss/index_factory.h +6 -7
- 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 +301 -310
- data/vendor/faiss/faiss/utils/distances.h +133 -118
- data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
- 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 +53 -48
- metadata +20 -2
@@ -11,28 +11,23 @@
|
|
11
11
|
|
12
12
|
#pragma once
|
13
13
|
|
14
|
-
|
15
|
-
|
16
14
|
namespace faiss {
|
17
15
|
|
18
16
|
struct Index;
|
19
17
|
struct IndexIVF;
|
20
18
|
struct VectorTransform;
|
21
19
|
|
22
|
-
|
23
20
|
/* cloning functions */
|
24
|
-
Index
|
21
|
+
Index* clone_index(const Index*);
|
25
22
|
|
26
23
|
/** Cloner class, useful to override classes with other cloning
|
27
24
|
* functions. The cloning function above just calls
|
28
25
|
* Cloner::clone_Index. */
|
29
26
|
struct Cloner {
|
30
|
-
virtual VectorTransform
|
31
|
-
virtual Index
|
32
|
-
virtual IndexIVF
|
27
|
+
virtual VectorTransform* clone_VectorTransform(const VectorTransform*);
|
28
|
+
virtual Index* clone_Index(const Index*);
|
29
|
+
virtual IndexIVF* clone_IndexIVF(const IndexIVF*);
|
33
30
|
virtual ~Cloner() {}
|
34
31
|
};
|
35
32
|
|
36
|
-
|
37
|
-
|
38
33
|
} // namespace faiss
|
@@ -8,19 +8,19 @@
|
|
8
8
|
#include <faiss/gpu/GpuAutoTune.h>
|
9
9
|
#include <typeinfo>
|
10
10
|
|
11
|
-
#include <faiss/
|
11
|
+
#include <faiss/IndexPreTransform.h>
|
12
12
|
#include <faiss/IndexReplicas.h>
|
13
13
|
#include <faiss/IndexShards.h>
|
14
|
-
#include <faiss/
|
14
|
+
#include <faiss/gpu/GpuIndex.h>
|
15
15
|
#include <faiss/gpu/GpuIndexFlat.h>
|
16
16
|
#include <faiss/gpu/GpuIndexIVFFlat.h>
|
17
17
|
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
18
18
|
#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
|
19
|
-
#include <faiss/impl/FaissAssert.h>
|
20
19
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
20
|
+
#include <faiss/impl/FaissAssert.h>
|
21
21
|
|
22
|
-
namespace faiss {
|
23
|
-
|
22
|
+
namespace faiss {
|
23
|
+
namespace gpu {
|
24
24
|
|
25
25
|
using namespace ::faiss;
|
26
26
|
|
@@ -28,68 +28,63 @@ using namespace ::faiss;
|
|
28
28
|
* Parameters to auto-tune on GpuIndex'es
|
29
29
|
**********************************************************/
|
30
30
|
|
31
|
-
#define DC(classname) auto ix = dynamic_cast<const classname
|
31
|
+
#define DC(classname) auto ix = dynamic_cast<const classname*>(index)
|
32
32
|
|
33
|
-
|
34
|
-
|
35
|
-
{
|
36
|
-
if (DC (IndexPreTransform)) {
|
33
|
+
void GpuParameterSpace::initialize(const Index* index) {
|
34
|
+
if (DC(IndexPreTransform)) {
|
37
35
|
index = ix->index;
|
38
36
|
}
|
39
|
-
if (DC
|
40
|
-
if (ix->count() == 0)
|
37
|
+
if (DC(IndexReplicas)) {
|
38
|
+
if (ix->count() == 0)
|
39
|
+
return;
|
41
40
|
index = ix->at(0);
|
42
41
|
}
|
43
|
-
if (DC
|
44
|
-
if (ix->count() == 0)
|
42
|
+
if (DC(IndexShards)) {
|
43
|
+
if (ix->count() == 0)
|
44
|
+
return;
|
45
45
|
index = ix->at(0);
|
46
46
|
}
|
47
|
-
if (DC
|
48
|
-
ParameterRange
|
47
|
+
if (DC(GpuIndexIVF)) {
|
48
|
+
ParameterRange& pr = add_range("nprobe");
|
49
49
|
for (int i = 0; i < 12; i++) {
|
50
50
|
size_t nprobe = 1 << i;
|
51
|
-
if (nprobe >= ix->getNumLists() ||
|
52
|
-
|
53
|
-
pr.values.push_back
|
51
|
+
if (nprobe >= ix->getNumLists() || nprobe > getMaxKSelection())
|
52
|
+
break;
|
53
|
+
pr.values.push_back(nprobe);
|
54
54
|
}
|
55
55
|
}
|
56
56
|
// not sure we should call the parent initializer
|
57
57
|
}
|
58
58
|
|
59
|
-
|
60
|
-
|
61
59
|
#undef DC
|
62
60
|
// non-const version
|
63
|
-
#define DC(classname) auto
|
64
|
-
|
61
|
+
#define DC(classname) auto* ix = dynamic_cast<classname*>(index)
|
65
62
|
|
66
|
-
|
67
|
-
|
68
|
-
|
69
|
-
{
|
70
|
-
if (DC
|
63
|
+
void GpuParameterSpace::set_index_parameter(
|
64
|
+
Index* index,
|
65
|
+
const std::string& name,
|
66
|
+
double val) const {
|
67
|
+
if (DC(IndexReplicas)) {
|
71
68
|
for (int i = 0; i < ix->count(); i++)
|
72
|
-
set_index_parameter
|
69
|
+
set_index_parameter(ix->at(i), name, val);
|
73
70
|
return;
|
74
71
|
}
|
75
72
|
if (name == "nprobe") {
|
76
|
-
if (DC
|
77
|
-
ix->setNumProbes
|
73
|
+
if (DC(GpuIndexIVF)) {
|
74
|
+
ix->setNumProbes(int(val));
|
78
75
|
return;
|
79
76
|
}
|
80
77
|
}
|
81
78
|
if (name == "use_precomputed_table") {
|
82
|
-
if (DC
|
83
|
-
ix->setPrecomputedCodes(bool
|
79
|
+
if (DC(GpuIndexIVFPQ)) {
|
80
|
+
ix->setPrecomputedCodes(bool(val));
|
84
81
|
return;
|
85
82
|
}
|
86
83
|
}
|
87
84
|
|
88
85
|
// maybe normal index parameters apply?
|
89
|
-
ParameterSpace::set_index_parameter
|
86
|
+
ParameterSpace::set_index_parameter(index, name, val);
|
90
87
|
}
|
91
88
|
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
} } // namespace
|
89
|
+
} // namespace gpu
|
90
|
+
} // namespace faiss
|
@@ -7,21 +7,23 @@
|
|
7
7
|
|
8
8
|
#pragma once
|
9
9
|
|
10
|
-
#include <faiss/Index.h>
|
11
10
|
#include <faiss/AutoTune.h>
|
11
|
+
#include <faiss/Index.h>
|
12
12
|
|
13
|
-
namespace faiss {
|
14
|
-
|
13
|
+
namespace faiss {
|
14
|
+
namespace gpu {
|
15
15
|
|
16
16
|
/// parameter space and setters for GPU indexes
|
17
|
-
struct GpuParameterSpace: faiss::ParameterSpace {
|
17
|
+
struct GpuParameterSpace : faiss::ParameterSpace {
|
18
18
|
/// initialize with reasonable parameters for the index
|
19
|
-
void initialize
|
19
|
+
void initialize(const faiss::Index* index) override;
|
20
20
|
|
21
21
|
/// set a combination of parameters on an index
|
22
|
-
void set_index_parameter
|
23
|
-
|
24
|
-
|
22
|
+
void set_index_parameter(
|
23
|
+
faiss::Index* index,
|
24
|
+
const std::string& name,
|
25
|
+
double val) const override;
|
25
26
|
};
|
26
27
|
|
27
|
-
}
|
28
|
+
} // namespace gpu
|
29
|
+
} // namespace faiss
|
@@ -5,51 +5,52 @@
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
6
6
|
*/
|
7
7
|
|
8
|
-
|
9
8
|
#include <faiss/gpu/GpuCloner.h>
|
9
|
+
#include <faiss/impl/FaissAssert.h>
|
10
10
|
#include <typeinfo>
|
11
11
|
|
12
|
-
#include <faiss/gpu/
|
13
|
-
|
14
|
-
#include <faiss/index_io.h>
|
12
|
+
#include <faiss/gpu/StandardGpuResources.h>
|
13
|
+
|
15
14
|
#include <faiss/IndexFlat.h>
|
16
15
|
#include <faiss/IndexIVF.h>
|
17
16
|
#include <faiss/IndexIVFFlat.h>
|
18
|
-
#include <faiss/IndexScalarQuantizer.h>
|
19
17
|
#include <faiss/IndexIVFPQ.h>
|
20
|
-
#include <faiss/IndexReplicas.h>
|
21
18
|
#include <faiss/IndexPreTransform.h>
|
19
|
+
#include <faiss/IndexReplicas.h>
|
20
|
+
#include <faiss/IndexScalarQuantizer.h>
|
22
21
|
#include <faiss/MetaIndexes.h>
|
22
|
+
#include <faiss/gpu/GpuIndex.h>
|
23
23
|
#include <faiss/gpu/GpuIndexFlat.h>
|
24
24
|
#include <faiss/gpu/GpuIndexIVFFlat.h>
|
25
25
|
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
26
26
|
#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
|
27
27
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
28
|
+
#include <faiss/impl/FaissAssert.h>
|
29
|
+
#include <faiss/index_io.h>
|
28
30
|
|
29
|
-
namespace faiss {
|
30
|
-
|
31
|
+
namespace faiss {
|
32
|
+
namespace gpu {
|
31
33
|
|
32
34
|
/**********************************************************
|
33
35
|
* Cloning to CPU
|
34
36
|
**********************************************************/
|
35
37
|
|
36
|
-
void ToCPUCloner::merge_index(Index
|
37
|
-
{
|
38
|
-
|
39
|
-
auto ifl2 = dynamic_cast<const IndexFlat *>(src);
|
38
|
+
void ToCPUCloner::merge_index(Index* dst, Index* src, bool successive_ids) {
|
39
|
+
if (auto ifl = dynamic_cast<IndexFlat*>(dst)) {
|
40
|
+
auto ifl2 = dynamic_cast<const IndexFlat*>(src);
|
40
41
|
FAISS_ASSERT(ifl2);
|
41
42
|
FAISS_ASSERT(successive_ids);
|
42
43
|
ifl->add(ifl2->ntotal, ifl2->xb.data());
|
43
|
-
} else if(auto ifl = dynamic_cast<IndexIVFFlat
|
44
|
-
auto ifl2 = dynamic_cast<IndexIVFFlat
|
44
|
+
} else if (auto ifl = dynamic_cast<IndexIVFFlat*>(dst)) {
|
45
|
+
auto ifl2 = dynamic_cast<IndexIVFFlat*>(src);
|
45
46
|
FAISS_ASSERT(ifl2);
|
46
47
|
ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
|
47
|
-
} else if(auto ifl = dynamic_cast<IndexIVFScalarQuantizer
|
48
|
-
auto ifl2 = dynamic_cast<IndexIVFScalarQuantizer
|
48
|
+
} else if (auto ifl = dynamic_cast<IndexIVFScalarQuantizer*>(dst)) {
|
49
|
+
auto ifl2 = dynamic_cast<IndexIVFScalarQuantizer*>(src);
|
49
50
|
FAISS_ASSERT(ifl2);
|
50
51
|
ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
|
51
|
-
} else if(auto ifl = dynamic_cast<IndexIVFPQ
|
52
|
-
auto ifl2 = dynamic_cast<IndexIVFPQ
|
52
|
+
} else if (auto ifl = dynamic_cast<IndexIVFPQ*>(dst)) {
|
53
|
+
auto ifl2 = dynamic_cast<IndexIVFPQ*>(src);
|
53
54
|
FAISS_ASSERT(ifl2);
|
54
55
|
ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
|
55
56
|
} else {
|
@@ -57,24 +58,22 @@ void ToCPUCloner::merge_index(Index *dst, Index *src, bool successive_ids)
|
|
57
58
|
}
|
58
59
|
}
|
59
60
|
|
60
|
-
|
61
|
-
|
62
|
-
|
63
|
-
if(auto ifl = dynamic_cast<const GpuIndexFlat *>(index)) {
|
64
|
-
IndexFlat *res = new IndexFlat();
|
61
|
+
Index* ToCPUCloner::clone_Index(const Index* index) {
|
62
|
+
if (auto ifl = dynamic_cast<const GpuIndexFlat*>(index)) {
|
63
|
+
IndexFlat* res = new IndexFlat();
|
65
64
|
ifl->copyTo(res);
|
66
65
|
return res;
|
67
|
-
} else if(auto ifl = dynamic_cast<const GpuIndexIVFFlat
|
68
|
-
IndexIVFFlat
|
66
|
+
} else if (auto ifl = dynamic_cast<const GpuIndexIVFFlat*>(index)) {
|
67
|
+
IndexIVFFlat* res = new IndexIVFFlat();
|
69
68
|
ifl->copyTo(res);
|
70
69
|
return res;
|
71
|
-
} else if(
|
72
|
-
|
73
|
-
IndexIVFScalarQuantizer
|
70
|
+
} else if (
|
71
|
+
auto ifl = dynamic_cast<const GpuIndexIVFScalarQuantizer*>(index)) {
|
72
|
+
IndexIVFScalarQuantizer* res = new IndexIVFScalarQuantizer();
|
74
73
|
ifl->copyTo(res);
|
75
74
|
return res;
|
76
|
-
} else if(auto ipq = dynamic_cast<const GpuIndexIVFPQ
|
77
|
-
IndexIVFPQ
|
75
|
+
} else if (auto ipq = dynamic_cast<const GpuIndexIVFPQ*>(index)) {
|
76
|
+
IndexIVFPQ* res = new IndexIVFPQ();
|
78
77
|
ipq->copyTo(res);
|
79
78
|
return res;
|
80
79
|
|
@@ -82,17 +81,17 @@ Index *ToCPUCloner::clone_Index(const Index *index)
|
|
82
81
|
// objective is to make a single component out of them
|
83
82
|
// (inverse op of ToGpuClonerMultiple)
|
84
83
|
|
85
|
-
} else if(auto ish = dynamic_cast<const IndexShards
|
84
|
+
} else if (auto ish = dynamic_cast<const IndexShards*>(index)) {
|
86
85
|
int nshard = ish->count();
|
87
86
|
FAISS_ASSERT(nshard > 0);
|
88
|
-
Index
|
89
|
-
for(int i = 1; i < ish->count(); i++) {
|
90
|
-
Index
|
87
|
+
Index* res = clone_Index(ish->at(0));
|
88
|
+
for (int i = 1; i < ish->count(); i++) {
|
89
|
+
Index* res_i = clone_Index(ish->at(i));
|
91
90
|
merge_index(res, res_i, ish->successive_ids);
|
92
91
|
delete res_i;
|
93
92
|
}
|
94
93
|
return res;
|
95
|
-
} else if(auto ipr = dynamic_cast<const IndexReplicas
|
94
|
+
} else if (auto ipr = dynamic_cast<const IndexReplicas*>(index)) {
|
96
95
|
// just clone one of the replicas
|
97
96
|
FAISS_ASSERT(ipr->count() > 0);
|
98
97
|
return clone_Index(ipr->at(0));
|
@@ -101,81 +100,97 @@ Index *ToCPUCloner::clone_Index(const Index *index)
|
|
101
100
|
}
|
102
101
|
}
|
103
102
|
|
104
|
-
faiss::Index
|
105
|
-
{
|
103
|
+
faiss::Index* index_gpu_to_cpu(const faiss::Index* gpu_index) {
|
106
104
|
ToCPUCloner cl;
|
107
105
|
return cl.clone_Index(gpu_index);
|
108
106
|
}
|
109
107
|
|
110
|
-
|
111
|
-
|
112
|
-
|
113
108
|
/**********************************************************
|
114
109
|
* Cloning to 1 GPU
|
115
110
|
**********************************************************/
|
116
111
|
|
117
|
-
ToGpuCloner::ToGpuCloner(
|
118
|
-
|
119
|
-
|
120
|
-
|
112
|
+
ToGpuCloner::ToGpuCloner(
|
113
|
+
GpuResourcesProvider* prov,
|
114
|
+
int device,
|
115
|
+
const GpuClonerOptions& options)
|
116
|
+
: GpuClonerOptions(options), provider(prov), device(device) {}
|
121
117
|
|
122
|
-
Index
|
123
|
-
|
124
|
-
if(auto ifl = dynamic_cast<const IndexFlat
|
118
|
+
Index* ToGpuCloner::clone_Index(const Index* index) {
|
119
|
+
using idx_t = Index::idx_t;
|
120
|
+
if (auto ifl = dynamic_cast<const IndexFlat*>(index)) {
|
125
121
|
GpuIndexFlatConfig config;
|
126
122
|
config.device = device;
|
127
123
|
config.useFloat16 = useFloat16;
|
128
124
|
config.storeTransposed = storeTransposed;
|
129
|
-
|
130
125
|
return new GpuIndexFlat(provider, ifl, config);
|
131
|
-
} else if
|
126
|
+
} else if (
|
127
|
+
dynamic_cast<const IndexScalarQuantizer*>(index) &&
|
128
|
+
static_cast<const IndexScalarQuantizer*>(index)->sq.qtype ==
|
129
|
+
ScalarQuantizer::QT_fp16) {
|
130
|
+
GpuIndexFlatConfig config;
|
131
|
+
config.device = device;
|
132
|
+
config.useFloat16 = true;
|
133
|
+
GpuIndexFlat* gif = new GpuIndexFlat(
|
134
|
+
provider, index->d, index->metric_type, config);
|
135
|
+
// transfer data by blocks
|
136
|
+
idx_t bs = 1024 * 1024;
|
137
|
+
for (idx_t i0 = 0; i0 < index->ntotal; i0 += bs) {
|
138
|
+
idx_t i1 = std::min(i0 + bs, index->ntotal);
|
139
|
+
std::vector<float> buffer((i1 - i0) * index->d);
|
140
|
+
index->reconstruct_n(i0, i1 - i0, buffer.data());
|
141
|
+
gif->add(i1 - i0, buffer.data());
|
142
|
+
}
|
143
|
+
assert(gif->getNumVecs() == index->ntotal);
|
144
|
+
return gif;
|
145
|
+
} else if (auto ifl = dynamic_cast<const faiss::IndexIVFFlat*>(index)) {
|
132
146
|
GpuIndexIVFFlatConfig config;
|
133
147
|
config.device = device;
|
134
148
|
config.indicesOptions = indicesOptions;
|
135
149
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
136
150
|
config.flatConfig.storeTransposed = storeTransposed;
|
137
151
|
|
138
|
-
GpuIndexIVFFlat
|
139
|
-
|
140
|
-
|
141
|
-
ifl->nlist,
|
142
|
-
ifl->metric_type,
|
143
|
-
config);
|
144
|
-
if(reserveVecs > 0 && ifl->ntotal == 0) {
|
152
|
+
GpuIndexIVFFlat* res = new GpuIndexIVFFlat(
|
153
|
+
provider, ifl->d, ifl->nlist, ifl->metric_type, config);
|
154
|
+
if (reserveVecs > 0 && ifl->ntotal == 0) {
|
145
155
|
res->reserveMemory(reserveVecs);
|
146
156
|
}
|
147
157
|
|
148
158
|
res->copyFrom(ifl);
|
149
159
|
return res;
|
150
|
-
} else if(
|
151
|
-
|
160
|
+
} else if (
|
161
|
+
auto ifl = dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(
|
162
|
+
index)) {
|
152
163
|
GpuIndexIVFScalarQuantizerConfig config;
|
153
164
|
config.device = device;
|
154
165
|
config.indicesOptions = indicesOptions;
|
155
166
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
156
167
|
config.flatConfig.storeTransposed = storeTransposed;
|
157
168
|
|
158
|
-
GpuIndexIVFScalarQuantizer
|
159
|
-
|
160
|
-
|
161
|
-
|
162
|
-
|
163
|
-
|
164
|
-
|
165
|
-
|
166
|
-
if(reserveVecs > 0 && ifl->ntotal == 0) {
|
169
|
+
GpuIndexIVFScalarQuantizer* res = new GpuIndexIVFScalarQuantizer(
|
170
|
+
provider,
|
171
|
+
ifl->d,
|
172
|
+
ifl->nlist,
|
173
|
+
ifl->sq.qtype,
|
174
|
+
ifl->metric_type,
|
175
|
+
ifl->by_residual,
|
176
|
+
config);
|
177
|
+
if (reserveVecs > 0 && ifl->ntotal == 0) {
|
167
178
|
res->reserveMemory(reserveVecs);
|
168
179
|
}
|
169
180
|
|
170
181
|
res->copyFrom(ifl);
|
171
182
|
return res;
|
172
|
-
} else if(auto ipq = dynamic_cast<const faiss::IndexIVFPQ
|
173
|
-
if(verbose)
|
183
|
+
} else if (auto ipq = dynamic_cast<const faiss::IndexIVFPQ*>(index)) {
|
184
|
+
if (verbose) {
|
174
185
|
printf(" IndexIVFPQ size %ld -> GpuIndexIVFPQ "
|
175
186
|
"indicesOptions=%d "
|
176
187
|
"usePrecomputed=%d useFloat16=%d reserveVecs=%ld\n",
|
177
|
-
ipq->ntotal,
|
178
|
-
|
188
|
+
ipq->ntotal,
|
189
|
+
indicesOptions,
|
190
|
+
usePrecomputed,
|
191
|
+
useFloat16,
|
192
|
+
reserveVecs);
|
193
|
+
}
|
179
194
|
GpuIndexIVFPQConfig config;
|
180
195
|
config.device = device;
|
181
196
|
config.indicesOptions = indicesOptions;
|
@@ -184,154 +199,144 @@ Index *ToGpuCloner::clone_Index(const Index *index)
|
|
184
199
|
config.useFloat16LookupTables = useFloat16;
|
185
200
|
config.usePrecomputedTables = usePrecomputed;
|
186
201
|
|
187
|
-
GpuIndexIVFPQ
|
202
|
+
GpuIndexIVFPQ* res = new GpuIndexIVFPQ(provider, ipq, config);
|
188
203
|
|
189
|
-
if(reserveVecs > 0 && ipq->ntotal == 0) {
|
204
|
+
if (reserveVecs > 0 && ipq->ntotal == 0) {
|
190
205
|
res->reserveMemory(reserveVecs);
|
191
206
|
}
|
192
207
|
|
193
208
|
return res;
|
194
209
|
} else {
|
210
|
+
// default: use CPU cloner
|
195
211
|
return Cloner::clone_Index(index);
|
196
212
|
}
|
197
213
|
}
|
198
214
|
|
199
|
-
|
200
|
-
|
201
|
-
|
202
|
-
|
203
|
-
|
204
|
-
{
|
215
|
+
faiss::Index* index_cpu_to_gpu(
|
216
|
+
GpuResourcesProvider* provider,
|
217
|
+
int device,
|
218
|
+
const faiss::Index* index,
|
219
|
+
const GpuClonerOptions* options) {
|
205
220
|
GpuClonerOptions defaults;
|
206
221
|
ToGpuCloner cl(provider, device, options ? *options : defaults);
|
207
222
|
return cl.clone_Index(index);
|
208
223
|
}
|
209
224
|
|
210
|
-
|
211
225
|
/**********************************************************
|
212
226
|
* Cloning to multiple GPUs
|
213
227
|
**********************************************************/
|
214
228
|
|
215
229
|
ToGpuClonerMultiple::ToGpuClonerMultiple(
|
216
|
-
|
217
|
-
|
218
|
-
|
219
|
-
|
220
|
-
{
|
230
|
+
std::vector<GpuResourcesProvider*>& provider,
|
231
|
+
std::vector<int>& devices,
|
232
|
+
const GpuMultipleClonerOptions& options)
|
233
|
+
: GpuMultipleClonerOptions(options) {
|
221
234
|
FAISS_ASSERT(provider.size() == devices.size());
|
222
|
-
for(int i = 0; i < provider.size(); i++) {
|
235
|
+
for (int i = 0; i < provider.size(); i++) {
|
223
236
|
sub_cloners.push_back(ToGpuCloner(provider[i], devices[i], options));
|
224
237
|
}
|
225
238
|
}
|
226
239
|
|
227
|
-
|
228
240
|
ToGpuClonerMultiple::ToGpuClonerMultiple(
|
229
|
-
|
230
|
-
|
231
|
-
|
232
|
-
|
233
|
-
|
234
|
-
|
235
|
-
|
236
|
-
|
237
|
-
|
238
|
-
long n, long i)
|
239
|
-
{
|
241
|
+
const std::vector<ToGpuCloner>& sub_cloners,
|
242
|
+
const GpuMultipleClonerOptions& options)
|
243
|
+
: GpuMultipleClonerOptions(options), sub_cloners(sub_cloners) {}
|
244
|
+
|
245
|
+
void ToGpuClonerMultiple::copy_ivf_shard(
|
246
|
+
const IndexIVF* index_ivf,
|
247
|
+
IndexIVF* idx2,
|
248
|
+
long n,
|
249
|
+
long i) {
|
240
250
|
if (shard_type == 2) {
|
241
251
|
long i0 = i * index_ivf->ntotal / n;
|
242
252
|
long i1 = (i + 1) * index_ivf->ntotal / n;
|
243
253
|
|
244
|
-
if(verbose)
|
245
|
-
printf("IndexShards shard %ld indices %ld:%ld\n",
|
246
|
-
i, i0, i1);
|
254
|
+
if (verbose)
|
255
|
+
printf("IndexShards shard %ld indices %ld:%ld\n", i, i0, i1);
|
247
256
|
index_ivf->copy_subset_to(*idx2, 2, i0, i1);
|
248
257
|
FAISS_ASSERT(idx2->ntotal == i1 - i0);
|
249
258
|
} else if (shard_type == 1) {
|
250
|
-
if(verbose)
|
251
|
-
printf("IndexShards shard %ld select modulo %ld = %ld\n",
|
252
|
-
i, n, i);
|
259
|
+
if (verbose)
|
260
|
+
printf("IndexShards shard %ld select modulo %ld = %ld\n", i, n, i);
|
253
261
|
index_ivf->copy_subset_to(*idx2, 1, n, i);
|
254
262
|
} else {
|
255
|
-
FAISS_THROW_FMT
|
263
|
+
FAISS_THROW_FMT("shard_type %d not implemented", shard_type);
|
256
264
|
}
|
257
|
-
|
258
265
|
}
|
259
266
|
|
260
|
-
Index
|
261
|
-
{
|
267
|
+
Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
|
262
268
|
long n = sub_cloners.size();
|
263
269
|
|
264
|
-
auto index_ivfpq =
|
265
|
-
|
266
|
-
auto index_ivfflat =
|
267
|
-
dynamic_cast<const faiss::IndexIVFFlat *>(index);
|
270
|
+
auto index_ivfpq = dynamic_cast<const faiss::IndexIVFPQ*>(index);
|
271
|
+
auto index_ivfflat = dynamic_cast<const faiss::IndexIVFFlat*>(index);
|
268
272
|
auto index_ivfsq =
|
269
|
-
|
270
|
-
auto index_flat =
|
271
|
-
|
272
|
-
|
273
|
-
|
274
|
-
|
275
|
-
|
276
|
-
"IndexFlat and IndexIVFPQ");
|
273
|
+
dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(index);
|
274
|
+
auto index_flat = dynamic_cast<const faiss::IndexFlat*>(index);
|
275
|
+
FAISS_THROW_IF_NOT_MSG(
|
276
|
+
index_ivfpq || index_ivfflat || index_flat || index_ivfsq,
|
277
|
+
"IndexShards implemented only for "
|
278
|
+
"IndexIVFFlat, IndexIVFScalarQuantizer, "
|
279
|
+
"IndexFlat and IndexIVFPQ");
|
277
280
|
|
278
281
|
std::vector<faiss::Index*> shards(n);
|
279
282
|
|
280
|
-
for(long i = 0; i < n; i++) {
|
283
|
+
for (long i = 0; i < n; i++) {
|
281
284
|
// make a shallow copy
|
282
|
-
if(reserveVecs)
|
283
|
-
sub_cloners[i].reserveVecs =
|
284
|
-
(reserveVecs + n - 1) / n;
|
285
|
+
if (reserveVecs)
|
286
|
+
sub_cloners[i].reserveVecs = (reserveVecs + n - 1) / n;
|
285
287
|
|
286
288
|
if (index_ivfpq) {
|
287
289
|
faiss::IndexIVFPQ idx2(
|
288
|
-
|
289
|
-
|
290
|
-
|
290
|
+
index_ivfpq->quantizer,
|
291
|
+
index_ivfpq->d,
|
292
|
+
index_ivfpq->nlist,
|
293
|
+
index_ivfpq->code_size,
|
294
|
+
index_ivfpq->pq.nbits);
|
291
295
|
idx2.metric_type = index_ivfpq->metric_type;
|
292
296
|
idx2.pq = index_ivfpq->pq;
|
293
297
|
idx2.nprobe = index_ivfpq->nprobe;
|
294
298
|
idx2.use_precomputed_table = 0;
|
295
299
|
idx2.is_trained = index->is_trained;
|
296
|
-
copy_ivf_shard
|
300
|
+
copy_ivf_shard(index_ivfpq, &idx2, n, i);
|
297
301
|
shards[i] = sub_cloners[i].clone_Index(&idx2);
|
298
302
|
} else if (index_ivfflat) {
|
299
303
|
faiss::IndexIVFFlat idx2(
|
300
|
-
|
301
|
-
|
304
|
+
index_ivfflat->quantizer,
|
305
|
+
index->d,
|
306
|
+
index_ivfflat->nlist,
|
307
|
+
index_ivfflat->metric_type);
|
302
308
|
idx2.nprobe = index_ivfflat->nprobe;
|
303
309
|
idx2.is_trained = index->is_trained;
|
304
|
-
copy_ivf_shard
|
310
|
+
copy_ivf_shard(index_ivfflat, &idx2, n, i);
|
305
311
|
shards[i] = sub_cloners[i].clone_Index(&idx2);
|
306
312
|
} else if (index_ivfsq) {
|
307
313
|
faiss::IndexIVFScalarQuantizer idx2(
|
308
|
-
|
309
|
-
|
310
|
-
|
311
|
-
|
314
|
+
index_ivfsq->quantizer,
|
315
|
+
index->d,
|
316
|
+
index_ivfsq->nlist,
|
317
|
+
index_ivfsq->sq.qtype,
|
318
|
+
index_ivfsq->metric_type,
|
319
|
+
index_ivfsq->by_residual);
|
312
320
|
|
313
321
|
idx2.nprobe = index_ivfsq->nprobe;
|
314
322
|
idx2.is_trained = index->is_trained;
|
315
323
|
idx2.sq = index_ivfsq->sq;
|
316
|
-
copy_ivf_shard
|
324
|
+
copy_ivf_shard(index_ivfsq, &idx2, n, i);
|
317
325
|
shards[i] = sub_cloners[i].clone_Index(&idx2);
|
318
326
|
} else if (index_flat) {
|
319
|
-
faiss::IndexFlat idx2
|
320
|
-
index->d, index->metric_type);
|
327
|
+
faiss::IndexFlat idx2(index->d, index->metric_type);
|
321
328
|
shards[i] = sub_cloners[i].clone_Index(&idx2);
|
322
329
|
if (index->ntotal > 0) {
|
323
330
|
long i0 = index->ntotal * i / n;
|
324
331
|
long i1 = index->ntotal * (i + 1) / n;
|
325
|
-
shards[i]->add
|
326
|
-
index_flat->xb.data() + i0 * index->d);
|
332
|
+
shards[i]->add(i1 - i0, index_flat->xb.data() + i0 * index->d);
|
327
333
|
}
|
328
334
|
}
|
329
335
|
}
|
330
336
|
|
331
337
|
bool successive_ids = index_flat != nullptr;
|
332
|
-
faiss::IndexShards
|
333
|
-
|
334
|
-
successive_ids);
|
338
|
+
faiss::IndexShards* res =
|
339
|
+
new faiss::IndexShards(index->d, true, successive_ids);
|
335
340
|
|
336
341
|
for (int i = 0; i < n; i++) {
|
337
342
|
res->add_shard(shards[i]);
|
@@ -341,33 +346,32 @@ Index * ToGpuClonerMultiple::clone_Index_to_shards (const Index *index)
|
|
341
346
|
return res;
|
342
347
|
}
|
343
348
|
|
344
|
-
Index
|
345
|
-
{
|
349
|
+
Index* ToGpuClonerMultiple::clone_Index(const Index* index) {
|
346
350
|
long n = sub_cloners.size();
|
347
351
|
if (n == 1)
|
348
352
|
return sub_cloners[0].clone_Index(index);
|
349
353
|
|
350
|
-
if(dynamic_cast<const IndexFlat
|
351
|
-
|
352
|
-
|
353
|
-
|
354
|
-
if(!shard) {
|
355
|
-
IndexReplicas
|
356
|
-
for(auto
|
354
|
+
if (dynamic_cast<const IndexFlat*>(index) ||
|
355
|
+
dynamic_cast<const faiss::IndexIVFFlat*>(index) ||
|
356
|
+
dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(index) ||
|
357
|
+
dynamic_cast<const faiss::IndexIVFPQ*>(index)) {
|
358
|
+
if (!shard) {
|
359
|
+
IndexReplicas* res = new IndexReplicas();
|
360
|
+
for (auto& sub_cloner : sub_cloners) {
|
357
361
|
res->addIndex(sub_cloner.clone_Index(index));
|
358
362
|
}
|
359
363
|
res->own_fields = true;
|
360
364
|
return res;
|
361
365
|
} else {
|
362
|
-
return clone_Index_to_shards
|
366
|
+
return clone_Index_to_shards(index);
|
363
367
|
}
|
364
|
-
} else if(auto miq = dynamic_cast<const MultiIndexQuantizer
|
368
|
+
} else if (auto miq = dynamic_cast<const MultiIndexQuantizer*>(index)) {
|
365
369
|
if (verbose) {
|
366
370
|
printf("cloning MultiIndexQuantizer: "
|
367
371
|
"will be valid only for search k=1\n");
|
368
372
|
}
|
369
|
-
const ProductQuantizer
|
370
|
-
IndexSplitVectors
|
373
|
+
const ProductQuantizer& pq = miq->pq;
|
374
|
+
IndexSplitVectors* splitv = new IndexSplitVectors(pq.d, true);
|
371
375
|
splitv->own_fields = true;
|
372
376
|
|
373
377
|
for (int m = 0; m < pq.M; m++) {
|
@@ -377,12 +381,13 @@ Index *ToGpuClonerMultiple::clone_Index(const Index *index)
|
|
377
381
|
long i1 = pq.M <= n ? (m + 1) * n / pq.M : i0 + 1;
|
378
382
|
std::vector<ToGpuCloner> sub_cloners_2;
|
379
383
|
sub_cloners_2.insert(
|
380
|
-
|
381
|
-
|
384
|
+
sub_cloners_2.begin(),
|
385
|
+
sub_cloners.begin() + i0,
|
386
|
+
sub_cloners.begin() + i1);
|
382
387
|
ToGpuClonerMultiple cm(sub_cloners_2, *this);
|
383
|
-
IndexFlatL2 idxc
|
384
|
-
idxc.add
|
385
|
-
Index
|
388
|
+
IndexFlatL2 idxc(pq.dsub);
|
389
|
+
idxc.add(pq.ksub, pq.centroids.data() + m * pq.d * pq.ksub);
|
390
|
+
Index* idx2 = cm.clone_Index(&idxc);
|
386
391
|
splitv->add_sub_index(idx2);
|
387
392
|
}
|
388
393
|
return splitv;
|
@@ -391,17 +396,39 @@ Index *ToGpuClonerMultiple::clone_Index(const Index *index)
|
|
391
396
|
}
|
392
397
|
}
|
393
398
|
|
394
|
-
|
395
|
-
|
396
|
-
|
397
|
-
|
398
|
-
|
399
|
-
const faiss::Index *index,
|
400
|
-
const GpuMultipleClonerOptions *options)
|
401
|
-
{
|
399
|
+
faiss::Index* index_cpu_to_gpu_multiple(
|
400
|
+
std::vector<GpuResourcesProvider*>& provider,
|
401
|
+
std::vector<int>& devices,
|
402
|
+
const faiss::Index* index,
|
403
|
+
const GpuMultipleClonerOptions* options) {
|
402
404
|
GpuMultipleClonerOptions defaults;
|
403
405
|
ToGpuClonerMultiple cl(provider, devices, options ? *options : defaults);
|
404
406
|
return cl.clone_Index(index);
|
405
407
|
}
|
406
408
|
|
407
|
-
|
409
|
+
GpuProgressiveDimIndexFactory::GpuProgressiveDimIndexFactory(int ngpu) {
|
410
|
+
FAISS_THROW_IF_NOT(ngpu >= 1);
|
411
|
+
devices.resize(ngpu);
|
412
|
+
vres.resize(ngpu);
|
413
|
+
|
414
|
+
for (int i = 0; i < ngpu; i++) {
|
415
|
+
vres[i] = new StandardGpuResources();
|
416
|
+
devices[i] = i;
|
417
|
+
}
|
418
|
+
ncall = 0;
|
419
|
+
}
|
420
|
+
|
421
|
+
GpuProgressiveDimIndexFactory::~GpuProgressiveDimIndexFactory() {
|
422
|
+
for (int i = 0; i < vres.size(); i++) {
|
423
|
+
delete vres[i];
|
424
|
+
}
|
425
|
+
}
|
426
|
+
|
427
|
+
Index* GpuProgressiveDimIndexFactory::operator()(int dim) {
|
428
|
+
IndexFlatL2 index(dim);
|
429
|
+
ncall++;
|
430
|
+
return index_cpu_to_gpu_multiple(vres, devices, &index, &options);
|
431
|
+
}
|
432
|
+
|
433
|
+
} // namespace gpu
|
434
|
+
} // namespace faiss
|