faiss 0.6.0 → 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 +4 -0
- data/ext/faiss/extconf.rb +2 -1
- data/ext/faiss/{index_rb.cpp → index.cpp} +1 -1
- data/ext/faiss/index_binary.cpp +1 -1
- data/ext/faiss/kmeans.cpp +1 -1
- data/ext/faiss/pca_matrix.cpp +1 -1
- data/ext/faiss/product_quantizer.cpp +1 -1
- data/ext/faiss/{utils_rb.cpp → utils.cpp} +1 -1
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +93 -80
- data/vendor/faiss/faiss/Clustering.cpp +39 -240
- data/vendor/faiss/faiss/Clustering.h +6 -0
- data/vendor/faiss/faiss/IVFlib.cpp +41 -21
- data/vendor/faiss/faiss/Index.cpp +6 -5
- data/vendor/faiss/faiss/Index.h +5 -5
- data/vendor/faiss/faiss/Index2Layer.cpp +37 -53
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +49 -37
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +36 -34
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +4 -1
- data/vendor/faiss/faiss/IndexBinary.cpp +5 -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 +84 -92
- 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 +87 -415
- data/vendor/faiss/faiss/IndexFastScan.cpp +72 -109
- data/vendor/faiss/faiss/IndexFastScan.h +25 -23
- data/vendor/faiss/faiss/IndexFlat.cpp +27 -20
- data/vendor/faiss/faiss/IndexFlat.h +21 -18
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +42 -19
- data/vendor/faiss/faiss/IndexHNSW.cpp +283 -145
- data/vendor/faiss/faiss/IndexHNSW.h +16 -2
- data/vendor/faiss/faiss/IndexIDMap.cpp +25 -21
- data/vendor/faiss/faiss/IndexIDMap.h +9 -7
- data/vendor/faiss/faiss/IndexIVF.cpp +465 -362
- data/vendor/faiss/faiss/IndexIVF.h +33 -12
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +77 -74
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +96 -93
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +4 -1
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +357 -238
- data/vendor/faiss/faiss/IndexIVFFastScan.h +42 -41
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +36 -68
- data/vendor/faiss/faiss/IndexIVFFlat.h +32 -0
- data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +53 -30
- data/vendor/faiss/faiss/IndexIVFFlatPanorama.h +3 -1
- data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +18 -15
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +71 -843
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +151 -121
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +3 -0
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +21 -17
- data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +26 -39
- data/vendor/faiss/faiss/IndexIVFRaBitQ.h +2 -1
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +475 -476
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +248 -93
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +41 -127
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +1 -1
- data/vendor/faiss/faiss/IndexLSH.cpp +36 -19
- data/vendor/faiss/faiss/IndexLattice.cpp +13 -13
- data/vendor/faiss/faiss/IndexNNDescent.cpp +36 -21
- data/vendor/faiss/faiss/IndexNNDescent.h +2 -2
- data/vendor/faiss/faiss/IndexNSG.cpp +39 -23
- data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +31 -11
- data/vendor/faiss/faiss/IndexPQ.cpp +128 -221
- 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 +11 -36
- data/vendor/faiss/faiss/IndexRaBitQ.h +2 -1
- data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +41 -277
- data/vendor/faiss/faiss/IndexRaBitQFastScan.h +183 -27
- data/vendor/faiss/faiss/IndexRefine.cpp +30 -25
- 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 +10 -9
- 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 +14 -7
- data/vendor/faiss/faiss/SuperKMeans.cpp +656 -0
- data/vendor/faiss/faiss/SuperKMeans.h +97 -0
- data/vendor/faiss/faiss/VectorTransform.cpp +237 -149
- data/vendor/faiss/faiss/VectorTransform.h +16 -16
- data/vendor/faiss/faiss/build.cpp +23 -0
- data/vendor/faiss/faiss/build.h +15 -0
- data/vendor/faiss/faiss/clone_index.cpp +48 -47
- 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/factory_tools.cpp +5 -0
- 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 +29 -25
- 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 -0
- 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 +16 -16
- data/vendor/faiss/faiss/impl/CodePacker.cpp +3 -3
- data/vendor/faiss/faiss/impl/CodePackerRaBitQ.cpp +1 -1
- data/vendor/faiss/faiss/impl/DistanceComputer.h +8 -8
- data/vendor/faiss/faiss/impl/FaissAssert.h +6 -3
- data/vendor/faiss/faiss/impl/FaissException.h +50 -3
- data/vendor/faiss/faiss/impl/HNSW.cpp +92 -317
- data/vendor/faiss/faiss/impl/HNSW.h +13 -34
- 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 +82 -77
- data/vendor/faiss/faiss/impl/NNDescent.cpp +62 -25
- data/vendor/faiss/faiss/impl/NNDescent.h +6 -2
- data/vendor/faiss/faiss/impl/NSG.cpp +38 -21
- data/vendor/faiss/faiss/impl/NSG.h +4 -4
- data/vendor/faiss/faiss/impl/Panorama.cpp +23 -6
- data/vendor/faiss/faiss/impl/Panorama.h +258 -87
- 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 +46 -32
- 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 +30 -23
- data/vendor/faiss/faiss/impl/Quantizer.h +2 -2
- data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +55 -49
- data/vendor/faiss/faiss/impl/RaBitQUtils.h +65 -0
- data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +296 -283
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +26 -23
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +1 -1
- data/vendor/faiss/faiss/impl/ResultHandler.h +99 -75
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +52 -4
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +27 -1
- data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +14 -11
- data/vendor/faiss/faiss/impl/VisitedTable.h +7 -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 +8 -3
- 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} +169 -2
- 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 -356
- 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} +282 -134
- 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 +1132 -45
- data/vendor/faiss/faiss/impl/index_read_utils.h +1 -1
- data/vendor/faiss/faiss/impl/index_write.cpp +95 -13
- 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 +37 -23
- data/vendor/faiss/faiss/impl/lattice_Zn.h +6 -6
- data/vendor/faiss/faiss/impl/mapped_io.cpp +6 -6
- 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/pq_code_distance/{pq_code_distance-avx2.cpp → pq_code_distance-avx2.h} +9 -13
- data/vendor/faiss/faiss/impl/pq_code_distance/{pq_code_distance-avx512.cpp → pq_code_distance-avx512.h} +9 -57
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +29 -111
- 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 +238 -5
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +5 -7
- data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +68 -0
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +311 -477
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +1 -1
- data/vendor/faiss/faiss/impl/scalar_quantizer/codecs.h +1 -1
- data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +3 -2
- data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +102 -11
- data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h +27 -1
- data/vendor/faiss/faiss/impl/scalar_quantizer/similarities.h +3 -3
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +148 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +167 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +59 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +163 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +311 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/training.cpp +192 -8
- data/vendor/faiss/faiss/impl/scalar_quantizer/training.h +12 -0
- data/vendor/faiss/faiss/impl/simd_dispatch.h +100 -66
- data/vendor/faiss/faiss/impl/simdlib/simdlib.h +57 -0
- data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_avx2.h +264 -172
- 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 +270 -218
- 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 +86 -18
- data/vendor/faiss/faiss/index_io.h +24 -0
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +66 -16
- 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 +13 -13
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +1 -1
- data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +1 -1
- data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +2 -2
- 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 +18 -2
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +1 -1
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +12 -3
- data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +7 -2
- data/vendor/faiss/faiss/utils/Heap.cpp +10 -10
- data/vendor/faiss/faiss/utils/NeuralNet.cpp +47 -36
- 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 +390 -560
- data/vendor/faiss/faiss/utils/distances.h +20 -1
- data/vendor/faiss/faiss/utils/distances_dispatch.h +117 -37
- 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 +5 -177
- data/vendor/faiss/faiss/utils/extra_distances.cpp +9 -8
- data/vendor/faiss/faiss/utils/extra_distances.h +32 -6
- 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 +2 -2
- 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 +57 -536
- 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 +5 -1
- data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +213 -4
- data/vendor/faiss/faiss/utils/simd_impl/distances_autovec-inl.h +163 -10
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +250 -4
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +7 -4
- 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 +2 -1
- 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 +17 -5
- data/vendor/faiss/faiss/utils/simd_levels.h +93 -1
- data/vendor/faiss/faiss/utils/sorting.cpp +48 -36
- data/vendor/faiss/faiss/utils/utils.cpp +5 -5
- data/vendor/faiss/faiss/utils/utils.h +3 -3
- metadata +119 -34
- data/vendor/faiss/faiss/impl/RaBitQStats.cpp +0 -29
- data/vendor/faiss/faiss/impl/RaBitQStats.h +0 -56
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +0 -224
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +0 -230
- 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 -235
- 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 -449
- 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 -365
- /data/ext/faiss/{utils_rb.h → utils.h} +0 -0
|
@@ -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
|
|
@@ -0,0 +1,40 @@
|
|
|
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
|
+
* Objective-C++ header. Runs L2/IP distance + top-k via Metal compute.
|
|
9
|
+
*/
|
|
10
|
+
|
|
11
|
+
#pragma once
|
|
12
|
+
|
|
13
|
+
#import <Metal/Metal.h>
|
|
14
|
+
|
|
15
|
+
#include <cstddef>
|
|
16
|
+
|
|
17
|
+
namespace faiss {
|
|
18
|
+
namespace gpu_metal {
|
|
19
|
+
|
|
20
|
+
/// Runs GPU search: distance matrix (L2 or IP) then top-k. Uses shared buffers
|
|
21
|
+
/// (queries, vectors, outDistances, outIndices). outIndices are int32
|
|
22
|
+
/// (0..nb-1). Maximum k supported by the GPU top-k kernel (256).
|
|
23
|
+
int getMetalFlatSearchMaxK();
|
|
24
|
+
|
|
25
|
+
/// Returns true on success; false if pipeline creation failed.
|
|
26
|
+
bool runFlatSearchGPU(
|
|
27
|
+
id<MTLDevice> device,
|
|
28
|
+
id<MTLCommandQueue> queue,
|
|
29
|
+
id<MTLBuffer> queries, // (nq * d) float, row-major
|
|
30
|
+
id<MTLBuffer> vectors, // (nb * d) float, row-major
|
|
31
|
+
int nq,
|
|
32
|
+
int nb,
|
|
33
|
+
int d,
|
|
34
|
+
int k,
|
|
35
|
+
bool isL2, // true = L2 squared, false = inner product
|
|
36
|
+
id<MTLBuffer> outDistances, // (nq * k) float
|
|
37
|
+
id<MTLBuffer> outIndices); // (nq * k) int32
|
|
38
|
+
|
|
39
|
+
} // namespace gpu_metal
|
|
40
|
+
} // namespace faiss
|
|
@@ -0,0 +1,51 @@
|
|
|
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
|
+
* Objective-C++ header (uses MetalResources).
|
|
9
|
+
*/
|
|
10
|
+
|
|
11
|
+
#pragma once
|
|
12
|
+
|
|
13
|
+
#include <faiss/Index.h>
|
|
14
|
+
#include <faiss/gpu_metal/MetalResources.h>
|
|
15
|
+
#include <memory>
|
|
16
|
+
|
|
17
|
+
namespace faiss {
|
|
18
|
+
namespace gpu_metal {
|
|
19
|
+
|
|
20
|
+
/// Configuration for Metal index (mirrors GpuIndexConfig roles).
|
|
21
|
+
struct MetalIndexConfig {
|
|
22
|
+
int device = 0;
|
|
23
|
+
};
|
|
24
|
+
|
|
25
|
+
/// Base class for Metal-backed indexes. Mirrors faiss::gpu::GpuIndex.
|
|
26
|
+
class MetalIndex : public faiss::Index {
|
|
27
|
+
public:
|
|
28
|
+
MetalIndex(
|
|
29
|
+
std::shared_ptr<MetalResources> resources,
|
|
30
|
+
int dims,
|
|
31
|
+
faiss::MetricType metric,
|
|
32
|
+
float metricArg,
|
|
33
|
+
MetalIndexConfig config = MetalIndexConfig());
|
|
34
|
+
|
|
35
|
+
int getDevice() const {
|
|
36
|
+
return config_.device;
|
|
37
|
+
}
|
|
38
|
+
std::shared_ptr<MetalResources> getResources() {
|
|
39
|
+
return resources_;
|
|
40
|
+
}
|
|
41
|
+
std::shared_ptr<const MetalResources> getResources() const {
|
|
42
|
+
return resources_;
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
protected:
|
|
46
|
+
std::shared_ptr<MetalResources> resources_;
|
|
47
|
+
MetalIndexConfig config_;
|
|
48
|
+
};
|
|
49
|
+
|
|
50
|
+
} // namespace gpu_metal
|
|
51
|
+
} // namespace faiss
|