Fix OMP exception safety in IndexHNSW search (#5133)

Summary:
Pull Request resolved: #5133

Exceptions thrown inside OpenMP parallel regions in IndexHNSW search methods call std::terminate because the OpenMP specification does not allow exceptions to escape parallel regions. This is triggered by corrupt serialized index data that causes allocation failures (e.g. VisitedTable construction with absurdly large ntotal) or other errors in DistanceComputer, but can occur with any exception thrown during the HNSW search loop.

Wrap the OpenMP parallel bodies in hnsw_search(), IndexHNSW::search_level_0(), and IndexHNSW2Level::search() with per-thread try/catch blocks that capture exceptions via std::exception_ptr and re-throw them on the main thread after the parallel region completes. This follows the same pattern established in IndexFlatCodes and IndexNNDescent/IndexNSG.

VisitedTable, SingleResultHandler, and DistanceComputer are moved from stack to heap allocation (std::unique_ptr) because the constructors are now inside try/catch blocks. If a stack-allocated object's constructor throws, the object never comes into existence, but OpenMP requires all threads to participate in the worksharing loop that follows — any reference to the uninitialized object would be undefined behavior. With std::unique_ptr, the variable starts as nullptr before the try block, remains nullptr if construction fails, and the interrupt flag ensures all loop iterations skip via continue without dereferencing it.

Reviewed By: mnorris11

Differential Revision: D101638521

fbshipit-source-id: 3920be5e41b1bff0490a14e2571b3dbeace64397

Author: scsiguy

faiss/IndexHNSW.cpp

@@ -8,6 +8,7 @@
 #include <faiss/IndexHNSW.h>
 
 #include <omp.h>
+#include <atomic>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
@@ -26,6 +27,7 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/FaissException.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
@@ -264,29 +266,48 @@ void hnsw_search(
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            VisitedTable vt(index->ntotal, hnsw.use_visited_hashset);
-            typename BlockResultHandler::SingleResultHandler res(bres);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(index->storage));
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
+                    res;
+            std::unique_ptr<DistanceComputer> dis;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        index->ntotal, hnsw.use_visited_hashset);
+                res = std::make_unique<
+                        typename BlockResultHandler::SingleResultHandler>(bres);
+                dis.reset(storage_distance_computer(index->storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
-                res.begin(i);
-                dis->set_query(x + i * index->d);
-
-                HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
-                n1 += stats.n1;
-                n2 += stats.n2;
-                ndis += stats.ndis;
-                nhops += stats.nhops;
-                res.end();
-                vt.advance();
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    res->begin(i);
+                    dis->set_query(x + i * index->d);
+
+                    HNSWStats stats =
+                            hnsw.search(*dis, index, *res, *vt, params);
+                    n1 += stats.n1;
+                    n2 += stats.n2;
+                    ndis += stats.ndis;
+                    nhops += stats.nhops;
+                    res->end();
+                    vt->advance();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 
@@ -441,37 +462,54 @@ void IndexHNSW::search_level_0(
     using RH = HeapBlockResultHandler<HNSW::C>;
     RH bres(n, distances, labels, k);
 
+    std::exception_ptr ex;
+    std::atomic<bool> interrupt{false};
 #pragma omp parallel
     {
-        std::unique_ptr<DistanceComputer> qdis(
-                storage_distance_computer(storage));
+        std::unique_ptr<DistanceComputer> qdis;
         HNSWStats search_stats;
-        VisitedTable vt(hnsw_ntotal, hnsw.use_visited_hashset);
-        RH::SingleResultHandler res(bres);
+        std::unique_ptr<VisitedTable> vt;
+        std::unique_ptr<RH::SingleResultHandler> res;
+        try {
+            qdis.reset(storage_distance_computer(storage));
+            vt = std::make_unique<VisitedTable>(
+                    hnsw_ntotal, hnsw.use_visited_hashset);
+            res = std::make_unique<RH::SingleResultHandler>(bres);
+        } catch (...) {
+            omp_capture_exception(ex, [&] { interrupt = true; });
+        }
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            res.begin(i);
-            qdis->set_query(x + i * d);
-
-            hnsw.search_level_0(
-                    *qdis.get(),
-                    res,
-                    nprobe,
-                    nearest + i * nprobe,
-                    nearest_d + i * nprobe,
-                    search_type,
-                    search_stats,
-                    vt,
-                    params);
-            res.end();
-            vt.advance();
+            if (interrupt.load(std::memory_order_relaxed)) {
+                continue;
+            }
+            try {
+                res->begin(i);
+                qdis->set_query(x + i * d);
+
+                hnsw.search_level_0(
+                        *qdis.get(),
+                        *res,
+                        nprobe,
+                        nearest + i * nprobe,
+                        nearest_d + i * nprobe,
+                        search_type,
+                        search_stats,
+                        *vt,
+                        params);
+                res->end();
+                vt->advance();
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
         }
 #pragma omp critical
         {
             hnsw_stats.combine(search_stats);
         }
     }
+    omp_rethrow_if_exception(ex);
     if (is_similarity_metric(this->metric_type)) {
 // we need to revert the negated distances
 #pragma omp parallel for
@@ -883,73 +921,88 @@ void IndexHNSW2Level::search(
                 labels,
                 false);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
             // visited table (not hash set) for tri-state flags.
-            VisitedTable vt(ntotal, /*use_hashset=*/false);
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
-
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<DistanceComputer> dis;
             constexpr int candidates_size = 1;
-            MinimaxHeap candidates(candidates_size);
+            std::unique_ptr<MinimaxHeap> candidates;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        ntotal, /*use_hashset=*/false);
+                dis.reset(storage_distance_computer(storage));
+                candidates = std::make_unique<MinimaxHeap>(candidates_size);
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops)
             for (idx_t i = 0; i < n; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
-
-                // mark all inverted list elements as visited
-
-                for (size_t j = 0; j < nprobe; j++) {
-                    idx_t key = coarse_assign[j + i * nprobe];
-                    if (key < 0) {
-                        break;
-                    }
-                    size_t list_length = index_ivfpq->get_list_size(key);
-                    const idx_t* ids = index_ivfpq->invlists->get_ids(key);
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    idx_t* idxi = labels + i * k;
+                    float* simi = distances + i * k;
+                    dis->set_query(x + i * d);
+
+                    // mark all inverted list elements as visited
+                    for (size_t j = 0; j < nprobe; j++) {
+                        idx_t key = coarse_assign[j + i * nprobe];
+                        if (key < 0) {
+                            break;
+                        }
+                        size_t list_length = index_ivfpq->get_list_size(key);
+                        const idx_t* ids = index_ivfpq->invlists->get_ids(key);
 
-                    for (size_t jj = 0; jj < list_length; jj++) {
-                        vt.set(ids[jj]);
+                        for (size_t jj = 0; jj < list_length; jj++) {
+                            vt->set(ids[jj]);
+                        }
                     }
-                }
 
-                candidates.clear();
+                    candidates->clear();
 
-                for (int j = 0; j < k; j++) {
-                    if (idxi[j] < 0) {
-                        break;
+                    for (int j = 0; j < k; j++) {
+                        if (idxi[j] < 0) {
+                            break;
+                        }
+                        candidates->push(
+                                static_cast<storage_idx_t>(idxi[j]), simi[j]);
                     }
-                    candidates.push(
-                            static_cast<storage_idx_t>(idxi[j]), simi[j]);
-                }
-
-                // reorder from sorted to heap
-                maxheap_heapify(k, simi, idxi, simi, idxi, k);
-
-                HNSWStats search_stats;
-                search_from_candidates_2(
-                        hnsw,
-                        *dis,
-                        static_cast<int>(k),
-                        idxi,
-                        simi,
-                        candidates,
-                        vt,
-                        search_stats,
-                        0,
-                        static_cast<int>(k));
-                n1 += search_stats.n1;
-                n2 += search_stats.n2;
-                ndis += search_stats.ndis;
-                nhops += search_stats.nhops;
 
-                vt.advance();
-                vt.advance();
+                    // reorder from sorted to heap
+                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
 
-                maxheap_reorder(k, simi, idxi);
+                    HNSWStats search_stats;
+                    search_from_candidates_2(
+                            hnsw,
+                            *dis,
+                            k,
+                            idxi,
+                            simi,
+                            *candidates,
+                            *vt,
+                            search_stats,
+                            0,
+                            k);
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    ndis += search_stats.ndis;
+                    nhops += search_stats.nhops;
+
+                    vt->advance();
+                    vt->advance();
+
+                    maxheap_reorder(k, simi, idxi);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
 
         hnsw_stats.combine({n1, n2, ndis, nhops});
     }