From b101346edf08cb4f7dc3ed58191caec1d32b60fa Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Sun, 1 Nov 2026 22:53:34 +0800 Subject: [PATCH 01/16] feat(core): add provisional combined_code and code_factors for rabitqlib.quant --- CMakeLists.txt | 1 + docs/development.md | 3 +- .../rabitqlib/core/provisional/quantize.hpp | 62 +++ src/core/provisional.cpp | 201 ++++++++++ .../provisional_quantize_test.cpp | 362 ++++++++++++++++++ 5 files changed, 628 insertions(+), 1 deletion(-) create mode 100644 include/rabitqlib/core/provisional/quantize.hpp create mode 100644 src/core/provisional.cpp create mode 100644 tests/unit/rabitqlib/quantization/provisional_quantize_test.cpp 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..ea507b72 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,7 +49,7 @@ 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 | | `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 | diff --git a/include/rabitqlib/core/provisional/quantize.hpp b/include/rabitqlib/core/provisional/quantize.hpp new file mode 100644 index 00000000..d5757720 --- /dev/null +++ b/include/rabitqlib/core/provisional/quantize.hpp @@ -0,0 +1,62 @@ +#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. + +/** + * 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/src/core/provisional.cpp b/src/core/provisional.cpp new file mode 100644 index 00000000..165b5413 --- /dev/null +++ b/src/core/provisional.cpp @@ -0,0 +1,201 @@ +#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 "rabitqlib/utils/input_checks.hpp" +#include "rabitqlib/utils/space.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"); + } + if (dim > static_cast(std::numeric_limits::max())) { + throw std::length_error("dim exceeds the supported index range"); + } +} + +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" + ); + } +} + +void check_vector(const float* values, size_t dim, const char* null_message) { + if (values == nullptr) { + throw std::invalid_argument(null_message); + } + 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, as quant::v1::encode does. Returns the squared norm. +float checked_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; +} + +bool overlaps( + const void* first, size_t first_bytes, const void* second, size_t second_bytes +) { + const auto first_address = reinterpret_cast(first); + const auto second_address = reinterpret_cast(second); + return first_address < second_address + second_bytes && + second_address < first_address + first_bytes; +} + +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_vector(residual, dim, "residual must not be null"); + if (code == nullptr) { + throw std::invalid_argument("code must not be null"); + } + if (overlaps(code, dim * sizeof(TP), residual, dim * sizeof(float))) { + throw std::invalid_argument( + "Writable buffers must not overlap other buffers or inputs" + ); + } + 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"); + } + (void)checked_squared_norm(residual, dim); + + // Quantize into owned storage so a failure leaves the caller's codes unchanged. + std::vector result(dim); + rabitq_impl::total_bits::combined_code( + residual, dim, total_bits, result.data(), t_const + ); + std::copy(result.begin(), result.end(), code); +} + +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_vector(residual, dim, "residual must not be null"); + if (centroid != nullptr) { + check_vector(centroid, dim, "centroid must not be null"); + } + if (code == nullptr) { + throw std::invalid_argument("code must not be null"); + } + const float squared_norm = checked_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); + std::vector centered(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"); + } + centered[i] = static_cast(widened[i]) + cb; + } + // The factor derivation divides by this inner product; only a code that quantizes + // this residual (its own code or a prefix of it) makes it positive. + if (squared_norm != 0 && !(dot_product(residual, centered.data(), dim) > 0)) { + throw std::invalid_argument( + "code does not quantize residual: must be positive" + ); + } + + 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 + ); + if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || + !std::isfinite(factors.error)) { + throw std::overflow_error("Quantization factors are not representable in float32"); + } + 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/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 From e6d40bea25c83ba1ca2d1f2dea90dc3b4c49b8f8 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Sun, 1 Nov 2026 23:19:27 +0800 Subject: [PATCH 02/16] feat(python): add provisional rabitqlib.quant module with Rotator and layer 1 --- docs/development.md | 3 +- docs/formats.md | 75 ++ python_bindings/CMakeLists.txt | 3 +- python_bindings/bindings_common.hpp | 88 ++ python_bindings/quant.py | 68 ++ python_bindings/quant.pyi | 152 +++ python_bindings/quant_bindings.cpp | 1064 +++++++++++++++++ python_bindings/rabitq_bindings.cpp | 2 + tests/python/fixtures/quant/README.md | 10 + .../quant/packed_codes_nbits1_n3_d96.bin | 1 + .../quant/packed_codes_nbits2_n3_d96.bin | 1 + .../quant/packed_codes_nbits3_n3_d96.bin | Bin 0 -> 108 bytes .../quant/packed_codes_nbits4_n3_d96.bin | 1 + .../quant/packed_codes_nbits5_n3_d96.bin | Bin 0 -> 180 bytes .../quant/packed_codes_nbits6_n3_d96.bin | 1 + .../quant/packed_codes_nbits7_n3_d96.bin | 1 + .../quant/packed_codes_nbits8_n3_d96.bin | Bin 0 -> 288 bytes ...otator_fht_kac_dim100_padded128_seed42.bin | Bin 0 -> 128 bytes .../quant/rotator_fht_kac_dim64_seed0.bin | Bin 0 -> 96 bytes tests/python/test_quant_editor_support.py | 118 ++ tests/python/test_quant_layer1.py | 261 ++++ tests/python/test_quant_rotator.py | 161 +++ tests/python/typing/public_api.py | 30 + 23 files changed, 2038 insertions(+), 2 deletions(-) create mode 100644 docs/formats.md create mode 100644 python_bindings/quant.py create mode 100644 python_bindings/quant.pyi create mode 100644 python_bindings/quant_bindings.cpp create mode 100644 tests/python/fixtures/quant/README.md create mode 100644 tests/python/fixtures/quant/packed_codes_nbits1_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits2_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits3_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits4_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits5_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits6_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits7_n3_d96.bin create mode 100644 tests/python/fixtures/quant/packed_codes_nbits8_n3_d96.bin create mode 100644 tests/python/fixtures/quant/rotator_fht_kac_dim100_padded128_seed42.bin create mode 100644 tests/python/fixtures/quant/rotator_fht_kac_dim64_seed0.bin create mode 100644 tests/python/test_quant_editor_support.py create mode 100644 tests/python/test_quant_layer1.py create mode 100644 tests/python/test_quant_rotator.py diff --git a/docs/development.md b/docs/development.md index ea507b72..5800fe56 100644 --- a/docs/development.md +++ b/docs/development.md @@ -50,12 +50,13 @@ factors estimate L2 distance or inner product. | `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 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/formats.md b/docs/formats.md new file mode 100644 index 00000000..a49e2fc4 --- /dev/null +++ b/docs/formats.md @@ -0,0 +1,75 @@ +# 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 `::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/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index e77b7ec0..e870c486 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) @@ -60,7 +61,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..7c5e32a7 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" @@ -251,6 +253,92 @@ 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. +template +RowView checked_output_rows(py::handle value, size_t rows, size_t cols) { + const std::string expected = "out 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..436cf66e --- /dev/null +++ b/python_bindings/quant.py @@ -0,0 +1,68 @@ +"""Core quantization building blocks: rotation, quantization, factors and packing. + +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) +""" + +import numpy as np + +from ._rabitqlib._quant import ( + ESTIMATE_DTYPE, + FACTORS_DTYPE, + QUERY_TERMS_DTYPE, + FormatError, + Rotator, + code_factors, + pack_codes, + padded_dim, + quantize, + 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.""" + +MISSING = np.uint32(0xFFFFFFFF) +"""Marks an empty result slot, as the indexes' search does.""" + +# Classes report this module, so pickle and help() find them here. +for _public in (FormatError, Rotator): + _public.__module__ = __name__ +del _public + +__all__ = [ + "ESTIMATE_DTYPE", + "FACTORS_DTYPE", + "FACTORS_VERSION", + "MISSING", + "PACKED_CODES_VERSION", + "QUERY_TERMS_DTYPE", + "FormatError", + "Rotator", + "code_factors", + "pack_codes", + "padded_dim", + "quantize", + "split", + "unpack_codes", +] diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi new file mode 100644 index 00000000..817176e7 --- /dev/null +++ b/python_bindings/quant.pyi @@ -0,0 +1,152 @@ +"""Core quantization building blocks: rotation, quantization, factors and packing. + +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", + "MISSING", + "PACKED_CODES_VERSION", + "QUERY_TERMS_DTYPE", + "FormatError", + "Rotator", + "code_factors", + "pack_codes", + "padded_dim", + "quantize", + "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 +MISSING: np.uint32 + +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: 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.""" + +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, +) -> 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: Literal["l2", "ip", "innerproduct"], + 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.""" diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp new file mode 100644 index 00000000..7bd20e57 --- /dev/null +++ b/python_bindings/quant_bindings.cpp @@ -0,0 +1,1064 @@ +#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 "bindings_common.hpp" +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/detail/fast_scale_table.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/index/detail/parallel_for.hpp" +#include "rabitqlib/utils/input_checks.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; +}; + +constexpr size_t kMaxRotatorDim = 65536; +constexpr size_t kMaxTotalBits = 9; + +static_assert(sizeof(quant::v1::Factors) == 3 * sizeof(float)); + +// --- Structured dtypes, byte-compatible with the planned C structs --------------- + +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(" 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; +} + +template +void check_code_values(const RowView& codes, size_t nbits, const char* name) { + const auto maximum = static_cast((1U << nbits) - 1); + for (size_t i = 0; i < codes.rows; ++i) { + const TP* row = codes.row(i); + for (size_t j = 0; j < codes.cols; ++j) { + if (row[j] > maximum) { + throw std::invalid_argument( + std::string(name) + " values must be below 2**nbits = " + + std::to_string(size_t{1} << nbits) + "; got " + std::to_string(row[j]) + + " at (" + std::to_string(i) + ", " + std::to_string(j) + ")" + ); + } + } + } +} + +template +py::array_t new_rows(size_t rows, size_t cols) { + return py::array_t(std::vector{ + static_cast(rows), static_cast(cols)}); +} + +// The output rows: a new array, or the caller's out= after its exact checks. +template +std::pair> output_rows( + const py::object& out, size_t rows, size_t cols +) { + if (out.is_none()) { + py::array array = new_rows(rows, cols); + auto view = make_row_view(array); + return {std::move(array), std::move(view)}; + } + auto view = checked_output_rows(out, rows, cols); + return {py::reinterpret_borrow(out), std::move(view)}; +} + +// A (d,) float32 vector, copied if misaligned. +py::array_t ensure_vector( + py::handle value, size_t d, const char* name +) { + auto vector = ensure_1d_array(value, name); + if (static_cast(vector.shape(0)) != d) { + throw std::invalid_argument( + std::string(name) + " must have shape (" + std::to_string(d) + + ",); got shape " + array_shape(vector) + ); + } + if (reinterpret_cast(vector.data()) % alignof(float) != 0) { + py::array_t aligned( + vector.shape(0) + ); + std::memcpy( + aligned.mutable_data(), vector.data(), static_cast(vector.nbytes()) + ); + return aligned; + } + return vector; +} + +// --- Layer 1 -------------------------------------------------------------------- + +size_t padded_dim(size_t dim) { + if (dim == 0) { + throw std::invalid_argument("dim must be positive"); + } + if (dim > std::numeric_limits::max() - 31) { + throw std::invalid_argument("dim is too large to pad"); + } + return round_up_to_multiple(dim, 32); +} + +template +void quantize_rows( + const float* residuals, + size_t rows, + size_t d, + size_t nbits, + double t_const, + const RowView& out, + size_t num_threads +) { + index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + quant::provisional::combined_code( + residuals + (i * d), d, nbits, out.row(i), t_const + ); + }); +} + +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; + } + + const auto run = [&](auto& result) { + auto [array, view] = std::move(result); + { + py::gil_scoped_release release; + check_coordinates(input.data(), rows * d, d, "residuals"); + quantize_rows(input.data(), rows, d, nbits, t_const, view, num_threads); + } + return array; + }; + if (nbits == kMaxTotalBits) { + auto result = output_rows(out, rows, d); + return run(result); + } + auto result = output_rows(out, rows, d); + return run(result); +} + +template +py::tuple split_rows(const RowView& codes, size_t nbits, size_t prefix_bits) { + const size_t ex_bits = nbits - prefix_bits; + const auto mask = static_cast((1U << ex_bits) - 1); + auto prefix = new_rows(codes.rows, codes.cols); + auto ex = new_rows(codes.rows, codes.cols); + uint8_t* prefix_data = prefix.mutable_data(); + uint8_t* ex_data = ex.mutable_data(); + { + py::gil_scoped_release release; + check_code_values(codes, nbits, "codes"); + for (size_t i = 0; i < codes.rows; ++i) { + const TP* row = codes.row(i); + for (size_t j = 0; j < codes.cols; ++j) { + prefix_data[(i * codes.cols) + j] = static_cast(row[j] >> ex_bits); + ex_data[(i * codes.cols) + j] = static_cast(row[j] & mask); + } + } + } + return py::make_tuple(prefix, ex); +} + +py::tuple split(py::handle codes, size_t nbits, size_t prefix_bits) { + 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); + } + const auto view = exact_dtype_rows(codes, "codes"); + checked_working_dim(view.cols, view.owner, "codes"); + return split_rows(view, nbits, prefix_bits); +} + +template +py::array code_factors_rows( + const py::array_t& residuals, + const RowView& codes, + size_t nbits, + MetricType metric, + const float* centroid, + size_t num_threads, + const py::object& out +) { + const size_t rows = codes.rows; + const size_t d = codes.cols; + const py::dtype dtype = factors_dtype(); + py::array result; + if (out.is_none()) { + result = py::array(dtype, std::vector{static_cast(rows)}); + } else { + const std::string expected = + "out must be a writable, C-contiguous NumPy array of " + "dtype quant.FACTORS_DTYPE and shape (" + + std::to_string(rows) + ",)"; + if (!py::isinstance(out)) { + throw std::invalid_argument(expected + "; got " + array_input_type(out)); + } + result = py::reinterpret_borrow(out); + if (!result.dtype().attr("__eq__")(dtype).cast() || result.ndim() != 1 || + static_cast(result.shape(0)) != rows || !result.writeable() || + (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || + reinterpret_cast(result.data()) % alignof(float) != 0) { + throw std::invalid_argument(expected + "; got shape " + array_shape(result)); + } + } + auto* factors = static_cast(const_cast(result.data())); + { + py::gil_scoped_release release; + check_coordinates(residuals.data(), rows * d, d, "residuals"); + if (centroid != nullptr) { + check_coordinates(centroid, d, d, "centroid"); + } + check_code_values(codes, nbits, "codes"); + const auto factors_of = [&](size_t i) { + return quant::provisional::code_factors( + residuals.data() + (i * d), centroid, d, codes.row(i), nbits, metric + ); + }; + // Report the lowest mismatched row, whatever the thread schedule. + std::atomic first_rejected{rows}; + index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + try { + factors[i] = factors_of(i); + } catch (const std::invalid_argument&) { + size_t current = first_rejected.load(); + while (i < current && !first_rejected.compare_exchange_weak(current, i)) {} + } + }); + if (const size_t i = first_rejected.load(); i < rows) { + try { + (void)factors_of(i); + } catch (const std::invalid_argument& error) { + throw std::invalid_argument( + "codes row " + std::to_string(i) + " for residuals row " + + std::to_string(i) + ": " + error.what() + ); + } + } + } + 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)); + + py::array_t centroid_array; + const float* centroid_data = nullptr; + if (!centroid.is_none()) { + centroid_array = ensure_vector(centroid, d, "centroid"); + centroid_data = centroid_array.data(); + } + + 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, centroid_data, 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"); + auto [array, packed] = output_rows(out, view.rows, d * nbits / 8); + { + py::gil_scoped_release release; + check_code_values(view, nbits, "codes"); + index_detail::parallel_for( + 0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) { + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + view.row(i), packed.row(i), d, nbits + ); + } + ); + } + return array; +} + +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; + auto [array, codes] = output_rows(out, view.rows, d); + { + py::gil_scoped_release release; + index_detail::parallel_for( + 0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) { + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + view.row(i), codes.row(i), d, nbits + ); + } + ); + } + return array; +} + +// --- 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; + +// CRC-32 with the IEEE polynomial, as zlib.crc32 computes it. +uint32_t crc32(const uint8_t* data, size_t size) { + static const std::array kTable = [] { + std::array table{}; + for (uint32_t i = 0; i < table.size(); ++i) { + uint32_t value = i; + for (int bit = 0; bit < 8; ++bit) { + value = (value & 1U) != 0 ? 0xEDB88320U ^ (value >> 1U) : value >> 1U; + } + table[i] = value; + } + return table; + }(); + uint32_t crc = 0xFFFFFFFFU; + for (size_t i = 0; i < size; ++i) { + crc = kTable[(crc ^ data[i]) & 0xFFU] ^ (crc >> 8U); + } + return crc ^ 0xFFFFFFFFU; +} + +// Checksum of a blob with its crc32 field (offset 32) taken as zero. +uint32_t blob_checksum(const uint8_t* blob, size_t size) { + std::vector copy(blob, blob + size); + std::fill_n(copy.begin() + 32, 4, uint8_t{0}); + return crc32(copy.data(), copy.size()); +} + +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) { + const std::string message = + "padded_dim must be an integer multiple of 32 in [dim, " + + std::to_string(kMaxRotatorDim) + "]"; + if (py::isinstance(value)) { + throw py::type_error(message); + } + py::object integer; + try { + integer = py::module_::import("operator").attr("index")(value); + } catch (const py::error_already_set&) { throw py::type_error(message); } + if (integer < py::int_(0) || integer > py::int_(kMaxRotatorDim)) { + throw std::invalid_argument(message); + } + return integer.cast(); + } + + 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)); + auto [array, view] = output_rows(out, rows, padded_dim_); + { + py::gil_scoped_release release; + check_coordinates(input.data(), rows * dim_, dim_, "x"); + index_detail::parallel_for( + 0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + rotator_->rotate(input.data() + (i * dim_), view.row(i)); + } + ); + } + return array; + } + + [[nodiscard]] py::bytes to_bytes() const { + const size_t body = rotator_->dump_bytes(); + std::vector blob(kRotatorHeaderBytes + body, 0); + std::memcpy(blob.data(), 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.data() + kRotatorHeaderBytes)); + put_u32(&blob[32], blob_checksum(blob.data(), blob.size())); + return {reinterpret_cast(blob.data()), blob.size()}; + } + + 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_ + ); + if (rotator->dump_bytes() != body) { + throw std::logic_error("Rotator body size disagrees with the blob format"); + } + 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. Rejected for d above 65536. +num_threads : int, optional + Worker threads; 0 (the default) uses every available CPU. 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"), + 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. + +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.)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) uses every available CPU. +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) uses every available CPU. +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) uses every available CPU. +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) uses every available CPU. +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); } + ) + ); +} 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/tests/python/fixtures/quant/README.md b/tests/python/fixtures/quant/README.md new file mode 100644 index 00000000..2e6febe8 --- /dev/null +++ b/tests/python/fixtures/quant/README.md @@ -0,0 +1,10 @@ +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) +``` 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 @@ +ªªªªUUUUªªªªUUUUªªªªUUUUªªªªUUUUªªªª \ 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 @@ +Pú¥Pú¥Pú¥Pú¥ÆÆÆÆÆÆÆÆ¥Pú¥Pú¥Pú¥Púú¥Pú¥Pú¥Pú¥Pllllllll \ 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 0000000000000000000000000000000000000000..3d32546eb00e379bed3536198a268de188ab9727 GIT binary patch literal 108 ycmWIG|Ftv#34dVV|G@wM1Haa@4{(H{0H~NBswmvB9ga{<2>==bQM9Ma5{>}GLpBuv literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..4d907155848f08a95bab869649e79e03a936666f GIT binary patch literal 180 zcmZSBohEpuT}&sC^nygSZn!ZHPÉ”_oºÅf±¼Çh³¾ÉÔ§ÖÀ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;†Qçr=È“é´Ê•e°»Æg²½Èi´¿Ê’²6&ddmMÊ*xm0.¢nèÌ’ßúdcŠñŽÒ.zîOP>¥p;Æ‘\'ò½HÞ©t?Ê•`+¶A ×b-øƒÙ¤oúÊ`k¶ÁÌbm¸ÃÎdo&dlMIÙÙ›o}²1ÅxGi=÷'(U ¸´†ÙwA3råˆÊ•`+ö�Lâ­xΙd/ú…P¦q<ÇRè³>É”_oºÅf±¼Çh³¾ÉMÉÙ››’²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 0000000000000000000000000000000000000000..1e6489ca779f25b76d54fbbae8d2759976489403 GIT binary patch literal 288 zcmZSJ7E{zQck+!$&Z}(cpS5(u?qe74Jpc5ULs(AJ#NI13wV=9f;+*B1_MW(W@8y^O zT%rovW{y7LNx2oxeKVJ=-*xKhgV)~}dBl};Eu8%#Q}U}?C(K^9anJEfcVB$|$0;JO zYGCaa7?WO9+cA0G$}Rgbly^;CuzK6UGuI!z{lUy9siJS?8W5dUSkpdf?uyO(PF}hH z{uir&w3?xfdr)jfab4$>`Kz`bIDPHmo9|4#63Tj(E?%Jt*=3DA(-*DXaroTLCm71t z?>c(n_Op+F*o9;@jO{!_;YcH8-OeND TZ$16+n@vzg-N@D>I4%DzD2y4`&vXa3Tp0-q;* V-&~WbPMvSBnwawDnoyY7QUI}f7SaF! literal 0 HcmV?d00001 diff --git a/tests/python/test_quant_editor_support.py b/tests/python/test_quant_editor_support.py new file mode 100644 index 00000000..6aa18fdb --- /dev/null +++ b/tests/python/test_quant_editor_support.py @@ -0,0 +1,118 @@ +"""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"], +) +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", "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__) + members = stub_functions(classes["Rotator"].body) + public = {key for key in vars(quant.Rotator) if not key.startswith("_")} + assert public <= set(members), public - set(members) + for member, node in members.items(): + runtime = getattr(quant.Rotator, member) + documentation = ast.get_docstring(node) + if documentation is not None: + assert normalized(documentation) in normalized(runtime.__doc__), member + if member in ("rotate", "from_bytes", "__init__"): + 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__), + ): + 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..05e7c3b9 --- /dev/null +++ b/tests/python/test_quant_layer1.py @@ -0,0 +1,261 @@ +"""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", 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=r"codes values must be below 2\*\*nbits = 8"): + quant.code_factors(r, codes, nbits=3, metric="l2") + 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) diff --git a/tests/python/test_quant_rotator.py b/tests/python/test_quant_rotator.py new file mode 100644 index 00000000..cdeab354 --- /dev/null +++ b/tests/python/test_quant_rotator.py @@ -0,0 +1,161 @@ +"""Python boundary coverage for quant.Rotator and the rotator blob format.""" + +import copy +import pickle +import zlib +from pathlib import Path + +import numpy as np +import pytest +from rabitqlib import quant + +FIXTURES = Path(__file__).resolve().parent / "fixtures" / "quant" + + +def vectors(n, dim, seed=0): + return np.random.default_rng(seed).standard_normal((n, dim)).astype(np.float32) + + +@pytest.mark.parametrize("method", ["fht_kac", "matrix"]) +def test_rotation_preserves_norms_and_round_trips(method): + rotator = quant.Rotator(100, method=method, seed=5) + assert (rotator.dim, rotator.padded_dim, rotator.method, rotator.seed) == ( + 100, + 128, + method, + 5, + ) + x = vectors(7, 100) + y = rotator.rotate(x) + assert y.dtype == np.float32 and y.shape == (7, 128) + np.testing.assert_allclose( + np.linalg.norm(y, axis=1), np.linalg.norm(x, axis=1), rtol=1e-4 + ) + for loaded in ( + quant.Rotator.from_bytes(rotator.to_bytes()), + quant.Rotator.from_bytes(bytearray(rotator.to_bytes())), + quant.Rotator.from_bytes(memoryview(rotator.to_bytes())), + quant.Rotator.from_bytes(np.frombuffer(rotator.to_bytes(), np.uint8)), + pickle.loads(pickle.dumps(rotator)), + ): + assert repr(loaded) == repr(rotator) + np.testing.assert_array_equal(loaded.rotate(x), y) + assert copy.copy(rotator) is rotator and copy.deepcopy(rotator) is rotator + + +def test_seeds_threads_and_out(): + x = vectors(50, 64) + first = quant.Rotator(64, seed=11) + np.testing.assert_array_equal(quant.Rotator(64, seed=11).rotate(x), first.rotate(x)) + assert not np.array_equal(quant.Rotator(64, seed=12).rotate(x), first.rotate(x)) + np.testing.assert_array_equal( + first.rotate(x, num_threads=1), first.rotate(x, num_threads=4) + ) + drawn = quant.Rotator(64) + assert isinstance(drawn.seed, int) + np.testing.assert_array_equal( + quant.Rotator(64, seed=drawn.seed).rotate(x), drawn.rotate(x) + ) + assert quant.Rotator(64, method="fht").method == "fht_kac" + out = np.empty((50, 64), np.float32) + assert first.rotate(x, out=out) is out + np.testing.assert_array_equal(out, first.rotate(x)) + assert first.rotate(np.empty((0, 64), np.float32)).shape == (0, 64) + + +def test_golden_blobs(): + for name, rotator in ( + ("rotator_fht_kac_dim64_seed0.bin", quant.Rotator(64, seed=0)), + ( + "rotator_fht_kac_dim100_padded128_seed42.bin", + quant.Rotator(100, padded_dim=128, seed=42), + ), + ): + golden = (FIXTURES / name).read_bytes() + assert rotator.to_bytes() == golden, name + assert repr(quant.Rotator.from_bytes(golden)) == repr(rotator) + + +def test_blob_header_fields(): + blob = quant.Rotator(100, padded_dim=128, seed=42).to_bytes() + assert blob[:8] == b"RABQBLOB" + fields = np.frombuffer(blob[8:64], dtype=" 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]) + 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.MISSING, np.uint32) + assert_type(quant.FACTORS_DTYPE, np.dtype[np.void]) + assert issubclass(quant.FormatError, ValueError) + 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] From 2eaffc5d5f5b99a307dfe93e082d872aca158b14 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Sun, 1 Nov 2026 23:34:15 +0800 Subject: [PATCH 03/16] feat(python): add quant layer 2 with Query, estimates, refinement and query terms --- python_bindings/quant.py | 13 +- python_bindings/quant.pyi | 100 ++- python_bindings/quant_bindings.cpp | 917 ++++++++++++++++++++++ tests/python/test_quant_editor_support.py | 47 +- tests/python/test_quant_layer2.py | 392 +++++++++ tests/python/typing/public_api.py | 24 + 6 files changed, 1477 insertions(+), 16 deletions(-) create mode 100644 tests/python/test_quant_layer2.py diff --git a/python_bindings/quant.py b/python_bindings/quant.py index 436cf66e..e00c6b42 100644 --- a/python_bindings/quant.py +++ b/python_bindings/quant.py @@ -1,4 +1,4 @@ -"""Core quantization building blocks: rotation, quantization, factors and packing. +"""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 @@ -18,6 +18,9 @@ >>> 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 """ import numpy as np @@ -27,11 +30,14 @@ FACTORS_DTYPE, QUERY_TERMS_DTYPE, FormatError, + Query, Rotator, code_factors, + estimate_from_ip, pack_codes, padded_dim, quantize, + query_terms, split, unpack_codes, ) @@ -46,7 +52,7 @@ """Marks an empty result slot, as the indexes' search does.""" # Classes report this module, so pickle and help() find them here. -for _public in (FormatError, Rotator): +for _public in (FormatError, Query, Rotator): _public.__module__ = __name__ del _public @@ -58,11 +64,14 @@ "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", "FormatError", + "Query", "Rotator", "code_factors", + "estimate_from_ip", "pack_codes", "padded_dim", "quantize", + "query_terms", "split", "unpack_codes", ] diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi index 817176e7..44f4477b 100644 --- a/python_bindings/quant.pyi +++ b/python_bindings/quant.pyi @@ -1,4 +1,4 @@ -"""Core quantization building blocks: rotation, quantization, factors and packing. +"""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. @@ -17,11 +17,14 @@ __all__ = [ "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", "FormatError", + "Query", "Rotator", "code_factors", + "estimate_from_ip", "pack_codes", "padded_dim", "quantize", + "query_terms", "split", "unpack_codes", ] @@ -77,6 +80,80 @@ class Rotator: 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 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.""" @@ -126,7 +203,7 @@ def code_factors( codes: npt.NDArray[np.uint8] | npt.NDArray[np.uint16], *, nbits: int, - metric: Literal["l2", "ip", "innerproduct"], + metric: _Metric, centroid: npt.ArrayLike | None = None, num_threads: int = 0, out: npt.NDArray[np.void] | None = None, @@ -150,3 +227,22 @@ def unpack_codes( out: npt.NDArray[np.uint8] | None = None, ) -> npt.NDArray[np.uint8]: """Unpack codes written by pack_codes; the inverse is lossless.""" + +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 index 7bd20e57..c795abe2 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -10,11 +10,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -25,6 +27,7 @@ #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" @@ -32,6 +35,7 @@ #include "rabitqlib/core/types.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; @@ -422,6 +426,618 @@ py::array unpack_codes( return array; } +// --- Layer 2 -------------------------------------------------------------------- + +constexpr size_t kRecordFloats = 3; + +// 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. +struct Records { + py::array owner; + const float* data = nullptr; + size_t rows = 0; + + [[nodiscard]] const float* record(size_t i) const { return data + (kRecordFloats * i); } +}; + +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().attr("__eq__")(dtype).cast()) { + 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(dtype.itemsize())) || + reinterpret_cast(array.data()) % alignof(float) != 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; +} + +// A new 1-D record array, or the caller's out= of exactly that dtype and shape. +std::pair record_output( + const py::object& out, const py::dtype& dtype, const char* dtype_name, size_t rows +) { + 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 quant." + 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().attr("__eq__")(dtype).cast() || result.ndim() != 1 || + static_cast(result.shape(0)) != rows || !result.writeable() || + (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || + reinterpret_cast(result.data()) % alignof(float) != 0) { + throw std::invalid_argument(expected + "; got shape " + array_shape(result)); + } + auto* data = static_cast(result.mutable_data()); + return {std::move(result), data}; +} + +// A new float32 vector, or the caller's out= of exactly that dtype and shape. +std::pair float_output(const py::object& out, size_t rows) { + if (out.is_none()) { + py::array_t result(static_cast(rows)); + auto* data = result.mutable_data(); + return {std::move(result), data}; + } + const std::string expected = + "out must be a writable, C-contiguous NumPy array of dtype float32 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.ndim() != 1 || static_cast(result.shape(0)) != rows || + !result.writeable() || + (rows > 1 && result.strides(0) != static_cast(sizeof(float))) || + reinterpret_cast(result.data()) % alignof(float) != 0) { + throw std::invalid_argument(expected + "; got shape " + array_shape(result)); + } + auto* data = static_cast(result.mutable_data()); + return {std::move(result), data}; +} + +// 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 selection(const py::object& indices, size_t rows) { + Selection selected; + if (indices.is_none()) { + selected.count = rows; + return 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; +} + +void check_finite_records( + const Records& records, + const Selection& selected, + size_t field_begin, + size_t field_end, + const char* message +) { + for (size_t k = 0; k < selected.count; ++k) { + const float* record = records.record(selected.row(k)); + for (size_t field = field_begin; field < field_end; ++field) { + if (!std::isfinite(record[field])) { + throw std::invalid_argument(message); + } + } + } +} + +template +void check_selected_code_values( + const RowView& codes, const Selection& selected, size_t nbits, const char* name +) { + const auto maximum = static_cast((1U << nbits) - 1); + for (size_t k = 0; k < selected.count; ++k) { + const size_t i = selected.row(k); + const TP* row = codes.row(i); + for (size_t j = 0; j < codes.cols; ++j) { + if (row[j] > maximum) { + 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(row[j]) + " at (" + std::to_string(i) + ", " + + std::to_string(j) + ")" + ); + } + } + } +} + +void check_record_rows( + const Records& records, size_t rows, const char* name, const char* of +) { + if (records.rows != rows) { + throw std::invalid_argument( + std::string(name) + " must have one record per row of " + of + ", " + + std::to_string(rows) + "; got shape " + array_shape(records.owner) + ); + } +} + +// 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; +}; + +QueryTerms compute_query_terms( + const float* query, const float* centroid, size_t d, MetricType metric +) { + const float distance_sqr = + centroid == nullptr ? l2norm_sqr(query, d) : euclidean_sqr(query, centroid, d); + QueryTerms terms; + terms.g_add = metric == METRIC_L2 + ? distance_sqr + : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, d)); + terms.g_error = std::sqrt(distance_sqr); + terms.g_k1xsumq = std::accumulate(query, query + d, 0.0F) * -0.5F; + 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; +} + +// One estimate record: distance, error, refine_input. +inline void write_estimate( + float* out, float ip, size_t bits, const float* factors, const QueryTerms& terms +) { + out[0] = quant::detail::full_est_dist_from_ip( + ip, bits, factors[0], factors[1], terms.g_add, terms.g_k1xsumq + ); + out[1] = factors[2] * terms.g_error; + out[2] = ip; +} + +// A rotated working-space vector: float32, shape (d,) with d a positive multiple of 32. +std::vector working_vector(py::handle value, const char* name) { + auto array = ensure_1d_array(value, name); + const auto d = static_cast(array.shape(0)); + if (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) + ); + } + std::vector vector(d); + std::memcpy(vector.data(), array.data(), d * sizeof(float)); + check_coordinates(vector.data(), d, d, name); + return vector; +} + +std::vector centroid_vector(const py::object& centroid, size_t d) { + if (centroid.is_none()) { + return {}; + } + auto array = ensure_vector(centroid, d, "centroid"); + std::vector vector(array.data(), array.data() + d); + check_coordinates(vector.data(), d, d, "centroid"); + return vector; +} + +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 = working_vector(rotated_query, "rotated_query"); + 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 Records records = + record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); + check_record_rows(records, rows, "factors", "ip"); + auto terms_array = py::array::ensure(terms); + if (!terms_array || + !terms_array.dtype().attr("__eq__")(query_terms_dtype()).cast() || + 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)) + ); + } + 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] = record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", rows); + const Selection all_rows = selection(py::none(), rows); + { + py::gil_scoped_release release; + if (!all_finite(ips.data(), rows)) { + throw std::invalid_argument("ip must contain only finite values"); + } + check_finite_records( + records, all_rows, 0, kRecordFloats, "factors must contain only finite values" + ); + index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + write_estimate( + estimates + (kRecordFloats * i), + ips.data()[i], + nbits, + records.record(i), + query + ); + }); + } + return result; +} + +// 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_; + IndexAccess access_; + + [[nodiscard]] size_t d() const { return query_.size(); } + + // 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(); + } + + void check_row_length( + const py::array& array, + size_t cols, + size_t expected, + const char* name, + const char* what + ) const { + 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 = working_vector(rotated_query, "rotated_query"); + auto center = centroid_vector(centroid, query.size()); + return {std::move(query), metric_type, std::move(center)}; + } + + QuantQuery(const QuantQuery& other) + : query_(other.query_) + , centroid_(other.centroid_) + , metric_(other.metric_) + , terms_(other.terms_) {} + QuantQuery(QuantQuery&& other) noexcept + : query_(std::move(other.query_)) + , centroid_(std::move(other.centroid_)) + , metric_(other.metric_) + , terms_(other.terms_) {} + QuantQuery& operator=(const QuantQuery&) = delete; + QuantQuery& operator=(QuantQuery&&) = delete; + ~QuantQuery() = default; + + [[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; + } + + [[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) { + 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 Records records = + record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); + check_record_rows(records, codes.rows, "factors", "bitsets"); + const Selection selected = selection(indices, codes.rows); + auto [result, estimates] = + record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", selected.count); + const auto ip_func = select_excode_ipfunc(1); + { + py::gil_scoped_release release; + check_finite_records( + records, + selected, + 0, + kRecordFloats, + "factors must contain only finite values" + ); + index_detail::parallel_for( + 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { + const size_t i = selected.row(k); + const float ip = ip_func(query_.data(), codes.row(i), d()); + write_estimate( + estimates + (kRecordFloats * k), + ip, + 1, + records.record(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 Records records = + record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); + check_record_rows(records, view.rows, "factors", "codes"); + const Selection selected = selection(indices, view.rows); + auto [result, estimates] = + record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", selected.count); + { + py::gil_scoped_release release; + check_selected_code_values(view, selected, nbits, "codes"); + check_finite_records( + records, + selected, + 0, + kRecordFloats, + "factors must contain only finite values" + ); + index_detail::parallel_for( + 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { + const size_t i = selected.row(k); + const auto* code = view.row(i); + float ip = 0; + for (size_t j = 0; j < d(); ++j) { + ip += query_[j] * static_cast(code[j]); + } + write_estimate( + estimates + (kRecordFloats * k), + ip, + nbits, + records.record(i), + terms_ + ); + } + ); + } + return result; + }; + if (nbits == kMaxTotalBits) { + return run(exact_dtype_rows(codes, "codes")); + } + return run(exact_dtype_rows(codes, "codes")); + }); + } + + 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, + bool packed + ) { + 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 (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 Records records = + record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); + check_record_rows(records, codes.rows, "factors", name); + const Selection selected = selection(indices, codes.rows); + const Records 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] = float_output(out, selected.count); + const Selection all_estimates = selection(py::none(), selected.count); + const auto ip_func = select_excode_ipfunc(ex_bits); + const auto scale = static_cast(1U << ex_bits); + { + py::gil_scoped_release release; + if (!packed) { + check_selected_code_values(codes, selected, ex_bits, name); + } + check_finite_records( + records, + selected, + 0, + kRecordFloats, + "factors must contain only finite values" + ); + check_finite_records( + estimates, + all_estimates, + 2, + 3, + "first['refine_input'] must contain only finite values" + ); + index_detail::parallel_for( + 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { + const size_t i = selected.row(k); + const uint8_t* code = codes.row(i); + float ip_ex = 0; + if (packed) { + ip_ex = ip_func(query_.data(), code, d()); + } else { + for (size_t j = 0; j < d(); ++j) { + ip_ex += query_[j] * static_cast(code[j]); + } + } + const float* factor = records.record(i); + distances[k] = quant::detail::full_est_dist_from_ip( + (scale * estimates.record(k)[2]) + ip_ex, + total_bits, + factor[0], + factor[1], + terms_.g_add, + terms_.g_k1xsumq + ); + } + ); + } + return result; + }); + } +}; + // --- Blob format (docs/formats.md) ------------------------------------------------ constexpr std::array kBlobMagic = {'R', 'A', 'B', 'Q', 'B', 'L', 'O', 'B'}; @@ -1061,4 +1677,305 @@ FormatError [](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 uses every available CPU. +out : numpy.ndarray of ESTIMATE_DTYPE, shape (n,), optional + +Returns +------- +estimates : numpy.ndarray of ESTIMATE_DTYPE, shape (n,))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 uses every available CPU.)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& self, + 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) { + return self.refine( + ex, + factors, + first, + ex_bits, + prefix_bits, + indices, + num_threads, + out, + false + ); + }, + 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& self, + py::handle packed, + 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) { + return self.refine( + packed, + factors, + first, + ex_bits, + prefix_bits, + indices, + num_threads, + out, + true + ); + }, + 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; on one build it equals refine. + +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("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/tests/python/test_quant_editor_support.py b/tests/python/test_quant_editor_support.py index 6aa18fdb..283ac99f 100644 --- a/tests/python/test_quant_editor_support.py +++ b/tests/python/test_quant_editor_support.py @@ -78,7 +78,16 @@ def test_module_docs_and_exports(): @pytest.mark.parametrize( "name", - ["padded_dim", "quantize", "split", "code_factors", "pack_codes", "unpack_codes"], + [ + "padded_dim", + "quantize", + "split", + "code_factors", + "pack_codes", + "unpack_codes", + "query_terms", + "estimate_from_ip", + ], ) def test_function_signatures_and_docs(name): node = stub_functions(installed_stub().body)[name] @@ -89,21 +98,34 @@ def test_function_signatures_and_docs(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", "Rotator"} + 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__) - members = stub_functions(classes["Rotator"].body) - public = {key for key in vars(quant.Rotator) if not key.startswith("_")} - assert public <= set(members), public - set(members) - for member, node in members.items(): - runtime = getattr(quant.Rotator, member) - documentation = ast.get_docstring(node) - if documentation is not None: - assert normalized(documentation) in normalized(runtime.__doc__), member - if member in ("rotate", "from_bytes", "__init__"): - assert runtime_parameters(runtime) == stub_parameters(node), member + checked = { + "Rotator": ("rotate", "from_bytes", "__init__"), + "Query": ( + "__init__", + "set_centroid", + "estimate_bitset", + "estimate_codes", + "refine", + "refine_packed", + ), + } + 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(): @@ -112,6 +134,7 @@ def test_documentation_examples(): 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 diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py new file mode 100644 index 00000000..0ec5d6b3 --- /dev/null +++ b/tests/python/test_quant_layer2.py @@ -0,0 +1,392 @@ +"""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 = "Index is busy: conflicting operation in progress" + + +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]) diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index 3efc8348..2529d59c 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -300,6 +300,30 @@ def check_quant(x: npt.NDArray[np.float32]) -> None: assert_type(quant.MISSING, np.uint32) 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], + ) + 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] From 39bea9f8301f75551f00e4a5eeba6fd145f4372e Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Sun, 1 Nov 2026 23:43:27 +0800 Subject: [PATCH 04/16] feat(python): add quant layer 3 with FastScan blocks and estimate_fastscan --- docs/formats.md | 19 ++ python_bindings/quant.py | 6 + python_bindings/quant.pyi | 21 ++ python_bindings/quant_bindings.cpp | 204 ++++++++++++++++++ tests/python/fixtures/quant/README.md | 13 ++ .../quant/fastscan_blocks_n33_d64.bin | Bin 0 -> 1280 bytes tests/python/test_quant_editor_support.py | 2 + tests/python/test_quant_layer3.py | 175 +++++++++++++++ tests/python/typing/public_api.py | 6 + 9 files changed, 446 insertions(+) create mode 100644 tests/python/fixtures/quant/fastscan_blocks_n33_d64.bin create mode 100644 tests/python/test_quant_layer3.py diff --git a/docs/formats.md b/docs/formats.md index a49e2fc4..7aa60efb 100644 --- a/docs/formats.md +++ b/docs/formats.md @@ -14,6 +14,7 @@ Raw arrays have no header: their layout is fixed by the version below, which the | --- | --- | --- | | `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. diff --git a/python_bindings/quant.py b/python_bindings/quant.py index e00c6b42..12845b5f 100644 --- a/python_bindings/quant.py +++ b/python_bindings/quant.py @@ -35,6 +35,7 @@ code_factors, estimate_from_ip, pack_codes, + pack_fastscan_blocks, padded_dim, quantize, query_terms, @@ -48,6 +49,9 @@ FACTORS_VERSION = 1 """Version of FACTORS_DTYPE, ESTIMATE_DTYPE and QUERY_TERMS_DTYPE.""" +FASTSCAN_BLOCKS_VERSION = 1 +"""Version of the pack_fastscan_blocks layout.""" + MISSING = np.uint32(0xFFFFFFFF) """Marks an empty result slot, as the indexes' search does.""" @@ -60,6 +64,7 @@ "ESTIMATE_DTYPE", "FACTORS_DTYPE", "FACTORS_VERSION", + "FASTSCAN_BLOCKS_VERSION", "MISSING", "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", @@ -69,6 +74,7 @@ "code_factors", "estimate_from_ip", "pack_codes", + "pack_fastscan_blocks", "padded_dim", "quantize", "query_terms", diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi index 44f4477b..75f968ab 100644 --- a/python_bindings/quant.pyi +++ b/python_bindings/quant.pyi @@ -13,6 +13,7 @@ __all__ = [ "ESTIMATE_DTYPE", "FACTORS_DTYPE", "FACTORS_VERSION", + "FASTSCAN_BLOCKS_VERSION", "MISSING", "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", @@ -22,6 +23,7 @@ __all__ = [ "code_factors", "estimate_from_ip", "pack_codes", + "pack_fastscan_blocks", "padded_dim", "quantize", "query_terms", @@ -34,6 +36,7 @@ ESTIMATE_DTYPE: np.dtype[np.void] QUERY_TERMS_DTYPE: np.dtype[np.void] PACKED_CODES_VERSION: int FACTORS_VERSION: int +FASTSCAN_BLOCKS_VERSION: int MISSING: np.uint32 class FormatError(ValueError): @@ -147,6 +150,15 @@ class Query: 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: @@ -228,6 +240,15 @@ def unpack_codes( ) -> 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, *, diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index c795abe2..34597f25 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -33,6 +33,10 @@ #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" @@ -736,6 +740,70 @@ py::array estimate_from_ip( return result; } +constexpr uint8_t reversed_byte(uint8_t value) { + uint8_t reversed = 0; + for (size_t bit = 0; bit < 8; ++bit) { + reversed = static_cast(reversed | (((value >> bit) & 1U) << (7 - bit))); + } + return reversed; +} + +// 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 Records records = + record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); + check_record_rows(records, codes.rows, "factors", "bitsets"); + const size_t n = codes.rows; + const size_t num_blocks = div_round_up(n, fastscan::kBatchSize); + const size_t block_bytes = BatchDataMap::data_bytes(d); + auto [array, blocks] = output_rows(out, num_blocks, block_bytes); + const Selection all_rows = selection(py::none(), n); + { + py::gil_scoped_release release; + check_finite_records( + records, all_rows, 0, kRecordFloats, "factors must contain only finite values" + ); + // The FastScan packer reads sign bytes with the first coordinate highest. + std::vector msb_first(n * codes.cols); + index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t /*thread*/) { + const uint8_t* row = codes.row(i); + for (size_t byte = 0; byte < codes.cols; ++byte) { + msb_first[(i * codes.cols) + byte] = reversed_byte(row[byte]); + } + }); + index_detail::parallel_for( + 0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) { + const size_t begin = b * fastscan::kBatchSize; + const size_t lanes = std::min(fastscan::kBatchSize, n - begin); + uint8_t* block = blocks.row(b); + fastscan::pack_codes( + d, msb_first.data() + (begin * codes.cols), lanes, block + ); + BatchDataMap map(reinterpret_cast(block), d); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const float* factor = + lane < lanes ? records.record(begin + lane) : nullptr; + map.f_add()[lane] = factor == nullptr ? 0.0F : factor[0]; + map.f_rescale()[lane] = factor == nullptr ? 0.0F : factor[1]; + map.f_error()[lane] = factor == nullptr ? 0.0F : factor[2]; + } + } + ); + } + return array; +} + // 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 { @@ -744,10 +812,33 @@ class QuantQuery { std::vector centroid_; // Empty means no centroid (zero). MetricType metric_; QueryTerms terms_; + // The FastScan query, built on the first estimate_fastscan call. It borrows + // query_, so copies and moves never take it over; they build their own. + std::unique_ptr> batch_query_; IndexAccess 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_->set_g_add(terms_.g_error); + } else { + batch_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_) { + batch_query_ = std::make_unique>( + query_.data(), d(), 0, metric_, true + ); + apply_batch_terms(); + } + return *batch_query_; + } + // The guard is taken before any input conversion, so conversion cannot race // set_centroid. template @@ -815,6 +906,9 @@ class QuantQuery { ); centroid_ = std::move(center); terms_ = terms; + if (batch_query_) { + apply_batch_terms(); + } } [[nodiscard]] QuantQuery copy() { @@ -1036,6 +1130,60 @@ class QuantQuery { 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] = + record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", num_points); + const SplitBatchQuery& query = batch_query(); + { + py::gil_scoped_release release; + index_detail::parallel_for( + 0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) { + std::array distances{}; + std::array lower_bounds{}; + std::array ips{}; + std::array errors{}; + const auto* block = reinterpret_cast(view.row(b)); + split_batch_estdist( + block, + query, + d(), + distances.data(), + lower_bounds.data(), + ips.data(), + true + ); + ConstBatchDataMap(block, d()) + .f_error() + .copy_to(errors.data(), errors.size()); + const size_t begin = b * fastscan::kBatchSize; + const size_t lanes = + std::min(fastscan::kBatchSize, num_points - begin); + for (size_t lane = 0; lane < lanes; ++lane) { + float* record = estimates + (kRecordFloats * (begin + lane)); + record[0] = distances[lane]; + record[1] = errors[lane] * terms_.g_error; + record[2] = ips[lane]; + } + } + ); + } + return result; + }); + } }; // --- Blob format (docs/formats.md) ------------------------------------------------ @@ -1738,6 +1886,35 @@ 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) uses every available CPU. +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. @@ -1963,6 +2140,33 @@ 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 uses every available CPU. +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( diff --git a/tests/python/fixtures/quant/README.md b/tests/python/fixtures/quant/README.md index 2e6febe8..4df4470b 100644 --- a/tests/python/fixtures/quant/README.md +++ b/tests/python/fixtures/quant/README.md @@ -8,3 +8,16 @@ These files freeze the `rabitqlib.quant` byte formats described in [docs/formats 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 0000000000000000000000000000000000000000..6fc6b2bd1e861278beaf5b4a8e3358367a113f3b GIT binary patch literal 1280 zcmc)Itxv;H6vp8bKEh`MxQU8o3b!Q~RJkMwlA0M3QdLDN=b*01%;Zu9?n;7M9SI2( zH3VX+!H`T;pvX)U-1w~oaqmB{n>^y=zIjjE^!@Wdo*vFFzUN1+!QIuhF3TYGyf}88 z@mTBf`s&?CXx9yb%&V$?e>{#N_ikV&qSP&lEORQ~Z?~f;Z5^43GEKdrZue;~)A_fV zD1*TBs<>`n?7BEJ6K)Y?SylP{y4@`vp3FpA6mjM_zOUP~r=A>`3EJbhU%#8X*d>4G zL?l5DQ6fz!gh()i#t1nkm|_Mw71_i-8ffAW9UMd999>)@M1&r0kUZ1zLJkQua!esP zy=bBXg)Tz$kYIozZlQ6H5gw7_4HJA|iZ9IY19`>bcyu9$oUjiHArcIsF+z?BrkFt# z?~i>n(8M7+IEKPGy0}D$2tC}WT`|wPIm(7L+xJ`6y<^QQ*1T%Xd)92<+xvE`{I)gQ c&ud+ay^rt9Kd{WR_Wuto=^+xeq{a2xUm|MJsQ>@~ literal 0 HcmV?d00001 diff --git a/tests/python/test_quant_editor_support.py b/tests/python/test_quant_editor_support.py index 283ac99f..65462890 100644 --- a/tests/python/test_quant_editor_support.py +++ b/tests/python/test_quant_editor_support.py @@ -87,6 +87,7 @@ def test_module_docs_and_exports(): "unpack_codes", "query_terms", "estimate_from_ip", + "pack_fastscan_blocks", ], ) def test_function_signatures_and_docs(name): @@ -112,6 +113,7 @@ def test_classes_match_runtime(): "estimate_codes", "refine", "refine_packed", + "estimate_fastscan", ), } for name, signatures in checked.items(): diff --git a/tests/python/test_quant_layer3.py b/tests/python/test_quant_layer3.py new file mode 100644 index 00000000..b4f49896 --- /dev/null +++ b/tests/python/test_quant_layer3.py @@ -0,0 +1,175 @@ +"""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: 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] From 074f54fd9cce6ac90ae4173baf309ddbf6e89061 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Sun, 1 Nov 2026 23:43:27 +0800 Subject: [PATCH 05/16] docs: add rabitqlib.quant guide and unreleased release note --- docs/docs/rabitq/quant_python.md | 131 +++++++++++++++++++++++++++++++ docs/docs/releases/unreleased.md | 10 +++ docs/mkdocs.yml | 2 + 3 files changed, 143 insertions(+) create mode 100644 docs/docs/rabitq/quant_python.md create mode 100644 docs/docs/releases/unreleased.md diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md new file mode 100644 index 00000000..332a84e1 --- /dev/null +++ b/docs/docs/rabitq/quant_python.md @@ -0,0 +1,131 @@ +# 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"] - 2.0 * est["error"] +candidates = np.flatnonzero(lower <= np.partition(est["distance"], 10)[10]) +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. The probabilistic lower bound is `distance - m * error` for a multiplier `m` of your choice. `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`, `code_factors`, `pack_codes`, `unpack_codes`, `pack_fastscan_blocks`) default to `num_threads=0`, every available CPU; 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. +- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. + +## 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/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 From f863b808a8dd9518633a35c944fa24596a8254a0 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 2 Nov 2026 00:40:32 +0800 Subject: [PATCH 06/16] fix(python): address quant review: reject overlapping out, stage code_factors output, copy-free rotator blobs --- docs/docs/rabitq/quant_python.md | 2 +- python_bindings/bindings_common.hpp | 9 +- python_bindings/quant.py | 6 - python_bindings/quant.pyi | 6 +- python_bindings/quant_bindings.cpp | 301 ++++++++++++++++++++-------- tests/python/test_quant_layer1.py | 32 +++ tests/python/test_quant_layer2.py | 30 ++- tests/python/test_quant_layer3.py | 24 +++ tests/python/test_quant_rotator.py | 5 + tests/python/typing/public_api.py | 1 - 10 files changed, 312 insertions(+), 104 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index 332a84e1..6250260e 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -113,7 +113,7 @@ If your `ip` is itself approximate, `error` covers only the code's quantization - Build-side calls (`Rotator.rotate`, `quantize`, `code_factors`, `pack_codes`, `unpack_codes`, `pack_fastscan_blocks`) default to `num_threads=0`, every available CPU; 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. +- 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`. If a call raises, a caller's `out=` is left unchanged. - Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. ## Byte formats diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 7c5e32a7..83494675 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -34,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 { @@ -44,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; @@ -65,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); } }; diff --git a/python_bindings/quant.py b/python_bindings/quant.py index 12845b5f..c31cac2b 100644 --- a/python_bindings/quant.py +++ b/python_bindings/quant.py @@ -23,8 +23,6 @@ True """ -import numpy as np - from ._rabitqlib._quant import ( ESTIMATE_DTYPE, FACTORS_DTYPE, @@ -52,9 +50,6 @@ FASTSCAN_BLOCKS_VERSION = 1 """Version of the pack_fastscan_blocks layout.""" -MISSING = np.uint32(0xFFFFFFFF) -"""Marks an empty result slot, as the indexes' search does.""" - # Classes report this module, so pickle and help() find them here. for _public in (FormatError, Query, Rotator): _public.__module__ = __name__ @@ -65,7 +60,6 @@ "FACTORS_DTYPE", "FACTORS_VERSION", "FASTSCAN_BLOCKS_VERSION", - "MISSING", "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", "FormatError", diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi index 75f968ab..b91bcc54 100644 --- a/python_bindings/quant.pyi +++ b/python_bindings/quant.pyi @@ -14,7 +14,6 @@ __all__ = [ "FACTORS_DTYPE", "FACTORS_VERSION", "FASTSCAN_BLOCKS_VERSION", - "MISSING", "PACKED_CODES_VERSION", "QUERY_TERMS_DTYPE", "FormatError", @@ -37,7 +36,6 @@ QUERY_TERMS_DTYPE: np.dtype[np.void] PACKED_CODES_VERSION: int FACTORS_VERSION: int FASTSCAN_BLOCKS_VERSION: int -MISSING: np.uint32 class FormatError(ValueError): """Corrupt or incompatible serialized bytes. A subclass of ValueError.""" @@ -76,7 +74,9 @@ class Rotator: def to_bytes(self) -> bytes: """Serialize the rotation as a versioned rotator blob (docs/formats.md).""" @staticmethod - def from_bytes(data: Any) -> Rotator: + 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.""" diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 34597f25..063e3dd7 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -52,7 +53,8 @@ struct FormatError : std::invalid_argument { using std::invalid_argument::invalid_argument; }; -constexpr size_t kMaxRotatorDim = 65536; +// 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; static_assert(sizeof(quant::v1::Factors) == 3 * sizeof(float)); @@ -127,10 +129,55 @@ py::array_t new_rows(size_t rows, size_t cols) { 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 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(std::string("out must not overlap ") + input.name); + } + } +} + // The output rows: a new array, or the caller's out= after its exact checks. template std::pair> output_rows( - const py::object& out, size_t rows, size_t cols + const py::object& out, size_t rows, size_t cols, Inputs inputs ) { if (out.is_none()) { py::array array = new_rows(rows, cols); @@ -138,9 +185,41 @@ std::pair> output_rows( return {std::move(array), std::move(view)}; } auto view = checked_output_rows(out, rows, cols); + check_output_disjoint(view.owner, inputs); return {py::reinterpret_borrow(out), std::move(view)}; } +// A new 1-D record array, or the caller's out= of exactly that dtype and shape. +std::pair record_output( + const py::object& out, + const py::dtype& dtype, + const char* dtype_name, + size_t rows, + Inputs inputs +) { + 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 quant." + 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().attr("__eq__")(dtype).cast() || result.ndim() != 1 || + static_cast(result.shape(0)) != rows || !result.writeable() || + (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || + reinterpret_cast(result.data()) % alignof(float) != 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}; +} + // A (d,) float32 vector, copied if misaligned. py::array_t ensure_vector( py::handle value, size_t d, const char* name @@ -230,10 +309,10 @@ py::array quantize( return array; }; if (nbits == kMaxTotalBits) { - auto result = output_rows(out, rows, d); + auto result = output_rows(out, rows, d, {{"residuals", input}}); return run(result); } - auto result = output_rows(out, rows, d); + auto result = output_rows(out, rows, d, {{"residuals", input}}); return run(result); } @@ -278,46 +357,37 @@ py::array code_factors_rows( const RowView& codes, size_t nbits, MetricType metric, - const float* centroid, + const float* centroid_data, + py::handle centroid, size_t num_threads, const py::object& out ) { const size_t rows = codes.rows; const size_t d = codes.cols; - const py::dtype dtype = factors_dtype(); - py::array result; - if (out.is_none()) { - result = py::array(dtype, std::vector{static_cast(rows)}); - } else { - const std::string expected = - "out must be a writable, C-contiguous NumPy array of " - "dtype quant.FACTORS_DTYPE and shape (" + - std::to_string(rows) + ",)"; - if (!py::isinstance(out)) { - throw std::invalid_argument(expected + "; got " + array_input_type(out)); - } - result = py::reinterpret_borrow(out); - if (!result.dtype().attr("__eq__")(dtype).cast() || result.ndim() != 1 || - static_cast(result.shape(0)) != rows || !result.writeable() || - (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || - reinterpret_cast(result.data()) % alignof(float) != 0) { - throw std::invalid_argument(expected + "; got shape " + array_shape(result)); - } - } - auto* factors = static_cast(const_cast(result.data())); + auto [result, data] = record_output( + out, + factors_dtype(), + "FACTORS_DTYPE", + rows, + {{"residuals", residuals}, {"codes", codes.owner}, {"centroid", centroid}} + ); { py::gil_scoped_release release; check_coordinates(residuals.data(), rows * d, d, "residuals"); - if (centroid != nullptr) { - check_coordinates(centroid, d, d, "centroid"); + if (centroid_data != nullptr) { + check_coordinates(centroid_data, d, d, "centroid"); } check_code_values(codes, nbits, "codes"); + // A caller's out= is written only after every row succeeds. + std::vector staged(out.is_none() ? 0 : rows); + auto* factors = + out.is_none() ? reinterpret_cast(data) : staged.data(); const auto factors_of = [&](size_t i) { return quant::provisional::code_factors( - residuals.data() + (i * d), centroid, d, codes.row(i), nbits, metric + residuals.data() + (i * d), centroid_data, d, codes.row(i), nbits, metric ); }; - // Report the lowest mismatched row, whatever the thread schedule. + // Report the lowest rejected row, whatever the thread schedule. std::atomic first_rejected{rows}; index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { try { @@ -337,6 +407,9 @@ py::array code_factors_rows( ); } } + if (!staged.empty()) { + std::memcpy(data, staged.data(), staged.size() * sizeof(quant::v1::Factors)); + } } return result; } @@ -357,11 +430,12 @@ py::array code_factors( checked_working_dim(static_cast(input.shape(1)), input, "residuals"); const auto rows = static_cast(input.shape(0)); - py::array_t centroid_array; + py::object centroid_array = py::none(); const float* centroid_data = nullptr; if (!centroid.is_none()) { - centroid_array = ensure_vector(centroid, d, "centroid"); - centroid_data = centroid_array.data(); + auto vector = ensure_vector(centroid, d, "centroid"); + centroid_data = vector.data(); + centroid_array = std::move(vector); } const auto run = [&](const auto& view) { @@ -372,7 +446,7 @@ py::array code_factors( ); } return code_factors_rows( - input, view, nbits, metric_type, centroid_data, num_threads, out + input, view, nbits, metric_type, centroid_data, centroid_array, num_threads, out ); }; if (nbits == kMaxTotalBits) { @@ -387,7 +461,8 @@ py::array pack_codes( 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"); - auto [array, packed] = output_rows(out, view.rows, d * nbits / 8); + auto [array, packed] = + output_rows(out, view.rows, d * nbits / 8, {{"codes", view.owner}}); { py::gil_scoped_release release; check_code_values(view, nbits, "codes"); @@ -416,7 +491,7 @@ py::array unpack_codes( ); } const size_t d = bits / nbits; - auto [array, codes] = output_rows(out, view.rows, d); + auto [array, codes] = output_rows(out, view.rows, d, {{"packed", view.owner}}); { py::gil_scoped_release release; index_detail::parallel_for( @@ -473,34 +548,10 @@ Records record_input( return records; } -// A new 1-D record array, or the caller's out= of exactly that dtype and shape. -std::pair record_output( - const py::object& out, const py::dtype& dtype, const char* dtype_name, size_t rows -) { - 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 quant." + 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().attr("__eq__")(dtype).cast() || result.ndim() != 1 || - static_cast(result.shape(0)) != rows || !result.writeable() || - (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || - reinterpret_cast(result.data()) % alignof(float) != 0) { - throw std::invalid_argument(expected + "; got shape " + array_shape(result)); - } - auto* data = static_cast(result.mutable_data()); - return {std::move(result), data}; -} - // A new float32 vector, or the caller's out= of exactly that dtype and shape. -std::pair float_output(const py::object& out, size_t rows) { +std::pair float_output( + const py::object& out, size_t rows, Inputs inputs +) { if (out.is_none()) { py::array_t result(static_cast(rows)); auto* data = result.mutable_data(); @@ -519,6 +570,7 @@ std::pair float_output(const py::object& out, size_t rows) { reinterpret_cast(result.data()) % alignof(float) != 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}; } @@ -717,7 +769,13 @@ py::array estimate_from_ip( throw std::invalid_argument("terms must contain only finite values"); } const QueryTerms query{values[0], values[1], values[2]}; - auto [result, estimates] = record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", rows); + auto [result, estimates] = record_output( + out, + estimate_dtype(), + "ESTIMATE_DTYPE", + rows, + {{"ip", ips}, {"factors", records.owner}, {"terms", terms_array}} + ); const Selection all_rows = selection(py::none(), rows); { py::gil_scoped_release release; @@ -767,7 +825,9 @@ py::array pack_fastscan_blocks( const size_t n = codes.rows; const size_t num_blocks = div_round_up(n, fastscan::kBatchSize); const size_t block_bytes = BatchDataMap::data_bytes(d); - auto [array, blocks] = output_rows(out, num_blocks, block_bytes); + auto [array, blocks] = output_rows( + out, num_blocks, block_bytes, {{"bitsets", codes.owner}, {"factors", records.owner}} + ); const Selection all_rows = selection(py::none(), n); { py::gil_scoped_release release; @@ -815,7 +875,7 @@ class QuantQuery { // The FastScan query, built on the first estimate_fastscan call. It borrows // query_, so copies and moves never take it over; they build their own. std::unique_ptr> batch_query_; - IndexAccess access_; + IndexAccess access_{"Query is busy: set_centroid cannot overlap another call"}; [[nodiscard]] size_t d() const { return query_.size(); } @@ -930,7 +990,7 @@ class QuantQuery { } static QuantQuery from_state(const py::tuple& state) { - if (state.size() != 3) { + 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]); @@ -956,8 +1016,15 @@ class QuantQuery { record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); check_record_rows(records, codes.rows, "factors", "bitsets"); const Selection selected = selection(indices, codes.rows); - auto [result, estimates] = - record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", selected.count); + auto [result, estimates] = record_output( + out, + estimate_dtype(), + "ESTIMATE_DTYPE", + selected.count, + {{"bitsets", codes.owner}, + {"factors", records.owner}, + {"indices", selected.owner}} + ); const auto ip_func = select_excode_ipfunc(1); { py::gil_scoped_release release; @@ -1002,8 +1069,15 @@ class QuantQuery { record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); check_record_rows(records, view.rows, "factors", "codes"); const Selection selected = selection(indices, view.rows); - auto [result, estimates] = - record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", selected.count); + auto [result, estimates] = record_output( + out, + estimate_dtype(), + "ESTIMATE_DTYPE", + selected.count, + {{"codes", view.owner}, + {"factors", records.owner}, + {"indices", selected.owner}} + ); { py::gil_scoped_release release; check_selected_code_values(view, selected, nbits, "codes"); @@ -1080,7 +1154,14 @@ class QuantQuery { array_shape(estimates.owner) ); } - auto [result, distances] = float_output(out, selected.count); + auto [result, distances] = float_output( + out, + selected.count, + {{name, codes.owner}, + {"factors", records.owner}, + {"first", estimates.owner}, + {"indices", selected.owner}} + ); const Selection all_estimates = selection(py::none(), selected.count); const auto ip_func = select_excode_ipfunc(ex_bits); const auto scale = static_cast(1U << ex_bits); @@ -1145,11 +1226,36 @@ class QuantQuery { "; got shape " + array_shape(view.owner) ); } - auto [result, estimates] = - record_output(out, estimate_dtype(), "ESTIMATE_DTYPE", num_points); + auto [result, estimates] = record_output( + out, + estimate_dtype(), + "ESTIMATE_DTYPE", + num_points, + {{"blocks", view.owner}} + ); const SplitBatchQuery& query = batch_query(); { py::gil_scoped_release release; + // Hand-built blocks get the finite-factor check of the other calls. + for (size_t b = 0; b < num_blocks; ++b) { + const ConstBatchDataMap map( + reinterpret_cast(view.row(b)), d() + ); + std::array values{}; + const size_t lanes = std::min( + fastscan::kBatchSize, num_points - (b * fastscan::kBatchSize) + ); + 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)) { + throw std::invalid_argument( + "blocks must contain only finite factors; block " + + std::to_string(b) + " does not" + ); + } + } + } index_detail::parallel_for( 0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) { std::array distances{}; @@ -1199,8 +1305,8 @@ constexpr uint32_t kMethodMatrix = 2; constexpr uint32_t kScalarFloat32 = 1; constexpr uint32_t kScalarFloat64 = 2; -// CRC-32 with the IEEE polynomial, as zlib.crc32 computes it. -uint32_t crc32(const uint8_t* data, size_t size) { +// CRC-32 with the IEEE polynomial, continued from crc as zlib.crc32(data, crc) does. +uint32_t crc32(uint32_t crc, const uint8_t* data, size_t size) { static const std::array kTable = [] { std::array table{}; for (uint32_t i = 0; i < table.size(); ++i) { @@ -1212,18 +1318,22 @@ uint32_t crc32(const uint8_t* data, size_t size) { } return table; }(); - uint32_t crc = 0xFFFFFFFFU; + crc ^= 0xFFFFFFFFU; for (size_t i = 0; i < size; ++i) { crc = kTable[(crc ^ data[i]) & 0xFFU] ^ (crc >> 8U); } return crc ^ 0xFFFFFFFFU; } -// Checksum of a blob with its crc32 field (offset 32) taken as zero. +// Checksum of a blob with its crc32 field (offset 32) taken as zero, without a copy. +// The caller has checked that the blob covers the common header. uint32_t blob_checksum(const uint8_t* blob, size_t size) { - std::vector copy(blob, blob + size); - std::fill_n(copy.begin() + 32, 4, uint8_t{0}); - return crc32(copy.data(), copy.size()); + constexpr size_t kChecksumOffset = 32; + constexpr std::array kZero{}; + uint32_t crc = crc32(0, blob, kChecksumOffset); + crc = crc32(crc, kZero.data(), kZero.size()); + const size_t rest = kChecksumOffset + kZero.size(); + return crc32(crc, blob + rest, size - rest); } void put_u32(uint8_t* out, uint32_t value) { @@ -1389,7 +1499,7 @@ class QuantRotator { 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)); - auto [array, view] = output_rows(out, rows, padded_dim_); + auto [array, view] = output_rows(out, rows, padded_dim_, {{"x", input}}); { py::gil_scoped_release release; check_coordinates(input.data(), rows * dim_, dim_, "x"); @@ -1404,8 +1514,19 @@ class QuantRotator { [[nodiscard]] py::bytes to_bytes() const { const size_t body = rotator_->dump_bytes(); - std::vector blob(kRotatorHeaderBytes + body, 0); - std::memcpy(blob.data(), kBlobMagic.data(), kBlobMagic.size()); + 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)); @@ -1418,9 +1539,9 @@ class QuantRotator { 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.data() + kRotatorHeaderBytes)); - put_u32(&blob[32], blob_checksum(blob.data(), blob.size())); - return {reinterpret_cast(blob.data()), blob.size()}; + rotator_->save(reinterpret_cast(blob + kRotatorHeaderBytes)); + put_u32(&blob[32], blob_checksum(blob, size)); + return result; } static QuantRotator from_bytes(py::handle data) { @@ -1574,7 +1695,8 @@ 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. Rejected for d above 65536. + 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) uses every available CPU. Results do not depend on the thread count. @@ -2118,7 +2240,8 @@ distances : numpy.ndarray of float32)doc") py::arg("out") = py::none(), (std::string(R"doc(Refine prefix estimates with packed extra bits. -Uses the dispatched SIMD kernel; on one build it equals refine. +Uses the dispatched SIMD kernel, so it agrees with refine within float32 +rounding. Parameters ---------- diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py index 05e7c3b9..6772fb4c 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -259,3 +259,35 @@ def test_golden_packed_codes(nbits): 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() diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py index 0ec5d6b3..0c94e7f9 100644 --- a/tests/python/test_quant_layer2.py +++ b/tests/python/test_quant_layer2.py @@ -10,7 +10,7 @@ import pytest from rabitqlib import quant -BUSY = "Index is busy: conflicting operation in progress" +BUSY = "Query is busy: set_centroid cannot overlap another call" def rows(n, d, seed=0, scale=1.0): @@ -390,3 +390,31 @@ def test_rejected_arguments(): 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) + ) + + +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 index b4f49896..133fd5cf 100644 --- a/tests/python/test_quant_layer3.py +++ b/tests/python/test_quant_layer3.py @@ -173,3 +173,27 @@ def test_rejected_arguments(): for call, message in cases: with pytest.raises(ValueError, match=message): call() + + +def test_hand_built_blocks_with_non_finite_factors_are_rejected(): + c, bitsets, factors, q = setup("l2", 40) + query = quant.Query(q, metric="l2", centroid=c) + nan = np.frombuffer(np.float32(np.nan).tobytes(), np.uint8) + add = 96 * 4 # Block 1 holds vectors 32-39 in lanes 0-7. + blocks = quant.pack_fastscan_blocks(bitsets, factors) + blocks[1, add : add + 4] = nan + with pytest.raises(ValueError, match="blocks must contain only finite factors"): + query.estimate_fastscan(blocks, num_points=40) + # Padding lanes are not read. + blocks = quant.pack_fastscan_blocks(bitsets, factors) + blocks[1, add + 4 * 31 : add + 4 * 32] = nan + query.estimate_fastscan(blocks, num_points=40) + + +def test_out_overlapping_an_input_is_rejected(): + _, bitsets, factors, _ = setup("l2", 40) + shared = np.zeros(2 * 768, np.uint8) + inputs = shared[: 40 * 12].reshape(40, 12) + inputs[:] = bitsets + with pytest.raises(ValueError, match="out must not overlap bitsets"): + quant.pack_fastscan_blocks(inputs, factors, out=shared.reshape(2, 768)) diff --git a/tests/python/test_quant_rotator.py b/tests/python/test_quant_rotator.py index cdeab354..d96a0686 100644 --- a/tests/python/test_quant_rotator.py +++ b/tests/python/test_quant_rotator.py @@ -159,3 +159,8 @@ def test_rejected_arguments(): rotator.rotate(np.full((1, 64), np.nan, np.float32)) with pytest.raises(ValueError, match="out must be a writable NumPy array"): rotator.rotate(vectors(2, 64), out=np.empty((2, 64), np.float64)) + shared = vectors(3, 64) + with pytest.raises(ValueError, match="out must not overlap x"): + rotator.rotate(shared[:2], out=shared[1:]) + with pytest.raises(ValueError, match="out must not overlap x"): + rotator.rotate(shared, out=shared) diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index 6c310b83..4f5726c3 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -297,7 +297,6 @@ def check_quant(x: npt.NDArray[np.float32]) -> None: npt.NDArray[np.uint8], ) assert_type(quant.padded_dim(100), int) - assert_type(quant.MISSING, np.uint32) 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") From 6dfdca5c447798674b90f91ad556ad242d76a93e Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 2 Nov 2026 16:46:22 +0800 Subject: [PATCH 07/16] refactor(python): share core validation and simplify quant bindings --- .../rabitqlib/core/provisional/quantize.hpp | 3 +- python_bindings/bindings_common.hpp | 30 +- python_bindings/quant_bindings.cpp | 964 ++++++++---------- src/core/api.cpp | 43 +- src/core/provisional.cpp | 75 +- src/core/validation.hpp | 71 ++ tests/python/test_quant_layer1.py | 7 +- 7 files changed, 538 insertions(+), 655 deletions(-) create mode 100644 src/core/validation.hpp diff --git a/include/rabitqlib/core/provisional/quantize.hpp b/include/rabitqlib/core/provisional/quantize.hpp index d5757720..d7c3933d 100644 --- a/include/rabitqlib/core/provisional/quantize.hpp +++ b/include/rabitqlib/core/provisional/quantize.hpp @@ -11,7 +11,8 @@ 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. +// 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 diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 83494675..3231e18c 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -132,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. diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 063e3dd7..869723a6 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -57,9 +58,17 @@ struct FormatError : std::invalid_argument { constexpr size_t kMaxRotatorDim = rotator_impl::FhtKacRotator::kMaxDim; constexpr size_t kMaxTotalBits = 9; -static_assert(sizeof(quant::v1::Factors) == 3 * sizeof(float)); +// --- 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; +}; -// --- Structured dtypes, byte-compatible with the planned C structs --------------- +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; @@ -73,14 +82,33 @@ py::dtype float_record_dtype(const std::array& names) { return {name_list, formats, offsets, 3 * sizeof(float)}; } -py::dtype factors_dtype() { return float_record_dtype({"add", "rescale", "error"}); } +// Each dtype is built once and shared, since every estimate call needs one. +template +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() { - return float_record_dtype({"distance", "error", "refine_input"}); + 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() { - return float_record_dtype({"g_add", "g_error", "g_k1xsumq"}); + 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 ------------------------------------------------------------ @@ -106,23 +134,97 @@ size_t checked_working_dim(size_t d, const py::array& array, const char* name) { 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 -void check_code_values(const RowView& codes, size_t nbits, const char* name) { +void check_code_values( + const RowView& codes, const Selection& selected, size_t nbits, const char* name +) { const auto maximum = static_cast((1U << nbits) - 1); - for (size_t i = 0; i < codes.rows; ++i) { + for (size_t k = 0; k < selected.count; ++k) { + const size_t i = selected.row(k); const TP* row = codes.row(i); for (size_t j = 0; j < codes.cols; ++j) { if (row[j] > maximum) { throw std::invalid_argument( - std::string(name) + " values must be below 2**nbits = " + - std::to_string(size_t{1} << nbits) + "; got " + std::to_string(row[j]) + - " at (" + std::to_string(i) + ", " + std::to_string(j) + ")" + std::string(name) + " values must be below 2**" + + std::to_string(nbits) + " = " + std::to_string(size_t{1} << nbits) + + "; got " + std::to_string(row[j]) + " at (" + std::to_string(i) + ", " + + std::to_string(j) + ")" ); } } } } +// 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{ @@ -174,73 +276,127 @@ void check_output_disjoint(const py::array& out, Inputs inputs) { } } -// The output rows: a new array, or the caller's out= after its exact checks. +// 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 -std::pair> output_rows( - const py::object& out, size_t rows, size_t cols, Inputs inputs -) { +RowView output_rows(const py::object& out, size_t rows, size_t cols, Inputs inputs) { if (out.is_none()) { - py::array array = new_rows(rows, cols); - auto view = make_row_view(array); - return {std::move(array), std::move(view)}; + return make_row_view(new_rows(rows, cols)); } auto view = checked_output_rows(out, rows, cols); check_output_disjoint(view.owner, inputs); - return {py::reinterpret_borrow(out), std::move(view)}; + return view; } -// A new 1-D record array, or the caller's out= of exactly that dtype and shape. -std::pair record_output( +// 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()); + 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 quant." + dtype_name + - " and shape (" + std::to_string(rows) + ",)"; + "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().attr("__eq__")(dtype).cast() || result.ndim() != 1 || + if (!result.dtype().equal(dtype) || result.ndim() != 1 || static_cast(result.shape(0)) != rows || !result.writeable() || - (rows > 1 && result.strides(0) != static_cast(dtype.itemsize())) || - reinterpret_cast(result.data()) % alignof(float) != 0) { + (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()); + auto* data = static_cast(result.mutable_data()); return {std::move(result), data}; } -// A (d,) float32 vector, copied if misaligned. -py::array_t ensure_vector( - py::handle value, size_t d, const char* name +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 ) { - auto vector = ensure_1d_array(value, name); - if (static_cast(vector.shape(0)) != d) { + 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( - std::string(name) + " must have shape (" + std::to_string(d) + - ",); got shape " + array_shape(vector) + expected + "; got dtype " + py::str(array.dtype()).cast() ); } - if (reinterpret_cast(vector.data()) % alignof(float) != 0) { - py::array_t aligned( - vector.shape(0) - ); - std::memcpy( - aligned.mutable_data(), vector.data(), static_cast(vector.nbytes()) + 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 aligned; } - return vector; + return records; +} + +// Factors are read only for the selected rows; call with the GIL released. +void check_finite_factors( + const Records& factors, const Selection& selected +) { + for (size_t k = 0; k < selected.count; ++k) { + const quant::v1::Factors& f = factors[selected.row(k)]; + if (!std::isfinite(f.add) || !std::isfinite(f.rescale) || !std::isfinite(f.error)) { + throw std::invalid_argument("factors must contain only finite values"); + } + } } // --- Layer 1 -------------------------------------------------------------------- @@ -249,27 +405,7 @@ size_t padded_dim(size_t dim) { if (dim == 0) { throw std::invalid_argument("dim must be positive"); } - if (dim > std::numeric_limits::max() - 31) { - throw std::invalid_argument("dim is too large to pad"); - } - return round_up_to_multiple(dim, 32); -} - -template -void quantize_rows( - const float* residuals, - size_t rows, - size_t d, - size_t nbits, - double t_const, - const RowView& out, - size_t num_threads -) { - index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { - quant::provisional::combined_code( - residuals + (i * d), d, nbits, out.row(i), t_const - ); - }); + return rotator_impl::padding_requirement(dim, RotatorType::FhtKacRotator); } py::array quantize( @@ -299,21 +435,26 @@ py::array quantize( t_const = *frozen; } - const auto run = [&](auto& result) { - auto [array, view] = std::move(result); + // 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"); - quantize_rows(input.data(), rows, d, nbits, t_const, view, num_threads); + index_detail::parallel_for( + 0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + quant::provisional::combined_code( + input.data() + (i * d), d, nbits, view.row(i), t_const + ); + } + ); } - return array; + return view.owner; }; if (nbits == kMaxTotalBits) { - auto result = output_rows(out, rows, d, {{"residuals", input}}); - return run(result); + return run(output_rows(out, rows, d, {{"residuals", input}})); } - auto result = output_rows(out, rows, d, {{"residuals", input}}); - return run(result); + return run(output_rows(out, rows, d, {{"residuals", input}})); } template @@ -324,9 +465,10 @@ py::tuple split_rows(const RowView& codes, size_t nbits, size_t prefix_bits) auto ex = new_rows(codes.rows, codes.cols); uint8_t* prefix_data = prefix.mutable_data(); uint8_t* ex_data = ex.mutable_data(); + const Selection every_row = all_rows(codes.rows); { py::gil_scoped_release release; - check_code_values(codes, nbits, "codes"); + check_code_values(codes, every_row, nbits, "codes"); for (size_t i = 0; i < codes.rows; ++i) { const TP* row = codes.row(i); for (size_t j = 0; j < codes.cols; ++j) { @@ -357,34 +499,30 @@ py::array code_factors_rows( const RowView& codes, size_t nbits, MetricType metric, - const float* centroid_data, - py::handle centroid, + 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] = record_output( + auto [result, data] = vector_output( out, factors_dtype(), - "FACTORS_DTYPE", + "quant.FACTORS_DTYPE", rows, - {{"residuals", residuals}, {"codes", codes.owner}, {"centroid", centroid}} + {{"residuals", residuals}, {"codes", codes.owner}} ); { py::gil_scoped_release release; check_coordinates(residuals.data(), rows * d, d, "residuals"); - if (centroid_data != nullptr) { - check_coordinates(centroid_data, d, d, "centroid"); - } - check_code_values(codes, nbits, "codes"); // A caller's out= is written only after every row succeeds. std::vector staged(out.is_none() ? 0 : rows); - auto* factors = - out.is_none() ? reinterpret_cast(data) : staged.data(); + quant::v1::Factors* factors = out.is_none() ? data : staged.data(); + // code_factors checks each row's code values and that the code quantizes the + // residual; the binding adds the row number to its message. const auto factors_of = [&](size_t i) { return quant::provisional::code_factors( - residuals.data() + (i * d), centroid_data, d, codes.row(i), nbits, metric + residuals.data() + (i * d), centroid.data(), d, codes.row(i), nbits, metric ); }; // Report the lowest rejected row, whatever the thread schedule. @@ -407,9 +545,7 @@ py::array code_factors_rows( ); } } - if (!staged.empty()) { - std::memcpy(data, staged.data(), staged.size() * sizeof(quant::v1::Factors)); - } + std::copy(staged.begin(), staged.end(), data); } return result; } @@ -429,14 +565,9 @@ py::array code_factors( const size_t d = checked_working_dim(static_cast(input.shape(1)), input, "residuals"); const auto rows = static_cast(input.shape(0)); - - py::object centroid_array = py::none(); - const float* centroid_data = nullptr; - if (!centroid.is_none()) { - auto vector = ensure_vector(centroid, d, "centroid"); - centroid_data = vector.data(); - centroid_array = std::move(vector); - } + // 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) { @@ -445,9 +576,7 @@ py::array code_factors( "; got shape " + array_shape(view.owner) ); } - return code_factors_rows( - input, view, nbits, metric_type, centroid_data, centroid_array, num_threads, out - ); + return code_factors_rows(input, view, nbits, metric_type, center, num_threads, out); }; if (nbits == kMaxTotalBits) { return run(exact_dtype_rows(codes, "codes")); @@ -461,11 +590,12 @@ py::array pack_codes( 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"); - auto [array, packed] = + const auto packed = output_rows(out, view.rows, d * nbits / 8, {{"codes", view.owner}}); + const Selection every_row = all_rows(view.rows); { py::gil_scoped_release release; - check_code_values(view, nbits, "codes"); + check_code_values(view, every_row, nbits, "codes"); index_detail::parallel_for( 0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) { quant::rabitq_impl::ex_bits::packing_rabitqplus_code( @@ -474,7 +604,7 @@ py::array pack_codes( } ); } - return array; + return packed.owner; } py::array unpack_codes( @@ -491,7 +621,7 @@ py::array unpack_codes( ); } const size_t d = bits / nbits; - auto [array, codes] = output_rows(out, view.rows, d, {{"packed", view.owner}}); + const auto codes = output_rows(out, view.rows, d, {{"packed", view.owner}}); { py::gil_scoped_release release; index_detail::parallel_for( @@ -502,161 +632,11 @@ py::array unpack_codes( } ); } - return array; + return codes.owner; } // --- Layer 2 -------------------------------------------------------------------- -constexpr size_t kRecordFloats = 3; - -// 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. -struct Records { - py::array owner; - const float* data = nullptr; - size_t rows = 0; - - [[nodiscard]] const float* record(size_t i) const { return data + (kRecordFloats * i); } -}; - -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().attr("__eq__")(dtype).cast()) { - 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(dtype.itemsize())) || - reinterpret_cast(array.data()) % alignof(float) != 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; -} - -// A new float32 vector, or the caller's out= of exactly that dtype and shape. -std::pair float_output( - const py::object& out, size_t rows, Inputs inputs -) { - if (out.is_none()) { - py::array_t result(static_cast(rows)); - auto* data = result.mutable_data(); - return {std::move(result), data}; - } - const std::string expected = - "out must be a writable, C-contiguous NumPy array of dtype float32 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.ndim() != 1 || static_cast(result.shape(0)) != rows || - !result.writeable() || - (rows > 1 && result.strides(0) != static_cast(sizeof(float))) || - reinterpret_cast(result.data()) % alignof(float) != 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}; -} - -// 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 selection(const py::object& indices, size_t rows) { - Selection selected; - if (indices.is_none()) { - selected.count = rows; - return 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; -} - -void check_finite_records( - const Records& records, - const Selection& selected, - size_t field_begin, - size_t field_end, - const char* message -) { - for (size_t k = 0; k < selected.count; ++k) { - const float* record = records.record(selected.row(k)); - for (size_t field = field_begin; field < field_end; ++field) { - if (!std::isfinite(record[field])) { - throw std::invalid_argument(message); - } - } - } -} - -template -void check_selected_code_values( - const RowView& codes, const Selection& selected, size_t nbits, const char* name -) { - const auto maximum = static_cast((1U << nbits) - 1); - for (size_t k = 0; k < selected.count; ++k) { - const size_t i = selected.row(k); - const TP* row = codes.row(i); - for (size_t j = 0; j < codes.cols; ++j) { - if (row[j] > maximum) { - 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(row[j]) + " at (" + std::to_string(i) + ", " + - std::to_string(j) + ")" - ); - } - } - } -} - -void check_record_rows( - const Records& records, size_t rows, const char* name, const char* of -) { - if (records.rows != rows) { - throw std::invalid_argument( - std::string(name) + " must have one record per row of " + of + ", " + - std::to_string(rows) + "; got shape " + array_shape(records.owner) - ); - } -} - // The query-side terms of the reference formula (docs/formats.md, core/estimator.hpp). struct QueryTerms { float g_add = 0; @@ -674,7 +654,8 @@ QueryTerms compute_query_terms( ? distance_sqr : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, d)); terms.g_error = std::sqrt(distance_sqr); - terms.g_k1xsumq = std::accumulate(query, query + d, 0.0F) * -0.5F; + terms.g_k1xsumq = + std::accumulate(query, query + d, 0.0F) * fastscan::detail::code_center(1); 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"); @@ -682,49 +663,33 @@ QueryTerms compute_query_terms( return terms; } -// One estimate record: distance, error, refine_input. -inline void write_estimate( - float* out, float ip, size_t bits, const float* factors, const QueryTerms& terms +inline Estimate estimate_of( + float ip, size_t bits, const quant::v1::Factors& f, const QueryTerms& terms ) { - out[0] = quant::detail::full_est_dist_from_ip( - ip, bits, factors[0], factors[1], terms.g_add, terms.g_k1xsumq - ); - out[1] = factors[2] * terms.g_error; - out[2] = ip; -} - -// A rotated working-space vector: float32, shape (d,) with d a positive multiple of 32. -std::vector working_vector(py::handle value, const char* name) { - auto array = ensure_1d_array(value, name); - const auto d = static_cast(array.shape(0)); - if (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) - ); - } - std::vector vector(d); - std::memcpy(vector.data(), array.data(), d * sizeof(float)); - check_coordinates(vector.data(), d, d, name); - return vector; + 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}; } -std::vector centroid_vector(const py::object& centroid, size_t d) { - if (centroid.is_none()) { - return {}; +// for unpacked codes, summed in order as quant::v1::compute does. This is +// the portable reference path, not a SIMD kernel. +template +inline float unpacked_ip(const float* query, const TP* code, size_t d) { + float ip = 0; + for (size_t j = 0; j < d; ++j) { + ip += query[j] * static_cast(code[j]); } - auto array = ensure_vector(centroid, d, "centroid"); - std::vector vector(array.data(), array.data() + d); - check_coordinates(vector.data(), d, d, "centroid"); - return vector; + return 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 = working_vector(rotated_query, "rotated_query"); + 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 @@ -748,12 +713,9 @@ py::array estimate_from_ip( check_nbits(nbits, 1, kMaxTotalBits); auto ips = ensure_1d_array(ip, "ip"); const auto rows = static_cast(ips.shape(0)); - const Records records = - record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); - check_record_rows(records, rows, "factors", "ip"); + const auto records = factors_input(factors, rows, "ip"); auto terms_array = py::array::ensure(terms); - if (!terms_array || - !terms_array.dtype().attr("__eq__")(query_terms_dtype()).cast() || + 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 " @@ -763,48 +725,44 @@ py::array estimate_from_ip( : array_input_type(terms)) ); } - std::array values{}; + // 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] = record_output( - out, - estimate_dtype(), - "ESTIMATE_DTYPE", - rows, - {{"ip", ips}, {"factors", records.owner}, {"terms", terms_array}} - ); - const Selection all_rows = selection(py::none(), rows); + auto [result, estimates] = + estimate_output(out, rows, {{"ip", ips}, {"factors", records.owner}}); + const Selection every_row = all_rows(rows); { py::gil_scoped_release release; if (!all_finite(ips.data(), rows)) { throw std::invalid_argument("ip must contain only finite values"); } - check_finite_records( - records, all_rows, 0, kRecordFloats, "factors must contain only finite values" - ); + check_finite_factors(records, every_row); index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { - write_estimate( - estimates + (kRecordFloats * i), - ips.data()[i], - nbits, - records.record(i), - query - ); + estimates[i] = estimate_of(ips.data()[i], nbits, records[i], query); }); } return result; } -constexpr uint8_t reversed_byte(uint8_t value) { - uint8_t reversed = 0; - for (size_t bit = 0; bit < 8; ++bit) { - reversed = static_cast(reversed | (((value >> bit) & 1U) << (7 - bit))); +// --- 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 reversed; -} + return table; +}(); // FastScan blocks (docs/formats.md): 32 vectors per block, the last one zero-padded. py::array pack_fastscan_blocks( @@ -819,51 +777,79 @@ py::array pack_fastscan_blocks( ); } const size_t d = codes.cols * 8; - const Records records = - record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); - check_record_rows(records, codes.rows, "factors", "bitsets"); + 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 size_t block_bytes = BatchDataMap::data_bytes(d); - auto [array, blocks] = output_rows( - out, num_blocks, block_bytes, {{"bitsets", codes.owner}, {"factors", records.owner}} + const auto blocks = output_rows( + out, + num_blocks, + BatchDataMap::data_bytes(d), + {{"bitsets", codes.owner}, {"factors", records.owner}} ); - const Selection all_rows = selection(py::none(), n); + const Selection every_row = all_rows(n); { py::gil_scoped_release release; - check_finite_records( - records, all_rows, 0, kRecordFloats, "factors must contain only finite values" - ); - // The FastScan packer reads sign bytes with the first coordinate highest. - std::vector msb_first(n * codes.cols); - index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t /*thread*/) { - const uint8_t* row = codes.row(i); - for (size_t byte = 0; byte < codes.cols; ++byte) { - msb_first[(i * codes.cols) + byte] = reversed_byte(row[byte]); - } - }); + check_finite_factors(records, every_row); + // One block of sign bytes per worker, in the order the FastScan packer reads: + // the first coordinate of each byte in its highest bit. + const size_t workers = + std::max(1, std::min(resolve_num_threads(num_threads), num_blocks)); + std::vector scratch(workers * fastscan::kBatchSize * codes.cols); index_detail::parallel_for( - 0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) { + 0, num_blocks, num_threads, [&](size_t b, size_t thread) { const size_t begin = b * fastscan::kBatchSize; const size_t lanes = std::min(fastscan::kBatchSize, n - begin); + uint8_t* msb_first = + scratch.data() + (thread * fastscan::kBatchSize * codes.cols); + 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() + (begin * codes.cols), lanes, block - ); + fastscan::pack_codes(d, msb_first, lanes, block); BatchDataMap map(reinterpret_cast(block), d); for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { - const float* factor = - lane < lanes ? records.record(begin + lane) : nullptr; - map.f_add()[lane] = factor == nullptr ? 0.0F : factor[0]; - map.f_rescale()[lane] = factor == nullptr ? 0.0F : factor[1]; - map.f_error()[lane] = factor == nullptr ? 0.0F : factor[2]; + 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 array; + 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 { @@ -872,31 +858,29 @@ class QuantQuery { std::vector centroid_; // Empty means no centroid (zero). MetricType metric_; QueryTerms terms_; - // The FastScan query, built on the first estimate_fastscan call. It borrows - // query_, so copies and moves never take it over; they build their own. - std::unique_ptr> batch_query_; - IndexAccess access_{"Query is busy: set_centroid cannot overlap another call"}; + 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_->set_g_add(terms_.g_error); + batch_query_.query->set_g_add(terms_.g_error); } else { - batch_query_->set_g_add(terms_.g_error, -terms_.g_add); + 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_) { - batch_query_ = std::make_unique>( + if (!batch_query_.query) { + batch_query_.query = std::make_unique>( query_.data(), d(), 0, metric_, true ); apply_batch_terms(); } - return *batch_query_; + return *batch_query_.query; } // The guard is taken before any input conversion, so conversion cannot race @@ -907,13 +891,13 @@ class QuantQuery { return score(); } - void check_row_length( + static void check_row_length( const py::array& array, size_t cols, size_t expected, const char* name, const char* what - ) const { + ) { if (cols != expected) { throw std::invalid_argument( std::string(name) + " must have " + what + " = " + @@ -936,25 +920,11 @@ class QuantQuery { py::handle rotated_query, const std::string& metric, const py::object& centroid ) { const MetricType metric_type = metric_from_string(metric); - auto query = working_vector(rotated_query, "rotated_query"); + 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)}; } - QuantQuery(const QuantQuery& other) - : query_(other.query_) - , centroid_(other.centroid_) - , metric_(other.metric_) - , terms_(other.terms_) {} - QuantQuery(QuantQuery&& other) noexcept - : query_(std::move(other.query_)) - , centroid_(std::move(other.centroid_)) - , metric_(other.metric_) - , terms_(other.terms_) {} - QuantQuery& operator=(const QuantQuery&) = delete; - QuantQuery& operator=(QuantQuery&&) = delete; - ~QuantQuery() = default; - [[nodiscard]] size_t dim() const { return d(); } [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } @@ -966,7 +936,7 @@ class QuantQuery { ); centroid_ = std::move(center); terms_ = terms; - if (batch_query_) { + if (batch_query_.query) { apply_batch_terms(); } } @@ -1012,14 +982,10 @@ class QuantQuery { return scored([&] { const auto codes = exact_dtype_rows(bitsets, "bitsets"); check_row_length(codes.owner, codes.cols, d() / 8, "bitsets", "d // 8"); - const Records records = - record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); - check_record_rows(records, codes.rows, "factors", "bitsets"); + const auto records = factors_input(factors, codes.rows, "bitsets"); const Selection selected = selection(indices, codes.rows); - auto [result, estimates] = record_output( + auto [result, estimates] = estimate_output( out, - estimate_dtype(), - "ESTIMATE_DTYPE", selected.count, {{"bitsets", codes.owner}, {"factors", records.owner}, @@ -1028,23 +994,12 @@ class QuantQuery { const auto ip_func = select_excode_ipfunc(1); { py::gil_scoped_release release; - check_finite_records( - records, - selected, - 0, - kRecordFloats, - "factors must contain only finite values" - ); + check_finite_factors(records, selected); index_detail::parallel_for( 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { const size_t i = selected.row(k); - const float ip = ip_func(query_.data(), codes.row(i), d()); - write_estimate( - estimates + (kRecordFloats * k), - ip, - 1, - records.record(i), - terms_ + estimates[k] = estimate_of( + ip_func(query_.data(), codes.row(i), d()), 1, records[i], terms_ ); } ); @@ -1065,14 +1020,10 @@ class QuantQuery { return scored([&] { const auto run = [&](const auto& view) { check_row_length(view.owner, view.cols, d(), "codes", "d"); - const Records records = - record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); - check_record_rows(records, view.rows, "factors", "codes"); + const auto records = factors_input(factors, view.rows, "codes"); const Selection selected = selection(indices, view.rows); - auto [result, estimates] = record_output( + auto [result, estimates] = estimate_output( out, - estimate_dtype(), - "ESTIMATE_DTYPE", selected.count, {{"codes", view.owner}, {"factors", records.owner}, @@ -1080,27 +1031,15 @@ class QuantQuery { ); { py::gil_scoped_release release; - check_selected_code_values(view, selected, nbits, "codes"); - check_finite_records( - records, - selected, - 0, - kRecordFloats, - "factors must contain only finite values" - ); + check_code_values(view, selected, nbits, "codes"); + check_finite_factors(records, selected); index_detail::parallel_for( 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { const size_t i = selected.row(k); - const auto* code = view.row(i); - float ip = 0; - for (size_t j = 0; j < d(); ++j) { - ip += query_[j] * static_cast(code[j]); - } - write_estimate( - estimates + (kRecordFloats * k), - ip, + estimates[k] = estimate_of( + unpacked_ip(query_.data(), view.row(i), d()), nbits, - records.record(i), + records[i], terms_ ); } @@ -1115,6 +1054,9 @@ class QuantQuery { }); } + // 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, @@ -1123,28 +1065,25 @@ class QuantQuery { size_t prefix_bits, const py::object& indices, size_t num_threads, - const py::object& out, - bool packed + 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"; + const char* name = Packed ? "packed" : "ex"; return scored([&] { const auto codes = exact_dtype_rows(ex, name); - if (packed) { + 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 Records records = - record_input(factors, factors_dtype(), "FACTORS_DTYPE", "factors"); - check_record_rows(records, codes.rows, "factors", name); + const auto records = factors_input(factors, codes.rows, name); const Selection selected = selection(indices, codes.rows); - const Records estimates = - record_input(first, estimate_dtype(), "ESTIMATE_DTYPE", "first"); + 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 ") + @@ -1154,54 +1093,43 @@ class QuantQuery { array_shape(estimates.owner) ); } - auto [result, distances] = float_output( + 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 Selection all_estimates = selection(py::none(), selected.count); const auto ip_func = select_excode_ipfunc(ex_bits); const auto scale = static_cast(1U << ex_bits); { py::gil_scoped_release release; - if (!packed) { - check_selected_code_values(codes, selected, ex_bits, name); + if constexpr (!Packed) { + check_code_values(codes, selected, ex_bits, name); + } + check_finite_factors(records, selected); + for (size_t k = 0; k < estimates.rows; ++k) { + if (!std::isfinite(estimates[k].refine_input)) { + throw std::invalid_argument( + "first['refine_input'] must contain only finite values" + ); + } } - check_finite_records( - records, - selected, - 0, - kRecordFloats, - "factors must contain only finite values" - ); - check_finite_records( - estimates, - all_estimates, - 2, - 3, - "first['refine_input'] must contain only finite values" - ); index_detail::parallel_for( 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { const size_t i = selected.row(k); - const uint8_t* code = codes.row(i); - float ip_ex = 0; - if (packed) { - ip_ex = ip_func(query_.data(), code, d()); - } else { - for (size_t j = 0; j < d(); ++j) { - ip_ex += query_[j] * static_cast(code[j]); - } - } - const float* factor = records.record(i); + const float ip_ex = + Packed ? ip_func(query_.data(), codes.row(i), d()) + : unpacked_ip(query_.data(), codes.row(i), d()); + const quant::v1::Factors& f = records[i]; distances[k] = quant::detail::full_est_dist_from_ip( - (scale * estimates.record(k)[2]) + ip_ex, + (scale * estimates[k].refine_input) + ip_ex, total_bits, - factor[0], - factor[1], + f.add, + f.rescale, terms_.g_add, terms_.g_k1xsumq ); @@ -1211,6 +1139,7 @@ class QuantQuery { return result; }); } + py::array estimate_fastscan( py::handle blocks, size_t num_points, size_t num_threads, const py::object& out ) { @@ -1226,13 +1155,8 @@ class QuantQuery { "; got shape " + array_shape(view.owner) ); } - auto [result, estimates] = record_output( - out, - estimate_dtype(), - "ESTIMATE_DTYPE", - num_points, - {{"blocks", view.owner}} - ); + auto [result, estimates] = + estimate_output(out, num_points, {{"blocks", view.owner}}); const SplitBatchQuery& query = batch_query(); { py::gil_scoped_release release; @@ -1279,10 +1203,8 @@ class QuantQuery { const size_t lanes = std::min(fastscan::kBatchSize, num_points - begin); for (size_t lane = 0; lane < lanes; ++lane) { - float* record = estimates + (kRecordFloats * (begin + lane)); - record[0] = distances[lane]; - record[1] = errors[lane] * terms_.g_error; - record[2] = ips[lane]; + estimates[begin + lane] = { + distances[lane], errors[lane] * terms_.g_error, ips[lane]}; } } ); @@ -1305,35 +1227,22 @@ constexpr uint32_t kMethodMatrix = 2; constexpr uint32_t kScalarFloat32 = 1; constexpr uint32_t kScalarFloat64 = 2; -// CRC-32 with the IEEE polynomial, continued from crc as zlib.crc32(data, crc) does. -uint32_t crc32(uint32_t crc, const uint8_t* data, size_t size) { - static const std::array kTable = [] { - std::array table{}; - for (uint32_t i = 0; i < table.size(); ++i) { - uint32_t value = i; - for (int bit = 0; bit < 8; ++bit) { - value = (value & 1U) != 0 ? 0xEDB88320U ^ (value >> 1U) : value >> 1U; - } - table[i] = value; - } - return table; - }(); - crc ^= 0xFFFFFFFFU; - for (size_t i = 0; i < size; ++i) { - crc = kTable[(crc ^ data[i]) & 0xFFU] ^ (crc >> 8U); - } - return crc ^ 0xFFFFFFFFU; -} - -// Checksum of a blob with its crc32 field (offset 32) taken as zero, without a copy. -// The caller has checked that the blob covers the common header. +// 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 std::array kZero{}; - uint32_t crc = crc32(0, blob, kChecksumOffset); - crc = crc32(crc, kZero.data(), kZero.size()); - const size_t rest = kChecksumOffset + kZero.size(); - return crc32(crc, blob + rest, size - rest); + 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) { @@ -1449,20 +1358,12 @@ class QuantRotator { } static size_t checked_padded_dim_argument(const py::object& value) { - const std::string message = + return checked_integer( + value, + kMaxRotatorDim, "padded_dim must be an integer multiple of 32 in [dim, " + - std::to_string(kMaxRotatorDim) + "]"; - if (py::isinstance(value)) { - throw py::type_error(message); - } - py::object integer; - try { - integer = py::module_::import("operator").attr("index")(value); - } catch (const py::error_already_set&) { throw py::type_error(message); } - if (integer < py::int_(0) || integer > py::int_(kMaxRotatorDim)) { - throw std::invalid_argument(message); - } - return integer.cast(); + std::to_string(kMaxRotatorDim) + "]" + ); } public: @@ -1499,7 +1400,7 @@ class QuantRotator { 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)); - auto [array, view] = output_rows(out, rows, padded_dim_, {{"x", input}}); + const auto view = output_rows(out, rows, padded_dim_, {{"x", input}}); { py::gil_scoped_release release; check_coordinates(input.data(), rows * dim_, dim_, "x"); @@ -1509,7 +1410,7 @@ class QuantRotator { } ); } - return array; + return view.owner; } [[nodiscard]] py::bytes to_bytes() const { @@ -1623,9 +1524,6 @@ class QuantRotator { auto rotator = rotator_impl::make_rotator_for_loading( result.dim_, result.type_, result.padded_dim_ ); - if (rotator->dump_bytes() != body) { - throw std::logic_error("Rotator body size disagrees with the blob format"); - } rotator->load(reinterpret_cast(blob + kRotatorHeaderBytes)); result.rotator_ = std::move(rotator); return result; @@ -2154,27 +2052,7 @@ estimates : numpy.ndarray of ESTIMATE_DTYPE)doc") ) .def( "refine", - [](QuantQuery& self, - 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) { - return self.refine( - ex, - factors, - first, - ex_bits, - prefix_bits, - indices, - num_threads, - out, - false - ); - }, + &QuantQuery::refine, py::arg("ex"), py::arg("factors"), py::arg("first"), @@ -2208,27 +2086,7 @@ distances : numpy.ndarray of float32)doc") ) .def( "refine_packed", - [](QuantQuery& self, - py::handle packed, - 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) { - return self.refine( - packed, - factors, - first, - ex_bits, - prefix_bits, - indices, - num_threads, - out, - true - ); - }, + &QuantQuery::refine, py::arg("packed"), py::arg("factors"), py::arg("first"), diff --git a/src/core/api.cpp b/src/core/api.cpp index 37c6ea92..1af0c56d 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -16,28 +16,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 +71,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. @@ -295,10 +274,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 +306,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; } diff --git a/src/core/provisional.cpp b/src/core/provisional.cpp index 165b5413..64a37f03 100644 --- a/src/core/provisional.cpp +++ b/src/core/provisional.cpp @@ -10,8 +10,7 @@ #include "rabitqlib/core/detail/rabitq_impl.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/provisional/quantize.hpp" -#include "rabitqlib/utils/input_checks.hpp" -#include "rabitqlib/utils/space.hpp" +#include "validation.hpp" namespace rabitqlib::quant::provisional { namespace { @@ -26,9 +25,7 @@ void check_dim(size_t dim) { if (dim == 0 || dim % 32 != 0) { throw std::invalid_argument("dim must be a positive multiple of 32"); } - if (dim > static_cast(std::numeric_limits::max())) { - throw std::length_error("dim exceeds the supported index range"); - } + validation::check_dim_range(dim); } template @@ -41,32 +38,12 @@ void check_total_bits(size_t total_bits) { } } -void check_vector(const float* values, size_t dim, const char* null_message) { - if (values == nullptr) { - throw std::invalid_argument(null_message); - } - if (!all_finite(values, dim)) { - throw std::invalid_argument("Input vectors must contain only finite values"); +// 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"); } -} - -// Rejects a residual whose squared norm overflows before the quantizer's -// float-to-integer conversions, as quant::v1::encode does. Returns the squared norm. -float checked_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; -} - -bool overlaps( - const void* first, size_t first_bytes, const void* second, size_t second_bytes -) { - const auto first_address = reinterpret_cast(first); - const auto second_address = reinterpret_cast(second); - return first_address < second_address + second_bytes && - second_address < first_address + first_bytes; + validation::check_finite_input(residual, dim); } template @@ -75,26 +52,19 @@ void combined_code_impl( ) { check_dim(dim); check_total_bits(total_bits); - check_vector(residual, dim, "residual must not be null"); + check_residual(residual, dim); if (code == nullptr) { throw std::invalid_argument("code must not be null"); } - if (overlaps(code, dim * sizeof(TP), residual, dim * sizeof(float))) { - throw std::invalid_argument( - "Writable buffers must not overlap other buffers or inputs" - ); - } - 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"); - } - (void)checked_squared_norm(residual, dim); + 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); - // Quantize into owned storage so a failure leaves the caller's codes unchanged. - std::vector result(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, result.data(), t_const + residual, dim, total_bits, code, t_const ); - std::copy(result.begin(), result.end(), code); } template @@ -109,30 +79,30 @@ v1::Factors code_factors_impl( check_dim(dim); check_total_bits(total_bits); validate_metric_type(metric); - check_vector(residual, dim, "residual must not be null"); + check_residual(residual, dim); if (centroid != nullptr) { - check_vector(centroid, dim, "centroid must not be null"); + validation::check_finite_input(centroid, dim); } if (code == nullptr) { throw std::invalid_argument("code must not be null"); } - const float squared_norm = checked_squared_norm(residual, dim); + const float squared_norm = 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); - std::vector centered(dim); + float centered_ip = 0; 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"); } - centered[i] = static_cast(widened[i]) + cb; + centered_ip += residual[i] * (static_cast(widened[i]) + cb); } // The factor derivation divides by this inner product; only a code that quantizes // this residual (its own code or a prefix of it) makes it positive. - if (squared_norm != 0 && !(dot_product(residual, centered.data(), dim) > 0)) { + if (squared_norm != 0 && !(centered_ip > 0)) { throw std::invalid_argument( "code does not quantize residual: must be positive" ); @@ -155,10 +125,7 @@ v1::Factors code_factors_impl( factors.error, metric ); - 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); return factors; } 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/test_quant_layer1.py b/tests/python/test_quant_layer1.py index 6772fb4c..178319cf 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -220,8 +220,13 @@ def test_rejected_arguments(): 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=r"codes values must be below 2\*\*nbits = 8"): + 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" ): From b0c962c383bd1f199c9a6bd77756cf5323f4cc64 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 2 Nov 2026 20:40:04 +0800 Subject: [PATCH 08/16] docs(quant): use the indexes' error interval and an upper-bound threshold for candidates --- docs/docs/rabitq/quant_python.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index 6250260e..a64676d4 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -56,8 +56,9 @@ Factors are structured arrays of `quant.FACTORS_DTYPE` (`add`, `rescale`, `error ```python qry = quant.Query(q, metric="l2", centroid=c) est = qry.estimate_bitset(bitsets, f_prefix) # ESTIMATE_DTYPE[1000] -lower = est["distance"] - 2.0 * est["error"] -candidates = np.flatnonzero(lower <= np.partition(est["distance"], 10)[10]) +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]] ``` @@ -72,7 +73,7 @@ top = candidates[np.argsort(dist)[:10]] | `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. The probabilistic lower bound is `distance - m * error` for a multiplier `m` of your choice. `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)`. +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. From 2adec41eb75d6c6aafbf706b449f6d413d6eb100 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 2 Nov 2026 21:46:14 +0800 Subject: [PATCH 09/16] docs(quant): describe num_threads=0 as the CPU count capped by OMP_NUM_THREADS --- docs/docs/rabitq/quant_python.md | 2 +- python_bindings/quant_bindings.cpp | 27 ++++++++++++++++++--------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index a64676d4..bc29c10e 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -111,7 +111,7 @@ If your `ip` is itself approximate, `error` covers only the code's quantization ## Threads, copies and pickling -- Build-side calls (`Rotator.rotate`, `quantize`, `code_factors`, `pack_codes`, `unpack_codes`, `pack_fastscan_blocks`) default to `num_threads=0`, every available CPU; their results do not depend on the thread count. Scoring calls (`estimate_*`, `refine*`, `estimate_from_ip`) default to 1. +- Build-side calls (`Rotator.rotate`, `quantize`, `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`. If a call raises, a caller's `out=` is left unchanged. diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 869723a6..4f4de1a2 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -1596,7 +1596,8 @@ fast_quantization : bool, optional 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) uses every available CPU. Results do not + 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. @@ -1664,7 +1665,8 @@ metric : {'l2', 'ip'} 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) uses every available CPU. + 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,). @@ -1694,7 +1696,8 @@ codes : numpy.ndarray of uint8, shape (n, d) nbits : int Width of codes, 1 to 8. Split 9-bit codes first. num_threads : int, optional - Worker threads; 0 (the default) uses every available CPU. + 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. @@ -1721,7 +1724,8 @@ packed : numpy.ndarray of uint8, shape (n, d * nbits // 8) nbits : int Width of the packed codes, 1 to 8. num_threads : int, optional - Worker threads; 0 (the default) uses every available CPU. + 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. @@ -1798,7 +1802,8 @@ 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) uses every available CPU. + 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. @@ -1898,7 +1903,8 @@ terms : numpy.void of QUERY_TERMS_DTYPE nbits : int Bit width of the codes, 1 to 9. num_threads : int, optional - Worker threads; 1 by default, 0 uses every available CPU. + 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 @@ -1926,7 +1932,8 @@ bitsets : numpy.ndarray of uint8, shape (n, d // 8) factors : numpy.ndarray of FACTORS_DTYPE, shape (n,) The 1-bit (prefix) factors. num_threads : int, optional - Worker threads; 0 (the default) uses every available CPU. + 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. @@ -1968,7 +1975,8 @@ 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 uses every available CPU.)doc"; + Worker threads; 1 by default, 0 selects the available CPU count, capped by + OMP_NUM_THREADS.)doc"; py::class_(quant, "Query", kQueryDoc) .def( @@ -2141,7 +2149,8 @@ blocks : numpy.ndarray of uint8, shape (ceil(num_points / 32), d * 4 + 384) num_points : int Number of vectors in the blocks; padding lanes are not returned. num_threads : int, optional - Worker threads; 1 by default, 0 uses every available CPU. + 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 From d7b61aefa710a148d944861871559f9d31ec0bb0 Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Mon, 2 Nov 2026 23:24:43 +0800 Subject: [PATCH 10/16] fix(python): compile quant bindings without FP contraction so scoring is thread-count independent --- python_bindings/CMakeLists.txt | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index e870c486..9bf0cea9 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -28,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 From 2ef5d96cdfdbefefcb128b2409134c32a666a50a Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 3 Nov 2026 16:28:09 +0800 Subject: [PATCH 11/16] fix(core): reject a code that does not quantize its residual at the divisor The provisional code_factors accepts caller-supplied codes, so the factor derivation can now meet a code whose centered inner product with the residual is not positive. code_factors checked this with its own sequential sum, but code_factor_terms divides by a separately dispatched dot product guarded only by assert(). The two sums round differently, so near zero the check could pass while the divisor was not positive: debug builds aborted the process and release builds returned wrong factors. code_factor_terms now throws std::invalid_argument when the divisor itself is not positive, and code_factors drops its duplicate sum. The library's own codes match the residual's signs, so every term of the divisor is nonnegative and they never reach the new error. A NaN still passes through, as the per-vector encoders do not check that inputs are finite (quantizer.md); a test pins that behavior. --- include/rabitqlib/core/detail/rabitq_impl.hpp | 10 ++- src/core/provisional.cpp | 13 +--- .../rabitqlib/quantization/factor_test.cpp | 68 +++++++++++++++++++ 3 files changed, 79 insertions(+), 12 deletions(-) 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/src/core/provisional.cpp b/src/core/provisional.cpp index 64a37f03..419c1dd9 100644 --- a/src/core/provisional.cpp +++ b/src/core/provisional.cpp @@ -86,27 +86,20 @@ v1::Factors code_factors_impl( if (code == nullptr) { throw std::invalid_argument("code must not be null"); } - const float squared_norm = validation::checked_residual_squared_norm(residual, dim); + (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); - float centered_ip = 0; 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"); } - centered_ip += residual[i] * (static_cast(widened[i]) + cb); - } - // The factor derivation divides by this inner product; only a code that quantizes - // this residual (its own code or a prefix of it) makes it positive. - if (squared_norm != 0 && !(centered_ip > 0)) { - throw std::invalid_argument( - "code does not quantize residual: must be positive" - ); } + // 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) { 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); From 6a77072aed591c0b0a12a62ca7467c7416e73867 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 3 Nov 2026 16:31:01 +0800 Subject: [PATCH 12/16] refactor(quant): share the v1 query terms and unpacked inner product The bindings computed the query terms and the inner product of unpacked codes with their own copies of quant::v1's arithmetic, so a later change to v1 would not reach rabitqlib.quant and the two would drift apart silently. The arithmetic now lives once in quant::detail (core/encoder.hpp, next to full_est_dist_from_ip): query_corrections gives g_add and g_k1xsumq, and unpacked_ip sums in order and reports the first code above the bit width. quant::v1::prepare_query and compute use them with unchanged arithmetic order, checks and messages, and the bindings use them for compute_query_terms and the unpacked scoring paths. A comment on the helpers states that a change to them changes v1 and needs the v1 parity check against a pre-change build. v1 encode, prepare_query and compute output is bitwise identical before and after (2016 values over dims 64/96/128/130, total_bits 1-9, both metrics, with and without a centroid, exact and fast quantization, both code layouts), and so are the outputs of every rabitqlib.quant call on fixed inputs. The bindings compile without FP contraction and api.cpp does not, so Python distances agree with v1's within float32 rounding, as the comment on estimate_of says. --- include/rabitqlib/core/encoder.hpp | 46 ++++++++++++++++++++++++++ python_bindings/quant_bindings.cpp | 52 +++++++++++++++++------------- src/core/api.cpp | 28 +++++++++------- 3 files changed, 92 insertions(+), 34 deletions(-) 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/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 4f4de1a2..f135f533 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -18,7 +18,6 @@ #include #include #include -#include #include #include #include @@ -193,6 +192,21 @@ void check_code_values( } } +// 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; +} + // 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) { @@ -644,18 +658,19 @@ struct QueryTerms { 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 = - centroid == nullptr ? l2norm_sqr(query, d) : euclidean_sqr(query, centroid, d); + metric == METRIC_L2 ? corrections.g_add + : (centroid == nullptr ? l2norm_sqr(query, d) + : euclidean_sqr(query, centroid, d)); QueryTerms terms; - terms.g_add = metric == METRIC_L2 - ? distance_sqr - : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, d)); + terms.g_add = corrections.g_add; terms.g_error = std::sqrt(distance_sqr); - terms.g_k1xsumq = - std::accumulate(query, query + d, 0.0F) * fastscan::detail::code_center(1); + 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"); @@ -663,6 +678,10 @@ QueryTerms compute_query_terms( 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. inline Estimate estimate_of( float ip, size_t bits, const quant::v1::Factors& f, const QueryTerms& terms ) { @@ -674,17 +693,6 @@ inline Estimate estimate_of( ip}; } -// for unpacked codes, summed in order as quant::v1::compute does. This is -// the portable reference path, not a SIMD kernel. -template -inline float unpacked_ip(const float* query, const TP* code, size_t d) { - float ip = 0; - for (size_t j = 0; j < d; ++j) { - ip += query[j] * static_cast(code[j]); - } - return ip; -} - py::object query_terms( py::handle rotated_query, const std::string& metric, const py::object& centroid ) { @@ -1037,7 +1045,7 @@ class QuantQuery { 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { const size_t i = selected.row(k); estimates[k] = estimate_of( - unpacked_ip(query_.data(), view.row(i), d()), + unpacked_row_ip(query_.data(), view, i), nbits, records[i], terms_ @@ -1121,9 +1129,9 @@ class QuantQuery { index_detail::parallel_for( 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { const size_t i = selected.row(k); - const float ip_ex = - Packed ? ip_func(query_.data(), codes.row(i), d()) - : unpacked_ip(query_.data(), codes.row(i), d()); + 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]; distances[k] = quant::detail::full_est_dist_from_ip( (scale * estimates[k].refine_input) + ip_ex, diff --git a/src/core/api.cpp b/src/core/api.cpp index 1af0c56d..7efac9ba 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -6,7 +6,6 @@ #include #include #include -#include #include #include @@ -165,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; } @@ -362,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"); } @@ -392,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; From 4b34bf7fe3684b6d8998dfeeefab5c11a28461fa Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 3 Nov 2026 16:36:17 +0800 Subject: [PATCH 13/16] perf(quant): run row loops in ranges and check inputs in parallel passes index_detail::parallel_for hands out one index at a time from a shared counter. That suits index queries, but a row of rabitqlib.quant costs nanoseconds to microseconds, so the threads mostly contended for the counter: every row loop was slower on 4 threads than on 1 (estimate_from_ip on 200k rows: 18.8 ms against 1.6 ms). Row loops now go out in ranges, about eight per worker: at least 64 rows for the cheap loops (scoring, packing, input checks) and at least one row for rotate, quantize and code_factors, whose rows cost enough that a small batch should still use every worker. The input checks that ran serially before the parallel loop (code values, finite factors, ip, refine_input, FastScan block factors) now run in a parallel pass of their own, still before anything is written, so a rejected call leaves out= unchanged. Factors of contiguous rows are tested once per range, and code values with a vectorizable maximum. When several rows fail a check, the error reported is the first failing row's in call order (the lowest row, or the earliest entry of indices=), whatever the thread count; the guide says so. code_factors reports its first rejected row without recomputing it. pack_fastscan_blocks allocates its sign-byte scratch per range. It used to index a per-worker buffer with the worker number of a separately resolved thread count, which could run past the buffer if the two counts differed. Outputs are unchanged: every rabitqlib.quant call returns bitwise identical arrays on fixed inputs. Benchmark: synthetic standard-normal data, 200k rows, d=128 (codes at nbits 1 and 4), plus 100 rows at d=4096; Xeon Gold 6418H, GCC 14, Release, runtime dispatch (AVX-512); median of 3 interleaved rounds of 7 timed calls each, pinned to 8 idle cores of one NUMA node on a shared host; PR head -> this commit, ms: call 1 thread 8 threads estimate_bitset 2.40 -> 2.33 24.68 -> 1.00 estimate_codes (nbits 4) 27.48 -> 22.13 34.80 -> 3.43 refine (ex_bits 3) 25.79 -> 22.28 34.97 -> 3.40 estimate_from_ip 1.61 -> 1.27 20.68 -> 1.47 estimate_fastscan 1.65 -> 1.80 2.57 -> 0.96 pack_codes (nbits 1) 15.88 -> 9.67 24.63 -> 2.02 code_factors (nbits 4) 67.54 -> 57.08 45.72 -> 30.65 rotate 102.20 -> 96.27 86.18 -> 36.85 quantize, 100 x 4096, 9b 297.87 -> 299.82 46.42 -> 48.09 --- docs/docs/rabitq/quant_python.md | 2 +- python_bindings/quant_bindings.cpp | 496 +++++++++++++++++++---------- tests/python/test_quant_layer1.py | 22 ++ tests/python/test_quant_layer2.py | 103 ++++++ tests/python/test_quant_layer3.py | 19 ++ 5 files changed, 479 insertions(+), 163 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index bc29c10e..c5fc0eeb 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -115,7 +115,7 @@ If your `ip` is itself approximate, `error` covers only the code's quantization - 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`. If a call raises, a caller's `out=` is left unchanged. -- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. +- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. 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 diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index f135f533..22db653d 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -15,9 +15,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -172,23 +174,38 @@ Selection selection(const py::object& indices, size_t rows) { } template -void check_code_values( - const RowView& codes, const Selection& selected, size_t nbits, const char* name +[[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); - for (size_t k = 0; k < selected.count; ++k) { - const size_t i = selected.row(k); - const TP* row = codes.row(i); - for (size_t j = 0; j < codes.cols; ++j) { - if (row[j] > maximum) { - 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(row[j]) + " at (" + std::to_string(i) + ", " + - std::to_string(j) + ")" - ); - } - } + 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 + ); } } @@ -207,6 +224,12 @@ float unpacked_row_ip(const float* query, const RowView& codes, size_t i) { 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) { @@ -401,16 +424,137 @@ Records factors_input(py::handle value, size_t rows, const c return records; } -// Factors are read only for the selected rows; call with the GIL released. -void check_finite_factors( - const Records& factors, const Selection& selected -) { - for (size_t k = 0; k < selected.count; ++k) { - const quant::v1::Factors& f = factors[selected.row(k)]; - if (!std::isfinite(f.add) || !std::isfinite(f.rescale) || !std::isfinite(f.error)) { - throw std::invalid_argument("factors must contain only finite values"); +// --- 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. +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); +} + +// Writes value_at(k) to out[k] for k in [0, count), for rows already validated. +template +void write_rows(T* out, size_t count, size_t num_threads, ValueAt value_at) { + parallel_rows(count, num_threads, kCheapRows, [&](size_t begin, size_t end) { + for (size_t k = begin; k < end; ++k) { + out[k] = value_at(k); + } + }); } // --- Layer 1 -------------------------------------------------------------------- @@ -455,13 +599,13 @@ py::array quantize( { py::gil_scoped_release release; check_coordinates(input.data(), rows * d, d, "residuals"); - index_detail::parallel_for( - 0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + 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; }; @@ -479,10 +623,9 @@ py::tuple split_rows(const RowView& codes, size_t nbits, size_t prefix_bits) auto ex = new_rows(codes.rows, codes.cols); uint8_t* prefix_data = prefix.mutable_data(); uint8_t* ex_data = ex.mutable_data(); - const Selection every_row = all_rows(codes.rows); { py::gil_scoped_release release; - check_code_values(codes, every_row, nbits, "codes"); + check_code_values(codes, nbits, "codes", 1); for (size_t i = 0; i < codes.rows; ++i) { const TP* row = codes.row(i); for (size_t j = 0; j < codes.cols; ++j) { @@ -526,39 +669,34 @@ py::array code_factors_rows( 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"); - // A caller's out= is written only after every row succeeds. - std::vector staged(out.is_none() ? 0 : rows); - quant::v1::Factors* factors = out.is_none() ? data : staged.data(); - // code_factors checks each row's code values and that the code quantizes the - // residual; the binding adds the row number to its message. - const auto factors_of = [&](size_t i) { - return quant::provisional::code_factors( - residuals.data() + (i * d), centroid.data(), d, codes.row(i), nbits, metric - ); - }; - // Report the lowest rejected row, whatever the thread schedule. - std::atomic first_rejected{rows}; - index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { - try { - factors[i] = factors_of(i); - } catch (const std::invalid_argument&) { - size_t current = first_rejected.load(); - while (i < current && !first_rejected.compare_exchange_weak(current, i)) {} - } - }); - if (const size_t i = first_rejected.load(); i < rows) { + // 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; 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 { - (void)factors_of(i); - } catch (const std::invalid_argument& error) { - throw std::invalid_argument( - "codes row " + std::to_string(i) + " for residuals row " + - std::to_string(i) + ": " + error.what() + 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()); } - } + }); std::copy(staged.begin(), staged.end(), data); } return result; @@ -606,17 +744,17 @@ py::array pack_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}}); - const Selection every_row = all_rows(view.rows); { py::gil_scoped_release release; - check_code_values(view, every_row, nbits, "codes"); - index_detail::parallel_for( - 0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) { + // 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; } @@ -638,13 +776,13 @@ py::array unpack_codes( const auto codes = output_rows(out, view.rows, d, {{"packed", view.owner}}); { py::gil_scoped_release release; - index_detail::parallel_for( - 0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) { + 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; } @@ -745,12 +883,23 @@ py::array estimate_from_ip( const Selection every_row = all_rows(rows); { py::gil_scoped_release release; - if (!all_finite(ips.data(), rows)) { - throw std::invalid_argument("ip must contain only finite values"); - } - check_finite_factors(records, every_row); - index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) { - estimates[i] = estimate_of(ips.data()[i], nbits, records[i], query); + 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; @@ -797,18 +946,23 @@ py::array pack_fastscan_blocks( const Selection every_row = all_rows(n); { py::gil_scoped_release release; - check_finite_factors(records, every_row); - // One block of sign bytes per worker, in the order the FastScan packer reads: - // the first coordinate of each byte in its highest bit. - const size_t workers = - std::max(1, std::min(resolve_num_threads(num_threads), num_blocks)); - std::vector scratch(workers * fastscan::kBatchSize * codes.cols); - index_detail::parallel_for( - 0, num_blocks, num_threads, [&](size_t b, size_t thread) { + // 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); - uint8_t* msb_first = - scratch.data() + (thread * fastscan::kBatchSize * codes.cols); 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) { @@ -816,7 +970,7 @@ py::array pack_fastscan_blocks( } } uint8_t* block = blocks.row(b); - fastscan::pack_codes(d, msb_first, lanes, block); + 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 = @@ -826,7 +980,7 @@ py::array pack_fastscan_blocks( map.f_error()[lane] = f.error; } } - ); + }); } return blocks.owner; } @@ -1002,15 +1156,21 @@ class QuantQuery { const auto ip_func = select_excode_ipfunc(1); { py::gil_scoped_release release; - check_finite_factors(records, selected); - index_detail::parallel_for( - 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { - const size_t i = selected.row(k); - estimates[k] = estimate_of( - ip_func(query_.data(), codes.row(i), d()), 1, records[i], terms_ - ); - } + 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; }); @@ -1039,19 +1199,20 @@ class QuantQuery { ); { py::gil_scoped_release release; - check_code_values(view, selected, nbits, "codes"); - check_finite_factors(records, selected); - index_detail::parallel_for( - 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { - const size_t i = selected.row(k); - estimates[k] = estimate_of( - unpacked_row_ip(query_.data(), view, i), - nbits, - records[i], - terms_ - ); - } - ); + 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; }; @@ -1115,34 +1276,32 @@ class QuantQuery { const auto scale = static_cast(1U << ex_bits); { py::gil_scoped_release release; - if constexpr (!Packed) { - check_code_values(codes, selected, ex_bits, name); - } - check_finite_factors(records, selected); - for (size_t k = 0; k < estimates.rows; ++k) { + 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" ); } - } - index_detail::parallel_for( - 0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) { - 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]; - distances[k] = 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 - ); - } - ); + }); + 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; }); @@ -1168,51 +1327,64 @@ class QuantQuery { 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. - for (size_t b = 0; b < num_blocks; ++b) { + const auto finite_factors = [&](size_t b) { const ConstBatchDataMap map( reinterpret_cast(view.row(b)), d() ); std::array values{}; - const size_t lanes = std::min( - fastscan::kBatchSize, num_points - (b * fastscan::kBatchSize) - ); 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)) { - throw std::invalid_argument( - "blocks must contain only finite factors; block " + - std::to_string(b) + " does not" - ); + if (!all_finite(values.data(), lanes_of(b))) { + return false; } } - } - index_detail::parallel_for( - 0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) { - std::array distances{}; - std::array lower_bounds{}; - std::array ips{}; - std::array errors{}; - const auto* block = reinterpret_cast(view.row(b)); - split_batch_estdist( - block, - query, - d(), - distances.data(), - lower_bounds.data(), - ips.data(), - true + 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" ); - ConstBatchDataMap(block, d()) - .f_error() - .copy_to(errors.data(), errors.size()); - const size_t begin = b * fastscan::kBatchSize; - const size_t lanes = - std::min(fastscan::kBatchSize, num_points - begin); - for (size_t lane = 0; lane < lanes; ++lane) { - estimates[begin + lane] = { - distances[lane], errors[lane] * terms_.g_error, ips[lane]}; + } + }); + // Writes the estimates of block b. + 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]}; + } + }; + parallel_rows( + num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) { + for (size_t b = first; b < last; ++b) { + score_block(b); } } ); @@ -1412,11 +1584,11 @@ class QuantRotator { { py::gil_scoped_release release; check_coordinates(input.data(), rows * dim_, dim_, "x"); - index_detail::parallel_for( - 0, rows, num_threads, [&](size_t i, size_t /*thread*/) { + 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; } diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py index 178319cf..e2b51f53 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -296,3 +296,25 @@ def test_code_factors_leaves_out_unchanged_on_rejection(): 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) diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py index 0c94e7f9..c34add87 100644 --- a/tests/python/test_quant_layer2.py +++ b/tests/python/test_quant_layer2.py @@ -414,6 +414,109 @@ def test_out_overlapping_an_input_is_rejected(): ) +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"): diff --git a/tests/python/test_quant_layer3.py b/tests/python/test_quant_layer3.py index 133fd5cf..8083570b 100644 --- a/tests/python/test_quant_layer3.py +++ b/tests/python/test_quant_layer3.py @@ -190,6 +190,25 @@ def test_hand_built_blocks_with_non_finite_factors_are_rejected(): query.estimate_fastscan(blocks, num_points=40) +def test_rejected_inputs_leave_out_unchanged(): + # Inputs are checked in a pass of their own, before out= is written. + c, bitsets, factors, q = setup("l2", 300) + query = quant.Query(q, metric="l2", centroid=c) + broken = factors.copy() + broken["rescale"][250] = np.nan + packed = np.zeros((10, 768), np.uint8) + with pytest.raises(ValueError, match="factors must contain only finite values"): + quant.pack_fastscan_blocks(bitsets, broken, num_threads=4, out=packed) + assert not packed.any() + blocks = quant.pack_fastscan_blocks(bitsets, factors) + add = 96 * 4 # Lane 0 of the add factors. + blocks[7, add : add + 4] = np.frombuffer(np.float32(np.nan).tobytes(), np.uint8) + estimates = np.zeros(300, quant.ESTIMATE_DTYPE) + with pytest.raises(ValueError, match="block 7 does not"): + query.estimate_fastscan(blocks, num_points=300, num_threads=4, out=estimates) + assert not estimates.view(np.float32).any() + + def test_out_overlapping_an_input_is_rejected(): _, bitsets, factors, _ = setup("l2", 40) shared = np.zeros(2 * 768, np.uint8) From af831f2fdb75adca99c6793f8f4781457304d072 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 3 Nov 2026 16:40:33 +0800 Subject: [PATCH 14/16] fix(quant): raise OverflowError for estimates that are not finite The scoring calls returned an inf or NaN distance when finite but extreme factors overflowed, where quant::v1::compute raises std::overflow_error. NumPy sorts NaN last, so a ranking silently dropped those vectors. estimate_bitset, estimate_codes, estimate_from_ip, estimate_fastscan, refine and refine_packed now raise OverflowError for a distance that is not finite, with v1's message, and for an estimated error that is not finite. code_factors now adds the row number to its OverflowError, as it already did to its ValueError. Each range is written and then tested at once, and only a failing range is recomputed row by row to report the row, so the check adds no measurable cost (same benchmark as the previous commit, 1 thread, previous commit -> this one, ms: estimate_bitset 2.33 -> 2.41, estimate_codes 22.13 -> 22.41, estimate_from_ip 1.27 -> 1.17, estimate_fastscan 1.80 -> 1.56; differences are within run-to-run noise). Overflow is found after the values are written, so the out= contract becomes: a call that raises leaves out= unchanged, except after OverflowError or MemoryError, when its contents are unspecified. Finding overflow before writing would need a staging copy on every call that passes out=, the cost out= exists to avoid. Input errors are still raised before anything is written. The guide states this, documents which values raise OverflowError, and notes that split, which returns a tuple of new arrays, takes no out=. --- docs/docs/rabitq/quant_python.md | 4 +- python_bindings/quant_bindings.cpp | 72 +++++++++++++++++++++++++----- tests/python/test_quant_layer1.py | 21 +++++++++ tests/python/test_quant_layer2.py | 58 ++++++++++++++++++++++++ 4 files changed, 142 insertions(+), 13 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index c5fc0eeb..82576c9c 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -114,8 +114,8 @@ If your `ip` is itself approximate, `error` covers only the code's quantization - Build-side calls (`Rotator.rotate`, `quantize`, `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`. If a call raises, a caller's `out=` is left unchanged. -- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. 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. +- Every call that returns an array, except `split`, accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. 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 diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 22db653d..1294ec95 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -530,6 +530,10 @@ void check_code_values( } // 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)); @@ -547,14 +551,40 @@ bool finite_factor_range( return selected.indices == nullptr && all_finite_records(factors.data, begin, end); } -// Writes value_at(k) to out[k] for k in [0, count), for rows already validated. +// 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) { - parallel_rows(count, num_threads, kCheapRows, [&](size_t begin, size_t end) { - for (size_t k = begin; k < end; ++k) { - out[k] = value_at(k); - } - }); + 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 -------------------------------------------------------------------- @@ -679,7 +709,8 @@ py::array code_factors_rows( 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; the binding adds the row number to its message. + // 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) + ": "; @@ -695,6 +726,8 @@ py::array code_factors_rows( ); } 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); @@ -820,6 +853,7 @@ QueryTerms compute_query_terms( // 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 ) { @@ -1355,7 +1389,7 @@ class QuantQuery { ); } }); - // Writes the estimates of block b. + // Whether block b writes finite estimates. const auto score_block = [&](size_t b) { const auto* block = reinterpret_cast(view.row(b)); std::array distances{}; @@ -1380,11 +1414,27 @@ class QuantQuery { estimates[begin + lane] = { distances[lane], errors[lane] * terms_.g_error, ips[lane]}; } + return all_finite_records(estimates, begin, begin + lanes); }; - parallel_rows( - num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) { + for_each_range( + num_blocks, + num_threads, + kCheapRows, + [&](size_t first, size_t last) { for (size_t b = first; b < last; ++b) { - score_block(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]); } } ); diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py index e2b51f53..a83fc886 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -318,3 +318,24 @@ def test_code_factors_names_the_lowest_rejected_row_across_threads(): 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 index c34add87..6c0f541f 100644 --- a/tests/python/test_quant_layer2.py +++ b/tests/python/test_quant_layer2.py @@ -414,6 +414,64 @@ def test_out_overlapping_an_input_is_rejected(): ) +@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) From ebde6e3b488384d3a91acba6349839dadbc8070e Mon Sep 17 00:00:00 2001 From: Tamim Ehsan Date: Thu, 5 Nov 2026 17:20:32 +0800 Subject: [PATCH 15/16] feat(quant): give split num_threads and an out=(prefix, ex) pair --- docs/docs/rabitq/quant_python.md | 4 +- python_bindings/bindings_common.hpp | 9 ++- python_bindings/quant.pyi | 2 + python_bindings/quant_bindings.cpp | 97 ++++++++++++++++++++++------- tests/python/test_quant_layer1.py | 60 ++++++++++++++++++ tests/python/typing/public_api.py | 4 ++ 6 files changed, 149 insertions(+), 27 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index 82576c9c..cff73522 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -111,10 +111,10 @@ If your `ip` is itself approximate, `error` covers only the code's quantization ## Threads, copies and pickling -- Build-side calls (`Rotator.rotate`, `quantize`, `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. +- 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, except `split`, accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. A call that raises leaves a caller's `out=` unchanged, except after `OverflowError` or `MemoryError`, when its contents are unspecified. +- Every call that returns an array accepts `out=` with the exact dtype and shape; for `split` it is a `(prefix, ex)` pair, whose two arrays must not overlap each other either. It must not share memory with any input, or the call raises `ValueError`. 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 diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 3231e18c..f88da6fc 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -329,10 +329,13 @@ RowView exact_dtype_rows(py::handle value, const char* name) { } // 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. +// 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 expected = "out must be a writable NumPy array of dtype " + +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"; diff --git a/python_bindings/quant.pyi b/python_bindings/quant.pyi index b91bcc54..9cdbf2f6 100644 --- a/python_bindings/quant.pyi +++ b/python_bindings/quant.pyi @@ -207,6 +207,8 @@ def split( *, 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.""" diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index 1294ec95..ef393bc9 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -296,7 +296,9 @@ std::pair byte_extent(const py::array& array) { // 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) { +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; @@ -308,7 +310,7 @@ void check_output_disjoint(const py::array& out, Inputs inputs) { const auto [low, high] = byte_extent(py::reinterpret_borrow(input.array)); if (low < out_high && out_low < high) { - throw std::invalid_argument(std::string("out must not overlap ") + input.name); + throw std::invalid_argument(name + " must not overlap " + input.name); } } } @@ -316,12 +318,18 @@ void check_output_disjoint(const py::array& out, Inputs inputs) { // 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) { +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); - check_output_disjoint(view.owner, inputs); + auto view = checked_output_rows(out, rows, cols, name); + check_output_disjoint(view.owner, inputs, name); return view; } @@ -646,38 +654,75 @@ py::array quantize( } template -py::tuple split_rows(const RowView& codes, size_t nbits, size_t prefix_bits) { +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()) { + if (!py::isinstance(out) || py::len(out) != 2) { + 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); - auto prefix = new_rows(codes.rows, codes.cols); - auto ex = new_rows(codes.rows, codes.cols); - uint8_t* prefix_data = prefix.mutable_data(); - uint8_t* ex_data = ex.mutable_data(); { py::gil_scoped_release release; - check_code_values(codes, nbits, "codes", 1); - for (size_t i = 0; i < codes.rows; ++i) { - const TP* row = codes.row(i); - for (size_t j = 0; j < codes.cols; ++j) { - prefix_data[(i * codes.cols) + j] = static_cast(row[j] >> ex_bits); - ex_data[(i * codes.cols) + j] = static_cast(row[j] & mask); + // 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, ex); + return py::make_tuple(prefix.owner, ex.owner); } -py::tuple split(py::handle codes, size_t nbits, size_t prefix_bits) { +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); + 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); + return split_rows(view, nbits, prefix_bits, num_threads, out); } template @@ -1845,6 +1890,8 @@ codes : numpy.ndarray, shape (n, d) 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 @@ -1856,12 +1903,18 @@ 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.)doc" + ex_bits = nbits - prefix_bits. With out=, the two arrays of out.)doc" ); quant.def( diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py index a83fc886..497c85f6 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -97,6 +97,66 @@ def test_split_prefix_factors(nbits): 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"), + ( + (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) diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index 4f5726c3..a0b27263 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -289,6 +289,10 @@ def check_quant(x: npt.NDArray[np.float32]) -> None: 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]) From 8f079221f6158c9393e9044e2f33a5674386152f Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 5 Nov 2026 20:30:47 +0800 Subject: [PATCH 16/16] fix(quant): reject a None part in split's out= pair split(out=(prefix, None)) allocated a new ex array instead of raising, so a partial pair was accepted although the message, the stub and the guide require two arrays. Either part being None now raises the same ValueError as any other malformed pair. The guide's out= sentence is reordered so that "either" follows the input-overlap rule it refers to. --- docs/docs/rabitq/quant_python.md | 2 +- python_bindings/quant_bindings.cpp | 4 +++- tests/python/test_quant_layer1.py | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md index cff73522..12df7922 100644 --- a/docs/docs/rabitq/quant_python.md +++ b/docs/docs/rabitq/quant_python.md @@ -114,7 +114,7 @@ If your `ip` is itself approximate, `error` covers only the code's quantization - 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; for `split` it is a `(prefix, ex)` pair, whose two arrays must not overlap each other either. It must not share memory with any input, or the call raises `ValueError`. A call that raises leaves a caller's `out=` unchanged, except after `OverflowError` or `MemoryError`, when its contents are unspecified. +- 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 diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp index ef393bc9..a713042d 100644 --- a/python_bindings/quant_bindings.cpp +++ b/python_bindings/quant_bindings.cpp @@ -666,7 +666,9 @@ py::tuple split_rows( py::object prefix_out = py::none(); py::object ex_out = py::none(); if (!out.is_none()) { - if (!py::isinstance(out) || py::len(out) != 2) { + // 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) diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py index 497c85f6..d043870e 100644 --- a/tests/python/test_quant_layer1.py +++ b/tests/python/test_quant_layer1.py @@ -127,6 +127,8 @@ def test_split_rejects_bad_out_and_leaves_out_unchanged(): (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\)",