diff --git a/.github/scripts/test_ci_workflow_contracts.py b/.github/scripts/test_ci_workflow_contracts.py index e3905b47..b4b9d2c3 100644 --- a/.github/scripts/test_ci_workflow_contracts.py +++ b/.github/scripts/test_ci_workflow_contracts.py @@ -89,6 +89,7 @@ def selections(build, event): ] artifacts = set() + verdicts = set() for prefix, tag, check_name in ( ("linux", "manylinux_x86_64", "Linux x86-64 wheels"), ( @@ -131,21 +132,64 @@ def selections(build, event): self.assertNotIn("continue-on-error", entry) gate = job(workflow, f"{prefix}-wheels") self.assertEqual(gate["name"], check_name) - self.assertEqual(needs(gate), [f"{prefix}-wheel-build"]) + # The gate always runs so that a build skipped because the + # eligibility check failed cannot satisfy the required check. + self.assertEqual( + needs(gate), ["prepare", "changes", f"{prefix}-wheel-build"] + ) + self.assertEqual(gate["if"], "${{ always() }}") (verdict_step,) = steps(gate) self.assertEqual( - verdict_step["env"]["BUILD_RESULT"], - f"${{{{ needs.{prefix}-wheel-build.result }}}}", + verdict_step["env"], + { + "PREPARE_RESULT": "${{ needs.prepare.result }}", + "BUILD_SELECTED": "${{ needs.prepare.outputs.build }}", + "CHANGES_RESULT": "${{ needs.changes.result }}", + "WHEELS_SELECTED": "${{ needs.changes.outputs.wheels }}", + "BUILD_RESULT": f"${{{{ needs.{prefix}-wheel-build.result }}}}", + }, ) - for result in ("success", "failure", "cancelled", "skipped", ""): - with self.subTest(result=result): - verdict = subprocess.run( - ["bash", "-e", "-c", verdict_step["run"]], - env=dict(os.environ, BUILD_RESULT=result), - capture_output=True, - check=False, - ) - self.assertEqual(verdict.returncode == 0, result == "success") + verdicts.add(verdict_step["run"]) + # One shared verdict script, exercised over every input combination. + (command,) = verdicts + results = ("success", "failure", "cancelled", "skipped", "") + for prepare in results: + for selected in ("true", "false", ""): + for changes in results: + for wheels in ("true", "false", ""): + for build in results: + if build == "success": + expected = True + elif build != "skipped": + expected = False + elif prepare == "skipped": + expected = True # excluded by the workflow_run filter + else: + expected = prepare == "success" and ( + selected == "false" + or (changes == "success" and wheels == "false") + ) + with self.subTest( + prepare=prepare, + selected=selected, + changes=changes, + wheels=wheels, + build=build, + ): + verdict = subprocess.run( + ["bash", "-e", "-c", command], + env=dict( + os.environ, + PREPARE_RESULT=prepare, + BUILD_SELECTED=selected, + CHANGES_RESULT=changes, + WHEELS_SELECTED=wheels, + BUILD_RESULT=build, + ), + capture_output=True, + check=False, + ) + self.assertEqual(verdict.returncode == 0, expected) def test_installed_typing_checks_are_required(self): config = tomllib.loads((ROOT / "pyproject.toml").read_text()) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a9b1fd92..f3f55c19 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -85,16 +85,43 @@ jobs: if-no-files-found: error linux-wheels: - # Preserve the required check name and include every selected ABI. + # Preserve the required check name and include every selected ABI. The gate + # always runs: a skipped build passes only when nothing selected it, never + # because the eligibility check failed. name: Linux x86-64 wheels - needs: linux-wheel-build - if: ${{ always() && !cancelled() && needs.linux-wheel-build.result != 'skipped' }} + needs: [prepare, changes, linux-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.linux-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 linux-arm64-wheel-build: name: Linux ARM64 wheel builds @@ -129,14 +156,39 @@ jobs: linux-arm64-wheels: name: Linux ARM64 wheels (CPython 3.11–3.14) - needs: linux-arm64-wheel-build - if: ${{ always() && !cancelled() && needs.linux-arm64-wheel-build.result != 'skipped' }} + needs: [prepare, changes, linux-arm64-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.linux-arm64-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 windows-wheel-build: # Windows wheels take longest, so each ABI builds on its own runner. @@ -174,14 +226,39 @@ jobs: windows-wheels: name: Windows x86-64 wheels - needs: windows-wheel-build - if: ${{ always() && !cancelled() && needs.windows-wheel-build.result != 'skipped' }} + needs: [prepare, changes, windows-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.windows-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 macos-wheel-build: name: macOS arm64 wheel builds @@ -218,14 +295,39 @@ jobs: macos-wheels: # Preserve the protected check name; PRs test the oldest and newest ABI. name: macOS arm64 wheels (CPython 3.11–3.14) - needs: macos-wheel-build - if: ${{ always() && !cancelled() && needs.macos-wheel-build.result != 'skipped' }} + needs: [prepare, changes, macos-wheel-build] + if: ${{ always() }} runs-on: ubuntu-latest + timeout-minutes: 5 steps: - name: Require all selected wheels env: + PREPARE_RESULT: ${{ needs.prepare.result }} + BUILD_SELECTED: ${{ needs.prepare.outputs.build }} + CHANGES_RESULT: ${{ needs.changes.result }} + WHEELS_SELECTED: ${{ needs.changes.outputs.wheels }} BUILD_RESULT: ${{ needs.macos-wheel-build.result }} - run: test "$BUILD_RESULT" = success + run: | + if [ "$BUILD_RESULT" = success ]; then + exit 0 + fi + if [ "$BUILD_RESULT" != skipped ]; then + exit 1 + fi + # A skipped build passes when the eligibility check excluded this run, + # ran and selected no build, or a pull request changed nothing + # wheel-related. A failed eligibility check keeps the gate red. + if [ "$PREPARE_RESULT" = skipped ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$BUILD_SELECTED" = false ]; then + exit 0 + fi + if [ "$PREPARE_RESULT" = success ] && [ "$CHANGES_RESULT" = success ] && + [ "$WHEELS_SELECTED" = false ]; then + exit 0 + fi + exit 1 release-tag: name: Create version tag diff --git a/CMakeLists.txt b/CMakeLists.txt index ad0097c0..579dd400 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -63,11 +63,13 @@ endif() if(RABITQ_ENABLE_SANITIZERS) # Keep sanitizer instrumentation consistent across public inline code and # compiled components, including their runtime linkage. + # float-cast-overflow and float-divide-by-zero are not part of "undefined"; + # quantization converts floats to integers, so enable them explicitly. target_compile_options(rabitq_compile_options INTERFACE - "$<$:-fsanitize=address,undefined;-fno-omit-frame-pointer>" + "$<$:-fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero;-fno-omit-frame-pointer>" ) target_link_options(rabitq_compile_options INTERFACE - $<$:-fsanitize=address,undefined> + $<$:-fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero> ) endif() @@ -281,6 +283,7 @@ install(DIRECTORY include/rabitqlib/core DESTINATION "${rabitq_include_dir}") install(FILES include/rabitqlib/utils/bitops.hpp include/rabitqlib/utils/cpu_features.hpp + include/rabitqlib/utils/input_checks.hpp include/rabitqlib/utils/memory.hpp include/rabitqlib/utils/openmp.hpp include/rabitqlib/utils/packed.hpp diff --git a/docs/development.md b/docs/development.md index 36b0b170..b9574089 100644 --- a/docs/development.md +++ b/docs/development.md @@ -165,11 +165,14 @@ Clang: It recompiles the library sources from `compile_commands.json` with that build's flags (extra arguments such as `-O1` are appended) and fails when copies of a weak function -differ between ISA flag sets. Only the Release (`-O3`) configuration is clean and checked. -Lower optimization levels emit helpers the compiler otherwise inlines: Debug (`-O0`) -emits many standard-library and Eigen functions per source, and GCC at `-O1`/`-O2` -still shares two Eigen assignment loops between the FastScan estimator sources. The -script only sees `src/`, needs GNU binutils, and does not cover MSVC or ARM-only builds. +differ between ISA flag sets. CI checks the Release (`-O3`) configuration; with GCC 14 +the check also passes at `-O1` and `-O2 -g`, which is what the sanitizer and +RelWithDebInfo configurations compile. Debug (`-O0`) still emits standard-library +helpers (`std::min`, `std::round`, `std::fill`) per source and is not clean. The FastScan +estimator correction is written as plain loops rather than Eigen expressions so that no +ISA source instantiates a shared Eigen helper; keep it that way, since a static library +only hides such sharing while its generic objects precede the ISA objects. The script +only sees `src/`, needs GNU binutils, and does not cover MSVC or ARM-only builds. ### Static analysis @@ -331,6 +334,16 @@ Keep validation at the public operation boundary and reuse only invariants estab for that operation. Buffer contents and per-record factors may change between calls: validate them again even when reusing a query or workspace. +**Index inputs** + +Every index entry point that accepts float data, centroids or queries calls +`check_coordinates` from `utils/input_checks.hpp` after its shape checks and before +mutating state: non-finite values and magnitudes above the clustering bound raise +`std::invalid_argument`. The core `quantize_*` encoders and `Lut` rely on that (the +LUT additionally rejects non-finite rotated queries) because quantization converts +floats to integers, which is undefined for NaN. Keep the check at the boundary rather +than inside per-vector kernels. + **Python conversion** `ensure_*` returns a normalized NumPy owner retained across GIL release; diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index f43dea45..f196cf40 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -146,7 +146,12 @@ auto routing = ivf.initializer_type(); using twice the standard RaBitQ error margin, and recomputes exact distances for the retained candidates. Returned distances are exact for those centroids, but routing is approximate: pruning can miss a true nearest centroid. -- `hnsw` searches a graph of float32 centroids approximately. +- `hnsw` searches a graph of float32 centroids approximately. It suits workloads + whose `nprobe` is much smaller than the cluster count. Graph construction can + leave a few centroids unreachable, so points stored in those clusters are not + returned at any `nprobe`, and an `nprobe` close to the cluster count raises an + error instead of probing every cluster. Use `flat` or `flat_rabitq` when a + query should probe most clusters. All choices support L2 and inner product. Flat RaBitQ uses the IVF-RaBitQ rotation and padded domain, retains float32 centroids for refinement, and stores its mean and diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index 894d25d8..47c9f7b1 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -14,6 +14,14 @@ as the indexed data. Arrays whose data pointers are not aligned for float32 are copied into aligned storage before native processing; inputs remain unchanged. +Data, centroids and queries must be finite. Coordinates are also limited to the +magnitude the clustering APIs accept, `sqrt(sqrt(FLT_MAX) / (64 * dim))` (about +`5e7` for `dim = 128`), so that squared distances and their products stay inside +float32. `build`, `add` and `search` raise `ValueError` (C++: +`std::invalid_argument`) for NaN, infinite, or larger coordinates before changing +the index. The C++ `quantize_*` encoders in `core/encoder.hpp` assume their +callers have validated the input. + These limits apply to the Python indexes. The C++ [stable core API](../rabitq/core_api.md) accepts working dimensions from 2, with multiples of 32 required for `SplitPacked`. C++ [MatrixRotator](../rabitq/rotator.md) has separate dimension and storage requirements; diff --git a/include/rabitqlib/core/encoder.hpp b/include/rabitqlib/core/encoder.hpp index 25a77f87..32d8bdaf 100644 --- a/include/rabitqlib/core/encoder.hpp +++ b/include/rabitqlib/core/encoder.hpp @@ -122,6 +122,8 @@ inline void quantize_compact_ex_bits( cur_ex_data.f_rescale_ex() = f_rescale_ex; } +// The per-vector encoders assume finite inputs; callers validate once per batch +// (see utils/input_checks.hpp) because quantization converts floats to integers. inline void quantize_split_single( const float* data, const float* centroid, diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp index bd8ecf74..4b13e22a 100644 --- a/include/rabitqlib/fastscan/detail/lut.hpp +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -12,6 +12,7 @@ #include "rabitqlib/fastscan/detail/lut_prepare.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib { @@ -68,6 +69,10 @@ class Lut { "Lut::reset requires uint16 scratch for high-accuracy lookup tables" ); } + // Table quantization converts query-derived values to integers. + if (!all_finite(rotated_query, padded_dim)) { + throw std::invalid_argument("Lut::reset requires a finite query"); + } lut_.resize(shape.storage_length); table_length_ = shape.table_length; build_lut(rotated_query, padded_dim, lut_scratch, lut_u16_scratch, use_hacc); diff --git a/include/rabitqlib/fastscan/detail/qg_estimator.hpp b/include/rabitqlib/fastscan/detail/qg_estimator.hpp index f350aaaa..a578e40e 100644 --- a/include/rabitqlib/fastscan/detail/qg_estimator.hpp +++ b/include/rabitqlib/fastscan/detail/qg_estimator.hpp @@ -4,14 +4,14 @@ #include #include -#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/fastscan/data_layout.hpp" #include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/fastscan/query.hpp" namespace rabitqlib::fastscan::detail { // Share QG accumulation and correction while callers select the accumulator. -// TU-local linkage keeps each ISA's Eigen expression compiled with its own flags. +// The correction is a plain loop so each ISA translation unit vectorizes it with +// its own flags without instantiating shared weak helper functions. template static inline void qg_batch_estdist_impl( const char* batch_data, @@ -25,25 +25,20 @@ static inline void qg_batch_estdist_impl( cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim ); - ConstRowMajorArrayMap ip_arr( - accumulated_values.data(), 1, fastscan::kBatchSize - ); std::array f_add_values; std::array f_rescale_values; cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - - RowMajorArrayMap estimated_distances_arr( - estimated_distances, 1, fastscan::kBatchSize - ); - estimated_distances_arr = - f_add_arr + query_wrapper.g_add() + - (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + - query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); + const T delta = query_wrapper.delta(); + const T sum_vl_lut = query_wrapper.sum_vl_lut(); + const T g_add = query_wrapper.g_add(); + const T k1xsumq = query_wrapper.k1xsumq(); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const T ip = + delta * static_cast(accumulated_values[lane]) + sum_vl_lut + k1xsumq; + estimated_distances[lane] = + f_add_values[lane] + g_add + f_rescale_values[lane] * ip; + } } } // namespace rabitqlib::fastscan::detail diff --git a/include/rabitqlib/utils/input_checks.hpp b/include/rabitqlib/utils/input_checks.hpp new file mode 100644 index 00000000..ba883018 --- /dev/null +++ b/include/rabitqlib/utils/input_checks.hpp @@ -0,0 +1,67 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace rabitqlib { + +// True when every float32 value is finite. The integer reduction lets portable +// builds check several values per instruction without relaxing floating-point +// compiler options, and never allocates. Internal linkage keeps every ISA-flagged +// translation unit that includes this header on its own copy. +static inline bool all_finite(const float* values, size_t count) { + static_assert( + sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 + ); + uint32_t nonfinite = 0; + for (size_t i = 0; i < count; ++i) { + uint32_t bits; + std::memcpy(&bits, values + i, sizeof(bits)); + nonfinite |= static_cast((bits & 0x7f800000U) == 0x7f800000U); + } + return nonfinite == 0; +} + +template +static inline bool all_finite(const T* values, size_t count) { + size_t nonfinite = 0; + for (size_t i = 0; i < count; ++i) { + nonfinite += static_cast(!std::isfinite(values[i])); + } + return nonfinite == 0; +} + +// Largest coordinate magnitude whose squared distances, their products in graph +// pruning, and rotated residuals stay inside float32; shared with clustering. +static inline float max_coordinate_magnitude(size_t dim) { + return static_cast(std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(dim)) + )); +} + +// Reject vectors the indexes cannot quantize. `what` names the input in the message. +static inline void check_coordinates( + const float* values, size_t count, size_t dim, const char* what +) { + if (!all_finite(values, count)) { + throw std::invalid_argument(std::string(what) + " must contain only finite values"); + } + const float bound = max_coordinate_magnitude(dim); + size_t out_of_range = 0; + for (size_t i = 0; i < count; ++i) { + out_of_range += static_cast(!(std::fabs(values[i]) <= bound)); + } + if (out_of_range != 0) { + throw std::invalid_argument( + std::string(what) + " exceeds the safe float32 coordinate range" + ); + } +} + +} // namespace rabitqlib diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 4823f37e..98d15e44 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -94,6 +94,8 @@ class IvfIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. Examples -------- @@ -140,6 +142,8 @@ class IvfIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. Examples -------- @@ -388,6 +392,8 @@ class HnswIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __init__( @@ -424,6 +430,8 @@ class HnswIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __repr__(self) -> str: @@ -641,6 +649,8 @@ class SymqgIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __init__( @@ -669,6 +679,8 @@ class SymqgIndex: ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. + Data, centroids and queries must be finite float32 values within the index + coordinate range; other values raise ValueError. """ def __repr__(self) -> str: diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index ae47f2c3..668c98ef 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -311,7 +311,9 @@ random_seed : int, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc"; +supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError.)doc"; py::class_(m, "HnswIndex", kIndexDoc) .def( diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index e03f8b52..b2bf292d 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -353,6 +353,8 @@ Notes ----- Build or load the index before searching or saving. Concurrent reads are supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError. Examples -------- diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index ee1d3fd1..e144b5c3 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -232,7 +232,9 @@ quantization_bits : {0, 4, 8}, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc"; +supported; an overlapping mutation raises RuntimeError. +Data, centroids and queries must be finite float32 values within the index +coordinate range; other values raise ValueError.)doc"; static constexpr char kSearchDoc[] = R"doc(Search a batch of query vectors. diff --git a/src/core/api.cpp b/src/core/api.cpp index fc8737d8..88cfdf9e 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -4,7 +4,6 @@ #include #include #include -#include #include #include #include @@ -17,6 +16,7 @@ #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" namespace rabitqlib::quant::v1 { @@ -94,18 +94,7 @@ void check_input(const float* input, size_t dim) { check_buffer( {input, checked_multiply(dim, sizeof(float))}, checked_multiply(dim, sizeof(float)) ); - static_assert( - sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 - ); - // An integer reduction lets portable builds check several float32 values - // per instruction without relaxing floating-point compiler options. - uint32_t nonfinite = 0; - for (size_t i = 0; i < dim; ++i) { - uint32_t bits; - std::memcpy(&bits, input + i, sizeof(bits)); - nonfinite |= static_cast((bits & 0x7f800000U) == 0x7f800000U); - } - if (nonfinite != 0) { + if (!all_finite(input, dim)) { throw std::invalid_argument("Input vectors must contain only finite values"); } } diff --git a/src/fastscan/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp index f046ce7b..d78fe836 100644 --- a/src/fastscan/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -4,7 +4,6 @@ #include #include -#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/fastscan/data_layout.hpp" #include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/fastscan/query.hpp" @@ -12,7 +11,10 @@ namespace rabitqlib::simd { namespace { // Compile the complete batch estimator at each ISA so accumulation conversion and -// factor correction use the same vector width as the scanning kernels. +// factor correction use the same vector width as the scanning kernels. The +// correction is written as plain loops so each backend vectorizes it with its own +// flags: an Eigen expression here would instantiate identically named weak helper +// functions in every ISA translation unit, and the linker keeps one copy for all. template inline void split_batch_estdist_impl( const char* batch_data, @@ -25,9 +27,6 @@ inline void split_batch_estdist_impl( ) { ConstBatchDataMap cur_batch(batch_data, padded_dim); std::array accumulated_values{}; - RowMajorArrayMap accumulated_arr( - accumulated_values.data(), 1, fastscan::kBatchSize - ); if (use_hacc) { AccumulateHacc( cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim @@ -44,27 +43,20 @@ inline void split_batch_estdist_impl( cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); cur_batch.f_error().copy_to(f_error_values.data(), f_error_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - ConstRowMajorArrayMap f_error_arr( - f_error_values.data(), 1, fastscan::kBatchSize - ); - RowMajorArrayMap estimated_distances_arr( - estimated_distances, 1, fastscan::kBatchSize - ); - RowMajorArrayMap ip_x0_qr_arr(ip_x0_qr, 1, fastscan::kBatchSize); - RowMajorArrayMap lower_bounds_arr(lower_bounds, 1, fastscan::kBatchSize); - - ip_x0_qr_arr = query_wrapper.delta() * (accumulated_arr.template cast()) + - query_wrapper.sum_vl_lut(); - - estimated_distances_arr = f_add_arr + query_wrapper.g_add() + - f_rescale_arr * (ip_x0_qr_arr + query_wrapper.k1xsumq()); - - lower_bounds_arr = estimated_distances_arr - f_error_arr * query_wrapper.g_error(); + const float delta = query_wrapper.delta(); + const float sum_vl_lut = query_wrapper.sum_vl_lut(); + const float g_add = query_wrapper.g_add(); + const float k1xsumq = query_wrapper.k1xsumq(); + const float g_error = query_wrapper.g_error(); + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { + const float ip = delta * static_cast(accumulated_values[lane]) + sum_vl_lut; + const float estimated = + f_add_values[lane] + g_add + f_rescale_values[lane] * (ip + k1xsumq); + ip_x0_qr[lane] = ip; + estimated_distances[lane] = estimated; + lower_bounds[lane] = estimated - f_error_values[lane] * g_error; + } } } // namespace diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index d70bc774..984638ee 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -35,6 +35,7 @@ #include "rabitqlib/index/detail/visited_pool.hpp" #include "rabitqlib/index/hnsw/hnsw_quant.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/path.hpp" #include "rabitqlib/utils/space.hpp" @@ -518,6 +519,8 @@ void HierarchicalNSW::construct( throw std::invalid_argument("HNSW cluster ID is out of range"); } } + check_coordinates(data, num_points * dim_, dim_, "HNSW construction data"); + check_coordinates(centroids, num_clusters * dim_, dim_, "HNSW construction centroids"); try { num_clusters_ = num_clusters; @@ -913,6 +916,12 @@ std::vector>> HierarchicalNSW::search( if (k == 0) { throw std::invalid_argument("HNSW k must be greater than zero"); } + if (num_queries != 0) { + if (queries == nullptr) { + throw std::invalid_argument("HNSW search queries must not be null"); + } + check_coordinates(queries, num_queries * dim_, dim_, "HNSW search queries"); + } std::vector>> results(num_queries); index_detail::parallel_for( 0, @@ -1290,6 +1299,7 @@ std::vector HierarchicalNSW::add( if (num_points > buffer::kSearchBufferMaxPointCount - num_stored_points_) { throw std::invalid_argument("HNSW add point count exceeds the supported ID range"); } + check_coordinates(data, num_points * dim_, dim_, "HNSW add data"); std::vector labels(num_points); // Either the caller names the cluster of every point, or each one goes to its // nearest centroid, ordered the way a query is routed. diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 4c7c08ae..e906a006 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -38,6 +38,7 @@ #include "rabitqlib/index/ivf/cluster.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/path.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -210,6 +211,8 @@ void IVF::construct( id_lists[cluster_id].push_back(static_cast(i)); cluster_sizes[cluster_id] += 1; } + check_coordinates(data, max_elements_ * dim_, dim_, "IVF construction data"); + check_coordinates(centroids, num_clusters_ * dim_, dim_, "IVF construction centroids"); allocate_storage(cluster_sizes); create_cluster_views(cluster_sizes); @@ -356,6 +359,7 @@ void IVF::add( if (exceeds_storage(new_num_points)) { throw std::invalid_argument("IVF configuration exceeds addressable storage"); } + check_coordinates(data, num_points * dim_, dim_, "IVF add data"); // Assign every new point to a cluster std::vector assigned(num_points); @@ -977,6 +981,7 @@ void IVF::search( bool use_hacc ) const { validate_search(query, k, nprobe, result_ids); + check_coordinates(query, dim_, dim_, "IVF search query"); nprobe = std::min(nprobe, num_clusters_); // Probe at most every cluster std::vector rotated_query(padded_dim_); @@ -1039,6 +1044,7 @@ void IVF::search_batch( if (num_queries > max_values / dim_ || num_queries > max_values / k) { throw std::length_error("IVF query batch is too large"); } + check_coordinates(queries, num_queries * dim_, dim_, "IVF search queries"); nprobe = std::min(nprobe, num_clusters_); const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); const size_t num_workers = diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 1aeb5551..e3253e09 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -32,6 +32,7 @@ #include "rabitqlib/index/symqg/detail/search_scratch.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/bitops.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/path.hpp" @@ -452,6 +453,7 @@ void QuantizedGraph::search_impl( const float* query, uint32_t k, uint32_t* result_ids, float* distances, size_t ef_search ) { validate_search(query, k, result_ids, distances, ef_search); + check_coordinates(query, dim_, dim_, "QuantizedGraph search query"); detail::SearchScratch scratch(padded_dim_, max_degree_, ef_search, k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; @@ -498,6 +500,7 @@ void QuantizedGraph::search_batch_with_ef( search_impl(queries, k, result_ids, distances, ef_search); return; } + check_coordinates(queries, num_queries * dim_, dim_, "QuantizedGraph search queries"); const int num_workers = rabitqlib::detail::openmp_team_size(num_threads, num_queries); index_detail::ParallelException failure; @@ -808,6 +811,7 @@ void detail::search_with_workspace( const float* k1xsumq ) { graph.validate_search(query, k, result_ids, distances, graph.default_ef_search_); + check_coordinates(query, graph.dim_, graph.dim_, "QuantizedGraph search query"); graph.search_with_scratch_impl(query, k, result_ids, distances, scratch, hint, k1xsumq); } diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index 338a5efd..dcf1bbc4 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -25,6 +25,7 @@ #include "rabitqlib/index/symqg/detail/pipnn.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -100,6 +101,7 @@ void QGBuilder::reset(const float* data) { if (data == nullptr) { throw std::invalid_argument("QuantizedGraph build data must not be null"); } + check_coordinates(data, num_points_ * dim_, dim_, "QuantizedGraph build data"); if (initialization_ == QGInitialization::PiPNN) { auto initial = detail::build_initial_graph( data, num_points_, dim_, max_degree_, qg_.metric_type_, num_workers_ diff --git a/tests/consumer/check_options.cmake b/tests/consumer/check_options.cmake index 5df01dc5..d01252d2 100644 --- a/tests/consumer/check_options.cmake +++ b/tests/consumer/check_options.cmake @@ -7,8 +7,10 @@ include("${OPTIONS_FILE}") # Catch optimization, architecture and warnings leaking through any level of # the target graph, including installed exports. set(required_options ${openmp_options}) +# Keep in sync with RABITQ_ENABLE_SANITIZERS in the top-level CMakeLists.txt. +set(sanitizer_flag -fsanitize=address,undefined,float-cast-overflow,float-divide-by-zero) if(expect_sanitizers) - list(APPEND required_options -fsanitize=address,undefined -fno-omit-frame-pointer) + list(APPEND required_options ${sanitizer_flag} -fno-omit-frame-pointer) endif() foreach(flag IN LISTS consumer_options) if(NOT flag STREQUAL "" AND NOT flag IN_LIST required_options) @@ -35,6 +37,6 @@ if(expect_msvc) endif() endforeach() endif() -if(expect_sanitizers AND NOT "-fsanitize=address,undefined" IN_LIST consumer_link_options) +if(expect_sanitizers AND NOT sanitizer_flag IN_LIST consumer_link_options) message(FATAL_ERROR "The instrumented library must propagate sanitizer runtime linkage") endif() diff --git a/tests/python/test_input_boundaries.py b/tests/python/test_input_boundaries.py index f9c5ff26..1fbb4dc4 100644 --- a/tests/python/test_input_boundaries.py +++ b/tests/python/test_input_boundaries.py @@ -277,3 +277,62 @@ def test_array_conversion_reports_target_dtype(built_index, value, description): == "queries must be convertible to a NumPy array of dtype float32; got " + description ) + + +@pytest.mark.parametrize( + "value,message", + [ + (np.nan, "must contain only finite values"), + (np.inf, "must contain only finite values"), + (-np.inf, "must contain only finite values"), + (1e30, "exceeds the safe float32 coordinate range"), + ], +) +def test_nonfinite_or_oversized_coordinates_are_rejected(built_index, value, message): + index, queries = built_index + before = index.num_points + bad_queries = np.array(queries, copy=True) + bad_queries[1, 3] = value + with pytest.raises(ValueError, match=message): + index.search(bad_queries, 3, 32) + extra = np.random.default_rng(7).normal(size=(2, 64)).astype(np.float32) + extra[0, 0] = value + if isinstance(index, IvfIndex): + with pytest.raises(ValueError, match=message): + index.add(extra) + elif isinstance(index, HnswIndex): + # The shared fixture is at capacity, which is checked before coordinates. + data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + spare = HnswIndex(64, 82, nbits=4) + spare.build( + data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32) + ) + with pytest.raises(ValueError, match=message): + spare.add(extra) + assert spare.num_points == 80 + assert index.num_points == before + ids, distances = index.search(queries, 3, 32) + assert ids.shape == (2, 3) and np.all(np.isfinite(distances)) + + +@pytest.mark.parametrize("index_kind", ["ivf", "hnsw", "symqg"]) +@pytest.mark.parametrize("value", [np.nan, np.inf, 1e30]) +def test_build_rejects_bad_coordinates_without_building(index_kind, value): + data = np.random.default_rng(3).normal(size=(80, 64)).astype(np.float32) + bad_data = data.copy() + bad_data[5, 3] = value + if index_kind == "symqg": + index = SymqgIndex(64, 32) + with pytest.raises(ValueError): + index.build(bad_data, 40) + else: + centroid = data.mean(axis=0, keepdims=True) + ids = np.zeros(80, dtype=np.uint32) + index = IvfIndex(64, 80, 1, 4) if index_kind == "ivf" else HnswIndex(64, 80) + with pytest.raises(ValueError): + index.build(bad_data, centroid, ids) + bad_centroid = centroid.copy() + bad_centroid[0, 0] = value + with pytest.raises(ValueError): + index.build(data, bad_centroid, ids) + assert not index.is_built diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index 8d2401c5..a529b9c1 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -1289,17 +1289,24 @@ TEST(QGKMeansTest, RejectsInvalidNativeParametersBeforeAllocation) { namespace rabitqlib::qgkmeans { TEST(QGAssignerTest, PropagatesSearchFailureFromWorkers) { + // Out-of-range queries are rejected by the per-query search entry that each + // worker calls, so the failure originates inside the parallel region. const std::vector centroids(33 * 64, 0.0F); const std::vector data(128 * 64, 1e19F); std::vector assignments(128, kPidMax); std::vector distances(128, -1.0F); QGAssigner assigner(64, 33, 32, 33, 33, 1, 2); - EXPECT_THROW( + try { assigner.assign( centroids.data(), data.data(), 128, assignments.data(), distances.data() - ), - std::runtime_error - ); + ); + ADD_FAILURE() << "expected the worker failure to propagate"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ( + error.what(), + "QuantizedGraph search query exceeds the safe float32 coordinate range" + ); + } EXPECT_EQ(assignments.front(), kPidMax); EXPECT_EQ(distances.front(), -1.0F); } diff --git a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp index 11b664be..bf1cf3ad 100644 --- a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp +++ b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include "rabitqlib/core/detail/simd/space_dispatch.hpp" @@ -262,5 +263,36 @@ TEST(FastScanLutPreparationTest, FloatLutPreservesStandardAndHaccCodes) { TEST(FastScanLutPreparationTest, DoubleLutRetainsOriginalPreparation) { check_prepared_lut(); } + +TEST(FastScanLutPreparationTest, RejectsNonFiniteQueriesForFloatAndDouble) { + constexpr size_t kPaddedDim = 64; + const std::vector query(kPaddedDim, 0.5F); + const std::vector double_query(kPaddedDim, 0.5); + const float poison[] = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + }; + for (float value : poison) { + auto bad = query; + bad[7] = value; + EXPECT_THROW((Lut(bad.data(), kPaddedDim)), std::invalid_argument); + EXPECT_THROW((Lut(bad.data(), kPaddedDim, true)), std::invalid_argument); + auto bad_double = double_query; + bad_double[7] = static_cast(value); + EXPECT_THROW((Lut(bad_double.data(), kPaddedDim)), std::invalid_argument); + EXPECT_THROW( + (Lut(bad_double.data(), kPaddedDim, true)), std::invalid_argument + ); + } + // A rejected reset keeps the previously prepared table. + Lut lut(query.data(), kPaddedDim); + std::vector scratch(kPaddedDim * 4); + auto bad = query; + bad[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(lut.reset(bad.data(), kPaddedDim, scratch.data()), std::invalid_argument); + EXPECT_EQ(lut.padded_dim(), kPaddedDim); +} + } // namespace } // namespace rabitqlib::fastscan diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 17ec4226..50d9e044 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -11,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +27,7 @@ #include "rabitqlib/index/hnsw/hnsw_quant.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "temp_file.hpp" @@ -1324,6 +1327,109 @@ TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { } } +TEST(HnswInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 16; + constexpr size_t kClusters = 2; + std::mt19937 gen(905); + std::normal_distribution normal; + std::vector data(kCount * kDim); + std::vector centroids(kClusters * kDim); + std::vector extra(2 * kDim); + std::vector queries(2 * kDim); + for (auto* values : {&data, ¢roids, &extra, &queries}) { + for (auto& value : *values) { + value = normal(gen); + } + } + std::vector cluster_ids(kCount); + for (size_t i = 0; i < kCount; ++i) { + cluster_ids[i] = static_cast(i % kClusters); + } + const float bound = max_coordinate_magnitude(kDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + HierarchicalNSW index(kCount + 2, kDim, 4, 8, 32); + for (float value : rejected) { + auto bad_data = data; + bad_data[kDim + 3] = value; + expect_message( + [&] { + index.construct( + kClusters, + centroids.data(), + kCount, + bad_data.data(), + cluster_ids.data(), + 1 + ); + }, + message_for(value, "HNSW construction data") + ); + auto bad_centroids = centroids; + bad_centroids.back() = value; + expect_message( + [&] { + index.construct( + kClusters, + bad_centroids.data(), + kCount, + data.data(), + cluster_ids.data(), + 1 + ); + }, + message_for(value, "HNSW construction centroids") + ); + EXPECT_EQ(index.num_stored_points(), 0U); + } + ASSERT_NO_THROW(index.construct( + kClusters, centroids.data(), kCount, data.data(), cluster_ids.data(), 1 + )); + ASSERT_EQ(index.num_stored_points(), kCount); + + for (float value : rejected) { + auto bad_extra = extra; + bad_extra[kDim] = value; + expect_message( + [&] { index.add(bad_extra.data(), 2, nullptr, false, 1); }, + message_for(value, "HNSW add data") + ); + EXPECT_EQ(index.num_stored_points(), kCount); + auto bad_queries = queries; + bad_queries[kDim + 1] = value; + expect_message( + [&] { index.search(bad_queries.data(), 2, 3, 10, 1); }, + message_for(value, "HNSW search queries") + ); + } + // The bound itself is accepted everywhere. + auto edge = extra; + edge[0] = bound; + edge[kDim] = -bound; + ASSERT_NO_THROW(index.add(edge.data(), 2, nullptr, false, 1)); + EXPECT_EQ(index.num_stored_points(), kCount + 2); + auto edge_queries = queries; + edge_queries[0] = bound; + EXPECT_EQ(index.search(edge_queries.data(), 2, 3, 10, 1).size(), 2U); +} + } // namespace } // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 6bcc640a..621a9d66 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -27,6 +27,7 @@ #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "temp_file.hpp" namespace rabitqlib::ivf { @@ -1444,6 +1445,109 @@ TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { } } +TEST(IvfInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kCount = 40; + const auto points = random_rows(kCount, 901); + const auto centroids = random_rows(kDynNumClusters, 902); + std::vector assigned(kCount); + for (size_t i = 0; i < kCount; ++i) { + assigned[i] = static_cast(i % kDynNumClusters); + } + const float bound = max_coordinate_magnitude(kDynDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + IVF index(kCount, kDynDim, kDynNumClusters, 4); + for (float value : rejected) { + auto bad_points = points; + bad_points[kDynDim + 3] = value; + expect_message( + [&] { + index.construct( + bad_points.data(), centroids.data(), assigned.data(), false, 1 + ); + }, + message_for(value, "IVF construction data") + ); + auto bad_centroids = centroids; + bad_centroids.back() = value; + expect_message( + [&] { + index.construct( + points.data(), bad_centroids.data(), assigned.data(), false, 1 + ); + }, + message_for(value, "IVF construction centroids") + ); + EXPECT_EQ(index.num_points(), 0U); + } + ASSERT_NO_THROW( + index.construct(points.data(), centroids.data(), assigned.data(), false, 1) + ); + ASSERT_EQ(index.num_points(), kCount); + + const auto extra = random_rows(2, 903); + const auto queries = random_rows(2, 904); + std::vector result_ids(2 * 3); + std::vector distances(2 * 3); + for (float value : rejected) { + auto bad_extra = extra; + bad_extra[kDynDim] = value; + expect_message( + [&] { index.add(bad_extra.data(), 2, nullptr, false, 1); }, + message_for(value, "IVF add data") + ); + EXPECT_EQ(index.num_points(), kCount); + auto bad_query = queries; + bad_query[0] = value; + expect_message( + [&] { + index.search(bad_query.data(), 3, 1, result_ids.data(), distances.data()); + }, + message_for(value, "IVF search query") + ); + auto bad_batch = queries; + bad_batch[kDynDim + 1] = value; + expect_message( + [&] { + index.search_batch( + bad_batch.data(), 2, 3, 1, result_ids.data(), distances.data() + ); + }, + message_for(value, "IVF search queries") + ); + } + // The bound itself is accepted everywhere. + auto edge = extra; + edge[0] = bound; + edge[kDynDim] = -bound; + ASSERT_NO_THROW(index.add(edge.data(), 2, nullptr, false, 1)); + EXPECT_EQ(index.num_points(), kCount + 2); + auto edge_query = queries; + edge_query[0] = bound; + ASSERT_NO_THROW( + index.search(edge_query.data(), 3, 1, result_ids.data(), distances.data()) + ); + ASSERT_NO_THROW( + index.search_batch(edge_query.data(), 2, 3, 1, result_ids.data(), distances.data()) + ); +} + } // namespace } // namespace rabitqlib::ivf diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 2359a5ed..32dae1a3 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -40,6 +40,7 @@ #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" #include "temp_file.hpp" @@ -2116,3 +2117,78 @@ TEST(QgSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { } } } // namespace rabitqlib::symqg + +namespace rabitqlib::symqg { +TEST(QgInputTest, RejectsNonFiniteAndOversizedCoordinatesBeforeMutation) { + constexpr size_t kCount = 33, kDim = 64; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.17F); + } + const float bound = max_coordinate_magnitude(kDim); + const std::array rejected{ + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 2.0F * bound, + }; + const auto message_for = [](float value, const char* what) { + return std::string(what) + (std::isfinite(value) + ? " exceeds the safe float32 coordinate range" + : " must contain only finite values"); + }; + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + + for (float value : rejected) { + auto bad_data = data; + bad_data[kDim + 5] = value; + for (auto init : {QGInitialization::Random, QGInitialization::PiPNN}) { + QuantizedGraph graph(kCount, kDim, 32); + expect_message( + [&] { QGBuilder(graph, 32, bad_data.data(), 1, init); }, + message_for(value, "QuantizedGraph build data") + ); + EXPECT_EQ(graph.num_points(), 0U); + } + } + + QuantizedGraph graph(kCount, kDim, 32); + QGBuilder builder(graph, 32, data.data(), 1, QGInitialization::Random); + builder.build(2); + graph.set_ef(10); + const std::vector queries(data.begin(), data.begin() + 2 * kDim); + std::vector result_ids(2 * 3); + std::vector distances(2 * 3); + for (float value : rejected) { + auto bad_query = queries; + bad_query[0] = value; + expect_message( + [&] { graph.search(bad_query.data(), 3, result_ids.data(), distances.data()); }, + message_for(value, "QuantizedGraph search query") + ); + auto bad_batch = queries; + bad_batch[kDim + 1] = value; + expect_message( + [&] { + graph.search_batch( + bad_batch.data(), 2, 3, result_ids.data(), distances.data(), 1 + ); + }, + message_for(value, "QuantizedGraph search queries") + ); + } + auto edge_query = queries; + edge_query[0] = bound; + ASSERT_NO_THROW( + graph.search(edge_query.data(), 3, result_ids.data(), distances.data()) + ); + ASSERT_NO_THROW( + graph.search_batch(edge_query.data(), 2, 3, result_ids.data(), distances.data(), 1) + ); +} +} // namespace rabitqlib::symqg diff --git a/tests/unit/rabitqlib/utils/input_checks_test.cpp b/tests/unit/rabitqlib/utils/input_checks_test.cpp new file mode 100644 index 00000000..f0abd877 --- /dev/null +++ b/tests/unit/rabitqlib/utils/input_checks_test.cpp @@ -0,0 +1,92 @@ +#include "rabitqlib/utils/input_checks.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace rabitqlib { +namespace { + +TEST(InputChecksTest, DetectsEveryNonFiniteValueForFloatAndDouble) { + const std::vector finite{ + 0.0F, + -0.0F, + std::numeric_limits::denorm_min(), + std::numeric_limits::max(), + -std::numeric_limits::max(), + 1.5F, + }; + EXPECT_TRUE(all_finite(finite.data(), finite.size())); + EXPECT_TRUE(all_finite(finite.data(), 0)); + const std::vector finite_double(finite.begin(), finite.end()); + EXPECT_TRUE(all_finite(finite_double.data(), finite_double.size())); + + const float poison[] = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::signaling_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + }; + for (float value : poison) { + for (size_t position : {size_t{0}, finite.size() / 2, finite.size() - 1}) { + auto values = finite; + values[position] = value; + EXPECT_FALSE(all_finite(values.data(), values.size())); + const std::vector doubles(values.begin(), values.end()); + EXPECT_FALSE(all_finite(doubles.data(), doubles.size())); + } + } +} + +TEST(InputChecksTest, CoordinateBoundMatchesClusteringAndAcceptsItself) { + constexpr size_t kDim = 128; + const double expected = std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(kDim)) + ); + const float bound = max_coordinate_magnitude(kDim); + EXPECT_NEAR(static_cast(bound), expected, expected * 1e-6); + EXPECT_LT(max_coordinate_magnitude(2 * kDim), bound); + + std::vector values(kDim, 0.25F); + values[0] = bound; + values[1] = -bound; + EXPECT_NO_THROW(check_coordinates(values.data(), values.size(), kDim, "data")); + EXPECT_NO_THROW(check_coordinates(values.data(), 0, kDim, "data")); + + const auto expect_message = [](const auto& call, const std::string& message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::invalid_argument& error) { EXPECT_EQ(error.what(), message); } + }; + values[kDim - 1] = std::nextafter(bound, std::numeric_limits::infinity()); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data exceeds the safe float32 coordinate range" + ); + values[kDim - 1] = -std::nextafter(bound, std::numeric_limits::infinity()); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "queries"); }, + "queries exceeds the safe float32 coordinate range" + ); + // Non-finite values are reported as such even when the range is also exceeded. + values[kDim - 1] = std::numeric_limits::infinity(); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data must contain only finite values" + ); + values[kDim - 1] = std::numeric_limits::quiet_NaN(); + expect_message( + [&] { check_coordinates(values.data(), values.size(), kDim, "data"); }, + "data must contain only finite values" + ); +} + +} // namespace +} // namespace rabitqlib