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
|
@@ -12,40 +12,46 @@
|
|
|
12
12
|
|
|
13
13
|
// Utilities for bit packing and unpacking CPU non-interleaved and GPU
|
|
14
14
|
// interleaved by 32 encodings
|
|
15
|
-
namespace faiss {
|
|
15
|
+
namespace faiss {
|
|
16
|
+
namespace gpu {
|
|
16
17
|
|
|
17
18
|
// Unpacks arbitrary bitwidth codes to a whole number of bytes per code
|
|
18
19
|
// The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
|
|
19
20
|
// (bit packed)
|
|
20
21
|
// The layout of the output is the same (byte packed to roundUp(bitsPerCode, 8)
|
|
21
22
|
// / 8 bytes)
|
|
22
|
-
std::vector<uint8_t> unpackNonInterleaved(
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
23
|
+
std::vector<uint8_t> unpackNonInterleaved(
|
|
24
|
+
std::vector<uint8_t> data,
|
|
25
|
+
int numVecs,
|
|
26
|
+
int dims,
|
|
27
|
+
int bitsPerCode);
|
|
26
28
|
|
|
27
29
|
// Unpacks arbitrary bitwidth codes to a whole number of bytes per scalar code
|
|
28
30
|
// The layout of the input is (v0 d0)(v1 d0) ... (v31 d0)(v0 d1) ...
|
|
29
31
|
// (bit packed)
|
|
30
32
|
// The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
|
|
31
33
|
// (byte packed)
|
|
32
|
-
std::vector<uint8_t> unpackInterleaved(
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
34
|
+
std::vector<uint8_t> unpackInterleaved(
|
|
35
|
+
std::vector<uint8_t> data,
|
|
36
|
+
int numVecs,
|
|
37
|
+
int dims,
|
|
38
|
+
int bitsPerCode);
|
|
36
39
|
|
|
37
40
|
// Packs data in the byte packed non-interleaved form to bit packed
|
|
38
41
|
// non-interleaved form
|
|
39
|
-
std::vector<uint8_t> packNonInterleaved(
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
42
|
+
std::vector<uint8_t> packNonInterleaved(
|
|
43
|
+
std::vector<uint8_t> data,
|
|
44
|
+
int numVecs,
|
|
45
|
+
int dims,
|
|
46
|
+
int bitsPerCode);
|
|
43
47
|
|
|
44
48
|
// Packs data in the byte packed non-interleaved form to bit packed
|
|
45
49
|
// interleaved form
|
|
46
|
-
std::vector<uint8_t> packInterleaved(
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
50
|
+
std::vector<uint8_t> packInterleaved(
|
|
51
|
+
std::vector<uint8_t> data,
|
|
52
|
+
int numVecs,
|
|
53
|
+
int dims,
|
|
54
|
+
int bitsPerCode);
|
|
55
|
+
|
|
56
|
+
} // namespace gpu
|
|
57
|
+
} // namespace faiss
|
|
@@ -5,39 +5,41 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#include <faiss/gpu/impl/RemapIndices.h>
|
|
10
9
|
#include <faiss/impl/FaissAssert.h>
|
|
11
10
|
|
|
12
|
-
namespace faiss {
|
|
11
|
+
namespace faiss {
|
|
12
|
+
namespace gpu {
|
|
13
13
|
|
|
14
14
|
// Utility function to translate (list id, offset) to a user index on
|
|
15
15
|
// the CPU. In a cpp in order to use OpenMP
|
|
16
16
|
void ivfOffsetToUserIndex(
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
17
|
+
Index::idx_t* indices,
|
|
18
|
+
int numLists,
|
|
19
|
+
int queries,
|
|
20
|
+
int k,
|
|
21
|
+
const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex) {
|
|
22
|
+
FAISS_ASSERT(numLists == listOffsetToUserIndex.size());
|
|
23
23
|
|
|
24
24
|
#pragma omp parallel for
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
25
|
+
for (int q = 0; q < queries; ++q) {
|
|
26
|
+
for (int r = 0; r < k; ++r) {
|
|
27
|
+
auto offsetIndex = indices[q * k + r];
|
|
28
28
|
|
|
29
|
-
|
|
29
|
+
if (offsetIndex < 0)
|
|
30
|
+
continue;
|
|
30
31
|
|
|
31
|
-
|
|
32
|
-
|
|
32
|
+
int listId = (int)(offsetIndex >> 32);
|
|
33
|
+
int listOffset = (int)(offsetIndex & 0xffffffff);
|
|
33
34
|
|
|
34
|
-
|
|
35
|
-
|
|
35
|
+
FAISS_ASSERT(listId < numLists);
|
|
36
|
+
auto& listIndices = listOffsetToUserIndex[listId];
|
|
36
37
|
|
|
37
|
-
|
|
38
|
-
|
|
38
|
+
FAISS_ASSERT(listOffset < listIndices.size());
|
|
39
|
+
indices[q * k + r] = listIndices[listOffset];
|
|
40
|
+
}
|
|
39
41
|
}
|
|
40
|
-
}
|
|
41
42
|
}
|
|
42
43
|
|
|
43
|
-
}
|
|
44
|
+
} // namespace gpu
|
|
45
|
+
} // namespace faiss
|
|
@@ -5,21 +5,22 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
11
10
|
#include <faiss/Index.h>
|
|
12
11
|
#include <vector>
|
|
13
12
|
|
|
14
|
-
namespace faiss {
|
|
13
|
+
namespace faiss {
|
|
14
|
+
namespace gpu {
|
|
15
15
|
|
|
16
16
|
/// Utility function to translate (list id, offset) to a user index on
|
|
17
17
|
/// the CPU. In a cpp in order to use OpenMP.
|
|
18
18
|
void ivfOffsetToUserIndex(
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
19
|
+
Index::idx_t* indices,
|
|
20
|
+
int numLists,
|
|
21
|
+
int queries,
|
|
22
|
+
int k,
|
|
23
|
+
const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex);
|
|
24
24
|
|
|
25
|
-
}
|
|
25
|
+
} // namespace gpu
|
|
26
|
+
} // namespace faiss
|
|
@@ -5,67 +5,62 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#include <faiss/impl/FaissAssert.h>
|
|
10
9
|
|
|
11
|
-
namespace faiss {
|
|
10
|
+
namespace faiss {
|
|
11
|
+
namespace gpu {
|
|
12
12
|
|
|
13
13
|
template <typename GpuIndex>
|
|
14
14
|
IndexWrapper<GpuIndex>::IndexWrapper(
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
15
|
+
int numGpus,
|
|
16
|
+
std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)>
|
|
17
|
+
init) {
|
|
18
|
+
FAISS_ASSERT(numGpus <= faiss::gpu::getNumDevices());
|
|
19
|
+
for (int i = 0; i < numGpus; ++i) {
|
|
20
|
+
auto res = std::unique_ptr<faiss::gpu::StandardGpuResources>(
|
|
21
|
+
new StandardGpuResources);
|
|
21
22
|
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
23
|
+
subIndex.emplace_back(init(res.get(), i));
|
|
24
|
+
resources.emplace_back(std::move(res));
|
|
25
|
+
}
|
|
25
26
|
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
27
|
+
if (numGpus > 1) {
|
|
28
|
+
// create proxy
|
|
29
|
+
replicaIndex =
|
|
30
|
+
std::unique_ptr<faiss::IndexReplicas>(new faiss::IndexReplicas);
|
|
30
31
|
|
|
31
|
-
|
|
32
|
-
|
|
32
|
+
for (auto& index : subIndex) {
|
|
33
|
+
replicaIndex->addIndex(index.get());
|
|
34
|
+
}
|
|
33
35
|
}
|
|
34
|
-
}
|
|
35
36
|
}
|
|
36
37
|
|
|
37
38
|
template <typename GpuIndex>
|
|
38
|
-
faiss::Index*
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
}
|
|
39
|
+
faiss::Index* IndexWrapper<GpuIndex>::getIndex() {
|
|
40
|
+
if ((bool)replicaIndex) {
|
|
41
|
+
return replicaIndex.get();
|
|
42
|
+
} else {
|
|
43
|
+
FAISS_ASSERT(!subIndex.empty());
|
|
44
|
+
return subIndex.front().get();
|
|
45
|
+
}
|
|
46
46
|
}
|
|
47
47
|
|
|
48
48
|
template <typename GpuIndex>
|
|
49
|
-
void
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
FAISS_ASSERT(!subIndex.empty());
|
|
59
|
-
f(subIndex.front().get());
|
|
60
|
-
}
|
|
49
|
+
void IndexWrapper<GpuIndex>::runOnIndices(std::function<void(GpuIndex*)> f) {
|
|
50
|
+
if ((bool)replicaIndex) {
|
|
51
|
+
replicaIndex->runOnIndex([f](int, faiss::Index* index) {
|
|
52
|
+
f(dynamic_cast<GpuIndex*>(index));
|
|
53
|
+
});
|
|
54
|
+
} else {
|
|
55
|
+
FAISS_ASSERT(!subIndex.empty());
|
|
56
|
+
f(subIndex.front().get());
|
|
57
|
+
}
|
|
61
58
|
}
|
|
62
59
|
|
|
63
60
|
template <typename GpuIndex>
|
|
64
|
-
void
|
|
65
|
-
|
|
66
|
-
runOnIndices([nprobe](GpuIndex* index) {
|
|
67
|
-
index->setNumProbes(nprobe);
|
|
68
|
-
});
|
|
61
|
+
void IndexWrapper<GpuIndex>::setNumProbes(int nprobe) {
|
|
62
|
+
runOnIndices([nprobe](GpuIndex* index) { index->setNumProbes(nprobe); });
|
|
69
63
|
}
|
|
70
64
|
|
|
71
|
-
}
|
|
65
|
+
} // namespace gpu
|
|
66
|
+
} // namespace faiss
|
|
@@ -5,7 +5,6 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
11
10
|
#include <faiss/IndexReplicas.h>
|
|
@@ -14,26 +13,29 @@
|
|
|
14
13
|
#include <memory>
|
|
15
14
|
#include <vector>
|
|
16
15
|
|
|
17
|
-
namespace faiss {
|
|
16
|
+
namespace faiss {
|
|
17
|
+
namespace gpu {
|
|
18
18
|
|
|
19
19
|
// If we want to run multi-GPU, create a proxy to wrap the indices.
|
|
20
20
|
// If we don't want multi-GPU, don't involve the proxy, so it doesn't
|
|
21
21
|
// affect the timings.
|
|
22
22
|
template <typename GpuIndex>
|
|
23
23
|
struct IndexWrapper {
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
24
|
+
std::vector<std::unique_ptr<faiss::gpu::StandardGpuResources>> resources;
|
|
25
|
+
std::vector<std::unique_ptr<GpuIndex>> subIndex;
|
|
26
|
+
std::unique_ptr<faiss::IndexReplicas> replicaIndex;
|
|
27
|
+
|
|
28
|
+
IndexWrapper(
|
|
29
|
+
int numGpus,
|
|
30
|
+
std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)>
|
|
31
|
+
init);
|
|
32
|
+
faiss::Index* getIndex();
|
|
33
|
+
|
|
34
|
+
void runOnIndices(std::function<void(GpuIndex*)> f);
|
|
35
|
+
void setNumProbes(int nprobe);
|
|
35
36
|
};
|
|
36
37
|
|
|
37
|
-
}
|
|
38
|
+
} // namespace gpu
|
|
39
|
+
} // namespace faiss
|
|
38
40
|
|
|
39
41
|
#include <faiss/gpu/perf/IndexWrapper-inl.h>
|
|
@@ -5,14 +5,13 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
|
-
#include <faiss/utils/random.h>
|
|
10
8
|
#include <faiss/Clustering.h>
|
|
11
9
|
#include <faiss/gpu/GpuIndexFlat.h>
|
|
12
10
|
#include <faiss/gpu/StandardGpuResources.h>
|
|
13
11
|
#include <faiss/gpu/perf/IndexWrapper.h>
|
|
14
12
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
|
15
13
|
#include <faiss/gpu/utils/Timer.h>
|
|
14
|
+
#include <faiss/utils/random.h>
|
|
16
15
|
#include <gflags/gflags.h>
|
|
17
16
|
#include <memory>
|
|
18
17
|
#include <vector>
|
|
@@ -29,87 +28,94 @@ DEFINE_bool(transposed, false, "transposed vector storage");
|
|
|
29
28
|
DEFINE_bool(verbose, false, "turn on clustering logging");
|
|
30
29
|
DEFINE_int64(seed, -1, "specify random seed");
|
|
31
30
|
DEFINE_int32(num_gpus, 1, "number of gpus to use");
|
|
32
|
-
DEFINE_int64(
|
|
31
|
+
DEFINE_int64(
|
|
32
|
+
min_paging_size,
|
|
33
|
+
-1,
|
|
34
|
+
"minimum size to use CPU -> GPU paged copies");
|
|
33
35
|
DEFINE_int64(pinned_mem, -1, "pinned memory allocation to use");
|
|
34
36
|
DEFINE_int32(max_points, -1, "max points per centroid");
|
|
35
37
|
|
|
36
38
|
using namespace faiss::gpu;
|
|
37
39
|
|
|
38
40
|
int main(int argc, char** argv) {
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
41
|
+
gflags::ParseCommandLineFlags(&argc, &argv, true);
|
|
42
|
+
|
|
43
|
+
cudaProfilerStop();
|
|
44
|
+
|
|
45
|
+
auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
|
|
46
|
+
printf("using seed %ld\n", seed);
|
|
47
|
+
|
|
48
|
+
std::vector<float> vecs((size_t)FLAGS_num * FLAGS_dim);
|
|
49
|
+
faiss::float_rand(vecs.data(), vecs.size(), seed);
|
|
50
|
+
|
|
51
|
+
printf("K-means metric %s dim %d centroids %d num train %d niter %d\n",
|
|
52
|
+
FLAGS_L2_metric ? "L2" : "IP",
|
|
53
|
+
FLAGS_dim,
|
|
54
|
+
FLAGS_k,
|
|
55
|
+
FLAGS_num,
|
|
56
|
+
FLAGS_niter);
|
|
57
|
+
printf("float16 math %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
|
|
58
|
+
printf("transposed storage %s\n",
|
|
59
|
+
FLAGS_transposed ? "enabled" : "disabled");
|
|
60
|
+
printf("verbose %s\n", FLAGS_verbose ? "enabled" : "disabled");
|
|
61
|
+
|
|
62
|
+
auto initFn = [](faiss::gpu::GpuResourcesProvider* res,
|
|
63
|
+
int dev) -> std::unique_ptr<faiss::gpu::GpuIndexFlat> {
|
|
64
|
+
if (FLAGS_pinned_mem >= 0) {
|
|
65
|
+
((faiss::gpu::StandardGpuResources*)res)
|
|
66
|
+
->setPinnedMemory(FLAGS_pinned_mem);
|
|
67
|
+
}
|
|
68
|
+
|
|
69
|
+
GpuIndexFlatConfig config;
|
|
70
|
+
config.device = dev;
|
|
71
|
+
config.useFloat16 = FLAGS_use_float16;
|
|
72
|
+
config.storeTransposed = FLAGS_transposed;
|
|
73
|
+
|
|
74
|
+
auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
|
|
75
|
+
FLAGS_L2_metric
|
|
76
|
+
? (faiss::gpu::GpuIndexFlat*)new faiss::gpu::
|
|
77
|
+
GpuIndexFlatL2(res, FLAGS_dim, config)
|
|
78
|
+
: (faiss::gpu::GpuIndexFlat*)new faiss::gpu::
|
|
79
|
+
GpuIndexFlatIP(res, FLAGS_dim, config));
|
|
80
|
+
|
|
81
|
+
if (FLAGS_min_paging_size >= 0) {
|
|
82
|
+
p->setMinPagingSize(FLAGS_min_paging_size);
|
|
83
|
+
}
|
|
84
|
+
return p;
|
|
85
|
+
};
|
|
86
|
+
|
|
87
|
+
IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
|
|
88
|
+
|
|
89
|
+
CUDA_VERIFY(cudaProfilerStart());
|
|
90
|
+
faiss::gpu::synchronizeAllDevices();
|
|
91
|
+
|
|
92
|
+
float gpuTime = 0.0f;
|
|
93
|
+
|
|
94
|
+
faiss::ClusteringParameters cp;
|
|
95
|
+
cp.niter = FLAGS_niter;
|
|
96
|
+
cp.verbose = FLAGS_verbose;
|
|
97
|
+
|
|
98
|
+
if (FLAGS_max_points > 0) {
|
|
99
|
+
cp.max_points_per_centroid = FLAGS_max_points;
|
|
77
100
|
}
|
|
78
|
-
return p;
|
|
79
|
-
};
|
|
80
|
-
|
|
81
|
-
IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
|
|
82
|
-
|
|
83
|
-
CUDA_VERIFY(cudaProfilerStart());
|
|
84
|
-
faiss::gpu::synchronizeAllDevices();
|
|
85
101
|
|
|
86
|
-
|
|
102
|
+
faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
|
|
87
103
|
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
104
|
+
// Time k-means
|
|
105
|
+
{
|
|
106
|
+
CpuTimer timer;
|
|
91
107
|
|
|
92
|
-
|
|
93
|
-
cp.max_points_per_centroid = FLAGS_max_points;
|
|
94
|
-
}
|
|
108
|
+
kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
|
|
95
109
|
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
CpuTimer timer;
|
|
101
|
-
|
|
102
|
-
kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
|
|
103
|
-
|
|
104
|
-
// There is a device -> host copy above, so no need to time
|
|
105
|
-
// additional synchronization with the GPU
|
|
106
|
-
gpuTime = timer.elapsedMilliseconds();
|
|
107
|
-
}
|
|
110
|
+
// There is a device -> host copy above, so no need to time
|
|
111
|
+
// additional synchronization with the GPU
|
|
112
|
+
gpuTime = timer.elapsedMilliseconds();
|
|
113
|
+
}
|
|
108
114
|
|
|
109
|
-
|
|
110
|
-
|
|
115
|
+
CUDA_VERIFY(cudaProfilerStop());
|
|
116
|
+
printf("k-means time %.3f ms\n", gpuTime);
|
|
111
117
|
|
|
112
|
-
|
|
118
|
+
CUDA_VERIFY(cudaDeviceSynchronize());
|
|
113
119
|
|
|
114
|
-
|
|
120
|
+
return 0;
|
|
115
121
|
}
|