benchmark core faiss prereqs

Summary:
1. Support `search_preassigned` in IVFFastScan
2. `try_extract_index_ivf` to search recursively and support `IndexRefine`
3. `get_InvertedListScanner` to fail where not available
4. Workaround an OpenMP issue with `IndexIVFSpectralHash`

Reviewed By: mdouze

Differential Revision: D51427241

fbshipit-source-id: 365e3f11d24e80f101f986fc358c28dcc00805fa

faiss/IVFlib.cpp

@@ -14,6 +14,7 @@
 #include <faiss/IndexIVFAdditiveQuantizer.h>
 #include <faiss/IndexIVFIndependentQuantizer.h>
 #include <faiss/IndexPreTransform.h>
+#include <faiss/IndexRefine.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
@@ -58,24 +59,29 @@ void check_compatible_for_merge(const Index* index0, const Index* index1) {
 }
 
 const IndexIVF* try_extract_index_ivf(const Index* index) {
- if (auto* pt = dynamic_cast<const IndexPreTransform*>(index)) {
- index = pt->index;
+ auto* ivf = dynamic_cast<const IndexIVF*>(index);
+ if (ivf != nullptr) {
+ return ivf;
  }
 
+ if (auto* pt = dynamic_cast<const IndexPreTransform*>(index)) {
+ return try_extract_index_ivf(pt->index);
+ }
  if (auto* idmap = dynamic_cast<const IndexIDMap*>(index)) {
- index = idmap->index;
+ return try_extract_index_ivf(idmap->index);
  }
  if (auto* idmap = dynamic_cast<const IndexIDMap2*>(index)) {
- index = idmap->index;
+ return try_extract_index_ivf(idmap->index);
  }
  if (auto* indep =
  dynamic_cast<const IndexIVFIndependentQuantizer*>(index)) {
- index = indep->index_ivf;
+ return try_extract_index_ivf(indep->index_ivf);
+ }
+ if (auto* refine = dynamic_cast<const IndexRefine*>(index)) {
+ return try_extract_index_ivf(refine->base_index);
  }
 
- auto* ivf = dynamic_cast<const IndexIVF*>(index);
-
- return ivf;
+ return nullptr;
 }
 
 IndexIVF* try_extract_index_ivf(Index* index) {

faiss/IndexIVF.cpp

@@ -904,7 +904,7 @@ void IndexIVF::range_search_preassigned(
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
  bool /*store_pairs*/,
  const IDSelector* /* sel */) const {
- return nullptr;
+ FAISS_THROW_MSG("get_InvertedListScanner not implemented");
 }
 
 void IndexIVF::reconstruct(idx_t key, float* recons) const {

faiss/IndexIVFAdditiveQuantizerFastScan.cpp

@@ -314,9 +314,11 @@ void IndexIVFAdditiveQuantizerFastScan::search(
 
  NormTableScaler scaler(norm_scale);
  if (metric_type == METRIC_L2) {
- search_dispatch_implem<true>(n, x, k, distances, labels, scaler);
+ search_dispatch_implem<true>(
+ n, x, k, distances, labels, nullptr, nullptr, scaler);
  } else {
- search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
+ search_dispatch_implem<false>(
+ n, x, k, distances, labels, nullptr, nullptr, scaler);
  }
 }