fix: respect CPU affinity and reduce clustering parallel overhead

CMakeLists.txt

@@ -58,6 +58,7 @@ set(RABITQ_COMMON_SOURCES
     src/index/ivf_search.cpp
     src/index/qg.cpp
     src/utils/cpu_features.cpp
+    src/utils/thread_count.cpp
     src/simd/dispatch.cpp
     src/simd/space_generic.cpp
     src/simd/fastscan_generic.cpp

docs/docs/clustering.md

@@ -77,13 +77,18 @@ Set Python keyword arguments or fields of the corresponding C++ parameter type:
 | --- | --- | --- |
 | `niter` | `25` | Maximum clustering iterations. |
 | `early_stop_threshold` | `0.0` | Stop from iteration 2 when the absolute relative objective change is at most this value; range `[0, 1]`. Zero stops only on an unchanged objective. |
-| `num_threads` | `0` | Use detected hardware threads; positive requests are capped at that count. |
+| `num_threads` | `0` | Use detected available logical CPUs; positive requests are capped at that count. See [threading](quick_start.md#threading-and-file-paths). |
 | `spherical` | `false` | Normalize centroids and assign by inner product; input must already be normalized. |
 | `seed` | `42` | Seed for centroid selection and rotation. |
 | `min_points_per_centroid` | `39` | Warn below this training-points-per-centroid ratio when verbose. |
 | `verbose` | `false` | Print training progress and warnings. |
 | `final_assignment` | `Approximate` | Use the selected method for final labels, or `Exact` for nearest-centroid labels. |
 
+`num_threads` sets an upper limit, not a fixed worker count for every stage.
+QGKMeans may use fewer workers to build small centroid graphs. That additional
+limit does not apply to training-vector assignment, which uses the thread budget
+after capping it to the available CPUs.
+
 Training requires `64 <= d <= 65536`, finite coordinates, and
 `abs(x) <= sqrt(sqrt(FLT_MAX) / (64 * d))`; scale larger inputs before training.
 Cluster-count requirements are listed above. Training does not normalize or

docs/docs/index/qg.md

@@ -145,7 +145,7 @@ Python provides `search_batch(queries, k, ef, num_threads=1)` starting with 0.5.
 The existing `search()` accepts the same two-dimensional array and remains an equivalent
 entry point. Both return `(ids, distances)` arrays of shape `(num_queries, k)`,
 including empty batches. Both APIs default to one worker; `num_threads=0`
-selects the available hardware thread count, capped by the number of queries.
+selects the available logical CPU count, capped by the number of queries.
 Queries use the original dimension, including when the internal rotation pads it.
 
 See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their

docs/docs/quick_start.md

@@ -153,11 +153,11 @@ counts, or [QGKMeans](clustering.md#qgkmeans) for graph assignment.
 
 ### Threading and file paths
 
-For all three indexes, `build` and `search` interpret `num_threads=0` as the
-detected hardware thread count. Larger requests are capped at that count;
-smaller positive requests are respected. Operations may use fewer workers when
-there are fewer work items. If hardware detection is unavailable, one thread is
-used. Python index methods default to one thread when `num_threads` is omitted.
+`num_threads=0` selects the detected available logical CPU count. Positive values
+set an upper limit, capped at that count; small workloads may use fewer workers.
+Python index methods default to one thread; clustering defaults to `0`.
+On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding
+before starting the program.
 
 Index save/load paths are UTF-8 strings on Windows and native path bytes on POSIX
 in C++; Python paths are Unicode strings on all platforms.

include/rabitqlib/clustering/detail/kmeans.hpp

@@ -17,6 +17,7 @@
 #include <random>
 #include <stdexcept>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "rabitqlib/defines.hpp"
@@ -327,8 +328,8 @@ class LloydKMeans : public Parameters {
             update_centroids(
                 x, n, sums, cluster_sizes, update_point_bins, update_worker_for_cluster
             );
-            const size_t nsplit = consolidate_centroids(x, n, sums, cluster_sizes);
-            const double shift = centroid_shift(previous_centroids);
+            const auto [nsplit, shift] =
+                consolidate_centroids(x, n, sums, cluster_sizes, previous_centroids);
             iteration_stats.push_back(KMeansIterationStats{
                 iteration + 1, obj, shift, nsplit});
 
@@ -457,15 +458,29 @@ class LloydKMeans : public Parameters {
         }
     }
 
-    size_t consolidate_centroids(
+    std::pair<size_t, double> consolidate_centroids(
         const float* x,
         size_t n,
         const std::vector<double>& sums,
-        std::vector<size_t>& cluster_sizes
+        std::vector<size_t>& cluster_sizes,
+        const std::vector<float>& previous_centroids
     ) {
         const uint32_t threads = effective_num_threads();
         static_cast<void>(threads);
-#pragma omp parallel for num_threads(threads)
+        // Means, normalization, and displacement share one worker team. Preserve
+        // the per-centroid arithmetic; empty rows are measured after refilling.
+        const auto cluster_shift = [&](size_t cluster) {
+            double result = 0;
+            for (size_t dim = 0; dim < d; ++dim) {
+                const size_t i = cluster * d + dim;
+                const double difference =
+                    static_cast<double>(centroids[i]) - previous_centroids[i];
+                result += difference * difference;
+            }
+            return result;
+        };
+        double shift = 0;
+#pragma omp parallel for num_threads(threads) reduction(+ : shift)
         for (std::ptrdiff_t loop_index = 0; loop_index < static_cast<std::ptrdiff_t>(k);
              ++loop_index) {
             const size_t cluster = static_cast<size_t>(loop_index);
@@ -476,16 +491,20 @@ class LloydKMeans : public Parameters {
                 centroids[cluster * d + dim] =
                     static_cast<float>(sums[cluster * d + dim] * scale);
             }
-        }
-        if (spherical) {
-            normalize_centroids();
+            if (spherical) {
+                normalize_centroid(centroids.data() + cluster * d);
+            }
+            // Empty centroids contribute only after their replacement is known.
+            if (cluster_sizes[cluster] != 0) {
+                shift += cluster_shift(cluster);
+            }
         }
 
         const size_t nsplit = static_cast<size_t>(
             std::count(cluster_sizes.begin(), cluster_sizes.end(), size_t{0})
         );
         if (nsplit == 0) {
-            return 0;
+            return {0, shift};
         }
 
         // Reuse the distance buffer after recording the pre-update objective.
@@ -537,7 +556,8 @@ class LloydKMeans : public Parameters {
             }
             --cluster_sizes[assignments[point]];
             cluster_sizes[cluster] = 1;
-            return 1;
+            shift += cluster_shift(cluster);
+            return {1, shift};
         }
 
 #pragma omp parallel for num_threads(threads) schedule(static)
@@ -566,8 +586,9 @@ class LloydKMeans : public Parameters {
             }
             --cluster_sizes[assignments[point]];
             cluster_sizes[cluster] = 1;
+            shift += cluster_shift(cluster);
         }
-        return nsplit;
+        return {nsplit, shift};
     }
 
     void normalize_centroid(float* centroid) const {
@@ -595,23 +616,6 @@ class LloydKMeans : public Parameters {
         }
     }
 
-    [[nodiscard]] double centroid_shift(const std::vector<float>& previous_centroids
-    ) const {
-        const uint32_t threads = effective_num_threads();
-        static_cast<void>(threads);
-        double shift = 0;
-#pragma omp parallel for num_threads(threads) reduction(+ : shift)
-        for (std::ptrdiff_t loop_index = 0;
-             loop_index < static_cast<std::ptrdiff_t>(centroids.size());
-             ++loop_index) {
-            const size_t index = static_cast<size_t>(loop_index);
-            const double difference =
-                static_cast<double>(centroids[index]) - previous_centroids[index];
-            shift += difference * difference;
-        }
-        return shift;
-    }
-
     [[nodiscard]] uint32_t effective_num_threads() const {
         return static_cast<uint32_t>(resolve_num_threads(num_threads));
     }

include/rabitqlib/clustering/qgkmeans.hpp

@@ -149,11 +149,17 @@ class QGAssigner {
         if (builder_) {
             builder_->reset(centroids);
         } else {
+            // Centroid graphs have much less work than the training-vector batch.
+            // Amortize their short phases without reducing assignment parallelism.
+            constexpr size_t kMinVerticesPerBuildThread = 128;
+            const size_t build_threads = std::min<size_t>(
+                num_threads_, std::max<size_t>(1, k_ / kMinVerticesPerBuildThread)
+            );
             builder_ = std::make_unique<symqg::QGBuilder>(
                 *graph_,
                 ef_construction_,
                 centroids,
-                num_threads_,
+                build_threads,
                 symqg::QGInitialization::PiPNN,
                 seed_,
                 true

include/rabitqlib/index/ivf/initializer.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <omp.h>
+
 #include <algorithm>
 #include <atomic>
 #include <cmath>
@@ -35,6 +37,28 @@ inline void parallel_for(size_t start, size_t end, size_t numThreads, Function f
             fn(id, 0);
         }
     } else {
+#if defined(_OPENMP) && _OPENMP >= 201307
+        if (omp_get_proc_bind() != omp_proc_bind_false) {
+            // Native threads inherit only the caller's place. Let OpenMP honor
+            // the configured placement of the whole team when binding is enabled.
+            std::vector<std::exception_ptr> errors(numThreads);
+#pragma omp parallel for num_threads(numThreads) schedule(dynamic)
+            for (size_t id = start; id < end; ++id) {
+                const size_t worker = static_cast<size_t>(omp_get_thread_num());
+                if (!errors[worker]) {
+                    try {
+                        fn(id, worker);
+                    } catch (...) { errors[worker] = std::current_exception(); }
+                }
+            }
+            for (const auto& error : errors) {
+                if (error) {
+                    std::rethrow_exception(error);
+                }
+            }
+            return;
+        }
+#endif
         std::vector<std::thread> threads;
         std::atomic<size_t> current(start);
 

include/rabitqlib/index/symqg/detail/pipnn.hpp

@@ -208,15 +208,11 @@ inline std::vector<Bucket> cluster(
     while (!active.empty()) {
         std::vector<std::vector<BucketJob>> next(active.size());
         std::vector<std::vector<Bucket>> finished(active.size());
-#pragma omp parallel for num_threads(active.front().depth == 0 ? 1 : threads) \
-    schedule(dynamic)
-        for (std::ptrdiff_t index = 0; index < static_cast<std::ptrdiff_t>(active.size());
-             ++index) {
-            const size_t i = static_cast<size_t>(index);
+        const auto process_bucket = [&](size_t i) {
             auto& job = active[i];
             if (job.ids.size() <= kLeafSize) {
                 finished[i].push_back(std::move(job.ids));
-                continue;
+                return;
             }
             auto children = partition(data, dim, job, metric, threads, scratch);
             std::sort(children.begin(), children.end(), [](const auto& a, const auto& b) {
@@ -265,6 +261,17 @@ inline std::vector<Bucket> cluster(
                 finished[i].push_back(std::move(small));
             }
             Bucket().swap(job.ids);
+        };
+        if (active.front().depth == 0) {
+            // A serialized outer OpenMP region would make the root's tile team
+            // nested, forcing the runtime to create workers again on every reset.
+            process_bucket(0);
+        } else {
+#pragma omp parallel for num_threads(threads) schedule(dynamic)
+            for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(active.size());
+                 ++i) {
+                process_bucket(static_cast<size_t>(i));
+            }
         }
         std::vector<BucketJob> new_active;
         for (size_t i = 0; i < active.size(); ++i) {

include/rabitqlib/index/symqg/qg.hpp

@@ -243,7 +243,7 @@ class QuantizedGraph<float> {
      * queries contains num_queries * dimension() floats; results and dists
      * each hold num_queries * knn elements. Empty batches do no work.
      * Scratch is reused within each worker. num_threads defaults to one;
-     * zero selects the available hardware thread count, capped by the batch.
+     * zero selects the available logical CPU count, capped by the batch.
      * Input and output buffers must not overlap. Concurrent searches require
      * separate outputs and no index mutation, including set_ef().
      */

include/rabitqlib/utils/tools.hpp

@@ -81,14 +81,14 @@ constexpr T round_up_to_multiple_of(size_t x, size_t multiple_of) {
     return multiple_of * (div_round_up(x, multiple_of));
 }
 
-// get number of threads of current sys
-inline size_t total_threads() {
-    const auto threads = std::thread::hardware_concurrency();
-    return threads == 0 ? 1 : threads;
-}
+// Available logical CPUs in the calling context, with a hardware-count fallback.
+size_t total_threads();
 
-// Zero selects all hardware threads; explicit requests cannot exceed that count.
+// Zero selects all available CPUs; explicit requests cannot exceed that count.
 inline size_t resolve_num_threads(size_t requested) {
+    if (requested == 1) {
+        return 1;
+    }
     const size_t available = total_threads();
     return requested == 0 ? available : std::min(requested, available);
 }

python_bindings/rabitq_bindings.cpp

@@ -3,6 +3,7 @@
 
 #include "rabitqlib/defines.hpp"
 #include "rabitqlib/utils/rotator.hpp"
+#include "rabitqlib/utils/tools.hpp"
 
 namespace py = pybind11;
 
@@ -15,6 +16,8 @@ void register_kmeans(py::module_& m);
 PYBIND11_MODULE(_rabitqlib, m) {
     m.doc() = "RabitQ Python bindings combined module";
 
+    m.def("_available_cpu_count", &rabitqlib::total_threads);
+
     // Register shared enums once
     py::enum_<rabitqlib::MetricType>(m, "MetricType")
         .value("L2", rabitqlib::METRIC_L2)

sample/python/README.md

@@ -2,8 +2,9 @@
 
 Run these commands from the repository root in your project Python environment
 with `rabitqlib>=0.5.2` and NumPy installed. Clustering and indexing both use RaBitQ
-Library; FAISS is needed only for `compare_with_faiss.py`. For unreleased changes,
-[install from a checkout](../../CONTRIBUTING.md#python-changes).
+Library; FAISS is needed only for `compare_with_faiss.py`. When running these
+examples from a checkout, [install the library from that same checkout](../../CONTRIBUTING.md#python-changes)
+so the examples and native extension use matching APIs.
 
 ## IVF: cluster, then index
 
@@ -34,7 +35,9 @@ behavior.
 
 For inner product, pass `--metric ip` to both stages and provide normalized
 training vectors. Use queries and ground truth prepared for the same metric.
-`--num-threads 0` uses the hardware thread count; each stage accepts its own limit.
+`--num-threads 0` uses the library's available logical CPU count; each stage accepts
+its own limit. See [threading](../../docs/docs/quick_start.md#threading-and-file-paths)
+for affinity and OpenMP placement behavior.
 
 ## HNSW
 

sample/python/kmeans_clustering.py

@@ -1,11 +1,11 @@
 """Train RaBitQKMeans or QGKMeans and save clusters for the index examples."""
 
 import argparse
-import os
 from time import time
 
 import numpy as np
 from rabitqlib import FinalAssignmentMode, QGKMeans, RaBitQKMeans
+from rabitqlib._rabitqlib import _available_cpu_count
 
 from utils import read_fvecs
 
@@ -26,9 +26,9 @@ def cluster_data(
     if method not in ("rabitq", "qg"):
         raise ValueError("method must be rabitq or qg")
     clustering_type = RaBitQKMeans if method == "rabitq" else QGKMeans
-    hardware_threads = os.cpu_count() or 1
+    available_threads = _available_cpu_count()
     threads = (
-        hardware_threads if num_threads == 0 else min(num_threads, hardware_threads)
+        available_threads if num_threads == 0 else min(num_threads, available_threads)
     )
     print(f"Clustering metric: {metric.upper()}, threads: {threads}")
     kmeans = clustering_type(
@@ -76,6 +76,6 @@ def main(args) -> None:
         "--num-threads",
         type=int,
         default=0,
-        help="Threads for clustering (0: hardware count; larger requests are capped)",
+        help="Threads for clustering (0: available CPUs; larger requests are capped)",
     )
     main(parser.parse_args())