├── .clang-format ├── .dockerignore ├── .github ├── ISSUE_TEMPLATE.md ├── actions │ ├── build_cmake │ │ └── action.yml │ └── build_conda │ │ └── action.yml └── workflows │ ├── autoclose.yml │ ├── build-pull-request.yml │ ├── build-release.yml │ ├── build.yml │ ├── nightly.yml │ ├── publish-docs.yml │ ├── retry_build.yml │ └── update-doxygen.yml ├── .gitignore ├── CHANGELOG.md ├── CMakeLists.txt ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── Doxyfile ├── INSTALL.md ├── LICENSE ├── README.md ├── benchs ├── CMakeLists.txt ├── README.md ├── bench_6bit_codec.cpp ├── bench_all_ivf │ ├── README.md │ ├── bench_all_ivf.py │ ├── bench_kmeans.py │ ├── cmp_with_scann.py │ ├── datasets_oss.py │ ├── make_groundtruth.py │ ├── parse_bench_all_ivf.py │ └── run_on_cluster_generic.bash ├── bench_big_batch_ivf.py ├── bench_cppcontrib_sa_decode.cpp ├── bench_for_interrupt.py ├── bench_fw │ ├── __init__.py │ ├── benchmark.py │ ├── benchmark_io.py │ ├── descriptors.py │ ├── index.py │ ├── optimize.py │ └── utils.py ├── bench_fw_codecs.py ├── bench_fw_ivf.py ├── bench_fw_notebook.ipynb ├── bench_fw_optimize.py ├── bench_fw_range.py ├── bench_gpu_1bn.py ├── bench_gpu_sift1m.py ├── bench_hamming_computer.cpp ├── bench_hamming_knn.py ├── bench_heap_replace.cpp ├── bench_hnsw.py ├── bench_hybrid_cpu_gpu.py ├── bench_index_flat.py ├── bench_index_pq.py ├── bench_ivf_fastscan.py ├── bench_ivf_fastscan_single_query.py ├── bench_ivf_selector.cpp ├── bench_ivfflat_cuvs.py ├── bench_ivfpq_cuvs.py ├── bench_pairwise_distances.py ├── bench_partition.py ├── bench_polysemous_1bn.py ├── bench_polysemous_sift1m.py ├── bench_pq_tables.py ├── bench_pq_transposed_centroid_table.py ├── bench_quantizer.py ├── bench_rabitq.py ├── bench_scalar_quantizer.py ├── bench_vector_ops.py ├── datasets.py ├── distributed_ondisk │ ├── README.md │ ├── combined_index.py │ ├── distributed_kmeans.py │ ├── distributed_query_demo.py │ ├── make_index_vslice.py │ ├── make_trained_index.py │ ├── merge_to_ondisk.py │ ├── run_on_cluster.bash │ └── search_server.py ├── kmeans_mnist.py └── link_and_code │ └── README.md ├── c_api ├── AutoTune_c.cpp ├── AutoTune_c.h ├── CMakeLists.txt ├── Clustering_c.cpp ├── Clustering_c.h ├── INSTALL.md ├── IndexBinaryIVF_c.cpp ├── IndexBinaryIVF_c.h ├── IndexBinary_c.cpp ├── IndexBinary_c.h ├── IndexFlat_c.cpp ├── IndexFlat_c.h ├── IndexIVFFlat_c.cpp ├── IndexIVFFlat_c.h ├── IndexIVF_c.cpp ├── IndexIVF_c.h ├── IndexLSH_c.cpp ├── IndexLSH_c.h ├── IndexPreTransform_c.cpp ├── IndexPreTransform_c.h ├── IndexReplicas_c.cpp ├── IndexReplicas_c.h ├── IndexScalarQuantizer_c.cpp ├── IndexScalarQuantizer_c.h ├── IndexShards_c.cpp ├── IndexShards_c.h ├── Index_c.cpp ├── Index_c.h ├── MetaIndexes_c.cpp ├── MetaIndexes_c.h ├── VectorTransform_c.cpp ├── VectorTransform_c.h ├── clone_index_c.cpp ├── clone_index_c.h ├── error_c.h ├── error_impl.cpp ├── error_impl.h ├── example_c.c ├── faiss_c.h ├── gpu │ ├── CMakeLists.txt │ ├── DeviceUtils_c.cpp │ ├── DeviceUtils_c.h │ ├── GpuAutoTune_c.cpp │ ├── GpuAutoTune_c.h │ ├── GpuClonerOptions_c.cpp │ ├── GpuClonerOptions_c.h │ ├── GpuIndex_c.cpp │ ├── GpuIndex_c.h │ ├── GpuIndicesOptions_c.h │ ├── GpuResources_c.cpp │ ├── GpuResources_c.h │ ├── StandardGpuResources_c.cpp │ ├── StandardGpuResources_c.h │ ├── example_gpu_c.c │ └── macros_impl.h ├── impl │ ├── AuxIndexStructures_c.cpp │ ├── AuxIndexStructures_c.h │ ├── io_c.cpp │ └── io_c.h ├── index_factory_c.cpp ├── index_factory_c.h ├── index_io_c.cpp ├── index_io_c.h ├── macros_impl.h └── utils │ ├── distances_c.cpp │ ├── distances_c.h │ ├── utils_c.cpp │ └── utils_c.h ├── cmake ├── FindMKL.cmake ├── faiss-config.cmake.in ├── link_to_faiss_lib.cmake └── thirdparty │ └── fetch_rapids.cmake ├── conda ├── conda_build_config.yaml ├── faiss-gpu-cuvs │ ├── build-lib.sh │ ├── build-pkg.sh │ ├── meta.yaml │ └── test_cpu_dispatch.sh ├── faiss-gpu │ ├── build-lib.sh │ ├── build-pkg.sh │ ├── meta.yaml │ └── test_cpu_dispatch.sh └── faiss │ ├── build-lib-arm64.sh │ ├── build-lib-osx.sh │ ├── build-lib.bat │ ├── build-lib.sh │ ├── build-pkg-arm64.sh │ ├── build-pkg-osx.sh │ ├── build-pkg.bat │ ├── build-pkg.sh │ ├── meta.yaml │ └── test_cpu_dispatch.sh ├── contrib ├── README.md ├── __init__.py ├── big_batch_search.py ├── client_server.py ├── clustering.py ├── datasets.py ├── evaluation.py ├── exhaustive_search.py ├── factory_tools.py ├── inspect_tools.py ├── ivf_tools.py ├── ondisk.py ├── rpc.py ├── torch │ ├── README.md │ ├── __init__.py │ ├── clustering.py │ └── quantization.py ├── torch_utils.py └── vecs_io.py ├── demos ├── CMakeLists.txt ├── README.md ├── demo_auto_tune.py ├── demo_client_server_ivf.py ├── demo_distributed_kmeans_torch.py ├── demo_imi_flat.cpp ├── demo_imi_pq.cpp ├── demo_ivfpq_indexing.cpp ├── demo_nndescent.cpp ├── demo_ondisk_ivf.py ├── demo_qinco.py ├── demo_residual_quantizer.cpp ├── demo_sift1M.cpp ├── demo_weighted_kmeans.cpp ├── index_pq_flat_separate_codes_from_codebook.py ├── offline_ivf │ ├── README.md │ ├── __init__.py │ ├── config_ssnpp.yaml │ ├── create_sharded_ssnpp_files.py │ ├── dataset.py │ ├── generate_config.py │ ├── offline_ivf.py │ ├── run.py │ ├── tests │ │ ├── test_iterate_input.py │ │ ├── test_offline_ivf.py │ │ └── testing_utils.py │ └── utils.py └── rocksdb_ivf │ ├── CMakeLists.txt │ ├── README.md │ ├── RocksDBInvertedLists.cpp │ ├── RocksDBInvertedLists.h │ └── demo_rocksdb_ivf.cpp ├── faiss ├── .flake8 ├── AutoTune.cpp ├── AutoTune.h ├── CMakeLists.txt ├── Clustering.cpp ├── Clustering.h ├── IVFlib.cpp ├── IVFlib.h ├── Index.cpp ├── Index.h ├── Index2Layer.cpp ├── Index2Layer.h ├── IndexAdditiveQuantizer.cpp ├── IndexAdditiveQuantizer.h ├── IndexAdditiveQuantizerFastScan.cpp ├── IndexAdditiveQuantizerFastScan.h ├── IndexBinary.cpp ├── IndexBinary.h ├── IndexBinaryFlat.cpp ├── IndexBinaryFlat.h ├── IndexBinaryFromFloat.cpp ├── IndexBinaryFromFloat.h ├── IndexBinaryHNSW.cpp ├── IndexBinaryHNSW.h ├── IndexBinaryHash.cpp ├── IndexBinaryHash.h ├── IndexBinaryIVF.cpp ├── IndexBinaryIVF.h ├── IndexFastScan.cpp ├── IndexFastScan.h ├── IndexFlat.cpp ├── IndexFlat.h ├── IndexFlatCodes.cpp ├── IndexFlatCodes.h ├── IndexHNSW.cpp ├── IndexHNSW.h ├── IndexIDMap.cpp ├── IndexIDMap.h ├── IndexIVF.cpp ├── IndexIVF.h ├── IndexIVFAdditiveQuantizer.cpp ├── IndexIVFAdditiveQuantizer.h ├── IndexIVFAdditiveQuantizerFastScan.cpp ├── IndexIVFAdditiveQuantizerFastScan.h ├── IndexIVFFastScan.cpp ├── IndexIVFFastScan.h ├── IndexIVFFlat.cpp ├── IndexIVFFlat.h ├── IndexIVFIndependentQuantizer.cpp ├── IndexIVFIndependentQuantizer.h ├── IndexIVFPQ.cpp ├── IndexIVFPQ.h ├── IndexIVFPQFastScan.cpp ├── IndexIVFPQFastScan.h ├── IndexIVFPQR.cpp ├── IndexIVFPQR.h ├── IndexIVFRaBitQ.cpp ├── IndexIVFRaBitQ.h ├── IndexIVFSpectralHash.cpp ├── IndexIVFSpectralHash.h ├── IndexLSH.cpp ├── IndexLSH.h ├── IndexLattice.cpp ├── IndexLattice.h ├── IndexNNDescent.cpp ├── IndexNNDescent.h ├── IndexNSG.cpp ├── IndexNSG.h ├── IndexNeuralNetCodec.cpp ├── IndexNeuralNetCodec.h ├── IndexPQ.cpp ├── IndexPQ.h ├── IndexPQFastScan.cpp ├── IndexPQFastScan.h ├── IndexPreTransform.cpp ├── IndexPreTransform.h ├── IndexRaBitQ.cpp ├── IndexRaBitQ.h ├── IndexRefine.cpp ├── IndexRefine.h ├── IndexReplicas.cpp ├── IndexReplicas.h ├── IndexRowwiseMinMax.cpp ├── IndexRowwiseMinMax.h ├── IndexScalarQuantizer.cpp ├── IndexScalarQuantizer.h ├── IndexShards.cpp ├── IndexShards.h ├── IndexShardsIVF.cpp ├── IndexShardsIVF.h ├── MatrixStats.cpp ├── MatrixStats.h ├── MetaIndexes.cpp ├── MetaIndexes.h ├── MetricType.h ├── VectorTransform.cpp ├── VectorTransform.h ├── clone_index.cpp ├── clone_index.h ├── cppcontrib │ ├── SaDecodeKernels.h │ ├── detail │ │ ├── CoarseBitType.h │ │ └── UintReader.h │ ├── docker_dev │ │ └── Dockerfile │ ├── factory_tools.cpp │ ├── factory_tools.h │ └── sa_decode │ │ ├── Level2-avx2-inl.h │ │ ├── Level2-inl.h │ │ ├── Level2-neon-inl.h │ │ ├── MinMax-inl.h │ │ ├── MinMaxFP16-inl.h │ │ ├── PQ-avx2-inl.h │ │ ├── PQ-inl.h │ │ └── PQ-neon-inl.h ├── gpu │ ├── CMakeLists.txt │ ├── GpuAutoTune.cpp │ ├── GpuAutoTune.h │ ├── GpuCloner.cpp │ ├── GpuCloner.h │ ├── GpuClonerOptions.h │ ├── GpuDistance.cu │ ├── GpuDistance.h │ ├── GpuFaissAssert.h │ ├── GpuIcmEncoder.cu │ ├── GpuIcmEncoder.h │ ├── GpuIndex.cu │ ├── GpuIndex.h │ ├── GpuIndexBinaryFlat.cu │ ├── GpuIndexBinaryFlat.h │ ├── GpuIndexCagra.cu │ ├── GpuIndexCagra.h │ ├── GpuIndexFlat.cu │ ├── GpuIndexFlat.h │ ├── GpuIndexIVF.cu │ ├── GpuIndexIVF.h │ ├── GpuIndexIVFFlat.cu │ ├── GpuIndexIVFFlat.h │ ├── GpuIndexIVFPQ.cu │ ├── GpuIndexIVFPQ.h │ ├── GpuIndexIVFScalarQuantizer.cu │ ├── GpuIndexIVFScalarQuantizer.h │ ├── GpuIndicesOptions.h │ ├── GpuResources.cpp │ ├── GpuResources.h │ ├── StandardGpuResources.cpp │ ├── StandardGpuResources.h │ ├── hipify.sh │ ├── impl │ │ ├── BinaryDistance.cu │ │ ├── BinaryDistance.cuh │ │ ├── BinaryFlatIndex.cu │ │ ├── BinaryFlatIndex.cuh │ │ ├── BroadcastSum.cu │ │ ├── BroadcastSum.cuh │ │ ├── CuvsCagra.cu │ │ ├── CuvsCagra.cuh │ │ ├── CuvsFlatIndex.cu │ │ ├── CuvsFlatIndex.cuh │ │ ├── CuvsIVFFlat.cu │ │ ├── CuvsIVFFlat.cuh │ │ ├── CuvsIVFPQ.cu │ │ ├── CuvsIVFPQ.cuh │ │ ├── Distance.cu │ │ ├── Distance.cuh │ │ ├── DistanceUtils.cuh │ │ ├── FlatIndex.cu │ │ ├── FlatIndex.cuh │ │ ├── GeneralDistance.cuh │ │ ├── GpuScalarQuantizer.cuh │ │ ├── IVFAppend.cu │ │ ├── IVFAppend.cuh │ │ ├── IVFBase.cu │ │ ├── IVFBase.cuh │ │ ├── IVFFlat.cu │ │ ├── IVFFlat.cuh │ │ ├── IVFFlatScan.cu │ │ ├── IVFFlatScan.cuh │ │ ├── IVFInterleaved.cu │ │ ├── IVFInterleaved.cuh │ │ ├── IVFPQ.cu │ │ ├── IVFPQ.cuh │ │ ├── IVFUtils.cu │ │ ├── IVFUtils.cuh │ │ ├── IVFUtilsSelect1.cu │ │ ├── IVFUtilsSelect2.cu │ │ ├── IcmEncoder.cu │ │ ├── IcmEncoder.cuh │ │ ├── IndexUtils.cu │ │ ├── IndexUtils.h │ │ ├── InterleavedCodes.cpp │ │ ├── InterleavedCodes.h │ │ ├── L2Norm.cu │ │ ├── L2Norm.cuh │ │ ├── L2Select.cu │ │ ├── L2Select.cuh │ │ ├── PQCodeDistances-inl.cuh │ │ ├── PQCodeDistances.cuh │ │ ├── PQCodeLoad.cuh │ │ ├── PQScanMultiPassNoPrecomputed-inl.cuh │ │ ├── PQScanMultiPassNoPrecomputed.cuh │ │ ├── PQScanMultiPassPrecomputed.cu │ │ ├── PQScanMultiPassPrecomputed.cuh │ │ ├── RemapIndices.cpp │ │ ├── RemapIndices.h │ │ ├── VectorResidual.cu │ │ ├── VectorResidual.cuh │ │ └── scan │ │ │ ├── IVFInterleavedImpl.cuh │ │ │ └── IVFInterleavedScanKernelTemplate.cu │ ├── perf │ │ ├── IndexWrapper-inl.h │ │ ├── IndexWrapper.h │ │ ├── PerfBinaryFlat.cu │ │ ├── PerfClustering.cpp │ │ ├── PerfFlat.cu │ │ ├── PerfIVFFlat.cu │ │ ├── PerfIVFPQ.cu │ │ ├── PerfIVFPQAdd.cpp │ │ ├── PerfSelect.cu │ │ ├── WriteIndex.cpp │ │ └── slow.py │ ├── test │ │ ├── CMakeLists.txt │ │ ├── TestCodePacking.cpp │ │ ├── TestGpuDistance.cu │ │ ├── TestGpuIcmEncoder.cpp │ │ ├── TestGpuIndexBinaryFlat.cpp │ │ ├── TestGpuIndexCagra.cu │ │ ├── TestGpuIndexFlat.cpp │ │ ├── TestGpuIndexIVFFlat.cpp │ │ ├── TestGpuIndexIVFPQ.cpp │ │ ├── TestGpuIndexIVFScalarQuantizer.cpp │ │ ├── TestGpuMemoryException.cpp │ │ ├── TestGpuResidualQuantizer.cpp │ │ ├── TestGpuSelect.cu │ │ ├── TestUtils.cpp │ │ ├── TestUtils.h │ │ ├── demo_ivfpq_indexing_gpu.cpp │ │ ├── test_cagra.py │ │ ├── test_contrib_gpu.py │ │ ├── test_cuvs.py │ │ ├── test_gpu_basics.py │ │ ├── test_gpu_index.py │ │ ├── test_gpu_index_ivfflat.py │ │ ├── test_gpu_index_ivfsq.py │ │ ├── test_gpu_index_refs.py │ │ ├── test_gpu_index_serialize.py │ │ ├── test_index_cpu_to_gpu.py │ │ ├── test_multi_gpu.py │ │ └── torch_test_contrib_gpu.py │ └── utils │ │ ├── BlockSelectFloat.cu │ │ ├── BlockSelectKernel.cuh │ │ ├── Comparators.cuh │ │ ├── ConversionOperators.cuh │ │ ├── CopyUtils.cuh │ │ ├── CuvsUtils.cu │ │ ├── CuvsUtils.h │ │ ├── DeviceDefs.cuh │ │ ├── DeviceTensor-inl.cuh │ │ ├── DeviceTensor.cuh │ │ ├── DeviceUtils.cu │ │ ├── DeviceUtils.h │ │ ├── DeviceVector.cuh │ │ ├── Float16.cuh │ │ ├── HostTensor-inl.cuh │ │ ├── HostTensor.cuh │ │ ├── Limits.cuh │ │ ├── LoadStoreOperators.cuh │ │ ├── MathOperators.cuh │ │ ├── MatrixMult-inl.cuh │ │ ├── MatrixMult.cuh │ │ ├── MergeNetworkBlock.cuh │ │ ├── MergeNetworkUtils.cuh │ │ ├── MergeNetworkWarp.cuh │ │ ├── NoTypeTensor.cuh │ │ ├── Pair.cuh │ │ ├── PtxUtils.cuh │ │ ├── ReductionOperators.cuh │ │ ├── Reductions.cuh │ │ ├── Select.cuh │ │ ├── StackDeviceMemory.cpp │ │ ├── StackDeviceMemory.h │ │ ├── StaticUtils.h │ │ ├── Tensor-inl.cuh │ │ ├── Tensor.cuh │ │ ├── ThrustUtils.cuh │ │ ├── Timer.cpp │ │ ├── Timer.h │ │ ├── Transpose.cuh │ │ ├── WarpPackedBits.cuh │ │ ├── WarpSelectFloat.cu │ │ ├── WarpSelectKernel.cuh │ │ ├── WarpShuffles.cuh │ │ ├── blockselect │ │ ├── BlockSelectFloat1.cu │ │ ├── BlockSelectFloat128.cu │ │ ├── BlockSelectFloat256.cu │ │ ├── BlockSelectFloat32.cu │ │ ├── BlockSelectFloat64.cu │ │ ├── BlockSelectFloatF1024.cu │ │ ├── BlockSelectFloatF2048.cu │ │ ├── BlockSelectFloatF512.cu │ │ ├── BlockSelectFloatT1024.cu │ │ ├── BlockSelectFloatT2048.cu │ │ ├── BlockSelectFloatT512.cu │ │ └── BlockSelectImpl.cuh │ │ └── warpselect │ │ ├── WarpSelectFloat1.cu │ │ ├── WarpSelectFloat128.cu │ │ ├── WarpSelectFloat256.cu │ │ ├── WarpSelectFloat32.cu │ │ ├── WarpSelectFloat64.cu │ │ ├── WarpSelectFloatF1024.cu │ │ ├── WarpSelectFloatF2048.cu │ │ ├── WarpSelectFloatF512.cu │ │ ├── WarpSelectFloatT1024.cu │ │ ├── WarpSelectFloatT2048.cu │ │ ├── WarpSelectFloatT512.cu │ │ └── WarpSelectImpl.cuh ├── impl │ ├── AdditiveQuantizer.cpp │ ├── AdditiveQuantizer.h │ ├── AuxIndexStructures.cpp │ ├── AuxIndexStructures.h │ ├── CodePacker.cpp │ ├── CodePacker.h │ ├── DistanceComputer.h │ ├── FaissAssert.h │ ├── FaissException.cpp │ ├── FaissException.h │ ├── HNSW.cpp │ ├── HNSW.h │ ├── IDSelector.cpp │ ├── IDSelector.h │ ├── LocalSearchQuantizer.cpp │ ├── LocalSearchQuantizer.h │ ├── LookupTableScaler.h │ ├── NNDescent.cpp │ ├── NNDescent.h │ ├── NSG.cpp │ ├── NSG.h │ ├── PolysemousTraining.cpp │ ├── PolysemousTraining.h │ ├── ProductAdditiveQuantizer.cpp │ ├── ProductAdditiveQuantizer.h │ ├── ProductQuantizer-inl.h │ ├── ProductQuantizer.cpp │ ├── ProductQuantizer.h │ ├── Quantizer.h │ ├── RaBitQuantizer.cpp │ ├── RaBitQuantizer.h │ ├── ResidualQuantizer.cpp │ ├── ResidualQuantizer.h │ ├── ResultHandler.h │ ├── ScalarQuantizer.cpp │ ├── ScalarQuantizer.h │ ├── ThreadedIndex-inl.h │ ├── ThreadedIndex.h │ ├── code_distance │ │ ├── code_distance-avx2.h │ │ ├── code_distance-avx512.h │ │ ├── code_distance-generic.h │ │ ├── code_distance-sve.h │ │ └── code_distance.h │ ├── index_read.cpp │ ├── index_read_utils.h │ ├── index_write.cpp │ ├── io.cpp │ ├── io.h │ ├── io_macros.h │ ├── kmeans1d.cpp │ ├── kmeans1d.h │ ├── lattice_Zn.cpp │ ├── lattice_Zn.h │ ├── mapped_io.cpp │ ├── mapped_io.h │ ├── maybe_owned_vector.h │ ├── platform_macros.h │ ├── pq4_fast_scan.cpp │ ├── pq4_fast_scan.h │ ├── pq4_fast_scan_search_1.cpp │ ├── pq4_fast_scan_search_qbs.cpp │ ├── residual_quantizer_encode_steps.cpp │ ├── residual_quantizer_encode_steps.h │ ├── simd_result_handlers.h │ ├── zerocopy_io.cpp │ └── zerocopy_io.h ├── index_factory.cpp ├── index_factory.h ├── index_io.h ├── invlists │ ├── BlockInvertedLists.cpp │ ├── BlockInvertedLists.h │ ├── DirectMap.cpp │ ├── DirectMap.h │ ├── InvertedLists.cpp │ ├── InvertedLists.h │ ├── InvertedListsIOHook.cpp │ ├── InvertedListsIOHook.h │ ├── OnDiskInvertedLists.cpp │ └── OnDiskInvertedLists.h ├── python │ ├── CMakeLists.txt │ ├── __init__.py │ ├── array_conversions.py │ ├── class_wrappers.py │ ├── extra_wrappers.py │ ├── faiss_example_external_module.swig │ ├── gpu_wrappers.py │ ├── loader.py │ ├── python_callbacks.cpp │ ├── python_callbacks.h │ ├── setup.py │ └── swigfaiss.swig └── utils │ ├── AlignedTable.h │ ├── Heap.cpp │ ├── Heap.h │ ├── NeuralNet.cpp │ ├── NeuralNet.h │ ├── WorkerThread.cpp │ ├── WorkerThread.h │ ├── approx_topk │ ├── approx_topk.h │ ├── avx2-inl.h │ ├── generic.h │ └── mode.h │ ├── approx_topk_hamming │ └── approx_topk_hamming.h │ ├── bf16.h │ ├── distances.cpp │ ├── distances.h │ ├── distances_fused │ ├── avx512.cpp │ ├── avx512.h │ ├── distances_fused.cpp │ ├── distances_fused.h │ ├── simdlib_based.cpp │ └── simdlib_based.h │ ├── distances_simd.cpp │ ├── extra_distances-inl.h │ ├── extra_distances.cpp │ ├── extra_distances.h │ ├── fp16-arm.h │ ├── fp16-fp16c.h │ ├── fp16-inl.h │ ├── fp16.h │ ├── hamming-inl.h │ ├── hamming.cpp │ ├── hamming.h │ ├── hamming_distance │ ├── avx2-inl.h │ ├── avx512-inl.h │ ├── common.h │ ├── generic-inl.h │ ├── hamdis-inl.h │ └── neon-inl.h │ ├── ordered_key_value.h │ ├── partitioning.cpp │ ├── partitioning.h │ ├── prefetch.h │ ├── quantize_lut.cpp │ ├── quantize_lut.h │ ├── random.cpp │ ├── random.h │ ├── simdlib.h │ ├── simdlib_avx2.h │ ├── simdlib_avx512.h │ ├── simdlib_emulated.h │ ├── simdlib_neon.h │ ├── simdlib_ppc64.h │ ├── sorting.cpp │ ├── sorting.h │ ├── transpose │ ├── transpose-avx2-inl.h │ └── transpose-avx512-inl.h │ ├── utils.cpp │ └── utils.h ├── misc └── test_blas.cpp ├── perf_tests ├── CMakeLists.txt ├── bench_hnsw.py ├── bench_no_multithreading_rcq_search.cpp ├── bench_scalar_quantizer_accuracy.cpp ├── bench_scalar_quantizer_decode.cpp ├── bench_scalar_quantizer_distance.cpp ├── bench_scalar_quantizer_encode.cpp ├── utils.cpp └── utils.h ├── tests ├── CMakeLists.txt ├── common_faiss_tests.py ├── external_module_test.py ├── test_NSG_compressed_graph.cpp ├── test_RCQ_cropping.cpp ├── test_approx_topk.cpp ├── test_autotune.py ├── test_binary_factory.py ├── test_binary_flat.cpp ├── test_binary_hashindex.py ├── test_binary_io.py ├── test_build_blocks.py ├── test_callback.cpp ├── test_callback_py.py ├── test_clone.py ├── test_clustering.py ├── test_code_distance.cpp ├── test_common_ivf_empty_index.cpp ├── test_contrib.py ├── test_contrib_with_scipy.py ├── test_cppcontrib_sa_decode.cpp ├── test_cppcontrib_uintreader.cpp ├── test_dealloc_invlists.cpp ├── test_disable_pq_sdc_tables.cpp ├── test_distances_simd.cpp ├── test_documentation.py ├── test_extra_distances.py ├── test_factory.py ├── test_factory_tools.cpp ├── test_fast_scan.py ├── test_fast_scan_ivf.py ├── test_fastscan_perf.cpp ├── test_graph_based.py ├── test_hamming.cpp ├── test_heap.cpp ├── test_hnsw.cpp ├── test_index.py ├── test_index_accuracy.py ├── test_index_binary.py ├── test_index_binary_from_float.py ├── test_index_composite.py ├── test_io.py ├── test_ivf_index.cpp ├── test_ivflib.py ├── test_ivfpq_codec.cpp ├── test_ivfpq_indexing.cpp ├── test_local_search_quantizer.py ├── test_lowlevel_ivf.cpp ├── test_mem_leak.cpp ├── test_merge.cpp ├── test_merge_index.py ├── test_meta_index.py ├── test_mmap.cpp ├── test_omp_threads.cpp ├── test_omp_threads_py.py ├── test_ondisk_ivf.cpp ├── test_pairs_decoding.cpp ├── test_params_override.cpp ├── test_partition.py ├── test_partitioning.cpp ├── test_pq_encoding.cpp ├── test_product_quantizer.py ├── test_rabitq.py ├── test_referenced_objects.py ├── test_refine.py ├── test_residual_quantizer.py ├── test_rowwise_minmax.py ├── test_search_params.py ├── test_simdlib.cpp ├── test_sliding_ivf.cpp ├── test_standalone_codec.py ├── test_swig_wrapper.py ├── test_threaded_index.cpp ├── test_transfer_invlists.cpp ├── test_util.h ├── test_utils.cpp ├── test_zerocopy.cpp ├── torch_test_contrib.py └── torch_test_neural_net.py └── tutorial ├── cpp ├── 1-Flat.cpp ├── 2-IVFFlat.cpp ├── 3-IVFPQ.cpp ├── 4-GPU.cpp ├── 5-Multiple-GPUs.cpp ├── 6-HNSW.cpp ├── 7-PQFastScan.cpp ├── 8-PQFastScanRefine.cpp ├── 9-RefineComparison.cpp └── CMakeLists.txt └── python ├── 1-Flat.py ├── 2-IVFFlat.py ├── 3-IVFPQ.py ├── 4-GPU.py ├── 5-Multiple-GPUs.py ├── 7-PQFastScan.py ├── 8-PQFastScanRefine.py └── 9-RefineComparison.py /.dockerignore: -------------------------------------------------------------------------------- 1 | sift1M -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | # Summary 2 | 3 | 6 | 7 | # Platform 8 | 9 | 10 | 11 | OS: 12 | 13 | Faiss version: 14 | 15 | Installed from: 16 | 17 | Faiss compilation options: 18 | 19 | Running on: 20 | - [ ] CPU 21 | - [ ] GPU 22 | 23 | Interface: 24 | - [ ] C++ 25 | - [ ] Python 26 | 27 | # Reproduction instructions 28 | 29 | 31 | 32 | 34 | -------------------------------------------------------------------------------- /.github/workflows/autoclose.yml: -------------------------------------------------------------------------------- 1 | name: Close Inactive Issues 2 | on: 3 | schedule: 4 | - cron: "30 1 * * *" 5 | 6 | jobs: 7 | close-issues: 8 | runs-on: ubuntu-latest 9 | permissions: 10 | issues: write 11 | pull-requests: write 12 | steps: 13 | - uses: actions/stale@v5 14 | with: 15 | only-labels: autoclose 16 | days-before-issue-stale: 7 17 | days-before-issue-close: 7 18 | stale-issue-label: "stale" 19 | stale-issue-message: "This issue is stale because it has been open for 7 days with no activity." 20 | close-issue-message: "This issue was closed because it has been inactive for 7 days since being marked as stale." 21 | days-before-pr-stale: -1 22 | days-before-pr-close: -1 23 | repo-token: ${{ secrets.GITHUB_TOKEN }} 24 | -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Build 2 | on: 3 | workflow_dispatch: 4 | pull_request: 5 | branches: 6 | - main 7 | push: 8 | tags: 9 | - 'v*' 10 | jobs: 11 | build-pull-request: 12 | uses: ./.github/workflows/build-pull-request.yml 13 | build-release: 14 | uses: ./.github/workflows/build-release.yml 15 | secrets: 16 | ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} 17 | if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') 18 | -------------------------------------------------------------------------------- /.github/workflows/publish-docs.yml: -------------------------------------------------------------------------------- 1 | name: Publish Docs 2 | on: 3 | page_build: 4 | branches: 5 | - gh-pages 6 | paths-ignore: 7 | - 'docs/**' 8 | workflow_run: 9 | workflows: [update-doxygen] 10 | types: 11 | - completed 12 | jobs: 13 | build_and_publish: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/checkout@v2 17 | - name: Set up Python 18 | uses: actions/setup-python@v2 19 | with: 20 | python-version: 3.8 21 | - name: Checkout gh-pages 22 | run: | 23 | git fetch origin gh-pages 24 | git checkout gh-pages 25 | - name: Install dependencies 26 | run: | 27 | python -m pip install --upgrade pip 28 | pip install -r requirements.txt 29 | - name: Generate html 30 | run: | 31 | make html 32 | git rm -rf docs 33 | mv _build/html docs 34 | touch docs/.nojekyll 35 | - name: Push changes 36 | run: | 37 | git config --global user.email "$GITHUB_ACTOR@users.noreply.github.com" 38 | git config --global user.name "$GITHUB_ACTOR" 39 | git add docs 40 | if [ -n "$(git status --porcelain)" ] 41 | then 42 | git commit docs -m "Sphinx rebuild ($(git rev-parse --short gh-pages))." 43 | git push origin gh-pages 44 | fi 45 | -------------------------------------------------------------------------------- /.github/workflows/retry_build.yml: -------------------------------------------------------------------------------- 1 | name: Retry Build 2 | on: 3 | workflow_dispatch: 4 | inputs: 5 | run_id: 6 | required: true 7 | jobs: 8 | rerun-on-failure: 9 | permissions: write-all 10 | runs-on: ubuntu-latest 11 | steps: 12 | - name: rerun ${{ inputs.run_id }} 13 | env: 14 | GH_REPO: ${{ github.repository }} 15 | GH_TOKEN: ${{ github.token }} 16 | GH_DEBUG: api 17 | run: | 18 | # status can be one of "queued", "in_progress", "completed", "waiting", "requested", "pending" 19 | # https://docs.github.com/en/rest/checks/runs 20 | # while not completed, sleep for 10 minutes 21 | while gh run view ${{ inputs.run_id }} --json status | grep -v completed 22 | do 23 | echo Workflow in progress - sleeping for 10 minutes then checking again 24 | sleep 10m 25 | done 26 | 27 | # Only retry if there are failed jobs 28 | if gh run view ${{ inputs.run_id }} --exit-status; then 29 | echo Workflow succeeded - no retry necessary. 30 | else 31 | echo Workflow failed - initiating retry. 32 | gh run rerun ${{ inputs.run_id }} --failed 33 | fi 34 | -------------------------------------------------------------------------------- /.github/workflows/update-doxygen.yml: -------------------------------------------------------------------------------- 1 | name: Update Doxygen 2 | on: 3 | push: 4 | branches: 5 | - main 6 | paths: 7 | - 'faiss/**' 8 | jobs: 9 | doxygen: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v2 13 | - name: Set up Python 14 | uses: actions/setup-python@v2 15 | with: 16 | python-version: 3.8 17 | - name: Install dependencies 18 | run: | 19 | sudo apt-get install -y doxygen 20 | python -m pip install --upgrade pip 21 | pip install breathe 22 | - name: Generate doxygen xml 23 | run: doxygen 24 | - name: Push changes 25 | run: | 26 | git config --global user.email "$GITHUB_ACTOR@users.noreply.github.com" 27 | git config --global user.name "$GITHUB_ACTOR" 28 | mkdir ./tmp 29 | mv xml ./tmp/xml 30 | git fetch origin gh-pages 31 | git checkout gh-pages 32 | git rm -rf xml cpp_api 33 | mv ./tmp/xml ./xml 34 | breathe-apidoc -o cpp_api xml 35 | git add xml cpp_api 36 | if [ -n "$(git status --porcelain)" ] 37 | then 38 | git commit -m "Update API docs ($(git rev-parse --short main))." 39 | git push origin gh-pages 40 | fi 41 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | *.swo 3 | *.o 4 | *.a 5 | *.dSYM 6 | *.so 7 | *.dylib 8 | *.pyc 9 | *~ 10 | /build/ 11 | /config.* 12 | /aclocal.m4 13 | /autom4te.cache/ 14 | /makefile.inc 15 | /bin/ 16 | /c_api/bin/ 17 | /c_api/gpu/bin/ 18 | /tests/test 19 | /tests/gtest/ 20 | faiss/python/swigfaiss_avx2.swig 21 | faiss/python/swigfaiss_avx512.swig 22 | faiss/python/swigfaiss_avx512_spr.swig 23 | faiss/python/swigfaiss_sve.swig 24 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | Facebook has adopted a Code of Conduct that we expect project participants to adhere to. Please [read the full text](https://code.fb.com/codeofconduct) so that you can understand what actions will and will not be tolerated. -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to Faiss 2 | 3 | We want to make contributing to this project as easy and transparent as 4 | possible. 5 | 6 | ## Our Development Process 7 | 8 | We mainly develop Faiss within Facebook. Sometimes, we will sync the 9 | github version of Faiss with the internal state. 10 | 11 | ## Pull Requests 12 | 13 | We welcome pull requests that add significant value to Faiss. If you plan to do 14 | a major development and contribute it back to Faiss, please contact us first before 15 | putting too much effort into it. 16 | 17 | 1. Fork the repo and create your branch from `main`. 18 | 2. If you've added code that should be tested, add tests. 19 | 3. If you've changed APIs, update the documentation. 20 | 4. Ensure the test suite passes. 21 | 5. Make sure your code lints. 22 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 23 | 24 | There is a Facebook internal test suite for Faiss, and we need to run 25 | all changes to Faiss through it. 26 | 27 | ## Contributor License Agreement ("CLA") 28 | 29 | In order to accept your pull request, we need you to submit a CLA. You only need 30 | to do this once to work on any of Facebook's open source projects. 31 | 32 | Complete your CLA here: 33 | 34 | ## Issues 35 | 36 | We use GitHub issues to track public bugs. Please ensure your description is 37 | clear and has sufficient instructions to be able to reproduce the issue. 38 | 39 | Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe 40 | disclosure of security bugs. In those cases, please go through the process 41 | outlined on that page and do not file a public issue. 42 | 43 | ## Coding Style 44 | 45 | * 4 spaces for indentation in C++ (no tabs) 46 | * 80 character line length (both for C++ and Python) 47 | * C++ language level: C++17 48 | 49 | ## License 50 | 51 | By contributing to Faiss, you agree that your contributions will be licensed 52 | under the LICENSE file in the root directory of this source tree. 53 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) Facebook, Inc. and its affiliates. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /benchs/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | 7 | 8 | add_executable(bench_ivf_selector EXCLUDE_FROM_ALL bench_ivf_selector.cpp) 9 | target_link_libraries(bench_ivf_selector PRIVATE faiss) 10 | 11 | -------------------------------------------------------------------------------- /benchs/bench_all_ivf/README.md: -------------------------------------------------------------------------------- 1 | # Benchmark of IVF variants 2 | 3 | This is a benchmark of IVF index variants, looking at compression vs. speed vs. accuracy. 4 | The results are in [this wiki chapter](https://github.com/facebookresearch/faiss/wiki/Indexing-1G-vectors) 5 | 6 | 7 | The code is organized as: 8 | 9 | - `datasets.py`: code to access the datafiles, compute the ground-truth and report accuracies 10 | 11 | - `bench_all_ivf.py`: evaluate one type of inverted file 12 | 13 | - `run_on_cluster_generic.bash`: call `bench_all_ivf.py` for all tested types of indices. 14 | Since the number of experiments is quite large the script is structured so that the benchmark can be run on a cluster. 15 | 16 | - `parse_bench_all_ivf.py`: make nice tradeoff plots from all the results. 17 | 18 | The code depends on Faiss and can use 1 to 8 GPUs to do the k-means clustering for large vocabularies. 19 | 20 | It was run in October 2018 for the results in the wiki. 21 | -------------------------------------------------------------------------------- /benchs/bench_all_ivf/make_groundtruth.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import logging 7 | 8 | # https://stackoverflow.com/questions/7016056/python-logging-not-outputting-anything 9 | logging.basicConfig() 10 | logger = logging.getLogger('faiss.contrib.exhaustive_search') 11 | logger.setLevel(logging.INFO) 12 | 13 | from faiss.contrib import datasets 14 | from faiss.contrib.exhaustive_search import knn_ground_truth 15 | from faiss.contrib import vecs_io 16 | 17 | ds = datasets.DatasetDeep1B(nb=int(1e9)) 18 | 19 | print("computing GT matches for", ds) 20 | 21 | D, I = knn_ground_truth( 22 | ds.get_queries(), 23 | ds.database_iterator(bs=65536), 24 | k=100 25 | ) 26 | 27 | vecs_io.ivecs_write("/tmp/tt.ivecs", I) 28 | -------------------------------------------------------------------------------- /benchs/bench_fw/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/faiss/f94f83531dfe9a538dcd51012aeecd31ff13affb/benchs/bench_fw/__init__.py -------------------------------------------------------------------------------- /benchs/bench_fw_optimize.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import argparse 7 | import logging 8 | import os 9 | 10 | from faiss.benchs.bench_fw.benchmark_io import BenchmarkIO 11 | from faiss.benchs.bench_fw.descriptors import DatasetDescriptor 12 | from faiss.benchs.bench_fw.optimize import Optimizer 13 | 14 | logging.basicConfig(level=logging.INFO) 15 | 16 | 17 | def bigann(bio): 18 | optimizer = Optimizer( 19 | distance_metric="L2", 20 | num_threads=32, 21 | run_local=False, 22 | ) 23 | optimizer.set_io(bio) 24 | query_vectors = DatasetDescriptor(namespace="std_q", tablename="bigann1M") 25 | xt = bio.get_dataset(query_vectors) 26 | optimizer.optimize( 27 | d=xt.shape[1], 28 | training_vectors=DatasetDescriptor( 29 | namespace="std_t", 30 | tablename="bigann1M", 31 | num_vectors=2_000_000, 32 | ), 33 | database_vectors_list=[ 34 | DatasetDescriptor( 35 | namespace="std_d", 36 | tablename="bigann1M", 37 | ), 38 | DatasetDescriptor(namespace="std_d", tablename="bigann10M"), 39 | ], 40 | query_vectors=query_vectors, 41 | min_accuracy=0.85, 42 | ) 43 | 44 | 45 | if __name__ == "__main__": 46 | parser = argparse.ArgumentParser() 47 | parser.add_argument("experiment") 48 | parser.add_argument("path") 49 | args = parser.parse_args() 50 | assert os.path.exists(args.path) 51 | path = os.path.join(args.path, args.experiment) 52 | if not os.path.exists(path): 53 | os.mkdir(path) 54 | bio = BenchmarkIO( 55 | path=path, 56 | ) 57 | if args.experiment == "bigann": 58 | bigann(bio) 59 | -------------------------------------------------------------------------------- /benchs/bench_hamming_knn.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import time 7 | import numpy as np 8 | import faiss 9 | 10 | if __name__ == "__main__": 11 | faiss.omp_set_num_threads(1) 12 | 13 | for d in 4, 8, 16, 13: 14 | nq = 10000 15 | nb = 30000 16 | print('Bits per vector = 8 *', d) 17 | xq = faiss.randint((nq, d // 4), seed=1234, vmax=256**4).view('uint8') 18 | xb = faiss.randint((nb, d // 4), seed=1234, vmax=256**4).view('uint8') 19 | for variant in "hc", "mc": 20 | print(f"{variant=:}", end="\t") 21 | for k in 1, 4, 16, 64, 256: 22 | times = [] 23 | for _run in range(5): 24 | t0 = time.time() 25 | D, I = faiss.knn_hamming(xq, xb, k, variant=variant) 26 | t1 = time.time() 27 | times.append(t1 - t0) 28 | print(f'| {k=:} t={np.mean(times):.3f} s ± {np.std(times):.3f} ', flush=True, end="") 29 | print() 30 | -------------------------------------------------------------------------------- /benchs/bench_index_pq.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from __future__ import print_function 7 | import faiss 8 | from datasets import load_sift1M, evaluate 9 | 10 | xb, xq, xt, gt = load_sift1M() 11 | nq, d = xq.shape 12 | 13 | k = 32 14 | 15 | for nbits in 4, 6, 8, 10, 12: 16 | index = faiss.IndexPQ(d, 8, nbits) 17 | index.train(xt) 18 | index.add(xb) 19 | 20 | t, r = evaluate(index, xq, gt, k) 21 | print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1])) 22 | del index 23 | -------------------------------------------------------------------------------- /benchs/bench_pairwise_distances.py: -------------------------------------------------------------------------------- 1 | #! /usr/bin/env python3 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | """small test script to benchmark the SIMD implementation of the 8 | distance computations for the additional metrics. Call eg. with L1 to 9 | get L1 distance computations. 10 | """ 11 | 12 | import faiss 13 | 14 | import sys 15 | import time 16 | 17 | d = 64 18 | nq = 4096 19 | nb = 16384 20 | 21 | print("sample") 22 | 23 | xq = faiss.randn((nq, d), 123) 24 | xb = faiss.randn((nb, d), 123) 25 | 26 | mt_name = "L2" if len(sys.argv) < 2 else sys.argv[1] 27 | 28 | mt = getattr(faiss, "METRIC_" + mt_name) 29 | 30 | print("distances") 31 | t0 = time.time() 32 | dis = faiss.pairwise_distances(xq, xb, mt) 33 | t1 = time.time() 34 | 35 | print("nq=%d nb=%d d=%d %s: %.3f s" % (nq, nb, d, mt_name, t1 - t0)) 36 | -------------------------------------------------------------------------------- /benchs/bench_polysemous_sift1m.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | from __future__ import print_function 8 | import time 9 | import numpy as np 10 | 11 | import faiss 12 | from datasets import load_sift1M, evaluate 13 | 14 | 15 | print("load data") 16 | xb, xq, xt, gt = load_sift1M() 17 | nq, d = xq.shape 18 | 19 | # index with 16 subquantizers, 8 bit each 20 | index = faiss.IndexPQ(d, 16, 8) 21 | index.do_polysemous_training = True 22 | index.verbose = True 23 | 24 | print("train") 25 | 26 | index.train(xt) 27 | 28 | print("add vectors to index") 29 | 30 | index.add(xb) 31 | 32 | nt = 1 33 | faiss.omp_set_num_threads(1) 34 | 35 | 36 | print("PQ baseline", end=' ') 37 | index.search_type = faiss.IndexPQ.ST_PQ 38 | t, r = evaluate(index, xq, gt, 1) 39 | print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1])) 40 | 41 | for ht in 64, 62, 58, 54, 50, 46, 42, 38, 34, 30: 42 | print("Polysemous", ht, end=' ') 43 | index.search_type = faiss.IndexPQ.ST_polysemous 44 | index.polysemous_ht = ht 45 | t, r = evaluate(index, xq, gt, 1) 46 | print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1])) 47 | -------------------------------------------------------------------------------- /benchs/datasets.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from __future__ import print_function 7 | import sys 8 | import time 9 | import numpy as np 10 | 11 | 12 | def ivecs_read(fname): 13 | a = np.fromfile(fname, dtype='int32') 14 | d = a[0] 15 | return a.reshape(-1, d + 1)[:, 1:].copy() 16 | 17 | 18 | def fvecs_read(fname): 19 | return ivecs_read(fname).view('float32') 20 | 21 | 22 | def load_sift1M(): 23 | print("Loading sift1M...", end='', file=sys.stderr) 24 | xt = fvecs_read("sift1M/sift_learn.fvecs") 25 | xb = fvecs_read("sift1M/sift_base.fvecs") 26 | xq = fvecs_read("sift1M/sift_query.fvecs") 27 | gt = ivecs_read("sift1M/sift_groundtruth.ivecs") 28 | print("done", file=sys.stderr) 29 | 30 | return xb, xq, xt, gt 31 | 32 | 33 | def evaluate(index, xq, gt, k): 34 | nq = xq.shape[0] 35 | t0 = time.time() 36 | D, I = index.search(xq, k) # noqa: E741 37 | t1 = time.time() 38 | 39 | recalls = {} 40 | i = 1 41 | while i <= k: 42 | recalls[i] = (I[:, :i] == gt[:, :1]).sum() / float(nq) 43 | i *= 10 44 | 45 | return (t1 - t0) * 1000.0 / nq, recalls 46 | -------------------------------------------------------------------------------- /benchs/distributed_ondisk/make_trained_index.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | import faiss 8 | 9 | deep1bdir = "/datasets01_101/simsearch/041218/deep1b/" 10 | workdir = "/checkpoint/matthijs/ondisk_distributed/" 11 | 12 | 13 | print('Load centroids') 14 | centroids = np.load(workdir + '1M_centroids.npy') 15 | ncent, d = centroids.shape 16 | 17 | 18 | print('apply random rotation') 19 | rrot = faiss.RandomRotationMatrix(d, d) 20 | rrot.init(1234) 21 | centroids = rrot.apply_py(centroids) 22 | 23 | print('make HNSW index as quantizer') 24 | quantizer = faiss.IndexHNSWFlat(d, 32) 25 | quantizer.hnsw.efSearch = 1024 26 | quantizer.hnsw.efConstruction = 200 27 | quantizer.add(centroids) 28 | 29 | print('build index') 30 | index = faiss.IndexPreTransform( 31 | rrot, 32 | faiss.IndexIVFScalarQuantizer( 33 | quantizer, d, ncent, faiss.ScalarQuantizer.QT_6bit 34 | ) 35 | ) 36 | 37 | def ivecs_mmap(fname): 38 | a = np.memmap(fname, dtype='int32', mode='r') 39 | d = a[0] 40 | return a.reshape(-1, d + 1)[:, 1:] 41 | 42 | def fvecs_mmap(fname): 43 | return ivecs_mmap(fname).view('float32') 44 | 45 | 46 | print('finish training index') 47 | xt = fvecs_mmap(deep1bdir + 'learn.fvecs') 48 | xt = np.ascontiguousarray(xt[:256 * 1000], dtype='float32') 49 | index.train(xt) 50 | 51 | print('write output') 52 | faiss.write_index(index, workdir + 'trained.faissindex') 53 | -------------------------------------------------------------------------------- /benchs/link_and_code/README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | README for the link & code implementation 4 | ========================================= 5 | 6 | What is this? 7 | ------------- 8 | 9 | Link & code is an indexing method that combines HNSW indexing with 10 | compression and exploits the neighborhood structure of the similarity 11 | graph to improve the reconstruction. It is described in 12 | 13 | ``` 14 | @inproceedings{link_and_code, 15 | author = {Matthijs Douze and Alexandre Sablayrolles and Herv\'e J\'egou}, 16 | title = {Link and code: Fast indexing with graphs and compact regression codes}, 17 | booktitle = {CVPR}, 18 | year = {2018} 19 | } 20 | ``` 21 | 22 | ArXiV [here](https://arxiv.org/abs/1804.09996) 23 | 24 | The necessary code for this paper was removed from Faiss in version 1.8.0. 25 | For a functioning verinsion, use Faiss 1.7.4. 26 | -------------------------------------------------------------------------------- /c_api/IndexLSH_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "IndexLSH_c.h" 11 | #include 12 | #include "macros_impl.h" 13 | 14 | using faiss::Index; 15 | using faiss::IndexLSH; 16 | 17 | DEFINE_DESTRUCTOR(IndexLSH) 18 | DEFINE_INDEX_DOWNCAST(IndexLSH) 19 | 20 | DEFINE_GETTER(IndexLSH, int, nbits) 21 | DEFINE_GETTER(IndexLSH, int, code_size) 22 | DEFINE_GETTER_PERMISSIVE(IndexLSH, int, rotate_data) 23 | DEFINE_GETTER_PERMISSIVE(IndexLSH, int, train_thresholds) 24 | 25 | int faiss_IndexLSH_new(FaissIndexLSH** p_index, idx_t d, int nbits) { 26 | try { 27 | *p_index = reinterpret_cast(new IndexLSH(d, nbits)); 28 | } 29 | CATCH_AND_HANDLE 30 | } 31 | 32 | int faiss_IndexLSH_new_with_options( 33 | FaissIndexLSH** p_index, 34 | idx_t d, 35 | int nbits, 36 | int rotate_data, 37 | int train_thresholds) { 38 | try { 39 | *p_index = reinterpret_cast(new IndexLSH( 40 | d, 41 | nbits, 42 | static_cast(rotate_data), 43 | static_cast(train_thresholds))); 44 | } 45 | CATCH_AND_HANDLE 46 | } 47 | -------------------------------------------------------------------------------- /c_api/IndexLSH_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef INDEX_LSH_C_H 11 | #define INDEX_LSH_C_H 12 | 13 | #include "Clustering_c.h" 14 | #include "Index_c.h" 15 | #include "faiss_c.h" 16 | 17 | #ifdef __cplusplus 18 | extern "C" { 19 | #endif 20 | 21 | /** The sign of each vector component is put in a binary signature */ 22 | FAISS_DECLARE_CLASS_INHERITED(IndexLSH, Index) 23 | FAISS_DECLARE_DESTRUCTOR(IndexLSH) 24 | FAISS_DECLARE_INDEX_DOWNCAST(IndexLSH) 25 | 26 | FAISS_DECLARE_GETTER(IndexLSH, int, nbits) 27 | FAISS_DECLARE_GETTER(IndexLSH, int, code_size) 28 | FAISS_DECLARE_GETTER(IndexLSH, int, rotate_data) 29 | FAISS_DECLARE_GETTER(IndexLSH, int, train_thresholds) 30 | 31 | int faiss_IndexLSH_new(FaissIndexLSH** p_index, idx_t d, int nbits); 32 | 33 | int faiss_IndexLSH_new_with_options( 34 | FaissIndexLSH** p_index, 35 | idx_t d, 36 | int nbits, 37 | int rotate_data, 38 | int train_thresholds); 39 | 40 | #ifdef __cplusplus 41 | } 42 | #endif 43 | 44 | #endif 45 | -------------------------------------------------------------------------------- /c_api/IndexPreTransform_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_INDEX_PRETRANSFORM_C_H 11 | #define FAISS_INDEX_PRETRANSFORM_C_H 12 | 13 | #include "Index_c.h" 14 | #include "VectorTransform_c.h" 15 | #include "faiss_c.h" 16 | 17 | #ifdef __cplusplus 18 | extern "C" { 19 | #endif 20 | 21 | /** Index that applies a LinearTransform transform on vectors before 22 | * handing them over to a sub-index */ 23 | FAISS_DECLARE_CLASS_INHERITED(IndexPreTransform, Index) 24 | FAISS_DECLARE_DESTRUCTOR(IndexPreTransform) 25 | FAISS_DECLARE_INDEX_DOWNCAST(IndexPreTransform) 26 | 27 | FAISS_DECLARE_GETTER(IndexPreTransform, FaissIndex*, index) 28 | FAISS_DECLARE_GETTER_SETTER(IndexPreTransform, int, own_fields) 29 | 30 | int faiss_IndexPreTransform_new(FaissIndexPreTransform** p_index); 31 | 32 | int faiss_IndexPreTransform_new_with( 33 | FaissIndexPreTransform** p_index, 34 | FaissIndex* index); 35 | 36 | int faiss_IndexPreTransform_new_with_transform( 37 | FaissIndexPreTransform** p_index, 38 | FaissVectorTransform* ltrans, 39 | FaissIndex* index); 40 | 41 | int faiss_IndexPreTransform_prepend_transform( 42 | FaissIndexPreTransform* index, 43 | FaissVectorTransform* ltrans); 44 | 45 | #ifdef __cplusplus 46 | } 47 | #endif 48 | 49 | #endif 50 | -------------------------------------------------------------------------------- /c_api/IndexReplicas_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include "IndexReplicas_c.h" 9 | #include 10 | #include "macros_impl.h" 11 | 12 | using faiss::Index; 13 | using faiss::IndexReplicas; 14 | 15 | DEFINE_DESTRUCTOR(IndexReplicas) 16 | 17 | DEFINE_GETTER(IndexReplicas, int, own_indices) 18 | DEFINE_SETTER(IndexReplicas, int, own_indices) 19 | 20 | int faiss_IndexReplicas_new(FaissIndexReplicas** p_index, idx_t d) { 21 | try { 22 | auto out = new IndexReplicas(d); 23 | *p_index = reinterpret_cast(out); 24 | } 25 | CATCH_AND_HANDLE 26 | } 27 | 28 | int faiss_IndexReplicas_new_with_options( 29 | FaissIndexReplicas** p_index, 30 | idx_t d, 31 | int threaded) { 32 | try { 33 | auto out = new IndexReplicas(d, static_cast(threaded)); 34 | *p_index = reinterpret_cast(out); 35 | } 36 | CATCH_AND_HANDLE 37 | } 38 | 39 | int faiss_IndexReplicas_add_replica( 40 | FaissIndexReplicas* index, 41 | FaissIndex* replica) { 42 | try { 43 | reinterpret_cast(index)->add_replica( 44 | reinterpret_cast(replica)); 45 | } 46 | CATCH_AND_HANDLE 47 | } 48 | 49 | int faiss_IndexReplicas_remove_replica( 50 | FaissIndexReplicas* index, 51 | FaissIndex* replica) { 52 | try { 53 | reinterpret_cast(index)->remove_replica( 54 | reinterpret_cast(replica)); 55 | } 56 | CATCH_AND_HANDLE 57 | } 58 | 59 | FaissIndex* faiss_IndexReplicas_at(FaissIndexReplicas* index, int i) { 60 | auto replica = reinterpret_cast(index)->at(i); 61 | return reinterpret_cast(replica); 62 | } 63 | -------------------------------------------------------------------------------- /c_api/IndexReplicas_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef INDEXREPLICAS_C_H 11 | #define INDEXREPLICAS_C_H 12 | 13 | #include "Index_c.h" 14 | #include "faiss_c.h" 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | /** Index that concatenates the results from several sub-indexes 21 | */ 22 | FAISS_DECLARE_CLASS_INHERITED(IndexReplicas, Index) 23 | FAISS_DECLARE_DESTRUCTOR(IndexReplicas) 24 | 25 | FAISS_DECLARE_GETTER_SETTER(IndexReplicas, int, own_fields) 26 | 27 | int faiss_IndexReplicas_new(FaissIndexReplicas** p_index, idx_t d); 28 | 29 | int faiss_IndexReplicas_new_with_options( 30 | FaissIndexReplicas** p_index, 31 | idx_t d, 32 | int threaded); 33 | 34 | int faiss_IndexReplicas_add_replica( 35 | FaissIndexReplicas* index, 36 | FaissIndex* replica); 37 | 38 | int faiss_IndexReplicas_remove_replica( 39 | FaissIndexReplicas* index, 40 | FaissIndex* replica); 41 | 42 | FaissIndex* faiss_IndexReplicas_at(FaissIndexReplicas* index, int i); 43 | 44 | #ifdef __cplusplus 45 | } 46 | #endif 47 | #endif 48 | -------------------------------------------------------------------------------- /c_api/IndexShards_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef INDEXSHARDS_C_H 11 | #define INDEXSHARDS_C_H 12 | 13 | #include "Index_c.h" 14 | #include "faiss_c.h" 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | /** Index that concatenates the results from several sub-indexes 21 | */ 22 | FAISS_DECLARE_CLASS_INHERITED(IndexShards, Index) 23 | FAISS_DECLARE_DESTRUCTOR(IndexShards) 24 | 25 | FAISS_DECLARE_GETTER_SETTER(IndexShards, int, own_fields) 26 | FAISS_DECLARE_GETTER_SETTER(IndexShards, int, successive_ids) 27 | 28 | int faiss_IndexShards_new(FaissIndexShards** p_index, idx_t d); 29 | 30 | int faiss_IndexShards_new_with_options( 31 | FaissIndexShards** p_index, 32 | idx_t d, 33 | int threaded, 34 | int successive_ids); 35 | 36 | int faiss_IndexShards_add_shard(FaissIndexShards* index, FaissIndex* shard); 37 | 38 | int faiss_IndexShards_remove_shard(FaissIndexShards* index, FaissIndex* shard); 39 | 40 | FaissIndex* faiss_IndexShards_at(FaissIndexShards* index, int i); 41 | 42 | #ifdef __cplusplus 43 | } 44 | #endif 45 | #endif 46 | -------------------------------------------------------------------------------- /c_api/clone_index_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | // I/O code for indexes 10 | 11 | #include "clone_index_c.h" 12 | #include 13 | #include "macros_impl.h" 14 | 15 | using faiss::Index; 16 | using faiss::IndexBinary; 17 | 18 | int faiss_clone_index(const FaissIndex* idx, FaissIndex** p_out) { 19 | try { 20 | auto out = faiss::clone_index(reinterpret_cast(idx)); 21 | *p_out = reinterpret_cast(out); 22 | } 23 | CATCH_AND_HANDLE 24 | } 25 | 26 | int faiss_clone_index_binary( 27 | const FaissIndexBinary* idx, 28 | FaissIndexBinary** p_out) { 29 | try { 30 | auto out = faiss::clone_binary_index( 31 | reinterpret_cast(idx)); 32 | *p_out = reinterpret_cast(out); 33 | } 34 | CATCH_AND_HANDLE 35 | } 36 | -------------------------------------------------------------------------------- /c_api/clone_index_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | // I/O code for indexes 10 | 11 | #ifndef FAISS_CLONE_INDEX_C_H 12 | #define FAISS_CLONE_INDEX_C_H 13 | 14 | #include 15 | #include "IndexBinary_c.h" 16 | #include "Index_c.h" 17 | #include "faiss_c.h" 18 | 19 | #ifdef __cplusplus 20 | extern "C" { 21 | #endif 22 | 23 | /* cloning functions */ 24 | 25 | /** Clone an index. This is equivalent to `faiss::clone_index` */ 26 | int faiss_clone_index(const FaissIndex*, FaissIndex** p_out); 27 | 28 | /** Clone a binary index. This is equivalent to `faiss::clone_index_binary` */ 29 | int faiss_clone_index_binary(const FaissIndexBinary*, FaissIndexBinary** p_out); 30 | 31 | #ifdef __cplusplus 32 | } 33 | #endif 34 | #endif 35 | -------------------------------------------------------------------------------- /c_api/error_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_ERROR_C_H 11 | #define FAISS_ERROR_C_H 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | /// An error code which depends on the exception thrown from the previous 18 | /// operation. See `faiss_get_last_error` to retrieve the error message. 19 | typedef enum FaissErrorCode { 20 | /// No error 21 | OK = 0, 22 | /// Any exception other than Faiss or standard C++ library exceptions 23 | UNKNOWN_EXCEPT = -1, 24 | /// Faiss library exception 25 | FAISS_EXCEPT = -2, 26 | /// Standard C++ library exception 27 | STD_EXCEPT = -4 28 | } FaissErrorCode; 29 | 30 | /** 31 | * Get the error message of the last failed operation performed by Faiss. 32 | * The given pointer is only invalid until another Faiss function is 33 | * called. 34 | */ 35 | const char* faiss_get_last_error(); 36 | 37 | #ifdef __cplusplus 38 | } 39 | #endif 40 | 41 | #endif 42 | -------------------------------------------------------------------------------- /c_api/error_impl.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "error_impl.h" 11 | #include 12 | #include 13 | #include "error_c.h" 14 | 15 | thread_local std::exception_ptr faiss_last_exception; 16 | 17 | const char* faiss_get_last_error() { 18 | if (faiss_last_exception) { 19 | try { 20 | std::rethrow_exception(faiss_last_exception); 21 | } catch (std::exception& e) { 22 | return e.what(); 23 | } 24 | } 25 | return nullptr; 26 | } 27 | -------------------------------------------------------------------------------- /c_api/error_impl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include 11 | 12 | /** global variable for holding the last exception thrown by 13 | * calls to Faiss functions through the C API 14 | */ 15 | extern thread_local std::exception_ptr faiss_last_exception; 16 | -------------------------------------------------------------------------------- /c_api/gpu/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | target_sources(faiss_c PRIVATE 7 | DeviceUtils_c.cpp 8 | GpuAutoTune_c.cpp 9 | GpuClonerOptions_c.cpp 10 | GpuIndex_c.cpp 11 | GpuResources_c.cpp 12 | StandardGpuResources_c.cpp 13 | ) 14 | 15 | file(GLOB FAISS_C_API_GPU_HEADERS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.h") 16 | faiss_install_headers("${FAISS_C_API_GPU_HEADERS}" c_api/gpu) 17 | 18 | if (FAISS_ENABLE_ROCM) 19 | target_link_libraries(faiss_c PUBLIC hip::host roc::hipblas) 20 | target_link_libraries(faiss_c_avx2 PUBLIC hip::host roc::hipblas) 21 | target_link_libraries(faiss_c_avx512 PUBLIC hip::host roc::hipblas) 22 | target_link_libraries(faiss_c_avx512_spr PUBLIC hip::host roc::hipblas) 23 | target_link_libraries(faiss_c_sve PUBLIC hip::host roc::hipblas) 24 | else() 25 | find_package(CUDAToolkit REQUIRED) 26 | target_link_libraries(faiss_c PUBLIC CUDA::cudart CUDA::cublas $<$:cuvs::cuvs>) 27 | target_link_libraries(faiss_c_avx2 PUBLIC CUDA::cudart CUDA::cublas $<$:cuvs::cuvs>) 28 | target_link_libraries(faiss_c_avx512 PUBLIC CUDA::cudart CUDA::cublas $<$:cuvs::cuvs>) 29 | target_link_libraries(faiss_c_avx512_spr PUBLIC CUDA::cudart CUDA::cublas $<$:cuvs::cuvs>) 30 | target_link_libraries(faiss_c_sve PUBLIC CUDA::cudart CUDA::cublas $<$:cuvs::cuvs>) 31 | endif() 32 | 33 | add_executable(example_gpu_c EXCLUDE_FROM_ALL example_gpu_c.c) 34 | target_link_libraries(example_gpu_c PRIVATE faiss_c) 35 | -------------------------------------------------------------------------------- /c_api/gpu/DeviceUtils_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "DeviceUtils_c.h" 11 | #include 12 | #include "macros_impl.h" 13 | 14 | /// Returns the number of available GPU devices 15 | int faiss_get_num_gpus(int* p_output) { 16 | try { 17 | int output = faiss::gpu::getNumDevices(); 18 | *p_output = output; 19 | } 20 | CATCH_AND_HANDLE 21 | } 22 | 23 | /// Starts the CUDA profiler (exposed via SWIG) 24 | int faiss_gpu_profiler_start() { 25 | try { 26 | faiss::gpu::profilerStart(); 27 | } 28 | CATCH_AND_HANDLE 29 | } 30 | 31 | /// Stops the CUDA profiler (exposed via SWIG) 32 | int faiss_gpu_profiler_stop() { 33 | try { 34 | faiss::gpu::profilerStop(); 35 | } 36 | CATCH_AND_HANDLE 37 | } 38 | 39 | /// Synchronizes the CPU against all devices (equivalent to 40 | /// cudaDeviceSynchronize for each device) 41 | int faiss_gpu_sync_all_devices() { 42 | try { 43 | faiss::gpu::synchronizeAllDevices(); 44 | } 45 | CATCH_AND_HANDLE 46 | } 47 | -------------------------------------------------------------------------------- /c_api/gpu/DeviceUtils_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_DEVICE_UTILS_C_H 11 | #define FAISS_DEVICE_UTILS_C_H 12 | 13 | #include 14 | #include 15 | #include "../faiss_c.h" 16 | 17 | #ifdef __cplusplus 18 | extern "C" { 19 | #endif 20 | 21 | /// Returns the number of available GPU devices 22 | int faiss_get_num_gpus(int* p_output); 23 | 24 | /// Starts the CUDA profiler (exposed via SWIG) 25 | int faiss_gpu_profiler_start(); 26 | 27 | /// Stops the CUDA profiler (exposed via SWIG) 28 | int faiss_gpu_profiler_stop(); 29 | 30 | /// Synchronizes the CPU against all devices (equivalent to 31 | /// cudaDeviceSynchronize for each device) 32 | int faiss_gpu_sync_all_devices(); 33 | 34 | #ifdef __cplusplus 35 | } 36 | #endif 37 | #endif 38 | -------------------------------------------------------------------------------- /c_api/gpu/GpuIndex_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "GpuIndex_c.h" 11 | #include 12 | #include "macros_impl.h" 13 | 14 | using faiss::gpu::GpuIndexConfig; 15 | 16 | DEFINE_GETTER(GpuIndexConfig, int, device) 17 | -------------------------------------------------------------------------------- /c_api/gpu/GpuIndex_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_GPU_INDEX_C_H 11 | #define FAISS_GPU_INDEX_C_H 12 | 13 | #include "../faiss_c.h" 14 | 15 | #ifdef __cplusplus 16 | extern "C" { 17 | #endif 18 | 19 | FAISS_DECLARE_CLASS(GpuIndexConfig) 20 | 21 | FAISS_DECLARE_GETTER(GpuIndexConfig, int, device) 22 | 23 | FAISS_DECLARE_CLASS_INHERITED(GpuIndex, Index) 24 | 25 | #ifdef __cplusplus 26 | } 27 | #endif 28 | 29 | #endif 30 | -------------------------------------------------------------------------------- /c_api/gpu/GpuIndicesOptions_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_GPU_INDICES_OPTIONS_C_H 11 | #define FAISS_GPU_INDICES_OPTIONS_C_H 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | /// How user vector index data is stored on the GPU 18 | typedef enum FaissIndicesOptions { 19 | /// The user indices are only stored on the CPU; the GPU returns 20 | /// (inverted list, offset) to the CPU which is then translated to 21 | /// the real user index. 22 | INDICES_CPU = 0, 23 | /// The indices are not stored at all, on either the CPU or 24 | /// GPU. Only (inverted list, offset) is returned to the user as the 25 | /// index. 26 | INDICES_IVF = 1, 27 | /// Indices are stored as 32 bit integers on the GPU, but returned 28 | /// as 64 bit integers 29 | INDICES_32_BIT = 2, 30 | /// Indices are stored as 64 bit integers on the GPU 31 | INDICES_64_BIT = 3, 32 | } FaissIndicesOptions; 33 | 34 | #ifdef __cplusplus 35 | } 36 | #endif 37 | #endif 38 | -------------------------------------------------------------------------------- /c_api/gpu/macros_impl.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef GPU_MACROS_IMPL_H 11 | #define GPU_MACROS_IMPL_H 12 | #include "../macros_impl.h" 13 | 14 | #undef DEFINE_GETTER 15 | #define DEFINE_GETTER(clazz, ty, name) \ 16 | ty faiss_##clazz##_##name(const Faiss##clazz* obj) { \ 17 | return static_cast( \ 18 | reinterpret_cast(obj)->name); \ 19 | } 20 | 21 | #undef DEFINE_SETTER 22 | #define DEFINE_SETTER(clazz, ty, name) \ 23 | void faiss_##clazz##_set_##name(Faiss##clazz* obj, ty val) { \ 24 | reinterpret_cast(obj)->name = val; \ 25 | } 26 | 27 | #undef DEFINE_SETTER_STATIC 28 | #define DEFINE_SETTER_STATIC(clazz, ty_to, ty_from, name) \ 29 | void faiss_##clazz##_set_##name(Faiss##clazz* obj, ty_from val) { \ 30 | reinterpret_cast(obj)->name = \ 31 | static_cast(val); \ 32 | } 33 | 34 | #undef DEFINE_DESTRUCTOR 35 | #define DEFINE_DESTRUCTOR(clazz) \ 36 | void faiss_##clazz##_free(Faiss##clazz* obj) { \ 37 | delete reinterpret_cast(obj); \ 38 | } 39 | 40 | #endif 41 | -------------------------------------------------------------------------------- /c_api/impl/io_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_IO_C_H 11 | #define FAISS_IO_C_H 12 | 13 | #include 14 | #include "../faiss_c.h" 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | FAISS_DECLARE_CLASS(IOReader) 21 | FAISS_DECLARE_DESTRUCTOR(IOReader) 22 | 23 | FAISS_DECLARE_CLASS(IOWriter) 24 | FAISS_DECLARE_DESTRUCTOR(IOWriter) 25 | 26 | /******************************************************* 27 | * Custom reader + writer 28 | * 29 | * Reader and writer which wraps a function pointer, 30 | * primarily for FFI use. 31 | *******************************************************/ 32 | 33 | FAISS_DECLARE_CLASS(CustomIOReader) 34 | FAISS_DECLARE_DESTRUCTOR(CustomIOReader) 35 | 36 | int faiss_CustomIOReader_new( 37 | FaissCustomIOReader** p_out, 38 | size_t (*func_in)(void* ptr, size_t size, size_t nitems)); 39 | 40 | FAISS_DECLARE_CLASS(CustomIOWriter) 41 | FAISS_DECLARE_DESTRUCTOR(CustomIOWriter) 42 | 43 | int faiss_CustomIOWriter_new( 44 | FaissCustomIOWriter** p_out, 45 | size_t (*func_in)(const void* ptr, size_t size, size_t nitems)); 46 | 47 | #ifdef __cplusplus 48 | } 49 | #endif 50 | #endif 51 | -------------------------------------------------------------------------------- /c_api/index_factory_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "index_factory_c.h" 11 | #include 12 | #include 13 | #include "macros_impl.h" 14 | 15 | using faiss::Index; 16 | 17 | /** Build an index with the sequence of processing steps described in 18 | * the string. 19 | */ 20 | int faiss_index_factory( 21 | FaissIndex** p_index, 22 | int d, 23 | const char* description, 24 | FaissMetricType metric) { 25 | try { 26 | *p_index = reinterpret_cast(faiss::index_factory( 27 | d, description, static_cast(metric))); 28 | } 29 | CATCH_AND_HANDLE 30 | } 31 | 32 | /** Build an index with the sequence of processing steps described in 33 | * the string. 34 | */ 35 | int faiss_index_binary_factory( 36 | FaissIndexBinary** p_index, 37 | int d, 38 | const char* description) { 39 | try { 40 | *p_index = reinterpret_cast( 41 | faiss::index_binary_factory(d, description)); 42 | } 43 | CATCH_AND_HANDLE 44 | } 45 | -------------------------------------------------------------------------------- /c_api/index_factory_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_INDEX_FACTORY_C_H 11 | #define FAISS_INDEX_FACTORY_C_H 12 | 13 | #include "IndexBinary_c.h" 14 | #include "Index_c.h" 15 | #include "faiss_c.h" 16 | 17 | #ifdef __cplusplus 18 | extern "C" { 19 | #endif 20 | 21 | /** Build an index with the sequence of processing steps described in 22 | * the string. 23 | */ 24 | int faiss_index_factory( 25 | FaissIndex** p_index, 26 | int d, 27 | const char* description, 28 | FaissMetricType metric); 29 | 30 | /** Build a binary index with the sequence of processing steps described in 31 | * the string. 32 | */ 33 | int faiss_index_binary_factory( 34 | FaissIndexBinary** p_index, 35 | int d, 36 | const char* description); 37 | 38 | #ifdef __cplusplus 39 | } 40 | #endif 41 | 42 | #endif 43 | -------------------------------------------------------------------------------- /c_api/utils/utils_c.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #include "utils_c.h" 11 | #include 12 | 13 | const char* faiss_get_version() { 14 | return VERSION_STRING; 15 | } 16 | -------------------------------------------------------------------------------- /c_api/utils/utils_c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c -*- 9 | 10 | #ifndef FAISS_UTILS_C_H 11 | #define FAISS_UTILS_C_H 12 | 13 | #include 14 | #include 15 | 16 | #ifdef __cplusplus 17 | extern "C" { 18 | #endif 19 | 20 | const char* faiss_get_version(); 21 | 22 | #ifdef __cplusplus 23 | } 24 | #endif 25 | 26 | #endif 27 | -------------------------------------------------------------------------------- /cmake/faiss-config.cmake.in: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # All rights reserved. 3 | # 4 | # This source code is licensed under the BSD-style license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | include("${CMAKE_CURRENT_LIST_DIR}/faiss-targets.cmake") 8 | -------------------------------------------------------------------------------- /cmake/thirdparty/fetch_rapids.cmake: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | # ============================================================================= 6 | # Copyright (c) 2023-2024, NVIDIA CORPORATION. 7 | # 8 | # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except 9 | # in compliance with the License. You may obtain a copy of the License at 10 | # 11 | # http://www.apache.org/licenses/LICENSE-2.0 12 | # 13 | # Unless required by applicable law or agreed to in writing, software distributed under the License 14 | # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 15 | # or implied. See the License for the specific language governing permissions and limitations under 16 | # the License. 17 | # ============================================================================= 18 | set(RAPIDS_VERSION "25.04") 19 | 20 | if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) 21 | file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake 22 | ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) 23 | endif() 24 | include(${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) 25 | -------------------------------------------------------------------------------- /conda/conda_build_config.yaml: -------------------------------------------------------------------------------- 1 | python: 2 | - 3.9 3 | - 3.10 4 | - 3.11 5 | - 3.12 # [not aarch64] 6 | -------------------------------------------------------------------------------- /conda/faiss-gpu-cuvs/build-lib.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so 11 | cmake -B _build \ 12 | -DBUILD_SHARED_LIBS=ON \ 13 | -DFAISS_ENABLE_C_API=ON \ 14 | -DBUILD_TESTING=OFF \ 15 | -DFAISS_OPT_LEVEL=avx512 \ 16 | -DFAISS_ENABLE_GPU=ON \ 17 | -DFAISS_ENABLE_CUVS=ON \ 18 | -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \ 19 | -DFAISS_ENABLE_PYTHON=OFF \ 20 | -DBLA_VENDOR=Intel10_64lp \ 21 | -DCMAKE_INSTALL_LIBDIR=lib \ 22 | -DCMAKE_BUILD_TYPE=Release . 23 | 24 | make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512 faiss_c faiss_c_avx2 faiss_c_avx512 25 | 26 | cmake --install _build --prefix $PREFIX 27 | cmake --install _build --prefix _libfaiss_stage/ 28 | -------------------------------------------------------------------------------- /conda/faiss-gpu-cuvs/build-pkg.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so 11 | cmake -B _build_python_${PY_VER} \ 12 | -Dfaiss_ROOT=_libfaiss_stage/ \ 13 | -DFAISS_OPT_LEVEL=avx512 \ 14 | -DFAISS_ENABLE_GPU=ON \ 15 | -DFAISS_ENABLE_CUVS=ON \ 16 | -DCMAKE_BUILD_TYPE=Release \ 17 | -DPython_EXECUTABLE=$PYTHON \ 18 | faiss/python 19 | 20 | make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512 21 | 22 | # Build actual python module. 23 | cd _build_python_${PY_VER}/ 24 | $PYTHON setup.py install --single-version-externally-managed --record=record.txt --prefix=$PREFIX 25 | -------------------------------------------------------------------------------- /conda/faiss-gpu-cuvs/test_cpu_dispatch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so 10 | FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so 11 | FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so 12 | -------------------------------------------------------------------------------- /conda/faiss-gpu/build-lib.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | # Workaround for CUDA 11.4.4 builds. Moves all necessary headers to include root. 10 | if [ -n "$FAISS_FLATTEN_CONDA_INCLUDES" ] && [ "$FAISS_FLATTEN_CONDA_INCLUDES" = "1" ]; then 11 | cp -r -n "$CONDA_PREFIX/x86_64-conda-linux-gnu/sysroot/usr/include/"* "$CONDA_PREFIX/include/" 12 | cp -r -n "$CONDA_PREFIX/x86_64-conda-linux-gnu/include/c++/11.2.0/"* "$CONDA_PREFIX/include/" 13 | cp -r -n "$CONDA_PREFIX/x86_64-conda-linux-gnu/include/c++/11.2.0/x86_64-conda-linux-gnu/"* "$CONDA_PREFIX/include/" 14 | fi 15 | 16 | # Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so 17 | cmake -B _build \ 18 | -DBUILD_SHARED_LIBS=ON \ 19 | -DFAISS_ENABLE_C_API=ON \ 20 | -DBUILD_TESTING=OFF \ 21 | -DFAISS_OPT_LEVEL=avx512 \ 22 | -DFAISS_ENABLE_GPU=ON \ 23 | -DFAISS_ENABLE_CUVS=OFF \ 24 | -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \ 25 | -DFAISS_ENABLE_PYTHON=OFF \ 26 | -DBLA_VENDOR=Intel10_64lp \ 27 | -DCMAKE_INSTALL_LIBDIR=lib \ 28 | -DCMAKE_BUILD_TYPE=Release . 29 | 30 | make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512 faiss_c faiss_c_avx2 faiss_c_avx512 31 | 32 | cmake --install _build --prefix $PREFIX 33 | cmake --install _build --prefix _libfaiss_stage/ 34 | -------------------------------------------------------------------------------- /conda/faiss-gpu/build-pkg.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so 11 | cmake -B _build_python_${PY_VER} \ 12 | -Dfaiss_ROOT=_libfaiss_stage/ \ 13 | -DFAISS_OPT_LEVEL=avx512 \ 14 | -DFAISS_ENABLE_GPU=ON \ 15 | -DFAISS_ENABLE_CUVS=OFF \ 16 | -DCMAKE_BUILD_TYPE=Release \ 17 | -DPython_EXECUTABLE=$PYTHON \ 18 | faiss/python 19 | 20 | make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512 21 | 22 | # Build actual python module. 23 | cd _build_python_${PY_VER}/ 24 | $PYTHON setup.py install --single-version-externally-managed --record=record.txt --prefix=$PREFIX 25 | -------------------------------------------------------------------------------- /conda/faiss-gpu/test_cpu_dispatch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so 10 | FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so 11 | FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so 12 | -------------------------------------------------------------------------------- /conda/faiss/build-lib-arm64.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build libfaiss.so 11 | cmake -B _build \ 12 | -DBUILD_SHARED_LIBS=ON \ 13 | -DFAISS_ENABLE_C_API=ON \ 14 | -DBUILD_TESTING=OFF \ 15 | -DFAISS_OPT_LEVEL=sve \ 16 | -DFAISS_ENABLE_GPU=OFF \ 17 | -DFAISS_ENABLE_PYTHON=OFF \ 18 | -DCMAKE_INSTALL_LIBDIR=lib \ 19 | -DCMAKE_BUILD_TYPE=Release . 20 | 21 | make -C _build -j$(nproc) faiss faiss_sve faiss_c faiss_c_sve 22 | 23 | cmake --install _build --prefix $PREFIX 24 | cmake --install _build --prefix _libfaiss_stage/ 25 | -------------------------------------------------------------------------------- /conda/faiss/build-lib-osx.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so 11 | cmake -B _build \ 12 | -DBUILD_SHARED_LIBS=ON \ 13 | -DFAISS_ENABLE_C_API=ON \ 14 | -DBUILD_TESTING=OFF \ 15 | -DFAISS_OPT_LEVEL=avx512 \ 16 | -DFAISS_ENABLE_GPU=OFF \ 17 | -DFAISS_ENABLE_PYTHON=OFF \ 18 | -DBLA_VENDOR=Intel10_64lp \ 19 | -DOpenMP_CXX_FLAGS=-fopenmp=libiomp5 \ 20 | -DOpenMP_CXX_LIB_NAMES=libiomp5 \ 21 | -DOpenMP_libiomp5_LIBRARY=$PREFIX/lib/libiomp5.dylib \ 22 | -DCMAKE_INSTALL_LIBDIR=lib \ 23 | -DCMAKE_BUILD_TYPE=Release . 24 | 25 | make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512 faiss_c faiss_c_avx2 faiss_c_avx512 26 | 27 | cmake --install _build --prefix $PREFIX 28 | cmake --install _build --prefix _libfaiss_stage/ 29 | -------------------------------------------------------------------------------- /conda/faiss/build-lib.bat: -------------------------------------------------------------------------------- 1 | @REM Copyright (c) Meta Platforms, Inc. and affiliates. 2 | @REM 3 | @REM This source code is licensed under the MIT license found in the 4 | @REM LICENSE file in the root directory of this source tree. 5 | 6 | :: Copyright (c) Facebook, Inc. and its affiliates. 7 | :: 8 | :: This source code is licensed under the MIT license found in the 9 | :: LICENSE file in the root directory of this source tree. 10 | 11 | :: Build libfaiss.so. 12 | cmake -B _build ^ 13 | -T v141 ^ 14 | -A x64 ^ 15 | -G "Visual Studio 16 2019" ^ 16 | -DBUILD_SHARED_LIBS=ON ^ 17 | -DBUILD_TESTING=OFF ^ 18 | -DFAISS_ENABLE_GPU=OFF ^ 19 | -DFAISS_ENABLE_PYTHON=OFF ^ 20 | -DBLA_VENDOR=Intel10_64_dyn ^ 21 | . 22 | if %errorlevel% neq 0 exit /b %errorlevel% 23 | 24 | cmake --build _build --config Release -j %CPU_COUNT% 25 | if %errorlevel% neq 0 exit /b %errorlevel% 26 | 27 | cmake --install _build --config Release --prefix %PREFIX% 28 | if %errorlevel% neq 0 exit /b %errorlevel% 29 | -------------------------------------------------------------------------------- /conda/faiss/build-lib.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so 11 | cmake -B _build \ 12 | -DBUILD_SHARED_LIBS=ON \ 13 | -DFAISS_ENABLE_C_API=ON \ 14 | -DBUILD_TESTING=OFF \ 15 | -DFAISS_OPT_LEVEL=avx512 \ 16 | -DFAISS_ENABLE_GPU=OFF \ 17 | -DFAISS_ENABLE_PYTHON=OFF \ 18 | -DBLA_VENDOR=Intel10_64lp \ 19 | -DCMAKE_INSTALL_LIBDIR=lib \ 20 | -DCMAKE_BUILD_TYPE=Release . 21 | 22 | make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512 faiss_c faiss_c_avx2 faiss_c_avx512 23 | 24 | cmake --install _build --prefix $PREFIX 25 | cmake --install _build --prefix _libfaiss_stage/ 26 | -------------------------------------------------------------------------------- /conda/faiss/build-pkg-arm64.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build swigfaiss.so 11 | cmake -B _build_python_${PY_VER} \ 12 | -Dfaiss_ROOT=_libfaiss_stage/ \ 13 | -DFAISS_ENABLE_GPU=OFF \ 14 | -DCMAKE_BUILD_TYPE=Release \ 15 | -DPython_EXECUTABLE=$PYTHON \ 16 | faiss/python 17 | 18 | make -C _build_python_${PY_VER} -j$(nproc) swigfaiss 19 | 20 | # Build actual python module. 21 | cd _build_python_${PY_VER}/ 22 | $PYTHON setup.py install --single-version-externally-managed --record=record.txt --prefix=$PREFIX 23 | -------------------------------------------------------------------------------- /conda/faiss/build-pkg-osx.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512 11 | cmake -B _build_python_${PY_VER} \ 12 | -Dfaiss_ROOT=_libfaiss_stage/ \ 13 | -DFAISS_OPT_LEVEL=avx512 \ 14 | -DFAISS_ENABLE_GPU=OFF \ 15 | -DOpenMP_CXX_FLAGS=-fopenmp=libiomp5 \ 16 | -DOpenMP_CXX_LIB_NAMES=libiomp5 \ 17 | -DOpenMP_libiomp5_LIBRARY=$PREFIX/lib/libiomp5.dylib \ 18 | -DCMAKE_BUILD_TYPE=Release \ 19 | -DPython_EXECUTABLE=$PYTHON \ 20 | faiss/python 21 | 22 | make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512 23 | 24 | # Build actual python module. 25 | cd _build_python_${PY_VER}/ 26 | $PYTHON setup.py install --single-version-externally-managed --record=record.txt --prefix=$PREFIX 27 | -------------------------------------------------------------------------------- /conda/faiss/build-pkg.bat: -------------------------------------------------------------------------------- 1 | @REM Copyright (c) Meta Platforms, Inc. and affiliates. 2 | @REM 3 | @REM This source code is licensed under the MIT license found in the 4 | @REM LICENSE file in the root directory of this source tree. 5 | 6 | :: Copyright (c) Facebook, Inc. and its affiliates. 7 | :: 8 | :: This source code is licensed under the MIT license found in the 9 | :: LICENSE file in the root directory of this source tree. 10 | 11 | :: Build vanilla version (no avx). 12 | cmake -B _build_python_%PY_VER% ^ 13 | -T v141 ^ 14 | -A x64 ^ 15 | -G "Visual Studio 16 2019" ^ 16 | -DFAISS_ENABLE_GPU=OFF ^ 17 | -DPython_EXECUTABLE=%PYTHON% ^ 18 | faiss/python 19 | if %errorlevel% neq 0 exit /b %errorlevel% 20 | 21 | cmake --build _build_python_%PY_VER% --config Release -j %CPU_COUNT% 22 | if %errorlevel% neq 0 exit /b %errorlevel% 23 | 24 | 25 | :: Build actual python module. 26 | cd _build_python_%PY_VER%/ 27 | %PYTHON% setup.py install --single-version-externally-managed --record=record.txt --prefix=%PREFIX% 28 | if %errorlevel% neq 0 exit /b %errorlevel% 29 | -------------------------------------------------------------------------------- /conda/faiss/build-pkg.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | 10 | # Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so 11 | cmake -B _build_python_${PY_VER} \ 12 | -Dfaiss_ROOT=_libfaiss_stage/ \ 13 | -DFAISS_OPT_LEVEL=avx512 \ 14 | -DFAISS_ENABLE_GPU=OFF \ 15 | -DCMAKE_BUILD_TYPE=Release \ 16 | -DPython_EXECUTABLE=$PYTHON \ 17 | faiss/python 18 | 19 | make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512 20 | 21 | # Build actual python module. 22 | cd _build_python_${PY_VER}/ 23 | $PYTHON setup.py install --single-version-externally-managed --record=record.txt --prefix=$PREFIX 24 | -------------------------------------------------------------------------------- /conda/faiss/test_cpu_dispatch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | set -e 8 | 9 | FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so 10 | FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so 11 | FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so 12 | -------------------------------------------------------------------------------- /contrib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/faiss/f94f83531dfe9a538dcd51012aeecd31ff13affb/contrib/__init__.py -------------------------------------------------------------------------------- /contrib/torch/README.md: -------------------------------------------------------------------------------- 1 | # The Torch contrib 2 | 3 | This contrib directory contains a few Pytorch routines that 4 | are useful for similarity search. They do not necessarily depend on Faiss. 5 | 6 | The code is designed to work with CPU and GPU tensors. 7 | -------------------------------------------------------------------------------- /contrib/torch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/faiss/f94f83531dfe9a538dcd51012aeecd31ff13affb/contrib/torch/__init__.py -------------------------------------------------------------------------------- /contrib/torch/clustering.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | """ 7 | This contrib module contains Pytorch code for k-means clustering 8 | """ 9 | import faiss 10 | import faiss.contrib.torch_utils 11 | import torch 12 | 13 | # the kmeans can produce both torch and numpy centroids 14 | from faiss.contrib.clustering import kmeans 15 | 16 | 17 | class DatasetAssign: 18 | """Wrapper for a tensor that offers a function to assign the vectors 19 | to centroids. All other implementations offer the same interface""" 20 | 21 | def __init__(self, x): 22 | self.x = x 23 | 24 | def count(self): 25 | return self.x.shape[0] 26 | 27 | def dim(self): 28 | return self.x.shape[1] 29 | 30 | def get_subset(self, indices): 31 | return self.x[indices] 32 | 33 | def perform_search(self, centroids): 34 | return faiss.knn(self.x, centroids, 1) 35 | 36 | def assign_to(self, centroids, weights=None): 37 | D, I = self.perform_search(centroids) 38 | 39 | I = I.ravel() 40 | D = D.ravel() 41 | nc, d = centroids.shape 42 | 43 | sum_per_centroid = torch.zeros_like(centroids) 44 | if weights is None: 45 | sum_per_centroid.index_add_(0, I, self.x) 46 | else: 47 | sum_per_centroid.index_add_(0, I, self.x * weights[:, None]) 48 | 49 | # the indices are still in numpy. 50 | return I.cpu().numpy(), D, sum_per_centroid 51 | 52 | 53 | class DatasetAssignGPU(DatasetAssign): 54 | 55 | def __init__(self, res, x): 56 | DatasetAssign.__init__(self, x) 57 | self.res = res 58 | 59 | def perform_search(self, centroids): 60 | return faiss.knn_gpu(self.res, self.x, centroids, 1) 61 | -------------------------------------------------------------------------------- /contrib/vecs_io.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import sys 7 | import numpy as np 8 | 9 | """ 10 | I/O functions in fvecs, bvecs, ivecs formats 11 | definition of the formats here: http://corpus-texmex.irisa.fr/ 12 | """ 13 | 14 | 15 | def ivecs_read(fname): 16 | a = np.fromfile(fname, dtype='int32') 17 | if sys.byteorder == 'big': 18 | a.byteswap(inplace=True) 19 | d = a[0] 20 | return a.reshape(-1, d + 1)[:, 1:].copy() 21 | 22 | 23 | def fvecs_read(fname): 24 | return ivecs_read(fname).view('float32') 25 | 26 | 27 | def ivecs_mmap(fname): 28 | assert sys.byteorder != 'big' 29 | a = np.memmap(fname, dtype='int32', mode='r') 30 | d = a[0] 31 | return a.reshape(-1, d + 1)[:, 1:] 32 | 33 | 34 | def fvecs_mmap(fname): 35 | return ivecs_mmap(fname).view('float32') 36 | 37 | 38 | def bvecs_mmap(fname): 39 | x = np.memmap(fname, dtype='uint8', mode='r') 40 | if sys.byteorder == 'big': 41 | da = x[:4][::-1].copy() 42 | d = da.view('int32')[0] 43 | else: 44 | d = x[:4].view('int32')[0] 45 | return x.reshape(-1, d + 4)[:, 4:] 46 | 47 | 48 | def ivecs_write(fname, m): 49 | n, d = m.shape 50 | m1 = np.empty((n, d + 1), dtype='int32') 51 | m1[:, 0] = d 52 | m1[:, 1:] = m 53 | if sys.byteorder == 'big': 54 | m1.byteswap(inplace=True) 55 | m1.tofile(fname) 56 | 57 | 58 | def fvecs_write(fname, m): 59 | m = m.astype('float32') 60 | ivecs_write(fname, m.view('int32')) 61 | -------------------------------------------------------------------------------- /demos/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | add_executable(demo_imi_flat EXCLUDE_FROM_ALL demo_imi_flat.cpp) 7 | target_link_libraries(demo_imi_flat PRIVATE faiss) 8 | 9 | add_executable(demo_imi_pq EXCLUDE_FROM_ALL demo_imi_pq.cpp) 10 | target_link_libraries(demo_imi_pq PRIVATE faiss) 11 | 12 | add_executable(demo_ivfpq_indexing EXCLUDE_FROM_ALL demo_ivfpq_indexing.cpp) 13 | target_link_libraries(demo_ivfpq_indexing PRIVATE faiss) 14 | 15 | add_executable(demo_nndescent EXCLUDE_FROM_ALL demo_nndescent.cpp) 16 | target_link_libraries(demo_nndescent PRIVATE faiss) 17 | 18 | add_executable(demo_sift1M EXCLUDE_FROM_ALL demo_sift1M.cpp) 19 | target_link_libraries(demo_sift1M PRIVATE faiss) 20 | 21 | add_executable(demo_weighted_kmeans EXCLUDE_FROM_ALL demo_weighted_kmeans.cpp) 22 | target_link_libraries(demo_weighted_kmeans PRIVATE faiss) 23 | 24 | add_executable(demo_residual_quantizer EXCLUDE_FROM_ALL demo_residual_quantizer.cpp) 25 | target_link_libraries(demo_residual_quantizer PRIVATE faiss) 26 | -------------------------------------------------------------------------------- /demos/README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | Demos for a few Faiss functionalities 4 | ===================================== 5 | 6 | 7 | demo_auto_tune.py 8 | ----------------- 9 | 10 | Demonstrates the auto-tuning functionality of Faiss 11 | 12 | 13 | demo_ondisk_ivf.py 14 | ------------------ 15 | 16 | Shows how to construct a Faiss index that stores the inverted file 17 | data on disk, eg. when it does not fit in RAM. The script works on a 18 | small dataset (sift1M) for demonstration and proceeds in stages: 19 | 20 | 0: train on the dataset 21 | 22 | 1-4: build 4 indexes, each containing 1/4 of the dataset. This can be 23 | done in parallel on several machines 24 | 25 | 5: merge the 4 indexes into one that is written directly to disk 26 | (needs not to fit in RAM) 27 | 28 | 6: load and test the index 29 | -------------------------------------------------------------------------------- /demos/offline_ivf/README.md: -------------------------------------------------------------------------------- 1 | 2 | # Offline IVF 3 | 4 | This folder contains the code for the offline ivf algorithm powered by faiss big batch search. 5 | 6 | Create a conda env: 7 | 8 | `conda create --name oivf python=3.10` 9 | 10 | `conda activate oivf` 11 | 12 | `conda install -c pytorch/label/nightly -c nvidia faiss-gpu=1.7.4` 13 | 14 | `conda install tqdm` 15 | 16 | `conda install pyyaml` 17 | 18 | `conda install -c conda-forge submitit` 19 | 20 | 21 | ## Run book 22 | 23 | 1. Optionally shard your dataset (see create_sharded_dataset.py) and create the corresponding yaml file `config_ssnpp.yaml`. You can use `generate_config.py` by specifying the root directory of your dataset and the files with the data shards 24 | 25 | `python generate_config` 26 | 27 | 2. Run the train index command 28 | 29 | `python run.py --command train_index --config config_ssnpp.yaml --xb ssnpp_1B` 30 | 31 | 32 | 3. Run the index-shard command so it produces sharded indexes, required for the search step 33 | 34 | `python run.py --command index_shard --config config_ssnpp.yaml --xb ssnpp_1B` 35 | 36 | 37 | 6. Send jobs to the cluster to run search 38 | 39 | `python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --cluster_run --partition ` 40 | 41 | 42 | Remarks about the `search` command: it is assumed that the database vectors are the query vectors when performing the search step. 43 | a. If the query vectors are different than the database vectors, it should be passed in the xq argument 44 | b. A new dataset needs to be prepared (step 1) before passing it to the query vectors argument `–xq` 45 | 46 | `python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --xq ` 47 | 48 | 49 | 6. We can always run the consistency-check for sanity checks! 50 | 51 | `python run.py --command consistency_check--config config_ssnpp.yaml --xb ssnpp_1B` 52 | 53 | -------------------------------------------------------------------------------- /demos/offline_ivf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/faiss/f94f83531dfe9a538dcd51012aeecd31ff13affb/demos/offline_ivf/__init__.py -------------------------------------------------------------------------------- /demos/offline_ivf/generate_config.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | import os 8 | import yaml 9 | 10 | # with ssnpp sharded data 11 | root = "/checkpoint/marialomeli/ssnpp_data" 12 | file_names = [f"ssnpp_{i:010}.npy" for i in range(20)] 13 | d = 256 14 | dt = np.dtype(np.uint8) 15 | 16 | 17 | def read_embeddings(fp): 18 | fl = os.path.getsize(fp) 19 | nb = fl // d // dt.itemsize 20 | print(nb) 21 | if fl == d * dt.itemsize * nb: # no header 22 | return ("raw", np.memmap(fp, shape=(nb, d), dtype=dt, mode="r")) 23 | else: # assume npy 24 | vecs = np.load(fp, mmap_mode="r") 25 | assert vecs.shape[1] == d 26 | assert vecs.dtype == dt 27 | return ("npy", vecs) 28 | 29 | 30 | cfg = {} 31 | files = [] 32 | size = 0 33 | for fn in file_names: 34 | fp = f"{root}/{fn}" 35 | assert os.path.exists(fp), f"{fp} is missing" 36 | ft, xb = read_embeddings(fp) 37 | files.append( 38 | {"name": fn, "size": xb.shape[0], "dtype": dt.name, "format": ft} 39 | ) 40 | size += xb.shape[0] 41 | 42 | cfg["size"] = size 43 | cfg["root"] = root 44 | cfg["d"] = d 45 | cfg["files"] = files 46 | print(yaml.dump(cfg)) 47 | -------------------------------------------------------------------------------- /demos/rocksdb_ivf/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | cmake_minimum_required(VERSION 3.17 FATAL_ERROR) 7 | project (ROCKSDB_IVF) 8 | set(CMAKE_BUILD_TYPE Debug) 9 | find_package(faiss REQUIRED) 10 | find_package(RocksDB REQUIRED) 11 | 12 | add_executable(demo_rocksdb_ivf demo_rocksdb_ivf.cpp RocksDBInvertedLists.cpp) 13 | target_link_libraries(demo_rocksdb_ivf faiss RocksDB::rocksdb) 14 | -------------------------------------------------------------------------------- /demos/rocksdb_ivf/README.md: -------------------------------------------------------------------------------- 1 | # Storing Faiss inverted lists in RocksDB 2 | 3 | Demo of storing the inverted lists of any IVF index in RocksDB or any similar key-value store which supports the prefix scan operation. 4 | 5 | # How to build 6 | 7 | We use conda to create the build environment for simplicity. Only tested on Linux x86. 8 | 9 | ``` 10 | conda create -n rocksdb_ivf 11 | conda activate rocksdb_ivf 12 | conda install pytorch::faiss-cpu conda-forge::rocksdb cmake make gxx_linux-64 sysroot_linux-64 13 | cd ~/faiss/demos/rocksdb_ivf 14 | cmake -B build . 15 | make -C build -j$(nproc) 16 | ``` 17 | 18 | # Run the example 19 | 20 | ``` 21 | cd ~/faiss/demos/rocksdb_ivf/build 22 | ./rocksdb_ivf test_db 23 | ``` 24 | -------------------------------------------------------------------------------- /faiss/.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | # Ignore flakes about ambiguous variable name `I`. 3 | ignore = E741 4 | -------------------------------------------------------------------------------- /faiss/IndexBinaryFromFloat.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef FAISS_INDEX_BINARY_FROM_FLOAT_H 11 | #define FAISS_INDEX_BINARY_FROM_FLOAT_H 12 | 13 | #include 14 | 15 | namespace faiss { 16 | 17 | struct Index; 18 | 19 | /** IndexBinary backed by a float Index. 20 | * 21 | * Supports adding vertices and searching them. 22 | * 23 | * All queries are symmetric because there is no distinction between codes and 24 | * vectors. 25 | */ 26 | struct IndexBinaryFromFloat : IndexBinary { 27 | Index* index = nullptr; 28 | 29 | bool own_fields = false; ///< Whether object owns the index pointer. 30 | 31 | IndexBinaryFromFloat(); 32 | 33 | explicit IndexBinaryFromFloat(Index* index); 34 | 35 | ~IndexBinaryFromFloat(); 36 | 37 | void add(idx_t n, const uint8_t* x) override; 38 | 39 | void reset() override; 40 | 41 | void search( 42 | idx_t n, 43 | const uint8_t* x, 44 | idx_t k, 45 | int32_t* distances, 46 | idx_t* labels, 47 | const SearchParameters* params = nullptr) const override; 48 | 49 | void train(idx_t n, const uint8_t* x) override; 50 | }; 51 | 52 | } // namespace faiss 53 | 54 | #endif // FAISS_INDEX_BINARY_FROM_FLOAT_H 55 | -------------------------------------------------------------------------------- /faiss/IndexBinaryHNSW.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #pragma once 11 | 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | 18 | /** The HNSW index is a normal random-access index with a HNSW 19 | * link structure built on top */ 20 | 21 | struct IndexBinaryHNSW : IndexBinary { 22 | typedef HNSW::storage_idx_t storage_idx_t; 23 | 24 | // the link structure 25 | HNSW hnsw; 26 | 27 | // the sequential storage 28 | bool own_fields; 29 | IndexBinary* storage; 30 | 31 | explicit IndexBinaryHNSW(); 32 | explicit IndexBinaryHNSW(int d, int M = 32); 33 | explicit IndexBinaryHNSW(IndexBinary* storage, int M = 32); 34 | 35 | ~IndexBinaryHNSW() override; 36 | 37 | DistanceComputer* get_distance_computer() const; 38 | 39 | void add(idx_t n, const uint8_t* x) override; 40 | 41 | /// Trains the storage if needed 42 | void train(idx_t n, const uint8_t* x) override; 43 | 44 | /// entry point for search 45 | void search( 46 | idx_t n, 47 | const uint8_t* x, 48 | idx_t k, 49 | int32_t* distances, 50 | idx_t* labels, 51 | const SearchParameters* params = nullptr) const override; 52 | 53 | void reconstruct(idx_t key, uint8_t* recons) const override; 54 | 55 | void reset() override; 56 | }; 57 | 58 | } // namespace faiss 59 | -------------------------------------------------------------------------------- /faiss/IndexIVFIndependentQuantizer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | /** An IVF index with a quantizer that has a different input dimension from the 16 | * payload size. The vectors to encode are obtained from the input vectors by a 17 | * VectorTransform. 18 | */ 19 | struct IndexIVFIndependentQuantizer : Index { 20 | /// quantizer is fed directly with the input vectors 21 | Index* quantizer = nullptr; 22 | 23 | /// transform before the IVF vectors are applied 24 | VectorTransform* vt = nullptr; 25 | 26 | /// the IVF index, controls nlist and nprobe 27 | IndexIVF* index_ivf = nullptr; 28 | 29 | /// whether *this owns the 3 fields 30 | bool own_fields = false; 31 | 32 | IndexIVFIndependentQuantizer( 33 | Index* quantizer, 34 | IndexIVF* index_ivf, 35 | VectorTransform* vt = nullptr); 36 | 37 | IndexIVFIndependentQuantizer() {} 38 | 39 | void train(idx_t n, const float* x) override; 40 | 41 | void add(idx_t n, const float* x) override; 42 | 43 | void search( 44 | idx_t n, 45 | const float* x, 46 | idx_t k, 47 | float* distances, 48 | idx_t* labels, 49 | const SearchParameters* params = nullptr) const override; 50 | 51 | void reset() override; 52 | 53 | ~IndexIVFIndependentQuantizer() override; 54 | }; 55 | 56 | } // namespace faiss 57 | -------------------------------------------------------------------------------- /faiss/IndexLattice.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | 17 | /** Index that encodes a vector with a series of Zn lattice quantizers 18 | */ 19 | struct IndexLattice : IndexFlatCodes { 20 | /// number of sub-vectors 21 | int nsq; 22 | /// dimension of sub-vectors 23 | size_t dsq; 24 | 25 | /// the lattice quantizer 26 | ZnSphereCodecAlt zn_sphere_codec; 27 | 28 | /// nb bits used to encode the scale, per subvector 29 | int scale_nbit, lattice_nbit; 30 | 31 | /// mins and maxes of the vector norms, per subquantizer 32 | std::vector trained; 33 | 34 | IndexLattice(idx_t d, int nsq, int scale_nbit, int r2); 35 | 36 | void train(idx_t n, const float* x) override; 37 | 38 | /* The standalone codec interface */ 39 | size_t sa_code_size() const override; 40 | 41 | void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override; 42 | 43 | void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; 44 | }; 45 | 46 | } // namespace faiss 47 | -------------------------------------------------------------------------------- /faiss/IndexNNDescent.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #pragma once 11 | 12 | #include 13 | 14 | #include 15 | #include 16 | #include 17 | 18 | namespace faiss { 19 | 20 | /** The NNDescent index is a normal random-access index with an NNDescent 21 | * link structure built on top */ 22 | 23 | struct IndexNNDescent : Index { 24 | // internal storage of vectors (32 bits) 25 | using storage_idx_t = NNDescent::storage_idx_t; 26 | 27 | /// Faiss results are 64-bit 28 | 29 | // the link structure 30 | NNDescent nndescent; 31 | 32 | // the sequential storage 33 | bool own_fields; 34 | Index* storage; 35 | 36 | explicit IndexNNDescent( 37 | int d = 0, 38 | int K = 32, 39 | MetricType metric = METRIC_L2); 40 | explicit IndexNNDescent(Index* storage, int K = 32); 41 | 42 | ~IndexNNDescent() override; 43 | 44 | void add(idx_t n, const float* x) override; 45 | 46 | /// Trains the storage if needed 47 | void train(idx_t n, const float* x) override; 48 | 49 | /// entry point for search 50 | void search( 51 | idx_t n, 52 | const float* x, 53 | idx_t k, 54 | float* distances, 55 | idx_t* labels, 56 | const SearchParameters* params = nullptr) const override; 57 | 58 | void reconstruct(idx_t key, float* recons) const override; 59 | 60 | void reset() override; 61 | }; 62 | 63 | /** Flat index topped with with a NNDescent structure to access elements 64 | * more efficiently. 65 | */ 66 | 67 | struct IndexNNDescentFlat : IndexNNDescent { 68 | IndexNNDescentFlat(); 69 | IndexNNDescentFlat(int d, int K, MetricType metric = METRIC_L2); 70 | }; 71 | 72 | } // namespace faiss 73 | -------------------------------------------------------------------------------- /faiss/IndexNeuralNetCodec.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | 17 | struct IndexNeuralNetCodec : IndexFlatCodes { 18 | NeuralNetCodec* net = nullptr; 19 | size_t M, nbits; 20 | 21 | explicit IndexNeuralNetCodec( 22 | int d = 0, 23 | int M = 0, 24 | int nbits = 0, 25 | MetricType metric = METRIC_L2); 26 | 27 | void train(idx_t n, const float* x) override; 28 | 29 | void sa_encode(idx_t n, const float* x, uint8_t* codes) const override; 30 | void sa_decode(idx_t n, const uint8_t* codes, float* x) const override; 31 | 32 | ~IndexNeuralNetCodec() {} 33 | }; 34 | 35 | struct IndexQINCo : IndexNeuralNetCodec { 36 | QINCo qinco; 37 | 38 | IndexQINCo( 39 | int d, 40 | int M, 41 | int nbits, 42 | int L, 43 | int h, 44 | MetricType metric = METRIC_L2); 45 | 46 | ~IndexQINCo() {} 47 | }; 48 | 49 | } // namespace faiss 50 | -------------------------------------------------------------------------------- /faiss/IndexPQFastScan.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | 17 | /** Fast scan version of IndexPQ. Works for 4-bit PQ for now. 18 | * 19 | * The codes are not stored sequentially but grouped in blocks of size bbs. 20 | * This makes it possible to compute distances quickly with SIMD instructions. 21 | * 22 | * Implementations: 23 | * 12: blocked loop with internal loop on Q with qbs 24 | * 13: same with reservoir accumulator to store results 25 | * 14: no qbs with heap accumulator 26 | * 15: no qbs with reservoir accumulator 27 | */ 28 | 29 | struct IndexPQFastScan : IndexFastScan { 30 | ProductQuantizer pq; 31 | 32 | IndexPQFastScan( 33 | int d, 34 | size_t M, 35 | size_t nbits, 36 | MetricType metric = METRIC_L2, 37 | int bbs = 32); 38 | 39 | IndexPQFastScan() = default; 40 | 41 | /// build from an existing IndexPQ 42 | explicit IndexPQFastScan(const IndexPQ& orig, int bbs = 32); 43 | 44 | void train(idx_t n, const float* x) override; 45 | 46 | void compute_codes(uint8_t* codes, idx_t n, const float* x) const override; 47 | 48 | void compute_float_LUT(float* lut, idx_t n, const float* x) const override; 49 | 50 | void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override; 51 | }; 52 | 53 | } // namespace faiss 54 | -------------------------------------------------------------------------------- /faiss/IndexShardsIVF.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | /** 16 | * IndexShards with a common coarse quantizer. All the indexes added should be 17 | * IndexIVFInterface indexes so that the search_precomputed can be called. 18 | */ 19 | struct IndexShardsIVF : public IndexShards, Level1Quantizer { 20 | explicit IndexShardsIVF( 21 | Index* quantizer, 22 | size_t nlist, 23 | bool threaded = false, 24 | bool successive_ids = true); 25 | 26 | void addIndex(Index* index) override; 27 | 28 | void add_with_ids(idx_t n, const component_t* x, const idx_t* xids) 29 | override; 30 | 31 | void train(idx_t n, const component_t* x) override; 32 | 33 | void search( 34 | idx_t n, 35 | const component_t* x, 36 | idx_t k, 37 | distance_t* distances, 38 | idx_t* labels, 39 | const SearchParameters* params = nullptr) const override; 40 | }; 41 | 42 | } // namespace faiss 43 | -------------------------------------------------------------------------------- /faiss/MatrixStats.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #pragma once 11 | 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | namespace faiss { 19 | 20 | /** Reports some statistics on a dataset and comments on them. 21 | * 22 | * It is a class rather than a function so that all stats can also be 23 | * accessed from code */ 24 | 25 | struct MatrixStats { 26 | MatrixStats(size_t n, size_t d, const float* x); 27 | std::string comments; 28 | 29 | // raw statistics 30 | size_t n = 0, d = 0; 31 | size_t n_collision = 0; 32 | size_t n_valid = 0; 33 | size_t n0 = 0; 34 | double min_norm2 = HUGE_VALF; 35 | double max_norm2 = 0; 36 | uint64_t hash_value = 0; 37 | 38 | struct PerDimStats { 39 | /// counts of various special entries 40 | size_t n = 0; 41 | size_t n_nan = 0; 42 | size_t n_inf = 0; 43 | size_t n0 = 0; 44 | 45 | /// to get min/max and stddev values 46 | float min = HUGE_VALF; 47 | float max = -HUGE_VALF; 48 | double sum = 0; 49 | double sum2 = 0; 50 | 51 | size_t n_valid = 0; 52 | double mean = NAN; 53 | double stddev = NAN; 54 | 55 | void add(float x); 56 | void compute_mean_std(); 57 | }; 58 | 59 | std::vector per_dim_stats; 60 | struct Occurrence { 61 | size_t first; 62 | size_t count; 63 | }; 64 | std::unordered_map occurrences; 65 | 66 | char* buf; 67 | size_t nbuf; 68 | void do_comment(const char* fmt, ...); 69 | }; 70 | 71 | } // namespace faiss 72 | -------------------------------------------------------------------------------- /faiss/MetricType.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef FAISS_METRIC_TYPE_H 11 | #define FAISS_METRIC_TYPE_H 12 | 13 | #include 14 | 15 | namespace faiss { 16 | 17 | /// The metric space for vector comparison for Faiss indices and algorithms. 18 | /// 19 | /// Most algorithms support both inner product and L2, with the flat 20 | /// (brute-force) indices supporting additional metric types for vector 21 | /// comparison. 22 | enum MetricType { 23 | METRIC_INNER_PRODUCT = 0, ///< maximum inner product search 24 | METRIC_L2 = 1, ///< squared L2 search 25 | METRIC_L1, ///< L1 (aka cityblock) 26 | METRIC_Linf, ///< infinity distance 27 | METRIC_Lp, ///< L_p distance, p is given by a faiss::Index 28 | /// metric_arg 29 | 30 | /// some additional metrics defined in scipy.spatial.distance 31 | METRIC_Canberra = 20, 32 | METRIC_BrayCurtis, 33 | METRIC_JensenShannon, 34 | 35 | /// sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i)) where a_i, b_i > 0 36 | METRIC_Jaccard, 37 | /// Squared Eucliden distance, ignoring NaNs 38 | METRIC_NaNEuclidean, 39 | /// abs(x | y): the distance to a hyperplane 40 | METRIC_ABS_INNER_PRODUCT, 41 | }; 42 | 43 | /// all vector indices are this type 44 | using idx_t = int64_t; 45 | 46 | /// this function is used to distinguish between min and max indexes since 47 | /// we need to support similarity and dis-similarity metrics in a flexible way 48 | constexpr bool is_similarity_metric(MetricType metric_type) { 49 | return ((metric_type == METRIC_INNER_PRODUCT) || 50 | (metric_type == METRIC_Jaccard)); 51 | } 52 | 53 | } // namespace faiss 54 | 55 | #endif 56 | -------------------------------------------------------------------------------- /faiss/clone_index.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | // I/O code for indexes 11 | 12 | #pragma once 13 | 14 | namespace faiss { 15 | 16 | struct Index; 17 | struct IndexIVF; 18 | struct VectorTransform; 19 | struct Quantizer; 20 | struct IndexBinary; 21 | 22 | /* cloning functions */ 23 | Index* clone_index(const Index*); 24 | 25 | /** Cloner class, useful to override classes with other cloning 26 | * functions. The cloning function above just calls 27 | * Cloner::clone_Index. */ 28 | struct Cloner { 29 | virtual VectorTransform* clone_VectorTransform(const VectorTransform*); 30 | virtual Index* clone_Index(const Index*); 31 | virtual IndexIVF* clone_IndexIVF(const IndexIVF*); 32 | virtual ~Cloner() {} 33 | }; 34 | 35 | Quantizer* clone_Quantizer(const Quantizer* quant); 36 | 37 | IndexBinary* clone_binary_index(const IndexBinary* index); 38 | 39 | } // namespace faiss 40 | -------------------------------------------------------------------------------- /faiss/cppcontrib/detail/CoarseBitType.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | namespace cppcontrib { 14 | namespace detail { 15 | 16 | template 17 | struct CoarseBitType {}; 18 | 19 | template <> 20 | struct CoarseBitType<8> { 21 | using bit_type = uint8_t; 22 | }; 23 | 24 | template <> 25 | struct CoarseBitType<16> { 26 | using bit_type = uint16_t; 27 | }; 28 | 29 | } // namespace detail 30 | } // namespace cppcontrib 31 | } // namespace faiss 32 | -------------------------------------------------------------------------------- /faiss/cppcontrib/docker_dev/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:22.04 2 | 3 | RUN apt update && apt install -y python3 python3-pip gcc g++ mc git swig sudo libomp-dev libopenblas-dev wget 4 | RUN pip3 install numpy==1.26.4 scipy pytest 5 | RUN cd /root && git clone https://github.com/facebookresearch/faiss 6 | RUN wget -qO- "https://cmake.org/files/v3.26/cmake-3.26.5-linux-x86_64.tar.gz" | sudo tar --strip-components=1 -xz -C /usr/local 7 | RUN cd /root/faiss && /usr/local/bin/cmake -B build -DFAISS_ENABLE_GPU=OFF -DBUILD_TESTING=ON -DCMAKE_BUILD_TYPE=Release . 8 | RUN cd /root/faiss && make -C build -j 8 faiss 9 | RUN cd /root/faiss && make -C build -j 8 swigfaiss 10 | RUN cd /root/faiss/build/faiss/python && python3 setup.py install 11 | -------------------------------------------------------------------------------- /faiss/cppcontrib/factory_tools.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #pragma once 11 | 12 | #include 13 | 14 | namespace faiss { 15 | 16 | struct Index; 17 | struct IndexBinary; 18 | 19 | std::string reverse_index_factory(const faiss::Index* index); 20 | std::string reverse_index_factory(const faiss::IndexBinary* index); 21 | 22 | } // namespace faiss 23 | -------------------------------------------------------------------------------- /faiss/gpu/GpuAutoTune.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | /// parameter space and setters for GPU indexes 17 | struct GpuParameterSpace : faiss::ParameterSpace { 18 | /// initialize with reasonable parameters for the index 19 | void initialize(const faiss::Index* index) override; 20 | 21 | /// set a combination of parameters on an index 22 | void set_index_parameter( 23 | faiss::Index* index, 24 | const std::string& name, 25 | double val) const override; 26 | }; 27 | 28 | } // namespace gpu 29 | } // namespace faiss 30 | -------------------------------------------------------------------------------- /faiss/gpu/GpuFaissAssert.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #ifndef GPU_FAISS_ASSERT_INCLUDED 9 | #define GPU_FAISS_ASSERT_INCLUDED 10 | 11 | #include 12 | #include 13 | 14 | /// 15 | /// Assertions 16 | /// 17 | 18 | #if defined(__CUDA_ARCH__) || defined(USE_AMD_ROCM) 19 | #define GPU_FAISS_ASSERT(X) assert(X) 20 | #define GPU_FAISS_ASSERT_MSG(X, MSG) assert(X) 21 | #define GPU_FAISS_ASSERT_FMT(X, FMT, ...) assert(X) 22 | #else 23 | #define GPU_FAISS_ASSERT(X) FAISS_ASSERT(X) 24 | #define GPU_FAISS_ASSERT_MSG(X, MSG) FAISS_ASSERT_MSG(X, MSG) 25 | #define GPU_FAISS_ASSERT_FMT(X, FMT, ...) FAISS_ASSERT_FMT(X, FMT, __VA_ARGS) 26 | #endif // __CUDA_ARCH__ 27 | 28 | #endif 29 | -------------------------------------------------------------------------------- /faiss/gpu/GpuIcmEncoder.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | #include 13 | 14 | namespace faiss { 15 | namespace gpu { 16 | 17 | class GpuResourcesProvider; 18 | struct IcmEncoderShards; 19 | 20 | /** Perform LSQ encoding on GPU. 21 | * 22 | * Split input vectors to different devices and call IcmEncoderImpl::encode 23 | * to encode them 24 | */ 25 | class GpuIcmEncoder : public lsq::IcmEncoder { 26 | public: 27 | GpuIcmEncoder( 28 | const LocalSearchQuantizer* lsq, 29 | const std::vector& provs, 30 | const std::vector& devices); 31 | 32 | ~GpuIcmEncoder(); 33 | 34 | GpuIcmEncoder(const GpuIcmEncoder&) = delete; 35 | GpuIcmEncoder& operator=(const GpuIcmEncoder&) = delete; 36 | 37 | void set_binary_term() override; 38 | 39 | void encode( 40 | int32_t* codes, 41 | const float* x, 42 | std::mt19937& gen, 43 | size_t n, 44 | size_t ils_iters) const override; 45 | 46 | private: 47 | std::unique_ptr shards; 48 | }; 49 | 50 | struct GpuIcmEncoderFactory : public lsq::IcmEncoderFactory { 51 | explicit GpuIcmEncoderFactory(int ngpus = 1); 52 | 53 | lsq::IcmEncoder* get(const LocalSearchQuantizer* lsq) override; 54 | 55 | std::vector provs; 56 | std::vector devices; 57 | }; 58 | 59 | } // namespace gpu 60 | } // namespace faiss 61 | -------------------------------------------------------------------------------- /faiss/gpu/GpuIndicesOptions.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | /// How user vector index data is stored on the GPU 14 | enum IndicesOptions { 15 | /// The user indices are only stored on the CPU; the GPU returns 16 | /// (inverted list, offset) to the CPU which is then translated to 17 | /// the real user index. 18 | INDICES_CPU = 0, 19 | /// The indices are not stored at all, on either the CPU or 20 | /// GPU. Only (inverted list, offset) is returned to the user as the 21 | /// index. 22 | INDICES_IVF = 1, 23 | /// Indices are stored as 32 bit integers on the GPU, but returned 24 | /// as 64 bit integers 25 | INDICES_32_BIT = 2, 26 | /// Indices are stored as 64 bit integers on the GPU 27 | INDICES_64_BIT = 3, 28 | }; 29 | 30 | } // namespace gpu 31 | } // namespace faiss 32 | -------------------------------------------------------------------------------- /faiss/gpu/impl/BinaryDistance.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | // Performs brute-force k-NN comparison between `vecs` and `query`, where they 14 | // are encoded as binary vectors 15 | void runBinaryDistance( 16 | Tensor& vecs, 17 | Tensor& query, 18 | Tensor& outK, 19 | Tensor& outV, 20 | int k, 21 | cudaStream_t stream); 22 | 23 | } // namespace gpu 24 | } // namespace faiss 25 | -------------------------------------------------------------------------------- /faiss/gpu/impl/BinaryFlatIndex.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | class GpuResources; 17 | 18 | /// Holder of GPU resources for a particular flat index 19 | class BinaryFlatIndex { 20 | public: 21 | BinaryFlatIndex(GpuResources* res, int dim, MemorySpace space); 22 | 23 | /// Returns the number of vectors we contain 24 | idx_t getSize() const; 25 | 26 | idx_t getDim() const; 27 | 28 | /// Reserve storage that can contain at least this many vectors 29 | void reserve(size_t numVecs, cudaStream_t stream); 30 | 31 | /// Returns a reference to our vectors currently in use 32 | Tensor& getVectorsRef(); 33 | 34 | void query( 35 | Tensor& vecs, 36 | int k, 37 | Tensor& outDistances, 38 | Tensor& outIndices); 39 | 40 | /// Add vectors to ourselves; the pointer passed can be on the host 41 | /// or the device 42 | void add(const unsigned char* data, idx_t numVecs, cudaStream_t stream); 43 | 44 | /// Free all storage 45 | void reset(); 46 | 47 | private: 48 | /// Collection of GPU resources that we use 49 | GpuResources* resources_; 50 | 51 | /// Dimensionality of our vectors 52 | const int dim_; 53 | 54 | /// How many vectors we have 55 | idx_t num_; 56 | 57 | /// The underlying expandable storage 58 | DeviceVector rawData_; 59 | 60 | /// Vectors currently in rawData_ 61 | DeviceTensor vectors_; 62 | }; 63 | 64 | } // namespace gpu 65 | } // namespace faiss 66 | -------------------------------------------------------------------------------- /faiss/gpu/impl/BroadcastSum.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | // output[x][i] += input[i] for all x 17 | void runSumAlongColumns( 18 | Tensor& input, 19 | Tensor& output, 20 | cudaStream_t stream); 21 | 22 | void runSumAlongColumns( 23 | Tensor& input, 24 | Tensor& output, 25 | cudaStream_t stream); 26 | 27 | // output[x][i] = input[i] for all x 28 | void runAssignAlongColumns( 29 | Tensor& input, 30 | Tensor& output, 31 | cudaStream_t stream); 32 | 33 | void runAssignAlongColumns( 34 | Tensor& input, 35 | Tensor& output, 36 | cudaStream_t stream); 37 | 38 | // output[i][x] += input[i] for all x 39 | // If zeroClamp, output[i][x] = max(output[i][x] + input[i], 0) for all x 40 | void runSumAlongRows( 41 | Tensor& input, 42 | Tensor& output, 43 | bool zeroClamp, 44 | cudaStream_t stream); 45 | 46 | void runSumAlongRows( 47 | Tensor& input, 48 | Tensor& output, 49 | bool zeroClamp, 50 | cudaStream_t stream); 51 | 52 | } // namespace gpu 53 | } // namespace faiss 54 | -------------------------------------------------------------------------------- /faiss/gpu/impl/IVFFlatScan.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | #include 16 | 17 | namespace faiss { 18 | namespace gpu { 19 | 20 | class GpuResources; 21 | 22 | void runIVFFlatScan( 23 | Tensor& queries, 24 | Tensor& listIds, 25 | DeviceVector& listData, 26 | DeviceVector& listIndices, 27 | IndicesOptions indicesOptions, 28 | DeviceVector& listLengths, 29 | idx_t maxListLength, 30 | int k, 31 | faiss::MetricType metric, 32 | bool useResidual, 33 | Tensor& residualBase, 34 | GpuScalarQuantizer* scalarQ, 35 | // output 36 | Tensor& outDistances, 37 | // output 38 | Tensor& outIndices, 39 | GpuResources* res); 40 | 41 | } // namespace gpu 42 | } // namespace faiss 43 | -------------------------------------------------------------------------------- /faiss/gpu/impl/IndexUtils.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | /// A collection of various utility functions for index implementation 17 | 18 | /// Returns the maximum k-selection value supported based on the CUDA SDK that 19 | /// we were compiled with. .cu files can use DeviceDefs.cuh, but this is for 20 | /// non-CUDA files 21 | int getMaxKSelection() { 22 | return GPU_MAX_SELECTION_K; 23 | } 24 | 25 | void validateKSelect(int k) { 26 | FAISS_THROW_IF_NOT_FMT( 27 | k > 0 && k <= getMaxKSelection(), 28 | "GPU index only supports min/max-K selection up to %d (requested %d)", 29 | getMaxKSelection(), 30 | k); 31 | } 32 | 33 | void validateNProbe(size_t nprobe) { 34 | FAISS_THROW_IF_NOT_FMT( 35 | nprobe > 0 && nprobe <= (size_t)getMaxKSelection(), 36 | "GPU IVF index only supports nprobe selection up to %d (requested %zu)", 37 | getMaxKSelection(), 38 | nprobe); 39 | } 40 | 41 | } // namespace gpu 42 | } // namespace faiss 43 | -------------------------------------------------------------------------------- /faiss/gpu/impl/IndexUtils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | namespace gpu { 14 | 15 | /// A collection of various utility functions for index implementation 16 | 17 | /// Returns the maximum k-selection value supported based on the CUDA SDK that 18 | /// we were compiled with. .cu files can use DeviceDefs.cuh, but this is for 19 | /// non-CUDA files 20 | int getMaxKSelection(); 21 | 22 | // Validate the k parameter for search 23 | void validateKSelect(int k); 24 | 25 | // Validate the nprobe parameter for search 26 | void validateNProbe(size_t nprobe); 27 | 28 | } // namespace gpu 29 | } // namespace faiss 30 | -------------------------------------------------------------------------------- /faiss/gpu/impl/InterleavedCodes.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | // Utilities for bit packing and unpacking CPU non-interleaved and GPU 14 | // interleaved by 32 encodings 15 | namespace faiss { 16 | namespace gpu { 17 | 18 | // Unpacks arbitrary bitwidth codes to a whole number of bytes per code 19 | // The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ... 20 | // (bit packed) 21 | // The layout of the output is the same (byte packed to roundUp(bitsPerCode, 8) 22 | // / 8 bytes) 23 | std::vector unpackNonInterleaved( 24 | std::vector data, 25 | int numVecs, 26 | int dims, 27 | int bitsPerCode); 28 | 29 | // Unpacks arbitrary bitwidth codes to a whole number of bytes per scalar code 30 | // The layout of the input is (v0 d0)(v1 d0) ... (v31 d0)(v0 d1) ... 31 | // (bit packed) 32 | // The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ... 33 | // (byte packed) 34 | std::vector unpackInterleaved( 35 | std::vector data, 36 | int numVecs, 37 | int dims, 38 | int bitsPerCode); 39 | 40 | // Packs data in the byte packed non-interleaved form to bit packed 41 | // non-interleaved form 42 | std::vector packNonInterleaved( 43 | std::vector data, 44 | int numVecs, 45 | int dims, 46 | int bitsPerCode); 47 | 48 | // Packs data in the byte packed non-interleaved form to bit packed 49 | // interleaved form 50 | std::vector packInterleaved( 51 | std::vector data, 52 | int numVecs, 53 | int dims, 54 | int bitsPerCode); 55 | 56 | } // namespace gpu 57 | } // namespace faiss 58 | -------------------------------------------------------------------------------- /faiss/gpu/impl/L2Norm.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | void runL2Norm( 17 | Tensor& input, 18 | bool inputRowMajor, 19 | Tensor& output, 20 | bool normSquared, 21 | cudaStream_t stream); 22 | 23 | void runL2Norm( 24 | Tensor& input, 25 | bool inputRowMajor, 26 | Tensor& output, 27 | bool normSquared, 28 | cudaStream_t stream); 29 | 30 | void runL2Norm( 31 | Tensor<__nv_bfloat16, 2, true>& input, 32 | bool inputRowMajor, 33 | Tensor& output, 34 | bool normSquared, 35 | cudaStream_t stream); 36 | 37 | } // namespace gpu 38 | } // namespace faiss 39 | -------------------------------------------------------------------------------- /faiss/gpu/impl/L2Select.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | namespace gpu { 14 | 15 | void runL2SelectMin( 16 | Tensor& productDistances, 17 | Tensor& centroidDistances, 18 | Tensor& outDistances, 19 | Tensor& outIndices, 20 | int k, 21 | cudaStream_t stream); 22 | 23 | } 24 | } // namespace faiss 25 | -------------------------------------------------------------------------------- /faiss/gpu/impl/PQCodeDistances.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | namespace gpu { 17 | 18 | class DeviceMemory; 19 | 20 | /// pqCentroids is of the form (sub q)(sub dim)(code id) 21 | /// Calculates the distance from the (query - centroid) residual to 22 | /// each sub-code vector, for the given list of query results in 23 | /// coarseIndices 24 | template 25 | void runPQCodeDistances( 26 | GpuResources* res, 27 | Tensor& pqCentroids, 28 | Tensor& queries, 29 | Tensor& coarseCentroids, 30 | Tensor& coarseDistances, 31 | Tensor& coarseIndices, 32 | NoTypeTensor<4, true>& outCodeDistances, 33 | bool useMMImplementation, 34 | bool l2Distance, 35 | bool useFloat16Lookup); 36 | 37 | } // namespace gpu 38 | } // namespace faiss 39 | 40 | #include 41 | -------------------------------------------------------------------------------- /faiss/gpu/impl/PQScanMultiPassNoPrecomputed.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | namespace gpu { 18 | 19 | class GpuResources; 20 | 21 | template 22 | void runPQScanMultiPassNoPrecomputed( 23 | Tensor& queries, 24 | Tensor& centroids, 25 | Tensor& pqCentroidsInnermostCode, 26 | Tensor& coarseDistances, 27 | Tensor& coarseIndices, 28 | bool useFloat16Lookup, 29 | bool useMMCodeDistance, 30 | bool interleavedCodeLayout, 31 | int bitsPerSubQuantizer, 32 | int numSubQuantizers, 33 | int numSubQuantizerCodes, 34 | DeviceVector& listCodes, 35 | DeviceVector& listIndices, 36 | IndicesOptions indicesOptions, 37 | DeviceVector& listLengths, 38 | idx_t maxListLength, 39 | int k, 40 | faiss::MetricType metric, 41 | // output 42 | Tensor& outDistances, 43 | // output 44 | Tensor& outIndices, 45 | GpuResources* res); 46 | 47 | } // namespace gpu 48 | } // namespace faiss 49 | 50 | #include 51 | -------------------------------------------------------------------------------- /faiss/gpu/impl/PQScanMultiPassPrecomputed.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | namespace gpu { 18 | 19 | class GpuResources; 20 | 21 | void runPQScanMultiPassPrecomputed( 22 | Tensor& queries, 23 | Tensor& precompTerm1, 24 | NoTypeTensor<3, true>& precompTerm2, 25 | NoTypeTensor<3, true>& precompTerm3, 26 | Tensor& ivfListIds, 27 | bool useFloat16Lookup, 28 | bool interleavedCodeLayout, 29 | int bitsPerSubQuantizer, 30 | int numSubQuantizers, 31 | int numSubQuantizerCodes, 32 | DeviceVector& listCodes, 33 | DeviceVector& listIndices, 34 | IndicesOptions indicesOptions, 35 | DeviceVector& listLengths, 36 | idx_t maxListLength, 37 | int k, 38 | // output 39 | Tensor& outDistances, 40 | // output 41 | Tensor& outIndices, 42 | GpuResources* res); 43 | 44 | } // namespace gpu 45 | } // namespace faiss 46 | -------------------------------------------------------------------------------- /faiss/gpu/impl/RemapIndices.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | namespace gpu { 13 | 14 | // Utility function to translate (list id, offset) to a user index on 15 | // the CPU. In a cpp in order to use OpenMP 16 | void ivfOffsetToUserIndex( 17 | idx_t* indices, 18 | idx_t numLists, 19 | idx_t queries, 20 | int k, 21 | const std::vector>& listOffsetToUserIndex) { 22 | FAISS_ASSERT(numLists == listOffsetToUserIndex.size()); 23 | 24 | #pragma omp parallel for 25 | for (idx_t q = 0; q < queries; ++q) { 26 | for (idx_t r = 0; r < k; ++r) { 27 | auto offsetIndex = indices[q * k + r]; 28 | 29 | if (offsetIndex < 0) { 30 | continue; 31 | } 32 | 33 | // FIXME: implicit limit on list and list offset length 34 | int listId = (int)(offsetIndex >> 32); 35 | int listOffset = (int)(offsetIndex & 0xffffffff); 36 | 37 | FAISS_ASSERT(listId < numLists); 38 | auto& listIndices = listOffsetToUserIndex[listId]; 39 | 40 | FAISS_ASSERT(listOffset < listIndices.size()); 41 | indices[q * k + r] = listIndices[listOffset]; 42 | } 43 | } 44 | } 45 | 46 | } // namespace gpu 47 | } // namespace faiss 48 | -------------------------------------------------------------------------------- /faiss/gpu/impl/RemapIndices.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | /// Utility function to translate (list id, offset) to a user index on 17 | /// the CPU. In a cpp in order to use OpenMP. 18 | void ivfOffsetToUserIndex( 19 | idx_t* indices, 20 | idx_t numLists, 21 | idx_t queries, 22 | int k, 23 | const std::vector>& listOffsetToUserIndex); 24 | 25 | } // namespace gpu 26 | } // namespace faiss 27 | -------------------------------------------------------------------------------- /faiss/gpu/impl/VectorResidual.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | namespace faiss { 16 | namespace gpu { 17 | 18 | // Calculates residual v_i - c_j for all v_i in vecs where j = vecToCentroid[i] 19 | void runCalcResidual( 20 | Tensor& vecs, 21 | Tensor& centroids, 22 | Tensor& vecToCentroid, 23 | Tensor& residuals, 24 | cudaStream_t stream); 25 | 26 | void runCalcResidual( 27 | Tensor& vecs, 28 | Tensor& centroids, 29 | Tensor& vecToCentroid, 30 | Tensor& residuals, 31 | cudaStream_t stream); 32 | 33 | // Gather vectors 34 | void runReconstruct( 35 | Tensor& listIds, 36 | Tensor& vecs, 37 | Tensor& out, 38 | cudaStream_t stream); 39 | 40 | void runReconstruct( 41 | Tensor& listIds, 42 | Tensor& vecs, 43 | Tensor& out, 44 | cudaStream_t stream); 45 | 46 | void runReconstruct( 47 | idx_t start, 48 | idx_t num, 49 | Tensor& vecs, 50 | Tensor& out, 51 | cudaStream_t stream); 52 | 53 | void runReconstruct( 54 | idx_t start, 55 | idx_t num, 56 | Tensor& vecs, 57 | Tensor& out, 58 | cudaStream_t stream); 59 | 60 | } // namespace gpu 61 | } // namespace faiss 62 | -------------------------------------------------------------------------------- /faiss/gpu/perf/IndexWrapper.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | namespace gpu { 18 | 19 | // If we want to run multi-GPU, create a proxy to wrap the indices. 20 | // If we don't want multi-GPU, don't involve the proxy, so it doesn't 21 | // affect the timings. 22 | template 23 | struct IndexWrapper { 24 | std::vector> resources; 25 | std::vector> subIndex; 26 | std::unique_ptr replicaIndex; 27 | 28 | IndexWrapper( 29 | int numGpus, 30 | std::function(GpuResourcesProvider*, int)> 31 | init); 32 | faiss::Index* getIndex(); 33 | 34 | void runOnIndices(std::function f); 35 | void setNumProbes(size_t nprobe); 36 | }; 37 | 38 | } // namespace gpu 39 | } // namespace faiss 40 | 41 | #include 42 | -------------------------------------------------------------------------------- /faiss/gpu/perf/slow.py: -------------------------------------------------------------------------------- 1 | #! /usr/bin/env python3 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the MIT license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # this is a slow computation to test whether ctrl-C handling works 8 | import faiss 9 | import numpy as np 10 | 11 | def test_slow(): 12 | d = 256 13 | index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 14 | 0, faiss.IndexFlatL2(d)) 15 | x = np.random.rand(10 ** 6, d).astype('float32') 16 | print('add') 17 | index.add(x) 18 | print('search') 19 | index.search(x, 10) 20 | print('done') 21 | 22 | 23 | if __name__ == '__main__': 24 | test_slow() 25 | -------------------------------------------------------------------------------- /faiss/gpu/test/test_gpu_index_ivfflat.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import unittest 7 | 8 | import faiss 9 | import numpy as np 10 | 11 | 12 | class TestGpuIndexIvfflat(unittest.TestCase): 13 | def test_reconstruct_n(self): 14 | index = faiss.index_factory(4, "IVF10,Flat") 15 | x = np.random.RandomState(123).rand(10, 4).astype('float32') 16 | index.train(x) 17 | index.add(x) 18 | res = faiss.StandardGpuResources() 19 | res.noTempMemory() 20 | config = faiss.GpuIndexIVFFlatConfig() 21 | config.use_cuvs = False 22 | index2 = faiss.GpuIndexIVFFlat(res, index, config) 23 | recons = index2.reconstruct_n(0, 10) 24 | 25 | np.testing.assert_array_equal(recons, x) 26 | -------------------------------------------------------------------------------- /faiss/gpu/utils/Comparators.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | template 17 | struct Comparator { 18 | __device__ static inline bool lt(T a, T b) { 19 | return a < b; 20 | } 21 | 22 | __device__ static inline bool gt(T a, T b) { 23 | return a > b; 24 | } 25 | }; 26 | 27 | template <> 28 | struct Comparator { 29 | __device__ static inline bool lt(half a, half b) { 30 | #if FAISS_USE_FULL_FLOAT16 31 | return __hlt(a, b); 32 | #else 33 | return __half2float(a) < __half2float(b); 34 | #endif // FAISS_USE_FULL_FLOAT16 35 | } 36 | 37 | __device__ static inline bool gt(half a, half b) { 38 | #if FAISS_USE_FULL_FLOAT16 39 | return __hgt(a, b); 40 | #else 41 | return __half2float(a) > __half2float(b); 42 | #endif // FAISS_USE_FULL_FLOAT16 43 | } 44 | }; 45 | 46 | } // namespace gpu 47 | } // namespace faiss 48 | -------------------------------------------------------------------------------- /faiss/gpu/utils/DeviceDefs.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | namespace gpu { 14 | 15 | #ifdef USE_AMD_ROCM 16 | 17 | #if __AMDGCN_WAVEFRONT_SIZE == 32u 18 | constexpr int kWarpSize = 32; 19 | #else 20 | constexpr int kWarpSize = 64; 21 | #endif 22 | 23 | // This is a memory barrier for intra-warp writes to shared memory. 24 | __forceinline__ __device__ void warpFence() { 25 | __threadfence_block(); 26 | } 27 | 28 | #define GPU_MAX_SELECTION_K 2048 29 | 30 | #else // USE_AMD_ROCM 31 | 32 | // We require at least CUDA 8.0 for compilation 33 | #if CUDA_VERSION < 8000 34 | #error "CUDA >= 8.0 is required" 35 | #endif 36 | 37 | // We validate this against the actual architecture in device initialization 38 | constexpr int kWarpSize = 32; 39 | 40 | // This is a memory barrier for intra-warp writes to shared memory. 41 | __forceinline__ __device__ void warpFence() { 42 | #if CUDA_VERSION >= 9000 43 | __syncwarp(); 44 | #else 45 | // For the time being, assume synchronicity. 46 | // __threadfence_block(); 47 | #endif 48 | } 49 | 50 | #if CUDA_VERSION > 9000 51 | // Based on the CUDA version (we assume what version of nvcc/ptxas we were 52 | // compiled with), the register allocation algorithm is much better, so only 53 | // enable the 2048 selection code if we are above 9.0 (9.2 seems to be ok) 54 | #define GPU_MAX_SELECTION_K 2048 55 | #else 56 | #define GPU_MAX_SELECTION_K 1024 57 | #endif 58 | 59 | #endif // USE_AMD_ROCM 60 | 61 | } // namespace gpu 62 | } // namespace faiss 63 | -------------------------------------------------------------------------------- /faiss/gpu/utils/MatrixMult.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | namespace gpu { 18 | 19 | class GpuResources; 20 | 21 | /// C = alpha * A * B + beta * C 22 | /// Expects row major layout, not fortran/blas column major! 23 | template 24 | void runMatrixMult( 25 | Tensor& c, 26 | bool transC, 27 | Tensor& a, 28 | bool transA, 29 | Tensor& b, 30 | bool transB, 31 | float alpha, 32 | float beta, 33 | cublasHandle_t handle, 34 | cudaStream_t stream); 35 | 36 | /// C_i = alpha * A_i * B_i + beta * C_i 37 | /// where `i` is the outermost dimension, via iterated gemm 38 | /// Expects row major layout, not fortran/blas column major! 39 | template 40 | void runIteratedMatrixMult( 41 | Tensor& c, 42 | bool transC, 43 | Tensor& a, 44 | bool transA, 45 | Tensor& b, 46 | bool transB, 47 | float alpha, 48 | float beta, 49 | cublasHandle_t handle, 50 | cudaStream_t stream); 51 | 52 | } // namespace gpu 53 | } // namespace faiss 54 | 55 | #include 56 | -------------------------------------------------------------------------------- /faiss/gpu/utils/MergeNetworkUtils.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | template 14 | inline __device__ void swap(bool swap, T& x, T& y) { 15 | T tmp = x; 16 | x = swap ? y : x; 17 | y = swap ? tmp : y; 18 | } 19 | 20 | template 21 | inline __device__ void assign(bool assign, T& x, T y) { 22 | x = assign ? y : x; 23 | } 24 | 25 | } // namespace gpu 26 | } // namespace faiss 27 | -------------------------------------------------------------------------------- /faiss/gpu/utils/ReductionOperators.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | namespace gpu { 17 | 18 | template 19 | struct Sum { 20 | __device__ inline T operator()(T a, T b) const { 21 | return Math::add(a, b); 22 | } 23 | 24 | inline __device__ T identity() const { 25 | return Math::zero(); 26 | } 27 | }; 28 | 29 | template 30 | struct Min { 31 | __device__ inline T operator()(T a, T b) const { 32 | return Math::lt(a, b) ? a : b; 33 | } 34 | 35 | inline __device__ T identity() const { 36 | return Limits::getMax(); 37 | } 38 | }; 39 | 40 | template 41 | struct Max { 42 | __device__ inline T operator()(T a, T b) const { 43 | return Math::gt(a, b) ? a : b; 44 | } 45 | 46 | inline __device__ T identity() const { 47 | return Limits::getMin(); 48 | } 49 | }; 50 | 51 | /// Used for producing segmented prefix scans; the value of the Pair 52 | /// denotes the start of a new segment for the scan 53 | template 54 | struct SegmentedReduce { 55 | inline __device__ SegmentedReduce(const ReduceOp& o) : op(o) {} 56 | 57 | __device__ inline Pair operator()( 58 | const Pair& a, 59 | const Pair& b) const { 60 | return Pair(b.v ? b.k : op(a.k, b.k), a.v || b.v); 61 | } 62 | 63 | inline __device__ Pair identity() const { 64 | return Pair(op.identity(), false); 65 | } 66 | 67 | ReduceOp op; 68 | }; 69 | 70 | } // namespace gpu 71 | } // namespace faiss 72 | -------------------------------------------------------------------------------- /faiss/gpu/utils/Timer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | KernelTimer::KernelTimer(cudaStream_t stream) 17 | : startEvent_(nullptr), 18 | stopEvent_(nullptr), 19 | stream_(stream), 20 | valid_(true) { 21 | CUDA_VERIFY(cudaEventCreate(&startEvent_)); 22 | CUDA_VERIFY(cudaEventCreate(&stopEvent_)); 23 | 24 | CUDA_VERIFY(cudaEventRecord(startEvent_, stream_)); 25 | } 26 | 27 | KernelTimer::~KernelTimer() { 28 | CUDA_VERIFY(cudaEventDestroy(startEvent_)); 29 | CUDA_VERIFY(cudaEventDestroy(stopEvent_)); 30 | } 31 | 32 | float KernelTimer::elapsedMilliseconds() { 33 | FAISS_ASSERT(valid_); 34 | 35 | CUDA_VERIFY(cudaEventRecord(stopEvent_, stream_)); 36 | CUDA_VERIFY(cudaEventSynchronize(stopEvent_)); 37 | 38 | auto time = 0.0f; 39 | CUDA_VERIFY(cudaEventElapsedTime(&time, startEvent_, stopEvent_)); 40 | valid_ = false; 41 | 42 | return time; 43 | } 44 | 45 | CpuTimer::CpuTimer() { 46 | start_ = std::chrono::steady_clock::now(); 47 | } 48 | 49 | float CpuTimer::elapsedMilliseconds() { 50 | auto end = std::chrono::steady_clock::now(); 51 | 52 | std::chrono::duration duration = end - start_; 53 | 54 | return duration.count(); 55 | } 56 | 57 | } // namespace gpu 58 | } // namespace faiss 59 | -------------------------------------------------------------------------------- /faiss/gpu/utils/Timer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | namespace gpu { 15 | 16 | /// Utility class for timing execution of a kernel 17 | class KernelTimer { 18 | public: 19 | /// Constructor starts the timer and adds an event into the current 20 | /// device stream 21 | KernelTimer(cudaStream_t stream = nullptr); 22 | 23 | /// Destructor releases event resources 24 | ~KernelTimer(); 25 | 26 | /// Adds a stop event then synchronizes on the stop event to get the 27 | /// actual GPU-side kernel timings for any kernels launched in the 28 | /// current stream. Returns the number of milliseconds elapsed. 29 | /// Can only be called once. 30 | float elapsedMilliseconds(); 31 | 32 | private: 33 | cudaEvent_t startEvent_; 34 | cudaEvent_t stopEvent_; 35 | cudaStream_t stream_; 36 | bool valid_; 37 | }; 38 | 39 | /// CPU wallclock elapsed timer 40 | class CpuTimer { 41 | public: 42 | /// Creates and starts a new timer 43 | CpuTimer(); 44 | 45 | /// Returns elapsed time in milliseconds 46 | float elapsedMilliseconds(); 47 | 48 | private: 49 | std::chrono::time_point start_; 50 | }; 51 | 52 | } // namespace gpu 53 | } // namespace faiss 54 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloat1.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 1, 1); 14 | BLOCK_SELECT_IMPL(float, false, 1, 1); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloat128.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 128, 3); 14 | BLOCK_SELECT_IMPL(float, false, 128, 3); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloat256.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 256, 4); 14 | BLOCK_SELECT_IMPL(float, false, 256, 4); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloat32.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 32, 2); 14 | BLOCK_SELECT_IMPL(float, false, 32, 2); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloat64.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 64, 3); 14 | BLOCK_SELECT_IMPL(float, false, 64, 3); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatF1024.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, false, 1024, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatF2048.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | namespace gpu { 13 | 14 | #if GPU_MAX_SELECTION_K >= 2048 15 | BLOCK_SELECT_IMPL(float, false, 2048, 8); 16 | #endif 17 | 18 | } // namespace gpu 19 | } // namespace faiss 20 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatF512.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, false, 512, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatT1024.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 1024, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatT2048.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | namespace gpu { 13 | 14 | #if GPU_MAX_SELECTION_K >= 2048 15 | BLOCK_SELECT_IMPL(float, true, 2048, 8); 16 | #endif 17 | 18 | } // namespace gpu 19 | } // namespace faiss 20 | -------------------------------------------------------------------------------- /faiss/gpu/utils/blockselect/BlockSelectFloatT512.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | BLOCK_SELECT_IMPL(float, true, 512, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloat1.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 1, 1); 14 | WARP_SELECT_IMPL(float, false, 1, 1); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloat128.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 128, 3); 14 | WARP_SELECT_IMPL(float, false, 128, 3); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloat256.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 256, 4); 14 | WARP_SELECT_IMPL(float, false, 256, 4); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloat32.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 32, 2); 14 | WARP_SELECT_IMPL(float, false, 32, 2); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloat64.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 64, 3); 14 | WARP_SELECT_IMPL(float, false, 64, 3); 15 | 16 | } // namespace gpu 17 | } // namespace faiss 18 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatF1024.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, false, 1024, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatF2048.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | namespace gpu { 13 | 14 | #if GPU_MAX_SELECTION_K >= 2048 15 | WARP_SELECT_IMPL(float, false, 2048, 8); 16 | #endif 17 | 18 | } // namespace gpu 19 | } // namespace faiss 20 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatF512.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, false, 512, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatT1024.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 1024, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatT2048.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | namespace gpu { 13 | 14 | #if GPU_MAX_SELECTION_K >= 2048 15 | WARP_SELECT_IMPL(float, true, 2048, 8); 16 | #endif 17 | 18 | } // namespace gpu 19 | } // namespace faiss 20 | -------------------------------------------------------------------------------- /faiss/gpu/utils/warpselect/WarpSelectFloatT512.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | namespace faiss { 11 | namespace gpu { 12 | 13 | WARP_SELECT_IMPL(float, true, 512, 8); 14 | 15 | } 16 | } // namespace faiss 17 | -------------------------------------------------------------------------------- /faiss/impl/CodePacker.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | /********************************************* 16 | * CodePacker 17 | * default of pack_all / unpack_all loops over the _1 versions 18 | */ 19 | 20 | void CodePacker::pack_all(const uint8_t* flat_codes, uint8_t* block) const { 21 | for (size_t i = 0; i < nvec; i++) { 22 | pack_1(flat_codes + code_size * i, i, block); 23 | } 24 | } 25 | 26 | void CodePacker::unpack_all(const uint8_t* block, uint8_t* flat_codes) const { 27 | for (size_t i = 0; i < nvec; i++) { 28 | unpack_1(block, i, flat_codes + code_size * i); 29 | } 30 | } 31 | 32 | /********************************************* 33 | * CodePackerFlat 34 | */ 35 | 36 | CodePackerFlat::CodePackerFlat(size_t code_size) { 37 | this->code_size = code_size; 38 | nvec = 1; 39 | block_size = code_size; 40 | } 41 | 42 | void CodePackerFlat::pack_all(const uint8_t* flat_codes, uint8_t* block) const { 43 | memcpy(block, flat_codes, code_size); 44 | } 45 | 46 | void CodePackerFlat::unpack_all(const uint8_t* block, uint8_t* flat_codes) 47 | const { 48 | memcpy(flat_codes, block, code_size); 49 | } 50 | 51 | void CodePackerFlat::pack_1( 52 | const uint8_t* flat_code, 53 | size_t offset, 54 | uint8_t* block) const { 55 | assert(offset == 0); 56 | pack_all(flat_code, block); 57 | } 58 | 59 | void CodePackerFlat::unpack_1( 60 | const uint8_t* block, 61 | size_t offset, 62 | uint8_t* flat_code) const { 63 | assert(offset == 0); 64 | unpack_all(block, flat_code); 65 | } 66 | 67 | } // namespace faiss 68 | -------------------------------------------------------------------------------- /faiss/impl/FaissException.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // -*- c++ -*- 9 | 10 | #ifndef FAISS_EXCEPTION_INCLUDED 11 | #define FAISS_EXCEPTION_INCLUDED 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | namespace faiss { 19 | 20 | /// Base class for Faiss exceptions 21 | class FaissException : public std::exception { 22 | public: 23 | explicit FaissException(const std::string& msg); 24 | 25 | FaissException( 26 | const std::string& msg, 27 | const char* funcName, 28 | const char* file, 29 | int line); 30 | 31 | /// from std::exception 32 | const char* what() const noexcept override; 33 | 34 | std::string msg; 35 | }; 36 | 37 | /// Handle multiple exceptions from worker threads, throwing an appropriate 38 | /// exception that aggregates the information 39 | /// The pair int is the thread that generated the exception 40 | void handleExceptions( 41 | std::vector>& exceptions); 42 | 43 | /** RAII object for a set of possibly transformed vectors (deallocated only if 44 | * they are indeed transformed) 45 | */ 46 | struct TransformedVectors { 47 | const float* x; 48 | bool own_x; 49 | TransformedVectors(const float* x_orig, const float* x) : x(x) { 50 | own_x = x_orig != x; 51 | } 52 | 53 | ~TransformedVectors() { 54 | if (own_x) { 55 | delete[] x; 56 | } 57 | } 58 | }; 59 | 60 | /// make typeids more readable 61 | std::string demangle_cpp_symbol(const char* name); 62 | 63 | } // namespace faiss 64 | 65 | #endif 66 | -------------------------------------------------------------------------------- /faiss/impl/Quantizer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | 14 | /** General interface for quantizer objects */ 15 | struct Quantizer { 16 | size_t d; ///< size of the input vectors 17 | size_t code_size; ///< bytes per indexed vector 18 | 19 | explicit Quantizer(size_t d = 0, size_t code_size = 0) 20 | : d(d), code_size(code_size) {} 21 | 22 | /** Train the quantizer 23 | * 24 | * @param x training vectors, size n * d 25 | */ 26 | virtual void train(size_t n, const float* x) = 0; 27 | 28 | /** Quantize a set of vectors 29 | * 30 | * @param x input vectors, size n * d 31 | * @param codes output codes, size n * code_size 32 | */ 33 | virtual void compute_codes(const float* x, uint8_t* codes, size_t n) 34 | const = 0; 35 | 36 | /** Decode a set of vectors 37 | * 38 | * @param codes input codes, size n * code_size 39 | * @param x output vectors, size n * d 40 | */ 41 | virtual void decode(const uint8_t* code, float* x, size_t n) const = 0; 42 | 43 | virtual ~Quantizer() {} 44 | }; 45 | 46 | } // namespace faiss 47 | -------------------------------------------------------------------------------- /faiss/impl/index_read_utils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // Utils for index_read 9 | 10 | #ifndef FAISS_INDEX_READ_UTILS_H 11 | #define FAISS_INDEX_READ_UTILS_H 12 | 13 | #include 14 | #include 15 | 16 | #pragma once 17 | 18 | namespace faiss { 19 | struct ProductQuantizer; 20 | struct ScalarQuantizer; 21 | 22 | void read_index_header(Index* idx, IOReader* f); 23 | void read_direct_map(DirectMap* dm, IOReader* f); 24 | void read_ivf_header( 25 | IndexIVF* ivf, 26 | IOReader* f, 27 | std::vector>* ids = nullptr); 28 | void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags); 29 | ArrayInvertedLists* set_array_invlist( 30 | IndexIVF* ivf, 31 | std::vector>& ids); 32 | void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f); 33 | void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f); 34 | 35 | } // namespace faiss 36 | 37 | #endif 38 | -------------------------------------------------------------------------------- /faiss/impl/kmeans1d.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | /** SMAWK algorithm. Find the row minima of a monotone matrix. 16 | * 17 | * Expose this for testing. 18 | * 19 | * @param nrows number of rows 20 | * @param ncols number of columns 21 | * @param x input matrix, size (nrows, ncols) 22 | * @param argmins argmin of each row 23 | */ 24 | void smawk( 25 | const idx_t nrows, 26 | const idx_t ncols, 27 | const float* x, 28 | idx_t* argmins); 29 | 30 | /** Exact 1D K-Means by dynamic programming 31 | * 32 | * From "Fast Exact k-Means, k-Medians and Bregman Divergence Clustering in 1D" 33 | * Allan Grønlund, Kasper Green Larsen, Alexander Mathiasen, Jesper Sindahl 34 | * Nielsen, Stefan Schneider, Mingzhou Song, ArXiV'17 35 | * 36 | * Section 2.2 37 | * 38 | * https://arxiv.org/abs/1701.07204 39 | * 40 | * @param x input 1D array 41 | * @param n input array length 42 | * @param nclusters number of clusters 43 | * @param centroids output centroids, size nclusters 44 | * @return imbalancce factor 45 | */ 46 | double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids); 47 | 48 | } // namespace faiss 49 | -------------------------------------------------------------------------------- /faiss/impl/mapped_io.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | namespace faiss { 18 | 19 | // holds a memory-mapped region over a file 20 | struct MmappedFileMappingOwner : public MaybeOwnedVectorOwner { 21 | MmappedFileMappingOwner(const std::string& filename); 22 | MmappedFileMappingOwner(FILE* f); 23 | ~MmappedFileMappingOwner(); 24 | 25 | void* data() const; 26 | size_t size() const; 27 | 28 | struct PImpl; 29 | std::unique_ptr p_impl; 30 | }; 31 | 32 | // A deserializer that supports memory-mapped files. 33 | // All de-allocations should happen as soon as the index gets destroyed, 34 | // after all underlying the MaybeOwnerVector objects are destroyed. 35 | struct MappedFileIOReader : IOReader { 36 | std::shared_ptr mmap_owner; 37 | 38 | size_t pos = 0; 39 | 40 | MappedFileIOReader(const std::shared_ptr& owner); 41 | 42 | // perform a copy 43 | size_t operator()(void* ptr, size_t size, size_t nitems) override; 44 | // perform a quasi-read that returns a mmapped address, owned by mmap_owner, 45 | // and updates the position 46 | size_t mmap(void** ptr, size_t size, size_t nitems); 47 | 48 | int filedescriptor() override; 49 | }; 50 | 51 | } // namespace faiss 52 | -------------------------------------------------------------------------------- /faiss/impl/zerocopy_io.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | namespace faiss { 12 | 13 | ZeroCopyIOReader::ZeroCopyIOReader(uint8_t* data, size_t size) 14 | : data_(data), rp_(0), total_(size) {} 15 | 16 | ZeroCopyIOReader::~ZeroCopyIOReader() {} 17 | 18 | size_t ZeroCopyIOReader::get_data_view(void** ptr, size_t size, size_t nitems) { 19 | if (size == 0) { 20 | return nitems; 21 | } 22 | 23 | size_t actual_size = size * nitems; 24 | if (rp_ + size * nitems > total_) { 25 | actual_size = total_ - rp_; 26 | } 27 | 28 | size_t actual_nitems = (actual_size + size - 1) / size; 29 | if (actual_nitems == 0) { 30 | return 0; 31 | } 32 | 33 | // get an address 34 | *ptr = (void*)(reinterpret_cast(data_ + rp_)); 35 | 36 | // alter pos 37 | rp_ += size * actual_nitems; 38 | 39 | return actual_nitems; 40 | } 41 | 42 | void ZeroCopyIOReader::reset() { 43 | rp_ = 0; 44 | } 45 | 46 | size_t ZeroCopyIOReader::operator()(void* ptr, size_t size, size_t nitems) { 47 | if (size * nitems == 0) { 48 | return 0; 49 | } 50 | 51 | if (rp_ >= total_) { 52 | return 0; 53 | } 54 | size_t nremain = (total_ - rp_) / size; 55 | if (nremain < nitems) { 56 | nitems = nremain; 57 | } 58 | memcpy(ptr, (data_ + rp_), size * nitems); 59 | rp_ += size * nitems; 60 | return nitems; 61 | } 62 | 63 | int ZeroCopyIOReader::filedescriptor() { 64 | return -1; // Indicating no file descriptor available for memory buffer 65 | } 66 | 67 | } // namespace faiss 68 | -------------------------------------------------------------------------------- /faiss/impl/zerocopy_io.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | #include 13 | 14 | namespace faiss { 15 | 16 | // ZeroCopyIOReader just maps the data from a given pointer. 17 | struct ZeroCopyIOReader : public faiss::IOReader { 18 | uint8_t* data_; 19 | size_t rp_ = 0; 20 | size_t total_ = 0; 21 | 22 | ZeroCopyIOReader(uint8_t* data, size_t size); 23 | ~ZeroCopyIOReader(); 24 | 25 | void reset(); 26 | size_t get_data_view(void** ptr, size_t size, size_t nitems); 27 | size_t operator()(void* ptr, size_t size, size_t nitems) override; 28 | 29 | int filedescriptor() override; 30 | }; 31 | 32 | } // namespace faiss 33 | -------------------------------------------------------------------------------- /faiss/index_factory.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | /** Build and index with the sequence of processing steps described in 16 | * the string. */ 17 | Index* index_factory( 18 | int d, 19 | const char* description, 20 | MetricType metric = METRIC_L2, 21 | // Whether to maintain inverted list within faiss index (only applicable 22 | // to IndexIVF*) 23 | bool own_invlists = true); 24 | 25 | /// set to > 0 to get more logs from index_factory 26 | FAISS_API extern int index_factory_verbose; 27 | 28 | IndexBinary* index_binary_factory( 29 | int d, 30 | const char* description, 31 | bool own_invlists = true); 32 | 33 | } // namespace faiss 34 | -------------------------------------------------------------------------------- /faiss/utils/WorkerThread.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace faiss { 17 | 18 | class WorkerThread { 19 | public: 20 | WorkerThread(); 21 | 22 | /// Stops and waits for the worker thread to exit, flushing all 23 | /// pending lambdas 24 | ~WorkerThread(); 25 | 26 | /// Request that the worker thread stop itself 27 | void stop(); 28 | 29 | /// Blocking waits in the current thread for the worker thread to 30 | /// stop 31 | void waitForThreadExit(); 32 | 33 | /// Adds a lambda to run on the worker thread; returns a future that 34 | /// can be used to block on its completion. 35 | /// Future status is `true` if the lambda was run in the worker 36 | /// thread; `false` if it was not run, because the worker thread is 37 | /// exiting or has exited. 38 | std::future add(std::function f); 39 | 40 | private: 41 | void startThread(); 42 | void threadMain(); 43 | void threadLoop(); 44 | 45 | /// Thread that all queued lambdas are run on 46 | std::thread thread_; 47 | 48 | /// Mutex for the queue and exit status 49 | std::mutex mutex_; 50 | 51 | /// Monitor for the exit status and the queue 52 | std::condition_variable monitor_; 53 | 54 | /// Whether or not we want the thread to exit 55 | bool wantStop_; 56 | 57 | /// Queue of pending lambdas to call 58 | std::deque, std::promise>> queue_; 59 | }; 60 | 61 | } // namespace faiss 62 | -------------------------------------------------------------------------------- /faiss/utils/approx_topk/mode.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | /// Represents the mode of use of approximate top-k computations 11 | /// that allows to trade accuracy vs speed. So, every options 12 | /// besides EXACT_TOPK increases the speed. 13 | /// 14 | /// B represents the number of buckets. 15 | /// D is the number of min-k elements to track within every bucket. 16 | /// 17 | /// Default option is EXACT_TOPK. 18 | /// APPROX_TOPK_BUCKETS_B16_D2 is worth starting from, if you'd like 19 | /// to experiment a bit. 20 | /// 21 | /// It seems that only the limited number of combinations are 22 | /// meaningful, because of the limited supply of SIMD registers. 23 | /// Also, certain combinations, such as B32_D1 and B16_D1, were concluded 24 | /// to be not very precise in benchmarks, so ones were not introduced. 25 | /// 26 | /// TODO: Consider d-ary SIMD heap. 27 | 28 | enum ApproxTopK_mode_t : int { 29 | EXACT_TOPK = 0, 30 | APPROX_TOPK_BUCKETS_B32_D2 = 1, 31 | APPROX_TOPK_BUCKETS_B8_D3 = 2, 32 | APPROX_TOPK_BUCKETS_B16_D2 = 3, 33 | APPROX_TOPK_BUCKETS_B8_D2 = 4, 34 | }; 35 | -------------------------------------------------------------------------------- /faiss/utils/bf16.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | namespace faiss { 13 | 14 | namespace { 15 | 16 | union fp32_bits { 17 | uint32_t as_u32; 18 | float as_f32; 19 | }; 20 | 21 | } // namespace 22 | 23 | inline uint16_t encode_bf16(const float f) { 24 | // Round off 25 | fp32_bits fp; 26 | fp.as_f32 = f; 27 | return static_cast((fp.as_u32 + 0x8000) >> 16); 28 | } 29 | 30 | inline float decode_bf16(const uint16_t v) { 31 | fp32_bits fp; 32 | fp.as_u32 = (uint32_t(v) << 16); 33 | return fp.as_f32; 34 | } 35 | 36 | } // namespace faiss 37 | -------------------------------------------------------------------------------- /faiss/utils/distances_fused/avx512.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // AVX512 might be not used, but this version provides ~2x speedup 9 | // over AVX2 kernel, say, for training PQx10 or PQx12, and speeds up 10 | // additional cases with larger dimensionalities. 11 | 12 | #pragma once 13 | 14 | #include 15 | #include 16 | 17 | #include 18 | 19 | #ifdef __AVX512F__ 20 | 21 | namespace faiss { 22 | 23 | // Returns true if the fused kernel is available and the data was processed. 24 | // Returns false if the fused kernel is not available. 25 | bool exhaustive_L2sqr_fused_cmax_AVX512( 26 | const float* x, 27 | const float* y, 28 | size_t d, 29 | size_t nx, 30 | size_t ny, 31 | Top1BlockResultHandler>& res, 32 | const float* y_norms); 33 | 34 | } // namespace faiss 35 | 36 | #endif 37 | -------------------------------------------------------------------------------- /faiss/utils/distances_fused/distances_fused.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | 12 | #include 13 | #include 14 | 15 | namespace faiss { 16 | 17 | bool exhaustive_L2sqr_fused_cmax( 18 | const float* x, 19 | const float* y, 20 | size_t d, 21 | size_t nx, 22 | size_t ny, 23 | Top1BlockResultHandler>& res, 24 | const float* y_norms) { 25 | if (nx == 0 || ny == 0) { 26 | // nothing to do 27 | return true; 28 | } 29 | 30 | #ifdef __AVX512F__ 31 | // avx512 kernel 32 | return exhaustive_L2sqr_fused_cmax_AVX512(x, y, d, nx, ny, res, y_norms); 33 | #elif defined(__AVX2__) || defined(__aarch64__) 34 | // avx2 or arm neon kernel 35 | return exhaustive_L2sqr_fused_cmax_simdlib(x, y, d, nx, ny, res, y_norms); 36 | #else 37 | // not supported, please use a general-purpose kernel 38 | return false; 39 | #endif 40 | } 41 | 42 | } // namespace faiss 43 | -------------------------------------------------------------------------------- /faiss/utils/distances_fused/distances_fused.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | // This file contains a fused kernel that combines distance computation 9 | // and the search for the CLOSEST point. Currently, this is done for small 10 | // dimensionality vectors when it is beneficial to avoid storing temporary 11 | // dot product information in RAM. This is particularly effective 12 | // when training PQx10 or PQx12 with the default parameters. 13 | // 14 | // InterruptCallback::check() is not used, because it is assumed that the 15 | // kernel takes a little time because of a tiny dimensionality. 16 | // 17 | // Later on, similar optimization can be implemented for large size vectors, 18 | // but a different kernel is needed. 19 | // 20 | 21 | #pragma once 22 | 23 | #include 24 | 25 | #include 26 | 27 | namespace faiss { 28 | 29 | // Returns true if the fused kernel is available and the data was processed. 30 | // Returns false if the fused kernel is not available. 31 | bool exhaustive_L2sqr_fused_cmax( 32 | const float* x, 33 | const float* y, 34 | size_t d, 35 | size_t nx, 36 | size_t ny, 37 | Top1BlockResultHandler>& res, 38 | const float* y_norms); 39 | 40 | } // namespace faiss 41 | -------------------------------------------------------------------------------- /faiss/utils/distances_fused/simdlib_based.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | #if defined(__AVX2__) || defined(__aarch64__) 16 | 17 | namespace faiss { 18 | 19 | // Returns true if the fused kernel is available and the data was processed. 20 | // Returns false if the fused kernel is not available. 21 | bool exhaustive_L2sqr_fused_cmax_simdlib( 22 | const float* x, 23 | const float* y, 24 | size_t d, 25 | size_t nx, 26 | size_t ny, 27 | Top1BlockResultHandler>& res, 28 | const float* y_norms); 29 | 30 | } // namespace faiss 31 | 32 | #endif 33 | -------------------------------------------------------------------------------- /faiss/utils/extra_distances.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | /** In this file are the implementations of extra metrics beyond L2 11 | * and inner product */ 12 | 13 | #include 14 | 15 | #include 16 | 17 | #include 18 | 19 | namespace faiss { 20 | 21 | struct FlatCodesDistanceComputer; 22 | 23 | void pairwise_extra_distances( 24 | int64_t d, 25 | int64_t nq, 26 | const float* xq, 27 | int64_t nb, 28 | const float* xb, 29 | MetricType mt, 30 | float metric_arg, 31 | float* dis, 32 | int64_t ldq = -1, 33 | int64_t ldb = -1, 34 | int64_t ldd = -1); 35 | 36 | void knn_extra_metrics( 37 | const float* x, 38 | const float* y, 39 | size_t d, 40 | size_t nx, 41 | size_t ny, 42 | MetricType mt, 43 | float metric_arg, 44 | size_t k, 45 | float* distances, 46 | int64_t* indexes); 47 | 48 | /** get a DistanceComputer that refers to this type of distance and 49 | * indexes a flat array of size nb */ 50 | FlatCodesDistanceComputer* get_extra_distance_computer( 51 | size_t d, 52 | MetricType mt, 53 | float metric_arg, 54 | size_t nb, 55 | const float* xb); 56 | 57 | } // namespace faiss 58 | 59 | #include 60 | -------------------------------------------------------------------------------- /faiss/utils/fp16-arm.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | inline uint16_t encode_fp16(float x) { 16 | float32x4_t fx4 = vdupq_n_f32(x); 17 | float16x4_t f16x4 = vcvt_f16_f32(fx4); 18 | uint16x4_t ui16x4 = vreinterpret_u16_f16(f16x4); 19 | return vduph_lane_u16(ui16x4, 3); 20 | } 21 | 22 | inline float decode_fp16(uint16_t x) { 23 | uint16x4_t ui16x4 = vdup_n_u16(x); 24 | float16x4_t f16x4 = vreinterpret_f16_u16(ui16x4); 25 | float32x4_t fx4 = vcvt_f32_f16(f16x4); 26 | return vdups_laneq_f32(fx4, 3); 27 | } 28 | 29 | } // namespace faiss 30 | -------------------------------------------------------------------------------- /faiss/utils/fp16-fp16c.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | #include 12 | 13 | namespace faiss { 14 | 15 | inline uint16_t encode_fp16(float x) { 16 | __m128 xf = _mm_set1_ps(x); 17 | __m128i xi = 18 | _mm_cvtps_ph(xf, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); 19 | return _mm_cvtsi128_si32(xi) & 0xffff; 20 | } 21 | 22 | inline float decode_fp16(uint16_t x) { 23 | __m128i xi = _mm_set1_epi16(x); 24 | __m128 xf = _mm_cvtph_ps(xi); 25 | return _mm_cvtss_f32(xf); 26 | } 27 | 28 | } // namespace faiss 29 | -------------------------------------------------------------------------------- /faiss/utils/fp16.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | #include 11 | 12 | #include 13 | 14 | #if defined(__F16C__) 15 | #include 16 | #elif defined(__aarch64__) 17 | #include 18 | #else 19 | #include 20 | #endif 21 | -------------------------------------------------------------------------------- /faiss/utils/hamming_distance/common.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #ifndef FAISS_hamming_common_h 9 | #define FAISS_hamming_common_h 10 | 11 | #include 12 | 13 | #include 14 | 15 | /* The Hamming distance type */ 16 | using hamdis_t = int32_t; 17 | 18 | namespace faiss { 19 | 20 | // trust the compiler to provide efficient popcount implementations 21 | inline int popcount32(uint32_t x) { 22 | return __builtin_popcount(x); 23 | } 24 | 25 | // popcount 26 | inline int popcount64(uint64_t x) { 27 | return __builtin_popcountl(x); 28 | } 29 | 30 | // This table was moved from .cpp to .h file, because 31 | // otherwise it was causing compilation errors while trying to 32 | // compile swig modules on Windows. 33 | // todo for C++17: switch to 'inline constexpr' 34 | static constexpr uint8_t hamdis_tab_ham_bytes[256] = { 35 | 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4, 36 | 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 37 | 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4, 38 | 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 39 | 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 40 | 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 41 | 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, 42 | 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 43 | 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 44 | 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 45 | 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8}; 46 | 47 | } // namespace faiss 48 | 49 | #endif 50 | -------------------------------------------------------------------------------- /faiss/utils/prefetch.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | // prefetches 11 | 12 | #ifdef __AVX__ 13 | 14 | // AVX 15 | 16 | #include 17 | 18 | inline void prefetch_L1(const void* address) { 19 | _mm_prefetch((const char*)address, _MM_HINT_T0); 20 | } 21 | inline void prefetch_L2(const void* address) { 22 | _mm_prefetch((const char*)address, _MM_HINT_T1); 23 | } 24 | inline void prefetch_L3(const void* address) { 25 | _mm_prefetch((const char*)address, _MM_HINT_T2); 26 | } 27 | 28 | #elif defined(__aarch64__) 29 | 30 | // ARM64 31 | 32 | #ifdef _MSC_VER 33 | 34 | // todo: arm on MSVC 35 | inline void prefetch_L1(const void* address) {} 36 | inline void prefetch_L2(const void* address) {} 37 | inline void prefetch_L3(const void* address) {} 38 | 39 | #else 40 | // arm on non-MSVC 41 | 42 | inline void prefetch_L1(const void* address) { 43 | __builtin_prefetch(address, 0, 3); 44 | } 45 | inline void prefetch_L2(const void* address) { 46 | __builtin_prefetch(address, 0, 2); 47 | } 48 | inline void prefetch_L3(const void* address) { 49 | __builtin_prefetch(address, 0, 1); 50 | } 51 | #endif 52 | 53 | #else 54 | 55 | // a generic platform 56 | 57 | #ifdef _MSC_VER 58 | 59 | inline void prefetch_L1(const void* address) {} 60 | inline void prefetch_L2(const void* address) {} 61 | inline void prefetch_L3(const void* address) {} 62 | 63 | #else 64 | 65 | inline void prefetch_L1(const void* address) { 66 | __builtin_prefetch(address, 0, 3); 67 | } 68 | inline void prefetch_L2(const void* address) { 69 | __builtin_prefetch(address, 0, 2); 70 | } 71 | inline void prefetch_L3(const void* address) { 72 | __builtin_prefetch(address, 0, 1); 73 | } 74 | 75 | #endif 76 | 77 | #endif 78 | -------------------------------------------------------------------------------- /faiss/utils/simdlib.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | 10 | /** Abstractions for 256-bit registers 11 | * 12 | * The objective is to separate the different interpretations of the same 13 | * registers (as a vector of uint8, uint16 or uint32), to provide printing 14 | * functions. 15 | */ 16 | 17 | #if defined(__AVX512F__) 18 | 19 | #include 20 | #include 21 | 22 | #elif defined(__AVX2__) 23 | 24 | #include 25 | 26 | #elif defined(__aarch64__) 27 | 28 | #include 29 | 30 | #elif defined(__PPC64__) 31 | 32 | #include 33 | 34 | #else 35 | 36 | // emulated = all operations are implemented as scalars 37 | #include 38 | 39 | // FIXME: make a SSE version 40 | // is this ever going to happen? We will probably rather implement AVX512 41 | 42 | #endif 43 | -------------------------------------------------------------------------------- /perf_tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | # @lint-ignore-every LINEWRAP 7 | project(faiss_perf_tests) 8 | set(BENCHMARK_ENABLE_TESTING OFF) 9 | 10 | include(FetchContent) 11 | FetchContent_Declare(googlebenchmark 12 | GIT_REPOSITORY https://github.com/google/benchmark.git 13 | GIT_TAG main) # need main for benchmark::benchmark 14 | FetchContent_MakeAvailable( 15 | googlebenchmark) 16 | 17 | 18 | find_package(Threads REQUIRED) 19 | find_package(OpenMP REQUIRED) 20 | find_package(gflags REQUIRED) 21 | 22 | add_library(faiss_perf_tests_utils 23 | utils.cpp 24 | ) 25 | # `#include ` or any other headers 26 | target_include_directories(faiss_perf_tests_utils PRIVATE 27 | ${PROJECT_SOURCE_DIR}/../..) 28 | 29 | include(../cmake/link_to_faiss_lib.cmake) 30 | 31 | link_to_faiss_lib(faiss_perf_tests_utils) 32 | 33 | set(FAISS_PERF_TEST_SRC 34 | bench_no_multithreading_rcq_search.cpp 35 | bench_scalar_quantizer_accuracy.cpp 36 | bench_scalar_quantizer_decode.cpp 37 | bench_scalar_quantizer_distance.cpp 38 | bench_scalar_quantizer_encode.cpp 39 | ) 40 | foreach(bench ${FAISS_PERF_TEST_SRC}) 41 | get_filename_component(bench_exec ${bench} NAME_WE) 42 | add_executable(${bench_exec} ${bench}) 43 | link_to_faiss_lib(${bench_exec}) 44 | target_link_libraries(${bench_exec} PRIVATE faiss_perf_tests_utils OpenMP::OpenMP_CXX benchmark::benchmark gflags) 45 | # `#include ` or any other headers 46 | target_include_directories(${bench_exec} PRIVATE 47 | ${PROJECT_SOURCE_DIR}/../..) 48 | 49 | endforeach() 50 | -------------------------------------------------------------------------------- /perf_tests/bench_scalar_quantizer_encode.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | 19 | using namespace faiss; 20 | DEFINE_uint32(d, 128, "dimension"); 21 | DEFINE_uint32(n, 2000, "dimension"); 22 | DEFINE_uint32(iterations, 20, "iterations"); 23 | 24 | static void bench_encode( 25 | benchmark::State& state, 26 | ScalarQuantizer::QuantizerType type, 27 | int d, 28 | int n) { 29 | std::vector x(d * n); 30 | 31 | float_rand(x.data(), d * n, 12345); 32 | ScalarQuantizer sq(d, type); 33 | 34 | omp_set_num_threads(1); 35 | size_t code_size = sq.code_size; 36 | 37 | sq.train(n, x.data()); 38 | state.counters["code_size"] = sq.code_size; 39 | std::vector codes(code_size * n); 40 | 41 | for (auto _ : state) { 42 | // encode 43 | sq.compute_codes(x.data(), codes.data(), n); 44 | } 45 | } 46 | 47 | int main(int argc, char** argv) { 48 | benchmark::Initialize(&argc, argv); 49 | gflags::AllowCommandLineReparsing(); 50 | gflags::ParseCommandLineFlags(&argc, &argv, true); 51 | int iterations = FLAGS_iterations; 52 | int d = FLAGS_d; 53 | int n = FLAGS_n; 54 | auto benchs = ::perf_tests::sq_types(); 55 | 56 | for (auto& [bench_name, quantizer_type] : benchs) { 57 | benchmark::RegisterBenchmark( 58 | bench_name.c_str(), bench_encode, quantizer_type, d, n) 59 | ->Iterations(iterations); 60 | } 61 | 62 | benchmark::RunSpecifiedBenchmarks(); 63 | benchmark::Shutdown(); 64 | } 65 | -------------------------------------------------------------------------------- /perf_tests/utils.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | namespace faiss::perf_tests { 10 | std::map sq_types() { 11 | static std::map 12 | sq_types = { 13 | {"QT_8bit", faiss::ScalarQuantizer::QT_8bit}, 14 | {"QT_4bit", faiss::ScalarQuantizer::QT_4bit}, 15 | {"QT_8bit_uniform", 16 | faiss::ScalarQuantizer::QT_8bit_uniform}, 17 | {"QT_4bit_uniform", 18 | faiss::ScalarQuantizer::QT_4bit_uniform}, 19 | {"QT_fp16", faiss::ScalarQuantizer::QT_fp16}, 20 | {"QT_8bit_direct", faiss::ScalarQuantizer::QT_8bit_direct}, 21 | {"QT_6bit", faiss::ScalarQuantizer::QT_6bit}, 22 | {"QT_bf16", faiss::ScalarQuantizer::QT_bf16}, 23 | {"QT_8bit_direct_signed", 24 | faiss::ScalarQuantizer::QT_8bit_direct_signed}}; 25 | return sq_types; 26 | } 27 | } // namespace faiss::perf_tests 28 | -------------------------------------------------------------------------------- /perf_tests/utils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #pragma once 9 | #include 10 | #include 11 | 12 | namespace faiss::perf_tests { 13 | 14 | std::map sq_types(); 15 | 16 | } // namespace faiss::perf_tests 17 | -------------------------------------------------------------------------------- /tests/test_autotune.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | 7 | import unittest 8 | import faiss 9 | 10 | 11 | class TestParameterSpace(unittest.TestCase): 12 | 13 | def test_nprobe(self): 14 | index = faiss.index_factory(32, "IVF32,Flat") 15 | ps = faiss.ParameterSpace() 16 | ps.set_index_parameter(index, "nprobe", 5) 17 | self.assertEqual(index.nprobe, 5) 18 | 19 | def test_nprobe_2(self): 20 | index = faiss.index_factory(32, "IDMap,IVF32,Flat") 21 | ps = faiss.ParameterSpace() 22 | ps.set_index_parameter(index, "nprobe", 5) 23 | index2 = faiss.downcast_index(index.index) 24 | self.assertEqual(index2.nprobe, 5) 25 | 26 | def test_nprobe_3(self): 27 | index = faiss.index_factory(32, "IVF32,SQ8,RFlat") 28 | ps = faiss.ParameterSpace() 29 | ps.set_index_parameter(index, "nprobe", 5) 30 | index2 = faiss.downcast_index(index.base_index) 31 | self.assertEqual(index2.nprobe, 5) 32 | 33 | def test_nprobe_4(self): 34 | index = faiss.index_factory(32, "PCAR32,IVF32,SQ8,RFlat") 35 | ps = faiss.ParameterSpace() 36 | 37 | ps.set_index_parameter(index, "nprobe", 5) 38 | index2 = faiss.downcast_index(index.base_index) 39 | index2 = faiss.downcast_index(index2.index) 40 | self.assertEqual(index2.nprobe, 5) 41 | 42 | def test_efSearch(self): 43 | index = faiss.index_factory(32, "IVF32_HNSW32,SQ8") 44 | ps = faiss.ParameterSpace() 45 | ps.set_index_parameter(index, "quantizer_efSearch", 5) 46 | index2 = faiss.downcast_index(index.quantizer) 47 | self.assertEqual(index2.hnsw.efSearch, 5) 48 | -------------------------------------------------------------------------------- /tests/test_binary_factory.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from __future__ import absolute_import, division, print_function 7 | 8 | import unittest 9 | import faiss 10 | 11 | 12 | class TestBinaryFactory(unittest.TestCase): 13 | 14 | def test_factory_IVF(self): 15 | 16 | index = faiss.index_binary_factory(16, "BIVF10") 17 | assert index.invlists is not None 18 | assert index.nlist == 10 19 | assert index.code_size == 2 20 | 21 | def test_factory_Flat(self): 22 | 23 | index = faiss.index_binary_factory(16, "BFlat") 24 | assert index.code_size == 2 25 | 26 | def test_factory_HNSW(self): 27 | 28 | index = faiss.index_binary_factory(256, "BHNSW32") 29 | assert index.code_size == 32 30 | 31 | def test_factory_IVF_HNSW(self): 32 | 33 | index = faiss.index_binary_factory(256, "BIVF1024_BHNSW32") 34 | assert index.code_size == 32 35 | assert index.nlist == 1024 36 | 37 | def test_factory_Hash(self): 38 | index = faiss.index_binary_factory(256, "BHash12") 39 | assert index.b == 12 40 | 41 | def test_factory_MultiHash(self): 42 | index = faiss.index_binary_factory(256, "BHash5x6") 43 | assert index.b == 6 44 | assert index.nhash == 5 45 | -------------------------------------------------------------------------------- /tests/test_binary_flat.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | TEST(BinaryFlat, accuracy) { 17 | // dimension of the vectors to index 18 | int d = 64; 19 | 20 | // size of the database we plan to index 21 | size_t nb = 1000; 22 | 23 | // make the index object and train it 24 | faiss::IndexBinaryFlat index(d); 25 | 26 | std::vector database(nb * (d / 8)); 27 | for (size_t i = 0; i < nb * (d / 8); i++) { 28 | database[i] = rand() % 0x100; 29 | } 30 | 31 | { // populating the database 32 | index.add(nb, database.data()); 33 | } 34 | 35 | size_t nq = 200; 36 | 37 | { // searching the database 38 | 39 | std::vector queries(nq * (d / 8)); 40 | for (size_t i = 0; i < nq * (d / 8); i++) { 41 | queries[i] = rand() % 0x100; 42 | } 43 | 44 | int k = 5; 45 | std::vector nns(k * nq); 46 | std::vector dis(k * nq); 47 | 48 | index.search(nq, queries.data(), k, dis.data(), nns.data()); 49 | 50 | for (size_t i = 0; i < nq; ++i) { 51 | faiss::HammingComputer8 hc(queries.data() + i * (d / 8), d / 8); 52 | hamdis_t dist_min = hc.hamming(database.data()); 53 | for (size_t j = 1; j < nb; ++j) { 54 | hamdis_t dist = hc.hamming(database.data() + j * (d / 8)); 55 | if (dist < dist_min) { 56 | dist_min = dist; 57 | } 58 | } 59 | EXPECT_EQ(dist_min, dis[k * i]); 60 | } 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /tests/test_callback.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | TEST(TestCallback, timeout) { 17 | int n = 1000; 18 | int k = 100; 19 | int d = 128; 20 | int niter = 1000000000; 21 | int seed = 42; 22 | 23 | std::vector vecs(n * d); 24 | faiss::float_rand(vecs.data(), vecs.size(), seed); 25 | 26 | auto index(new faiss::IndexFlat(d)); 27 | 28 | faiss::ClusteringParameters cp; 29 | cp.niter = niter; 30 | cp.verbose = false; 31 | 32 | faiss::Clustering kmeans(d, k, cp); 33 | 34 | faiss::TimeoutCallback::reset(0.010); 35 | EXPECT_THROW(kmeans.train(n, vecs.data(), *index), faiss::FaissException); 36 | delete index; 37 | } 38 | -------------------------------------------------------------------------------- /tests/test_callback_py.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import unittest 7 | import numpy as np 8 | import faiss 9 | 10 | 11 | class TestCallbackPy(unittest.TestCase): 12 | def setUp(self) -> None: 13 | super().setUp() 14 | 15 | def test_timeout(self) -> None: 16 | n = 1000 17 | k = 100 18 | d = 128 19 | niter = 1_000_000_000 20 | 21 | x = np.random.rand(n, d).astype('float32') 22 | index = faiss.IndexFlat(d) 23 | 24 | cp = faiss.ClusteringParameters() 25 | cp.niter = niter 26 | cp.verbose = False 27 | 28 | kmeans = faiss.Clustering(d, k, cp) 29 | 30 | with self.assertRaises(RuntimeError): 31 | with faiss.TimeoutGuard(0.010): 32 | kmeans.train(x, index) 33 | -------------------------------------------------------------------------------- /tests/test_documentation.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import unittest 7 | import faiss 8 | 9 | 10 | class TestDocumentation(unittest.TestCase): 11 | 12 | def test_io_error(self): 13 | index = faiss.IndexFlatL2(32) 14 | 15 | self.assertTrue("Adds vectors to the index" in index.add.__doc__) 16 | -------------------------------------------------------------------------------- /tests/test_factory_tools.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | 12 | namespace faiss { 13 | 14 | TEST(TestFactoryTools, TestReverseIndexFactory) { 15 | for (const char* factory : { 16 | "Flat", 17 | "IMI2x5,PQ8x8", 18 | "IVF32_HNSW32,SQ8", 19 | "IVF8,Flat", 20 | "IVF8,SQ4", 21 | "IVF8,PQ4x8", 22 | "LSHrt", 23 | "PQ4x8", 24 | "HNSW32", 25 | "SQ8", 26 | "SQfp16", 27 | "NSG24,Flat", 28 | "NSG16,SQ8", 29 | }) { 30 | std::unique_ptr index{index_factory(64, factory)}; 31 | ASSERT_TRUE(index); 32 | EXPECT_EQ(factory, reverse_index_factory(index.get())); 33 | } 34 | using Case = std::pair; 35 | for (auto [src, dst] : { 36 | Case{"SQ8,RFlat", "SQ8,Refine(Flat)"}, 37 | Case{"NSG", "NSG32,Flat"}, 38 | Case{"NSG,PQ8", "NSG32,PQ8x8"}, 39 | }) { 40 | std::unique_ptr index{index_factory(64, src)}; 41 | ASSERT_TRUE(index); 42 | EXPECT_EQ(dst, reverse_index_factory(index.get())); 43 | } 44 | } 45 | 46 | } // namespace faiss 47 | -------------------------------------------------------------------------------- /tests/test_heap.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | using namespace faiss; 14 | 15 | TEST(Heap, addn_with_ids) { 16 | size_t n = 1000; 17 | size_t k = 1; 18 | std::vector heap_labels(n, -1); 19 | std::vector heap_distances(n, 0); 20 | float_minheap_array_t heaps = { 21 | n, k, heap_labels.data(), heap_distances.data()}; 22 | heaps.heapify(); 23 | std::vector labels(n, 1); 24 | std::vector distances(n, 0.0f); 25 | std::vector subset(n); 26 | std::iota(subset.begin(), subset.end(), 0); 27 | heaps.addn_with_ids(1, distances.data(), labels.data(), 1); 28 | heaps.reorder(); 29 | EXPECT_TRUE( 30 | std::all_of(heap_labels.begin(), heap_labels.end(), [](int64_t i) { 31 | return i == 1; 32 | })); 33 | } 34 | 35 | TEST(Heap, addn_query_subset_with_ids) { 36 | size_t n = 20000000; // more than 2^24 37 | size_t k = 1; 38 | std::vector heap_labels(n, -1); 39 | std::vector heap_distances(n, 0); 40 | float_minheap_array_t heaps = { 41 | n, k, heap_labels.data(), heap_distances.data()}; 42 | heaps.heapify(); 43 | std::vector labels(n, 1); 44 | std::vector distances(n, 0.0f); 45 | std::vector subset(n); 46 | std::iota(subset.begin(), subset.end(), 0); 47 | heaps.addn_query_subset_with_ids( 48 | n, subset.data(), 1, distances.data(), labels.data(), 1); 49 | heaps.reorder(); 50 | EXPECT_TRUE( 51 | std::all_of(heap_labels.begin(), heap_labels.end(), [](int64_t i) { 52 | return i == 1; 53 | })); 54 | } 55 | -------------------------------------------------------------------------------- /tests/test_omp_threads.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | 12 | TEST(Threading, openmp) { 13 | EXPECT_TRUE(faiss::check_openmp()); 14 | } 15 | -------------------------------------------------------------------------------- /tests/test_omp_threads_py.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from __future__ import absolute_import, division, print_function, unicode_literals 7 | 8 | import faiss 9 | import unittest 10 | 11 | 12 | class TestOpenMP(unittest.TestCase): 13 | 14 | def test_openmp(self): 15 | assert faiss.check_openmp() 16 | -------------------------------------------------------------------------------- /tests/test_partitioning.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | #include 12 | 13 | using namespace faiss; 14 | 15 | typedef AlignedTable AlignedTableUint16; 16 | 17 | // TODO: This test fails when Faiss is compiled with 18 | // GCC 13.2 from conda-forge with AVX2 enabled. This may be 19 | // a GCC bug that needs to be investigated further. 20 | // As of 16-AUG-2023 the Faiss conda packages are built 21 | // with GCC 11.2, so the published binaries are not affected. 22 | TEST(TestPartitioning, TestPartitioningBigRange) { 23 | auto n = 1024; 24 | AlignedTableUint16 tab(n); 25 | for (auto i = 0; i < n; i++) { 26 | tab[i] = i * 64; 27 | } 28 | int32_t hist[16]{}; 29 | simd_histogram_16(tab.get(), n, 0, 12, hist); 30 | for (auto i = 0; i < 16; i++) { 31 | ASSERT_EQ(hist[i], 64); 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /tests/test_rowwise_minmax.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | 8 | import faiss 9 | import unittest 10 | 11 | from common_faiss_tests import get_dataset_2 12 | 13 | 14 | class TestIndexRowwiseMinmax(unittest.TestCase): 15 | def compare_train_vs_train_inplace(self, factory_key): 16 | d = 96 17 | nb = 1000 18 | nq = 0 19 | nt = 2000 20 | 21 | xt, x, _ = get_dataset_2(d, nt, nb, nq) 22 | 23 | assert x.size > 0 24 | 25 | codec = faiss.index_factory(d, factory_key) 26 | 27 | # use the regular .train() 28 | codec.train(xt) 29 | codes_train = codec.sa_encode(x) 30 | 31 | decoded = codec.sa_decode(codes_train) 32 | 33 | # use .train_inplace() 34 | xt_cloned = np.copy(xt) 35 | codec.train_inplace(xt_cloned) 36 | codes_train_inplace = codec.sa_encode(x) 37 | 38 | # compare .train and .train_inplace codes 39 | n_diff = (codes_train != codes_train_inplace).sum() 40 | self.assertEqual(n_diff, 0) 41 | 42 | # make sure that the array used for .train_inplace got affected 43 | n_diff_xt = (xt_cloned != xt).sum() 44 | self.assertNotEqual(n_diff_xt, 0) 45 | 46 | # make sure that the reconstruction error is not crazy 47 | reconstruction_err = ((x - decoded) ** 2).sum() 48 | 49 | self.assertLess(reconstruction_err, 0.6) 50 | 51 | def test_fp32(self) -> None: 52 | self.compare_train_vs_train_inplace("MinMax,SQ8") 53 | 54 | def test_fp16(self) -> None: 55 | self.compare_train_vs_train_inplace("MinMaxFP16,SQ8") 56 | -------------------------------------------------------------------------------- /tests/test_util.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #ifndef FAISS_TEST_UTIL_H 9 | #define FAISS_TEST_UTIL_H 10 | 11 | #include 12 | #include 13 | 14 | struct Tempfilename { 15 | pthread_mutex_t* mutex; 16 | std::string filename; 17 | 18 | Tempfilename(pthread_mutex_t* mutex, std::string filename_template) { 19 | this->mutex = mutex; 20 | this->filename = filename_template; 21 | pthread_mutex_lock(mutex); 22 | int fd = mkstemp(&this->filename[0]); 23 | close(fd); 24 | pthread_mutex_unlock(mutex); 25 | } 26 | 27 | ~Tempfilename() { 28 | if (access(filename.c_str(), F_OK)) { 29 | unlink(filename.c_str()); 30 | } 31 | } 32 | 33 | const char* c_str() { 34 | return filename.c_str(); 35 | } 36 | }; 37 | 38 | #endif // FAISS_TEST_UTIL_H 39 | -------------------------------------------------------------------------------- /tests/test_utils.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | 10 | #include 11 | #include 12 | 13 | TEST(TestUtils, get_version) { 14 | std::string version = std::to_string(FAISS_VERSION_MAJOR) + "." + 15 | std::to_string(FAISS_VERSION_MINOR) + "." + 16 | std::to_string(FAISS_VERSION_PATCH); 17 | 18 | EXPECT_EQ(version, faiss::get_version()); 19 | } 20 | -------------------------------------------------------------------------------- /tutorial/cpp/6-HNSW.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. 3 | * 4 | * This source code is licensed under the MIT license found in the 5 | * LICENSE file in the root directory of this source tree. 6 | */ 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | using idx_t = faiss::idx_t; 16 | 17 | int main() { 18 | int d = 64; // dimension 19 | int nb = 100000; // database size 20 | int nq = 10000; // nb of queries 21 | 22 | std::mt19937 rng; 23 | std::uniform_real_distribution<> distrib; 24 | 25 | float* xb = new float[d * nb]; 26 | float* xq = new float[d * nq]; 27 | 28 | for (int i = 0; i < nb; i++) { 29 | for (int j = 0; j < d; j++) 30 | xb[d * i + j] = distrib(rng); 31 | xb[d * i] += i / 1000.; 32 | } 33 | 34 | for (int i = 0; i < nq; i++) { 35 | for (int j = 0; j < d; j++) 36 | xq[d * i + j] = distrib(rng); 37 | xq[d * i] += i / 1000.; 38 | } 39 | 40 | int k = 4; 41 | 42 | faiss::IndexHNSWFlat index(d, 32); 43 | index.add(nb, xb); 44 | 45 | { // search xq 46 | idx_t* I = new idx_t[k * nq]; 47 | float* D = new float[k * nq]; 48 | 49 | index.search(nq, xq, k, D, I); 50 | 51 | printf("I=\n"); 52 | for (int i = nq - 5; i < nq; i++) { 53 | for (int j = 0; j < k; j++) 54 | printf("%5zd ", I[i * k + j]); 55 | printf("\n"); 56 | } 57 | 58 | printf("D=\n"); 59 | for (int i = nq - 5; i < nq; i++) { 60 | for (int j = 0; j < k; j++) 61 | printf("%5f ", D[i * k + j]); 62 | printf("\n"); 63 | } 64 | 65 | delete[] I; 66 | delete[] D; 67 | } 68 | 69 | delete[] xb; 70 | delete[] xq; 71 | 72 | return 0; 73 | } 74 | -------------------------------------------------------------------------------- /tutorial/cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | add_executable(1-Flat EXCLUDE_FROM_ALL 1-Flat.cpp) 7 | target_link_libraries(1-Flat PRIVATE faiss) 8 | 9 | add_executable(2-IVFFlat EXCLUDE_FROM_ALL 2-IVFFlat.cpp) 10 | target_link_libraries(2-IVFFlat PRIVATE faiss) 11 | 12 | add_executable(3-IVFPQ EXCLUDE_FROM_ALL 3-IVFPQ.cpp) 13 | target_link_libraries(3-IVFPQ PRIVATE faiss) 14 | 15 | add_executable(4-GPU EXCLUDE_FROM_ALL 4-GPU.cpp) 16 | target_link_libraries(4-GPU PRIVATE faiss) 17 | 18 | add_executable(5-Multiple-GPUs EXCLUDE_FROM_ALL 5-Multiple-GPUs.cpp) 19 | target_link_libraries(5-Multiple-GPUs PRIVATE faiss) 20 | 21 | add_executable(6-HNSW EXCLUDE_FROM_ALL 6-HNSW.cpp) 22 | target_link_libraries(6-HNSW PRIVATE faiss) 23 | 24 | add_executable(7-PQFastScan EXCLUDE_FROM_ALL 7-PQFastScan.cpp) 25 | target_link_libraries(7-PQFastScan PRIVATE faiss) 26 | 27 | add_executable(8-PQFastScanRefine EXCLUDE_FROM_ALL 8-PQFastScanRefine.cpp) 28 | target_link_libraries(8-PQFastScanRefine PRIVATE faiss) 29 | 30 | add_executable(9-RefineComparison EXCLUDE_FROM_ALL 9-RefineComparison.cpp) 31 | target_link_libraries(9-RefineComparison PRIVATE faiss) 32 | -------------------------------------------------------------------------------- /tutorial/python/1-Flat.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | 8 | d = 64 # dimension 9 | nb = 100000 # database size 10 | nq = 10000 # nb of queries 11 | np.random.seed(1234) # make reproducible 12 | xb = np.random.random((nb, d)).astype('float32') 13 | xb[:, 0] += np.arange(nb) / 1000. 14 | xq = np.random.random((nq, d)).astype('float32') 15 | xq[:, 0] += np.arange(nq) / 1000. 16 | 17 | import faiss # make faiss available 18 | index = faiss.IndexFlatL2(d) # build the index 19 | print(index.is_trained) 20 | index.add(xb) # add vectors to the index 21 | print(index.ntotal) 22 | 23 | k = 4 # we want to see 4 nearest neighbors 24 | D, I = index.search(xb[:5], k) # sanity check 25 | print(I) 26 | print(D) 27 | D, I = index.search(xq, k) # actual search 28 | print(I[:5]) # neighbors of the 5 first queries 29 | print(I[-5:]) # neighbors of the 5 last queries 30 | -------------------------------------------------------------------------------- /tutorial/python/2-IVFFlat.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | 8 | d = 64 # dimension 9 | nb = 100000 # database size 10 | nq = 10000 # nb of queries 11 | np.random.seed(1234) # make reproducible 12 | xb = np.random.random((nb, d)).astype('float32') 13 | xb[:, 0] += np.arange(nb) / 1000. 14 | xq = np.random.random((nq, d)).astype('float32') 15 | xq[:, 0] += np.arange(nq) / 1000. 16 | 17 | import faiss 18 | 19 | nlist = 100 20 | k = 4 21 | quantizer = faiss.IndexFlatL2(d) # the other index 22 | index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2) 23 | # here we specify METRIC_L2, by default it performs inner-product search 24 | 25 | assert not index.is_trained 26 | index.train(xb) 27 | assert index.is_trained 28 | 29 | index.add(xb) # add may be a bit slower as well 30 | D, I = index.search(xq, k) # actual search 31 | print(I[-5:]) # neighbors of the 5 last queries 32 | index.nprobe = 10 # default nprobe is 1, try a few more 33 | D, I = index.search(xq, k) 34 | print(I[-5:]) # neighbors of the 5 last queries 35 | -------------------------------------------------------------------------------- /tutorial/python/3-IVFPQ.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | 8 | d = 64 # dimension 9 | nb = 100000 # database size 10 | nq = 10000 # nb of queries 11 | np.random.seed(1234) # make reproducible 12 | xb = np.random.random((nb, d)).astype('float32') 13 | xb[:, 0] += np.arange(nb) / 1000. 14 | xq = np.random.random((nq, d)).astype('float32') 15 | xq[:, 0] += np.arange(nq) / 1000. 16 | 17 | import faiss 18 | 19 | nlist = 100 20 | m = 8 21 | k = 4 22 | quantizer = faiss.IndexFlatL2(d) # this remains the same 23 | index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) 24 | # 8 specifies that each sub-vector is encoded as 8 bits 25 | index.train(xb) 26 | index.add(xb) 27 | D, I = index.search(xb[:5], k) # sanity check 28 | print(I) 29 | print(D) 30 | index.nprobe = 10 # make comparable with experiment above 31 | D, I = index.search(xq, k) # search 32 | print(I[-5:]) 33 | -------------------------------------------------------------------------------- /tutorial/python/5-Multiple-GPUs.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import numpy as np 7 | 8 | d = 64 # dimension 9 | nb = 100000 # database size 10 | nq = 10000 # nb of queries 11 | np.random.seed(1234) # make reproducible 12 | xb = np.random.random((nb, d)).astype('float32') 13 | xb[:, 0] += np.arange(nb) / 1000. 14 | xq = np.random.random((nq, d)).astype('float32') 15 | xq[:, 0] += np.arange(nq) / 1000. 16 | 17 | import faiss # make faiss available 18 | 19 | ngpus = faiss.get_num_gpus() 20 | 21 | print("number of GPUs:", ngpus) 22 | 23 | cpu_index = faiss.IndexFlatL2(d) 24 | 25 | gpu_index = faiss.index_cpu_to_all_gpus( # build the index 26 | cpu_index 27 | ) 28 | 29 | gpu_index.add(xb) # add vectors to the index 30 | print(gpu_index.ntotal) 31 | 32 | k = 4 # we want to see 4 nearest neighbors 33 | D, I = gpu_index.search(xq, k) # actual search 34 | print(I[:5]) # neighbors of the 5 first queries 35 | print(I[-5:]) # neighbors of the 5 last queries 36 | -------------------------------------------------------------------------------- /tutorial/python/7-PQFastScan.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import faiss 7 | import numpy as np 8 | 9 | d = 64 # dimension 10 | nb = 100000 # database size 11 | nq = 10000 # nb of queries 12 | np.random.seed(1234) # make reproducible 13 | xb = np.random.random((nb, d)).astype('float32') # 64-dim *nb queries 14 | xb[:, 0] += np.arange(nb) / 1000. 15 | xq = np.random.random((nq, d)).astype('float32') 16 | xq[:, 0] += np.arange(nq) / 1000. 17 | 18 | m = 8 # 8 specifies that the number of sub-vector is 8 19 | k = 4 # number of dimension in etracted vector 20 | n_bit = 4 # 4 specifies that each sub-vector is encoded as 4 bits 21 | bbs = 32 # build block size ( bbs % 32 == 0 ) for PQ 22 | index = faiss.IndexPQFastScan(d, m, n_bit, faiss.METRIC_L2, bbs) 23 | # construct FastScan Index 24 | 25 | assert not index.is_trained 26 | index.train(xb) # Train vectors data index within mockup database 27 | assert index.is_trained 28 | 29 | index.add(xb) 30 | D, I = index.search(xb[:5], k) # sanity check 31 | print(I) 32 | print(D) 33 | index.nprobe = 10 # make comparable with experiment above 34 | D, I = index.search(xq, k) # search 35 | print(I[-5:]) # neighbors of the 5 last queries 36 | -------------------------------------------------------------------------------- /tutorial/python/8-PQFastScanRefine.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import faiss 7 | import numpy as np 8 | 9 | d = 64 # dimension 10 | nb = 100000 # database size 11 | nq = 10000 # nb of queries 12 | np.random.seed(1234) # make reproducible 13 | xb = np.random.random((nb, d)).astype('float32') # 64-dim *nb queries 14 | xb[:, 0] += np.arange(nb) / 1000. 15 | xq = np.random.random((nq, d)).astype('float32') 16 | xq[:, 0] += np.arange(nq) / 1000. 17 | 18 | m = 8 # 8 specifies that the number of sub-vector is 8 19 | k = 4 # number of dimension in etracted vector 20 | n_bit = 4 # 4 specifies that each sub-vector is encoded as 4 bits 21 | bbs = 32 # build block size ( bbs % 32 == 0 ) for PQ 22 | 23 | index = faiss.IndexPQFastScan(d, m, n_bit, faiss.METRIC_L2) 24 | index_refine = faiss.IndexRefineFlat(index) 25 | # construct FastScan and run index refinement 26 | 27 | assert not index_refine.is_trained 28 | index_refine.train(xb) # Train vectors data index within mockup database 29 | assert index_refine.is_trained 30 | 31 | index_refine.add(xb) 32 | params = faiss.IndexRefineSearchParameters(k_factor=3) 33 | D, I = index_refine.search(xq[:5], 10, params=params) 34 | print(I) 35 | print(D) 36 | index.nprobe = 10 # make comparable with experiment above 37 | D, I = index.search(xq[:5], k) # search 38 | print(I[-5:]) 39 | -------------------------------------------------------------------------------- /tutorial/python/9-RefineComparison.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import faiss 7 | 8 | from faiss.contrib.evaluation import knn_intersection_measure 9 | from faiss.contrib import datasets 10 | 11 | # 64-dim vectors, 50000 vectors in the training, 100000 in database, 12 | # 10000 in queries, dtype ('float32') 13 | ds = datasets.SyntheticDataset(64, 50000, 100000, 10000) 14 | d = 64 # dimension 15 | 16 | # Constructing the refine PQ index with SQfp16 with index factory 17 | index_fp16 = faiss.index_factory(d, 'PQ32x4fs,Refine(SQfp16)') 18 | index_fp16.train(ds.get_train()) 19 | index_fp16.add(ds.get_database()) 20 | 21 | # Constructing the refine PQ index with SQ8 22 | index_sq8 = faiss.index_factory(d, 'PQ32x4fs,Refine(SQ8)') 23 | index_sq8.train(ds.get_train()) 24 | index_sq8.add(ds.get_database()) 25 | 26 | # Parameterization on k factor while doing search for index refinement 27 | k_factor = 3.0 28 | params = faiss.IndexRefineSearchParameters(k_factor=k_factor) 29 | 30 | # Perform index search using different index refinement 31 | D_fp16, I_fp16 = index_fp16.search(ds.get_queries(), 100, params=params) 32 | D_sq8, I_sq8 = index_sq8.search(ds.get_queries(), 100, params=params) 33 | 34 | # Calculating knn intersection measure for different index types on refinement 35 | KIM_fp16 = knn_intersection_measure(I_fp16, ds.get_groundtruth()) 36 | KIM_sq8 = knn_intersection_measure(I_sq8, ds.get_groundtruth()) 37 | 38 | # KNN intersection measure accuracy shows that choosing SQ8 impacts accuracy 39 | assert (KIM_fp16 > KIM_sq8) 40 | 41 | print(I_sq8[:5]) 42 | print(I_fp16[:5]) 43 | --------------------------------------------------------------------------------