diff --git a/.gitignore b/.gitignore index 80062aa4..9fcd60e7 100644 --- a/.gitignore +++ b/.gitignore @@ -48,3 +48,6 @@ __pycache__/ # macOS garbages .DS_Store + +# Generated documentation +/docs/site/ diff --git a/AGENTS.md b/AGENTS.md index 22aef9fe..33a50d54 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,6 +17,7 @@ factors estimate L2 distance or inner product. | --- | --- | | `include/rabitqlib/quantization/` | Encoding, packing, reconstruction, and byte layouts | | `include/rabitqlib/fastscan/` | FastScan interfaces and high-accuracy scanning | +| `include/rabitqlib/clustering/` | RaBitQKMeans flat assignment, QGKMeans graph assignment, and shared Lloyd training | | `include/rabitqlib/index/{ivf,hnsw,symqg}/` | Index construction, persistence, and search | | `include/rabitqlib/index/{query,estimator}.hpp` | Query state and distance estimation | | `include/rabitqlib/simd/`, `src/simd/` | Kernel declarations, implementations, and dispatch | diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9561e667..d22554b6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -175,7 +175,6 @@ another. Each task can be a separate PR. | --- | --- | --- | | Explain Python search results | [Quick start](docs/docs/quick_start.md), [IVF binding](python_bindings/ivf_bindings.cpp), and [IVF tests](tests/python/test_ivf.py) | A short explanation covers result shapes, how IDs map to input rows, and what L2 distances represent, checked against the implementation. The example runs and the strict docs build passes. | | Add a self-contained IVF save/load example | [Python examples](sample/python/), [README quick start](README.md#python-quick-start), and [IVF tests](tests/python/test_ivf.py) | A deterministic script builds from synthetic data, saves to a temporary directory, reloads, and checks that search IDs and distances match. It needs no dataset download, cleans up its temporary files, passes Python checks, and is linked from the README. | -| Explain the quick start's cluster assignment | [Quick start](docs/docs/quick_start.md) and [IVF guide](docs/docs/index/ivf.md) | The tutorial explains that its round-robin assignment is for a small runnable example, explains why real IVF workloads use clustering, and links to the existing clustering workflow. The strict docs build passes. | ## Feedback and further reading diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index efe9f676..a75b9d69 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -310,10 +310,11 @@ covers its arithmetic kernels, not every scalar loop in construction and search. | IVF and float SymphonyQG batch correction | Complete estimator runs in the selected backend; non-float template paths remain generic | | FHT/Kac rotation | Complete rotation and scaling run in selected ISA translation units; shared AVX intrinsics preserve the FFHT butterfly order | | Float matrix rotation and PiPNN construction | Matrix products, row norms, and lower-triangle pairwise distances use isolated matrix backends | +| RaBitQKMeans and QGKMeans centroid updates | Indexed float32 rows accumulate into float64 sums through the matrix backends, preserving per-cluster point order without additional workspace | | HNSW search and IVF centroid routing | Cached HNSW search selection; centroid routing uses the common raw-distance dispatcher | | Quantization orchestration, reconstruction, non-float utilities | Template/control code remains generic; no blanket native tuning or reduction-order rewrite | | Graph scheduling, candidate queues, I/O, allocation, random initialization | Generic control code; IVF one-bit candidate insertion stays in a small compiled function to avoid inlining-induced register spills; thread scheduling and seeds remain caller-owned | -| Example KMeans training | Uses external FAISS, whose build and dispatch are independent of this package | +| Example KMeans training | Explicit RaBitQKMeans flat or QGKMeans graph assignment; FAISS is an optional comparison dependency only | Portable wheels disable `RABITQ_ENABLE_NATIVE_OPTIMIZATION`. ARM64 builds exclude all x86 source groups and use standard AArch64 NEON intrinsics, without Apple-only APIs. diff --git a/README.md b/README.md index d31d8c6a..d289c1c7 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@

A research-backed C++17 library with Python bindings for 1-bit and multi-bit
- vector quantization, IVF, HNSW, and SymphonyQG. + vector quantization, IVF, HNSW, SymphonyQG, and clustering.

@@ -63,17 +63,16 @@ Build and search a small IVF index using synthetic data: ```python import numpy as np -from rabitqlib import IvfIndex +from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans rng = np.random.default_rng(42) data = rng.standard_normal((500, 64)).astype(np.float32) queries = rng.standard_normal((5, 64)).astype(np.float32) -# Assign vectors to five clusters and calculate their centroids. -cluster_ids = (np.arange(len(data)) % 5).astype(np.uint32) -centroids = np.stack( - [data[cluster_ids == cluster].mean(axis=0) for cluster in range(5)] -).astype(np.float32) +clustering = RaBitQKMeans( + 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact +) +clustering.train(data) index = IvfIndex( dim=64, @@ -82,7 +81,7 @@ index = IvfIndex( nbits=4, metric="l2", ) -index.build(data, centroids, cluster_ids) +index.build(data, clustering.centroids, clustering.assignments) ids, distances = index.search(queries, k=10, nprobe=5) print(ids.shape, distances.shape) # (5, 10) (5, 10) @@ -93,16 +92,20 @@ 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. Omitting `num_threads` in Python still defaults to one thread. +used. Python index methods default to one thread when `num_threads` is omitted. -Python bindings are also available for `HnswIndex` and `SymqgIndex`. See the -[Python examples](sample/python/README.md) for index construction, querying, and -index persistence. IVF and HNSW examples run Faiss clustering in a separate process, -then pass saved clusters to RaBitQ indexing so their OpenMP runtimes stay separate. +IVF and HNSW examples save clusters for reuse across index configurations. +Choose RaBitQKMeans for flat assignment or QGKMeans for graph assignment; +FAISS is needed only for the optional clustering comparison. 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. +See the [Python examples](sample/python/README.md) for IVF, HNSW, and SymphonyQG. +For clustering, use [RaBitQKMeans](docs/docs/clustering.md#rabitqkmeans) for small cluster +counts or [QGKMeans](docs/docs/clustering.md#qgkmeans) for graph assignment. The +[FAISS comparison](sample/python/compare_with_faiss.py) benchmarks both methods. +

Build the Python bindings from source @@ -145,6 +148,7 @@ depend on the dataset, index configuration, and search parameters. | **Accurate estimates** | An asymptotically optimal theoretical error bound supports reliable ordering and reranking. | | **Native CPU backends** | Runtime AVX2/AVX-512 selection on x86-64; NEON distance, packed-code, FastScan, rotation, query preparation, and HNSW search kernels on ARM64, with portable scalar fallbacks. | | **Ready for ANN search** | Use the quantizer directly or build complete IVF, HNSW, and [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) indexes. | +| **Native clustering** | Choose [RaBitQKMeans](docs/docs/clustering.md#rabitqkmeans) with flat RaBitQ assignment or [QGKMeans](docs/docs/clustering.md#qgkmeans) with SymphonyQG assignment; neither needs an external k-means package. | The library supports Euclidean distance and inner product. Cosine search is available by normalizing vectors before using inner product. diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md new file mode 100644 index 00000000..57d16c04 --- /dev/null +++ b/docs/docs/clustering.md @@ -0,0 +1,189 @@ +# Clustering + +RaBitQ Library provides two implementations of Lloyd's k-means. Both compute +centroid means from the original float32 data, accumulate in float64, and require +no external clustering package. Choose the class explicitly; neither switches +algorithms based on the cluster count. + +## Choose a method + +| Class | Assignment | Recommended use | Cluster count | +| --- | --- | --- | --- | +| [RaBitQKMeans](#rabitqkmeans) | Flat RaBitQ scan | Small cluster counts | `n >= k >= 1` | +| [QGKMeans](#qgkmeans) | SymphonyQG centroid graph | Larger cluster counts | `n >= k > graph_degree` (default 32) | + +## Python + +Both classes have the same training and result interface: + +```python +from rabitqlib import QGKMeans, RaBitQKMeans + +kmeans = RaBitQKMeans(d, 16, niter=25, num_threads=8) +# For larger cluster counts, use instead: +# kmeans = QGKMeans(d, 4096, niter=25, num_threads=8) +last_iteration_obj = kmeans.train(x) + +centroids = kmeans.centroids # shape (k, d) +assignments = kmeans.assignments # shape (n,) +distances = kmeans.distances # shape (n,) +final_obj = kmeans.final_obj +``` + +`x` has shape `(n, d)` and is converted to a C-contiguous float32 array. + +## C++ + +```cpp +#include +#include + +rabitqlib::rabitqkmeans::RaBitQKMeansParameters parameters; +parameters.niter = 25; +parameters.num_threads = 8; +rabitqlib::rabitqkmeans::RaBitQKMeans kmeans(dimension, 16, parameters); + +// For graph assignment, use qgkmeans::QGKMeansParameters and qgkmeans::QGKMeans. +kmeans.train(num_points, data); // row-major float32 input + +const auto& centroids = kmeans.centroids; +const auto& assignments = kmeans.assignments; +const auto& distances = kmeans.distances; +const double final_obj = kmeans.final_obj; +``` + +Complete examples: [RaBitQKMeans](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/rabitqkmeans.cpp) +(`bin/rabitqkmeans`) and [QGKMeans](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/qgkmeans.cpp) +(`bin/qgkmeans`). Run either without arguments for a synthetic demo, or save +clusters for the C++ indexing examples: + +```bash +./bin/qgkmeans data.fvecs 4096 centroids.fvecs labels.ivecs +./bin/rabitqkmeans data.fvecs 16 centroids.fvecs labels.ivecs +``` + +File mode uses exact final assignment and the default training parameters. +Append `ip` for spherical clustering of normalized input; the default is `l2`. +`example.sh` runs the C++ clustering and indexing workflow without Python. + +## Parameters and input + +Set Python keyword arguments or fields of the corresponding C++ parameter type: + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `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. | +| `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. | + +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 +mutate the input; provide normalized vectors for spherical mode. + +Empty clusters are reseeded from distant points; duplicate data can still leave +empty clusters in the final assignment. Thread scheduling can affect graph +construction and floating-point reductions, so a fixed seed does not guarantee +identical multithreaded results. + +## Method details + +### RaBitQKMeans + +The flat assignment design is inspired by +[SuperKMeans](https://github.com/cwida/SuperKMeans). Training points are rotated +and encoded once around a fixed dataset mean. Each centroid becomes a query +that scans all packed point batches, and each point receives the centroid with +the best estimated distance. + +Each iteration compares all `n * k` point-centroid pairs. The batch size of 32 +is not a cluster-count limit. This method uses one-bit point codes throughout; +it has no graph parameters or `quantization_bits` setting. + +### QGKMeans + +QGKMeans builds a SymphonyQG graph over the centroids and uses previous +assignments as search hints. Its additional parameters are: + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `quantization_bits` | `0` | Graph centroid storage: raw, or `4`/`8` bits. | +| `graph_degree` | `32` | A positive multiple of 32, strictly smaller than `k`. | +| `ef_build`, `ef_search` | `240`, `16` | Positive graph construction and search window sizes. | +| `graph_build_iterations` | `1` | PiPNN initialization followed by one refinement; values ≥2 request that many graph passes, ending in refinement. Independent of `niter`. | + +The seed also controls fallback neighbors; PiPNN uses fixed internal seeds. +`QGAssigner` in `qgkmeans.hpp` supports repeated vector-to-centroid assignment +without the clustering loop. + +## Results and final assignment + +Python result properties are `None` until a successful fit; a failed fit +preserves the previous successful results. + +`distances` contains distances to the returned **assigned** centroids, computed +from the original float32 vectors and centroids. For squared L2, `final_obj` is +their sum (WCSS); spherical mode sums `1 - inner_product`. `obj` and +`iteration_stats` record objectives before each centroid update; Python +`train()` returns the last such objective. + +`FinalAssignmentMode.Approximate` (default) uses the selected class's assignment +algorithm. The legacy name `FinalAssignmentMode.SymphonyQG` remains an alias for +`Approximate`. Labels need not identify the exact nearest centroids, but reported +distances and WCSS remain exact for those labels. Quantized QGKMeans final +assignment retains a previous label when it is closer. Approximate training +objectives need not decrease monotonically. + +To require exact nearest-centroid **final** assignments with either class: + +```python +from rabitqlib import FinalAssignmentMode, RaBitQKMeans + +kmeans = RaBitQKMeans(d, 16, final_assignment=FinalAssignmentMode.Exact) +kmeans.train(x) +``` + +In C++, set `parameters.final_assignment` to +`rabitqlib::rabitqkmeans::FinalAssignmentMode::Exact` or +`rabitqlib::qgkmeans::FinalAssignmentMode::Exact`; both names refer to the same +enum. Training still uses the selected approximate assignment algorithm. + +Exact final assignment screens all centroids with blocked GEMM and refines +possible winners in float64; ties select the smallest centroid index. Returned +distances remain float32. Work is proportional to `n * k * d`; large coordinate +offsets or near ties can require more refinement. + +## Compare with FAISS + +The [comparison script](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/python/compare_with_faiss.py) +compares both methods with FAISS and requires `faiss-cpu`. Defaults are one +million 256-dimensional vectors, 4,000 centroids, and up to 25 iterations. +Load data with `--input data.npy`, or `--input data.bin --d D` for raw row-major +float32 data. Choose a thread count within the available hardware limit. + +```bash +python sample/python/compare_with_faiss.py --threads 48 --num-seeds 2 --repeats 2 +``` + +All three methods receive the same full dataset, iteration limit, cluster count, +thread budget, and seed values; FAISS subsampling is disabled. This command runs +four fits per method with rotating order. Different initializers mean identical +seeds need not produce identical centroids. + +Only `train()` is timed, including the library methods' default approximate +final assignment. The script reports actual iterations, timing medians/ranges, +exhaustive nearest-centroid objectives, and both library methods' returned-label +objectives and label mismatches. Equal-distance ties can produce mismatches +without changing the objective. `--spherical` normalizes the input copy and +compares sums of `1 - inner_product`. + +Separate worker processes isolate OpenMP runtimes and share a temporary float32 +input file; allow about `4 * n * d` bytes of temporary disk space. Input staging, +process startup, and exhaustive evaluation are outside training time. For Linux +CPU placement, set `OMP_PLACES=threads OMP_PROC_BIND=close`; on NUMA systems, +use `numactl --cpunodebind=0 --membind=0` with a thread count that fits the node. diff --git a/docs/docs/index.md b/docs/docs/index.md index 8a1edcca..b9bd8bac 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -46,16 +46,16 @@ Build an IVF index and search a batch of queries: ```python import numpy as np -from rabitqlib import IvfIndex +from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans rng = np.random.default_rng(42) data = rng.standard_normal((500, 64)).astype(np.float32) queries = rng.standard_normal((5, 64)).astype(np.float32) -cluster_ids = (np.arange(len(data)) % 5).astype(np.uint32) -centroids = np.stack( - [data[cluster_ids == cluster].mean(axis=0) for cluster in range(5)] -).astype(np.float32) +clustering = RaBitQKMeans( + 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact +) +clustering.train(data) index = IvfIndex( dim=64, @@ -64,7 +64,7 @@ index = IvfIndex( nbits=4, metric="l2", ) -index.build(data, centroids, cluster_ids) +index.build(data, clustering.centroids, clustering.assignments) ids, distances = index.search(queries, k=10, nprobe=5) print(ids.shape, distances.shape) # (5, 10) (5, 10) diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 1ecb6afd..f3e363ab 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -28,7 +28,7 @@ HierarchicalNSW::construct(size_t cluster_num, - **data**: Pointer to the raw data vectors. - **data_num**: The number of data vectors. -- **centroids**: Centroids computed by K-means clustering on the raw data vectors (we recommend `cluster_num = 16`). +- **centroids**: Centroids computed by [RaBitQKMeans](../clustering.md#rabitqkmeans); the examples use `cluster_num = 16`. - **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`. - **num_threads**: Number of threads to use (default: 0, which auto-selects). - **faster**: If `true`, enables the faster quantizer. diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 547df2b3..5ef43825 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -10,24 +10,19 @@ width, number of clusters, and `nprobe`. The algorithm includes two phases: indexing and querying. ## Index Construction -The first step is to run a clustering algorithm to partition raw data vectors (`*.fvecs` format) into different buckets. -The algorithm performs KMeans clustering on raw vectors based on [Faiss](https://github.com/facebookresearch/faiss) (see `python/ivf.py`). -To run the algorithm, you need to execute the command in `shell`: + +Use [QGKMeans](../clustering.md#qgkmeans) to partition the input vectors (`.fvecs`). +Build the C++ examples, then run `bin/qgkmeans` to save centroids and cluster IDs. +For small cluster counts, use `bin/rabitqkmeans` for +[RaBitQKMeans](../clustering.md#rabitqkmeans) flat assignment. Both use exact final +assignment. The optional metric defaults to `l2`; append `ip` for normalized input. + +For example, split SIFT into 4,096 clusters using squared L2 distance: ```shell -python python/ivf.py /path/to/raw/data \ - number_of_clusters \ - /path/to/output/centroids \ - /path/to/output/cluster_ids \ - distance_metric -``` -For example, the following command splits the sift vector data into 4096 clusters using Euclidean (l2) distance: -```shell -python python/ivf.py /data/sift/sift_base.fvecs \ - 4096 \ - /data/sift/sift_centroids_4096_l2.fvecs \ - /data/sift/sift_clusterids_4096_l2.ivecs \ - l2 +./bin/qgkmeans /data/sift/sift_base.fvecs 4096 \ + /data/sift/sift_centroids_4096_l2.fvecs \ + /data/sift/sift_clusterids_4096_l2.ivecs ``` After files are prepared, you need to load them into memory: @@ -72,7 +67,7 @@ void IVF::construct( ``` - **data**: Pointer to the raw data vectors. -- **centroids**: Centroids computed by K-means clustering on the raw data vectors (we recommend to tune cluster_num around 4 * the square root of the dataset following Faiss). +- **centroids**: Clustering centroids; tune `cluster_num` for the dataset and search budget. - **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`. - **faster**: If true, enable fast implementations for RaBitQ (By default, it is set as `false` to pursue better accuracy.). - **num_threads**: Maximum number of OpenMP threads used to quantize clusters. diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 85316697..21d9b47b 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -10,7 +10,7 @@ Set `quantization_bits=0` for raw vectors (default), or `4`/`8` for QG-quant. Quantized vectors share a rotated global centroid. Refinement temporarily reconstructs source vectors and estimates distances to stored target codes; reverse edges are rescored because these estimates are directional. Raw refinement -uses owned raw vectors. No full reconstructed-vector cache is retained. +uses owned raw vectors. No full reconstructed-vector cache is retained by default. ## Index Construction @@ -39,29 +39,38 @@ window; `ef` controls the query search window. Python defaults to one thread. ### C++ -The C++ API uses the float-only `rabitqlib::symqg::QuantizedGraph` and `QGBuilder`. -`QuantizedGraph` retains the C++ template API, with its implementation -compiled in `src/index/qg.cpp`. C++17 variable declarations can also omit `` -through class template argument deduction: +The C++ API provides `rabitqlib::symqg::QuantizedGraph` and `QGBuilder`. +Only float vectors are supported; C++17 declarations may omit `` through +class template argument deduction: ```cpp QuantizedGraph( size_t num, size_t dim, size_t max_deg, MetricType metric_type = METRIC_L2, RotatorType rotator_type = RotatorType::FhtKacRotator, - size_t quantization_bits = 0 + size_t quantization_bits = 0, uint32_t seed = std::random_device{}() ); QGBuilder( QuantizedGraph& index, uint32_t ef_build, const float* data, size_t num_threads = std::numeric_limits::max(), - QGInitialization init = QGInitialization::PiPNN + QGInitialization init = QGInitialization::PiPNN, + uint32_t seed = std::random_device{}(), bool cache_vectors = false ); ``` -`data` contains `num * dim` floats; `max_deg` has the same constraints as Python's -`max_degree`. C++ defaults to all available threads. Pass -`QGInitialization::Random` as the final builder argument to use random initialization. -The builder handles initialization internally. +`data` contains `num * dim` floats; `max_deg` has Python's `max_degree` +constraints. C++ defaults to all available threads. `build()` runs one refinement +for PiPNN or three passes for `QGInitialization::Random`; `build(n)` requests +`n >= 2` passes. The older overload with a numeric fifth argument uses that value +as the random-initialization seed. + +`builder.reset(data)` reinitializes vectors of the same shape and reuses workspace. +Input may be released after reset returns; rebuild before querying or saving. +`cache_vectors=true` retains transformed vectors and query factors +(`4 * num * padded_dim` bytes plus factors). [QGKMeans](../clustering.md#qgkmeans) enables +this for its centroid graph; the cache is not persisted. The graph seed controls +rotation, and the builder seed controls random initialization and fallback edges. +PiPNN uses fixed internal seeds. ```cpp using namespace rabitqlib::symqg; diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index 358253a3..9d5297e5 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -55,17 +55,16 @@ require a dataset download: ```python import numpy as np -from rabitqlib import IvfIndex +from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans rng = np.random.default_rng(42) data = rng.standard_normal((500, 64)).astype(np.float32) queries = rng.standard_normal((5, 64)).astype(np.float32) -# Assign vectors to five clusters and calculate their centroids. -cluster_ids = (np.arange(len(data)) % 5).astype(np.uint32) -centroids = np.stack( - [data[cluster_ids == cluster].mean(axis=0) for cluster in range(5)] -).astype(np.float32) +clustering = RaBitQKMeans( + 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact +) +clustering.train(data) index = IvfIndex( dim=64, @@ -74,7 +73,7 @@ index = IvfIndex( nbits=4, metric="l2", ) -index.build(data, centroids, cluster_ids) +index.build(data, clustering.centroids, clustering.assignments) ids, distances = index.search(queries, k=10, nprobe=5) print(ids.shape, distances.shape) # (5, 10) (5, 10) @@ -101,9 +100,10 @@ The `metric` argument accepts `"l2"` and `"ip"` (also spelled `"innerproduct"`). To search by cosine similarity, normalize database and query vectors first and use `metric="ip"`. -Python bindings are also available for `HnswIndex` and `SymqgIndex`. The -[Python examples](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/sample/python) -cover construction, querying, and index persistence. +See the [Python examples](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/sample/python) +for IVF, HNSW, and SymphonyQG. For clustering, choose +[RaBitQKMeans](clustering.md#rabitqkmeans) for flat assignment, recommended for small cluster +counts, or [QGKMeans](clustering.md#qgkmeans) for graph assignment.
Build the Python bindings from source diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 9f931cb2..894d3ac6 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -24,6 +24,7 @@ nav: - SymphonyQG: index/qg.md - Integrations: - Zvec: integrations/zvec.md + - Clustering: clustering.md theme: name: material diff --git a/example.sh b/example.sh index 499c02ff..1582232c 100755 --- a/example.sh +++ b/example.sh @@ -14,15 +14,15 @@ tar -xzvf ./data/gist/gist.tar.gz -C ./data/gist ./bin/symqg_querying ./data/gist/symqg_32.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs -# indexing and querying for RabitQ+ with ivf, please refer to python/ivf.py for more information about clustering -python ./python/ivf.py ./data/gist/gist_base.fvecs 4096 ./data/gist/gist_centroids_4096.fvecs ./data/gist/gist_clusterids_4096.ivecs +# Cluster with QGKMeans, then build and query an IVF index. +./bin/qgkmeans ./data/gist/gist_base.fvecs 4096 ./data/gist/gist_centroids_4096.fvecs ./data/gist/gist_clusterids_4096.ivecs ./bin/ivf_rabitq_indexing ./data/gist/gist_base.fvecs ./data/gist/gist_centroids_4096.fvecs ./data/gist/gist_clusterids_4096.ivecs 3 ./data/gist/ivf_4096_3.index ./bin/ivf_rabitq_querying ./data/gist/ivf_4096_3.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs -# indexing and querying for RabitQ+ with hnsw, do clustering first -python ./python/ivf.py ./data/gist/gist_base.fvecs 16 ./data/gist/gist_centroids_16.fvecs ./data/gist/gist_clusterids_16.ivecs +# Cluster with RaBitQKMeans, then build and query an HNSW index. +./bin/rabitqkmeans ./data/gist/gist_base.fvecs 16 ./data/gist/gist_centroids_16.fvecs ./data/gist/gist_clusterids_16.ivecs ./bin/hnsw_rabitq_indexing ./data/gist/gist_base.fvecs ./data/gist/gist_centroids_16.fvecs ./data/gist/gist_clusterids_16.ivecs 16 200 5 ./data/gist/hnsw_5.index diff --git a/example_py.sh b/example_py.sh index 661e7a2b..d47b17c3 100755 --- a/example_py.sh +++ b/example_py.sh @@ -1,5 +1,5 @@ # compiling -pip install . +python -m pip install . # Download the dataset mkdir -p ./data/gist @@ -14,14 +14,14 @@ python ./sample/python/symqg_indexing.py --max-degree 32 --ef-construction 400 python ./sample/python/symqg_querying.py ./data/gist/symqg_32.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs # indexing and querying for RabitQ+ with ivf -python ./sample/python/faiss_clustering.py --num-clusters 4096 ./data/gist/gist_base.fvecs ./data/gist/clusters_4096_l2.npz +python ./sample/python/kmeans_clustering.py --method qg --num-clusters 4096 ./data/gist/gist_base.fvecs ./data/gist/clusters_4096_l2.npz python ./sample/python/ivf_rabitq_indexing.py --total-bits 5 --clusters ./data/gist/clusters_4096_l2.npz ./data/gist/gist_base.fvecs ./data/gist/ivf_4096_5.index python ./sample/python/ivf_rabitq_querying.py ./data/gist/ivf_4096_5.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs # indexing and querying for RabitQ+ with hnsw -python ./sample/python/faiss_clustering.py --num-clusters 16 ./data/gist/gist_base.fvecs ./data/gist/clusters_16_l2.npz +python ./sample/python/kmeans_clustering.py --method rabitq --num-clusters 16 ./data/gist/gist_base.fvecs ./data/gist/clusters_16_l2.npz python ./sample/python/hnsw_rabitq_indexing.py --total-bits 5 --clusters ./data/gist/clusters_16_l2.npz --degree 16 --ef-construction 200 ./data/gist/gist_base.fvecs ./data/gist/hnsw_5.index diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp new file mode 100644 index 00000000..ef783719 --- /dev/null +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -0,0 +1,594 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::clustering { + +enum class FinalAssignmentMode : uint8_t { + Approximate = 0, + SymphonyQG = Approximate, // Compatibility name for the graph assignment path. + Exact = 1, +}; + +struct KMeansIterationStats { + size_t iteration = 0; + double obj = 0; + double shift = 0; + size_t nsplit = 0; +}; + +namespace detail { + +[[nodiscard]] inline bool should_terminate_by_objective( + double previous_objective, double objective, double early_stop_threshold +) { + if (previous_objective == 0.0) { + return objective == 0.0; + } + const double change = + std::abs(previous_objective - objective) / std::abs(previous_objective); + return change <= early_stop_threshold; +} + +// Inputs have already passed clustering validation. GEMM only screens candidates; +// final distances and tie breaking retain the scalar double-precision arithmetic. +inline void exact_assign( + const float* x, + const float* centroids, + size_t n, + size_t d, + size_t k, + bool spherical, + uint32_t num_threads, + PID* assignments, + float* distances +) { + constexpr size_t kPointBlock = 64; + constexpr size_t kCentroidBlock = 256; + const size_t blocks = 1 + (n - 1) / kPointBlock; + const auto threads = static_cast(std::min(num_threads, blocks)); + std::vector centroid_norms(k); + std::vector> products( + threads, std::vector(kPointBlock * std::min(k, kCentroidBlock)) + ); + std::vector errors(threads); + +#pragma omp parallel for num_threads(threads) schedule(static) + for (std::ptrdiff_t cluster = 0; cluster < static_cast(k); ++cluster) { + double norm = 0.0; + for (size_t dim = 0; dim < d; ++dim) { + const double value = centroids[static_cast(cluster) * d + dim]; + norm += value * value; + } + centroid_norms[static_cast(cluster)] = norm; + } + + // For d <= 65536, this exceeds the float dot-product error bound + // gamma_d * sum(abs(x_i*c_i)), using 2*abs(x_i*c_i) <= x_i^2+c_i^2. + // The slack also covers double norm/distance rounding and underflow. + const double error_scale = 8.0 * static_cast(d) * + static_cast(std::numeric_limits::epsilon()); + const double error_floor = + 8.0 * static_cast(d) * std::numeric_limits::min() + + 8.0 * std::numeric_limits::epsilon(); +#pragma omp parallel num_threads(threads) + { + // Set only this OpenMP task's ICV, not Eigen's process-global thread limit. + // This also prevents nested GEMM teams when the outer team has one worker. + omp_set_num_threads(1); + const auto thread = static_cast(omp_get_thread_num()); + float* dots = products[thread].data(); + std::array point_norms{}; + std::array best_distances{}; +#pragma omp for schedule(static) + for (std::ptrdiff_t block = 0; block < static_cast(blocks); + ++block) { + if (errors[thread]) { + continue; + } + try { + const size_t first_point = static_cast(block) * kPointBlock; + const size_t rows = std::min(kPointBlock, n - first_point); + for (size_t row = 0; row < rows; ++row) { + const float* vector = x + (first_point + row) * d; + double norm = 0.0; + for (size_t dim = 0; dim < d; ++dim) { + const double value = vector[dim]; + norm += value * value; + } + point_norms[row] = norm; + best_distances[row] = std::numeric_limits::infinity(); + } + for (size_t first_cluster = 0; first_cluster < k; + first_cluster += kCentroidBlock) { + const size_t cols = std::min(kCentroidBlock, k - first_cluster); + simd::matrix_product_transposed( + x + first_point * d, + centroids + first_cluster * d, + dots, + rows, + d, + cols + ); + for (size_t row = 0; row < rows; ++row) { + const size_t point = first_point + row; + const float* vector = x + point * d; + for (size_t col = 0; col < cols; ++col) { + const size_t cluster = first_cluster + col; + const double norm_sum = + point_norms[row] + centroid_norms[cluster]; + const double dot = dots[row * cols + col]; + const double estimate = + spherical ? 1.0 - dot : norm_sum - 2.0 * dot; + const double error = error_scale * norm_sum + error_floor; + if (estimate > best_distances[row] + error) { + continue; + } + const float* centroid = centroids + cluster * d; + double distance = 0.0; + if (spherical) { + for (size_t dim = 0; dim < d; ++dim) { + distance -= + static_cast(vector[dim]) * centroid[dim]; + } + distance += 1.0; + } else { + for (size_t dim = 0; dim < d; ++dim) { + const double difference = + static_cast(vector[dim]) - centroid[dim]; + distance += difference * difference; + } + } + // Visit in ascending ID order, retaining the first exact tie. + if (distance < best_distances[row]) { + assignments[point] = static_cast(cluster); + best_distances[row] = distance; + } + } + } + } + for (size_t row = 0; row < rows; ++row) { + distances[first_point + row] = static_cast(best_distances[row]); + } + } catch (...) { errors[thread] = std::current_exception(); } + } + } + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } +} + +template +class LloydKMeans : public Parameters { + public: + using Parameters::early_stop_threshold; + using Parameters::final_assignment; + using Parameters::min_points_per_centroid; + using Parameters::niter; + using Parameters::num_threads; + using Parameters::seed; + using Parameters::spherical; + using Parameters::verbose; + + LloydKMeans(size_t d, size_t k, const Parameters& cp) : Parameters(cp), d(d), k(k) { + validate_parameters(); + } + + void train(size_t n, const float* x) { + validate_parameters(); + if (x == nullptr) { + throw std::invalid_argument( + std::string(Assignment::kName) + " x must not be null" + ); + } + if (n < k) { + throw std::invalid_argument( + std::string(Assignment::kName) + " requires n >= k" + ); + } + if (n > static_cast(std::numeric_limits::max()) / d) { + throw std::invalid_argument( + std::string(Assignment::kName) + " input dimensions overflow size_t" + ); + } + const uint32_t threads = effective_num_threads(); + static_cast(threads); + // Leave headroom for products of squared distances in graph pruning, + // as well as quantized reconstruction and rotated residuals. + const double max_coordinate = std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(d)) + ); + bool in_range = true; + bool finite = true; +#pragma omp parallel for num_threads(threads) reduction(&& : finite, in_range) schedule(static) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n * d); + ++loop_index) { + const size_t i = static_cast(loop_index); + finite = finite && std::isfinite(x[i]); + in_range = in_range && std::abs(static_cast(x[i])) <= max_coordinate; + } + if (!finite) { + throw std::invalid_argument( + std::string(Assignment::kName) + " x must contain only finite values" + ); + } + if (!in_range) { + throw std::invalid_argument( + std::string(Assignment::kName) + + " x exceeds the safe float32 coordinate range" + ); + } + if (verbose && min_points_per_centroid > 0 && n / k < min_points_per_centroid) { + std::cerr << "WARNING " << Assignment::kName << " clustering " << n + << " points to " << k << " centroids: please provide at least " + << min_points_per_centroid << " training points per centroid\n"; + } + + // Retraining may borrow the previous centroids, including a subrange. + // Preserve that input through initialization, updates, and final assignment. + std::vector input_copy; + const std::less before; + if (!centroids.empty() && !before(x, centroids.data()) && + before(x, centroids.data() + centroids.size())) { + input_copy.assign(x, x + n * d); + x = input_copy.data(); + } + + Assignment assigner(d, k, x, n, threads, *this); + + centroids.resize(k * d); + assignments.resize(n); + distances.resize(n); + iteration_stats.clear(); + iteration_stats.reserve(niter); + initialize_centroids(x, n); + if (spherical) { + normalize_centroids(); + } + + std::vector previous_centroids(centroids.size()); + std::vector cluster_sizes(k); + std::vector sums(centroids.size()); + const size_t update_workers = std::min(k, static_cast(threads)); + std::vector> update_point_bins(update_workers); + std::vector update_worker_for_cluster(k); + for (size_t worker = 0; worker < update_workers; ++worker) { + const size_t first_cluster = k * worker / update_workers; + const size_t last_cluster = k * (worker + 1) / update_workers; + for (size_t cluster = first_cluster; cluster < last_cluster; ++cluster) { + update_worker_for_cluster[cluster] = worker; + } + update_point_bins[worker].reserve(n / update_workers); + } + double previous_obj = std::numeric_limits::infinity(); + + for (size_t iteration = 0; iteration < niter; ++iteration) { + previous_centroids = centroids; + assigner.assign( + previous_centroids.data(), assignments.data(), distances.data() + ); + + const double obj = std::accumulate(distances.begin(), distances.end(), 0.0); + std::fill(sums.begin(), sums.end(), 0.0); + std::fill(cluster_sizes.begin(), cluster_sizes.end(), size_t{0}); + 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); + iteration_stats.push_back(KMeansIterationStats{ + iteration + 1, obj, shift, nsplit}); + + if (verbose) { + std::cout << Assignment::kName << " iteration " << iteration + 1 << "/" + << niter << " | objective=" << obj << " | shift=" << shift + << " | splits=" << nsplit << '\n'; + } + + const bool converged = detail::should_terminate_by_objective( + previous_obj, obj, early_stop_threshold + ); + previous_obj = obj; + if (iteration > 0 && converged) { + break; + } + } + + // Return assignments for the final centroids, not the centroids from + // the beginning of the last update step. + if (final_assignment == FinalAssignmentMode::Exact) { + detail::exact_assign( + x, + centroids.data(), + n, + d, + k, + spherical, + threads, + assignments.data(), + distances.data() + ); + } else { + assigner.assign(centroids.data(), assignments.data(), distances.data(), true); + } + final_obj = std::accumulate(distances.begin(), distances.end(), 0.0); + } + + size_t d; + size_t k; + std::vector centroids; + std::vector assignments; + std::vector distances; + std::vector iteration_stats; + /// Objective of the final returned assignments against the final centroids. + double final_obj = 0.0; + + private: + void validate_parameters() const { + if (d == 0 || k == 0 || niter == 0) { + throw std::invalid_argument( + std::string(Assignment::kName) + " d, k, and niter must be positive" + ); + } + if (d < 64 || d > rotator_impl::FhtKacRotator::kMaxDim) { + throw std::invalid_argument( + std::string(Assignment::kName) + + " currently requires dimensions in the range [64, 65536]" + ); + } + if (num_threads > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument( + std::string(Assignment::kName) + " num_threads must fit in an OpenMP int" + ); + } + if (final_assignment != FinalAssignmentMode::Approximate && + final_assignment != FinalAssignmentMode::Exact) { + throw std::invalid_argument( + std::string(Assignment::kName) + " final_assignment mode is invalid" + ); + } + Assignment::validate_parameters(k, *this); + if (!std::isfinite(early_stop_threshold) || early_stop_threshold < 0.0 || + early_stop_threshold > 1.0) { + throw std::invalid_argument( + std::string(Assignment::kName) + + " early_stop_threshold must be in the range [0, 1]" + ); + } + } + + void initialize_centroids(const float* x, size_t n) { + std::vector indices(n); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::mt19937 rng(seed); + std::shuffle(indices.begin(), indices.end(), rng); + for (size_t cluster = 0; cluster < k; ++cluster) { + std::memcpy( + centroids.data() + cluster * d, x + indices[cluster] * d, d * sizeof(float) + ); + } + } + + void update_centroids( + const float* x, + size_t n, + std::vector& sums, + std::vector& cluster_sizes, + std::vector>& point_bins, + const std::vector& worker_for_cluster + ) { + for (auto& bin : point_bins) { + bin.clear(); + } + // Keep point order within each cluster, preserving the original sums. + for (size_t point = 0; point < n; ++point) { + point_bins[worker_for_cluster[assignments[point]]].push_back(point); + } + const uint32_t threads = effective_num_threads(); + static_cast(threads); +#pragma omp parallel for num_threads(threads) schedule(static) + for (std::ptrdiff_t loop_worker = 0; + loop_worker < static_cast(point_bins.size()); + ++loop_worker) { + const size_t worker = static_cast(loop_worker); + simd::accumulate_cluster_sums( + x, + point_bins[worker].data(), + point_bins[worker].size(), + d, + assignments.data(), + sums.data(), + cluster_sizes.data() + ); + } + } + + size_t consolidate_centroids( + const float* x, + size_t n, + const std::vector& sums, + std::vector& cluster_sizes + ) { + const uint32_t threads = effective_num_threads(); + static_cast(threads); +#pragma omp parallel for num_threads(threads) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); + ++loop_index) { + const size_t cluster = static_cast(loop_index); + const double scale = cluster_sizes[cluster] == 0 + ? 0.0 + : 1.0 / static_cast(cluster_sizes[cluster]); + for (size_t dim = 0; dim < d; ++dim) { + centroids[cluster * d + dim] = + static_cast(sums[cluster * d + dim] * scale); + } + } + if (spherical) { + normalize_centroids(); + } + + const size_t nsplit = static_cast( + std::count(cluster_sizes.begin(), cluster_sizes.end(), size_t{0}) + ); + if (nsplit == 0) { + return 0; + } + + // Reuse the distance buffer after recording the pre-update objective. + // Keep populated means intact and give unused centroids poorly fitted + // training points. Unlike coordinate perturbations, this also works at zero. + const auto distance_func = + spherical ? dot_product_dis : euclidean_sqr; + const auto less = [&](size_t left, size_t right) { + return distances[left] == distances[right] ? left > right + : distances[left] < distances[right]; + }; + if (nsplit == 1) { + // A single empty centroid needs only one eligible maximum per worker. + std::vector best_points(threads, n); +#pragma omp parallel num_threads(threads) + { + size_t best = n; +#pragma omp for schedule(static) + for (std::ptrdiff_t loop_index = 0; + loop_index < static_cast(n); + ++loop_index) { + const size_t point = static_cast(loop_index); + if (cluster_sizes[assignments[point]] <= 1) { + continue; + } + distances[point] = distance_func( + x + point * d, centroids.data() + assignments[point] * d, d + ); + if (best == n || less(best, point)) { + best = point; + } + } + best_points[static_cast(omp_get_thread_num())] = best; + } + size_t point = n; + for (const size_t candidate : best_points) { + if (candidate != n && (point == n || less(point, candidate))) { + point = candidate; + } + } + assert(point != n); // n >= k guarantees a donor with at least two points. + const size_t cluster = static_cast( + std::find(cluster_sizes.begin(), cluster_sizes.end(), size_t{0}) - + cluster_sizes.begin() + ); + std::copy_n(x + point * d, d, centroids.data() + cluster * d); + if (spherical) { + normalize_centroid(centroids.data() + cluster * d); + } + --cluster_sizes[assignments[point]]; + cluster_sizes[cluster] = 1; + return 1; + } + +#pragma omp parallel for num_threads(threads) schedule(static) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); + ++loop_index) { + const size_t point = static_cast(loop_index); + distances[point] = + distance_func(x + point * d, centroids.data() + assignments[point] * d, d); + } + std::vector candidates(n); + std::iota(candidates.begin(), candidates.end(), size_t{0}); + std::make_heap(candidates.begin(), candidates.end(), less); + for (size_t cluster = 0; cluster < k; ++cluster) { + if (cluster_sizes[cluster] != 0) { + continue; + } + size_t point = 0; + do { + point = candidates.front(); + std::pop_heap(candidates.begin(), candidates.end(), less); + candidates.pop_back(); + } while (cluster_sizes[assignments[point]] <= 1); + std::copy_n(x + point * d, d, centroids.data() + cluster * d); + if (spherical) { + normalize_centroid(centroids.data() + cluster * d); + } + --cluster_sizes[assignments[point]]; + cluster_sizes[cluster] = 1; + } + return nsplit; + } + + void normalize_centroid(float* centroid) const { + double squared_norm = 0; + for (size_t dim = 0; dim < d; ++dim) { + squared_norm += static_cast(centroid[dim]) * centroid[dim]; + } + if (squared_norm == 0) { + return; + } + const double scale = 1.0 / std::sqrt(squared_norm); + for (size_t dim = 0; dim < d; ++dim) { + centroid[dim] = static_cast(centroid[dim] * scale); + } + } + + void normalize_centroids() { + const uint32_t threads = effective_num_threads(); + static_cast(threads); +#pragma omp parallel for num_threads(threads) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); + ++loop_index) { + const size_t cluster = static_cast(loop_index); + normalize_centroid(centroids.data() + cluster * d); + } + } + + [[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)); + } +}; + +} // namespace detail + +} // namespace rabitqlib::clustering diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp new file mode 100644 index 00000000..3de60e81 --- /dev/null +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -0,0 +1,408 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/clustering/detail/kmeans.hpp" +#include "rabitqlib/defines.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" +#include "rabitqlib/utils/visited_set.hpp" + +namespace rabitqlib::qgkmeans { + +using clustering::FinalAssignmentMode; +using QGKMeansIterationStats = clustering::KMeansIterationStats; + +namespace detail { + +using clustering::detail::exact_assign; +using clustering::detail::should_terminate_by_objective; +class QGAssignment; + +} // namespace detail + +struct QGKMeansParameters { + /// Number of clustering iterations. + size_t niter = 25; + bool verbose = false; + /// Normalize centroids after each update. + bool spherical = false; + uint32_t seed = 42; + /// Warn below this number of training vectors per centroid. + size_t min_points_per_centroid = 39; + /// Stop when the absolute relative objective (WCSS) change is at most this value. + double early_stop_threshold = 0.0; + FinalAssignmentMode final_assignment = FinalAssignmentMode::Approximate; + + // SymphonyQG assignment parameters. + uint32_t graph_degree = 32; + uint32_t ef_build = 240; + uint32_t ef_search = 16; + /// Graph-building passes after PiPNN initialization; the last refines the graph. + uint32_t graph_build_iterations = 1; + uint32_t num_threads = 0; + /// Centroid storage in QG: 0 for raw float32, or 4/8-bit RaBitQ. + uint32_t quantization_bits = 0; +}; + +/** + * Builds a SymphonyQG over a fixed number of centroids and assigns vectors to + * them. The graph allocation and per-thread search scratch are reused between + * calls, and the previous assignment is used as the next call's search hint. + */ +class QGAssigner { + friend class detail::QGAssignment; + + private: + struct ThreadScratch { + std::vector rotated_query; + std::vector estimated_distances; + std::vector lookup_table; + BatchQuery batch_query; + buffer::SearchBuffer search_pool; + buffer::SearchBuffer result_pool; + + ThreadScratch(size_t padded_dim, size_t M, size_t ef_search) + : rotated_query(padded_dim) + , estimated_distances(M) + , lookup_table(padded_dim << 2) + , search_pool(ef_search) + , result_pool(1) {} + }; + + public: + QGAssigner( + size_t d, + size_t k, + uint32_t M, + uint32_t ef_construction, + uint32_t ef_search, + uint32_t niter, + uint32_t num_threads = 0, + uint32_t seed = 42, + MetricType metric_type = METRIC_L2, + uint32_t quantization_bits = 0 + ) + : d_(d) + , k_(k) + , M_(M) + , ef_construction_(ef_construction) + , ef_search_(ef_search) + , niter_(niter) + , num_threads_(num_threads) + , seed_(seed) { + validate_configuration(); + num_threads_ = static_cast(resolve_num_threads(num_threads_)); + + graph_ = std::make_unique>( + k_, d_, M_, metric_type, RotatorType::FhtKacRotator, quantization_bits, seed_ + ); + graph_->set_ef(ef_search_); + + scratch_.reserve(num_threads_); + for (uint32_t thread = 0; thread < num_threads_; ++thread) { + scratch_.emplace_back( + std::make_unique(graph_->padded_dim(), M_, ef_search_) + ); + } + } + + // Optional n-element cache: fill it on the first call, then reuse it only + // for the same input vectors and this assigner. A null pointer disables it. + void assign( + const float* centroids, + const float* x, + size_t n, + PID* labels, + float* distances, + float* query_sums = nullptr, + bool sums_ready = false + ) { + if (centroids == nullptr || x == nullptr || labels == nullptr || + distances == nullptr) { + throw std::invalid_argument( + "QGAssigner input and output pointers must not be null" + ); + } + + if (n > static_cast(std::numeric_limits::max()) / d_) { + throw std::invalid_argument("QGAssigner input dimensions overflow ptrdiff_t"); + } + + sums_ready = query_sums != nullptr && sums_ready; + if (builder_) { + builder_->reset(centroids); + } else { + builder_ = std::make_unique( + *graph_, + ef_construction_, + centroids, + num_threads_, + symqg::QGInitialization::PiPNN, + seed_, + true + ); + } + if (niter_ == 1) { + builder_->build(); + } else { + builder_->build(niter_); + } + const auto distance_func = graph_->metric_type() == METRIC_IP + ? dot_product_dis + : euclidean_sqr; + + // Stable grouping changes only the order of independent full searches. + std::vector query_order; + if (previous_labels_.size() == n && n <= std::numeric_limits::max()) { + query_order.resize(n); + std::vector offsets(k_ + 1, 0); + for (PID label : previous_labels_) { + ++offsets[label + 1]; + } + std::partial_sum(offsets.begin(), offsets.end(), offsets.begin()); + for (size_t point = 0; point < n; ++point) { + query_order[offsets[previous_labels_[point]]++] = + static_cast(point); + } + } + std::vector errors(num_threads_); +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic, 64) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); + ++loop_index) { + const size_t i = query_order.empty() + ? static_cast(loop_index) + : query_order[static_cast(loop_index)]; + const int thread_id = omp_get_thread_num(); + if (errors[static_cast(thread_id)]) { + continue; + } + try { + ThreadScratch& scratch = *scratch_[static_cast(thread_id)]; + // Keep visited storage with the worker, as in QuantizedGraph::search. + thread_local VisitedSet visited; + thread_local size_t visited_size = 0; + thread_local size_t visited_capacity = 0; + const size_t capacity = + std::min(static_cast(ef_search_) * ef_search_, k_ / 10); + if (visited_size != k_ || visited_capacity != capacity) { + visited.initialize(k_, capacity); + visited_size = k_; + visited_capacity = capacity; + } + const PID hint = + previous_labels_.size() == n ? previous_labels_[i] : kPidMax; + graph_->search_with_scratch( + x + i * d_, + 1, + labels + i, + distances + i, + scratch.rotated_query.data(), + scratch.estimated_distances.data(), + scratch.lookup_table.data(), + scratch.batch_query, + scratch.search_pool, + scratch.result_pool, + visited, + hint, + sums_ready ? query_sums + i : nullptr + ); + if (query_sums != nullptr && !sums_ready) { + query_sums[i] = scratch.batch_query.k1xsumq(); + } + if (graph_->is_quantized()) { + // Keep objectives and returned distances relative to the original + // centroids; graph search scores are quantized estimates. + distances[i] = + distance_func(x + i * d_, centroids + labels[i] * d_, d_); + } + } catch (...) { + errors[static_cast(thread_id)] = std::current_exception(); + } + } + + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } + + previous_labels_.assign(labels, labels + n); + } + + void clear_hints() { previous_labels_.clear(); } + + private: + void validate_configuration() const { + if (num_threads_ > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("QGAssigner num_threads must fit in an OpenMP int"); + } + if (d_ < 64 || d_ > rotator_impl::FhtKacRotator::kMaxDim) { + throw std::invalid_argument( + "QGKMeans currently requires dimensions in the range [64, 65536]" + ); + } + if (M_ == 0 || M_ % fastscan::kBatchSize != 0) { + throw std::invalid_argument( + "QGKMeans graph degree must be a positive multiple of 32" + ); + } + if (k_ <= M_) { + throw std::invalid_argument( + "QGKMeans requires more centroids than the graph degree" + ); + } + if (ef_construction_ == 0 || ef_search_ == 0) { + throw std::invalid_argument("QGKMeans ef_build and ef_search must be positive"); + } + if (niter_ == 0) { + throw std::invalid_argument("QGKMeans graph_build_iterations must be positive"); + } + } + + size_t d_; + size_t k_; + uint32_t M_; + uint32_t ef_construction_; + uint32_t ef_search_; + uint32_t niter_; + uint32_t num_threads_; + uint32_t seed_; + std::unique_ptr> graph_; + std::unique_ptr builder_; + std::vector> scratch_; + std::vector previous_labels_; +}; + +namespace detail { + +// Retain graph-specific state for one synchronous training call. +class QGAssignment { + size_t d_; + const float* x_; + size_t n_; + uint32_t num_threads_; + MetricType metric_type_; + bool quantized_; + QGAssigner assigner_; + bool cache_query_sums_; + std::vector query_sums_; + bool query_sums_ready_ = false; + + public: + static constexpr const char* kName = "QGKMeans"; + + static void validate_parameters(size_t, const QGKMeansParameters& cp) { + if (cp.quantization_bits != 0 && cp.quantization_bits != 4 && + cp.quantization_bits != 8) { + throw std::invalid_argument("QGKMeans quantization_bits must be 0, 4, or 8"); + } + } + + QGAssignment( + size_t d, + size_t k, + const float* x, + size_t n, + uint32_t num_threads, + const QGKMeansParameters& cp + ) + : d_(d) + , x_(x) + , n_(n) + , num_threads_(num_threads) + , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) + , quantized_(cp.quantization_bits != 0) + , assigner_( + d, + k, + cp.graph_degree, + cp.ef_build, + cp.ef_search, + cp.graph_build_iterations, + cp.num_threads, + cp.seed, + metric_type_, + cp.quantization_bits + ) + // Inputs and the assigner's rotator are immutable for this training call. + , cache_query_sums_( + assigner_.graph_->padded_dim() >= 1024 && + (cp.niter > 1 || cp.final_assignment != FinalAssignmentMode::Exact) + ) {} + + void assign(const float* centroids, PID* labels, float* distances, bool final = false) { + if (cache_query_sums_ && query_sums_.empty()) { + // Allocate after centroid initialization releases its point permutation. + query_sums_.resize(n_); + } + std::vector previous_labels; + if (final && quantized_) { + previous_labels.assign(labels, labels + n_); + } + assigner_.assign( + centroids, + x_, + n_, + labels, + distances, + query_sums_.empty() ? nullptr : query_sums_.data(), + query_sums_ready_ + ); + query_sums_ready_ = true; + if (previous_labels.empty()) { + return; + } + + // Protect the returned assignments without changing exploration during + // Lloyd iterations: uphill approximate moves can improve final centroids. + const auto distance_func = + metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; +#pragma omp parallel for num_threads(num_threads_) schedule(static) + for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n_); + ++loop_index) { + const size_t point = static_cast(loop_index); + const PID previous = previous_labels[point]; + if (previous == labels[point]) { + continue; + } + const float distance = + distance_func(x_ + point * d_, centroids + previous * d_, d_); + if (distance < distances[point]) { + labels[point] = previous; + distances[point] = distance; + } + } + } +}; + +} // namespace detail + +/** K-means with SymphonyQG graph assignment. */ +class QGKMeans + : public clustering::detail::LloydKMeans { + using Base = clustering::detail::LloydKMeans; + + public: + QGKMeans(size_t d, size_t k) : QGKMeans(d, k, QGKMeansParameters{}) {} + + QGKMeans(size_t d, size_t k, const QGKMeansParameters& cp) : Base(d, k, cp) {} +}; + +} // namespace rabitqlib::qgkmeans diff --git a/include/rabitqlib/clustering/rabitqkmeans.hpp b/include/rabitqlib/clustering/rabitqkmeans.hpp new file mode 100644 index 00000000..d7e54f6b --- /dev/null +++ b/include/rabitqlib/clustering/rabitqkmeans.hpp @@ -0,0 +1,243 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/clustering/detail/kmeans.hpp" +#include "rabitqlib/defines.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::rabitqkmeans { + +using clustering::FinalAssignmentMode; +using RaBitQKMeansIterationStats = clustering::KMeansIterationStats; + +struct RaBitQKMeansParameters { + /// Number of clustering iterations. + size_t niter = 25; + bool verbose = false; + /// Normalize centroids after each update. + bool spherical = false; + uint32_t seed = 42; + /// Warn below this number of training vectors per centroid. + size_t min_points_per_centroid = 39; + /// Stop when the absolute relative objective (WCSS) change is at most this value. + double early_stop_threshold = 0.0; + FinalAssignmentMode final_assignment = FinalAssignmentMode::Approximate; + uint32_t num_threads = 0; +}; + +namespace detail { + +// Encode immutable training points once; each moving centroid supplies one LUT. +// RaBitQKMeans validates inputs and keeps the borrowed x buffer alive through training. +class RaBitQAssigner { + size_t d_; + size_t k_; + const float* x_; + size_t n_; + size_t num_batches_; + uint32_t num_threads_; + MetricType metric_type_; + size_t padded_dim_; + rotator_impl::FhtKacRotator rotator_; + std::vector center_; + std::vector batch_data_; + std::vector> centroid_queries_; + std::vector rotated_query_; + std::vector lookup_table_; + + public: + static constexpr const char* kName = "RaBitQKMeans"; + + static void validate_parameters(size_t k, const RaBitQKMeansParameters&) { + if (k - 1 > std::numeric_limits::max()) { + throw std::invalid_argument( + "RaBitQKMeans k exceeds the representable assignment ID range" + ); + } + } + + RaBitQAssigner( + size_t d, + size_t k, + const float* x, + size_t n, + uint32_t num_threads, + const RaBitQKMeansParameters& cp + ) + : d_(d) + , k_(k) + , x_(x) + , n_(n) + , num_batches_(1 + (n - 1) / fastscan::kBatchSize) + , num_threads_(static_cast(std::min(num_threads, num_batches_))) + , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) + , padded_dim_(rotator_impl::padding_requirement(d, RotatorType::FhtKacRotator)) + , rotator_(d, padded_dim_, cp.seed) + , center_(d) + , batch_data_(num_batches_ * QGBatchDataMap::data_bytes(padded_dim_), 0) + , centroid_queries_(k) + , rotated_query_(padded_dim_) + , lookup_table_(padded_dim_ * 4) { + { + std::vector> partial_sums( + num_threads_, std::vector(d_, 0.0) + ); +#pragma omp parallel num_threads(num_threads_) + { + auto& sums = partial_sums[static_cast(omp_get_thread_num())]; +#pragma omp for schedule(static) + for (std::ptrdiff_t point = 0; point < static_cast(n_); + ++point) { + const float* vector = x_ + static_cast(point) * d_; + for (size_t dim = 0; dim < d_; ++dim) { + sums[dim] += static_cast(vector[dim]); + } + } + } + for (size_t dim = 0; dim < d_; ++dim) { + double sum = 0.0; + for (const auto& partial : partial_sums) { + sum += partial[dim]; + } + center_[dim] = static_cast(sum / static_cast(n_)); + } + } + + std::vector rotated_center(padded_dim_); + rotator_.rotate(center_.data(), rotated_center.data()); + std::vector> rotated_blocks( + num_threads_, std::vector(fastscan::kBatchSize * padded_dim_) + ); + std::vector errors(num_threads_); +#pragma omp parallel for num_threads(num_threads_) schedule(static) + for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); + ++batch) { + const auto thread = static_cast(omp_get_thread_num()); + if (errors[thread]) { + continue; + } + try { + const auto block = static_cast(batch); + const size_t first = block * fastscan::kBatchSize; + const size_t count = std::min(fastscan::kBatchSize, n_ - first); + auto& rotated = rotated_blocks[thread]; + for (size_t point = 0; point < count; ++point) { + rotator_.rotate( + x_ + (first + point) * d_, rotated.data() + point * padded_dim_ + ); + } + quant::quantize_qg_batch( + rotated.data(), + rotated_center.data(), + count, + padded_dim_, + batch_data_.data() + + block * QGBatchDataMap::data_bytes(padded_dim_), + metric_type_ + ); + } catch (...) { errors[thread] = std::current_exception(); } + } + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } + } + + void assign(const float* centroids, PID* labels, float* distances, bool = false) { + const auto distance_func = + metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; + for (size_t cluster = 0; cluster < k_; ++cluster) { + const float* centroid = centroids + cluster * d_; + rotator_.rotate(centroid, rotated_query_.data()); + auto& query = centroid_queries_[cluster]; + query.reset( + rotated_query_.data(), padded_dim_, lookup_table_.data(), metric_type_ + ); + query.set_g_add(distance_func(centroid, center_.data(), d_)); + } + + std::vector errors(num_threads_); +#pragma omp parallel for num_threads(num_threads_) schedule(static) + for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); + ++batch) { + const auto thread = static_cast(omp_get_thread_num()); + if (errors[thread]) { + continue; + } + try { + const auto block = static_cast(batch); + const size_t first = block * fastscan::kBatchSize; + const size_t count = std::min(fastscan::kBatchSize, n_ - first); + const char* codes = batch_data_.data() + + block * QGBatchDataMap::data_bytes(padded_dim_); + std::array estimated_distances{}; + std::array best_distances; + best_distances.fill(std::numeric_limits::infinity()); + std::array best_labels{}; + for (size_t cluster = 0; cluster < k_; ++cluster) { + qg_batch_estdist( + codes, + centroid_queries_[cluster], + padded_dim_, + estimated_distances.data() + ); + for (size_t point = 0; point < count; ++point) { + // Ascending centroid order retains the lowest ID for exact score + // ties. + if (estimated_distances[point] < best_distances[point]) { + best_distances[point] = estimated_distances[point]; + best_labels[point] = static_cast(cluster); + } + } + } + for (size_t point = 0; point < count; ++point) { + const size_t id = first + point; + labels[id] = best_labels[point]; + // Scores choose the centroid; objectives retain original-space + // distances. + distances[id] = distance_func( + x_ + id * d_, centroids + best_labels[point] * d_, d_ + ); + } + } catch (...) { errors[thread] = std::current_exception(); } + } + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } + } +}; + +} // namespace detail + +/** K-means with a graph-free RaBitQ scan, recommended for small cluster counts. */ +class RaBitQKMeans : public clustering::detail:: + LloydKMeans { + using Base = + clustering::detail::LloydKMeans; + + public: + RaBitQKMeans(size_t d, size_t k) : RaBitQKMeans(d, k, RaBitQKMeansParameters{}) {} + + RaBitQKMeans(size_t d, size_t k, const RaBitQKMeansParameters& cp) : Base(d, k, cp) {} +}; + +} // namespace rabitqlib::rabitqkmeans diff --git a/include/rabitqlib/index/lut.hpp b/include/rabitqlib/index/lut.hpp index 242d9410..a37559d5 100644 --- a/include/rabitqlib/index/lut.hpp +++ b/include/rabitqlib/index/lut.hpp @@ -51,30 +51,55 @@ class Lut { explicit Lut(const T* rotated_query, size_t padded_dim, bool use_hacc = false) : table_length_(checked_table_length(padded_dim)) , lut_(checked_storage_length(padded_dim, use_hacc)) { - // quantize float lut std::vector lut_float(table_length_); - fastscan::pack_lut(padded_dim, rotated_query, lut_float.data()); + std::vector lut_u16(use_hacc ? table_length_ : 0); + reset(rotated_query, padded_dim, lut_float.data(), lut_u16.data(), use_hacc); + } + + /** + * Reinitialize the lookup table while reusing its allocation. + * lut_float_scratch must contain space for padded_dim * 4 floats. + * lut_u16_scratch is additionally required when use_hacc is true. + */ + void reset( + const T* rotated_query, + size_t padded_dim, + float* lut_float_scratch, + uint16_t* lut_u16_scratch = nullptr, + bool use_hacc = false + ) { + const size_t required_size = checked_storage_length(padded_dim, use_hacc); + if (rotated_query == nullptr || lut_float_scratch == nullptr) { + throw std::invalid_argument("Lut::reset requires query and float scratch"); + } + if (use_hacc && lut_u16_scratch == nullptr) { + throw std::invalid_argument( + "Lut::reset requires uint16 scratch for high-accuracy lookup tables" + ); + } + table_length_ = checked_table_length(padded_dim); + lut_.resize(required_size); + + fastscan::pack_lut(padded_dim, rotated_query, lut_float_scratch); T vl_lut; T vr_lut; - data_range(lut_float.data(), table_length_, vl_lut, vr_lut); + data_range(lut_float_scratch, table_length_, vl_lut, vr_lut); if (use_hacc) { delta_ = (vr_lut - vl_lut) / ((1 << kNumBitsHacc) - 1); - - // quantize float lut into uint16 then change to split table - std::vector lut_u16(table_length_); scalar_quantize( - lut_u16.data(), lut_float.data(), table_length_, vl_lut, delta_ + lut_u16_scratch, lut_float_scratch, table_length_, vl_lut, delta_ ); - fastscan::transfer_lut_hacc(lut_u16.data(), padded_dim, lut_.data()); + fastscan::transfer_lut_hacc(lut_u16_scratch, padded_dim, lut_.data()); } else { delta_ = (vr_lut - vl_lut) / ((1 << kNumBits) - 1); - scalar_quantize(lut_.data(), lut_float.data(), table_length_, vl_lut, delta_); + scalar_quantize(lut_.data(), lut_float_scratch, table_length_, vl_lut, delta_); } - size_t num_table = table_length_ / 16; - sum_vl_lut_ = vl_lut * static_cast(num_table); + const size_t num_tables = table_length_ / 16; + sum_vl_lut_ = vl_lut * static_cast(num_tables); } + [[nodiscard]] const uint8_t* lut() const { return lut_.data(); }; [[nodiscard]] T delta() const { return delta_; }; [[nodiscard]] T sum_vl() const { return sum_vl_lut_; }; diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index 3440d765..ec4f1460 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -21,9 +21,11 @@ class BatchQuery { Lut lookup_table_; T G_add_ = 0; T G_k1xSumq_ = 0; // G_k1xSumq - MetricType metric_type_; + MetricType metric_type_ = METRIC_L2; public: + explicit BatchQuery() = default; + explicit BatchQuery( const T* rotated_query, size_t padded_dim, MetricType metric_type = METRIC_L2 ) @@ -49,6 +51,32 @@ class BatchQuery { G_add_ = metric_type_ == METRIC_IP ? dist - T{1} : dist; } + /** + * Reinitialize the query while reusing the lookup-table allocation. + * An optional correction must come from this exact rotated query. + */ + void reset( + const T* rotated_query, + size_t padded_dim, + float* lut_float_scratch, + MetricType metric_type = METRIC_L2, + const T* k1xsumq = nullptr + ) { + metric_type_ = metric_type; + lookup_table_.reset(rotated_query, padded_dim, lut_float_scratch); + + if (k1xsumq != nullptr) { + G_k1xSumq_ = *k1xsumq; + } else { + constexpr float c_1 = -((1 << 1) - 1) / 2.F; + const T sumq = std::accumulate( + rotated_query, rotated_query + padded_dim, static_cast(0) + ); + G_k1xSumq_ = sumq * c_1; + } + G_add_ = 0; + } + [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } }; diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 2758bf62..67c07be5 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "rabitqlib/defines.hpp" @@ -30,6 +31,7 @@ class QuantizedQuery { size_t padded_dim, MetricType metric_type ); + QuantizedQuery(const float* rotated_query, const QuantizedQuery& prepared); [[nodiscard]] const float* rotated_query() const; [[nodiscard]] float k1xsumq() const; [[nodiscard]] float g_add() const; @@ -73,6 +75,13 @@ class QuantizedGraph { size_t row_offset_ = 0; // length of entire row size_t ef_ = 0; bool ready_ = false; + uint32_t seed_ = 42; + + // Owned by QGBuilder and never persisted. + struct BuildCache { + std::vector rotated_vectors; + std::vector queries; + }; [[nodiscard]] static size_t checked_add(size_t lhs, size_t rhs); @@ -90,23 +99,36 @@ class QuantizedGraph { void set_quantization_centroid(const float* centroid); - [[nodiscard]] char* get_row_data(PID data_id); + [[nodiscard]] char* get_row_data(PID data_id) { + return reinterpret_cast(get_vector(data_id)); + } - [[nodiscard]] const char* get_row_data(PID data_id) const; + [[nodiscard]] const char* get_row_data(PID data_id) const { + return reinterpret_cast(get_vector(data_id)); + } - [[nodiscard]] float* get_vector(PID data_id); + [[nodiscard]] float* get_vector(PID data_id) { + return data_.data() + ((row_offset_ / sizeof(float)) * data_id); + } - [[nodiscard]] const float* get_vector(PID data_id) const; + [[nodiscard]] const float* get_vector(PID data_id) const { + return data_.data() + ((row_offset_ / sizeof(float)) * data_id); + } - [[nodiscard]] char* get_quantized_vector(PID data_id); + [[nodiscard]] char* get_quantized_vector(PID data_id) { return get_row_data(data_id); } - [[nodiscard]] const char* get_quantized_vector(PID data_id) const; + [[nodiscard]] const char* get_quantized_vector(PID data_id) const { + return get_row_data(data_id); + } - void prepare_query(const float*, std::vector&, std::optional&) - const; + void validate_search(const float*, uint32_t, const uint32_t*, const float*) const; - const float* - prepare_build_query(PID, std::vector&, std::optional&) const; + const float* prepare_build_query( + PID, + std::vector&, + std::optional&, + const BuildCache* = nullptr + ) const; float point_distance(const float*, const QuantizedQuery*, PID) const; @@ -114,25 +136,47 @@ class QuantizedGraph { void reconstruct_quantized_vector(PID, float*) const; - [[nodiscard]] char* get_batch_data(PID data_id); + [[nodiscard]] char* get_batch_data(PID data_id) { + return get_row_data(data_id) + batch_data_offset_; + } - [[nodiscard]] const char* get_batch_data(PID data_id) const; + [[nodiscard]] const char* get_batch_data(PID data_id) const { + return get_row_data(data_id) + batch_data_offset_; + } - [[nodiscard]] rabitqlib::detail::PackedArrayView get_neighbors(PID data_id); + [[nodiscard]] rabitqlib::detail::PackedArrayView get_neighbors(PID data_id) { + return rabitqlib::detail::PackedArrayView( + get_row_data(data_id) + neighbor_offset_ + ); + } [[nodiscard]] rabitqlib::detail::ConstPackedArrayView get_neighbors(PID data_id - ) const; + ) const { + return rabitqlib::detail::ConstPackedArrayView( + get_row_data(data_id) + neighbor_offset_ + ); + } - void - find_candidates(PID, size_t, std::vector>&, VisitedSet&, const std::vector&) - const; + void find_candidates( + PID, + size_t, + std::vector>&, + VisitedSet&, + const std::vector&, + const BuildCache* = nullptr + ) const; - void update_qg(PID, const std::vector>&); + void update_qg( + PID, + const std::vector>&, + const BuildCache* = nullptr, + std::vector* = nullptr + ); void update_results(buffer::SearchBuffer&, VisitedSet&, const float*, const QuantizedQuery*); - void scan_neighbors( + inline void scan_neighbors( const BatchQuery&, PID, float*, @@ -148,7 +192,8 @@ class QuantizedGraph { size_t max_deg, MetricType metric_type = METRIC_L2, RotatorType rotator_type = RotatorType::FhtKacRotator, - size_t quantization_bits = 0 + size_t quantization_bits = 0, + uint32_t seed = std::random_device{}() ); explicit QuantizedGraph(); @@ -189,6 +234,38 @@ class QuantizedGraph { uint32_t* __restrict__ results, float* __restrict__ dists ); + /** + * Search using caller-owned scratch buffers. + * + * This avoids per-query allocations and can use a previous result as a + * warm-start hint. The search and result buffers must be sized for ef and + * knn respectively. rotated_query_scratch, est_dist_scratch, and + * lut_float_scratch require padded_dim(), degree_bound(), and + * padded_dim() * 4 elements respectively. The visited set must be + * initialized before the first search; an uninitialized set raises + * std::invalid_argument. Concurrent searches require separate scratch + * buffers, including visited sets. + * + * An optional k1xsumq must come from BatchQuery::k1xsumq() for the same + * query and this graph's rotation. A null pointer recomputes the sum. + */ + void search_with_scratch( + const float* __restrict__ query, + uint32_t knn, + uint32_t* __restrict__ results, + float* __restrict__ dists, + float* __restrict__ rotated_query_scratch, + float* __restrict__ est_dist_scratch, + float* __restrict__ lut_float_scratch, + BatchQuery& batch_query, + buffer::SearchBuffer& search_pool, + buffer::SearchBuffer& result_pool, + VisitedSet& visited, + PID hint = kPidMax, + const float* k1xsumq = nullptr + ); + + [[nodiscard]] size_t padded_dim() const { return this->padded_dim_; } }; // Preserve C++17 deduction for callers that omit the float template argument. @@ -199,7 +276,8 @@ QuantizedGraph( size_t, MetricType = METRIC_L2, RotatorType = RotatorType::FhtKacRotator, - size_t = 0 + size_t = 0, + uint32_t = std::random_device{}() ) ->QuantizedGraph; diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 42bce716..cab3ecda 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -47,6 +48,15 @@ class QGBuilder { std::vector pruned_neighbors_; // recorded pruned neighbors std::vector visited_list_; // per-thread visited sets std::vector degrees_; // record degree of qg + uint32_t seed_; // base seed for graph initialization + std::optional::BuildCache> cache_; + std::vector candidates_; + std::vector reverse_buffer_; + std::vector> rotated_neighbors_; + + const QuantizedGraph::BuildCache* build_cache() const { + return cache_ ? &*cache_ : nullptr; + } void random_init(); void search_new_neighbors(bool refine); void heuristic_prune(PID, CandidateList&, CandidateList&, bool); @@ -59,13 +69,24 @@ class QGBuilder { void initialize_storage(const float* data); - QGBuilder(QuantizedGraph& index, uint32_t ef_build, size_t num_threads) + QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + size_t num_threads, + uint32_t seed = std::random_device{}(), + bool cache_vectors = false + ) : qg_{index} , ef_build_{ef_build} , num_threads_{resolve_num_threads(num_threads)} , num_nodes_{qg_.num_vertices()} , dim_{qg_.dimension()} - , degree_bound_(qg_.degree_bound()) {} + , degree_bound_(qg_.degree_bound()) + , seed_(seed) { + if (cache_vectors) { + cache_.emplace(); + } + } public: explicit QGBuilder( @@ -73,22 +94,50 @@ class QGBuilder { uint32_t ef_build, const float* data, size_t num_threads = std::numeric_limits::max(), - QGInitialization init = QGInitialization::PiPNN + QGInitialization init = QGInitialization::PiPNN, + uint32_t seed = std::random_device{}(), + bool cache_vectors = false ) - : QGBuilder(index, ef_build, num_threads) { - if (data == nullptr) { - throw std::invalid_argument("QGBuilder data must not be null"); - } + : QGBuilder(index, ef_build, num_threads, seed, cache_vectors) { if (init != QGInitialization::PiPNN && init != QGInitialization::Random) { throw std::invalid_argument("Unknown QG initialization"); } - if (init == QGInitialization::PiPNN) { - auto seed = detail::build_initial_graph( + seeded_ = init == QGInitialization::PiPNN; + reset(data); + } + + // Preserve the seeded random-initialization overload used before PiPNN. + QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + const float* data, + size_t num_threads, + uint32_t seed, + bool cache_vectors = false + ) + : QGBuilder( + index, + ef_build, + data, + num_threads, + QGInitialization::Random, + seed, + cache_vectors + ) {} + + /** Replace vectors and restart the selected initializer, retaining workspace. + * Cached vectors and queries are owned; input may be released after reset. + */ + void reset(const float* data) { + if (data == nullptr) { + throw std::invalid_argument("QGBuilder data must not be null"); + } + if (seeded_) { + auto initial = detail::build_initial_graph( data, num_nodes_, dim_, degree_bound_, qg_.metric_type_, num_threads_ ); - initialize_seed(data, seed.offsets, seed.neighbors); - } else if (init == QGInitialization::Random) { - seeded_ = false; + initialize_seed(data, initial.offsets, initial.neighbors); + } else { initialize_storage(data); random_init(); } @@ -104,7 +153,9 @@ class QGBuilder { offsets.back() != neighbors.size()) { throw std::invalid_argument("Seed graph offsets must delimit every vertex"); } - for (size_t i = 0; i < num_nodes_; ++i) { + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || offsets[i + 1] - offsets[i] > degree_bound_) { throw std::invalid_argument( @@ -143,7 +194,12 @@ class QGBuilder { row.emplace_back(id, 0.0F); } degrees_[i] = row.size(); - qg_.update_qg(i, row); + qg_.update_qg( + i, + row, + build_cache(), + &rotated_neighbors_[static_cast(omp_get_thread_num())] + ); } } } @@ -208,20 +264,63 @@ inline void QGBuilder::initialize_storage(const float* data) { if (data == nullptr) { throw std::invalid_argument("QGBuilder data must not be null"); } - // Allocate refinement scratch only after PiPNN's dense workspace is released. + // Allocate refinement scratch after PiPNN releases its dense workspace. + // Clear old graph-dependent state on reset while retaining vector capacities. new_neighbors_.resize(num_nodes_); pruned_neighbors_.resize(num_nodes_); - visited_list_ = std::vector( - num_threads_, - VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) - ); + reverse_buffer_.resize(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); + new_neighbors_[i].clear(); + pruned_neighbors_[i].clear(); + reverse_buffer_[i].clear(); + } + if (visited_list_.empty()) { + visited_list_ = std::vector( + num_threads_, + VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + ); + } degrees_.assign(num_nodes_, degree_bound_); + candidates_.resize(num_threads_); + rotated_neighbors_.resize(num_threads_); std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); qg_.ready_ = false; qg_.set_quantization_centroid(centroid.data()); qg_.copy_vectors(data, num_threads_); + if (cache_) { + const size_t padded_dim = qg_.padded_dim_; + cache_->rotated_vectors.resize(num_nodes_ * padded_dim); +#pragma omp parallel for num_threads(num_threads_) schedule(static) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t id = static_cast(index); + float* row = cache_->rotated_vectors.data() + id * padded_dim; + if (qg_.is_quantized()) { + qg_.reconstruct_quantized_vector(id, row); + } else { + qg_.rotator_->rotate(qg_.get_vector(id), row); + } + } + cache_->queries.clear(); + if (qg_.is_quantized()) { + cache_->queries.reserve(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t id = static_cast(index); + cache_->queries.emplace_back( + cache_->rotated_vectors.data() + id * padded_dim, + qg_.centroid_.data(), + padded_dim, + qg_.metric_type_ + ); + } + } + } + PID entry_point = 0; if (qg_.is_quantized()) { QuantizedQuery query( @@ -265,7 +364,8 @@ inline void QGBuilder::add_pruned_edges( while (new_result.size() < degree_bound_ && start < pruned_list.size()) { const auto& cur = pruned_list[start]; bool occlude = false; - const float* cur_data = qg_.prepare_build_query(cur.id, reconstructed, prepared); + const float* cur_data = + qg_.prepare_build_query(cur.id, reconstructed, prepared, build_cache()); float dik_sqr = cur.distance; if (nei_set.find(cur.id) != nei_set.end()) { @@ -330,7 +430,7 @@ inline void QGBuilder::heuristic_prune( pruned_results.emplace_back(pool[start]); // add current candidate to result const float* data_j = - qg_.prepare_build_query(candidate_id, reconstructed, prepared); + qg_.prepare_build_query(candidate_id, reconstructed, prepared, build_cache()); // i : current vertex // j : neighbor added in this iter @@ -361,24 +461,26 @@ inline void QGBuilder::heuristic_prune( * @param refine refine = true means recording pruned candidates */ inline void QGBuilder::search_new_neighbors(bool refine) { -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const size_t i = static_cast(index); PID cur_id = i; auto tid = omp_get_thread_num(); - CandidateList candidates; + CandidateList& candidates = candidates_[tid]; + candidates.clear(); VisitedSet& vis = visited_list_[tid]; candidates.reserve(2 * kMaxCandidatePoolSize); vis.clear(); - qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_); + qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_, build_cache()); // Seeded construction keeps its initial edges only in QG. Materialize // their scores on demand, after the caller can release the input/CSR. if (new_neighbors_[cur_id].empty() && degrees_[cur_id] != 0) { std::vector reconstructed; std::optional prepared; - const float* source = qg_.prepare_build_query(cur_id, reconstructed, prepared); + const float* source = + qg_.prepare_build_query(cur_id, reconstructed, prepared, build_cache()); const auto ids = qg_.get_neighbors(cur_id); for (size_t j = 0; j < degrees_[cur_id]; ++j) { if (ids[j] != cur_id && !vis.get(ids[j])) { @@ -414,11 +516,13 @@ inline void QGBuilder::search_new_neighbors(bool refine) { inline void QGBuilder::add_reverse_edges(bool refine) { std::vector locks(num_nodes_); - std::vector reverse_buffer(num_nodes_); + for (auto& row : reverse_buffer_) { + row.clear(); + } // Keep new_neighbors_ read-only while reverse candidates are collected. Mutating a // destination row here races with another worker reading that row as its source. -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const PID data_id = static_cast(index); @@ -439,23 +543,23 @@ inline void QGBuilder::add_reverse_edges(bool refine) { std::lock_guard lock(locks[destination]); const size_t missing_slots = degree_bound_ - std::min(degree_bound_, destination_neighbors.size()); - if (reverse_buffer[destination].size() < + if (reverse_buffer_[destination].size() < kMaxCandidatePoolSize + missing_slots) { - reverse_buffer[destination].emplace_back(data_id, nei.distance); + reverse_buffer_[destination].emplace_back(data_id, nei.distance); } } } -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const PID data_id = static_cast(index); - CandidateList& tmp_pool = reverse_buffer[data_id]; + CandidateList& tmp_pool = reverse_buffer_[data_id]; if (qg_.is_quantized() && !tmp_pool.empty()) { // RaBitQ estimates are directional: score destination -> source afresh. std::vector reconstructed; std::optional prepared; - qg_.prepare_build_query(data_id, reconstructed, prepared); + qg_.prepare_build_query(data_id, reconstructed, prepared, build_cache()); for (auto& candidate : tmp_pool) { candidate.distance = qg_.quantized_distance(*prepared, candidate.id); } @@ -472,14 +576,18 @@ inline void QGBuilder::add_reverse_edges(bool refine) { inline void QGBuilder::random_init() { const PID min_id = 0; const PID max_id = num_nodes_ - 1; -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const size_t i = static_cast(index); + new_neighbors_[i].clear(); + pruned_neighbors_[i].clear(); + std::mt19937 rng(seed_ ^ static_cast(i)); + std::uniform_int_distribution distribution(min_id, max_id); std::unordered_set neighbor_set; neighbor_set.reserve(degree_bound_); while (neighbor_set.size() < degree_bound_) { - PID rand_id = rand_integer(min_id, max_id); + const PID rand_id = distribution(rng); if (rand_id != i) { neighbor_set.emplace(rand_id); } @@ -487,7 +595,8 @@ inline void QGBuilder::random_init() { std::vector reconstructed; std::optional prepared; - const float* cur_data = qg_.prepare_build_query(i, reconstructed, prepared); + const float* cur_data = + qg_.prepare_build_query(i, reconstructed, prepared, build_cache()); new_neighbors_[i].reserve(degree_bound_); for (PID cur_neigh : neighbor_set) { new_neighbors_[i].emplace_back( @@ -497,7 +606,12 @@ inline void QGBuilder::random_init() { } degrees_[i] = new_neighbors_[i].size(); - qg_.update_qg(i, new_neighbors_[i]); + qg_.update_qg( + i, + new_neighbors_[i], + build_cache(), + &rotated_neighbors_[static_cast(omp_get_thread_num())] + ); } } @@ -507,7 +621,7 @@ inline void QGBuilder::random_init() { * */ inline void QGBuilder::graph_refine() { -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const size_t i = static_cast(index); @@ -526,6 +640,8 @@ inline void QGBuilder::graph_refine() { std::sort(pruned_list.begin(), pruned_list.end()); // use binary search to get refined results + CandidateList successful_result; + successful_result.reserve(degree_bound_); float left = 0.5; float right = 1.0; size_t iter = 0; @@ -536,14 +652,23 @@ inline void QGBuilder::graph_refine() { left = mid; } else { right = mid; + successful_result = new_result; } } - // update neighbors with larger cosine value since we want to retain more edges - add_pruned_edges(cur_neighbors, pruned_list, new_result, right); + // Every successful right endpoint was already evaluated at that exact threshold. + if (!successful_result.empty()) { + new_result.swap(successful_result); + } else { + add_pruned_edges(cur_neighbors, pruned_list, new_result, right); + } // if the vertex still doesn't have enough neighbors, use random vertices if (new_result.size() < degree_bound_) { + std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); + std::uniform_int_distribution distribution( + 0, static_cast(num_nodes_) - 1 + ); std::unordered_set ids; ids.reserve(degree_bound_); for (auto& neighbor : new_result) { @@ -551,9 +676,10 @@ inline void QGBuilder::graph_refine() { } std::vector reconstructed; std::optional prepared; - const float* source = qg_.prepare_build_query(i, reconstructed, prepared); + const float* source = + qg_.prepare_build_query(i, reconstructed, prepared, build_cache()); while (new_result.size() < degree_bound_) { - PID rand_id = rand_integer(0, static_cast(num_nodes_) - 1); + const PID rand_id = distribution(rng); if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { new_result.emplace_back( rand_id, @@ -570,7 +696,9 @@ inline void QGBuilder::graph_refine() { inline void QGBuilder::iter(bool refine) { if (refine) { - for (size_t i = 0; i < num_nodes_; ++i) { + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); pruned_neighbors_[i].clear(); pruned_neighbors_[i].reserve(kMaxPrunedSize); } @@ -586,11 +714,16 @@ inline void QGBuilder::iter(bool refine) { } // update qg -#pragma omp parallel for schedule(dynamic) num_threads(num_threads_) +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); ++index) { const size_t i = static_cast(index); - qg_.update_qg(i, new_neighbors_[i]); + qg_.update_qg( + i, + new_neighbors_[i], + build_cache(), + &rotated_neighbors_[static_cast(omp_get_thread_num())] + ); degrees_[i] = new_neighbors_[i].size(); } } diff --git a/include/rabitqlib/simd/matrix_dispatch.hpp b/include/rabitqlib/simd/matrix_dispatch.hpp index aadbd8c1..37a6b642 100644 --- a/include/rabitqlib/simd/matrix_dispatch.hpp +++ b/include/rabitqlib/simd/matrix_dispatch.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include namespace rabitqlib::simd { // Dense row-major float32 kernels; inputs and outputs must not overlap. @@ -120,4 +121,47 @@ void pairwise_distances_lower_avx512( size_t dim, bool inner_product ); + +// Accumulate input rows in point_ids order into double cluster sums and counts. +// Inputs and outputs must not overlap. Concurrent calls require disjoint clusters. +void accumulate_cluster_sums( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +); + +void accumulate_cluster_sums_generic( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +); + +void accumulate_cluster_sums_avx2( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +); + +void accumulate_cluster_sums_avx512( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +); + } // namespace rabitqlib::simd diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 61ab31ff..0e5f3215 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -143,14 +143,14 @@ class FhtKacRotator : public Rotator { public: static constexpr size_t kMaxDim = 65536; - explicit FhtKacRotator(size_t dim, size_t padded_dim) + explicit FhtKacRotator( + size_t dim, size_t padded_dim, uint32_t seed = std::random_device{}() + ) : Rotator(dim, padded_dim) { if (dim < 64 || padded_dim < dim || padded_dim % 64 != 0 || padded_dim > kMaxDim) { throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); } flip_.resize((padded_dim / kByteLen) * 4); - std::random_device rd; - std::seed_seq seed{rd(), rd(), rd(), rd()}; std::mt19937 gen(seed); // Uniform distribution in the range [0, 255] @@ -214,7 +214,10 @@ class FhtKacRotator : public Rotator { // for given dim & type, set rotator, return padded dimension template Rotator* choose_rotator( - size_t dim, RotatorType type = RotatorType::FhtKacRotator, size_t padded_dim = 0 + size_t dim, + RotatorType type = RotatorType::FhtKacRotator, + size_t padded_dim = 0, + uint32_t seed = std::random_device{}() ) { if (dim == 0) { throw std::invalid_argument("Rotator dimension must be positive"); @@ -237,7 +240,7 @@ Rotator* choose_rotator( if (!std::is_same_v) { throw std::invalid_argument("FhtKacRotator only supports float"); } - return ::new rotator_impl::FhtKacRotator(dim, padded_dim); + return ::new rotator_impl::FhtKacRotator(dim, padded_dim, seed); } if (type == RotatorType::MatrixRotator) { diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/utils/visited_set_hash.hpp index d6aca331..425aa934 100644 --- a/include/rabitqlib/utils/visited_set_hash.hpp +++ b/include/rabitqlib/utils/visited_set_hash.hpp @@ -81,6 +81,8 @@ class HashBasedVisitedSet { stl_hash_.clear(); } + [[nodiscard]] bool initialized() const noexcept { return !table_.empty(); } + [[nodiscard]] bool get(PID data_id) const { PID val = this->table_[hash1(data_id)]; if (val == data_id) { diff --git a/pyproject.toml b/pyproject.toml index df89b4c8..9008db1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "scikit_build_core.build" [project] name = "rabitqlib" version = "0.4.0" -description = "RaBitQ Python bindings for HNSW, IVF, and SymQG" +description = "RaBitQ Python bindings for indexes and clustering" readme = "README.md" requires-python = ">=3.11" classifiers = [ diff --git a/python/ivf.py b/python/ivf.py index dadc3c69..5bb589d4 100644 --- a/python/ivf.py +++ b/python/ivf.py @@ -1,37 +1,38 @@ -import sys +import argparse from time import time -import faiss +from rabitqlib import FinalAssignmentMode, QGKMeans, RaBitQKMeans from utils.io import read_fvecs, write_fvecs, write_ivecs if __name__ == "__main__": - if len(sys.argv) < 5: - print(f"Usage: {sys.argv[0]} ") - print("arg1: path for data file, format .fvecs") - print("arg2: number of clusters") - print("arg3: path for centroid vectors") - print("arg4: path for cluster ids") - print("arg5: distance metric") - exit(1) - - # path - data_path = sys.argv[1] - K = int(sys.argv[2]) - centroids_path = sys.argv[3] - cluster_id_path = sys.argv[4] - if len(sys.argv) == 6: - distance_metric = sys.argv[5].lower() - if distance_metric == "l2": - metric = faiss.METRIC_L2 - print("Using L2 metric") - elif distance_metric in ["ip", "innerproduct"]: - metric = faiss.METRIC_INNER_PRODUCT - print("Using InnerProduct metric") - else: - raise ValueError("Unsupported distance metric. Use 'l2' or 'ip'.") - else: - metric = faiss.METRIC_L2 # by default, L2 metric - print("Using L2 metric by default") + parser = argparse.ArgumentParser( + description="Cluster vectors and save fvecs/ivecs outputs." + ) + parser.add_argument("data_path") + parser.add_argument("num_clusters", type=int) + parser.add_argument("centroids_path") + parser.add_argument("cluster_id_path") + parser.add_argument( + "metric", + nargs="?", + default="l2", + type=str.lower, + choices=["l2", "ip", "innerproduct"], + ) + parser.add_argument( + "--method", + choices=["rabitq", "qg"], + default="qg", + help="rabitq for few centroids; qg for many", + ) + args = parser.parse_args() + data_path = args.data_path + K = args.num_clusters + centroids_path = args.centroids_path + cluster_id_path = args.cluster_id_path + metric = "l2" if args.metric == "l2" else "ip" + clustering_type = RaBitQKMeans if args.method == "rabitq" else QGKMeans + print(f"Using {metric.upper()} metric with {clustering_type.__name__}") X = read_fvecs(data_path) @@ -40,15 +41,20 @@ t1 = time() # cluster data vectors - index = faiss.index_factory(dim, f"IVF{K},Flat", metric) - index.verbose = True - index.train(X) + kmeans = clustering_type( + dim, + K, + spherical=metric == "ip", + final_assignment=FinalAssignmentMode.Exact, + verbose=True, + ) + kmeans.train(X) t2 = time() print(f"Time for training ivf {t2 - t1} secs") - centroids = index.quantizer.reconstruct_n(0, index.nlist) - _, cluster_id = index.quantizer.search(X, 1) + centroids = kmeans.centroids + cluster_id = kmeans.assignments.reshape(-1, 1) write_ivecs(cluster_id_path, cluster_id) write_fvecs(centroids_path, centroids) diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index a1067da1..082f0e4e 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -20,6 +20,7 @@ pybind11_add_module(_rabitqlib hnsw_bindings.cpp ivf_bindings.cpp symqg_bindings.cpp + kmeans_bindings.cpp ) target_include_directories(_rabitqlib PRIVATE diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index 41143fb9..ca99e806 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -1,9 +1,94 @@ -"""Python bindings for RaBitQLib's three index types. +"""Python bindings for RaBitQLib indexes and clustering. This module re-exports symbols from the compiled `rabitqlib` extension so code can import using `from python_bindings import ...` if needed. """ -from ._rabitqlib import HnswIndex, IvfIndex, SymqgIndex +from ._rabitqlib import ( + FinalAssignmentMode, + HnswIndex, + IvfIndex, + QGKMeansIterationStats, + QGKMeansParameters, + RaBitQKMeansIterationStats, + RaBitQKMeansParameters, + SymqgIndex, + _QGKMeans, + _RaBitQKMeans, +) -__all__ = ["HnswIndex", "IvfIndex", "SymqgIndex"] + +class _KMeans: + """Shared result handling for the two clustering methods.""" + + def __init__(self, d, k, **kwargs): + self.d = int(d) + self.k = int(k) + self.cp = self._parameters_type() + for name, value in kwargs.items(): + getattr(self.cp, name) + setattr(self.cp, name, value) + self._impl = None + + def train(self, x): + """Train on ``(n, d)`` data and return the last pre-update objective. + + ``final_obj`` measures the returned assignments against the final centroids. + A failed fit preserves the results of the previous successful fit. + """ + implementation = self._implementation_type(self.d, self.k, self.cp) + objective = implementation.train(x) + self._impl = implementation + return objective + + @property + def centroids(self): + return None if self._impl is None else self._impl.centroids + + @property + def assignments(self): + return None if self._impl is None else self._impl.assignments + + @property + def distances(self): + return None if self._impl is None else self._impl.distances + + @property + def obj(self): + return None if self._impl is None else self._impl.obj + + @property + def final_obj(self): + return None if self._impl is None else self._impl.final_obj + + @property + def iteration_stats(self): + return None if self._impl is None else self._impl.iteration_stats + + +class QGKMeans(_KMeans): + """SymphonyQG graph assignment; recommended for many centroids.""" + + _parameters_type = QGKMeansParameters + _implementation_type = _QGKMeans + + +class RaBitQKMeans(_KMeans): + """Flat RaBitQ assignment; recommended for few centroids.""" + + _parameters_type = RaBitQKMeansParameters + _implementation_type = _RaBitQKMeans + + +__all__ = [ + "FinalAssignmentMode", + "HnswIndex", + "IvfIndex", + "QGKMeans", + "QGKMeansIterationStats", + "QGKMeansParameters", + "RaBitQKMeans", + "RaBitQKMeansIterationStats", + "RaBitQKMeansParameters", + "SymqgIndex", +] diff --git a/python_bindings/kmeans_bindings.cpp b/python_bindings/kmeans_bindings.cpp new file mode 100644 index 00000000..a889934c --- /dev/null +++ b/python_bindings/kmeans_bindings.cpp @@ -0,0 +1,154 @@ +#include +#include +#include +#include +#include +#include +#include // IWYU pragma: keep (vector conversion for iteration_stats) + +#include +#include +#include +#include + +#include "bindings_common.hpp" +#include "rabitqlib/clustering/qgkmeans.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" + +namespace py = pybind11; + +namespace rabitqlib::python_bindings { + +template +py::array_t vector_to_array(const std::vector& values) { + py::array_t result(static_cast(values.size())); + if (!values.empty()) { + std::memcpy(result.mutable_data(), values.data(), values.size() * sizeof(T)); + } + return result; +} + +template +py::array_t centroids_to_array(const Clustering& clustering, const char* name) { + if (clustering.centroids.empty()) { + throw std::logic_error(std::string(name) + " has not been trained"); + } + const auto shape = std::vector{ + static_cast(clustering.k), static_cast(clustering.d)}; + py::array_t result(shape); + std::memcpy( + result.mutable_data(), + clustering.centroids.data(), + clustering.centroids.size() * sizeof(float) + ); + return result; +} + +template +double train(Clustering& clustering, py::handle x) { + auto x_array = ensure_2d_array(x, "x"); + if (static_cast(x_array.shape(1)) != clustering.d) { + throw std::invalid_argument("x dimension does not match d"); + } + + { + py::gil_scoped_release release; + clustering.train(static_cast(x_array.shape(0)), x_array.data()); + } + return clustering.iteration_stats.back().obj; +} + +template +void bind_training_parameters(py::class_& binding) { + binding.def(py::init<>()) + .def_readwrite("niter", &Parameters::niter) + .def_readwrite("verbose", &Parameters::verbose) + .def_readwrite("spherical", &Parameters::spherical) + .def_readwrite("seed", &Parameters::seed) + .def_readwrite("min_points_per_centroid", &Parameters::min_points_per_centroid) + .def_readwrite("early_stop_threshold", &Parameters::early_stop_threshold) + .def_readwrite("num_threads", &Parameters::num_threads) + .def_readwrite("final_assignment", &Parameters::final_assignment); +} + +template +void bind_clustering(py::module_& m, const char* native_name, const char* public_name) { + py::class_(m, native_name) + .def(py::init()) + .def("train", &train, py::arg("x")) + .def_readonly("d", &Clustering::d) + .def_readonly("k", &Clustering::k) + .def_readonly("final_obj", &Clustering::final_obj) + .def_property_readonly( + "centroids", + [public_name](const Clustering& clustering) { + return centroids_to_array(clustering, public_name); + } + ) + .def_property_readonly( + "assignments", + [](const Clustering& clustering) { + return vector_to_array(clustering.assignments); + } + ) + .def_property_readonly( + "distances", + [](const Clustering& clustering) { + return vector_to_array(clustering.distances); + } + ) + .def_property_readonly( + "obj", + [](const Clustering& clustering) { + std::vector values; + values.reserve(clustering.iteration_stats.size()); + for (const auto& stats : clustering.iteration_stats) { + values.push_back(stats.obj); + } + return vector_to_array(values); + } + ) + .def_readonly("iteration_stats", &Clustering::iteration_stats); +} + +} // namespace rabitqlib::python_bindings + +void register_kmeans(py::module_& m) { + using namespace rabitqlib::python_bindings; + using rabitqlib::qgkmeans::FinalAssignmentMode; + using rabitqlib::qgkmeans::QGKMeans; + using rabitqlib::qgkmeans::QGKMeansIterationStats; + using rabitqlib::qgkmeans::QGKMeansParameters; + using rabitqlib::rabitqkmeans::RaBitQKMeans; + using rabitqlib::rabitqkmeans::RaBitQKMeansParameters; + + py::enum_(m, "FinalAssignmentMode") + .value("Approximate", FinalAssignmentMode::Approximate) + .value("SymphonyQG", FinalAssignmentMode::SymphonyQG) + .value("Exact", FinalAssignmentMode::Exact); + + py::class_ qg_parameters(m, "QGKMeansParameters"); + bind_training_parameters(qg_parameters); + qg_parameters.def_readwrite("graph_degree", &QGKMeansParameters::graph_degree) + .def_readwrite("ef_build", &QGKMeansParameters::ef_build) + .def_readwrite("ef_search", &QGKMeansParameters::ef_search) + .def_readwrite( + "graph_build_iterations", &QGKMeansParameters::graph_build_iterations + ) + .def_readwrite("quantization_bits", &QGKMeansParameters::quantization_bits); + + py::class_ flat_parameters(m, "RaBitQKMeansParameters"); + bind_training_parameters(flat_parameters); + + py::class_(m, "QGKMeansIterationStats") + .def_readonly("iteration", &QGKMeansIterationStats::iteration) + .def_readonly("obj", &QGKMeansIterationStats::obj) + .def_readonly("shift", &QGKMeansIterationStats::shift) + .def_readonly("nsplit", &QGKMeansIterationStats::nsplit); + m.attr("RaBitQKMeansIterationStats") = m.attr("QGKMeansIterationStats"); + + bind_clustering(m, "_QGKMeans", "QGKMeans"); + bind_clustering( + m, "_RaBitQKMeans", "RaBitQKMeans" + ); +} diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index 33a8dd24..d783867f 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -1,3 +1,4 @@ +#include #include #include "rabitqlib/defines.hpp" @@ -9,6 +10,7 @@ namespace py = pybind11; void register_hnsw(py::module_& m); void register_ivf(py::module_& m); void register_symqg(py::module_& m); +void register_kmeans(py::module_& m); PYBIND11_MODULE(_rabitqlib, m) { m.doc() = "RabitQ Python bindings combined module"; @@ -28,4 +30,5 @@ PYBIND11_MODULE(_rabitqlib, m) { register_hnsw(m); register_ivf(m); register_symqg(m); + register_kmeans(m); } diff --git a/sample/cpp/CMakeLists.txt b/sample/cpp/CMakeLists.txt index a6f3cc00..7045e9d6 100644 --- a/sample/cpp/CMakeLists.txt +++ b/sample/cpp/CMakeLists.txt @@ -2,6 +2,8 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin) add_executable(symqg_indexing symqg_indexing.cpp) add_executable(symqg_querying symqg_querying.cpp) +add_executable(qgkmeans qgkmeans.cpp) +add_executable(rabitqkmeans rabitqkmeans.cpp) add_executable(ivf_rabitq_indexing ivf_rabitq_indexing.cpp) add_executable(ivf_rabitq_querying ivf_rabitq_querying.cpp) @@ -12,6 +14,8 @@ add_executable(hnsw_rabitq_querying hnsw_rabitq_querying.cpp) foreach(RABITQ_SAMPLE_TARGET symqg_indexing symqg_querying + qgkmeans + rabitqkmeans ivf_rabitq_indexing ivf_rabitq_querying hnsw_rabitq_indexing diff --git a/sample/cpp/ivf_rabitq_indexing.cpp b/sample/cpp/ivf_rabitq_indexing.cpp index da53485e..a4fddbb7 100644 --- a/sample/cpp/ivf_rabitq_indexing.cpp +++ b/sample/cpp/ivf_rabitq_indexing.cpp @@ -16,8 +16,8 @@ int run(int argc, char** argv) { if (argc < 6) { std::cerr << "Usage: " << argv[0] << " \n" << "arg1: path for data file, format .fvecs\n" - << "arg2: path for centroids file generated by ivf.py\n" - << "arg3: path for cluster ids file generated by ivf.py\n" + << "arg2: path for centroids file, format .fvecs\n" + << "arg3: path for cluster ids file, format .ivecs\n" << "arg4: total number of bits for quantization\n" << "arg5: path for saving index\n" << "arg6: metric type (\"l2\" or \"ip\"), l2 by default\n" diff --git a/sample/cpp/kmeans_example.hpp b/sample/cpp/kmeans_example.hpp new file mode 100644 index 00000000..9d9b241b --- /dev/null +++ b/sample/cpp/kmeans_example.hpp @@ -0,0 +1,72 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/io.hpp" + +namespace kmeans_example { + +template +void write_vecs(const char* path, const T* data, size_t rows, size_t cols) { + std::ofstream output; + output.exceptions(std::ios::failbit | std::ios::badbit); + output.open(rabitqlib::io_impl::filesystem_path(path), std::ios::binary); + const auto dimension = static_cast(cols); + for (size_t row = 0; row < rows; ++row) { + output.write(reinterpret_cast(&dimension), sizeof(dimension)); + output.write( + reinterpret_cast(data + row * cols), + static_cast(cols * sizeof(T)) + ); + } + output.close(); +} + +template +int cluster_file(int argc, char** argv) { + if (argc != 5 && argc != 6) { + std::cerr + << "Usage: " << argv[0] + << " [l2|ip]\n" + << "Without arguments, run a synthetic clustering demo.\n"; + return 1; + } + const std::string count(argv[2]); + size_t num_clusters = 0; + const auto parsed = + std::from_chars(count.data(), count.data() + count.size(), num_clusters); + if (parsed.ec != std::errc{} || parsed.ptr != count.data() + count.size() || + num_clusters == 0) { + throw std::invalid_argument("num_clusters must be a positive integer"); + } + const std::string metric = argc == 6 ? argv[5] : "l2"; + if (metric != "l2" && metric != "ip") { + throw std::invalid_argument("metric must be l2 or ip; ip requires normalized input" + ); + } + + rabitqlib::RowMajorArray data; + rabitqlib::load_vecs(argv[1], data); + Parameters parameters; + parameters.spherical = metric == "ip"; + parameters.final_assignment = decltype(parameters.final_assignment)::Exact; + parameters.verbose = true; + KMeans clustering(data.cols(), num_clusters, parameters); + clustering.train(data.rows(), data.data()); + + write_vecs(argv[3], clustering.centroids.data(), num_clusters, data.cols()); + write_vecs(argv[4], clustering.assignments.data(), data.rows(), 1); + std::cout << "Completed " << clustering.iteration_stats.size() + << " iterations; final objective=" << clustering.final_obj << '\n'; + return 0; +} + +} // namespace kmeans_example diff --git a/sample/cpp/qgkmeans.cpp b/sample/cpp/qgkmeans.cpp new file mode 100644 index 00000000..fd84b5e3 --- /dev/null +++ b/sample/cpp/qgkmeans.cpp @@ -0,0 +1,56 @@ +#include "rabitqlib/clustering/qgkmeans.hpp" + +#include +#include +#include +#include +#include + +#include "kmeans_example.hpp" + +int run_demo() { + constexpr size_t kNumClusters = 33; + constexpr size_t kPointsPerCluster = 40; + constexpr size_t kDimension = 64; + constexpr size_t kNumPoints = kNumClusters * kPointsPerCluster; + + std::mt19937 rng(42); + std::normal_distribution noise(0.0F, 0.05F); + std::vector data(kNumPoints * kDimension); + for (size_t point = 0; point < kNumPoints; ++point) { + const float center = static_cast(point / kPointsPerCluster); + for (size_t dim = 0; dim < kDimension; ++dim) { + data[point * kDimension + dim] = center + noise(rng); + } + } + + rabitqlib::qgkmeans::QGKMeansParameters parameters; + parameters.niter = 5; + parameters.graph_degree = 32; + parameters.ef_build = 64; + parameters.ef_search = 32; + parameters.quantization_bits = 8; + parameters.verbose = true; + + rabitqlib::qgkmeans::QGKMeans clustering(kDimension, kNumClusters, parameters); + clustering.train(kNumPoints, data.data()); + + std::cout << "Completed " << clustering.iteration_stats.size() + << " iterations; final assignment count=" << clustering.assignments.size() + << '\n'; + return 0; +} + +int main(int argc, char** argv) { + try { + if (argc == 1) { + return run_demo(); + } + return kmeans_example::cluster_file< + rabitqlib::qgkmeans::QGKMeans, + rabitqlib::qgkmeans::QGKMeansParameters>(argc, argv); + } catch (const std::exception& error) { + std::cerr << "Error: " << error.what() << '\n'; + return 1; + } +} diff --git a/sample/cpp/rabitqkmeans.cpp b/sample/cpp/rabitqkmeans.cpp new file mode 100644 index 00000000..5ec918db --- /dev/null +++ b/sample/cpp/rabitqkmeans.cpp @@ -0,0 +1,52 @@ +#include "rabitqlib/clustering/rabitqkmeans.hpp" + +#include +#include +#include +#include +#include + +#include "kmeans_example.hpp" + +int run_demo() { + constexpr size_t kNumClusters = 16; + constexpr size_t kPointsPerCluster = 40; + constexpr size_t kDimension = 64; + constexpr size_t kNumPoints = kNumClusters * kPointsPerCluster; + + std::mt19937 rng(42); + std::normal_distribution noise(0.0F, 0.05F); + std::vector data(kNumPoints * kDimension); + for (size_t point = 0; point < kNumPoints; ++point) { + const float center = static_cast(point / kPointsPerCluster); + for (size_t dim = 0; dim < kDimension; ++dim) { + data[point * kDimension + dim] = center + noise(rng); + } + } + + rabitqlib::rabitqkmeans::RaBitQKMeansParameters parameters; + parameters.niter = 5; + parameters.verbose = true; + + rabitqlib::rabitqkmeans::RaBitQKMeans clustering(kDimension, kNumClusters, parameters); + clustering.train(kNumPoints, data.data()); + + std::cout << "Completed " << clustering.iteration_stats.size() + << " iterations; final assignment count=" << clustering.assignments.size() + << '\n'; + return 0; +} + +int main(int argc, char** argv) { + try { + if (argc == 1) { + return run_demo(); + } + return kmeans_example::cluster_file< + rabitqlib::rabitqkmeans::RaBitQKMeans, + rabitqlib::rabitqkmeans::RaBitQKMeansParameters>(argc, argv); + } catch (const std::exception& error) { + std::cerr << "Error: " << error.what() << '\n'; + return 1; + } +} diff --git a/sample/python/README.md b/sample/python/README.md index 3f00b37f..577b2374 100644 --- a/sample/python/README.md +++ b/sample/python/README.md @@ -1,17 +1,20 @@ # Python examples -Run these commands from the repository root in your project Python environment. -Install `rabitqlib` and NumPy; the clustering step additionally needs `faiss-cpu`. +Run these commands from the repository root in your project Python environment +with `rabitqlib` and NumPy installed. Clustering and indexing both use RaBitQ +Library; FAISS is needed only for `compare_with_faiss.py`. ## IVF: cluster, then index -Faiss clustering and RaBitQ indexing run in separate Python processes. This keeps -their OpenMP runtimes separate, including when both packages come from pip wheels. -The intermediate `.npz` contains float32 `centroids`, uint32 `cluster_ids`, and the -`metric` (`l2` or `ip`). It can be reused for different RaBitQ bit widths. +`kmeans_clustering.py` uses `--method qg` (the default) for QGKMeans graph +assignment or `--method rabitq` for RaBitQKMeans flat assignment. Both make exact +final assignments. The saved `.npz` contains float32 `centroids`, uint32 +`cluster_ids`, and `metric` (`l2` or `ip`), and can be reused across index +configurations. Choose the method explicitly; the script does not switch based +on the cluster count. ```sh -python sample/python/faiss_clustering.py --num-clusters 4096 \ +python sample/python/kmeans_clustering.py --method qg --num-clusters 4096 \ data/gist/gist_base.fvecs data/gist/clusters_4096_l2.npz python sample/python/ivf_rabitq_indexing.py --total-bits 5 \ @@ -22,21 +25,23 @@ python sample/python/ivf_rabitq_querying.py \ data/gist/ivf_4096_5.index data/gist/gist_query.fvecs data/gist/gist_groundtruth.ivecs ``` -The indexing command now takes `--clusters` instead of `--num-clusters`; its cluster -count comes from the saved centroids. Always use the same base vectors in the same -row order for both stages. The loader checks the metric, dimensions, point count, -array types, and cluster ID range; it cannot detect reordered or replaced vectors. +Use the same base vectors in the same row order for clustering and indexing. +The loader checks the metric, dimensions, count, types, and ID range, but cannot +detect replaced or reordered vectors. See [QGKMeans](../../docs/docs/clustering.md#qgkmeans) +and [RaBitQKMeans](../../docs/docs/clustering.md#rabitqkmeans) for input limits and assignment +behavior. -For inner product, pass `--metric ip` to **both** clustering and indexing, and use -ground truth for that metric when evaluating recall. `--num-threads 0` selects the -hardware thread count; larger requests are capped. Each stage has its own setting. +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. ## HNSW -HNSW uses the same saved-cluster format. Choose its cluster count in the Faiss stage: +Use [RaBitQKMeans](../../docs/docs/clustering.md#rabitqkmeans) for this small-cluster +example, with the same clustering file format: ```sh -python sample/python/faiss_clustering.py --num-clusters 16 \ +python sample/python/kmeans_clustering.py --method rabitq --num-clusters 16 \ data/gist/gist_base.fvecs data/gist/clusters_16_l2.npz python sample/python/hnsw_rabitq_indexing.py --total-bits 5 --degree 16 \ @@ -44,10 +49,5 @@ python sample/python/hnsw_rabitq_indexing.py --total-bits 5 --degree 16 \ data/gist/gist_base.fvecs data/gist/hnsw_5.index ``` -HNSW querying is unchanged. SymphonyQG does not need Faiss clustering; run -`symqg_indexing.py` and `symqg_querying.py` directly. All querying scripts and the -shared file/recall utilities are independent of Faiss. - -Run the stages as separate commands, not by importing both libraries into one -Python process or notebook kernel. This separation fixes the example workflow; -it does not change native-runtime coexistence for applications importing both. +Query with `hnsw_rabitq_querying.py`. SymphonyQG builds directly from vectors; +use `symqg_indexing.py` and `symqg_querying.py` without a clustering stage. diff --git a/sample/python/compare_with_faiss.py b/sample/python/compare_with_faiss.py new file mode 100644 index 00000000..f57f42d1 --- /dev/null +++ b/sample/python/compare_with_faiss.py @@ -0,0 +1,306 @@ +"""Repeated, matched-settings comparison of QGKMeans, RaBitQKMeans, and FAISS.""" + +import argparse +import multiprocessing +import statistics +import time +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--input", + type=str, + help="optional float32 .npy or row-major .bin matrix", + ) + parser.add_argument("--n", type=int, default=1_000_000) + parser.add_argument("--d", type=int, default=256) + parser.add_argument("--k", type=int, default=4_000) + parser.add_argument("--niter", type=int, default=25) + parser.add_argument("--threads", type=int, default=32) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--num-seeds", type=int, default=1, help="number of consecutive training seeds" + ) + parser.add_argument( + "--repeats", + type=int, + default=1, + help="trials per seed, rotating train order", + ) + parser.add_argument( + "--spherical", + action="store_true", + help="normalize centroids and compare by inner product", + ) + return parser.parse_args() + + +def initialize_worker(input_path): + # Both workers map the same temporary input; only FAISS normalizes it, before + # training starts. Spawn keeps their native OpenMP runtimes independent. + global worker_data + worker_data = np.load(input_path, mmap_mode="r+") + + +def normalize_input(threads): + import faiss + + faiss.omp_set_num_threads(threads) + faiss.normalize_L2(worker_data) + worker_data.flush() + + +def evaluate_centroids(centroids, spherical, assignments, save_nearest, threads): + import faiss + + faiss.omp_set_num_threads(threads) + x = worker_data + nearest_out = np.empty(len(x), dtype=np.int64) if save_nearest else None + index = ( + faiss.IndexFlatIP(x.shape[1]) if spherical else faiss.IndexFlatL2(x.shape[1]) + ) + index.add(np.ascontiguousarray(centroids, dtype=np.float32)) + total = 0.0 + mismatches = 0 + for begin in range(0, x.shape[0], 100_000): + end = begin + 100_000 + distances, nearest = index.search(x[begin:end], 1) + if nearest_out is not None: + nearest_out[begin:end] = nearest[:, 0] + if spherical: + total += distances.shape[0] - distances.sum(dtype=np.float64) + else: + total += distances.sum(dtype=np.float64) + if assignments is not None: + mismatches += int(np.count_nonzero(assignments[begin:end] != nearest[:, 0])) + return total, mismatches, nearest_out + + +def exact_objective( + worker, threads, centroids, spherical, assignments=None, nearest_out=None +): + """Evaluate all methods in the FAISS worker, outside training time.""" + total, mismatches, nearest = worker.submit( + evaluate_centroids, + centroids, + spherical, + assignments, + nearest_out is not None, + threads, + ).result() + if nearest_out is not None: + nearest_out[:] = nearest + return total, mismatches + + +def train_method(method, args, seed): + x = worker_data + n, d = x.shape + common = dict( + niter=args.niter, + seed=seed, + spherical=args.spherical, + min_points_per_centroid=1, + early_stop_threshold=0.0, + ) + if method != "FAISS": + from rabitqlib import QGKMeans, RaBitQKMeans + + cls = QGKMeans if method == "QGKMeans" else RaBitQKMeans + kmeans = cls(d, args.k, num_threads=args.threads, **common) + else: + import faiss + + faiss.omp_set_num_threads(args.threads) + # Prevent subsampling so all methods train on all vectors. + kmeans = faiss.Kmeans( + d, + args.k, + nredo=1, + verbose=False, + gpu=False, + max_points_per_centroid=n, + **common, + ) + start = time.perf_counter() + kmeans.train(x) + elapsed = time.perf_counter() - start + result = dict( + seconds=elapsed, + iterations=len(kmeans.iteration_stats), + centroids=kmeans.centroids, + ) + if method != "FAISS": + result.update(assignments=kmeans.assignments, final_obj=kmeans.final_obj) + return result + + +def main(): + args = parse_args() + if args.num_seeds < 1 or args.repeats < 1: + raise ValueError("--num-seeds and --repeats must be positive") + + if args.input: + if args.input.endswith(".npy"): + x = np.load(args.input) + else: + x = np.fromfile(args.input, dtype=np.float32) + if x.size % args.d != 0: + raise ValueError("raw --input size must be divisible by --d") + x = x.reshape(-1, args.d) + x = np.ascontiguousarray(x, dtype=np.float32) + if x.ndim != 2: + raise ValueError("--input must contain a two-dimensional array") + else: + rng = np.random.default_rng(args.seed) + x = rng.standard_normal((args.n, args.d), dtype=np.float32) + + n, d = x.shape + if n < args.k: + raise ValueError(f"number of vectors ({n}) must be at least k ({args.k})") + + # Stage once, share pages across workers, and delete after both exit. Native + # imports must stay in workers: even importing FAISS can pin the parent CPU. + with TemporaryDirectory(prefix="qgkmeans-comparison-") as directory: + input_path = Path(directory) / "input.npy" + np.save(input_path, x) + del x + context = multiprocessing.get_context("spawn") + with ( + ProcessPoolExecutor( + max_workers=1, + mp_context=context, + initializer=initialize_worker, + initargs=(input_path,), + ) as faiss_worker, + ProcessPoolExecutor( + max_workers=1, + mp_context=context, + initializer=initialize_worker, + initargs=(input_path,), + ) as rabitq_worker, + ): + if args.spherical: + faiss_worker.submit(normalize_input, args.threads).result() + compare_trials(args, n, d, rabitq_worker, faiss_worker) + + +def compare_trials(args, n, d, rabitq_worker, faiss_worker): + metric = "cosine distance" if args.spherical else "squared L2" + print( + f"n={n}, d={d}, k={args.k}, requested niter={args.niter}, " + f"threads={args.threads}, metric={metric}, " + f"seeds={args.num_seeds}, repeats/seed={args.repeats}" + ) + print("Input data stays fixed across trials; training seeds change.") + print("The same seed does not give identical initial centroids across methods.") + print( + "Training includes the library methods' default approximate final assignment; " + "exhaustive evaluation is outside training time." + ) + methods = ["QGKMeans", "RaBitQKMeans", "FAISS"] + results = {method: [] for method in methods} + for seed_offset in range(args.num_seeds): + seed = args.seed + seed_offset + cache = {} + nearest = { + method: np.empty(n, dtype=np.int64) if args.repeats > 1 else None + for method in methods + if method != "FAISS" + } + for repeat in range(args.repeats): + offset = (seed_offset * args.repeats + repeat) % len(methods) + order = methods[offset:] + methods[:offset] + trained = {} + for method in order: + worker = faiss_worker if method == "FAISS" else rabitq_worker + trained[method] = worker.submit( + train_method, method, args, seed + ).result() + + print(f"seed={seed}, repeat={repeat + 1}, train order={', '.join(order)}") + print( + f"{'method':<12} {'train (s)':>12} {'iterations':>11} " + f"{'exhaustive objective':>21} {'per vector':>14}" + ) + for method in methods: + result = trained[method] + centroids = result["centroids"] + assignments = result.get("assignments") + previous = cache.get(method) + if previous is not None and np.array_equal(centroids, previous[0]): + objective = previous[1] + mismatches = ( + int(np.count_nonzero(assignments != nearest[method])) + if assignments is not None + else 0 + ) + else: + objective, mismatches = exact_objective( + faiss_worker, + args.threads, + centroids, + args.spherical, + assignments, + nearest.get(method), + ) + cache[method] = (centroids.copy(), objective) + results[method].append((result["seconds"], objective)) + print( + f"{method:<12} {result['seconds']:>12.6f} " + f"{result['iterations']:>11} {objective:>21.6g} " + f"{objective / n:>14.6g}" + ) + if assignments is not None: + print( + f"{method} returned-label objective: {result['final_obj']:.6g}; " + f"label mismatches vs exhaustive: {mismatches}/{n} " + f"({mismatches / n:.4%})" + ) + + print("Training time over all trials (median [min, max], seconds):") + for method in methods: + times = [seconds for seconds, _ in results[method]] + print( + f" {method}: {statistics.median(times):.6f} " + f"[{min(times):.6f}, {max(times):.6f}]" + ) + for method in methods[:-1]: + paired = list(zip(results[method], results["FAISS"])) + ratios = [seconds / baseline for (seconds, _), (baseline, _) in paired] + print( + f"{method}/FAISS paired training-time ratio: " + f"median {statistics.median(ratios):.3f}x " + f"[{min(ratios):.3f}x, {max(ratios):.3f}x]" + ) + differences = [objective - baseline for (_, objective), (_, baseline) in paired] + print( + f"{method} exhaustive objective difference vs FAISS (absolute): " + f"median {statistics.median(differences):+.6g} " + f"[{min(differences):+.6g}, {max(differences):+.6g}]" + ) + if all(baseline != 0 for _, (_, baseline) in paired): + percentages = [ + (objective / baseline - 1) * 100 + for (_, objective), (_, baseline) in paired + ] + print( + f"{method} exhaustive objective difference vs FAISS (%): " + f"median {statistics.median(percentages):+.4f}% " + f"[{min(percentages):+.4f}%, {max(percentages):+.4f}%]" + ) + else: + print( + f"{method} exhaustive objective difference vs FAISS (%): N/A " + "(at least one FAISS objective is zero)" + ) + + +if __name__ == "__main__": + main() diff --git a/sample/python/hnsw_rabitq_indexing.py b/sample/python/hnsw_rabitq_indexing.py index 1aec6250..7f2e4f6e 100644 --- a/sample/python/hnsw_rabitq_indexing.py +++ b/sample/python/hnsw_rabitq_indexing.py @@ -22,7 +22,7 @@ def main(args=None) -> None: data = read_fvecs(args.data_file) print(f"Data shape: {data.shape}") - # 2. Load clusters computed in a separate Faiss process. + # 2. Load clusters computed by RaBitQKMeans or QGKMeans. centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") @@ -64,7 +64,7 @@ def main(args=None) -> None: parser.add_argument( "--clusters", required=True, - help="Clustering file produced by faiss_clustering.py for this data and metric", + help="Clustering file produced by kmeans_clustering.py for this data and metric", ) parser.add_argument( "--degree", diff --git a/sample/python/ivf_rabitq_indexing.py b/sample/python/ivf_rabitq_indexing.py index fd7781cc..bf52d81f 100644 --- a/sample/python/ivf_rabitq_indexing.py +++ b/sample/python/ivf_rabitq_indexing.py @@ -23,7 +23,7 @@ def main(args=None) -> None: print(f"\tN: {n}") print(f"\tDIM: {dim}") - # 2. Load clusters computed in a separate Faiss process. + # 2. Load clusters computed by RaBitQKMeans or QGKMeans. centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") @@ -65,7 +65,7 @@ def main(args=None) -> None: parser.add_argument( "--clusters", required=True, - help="Clustering file produced by faiss_clustering.py for this data and metric", + help="Clustering file produced by kmeans_clustering.py for this data and metric", ) parser.add_argument( "--total-bits", diff --git a/sample/python/faiss_clustering.py b/sample/python/kmeans_clustering.py similarity index 62% rename from sample/python/faiss_clustering.py rename to sample/python/kmeans_clustering.py index f14836ac..94257da7 100644 --- a/sample/python/faiss_clustering.py +++ b/sample/python/kmeans_clustering.py @@ -1,17 +1,21 @@ -"""Run Faiss in its own process and save clusters for the RaBitQ examples.""" +"""Train RaBitQKMeans or QGKMeans and save clusters for the index examples.""" import argparse import os from time import time -import faiss import numpy as np +from rabitqlib import FinalAssignmentMode, QGKMeans, RaBitQKMeans from utils import read_fvecs def cluster_data( - data: np.ndarray, num_clusters: int, metric: str, num_threads: int + data: np.ndarray, + num_clusters: int, + metric: str, + num_threads: int, + method: str = "qg", ) -> tuple[np.ndarray, np.ndarray]: if not 1 <= num_clusters <= len(data): raise ValueError("num_clusters must be between 1 and the number of data points") @@ -19,27 +23,32 @@ def cluster_data( raise ValueError("num_threads must be non-negative") if metric not in ("l2", "ip"): raise ValueError("metric must be l2 or ip") + 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 threads = ( hardware_threads if num_threads == 0 else min(num_threads, hardware_threads) ) - faiss.omp_set_num_threads(threads) print(f"Clustering metric: {metric.upper()}, threads: {threads}") - faiss_metric = faiss.METRIC_L2 if metric == "l2" else faiss.METRIC_INNER_PRODUCT - index = faiss.index_factory(data.shape[1], f"IVF{num_clusters},Flat", faiss_metric) - index.verbose = True + kmeans = clustering_type( + data.shape[1], + num_clusters, + spherical=metric == "ip", + num_threads=threads, + final_assignment=FinalAssignmentMode.Exact, + verbose=True, + ) start = time() - index.train(data) - print(f"IVF training time: {time() - start:.2f}s") - centroids = index.quantizer.reconstruct_n(0, index.nlist) - _, labels = index.quantizer.search(data, 1) - return centroids, labels.ravel().astype(np.uint32) + kmeans.train(data) + print(f"{clustering_type.__name__} training time: {time() - start:.2f}s") + return kmeans.centroids, kmeans.assignments def main(args) -> None: data = read_fvecs(args.data_file) centroids, cluster_ids = cluster_data( - data, args.num_clusters, args.metric, args.num_threads + data, args.num_clusters, args.metric, args.num_threads, args.method ) # Open the exact requested path; np.savez otherwise appends .npz automatically. with open(args.clusters_file, "wb") as output: @@ -57,6 +66,12 @@ def main(args) -> None: ) parser.add_argument("--num-clusters", type=int, default=256) parser.add_argument("--metric", choices=["l2", "ip"], default="l2") + parser.add_argument( + "--method", + choices=["rabitq", "qg"], + default="qg", + help="rabitq for few centroids; qg for many (requires k > graph_degree)", + ) parser.add_argument( "--num-threads", type=int, diff --git a/sample/python/utils.py b/sample/python/utils.py index 30587a96..427f42b3 100644 --- a/sample/python/utils.py +++ b/sample/python/utils.py @@ -35,14 +35,14 @@ def compute_recall(ids: np.ndarray, gt: np.ndarray, topk: int) -> float: # ────────────────────────────────────────────── -# Saved clustering results (no Faiss dependency) +# Saved clustering results # ────────────────────────────────────────────── def load_clusters( filename: str, data_shape: tuple[int, int], metric: str ) -> tuple[np.ndarray, np.ndarray]: - """Load results from faiss_clustering.py for the same data in the same order.""" + """Load results from kmeans_clustering.py for the same data in the same order.""" with np.load(filename, allow_pickle=False) as saved: if not {"centroids", "cluster_ids", "metric"}.issubset(saved.files): raise ValueError( diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 696b0b09..9a93b80c 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -44,6 +44,8 @@ QuantizedQuery::QuantizedQuery( g_add_ = metric_type == METRIC_IP ? -dot_product(rotated_query, centroid, padded_dim) : euclidean_sqr(rotated_query, centroid, padded_dim); } +QuantizedQuery::QuantizedQuery(const float* rotated_query, const QuantizedQuery& prepared) + : rotated_query_(rotated_query), k1xsumq_(prepared.k1xsumq_), g_add_(prepared.g_add_) {} const float* QuantizedQuery::rotated_query() const { return rotated_query_; } float QuantizedQuery::k1xsumq() const { return k1xsumq_; } float QuantizedQuery::g_add() const { return g_add_; } @@ -66,52 +68,6 @@ size_t QuantizedGraph::padded_dimension(size_t dim) { return (checked_add(dim, 63) / 64) * 64; } -char* QuantizedGraph::get_row_data(PID data_id) { - return reinterpret_cast(get_vector(data_id)); -} - -const char* QuantizedGraph::get_row_data(PID data_id) const { - return reinterpret_cast(get_vector(data_id)); -} - -float* QuantizedGraph::get_vector(PID data_id) { - return data_.data() + ((row_offset_ / sizeof(float)) * data_id); -} - -const float* QuantizedGraph::get_vector(PID data_id) const { - return data_.data() + ((row_offset_ / sizeof(float)) * data_id); -} - -char* QuantizedGraph::get_quantized_vector(PID data_id) { - return get_row_data(data_id); -} - -const char* QuantizedGraph::get_quantized_vector(PID data_id) const { - return get_row_data(data_id); -} - -char* QuantizedGraph::get_batch_data(PID data_id) { - return get_row_data(data_id) + batch_data_offset_; -} - -const char* QuantizedGraph::get_batch_data(PID data_id) const { - return get_row_data(data_id) + batch_data_offset_; -} - -rabitqlib::detail::PackedArrayView QuantizedGraph::get_neighbors(PID data_id) { - return rabitqlib::detail::PackedArrayView( - get_row_data(data_id) + neighbor_offset_ - ); -} - -rabitqlib::detail::ConstPackedArrayView QuantizedGraph::get_neighbors( - PID data_id -) const { - return rabitqlib::detail::ConstPackedArrayView( - get_row_data(data_id) + neighbor_offset_ - ); -} - size_t QuantizedGraph::num_vertices() const { return this->num_points_; } size_t QuantizedGraph::dimension() const { return this->dim_; } @@ -148,7 +104,8 @@ QuantizedGraph::QuantizedGraph( size_t max_deg, MetricType metric_type, RotatorType rotator_type, - size_t quantization_bits + size_t quantization_bits, + uint32_t seed ) : num_points_(num) , degree_bound_(max_deg) @@ -157,7 +114,8 @@ QuantizedGraph::QuantizedGraph( , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) , metric_type_(metric_type) , rotator_type_(rotator_type) - , quantization_bits_(quantization_bits) { + , quantization_bits_(quantization_bits) + , seed_(seed) { validate_configuration(); initialize(); } @@ -447,75 +405,48 @@ void QuantizedGraph::search( uint32_t* __restrict__ results, float* __restrict__ dists ) { - if (!ready_ || rotator_ == nullptr) { - throw std::logic_error("QuantizedGraph must be built or loaded before search"); - } - if (query == nullptr || results == nullptr || dists == nullptr) { - throw std::invalid_argument("QuantizedGraph search buffers must not be null"); - } - if (k == 0 || k > num_points_) { - throw std::invalid_argument("QuantizedGraph k must be between 1 and num_points"); - } - if (ef_ < k) { - throw std::invalid_argument("QuantizedGraph ef must be at least k"); - } - + validate_search(query, k, results, dists); std::vector rotated_query(padded_dim_); - std::optional quantized_query; - prepare_query(query, rotated_query, quantized_query); - BatchQuery q_obj(rotated_query.data(), padded_dim_, metric_type_); - + std::vector est_dist(degree_bound_); + std::vector lut_float(padded_dim_ * 4); + BatchQuery batch_query; buffer::SearchBuffer search_pool(ef_); - // init search buffer - search_pool.insert(this->entry_point_, std::numeric_limits::max()); - - buffer::SearchBuffer res_pool(k); // result buffer + buffer::SearchBuffer result_pool(k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; if (visited_size != num_points_) { visited.initialize(num_points_, num_points_ / 10); visited_size = num_points_; } - visited.clear(); - auto* vis = &visited; - - std::vector est_dist(degree_bound_); // estimated distances - - while (search_pool.has_next()) { - PID cur_node = search_pool.pop(); - if (vis->get(cur_node)) { - continue; - } - vis->set(cur_node); - - const float vertex_distance = - point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_node); - q_obj.set_g_add(vertex_distance); - - scan_neighbors( - q_obj, cur_node, est_dist.data(), search_pool, *vis, this->degree_bound_ - ); - res_pool.insert(cur_node, vertex_distance); - } - - update_results(res_pool, *vis, query, quantized_query ? &*quantized_query : nullptr); - if (res_pool.size() != k) { - throw std::runtime_error("QuantizedGraph search could not produce k results"); - } - res_pool.copy_results(results, dists); + search_with_scratch( + query, + k, + results, + dists, + rotated_query.data(), + est_dist.data(), + lut_float.data(), + batch_query, + search_pool, + result_pool, + visited + ); } -void QuantizedGraph::prepare_query( - const float* query, - std::vector& rotated_query, - std::optional& quantized_query +void QuantizedGraph::validate_search( + const float* query, uint32_t k, const uint32_t* results, const float* dists ) const { - rotator_->rotate(query, rotated_query.data()); - - if (quantization_bits_ != 0) { - quantized_query.emplace( - rotated_query.data(), centroid_.data(), padded_dim_, metric_type_ - ); + if (!ready_ || rotator_ == nullptr) { + throw std::logic_error("QuantizedGraph must be built or loaded before search"); + } + if (query == nullptr || results == nullptr || dists == nullptr) { + throw std::invalid_argument("QuantizedGraph search buffers must not be null"); + } + if (k == 0 || k > num_points_) { + throw std::invalid_argument("QuantizedGraph k must be between 1 and num_points"); + } + if (ef_ < k) { + throw std::invalid_argument("QuantizedGraph ef must be at least k"); } } @@ -530,7 +461,7 @@ float QuantizedGraph::point_distance( // Scan a data row and store estimated neighbor distances. The caller scores the current // vertex from either its raw vector (vanilla QG) or its 4/8-bit code (qg-quant). -void QuantizedGraph::scan_neighbors( +inline void QuantizedGraph::scan_neighbors( const BatchQuery& q_obj, PID data_id, float* est_dist, @@ -624,7 +555,7 @@ void QuantizedGraph::initialize_layout() { void QuantizedGraph::initialize() { padded_dim_ = padded_dimension(dim_); - rotator_.reset(choose_rotator(dim_, rotator_type_, padded_dim_)); + rotator_.reset(choose_rotator(dim_, rotator_type_, padded_dim_, seed_)); assert(padded_dim_ % 64 == 0); assert(padded_dim_ >= dim_); @@ -691,9 +622,17 @@ void QuantizedGraph::reconstruct_quantized_vector(PID data_id, float* rec // Construction sources come from owned raw rows or from the existing RaBitQ codes. // Reconstructed sources are already rotated; never pass them through prepare_query. const float* QuantizedGraph::prepare_build_query( - PID id, std::vector& rotated, std::optional& prepared + PID id, + std::vector& rotated, + std::optional& prepared, + const BuildCache* cache ) const { if (is_quantized()) { + if (cache != nullptr) { + const float* row = cache->rotated_vectors.data() + id * padded_dim_; + prepared.emplace(row, cache->queries[id]); + return row; + } rotated.resize(padded_dim_); reconstruct_quantized_vector(id, rotated.data()); prepared.emplace(rotated.data(), centroid_.data(), padded_dim_, metric_type_); @@ -709,15 +648,19 @@ void QuantizedGraph::find_candidates( size_t search_ef, std::vector>& results, VisitedSet& vis, - const std::vector& degrees + const std::vector& degrees, + const BuildCache* cache ) const { - std::vector rotated_query(padded_dim_); + std::vector rotated_query(cache == nullptr ? padded_dim_ : 0); std::optional quantized_query; - const float* query = prepare_build_query(cur_id, rotated_query, quantized_query); - if (!is_quantized()) { + const float* query = prepare_build_query(cur_id, rotated_query, quantized_query, cache); + if (cache == nullptr && !is_quantized()) { rotator_->rotate(query, rotated_query.data()); } - BatchQuery q_obj(rotated_query.data(), padded_dim_, metric_type_); + const float* rotated = cache == nullptr + ? rotated_query.data() + : cache->rotated_vectors.data() + cur_id * padded_dim_; + BatchQuery q_obj(rotated, padded_dim_, metric_type_); // insert entry point to initialize search buffer buffer::SearchBuffer tmp_pool(search_ef); @@ -744,7 +687,10 @@ void QuantizedGraph::find_candidates( // based on new neighbor lists to update quantization code and factors void QuantizedGraph::update_qg( - PID cur_id, const std::vector>& new_neighbors + PID cur_id, + const std::vector>& new_neighbors, + const BuildCache* cache, + std::vector* scratch ) { size_t cur_degree = new_neighbors.size(); @@ -752,16 +698,24 @@ void QuantizedGraph::update_qg( return; } // copy neighbors - auto neighbors = get_neighbors(cur_id); + auto neighbor_ptr = get_neighbors(cur_id); for (size_t i = 0; i < cur_degree; ++i) { - neighbors[i] = new_neighbors[i].id; + neighbor_ptr[i] = new_neighbors[i].id; } // rotated data - std::vector rotated_data(cur_degree * padded_dim_); - std::vector rotated_centroid(padded_dim_); + std::vector local_scratch; + std::vector& rotated_data = scratch == nullptr ? local_scratch : *scratch; + rotated_data.resize((cur_degree + 1) * padded_dim_); + float* rotated_centroid = rotated_data.data() + cur_degree * padded_dim_; for (size_t i = 0; i < cur_degree; ++i) { - if (quantization_bits_ == 0) { + if (cache != nullptr) { + std::copy_n( + cache->rotated_vectors.data() + new_neighbors[i].id * padded_dim_, + padded_dim_, + rotated_data.data() + i * padded_dim_ + ); + } else if (quantization_bits_ == 0) { const float* neighbor_vec = get_vector(new_neighbors[i].id); this->rotator_->rotate(neighbor_vec, &rotated_data[i * padded_dim_]); } else { @@ -770,10 +724,16 @@ void QuantizedGraph::update_qg( ); } } - if (quantization_bits_ == 0) { - this->rotator_->rotate(get_vector(cur_id), rotated_centroid.data()); + if (cache != nullptr) { + std::copy_n( + cache->rotated_vectors.data() + cur_id * padded_dim_, + padded_dim_, + rotated_centroid + ); + } else if (quantization_bits_ == 0) { + this->rotator_->rotate(get_vector(cur_id), rotated_centroid); } else { - reconstruct_quantized_vector(cur_id, rotated_centroid.data()); + reconstruct_quantized_vector(cur_id, rotated_centroid); } // quantize batches for current vertex @@ -782,7 +742,7 @@ void QuantizedGraph::update_qg( const size_t batch_size = std::min(cur_degree - i, fastscan::kBatchSize); quant::quantize_qg_batch( rotated_data.data() + (i * padded_dim_), - rotated_centroid.data(), + rotated_centroid, batch_size, padded_dim_, batch_data, @@ -805,4 +765,73 @@ void QuantizedGraph::update_qg( batch_data += QGBatchDataMap::data_bytes(padded_dim_); } } +void QuantizedGraph::search_with_scratch( + const float* __restrict__ query, + uint32_t knn, + uint32_t* __restrict__ results, + float* __restrict__ dists, + float* __restrict__ rotated_query_scratch, + float* __restrict__ est_dist_scratch, + float* __restrict__ lut_float_scratch, + BatchQuery& batch_query, + buffer::SearchBuffer& search_pool, + buffer::SearchBuffer& result_pool, + VisitedSet& visited, + PID hint, + const float* k1xsumq +) { + validate_search(query, knn, results, dists); + if (rotated_query_scratch == nullptr || est_dist_scratch == nullptr || + lut_float_scratch == nullptr) { + throw std::invalid_argument("QuantizedGraph search buffers must not be null"); + } + if (hint != kPidMax && hint >= num_points_) { + throw std::invalid_argument("SymphonyQG search hint is outside the graph"); + } + if (!visited.initialized()) { + throw std::invalid_argument("SymphonyQG visited set must be initialized"); + } + + rotator_->rotate(query, rotated_query_scratch); + std::optional quantized_query; + if (quantization_bits_ != 0) { + quantized_query.emplace( + rotated_query_scratch, centroid_.data(), padded_dim_, metric_type_ + ); + } + const auto* quantized = quantized_query ? &*quantized_query : nullptr; + batch_query.reset( + rotated_query_scratch, padded_dim_, lut_float_scratch, metric_type_, k1xsumq + ); + search_pool.clear(); + result_pool.clear(); + visited.clear(); + + if (hint != kPidMax && hint != entry_point_) { + search_pool.insert(hint, point_distance(query, quantized, hint)); + } + search_pool.insert(entry_point_, std::numeric_limits::max()); + + while (search_pool.has_next()) { + const PID current = search_pool.pop(); + if (visited.get(current)) { + continue; + } + visited.set(current); + + const float vertex_distance = point_distance(query, quantized, current); + batch_query.set_g_add(vertex_distance); + scan_neighbors( + batch_query, current, est_dist_scratch, search_pool, visited, degree_bound_ + ); + result_pool.insert(current, vertex_distance); + } + + update_results(result_pool, visited, query, quantized); + if (result_pool.size() != knn) { + throw std::runtime_error("QuantizedGraph search could not produce k results"); + } + result_pool.copy_results(results, dists); +} + } // namespace rabitqlib::symqg diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index b8cbc193..14630c2f 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -70,6 +70,24 @@ Function resolve_optional_kernel(Function preferred, Function fallback, bool sup #define RABITQ_RESOLVE2(avx2, fallback) fallback #endif +const auto kAccumulateClusterSumsFn = RABITQ_RESOLVE( + accumulate_cluster_sums_avx512, + accumulate_cluster_sums_avx2, + accumulate_cluster_sums_generic +); + +void accumulate_cluster_sums( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +) { + kAccumulateClusterSumsFn(x, point_ids, count, dim, labels, sums, counts); +} + const auto kMatrixProductFn = RABITQ_RESOLVE(matrix_product_avx512, matrix_product_avx2, matrix_product_generic); diff --git a/src/simd/matrix_avx2.cpp b/src/simd/matrix_avx2.cpp index 00410c0f..38633019 100644 --- a/src/simd/matrix_avx2.cpp +++ b/src/simd/matrix_avx2.cpp @@ -6,6 +6,18 @@ #include "matrix_kernels.hpp" namespace rabitqlib::simd { +void accumulate_cluster_sums_avx2( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +) { + accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); +} + void matrix_product_avx2( const float* left, const float* right, diff --git a/src/simd/matrix_avx512.cpp b/src/simd/matrix_avx512.cpp index f35c0287..88d75c5e 100644 --- a/src/simd/matrix_avx512.cpp +++ b/src/simd/matrix_avx512.cpp @@ -6,6 +6,18 @@ #include "matrix_kernels.hpp" namespace rabitqlib::simd { +void accumulate_cluster_sums_avx512( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +) { + accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); +} + void matrix_product_avx512( const float* left, const float* right, diff --git a/src/simd/matrix_generic.cpp b/src/simd/matrix_generic.cpp index 368707e5..0e9e83a2 100644 --- a/src/simd/matrix_generic.cpp +++ b/src/simd/matrix_generic.cpp @@ -6,6 +6,18 @@ #include "matrix_kernels.hpp" namespace rabitqlib::simd { +void accumulate_cluster_sums_generic( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +) { + accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); +} + void matrix_product_generic( const float* left, const float* right, diff --git a/src/simd/matrix_kernels.hpp b/src/simd/matrix_kernels.hpp index d8cac1e0..9c34add2 100644 --- a/src/simd/matrix_kernels.hpp +++ b/src/simd/matrix_kernels.hpp @@ -2,6 +2,7 @@ #include #include +#include // Keep Eigen's ISA-dependent template helpers private to each backend. The // default also lets this header compile on its own for static analysis. @@ -24,6 +25,27 @@ using ConstRowMajorMatrixMap = kernel_eigen::Map; using RowMajorMatrixMap = kernel_eigen::Map; using VectorMap = kernel_eigen::Map>; +inline void accumulate_cluster_sums_impl( + const float* x, + const size_t* point_ids, + size_t count, + size_t dim, + const uint32_t* labels, + double* sums, + size_t* counts +) { + for (size_t index = 0; index < count; ++index) { + const size_t point = point_ids[index]; + const uint32_t cluster = labels[point]; + ++counts[cluster]; + const float* vector = x + point * dim; + double* sum = sums + static_cast(cluster) * dim; + for (size_t coordinate = 0; coordinate < dim; ++coordinate) { + sum[coordinate] += static_cast(vector[coordinate]); + } + } +} + inline void matrix_product_impl( const float* left, const float* right, diff --git a/tests/README.md b/tests/README.md index 4a6d09ff..e5fa9b6b 100644 --- a/tests/README.md +++ b/tests/README.md @@ -113,12 +113,15 @@ python -m pytest tests/python -ra -q Confirm the printed extension path belongs to the intended environment. The suite includes all three indexes, quantization, persistence, and Unicode paths. -The example integration tests additionally need `faiss-cpu`. Wheel CI installs it -as a test dependency and runs Faiss clustering and RaBitQ indexing/querying in -separate subprocesses, with guards against importing both libraries together. -These tests cover L2/IP and the saved clustering file used by the -[Python examples](../sample/python/README.md). They skip the Faiss-dependent cases -when Faiss is absent from a local test environment. +The [Python example tests](python/test_examples.py) cover RaBitQKMeans and +QGKMeans clustering, L2/IP indexing, saved cluster files, and queries without +importing FAISS. The [FAISS comparison tests](python/test_compare_with_faiss.py) +cover both clustering methods and skip FAISS-dependent cases when `faiss-cpu` +is absent; wheel CI includes it as a test dependency. + +The [C++ clustering example tests](python/test_cpp_clustering_examples.py) check +saved vectors, exact labels, and compatibility with the C++ index builders. +Build the corresponding examples first; cases skip when their executables are absent. ## Installed CMake package test diff --git a/tests/dispatch/fastscan_no_simd_test.cpp b/tests/dispatch/fastscan_no_simd_test.cpp index 7fcf6e05..407d0c79 100644 --- a/tests/dispatch/fastscan_no_simd_test.cpp +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -4,6 +4,7 @@ #include #include +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" #include "rabitqlib/utils/cpu_features.hpp" @@ -36,6 +37,33 @@ TEST(FastScanDispatchTest, UnsupportedCpuThrowsWithoutAccessingInputsOrOutput) { EXPECT_EQ(result, original); } +TEST(FastScanDispatchTest, PointCodeAssignmentPropagatesUnsupportedCpuFromWorkers) { + constexpr size_t kDim = 64, kPoints = 65, kClusters = 2; + const std::array data{}; + const std::array centroids{}; + std::array labels; + labels.fill(kPidMax); + std::array distances; + distances.fill(-1.0F); + rabitqkmeans::detail::RaBitQAssigner assigner( + kDim, kClusters, data.data(), kPoints, 2, rabitqkmeans::RaBitQKMeansParameters{} + ); + try { + assigner.assign(centroids.data(), labels.data(), distances.data()); + FAIL() << "Point-code assignment must propagate an unsupported FastScan backend"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "Standard FastScan accumulation requires AVX2/FMA, AVX-512, or ARM NEON; " + "no supported SIMD backend is available" + ); + } + for (size_t point = 0; point < kPoints; ++point) { + EXPECT_EQ(labels[point], kPidMax); + EXPECT_EQ(distances[point], -1.0F); + } +} + TEST(FastScanDispatchTest, InvalidDimensionStillRaisesInvalidArgument) { EXPECT_THROW(accumulate(nullptr, nullptr, nullptr, 0), std::invalid_argument); EXPECT_THROW(accumulate(nullptr, nullptr, nullptr, 15), std::invalid_argument); diff --git a/tests/python/test_compare_with_faiss.py b/tests/python/test_compare_with_faiss.py new file mode 100644 index 00000000..3aed0d83 --- /dev/null +++ b/tests/python/test_compare_with_faiss.py @@ -0,0 +1,190 @@ +"""Checks that repeated comparisons reuse only identical exhaustive evaluations.""" + +import importlib.util +import json +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +SCRIPT = Path(__file__).resolve().parents[2] / "sample/python/compare_with_faiss.py" +# Detect availability without loading FAISS's native runtime in the pytest process. +requires_faiss = pytest.mark.skipif( + importlib.util.find_spec("faiss") is None, + reason="faiss-cpu is required for the optional clustering comparison", +) + + +def run_comparison(monkeypatch, capsys, *, force_recompute): + spec = importlib.util.spec_from_file_location("compare_with_faiss", SCRIPT) + module = importlib.util.module_from_spec(spec) + monkeypatch.syspath_prepend(str(SCRIPT.parent)) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + calls = 0 + original = module.exact_objective + + def counted(*args, **kwargs): + nonlocal calls + calls += 1 + return original(*args, **kwargs) + + module.exact_objective = counted + if force_recompute: + + class NoCacheNumpy: + def __getattr__(self, name): + return getattr(np, name) + + @staticmethod + def array_equal(left, right): + return False + + module.np = NoCacheNumpy() + + monkeypatch.setattr( + sys, + "argv", + [ + str(SCRIPT), + "--n", + "1200", + "--d", + "64", + "--k", + "64", + "--niter", + "2", + "--threads", + "4", + "--seed", + "23", + "--num-seeds", + "1", + "--repeats", + "2", + ], + ) + module.main() + quality = [] + for line in capsys.readouterr().out.splitlines(): + fields = line.split() + if ( + fields + and fields[0] in {"QGKMeans", "RaBitQKMeans", "FAISS"} + and len(fields) == 5 + ): + quality.append((fields[0], *fields[2:])) + elif line.startswith( + ( + "QGKMeans returned-label objective:", + "RaBitQKMeans returned-label objective:", + ) + ): + quality.append(line) + return calls, quality + + +@requires_faiss +def test_repeated_comparison_reuses_identical_exhaustive_results(monkeypatch, capsys): + cached_calls, cached_quality = run_comparison( + monkeypatch, capsys, force_recompute=False + ) + full_calls, full_quality = run_comparison(monkeypatch, capsys, force_recompute=True) + assert cached_calls == 3 + assert full_calls == 6 + assert cached_quality == full_quality + + +def test_default_comparison_trains_once_per_method(monkeypatch): + spec = importlib.util.spec_from_file_location("compare_with_faiss", SCRIPT) + module = importlib.util.module_from_spec(spec) + monkeypatch.syspath_prepend(str(SCRIPT.parent)) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr(sys, "argv", [str(SCRIPT)]) + args = module.parse_args() + assert args.num_seeds == 1 + assert args.repeats == 1 + + +@requires_faiss +@pytest.mark.parametrize("spherical", [False, True]) +def test_native_runtimes_are_isolated(tmp_path, spherical): + # This guard is inherited by spawned workers. It catches mixed imports even + # on machines where the two runtimes happen to coexist without a slowdown. + log = tmp_path / "imports.jsonl" + (tmp_path / "sitecustomize.py").write_text( + """ +import json +import os +import sys +from pathlib import Path + +class NativeImportGuard: + def find_spec(self, fullname, path=None, target=None): + if fullname in {"faiss", "rabitqlib"}: + other = "rabitqlib" if fullname == "faiss" else "faiss" + assert other not in sys.modules, "mixed native runtimes" + row = {"pid": os.getpid(), "library": fullname} + with Path(os.environ["NATIVE_IMPORT_LOG"]).open("a") as stream: + stream.write(json.dumps(row) + "\\n") +sys.meta_path.insert(0, NativeImportGuard()) +""" + ) + original = np.random.default_rng(23).standard_normal((1200, 64), dtype=np.float32) + input_path = tmp_path / "input.npy" + np.save(input_path, original) + env = dict( + os.environ, + PYTHONPATH=str(tmp_path), + NATIVE_IMPORT_LOG=str(log), + OMP_PROC_BIND="close", + OMP_PLACES="threads", + ) + result = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "--input", + str(input_path), + "--k", + "64", + "--niter", + "2", + "--threads", + "4", + "--repeats", + "2", + *(["--spherical"] if spherical else []), + ], + env=env, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + imports = [json.loads(line) for line in log.read_text().splitlines()] + assert {row["library"] for row in imports} == {"faiss", "rabitqlib"} + assert len({row["pid"] for row in imports}) == 2 + for method in ("QGKMeans", "RaBitQKMeans", "FAISS"): + rows = [ + line.split() + for line in result.stdout.splitlines() + if line.split() and line.split()[0] == method and len(line.split()) == 5 + ] + assert len(rows) == 2 + assert all(int(row[2]) == 2 for row in rows) + assert all(np.isfinite(float(row[3])) for row in rows) + if method != "FAISS": + assert f"{method} returned-label objective:" in result.stdout + assert f"{method}/FAISS paired training-time ratio:" in result.stdout + assert ( + f"{method} exhaustive objective difference vs FAISS (%):" + in result.stdout + ) + assert "seeds=1, repeats/seed=2" in result.stdout + np.testing.assert_array_equal(np.load(input_path), original) diff --git a/tests/python/test_cpp_clustering_examples.py b/tests/python/test_cpp_clustering_examples.py new file mode 100644 index 00000000..22dee2c5 --- /dev/null +++ b/tests/python/test_cpp_clustering_examples.py @@ -0,0 +1,121 @@ +"""File-based C++ clustering examples interoperate with C++ index builders.""" + +import os +import subprocess +from pathlib import Path + +import numpy as np +import pytest + +BIN = Path(__file__).resolve().parents[2] / "bin" + + +def run_cpp(name, *args): + executable = BIN / (name + (".exe" if os.name == "nt" else "")) + if not executable.is_file(): + pytest.skip(f"Build the C++ example target {name} first") + return subprocess.run( + [str(executable), *map(str, args)], + env=dict(os.environ, OMP_THREAD_LIMIT="2", OMP_NUM_THREADS="2"), + capture_output=True, + text=True, + timeout=60, + ) + + +def read_vecs(path, rows, cols, dtype): + words = np.fromfile(path, dtype=np.uint32).reshape(rows, cols + 1) + np.testing.assert_array_equal(words[:, 0], cols) + return words[:, 1:].copy().view(dtype) + + +@pytest.mark.parametrize("method,k", [("qgkmeans", 33), ("rabitqkmeans", 16)]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_cpp_clustering_file_outputs(tmp_path, method, k, metric): + data = np.random.default_rng(42).standard_normal((512, 65)).astype(np.float32) + if metric == "ip": + data /= np.linalg.norm(data, axis=1, keepdims=True) + words = np.empty((len(data), data.shape[1] + 1), dtype=np.uint32) + words[:, 0] = data.shape[1] + words[:, 1:] = data.view(np.uint32) + source = tmp_path / "data.fvecs" + words.tofile(source) + centroids_path = tmp_path / "centroids.fvecs" + labels_path = tmp_path / "labels.ivecs" + result = run_cpp( + method, + source, + k, + centroids_path, + labels_path, + *([metric] if metric == "ip" else []), + ) + assert result.returncode == 0, result.stdout + result.stderr + centroids = read_vecs(centroids_path, k, data.shape[1], np.float32) + labels = read_vecs(labels_path, len(data), 1, np.uint32).ravel() + assert np.isfinite(centroids).all() + original, centers = data.astype(np.float64), centroids.astype(np.float64) + if metric == "ip": + np.testing.assert_allclose(np.linalg.norm(centers, axis=1), 1.0, atol=1e-6) + distances = 1 - original @ centers.T + else: + distances = np.sum((original[:, None, :] - centers) ** 2, axis=2) + np.testing.assert_array_equal(labels, distances.argmin(axis=1)) + np.testing.assert_array_equal(np.fromfile(source, dtype=np.uint32), words.ravel()) + + index_path = tmp_path / "index.bin" + if method == "qgkmeans": + result = run_cpp( + "ivf_rabitq_indexing", + source, + centroids_path, + labels_path, + 3, + index_path, + metric, + ) + else: + result = run_cpp( + "hnsw_rabitq_indexing", + source, + centroids_path, + labels_path, + 16, + 32, + 5, + index_path, + metric, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert index_path.stat().st_size > 0 + + +@pytest.mark.parametrize("method", ["qgkmeans", "rabitqkmeans"]) +def test_cpp_clustering_synthetic_demo(method): + result = run_cpp(method) + assert result.returncode == 0, result.stdout + result.stderr + assert "final assignment count=" in result.stdout + + +@pytest.mark.parametrize("method", ["qgkmeans", "rabitqkmeans"]) +@pytest.mark.parametrize("count", ["0", "-1", "16x", "184467440737095516160"]) +def test_cpp_clustering_rejects_invalid_count(tmp_path, method, count): + result = run_cpp(method, "unused.fvecs", count, tmp_path / "c", tmp_path / "ids") + assert result.returncode == 1 + assert "num_clusters must be a positive integer" in result.stderr + assert not (tmp_path / "c").exists() + assert not (tmp_path / "ids").exists() + + +@pytest.mark.parametrize("method", ["qgkmeans", "rabitqkmeans"]) +def test_cpp_clustering_usage_and_errors(tmp_path, method): + result = run_cpp(method, "unused.fvecs") + assert result.returncode == 1 + assert "Usage:" in result.stderr + args = [tmp_path / "missing.fvecs", 33, tmp_path / "c", tmp_path / "ids"] + result = run_cpp(method, *args, "cosine") + assert result.returncode == 1 + assert "metric must be l2 or ip" in result.stderr + result = run_cpp(method, *args) + assert result.returncode == 1 + assert "File does not exist:" in result.stderr diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index abd132df..f8ed6851 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -1,4 +1,4 @@ -"""Example stages keep Faiss and RaBitQ in different Python processes.""" +"""Clustering, indexing and querying examples run without importing Faiss.""" import importlib.util import os @@ -12,7 +12,6 @@ from rabitqlib import HnswIndex, IvfIndex, SymqgIndex EXAMPLES = Path(__file__).resolve().parents[2] / "sample" / "python" -HAS_FAISS = importlib.util.find_spec("faiss") is not None # Fail on an accidental import, even when a runtime combination happens to coexist. BOOTSTRAP = """ @@ -36,6 +35,7 @@ def run_example(name, *args, blocked="faiss"): env = os.environ.copy() # Reproduce Windows redirected output even on hosts using UTF-8 by default. env["PYTHONIOENCODING"] = "cp1252:strict" + env["OMP_NUM_THREADS"] = "2" for key in ( "DYLD_LIBRARY_PATH", "DYLD_INSERT_LIBRARIES", @@ -167,33 +167,36 @@ def check_saved_index_and_query(tmp_path, kind, metric, data): assert "Recall" in output -@pytest.mark.skipif( - not HAS_FAISS, reason="Faiss is required only for clustering examples" -) @pytest.mark.parametrize("metric", ["l2", "ip"]) @pytest.mark.parametrize("kind", ["ivf", "hnsw"]) -def test_separate_clustering_and_indexing(tmp_path, example_data, metric, kind): +@pytest.mark.parametrize("method,count", [("rabitq", 2), ("qg", 33)]) +def test_separate_clustering_and_indexing( + tmp_path, example_data, metric, kind, method, count +): data_path, data = example_data clusters = tmp_path / "clusters.npz" run_example( - "faiss_clustering.py", + "kmeans_clustering.py", data_path, clusters, "--num-clusters", - 2, + count, "--metric", metric, "--num-threads", 2, - blocked="rabitqlib", + *(["--method", method] if method == "rabitq" else []), ) with np.load(clusters, allow_pickle=False) as saved: - assert saved["centroids"].shape == (2, data.shape[1]) + assert saved["centroids"].shape == (count, data.shape[1]) assert saved["centroids"].dtype == np.float32 assert saved["cluster_ids"].shape == (len(data),) assert saved["cluster_ids"].dtype == np.uint32 - assert (saved["cluster_ids"] < 2).all() + assert (saved["cluster_ids"] < count).all() assert saved["metric"].item() == metric + check_exact_cluster_labels( + data, saved["centroids"], saved["cluster_ids"], metric + ) run_example( f"{kind}_rabitq_indexing.py", data_path, @@ -210,6 +213,72 @@ def test_separate_clustering_and_indexing(tmp_path, example_data, metric, kind): check_saved_index_and_query(tmp_path, kind, metric, data) +def check_exact_cluster_labels(data, centroids, labels, metric): + original = data.astype(np.float64) + centers = centroids.astype(np.float64) + if metric == "ip": + np.testing.assert_allclose(np.linalg.norm(centers, axis=1), 1.0, atol=1e-6) + distances = 1.0 - original @ centers.T + else: + delta = original[:, None, :] - centers[None, :, :] + distances = np.einsum("nkd,nkd->nk", delta, delta) + np.testing.assert_array_equal(labels, np.argmin(distances, axis=1)) + + +@pytest.mark.parametrize("argument", [None, "l2", "ip", "innerproduct"]) +def test_legacy_clustering_formats_without_faiss(tmp_path, example_data, argument): + data_path, data = example_data + centroids_path = tmp_path / "centroids.fvecs" + labels_path = tmp_path / "labels.ivecs" + run_example( + EXAMPLES.parents[1] / "python" / "ivf.py", + data_path, + 16, + centroids_path, + labels_path, + *([] if argument is None else [argument]), + "--method", + "rabitq", + ) + centroid_words = np.fromfile(centroids_path, dtype=np.int32).reshape(16, -1) + np.testing.assert_array_equal(centroid_words[:, 0], data.shape[1]) + centroids = centroid_words[:, 1:].copy().view(np.float32) + label_words = np.fromfile(labels_path, dtype=np.int32).reshape(len(data), 2) + np.testing.assert_array_equal(label_words[:, 0], 1) + labels = label_words[:, 1] + assert np.isfinite(centroids).all() + assert ((0 <= labels) & (labels < 16)).all() + metric = "ip" if argument in ("ip", "innerproduct") else "l2" + check_exact_cluster_labels(data, centroids, labels, metric) + + +@pytest.mark.parametrize( + "clusters,dimension,message", + [(0, 64, "num_clusters"), (129, 64, "num_clusters"), (2, 63, "dimension")], +) +def test_clustering_rejects_invalid_sizes(clusters, dimension, message, monkeypatch): + 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) + data = np.zeros((128, dimension), dtype=np.float32) + with pytest.raises(ValueError, match=message): + module.cluster_data(data, clusters, "l2", 1, method="rabitq") + + +def test_clustering_default_requires_graph_compatible_cluster_count(monkeypatch): + 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) + with pytest.raises(ValueError, match="more centroids than the graph degree"): + module.cluster_data(np.zeros((65, 64), dtype=np.float32), 2, "l2", 1) + + @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_symqg_examples_do_not_import_faiss(tmp_path, example_data, metric): data_path, data = example_data diff --git a/tests/python/test_import.py b/tests/python/test_import.py index 209b2c53..cf07c84a 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -1,7 +1,7 @@ """Tests for module import and class construction.""" import pytest -from rabitqlib import HnswIndex, IvfIndex, SymqgIndex +from rabitqlib import HnswIndex, IvfIndex, QGKMeans, SymqgIndex # ── import ──────────────────────────────────────────────────────────────────── @@ -10,6 +10,7 @@ def test_all_classes_importable(): assert HnswIndex is not None assert IvfIndex is not None assert SymqgIndex is not None + assert QGKMeans is not None # ── valid construction ──────────────────────────────────────────────────────── diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py new file mode 100644 index 00000000..ff855770 --- /dev/null +++ b/tests/python/test_qgkmeans.py @@ -0,0 +1,528 @@ +"""Tests for the FAISS-style QGKMeans Python API.""" + +import numpy as np +import pytest +from conftest import DIM +from rabitqlib import FinalAssignmentMode, QGKMeans, QGKMeansParameters, RaBitQKMeans + +K = 33 + + +def qgkmeans(**kwargs): + parameters = { + "niter": 2, + "num_threads": 1, + "min_points_per_centroid": 1, + "ef_build": K, + "ef_search": K, + "graph_build_iterations": 2, + } + parameters.update(kwargs) + return QGKMeans(DIM, K, **parameters) + + +def test_parameters_use_faiss_names_and_defaults(): + cp = QGKMeansParameters() + assert cp.niter == 25 + assert cp.graph_build_iterations == 1 + assert cp.min_points_per_centroid == 39 + assert cp.early_stop_threshold == 0.0 + assert cp.quantization_bits == 0 + + +def test_constructor_accepts_faiss_style_keyword_parameters(): + clustering = qgkmeans(verbose=True, seed=123) + assert clustering.d == DIM + assert clustering.k == K + assert clustering.cp.niter == 2 + assert clustering.cp.verbose + assert clustering.cp.seed == 123 + assert clustering.centroids is None + assert clustering.obj is None + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_unknown_parameter_raises(model_type): + with pytest.raises(AttributeError): + model_type(DIM, K, unknown_parameter=True) + + +def test_train_populates_faiss_style_outputs(base_data): + clustering = qgkmeans() + final_obj = clustering.train(base_data) + + assert clustering.centroids.shape == (K, DIM) + assert clustering.centroids.dtype == np.float32 + assert clustering.assignments.shape == (base_data.shape[0],) + assert clustering.distances.shape == (base_data.shape[0],) + assert clustering.obj.shape == (clustering.cp.niter,) + assert len(clustering.iteration_stats) == clustering.cp.niter + assert final_obj == clustering.obj[-1] + assert clustering.iteration_stats[-1].obj == clustering.obj[-1] + + +def test_parameters_can_be_changed_through_cp(base_data): + clustering = qgkmeans() + clustering.cp.niter = 1 + clustering.train(base_data) + + assert clustering.obj.shape == (1,) + assert len(clustering.iteration_stats) == 1 + + +@pytest.mark.parametrize("bits", [0, 8]) +def test_exact_final_assignment_has_no_nearest_centroid_violations(base_data, bits): + clustering = qgkmeans( + final_assignment=FinalAssignmentMode.Exact, quantization_bits=bits + ) + clustering.train(base_data) + + differences = ( + base_data[:, np.newaxis, :] - clustering.centroids[np.newaxis, :, :] + ).astype(np.float64) + exact_distances = np.einsum("nkd,nkd->nk", differences, differences) + exact_assignments = np.argmin(exact_distances, axis=1) + selected_distances = exact_distances[ + np.arange(base_data.shape[0]), exact_assignments + ] + + np.testing.assert_array_equal(clustering.assignments, exact_assignments) + np.testing.assert_allclose(clustering.distances, selected_distances, rtol=1e-6) + + +@pytest.mark.parametrize("spherical", [False, True]) +@pytest.mark.parametrize("threads", [1, 4]) +def test_exact_final_assignment_handles_matrix_tile_tails(spherical, threads): + x = np.random.default_rng(341).standard_normal((321, 65)).astype(np.float32) + if spherical: + x /= np.linalg.norm(x, axis=1, keepdims=True) + clustering = QGKMeans( + 65, + 257, + niter=1, + num_threads=threads, + min_points_per_centroid=1, + ef_build=40, + ef_search=40, + graph_build_iterations=1, + spherical=spherical, + final_assignment=FinalAssignmentMode.Exact, + ) + clustering.train(x) + + original = x.astype(np.float64) + centroids = clustering.centroids.astype(np.float64) + if spherical: + exact_distances = 1.0 - original @ centroids.T + else: + differences = original[:, None, :] - centroids[None, :, :] + exact_distances = np.einsum("nkd,nkd->nk", differences, differences) + exact_assignments = np.argmin(exact_distances, axis=1) + selected = exact_distances[np.arange(x.shape[0]), exact_assignments] + np.testing.assert_array_equal(clustering.assignments, exact_assignments) + np.testing.assert_allclose(clustering.distances, selected, rtol=1e-6, atol=1e-7) + np.testing.assert_allclose( + clustering.final_obj, clustering.distances.sum(dtype=np.float64), rtol=1e-14 + ) + + +def test_exact_final_assignment_preserves_shifted_l2_accuracy(): + x = np.random.default_rng(85).uniform(-0.125, 0.125, size=(129, 65)) + x = (x + 100000.0).astype(np.float32) + clustering = QGKMeans( + 65, + K, + niter=1, + num_threads=4, + min_points_per_centroid=1, + final_assignment=FinalAssignmentMode.Exact, + ) + clustering.train(x) + + differences = ( + x.astype(np.float64)[:, None, :] + - clustering.centroids.astype(np.float64)[None, :, :] + ) + exact_distances = np.einsum("nkd,nkd->nk", differences, differences) + exact_assignments = np.argmin(exact_distances, axis=1) + selected = exact_distances[np.arange(x.shape[0]), exact_assignments] + np.testing.assert_array_equal(clustering.assignments, exact_assignments) + np.testing.assert_array_equal(clustering.distances, selected.astype(np.float32)) + + +@pytest.mark.parametrize("spherical", [False, True]) +def test_exact_final_assignment_uses_first_duplicate_centroid(spherical): + x = np.zeros((65, 65), dtype=np.float32) + x[:, 0] = 1.0 + clustering = QGKMeans( + 65, + K, + niter=1, + num_threads=4, + min_points_per_centroid=1, + spherical=spherical, + final_assignment=FinalAssignmentMode.Exact, + ) + clustering.train(x) + + np.testing.assert_array_equal(clustering.centroids, x[:K]) + np.testing.assert_array_equal(clustering.assignments, 0) + np.testing.assert_array_equal(clustering.distances, 0.0) + assert clustering.final_obj == 0.0 + + +def test_spherical_mode_returns_normalized_centroids(base_data): + normalized = base_data / np.linalg.norm(base_data, axis=1, keepdims=True) + clustering = qgkmeans( + spherical=True, + final_assignment=FinalAssignmentMode.Exact, + ) + clustering.train(normalized) + + np.testing.assert_allclose( + np.linalg.norm(clustering.centroids, axis=1), + np.ones(K), + rtol=1e-5, + atol=1e-6, + ) + cosine_distances = 1.0 - normalized @ clustering.centroids.T + exact_assignments = np.argmin(cosine_distances, axis=1) + selected_distances = cosine_distances[ + np.arange(normalized.shape[0]), exact_assignments + ] + np.testing.assert_array_equal(clustering.assignments, exact_assignments) + np.testing.assert_allclose(clustering.distances, selected_distances, atol=1e-6) + + +def test_train_rejects_wrong_dimension(): + clustering = qgkmeans() + with pytest.raises(Exception, match="dimension"): + clustering.train(np.zeros((K, DIM + 1), dtype=np.float32)) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("spherical", [False, True]) +def test_graph_returns_exact_selected_distances(bits, spherical): + x = np.random.default_rng(42).standard_normal((100, 65)).astype(np.float32) + if spherical: + x /= np.linalg.norm(x, axis=1, keepdims=True) + clustering = QGKMeans( + 65, + K, + quantization_bits=bits, + spherical=spherical, + niter=2, + num_threads=2, + min_points_per_centroid=1, + ef_build=K, + ef_search=K, + graph_build_iterations=2, + ) + clustering.train(x) + assert len(clustering.iteration_stats) == 2 + assert np.isfinite(clustering.obj).all() + assert np.isfinite(clustering.centroids).all() + selected = clustering.centroids[clustering.assignments].astype(np.float64) + original = x.astype(np.float64) + if spherical: + expected = 1.0 - np.einsum("ij,ij->i", original, selected) + np.testing.assert_allclose( + np.linalg.norm(clustering.centroids, axis=1), 1, atol=1e-6 + ) + else: + delta = original - selected + expected = np.einsum("ij,ij->i", delta, delta) + np.testing.assert_allclose(clustering.distances, expected, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("bits", [1, 7, 9, 32]) +def test_train_rejects_unsupported_quantization_bits(base_data, bits): + clustering = qgkmeans(quantization_bits=bits) + with pytest.raises( + ValueError, match="^QGKMeans quantization_bits must be 0, 4, or 8$" + ): + clustering.train(base_data) + + +@pytest.mark.parametrize("invalid", [np.nan, np.inf, -np.inf]) +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_nonfinite_training_input_preserves_previous_fit( + base_data, invalid, model_type +): + clustering = model_type(DIM, K, niter=2, num_threads=1) + clustering.train(base_data) + expected = clustering.centroids + invalid_data = base_data.copy() + invalid_data[-1, -1] = invalid + with pytest.raises( + ValueError, match=f"^{model_type.__name__} x must contain only finite values$" + ): + clustering.train(invalid_data) + np.testing.assert_array_equal(clustering.centroids, expected) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_failed_first_fit_keeps_results_unavailable(model_type): + clustering = model_type(DIM, K, niter=2, num_threads=1) + with pytest.raises(ValueError, match="^x dimension does not match d$"): + clustering.train(np.zeros((K, DIM + 1), dtype=np.float32)) + assert clustering.centroids is None + assert clustering.obj is None + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_failed_refit_keeps_previous_results(base_data, model_type): + clustering = model_type(DIM, K, niter=2, num_threads=1) + clustering.train(base_data) + expected_centroids = clustering.centroids + expected_assignments = clustering.assignments + with pytest.raises(ValueError, match="^x dimension does not match d$"): + clustering.train(np.zeros((K, DIM + 1), dtype=np.float32)) + np.testing.assert_array_equal(clustering.centroids, expected_centroids) + np.testing.assert_array_equal(clustering.assignments, expected_assignments) + + +def test_zero_input_converges_without_exhausting_iterations(): + clustering = qgkmeans(niter=10) + clustering.train(np.zeros((100, DIM), dtype=np.float32)) + assert len(clustering.iteration_stats) < 10 + np.testing.assert_array_equal(clustering.distances, 0) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +def test_final_objective_measures_returned_assignments(base_data, bits): + clustering = qgkmeans(niter=1, quantization_bits=bits) + assert clustering.final_obj is None + returned = clustering.train(base_data) + assert returned == clustering.obj[-1] + assert clustering.final_obj == clustering.distances.sum(dtype=np.float64) + assert clustering.final_obj < returned + + +def test_native_centroids_before_training_raise(): + from rabitqlib._rabitqlib import _QGKMeans + + clustering = _QGKMeans(DIM, K, QGKMeansParameters()) + with pytest.raises(RuntimeError, match="^QGKMeans has not been trained$"): + _ = clustering.centroids + + +@pytest.mark.parametrize("spherical", [False, True]) +def test_empty_cluster_recovery_preserves_constant_data(spherical): + x = np.zeros((100, DIM), dtype=np.float32) + x[:, 0] = 1 + clustering = qgkmeans(niter=10, spherical=spherical) + clustering.train(x) + assert len(clustering.iteration_stats) < 10 + assert clustering.final_obj == 0 + np.testing.assert_array_equal(clustering.centroids, x[:K]) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("threads", [1, 4]) +def test_single_empty_cluster_preserves_duplicate_centroids(bits, threads): + # n == k includes every input row in initialization. Exactly one duplicate + # centroid is empty, and its only eligible donor contains the identical pair. + x = np.eye(DIM, dtype=np.float32)[:K].copy() + x[1] = x[0] + clustering = qgkmeans( + niter=1, + num_threads=threads, + quantization_bits=bits, + final_assignment=FinalAssignmentMode.Exact, + ) + clustering.train(x) + + assert clustering.iteration_stats[0].nsplit == 1 + expected = sorted(tuple(row) for row in x) + actual = sorted(tuple(row) for row in clustering.centroids) + np.testing.assert_array_equal(actual, expected) + assert clustering.final_obj == 0 + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("spherical", [False, True]) +def test_final_objective_does_not_regress_after_last_update(bits, spherical): + x = np.random.default_rng(739).standard_normal((2048, 65)).astype(np.float32) + if spherical: + x /= np.linalg.norm(x, axis=1, keepdims=True) + clustering = QGKMeans( + 65, + 256, + niter=4, + num_threads=2, + min_points_per_centroid=1, + quantization_bits=bits, + spherical=spherical, + ) + clustering.train(x) + assert clustering.final_obj <= clustering.obj[-1] * (1 + 1e-6) + assert clustering.final_obj == clustering.distances.sum(dtype=np.float64) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("spherical", [False, True]) +def test_pipnn_default_rebuilds_and_returns_exact_selected_distances(bits, spherical): + x = np.random.default_rng(73).standard_normal((100, 65)).astype(np.float32) + if spherical: + x /= np.linalg.norm(x, axis=1, keepdims=True) + model = QGKMeans( + 65, + K, + niter=3, + num_threads=2, + min_points_per_centroid=1, + quantization_bits=bits, + spherical=spherical, + ) + for shift in (0.0, 0.02): + data = x + np.float32(shift) + if spherical: + data /= np.linalg.norm(data, axis=1, keepdims=True) + model.train(data) + chosen = model.centroids[model.assignments].astype(np.float64) + original = data.astype(np.float64) + if spherical: + expected = 1 - np.einsum("ij,ij->i", original, chosen) + else: + delta = original - chosen + expected = np.einsum("ij,ij->i", delta, delta) + np.testing.assert_allclose(model.distances, expected, rtol=1e-5, atol=1e-6) + np.testing.assert_allclose(model.final_obj, expected.sum(), rtol=1e-6) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +def test_finite_out_of_range_input_preserves_previous_fit(base_data, bits): + model = qgkmeans(quantization_bits=bits) + model.train(base_data) + expected = model.centroids + with pytest.raises( + ValueError, match="^QGKMeans x exceeds the safe float32 coordinate range$" + ): + model.train(np.full_like(base_data, 1e19)) + np.testing.assert_array_equal(model.centroids, expected) + + +@pytest.mark.parametrize("verbose", [False, True]) +def test_training_size_warning_is_opt_in(capfd, verbose): + model = qgkmeans(niter=1, min_points_per_centroid=39, verbose=verbose) + model.train(np.zeros((K, DIM), dtype=np.float32)) + captured = capfd.readouterr() + assert ("please provide at least" in captured.err) == verbose + + +@pytest.mark.parametrize("dim", [4096, 4097, 65535, 65536]) +@pytest.mark.parametrize("bits", [0, 4, 8]) +def test_large_dimensions(dim, bits): + x = np.zeros((33, dim), dtype=np.float32) + x[np.arange(33), np.arange(33)] = 1 + x[:, -1] = np.arange(33) / 64 + model = QGKMeans( + dim, + 33, + niter=1, + num_threads=2, + quantization_bits=bits, + ) + model.train(x) + assert model.centroids.shape == (33, dim) + assert model.final_obj == 0 + np.testing.assert_array_equal(x, model.centroids[model.assignments]) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("threads", [1, 4]) +def test_grouped_training_preserves_centroid_means_with_dimension_tail(bits, threads): + x = np.repeat(np.eye(K, 65, dtype=np.float32) * 16, 3, axis=0) + x[:, -1] = np.tile(np.array([-0.125, 0, 0.125], dtype=np.float32), K) + model = QGKMeans( + 65, + K, + niter=8, + num_threads=threads, + min_points_per_centroid=1, + quantization_bits=bits, + ef_build=K, + ef_search=K, + ) + model.train(x) + labels = model.assignments + counts = np.bincount(labels, minlength=K) + assert np.all(counts > 0) + expected = np.array( + [x[labels == cluster].mean(axis=0, dtype=np.float64) for cluster in range(K)], + dtype=np.float32, + ) + np.testing.assert_array_equal(model.centroids, expected) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("spherical", [False, True]) +def test_refit_recomputes_sums_for_mutated_storage_seed_and_count(bits, spherical): + rng = np.random.default_rng(823) + storage = np.zeros((257, 1024), dtype=np.float32) + storage[:, :8] = rng.standard_normal((257, 8)).astype(np.float32) + original_pointer = storage.ctypes.data + + def make_model(seed): + return QGKMeans( + 1024, + 65, + niter=3, + num_threads=1, + seed=seed, + min_points_per_centroid=1, + ef_build=128, + ef_search=1, + quantization_bits=bits, + spherical=spherical, + ) + + model = make_model(42) + model.train(storage[:193]) + for count, seed in [(193, 42), (193, 77), (257, 77), (129, 77)]: + # Keep both the backing storage and view's starting pointer unchanged. + storage *= np.float32(0.375) + storage[:, :8] += rng.standard_normal((257, 8)).astype(np.float32) + assert storage.ctypes.data == original_pointer + data = storage[:count] + assert data.ctypes.data == original_pointer + before = data.copy() + model.cp.seed = seed + model.train(data) + fresh = make_model(seed) + fresh.train(data) + np.testing.assert_array_equal( + model.centroids.view(np.uint32), fresh.centroids.view(np.uint32) + ) + np.testing.assert_array_equal(model.assignments, fresh.assignments) + np.testing.assert_array_equal( + model.distances.view(np.uint32), fresh.distances.view(np.uint32) + ) + np.testing.assert_array_equal( + model.obj.view(np.uint64), fresh.obj.view(np.uint64) + ) + assert model.final_obj == fresh.final_obj + np.testing.assert_array_equal(data, before) + + +def test_approximate_final_assignment_keeps_graph_mode_alias(): + assert FinalAssignmentMode.Approximate == FinalAssignmentMode.SymphonyQG + assert QGKMeansParameters().final_assignment == FinalAssignmentMode.Approximate + + +@pytest.mark.parametrize( + "clusters,degree", + [(1, 32), (2, 32), (16, 32), (31, 32), (32, 32), (33, 64), (64, 64)], +) +@pytest.mark.parametrize( + "mode", [FinalAssignmentMode.Approximate, FinalAssignmentMode.Exact] +) +def test_graph_rejects_clusters_at_or_below_graph_degree(clusters, degree, mode): + model = QGKMeans( + 64, clusters, niter=1, num_threads=1, graph_degree=degree, final_assignment=mode + ) + with pytest.raises( + ValueError, match="^QGKMeans requires more centroids than the graph degree$" + ): + model.train(np.zeros((65, 64), dtype=np.float32)) diff --git a/tests/python/test_rabitqkmeans.py b/tests/python/test_rabitqkmeans.py new file mode 100644 index 00000000..c2f87e1b --- /dev/null +++ b/tests/python/test_rabitqkmeans.py @@ -0,0 +1,161 @@ +"""Tests for flat RaBitQ k-means, independently of the cluster count.""" + +import numpy as np +import pytest +from rabitqlib import ( + FinalAssignmentMode, + QGKMeans, + QGKMeansParameters, + RaBitQKMeans, + RaBitQKMeansIterationStats, + RaBitQKMeansParameters, +) + + +def test_flat_parameters_and_outputs(): + assert RaBitQKMeans is not QGKMeans + assert RaBitQKMeansParameters is not QGKMeansParameters + parameters = RaBitQKMeansParameters() + assert parameters.niter == 25 + assert parameters.min_points_per_centroid == 39 + assert parameters.early_stop_threshold == 0.0 + assert parameters.final_assignment == FinalAssignmentMode.Approximate + model = RaBitQKMeans(65, 33, niter=2, num_threads=1, seed=123) + assert model.cp.seed == 123 + assert model.centroids is None + assert model.obj is None + model.cp.niter = 1 + data = np.random.default_rng(751).standard_normal((97, 65), dtype=np.float32) + objective = model.train(data) + assert model.centroids.shape == (33, 65) + assert model.centroids.dtype == np.float32 + assert model.assignments.shape == (len(data),) + assert model.assignments.dtype == np.uint32 + assert model.distances.shape == (len(data),) + assert model.obj.shape == (1,) + assert objective == model.obj[-1] + assert model.final_obj == model.distances.sum(dtype=np.float64) + assert isinstance(model.iteration_stats[0], RaBitQKMeansIterationStats) + + +@pytest.mark.parametrize( + "parameter", + [ + "graph_degree", + "ef_build", + "ef_search", + "graph_build_iterations", + "quantization_bits", + ], +) +def test_flat_rejects_graph_only_parameters(parameter): + with pytest.raises(AttributeError, match=parameter): + RaBitQKMeans(65, 16, **{parameter: 0}) + + +@pytest.mark.parametrize( + "clusters,message", [(0, "must be positive"), (98, "requires n >= k")] +) +def test_flat_rejects_invalid_cluster_counts(clusters, message): + model = RaBitQKMeans(65, clusters, niter=1, num_threads=1) + with pytest.raises(ValueError, match=message): + model.train(np.zeros((97, 65), dtype=np.float32)) + + +def test_flat_native_centroids_before_training_raise(): + from rabitqlib._rabitqlib import _RaBitQKMeans + + model = _RaBitQKMeans(65, 16, RaBitQKMeansParameters()) + with pytest.raises(RuntimeError, match="^RaBitQKMeans has not been trained$"): + _ = model.centroids + + +@pytest.mark.parametrize("clusters", [1, 2, 16, 31, 32, 33, 65]) +@pytest.mark.parametrize("spherical", [False, True]) +@pytest.mark.parametrize("exact", [False, True]) +def test_flat_train_refit_and_return_exact_selected_distances( + clusters, spherical, exact +): + rng = np.random.default_rng(573) + x = rng.standard_normal((97, 65), dtype=np.float32) + mode = FinalAssignmentMode.Exact if exact else FinalAssignmentMode.Approximate + parameters = dict( + niter=3, + num_threads=2, + spherical=spherical, + min_points_per_centroid=1, + final_assignment=mode, + ) + model = RaBitQKMeans(65, clusters, **parameters) + for count, seed in ((97, 42), (65, 77)): + x += rng.standard_normal(x.shape, dtype=np.float32) * np.float32(0.125) + fit = x[:count] + original = x.copy() + model.cp.seed = seed + model.train(fit) + fresh = RaBitQKMeans(65, clusters, seed=seed, **parameters) + fresh.train(fit) + np.testing.assert_array_equal(x, original) + np.testing.assert_array_equal(model.centroids, fresh.centroids) + np.testing.assert_array_equal(model.assignments, fresh.assignments) + np.testing.assert_array_equal(model.obj, fresh.obj) + centers = model.centroids.astype(np.float64) + data = fit.astype(np.float64) + if spherical: + np.testing.assert_allclose(np.linalg.norm(centers, axis=1), 1.0, atol=1e-6) + distances = 1.0 - data @ centers.T + else: + delta = data[:, None, :] - centers[None, :, :] + distances = np.einsum("nkd,nkd->nk", delta, delta) + selected = distances[np.arange(len(fit)), model.assignments] + np.testing.assert_allclose(model.distances, selected, rtol=1e-6, atol=1e-5) + np.testing.assert_allclose( + model.final_obj, selected.sum(), rtol=1e-6, atol=1e-5 + ) + if exact: + np.testing.assert_array_equal( + model.assignments, np.argmin(distances, axis=1) + ) + + +@pytest.mark.parametrize("spherical", [False, True]) +@pytest.mark.parametrize("clusters", [16, 65]) +def test_flat_duplicates_and_zero(clusters, spherical): + for value in (0.0, 1.0): + x = np.zeros((97, 64), dtype=np.float32) + x[:, 0] = value + model = RaBitQKMeans( + 64, + clusters, + niter=3, + num_threads=2, + spherical=spherical, + min_points_per_centroid=1, + final_assignment=FinalAssignmentMode.Approximate, + ) + model.train(x) + np.testing.assert_array_equal(model.centroids, x[:clusters]) + np.testing.assert_array_equal(model.assignments, 0) + expected = 1.0 - value if spherical else 0.0 + np.testing.assert_array_equal(model.distances, expected) + + +@pytest.mark.parametrize("dimension", [63, 65537]) +def test_flat_keeps_dimension_limits(dimension): + model = RaBitQKMeans(dimension, 2, niter=1, num_threads=1) + with pytest.raises(ValueError, match="dimension"): + model.train(np.zeros((2, dimension), dtype=np.float32)) + + +def test_flat_refit_crosses_previous_cluster_threshold(): + data = np.random.default_rng(514).standard_normal((129, 65), dtype=np.float32) + model = RaBitQKMeans(65, 32, niter=2, num_threads=1) + for clusters in (32, 33, 16): + model.k = clusters + model.train(data) + fresh = RaBitQKMeans(65, clusters, niter=2, num_threads=1) + fresh.train(data) + np.testing.assert_array_equal(model.centroids, fresh.centroids) + np.testing.assert_array_equal(model.assignments, fresh.assignments) + np.testing.assert_array_equal(model.distances, fresh.distances) + np.testing.assert_array_equal(model.obj, fresh.obj) diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp new file mode 100644 index 00000000..36e8902a --- /dev/null +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -0,0 +1,1352 @@ +#include "rabitqlib/clustering/qgkmeans.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::qgkmeans { +namespace { + +TEST(QGKMeansTest, UsesTwentyFiveIterationsByDefault) { + EXPECT_EQ(QGKMeansParameters{}.niter, 25U); +} + +TEST(QGKMeansTest, UsesRawGraphStorageByDefaultAndRejectsUnsupportedBits) { + EXPECT_EQ(QGKMeansParameters{}.quantization_bits, 0U); + for (const uint32_t bits : {1U, 7U, 9U, 32U}) { + QGKMeansParameters parameters; + parameters.quantization_bits = bits; + try { + QGKMeans clustering(65, 33, parameters); + FAIL() << "Accepted unsupported quantization bits " << bits; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "QGKMeans quantization_bits must be 0, 4, or 8"); + } + } +} + +TEST(QGAssignerTest, QuantizedSearchReturnsExactDistancesToOriginalCentroids) { + constexpr size_t kDimension = 65; // Exercise the padded quantized domain. + constexpr size_t kClusters = 33; + constexpr size_t kQueries = 20; + std::mt19937 rng(123); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector centroids(kClusters * kDimension); + std::vector queries(kQueries * kDimension); + std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); + std::generate(queries.begin(), queries.end(), [&] { return distribution(rng); }); + + for (const uint32_t bits : {4U, 8U}) { + SCOPED_TRACE(bits); + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(static_cast(metric)); + QGAssigner assigner(kDimension, kClusters, 32, 33, 33, 2, 1, 42, metric, bits); + std::vector labels(kQueries); + std::vector distances(kQueries); + assigner.assign( + centroids.data(), queries.data(), kQueries, labels.data(), distances.data() + ); + + symqg::QuantizedGraph graph( + kClusters, kDimension, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + symqg::QGBuilder builder(graph, 33, centroids.data(), 1, 42); + builder.build(2); + graph.set_ef(33); + bool saw_estimate_difference = false; + for (size_t i = 0; i < kQueries; ++i) { + PID expected_label = kPidMax; + float estimated_distance = 0; + graph.search( + queries.data() + i * kDimension, 1, &expected_label, &estimated_distance + ); + ASSERT_EQ(labels[i], expected_label); + double exact_distance = metric == METRIC_L2 ? 0.0 : 1.0; + for (size_t dim = 0; dim < kDimension; ++dim) { + const double x = queries[i * kDimension + dim]; + const double c = centroids[labels[i] * kDimension + dim]; + exact_distance += metric == METRIC_L2 ? (x - c) * (x - c) : -x * c; + } + EXPECT_NEAR(distances[i], exact_distance, 1e-5); + saw_estimate_difference |= + std::abs(estimated_distance - exact_distance) > 1e-5; + } + EXPECT_TRUE(saw_estimate_difference); + } + } +} + +TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) { + constexpr size_t kDimension = 65; + constexpr size_t kClusters = 65; + constexpr size_t kQueries = 129; + constexpr uint32_t kEf = 8; + std::mt19937 rng(823); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector initial(kClusters * kDimension), queries(kQueries * kDimension); + std::generate(initial.begin(), initial.end(), [&] { return distribution(rng); }); + std::generate(queries.begin(), queries.end(), [&] { return distribution(rng); }); + for (const uint32_t bits : {0U, 4U, 8U}) { + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE( + ::testing::Message() << bits << "/" << metric << "/" << threads + ); + auto centroids = initial; + QGAssigner assigner( + kDimension, kClusters, 32, 64, kEf, 1, threads, 42, metric, bits + ); + symqg::QuantizedGraph graph( + kClusters, kDimension, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + symqg::QGBuilder builder( + graph, + 64, + centroids.data(), + threads, + symqg::QGInitialization::PiPNN, + 42, + true + ); + graph.set_ef(kEf); + std::vector rotated(graph.padded_dim()), estimates(32), + lut(graph.padded_dim() * 4); + BatchQuery batch; + buffer::SearchBuffer search_pool(kEf), result_pool(1); + VisitedSet visited; + visited.initialize(kClusters, std::min(size_t{kEf} * kEf, kClusters / 10)); + std::vector previous; + for (const size_t count : + {kQueries, kQueries, kQueries - 2, kQueries - 2}) { + for (float& value : centroids) { + value += 0.01F; + } + builder.reset(centroids.data()); + builder.build(); + std::vector actual(count), expected(count); + std::vector actual_distances(count), expected_distances(count); + assigner.assign( + centroids.data(), + queries.data(), + count, + actual.data(), + actual_distances.data() + ); + const auto distance = + metric == METRIC_IP ? dot_product_dis : euclidean_sqr; + for (size_t point = 0; point < count; ++point) { + const PID hint = + previous.size() == count ? previous[point] : kPidMax; + graph.search_with_scratch( + queries.data() + point * kDimension, + 1, + &expected[point], + &expected_distances[point], + rotated.data(), + estimates.data(), + lut.data(), + batch, + search_pool, + result_pool, + visited, + hint + ); + if (bits != 0) { + expected_distances[point] = distance( + queries.data() + point * kDimension, + centroids.data() + expected[point] * kDimension, + kDimension + ); + } + } + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_distances, expected_distances); + previous = expected; + } + } + } + } +} + +TEST(QGKMeansTest, UsesSymphonyQGForFinalAssignmentByDefault) { + EXPECT_EQ(QGKMeansParameters{}.final_assignment, FinalAssignmentMode::SymphonyQG); +} + +TEST(QGKMeansTest, UsesFaissTrainingSizeWarningThresholdByDefault) { + EXPECT_EQ(QGKMeansParameters{}.min_points_per_centroid, 39U); +} + +QGKMeansParameters exact_search_parameters() { + QGKMeansParameters parameters; + parameters.niter = 4; + parameters.graph_degree = 32; + parameters.ef_build = 33; + parameters.ef_search = 33; + parameters.graph_build_iterations = 2; + parameters.num_threads = 1; + parameters.min_points_per_centroid = 1; + return parameters; +} + +QGKMeans train_qgkmeans( + size_t d, size_t k, const QGKMeansParameters& parameters, const float* x, size_t n +) { + QGKMeans clustering(d, k, parameters); + clustering.train(n, x); + return clustering; +} + +std::vector separated_points(size_t num_points, size_t dimension) { + std::vector data(num_points * dimension); + for (size_t point = 0; point < num_points; ++point) { + for (size_t dim = 0; dim < dimension; ++dim) { + data[point * dimension + dim] = + static_cast(point * 10) + static_cast(dim) / 64.0F; + } + } + return data; +} + +TEST(QGAssignerTest, ReusesVisitedStateAcrossQueriesAndAssignerSizes) { + constexpr size_t kDimension = 64; + constexpr size_t kQueries = 256; + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + for (const size_t clusters : {33U, 65U, 33U}) { + SCOPED_TRACE(clusters); + const auto centroids = separated_points(clusters, kDimension); + std::vector queries(kQueries * kDimension); + std::vector labels(kQueries); + std::vector distances(kQueries); + QGAssigner assigner(kDimension, clusters, 32, 128, 128, 2, threads, 42); + for (size_t round = 0; round < 2; ++round) { + for (size_t point = 0; point < kQueries; ++point) { + const size_t cluster = (point + round) % clusters; + std::copy_n( + centroids.data() + cluster * kDimension, + kDimension, + queries.data() + point * kDimension + ); + } + assigner.assign( + centroids.data(), + queries.data(), + kQueries, + labels.data(), + distances.data() + ); + for (size_t point = 0; point < kQueries; ++point) { + EXPECT_EQ(labels[point], (point + round) % clusters); + EXPECT_EQ(distances[point], 0.0F); + } + } + } + } +} + +TEST(QGKMeansTest, WarnsWhenThereAreTooFewPointsPerCentroid) { + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + const std::vector data = separated_points(kClusters, kDimension); + QGKMeansParameters parameters = exact_search_parameters(); + parameters.niter = 1; + parameters.verbose = true; + parameters.min_points_per_centroid = 39; + + testing::internal::CaptureStderr(); + QGKMeans clustering(kDimension, kClusters, parameters); + clustering.train(kClusters, data.data()); + const std::string warning = testing::internal::GetCapturedStderr(); + + EXPECT_NE( + warning.find("please provide at least 39 training points per centroid"), + std::string::npos + ); +} + +TEST(QGKMeansTest, UsesFaissEarlyStopThresholdByDefault) { + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + const std::vector data = separated_points(kClusters, kDimension); + QGKMeansParameters parameters = exact_search_parameters(); + + EXPECT_DOUBLE_EQ(parameters.early_stop_threshold, 0.0); +} + +TEST(QGKMeansTest, StopsAtFaissEarlyStopThreshold) { + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + const std::vector data = separated_points(kClusters * 2, kDimension); + QGKMeansParameters parameters = exact_search_parameters(); + parameters.early_stop_threshold = 1.0; + + const QGKMeans result = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kClusters * 2); + + ASSERT_EQ(result.iteration_stats.size(), 2U); +} + +TEST(QGKMeansTest, ExactFinalAssignmentReturnsNearestCentroidForEveryPoint) { + constexpr size_t kPoints = 100; + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + std::vector data(kPoints * kDimension); + std::mt19937 rng(123); + std::uniform_real_distribution distribution(-10.0F, 10.0F); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + + QGKMeansParameters parameters = exact_search_parameters(); + parameters.niter = 2; + parameters.final_assignment = FinalAssignmentMode::Exact; + const QGKMeans result = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + + EXPECT_DOUBLE_EQ( + result.final_obj, + std::accumulate(result.distances.begin(), result.distances.end(), 0.0) + ); + for (size_t point = 0; point < kPoints; ++point) { + PID expected = 0; + double expected_distance = std::numeric_limits::infinity(); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + double distance = 0.0; + for (size_t dim = 0; dim < kDimension; ++dim) { + const double difference = + static_cast(data[point * kDimension + dim]) - + result.centroids[cluster * kDimension + dim]; + distance += difference * difference; + } + if (distance < expected_distance) { + expected = static_cast(cluster); + expected_distance = distance; + } + } + EXPECT_EQ(result.assignments[point], expected); + EXPECT_FLOAT_EQ(result.distances[point], static_cast(expected_distance)); + } +} + +void expect_exact_assignment_matches_scalar( + const std::vector& data, + const std::vector& centroids, + size_t dimension, + bool spherical +) { + const size_t points = data.size() / dimension; + const size_t clusters = centroids.size() / dimension; + std::vector expected_labels(points); + std::vector expected_distances(points); + for (size_t point = 0; point < points; ++point) { + double best = std::numeric_limits::infinity(); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + double distance = 0.0; + for (size_t dim = 0; dim < dimension; ++dim) { + const double x = data[point * dimension + dim]; + const double c = centroids[cluster * dimension + dim]; + if (spherical) { + distance -= x * c; + } else { + const double difference = x - c; + distance += difference * difference; + } + } + if (spherical) { + distance += 1.0; + } + if (distance < best) { + best = distance; + expected_labels[point] = static_cast(cluster); + } + } + expected_distances[point] = static_cast(best); + } + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + std::vector labels(points, kPidMax); + std::vector distances(points, std::numeric_limits::quiet_NaN()); + detail::exact_assign( + data.data(), + centroids.data(), + points, + dimension, + clusters, + spherical, + threads, + labels.data(), + distances.data() + ); + EXPECT_EQ(labels, expected_labels); + EXPECT_EQ(distances, expected_distances); + } +} + +TEST(QGKMeansTest, ExactAssignmentMatchesScalarAcrossMatrixTileTails) { + std::mt19937 rng(741); + std::uniform_real_distribution distribution(-2.0F, 2.0F); + for (const size_t size : {65U, 129U}) { + SCOPED_TRACE(size); + std::vector data(size * size); + std::vector centroids(257 * size); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + std::generate(centroids.begin(), centroids.end(), [&] { + return distribution(rng); + }); + for (const bool spherical : {false, true}) { + SCOPED_TRACE(spherical); + expect_exact_assignment_matches_scalar(data, centroids, size, spherical); + } + } +} + +TEST(QGKMeansTest, ExactAssignmentMatchesScalarAtMaximumDimension) { + constexpr size_t kDimension = 65536; + std::mt19937 rng(963); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector data(3 * kDimension); + std::vector centroids(3 * kDimension); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); + for (const bool spherical : {false, true}) { + SCOPED_TRACE(spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + } +} + +TEST(QGKMeansTest, ExactAssignmentPreservesCallerOpenMPSettings) { + struct SavedOpenMPSettings { + int threads = omp_get_max_threads(); +#if _OPENMP >= 200805 + int levels = omp_get_max_active_levels(); +#else + int nested = omp_get_nested(); +#endif + ~SavedOpenMPSettings() { + omp_set_num_threads(threads); +#if _OPENMP >= 200805 + omp_set_max_active_levels(levels); +#else + omp_set_nested(nested); +#endif + } + }; + const SavedOpenMPSettings saved; + omp_set_num_threads(3); +#if _OPENMP >= 200805 + omp_set_max_active_levels(2); +#else + omp_set_nested(1); +#endif + constexpr size_t kDimension = 65; + const std::vector centroids(3 * kDimension, 1.0F); + // Three rows reduce the requested team to one; 65 rows use two blocks. + for (const size_t points : {3U, 65U}) { + SCOPED_TRACE(points); + const std::vector data(points * kDimension, 0.0F); + std::vector labels(points); + std::vector distances(points); + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + detail::exact_assign( + data.data(), + centroids.data(), + points, + kDimension, + 3, + false, + threads, + labels.data(), + distances.data() + ); + EXPECT_EQ(omp_get_max_threads(), 3); +#if _OPENMP >= 200805 + EXPECT_EQ(omp_get_max_active_levels(), 2); +#else + EXPECT_NE(omp_get_nested(), 0); +#endif + } + } +} + +TEST(QGKMeansTest, ExactAssignmentPreservesLowestIndexForTiesAcrossTiles) { + constexpr size_t kDimension = 65; + std::vector data(65 * kDimension, 0.0F); + std::vector centroids(257 * kDimension, 4.0F); + // Centroids 3 and 256 tie, while preceding centroids are worse for every row. + std::fill_n(centroids.data() + 3 * kDimension, kDimension, 0.0F); + std::fill_n(centroids.data() + 256 * kDimension, kDimension, 0.0F); + centroids[3 * kDimension] = -1.0F; + centroids[256 * kDimension] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + std::fill(data.begin(), data.end(), 1.0F); + std::fill(centroids.begin(), centroids.end(), 0.0F); + centroids[3 * kDimension] = 1.0F; + centroids[256 * kDimension] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDimension, true); +} + +TEST(QGKMeansTest, ExactAssignmentRefinesLargeOffsetL2Distances) { + constexpr size_t kDimension = 129; + std::mt19937 rng(51); + std::uniform_real_distribution distribution(-0.125F, 0.125F); + std::vector data(65 * kDimension); + std::vector centroids(257 * kDimension); + auto shifted = [&] { return 100000.0F + distribution(rng); }; + std::generate(data.begin(), data.end(), shifted); + std::generate(centroids.begin(), centroids.end(), shifted); + // Include a zero-distance winner at the tail of the centroid matrix. + std::copy_n(data.data(), kDimension, centroids.data() + 256 * kDimension); + expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); +} + +TEST(QGKMeansTest, ExactAssignmentResolvesFloatNearTiesAndTinyDistances) { + constexpr size_t kDimension = 65; + std::vector data(65 * kDimension, 0.0F); + std::vector centroids(257 * kDimension, 2.0F); + std::fill_n(centroids.data(), kDimension, 0.0F); + std::fill_n(centroids.data() + 256 * kDimension, kDimension, 0.0F); + // Both squared distances round to 1 in float, but centroid 256 is nearer. + centroids[0] = 1.0F; + centroids[1] = 0.0001F; + centroids[256 * kDimension] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + for (float& value : centroids) { + value *= 1e-22F; + } + expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + std::fill(centroids.begin(), centroids.end(), 0.0F); + for (const bool spherical : {false, true}) { + SCOPED_TRACE(spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + } +} + +TEST(QGKMeansTest, ExactAssignmentRefinesInnerProductCancellation) { + constexpr size_t kDimension = 129; + std::vector data(65 * kDimension, 100000.0F); + std::vector centroids(257 * kDimension); + for (size_t point = 0; point < 65; ++point) { + data[point * kDimension + kDimension - 1] = 1.0F; + } + for (size_t cluster = 0; cluster < 257; ++cluster) { + for (size_t dim = 0; dim + 1 < kDimension; ++dim) { + centroids[cluster * kDimension + dim] = dim % 2 == 0 ? 100000.0F : -100000.0F; + } + centroids[cluster * kDimension + kDimension - 1] = + 1.0F + static_cast(cluster) / 4096.0F; + } + expect_exact_assignment_matches_scalar(data, centroids, kDimension, true); +} + +TEST(QGKMeansTest, SphericalModeNormalizesCentroidsAndUsesInnerProduct) { + constexpr size_t kPoints = 100; + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + std::vector data(kPoints * kDimension); + std::mt19937 rng(456); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + for (size_t point = 0; point < kPoints; ++point) { + double squared_norm = 0.0; + for (size_t dim = 0; dim < kDimension; ++dim) { + squared_norm += static_cast(data[point * kDimension + dim]) * + data[point * kDimension + dim]; + } + const float scale = 1.0F / std::sqrt(static_cast(squared_norm)); + for (size_t dim = 0; dim < kDimension; ++dim) { + data[point * kDimension + dim] *= scale; + } + } + + QGKMeansParameters parameters = exact_search_parameters(); + parameters.niter = 2; + parameters.spherical = true; + parameters.final_assignment = FinalAssignmentMode::Exact; + const QGKMeans result = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + double squared_norm = 0.0; + for (size_t dim = 0; dim < kDimension; ++dim) { + squared_norm += + static_cast(result.centroids[cluster * kDimension + dim]) * + result.centroids[cluster * kDimension + dim]; + } + EXPECT_NEAR(squared_norm, 1.0, 1e-5); + } + + EXPECT_DOUBLE_EQ( + result.final_obj, + std::accumulate(result.distances.begin(), result.distances.end(), 0.0) + ); + for (size_t point = 0; point < kPoints; ++point) { + PID expected = 0; + double expected_distance = std::numeric_limits::infinity(); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + double distance = 1.0; + for (size_t dim = 0; dim < kDimension; ++dim) { + distance -= static_cast(data[point * kDimension + dim]) * + result.centroids[cluster * kDimension + dim]; + } + if (distance < expected_distance) { + expected = static_cast(cluster); + expected_distance = distance; + } + } + EXPECT_EQ(result.assignments[point], expected); + EXPECT_NEAR(result.distances[point], expected_distance, 1e-6); + } +} + +TEST(QGKMeansTest, ObjectiveTerminationMatchesFaissAbsoluteRelativeChange) { + EXPECT_TRUE(detail::should_terminate_by_objective(100.0, 101.0, 0.01)); + EXPECT_TRUE(detail::should_terminate_by_objective(100.0, 99.0, 0.01)); + EXPECT_FALSE(detail::should_terminate_by_objective(100.0, 98.0, 0.01)); + EXPECT_TRUE(detail::should_terminate_by_objective(0.0, 0.0, 0.0)); + EXPECT_FALSE(detail::should_terminate_by_objective(0.0, 1.0, 1.0)); +} + +TEST(QGKMeansTest, TrainsAndReturnsAssignmentsForFinalCentroids) { + constexpr size_t kClusters = 33; + constexpr size_t kDimension = 64; + std::vector data(kClusters * kDimension); + for (size_t point = 0; point < kClusters; ++point) { + for (size_t dim = 0; dim < kDimension; ++dim) { + data[point * kDimension + dim] = + static_cast(point * 10) + static_cast(dim) / 64.0F; + } + } + + QGKMeansParameters parameters = exact_search_parameters(); + parameters.niter = 1; + parameters.num_threads = 2; + + const QGKMeans result = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kClusters); + + ASSERT_EQ(result.centroids.size(), kClusters * kDimension); + ASSERT_EQ(result.assignments.size(), kClusters); + ASSERT_EQ(result.distances.size(), kClusters); + ASSERT_EQ(result.iteration_stats.size(), 1U); + for (size_t point = 0; point < kClusters; ++point) { + EXPECT_LT(result.assignments[point], kClusters); + EXPECT_TRUE(std::isfinite(result.distances[point])); + EXPECT_FLOAT_EQ(result.distances[point], 0.0F); + const size_t cluster = result.assignments[point]; + for (size_t dim = 0; dim < kDimension; ++dim) { + EXPECT_FLOAT_EQ( + result.centroids[cluster * kDimension + dim], data[point * kDimension + dim] + ); + } + } +} + +TEST(QGKMeansTest, RetrainsOnOwnedCentroidsLikeAnIndependentCopy) { + constexpr size_t kDimension = 65, kInitialClusters = 35, kClusters = 33; + const auto data = separated_points(kInitialClusters, kDimension); + for (const uint32_t bits : {0U, 4U, 8U}) { + for (const auto mode : + {FinalAssignmentMode::SymphonyQG, FinalAssignmentMode::Exact}) { + for (const size_t offset : {0U, 1U}) { + SCOPED_TRACE( + ::testing::Message() + << bits << "/" << static_cast(mode) << "/" << offset + ); + auto parameters = exact_search_parameters(); + parameters.niter = 2; + parameters.quantization_bits = bits; + parameters.final_assignment = mode; + QGKMeans clustering(kDimension, kInitialClusters, parameters); + clustering.train(kInitialClusters, data.data()); + const auto original = clustering.centroids; + const float* input = original.data() + offset * kDimension; + const size_t count = kInitialClusters - offset; + + QGKMeans reference(kDimension, kClusters, parameters); + reference.train(count, input); + // Shrink the output while training on its former contents, including + // a subrange that does not start at centroids.data(). + clustering.k = kClusters; + clustering.train(count, clustering.centroids.data() + offset * kDimension); + + EXPECT_EQ(clustering.centroids, reference.centroids); + EXPECT_EQ(clustering.assignments, reference.assignments); + EXPECT_EQ(clustering.distances, reference.distances); + EXPECT_DOUBLE_EQ(clustering.final_obj, reference.final_obj); + double objective = 0; + for (size_t point = 0; point < count; ++point) { + ASSERT_LT(clustering.assignments[point], kClusters); + double distance = 0; + for (size_t dim = 0; dim < kDimension; ++dim) { + const double difference = + static_cast(input[point * kDimension + dim]) - + clustering.centroids + [clustering.assignments[point] * kDimension + dim]; + distance += difference * difference; + } + EXPECT_NEAR( + clustering.distances[point], + distance, + 1e-5 * std::max(1.0, distance) + ); + objective += distance; + } + EXPECT_NEAR( + clustering.final_obj, objective, 1e-5 * std::max(1.0, objective) + ); + } + } + } +} + +TEST(QGKMeansTest, EmptyClusterRecoveryDoesNotPerturbConstantData) { + constexpr size_t kPoints = 100, kDimension = 64; + auto parameters = exact_search_parameters(); + parameters.niter = 10; + for (bool spherical : {false, true}) { + SCOPED_TRACE(spherical); + parameters.spherical = spherical; + std::vector data(kPoints * kDimension, 0.0F); + for (size_t point = 0; point < kPoints; ++point) { + data[point * kDimension] = 1.0F; + } + const auto clustering = + train_qgkmeans(kDimension, 33, parameters, data.data(), kPoints); + EXPECT_LT(clustering.iteration_stats.size(), parameters.niter); + EXPECT_EQ(clustering.final_obj, 0.0); + for (size_t cluster = 0; cluster < 33; ++cluster) { + for (size_t dim = 0; dim < kDimension; ++dim) { + EXPECT_EQ(clustering.centroids[cluster * kDimension + dim], data[dim]); + } + } + } +} + +TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { + constexpr size_t kDimension = 64, kClusters = 33, kPoints = kClusters + 8; + const auto normalize = [=](float* row) { + double squared_norm = 0; + for (size_t dim = 0; dim < kDimension; ++dim) { + squared_norm += static_cast(row[dim]) * row[dim]; + } + if (squared_norm != 0) { + const double scale = 1.0 / std::sqrt(squared_norm); + for (size_t dim = 0; dim < kDimension; ++dim) { + row[dim] = static_cast(row[dim] * scale); + } + } + }; + std::vector indices(kPoints); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::mt19937 initial_rng(42); + std::shuffle(indices.begin(), indices.end(), initial_rng); + + for (bool spherical : {false, true}) { + for (size_t empty_clusters : {1U, 3U}) { + SCOPED_TRACE(::testing::Message() << spherical << "/" << empty_clusters); + // Force duplicate initial seeds while all other initial centroids are + // distinct basis vectors. Only the duplicate seeds have extra points. + std::vector data(kPoints * kDimension); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + const size_t axis = cluster <= empty_clusters ? 0 : cluster; + data[indices[cluster] * kDimension + axis] = 1.0F; + } + for (size_t extra = 0; extra < 8; ++extra) { + float* row = data.data() + indices[kClusters + extra] * kDimension; + row[0] = 1.0F; + row[kClusters] = static_cast(4 - extra / 2) / 8.0F * + (extra % 2 == 0 ? 1.0F : -1.0F); + if (spherical) { + normalize(row); + } + } + std::vector initial(kClusters * kDimension); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + std::copy_n( + data.data() + indices[cluster] * kDimension, + kDimension, + initial.data() + cluster * kDimension + ); + } + QGAssigner reference( + kDimension, + kClusters, + 32, + 33, + 33, + 2, + 1, + 42, + spherical ? METRIC_IP : METRIC_L2 + ); + std::vector labels(kPoints); + std::vector distances(kPoints); + reference.assign( + initial.data(), data.data(), kPoints, labels.data(), distances.data() + ); + + // Derive occupied means from the public initial assignment. Empty + // centroids receive selected rows without changing these donor means. + std::vector counts(kClusters); + std::vector sums(kClusters * kDimension); + for (size_t point = 0; point < kPoints; ++point) { + ++counts[labels[point]]; + for (size_t dim = 0; dim < kDimension; ++dim) { + sums[labels[point] * kDimension + dim] += + data[point * kDimension + dim]; + } + } + ASSERT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), empty_clusters); + std::vector means(kClusters * kDimension); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + if (counts[cluster] != 0) { + for (size_t dim = 0; dim < kDimension; ++dim) { + means[cluster * kDimension + dim] = static_cast( + sums[cluster * kDimension + dim] / + static_cast(counts[cluster]) + ); + } + if (spherical) { + normalize(means.data() + cluster * kDimension); + } + } + } + const auto distance = spherical ? dot_product_dis : euclidean_sqr; + for (size_t point = 0; point < kPoints; ++point) { + distances[point] = distance( + data.data() + point * kDimension, + means.data() + labels[point] * kDimension, + kDimension + ); + } + std::vector candidates(kPoints); + std::iota(candidates.begin(), candidates.end(), size_t{0}); + std::sort(candidates.begin(), candidates.end(), [&](size_t left, size_t right) { + return distances[left] == distances[right] + ? left < right + : distances[left] > distances[right]; + }); + // The symmetric farthest pair also exercises deterministic ID tie-breaking. + ASSERT_EQ(distances[candidates[0]], distances[candidates[1]]); + size_t next_candidate = 0; + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + if (counts[cluster] != 0) { + continue; + } + while (counts[labels[candidates[next_candidate]]] <= 1) { + ++next_candidate; + ASSERT_LT(next_candidate, candidates.size()); + } + const size_t point = candidates[next_candidate++]; + std::copy_n( + data.data() + point * kDimension, + kDimension, + means.data() + cluster * kDimension + ); + if (spherical) { + normalize(means.data() + cluster * kDimension); + } + --counts[labels[point]]; + counts[cluster] = 1; + } + ASSERT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), 0); + for (uint32_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + auto parameters = exact_search_parameters(); + parameters.niter = 1; + parameters.spherical = spherical; + parameters.num_threads = threads; + parameters.final_assignment = FinalAssignmentMode::Exact; + const auto clustering = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + ASSERT_EQ(clustering.iteration_stats.front().nsplit, empty_clusters); + // Parallel graph construction may choose a different duplicate + // seed label, but occupied means and selected seeds must match. + std::vector> expected_rows, actual_rows; + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + const size_t offset = cluster * kDimension; + expected_rows.emplace_back( + means.data() + offset, means.data() + offset + kDimension + ); + actual_rows.emplace_back( + clustering.centroids.data() + offset, + clustering.centroids.data() + offset + kDimension + ); + } + std::sort(expected_rows.begin(), expected_rows.end()); + std::sort(actual_rows.begin(), actual_rows.end()); + EXPECT_EQ(actual_rows, expected_rows); + std::fill(counts.begin(), counts.end(), size_t{0}); + for (PID label : clustering.assignments) { + ++counts[label]; + } + EXPECT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), 0); + } + } + } +} + +TEST(QGKMeansTest, RejectsNonfiniteTrainingDataWithoutReplacingResults) { + constexpr size_t kDimension = 64; + const auto data = separated_points(66, kDimension); + auto parameters = exact_search_parameters(); + parameters.niter = 1; + QGKMeans clustering(kDimension, 33, parameters); + clustering.train(66, data.data()); + const auto expected_centroids = clustering.centroids; + const auto expected_assignments = clustering.assignments; + for (float invalid : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + auto invalid_data = data; + invalid_data.back() = invalid; + try { + clustering.train(66, invalid_data.data()); + FAIL() << "Accepted nonfinite input"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "QGKMeans x must contain only finite values"); + } + EXPECT_EQ(clustering.centroids, expected_centroids); + EXPECT_EQ(clustering.assignments, expected_assignments); + } +} + +TEST(QGKMeansTest, AccumulatesShiftedClusterMeansAccurately) { + constexpr size_t kDimension = 64, kClusters = 33, kPoints = 10000; + std::mt19937 rng(123); + std::uniform_real_distribution distribution(-1, 1); + std::vector data(kPoints * kDimension); + for (auto& value : data) { + value = 10000.0F + distribution(rng); + } + // Replay the public seeded initialization to obtain the first assignment; + // compute its expected means independently in double precision. + std::vector indices(kPoints); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::mt19937 initial_rng(42); + std::shuffle(indices.begin(), indices.end(), initial_rng); + std::vector initial_centroids(kClusters * kDimension); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + std::copy_n( + data.data() + indices[cluster] * kDimension, + kDimension, + initial_centroids.data() + cluster * kDimension + ); + } + QGAssigner assigner(kDimension, kClusters, 32, 33, 33, 2, 1, 42); + std::vector labels(kPoints); + std::vector distances(kPoints); + assigner.assign( + initial_centroids.data(), data.data(), kPoints, labels.data(), distances.data() + ); + std::vector sums(kClusters * kDimension); + std::vector counts(kClusters); + for (size_t point = 0; point < kPoints; ++point) { + ++counts[labels[point]]; + for (size_t dim = 0; dim < kDimension; ++dim) { + sums[labels[point] * kDimension + dim] += data[point * kDimension + dim]; + } + } + auto parameters = exact_search_parameters(); + parameters.niter = 1; + const auto clustering = + train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + EXPECT_EQ( + clustering.iteration_stats.front().nsplit, + std::count(counts.begin(), counts.end(), size_t{0}) + ); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + if (counts[cluster] == 0) { + continue; + } + for (size_t dim = 0; dim < kDimension; ++dim) { + EXPECT_FLOAT_EQ( + clustering.centroids[cluster * kDimension + dim], + static_cast(sums[cluster * kDimension + dim] / counts[cluster]) + ); + } + } +} + +TEST(QGKMeansTest, FinalAssignmentsRetainACloserPreviousCluster) { + constexpr size_t kDimension = 65, kClusters = 256, kPoints = 2048; + const auto normalize = [=](float* row) { + double norm = 0; + for (size_t dim = 0; dim < kDimension; ++dim) { + norm += static_cast(row[dim]) * row[dim]; + } + const double scale = 1.0 / std::sqrt(norm); + for (size_t dim = 0; dim < kDimension; ++dim) { + row[dim] = static_cast(row[dim] * scale); + } + }; + for (bool spherical : {false, true}) { + std::mt19937 rng(123); + std::normal_distribution normal; + std::vector data(kPoints * kDimension); + std::generate(data.begin(), data.end(), [&] { return normal(rng); }); + if (spherical) { + for (size_t point = 0; point < kPoints; ++point) { + normalize(data.data() + point * kDimension); + } + } + std::vector indices(kPoints); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::mt19937 init_rng(42); + std::shuffle(indices.begin(), indices.end(), init_rng); + std::vector initial(kClusters * kDimension); + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + float* row = initial.data() + cluster * kDimension; + std::copy_n(data.data() + indices[cluster] * kDimension, kDimension, row); + if (spherical) { + normalize(row); + } + } + const auto metric = spherical ? METRIC_IP : METRIC_L2; + for (uint32_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(::testing::Message() << spherical << "/" << bits); + QGAssigner reference( + kDimension, + kClusters, + 32, + 240, + 16, + QGKMeansParameters{}.graph_build_iterations, + 1, + 42, + metric, + bits + ); + std::vector previous(kPoints); + std::vector distances(kPoints); + reference.assign( + initial.data(), data.data(), kPoints, previous.data(), distances.data() + ); + QGKMeansParameters parameters; + parameters.niter = 1; + parameters.num_threads = 1; + parameters.min_points_per_centroid = 1; + parameters.spherical = spherical; + parameters.quantization_bits = bits; + QGKMeans clustering(kDimension, kClusters, parameters); + clustering.train(kPoints, data.data()); + const auto distance = spherical ? dot_product_dis : euclidean_sqr; + for (size_t point = 0; point < kPoints; ++point) { + const float previous_distance = distance( + data.data() + point * kDimension, + clustering.centroids.data() + previous[point] * kDimension, + kDimension + ); + EXPECT_LE(clustering.distances[point], previous_distance); + } + } + } +} + +TEST(QGKMeansTest, StopsWhenZeroDataCannotImprove) { + auto parameters = exact_search_parameters(); + parameters.niter = 10; + const std::vector data(100 * 64, 0.0F); + const auto clustering = train_qgkmeans(64, 33, parameters, data.data(), 100); + EXPECT_LT(clustering.iteration_stats.size(), parameters.niter); + EXPECT_EQ( + std::accumulate(clustering.distances.begin(), clustering.distances.end(), 0.0), 0.0 + ); +} + +TEST(QGKMeansTest, TrainsWithMovingCentroidsAndAnApproximateGraph) { + constexpr size_t kDim = 65, kCount = 2048, kClusters = 256; + std::mt19937 rng(739); + std::normal_distribution normal; + std::vector data(kCount * kDim); + std::generate(data.begin(), data.end(), [&] { return normal(rng); }); + for (bool spherical : {false, true}) { + if (spherical) { + for (size_t point = 0; point < kCount; ++point) { + double norm = 0; + for (size_t dim = 0; dim < kDim; ++dim) { + const double value = data[point * kDim + dim]; + norm += value * value; + } + for (size_t dim = 0; dim < kDim; ++dim) { + data[point * kDim + dim] /= static_cast(std::sqrt(norm)); + } + } + } + for (uint32_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(::testing::Message() << spherical << "/" << bits); + QGKMeansParameters parameters; + parameters.niter = 5; + parameters.num_threads = 2; + parameters.min_points_per_centroid = 1; + parameters.spherical = spherical; + parameters.quantization_bits = bits; + QGKMeans clustering(kDim, kClusters, parameters); + clustering.train(kCount, data.data()); + double previous = std::numeric_limits::infinity(); + for (const auto& stats : clustering.iteration_stats) { + EXPECT_TRUE(std::isfinite(stats.obj)); + if (bits == 0) { + EXPECT_LE(stats.obj, previous + 1e-6 * std::abs(previous)); + } + previous = stats.obj; + } + EXPECT_LE(clustering.final_obj, previous + 1e-6 * std::abs(previous)); + } + } +} + +TEST(QGAssignerTest, RejectsTooFewClustersForGraphDegree) { + EXPECT_THROW(QGAssigner(64, 32, 32, 64, 16, 1, 1), std::invalid_argument); +} + +TEST(QGKMeansTest, RejectsTooFewClustersForGraphDegreeInBothFinalModes) { + constexpr size_t kDim = 64, kPoints = 65; + const std::vector data(kPoints * kDim, 0.0F); + for (const auto mode : {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + for (const uint32_t degree : {32U, 64U}) { + for (const size_t clusters : {1U, 2U, 16U, 31U, 32U, 33U, 64U}) { + if (clusters > degree) { + continue; + } + QGKMeansParameters parameters; + parameters.niter = 1; + parameters.num_threads = 1; + parameters.graph_degree = degree; + parameters.final_assignment = mode; + try { + QGKMeans model(kDim, clusters, parameters); + model.train(kPoints, data.data()); + FAIL() << "Accepted k=" << clusters << " for degree=" << degree; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ( + error.what(), + "QGKMeans requires more centroids than the graph degree" + ); + } + } + } + } +} + +TEST(QGKMeansTest, RejectsInvalidNativeParametersBeforeAllocation) { + QGKMeansParameters parameters; + parameters.num_threads = std::numeric_limits::max(); + EXPECT_THROW((QGKMeans(64, 33, parameters)), std::invalid_argument); + EXPECT_THROW( + (QGAssigner(64, 33, 32, 64, 16, 2, parameters.num_threads)), std::invalid_argument + ); + parameters.num_threads = 1; + parameters.final_assignment = static_cast(255); + EXPECT_THROW((QGKMeans(64, 33, parameters)), std::invalid_argument); + parameters.final_assignment = FinalAssignmentMode::Exact; + QGKMeans clustering(64, 33, parameters); + const float value = 0; + EXPECT_THROW( + clustering.train(std::numeric_limits::max() / 64 + 1, &value), + std::invalid_argument + ); +} + +} // namespace +} // namespace rabitqlib::qgkmeans + +namespace rabitqlib::qgkmeans { +TEST(QGAssignerTest, PropagatesSearchFailureFromWorkers) { + const std::vector centroids(33 * 64, 0.0F); + const std::vector data(128 * 64, 1e19F); + std::vector labels(128, kPidMax); + std::vector distances(128, -1.0F); + QGAssigner assigner(64, 33, 32, 33, 33, 1, 2); + EXPECT_THROW( + assigner.assign( + centroids.data(), data.data(), 128, labels.data(), distances.data() + ), + std::runtime_error + ); + EXPECT_EQ(labels.front(), kPidMax); + EXPECT_EQ(distances.front(), -1.0F); +} + +TEST(QGKMeansTest, RejectsFiniteOutOfRangeInputBeforeConstruction) { + for (uint32_t bits : {0U, 4U, 8U}) { + QGKMeansParameters parameters; + parameters.niter = 1; + parameters.num_threads = 2; + parameters.quantization_bits = bits; + QGKMeans model(64, 33, parameters); + const std::vector data(66 * 64, 1e19F); + EXPECT_THROW(model.train(66, data.data()), std::invalid_argument); + EXPECT_TRUE(model.centroids.empty()); + } +} +} // namespace rabitqlib::qgkmeans + +namespace rabitqlib::qgkmeans { +TEST(QGKMeansTest, SupportsLargerPaddedDimensions) { + constexpr size_t k = 33; + for (const size_t d : {4097U, 65535U, 65536U}) { + SCOPED_TRACE(d); + std::vector data(k * d, 0.0F); + for (size_t i = 0; i < k; ++i) { + data[i * d + i] = 1.0F; + data[i * d + d - 1] = static_cast(i) / 64.0F; + } + for (uint32_t bits : {0U, 4U, 8U}) { + QGKMeansParameters parameters; + parameters.niter = 1; + parameters.num_threads = 2; + parameters.quantization_bits = bits; + QGKMeans model(d, k, parameters); + model.train(k, data.data()); + EXPECT_DOUBLE_EQ(model.final_obj, 0.0); + EXPECT_EQ(model.assignments.size(), k); + } + } + EXPECT_THROW(QGAssigner(65537, 33, 32, 33, 33, 1, 1), std::invalid_argument); +} + +namespace { +void normalize_reference_centroids(std::vector& rows, size_t dim) { + for (size_t offset = 0; offset < rows.size(); offset += dim) { + double norm = 0; + for (size_t j = 0; j < dim; ++j) { + norm += static_cast(rows[offset + j]) * rows[offset + j]; + } + if (norm == 0) + continue; + const double scale = 1.0 / std::sqrt(norm); + for (size_t j = 0; j < dim; ++j) { + rows[offset + j] = static_cast(rows[offset + j] * scale); + } + } +} + +void check_uncached_reference(size_t clusters, uint32_t ef_search, size_t intrinsic_dim) { + constexpr size_t kDimension = 1024; + constexpr size_t kPoints = 257; + constexpr size_t kIterations = 3; + std::mt19937 data_rng(823); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector x(kPoints * kDimension); + std::generate(x.begin(), x.end(), [&] { return distribution(data_rng); }); + // Low intrinsic dimension makes points change clusters during Lloyd updates. + // Padding keeps the scalar-sum cache enabled without concentrating distances. + for (size_t point = 0; point < kPoints; ++point) { + std::fill_n( + x.data() + point * kDimension + intrinsic_dim, kDimension - intrinsic_dim, 0.0F + ); + } + for (const uint32_t bits : {0U, 4U, 8U}) { + for (const bool spherical : {false, true}) { + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE( + ::testing::Message() << clusters << "/" << ef_search << "/" << bits + << "/" << spherical << "/" << threads + ); + QGKMeansParameters parameters; + parameters.niter = kIterations; + parameters.num_threads = threads; + parameters.seed = 42; + parameters.min_points_per_centroid = 1; + parameters.quantization_bits = bits; + parameters.spherical = spherical; + parameters.ef_build = 128; + parameters.ef_search = ef_search; + QGKMeans model(kDimension, clusters, parameters); + model.train(kPoints, x.data()); + ASSERT_EQ(model.iteration_stats.size(), kIterations); + + std::vector initial_ids(kPoints); + std::iota(initial_ids.begin(), initial_ids.end(), size_t{0}); + std::mt19937 initialization_rng(parameters.seed); + std::shuffle(initial_ids.begin(), initial_ids.end(), initialization_rng); + std::vector centroids(clusters * kDimension); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + std::copy_n( + x.data() + initial_ids[cluster] * kDimension, + kDimension, + centroids.data() + cluster * kDimension + ); + } + if (spherical) + normalize_reference_centroids(centroids, kDimension); + QGAssigner uncached( + kDimension, + clusters, + 32, + parameters.ef_build, + ef_search, + 1, + threads, + parameters.seed, + spherical ? METRIC_IP : METRIC_L2, + bits + ); + std::vector labels(kPoints); + std::vector distances(kPoints); + for (size_t iteration = 0; iteration < kIterations; ++iteration) { + SCOPED_TRACE(iteration); + // This public call intentionally never receives a prepared sum. + uncached.assign( + centroids.data(), x.data(), kPoints, labels.data(), distances.data() + ); + const double objective = + std::accumulate(distances.begin(), distances.end(), 0.0); + EXPECT_EQ(model.iteration_stats[iteration].obj, objective); + if (iteration == 0) { + EXPECT_FALSE(std::is_sorted(labels.begin(), labels.end())); + } + std::vector sums(clusters * kDimension, 0.0); + std::vector counts(clusters, 0); + for (size_t point = 0; point < kPoints; ++point) { + ++counts[labels[point]]; + for (size_t j = 0; j < kDimension; ++j) { + sums[labels[point] * kDimension + j] += + static_cast(x[point * kDimension + j]); + } + } + for (size_t cluster = 0; cluster < clusters; ++cluster) { + // Keep this an independent Lloyd reference, without empty repair. + ASSERT_GT(counts[cluster], 0U); + const double scale = 1.0 / static_cast(counts[cluster]); + for (size_t j = 0; j < kDimension; ++j) { + centroids[cluster * kDimension + j] = + static_cast(sums[cluster * kDimension + j] * scale); + } + } + if (spherical) + normalize_reference_centroids(centroids, kDimension); + } + // Compare state before the final assignment, whose safeguard is separate. + EXPECT_EQ(model.centroids, centroids); + } + } + } +} +TEST(QGKMeansQuerySumTest, FullGraphTrainingMatchesUncachedLloydReference) { + check_uncached_reference(33, 33, 1024); +} +TEST(QGKMeansQuerySumTest, PartialGraphTrainingMatchesUncachedLloydReference) { + check_uncached_reference(65, 1, 8); +} +} // namespace + +} // namespace rabitqlib::qgkmeans diff --git a/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp new file mode 100644 index 00000000..5a967e68 --- /dev/null +++ b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp @@ -0,0 +1,467 @@ +#include "rabitqlib/clustering/rabitqkmeans.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::rabitqkmeans { +namespace { + +void normalize_reference_centroids(std::vector& rows, size_t dim) { + for (size_t offset = 0; offset < rows.size(); offset += dim) { + double norm = 0; + for (size_t j = 0; j < dim; ++j) { + norm += static_cast(rows[offset + j]) * rows[offset + j]; + } + if (norm == 0) + continue; + const double scale = 1.0 / std::sqrt(norm); + for (size_t j = 0; j < dim; ++j) { + rows[offset + j] = static_cast(rows[offset + j] * scale); + } + } +} + +std::vector scalar_rabitq_labels( + const std::vector& centroids, + const std::vector& data, + size_t dim, + uint32_t seed, + MetricType metric +) { + const size_t clusters = centroids.size() / dim; + const size_t points = data.size() / dim; + const size_t padded = + rotator_impl::padding_requirement(dim, RotatorType::FhtKacRotator); + rotator_impl::FhtKacRotator rotator(dim, padded, seed); + std::vector center(dim), rotated_center(padded), rotated(points * padded); + for (size_t j = 0; j < dim; ++j) { + double sum = 0; + for (size_t point = 0; point < points; ++point) { + sum += data[point * dim + j]; + } + center[j] = static_cast(sum / static_cast(points)); + } + rotator.rotate(center.data(), rotated_center.data()); + for (size_t point = 0; point < points; ++point) { + rotator.rotate(data.data() + point * dim, rotated.data() + point * padded); + } + const size_t batches = (points + 31) / 32; + const size_t bytes = QGBatchDataMap::data_bytes(padded); + std::vector storage(batches * bytes); + for (size_t batch = 0; batch < batches; ++batch) { + const size_t first = batch * 32; + quant::quantize_qg_batch( + rotated.data() + first * padded, + rotated_center.data(), + std::min(size_t{32}, points - first), + padded, + storage.data() + batch * bytes, + metric + ); + } + const auto distance = + metric == METRIC_IP ? dot_product_dis : euclidean_sqr; + std::vector labels(points); + std::vector best(points, std::numeric_limits::infinity()), query(padded); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + const float* centroid = centroids.data() + cluster * dim; + rotator.rotate(centroid, query.data()); + BatchQuery prepared(query.data(), padded, metric); + // This term depends on the moving centroid and cannot be omitted from argmin. + prepared.set_g_add(distance(centroid, center.data(), dim)); + for (size_t point = 0; point < points; ++point) { + ConstQGBatchDataMap batch(storage.data() + (point / 32) * bytes, padded); + const size_t index = point % 32; + const auto lane = std::find( + fastscan::kPerm0.begin(), + fastscan::kPerm0.end(), + static_cast(index % 16) + ); + const size_t packed_lane = static_cast(lane - fastscan::kPerm0.begin()); + int32_t accumulator = 0; + // Independent scalar accumulation over the encoded point's nibble sequence. + for (size_t group = 0; group < padded / 4; ++group) { + const uint8_t packed = batch.bin_code()[group * 16 + packed_lane]; + const size_t code = (packed >> (index / 16 * 4)) & 15U; + accumulator += prepared.lut()[group * 16 + code]; + } + const float ip = prepared.delta() * static_cast(accumulator) + + prepared.sum_vl_lut() + prepared.k1xsumq(); + const float estimate = + batch.f_add()[index] + prepared.g_add() + batch.f_rescale()[index] * ip; + if (estimate < best[point]) { + best[point] = estimate; + labels[point] = static_cast(cluster); + } + } + } + return labels; +} + +double scalar_centroid_distance(const float* x, const float* c, size_t dim, bool ip) { + double distance = 0; + for (size_t j = 0; j < dim; ++j) { + const double left = x[j], right = c[j]; + const double difference = left - right; + distance += ip ? -left * right : difference * difference; + } + return ip ? 1.0 + distance : distance; +} + +TEST(RaBitQKMeansTest, ValidatesCommonParametersBeforeAllocation) { + EXPECT_EQ(RaBitQKMeansParameters{}.niter, 25U); + EXPECT_EQ(RaBitQKMeansParameters{}.final_assignment, FinalAssignmentMode::Approximate); + for (const size_t dimension : {0U, 63U, 65537U}) { + EXPECT_THROW((RaBitQKMeans(dimension, 1)), std::invalid_argument); + } + EXPECT_THROW((RaBitQKMeans(64, 0)), std::invalid_argument); + if constexpr (sizeof(size_t) > sizeof(PID)) { + const size_t max_clusters = + static_cast(std::numeric_limits::max()) + 1; + EXPECT_NO_THROW((RaBitQKMeans(64, max_clusters))); + EXPECT_THROW((RaBitQKMeans(64, max_clusters + 1)), std::invalid_argument); + } + RaBitQKMeansParameters parameters; + parameters.niter = 0; + EXPECT_THROW((RaBitQKMeans(64, 1, parameters)), std::invalid_argument); + parameters.niter = 1; + parameters.final_assignment = static_cast(255); + EXPECT_THROW((RaBitQKMeans(64, 1, parameters)), std::invalid_argument); + parameters.final_assignment = FinalAssignmentMode::Approximate; + parameters.num_threads = std::numeric_limits::max(); + EXPECT_THROW((RaBitQKMeans(64, 1, parameters)), std::invalid_argument); + parameters.num_threads = 1; + RaBitQKMeans model(64, 1, parameters); + const float value = 0.0F; + EXPECT_THROW(model.train(1, nullptr), std::invalid_argument); + EXPECT_THROW(model.train(0, &value), std::invalid_argument); + EXPECT_THROW( + model.train(std::numeric_limits::max() / 64 + 1, &value), + std::invalid_argument + ); +} + +TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroids) { + constexpr uint32_t kSeed = 53; + std::mt19937 rng(157); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + size_t approximate_choices_above_batch_size = 0; + for (const size_t dim : {64U, 65U, 1025U}) { + for (const size_t points : {1U, 31U, 32U, 33U, 65U, 129U}) { + std::vector data(points * dim); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + const auto original = data; + for (const size_t clusters : {1U, 2U, 16U, 31U, 32U, 33U, 65U}) { + if (clusters > points) { + continue; + } + std::vector initial(clusters * dim); + std::generate(initial.begin(), initial.end(), [&] { + return distribution(rng); + }); + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + for (const uint32_t threads : {1U, 4U}) { + SCOPED_TRACE( + ::testing::Message() + << dim << '/' << points << '/' << clusters << '/' + << static_cast(metric) << '/' << threads + ); + auto centers = initial; + RaBitQKMeansParameters parameters; + parameters.seed = kSeed; + parameters.spherical = metric == METRIC_IP; + detail::RaBitQAssigner assigner( + dim, clusters, data.data(), points, threads, parameters + ); + for (size_t pass = 0; pass < 2; ++pass) { + for (float& value : centers) { + value = value * 0.75F + 0.125F; + } + const auto expected = + scalar_rabitq_labels(centers, data, dim, kSeed, metric); + std::vector labels(points, kPidMax); + std::vector distances(points); + assigner.assign( + centers.data(), labels.data(), distances.data() + ); + EXPECT_EQ(labels, expected); + for (size_t point = 0; point < points; ++point) { + ASSERT_LT(labels[point], clusters); + const float* vector = data.data() + point * dim; + const auto exact_distance = metric == METRIC_IP + ? dot_product_dis + : euclidean_sqr; + EXPECT_EQ( + distances[point], + exact_distance( + vector, centers.data() + labels[point] * dim, dim + ) + ); + PID nearest = 0; + double best = std::numeric_limits::infinity(); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + const double exact = scalar_centroid_distance( + vector, + centers.data() + cluster * dim, + dim, + metric == METRIC_IP + ); + if (exact < best) { + best = exact; + nearest = static_cast(cluster); + } + } + if (clusters > 32 && labels[point] != nearest) { + ++approximate_choices_above_batch_size; + } + } + } + } + } + } + EXPECT_EQ(data, original); + } + } + EXPECT_GT(approximate_choices_above_batch_size, 0U); +} + +TEST(RaBitQAssignerTest, ZeroAndDuplicateResidualsChooseFirstLogicalCentroid) { + constexpr size_t kDim = 65, kPoints = 65; + for (const size_t clusters : {1U, 2U, 16U, 31U, 32U, 33U, 65U}) { + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + for (const float value : {0.0F, 0.25F}) { + std::vector centers(clusters * kDim, value), + data(kPoints * kDim, value); + RaBitQKMeansParameters parameters; + parameters.spherical = metric == METRIC_IP; + detail::RaBitQAssigner assigner( + kDim, clusters, data.data(), kPoints, 2, parameters + ); + std::vector labels(kPoints, kPidMax); + std::vector distances(kPoints); + assigner.assign(centers.data(), labels.data(), distances.data()); + EXPECT_EQ(labels, (std::vector(kPoints, 0))); + const double expected = scalar_centroid_distance( + data.data(), centers.data(), kDim, metric == METRIC_IP + ); + for (const float distance : distances) { + EXPECT_NEAR(distance, expected, 1e-6); + } + } + } + } +} + +TEST(RaBitQAssignerTest, KeepsCentroidDependentCorrectionForZeroResidualPointCodes) { + constexpr size_t kDim = 65, kPoints = 33; + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + std::vector data(kPoints * kDim), centers(2 * kDim); + if (metric == METRIC_IP) { + for (size_t point = 0; point < kPoints; ++point) { + data[point * kDim] = 1.0F; + } + centers[1] = 1.0F; + centers[kDim] = 1.0F; + } else { + centers[0] = 1.0F; + } + RaBitQKMeansParameters parameters; + parameters.spherical = metric == METRIC_IP; + detail::RaBitQAssigner assigner(kDim, 2, data.data(), kPoints, 2, parameters); + std::vector labels(kPoints, kPidMax); + std::vector distances(kPoints); + assigner.assign(centers.data(), labels.data(), distances.data()); + EXPECT_EQ(labels, (std::vector(kPoints, 1))); + EXPECT_EQ(distances, (std::vector(kPoints, 0.0F))); + } +} + +TEST(RaBitQKMeansTest, FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts) { + constexpr size_t kDim = 65, kPoints = 67; + std::mt19937 rng(675); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector data(kPoints * kDim); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + for (const size_t clusters : {1U, 2U, 16U, 31U, 32U, 33U, 65U}) { + for (const bool spherical : {false, true}) { + std::vector ids(kPoints); + std::iota(ids.begin(), ids.end(), size_t{0}); + std::mt19937 initialization(42); + std::shuffle(ids.begin(), ids.end(), initialization); + std::vector initial(clusters * kDim); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + std::copy_n( + data.data() + ids[cluster] * kDim, kDim, initial.data() + cluster * kDim + ); + } + if (spherical) { + normalize_reference_centroids(initial, kDim); + } + const auto metric = spherical ? METRIC_IP : METRIC_L2; + const auto expected = scalar_rabitq_labels(initial, data, kDim, 42, metric); + const auto distance = spherical ? dot_product_dis : euclidean_sqr; + double expected_objective = 0; + for (size_t point = 0; point < kPoints; ++point) { + expected_objective += distance( + data.data() + point * kDim, + initial.data() + expected[point] * kDim, + kDim + ); + } + std::vector counts(clusters); + std::vector sums(clusters * kDim); + for (size_t point = 0; point < kPoints; ++point) { + ++counts[expected[point]]; + for (size_t j = 0; j < kDim; ++j) { + sums[expected[point] * kDim + j] += data[point * kDim + j]; + } + } + std::vector means(clusters * kDim); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + if (counts[cluster] != 0) { + const double scale = 1.0 / static_cast(counts[cluster]); + for (size_t j = 0; j < kDim; ++j) { + means[cluster * kDim + j] = + static_cast(sums[cluster * kDim + j] * scale); + } + } + } + if (spherical) { + normalize_reference_centroids(means, kDim); + } + for (const auto mode : + {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + RaBitQKMeansParameters parameters; + parameters.niter = 1; + parameters.num_threads = 2; + parameters.min_points_per_centroid = 1; + parameters.spherical = spherical; + parameters.final_assignment = mode; + RaBitQKMeans model(kDim, clusters, parameters); + model.train(kPoints, data.data()); + ASSERT_EQ(model.iteration_stats.size(), 1U); + EXPECT_DOUBLE_EQ(model.iteration_stats.front().obj, expected_objective); + // Codes only drive assignment; populated means use original float data. + for (size_t cluster = 0; cluster < clusters; ++cluster) { + if (counts[cluster] != 0) { + for (size_t j = 0; j < kDim; ++j) { + EXPECT_EQ( + model.centroids[cluster * kDim + j], + means[cluster * kDim + j] + ); + } + } + } + const auto approximate = + scalar_rabitq_labels(model.centroids, data, kDim, 42, metric); + for (size_t point = 0; point < kPoints; ++point) { + PID nearest = 0; + double best = std::numeric_limits::infinity(); + for (size_t cluster = 0; cluster < clusters; ++cluster) { + const double exact = scalar_centroid_distance( + data.data() + point * kDim, + model.centroids.data() + cluster * kDim, + kDim, + spherical + ); + if (exact < best) { + best = exact; + nearest = static_cast(cluster); + } + } + EXPECT_EQ( + model.assignments[point], + mode == FinalAssignmentMode::Exact ? nearest : approximate[point] + ); + const double selected = scalar_centroid_distance( + data.data() + point * kDim, + model.centroids.data() + model.assignments[point] * kDim, + kDim, + spherical + ); + EXPECT_NEAR( + model.distances[point], + selected, + 2e-6 * std::max(1.0, std::abs(selected)) + ); + } + EXPECT_DOUBLE_EQ( + model.final_obj, + std::accumulate(model.distances.begin(), model.distances.end(), 0.0) + ); + } + } + } +} + +TEST(RaBitQKMeansTest, PointCodesPreserveBorrowedCentroidsDuringRetraining) { + constexpr size_t kDim = 65, kPoints = 129; + std::mt19937 rng(945); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector data(kPoints * kDim); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + for (const size_t kInitialClusters : {32U, 65U}) { + const size_t kClusters = kInitialClusters == 32 ? 16 : 33; + for (const bool spherical : {false, true}) { + for (const auto mode : + {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + for (const size_t offset : {0U, 1U}) { + SCOPED_TRACE( + ::testing::Message() + << spherical << '/' << static_cast(mode) << '/' << offset + ); + RaBitQKMeansParameters parameters; + parameters.niter = 2; + parameters.num_threads = 2; + parameters.min_points_per_centroid = 1; + parameters.spherical = spherical; + parameters.final_assignment = mode; + RaBitQKMeans model(kDim, kInitialClusters, parameters); + model.train(kPoints, data.data()); + const auto original = model.centroids; + const size_t count = kInitialClusters - offset; + parameters.seed = 77; + RaBitQKMeans reference(kDim, kClusters, parameters); + reference.train(count, original.data() + offset * kDim); + model.k = kClusters; + model.seed = parameters.seed; + // The point-code cache must bind to the preserved copy, because + // both updates and returned distances still read original input. + model.train(count, model.centroids.data() + offset * kDim); + EXPECT_EQ(model.centroids, reference.centroids); + EXPECT_EQ(model.assignments, reference.assignments); + EXPECT_EQ(model.distances, reference.distances); + EXPECT_DOUBLE_EQ(model.final_obj, reference.final_obj); + ASSERT_EQ( + model.iteration_stats.size(), reference.iteration_stats.size() + ); + for (size_t i = 0; i < model.iteration_stats.size(); ++i) { + EXPECT_DOUBLE_EQ( + model.iteration_stats[i].obj, reference.iteration_stats[i].obj + ); + } + } + } + } + } +} + +} // namespace +} // namespace rabitqlib::rabitqkmeans diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp index aeffe09f..fb78555e 100644 --- a/tests/unit/rabitqlib/index/pipnn_test.cpp +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -1,13 +1,19 @@ #include "rabitqlib/index/symqg/detail/pipnn.hpp" #include +#include #include #include +#include #include #include +#include #include +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/space.hpp" + namespace rabitqlib::symqg::detail { namespace { TEST(PipnnTest, DenseDistancesMatchReference) { @@ -109,6 +115,14 @@ TEST(PipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { } } +TEST(PipnnTest, PreservesCallerThreadCount) { + const int previous = omp_get_max_threads(); + std::vector data(65 * 65, 0.5F); + const size_t requested = previous == 1 ? 2 : 1; + build_initial_graph(data.data(), 65, 65, 32, METRIC_L2, requested); + EXPECT_EQ(omp_get_max_threads(), previous); +} + TEST(PipnnTest, RejectsInvalidConfiguration) { std::vector data(33 * 7, 0); EXPECT_THROW(build_initial_graph(nullptr, 33, 7, 32), std::invalid_argument); diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 90155ae5..3272c9d4 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -1,6 +1,7 @@ #include "rabitqlib/index/symqg/qg.hpp" #include +#include #include #include @@ -13,7 +14,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -29,12 +32,117 @@ #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/quantization/data_layout.hpp" #include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/simd/estimator_dispatch.hpp" +#include "rabitqlib/utils/buffer.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { struct QGConstructionTestAccess { + static float build_distance(QGBuilder& builder, PID source_id, PID target_id) { + std::vector reconstructed; + std::optional prepared; + const float* source = builder.qg_.prepare_build_query( + source_id, reconstructed, prepared, builder.build_cache() + ); + return builder.qg_.point_distance( + source, prepared ? &*prepared : nullptr, target_id + ); + } + static CandidateList probe_pruned_edges( + QGBuilder& builder, + const CandidateList& current, + const CandidateList& pruned, + float threshold + ) { + CandidateList result; + builder.add_pruned_edges(current, pruned, result, threshold); + return result; + } + static std::array check_refinement_endpoint( + QGBuilder& builder, float cosine, bool duplicate_id + ) { + // Identical current vectors make every pair use the same cosine. Choose + // radial scores from the graph's own distance, including quantized modes. + const float radial = build_distance(builder, 32, 1) / (2.0F * (1.0F - cosine)); + EXPECT_GT(radial, 0.0F); + CandidateList current; + for (PID id = 1; id < 32; ++id) { + current.emplace_back(id, radial); + } + CandidateList pruned{AnnCandidate(duplicate_id ? 1 : 32, radial)}; + float left = 0.5F; + float right = 1.0F; + size_t successful = 0; + size_t unsuccessful = 0; + bool last_successful = false; + for (size_t iteration = 0; iteration < kMaxBsIter; ++iteration) { + const float mid = (left + right) / 2.0F; + const auto probe = probe_pruned_edges(builder, current, pruned, mid); + last_successful = probe.size() == 32; + if (last_successful) { + ++successful; + right = mid; + } else { + ++unsuccessful; + left = mid; + } + } + // The original algorithm always reevaluates the final right endpoint. + auto expected = probe_pruned_edges(builder, current, pruned, right); + if (expected.size() < 32) { + // Only one slot is missing; independently reproduce its seeded choice. + std::mt19937 random(builder.seed_ ^ 0x80000000U); + std::uniform_int_distribution distribution( + 0, static_cast(builder.num_nodes_) - 1 + ); + PID selected = 0; + do { + selected = distribution(random); + } while (selected < 32); + expected.emplace_back(selected, build_distance(builder, 0, selected)); + } + for (auto& row : builder.new_neighbors_) { + row.assign(32, AnnCandidate(1, 0.0F)); + } + builder.new_neighbors_[0] = current; + builder.pruned_neighbors_[0] = pruned; + builder.graph_refine(); + const auto& actual = builder.new_neighbors_[0]; + EXPECT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < std::min(actual.size(), expected.size()); ++i) { + EXPECT_EQ(actual[i].id, expected[i].id); + uint32_t actual_bits = 0; + uint32_t expected_bits = 0; + std::memcpy(&actual_bits, &actual[i].distance, sizeof(actual_bits)); + std::memcpy(&expected_bits, &expected[i].distance, sizeof(expected_bits)); + EXPECT_EQ(actual_bits, expected_bits); + } + return {successful, unsuccessful, static_cast(last_successful)}; + } + + static void disconnect(QuantizedGraph& graph) { + for (PID i = 0; i < 66; ++i) { + std::vector> row; + for (PID j = (i / 33) * 33; j < (i / 33 + 1) * 33; ++j) { + if (i != j) { + row.emplace_back(j, 0.0F); + } + } + graph.update_qg(i, row); + } + } + + static auto rows(const QuantizedGraph& graph) { + std::vector result; + for (PID id = 0; id < graph.num_points_; ++id) { + const char* row = graph.get_quantized_vector(id); + result.insert(result.end(), row, row + graph.row_offset_); + } + return result; + } static void check_contiguous_rows(QuantizedGraph& graph) { const size_t vector_bytes = graph.is_quantized() @@ -177,6 +285,451 @@ struct QGConstructionTestAccess { namespace { +TEST(QGConstructionTest, RefinementMatchesOriginalFinalEndpointAndSeededFallback) { + constexpr size_t kDim = 65; + constexpr size_t kCount = 65; + std::vector data(kCount * kDim, 0.0F); + for (PID id = 1; id < 32; ++id) { + data[id * kDim] = 1.0F; + } + data[32 * kDim + 1] = 1.0F; + for (PID id = 33; id < kCount; ++id) { + data[id * kDim + id % kDim] = 0.25F; + } + for (size_t bits : {0U, 4U, 8U}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + for (bool cache_vectors : {false, true}) { + for (size_t threads : {1U, 4U}) { + SCOPED_TRACE( + ::testing::Message() << bits << "/" << static_cast(metric) + << "/" << cache_vectors << "/" << threads + ); + QuantizedGraph graph( + kCount, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder builder( + graph, + 32, + data.data(), + threads, + QGInitialization::Random, + 42, + cache_vectors + ); + // All probes succeed, ending at the final successful threshold. + EXPECT_EQ( + QGConstructionTestAccess::check_refinement_endpoint( + builder, 0.25F, false + ), + (std::array{5, 0, 1}) + ); + // A later failed probe must not replace the saved right endpoint. + EXPECT_EQ( + QGConstructionTestAccess::check_refinement_endpoint( + builder, 0.8F, false + ), + (std::array{2, 3, 0}) + ); + // No probe succeeds; the previously untested endpoint 1.0 does. + EXPECT_EQ( + QGConstructionTestAccess::check_refinement_endpoint( + builder, 0.9921875F, false + ), + (std::array{0, 5, 0}) + ); + // Duplicate IDs retain the existing early exit and seeded fill. + EXPECT_EQ( + QGConstructionTestAccess::check_refinement_endpoint( + builder, 0.8F, true + ), + (std::array{0, 5, 0}) + ); + } + } + } + } +} + +TEST(BatchQueryTest, PreparedSumPreservesLookupTableAndMetric) { + BatchQuery reference, prepared; + for (const size_t dimension : {64U, 1024U}) { + std::vector query(dimension), scratch(dimension * 4); + for (size_t i = 0; i < dimension; ++i) { + query[i] = std::sin(static_cast(i) * 0.37F); + } + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + reference.reset(query.data(), dimension, scratch.data(), metric); + const float prepared_sum = reference.k1xsumq(); + prepared.set_g_add(17.0F); + prepared.reset(query.data(), dimension, scratch.data(), metric, &prepared_sum); + EXPECT_EQ(prepared.g_add(), 0.0F); + EXPECT_EQ(prepared.k1xsumq(), reference.k1xsumq()); + EXPECT_EQ(prepared.delta(), reference.delta()); + EXPECT_EQ(prepared.sum_vl_lut(), reference.sum_vl_lut()); + EXPECT_TRUE( + std::equal(reference.lut(), reference.lut() + dimension * 4, prepared.lut()) + ); + reference.set_g_add(1.5F); + prepared.set_g_add(1.5F); + EXPECT_EQ(prepared.g_add(), reference.g_add()); + } + } +} + +TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { + constexpr size_t kNumPoints = 257; + for (const size_t kDim : {64U, 65U, 1025U, 4097U}) { + SCOPED_TRACE(kDim); + constexpr size_t kDegree = 32; + constexpr size_t kEfSearch = 1; + const size_t kNumQueries = kDim == 64 ? 256 : 16; + std::mt19937 rng(123); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector data(kNumPoints * kDim); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + + for (const size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(bits); + for (const MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(static_cast(metric)); + QuantizedGraph graph( + kNumPoints, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder builder(graph, kDegree, data.data(), 1, 42); + graph.set_ef(kEfSearch); + + std::vector rotated_query(graph.padded_dim()); + std::vector estimated_distances(kDegree); + std::vector lookup_table(graph.padded_dim() * 4); + BatchQuery batch_query; + buffer::SearchBuffer search_pool(kEfSearch); + buffer::SearchBuffer result_pool(1); + VisitedSet visited(kNumPoints, kNumPoints / 10); + + PID label = kPidMax; + float distance = 0.0F; + EXPECT_THROW( + graph.search_with_scratch( + data.data(), + 1, + &label, + &distance, + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited + ), + std::logic_error + ); + builder.build(2); + + for (size_t query_index = 0; query_index < kNumQueries; ++query_index) { + SCOPED_TRACE(query_index); + std::vector query(kDim); + std::generate(query.begin(), query.end(), [&] { + return distribution(rng); + }); + PID regular_label = kPidMax; + float regular_distance = 0.0F; + graph.search(query.data(), 1, ®ular_label, ®ular_distance); + + PID scratch_label = kPidMax; + float scratch_distance = 0.0F; + graph.search_with_scratch( + query.data(), + 1, + &scratch_label, + &scratch_distance, + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited + ); + + EXPECT_EQ(scratch_label, regular_label); + EXPECT_FLOAT_EQ(scratch_distance, regular_distance); + const float expected_distance = + metric == METRIC_IP + ? dot_product_dis( + query.data(), data.data() + scratch_label * kDim, kDim + ) + : euclidean_sqr( + query.data(), data.data() + scratch_label * kDim, kDim + ); + if (bits == 0) { + EXPECT_NEAR(scratch_distance, expected_distance, 1e-5F); + } + } + + if (metric == METRIC_L2 && bits == 0) { + const PID hint = static_cast(kNumPoints - 1); + PID hinted_label = kPidMax; + float hinted_distance = 1.0F; + graph.search_with_scratch( + data.data() + hint * kDim, + 1, + &hinted_label, + &hinted_distance, + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited, + hint + ); + EXPECT_EQ(hinted_label, hint); + EXPECT_FLOAT_EQ(hinted_distance, 0.0F); + } + } + } + } +} + +TEST(QuantizedGraphSearchTest, SingleResultMatchesTopKWithTiesAndHints) { + constexpr size_t kCount = 65; + for (size_t dim : {64U, 65U}) { + for (size_t bits : {0U, 4U, 8U}) { + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (bool tied : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << dim << "/" << bits << "/" << metric << "/" << tied + ); + std::vector data(kCount * dim, 0.25F); + if (!tied) { + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + } + QuantizedGraph graph( + kCount, dim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder builder(graph, 64, data.data(), 1, 42); + builder.build(); + graph.set_ef(kCount); + std::vector rotated(graph.padded_dim()); + std::vector estimates(graph.degree_bound()); + std::vector lut(graph.padded_dim() * 4); + BatchQuery query; + buffer::SearchBuffer search_pool(kCount); + buffer::SearchBuffer one_result(1), two_results(2); + VisitedSet visited(kCount); + for (PID hint : {kPidMax, graph.entry_point(), PID{64}}) { + for (size_t i = 0; i < 4; ++i) { + PID single = kPidMax; + float distance = 0; + std::array expected{}; + std::array expected_distances{}; + graph.search_with_scratch( + data.data() + i * dim, + 2, + expected.data(), + expected_distances.data(), + rotated.data(), + estimates.data(), + lut.data(), + query, + search_pool, + two_results, + visited, + hint + ); + graph.search_with_scratch( + data.data() + i * dim, + 1, + &single, + &distance, + rotated.data(), + estimates.data(), + lut.data(), + query, + search_pool, + one_result, + visited, + hint + ); + EXPECT_EQ(single, expected[0]); + EXPECT_EQ(distance, expected_distances[0]); + } + } + } + } + } + } +} + +TEST(QuantizedGraphSearchTest, ScratchSearchRejectsDefaultVisitedSet) { + constexpr size_t kNumPoints = 66; + constexpr size_t kDim = 64; + constexpr size_t kDegree = 32; + constexpr size_t kEfSearch = 32; + std::mt19937 rng(123); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector data(kNumPoints * kDim); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + + QuantizedGraph graph(kNumPoints, kDim, kDegree); + QGBuilder builder(graph, kDegree, data.data(), 1, 42); + builder.build(2); + graph.set_ef(kEfSearch); + + std::vector rotated_query(graph.padded_dim()); + std::vector estimated_distances(kDegree); + std::vector lookup_table(graph.padded_dim() * 4); + BatchQuery batch_query; + buffer::SearchBuffer search_pool(kEfSearch); + buffer::SearchBuffer result_pool(1); + VisitedSet visited; + EXPECT_FALSE(visited.initialized()); + try { + PID label = kPidMax; + float distance = 0.0F; + graph.search_with_scratch( + data.data(), + 1, + &label, + &distance, + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited + ); + FAIL() << "Expected an uninitialized visited set to be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "SymphonyQG visited set must be initialized"); + } + visited.initialize(kNumPoints, kNumPoints / 10); + + for (size_t query_index = 0; query_index < 2; ++query_index) { + const float* query = data.data() + query_index * kDim; + PID regular_label = kPidMax; + PID scratch_label = kPidMax; + float regular_distance = 0.0F; + float scratch_distance = 0.0F; + graph.search(query, 1, ®ular_label, ®ular_distance); + graph.search_with_scratch( + query, + 1, + &scratch_label, + &scratch_distance, + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited + ); + EXPECT_TRUE(visited.initialized()); + EXPECT_EQ(scratch_label, regular_label); + EXPECT_FLOAT_EQ(scratch_distance, regular_distance); + } +} + +TEST(QuantizedGraphSearchTest, LookupTableResetMatchesFreshAcrossSizesAndModes) { + Lut reused; + for (const size_t dim : {64U, 4160U, 128U, 64U}) { + SCOPED_TRACE(dim); + std::vector query(dim), scratch(dim * 4); + std::vector hacc_scratch(dim * 4); + for (const bool hacc : {false, true, false}) { + SCOPED_TRACE(hacc); + for (const bool zero : {false, true}) { + for (size_t i = 0; i < dim; ++i) { + query[i] = zero ? 0.0F : static_cast(i % 23) - 11.0F; + } + const Lut fresh(query.data(), dim, hacc); + reused.reset(query.data(), dim, scratch.data(), hacc_scratch.data(), hacc); + EXPECT_FLOAT_EQ(reused.delta(), fresh.delta()); + EXPECT_FLOAT_EQ(reused.sum_vl(), fresh.sum_vl()); + const size_t bytes = dim * 4 * (hacc ? 2 : 1); + EXPECT_TRUE(std::equal(fresh.lut(), fresh.lut() + bytes, reused.lut())); + EXPECT_TRUE(std::isfinite(reused.delta())); + EXPECT_TRUE(std::isfinite(reused.sum_vl())); + } + } + } +} + +TEST(QGConstructionTest, CopiedCacheOwnsItsQueryStorage) { + constexpr size_t kCount = 65, kDim = 65; + for (size_t bits : {4U, 8U}) { + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i)); + } + QuantizedGraph graph( + kCount, kDim, 32, METRIC_L2, RotatorType::FhtKacRotator, bits, 42 + ); + auto original = std::make_unique(graph, 64, data.data(), 1, 42, true); + QGBuilder copy(*original); + original.reset(); + std::fill(data.begin(), data.end(), std::numeric_limits::quiet_NaN()); + std::vector().swap(data); + copy.build(); + for (const auto& row : QGConstructionTestAccess::neighbors(copy)) { + ASSERT_EQ(row.size(), 32U); + for (const auto& neighbor : row) { + EXPECT_TRUE(std::isfinite(neighbor.distance)); + } + } + } +} + +TEST(QGConstructionTest, CachedResetMatchesFreshConstruction) { + constexpr size_t kCount = 65; + for (auto init : {QGInitialization::Random, QGInitialization::PiPNN}) { + for (size_t dim : {65U, 1025U}) { + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE( + ::testing::Message() << dim << "/" << metric << "/" << bits + ); + std::mt19937 rng(123); + std::normal_distribution normal; + std::vector data(kCount * dim); + std::generate(data.begin(), data.end(), [&] { return normal(rng); }); + QuantizedGraph cached( + kCount, dim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder builder(cached, 64, data.data(), 1, init, 42, true); + for (size_t iteration = 0; iteration < 3; ++iteration) { + SCOPED_TRACE(iteration); + if (iteration > 0) { + for (float& value : data) { + value += 0.2F * normal(rng); + } + builder.reset(data.data()); + } + QuantizedGraph fresh( + kCount, dim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGBuilder fresh_builder(fresh, 64, data.data(), 1, init, 42); + builder.build(); + fresh_builder.build(); + EXPECT_EQ(cached.entry_point(), fresh.entry_point()); + EXPECT_EQ( + QGConstructionTestAccess::rows(cached), + QGConstructionTestAccess::rows(fresh) + ); + } + } + } + } + } +} + TEST(QuantizedGraphCompatibilityTest, SupportsExplicitAndDeducedFloatTypes) { using Graph = rabitqlib::symqg::QuantizedGraph; static_assert(std::is_same_v>); @@ -1022,3 +1575,42 @@ TEST(QGSearchTest, ParallelQueriesMatchSerialAcrossIndexesAndSettings) { } // namespace } // namespace rabitqlib::symqg + +namespace rabitqlib::symqg { +TEST(QuantizedGraphSearchTest, ScratchSearchRejectsIncompleteResults) { + QuantizedGraph graph(66, 64, 32, METRIC_L2, RotatorType::FhtKacRotator, 0, 42); + std::vector data(66 * 64, 0.0F); + QGBuilder builder(graph, 66, data.data(), 1, 42); + builder.build(2); + QGConstructionTestAccess::disconnect(graph); + graph.set_ef(66); + std::vector ids(66, kPidMax); + std::vector distances(66, -1.0F), rotated(64), estimates(32), lut(256); + BatchQuery query; + buffer::SearchBuffer search(66), result(66); + VisitedSet visited(66, 66); + EXPECT_THROW( + graph.search(data.data(), 66, ids.data(), distances.data()), std::runtime_error + ); + EXPECT_THROW( + graph.search_with_scratch( + data.data(), + 66, + ids.data(), + distances.data(), + rotated.data(), + estimates.data(), + lut.data(), + query, + search, + result, + visited + ), + std::runtime_error + ); + EXPECT_TRUE(std::all_of(ids.begin(), ids.end(), [](PID id) { return id == kPidMax; })); + EXPECT_TRUE(std::all_of(distances.begin(), distances.end(), [](float d) { + return d == -1.0F; + })); +} +} // namespace rabitqlib::symqg diff --git a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp index 6330dcba..e853395e 100644 --- a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp +++ b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp @@ -5,6 +5,10 @@ #include #include #include +#include +#include +#include +#include #include #include "rabitqlib/utils/cpu_features.hpp" @@ -109,5 +113,176 @@ TEST(MatrixDispatchTest, AllBackendsMatchScalarForRectangularUnalignedInputs) { } } } +using AccumulationBackend = std::pair; + +std::vector accumulation_backends() { + std::vector backends{ + {"dispatch", accumulate_cluster_sums}, + {"generic", accumulate_cluster_sums_generic}}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + backends.emplace_back("avx2", accumulate_cluster_sums_avx2); + } + if (cpu::has_avx512_core()) { + backends.emplace_back("avx512", accumulate_cluster_sums_avx512); + } +#endif + return backends; +} + +TEST(MatrixDispatchTest, IndexedAccumulationMatchesOrderedReferenceWithUnalignedTails) { + constexpr size_t kPoints = 97; + constexpr size_t kClusters = 7; + const auto backends = accumulation_backends(); + for (const size_t dim : + {1U, + 3U, + 4U, + 7U, + 8U, + 15U, + 16U, + 31U, + 32U, + 63U, + 64U, + 65U, + 127U, + 128U, + 129U, + 1536U, + 1537U}) { + SCOPED_TRACE(dim); + std::vector data(kPoints * dim + 2, 12345.0F); + std::vector labels(kPoints + 2, 99); + for (size_t point = 0; point < kPoints; ++point) { + labels[point + 1] = static_cast((point * 3) % kClusters); + } + for (const bool zero : {false, true}) { + SCOPED_TRACE(zero); + for (size_t index = 0; index < kPoints * dim; ++index) { + data[index + 1] = + zero ? (index % 2 == 0 ? 0.0F : -0.0F) + : std::ldexp( + static_cast(static_cast(index % 37) - 18), + static_cast(index % 21) - 10 + ); + } + for (const size_t stride : {1U, 3U, 100U}) { + SCOPED_TRACE(stride); + std::vector point_ids; + for (size_t point = 0; point < kPoints; point += stride) { + point_ids.push_back(point); + } + if (point_ids.size() > 2) { + point_ids.push_back(point_ids[1]); + std::reverse(point_ids.begin(), point_ids.end()); + } + std::vector initial(kClusters * dim + 2, 12345.0); + std::vector initial_counts(kClusters + 2, 12345); + for (size_t index = 0; index < kClusters * dim; ++index) { + initial[index + 1] = zero ? (index % 2 == 0 ? 0.0 : -0.0) + : static_cast(index % 11) * 0.125; + } + for (size_t cluster = 0; cluster < kClusters; ++cluster) { + initial_counts[cluster + 1] = cluster + 2; + } + auto expected = initial; + auto expected_counts = initial_counts; + for (const size_t point : point_ids) { + const size_t cluster = labels[point + 1]; + ++expected_counts[cluster + 1]; + for (size_t coordinate = 0; coordinate < dim; ++coordinate) { + expected[1 + cluster * dim + coordinate] += + static_cast(data[1 + point * dim + coordinate]); + } + } + for (const auto& backend : backends) { + SCOPED_TRACE(backend.first); + auto actual = initial; + auto actual_counts = initial_counts; + const size_t split = point_ids.size() / 2; + backend.second( + data.data() + 1, + point_ids.data(), + split, + dim, + labels.data() + 1, + actual.data() + 1, + actual_counts.data() + 1 + ); + backend.second( + data.data() + 1, + point_ids.data() + split, + point_ids.size() - split, + dim, + labels.data() + 1, + actual.data() + 1, + actual_counts.data() + 1 + ); + EXPECT_EQ(actual_counts, expected_counts); + EXPECT_EQ( + std::memcmp( + actual.data(), expected.data(), actual.size() * sizeof(double) + ), + 0 + ); + } + } + } + } +} + +TEST(MatrixDispatchTest, IndexedAccumulationPreservesPointOrderUnderCancellation) { + constexpr size_t kDimension = 65; + const std::vector values{1e7F, 1e7F, 1e-10F, -1e7F, -1e7F, 1e-10F}; + const std::vector labels{0, 1, 0, 1, 0, 1}; + std::vector data(values.size() * kDimension); + std::vector point_ids(values.size()); + std::iota(point_ids.begin(), point_ids.end(), size_t{0}); + for (size_t point = 0; point < values.size(); ++point) { + std::fill_n(data.data() + point * kDimension, kDimension, values[point]); + } + for (const auto& backend : accumulation_backends()) { + SCOPED_TRACE(backend.first); + std::vector sums(2 * kDimension); + std::vector counts(2); + backend.second( + data.data(), + point_ids.data(), + 3, + kDimension, + labels.data(), + sums.data(), + counts.data() + ); + backend.second( + data.data(), + point_ids.data() + 3, + 3, + kDimension, + labels.data(), + sums.data(), + counts.data() + ); + EXPECT_EQ(counts, (std::vector{3, 3})); + for (size_t coordinate = 0; coordinate < kDimension; ++coordinate) { + EXPECT_DOUBLE_EQ(sums[coordinate], 0.0); + EXPECT_DOUBLE_EQ(sums[kDimension + coordinate], static_cast(1e-10F)); + } + } +} + +TEST(MatrixDispatchTest, EmptyIndexedAccumulationLeavesOutputsUntouched) { + for (const auto& backend : accumulation_backends()) { + SCOPED_TRACE(backend.first); + double sum = -7.5; + size_t count = 23; + backend.second(nullptr, nullptr, 0, 65, nullptr, &sum, &count); + EXPECT_DOUBLE_EQ(sum, -7.5); + EXPECT_EQ(count, 23U); + } +} + } // namespace } // namespace rabitqlib::simd