merge upstream

facebookresearch · Nov 21, 2024 · 096a826 · 096a826
2 parents 3201ee6 + eaab46c
commit 096a826
Show file tree

Hide file tree

Showing 78 changed files with 928 additions and 173 deletions.
diff --git a/benchs/bench_cppcontrib_sa_decode.cpp b/benchs/bench_cppcontrib_sa_decode.cpp
@@ -20,9 +20,7 @@
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexPQ.h>
-#include <faiss/impl/io.h>
 #include <faiss/index_factory.h>
-#include <faiss/index_io.h>
 
 #include <faiss/IndexRowwiseMinMax.h>
 #include <faiss/cppcontrib/SaDecodeKernels.h>

diff --git a/benchs/bench_fw/descriptors.py b/benchs/bench_fw/descriptors.py
@@ -214,6 +214,7 @@ class CodecDescriptor(IndexBaseDescriptor):
     factory: Optional[str] = None
     construction_params: Optional[List[Dict[str, int]]] = None
     training_vectors: Optional[DatasetDescriptor] = None
+    FILENAME_PREFIX: str = "xt"
 
     def __post_init__(self):
         self.get_name()
@@ -254,7 +255,7 @@ def name_from_factory(self) -> str:
         name += f"d_{self.d}.{self.metric.upper()}."
         if self.factory != "Flat":
             assert self.training_vectors is not None
-            name += self.training_vectors.get_filename("xt")
+            name += self.training_vectors.get_filename(CodecDescriptor.FILENAME_PREFIX)
         name += IndexBaseDescriptor.param_dict_list_to_name(self.construction_params)
         return name
 
@@ -278,6 +279,7 @@ def alias(self, benchmark_io: BenchmarkIO):
 class IndexDescriptor(IndexBaseDescriptor):
     codec_desc: Optional[CodecDescriptor] = None
     database_desc: Optional[DatasetDescriptor] = None
+    FILENAME_PREFIX: str = "xb"
 
     def __hash__(self):
         return hash(str(self))
@@ -290,14 +292,14 @@ def is_built(self):
 
     def get_name(self) -> str:
         if self.desc_name is None:
-            self.desc_name = self.codec_desc.get_name() + self.database_desc.get_filename(prefix="xb")
+            self.desc_name = self.codec_desc.get_name() + self.database_desc.get_filename(prefix=IndexDescriptor.FILENAME_PREFIX)
 
         return self.desc_name
 
     def flat_name(self):
         if self.flat_desc_name is not None:
             return self.flat_desc_name
-        self.flat_desc_name = self.codec_desc.flat_name() + self.database_desc.get_filename(prefix="xb")
+        self.flat_desc_name = self.codec_desc.flat_name() + self.database_desc.get_filename(prefix=IndexDescriptor.FILENAME_PREFIX)
         return self.flat_desc_name
 
     # alias is used to refer when index is uploaded to blobstore and refered again
@@ -313,6 +315,7 @@ class KnnDescriptor(IndexBaseDescriptor):
     query_dataset: Optional[DatasetDescriptor] = None
     search_params: Optional[Dict[str, int]] = None
     reconstruct: bool = False
+    FILENAME_PREFIX: str = "q"
     # range metric definitions
     # key: name
     # value: one of the following:
@@ -340,7 +343,7 @@ def __hash__(self):
     def get_name(self):
         name = self.index_desc.get_name()
         name += IndexBaseDescriptor.param_dict_to_name(self.search_params)
-        name += self.query_dataset.get_filename("q")
+        name += self.query_dataset.get_filename(KnnDescriptor.FILENAME_PREFIX)
         name += f"k_{self.k}."
         name += f"t_{self.num_threads}."
         if self.reconstruct:
@@ -353,7 +356,7 @@ def flat_name(self):
         if self.flat_desc_name is not None:
             return self.flat_desc_name
         name = self.index_desc.flat_name()
-        name += self.query_dataset.get_filename("q")
+        name += self.query_dataset.get_filename(KnnDescriptor.FILENAME_PREFIX)
         name += f"k_{self.k}."
         name += f"t_{self.num_threads}."
         if self.reconstruct:

diff --git a/benchs/bench_hamming_computer.cpp b/benchs/bench_hamming_computer.cpp
@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <omp.h>
 #include <cstdio>
 #include <vector>
 

diff --git a/benchs/bench_heap_replace.cpp b/benchs/bench_heap_replace.cpp
@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <omp.h>
 #include <cstdio>
 
 #include <faiss/impl/FaissAssert.h>

diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py
@@ -56,6 +56,13 @@ def swig_ptr_from_FloatTensor(x):
     return faiss.cast_integer_to_float_ptr(
         x.untyped_storage().data_ptr() + x.storage_offset() * 4)
 
+def swig_ptr_from_BFloat16Tensor(x):
+    """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
+    assert x.is_contiguous()
+    assert x.dtype == torch.bfloat16
+    return faiss.cast_integer_to_void_ptr(
+        x.untyped_storage().data_ptr() + x.storage_offset() * 2)
+
 
 def swig_ptr_from_IntTensor(x):
     """ gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
@@ -606,8 +613,11 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRI
     elif xb.dtype == torch.float16:
         xb_type = faiss.DistanceDataType_F16
         xb_ptr = swig_ptr_from_HalfTensor(xb)
+    elif xb.dtype == torch.bfloat16:
+        xb_type = faiss.DistanceDataType_BF16
+        xb_ptr = swig_ptr_from_BFloat16Tensor(xb)
     else:
-        raise TypeError('xb must be f32 or f16')
+        raise TypeError('xq must be float32, float16 or bfloat16')
 
     nq, d2 = xq.size()
     assert d2 == d
@@ -625,8 +635,11 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRI
     elif xq.dtype == torch.float16:
         xq_type = faiss.DistanceDataType_F16
         xq_ptr = swig_ptr_from_HalfTensor(xq)
+    elif xq.dtype == torch.bfloat16:
+        xq_type = faiss.DistanceDataType_BF16
+        xq_ptr = swig_ptr_from_BFloat16Tensor(xq)
     else:
-        raise TypeError('xq must be f32 or f16')
+        raise TypeError('xq must be float32, float16 or bfloat16')
 
     if D is None:
         D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)

diff --git a/demos/demo_imi_flat.cpp b/demos/demo_imi_flat.cpp
@@ -15,7 +15,6 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexPQ.h>
-#include <faiss/index_io.h>
 
 double elapsed() {
     struct timeval tv;

diff --git a/demos/demo_sift1M.cpp b/demos/demo_sift1M.cpp
@@ -12,8 +12,6 @@
 #include <cstring>
 
 #include <sys/stat.h>
-#include <sys/types.h>
-#include <unistd.h>
 
 #include <sys/time.h>
 

diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt
@@ -125,6 +125,7 @@ set(FAISS_HEADERS
   IndexIVFPQR.h
   IndexIVFSpectralHash.h
   IndexLSH.h
+  IndexNeuralNetCodec.h
   IndexLattice.h
   IndexNNDescent.h
   IndexNSG.h
@@ -168,6 +169,7 @@ set(FAISS_HEADERS
   impl/ScalarQuantizer.h
   impl/ThreadedIndex-inl.h
   impl/ThreadedIndex.h
+  impl/index_read_utils.h
   impl/io.h
   impl/io_macros.h
   impl/kmeans1d.h
@@ -179,13 +181,17 @@ set(FAISS_HEADERS
   impl/code_distance/code_distance.h
   impl/code_distance/code_distance-generic.h
   impl/code_distance/code_distance-avx2.h
+  impl/code_distance/code_distance-avx512.h
+  impl/code_distance/code_distance-sve.h
   invlists/BlockInvertedLists.h
   invlists/DirectMap.h
   invlists/InvertedLists.h
   invlists/InvertedListsIOHook.h
+  invlists/OnDiskInvertedLists.h
   utils/AlignedTable.h
   utils/bf16.h
   utils/Heap.h
+  utils/NeuralNet.h
   utils/WorkerThread.h
   utils/distances.h
   utils/extra_distances-inl.h
@@ -204,8 +210,10 @@ set(FAISS_HEADERS
   utils/sorting.h
   utils/simdlib.h
   utils/simdlib_avx2.h
+  utils/simdlib_avx512.h
   utils/simdlib_emulated.h
   utils/simdlib_neon.h
+  utils/simdlib_ppc64.h
   utils/utils.h
   utils/distances_fused/avx512.h
   utils/distances_fused/distances_fused.h
@@ -216,6 +224,7 @@ set(FAISS_HEADERS
   utils/approx_topk/mode.h
   utils/approx_topk_hamming/approx_topk_hamming.h
   utils/transpose/transpose-avx2-inl.h
+  utils/transpose/transpose-avx512-inl.h
   utils/hamming_distance/common.h
   utils/hamming_distance/generic-inl.h
   utils/hamming_distance/hamdis-inl.h

diff --git a/faiss/IndexAdditiveQuantizer.cpp b/faiss/IndexAdditiveQuantizer.cpp
@@ -16,7 +16,6 @@
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
-#include <faiss/utils/utils.h>
 
 namespace faiss {
 

diff --git a/faiss/IndexAdditiveQuantizerFastScan.cpp b/faiss/IndexAdditiveQuantizerFastScan.cpp
@@ -8,11 +8,8 @@
 #include <faiss/IndexAdditiveQuantizerFastScan.h>
 
 #include <cassert>
-#include <climits>
 #include <memory>
 
-#include <omp.h>
-
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
 #include <faiss/impl/LookupTableScaler.h>

diff --git a/faiss/IndexBinaryFlat.cpp b/faiss/IndexBinaryFlat.cpp
@@ -14,7 +14,6 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 #include <cstring>
 
 namespace faiss {

diff --git a/faiss/IndexBinaryHNSW.cpp b/faiss/IndexBinaryHNSW.cpp
@@ -15,11 +15,6 @@
 #include <cstring>
 #include <memory>
 
-#include <queue>
-#include <unordered_set>
-
-#include <sys/stat.h>
-#include <sys/types.h>
 #include <cstdint>
 
 #include <faiss/IndexBinaryFlat.h>

diff --git a/faiss/IndexBinaryHash.cpp b/faiss/IndexBinaryHash.cpp
@@ -19,7 +19,6 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 

diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp
@@ -17,11 +17,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/utils/distances.h>
-#include <faiss/utils/extra_distances.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/random.h>
-#include <faiss/utils/utils.h>
 
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>

diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp
@@ -16,7 +16,6 @@
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/prefetch.h>
 #include <faiss/utils/sorting.h>
-#include <faiss/utils/utils.h>
 #include <cstring>
 
 namespace faiss {

diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp
@@ -19,10 +19,7 @@
 #include <memory>
 #include <queue>
 #include <random>
-#include <unordered_set>
 
-#include <sys/stat.h>
-#include <sys/types.h>
 #include <cstdint>
 
 #include <faiss/Index2Layer.h>
@@ -31,7 +28,6 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/utils/distances.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
 

diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp
@@ -12,7 +12,6 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
-#include <limits>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>

diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp
@@ -18,7 +18,6 @@
 #include <cinttypes>
 #include <cstdio>
 #include <limits>
-#include <memory>
 
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>

diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/faiss/IndexIVFAdditiveQuantizer.cpp
@@ -16,7 +16,6 @@
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
-#include <faiss/utils/utils.h>
 
 namespace faiss {
 

diff --git a/faiss/IndexIVFAdditiveQuantizerFastScan.cpp b/faiss/IndexIVFAdditiveQuantizerFastScan.cpp
@@ -11,8 +11,6 @@
 #include <cinttypes>
 #include <cstdio>
 
-#include <omp.h>
-
 #include <memory>
 
 #include <faiss/impl/AuxIndexStructures.h>
@@ -23,7 +21,6 @@
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/quantize_lut.h>
-#include <faiss/utils/simdlib.h>
 #include <faiss/utils/utils.h>
 
 namespace faiss {

diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp
@@ -23,7 +23,6 @@
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
 #include <faiss/invlists/BlockInvertedLists.h>
-#include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/quantize_lut.h>
 #include <faiss/utils/utils.h>

diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
@@ -22,7 +22,6 @@
 #include <faiss/utils/utils.h>
 
 #include <faiss/Clustering.h>
-#include <faiss/IndexFlat.h>
 
 #include <faiss/utils/hamming.h>
 

diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp
@@ -11,21 +11,17 @@
 #include <cinttypes>
 #include <cstdio>
 
-#include <omp.h>
-
 #include <memory>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/simdlib.h>
-#include <faiss/utils/utils.h>
 
 #include <faiss/invlists/BlockInvertedLists.h>
 
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
-#include <faiss/utils/quantize_lut.h>
 
 namespace faiss {
 

diff --git a/faiss/IndexIVFSpectralHash.cpp b/faiss/IndexIVFSpectralHash.cpp
@@ -19,7 +19,6 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 
 namespace faiss {
 

diff --git a/faiss/IndexLSH.cpp b/faiss/IndexLSH.cpp
@@ -15,7 +15,6 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 
 namespace faiss {