mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
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:
committed by
meta-codesync[bot]
parent
255250985a
commit
06ef481afa
+98
-82
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user