diff --git a/CMakeLists.txt b/CMakeLists.txt index 7bc8906c..0d93f4fb 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md index 628c6842..719ce814 100644 --- a/docs/docs/clustering.md +++ b/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 diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index abbe3428..35ee0c27 100644 --- a/docs/docs/index/qg.md +++ b/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 diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index bc8529f2..7cd37f80 100644 --- a/docs/docs/quick_start.md +++ b/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. diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index f8efac06..b90d7db8 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #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 consolidate_centroids( const float* x, size_t n, const std::vector& sums, - std::vector& cluster_sizes + std::vector& cluster_sizes, + const std::vector& previous_centroids ) { const uint32_t threads = effective_num_threads(); static_cast(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(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(k); ++loop_index) { const size_t cluster = static_cast(loop_index); @@ -476,16 +491,20 @@ class LloydKMeans : public Parameters { centroids[cluster * d + dim] = static_cast(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( 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& previous_centroids - ) const { - const uint32_t threads = effective_num_threads(); - static_cast(threads); - double shift = 0; -#pragma omp parallel for num_threads(threads) reduction(+ : shift) - for (std::ptrdiff_t loop_index = 0; - loop_index < static_cast(centroids.size()); - ++loop_index) { - const size_t index = static_cast(loop_index); - const double difference = - static_cast(centroids[index]) - previous_centroids[index]; - shift += difference * difference; - } - return shift; - } - [[nodiscard]] uint32_t effective_num_threads() const { return static_cast(resolve_num_threads(num_threads)); } diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index 3de60e81..88b96e3e 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/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( + num_threads_, std::max(1, k_ / kMinVerticesPerBuildThread) + ); builder_ = std::make_unique( *graph_, ef_construction_, centroids, - num_threads_, + build_threads, symqg::QGInitialization::PiPNN, seed_, true diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index b56a67e7..8885369d 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include #include #include @@ -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 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(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 threads; std::atomic current(start); diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp index 7dc86bcb..0e519d00 100644 --- a/include/rabitqlib/index/symqg/detail/pipnn.hpp +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -208,15 +208,11 @@ inline std::vector cluster( while (!active.empty()) { std::vector> next(active.size()); std::vector> 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(active.size()); - ++index) { - const size_t i = static_cast(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 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(active.size()); + ++i) { + process_bucket(static_cast(i)); + } } std::vector new_active; for (size_t i = 0; i < active.size(); ++i) { diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index a68e8ffc..21e2f370 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -243,7 +243,7 @@ class QuantizedGraph { * 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(). */ diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index ff6e964b..a653b463 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/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); } diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index d783867f..3fbe3db4 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/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_(m, "MetricType") .value("L2", rabitqlib::METRIC_L2) diff --git a/sample/python/README.md b/sample/python/README.md index 3c828b2d..85977a04 100644 --- a/sample/python/README.md +++ b/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 diff --git a/sample/python/kmeans_clustering.py b/sample/python/kmeans_clustering.py index 94257da7..367ec8e7 100644 --- a/sample/python/kmeans_clustering.py +++ b/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()) diff --git a/src/utils/thread_count.cpp b/src/utils/thread_count.cpp new file mode 100644 index 00000000..2f558393 --- /dev/null +++ b/src/utils/thread_count.cpp @@ -0,0 +1,67 @@ +#include + +#include +#include +#include +#include +#include + +#include "rabitqlib/utils/tools.hpp" + +#if defined(__linux__) +#include +#endif + +namespace rabitqlib { + +size_t total_threads() { +#if defined(_OPENMP) && _OPENMP >= 201511 + if (omp_get_proc_bind() != omp_proc_bind_false && omp_get_num_places() > 0) { + // The runtime may pin the caller to one place. Count its whole partition, + // including overlapping places only once, rather than that single mask. + std::vector places(static_cast(omp_get_partition_num_places())); + omp_get_partition_place_nums(places.data()); + std::vector cpus; + for (int place : places) { + const size_t begin = cpus.size(); + cpus.resize(begin + static_cast(omp_get_place_num_procs(place))); + omp_get_place_proc_ids(place, cpus.data() + begin); + } + std::sort(cpus.begin(), cpus.end()); + cpus.erase(std::unique(cpus.begin(), cpus.end()), cpus.end()); + if (!cpus.empty()) { + return cpus.size(); + } + } +#elif defined(_OPENMP) && _OPENMP >= 201307 + if (omp_get_proc_bind() != omp_proc_bind_false) { + // OpenMP 4.0 has binding but no place-partition query routines. + const int count = omp_get_num_procs(); + if (count > 0) { + return static_cast(count); + } + } +#endif +#if defined(__linux__) + // The kernel's CPU mask can exceed CPU_SETSIZE on large systems. + std::vector mask(1); + while (true) { + const size_t bytes = mask.size() * sizeof(cpu_set_t); + if (sched_getaffinity(0, bytes, mask.data()) == 0) { + const int count = CPU_COUNT_S(bytes, mask.data()); + if (count > 0) { + return static_cast(count); + } + break; + } + if (errno != EINVAL) { + break; + } + mask.resize(mask.size() * 2); + } +#endif + const auto threads = std::thread::hardware_concurrency(); + return threads == 0 ? 1 : threads; +} + +} // namespace rabitqlib diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 144499b8..b19db71c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -40,6 +40,12 @@ target_link_libraries(rabitq_tests include(GoogleTest) gtest_discover_tests(rabitq_tests DISCOVERY_MODE PRE_TEST) +# Binding is initialized before main; test it in its own process. +add_test(NAME OpenMPBinding COMMAND rabitq_tests + --gtest_filter=ThreadCount.BoundParallelForUsesRuntimePlacementAndPropagatesErrors) +set_tests_properties(OpenMPBinding PROPERTIES + ENVIRONMENT "OMP_PROC_BIND=spread;OMP_PLACES=threads;OMP_NUM_THREADS=4") + # Substitute CPU predicates only in this executable to exercise real cached dispatch # on a CPU without SIMD. The static library's detector object is not linked here. add_executable(fastscan_no_simd_test dispatch/fastscan_no_simd_test.cpp) diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index cc20e8ce..2a6b4ead 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -13,6 +13,63 @@ EXAMPLES = Path(__file__).resolve().parents[2] / "sample" / "python" + +@pytest.mark.parametrize("binding", ["none", "threads", "cores", "overlap"]) +def test_native_cpu_count_respects_launch_affinity_and_openmp_places(binding): + if not hasattr(os, "sched_getaffinity"): + pytest.skip("Linux affinity is required") + cpus = sorted(os.sched_getaffinity(0))[:4] + expected = min(len(cpus), 2) if binding == "overlap" else len(cpus) + env = os.environ.copy() + for name in ("OMP_PROC_BIND", "OMP_PLACES", "GOMP_CPU_AFFINITY"): + env.pop(name, None) + if binding != "none": + env["OMP_PROC_BIND"] = "spread" + if binding == "overlap": + places = [cpus[0], cpus[0], cpus[-1]] + env["OMP_PLACES"] = ",".join(f"{{{cpu}}}" for cpu in places) + else: + env["OMP_PLACES"] = binding + code = """ +import os +import sys +cpus = {int(cpu) for cpu in sys.argv[1].split(',')} +os.sched_setaffinity(0, cpus) +from rabitqlib._rabitqlib import _available_cpu_count +assert _available_cpu_count() == int(sys.argv[2]), _available_cpu_count() +""" + result = subprocess.run( + [sys.executable, "-c", code, ",".join(map(str, cpus)), str(expected)], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("requested, expected", [(0, 2), (1, 1), (96, 2)]) +def test_clustering_example_uses_native_cpu_budget(monkeypatch, requested, expected): + monkeypatch.syspath_prepend(str(EXAMPLES)) + spec = importlib.util.spec_from_file_location( + "kmeans_clustering", EXAMPLES / "kmeans_clustering.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "_available_cpu_count", lambda: 2) + calls = [] + + def clustering(*args, **kwargs): + calls.append(kwargs["num_threads"]) + return SimpleNamespace( + train=lambda data: None, centroids=None, assignments=None + ) + + monkeypatch.setattr(module, "RaBitQKMeans", clustering) + module.cluster_data(np.zeros((3, 64)), 2, "l2", requested, "rabitq") + assert calls == [expected] + + # Fail on an accidental import, even when a runtime combination happens to coexist. BOOTSTRAP = """ import runpy diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py index cc8d62b0..5c08b757 100644 --- a/tests/python/test_qgkmeans.py +++ b/tests/python/test_qgkmeans.py @@ -8,6 +8,17 @@ K = 33 +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("spherical", [False, True]) +def test_refilled_duplicate_centroids_have_zero_shift(model_type, spherical): + data = np.full((65, 64), 0.125, dtype=np.float32) + model = model_type(64, 33, niter=1, spherical=spherical, num_threads=2) + model.train(data) + assert model.iteration_stats[0].nsplit == 32 + assert model.iteration_stats[0].shift == 0.0 + np.testing.assert_array_equal(model.centroids, data[:33]) + + def qgkmeans(**kwargs): parameters = { "niter": 2, diff --git a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp new file mode 100644 index 00000000..51682852 --- /dev/null +++ b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp @@ -0,0 +1,87 @@ +#include + +#include +#include +#include + +#include "rabitqlib/clustering/detail/kmeans.hpp" +#include "rabitqlib/clustering/qgkmeans.hpp" + +namespace rabitqlib::clustering::detail { +namespace { +struct RecordedParameters : qgkmeans::QGKMeansParameters { + // Owned by the test for the duration of synchronous training. + std::vector* initial_centroids = nullptr; + size_t occupied_clusters = 0; +}; + +class RecordedAssignment { + size_t dim_; + size_t clusters_; + size_t points_; + RecordedParameters parameters_; + + public: + static constexpr const char* kName = "RecordedAssignment"; + static void validate_parameters(size_t, const RecordedParameters&) {} + + RecordedAssignment( + size_t d, + size_t k, + const float*, + size_t n, + uint32_t, + const RecordedParameters& parameters + ) + : dim_(d), clusters_(k), points_(n), parameters_(parameters) {} + + void assign(const float* centroids, PID* labels, float* distances, bool final = false) { + if (!final) { + parameters_.initial_centroids->assign(centroids, centroids + clusters_ * dim_); + } + for (size_t point = 0; point < points_; ++point) { + labels[point] = static_cast(point % parameters_.occupied_clusters); + distances[point] = 0; + } + } +}; + +TEST(CentroidUpdateTest, ShiftIncludesNormalizedAndRefilledCentroids) { + constexpr size_t kDim = 64, kClusters = 3, kPoints = 12; + std::vector data(kPoints * kDim, 0); + for (size_t point = 0; point < kPoints; ++point) { + data[point * kDim + point] = 1; + } + for (bool spherical : {false, true}) { + for (size_t empty : {0U, 1U, 2U}) { + for (uint32_t threads : {1U, 4U}) { + SCOPED_TRACE( + ::testing::Message() << spherical << '/' << empty << '/' << threads + ); + std::vector initial; + RecordedParameters parameters; + parameters.niter = 1; + parameters.spherical = spherical; + parameters.num_threads = threads; + parameters.initial_centroids = &initial; + parameters.occupied_clusters = kClusters - empty; + LloydKMeans model( + kDim, kClusters, parameters + ); + model.train(kPoints, data.data()); + ASSERT_EQ(initial.size(), model.centroids.size()); + double expected = 0; + for (size_t i = 0; i < initial.size(); ++i) { + const double delta = + static_cast(model.centroids[i]) - initial[i]; + expected += delta * delta; + } + ASSERT_EQ(model.iteration_stats.size(), 1U); + EXPECT_EQ(model.iteration_stats.front().nsplit, empty); + EXPECT_NEAR(model.iteration_stats.front().shift, expected, 1e-12); + } + } + } +} +} // namespace +} // namespace rabitqlib::clustering::detail diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp index fb78555e..4cd9a48e 100644 --- a/tests/unit/rabitqlib/index/pipnn_test.cpp +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -130,5 +130,29 @@ TEST(PipnnTest, RejectsInvalidConfiguration) { EXPECT_THROW(build_initial_graph(data.data(), 33, 7, 16), std::invalid_argument); EXPECT_THROW(build_initial_graph(data.data(), 32, 7, 32), std::invalid_argument); } + +TEST(PipnnTest, RootTileTailAndRecursiveBucketsCoverDuplicateInputs) { + constexpr size_t kCount = 4097, kDim = 7; + std::vector data(kCount * kDim, 0.5F); + for (size_t threads : {1U, 2U}) { + pipnn_impl::ScratchPool scratch(kCount, kDim, threads); + auto leaves = + pipnn_impl::cluster(data.data(), kCount, kDim, METRIC_L2, threads, scratch); + std::vector memberships(kCount, 0); + for (auto& leaf : leaves) { + EXPECT_LE(leaf.size(), pipnn_impl::kLeafSize); + std::sort(leaf.begin(), leaf.end()); + EXPECT_EQ(std::unique(leaf.begin(), leaf.end()), leaf.end()); + for (PID id : leaf) { + ASSERT_LT(id, kCount); + ++memberships[id]; + } + } + for (size_t membership : memberships) { + EXPECT_GE(membership, 1U); + EXPECT_LE(membership, 30U); + } + } +} } // namespace } // namespace rabitqlib::symqg::detail diff --git a/tests/unit/rabitqlib/utils/thread_count_test.cpp b/tests/unit/rabitqlib/utils/thread_count_test.cpp index a45e9f68..932f0886 100644 --- a/tests/unit/rabitqlib/utils/thread_count_test.cpp +++ b/tests/unit/rabitqlib/utils/thread_count_test.cpp @@ -1,29 +1,108 @@ #include +#include + +#if defined(__linux__) +#include +#endif #include #include #include +#include #include #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/utils/tools.hpp" -TEST(ThreadCount, AutomaticAndOversizedRequestsUseHardwareCount) { - const size_t hardware = rabitqlib::total_threads(); - ASSERT_GE(hardware, 1U); - EXPECT_EQ(rabitqlib::resolve_num_threads(0), hardware); +TEST(ThreadCount, AutomaticAndOversizedRequestsUseAvailableCount) { + const size_t available = rabitqlib::total_threads(); + ASSERT_GE(available, 1U); + EXPECT_EQ(rabitqlib::resolve_num_threads(0), available); EXPECT_EQ(rabitqlib::resolve_num_threads(1), 1U); - EXPECT_EQ(rabitqlib::resolve_num_threads(hardware), hardware); - EXPECT_EQ(rabitqlib::resolve_num_threads(hardware + 1), hardware); - EXPECT_EQ(rabitqlib::resolve_num_threads(std::numeric_limits::max()), hardware); - if (hardware > 1) { - EXPECT_EQ(rabitqlib::resolve_num_threads(hardware - 1), hardware - 1); + EXPECT_EQ(rabitqlib::resolve_num_threads(available), available); + EXPECT_EQ(rabitqlib::resolve_num_threads(available + 1), available); + EXPECT_EQ( + rabitqlib::resolve_num_threads(std::numeric_limits::max()), available + ); + if (available > 1) { + EXPECT_EQ(rabitqlib::resolve_num_threads(available - 1), available - 1); + } +} + +#if defined(__linux__) +TEST(ThreadCount, AutomaticAndExplicitRequestsRespectAffinity) { +#if defined(_OPENMP) && _OPENMP >= 201307 + if (omp_get_proc_bind() != omp_proc_bind_false) { + GTEST_SKIP() << "OpenMP owns placement when binding is enabled"; } +#endif + cpu_set_t original; + if (sched_getaffinity(0, sizeof(original), &original) != 0) { + GTEST_SKIP() << "Cannot read affinity"; + } + struct RestoreAffinity { + cpu_set_t mask; + ~RestoreAffinity() { EXPECT_EQ(sched_setaffinity(0, sizeof(mask), &mask), 0); } + } restore{original}; + cpu_set_t restricted; + CPU_ZERO(&restricted); + size_t available = 0; + for (size_t cpu = 0; cpu < CPU_SETSIZE && available < 2; ++cpu) { + if (CPU_ISSET(cpu, &original)) { + CPU_SET(cpu, &restricted); + ++available; + } + } + ASSERT_GT(available, 0U); + ASSERT_EQ(sched_setaffinity(0, sizeof(restricted), &restricted), 0); + EXPECT_EQ(rabitqlib::total_threads(), available); + for (size_t requested : {size_t{0}, size_t{1}, std::numeric_limits::max()}) { + const size_t expected = requested == 1 ? 1 : available; + EXPECT_EQ(rabitqlib::resolve_num_threads(requested), expected); + rabitqlib::ivf::parallel_for(0, 19, requested, [&](size_t, size_t worker) { + EXPECT_LT(worker, expected); + }); + } +} +#endif + +TEST(ThreadCount, BoundParallelForUsesRuntimePlacementAndPropagatesErrors) { +#if defined(_OPENMP) && _OPENMP >= 201511 + if (omp_get_proc_bind() == omp_proc_bind_false) { + GTEST_SKIP() << "Run with OpenMP binding enabled"; + } + const size_t available = rabitqlib::total_threads(); + ASSERT_GT(available, 0U); + // A binding policy can be set even when the runtime exposes no CPU places + // (for example, LLVM libomp on macOS). + const int places = omp_get_num_places(); + const size_t threads = std::min(available, 4); + rabitqlib::ivf::parallel_for(0, 257, threads, [&](size_t, size_t worker) { + EXPECT_LT(worker, threads); + if (places > 0) { + EXPECT_GE(omp_get_place_num(), 0); + EXPECT_LT(omp_get_place_num(), places); + } else { + EXPECT_EQ(omp_get_place_num(), -1); + } + if (threads > 1) { + EXPECT_GT(omp_get_level(), 0); + } + }); + EXPECT_THROW( + rabitqlib::ivf::parallel_for( + 0, 257, threads, [](size_t, size_t) { throw std::runtime_error("failed"); } + ), + std::runtime_error + ); +#else + GTEST_SKIP() << "OpenMP places require OpenMP 4.5"; +#endif } TEST(ThreadCount, ParallelForVisitsEachItemOnceWithBoundedWorkerIds) { - const size_t hardware = rabitqlib::total_threads(); - const size_t count = 2 * hardware + 3; + const size_t available = rabitqlib::total_threads(); + const size_t count = 2 * available + 3; for (size_t requested : {size_t{0}, size_t{1}, std::numeric_limits::max()}) { SCOPED_TRACE(requested); std::vector> visits(count); @@ -31,7 +110,7 @@ TEST(ThreadCount, ParallelForVisitsEachItemOnceWithBoundedWorkerIds) { visited.store(0); } rabitqlib::ivf::parallel_for(0, count, requested, [&](size_t i, size_t worker) { - EXPECT_LT(worker, requested == 1 ? 1U : hardware); + EXPECT_LT(worker, requested == 1 ? 1U : available); visits[i].fetch_add(1); }); for (const auto& visited : visits) {