Fix uncaught exceptions in IndexIVF::range_search_preassigned OpenMP region (#5053)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5053

C++ exceptions thrown inside `#pragma omp parallel` regions that are not
caught within the region call `std::terminate` — they cannot propagate
across thread boundaries.

`IndexIVF::range_search_preassigned` had the same class of bugs fixed in
`search_preassigned` by D99455250:

1. **`scan_list_func` lambda**:
   `FAISS_THROW_IF_NOT_FMT(key < nlist)` was above the try-catch block,
   so a corrupt key >= nlist would throw uncaught and call
   `std::terminate`.

2. **Outer parallel region**:
   `get_InvertedListScanner()`, `scanner->set_query()`, and
   `FAISS_THROW_IF_NOT(scanner.get())` had no try-catch at all.

Fixes:

1. Moved the existing try-catch in `scan_list_func` up to also cover the
   key validation.

2. Wrapped the entire `#pragma omp parallel` body in a try-catch that
   uses the existing `interrupt`/`exception_string`/`exception_mutex`
   pattern to safely propagate exceptions out of the parallel region.

Reviewed By: mnorris11

Differential Revision: D99879998

fbshipit-source-id: e768372cbaf8a22a9459fc3fd9b9df6e019897a6
This commit is contained in:
Justin Gibbs
2026-04-09 15:30:20 -07:00
committed by meta-codesync[bot]
parent 255250985a
commit 06ef481afa
2 changed files with 135 additions and 82 deletions
+98 -82
View File
@@ -828,103 +828,119 @@ void IndexIVF::range_search_preassigned(
#pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis)
{
RangeSearchPartialResult pres(result);
std::unique_ptr<InvertedListScanner> scanner(
get_InvertedListScanner(store_pairs, sel, params));
FAISS_THROW_IF_NOT(scanner.get());
all_pres[omp_get_thread_num()] = &pres;
try {
RangeSearchPartialResult pres(result);
std::unique_ptr<InvertedListScanner> scanner(
get_InvertedListScanner(store_pairs, sel, params));
FAISS_THROW_IF_NOT(scanner.get());
all_pres[omp_get_thread_num()] = &pres;
// prepare the list scanning function
// prepare the list scanning function
auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
idx_t key = keys[i * cur_nprobe + ik]; /* select the list */
if (key < 0) {
return;
}
FAISS_THROW_IF_NOT_FMT(
key < (idx_t)nlist,
"Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
key,
ik,
nlist);
auto scan_list_func = [&](size_t i,
size_t ik,
RangeQueryResult& qres) {
try {
idx_t key = keys[i * cur_nprobe + ik]; /* select the list */
if (key < 0) {
return;
}
if (invlists->is_empty(key, inverted_list_context)) {
return;
}
FAISS_THROW_IF_NOT_FMT(
key < (idx_t)nlist,
"Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
key,
ik,
nlist);
try {
size_t list_size = 0;
scanner->set_list(key, coarse_dis[i * cur_nprobe + ik]);
if (invlists->use_iterator) {
std::unique_ptr<InvertedListsIterator> it(
invlists->get_iterator(key, inverted_list_context));
if (invlists->is_empty(key, inverted_list_context)) {
return;
}
scanner->iterate_codes_range(
it.get(), radius, qres, list_size);
} else {
InvertedLists::ScopedCodes scodes(invlists, key);
InvertedLists::ScopedIds ids(invlists, key);
list_size = invlists->list_size(key);
size_t list_size = 0;
scanner->set_list(key, coarse_dis[i * cur_nprobe + ik]);
if (invlists->use_iterator) {
std::unique_ptr<InvertedListsIterator> it(
invlists->get_iterator(
key, inverted_list_context));
scanner->scan_codes_range(
list_size, scodes.get(), ids.get(), radius, qres);
scanner->iterate_codes_range(
it.get(), radius, qres, list_size);
} else {
InvertedLists::ScopedCodes scodes(invlists, key);
InvertedLists::ScopedIds ids(invlists, key);
list_size = invlists->list_size(key);
scanner->scan_codes_range(
list_size,
scodes.get(),
ids.get(),
radius,
qres);
}
nlistv++;
ndis += list_size;
} catch (const std::exception& e) {
std::lock_guard<std::mutex> lock(exception_mutex);
exception_string = demangle_cpp_symbol(typeid(e).name()) +
" " + e.what();
interrupt = true;
}
nlistv++;
ndis += list_size;
} catch (const std::exception& e) {
std::lock_guard<std::mutex> lock(exception_mutex);
exception_string =
demangle_cpp_symbol(typeid(e).name()) + " " + e.what();
interrupt = true;
}
};
};
if (parallel_mode == 0) {
if (parallel_mode == 0) {
#pragma omp for
for (idx_t i = 0; i < nx; i++) {
scanner->set_query(x + i * d);
RangeQueryResult& qres = pres.new_result(i);
for (idx_t ik = 0; ik < cur_nprobe; ik++) {
scan_list_func(i, ik, qres);
}
}
} else if (parallel_mode == 1) {
for (idx_t i = 0; i < nx; i++) {
scanner->set_query(x + i * d);
RangeQueryResult& qres = pres.new_result(i);
#pragma omp for schedule(dynamic)
for (int64_t ik = 0; ik < cur_nprobe; ik++) {
scan_list_func(i, ik, qres);
}
}
} else if (parallel_mode == 2) {
RangeQueryResult* qres = nullptr;
#pragma omp for schedule(dynamic)
for (idx_t iik = 0; iik < nx * (idx_t)cur_nprobe; iik++) {
idx_t i = iik / (idx_t)cur_nprobe;
idx_t ik = iik % (idx_t)cur_nprobe;
if (qres == nullptr || qres->qno != i) {
qres = &pres.new_result(i);
for (idx_t i = 0; i < nx; i++) {
scanner->set_query(x + i * d);
RangeQueryResult& qres = pres.new_result(i);
for (idx_t ik = 0; ik < cur_nprobe; ik++) {
scan_list_func(i, ik, qres);
}
}
scan_list_func(i, ik, *qres);
} else if (parallel_mode == 1) {
for (idx_t i = 0; i < nx; i++) {
scanner->set_query(x + i * d);
RangeQueryResult& qres = pres.new_result(i);
#pragma omp for schedule(dynamic)
for (int64_t ik = 0; ik < cur_nprobe; ik++) {
scan_list_func(i, ik, qres);
}
}
} else if (parallel_mode == 2) {
RangeQueryResult* qres = nullptr;
#pragma omp for schedule(dynamic)
for (idx_t iik = 0; iik < nx * (idx_t)cur_nprobe; iik++) {
idx_t i = iik / (idx_t)cur_nprobe;
idx_t ik = iik % (idx_t)cur_nprobe;
if (qres == nullptr || qres->qno != i) {
qres = &pres.new_result(i);
scanner->set_query(x + i * d);
}
scan_list_func(i, ik, *qres);
}
} else {
FAISS_THROW_FMT(
"parallel_mode %d not supported\n", parallel_mode);
}
} else {
FAISS_THROW_FMT("parallel_mode %d not supported\n", parallel_mode);
}
if (parallel_mode == 0) {
pres.finalize();
} else {
if (parallel_mode == 0) {
pres.finalize();
} else {
#pragma omp barrier
#pragma omp single
RangeSearchPartialResult::merge(all_pres, false);
RangeSearchPartialResult::merge(all_pres, false);
#pragma omp barrier
}
} catch (const std::exception& e) {
std::lock_guard<std::mutex> lock(exception_mutex);
exception_string =
demangle_cpp_symbol(typeid(e).name()) + " " + e.what();
interrupt = true;
}
}
+37
View File
@@ -8,6 +8,7 @@
#include <omp.h>
#include <algorithm>
#include <cstddef>
#include <limits>
#include <map>
#include <random>
#include <set>
@@ -16,6 +17,7 @@
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
namespace {
@@ -292,3 +294,38 @@ TEST(IVF, search_preassigned_out_of_range_key) {
false),
faiss::FaissException);
}
// Test: range_search_preassigned with out-of-range keys throws a catchable
// FaissException instead of calling std::terminate from an uncaught
// exception inside the OpenMP parallel region.
TEST(IVF, range_search_preassigned_out_of_range_key) {
int d = 4;
int nlist = 2;
faiss::IndexFlatL2 quantizer(d);
faiss::IndexIVFFlat idx(&quantizer, d, nlist);
idx.own_fields = false;
std::vector<float> train_data(nlist * d, 0.0f);
for (int i = 0; i < nlist * d; i++) {
train_data[i] = static_cast<float>(i);
}
idx.train(nlist, train_data.data());
idx.add(nlist, train_data.data());
std::vector<float> xq(d, 1.0f);
faiss::RangeSearchResult result(1);
faiss::idx_t bad_key = nlist; // out of range
float coarse_dis = 0.0f;
EXPECT_THROW(
idx.range_search_preassigned(
1,
xq.data(),
std::numeric_limits<float>::max(),
&bad_key,
&coarse_dis,
&result,
false),
faiss::FaissException);
}