From dfce6e948e153b783d8bd0785c1c30432e84a166 Mon Sep 17 00:00:00 2001 From: Justin Gibbs Date: Fri, 24 Apr 2026 12:53:56 -0700 Subject: [PATCH] Fix OMP exception safety in IndexHNSW search (#5133) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/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 --- faiss/IndexHNSW.cpp | 223 +++++++++++------- tests/test_omp_exception_safety.cpp | 336 ++++++++++++++++++++++++++++ 2 files changed, 474 insertions(+), 85 deletions(-) diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index b92132d19..ee6269735 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -26,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -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 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 dis( - storage_distance_computer(index->storage)); + std::unique_ptr vt; + std::unique_ptr + res; + std::unique_ptr dis; + try { + vt = std::make_unique( + 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); + 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(); + 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; RH bres(n, distances, labels, k); + std::exception_ptr ex; + std::atomic interrupt{false}; #pragma omp parallel { - std::unique_ptr qdis( - storage_distance_computer(storage)); + std::unique_ptr qdis; HNSWStats search_stats; - VisitedTable vt(hnsw_ntotal, hnsw.use_visited_hashset); - RH::SingleResultHandler res(bres); + std::unique_ptr vt; + std::unique_ptr res; + try { + qdis.reset(storage_distance_computer(storage)); + vt = std::make_unique( + hnsw_ntotal, hnsw.use_visited_hashset); + res = std::make_unique(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); + 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(); + 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 interrupt{false}; #pragma omp parallel { // visited table (not hash set) for tri-state flags. - VisitedTable vt(ntotal, /*use_hashset=*/false); - std::unique_ptr dis( - storage_distance_computer(storage)); - + std::unique_ptr vt; + std::unique_ptr dis; constexpr int candidates_size = 1; - MinimaxHeap candidates(candidates_size); + std::unique_ptr candidates; + try { + vt = std::make_unique( + ntotal, /*use_hashset=*/false); + dis.reset(storage_distance_computer(storage)); + candidates = std::make_unique(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); - - for (size_t jj = 0; jj < list_length; jj++) { - vt.set(ids[jj]); - } + 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); - candidates.clear(); + // 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 (int j = 0; j < k; j++) { - if (idxi[j] < 0) { - break; + for (size_t jj = 0; jj < list_length; jj++) { + vt->set(ids[jj]); + } } - candidates.push( - static_cast(idxi[j]), simi[j]); + + candidates->clear(); + + for (int j = 0; j < k; j++) { + if (idxi[j] < 0) { + break; + } + candidates->push( + static_cast(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, + 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; }); } - - // reorder from sorted to heap - maxheap_heapify(k, simi, idxi, simi, idxi, k); - - HNSWStats search_stats; - search_from_candidates_2( - hnsw, - *dis, - static_cast(k), - idxi, - simi, - candidates, - vt, - search_stats, - 0, - static_cast(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); } } + omp_rethrow_if_exception(ex); hnsw_stats.combine({n1, n2, ndis, nhops}); } diff --git a/tests/test_omp_exception_safety.cpp b/tests/test_omp_exception_safety.cpp index 8c5f845c9..a34cd9ae1 100644 --- a/tests/test_omp_exception_safety.cpp +++ b/tests/test_omp_exception_safety.cpp @@ -13,9 +13,12 @@ #include +#include #include #include +#include #include +#include #include #include #include @@ -381,3 +384,336 @@ TEST(OMPExceptionSafety, nsg_search) { index.search(1, xq.data(), 1, distances.data(), labels.data()), FaissException); } + +// --------------------------------------------------------------------------- +// IndexHNSW OMP exception safety tests. +// +// There are 3 OMP regions wrapped with exception capture in IndexHNSW.cpp: +// 1. hnsw_search() — primary search path +// 2. IndexHNSW::search_level_0 — level-0 search +// 3. IndexHNSW2Level::search — mixed IVFPQ+HNSW path +// +// Each region has two throw locations: +// (a) Init phase — DistanceComputer/VisitedTable construction +// (b) Loop body — set_query, graph traversal, etc. +// --------------------------------------------------------------------------- + +// DistanceComputer that throws from set_query(). Exercises loop-body +// try/catch after successful init-phase construction. +struct ThrowingDC : DistanceComputer { + void set_query(const float*) override { + throw std::runtime_error("corrupt storage"); + } + float operator()(idx_t) override { + return 0; + } + float symmetric_dis(idx_t, idx_t) override { + return 0; + } +}; + +// Storage whose get_distance_computer() returns a ThrowingDC. +// Exercises loop-body throws (DC construction succeeds, set_query throws). +struct LoopThrowingStorage : IndexFlatL2 { + using IndexFlatL2::IndexFlatL2; + DistanceComputer* get_distance_computer() const override { + return new ThrowingDC(); + } +}; + +// Storage whose get_distance_computer() itself throws. +// Exercises init-phase throws (DC construction fails). +struct InitThrowingStorage : IndexFlatL2 { + using IndexFlatL2::IndexFlatL2; + DistanceComputer* get_distance_computer() const override { + throw std::runtime_error("corrupt storage init"); + } +}; + +// Helper: build a valid IndexHNSWFlat with data, then swap storage. +struct HNSWFixture { + static constexpr int d = 4; + static constexpr int nb = 10; + static constexpr int M = 16; + + IndexHNSWFlat index{d, M}; + std::vector xb; + std::vector xq; + + HNSWFixture() : xb(d * nb), xq(d) { + std::mt19937 rng(42); + std::uniform_real_distribution dist; + for (auto& v : xb) { + v = dist(rng); + } + index.add(nb, xb.data()); + for (auto& v : xq) { + v = dist(rng); + } + } + + void swap_storage(Index* new_storage) { + index.own_fields = false; + index.storage = new_storage; + } + + void restore_storage(Index* old_storage) { + index.storage = old_storage; + index.own_fields = true; + } +}; + +// --------------------------------------------------------------------------- +// hnsw_search() init-phase throw: storage_distance_computer() throws +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_search_init) { + HNSWFixture f; + auto throwing = std::make_unique(f.d); + throwing->add(f.nb, f.xb.data()); + auto* old = f.index.storage; + f.swap_storage(throwing.get()); + + std::vector distances(1); + std::vector labels(1); + EXPECT_THROW( + f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()), + std::runtime_error); + + f.restore_storage(old); +} + +// --------------------------------------------------------------------------- +// hnsw_search() loop-body throw: set_query() throws +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_search_loop) { + HNSWFixture f; + auto throwing = std::make_unique(f.d); + throwing->add(f.nb, f.xb.data()); + auto* old = f.index.storage; + f.swap_storage(throwing.get()); + + std::vector distances(1); + std::vector labels(1); + EXPECT_THROW( + f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()), + std::runtime_error); + + f.restore_storage(old); +} + +// --------------------------------------------------------------------------- +// search_level_0() init-phase throw: storage_distance_computer() throws +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_search_level_0_init) { + HNSWFixture f; + auto throwing = std::make_unique(f.d); + throwing->add(f.nb, f.xb.data()); + auto* old = f.index.storage; + f.swap_storage(throwing.get()); + + idx_t k = 1; + int nprobe = 1; + std::vector nearest(nprobe, 0); + std::vector nearest_d(nprobe, 0.0f); + std::vector distances(k); + std::vector labels(k); + + EXPECT_THROW( + f.index.search_level_0( + 1, + f.xq.data(), + k, + nearest.data(), + nearest_d.data(), + distances.data(), + labels.data(), + nprobe), + std::runtime_error); + + f.restore_storage(old); +} + +// --------------------------------------------------------------------------- +// search_level_0() loop-body throw: set_query() throws +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_search_level_0_loop) { + HNSWFixture f; + auto throwing = std::make_unique(f.d); + throwing->add(f.nb, f.xb.data()); + auto* old = f.index.storage; + f.swap_storage(throwing.get()); + + idx_t k = 1; + int nprobe = 1; + std::vector nearest(nprobe, 0); + std::vector nearest_d(nprobe, 0.0f); + std::vector distances(k); + std::vector labels(k); + + EXPECT_THROW( + f.index.search_level_0( + 1, + f.xq.data(), + k, + nearest.data(), + nearest_d.data(), + distances.data(), + labels.data(), + nprobe), + std::runtime_error); + + f.restore_storage(old); +} + +// IndexIVFPQ subclass with a counted throw mechanism. The first +// `throw_after` calls to get_distance_computer() delegate to the +// real implementation; subsequent calls throw (init) or return +// ThrowingDC (loop). This lets search_preassigned succeed while +// the HNSW2Level OMP region's storage_distance_computer() fails. +struct CountedThrowIVFPQ : IndexIVFPQ { + mutable std::atomic calls_remaining; + bool throw_from_init; + + CountedThrowIVFPQ( + Index* quantizer, + size_t d, + size_t nlist, + size_t M, + size_t nbits, + MetricType metric, + int throw_after, + bool init_throw) + : IndexIVFPQ(quantizer, d, nlist, M, nbits, metric), + calls_remaining(throw_after), + throw_from_init(init_throw) {} + + DistanceComputer* get_distance_computer() const override { + if (calls_remaining.fetch_sub(1) <= 0) { + if (throw_from_init) { + throw std::runtime_error("IVFPQ init throw"); + } + return new ThrowingDC(); + } + return IndexIVFPQ::get_distance_computer(); + } +}; + +struct HNSW2LevelFixture { + static constexpr int d = 8; + static constexpr int nlist = 4; + static constexpr int m_pq = 4; + static constexpr int nb = 512; + static constexpr int M = 16; + + IndexFlatL2 quantizer{d}; + IndexHNSW2Level index{&quantizer, nlist, m_pq, M}; + std::vector xb; + std::vector xq; + + HNSW2LevelFixture() : xb(nb * d), xq(d) { + index.own_fields = false; + std::mt19937 rng(42); + std::uniform_real_distribution dist; + for (auto& v : xb) { + v = dist(rng); + } + index.train(nb, xb.data()); + index.add(nb, xb.data()); + // Convert Index2Layer storage to IndexIVFPQ for the mixed + // search path. After this, storage is an IndexIVFPQ owned + // by the HNSW2Level index. + index.flip_to_ivf(); + index.own_fields = true; + for (auto& v : xq) { + v = dist(rng); + } + } +}; + +// --------------------------------------------------------------------------- +// IndexHNSW2Level::search() mixed path init-phase throw: +// storage_distance_computer() throws during OMP init, after +// search_preassigned has already succeeded. +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_2level_mixed_init) { + HNSW2LevelFixture f; + + // The replacement IVFPQ shares the real index's trained state + // (codebooks, inverted lists). We copy the real IVFPQ's invlists + // and PQ so search_preassigned succeeds, then the counted throw + // fires during the HNSW2Level OMP init phase. + auto* real_ivfpq = dynamic_cast(f.index.storage); + ASSERT_NE(real_ivfpq, nullptr); + + CountedThrowIVFPQ throwing( + &f.quantizer, + f.d, + f.nlist, + f.m_pq, + 8, + METRIC_L2, + /*throw_after=*/100, + /*init_throw=*/true); + throwing.own_fields = false; + throwing.is_trained = true; + throwing.ntotal = real_ivfpq->ntotal; + throwing.pq = real_ivfpq->pq; + throwing.code_size = real_ivfpq->code_size; + throwing.nprobe = real_ivfpq->nprobe; + throwing.replace_invlists(real_ivfpq->invlists, false); + + auto* old = f.index.storage; + f.index.storage = &throwing; + f.index.own_fields = false; + + std::vector distances(1); + std::vector labels(1); + EXPECT_THROW( + f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()), + std::exception); + + f.index.storage = old; + f.index.own_fields = true; +} + +// --------------------------------------------------------------------------- +// IndexHNSW2Level::search() mixed path loop-body throw: +// set_query() throws inside OMP for loop, after search_preassigned +// and OMP init have succeeded. +// --------------------------------------------------------------------------- +TEST(OMPExceptionSafety, hnsw_2level_mixed_loop) { + HNSW2LevelFixture f; + + auto* real_ivfpq = dynamic_cast(f.index.storage); + ASSERT_NE(real_ivfpq, nullptr); + + CountedThrowIVFPQ throwing( + &f.quantizer, + f.d, + f.nlist, + f.m_pq, + 8, + METRIC_L2, + /*throw_after=*/100, + /*init_throw=*/false); + throwing.own_fields = false; + throwing.is_trained = true; + throwing.ntotal = real_ivfpq->ntotal; + throwing.pq = real_ivfpq->pq; + throwing.code_size = real_ivfpq->code_size; + throwing.nprobe = real_ivfpq->nprobe; + throwing.replace_invlists(real_ivfpq->invlists, false); + + auto* old = f.index.storage; + f.index.storage = &throwing; + f.index.own_fields = false; + + std::vector distances(1); + std::vector labels(1); + EXPECT_THROW( + f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()), + std::exception); + + f.index.storage = old; + f.index.own_fields = true; +}