faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -7,18 +7,21 @@
 
 #include <faiss/IndexFlatCodes.h>
 
+#include <atomic>
+
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/CodePacker.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
-IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
-        : Index(d, metric), code_size(code_size) {}
+IndexFlatCodes::IndexFlatCodes(size_t code_size_, idx_t d_, MetricType metric)
+        : Index(d_, metric), code_size(code_size_) {}
 
 IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
 
@@ -131,12 +134,16 @@ struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
     std::vector<float> vec_buffer;
     const float* query = nullptr;
 
-    GenericFlatCodesDistanceComputer(const IndexFlatCodes* codec, const VD& vd)
-            : FlatCodesDistanceComputer(codec->codes.data(), codec->code_size),
-              codec(*codec),
-              vd(vd),
-              code_buffer(codec->code_size * 4),
-              vec_buffer(codec->d * 4) {}
+    GenericFlatCodesDistanceComputer(
+            const IndexFlatCodes* codec_,
+            const VD& vd_)
+            : FlatCodesDistanceComputer(
+                      codec_->codes.data(),
+                      codec_->code_size),
+              codec(*codec_),
+              vd(vd_),
+              code_buffer(codec_->code_size * 4),
+              vec_buffer(codec_->d * 4) {}
 
     void set_query(const float* x) override {
         query = x;
@@ -199,23 +206,39 @@ struct Run_search_with_decompress {
         using SingleResultHandler =
                 typename BlockResultHandler::SingleResultHandler;
         using DC = GenericFlatCodesDistanceComputer<VectorDistance>;
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel // if (res.nq > 100)
         {
-            std::unique_ptr<DC> dc(new DC(&index, vd));
-            SingleResultHandler resi(res);
+            std::unique_ptr<DC> dc;
+            std::unique_ptr<SingleResultHandler> resi;
+            try {
+                dc = std::make_unique<DC>(&index, vd);
+                resi = std::make_unique<SingleResultHandler>(res);
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 #pragma omp for
-            for (int64_t q = 0; q < res.nq; q++) {
-                resi.begin(q);
-                dc->set_query(xq + vd.d * q);
-                for (size_t i = 0; i < ntotal; i++) {
-                    if (res.is_in_selection(i)) {
-                        float dis = (*dc)(i);
-                        resi.add_result(dis, i);
+            for (int64_t q = 0; q < static_cast<int64_t>(res.nq); q++) {
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    resi->begin(q);
+                    dc->set_query(xq + vd.d * q);
+                    for (size_t i = 0; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            float dis = (*dc)(i);
+                            resi->add_result(dis, i);
+                        }
                     }
+                    resi->end();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
                 }
-                resi.end();
             }
         }
+        omp_rethrow_if_exception(ex);
     }
 };
 
@@ -262,7 +285,7 @@ void IndexFlatCodes::search(
 }
 
 void IndexFlatCodes::range_search(
-        idx_t n,
+        idx_t /* n */,
         const float* x,
         float radius,
         RangeSearchResult* result,