faiss 0.5.3 → 0.6.1
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 +12 -0
- data/ext/faiss/ext.cpp +1 -1
- data/ext/faiss/extconf.rb +4 -4
- data/ext/faiss/index.cpp +63 -45
- data/ext/faiss/index_binary.cpp +37 -27
- data/ext/faiss/kmeans.cpp +9 -8
- data/ext/faiss/pca_matrix.cpp +9 -7
- data/ext/faiss/product_quantizer.cpp +13 -11
- data/ext/faiss/utils.cpp +4 -2
- data/ext/faiss/utils.h +4 -0
- data/lib/faiss/version.rb +1 -1
- data/lib/faiss.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +214 -82
- data/vendor/faiss/faiss/AutoTune.h +14 -1
- data/vendor/faiss/faiss/Clustering.cpp +97 -249
- data/vendor/faiss/faiss/Clustering.h +18 -0
- data/vendor/faiss/faiss/IVFlib.cpp +67 -44
- data/vendor/faiss/faiss/Index.cpp +25 -12
- data/vendor/faiss/faiss/Index.h +26 -4
- data/vendor/faiss/faiss/Index2Layer.cpp +37 -53
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +68 -61
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +36 -34
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +4 -1
- data/vendor/faiss/faiss/IndexBinary.cpp +6 -3
- data/vendor/faiss/faiss/IndexBinary.h +4 -4
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +1 -1
- data/vendor/faiss/faiss/IndexBinaryFlat.h +1 -1
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -4
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +92 -95
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +9 -3
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +45 -236
- data/vendor/faiss/faiss/IndexBinaryHash.h +6 -6
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +120 -414
- data/vendor/faiss/faiss/IndexFastScan.cpp +105 -129
- data/vendor/faiss/faiss/IndexFastScan.h +35 -24
- data/vendor/faiss/faiss/IndexFlat.cpp +216 -152
- data/vendor/faiss/faiss/IndexFlat.h +32 -14
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +88 -41
- data/vendor/faiss/faiss/IndexFlatCodes.h +7 -1
- data/vendor/faiss/faiss/IndexHNSW.cpp +299 -187
- data/vendor/faiss/faiss/IndexHNSW.h +30 -14
- data/vendor/faiss/faiss/IndexIDMap.cpp +26 -22
- data/vendor/faiss/faiss/IndexIDMap.h +9 -7
- data/vendor/faiss/faiss/IndexIVF.cpp +535 -405
- data/vendor/faiss/faiss/IndexIVF.h +47 -16
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +77 -74
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +105 -99
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +6 -3
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +379 -249
- data/vendor/faiss/faiss/IndexIVFFastScan.h +65 -60
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +41 -124
- data/vendor/faiss/faiss/IndexIVFFlat.h +32 -0
- data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +89 -138
- data/vendor/faiss/faiss/IndexIVFFlatPanorama.h +3 -1
- data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +18 -15
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +77 -907
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +184 -122
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +3 -0
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +23 -18
- data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +59 -60
- data/vendor/faiss/faiss/IndexIVFRaBitQ.h +4 -3
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +564 -416
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +269 -111
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +41 -127
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +1 -1
- data/vendor/faiss/faiss/IndexLSH.cpp +44 -25
- data/vendor/faiss/faiss/IndexLattice.cpp +41 -36
- data/vendor/faiss/faiss/IndexNNDescent.cpp +37 -21
- data/vendor/faiss/faiss/IndexNNDescent.h +2 -2
- data/vendor/faiss/faiss/IndexNSG.cpp +40 -23
- data/vendor/faiss/faiss/IndexNSG.h +0 -2
- data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +32 -12
- data/vendor/faiss/faiss/IndexPQ.cpp +129 -213
- data/vendor/faiss/faiss/IndexPQ.h +3 -2
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +20 -14
- data/vendor/faiss/faiss/IndexPQFastScan.h +3 -0
- data/vendor/faiss/faiss/IndexPreTransform.cpp +25 -18
- data/vendor/faiss/faiss/IndexPreTransform.h +1 -1
- data/vendor/faiss/faiss/IndexRaBitQ.cpp +31 -43
- data/vendor/faiss/faiss/IndexRaBitQ.h +4 -3
- data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +135 -317
- data/vendor/faiss/faiss/IndexRaBitQFastScan.h +192 -34
- data/vendor/faiss/faiss/IndexRefine.cpp +30 -55
- data/vendor/faiss/faiss/IndexRefine.h +4 -4
- data/vendor/faiss/faiss/IndexReplicas.cpp +6 -6
- data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +15 -14
- data/vendor/faiss/faiss/IndexRowwiseMinMax.h +1 -1
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +82 -14
- data/vendor/faiss/faiss/IndexShards.cpp +13 -13
- data/vendor/faiss/faiss/IndexShardsIVF.cpp +21 -15
- data/vendor/faiss/faiss/MatrixStats.cpp +5 -4
- data/vendor/faiss/faiss/MetaIndexes.cpp +19 -17
- data/vendor/faiss/faiss/MetaIndexes.h +1 -1
- data/vendor/faiss/faiss/MetricType.h +29 -6
- data/vendor/faiss/faiss/SuperKMeans.cpp +656 -0
- data/vendor/faiss/faiss/SuperKMeans.h +97 -0
- data/vendor/faiss/faiss/VectorTransform.cpp +349 -141
- data/vendor/faiss/faiss/VectorTransform.h +39 -16
- data/vendor/faiss/faiss/build.cpp +23 -0
- data/vendor/faiss/faiss/build.h +15 -0
- data/vendor/faiss/faiss/clone_index.cpp +55 -51
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +47 -47
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +11 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +38 -38
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +11 -0
- data/vendor/faiss/faiss/{cppcontrib/factory_tools.cpp → factory_tools.cpp} +6 -1
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +1 -1
- data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +6 -5
- data/vendor/faiss/faiss/gpu/GpuResources.h +1 -1
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +9 -9
- data/vendor/faiss/faiss/gpu/StandardGpuResources.h +4 -3
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +46 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +56 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +78 -1
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +72 -0
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +23 -0
- data/vendor/faiss/faiss/gpu/utils/CuvsFilterConvert.h +1 -1
- data/vendor/faiss/faiss/gpu/utils/CuvsUtils.h +21 -10
- data/vendor/faiss/faiss/gpu_metal/GpuIndexFlat.h +22 -0
- data/vendor/faiss/faiss/gpu_metal/MetalCloner.h +35 -0
- data/vendor/faiss/faiss/gpu_metal/MetalFlatKernels.h +40 -0
- data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +51 -0
- data/vendor/faiss/faiss/gpu_metal/MetalIndexFlat.h +65 -0
- data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +66 -0
- data/vendor/faiss/faiss/gpu_metal/MetalResources.h +79 -0
- data/vendor/faiss/faiss/gpu_metal/StandardMetalResources.h +35 -0
- data/vendor/faiss/faiss/impl/AdSampling.cpp +103 -0
- data/vendor/faiss/faiss/impl/AdSampling.h +35 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +64 -34
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +1 -0
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +10 -9
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +3 -28
- data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp +244 -0
- data/vendor/faiss/faiss/impl/ClusteringHelpers.h +94 -0
- data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +367 -0
- data/vendor/faiss/faiss/impl/ClusteringInitialization.h +107 -0
- data/vendor/faiss/faiss/impl/CodePacker.cpp +7 -3
- data/vendor/faiss/faiss/impl/CodePacker.h +11 -3
- data/vendor/faiss/faiss/impl/CodePackerRaBitQ.cpp +83 -0
- data/vendor/faiss/faiss/impl/CodePackerRaBitQ.h +47 -0
- data/vendor/faiss/faiss/impl/DistanceComputer.h +8 -8
- data/vendor/faiss/faiss/impl/FaissAssert.h +64 -3
- data/vendor/faiss/faiss/impl/FaissException.h +50 -3
- data/vendor/faiss/faiss/impl/HNSW.cpp +117 -351
- data/vendor/faiss/faiss/impl/HNSW.h +21 -40
- data/vendor/faiss/faiss/impl/IDSelector.cpp +15 -11
- data/vendor/faiss/faiss/impl/IDSelector.h +8 -8
- data/vendor/faiss/faiss/impl/InvertedListScannerStats.h +26 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +114 -102
- data/vendor/faiss/faiss/impl/NNDescent.cpp +63 -26
- data/vendor/faiss/faiss/impl/NNDescent.h +6 -2
- data/vendor/faiss/faiss/impl/NSG.cpp +44 -26
- data/vendor/faiss/faiss/impl/NSG.h +20 -10
- data/vendor/faiss/faiss/impl/Panorama.cpp +76 -52
- data/vendor/faiss/faiss/impl/Panorama.h +265 -78
- data/vendor/faiss/faiss/impl/PdxLayout.cpp +93 -0
- data/vendor/faiss/faiss/impl/PdxLayout.h +41 -0
- data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +62 -37
- data/vendor/faiss/faiss/impl/PolysemousTraining.h +3 -3
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +35 -35
- data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +21 -16
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +99 -80
- data/vendor/faiss/faiss/impl/Quantizer.h +2 -2
- data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +135 -37
- data/vendor/faiss/faiss/impl/RaBitQUtils.h +148 -21
- data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +298 -301
- data/vendor/faiss/faiss/impl/RaBitQuantizer.h +3 -10
- data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp +15 -41
- data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.h +0 -4
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +40 -32
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +1 -1
- data/vendor/faiss/faiss/impl/ResultHandler.h +218 -113
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +119 -2362
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +27 -3
- data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +14 -11
- data/vendor/faiss/faiss/impl/VisitedTable.cpp +42 -0
- data/vendor/faiss/faiss/impl/VisitedTable.h +76 -0
- data/vendor/faiss/faiss/impl/approx_topk/approx_topk.h +276 -0
- data/vendor/faiss/faiss/impl/approx_topk/avx2.cpp +68 -0
- data/vendor/faiss/faiss/{utils → impl}/approx_topk/generic.h +15 -8
- data/vendor/faiss/faiss/impl/approx_topk/neon.cpp +68 -0
- data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab-inl.h +169 -0
- data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab.h +117 -0
- data/vendor/faiss/faiss/impl/approx_topk/simdlib256-inl.h +146 -0
- data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h +73 -0
- data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHash_impl.h +270 -0
- data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +460 -0
- data/vendor/faiss/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h +159 -0
- data/vendor/faiss/faiss/impl/binary_hamming/IndexPQ_impl.h +92 -0
- data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +26 -0
- data/vendor/faiss/faiss/impl/binary_hamming/avx512.cpp +26 -0
- data/vendor/faiss/faiss/impl/binary_hamming/dispatch.h +143 -0
- data/vendor/faiss/faiss/impl/binary_hamming/neon.cpp +26 -0
- data/vendor/faiss/faiss/impl/binary_hamming/rvv.cpp +26 -0
- data/vendor/faiss/faiss/impl/expanded_scanners.h +163 -0
- data/vendor/faiss/faiss/impl/{FastScanDistancePostProcessing.h → fast_scan/FastScanDistancePostProcessing.h} +13 -6
- data/vendor/faiss/faiss/impl/{LookupTableScaler.h → fast_scan/LookupTableScaler.h} +16 -5
- data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops.h +237 -0
- data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops_512.h +185 -0
- data/vendor/faiss/faiss/impl/fast_scan/decompose_qbs.h +229 -0
- data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +268 -0
- data/vendor/faiss/faiss/impl/{pq4_fast_scan.cpp → fast_scan/fast_scan.cpp} +176 -4
- data/vendor/faiss/faiss/impl/fast_scan/fast_scan.h +341 -0
- data/vendor/faiss/faiss/impl/fast_scan/impl-avx2.cpp +36 -0
- data/vendor/faiss/faiss/impl/fast_scan/impl-avx512.cpp +40 -0
- data/vendor/faiss/faiss/impl/fast_scan/impl-neon.cpp +120 -0
- data/vendor/faiss/faiss/impl/fast_scan/impl-riscv.cpp +104 -0
- data/vendor/faiss/faiss/impl/fast_scan/kernels_simd256.h +213 -0
- data/vendor/faiss/faiss/impl/{pq4_fast_scan_search_qbs.cpp → fast_scan/kernels_simd512.h} +26 -348
- data/vendor/faiss/faiss/impl/fast_scan/rabitq_dispatching.h +90 -0
- data/vendor/faiss/faiss/impl/fast_scan/rabitq_result_handler.h +108 -0
- data/vendor/faiss/faiss/impl/{simd_result_handlers.h → fast_scan/simd_result_handlers.h} +290 -142
- data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +54 -0
- data/vendor/faiss/faiss/impl/hnsw/LockVector.h +64 -0
- data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp +91 -0
- data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h +64 -0
- data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +104 -0
- data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +111 -0
- data/vendor/faiss/faiss/impl/index_read.cpp +1950 -505
- data/vendor/faiss/faiss/impl/index_read_utils.h +1 -2
- data/vendor/faiss/faiss/impl/index_write.cpp +112 -21
- data/vendor/faiss/faiss/impl/io.cpp +6 -6
- data/vendor/faiss/faiss/impl/io_macros.h +33 -16
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +10 -10
- data/vendor/faiss/faiss/impl/lattice_Zn.cpp +81 -40
- data/vendor/faiss/faiss/impl/lattice_Zn.h +6 -6
- data/vendor/faiss/faiss/impl/mapped_io.cpp +15 -8
- data/vendor/faiss/faiss/impl/platform_macros.h +11 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQScanner_impl.h +549 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +245 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h +105 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +106 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +21 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +21 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +21 -0
- data/vendor/faiss/faiss/impl/{code_distance/code_distance-avx2.h → pq_code_distance/pq_code_distance-avx2.h} +43 -220
- data/vendor/faiss/faiss/impl/{code_distance/code_distance-avx512.h → pq_code_distance/pq_code_distance-avx512.h} +25 -112
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +59 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.h +96 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +256 -0
- data/vendor/faiss/faiss/impl/{code_distance/code_distance-sve.h → pq_code_distance/pq_code_distance-sve.cpp} +57 -146
- data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +68 -0
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +320 -483
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +1 -1
- data/vendor/faiss/faiss/impl/scalar_quantizer/codecs.h +121 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +137 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +371 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h +190 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/similarities.h +94 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +603 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +597 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +388 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +630 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +311 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/training.cpp +387 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/training.h +54 -0
- data/vendor/faiss/faiss/impl/simd_dispatch.h +173 -0
- data/vendor/faiss/faiss/impl/simdlib/simdlib.h +57 -0
- data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_avx2.h +274 -171
- data/vendor/faiss/faiss/impl/simdlib/simdlib_avx512.h +414 -0
- data/vendor/faiss/faiss/impl/simdlib/simdlib_dispatch.h +44 -0
- data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_emulated.h +231 -166
- data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_neon.h +275 -217
- data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_ppc64.h +201 -160
- data/vendor/faiss/faiss/impl/svs_io.cpp +12 -3
- data/vendor/faiss/faiss/impl/svs_io.h +8 -2
- data/vendor/faiss/faiss/index_factory.cpp +115 -28
- data/vendor/faiss/faiss/index_io.h +53 -3
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +73 -20
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +24 -14
- data/vendor/faiss/faiss/invlists/DirectMap.h +4 -3
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +157 -73
- data/vendor/faiss/faiss/invlists/InvertedLists.h +86 -23
- data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +4 -4
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +14 -14
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +1 -1
- data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +9 -19
- data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +2 -2
- data/vendor/faiss/faiss/svs/IndexSVSFlat.h +2 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +350 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVF.h +128 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.cpp +40 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.h +43 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +225 -0
- data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.h +71 -0
- data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +25 -1
- data/vendor/faiss/faiss/svs/IndexSVSVamana.h +19 -2
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +1 -1
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +19 -2
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +14 -0
- data/vendor/faiss/faiss/utils/Heap.cpp +56 -10
- data/vendor/faiss/faiss/utils/Heap.h +21 -0
- data/vendor/faiss/faiss/utils/NeuralNet.cpp +54 -40
- data/vendor/faiss/faiss/utils/NeuralNet.h +1 -1
- data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +10 -4
- data/vendor/faiss/faiss/utils/distances.cpp +507 -559
- data/vendor/faiss/faiss/utils/distances.h +118 -1
- data/vendor/faiss/faiss/utils/distances_dispatch.h +250 -0
- data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +8 -7
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +33 -14
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +12 -1
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +16 -293
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based_neon.cpp +57 -0
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_kernel-inl.h +290 -0
- data/vendor/faiss/faiss/utils/distances_simd.cpp +72 -3681
- data/vendor/faiss/faiss/utils/extra_distances.cpp +60 -102
- data/vendor/faiss/faiss/utils/extra_distances.h +79 -7
- data/vendor/faiss/faiss/utils/hamming-inl.h +13 -11
- data/vendor/faiss/faiss/utils/hamming.cpp +66 -517
- data/vendor/faiss/faiss/utils/hamming.h +92 -2
- data/vendor/faiss/faiss/utils/hamming_distance/common.h +287 -10
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +15 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512.cpp +15 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx2.h +142 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +234 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-generic.h +368 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-neon.h +322 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-rvv.h +39 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer.h +146 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h +481 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_neon.cpp +15 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_rvv.cpp +15 -0
- data/vendor/faiss/faiss/utils/partitioning.cpp +66 -987
- data/vendor/faiss/faiss/utils/partitioning.h +31 -0
- data/vendor/faiss/faiss/utils/popcount.h +29 -0
- data/vendor/faiss/faiss/utils/pq_code_distance.h +251 -0
- data/vendor/faiss/faiss/utils/prefetch.h +2 -2
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +30 -30
- data/vendor/faiss/faiss/utils/quantize_lut.h +1 -1
- data/vendor/faiss/faiss/utils/rabitq_simd.h +124 -343
- data/vendor/faiss/faiss/utils/random.cpp +6 -6
- data/vendor/faiss/faiss/utils/simd_impl/IVFFlatScanner-inl.h +51 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_aarch64.cpp +154 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +777 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_autovec-inl.h +306 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +1431 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +1095 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +189 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_simdlib256.h +195 -0
- data/vendor/faiss/faiss/utils/simd_impl/distances_sse-inl.h +392 -0
- data/vendor/faiss/faiss/utils/{distances_fused/simdlib_based.h → simd_impl/exhaustive_L2sqr_blas_cmax.h} +5 -10
- data/vendor/faiss/faiss/utils/simd_impl/hamming_impl.h +481 -0
- data/vendor/faiss/faiss/utils/simd_impl/partitioning_avx2.cpp +14 -0
- data/vendor/faiss/faiss/utils/simd_impl/partitioning_neon.cpp +14 -0
- data/vendor/faiss/faiss/utils/simd_impl/partitioning_simdlib256.h +1085 -0
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +355 -0
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +477 -0
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +55 -0
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +55 -0
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h +32 -0
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h +43 -0
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +57 -0
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +45 -0
- data/vendor/faiss/faiss/utils/simd_levels.cpp +334 -0
- data/vendor/faiss/faiss/utils/simd_levels.h +183 -0
- data/vendor/faiss/faiss/utils/sorting.cpp +48 -36
- data/vendor/faiss/faiss/utils/utils.cpp +21 -14
- data/vendor/faiss/faiss/utils/utils.h +3 -3
- metadata +156 -42
- data/vendor/faiss/faiss/impl/RaBitQStats.cpp +0 -29
- data/vendor/faiss/faiss/impl/RaBitQStats.h +0 -56
- data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h +0 -81
- data/vendor/faiss/faiss/impl/code_distance/code_distance.h +0 -186
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +0 -216
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +0 -224
- data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h +0 -84
- data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h +0 -196
- data/vendor/faiss/faiss/utils/approx_topk/mode.h +0 -34
- data/vendor/faiss/faiss/utils/distances_fused/avx512.h +0 -36
- data/vendor/faiss/faiss/utils/extra_distances-inl.h +0 -228
- data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +0 -462
- data/vendor/faiss/faiss/utils/hamming_distance/avx512-inl.h +0 -490
- data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +0 -450
- data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +0 -87
- data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +0 -524
- data/vendor/faiss/faiss/utils/simdlib.h +0 -42
- data/vendor/faiss/faiss/utils/simdlib_avx512.h +0 -296
- /data/vendor/faiss/faiss/{cppcontrib/factory_tools.h → factory_tools.h} +0 -0
|
@@ -28,6 +28,7 @@
|
|
|
28
28
|
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
|
29
29
|
|
|
30
30
|
#include <variant>
|
|
31
|
+
#include <vector>
|
|
31
32
|
#include "faiss/Index.h"
|
|
32
33
|
|
|
33
34
|
namespace faiss {
|
|
@@ -193,14 +194,14 @@ struct GpuIndexCagraConfig : public GpuIndexConfig {
|
|
|
193
194
|
|
|
194
195
|
enum class search_algo {
|
|
195
196
|
/// For large batch sizes.
|
|
196
|
-
SINGLE_CTA,
|
|
197
|
+
SINGLE_CTA = 0,
|
|
197
198
|
/// For small batch sizes.
|
|
198
|
-
MULTI_CTA,
|
|
199
|
-
MULTI_KERNEL,
|
|
200
|
-
AUTO
|
|
199
|
+
MULTI_CTA = 1,
|
|
200
|
+
MULTI_KERNEL = 2,
|
|
201
|
+
AUTO = 100
|
|
201
202
|
};
|
|
202
203
|
|
|
203
|
-
enum class hash_mode { HASH, SMALL, AUTO };
|
|
204
|
+
enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 };
|
|
204
205
|
|
|
205
206
|
struct SearchParametersCagra : SearchParameters {
|
|
206
207
|
/// Maximum number of queries to search at the same time (batch size). Auto
|
|
@@ -23,9 +23,9 @@
|
|
|
23
23
|
|
|
24
24
|
#if defined USE_NVIDIA_CUVS
|
|
25
25
|
#include <raft/core/device_resources.hpp>
|
|
26
|
-
#include <rmm/mr/
|
|
27
|
-
#include <rmm/mr/
|
|
28
|
-
#include <rmm/mr/
|
|
26
|
+
#include <rmm/mr/managed_memory_resource.hpp>
|
|
27
|
+
#include <rmm/mr/per_device_resource.hpp>
|
|
28
|
+
#include <rmm/mr/pinned_host_memory_resource.hpp>
|
|
29
29
|
#include <memory>
|
|
30
30
|
#endif
|
|
31
31
|
|
|
@@ -93,7 +93,7 @@ StandardGpuResourcesImpl::StandardGpuResourcesImpl()
|
|
|
93
93
|
:
|
|
94
94
|
#if defined USE_NVIDIA_CUVS
|
|
95
95
|
mmr_(new rmm::mr::managed_memory_resource),
|
|
96
|
-
pmr_(new rmm::mr::
|
|
96
|
+
pmr_(new rmm::mr::pinned_host_memory_resource),
|
|
97
97
|
#endif
|
|
98
98
|
pinnedMemAlloc_(nullptr),
|
|
99
99
|
pinnedMemAllocSize_(0),
|
|
@@ -164,7 +164,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
|
|
|
164
164
|
|
|
165
165
|
if (pinnedMemAlloc_) {
|
|
166
166
|
#if defined USE_NVIDIA_CUVS
|
|
167
|
-
pmr_->
|
|
167
|
+
pmr_->deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
|
|
168
168
|
#else
|
|
169
169
|
auto err = cudaFreeHost(pinnedMemAlloc_);
|
|
170
170
|
FAISS_ASSERT_FMT(
|
|
@@ -350,7 +350,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
|
|
|
350
350
|
// pinned memory allocation
|
|
351
351
|
if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
|
|
352
352
|
try {
|
|
353
|
-
pinnedMemAlloc_ = pmr_->
|
|
353
|
+
pinnedMemAlloc_ = pmr_->allocate_sync(pinnedMemSize_);
|
|
354
354
|
} catch (const std::bad_alloc& rmm_ex) {
|
|
355
355
|
FAISS_THROW_MSG("CUDA memory allocation error");
|
|
356
356
|
}
|
|
@@ -549,7 +549,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
|
|
|
549
549
|
rmm::mr::device_memory_resource* current_mr =
|
|
550
550
|
rmm::mr::get_per_device_resource(
|
|
551
551
|
rmm::cuda_device_id{adjReq.device});
|
|
552
|
-
p = current_mr->
|
|
552
|
+
p = current_mr->allocate(adjReq.stream, adjReq.size);
|
|
553
553
|
adjReq.mr = current_mr;
|
|
554
554
|
} catch (const std::bad_alloc& rmm_ex) {
|
|
555
555
|
FAISS_THROW_MSG("CUDA memory allocation error");
|
|
@@ -584,7 +584,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
|
|
|
584
584
|
// TODO: change this to use the current device resource once RMM has
|
|
585
585
|
// a way to retrieve a "guaranteed" managed memory resource for a
|
|
586
586
|
// device.
|
|
587
|
-
p = mmr_->
|
|
587
|
+
p = mmr_->allocate(adjReq.stream, adjReq.size);
|
|
588
588
|
adjReq.mr = mmr_.get();
|
|
589
589
|
} catch (const std::bad_alloc& rmm_ex) {
|
|
590
590
|
FAISS_THROW_MSG("CUDA memory allocation error");
|
|
@@ -648,7 +648,7 @@ void StandardGpuResourcesImpl::deallocMemory(int device, void* p) {
|
|
|
648
648
|
req.space == MemorySpace::Device ||
|
|
649
649
|
req.space == MemorySpace::Unified) {
|
|
650
650
|
#if defined USE_NVIDIA_CUVS
|
|
651
|
-
req.mr->
|
|
651
|
+
req.mr->deallocate(req.stream, p, req.size);
|
|
652
652
|
#else
|
|
653
653
|
auto err = cudaFree(p);
|
|
654
654
|
FAISS_ASSERT_FMT(
|
|
@@ -25,7 +25,8 @@
|
|
|
25
25
|
|
|
26
26
|
#if defined USE_NVIDIA_CUVS
|
|
27
27
|
#include <raft/core/device_resources.hpp>
|
|
28
|
-
#include <rmm/mr/
|
|
28
|
+
#include <rmm/mr/device_memory_resource.hpp>
|
|
29
|
+
#include <rmm/mr/pinned_host_memory_resource.hpp>
|
|
29
30
|
#endif
|
|
30
31
|
|
|
31
32
|
#include <faiss/gpu/GpuResources.h>
|
|
@@ -172,8 +173,8 @@ class StandardGpuResourcesImpl : public GpuResources {
|
|
|
172
173
|
// managed_memory_resource
|
|
173
174
|
std::unique_ptr<rmm::mr::device_memory_resource> mmr_;
|
|
174
175
|
|
|
175
|
-
//
|
|
176
|
-
std::unique_ptr<rmm::mr::
|
|
176
|
+
// pinned_host_memory_resource
|
|
177
|
+
std::unique_ptr<rmm::mr::pinned_host_memory_resource> pmr_;
|
|
177
178
|
#endif
|
|
178
179
|
|
|
179
180
|
/// Pinned memory allocation for use with this GPU
|
|
@@ -11,6 +11,7 @@
|
|
|
11
11
|
#include <faiss/gpu/impl/IndexUtils.h>
|
|
12
12
|
#include <faiss/gpu/test/TestUtils.h>
|
|
13
13
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
|
14
|
+
#include <faiss/impl/IDSelector.h>
|
|
14
15
|
#include <gtest/gtest.h>
|
|
15
16
|
#include <sstream>
|
|
16
17
|
#include <unordered_map>
|
|
@@ -761,6 +762,51 @@ TEST(TestCuvsGpuIndexFlat, SearchAndReconstruct) {
|
|
|
761
762
|
}
|
|
762
763
|
#endif
|
|
763
764
|
|
|
765
|
+
void testIDSelectorFlat(faiss::MetricType metricType) {
|
|
766
|
+
int numAdd = faiss::gpu::randVal(2000, 5000);
|
|
767
|
+
int dim = faiss::gpu::randVal(64, 200);
|
|
768
|
+
int numQuery = faiss::gpu::randVal(32, 100);
|
|
769
|
+
int k = faiss::gpu::randVal(10, 30);
|
|
770
|
+
int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
|
|
771
|
+
|
|
772
|
+
faiss::gpu::StandardGpuResources res;
|
|
773
|
+
res.noTempMemory();
|
|
774
|
+
|
|
775
|
+
auto queryVecs = faiss::gpu::randVecs(numQuery, dim);
|
|
776
|
+
faiss::gpu::GpuIndexFlatConfig config;
|
|
777
|
+
config.device = device;
|
|
778
|
+
config.use_cuvs = true;
|
|
779
|
+
for (bool useFloat16 : {false, true}) {
|
|
780
|
+
config.useFloat16 = useFloat16;
|
|
781
|
+
faiss::gpu::GpuIndexFlat gpuIndex(&res, dim, metricType, config);
|
|
782
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
|
|
783
|
+
gpuIndex.add(numAdd, addVecs.data());
|
|
784
|
+
|
|
785
|
+
faiss::gpu::TestIDSelectorStruct selector_struct(numAdd);
|
|
786
|
+
faiss::SearchParameters search_params;
|
|
787
|
+
for (auto& [selectorName, selector] : selector_struct.selector_map) {
|
|
788
|
+
search_params.sel = selector.get();
|
|
789
|
+
faiss::gpu::testIDSelectorSearch(
|
|
790
|
+
&gpuIndex,
|
|
791
|
+
&search_params,
|
|
792
|
+
queryVecs,
|
|
793
|
+
numQuery,
|
|
794
|
+
k,
|
|
795
|
+
selectorName);
|
|
796
|
+
}
|
|
797
|
+
}
|
|
798
|
+
}
|
|
799
|
+
|
|
800
|
+
#if defined USE_NVIDIA_CUVS
|
|
801
|
+
TEST(TestCuvsGpuIndexFlat, IDSelector_L2) {
|
|
802
|
+
testIDSelectorFlat(faiss::METRIC_L2);
|
|
803
|
+
}
|
|
804
|
+
|
|
805
|
+
TEST(TestCuvsGpuIndexFlat, IDSelector_IP) {
|
|
806
|
+
testIDSelectorFlat(faiss::METRIC_INNER_PRODUCT);
|
|
807
|
+
}
|
|
808
|
+
#endif
|
|
809
|
+
|
|
764
810
|
int main(int argc, char** argv) {
|
|
765
811
|
testing::InitGoogleTest(&argc, argv);
|
|
766
812
|
|
|
@@ -28,6 +28,8 @@
|
|
|
28
28
|
#include <faiss/gpu/StandardGpuResources.h>
|
|
29
29
|
#include <faiss/gpu/test/TestUtils.h>
|
|
30
30
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
|
31
|
+
#include <faiss/impl/IDSelector.h>
|
|
32
|
+
#include <faiss/utils/distances.h>
|
|
31
33
|
#include <gtest/gtest.h>
|
|
32
34
|
#include <cmath>
|
|
33
35
|
#include <sstream>
|
|
@@ -354,6 +356,11 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
|
|
|
354
356
|
}
|
|
355
357
|
|
|
356
358
|
TEST(TestGpuIndexIVFFlat, LargeBatch) {
|
|
359
|
+
// With low-dim vectors, CPU will use non-BLAS. Force the CPU to use
|
|
360
|
+
// the BLAS for consistent comparison.
|
|
361
|
+
int saved_threshold = faiss::distance_compute_blas_threshold;
|
|
362
|
+
faiss::distance_compute_blas_threshold = 1;
|
|
363
|
+
|
|
357
364
|
Options opt;
|
|
358
365
|
opt.dim = 3;
|
|
359
366
|
opt.numQuery = 100000;
|
|
@@ -364,6 +371,8 @@ TEST(TestGpuIndexIVFFlat, LargeBatch) {
|
|
|
364
371
|
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
365
372
|
queryTest(opt, faiss::METRIC_L2, false);
|
|
366
373
|
#endif
|
|
374
|
+
|
|
375
|
+
faiss::distance_compute_blas_threshold = saved_threshold;
|
|
367
376
|
}
|
|
368
377
|
|
|
369
378
|
// float16 coarse quantizer
|
|
@@ -909,6 +918,53 @@ TEST(TestGpuIndexIVFFlat, Reconstruct_n) {
|
|
|
909
918
|
EXPECT_EQ(gpuVals, cpuVals);
|
|
910
919
|
}
|
|
911
920
|
|
|
921
|
+
void testIDSelectorIVFFlat(faiss::MetricType metricType) {
|
|
922
|
+
Options opt;
|
|
923
|
+
|
|
924
|
+
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
925
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
|
926
|
+
|
|
927
|
+
faiss::gpu::StandardGpuResources res;
|
|
928
|
+
res.noTempMemory();
|
|
929
|
+
|
|
930
|
+
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
931
|
+
config.device = opt.device;
|
|
932
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
933
|
+
config.use_cuvs = true;
|
|
934
|
+
|
|
935
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
936
|
+
&res, opt.dim, opt.numCentroids, metricType, config);
|
|
937
|
+
gpuIndex.nprobe = opt.nprobe;
|
|
938
|
+
|
|
939
|
+
gpuIndex.train(opt.numTrain, trainVecs.data());
|
|
940
|
+
gpuIndex.add(opt.numAdd, addVecs.data());
|
|
941
|
+
|
|
942
|
+
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
|
943
|
+
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
|
|
944
|
+
faiss::SearchParametersIVF search_params;
|
|
945
|
+
search_params.nprobe = opt.nprobe;
|
|
946
|
+
for (auto& [selectorName, selector] : selector_struct.selector_map) {
|
|
947
|
+
search_params.sel = selector.get();
|
|
948
|
+
faiss::gpu::testIDSelectorSearch(
|
|
949
|
+
&gpuIndex,
|
|
950
|
+
&search_params,
|
|
951
|
+
queryVecs,
|
|
952
|
+
opt.numQuery,
|
|
953
|
+
opt.k,
|
|
954
|
+
selectorName);
|
|
955
|
+
}
|
|
956
|
+
}
|
|
957
|
+
|
|
958
|
+
#if defined USE_NVIDIA_CUVS
|
|
959
|
+
TEST(TestCuvsGpuIndexIVFFlat, IDSelector_L2) {
|
|
960
|
+
testIDSelectorIVFFlat(faiss::METRIC_L2);
|
|
961
|
+
}
|
|
962
|
+
|
|
963
|
+
TEST(TestCuvsGpuIndexIVFFlat, IDSelector_IP) {
|
|
964
|
+
testIDSelectorIVFFlat(faiss::METRIC_INNER_PRODUCT);
|
|
965
|
+
}
|
|
966
|
+
#endif
|
|
967
|
+
|
|
912
968
|
int main(int argc, char** argv) {
|
|
913
969
|
testing::InitGoogleTest(&argc, argv);
|
|
914
970
|
|
|
@@ -6,11 +6,14 @@
|
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
8
|
#include <faiss/IndexFlat.h>
|
|
9
|
+
#include <faiss/IndexIVF.h>
|
|
9
10
|
#include <faiss/IndexIVFPQ.h>
|
|
10
11
|
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
|
11
12
|
#include <faiss/gpu/StandardGpuResources.h>
|
|
12
13
|
#include <faiss/gpu/test/TestUtils.h>
|
|
13
14
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
|
15
|
+
#include <faiss/impl/IDSelector.h>
|
|
16
|
+
#include <faiss/utils/distances.h>
|
|
14
17
|
#include <gtest/gtest.h>
|
|
15
18
|
#include <cmath>
|
|
16
19
|
#include <sstream>
|
|
@@ -101,7 +104,11 @@ struct Options {
|
|
|
101
104
|
}
|
|
102
105
|
|
|
103
106
|
float getCompareEpsilon() const {
|
|
104
|
-
|
|
107
|
+
// With very low dimensionality (e.g., dim=4, codes=2 giving
|
|
108
|
+
// dimPerSubQuantizer=2), L2 distances can be very small
|
|
109
|
+
// (near-zero), causing relative error comparisons to be
|
|
110
|
+
// unstable despite tiny absolute differences.
|
|
111
|
+
return (dim <= 8) ? 0.15f : 0.035f;
|
|
105
112
|
}
|
|
106
113
|
|
|
107
114
|
float getPctMaxDiff1() const {
|
|
@@ -190,6 +197,11 @@ TEST(TestGpuIndexIVFPQ, Query_IP) {
|
|
|
190
197
|
|
|
191
198
|
// Large batch sizes (>= 65536) should also work
|
|
192
199
|
TEST(TestGpuIndexIVFPQ, LargeBatch) {
|
|
200
|
+
// With low-dim vectors, CPU will use non-BLAS. Force the CPU to use
|
|
201
|
+
// the BLAS for consistent comparison.
|
|
202
|
+
int saved_threshold = faiss::distance_compute_blas_threshold;
|
|
203
|
+
faiss::distance_compute_blas_threshold = 1;
|
|
204
|
+
|
|
193
205
|
for (bool usePrecomputed : {false, true}) {
|
|
194
206
|
Options opt;
|
|
195
207
|
|
|
@@ -202,6 +214,8 @@ TEST(TestGpuIndexIVFPQ, LargeBatch) {
|
|
|
202
214
|
|
|
203
215
|
queryTest(opt, faiss::MetricType::METRIC_L2);
|
|
204
216
|
}
|
|
217
|
+
|
|
218
|
+
faiss::distance_compute_blas_threshold = saved_threshold;
|
|
205
219
|
}
|
|
206
220
|
|
|
207
221
|
void testMMCodeDistance(faiss::MetricType mt) {
|
|
@@ -697,6 +711,11 @@ TEST(TestGpuIndexIVFPQ, Query_IP_Cuvs) {
|
|
|
697
711
|
|
|
698
712
|
// Large batch sizes (>= 65536) should also work
|
|
699
713
|
TEST(TestGpuIndexIVFPQ, LargeBatch_Cuvs) {
|
|
714
|
+
// See LargeBatch comment: force CPU BLAS path to match GPU GEMM
|
|
715
|
+
// decomposition for consistent L2 distance computation.
|
|
716
|
+
int saved_threshold = faiss::distance_compute_blas_threshold;
|
|
717
|
+
faiss::distance_compute_blas_threshold = 1;
|
|
718
|
+
|
|
700
719
|
Options opt;
|
|
701
720
|
|
|
702
721
|
// override for large sizes
|
|
@@ -711,6 +730,8 @@ TEST(TestGpuIndexIVFPQ, LargeBatch_Cuvs) {
|
|
|
711
730
|
opt.bitsPerCode = 8;
|
|
712
731
|
|
|
713
732
|
queryTest(opt, faiss::MetricType::METRIC_L2);
|
|
733
|
+
|
|
734
|
+
faiss::distance_compute_blas_threshold = saved_threshold;
|
|
714
735
|
}
|
|
715
736
|
|
|
716
737
|
TEST(TestGpuIndexIVFPQ, CopyFrom_Cuvs) {
|
|
@@ -877,6 +898,62 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
|
|
|
877
898
|
#endif
|
|
878
899
|
}
|
|
879
900
|
|
|
901
|
+
void testIDSelectorIVFPQ(faiss::MetricType metricType) {
|
|
902
|
+
Options opt;
|
|
903
|
+
|
|
904
|
+
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
905
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
|
906
|
+
|
|
907
|
+
faiss::IndexFlatL2 quantizerL2(opt.dim);
|
|
908
|
+
faiss::IndexFlatIP quantizerIP(opt.dim);
|
|
909
|
+
faiss::Index* quantizer = metricType == faiss::METRIC_L2
|
|
910
|
+
? (faiss::Index*)&quantizerL2
|
|
911
|
+
: (faiss::Index*)&quantizerIP;
|
|
912
|
+
|
|
913
|
+
faiss::IndexIVFPQ cpuIndex(
|
|
914
|
+
quantizer, opt.dim, opt.numCentroids, opt.codes, opt.bitsPerCode);
|
|
915
|
+
cpuIndex.metric_type = metricType;
|
|
916
|
+
cpuIndex.train(opt.numTrain, trainVecs.data());
|
|
917
|
+
cpuIndex.add(opt.numAdd, addVecs.data());
|
|
918
|
+
|
|
919
|
+
faiss::gpu::StandardGpuResources res;
|
|
920
|
+
res.noTempMemory();
|
|
921
|
+
|
|
922
|
+
faiss::gpu::GpuIndexIVFPQConfig config;
|
|
923
|
+
config.device = opt.device;
|
|
924
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
925
|
+
config.interleavedLayout = true;
|
|
926
|
+
config.use_cuvs = true;
|
|
927
|
+
|
|
928
|
+
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
|
929
|
+
gpuIndex.nprobe = opt.nprobe;
|
|
930
|
+
|
|
931
|
+
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
|
932
|
+
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
|
|
933
|
+
faiss::SearchParametersIVF search_params;
|
|
934
|
+
search_params.nprobe = opt.nprobe;
|
|
935
|
+
for (auto& [selectorName, selector] : selector_struct.selector_map) {
|
|
936
|
+
search_params.sel = selector.get();
|
|
937
|
+
faiss::gpu::testIDSelectorSearch(
|
|
938
|
+
&gpuIndex,
|
|
939
|
+
&search_params,
|
|
940
|
+
queryVecs,
|
|
941
|
+
opt.numQuery,
|
|
942
|
+
opt.k,
|
|
943
|
+
selectorName);
|
|
944
|
+
}
|
|
945
|
+
}
|
|
946
|
+
|
|
947
|
+
#if defined USE_NVIDIA_CUVS
|
|
948
|
+
TEST(TestCuvsGpuIndexIVFPQ, IDSelector_L2) {
|
|
949
|
+
testIDSelectorIVFPQ(faiss::METRIC_L2);
|
|
950
|
+
}
|
|
951
|
+
|
|
952
|
+
TEST(TestCuvsGpuIndexIVFPQ, IDSelector_IP) {
|
|
953
|
+
testIDSelectorIVFPQ(faiss::METRIC_INNER_PRODUCT);
|
|
954
|
+
}
|
|
955
|
+
#endif
|
|
956
|
+
|
|
880
957
|
int main(int argc, char** argv) {
|
|
881
958
|
testing::InitGoogleTest(&argc, argv);
|
|
882
959
|
|
|
@@ -6,6 +6,7 @@
|
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
8
|
#include <faiss/gpu/test/TestUtils.h>
|
|
9
|
+
#include <faiss/impl/IDSelector.h>
|
|
9
10
|
#include <faiss/utils/random.h>
|
|
10
11
|
#include <gtest/gtest.h>
|
|
11
12
|
#include <time.h>
|
|
@@ -429,5 +430,76 @@ void compareLists(
|
|
|
429
430
|
}
|
|
430
431
|
}
|
|
431
432
|
|
|
433
|
+
TestIDSelectorStruct::TestIDSelectorStruct(int numAdd) {
|
|
434
|
+
// Range selector [20%, 80%] of the database
|
|
435
|
+
size_t min_id = numAdd / 5;
|
|
436
|
+
size_t max_id = numAdd * 4 / 5;
|
|
437
|
+
selector_map["Range"] =
|
|
438
|
+
std::make_unique<faiss::IDSelectorRange>(min_id, max_id);
|
|
439
|
+
|
|
440
|
+
// Array selector (every 3rd element)
|
|
441
|
+
array_ids.clear();
|
|
442
|
+
for (int i = 0; i < numAdd; i += 3) {
|
|
443
|
+
array_ids.push_back(i);
|
|
444
|
+
}
|
|
445
|
+
selector_map["Array"] = std::make_unique<faiss::IDSelectorArray>(
|
|
446
|
+
array_ids.size(), array_ids.data());
|
|
447
|
+
|
|
448
|
+
// Batch selector (every 5th element)
|
|
449
|
+
batch_ids.clear();
|
|
450
|
+
for (int i = 1; i < numAdd; i += 5) {
|
|
451
|
+
batch_ids.push_back(i);
|
|
452
|
+
}
|
|
453
|
+
selector_map["Batch"] = std::make_unique<faiss::IDSelectorBatch>(
|
|
454
|
+
batch_ids.size(), batch_ids.data());
|
|
455
|
+
|
|
456
|
+
// Bitmap selector (every 4th element selected)
|
|
457
|
+
size_t bitmap_size = (numAdd + 7) / 8;
|
|
458
|
+
bitmap.resize(bitmap_size, 0);
|
|
459
|
+
for (int i = 0; i < numAdd; i += 4) {
|
|
460
|
+
int byte_idx = i / 8;
|
|
461
|
+
int bit_idx = i % 8;
|
|
462
|
+
bitmap[byte_idx] |= (1 << bit_idx);
|
|
463
|
+
}
|
|
464
|
+
selector_map["Bitmap"] =
|
|
465
|
+
std::make_unique<faiss::IDSelectorBitmap>(numAdd, bitmap.data());
|
|
466
|
+
|
|
467
|
+
selector_map["Not"] =
|
|
468
|
+
std::make_unique<faiss::IDSelectorNot>(selector_map["Range"].get());
|
|
469
|
+
selector_map["And"] = std::make_unique<faiss::IDSelectorAnd>(
|
|
470
|
+
selector_map["Range"].get(), selector_map["Array"].get());
|
|
471
|
+
selector_map["Or"] = std::make_unique<faiss::IDSelectorOr>(
|
|
472
|
+
selector_map["Range"].get(), selector_map["Batch"].get());
|
|
473
|
+
selector_map["XOr"] = std::make_unique<faiss::IDSelectorXOr>(
|
|
474
|
+
selector_map["Not"].get(), selector_map["Array"].get());
|
|
475
|
+
}
|
|
476
|
+
|
|
477
|
+
void testIDSelectorSearch(
|
|
478
|
+
faiss::Index* index,
|
|
479
|
+
faiss::SearchParameters* search_params,
|
|
480
|
+
const std::vector<float>& queryVecs,
|
|
481
|
+
int numQuery,
|
|
482
|
+
int k,
|
|
483
|
+
const std::string& selectorName) {
|
|
484
|
+
FAISS_ASSERT(search_params && search_params->sel);
|
|
485
|
+
std::vector<float> distances(numQuery * k, 0);
|
|
486
|
+
std::vector<faiss::idx_t> labels(numQuery * k, -1);
|
|
487
|
+
index->search(
|
|
488
|
+
numQuery,
|
|
489
|
+
queryVecs.data(),
|
|
490
|
+
k,
|
|
491
|
+
distances.data(),
|
|
492
|
+
labels.data(),
|
|
493
|
+
search_params);
|
|
494
|
+
faiss::IDSelector* selector = search_params->sel;
|
|
495
|
+
for (int i = 0; i < numQuery * k; ++i) {
|
|
496
|
+
if (labels[i] >= 0) {
|
|
497
|
+
EXPECT_TRUE(selector->is_member(labels[i]))
|
|
498
|
+
<< "Label " << labels[i] << " @ " << i << " not in "
|
|
499
|
+
<< selectorName << " selector";
|
|
500
|
+
}
|
|
501
|
+
}
|
|
502
|
+
}
|
|
503
|
+
|
|
432
504
|
} // namespace gpu
|
|
433
505
|
} // namespace faiss
|
|
@@ -13,6 +13,7 @@
|
|
|
13
13
|
#include <gtest/gtest.h>
|
|
14
14
|
#include <cstring>
|
|
15
15
|
#include <initializer_list>
|
|
16
|
+
#include <map>
|
|
16
17
|
#include <memory>
|
|
17
18
|
#include <string>
|
|
18
19
|
#include <vector>
|
|
@@ -141,5 +142,27 @@ void testIVFEquality(A& cpuIndex, B& gpuIndex) {
|
|
|
141
142
|
}
|
|
142
143
|
}
|
|
143
144
|
|
|
145
|
+
/// Run search with the given search_params and verify all returned labels are
|
|
146
|
+
/// members of the selector. Works with any Index that supports search with
|
|
147
|
+
/// SearchParameters.
|
|
148
|
+
void testIDSelectorSearch(
|
|
149
|
+
faiss::Index* index,
|
|
150
|
+
faiss::SearchParameters* search_params,
|
|
151
|
+
const std::vector<float>& queryVecs,
|
|
152
|
+
int numQuery,
|
|
153
|
+
int k,
|
|
154
|
+
const std::string& selectorName);
|
|
155
|
+
|
|
156
|
+
// Structure to hold all IDSelector instances
|
|
157
|
+
struct TestIDSelectorStruct {
|
|
158
|
+
// Storage for selectors that need it
|
|
159
|
+
std::vector<faiss::idx_t> array_ids;
|
|
160
|
+
std::vector<faiss::idx_t> batch_ids;
|
|
161
|
+
std::vector<uint8_t> bitmap;
|
|
162
|
+
std::map<std::string, std::unique_ptr<faiss::IDSelector>> selector_map;
|
|
163
|
+
|
|
164
|
+
explicit TestIDSelectorStruct(int numAdd);
|
|
165
|
+
};
|
|
166
|
+
|
|
144
167
|
} // namespace gpu
|
|
145
168
|
} // namespace faiss
|
|
@@ -36,6 +36,6 @@ namespace faiss::gpu {
|
|
|
36
36
|
void convert_to_bitset(
|
|
37
37
|
faiss::gpu::GpuResources* res,
|
|
38
38
|
const faiss::IDSelector& selector,
|
|
39
|
-
cuvs::core::bitset_view<uint32_t,
|
|
39
|
+
cuvs::core::bitset_view<uint32_t, int64_t> bitset,
|
|
40
40
|
int num_threads = 0);
|
|
41
41
|
} // namespace faiss::gpu
|
|
@@ -27,32 +27,32 @@
|
|
|
27
27
|
#include <faiss/gpu/GpuResources.h>
|
|
28
28
|
#include <faiss/gpu/utils/Tensor.cuh>
|
|
29
29
|
|
|
30
|
-
#include <cuvs/distance/distance.
|
|
30
|
+
#include <cuvs/distance/distance.hpp>
|
|
31
31
|
|
|
32
32
|
#pragma GCC visibility push(default)
|
|
33
33
|
namespace faiss {
|
|
34
34
|
namespace gpu {
|
|
35
35
|
|
|
36
|
-
inline
|
|
36
|
+
inline cuvs::distance::DistanceType metricFaissToCuvs(
|
|
37
37
|
MetricType metric,
|
|
38
38
|
bool exactDistance) {
|
|
39
39
|
switch (metric) {
|
|
40
40
|
case MetricType::METRIC_INNER_PRODUCT:
|
|
41
|
-
return
|
|
41
|
+
return cuvs::distance::DistanceType::InnerProduct;
|
|
42
42
|
case MetricType::METRIC_L2:
|
|
43
|
-
return
|
|
43
|
+
return cuvs::distance::DistanceType::L2Expanded;
|
|
44
44
|
case MetricType::METRIC_L1:
|
|
45
|
-
return
|
|
45
|
+
return cuvs::distance::DistanceType::L1;
|
|
46
46
|
case MetricType::METRIC_Linf:
|
|
47
|
-
return
|
|
47
|
+
return cuvs::distance::DistanceType::Linf;
|
|
48
48
|
case MetricType::METRIC_Lp:
|
|
49
|
-
return
|
|
49
|
+
return cuvs::distance::DistanceType::LpUnexpanded;
|
|
50
50
|
case MetricType::METRIC_Canberra:
|
|
51
|
-
return
|
|
51
|
+
return cuvs::distance::DistanceType::Canberra;
|
|
52
52
|
case MetricType::METRIC_BrayCurtis:
|
|
53
|
-
return
|
|
53
|
+
return cuvs::distance::DistanceType::BrayCurtis;
|
|
54
54
|
case MetricType::METRIC_JensenShannon:
|
|
55
|
-
return
|
|
55
|
+
return cuvs::distance::DistanceType::JensenShannon;
|
|
56
56
|
default:
|
|
57
57
|
RAFT_FAIL("Distance type not supported");
|
|
58
58
|
}
|
|
@@ -71,6 +71,17 @@ idx_t inplaceGatherFilteredRows(
|
|
|
71
71
|
GpuResources* res,
|
|
72
72
|
Tensor<float, 2, true>& vecs,
|
|
73
73
|
Tensor<idx_t, 1, true>& indices);
|
|
74
|
+
|
|
75
|
+
/// Copy uint32_t indices to idx_t, replacing any index >= n with -1.
|
|
76
|
+
/// cuVS CAGRA returns sentinel values (e.g. INT32_MAX) for result slots
|
|
77
|
+
/// where filtered search couldn't find a valid neighbor.
|
|
78
|
+
void sanitizeCuvsIndices(
|
|
79
|
+
GpuResources* res,
|
|
80
|
+
uint32_t* src,
|
|
81
|
+
idx_t* dst,
|
|
82
|
+
size_t count,
|
|
83
|
+
idx_t n);
|
|
84
|
+
|
|
74
85
|
} // namespace gpu
|
|
75
86
|
} // namespace faiss
|
|
76
87
|
#pragma GCC visibility pop
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
// @lint-ignore-every LICENSELINT
|
|
2
|
+
/**
|
|
3
|
+
* Copyright (c) Meta Platforms, Inc. and its affiliates.
|
|
4
|
+
*
|
|
5
|
+
* This source code is licensed under the MIT license found in the
|
|
6
|
+
* LICENSE file in the root directory of this source tree.
|
|
7
|
+
*
|
|
8
|
+
* Unified name for flat GPU index when Metal backend is built.
|
|
9
|
+
* Include this when using the Metal backend for API parity with
|
|
10
|
+
* faiss::gpu::GpuIndexFlat.
|
|
11
|
+
*/
|
|
12
|
+
|
|
13
|
+
#pragma once
|
|
14
|
+
|
|
15
|
+
#include <faiss/gpu_metal/MetalIndexFlat.h>
|
|
16
|
+
|
|
17
|
+
namespace faiss {
|
|
18
|
+
|
|
19
|
+
/// When FAISS is built with Metal backend, GpuIndexFlat is MetalIndexFlat.
|
|
20
|
+
using GpuIndexFlat = gpu_metal::MetalIndexFlat;
|
|
21
|
+
|
|
22
|
+
} // namespace faiss
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
// @lint-ignore-every LICENSELINT
|
|
2
|
+
/**
|
|
3
|
+
* Copyright (c) Meta Platforms, Inc. and its affiliates.
|
|
4
|
+
*
|
|
5
|
+
* This source code is licensed under the MIT license found in the
|
|
6
|
+
* LICENSE file in the root directory of this source tree.
|
|
7
|
+
*
|
|
8
|
+
* Clone CPU <-> Metal GPU. Mirrors GpuCloner roles for Metal backend.
|
|
9
|
+
*/
|
|
10
|
+
|
|
11
|
+
#pragma once
|
|
12
|
+
|
|
13
|
+
#include <faiss/Index.h>
|
|
14
|
+
|
|
15
|
+
namespace faiss {
|
|
16
|
+
namespace gpu_metal {
|
|
17
|
+
|
|
18
|
+
class StandardMetalResources;
|
|
19
|
+
|
|
20
|
+
/// Returns the number of Metal "devices" (1 if Metal is available, else 0).
|
|
21
|
+
int get_num_gpus();
|
|
22
|
+
|
|
23
|
+
/// Clone a CPU index to Metal GPU. Supports IndexFlat, IndexFlatL2,
|
|
24
|
+
/// IndexFlatIP. device must be 0. Caller owns the returned index.
|
|
25
|
+
faiss::Index* index_cpu_to_metal_gpu(
|
|
26
|
+
StandardMetalResources* res,
|
|
27
|
+
int device,
|
|
28
|
+
const faiss::Index* index);
|
|
29
|
+
|
|
30
|
+
/// Copy a Metal index back to CPU. Supports MetalIndexFlat -> IndexFlat.
|
|
31
|
+
/// Caller owns the returned index.
|
|
32
|
+
faiss::Index* index_metal_gpu_to_cpu(const faiss::Index* index);
|
|
33
|
+
|
|
34
|
+
} // namespace gpu_metal
|
|
35
|
+
} // namespace faiss
|