diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp index 855665c617..00bc6c823e 100644 --- a/faiss/IndexIVFFastScan.cpp +++ b/faiss/IndexIVFFastScan.cpp @@ -643,7 +643,7 @@ void IndexIVFFastScan::range_search_dispatch_implem( { RangeSearchPartialResult pres(&rres); -#pragma omp parallel for reduction(+ : ndis, nlist_visited) +#pragma omp for reduction(+ : ndis, nlist_visited) for (int slice = 0; slice < nslice; slice++) { idx_t i0 = n * slice / nslice; idx_t i1 = n * (slice + 1) / nslice; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a1c3266961..8522fa613d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -31,6 +31,7 @@ set(FAISS_TEST_SRC test_code_distance.cpp test_hnsw.cpp test_partitioning.cpp + test_fastscan_perf.cpp ) add_executable(faiss_test ${FAISS_TEST_SRC}) diff --git a/tests/test_fastscan_perf.cpp b/tests/test_fastscan_perf.cpp new file mode 100644 index 0000000000..82321b2ede --- /dev/null +++ b/tests/test_fastscan_perf.cpp @@ -0,0 +1,70 @@ +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +TEST(TestFastScan, knnVSrange) { + // small vectors and database + int d = 64; + size_t nb = 1000; + + // ivf centroids + size_t nlist = 4; + + // more than 2 threads to surface + // problems related to multi-threading + omp_set_num_threads(8); + + // random database, also used as queries + std::vector database(nb * d); + std::mt19937 rng; + std::uniform_real_distribution<> distrib; + for (size_t i = 0; i < nb * d; i++) { + database[i] = distrib(rng); + } + + // build index + faiss::IndexFlatL2 coarse_quantizer(d); + faiss::IndexIVFPQFastScan index( + &coarse_quantizer, d, nlist, d / 2, 4, faiss::METRIC_L2, 32); + index.pq.cp.niter = 10; // speed up train + index.nprobe = nlist; + index.train(nb, database.data()); + index.add(nb, database.data()); + + std::vector distances(nb); + std::vector labels(nb); + auto t = std::chrono::high_resolution_clock::now(); + index.search(nb, database.data(), 1, distances.data(), labels.data()); + auto knn_time = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - t) + .count(); + + faiss::RangeSearchResult rsr(nb); + t = std::chrono::high_resolution_clock::now(); + index.range_search(nb, database.data(), 1.0, &rsr); + auto range_time = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - t) + .count(); + + // we expect the perf of knn and range search + // to be similar, at least within a factor of 2 + ASSERT_LT(range_time, knn_time * 2); + ASSERT_LT(knn_time, range_time * 2); +}