diff --git a/CMakeLists.txt b/CMakeLists.txt index d63881f3..7b1d103f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -85,6 +85,7 @@ endfunction() set(RABITQ_CORE_SOURCES src/core/api.cpp src/core/fast_scale_table.cpp + src/core/provisional.cpp src/utils/cpu_features.cpp src/utils/thread_count.cpp src/simd/dispatch.cpp diff --git a/docs/development.md b/docs/development.md index 77a3f6a1..5800fe56 100644 --- a/docs/development.md +++ b/docs/development.md @@ -41,6 +41,7 @@ factors estimate L2 distance or inner product. | Path | Responsibility | | --- | --- | | `include/rabitqlib/core/api.hpp` | Stable float32 v1 API; code/factors separate, caller workspace, source compatibility | +| `include/rabitqlib/core/provisional/` | Provisional quantization API (`quant::provisional`) behind `rabitqlib.quant`; may change in any minor release and is not part of v1 | | `include/rabitqlib/core/` | Per-vector encoding, query/estimation, rotation and per-vector byte layouts | | `include/rabitqlib/fastscan/` | All 32-vector batch layouts, encoding adapters, packing, LUT/estimation and FastScan/HACC | | `include/rabitqlib/core/types.hpp`, `core/detail/eigen.hpp` | Numerical configuration/types and internal Eigen aliases, respectively | @@ -48,13 +49,14 @@ factors estimate L2 distance or inner product. | `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/{core,fastscan,index}/detail/` | Module-owned implementation headers; no public stability promise | -| `src/core/` | Compiled stable core v1 API implementation and the frozen `faster_config` scale table | +| `src/core/` | Compiled stable core v1 API and provisional API implementations, and the frozen `faster_config` scale table | +| `docs/formats.md` | Versioned byte formats of `rabitqlib.quant`: raw arrays and self-describing blobs | | `src/simd/` | Core kernels, dispatch, and the private shared resolver | | `src/fastscan/` | FastScan kernels, batch estimators, and paired backend selection | | `src/index/` | Compiled IVF/HNSW/SymphonyQG operations and private search kernels | | `src/utils/cpu_features.cpp` | Generic CPU feature detection | | `include/rabitqlib/utils/` | Shared allocation, I/O and numerical helpers; index buffers live in `index/detail/` | -| `python_bindings/` | pybind11 extension and index wrappers | +| `python_bindings/` | pybind11 extension, index wrappers and the provisional `rabitqlib.quant` module (`quant_bindings.cpp`, `quant.py`) | | `sample/cpp/`, `sample/python/` | Usage examples | | `tests/unit/`, `tests/python/` | C++ reference/unit tests and Python boundary/integration tests | | `docs/docs/` | MkDocs site sources | diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md new file mode 100644 index 00000000..12df7922 --- /dev/null +++ b/docs/docs/rabitq/quant_python.md @@ -0,0 +1,132 @@ +# Python quantization building blocks + +!!! warning "Provisional" + `rabitqlib.quant` is provisional in the sense of [PEP 411](https://peps.python.org/pep-0411/): any minor release may change or remove part of its API, with a note in the release notes, until it is declared stable. Byte formats are versioned from the first release, and bytes of an unknown version raise `quant.FormatError` instead of giving wrong results. + +`rabitqlib.quant` exposes the pieces the indexes are built from: rotation, quantization, correction factors, packed codes and distance estimation. Use it to store RaBitQ codes in your own database or file format and to score them with your own selection logic. For a ready-made index, use [IVF-RaBitQ](../index/ivf.md), [HNSW-RaBitQ](../index/hnsw.md) or [SymphonyQG](../index/qg.md) instead. + +```python +from rabitqlib import quant +``` + +`import rabitqlib` does not import the module; import it explicitly. + +## The working space + +Every layer call works on rotated, padded vectors: float32 arrays of shape `(n, d)` with `d` a positive multiple of 32. `Rotator` maps original vectors into that space; callers may bring their own rotation instead. Rotate the data, the centroids and the queries with the same rotator. + +```python +import numpy as np +from rabitqlib import quant + +rng = np.random.default_rng(0) +data = rng.standard_normal((1000, 100)).astype(np.float32) +query = rng.standard_normal(100).astype(np.float32) + +rot = quant.Rotator(100, seed=42) # padded_dim = 128 +x = rot.rotate(data) # float32[1000, 128] +q = rot.rotate(query[None])[0] +c = x.mean(axis=0) # a centroid in the working space +``` + +A seeded `"fht_kac"` rotation (the default, for 64 dimensions and up) is identical on every platform. A `"matrix"` rotation drawn from a seed is reproducible only with the same standard library and build, so persist rotators with `to_bytes()` and load them with `Rotator.from_bytes(data)`. Rotators also pickle. + +## Encoding + +```python +residuals = x - c +codes = quant.quantize(residuals, nbits=4) # uint8[1000, 128] +prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) # 1 sign bit + 3 extra bits +f_full = quant.code_factors(residuals, codes, nbits=4, metric="l2", centroid=c) +f_prefix = quant.code_factors(residuals, prefix, nbits=1, metric="l2", centroid=c) +bitsets = quant.pack_codes(prefix, nbits=1) # uint8[1000, 16] +packed_ex = quant.pack_codes(ex, nbits=3) # uint8[1000, 48] +``` + +- `nbits` is the total code width, 1–9, including the sign bit. Codes are unsigned integers in `[0, 2**nbits)`: `uint8` up to 8 bits, `uint16` at 9. +- `split` keeps the top `prefix_bits` bits as `prefix` and the low `ex_bits = nbits - prefix_bits` bits as `ex`. The prefix is itself a valid code of the same residual, so `code_factors` with `nbits=prefix_bits` gives its factors; `ex` has no factors of its own, because refinement pairs it with the full-code factors. +- `metric` is required on `code_factors`, `Query` and `query_terms`: factors for the wrong metric give wrong distances without an error. `centroid` is in the working space; `None` means zero. +- `quantize(..., fast_quantization=True)` uses the library's precomputed scale table instead of the exact search, as the indexes' `faster_config` does. +- `pack_codes` writes the layout of [compact code storage](../compact_code.md) for 1–8 bits; at `nbits=1` it equals `numpy.packbits(prefix, axis=1, bitorder="little")`. `unpack_codes` reverses it losslessly. Split 9-bit codes before packing. + +Factors are structured arrays of `quant.FACTORS_DTYPE` (`add`, `rescale`, `error`), three float32 fields per record. Consumers that do not take structured dtypes can use `factors.view(np.float32).reshape(-1, 3)`. + +## Estimation + +```python +qry = quant.Query(q, metric="l2", centroid=c) +est = qry.estimate_bitset(bitsets, f_prefix) # ESTIMATE_DTYPE[1000] +lower = est["distance"] - est["error"] +upper = est["distance"] + est["error"] +candidates = np.flatnonzero(lower <= np.partition(upper, 9)[9]) # 10th-smallest upper bound +dist = qry.refine_packed(packed_ex, f_full, est[candidates], indices=candidates, ex_bits=3) +top = candidates[np.argsort(dist)[:10]] +``` + +`Query` copies the rotated query. `set_centroid(c)` switches to another list's factors without preparing the query again. + +| Call | Reads | Returns | +| --- | --- | --- | +| `estimate_bitset(bitsets, factors)` | `pack_codes(prefix, nbits=1)` with the 1-bit factors, scored with the float query | `ESTIMATE_DTYPE` | +| `estimate_codes(codes, factors, nbits=…)` | Unpacked codes of any width with their factors | `ESTIMATE_DTYPE` | +| `refine(ex, f_full, first, ex_bits=…)` | Unpacked extra bits, the full-code factors and the first-stage estimates | `float32` | +| `refine_packed(packed, f_full, first, ex_bits=…)` | The same with `pack_codes(ex, nbits=ex_bits)` | `float32` | +| `estimate_fastscan(blocks, num_points=…)` | `pack_fastscan_blocks` output | `ESTIMATE_DTYPE` | + +An estimate has three fields: `distance`, the `error` term, and `refine_input`, the prefix inner product that refinement reads. `distance ± m * error` is a probabilistic interval for the true distance. `error` already includes the encoder's confidence parameter ε₀ = 1.9, so `m = 1` is the interval the indexes use, and about 6% of true distances fall outside it; a larger `m` widens it, refining more candidates for fewer misses. To keep the top `k` within that confidence, refine every vector whose lower bound is at most the `k`-th smallest upper bound, as in the example above; comparing lower bounds with the `k`-th smallest estimate instead drops true neighbours. `refine` and `refine_packed` take `prefix_bits=` (default 1) for the width the first stage scored, and agree with each other within float32 rounding: the packed path uses the dispatched SIMD kernel. The scoring calls accept `indices=` to read just those rows in place; `first` and the output then line up with `indices`. Indices must be integers in `[0, n)`. + +Distances follow the indexes. L2 is the estimated squared distance; IP is `1 - `. Lower is better. + +### FastScan blocks + +```python +blocks = quant.pack_fastscan_blocks(bitsets, f_prefix) # uint8[32, 128 * 4 + 384] +est_fs = qry.estimate_fastscan(blocks, num_points=len(bitsets)) +``` + +A block always holds 32 vectors, the library's fixed FastScan batch; the last block is zero-padded. `estimate_fastscan` scores with a lookup table built from a quantized query on the first call, so its estimates agree with `estimate_bitset` only within that table's error. FastScan accumulation needs AVX2/FMA, AVX-512 or NEON, and raises `RuntimeError` on a CPU with none of them. + +### The reference formula + +Every estimate is an inner product followed by one formula. For an `nbits`-bit code with unsigned values `u`, factors `f` and the query terms of `query_terms`: + +```text +ip = +distance = f.add + g_add + f.rescale * (ip + g_k1xsumq * (2^nbits - 1)) +error = f.error * g_error + +g_k1xsumq = -0.5 * sum(q) +L2: g_add = |q - c|^2 IP: g_add = - both: g_error = |q - c| +``` + +For IP, the `1` of `1 - ` is folded into `f.add`, not into `g_add`. For a split code, `ip_full = 2^ex_bits * ip_prefix + ip_ex`, so refinement is the same formula with the full-code factors and `nbits = prefix_bits + ex_bits`. `quant.estimate_from_ip(ip, factors, terms, nbits=…)` applies the formula to inner products from your own kernel: + +```python +terms = quant.query_terms(q, metric="l2", centroid=c) # QUERY_TERMS_DTYPE record +ip = codes.astype(np.float32) @ q +est4 = quant.estimate_from_ip(ip, f_full, terms, nbits=4) +``` + +If your `ip` is itself approximate, `error` covers only the code's quantization error, not the error in `ip`. + +## Threads, copies and pickling + +- Build-side calls (`Rotator.rotate`, `quantize`, `split`, `code_factors`, `pack_codes`, `unpack_codes`, `pack_fastscan_blocks`) default to `num_threads=0`, the available CPU count capped by `OMP_NUM_THREADS`, as for the indexes (see [threading](../index/usage.md#threading-and-file-paths)); their results do not depend on the thread count. Scoring calls (`estimate_*`, `refine*`, `estimate_from_ip`) default to 1. +- The GIL is released during computation. `Rotator` is immutable and safe to share. Scoring calls on one `Query` may run concurrently; a `set_centroid` during one raises `RuntimeError`. Use `qry.copy()` for one query per thread. +- `Query` supports `copy.copy`, `copy.deepcopy` and pickling of its query vector, metric and centroid; a copy never shares storage with the original. Queries have no `to_bytes`, because they are not a stored format. +- Every call that returns an array accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. For `split`, `out=` is a `(prefix, ex)` pair of arrays, which must not overlap each other either. A call that raises leaves a caller's `out=` unchanged, except after `OverflowError` or `MemoryError`, when its contents are unspecified. +- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. Factors and query terms that are not representable in float32 raise `OverflowError`, and so does an estimated distance, as in `quant::v1`, or an estimated error. Inputs are checked before anything is computed; within each check, the error reported is that of the first failing row in call order (the lowest row, or the earliest entry of `indices=`), whatever the thread count. + +## Byte formats + +The layouts are specified in [`docs/formats.md`](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/docs/formats.md) in the repository. Raw arrays carry no header; record the version constant next to them: + +| Constant | Covers | +| --- | --- | +| `quant.PACKED_CODES_VERSION` | `pack_codes` output, including the 1-bit bitset | +| `quant.FACTORS_VERSION` | `FACTORS_DTYPE`, `ESTIMATE_DTYPE` and `QUERY_TERMS_DTYPE` | +| `quant.FASTSCAN_BLOCKS_VERSION` | `pack_fastscan_blocks` output | + +`Rotator.to_bytes()` writes a self-describing blob with a magic, a kind, a version and a CRC-32. A change to any layout is a new version, never an edit of an existing one, and golden files in the test suite freeze each version. While the module is provisional, a release may stop reading an old version; reading it then raises `FormatError`. A promise that every released version stays readable comes with stabilisation. + +Byte layouts, rotator blobs and seeded FHT/Kac rotations are identical everywhere. Rotated values, codes, factors and distances are identical with the same wheel and CPU features, and otherwise agree within float32 rounding, because kernels are chosen at runtime by CPU. diff --git a/docs/docs/releases/unreleased.md b/docs/docs/releases/unreleased.md new file mode 100644 index 00000000..d2ec380a --- /dev/null +++ b/docs/docs/releases/unreleased.md @@ -0,0 +1,10 @@ +# Unreleased + +## Python changes + +- New provisional module `rabitqlib.quant` with the quantization building blocks behind the indexes: `Rotator` (seeded FHT/Kac or matrix rotation, persisted with `to_bytes`/`from_bytes`), `quantize`, `split`, `code_factors`, `pack_codes`/`unpack_codes`, the `Query` estimators and refinement, `query_terms`/`estimate_from_ip`, and FastScan blocks with `pack_fastscan_blocks`/`Query.estimate_fastscan`. It is provisional in the sense of PEP 411: any minor release may change it, with a note here. Import it explicitly with `from rabitqlib import quant`. See [Python building blocks](../rabitq/quant_python.md). +- Its byte formats (packed codes, factor records, FastScan blocks and rotator blobs) are versioned from the start and frozen by golden files; bytes of an unknown version raise `quant.FormatError`, a subclass of `ValueError`. + +## C++ changes + +- New provisional header `rabitqlib/core/provisional/quantize.hpp` with `rabitqlib::quant::provisional::combined_code` and `code_factors`, which back `rabitqlib.quant`. It is not part of the stable v1 API and may change in any minor release. diff --git a/docs/formats.md b/docs/formats.md new file mode 100644 index 00000000..7aa60efb --- /dev/null +++ b/docs/formats.md @@ -0,0 +1,94 @@ +# Byte formats + +This is the specification of the byte formats `rabitqlib.quant` reads and writes. It is shared with the planned index format revision, which will add its own blob kinds. While `rabitqlib.quant` is provisional, a release may stop reading an old format version; reading it then raises `quant.FormatError` instead of returning wrong results. A promise that every released version stays readable comes with stabilisation. + +Every format has a version. A change to any layout is a new version, never an edit of an existing one. The golden files in `tests/python/fixtures/quant/` freeze each version and are never regenerated. + +All multi-byte fields are little-endian. + +## Raw arrays + +Raw arrays have no header: their layout is fixed by the version below, which the module exports as a constant. A caller who stores raw arrays records the constant next to them. + +| Constant | Version | Format | +| --- | --- | --- | +| `PACKED_CODES_VERSION` | 1 | `pack_codes` output: one row of `d * nbits / 8` bytes per vector, for `nbits` 1–8 and `d` a positive multiple of 32 | +| `FACTORS_VERSION` | 1 | `FACTORS_DTYPE`, `ESTIMATE_DTYPE` and `QUERY_TERMS_DTYPE`: three `` layout. Vectors `32b` to `32b + 31` go to block `b`; in the last block, lanes past `n` have zero codes and zero factors. + +| Offset | Size | Content | +| --- | --- | --- | +| 0 | `d * 4` | Sign codes, `d / 8` groups of 32 bytes | +| `d * 4` | 128 | `add` of lanes 0–31, float32 | +| `d * 4 + 128` | 128 | `rescale` of lanes 0–31, float32 | +| `d * 4 + 256` | 128 | `error` of lanes 0–31, float32 | + +Group `i` covers coordinates `8i` to `8i + 7`. Let `a[v]` be byte `i` of vector `v`'s bitset with its bit order reversed, so coordinate `8i` is in bit 7. With `perm = [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15]`, for `j` in 0–15, `p = perm[j]`: + +- byte `j` is `(a[p] >> 4) | (a[p + 16] & 0xF0)`: coordinates `8i` to `8i + 3` of vectors `p` (low nibble) and `p + 16` (high nibble); +- byte `j + 16` is `(a[p] & 0x0F) | (a[p + 16] << 4)`: coordinates `8i + 4` to `8i + 7` of the same vectors. + +This is the order `fastscan::pack_codes` writes. The high-accuracy lookup table that scores the blocks is internal and not part of the format. + +## Blobs + +A blob is a self-describing byte string, such as `Rotator.to_bytes()`. It starts with a 40-byte common header, followed by the fields of its kind and then its body. + +### Common header + +| Offset | Size | Type | Field | Rule | +| --- | --- | --- | --- | --- | +| 0 | 8 | ASCII | `magic` = `"RABQBLOB"` | Anything else raises `FormatError` | +| 8 | 4 | u32 | `kind` | 1 = rotator; 2 = codec (reserved); 3 and up reserved for index formats. A kind the reader does not expect raises `FormatError` | +| 12 | 4 | u32 | `version` | Layout version of this kind; an unknown version raises `FormatError` | +| 16 | 4 | u32 | `header_bytes` | Common header plus kind fields, a multiple of 8; it must equal the size the kind and version define | +| 20 | 4 | u32 | `flags` | Kind-specific; unknown bits raise `FormatError` | +| 24 | 8 | u64 | `body_bytes` | The blob is exactly `header_bytes + body_bytes` long | +| 32 | 4 | u32 | `crc32` | CRC-32 as `zlib.crc32` computes it, over the whole blob with this field taken as zero; a mismatch raises `FormatError` | +| 36 | 4 | — | reserved | Writers write zero; readers reject anything else | + +Readers check, in order: the length covers the common header, the magic, the kind, the version, `header_bytes`, the total length, the checksum, and then the flags, reserved fields and kind fields. Sizes are validated before anything is allocated. The body starts at a multiple of 8, so an aligned buffer can be read in place. + +### Rotator, kind 1, version 1 + +`header_bytes` is 64. + +| Offset | Size | Type | Field | Rule | +| --- | --- | --- | --- | --- | +| 40 | 4 | u32 | `method` | 1 = FHT/Kac, 2 = matrix | +| 44 | 4 | u32 | `scalar` | 1 = float32. 2 = float64 is reserved and rejected in version 1 | +| 48 | 4 | u32 | `dim` | 1–65536 for matrix; 64–65536 for FHT/Kac | +| 52 | 4 | u32 | `padded_dim` | A multiple of 32 in `[dim, 65536]` | +| 56 | 4 | u32 | `seed` | Informational. Valid only when flag bit 0 (`seed_known`) is set; zero otherwise | +| 60 | 4 | — | reserved | Writers write zero; readers reject anything else | + +Flag bit 0, `seed_known`, is the only defined flag. + +The body is what `Rotator::save()` writes, and `body_bytes` must equal its size for the stored parameters: + +| Method | Body | `body_bytes` | +| --- | --- | --- | +| FHT/Kac | The sign-flip bytes | `padded_dim / 2` | +| Matrix | The rotation matrix as row-major float32 | `4 * dim * padded_dim` | + +The body defines the rotation; `seed` only records where it came from. A seeded FHT/Kac rotation is identical on every platform, while a matrix drawn from a seed is reproducible only with the same standard library and build, which is why the matrix itself is stored. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index f80cb19a..ec9001d2 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -14,6 +14,7 @@ nav: - C++ Integration: cpp_integration.md # Group by minor version: nested groups stay collapsed except the current page's. - Release notes: + - Unreleased: releases/unreleased.md - "0.8": - "0.8.2": releases/0.8.2.md - "0.8.1": releases/0.8.1.md @@ -29,6 +30,7 @@ nav: - Rotation: rabitq/rotator.md - Quantization: rabitq/quantizer.md - Stable core API: rabitq/core_api.md + - Python building blocks: rabitq/quant_python.md - Modular migration: modular_migration.md - Estimation: rabitq/estimator.md - Reranking: rabitq/reranking.md diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index efc0374b..8c145b40 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -2,7 +2,6 @@ #include #include -#include #include #include #include @@ -69,7 +68,14 @@ inline CodeFactorTerms code_factor_terms( const T l2_norm = std::sqrt(l2_sqr); const T ip_resi_xucb = dot_product(residual, centered, dim); const T ip_cent_xucb = dot_product(centroid, centered, dim); - assert(ip_resi_xucb > 0); + // The library's own codes match the residual's signs, so every term is + // nonnegative; only a caller-supplied code can make the divisor nonpositive. A NaN + // passes: the per-vector encoders do not check that inputs are finite. + if (ip_resi_xucb <= 0) { + throw std::invalid_argument( + "code does not quantize residual: must be positive" + ); + } // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. const T normalized_error = diff --git a/include/rabitqlib/core/encoder.hpp b/include/rabitqlib/core/encoder.hpp index aa042052..2009aa8f 100644 --- a/include/rabitqlib/core/encoder.hpp +++ b/include/rabitqlib/core/encoder.hpp @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -11,6 +12,7 @@ #include "rabitqlib/core/detail/rabitq_impl.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/types.hpp" // IWYU pragma: export +#include "rabitqlib/utils/space.hpp" namespace rabitqlib::quant { @@ -311,6 +313,50 @@ inline T full_est_dist_from_ip( return f_add + g_add + (f_rescale * (ip + g_k1xsumq * static_cast((1 << bits) - 1))); } + +// query_corrections and unpacked_ip back the stable quant::v1 prepare_query and compute +// as well as rabitqlib.quant. A change to either changes v1 results: validate v1 +// output against a saved pre-change build ("Stable core API changes" in +// docs/development.md), not only the provisional layer's tests. + +// The query terms of full_est_dist_from_ip for a rotated query and a centroid (nullptr +// means zero): g_add is |q - c|^2 for L2 and - for IP; g_k1xsumq is +// -0.5 * sum(q). Callers check that both are finite. +struct QueryCorrections { + float g_add; + float g_k1xsumq; +}; + +inline QueryCorrections query_corrections( + const float* query, const float* centroid, size_t dim, MetricType metric +) { + QueryCorrections corrections{}; + corrections.g_k1xsumq = std::accumulate(query, query + dim, 0.0F) * -0.5F; + corrections.g_add = + metric == METRIC_L2 + ? (centroid == nullptr ? l2norm_sqr(query, dim) + : euclidean_sqr(query, centroid, dim)) + : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, dim)); + return corrections; +} + +// for unpacked codes, summed in order: the reference inner product of +// quant::v1::compute. code_at(i) returns the i-th code value. Returns dim, or the index +// of the first code above maximum, in which case ip is incomplete. +template +inline size_t unpacked_ip( + const float* query, size_t dim, uint32_t maximum, CodeAt code_at, float& ip +) { + ip = 0; + for (size_t i = 0; i < dim; ++i) { + const uint32_t value = code_at(i); + if (value > maximum) { + return i; + } + ip += query[i] * static_cast(value); + } + return dim; +} } // namespace detail template diff --git a/include/rabitqlib/core/provisional/quantize.hpp b/include/rabitqlib/core/provisional/quantize.hpp new file mode 100644 index 00000000..d7c3933d --- /dev/null +++ b/include/rabitqlib/core/provisional/quantize.hpp @@ -0,0 +1,63 @@ +#pragma once + +#include +#include + +#include "rabitqlib/core/api.hpp" // v1::Factors +#include "rabitqlib/core/metric.hpp" // IWYU pragma: export + +namespace rabitqlib::quant::provisional { + +// Provisional quantization building blocks for rabitqlib.quant. Any minor release may +// change or remove them; they are not part of the stable quant::v1 API. Vectors are in +// the rotated, padded working space. Inputs are borrowed only for the call, outputs are +// written only after every check passes, and the calls allocate temporaries. If such an +// allocation fails, combined_code's code contents are unspecified. + +/** + * Unsigned total_bits-bit codes of one residual (vector minus centroid), the codes + * quant::v1::encode and the indexes store: the sign bit followed by the extra bits. + * t_const <= 0 runs the exact search; t_const > 0 is a fast-mode scale (faster_config). + * + * Requires dim to be a positive multiple of 32, residual to hold dim finite values, + * total_bits in [1, 8] for uint8_t codes or [1, 9] for uint16_t codes, and code to + * hold dim values that do not overlap residual. + */ +void combined_code( + const float* residual, size_t dim, size_t total_bits, uint8_t* code, double t_const = -1 +); +void combined_code( + const float* residual, + size_t dim, + size_t total_bits, + uint16_t* code, + double t_const = -1 +); + +/** + * Factors of a total_bits-bit code for residual. Pass the full code with its width for + * the full-code factors, or a prefix (code >> ex_bits) with the prefix width for the + * prefix factors. centroid == nullptr means zero; it only affects add. + * + * Requires dim to be a positive multiple of 32, residual and centroid to hold dim finite + * values, every code value to be below 2^total_bits, and the code to quantize residual: + * must be positive unless residual is zero. + */ +v1::Factors code_factors( + const float* residual, + const float* centroid, + size_t dim, + const uint8_t* code, + size_t total_bits, + MetricType metric +); +v1::Factors code_factors( + const float* residual, + const float* centroid, + size_t dim, + const uint16_t* code, + size_t total_bits, + MetricType metric +); + +} // namespace rabitqlib::quant::provisional diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index e77b7ec0..9bf0cea9 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -17,6 +17,7 @@ pybind11_add_module(_rabitqlib ivf_bindings.cpp symqg_bindings.cpp kmeans_bindings.cpp + quant_bindings.cpp ) add_library(rabitqlib::python ALIAS _rabitqlib) @@ -27,6 +28,13 @@ target_include_directories(_rabitqlib PRIVATE rabitq_configure_target(_rabitqlib) # Release preparation keeps this in sync with the Python package version. target_compile_definitions(_rabitqlib PRIVATE RABITQ_VERSION="${PROJECT_VERSION}") +# rabitqlib.quant results must not depend on the thread count or call site. GCC +# contracts a * b + c into FMA per compiled loop (the default -ffp-contract=fast on +# aarch64 or with -mfma), so the serial and OpenMP loops could round the inlined +# estimator differently. MSVC does not contract by default. +set_source_files_properties(quant_bindings.cpp PROPERTIES COMPILE_OPTIONS + "$<$:-ffp-contract=off>" +) # 4. Imported targets provide the matching Python, NumPy, and OpenMP settings. target_link_libraries(_rabitqlib PRIVATE @@ -60,7 +68,7 @@ endif() # 5. Map files to a 'rabitqlib' directory # inside the target wheel, completely bypassing your local layout. install(TARGETS _rabitqlib DESTINATION rabitqlib) -install(FILES __init__.py __init__.pyi py.typed DESTINATION rabitqlib) +install(FILES __init__.py __init__.pyi py.typed quant.py quant.pyi DESTINATION rabitqlib) if(RABITQ_BUILD_BENCHMARKS) # Benchmark the production conversion helper without importing or replacing diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index c80b0a63..f88da6fc 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -7,12 +7,14 @@ #include #include #include +#include #include #include #include #include #include #include +#include #include #include "rabitqlib/core/metric.hpp" @@ -32,6 +34,7 @@ class IndexAccess { private: size_t readers_ = 0; bool writing_ = false; + const char* busy_message_ = "Index is busy: conflicting operation in progress"; public: class Guard { @@ -42,9 +45,7 @@ class IndexAccess { public: Guard(IndexAccess& access, bool writing) : access_(access), writing_(writing) { if (access.writing_ || (writing && access.readers_ != 0)) { - throw std::runtime_error( - "Index is busy: conflicting operation in progress" - ); + throw std::runtime_error(access.busy_message_); } if (writing) { access.writing_ = true; @@ -63,6 +64,10 @@ class IndexAccess { } }; + IndexAccess() = default; + // busy_message must outlive the object, as a string literal does. + explicit IndexAccess(const char* busy_message) : busy_message_(busy_message) {} + [[nodiscard]] Guard read() { return Guard(*this, false); } [[nodiscard]] Guard write() { return Guard(*this, true); } }; @@ -127,27 +132,35 @@ inline void ensure_dimension(const py::array& array, size_t expected, const char } } -// A seed is an integer in [0, max_seed]; bool is rejected like other integer options. -inline uint64_t checked_seed(py::handle value, uint64_t max_seed) { - const std::string range = - "seed must be an integer in [0, " + std::to_string(max_seed) + "]"; +// An integer option in [0, max_value]. bool is rejected; a non-integer raises TypeError +// and an out-of-range value ValueError, both with message. +inline uint64_t checked_integer( + py::handle value, uint64_t max_value, const std::string& message +) { if (PyBool_Check(value.ptr())) { - throw py::type_error(range); + throw py::type_error(message); } auto integer = py::reinterpret_steal(PyNumber_Index(value.ptr())); if (!integer) { PyErr_Clear(); - throw py::type_error(range); + throw py::type_error(message); } - const unsigned long long seed = PyLong_AsUnsignedLongLong(integer.ptr()); + const unsigned long long result = PyLong_AsUnsignedLongLong(integer.ptr()); if (PyErr_Occurred() != nullptr) { PyErr_Clear(); - throw std::invalid_argument(range); + throw std::invalid_argument(message); } - if (seed > max_seed) { - throw std::invalid_argument(range); + if (result > max_value) { + throw std::invalid_argument(message); } - return seed; + return result; +} + +// A seed is an integer in [0, max_seed]. +inline uint64_t checked_seed(py::handle value, uint64_t max_seed) { + return checked_integer( + value, max_seed, "seed must be an integer in [0, " + std::to_string(max_seed) + "]" + ); } // Call only after rank validation. @@ -251,6 +264,95 @@ inline py::array_t ensure_1d_array return ensure_array_rank(value, name, 1); } +// A (rows, cols) view whose rows are contiguous but whose row stride is free, such as +// a column slice of a record buffer. owner keeps the NumPy storage alive across GIL +// release. Element alignment is guaranteed, so typed reads are valid. +template +struct RowView { + py::array owner; + char* data = nullptr; + py::ssize_t row_stride = 0; + size_t rows = 0; + size_t cols = 0; + + [[nodiscard]] T* row(size_t i) const { + return reinterpret_cast(data + (static_cast(i) * row_stride)); + } +}; + +template +bool rows_are_usable(const py::array& array) { + const bool contiguous_rows = array.shape(1) <= 1 || array.strides(1) == sizeof(T); + const bool aligned = reinterpret_cast(array.data()) % alignof(T) == 0 && + array.strides(0) % static_cast(alignof(T)) == 0; + return contiguous_rows && aligned; +} + +template +RowView make_row_view(py::array array) { + RowView view; + view.rows = static_cast(array.shape(0)); + view.cols = static_cast(array.shape(1)); + view.row_stride = array.strides(0); + view.data = static_cast(const_cast(array.data())); + view.owner = std::move(array); + return view; +} + +// A 2D input of exactly dtype T, never cast. Rows with a free stride are read in place; +// rows that are not contiguous, or misaligned storage, are copied. +template +RowView exact_dtype_rows(py::handle value, const char* name) { + auto array = py::array::ensure(value); + if (!array) { + throw std::invalid_argument( + std::string(name) + " must be convertible to a NumPy array; got " + + array_input_type(value) + ); + } + if (!py::isinstance>(array)) { + throw std::invalid_argument( + std::string(name) + " must have dtype " + + py::str(py::dtype::of()).cast() + "; got dtype " + + py::str(array.dtype()).cast() + ); + } + if (array.ndim() != 2) { + throw std::invalid_argument( + std::string(name) + " must be a 2D NumPy array; got shape " + array_shape(array) + ); + } + if (!rows_are_usable(array)) { + array = py::array(array.attr("copy")(py::arg("order") = "C")); + } + return make_row_view(std::move(array)); +} + +// A caller-supplied output: exactly dtype T and shape (rows, cols), writable, with +// contiguous, aligned and non-overlapping rows. It is never converted or copied. name +// is how messages refer to it. +template +RowView checked_output_rows( + py::handle value, size_t rows, size_t cols, const std::string& name = "out" +) { + const std::string expected = name + " must be a writable NumPy array of dtype " + + py::str(py::dtype::of()).cast() + + " and shape (" + std::to_string(rows) + ", " + + std::to_string(cols) + ") with contiguous, aligned rows"; + if (!py::isinstance>(value)) { + throw std::invalid_argument(expected + "; got " + array_input_type(value)); + } + auto array = py::reinterpret_borrow(value); + const auto row_bytes = static_cast(cols * sizeof(T)); + if (array.ndim() != 2 || static_cast(array.shape(0)) != rows || + static_cast(array.shape(1)) != cols || !array.writeable() || + !rows_are_usable(array) || + (rows > 1 && std::abs(array.strides(0)) < row_bytes)) { + throw std::invalid_argument(expected + "; got shape " + array_shape(array)); + } + return make_row_view(std::move(array)); +} + struct SearchResults { py::array_t ids; py::array_t distances; diff --git a/python_bindings/quant.py b/python_bindings/quant.py new file mode 100644 index 00000000..c31cac2b --- /dev/null +++ b/python_bindings/quant.py @@ -0,0 +1,77 @@ +"""Core quantization building blocks: rotation, quantization, factors, packing and estimation. + +This module is provisional in the sense of PEP 411: any minor release may change +or remove its API, with a changelog note, until it is declared stable. Byte +formats are versioned; bytes of an unknown version raise ``FormatError``. + +The layer calls work on rotated, padded vectors: arrays of shape ``(n, d)`` +where ``d`` is a positive multiple of 32. ``Rotator`` maps original vectors +into that space, or callers bring their own rotation. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import quant +>>> rot = quant.Rotator(128, seed=0) +>>> residuals = rot.rotate(np.ones((4, 128), dtype=np.float32)) +>>> codes = quant.quantize(residuals, nbits=4) +>>> factors = quant.code_factors(residuals, codes, nbits=4, metric="l2") +>>> quant.pack_codes(codes, nbits=4).shape +(4, 64) +>>> q = quant.Query(residuals[0], metric="l2") +>>> q.estimate_codes(codes, factors, nbits=4).dtype == quant.ESTIMATE_DTYPE +True +""" + +from ._rabitqlib._quant import ( + ESTIMATE_DTYPE, + FACTORS_DTYPE, + QUERY_TERMS_DTYPE, + FormatError, + Query, + Rotator, + code_factors, + estimate_from_ip, + pack_codes, + pack_fastscan_blocks, + padded_dim, + quantize, + query_terms, + split, + unpack_codes, +) + +PACKED_CODES_VERSION = 1 +"""Version of the pack_codes layout, including the 1-bit sign bitset.""" + +FACTORS_VERSION = 1 +"""Version of FACTORS_DTYPE, ESTIMATE_DTYPE and QUERY_TERMS_DTYPE.""" + +FASTSCAN_BLOCKS_VERSION = 1 +"""Version of the pack_fastscan_blocks layout.""" + +# Classes report this module, so pickle and help() find them here. +for _public in (FormatError, Query, Rotator): + _public.__module__ = __name__ +del _public + +__all__ = [ + "ESTIMATE_DTYPE", + "FACTORS_DTYPE", + "FACTORS_VERSION", + "FASTSCAN_BLOCKS_VERSION", + "PACKED_CODES_VERSION", + "QUERY_TERMS_DTYPE", + "FormatError", + "Query", + "Rotator", + "code_factors", + "estimate_from_ip", + "pack_codes", + "pack_fastscan_blocks", + "padded_dim", + "quantize", + "query_terms", + "split", + "unpack_codes", +] diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi new file mode 100644 index 00000000..9cdbf2f6 --- /dev/null +++ b/python_bindings/quant.pyi @@ -0,0 +1,271 @@ +"""Core quantization building blocks: rotation, quantization, factors, packing and estimation. + +This module is provisional in the sense of PEP 411: any minor release may change +or remove its API, with a changelog note, until it is declared stable. +""" + +from typing import Any, Literal, overload + +import numpy as np +import numpy.typing as npt + +__all__ = [ + "ESTIMATE_DTYPE", + "FACTORS_DTYPE", + "FACTORS_VERSION", + "FASTSCAN_BLOCKS_VERSION", + "PACKED_CODES_VERSION", + "QUERY_TERMS_DTYPE", + "FormatError", + "Query", + "Rotator", + "code_factors", + "estimate_from_ip", + "pack_codes", + "pack_fastscan_blocks", + "padded_dim", + "quantize", + "query_terms", + "split", + "unpack_codes", +] + +FACTORS_DTYPE: np.dtype[np.void] +ESTIMATE_DTYPE: np.dtype[np.void] +QUERY_TERMS_DTYPE: np.dtype[np.void] +PACKED_CODES_VERSION: int +FACTORS_VERSION: int +FASTSCAN_BLOCKS_VERSION: int + +class FormatError(ValueError): + """Corrupt or incompatible serialized bytes. A subclass of ValueError.""" + +class Rotator: + """A seeded rotation from the original dimension to the padded working space.""" + + def __init__( + self, + dim: int, + *, + method: Literal["fht_kac", "fht", "matrix"] = "fht_kac", + padded_dim: int | None = None, + seed: int | None = None, + ) -> None: ... + @property + def dim(self) -> int: + """Original vector dimension.""" + @property + def padded_dim(self) -> int: + """Working dimension of rotated vectors.""" + @property + def method(self) -> Literal["fht_kac", "matrix"]: + """Canonical method name, 'fht_kac' or 'matrix'.""" + @property + def seed(self) -> int | None: + """The seed the rotation was drawn from, or None if unknown.""" + def rotate( + self, + x: npt.ArrayLike, + *, + num_threads: int = 0, + out: npt.NDArray[np.float32] | None = None, + ) -> npt.NDArray[np.float32]: + """Rotate vectors into the padded working space.""" + def to_bytes(self) -> bytes: + """Serialize the rotation as a versioned rotator blob (docs/formats.md).""" + @staticmethod + def from_bytes( + data: bytes | bytearray | memoryview | npt.NDArray[Any], + ) -> Rotator: + """Load a rotator blob written by to_bytes.""" + def __copy__(self) -> Rotator: + """Return self; rotators are immutable.""" + def __deepcopy__(self, memo: Any) -> Rotator: + """Return self; rotators are immutable.""" + +_Metric = Literal["l2", "ip", "innerproduct"] + +class Query: + """A rotated, padded query for scoring codes in the working space.""" + + def __init__( + self, + rotated_query: npt.ArrayLike, + *, + metric: _Metric, + centroid: npt.ArrayLike | None = None, + ) -> None: ... + @property + def dim(self) -> int: + """Working dimension d of the query.""" + @property + def metric(self) -> Literal["l2", "ip"]: + """The metric, 'l2' or 'ip'.""" + def set_centroid(self, centroid: npt.ArrayLike | None) -> None: + """Switch to the factors of another centroid; None is zero.""" + def estimate_bitset( + self, + bitsets: npt.NDArray[np.uint8], + factors: npt.NDArray[np.void], + *, + indices: npt.ArrayLike | None = None, + num_threads: int = 1, + out: npt.NDArray[np.void] | None = None, + ) -> npt.NDArray[np.void]: + """Estimate distances from packed 1-bit codes with the float query.""" + def estimate_codes( + self, + codes: npt.NDArray[np.uint8] | npt.NDArray[np.uint16], + factors: npt.NDArray[np.void], + *, + nbits: int, + indices: npt.ArrayLike | None = None, + num_threads: int = 1, + out: npt.NDArray[np.void] | None = None, + ) -> npt.NDArray[np.void]: + """Estimate distances from unpacked codes of any width.""" + def refine( + self, + ex: npt.NDArray[np.uint8], + factors: npt.NDArray[np.void], + first: npt.NDArray[np.void], + *, + ex_bits: int, + prefix_bits: int = 1, + indices: npt.ArrayLike | None = None, + num_threads: int = 1, + out: npt.NDArray[np.float32] | None = None, + ) -> npt.NDArray[np.float32]: + """Refine prefix estimates with the unpacked extra bits.""" + def refine_packed( + self, + packed: npt.NDArray[np.uint8], + factors: npt.NDArray[np.void], + first: npt.NDArray[np.void], + *, + ex_bits: int, + prefix_bits: int = 1, + indices: npt.ArrayLike | None = None, + num_threads: int = 1, + out: npt.NDArray[np.float32] | None = None, + ) -> npt.NDArray[np.float32]: + """Refine prefix estimates with packed extra bits.""" + def estimate_fastscan( + self, + blocks: npt.NDArray[np.uint8], + *, + num_points: int, + num_threads: int = 1, + out: npt.NDArray[np.void] | None = None, + ) -> npt.NDArray[np.void]: + """Estimate distances for FastScan blocks with a 4-bit query table.""" + def copy(self) -> Query: + """Return an independent copy of the query.""" + def __copy__(self) -> Query: + """Return an independent copy of the query.""" + def __deepcopy__(self, memo: Any) -> Query: + """Return an independent copy of the query.""" + +def padded_dim(dim: int) -> int: + """Return the working dimension for dim: dim rounded up to a multiple of 32.""" + +@overload +def quantize( + residuals: npt.ArrayLike, + *, + nbits: Literal[1, 2, 3, 4, 5, 6, 7, 8], + fast_quantization: bool = False, + num_threads: int = 0, + out: npt.NDArray[np.uint8] | None = None, +) -> npt.NDArray[np.uint8]: + """Quantize rotated residuals into unsigned nbits-bit codes.""" + +@overload +def quantize( + residuals: npt.ArrayLike, + *, + nbits: Literal[9], + fast_quantization: bool = False, + num_threads: int = 0, + out: npt.NDArray[np.uint16] | None = None, +) -> npt.NDArray[np.uint16]: + """Quantize rotated residuals into unsigned nbits-bit codes.""" + +@overload +def quantize( + residuals: npt.ArrayLike, + *, + nbits: int, + fast_quantization: bool = False, + num_threads: int = 0, + out: npt.NDArray[np.uint8] | npt.NDArray[np.uint16] | None = None, +) -> npt.NDArray[np.uint8] | npt.NDArray[np.uint16]: + """Quantize rotated residuals into unsigned nbits-bit codes.""" + +def split( + codes: npt.NDArray[np.uint8] | npt.NDArray[np.uint16], + *, + nbits: int, + prefix_bits: int, + num_threads: int = 0, + out: tuple[npt.NDArray[np.uint8], npt.NDArray[np.uint8]] | None = None, +) -> tuple[npt.NDArray[np.uint8], npt.NDArray[np.uint8]]: + """Split nbits-bit codes into a prefix and extra bits.""" + +def code_factors( + residuals: npt.ArrayLike, + codes: npt.NDArray[np.uint8] | npt.NDArray[np.uint16], + *, + nbits: int, + metric: _Metric, + centroid: npt.ArrayLike | None = None, + num_threads: int = 0, + out: npt.NDArray[np.void] | None = None, +) -> npt.NDArray[np.void]: + """Return the factors of nbits-bit codes for their residuals.""" + +def pack_codes( + codes: npt.NDArray[np.uint8], + *, + nbits: int, + num_threads: int = 0, + out: npt.NDArray[np.uint8] | None = None, +) -> npt.NDArray[np.uint8]: + """Pack nbits-bit codes into the packed-codes format (version 1).""" + +def unpack_codes( + packed: npt.NDArray[np.uint8], + *, + nbits: int, + num_threads: int = 0, + out: npt.NDArray[np.uint8] | None = None, +) -> npt.NDArray[np.uint8]: + """Unpack codes written by pack_codes; the inverse is lossless.""" + +def pack_fastscan_blocks( + bitsets: npt.NDArray[np.uint8], + factors: npt.NDArray[np.void], + *, + num_threads: int = 0, + out: npt.NDArray[np.uint8] | None = None, +) -> npt.NDArray[np.uint8]: + """Pack bitsets and their factors into FastScan blocks (version 1).""" + +def query_terms( + rotated_query: npt.ArrayLike, + *, + metric: _Metric, + centroid: npt.ArrayLike | None = None, +) -> np.void: + """Return the query-side terms of the reference estimator formula.""" + +def estimate_from_ip( + ip: npt.ArrayLike, + factors: npt.NDArray[np.void], + terms: np.void | npt.NDArray[np.void], + *, + nbits: int, + num_threads: int = 1, + out: npt.NDArray[np.void] | None = None, +) -> npt.NDArray[np.void]: + """Apply the reference estimator formula to the caller's inner products.""" diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp new file mode 100644 index 00000000..a713042d --- /dev/null +++ b/python_bindings/quant_bindings.cpp @@ -0,0 +1,2460 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include // IWYU pragma: keep (std::optional seed) + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "bindings_common.hpp" +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/detail/fast_scale_table.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/core/provisional/quantize.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/core/types.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/utils/input_checks.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace py = pybind11; + +namespace rabitqlib::python_bindings { +namespace { + +// Raised for corrupt or incompatible serialized bytes; a ValueError in Python. +struct FormatError : std::invalid_argument { + using std::invalid_argument::invalid_argument; +}; + +// The blob format and both rotation methods share the FHT/Kac limit. +constexpr size_t kMaxRotatorDim = rotator_impl::FhtKacRotator::kMaxDim; +constexpr size_t kMaxTotalBits = 9; + +// --- Records and structured dtypes, byte-compatible with the planned C structs ----- + +// An estimate record, the layout of ESTIMATE_DTYPE; factor records are v1::Factors. +struct Estimate { + float distance; + float error; + float refine_input; +}; + +static_assert(sizeof(quant::v1::Factors) == 3 * sizeof(float)); +static_assert(sizeof(Estimate) == 3 * sizeof(float)); + +py::dtype float_record_dtype(const std::array& names) { + py::list name_list; + py::list formats; + py::list offsets; + for (size_t i = 0; i < names.size(); ++i) { + name_list.append(names[i]); + formats.append(" +py::dtype stored_dtype(py::gil_safe_call_once_and_store& storage, Names names) { + return storage.call_once_and_store_result( + [&] { return float_record_dtype(names()); } + ).get_stored(); +} + +py::dtype factors_dtype() { + PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; + return stored_dtype(storage, [] { + return std::array{"add", "rescale", "error"}; + }); +} + +py::dtype estimate_dtype() { + PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; + return stored_dtype(storage, [] { + return std::array{"distance", "error", "refine_input"}; + }); +} + +py::dtype query_terms_dtype() { + PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store storage; + return stored_dtype(storage, [] { + return std::array{"g_add", "g_error", "g_k1xsumq"}; + }); +} + +// --- Argument checks ------------------------------------------------------------ + +void check_nbits(size_t nbits, size_t low, size_t high, const char* name = "nbits") { + if (nbits < low || nbits > high) { + throw std::invalid_argument( + std::string(name) + " must be in [" + std::to_string(low) + ", " + + std::to_string(high) + "]; got " + std::to_string(nbits) + ); + } +} + +// Layer calls work on rotated, padded vectors: d is a positive multiple of 32. +size_t checked_working_dim(size_t d, const py::array& array, const char* name) { + if (d == 0 || d % 32 != 0) { + throw std::invalid_argument( + std::string(name) + + " must have a last dimension that is a positive multiple of 32; got shape " + + array_shape(array) + ); + } + return d; +} + +// The rows a call reads: all of them, or the caller's indices, checked against n. +struct Selection { + py::object owner; // Keeps the converted indices alive; none for all rows. + const int64_t* indices = nullptr; + size_t count = 0; + + [[nodiscard]] size_t row(size_t k) const { + return indices == nullptr ? k : static_cast(indices[k]); + } +}; + +Selection all_rows(size_t rows) { + Selection selected; + selected.count = rows; + return selected; +} + +Selection selection(const py::object& indices, size_t rows) { + if (indices.is_none()) { + return all_rows(rows); + } + Selection selected; + auto array = ensure_1d_integer_array(indices, "indices"); + selected.indices = array.data(); + selected.count = static_cast(array.shape(0)); + selected.owner = std::move(array); + for (size_t k = 0; k < selected.count; ++k) { + const int64_t i = selected.indices[k]; + if (i < 0 || static_cast(i) >= rows) { + throw std::invalid_argument( + "indices must be in [0, " + std::to_string(rows) + "); got " + + std::to_string(i) + " at position " + std::to_string(k) + ); + } + } + return selected; +} + +template +[[noreturn]] void reject_code_value( + const RowView& codes, size_t i, size_t j, size_t nbits, const char* name +) { + throw std::invalid_argument( + std::string(name) + " values must be below 2**" + std::to_string(nbits) + " = " + + std::to_string(size_t{1} << nbits) + "; got " + std::to_string(codes.row(i)[j]) + + " at (" + std::to_string(i) + ", " + std::to_string(j) + ")" + ); +} + +template +void check_code_row(const RowView& codes, size_t i, size_t nbits, const char* name) { + const auto maximum = static_cast((1U << nbits) - 1); + const TP* row = codes.row(i); + // A vectorizable maximum first; the position is searched only for a rejected row. + TP largest = 0; + for (size_t j = 0; j < codes.cols; ++j) { + largest = std::max(largest, row[j]); + } + if (largest > maximum) { + reject_code_value( + codes, + i, + static_cast( + std::find_if( + row, row + codes.cols, [&](TP value) { return value > maximum; } + ) - + row + ), + nbits, + name + ); + } +} + +// of row i for unpacked codes already checked against their bit width. +template +float unpacked_row_ip(const float* query, const RowView& codes, size_t i) { + const TP* row = codes.row(i); + float ip = 0; + (void)quant::detail::unpacked_ip( + query, + codes.cols, + std::numeric_limits::max(), + [row](size_t at) { return row[at]; }, + ip + ); + return ip; +} + +void check_finite_factor(const quant::v1::Factors& f) { + if (!std::isfinite(f.add) || !std::isfinite(f.rescale) || !std::isfinite(f.error)) { + throw std::invalid_argument("factors must contain only finite values"); + } +} + +// A float32 vector copied into owned storage and checked like index inputs. Its length +// must be expected_d, or with expected_d == 0 a positive multiple of 32. +std::vector owned_vector(py::handle value, const char* name, size_t expected_d) { + auto array = ensure_1d_array(value, name); + const auto d = static_cast(array.shape(0)); + if (expected_d == 0 && (d == 0 || d % 32 != 0)) { + throw std::invalid_argument( + std::string(name) + + " must have a length that is a positive multiple of 32; got shape " + + array_shape(array) + ); + } + if (expected_d != 0 && d != expected_d) { + throw std::invalid_argument( + std::string(name) + " must have shape (" + std::to_string(expected_d) + + ",); got shape " + array_shape(array) + ); + } + std::vector vector(d); + std::memcpy(vector.data(), array.data(), d * sizeof(float)); + check_coordinates(vector.data(), d, d, name); + return vector; +} + +// A centroid in the working space, or an empty vector for none. +std::vector centroid_vector(const py::object& centroid, size_t d) { + return centroid.is_none() ? std::vector{} + : owned_vector(centroid, "centroid", d); +} + +// --- Outputs ---------------------------------------------------------------------- + +template +py::array_t new_rows(size_t rows, size_t cols) { + return py::array_t(std::vector{ + static_cast(rows), static_cast(cols)}); +} + +// An input that an out= must not overlap: its name and the array the call reads. +struct NamedInput { + const char* name; + py::handle array; // Null or None when the input is absent. +}; +using Inputs = std::initializer_list; + +// The bytes an array spans, from its lowest to past its highest element. +std::pair byte_extent(const py::array& array) { + if (array.size() == 0) { + return {0, 0}; + } + auto low = reinterpret_cast(array.data()); + auto high = low; + for (py::ssize_t axis = 0; axis < array.ndim(); ++axis) { + const py::ssize_t span = (array.shape(axis) - 1) * array.strides(axis); + if (span < 0) { + low -= static_cast(-span); + } else { + high += static_cast(span); + } + } + return {low, high + static_cast(array.itemsize())}; +} + +// Reject an out= that shares memory with an input, as the C++ API rejects overlapping +// buffers: rows written first would change rows read later. Exact aliasing is rejected +// too, so the rule does not depend on how a call orders its reads and writes. +void check_output_disjoint( + const py::array& out, Inputs inputs, const std::string& name = "out" +) { + const auto [out_low, out_high] = byte_extent(out); + if (out_low == out_high) { + return; + } + for (const NamedInput& input : inputs) { + if (!input.array || input.array.is_none()) { + continue; + } + const auto [low, high] = + byte_extent(py::reinterpret_borrow(input.array)); + if (low < out_high && out_low < high) { + throw std::invalid_argument(name + " must not overlap " + input.name); + } + } +} + +// The output rows: a new array, or the caller's out= after its exact checks. The +// view's owner is the array to return. +template +RowView output_rows( + const py::object& out, + size_t rows, + size_t cols, + Inputs inputs, + const std::string& name = "out" +) { + if (out.is_none()) { + return make_row_view(new_rows(rows, cols)); + } + auto view = checked_output_rows(out, rows, cols, name); + check_output_disjoint(view.owner, inputs, name); + return view; +} + +// A new 1-D array of dtype, or the caller's out= of exactly that dtype and shape. +// dtype_name is how messages name the dtype, such as "quant.ESTIMATE_DTYPE". +template +std::pair vector_output( + const py::object& out, + const py::dtype& dtype, + const char* dtype_name, + size_t rows, + Inputs inputs +) { + static_assert(alignof(T) == alignof(float)); + if (out.is_none()) { + py::array result(dtype, std::vector{static_cast(rows)}); + auto* data = static_cast(result.mutable_data()); + return {std::move(result), data}; + } + const std::string expected = std::string("out must be a writable, C-contiguous ") + + "NumPy array of dtype " + dtype_name + " and shape (" + + std::to_string(rows) + ",)"; + if (!py::isinstance(out)) { + throw std::invalid_argument(expected + "; got " + array_input_type(out)); + } + auto result = py::reinterpret_borrow(out); + if (!result.dtype().equal(dtype) || result.ndim() != 1 || + static_cast(result.shape(0)) != rows || !result.writeable() || + (rows > 1 && result.strides(0) != static_cast(sizeof(T))) || + reinterpret_cast(result.data()) % alignof(T) != 0) { + throw std::invalid_argument(expected + "; got shape " + array_shape(result)); + } + check_output_disjoint(result, inputs); + auto* data = static_cast(result.mutable_data()); + return {std::move(result), data}; +} + +std::pair estimate_output( + const py::object& out, size_t rows, Inputs inputs +) { + return vector_output( + out, estimate_dtype(), "quant.ESTIMATE_DTYPE", rows, inputs + ); +} + +// --- Record inputs ---------------------------------------------------------------- + +// A 1-D array of one of the float-record dtypes, read as contiguous aligned records. +// Strided or misaligned input is copied; other dtypes are rejected, never cast. +template +struct Records { + py::array owner; + const T* data = nullptr; + size_t rows = 0; + + [[nodiscard]] const T& operator[](size_t i) const { return data[i]; } +}; + +template +Records record_input( + py::handle value, const py::dtype& dtype, const char* dtype_name, const char* name +) { + const std::string expected = + std::string(name) + " must be a 1D NumPy array of dtype quant." + dtype_name; + auto array = py::array::ensure(value); + if (!array) { + throw std::invalid_argument(expected + "; got " + array_input_type(value)); + } + if (!array.dtype().equal(dtype)) { + throw std::invalid_argument( + expected + "; got dtype " + py::str(array.dtype()).cast() + ); + } + if (array.ndim() != 1) { + throw std::invalid_argument(expected + "; got shape " + array_shape(array)); + } + const auto rows = static_cast(array.shape(0)); + if ((rows > 1 && array.strides(0) != static_cast(sizeof(T))) || + reinterpret_cast(array.data()) % alignof(T) != 0) { + array = py::array(array.attr("copy")()); + } + Records records; + records.rows = rows; + records.data = static_cast(array.data()); + records.owner = std::move(array); + return records; +} + +// The factors argument: FACTORS_DTYPE records, one per row of the named input. +Records factors_input(py::handle value, size_t rows, const char* of) { + auto records = record_input( + value, factors_dtype(), "FACTORS_DTYPE", "factors" + ); + if (records.rows != rows) { + throw std::invalid_argument( + std::string("factors must have one record per row of ") + of + ", " + + std::to_string(rows) + "; got shape " + array_shape(records.owner) + ); + } + return records; +} + +// --- Parallel rows ---------------------------------------------------------------- + +// parallel_for hands out one index at a time from a shared counter, which costs more +// than a cheap row. Rows go out in ranges instead, about eight per worker and at least +// min_rows long: kCheapRows for rows of up to a few hundred nanoseconds, such as +// scoring, packing and checks, and kCostlyRows for rotate, quantize and code_factors, +// whose rows cost enough that a small batch should still use every worker. +constexpr size_t kCheapRows = 64; +constexpr size_t kCostlyRows = 1; + +// Runs task(begin, end) on consecutive ranges covering [0, rows). Call with the GIL +// released. +template +void parallel_rows(size_t rows, size_t num_threads, size_t min_rows, Task task) { + if (rows == 0) { + return; + } + const size_t workers = resolve_num_threads(num_threads); + const size_t task_rows = std::max(min_rows, div_round_up(rows, workers * 8)); + index_detail::parallel_for( + 0, div_round_up(rows, task_rows), workers, [&](size_t t, size_t /*worker*/) { + const size_t begin = t * task_rows; + task(begin, std::min(rows, begin + task_rows)); + } + ); +} + +// The exception of the lowest failing row of a parallel loop, so the error a call +// raises does not depend on the thread schedule. Rows after a recorded failure are +// skipped. +class FirstFailure { + public: + explicit FirstFailure(size_t rows) : row_(rows) {} + + // A stale value only delays the skip; record() decides under the lock. + [[nodiscard]] bool after(size_t row) const { + return row > row_.load(std::memory_order_relaxed); + } + + // Call from the catch block of the failing row. + void record(size_t row) { + const std::lock_guard lock(mutex_); + if (row < row_.load()) { + row_.store(row); + error_ = std::current_exception(); + } + } + + void rethrow() const { + if (error_) { + std::rethrow_exception(error_); + } + } + + private: + std::atomic row_; + std::mutex mutex_; + std::exception_ptr error_; +}; + +// Runs row(i) for every i in [0, rows) and rethrows the lowest failing row's exception. +// range(begin, end) first handles a whole range without per-row checks, so its loops +// can vectorize, and returns false if any row in it fails; only such a range runs +// through row(i), whose checks find that row. range must not throw. +template +void for_each_range( + size_t rows, size_t num_threads, size_t min_rows, Range range, Row row +) { + FirstFailure failure(rows); + parallel_rows(rows, num_threads, min_rows, [&](size_t begin, size_t end) { + if (failure.after(begin) || range(begin, end)) { + return; + } + for (size_t i = begin; i < end && !failure.after(i); ++i) { + try { + row(i); + } catch (...) { + failure.record(i); + return; + } + } + }); + failure.rethrow(); +} + +template +void for_each_row(size_t rows, size_t num_threads, size_t min_rows, Row row) { + for_each_range( + rows, + num_threads, + min_rows, + [](size_t /*begin*/, size_t /*end*/) { return false; }, + row + ); +} + +template +void check_code_values( + const RowView& codes, size_t nbits, const char* name, size_t num_threads +) { + for_each_row(codes.rows, num_threads, kCheapRows, [&](size_t i) { + check_code_row(codes, i, nbits, name); + }); +} + +// Whether every float of records [begin, end) is finite. +bool all_finite_records(const float* values, size_t begin, size_t end) { + return all_finite(values + begin, end - begin); +} + +template +bool all_finite_records(const T* records, size_t begin, size_t end) { + static_assert(sizeof(T) == 3 * sizeof(float)); + return all_finite(reinterpret_cast(records + begin), 3 * (end - begin)); +} + +// A vectorized test that the selected rows [begin, end) have finite factors, for every +// row selected; false means only that the rows must be checked one by one. +bool finite_factor_range( + const Records& factors, + const Selection& selected, + size_t begin, + size_t end +) { + return selected.indices == nullptr && all_finite_records(factors.data, begin, end); +} + +// quant::v1::compute rejects the distances that are not finite; so do the scoring calls. +float checked(float distance) { + if (!std::isfinite(distance)) { + throw std::overflow_error("Estimated distance is not representable in float32"); + } + return distance; +} + +Estimate checked(const Estimate& estimate) { + (void)checked(estimate.distance); + if (!std::isfinite(estimate.error)) { + throw std::overflow_error("Estimated error is not representable in float32"); + } + return estimate; +} + +// Writes value_at(k) to out[k] for k in [0, count), for rows already validated, and +// raises OverflowError for the lowest row whose value is not finite. Ranges are +// written without per-row checks and tested at once; only a failing range is +// recomputed row by row. out may then be partly written. +template +void write_rows(T* out, size_t count, size_t num_threads, ValueAt value_at) { + for_each_range( + count, + num_threads, + kCheapRows, + [&](size_t begin, size_t end) { + for (size_t k = begin; k < end; ++k) { + out[k] = value_at(k); + } + return all_finite_records(out, begin, end); + }, + [&](size_t k) { out[k] = checked(value_at(k)); } + ); +} + +// --- Layer 1 -------------------------------------------------------------------- + +size_t padded_dim(size_t dim) { + if (dim == 0) { + throw std::invalid_argument("dim must be positive"); + } + return rotator_impl::padding_requirement(dim, RotatorType::FhtKacRotator); +} + +py::array quantize( + py::handle residuals, + size_t nbits, + bool fast_quantization, + size_t num_threads, + const py::object& out +) { + check_nbits(nbits, 1, kMaxTotalBits); + auto input = ensure_2d_array(residuals, "residuals"); + const auto rows = static_cast(input.shape(0)); + const size_t d = + checked_working_dim(static_cast(input.shape(1)), input, "residuals"); + + double t_const = -1; + if (fast_quantization && nbits > 1) { + const std::optional frozen = + quant::rabitq_impl::ex_bits::frozen_const_scaling_factor(d, nbits - 1); + if (!frozen.has_value()) { + throw std::invalid_argument( + "fast_quantization requires a working dimension of at most " + + std::to_string(quant::rabitq_impl::ex_bits::kFastScaleMaxDim) + + " (the frozen scale table); got shape " + array_shape(input) + ); + } + t_const = *frozen; + } + + // The whole-array pass applies the index coordinate bound; combined_code's own + // per-row checks guard its public contract and cost O(d) against the code search. + const auto run = [&](const auto& view) { + { + py::gil_scoped_release release; + check_coordinates(input.data(), rows * d, d, "residuals"); + parallel_rows(rows, num_threads, kCostlyRows, [&](size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + quant::provisional::combined_code( + input.data() + (i * d), d, nbits, view.row(i), t_const + ); + } + }); + } + return view.owner; + }; + if (nbits == kMaxTotalBits) { + return run(output_rows(out, rows, d, {{"residuals", input}})); + } + return run(output_rows(out, rows, d, {{"residuals", input}})); +} + +template +py::tuple split_rows( + const RowView& codes, + size_t nbits, + size_t prefix_bits, + size_t num_threads, + const py::object& out +) { + // out= is a (prefix, ex) pair: each part is checked like any out=, and the two + // must not overlap each other. + py::object prefix_out = py::none(); + py::object ex_out = py::none(); + if (!out.is_none()) { + // A None part would otherwise get a new array, as if out= were partial. + if (!py::isinstance(out) || py::len(out) != 2 || + out[py::int_(0)].is_none() || out[py::int_(1)].is_none()) { + throw std::invalid_argument( + "out must be a (prefix, ex) tuple of two arrays; got " + + array_input_type(out) + ); + } + prefix_out = out[py::int_(0)]; + ex_out = out[py::int_(1)]; + } + const auto prefix = output_rows( + prefix_out, codes.rows, codes.cols, {{"codes", codes.owner}}, "out[0]" + ); + const auto ex = output_rows( + ex_out, + codes.rows, + codes.cols, + {{"codes", codes.owner}, {"out[0]", prefix.owner}}, + "out[1]" + ); + const size_t ex_bits = nbits - prefix_bits; + const auto mask = static_cast((1U << ex_bits) - 1); + { + py::gil_scoped_release release; + // Checked in a pass of its own, so a rejected call leaves out= unchanged. + check_code_values(codes, nbits, "codes", num_threads); + parallel_rows(codes.rows, num_threads, kCheapRows, [&](size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + const TP* row = codes.row(i); + uint8_t* prefix_row = prefix.row(i); + uint8_t* ex_row = ex.row(i); + for (size_t j = 0; j < codes.cols; ++j) { + prefix_row[j] = static_cast(row[j] >> ex_bits); + ex_row[j] = static_cast(row[j] & mask); + } + } + }); + } + return py::make_tuple(prefix.owner, ex.owner); +} + +py::tuple split( + py::handle codes, + size_t nbits, + size_t prefix_bits, + size_t num_threads, + const py::object& out +) { + check_nbits(nbits, 2, kMaxTotalBits); + check_nbits(prefix_bits, 1, nbits - 1, "prefix_bits"); + if (nbits == kMaxTotalBits) { + const auto view = exact_dtype_rows(codes, "codes"); + checked_working_dim(view.cols, view.owner, "codes"); + return split_rows(view, nbits, prefix_bits, num_threads, out); + } + const auto view = exact_dtype_rows(codes, "codes"); + checked_working_dim(view.cols, view.owner, "codes"); + return split_rows(view, nbits, prefix_bits, num_threads, out); +} + +template +py::array code_factors_rows( + const py::array_t& residuals, + const RowView& codes, + size_t nbits, + MetricType metric, + const std::vector& centroid, + size_t num_threads, + const py::object& out +) { + const size_t rows = codes.rows; + const size_t d = codes.cols; + auto [result, data] = vector_output( + out, + factors_dtype(), + "quant.FACTORS_DTYPE", + rows, + {{"residuals", residuals}, {"codes", codes.owner}} + ); + const bool caller_out = !out.is_none(); + { + py::gil_scoped_release release; + check_coordinates(residuals.data(), rows * d, d, "residuals"); + // code_factors checks while it computes, so a caller's out= is written only after + // every row succeeds. + std::vector staged(caller_out ? rows : 0); + quant::v1::Factors* factors = caller_out ? staged.data() : data; + for_each_row(rows, num_threads, kCostlyRows, [&](size_t i) { + // code_factors checks the row's code values and that the code quantizes the + // residual, and rejects factors that overflow float32; the binding adds the + // row number to its message. + const auto row = [i] { + return "codes row " + std::to_string(i) + " for residuals row " + + std::to_string(i) + ": "; + }; + try { + factors[i] = quant::provisional::code_factors( + residuals.data() + (i * d), + centroid.data(), + d, + codes.row(i), + nbits, + metric + ); + } catch (const std::invalid_argument& error) { + throw std::invalid_argument(row() + error.what()); + } catch (const std::overflow_error& error) { + throw std::overflow_error(row() + error.what()); + } + }); + std::copy(staged.begin(), staged.end(), data); + } + return result; +} + +py::array code_factors( + py::handle residuals, + py::handle codes, + size_t nbits, + const std::string& metric, + const py::object& centroid, + size_t num_threads, + const py::object& out +) { + check_nbits(nbits, 1, kMaxTotalBits); + const MetricType metric_type = metric_from_string(metric); + auto input = ensure_2d_array(residuals, "residuals"); + const size_t d = + checked_working_dim(static_cast(input.shape(1)), input, "residuals"); + const auto rows = static_cast(input.shape(0)); + // One zero vector for every row when there is no centroid. + std::vector center = centroid_vector(centroid, d); + center.resize(d, 0.0F); + + const auto run = [&](const auto& view) { + if (view.rows != rows || view.cols != d) { + throw std::invalid_argument( + "codes must have the shape of residuals, " + array_shape(input) + + "; got shape " + array_shape(view.owner) + ); + } + return code_factors_rows(input, view, nbits, metric_type, center, num_threads, out); + }; + if (nbits == kMaxTotalBits) { + return run(exact_dtype_rows(codes, "codes")); + } + return run(exact_dtype_rows(codes, "codes")); +} + +py::array pack_codes( + py::handle codes, size_t nbits, size_t num_threads, const py::object& out +) { + check_nbits(nbits, 1, 8); + const auto view = exact_dtype_rows(codes, "codes"); + const size_t d = checked_working_dim(view.cols, view.owner, "codes"); + const auto packed = + output_rows(out, view.rows, d * nbits / 8, {{"codes", view.owner}}); + { + py::gil_scoped_release release; + // Checked in a pass of its own, so a rejected call leaves out= unchanged. + check_code_values(view, nbits, "codes", num_threads); + parallel_rows(view.rows, num_threads, kCheapRows, [&](size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + view.row(i), packed.row(i), d, nbits + ); + } + }); + } + return packed.owner; +} + +py::array unpack_codes( + py::handle packed, size_t nbits, size_t num_threads, const py::object& out +) { + check_nbits(nbits, 1, 8); + const auto view = exact_dtype_rows(packed, "packed"); + const size_t bits = view.cols * 8; + if (view.cols == 0 || bits % nbits != 0 || (bits / nbits) % 32 != 0) { + throw std::invalid_argument( + "packed must have d * nbits / 8 bytes per row for a positive multiple of 32 " + "d; got shape " + + array_shape(view.owner) + " with nbits " + std::to_string(nbits) + ); + } + const size_t d = bits / nbits; + const auto codes = output_rows(out, view.rows, d, {{"packed", view.owner}}); + { + py::gil_scoped_release release; + parallel_rows(view.rows, num_threads, kCheapRows, [&](size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + view.row(i), codes.row(i), d, nbits + ); + } + }); + } + return codes.owner; +} + +// --- Layer 2 -------------------------------------------------------------------- + +// The query-side terms of the reference formula (docs/formats.md, core/estimator.hpp). +struct QueryTerms { + float g_add = 0; + float g_error = 0; + float g_k1xsumq = 0; +}; + +// g_add and g_k1xsumq are quant::v1::prepare_query's; g_error is |q - c|. +QueryTerms compute_query_terms( + const float* query, const float* centroid, size_t d, MetricType metric +) { + const auto corrections = quant::detail::query_corrections(query, centroid, d, metric); + const float distance_sqr = + metric == METRIC_L2 ? corrections.g_add + : (centroid == nullptr ? l2norm_sqr(query, d) + : euclidean_sqr(query, centroid, d)); + QueryTerms terms; + terms.g_add = corrections.g_add; + terms.g_error = std::sqrt(distance_sqr); + terms.g_k1xsumq = corrections.g_k1xsumq; + if (!std::isfinite(terms.g_add) || !std::isfinite(terms.g_error) || + !std::isfinite(terms.g_k1xsumq)) { + throw std::overflow_error("Query terms are not representable in float32"); + } + return terms; +} + +// The estimate of one row by the reference formula. Unpacked codes reach ip through +// quant::detail::unpacked_ip, the inner product of quant::v1::compute. The arithmetic +// is v1's, but this file is compiled without FP contraction and api.cpp is not, so +// distances agree with v1's within float32 rounding rather than bitwise. +// It is not checked here: write_rows rejects an estimate that is not finite. +inline Estimate estimate_of( + float ip, size_t bits, const quant::v1::Factors& f, const QueryTerms& terms +) { + return { + quant::detail::full_est_dist_from_ip( + ip, bits, f.add, f.rescale, terms.g_add, terms.g_k1xsumq + ), + f.error * terms.g_error, + ip}; +} + +py::object query_terms( + py::handle rotated_query, const std::string& metric, const py::object& centroid +) { + const MetricType metric_type = metric_from_string(metric); + const auto query = owned_vector(rotated_query, "rotated_query", 0); + const auto center = centroid_vector(centroid, query.size()); + const QueryTerms terms = compute_query_terms( + query.data(), center.empty() ? nullptr : center.data(), query.size(), metric_type + ); + py::array result(query_terms_dtype(), std::vector{1}); + auto* data = static_cast(result.mutable_data()); + data[0] = terms.g_add; + data[1] = terms.g_error; + data[2] = terms.g_k1xsumq; + return result.attr("__getitem__")(0); +} + +py::array estimate_from_ip( + py::handle ip, + py::handle factors, + py::handle terms, + size_t nbits, + size_t num_threads, + const py::object& out +) { + check_nbits(nbits, 1, kMaxTotalBits); + auto ips = ensure_1d_array(ip, "ip"); + const auto rows = static_cast(ips.shape(0)); + const auto records = factors_input(factors, rows, "ip"); + auto terms_array = py::array::ensure(terms); + if (!terms_array || !terms_array.dtype().equal(query_terms_dtype()) || + terms_array.size() != 1) { + throw std::invalid_argument( + "terms must be one record of dtype quant.QUERY_TERMS_DTYPE, as query_terms " + "returns; got " + + (terms_array ? "dtype " + py::str(terms_array.dtype()).cast() + + " and shape " + array_shape(terms_array) + : array_input_type(terms)) + ); + } + // Copied before anything is written, so terms may share memory with out. + std::array values{}; + std::memcpy(values.data(), terms_array.data(), sizeof(values)); + if (!all_finite(values.data(), values.size())) { + throw std::invalid_argument("terms must contain only finite values"); + } + const QueryTerms query{values[0], values[1], values[2]}; + auto [result, estimates] = + estimate_output(out, rows, {{"ip", ips}, {"factors", records.owner}}); + const Selection every_row = all_rows(rows); + { + py::gil_scoped_release release; + for_each_range( + rows, + num_threads, + kCheapRows, + [&](size_t begin, size_t end) { + return all_finite(ips.data() + begin, end - begin) && + finite_factor_range(records, every_row, begin, end); + }, + [&](size_t i) { + if (!std::isfinite(ips.data()[i])) { + throw std::invalid_argument("ip must contain only finite values"); + } + check_finite_factor(records[i]); + } + ); + write_rows(estimates, rows, num_threads, [&](size_t i) { + return estimate_of(ips.data()[i], nbits, records[i], query); + }); + } + return result; +} + +// --- Layer 3 -------------------------------------------------------------------- + +// Bit-reversal of every byte value. +constexpr std::array kReversedBytes = [] { + std::array table{}; + for (size_t value = 0; value < table.size(); ++value) { + uint8_t reversed = 0; + for (size_t bit = 0; bit < 8; ++bit) { + reversed = + static_cast(reversed | (((value >> bit) & 1U) << (7 - bit))); + } + table[value] = reversed; + } + return table; +}(); + +// FastScan blocks (docs/formats.md): 32 vectors per block, the last one zero-padded. +py::array pack_fastscan_blocks( + py::handle bitsets, py::handle factors, size_t num_threads, const py::object& out +) { + const auto codes = exact_dtype_rows(bitsets, "bitsets"); + if (codes.cols == 0 || (codes.cols * 8) % 32 != 0) { + throw std::invalid_argument( + "bitsets must have d // 8 bytes per row for a positive multiple of 32 d; got " + "shape " + + array_shape(codes.owner) + ); + } + const size_t d = codes.cols * 8; + const auto records = factors_input(factors, codes.rows, "bitsets"); + const size_t n = codes.rows; + const size_t num_blocks = div_round_up(n, fastscan::kBatchSize); + const auto blocks = output_rows( + out, + num_blocks, + BatchDataMap::data_bytes(d), + {{"bitsets", codes.owner}, {"factors", records.owner}} + ); + const Selection every_row = all_rows(n); + { + py::gil_scoped_release release; + // Checked in a pass of its own, so a rejected call leaves out= unchanged. + for_each_range( + n, + num_threads, + kCheapRows, + [&](size_t begin, size_t end) { + return finite_factor_range(records, every_row, begin, end); + }, + [&](size_t i) { check_finite_factor(records[i]); } + ); + parallel_rows(num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) { + // One block of sign bytes in the order the FastScan packer reads: the first + // coordinate of each byte in its highest bit. + std::vector msb_first(fastscan::kBatchSize * codes.cols); + for (size_t b = first; b < last; ++b) { + const size_t begin = b * fastscan::kBatchSize; + const size_t lanes = std::min(fastscan::kBatchSize, n - begin); + for (size_t lane = 0; lane < lanes; ++lane) { + const uint8_t* row = codes.row(begin + lane); + for (size_t byte = 0; byte < codes.cols; ++byte) { + msb_first[(lane * codes.cols) + byte] = kReversedBytes[row[byte]]; + } + } + uint8_t* block = blocks.row(b); + fastscan::pack_codes(d, msb_first.data(), lanes, block); + BatchDataMap map(reinterpret_cast(block), d); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const quant::v1::Factors f = + lane < lanes ? records[begin + lane] : quant::v1::Factors{0, 0, 0}; + map.f_add()[lane] = f.add; + map.f_rescale()[lane] = f.rescale; + map.f_error()[lane] = f.error; + } + } + }); + } + return blocks.owner; +} + +// --- Query ------------------------------------------------------------------------ + +// The FastScan query of a Query, built on the first estimate_fastscan call. It borrows +// the Query's own vector, so a copied or moved Query starts without one and builds its +// own. +struct LazyBatchQuery { + std::unique_ptr> query; + + LazyBatchQuery() = default; + LazyBatchQuery(const LazyBatchQuery& /*other*/) {} + LazyBatchQuery(LazyBatchQuery&& /*other*/) noexcept {} + LazyBatchQuery& operator=(const LazyBatchQuery&) = delete; + LazyBatchQuery& operator=(LazyBatchQuery&&) = delete; + ~LazyBatchQuery() = default; +}; + +// A Query's guard. A copy gets a fresh, idle guard rather than the original's state. +struct QueryAccess : IndexAccess { + QueryAccess() + : IndexAccess("Query is busy: set_centroid cannot overlap another call") {} + QueryAccess(const QueryAccess& /*other*/) : QueryAccess() {} + QueryAccess(QueryAccess&& /*other*/) noexcept : QueryAccess() {} + QueryAccess& operator=(const QueryAccess&) = delete; + QueryAccess& operator=(QueryAccess&&) = delete; + ~QueryAccess() = default; +}; + +// A rotated query in the working space, with the terms of its current centroid. It owns +// copies of the query and centroid, so copies and unpickled queries never share storage. +class QuantQuery { + private: + std::vector query_; + std::vector centroid_; // Empty means no centroid (zero). + MetricType metric_; + QueryTerms terms_; + LazyBatchQuery batch_query_; + QueryAccess access_; + + [[nodiscard]] size_t d() const { return query_.size(); } + + // SplitBatchQuery::set_g_add uses the - IP convention of the terms here. + void apply_batch_terms() { + if (metric_ == METRIC_L2) { + batch_query_.query->set_g_add(terms_.g_error); + } else { + batch_query_.query->set_g_add(terms_.g_error, -terms_.g_add); + } + } + + // Called with the GIL held, so concurrent readers cannot both build it. + const SplitBatchQuery& batch_query() { + if (!batch_query_.query) { + batch_query_.query = std::make_unique>( + query_.data(), d(), 0, metric_, true + ); + apply_batch_terms(); + } + return *batch_query_.query; + } + + // The guard is taken before any input conversion, so conversion cannot race + // set_centroid. + template + py::array scored(Score&& score) { + const auto operation = access_.read(); + return score(); + } + + static void check_row_length( + const py::array& array, + size_t cols, + size_t expected, + const char* name, + const char* what + ) { + if (cols != expected) { + throw std::invalid_argument( + std::string(name) + " must have " + what + " = " + + std::to_string(expected) + " columns for this query; got shape " + + array_shape(array) + ); + } + } + + public: + QuantQuery(std::vector query, MetricType metric, std::vector centroid) + : query_(std::move(query)) + , centroid_(std::move(centroid)) + , metric_(metric) + , terms_(compute_query_terms( + query_.data(), centroid_.empty() ? nullptr : centroid_.data(), d(), metric + )) {} + + static QuantQuery create( + py::handle rotated_query, const std::string& metric, const py::object& centroid + ) { + const MetricType metric_type = metric_from_string(metric); + auto query = owned_vector(rotated_query, "rotated_query", 0); + auto center = centroid_vector(centroid, query.size()); + return {std::move(query), metric_type, std::move(center)}; + } + + [[nodiscard]] size_t dim() const { return d(); } + [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } + + void set_centroid(const py::object& centroid) { + const auto operation = access_.write(); + auto center = centroid_vector(centroid, d()); + const QueryTerms terms = compute_query_terms( + query_.data(), center.empty() ? nullptr : center.data(), d(), metric_ + ); + centroid_ = std::move(center); + terms_ = terms; + if (batch_query_.query) { + apply_batch_terms(); + } + } + + [[nodiscard]] QuantQuery copy() { + const auto operation = access_.read(); + return QuantQuery(*this); + } + + py::tuple state() { + const auto operation = access_.read(); + py::array_t query(static_cast(d())); + std::memcpy(query.mutable_data(), query_.data(), d() * sizeof(float)); + py::object centroid = py::none(); + if (!centroid_.empty()) { + py::array_t center(static_cast(d())); + std::memcpy(center.mutable_data(), centroid_.data(), d() * sizeof(float)); + centroid = std::move(center); + } + return py::make_tuple(query, metric(), centroid); + } + + static QuantQuery from_state(const py::tuple& state) { + if (state.size() != 3 || !py::isinstance(state[1])) { + throw std::invalid_argument("Invalid Query pickle state"); + } + return create(state[0], state[1].cast(), state[2]); + } + + [[nodiscard]] std::string repr() { + const auto operation = access_.read(); + return "Query(dim=" + std::to_string(d()) + ", metric='" + metric() + + "', centroid=" + (centroid_.empty() ? "None" : "set") + ")"; + } + + py::array estimate_bitset( + py::handle bitsets, + py::handle factors, + const py::object& indices, + size_t num_threads, + const py::object& out + ) { + return scored([&] { + const auto codes = exact_dtype_rows(bitsets, "bitsets"); + check_row_length(codes.owner, codes.cols, d() / 8, "bitsets", "d // 8"); + const auto records = factors_input(factors, codes.rows, "bitsets"); + const Selection selected = selection(indices, codes.rows); + auto [result, estimates] = estimate_output( + out, + selected.count, + {{"bitsets", codes.owner}, + {"factors", records.owner}, + {"indices", selected.owner}} + ); + const auto ip_func = select_excode_ipfunc(1); + { + py::gil_scoped_release release; + for_each_range( + selected.count, + num_threads, + kCheapRows, + [&](size_t begin, size_t end) { + return finite_factor_range(records, selected, begin, end); + }, + [&](size_t k) { check_finite_factor(records[selected.row(k)]); } + ); + write_rows(estimates, selected.count, num_threads, [&](size_t k) { + const size_t i = selected.row(k); + return estimate_of( + ip_func(query_.data(), codes.row(i), d()), 1, records[i], terms_ + ); + }); + } + return result; + }); + } + + py::array estimate_codes( + py::handle codes, + py::handle factors, + size_t nbits, + const py::object& indices, + size_t num_threads, + const py::object& out + ) { + check_nbits(nbits, 1, kMaxTotalBits); + return scored([&] { + const auto run = [&](const auto& view) { + check_row_length(view.owner, view.cols, d(), "codes", "d"); + const auto records = factors_input(factors, view.rows, "codes"); + const Selection selected = selection(indices, view.rows); + auto [result, estimates] = estimate_output( + out, + selected.count, + {{"codes", view.owner}, + {"factors", records.owner}, + {"indices", selected.owner}} + ); + { + py::gil_scoped_release release; + for_each_row(selected.count, num_threads, kCheapRows, [&](size_t k) { + const size_t i = selected.row(k); + check_code_row(view, i, nbits, "codes"); + check_finite_factor(records[i]); + }); + write_rows(estimates, selected.count, num_threads, [&](size_t k) { + const size_t i = selected.row(k); + return estimate_of( + unpacked_row_ip(query_.data(), view, i), + nbits, + records[i], + terms_ + ); + }); + } + return result; + }; + if (nbits == kMaxTotalBits) { + return run(exact_dtype_rows(codes, "codes")); + } + return run(exact_dtype_rows(codes, "codes")); + }); + } + + // refine (Packed = false) reads unpacked extra bits; refine_packed reads + // pack_codes(ex, nbits=ex_bits) with the dispatched kernel. + template + py::array refine( + py::handle ex, + py::handle factors, + py::handle first, + size_t ex_bits, + size_t prefix_bits, + const py::object& indices, + size_t num_threads, + const py::object& out + ) { + check_nbits(prefix_bits, 1, kMaxTotalBits - 1, "prefix_bits"); + check_nbits(ex_bits, 1, kMaxTotalBits - prefix_bits, "ex_bits"); + const size_t total_bits = prefix_bits + ex_bits; + const char* name = Packed ? "packed" : "ex"; + return scored([&] { + const auto codes = exact_dtype_rows(ex, name); + if constexpr (Packed) { + check_row_length( + codes.owner, codes.cols, d() * ex_bits / 8, name, "d * ex_bits // 8" + ); + } else { + check_row_length(codes.owner, codes.cols, d(), name, "d"); + } + const auto records = factors_input(factors, codes.rows, name); + const Selection selected = selection(indices, codes.rows); + const auto estimates = + record_input(first, estimate_dtype(), "ESTIMATE_DTYPE", "first"); + if (estimates.rows != selected.count) { + throw std::invalid_argument( + std::string("first must have one record per ") + + (indices.is_none() ? "row of " + std::string(name) + : "entry of indices") + + ", " + std::to_string(selected.count) + "; got shape " + + array_shape(estimates.owner) + ); + } + auto [result, distances] = vector_output( + out, + py::dtype::of(), + "float32", + selected.count, + {{name, codes.owner}, + {"factors", records.owner}, + {"first", estimates.owner}, + {"indices", selected.owner}} + ); + const auto ip_func = select_excode_ipfunc(ex_bits); + const auto scale = static_cast(1U << ex_bits); + { + py::gil_scoped_release release; + for_each_row(selected.count, num_threads, kCheapRows, [&](size_t k) { + const size_t i = selected.row(k); + if constexpr (!Packed) { + check_code_row(codes, i, ex_bits, name); + } + check_finite_factor(records[i]); + if (!std::isfinite(estimates[k].refine_input)) { + throw std::invalid_argument( + "first['refine_input'] must contain only finite values" + ); + } + }); + write_rows(distances, selected.count, num_threads, [&](size_t k) { + const size_t i = selected.row(k); + const float ip_ex = Packed ? ip_func(query_.data(), codes.row(i), d()) + : unpacked_row_ip(query_.data(), codes, i); + const quant::v1::Factors& f = records[i]; + return quant::detail::full_est_dist_from_ip( + (scale * estimates[k].refine_input) + ip_ex, + total_bits, + f.add, + f.rescale, + terms_.g_add, + terms_.g_k1xsumq + ); + }); + } + return result; + }); + } + + py::array estimate_fastscan( + py::handle blocks, size_t num_points, size_t num_threads, const py::object& out + ) { + return scored([&] { + const auto view = exact_dtype_rows(blocks, "blocks"); + const size_t block_bytes = BatchDataMap::data_bytes(d()); + const size_t num_blocks = div_round_up(num_points, fastscan::kBatchSize); + if (view.rows != num_blocks || view.cols != block_bytes) { + throw std::invalid_argument( + "blocks must have shape (" + std::to_string(num_blocks) + ", " + + std::to_string(block_bytes) + ") for num_points " + + std::to_string(num_points) + " and d " + std::to_string(d()) + + "; got shape " + array_shape(view.owner) + ); + } + auto [result, estimates] = + estimate_output(out, num_points, {{"blocks", view.owner}}); + const SplitBatchQuery& query = batch_query(); + { + py::gil_scoped_release release; + const auto lanes_of = [&](size_t b) { + return std::min( + fastscan::kBatchSize, num_points - (b * fastscan::kBatchSize) + ); + }; + // Hand-built blocks get the finite-factor check of the other calls. + const auto finite_factors = [&](size_t b) { + const ConstBatchDataMap map( + reinterpret_cast(view.row(b)), d() + ); + std::array values{}; + for (const auto& field : + {map.f_add(), map.f_rescale(), map.f_error()}) { + field.copy_to(values.data(), values.size()); + if (!all_finite(values.data(), lanes_of(b))) { + return false; + } + } + return true; + }; + for_each_row(num_blocks, num_threads, kCheapRows, [&](size_t b) { + if (!finite_factors(b)) { + throw std::invalid_argument( + "blocks must contain only finite factors; block " + + std::to_string(b) + " does not" + ); + } + }); + // Whether block b writes finite estimates. + const auto score_block = [&](size_t b) { + const auto* block = reinterpret_cast(view.row(b)); + std::array distances{}; + std::array lower_bounds{}; + std::array ips{}; + std::array errors{}; + ConstBatchDataMap(block, d()) + .f_error() + .copy_to(errors.data(), errors.size()); + split_batch_estdist( + block, + query, + d(), + distances.data(), + lower_bounds.data(), + ips.data(), + true + ); + const size_t begin = b * fastscan::kBatchSize; + const size_t lanes = lanes_of(b); + for (size_t lane = 0; lane < lanes; ++lane) { + estimates[begin + lane] = { + distances[lane], errors[lane] * terms_.g_error, ips[lane]}; + } + return all_finite_records(estimates, begin, begin + lanes); + }; + for_each_range( + num_blocks, + num_threads, + kCheapRows, + [&](size_t first, size_t last) { + for (size_t b = first; b < last; ++b) { + if (!score_block(b)) { + return false; + } + } + return true; + }, + [&](size_t b) { + if (score_block(b)) { + return; + } + const size_t begin = b * fastscan::kBatchSize; + for (size_t lane = 0; lane < lanes_of(b); ++lane) { + (void)checked(estimates[begin + lane]); + } + } + ); + } + return result; + }); + } +}; + +// --- Blob format (docs/formats.md) ------------------------------------------------ + +constexpr std::array kBlobMagic = {'R', 'A', 'B', 'Q', 'B', 'L', 'O', 'B'}; +constexpr size_t kCommonHeaderBytes = 40; +constexpr uint32_t kRotatorKind = 1; +constexpr uint32_t kRotatorVersion = 1; +constexpr size_t kRotatorHeaderBytes = 64; +constexpr uint32_t kSeedKnown = 1; +constexpr uint32_t kMethodFhtKac = 1; +constexpr uint32_t kMethodMatrix = 2; +constexpr uint32_t kScalarFloat32 = 1; +constexpr uint32_t kScalarFloat64 = 2; + +// Checksum of a blob with its crc32 field (offset 32) taken as zero, as zlib.crc32 +// computes it, without copying the blob. The caller has checked that the blob covers +// the common header. Requires the GIL; zlib.crc32 releases it for large inputs. +uint32_t blob_checksum(const uint8_t* blob, size_t size) { + constexpr size_t kChecksumOffset = 32; + constexpr size_t kChecksumEnd = kChecksumOffset + 4; + const py::object crc32 = py::module_::import("zlib").attr("crc32"); + const auto range = [&](size_t begin, size_t end) { + return py::memoryview::from_memory( + blob + begin, static_cast(end - begin) + ); + }; + py::object crc = crc32(range(0, kChecksumOffset)); + crc = crc32(py::bytes(std::string(kChecksumEnd - kChecksumOffset, '\0')), crc); + crc = crc32(range(kChecksumEnd, size), crc); + return crc.cast(); +} + +void put_u32(uint8_t* out, uint32_t value) { + for (size_t i = 0; i < 4; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +void put_u64(uint8_t* out, uint64_t value) { + for (size_t i = 0; i < 8; ++i) { + out[i] = static_cast(value >> (8 * i)); + } +} + +uint32_t get_u32(const uint8_t* in) { + uint32_t value = 0; + for (size_t i = 0; i < 4; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +uint64_t get_u64(const uint8_t* in) { + uint64_t value = 0; + for (size_t i = 0; i < 8; ++i) { + value |= static_cast(in[i]) << (8 * i); + } + return value; +} + +// A contiguous, read-only byte view of any buffer-protocol object, held for one call. +class ByteView { + private: + py::buffer_info info_; + + public: + explicit ByteView(py::handle value) { + const std::string message = + "data must be a contiguous bytes-like object; got " + array_input_type(value); + if (!py::isinstance(value)) { + throw py::type_error(message); + } + info_ = py::reinterpret_borrow(value).request(); + // C-contiguous: each stride is the size of everything after it. + py::ssize_t expected = info_.itemsize; + for (py::ssize_t axis = info_.ndim - 1; axis >= 0; --axis) { + const auto index = static_cast(axis); + if (info_.shape[index] > 1 && info_.strides[index] != expected) { + throw py::type_error(message); + } + expected *= info_.shape[index]; + } + } + + [[nodiscard]] const uint8_t* data() const { + return static_cast(info_.ptr); + } + [[nodiscard]] size_t size() const { + return static_cast(info_.size * info_.itemsize); + } +}; + +// --- Rotator -------------------------------------------------------------------- + +RotatorType rotator_method(const std::string& method) { + if (method == "fht_kac" || method == "fht") { + return RotatorType::FhtKacRotator; + } + if (method == "matrix") { + return RotatorType::MatrixRotator; + } + throw std::invalid_argument("method must be 'fht_kac' (alias 'fht') or 'matrix'"); +} + +const char* method_name(RotatorType type) { + return type == RotatorType::FhtKacRotator ? "fht_kac" : "matrix"; +} + +size_t rotator_body_bytes(RotatorType type, size_t dim, size_t padded_dim) { + return type == RotatorType::FhtKacRotator ? padded_dim / 2 + : sizeof(float) * dim * padded_dim; +} + +// Python's quant.Rotator: immutable, so copies and threads may share it. +class QuantRotator { + private: + std::shared_ptr> rotator_; + RotatorType type_ = RotatorType::FhtKacRotator; + size_t dim_ = 0; + size_t padded_dim_ = 0; + std::optional seed_; + + QuantRotator() = default; + + static void check_dimensions(RotatorType type, size_t dim, size_t padded_dim) { + if (dim == 0 || dim > kMaxRotatorDim) { + throw std::invalid_argument( + "dim must be in [1, " + std::to_string(kMaxRotatorDim) + "]; got " + + std::to_string(dim) + ); + } + if (type == RotatorType::FhtKacRotator && dim < 64) { + throw std::invalid_argument( + "dim must be at least 64 for method 'fht_kac'; got " + std::to_string(dim) + ); + } + if (padded_dim % 32 != 0 || padded_dim < dim || padded_dim > kMaxRotatorDim) { + throw std::invalid_argument( + "padded_dim must be a multiple of 32 in [dim, " + + std::to_string(kMaxRotatorDim) + "]; got " + std::to_string(padded_dim) + ); + } + } + + static size_t checked_padded_dim_argument(const py::object& value) { + return checked_integer( + value, + kMaxRotatorDim, + "padded_dim must be an integer multiple of 32 in [dim, " + + std::to_string(kMaxRotatorDim) + "]" + ); + } + + public: + QuantRotator( + size_t dim, + const std::string& method, + const py::object& padded, + const py::object& seed + ) + : type_(rotator_method(method)), dim_(dim) { + padded_dim_ = padded.is_none() ? (dim == 0 ? 0 : padded_dim(dim)) + : checked_padded_dim_argument(padded); + check_dimensions(type_, dim_, padded_dim_); + seed_ = seed.is_none() ? std::random_device{}() + : static_cast(checked_seed( + seed, std::numeric_limits::max() + )); + py::gil_scoped_release release; + // A matrix rotation runs a QR factorization; keep the GIL released for it. + std::unique_ptr> created( + choose_rotator(dim_, type_, padded_dim_, *seed_) + ); + rotator_ = std::move(created); + } + + [[nodiscard]] size_t dim() const { return dim_; } + [[nodiscard]] size_t padded_dim_value() const { return padded_dim_; } + [[nodiscard]] std::string method() const { return method_name(type_); } + [[nodiscard]] std::optional seed() const { return seed_; } + + [[nodiscard]] py::array rotate( + py::handle x, size_t num_threads, const py::object& out + ) const { + auto input = ensure_2d_array(x, "x"); + ensure_dimension(input, dim_, "x dimension does not match dim"); + const auto rows = static_cast(input.shape(0)); + const auto view = output_rows(out, rows, padded_dim_, {{"x", input}}); + { + py::gil_scoped_release release; + check_coordinates(input.data(), rows * dim_, dim_, "x"); + parallel_rows(rows, num_threads, kCostlyRows, [&](size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + rotator_->rotate(input.data() + (i * dim_), view.row(i)); + } + }); + } + return view.owner; + } + + [[nodiscard]] py::bytes to_bytes() const { + const size_t body = rotator_->dump_bytes(); + const size_t size = kRotatorHeaderBytes + body; + // Written in place: a matrix rotator's body is up to 16 GiB, so no staging copy. + // Not py::bytes(nullptr, size): it turns a failed allocation into RuntimeError + // while leaving Python's MemoryError set. Raise that MemoryError instead. + auto result = py::reinterpret_steal( + PYBIND11_BYTES_FROM_STRING_AND_SIZE(nullptr, static_cast(size)) + ); + if (!result) { + throw py::error_already_set(); + } + auto* blob = reinterpret_cast(PYBIND11_BYTES_AS_STRING(result.ptr())); + std::memset(blob, 0, kRotatorHeaderBytes); + std::memcpy(blob, kBlobMagic.data(), kBlobMagic.size()); + put_u32(&blob[8], kRotatorKind); + put_u32(&blob[12], kRotatorVersion); + put_u32(&blob[16], static_cast(kRotatorHeaderBytes)); + put_u32(&blob[20], seed_.has_value() ? kSeedKnown : 0); + put_u64(&blob[24], body); + put_u32( + &blob[40], type_ == RotatorType::FhtKacRotator ? kMethodFhtKac : kMethodMatrix + ); + put_u32(&blob[44], kScalarFloat32); + put_u32(&blob[48], static_cast(dim_)); + put_u32(&blob[52], static_cast(padded_dim_)); + put_u32(&blob[56], seed_.value_or(0)); + rotator_->save(reinterpret_cast(blob + kRotatorHeaderBytes)); + put_u32(&blob[32], blob_checksum(blob, size)); + return result; + } + + static QuantRotator from_bytes(py::handle data) { + const ByteView bytes(data); + const uint8_t* blob = bytes.data(); + const size_t size = bytes.size(); + if (size < kCommonHeaderBytes) { + throw FormatError("data is too short for a rabitqlib blob"); + } + if (std::memcmp(blob, kBlobMagic.data(), kBlobMagic.size()) != 0) { + throw FormatError("data is not a rabitqlib blob (bad magic)"); + } + const uint32_t kind = get_u32(blob + 8); + if (kind != kRotatorKind) { + throw FormatError( + "data is not a rotator blob; got kind " + std::to_string(kind) + ); + } + const uint32_t version = get_u32(blob + 12); + if (version != kRotatorVersion) { + throw FormatError( + "data has unsupported rotator format version " + std::to_string(version) + + "; this release reads version " + std::to_string(kRotatorVersion) + ); + } + if (get_u32(blob + 16) != kRotatorHeaderBytes) { + throw FormatError("data has an invalid header size for rotator version 1"); + } + const uint64_t body = get_u64(blob + 24); + if (size < kRotatorHeaderBytes || body != size - kRotatorHeaderBytes) { + throw FormatError( + "data length does not match its header (truncated or padded)" + ); + } + if (get_u32(blob + 32) != blob_checksum(blob, size)) { + throw FormatError("data checksum mismatch (corrupt blob)"); + } + const uint32_t flags = get_u32(blob + 20); + if ((flags & ~kSeedKnown) != 0) { + throw FormatError("data has unknown flags"); + } + if (get_u32(blob + 36) != 0 || get_u32(blob + 60) != 0) { + throw FormatError("data has non-zero reserved fields"); + } + const uint32_t method = get_u32(blob + 40); + if (method != kMethodFhtKac && method != kMethodMatrix) { + throw FormatError( + "data has an unknown rotator method " + std::to_string(method) + ); + } + const uint32_t scalar = get_u32(blob + 44); + if (scalar == kScalarFloat64) { + throw FormatError("data stores a float64 rotator, which version 1 cannot load"); + } + if (scalar != kScalarFloat32) { + throw FormatError("data has an unknown scalar type " + std::to_string(scalar)); + } + QuantRotator result; + result.type_ = method == kMethodFhtKac ? RotatorType::FhtKacRotator + : RotatorType::MatrixRotator; + result.dim_ = get_u32(blob + 48); + result.padded_dim_ = get_u32(blob + 52); + try { + check_dimensions(result.type_, result.dim_, result.padded_dim_); + } catch (const std::invalid_argument& error) { + throw FormatError( + std::string("data has invalid rotator dimensions: ") + error.what() + ); + } + const uint32_t seed = get_u32(blob + 56); + if ((flags & kSeedKnown) != 0) { + result.seed_ = seed; + } else if (seed != 0) { + throw FormatError("data has a seed without the seed_known flag"); + } + if (body != rotator_body_bytes(result.type_, result.dim_, result.padded_dim_)) { + throw FormatError("data body size does not match the rotator parameters"); + } + auto rotator = rotator_impl::make_rotator_for_loading( + result.dim_, result.type_, result.padded_dim_ + ); + rotator->load(reinterpret_cast(blob + kRotatorHeaderBytes)); + result.rotator_ = std::move(rotator); + return result; + } + + [[nodiscard]] std::string repr() const { + return "Rotator(dim=" + std::to_string(dim_) + + ", padded_dim=" + std::to_string(padded_dim_) + ", method='" + method() + + "', seed=" + (seed_.has_value() ? std::to_string(*seed_) : "None") + ")"; + } +}; + +} // namespace +} // namespace rabitqlib::python_bindings + +void register_quant(py::module_& m) { + using namespace rabitqlib::python_bindings; + + auto quant = m.def_submodule( + "_quant", + "Native backend of rabitqlib.quant. Import rabitqlib.quant instead of this module." + ); + py::register_exception( + quant, "FormatError", py::module_::import("builtins").attr("ValueError") + ); + quant.attr("FormatError").attr("__doc__") = + "Corrupt or incompatible serialized bytes. A subclass of ValueError."; + + quant.attr("FACTORS_DTYPE") = factors_dtype(); + quant.attr("ESTIMATE_DTYPE") = estimate_dtype(); + quant.attr("QUERY_TERMS_DTYPE") = query_terms_dtype(); + + quant.def( + "padded_dim", + &padded_dim, + py::arg("dim"), + R"doc(Return the working dimension for dim: dim rounded up to a multiple of 32. + +Parameters +---------- +dim : int + Original vector dimension, positive. + +Returns +------- +padded_dim : int + The smallest multiple of 32 that is at least dim.)doc" + ); + + quant.def( + "quantize", + &quantize, + py::arg("residuals"), + py::kw_only(), + py::arg("nbits"), + py::arg("fast_quantization") = false, + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Quantize rotated residuals into unsigned nbits-bit codes. + +Parameters +---------- +residuals : array_like of float32, shape (n, d) + Rotated, padded residuals (vector minus centroid); d is a positive + multiple of 32. Converted to contiguous float32. +nbits : int + Total bits per coordinate, 1 to 9, including the sign bit. +fast_quantization : bool, optional + Use the frozen fast-mode scale instead of the exact search. Default is + False. It has no effect at nbits=1, where no scale is searched, and is + rejected for d above 65536 otherwise. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. Results do not + depend on the thread count. +out : numpy.ndarray, optional + Output of the exact dtype and shape, with contiguous rows. + +Returns +------- +codes : numpy.ndarray, shape (n, d) + Codes in [0, 2**nbits): uint8, or uint16 when nbits is 9.)doc" + ); + + quant.def( + "split", + &split, + py::arg("codes"), + py::kw_only(), + py::arg("nbits"), + py::arg("prefix_bits"), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Split nbits-bit codes into a prefix and extra bits. + +Parameters +---------- +codes : numpy.ndarray, shape (n, d) + Codes from quantize: uint8, or uint16 when nbits is 9. Never cast; rows + may have any stride. +nbits : int + Width of codes, 2 to 9. +prefix_bits : int + Width of the prefix, 1 to nbits - 1. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : tuple of two numpy.ndarray, optional + (prefix, ex) outputs, each of dtype uint8 and shape (n, d) with contiguous + rows. They must not overlap codes or each other. + +Returns +------- +prefix, ex : tuple of numpy.ndarray of uint8, shape (n, d) + prefix = codes >> ex_bits and ex = codes & (2**ex_bits - 1), where + ex_bits = nbits - prefix_bits. With out=, the two arrays of out.)doc" + ); + + quant.def( + "code_factors", + &code_factors, + py::arg("residuals"), + py::arg("codes"), + py::kw_only(), + py::arg("nbits"), + py::arg("metric"), + py::arg("centroid") = py::none(), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Return the factors of nbits-bit codes for their residuals. + +Pass the full codes with their width for the full-code factors, or a prefix +from split with nbits=prefix_bits for the prefix factors. + +Parameters +---------- +residuals : array_like of float32, shape (n, d) + The rotated, padded residuals the codes were quantized from. +codes : numpy.ndarray, shape (n, d) + uint8 codes, or uint16 when nbits is 9. Never cast; rows may have any + stride. Each row must quantize its residual. +nbits : int + Width of codes, 1 to 9. +metric : {'l2', 'ip'} + Distance metric. Required: factors for the wrong metric give wrong + distances without an error. +centroid : array_like of float32, shape (d,), optional + The centroid in working space; None means zero. It only affects add. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray, optional + A C-contiguous FACTORS_DTYPE array of shape (n,). + +Returns +------- +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The add, rescale and error factors of each code.)doc" + ); + + quant.def( + "pack_codes", + &pack_codes, + py::arg("codes"), + py::kw_only(), + py::arg("nbits"), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Pack nbits-bit codes into the packed-codes format (version 1). + +At nbits=1 this is the sign bitset, equal to +numpy.packbits(codes, axis=1, bitorder="little"). + +Parameters +---------- +codes : numpy.ndarray of uint8, shape (n, d) + Codes in [0, 2**nbits). Never cast; rows may have any stride. +nbits : int + Width of codes, 1 to 8. Split 9-bit codes first. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray, optional + Output of the exact dtype and shape, with contiguous rows. + +Returns +------- +packed : numpy.ndarray of uint8, shape (n, d * nbits // 8))doc" + ); + + quant.def( + "unpack_codes", + &unpack_codes, + py::arg("packed"), + py::kw_only(), + py::arg("nbits"), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Unpack codes written by pack_codes; the inverse is lossless. + +Parameters +---------- +packed : numpy.ndarray of uint8, shape (n, d * nbits // 8) + Packed rows; d, a positive multiple of 32, follows from the row size. + Never cast; rows may have any stride. +nbits : int + Width of the packed codes, 1 to 8. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray, optional + Output of the exact dtype and shape, with contiguous rows. + +Returns +------- +codes : numpy.ndarray of uint8, shape (n, d))doc" + ); + + static constexpr char kRotatorDoc[] = + R"doc(A seeded rotation from the original dimension to the padded working space. + +Parameters +---------- +dim : int + Original vector dimension: 64 to 65536 for 'fht_kac', 1 to 65536 for + 'matrix'. +method : {'fht_kac', 'fht', 'matrix'}, optional + FHT/Kac rotation (default; 'fht' is an alias) or a dense random matrix. +padded_dim : int, optional + Working dimension, a multiple of 32 in [dim, 65536]. Default is + padded_dim(dim). +seed : int, optional + Seed in [0, 2**32 - 1]; None (the default) draws one. A seeded FHT/Kac + rotation is identical on every platform; a matrix rotation only with the + same standard library and build, so persist rotators with to_bytes. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import quant +>>> rot = quant.Rotator(128, seed=0) +>>> rot.rotate(np.ones((2, 128), dtype=np.float32)).shape +(2, 128))doc"; + + py::class_(quant, "Rotator", kRotatorDoc) + .def( + py::init(), + py::arg("dim"), + py::kw_only(), + py::arg("method") = "fht_kac", + py::arg("padded_dim") = py::none(), + py::arg("seed") = py::none(), + kRotatorDoc + ) + .def_property_readonly( + "dim", &QuantRotator::dim, R"doc(Original vector dimension.)doc" + ) + .def_property_readonly( + "padded_dim", + &QuantRotator::padded_dim_value, + R"doc(Working dimension of rotated vectors.)doc" + ) + .def_property_readonly( + "method", + &QuantRotator::method, + R"doc(Canonical method name, 'fht_kac' or 'matrix'.)doc" + ) + .def_property_readonly( + "seed", + &QuantRotator::seed, + R"doc(The seed the rotation was drawn from, or None if unknown.)doc" + ) + .def( + "rotate", + &QuantRotator::rotate, + py::arg("x"), + py::kw_only(), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Rotate vectors into the padded working space. + +Parameters +---------- +x : array_like of float32, shape (n, dim) + Vectors in the original space. Converted to contiguous float32. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray, optional + A float32 array of shape (n, padded_dim) with contiguous rows. + +Returns +------- +rotated : numpy.ndarray of float32, shape (n, padded_dim))doc" + ) + .def( + "to_bytes", + &QuantRotator::to_bytes, + R"doc(Serialize the rotation as a versioned rotator blob (docs/formats.md).)doc" + ) + .def_static( + "from_bytes", + &QuantRotator::from_bytes, + py::arg("data"), + R"doc(Load a rotator blob written by to_bytes. + +Parameters +---------- +data : bytes-like + Any contiguous buffer, read without an intermediate copy. + +Raises +------ +FormatError + If the bytes are corrupt, truncated or of an unsupported version.)doc" + ) + .def( + "__copy__", + [](const py::object& self) { return self; }, + "Return self; rotators are immutable." + ) + .def( + "__deepcopy__", + [](const py::object& self, const py::object& /*memo*/) { return self; }, + py::arg("memo"), + "Return self; rotators are immutable." + ) + .def("__repr__", &QuantRotator::repr, "Return the rotator's parameters.") + .def( + py::pickle( + [](const QuantRotator& self) { return self.to_bytes(); }, + [](const py::bytes& state) { return QuantRotator::from_bytes(state); } + ) + ); + quant.def( + "query_terms", + &query_terms, + py::arg("rotated_query"), + py::kw_only(), + py::arg("metric"), + py::arg("centroid") = py::none(), + R"doc(Return the query-side terms of the reference estimator formula. + +g_add is |q - c|**2 for L2 and - for IP, g_error is |q - c|, and +g_k1xsumq is -0.5 * sum(q). The 1 of the IP distance 1 - is folded +into the factors' add, not into g_add. + +Parameters +---------- +rotated_query : array_like of float32, shape (d,) + A rotated, padded query; d is a positive multiple of 32. +metric : {'l2', 'ip', 'innerproduct'} + The metric the factors were computed for. +centroid : array_like of float32, shape (d,), optional + The working-space centroid the factors were computed with; None is zero. + +Returns +------- +terms : numpy.void of QUERY_TERMS_DTYPE)doc" + ); + + quant.def( + "estimate_from_ip", + &estimate_from_ip, + py::arg("ip"), + py::arg("factors"), + py::arg("terms"), + py::kw_only(), + py::arg("nbits"), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + R"doc(Apply the reference estimator formula to the caller's inner products. + +distance = add + g_add + rescale * (ip + g_k1xsumq * (2**nbits - 1)), +error = error_factor * g_error, and refine_input = ip. + +Parameters +---------- +ip : array_like of float32, shape (n,) + Inner products with the unsigned nbits-bit codes. +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The factors of the same codes. +terms : numpy.void of QUERY_TERMS_DTYPE + The query terms, as query_terms returns them. +nbits : int + Bit width of the codes, 1 to 9. +num_threads : int, optional + Worker threads; 1 by default, 0 selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray of ESTIMATE_DTYPE, shape (n,), optional + +Returns +------- +estimates : numpy.ndarray of ESTIMATE_DTYPE, shape (n,))doc" + ); + + quant.def( + "pack_fastscan_blocks", + &pack_fastscan_blocks, + py::arg("bitsets"), + py::arg("factors"), + py::kw_only(), + py::arg("num_threads") = 0, + py::arg("out") = py::none(), + R"doc(Pack bitsets and their factors into FastScan blocks (version 1). + +Each block holds 32 vectors, the library's fixed FastScan batch; the last block +is zero-padded when n % 32 != 0. + +Parameters +---------- +bitsets : numpy.ndarray of uint8, shape (n, d // 8) + pack_codes(prefix, nbits=1) output. +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The 1-bit (prefix) factors. +num_threads : int, optional + Worker threads; 0 (the default) selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray of uint8, optional + Shape (ceil(n / 32), d * 4 + 384) with contiguous rows. + +Returns +------- +blocks : numpy.ndarray of uint8, shape (ceil(n / 32), d * 4 + 384))doc" + ); + + static constexpr const char* kQueryDoc = + R"doc(A rotated, padded query for scoring codes in the working space. + +The query and centroid are copied. set_centroid switches to another +centroid's factors without preparing the query again. Scoring calls may run +concurrently; set_centroid during one raises RuntimeError, so use copy() to +get one query per thread. + +Parameters +---------- +rotated_query : array_like of float32, shape (d,) + A rotated, padded query; d is a positive multiple of 32. +metric : {'l2', 'ip', 'innerproduct'} + The metric the factors were computed for. +centroid : array_like of float32, shape (d,), optional + The working-space centroid the factors were computed with; None is zero. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import quant +>>> x = np.random.default_rng(0).standard_normal((4, 64)).astype(np.float32) +>>> codes = quant.quantize(x, nbits=1) +>>> factors = quant.code_factors(x, codes, nbits=1, metric="l2") +>>> q = quant.Query(x[0], metric="l2") +>>> q.estimate_bitset(quant.pack_codes(codes, nbits=1), factors).shape +(4,))doc"; + + static constexpr const char* kIndicesDoc = R"doc( +indices : array_like of int, shape (m,), optional + Rows to score, read in place; the output lines up with indices. Must be + integers in [0, n). +num_threads : int, optional + Worker threads; 1 by default, 0 selects the available CPU count, capped by + OMP_NUM_THREADS.)doc"; + + py::class_(quant, "Query", kQueryDoc) + .def( + py::init(&QuantQuery::create), + py::arg("rotated_query"), + py::kw_only(), + py::arg("metric"), + py::arg("centroid") = py::none(), + kQueryDoc + ) + .def_property_readonly( + "dim", &QuantQuery::dim, R"doc(Working dimension d of the query.)doc" + ) + .def_property_readonly( + "metric", &QuantQuery::metric, R"doc(The metric, 'l2' or 'ip'.)doc" + ) + .def( + "set_centroid", + &QuantQuery::set_centroid, + py::arg("centroid"), + R"doc(Switch to the factors of another centroid; None is zero. + +Parameters +---------- +centroid : array_like of float32, shape (d,) or None + The working-space centroid the next factors were computed with.)doc" + ) + .def( + "estimate_bitset", + &QuantQuery::estimate_bitset, + py::arg("bitsets"), + py::arg("factors"), + py::kw_only(), + py::arg("indices") = py::none(), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + (std::string( + R"doc(Estimate distances from packed 1-bit codes with the float query. + +Parameters +---------- +bitsets : numpy.ndarray of uint8, shape (n, d // 8) + pack_codes(prefix, nbits=1) output. +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The 1-bit (prefix) factors.)doc" + ) + + kIndicesDoc + R"doc( +out : numpy.ndarray of ESTIMATE_DTYPE, optional + +Returns +------- +estimates : numpy.ndarray of ESTIMATE_DTYPE)doc") + .c_str() + ) + .def( + "estimate_codes", + &QuantQuery::estimate_codes, + py::arg("codes"), + py::arg("factors"), + py::kw_only(), + py::arg("nbits"), + py::arg("indices") = py::none(), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + (std::string(R"doc(Estimate distances from unpacked codes of any width. + +Parameters +---------- +codes : numpy.ndarray of uint8 (nbits 1 to 8) or uint16 (nbits 9), shape (n, d) +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The factors of these nbits-bit codes. +nbits : int + Bit width of the codes, 1 to 9.)doc") + + kIndicesDoc + R"doc( +out : numpy.ndarray of ESTIMATE_DTYPE, optional + +Returns +------- +estimates : numpy.ndarray of ESTIMATE_DTYPE)doc") + .c_str() + ) + .def( + "refine", + &QuantQuery::refine, + py::arg("ex"), + py::arg("factors"), + py::arg("first"), + py::kw_only(), + py::arg("ex_bits"), + py::arg("prefix_bits") = 1, + py::arg("indices") = py::none(), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + (std::string(R"doc(Refine prefix estimates with the unpacked extra bits. + +Parameters +---------- +ex : numpy.ndarray of uint8, shape (n, d) + The low ex_bits bits of each code, from split. +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The full-code factors (prefix_bits + ex_bits bits). +first : numpy.ndarray of ESTIMATE_DTYPE + The prefix estimates of the scored rows; only refine_input is read. +ex_bits : int + Width of the extra bits, 1 to 9 - prefix_bits. +prefix_bits : int, optional + Width of the prefix the first estimates scored, 1 by default.)doc") + + kIndicesDoc + R"doc( +out : numpy.ndarray of float32, optional + +Returns +------- +distances : numpy.ndarray of float32)doc") + .c_str() + ) + .def( + "refine_packed", + &QuantQuery::refine, + py::arg("packed"), + py::arg("factors"), + py::arg("first"), + py::kw_only(), + py::arg("ex_bits"), + py::arg("prefix_bits") = 1, + py::arg("indices") = py::none(), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + (std::string(R"doc(Refine prefix estimates with packed extra bits. + +Uses the dispatched SIMD kernel, so it agrees with refine within float32 +rounding. + +Parameters +---------- +packed : numpy.ndarray of uint8, shape (n, d * ex_bits // 8) + pack_codes(ex, nbits=ex_bits) output. +factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) + The full-code factors (prefix_bits + ex_bits bits). +first : numpy.ndarray of ESTIMATE_DTYPE + The prefix estimates of the scored rows; only refine_input is read. +ex_bits : int + Width of the extra bits, 1 to 9 - prefix_bits. +prefix_bits : int, optional + Width of the prefix the first estimates scored, 1 by default.)doc") + + kIndicesDoc + R"doc( +out : numpy.ndarray of float32, optional + +Returns +------- +distances : numpy.ndarray of float32)doc") + .c_str() + ) + .def( + "estimate_fastscan", + &QuantQuery::estimate_fastscan, + py::arg("blocks"), + py::kw_only(), + py::arg("num_points"), + py::arg("num_threads") = 1, + py::arg("out") = py::none(), + R"doc(Estimate distances for FastScan blocks with a 4-bit query table. + +The query table is built on the first call. Its quantization makes these +estimates agree with estimate_bitset only within that table's error. + +Parameters +---------- +blocks : numpy.ndarray of uint8, shape (ceil(num_points / 32), d * 4 + 384) + pack_fastscan_blocks output. +num_points : int + Number of vectors in the blocks; padding lanes are not returned. +num_threads : int, optional + Worker threads; 1 by default, 0 selects the available CPU count, capped by + OMP_NUM_THREADS. +out : numpy.ndarray of ESTIMATE_DTYPE, shape (num_points,), optional + +Returns +------- +estimates : numpy.ndarray of ESTIMATE_DTYPE, shape (num_points,))doc" + ) + .def("copy", &QuantQuery::copy, "Return an independent copy of the query.") + .def("__copy__", &QuantQuery::copy, "Return an independent copy of the query.") + .def( + "__deepcopy__", + [](QuantQuery& self, const py::object& /*memo*/) { return self.copy(); }, + py::arg("memo"), + "Return an independent copy of the query." + ) + .def("__repr__", &QuantQuery::repr, "Return the query's parameters.") + .def( + py::pickle( + [](QuantQuery& self) { return self.state(); }, + [](const py::tuple& state) { return QuantQuery::from_state(state); } + ) + ); +} diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index 1d171ad4..c609ca06 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -12,6 +12,7 @@ void register_hnsw(py::module_& m); void register_ivf(py::module_& m); void register_symqg(py::module_& m); void register_kmeans(py::module_& m); +void register_quant(py::module_& m); PYBIND11_MODULE(_rabitqlib, m) { m.doc() = "RabitQ Python bindings combined module"; @@ -36,4 +37,5 @@ PYBIND11_MODULE(_rabitqlib, m) { register_ivf(m); register_symqg(m); register_kmeans(m); + register_quant(m); } diff --git a/src/core/api.cpp b/src/core/api.cpp index 37c6ea92..7efac9ba 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -6,7 +6,6 @@ #include #include #include -#include #include #include @@ -16,28 +15,25 @@ #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/pack_excode.hpp" -#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" +#include "validation.hpp" namespace rabitqlib::quant::v1 { namespace { +using validation::check_disjoint; + constexpr size_t kWorkspaceAlignment = alignof(std::max_align_t); void validate(const Parameters& parameters) { if (parameters.dim < 2) { throw std::invalid_argument("dim must be at least 2"); } - if (parameters.dim > static_cast(std::numeric_limits::max())) { - throw std::length_error("dim exceeds the supported index range"); - } + validation::check_dim_range(parameters.dim); if (parameters.total_bits < 1 || parameters.total_bits > 9) { throw std::invalid_argument("total_bits must be in [1, 9]"); } validate_metric_type(parameters.metric); - if (!std::isfinite(parameters.t_const) || - parameters.t_const > std::numeric_limits::max() / 2.0) { - throw std::invalid_argument("t_const must be finite and at most INT_MAX / 2"); - } + validation::check_t_const(parameters.t_const); } size_t checked_multiply(size_t count, size_t width) { @@ -74,29 +70,11 @@ void check_buffer(ConstBuffer buffer, size_t required) { } } -void check_disjoint( - const void* first, size_t first_bytes, const void* second, size_t second_bytes -) { - if (first_bytes == 0 || second_bytes == 0) { - return; - } - const auto first_address = reinterpret_cast(first); - const auto second_address = reinterpret_cast(second); - if (first_address < second_address + second_bytes && - second_address < first_address + first_bytes) { - throw std::invalid_argument( - "Writable buffers must not overlap other buffers or inputs" - ); - } -} - void check_input(const float* input, size_t dim) { check_buffer( {input, checked_multiply(dim, sizeof(float))}, checked_multiply(dim, sizeof(float)) ); - if (!all_finite(input, dim)) { - throw std::invalid_argument("Input vectors must contain only finite values"); - } + validation::check_finite_input(input, dim); } // Parameters and arithmetic ranges were validated by requirements/prepare_query. @@ -186,11 +164,15 @@ uint16_t read_u16(const uint8_t* bytes) { return static_cast(bytes[0] | (uint16_t{bytes[1]} << 8U)); } +[[noreturn]] void reject_unpacked_code() { + throw std::invalid_argument("Unpacked code exceeds the configured bit width"); +} + // The i-th UnpackedU16 code, rejected if it does not fit the configured bit width. uint16_t checked_unpacked_code(const uint8_t* bytes, size_t i, uint16_t maximum) { const uint16_t value = read_u16(bytes + 2 * i); if (value > maximum) { - throw std::invalid_argument("Unpacked code exceeds the configured bit width"); + reject_unpacked_code(); } return value; } @@ -295,10 +277,7 @@ Factors encode( centroid, 1, static_cast(parameters.dim) ); const float residual_squared_norm = - l2norm_sqr(storage.scratch.residual, parameters.dim); - if (!std::isfinite(residual_squared_norm)) { - throw std::overflow_error("Residual squared norm is not representable in float32"); - } + validation::checked_residual_squared_norm(storage.scratch.residual, parameters.dim); Factors factors; if (parameters.total_bits > 1) { rabitq_impl::ex_bits::code_with_factor_from_residual( @@ -330,10 +309,7 @@ Factors encode( &storage.scratch ); } - if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || - !std::isfinite(factors.error)) { - throw std::overflow_error("Quantization factors are not representable in float32"); - } + validation::check_factors_representable(factors); write_codes(parameters, output, storage); return factors; } @@ -389,12 +365,11 @@ QueryState prepare_query( state.total_bits_ = parameters.total_bits; state.metric_ = parameters.metric; state.extra_ip_ = select_excode_ipfunc(parameters.total_bits - 1); - state.g_k1xsumq_ = std::accumulate(query, query + parameters.dim, 0.0F) * -0.5F; - state.g_add_ = - parameters.metric == METRIC_L2 - ? (centroid == nullptr ? l2norm_sqr(query, parameters.dim) - : euclidean_sqr(query, centroid, parameters.dim)) - : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, parameters.dim)); + const auto corrections = quant::detail::query_corrections( + query, centroid, parameters.dim, parameters.metric + ); + state.g_k1xsumq_ = corrections.g_k1xsumq; + state.g_add_ = corrections.g_add; if (!std::isfinite(state.g_add_) || !std::isfinite(state.g_k1xsumq_)) { throw std::overflow_error("Query factors are not representable in float32"); } @@ -419,10 +394,12 @@ float QueryState::compute_checked( float ip = 0; if (code.layout == CodeLayout::UnpackedU16) { const auto* bytes = static_cast(code.codes.data); - const auto maximum = static_cast((1U << parameters.total_bits) - 1); - for (size_t i = 0; i < parameters.dim; ++i) { - ip += query_wrapper.query_[i] * - static_cast(checked_unpacked_code(bytes, i, maximum)); + const auto maximum = static_cast((1U << parameters.total_bits) - 1); + const auto code_at = [bytes](size_t i) { return read_u16(bytes + (2 * i)); }; + if (quant::detail::unpacked_ip( + query_wrapper.query_, parameters.dim, maximum, code_at, ip + ) != parameters.dim) { + reject_unpacked_code(); } } else { const size_t ex_bits = parameters.total_bits - 1; diff --git a/src/core/provisional.cpp b/src/core/provisional.cpp new file mode 100644 index 00000000..419c1dd9 --- /dev/null +++ b/src/core/provisional.cpp @@ -0,0 +1,161 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/provisional/quantize.hpp" +#include "validation.hpp" + +namespace rabitqlib::quant::provisional { +namespace { + +// Widest code a TP element holds, capped at one sign bit plus eight extra bits. +template +constexpr size_t max_total_bits() { + return std::min(9, std::numeric_limits::digits); +} + +void check_dim(size_t dim) { + if (dim == 0 || dim % 32 != 0) { + throw std::invalid_argument("dim must be a positive multiple of 32"); + } + validation::check_dim_range(dim); +} + +template +void check_total_bits(size_t total_bits) { + if (total_bits < 1 || total_bits > max_total_bits()) { + throw std::invalid_argument( + max_total_bits() == 8 ? "total_bits must be in [1, 8] for uint8_t codes" + : "total_bits must be in [1, 9] for uint16_t codes" + ); + } +} + +// A residual is required; a centroid is optional and checked only when present. +void check_residual(const float* residual, size_t dim) { + if (residual == nullptr) { + throw std::invalid_argument("residual must not be null"); + } + validation::check_finite_input(residual, dim); +} + +template +void combined_code_impl( + const float* residual, size_t dim, size_t total_bits, TP* code, double t_const +) { + check_dim(dim); + check_total_bits(total_bits); + check_residual(residual, dim); + if (code == nullptr) { + throw std::invalid_argument("code must not be null"); + } + validation::check_disjoint(code, dim * sizeof(TP), residual, dim * sizeof(float)); + validation::check_t_const(t_const); + (void)validation::checked_residual_squared_norm(residual, dim); + + // Every check has passed; the quantizer can now fail only on allocation, after + // which the contents of code are unspecified. + rabitq_impl::total_bits::combined_code( + residual, dim, total_bits, code, t_const + ); +} + +template +v1::Factors code_factors_impl( + const float* residual, + const float* centroid, + size_t dim, + const TP* code, + size_t total_bits, + MetricType metric +) { + check_dim(dim); + check_total_bits(total_bits); + validate_metric_type(metric); + check_residual(residual, dim); + if (centroid != nullptr) { + validation::check_finite_input(centroid, dim); + } + if (code == nullptr) { + throw std::invalid_argument("code must not be null"); + } + (void)validation::checked_residual_squared_norm(residual, dim); + + // The detail function takes int codes and their offset-binary constant. + const auto maximum = static_cast((1U << total_bits) - 1); + const float cb = -static_cast(maximum) / 2.0F; + std::vector widened(dim); + for (size_t i = 0; i < dim; ++i) { + widened[i] = static_cast(code[i]); + if (widened[i] > maximum) { + throw std::invalid_argument("Unpacked code exceeds the configured bit width"); + } + } + // The factor derivation rejects a code that does not quantize this residual (its + // own code or a prefix of it), testing the inner product it divides by. + + std::vector zero; + if (centroid == nullptr) { + zero.assign(dim, 0.0F); + centroid = zero.data(); + } + v1::Factors factors; + rabitq_impl::xy_bits::code_factors( + residual, + centroid, + dim, + widened.data(), + cb, + factors.add, + factors.rescale, + factors.error, + metric + ); + validation::check_factors_representable(factors); + return factors; +} + +} // namespace + +void combined_code( + const float* residual, size_t dim, size_t total_bits, uint8_t* code, double t_const +) { + combined_code_impl(residual, dim, total_bits, code, t_const); +} + +void combined_code( + const float* residual, size_t dim, size_t total_bits, uint16_t* code, double t_const +) { + combined_code_impl(residual, dim, total_bits, code, t_const); +} + +v1::Factors code_factors( + const float* residual, + const float* centroid, + size_t dim, + const uint8_t* code, + size_t total_bits, + MetricType metric +) { + return code_factors_impl(residual, centroid, dim, code, total_bits, metric); +} + +v1::Factors code_factors( + const float* residual, + const float* centroid, + size_t dim, + const uint16_t* code, + size_t total_bits, + MetricType metric +) { + return code_factors_impl(residual, centroid, dim, code, total_bits, metric); +} + +} // namespace rabitqlib::quant::provisional diff --git a/src/core/validation.hpp b/src/core/validation.hpp new file mode 100644 index 00000000..44edb156 --- /dev/null +++ b/src/core/validation.hpp @@ -0,0 +1,71 @@ +#pragma once + +// Input checks shared by the stable v1 API (api.cpp) and the provisional API +// (provisional.cpp), so both reject the same inputs with the same messages. This header +// is private to the core sources and is not installed. + +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/utils/input_checks.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::quant::validation { + +// Eigen indexes the working vectors with long. +inline void check_dim_range(size_t dim) { + if (dim > static_cast(std::numeric_limits::max())) { + throw std::length_error("dim exceeds the supported index range"); + } +} + +inline void check_t_const(double t_const) { + if (!std::isfinite(t_const) || t_const > std::numeric_limits::max() / 2.0) { + throw std::invalid_argument("t_const must be finite and at most INT_MAX / 2"); + } +} + +inline void check_disjoint( + const void* first, size_t first_bytes, const void* second, size_t second_bytes +) { + if (first_bytes == 0 || second_bytes == 0) { + return; + } + const auto first_address = reinterpret_cast(first); + const auto second_address = reinterpret_cast(second); + if (first_address < second_address + second_bytes && + second_address < first_address + first_bytes) { + throw std::invalid_argument( + "Writable buffers must not overlap other buffers or inputs" + ); + } +} + +inline void check_finite_input(const float* values, size_t dim) { + if (!all_finite(values, dim)) { + throw std::invalid_argument("Input vectors must contain only finite values"); + } +} + +// Rejects a residual whose squared norm overflows before the quantizer's +// float-to-integer conversions. Returns the squared norm. +inline float checked_residual_squared_norm(const float* residual, size_t dim) { + const float squared_norm = l2norm_sqr(residual, dim); + if (!std::isfinite(squared_norm)) { + throw std::overflow_error("Residual squared norm is not representable in float32"); + } + return squared_norm; +} + +inline void check_factors_representable(const v1::Factors& factors) { + if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || + !std::isfinite(factors.error)) { + throw std::overflow_error("Quantization factors are not representable in float32"); + } +} + +} // namespace rabitqlib::quant::validation diff --git a/tests/python/fixtures/quant/README.md b/tests/python/fixtures/quant/README.md new file mode 100644 index 00000000..4df4470b --- /dev/null +++ b/tests/python/fixtures/quant/README.md @@ -0,0 +1,23 @@ +These files freeze the `rabitqlib.quant` byte formats described in [docs/formats.md](../../../../docs/formats.md). They were written once by the code that introduced each format and are never regenerated: a test failure against them means a format changed, which needs a new format version rather than new fixtures. + +`rotator_fht_kac_dim64_seed0.bin` and `rotator_fht_kac_dim100_padded128_seed42.bin` are rotator blobs (kind 1, version 1) from `quant.Rotator(64, seed=0).to_bytes()` and `quant.Rotator(100, padded_dim=128, seed=42).to_bytes()`. The first body equals the frozen FHT/Kac sign flips for `(64, 64, seed 0)` in `tests/unit/rabitqlib/utils/rotator_test.cpp` (FNV-1a hash `0x7fcf0925975a6086`). + +`packed_codes_nbits{1..8}_n3_d96.bin` are `quant.pack_codes(codes, nbits=nbits).tobytes()` (packed-codes version 1) for three rows of 96 coordinates, which covers one full 64-coordinate block and a 32-coordinate tail, with + +```python +i, j = np.meshgrid(np.arange(3), np.arange(96), indexing="ij") +codes = ((i * 37 + j * 11 + (j // 32) * 5) % (1 << nbits)).astype(np.uint8) +``` + +`fastscan_blocks_n33_d64.bin` is `quant.pack_fastscan_blocks(bitsets, factors).tobytes()` (FastScan blocks version 1): two blocks, the second with one vector and 31 zero-padded lanes, with + +```python +i, byte = np.arange(33)[:, None], np.arange(8)[None, :] +bitsets = ((i * 37 + byte * 11 + 3) % 256).astype(np.uint8) +factors = np.zeros(33, quant.FACTORS_DTYPE) +factors["add"] = np.arange(33) * 0.5 +factors["rescale"] = -(np.arange(33) + 1) * 0.25 +factors["error"] = np.arange(33) / 8 +``` + +It was checked against an independent NumPy implementation of the layout in `docs/formats.md` before it was frozen. diff --git a/tests/python/fixtures/quant/fastscan_blocks_n33_d64.bin b/tests/python/fixtures/quant/fastscan_blocks_n33_d64.bin new file mode 100644 index 00000000..6fc6b2bd Binary files /dev/null and b/tests/python/fixtures/quant/fastscan_blocks_n33_d64.bin differ diff --git a/tests/python/fixtures/quant/packed_codes_nbits1_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits1_n3_d96.bin new file mode 100644 index 00000000..c147096d --- /dev/null +++ b/tests/python/fixtures/quant/packed_codes_nbits1_n3_d96.bin @@ -0,0 +1 @@ +UUUUUUUUUUUUUUUU \ No newline at end of file diff --git a/tests/python/fixtures/quant/packed_codes_nbits2_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits2_n3_d96.bin new file mode 100644 index 00000000..0a89ef2f --- /dev/null +++ b/tests/python/fixtures/quant/packed_codes_nbits2_n3_d96.bin @@ -0,0 +1 @@ +PPPPƥPPPPPPPPllllllll \ No newline at end of file diff --git a/tests/python/fixtures/quant/packed_codes_nbits3_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits3_n3_d96.bin new file mode 100644 index 00000000..3d32546e Binary files /dev/null and b/tests/python/fixtures/quant/packed_codes_nbits3_n3_d96.bin differ diff --git a/tests/python/fixtures/quant/packed_codes_nbits4_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits4_n3_d96.bin new file mode 100644 index 00000000..93f71e2a --- /dev/null +++ b/tests/python/fixtures/quant/packed_codes_nbits4_n3_d96.bin @@ -0,0 +1 @@ +;L];L]Հ;LՀ;L*Հ;L*Հ;LՀ;LՀ;L*Հ;L*Հ;L*Հ;L*Հ;L*Հ;L*Հ;L*Հ;L*Հ;L*Հ;*Հ; \ No newline at end of file diff --git a/tests/python/fixtures/quant/packed_codes_nbits5_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits5_n3_d96.bin new file mode 100644 index 00000000..4d907155 Binary files /dev/null and b/tests/python/fixtures/quant/packed_codes_nbits5_n3_d96.bin differ diff --git a/tests/python/fixtures/quant/packed_codes_nbits6_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits6_n3_d96.bin new file mode 100644 index 00000000..6943bc5b --- /dev/null +++ b/tests/python/fixtures/quant/packed_codes_nbits6_n3_d96.bin @@ -0,0 +1 @@ +@ ֡l7X#Dڥp;Qr=ȓʕegiJv\bNFzmfR~B}p;Ƒ\'Hީt?ʕ`+A b-٤o`kbmdo^@ԡPZ`ʕ`+LxΙd/Pqɔ_ofhԧb-4ģnD) \ No newline at end of file diff --git a/tests/python/fixtures/quant/packed_codes_nbits7_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits7_n3_d96.bin new file mode 100644 index 00000000..6d98628e --- /dev/null +++ b/tests/python/fixtures/quant/packed_codes_nbits7_n3_d96.bin @@ -0,0 +1 @@ +@ ֡l7X#Dڥp;Qr=ȓʕegiʒ6&ddmM*xm0.n̒dc.zOP>p;Ƒ\'Hީt?ʕ`+A b-٤o`kbmdo&dlMIٛo}1xGi='(U wA3rʕ`+LxΙd/Pqɔ_ofhMٛ6\Z컠rħ "z \ No newline at end of file diff --git a/tests/python/fixtures/quant/packed_codes_nbits8_n3_d96.bin b/tests/python/fixtures/quant/packed_codes_nbits8_n3_d96.bin new file mode 100644 index 00000000..1e6489ca Binary files /dev/null and b/tests/python/fixtures/quant/packed_codes_nbits8_n3_d96.bin differ diff --git a/tests/python/fixtures/quant/rotator_fht_kac_dim100_padded128_seed42.bin b/tests/python/fixtures/quant/rotator_fht_kac_dim100_padded128_seed42.bin new file mode 100644 index 00000000..7a29a3ce Binary files /dev/null and b/tests/python/fixtures/quant/rotator_fht_kac_dim100_padded128_seed42.bin differ diff --git a/tests/python/fixtures/quant/rotator_fht_kac_dim64_seed0.bin b/tests/python/fixtures/quant/rotator_fht_kac_dim64_seed0.bin new file mode 100644 index 00000000..bd983df5 Binary files /dev/null and b/tests/python/fixtures/quant/rotator_fht_kac_dim64_seed0.bin differ diff --git a/tests/python/test_quant_editor_support.py b/tests/python/test_quant_editor_support.py new file mode 100644 index 00000000..65462890 --- /dev/null +++ b/tests/python/test_quant_editor_support.py @@ -0,0 +1,143 @@ +"""Installed-package contracts for the rabitqlib.quant stub and runtime help.""" + +import ast +import doctest +import re +from pathlib import Path + +import pytest +import rabitqlib +from rabitqlib import quant + + +def installed_stub(): + stub = Path(rabitqlib.__file__).parent / "quant.pyi" + assert stub.is_file() + return ast.parse(stub.read_text()) + + +def normalized(text): + return " ".join(text.split()) + + +def stub_functions(body): + """Each name's first declaration; overloads share their parameters and docs.""" + functions = {} + for node in body: + if isinstance(node, ast.FunctionDef): + functions.setdefault(node.name, node) + return functions + + +def stub_parameters(node): + args = node.args + positional = [arg.arg for arg in args.posonlyargs + args.args if arg.arg != "self"] + defaults = { + arg.arg: ast.literal_eval(value) + for arg, value in zip(args.kwonlyargs, args.kw_defaults) + if value is not None + } + return positional, [arg.arg for arg in args.kwonlyargs], defaults + + +def runtime_parameters(member): + """The same, from pybind11's first doc line.""" + signature = member.__doc__.splitlines()[0] + names = [ + n for n in re.findall(r"(?:\(|, )([A-Za-z_]\w*):", signature) if n != "self" + ] + _, star, keyword_part = signature.partition(", *, ") + keyword = re.findall(r"(?:^|, )([A-Za-z_]\w*):", keyword_part) if star else [] + defaults = { + found[1]: ast.literal_eval(found[2]) + for found in re.finditer(r"\b([A-Za-z_]\w*): [^=,]+ = ([^,)]+)", signature) + if found[1] in keyword + } + return [n for n in names if n not in keyword], keyword, defaults + + +def test_module_docs_and_exports(): + tree = installed_stub() + assert "PEP 411" in quant.__doc__ + assert normalized(ast.get_docstring(tree)) in normalized(quant.__doc__) + exports = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "__all__" for t in node.targets) + ) + assert ast.literal_eval(exports.value) == quant.__all__ + declared = { + n.name for n in tree.body if isinstance(n, (ast.ClassDef, ast.FunctionDef)) + } + declared.update(n.target.id for n in tree.body if isinstance(n, ast.AnnAssign)) + assert set(quant.__all__) == declared + for name in quant.__all__: + assert hasattr(quant, name), name + + +@pytest.mark.parametrize( + "name", + [ + "padded_dim", + "quantize", + "split", + "code_factors", + "pack_codes", + "unpack_codes", + "query_terms", + "estimate_from_ip", + "pack_fastscan_blocks", + ], +) +def test_function_signatures_and_docs(name): + node = stub_functions(installed_stub().body)[name] + runtime = getattr(quant, name) + assert runtime_parameters(runtime) == stub_parameters(node), name + assert normalized(ast.get_docstring(node)) in normalized(runtime.__doc__), name + + +def test_classes_match_runtime(): + classes = {n.name: n for n in installed_stub().body if isinstance(n, ast.ClassDef)} + assert set(classes) == {"FormatError", "Query", "Rotator"} + for name, cls in classes.items(): + runtime = getattr(quant, name) + assert runtime.__module__ == "rabitqlib.quant" + assert normalized(ast.get_docstring(cls)) in normalized(runtime.__doc__) + checked = { + "Rotator": ("rotate", "from_bytes", "__init__"), + "Query": ( + "__init__", + "set_centroid", + "estimate_bitset", + "estimate_codes", + "refine", + "refine_packed", + "estimate_fastscan", + ), + } + for name, signatures in checked.items(): + cls = getattr(quant, name) + members = stub_functions(classes[name].body) + public = {key for key in vars(cls) if not key.startswith("_")} + assert public <= set(members), public - set(members) + for member, node in members.items(): + runtime = getattr(cls, member) + documentation = ast.get_docstring(node) + if documentation is not None: + assert normalized(documentation) in normalized(runtime.__doc__), member + if member in signatures: + assert runtime_parameters(runtime) == stub_parameters(node), member + + +def test_documentation_examples(): + runner = doctest.DocTestRunner() + parser = doctest.DocTestParser() + for name, documentation in ( + ("quant", quant.__doc__), + ("Rotator", quant.Rotator.__doc__), + ("Query", quant.Query.__doc__), + ): + runner.run(parser.get_doctest(documentation, {}, name, None, 0)) + assert runner.failures == 0 + assert runner.tries > 0 diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py new file mode 100644 index 00000000..d043870e --- /dev/null +++ b/tests/python/test_quant_layer1.py @@ -0,0 +1,403 @@ +"""Python boundary coverage for rabitqlib.quant layer 1: quantize, split, factors, packing.""" + +from pathlib import Path + +import numpy as np +import pytest +from rabitqlib import quant + +D = 96 # one full 64-coordinate block plus a 32-coordinate tail +EPSILON = 1.9 # folded into the error factor (kConstEpsilon) + + +def residuals(n=8, d=D, seed=0): + return np.random.default_rng(seed).standard_normal((n, d)).astype(np.float32) + + +def code_dtype(nbits): + return np.uint16 if nbits == 9 else np.uint8 + + +def reference_factors(r, codes, nbits, metric, centroid): + """The published factor formulas (docs/docs/rabitq/estimator.md) in float64.""" + r = r.astype(np.float64) + c = np.zeros(r.shape[1]) if centroid is None else centroid.astype(np.float64) + xu = codes.astype(np.float64) - (2**nbits - 1) / 2 + l2 = np.sum(r * r, axis=1) + ip_r = np.sum(r * xu, axis=1) + ip_c = xu @ c + # error = scale * sqrt((ratio - 1) / (d - 1)) with ratio = |r|^2 |xu|^2 / ^2. + ratio = l2 * np.sum(xu * xu, axis=1) / ip_r**2 + scale = (2 if metric == "l2" else 1) * EPSILON * np.sqrt(l2) + if metric == "l2": + return l2 + 2 * l2 * ip_c / ip_r, -2 * l2 / ip_r, ratio, scale + return 1 - r @ c + l2 * ip_c / ip_r, -l2 / ip_r, ratio, scale + + +def test_padded_dim(): + assert [quant.padded_dim(d) for d in (1, 32, 33, 100, 128)] == [ + 32, + 32, + 64, + 128, + 128, + ] + with pytest.raises(ValueError, match="dim must be positive"): + quant.padded_dim(0) + + +@pytest.mark.parametrize("nbits", range(1, 10)) +def test_quantize_codes_and_factors(nbits): + r = residuals() + codes = quant.quantize(r, nbits=nbits) + assert codes.dtype == code_dtype(nbits) and codes.shape == r.shape + assert codes.max() < 2**nbits + # The top bit of each code is the sign of its residual. + np.testing.assert_array_equal(codes >> (nbits - 1), (r > 0).astype(codes.dtype)) + centroid = residuals(1, seed=1)[0] + for metric in ("l2", "ip"): + factors = quant.code_factors( + r, codes, nbits=nbits, metric=metric, centroid=centroid + ) + assert factors.dtype == quant.FACTORS_DTYPE and factors.shape == (len(r),) + add, rescale, ratio, scale = reference_factors( + r, codes, nbits, metric, centroid + ) + np.testing.assert_allclose(factors["add"], add, rtol=1e-4, atol=1e-4) + np.testing.assert_allclose(factors["rescale"], rescale, rtol=1e-4) + # The float32 ratio - 1 cancels at high widths, so compare the ratio itself. + error = factors["error"].astype(np.float64) + np.testing.assert_allclose((error / scale) ** 2 * (D - 1) + 1, ratio, rtol=1e-5) + + +def test_fast_quantization_and_thread_counts(): + r = residuals(n=40) + exact = quant.quantize(r, nbits=5, num_threads=1) + np.testing.assert_array_equal(quant.quantize(r, nbits=5, num_threads=0), exact) + fast = quant.quantize(r, nbits=5, fast_quantization=True) + assert fast.dtype == np.uint8 and fast.max() < 32 + # The sign bit does not depend on the scale. + np.testing.assert_array_equal(fast >> 4, exact >> 4) + np.testing.assert_array_equal( + quant.quantize(r, nbits=5, fast_quantization=True, num_threads=3), fast + ) + + +@pytest.mark.parametrize("nbits", [2, 4, 9]) +def test_split_prefix_factors(nbits): + r = residuals() + codes = quant.quantize(r, nbits=nbits) + for prefix_bits in range(1, nbits): + prefix, ex = quant.split(codes, nbits=nbits, prefix_bits=prefix_bits) + ex_bits = nbits - prefix_bits + assert prefix.dtype == ex.dtype == np.uint8 + np.testing.assert_array_equal(prefix, codes >> ex_bits) + np.testing.assert_array_equal(ex, codes & (2**ex_bits - 1)) + # A prefix is itself a code of the prefix width for the same residual. + quant.code_factors(r, prefix, nbits=prefix_bits, metric="l2") + + +@pytest.mark.parametrize("nbits", [4, 9]) +def test_split_threads_and_out(nbits): + # 300 rows span several parallel ranges. + codes = quant.quantize(residuals(300, seed=12), nbits=nbits) + prefix, ex = quant.split(codes, nbits=nbits, prefix_bits=2, num_threads=1) + for threads in (0, 4): + parts = quant.split(codes, nbits=nbits, prefix_bits=2, num_threads=threads) + np.testing.assert_array_equal(parts[0], prefix) + np.testing.assert_array_equal(parts[1], ex) + out = (np.empty((300, D), np.uint8), np.empty((300, D), np.uint8)) + result = quant.split(codes, nbits=nbits, prefix_bits=2, num_threads=4, out=out) + assert result[0] is out[0] and result[1] is out[1] + np.testing.assert_array_equal(out[0], prefix) + np.testing.assert_array_equal(out[1], ex) + # Strided output rows are accepted, as for the other out= arguments. Overlap is + # judged by byte extent, so the two parts come from separate buffers. + wide = (np.zeros((300, 2 * D), np.uint8), np.zeros((300, 2 * D), np.uint8)) + quant.split(codes, nbits=nbits, prefix_bits=2, out=(wide[0][:, :D], wide[1][:, D:])) + np.testing.assert_array_equal(wide[0][:, :D], prefix) + np.testing.assert_array_equal(wide[1][:, D:], ex) + assert not wide[0][:, D:].any() and not wide[1][:, :D].any() + + +def test_split_rejects_bad_out_and_leaves_out_unchanged(): + codes = quant.quantize(residuals(300, seed=13), nbits=4) + good = np.zeros((300, D), np.uint8) + cases = [ + (np.zeros((300, D), np.uint8), r"out must be a \(prefix, ex\) tuple"), + ((good,), r"out must be a \(prefix, ex\) tuple"), + ([good, good.copy()], r"out must be a \(prefix, ex\) tuple"), + ((good, None), r"out must be a \(prefix, ex\) tuple"), + ((None, good), r"out must be a \(prefix, ex\) tuple"), + ( + (np.zeros((300, D), np.int8), good), + r"out\[0\] must be a writable NumPy array of dtype uint8 and shape \(300, 96\)", + ), + ( + (good, np.zeros((299, D), np.uint8)), + r"out\[1\] must be a writable NumPy array of dtype uint8", + ), + ((good, good), r"out\[1\] must not overlap out\[0\]"), + ((good, good[:, ::-1]), r"out\[1\] must be a writable NumPy array"), + ] + for out, message in cases: + with pytest.raises(ValueError, match=message): + quant.split(codes, nbits=4, prefix_bits=1, out=out) + shared = np.zeros((300, 2 * D), np.uint8) + view = shared[:, :D] + view[:] = codes + with pytest.raises(ValueError, match=r"out\[0\] must not overlap codes"): + quant.split( + view, nbits=4, prefix_bits=1, out=(shared[:, D // 2 : D + D // 2], good) + ) + # Codes are checked in a pass of their own, before out= is written. + bad = codes.copy() + bad[250, 3] = 16 + out = (np.zeros((300, D), np.uint8), np.zeros((300, D), np.uint8)) + with pytest.raises(ValueError, match=r"got 16 at \(250, 3\)$"): + quant.split(bad, nbits=4, prefix_bits=1, num_threads=4, out=out) + assert not out[0].any() and not out[1].any() + + +@pytest.mark.parametrize("nbits", range(1, 9)) +def test_pack_round_trip(nbits): + codes = np.random.default_rng(nbits).integers(0, 2**nbits, (5, D), dtype=np.uint8) + packed = quant.pack_codes(codes, nbits=nbits) + assert packed.dtype == np.uint8 and packed.shape == (5, D * nbits // 8) + np.testing.assert_array_equal(quant.unpack_codes(packed, nbits=nbits), codes) + if nbits == 1: + np.testing.assert_array_equal( + packed, np.packbits(codes, axis=1, bitorder="little") + ) + + +def test_empty_inputs(): + r = np.empty((0, D), dtype=np.float32) + codes = quant.quantize(r, nbits=4) + assert codes.shape == (0, D) + assert quant.code_factors(r, codes, nbits=4, metric="ip").shape == (0,) + assert quant.pack_codes(codes, nbits=4).shape == (0, D // 2) + assert quant.unpack_codes(np.empty((0, D // 2), np.uint8), nbits=4).shape == (0, D) + prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) + assert prefix.shape == ex.shape == (0, D) + + +def test_strided_noncontiguous_misaligned_and_memmapped_codes(tmp_path): + codes = np.random.default_rng(3).integers(0, 16, (6, 128), dtype=np.uint8) + expected = quant.pack_codes(codes, nbits=4) + # A column slice of a wider record buffer is read through its row stride. + records = np.zeros((6, 200), dtype=np.uint8) + records[:, 9 : 9 + 128] = codes + np.testing.assert_array_equal( + quant.pack_codes(records[:, 9 : 9 + 128], nbits=4), expected + ) + # Rows that are not contiguous are copied. + np.testing.assert_array_equal( + quant.pack_codes(codes[:, ::2], nbits=4), + quant.pack_codes(np.ascontiguousarray(codes[:, ::2]), nbits=4), + ) + np.testing.assert_array_equal( + quant.pack_codes(codes[::-1], nbits=4), expected[::-1] + ) + mapped = np.memmap( + tmp_path / "codes.bin", dtype=np.uint8, mode="w+", shape=codes.shape + ) + mapped[:] = codes + np.testing.assert_array_equal(quant.pack_codes(mapped, nbits=4), expected) + # Misaligned uint16 codes are copied before typed reads. + r = residuals() + wide = quant.quantize(r, nbits=9) + storage = np.zeros(wide.nbytes + 1, dtype=np.uint8) + misaligned = np.frombuffer(storage[1:].data, dtype=np.uint16).reshape(wide.shape) + misaligned[:] = wide + np.testing.assert_array_equal( + quant.code_factors(r, misaligned, nbits=9, metric="l2"), + quant.code_factors(r, wide, nbits=9, metric="l2"), + ) + + +def test_out_arguments(): + r = residuals() + out = np.empty(r.shape, dtype=np.uint8) + assert quant.quantize(r, nbits=4, out=out) is out + np.testing.assert_array_equal(out, quant.quantize(r, nbits=4)) + factors = np.empty(len(r), dtype=quant.FACTORS_DTYPE) + assert quant.code_factors(r, out, nbits=4, metric="l2", out=factors) is factors + packed = np.empty((len(r), D // 2), dtype=np.uint8) + assert quant.pack_codes(out, nbits=4, out=packed) is packed + unpacked = np.empty_like(out) + assert quant.unpack_codes(packed, nbits=4, out=unpacked) is unpacked + np.testing.assert_array_equal(unpacked, out) + # A row-strided output is written through its stride. + wide = np.zeros((len(r), 2 * D), dtype=np.uint8) + quant.quantize(r, nbits=4, out=wide[:, :D]) + np.testing.assert_array_equal(wide[:, :D], out) + + readonly = np.empty_like(out) + readonly.flags.writeable = False + for bad in ( + np.empty(r.shape, dtype=np.uint16), + np.empty((len(r), D + 32), dtype=np.uint8), + readonly, + np.empty((len(r), 2 * D), dtype=np.uint8)[:, ::2], + out.tolist(), + ): + with pytest.raises(ValueError, match="out must be a writable NumPy array"): + quant.quantize(r, nbits=4, out=bad) + with pytest.raises(ValueError, match="out must be a writable, C-contiguous"): + quant.code_factors( + r, out, nbits=4, metric="l2", out=np.empty(len(r), np.float32) + ) + + +def test_rejected_arguments(): + r = residuals() + codes = quant.quantize(r, nbits=4) + with pytest.raises(ValueError, match=r"nbits must be in \[1, 9\]; got 0"): + quant.quantize(r, nbits=0) + with pytest.raises(ValueError, match="positive multiple of 32"): + quant.quantize(r[:, :40], nbits=4) + for bad in (np.nan, np.inf): + nonfinite = r.copy() + nonfinite[2, 5] = bad + with pytest.raises( + ValueError, match="residuals must contain only finite values" + ): + quant.quantize(nonfinite, nbits=4) + with pytest.raises( + ValueError, match="residuals exceeds the safe float32 coordinate range" + ): + quant.quantize(np.full((1, D), 1e30, np.float32), nbits=4) + with pytest.raises(ValueError, match="metric must be 'l2' or 'ip'"): + quant.code_factors(r, codes, nbits=4, metric="cosine") + with pytest.raises(TypeError): + quant.code_factors(r, codes, nbits=4) # metric is required + with pytest.raises( + ValueError, match="codes must have dtype uint8; got dtype int64" + ): + quant.code_factors(r, codes.astype(np.int64), nbits=4, metric="l2") + with pytest.raises(ValueError, match="codes must have dtype uint8; got dtype bool"): + quant.pack_codes(codes > 7, nbits=1) + with pytest.raises(ValueError, match="codes must have dtype uint16"): + quant.code_factors(r, codes, nbits=9, metric="l2") + with pytest.raises(ValueError, match="codes must have the shape of residuals"): + quant.code_factors(r, codes[:4], nbits=4, metric="l2") + with pytest.raises( + ValueError, + match="codes row 0 for residuals row 0: Unpacked code exceeds the configured", + ): + quant.code_factors(r, codes, nbits=3, metric="l2") + with pytest.raises(ValueError, match=r"codes values must be below 2\*\*3 = 8"): + quant.pack_codes(codes, nbits=3) + with pytest.raises( + ValueError, match="codes row 0 for residuals row 0: code does not" + ): + quant.code_factors(-r, codes, nbits=4, metric="l2") + with pytest.raises(ValueError, match=r"centroid must have shape \(96,\)"): + quant.code_factors(r, codes, nbits=4, metric="l2", centroid=np.zeros(64)) + with pytest.raises(ValueError, match="centroid must contain only finite values"): + quant.code_factors(r, codes, nbits=4, metric="l2", centroid=np.full(D, np.nan)) + with pytest.raises(ValueError, match=r"prefix_bits must be in \[1, 3\]"): + quant.split(codes, nbits=4, prefix_bits=4) + with pytest.raises(ValueError, match=r"nbits must be in \[1, 8\]; got 9"): + quant.pack_codes(codes, nbits=9) + with pytest.raises( + ValueError, match="packed must have d \\* nbits / 8 bytes per row" + ): + quant.unpack_codes(np.zeros((2, 10), np.uint8), nbits=3) + with pytest.raises( + ValueError, match="fast_quantization requires a working dimension" + ): + quant.quantize( + np.zeros((1, 65568), np.float32), nbits=2, fast_quantization=True + ) + + +@pytest.mark.parametrize("nbits", range(1, 9)) +def test_golden_packed_codes(nbits): + """Packed-codes version 1 is frozen; see tests/python/fixtures/quant/README.md.""" + assert quant.PACKED_CODES_VERSION == 1 + fixtures = Path(__file__).resolve().parent / "fixtures" / "quant" + golden = (fixtures / f"packed_codes_nbits{nbits}_n3_d96.bin").read_bytes() + i, j = np.meshgrid(np.arange(3), np.arange(96), indexing="ij") + codes = ((i * 37 + j * 11 + (j // 32) * 5) % (1 << nbits)).astype(np.uint8) + assert quant.pack_codes(codes, nbits=nbits).tobytes() == golden + unpacked = quant.unpack_codes( + np.frombuffer(golden, np.uint8).reshape(3, -1), nbits=nbits + ) + np.testing.assert_array_equal(unpacked, codes) + + +def test_out_overlapping_an_input_is_rejected(): + rng = np.random.default_rng(7) + buffer = rng.standard_normal((5, 64)).astype(np.float32) + with pytest.raises(ValueError, match="out must not overlap residuals"): + quant.quantize(buffer[:4], nbits=2, out=buffer.view(np.uint8)[1:, :64]) + codes = quant.quantize(buffer, nbits=4) + with pytest.raises(ValueError, match="out must not overlap codes"): + quant.pack_codes(codes[:4], nbits=4, out=codes[1:, :32]) + packed = np.zeros((4, 64), np.uint8) + with pytest.raises(ValueError, match="out must not overlap packed"): + quant.unpack_codes(packed[:2, :32], nbits=4, out=packed[1:3]) + shared = np.zeros(5 * 64 + 15, np.float32) + x = shared[: 5 * 64].reshape(5, 64) + x[:] = buffer + out = shared[5 * 64 - 15 : 5 * 64].view(quant.FACTORS_DTYPE) + with pytest.raises(ValueError, match="out must not overlap residuals"): + quant.code_factors(x, codes, nbits=4, metric="l2", out=out) + # Disjoint caller buffers are accepted. + assert quant.quantize(buffer, nbits=4, out=np.empty_like(codes)) is not None + + +def test_code_factors_leaves_out_unchanged_on_rejection(): + rng = np.random.default_rng(8) + x = rng.standard_normal((6, 64)).astype(np.float32) + codes = quant.quantize(x, nbits=4) + codes[4] = 15 - codes[4] + out = np.zeros(6, quant.FACTORS_DTYPE) + with pytest.raises(ValueError, match="codes row 4 for residuals row 4"): + quant.code_factors(x, codes, nbits=4, metric="l2", out=out) + assert not out.view(np.float32).any() + + +def test_rejected_codes_leave_pack_out_unchanged(): + # The codes are checked in a pass of their own, before out= is written. + codes = quant.quantize(residuals(300, seed=11), nbits=4) + codes[200, 7] = 16 + out = np.zeros((300, D * 4 // 8), np.uint8) + with pytest.raises(ValueError, match=r"got 16 at \(200, 7\)$"): + quant.pack_codes(codes, nbits=4, num_threads=4, out=out) + assert not out.any() + + +def test_code_factors_names_the_lowest_rejected_row_across_threads(): + # 1000 rows span several parallel ranges; the lowest rejected row is reported. + x = residuals(1000, seed=10) + codes = quant.quantize(x, nbits=4) + codes[[700, 300]] = 15 - codes[[700, 300]] + for threads in (1, 4): + with pytest.raises( + ValueError, match="^codes row 300 for residuals row 300: code does not" + ): + quant.code_factors(x, codes, nbits=4, metric="l2", num_threads=threads) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_code_factors_reports_the_overflowing_row(metric): + # With every code 1, is half the residual's sum. The + # +-2**25 at coordinates 0 and 64 share an accumulator slot in every dot-product + # backend and cancel exactly, and coordinate 65 keeps its tiny value apart, so the + # sum is a positive 1e-30 and the rescale factor overflows float32. + x = residuals(6, seed=9) + codes = quant.quantize(x, nbits=1) + for row in (3, 5): + x[row] = 0 + x[row, [0, 64, 65]] = [2.0**25, -(2.0**25), 1e-30] + codes[row] = 1 + out = np.zeros(6, quant.FACTORS_DTYPE) + with pytest.raises( + OverflowError, + match="^codes row 3 for residuals row 3: Quantization factors are not " + "representable in float32$", + ): + quant.code_factors(x, codes, nbits=1, metric=metric, num_threads=4, out=out) diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py new file mode 100644 index 00000000..6c0f541f --- /dev/null +++ b/tests/python/test_quant_layer2.py @@ -0,0 +1,581 @@ +"""Python boundary coverage for quant layer 2: Query, estimates and refinement.""" + +import copy +import gc +import pickle +from concurrent.futures import ThreadPoolExecutor +from threading import Event + +import numpy as np +import pytest +from rabitqlib import quant + +BUSY = "Query is busy: set_centroid cannot overlap another call" + + +def rows(n, d, seed=0, scale=1.0): + rng = np.random.default_rng(seed) + return (rng.standard_normal((n, d)) * scale).astype(np.float32) + + +def setup(metric, nbits, *, n=40, d=96, centroid=True, seed=0): + """Codes, factors and a query for vectors x quantized around centroid c.""" + x = rows(n, d, seed) + c = rows(1, d, seed + 1, 0.3)[0] if centroid else None + residuals = x - (c if centroid else 0) + codes = quant.quantize(residuals, nbits=nbits) + factors = quant.code_factors( + residuals, codes, nbits=nbits, metric=metric, centroid=c + ) + q = rows(1, d, seed + 2)[0] + return x, c, residuals, codes, factors, q + + +def oracle(q, c, codes, factors, nbits, metric): + """The reference formula in float64.""" + q64 = q.astype(np.float64) + c64 = np.zeros_like(q64) if c is None else c.astype(np.float64) + ip = codes.astype(np.float64) @ q64 + g_add = ((q64 - c64) ** 2).sum() if metric == "l2" else -(q64 @ c64) + g_error = np.sqrt(((q64 - c64) ** 2).sum()) + k = -0.5 * q64.sum() * (2**nbits - 1) + distance = factors["add"] + g_add + factors["rescale"] * (ip + k) + return distance, factors["error"] * g_error, ip + + +def close(actual, expected, scale): + np.testing.assert_allclose(actual, expected, rtol=0, atol=2e-5 * scale) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("nbits", range(1, 10)) +@pytest.mark.parametrize("centroid", [True, False]) +def test_estimate_codes_matches_the_formula(metric, nbits, centroid): + _, c, _, codes, factors, q = setup(metric, nbits, centroid=centroid) + query = quant.Query(q, metric=metric, centroid=c) + estimates = query.estimate_codes(codes, factors, nbits=nbits) + assert estimates.dtype == quant.ESTIMATE_DTYPE and estimates.shape == (40,) + distance, error, ip = oracle(q, c, codes, factors, nbits, metric) + scale = np.abs(factors["rescale"]).max() * np.abs(q).sum() * 2**nbits + close(estimates["distance"], distance, scale) + close(estimates["refine_input"], ip, np.abs(q).sum() * 2**nbits) + np.testing.assert_allclose(estimates["error"], error, rtol=1e-5) + terms = quant.query_terms(q, metric=metric, centroid=c) + assert isinstance(terms, np.void) and terms.dtype == quant.QUERY_TERMS_DTYPE + from_ip = quant.estimate_from_ip( + estimates["refine_input"], factors, terms, nbits=nbits + ) + np.testing.assert_array_equal(from_ip, estimates) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_estimates_track_true_distances(metric): + x, c, _, codes, factors, q = setup(metric, 7, n=200, d=128) + estimates = quant.Query(q, metric=metric, centroid=c).estimate_codes( + codes, factors, nbits=7 + ) + # IP folds the 1 of 1 - into the factors, not into g_add. + true = ((x - q) ** 2).sum(axis=1) if metric == "l2" else 1 - x @ q + assert np.corrcoef(estimates["distance"], true)[0, 1] > 0.999 + # The estimate is unbiased: a missing or doubled 1 would shift it by 1. + scale = np.linalg.norm(q) * np.linalg.norm(x - c, axis=1).mean() + difference = estimates["distance"] - true + assert abs(difference.mean()) < 1e-3 * scale + assert np.abs(difference).mean() < 4e-3 * scale + terms = quant.query_terms(q, metric=metric, centroid=c) + expected_add = ((q - c) ** 2).sum() if metric == "l2" else -(q @ c) + np.testing.assert_allclose(terms["g_add"], expected_add, rtol=1e-5) + np.testing.assert_allclose(terms["g_error"], np.linalg.norm(q - c), rtol=1e-5) + np.testing.assert_allclose(terms["g_k1xsumq"], -0.5 * q.sum(), rtol=1e-5) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("prefix_bits", [1, 2]) +def test_two_stage_scoring_with_one_query(metric, prefix_bits): + nbits = 5 + _, c, residuals, codes, f_full, q = setup(metric, nbits) + prefix, ex = quant.split(codes, nbits=nbits, prefix_bits=prefix_bits) + ex_bits = nbits - prefix_bits + f_prefix = quant.code_factors( + residuals, prefix, nbits=prefix_bits, metric=metric, centroid=c + ) + query = quant.Query(q, metric=metric, centroid=c) + first = query.estimate_codes(prefix, f_prefix, nbits=prefix_bits) + full = query.estimate_codes(codes, f_full, nbits=nbits) + refined = query.refine(ex, f_full, first, ex_bits=ex_bits, prefix_bits=prefix_bits) + packed = query.refine_packed( + quant.pack_codes(ex, nbits=ex_bits), + f_full, + first, + ex_bits=ex_bits, + prefix_bits=prefix_bits, + ) + assert refined.dtype == np.float32 and refined.shape == (40,) + scale = np.abs(f_full["rescale"]).max() * np.abs(q).sum() * 2**nbits + close(refined, full["distance"], scale) + close(packed, refined, scale) + if prefix_bits == 1: + bitset = query.estimate_bitset(quant.pack_codes(prefix, nbits=1), f_prefix) + close(bitset["distance"], first["distance"], scale) + np.testing.assert_array_equal(bitset["error"], first["error"]) + # Scoring several widths in turn does not change the query. + fresh = quant.Query(q, metric=metric, centroid=c) + np.testing.assert_array_equal( + fresh.estimate_codes(codes, f_full, nbits=nbits), full + ) + np.testing.assert_array_equal( + query.estimate_codes(prefix, f_prefix, nbits=prefix_bits), first + ) + + +def test_indices_read_rows_in_place(): + _, c, residuals, codes, factors, q = setup("l2", 4) + prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) + f_prefix = quant.code_factors(residuals, prefix, nbits=1, metric="l2", centroid=c) + bitsets = quant.pack_codes(prefix, nbits=1) + query = quant.Query(q, metric="l2", centroid=c) + everything = query.estimate_bitset(bitsets, f_prefix) + chosen = np.array([7, 0, 39, 7], dtype=np.uint32) + some = query.estimate_bitset(bitsets, f_prefix, indices=chosen) + np.testing.assert_array_equal(some, everything[chosen]) + np.testing.assert_array_equal( + query.estimate_codes(codes, factors, nbits=4, indices=[3, 5]), + query.estimate_codes(codes, factors, nbits=4)[[3, 5]], + ) + refined_all = query.refine(ex, factors, everything, ex_bits=3) + refined = query.refine(ex, factors, everything[chosen], ex_bits=3, indices=chosen) + np.testing.assert_array_equal(refined, refined_all[chosen]) + packed = quant.pack_codes(ex, nbits=3) + np.testing.assert_array_equal( + query.refine_packed( + packed, factors, everything[chosen], ex_bits=3, indices=chosen + ), + query.refine_packed(packed, factors, everything, ex_bits=3)[chosen], + ) + assert query.estimate_bitset(bitsets, f_prefix, indices=[]).shape == (0,) + for bad, message in ( + (np.array([True, False]), "indices must contain integers"), + (np.array([1.0, 2.0]), "indices must contain integers"), + ([40], r"indices must be in \[0, 40\); got 40 at position 0"), + ([0, -1], r"indices must be in \[0, 40\); got -1 at position 1"), + ): + with pytest.raises(ValueError, match=message): + query.estimate_bitset(bitsets, f_prefix, indices=bad) + with pytest.raises(ValueError, match="first must have one record per entry"): + query.refine(ex, factors, everything, ex_bits=3, indices=chosen) + + +def test_set_centroid_copy_and_pickle(): + _, c, _, codes, factors, q = setup("ip", 4) + query = quant.Query(q, metric="ip") + assert repr(query) == "Query(dim=96, metric='ip', centroid=None)" + assert (query.dim, query.metric) == (96, "ip") + query.set_centroid(c) + assert repr(query) == "Query(dim=96, metric='ip', centroid=set)" + expected = quant.Query(q, metric="ip", centroid=c).estimate_codes( + codes, factors, nbits=4 + ) + np.testing.assert_array_equal( + query.estimate_codes(codes, factors, nbits=4), expected + ) + copies = [ + query.copy(), + copy.copy(query), + copy.deepcopy(query), + pickle.loads(pickle.dumps(query)), + ] + assert all(other is not query for other in copies) + del query + gc.collect() + for other in copies: + np.testing.assert_array_equal( + other.estimate_codes(codes, factors, nbits=4), expected + ) + # Copies are independent. + copies[0].set_centroid(None) + np.testing.assert_array_equal( + copies[1].estimate_codes(codes, factors, nbits=4), expected + ) + np.testing.assert_array_equal( + copies[0].estimate_codes(codes, factors, nbits=4), + quant.Query(q, metric="ip").estimate_codes(codes, factors, nbits=4), + ) + # The caller's vectors are copied, not referenced. + q_mutable = q.copy() + owned = quant.Query(q_mutable, metric="ip", centroid=c) + q_mutable[:] = 0 + np.testing.assert_array_equal( + owned.estimate_codes(codes, factors, nbits=4), expected + ) + + +class PausedArray: + """Hold an operation open deterministically during Python input conversion.""" + + def __init__(self, array): + self.array = array + self.entered = Event() + self.resume = Event() + + def __array__(self, dtype=None, copy=None): + self.entered.set() + if not self.resume.wait(timeout=15): + raise RuntimeError("test did not release input conversion") + return np.array(self.array, dtype=dtype, copy=True) + + +def test_set_centroid_during_scoring_is_rejected(): + _, c, _, codes, factors, q = setup("l2", 4) + query = quant.Query(q, metric="l2", centroid=c) + expected = query.estimate_codes(codes, factors, nbits=4) + paused = PausedArray(codes) + with ThreadPoolExecutor(max_workers=1) as pool: + running = pool.submit(query.estimate_codes, paused, factors, nbits=4) + try: + assert paused.entered.wait(timeout=10) + np.testing.assert_array_equal( + query.estimate_codes(codes, factors, nbits=4), expected + ) + assert repr(query.copy()) == repr(query) + with pytest.raises(RuntimeError, match=f"^{BUSY}$"): + query.set_centroid(None) + finally: + paused.resume.set() + np.testing.assert_array_equal(running.result(timeout=10), expected) + query.set_centroid(None) + + +def test_threads_out_and_empty_inputs(): + _, c, _, codes, factors, q = setup("l2", 8, n=300) + query = quant.Query(q, metric="l2", centroid=c) + single = query.estimate_codes(codes, factors, nbits=8) + np.testing.assert_array_equal( + query.estimate_codes(codes, factors, nbits=8, num_threads=4), single + ) + out = np.empty(300, quant.ESTIMATE_DTYPE) + assert query.estimate_codes(codes, factors, nbits=8, out=out) is out + np.testing.assert_array_equal(out, single) + prefix, ex = quant.split(codes, nbits=8, prefix_bits=1) + distances = np.empty(300, np.float32) + assert query.refine(ex, factors, single, ex_bits=7, out=distances) is distances + terms = quant.query_terms(q, metric="l2", centroid=c) + estimates = np.empty(300, quant.ESTIMATE_DTYPE) + assert ( + quant.estimate_from_ip( + single["refine_input"], factors, terms, nbits=8, out=estimates + ) + is estimates + ) + # Strided factors are read correctly. + doubled = np.repeat(factors, 2)[::2] + np.testing.assert_array_equal(query.estimate_codes(codes, doubled, nbits=8), single) + empty_codes = np.empty((0, 96), np.uint8) + empty_factors = np.empty(0, quant.FACTORS_DTYPE) + assert query.estimate_codes(empty_codes, empty_factors, nbits=8).shape == (0,) + assert query.estimate_bitset(np.empty((0, 12), np.uint8), empty_factors).shape == ( + 0, + ) + assert query.refine( + empty_codes, empty_factors, np.empty(0, quant.ESTIMATE_DTYPE), ex_bits=7 + ).shape == (0,) + + +def test_rejected_arguments(): + _, c, _, codes, factors, q = setup("l2", 4) + prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) + query = quant.Query(q, metric="l2", centroid=c) + first = query.estimate_codes(codes, factors, nbits=4) + cases = [ + (lambda: quant.Query(q[:40], metric="l2"), "rotated_query must have a length"), + (lambda: quant.Query(q, metric="cosine"), "metric"), + ( + lambda: quant.Query(q, metric="l2", centroid=c[:32]), + "centroid must have shape", + ), + ( + lambda: quant.Query(np.full(96, np.nan, np.float32), metric="l2"), + "rotated_query must contain only finite values", + ), + ( + lambda: query.estimate_codes(codes, factors, nbits=10), + r"nbits must be in \[1, 9\]", + ), + ( + lambda: query.estimate_codes(codes[:, :64], factors, nbits=4), + "codes must have d = 96 columns", + ), + ( + lambda: query.estimate_codes(codes.astype(np.uint16), factors, nbits=4), + "codes must have dtype uint8", + ), + ( + lambda: query.estimate_codes(codes, factors, nbits=3), + r"codes values must be below 2\*\*3", + ), + ( + lambda: query.estimate_codes(codes, factors[:5], nbits=4), + "factors must have one record per row of codes, 40", + ), + ( + lambda: query.estimate_codes(codes, first, nbits=4), + "factors must be a 1D NumPy array of dtype quant.FACTORS_DTYPE", + ), + ( + lambda: query.estimate_codes( + codes, factors.view(np.float32).reshape(-1, 3), nbits=4 + ), + "factors must be a 1D NumPy array of dtype quant.FACTORS_DTYPE", + ), + ( + lambda: query.estimate_bitset(codes, factors), + "bitsets must have d // 8 = 12 columns", + ), + ( + lambda: query.refine(ex, factors, first, ex_bits=2), + r"ex values must be below 2\*\*2", + ), + ( + lambda: query.refine(ex, factors, first, ex_bits=9), + r"ex_bits must be in \[1, 8\]", + ), + ( + lambda: query.refine(ex, factors, first, ex_bits=3, prefix_bits=0), + r"prefix_bits must be in \[1, 8\]", + ), + ( + lambda: query.refine(ex, factors, first[:3], ex_bits=3), + "first must have one record per row of ex, 40", + ), + ( + lambda: query.refine(ex, factors, factors, ex_bits=3), + "first must be a 1D NumPy array of dtype quant.ESTIMATE_DTYPE", + ), + ( + lambda: query.refine_packed(ex, factors, first, ex_bits=3), + r"packed must have d \* ex_bits // 8 = 36 columns", + ), + ( + lambda: query.estimate_codes( + codes, factors, nbits=4, out=np.empty(40, quant.FACTORS_DTYPE) + ), + "out must be a writable, C-contiguous NumPy array of dtype " + r"quant.ESTIMATE_DTYPE and shape \(40,\)", + ), + ( + lambda: query.refine(ex, factors, first, ex_bits=3, out=np.empty(40)), + "out must be a writable, C-contiguous NumPy array of dtype float32", + ), + ( + lambda: quant.estimate_from_ip( + np.zeros(40, np.float32), factors, first[0], nbits=4 + ), + "terms must be one record of dtype quant.QUERY_TERMS_DTYPE", + ), + ( + lambda: quant.estimate_from_ip( + np.full(40, np.inf, np.float32), + factors, + quant.query_terms(q, metric="l2"), + nbits=4, + ), + "ip must contain only finite values", + ), + ] + for call, message in cases: + with pytest.raises(ValueError, match=message): + call() + broken = factors.copy() + broken["rescale"][3] = np.nan + with pytest.raises(ValueError, match="factors must contain only finite values"): + query.estimate_codes(codes, broken, nbits=4) + # A non-finite factor outside the selected rows is not read. + query.estimate_codes(codes, broken, nbits=4, indices=[0, 1]) + + +def test_out_overlapping_an_input_is_rejected(): + _, c, _, codes, factors, q = setup("l2", 4) + prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) + query = quant.Query(q, metric="l2", centroid=c) + shared = np.zeros(80, quant.ESTIMATE_DTYPE) + shared_factors = shared.view(quant.FACTORS_DTYPE)[:40] + shared_factors[:] = factors + with pytest.raises(ValueError, match="out must not overlap factors"): + query.estimate_codes(codes, shared_factors, nbits=4, out=shared[20:60]) + first = query.estimate_codes(codes, factors, nbits=4) + with pytest.raises(ValueError, match="out must not overlap first"): + query.refine(ex, factors, first, ex_bits=3, out=first.view(np.float32)[:40]) + terms = quant.query_terms(q, metric="l2", centroid=c) + shared_ip = np.zeros(3 * 40, np.float32) + ip = shared_ip[:40] + ip[:] = first["refine_input"] + with pytest.raises(ValueError, match="out must not overlap ip"): + quant.estimate_from_ip( + ip, factors, terms, nbits=4, out=shared_ip.view(quant.ESTIMATE_DTYPE) + ) + + +@pytest.mark.parametrize( + "field,message", + [ + ("rescale", "Estimated distance is not representable in float32"), + ("error", "Estimated error is not representable in float32"), + ], +) +def test_estimates_that_overflow_are_rejected(field, message): + # With q all ones and every code 1, each centered inner product is 48 and |q| is + # sqrt(96), so a factor of 3e38 overflows that row's distance or error. + n, d = 300, 96 + q = np.ones(d, np.float32) + query = quant.Query(q, metric="l2") + codes = np.ones((n, d), np.uint8) + bitsets = quant.pack_codes(codes, nbits=1) + factors = np.zeros(n, quant.FACTORS_DTYPE) + factors[field][[250, 120]] = 3e38 + terms = quant.query_terms(q, metric="l2") + blocks = quant.pack_fastscan_blocks(bitsets, factors) + calls = [ + lambda out: query.estimate_bitset(bitsets, factors, num_threads=4, out=out), + lambda out: query.estimate_codes( + codes, factors, nbits=1, num_threads=4, out=out + ), + lambda out: quant.estimate_from_ip( + np.full(n, d, np.float32), factors, terms, nbits=1, num_threads=4, out=out + ), + lambda out: query.estimate_fastscan( + blocks, num_points=n, num_threads=4, out=out + ), + ] + for call in calls: + # After OverflowError the contents of out= are unspecified. + with pytest.raises(OverflowError, match=f"^{message}$"): + call(np.zeros(n, quant.ESTIMATE_DTYPE)) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_refined_distances_that_overflow_are_rejected(packed): + # Prefix and extra bits all 1 give total codes 3, centered at 1.5: the inner product + # with q all ones is 144, so a rescale of 3e38 overflows the distance. + n, d = 300, 96 + query = quant.Query(np.ones(d, np.float32), metric="l2") + ex = np.ones((n, d), np.uint8) + factors = np.zeros(n, quant.FACTORS_DTYPE) + factors["rescale"][[250, 120]] = 3e38 + first = np.zeros(n, quant.ESTIMATE_DTYPE) + first["refine_input"] = d + refine = query.refine_packed if packed else query.refine + codes = quant.pack_codes(ex, nbits=1) if packed else ex + with pytest.raises( + OverflowError, match="^Estimated distance is not representable in float32$" + ): + refine( + codes, factors, first, ex_bits=1, num_threads=4, out=np.zeros(n, np.float32) + ) + + +def test_rejections_name_the_lowest_failing_row_across_threads(): + _, c, residuals, codes, factors, q = setup("l2", 4, n=1000) + query = quant.Query(q, metric="l2", centroid=c) + bad = codes.copy() + bad[700, 5] = 16 + bad[300, 9] = 16 + prefix, _ = quant.split(codes, nbits=4, prefix_bits=1) + bitsets = quant.pack_codes(prefix, nbits=1) + broken = quant.code_factors(residuals, prefix, nbits=1, metric="l2", centroid=c) + broken["error"][[800, 500]] = np.nan + out = np.zeros(1000, quant.ESTIMATE_DTYPE) + for threads in (1, 4): + with pytest.raises( + ValueError, + match=r"^codes values must be below 2\*\*4 = 16; got 16 at \(300, 9\)$", + ): + query.estimate_codes(bad, factors, nbits=4, num_threads=threads, out=out) + # The vectorized range check fails and the row checks name the cause. + with pytest.raises( + ValueError, match="^factors must contain only finite values$" + ): + query.estimate_bitset(bitsets, broken, num_threads=threads, out=out) + # A rejected input leaves out= unchanged. + assert not out.view(np.float32).any() + # With indices=, the first failing entry is reported, not the lowest row. + with pytest.raises(ValueError, match=r"got 16 at \(700, 5\)$"): + query.estimate_codes(bad, factors, nbits=4, indices=[700, 300]) + + +def test_rejected_inputs_leave_out_unchanged(): + # Inputs are checked in a pass of their own, before out= is written; the failing + # rows lie past the first parallel range. + _, c, _, codes, factors, q = setup("l2", 4, n=300) + query = quant.Query(q, metric="l2", centroid=c) + _, ex = quant.split(codes, nbits=4, prefix_bits=1) + first = query.estimate_codes(codes, factors, nbits=4) + terms = quant.query_terms(q, metric="l2", centroid=c) + broken = factors.copy() + broken["add"][200] = np.nan + bad_ex = ex.copy() + bad_ex[100, 3] = 8 + bad_first = first.copy() + bad_first["refine_input"][50] = np.nan + bad_ip = first["refine_input"].copy() + bad_ip[150] = np.inf + distances = np.zeros(300, np.float32) + estimates = np.zeros(300, quant.ESTIMATE_DTYPE) + finite = "factors must contain only finite values" + cases = [ + ( + lambda: query.refine( + bad_ex, factors, first, ex_bits=3, num_threads=4, out=distances + ), + r"ex values must be below 2\*\*3 = 8; got 8 at \(100, 3\)", + ), + ( + lambda: query.refine( + ex, broken, first, ex_bits=3, num_threads=4, out=distances + ), + finite, + ), + ( + lambda: query.refine( + ex, factors, bad_first, ex_bits=3, num_threads=4, out=distances + ), + r"first\['refine_input'\] must contain only finite values", + ), + ( + lambda: query.refine_packed( + quant.pack_codes(ex, nbits=3), + broken, + first, + ex_bits=3, + num_threads=4, + out=distances, + ), + finite, + ), + ( + lambda: quant.estimate_from_ip( + bad_ip, factors, terms, nbits=4, num_threads=4, out=estimates + ), + "ip must contain only finite values", + ), + ( + lambda: quant.estimate_from_ip( + first["refine_input"], + broken, + terms, + nbits=4, + num_threads=4, + out=estimates, + ), + finite, + ), + ] + for call, message in cases: + with pytest.raises(ValueError, match=message): + call() + assert not distances.any() + assert not estimates.view(np.float32).any() + + +def test_corrupt_pickle_state_is_rejected(): + query = quant.Query.__new__(quant.Query) + with pytest.raises(ValueError, match="Invalid Query pickle state"): + query.__setstate__((np.zeros(64, np.float32), 5, None)) diff --git a/tests/python/test_quant_layer3.py b/tests/python/test_quant_layer3.py new file mode 100644 index 00000000..8083570b --- /dev/null +++ b/tests/python/test_quant_layer3.py @@ -0,0 +1,218 @@ +"""Python boundary coverage for quant layer 3: FastScan blocks and estimates.""" + +import copy +import gc +import pickle +from pathlib import Path + +import numpy as np +import pytest +from rabitqlib import quant + +FIXTURES = Path(__file__).resolve().parent / "fixtures" / "quant" +PERM = [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15] + + +def reference_blocks(bitsets, factors): + """The FastScan blocks version 1 layout of docs/formats.md, in NumPy.""" + n, cols = bitsets.shape + msb_first = np.packbits( + np.unpackbits(bitsets, axis=1, bitorder="little"), axis=1, bitorder="big" + ) + blocks = [] + for b in range(-(-n // 32)): + count = min(32, n - 32 * b) + lanes = np.zeros((32, cols), np.uint8) + lanes[:count] = msb_first[32 * b : 32 * b + count] + codes = np.zeros((cols, 32), np.uint8) + for j, p in enumerate(PERM): + a = lanes[p].astype(np.uint16) + c = lanes[p + 16].astype(np.uint16) + codes[:, j] = (a >> 4) | (c & 0xF0) + codes[:, j + 16] = ((a & 0x0F) | (c << 4)) & 0xFF + record = np.zeros((3, 32), " None: + rotator = quant.Rotator(64, method="fht_kac", seed=0) + assert_type(rotator.padded_dim, int) + assert_type(rotator.seed, int | None) + assert_type(rotator.method, Literal["fht_kac", "matrix"]) + residuals = rotator.rotate(x, num_threads=1) + assert_type(residuals, npt.NDArray[np.float32]) + assert_type(quant.Rotator.from_bytes(rotator.to_bytes()), quant.Rotator) + codes = quant.quantize(residuals, nbits=4) + assert_type(codes, npt.NDArray[np.uint8]) + assert_type(quant.quantize(residuals, nbits=9), npt.NDArray[np.uint16]) + prefix, ex = quant.split(codes, nbits=4, prefix_bits=1) + assert_type(prefix, npt.NDArray[np.uint8]) + assert_type( + quant.split(codes, nbits=4, prefix_bits=1, num_threads=2, out=(prefix, ex)), + tuple[npt.NDArray[np.uint8], npt.NDArray[np.uint8]], + ) + factors = quant.code_factors(residuals, codes, nbits=4, metric="l2") + assert_type(factors, npt.NDArray[np.void]) + assert_type(quant.pack_codes(ex, nbits=3), npt.NDArray[np.uint8]) + assert_type( + quant.unpack_codes(quant.pack_codes(ex, nbits=3), nbits=3), + npt.NDArray[np.uint8], + ) + assert_type(quant.padded_dim(100), int) + assert_type(quant.FACTORS_DTYPE, np.dtype[np.void]) + assert issubclass(quant.FormatError, ValueError) + f_prefix = quant.code_factors(residuals, prefix, nbits=1, metric="l2") + query = quant.Query(residuals[0], metric="l2", centroid=None) + assert_type(query.metric, Literal["l2", "ip"]) + bitsets = quant.pack_codes(prefix, nbits=1) + first = query.estimate_bitset(bitsets, f_prefix, indices=[0], num_threads=1) + assert_type(first, npt.NDArray[np.void]) + assert_type(query.estimate_codes(codes, factors, nbits=4), npt.NDArray[np.void]) + assert_type( + query.refine(ex, factors, first, ex_bits=3, indices=[0]), + npt.NDArray[np.float32], + ) + assert_type( + query.refine_packed(quant.pack_codes(ex, nbits=3), factors, first, ex_bits=3), + npt.NDArray[np.float32], + ) + query.set_centroid(None) + assert_type(query.copy(), quant.Query) + terms = quant.query_terms(residuals[0], metric="ip") + assert_type(terms, np.void) + assert_type( + quant.estimate_from_ip(first["refine_input"], f_prefix, terms, nbits=1), + npt.NDArray[np.void], + ) + blocks = quant.pack_fastscan_blocks(bitsets, f_prefix, num_threads=0) + assert_type(blocks, npt.NDArray[np.uint8]) + assert_type( + query.estimate_fastscan(blocks, num_points=len(bitsets)), npt.NDArray[np.void] + ) + assert_type(quant.FASTSCAN_BLOCKS_VERSION, int) + quant.Query(residuals[0], metric="cosine") # type: ignore[arg-type] + quant.code_factors(residuals, codes, nbits=4, metric="cosine") # type: ignore[arg-type] + quant.quantize(residuals, 4) # type: ignore[call-overload] + quant.Rotator(64, method="hadamard") # type: ignore[arg-type] diff --git a/tests/unit/rabitqlib/quantization/factor_test.cpp b/tests/unit/rabitqlib/quantization/factor_test.cpp index a297e494..184146ba 100644 --- a/tests/unit/rabitqlib/quantization/factor_test.cpp +++ b/tests/unit/rabitqlib/quantization/factor_test.cpp @@ -89,6 +89,74 @@ TEST(RabitqFactorTest, CollinearMetricFactorsFloat) { check_collinear_factors(); } +template +void check_rejected_foreign_code() { + constexpr size_t kDim = 64; + std::vector residual(kDim), centroid(kDim, T{0}); + std::vector codes(kDim); + T add, rescale, error; + for (const auto metric : {METRIC_L2, METRIC_IP}) { + // A code of the opposite signs, then one orthogonal to the residual: the + // factor derivation divides by , which must be > 0. + for (const bool orthogonal : {false, true}) { + for (size_t i = 0; i < kDim; ++i) { + residual[i] = orthogonal && i % 2 == 1 ? T{-1} : T{1}; + codes[i] = orthogonal ? 1 : 0; + } + try { + rabitq_impl::xy_bits::code_factors( + residual.data(), + centroid.data(), + kDim, + codes.data(), + -0.5F, + add, + rescale, + error, + metric + ); + FAIL() << "Expected rejection of a code that does not quantize residual"; + } catch (const std::invalid_argument& exception) { + EXPECT_STREQ( + exception.what(), + "code does not quantize residual: must be " + "positive" + ); + } + } + } +} + +TEST(RabitqFactorTest, NanInputIsNotCheckedByTheEncoders) { + // The per-vector encoders do not check that inputs are finite (quantizer.md); the + // rejection of a nonpositive divisor must not turn a NaN into an exception. + constexpr size_t kDim = 64; + std::vector data(kDim, 0.5F), centroid(kDim, 0), residual(kDim), centered(kDim); + data[7] = std::numeric_limits::quiet_NaN(); + std::vector codes(kDim); + float add, rescale, error; + rabitq_impl::one_bit::one_bit_code_with_factor_into( + data.data(), + centroid.data(), + kDim, + codes.data(), + add, + rescale, + error, + METRIC_L2, + residual.data(), + centered.data() + ); + EXPECT_TRUE(std::isnan(add)); + EXPECT_TRUE(std::isnan(rescale)); +} + +TEST(RabitqFactorTest, ForeignCodeIsRejectedFloat) { check_rejected_foreign_code(); } + +TEST(RabitqFactorTest, ForeignCodeIsRejectedDouble) { + check_rejected_foreign_code(); +} + TEST(RabitqFactorTest, InvalidMetricMessagesArePreserved) { constexpr size_t kDim = 64; const auto invalid_metric = static_cast(2); diff --git a/tests/unit/rabitqlib/quantization/provisional_quantize_test.cpp b/tests/unit/rabitqlib/quantization/provisional_quantize_test.cpp new file mode 100644 index 00000000..8edc0f7f --- /dev/null +++ b/tests/unit/rabitqlib/quantization/provisional_quantize_test.cpp @@ -0,0 +1,362 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/provisional/quantize.hpp" + +namespace { +namespace api = rabitqlib::quant::v1; +namespace provisional = rabitqlib::quant::provisional; +namespace rabitq_impl = rabitqlib::quant::rabitq_impl; +using rabitqlib::METRIC_IP; +using rabitqlib::METRIC_L2; +using rabitqlib::MetricType; + +std::vector gaussian(size_t dim, std::mt19937& generator) { + std::normal_distribution distribution(0.0F, 1.0F); + std::vector values(dim); + for (auto& value : values) { + value = distribution(generator); + } + return values; +} + +struct Encoded { + std::vector codes; + api::Factors factors; +}; + +// quant::v1::encode with UnpackedU16 codes, the reference for codes and factors. +Encoded v1_encode( + const std::vector& input, + const std::vector& centroid, + size_t total_bits, + MetricType metric, + double t_const +) { + const api::Parameters parameters{input.size(), total_bits, metric, t_const}; + const auto sizes = api::requirements(parameters, api::CodeLayout::UnpackedU16); + std::vector workspace( + (sizes.workspace_bytes + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t) + ); + std::vector bytes(sizes.code_bytes); + const api::CodeOutput output{ + api::CodeLayout::UnpackedU16, {bytes.data(), bytes.size()}, {}}; + Encoded result; + result.factors = api::encode( + parameters, + input.data(), + centroid.data(), + output, + {workspace.data(), sizes.workspace_bytes} + ); + result.codes.resize(input.size()); + for (size_t i = 0; i < input.size(); ++i) { + result.codes[i] = static_cast(bytes[2 * i] | (bytes[2 * i + 1] << 8U)); + } + return result; +} + +// The detail code_factors on int codes, the bit-exact reference for the wrapper. +template +api::Factors detail_factors( + const std::vector& residual, + const std::vector& centroid, + const std::vector& code, + size_t total_bits, + MetricType metric +) { + const std::vector widened(code.begin(), code.end()); + const float cb = -static_cast((1U << total_bits) - 1) / 2.0F; + api::Factors factors; + rabitq_impl::xy_bits::code_factors( + residual.data(), + centroid.data(), + residual.size(), + widened.data(), + cb, + factors.add, + factors.rescale, + factors.error, + metric + ); + return factors; +} + +void expect_same_factors(const api::Factors& actual, const api::Factors& expected) { + EXPECT_EQ(actual.add, expected.add); + EXPECT_EQ(actual.rescale, expected.rescale); + EXPECT_EQ(actual.error, expected.error); +} + +template +void check_code_and_factors( + const std::vector& input, + const std::vector& centroid, + size_t total_bits, + MetricType metric, + double t_const +) { + const size_t dim = input.size(); + // The residual v1::encode forms internally; elementwise subtraction is exact IEEE. + std::vector residual(dim); + for (size_t i = 0; i < dim; ++i) { + residual[i] = input[i] - centroid[i]; + } + const Encoded reference = v1_encode(input, centroid, total_bits, metric, t_const); + + std::vector code(dim); + provisional::combined_code(residual.data(), dim, total_bits, code.data(), t_const); + std::vector detail_code(dim); + rabitq_impl::total_bits::combined_code( + residual.data(), dim, total_bits, detail_code.data(), t_const + ); + EXPECT_EQ(code, detail_code); + EXPECT_EQ(std::vector(code.begin(), code.end()), reference.codes); + + const api::Factors factors = provisional::code_factors( + residual.data(), centroid.data(), dim, code.data(), total_bits, metric + ); + expect_same_factors( + factors, detail_factors(residual, centroid, code, total_bits, metric) + ); + // v1::encode derives rescale from the normalized code in a different rounding + // order; add and error come from the same terms. + EXPECT_EQ(factors.add, reference.factors.add); + EXPECT_EQ(factors.error, reference.factors.error); + EXPECT_NEAR( + factors.rescale, + reference.factors.rescale, + 1e-5F * std::abs(reference.factors.rescale) + ); + + // Every prefix of the code is itself a code of that width for the same residual. + for (size_t prefix_bits = 1; prefix_bits < total_bits; ++prefix_bits) { + SCOPED_TRACE(::testing::Message() << "prefix_bits " << prefix_bits); + std::vector prefix(dim); + for (size_t i = 0; i < dim; ++i) { + prefix[i] = static_cast(code[i] >> (total_bits - prefix_bits)); + } + expect_same_factors( + provisional::code_factors( + residual.data(), centroid.data(), dim, prefix.data(), prefix_bits, metric + ), + detail_factors(residual, centroid, prefix, prefix_bits, metric) + ); + } +} + +TEST(ProvisionalQuantizeTest, MatchesV1EncodeAndTheDetailFunctions) { + std::mt19937 generator(7); + for (size_t dim : {32U, 64U, 96U, 1056U}) { + const std::vector input = gaussian(dim, generator); + const std::vector centroid = gaussian(dim, generator); + for (size_t total_bits = 1; total_bits <= 9; ++total_bits) { + const double fast = rabitqlib::quant::faster_config(dim, total_bits).t_const; + for (double t_const : {-1.0, fast}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE( + ::testing::Message() + << "dim " << dim << " total_bits " << total_bits << " t_const " + << t_const << " metric " << metric + ); + check_code_and_factors( + input, centroid, total_bits, metric, t_const + ); + if (total_bits <= 8) { + check_code_and_factors( + input, centroid, total_bits, metric, t_const + ); + } + } + } + } + } +} + +TEST(ProvisionalQuantizeTest, NullCentroidMeansZero) { + std::mt19937 generator(11); + constexpr size_t kDim = 64; + const std::vector residual = gaussian(kDim, generator); + const std::vector zero(kDim, 0.0F); + std::vector code(kDim); + provisional::combined_code(residual.data(), kDim, 4, code.data()); + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + expect_same_factors( + provisional::code_factors( + residual.data(), nullptr, kDim, code.data(), 4, metric + ), + provisional::code_factors( + residual.data(), zero.data(), kDim, code.data(), 4, metric + ) + ); + } +} + +TEST(ProvisionalQuantizeTest, ZeroAndUnderflowingResiduals) { + constexpr size_t kDim = 64; + std::mt19937 generator(3); + const std::vector centroid = gaussian(kDim, generator); + std::vector tiny(kDim, 0.0F); + // A nonzero residual whose squared norm underflows to zero keeps its sign bits. + tiny[3] = 1e-30F; + tiny[40] = -1e-30F; + for (const auto& residual : {std::vector(kDim, 0.0F), tiny}) { + for (size_t total_bits = 1; total_bits <= 9; ++total_bits) { + std::vector input(kDim); + for (size_t i = 0; i < kDim; ++i) { + input[i] = centroid[i] + residual[i]; + } + std::vector code(kDim, 0xFFFF); + provisional::combined_code(residual.data(), kDim, total_bits, code.data()); + for (size_t i = 0; i < kDim; ++i) { + EXPECT_EQ(code[i], residual[i] > 0 ? 1U << (total_bits - 1) : 0U); + } + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + const api::Factors factors = provisional::code_factors( + residual.data(), centroid.data(), kDim, code.data(), total_bits, metric + ); + EXPECT_EQ(factors.add, metric == METRIC_L2 ? 0.0F : 1.0F); + EXPECT_EQ(factors.rescale, 0.0F); + EXPECT_EQ(factors.error, 0.0F); + } + } + } +} + +template +void expect_rejected(Operation operation, const std::string& message) { + try { + operation(); + ADD_FAILURE() << "accepted; expected: " << message; + } catch (const Exception& error) { EXPECT_EQ(error.what(), message); } +} + +TEST(ProvisionalQuantizeTest, RejectsInvalidArgumentsWithoutWritingOutput) { + constexpr size_t kDim = 64; + std::mt19937 generator(5); + std::vector residual = gaussian(kDim, generator); + std::vector code8(kDim, 0x5A); + std::vector code16(kDim, 0x5A5A); + const auto code8_untouched = code8; + const auto code16_untouched = code16; + const auto encode8 = [&](size_t dim, size_t total_bits, double t_const = -1) { + provisional::combined_code(residual.data(), dim, total_bits, code8.data(), t_const); + }; + + for (size_t dim : {0U, 31U, 33U}) { + expect_rejected( + [&] { encode8(dim, 1); }, "dim must be a positive multiple of 32" + ); + } + for (size_t total_bits : {0U, 9U}) { + expect_rejected( + [&] { encode8(kDim, total_bits); }, + "total_bits must be in [1, 8] for uint8_t codes" + ); + } + expect_rejected( + [&] { provisional::combined_code(residual.data(), kDim, 10, code16.data()); }, + "total_bits must be in [1, 9] for uint16_t codes" + ); + for (double t_const : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + static_cast(std::numeric_limits::max())}) { + expect_rejected( + [&] { encode8(kDim, 4, t_const); }, + "t_const must be finite and at most INT_MAX / 2" + ); + } + expect_rejected( + [&] { provisional::combined_code(nullptr, kDim, 4, code8.data()); }, + "residual must not be null" + ); + expect_rejected( + [&] { + provisional::combined_code( + residual.data(), kDim, 4, static_cast(nullptr) + ); + }, + "code must not be null" + ); + expect_rejected( + [&] { + provisional::combined_code( + residual.data(), kDim, 4, reinterpret_cast(residual.data() + 8) + ); + }, + "Writable buffers must not overlap other buffers or inputs" + ); + + std::vector bad = residual; + bad[17] = std::numeric_limits::quiet_NaN(); + expect_rejected( + [&] { provisional::combined_code(bad.data(), kDim, 4, code8.data()); }, + "Input vectors must contain only finite values" + ); + std::vector huge(kDim, 1e30F); + expect_rejected( + [&] { provisional::combined_code(huge.data(), kDim, 4, code8.data()); }, + "Residual squared norm is not representable in float32" + ); + EXPECT_EQ(code8, code8_untouched); + EXPECT_EQ(code16, code16_untouched); + + std::vector code(kDim); + provisional::combined_code(residual.data(), kDim, 4, code.data()); + const auto factors = [&](const float* r, + const float* c, + const uint8_t* q, + size_t bits, + MetricType metric = METRIC_L2) { + return provisional::code_factors(r, c, kDim, q, bits, metric); + }; + expect_rejected( + [&] { + provisional::code_factors( + residual.data(), nullptr, 48, code.data(), 4, METRIC_L2 + ); + }, + "dim must be a positive multiple of 32" + ); + expect_rejected( + [&] { + factors(residual.data(), nullptr, code.data(), 4, static_cast(7)); + }, + "Metric type must be METRIC_L2 or METRIC_IP" + ); + expect_rejected( + [&] { factors(residual.data(), bad.data(), code.data(), 4); }, + "Input vectors must contain only finite values" + ); + expect_rejected( + [&] { factors(residual.data(), nullptr, nullptr, 4); }, "code must not be null" + ); + expect_rejected( + [&] { factors(residual.data(), nullptr, code.data(), 3); }, + "Unpacked code exceeds the configured bit width" + ); + // The code of the negated residual points the other way. + std::vector negated(kDim); + for (size_t i = 0; i < kDim; ++i) { + negated[i] = -residual[i]; + } + expect_rejected( + [&] { factors(negated.data(), nullptr, code.data(), 4); }, + "code does not quantize residual: must be positive" + ); +} + +} // namespace