MB-61930: Add a num_threads clause to every openMP pragma.

c_api/Index_c.cpp

@@ -11,6 +11,7 @@
 #include "Index_c.h"
 #include <faiss/Index.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/OMPConfig.h>
 #include "macros_impl.h"
 
 extern "C" {
@@ -222,4 +223,9 @@ int faiss_Index_sa_decode(
     }
     CATCH_AND_HANDLE
 }
+
+void faiss_set_omp_threads(unsigned int n) {
+    faiss::set_num_omp_threads(n);
+}
+
 }

c_api/Index_c.h

@@ -275,6 +275,8 @@ int faiss_Index_sa_decode(
         const uint8_t* bytes,
         float* x);
 
+void faiss_set_omp_threads(unsigned int n);
+
 #ifdef __cplusplus
 }
 #endif

faiss/AutoTune.cpp

@@ -40,6 +40,8 @@
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexBinaryIVF.h>
 
+#include <faiss/OMPConfig.h>
+
 namespace faiss {
 
 AutoTuneCriterion::AutoTuneCriterion(idx_t nq, idx_t nnn)
@@ -89,7 +91,7 @@ double IntersectionCriterion::evaluate(const float* /*D*/, const idx_t* I)
             (gt_I.size() == gt_nnn * nq && gt_nnn >= R && nnn >= R),
             "ground truth not initialized");
     int64_t n_ok = 0;
-#pragma omp parallel for reduction(+ : n_ok)
+#pragma omp parallel for reduction(+ : n_ok) num_threads(num_omp_threads)
     for (idx_t q = 0; q < nq; q++) {
         n_ok += ranklist_intersection_size(
                 R, &gt_I[q * gt_nnn], R, I + q * nnn);
@@ -698,7 +700,7 @@ void ParameterSpace::explore(
 
         do {
             if (thread_over_batches) {
-#pragma omp parallel for
+#pragma omp parallel for num_threads(num_omp_threads)
                 for (idx_t q0 = 0; q0 < nq; q0 += batchsize) {
                     size_t q1 = q0 + batchsize;
                     if (q1 > nq)

faiss/CMakeLists.txt

@@ -47,6 +47,7 @@ set(FAISS_SRC
   IndexShardsIVF.cpp
   MatrixStats.cpp
   MetaIndexes.cpp
+  OMPConfig.cpp
   VectorTransform.cpp
   clone_index.cpp
   index_factory.cpp
@@ -138,6 +139,7 @@ set(FAISS_HEADERS
   MatrixStats.h
   MetaIndexes.h
   MetricType.h
+  OMPConfig.h
   VectorTransform.h
   clone_index.h
   index_factory.h

faiss/Clustering.cpp

@@ -17,6 +17,7 @@
 #include <cstring>
 
 #include <omp.h>
+#include <faiss/OMPConfig.h>
 
 #include <faiss/IndexFlat.h>
 #include <faiss/impl/FaissAssert.h>
@@ -152,7 +153,7 @@ void compute_centroids(
 
     size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);
 
-#pragma omp parallel
+#pragma omp parallel num_threads(num_omp_threads)
     {
         int nt = omp_get_num_threads();
         int rank = omp_get_thread_num();
@@ -192,7 +193,7 @@ void compute_centroids(
         }
     }
 
-#pragma omp parallel for
+#pragma omp parallel for num_threads(num_omp_threads)
     for (idx_t ci = 0; ci < k; ci++) {
         if (hassign[ci] == 0) {
             continue;

faiss/IVFlib.cpp

@@ -469,7 +469,7 @@ void ivf_residual_add_from_flat_codes(
     const ResidualQuantizer& rq = index->rq;
 
     // populate inverted lists
-#pragma omp parallel if (nb > 10000)
+#pragma omp parallel if (nb > 10000) num_threads(num_omp_threads)
     {
         std::vector<uint8_t> tmp_code(index->code_size);
         std::vector<float> tmp(rq.d);

faiss/Index.cpp

@@ -8,6 +8,7 @@
 // -*- c++ -*-
 
 #include <faiss/Index.h>
+#include <faiss/OMPConfig.h>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
@@ -57,7 +58,7 @@ void Index::reconstruct(idx_t, float*) const {
 void Index::reconstruct_batch(idx_t n, const idx_t* keys, float* recons) const {
     std::mutex exception_mutex;
     std::string exception_string;
-#pragma omp parallel for if (n > 1000)
+#pragma omp parallel for if (n > 1000) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         try {
             reconstruct(keys[i], &recons[i * d]);
@@ -72,7 +73,7 @@ void Index::reconstruct_batch(idx_t n, const idx_t* keys, float* recons) const {
 }
 
 void Index::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
-#pragma omp parallel for if (ni > 1000)
+#pragma omp parallel for if (ni > 1000) num_threads(num_omp_threads)
     for (idx_t i = 0; i < ni; i++) {
         reconstruct(i0 + i, recons + i * d);
     }
@@ -116,7 +117,7 @@ void Index::compute_residual_n(
         const float* xs,
         float* residuals,
         const idx_t* keys) const {
-#pragma omp parallel for
+#pragma omp parallel for num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; ++i) {
         compute_residual(&xs[i * d], &residuals[i * d], keys[i]);
     }

faiss/Index2Layer.cpp

@@ -322,7 +322,7 @@ void Index2Layer::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
 }
 
 void Index2Layer::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
-#pragma omp parallel
+#pragma omp parallel num_threads(num_omp_threads)
     {
         std::vector<float> residual(d);
 

faiss/IndexAdditiveQuantizer.cpp

@@ -137,7 +137,7 @@ void search_with_decompress(
 
     using SingleResultHandler = typename ResultHandler::SingleResultHandler;
 
-#pragma omp parallel for if(res.nq > 100)
+#pragma omp parallel for if(res.nq > 100) num_threads(num_omp_threads)
     for (int64_t q = 0; q < res.nq; q++) {
         SingleResultHandler resi(res);
         resi.begin(q);
@@ -170,7 +170,7 @@ void search_with_LUT(
 
     aq.compute_LUT(nq, xq, LUT.get());
 
-#pragma omp parallel for if(nq > 100)
+#pragma omp parallel for if(nq > 100) num_threads(num_omp_threads)
     for (int64_t q = 0; q < nq; q++) {
         SingleResultHandler resi(res);
         resi.begin(q);
@@ -571,7 +571,7 @@ void ResidualCoarseQuantizer::search(
     rq.refine_beam(
             n, 1, x, beam_size, codes.data(), nullptr, beam_distances.data());
 
-#pragma omp parallel for if (n > 4000)
+#pragma omp parallel for if (n > 4000) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         memcpy(distances + i * k,
                beam_distances.data() + beam_size * i,

faiss/IndexAdditiveQuantizerFastScan.cpp

@@ -13,6 +13,7 @@
 
 #include <omp.h>
 
+
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
 #include <faiss/impl/LookupTableScaler.h>
@@ -131,7 +132,7 @@ void IndexAdditiveQuantizerFastScan::estimate_norm_scale(
     // TODO: try max of scales
     double scale = 0;
 
-#pragma omp parallel for reduction(+ : scale)
+#pragma omp parallel for reduction(+ : scale) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         const float* lut = dis_tables.data() + i * M * ksub;
         scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut);

faiss/IndexBinaryHNSW.cpp

@@ -10,6 +10,7 @@
 #include <faiss/IndexBinaryHNSW.h>
 
 #include <omp.h>
+
 #include <cassert>
 #include <cmath>
 #include <cstdio>
@@ -116,7 +117,7 @@ void hnsw_add_vertices(
                 std::swap(order[j], order[j + rng2.rand_int(i1 - j)]);
             }
 
-#pragma omp parallel
+#pragma omp parallel num_threads(num_omp_threads)
             {
                 VisitedTable vt(ntotal);
 
@@ -201,7 +202,7 @@ void IndexBinaryHNSW::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
-#pragma omp parallel
+#pragma omp parallel num_threads(num_omp_threads)
     {
         VisitedTable vt(ntotal);
         std::unique_ptr<DistanceComputer> dis(get_distance_computer());
@@ -219,7 +220,7 @@ void IndexBinaryHNSW::search(
         }
     }
 
-#pragma omp parallel for
+#pragma omp parallel for num_threads(num_omp_threads)
     for (int i = 0; i < n * k; ++i) {
         distances[i] = std::round(((float*)distances)[i]);
     }

faiss/IndexBinaryHash.cpp

@@ -221,7 +221,7 @@ void IndexBinaryHash::range_search(
             !params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
-#pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
+#pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist) num_threads(num_omp_threads)
     {
         RangeSearchPartialResult pres(result);
 
@@ -255,7 +255,7 @@ void IndexBinaryHash::search(
     using HeapForL2 = CMax<int32_t, idx_t>;
     size_t nlist = 0, ndis = 0, n0 = 0;
 
-#pragma omp parallel for if (n > 100) reduction(+ : nlist, ndis, n0)
+#pragma omp parallel for if (n > 100) reduction(+ : nlist, ndis, n0) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         int32_t* simi = distances + k * i;
         idx_t* idxi = labels + k * i;
@@ -442,7 +442,7 @@ void IndexBinaryMultiHash::range_search(
             !params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
-#pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
+#pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist) num_threads(num_omp_threads)
     {
         RangeSearchPartialResult pres(result);
 
@@ -476,7 +476,7 @@ void IndexBinaryMultiHash::search(
     using HeapForL2 = CMax<int32_t, idx_t>;
     size_t nlist = 0, ndis = 0, n0 = 0;
 
-#pragma omp parallel for if (n > 100) reduction(+ : nlist, ndis, n0)
+#pragma omp parallel for if (n > 100) reduction(+ : nlist, ndis, n0) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         int32_t* simi = distances + k * i;
         idx_t* idxi = labels + k * i;

faiss/IndexBinaryIVF.cpp

@@ -10,6 +10,7 @@
 #include <faiss/IndexBinaryIVF.h>
 
 #include <omp.h>
+
 #include <cinttypes>
 #include <cstdio>
 
@@ -391,7 +392,7 @@ void search_knn_hamming_heap(
     using HeapForIP = CMin<int32_t, idx_t>;
     using HeapForL2 = CMax<int32_t, idx_t>;
 
-#pragma omp parallel if (n > 1) reduction(+ : nlistv, ndis, nheap)
+#pragma omp parallel if (n > 1) reduction(+ : nlistv, ndis, nheap) num_threads(num_omp_threads)
     {
         std::unique_ptr<BinaryInvertedListScanner> scanner(
                 ivf.get_InvertedListScanner(store_pairs));
@@ -494,7 +495,7 @@ void search_knn_hamming_count(
 
     size_t nlistv = 0, ndis = 0;
 
-#pragma omp parallel for reduction(+ : nlistv, ndis)
+#pragma omp parallel for reduction(+ : nlistv, ndis) num_threads(num_omp_threads)
     for (int64_t i = 0; i < nx; i++) {
         const idx_t* keysi = keys + i * nprobe;
         HCounterState<HammingComputer>& csi = cs[i];
@@ -920,7 +921,7 @@ void IndexBinaryIVF::range_search_preassigned(
 
     std::vector<RangeSearchPartialResult*> all_pres(omp_get_max_threads());
 
-#pragma omp parallel reduction(+ : nlistv, ndis)
+#pragma omp parallel reduction(+ : nlistv, ndis) num_threads(num_omp_threads)
     {
         RangeSearchPartialResult pres(res);
         std::unique_ptr<BinaryInvertedListScanner> scanner(

faiss/IndexFastScan.cpp

@@ -13,6 +13,7 @@
 
 #include <omp.h>
 
+
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/LookupTableScaler.h>
@@ -348,7 +349,7 @@ void IndexFastScan::search_implem_234(
         }
     }
 
-#pragma omp parallel for if (n > 1000)
+#pragma omp parallel for if (n > 1000) num_threads(num_omp_threads)
     for (int64_t i = 0; i < n; i++) {
         int64_t* heap_ids = labels + i * k;
         float* heap_dis = distances + i * k;

faiss/IndexFlat.cpp

@@ -227,7 +227,7 @@ void IndexFlat1D::search(
             perm.size() == ntotal, "Call update_permutation before search");
     const float* xb = get_xb();
 
-#pragma omp parallel for if (n > 10000)
+#pragma omp parallel for if (n > 10000) num_threads(num_omp_threads)
     for (idx_t i = 0; i < n; i++) {
         float q = x[i]; // query
         float* D = distances + i * k;