diff --git a/.clang-tidy b/.clang-tidy index 5ba7be3b..4419d66d 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -19,12 +19,30 @@ Checks: > performance-*, portability-*, readability-duplicate-include, + readability-identifier-naming, readability-inconsistent-declaration-parameter-name, readability-misleading-indentation, readability-redundant-control-flow, -portability-simd-intrinsics, -clang-analyzer-security.insecureAPI.DeprecatedOrUnsafeBufferHandling -# Keep the selected first-party baseline regression-free. Repository-wide -# naming checks remain deferred because they would change public APIs. -WarningsAsErrors: "*" +# Keep the established checks blocking. Naming starts as warnings while existing +# internal/parameter names are migrated; do not bulk-fix public interfaces. +WarningsAsErrors: "*,-readability-identifier-naming" + +CheckOptions: + # Check locals, parameters and private fields. Public member/function names, + # types, enums and macros retain their current compatibility contracts. + readability-identifier-naming.LocalVariableCase: lower_case + readability-identifier-naming.ParameterCase: lower_case + # M is the existing HNSW graph-degree parameter, including its public constructor. + readability-identifier-naming.ParameterIgnoredRegexp: '^M$' + readability-identifier-naming.LocalConstantCase: lower_case + # Runtime const values stay snake_case. Also accept existing named constants + # such as kBatchSize; const/constexpr alone does not express that distinction. + readability-identifier-naming.LocalConstantIgnoredRegexp: '^k[A-Z][a-zA-Z0-9]*$' + readability-identifier-naming.PrivateMemberCase: lower_case + readability-identifier-naming.PrivateMemberSuffix: '_' + # Preserve these existing graph parameters and formula coefficient names. + # Keep exceptions explicit rather than exempting all capitalized identifiers. + readability-identifier-naming.PrivateMemberIgnoredRegexp: '^(M|maxM|maxM0|G_add|G_error|G_k1xSumq|G_kbaseSumq|G_kbxSumq)_$' diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 64282272..22c7f30a 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -254,6 +254,48 @@ jobs: - name: Test downstream consumer run: ctest --test-dir build-consumer --output-on-failure + core-consumer: + name: Core and FastScan consumers (${{ matrix.os }}, OpenMP ${{ matrix.openmp }}) + needs: changes + if: ${{ !cancelled() && (needs.changes.result != 'success' || needs.changes.outputs.cpp == 'true') }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, ubuntu-24.04-arm, macos-15, windows-2025-vs2026] + openmp: ["ON", "OFF"] + runs-on: ${{ matrix.os }} + timeout-minutes: 45 + steps: + - uses: actions/checkout@v7 + - name: Install Apple OpenMP + if: runner.os == 'macOS' && matrix.openmp == 'ON' + run: brew install libomp + - name: Test external core-only integrations (Linux) + if: runner.os == 'Linux' + run: >- + cmake -DRABITQ_SOURCE_DIR="$PWD" + -DWORK_DIR="$RUNNER_TEMP/rabitq-core" + -DRABITQ_USE_OPENMP=${{ matrix.openmp }} + -P tests/core_consumer/run.cmake + - name: Test external core-only integrations (macOS) + if: runner.os == 'macOS' + run: >- + cmake -DRABITQ_SOURCE_DIR="$PWD" + -DWORK_DIR="$RUNNER_TEMP/rabitq-core" + -DRABITQ_USE_OPENMP=${{ matrix.openmp }} + -DOpenMP_ROOT="$(brew --prefix)/opt/libomp" + -P tests/core_consumer/run.cmake + - name: Test external core-only integrations (Windows) + if: runner.os == 'Windows' + env: + CMAKE_GENERATOR: Visual Studio 18 2026 + CMAKE_GENERATOR_PLATFORM: x64 + run: >- + cmake "-DRABITQ_SOURCE_DIR=$PWD" + "-DWORK_DIR=$env:RUNNER_TEMP/rabitq-core" + -DRABITQ_USE_OPENMP=${{ matrix.openmp }} + -P tests/core_consumer/run.cmake + linux-arm64: name: Native Linux ARM64 C++ needs: changes diff --git a/AGENTS.md b/AGENTS.md index 5c3452f0..8f10047d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,24 +15,69 @@ factors estimate L2 distance or inner product. | Path | Responsibility | | --- | --- | -| `include/rabitqlib/quantization/` | Encoding, packing, reconstruction, and byte layouts | -| `include/rabitqlib/fastscan/` | FastScan interfaces and high-accuracy scanning | +| `include/rabitqlib/core/api.hpp` | Stable float32 v1 API; code/factors separate, caller workspace, source compatibility | +| `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 | +| `include/rabitqlib/index/types.hpp` | Point IDs and ANN candidate values; no Eigen dependency | | `include/rabitqlib/clustering/` | RaBitQKMeans flat assignment, QGKMeans graph assignment, and shared Lloyd training | | `include/rabitqlib/index/{ivf,hnsw,symqg}/` | Index construction, persistence, and search | -| `include/rabitqlib/index/{query,estimator}.hpp` | Query state and distance estimation | -| `include/rabitqlib/simd/`, `src/simd/` | Kernel declarations, implementations, and dispatch | -| `src/index/` | Compiled SymphonyQG implementation, HNSW search kernels, and IVF candidate insertion | +| `include/rabitqlib/{core,fastscan,index}/detail/` | Module-owned implementation headers; no public stability promise | +| `src/core/` | Compiled stable core v1 API implementation | +| `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/` | Rotation, allocation, buffers, I/O, and helpers | +| `include/rabitqlib/utils/` | Shared allocation, I/O and numerical helpers; index buffers live in `index/detail/` | | `python_bindings/` | pybind11 extension and index wrappers | | `sample/cpp/`, `sample/python/` | Usage examples | -| `tests/unit/`, `tests/integration/`, `tests/python/` | C++ and Python tests | +| `tests/unit/`, `tests/python/` | C++ reference/unit tests and Python boundary/integration tests | | `docs/docs/` | MkDocs site sources | See [README.md](README.md) for usage, [DEVELOPMENT.md](DEVELOPMENT.md) for tool setup and [implementation recipes](DEVELOPMENT.md#implementation-recipes), and [tests/README.md](tests/README.md) for build and test commands. +CMake exposes `rabitqlib::core`, `rabitqlib::fastscan` (depending on core), and +`rabitqlib::index` (depending on FastScan). The old `rabitqlib::rabitqlib` target and +forwarding header paths have been removed in the 0.7.0 breaking refactor. +Optional Python bindings expose the build-tree MODULE alias +`rabitqlib::python` and keep `import rabitqlib`. For pure core, disable both +`RABITQ_BUILD_INDEX` and `RABITQ_BUILD_FASTSCAN`. Canonical core headers must not include +FastScan or index headers; canonical FastScan headers must not include index headers. +Install only enabled components and their required header dependencies. Do not restore +umbrella targets/forwarders to bypass migration failures. Shared index scheduling belongs +in `index/detail/parallel_for.hpp`; HNSW must not import IVF to access this utility. +The former `defines.hpp` is removed: include metric, scalar-quantizer, Eigen or index +types explicitly. Core must not expose ANN IDs/candidates or QG batch adapters. +`BatchDataMap`, `QGBatchDataMap`, their encoders and batch packing/permutation belong +to FastScan, including their uses from IVF, clustering and SymphonyQG. Core must not +declare `fastscan::` functions/constants or embed these layouts. Generic unaligned +value/array access lives in `utils/packed.hpp`. `exact_nn` belongs to +`index/detail/exact_nn.hpp`. +See [the migration contract](docs/docs/modular_migration.md). + +IVF/HNSW/SymphonyQG public index headers must not import encoding, scoring, rotator +implementation, search scratch, Eigen, or hnswlib headers. Their non-template operations +live in `src/index/{ivf,hnsw,qg}.cpp`; +HNSW's per-ISA search template definitions live in `src/index/hnsw_search_impl.hpp`. +Keep lightweight numerical values in `core/types.hpp` and IVF initializer choices in +`index/ivf/types.hpp`. Preserve the explicit advanced `hnsw_quant.hpp` query-cache header. + +Core/FastScan support `RABITQ_USE_OPENMP=OFF`; ANN/Python require ON. The option defaults +to ON and is exported to consumers. Optional numerical OpenMP use must include +`utils/openmp.hpp` and guard both runtime calls and pragmas. Disabled builds use serial +helpers and their installed package must neither find nor link OpenMP. The standalone +core benchmark lives in `benchmarks/core/` and links only core. + +The v1 contract is documented in [Stable Core API](docs/docs/rabitq/core_api.md). +Keep `tests/fixtures/core_api_v1/consumer.cpp` unchanged when extending v1; add tests +instead of editing away compatibility failures. Changes to v1 names, signatures, +aggregate fields, enum values, buffer/code semantics or allocation guarantees require +an explicit API-major decision. The stable header must not expose Eigen, record maps, +ANN or ISA-specific types. Keep the existing algorithm and exact fallback authoritative; +workspace APIs must never silently fall back to heap allocation or spawn worker threads. + ## Required working boundaries - Inspect relevant code and tests first. For multi-step work, state brief, verifiable success @@ -112,13 +157,18 @@ Recommended: - Public/generic code calls centralized dispatch entry points. Keep ISA-specific translation units and their flags in `CMakeLists.txt` synchronized with feature predicates in - `src/simd/dispatch.cpp` and detection in `src/utils/cpu_features.cpp`, including HNSW source groups. + `src/simd/dispatch.cpp`, `src/fastscan/fastscan_dispatch.cpp`, `src/index/hnsw_dispatch.cpp` and detection in + `src/utils/cpu_features.cpp`, including HNSW source groups. ARM64 builds must exclude all x86 ISA source groups; ARM kernels use portable AArch64 NEON intrinsics. -- Use the shared resolver in `src/simd/dispatch.cpp` for cached selection, including HNSW. +- Use the private shared resolver in `src/simd/dispatch_resolver.hpp` for cached selection. + Core dispatch lives in `src/simd/dispatch.cpp`; batch scoring dispatch lives in + `src/fastscan/fastscan_dispatch.cpp`; HNSW dispatch lives in + `src/index/hnsw_dispatch.cpp` so core never references index search symbols. Keep calculations in backend source files; see the dispatch coverage table in [DEVELOPMENT.md](DEVELOPMENT.md#dispatch-conventions-and-coverage). -- Dispatch resolves function pointers during static initialization. Detection must use safe generic +- Dispatch caches function pointers in thread-safe function-local statics so calls before `main()` + are safe regardless of translation-unit initialization order. Detection must use safe generic code; never execute a high-ISA kernel to find out whether the CPU supports it. - Semantic kernel changes must cover every implementation and a backend-independent reference test, including dimensions around vector-width/packing boundaries and supported bit widths. diff --git a/CMakeLists.txt b/CMakeLists.txt index 8474c083..09247885 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,6 @@ cmake_minimum_required(VERSION 3.20) -project(RaBitQLib VERSION 0.6.0 LANGUAGES CXX) +project(RaBitQLib VERSION 0.7.0 LANGUAGES CXX) include(CMakePackageConfigHelpers) include(GNUInstallDirs) @@ -10,9 +10,13 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) +option(RABITQ_BUILD_FASTSCAN "Build batch scoring and FastScan SIMD kernels" ON) +option(RABITQ_BUILD_INDEX "Build ANN indexes and clustering" ON) option(RABITQ_BUILD_SAMPLES "Build C++ sample programs" ON) option(RABITQ_BUILD_TESTS "Build C++ tests" OFF) +option(RABITQ_BUILD_BENCHMARKS "Build component quantization and hot-path benchmarks" OFF) option(RABITQ_BUILD_PYTHON_BINDINGS "Build Python bindings" OFF) +option(RABITQ_USE_OPENMP "Use OpenMP for parallel helpers (required by ANN indexes)" ON) option(RABITQ_ENABLE_SANITIZERS "Enable AddressSanitizer and UndefinedBehaviorSanitizer" OFF) # Enable native tuning for local builds; release wheels explicitly set this OFF. option( @@ -21,12 +25,33 @@ option( ON ) -include(cmake/rabitqOpenMP.cmake) -find_package(OpenMP REQUIRED) +if(RABITQ_BUILD_INDEX AND NOT RABITQ_BUILD_FASTSCAN) + message(FATAL_ERROR "RABITQ_BUILD_INDEX requires RABITQ_BUILD_FASTSCAN=ON") +endif() + +if(RABITQ_BUILD_PYTHON_BINDINGS AND NOT RABITQ_BUILD_INDEX) + message(FATAL_ERROR "RABITQ_BUILD_PYTHON_BINDINGS requires RABITQ_BUILD_INDEX=ON") +endif() + +if(RABITQ_BUILD_INDEX AND NOT RABITQ_USE_OPENMP) + message(FATAL_ERROR "ANN indexes require RABITQ_USE_OPENMP=ON; disable RABITQ_BUILD_INDEX for an OpenMP-free numerical library") +endif() + +if(RABITQ_USE_OPENMP) + include(cmake/rabitqOpenMP.cmake) + find_package(OpenMP REQUIRED COMPONENTS CXX) +endif() add_library(rabitq_compile_options INTERFACE) target_compile_features(rabitq_compile_options INTERFACE cxx_std_17) -target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX) +target_compile_definitions(rabitq_compile_options INTERFACE + RABITQ_USE_OPENMP=$ + # Eigen-backed TLS/storage may be allocated by native code and freed by a + # generic consumer. Keep its layout and aligned-allocation policy identical. + EIGEN_MAX_ALIGN_BYTES=64 EIGEN_MAX_STATIC_ALIGN_BYTES=64) +if(RABITQ_USE_OPENMP) + target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX) +endif() if(MSVC) # MSVC spells the GNU restrict extension differently. Public headers use # __restrict__, so consumers need the same definition as the library. @@ -37,7 +62,7 @@ if(MSVC) endif() if(RABITQ_ENABLE_SANITIZERS) # Keep sanitizer instrumentation consistent across public inline code and - # the core (Eigen's allocation policy depends on __SANITIZE_ADDRESS__). + # compiled components, including their runtime linkage. target_compile_options(rabitq_compile_options INTERFACE "$<$:-fsanitize=address,undefined;-fno-omit-frame-pointer>" ) @@ -55,39 +80,31 @@ function(rabitq_configure_target target) ) endfunction() -set(RABITQ_COMMON_SOURCES - src/index/ivf_initializer.cpp - src/index/ivf_search.cpp - src/index/qg.cpp +set(RABITQ_CORE_SOURCES + src/core/api.cpp src/utils/cpu_features.cpp src/utils/thread_count.cpp src/simd/dispatch.cpp src/simd/space_generic.cpp - src/simd/fastscan_generic.cpp src/simd/quantization_generic.cpp - src/simd/estimator_generic.cpp src/simd/matrix_generic.cpp ) -set(RABITQ_AVX2_SOURCES +set(RABITQ_CORE_AVX2_SOURCES src/simd/quantization_avx2.cpp src/simd/pack_excode_avx2.cpp src/simd/space_excode_avx2.cpp src/simd/space_avx2.cpp - src/simd/fastscan_avx2.cpp - src/simd/estimator_avx2.cpp src/simd/matrix_avx2.cpp src/simd/warmup_avx2.cpp src/simd/rotator_avx2.cpp ) -set(RABITQ_AVX512_SOURCES +set(RABITQ_CORE_AVX512_SOURCES src/simd/quantization_avx512.cpp src/simd/pack_excode_avx512.cpp src/simd/space_excode_avx512.cpp src/simd/space_avx512.cpp - src/simd/fastscan_avx512.cpp - src/simd/estimator_avx512.cpp src/simd/matrix_avx512.cpp src/simd/rotator_avx512.cpp ) @@ -96,6 +113,16 @@ set(RABITQ_AVX512_POPCNT_SOURCES src/simd/warmup_avx512.cpp ) +set(RABITQ_FASTSCAN_AVX2_SOURCES + src/fastscan/fastscan_avx2.cpp + src/fastscan/estimator_avx2.cpp +) + +set(RABITQ_FASTSCAN_AVX512_SOURCES + src/fastscan/fastscan_avx512.cpp + src/fastscan/estimator_avx512.cpp +) + set(RABITQ_HNSW_AVX2_SOURCES src/index/hnsw_search_avx2.cpp ) @@ -116,9 +143,15 @@ endif() string(TOLOWER "${RABITQ_TARGET_PROCESSOR}" RABITQ_TARGET_PROCESSOR) if(RABITQ_TARGET_PROCESSOR MATCHES "^(x86_64|amd64)$") set(RABITQ_ARCH_SOURCES - ${RABITQ_AVX2_SOURCES} - ${RABITQ_AVX512_SOURCES} + ${RABITQ_CORE_AVX2_SOURCES} + ${RABITQ_CORE_AVX512_SOURCES} ${RABITQ_AVX512_POPCNT_SOURCES} + ) + set(RABITQ_FASTSCAN_ARCH_SOURCES + ${RABITQ_FASTSCAN_AVX2_SOURCES} + ${RABITQ_FASTSCAN_AVX512_SOURCES} + ) + set(RABITQ_INDEX_ARCH_SOURCES ${RABITQ_HNSW_AVX2_SOURCES} ${RABITQ_HNSW_AVX512_CORE_SOURCES} ${RABITQ_HNSW_AVX512_POPCNT_SOURCES} @@ -126,22 +159,22 @@ if(RABITQ_TARGET_PROCESSOR MATCHES "^(x86_64|amd64)$") elseif(RABITQ_TARGET_PROCESSOR MATCHES "^(arm64|aarch64)$") set(RABITQ_ARCH_SOURCES src/simd/space_neon.cpp - src/simd/fastscan_neon.cpp src/simd/space_packed_neon.cpp src/simd/rotator_neon.cpp src/simd/query_neon.cpp - src/index/hnsw_search_neon.cpp ) + set(RABITQ_FASTSCAN_ARCH_SOURCES src/fastscan/fastscan_neon.cpp) + set(RABITQ_INDEX_ARCH_SOURCES src/index/hnsw_search_neon.cpp) else() message(FATAL_ERROR "Supported targets are x86-64 and ARM64 (single architecture)") endif() add_library(rabitq_core STATIC - ${RABITQ_COMMON_SOURCES} + ${RABITQ_CORE_SOURCES} ${RABITQ_ARCH_SOURCES} src/simd/portable_generic.cpp - src/index/hnsw_search_generic.cpp ) +add_library(rabitqlib::core ALIAS rabitq_core) target_include_directories(rabitq_core PUBLIC $ @@ -156,22 +189,26 @@ rabitq_configure_target(rabitq_core) if(MSVC) set_source_files_properties( - ${RABITQ_AVX2_SOURCES} ${RABITQ_HNSW_AVX2_SOURCES} + ${RABITQ_CORE_AVX2_SOURCES} ${RABITQ_FASTSCAN_AVX2_SOURCES} + ${RABITQ_HNSW_AVX2_SOURCES} PROPERTIES COMPILE_OPTIONS "/arch:AVX2" ) # MSVC enables the complete F/CD/VL/BW/DQ group. Keep the optional # backend's runtime guard in cpu_features.cpp aligned with this flag. set_source_files_properties( - ${RABITQ_AVX512_SOURCES} ${RABITQ_AVX512_POPCNT_SOURCES} + ${RABITQ_CORE_AVX512_SOURCES} ${RABITQ_FASTSCAN_AVX512_SOURCES} + ${RABITQ_AVX512_POPCNT_SOURCES} ${RABITQ_HNSW_AVX512_CORE_SOURCES} ${RABITQ_HNSW_AVX512_POPCNT_SOURCES} PROPERTIES COMPILE_OPTIONS "/arch:AVX512" ) else() - set_source_files_properties(${RABITQ_AVX2_SOURCES} PROPERTIES - COMPILE_OPTIONS "-mavx2;-mfma" + set_source_files_properties( + ${RABITQ_CORE_AVX2_SOURCES} ${RABITQ_FASTSCAN_AVX2_SOURCES} + PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma" ) - set_source_files_properties(${RABITQ_AVX512_SOURCES} PROPERTIES - COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mfma" + set_source_files_properties( + ${RABITQ_CORE_AVX512_SOURCES} ${RABITQ_FASTSCAN_AVX512_SOURCES} + PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mfma" ) set_source_files_properties(${RABITQ_AVX512_POPCNT_SOURCES} PROPERTIES COMPILE_OPTIONS "-mavx512f;-mavx512bw;-mavx512dq;-mavx512vpopcntdq;-mfma" @@ -187,24 +224,83 @@ else() ) endif() -add_library(rabitq_headers INTERFACE) -add_library(rabitqlib::rabitqlib ALIAS rabitq_headers) -target_include_directories(rabitq_headers INTERFACE - $ - $ -) -target_link_libraries(rabitq_headers INTERFACE rabitq_core) -set_target_properties(rabitq_headers PROPERTIES EXPORT_NAME rabitqlib) +set(RABITQ_INSTALL_TARGETS rabitq_core rabitq_compile_options) +if(RABITQ_BUILD_FASTSCAN) + add_library(rabitq_fastscan STATIC + src/fastscan/fastscan_dispatch.cpp + src/fastscan/fastscan_generic.cpp + src/fastscan/estimator_generic.cpp + ${RABITQ_FASTSCAN_ARCH_SOURCES} + ) + add_library(rabitqlib::fastscan ALIAS rabitq_fastscan) + target_link_libraries(rabitq_fastscan PUBLIC rabitq_core) + set_target_properties(rabitq_fastscan PROPERTIES + EXPORT_NAME fastscan + POSITION_INDEPENDENT_CODE ON + ) + rabitq_configure_target(rabitq_fastscan) + list(APPEND RABITQ_INSTALL_TARGETS rabitq_fastscan) +endif() +if(RABITQ_BUILD_INDEX) + add_library(rabitq_index STATIC + src/index/hnsw.cpp + src/index/ivf.cpp + src/index/ivf_initializer.cpp + src/index/ivf_search.cpp + src/index/qg.cpp + src/index/qg_builder.cpp + src/index/pipnn.cpp + src/index/hnsw_dispatch.cpp + src/index/hnsw_search_generic.cpp + ${RABITQ_INDEX_ARCH_SOURCES} + ) + add_library(rabitqlib::index ALIAS rabitq_index) + target_link_libraries(rabitq_index PUBLIC rabitq_fastscan) + set_target_properties(rabitq_index PROPERTIES + EXPORT_NAME index + POSITION_INDEPENDENT_CODE ON + ) + rabitq_configure_target(rabitq_index) + + list(APPEND RABITQ_INSTALL_TARGETS rabitq_index) +endif() set_target_properties(rabitq_compile_options PROPERTIES EXPORT_NAME compile_options) -install(TARGETS rabitq_headers rabitq_core rabitq_compile_options +install(TARGETS ${RABITQ_INSTALL_TARGETS} EXPORT rabitqlibTargets ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} ) -install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +# Install only the enabled components and their header dependencies. Detail headers +# remain necessary for template implementations, but are not supported public APIs. +set(rabitq_include_dir "${CMAKE_INSTALL_INCLUDEDIR}/rabitqlib") +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/memory.hpp + include/rabitqlib/utils/openmp.hpp + include/rabitqlib/utils/packed.hpp + include/rabitqlib/utils/path.hpp + include/rabitqlib/utils/space.hpp + include/rabitqlib/utils/tools.hpp + include/rabitqlib/utils/warmup_space.hpp + DESTINATION "${rabitq_include_dir}/utils") +install(DIRECTORY include/rabitqlib/third/Eigen + DESTINATION "${rabitq_include_dir}/third") +if(RABITQ_BUILD_FASTSCAN) + install(DIRECTORY include/rabitqlib/fastscan DESTINATION "${rabitq_include_dir}") +endif() +if(RABITQ_BUILD_INDEX) + install(DIRECTORY include/rabitqlib/index include/rabitqlib/clustering + DESTINATION "${rabitq_include_dir}") + install(FILES include/rabitqlib/utils/array.hpp include/rabitqlib/utils/io.hpp + include/rabitqlib/utils/stopw.hpp DESTINATION "${rabitq_include_dir}/utils") + install(DIRECTORY include/rabitqlib/third/hnswlib + DESTINATION "${rabitq_include_dir}/third") +endif() set(rabitqlib_cmake_dir "${CMAKE_INSTALL_LIBDIR}/cmake/rabitqlib") configure_package_config_file( @@ -229,15 +325,23 @@ install(FILES DESTINATION "${rabitqlib_cmake_dir}" ) +if(RABITQ_BUILD_TESTS OR RABITQ_BUILD_BENCHMARKS) + enable_testing() +endif() + if(RABITQ_BUILD_SAMPLES) add_subdirectory(sample/cpp) endif() +if(RABITQ_BUILD_BENCHMARKS) + add_subdirectory(benchmarks/core) + add_subdirectory(benchmarks/hotpaths) +endif() + if(RABITQ_BUILD_PYTHON_BINDINGS) add_subdirectory(python_bindings) endif() if(RABITQ_BUILD_TESTS) - enable_testing() add_subdirectory(tests) endif() diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 53120e41..23b21fb7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -40,7 +40,10 @@ Run any examples you add or change, including examples inside Markdown. Source builds support Linux and Windows x86-64 with AVX2/FMA (AVX-512 optional), and Linux and macOS ARM64 with NEON. They need CMake 3.20 or newer, -a C++17 compiler, and OpenMP. Windows x86-64 uses Visual Studio 2026. See the +a C++17 compiler, and OpenMP for ANN/Python. Core/FastScan can disable OpenMP +with `RABITQ_USE_OPENMP=OFF`; pure core also has scalar kernels for CPUs without +the FastScan SIMD baseline. See [component builds](docs/docs/cpp_integration.md#core-only-c-integration). +Windows x86-64 uses Visual Studio 2026. See the [platform-specific build instructions](tests/README.md#quick-start) for [Linux ARM64](tests/README.md#linux-arm64) and [Apple Silicon](tests/README.md#macos-arm64). @@ -75,7 +78,7 @@ Install the checkout and test dependencies, then run one test to check your setu ```bash python -m pip install ".[test]" -python -m pytest tests/python/test_ivf.py::test_search_output_shape -q +python -m pytest tests/python/test_ivf.py::test_search_output_contract -q ``` On Apple Silicon, pass the OpenMP prefix and disable native tuning when building @@ -85,7 +88,7 @@ the extension in the same environment: python -m pip install ".[test]" \ -Ccmake.define.OpenMP_ROOT="$(brew --prefix libomp)" \ -Ccmake.define.RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF -python -m pytest tests/python/test_ivf.py::test_search_output_shape -q +python -m pytest tests/python/test_ivf.py::test_search_output_contract -q ``` This builds the C++ extension, so the source-build prerequisites above apply. @@ -175,8 +178,7 @@ another. Each task can be a separate PR. | Task | Where to start | Done when | | --- | --- | --- | -| Explain Python search results | [Quick start](docs/docs/quick_start.md), [IVF binding](python_bindings/ivf_bindings.cpp), and [IVF tests](tests/python/test_ivf.py) | A short explanation covers result shapes, how IDs map to input rows, and what L2 distances represent, checked against the implementation. The example runs and the strict docs build passes. | -| Add a self-contained IVF save/load example | [Python examples](sample/python/), [README quick start](README.md#python-quick-start), and [IVF tests](tests/python/test_ivf.py) | A deterministic script builds from synthetic data, saves to a temporary directory, reloads, and checks that search IDs and distances match. It needs no dataset download, cleans up its temporary files, passes Python checks, and is linked from the README. | +| Automate the documentation quick-start check | [Quick start](docs/docs/quick_start.md), [documentation workflow](.github/workflows/docs.yml), and [Python tests](tests/python/) | A check executes the documented Python example blocks in order against the current package, verifies the save/load assertions, and cleans up temporary files without keeping a second copy of the example. | ## Feedback and further reading diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index e7cd286d..606276b2 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -124,36 +124,47 @@ for the LLVM 22 example above, that is `libomp-22-dev`. An older installed ```bash cmake -S . -B build-includes -G Ninja \ -DCMAKE_CXX_COMPILER="${CXX:-clang++}" \ - -DRABITQ_BUILD_SAMPLES=OFF \ - -DRABITQ_BUILD_TESTS=OFF \ - -DRABITQ_BUILD_PYTHON_BINDINGS=OFF \ + -DPython_EXECUTABLE="$(python -c 'import sys; print(sys.executable)')" \ + -Dpybind11_DIR="$(python -m pybind11 --cmakedir)" \ + -DRABITQ_BUILD_SAMPLES=ON \ + -DRABITQ_BUILD_TESTS=ON \ + -DRABITQ_BUILD_BENCHMARKS=ON \ + -DRABITQ_BUILD_PYTHON_BINDINGS=ON \ -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ -DCMAKE_BUILD_TYPE=Release ./scripts/check-includes.sh build-includes ``` -The script checks library sources using their compilation database and checks -headers as main files, since this clang-tidy check does not report findings in -included headers. Private headers are checked with AVX2 and AVX-512 flags. New -untracked files are included, and tracked files deleted from the working tree are skipped. -Vendored files are excluded. The script also ignores suggestions to include -Eigen and hnswlib implementation headers behind their existing public headers; -these vendor snapshots lack the export annotations needed by include-cleaner. -`INCLUDE_JOBS` controls parallelism (default: 2). -The script fails for findings or analyzer errors and does not modify files. - -Review suggestions before applying them, especially for templates and public -forwarding headers. For a source file, automatic fixes can be applied with: - -```bash -"${CLANG_TIDY:-clang-tidy}" -p build-includes \ - --config='{CheckOptions: {misc-include-cleaner.IgnoreHeaders: "rabitqlib/third/Eigen/src/.*;rabitqlib/third/hnswlib/(hnswalg|space_l2)[.]h"}}' \ - --checks='-*,misc-include-cleaner' --fix src/simd/dispatch.cpp -``` - -Review the diff, run formatting, and rebuild and test affected code after fixes. -Findings can vary between LLVM releases. This include check is independent of -the existing general clang-tidy job. +The script checks every configured first-party source, including bindings, samples, +tests and benchmarks. Disable optional targets above for a smaller check; skipped +files are listed explicitly. Installed-package consumers and the historical fixture +generator have separate build configurations and must be checked in those contexts. +The historical generator requires the pinned release headers described in +[its README](tests/fixtures/persistence/README.md). + +Headers are checked as main files with `-x c++-header`, preserving `#pragma once` +for mutually including headers. Include Cleaner does not report findings in included +headers, so a source-only pass is insufficient. Public headers inherit configured +component flags; numerical headers also run with OpenMP disabled. Private headers +inherit the flags of their configured source consumers through local includes, +covering the enabled ISA variants without maintaining a second backend list. +Run a separate AArch64 compilation database to cover NEON; an x86 check lists those +files as skipped. New untracked files are included, and tracked files deleted from +the working tree are skipped. `INCLUDE_JOBS` controls parallelism (default: 2), and +`PYTHON` can select the existing project interpreter for the driver. + +Vendored files are excluded. The check ignores suggestions to include Eigen/hnswlib +implementation headers behind their public entry points, and pybind11's internal +`detail/common.h` behind `PYBIND11_MODULE`. These vendor headers lack the export +annotations needed by Include Cleaner. First-party export annotations describe +intentional API re-exports; `keep` annotations retain configuration-only includes +and explicit header-compilation tests. Neither is a general unused-include waiver. + +The script fails for include findings, duplicate includes or analyzer errors and +does not modify files. Review suggestions across supported configurations before +applying them. Preserve public re-exports, frozen source consumers and conditional +includes. Format changes and rebuild/test affected targets after fixes. Findings can +vary between LLVM releases; this check is independent of the general clang-tidy job. ### Focused static analysis @@ -248,12 +259,33 @@ This matches CI's script set. Also pass any changed shell scripts outside ## Implementation recipes +### Validation, ownership, and numerical preparation + +Keep validation at the public operation boundary and reuse only invariants established +for that operation. Buffer contents and per-record factors may change between calls: +validate them again even when reusing a query or workspace. + +| Boundary | Ownership and checks | Internal numerical work | +| --- | --- | --- | +| Python conversion | `ensure_*` returns a normalized NumPy owner retained across GIL release; `checked_cluster_ids` validates every ID while making the owned PID copy required by builders | Native calls borrow those arrays only within their established lifetime | +| Stable core v1 | Callers own codes, factors, queries and workspace; each operation checks its buffers and ranges | Prepared queries retain validated dimensions, bit width, metric and kernel selection; conversion checks both layouts and each logical code before writing output | +| FastScan LUT preparation | `Lut` owns prepared bytes; constructor/reset check shape once, then check borrowed scratch and resize owned storage | Private `build_lut` performs preparation with the established shape and scratch; reset reuses capacity | +| Query correction | `BatchQuery` owns its LUT and scalar corrections; `SplitBatchQuery` additionally borrows the rotated query for reranking | Complete LUT reset before updating the query's metric and correction metadata | + +Keep the v1 `compute` checks for each code view, finite factors, unpacked code values, +and finite output. Keep scratch APIs allocation-bounded as documented; pointer reuse +alone is not permission to skip validation. The FastScan scratch pointers are borrowed +only during preparation and must never become retained query storage. Test rejected +resets, scratch reuse, and changed records alongside normal numerical parity. + ### Add or change a SIMD kernel -1. Define the backend-neutral entry point under `include/rabitqlib/simd/`. -2. Update all applicable declarations and AVX2/AVX-512 implementations in `src/simd/` or - `src/index/`. -3. Register selection and unsupported-CPU behavior in `src/simd/dispatch.cpp` when dispatched. +1. Define the backend-neutral entry point under the owning module's `detail/` headers. +2. Update all applicable declarations and AVX2/AVX-512 implementations in `src/simd/` (core), + `src/fastscan/` (FastScan), or `src/index/` (indexes). +3. Register selection and unsupported-CPU behavior in `src/simd/dispatch.cpp` when dispatched + (batch scoring uses `src/fastscan/fastscan_dispatch.cpp`, HNSW search uses + `src/index/hnsw_dispatch.cpp`). 4. If adding a translation unit, put it in the correct ISA source group and flags in `CMakeLists.txt`. 5. Add differential tests against a scalar or simple reference, including boundary dimensions, @@ -264,16 +296,21 @@ must happen in generic code first. #### Dispatch conventions and coverage -- Keep backend-neutral declarations in `include/rabitqlib/simd/*_dispatch.hpp` and - implementations in `src/simd/*_{generic,neon,avx2,avx512}.cpp`. HNSW keeps its existing +- Keep core backend declarations in `include/rabitqlib/core/detail/simd/`, FastScan + declarations in `fastscan/detail/`, and HNSW declarations in `index/detail/`. Keep + core implementations in `src/simd/*_{generic,neon,avx2,avx512}.cpp` and FastScan + implementations in `src/fastscan/`. HNSW keeps its existing `src/index/` implementations and compatibility namespaces. -- Select a cached function pointer through `resolve_kernel` in `src/simd/dispatch.cpp`. +- Select a cached function pointer through the private shared resolver in + `src/simd/dispatch_resolver.hpp`. Core, FastScan and HNSW have separate dispatch translation units. The x86 order is AVX-512, AVX2/FMA, then the portable generic implementation. ARM64 selects NEON for raw float distances, FastScan, packed-code dot products (1–8 extra bits), masked sums, binary warmup, HNSW search, and FHT/Kac rotation including sign flipping, query-bit transposition, uint8/uint16 scalar quantization, and FastScan LUT construction. Extra-code packing retains a portable scalar implementation. - Public wrappers do not repeat feature checks. + Cache selection in thread-safe function-local statics, including tables; public calls must + work before `main()` regardless of translation-unit initialization order. Public wrappers + do not repeat feature checks. Standard FastScan accumulation requires AVX2/FMA, AVX-512, or NEON and throws `std::runtime_error` when no supported SIMD backend is available. Its scalar kernel is retained as a correctness reference only, never as a runtime fallback. @@ -287,6 +324,25 @@ must happen in generic code first. Missing group features select the AVX2 fallback; they do not raise the library's minimum CPU requirements. Native tuning defaults to on for local builds and must be explicitly disabled for distributed builds; release wheels already set it to off. +- FastScan selects one immutable private `Backend` in `src/fastscan/fastscan_backend.hpp` + and `fastscan_dispatch.cpp`. LUT production, HACC conversion, accumulation and batch + correction must be chosen as a compatible group. Standard LUTs share consecutive + 16-entry tables; HACC has paired 32-byte planes (AVX2/generic/NEON) or grouped 64-byte + planes (AVX-512). These are transient query layouts, not persisted-code formats. + Do not select producer and consumer independently when adding a backend. +- Dispatch once at a generic operation boundary. Warm wrappers cache their selected + function pointer; keep feature selection and backend-table access in cold initialization. + A backend may directly compose an accumulator only when its LUT layout is established. + AVX-512 batch correction does so. Explicit AVX2/generic correction entry points also + accept host-prepared queries and intentionally retain host accumulation dispatch; + their function name alone does not specify the LUT format. Estimator templates require + explicit accumulator arguments so new call sites must make this choice deliberately. + Never bypass public dimension checks, unsupported-CPU errors or initialization safety. +- Keep performance regression workloads in [component hot-path benchmarks](benchmarks/hotpaths/README.md). + Validate numerical results first, then inspect generated code and alternate comparable + baseline/candidate runs. Keep failed experiments and quality results. Fusion, inlining, + scratch reuse and dispatch changes need measurements across relevant vector sizes; + a short-vector improvement does not justify applying it to wider vectors. - Put calculations and scratch-storage helpers outside the dispatcher. Shared implementation headers use internal linkage so independently compiled backends retain their own bodies. The AVX FHT intrinsics in `src/simd/fht_kernels.hpp` use the private kernel namespace @@ -327,15 +383,68 @@ macOS and Linux ARM64 have native C++ and installed-wheel CI. The reference suit also exercise portable scalar kernels on x86, including byte-for-byte packing parity with the selected x86 backend. +### Module build boundaries + +Core must not depend on FastScan scoring or index headers, objects or symbols. FastScan +must not depend on indexes. The canonical headers are under `core/`, `fastscan/` and +`index/{ivf,hnsw,symqg}/`. Old `quantization/`, `utils/rotator.hpp` and generic +`index/{query,lut,estimator}.hpp` forwarders have been removed. Migrate consumers to +their owning module; see the [migration guide](docs/docs/modular_migration.md). +Core-only installs exclude FastScan/index/clustering headers and vendored hnswlib. +Installed detail headers support templates and are not public APIs. Internal-only +non-template dispatch declarations stay in `src/`. +`defines.hpp` has been removed. Keep metrics in `core/metric.hpp`, scalar quantizer, +rotation and quantization configuration in `core/types.hpp`, Eigen aliases in `core/detail/eigen.hpp`, and ANN values +in `index/types.hpp`. Do not introduce an umbrella that couples them again. +IVF initializer choices live in `index/ivf/types.hpp`. IVF/HNSW/SymphonyQG public index headers +contain declarations and small accessors; their operations compile in +`src/index/{ivf,hnsw,qg}.cpp`. Backend translation units include HNSW's private +`src/index/hnsw_search_impl.hpp` to retain per-ISA template instantiation. +The installed consumer compiles and links each index header independently while poisoning +encoding, scoring, rotation implementation and vendor headers to enforce this boundary. +All CMake components export 64-byte Eigen alignment limits so native and generic +translation units share compatible storage and TLS allocation rules. Retain these +definitions in installed exports; SIMD backends keep their private Eigen namespaces. +Both two-factor and three-factor batch layouts/encoders belong to +`fastscan/{data_layout,encoder}.hpp`. Index-only nearest-neighbor helpers +belong in `index/detail/`, not generic numerical `utils/space.hpp`. +FastScan's 32-vector packing, inverse packing, batch size and permutation are owned by +`fastscan/pack_codes.hpp`; the batch encoder adapter lives in `fastscan/detail/encoder.hpp`. +Core retains per-vector packing and encoding loops independent of that batch format. +`utils/packed.hpp` supplies generic unaligned byte access for both numerical layouts and +index records. The `CoreModuleBoundary` test rejects format-specific declarations copied +into core even if they introduce no upward header include. + +Keep core's and FastScan's architecture source groups explicit in CMake. The NEON LUT +implementation belongs to `fastscan_neon.cpp`, while scalar-query preprocessing belongs +to `query_neon.cpp`. All dispatchers share the private resolver and CPU predicates. +Eigen remains bundled. `RABITQ_USE_OPENMP=OFF` removes OpenMP from core/FastScan builds; +ANN still requires it. Guard optional numerical OpenMP calls/pragmas with the exported +`RABITQ_USE_OPENMP` definition via `utils/openmp.hpp`; do not provide global fake `omp_*` +functions. Test both configurations, including installed consumers. Python is a MODULE extension with an optional +build-tree alias; it is not a C++ link library or a C++ package component. +Validate these boundaries with the +[external-consumer checks](tests/README.md#core-only-cmake-package-tests). + +### Stable core API changes + +`core/api.hpp` is the source-stable float32 v1 boundary; the implementation is +`src/core/api.cpp`. See [its contract](docs/docs/rabitq/core_api.md). Keep the frozen +consumer, API unit tests and both `core_api_no_alloc_*` processes in the validation +matrix. The scalar allocation process forces cached dispatch onto the exact event +fallback. Do not replace that fallback with an approximate winner to avoid allocations. +Legacy and caller-backed paths share factor/search implementations. Validate numerical +parity against a saved pre-change build when changing those shared implementations. + ### Change quantization or packing Check all of these together: -- `quantization/rabitq.hpp` and `quantization/rabitq_impl.hpp`; -- `quantization/data_layout.hpp` and `quantization/pack_excode.hpp`; +- `core/encoder.hpp` and `core/detail/rabitq_impl.hpp`; +- `core/data_layout.hpp` and `core/pack_excode.hpp`; - SIMD pack, inner-product, and FastScan kernels; -- query-side factors in `index/query.hpp`; -- estimators in `index/estimator.hpp`; +- query-side factors in `core/query.hpp` and `fastscan/query.hpp`; +- estimators in `core/estimator.hpp` and `fastscan/estimator.hpp`; - IVF, HNSW, and SymphonyQG consumers; - C++ tests for factor finiteness, reconstruction, sign convention, pack/unpack, and estimates. diff --git a/README.md b/README.md index 36038986..cca6d7b5 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,9 @@ ## News +- **v0.7.0:** Modular C++ components, stable core API v1, optional + OpenMP for numerical components, and unified index counts. This documentation + targets 0.7.0; see the [upgrade notes](docs/docs/releases/0.7.0.md). - **v0.6.0:** Concurrent Python index searches, configurable IVF centroid routing, 32-coordinate padding, and Eigen 5.0.1. See the [upgrade notes](docs/docs/releases/0.6.0.md). @@ -40,47 +43,22 @@ ## Install ```bash -python -m pip install --upgrade "rabitqlib>=0.6.0" +python -m pip install --upgrade "rabitqlib>=0.7.0" ``` Wheels cover the platforms above. x86-64 requires AVX2/FMA and optionally uses -AVX-512; ARM64 uses NEON. See [platform requirements and source installation](docs/docs/quick_start.md#requirements). +AVX-512; ARM64 uses NEON. See [platform requirements and source installation](docs/docs/installation.md#requirements). ## Python quick start -Build and search a small IVF index using synthetic data: - -```python -import numpy as np -from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans - -rng = np.random.default_rng(42) -data = rng.standard_normal((500, 64)).astype(np.float32) -queries = rng.standard_normal((5, 64)).astype(np.float32) - -clustering = RaBitQKMeans( - 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact -) -clustering.train(data) - -index = IvfIndex( - dim=64, - max_elements=len(data), - num_clusters=5, - nbits=4, - metric="l2", -) -index.build(data, clustering.centroids, clustering.assignments) - -ids, distances = index.search(queries, k=10, nprobe=5) -print(ids.shape, distances.shape) # (5, 10) (5, 10) -print(ids[0]) -``` +Follow the [complete Python quick start](docs/docs/quick_start.md) to build and +search an IVF index, interpret its results, and save/load it using synthetic data. +Native [clustering](docs/docs/clustering.md) needs no external k-means package. -See the [quick start](docs/docs/quick_start.md) for index updates and -[Python examples](sample/python/README.md) for all three indexes. Native -[clustering](docs/docs/clustering.md) needs no external k-means package. -See [threading and save/load paths](docs/docs/quick_start.md#threading-and-file-paths) +Then see the [IVF](docs/docs/index/ivf.md), [HNSW](docs/docs/index/hnsw.md), and +[SymphonyQG](docs/docs/index/qg.md) guides, or the +[dataset-based Python examples](sample/python/README.md). +See [threading and save/load paths](docs/docs/index/usage.md#threading-and-file-paths) for runtime conventions. ## Choose the right building block @@ -96,6 +74,22 @@ IVF and HNSW support adding and removing vectors; HNSW also supports explicit capacity resizing. See their guides for update costs and file compatibility. The library supports L2 and inner product; normalize vectors for cosine search. +C++ systems with their own indexes can link `rabitqlib::core` for encoding, +single-vector estimation and rotation, or `rabitqlib::fastscan` for batch scoring. +Configure `RABITQ_BUILD_INDEX=OFF` to omit ANN implementations; also disable +`RABITQ_BUILD_FASTSCAN` for pure core. Set `RABITQ_USE_OPENMP=OFF` for core/FastScan +without an OpenMP dependency; ANN builds retain OpenMP. `rabitqlib::index` +provides indexes and clustering. The 0.7.0 refactor removes the former +`rabitqlib::rabitqlib` target and forwarding headers; see the +[migration guide](docs/docs/modular_migration.md). Optional Python bindings retain +`import rabitqlib`. +For a source-stable, storage-neutral interface with caller-backed workspace, use +[`rabitqlib/core/api.hpp`](docs/docs/rabitq/core_api.md). Its v1 entry points make no heap +allocations on valid calls and keep code and factors separate. +See [modular integration](docs/docs/cpp_integration.md#core-only-c-integration), the +[core example](sample/cpp/core/) and [FastScan example](sample/cpp/fastscan/). +The [core benchmark](benchmarks/core/README.md) runs independently of ANN and Python. + ## Why RaBitQ? - **Compact codes:** Choose 1-bit or multi-bit quantization for your memory and accuracy needs. @@ -180,11 +174,10 @@ another integration example. ## C++ quick start Requires CMake 3.20+, a C++17 compiler with OpenMP, and an x86-64 CPU with -AVX2/FMA or an ARM64 CPU. Build the library and examples: +AVX2/FMA or an ARM64 CPU. First [get the 0.7.0 source](docs/docs/installation.md#get-the-source), +then build the library and examples from the repository root: ```bash -git clone https://github.com/VectorDB-NTU/RaBitQ-Library.git -cd RaBitQ-Library cmake -S . -B build -DCMAKE_BUILD_TYPE=Release cmake --build build --parallel ``` @@ -194,13 +187,15 @@ other CPUs. See [platform-specific build instructions](tests/README.md#prerequis ### Use RaBitQ-Library in another C++ project -Add the repository as a submodule at `third_party/rabitqlib`, pin a release or -commit, and link its CMake target: +The component targets below are introduced in 0.7.0. Use a matching checkout, +as described in the [migration guide](docs/docs/modular_migration.md). Add that +checkout as a submodule at `third_party/rabitqlib`, pin its revision, and link the +component you need: ```cmake set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) add_subdirectory(third_party/rabitqlib) -target_link_libraries(my_program PRIVATE rabitqlib::rabitqlib) +target_link_libraries(my_program PRIVATE rabitqlib::index) ``` The target propagates C++17, OpenMP, and required public-header definitions. @@ -208,14 +203,16 @@ Optimization and native-tuning options apply only to RaBitQ's own targets; configure your application's optimization flags separately. Explicit sanitizer builds propagate instrumentation and runtime linkage for header compatibility. -The C++ API and ABI are evolving; update the pinned revision deliberately. -For installed packages, use `find_package(rabitqlib CONFIG REQUIRED)` and the -same target. Both approaches require OpenMP. +The [core v1 API](docs/docs/rabitq/core_api.md) has a source-compatibility contract; +other C++ interfaces may evolve. No cross-version C++ binary ABI is promised, so +rebuild consumers when updating the library. For an installed ANN package, use +`find_package(rabitqlib CONFIG REQUIRED COMPONENTS index)` and the same target. +ANN requires OpenMP; core/FastScan can be built with `RABITQ_USE_OPENMP=OFF`. -See [submodule setup and version pinning](docs/docs/quick_start.md#use-in-another-c-project), -[C++ installation](docs/docs/quick_start.md#install-the-c-library), and the +See [submodule setup and version pinning](docs/docs/cpp_integration.md#use-in-another-c-project), +[C++ installation](docs/docs/cpp_integration.md#install-the-c-library), and the [build and test guide](tests/README.md) for complete workflows. -Examples: [C++ indexes and quantization](docs/docs/quick_start.md#c-examples). +Examples: [C++ indexes and quantization](docs/docs/cpp_integration.md#c-examples). Benchmarks: [FAISS clustering comparison](docs/docs/clustering.md#compare-with-faiss) and [GIST workflow](example.sh). diff --git a/benchmarks/README.md b/benchmarks/README.md index 85d2cf07..e758ae74 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,5 +1,12 @@ # Benchmarks +For embedding the quantizer in another system, use the independent +[core C++ benchmark](core/README.md). It needs only core and reports per-repeat timings, +distance-estimate error and quantized full-scan recall. The workflows below benchmark +complete ANN/clustering systems. For dispatch, encoding, FastScan, sidecar loading and +Python ID-conversion regression measurements, use the +[component hot-path benchmarks](hotpaths/README.md). + See the [benchmark overview](../docs/docs/benchmarks/bench.md) for datasets, installation, measurement definitions, file locations, and reproduction options. Results: [IVF](../docs/docs/benchmarks/ivf.md), @@ -7,7 +14,10 @@ Results: [IVF](../docs/docs/benchmarks/ivf.md), [k-means](../docs/docs/benchmarks/kmeans.md). From the repository root, after installing the documented system prerequisites -and activating a Python 3.13 environment: +and activating a Python 3.13 environment, these commands measure the active checkout. +To repeat the published 0.5.2 software configuration, first use the +[pinned historical checkout](../docs/docs/benchmarks/bench.md#package-installation). +Keep measurements of different revisions in separate output directories. ```bash python benchmarks/install_baselines.py --jobs 12 diff --git a/benchmarks/build_indexes.py b/benchmarks/build_indexes.py index 24be0c92..1c492f6f 100644 --- a/benchmarks/build_indexes.py +++ b/benchmarks/build_indexes.py @@ -4,6 +4,7 @@ import json import subprocess import sys +from dataclasses import dataclass from pathlib import Path ROOT = Path(__file__).resolve().parents[1] @@ -25,48 +26,90 @@ } DATASET_SHAPES = {"gist1m": (1000000, 960), "dbpedia": (999000, 1536)} -CASES = ( - ("rabitq-ivf-raw", "rabitq-ivf", ("--bits", "32")), - ("rabitq-ivf-1p4", "rabitq-ivf", ("--bits", "5")), - ("rabitq-ivf-1p8", "rabitq-ivf", ("--bits", "9")), - ("faiss-rabitqfs-raw", "faiss-ivf-rabitq-fs", ("--bits", "1")), - ("faiss-rabitqfs-1p4", "faiss-ivf-rabitq-fs", ("--bits", "5")), - ("faiss-rabitqfs-1p8", "faiss-ivf-rabitq-fs", ("--bits", "9")), - ("faiss-pqfs", "faiss-ivf-pqfs", ()), - ("faiss-opq-pqfs", "faiss-ivf-pqfs", ("--opq",)), - ("rabitq-symqg-raw", "rabitq-symqg", ("--quantization-bits", "0")), - ("rabitq-symqg-4", "rabitq-symqg", ("--quantization-bits", "4")), - ("rabitq-symqg-8", "rabitq-symqg", ("--quantization-bits", "8")), - ("rabitq-hnsw-1p4", "rabitq-hnsw", ("--bits", "5")), - ("rabitq-hnsw-1p8", "rabitq-hnsw", ("--bits", "9")), - ("faiss-hnsw-flat", "faiss-hnsw-flat", ()), - ("faiss-hnsw-sq4", "faiss-hnsw-sq4", ()), - ("faiss-hnsw-sq8", "faiss-hnsw-sq8", ()), -) -LABELS = ( - "IVF RaBitQ raw", - "IVF RaBitQ 1+4", - "IVF RaBitQ 1+8", - "Faiss IVF RaBitQfs raw + RFlat", - "Faiss IVF RaBitQfs 1+4", - "Faiss IVF RaBitQfs 1+8", - "Faiss IVF PQfs + RFlat", - "Faiss OPQ-IVF PQfs + RFlat", - "SymphonyQG raw", - "SymphonyQG 4-bit", - "SymphonyQG 8-bit", - "HNSW RaBitQ 1+4", - "HNSW RaBitQ 1+8", - "Faiss HNSW Flat", - "Faiss HNSW SQ4", - "Faiss HNSW SQ8", +@dataclass(frozen=True) +class Case: + name: str + label: str + method: str + family: str + parameters: dict + + def cli_args(self): + args = [] + for key, value in self.parameters.items(): + args.append("--" + key.replace("_", "-")) + if value is not True: + args.append(str(value)) + return args + + +CASES = ( + Case("rabitq-ivf-raw", "IVF RaBitQ raw", "rabitq-ivf", "ivf", {"bits": 32}), + Case("rabitq-ivf-1p4", "IVF RaBitQ 1+4", "rabitq-ivf", "ivf", {"bits": 5}), + Case("rabitq-ivf-1p8", "IVF RaBitQ 1+8", "rabitq-ivf", "ivf", {"bits": 9}), + Case( + "faiss-rabitqfs-raw", + "Faiss IVF RaBitQfs raw + RFlat", + "faiss-ivf-rabitq-fs", + "ivf", + {"bits": 1}, + ), + Case( + "faiss-rabitqfs-1p4", + "Faiss IVF RaBitQfs 1+4", + "faiss-ivf-rabitq-fs", + "ivf", + {"bits": 5}, + ), + Case( + "faiss-rabitqfs-1p8", + "Faiss IVF RaBitQfs 1+8", + "faiss-ivf-rabitq-fs", + "ivf", + {"bits": 9}, + ), + Case("faiss-pqfs", "Faiss IVF PQfs + RFlat", "faiss-ivf-pqfs", "ivf", {}), + Case( + "faiss-opq-pqfs", + "Faiss OPQ-IVF PQfs + RFlat", + "faiss-ivf-pqfs", + "ivf", + {"opq": True}, + ), + Case( + "rabitq-symqg-raw", + "SymphonyQG raw", + "rabitq-symqg", + "graphs", + {"quantization_bits": 0}, + ), + Case( + "rabitq-symqg-4", + "SymphonyQG 4-bit", + "rabitq-symqg", + "graphs", + {"quantization_bits": 4}, + ), + Case( + "rabitq-symqg-8", + "SymphonyQG 8-bit", + "rabitq-symqg", + "graphs", + {"quantization_bits": 8}, + ), + Case("rabitq-hnsw-1p4", "HNSW RaBitQ 1+4", "rabitq-hnsw", "graphs", {"bits": 5}), + Case("rabitq-hnsw-1p8", "HNSW RaBitQ 1+8", "rabitq-hnsw", "graphs", {"bits": 9}), + Case("faiss-hnsw-flat", "Faiss HNSW Flat", "faiss-hnsw-flat", "graphs", {}), + Case("faiss-hnsw-sq4", "Faiss HNSW SQ4", "faiss-hnsw-sq4", "graphs", {}), + Case("faiss-hnsw-sq8", "Faiss HNSW SQ8", "faiss-hnsw-sq8", "graphs", {}), ) def validate_record(record, dataset, name): """Reject results that do not match the published ANN configuration.""" - method, extra = next((m, e) for n, m, e in CASES if n == name) + case = next(case for case in CASES if case.name == name) + method = case.method base, queries, truth, metric = DATASETS[dataset] n, dim = DATASET_SHAPES[dataset] expected = { @@ -80,17 +123,15 @@ def validate_record(record, dataset, name): "socket": 0, "k": 10, } - if "ivf" in method: - expected.update(nlist=4096, opq="--opq" in extra) + if case.family == "ivf": + expected.update(nlist=4096, opq=False) if "hnsw" in method: expected.update(m=16, ef_construction=200) if method == "rabitq-hnsw": expected["hnsw_clusters"] = 16 if method == "rabitq-symqg": expected.update(degree=32, ef_construction=200) - for option in ("--bits", "--quantization-bits"): - if option in extra: - expected[option[2:].replace("-", "_")] = int(extra[extra.index(option) + 1]) + expected.update(case.parameters) cfg = record["config"] if ( any(cfg.get(key) != value for key, value in expected.items()) @@ -115,14 +156,15 @@ def main(): parser.add_argument( "--output-dir", type=Path, default=Path("results/benchmarks/ann") ) - parser.add_argument("--methods", nargs="+", choices=[case[0] for case in CASES]) + parser.add_argument("--methods", nargs="+", choices=[case.name for case in CASES]) parser.add_argument("--threads", type=int, default=48) args = parser.parse_args() failed = [] for dataset in args.datasets: base, queries, truth, metric = DATASETS[dataset] - for name, method, extra in CASES: + for case in CASES: + name, method = case.name, case.method if args.methods is not None and name not in args.methods: continue output = args.output_dir / dataset / f"{name}.json" @@ -177,7 +219,7 @@ def main(): str(index_dir), "--output", str(output), - *extra, + *case.cli_args(), ] print(f"START {dataset} {name}", flush=True) with log.open("w") as stream: diff --git a/benchmarks/core/CMakeLists.txt b/benchmarks/core/CMakeLists.txt new file mode 100644 index 00000000..43902191 --- /dev/null +++ b/benchmarks/core/CMakeLists.txt @@ -0,0 +1,25 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitqlib_core_benchmark LANGUAGES CXX) +if(NOT TARGET rabitqlib::core) + find_package(rabitqlib CONFIG REQUIRED COMPONENTS core) +endif() +add_executable(rabitq_core_benchmark main.cpp) +target_link_libraries(rabitq_core_benchmark PRIVATE rabitqlib::core) +add_executable(rabitq_core_scoring_benchmark scoring.cpp scoring_legacy.cpp) +target_link_libraries(rabitq_core_scoring_benchmark PRIVATE rabitqlib::core) +add_executable(rabitq_core_scoring_check_test scoring_check_test.cpp) +target_link_libraries(rabitq_core_scoring_check_test PRIVATE rabitqlib::core) +# The out-of-line comparator requires a genuine translation-unit boundary. +set_target_properties(rabitq_core_scoring_benchmark PROPERTIES INTERPROCEDURAL_OPTIMIZATION OFF) +if(COMMAND rabitq_configure_target) + rabitq_configure_target(rabitq_core_benchmark) + rabitq_configure_target(rabitq_core_scoring_benchmark) + rabitq_configure_target(rabitq_core_scoring_check_test) +endif() +enable_testing() +add_test(NAME CoreScoringCheck COMMAND rabitq_core_scoring_check_test) +add_test(NAME CoreBenchmarkSmoke COMMAND rabitq_core_benchmark 96 4 256 8 1) +add_test(NAME CoreBenchmarkOneBit COMMAND rabitq_core_benchmark 64 1 32 2 1) +add_test(NAME CoreScoringPacked COMMAND rabitq_core_scoring_benchmark 128 4 packed hot 4096 1 1) +add_test(NAME CoreScoringUnpacked COMMAND rabitq_core_scoring_benchmark 65 9 unpacked random 4096 1 1) +add_test(NAME CoreScoringOneBit COMMAND rabitq_core_scoring_benchmark 32 1 packed sequential 4096 1 1) diff --git a/benchmarks/core/README.md b/benchmarks/core/README.md new file mode 100644 index 00000000..efa244cf --- /dev/null +++ b/benchmarks/core/README.md @@ -0,0 +1,118 @@ +# Standalone core benchmark + +This C++17 executable links only `rabitqlib::core`. It measures stable v1 encoding, +query preparation, packed scoring and packed-to-unpacked conversion using caller-owned +storage. No ANN implementation, Python environment or dataset download is required. + +Build from the checkout: + +```bash +cmake -S . -B build-core-bench -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_INDEX=OFF -DRABITQ_BUILD_FASTSCAN=OFF \ + -DRABITQ_BUILD_SAMPLES=OFF -DRABITQ_USE_OPENMP=OFF \ + -DRABITQ_BUILD_BENCHMARKS=ON -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF +cmake --build build-core-bench --parallel +ctest --test-dir build-core-bench --output-on-failure +./build-core-bench/benchmarks/core/rabitq_core_benchmark > core-results.csv +``` + +For an installed package, configure this directory directly: + +```bash +cmake -S benchmarks/core -B build-core-bench-installed \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_PREFIX_PATH="" +cmake --build build-core-bench-installed --parallel +ctest --test-dir build-core-bench-installed --output-on-failure +./build-core-bench-installed/rabitq_core_benchmark 128 4 4096 64 7 > core-results.csv +``` + +Multi-configuration generators also need `--config Release` / `-C Release`, and the +executable is inside `Release/`. Positional arguments are dimension, total bits, +number of vectors, number of queries and measured repetitions. Defaults are +`128 4 4096 64 5`. This benchmark's FHT/Kac rotation requires a dimension at least 64 +and divisible by 32; bits are 1–9 and at least ten vectors are required. These are +benchmark constraints, not additional restrictions on the stable unpacked API. + +The deterministic uniform dataset uses a fixed LCG seed 42; one shared FHT/Kac rotator +is initialized with seed 42 for both data and queries. Codes use SplitPacked, L2 and a +zero centroid. All storage allocation, rotation, initial encoding and query preparation +happen before measurement. Each operation has one untimed warmup, followed by the +requested number of timed passes. Query preparation is timed explicitly in its own row. +The executable is single-threaded even in an OpenMP-enabled build. + +The output includes raw time, operations/s, reciprocal throughput in ns/operation and +a checksum for every repetition. Comment-prefixed metadata records parameters, available +ISA capabilities, OpenMP build choice, distance RMSE, maximum absolute distance error +and recall@10 of an exhaustive quantized scan. Ground truth is double-precision L2 on +the actual rotated float32 vectors. There is no ANN graph/pruning or reranking. Recall +therefore measures this quantizer/workload, not an index's search recall or end-to-end QPS. + +The scoring pass traverses all encoded vectors for each prepared query, unlike the +earlier single-code cache-resident microbenchmark. Timing excludes top-k selection and +quality computation. Do not compare these operations/s directly with index QPS. Quality +values describe approximation error; they are not tolerances allowing semantic changes. + +For comparisons, retain all raw repetitions and record CPU model, compiler/version, +build flags, affinity/SMT policy, OS and dataset parameters beside the CSV. Keep these +fixed and alternate baseline/candidate runs; compare quality before claiming a speedup. +Linux users can pin the executable with `taskset -c `. The program reports available +ISA features, not a trace of every selected kernel. CI only smoke-tests functionality; +shared CI timing is not a performance gate. + +## Scoring-call overhead + +The stable v1 `compute` takes `CodeView` by value. Its thin inline entry forwards +that local view by reference to the private checked implementation, avoiding an +aggregate copy across the compiled call boundary in optimized direct calls. +Validation stays in the checked implementation, and taking the public function's +address remains supported. Use matching headers/library and rebuild consumers +when upgrading. Performance measured for optimized direct calls does not imply +the same gain through function pointers or in unoptimized builds. The checked +interface can cost more than an unchecked inlined estimator. + +`rabitq_core_scoring_benchmark` compares the checked v1 call with three benchmark-only +unchecked references: the same kernels/formula inlined into the loop, a separate +translation-unit call taking `CodeView` by value, and a separate call taking it by +reference. These references are not library +APIs and only receive valid encoded inputs. They help distinguish validation overhead +from the combined effects of call boundaries, argument passing and inlining; subtracting +their timings does not isolate one CPU instruction's cost. + +```bash +./build-core-bench/benchmarks/core/rabitq_core_scoring_benchmark \ + 128 4 packed hot 1000000 7 +./build-core-bench/benchmarks/core/rabitq_core_scoring_benchmark \ + 768 9 packed random 262144 7 64 +./build-core-bench/benchmarks/core/rabitq_core_scoring_benchmark \ + 65 4 unpacked sequential 32768 7 64 +``` + +Arguments are dimension, total bits, `packed|unpacked`, `hot|sequential|random`, minimum +iterations, repetitions and optional code/factor working-set target in MiB (default 64). +An optional final `all|v1` argument selects timed paths; `v1` still checks every prototype +against all three references before timing and is useful for a large baseline/candidate matrix. +No arguments select `128 4 packed hot 1000000 5`. Packed dimensions must be multiples +of 32; unpacked dimensions may be any supported dimension. Build without LTO or native +tuning for the intended comparison. The target disables CMake IPO, but callers must +also avoid injecting manual LTO flags. + +Hot mode repeatedly scores one code. The other modes use a power-of-two record count +fitting the byte target, with sequential IDs or a fixed shuffled permutation. Metadata +reports actual code/factor bytes and the additional order-array bytes; this is a defined +access pattern, not a promise that accesses miss a particular cache. Each timed pass +visits at least the entire bank, even if the requested minimum iterations is smaller. + +Data/query values use an LCG with seed 42; the permutation uses seed 91. At most 32 +valid code/factor prototypes are encoded, then replicated at distinct memory addresses. +This models access and call costs, not dataset-dependent ANN quality. Inputs are already +in the numerical working domain; no rotation or top-k retrieval is timed or evaluated. +The metric is L2 with a zero centroid. Prototype scores must match all three references +exactly before timing. Raw timings and double-precision checksums are printed for every +path and repetition; path order rotates and each timed pass has its own warmup. + +Compare baseline/candidate binaries with identical harness source, compiler flags and +parameters. Alternate process order, retain all samples, and verify checksums alongside +performance. Record machine/affinity/SMT details as for the main benchmark. The unchecked +paths are useful controls, but their generated code need not match a real downstream +application. Neither these loops nor reciprocal throughput report ANN QPS or isolated +call latency. diff --git a/benchmarks/core/main.cpp b/benchmarks/core/main.cpp new file mode 100644 index 00000000..745e9ffa --- /dev/null +++ b/benchmarks/core/main.cpp @@ -0,0 +1,223 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace { +namespace api = rabitqlib::quant::v1; +volatile double sink = 0; + +size_t number(const char* text) { + size_t result = 0; + const char* end = text + std::strlen(text); + const auto parsed = std::from_chars(text, end, result); + if (parsed.ec != std::errc{} || parsed.ptr != end || result == 0) { + throw std::invalid_argument("Arguments must be positive decimal integers"); + } + return result; +} + +size_t product(size_t a, size_t b) { + if (a > std::numeric_limits::max() / b) { + throw std::length_error("Benchmark dimensions overflow size_t"); + } + return a * b; +} + +template +void measure(const char* operation, size_t operations, size_t repeats, Function work) { + sink = work(); // One untimed warmup; setup and allocation are outside work. + for (size_t repeat = 0; repeat < repeats; ++repeat) { + const auto begin = std::chrono::steady_clock::now(); + const double checksum = work(); + const double seconds = + std::chrono::duration(std::chrono::steady_clock::now() - begin).count(); + sink = checksum; + std::cout << operation << ',' << repeat << ',' << operations << ',' << seconds + << ',' << static_cast(operations) / seconds << ',' + << seconds * 1e9 / static_cast(operations) << ',' << checksum + << '\n'; + } +} + +int run(int argc, char** argv) { + if (argc != 1 && argc != 6) { + throw std::invalid_argument( + "Usage: rabitq_core_benchmark [dim bits vectors queries repetitions]" + ); + } + const size_t dim = argc == 1 ? 128 : number(argv[1]); + const size_t bits = argc == 1 ? 4 : number(argv[2]); + const size_t count = argc == 1 ? 4096 : number(argv[3]); + const size_t queries = argc == 1 ? 64 : number(argv[4]); + const size_t repeats = argc == 1 ? 5 : number(argv[5]); + constexpr size_t k = 10; + if (dim < 64 || dim % 32 != 0) { + throw std::invalid_argument( + "This FHT/Kac benchmark requires dim >= 64 and divisible by 32" + ); + } + if (count < k) { + throw std::invalid_argument("At least ten vectors are needed for recall@10"); + } + const api::Parameters parameters{dim, bits, rabitqlib::METRIC_L2}; + const auto sizes = api::requirements(parameters, api::CodeLayout::SplitPacked); + const size_t comparisons = product(count, queries); + std::vector data(product(count, dim)), query(product(queries, dim)), input(dim); + std::unique_ptr> rotation( + rabitqlib::choose_rotator(dim, rabitqlib::RotatorType::FhtKacRotator, 0, 42) + ); + uint32_t seed = 42; + for (auto* matrix : {&data, &query}) { + for (size_t row = 0; row < matrix->size() / dim; ++row) { + for (float& value : input) { + seed = seed * 1664525U + 1013904223U; + value = static_cast(seed >> 8U) / 8388608.0F - 1.0F; + } + rotation->rotate(input.data(), matrix->data() + row * dim); + } + } + std::vector scratch( + (sizes.workspace_bytes + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t) + ); + const api::Buffer workspace{scratch.data(), sizes.workspace_bytes}; + std::vector signs(product(count, sizes.code_bytes)); + // At one bit the unused extra buffer has a harmless one-byte stride. + const size_t extra_stride = std::max(size_t{1}, sizes.extra_bytes); + std::vector extra(product(count, extra_stride)); + std::vector factors(count); + std::vector states(queries); + std::vector unpacked(product(dim, sizeof(uint16_t))); + const auto code = [&](size_t row) { + return api::CodeView{ + api::CodeLayout::SplitPacked, + {signs.data() + row * sizes.code_bytes, sizes.code_bytes}, + {extra.data() + row * extra_stride, sizes.extra_bytes}}; + }; + const auto encode = [&](size_t row) { + return api::encode( + parameters, + data.data() + row * dim, + nullptr, + {api::CodeLayout::SplitPacked, + {signs.data() + row * sizes.code_bytes, sizes.code_bytes}, + {extra.data() + row * extra_stride, sizes.extra_bytes}}, + workspace + ); + }; + for (size_t row = 0; row < count; ++row) { + factors[row] = encode(row); + } + for (size_t row = 0; row < queries; ++row) { + states[row] = api::prepare_query(parameters, query.data() + row * dim, nullptr); + } + + double squared_error = 0, max_error = 0; + size_t hits = 0; + std::vector> exact(count), estimated(count); + for (size_t q = 0; q < queries; ++q) { + for (size_t row = 0; row < count; ++row) { + double distance = 0; + for (size_t j = 0; j < dim; ++j) { + const double delta = + static_cast(data[row * dim + j]) - query[q * dim + j]; + distance += delta * delta; + } + const double approximation = api::compute(states[q], code(row), factors[row]); + exact[row] = {distance, row}; + estimated[row] = {approximation, row}; + const double error = approximation - distance; + squared_error += error * error; + max_error = std::max(max_error, std::abs(error)); + } + std::partial_sort(exact.begin(), exact.begin() + k, exact.end()); + std::partial_sort(estimated.begin(), estimated.begin() + k, estimated.end()); + for (size_t i = 0; i < k; ++i) { + for (size_t j = 0; j < k; ++j) { + hits += estimated[i].second == exact[j].second ? 1 : 0; + } + } + } + std::cout + << std::setprecision(12) << "# dim=" << dim << " bits=" << bits + << " vectors=" << count << " queries=" << queries << " repetitions=" << repeats + << '\n' + << "# metric=L2 layout=SplitPacked threads=1 seed=42 rotation=FHT/Kac\n" + << "# rotation/setup/allocations excluded; full scan has no reranking\n" + << "# avx2=" << rabitqlib::cpu::has_avx2() + << " avx512_core=" << rabitqlib::cpu::has_avx512_core() + << " neon=" << rabitqlib::cpu::has_neon() << " openmp_build=" << RABITQ_USE_OPENMP + << '\n' + << "# rmse=" << std::sqrt(squared_error / static_cast(comparisons)) + << " max_abs_error=" << max_error << " flat_recall_at_10=" + << static_cast(hits) / static_cast(product(queries, k)) << '\n' + << "operation,repeat,operations,seconds,ops_per_second,ns_per_operation,checksum\n"; + measure("encode_packed", count, repeats, [&] { + double checksum = 0; + for (size_t row = 0; row < count; ++row) { + factors[row] = encode(row); + checksum += factors[row].add; + } + return checksum; + }); + const size_t preparations = product(queries, size_t{256}); + measure("prepare_query", preparations, repeats, [&] { + double checksum = 0; + for (size_t row = 0; row < preparations; ++row) { + const size_t q = row % queries; + states[q] = api::prepare_query(parameters, query.data() + q * dim, nullptr); + checksum += static_cast(states[q].dimension()); + } + return checksum; + }); + measure("compute_packed", comparisons, repeats, [&] { + double checksum = 0; + for (size_t q = 0; q < queries; ++q) { + for (size_t row = 0; row < count; ++row) { + checksum += api::compute(states[q], code(row), factors[row]); + } + } + return checksum; + }); + measure("convert_to_unpacked", count, repeats, [&] { + double checksum = 0; + for (size_t row = 0; row < count; ++row) { + api::convert_codes( + parameters, + code(row), + {api::CodeLayout::UnpackedU16, {unpacked.data(), unpacked.size()}, {}}, + workspace + ); + checksum += unpacked.front(); + } + return checksum; + }); + return 0; +} +} // namespace + +int main(int argc, char** argv) { + try { + return run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/benchmarks/core/scoring.cpp b/benchmarks/core/scoring.cpp new file mode 100644 index 00000000..c4439a1f --- /dev/null +++ b/benchmarks/core/scoring.cpp @@ -0,0 +1,301 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" +#include "scoring_legacy.hpp" + +namespace { +namespace api = rabitqlib::quant::v1; +volatile double sink = 0; + +size_t number(const char* text) { + size_t value = 0; + const char* end = text + std::strlen(text); + const auto parsed = std::from_chars(text, end, value); + if (parsed.ec != std::errc{} || parsed.ptr != end || value == 0) { + throw std::invalid_argument("Expected a positive decimal integer"); + } + return value; +} + +size_t product(size_t a, size_t b) { + if (b != 0 && a > std::numeric_limits::max() / b) { + throw std::length_error("Benchmark buffer size overflows size_t"); + } + return a * b; +} + +template +double score_loop( + Score score, + size_t iterations, + const std::vector& order, + const std::vector& codes, + const std::vector& extras, + const std::vector& factors, + const api::Requirements& sizes, + api::CodeLayout layout +) { + double sum = 0; + const size_t extra_stride = std::max(size_t{1}, sizes.extra_bytes); + for (size_t i = 0; i < iterations; ++i) { + const size_t row = order[i & (order.size() - 1)]; + sum += score( + {layout, + {codes.data() + row * sizes.code_bytes, sizes.code_bytes}, + {extras.data() + row * extra_stride, sizes.extra_bytes}}, + factors[row] + ); + } + return sum; +} + +int run(int argc, char** argv) { + if (argc != 1 && argc != 7 && argc != 8 && argc != 9) { + throw std::invalid_argument( + "Usage: rabitq_core_scoring_benchmark [dim bits packed|unpacked " + "hot|sequential|random iterations repetitions [working_set_mib [all|v1]]]" + ); + } + const size_t dim = argc == 1 ? 128 : number(argv[1]); + const size_t bits = argc == 1 ? 4 : number(argv[2]); + const std::string layout_name = argc == 1 ? "packed" : argv[3]; + const std::string pattern = argc == 1 ? "hot" : argv[4]; + const size_t requested_iterations = argc == 1 ? 1000000 : number(argv[5]); + const size_t repeats = argc == 1 ? 5 : number(argv[6]); + const size_t mib = argc >= 8 ? number(argv[7]) : 64; + const std::string paths = argc == 9 ? argv[8] : "all"; + if (paths != "all" && paths != "v1") { + throw std::invalid_argument("Paths must be all or v1"); + } + if (layout_name != "packed" && layout_name != "unpacked") { + throw std::invalid_argument("Layout must be packed or unpacked"); + } + if (pattern != "hot" && pattern != "sequential" && pattern != "random") { + throw std::invalid_argument("Pattern must be hot, sequential or random"); + } + const auto layout = layout_name == "packed" ? api::CodeLayout::SplitPacked + : api::CodeLayout::UnpackedU16; + const api::Parameters parameters{dim, bits}; + const auto sizes = api::requirements(parameters, layout); + const size_t extra_stride = std::max(size_t{1}, sizes.extra_bytes); + const size_t row_bytes = sizes.code_bytes + extra_stride + sizeof(api::Factors); + const size_t desired_rows = product(mib, size_t{1} << 20U) / row_bytes; + size_t rows = 1; + if (pattern != "hot") { + while (rows <= desired_rows / 2) { + rows *= 2; + } + } + const size_t iterations = std::max(requested_iterations, rows); + std::vector codes(product(rows, sizes.code_bytes)); + std::vector extras(product(rows, extra_stride)); + std::vector factors(rows); + std::vector order(rows); + std::iota(order.begin(), order.end(), size_t{0}); + if (pattern == "random") { + std::mt19937 random(91); + std::shuffle(order.begin(), order.end(), random); + } + std::vector memory( + (sizes.workspace_bytes + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t) + ); + const api::Buffer work{memory.data(), sizes.workspace_bytes}; + std::vector input(dim), query(dim); + uint32_t random = 42; + const auto next = [&] { + random = random * 1664525U + 1013904223U; + return static_cast(random >> 8U) / 8388608.0F - 1.0F; + }; + for (auto& value : query) { + value = next(); + } + const size_t prototypes = std::min(rows, size_t{32}); + for (size_t row = 0; row < prototypes; ++row) { + for (auto& value : input) { + value = next(); + } + factors[row] = api::encode( + parameters, + input.data(), + nullptr, + {layout, + {codes.data() + row * sizes.code_bytes, sizes.code_bytes}, + {extras.data() + row * extra_stride, sizes.extra_bytes}}, + work + ); + } + for (size_t row = prototypes; row < rows; ++row) { + const size_t source = row % prototypes; + std::memcpy( + codes.data() + row * sizes.code_bytes, + codes.data() + source * sizes.code_bytes, + sizes.code_bytes + ); + std::memcpy( + extras.data() + row * extra_stride, + extras.data() + source * extra_stride, + sizes.extra_bytes + ); + factors[row] = factors[source]; + } + const auto prepared = api::prepare_query(parameters, query.data(), nullptr); + const scoring_benchmark::Query legacy{ + query.data(), + dim, + bits, + rabitqlib::l2norm_sqr(query.data(), dim), + std::accumulate(query.begin(), query.end(), 0.0F) * -0.5F, + rabitqlib::select_excode_ipfunc(bits - 1)}; + const double eps = std::numeric_limits::epsilon(); + const double gamma = (dim + 8) * eps / (1 - (dim + 8) * eps); + const double max_code = static_cast((1U << bits) - 1); + const double absolute_query = + std::accumulate(query.begin(), query.end(), 0.0, [](double sum, float value) { + return sum + std::abs(value); + }); + for (size_t row = 0; row < prototypes; ++row) { + const api::CodeView code{ + layout, + {codes.data() + row * sizes.code_bytes, sizes.code_bytes}, + {extras.data() + row * extra_stride, sizes.extra_bytes}}; + const float checked = api::compute(prepared, code, factors[row]); + // The installed library and this consumer can use different reduction + // orders. Bound their float32 rounding error before timing either path. + const double tolerance = + 2 * gamma * std::abs(factors[row].rescale) * max_code * absolute_query + + 8 * eps * + (1 + std::abs(factors[row].add) + std::abs(legacy.add) + std::abs(checked)); + if (!scoring_benchmark::scores_match( + checked, + scoring_benchmark::unchecked_inline(legacy, code, factors[row]), + tolerance + ) || + !scoring_benchmark::scores_match( + checked, + scoring_benchmark::unchecked_call(legacy, code, factors[row]), + tolerance + ) || + !scoring_benchmark::scores_match( + checked, + scoring_benchmark::unchecked_ref(legacy, code, factors[row]), + tolerance + )) { + throw std::runtime_error("Checked and unchecked score results differ"); + } + } + std::cout + << std::setprecision(17) + << "# seed=42 metric=L2 centroid=zero threads=1 setup/encoding excluded\n" + << "# rows=" << rows << " code_factor_bytes=" << product(rows, row_bytes) + << " order_bytes=" << product(rows, sizeof(size_t)) << " prototypes=" << prototypes + << '\n' + << "# avx2=" << rabitqlib::cpu::has_avx2() + << " avx512_core=" << rabitqlib::cpu::has_avx512_core() + << " neon=" << rabitqlib::cpu::has_neon() << '\n' + << "dim,bits,layout,pattern,path,repeat,operations,seconds,ns_per_score,checksum\n"; + const std::array names{ + "v1", "unchecked_inline", "unchecked_call", "unchecked_ref"}; + const auto execute = [&](size_t path, size_t count) { + if (path == 0) { + return score_loop( + [&](api::CodeView c, const api::Factors& f) { + return api::compute(prepared, c, f); + }, + count, + order, + codes, + extras, + factors, + sizes, + layout + ); + } + if (path == 1) { + return score_loop( + [&](api::CodeView c, const api::Factors& f) { + return scoring_benchmark::unchecked_inline(legacy, c, f); + }, + count, + order, + codes, + extras, + factors, + sizes, + layout + ); + } + if (path == 2) { + return score_loop( + [&](api::CodeView c, const api::Factors& f) { + return scoring_benchmark::unchecked_call(legacy, c, f); + }, + count, + order, + codes, + extras, + factors, + sizes, + layout + ); + } + return score_loop( + [&](api::CodeView c, const api::Factors& f) { + return scoring_benchmark::unchecked_ref(legacy, c, f); + }, + count, + order, + codes, + extras, + factors, + sizes, + layout + ); + }; + for (size_t repeat = 0; repeat < repeats; ++repeat) { + const size_t path_count = paths == "all" ? names.size() : 1; + for (size_t offset = 0; offset < path_count; ++offset) { + const size_t path = (repeat + offset) % path_count; + sink = execute(path, std::max(rows, iterations)); + const auto start = std::chrono::steady_clock::now(); + const double checksum = execute(path, iterations); + const double seconds = + std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); + sink = checksum; + std::cout << dim << ',' << bits << ',' << layout_name << ',' << pattern << ',' + << names[path] << ',' << repeat << ',' << iterations << ',' << seconds + << ',' << seconds * 1e9 / static_cast(iterations) << ',' + << checksum << '\n'; + } + } + return 0; +} +} // namespace + +int main(int argc, char** argv) { + try { + return run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/benchmarks/core/scoring_check_test.cpp b/benchmarks/core/scoring_check_test.cpp new file mode 100644 index 00000000..2e36b9a9 --- /dev/null +++ b/benchmarks/core/scoring_check_test.cpp @@ -0,0 +1,23 @@ +#include +#include + +#include "scoring_legacy.hpp" + +int main() { + using scoring_benchmark::scores_match; + if (!scores_match(1, 1, 0) || !scores_match(1, 1.125F, 0.125) || + !scores_match(1.125F, 1, 0.125) || scores_match(1, 1.25F, 0.125) || + scores_match(1, 1, -1)) { + return 1; + } + for (float invalid : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + if (scores_match(invalid, 1, 1) || scores_match(1, invalid, 1) || + scores_match(invalid, invalid, 1) || scores_match(1, 1, invalid)) { + return 1; + } + } + return 0; +} diff --git a/benchmarks/core/scoring_legacy.cpp b/benchmarks/core/scoring_legacy.cpp new file mode 100644 index 00000000..53930d6e --- /dev/null +++ b/benchmarks/core/scoring_legacy.cpp @@ -0,0 +1,12 @@ +#include "scoring_legacy.hpp" + +#include "rabitqlib/core/api.hpp" + +namespace scoring_benchmark { +float unchecked_call(const Query& query, api::CodeView code, const api::Factors& f) { + return unchecked_inline(query, code, f); +} +float unchecked_ref(const Query& query, const api::CodeView& code, const api::Factors& f) { + return unchecked_inline(query, code, f); +} +} // namespace scoring_benchmark diff --git a/benchmarks/core/scoring_legacy.hpp b/benchmarks/core/scoring_legacy.hpp new file mode 100644 index 00000000..a42b8d32 --- /dev/null +++ b/benchmarks/core/scoring_legacy.hpp @@ -0,0 +1,59 @@ +#pragma once + +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace scoring_benchmark { +namespace api = rabitqlib::quant::v1; + +inline bool scores_match(float actual, float expected, double tolerance) { + return std::isfinite(actual) && std::isfinite(expected) && std::isfinite(tolerance) && + tolerance >= 0 && std::abs(static_cast(actual) - expected) <= tolerance; +} + +struct Query { + const float* values; + size_t dim; + size_t bits; + float add; + float half_sum; + float (*extra_ip)(const float*, const uint8_t*, size_t); +}; + +// Same full-distance expression and existing kernels, with no public-API checks. +// The benchmark always supplies validated buffers; this is not a library entry point. +inline float unchecked_inline( + const Query& query, api::CodeView code, const api::Factors& f +) { + const auto* bytes = static_cast(code.codes.data); + float ip = 0; + if (code.layout == api::CodeLayout::UnpackedU16) { + for (size_t i = 0; i < query.dim; ++i) { + const uint16_t value = + static_cast(bytes[2 * i] | (uint16_t{bytes[2 * i + 1]} << 8U)); + ip += query.values[i] * static_cast(value); + } + } else { + const size_t extra_bits = query.bits - 1; + ip = static_cast(1U << extra_bits) * + rabitqlib::mask_ip_x0_q(query.values, bytes, query.dim); + if (extra_bits > 0) { + ip += query.extra_ip( + query.values, static_cast(code.extra.data), query.dim + ); + } + } + return rabitqlib::quant::detail::full_est_dist_from_ip( + ip, query.bits, f.add, f.rescale, query.add, query.half_sum + ); +} + +// Compiled in a separate translation unit, without LTO, to retain a call boundary. +float unchecked_call(const Query& query, api::CodeView code, const api::Factors& f); +float unchecked_ref(const Query& query, const api::CodeView& code, const api::Factors& f); +} // namespace scoring_benchmark diff --git a/benchmarks/hotpaths/CMakeLists.txt b/benchmarks/hotpaths/CMakeLists.txt new file mode 100644 index 00000000..9f356d99 --- /dev/null +++ b/benchmarks/hotpaths/CMakeLists.txt @@ -0,0 +1,24 @@ +# Keep each benchmark's dependency closure at the boundary it measures. +function(rabitq_hotpath_target name component) + add_executable(rabitq_${name}_hotpaths ${name}.cpp) + target_link_libraries(rabitq_${name}_hotpaths PRIVATE rabitqlib::${component}) + rabitq_configure_target(rabitq_${name}_hotpaths) + set_target_properties(rabitq_${name}_hotpaths PROPERTIES INTERPROCEDURAL_OPTIMIZATION OFF) +endfunction() +rabitq_hotpath_target(core core) +add_test(NAME CoreHotpathEncoding COMMAND rabitq_core_hotpaths 64 32 1 2 1 encode) +add_test(NAME CoreHotpathL2 COMMAND rabitq_core_hotpaths 128 32 1 2 1 l2) +if(RABITQ_BUILD_FASTSCAN) + rabitq_hotpath_target(fastscan fastscan) + add_test(NAME FastScanHotpathStandard COMMAND rabitq_fastscan_hotpaths 128 2 1 standard) + add_test(NAME FastScanHotpathHacc COMMAND rabitq_fastscan_hotpaths 128 2 1 hacc) +endif() +if(RABITQ_BUILD_INDEX) + rabitq_hotpath_target(load index) + add_test(NAME SidecarHotpathPrepare COMMAND rabitq_load_hotpaths prepare + "${CMAKE_CURRENT_BINARY_DIR}/sidecar-smoke" 128 64 1) + add_test(NAME SidecarHotpathLoad COMMAND rabitq_load_hotpaths load + "${CMAKE_CURRENT_BINARY_DIR}/sidecar-smoke" 128 64 2) + set_tests_properties(SidecarHotpathPrepare PROPERTIES FIXTURES_SETUP sidecar_hotpath) + set_tests_properties(SidecarHotpathLoad PROPERTIES FIXTURES_REQUIRED sidecar_hotpath) +endif() diff --git a/benchmarks/hotpaths/README.md b/benchmarks/hotpaths/README.md new file mode 100644 index 00000000..f8d30c90 --- /dev/null +++ b/benchmarks/hotpaths/README.md @@ -0,0 +1,108 @@ +# Component hot-path benchmarks + +These deterministic benchmarks retain the small workloads that exposed modularization +costs. They complement the [stable core benchmark](../core/README.md) and +[full ANN benchmarks](../README.md); their throughput is not end-to-end ANN QPS. + +Enable `RABITQ_BUILD_BENCHMARKS=ON` in a Release build. Targets follow enabled components: + +| Executable | Dependency | Measurement | +| --- | --- | --- | +| `rabitq_core_hotpaths` | core only | One-bit compact encoding or raw L2 calls | +| `rabitq_fastscan_hotpaths` | FastScan | Cached standard/HACC batch scoring | +| `rabitq_load_hotpaths` | index | HNSW centroid-sidecar load and validation | +| `rabitq_ids_benchmark` module + `ids.py` | Python bindings | Production cluster-ID conversion helper | + +No additional production dependencies are introduced. Core/FastScan benchmarks support +OpenMP OFF. Python benchmarks use the existing Python/NumPy/pybind11 environment and +are available only with `RABITQ_BUILD_PYTHON_BINDINGS=ON`. The helper module is not installed +in the package. `ctest` smoke tests verify functionality, never shared-host speed thresholds. + +## Build and run + +```bash +cmake -S . -B build-hotpaths -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_BENCHMARKS=ON -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF +cmake --build build-hotpaths --parallel +ctest --test-dir build-hotpaths -R 'Hotpath' --output-on-failure +``` + +Set both `RABITQ_BUILD_INDEX=OFF` and `RABITQ_BUILD_FASTSCAN=OFF` for pure core. +For the Python helper, additionally enable bindings and select your existing interpreter +with `Python_EXECUTABLE` and its `pybind11_DIR`, as described in +[development setup](../../DEVELOPMENT.md). Multi-configuration generators need +`--config Release`, `ctest -C Release`, and the `Release/` executable subdirectory. + +Examples (from the repository root): + +```bash +# dim rows threads iterations repetitions mode +./build-hotpaths/benchmarks/hotpaths/rabitq_core_hotpaths 64 4096 1 1024 3 encode +./build-hotpaths/benchmarks/hotpaths/rabitq_core_hotpaths 128 1 1 2000000 5 l2 + +# dim batches repetitions standard|hacc: one batch contains 32 vectors +./build-hotpaths/benchmarks/hotpaths/rabitq_fastscan_hotpaths 128 50000000 3 hacc + +# Prepare once; both versions must load the same file. The writer appends .hnsw. +./build-hotpaths/benchmarks/hotpaths/rabitq_load_hotpaths prepare "$TMPDIR/centroids" 128 16384 1 +./build-hotpaths/benchmarks/hotpaths/rabitq_load_hotpaths load "$TMPDIR/centroids" 128 16384 7 + +python benchmarks/hotpaths/ids.py --module-dir build-hotpaths/python_bindings +``` + +Use a task-specific writable `TMPDIR`. `prepare` reuses and validates an existing +fixture instead of replacing it. It rejects a file whose centroid contents do not match +the deterministic input. Original production indexes are not benchmark fixtures. + +## Timing and correctness contract + +- Inputs use an LCG with seed 42 (query/centroid seed 91 where applicable), metric L2. + Dimensions for encoding/scoring are multiples of 32. Use 64, 128 and 960 dimensions; + encoding rows 32/4096 and threads 1/4 retain the cases from the original investigation. +- Core encoding times the existing serial template at one thread and OpenMP template + above one thread. It includes scratch reuse and encoder scheduling, excludes input/output + allocation, and verifies serial/parallel code and factor hashes before timing. +- L2 checks a float64 reference with absolute tolerance `2e-6 * dim`. Results are stored + per row; this benchmark includes those stores and is not an isolated instruction latency. +- FastScan repeatedly scores one 32-vector batch in cache. It checks generic correction + within `2e-6 * dim`, reports one-bit filtering RMSE and recall@10 against float64 L2, + and hashes estimated distances, lower bounds and inner products. The query uses three + extra bits for preparation, but this benchmark times only the one-bit filter, not reranking. +- Loading times `load` only, including its allocations and validation. Initial constructor, + destructor, centroid verification and reference search are outside timing. The first + untimed load warms the filesystem cache. Each subsequent load verifies every centroid + and hashes the same eight search results. +- Python times conversion with its owned output and Python call overhead, for native + uint32 and int64 arrays. It excludes input generation and prints the helper's path/hash. + Its `hash` column is only the returned final ID, not a full-array correctness proof; + the binding regression tests provide boundary coverage. It does not measure full builds. +- C++ results are hashed outside timing; each case has one untimed warmup. IPO is disabled + for these targets; do not inject manual LTO flags. Do not run builds/profilers during timing. + +Use an identical compiler, Release flags, ISA policy, inputs, affinity and thread placement +for before/after. Native tuning is optional for local experiments and must be identical on +both sides; portable correctness validation uses native OFF. Record CPU, compiler/version, +flags, revisions, OpenMP settings and NUMA policy. Pin memory as well as CPU **before loading** +on NUMA hosts. Available ISA metadata does not prove all backends executed. + +## Alternating comparisons and retained evidence + +`compare.py` runs commands without a shell, alternates process order, checks matching CSV +output fingerprints, and saves every stdout/stderr file plus status and medians. It requires +a new output directory so earlier failures are not overwritten. For example: + +```bash +python benchmarks/hotpaths/compare.py \ + --before './before/benchmarks/hotpaths/rabitq_fastscan_hotpaths 128 50000000 3 hacc' \ + --after './after/benchmarks/hotpaths/rabitq_fastscan_hotpaths 128 50000000 3 hacc' \ + --pairs 3 --output "$TMPDIR/hacc-comparison" \ + --notes 'before/after revisions; CPU; compiler/version; flags; threads=1; CPU/memory binding' +``` + +Launch the runner under your host's CPU/memory binding tool; children inherit its placement. +The notes are required protocol metadata, not automatic verification of NUMA placement. +Review quality metadata and individual repetitions alongside medians. A hash difference stops +comparison for investigation; it does not prove a numerical bug across different compilers/ISAs. +Repeated matching hashes also do not establish correctness by themselves. Retain failed runs, +report regressions, and do not infer index QPS gains from these microbenchmarks. diff --git a/benchmarks/hotpaths/common.hpp b/benchmarks/hotpaths/common.hpp new file mode 100644 index 00000000..9de5ddb5 --- /dev/null +++ b/benchmarks/hotpaths/common.hpp @@ -0,0 +1,93 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/utils/cpu_features.hpp" + +namespace bench { +using Clock = std::chrono::steady_clock; +inline volatile uint64_t sink = 0; + +inline size_t number(const char* text) { + size_t value = 0; + const char* end = text + std::strlen(text); + const auto parsed = std::from_chars(text, end, value); + if (parsed.ec != std::errc{} || parsed.ptr != end || value == 0) { + throw std::invalid_argument("Arguments must be positive decimal integers"); + } + return value; +} + +inline size_t product(size_t a, size_t b) { + if (a > std::numeric_limits::max() / b) { + throw std::length_error("Benchmark size overflow"); + } + return a * b; +} + +inline std::vector values(size_t count, uint32_t seed = 42) { + std::vector result(count); + for (auto& value : result) { + seed = seed * 1664525U + 1013904223U; + value = static_cast(seed >> 8) / 8388608.0F - 1.0F; + } + return result; +} + +inline uint64_t hash( + const void* data, size_t bytes, uint64_t seed = 14695981039346656037ULL +) { + const auto* source = static_cast(data); + for (size_t i = 0; i < bytes; ++i) { + seed = (seed ^ source[i]) * 1099511628211ULL; + } + return seed; +} + +inline void metadata(size_t dim, size_t rows, size_t threads) { + std::cout << std::setprecision(12) << "# dim=" << dim << " rows=" << rows + << " threads=" << threads << " seed=42 metric=L2 openmp=" << RABITQ_USE_OPENMP + << " avx2=" << rabitqlib::cpu::has_avx2() + << " avx512=" << rabitqlib::cpu::has_avx512_core() + << " neon=" << rabitqlib::cpu::has_neon() << '\n'; + std::cout << "operation,repetition,operations,seconds,ns_per_operation,hash\n"; +} + +inline void row( + const char* name, size_t rep, size_t operations, double seconds, uint64_t h +) { + sink = h; + std::cout << name << ',' << rep << ',' << operations << ',' << seconds << ',' + << seconds * 1e9 / static_cast(operations) << ',' << h << '\n'; +} + +// Hashing is deliberately outside the timed region. Each work invocation must +// leave observable results, and the benchmark targets are built without IPO. +template +void measure( + const char* name, size_t operations, size_t reps, Work work, Checksum checksum +) { + work(); + const uint64_t expected = checksum(); + for (size_t rep = 0; rep < reps; ++rep) { + const auto start = Clock::now(); + work(); + const double seconds = std::chrono::duration(Clock::now() - start).count(); + const uint64_t actual = checksum(); + if (actual != expected) { + throw std::runtime_error("Benchmark output changed between repetitions"); + } + row(name, rep, operations, seconds, actual); + } +} +} // namespace bench diff --git a/benchmarks/hotpaths/compare.py b/benchmarks/hotpaths/compare.py new file mode 100644 index 00000000..df8edd95 --- /dev/null +++ b/benchmarks/hotpaths/compare.py @@ -0,0 +1,139 @@ +"""Alternate two commands, retain raw evidence, and compare output fingerprints.""" + +import argparse +import csv +import io +import json +import os +import platform +import shlex +import statistics +import subprocess +import time +from pathlib import Path + + +def records(text): + rows = list( + csv.DictReader( + io.StringIO( + "\n".join( + line for line in text.splitlines() if not line.startswith("#") + ) + ) + ) + ) + if ( + not rows + or not {"operation", "repetition", "operations", "hash", "ns_per_operation"} + <= rows[0].keys() + ): + raise ValueError("Command did not produce hot-path benchmark CSV") + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--before", required=True, help="command, parsed without a shell" + ) + parser.add_argument( + "--after", required=True, help="command, parsed without a shell" + ) + parser.add_argument("--pairs", type=int, default=3) + parser.add_argument( + "--output", type=Path, required=True, help="new evidence directory" + ) + parser.add_argument( + "--notes", required=True, help="revisions, compiler/flags and NUMA policy" + ) + args = parser.parse_args() + if args.pairs < 1: + parser.error("--pairs must be positive") + args.output.mkdir(parents=True, exist_ok=False) + cpuinfo = Path("/proc/cpuinfo") + report = { + "platform": platform.platform(), + "cpu": platform.processor(), + "affinity": sorted(os.sched_getaffinity(0)) + if hasattr(os, "sched_getaffinity") + else None, + "notes": args.notes, + "commands": { + "before": shlex.split(args.before), + "after": shlex.split(args.after), + }, + "environment": { + key: os.environ.get(key) + for key in ("OMP_NUM_THREADS", "OMP_PROC_BIND", "OMP_PLACES") + }, + "runs": [], + "status": "running", + } + if cpuinfo.exists(): + (args.output / "cpuinfo.txt").write_text(cpuinfo.read_text()) + samples = {"before": {}, "after": {}} + + def save(): + (args.output / "results.json").write_text(json.dumps(report, indent=2) + "\n") + + save() + try: + for pair in range(args.pairs): + fingerprints = {} + for version in ( + ("before", "after") if pair % 2 == 0 else ("after", "before") + ): + prefix = args.output / f"pair-{pair}-{version}" + start = time.time() + # Stream into files so failures/interruptions retain partial output. + with ( + prefix.with_suffix(".csv").open("w") as stdout, + prefix.with_suffix(".stderr").open("w") as stderr, + ): + result = subprocess.run( + report["commands"][version], + stdout=stdout, + stderr=stderr, + check=False, + ) + report["runs"].append( + { + "pair": pair, + "version": version, + "start": start, + "returncode": result.returncode, + } + ) + save() + if result.returncode: + raise RuntimeError(f"{version} failed; inspect {prefix}.stderr") + rows = records(prefix.with_suffix(".csv").read_text()) + fingerprints[version] = [ + (r["operation"], r["repetition"], r["operations"], r["hash"]) + for r in rows + ] + for row in rows: + samples[version].setdefault(row["operation"], []).append( + float(row["ns_per_operation"]) + ) + if fingerprints["before"] != fingerprints["after"]: + raise RuntimeError( + f"Output fingerprints differ in pair {pair}; review numerical results before timing" + ) + report["medians_ns"] = { + v: {op: statistics.median(xs) for op, xs in ops.items()} + for v, ops in samples.items() + } + report["status"] = "passed" + except BaseException as error: + report["status"] = "failed" + report["error"] = str(error) + raise + finally: + save() + print(json.dumps(report["medians_ns"], indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/hotpaths/core.cpp b/benchmarks/hotpaths/core.cpp new file mode 100644 index 00000000..63a7933e --- /dev/null +++ b/benchmarks/hotpaths/core.cpp @@ -0,0 +1,133 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/utils/openmp.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace { +void run(int argc, char** argv) { + if (argc != 7) { + throw std::invalid_argument( + "Usage: rabitq_core_hotpaths dim rows threads iterations repetitions " + "mode(encode|l2)" + ); + } + const size_t dim = bench::number(argv[1]), rows = bench::number(argv[2]); + const size_t threads = bench::number(argv[3]), iterations = bench::number(argv[4]); + const size_t reps = bench::number(argv[5]); + const std::string mode(argv[6]); + if (dim % 32 != 0 || threads > static_cast(std::numeric_limits::max()) || + (mode != "encode" && mode != "l2")) { + throw std::invalid_argument( + "Require dim divisible by 32, valid thread count and mode" + ); + } +#if RABITQ_USE_OPENMP + omp_set_dynamic(0); + omp_set_num_threads(static_cast(threads)); +#else + if (threads != 1) { + throw std::invalid_argument("OpenMP OFF requires threads=1"); + } +#endif + auto data = bench::values(bench::product(rows, dim)); + const auto centroid = bench::values(dim, 91); + bench::metadata(dim, rows, mode == "l2" ? 1 : threads); + if (mode == "l2") { + std::vector distances(rows); + const auto work = [&] { + for (size_t i = 0; i < iterations; ++i) { + for (size_t row = 0; row < rows; ++row) { + distances[row] = rabitqlib::euclidean_sqr( + data.data() + row * dim, centroid.data(), dim + ); + } + } + }; + work(); + for (size_t row = 0; row < rows; ++row) { + double reference = 0; + for (size_t d = 0; d < dim; ++d) { + const double delta = static_cast(data[row * dim + d]) - centroid[d]; + reference += delta * delta; + } + if (!(std::abs(distances[row] - reference) <= 2e-6 * static_cast(dim) + )) { + throw std::runtime_error("L2 differs from float64 reference"); + } + } + bench::measure("l2", bench::product(rows, iterations), reps, work, [&] { + return bench::hash(distances.data(), distances.size() * sizeof(float)); + }); + return; + } + std::vector codes(bench::product(rows, dim / 8)); + std::vector add(rows), rescale(rows), error(rows); + const auto encode = [&](bool parallel) { + using namespace rabitqlib::quant::rabitq_impl::one_bit; + if (parallel) { + one_bit_compact_codes( + data.data(), + centroid.data(), + rows, + dim, + codes.data(), + add.data(), + rescale.data(), + error.data() + ); + } else { + one_bit_compact_codes( + data.data(), + centroid.data(), + rows, + dim, + codes.data(), + add.data(), + rescale.data(), + error.data() + ); + } + }; + const auto checksum = [&] { + uint64_t h = bench::hash(codes.data(), codes.size()); + for (const auto* factors : {&add, &rescale, &error}) { + h = bench::hash(factors->data(), factors->size() * sizeof(float), h); + } + return h; + }; + encode(false); + const uint64_t serial_hash = checksum(); + encode(threads > 1); + if (checksum() != serial_hash) { + throw std::runtime_error("Serial and parallel encodings differ"); + } + bench::measure( + "encode", + bench::product(rows, iterations), + reps, + [&] { + for (size_t i = 0; i < iterations; ++i) + encode(threads > 1); + }, + checksum + ); +} +} // namespace +int main(int argc, char** argv) { + try { + run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/benchmarks/hotpaths/fastscan.cpp b/benchmarks/hotpaths/fastscan.cpp new file mode 100644 index 00000000..5dedb531 --- /dev/null +++ b/benchmarks/hotpaths/fastscan.cpp @@ -0,0 +1,125 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace { +void run(int argc, char** argv) { + if (argc != 5) { + throw std::invalid_argument( + "Usage: rabitq_fastscan_hotpaths dim batches repetitions mode(standard|hacc)" + ); + } + const size_t dim = bench::number(argv[1]), batches = bench::number(argv[2]); + const size_t reps = bench::number(argv[3]); + const std::string mode(argv[4]); + if (dim % 32 != 0 || (mode != "standard" && mode != "hacc")) { + throw std::invalid_argument("Require dim divisible by 32 and standard or hacc"); + } + constexpr size_t kRows = 32; + const bool hacc = mode == "hacc"; + auto data = bench::values(bench::product(kRows, dim)); + const auto query = bench::values(dim, 91); + const std::vector centroid(dim, 0); + std::vector storage(rabitqlib::BatchDataMap::data_bytes(dim)); + rabitqlib::quant::quantize_one_batch( + data.data(), centroid.data(), kRows, dim, storage.data() + ); + rabitqlib::SplitBatchQuery prepared( + query.data(), dim, 3, rabitqlib::METRIC_L2, hacc + ); + prepared.set_g_add(std::sqrt(rabitqlib::l2norm_sqr(query.data(), dim))); + std::array distance{}, lower{}, ip{}, ref_distance{}, ref_lower{}, + ref_ip{}; + rabitqlib::simd::split_batch_estdist_generic( + storage.data(), + prepared, + dim, + ref_distance.data(), + ref_lower.data(), + ref_ip.data(), + hacc + ); + const auto score = [&] { + rabitqlib::split_batch_estdist( + storage.data(), prepared, dim, distance.data(), lower.data(), ip.data(), hacc + ); + }; + score(); + for (size_t row = 0; row < kRows; ++row) { + const double tolerance = 2e-6 * static_cast(dim); + if (!(std::abs(distance[row] - ref_distance[row]) <= tolerance) || + !(std::abs(lower[row] - ref_lower[row]) <= tolerance) || + !(std::abs(ip[row] - ref_ip[row]) <= tolerance)) { + throw std::runtime_error("Batch estimator differs from generic correction"); + } + } + // Numerical quality is separate from timing. This is one-bit filtering of + // 32 synthetic vectors, not complete multi-bit ANN search. + std::array exact{}; + double squared_error = 0; + for (size_t row = 0; row < kRows; ++row) { + for (size_t d = 0; d < dim; ++d) { + const double delta = static_cast(data[row * dim + d]) - query[d]; + exact[row] += delta * delta; + } + squared_error += std::pow(exact[row] - distance[row], 2); + } + std::array expected_order{}, actual_order{}; + std::iota(expected_order.begin(), expected_order.end(), 0); + actual_order = expected_order; + std::sort(expected_order.begin(), expected_order.end(), [&](size_t a, size_t b) { + return exact[a] < exact[b]; + }); + std::sort(actual_order.begin(), actual_order.end(), [&](size_t a, size_t b) { + return distance[a] < distance[b]; + }); + size_t hits = 0; + for (size_t i = 0; i < 10; ++i) { + hits += std::find( + expected_order.begin(), expected_order.begin() + 10, actual_order[i] + ) != expected_order.begin() + 10; + } + bench::metadata(dim, kRows, 1); + std::cout << "# one_bit_filter_rmse=" << std::sqrt(squared_error / kRows) + << " recall10=" << static_cast(hits) / 10 << '\n'; + bench::measure( + mode.c_str(), + bench::product(batches, kRows), + reps, + [&] { + for (size_t batch = 0; batch < batches; ++batch) + score(); + }, + [&] { + uint64_t h = bench::hash(distance.data(), sizeof(distance)); + h = bench::hash(lower.data(), sizeof(lower), h); + return bench::hash(ip.data(), sizeof(ip), h); + } + ); +} +} // namespace +int main(int argc, char** argv) { + try { + run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/benchmarks/hotpaths/ids.cpp b/benchmarks/hotpaths/ids.cpp new file mode 100644 index 00000000..27469b3c --- /dev/null +++ b/benchmarks/hotpaths/ids.cpp @@ -0,0 +1,21 @@ +#include +#include +#include + +#include +#include + +#include "bindings_common.hpp" + +PYBIND11_MODULE(rabitq_ids_benchmark, module) { + module.def("convert", [](const pybind11::object& input, size_t clusters) { + const auto array = pybind11::array::ensure(input); + if (!array || array.ndim() != 1 || array.size() == 0) { + throw std::invalid_argument("Benchmark requires a nonempty 1D array"); + } + const auto result = rabitqlib::python_bindings::checked_cluster_ids( + input, static_cast(array.size()), clusters + ); + return result.back(); + }); +} diff --git a/benchmarks/hotpaths/ids.py b/benchmarks/hotpaths/ids.py new file mode 100644 index 00000000..a0dec69f --- /dev/null +++ b/benchmarks/hotpaths/ids.py @@ -0,0 +1,56 @@ +"""Time the production ID-conversion helper in a benchmark-only extension.""" + +import argparse +import hashlib +import importlib +import sys +import time +from pathlib import Path + +import numpy as np + + +def positive(text): + value = int(text) + if value < 1: + raise argparse.ArgumentTypeError("must be positive") + return value + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--module-dir", type=Path, required=True) + parser.add_argument("--rows", type=positive, default=1_000_000) + parser.add_argument("--iterations", type=positive, default=5) + parser.add_argument("--repetitions", type=positive, default=21) + args = parser.parse_args() + sys.path.insert(0, str(args.module_dir.resolve())) + module = importlib.import_module("rabitq_ids_benchmark") + extension = Path(module.__file__).resolve() + if extension.parent != args.module_dir.resolve(): + raise RuntimeError("Imported benchmark module from an unexpected directory") + print( + f"# extension={extension} sha256={hashlib.sha256(extension.read_bytes()).hexdigest()}" + ) + print(f"# numpy={np.__version__} rows={args.rows} threads=1 clusters=64") + print("operation,repetition,operations,seconds,ns_per_operation,hash") + for dtype in (np.uint32, np.int64): + ids = np.arange(args.rows, dtype=dtype) % 64 + expected = int(ids[-1]) + if module.convert(ids, 64) != expected: + raise RuntimeError("Incorrect converted ID") + for rep in range(args.repetitions): + start = time.perf_counter_ns() + for _ in range(args.iterations): + last = module.convert(ids, 64) + elapsed = (time.perf_counter_ns() - start) / 1e9 + if last != expected: + raise RuntimeError("Incorrect converted ID") + print( + f"ids_{np.dtype(dtype)},{rep},{args.iterations},{elapsed:.12g}," + f"{elapsed * 1e9 / args.iterations:.12g},{last}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/hotpaths/load.cpp b/benchmarks/hotpaths/load.cpp new file mode 100644 index 00000000..93ff78d0 --- /dev/null +++ b/benchmarks/hotpaths/load.cpp @@ -0,0 +1,85 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/path.hpp" + +namespace { +void run(int argc, char** argv) { + if (argc != 6) { + throw std::invalid_argument( + "Usage: rabitq_load_hotpaths prepare|load path dim centroids repetitions" + ); + } + const std::string mode(argv[1]), path(argv[2]); + const size_t dim = bench::number(argv[3]), count = bench::number(argv[4]); + const size_t reps = bench::number(argv[5]); + if (mode != "prepare" && mode != "load") { + throw std::invalid_argument("Mode must be prepare or load"); + } + if (count < 8) + throw std::invalid_argument("At least eight centroids are required"); + const auto data = bench::values(bench::product(count, dim)); + if (mode == "prepare") { + // Repeated smoke tests reuse their fixture; never replace an existing file. + if (!std::filesystem::exists(rabitqlib::io_impl::filesystem_path(path + ".hnsw"))) { + rabitqlib::ivf::HNSWInitializer index(dim, count); + index.add_vectors(data.data(), 1); + std::ofstream unused; + index.save(unused, path.c_str()); + return; + } + } + bench::metadata(dim, count, 1); + uint64_t expected = 0; + for (size_t rep = 0; rep <= reps; ++rep) { + rabitqlib::ivf::HNSWInitializer index(dim, count); + std::ifstream unused; + const auto start = bench::Clock::now(); + index.load(unused, path.c_str()); + const double seconds = + std::chrono::duration(bench::Clock::now() - start).count(); + uint64_t h = 14695981039346656037ULL; + // Validate every stored vector outside timing, including reused fixtures. + for (size_t row = 0; row < count; ++row) { + const float* stored = index.centroid(static_cast(row)); + if (std::memcmp(stored, data.data() + row * dim, dim * sizeof(float)) != 0) { + throw std::runtime_error("Fixture differs from deterministic input"); + } + h = bench::hash(stored, dim * sizeof(float), h); + } + std::vector> candidates(8); + index.centroids_distances(data.data(), candidates.size(), candidates); + for (const auto& candidate : candidates) { + h = bench::hash(&candidate.id, sizeof(candidate.id), h); + h = bench::hash(&candidate.distance, sizeof(candidate.distance), h); + } + if (rep == 0) + expected = h; + else { + if (h != expected) + throw std::runtime_error("Loaded search results changed"); + bench::row("sidecar_load", rep - 1, 1, seconds, h); + } + } +} +} // namespace +int main(int argc, char** argv) { + try { + run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/benchmarks/plot_results.py b/benchmarks/plot_results.py index c0300bbc..3945ee26 100644 --- a/benchmarks/plot_results.py +++ b/benchmarks/plot_results.py @@ -14,10 +14,10 @@ from matplotlib.lines import Line2D if __package__: - from .build_indexes import CASES, DATASETS, LABELS, validate_record + from .build_indexes import CASES, DATASETS, validate_record from .query_index import settings else: - from build_indexes import CASES, DATASETS, LABELS, validate_record + from build_indexes import CASES, DATASETS, validate_record from query_index import settings @@ -55,8 +55,8 @@ def main(): parser.add_argument( "--methods", nargs="+", - choices=[case[0] for case in CASES], - default=[case[0] for case in CASES], + choices=[case.name for case in CASES], + default=[case.name for case in CASES], ) args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) @@ -103,8 +103,7 @@ def main(): for dataset, family in product(args.datasets, ("ivf", "graphs")): if not any( - name in args.methods and ("ivf" if "ivf" in method else "graphs") == family - for name, method, _ in CASES + case.name in args.methods and case.family == family for case in CASES ): continue fig, ax = plt.subplots(figsize=(12, 9)) @@ -123,11 +122,8 @@ def main(): ) handles = [] recalls = [] - for ordinal, ((name, method, _), label) in enumerate( - (case, label) - for case, label in zip(CASES, LABELS, strict=True) - if ("ivf" if "ivf" in case[1] else "graphs") == family - ): + for ordinal, case in enumerate(case for case in CASES if case.family == family): + name = case.name if (dataset, name) not in measurements: continue points = [] @@ -140,7 +136,14 @@ def main(): marker = markers[ordinal % len(markers)] style = "-" if name.startswith("rabitq-") else "--" handles.append( - Line2D([], [], color=color, marker=marker, linestyle=style, label=label) + Line2D( + [], + [], + color=color, + marker=marker, + linestyle=style, + label=case.label, + ) ) if points: x, y = zip(*points) @@ -189,7 +192,8 @@ def main(): plt.close(fig) print(f"Wrote {path}") for dataset in DATASETS: - for name, _, _ in CASES: + for case in CASES: + name = case.name for setting, recall, qps in measurements.get((dataset, name), []): rows.append( ( diff --git a/benchmarks/reproduce.py b/benchmarks/reproduce.py index e97ca23c..8033f911 100644 --- a/benchmarks/reproduce.py +++ b/benchmarks/reproduce.py @@ -16,9 +16,7 @@ def main(): - families = { - name: "ivf" if "ivf" in method else "graphs" for name, method, _ in CASES - } + families = {case.name: case.family for case in CASES} families.update(dict.fromkeys(KMEANS_METHODS, "kmeans")) parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( diff --git a/benchmarks/run.py b/benchmarks/run.py index de98f467..9a8e2034 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -55,7 +55,7 @@ def read_matrix(path, limit=None): matrix = np.memmap(path, dtype="= 0 for value in values): raise ValueError(f"invalid indexing measurements: {path}") - family = "ivf" if "ivf" in method else "graphs" - rows[dataset, name] = (dataset, family, name, *values, 48) + rows[dataset, name] = (dataset, case.family, name, *values, 48) pages = {} for family in ("ivf", "graphs"): - if not any( - ("ivf" if "ivf" in method else "graphs") == family for _, method, _ in cases - ): + if not any(case.family == family for case in cases): continue path = docs_dir / f"{family}.md" text = path.read_text() @@ -72,10 +70,10 @@ def update(indexing_dir, docs_dir, datasets=None, methods=None): "| Method | Train s | Construct s | Total s | Save s | Index GiB |", "| --- | ---: | ---: | ---: | ---: | ---: |", ] - for (name, method, _), label in zip(CASES, LABELS, strict=True): - if ("ivf" if "ivf" in method else "graphs") != family: + for case in CASES: + if case.family != family: continue - row = rows.get((dataset, name)) + row = rows.get((dataset, case.name)) if row is None: cells = ["Pending", *["—"] * 4] else: @@ -84,7 +82,7 @@ def update(indexing_dir, docs_dir, datasets=None, methods=None): float(row[7]) / 2**30, ] cells = [f"{value:.2f}" for value in values] - table.append("| " + label + " | " + " | ".join(cells) + " |") + table.append("| " + case.label + " | " + " | ".join(cells) + " |") text = replace_block(text, f"indexing:{dataset}", "\n".join(table)) pages[path] = text @@ -92,7 +90,9 @@ def update(indexing_dir, docs_dir, datasets=None, methods=None): writer = csv.writer(stream, lineterminator="\n") writer.writerow(FIELDS) writer.writerows( - rows[key] for key in ((d, c[0]) for d in DATASETS for c in CASES) if key in rows + rows[key] + for key in ((d, c.name) for d in DATASETS for c in CASES) + if key in rows ) pages[csv_path] = stream.getvalue() for path, text in pages.items(): @@ -107,7 +107,7 @@ def main(): parser.add_argument("--indexing-dir", type=Path, required=True) parser.add_argument("--docs-dir", type=Path, default=Path("docs/docs/benchmarks")) parser.add_argument("--datasets", nargs="+", choices=DATASETS) - parser.add_argument("--methods", nargs="+", choices=[case[0] for case in CASES]) + parser.add_argument("--methods", nargs="+", choices=[case.name for case in CASES]) args = parser.parse_args() update(args.indexing_dir, args.docs_dir, args.datasets, args.methods) diff --git a/cmake/rabitqlibConfig.cmake.in b/cmake/rabitqlibConfig.cmake.in index 8b9840ab..2d59d4ae 100644 --- a/cmake/rabitqlibConfig.cmake.in +++ b/cmake/rabitqlibConfig.cmake.in @@ -1,9 +1,15 @@ @PACKAGE_INIT@ include(CMakeFindDependencyMacro) -include("${CMAKE_CURRENT_LIST_DIR}/rabitqOpenMP.cmake") -find_dependency(OpenMP REQUIRED) +set(rabitqlib_USE_OPENMP @RABITQ_USE_OPENMP@) +if(rabitqlib_USE_OPENMP) + include("${CMAKE_CURRENT_LIST_DIR}/rabitqOpenMP.cmake") + find_dependency(OpenMP REQUIRED COMPONENTS CXX) +endif() include("${CMAKE_CURRENT_LIST_DIR}/rabitqlibTargets.cmake") +set(rabitqlib_core_FOUND TRUE) +set(rabitqlib_fastscan_FOUND @RABITQ_BUILD_FASTSCAN@) +set(rabitqlib_index_FOUND @RABITQ_BUILD_INDEX@) check_required_components(rabitqlib) diff --git a/docs/docs/benchmarks/bench.md b/docs/docs/benchmarks/bench.md index 71ed0299..e5084078 100644 --- a/docs/docs/benchmarks/bench.md +++ b/docs/docs/benchmarks/bench.md @@ -109,6 +109,10 @@ checks are outside construction timing. Each method has one build or fit per dataset; results do not measure run-to-run variation. +The committed CSVs support checking the displayed tables and plotted points, +but do not include the original raw run logs or measurement binaries. They are +archived results, not an independent reproduction or results from the current +source tree. The experiments in the cited papers are separate measurements. ## Data and reproduction diff --git a/docs/docs/benchmarks/graphs.md b/docs/docs/benchmarks/graphs.md index b164ae46..8a7de520 100644 --- a/docs/docs/benchmarks/graphs.md +++ b/docs/docs/benchmarks/graphs.md @@ -1,5 +1,9 @@ # Graph benchmark +These are historical measurements labeled RaBitQ-Library 0.5.2. See the +[recorded snapshot and reproduction limits](bench.md#package-installation). +Measurements of the modular source tree require a separate run. + Settings and results for eight configurations. See the [overview](bench.md) for hardware, datasets, and measurement definitions. diff --git a/docs/docs/benchmarks/ivf.md b/docs/docs/benchmarks/ivf.md index ebfef203..98bc73be 100644 --- a/docs/docs/benchmarks/ivf.md +++ b/docs/docs/benchmarks/ivf.md @@ -1,5 +1,9 @@ # IVF benchmark +These are historical measurements labeled RaBitQ-Library 0.5.2. See the +[recorded snapshot and reproduction limits](bench.md#package-installation). +Measurements of the modular source tree require a separate run. + Settings and results for eight configurations. See the [overview](bench.md) for hardware, datasets, and measurement definitions. diff --git a/docs/docs/benchmarks/kmeans.md b/docs/docs/benchmarks/kmeans.md index a62923c3..973cba0f 100644 --- a/docs/docs/benchmarks/kmeans.md +++ b/docs/docs/benchmarks/kmeans.md @@ -1,5 +1,9 @@ # K-means benchmark +These are historical measurements labeled RaBitQ-Library 0.5.2. See the +[recorded snapshot and reproduction limits](bench.md#package-installation). +Measurements of the modular source tree require a separate run. + Compare **QGKMeans**, **Faiss K-means**, and **SuperKMeans** using their Python APIs. Metrics are total time, WCSS, and recall@10 with 1% of clusters probed. @@ -9,7 +13,7 @@ Metrics are total time, WCSS, and recall@10 with 1% of clusters probed. | --- | --- | | Threads | 48; methods run sequentially in separate processes | | Clustering | Flat L2, 4,096 clusters | -| Training | All raw float32 base vectors; no sampling, quantization, or hierarchy | +| Training | All float32 base vectors; no input subsampling or prequantization; nonhierarchical training | | Iterations | At most 25; stopping thresholds below | | Final assignment | Exact nearest centroid after the last update | | Seed / fits | 42 / one fit per method and dataset | @@ -25,7 +29,8 @@ search. Centroids are not normalized. | `faiss.Kmeans` | CPU, `nredo=1`, `max_points_per_centroid=N`, `spherical=False`, `early_stop_threshold=0`; exact final `index.search(x, 1)` | | `superkmeans.SuperKMeans` | v0.2.0, `quantizer="f32"`, `hierarchical=False`, `sampling_fraction=1`, `angular=False`, `use_blas_only=False`, `early_termination=True`, `tol=1e-4`; exact final `assign(x, centroids)` | -`N` is the full base count. Methods retain their own training assignment algorithms +`N` is the full base count. These input settings do not disable internal +quantization, such as QGKMeans' one-bit neighbor codes. Methods retain their own training assignment algorithms and initializers; the same seed does not guarantee identical starting centroids. ## Metrics diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md index 719ce814..11ceea03 100644 --- a/docs/docs/clustering.md +++ b/docs/docs/clustering.md @@ -5,8 +5,7 @@ centroid means from the original float32 data, accumulate in float64, and requir no external clustering package. Choose the class explicitly; neither switches algorithms based on the cluster count. -Both classes require RaBitQ Library 0.5.0 or newer; see -[installation options](quick_start.md#install). +For installation, see [Quick start](quick_start.md#install). ## Choose a method @@ -37,6 +36,9 @@ final_obj = kmeans.final_obj ## C++ +Clustering belongs to `rabitqlib::index`, including the flat RaBitQ method. +Link that target; a core-only installation does not include clustering headers. + ```cpp #include #include @@ -77,7 +79,7 @@ Set Python keyword arguments or fields of the corresponding C++ parameter type: | --- | --- | --- | | `niter` | `25` | Maximum clustering iterations. | | `early_stop_threshold` | `0.0` | Stop from iteration 2 when the absolute relative objective change is at most this value; range `[0, 1]`. Zero stops only on an unchanged objective. | -| `num_threads` | `0` | Use detected available logical CPUs; positive requests are capped at that count. See [threading](quick_start.md#threading-and-file-paths). | +| `num_threads` | `0` | Use detected available logical CPUs; positive requests are capped at that count. See [threading](index/usage.md#threading-and-file-paths). | | `spherical` | `false` | Normalize centroids and assign by inner product; input must already be normalized. | | `seed` | `42` | Seed for centroid selection and rotation. | | `min_points_per_centroid` | `39` | Warn below this training-points-per-centroid ratio when verbose. | @@ -166,11 +168,11 @@ possible winners in float64; ties select the smallest centroid index. Returned distances remain float32. Work is proportional to `n * k * d`; large coordinate offsets or near ties can require more refinement. -## Compare with FAISS +## Compare with Faiss The [comparison script](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/python/compare_with_faiss.py) compares both methods with FAISS and requires `faiss-cpu`. Defaults are one -million 256-dimensional vectors, 4,000 centroids, and up to 25 iterations. +million 500-dimensional vectors, 4,000 centroids, and up to 25 iterations. Load data with `--input data.npy`, or `--input data.bin --d D` for raw row-major float32 data. Choose a thread count within the available hardware limit. @@ -180,8 +182,9 @@ python sample/python/compare_with_faiss.py --threads 48 All three methods receive the same full dataset, iteration limit, cluster count, thread budget, and seed values; FAISS subsampling is disabled. This command runs -four fits per method with rotating order. Different initializers mean identical -seeds need not produce identical centroids. +one fit per method with seed 42. Add `--repeats 4` for four fits per method with +rotating order, or `--num-seeds N` to use consecutive seeds. Different initializers +mean identical seeds need not produce identical centroids. Only `train()` is timed, including the library methods' default approximate final assignment. The script reports actual iterations, timing medians/ranges, diff --git a/docs/docs/compact_code.md b/docs/docs/compact_code.md index ce88749d..61e2e698 100644 --- a/docs/docs/compact_code.md +++ b/docs/docs/compact_code.md @@ -1,4 +1,4 @@ -# Compact Storage of Codes +# Compact code storage This section describes packing **extended codes** with 1–8 bits per dimension. IVF and HNSW store a separate sign bit, giving total quantized widths of 1–9 bits; @@ -7,7 +7,7 @@ are multiples of 32. Pass the padded dimension and already padded, quantized codes: the packing function does not add padding itself. For rotated vectors, use the dimension returned by `rotator->size()`; see [Random Rotation](rabitq/rotator.md). The implementation is in -`rabitqlib/quantization/pack_excode.hpp`. +`rabitqlib/core/pack_excode.hpp`. ## How to read the layouts @@ -75,8 +75,6 @@ the detail represents one stored bit; a code can cross a byte boundary. Use `unpacking_rabitqplus_code` with the same dimension and extended bit width to recover the codes; the packed layout is independent of the selected SIMD backend. - - ## 1-bit Each complete block occupies eight bytes. Code `c[i]` is stored in bit @@ -199,7 +197,7 @@ Each code is copied directly: byte `i` is `c[i]`. A complete block occupies ## Packing API example ```cpp -#include +#include #include #include #include diff --git a/docs/docs/cpp_integration.md b/docs/docs/cpp_integration.md new file mode 100644 index 00000000..645c78ea --- /dev/null +++ b/docs/docs/cpp_integration.md @@ -0,0 +1,255 @@ +# C++ integration + +These component targets and the stable core v1 API are introduced in 0.7.0. +First [get the source](installation.md#get-the-source); see the +[migration guide](modular_migration.md). + +## Choose a component + +For a stable, storage-neutral interface, start with [core API v1](rabitq/core_api.md) +and the `sample/cpp/stable_core` example. It accepts caller-owned code buffers and +workspace and keeps factors separate from codes. It does not rotate inputs; +callers provide data, centroids and queries in a shared working coordinate system. +Advanced templates have separate contracts; the v1 guarantees do not cover them. + +Use `rabitqlib::core` for encoding, single-vector estimation and rotation, or +`rabitqlib::fastscan` for batch scoring, without ANN implementations. Eigen remains bundled; +OpenMP is optional for these components and enabled by default. +The C++ libraries are static; `BUILD_SHARED_LIBS` does not change their type. + +| Target | Functionality | +| --- | --- | +| `rabitqlib::core` | Encoding, single-vector query/estimation, rotation, packing and essential SIMD | +| `rabitqlib::fastscan` | Batch layout/packing/encoding, query/LUT/estimation and FastScan/HACC SIMD; depends on core | +| `rabitqlib::index` | IVF, HNSW, SymphonyQG and clustering; depends on FastScan | + +The default C++ build includes all three targets. For numerical components only, +see [core](#core-only-c-integration) or [FastScan](#fastscan-integration). + +## Build the library + +On Linux, build the library and examples from the root of that checkout: + +```bash +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build --parallel +``` + +CMake enables native CPU tuning by default where supported by the compiler. +For portable binaries within a supported OS and architecture, configure with +`-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`; runtime kernel selection remains active. +See the [platform-specific build commands](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/tests/README.md#quick-start) +for Linux ARM64, Windows, and Apple Silicon. + +## C++ examples + +ANN example executables are written to `bin/`. Their source demonstrates complete +indexing and querying workflows: + +- [IVF RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/ivf_rabitq_indexing.cpp) +- [HNSW RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/hnsw_rabitq_indexing.cpp) +- [SymphonyQG](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/symqg_indexing.cpp) +- [Core quantization and distance estimation](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/core/main.cpp) + +The core example is built as the `core_example` CMake target. For the GIST benchmark workflow, see +[`example.sh`](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/example.sh). + +## Use in another C++ project + +The [core quantization v1 API](rabitq/core_api.md) has a C++ source-compatibility contract. +Other C++ APIs may evolve, and no cross-version C++ binary ABI is promised. +Pin the `v0.7.0` release and include RaBitQ-Library as a Git submodule: + +```bash +git submodule add https://github.com/VectorDB-NTU/RaBitQ-Library.git third_party/rabitqlib +git submodule update --init --recursive +git -C third_party/rabitqlib checkout v0.7.0 +git add third_party/rabitqlib +``` + +Add the library and link its namespaced target in your `CMakeLists.txt`: + +```cmake +set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) +add_subdirectory(third_party/rabitqlib) +target_link_libraries(my_program PRIVATE rabitqlib::index) +``` + +Update the pinned revision deliberately when adopting upstream changes: + +```bash +git -C third_party/rabitqlib fetch +git -C third_party/rabitqlib checkout +git add third_party/rabitqlib +``` + +## Install the C++ library + +Installation is useful for package managers, container images, and shared server +environments. Disable native optimization for use on other CPUs: + +```bash +cmake -S . -B build \ + -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX="$HOME/.local" +cmake --build build --parallel +cmake --install build +``` + +Consume the installed package with: + +```cmake +find_package(rabitqlib CONFIG REQUIRED) +target_link_libraries(my_program PRIVATE rabitqlib::index) +``` + +For a non-system prefix, point CMake to the installation: + +```bash +cmake -S . -B build -DCMAKE_PREFIX_PATH="$HOME/.local" +cmake --build build --parallel +``` + +Full-library submodule and installed-package integration require OpenMP on the consuming +system. The [downstream consumer test](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/tests/consumer) +provides a complete installed-package example. + +## Core-only C++ integration + +To build and install only core: + +```bash +cmake -S . -B build-core -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_INDEX=OFF -DRABITQ_BUILD_FASTSCAN=OFF \ + -DRABITQ_USE_OPENMP=OFF \ + -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DCMAKE_INSTALL_PREFIX="$HOME/.local/rabitqlib-core" +cmake --build build-core --parallel +cmake --install build-core +``` + +This configuration does not discover or link OpenMP. Legacy parallel batch helpers run +serially, including when a caller requests more workers; v1 numerical calls remain serial. +Enable FastScan with `RABITQ_BUILD_FASTSCAN=ON` if needed. ANN and Python builds currently +require `RABITQ_USE_OPENMP=ON`. Use one consistent OpenMP choice for all RaBitQ headers +and linked targets within an application. + +Use a fresh installation prefix when changing between full and core-only installations. +Pure core mode omits FastScan/index targets and their examples/tests; enabling Python bindings +in that mode is a configuration error. Core tests are available with `RABITQ_BUILD_TESTS=ON`. + +### Headers + +Use canonical module headers in new code: + +| Headers | Owning module | +| --- | --- | +| `core/encoder.hpp`, `core/data_layout.hpp`, `core/pack_excode.hpp` | Per-vector encoding and layouts | +| `fastscan/encoder.hpp`, `fastscan/data_layout.hpp`, `fastscan/pack_codes.hpp` | All 32-vector batch layouts, adapters and packing | +| `core/query.hpp`, `core/estimator.hpp`, `core/rotator.hpp` | Core single-vector query, estimation and rotation | +| `fastscan/query.hpp`, `fastscan/estimator.hpp`, `fastscan/fastscan.hpp` | Batch scoring | + +Paths are relative to `rabitqlib/`. Old forwarding headers and the full-library alias +are removed in 0.7.0; follow the [migration guide](modular_migration.md). +Installations contain only enabled module headers and their dependencies. In particular, +core-only installations do not carry FastScan, ANN, clustering or vendored hnswlib headers. + +### Link core + +Installed package: + +```cmake +find_package(rabitqlib CONFIG REQUIRED COMPONENTS core) +target_link_libraries(my_program PRIVATE rabitqlib::core) +``` + +This also works against a full installation without linking the index archive. For ANN +use, request `COMPONENTS index` and link `rabitqlib::index`. +For changes to older targets and headers, see the [migration guide](modular_migration.md). + +Source-tree integration: + +```cmake +set(RABITQ_BUILD_INDEX OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_FASTSCAN OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) +set(RABITQ_ENABLE_NATIVE_OPTIMIZATION OFF CACHE BOOL "" FORCE) +add_subdirectory(third_party/rabitqlib) +target_link_libraries(my_program PRIVATE rabitqlib::core) +``` + +For FetchContent, the same options apply. This example consumes a local checkout, +which is also the configuration exercised by the integration tests: + +```cmake +include(FetchContent) +set(RABITQ_BUILD_INDEX OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_FASTSCAN OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) +set(RABITQ_ENABLE_NATIVE_OPTIMIZATION OFF CACHE BOOL "" FORCE) +FetchContent_Declare(rabitqlib SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/third_party/rabitqlib") +FetchContent_MakeAvailable(rabitqlib) +target_link_libraries(my_program PRIVATE rabitqlib::core) +``` + +To fetch remotely, replace `SOURCE_DIR` with `GIT_REPOSITORY` and a pinned `GIT_TAG`. +The local test does not validate network access or any package-manager integration. + +### Run the core example + +The standalone example in `sample/cpp/core` +rotates a 65-dimensional vector and query using the same seeded rotator, allocates +buffers using the existing data maps, encodes four-bit RaBitQ and estimates squared L2 +distance relative to a zero centroid. Query state borrows its rotated input; keep that +input alive through estimation. All internal buffers use the rotator's padded dimension. +No ANN index is constructed. + +```bash +cmake -S sample/cpp/core -B build-core-example \ + -DCMAKE_PREFIX_PATH="$HOME/.local/rabitqlib-core" +cmake --build build-core-example +ctest --test-dir build-core-example --output-on-failure +``` + +With project samples enabled, the same program is built as `core_example` under +`build-core/sample/cpp/` (under `Release/` there for multi-configuration generators). + +## FastScan integration + +For batch scoring without ANN, configure `RABITQ_BUILD_INDEX=OFF` and +`RABITQ_BUILD_FASTSCAN=ON`, then consume the installed package with: + +```cmake +find_package(rabitqlib CONFIG REQUIRED COMPONENTS fastscan) +target_link_libraries(my_program PRIVATE rabitqlib::fastscan) +``` + +The standalone FastScan example in `sample/cpp/fastscan` of the same source checkout +encodes 49 vectors in two batches and scores them with HACC and extra-bit refinement, +for both L2 and IP with a nonzero centroid. It reuses storage for 32 physical lanes +and consumes only the 17 logical results of the final batch. Build it with the same +commands as the core example, substituting `sample/cpp/fastscan` and its installation +prefix. The in-tree sample target is `fastscan_example`. + +Index builds require FastScan; requesting `RABITQ_BUILD_INDEX=ON` with FastScan +disabled fails configuration. + +## Python bindings + +Python bindings are optional (`RABITQ_BUILD_PYTHON_BINDINGS=ON`) and require indexes. The +existing `_rabitqlib` extension has a `rabitqlib::python` build-tree alias; build the +`_rabitqlib` target or install the wheel and use `import rabitqlib`. It is a Python +MODULE, not a C++ link library, and is not exported as a C++ `find_package` component. + +## Run the C++ tests on Linux + +```bash +cmake -S . -B build -DRABITQ_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Release -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF +cmake --build build --parallel +ctest --test-dir build --output-on-failure +``` + +GoogleTest is downloaded during test configuration. diff --git a/docs/docs/index.md b/docs/docs/index.md index b023f1c9..77a8831c 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -1,17 +1,26 @@ -# Compact vector search with RaBitQ +--- +title: RaBitQ Library +hide: + - toc +---
-**RaBitQ Library** is a C++17 library with Python bindings for compact, -accurate vector quantization and approximate nearest-neighbor search. +

RaBitQ Library · C++17 & Python

-Build with the low-level quantizer or use complete IVF, HNSW, and SymphonyQG -indexes on Linux x86-64 and ARM64, Windows x86-64, or macOS ARM64 -(Apple Silicon). +# Compact vectors.
Accurate search. { #compact-vector-search-with-rabitq } + +

Vector quantization and approximate nearest-neighbor +search. Start with a quantizer, or build with complete IVF, HNSW, and SymphonyQG indexes.

+ +
[Get started](quick_start.md){ .md-button .md-button--primary } -[View on GitHub](https://github.com/VectorDB-NTU/RaBitQ-Library){ .md-button } -[Install from PyPI](https://pypi.org/project/rabitqlib/){ .md-button } +[Explore C++ integration →](cpp_integration.md) + +
+ +

Linux x86-64 & ARM64 · Windows x86-64 · macOS ARM64

@@ -34,60 +43,37 @@ indexes on Linux x86-64 and ARM64, Windows x86-64, or macOS ARM64 -## New in 0.6.0 +## New in 0.7.0 -- **Concurrent Python queries:** Native searches release the GIL and use independent - parameters; conflicting updates raise an exception. -- **IVF centroid routing:** Choose Flat, Flat RaBitQ, HNSW, or automatic selection. -- **32-coordinate padding:** New indexes pad to multiples of 32; loading preserves - the saved rotation and padding. -- **Build integration:** Eigen 5.0.1; CMake optimization flags no longer propagate - to downstream targets. +- **Modular C++ integration:** Link core, FastScan, or indexes independently. +- **Stable core API v1:** Encode and estimate distances with caller-owned buffers + and workspace, with no heap allocation on valid calls. +- **Optional OpenMP:** Build core and FastScan without an OpenMP dependency. +- **Unified index counts:** Distinguish capacity, stored records, and live points + consistently in C++ and Python. -Read the [0.6.0 upgrade notes](releases/0.6.0.md) before rebuilding indexes, -sharing files with older installations, or updating a downstream C++ library. +This documentation describes 0.7.0. Read the +[0.7.0 upgrade notes](releases/0.7.0.md) before updating an existing application. ## Start with Python See [installation options](quick_start.md#install) for supported platforms and source builds. ```bash -python -m pip install "rabitqlib>=0.6.0" +python -m pip install "rabitqlib>=0.7.0" ``` -Build an IVF index and search a batch of queries: - -```python -import numpy as np -from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans +The [complete quick start](quick_start.md) builds and searches an IVF index with +synthetic data, explains the results, and demonstrates save/load. No external +dataset or clustering package is needed. -rng = np.random.default_rng(42) -data = rng.standard_normal((500, 64)).astype(np.float32) -queries = rng.standard_normal((5, 64)).astype(np.float32) +[Start the Python example](quick_start.md#build-and-search-an-ivf-index){ .md-button .md-button--primary } -clustering = RaBitQKMeans( - 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact -) -clustering.train(data) +## C++ components -index = IvfIndex( - dim=64, - max_elements=len(data), - num_clusters=5, - nbits=4, - metric="l2", -) -index.build(data, clustering.centroids, clustering.assignments) - -ids, distances = index.search(queries, k=10, nprobe=5) -print(ids.shape, distances.shape) # (5, 10) (5, 10) -``` - -Starting with 0.5.1, IVF also supports [adding and removing vectors](index/ivf.md#updating-an-index) -without rebuilding the index or retaining the original dataset. -Starting with 0.5.2, HNSW supports [adding and removing vectors and resizing capacity](index/hnsw.md#updating-an-index). - -[Continue to the complete quick start](quick_start.md){ .md-button .md-button--primary } +The [0.7.0 component model](modular_migration.md) provides independent core and +FastScan targets alongside the index library, plus a [stable core v1 API](rabitq/core_api.md). +Use a matching source checkout or 0.7.0 installation for these interfaces. ## Choose an index @@ -105,26 +91,18 @@ optimize its access pattern. These are typical relative profiles, not fixed guarantees. Actual memory, latency, and recall depend on vector dimension, quantization width, graph degree, and search parameters. +See [memory estimates](index/usage.md#estimate-storage) for the main storage costs, +including SymphonyQG's copies of neighbor codes. ## Why RaBitQ? -- **High accuracy with tiny codes.** RaBitQ provides strong similarity - estimates across different bit widths and remains effective with a one-bit - code per padded dimension plus per-vector factors. -- **Fast distance estimation.** IVF and SymphonyQG use - [FastScan](https://arxiv.org/abs/1704.07355) for batched estimates; HNSW - uses single-code SIMD kernels. -- **Theoretical error bounds.** An asymptotically optimal error bound supports - reliable ordering and reranking. -- **Multiple integration points.** Use the quantizer directly or select a - complete vector-search index. - -The library supports Euclidean distance and inner product. Cosine similarity -can be implemented by normalizing vectors and using inner product. -It implements the [1-bit RaBitQ](https://arxiv.org/abs/2405.12497) and -[multi-bit RaBitQ](https://doi.org/10.1145/3725413) research from the -[VectorDB Group](https://vectordb-ntu.github.io/) at Nanyang Technological -University. +RaBitQ combines compact codes with distance estimates and probabilistic error +intervals. The [algorithm overview](rabitq/rabitq.md) explains the rotation, +encoding and scoring workflow; the [benchmarks](benchmarks/bench.md) report +measured memory, recall and throughput for specific configurations. + +The library supports squared L2 and inner-product distance (`1 - dot`). For cosine +search, normalize data and queries before using inner product. ## Used across the vector-search ecosystem diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index a07bef15..cbf7e36e 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -1,12 +1,74 @@ # HNSW RaBitQ + +C++ callers include `rabitqlib/index/hnsw/hnsw.hpp` and link `rabitqlib::index`. [HNSW](https://arxiv.org/abs/1603.09320) is a popular graph-based index. Under comparable quantization settings, HNSW RaBitQ typically uses more memory than IVF RaBitQ because it stores graph structures in addition to quantized vectors. The actual comparison depends on vector dimension, bit width, graph degree, and index parameters. -This document describes how the library integrates HNSW with RaBitQ to support efficient vector search. -## Index Construction +## Minimal workflow + +### Python + +Run the [Quick start](../quick_start.md#build-and-search-an-ivf-index) to prepare +`data`, `queries` and `clustering`, then use those same arrays with HNSW: + +```python +from rabitqlib import HnswIndex + +index = HnswIndex( + dim=data.shape[1], max_elements=len(data), M=16, ef_construction=100, nbits=4, +) +index.build(data, clustering.centroids, clustering.assignments) +ids, distances = index.search(queries, k=10, ef=100) +index.save("hnsw.index") +loaded = HnswIndex.load("hnsw.index") +``` + +### C++ + +Link `rabitqlib::index`. Here `data`, `centroids` and `queries` are row-major +`std::vector` buffers with original dimension `dim`; `cluster_ids` is a +`std::vector` with one centroid ID per data row. Supply at least +ten points for this example. + +```cpp +#include +#include + +// Constructor order: capacity, dimension, total bits, M, ef_construction. +rabitqlib::hnsw::HierarchicalNSW index(data.size() / dim, dim, 4, 16, 100); +index.construct( + centroids.size() / dim, centroids.data(), data.size() / dim, data.data(), + cluster_ids.data(), 1 +); +auto results = index.search(queries.data(), queries.size() / dim, 10, 100, 1); +index.save("hnsw.index"); +rabitqlib::hnsw::HierarchicalNSW loaded; +loaded.load("hnsw.index"); +``` + +C++ results are per-query vectors of `(distance, ID)` pairs. For build setup, +see [C++ integration](../cpp_integration.md). + +## Parameters and limits + +| Parameter | Contract | +| --- | --- | +| `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | +| `max_elements` | Capacity; build accepts 1 through this many rows. Reserve extra capacity or call `resize` before adding beyond it. | +| `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit; Python defaults to 8. Raw-vector storage is not supported. | +| `M`, `ef_construction` | Graph connectivity and construction search budget; Python defaults are 16 and 200. | +| `metric`, `random_seed` | Python defaults are `"l2"` and 100; `"ip"` is also supported. | +| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, and positive `ef` is raised to at least `k`. | + +Results use squared L2 or `1 - dot` estimates, ordered from smallest to largest. +Python fills missing results with ID `2**32 - 1` and distance `+inf`; C++ returns +only the available pairs. See [result handling](../quick_start.md#read-the-results) +and the shared [input, count and concurrency contracts](usage.md). + +## Index construction The graph is built by incrementally inserting raw data vectors. Raw vectors are used to construct and prune graph links, but the completed index stores @@ -14,7 +76,10 @@ the links, centroids, cluster IDs, labels, RaBitQ codes, and factors rather than a copy of the raw dataset. Total bit widths from 1 through 9 are supported. -Users can invoke: +`construct()` is synchronous. Keep its input buffers valid for the entire call; +the index retains no pointers to them after it returns or throws. The borrowed +input pointer and search candidate helpers are private implementation details; +see the [C++ migration guide](../modular_migration.md#public-and-implementation-interfaces). ```cpp HierarchicalNSW::construct(size_t cluster_num, @@ -26,35 +91,15 @@ HierarchicalNSW::construct(size_t cluster_num, bool faster = false); ``` -- **data**: Pointer to the raw data vectors. +- **data**: Vectors in the original input coordinates. - **data_num**: The number of data vectors. - **centroids**: Centroids computed by [RaBitQKMeans](../clustering.md#rabitqkmeans); the examples use `cluster_num = 16`. - **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`. - **num_threads**: Number of threads to use (default: 0, which auto-selects). -- **faster**: If `true`, enables the faster quantizer. - - -During construction, we first rotate the centroids and then insert each element one by one. For each element: - -1. Update the graph structure (edges) by searching with raw vectors and pruning. -2. Quantize the rotated vector and store its quantization code. - - -### Data Layout - -Each indexed element is stored in the following layout: - -``` -[number of edges] -[edges] -[cluster ID] (its high bit marks a removed point, see Removing points) -[external label] -[BinData (1-bit * dim + factors)] -[ExData (ex-bits * dim + factors)] -``` +- **faster**: Use calibrated scaling instead of per-vector scale search; see [encoding paths](../rabitq/quantizer.md#implementations). ## Querying -Users can invoke: + ```cpp std::vector>> HierarchicalNSW::search(const float* queries, size_t query_num, @@ -63,7 +108,7 @@ std::vector>> HierarchicalNSW::search(const fl size_t thread_num); ``` -- **queries**: Pointer to the raw query vectors. +- **queries**: Query vectors in the original input coordinates. - **query_num**: The number of query vectors. - **TOPK**: The number of nearest neighbors to search. - **efSearch**: The size of the candidate set for searching HNSW base layer. @@ -71,48 +116,11 @@ std::vector>> HierarchicalNSW::search(const fl The search window belongs to each call, so concurrent C++ searches can use different `efSearch` values and separate outputs while the index is unchanged. -Python search releases the GIL. Conflicting updates raise `RuntimeError` -immediately; see the [Python concurrency contract](../quick_start.md#threading-and-file-paths). - -We first pre-process the query: - -1. Rotate the raw query vector. -2. Compute distances between the rotated query and all rotated centroids. -3. Encapsulate the query into a `query_wrapper` for subsequent search. - -### Upper Layers - -In the upper layers of HNSW, we compute the 1-bit estimated distance (using `BinData`) to quickly locate the entry point for the next layer. - -### Base Layer - -In the base layer, we apply an adaptive re-ranking strategy: - -- **candidate_set**: Elements to be visited. -- **boundedKNN**: Current best TOPK candidates. - -Repeat until `candidate_set` is empty: - -1. Extract the nearest element `e` from `candidate_set`. -2. Visit all unvisited neighbors of `e`. -3. For each neighbor: - - Compute the 1-bit lower-bound distance (along with 1-bit estimated distance) using `BinData`. - - If `boundedKNN` has fewer than TOPK elements, or if the 1-bit lower-bound is smaller than the full-bits estimated distance of the current farthest element in `boundedKNN`: - 1. Refine the distance estimate using `ExData` to obtain the full-bits estimated distance. - 2. Update `boundedKNN`. - - Insert the neighbor into `candidate_set` with its (possibly refined) estimated distance. +Python follows the shared [concurrency contract](usage.md#threading-and-file-paths). -The search terminates when `candidate_set` is empty. +## Updating an index -A point hidden by `remove` never enters `boundedKNN`, so search never returns it, -but it still enters `candidate_set` and the search keeps walking through it. A -removed entry point keeps the estimate already computed for it (the full-bits one -when nbits > 1). A removed neighbor is not refined and enters `candidate_set` with -its 1-bit estimate. - -## Updating an Index - -Starting with 0.5.2, a constructed or loaded index supports `add()`, `resize()`, +A constructed or loaded index supports `add()`, `resize()`, and `remove()` without the original data. Adding points and resizing preserve the index file format; files with removals have [compatibility restrictions](#how-removal-is-stored). @@ -130,66 +138,54 @@ void HierarchicalNSW::resize(size_t new_max_elements); ``` - **data**, **n**: `n` new vectors in the original coordinates. Vector `i` receives - the label `num_points() + i`, so labels stay dense, and `add` returns them. + the label `num_stored_points() + i`, so labels stay dense, and `add` returns them. - **cluster_ids**: The cluster of each new vector, in `[0, num_clusters)`. When it is `nullptr`, each vector goes to its nearest centroid, chosen the same way a query is routed. - **faster**, **num_threads**: Same meaning as in `construct`, but `add` defaults to one thread; pass `0` to auto-select. A point keeps the label - `num_points() + i` whatever order the threads finish in, because the slots are + `num_stored_points() + i` whatever order the threads finish in, because the slots are reserved as one block before the inserts start. -- **new_max_elements**: The new capacity. It must be at least `num_points()`. +- **new_max_elements**: The new capacity. It must be at least `num_stored_points()`. In C++, initialize the index through its parameterized constructor or `load` before calling `resize`. In Python: ```python +# Grow before adding; preserve any capacity already available. +index.resize(max(index.max_elements, index.num_stored_points + len(vectors))) new_ids = index.add(vectors) # route to the nearest centroids # Alternatively, choose the clusters instead of the call above: # new_ids = index.add(vectors, cluster_ids=labels) -index.resize(index.max_elements + 100_000) ``` -In C++, callers must prevent `add` or `resize` from overlapping with searches -on the same index. Python detects these conflicts and raises `RuntimeError` -immediately. +Updates require exclusive access; see the +[concurrency contract](usage.md#threading-and-file-paths). `add` prepares the result IDs, quantized vectors, and link-list allocations before changing the graph or point count. A failure during preparation leaves those unchanged. A failure during graph linking can leave the batch counted in -`num_points()` with only some points reachable through the graph. The index -remains safe to search, save, load, and destroy; inspect `num_points()` before +`num_stored_points()` with only some points reachable through the graph. The index +remains safe to search, save, load, and destroy; inspect `num_stored_points()` before retrying, or rebuild to restore full reachability. ### Capacity -`add` never grows the index: it throws once `num_points()` reaches -`max_elements()`, and `resize` is the only way to raise that. An index built with -spare capacity absorbs points with no reallocation at all, so sizing -`max_elements` ahead of time is the cheap path. +See the shared [index count contract](usage.md#index-counts). Removed +records still occupy capacity. + +`add` never grows the index: it throws once `num_stored_points()` reaches +`max_elements()`, and `resize` is the only way to raise that. Reserving capacity +avoids reallocating and copying the whole base-layer record storage during `add`. +It does not make insertion allocation-free: batch scratch, upper-layer link lists, +and label-map entries still require allocations, and the label map may rehash. `resize` rebuilds the base-layer storage and copies it, so it needs memory for both copies while it runs and costs time proportional to the whole index. It invalidates every pointer into the index. The per-node upper-layer link lists are not copied, only their pointer table. -### How added points are linked - -`construct` borrows the caller's vectors and links points using exact distances -between them. That pointer is not retained, so `add` cannot use it, and instead -scores candidates with the estimator the search already uses: the vector being -compared against is rebuilt from its stored RaBitQ code, sign bit and extra bits -together, and the other side is scored straight from its codes. This is the same -approach quantized SymphonyQG construction takes. - -Graph quality therefore depends on how a point arrived. - -Save/load preserves the existing graph and rotation, but not the random generator -state used to choose levels for future insertions. Adding the same vectors to an -original index and a loaded copy can therefore produce different graphs and -search results. - ### Recall after `add` The centroids and the rotation are fixed when the index is constructed, so `add` @@ -212,7 +208,7 @@ size_t HierarchicalNSW::remove(const PID* labels, size_t n); In Python: ```python -removed = index.remove(ids) # ids in [0, num_points) +removed = index.remove(ids) # ids in [0, num_stored_points) ``` A removed point keeps its codes, its links, and its place in the graph. Search @@ -220,24 +216,27 @@ still walks through it and `add` may link new points to it, which keeps the grap connected, as in hnswlib. Search never returns it, so a query can get fewer than `k` results once points are removed: in Python the missing slots hold `kPidMax` (`2**32 - 1`) with an infinite distance. Removed points still count in -`num_points()` and cannot be restored. In C++, callers must prevent `remove` -from overlapping with searches or updates on the same index. Python detects -these conflicts and raises `RuntimeError` immediately. +`num_stored_points()` and cannot be restored. Removal requires the same +exclusive access as other updates. #### Choosing `ef` after removals Removed points still take slots in the `efSearch` candidate set, so a search with -`efSearch` close to `k` finds fewer live points as removals accumulate. `search` -uses the `efSearch` it is given. To keep about the same number of live candidates -as before the removals, scale it by the share of live points: +`efSearch` close to `k` finds fewer live points as removals accumulate. C++ search +uses `max(efSearch, k)`. Python uses `max(k, 10)` when `ef` is omitted or zero; +an explicit positive `ef` is also raised to at least `k`. +To keep about the same number of live candidates as before the removals, scale +the effective window by the share of live points: ```text -efSearch * num_points() / (num_points() - removed points) +max(efSearch, k) * num_stored_points() / (num_stored_points() - removed points) ``` Callers that track how many points they removed can apply this directly. -#### How removal is stored +## Persistence + +### How removal is stored `remove` sets the high bit of the point's stored cluster ID. `construct` and `load` limit the number of clusters to 2^31, so no real cluster ID has that bit, and @@ -254,3 +253,87 @@ New HNSW indexes pad dimensions to multiples of 32 for all supported bit widths. Loading preserves the saved padded dimension and rotation, so older indexes using multiples of 64 remain readable. Older library versions cannot load new files whose padded dimension is not a multiple of 64. + +## Implementation details + +### Construction + +During construction, we first rotate the centroids and then insert each element one by one. For each element: + +1. Update the graph structure (edges) by searching with raw vectors and pruning. +2. Quantize the rotated vector and store its quantization code. + +### Data layout + +Each indexed element is stored in the following layout: + +``` +[number of edges] +[edges] +[cluster ID] (its high bit marks a removed point, see Removing points) +[external label] +[BinData (1-bit * dim + factors)] +[ExData (ex-bits * dim + factors)] +``` + +### Search algorithm + +We first pre-process the query: + +1. Rotate the raw query vector. +2. Compute distances between the rotated query and all rotated centroids. +3. Encapsulate the query into `SplitSingleQuery` for subsequent search. + +#### Upper Layers + +In the upper layers of HNSW, we compute the 1-bit estimated distance (using `BinData`) to quickly locate the entry point for the next layer. + +#### Base Layer + +In the base layer, we apply an adaptive re-ranking strategy: + +- **candidate_set**: A distance-sorted exploration window with capacity + `max(efSearch, TOPK)`, containing both expanded and unexpanded candidates. +- **boundedKNN**: Current best TOPK candidates. + +Repeat while the exploration window has an unexpanded candidate: + +1. Select the nearest unexpanded element `e` and mark it expanded. It remains + in the window unless displaced by a later insertion. +2. Visit all unvisited neighbors of `e`. +3. For each neighbor: + - Compute the 1-bit lower-bound distance (along with 1-bit estimated distance) using `BinData`. + - For a nonremoved neighbor, if `boundedKNN` has fewer than TOPK elements + or the neighbor's lower bound is smaller than the worst result's estimated distance: + 1. Refine the distance estimate using `ExData` when extra bits are present. + 2. Update `boundedKNN`. + - Insert the neighbor into `candidate_set` only if the window has space or + its (possibly refined) estimated distance is no greater than the worst + distance in the window. Inserting into a full window displaces its last entry. + +The search stops when no unexpanded candidates remain in the window, even though +the window still contains entries. `efSearch` controls this exploration budget +separately from the TOPK result count; increasing it allows a wider search at +additional cost, rather than requiring traversal of the entire graph. + +A point hidden by `remove` never enters `boundedKNN`, so search never returns it, +but it can still enter `candidate_set` under the same window admission rule, +allowing search to walk through it. A removed entry point keeps the estimate +already computed for it (the full-bits one when nbits > 1). A removed neighbor +is not refined and uses its 1-bit estimate for window admission. + +### How added points are linked + +`construct` borrows the caller's vectors and links points using exact distances +between them. That pointer is not retained, so `add` cannot use it, and instead +scores candidates with the estimator the search already uses: the vector being +compared against is rebuilt from its stored RaBitQ code, sign bit and extra bits +together, and the other side is scored straight from its codes. This is the same +approach quantized SymphonyQG construction takes. + +Graph quality therefore depends on how a point arrived. + +Save/load preserves the existing graph and rotation, but not the random generator +state used to choose levels for future insertions. Adding the same vectors to an +original index and a loaded copy can therefore produce different graphs and +search results. diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index bdaed502..2c1a614b 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -1,5 +1,7 @@ # IVF RaBitQ +C++ callers include `rabitqlib/index/ivf/ivf.hpp` and link `rabitqlib::index`. + [IVF](https://dl.acm.org/doi/10.1109/TPAMI.2010.57) is a classical clustering-based ANN method. This implementation stores RaBitQ codes and vector IDs, optionally retaining raw vectors for reranking, and uses @@ -7,9 +9,59 @@ vector IDs, optionally retaining raw vectors for reranking, and uses batch of distances. Its actual memory, latency, and recall depend on the bit width, number of clusters, and `nprobe`. -The algorithm includes two phases: indexing and querying. +## Minimal workflow + +### Python + +The [complete Quick Start](../quick_start.md) is the runnable IVF example: +train centroids, call `IvfIndex.build`, search, then save and reload the index. +The update examples below reuse its built `index`. + +### C++ + +Link `rabitqlib::index`. In this snippet, `data`, `centroids` and `queries` are +row-major `std::vector` buffers with original dimension `dim`; +`cluster_ids` is a `std::vector` with one centroid ID per data row. +The example requests ten neighbors from five clusters, so supply at least ten +points and five centroids. + +```cpp +#include +#include + +rabitqlib::ivf::IVF index(data.size() / dim, dim, centroids.size() / dim, 4); +index.construct(data.data(), centroids.data(), cluster_ids.data(), false, 1); +const size_t num_queries = queries.size() / dim; +std::vector ids(num_queries * 10); +std::vector distances(num_queries * 10); +index.search_batch(queries.data(), num_queries, 10, 5, ids.data(), distances.data()); +index.save("ivf.index"); +``` + +See [C++ integration](../cpp_integration.md) for build setup and +[Index Construction](#index-construction) for loading `.fvecs` inputs. -## Index Construction +## Parameters and limits + +| Parameter | Contract | +| --- | --- | +| `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | +| `max_elements` | Python IVF build requires exactly this many rows; `add` grows storage automatically. C++ calls this constructor argument `num_points`. | +| `num_clusters` | Must match the centroid count; every cluster ID must be in `[0, num_clusters)`. | +| `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit, or 32 for owned raw-vector reranking. | +| `metric` | Python `"l2"` or `"ip"`; defaults to `"l2"`. | +| `initializer` | Defaults to `"auto"`; see [centroid routing](#centroid-routing). | +| `k`, `nprobe` | Positive result count and cluster probe count; Python requires `k <= max_elements`. | +| `high_accuracy` | Python defaults to `None`: HACC for 4–9 bits, standard FastScan otherwise. | + +Results contain IDs and squared L2 or `1 - dot` distances, ordered from smallest +to largest. Quantized distances are estimates. Unfilled slots use ID `2**32 - 1` +and distance `+inf`; see [result handling](../quick_start.md#read-the-results). +See the shared [input and concurrency contracts](usage.md). + +## Index construction + +See the shared [index count contract](usage.md#index-counts). Use [QGKMeans](../clustering.md#qgkmeans) to partition the input vectors (`.fvecs`). Build the C++ examples, then run `bin/qgkmeans` to save centroids and cluster IDs. @@ -25,10 +77,19 @@ For example, split SIFT into 4,096 clusters using squared L2 distance: /data/sift/sift_clusterids_4096_l2.ivecs ``` -After files are prepared, you need to load them into memory: +After files are prepared, load them into memory. The snippet below uses +`const char*` paths `data_file`, `centroids_file`, and `cids_file`. Include the +Eigen aliases and I/O helpers explicitly; `ivf.hpp` does not import them. + ```c++ +#include +#include +#include +#include + +using PID = rabitqlib::PID; using data_type = rabitqlib::RowMajorArray; -using gt_type = rabitqlib::RowMajorArray; +using gt_type = rabitqlib::RowMajorArray; data_type data; data_type centroids; @@ -39,11 +100,13 @@ rabitqlib::load_vecs(centroids_file, centroids); rabitqlib::load_vecs(cids_file, cids); ``` +See the [complete indexing example](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/ivf_rabitq_indexing.cpp) +for argument parsing and the surrounding program. + Then initialize an IVF object with the number of data points, vector dimension, number of clusters, and total bits per dimension. The C++ index accepts total quantized bit widths from 1 through 9, or `total_bits=32` -(Python `nbits=32`) to store raw float32 vectors for reranking. Raw storage -and automatic FastScan precision selection are available starting with version 0.3.3. +(Python `nbits=32`) to store raw float32 vectors for reranking. ```c++ using index_type = rabitqlib::ivf::IVF; @@ -57,8 +120,6 @@ index_type ivf(num_points, dim, k, total_bits); ### Centroid routing -Configurable centroid routing requires 0.6.0 or newer. - The optional final C++ constructor argument selects `InitializerType::Auto`, `Flat`, `FlatRaBitQ`, or `HNSW` from the `rabitqlib::ivf` namespace. Python accepts `initializer="auto"`, `"flat"`, `"flat_rabitq"`, or `"hnsw"`: @@ -95,12 +156,14 @@ buffer for fewer full-vector distance evaluations. Performance depends on cluster count, dimension, `nprobe`, and data distribution. Measure recall and latency on your workload, and use an explicit initializer to override `auto`. -Existing constructor calls remain source-compatible. C++ consumers must rebuild -against matching headers and the core library because the IVF object layout has -changed. Newly built indexes with 5,000–59,999 clusters now use Flat RaBitQ routing. In particular, the 5,000–19,999 range changes -from exact to approximate centroid selection. To reproduce the previous selection -policy when rebuilding, explicitly choose `flat` below 20,000 clusters and `hnsw` -otherwise. Loading an old index preserves its historical routing instead. +The optional routing argument introduced in 0.6.0 preserves existing constructor +call forms; this does not imply compatibility with all 0.7.0 C++ changes. See the +[migration guide](../modular_migration.md) and rebuild with matching headers and libraries. + +Compared with releases before 0.6.0, `auto` changes routing for 5,000–59,999 +clusters. The 5,000–19,999 range changes from exact to approximate selection. +To retain the earlier policy when rebuilding, choose `flat` below 20,000 clusters +and `hnsw` otherwise. Loading an old index preserves its historical routing. Routing is fixed for an index and is also used by `add` when cluster IDs are omitted. `initializer_type()` (Python `initializer`) reports the resolved type, @@ -108,6 +171,7 @@ including after loading; loading never re-runs the current `auto` policy for a new-format file. Finally, call the construct API: + ```c++ void IVF::construct( const float* data, @@ -118,72 +182,29 @@ void IVF::construct( ); ``` -- **data**: Pointer to the raw data vectors. -- **centroids**: Clustering centroids; tune `cluster_num` for the dataset and search budget. -- **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`. -- **faster**: If true, enable fast implementations for RaBitQ (By default, it is set as `false` to pursue better accuracy.). +- **data**: Vectors in the original input coordinates. +- **centroids**: One centroid per configured cluster, in the original input coordinates. +- **cluster_ids**: One ID per input vector, in `[0, num_clusters())`. +- **faster**: Use a calibrated scaling factor instead of per-vector scale search; defaults to `false`. See [encoding paths](../rabitq/quantizer.md#implementations). - **num_threads**: Maximum number of OpenMP threads used to quantize clusters. For example: -```c++ -ivf.construct(data.data(), centroids.data(), cids.data(), true); -``` - -During construction, we process clusters in parallel, rotate their centroids and -vectors, and compute one-bit codes and factors. Quantized mode also computes -`total_bits - 1` extra-bit codes and factors; raw mode stores the original -float32 vectors instead. In raw mode, the index owns a copy of the input, so -the original data can be released after construction. - -Construction reuses a rotation buffer of at most 32 vectors per worker. This -temporary storage uses at most `workers * 32 * padded_dim * sizeof(float)` bytes, -independently of cluster sizes. -New indexes pad the rotated dimension to a multiple of 32 for every supported -bit width, including raw mode (32). This applies to both rotator types. -Extra-bit layouts retain full 64-coordinate blocks and store any final -32-coordinate tail compactly. - -After construction, you can directly save the index file to disk: -```c++ -ivf.save(output_index_file); -``` -New files use the `RABQIDX1` magic, version 2, storage flags, an explicit -padded dimension, and a uint32 initializer type (1 = Flat, 2 = Flat RaBitQ, -3 = HNSW) before the index metadata. `Auto` is resolved before saving. Loading -restores the routing type, storage mode, rotation state, and padded dimension. -Raw files also contain the original vectors, so querying does not require an external dataset. - -Version-1 `RABQIDX1` files remain readable with their stored padding. Legacy -unversioned quantized files and raw v1 files retain their original 64-dimension -padding. All three historical formats infer routing with the original 20,000 -cluster threshold. HNSW routing still requires the accompanying `.hnsw` file. -Flat RaBitQ loads its saved codes without random re-quantization. Saving a loaded -legacy index writes the new header without changing its padded dimension or encoded data. Older library -versions cannot read the new format. - -### Data Layout -The main data layout for our IVF is organized as follows: ```c++ -[batch data] // 1-bit code and factors -[ex_data] // code for remaining bits, or original float32 vectors -[ids] // PID of vectors (organized by clusters) -[cluster_lst] // List of clusters' metadata in IVF +ivf.construct(data.data(), centroids.data(), cids.data(), true); ``` -In raw mode, `ex_data` holds `4 * num_points * dim` bytes in cluster order, -without padding or extra-bit factors. The one-bit filtering codes and factors -are additional storage. `nbits()` (Python `nbits`) reports 32 in raw mode; -the one-bit filtering code is not included in that value. ## Querying -Currently, querying requires the index to be loaded in memory. If you want to use a previously saved index on the disk, firstly load it into memory: + +To search a saved index, load it first: ```c++ using index_type = rabitqlib::ivf::IVF; index_type ivf; ivf.load(index_file); ``` -Once the index is loaded, you can call the search function for queries: +Search a constructed or loaded index with: + ```c++ void IVF::search( const float* __restrict__ query, @@ -220,14 +241,7 @@ The querying examples also default to automatic selection; Python's `--use-hacc` flag forces HACC, and the C++ example accepts an optional final `true` or `false` argument. -During search, we rotate the query and select the `nprobe` closest centroids. -FastScan filters candidates using one-bit codes, then reranking uses either the -remaining quantized bits or the stored raw vectors. Raw reranking computes -squared L2 or `1 - dot(query, vector)` in the original coordinates; the search -API is unchanged. Cluster selection and filtering remain approximate in both -modes. Search returns the top `k` results after scanning the selected clusters. - -Starting with 0.5.2, C++ `search_batch` distributes multiple queries across workers. +C++ `search_batch` distributes queries across workers. Input queries are contiguous row-major float32 vectors with shape `(num_queries, dimension())`; output buffers have shape `(num_queries, k)`. @@ -246,17 +260,13 @@ need separate output buffers, and must not overlap index updates. Python's existing `index.search(queries, k, nprobe, num_threads=4)` uses this batch path with the same output and precision rules. -Python search releases the GIL. Concurrent searches may use independent -`nprobe`, `high_accuracy`, and `num_threads` values. Conflicting updates raise -`RuntimeError` immediately; see the -[Python concurrency contract](../quick_start.md#threading-and-file-paths). +Concurrent Python searches may use independent `nprobe`, `high_accuracy`, and +`num_threads` values. See the [concurrency contract](usage.md#threading-and-file-paths). -## Updating an Index +## Updating an index -Starting with 0.5.1, C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` -on a constructed or loaded index without the original dataset. HNSW provides -its own [update APIs](hnsw.md#updating-an-index) starting with 0.5.2. -IVF updates preserve the index file format. +C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` on a constructed +or loaded index without the original dataset. Updates preserve the file format. ```c++ void IVF::add( @@ -271,13 +281,13 @@ size_t IVF::remove(const PID* ids_to_remove, size_t n); ``` - **data**, **n**: `n` new vectors, quantized like `construct` does. Vector `i` - receives the PID `max_elements() + i`, using the count before the call. Existing + receives the PID `num_stored_points() + i`, using the count before the call. Existing IDs stay unchanged, and removed IDs are never reused. - **cluster_ids**: The cluster of each new vector, in `[0, cluster_num)`. When it is `nullptr`, vectors use the same centroid routing as queries. Routing uses the initializer selected at construction or restored from the index file. - **faster**, **num_threads**: Same as in `construct`. -- **ids_to_remove**: PIDs to remove. Every PID must be below `max_elements()`; nothing is +- **ids_to_remove**: PIDs to remove. Every PID must be below `num_stored_points()`; nothing is removed if one is not. `remove` returns how many were newly removed and can be repeated safely. @@ -297,11 +307,10 @@ one-dimensional integer array or sequence with one valid cluster ID per vector; IDs and returns the number newly removed. Duplicate IDs count once. Python `num_threads` defaults to 1; `fast_quantization` defaults to `False`. -The stored point count (`max_elements()` in C++, `index.max_elements` in Python) -increases after `add()` and includes removed points. Rebuilding with C++ -`construct` or Python `build` expects that many rows. In C++, callers must prevent -`add()` or `remove()` from overlapping with searches or other updates on the same -index. Python detects these conflicts and raises `RuntimeError` immediately. +After appending, rebuilding with `construct` or `build` requires +`num_stored_points` rows, including removed records. Updates require exclusive +access; see the +[concurrency contract](usage.md#threading-and-file-paths). ### Cost of `add` @@ -312,13 +321,9 @@ unchanged if it throws. Every call copies the whole storage, so: - Each call has a fixed cost that grows with the size of the whole index, on top of a cost per new vector. **Add many vectors per call rather than one at a time.** -On an index of 1,000,000 vectors with 128 dimensions, 1,024 clusters and 1 bit -(8 threads, portable build, Intel i7-12700KF), adding 1,000 vectors in one call took -about 5 ms, and adding the same vectors one call each took about 3 s (2.8 s and 3.7 s -in two runs). For that 33 MB file, a process that had just loaded it held about -65 MB, and peaked at about 100 MB during one `add` of 1,000 vectors. - There is no spare capacity: every `add` moves the data. +Measure update latency and peak memory with your index size, bit width and batch +size; serialized file size alone does not determine peak process memory. ### Recall after `add` @@ -331,7 +336,7 @@ recall matters more than the cost of a rebuild. ### How removal is stored `remove` hides points from search but keeps their storage. Removed points still -count in `max_elements()` and cannot be restored. Each nonempty call scans all +count in `num_stored_points()` and cannot be restored. Each nonempty call scans all stored point IDs, so batch removals when possible. A removed point has its `f_add` value set to `+inf`. That value is the constant @@ -351,3 +356,67 @@ reinterpreted. Release 0.5.0 also excludes removed points when reading its supported legacy formats. However, current saves always use `RABQIDX1` v2, including when resaving a loaded legacy index. Older releases cannot read these files, regardless of whether they contain removals. + +## Persistence + +Save a constructed index with: + +```c++ +ivf.save(output_index_file); +``` +New files use the `RABQIDX1` magic, version 2, storage flags, an explicit +padded dimension, and a uint32 initializer type (1 = Flat, 2 = Flat RaBitQ, +3 = HNSW) before the index metadata. `Auto` is resolved before saving. Loading +restores the routing type, storage mode, rotation state, and padded dimension. +Raw files also contain the original vectors, so querying does not require an external dataset. + +Version-1 `RABQIDX1` files remain readable with their stored padding. Legacy +unversioned quantized files and raw v1 files retain their original 64-dimension +padding. All three historical formats infer routing with the original 20,000 +cluster threshold. HNSW routing still requires the accompanying `.hnsw` file. +Flat RaBitQ loads its saved codes without random re-quantization. Saving a loaded +legacy index writes the new header without changing its padded dimension or encoded data. Older library +versions cannot read the new format. + +## Implementation details + +### Construction + +During construction, we process clusters in parallel, rotate their centroids and +vectors, and compute one-bit codes and factors. Quantized mode also computes +`total_bits - 1` extra-bit codes and factors; raw mode stores the original +float32 vectors instead. In raw mode, the index owns a copy of the input, so +the original data can be released after construction. + +Construction reuses a rotation buffer of at most 32 vectors per worker. This +temporary storage uses at most `workers * 32 * padded_dim * sizeof(float)` bytes, +independently of cluster sizes. + +New indexes pad the rotated dimension to a multiple of 32 for every supported +bit width, including raw mode (32). This applies to both rotator types. +Extra-bit layouts retain full 64-coordinate blocks and store any final +32-coordinate tail compactly. + +### Data layout + +IVF stores the following data: + +```c++ +[batch data] // 1-bit code and factors +[ex_data] // code for remaining bits, or original float32 vectors +[ids] // PID of vectors (organized by clusters) +[cluster_lst] // List of clusters' metadata in IVF +``` +In raw mode, `ex_data` holds `4 * num_points * dim` bytes in cluster order, +without padding or extra-bit factors. The one-bit filtering codes and factors +are additional storage. `nbits()` (Python `nbits`) reports 32 in raw mode; +the one-bit filtering code is not included in that value. + +### Search algorithm + +During search, we rotate the query and select the `nprobe` closest centroids. +FastScan filters candidates using one-bit codes, then reranking uses either the +remaining quantized bits or the stored raw vectors. Raw reranking computes +squared L2 or `1 - dot(query, vector)` in the original coordinates; the search +API is unchanged. Cluster selection and filtering remain approximate in both +modes. Search returns the top `k` results after scanning the selected clusters. diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 7cb4a0db..50aeea61 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -6,47 +6,89 @@ from [NGT](https://github.com/yahoojapan/NGT), with batched RaBitQ distance esti FastScan estimates neighbor distances; visited vertices are scored using stored raw vectors or quantized codes. -Set `quantization_bits=0` for raw vectors (default), or `4`/`8` for quantized SymphonyQG. -All three modes pad the rotated domain to a multiple of 32. Loading an older -index preserves its saved padded dimension, including historical multiples of 64. -Older library versions that enforce 64-dimension padding cannot load new files -whose padded dimension is not a multiple of 64. -Quantized vectors share a rotated global centroid. Refinement temporarily -reconstructs source vectors and estimates distances to stored target codes; -reverse edges are rescored because these estimates are directional. Raw refinement -uses owned raw vectors. No full reconstructed-vector cache is retained by default. - -## Index Construction +The [original paper](https://dl.acm.org/doi/abs/10.1145/3709730) uses raw vectors for +precise candidate distances during traversal. This library's optional 4/8-bit +storage and PiPNN initialization are implementation extensions; the paper's +reported performance should not be attributed to these configurations. -[PiPNN](https://dl.acm.org/doi/abs/10.1145/3770855.3817891) is the default initializer -for fast indexing, followed by one SymphonyQG refinement iteration. Random -initialization remains available and uses three iterations by default. +## Minimal workflow ### Python +Use `data` and `queries` from the [Quick start](../quick_start.md), or supply your +own float32 arrays. SymphonyQG builds directly from vectors and needs no clustering +step. The degree-32 example requires at least 33 data points. + ```python from rabitqlib import SymqgIndex # data and queries are float32 arrays with shape (count, dim). index = SymqgIndex(dim=data.shape[1], max_degree=32, metric="l2", quantization_bits=0) -index.build(data, ef_construction=200, num_threads=32, init="pipnn") +index.build(data, ef_construction=200, num_threads=1, init="pipnn") index.save("qg_example.index") loaded = SymqgIndex.load("qg_example.index") ids, distances = loaded.search_batch(queries, k=10, ef=100, num_threads=1) ``` -This example uses `search_batch()`, which requires 0.5.2 or newer. On earlier -versions, use `search()` with the same arguments. - -`init` defaults to `"pipnn"`; use `"random"` for random initialization. -Supported metrics are `"l2"` and `"ip"`. `max_degree` must be a multiple of 32 -and smaller than the point count. `ef_construction` controls the build search -window; `ef` controls the query search window. Python defaults to one thread. +`search()` accepts the same arguments and is equivalent to `search_batch()`. ### C++ -The C++ API provides `rabitqlib::symqg::QuantizedGraph` and `QGBuilder`. +Link `rabitqlib::index`. Here `data` is a row-major `std::vector` with +`rows * cols` values, and `query` contains `cols` values; use at least 33 rows. +See [C++ integration](../cpp_integration.md) for build setup. + +```cpp +#include +#include +#include + +using namespace rabitqlib::symqg; + +// data contains rows * cols floats. +QuantizedGraph qg(rows, cols, 32); +{ + QGBuilder builder(qg, 200, data.data(), 1, QGInitialization::PiPNN); + builder.build(); +} // release builder scratch before saving +qg.set_ef(100); +std::vector ids(10); +std::vector distances(10); +qg.search(query.data(), 10, ids.data(), distances.data()); +qg.save("qg_example.index"); +``` + +## Parameters and limits + +| Parameter | Contract | +| --- | --- | +| `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | +| `max_degree` | Required positive multiple of 32, strictly smaller than the point count. | +| `quantization_bits` | 0 for owned raw vectors (default), or 4/8 for quantized storage. | +| `metric` | Python `"l2"` or `"ip"`; defaults to `"l2"`. | +| `ef_construction` | Required positive construction search window. | +| `init` | Python `"pipnn"` (default) or `"random"`. | +| `k`, `ef` | Python requires `1 <= k <= num_points` and `ef >= k`. | +| `num_threads` | Python defaults to 1; 0 selects available logical CPUs. | + +Both Python search names return `(ids, distances)` arrays. Distances are squared +L2 or `1 - dot`, ordered from smallest to largest; quantized storage returns +estimates. SymphonyQG has no incremental add/remove API. See shared +[input, count and concurrency contracts](usage.md). + +## Index construction + +See the shared [index count contract](usage.md#index-counts). C++ +`num_vertices()` returns the configured graph size, even before building. + +[PiPNN](https://dl.acm.org/doi/abs/10.1145/3770855.3817891) is the default initializer +for fast indexing, followed by one SymphonyQG refinement iteration. Random +initialization remains available and uses three iterations by default. + +Include `rabitqlib/index/symqg/qg.hpp` and `qg_builder.hpp` from the same directory, +and link `rabitqlib::index`. The C++ API provides +`rabitqlib::symqg::QuantizedGraph` and `QGBuilder`. Only float vectors are supported; C++17 declarations may omit `` through class template argument deduction: @@ -71,44 +113,26 @@ for PiPNN or three passes for `QGInitialization::Random`; `build(n)` requests `n >= 2` passes. The older overload with a numeric fifth argument uses that value as the random-initialization seed. -`builder.reset(data)` reinitializes vectors of the same shape and reuses workspace. -Input may be released after reset returns; rebuild before querying or saving. -`cache_vectors=true` retains transformed vectors and query factors -(`4 * num * padded_dim` bytes plus factors). [QGKMeans](../clustering.md#qgkmeans) enables -this for its centroid graph; the cache is not persisted. The graph seed controls -rotation, and the builder seed controls random initialization and fallback edges. -PiPNN uses fixed internal seeds. - -```cpp -using namespace rabitqlib::symqg; - -// data contains rows * cols floats. -QuantizedGraph qg(rows, cols, 32); -{ - QGBuilder builder(qg, 200, data.data(), 32, QGInitialization::PiPNN); - builder.build(); -} // release builder scratch before saving -qg.save("qg_example.index"); -``` +### Input lifetime and rebuilds The caller can release input vectors after the `QGBuilder` constructor returns. -Python retains the caller's array during `build`. Complete the build before -querying or saving. Initialization does not change query-distance conventions or -save/load formats. +Python retains the caller's array during `build`. Complete construction before +querying or saving. -### Data Layout +`builder.reset(data)` reinitializes vectors of the same shape and reuses workspace. +Input may be released after reset returns; rebuild before querying or saving. -Each row contains: +### Vector cache -```text -[Raw vector or packed 4/8-bit code + factors] -[One-bit neighbor codes + factors] -[Neighbor IDs] -``` +`cache_vectors=true` retains transformed vectors and query factors +(`4 * num * padded_dim` bytes plus factors). [QGKMeans](../clustering.md#qgkmeans) enables +this for its centroid graph; the cache is not persisted. -Neighbor codes use FastScan batches of 32, which determines the degree alignment. -Quantized storage reduces each vector's size but not its neighborhood codes; -those codes and temporary refinement pools still consume substantial memory. +### Random seeds + +The graph seed controls FHT/Kac rotation; MatrixRotator samples its own matrix. +The builder seed controls random initialization and fallback edges. PiPNN uses +fixed internal seeds. ## Querying @@ -125,7 +149,7 @@ std::vector distances(10); qg.search(query.data(), 10, ids.data(), distances.data()); ``` -Starting with 0.5.2, C++ `search_batch()` accepts contiguous row-major batches +C++ `search_batch()` accepts contiguous row-major batches and reuses scratch storage within each worker, preserving the results of independent `search()` calls: @@ -145,21 +169,18 @@ do no work and may use null buffers. Inputs and outputs must not overlap. Concurrent searches need separate outputs; do not modify the index or call `set_ef()` while searches are running. -Python provides `search_batch(queries, k, ef, num_threads=1)` starting with 0.5.2. -The existing `search()` accepts the same two-dimensional array and remains an equivalent -entry point. Both return `(ids, distances)` arrays of shape `(num_queries, k)`, -including empty batches. Both APIs default to one worker; `num_threads=0` +Python `search_batch(queries, k, ef, num_threads=1)` and `search()` accept +the same two-dimensional query array. Both return `(ids, distances)` arrays +of shape `(num_queries, k)`, including empty batches. Both APIs default to one worker; `num_threads=0` selects the available logical CPU count, capped by the number of queries. Queries use the original dimension, including when the internal rotation pads it. -Python search releases the GIL and uses the `ef` passed to each call without -changing a shared default. Concurrent searches may use different windows. -A concurrent `build()` on the same index raises `RuntimeError` immediately; -see the [Python concurrency contract](../quick_start.md#threading-and-file-paths). +Python searches use per-call `ef` values. For access rules, see the +[concurrency contract](usage.md#threading-and-file-paths). -Starting with 0.6.0, C++ also provides +For C++, use `search_batch_with_ef(queries, num_queries, k, ids, distances, ef, num_threads=1)` -for concurrent calls with independent windows: +to run concurrent calls with independent windows: ```cpp qg.search_batch_with_ef( @@ -173,3 +194,31 @@ All concurrent searches require separate outputs and an unchanged index. See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their Python counterparts for complete examples. + +## Persistence + +All three modes pad the rotated domain to a multiple of 32. Loading an older +index preserves its saved padded dimension, including historical multiples of 64. +Older library versions that enforce 64-dimension padding cannot load new files +whose padded dimension is not a multiple of 64. + +## Implementation details + +Quantized vectors share a rotated global centroid. Refinement temporarily +reconstructs source vectors and estimates distances to stored target codes; +reverse edges are rescored because these estimates are directional. Raw refinement +uses owned raw vectors. No full reconstructed-vector cache is retained by default. + +### Data layout + +Each row contains: + +```text +[Raw vector or packed 4/8-bit code + factors] +[One-bit neighbor codes + factors] +[Neighbor IDs] +``` + +Neighbor codes use FastScan batches of 32, which determines the degree alignment. +Quantized storage reduces each vector's size but not its neighborhood codes; +those codes and temporary refinement pools still consume substantial memory. diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md new file mode 100644 index 00000000..c9062254 --- /dev/null +++ b/docs/docs/index/usage.md @@ -0,0 +1,140 @@ +# Index usage contracts + +Shared input, count, concurrency and persistence conventions for IVF, HNSW and +SymphonyQG. Start with the [Python quick start](../quick_start.md) for a complete example. + +## Input requirements + +`IvfIndex`, `HnswIndex`, and `SymqgIndex` accept data and query arrays with shape +`(count, dim)`, converted to contiguous float32. Their FHT/Kac rotation requires +`64 <= dim <= 65536`. Supply the original dimension; the indexes pad internally +to a multiple of 32. Queries and centroids must have the same original dimension +as the indexed data. + +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; +the Python indexes do not expose a rotator choice. + +## Index counts + +The unified count properties below are introduced in **0.7.0**. In particular, +HNSW's `num_points` now excludes removed records; see the +[count migration guide](../modular_migration.md#unified-index-counts) when upgrading. + +IVF, HNSW, and SymphonyQG expose these read-only Python properties +(C++ uses the same names with `()`): + +| Property | Meaning | +| --- | --- | +| `max_elements` | Logical record capacity, excluding SIMD padding | +| `num_points` | Nonremoved points | +| `num_stored_points` | Stored records, including removed ones | + +`0 <= num_points <= num_stored_points <= max_elements`. Before building, both +point counts are zero; configured capacity may be nonzero. + +At build time, IVF requires exactly `max_elements` input rows. HNSW accepts +1 through `max_elements` rows, so it can reserve capacity for later additions. + +Removal decreases only `num_points`; slots and IDs remain occupied. Append IDs +start at the previous `num_stored_points`. Use this stored count for ID bounds, +and `num_points` for a live count. Counts survive save/load. + +IVF grows automatically and has no spare logical capacity after build or append. +HNSW's `resize` changes only capacity. SymphonyQG has no incremental updates; +all three counts are equal after build. + +## Estimate storage + +The following estimates describe the main float32 index storage, not total process +memory or peak construction memory. Let `d` be the original dimension, `D` the +stored padded dimension, and `B` the total quantized bits per dimension. New +indexes use `D = 32 * ceil(d / 32)`; loaded indexes retain their saved padding. + +For IVF/HNSW, codes and factors occupy `V(B)` bytes per vector: + +- One bit: `V(1) = D / 8 + 12` (three float32 factors). +- Two through nine bits: `V(B) = D * B / 8 + 20` (three binary and two extra-code factors). + +| Index / mode | Main bytes per stored point or capacity slot | +| --- | --- | +| IVF quantized | Approximately `V(B) + 4`, including the point ID | +| IVF raw (`nbits=32`) | Approximately `D / 8 + 12 + 4*d + 4`, including filtering codes, raw vector and ID | +| HNSW | `V(B) + 8*M + 12` per **capacity slot** for base-layer records: up to `2*M` neighbor IDs, a link count, cluster ID and label | +| SymphonyQG raw | `4*d + R*(D/8 + 12)`, where `R = max_degree` | +| SymphonyQG quantized | `D*B/8 + 8 + R*(D/8 + 12)`, with `B` equal to 4 or 8 | + +IVF allocates one-bit codes and their three factors in full batches of 32 +**per cluster**. For cluster sizes `n_c`, that part costs +`sum(32 * ceil(n_c / 32)) * (D/8 + 12)` bytes. Extra codes/raw vectors and IDs +use the stored point count. Small clusters therefore increase padding overhead. +Centroids, routing structures and cluster metadata are additional. + +HNSW allocates base-layer records for `max_elements`, including unused capacity. +Upper-layer edges, label maps, locks, visited lists and other bookkeeping are +additional. Removing points does not reclaim their storage in either IVF or HNSW. + +SymphonyQG stores a one-bit code and two factors **for each neighbor in each row**, +plus a four-byte neighbor ID. This accounts for `R*(D/8 + 12)` bytes per row, +even in quantized mode: reducing the vector's own code width does not shrink +these neighborhood copies. + +For example, with `d=D=128`, four-bit storage, `M=16` and `R=32`, the terms above +are about 88 bytes per IVF point before batch padding, 224 bytes per HNSW capacity +slot, and 968 bytes per SymphonyQG row. These are layout calculations, not measured +RSS. All indexes also retain rotation state; construction, query scratch, +allocator overhead and update workspaces add to memory use. Measure process +memory on the intended workload when sizing a deployment. + +These formulas follow `BatchDataMap` / `QGBatchDataMap` in +`fastscan/data_layout.hpp`, `BinDataMap` / `ExDataMap` in `core/data_layout.hpp`, +and the index record layouts described in the individual guides. Applications +allocating code buffers should call the corresponding `data_bytes()` helpers +rather than hardcode these estimates. + +## Threading and file paths + +The library releases the Python GIL during native index search, +construction, updates, and file I/O. The following contract applies to +`IvfIndex`, `HnswIndex`, and `SymqgIndex`: + +| Operation on the same index | Concurrent access | +| --- | --- | +| `search`, `search_batch`, `save`, and property reads | May run together | +| `build`, `add`, `remove`, and `resize`, where available | Require exclusive access | + +A conflicting call immediately raises +`RuntimeError("Index is busy: conflicting operation in progress")`. +It does not wait or queue an update. Access protection includes Python input +conversion and is released when the operation returns or raises an exception. +Different index objects can operate independently. Concurrent saves must use +different output paths, including any sidecar files. + +Search parameters such as `ef`, `nprobe`, precision, and `num_threads` belong to +each call. When using a Python thread pool, consider `num_threads=1` to avoid +starting multiple native worker pools. Keep input arrays and any shared backing +storage unchanged for the entire call; do not resize or modify them from another +thread. Returned result arrays own their storage. + +These access checks belong to the Python wrappers. C++ callers must synchronize +index updates themselves and use separate search output buffers. + +`num_threads=0` selects the detected available logical CPU count. Positive values +set an upper limit, capped at that count; small workloads may use fewer workers. +Python index methods default to one thread; clustering defaults to `0`. +On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding +before starting the program. + +Index save/load paths are UTF-8 strings on Windows and native path bytes on POSIX +in C++; Python paths are Unicode strings on all platforms. + +## Index file compatibility + +The compatibility suite loads frozen files written by v0.5.2, compares search +IDs and distances, and verifies save/load migration and rejection of truncated +or unsupported versions. See the +[fixture provenance and coverage](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/tests/fixtures/persistence). +Compatibility is format-specific; preserving old readers' ability to open new +files is not guaranteed. Consult each index's format documentation before +downgrading the library. diff --git a/docs/docs/installation.md b/docs/docs/installation.md new file mode 100644 index 00000000..c44542e4 --- /dev/null +++ b/docs/docs/installation.md @@ -0,0 +1,87 @@ +# Installation + +Platform requirements and source installation for RaBitQ Library **0.7.0**. +For the first Python example, see [Quick Start](quick_start.md). + +## Requirements + +| Platform | CPU baseline | Python wheel targets | +| --- | --- | --- | +| Linux x86-64 | AVX2 and FMA | CPython 3.11–3.14 | +| Linux ARM64 (AArch64) | NEON | CPython 3.11–3.14 | +| Windows x86-64 | AVX2 and FMA | CPython 3.11–3.14 | +| macOS 14+ ARM64 (Apple Silicon) | NEON | CPython 3.11–3.14 | + +Source builds require a C++17 compiler and CMake 3.20 or newer. Full ANN/Python builds +also require OpenMP; core/FastScan can disable it with `RABITQ_USE_OPENMP=OFF`. +Windows uses MSVC (Visual Studio 2026 with the Desktop development with C++ +workload and CMake 4.2+ for its generator). Apple Silicon uses AppleClang and +an external OpenMP runtime such as Homebrew `libomp`. + +Linux ARM64 source builds and repaired wheels run on native AArch64 CI. +Linux ARM64 wheels carry `manylinux_2_27_aarch64` and +`manylinux_2_28_aarch64` tags. + +
+CPU dispatch details + +On x86-64, optional AVX-512 kernels require AVX2, FMA, and AVX-512F/BW/DQ; +MSVC builds also require AVX-512VL/CD. Popcount-specific kernels additionally +require AVX-512 VPOPCNTDQ. Detection checks CPU features and OS support for +the required register state. Missing AVX-512 features select the AVX2 backend. + +On ARM64, the build excludes x86 kernels and uses NEON and portable scalar +implementations. Standard FastScan requires AVX2/FMA, a supported AVX-512 +backend, or NEON; scalar fallbacks do not remove that requirement. A pure-core +build can use scalar kernels without the FastScan CPU baseline. + +
+ +## Python wheels + +```bash +python -m pip install "rabitqlib>=0.7.0" +``` + +For changes from 0.6.0, see the [0.7.0 upgrade notes](releases/0.7.0.md). + +Wheels target the platforms above and require no compiler or CMake. +Linux ARM64 and macOS ARM64 wheels bundle OpenMP, so wheel users do not +need a separate OpenMP installation. +Release wheels disable native CPU tuning and select supported kernels at runtime. + +## Get the source + +For Python source builds or C++ use, obtain the 0.7.0 checkout once. With Git +installed, run: + +```bash +git clone --branch v0.7.0 --depth 1 https://github.com/VectorDB-NTU/RaBitQ-Library.git +cd RaBitQ-Library +``` + +Run the source-build commands below from this repository root. If you already +have the matching checkout, use it directly. + +## Build the Python bindings from source + +Source builds require Python 3.11 or newer, a C++17 compiler, CMake 3.20 or +newer, and OpenMP. First [get the source](#get-the-source), then on Ubuntu or +Debian run: + +```bash +sudo apt-get update +sudo apt-get install -y build-essential cmake libomp-dev +python -m pip install . +``` + +On Apple Silicon, use a native ARM64 Python environment. From the repository root: + +```bash +brew install cmake ninja libomp +python -m pip install . -Ccmake.define.OpenMP_ROOT="$(brew --prefix libomp)" +``` + +On Windows, install the build tools listed above, then run `python -m pip install .` +from the repository root. For a portable source build on any supported platform, +also pass `-Ccmake.define.RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`. diff --git a/docs/docs/integrations/zvec.md b/docs/docs/integrations/zvec.md index d66a6b51..802207a5 100644 --- a/docs/docs/integrations/zvec.md +++ b/docs/docs/integrations/zvec.md @@ -18,9 +18,11 @@ This overview follows zvec commit Its [bundled dependency][dependency] includes local patches for lookup-table reuse and SIMD error handling; check the pinned version before adapting its code. -To integrate RaBitQ-Library, start with the [quantizer](../rabitq/quantizer.md) -and [estimator](../rabitq/estimator.md) guides. For a complete standalone index, -see the [quick start](../quick_start.md). +For a new per-vector integration, start with the +[stable core API v1](../rabitq/core_api.md#minimal-example). For FastScan or existing +record layouts like those used here, consult the advanced +[quantizer](../rabitq/quantizer.md) and [estimator](../rabitq/estimator.md) references. +For a complete standalone index, see the [quick start](../quick_start.md). [ivf-reformer]: https://github.com/alibaba/zvec/blob/44d8bc6f7c3d656efde2641c9410432ea14c1657/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc [ivf-entity]: https://github.com/alibaba/zvec/blob/44d8bc6f7c3d656efde2641c9410432ea14c1657/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc diff --git a/docs/docs/kernel_ip.md b/docs/docs/kernel_ip.md index 73c654c1..dd826e48 100644 --- a/docs/docs/kernel_ip.md +++ b/docs/docs/kernel_ip.md @@ -1,26 +1,26 @@ -# Kernel - Inner Product +# Inner-product kernels > Inner-product kernels are selected at runtime. The library provides AVX2/FMA > and AVX-512 implementations on x86-64, plus NEON and portable scalar > implementations on ARM64. AVX-512 VPOPCNTDQ is used by popcount-specific > x86 paths when available. -This part introduces how to compute the inner product between quantization codes and rotated query vectors i.e., $\left< \mathbf{x}_0,\mathbf{q}_r'\right>$ and $\left< \mathbf{x}_u,\mathbf{q}_r'\right>$. The implementation includes two types: +These kernels compute inner products of sign or multi-bit codes with a rotated +query. The notation follows the [Estimator](rabitq/estimator.md). -1. The inner product between binary codes and $\mathbf{q}_r'$. -2. The inner product between multi-bit codes and $\mathbf{q}_r'$. +## The kernel for binary codes +### Single code -## The Kernel for Binary Codes - -### Single Code -We compute the inner product between a single binary vector $\mathbf{x}_0$ and a floating-point vector $\mathbf{q}_r'$ via bitwise-and `&` and `popcnt`. We first quantize $\mathbf{q}_r'$ into a vector of 4-bit unsigned integers (based on a fast version of RaBitQ). +For a sign code $\mathbf{x}_0$, first quantize the rotated query +$\mathbf{q}_r'$ into four-bit unsigned integers using the fast RaBitQ encoder. +Then use bitwise-and `&` and `popcnt` to estimate the inner product. $$ \mathbf{q}_r'\approx \Delta_q \mathbf{q}_u + v_q\cdot \mathbf{1}_D $$ -Let $\mathbf{q}_u^{(i)}$ be the $i$-th bit of $\mathbf{q}_u$. Then with the quantized vector $\mathbf{q}_u$, we can compute the inner product based on the following formula: +Let $\mathbf{q}_u^{(i)}$ be bit plane $i$ of $\mathbf{q}_u$. Then: $$ \begin{align} @@ -32,23 +32,33 @@ $$ \end{align} $$ +### Batch code -### Batch Code -We compute the inner product between a batch of binary vectors $\mathbf{x}_0$ and a floating-point vector $\mathbf{q}_r'$ via `FastScan`. Here we provide a brief introduction and refer readers to a detailed [tutorial by Faiss](https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)). +FastScan scores batches of sign codes using query lookup tables. See the +[Faiss tutorial](https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)) +for background. -For a $D$-bit binary, we split it into $M=D/4$ segments. We prepare look-up-tables for each segment. -- $\mathrm{LUT}[m][mask]$: the inner product with $\mathbf{q}_r'$ in the $m$-th segment, i.e., the $(4m)$-th to the $(4m+3)$-th dimensions, when the code in the $m$-th segment equals to $mask$. +For a $D$-bit binary code, we split it into $M=D/4$ segments. We prepare look-up-tables for each segment. +- $\mathrm{LUT}[m][mask]$: the inner product with $\mathbf{q}_r'$ in the $m$-th segment, i.e., the $(4m)$-th to the $(4m+3)$-th dimensions, when the code in the $m$-th segment equals to $mask$. Based on the look-up-tables, we can compute the inner product as follows: $$ \begin{align} -\left< \mathbf{x}_0,\mathbf{q}_r'\right>= \sum_{m=0}^{M-1} \mathrm{LUT}[m][\mathbf{x}_{0}[4m:4m+3]] +\left< \mathbf{x}_0,\mathbf{q}_r'\right>= \sum_{m=0}^{M-1} \mathrm{LUT}[m][\mathbf{x}_{0}[4m:4m+3]] \end{align} $$ +This equality describes real-valued lookup tables. Standard FastScan quantizes +the entries to eight bits for SIMD accumulation. For one-bit data codes, that +additional error is typically negligible relative to the RaBitQ encoding error. +As the data-code bit width increases, the encoding error shrinks and LUT precision +becomes more significant. High-accuracy FastScan (HACC) uses 16-bit LUT entries +to reduce the error carried from the sign-code inner product into extra-bit +refinement. IVF selects HACC by default for 4–9 total bits and standard FastScan +for 1–3 bits or raw-vector reranking; see [precision selection](index/ivf.md#querying). -## The Kernel for Multi-bit Codes +## The kernel for multi-bit codes For multi-bit codes, the packed unsigned values are unpacked and accumulated by the selected AVX2/FMA, AVX-512, NEON, or scalar kernel. The dispatch table supports zero diff --git a/docs/docs/modular_migration.md b/docs/docs/modular_migration.md new file mode 100644 index 00000000..f9faaf31 --- /dev/null +++ b/docs/docs/modular_migration.md @@ -0,0 +1,190 @@ +# Migrating to 0.7.0 + +This guide covers upgrading from **0.6.0 to 0.7.0**. See the +[release overview](releases/0.7.0.md). + +The former full-library CMake alias and forwarding header paths are removed. +Version 0.7.0 introduces the stable float32 core API v1, reusing the existing +quantization algorithms and code interpretation. Persisted index formats remain +compatible with 0.6.0. Index count semantics are unified as described below. +Rebuild C++ consumers against 0.7.0; no cross-version C++ binary ABI is promised. + +## Choose a component explicitly + +```cmake +find_package(rabitqlib CONFIG REQUIRED COMPONENTS index) +target_link_libraries(my_app PRIVATE rabitqlib::index) +``` + +Replace `rabitqlib::rabitqlib` with the smallest component needed: + +| Target | Contains | Dependency | +| --- | --- | --- | +| `rabitqlib::core` | Encoding, single-vector estimation, rotation and packing | Numerical helpers and essential SIMD | +| `rabitqlib::fastscan` | Batch layouts, packing/encoding adapters and FastScan/HACC scoring | core | +| `rabitqlib::index` | IVF, HNSW, SymphonyQG and clustering | fastscan and core | +| `rabitqlib::python` | Existing Python ANN extension; build-tree alias only | index | + +The old implementation target `rabitq_headers` is also removed. C++ libraries remain +static. Python still uses `import rabitqlib`; it is not an exported C++ package component. + +For a numerical library without ANN or OpenMP: + +```bash +cmake -S . -B build-core -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_INDEX=OFF -DRABITQ_BUILD_FASTSCAN=OFF \ + -DRABITQ_USE_OPENMP=OFF -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DCMAKE_INSTALL_PREFIX="" +cmake --build build-core --parallel +cmake --install build-core +``` + +Enable FastScan separately if needed. ANN/Python still require OpenMP. A core-only +installation contains core and its header dependencies, including bundled Eigen; +it does not install FastScan, index, clustering or vendored hnswlib headers. Source-tree +presence of disabled modules does not make them part of the installed component. + +Use a fresh install prefix during migration. CMake installation does not remove files +left by an older package, and stale headers can conceal a missing migration. Point each +consumer to the intended prefix and rebuild its CMake cache/objects accordingly. + +## Replace old includes + +Paths below are relative to `rabitqlib/`. Include the modules actually used instead of +relying on transitive declarations from a former umbrella header. + +| Removed header | Replacement | +| --- | --- | +| `quantization/rabitq.hpp` | `core/encoder.hpp` for per-vector encoding; `fastscan/encoder.hpp` for batch encoding; `fastscan/fastscan.hpp` for accumulation/LUT functions | +| `quantization/data_layout.hpp` | `core/data_layout.hpp` for per-vector maps; `fastscan/data_layout.hpp` for `BatchDataMap`, `QGBatchDataMap` and their const variants | +| `quantization/pack_excode.hpp` | `core/pack_excode.hpp` | +| `utils/rotator.hpp` | `core/rotator.hpp` | +| `index/query.hpp` | `core/query.hpp` for single vectors; `fastscan/query.hpp` for batches | +| `index/estimator.hpp` | `core/estimator.hpp` for single vectors; `fastscan/estimator.hpp` for batches | +| `index/lut.hpp` | Use the batch-query interfaces in `fastscan/query.hpp` | +| `defines.hpp` | Include `core/metric.hpp`, `core/types.hpp` or `index/types.hpp` according to the types used; internal Eigen aliases live in `core/detail/eigen.hpp` | +| `core/pack_codes.hpp` | `fastscan/pack_codes.hpp` for 32-vector batch packing/unpacking | + +New integrations should prefer [`core/api.hpp`](rabitq/core_api.md) and its versioned, +storage-neutral interface. Migrating an existing low-level caller to canonical headers +does not require changing its code bytes or factors. Batch scoring requires linking +FastScan; advanced core headers do not import that component for convenience. + +## Public and implementation interfaces + +### Encoding and shared types + +The advanced two-factor batch adapter has also moved: include `fastscan/data_layout.hpp` +for `QGBatchDataMap` / `ConstQGBatchDataMap`, and `fastscan/encoder.hpp` for +`quant::quantize_qg_batch`. It is shared by FastScan, clustering and SymphonyQG, so it +does not belong exclusively to the index module. Its code packing, factor order and +record sizes are unchanged. The three-factor `BatchDataMap` / `ConstBatchDataMap` also +belong to `fastscan/data_layout.hpp`. Include `fastscan/encoder.hpp` for +`quantize_one_batch` and `quantize_split_batch`, and link `rabitqlib::fastscan` when +using this format. The fixed batch size, permutation and `pack_codes`/`unpack_codes` +are all in `fastscan/pack_codes.hpp`; no FastScan declarations remain in core. + +Core retains single-vector `BinDataMap`, `ExDataMap` and `BaseDataMap`, per-vector +packing and generic encoding loops that do not impose a 32-vector format. Shared +unaligned value/array access resides in `utils/packed.hpp`. The v1 `SplitPacked` layout +is a per-vector sign/extra representation and remains part of core. + +`index/types.hpp` supplies `PID`, `kPidMax` and `AnnCandidate` without importing Eigen. +Identifier width, candidate layout and marker semantics are unchanged. The advanced +`exact_nn` helper now requires `index/detail/exact_nn.hpp` instead of `utils/space.hpp`. +Metric definitions move to `core/metric.hpp`; `ScalarQuantizerType` now resides in +`core/types.hpp`. The unused `BIT_ID` macro and internal `LOWBIT` macro are removed. +The new `core/api.hpp` exposes the stable v1 interface; it was not present in 0.6.0. + +Only `core/api.hpp`'s documented v1 surface has the new source-stability guarantee. +Other core/fastscan/index module interfaces remain available as advanced interfaces; +they are not implicitly frozen by that promise. + +Query metric arguments use `MetricType`. `SplitSingleQuery` no longer accepts +integer arguments; pass `METRIC_L2` or `METRIC_IP`. + +### Index headers + +IVF, HNSW and SymphonyQG operations now compile into `rabitqlib::index`. Their +main headers no longer import numerical implementations, search scratch, Eigen +or hnswlib. This implementation move preserves method signatures and file formats. + +Include the types you use explicitly: + +| Types | Header | +| --- | --- | +| `RotatorType`, `quant::RabitqConfig`, scalar quantizer choices | `core/types.hpp` | +| `ivf::InitializerType` | `index/ivf/types.hpp` | +| SymphonyQG query and buffer helpers | `fastscan/query.hpp`, `index/detail/buffer.hpp`, `index/detail/visited_set.hpp` | +| Internal `SearchScratch` | `index/symqg/detail/search_scratch.hpp` | + +Explicit includes of `core/rotator.hpp`, `core/encoder.hpp` and +`index/ivf/initializer.hpp` still expose their associated types. HNSW backend +search templates are private in `src/index/hnsw_search_impl.hpp`. + +### Removed HNSW members + +HNSW no longer exposes `HierarchicalNSW::rawDataPtr_`, `ResultRecord`, `Candidate`, +or `BoundedKNN` as public interfaces. The construction pointer and search helpers +are private implementation details. This is a C++ source compatibility break: +code that accessed those members must stop doing so. Pass construction input through +`construct()` and consume the distance/label pairs returned by `search()`; there are +no public replacements for the internal candidate types. Construction borrows its +input buffers only until the synchronous call returns or throws. This visibility +change does not alter persisted indexes or the Python API. + +### Internal headers + +Internal headers have explicit ownership: + +| Implementation | Location | +| --- | --- | +| Exact encoder scratch and quantization implementation | `core/detail/` | +| Eigen aliases used by numerical implementations | `core/detail/eigen.hpp` | +| Core SIMD declarations | `core/detail/simd/` | +| FastScan SIMD declarations and LUT implementation | `fastscan/detail/` | +| FastScan batch-encoder adapter | `fastscan/detail/encoder.hpp` | +| HNSW dispatch, search buffers, visited sets and shared scheduling | `index/detail/` | +| Non-template dispatcher implementation and private resolver | `src/simd/`, `src/fastscan/` and `src/index/` | + +The old `simd/*` public paths, `fastscan/{simd,estimator_dispatch,lut}.hpp`, +`quantization/rabitq_impl.hpp` and index-specific `utils/{buffer,hashset,visited_*}.hpp` +paths are removed. Callers using these internals should migrate to supported module +operations; direct use of `detail/` binds a caller to that implementation revision. +Some detail headers must be installed because public templates include them. Installation +does not turn them into a stable extension API. + +HNSW and IVF share `index/detail/parallel_for.hpp`. Callers relying on HNSW +to import the IVF initializer must include it explicitly. Scheduling behavior +is unchanged; custom thread-pool and SIMD registration remain unsupported. + +## Unified index counts + +All three indexes now share the [count contract](index/usage.md#index-counts) +in C++ and Python. + +Compared with the 0.6.0 Python API, IVF adds `num_points` and `num_stored_points`, +HNSW adds `num_stored_points`, and SymphonyQG adds `max_elements` and +`num_stored_points`. + +**HNSW `num_points` now excludes removed records.** Use `num_stored_points` for +ID bounds, append offsets, and minimum resize capacity. Use `num_points` for a +live count, without subtracting a separate deletion counter. + +The IVF constructor keyword `max_elements` still specifies the required initial +build size. SymphonyQG's C++ `num_vertices()` still returns the configured graph +size for builders. Existing index files need no conversion: loading recovers +counts from their removal markers. + +## Downstream rollout + +Pin 0.7.0 in downstream migration PRs and validate the affected integration +entry points with their maintainers. +Projects that have not migrated can continue pinning their supported old release. + +The conformance suite preserves the frozen v1 source consumer and historical index +fixtures. It also verifies component link closure, absent retired targets/headers, +standalone installed headers, and independent add_subdirectory/FetchContent/package +consumers. Historical release notes and the v0.5.2 fixture generator deliberately retain +their old target names: those documents/tools describe or build the historical release. diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index 24d472c2..c33ac4a9 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -1,58 +1,24 @@ -# Quick Start +# Quick start -RaBitQ Library provides Python bindings for complete vector-search indexes and -a C++17 API for both indexes and low-level quantization. +Build, search, save and reload a small IVF index with Python. This guide describes +**0.7.0** and uses synthetic data, so no dataset download is required. +For C++, go directly to [C++ integration](cpp_integration.md). -## Requirements - -| Platform | CPU baseline | Python wheel targets | -| --- | --- | --- | -| Linux x86-64 | AVX2 and FMA | CPython 3.11–3.14 | -| Linux ARM64 (AArch64) | NEON | CPython 3.11–3.14 | -| Windows x86-64 | AVX2 and FMA | CPython 3.11–3.14 | -| macOS 14+ ARM64 (Apple Silicon) | NEON | CPython 3.11–3.14 | - -Source builds require a C++17 compiler, OpenMP, and CMake 3.20 or newer. -Windows uses MSVC (Visual Studio 2026 with the Desktop development with C++ -workload and CMake 4.2+ for its generator). Apple Silicon uses AppleClang and -an external OpenMP runtime such as Homebrew `libomp`. - -Linux ARM64 source builds and repaired wheels run on native AArch64 CI. -Linux ARM64 wheels carry `manylinux_2_27_aarch64` and -`manylinux_2_28_aarch64` tags. - -
-CPU dispatch details - -On x86-64, optional AVX-512 kernels require AVX2, FMA, and AVX-512F/BW/DQ; -MSVC builds also require AVX-512VL/CD. Popcount-specific kernels additionally -require AVX-512 VPOPCNTDQ. Detection checks CPU features and OS support for -the required register state. Missing AVX-512 features select the AVX2 backend. - -On ARM64, the build excludes x86 kernels and uses NEON and portable scalar -implementations. Standard FastScan requires AVX2/FMA, a supported AVX-512 -backend, or NEON; scalar fallbacks do not remove that requirement. - -
- -## Python - -### Install +## Install ```bash -python -m pip install "rabitqlib>=0.6.0" +python -m pip install "rabitqlib>=0.7.0" ``` -Configurable IVF routing and the concurrency contract below require 0.6.0; -see the [upgrade notes](releases/0.6.0.md). For unreleased changes, -[install from a checkout](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/CONTRIBUTING.md#python-changes). +Wheels support CPython 3.11–3.14 on Linux x86-64/ARM64, Windows x86-64 and +macOS 14+ ARM64. x86-64 requires AVX2/FMA; ARM64 uses NEON. +See [installation](installation.md) for platform details and source builds. -Wheels target the platforms above and require no compiler or CMake. -Linux ARM64 and macOS ARM64 wheels bundle OpenMP, so wheel users do not -need a separate OpenMP installation. -Release wheels disable native CPU tuning and select supported kernels at runtime. +Supply data and queries with shape `(count, dim)` in the original dimension. +Python indexes convert them to contiguous float32 and accept `64 <= dim <= 65536`; +they rotate and pad internally. See [input requirements](index/usage.md#input-requirements). -### Build and search an IVF index +## Build and search an IVF index The following complete example uses deterministic synthetic data and does not require a dataset download: @@ -70,6 +36,7 @@ clustering = RaBitQKMeans( ) clustering.train(data) +# IVF requires exactly max_elements rows when building. index = IvfIndex( dim=64, max_elements=len(data), @@ -84,250 +51,83 @@ print(ids.shape, distances.shape) # (5, 10) (5, 10) print(ids[0]) ``` -For raw-vector reranking, use `nbits=32` in the constructor above. The index -copies the original float32 vectors instead of storing extra-bit codes, while -retaining one-bit codes for filtering. Build, search, and save/load use the -same APIs; the original data can be released after construction. - -IVF selects FastScan precision automatically: HACC for 4–9-bit codes, standard -FastScan for 1–3-bit codes and raw vectors. To override that choice: - -```python -ids, distances = index.search(queries, k=10, nprobe=5, high_accuracy=True) -ids, distances = index.search(queries, k=10, nprobe=5, high_accuracy=False) -# Omit high_accuracy, or pass None, to use automatic selection. -``` - -### Add and remove IVF vectors - -Starting with 0.5.1, a built or loaded IVF index can accept new vectors and exclude -existing vectors from search without the original dataset: +## Read the results -```python -new_vectors = rng.standard_normal((50, 64)).astype(np.float32) -new_ids = index.add(new_vectors) # automatic cluster routing; ids 500..549 -removed = index.remove(new_ids[:10]) # returns 10; storage is retained -``` +Both outputs have shape `(number_of_queries, k)`. IDs are zero-based row numbers +in the data supplied to `build`; added vectors receive new IDs without renumbering +existing ones. Results are ordered by increasing distance: `metric="l2"` estimates +squared Euclidean distance, while `metric="ip"` estimates `1 - dot(data, query)`. +Smaller values are better for both metrics. Quantized distances are estimates, +not exact distances to the original vectors. -Batch additions because each `add()` copies the index storage. Removed IDs are -not reused, and both additions and removals survive save/load. See the -[IVF update guide](index/ivf.md#updating-an-index) for costs and limits. +IVF and HNSW may return fewer than `k` valid neighbors, for example after removals +or when the probed IVF clusters contain too few points. Unfilled slots contain +ID `2**32 - 1` and distance `+inf`. Filter them with +`valid = ids != np.iinfo(np.uint32).max` before using IDs to index your data. -The `metric` argument accepts `"l2"` and `"ip"` (also spelled -`"innerproduct"`). To search by cosine similarity, normalize database and -query vectors first and use `metric="ip"`. +The `metric` argument accepts `"l2"` and `"ip"` (also `"innerproduct"`). For +cosine search, normalize database and query vectors first and use `metric="ip"`. -### Add and remove HNSW vectors +## Measure recall -With 0.5.2 or newer, reuse `data`, `queries`, and `clustering` from the IVF -example above to build and update an HNSW index: +For this small L2 example, compute exact neighbors from the original vectors, +then compare several probe counts on the same built index. Recall@10 is the mean +fraction of exact top-10 IDs returned for each query. ```python -from rabitqlib import HnswIndex - -hnsw = HnswIndex( - dim=64, max_elements=len(data), M=16, ef_construction=100, nbits=4, -) -hnsw.build(data, clustering.centroids, clustering.assignments) - -# HNSW needs explicit capacity growth before adding beyond max_elements. -new_vectors = rng.standard_normal((50, 64)).astype(np.float32) -hnsw.resize(hnsw.max_elements + len(new_vectors)) -new_ids = hnsw.add(new_vectors) # automatic nearest-centroid routing -removed = hnsw.remove(new_ids[:10]) # returns 10 -ids, distances = hnsw.search(queries, k=10, ef=100) +k = 10 +# Float64 exhaustive distances are practical for this tiny dataset only. +delta = queries.astype(np.float64)[:, None, :] - data.astype(np.float64)[None, :, :] +exact_distances = np.sum(delta * delta, axis=2) +exact_ids = np.argsort(exact_distances, axis=1, kind="stable")[:, :k] + +for nprobe in (1, 2, 5): + candidate_ids, _ = index.search(queries, k=k, nprobe=nprobe) + recall = np.mean([ + np.intersect1d(found, expected).size / k + for found, expected in zip(candidate_ids, exact_ids) + ]) + print(f"nprobe={nprobe}: Recall@{k}={recall:.3f}") ``` -Removed points remain in the graph and count toward capacity, but never appear -in results. Additions and removals survive save/load; files containing HNSW -removals require 0.5.2 or newer. Do not update an index while it is being searched. -See the [HNSW update guide](index/hnsw.md#updating-an-index) for recall and capacity -considerations. +Missing-result IDs cannot match any ground-truth ID and count as misses. +Probing more clusters searches more candidates, but even probing all five does +not guarantee recall of 1: quantized scoring and filtering remain approximate. -See the [Python examples](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/sample/python) -for IVF, HNSW, and SymphonyQG. For clustering, choose -[RaBitQKMeans](clustering.md#rabitqkmeans) for flat assignment, recommended for small cluster -counts, or [QGKMeans](clustering.md#qgkmeans) for graph assignment. - -### Threading and file paths - -Starting with 0.6.0, the library releases the Python GIL during native index search, -construction, updates, and file I/O. The following contract applies to -`IvfIndex`, `HnswIndex`, and `SymqgIndex`: - -| Operation on the same index | Concurrent access | -| --- | --- | -| `search`, `search_batch`, `save`, and property reads | May run together | -| `build`, `add`, `remove`, and `resize`, where available | Require exclusive access | - -A conflicting call immediately raises -`RuntimeError("Index is busy: conflicting operation in progress")`. -It does not wait or queue an update. Access protection includes Python input -conversion and is released when the operation returns or raises an exception. -Different index objects can operate independently. Concurrent saves must use -different output paths, including any sidecar files. - -Search parameters such as `ef`, `nprobe`, precision, and `num_threads` belong to -each call. When using a Python thread pool, consider `num_threads=1` to avoid -starting multiple native worker pools. Keep input arrays and any shared backing -storage unchanged for the entire call; do not resize or modify them from another -thread. Returned result arrays own their storage. - -These access checks belong to the Python wrappers. C++ callers must synchronize -index updates themselves and use separate search output buffers. - -`num_threads=0` selects the detected available logical CPU count. Positive values -set an upper limit, capped at that count; small workloads may use fewer workers. -Python index methods default to one thread; clustering defaults to `0`. -On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding -before starting the program. - -Index save/load paths are UTF-8 strings on Windows and native path bytes on POSIX -in C++; Python paths are Unicode strings on all platforms. - -### Index file compatibility - -The compatibility suite loads frozen files written by v0.5.2, compares search -IDs and distances, and verifies save/load migration and rejection of truncated -or unsupported versions. See the -[fixture provenance and coverage](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/tests/fixtures/persistence). -Compatibility is format-specific; preserving old readers' ability to open new -files is not guaranteed. Consult each index's format documentation before -downgrading the library. - -### Build the Python bindings from source - -Source builds require Python 3.11 or newer, a C++17 compiler, CMake 3.20 or -newer, and OpenMP. To install the current development version on Ubuntu or -Debian: - -```bash -sudo apt-get update -sudo apt-get install -y build-essential cmake libomp-dev -git clone https://github.com/VectorDB-NTU/RaBitQ-Library.git -cd RaBitQ-Library -python -m pip install . -``` - -On Apple Silicon, use a native ARM64 Python environment. From the repository root: - -```bash -brew install cmake ninja libomp -python -m pip install . -Ccmake.define.OpenMP_ROOT="$(brew --prefix libomp)" -``` +This small example demonstrates the calculation. For tuning, use representative +held-out queries and compare recall with latency; see the [benchmark guide](benchmarks/bench.md). +The exhaustive array above is only suitable for small datasets. IP ground truth +must rank `1 - dot` instead of squared L2. -On Windows, install the build tools listed above, then run `python -m pip install .` -from the repository root. For a portable source build on any supported platform, -also pass `-Ccmake.define.RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`. +## Save and load -## C++ +Continue with the same `index` and `queries`. This example uses a temporary +directory that is removed when the block finishes; use your own persistent path +when keeping an index. -On Linux, clone the repository and build the library and examples: - -```bash -git clone https://github.com/VectorDB-NTU/RaBitQ-Library.git -cd RaBitQ-Library - -cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -cmake --build build --parallel -``` - -CMake enables native CPU tuning by default where supported by the compiler. -For portable binaries within a supported OS and architecture, configure with -`-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`; runtime kernel selection remains active. -See the [platform-specific build commands](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/tests/README.md#quick-start) -for Linux ARM64, Windows, and Apple Silicon. - -### C++ examples - -Example executables are written to `bin/`. Their source demonstrates complete -indexing and querying workflows: - -- [IVF RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/ivf_rabitq_indexing.cpp) -- [HNSW RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/hnsw_rabitq_indexing.cpp) -- [SymphonyQG](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/symqg_indexing.cpp) -- [Low-level quantization](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/quantizer.cpp) - -The low-level quantization example is provided as source and is not currently a -CMake target. For the GIST benchmark workflow, see -[`example.sh`](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/example.sh). - -### Use in another C++ project - -The C++ API and ABI are still evolving. For reproducible builds, pin a release -or commit and include RaBitQ-Library as a Git submodule: - -```bash -git submodule add https://github.com/VectorDB-NTU/RaBitQ-Library.git third_party/rabitqlib -git submodule update --init --recursive -git -C third_party/rabitqlib checkout -git add third_party/rabitqlib -``` - -Add the library and link its namespaced target in your `CMakeLists.txt`: - -```cmake -set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) -add_subdirectory(third_party/rabitqlib) -target_link_libraries(my_program PRIVATE rabitqlib::rabitqlib) -``` - -Update the pinned revision deliberately when adopting upstream changes: - -```bash -git -C third_party/rabitqlib fetch -git -C third_party/rabitqlib checkout -git add third_party/rabitqlib -``` - -### Install the C++ library - -Installation is useful for package managers, container images, and shared server -environments. Disable native optimization for use on other CPUs: - -```bash -cmake -S . -B build \ - -DRABITQ_BUILD_SAMPLES=OFF \ - -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_INSTALL_PREFIX="$HOME/.local" -cmake --build build --parallel -cmake --install build -``` - -Consume the installed package with: - -```cmake -find_package(rabitqlib CONFIG REQUIRED) -target_link_libraries(my_program PRIVATE rabitqlib::rabitqlib) -``` - -For a non-system prefix, point CMake to the installation: - -```bash -cmake -S . -B build -DCMAKE_PREFIX_PATH="$HOME/.local" -cmake --build build --parallel -``` - -Both submodule and installed-package integration require OpenMP on the consuming -system. The [downstream consumer test](https://github.com/VectorDB-NTU/RaBitQ-Library/tree/main/tests/consumer) -provides a complete installed-package example. - -### Run the C++ tests on Linux - -```bash -cmake -S . -B build -DRABITQ_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Release -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF -cmake --build build --parallel -ctest --test-dir build --output-on-failure +```python +from pathlib import Path +from tempfile import TemporaryDirectory + +with TemporaryDirectory() as directory: + path = str(Path(directory) / "ivf.index") + index.save(path) + loaded = IvfIndex.load(path) + loaded_ids, loaded_distances = loaded.search(queries, k=10, nprobe=5) + np.testing.assert_array_equal(loaded_ids, ids) + np.testing.assert_allclose(loaded_distances, distances) ``` -GoogleTest is downloaded during test configuration. +The loaded index owns its data and restores the saved rotation. IVF indexes using +HNSW centroid routing also need the accompanying `.hnsw` file; keep it with the +main index. This small example uses Flat routing and has no sidecar. +See [file compatibility](index/usage.md#index-file-compatibility) before downgrading. ## Next steps -- Learn how the [RaBitQ quantizer](rabitq/rabitq.md) works. -- Select an index: [IVF](index/ivf.md), [HNSW](index/hnsw.md), or - [SymphonyQG](index/qg.md). -- Review the - [contribution workflow](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/CONTRIBUTING.md) - before opening a pull request. +- [IVF](index/ivf.md): raw-vector reranking, precision, centroid routing and updates. +- [HNSW](index/hnsw.md) and [SymphonyQG](index/qg.md): graph-index workflows and parameters. +- [Index usage contracts](index/usage.md): counts, threading, input lifetimes and file paths. +- [Clustering](clustering.md): choose flat RaBitQKMeans or graph-based QGKMeans. +- [C++ integration](cpp_integration.md): builds, installation and component selection. +- [Stable core API](rabitq/core_api.md): integrate quantization into your own storage. diff --git a/docs/docs/rabitq/core_api.md b/docs/docs/rabitq/core_api.md new file mode 100644 index 00000000..a72e5084 --- /dev/null +++ b/docs/docs/rabitq/core_api.md @@ -0,0 +1,282 @@ +# Stable core quantization API v1 + +This interface is introduced in [RaBitQ Library 0.7.0](../releases/0.7.0.md). + +Include `rabitqlib/core/api.hpp` and link `rabitqlib::core`. The stable C++17 source +interface lives in `rabitqlib::quant::v1`; it uses float32 inputs, caller-owned buffers +and separate factor values. The public header does not include Eigen, index containers, +record maps, or ISA-specific interfaces. + +This interface reuses the existing RaBitQ encoder and full-distance formula. It does +not rotate inputs, create a thread pool, or take ownership of application memory. +These entry points execute serially. Core/FastScan can be built without OpenMP using +`RABITQ_BUILD_INDEX=OFF` and `RABITQ_USE_OPENMP=OFF`; the default full build retains it. + +## Minimal example + +The call sequence is `requirements` → `encode` → `prepare_query` → `compute`. +This complete example generates 128-dimensional data and a query with a fixed +seed, applies the same seeded [FHT/Kac rotation](rotator.md) to both, and uses +unpacked four-bit codes with a zero centroid. It compares the estimated squared +L2 distance with the true distance computed from the original vectors. +Rotation is explicit caller-side preprocessing; the v1 API does not perform it. + +```cpp +#include +#include +#include +#include +#include + +#include +#include + +int main() { + namespace api = rabitqlib::quant::v1; + constexpr size_t kDim = 128; + constexpr size_t kTotalBits = 4; // Bits per dimension: 1 sign bit + 3 extra bits. + std::mt19937 generator(42); + std::normal_distribution normal(0.0F, 1.0F); + std::vector input(kDim), query(kDim); + for (size_t i = 0; i < kDim; ++i) { + input[i] = normal(generator); + query[i] = normal(generator); + } + + std::unique_ptr> rotator(rabitqlib::choose_rotator( + kDim, rabitqlib::RotatorType::FhtKacRotator, 0, 42 + )); + std::vector rotated_input(rotator->size()), rotated_query(rotator->size()); + rotator->rotate(input.data(), rotated_input.data()); + rotator->rotate(query.data(), rotated_query.data()); + // The zero centroid stays zero under rotation. + const api::Parameters parameters{rotator->size(), kTotalBits, rabitqlib::METRIC_L2}; + // Each 4-bit code occupies a uint16_t here; SplitPacked stores codes compactly. + const auto layout = api::CodeLayout::UnpackedU16; + const auto sizes = api::requirements(parameters, layout); + + // Allocate during setup, with space to satisfy the reported alignment. + std::vector codes(sizes.code_bytes); + std::vector scratch(sizes.workspace_bytes + sizes.workspace_alignment - 1); + void* workspace = scratch.data(); + size_t available = scratch.size(); + std::align(sizes.workspace_alignment, sizes.workspace_bytes, workspace, available); + + const auto factors = api::encode( + parameters, rotated_input.data(), nullptr, + {layout, {codes.data(), codes.size()}, {}}, + {workspace, sizes.workspace_bytes} + ); + // QueryState borrows rotated_query; keep it alive and unchanged through compute(). + const auto state = api::prepare_query(parameters, rotated_query.data(), nullptr); + const float distance = api::compute( + state, {layout, {codes.data(), codes.size()}, {}}, factors + ); + double true_distance = 0.0; + for (size_t i = 0; i < kDim; ++i) { + const double difference = static_cast(input[i]) - query[i]; + true_distance += difference * difference; + } + std::cout << "Estimated squared L2 distance: " << distance << '\n' + << "True squared L2 distance: " << true_distance << '\n'; +} +``` + +Link this program to `rabitqlib::core`; see +[C++ integration](../cpp_integration.md#core-only-c-integration) for setup. +The application owns all storage. `encode` uses the workspace only during its +call; scoring needs the codes, factors and live query, but not that workspace. +For a nonzero centroid, rotate it with the same rotator before encoding and +query preparation. Fixed seeds make repeated runs reproducible within a given +toolchain; random-distribution implementations and floating-point reductions may +differ across platforms. This single pair illustrates usage, not a quality benchmark. +Reuse allocated storage for subsequent calls with the same requirements. +For packed codes, use `SplitPacked` and provide both buffers described under +[code layouts and sizes](#code-layouts-and-sizes). + +## Stability policy + +`RABITQLIB_QUANT_API_VERSION_MAJOR` is `1` and `RABITQLIB_QUANT_API_VERSION_MINOR` is +`0`. API v1.0 is independent of the package's 0.7.0 release version. + +### What remains stable + +Within API major 1, existing declarations and documented behavior remain +source-compatible. This covers names, signatures, argument order/defaults, +aggregate fields, enum values, supported formats/ranges, code interpretation, +metrics, ownership and allocation guarantees. + +Minor revisions may add functions or separate types, but cannot remove or +repurpose v1 declarations, including deprecated ones, or add aggregate fields. +Breaking changes require a new API major and migration guidance. + +### Compatibility limits + +This is **source compatibility**, not a C ABI or a cross-version C++ binary ABI. Rebuild +consumers when upgrading the library/toolchain. Do not serialize C++ structs or query +state using `sizeof`/`memcpy`. Workspace sizes may change: query them again during setup, +and do not persist or hardcode them. Compiler warnings from future compilers are not +part of the compatibility guarantee. + +The stable surface is the v1 namespace, its version macros, and the existing +`MetricType` / `METRIC_L2` / `METRIC_IP` values used by it. Other core templates, +`detail/`, SIMD entry points, rotator implementations, FastScan and ANN +interfaces are not implicitly frozen by including this header. Existing interfaces +remain available through canonical module headers; old forwarding paths are removed as +described in the [migration guide](../modular_migration.md). Advanced interfaces have +operation-specific allocation and threading behavior; the v1 guarantees below do not +automatically apply to them. + +## Data and numerical contract + +`Parameters` contains: + +| Field | Meaning | +| --- | --- | +| `dim` | Number of coordinates actually supplied, at least 2; no implicit padding | +| `total_bits` | 1 through 9, including the sign bit | +| `metric` | `METRIC_L2` or `METRIC_IP` | +| `t_const` | Finite, at most `INT_MAX / 2`; nonpositive selects the existing exact search, positive selects the existing fast scaling algorithm | + +Inputs and a non-null centroid contain at least `dim` aligned float32 values. Pass +`nullptr` for an all-zero centroid. Data, query and centroid must use the same transformed +coordinate system. If preserving distances from an original coordinate system, use a +compatible metric-preserving rotation/padding pipeline. The quantizer does not infer +or validate the transform. The quality assumptions of RaBitQ still apply; arbitrary +preprocessing does not acquire a quality guarantee merely by satisfying the API types. + +Values must be finite and intermediate computations representable in float32. Use the +usual round-to-nearest floating-point environment. No implicit normalization, NaN repair, +precision-mode change or approximate-search fallback is performed. The exact encoder +keeps its SIMD-certified search and original exact event-sweep fallback. A positive +`t_const` explicitly selects the pre-existing speed/accuracy tradeoff; the API does not +choose it automatically. + +`Factors {add, rescale, error}` are the values returned by the existing full-code encoder. +They may be stored in separate columns, inside user records, or passed as values. No +byte offset relative to the code is prescribed. `error` is the existing full-code error +factor, not a guaranteed deterministic lower bound; `compute` does not use it for pruning. + +`prepare_query` borrows the supplied float query and computes metric/centroid terms. +It does not quantize or copy the query and needs no workspace. The centroid is read +synchronously and not retained. The query must remain alive and unchanged while any +returned/copied `QueryState` is used. Default-constructed state is invalid until assigned +from `prepare_query`. To change a query or centroid, prepare a new state. + +For logical integer codes `u[i]` and `B = total_bits`, the full estimator is: + +```text +L2: g_add = squared_norm(query - centroid) +IP: g_add = -dot(query, centroid) + +distance = factors.add + g_add + + factors.rescale * (dot(query, u) - ((2^B - 1) / 2) * sum(query)) +``` + +L2 results estimate squared Euclidean distance and are not clamped to zero. IP results +estimate `1 - dot(data, query)`, not the raw dot product. Codes, factors and query state +must agree on dimension, bit width, metric and centroid; that metadata is not embedded +in the buffers. This is full float-query estimation even at one bit, not the legacy +four-bit query approximation used by `split_single_estdist` filtering. + +Different layouts/ISAs can use different float reduction orders. Bit-identical distances +across layouts, compilers or CPUs are not promised. Numerical checks must account for +absolute summands and cancellation, rather than using only a relative error against a +possibly near-zero result. Layout conversion itself preserves every logical integer +code exactly and does not change factors. + +## Code layouts and sizes + +Call `requirements(parameters, layout)` during setup. It returns code/extra buffer sizes +and the workspace size/alignment. Capacities may be larger; only required prefixes are +accessed. Code buffers support byte alignment, including unaligned fields in user records. +The workspace must satisfy the reported alignment. + +| Layout | `codes` | `extra` | Dimension requirement | +| --- | --- | --- | --- | +| `UnpackedU16` | `2 * dim` bytes; one little-endian unsigned 16-bit value per coordinate, in `[0, 2^B)` | Unused | Any `dim >= 2` | +| `SplitPacked` | `dim / 8` bytes of sign codes in the existing full-64/final-32 word order | `dim * (B-1) / 8` bytes in the existing extra-code packing; unused at one bit | Multiple of 32 | + +For `SplitPacked`, sign and extra buffers can be anywhere in memory; they need not be +adjacent. No factor/header is appended. Bit/coordinate order matches existing library +pack/unpack kernels. The existing packed convention targets supported little-endian +x86-64/ARM64 systems. `UnpackedU16` byte order is explicitly little-endian, independent +of the alignment of the supplied byte buffer. + +`convert_codes` converts either layout to either layout with the same dimension/bit width. +It does not rotate, requantize, pad, or alter factors. Source and destination must not +overlap, even for the same layout. To use an application-specific layout, consume logical +codes or convert explicitly; the API does not interpret arbitrary external bit fields. +Packing/unpacking should normally happen during ingestion or construction, not for every +candidate distance. Packed scoring uses the existing SIMD kernels; unpacked scoring is +a general float32 reference path and is not promised equal throughput. + +## Ownership, overlap and allocation + +`Buffer` and `ConstBuffer` are non-owning pointer/capacity pairs. Requirements do not +allocate. `encode` and `convert_codes` consume an aligned caller workspace synchronously +and retain no pointers to it. All writable code regions and the workspace must be mutually +disjoint and disjoint from their inputs, using the required ranges rather than surplus +capacity. Read-only input/centroid regions may alias each other. In-place encoding and +conversion are not supported. Unused `extra` buffers are ignored. + +On valid successful calls, **requirements, encode, convert_codes, prepare_query and +compute make no heap allocations**, including their first call, larger dimensions within +the supplied workspace, and exact fallback searches. There is no warmup requirement, +TLS growth, hidden vector allocation, or heap fallback for an undersized workspace. +The guarantee assumes supported standard-library/runtime configurations, not instrumentation +that allocates on every operation. + +Allocate application storage and initialize any rotation/configuration objects before +entering the hot path. Allocation failures during setup and allocation by exception +machinery for invalid calls are outside this guarantee. Other numerical APIs have their own allocation behavior; for example, +`MatrixRotator` uses temporary storage. Callers may supply their own rotation pipeline. + +Concurrent encodes/conversions use separate workspaces and writable outputs. Read-only +codes, factors, parameters and prepared query state can be shared; callers synchronize +mutations and ensure borrowed query lifetimes. + +## Errors + +| Error type | Conditions | +| --- | --- | +| `std::invalid_argument` | Invalid dimensions/bit width/metric/layout/scaling constant; null/short/misaligned buffers; overlapping writable regions; non-finite input/factors; invalid query state or logical code value | +| `std::length_error` | Dimension/index or required byte counts cannot be represented | +| `std::overflow_error` | Required float32 norms, factors or estimated result are not finite | + +Input pointers must reference accessible live storage; capacity claims do not let the +library detect dangling pointers or fictitious allocations. Error categories are stable; +exact diagnostic wording is not. Validation/arithmetic errors leave encoded output and +conversion destination bytes unchanged; workspace contents are unspecified after any call. + +## Example and verification + +The standalone example (`sample/cpp/stable_core` in the same source checkout) +uses only `api.hpp`, application-owned records/factor columns, a 65-coordinate working +space and an external transform. Its sign/permutation transform demonstrates integration; +it is not a representative random-rotation quality benchmark. + +```bash +cmake -S sample/cpp/stable_core -B build-stable-example \ + -DCMAKE_PREFIX_PATH="" +cmake --build build-stable-example +ctest --test-dir build-stable-example --output-on-failure +``` + +The frozen v1 consumer (`tests/fixtures/core_api_v1/consumer.cpp`) is compiled +against source-tree and installed packages in CI. The internal workspace resource +uses only caller memory with no upstream allocator. A separate test checks the +FHT rotator's first `rotate` call after construction for allocations. + +The unit suite covers bits 1–9, both metrics, dimensions around packing boundaries, +non-padded logical codes, zero residuals, buffer guards, code/factor parity, legacy packed +estimation and independent workspaces on concurrent callers. Separate allocation test +processes exercise normal dispatch and forced scalar fallback without warming the API. +They intercept C++ allocation on all tested platforms and also C allocation entry points +on Linux. Existing index, persistence and Python tests protect the shared implementation. + +Zero allocation does not imply zero call overhead. The +[scoring performance reference](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/benchmarks/core/README.md#scoring-call-overhead) +explains the call implementation, comparison protocol and numerical checks for hot, +sequential and random code access. Scoring throughput does not predict end-to-end +ANN query performance. diff --git a/docs/docs/rabitq/estimator.md b/docs/docs/rabitq/estimator.md index 1c0d6601..b06a80eb 100644 --- a/docs/docs/rabitq/estimator.md +++ b/docs/docs/rabitq/estimator.md @@ -1,68 +1,97 @@ - # Estimator -This part presents how to use the quantized vectors to estimate Euclidean distances and inner product with minimal efforts. + +Estimate squared L2 or inner-product distance from quantized codes and stored +factors. For API usage, see [Stable core API](core_api.md) or the advanced +[Quantizer examples](quantizer.md). | Notation | Description | |----------|----------------------------------------| -| $\mathbf{o}_r$, $\mathbf{q}_r$ | The raw data vector and the raw query vector. | -| $\mathbf{c}$ | The center vector. | -| $D$ | The dimension of vectors. | -| $B$ | The number of bits per dimension used for the quantization. | +| $\mathbf{o}_r$, $\mathbf{q}_r$ | The original data and query vectors, zero-padded to the working dimension. | +| $\mathbf{c}$ | The centroid in the same coordinate system as the original vectors. | +| $D$ | The working (padded) dimension, $D>1$. | +| $B$ | Total bits per dimension, including the sign bit. | | $c_B$| A constant value $c_B=- \frac{2^B-1}{2}$. | | $\mathbf{1}_D$ | The all-one vector of dimension $D$. | -| $P$ | The random rotation matrix. | -| $\mathbf{\bar o}$ | The quantized normalized data vector. | -| $\mathbf{q}_r'$ | The reversely rotated raw query vector: $\mathbf{q}_r'=P^{-1}\mathbf{q}_r$. | +| $P$ | Orthogonal map from the quantized coordinate system back to the original padded domain. | +| $\mathbf{o}$, $\mathbf{q}$ | Normalized residuals: $(\mathbf{o}_r-\mathbf{c})/\|\mathbf{o}_r-\mathbf{c}\|$ and $(\mathbf{q}_r-\mathbf{c})/\|\mathbf{q}_r-\mathbf{c}\|$. | +| $\mathbf{\bar o}$ | Unit-length quantized residual direction, expressed in the original padded domain. | +| $\mathbf{q}_r'$ | The query in the quantized coordinate system: $\mathbf{q}_r'=P^{-1}\mathbf{q}_r$. | | $S_q$ | $S_q=\sum_{i=1}^{D} \mathbf{q}'_{r}[i].$ | -| $\Delta_x$ | The rescaling factor of a quantized vector. | -| $\mathbf{x}_u$ | The uint representation of the quantized vectors.| +| $\Delta_x$ | Residual-norm scale defined below; distinct from the estimator factor `F_rescale`. | +| $\mathbf{x}_u$ | Unsigned integer codes.| + +## Quantization + +The derivation assumes nonzero residual norms. Zero residuals are handled by +the encoder; callers need not normalize residuals themselves. + +Here $P^{-1}$ denotes the forward preprocessing transform applied by +[the rotator](rotator.md). Apply that same transform to data, centroids and +queries; no separate inverse-rotation call is needed for queries. +For unsigned codes $\mathbf{x}_u$, define -## Quantization -Recall that the quantization algorithm of RaBitQ receives a raw data vector $\mathbf{o}_r$, a center vector $\mathbf{c}$ and a sample of random rotation matrices $P$ as inputs and returns a value $\Delta_x$ and an uint vector $\mathbf{x}_u$ as output, such that +$$ +\mathbf{z}=\mathbf{x}_u+c_B\mathbf{1}_D,\qquad +\mathbf{\bar o}=\frac{P\mathbf{z}}{\|\mathbf{z}\|},\qquad +\Delta_x=\frac{\|\mathbf{o}_r-\mathbf{c}\|}{\|\mathbf{z}\|}. +$$ + +Then the norm-preserving reconstruction satisfies $$ \begin{align} -\mathbf{o}_r - \mathbf{c} \approx \| \mathbf{o}_r-\mathbf{c} \| \cdot \mathbf{\bar o} = \Delta_x P(\mathbf{x}_u+c_B\mathbf{1}_D) +\mathbf{o}_r - \mathbf{c} \approx \| \mathbf{o}_r-\mathbf{c} \| \cdot \mathbf{\bar o} = \Delta_x P(\mathbf{x}_u+c_B\mathbf{1}_D) \end{align} $$ -where $c_B=- \frac{2^B-1}{2}$ and $B$ is the number of bits used for the quantization. $\mathbf{1}_D$ is the all-one vector of dimension $D$. $P$ is a sample of random rotation matrices. - ## Estimator + The following derivation covers squared Euclidean distance and inner product. The public metric choices are `METRIC_L2` and `METRIC_IP`; there is no separate cosine metric. For cosine search, normalize both data and query vectors, then use inner product. The implementation represents inner-product distance as -$1-\langle\mathbf{o},\mathbf{q}\rangle$, which preserves the same ranking while +$1-\langle\mathbf{o}_r,\mathbf{q}_r\rangle$, which preserves the same ranking while keeping the nearest-neighbor convention that smaller values are better. -### Estimator of Euclidean Distance +Here $\epsilon_0$ controls the width of the probabilistic error interval. The current +encoder uses the internal constant `kConstEpsilon = 1.9`; it is not a public +configuration parameter. These intervals support approximate pruning, not a +deterministic guarantee for every vector or query. Validate recall for the target workload. +The expressions below describe the data-code error with an unquantized query +under the random-rotation model. They do not include floating-point roundoff or +additional query/LUT quantization error. The one-bit probabilistic bound is +derived in [RaBitQ, Section 3.2](https://arxiv.org/pdf/2405.12497). +The default [FHT/Kac rotation](rotator.md) efficiently approximates the random +orthogonal rotation, following the Fast JLT implementation approach in SymphonyQG. +For one-bit estimation, query/LUT quantization error is typically negligible +relative to the data-code error. For higher-bit batch estimation, +[high-accuracy FastScan](../kernel_ip.md#batch-code) reduces the LUT error +carried into refinement. + +### Estimator of Euclidean distance + $$ \begin{align} -&\| \mathbf{o}_r-\mathbf{q}_r\|^2 +&\| \mathbf{o}_r-\mathbf{q}_r\|^2 \\=& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2 \left< \mathbf{o}_r - \mathbf{c}, \mathbf{q}_r - \mathbf{c}\right> \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \left<\mathbf{ o},\mathbf{q} \right> \\ \approx& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{\left<\mathbf{\bar o},\mathbf{q} \right>}{ \left<\mathbf{\bar o},\mathbf{o} \right>} \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 - 2\| \mathbf{o}_r-\mathbf{c}\| \cdot \frac{\left<\mathbf{\bar o},\mathbf{q}_r-\mathbf{c} \right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} \\ =& \|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left< \| \mathbf{o}_r-\mathbf{c}\|\cdot P^{-1}\mathbf{\bar o}, \mathbf{q}_r' \right> \\=&\|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] -\\ \| \mathbf{o}_r-\mathbf{q}_r\|^2 \approx&\|\mathbf{o}_r -\mathbf{c}\|^2 + \|\mathbf{q}_r -\mathbf{c}\|^2 +2\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{2\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \end{align} $$ -The error bound is given by +The probabilistic error interval is $$ \begin{align} -% \pm \ 2 \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{c_{error}}{\sqrt{D}\cdot 2^{B-1}}\\ \pm \ 2 \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}} \end{align} $$ -Here $\epsilon_0$ is a tunable parameter which controls the confidence level of the error bound. By default, $\epsilon_0=1.9$ guarantees nearly perfect confidence. - -We store the following variables such that the estimator can be computed easily. +Store the following factors (`F_` for data and `G_` for the query). | Name (Type) of Variable | Description | | ----------------- | ----------- | @@ -73,13 +102,11 @@ We store the following variables such that the estimator can be computed easily. | `G_kBxSumq (float)`| $c_BS_q$| | `G_error (float)`| $\| \mathbf{q}_r-\mathbf{c}\|$| - - +### Estimator of inner product -### Estimator of Inner Product When inner product is used as the metric of vector search, it targets the data vector with the *maximum* inner product with the query. The implementation -uses $1-\langle\mathbf{o},\mathbf{q}\rangle$ so that smaller distances rank +uses $1-\langle\mathbf{o}_r,\mathbf{q}_r\rangle$ so that smaller distances rank higher-inner-product vectors first. The constant 1 does not affect ordering. $$ @@ -88,24 +115,18 @@ $$ \\=&1 -\left< \mathbf{o}_r-\mathbf{c} + \mathbf{c},\mathbf{q}_r-\mathbf{c} + \mathbf{c}\right> \\=&1 -\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> - \left< \mathbf{o}_r-\mathbf{c},\mathbf{q}_r-\mathbf{c} \right> \\ \approx &1 -\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> +\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] -\\ 1-\left< \mathbf{o}_r,\mathbf{q}_r\right>\approx &1 -\left< \mathbf{q}_r,\mathbf{c}\right> -\left< \mathbf{o}_r-\mathbf{c},\mathbf{c}\right> +\| \mathbf{o}_r-\mathbf{c}\| \frac{\left< \mathbf{\bar o}, \mathbf{c}\right>}{\left<\mathbf{\bar o},\mathbf{o} \right>} - \frac{\Delta_x}{\left<\mathbf{\bar o},\mathbf{o} \right>}\cdot \left[ \left< \mathbf{x}_u, \mathbf{q}_r' \right>+c_B S_q\right] \end{align} $$ -The error bound is given by +The probabilistic error interval is $$ \begin{align} -% \pm \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \frac{c_{error}}{\sqrt{D}\cdot 2^{B-1}} \pm \|\mathbf{o}_r-\mathbf{c}\|\cdot \| \mathbf{q}_r-\mathbf{c}\| \cdot \sqrt{\frac{1 - \left< \mathbf{\bar o},\mathbf{o}\right>^2}{\left< \mathbf{\bar o},\mathbf{o}\right>^2}} \frac{\epsilon_{0}}{\sqrt{D-1}} \end{align} $$ - - -Here $\epsilon_0$ is a tunable parameter which controls the confidence level of the error bound. By default, $\epsilon_0=1.9$ guarantees nearly perfect confidence. - -We store the following variables such that the estimator can be computed easily. The variables `F_` indicates the factors for the data vector. The variables `G_` indicates the factors for the query vector. +Store the following factors (`F_` for data and `G_` for the query). | Name (Type) of Variable | Description | | ----------------- | ----------- | @@ -116,40 +137,53 @@ We store the following variables such that the estimator can be computed easily. | `G_kBxSumq (float)`| $c_BS_q$| | `G_k1xSumq (float)`| $c_1S_q$| | `G_error (float)`| $\| \mathbf{q}_r-\mathbf{c}\|$| - ## Implementation -### Distance Estimation -Based on the above derivation, we can implement the estimator of Euclidean distance and inner product using exactly the same implementation. A pseudo code is given below. Let `ip` be the inner product between the binary code and the randomly rotated query vector. + +### Distance estimation + +Both metrics use the same scoring formula with metric-specific factors. +Here `ip` is the inner product between the logical +unsigned $B$-bit code and the rotated query vector; it is a binary code only +when $B=1$. ```cpp // Compute the estimated distance -// Note that G_add is dependent on the center vector. -float est_dist = F_add + G_add + F_rescale * (ip + G_kBxSumq) +// Note that G_add is dependent on the centroid. +float est_dist = F_add + G_add + F_rescale * (ip + G_kBxSumq); // Compute the error bound -// Note that G_error is dependent on the center vector. -float error_bound = F_error * G_error +// Note that G_error is dependent on the centroid. +float error_bound = F_error * G_error; // Compute the lower and upper bounds of the estimated distance -float lb_dist = est_dist - error_bound -float ub_dist = est_dist + error_bound +float lb_dist = est_dist - error_bound; +float ub_dist = est_dist + error_bound; ``` +### Incremental distance estimation -### Incremental Distance Estimation - -RaBitQ supports incremental distance estimation. We split the code into two parts: the binary code (the most significant bits) and the extended code (the remaining $B-1$ bits). -Incremental distance estimation supports to first estimate a coarse distance based on the binary code. If the accuracy is insufficient, we then access the extended code to boost the accuracy. -For this, we prepare factors for both parts. The binary layout stores `F_add`, +Split codes separate the sign bit from the remaining $B-1$ extra bits. +Score the sign code first, then refine selected candidates with the extra bits. +Each part has its own factors. The binary layout stores `F_add`, `F_rescale`, and `F_error`; the extended layout stores `F_add_ex` and -`F_rescale_ex`. The implementation obtains the refined error bound by scaling -the binary error factor by $2^{B-1}$ rather than storing a separate -`F_error_ex`. -The factors for the query includes `G_add`, `G_error`, `G_kBxSumq` and `G_k1xSumq`. -Let `ip_bin` be the inner product between the binary code and the randomly rotated query vector and `ip_ex` be the inner product between the ex-code and the randomly rotated query vector. - +`F_rescale_ex`. For $B>1$, the logical codes satisfy +$\mathbf{x}_u=2^{B-1}\mathbf{x}_{\mathrm{bin}}+\mathbf{x}_{\mathrm{ex}}$, +which gives the refined distance formula below. + +The split-code implementation uses the binary error factor divided by +$2^{B-1}$ as its refined pruning margin. This is an approximation, not an +identity for the full-code error factor in the tables above: that factor depends +on the full-code direction's own $\langle\mathbf{\bar o},\mathbf{o}\rangle$. +The split layout does not store a separate `F_error_ex`. +The query factors are `G_add`, `G_error`, `G_kBxSumq` and `G_k1xSumq`. +Here `ip_bin` and `ip_ex` are the inner products of the sign and extra-bit +codes, respectively, with the rotated query. +These equations use exact code/query inner products. Binary filtering may instead +approximate `ip_bin` through four-bit query codes or FastScan lookup tables; +see [inner-product kernels](../kernel_ip.md). The margins do not separately bound +that additional error. ```cpp // 1-bit dist diff --git a/docs/docs/rabitq/quantizer.md b/docs/docs/rabitq/quantizer.md index 02a7d98b..9843ab60 100644 --- a/docs/docs/rabitq/quantizer.md +++ b/docs/docs/rabitq/quantizer.md @@ -1,66 +1,68 @@ # Quantizer -RaBitQLib includes two versions of implementations for the RaBitQ algorithm and designs various data formats to ease practical deployment. Specifically, the two implementaions offer different trade-offs as follows. - -1. Optimal accuracy with longer quantization time. -2. Nearly-optimal accuracy with significantly fast quantization. - -Various advanced data formats are provided to support the following needs. - -1. Drop-in replacement for uniform scalar quantization. -2. Efficient distance estimation for single vectors. -3. Efficient distance estimation for batched vectors. -4. Incremental distance estimation for splitted single vectors. -5. Incremental distance estimation for splitted batched vectors. - -These data formats map raw floating-point vectors into `uint8` or `uint32` +For a new per-vector quantization integration, start with the +[stable core API v1](core_api.md#minimal-example). This page is an advanced +reference for the existing encoding interfaces and five data formats below. + +| Your task | Start here | +| --- | --- | +| Encode and score individual vectors in your own storage | [Stable core API v1](core_api.md): caller-owned buffers, separate factors and a source-stability contract | +| Score batches with FastScan | [FastScan integration](../cpp_integration.md#fastscan-integration) and the batch formats below | +| Integrate existing scalar, split-code or record layouts | [Data format reference](#data-format) below; these advanced interfaces are outside the v1 stability guarantee | +| Build a complete vector-search index | [Quick start](../quick_start.md) | + +These data formats map floating-point vectors into `uint8` or `uint32` codes. For the packed extended-code layouts, see -`rabitqlib/quantization/pack_excode.hpp`. +`rabitqlib/core/pack_excode.hpp`. The examples below focus on layouts and omit rotation. In an index, rotate data, centroids, and queries with the same rotator first and use its padded dimension. -Query snippets use `rabitqlib/index/estimator.hpp`; the full-vector example also -needs `` and ``. - -RaBitQ quantizer is included in `rabitq_impl.hpp` and `rabitq.hpp`. -```css -. -├── rabitqlib -│ ├── ... -│ └── quantization -│ ├── ... -│ ├── rabitq_impl.hpp -│ └── rabitq.hpp -└── ... -``` +Single-vector query snippets use `rabitqlib/core/estimator.hpp` and link +`rabitqlib::core`; batch-query snippets use `rabitqlib/fastscan/estimator.hpp` and link +`rabitqlib::fastscan`. Both batch formats use `fastscan/encoder.hpp` and +`fastscan/data_layout.hpp`; their 32-vector packing is in `fastscan/pack_codes.hpp`. +The old `quantization/` and generic `index/` forwarding paths have been removed. +See the [migration guide](../modular_migration.md) and +[modular C++ integration](../cpp_integration.md#core-only-c-integration). ## Implementations -Let $B$ be the bit-width for each dimension. Both implementations of RaBitQ include two steps. -1. Compute a **binary code** by recording the sign of every coordinate. -2. Compute an **ex-code** of $B-1$ bits (when $B>1$). +Let $B$ be the total bits per dimension. Encoding produces two parts: -The binary code is easily computed by the function `one_bit_code` in `rabitq_impl.hpp`. +1. A **sign code**: one bit per coordinate. +2. An **extra-bit code**: the remaining $B-1$ bits when $B>1$. -The computation of ex-codes includes two versions of implementation. +The internal sign encoder is in `core/detail/rabitq_impl.hpp`. Use the public +entry points in `core/encoder.hpp` or `fastscan/encoder.hpp`; implementation +headers are not a stable integration surface. -In the first implementation, we compute the ex-codes of RaBitQ based on the algorithm described in the RaBitQ [SIGMOD paper](https://doi.org/10.1145/3725413) (Section 3.2.2). For a vector, to minimize the quantization error, the algorithm tries many different rescaling factors. For each rescaling factor, it rescales the vector and performs rounding (i.e., scalar quantization) to generate a quantization code. Then it finds out the factor and codes which minimizes the quantization error. Note that in the library, the range of enumeration is approriately shrinked, which brings better efficiency without affecting the accuracy. +Choose between two extra-code encoding paths: -In the second implementation, instead of enumerating different rescaling factors, it directly rounds every vector based on the **expected optimal factor**. Specifically, recall that all data vectors are randomly rotated before quantization. The expected optimal factor is computed as follows. We sample several random vectors which follow uniform distribution on the unit sphere and use the first implementation to quantize them. We record the optimal factor for each and take the average of the optimal factors as the expected optimal factor. This implementation introduces some accuracy decrease while significantly speeds up the quantization. +The default path searches the configured rescaling interval +for the best quantized direction, following the RaBitQ [SIGMOD paper](https://doi.org/10.1145/3725413) +(Section 3.2.2). The current implementation can use SIMD to evaluate and prune +intervals; a scalar event sweep remains the reference and fallback when the SIMD +search cannot certify a winner. This path preserves the reference tie and rounding rules. +The fast path reuses a calibrated scaling factor instead of searching per vector. +`faster_config` samples uniformly distributed unit vectors, finds their scaling +factors with the default path, and averages them. Calibrate once during setup +and reuse the result for the same padded dimension and total bit width. This +reduces encoding work at a possible accuracy cost; evaluate it on your data. +## Data format -## Data Format +### Format 1: Scalar codes -### Format 1 - Drop-in Replacement of Uniform Scalar Quantization +Use scalar codes when your application manages its own storage. The encoder +returns unpacked integer codes, a scale and an offset; packing is separate. -This format allows RaBitQ to be used as a direct replacement for uniform scalar quantization, which offers higher accuracy under the same bit-width. The improvement of accuracy is significant when the bit-width is small ($B < 6$). ```cpp #include #include #include -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/encoder.hpp" int main() { size_t dim = 768; @@ -90,27 +92,28 @@ int main() { ); // Note that with this interface, the codes are not compactly stored. - // To compactly store the codes, please refer to - // `rabitqlib/quantization/pack_excode.hpp` + // For packed storage, see + // `rabitqlib/core/pack_excode.hpp` return 0; } ``` - - -### Format 2 - Distance Estimation for a Single Vector - +### Format 2: Single-vector scoring #### Indexing -This format is designed for computing distance metrics between data vectors and query vectors efficiently. It includes precomputed factors to ease similarity calculations. +Encode one vector into unpacked integer codes and three distance-estimation +factors. The application stores the codes and factors separately or together. + ```cpp #include +#include +#include #include #include -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/encoder.hpp" int main() { size_t dim = 768; @@ -161,24 +164,25 @@ int main() { ); // Note that with this interface, the codes are not compactly stored. - // To compactly store the codes, please refer to - // `rabitqlib/quantization/pack_excode.hpp` + // For packed storage, see + // `rabitqlib/core/pack_excode.hpp` return 0; } ``` #### Querying -After quantization, we can use the pre-computed quantization codes and factors to get estimated distance for a given query.Here, we separately stored factors and quantization codes, and user may choose to compact them all together to get a improve space locality (e.g., implementation in our index). Also, we did not rotate the vectors in this example for simplicity of illustration. + +Score a query using the encoded vector and its factors. +Insert the following before `return 0` in the indexing example above. + ```cpp -... - ...... // we omit the quantization code here std::vector query(dim); for (size_t i = 0; i < dim; ++i) { query[i] = dist(gen); } - // please refer to estimator.md for defination of these factors + // See the Estimation guide for these query factors. float c_1 = -static_cast((1 << 1) - 1) / 2.F; float k1xsumq = c_1 * std::accumulate(query.begin(), query.end(), 0.F); float g_add = rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim); @@ -201,23 +205,23 @@ After quantization, we can use the pre-computed quantization codes and factors t ``` - -### Format 3 - Distance Estimation for a Batch of Vectors (for QG) -A variant of this data format is used in QG. +### Format 3: Batch scoring (QG) #### Indexing -This format is designed for scenarios where distances need to be computed between a query vector and multiple data vectors (quantized into 1-bit per dimension) simultaneously, providing significant performance improvements with [FastScan](https://arxiv.org/abs/1704.07355) for batch processing. -In practical implementation of our SymphonyQG index, the data is compactly stored. Please refer to `QGBatchDataMap` in `rabitqlib/quantization/data_layout.hpp` for detailed information. - +Encode up to 32 vectors for batched one-bit distance estimation with +[FastScan](https://arxiv.org/abs/1704.07355). +SymphonyQG and clustering use the two-factor `QGBatchDataMap` layout from +`rabitqlib/fastscan/data_layout.hpp`. Its encoder lives in +`rabitqlib/fastscan/encoder.hpp`; both are owned by the FastScan component. ```cpp #include #include #include -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" int main() { size_t dim = 768; @@ -248,7 +252,8 @@ int main() { ``` #### Querying -During querying, a query is pre-processed as follows. Then FastScan can be called to estimate distance batch by batch. The detailed implementation is in `rabitqlib/index/estimator.hpp`. + +Prepare the query lookup tables once, then score each encoded batch. Insert the following before `return 0` in the indexing example, reusing its initialized `batch_data`, `centroid`, `dim`, and `batch_size`. @@ -267,10 +272,8 @@ initialized `batch_data`, `centroid`, `dim`, and `batch_size`. std::vector est_distance(batch_size); // store the estimated distances - // We suggest users to customize the kernel if some outputs are not needed. - // Eg., QG does not need error bounds, thus we only get the estimated distance here. - // If users want to maintain the information of error bound, please refer to our - // implementation of IVF index. + // QG scoring returns estimates only. Use the IVF batch format below + // when error intervals are needed. rabitqlib::qg_batch_estdist( batch_data.data(), processed_query, dim, est_distance.data() @@ -278,19 +281,21 @@ initialized `batch_data`, `centroid`, `dim`, and `batch_size`. ``` -### Format 4 - Incremental Distance Estimation for Split Single Vectors (for HNSW) -This format can be used in HNSW. - +### Format 4: Split single-vector scoring (HNSW) #### Indexing -This format supports computing distances incrementally when vectors are split across multiple memory locations or when only partial vector information is available at a time. +Store a vector as separate sign and extra-bit codes, with factors for each. +Score the sign code first, then refine selected candidates using the extra bits. +HNSW uses this layout. ```cpp +#include #include #include -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" int main() { size_t dim = 768; @@ -341,19 +346,21 @@ int main() { #### Querying +Insert this before `return 0` in the indexing example, reusing its data and dimensions. + ```cpp -... - size_t dim = 768; // the dimensionality - size_t bits = 5; // the bit-width of DATA vectors std::vector query(dim); + for (auto& value : query) { + value = dist(gen); + } // Queries use four bits, independently of the data code width. - rabitqlib::quant::RabitqConfig config = rabitqlib::quant::faster_config( + rabitqlib::quant::RabitqConfig query_config = rabitqlib::quant::faster_config( dim, rabitqlib::SplitSingleQuery::kNumBits ); rabitqlib::SplitSingleQuery processed_query( - query.data(), dim, bits - 1, config, rabitqlib::METRIC_L2 + query.data(), dim, bits - 1, query_config, rabitqlib::METRIC_L2 ); // This standalone example omits rotation, so query and centroid are both raw. @@ -400,17 +407,21 @@ int main() { ); ``` -### Format 5 - Incremental Distance Estimation for Split Batched Vectors (for IVF) -This data format is used in IVF. +### Format 5: Split batch scoring (IVF) #### Indexing -This format combines the benefits of batched processing with incremental computation, allowing efficient distance estimation when dealing with large collections of split vectors. + +Batch the sign codes for FastScan and store each vector's extra-bit codes +separately. IVF uses this layout to filter a batch before refining candidates. ```cpp +#include #include #include -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" int main() { size_t dim = 768; @@ -468,10 +479,13 @@ int main() { #### Querying +Insert this before `return 0` in the indexing example, reusing its buffers and dimensions. + ```cpp - size_t dim = 768; // the dimensionality - size_t bits = 5; // the bit-width of DATA vectors std::vector query(dim); + for (auto& value : query) { + value = dist(gen); + } // The flag use_hacc controls the precision of FastScan. // `use_hacc = false` - each number in LUTs is quantized into 8 bits. @@ -490,16 +504,12 @@ int main() { std::sqrt(rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim)); processed_query.set_g_add(distance_norm); - size_t batch_size = 32; - std::vector est_distance(batch_size); // store the estimated distances std::vector low_distance(batch_size); // store the lower bound std::vector ip_x0_qr(batch_size ); // store the intermediate result of inner product - // We suggest users to customize the kernel if some outputs are not needed. - // For example, QG does not need error bounds, see `qg_batch_estdist` in - // `rabitqlib/index/estimator.hpp` for details. + // For estimates without error intervals, see the QG batch format above. rabitqlib::split_batch_estdist( batch_data.data(), @@ -515,7 +525,8 @@ int main() { auto ip_func = rabitqlib::select_excode_ipfunc(bits - 1); size_t i = 15; - rabitqlib::split_distance_boosting( + // Replace this candidate's one-bit estimate with the refined distance. + est_distance[i] = rabitqlib::split_distance_boosting( ex_data.data() + (i * rabitqlib::ExDataMap::data_bytes(dim, bits - 1)), ip_func, processed_query, @@ -525,3 +536,61 @@ int main() { ); ``` + +## Batch API contracts + +Batch query preparation and scoring require a positive dimension divisible by 16. +Index pipelines use multiples of 32, which also satisfy the packed extra-code +requirements. Use the same working dimension for encoding, query preparation and +scoring; none of these batch calls rotates or pads an input vector for you. + +Single-batch encoders accept 0 through 32 logical vectors. Zero vectors leave the +output unchanged; more than 32 throws `std::invalid_argument` before writing. +For larger inputs, iterate in chunks of `fastscan::kBatchSize`, as shown in +`sample/cpp/fastscan` in the same checkout. Every batch, including a tail, needs +full `BatchDataMap` or `QGBatchDataMap` storage and 32 score outputs. Read only +its logical results. Buffers belong to the caller; encoders retain no input pointers. +`quantize_one_batch` uses OpenMP when enabled and otherwise runs serially. +Allocation exceptions propagate to the caller, including from worker threads; a +failed encoding may leave partial output, so retry or discard that batch. + +`BatchQuery` and QG encoding/scoring support `float` and `double`. LUT preparation +uses the same scalar type as the query; float uses SIMD dispatch, while double +uses generic arithmetic without narrowing to float. Byte LUTs still quantize the +query, so double does not make estimates exact. Split batch scoring and +`SplitSingleQuery` use float32 query paths; the latter does not support +`SplitSingleQuery`. + +A default `BatchQuery` must be prepared with `reset` before scoring. Its reset +scratch needs **`padded_dim * 4` elements of `T`**, not bytes. Reset borrows query, +scratch and optional precomputed correction only for the call; it owns the resulting +LUT and clears `g_add`. Reuse the scratch between calls, then call `set_g_add` for +the current centroid. `BatchQuery::set_g_add` takes squared L2 distance, or +`1 - dot(query, centroid)` for IP. + +`SplitBatchQuery` owns its LUT but borrows the rotated query for later refinement; +keep that vector alive and unchanged. Its `set_g_add` takes the **L2 norm** of +`query - centroid` for either metric, plus `dot(query, centroid)` for IP. All operands +must use the same rotated, padded domain. The standalone FastScan example demonstrates +both metrics with a nonzero centroid. + +For repeated queries, default-construct `SplitBatchQuery` and use `reset`. +LUT scratch needs `padded_dim * 4` elements of `T`; HACC also needs that many +`uint16_t` elements (pass `nullptr` in standard mode). Scratch is call-only and +must not overlap other buffers. The query remains borrowed until the next successful +reset or destruction. `ex_bits` must be in `[0, 8]`. + +Only first preparation or capacity growth may allocate. Success clears `g_add` +and `g_error`; call `set_g_add` again before scoring. Invalid arguments or allocation +failure preserve the previous state. + +Single-vector estimators also have named-result functions (`estimate_split_single`, +`estimate_split_single_full`, `estimate_xy_base`, and `estimate_xy_full`) returning +`distance`, `lower_bound`, and `inner_product` fields instead of taking three output +references. Both forms use the same numerical implementation. + +Scoring rejects unprepared queries, mismatched dimensions and mismatched HACC modes +with `std::invalid_argument` before touching output. Pass `prepared.padded_dim()` +and, for split scoring, `prepared.use_hacc()` to reuse the prepared shape. The batch +buffer must still contain codes encoded for that shape: raw buffers carry no metadata +or capacity, so the caller must allocate them using the corresponding data map. diff --git a/docs/docs/rabitq/rabitq.md b/docs/docs/rabitq/rabitq.md index 06206ec0..47c5c355 100644 --- a/docs/docs/rabitq/rabitq.md +++ b/docs/docs/rabitq/rabitq.md @@ -1,28 +1,22 @@ # RaBitQ -The RaBitQ algorithm is a drop-in replacement of binary quantization and (uniform) scalar quantization, with its [1-bit version](https://arxiv.org/abs/2405.12497) (released in May 2024) and [multi-bit version](https://doi.org/10.1145/3725413) (published at SIGMOD 2025), respectively. - - - - - -The key advantages of RaBitQ include - -- **High Accuracy with Tiny Space** - RaBitQ provides the state-of-the-art estimates under diverse space budgets. Its smallest representation uses a one-bit code per padded dimension plus per-vector factors. -- **Fast Distance Estimation** - The implementation uses bitwise operations and, for batched scans in IVF and SymphonyQG, [FastScan](https://arxiv.org/abs/1704.07355). -- **Theoretical Error Bound** - RaBitQ provides an asymptotically optimal error bound for the estimation of distances and inner product. The error bound can be used for reliable ordering and [reranking](reranking.md). - - -## Workflow - -The RaBitQ algorithm includes two steps: - -1. **Random Rotation** - Sample a random rotation and apply it to all vectors (including the raw data vectors, the center vector and the raw query vectors). See [Rotator](rotator.md) for more details. - -2. **Quantization** - After the random rotation, the quantization algorithm quantizes a vector of floating-point numbers into a vector of low-bit unsigned integers. See [Quantizer](quantizer.md) for more details. - - -After quantization, the library supports Euclidean distance and inner product -using the code vector $\mathbf{x}_u$ and its factors. Cosine similarity is -obtained by normalizing data and query vectors before using inner-product -search. See [Estimator](estimator.md) for details. +RaBitQ is an alternative to binary and uniform scalar quantization. Its +[one-bit](https://arxiv.org/abs/2405.12497) and +[multi-bit](https://doi.org/10.1145/3725413) variants encode each padded coordinate +with a low-bit unsigned integer and store per-vector factors for distance estimation. + +## Workflow + +1. **Rotate.** Apply the same sampled transformation to data, centroids and queries. + See [Rotation](rotator.md). +2. **Encode.** Quantize each rotated data vector relative to its centroid. + See [Quantization](quantizer.md). +3. **Estimate.** Combine the codes, factors and rotated query to estimate squared + L2 or inner-product distance (`1 - dot`). See [Estimation](estimator.md). + +For cosine search, normalize data and queries before rotation and use inner product. +The probabilistic error intervals can guide [reranking](reranking.md), but do not +make approximate search exact. Recall depends on the data and search settings. + +For integration into your own storage, start with the [stable core API](core_api.md). +For a complete index, start with the [Quick start](../quick_start.md). diff --git a/docs/docs/rabitq/reranking.md b/docs/docs/rabitq/reranking.md index a8732f72..0a5b3452 100644 --- a/docs/docs/rabitq/reranking.md +++ b/docs/docs/rabitq/reranking.md @@ -1,28 +1,45 @@ # Reranking -Reranking is a technique widely adopted for improving recall of vector search. During searching, ANN algorithms usually -1. shortlist a set of candidates based on an index and their quantization codes (e.g., **in memory**); -2. retrieve the raw vectors (e.g., **from disks**) for those candidates which have the smallest estimated distances; -3. computes the exact distances for the candidates to find the nearest neighbors. +Reranking refines shortlisted candidates by computing distances from their +original vectors. Those vectors may be kept in memory or retrieved from disk. -RaBitQ has a unique advantage in reranking due to its theoretical error bound. It can skip reranking a candidate if the lower bound of its estimated distance is larger than the upper bound of the distance of the nearest neighbors. +RaBitQ error intervals can guide approximate reranking: a candidate can be +skipped when its estimated lower bound exceeds the relevant upper bound among +the current nearest-neighbor candidates. These are probabilistic intervals, not +deterministic bounds for every query. -Based on error bounds, it is possible to rerank **fewer than $K$ vectors** to achieve nearly perfect recall - it only reranks the vectors on the boundaries of KNNs. +The strategy may rerank fewer than $K$ vectors when only boundary candidates +need refinement. It does not guarantee exact top-k membership or exact distances +for results that were never reranked; measure the recall/cost trade-off for the workload. This is an application-level strategy, not a built-in disk-retrieval API. HNSW stores quantized vectors, so applications must retain raw vectors separately for exact reranking. [IVF raw mode](../index/ivf.md#index-construction) (`nbits=32`) instead owns the original vectors and reranks them in memory. -## Algorithm Description -Let $K$ be the number of nearest neighbors we target. After receiving the candidates and their estimated distances from an index, e.g., HNSW RaBitQ, we perform the following strategy of reranking to minimize the number of retrieved raw vectors from disks. - -1. Sort the candidates with respect to their estimated distances. -2. Initialize a max-heap (sorted by **upper bounds of distances**) `KNNs` with the top-$K$ candidates which have the smallest estimated distances. -3. Enumerate the remaining candidates. For a new candidate A, - 1. [Condition 1 - A's lower bound $>$ the maximum upper bound in `KNNs`.] - 1. Drop the new candidate and move on. - 2. [Condition 2 - A's lower bound $\le$ the maximum upper bound in `KNNs`, and the top candidate in `KNNs` has been reranked.] - 1. Rerank the new candidate, update `KNNs` and move on. - 3. [Condition 3 - A's lower bound $\le$ the maximum upper bound in `KNNs`, and the top candidate in `KNNs` has NOT been reranked.] - 1. Rerank the top candidate in `KNNs`, update `KNNs` and repeat the procedure for candidate A. +## Algorithm description + +Given at least `K` candidates with distance estimates and error intervals, keep +`K` candidates in a max-heap ordered by their upper bounds. The heap root is the +candidate with the **largest upper bound**, not the nearest candidate. +Refining a candidate computes its exact distance and sets both bounds to that value. + +```text +Sort candidates by estimated distance. +Initialize heap H with the first K candidates, keyed by upper bound. +For each remaining candidate A: + Repeat: + W = candidate with the largest upper bound in H + If A.lower_bound > W.upper_bound: + Skip A and continue with the next candidate. + If W has already been refined: + Refine A. + Insert A into H, retaining the K smallest upper bounds. + Continue with the next candidate. + Refine W and restore heap order. +``` + +Candidates remaining in the heap may still have estimated distances. If the +application needs exact distances for every returned candidate, refine those +candidates too. This does not recover candidates excluded earlier by approximate +search or probabilistic pruning. diff --git a/docs/docs/rabitq/rotator.md b/docs/docs/rabitq/rotator.md index 7d91c6e3..5bc005a6 100644 --- a/docs/docs/rabitq/rotator.md +++ b/docs/docs/rabitq/rotator.md @@ -1,12 +1,18 @@ -# Random Rotation -Random rotation (i.e., Johnson Lindenstrauss Transformation) is a crucial step to ensure robust performance and theoretical error bounds of RaBitQ. It is applied to all vectors (including raw data vectors, center vectors and raw query vectors) as a preprocessing step. This section describes the usage of the random rotation. - -RaBitQLib provides two types of random rotation. All implementations sample and store a random rotation at first. Then they apply the sampled random rotation to every input vector and return the rotated vector. - -By default, the library uses the `FFHT + Kac’s Walk` method. - -The default rotator accepts input dimensions from 64 through 65,536 inclusive, -including 16,384. It pads to a multiple of 32 and uses AVX or NEON kernels +# Random rotation +Random rotation is the preprocessing step used by RaBitQ's distance estimators. +Apply the same sampled transformation to data, centroids and queries. +RaBitQ Library provides FHT/Kac rotation (the default) and a dense random +orthogonal transformation. + +The default FHT/Kac transform is an efficient structured approximation to a +random orthogonal rotation. [SymphonyQG, Section 3.1.4](https://dl.acm.org/doi/abs/10.1145/3709730) +describes using a Fast Hadamard-based Fast Johnson–Lindenstrauss Transform to +reduce rotation work from $O(D^2)$ to $O(D\log D)$. This library combines +Hadamard transforms, random sign flips and the Givens steps described below. +The transform preserves norms; its sampling distribution approximates the +uniform random rotation used in the RaBitQ analysis. + +The default FHT/Kac rotator supports `float` and accepts input dimensions from 64 through 65,536 inclusive. It pads to a multiple of 32 and uses AVX or NEON kernels for power-of-two blocks through `2^16`. Allocate the output using `rotator->size()`, which can exceed the input dimension. @@ -20,24 +26,14 @@ When restoring a standalone rotator saved with earlier 64-dimension padding, pass that original padded dimension to `choose_rotator` before loading its state. The saved sign bytes do not contain dimension metadata. -The implementation can be found in `rotator.hpp`. - -```css -. -├── rabitqlib -│ ├── ... -│ └── utils -│ ├── ... -│ └── rotator.hpp -└── ... -``` +The rotation interfaces are in `rabitqlib/core/rotator.hpp`. ### Example ```cpp #include #include -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/core/rotator.hpp" int main() { const size_t dim = 769; // deliberately not a multiple of 32 @@ -55,12 +51,16 @@ Use the same rotator for data, centroids, and queries. Pass `rotator->size()` to quantization and estimation routines that operate on rotated vectors. For a random orthogonal transformation, select `RotatorType::MatrixRotator`. -It keeps the input dimension by default; pass a padded dimension explicitly if -subsequent packed-code operations require a multiple of 32. Its storage and -computation cost are quadratic in the dimension. - -## FFHT + Kac’s Walk -### Description +It supports `float` and `double` and keeps the input dimension by default. Pass a +padded dimension explicitly if subsequent packed-code operations require a multiple +of 32. Stored space and per-vector work are proportional to `dim * padded_dim` +(quadratic when both dimensions are equal). +Constructing the random matrix also performs QR factorization. The `seed` argument +of `choose_rotator` controls FHT/Kac signs only; MatrixRotator samples its own matrix. +Save and restore the actual rotator state rather than relying on that seed to +recreate a matrix rotation. + +## FFHT + Kac’s walk On x86-64, transforms use the shared AVX intrinsics in `src/simd/fht_kernels.hpp` within the AVX2 or AVX-512 backend selected at runtime. GCC, Clang, and MSVC use the same implementation without inline assembly or MASM. ARM64 uses @@ -68,27 +68,30 @@ NEON kernels in `src/simd/rotator_neon.cpp`; a portable scalar implementation is also available. The backends preserve the FFHT butterfly order, normalization, four sign-flip passes, and saved rotation state. -This method is a combination of the well-known Fast Johnson-Lindenstrauss Transformation algorithms based on [Fast Hadamard Transform](https://www.cs.princeton.edu/~chazelle/pubs/FJLT-sicomp09.pdf) and ideas in [Kac’s Walk](https://projecteuclid.org/journals/annals-of-applied-probability/volume-27/issue-1/Kacs-walk-on-n-sphere-mixes-in-nlog-n-steps/10.1214/16-AAP1214.full). -It first samples 4 sequences of random signs (i.e., Rademacher random variables). Then for each vector, it repeats the following procedures 4 times. +This method is a combination of the well-known Fast Johnson-Lindenstrauss Transformation algorithms based on [Fast Hadamard Transform](https://www.cs.princeton.edu/~chazelle/pubs/FJLT-sicomp09.pdf) and ideas in [Kac’s Walk](https://projecteuclid.org/journals/annals-of-applied-probability/volume-27/issue-1/Kacs-walk-on-n-sphere-mixes-in-nlog-n-steps/10.1214/16-AAP1214.full). +It first samples 4 sequences of random signs (i.e., Rademacher random variables). +Let `trunc_dim` be the largest power of two not exceeding the original input +dimension (`dim`), and `padded_dim` the working dimension. -1. Flip its coordinates with the $i$-th sequence of sampled random signs. -2. Apply FFHT on the first/last $2^k$ coordinates (alternately), where $2^k$ is the maximum power of 2 that is less than or equal to the dimensionality of the vector. -3. Apply Givens rotation with a fixed angle $\theta = \frac{\pi}{4}$ to the 1st and the 2nd halves of coordinates. +When `trunc_dim == padded_dim`, each of the four rounds only flips signs and +applies normalized FFHT to the whole vector; no Givens/Kac step is performed. +Otherwise, each round performs the following steps: -The following table summarizes the space and time complexity of this method. +1. Flip its coordinates with the $i$-th sequence of sampled random signs. +2. Apply normalized FFHT on the first/last `trunc_dim` coordinates (alternately) of the padded vector. +3. Apply Givens rotation with a fixed angle $\theta = \frac{\pi}{4}$ to the 1st and the 2nd halves of coordinates. | Space Consumption | Time Complexity | | ----------------- | --------------- | | $4D$ binary values ($4D$ bits) | $O(D\log D)$ | -This implementation is based on the [FFHT library](https://github.com/FALCONN-LIB/FFHT) developed by Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn and Ludwig Schmidt. - - -## Random Orthogonal Transformation -### Description -This method is the classical Johnson-Lindenstrauss Transformation. It first samples a random gaussian matrix and orthogonalizes it with QR decomposition. Then it multiplies the matrix to every vector. +This implementation is based on the [FFHT library](https://github.com/FALCONN-LIB/FFHT) developed by Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn and Ludwig Schmidt. -The following table summarizes the space and time complexity of this method. +## Random orthogonal transformation +This method samples a Gaussian matrix and orthogonalizes it with QR decomposition. +It then multiplies each vector by the stored transform. The table describes stored +space and per-vector rotation work when input and output dimensions are both $D$; +QR construction is separate setup work. | Space Consumption | Time Complexity | | ----------------- | --------------- | diff --git a/docs/docs/releases/0.6.0.md b/docs/docs/releases/0.6.0.md index f49cbc5b..c8303887 100644 --- a/docs/docs/releases/0.6.0.md +++ b/docs/docs/releases/0.6.0.md @@ -18,7 +18,7 @@ Keep input arrays unchanged until each call returns. HNSW's Python `ef` remains optional; SymphonyQG's remains required. With a Python thread pool, consider `num_threads=1` to avoid multiplying native worker pools. Concurrent saves need distinct paths, including sidecars. See the -[threading contract](../quick_start.md#threading-and-file-paths). +[threading contract](../index/usage.md#threading-and-file-paths). C++ callers must synchronize mutations themselves. SymphonyQG adds `search_batch_with_ef()` for per-call windows; existing search APIs remain available. @@ -55,7 +55,7 @@ Tests cover frozen v0.5.2 IVF, HNSW, and raw/quantized SymphonyQG files under L2 and inner product, including queries, resave/load migration, and malformed files. This coverage does not guarantee compatibility with every historical format. Keep original files if you may need to downgrade. See the -[compatibility overview](../quick_start.md#index-file-compatibility). +[compatibility overview](../index/usage.md#index-file-compatibility). ### C++ and CMake integration diff --git a/docs/docs/releases/0.7.0.md b/docs/docs/releases/0.7.0.md new file mode 100644 index 00000000..01181901 --- /dev/null +++ b/docs/docs/releases/0.7.0.md @@ -0,0 +1,39 @@ +# RaBitQ Library 0.7.0 + +This release introduces modular C++ components, stable core API v1, optional +OpenMP for numerical components, and unified index counts. + +## Upgrading from 0.6.0 + +### C++ components and headers + +Choose `rabitqlib::core` for per-vector quantization, estimation and rotation, +`rabitqlib::fastscan` for batch scoring, or `rabitqlib::index` for ANN and clustering. +The former `rabitqlib::rabitqlib` target and forwarding headers are removed. +Update includes and link targets using the [migration guide](../modular_migration.md). +Rebuild C++ consumers with matching headers and libraries; no cross-version +C++ binary ABI is promised. + +Core and FastScan support `RABITQ_USE_OPENMP=OFF`. ANN and Python retain their +OpenMP requirement. Use a fresh installation prefix when changing components +so stale headers from an older installation cannot mask migration errors. + +### Stable core API v1 + +The new [core API](../rabitq/core_api.md) provides a source-stable float32 interface +with caller-owned code buffers and workspace, separate factors, and no heap +allocation on valid calls. API version 1.0 is independent of package version 0.7.0. +Advanced templates and index APIs are not covered by this stability guarantee. + +### Index counts in C++ and Python + +IVF, HNSW and SymphonyQG share `max_elements`, `num_stored_points` and `num_points` +for capacity, stored records and live points respectively. **HNSW's `num_points` +now excludes removed records.** Use `num_stored_points` for ID bounds, append +offsets and minimum resize capacity. See the +[count migration details](../modular_migration.md#unified-index-counts). + +Python continues to use `import rabitqlib` and the existing index class names. +Saved index formats are retained from 0.6.0; loading recovers counts from existing +removal markers. These changes do not require index-file conversion. Restrictions +for readers older than 0.6.0 still apply; see the [0.6.0 notes](0.6.0.md). diff --git a/docs/docs/stylesheets/extra.css b/docs/docs/stylesheets/extra.css index f4a384c7..24424fad 100644 --- a/docs/docs/stylesheets/extra.css +++ b/docs/docs/stylesheets/extra.css @@ -1,74 +1,263 @@ +/* Local system fonts, including SF on Apple platforms. */ :root { - --rabitq-gradient-start: #4051b5; - --rabitq-gradient-end: #7e57c2; + --md-text-font: -apple-system, BlinkMacSystemFont, "Segoe UI", Helvetica, Arial, sans-serif; + --md-code-font: "SFMono-Regular", Consolas, "Liberation Mono", monospace; } -.hero { - margin: 0 0 2rem; - padding: 2.5rem; - color: #fff; - background: linear-gradient( - 135deg, - var(--rabitq-gradient-start), - var(--rabitq-gradient-end) - ); - border-radius: 0.8rem; - box-shadow: 0 0.5rem 1.5rem rgb(32 45 94 / 20%); +[data-md-color-scheme="default"] { + color-scheme: light; + --md-default-bg-color: #fff; + --md-default-fg-color: #1d1d1f; + --md-default-fg-color--light: #626268; + --md-default-fg-color--lighter: #86868b; + --md-default-fg-color--lightest: #e5e5e9; + --md-primary-fg-color: #1d1d1f; + --md-primary-bg-color: #fff; + --md-accent-fg-color: #0055aa; + --md-typeset-a-color: #0066cc; + --md-code-bg-color: #f5f5f7; + --md-code-fg-color: #28282c; + --md-footer-bg-color: #f5f5f7; + --md-footer-bg-color--dark: #f5f5f7; + --md-footer-fg-color: #1d1d1f; + --md-footer-fg-color--light: #626268; + --md-footer-fg-color--lighter: #626268; + --rq-surface: #f5f5f7; + --rq-line: #dedee3; + --rq-header: rgb(255 255 255 / 94%); + --rq-active: #eaf3fc; + --rq-button: #0066cc; + --rq-button-hover: #0055aa; } -.hero > :first-child { - margin-top: 0; +[data-md-color-scheme="slate"] { + color-scheme: dark; + --md-default-bg-color: #161617; + --md-default-fg-color: #f5f5f7; + --md-default-fg-color--light: #b0b0b7; + --md-default-fg-color--lighter: #86868b; + --md-default-fg-color--lightest: #37373b; + --md-primary-fg-color: #f5f5f7; + --md-primary-bg-color: #161617; + --md-accent-fg-color: #8ec5ff; + --md-typeset-a-color: #70b5ff; + --md-code-bg-color: #232325; + --md-code-fg-color: #e8e8ed; + --md-footer-bg-color: #1d1d1f; + --md-footer-bg-color--dark: #1d1d1f; + --md-footer-fg-color: #f5f5f7; + --md-footer-fg-color--light: #b0b0b7; + --md-footer-fg-color--lighter: #b0b0b7; + --rq-surface: #232325; + --rq-line: #37373b; + --rq-header: rgb(22 22 23 / 94%); + --rq-active: #203044; + --rq-button: #0071e3; + --rq-button-hover: #0066cc; } -.hero .md-button { - margin-top: 0.5rem; - margin-right: 0.35rem; - color: #fff; - border-color: rgb(255 255 255 / 70%); +body { -webkit-font-smoothing: antialiased; } + +.md-header { + color: var(--md-default-fg-color); + background: var(--rq-header); + border-bottom: 1px solid var(--rq-line); + box-shadow: none; } -.hero .md-button:hover, -.hero .md-button:focus { - color: var(--md-primary-fg-color); - background-color: #fff; - border-color: #fff; +.md-header__title { font-size: .85rem; letter-spacing: -.025em; } +.md-header__button.md-logo { color: var(--md-typeset-a-color); } +.md-grid { max-width: 68rem; } +.md-main__inner { margin-top: 2rem; } +.md-typeset { font-size: .8rem; line-height: 1.7; } + +.md-typeset h1, +.md-typeset h2, +.md-typeset h3 { + color: var(--md-default-fg-color); + font-weight: 600; + letter-spacing: -.02em; + line-height: 1.25; + overflow-wrap: break-word; } -.hero .md-button--primary { - color: var(--md-primary-fg-color); - background-color: #fff; - border-color: #fff; +.md-typeset h1 { margin-bottom: .85em; font-size: 2em; } +.md-typeset h2 { + margin: 2em 0 .7em; + font-size: 1.4em; } +.md-typeset h3 { margin-top: 1.7em; font-size: 1.15em; } +.md-typeset a { text-underline-offset: .18em; } +.md-typeset a:hover { text-decoration: underline; } +:where(a, button, input, label[tabindex]):focus-visible { + outline: 2px solid var(--md-typeset-a-color); + outline-offset: 3px; +} + +.md-typeset .headerlink { font-weight: 400; } +.md-typeset code { border-radius: .25rem; font-size: .82em; } +.md-typeset pre > code, +.md-typeset .highlight pre > code { + padding: .9rem 1rem; + border-radius: .4rem; + line-height: 1.7; +} +.md-typeset .highlight { margin: 1.2em 0; } +.md-typeset :is(.md-clipboard, .md-code__button) { color: var(--md-default-fg-color--light); } +.md-typeset :is(.md-clipboard, .md-code__button):is(:hover, :focus-visible) { color: var(--md-typeset-a-color); } + +.md-typeset__table { display: block; } +.md-typeset table:not([class]) { + display: table; + width: 100%; + border: 1px solid var(--rq-line); + border-radius: .4rem; + font-size: .72rem; + box-shadow: none; + overflow: hidden; +} +.md-typeset table:not([class]) th { + min-width: 0; + background: var(--rq-surface); + font-weight: 600; +} +.md-typeset table:not([class]) :is(th, td) { padding: .65rem .75rem; } +.md-typeset blockquote { + padding-left: 1rem; + color: var(--md-default-fg-color--light); + border-left: 2px solid var(--rq-line); +} +.md-typeset :is(.admonition, details) { border-radius: .4rem; box-shadow: none; } + +.md-typeset .md-button { + padding: .55em 1.3em; + border: 1px solid var(--rq-line); + border-radius: .4rem; + font-weight: 500; + transition: background-color 150ms, border-color 150ms; +} +.md-typeset .md-button--primary { + color: #fff; + background: var(--rq-button); + border-color: var(--rq-button); +} +.md-typeset .md-button:is(:hover, :focus-visible) { + color: #fff; + background: var(--rq-button-hover); + border-color: var(--rq-button-hover); + text-decoration: none; +} + +.md-nav { font-size: .69rem; line-height: 1.5; } +.md-nav__title { color: var(--md-default-fg-color); font-weight: 600; box-shadow: none; } +.md-nav--primary > .md-nav__list > .md-nav__item > .md-nav__link { font-weight: 600; } +.md-nav__item .md-nav__link--active { color: var(--md-typeset-a-color); font-weight: 600; } +.md-nav--secondary { + color: var(--md-default-fg-color--light); + font-size: .65rem; + border-left: 1px solid var(--rq-line); +} +.md-nav--secondary .md-nav__title { background: transparent; } +.md-footer { margin-top: 2rem; border-top: 1px solid var(--rq-line); } +.md-footer__title { font-size: .9rem; } + +/* Homepage introduction; the content below remains ordinary documentation. */ +.hero { + margin: 0 0 1.8rem; + padding: .6rem 0 1.8rem; + border-bottom: 1px solid var(--rq-line); +} +.hero .eyebrow { + margin: 0 0 1rem; + color: var(--md-typeset-a-color); + font-size: .68rem; + font-weight: 600; +} +.hero h1 { + margin: 0 0 .7rem; + font-size: clamp(1.8rem, 3.4vw, 2.6rem); + letter-spacing: -.035em; + line-height: 1.15; +} +.hero .hero-description { + max-width: 31rem; + color: var(--md-default-fg-color--light); + font-size: .9rem; + line-height: 1.65; +} +.hero .hero-actions { + display: flex; + flex-wrap: wrap; + align-items: center; + gap: .6rem 1.2rem; + margin: 1.4rem 0; +} +.hero .hero-actions p { display: contents; } +.hero .hero-platforms { margin: 0; color: var(--md-default-fg-color--light); font-size: .68rem; } .feature-grid { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); - gap: 1rem; - margin: 1.5rem 0 2.5rem; + gap: 1.8rem; + margin: 0 0 2.4rem; +} +.feature-card { min-width: 0; } +.feature-card h2 { + margin: 0 0 .5rem; + padding: 0; + border: 0; + font-size: .82rem; + letter-spacing: -.02em; } +.feature-card p { margin: 0; color: var(--md-default-fg-color--light); font-size: .72rem; line-height: 1.65; } -.feature-card { - padding: 1.2rem; - border: 0.05rem solid var(--md-default-fg-color--lightest); - border-radius: 0.6rem; +@media screen and (min-width: 60em) { + .md-search__form { + color: var(--md-default-fg-color); + background: var(--rq-surface); + border: 1px solid var(--rq-line); + border-radius: .5rem; + } + .md-search__input::placeholder, + .md-search__input + .md-search__icon { color: var(--md-default-fg-color--light); } } -.feature-card h2 { - margin-top: 0; - font-size: 1rem; +@media screen and (min-width: 76.25em) { + .md-sidebar--primary { padding-right: .6rem; } + .md-nav--primary .md-nav__link { + padding: .2rem .5rem; + margin-top: .05rem; + border-radius: .35rem; + } + .md-nav--primary .md-nav__link--active { background: var(--rq-active); } + .md-nav--primary .md-nav__item--section > .md-nav__link { + margin-top: .8rem; + color: var(--md-default-fg-color); + font-weight: 600; + } } -.feature-card p { - margin-bottom: 0; - color: var(--md-default-fg-color--light); +@media screen and (max-width: 76.2344em) { + .md-nav--primary .md-nav__title, + .md-nav--primary .md-nav__title[for="__drawer"], + .md-nav--primary .md-nav__source { + color: var(--md-default-fg-color); + background: var(--rq-surface); + } + .md-nav--primary .md-nav__source { border-bottom: 1px solid var(--rq-line); } } @media screen and (max-width: 44.9844em) { - .hero { - padding: 1.5rem; - } + .md-main__inner { margin-top: 1rem; } + .md-typeset h1 { font-size: 1.8em; } + .hero { padding-top: .4rem; padding-bottom: 1.4rem; } + .hero h1 { font-size: clamp(1.65rem, 7.5vw, 2.4rem); } + .hero .hero-description { font-size: .86rem; } + .feature-grid { grid-template-columns: 1fr; gap: 1rem; } + .md-typeset pre > code, + .md-typeset .highlight pre > code { padding: .8rem; } +} - .feature-grid { - grid-template-columns: 1fr; - } +@media (prefers-reduced-motion: reduce) { + .md-typeset .md-button { transition: none; } } diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index f4168c7a..8d2d6707 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -8,19 +8,25 @@ copyright: Copyright © VectorDB Group, Nanyang Technological University nav: - Home: index.md - - Quick Start: quick_start.md - - Release Notes: + - Quick start: quick_start.md + - Installation: installation.md + - C++ Integration: cpp_integration.md + - Release notes: + - "0.7.0": releases/0.7.0.md - "0.6.0": releases/0.6.0.md - RaBitQ: - Overview: rabitq/rabitq.md - Rotation: rabitq/rotator.md - Quantization: rabitq/quantizer.md + - Stable core API: rabitq/core_api.md + - Modular migration: modular_migration.md - Estimation: rabitq/estimator.md - Reranking: rabitq/reranking.md - Kernels: - - Compact Code Storage: compact_code.md - - Inner Product: kernel_ip.md - - Vector Search: + - Compact code storage: compact_code.md + - Inner product: kernel_ip.md + - Vector search: + - Shared contracts: index/usage.md - IVF RaBitQ: index/ivf.md - HNSW RaBitQ: index/hnsw.md - SymphonyQG: index/qg.md @@ -29,28 +35,29 @@ nav: - Clustering: clustering.md - Benchmarks: - Overview: benchmarks/bench.md - - IVF Indexes: benchmarks/ivf.md - - Graph Indexes: benchmarks/graphs.md + - IVF indexes: benchmarks/ivf.md + - Graph indexes: benchmarks/graphs.md - K-means: benchmarks/kmeans.md theme: name: material language: en + font: false icon: logo: material/vector-polygon repo: fontawesome/brands/github palette: - media: "(prefers-color-scheme: light)" scheme: default - primary: indigo - accent: deep purple + primary: custom + accent: custom toggle: icon: material/brightness-7 name: Switch to dark mode - media: "(prefers-color-scheme: dark)" scheme: slate - primary: indigo - accent: deep purple + primary: custom + accent: custom toggle: icon: material/brightness-4 name: Switch to light mode diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index b90d7db8..12ab4b4b 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -20,10 +20,10 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" -#include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -304,15 +304,17 @@ class LloydKMeans : public Parameters { std::vector cluster_sizes(k); std::vector sums(centroids.size()); const size_t update_workers = std::min(k, static_cast(threads)); - std::vector> update_point_bins(update_workers); + std::vector update_partitions(update_workers); std::vector update_worker_for_cluster(k); for (size_t worker = 0; worker < update_workers; ++worker) { - const size_t first_cluster = k * worker / update_workers; - const size_t last_cluster = k * (worker + 1) / update_workers; - for (size_t cluster = first_cluster; cluster < last_cluster; ++cluster) { + auto& partition = update_partitions[worker]; + partition.first_cluster = k * worker / update_workers; + partition.last_cluster = k * (worker + 1) / update_workers; + for (size_t cluster = partition.first_cluster; cluster < partition.last_cluster; + ++cluster) { update_worker_for_cluster[cluster] = worker; } - update_point_bins[worker].reserve(n / update_workers); + partition.points.reserve(n / update_workers); } double previous_obj = std::numeric_limits::infinity(); @@ -323,10 +325,8 @@ class LloydKMeans : public Parameters { previous_centroids.swap(centroids); const double obj = std::accumulate(distances.begin(), distances.end(), 0.0); - std::fill(sums.begin(), sums.end(), 0.0); - std::fill(cluster_sizes.begin(), cluster_sizes.end(), size_t{0}); update_centroids( - x, n, sums, cluster_sizes, update_point_bins, update_worker_for_cluster + x, n, sums, cluster_sizes, update_partitions, update_worker_for_cluster ); const auto [nsplit, shift] = consolidate_centroids(x, n, sums, cluster_sizes, previous_centroids); @@ -379,6 +379,12 @@ class LloydKMeans : public Parameters { double final_obj = 0.0; private: + struct UpdatePartition { + size_t first_cluster; + size_t last_cluster; // Exclusive; also bounds the accumulators we clear. + std::vector points; + }; + void validate_parameters() const { if (d == 0 || k == 0 || niter == 0) { throw std::invalid_argument( @@ -429,27 +435,39 @@ class LloydKMeans : public Parameters { size_t n, std::vector& sums, std::vector& cluster_sizes, - std::vector>& point_bins, + std::vector& partitions, const std::vector& worker_for_cluster ) { - for (auto& bin : point_bins) { - bin.clear(); + for (auto& partition : partitions) { + partition.points.clear(); } // Keep point order within each cluster, preserving the original sums. for (size_t point = 0; point < n; ++point) { - point_bins[worker_for_cluster[assignments[point]]].push_back(point); + partitions[worker_for_cluster[assignments[point]]].points.push_back(point); } const uint32_t threads = effective_num_threads(); static_cast(threads); #pragma omp parallel for num_threads(threads) schedule(static) for (std::ptrdiff_t loop_worker = 0; - loop_worker < static_cast(point_bins.size()); + loop_worker < static_cast(partitions.size()); ++loop_worker) { - const size_t worker = static_cast(loop_worker); + const auto& partition = partitions[static_cast(loop_worker)]; + // The same interval defines point ownership and accumulator clearing, + // including empty clusters, independently of the physical thread team. + const size_t first_cluster = partition.first_cluster; + const size_t last_cluster = partition.last_cluster; + std::fill_n( + sums.data() + first_cluster * d, (last_cluster - first_cluster) * d, 0.0 + ); + std::fill_n( + cluster_sizes.data() + first_cluster, + last_cluster - first_cluster, + size_t{0} + ); simd::accumulate_cluster_sums( x, - point_bins[worker].data(), - point_bins[worker].size(), + partition.points.data(), + partition.points.size(), d, assignments.data(), sums.data(), @@ -507,6 +525,16 @@ class LloydKMeans : public Parameters { return {0, shift}; } + const auto replace_empty_centroid = [&](size_t cluster, size_t point) { + std::copy_n(x + point * d, d, centroids.data() + cluster * d); + if (spherical) { + normalize_centroid(centroids.data() + cluster * d); + } + --cluster_sizes[assignments[point]]; + cluster_sizes[cluster] = 1; + shift += cluster_shift(cluster); + }; + // Reuse the distance buffer after recording the pre-update objective. // Keep populated means intact and give unused centroids poorly fitted // training points. Unlike coordinate perturbations, this also works at zero. @@ -550,13 +578,7 @@ class LloydKMeans : public Parameters { std::find(cluster_sizes.begin(), cluster_sizes.end(), size_t{0}) - cluster_sizes.begin() ); - std::copy_n(x + point * d, d, centroids.data() + cluster * d); - if (spherical) { - normalize_centroid(centroids.data() + cluster * d); - } - --cluster_sizes[assignments[point]]; - cluster_sizes[cluster] = 1; - shift += cluster_shift(cluster); + replace_empty_centroid(cluster, point); return {1, shift}; } @@ -580,13 +602,7 @@ class LloydKMeans : public Parameters { std::pop_heap(candidates.begin(), candidates.end(), less); candidates.pop_back(); } while (cluster_sizes[assignments[point]] <= 1); - std::copy_n(x + point * d, d, centroids.data() + cluster * d); - if (spherical) { - normalize_centroid(centroids.data() + cluster * d); - } - --cluster_sizes[assignments[point]]; - cluster_sizes[cluster] = 1; - shift += cluster_shift(cluster); + replace_empty_centroid(cluster, point); } return {nsplit, shift}; } diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index 88b96e3e..3882b3f4 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -12,17 +12,18 @@ #include #include -#include "rabitqlib/clustering/detail/kmeans.hpp" -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/query.hpp" +#include "rabitqlib/clustering/detail/kmeans.hpp" // IWYU pragma: export +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/symqg/detail/search_scratch.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" -#include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" -#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::qgkmeans { @@ -61,6 +62,63 @@ struct QGKMeansParameters { uint32_t quantization_bits = 0; }; +namespace detail { + +// One training call owns the storage; the public assign adapter can instead +// borrow caller storage for one batch. State is read-only while workers run. +class QuerySumCache { + enum class State { Disabled, Pending, Ready }; + State state_ = State::Disabled; + std::vector storage_; + float* values_ = nullptr; + + public: + QuerySumCache(size_t padded_dim, const QGKMeansParameters& parameters) { + if (padded_dim >= 1024 && + (parameters.niter > 1 || + parameters.final_assignment != FinalAssignmentMode::Exact)) { + state_ = State::Pending; + } + } + + QuerySumCache(float* values, bool ready) : values_(values) { + if (values != nullptr) { + state_ = ready ? State::Ready : State::Pending; + } + } + + QuerySumCache(const QuerySumCache&) = delete; + QuerySumCache& operator=(const QuerySumCache&) = delete; + + // Called after centroid initialization releases its point permutation. + // Training keeps n and the input/rotation fixed throughout this lifetime. + void prepare(size_t n) { + if (state_ == State::Pending && values_ == nullptr) { + storage_.resize(n); + values_ = storage_.data(); + } + } + + [[nodiscard]] const float* cached(size_t i) const { + return state_ == State::Ready ? values_ + i : nullptr; + } + + void store(size_t i, float value) const { + if (state_ == State::Pending) { + values_[i] = value; + } + } + + // Commit only after every worker and the assignment's state update succeed. + void complete() { + if (state_ == State::Pending) { + state_ = State::Ready; + } + } +}; + +} // namespace detail + /** * Builds a SymphonyQG over a fixed number of centroids and assigns vectors to * them. The graph allocation and per-thread search scratch are reused between @@ -69,23 +127,6 @@ struct QGKMeansParameters { class QGAssigner { friend class detail::QGAssignment; - private: - struct ThreadScratch { - std::vector rotated_query; - std::vector estimated_distances; - std::vector lookup_table; - BatchQuery batch_query; - buffer::SearchBuffer search_pool; - buffer::SearchBuffer result_pool; - - ThreadScratch(size_t padded_dim, size_t M, size_t ef_search) - : rotated_query(padded_dim) - , estimated_distances(M) - , lookup_table(padded_dim << 2) - , search_pool(ef_search) - , result_pool(1) {} - }; - public: QGAssigner( size_t d, @@ -117,9 +158,9 @@ class QGAssigner { scratch_.reserve(num_threads_); for (uint32_t thread = 0; thread < num_threads_; ++thread) { - scratch_.emplace_back( - std::make_unique(graph_->padded_dim(), M_, ef_search_) - ); + scratch_.emplace_back(std::make_unique( + graph_->padded_dim(), M_, ef_search_, 1 + )); } } @@ -133,6 +174,21 @@ class QGAssigner { float* distances, float* query_sums = nullptr, bool sums_ready = false + ) { + detail::QuerySumCache cache(query_sums, sums_ready); + assign_with_cache(centroids, x, n, labels, distances, cache); + } + + void clear_hints() { previous_labels_.clear(); } + + private: + void assign_with_cache( + const float* centroids, + const float* x, + size_t n, + PID* labels, + float* distances, + detail::QuerySumCache& query_sums ) { if (centroids == nullptr || x == nullptr || labels == nullptr || distances == nullptr) { @@ -145,7 +201,6 @@ class QGAssigner { throw std::invalid_argument("QGAssigner input dimensions overflow ptrdiff_t"); } - sums_ready = query_sums != nullptr && sums_ready; if (builder_) { builder_->reset(centroids); } else { @@ -200,7 +255,7 @@ class QGAssigner { continue; } try { - ThreadScratch& scratch = *scratch_[static_cast(thread_id)]; + auto& scratch = *scratch_[static_cast(thread_id)]; // Keep visited storage with the worker, as in QuantizedGraph::search. thread_local VisitedSet visited; thread_local size_t visited_size = 0; @@ -214,24 +269,17 @@ class QGAssigner { } const PID hint = previous_labels_.size() == n ? previous_labels_[i] : kPidMax; - graph_->search_with_scratch( + symqg::detail::search_with_workspace( + *graph_, x + i * d_, 1, labels + i, distances + i, - scratch.rotated_query.data(), - scratch.estimated_distances.data(), - scratch.lookup_table.data(), - scratch.batch_query, - scratch.search_pool, - scratch.result_pool, - visited, + scratch.view(visited), hint, - sums_ready ? query_sums + i : nullptr + query_sums.cached(i) ); - if (query_sums != nullptr && !sums_ready) { - query_sums[i] = scratch.batch_query.k1xsumq(); - } + query_sums.store(i, scratch.batch_query.k1xsumq()); if (graph_->is_quantized()) { // Keep objectives and returned distances relative to the original // centroids; graph search scores are quantized estimates. @@ -250,11 +298,9 @@ class QGAssigner { } previous_labels_.assign(labels, labels + n); + query_sums.complete(); } - void clear_hints() { previous_labels_.clear(); } - - private: void validate_configuration() const { if (num_threads_ > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("QGAssigner num_threads must fit in an OpenMP int"); @@ -292,7 +338,7 @@ class QGAssigner { uint32_t seed_; std::unique_ptr> graph_; std::unique_ptr builder_; - std::vector> scratch_; + std::vector> scratch_; std::vector previous_labels_; }; @@ -307,9 +353,7 @@ class QGAssignment { MetricType metric_type_; bool quantized_; QGAssigner assigner_; - bool cache_query_sums_; - std::vector query_sums_; - bool query_sums_ready_ = false; + QuerySumCache query_sums_; public: static constexpr const char* kName = "QGKMeans"; @@ -347,31 +391,15 @@ class QGAssignment { metric_type_, cp.quantization_bits ) - // Inputs and the assigner's rotator are immutable for this training call. - , cache_query_sums_( - assigner_.graph_->padded_dim() >= 1024 && - (cp.niter > 1 || cp.final_assignment != FinalAssignmentMode::Exact) - ) {} + , query_sums_(assigner_.graph_->padded_dim(), cp) {} void assign(const float* centroids, PID* labels, float* distances, bool final = false) { - if (cache_query_sums_ && query_sums_.empty()) { - // Allocate after centroid initialization releases its point permutation. - query_sums_.resize(n_); - } + query_sums_.prepare(n_); std::vector previous_labels; if (final && quantized_) { previous_labels.assign(labels, labels + n_); } - assigner_.assign( - centroids, - x_, - n_, - labels, - distances, - query_sums_.empty() ? nullptr : query_sums_.data(), - query_sums_ready_ - ); - query_sums_ready_ = true; + assigner_.assign_with_cache(centroids, x_, n_, labels, distances, query_sums_); if (previous_labels.empty()) { return; } diff --git a/include/rabitqlib/clustering/rabitqkmeans.hpp b/include/rabitqlib/clustering/rabitqkmeans.hpp index d7e54f6b..5717b694 100644 --- a/include/rabitqlib/clustering/rabitqkmeans.hpp +++ b/include/rabitqlib/clustering/rabitqkmeans.hpp @@ -11,14 +11,15 @@ #include #include -#include "rabitqlib/clustering/detail/kmeans.hpp" -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/estimator.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/clustering/detail/kmeans.hpp" // IWYU pragma: export +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::rabitqkmeans { diff --git a/include/rabitqlib/core/api.hpp b/include/rabitqlib/core/api.hpp new file mode 100644 index 00000000..a5f551e0 --- /dev/null +++ b/include/rabitqlib/core/api.hpp @@ -0,0 +1,119 @@ +#pragma once + +#include +#include + +#include "rabitqlib/core/metric.hpp" // IWYU pragma: export + +#define RABITQLIB_QUANT_API_VERSION_MAJOR 1 +#define RABITQLIB_QUANT_API_VERSION_MINOR 0 + +namespace rabitqlib::quant::v1 { + +// Stable C++ source API. This header exposes no index, Eigen, or SIMD implementation. +// See docs/docs/rabitq/core_api.md for the complete ownership/numerical contract. +enum class CodeLayout : std::uint8_t { UnpackedU16 = 0, SplitPacked = 1 }; + +struct Parameters { + // Working dimension >= 2; SplitPacked additionally requires dim % 32 == 0. + size_t dim = 0; + size_t total_bits = 1; // One sign bit plus 0..8 magnitude bits. + MetricType metric = METRIC_L2; + double t_const = -1; // <= 0: existing exact search; > 0: existing fast quantizer. +}; + +struct Buffer { + void* data = nullptr; + size_t bytes = 0; +}; + +struct ConstBuffer { + const void* data = nullptr; + size_t bytes = 0; +}; + +struct CodeOutput { + CodeLayout layout = CodeLayout::UnpackedU16; + Buffer codes; + Buffer extra; // Independent storage; used only by SplitPacked with total_bits > 1. +}; + +struct CodeView { + CodeLayout layout = CodeLayout::UnpackedU16; + ConstBuffer codes; + ConstBuffer extra; +}; + +// Values, not a record format. Callers may store these in separate columns/records. +struct Factors { + float add = 0; + float rescale = 0; + float error = 0; +}; + +struct Requirements { + size_t code_bytes; + size_t extra_bytes; + size_t workspace_bytes; + size_t workspace_alignment; +}; + +class QueryState; + +Requirements requirements(const Parameters& parameters, CodeLayout layout); + +// Input/centroid contain dim float32 values. nullptr centroid means zero. +// Output and workspace must be disjoint from inputs and each other. +Factors encode( + const Parameters& parameters, + const float* input, + const float* centroid, + CodeOutput output, + Buffer workspace +); + +// Lossless code conversion; factors are unchanged. In-place conversion is not supported. +void convert_codes( + const Parameters& parameters, CodeView source, CodeOutput destination, Buffer workspace +); + +// Borrows the transformed query without copying it. It must remain alive and unchanged +// while the returned/copied state is used. The centroid is not retained. +QueryState prepare_query( + const Parameters& parameters, const float* query, const float* centroid +); + +// Full float-query estimation, including at one bit; no quantized-query filter is implied. +inline float compute(const QueryState& query, CodeView code, const Factors& factors); + +class QueryState { + public: + QueryState() = default; + [[nodiscard]] size_t dimension() const noexcept { return dim_; } + [[nodiscard]] size_t total_bits() const noexcept { return total_bits_; } + [[nodiscard]] MetricType metric_type() const noexcept { return metric_; } + + private: + const float* query_ = nullptr; + size_t dim_ = 0; + size_t total_bits_ = 0; + MetricType metric_ = METRIC_L2; + float g_add_ = 0; + float k1xsumq_ = 0; + float (*extra_ip_)(const float*, const uint8_t*, size_t) = nullptr; + + static float compute_checked( + const QueryState& query, const CodeView& code, const Factors& factors + ); + + friend QueryState prepare_query(const Parameters&, const float*, const float*); + friend float compute(const QueryState&, CodeView, const Factors&); +}; + +inline float compute(const QueryState& query, CodeView code, const Factors& factors) { + // Keep the public value semantics; the compiled implementation borrows this + // local view instead of passing the entire aggregate through its call ABI. + return QueryState::compute_checked(query, code, factors); +} + +} // namespace rabitqlib::quant::v1 diff --git a/include/rabitqlib/core/data_layout.hpp b/include/rabitqlib/core/data_layout.hpp new file mode 100644 index 00000000..005a41da --- /dev/null +++ b/include/rabitqlib/core/data_layout.hpp @@ -0,0 +1,175 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/utils/packed.hpp" + +namespace rabitqlib { +namespace detail { +// Byte offsets shared by writable and read-only single-vector views. +// The one-bit layout is the base-code layout with bits == 1. +template +struct SingleCodeOffsets { + static size_t add(size_t padded_dim, size_t bits = 1) { return padded_dim * bits / 8; } + static size_t rescale(size_t padded_dim, size_t bits = 1) { + return add(padded_dim, bits) + sizeof(T); + } + static size_t error(size_t padded_dim, size_t bits = 1) { + return rescale(padded_dim, bits) + sizeof(T); + } +}; +} // namespace detail + +template +struct ExDataMap { + public: + explicit ExDataMap(char* data, size_t padded_dim, size_t ex_bits) + : ex_code_(reinterpret_cast(data)) + , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) + , f_recale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) {} + + static size_t data_bytes(size_t padded_dim, size_t ex_bits) { + if (ex_bits == 0) { + return 0; + } + const size_t max_size = std::numeric_limits::max(); + if (padded_dim > max_size / ex_bits) { + throw std::length_error("ExDataMap storage size exceeds size_t"); + } + const size_t code_bytes = detail::SingleCodeOffsets::add(padded_dim, ex_bits); + const size_t factor_bytes = sizeof(T) * 2; + if (code_bytes > max_size - factor_bytes) { + throw std::length_error("ExDataMap storage size exceeds size_t"); + } + return code_bytes + factor_bytes; + } + + [[nodiscard]] uint8_t* ex_code() { return ex_code_; } + [[nodiscard]] detail::PackedValueRef f_add_ex() { return f_add_ex_; } + [[nodiscard]] detail::PackedValueRef f_rescale_ex() { return f_recale_ex_; } + + private: + uint8_t* ex_code_; + detail::PackedValueRef f_add_ex_; + detail::PackedValueRef f_recale_ex_; +}; + +template +struct ConstExDataMap { + public: + explicit ConstExDataMap(const char* data, size_t padded_dim, size_t ex_bits) + : ex_code_(reinterpret_cast(data)) + , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) + , f_recale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) {} + + [[nodiscard]] const uint8_t* ex_code() const { return ex_code_; } + [[nodiscard]] T f_add_ex() const { return detail::load_packed_value(f_add_ex_); } + [[nodiscard]] T f_rescale_ex() const { + return detail::load_packed_value(f_recale_ex_); + } + + private: + const uint8_t* ex_code_; + const char* f_add_ex_; + const char* f_recale_ex_; +}; + +template +struct BaseDataMap { + public: + explicit BaseDataMap(char* data, size_t padded_dim, size_t base_bits) + : base_code_(reinterpret_cast(data)) + , f_add_(data + detail::SingleCodeOffsets::add(padded_dim, base_bits)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim, base_bits)) + , f_error_(data + detail::SingleCodeOffsets::error(padded_dim, base_bits)) {} + + static size_t data_bytes(size_t padded_dim, size_t base_bits) { + return detail::SingleCodeOffsets::error(padded_dim, base_bits) + sizeof(T); + } + + [[nodiscard]] uint8_t* base_code() { return base_code_; } + [[nodiscard]] detail::PackedValueRef f_add() { return f_add_; } + [[nodiscard]] detail::PackedValueRef f_rescale() { return f_rescale_; } + [[nodiscard]] detail::PackedValueRef f_error() { return f_error_; } + + private: + uint8_t* base_code_; + detail::PackedValueRef f_add_; + detail::PackedValueRef f_rescale_; + detail::PackedValueRef f_error_; +}; + +template +struct ConstBaseDataMap { + public: + explicit ConstBaseDataMap(const char* data, size_t padded_dim, size_t base_bits) + : base_code_(reinterpret_cast(data)) + , f_add_(data + detail::SingleCodeOffsets::add(padded_dim, base_bits)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim, base_bits)) + , f_error_(data + detail::SingleCodeOffsets::error(padded_dim, base_bits)) {} + + [[nodiscard]] const uint8_t* base_code() const { return base_code_; } + [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } + [[nodiscard]] T f_rescale() const { return detail::load_packed_value(f_rescale_); } + [[nodiscard]] T f_error() const { return detail::load_packed_value(f_error_); } + + private: + const uint8_t* base_code_; + const char* f_add_; + const char* f_rescale_; + const char* f_error_; +}; + +template +struct BinDataMap { + public: + explicit BinDataMap(char* data, size_t padded_dim) + : bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::SingleCodeOffsets::add(padded_dim)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim)) + , f_error_(data + detail::SingleCodeOffsets::error(padded_dim)) {} + + [[nodiscard]] uint8_t* bin_code() { return bin_code_; } + [[nodiscard]] detail::PackedValueRef f_add() { return f_add_; } + [[nodiscard]] detail::PackedValueRef f_rescale() { return f_rescale_; } + [[nodiscard]] detail::PackedValueRef f_error() { return f_error_; } + + static size_t data_bytes(size_t padded_dim) { + return detail::SingleCodeOffsets::error(padded_dim) + sizeof(T); + } + + private: + uint8_t* bin_code_; + detail::PackedValueRef f_add_; + detail::PackedValueRef f_rescale_; + detail::PackedValueRef f_error_; +}; + +template +struct ConstBinDataMap { + public: + explicit ConstBinDataMap(const char* data, size_t padded_dim) + : bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::SingleCodeOffsets::add(padded_dim)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim)) + , f_error_(data + detail::SingleCodeOffsets::error(padded_dim)) {} + + [[nodiscard]] const uint8_t* bin_code() const { return bin_code_; } + [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } + [[nodiscard]] T f_rescale() const { return detail::load_packed_value(f_rescale_); } + [[nodiscard]] T f_error() const { return detail::load_packed_value(f_error_); } + + static size_t data_bytes(size_t padded_dim) { + return BinDataMap::data_bytes(padded_dim); + } + + private: + const uint8_t* bin_code_; + const char* f_add_; + const char* f_rescale_; + const char* f_error_; +}; +} // namespace rabitqlib diff --git a/include/rabitqlib/core/detail/eigen.hpp b/include/rabitqlib/core/detail/eigen.hpp new file mode 100644 index 00000000..4dfbac57 --- /dev/null +++ b/include/rabitqlib/core/detail/eigen.hpp @@ -0,0 +1,35 @@ +#pragma once + +// Retain Dense as the existing entry point and explicitly export the alias types. +#include "rabitqlib/third/Eigen/Core" // IWYU pragma: export +#include "rabitqlib/third/Eigen/Dense" // IWYU pragma: export + +namespace rabitqlib { +template +using RowMajorMatrix = Eigen::Matrix; + +template +using RowMajorMatrixMap = Eigen::Map>; + +template +using ConstRowMajorMatrixMap = Eigen::Map>; + +template +using RowMajorArray = Eigen::Array; + +template +using Vector = Eigen::Matrix; + +template +using RowMajorArrayMap = Eigen::Map>; + +template +using ConstRowMajorArrayMap = Eigen::Map>; + +template +using VectorMap = Eigen::Map>; + +template +using ConstVectorMap = Eigen::Map>; + +} // namespace rabitqlib diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp similarity index 68% rename from include/rabitqlib/quantization/rabitq_impl.hpp rename to include/rabitqlib/core/detail/rabitq_impl.hpp index 2425f9c9..9d4ef7a8 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -9,16 +9,22 @@ #include #include #include -#include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/pack_excode.hpp" -#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/scratch.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/core/types.hpp" +#include "rabitqlib/utils/openmp.hpp" // IWYU pragma: keep (configuration macros) #include "rabitqlib/utils/space.hpp" +#if RABITQ_USE_OPENMP +#include +#endif + namespace rabitqlib::quant::rabitq_impl { constexpr float kConstEpsilon = 1.9; @@ -26,33 +32,18 @@ constexpr float kConstEpsilon = 1.9; namespace one_bit { /** - * @brief The one_bit_code function maps a data vector to a binary code by recording the - * sign of every coordinate in an int array. - * - * @param data Input data vector to be quantized - * @param centroid Input center vector used as reference for quantization - * @param dim Dimensionality of the vectors - * @param binary_code Output binary code where each element is 0 or 1 based on whether - * the corresponding residual element (data - centroid) is positive - * - * @return RowMajorArray The residual vector (data - centroid) - * - * This function computes a binary code by recording the sign of every coordinate - * in the residual vector (data - centroid). It returns the residual vector for - * further processing. + * Return data - centroid and write its sign code: 1 for strictly positive + * coordinates, 0 otherwise (including zero). All buffers contain dim elements. */ template inline RowMajorArray one_bit_code( const T* data, const T* centroid, size_t dim, int* binary_code ) { - // map pointer to array ConstRowMajorArrayMap data_arr(data, 1, dim); ConstRowMajorArrayMap cent_arr(centroid, 1, dim); - // residual vector RowMajorArray residual_arr = data_arr - cent_arr; - // unsigned representation RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); x_u = (residual_arr > 0).template cast(); @@ -60,26 +51,17 @@ inline RowMajorArray one_bit_code( } /** - * @brief The one_bit_code_with_factor function maps a data vector to a binary code and - * computes factors for distance estimation. - * - * @param data Input data vector to be quantized - * @param centroid Input center vector used as reference for quantization - * @param dim Dimensionality of the vectors - * @param binary_code Output binary code where each element is 0 or 1 based on whether - * the corresponding residual element (data - centroid) is positive - * @param f_add Output factor used in distance estimation to add to the final result - * @param f_rescale Output scaling factor used in distance estimation - * @param f_error Output error bound for the distance estimation - * @param metric_type Type of distance metric to use (L2 or inner product) - * - * This function computes a binary code by recording the sign of every coordinate - * in the residual vector (data - centroid) and calculates the necessary factors - * for accurate distance estimation. The factors are used in the estimator to - * approximate distances between the original vectors. + * Write the sign code and estimation factors for data - centroid, using the + * same zero convention as one_bit_code. Inputs, binary_code, residual and + * centered scratch each contain dim elements. + * f_add/f_rescale supply the additive/scaling terms for squared L2 or 1 - dot. + * f_error is query-independent: multiplying it by the query residual norm gives + * a probabilistic error margin, not a deterministic bound for every query. + * See docs/docs/rabitq/estimator.md, "Estimator of Euclidean Distance" and + * "Estimator of Inner Product", for the factor derivation. */ template -inline void one_bit_code_with_factor( +inline void one_bit_code_with_factor_into( const T* data, const T* centroid, size_t dim, @@ -87,23 +69,33 @@ inline void one_bit_code_with_factor( T& f_add, T& f_rescale, T& f_error, - MetricType metric_type = METRIC_L2 + MetricType metric_type, + T* residual, + T* centered ) { - // Scratch stays private to each thread and is reused across vectors and batches. - thread_local RowMajorArray residual_arr; - thread_local RowMajorArray xu_cb; - ConstRowMajorArrayMap data_arr(data, 1, dim); - ConstRowMajorArrayMap cent_arr(centroid, 1, dim); - residual_arr = data_arr - cent_arr; - RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); - x_u = (residual_arr > 0).template cast(); - - // xu_cb = x_u + cb, xu_cb has same direction and different length with x_bar - float cb = -((1 << 1) - 1) / 2.F; - xu_cb = x_u.template cast() + cb; + if (std::is_same_v && dim <= 128) { + // Fuse short float vectors; zero residuals retain sign 0. + for (size_t i = 0; i < dim; ++i) { + const T value = data[i] - centroid[i]; + residual[i] = value; + binary_code[i] = static_cast(value > 0); + centered[i] = value > 0 ? T{0.5} : T{-0.5}; + } + } else { + // Wider vectors benefit from Eigen's packet width rather than the + // compiler's narrower vectorization of the fused multi-output loop. + RowMajorArrayMap residual_arr(residual, 1, static_cast(dim)); + RowMajorArrayMap xu_cb(centered, 1, static_cast(dim)); + ConstRowMajorArrayMap data_arr(data, 1, dim); + ConstRowMajorArrayMap cent_arr(centroid, 1, dim); + residual_arr = data_arr - cent_arr; + RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); + x_u = (residual_arr > 0).template cast(); + xu_cb = x_u.template cast() + T{-0.5}; + } // distance to centroid - T l2_sqr = l2norm_sqr(residual_arr.data(), dim); + T l2_sqr = l2norm_sqr(residual, dim); T l2_norm = std::sqrt(l2_sqr); if (l2_sqr == 0) { @@ -118,26 +110,20 @@ inline void one_bit_code_with_factor( f_error = 0; return; } - - // dot product between residual and xu_cb - T ip_resi_xucb = dot_product(residual_arr.data(), xu_cb.data(), dim); - // dot product between centroid and xu_cb - T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); + T ip_resi_xucb = dot_product(residual, centered, dim); + T ip_cent_xucb = dot_product(centroid, centered, dim); // A nonzero residual and its sign code have a strictly positive inner product. assert(ip_resi_xucb > 0); - // We use unnormalized vector to get error factor. To be more specific, - // sqrt((1 - ^2) / ^2) / sqrt(dim - 1) = 3rd item in following - // expression + // tan^2(angle(residual, centered)) / (dim - 1), using unnormalized vectors. // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero. const T normalized_error = - (((l2_sqr * l2norm_sqr(xu_cb.data(), dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1 - ) / + (((l2_sqr * l2norm_sqr(centered, dim)) / (ip_resi_xucb * ip_resi_xucb)) - 1) / (dim - 1); T tmp_error = l2_norm * kConstEpsilon * std::sqrt(std::max(normalized_error, T{0})); - // 3 factors used for distance estimation, please refer to document for more info. + // Factor derivation: docs/docs/rabitq/estimator.md, metric-specific estimators. // For f_rescale and 2nd item of f_add, we use the dot product of raw residual (rather // than the normalized one) as the denominator, thus we need to multiply another l2norm. // For (ip_cent_xucb / ip_resi_xucb), the norm of xucb does not matter since it is also @@ -147,7 +133,7 @@ inline void one_bit_code_with_factor( f_rescale = -2 * l2_sqr / ip_resi_xucb; f_error = 2 * tmp_error; } else if (metric_type == METRIC_IP) { - f_add = 1 - dot_product(residual_arr.data(), centroid, dim) + + f_add = 1 - dot_product(residual, centroid, dim) + (l2_sqr * ip_cent_xucb / ip_resi_xucb); f_rescale = -l2_sqr / ip_resi_xucb; f_error = 1 * tmp_error; @@ -156,26 +142,51 @@ inline void one_bit_code_with_factor( } } +template +inline void one_bit_code_with_factor( + const T* data, + const T* centroid, + size_t dim, + int* binary_code, + T& f_add, + T& f_rescale, + T& f_error, + MetricType metric_type = METRIC_L2 +) { + // Legacy calls retain their existing thread-local allocation/reuse behavior. + thread_local RowMajorArray residual; + thread_local RowMajorArray centered; + // Eigen's resize frees the old allocation before allocating its replacement. + // Allocate first so failure leaves the thread-local owner safe to reuse/destroy. + if (residual.size() != static_cast(dim)) { + RowMajorArray replacement(1, static_cast(dim)); + residual.swap(replacement); + } + if (centered.size() != static_cast(dim)) { + RowMajorArray replacement(1, static_cast(dim)); + centered.swap(replacement); + } + one_bit_code_with_factor_into( + data, + centroid, + dim, + binary_code, + f_add, + f_rescale, + f_error, + metric_type, + residual.data(), + centered.data() + ); +} + /** - * @brief The one_bit_compact_code function maps a data vector to a compact binary code and - * computes factors for distance estimation. - * - * @param data Input data vector to be quantized - * @param centroid Input center vector used as reference for quantization - * @param padded_dim Dimensionality of the vectors (padded to a multiple of the bit-packing - * size) - * @param compact_code Output compact binary code where the binary values are packed into - * integers in big-endian order - * @param f_add Output factor used in distance estimation to add to the final result - * @param f_rescale Output scaling factor used in distance estimation - * @param f_error Output error bound for the distance estimation - * @param metric_type Type of distance metric to use (L2 or inner product) - * - * This function computes a compact binary code by recording the sign of every coordinate - * in the residual vector (data - centroid) and packs the binary values into integers - * for efficient storage and computation. It also calculates the necessary factors - * for accurate distance estimation. The factors are used in the estimator to - * approximate distances between the original vectors. + * Pack the sign code into padded_dim / 8 bytes; padded_dim is a positive + * multiple of 8. Both inputs contain padded_dim elements. The first coordinate + * uses the highest used bit of each word; full and partial words use native + * byte order. Zero residual coordinates have sign bit 0. + * Factors have the same meaning as in one_bit_code_with_factor_into, including + * the query-independent, probabilistic error factor f_error. */ template inline void one_bit_compact_code_to_bytes( @@ -190,8 +201,6 @@ inline void one_bit_compact_code_to_bytes( ) { thread_local std::vector binary_code; binary_code.resize(padded_dim); - - // get binary code one_bit_code_with_factor( data, centroid, @@ -248,9 +257,7 @@ inline void one_bit_compact_codes( FError f_error, MetricType metric_type = METRIC_L2 ) { -#pragma omp parallel for if (Parallel) - for (std::ptrdiff_t index = 0; index < static_cast(num); ++index) { - const size_t i = static_cast(index); + const auto encode_one = [&](size_t i) { T add; T rescale; T error; @@ -267,60 +274,40 @@ inline void one_bit_compact_codes( f_add[i] = add; f_rescale[i] = rescale; f_error[i] = error; + }; +#if RABITQ_USE_OPENMP + if constexpr (Parallel) { + std::exception_ptr failure; +#pragma omp parallel for + for (std::ptrdiff_t index = 0; index < static_cast(num); ++index) { + try { + encode_one(static_cast(index)); + } catch (...) { +#pragma omp critical(rabitq_compact_codes_error) + { + if (!failure) { + failure = std::current_exception(); + } + } + } + } + if (failure) { + std::rethrow_exception(failure); + } + } else +#endif + { + // No serialized OpenMP region: an outer caller can catch scratch failures. + for (size_t i = 0; i < num; ++i) { + encode_one(i); + } } } -// Encoding requires a positive padded_dim divisible by 8; index pipelines use multiples -// of 32. packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail -// padding. -template < - typename T, - bool Parallel = false, - typename FAdd, - typename FRescale, - typename FError> -inline void one_bit_batch_code( - const T* data, - const T* centroid, - size_t num, - size_t padded_dim, - uint8_t* packed_code, - FAdd f_add, - FRescale f_recale, - FError f_error, - MetricType metric_type = METRIC_L2 -) { - std::vector compact_codes(num * padded_dim / 8); - - one_bit_compact_codes( - data, - centroid, - num, - padded_dim, - compact_codes.data(), - f_add, - f_recale, - f_error, - metric_type - ); - fastscan::pack_codes(padded_dim, compact_codes.data(), num, packed_code); -} - } // namespace one_bit namespace ex_bits { -/** - * @brief Get quantization code and factor (ipnorm_inv) for RaBitQ+ - * @note Currently assume all number in o_prime are positive. For negative - * values, the quantization code needs to be flipped. - * - * @param o_prime Rotated & normalized data vector, length of dim - * @param code Quantization code of RaBitQ+, length of dim, stored as uint8 - * @param dim Dimension of rotated vec - * @param ex_bits Number of bits used for code - * @return ipnorm_inv Returned factor - */ - +// Lower endpoint of the scale search as a fraction of t_end, indexed by ex_bits. constexpr std::array kTightStart = { 0, 0.15, @@ -349,45 +336,19 @@ inline int quantized_level_at_scale(double magnitude, double t, int max_code) { return code; } -// Input is the library's finite, nonnegative normalized magnitude vector. -template -inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) { - if (ex_bits >= kTightStart.size()) { - throw std::invalid_argument("ex_bits exceeds the supported search table"); - } - if (dim == 0) - return 0; - const double max_o = static_cast(*std::max_element(o_abs, o_abs + dim)); - if (max_o == 0) - return 0; - - const int max_code = (1 << ex_bits) - 1; - const double t_end = static_cast(max_code + 10) / max_o; - const double t_start = t_end * kTightStart[ex_bits]; - - // Both paths maximize cosine similarity between the input and quantized - // directions over the configured scale interval. Here a_i = o_abs[i] is a - // rotated, normalized residual magnitude, and q_i = code_i + 0.5. - // With ||a|| = 1, cosine(a, q) = dot(a, q / ||q||) = N / sqrt(S), where - // N = sum(a_i * q_i) and S = sum(q_i * q_i). - // SIMD evaluates whole vectors at selected scales and prunes intervals; - // its search and bounds are in src/simd/rescale_search.hpp. - if constexpr (std::is_same_v) { - const double t = simd::best_rescale_factor(o_abs, dim, max_code, t_start, t_end); - if (t >= 0) - return t; - } - - // Scalar event sweep, also used if SIMD is unavailable or cannot certify a - // winner within its work limit. Each heap entry holds a coordinate's next - // code-changing scale (code_i + 1) / a_i. Visiting only these events skips - // ranges with constant codes and updates N and S in O(1) per coordinate. - // Heap traversal is sequential and branch-dependent, so it offers little - // SIMD parallelism; the SIMD path instead prunes ranges of candidate codes. - using Event = std::pair; - std::vector next_t; +template +inline double event_sweep( + const T* o_abs, + size_t dim, + int max_code, + double t_start, + double t_end, + Heap& next_t, + Codes& cur_o_bar +) { + using Event = RescaleEvent; next_t.reserve(dim); - std::vector cur_o_bar(dim); + cur_o_bar.resize(dim); double sqr_denominator = static_cast(dim) * 0.25; double numerator = 0; @@ -459,9 +420,77 @@ inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) { return best_t; } +// Input is the library's finite, nonnegative normalized magnitude vector. +template +inline double best_rescale_factor_with_scratch( + const T* o_abs, size_t dim, size_t ex_bits, Scratch* scratch +) { + if (ex_bits >= kTightStart.size()) { + throw std::invalid_argument("ex_bits exceeds the supported search table"); + } + if (dim == 0) + return 0; + const double max_o = static_cast(*std::max_element(o_abs, o_abs + dim)); + if (max_o == 0) + return 0; + + const int max_code = (1 << ex_bits) - 1; + const double t_end = static_cast(max_code + 10) / max_o; + const double t_start = t_end * kTightStart[ex_bits]; + + // Both paths maximize cosine similarity between the input and quantized + // directions over the configured scale interval. Here a_i = o_abs[i] is a + // rotated, normalized residual magnitude, and q_i = code_i + 0.5. + // With ||a|| = 1, cosine(a, q) = dot(a, q / ||q||) = N / sqrt(S), where + // N = sum(a_i * q_i) and S = sum(q_i * q_i). + // SIMD evaluates whole vectors at selected scales and prunes intervals; + // its search and bounds are in src/simd/rescale_search.hpp. + if constexpr (std::is_same_v) { + const double t = + scratch == nullptr + ? simd::best_rescale_factor(o_abs, dim, max_code, t_start, t_end) + : simd::best_rescale_factor_with_workspace( + o_abs, + dim, + max_code, + t_start, + t_end, + scratch->magnitudes, + scratch->reciprocals + ); + if (t >= 0) + return t; + } + + // Scalar event sweep, also used if SIMD is unavailable or cannot certify a + // winner within its work limit. Each heap entry holds a coordinate's next + // code-changing scale (code_i + 1) / a_i. Visiting only these events skips + // ranges with constant codes and updates N and S in O(1) per coordinate. + // Heap traversal is sequential and branch-dependent, so it offers little + // SIMD parallelism; the SIMD path instead prunes ranges of candidate codes. + using Event = RescaleEvent; + if (scratch != nullptr) { + ScratchAllocator allocator(scratch->resource, dim); + auto [next_t, cur_o_bar] = event_sweep_buffers(allocator); + return event_sweep(o_abs, dim, max_code, t_start, t_end, next_t, cur_o_bar); + } + std::vector next_t; + std::vector cur_o_bar; + return event_sweep(o_abs, dim, max_code, t_start, t_end, next_t, cur_o_bar); +} + +template +inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) { + return best_rescale_factor_with_scratch( + o_abs, dim, ex_bits, static_cast*>(nullptr) + ); +} + template -inline T quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits) { - const double t = best_rescale_factor(o_abs, dim, ex_bits); +inline T quantize_ex_with_scratch( + const T* o_abs, TP* code, size_t dim, size_t ex_bits, Scratch* scratch +) { + const double t = best_rescale_factor_with_scratch(o_abs, dim, ex_bits, scratch); const int max_code = (1 << ex_bits) - 1; double ipnorm = 0; for (size_t i = 0; i < dim; ++i) { @@ -477,6 +506,13 @@ inline T quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits) { return std::isnormal(ipnorm_inv) ? ipnorm_inv : static_cast(1); } +template +inline T quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits) { + return quantize_ex_with_scratch( + o_abs, code, dim, ex_bits, static_cast*>(nullptr) + ); +} + // For given dim and ex_bits, use random vectors to get the const rescale factor inline double get_const_scaling_factors(size_t dim, size_t ex_bits) { constexpr long kConstNum = 100; @@ -500,17 +536,15 @@ T faster_quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits, doubl constexpr double kEps = 1e-5; double ipnorm = 0; - std::vector tmp_code(dim); for (size_t i = 0; i < dim; i++) { - // compute and store code - tmp_code[i] = static_cast((t_const * o_abs[i]) + kEps); - if (tmp_code[i] >= (1 << ex_bits)) { - tmp_code[i] = (1 << ex_bits) - 1; + int tmp_code = static_cast((t_const * o_abs[i]) + kEps); + if (tmp_code >= (1 << ex_bits)) { + tmp_code = (1 << ex_bits) - 1; } - code[i] = static_cast(tmp_code[i]); + code[i] = static_cast(tmp_code); // ip * norm = unnormalized ip - ipnorm += (tmp_code[i] + 0.5) * o_abs[i]; + ipnorm += (tmp_code + 0.5) * o_abs[i]; } T ipnorm_inv = static_cast(1 / ipnorm); // 1 / (ip*norm) @@ -522,22 +556,30 @@ T faster_quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits, doubl } template -inline T ex_bits_code( - const T* residual, size_t dim, size_t ex_bits, TP* ex_code, double t_const = -1 +inline T ex_bits_code_with_scratch( + const T* residual, + size_t dim, + size_t ex_bits, + TP* ex_code, + double t_const, + Scratch* scratch ) { ConstRowMajorArrayMap res_arr(residual, 1, dim); - - // quantize data T ipnorm_inv = 1; const T residual_norm = std::sqrt(l2norm_sqr(residual, dim)); if (residual_norm == 0) { std::fill(ex_code, ex_code + dim, static_cast(0)); } else { - RowMajorArray abs_res = res_arr.abs() / residual_norm; + RowMajorArray owned_absolute; + auto abs_res = scratch_array( + scratch == nullptr ? nullptr : scratch->absolute, dim, owned_absolute + ); + abs_res = res_arr.abs() / residual_norm; if (t_const > 0) { ipnorm_inv = faster_quantize_ex(abs_res.data(), ex_code, dim, ex_bits, t_const); } else { - ipnorm_inv = quantize_ex(abs_res.data(), ex_code, dim, ex_bits); + ipnorm_inv = + quantize_ex_with_scratch(abs_res.data(), ex_code, dim, ex_bits, scratch); } } @@ -554,7 +596,16 @@ inline T ex_bits_code( } template -inline void ex_bits_code_with_factor( +inline T ex_bits_code( + const T* residual, size_t dim, size_t ex_bits, TP* ex_code, double t_const = -1 +) { + return ex_bits_code_with_scratch( + residual, dim, ex_bits, ex_code, t_const, static_cast*>(nullptr) + ); +} + +template +inline void ex_bits_code_with_factor_with_scratch( const T* data, const T* centroid, size_t dim, @@ -563,14 +614,17 @@ inline void ex_bits_code_with_factor( T& f_add_ex, T& f_rescale_ex, T& f_error_ex, - MetricType metric_type = METRIC_L2, - double t_const = -1 + MetricType metric_type, + double t_const, + Scratch* scratch ) { ConstRowMajorArrayMap data_arr(data, 1, dim); ConstRowMajorArrayMap cent_arr(centroid, 1, dim); - - // residual vector - RowMajorArray residual_arr = data_arr - cent_arr; + RowMajorArray owned_residual; + auto residual_arr = scratch_array( + scratch == nullptr ? nullptr : scratch->residual, dim, owned_residual + ); + residual_arr = data_arr - cent_arr; const T l2_sqr = l2norm_sqr(residual_arr.data(), dim); if (l2_sqr == 0) { @@ -587,20 +641,26 @@ inline void ex_bits_code_with_factor( return; } - T ipnorm_inv = ex_bits_code(residual_arr.data(), dim, ex_bits, ex_code, t_const); - - // get factors - RowMajorArray total_code = - RowMajorArrayMap(ex_code, 1, dim).template cast(); + T ipnorm_inv = ex_bits_code_with_scratch( + residual_arr.data(), dim, ex_bits, ex_code, t_const, scratch + ); + RowMajorArray owned_total; + auto total_code = + scratch_array(scratch == nullptr ? nullptr : scratch->total, dim, owned_total); + total_code = RowMajorArrayMap(ex_code, 1, dim).template cast(); for (size_t i = 0; i < dim; ++i) { total_code(0, static_cast(i)) += static_cast(residual_arr.data()[i] > 0) << ex_bits; } - // Factors are similar to those in one_bit_code_with_factor(), please refer to document - // for detailed info. + // Full-code factors use the metric-specific derivation in + // docs/docs/rabitq/estimator.md, with xu_cb centered at this code width. float cb = -(static_cast(1 << ex_bits) - 0.5F); - RowMajorArray xu_cb = total_code.template cast() + cb; + RowMajorArray owned_centered; + auto xu_cb = scratch_array( + scratch == nullptr ? nullptr : scratch->centered, dim, owned_centered + ); + xu_cb = total_code.template cast() + cb; T l2_norm = std::sqrt(l2_sqr); @@ -631,6 +691,34 @@ inline void ex_bits_code_with_factor( } } +template +inline void ex_bits_code_with_factor( + const T* data, + const T* centroid, + size_t dim, + size_t ex_bits, + TP* ex_code, + T& f_add_ex, + T& f_rescale_ex, + T& f_error_ex, + MetricType metric_type = METRIC_L2, + double t_const = -1 +) { + ex_bits_code_with_factor_with_scratch( + data, + centroid, + dim, + ex_bits, + ex_code, + f_add_ex, + f_rescale_ex, + f_error_ex, + metric_type, + t_const, + static_cast*>(nullptr) + ); +} + template inline void ex_bits_compact_code( const T* data, @@ -731,7 +819,7 @@ static inline void rabitq_scalar_impl( } template -static inline void rabitq_full_impl( +static inline void rabitq_full_impl_with_scratch( const T* data, const T* centroid, size_t dim, @@ -740,16 +828,33 @@ static inline void rabitq_full_impl( T& f_add, T& f_rescale, T& f_error, - MetricType metric_type = METRIC_L2, - double t_const = -1 + MetricType metric_type, + double t_const, + Scratch* scratch ) { - std::vector binary_code(dim); + std::vector owned_binary; + int* binary_code = nullptr; + if (scratch == nullptr) { + owned_binary.resize(dim); + binary_code = owned_binary.data(); + } else { + binary_code = scratch->binary; + } size_t ex_bits = total_bits - 1; - rabitq_impl::one_bit::one_bit_code(data, centroid, dim, binary_code.data()); - if (ex_bits > 0) { - ex_bits::ex_bits_code_with_factor( + if (scratch == nullptr) { + rabitq_impl::one_bit::one_bit_code(data, centroid, dim, binary_code); + } else { + ConstRowMajorArrayMap input(data, 1, static_cast(dim)); + ConstRowMajorArrayMap center(centroid, 1, static_cast(dim)); + RowMajorArrayMap residual(scratch->residual, 1, static_cast(dim)); + residual = input - center; + RowMajorArrayMap signs(binary_code, 1, static_cast(dim)); + signs = (residual > 0).template cast(); + } + + ex_bits::ex_bits_code_with_factor_with_scratch( data, centroid, dim, @@ -759,12 +864,28 @@ static inline void rabitq_full_impl( f_rescale, f_error, metric_type, - t_const + t_const, + scratch ); } else { - one_bit::one_bit_code_with_factor( - data, centroid, dim, binary_code.data(), f_add, f_rescale, f_error, metric_type - ); + if (scratch == nullptr) { + one_bit::one_bit_code_with_factor( + data, centroid, dim, binary_code, f_add, f_rescale, f_error, metric_type + ); + } else { + one_bit::one_bit_code_with_factor_into( + data, + centroid, + dim, + binary_code, + f_add, + f_rescale, + f_error, + metric_type, + scratch->residual, + scratch->centered + ); + } } for (size_t i = 0; i < dim; ++i) { @@ -772,6 +893,34 @@ static inline void rabitq_full_impl( total_code[i] = ex_bits > 0 ? total_code[i] + sign_code : sign_code; } } +template +static inline void rabitq_full_impl( + const T* data, + const T* centroid, + size_t dim, + size_t total_bits, + TP* total_code, + T& f_add, + T& f_rescale, + T& f_error, + MetricType metric_type = METRIC_L2, + double t_const = -1 +) { + rabitq_full_impl_with_scratch( + data, + centroid, + dim, + total_bits, + total_code, + f_add, + f_rescale, + f_error, + metric_type, + t_const, + static_cast*>(nullptr) + ); +} + } // namespace total_bits namespace xy_bits { @@ -796,7 +945,11 @@ inline void validate_bit_size(size_t base_bits, size_t ex_bits) { } /** - * @brief Derive the 3 estimation factors for an arbitrary integer code. + * @brief Derive estimation factors for a residual-compatible integer code. + * + * For a nonzero residual, dot(residual, code + cb) must be strictly positive; + * the code's integer range alone does not guarantee this. A zero residual + * produces zero rescale/error factors and a metric-specific additive term. * * one_bit_code_with_factor() and ex_bits_code_with_factor() both compute the * same three quantities from xu_cb = code + cb, differing only in how the diff --git a/include/rabitqlib/core/detail/scratch.hpp b/include/rabitqlib/core/detail/scratch.hpp new file mode 100644 index 00000000..9832ded1 --- /dev/null +++ b/include/rabitqlib/core/detail/scratch.hpp @@ -0,0 +1,105 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/eigen.hpp" + +namespace rabitqlib::quant::rabitq_impl { +// Every scratch array has one entry per coordinate. Walk the same declarations +// either to derive a reusable size budget or to bind a caller's bounded resource. +class ScratchAllocator { + std::pmr::memory_resource* resource_ = nullptr; + size_t dim_ = 0; + size_t bytes_per_coordinate_ = 0; + size_t alignment_padding_ = 0; + + template + void measure() { + static_assert(alignof(T) <= alignof(std::max_align_t)); + const size_t maximum = std::numeric_limits::max(); + if (sizeof(T) > maximum - bytes_per_coordinate_ || + alignof(T) - 1 > maximum - alignment_padding_) { + throw std::length_error("Quantization workspace size overflows size_t"); + } + bytes_per_coordinate_ += sizeof(T); + // Each allocation may need at most alignof(T)-1 bytes before its data. + alignment_padding_ += alignof(T) - 1; + } + + public: + ScratchAllocator() = default; + ScratchAllocator(std::pmr::memory_resource* resource, size_t dim) + : resource_(resource), dim_(dim) {} + + [[nodiscard]] size_t bytes(size_t dim) const { + if (bytes_per_coordinate_ != 0 && + dim > (std::numeric_limits::max() - alignment_padding_) / + bytes_per_coordinate_) { + throw std::length_error("Quantization workspace size overflows size_t"); + } + return dim * bytes_per_coordinate_ + alignment_padding_; + } + [[nodiscard]] std::pmr::memory_resource* resource() const { return resource_; } + + template + T* array() { + if (resource_ == nullptr) { + measure(); + return nullptr; + } + auto* result = std::pmr::polymorphic_allocator(resource_).allocate(dim_); + // Establish C++17 object lifetimes without initializing scalar values. + std::uninitialized_default_construct_n(result, dim_); + return result; + } + + template + std::pmr::vector vector() { + std::pmr::vector result( + resource_ == nullptr ? std::pmr::null_memory_resource() : resource_ + ); + if (resource_ == nullptr) { + measure(); + } else { + result.reserve(dim_); + } + return result; + } +}; + +using RescaleEvent = std::pair; + +// This declaration supplies both requirements()'s fallback budget and the +// actual event-sweep containers. Sizing returns empty, non-allocating vectors. +inline auto event_sweep_buffers(ScratchAllocator& allocator) { + return std::pair{allocator.vector(), allocator.vector()}; +} + +// Internal non-owning scratch; never part of the stable public API. +template +struct Scratch { + T* residual; + T* absolute; + T* centered; + int* binary; + int* total; + double* magnitudes; + double* reciprocals; + std::pmr::memory_resource* resource; +}; + +template +inline RowMajorArrayMap scratch_array(T* external, size_t dim, RowMajorArray& owned) { + if (external == nullptr) { + owned.resize(1, static_cast(dim)); + external = owned.data(); + } + return RowMajorArrayMap(external, 1, static_cast(dim)); +} +} // namespace rabitqlib::quant::rabitq_impl diff --git a/include/rabitqlib/simd/matrix_dispatch.hpp b/include/rabitqlib/core/detail/simd/matrix_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/matrix_dispatch.hpp rename to include/rabitqlib/core/detail/simd/matrix_dispatch.hpp diff --git a/include/rabitqlib/simd/pack_excode_dispatch.hpp b/include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/pack_excode_dispatch.hpp rename to include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp diff --git a/include/rabitqlib/core/detail/simd/quantization_dispatch.hpp b/include/rabitqlib/core/detail/simd/quantization_dispatch.hpp new file mode 100644 index 00000000..9fb71a6c --- /dev/null +++ b/include/rabitqlib/core/detail/simd/quantization_dispatch.hpp @@ -0,0 +1,60 @@ +#pragma once + +#include + +namespace rabitqlib::simd { + +// Search normalized, nonnegative float magnitudes with max_code in [0, 255]. +// A negative result requests the event sweep when the bounded search cannot +// certify a winner within its work limit. +double best_rescale_factor( + const float* magnitudes, size_t dim, int max_code, double start, double end +); +double best_rescale_factor_generic( + const float* magnitudes, size_t dim, int max_code, double start, double end +); +double best_rescale_factor_avx2( + const float* magnitudes, size_t dim, int max_code, double start, double end +); +double best_rescale_factor_avx512( + const float* magnitudes, size_t dim, int max_code, double start, double end +); + +// Internal explicit scratch entry points; same kernels and exact fallback contract. +double best_rescale_factor_with_workspace( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +); +double best_rescale_factor_with_workspace_generic( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +); +double best_rescale_factor_with_workspace_avx2( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +); +double best_rescale_factor_with_workspace_avx512( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +); +} // namespace rabitqlib::simd diff --git a/include/rabitqlib/simd/rotator_dispatch.hpp b/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/rotator_dispatch.hpp rename to include/rabitqlib/core/detail/simd/rotator_dispatch.hpp diff --git a/include/rabitqlib/simd/space_dispatch.hpp b/include/rabitqlib/core/detail/simd/space_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/space_dispatch.hpp rename to include/rabitqlib/core/detail/simd/space_dispatch.hpp diff --git a/include/rabitqlib/simd/warmup_dispatch.hpp b/include/rabitqlib/core/detail/simd/warmup_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/warmup_dispatch.hpp rename to include/rabitqlib/core/detail/simd/warmup_dispatch.hpp diff --git a/include/rabitqlib/quantization/rabitq.hpp b/include/rabitqlib/core/encoder.hpp similarity index 63% rename from include/rabitqlib/quantization/rabitq.hpp rename to include/rabitqlib/core/encoder.hpp index 0c91a813..50445f61 100644 --- a/include/rabitqlib/quantization/rabitq.hpp +++ b/include/rabitqlib/core/encoder.hpp @@ -4,24 +4,16 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq_impl.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/types.hpp" // IWYU pragma: export namespace rabitqlib::quant { -struct RabitqConfig { - double t_const = -1; - explicit RabitqConfig() = default; - RabitqConfig(RabitqConfig const&) = default; - RabitqConfig(RabitqConfig&&) = default; - RabitqConfig& operator=(const RabitqConfig&) = default; - RabitqConfig& operator=(RabitqConfig&&) = default; -}; - -// init config for faster quantization rabitq, for total_bits > 0, store const rescaling -// factor in RabitqConfig +// Prepare a constant rescaling factor for total_bits > 1. The one-bit path +// needs no extra-code scaling and leaves the default configuration unchanged. inline RabitqConfig faster_config(size_t dim, size_t total_bits) { RabitqConfig config; if (total_bits > 1) { @@ -31,104 +23,6 @@ inline RabitqConfig faster_config(size_t dim, size_t total_bits) { return config; } -template -inline void quantize_one_batch( - const T* data, - const T* centroid, - size_t num, - size_t padded_dim, - char* batch_data, - MetricType metric_type = METRIC_L2 -) { - BatchDataMap this_batch(batch_data, padded_dim); - - rabitq_impl::one_bit::one_bit_batch_code( - data, - centroid, - num, - padded_dim, - this_batch.bin_code(), - this_batch.f_add(), - this_batch.f_rescale(), - this_batch.f_error(), - metric_type - ); -} - -template -inline void quantize_one_batch( - const T* data, - size_t num, - size_t padded_dim, - char* batch_data, - MetricType metric_type = METRIC_L2 -) { - std::vector centroid(padded_dim, 0); - BatchDataMap this_batch(batch_data, padded_dim); - - rabitq_impl::one_bit::one_bit_batch_code( - data, - centroid.data(), - num, - padded_dim, - this_batch.bin_code(), - this_batch.f_add(), - this_batch.f_rescale(), - this_batch.f_error(), - metric_type - ); -} - -template -static inline void quantize_qg_batch( - const T* data, - const T* centroid, - size_t num, - size_t padded_dim, - char* batch_data, - MetricType metric_type = METRIC_L2 -) { - std::vector f_error(fastscan::kBatchSize); // we dont need this factor for qg - QGBatchDataMap cur_batch(batch_data, padded_dim); - - rabitq_impl::one_bit::one_bit_batch_code( - data, - centroid, - num, - padded_dim, - cur_batch.bin_code(), - cur_batch.f_add(), - cur_batch.f_rescale(), - f_error.data(), - metric_type - ); -} - -template -static inline void quantize_qg_batch( - const T* data, - size_t num, - size_t padded_dim, - char* batch_data, - MetricType metric_type = METRIC_L2 -) { - std::vector centroid(padded_dim, 0); - std::vector f_error(fastscan::kBatchSize); // we dont need this factor for qg - QGBatchDataMap cur_batch(batch_data, padded_dim); - - rabitq_impl::one_bit::one_bit_batch_code( - data, - centroid.data(), - num, - padded_dim, - cur_batch.bin_code(), - cur_batch.f_add(), - cur_batch.f_rescale(), - f_error.data(), - metric_type - ); -} - template inline void quantize_compact_one_bit( const T* data, @@ -229,31 +123,6 @@ inline void quantize_compact_ex_bits( cur_ex_data.f_rescale_ex() = f_rescale_ex; } -inline void quantize_split_batch( - const float* data, - const float* centroid, - size_t num_points, - size_t padded_dim, - size_t ex_bits, - char* batch_data, - char* ex_data, - MetricType metric_type = METRIC_L2, - RabitqConfig config = RabitqConfig() -) { - quantize_one_batch(data, centroid, num_points, padded_dim, batch_data, metric_type); - - if (ex_bits > 0) { - for (size_t i = 0; i < num_points; ++i) { - quantize_compact_ex_bits( - data, centroid, padded_dim, ex_bits, ex_data, metric_type, config - ); - - ex_data += ExDataMap::data_bytes(padded_dim, ex_bits); - data += padded_dim; - } - } -} - inline void quantize_split_single( const float* data, const float* centroid, @@ -420,11 +289,21 @@ inline void reconstruct_full_vec( code_center); } +namespace detail { +template +inline TF full_est_dist_from_ip( + TF ip, size_t bits, TF f_add, TF f_rescale, TF g_add, TF k1xsumq +) { + return f_add + g_add + + (f_rescale * (ip + k1xsumq * static_cast((1 << bits) - 1))); +} +} // namespace detail + template inline TF full_est_dist( const TI* quantized_vec, const TF* query, - TF (*ip_func_)(const TF*, const TI*, size_t), + TF (*ip_func)(const TF*, const TI*, size_t), size_t dim, size_t bits, TF f_add, @@ -432,10 +311,8 @@ inline TF full_est_dist( TF g_add, TF k1xsumq ) { - TF est_dist = f_add + g_add + - (f_rescale * (ip_func_(query, quantized_vec, dim) + - k1xsumq * static_cast((1 << bits) - 1))); - - return est_dist; + return detail::full_est_dist_from_ip( + ip_func(query, quantized_vec, dim), bits, f_add, f_rescale, g_add, k1xsumq + ); } } // namespace rabitqlib::quant diff --git a/include/rabitqlib/index/estimator.hpp b/include/rabitqlib/core/estimator.hpp similarity index 53% rename from include/rabitqlib/index/estimator.hpp rename to include/rabitqlib/core/estimator.hpp index 775aaee8..a3fbac35 100644 --- a/include/rabitqlib/index/estimator.hpp +++ b/include/rabitqlib/core/estimator.hpp @@ -1,60 +1,36 @@ #pragma once -#include #include #include -#include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/query.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/warmup_space.hpp" namespace rabitqlib { -/** - * @brief Use FastScan to estimate batch distance - * - * @param batch_data batch data, refer to BatchDataMap in data_layout.hpp - * @param q_obj query object - * @param padded_dim dim, must be multiple of 16 - * @param est_distance estimated distance - * @param low_distance lower bound of distance - * @param ip_x0_qr intermediate result for re-ranking - * @param use_hacc if use high accuracy fastscan - */ -inline void split_batch_estdist( - const char* batch_data, - const SplitBatchQuery& q_obj, - size_t padded_dim, - float* est_distance, - float* low_distance, - float* ip_x0_qr, - bool use_hacc -) { - simd::split_batch_estdist( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc - ); -} +namespace detail { +struct DistanceEstimate { + float distance; + float lower_bound; + float inner_product; +}; +} // namespace detail /** - * @brief Use ex-data bits to get more accurate distance + * @brief Refine a split-code estimate using the remaining ex_bits bits. * - * @tparam Query - * @param ex_data ex data, refer to ExDataMap in data_layout.hpp - * @param ip_func_ inner product function for compactly stored ex codes - * @param q_obj - * @param padded_dim - * @param ex_bits - * @param ip_x0_qr intermediate result generated by 1-bit distance estimation - * @return float + * Query and codes must share the rotation, padded dimension, metric and bit + * configuration. Set the query's centroid correction before scoring. + * ex_data uses ExDataMap from core/data_layout.hpp; ip_func must decode + * its packed ex_bits layout. ip_x0_qr is the base-code inner product for this + * same vector and query, as returned by the matching base-distance estimator. + * Returns the refined squared-L2 or 1 - dot estimate, without an error margin. */ template inline float split_distance_boosting( const char* ex_data, - float (*ip_func_)(const float*, const uint8_t*, size_t), + float (*ip_func)(const float*, const uint8_t*, size_t), const Query& q_obj, size_t padded_dim, size_t ex_bits, @@ -66,7 +42,7 @@ inline float split_distance_boosting( cur_ex.f_add_ex() + q_obj.g_add() + (cur_ex.f_rescale_ex() * (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func_(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); return ex_dist; } @@ -75,7 +51,7 @@ template inline void split_single_fulldist_direct( const char* bin_data, const char* ex_data, - float (*ip_func_)(const float*, const uint8_t*, size_t), + float (*ip_func)(const float*, const uint8_t*, size_t), const Query& q_obj, size_t padded_dim, size_t ex_bits, @@ -88,7 +64,6 @@ inline void split_single_fulldist_direct( ConstBinDataMap cur_bin(bin_data, padded_dim); ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); - // [TODO: optimize this function] // Direct HNSW kernels retain their legacy word-typed ABI. Generic kernels use the // byte-oriented overload and do not manufacture a typed view of packed storage. ip_x0_qr = Kernel::mask_ip_x0_q( @@ -101,52 +76,15 @@ inline void split_single_fulldist_direct( cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func_(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); low_dist = est_dist - (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); } -/** - * @brief Batch distance estimation for qg. Here, we do not need intermediate results and - * lower bound - */ -template -inline void qg_batch_estdist( - const char* batch_data, const BatchQuery& q_obj, size_t padded_dim, T* est_distance -) { - static_assert(std::is_same_v, "FastScan accumulation requires int32_t"); - if constexpr (std::is_same_v) { - simd::qg_batch_estdist(batch_data, q_obj, padded_dim, est_distance); - return; - } - - ConstQGBatchDataMap cur_batch(batch_data, padded_dim); - std::array accu_res{}; - fastscan::accumulate(cur_batch.bin_code(), q_obj.lut(), accu_res.data(), padded_dim); - ConstRowMajorArrayMap ip_arr(accu_res.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 est_dist_arr(est_distance, 1, fastscan::kBatchSize); - - est_dist_arr = f_add_arr + q_obj.g_add() + - (f_rescale_arr * (q_obj.delta() * (ip_arr.template cast()) + - q_obj.sum_vl_lut() + q_obj.k1xsumq())); -} - -/** - * @brief Full bits distance estimation for a single vector. - */ inline void split_single_fulldist( const char* bin_data, const char* ex_data, - float (*ip_func_)(const float*, const uint8_t*, size_t), + float (*ip_func)(const float*, const uint8_t*, size_t), const SplitSingleQuery& q_obj, size_t padded_dim, size_t ex_bits, @@ -159,21 +97,17 @@ inline void split_single_fulldist( ConstBinDataMap cur_bin(bin_data, padded_dim); ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); - // [TODO: optimize this function] ip_x0_qr = mask_ip_x0_q(q_obj.rotated_query(), cur_bin.bin_code(), padded_dim); est_dist = cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func_(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); low_dist = est_dist - (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); } -/** - * @brief 1-bit distance estimation for a single vector (no FastScan) - */ inline void split_single_estdist( const char* bin_data, const SplitSingleQuery& q_obj, @@ -287,4 +221,160 @@ inline void xy_single_full_dist( low_dist = est_dist - (cur_base.f_error() * g_error / static_cast(1 << ex_bits)); } +// Named-result functions keep the numerical implementation and legacy output-reference +// interfaces intact. Internal callers do not need to remember output argument order. +template +inline detail::DistanceEstimate estimate_split_single_full_direct( + const char* bin_data, + const char* ex_data, + float (*ip_func)(const float*, const uint8_t*, size_t), + const Query& q_obj, + size_t padded_dim, + size_t ex_bits, + float g_add, + float g_error +) { + detail::DistanceEstimate result; + split_single_fulldist_direct( + bin_data, + ex_data, + ip_func, + q_obj, + padded_dim, + ex_bits, + result.distance, + result.lower_bound, + result.inner_product, + g_add, + g_error + ); + return result; +} + +inline detail::DistanceEstimate estimate_split_single_full( + const char* bin_data, + const char* ex_data, + float (*ip_func)(const float*, const uint8_t*, size_t), + const SplitSingleQuery& q_obj, + size_t padded_dim, + size_t ex_bits, + float g_add, + float g_error +) { + detail::DistanceEstimate result; + split_single_fulldist( + bin_data, + ex_data, + ip_func, + q_obj, + padded_dim, + ex_bits, + result.distance, + result.lower_bound, + result.inner_product, + g_add, + g_error + ); + return result; +} + +inline detail::DistanceEstimate estimate_split_single( + const char* bin_data, + const SplitSingleQuery& q_obj, + size_t padded_dim, + float g_add = 0, + float g_error = 0 +) { + detail::DistanceEstimate result; + split_single_estdist( + bin_data, + q_obj, + padded_dim, + result.inner_product, + result.distance, + result.lower_bound, + g_add, + g_error + ); + return result; +} + +template +inline detail::DistanceEstimate estimate_split_single_direct( + const char* bin_data, + const SplitSingleQuery& q_obj, + size_t padded_dim, + float g_add = 0, + float g_error = 0 +) { + detail::DistanceEstimate result; + split_single_estdist_direct( + bin_data, + q_obj, + padded_dim, + result.inner_product, + result.distance, + result.lower_bound, + g_add, + g_error + ); + return result; +} + +inline detail::DistanceEstimate estimate_xy_base( + const char* base_data, + float (*base_ip_func)(const float*, const uint8_t*, size_t), + const SplitSingleQuery& q_obj, + size_t padded_dim, + size_t base_bits, + float g_add = 0, + float g_error = 0 +) { + detail::DistanceEstimate result; + xy_single_base_dist( + base_data, + base_ip_func, + q_obj, + padded_dim, + base_bits, + result.inner_product, + result.distance, + result.lower_bound, + g_add, + g_error + ); + return result; +} + +inline detail::DistanceEstimate estimate_xy_full( + const char* base_data, + const char* ex_data, + float (*base_ip_func)(const float*, const uint8_t*, size_t), + float (*ex_ip_func)(const float*, const uint8_t*, size_t), + const SplitSingleQuery& q_obj, + size_t padded_dim, + size_t base_bits, + size_t ex_bits, + float g_add = 0, + float g_error = 0 +) { + detail::DistanceEstimate result; + xy_single_full_dist( + base_data, + ex_data, + base_ip_func, + ex_ip_func, + q_obj, + padded_dim, + base_bits, + ex_bits, + result.distance, + result.lower_bound, + result.inner_product, + g_add, + g_error + ); + return result; +} + } // namespace rabitqlib diff --git a/include/rabitqlib/core/metric.hpp b/include/rabitqlib/core/metric.hpp new file mode 100644 index 00000000..e3d224c8 --- /dev/null +++ b/include/rabitqlib/core/metric.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include +#include + +namespace rabitqlib { +enum MetricType : std::uint8_t { METRIC_L2, METRIC_IP }; + +inline void validate_metric_type(MetricType metric_type) { + if (metric_type != METRIC_L2 && metric_type != METRIC_IP) { + throw std::invalid_argument("Metric type must be METRIC_L2 or METRIC_IP"); + } +} + +} // namespace rabitqlib diff --git a/include/rabitqlib/quantization/pack_excode.hpp b/include/rabitqlib/core/pack_excode.hpp similarity index 88% rename from include/rabitqlib/quantization/pack_excode.hpp rename to include/rabitqlib/core/pack_excode.hpp index 15b8ae02..217ea5b0 100644 --- a/include/rabitqlib/quantization/pack_excode.hpp +++ b/include/rabitqlib/core/pack_excode.hpp @@ -5,7 +5,7 @@ #include #include -#include "rabitqlib/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" namespace rabitqlib::quant::rabitq_impl::ex_bits { inline void packing_1bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { @@ -51,10 +51,11 @@ inline void packing_8bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t } /** - * @brief Packing ex_bits code to save space. For example, two 4-bit code will be - * stored as 1 uint8. To compute inner product with the support of SIMD, the - * packed codes need to be stored in different patterns. For details, please check the - * code and comments for certain number of bits. + * @brief Pack one vector's extra-bit codes into dim * ex_bits / 8 bytes. + * Layouts depend on bit width, not the selected SIMD backend. See + * docs/docs/compact_code.md, "How to read the layouts" and "Final 32-coordinate + * tail", for coordinate-to-bit mappings. The scalar mapping is implemented by + * packed_bit and tail_code in src/simd/portable_generic.cpp. * * @param o_raw unpacked code, code for each dim is represented by uint8 * @param o_compact compact format of code diff --git a/include/rabitqlib/core/query.hpp b/include/rabitqlib/core/query.hpp new file mode 100644 index 00000000..e053dc47 --- /dev/null +++ b/include/rabitqlib/core/query.hpp @@ -0,0 +1,109 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib { +// The binary-query preparation path currently supports float inputs. +// Use SplitSingleQuery; the template spelling does not imply double support. +template +class SplitSingleQuery { + private: + const T* rotated_query_; // Borrowed; keep it alive and unchanged during scoring. + std::vector query_bin_; + T G_add_ = 0; + T G_k1xSumq_; + T G_kbaseSumq_; + T G_kbxSumq_; + T G_error_ = 0; + T delta_; + T vl_; + MetricType metric_type_ = METRIC_L2; + + public: + static constexpr size_t kNumBits = 4; + explicit SplitSingleQuery( + const T* rotated_query, + size_t padded_dim, + size_t ex_bits, + quant::RabitqConfig config, + MetricType metric_type = METRIC_L2, + size_t base_bits = 1 + ) + : rotated_query_(rotated_query) + , query_bin_(div_round_up(padded_dim, 64) * kNumBits, 0) { + if (base_bits < 1 || base_bits > 8 || ex_bits > 8 || base_bits + ex_bits > 9) { + throw std::invalid_argument( + "SplitSingleQuery requires base_bits in [1, 8], ex_bits in [0, 8], " + "and base_bits + ex_bits <= 9" + ); + } + + validate_metric_type(metric_type); + + float c_1 = -static_cast((1 << 1) - 1) / 2.F; + float c_base = -static_cast((1U << base_bits) - 1) / 2.F; + float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.F; + T sumq = + std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); + + G_k1xSumq_ = sumq * c_1; + G_kbaseSumq_ = sumq * c_base; + G_kbxSumq_ = sumq * c_b; + + metric_type_ = metric_type; + + std::vector quant_query(padded_dim); + quant::quantize_scalar( + rotated_query, padded_dim, kNumBits, quant_query.data(), delta_, vl_, config + ); + rabitqlib::new_transpose_bin_512( + quant_query.data(), query_bin_.data(), padded_dim, kNumBits + ); + } + + [[nodiscard]] size_t num_bits() const { return kNumBits; } + + [[nodiscard]] const uint64_t* query_bin() const { return query_bin_.data(); } + + [[nodiscard]] const T* rotated_query() const { return rotated_query_; } + + [[nodiscard]] T delta() const { return delta_; } + + [[nodiscard]] T vl() const { return vl_; } + + [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } + + // Equals k1xsumq() when base_bits == 1. + [[nodiscard]] T kbase_sumq() const { return G_kbaseSumq_; } + + [[nodiscard]] T kbxsumq() const { return G_kbxSumq_; } + + [[nodiscard]] T g_add() const { return G_add_; } + + [[nodiscard]] T g_error() const { return G_error_; } + + // residual_norm is ||query - centroid||, not squared distance. For IP also + // supply query_centroid_ip = dot(query, centroid). Set before scoring. + void set_g_add(T residual_norm, T query_centroid_ip = 0) { + if (metric_type_ == METRIC_L2) { + G_add_ = residual_norm * residual_norm; + G_error_ = residual_norm; + } else if (metric_type_ == METRIC_IP) { + G_add_ = -query_centroid_ip; + G_error_ = residual_norm; + } + } + + void set_g_error(T residual_norm) { G_error_ = residual_norm; } +}; + +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/core/rotator.hpp similarity index 79% rename from include/rabitqlib/utils/rotator.hpp rename to include/rabitqlib/core/rotator.hpp index 47dbfd10..49297aa7 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/core/rotator.hpp @@ -8,22 +8,22 @@ #include #include #include +#include #include #include #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/simd/matrix_dispatch.hpp" -#include "rabitqlib/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/types.hpp" // IWYU pragma: export +#include "rabitqlib/third/Eigen/QR" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" namespace rabitqlib { -enum class RotatorType : uint8_t { MatrixRotator, FhtKacRotator }; - -// abstract rotator template class Rotator { protected: @@ -39,14 +39,13 @@ class Rotator { virtual void save(std::ofstream&) const = 0; // Buffer I/O virtual void load(const char* data) = 0; - virtual void save(char* data) const = 0; // dump to buffer + virtual void save(char* data) const = 0; virtual size_t dump_bytes() const = 0; [[nodiscard]] size_t size() const { return this->padded_dim_; } }; namespace rotator_impl { -// get padding requirement for different rotator inline size_t padding_requirement(size_t dim, RotatorType type) { if (type == RotatorType::MatrixRotator) { return dim; @@ -60,13 +59,38 @@ inline size_t padding_requirement(size_t dim, RotatorType type) { throw std::invalid_argument("Invalid rotator type in padding_requirement()"); } +inline size_t checked_padded_dimension(size_t dim, RotatorType type, size_t padded_dim) { + if (dim == 0) { + throw std::invalid_argument("Rotator dimension must be positive"); + } + if (padded_dim == 0) { + padded_dim = padding_requirement(dim, type); + } + + if (padded_dim < dim) { + throw std::invalid_argument( + "Padded rotator dimension must not be smaller than input" + ); + } + + if (padded_dim != padding_requirement(padded_dim, type)) { + throw std::invalid_argument("Invalid padded dimension for the rotator type"); + } + + return padded_dim; +} + +// Internal loading path: allocate storage, then restore the saved rotation before use. +struct AllocateOnly {}; + template class MatrixRotator : public Rotator { private: - RowMajorMatrix rand_mat_; // Rotation matrix + RowMajorMatrix rand_mat_; + public: explicit MatrixRotator(size_t dim, size_t padded_dim) - : Rotator(dim, padded_dim), rand_mat_(dim, padded_dim) { + : MatrixRotator(dim, padded_dim, AllocateOnly{}) { RowMajorMatrix rand = random_gaussian_matrix(padded_dim, padded_dim); Eigen::HouseholderQR> qr(rand); RowMajorMatrix q_inv = qr.householderQ().transpose( @@ -76,15 +100,13 @@ class MatrixRotator : public Rotator { // the vector to be rotated to padded dimension std::memcpy(&rand_mat_(0, 0), &q_inv(0, 0), sizeof(T) * dim * padded_dim); } + MatrixRotator(size_t dim, size_t padded_dim, AllocateOnly) + : Rotator(dim, padded_dim), rand_mat_(dim, padded_dim) {} + MatrixRotator() = default; ~MatrixRotator() override = default; - MatrixRotator& operator=(const MatrixRotator& other) { - this->dim_ = other.dim_; - this->padded_dim_ = other.padded_dim_; - this->rand_mat_ = other.rand_mat_; - return *this; - } + MatrixRotator& operator=(const MatrixRotator&) = default; void load(std::ifstream& input) override { input.read( @@ -146,21 +168,20 @@ class FhtKacRotator : public Rotator { explicit FhtKacRotator( size_t dim, size_t padded_dim, uint32_t seed = std::random_device{}() ) - : Rotator(dim, padded_dim) { - if (dim < 64 || padded_dim < dim || padded_dim % 32 != 0 || padded_dim > kMaxDim) { - throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); - } - flip_.resize((padded_dim / kByteLen) * 4); + : FhtKacRotator(dim, padded_dim, AllocateOnly{}) { std::mt19937 gen(seed); - - // Uniform distribution in the range [0, 255] std::uniform_int_distribution dist(0, 255); - - // Generate a single random uint8_t value for (auto& i : flip_) { i = static_cast(dist(gen)); } + } + FhtKacRotator(size_t dim, size_t padded_dim, AllocateOnly) + : Rotator(dim, padded_dim) { + if (dim < 64 || padded_dim < dim || padded_dim % 32 != 0 || padded_dim > kMaxDim) { + throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); + } + flip_.resize((padded_dim / kByteLen) * 4); size_t bottom_log_dim = floor_log2(dim); trunc_dim_ = size_t{1} << bottom_log_dim; fac_ = 1.0F / std::sqrt(static_cast(trunc_dim_)); @@ -209,9 +230,20 @@ class FhtKacRotator : public Rotator { ); } }; +inline std::unique_ptr> make_rotator_for_loading( + size_t dim, RotatorType type, size_t padded_dim +) { + padded_dim = checked_padded_dimension(dim, type, padded_dim); + if (type == RotatorType::FhtKacRotator) { + return std::make_unique(dim, padded_dim, AllocateOnly{}); + } + return std::make_unique>(dim, padded_dim, AllocateOnly{}); +} } // namespace rotator_impl -// for given dim & type, set rotator, return padded dimension +// Return an owning pointer; wrap it in RAII and use size() for the output dimension. +// FHT/Kac supports float and uses seed. Matrix rotation supports float/double; +// it samples its own matrix and does not use this factory's seed argument. template Rotator* choose_rotator( size_t dim, @@ -219,28 +251,14 @@ Rotator* choose_rotator( size_t padded_dim = 0, uint32_t seed = std::random_device{}() ) { - if (dim == 0) { - throw std::invalid_argument("Rotator dimension must be positive"); - } - if (padded_dim == 0) { - padded_dim = rotator_impl::padding_requirement(dim, type); - } - - if (padded_dim < dim) { - throw std::invalid_argument( - "Padded rotator dimension must not be smaller than input" - ); - } - - if (padded_dim != rotator_impl::padding_requirement(padded_dim, type)) { - throw std::invalid_argument("Invalid padded dimension for the rotator type"); - } + padded_dim = rotator_impl::checked_padded_dimension(dim, type, padded_dim); if (type == RotatorType::FhtKacRotator) { - if (!std::is_same_v) { + if constexpr (std::is_same_v) { + return ::new rotator_impl::FhtKacRotator(dim, padded_dim, seed); + } else { throw std::invalid_argument("FhtKacRotator only supports float"); } - return ::new rotator_impl::FhtKacRotator(dim, padded_dim, seed); } if (type == RotatorType::MatrixRotator) { diff --git a/include/rabitqlib/core/types.hpp b/include/rabitqlib/core/types.hpp new file mode 100644 index 00000000..d1638c52 --- /dev/null +++ b/include/rabitqlib/core/types.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include + +namespace rabitqlib { +enum ScalarQuantizerType : std::uint8_t { RECONSTRUCTION, UNBIASED_ESTIMATION, PLAIN }; +enum class RotatorType : uint8_t { MatrixRotator, FhtKacRotator }; + +namespace quant { +struct RabitqConfig { + double t_const = -1; + explicit RabitqConfig() = default; + RabitqConfig(RabitqConfig const&) = default; + RabitqConfig(RabitqConfig&&) = default; + RabitqConfig& operator=(const RabitqConfig&) = default; + RabitqConfig& operator=(RabitqConfig&&) = default; +}; +} // namespace quant +} // namespace rabitqlib diff --git a/include/rabitqlib/defines.hpp b/include/rabitqlib/defines.hpp deleted file mode 100644 index ef7402ee..00000000 --- a/include/rabitqlib/defines.hpp +++ /dev/null @@ -1,77 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "rabitqlib/third/Eigen/Dense" -#include "rabitqlib/utils/bitops.hpp" // IWYU pragma: keep; used by BIT_ID - -#define BIT_ID(x) (rabitqlib::bitops::popcount32((x)-1)) -#define LOWBIT(x) ((x) & (-(x))) - -namespace rabitqlib { - -using PID = uint32_t; - -constexpr uint32_t kPidMax = 0xFFFFFFFF; - -template -using RowMajorMatrix = Eigen::Matrix; - -template -using RowMajorMatrixMap = Eigen::Map>; - -template -using ConstRowMajorMatrixMap = Eigen::Map>; - -template -using RowMajorArray = Eigen::Array; - -template -using Vector = Eigen::Matrix; - -template -using RowMajorArrayMap = Eigen::Map>; - -template -using ConstRowMajorArrayMap = Eigen::Map>; - -template -using VectorMap = Eigen::Map>; - -template -using ConstVectorMap = Eigen::Map>; - -template -struct AnnCandidate { - TP id = 0; - T distance = std::numeric_limits::max(); - - AnnCandidate() = default; - explicit AnnCandidate(TP vec_id, T dis) : id(vec_id), distance(dis) {} - - friend bool operator<(const AnnCandidate& first, const AnnCandidate& second) { - return first.distance < second.distance; - } - friend bool operator>(const AnnCandidate& first, const AnnCandidate& second) { - return first.distance > second.distance; - } - friend bool operator>=(const AnnCandidate& first, const AnnCandidate& second) { - return first.distance >= second.distance; - } - friend bool operator<=(const AnnCandidate& first, const AnnCandidate& second) { - return first.distance <= second.distance; - } -}; - -enum MetricType : std::uint8_t { METRIC_L2, METRIC_IP }; - -inline void validate_metric_type(MetricType metric_type) { - if (metric_type != METRIC_L2 && metric_type != METRIC_IP) { - throw std::invalid_argument("Metric type must be METRIC_L2 or METRIC_IP"); - } -} - -enum ScalarQuantizerType : std::uint8_t { RECONSTRUCTION, UNBIASED_ESTIMATION, PLAIN }; -} // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/data_layout.hpp b/include/rabitqlib/fastscan/data_layout.hpp new file mode 100644 index 00000000..ac305ff2 --- /dev/null +++ b/include/rabitqlib/fastscan/data_layout.hpp @@ -0,0 +1,128 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/utils/packed.hpp" + +namespace rabitqlib { +namespace detail { +// Standard and QG batches share the code/add/rescale prefix; only standard +// batches store the trailing error factors. Counts include all physical lanes. +template +struct BatchCodeOffsets { + static size_t add(size_t padded_dim) { return padded_dim * fastscan::kBatchSize / 8; } + static size_t rescale(size_t padded_dim) { + return add(padded_dim) + sizeof(T) * fastscan::kBatchSize; + } + static size_t error(size_t padded_dim) { + return rescale(padded_dim) + sizeof(T) * fastscan::kBatchSize; + } +}; +} // namespace detail + +template +struct BatchDataMap { + public: + explicit BatchDataMap(char* data, size_t padded_dim) + : batch_bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) + , f_error_(data + detail::BatchCodeOffsets::error(padded_dim)) {} + + [[nodiscard]] uint8_t* bin_code() { return batch_bin_code_; } + [[nodiscard]] detail::PackedArrayView f_add() { return f_add_; } + [[nodiscard]] detail::PackedArrayView f_rescale() { return f_rescale_; } + [[nodiscard]] detail::PackedArrayView f_error() { return f_error_; } + + static size_t data_bytes(size_t padded_dim) { + return detail::BatchCodeOffsets::error(padded_dim) + + sizeof(T) * fastscan::kBatchSize; + } + + private: + uint8_t* batch_bin_code_; + detail::PackedArrayView f_add_; + detail::PackedArrayView f_rescale_; + detail::PackedArrayView f_error_; +}; + +template +struct ConstBatchDataMap { + public: + explicit ConstBatchDataMap(const char* data, size_t padded_dim) + : batch_bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) + , f_error_(data + detail::BatchCodeOffsets::error(padded_dim)) {} + + [[nodiscard]] const uint8_t* bin_code() const { return batch_bin_code_; } + [[nodiscard]] detail::ConstPackedArrayView f_add() const { return f_add_; } + [[nodiscard]] detail::ConstPackedArrayView f_rescale() const { return f_rescale_; } + [[nodiscard]] detail::ConstPackedArrayView f_error() const { return f_error_; } + + private: + const uint8_t* batch_bin_code_; + detail::ConstPackedArrayView f_add_; + detail::ConstPackedArrayView f_rescale_; + detail::ConstPackedArrayView f_error_; +}; + +// Two-factor batch layout used by FastScan consumers, including SymphonyQG. +template +struct QGBatchDataMap { + public: + explicit QGBatchDataMap(char* data, size_t padded_dim) + : batch_bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) {} + + [[nodiscard]] uint8_t* bin_code() { return batch_bin_code_; } + [[nodiscard]] detail::PackedArrayView f_add() { return f_add_; } + [[nodiscard]] detail::PackedArrayView f_rescale() { return f_rescale_; } + + static size_t data_bytes(size_t padded_dim) { + const size_t max_size = std::numeric_limits::max(); + if (padded_dim > max_size / fastscan::kBatchSize) { + throw std::length_error("QGBatchDataMap storage size exceeds size_t"); + } + const size_t code_bytes = detail::BatchCodeOffsets::add(padded_dim); + const size_t factor_bytes = sizeof(T) * fastscan::kBatchSize * 2; + if (code_bytes > max_size - factor_bytes) { + throw std::length_error("QGBatchDataMap storage size exceeds size_t"); + } + return code_bytes + factor_bytes; + } + + private: + uint8_t* batch_bin_code_; + detail::PackedArrayView f_add_; + detail::PackedArrayView f_rescale_; +}; + +template +struct ConstQGBatchDataMap { + public: + explicit ConstQGBatchDataMap(const char* data, size_t padded_dim) + : batch_bin_code_(reinterpret_cast(data)) + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) {} + + [[nodiscard]] const uint8_t* bin_code() const { return batch_bin_code_; } + [[nodiscard]] detail::ConstPackedArrayView f_add() const { return f_add_; } + [[nodiscard]] detail::ConstPackedArrayView f_rescale() const { return f_rescale_; } + + static size_t data_bytes(size_t padded_dim) { + return QGBatchDataMap::data_bytes(padded_dim); + } + + private: + const uint8_t* batch_bin_code_; + detail::ConstPackedArrayView f_add_; + detail::ConstPackedArrayView f_rescale_; +}; + +} // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/detail/encoder.hpp b/include/rabitqlib/fastscan/detail/encoder.hpp new file mode 100644 index 00000000..f91e74cd --- /dev/null +++ b/include/rabitqlib/fastscan/detail/encoder.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include +#include +#include + +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" + +namespace rabitqlib::quant::rabitq_impl::one_bit { +// Encoding requires a positive padded_dim divisible by 8; index pipelines use multiples +// of 32. packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail +// padding. +template < + typename T, + bool Parallel = false, + typename FAdd, + typename FRescale, + typename FError> +inline void one_bit_batch_code( + const T* data, + const T* centroid, + size_t num, + size_t padded_dim, + uint8_t* packed_code, + FAdd f_add, + FRescale f_recale, + FError f_error, + MetricType metric_type = METRIC_L2 +) { + std::vector compact_codes(num * padded_dim / 8); + + one_bit_compact_codes( + data, + centroid, + num, + padded_dim, + compact_codes.data(), + f_add, + f_recale, + f_error, + metric_type + ); + fastscan::pack_codes(padded_dim, compact_codes.data(), num, packed_code); +} + +} // namespace rabitqlib::quant::rabitq_impl::one_bit diff --git a/include/rabitqlib/simd/estimator_dispatch.hpp b/include/rabitqlib/fastscan/detail/estimator_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/estimator_dispatch.hpp rename to include/rabitqlib/fastscan/detail/estimator_dispatch.hpp diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp new file mode 100644 index 00000000..1b74d77e --- /dev/null +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -0,0 +1,159 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib { + +template +class Lut { + static constexpr size_t kNumBits = 8; + static constexpr size_t kNumBitsHacc = 16; + static_assert(std::is_floating_point_v, "T must be a floating-point type in Lut"); + + private: + struct Shape { + size_t table_length; + size_t storage_length; + }; + + [[nodiscard]] static Shape checked_shape(size_t padded_dim, bool use_hacc) { + if (padded_dim == 0 || padded_dim % 16 != 0) { + throw std::invalid_argument( + "FastScan dimension must be a positive multiple of 16" + ); + } + if (padded_dim > std::numeric_limits::max() / 4) { + throw std::length_error("FastScan lookup table is too large"); + } + const size_t table_length = padded_dim * 4; + const size_t copies = use_hacc ? 2 : 1; + if (table_length > std::numeric_limits::max() / copies) { + throw std::length_error("FastScan lookup table is too large"); + } + return {table_length, table_length * copies}; + } + + size_t table_length_ = 0; + std::vector lut_; + T delta_ = 0; + T sum_vl_lut_ = 0; + + // Both entry points validate the shape once. Only this method manages owned + // storage; numerical preparation below receives valid shape and scratch. + void reset_with_shape( + const T* rotated_query, + size_t padded_dim, + T* lut_scratch, + uint16_t* lut_u16_scratch, + bool use_hacc, + Shape shape + ) { + if (rotated_query == nullptr || lut_scratch == nullptr) { + throw std::invalid_argument("Lut::reset requires query and LUT scratch"); + } + if (use_hacc && lut_u16_scratch == nullptr) { + throw std::invalid_argument( + "Lut::reset requires uint16 scratch for high-accuracy lookup tables" + ); + } + lut_.resize(shape.storage_length); + table_length_ = shape.table_length; + build_lut(rotated_query, padded_dim, lut_scratch, lut_u16_scratch, use_hacc); + } + + template + void quantize_entries(Code* output, const T* values, T lo) const { + if constexpr (std::is_same_v) { + scalar_quantize(output, values, table_length_, lo, delta_); + } else { + // Match the portable scalar quantizer without narrowing the query to float. + const T reciprocal = delta_ == 0 ? T{0} : T{1} / delta_; + for (size_t i = 0; i < table_length_; ++i) { + const T rounded = std::round((values[i] - lo) * reciprocal); + output[i] = static_cast(std::clamp( + rounded, T{0}, static_cast(std::numeric_limits::max()) + )); + } + } + } + + void build_lut( + const T* rotated_query, + size_t padded_dim, + T* lut_scratch, + uint16_t* lut_u16_scratch, + bool use_hacc + ) { + fastscan::pack_lut(padded_dim, rotated_query, lut_scratch); + T vl_lut; + T vr_lut; + data_range(lut_scratch, table_length_, vl_lut, vr_lut); + + if (use_hacc) { + delta_ = (vr_lut - vl_lut) / ((1 << kNumBitsHacc) - 1); + quantize_entries(lut_u16_scratch, lut_scratch, vl_lut); + fastscan::transfer_lut_hacc(lut_u16_scratch, padded_dim, lut_.data()); + } else { + delta_ = (vr_lut - vl_lut) / ((1 << kNumBits) - 1); + quantize_entries(lut_.data(), lut_scratch, vl_lut); + } + + const size_t num_tables = table_length_ / 16; + sum_vl_lut_ = vl_lut * static_cast(num_tables); + } + + // Construct owned storage at its known size. Growing an initially empty + // vector through reset adds work to every fresh query's preparation. + Lut(const T* rotated_query, size_t padded_dim, bool use_hacc, Shape shape) + : table_length_(shape.table_length), lut_(shape.storage_length) { + std::vector lut_values(shape.table_length); + std::vector lut_u16(use_hacc ? shape.table_length : 0); + reset_with_shape( + rotated_query, padded_dim, lut_values.data(), lut_u16.data(), use_hacc, shape + ); + } + + public: + explicit Lut() = default; + explicit Lut(const T* rotated_query, size_t padded_dim, bool use_hacc = false) + : Lut(rotated_query, padded_dim, use_hacc, checked_shape(padded_dim, use_hacc)) {} + + /** + * Reinitialize owned LUT bytes while reusing their allocation. Scratch buffers + * are borrowed only for this call and may be overwritten/released afterwards. + * lut_scratch must contain padded_dim * 4 values of T; high-accuracy mode + * additionally requires padded_dim * 4 uint16_t values in lut_u16_scratch. + */ + void reset( + const T* rotated_query, + size_t padded_dim, + T* lut_scratch, + uint16_t* lut_u16_scratch = nullptr, + bool use_hacc = false + ) { + const Shape shape = checked_shape(padded_dim, use_hacc); + reset_with_shape( + rotated_query, padded_dim, lut_scratch, lut_u16_scratch, use_hacc, shape + ); + } + + // Empty (including moved-from) storage is not a prepared query. + [[nodiscard]] size_t padded_dim() const { return lut_.empty() ? 0 : table_length_ / 4; } + [[nodiscard]] bool use_hacc() const { return lut_.size() > table_length_; } + + [[nodiscard]] const uint8_t* lut() const { return lut_.data(); }; + [[nodiscard]] T delta() const { return delta_; }; + [[nodiscard]] T sum_vl() const { return sum_vl_lut_; }; +}; +} // namespace rabitqlib diff --git a/include/rabitqlib/simd/fastscan_dispatch.hpp b/include/rabitqlib/fastscan/detail/simd.hpp similarity index 100% rename from include/rabitqlib/simd/fastscan_dispatch.hpp rename to include/rabitqlib/fastscan/detail/simd.hpp diff --git a/include/rabitqlib/fastscan/encoder.hpp b/include/rabitqlib/fastscan/encoder.hpp new file mode 100644 index 00000000..7a11deac --- /dev/null +++ b/include/rabitqlib/fastscan/encoder.hpp @@ -0,0 +1,136 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/encoder.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" + +namespace rabitqlib::quant { +// Single-batch entry points accept 0..32 rows. Zero rows leave output unchanged. +// Allocate full BatchDataMap/QGBatchDataMap storage even for a partial batch; +// only the first num factors/results are meaningful. Inputs are borrowed for the call. +// Parallel=true enables OpenMP when built with it, otherwise encoding remains serial. +template +inline void quantize_one_batch( + const T* data, + const T* centroid, + size_t num, + size_t padded_dim, + char* batch_data, + MetricType metric_type = METRIC_L2 +) { + if (num > fastscan::kBatchSize) { + throw std::invalid_argument( + "FastScan single-batch encoding accepts at most 32 vectors" + ); + } + if (num == 0) { + return; + } + BatchDataMap this_batch(batch_data, padded_dim); + + rabitq_impl::one_bit::one_bit_batch_code( + data, + centroid, + num, + padded_dim, + this_batch.bin_code(), + this_batch.f_add(), + this_batch.f_rescale(), + this_batch.f_error(), + metric_type + ); +} + +template +inline void quantize_one_batch( + const T* data, + size_t num, + size_t padded_dim, + char* batch_data, + MetricType metric_type = METRIC_L2 +) { + std::vector centroid(padded_dim, 0); + quantize_one_batch( + data, centroid.data(), num, padded_dim, batch_data, metric_type + ); +} + +template +static inline void quantize_qg_batch( + const T* data, + const T* centroid, + size_t num, + size_t padded_dim, + char* batch_data, + MetricType metric_type = METRIC_L2 +) { + if (num > fastscan::kBatchSize) { + throw std::invalid_argument( + "FastScan single-batch encoding accepts at most 32 vectors" + ); + } + if (num == 0) { + return; + } + std::array f_error{}; // we dont need this factor for qg + QGBatchDataMap cur_batch(batch_data, padded_dim); + + rabitq_impl::one_bit::one_bit_batch_code( + data, + centroid, + num, + padded_dim, + cur_batch.bin_code(), + cur_batch.f_add(), + cur_batch.f_rescale(), + f_error.data(), + metric_type + ); +} + +template +static inline void quantize_qg_batch( + const T* data, + size_t num, + size_t padded_dim, + char* batch_data, + MetricType metric_type = METRIC_L2 +) { + std::vector centroid(padded_dim, 0); + quantize_qg_batch(data, centroid.data(), num, padded_dim, batch_data, metric_type); +} + +inline void quantize_split_batch( + const float* data, + const float* centroid, + size_t num_points, + size_t padded_dim, + size_t ex_bits, + char* batch_data, + char* ex_data, + MetricType metric_type = METRIC_L2, + RabitqConfig config = RabitqConfig() +) { + quantize_one_batch(data, centroid, num_points, padded_dim, batch_data, metric_type); + + if (ex_bits > 0) { + for (size_t i = 0; i < num_points; ++i) { + quantize_compact_ex_bits( + data, centroid, padded_dim, ex_bits, ex_data, metric_type, config + ); + + ex_data += ExDataMap::data_bytes(padded_dim, ex_bits); + data += padded_dim; + } + } +} + +} // namespace rabitqlib::quant diff --git a/include/rabitqlib/fastscan/estimator.hpp b/include/rabitqlib/fastscan/estimator.hpp new file mode 100644 index 00000000..f7ed71b5 --- /dev/null +++ b/include/rabitqlib/fastscan/estimator.hpp @@ -0,0 +1,95 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" + +namespace rabitqlib { +/** + * @brief Use FastScan to estimate batch distance + * + * @param batch_data codes and factors in fastscan/data_layout.hpp's BatchDataMap layout + * @param q_obj prepared query with the matching centroid correction already set + * @param padded_dim dim, must be multiple of 16 + * @param est_distance estimated distance + * @param low_distance lower endpoint of the probabilistic distance interval; + * not a deterministic bound for every query + * @param ip_x0_qr intermediate result for re-ranking + * @param use_hacc must match the mode used to prepare q_obj + * All three output buffers need 32 floats, including for a partial batch. + * Throws std::invalid_argument for an unprepared or mismatched query. + * The margin is f_error * g_error; see docs/docs/rabitq/estimator.md, + * "Distance Estimation", for the factor definitions and probabilistic limits. + */ +inline void split_batch_estdist( + const char* batch_data, + const SplitBatchQuery& q_obj, + size_t padded_dim, + float* est_distance, + float* low_distance, + float* ip_x0_qr, + bool use_hacc +) { + if (padded_dim == 0 || q_obj.padded_dim() != padded_dim) { + throw std::invalid_argument( + "FastScan scoring requires a prepared query with matching dimension" + ); + } + if (q_obj.use_hacc() != use_hacc) { + throw std::invalid_argument("FastScan scoring mode must match the prepared query"); + } + simd::split_batch_estdist( + batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + ); +} + +/** + * Estimate distances from QGBatchDataMap without producing refinement or error terms. + * est_distance needs 32 values of T, even for a partial batch. + * Throws std::invalid_argument for an unprepared or mismatched query. + */ +template +inline void qg_batch_estdist( + const char* batch_data, const BatchQuery& q_obj, size_t padded_dim, T* est_distance +) { + static_assert(std::is_same_v, "FastScan accumulation requires int32_t"); + if (padded_dim == 0 || q_obj.padded_dim() != padded_dim) { + throw std::invalid_argument( + "FastScan scoring requires a prepared query with matching dimension" + ); + } + if constexpr (std::is_same_v) { + simd::qg_batch_estdist(batch_data, q_obj, padded_dim, est_distance); + return; + } + + ConstQGBatchDataMap cur_batch(batch_data, padded_dim); + std::array accu_res{}; + fastscan::accumulate(cur_batch.bin_code(), q_obj.lut(), accu_res.data(), padded_dim); + ConstRowMajorArrayMap ip_arr(accu_res.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 est_dist_arr(est_distance, 1, fastscan::kBatchSize); + + est_dist_arr = f_add_arr + q_obj.g_add() + + (f_rescale_arr * (q_obj.delta() * (ip_arr.template cast()) + + q_obj.sum_vl_lut() + q_obj.k1xsumq())); +} + +} // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index ed5757e8..423e0bf5 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -3,17 +3,12 @@ #pragma once -#include #include #include #include -#include "rabitqlib/defines.hpp" - namespace rabitqlib::fastscan { -constexpr static size_t kBatchSize = 32; // number of vectors in each batch - constexpr static std::array kPos = { 3 /*0000*/, 3 /*0001*/, @@ -33,127 +28,6 @@ constexpr static std::array kPos = { 3 /*1111*/, }; // all possible combination for a 4 bit string -constexpr static std::array kPerm0 = { - 0, - 8, - 1, - 9, - 2, - 10, - 3, - 11, - 4, - 12, - 5, - 13, - 6, - 14, - 7, - 15}; // data order of packed quantization code, please refer to code and the link - // offered above for detailed information - -template -static inline void get_column( - const T* src, size_t rows, size_t cols, size_t row, size_t col, TA& dest -) { - size_t k = 0; - size_t max_k = std::min(rows - row, dest.size()); - for (; k < max_k; ++k) { - dest[k] = src[((k + row) * cols) + col]; - } - if (k < dest.size()) { - std::fill(dest.begin() + k, dest.end(), 0); - } -} - -/** - * @brief Pack quantization codes into blocks; the data organization is illustrated in - * the link and kPerm0. Since we pack codes as 32-sized groups, if the num is not a multiple - * of 32, we have to use some space for these absent data - * - * @param padded_dim dimension of quantized data (i.e., quantization code) - * @param quantization_code quantizaiton code, stored as uint8 - * @param num number of quantization code - * @param blocks packed quantization code - */ -inline void pack_codes( - size_t padded_dim, const uint8_t* quantization_code, size_t num, uint8_t* blocks -) { - // consider codes is a matrix - // rows = number of vectors - // cols = number of uint8_t of one vector's code - size_t cols = padded_dim / 8; - - // Full batches can gather directly in output order, without a temporary column. - size_t row = 0; - for (; num - row >= kBatchSize; row += kBatchSize) { - for (size_t i = 0; i < cols; ++i) { - for (size_t j = 0; j < 16; ++j) { - const size_t first = row + kPerm0[j]; - const uint8_t a = quantization_code[first * cols + i]; - const uint8_t b = quantization_code[(first + 16) * cols + i]; - blocks[j] = (a >> 4) | (b & 0xF0); - blocks[j + 16] = (a & 0x0F) | (b << 4); - } - blocks += 32; - } - } - - // Only the final partial batch needs a zero-padded column. - if (row < num) { - std::array col; - for (size_t i = 0; i < cols; ++i) { - get_column(quantization_code, num, cols, row, i, col); - for (size_t j = 0; j < 16; ++j) { - // the lower 4 bits represent vector 0 to 15 - // the upper 4 bits represent vector 16 to 31 - const uint8_t a = col[kPerm0[j]]; - const uint8_t b = col[kPerm0[j] + 16]; - blocks[j] = (a >> 4) | (b & 0xF0); - blocks[j + 16] = (a & 0x0F) | (b << 4); - } - blocks += 32; - } - } -} - -/** - * @brief Inverse of pack_codes: recover the compact quantization codes of the first `num` - * vectors from their packed blocks. Padding lanes of a partial final batch are ignored. - * - * @param padded_dim dimension of quantized data (i.e., quantization code) - * @param blocks packed quantization code, as produced by pack_codes - * @param num number of quantization code to recover - * @param quantization_code recovered codes, num * padded_dim / 8 bytes - */ -inline void unpack_codes( - size_t padded_dim, const uint8_t* blocks, size_t num, uint8_t* quantization_code -) { - const size_t cols = padded_dim / 8; - - for (size_t row = 0; row < num; row += kBatchSize) { - for (size_t i = 0; i < cols; ++i) { - for (size_t j = 0; j < 16; ++j) { - // blocks[j] holds the high nibbles of vector kPerm0[j] and the low nibbles - // of vector kPerm0[j] + 16 packed together; blocks[j + 16] holds the rest. - const size_t first = row + kPerm0[j]; - const auto a = static_cast( - ((blocks[j] & 0x0F) << 4) | (blocks[j + 16] & 0x0F) - ); - const auto b = - static_cast((blocks[j] & 0xF0) | (blocks[j + 16] >> 4)); - if (first < num) { - quantization_code[first * cols + i] = a; - } - if (first + 16 < num) { - quantization_code[(first + 16) * cols + i] = b; - } - } - blocks += 32; - } - } -} - // Accumulate one batch into int32_t results. Throws if a result exceeds int32_t. // dim must be a positive multiple of 16. // Requires AVX2/FMA, AVX-512, or ARM NEON; otherwise throws std::runtime_error. @@ -174,7 +48,8 @@ inline void pack_lut(size_t dim, const T* __restrict__ query, T* __restrict__ lu for (size_t i = 0; i < num_codebook; ++i) { lut[0] = zero; for (size_t j = 1; j < 16; ++j) { - lut[j] = lut[j - LOWBIT(j)] + query[kPos[j]]; + const size_t lowbit = j & (size_t{0} - j); + lut[j] = lut[j - lowbit] + query[kPos[j]]; } lut += 16; query += 4; diff --git a/include/rabitqlib/fastscan/highacc_fastscan.hpp b/include/rabitqlib/fastscan/highacc_fastscan.hpp index 6efee5df..390c1f44 100644 --- a/include/rabitqlib/fastscan/highacc_fastscan.hpp +++ b/include/rabitqlib/fastscan/highacc_fastscan.hpp @@ -5,11 +5,11 @@ namespace rabitqlib::fastscan { /** - * @brief Change u16 lookup table to u8. Since we use more bits (higher accuracy) - * to quantize data vector by rabitq+, we also needs to increase the accuracy of data in - * lut. - * We split the higher & lower 8 bits of a u16 into two sub luts. - **/ + * @brief Store a 16-bit query LUT as low/high byte planes for HACC. + * Input contains dim * 4 uint16_t entries; output needs dim * 8 bytes. + * The runtime-selected backend determines the byte-plane layout. Pair this + * dispatched producer with accumulate_hacc rather than assuming an ISA layout. + */ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut); // Accumulate one batch into int32_t results. Throws if a result exceeds int32_t. diff --git a/include/rabitqlib/fastscan/pack_codes.hpp b/include/rabitqlib/fastscan/pack_codes.hpp new file mode 100644 index 00000000..7d152d10 --- /dev/null +++ b/include/rabitqlib/fastscan/pack_codes.hpp @@ -0,0 +1,122 @@ +// The implementation is largely based on the implementation of Faiss. +// https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan) + +#pragma once + +#include +#include +#include +#include + +namespace rabitqlib::fastscan { + +constexpr static size_t kBatchSize = 32; // number of vectors in each batch + +// Interleave lanes 0..7 and 8..15 for the 16-bit SIMD accumulators. +constexpr static std::array kPerm0 = { + 0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15}; + +template +static inline void get_column( + const T* src, size_t rows, size_t cols, size_t row, size_t col, TA& dest +) { + size_t k = 0; + size_t max_k = std::min(rows - row, dest.size()); + for (; k < max_k; ++k) { + dest[k] = src[((k + row) * cols) + col]; + } + if (k < dest.size()) { + std::fill(dest.begin() + k, dest.end(), 0); + } +} + +/** + * @brief Pack sign codes into 32-vector blocks, zero-filling missing tail lanes. + * For each input byte column, output bytes j and j + 16 pair vectors kPerm0[j] + * and kPerm0[j] + 16. The first output holds both inputs' high nibbles, the second + * their low nibbles; the first vector always occupies the output's low nibble. + * + * @param padded_dim dimension of quantized data (i.e., quantization code) + * @param quantization_code num rows of padded_dim / 8 compact sign-code bytes + * @param num number of logical vectors + * @param blocks output, ceil(num / 32) * 32 * padded_dim / 8 bytes + */ +inline void pack_codes( + size_t padded_dim, const uint8_t* quantization_code, size_t num, uint8_t* blocks +) { + // consider codes is a matrix + // rows = number of vectors + // cols = number of uint8_t of one vector's code + size_t cols = padded_dim / 8; + + // Full batches can gather directly in output order, without a temporary column. + size_t row = 0; + for (; num - row >= kBatchSize; row += kBatchSize) { + for (size_t i = 0; i < cols; ++i) { + for (size_t j = 0; j < 16; ++j) { + const size_t first = row + kPerm0[j]; + const uint8_t a = quantization_code[first * cols + i]; + const uint8_t b = quantization_code[(first + 16) * cols + i]; + blocks[j] = (a >> 4) | (b & 0xF0); + blocks[j + 16] = (a & 0x0F) | (b << 4); + } + blocks += 32; + } + } + + // Only the final partial batch needs a zero-padded column. + if (row < num) { + std::array col; + for (size_t i = 0; i < cols; ++i) { + get_column(quantization_code, num, cols, row, i, col); + for (size_t j = 0; j < 16; ++j) { + // the lower 4 bits represent vector 0 to 15 + // the upper 4 bits represent vector 16 to 31 + const uint8_t a = col[kPerm0[j]]; + const uint8_t b = col[kPerm0[j] + 16]; + blocks[j] = (a >> 4) | (b & 0xF0); + blocks[j + 16] = (a & 0x0F) | (b << 4); + } + blocks += 32; + } + } +} + +/** + * @brief Inverse of pack_codes: recover the compact quantization codes of the first `num` + * vectors from their packed blocks. Padding lanes of a partial final batch are ignored. + * + * @param padded_dim dimension of quantized data (i.e., quantization code) + * @param blocks packed quantization code, as produced by pack_codes + * @param num number of quantization code to recover + * @param quantization_code recovered codes, num * padded_dim / 8 bytes + */ +inline void unpack_codes( + size_t padded_dim, const uint8_t* blocks, size_t num, uint8_t* quantization_code +) { + const size_t cols = padded_dim / 8; + + for (size_t row = 0; row < num; row += kBatchSize) { + for (size_t i = 0; i < cols; ++i) { + for (size_t j = 0; j < 16; ++j) { + // blocks[j] holds both input vectors' high nibbles; blocks[j + 16] + // holds both low nibbles. Undo the pairing documented in pack_codes. + const size_t first = row + kPerm0[j]; + const auto a = static_cast( + ((blocks[j] & 0x0F) << 4) | (blocks[j + 16] & 0x0F) + ); + const auto b = + static_cast((blocks[j] & 0xF0) | (blocks[j + 16] >> 4)); + if (first < num) { + quantization_code[first * cols + i] = a; + } + if (first + 16 < num) { + quantization_code[(first + 16) * cols + i] = b; + } + } + blocks += 32; + } + } +} + +} // namespace rabitqlib::fastscan diff --git a/include/rabitqlib/fastscan/query.hpp b/include/rabitqlib/fastscan/query.hpp new file mode 100644 index 00000000..bfe1f1d1 --- /dev/null +++ b/include/rabitqlib/fastscan/query.hpp @@ -0,0 +1,208 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/detail/lut.hpp" + +namespace rabitqlib { +/** + * @brief Own a standard FastScan LUT; the query input is not retained. + */ +template +class BatchQuery { + private: + Lut lookup_table_; // Owns the prepared bytes; reset scratch is never retained. + T G_add_ = 0; + T G_k1xSumq_ = 0; // G_k1xSumq + MetricType metric_type_ = METRIC_L2; + + static Lut make_lut( + const T* rotated_query, size_t padded_dim, MetricType metric_type + ) { + validate_metric_type(metric_type); + return Lut(rotated_query, padded_dim); + } + + static T query_correction(const T* rotated_query, size_t padded_dim) { + constexpr float c_1 = -((1 << 1) - 1) / 2.F; + const T sumq = + std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); + return sumq * c_1; + } + + public: + explicit BatchQuery() = default; + + explicit BatchQuery( + const T* rotated_query, size_t padded_dim, MetricType metric_type = METRIC_L2 + ) + : lookup_table_(make_lut(rotated_query, padded_dim, metric_type)) + , metric_type_(metric_type) { + G_k1xSumq_ = query_correction(rotated_query, padded_dim); + } + + [[nodiscard]] T delta() const { return lookup_table_.delta(); } + + [[nodiscard]] T sum_vl_lut() const { return lookup_table_.sum_vl(); } + + [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } + + [[nodiscard]] T g_add() const { return G_add_; } + + // centroid_distance is squared L2, or 1 - dot(query, centroid) for IP. + void set_g_add(T centroid_distance) { + // IP encoder factors already include the constant 1. + G_add_ = metric_type_ == METRIC_IP ? centroid_distance - T{1} : centroid_distance; + } + + /** + * Reinitialize the query while reusing the lookup-table allocation. + * lut_scratch needs padded_dim * 4 values of T (float for BatchQuery). + * Input and scratch are borrowed only for this call; neither is retained. + * An optional correction must come from this exact rotated query. + * A successful reset clears g_add; set the centroid correction before scoring. + */ + void reset( + const T* rotated_query, + size_t padded_dim, + T* lut_scratch, + MetricType metric_type = METRIC_L2, + const T* k1xsumq = nullptr + ) { + validate_metric_type(metric_type); + // Reject invalid input before changing the correction metadata. + lookup_table_.reset(rotated_query, padded_dim, lut_scratch); + metric_type_ = metric_type; + + G_k1xSumq_ = + k1xsumq != nullptr ? *k1xsumq : query_correction(rotated_query, padded_dim); + G_add_ = 0; + } + + [[nodiscard]] size_t padded_dim() const { return lookup_table_.padded_dim(); } + + [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } +}; + +template +class SplitBatchQuery { + private: + static void validate_parameters(size_t ex_bits, MetricType metric_type) { + if (ex_bits > 8) { + throw std::invalid_argument("SplitBatchQuery requires ex_bits in [0, 8]"); + } + validate_metric_type(metric_type); + } + + static Lut make_lut( + const T* rotated_query, + size_t padded_dim, + size_t ex_bits, + MetricType metric_type, + bool use_hacc + ) { + validate_parameters(ex_bits, metric_type); + return Lut(rotated_query, padded_dim, use_hacc); + } + + // Borrowed until the next successful reset or destruction. + const T* rotated_query_ = nullptr; + Lut lookup_table_; + T G_add_ = 0; + T G_error_ = 0; + T G_k1xSumq_ = 0; + T G_kbxSumq_ = 0; + MetricType metric_type_ = METRIC_L2; + + void prepare_coefficients(const T* rotated_query, size_t padded_dim, size_t ex_bits) { + const float c_1 = -static_cast((1 << 1) - 1) / 2.F; + const float c_b = -static_cast((1 << (ex_bits + 1)) - 1) / 2.F; + const T sumq = + std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); + + G_k1xSumq_ = sumq * c_1; + G_kbxSumq_ = sumq * c_b; + } + + public: + explicit SplitBatchQuery() = default; + + explicit SplitBatchQuery( + const T* rotated_query, + size_t padded_dim, + size_t ex_bits, + MetricType metric_type = METRIC_L2, + bool use_hacc = true + ) + : rotated_query_(rotated_query) + , lookup_table_(make_lut(rotated_query, padded_dim, ex_bits, metric_type, use_hacc)) + , metric_type_(metric_type) { + prepare_coefficients(rotated_query, padded_dim, ex_bits); + } + + /** + * Reuse owned LUT storage; only first preparation or growth may allocate. + * Scratch: padded_dim * 4 values of T, plus the same count of uint16_t for + * HACC (nullptr otherwise). Scratch is call-only; buffers must not overlap + * each other, the query or LUT. The query remains borrowed for refinement. + * ex_bits: [0, 8]. Success clears g_add/g_error; call set_g_add before scoring. + * Invalid arguments or allocation failure preserve the previous state. + */ + void reset( + const T* rotated_query, + size_t padded_dim, + size_t ex_bits, + T* lut_scratch, + uint16_t* lut_u16_scratch, + MetricType metric_type = METRIC_L2, + bool use_hacc = true + ) { + validate_parameters(ex_bits, metric_type); + // Validate and grow before changing the borrowed pointer or corrections. + lookup_table_.reset( + rotated_query, padded_dim, lut_scratch, lut_u16_scratch, use_hacc + ); + rotated_query_ = rotated_query; + metric_type_ = metric_type; + prepare_coefficients(rotated_query, padded_dim, ex_bits); + G_add_ = 0; + G_error_ = 0; + } + [[nodiscard]] const T* rotated_query() const { return rotated_query_; } + + [[nodiscard]] T delta() const { return lookup_table_.delta(); } + + [[nodiscard]] T sum_vl_lut() const { return lookup_table_.sum_vl(); } + + [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } + + [[nodiscard]] T kbxsumq() const { return G_kbxSumq_; } + + [[nodiscard]] T g_add() const { return G_add_; } + + [[nodiscard]] T g_error() const { return G_error_; } + + // residual_norm is ||query - centroid|| for both metrics (not squared distance). + // For IP supply query_centroid_ip = dot(query, centroid); estimates use 1 - dot. + void set_g_add(T residual_norm, T query_centroid_ip = 0) { + if (metric_type_ == METRIC_L2) { + G_add_ = residual_norm * residual_norm; + G_error_ = residual_norm; + } else if (metric_type_ == METRIC_IP) { + G_add_ = -query_centroid_ip; + G_error_ = residual_norm; + } + } + + [[nodiscard]] bool use_hacc() const { return lookup_table_.use_hacc(); } + + [[nodiscard]] size_t padded_dim() const { return lookup_table_.padded_dim(); } + + [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } +}; + +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/buffer.hpp b/include/rabitqlib/index/detail/buffer.hpp similarity index 99% rename from include/rabitqlib/utils/buffer.hpp rename to include/rabitqlib/index/detail/buffer.hpp index 2da75708..5b3b551a 100644 --- a/include/rabitqlib/utils/buffer.hpp +++ b/include/rabitqlib/index/detail/buffer.hpp @@ -6,7 +6,7 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/memory.hpp" namespace rabitqlib::buffer { diff --git a/include/rabitqlib/index/detail/exact_nn.hpp b/include/rabitqlib/index/detail/exact_nn.hpp new file mode 100644 index 00000000..f2f15c12 --- /dev/null +++ b/include/rabitqlib/index/detail/exact_nn.hpp @@ -0,0 +1,69 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/openmp.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib { +template +inline PID exact_nn( + const T* data, + const T* query, + size_t num_points, + size_t dim, + size_t num_threads, + T (*dist_func)(const T*, const T*, size_t) +) { +#if RABITQ_USE_OPENMP + const auto thread_count = static_cast(std::max( + 1, + std::min( + {resolve_num_threads(num_threads), + std::max(num_points, 1), + static_cast(std::numeric_limits::max())} + ) + )); +#else + (void)num_threads; + constexpr int thread_count = 1; +#endif + std::vector> best_entries(thread_count); + +#if RABITQ_USE_OPENMP +#pragma omp parallel for schedule(dynamic) num_threads(thread_count) +#endif + for (std::ptrdiff_t index = 0; index < static_cast(num_points); + ++index) { + const size_t i = static_cast(index); +#if RABITQ_USE_OPENMP + const int tid = omp_get_thread_num(); +#else + constexpr int tid = 0; +#endif + AnnCandidate& cur_entry = best_entries[tid]; + const T* cur_data = data + (dim * i); + + T distance = dist_func(cur_data, query, dim); + if (distance < cur_entry.distance) { + cur_entry.id = static_cast(i); + cur_entry.distance = distance; + } + } + + PID nearest_neighbor = 0; + T min_dist = std::numeric_limits::max(); + for (auto& candi : best_entries) { + if (candi.distance < min_dist) { + nearest_neighbor = candi.id; + min_dist = candi.distance; + } + } + return nearest_neighbor; +} + +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/hashset.hpp b/include/rabitqlib/index/detail/hashset.hpp similarity index 76% rename from include/rabitqlib/utils/hashset.hpp rename to include/rabitqlib/index/detail/hashset.hpp index 2633417c..c0ce5c43 100644 --- a/include/rabitqlib/utils/hashset.hpp +++ b/include/rabitqlib/index/detail/hashset.hpp @@ -1,7 +1,7 @@ // This compatibility header preserves the original visited-set API. #pragma once -#include "rabitqlib/utils/visited_set_hash.hpp" +#include "rabitqlib/index/detail/visited_set_hash.hpp" namespace rabitqlib { using HashBasedBooleanSet = HashBasedVisitedSet; diff --git a/include/rabitqlib/index/detail/hnsw_dispatch.hpp b/include/rabitqlib/index/detail/hnsw_dispatch.hpp new file mode 100644 index 00000000..3b2f9d53 --- /dev/null +++ b/include/rabitqlib/index/detail/hnsw_dispatch.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/index/types.hpp" + +namespace rabitqlib::hnsw { +class HierarchicalNSW; +namespace detail { +// Compatibility adapter; individual backends return ordered results only. +std::priority_queue> search_knn( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); + +// Ordered results used by batch search and backend parity tests. +std::vector> search_knn_sorted( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); +std::vector> search_knn_sorted_generic( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); +std::vector> search_knn_sorted_neon( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); +std::vector> search_knn_sorted_avx2( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); +std::vector> search_knn_sorted_avx512_core( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); +std::vector> search_knn_sorted_avx512_popcnt( + HierarchicalNSW&, const float*, size_t, size_t ef = 0 +); + +} // namespace detail +} // namespace rabitqlib::hnsw diff --git a/include/rabitqlib/index/detail/parallel_exception.hpp b/include/rabitqlib/index/detail/parallel_exception.hpp new file mode 100644 index 00000000..ac713ad6 --- /dev/null +++ b/include/rabitqlib/index/detail/parallel_exception.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include +#include + +namespace rabitqlib::index_detail { + +// Shared by one OpenMP team. Capture inside each worksharing iteration and +// rethrow only after the team has joined. Successful iterations never lock. +class ParallelException { + public: + bool failed() const noexcept { return failed_.load(std::memory_order_relaxed); } + + void capture() noexcept { + if (!failed_.exchange(true, std::memory_order_relaxed)) { + error_ = std::current_exception(); + } + } + + void rethrow() const { + if (error_) { + std::rethrow_exception(error_); + } + } + + private: + std::atomic failed_{false}; + std::exception_ptr error_; +}; + +} // namespace rabitqlib::index_detail diff --git a/include/rabitqlib/index/detail/parallel_for.hpp b/include/rabitqlib/index/detail/parallel_for.hpp new file mode 100644 index 00000000..6ae8fe53 --- /dev/null +++ b/include/rabitqlib/index/detail/parallel_for.hpp @@ -0,0 +1,106 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/utils/openmp.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::index_detail { +template +inline void parallel_for(size_t start, size_t end, size_t num_threads, Function fn) { + if (start >= end) { + return; + } + num_threads = std::min(resolve_num_threads(num_threads), end - start); + + if (num_threads == 1) { + for (size_t id = start; id < end; id++) { + fn(id, 0); + } + } else { +#if defined(_OPENMP) && _OPENMP >= 201307 + if (omp_get_proc_bind() != omp_proc_bind_false) { + // Native threads inherit only the caller's place. Let OpenMP honor + // the configured placement of the whole team when binding is enabled. + std::vector errors(num_threads); +#pragma omp parallel for num_threads(num_threads) schedule(dynamic) + for (size_t id = start; id < end; ++id) { + const size_t worker = static_cast(omp_get_thread_num()); + if (!errors[worker]) { + try { + fn(id, worker); + } catch (...) { errors[worker] = std::current_exception(); } + } + } + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } + return; + } +#endif + std::vector threads; + std::atomic current(start); + + // keep track of exceptions in threads + // https://stackoverflow.com/a/32428427/1713196 + std::exception_ptr last_exception = nullptr; + std::mutex last_except_mutex; + + const auto join_threads = [&threads] { + for (auto& thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + }; + + try { + threads.reserve(num_threads); + for (size_t thread_id = 0; thread_id < num_threads; ++thread_id) { + threads.emplace_back([&, thread_id] { + while (true) { + size_t id = current.fetch_add(1); + + if (id >= end) { + break; + } + + try { + fn(id, thread_id); + } catch (...) { + std::unique_lock last_except_lock(last_except_mutex + ); + last_exception = std::current_exception(); + /* + * This will work even when current is the largest value that + * size_t can fit, because fetch_add returns the previous value + * before the increment (what will result in overflow + * and produce 0 instead of current + 1). + */ + current = end; + break; + } + } + }); + } + } catch (...) { + current = end; + join_threads(); + throw; + } + join_threads(); + if (last_exception) { + std::rethrow_exception(last_exception); + } + } +} + +} // namespace rabitqlib::index_detail diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/index/detail/visited_pool.hpp similarity index 53% rename from include/rabitqlib/utils/visited_pool.hpp rename to include/rabitqlib/index/detail/visited_pool.hpp index 27abf56e..2e3f4240 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/index/detail/visited_pool.hpp @@ -1,13 +1,15 @@ #pragma once #include #include +#include #include +#include -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" namespace rabitqlib { class VisitedListPool { - std::deque pool_; + std::deque> pool_; std::mutex poolguard_; size_t numelements_; @@ -15,36 +17,30 @@ class VisitedListPool { VisitedListPool(size_t initpoolsize, size_t max_elements) { numelements_ = max_elements; for (size_t i = 0; i < initpoolsize; i++) { - pool_.push_front(new VisitedSet(numelements_, numelements_ / 10)); + pool_.push_front(std::make_unique(numelements_, numelements_ / 10)); } } - VisitedSet* get_free_vislist() { - VisitedSet* rez; + std::unique_ptr get_free_vislist() { + std::unique_ptr rez; { std::unique_lock lock(poolguard_); if (pool_.size() > 0) { - rez = pool_.front(); + rez = std::move(pool_.front()); pool_.pop_front(); } else { - rez = new VisitedSet(numelements_, numelements_ / 10); + rez = std::make_unique(numelements_, numelements_ / 10); } } rez->clear(); return rez; } - void release_vis_list(VisitedSet* vl) { + void release_vis_list(std::unique_ptr vl) { std::unique_lock lock(poolguard_); - pool_.push_front(vl); + pool_.push_front(std::move(vl)); } - ~VisitedListPool() { - while (pool_.size() > 0) { - VisitedSet* rez = pool_.front(); - pool_.pop_front(); - ::delete rez; - } - } + ~VisitedListPool() = default; }; } // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set.hpp b/include/rabitqlib/index/detail/visited_set.hpp similarity index 85% rename from include/rabitqlib/utils/visited_set.hpp rename to include/rabitqlib/index/detail/visited_set.hpp index 6ca6c14d..b8e8ed33 100644 --- a/include/rabitqlib/utils/visited_set.hpp +++ b/include/rabitqlib/index/detail/visited_set.hpp @@ -3,9 +3,10 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/visited_set_epoch.hpp" -#include "rabitqlib/utils/visited_set_hash.hpp" +#include "rabitqlib/index/detail/visited_set_epoch.hpp" +#include "rabitqlib/index/detail/visited_set_fwd.hpp" // IWYU pragma: export +#include "rabitqlib/index/detail/visited_set_hash.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib { namespace detail { @@ -42,5 +43,4 @@ struct CheckVisitedSet { static_assert(detail::CheckVisitedSet::kValue, ""); static_assert(detail::CheckVisitedSet::kValue, ""); -using VisitedSet = HashBasedVisitedSet; } // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set_epoch.hpp b/include/rabitqlib/index/detail/visited_set_epoch.hpp similarity index 97% rename from include/rabitqlib/utils/visited_set_epoch.hpp rename to include/rabitqlib/index/detail/visited_set_epoch.hpp index c32e5734..756288ed 100644 --- a/include/rabitqlib/utils/visited_set_epoch.hpp +++ b/include/rabitqlib/index/detail/visited_set_epoch.hpp @@ -5,7 +5,7 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/memory.hpp" namespace rabitqlib { diff --git a/include/rabitqlib/index/detail/visited_set_fwd.hpp b/include/rabitqlib/index/detail/visited_set_fwd.hpp new file mode 100644 index 00000000..3c9addfd --- /dev/null +++ b/include/rabitqlib/index/detail/visited_set_fwd.hpp @@ -0,0 +1,6 @@ +#pragma once + +namespace rabitqlib { +class HashBasedVisitedSet; +using VisitedSet = HashBasedVisitedSet; +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/index/detail/visited_set_hash.hpp similarity index 98% rename from include/rabitqlib/utils/visited_set_hash.hpp rename to include/rabitqlib/index/detail/visited_set_hash.hpp index 425aa934..6c19fafc 100644 --- a/include/rabitqlib/utils/visited_set_hash.hpp +++ b/include/rabitqlib/index/detail/visited_set_hash.hpp @@ -25,7 +25,7 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/memory.hpp" namespace rabitqlib { diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 0c9f198a..bb33c14b 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -3,42 +3,38 @@ #include #include -#include -#include #include #include -#include #include -#include #include -#include -#include #include #include #include #include -#include #include #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/estimator.hpp" -#include "rabitqlib/index/ivf/initializer.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/simd/hnsw_dispatch.hpp" -#include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/memory.hpp" -#include "rabitqlib/utils/path.hpp" -#include "rabitqlib/utils/rotator.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/tools.hpp" -#include "rabitqlib/utils/visited_pool.hpp" -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/core/metric.hpp" // IWYU pragma: export +#include "rabitqlib/core/types.hpp" // IWYU pragma: export +#include "rabitqlib/index/detail/hnsw_dispatch.hpp" +#include "rabitqlib/index/types.hpp" // IWYU pragma: export + +namespace rabitqlib { +template +class Rotator; +template +class SplitSingleQuery; +class VisitedListPool; +namespace detail { +struct DistanceEstimate; +} +} // namespace rabitqlib namespace rabitqlib::hnsw { +namespace detail { +struct QuantQuery; +} template using maxheap = std::priority_queue; @@ -48,35 +44,59 @@ using minheap = std::priority_queue, std::greater>; class HierarchicalNSW { public: - explicit HierarchicalNSW(){}; + explicit HierarchicalNSW(); + // dim is the original input dimension; total_bits includes the sign bit. explicit HierarchicalNSW( - size_t, size_t, size_t, size_t, size_t, size_t = 100, MetricType = METRIC_L2 + size_t max_elements, + size_t dim, + size_t total_bits, + size_t M, + size_t ef_construction, + size_t random_seed = 100, + MetricType metric_type = METRIC_L2 ); ~HierarchicalNSW(); [[nodiscard]] size_t dimension() const { return dim_; } - [[nodiscard]] size_t num_clusters() const { return num_cluster_; } + [[nodiscard]] size_t num_clusters() const { return num_clusters_; } [[nodiscard]] size_t nbits() const { return ex_bits_ + 1; } [[nodiscard]] size_t M() const { return M_; } [[nodiscard]] size_t ef_construction() const { return ef_construction_; } [[nodiscard]] MetricType metric_type() const { return metric_type_; } + // Logical capacity, nonremoved points, and stored records (including removals). [[nodiscard]] size_t max_elements() const { return max_elements_; } - [[nodiscard]] size_t num_points() const { return cur_element_count_; } + [[nodiscard]] size_t num_points() const { + return cur_element_count_.load() - num_removed_; + } + [[nodiscard]] size_t num_stored_points() const { return cur_element_count_; } [[nodiscard]] PID cluster_id_of(PID internal_id) const { return get_clusterid_by_internalid(internal_id); } - void save(const char*) const; - void load(const char*); - - void construct(size_t, const float*, size_t, const float*, PID*, size_t, bool); + void save(const char* filename) const; + void load(const char* filename); + + // Centroids and data use the original dimension; cluster_ids has num_points entries. + // Input buffers must remain valid until this synchronous call returns or throws; + // the index retains no pointers to them afterward. + // faster selects fixed-scale extra-bit quantization instead of exact scale search, + // trading quantization accuracy for speed. One-bit storage is unaffected. + void construct( + size_t num_clusters, + const float* centroids, + size_t num_points, + const float* data, + PID* cluster_ids, + size_t num_threads = 0, + bool faster = false + ); /** * @brief Insert points into a constructed or loaded index. * * Linked with the stored quantized codes, not the exact vectors construct uses, * which are not kept. Centroids and rotation are never retrained. Point `i` gets - * the label `num_points() + i` whatever `num_threads` is. Throws rather than + * the label `num_stored_points() + i` whatever `num_threads` is. Throws rather than * growing past `max_elements()`. Not safe to call while another thread searches. * Preparation failures leave the graph and point count unchanged. A failure * during linking may leave a partially linked batch counted in the index, @@ -86,25 +106,31 @@ class HierarchicalNSW { * @param n Number of new points * @param cluster_ids Cluster of each point in `[0, num_clusters)`, or nullptr to * route each one to its nearest centroid - * @param faster Same meaning as in `construct` + * @param faster Use fixed-scale extra-bit quantization for speed at a possible + * loss of accuracy; has no effect on one-bit storage * @param num_threads Threads used to route and insert * @return The label given to each new point */ std::vector add( - const float*, size_t, const PID* = nullptr, bool = false, size_t = 1 + const float* data, + size_t n, + const PID* cluster_ids = nullptr, + bool faster = false, + size_t num_threads = 1 ); /** - * @brief Grow the number of points the index can hold. + * @brief Change capacity, allowing growth or shrinkage down to num_stored_points(). * - * Reallocates and copies the base layer, so it needs room for both copies and - * costs time proportional to the whole index. Unchanged if it throws, and - * invalidates every pointer into the index. + * An unchanged capacity is a no-op. Otherwise reallocates the base layer and + * copies stored records, requiring room for both allocations and invalidating + * pointers into the base layer. Work includes initializing the new capacity. + * The index is unchanged if this throws. * Requires initialization through the parameterized constructor or `load`. * - * @param new_max_elements The new capacity, at least `num_points()` + * @param new_max_elements The new capacity, at least `num_stored_points()` */ - void resize(size_t); + void resize(size_t new_max_elements); /** * @brief Exclude points from all later search results. @@ -112,7 +138,8 @@ class HierarchicalNSW { * A removed point keeps its codes and its place in the graph: search and `add` * still traverse it and may link new points to it, but search never returns it, * so a query may get fewer than k results. A removed point still counts in - * `num_points()`. Removal is stored in the index, so it survives save and load, + * `num_stored_points()` and occupies capacity, but no longer counts in num_points(). + * Removal is stored in the index, so it survives save and load, * and is idempotent. Removed points cannot be restored. Every label is checked * before anything changes. Not safe to call while another thread searches or adds. * @@ -120,13 +147,21 @@ class HierarchicalNSW { * @param n Number of labels * @return Number of points newly removed */ - size_t remove(const PID*, size_t); + size_t remove(const PID* labels, size_t n); + // Queries use the original dimension. Each result is (distance, external label), + // where distance estimates squared L2 or 1 - dot according to the index metric. std::vector>> search( - const float*, size_t, size_t, size_t, size_t + const float* queries, + size_t num_queries, + size_t k, + size_t ef_search, + size_t num_threads ); - const float* rawDataPtr_{nullptr}; + private: + // Borrows construct() input only while its workers run; cleared on success or failure. + const float* construction_data_{nullptr}; struct ResultRecord { float est_dist; @@ -150,7 +185,6 @@ class HierarchicalNSW { // Insert a candidate in sorted order (ascending by est_dist). void insert(const Candidate& cand) { - // Find insertion position using binary search. auto it = std::upper_bound( queue_.begin(), queue_.end(), @@ -179,34 +213,30 @@ class HierarchicalNSW { std::vector queue_; }; - private: friend struct HnswPruningTestAccess; - friend maxheap> detail::search_knn_neon( + friend std::vector> detail::search_knn_sorted_neon( HierarchicalNSW&, const float*, size_t, size_t ); - friend maxheap> detail::search_knn_generic( + friend std::vector> detail::search_knn_sorted_generic( HierarchicalNSW&, const float*, size_t, size_t ); - friend maxheap> detail::search_knn_avx2( + friend std::vector> detail::search_knn_sorted_avx2( HierarchicalNSW&, const float*, size_t, size_t ); - friend maxheap> detail::search_knn_avx512_core( + friend std::vector> detail::search_knn_sorted_avx512_core( HierarchicalNSW&, const float*, size_t, size_t ); - friend maxheap> detail::search_knn_avx512_popcnt( + friend std::vector> detail::search_knn_sorted_avx512_popcnt( HierarchicalNSW&, const float*, size_t, size_t ); static constexpr PID kMaxLabelOperationLock = 65536; // `remove` sets the high bit of a point's stored cluster ID. construct and load - // keep num_cluster_ at or below 2^31, so a real cluster ID never has it. + // keep num_clusters_ at or below 2^31, so a real cluster ID never has it. static constexpr PID kRemovedMask = 0x80000000U; - static_assert( - buffer::kSearchBufferMaxPointCount <= kRemovedMask, - "cluster IDs must leave the removal bit free" - ); size_t max_elements_{0}; - mutable std::atomic cur_element_count_{0}; // current number of elements + mutable std::atomic cur_element_count_{0}; + size_t num_removed_{0}; // Recovered from stored cluster-ID markers when loading. size_t size_data_per_element_{0}; size_t size_links_per_element_{0}; size_t M_{0}; @@ -216,8 +246,8 @@ class HierarchicalNSW { size_t ef_{0}; MetricType metric_type_{METRIC_L2}; - double mult_{0.0}, revSize_{0.0}; - int maxlevel_{0}; + double mult_{0.0}; + int max_level_{0}; // Locks operations with element by label value mutable std::vector label_op_locks_; @@ -225,33 +255,39 @@ class HierarchicalNSW { std::mutex global_; std::vector link_list_locks_; - PID enterpoint_node_{0}; + PID entry_point_id_{0}; size_t size_links_level0_{0}; - size_t offsetBinData_{0}, offsetExData_{0}, label_offset_{0}; + size_t bin_data_offset_{0}, ex_data_offset_{0}, label_offset_{0}; size_t size_bin_data_{0}, size_ex_data_{0}; size_t ex_bits_{0}; + struct AlignedDeleter { + void operator()(char* pointer) const noexcept; + }; + struct LinkDeleter { + void operator()(char* pointer) const noexcept; + }; + using AlignedStorage = std::unique_ptr; + using LinkStorage = std::unique_ptr; + // Layout: (# of edges + edges) + (cluster_id) + (External_id) + (BinData) + (ExData) - char* data_level0_memory_{nullptr}; - char** linkLists_{nullptr}; - std::vector element_levels_; // keeps level of each element + AlignedStorage data_level0_memory_; + // Reserve capacity without initializing unused slots; size tracks live nodes. + std::vector link_lists_; + std::vector element_levels_; - size_t num_cluster_{0}; + size_t num_clusters_{0}; size_t dim_{0}; size_t padded_dim_{0}; - char* centroids_memory_{nullptr}; + AlignedStorage centroids_memory_; - mutable std::mutex label_lookup_lock_; // lock for label_lookup_ + mutable std::mutex label_lookup_lock_; std::unordered_map label_lookup_; std::default_random_engine level_generator_; - std::default_random_engine update_probability_generator_; - - mutable std::atomic metric_distance_computations_{0}; - mutable std::atomic metric_hops_{0}; std::unique_ptr visited_list_pool_{nullptr}; @@ -261,42 +297,14 @@ class HierarchicalNSW { quant::RabitqConfig query_config_; - struct EstimateRecord { - float ip_x0_qr; - float est_dist; - float low_dist; - - bool operator<(const EstimateRecord& other) const { - return this->est_dist < other.est_dist; - } - }; + using EstimateRecord = rabitqlib::detail::DistanceEstimate; float (*raw_dist_func_ )(const float* __restrict__, const float* __restrict__, size_t){nullptr}; - void free_memory() { - memory::aligned_deallocate(data_level0_memory_); - data_level0_memory_ = nullptr; - for (PID i = 0; i < cur_element_count_; i++) { - if (element_levels_[i] > 0) { - free(linkLists_[i]); - } - } - free(reinterpret_cast(linkLists_)); - linkLists_ = nullptr; - cur_element_count_ = 0; - - memory::aligned_deallocate(centroids_memory_); - centroids_memory_ = nullptr; - - rotator_.reset(); - label_lookup_.clear(); - } - void set_ef(size_t ef) { ef_ = ef; } - std::mutex& get_lable_op_mutex(PID label) const { - // calculate hash + std::mutex& get_label_op_mutex(PID label) const { size_t lock_id = label & (kMaxLabelOperationLock - 1); return label_op_locks_[lock_id]; } @@ -305,7 +313,8 @@ class HierarchicalNSW { PID return_label; memcpy( &return_label, - (data_level0_memory_ + (internal_id * size_data_per_element_) + label_offset_), + (data_level0_memory_.get() + (internal_id * size_data_per_element_) + + label_offset_), sizeof(PID) ); return return_label; @@ -313,7 +322,8 @@ class HierarchicalNSW { void set_external_label(PID internal_id, PID label) const { memcpy( - (data_level0_memory_ + (internal_id * size_data_per_element_) + label_offset_), + (data_level0_memory_.get() + (internal_id * size_data_per_element_) + + label_offset_), &label, sizeof(PID) ); @@ -321,26 +331,29 @@ class HierarchicalNSW { PID* get_external_label_pt(PID internal_id) const { return reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + label_offset_ + data_level0_memory_.get() + (internal_id * size_data_per_element_) + + label_offset_ ); } char* get_bindata_by_internalid(PID internal_id) const { return reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + offsetBinData_ + data_level0_memory_.get() + (internal_id * size_data_per_element_) + + bin_data_offset_ ); } char* get_exdata_by_internalid(PID internal_id) const { return reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + offsetExData_ + data_level0_memory_.get() + (internal_id * size_data_per_element_) + + ex_data_offset_ ); } // The stored cluster ID with the removal mark, as save and load see it. PID get_raw_clusterid(PID internal_id) const { return *(reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + + data_level0_memory_.get() + (internal_id * size_data_per_element_) + size_links_level0_ )); } @@ -355,16 +368,12 @@ class HierarchicalNSW { char* get_clusterid_pt(PID internal_id) const { return reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + + data_level0_memory_.get() + (internal_id * size_data_per_element_) + size_links_level0_ ); } - int get_random_level(double reverse_size) { - std::uniform_real_distribution distribution(0.0, 1.0); - double r = -log(distribution(level_generator_)) * reverse_size; - return static_cast(r); - } + int get_random_level(double reverse_size); size_t get_max_elements() const { return max_elements_; } @@ -372,13 +381,13 @@ class HierarchicalNSW { PID* get_linklist(PID internal_id, int level) const { return reinterpret_cast( - linkLists_[internal_id] + ((level - 1) * size_links_per_element_) + link_lists_[internal_id].get() + ((level - 1) * size_links_per_element_) ); } PID* get_linklist0(PID internal_id) const { return reinterpret_cast( - data_level0_memory_ + (internal_id * size_data_per_element_) + data_level0_memory_.get() + (internal_id * size_data_per_element_) ); } @@ -400,20 +409,20 @@ class HierarchicalNSW { get_full_est_direct(std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord&) const; - maxheap> search_knn(const float*, size_t, size_t); - template - maxheap> search_knn_direct(const float*, size_t, size_t); + std::vector> search_knn_sorted_direct( + const float*, size_t, size_t + ); template - void searchBaseLayerST_AdaptiveRerankOptDirect( + void search_base_layer_direct( PID ep_id, size_t ef, - size_t TOPK, + size_t k, SplitSingleQuery& query_wrapper, std::vector& q_to_centroids, const float* query, - BoundedKNN& boundedKNN + BoundedKNN& nearest_candidates ); // Construction @@ -422,7 +431,7 @@ class HierarchicalNSW { PID label1 = get_external_label(obj1); PID label2 = get_external_label(obj2); return raw_dist_func_( - rawDataPtr_ + (label1 * dim_), rawDataPtr_ + (label2 * dim_), dim_ + construction_data_ + (label1 * dim_), construction_data_ + (label2 * dim_), dim_ ); } @@ -432,1243 +441,24 @@ class HierarchicalNSW { PID mutually_connect_new_element(PID, maxheap>&, int); - void get_neighbors_by_heuristic2(maxheap>&, size_t); + template + PID mutually_connect( + PID, maxheap>&, int, Distance, ReverseEdgeDistance + ); - // The same steps once the original data is gone, defined in hnsw_quant.hpp. + // The same steps once the original data is gone, implemented in hnsw.cpp. void reconstruct_rotated(PID, float*) const; - float get_quant_dist(PID target, PID query) const; - void add_point_quant(PID, int, std::unique_ptr&); - maxheap> search_base_layer_quant(PID, PID, int); - PID mutually_connect_quant(PID, maxheap>&, int); - void get_neighbors_by_heuristic2_quant(maxheap>&, size_t); + float get_quant_dist(PID target, PID query, detail::QuantQuery&) const; + void add_point_quant(PID, int, LinkStorage&, detail::QuantQuery&); + maxheap> + search_base_layer_quant(PID, PID, int, detail::QuantQuery&); + PID + mutually_connect_quant(PID, maxheap>&, int, detail::QuantQuery&); }; -inline HierarchicalNSW::HierarchicalNSW( - size_t max_elements, - size_t dim, - size_t total_bits, - size_t M, - size_t ef_construction, - size_t random_seed, - MetricType metric_type -) - : metric_type_(metric_type) - , label_op_locks_(kMaxLabelOperationLock) - , link_list_locks_(max_elements) - , element_levels_(max_elements) - , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) { - validate_metric_type(metric_type); - if (M < 2) { - throw std::invalid_argument("HNSW M must be at least 2"); - } - max_elements_ = max_elements; - dim_ = dim; - rotator_.reset(choose_rotator( - dim, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 32) - )); - padded_dim_ = rotator_->size(); - /* check size */ - assert(padded_dim_ % 32 == 0); - assert(padded_dim_ >= dim_); - ex_bits_ = total_bits - 1; - - if (total_bits < 1 || total_bits > 9) { - throw std::invalid_argument("HNSW quantization bits must be in [1, 9]"); - }; - - assert(padded_dim_ % 32 == 0); - - ip_func_ = select_excode_ipfunc(ex_bits_); - - if (M <= 10000) { - M_ = M; - } else { - M_ = 10000; - } - - maxM_ = M_; - maxM0_ = M_ * 2; - ef_construction_ = std::max(ef_construction, M_); - ef_ = 10; - - size_bin_data_ = BinDataMap::data_bytes(padded_dim_); - size_ex_data_ = ExDataMap::data_bytes(padded_dim_, ex_bits_); - size_links_level0_ = maxM0_ * sizeof(PID) + sizeof(PID); - label_offset_ = - size_links_level0_ + sizeof(PID); // (# of edges + edges) + (cluster_id) - offsetBinData_ = label_offset_ + - sizeof(PID); // (# of edges + edges) + (cluster_id) + (external label) - offsetExData_ = offsetBinData_ + size_bin_data_; // (# of edges + edges) + (cluster_id) - // + (external label) + (BinData) - size_data_per_element_ = offsetExData_ + size_ex_data_; - data_level0_memory_ = - memory::huge_page_allocate(max_elements_ * size_data_per_element_); - if (data_level0_memory_ == nullptr) { - throw std::runtime_error("Not enough memory"); - } - - level_generator_.seed(random_seed); - update_probability_generator_.seed(random_seed + 1); - - cur_element_count_ = 0; - - visited_list_pool_ = std::make_unique(1, max_elements_); - - // initializations for special treatment of the first node - enterpoint_node_ = -1; - maxlevel_ = -1; - - linkLists_ = reinterpret_cast(malloc(sizeof(void*) * max_elements_)); - if (linkLists_ == nullptr) { - throw std::runtime_error("Not enough memory: HNSW failed to allocate linklists"); - } - size_links_per_element_ = maxM_ * sizeof(PID) + sizeof(PID); - mult_ = 1 / log(1.0 * static_cast(M_)); - revSize_ = 1.0 / mult_; - - this->query_config_ = - quant::faster_config(padded_dim_, SplitSingleQuery::kNumBits); -} - -inline HierarchicalNSW::~HierarchicalNSW() { free_memory(); } - -inline void HierarchicalNSW::save(const char* filename) const { - std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); - if (!output.is_open()) { - throw std::runtime_error("HNSW: cannot open index file for writing"); - } - - output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); - output.write(reinterpret_cast(&cur_element_count_), sizeof(size_t)); - - output.write(reinterpret_cast(&dim_), sizeof(size_t)); - output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); - output.write(reinterpret_cast(&num_cluster_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); - output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - - output.write(reinterpret_cast(&size_bin_data_), sizeof(size_t)); - output.write(reinterpret_cast(&size_ex_data_), sizeof(size_t)); - output.write(reinterpret_cast(&size_links_level0_), sizeof(size_t)); - output.write(reinterpret_cast(&offsetBinData_), sizeof(size_t)); - output.write(reinterpret_cast(&offsetExData_), sizeof(size_t)); - output.write(reinterpret_cast(&label_offset_), sizeof(PID)); - output.write(reinterpret_cast(&size_data_per_element_), sizeof(size_t)); - output.write(reinterpret_cast(&size_links_per_element_), sizeof(size_t)); - - output.write(reinterpret_cast(&maxlevel_), sizeof(int)); - output.write(reinterpret_cast(&enterpoint_node_), sizeof(PID)); - - output.write(reinterpret_cast(&M_), sizeof(size_t)); - output.write(reinterpret_cast(&maxM_), sizeof(size_t)); - output.write(reinterpret_cast(&maxM0_), sizeof(size_t)); - output.write(reinterpret_cast(&mult_), sizeof(double)); - output.write(reinterpret_cast(&ef_construction_), sizeof(size_t)); - - output.write( - reinterpret_cast(centroids_memory_), - static_cast(num_cluster_ * padded_dim_ * sizeof(float)) - ); - - output.write( - reinterpret_cast(data_level0_memory_), - static_cast(cur_element_count_ * size_data_per_element_) - ); - - for (size_t i = 0; i < cur_element_count_; i++) { - unsigned int link_list_size = - element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; - output.write(reinterpret_cast(&link_list_size), sizeof(unsigned int)); - if (link_list_size != 0) { - output.write(reinterpret_cast(linkLists_[i]), link_list_size); - } - } - - rotator_->save(output); - output.close(); - if (!output) { - throw std::runtime_error("HNSW: failed to write index file"); - } -} - -inline void HierarchicalNSW::load(const char* filename) { - if (filename == nullptr || filename[0] == '\0') { - throw std::invalid_argument("HNSW: index filename must not be empty"); - } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); - if (!input.is_open()) { - throw std::runtime_error("HNSW: cannot open index file"); - } - - auto invalid_file = []() -> void { - throw std::runtime_error("HNSW: invalid or truncated index file"); - }; - auto read_header = [&](auto& value) { - input.read(reinterpret_cast(&value), sizeof(value)); - if (!input) { - invalid_file(); - } - }; - auto checked_product = [&](size_t left, size_t right) { - if (right != 0 && left > std::numeric_limits::max() / right) { - invalid_file(); - } - return left * right; - }; - - // Keep the current index intact until the complete file has been validated. - HierarchicalNSW loaded; - size_t element_count = 0; - PID label_offset = 0; - read_header(loaded.max_elements_); - read_header(element_count); - read_header(loaded.dim_); - read_header(loaded.padded_dim_); - read_header(loaded.num_cluster_); - read_header(loaded.ex_bits_); - read_header(loaded.metric_type_); - read_header(loaded.size_bin_data_); - read_header(loaded.size_ex_data_); - read_header(loaded.size_links_level0_); - read_header(loaded.offsetBinData_); - read_header(loaded.offsetExData_); - read_header(label_offset); - loaded.label_offset_ = label_offset; - read_header(loaded.size_data_per_element_); - read_header(loaded.size_links_per_element_); - read_header(loaded.maxlevel_); - read_header(loaded.enterpoint_node_); - read_header(loaded.M_); - read_header(loaded.maxM_); - read_header(loaded.maxM0_); - read_header(loaded.mult_); - read_header(loaded.ef_construction_); - - if (loaded.max_elements_ == 0 || - loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || - element_count > loaded.max_elements_ || - loaded.num_cluster_ > buffer::kSearchBufferMaxPointCount || - (element_count != 0 && loaded.num_cluster_ == 0) || loaded.dim_ < 64 || - loaded.dim_ > rotator_impl::FhtKacRotator::kMaxDim || - loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || - loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || - (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || - loaded.M_ < 2 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || - loaded.maxM0_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || - (element_count == 0 && - (loaded.maxlevel_ != -1 || loaded.enterpoint_node_ != kPidMax)) || - (element_count != 0 && - (loaded.maxlevel_ < 0 || loaded.enterpoint_node_ >= element_count))) { - invalid_file(); - } - - const size_t expected_level0 = (loaded.maxM0_ + 1) * sizeof(PID); - const size_t expected_bin = BinDataMap::data_bytes(loaded.padded_dim_); - const size_t expected_ex = - ExDataMap::data_bytes(loaded.padded_dim_, loaded.ex_bits_); - const size_t expected_label = expected_level0 + sizeof(PID); - const size_t expected_bin_offset = expected_label + sizeof(PID); - const size_t expected_ex_offset = expected_bin_offset + expected_bin; - const size_t expected_stride = expected_ex_offset + expected_ex; - if (loaded.size_bin_data_ != expected_bin || loaded.size_ex_data_ != expected_ex || - loaded.size_links_level0_ != expected_level0 || - loaded.label_offset_ != expected_label || - loaded.offsetBinData_ != expected_bin_offset || - loaded.offsetExData_ != expected_ex_offset || - loaded.size_data_per_element_ != expected_stride || - loaded.size_links_per_element_ != (loaded.maxM_ + 1) * sizeof(PID) || - !std::isfinite(loaded.mult_) || loaded.mult_ <= 0.0) { - invalid_file(); - } - - const size_t centroids_bytes = checked_product( - checked_product(loaded.num_cluster_, loaded.padded_dim_), sizeof(float) - ); - const size_t data_capacity_bytes = - checked_product(loaded.max_elements_, loaded.size_data_per_element_); - const size_t data_bytes = checked_product(element_count, loaded.size_data_per_element_); - const size_t link_pointer_bytes = checked_product(loaded.max_elements_, sizeof(char*)); - const size_t link_headers_bytes = checked_product(element_count, sizeof(unsigned int)); - loaded.rotator_.reset( - choose_rotator(loaded.dim_, RotatorType::FhtKacRotator, loaded.padded_dim_) - ); - const size_t rotator_bytes = loaded.rotator_->dump_bytes(); - - const auto payload_start = input.tellg(); - input.seekg(0, std::ios::end); - const auto payload_end = input.tellg(); - if (!input || payload_start < 0 || payload_end < payload_start) { - invalid_file(); - } - size_t remaining = static_cast(payload_end - payload_start); - input.seekg(payload_start); - if (centroids_bytes > remaining || data_bytes > remaining - centroids_bytes || - link_headers_bytes > remaining - centroids_bytes - data_bytes || - rotator_bytes > remaining - centroids_bytes - data_bytes - link_headers_bytes) { - invalid_file(); - } - auto read_payload = [&](char* destination, size_t bytes) { - if (bytes > remaining || - bytes > static_cast(std::numeric_limits::max())) { - invalid_file(); - } - if (bytes != 0) { - input.read(destination, static_cast(bytes)); - if (!input) { - invalid_file(); - } - } - remaining -= bytes; - }; - - if (centroids_bytes != 0) { - loaded.centroids_memory_ = memory::huge_page_allocate(centroids_bytes); - if (loaded.centroids_memory_ == nullptr) { - throw std::bad_alloc(); - } - read_payload(loaded.centroids_memory_, centroids_bytes); - } - loaded.data_level0_memory_ = memory::huge_page_allocate(data_capacity_bytes); - if (loaded.data_level0_memory_ == nullptr) { - throw std::bad_alloc(); - } - read_payload(loaded.data_level0_memory_, data_bytes); - - loaded.element_levels_ = std::vector(loaded.max_elements_); - loaded.link_list_locks_ = std::vector(loaded.max_elements_); - loaded.label_op_locks_ = std::vector(kMaxLabelOperationLock); - loaded.linkLists_ = reinterpret_cast(calloc(1, link_pointer_bytes)); - if (loaded.linkLists_ == nullptr) { - throw std::bad_alloc(); - } - for (size_t i = 0; i < element_count; ++i) { - const char* row = loaded.data_level0_memory_ + i * loaded.size_data_per_element_; - PID cluster_id = 0; - PID label = 0; - unsigned short level0_degree = 0; - std::memcpy(&level0_degree, row, sizeof(level0_degree)); - std::memcpy(&cluster_id, row + loaded.size_links_level0_, sizeof(cluster_id)); - std::memcpy(&label, row + loaded.label_offset_, sizeof(label)); - if ((cluster_id & ~kRemovedMask) >= loaded.num_cluster_ || - level0_degree > loaded.maxM0_) { - invalid_file(); - } - for (size_t edge = 0; edge < level0_degree; ++edge) { - PID neighbor = 0; - std::memcpy(&neighbor, row + sizeof(PID) + edge * sizeof(PID), sizeof(PID)); - if (neighbor >= element_count) { - invalid_file(); - } - } - if (!loaded.label_lookup_.emplace(label, static_cast(i)).second) { - invalid_file(); - } - - unsigned int link_list_size = 0; - read_payload(reinterpret_cast(&link_list_size), sizeof(link_list_size)); - if (link_list_size != 0) { - if (link_list_size % loaded.size_links_per_element_ != 0 || - link_list_size / loaded.size_links_per_element_ > - static_cast(loaded.maxlevel_) || - link_list_size > remaining) { - invalid_file(); - } - loaded.linkLists_[i] = reinterpret_cast(malloc(link_list_size)); - if (loaded.linkLists_[i] == nullptr) { - throw std::bad_alloc(); - } - loaded.element_levels_[i] = - static_cast(link_list_size / loaded.size_links_per_element_); - } - loaded.cur_element_count_ = i + 1; - if (link_list_size != 0) { - read_payload(loaded.linkLists_[i], link_list_size); - for (int level = 0; level < loaded.element_levels_[i]; ++level) { - const char* link = - loaded.linkLists_[i] + level * loaded.size_links_per_element_; - unsigned short degree = 0; - std::memcpy(°ree, link, sizeof(degree)); - if (degree > loaded.maxM_) { - invalid_file(); - } - for (size_t edge = 0; edge < degree; ++edge) { - PID neighbor = 0; - std::memcpy( - &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) - ); - if (neighbor >= element_count) { - invalid_file(); - } - } - } - } - } - for (size_t i = 0; i < element_count; ++i) { - for (int level = 1; level <= loaded.element_levels_[i]; ++level) { - const char* link = loaded.linkLists_[i] + (static_cast(level) - 1) * - loaded.size_links_per_element_; - unsigned short degree = 0; - std::memcpy(°ree, link, sizeof(degree)); - for (size_t edge = 0; edge < degree; ++edge) { - PID neighbor = 0; - std::memcpy( - &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) - ); - if (loaded.element_levels_[neighbor] < level) { - invalid_file(); - } - } - } - } - if (element_count != 0 && - loaded.element_levels_[loaded.enterpoint_node_] != loaded.maxlevel_) { - invalid_file(); - } - - if (rotator_bytes > remaining) { - invalid_file(); - } - loaded.rotator_->load(input); - if (!input) { - invalid_file(); - } - remaining -= rotator_bytes; - if (remaining != 0) { - invalid_file(); - } - - loaded.visited_list_pool_ = std::make_unique(1, loaded.max_elements_); - loaded.raw_dist_func_ = - (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; - loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); - loaded.revSize_ = 1.0 / loaded.mult_; - loaded.ef_ = 10; - loaded.query_config_ = - quant::faster_config(loaded.padded_dim_, SplitSingleQuery::kNumBits); - - // All swaps below are non-allocating. The temporary releases the old state. - using std::swap; - swap(rawDataPtr_, loaded.rawDataPtr_); - swap(max_elements_, loaded.max_elements_); - const size_t old_count = cur_element_count_.load(); - cur_element_count_ = loaded.cur_element_count_.load(); - loaded.cur_element_count_ = old_count; - swap(size_data_per_element_, loaded.size_data_per_element_); - swap(size_links_per_element_, loaded.size_links_per_element_); - swap(M_, loaded.M_); - swap(maxM_, loaded.maxM_); - swap(maxM0_, loaded.maxM0_); - swap(ef_construction_, loaded.ef_construction_); - swap(ef_, loaded.ef_); - swap(metric_type_, loaded.metric_type_); - swap(mult_, loaded.mult_); - swap(revSize_, loaded.revSize_); - swap(maxlevel_, loaded.maxlevel_); - label_op_locks_.swap(loaded.label_op_locks_); - link_list_locks_.swap(loaded.link_list_locks_); - swap(enterpoint_node_, loaded.enterpoint_node_); - swap(size_links_level0_, loaded.size_links_level0_); - swap(offsetBinData_, loaded.offsetBinData_); - swap(offsetExData_, loaded.offsetExData_); - swap(label_offset_, loaded.label_offset_); - swap(size_bin_data_, loaded.size_bin_data_); - swap(size_ex_data_, loaded.size_ex_data_); - swap(ex_bits_, loaded.ex_bits_); - swap(data_level0_memory_, loaded.data_level0_memory_); - swap(linkLists_, loaded.linkLists_); - element_levels_.swap(loaded.element_levels_); - swap(num_cluster_, loaded.num_cluster_); - swap(dim_, loaded.dim_); - swap(padded_dim_, loaded.padded_dim_); - swap(centroids_memory_, loaded.centroids_memory_); - label_lookup_.swap(loaded.label_lookup_); - swap(visited_list_pool_, loaded.visited_list_pool_); - swap(ip_func_, loaded.ip_func_); - swap(rotator_, loaded.rotator_); - swap(query_config_, loaded.query_config_); - swap(raw_dist_func_, loaded.raw_dist_func_); -} - -inline void HierarchicalNSW::construct( - size_t cluster_num, - const float* centroids, - size_t data_num, - const float* data, - PID* cluster_ids, - size_t num_threads = 0, - bool faster = false -) { - if (num_cluster_ != 0 || cur_element_count_ != 0) { - throw std::logic_error("HNSW index is already constructed or loaded"); - } - if (cluster_num == 0 || cluster_num > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("HNSW cluster count is out of range"); - } - if (data_num == 0 || data_num > max_elements_ || - data_num > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("HNSW point count is out of range"); - } - if (centroids == nullptr || data == nullptr || cluster_ids == nullptr) { - throw std::invalid_argument("HNSW construction inputs must not be null"); - } - if (cluster_num > std::numeric_limits::max() / sizeof(float) / padded_dim_) { - throw std::invalid_argument("HNSW centroids exceed addressable storage"); - } - for (size_t i = 0; i < data_num; ++i) { - if (cluster_ids[i] >= cluster_num) { - throw std::invalid_argument("HNSW cluster ID is out of range"); - } - } - - try { - num_cluster_ = cluster_num; - const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); - centroids_memory_ = memory::huge_page_allocate(centroids_bytes); - if (centroids_memory_ == nullptr) { - throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids" - ); - } - - for (size_t i = 0; i < cluster_num; ++i) { - this->rotator_->rotate( - centroids + (i * dim_), - reinterpret_cast(centroids_memory_) + (i * padded_dim_) - ); - } - - quant::RabitqConfig config; - if (faster) { - config = quant::faster_config(padded_dim_, ex_bits_ + 1); - } - - rawDataPtr_ = data; - rabitqlib::ivf::parallel_for( - 0, - data_num, - num_threads, - [&](size_t idx, size_t /*threadId*/) { - add_point(idx, cluster_ids[idx], config); - } - ); - } catch (...) { - // parallel_for joins every worker before propagating an exception. - // Discard the partial graph, retaining capacity and rotation for a retry. - for (size_t i = 0; i < cur_element_count_; ++i) { - std::free(linkLists_[i]); - linkLists_[i] = nullptr; - element_levels_[i] = 0; - } - cur_element_count_ = 0; - label_lookup_.clear(); - enterpoint_node_ = kPidMax; - maxlevel_ = -1; - memory::aligned_deallocate(centroids_memory_); - centroids_memory_ = nullptr; - num_cluster_ = 0; - rawDataPtr_ = nullptr; - throw; - } - rawDataPtr_ = nullptr; -} - -inline void HierarchicalNSW::add_point( - PID label, PID cluster_id, const quant::RabitqConfig& config -) { - std::unique_lock lock_label(get_lable_op_mutex(label)); - - int level = -1; - int curlevel = 0; - PID cur_c = 0; - { - std::unique_lock lock_table(label_lookup_lock_); - if (label_lookup_.find(label) != label_lookup_.end()) { - throw std::runtime_error( - "Currently not support replacement of existing elements, only support " - "inserting elements with distinct labels" - ); - } - - if (cur_element_count_ >= max_elements_) { - throw std::runtime_error("The number of elements exceeds the specified limit"); - } - - cur_c = cur_element_count_; - linkLists_[cur_c] = nullptr; - cur_element_count_++; - label_lookup_[label] = cur_c; - curlevel = get_random_level(mult_); - } - - std::unique_lock lock_el(link_list_locks_[cur_c]); - if (level > 0) { - curlevel = level; - } - - element_levels_[cur_c] = curlevel; - std::unique_lock templock(global_); - int maxlevelcopy = maxlevel_; - PID curr_obj = enterpoint_node_; - if (curlevel <= maxlevelcopy) { - templock.unlock(); - } - - // initialize the current memory. - memset( - data_level0_memory_ + (cur_c * size_data_per_element_), 0, size_data_per_element_ - ); - - // Initialisation of label and cluster id - memcpy(get_external_label_pt(cur_c), &label, sizeof(PID)); - memcpy(get_clusterid_pt(cur_c), &cluster_id, sizeof(PID)); - - // Quantize raw data and initialize quantized data - std::vector rotated_data(padded_dim_); - rotator_->rotate(rawDataPtr_ + (label * dim_), rotated_data.data()); - quant::quantize_split_single( - rotated_data.data(), - reinterpret_cast(centroids_memory_) + (cluster_id * padded_dim_), - padded_dim_, - ex_bits_, - get_bindata_by_internalid(cur_c), - get_exdata_by_internalid(cur_c), - metric_type_, - config - ); - - // If the current vertex is at level >0, it needs some space to store the extra edges. - if (curlevel > 0) { - linkLists_[cur_c] = - static_cast(malloc((size_links_per_element_ * curlevel) + 1)); - if (linkLists_[cur_c] == nullptr) { - throw std::runtime_error( - "Not enough memory: add_point failed to allocate linklist" - ); - } - memset(linkLists_[cur_c], 0, (size_links_per_element_ * curlevel) + 1); - } - - if (static_cast(curr_obj) != -1) { - if (curlevel < maxlevelcopy) { - float curdist = get_data_dist(curr_obj, cur_c); - for (int level = maxlevelcopy; level > curlevel; level--) { - bool changed = true; - while (changed) { - changed = false; - unsigned int* data; - std::unique_lock lock(link_list_locks_[curr_obj]); - data = get_linklist(curr_obj, level); - int size = get_list_count(data); - - auto* datal = static_cast(data + 1); - for (int i = 0; i < size; i++) { - PID cand = datal[i]; - if (cand > max_elements_) { - throw std::runtime_error("cand error"); - } - float d = get_data_dist(cand, cur_c); - if (d < curdist) { - curdist = d; - curr_obj = cand; - changed = true; - } - } - } - } - } - - for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { - maxheap> top_candidates = - search_base_layer(curr_obj, cur_c, level); - curr_obj = mutually_connect_new_element(cur_c, top_candidates, level); - } - } else { - // Do nothing for the first element - enterpoint_node_ = 0; - maxlevel_ = curlevel; - } - - // Releasing lock for the maximum level - if (curlevel > maxlevelcopy) { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; - } -} - -inline maxheap> HierarchicalNSW::search_base_layer( - PID ep_id, PID cur_c, int layer -) { - std::unique_ptr vl(visited_list_pool_->get_free_vislist()); - - maxheap> top_candidates; - minheap> candidate_set; - - float lower_bound = get_data_dist(ep_id, cur_c); - top_candidates.emplace(lower_bound, ep_id); - candidate_set.emplace(lower_bound, ep_id); - vl->set(ep_id); - - while (!candidate_set.empty()) { - std::pair curr_el_pair = candidate_set.top(); - if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { - break; - } - candidate_set.pop(); - - PID cur_node_num = curr_el_pair.second; - - std::unique_lock lock(link_list_locks_[cur_node_num]); - - int* data; - if (layer == 0) { - data = reinterpret_cast(get_linklist0(cur_node_num)); - } else { - data = reinterpret_cast(get_linklist(cur_node_num, layer)); - } - size_t size = get_list_count(reinterpret_cast(data)); - auto* datal = reinterpret_cast(data + 1); - - rabitqlib::memory::mem_prefetch_l1( - reinterpret_cast( - rawDataPtr_ + (get_external_label(*datal) * dim_) - ), - padded_dim_ / 16 - ); - - rabitqlib::memory::mem_prefetch_l1( - reinterpret_cast( - rawDataPtr_ + (get_external_label(*(datal + 1)) * dim_) - ), - padded_dim_ / 16 - ); - - for (size_t j = 0; j < size; j++) { - PID candidate_id = *(datal + j); - if (vl->get(candidate_id)) { - continue; - } - vl->set(candidate_id); - - if (j < size - 1) { - rabitqlib::memory::mem_prefetch_l1( - reinterpret_cast( - rawDataPtr_ + (get_external_label(*(datal + j + 1)) * dim_) - ), - padded_dim_ / 16 - ); - } - - float dist1 = get_data_dist(candidate_id, cur_c); - if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { - candidate_set.emplace(dist1, candidate_id); - top_candidates.emplace(dist1, candidate_id); - if (top_candidates.size() > ef_construction_) { - top_candidates.pop(); - } - if (!top_candidates.empty()) { - lower_bound = top_candidates.top().first; - } - } - } - } - visited_list_pool_->release_vis_list(vl.get()); - vl.release(); - return top_candidates; -} - -inline PID HierarchicalNSW::mutually_connect_new_element( - PID cur_c, maxheap>& top_candidates, int level -) { - size_t max_m = level > 0 ? maxM_ : maxM0_; - get_neighbors_by_heuristic2(top_candidates, M_); - if (top_candidates.size() > M_) { - throw std::runtime_error( - "Should be not be more than M_ candidates returned by the heuristic" - ); - } - - std::vector selected_neighbors; - selected_neighbors.reserve(M_); - while (top_candidates.size() > 0) { - selected_neighbors.push_back(top_candidates.top().second); - top_candidates.pop(); - } - - PID next_closest_entry_point = selected_neighbors.back(); - - { - PID* ll_cur; - if (level == 0) { - ll_cur = get_linklist0(cur_c); - } else { - ll_cur = get_linklist(cur_c, level); - } - - if (*ll_cur > 0) { - throw std::runtime_error( - "The newly inserted element should have blank link list" - ); - } - - set_list_count(ll_cur, selected_neighbors.size()); - auto* data = static_cast(ll_cur + 1); - for (size_t idx = 0; idx < selected_neighbors.size(); idx++) { - if (data[idx] != 0) { - throw std::runtime_error("Possible memory corruption"); - } - if (level > element_levels_[selected_neighbors[idx]]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); - } - - data[idx] = selected_neighbors[idx]; - } - } - - for (auto selected_neighbor : selected_neighbors) { - std::unique_lock lock(link_list_locks_[selected_neighbor]); - - PID* ll_other; - if (level == 0) { - ll_other = get_linklist0(selected_neighbor); - } else { - ll_other = get_linklist(selected_neighbor, level); - } - - size_t sz_link_list_other = get_list_count(ll_other); - - if (sz_link_list_other > max_m) { - throw std::runtime_error("Bad value of sz_link_list_other"); - } - if (selected_neighbor == cur_c) { - throw std::runtime_error("Trying to connect an element to itself"); - } - if (level > element_levels_[selected_neighbor]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); - } - - auto* data = static_cast(ll_other + 1); - - bool is_cur_c_present = false; - for (size_t j = 0; j < sz_link_list_other; j++) { - if (data[j] == cur_c) { - is_cur_c_present = true; - break; - } - } - - if (!is_cur_c_present) { - if (sz_link_list_other < max_m) { - data[sz_link_list_other] = cur_c; - set_list_count(ll_other, sz_link_list_other + 1); - } else { - float d_max = get_data_dist(selected_neighbor, cur_c); - maxheap> candidates; - candidates.emplace(d_max, cur_c); - for (size_t j = 0; j < sz_link_list_other; j++) { - candidates.emplace(get_data_dist(data[j], selected_neighbor), data[j]); - } - - get_neighbors_by_heuristic2(candidates, max_m); - - int indx = 0; - while (candidates.size() > 0) { - data[indx] = candidates.top().second; - candidates.pop(); - indx++; - } - - set_list_count(ll_other, indx); - } - } - } - - return next_closest_entry_point; -} - -inline void HierarchicalNSW::get_neighbors_by_heuristic2( - maxheap>& top_candidates, size_t M -) { - if (top_candidates.size() < M) { - return; - } - - minheap> queue_closest; - std::vector> return_list; - while (top_candidates.size() > 0) { - queue_closest.emplace(top_candidates.top()); - top_candidates.pop(); - } - - while (queue_closest.size() > 0) { - if (return_list.size() >= M) { - break; - } - std::pair current_pair = queue_closest.top(); - float dist_to_query = current_pair.first; - queue_closest.pop(); - bool good = true; - - for (std::pair second_pair : return_list) { - float curdist = get_data_dist(second_pair.second, current_pair.second); - if (curdist < dist_to_query) { - good = false; - break; - } - } - if (good) { - return_list.push_back(current_pair); - } - } - - for (std::pair current_pair : return_list) { - top_candidates.emplace(current_pair); - } -} - -template -inline void HierarchicalNSW::get_bin_est_direct( - std::vector& q_to_centroids, - SplitSingleQuery& query_wrapper, - PID currObj, - HierarchicalNSW::EstimateRecord& res -) { - if (metric_type_ == METRIC_IP) { - float norm = q_to_centroids[get_clusterid_by_internalid(currObj)]; - float error = q_to_centroids[get_clusterid_by_internalid(currObj) + num_cluster_]; - split_single_estdist_direct( - get_bindata_by_internalid(currObj), - query_wrapper, - padded_dim_, - res.ip_x0_qr, - res.est_dist, - res.low_dist, - -norm, - error - ); - } else { - // L2 distance - float norm = q_to_centroids[get_clusterid_by_internalid(currObj)]; - split_single_estdist_direct( - get_bindata_by_internalid(currObj), - query_wrapper, - padded_dim_, - res.ip_x0_qr, - res.est_dist, - res.low_dist, - norm * norm, - norm - ); - } -} - -template -inline void HierarchicalNSW::get_full_est_direct( - std::vector& q_to_centroids, - SplitSingleQuery& query_wrapper, - PID currObj, - HierarchicalNSW::EstimateRecord& res -) const { - if (metric_type_ == METRIC_IP) { - float norm = q_to_centroids[get_clusterid_by_internalid(currObj)]; - float error = q_to_centroids[get_clusterid_by_internalid(currObj) + num_cluster_]; - split_single_fulldist_direct( - get_bindata_by_internalid(currObj), - get_exdata_by_internalid(currObj), - ip_func_, - query_wrapper, - padded_dim_, - ex_bits_, - res.est_dist, - res.low_dist, - res.ip_x0_qr, - -norm, - error - ); - } else { - // L2 distance - float norm = q_to_centroids[get_clusterid_by_internalid(currObj)]; - split_single_fulldist_direct( - get_bindata_by_internalid(currObj), - get_exdata_by_internalid(currObj), - ip_func_, - query_wrapper, - padded_dim_, - ex_bits_, - res.est_dist, - res.low_dist, - res.ip_x0_qr, - norm * norm, - norm - ); - } -} - -inline std::vector>> HierarchicalNSW::search( - const float* queries, size_t query_num, size_t TOPK, size_t efSearch, size_t thread_num -) { - std::vector>> results(query_num); - rabitqlib::ivf::parallel_for( - 0, - query_num, - thread_num, - [&](size_t idx, size_t /*threadId*/) { - std::vector rotated_query(padded_dim_); - this->rotator_->rotate(queries + (idx * dim_), rotated_query.data()); - maxheap> knn = - search_knn(rotated_query.data(), TOPK, std::max(efSearch, TOPK)); - while (knn.size()) { - results[idx].emplace_back(knn.top()); - knn.pop(); - } - std::reverse(results[idx].begin(), results[idx].end()); - } - ); - return results; -} - -inline size_t HierarchicalNSW::remove(const PID* labels, size_t n) { - if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || - rotator_ == nullptr || num_cluster_ == 0) { - throw std::logic_error("HNSW index must be constructed or loaded before remove"); - } - if (n == 0) { - return 0; - } - if (labels == nullptr) { - throw std::invalid_argument("HNSW remove labels must not be null"); - } - - // Resolve every label before marking any point, so an unknown label removes - // nothing. - std::vector internal_ids(n); - { - std::unique_lock lock_table(label_lookup_lock_); - for (size_t i = 0; i < n; ++i) { - const auto found = label_lookup_.find(labels[i]); - if (found == label_lookup_.end()) { - throw std::invalid_argument("HNSW remove label is not in the index"); - } - internal_ids[i] = found->second; - } - } - - size_t removed = 0; - for (const PID id : internal_ids) { - if (!is_removed(id)) { - const PID marked = get_raw_clusterid(id) | kRemovedMask; - std::memcpy(get_clusterid_pt(id), &marked, sizeof(PID)); - ++removed; - } - } - return removed; -} - -inline maxheap> HierarchicalNSW::search_knn( - const float* rotated_query, size_t TOPK, size_t ef -) { - return detail::search_knn(*this, rotated_query, TOPK, ef); -} - -template -inline maxheap> HierarchicalNSW::search_knn_direct( - const float* rotated_query, size_t TOPK, size_t ef -) { - maxheap> result; - if (cur_element_count_ == 0) { - return result; - } - - SplitSingleQuery query_wrapper( - rotated_query, padded_dim_, ex_bits_, query_config_, metric_type_ - ); - - // Preprocess - get the distance from query to all centroids - std::vector q_to_centroids(num_cluster_); - - if (metric_type_ == METRIC_L2) { - for (size_t i = 0; i < num_cluster_; i++) { - q_to_centroids[i] = std::sqrt(raw_dist_func_( - rotated_query, - reinterpret_cast(centroids_memory_) + (i * padded_dim_), - padded_dim_ - )); - } - } else if (metric_type_ == METRIC_IP) { - q_to_centroids.resize(2 * num_cluster_); - // first half as g_add, second half as g_error - for (size_t i = 0; i < num_cluster_; i++) { - q_to_centroids[i] = dot_product( - rotated_query, - reinterpret_cast(centroids_memory_) + (i * padded_dim_), - padded_dim_ - ); - q_to_centroids[i + num_cluster_] = std::sqrt(euclidean_sqr( - rotated_query, - reinterpret_cast(centroids_memory_) + (i * padded_dim_), - padded_dim_ - )); - } - } - - PID curr_obj = enterpoint_node_; - EstimateRecord curest; - - get_bin_est_direct(q_to_centroids, query_wrapper, curr_obj, curest); - - for (int level = maxlevel_; level > 0; level--) { - bool changed = true; - while (changed) { - changed = false; - unsigned int* data; - - data = static_cast(get_linklist(curr_obj, level)); - int size = get_list_count(data); - - PID* datal = static_cast(data + 1); - for (int i = 0; i < size; i++) { - PID cand = datal[i]; - if (cand > max_elements_) { - throw std::runtime_error("cand error"); - } - - EstimateRecord candest; - get_bin_est_direct(q_to_centroids, query_wrapper, cand, candest); - - if (candest.est_dist < curest.est_dist) { - curest = candest; - curr_obj = cand; - changed = true; - } - } - } - } - - BoundedKNN boundedKnn(TOPK); - searchBaseLayerST_AdaptiveRerankOptDirect( - curr_obj, - std::max(ef, TOPK), - TOPK, - query_wrapper, - q_to_centroids, - rotated_query, - boundedKnn - ); - for (auto& candidate : boundedKnn.candidates()) { - result.emplace(candidate.record.est_dist, get_external_label(candidate.id)); - } - return result; -} - struct EstimateRecord { float est_dist; float low_dist; }; -template -inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( - PID ep_id, - size_t ef, - size_t TOPK, - SplitSingleQuery& query_wrapper, - std::vector& q_to_centroids, - [[maybe_unused]] const float* query, - BoundedKNN& boundedKNN -) { - VisitedSet* vl = visited_list_pool_->get_free_vislist(); - - // Use our bounded priority queue instead of the maxheap. - buffer::SearchBuffer candidate_set(ef); - - float distk = 1e10; - - EstimateRecord start_estimate_record; - if (ex_bits_ > 0) { - get_full_est_direct( - q_to_centroids, query_wrapper, ep_id, start_estimate_record - ); - } else { - get_bin_est_direct( - q_to_centroids, query_wrapper, ep_id, start_estimate_record - ); - } - float est_dist = start_estimate_record.est_dist; - float low_dist = start_estimate_record.low_dist; - - // Insert initial candidate. A removed point is still explored, never returned. - if (!is_removed(ep_id)) { - boundedKNN.insert({ResultRecord(est_dist, low_dist), ep_id}); - distk = est_dist; - } - candidate_set.insert(ep_id, est_dist); - - vl->set(ep_id); - - const size_t prefetch_size = (((padded_dim_ / 8) + 63) / 64) + 1; - const size_t prefetch_lookahead = 4; // Number of neighbors to prefetch in advance. - - while (candidate_set.has_next()) { - // Step 1 - get the next node to explore. - PID current_node_id = candidate_set.pop(); - int* data = (int*)get_linklist0(current_node_id); - size_t size = get_list_count((PID*)data); - - for (size_t p = 0; p < prefetch_lookahead; ++p) { - rabitqlib::memory::mem_prefetch_l1( - get_bindata_by_internalid(*(data + 1 + p)), prefetch_size - ); - } - // Iterate over neighbors. (List starts at index 1.) - for (size_t j = 1; j <= size; j++) { - int candidate_id = *(data + j); - - if (j + prefetch_lookahead <= size) { - rabitqlib::memory::mem_prefetch_l1( - get_bindata_by_internalid(*(data + j + prefetch_lookahead)), - prefetch_size - ); - } - - if (vl->get(candidate_id)) { - continue; - } - vl->set(candidate_id); - - EstimateRecord candest; - get_bin_est_direct( - q_to_centroids, query_wrapper, candidate_id, candest - ); - - // A removed neighbor keeps its binary estimate and stays in candidate_set. - bool flag_update_KNNs = !is_removed(static_cast(candidate_id)) && - (boundedKNN.size() < TOPK || candest.low_dist < distk); - - if (flag_update_KNNs) { - // Compute the full estimate if promising. - if (ex_bits_ > 0) { - get_full_est_direct( - q_to_centroids, query_wrapper, candidate_id, candest - ); - } - Candidate cand{ - ResultRecord(candest.est_dist, candest.low_dist), - static_cast(candidate_id)}; - boundedKNN.insert(cand); - distk = boundedKNN.worst().record.est_dist; - } - - if (!candidate_set.is_full(candest.est_dist)) { - candidate_set.insert(candidate_id, candest.est_dist); - } - - rabitqlib::memory::mem_prefetch_l2( - (char*)get_linklist0(candidate_set.next_id()), 2 - ); - } - } - - visited_list_pool_->release_vis_list(vl); -} - } // namespace rabitqlib::hnsw - -// Included last: these definitions need the complete class. -#include "rabitqlib/index/hnsw/hnsw_quant.hpp" // IWYU pragma: export diff --git a/include/rabitqlib/index/hnsw/hnsw_quant.hpp b/include/rabitqlib/index/hnsw/hnsw_quant.hpp index a138240f..d9ce54b6 100644 --- a/include/rabitqlib/index/hnsw/hnsw_quant.hpp +++ b/include/rabitqlib/index/hnsw/hnsw_quant.hpp @@ -1,627 +1,46 @@ -// Insertion into a built index: the construct path with get_quant_dist in place -// of get_data_dist, because raw input is borrowed only during construct. -// -// add routes, draws every level, allocates every link list and quantizes every -// point before reserving one block of slots and linking in parallel. Levels come -// first because the generator is shared; the slots are one block so labels do -// not depend on which thread finishes first. #pragma once -#include -#include -#include -#include -#include -#include -#include -#include -#include #include -#include #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/hnsw/hnsw.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/pack_excode.hpp" -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" // IWYU pragma: export +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw { namespace detail { -// The reconstructed query of get_quant_dist, reused while the query does not -// move. Cleared at the start of every insert so it cannot outlive its index. +// Owned by one add() worker. Codes and index identity remain fixed for its lifetime. +// Reuse the prepared query while its ID is unchanged, including across insertions. struct QuantQuery { PID id = kPidMax; std::vector rotated; std::optional> prepared; - void clear() { - id = kPidMax; - prepared.reset(); - } -}; - -inline QuantQuery& quant_query() { - thread_local QuantQuery query; - return query; -} -} // namespace detail - -// Code per dimension is (sign << ex_bits) + ex, up to nine bits, scaled by -// f_rescale_ex. -inline void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { - thread_local std::vector combined; - thread_local std::vector ex_codes; - combined.resize(padded_dim_); - - const ConstBinDataMap bin(get_bindata_by_internalid(id), padded_dim_); - const uint8_t* sign_code = bin.bin_code(); - for (size_t word = 0; word * 64 < padded_dim_; ++word) { - uint64_t bits = 0; - const size_t base = word * 64; - const size_t span = std::min(64, padded_dim_ - base); - std::memcpy(&bits, sign_code + (word * sizeof(uint64_t)), span / 8); - for (size_t j = 0; j < span; ++j) { - combined[base + j] = static_cast((bits >> (span - 1 - j)) & 1ULL); - } - } - - const auto* centroid = reinterpret_cast(centroids_memory_) + - (get_clusterid_by_internalid(id) * padded_dim_); - - if (ex_bits_ == 0) { - quant::reconstruct_full_vec( - combined.data(), centroid, padded_dim_, 1, bin.f_rescale(), out, metric_type_ - ); - return; - } - - ex_codes.resize(padded_dim_); - const ConstExDataMap ex(get_exdata_by_internalid(id), padded_dim_, ex_bits_); - quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( - ex.ex_code(), ex_codes.data(), padded_dim_, ex_bits_ - ); - for (size_t d = 0; d < padded_dim_; ++d) { - combined[d] = static_cast((combined[d] << ex_bits_) + ex_codes[d]); - } - - quant::reconstruct_full_vec( - combined.data(), - centroid, - padded_dim_, - ex_bits_ + 1, - ex.f_rescale_ex(), - out, - metric_type_ - ); -} - -// Only `query` is reconstructed; `target` is scored from its codes, as QGBuilder -// does. Callers hold `query` fixed across a loop, so it is prepared once. -inline float HierarchicalNSW::get_quant_dist(PID target, PID query) const { - detail::QuantQuery& cached = detail::quant_query(); - if (cached.id != query || !cached.prepared.has_value()) { - cached.rotated.resize(padded_dim_); - reconstruct_rotated(query, cached.rotated.data()); - cached.prepared.emplace( - cached.rotated.data(), padded_dim_, ex_bits_, query_config_, metric_type_ - ); - cached.id = query; - } - - const auto* centroid = reinterpret_cast(centroids_memory_) + - (get_clusterid_by_internalid(target) * padded_dim_); - const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_); - const float g_error = std::sqrt(dist_sqr); - const float g_add = (metric_type_ == METRIC_IP) - ? -dot_product(cached.rotated.data(), centroid, padded_dim_) - : dist_sqr; - - float est_dist = 0; - float low_dist = 0; - float ip_x0_qr = 0; - if (ex_bits_ == 0) { - split_single_estdist( - get_bindata_by_internalid(target), - *cached.prepared, - padded_dim_, - ip_x0_qr, - est_dist, - low_dist, - g_add, - g_error - ); - } else { - split_single_fulldist( - get_bindata_by_internalid(target), - get_exdata_by_internalid(target), - ip_func_, - *cached.prepared, - padded_dim_, - ex_bits_, - est_dist, - low_dist, - ip_x0_qr, - g_add, - g_error - ); - } - return est_dist; -} - -inline maxheap> HierarchicalNSW::search_base_layer_quant( - PID ep_id, PID cur_c, int layer -) { - std::unique_ptr vl(visited_list_pool_->get_free_vislist()); - - maxheap> top_candidates; - minheap> candidate_set; - - float lower_bound = get_quant_dist(ep_id, cur_c); - top_candidates.emplace(lower_bound, ep_id); - candidate_set.emplace(lower_bound, ep_id); - vl->set(ep_id); - - while (!candidate_set.empty()) { - std::pair curr_el_pair = candidate_set.top(); - if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { - break; - } - candidate_set.pop(); - - PID cur_node_num = curr_el_pair.second; - std::unique_lock lock(link_list_locks_[cur_node_num]); - - PID* data = - (layer == 0) ? get_linklist0(cur_node_num) : get_linklist(cur_node_num, layer); - size_t size = get_list_count(data); - auto* datal = data + 1; - - for (size_t j = 0; j < size; j++) { - PID candidate_id = datal[j]; - if (candidate_id >= cur_element_count_) { - throw std::runtime_error("cand error"); - } - if (vl->get(candidate_id)) { - continue; - } - vl->set(candidate_id); - - float dist1 = get_quant_dist(candidate_id, cur_c); - if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { - candidate_set.emplace(dist1, candidate_id); - top_candidates.emplace(dist1, candidate_id); - if (top_candidates.size() > ef_construction_) { - top_candidates.pop(); - } - if (!top_candidates.empty()) { - lower_bound = top_candidates.top().first; - } - } - } - } - visited_list_pool_->release_vis_list(vl.get()); - vl.release(); - return top_candidates; -} - -inline void HierarchicalNSW::get_neighbors_by_heuristic2_quant( - maxheap>& top_candidates, size_t M -) { - if (top_candidates.size() < M) { - return; - } - - minheap> queue_closest; - std::vector> return_list; - while (top_candidates.size() > 0) { - queue_closest.emplace(top_candidates.top()); - top_candidates.pop(); - } - - while (queue_closest.size() > 0) { - if (return_list.size() >= M) { - break; - } - std::pair current_pair = queue_closest.top(); - float dist_to_query = current_pair.first; - queue_closest.pop(); - bool good = true; - - for (std::pair second_pair : return_list) { - float curdist = get_quant_dist(second_pair.second, current_pair.second); - if (curdist < dist_to_query) { - good = false; - break; - } - } - if (good) { - return_list.push_back(current_pair); - } - } - - for (std::pair current_pair : return_list) { - top_candidates.emplace(current_pair); - } -} + QuantQuery() = default; + // prepared borrows rotated; copying would retain the source's storage pointer. + QuantQuery(const QuantQuery&) = delete; + QuantQuery& operator=(const QuantQuery&) = delete; -inline PID HierarchicalNSW::mutually_connect_quant( - PID cur_c, maxheap>& top_candidates, int level -) { - size_t max_m = level > 0 ? maxM_ : maxM0_; - get_neighbors_by_heuristic2_quant(top_candidates, M_); - if (top_candidates.size() > M_) { - throw std::runtime_error( - "Should be not be more than M_ candidates returned by the heuristic" - ); + QuantQuery(QuantQuery&& other) noexcept + : id(std::exchange(other.id, kPidMax)) + , rotated(std::move(other.rotated)) + , prepared(std::move(other.prepared)) { + other.prepared.reset(); } - std::vector selected_neighbors; - selected_neighbors.reserve(M_); - while (top_candidates.size() > 0) { - selected_neighbors.push_back(top_candidates.top().second); - top_candidates.pop(); - } - if (selected_neighbors.empty()) { - throw std::runtime_error("No neighbors selected for the new element"); - } - PID next_closest_entry_point = selected_neighbors.back(); - - { - PID* ll_cur = (level == 0) ? get_linklist0(cur_c) : get_linklist(cur_c, level); - if (*ll_cur > 0) { - throw std::runtime_error( - "The newly inserted element should have blank link list" - ); - } - - set_list_count(ll_cur, selected_neighbors.size()); - auto* data = ll_cur + 1; - for (size_t idx = 0; idx < selected_neighbors.size(); idx++) { - if (data[idx] != 0) { - throw std::runtime_error("Possible memory corruption"); - } - if (level > element_levels_[selected_neighbors[idx]]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); - } - data[idx] = selected_neighbors[idx]; + QuantQuery& operator=(QuantQuery&& other) noexcept { + if (this != &other) { + // std::allocator-backed vector moves retain the borrowed data address. + rotated = std::move(other.rotated); + prepared = std::move(other.prepared); + id = std::exchange(other.id, kPidMax); + other.prepared.reset(); } + return *this; } - - for (auto selected_neighbor : selected_neighbors) { - std::unique_lock lock(link_list_locks_[selected_neighbor]); - - PID* ll_other = (level == 0) ? get_linklist0(selected_neighbor) - : get_linklist(selected_neighbor, level); - size_t sz_link_list_other = get_list_count(ll_other); - - if (sz_link_list_other > max_m) { - throw std::runtime_error("Bad value of sz_link_list_other"); - } - if (selected_neighbor == cur_c) { - throw std::runtime_error("Trying to connect an element to itself"); - } - if (level > element_levels_[selected_neighbor]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); - } - - auto* data = ll_other + 1; - - bool is_cur_c_present = false; - for (size_t j = 0; j < sz_link_list_other; j++) { - if (data[j] == cur_c) { - is_cur_c_present = true; - break; - } - } - - if (!is_cur_c_present) { - if (sz_link_list_other < max_m) { - data[sz_link_list_other] = cur_c; - set_list_count(ll_other, sz_link_list_other + 1); - } else { - float d_max = get_quant_dist(cur_c, selected_neighbor); - maxheap> candidates; - candidates.emplace(d_max, cur_c); - for (size_t j = 0; j < sz_link_list_other; j++) { - candidates.emplace(get_quant_dist(data[j], selected_neighbor), data[j]); - } - - get_neighbors_by_heuristic2_quant(candidates, max_m); - - int indx = 0; - while (candidates.size() > 0) { - data[indx] = candidates.top().second; - candidates.pop(); - indx++; - } - set_list_count(ll_other, indx); - } - } - } - - return next_closest_entry_point; -} - -// The slot, the level and the link list are all decided by the caller, so this -// can run on many threads at once and the labels do not depend on their order. -inline void HierarchicalNSW::add_point_quant( - PID cur_c, int curlevel, std::unique_ptr& link_list -) { - detail::quant_query().clear(); - - std::unique_lock lock_el(link_list_locks_[cur_c]); - - std::unique_lock templock(global_); - int maxlevelcopy = maxlevel_; - PID curr_obj = enterpoint_node_; - if (curlevel <= maxlevelcopy) { - templock.unlock(); - } - const PID previous_entry_point = curr_obj; - - linkLists_[cur_c] = link_list.release(); - element_levels_[cur_c] = curlevel; - - try { - if (static_cast(curr_obj) != -1) { - if (curlevel < maxlevelcopy) { - float curdist = get_quant_dist(curr_obj, cur_c); - for (int level = maxlevelcopy; level > curlevel; level--) { - bool changed = true; - while (changed) { - changed = false; - std::unique_lock lock(link_list_locks_[curr_obj]); - PID* data = get_linklist(curr_obj, level); - int size = get_list_count(data); - auto* datal = data + 1; - for (int i = 0; i < size; i++) { - PID cand = datal[i]; - if (cand >= cur_element_count_) { - throw std::runtime_error("cand error"); - } - float d = get_quant_dist(cand, cur_c); - if (d < curdist) { - curdist = d; - curr_obj = cand; - changed = true; - } - } - } - } - } - - for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { - maxheap> top_candidates = - search_base_layer_quant(curr_obj, cur_c, level); - curr_obj = mutually_connect_quant(cur_c, top_candidates, level); - } - } else { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; - } - - if (curlevel > maxlevelcopy) { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; - } - } catch (...) { - // Upper layers may already expose this node. Keep a route to the old - // graph if base-layer linking was interrupted before choosing neighbors. - PID* base_links = get_linklist0(cur_c); - if (previous_entry_point != kPidMax && get_list_count(base_links) == 0) { - base_links[1] = previous_entry_point; - set_list_count(base_links, 1); - } - // An interrupted promotion has no incoming links above the previous - // maximum level. Discard those levels so save/load sees a valid graph. - if (curlevel > maxlevelcopy) { - element_levels_[cur_c] = std::max(0, maxlevelcopy); - if (element_levels_[cur_c] == 0) { - std::free(linkLists_[cur_c]); - linkLists_[cur_c] = nullptr; - } - } - throw; - } -} - -// Grows the capacity. Every pointer into the level-0 block is invalidated, so -// this must not run while another thread searches. -inline void HierarchicalNSW::resize(size_t new_max_elements) { - if (size_data_per_element_ == 0) { - throw std::logic_error("HNSW index must be initialized before resize"); - } - if (new_max_elements < cur_element_count_) { - throw std::invalid_argument("HNSW resize: capacity below the element count"); - } - if (new_max_elements > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("HNSW resize: capacity exceeds the supported ID range"); - } - if (new_max_elements == max_elements_) { - return; - } - if (new_max_elements > std::numeric_limits::max() / size_data_per_element_) { - throw std::invalid_argument("HNSW resize: capacity exceeds addressable storage"); - } - - // Everything new is built before anything old is released, so a failure - // leaves the index as it was. realloc is not an option: the level-0 block - // comes from aligned_alloc, which realloc does not keep aligned. - const size_t level0_bytes = new_max_elements * size_data_per_element_; - std::unique_ptr level0( - memory::huge_page_allocate(level0_bytes), - [](char* ptr) { memory::aligned_deallocate(ptr); } - ); - if (level0 == nullptr) { - throw std::bad_alloc(); - } - std::memset(level0.get(), 0, level0_bytes); - std::memcpy( - level0.get(), data_level0_memory_, cur_element_count_ * size_data_per_element_ - ); - - std::unique_ptr links( - static_cast(std::calloc(new_max_elements, sizeof(char*))), - [](char** ptr) { std::free(ptr); } - ); - if (links == nullptr) { - throw std::bad_alloc(); - } - std::memcpy(links.get(), linkLists_, cur_element_count_ * sizeof(char*)); - - std::vector levels(element_levels_); - levels.resize(new_max_elements); - std::vector locks(new_max_elements); - auto pool = std::make_unique(1, new_max_elements); - - memory::aligned_deallocate(data_level0_memory_); - data_level0_memory_ = level0.release(); - std::free(reinterpret_cast(linkLists_)); - linkLists_ = links.release(); - element_levels_.swap(levels); - link_list_locks_.swap(locks); - visited_list_pool_ = std::move(pool); - max_elements_ = new_max_elements; -} - -inline std::vector HierarchicalNSW::add( - const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads -) { - if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || - rotator_ == nullptr || num_cluster_ == 0) { - throw std::logic_error("HNSW index must be constructed or loaded before add"); - } - if (n == 0) { - return {}; - } - if (data == nullptr) { - throw std::invalid_argument("HNSW add inputs must not be null"); - } - if (n > max_elements_ - cur_element_count_) { - throw std::invalid_argument("HNSW add: not enough capacity"); - } - if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { - throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); - } - std::vector labels(n); - // Either the caller names the cluster of every point, or each one goes to its - // nearest centroid, ordered the way a query is routed. - std::vector assigned(n, 0); - if (cluster_ids != nullptr) { - for (size_t i = 0; i < n; ++i) { - if (cluster_ids[i] >= num_cluster_) { - throw std::invalid_argument("HNSW cluster ID is out of range"); - } - assigned[i] = cluster_ids[i]; - } - } else { - const auto* centroids = reinterpret_cast(centroids_memory_); - rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - std::vector rotated(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated.data()); - float best = std::numeric_limits::max(); - for (size_t c = 0; c < num_cluster_; ++c) { - const float dist = raw_dist_func_( - rotated.data(), centroids + (c * padded_dim_), padded_dim_ - ); - if (dist < best) { - best = dist; - assigned[i] = static_cast(c); - } - } - }); - } - - quant::RabitqConfig config; - if (faster) { - config = quant::faster_config(padded_dim_, ex_bits_ + 1); - } - - // Drawn here because the generator is shared: the inserts below run in - // parallel and must not touch it. - std::vector levels(n); - for (size_t i = 0; i < n; ++i) { - levels[i] = get_random_level(mult_); - } - - // Every link list is allocated before any of them is installed, so running - // out of memory leaves the index untouched. - std::vector> link_lists; - link_lists.reserve(n); - for (size_t i = 0; i < n; ++i) { - char* raw = nullptr; - if (levels[i] > 0) { - raw = - static_cast(std::calloc(1, (size_links_per_element_ * levels[i]) + 1) - ); - if (raw == nullptr) { - throw std::runtime_error( - "Not enough memory: HNSW add failed to allocate a link list" - ); - } - } - link_lists.emplace_back(raw, std::free); - } - - // Prepare every code before publishing the slots. If quantization or worker - // creation fails, the graph and point count are unchanged. - { - std::unique_lock lock_table(label_lookup_lock_); - if (cur_element_count_ + n > max_elements_) { - throw std::runtime_error("The number of elements exceeds the specified limit"); - } - const auto first = static_cast(cur_element_count_.load()); - rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - const auto id = static_cast(first + i); - labels[i] = id; - std::memset( - data_level0_memory_ + (id * size_data_per_element_), - 0, - size_data_per_element_ - ); - std::memcpy(get_external_label_pt(id), &id, sizeof(PID)); - std::memcpy(get_clusterid_pt(id), &assigned[i], sizeof(PID)); - std::vector rotated_data(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated_data.data()); - quant::quantize_split_single( - rotated_data.data(), - reinterpret_cast(centroids_memory_) + (assigned[i] * padded_dim_), - padded_dim_, - ex_bits_, - get_bindata_by_internalid(id), - get_exdata_by_internalid(id), - metric_type_, - config - ); - }); - size_t registered = 0; - try { - for (; registered < n; ++registered) { - label_lookup_[labels[registered]] = labels[registered]; - } - } catch (...) { - for (size_t i = 0; i < registered; ++i) { - label_lookup_.erase(labels[i]); - } - throw; - } - // Unstarted inserts are valid level-0 nodes with owned codes. Their - // upper-layer lists remain under RAII until an insert takes ownership. - for (size_t i = 0; i < n; ++i) { - element_levels_[labels[i]] = 0; - linkLists_[labels[i]] = nullptr; - } - cur_element_count_ += n; - } - - rabitqlib::ivf::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - add_point_quant(labels[i], levels[i], link_lists[i]); - }); - - return labels; -} +}; +} // namespace detail } // namespace rabitqlib::hnsw diff --git a/include/rabitqlib/index/ivf/cluster.hpp b/include/rabitqlib/index/ivf/cluster.hpp index 6917b16d..e20eae60 100644 --- a/include/rabitqlib/index/ivf/cluster.hpp +++ b/include/rabitqlib/index/ivf/cluster.hpp @@ -2,15 +2,11 @@ #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::ivf { -/** - * @brief Cluster is used for ivf index with rabitq+. Components are only used for record - * the addresses for different part of data. - * - */ +// Non-owning view of one IVF cluster's codes, reranking data and point IDs. class Cluster { private: size_t num_; // Num of vectors in this cluster @@ -20,9 +16,11 @@ class Cluster { public: explicit Cluster(size_t, char*, char*, PID*); - Cluster(const Cluster& other); - Cluster(Cluster&& other) noexcept; - ~Cluster() {} + Cluster(const Cluster&) = default; + Cluster(Cluster&&) noexcept = default; + Cluster& operator=(const Cluster&) = delete; + Cluster& operator=(Cluster&&) = delete; + ~Cluster() = default; [[nodiscard]] char* batch_data() const { return this->batch_data_; } @@ -36,15 +34,4 @@ class Cluster { inline Cluster::Cluster(size_t num, char* batch_data, char* ex_data, PID* ids) : num_(num), batch_data_(batch_data), ex_data_(ex_data), ids_(ids) {} -inline Cluster::Cluster(const Cluster& other) - : num_(other.num_) - , batch_data_(other.batch_data_) - , ex_data_(other.ex_data_) - , ids_(other.ids_) {} - -inline Cluster::Cluster(Cluster&& other) noexcept - : num_(other.num_) - , batch_data_(other.batch_data_) - , ex_data_(other.ex_data_) - , ids_(other.ids_) {} -} // namespace rabitqlib::ivf \ No newline at end of file +} // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp new file mode 100644 index 00000000..179e83bd --- /dev/null +++ b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp @@ -0,0 +1,151 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/third/hnswlib/hnswlib.h" +#include "rabitqlib/utils/path.hpp" + +namespace rabitqlib::ivf::detail { + +// hnswlib's native sidecar has no version tag. Check the existing writer's +// header and record lengths before its loader uses offsets or allocates links. +// As with the main index file, callers must not mutate the file during loading. +inline void validate_hnsw_sidecar_file( + const std::string& path, size_t expected_count, size_t vector_bytes +) { + std::ifstream input(io_impl::filesystem_path(path), std::ios::binary | std::ios::ate); + if (!input) { + throw std::runtime_error("Cannot open HNSW centroid sidecar"); + } + const auto end = input.tellg(); + input.seekg(0); + const auto read = [&](auto& value) { + if (!input.read(reinterpret_cast(&value), sizeof(value))) { + throw std::runtime_error("Truncated HNSW centroid sidecar"); + } + }; + size_t offset_level0 = 0, capacity = 0, count = 0, record_bytes = 0; + size_t label_offset = 0, data_offset = 0; + int max_level = 0; + hnswlib::tableint entry = 0; + size_t max_m = 0, max_m0 = 0, m = 0, ef_construction = 0; + double mult = 0; + read(offset_level0); + read(capacity); + read(count); + read(record_bytes); + read(label_offset); + read(data_offset); + read(max_level); + read(entry); + read(max_m); + read(max_m0); + read(m); + read(mult); + read(ef_construction); + if (count != expected_count || count == 0 || capacity < count || + count > std::numeric_limits::max()) { + throw std::runtime_error("HNSW centroid count does not match IVF cluster count"); + } + if (m < 2 || m > 10000 || max_m != m || max_m0 != 2 * m || !std::isfinite(mult) || + mult <= 0 || ef_construction < m) { + throw std::runtime_error("Invalid HNSW centroid graph parameters"); + } + const size_t upper_stride = + sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); + const size_t base_bytes = + sizeof(hnswlib::linklistsizeint) + max_m0 * sizeof(hnswlib::tableint); + if (offset_level0 != 0 || data_offset != base_bytes || data_offset > label_offset || + label_offset > record_bytes || label_offset - data_offset != vector_bytes || + record_bytes - label_offset != sizeof(hnswlib::labeltype)) { + throw std::runtime_error("HNSW centroid layout does not match IVF dimension"); + } + if (max_level < 0 || entry >= count) { + throw std::runtime_error("Invalid HNSW centroid entry point or maximum level"); + } + auto remaining = end - input.tellg(); + if (remaining < 0 || count > static_cast(remaining) / record_bytes) { + throw std::runtime_error("Truncated HNSW centroid base records"); + } + const auto base_records_bytes = static_cast(count * record_bytes); + input.seekg(base_records_bytes, std::ios::cur); + remaining -= base_records_bytes; + size_t highest_level = 0, entry_level = 0; + for (size_t i = 0; i < count; ++i) { + unsigned int link_bytes = 0; // Native type written by hnswlib::saveIndex. + if (remaining < static_cast(sizeof(link_bytes))) { + throw std::runtime_error("Truncated HNSW centroid sidecar"); + } + read(link_bytes); + remaining -= static_cast(sizeof(link_bytes)); + if (link_bytes % upper_stride != 0 || + link_bytes > static_cast(remaining)) { + throw std::runtime_error("Invalid HNSW centroid upper-link length"); + } + const size_t levels = link_bytes / upper_stride; + highest_level = std::max(highest_level, levels); + if (i == entry) { + entry_level = levels; + } + // Most nodes have no upper links. Keep the stream's read buffer intact + // and track the validated byte budget instead of querying each position. + if (link_bytes != 0) { + input.seekg(static_cast(link_bytes), std::ios::cur); + } + remaining -= static_cast(link_bytes); + } + if (!input || remaining != 0 || highest_level != static_cast(max_level) || + entry_level != highest_level) { + throw std::runtime_error("Invalid HNSW centroid level structure"); + } +} + +inline void validate_hnsw_centroid_graph(const hnswlib::HierarchicalNSW& graph) { + const size_t count = graph.cur_element_count; + if (graph.enterpoint_node_ >= count || graph.maxlevel_ < 0 || + graph.element_levels_[graph.enterpoint_node_] != graph.maxlevel_ || + graph.num_deleted_ != 0) { + throw std::runtime_error("Invalid HNSW centroid entry point or deletion state"); + } + for (size_t i = 0; i < count; ++i) { + if (graph.element_levels_[i] < 0 || graph.element_levels_[i] > graph.maxlevel_) { + throw std::runtime_error("Invalid HNSW centroid node level"); + } + for (int level = 0; level <= graph.element_levels_[i]; ++level) { + const char* links = + level == 0 ? graph.data_level0_memory_ + i * graph.size_data_per_element_ + + graph.offsetLevel0_ + : graph.linkLists_[i] + (static_cast(level) - 1) * + graph.size_links_per_element_; + uint16_t degree = 0; // The high two bytes hold flags, not the degree. + std::memcpy(°ree, links, sizeof(degree)); + if (degree > (level == 0 ? graph.maxM0_ : graph.maxM_)) { + throw std::runtime_error("Invalid HNSW centroid neighbor count"); + } + for (size_t j = 0; j < degree; ++j) { + hnswlib::tableint neighbor = 0; + std::memcpy( + &neighbor, + links + sizeof(hnswlib::linklistsizeint) + j * sizeof(neighbor), + sizeof(neighbor) + ); + if (neighbor >= count || graph.element_levels_[neighbor] < level) { + throw std::runtime_error( + "Invalid HNSW centroid neighbor or target level" + ); + } + } + } + } +} + +} // namespace rabitqlib::ivf::detail diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 733b346f..79213181 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -1,14 +1,9 @@ #pragma once -#include - #include -#include #include #include -#include #include -#include #include #include #include @@ -16,19 +11,19 @@ #include #include #include -#include #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/ivf/detail/hnsw_sidecar.hpp" +#include "rabitqlib/index/ivf/types.hpp" // IWYU pragma: export +#include "rabitqlib/index/types.hpp" #include "rabitqlib/third/hnswlib/hnswlib.h" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/tools.hpp" namespace rabitqlib::ivf { -enum class InitializerType : uint32_t { Auto = 0, Flat = 1, FlatRaBitQ = 2, HNSW = 3 }; - inline InitializerType resolve_initializer_type(InitializerType type, size_t clusters) { switch (type) { case InitializerType::Auto: @@ -44,97 +39,6 @@ inline InitializerType resolve_initializer_type(InitializerType type, size_t clu throw std::invalid_argument("Invalid IVF initializer type"); } -template -inline void parallel_for(size_t start, size_t end, size_t numThreads, Function fn) { - if (start >= end) { - return; - } - numThreads = std::min(resolve_num_threads(numThreads), end - start); - - if (numThreads == 1) { - for (size_t id = start; id < end; id++) { - fn(id, 0); - } - } else { -#if defined(_OPENMP) && _OPENMP >= 201307 - if (omp_get_proc_bind() != omp_proc_bind_false) { - // Native threads inherit only the caller's place. Let OpenMP honor - // the configured placement of the whole team when binding is enabled. - std::vector errors(numThreads); -#pragma omp parallel for num_threads(numThreads) schedule(dynamic) - for (size_t id = start; id < end; ++id) { - const size_t worker = static_cast(omp_get_thread_num()); - if (!errors[worker]) { - try { - fn(id, worker); - } catch (...) { errors[worker] = std::current_exception(); } - } - } - for (const auto& error : errors) { - if (error) { - std::rethrow_exception(error); - } - } - return; - } -#endif - std::vector threads; - std::atomic current(start); - - // keep track of exceptions in threads - // https://stackoverflow.com/a/32428427/1713196 - std::exception_ptr last_exception = nullptr; - std::mutex last_except_mutex; - - const auto join_threads = [&threads] { - for (auto& thread : threads) { - if (thread.joinable()) { - thread.join(); - } - } - }; - - try { - threads.reserve(numThreads); - for (size_t thread_id = 0; thread_id < numThreads; ++thread_id) { - threads.emplace_back([&, thread_id] { - while (true) { - size_t id = current.fetch_add(1); - - if (id >= end) { - break; - } - - try { - fn(id, thread_id); - } catch (...) { - std::unique_lock last_except_lock(last_except_mutex - ); - last_exception = std::current_exception(); - /* - * This will work even when current is the largest value that - * size_t can fit, because fetch_add returns the previous value - * before the increment (what will result in overflow - * and produce 0 instead of current + 1). - */ - current = end; - break; - } - } - }); - } - } catch (...) { - current = end; - join_threads(); - throw; - } - join_threads(); - if (last_exception) { - std::rethrow_exception(last_exception); - } - } -} - /** * @brief For IVF centroids, an initializer finds the candidate clusters. */ @@ -330,9 +234,14 @@ class HNSWInitializer : public Initializer { void add_vectors(const float* cent, size_t num_threads) override { size_t start = 0; size_t rows = num_cluster_; - parallel_for(start, rows, num_threads, [&](size_t row, size_t /*thread_id*/) { - alg_hnsw_->addPoint(cent + (row * dim_), row); - }); + rabitqlib::index_detail::parallel_for( + start, + rows, + num_threads, + [&](size_t row, size_t /*thread_id*/) { + alg_hnsw_->addPoint(cent + (row * dim_), row); + } + ); rebuild_label_lookup(); } @@ -349,6 +258,9 @@ class HNSWInitializer : public Initializer { alg_hnsw_->setEf(std::max(size_t{768}, 2 * nprobe)); std::priority_queue> result = alg_hnsw_->searchKnn(query, nprobe); + if (result.size() != nprobe) { + throw std::runtime_error("HNSW centroid graph returned too few candidates"); + } for (size_t i = 0; i < nprobe; ++i) { candidates[i].distance = metric_type_ == METRIC_IP @@ -369,7 +281,17 @@ class HNSWInitializer : public Initializer { void load(std::ifstream&, const char* filename) override { std::string hnsw(filename); hnsw += ".hnsw"; + detail::validate_hnsw_sidecar_file(hnsw, num_cluster_, space_->get_data_size()); alg_hnsw_->loadIndex(hnsw, space_, num_cluster_); + const auto& graph = *alg_hnsw_; + if (graph.offsetData_ > graph.label_offset_ || + graph.label_offset_ > graph.size_data_per_element_ || + graph.label_offset_ - graph.offsetData_ != space_->get_data_size() || + graph.size_data_per_element_ - graph.label_offset_ != + sizeof(hnswlib::labeltype)) { + throw std::runtime_error("HNSW centroid layout does not match IVF dimension"); + } + detail::validate_hnsw_centroid_graph(graph); rebuild_label_lookup(); } diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 93256f44..b3362066 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -1,44 +1,32 @@ #pragma once -#include - -#include -#include -#include -#include -#include #include #include -#include -#include -#include -#include #include #include #include -#include -#include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/core/metric.hpp" // IWYU pragma: export +#include "rabitqlib/core/types.hpp" // IWYU pragma: export #include "rabitqlib/index/ivf/cluster.hpp" -#include "rabitqlib/index/ivf/initializer.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/index/ivf/types.hpp" // IWYU pragma: export +#include "rabitqlib/index/types.hpp" // IWYU pragma: export #include "rabitqlib/utils/memory.hpp" -#include "rabitqlib/utils/path.hpp" -#include "rabitqlib/utils/rotator.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib { +template +class Rotator; +template +class SplitBatchQuery; +namespace buffer { +template +class SearchBuffer; +} +} // namespace rabitqlib namespace rabitqlib::ivf { -namespace detail { -void insert_candidates(buffer::SearchBuffer&, const PID*, const float*, size_t); -} // namespace detail +class Initializer; class IVF { private: @@ -48,22 +36,23 @@ class IVF { using IdStorage = std::vector>; InitializerType initializer_type_ = InitializerType::Auto; - std::unique_ptr initer_; // initializer for candidate clusters - ByteStorage batch_storage_; // 1-bit code and factors - ByteStorage ex_storage_; // extra-bit codes and packed factors - FloatStorage raw_storage_; // original vectors for raw reranking - IdStorage id_storage_; // point IDs organized by cluster - size_t num_ = 0; // num of data points - size_t dim_ = 0; // dimension of data points - size_t padded_dim_ = 0; // dimension after padding - size_t num_cluster_ = 0; // num of centroids (clusters) - bool raw_reranking_ = false; // raw vectors replace extra-bit codes - bool ready_ = false; // construction or loading completed - size_t ex_bits_ = 0; // total bits = ex_bits_ + 1 - RotatorType type_ = RotatorType::FhtKacRotator; // type of rotator - std::unique_ptr> rotator_; // Data Rotator - std::vector cluster_lst_; // List of clusters in ivf - MetricType metric_type_ = rabitqlib::METRIC_L2; // metric type + std::unique_ptr initializer_; // initializer for candidate clusters + ByteStorage batch_storage_; // 1-bit code and factors + ByteStorage ex_storage_; // extra-bit codes and packed factors + FloatStorage raw_storage_; // original vectors for raw reranking + IdStorage id_storage_; // point IDs organized by cluster + size_t max_elements_ = 0; // Logical capacity; construct fills it, add grows it. + size_t num_removed_ = 0; // Recovered from removal markers when loading. + size_t dim_ = 0; + size_t padded_dim_ = 0; + size_t num_clusters_ = 0; + bool raw_reranking_ = false; // raw vectors replace extra-bit codes + bool ready_ = false; // construction or loading completed + size_t ex_bits_ = 0; // total bits = ex_bits_ + 1 + RotatorType rotator_type_ = RotatorType::FhtKacRotator; + std::unique_ptr> rotator_; + std::vector clusters_; + MetricType metric_type_ = rabitqlib::METRIC_L2; float (*ip_func_)(const float*, const uint8_t*, size_t) = nullptr; static constexpr uint64_t kRawFormatMagic = 0x3157415251464252ULL; @@ -79,33 +68,16 @@ class IVF { const; // True if `n` points cannot be addressed by the storage size computations - [[nodiscard]] bool exceeds_storage(size_t n) const { - const size_t max_size = std::numeric_limits::max(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); - const size_t rerank_bytes = rerank_vector_bytes(); - return n > max_size / sizeof(PID) || n > max_size / sizeof(float) / padded_dim_ || - n > max_size / batch_bytes || - (rerank_bytes != 0 && n > max_size / rerank_bytes); - } + [[nodiscard]] bool exceeds_storage(size_t n) const; - [[nodiscard]] size_t ids_bytes() const { return sizeof(PID) * num_; } + [[nodiscard]] size_t ids_bytes() const { return sizeof(PID) * max_elements_; } // get num of bytes used for 1-bit code and corresponding factors - [[nodiscard]] size_t batch_data_bytes(const std::vector& cluster_sizes) const { - assert(cluster_sizes.size() == num_cluster_); // num of clusters - size_t total_blocks = 0; - for (auto size : cluster_sizes) { - total_blocks += div_round_up(size, fastscan::kBatchSize); - } - return total_blocks * BatchDataMap::data_bytes(padded_dim_); - } + [[nodiscard]] size_t batch_data_bytes(const std::vector& cluster_sizes) const; - [[nodiscard]] size_t rerank_vector_bytes() const { - return raw_reranking_ ? sizeof(float) * dim_ - : ExDataMap::data_bytes(padded_dim_, ex_bits_); - } + [[nodiscard]] inline size_t rerank_vector_bytes() const; - [[nodiscard]] size_t ex_data_bytes() const { return rerank_vector_bytes() * num_; } + [[nodiscard]] inline size_t ex_data_bytes() const; [[nodiscard]] char* batch_data() { return reinterpret_cast(batch_storage_.data()); @@ -133,69 +105,58 @@ class IVF { void init_clusters(const std::vector&); - void free_memory() { - initer_.reset(); - ByteStorage().swap(batch_storage_); - ByteStorage().swap(ex_storage_); - FloatStorage().swap(raw_storage_); - IdStorage().swap(id_storage_); - } + void free_memory(); - void swap(IVF& other) noexcept { - using std::swap; - swap(initializer_type_, other.initializer_type_); - swap(initer_, other.initer_); - swap(batch_storage_, other.batch_storage_); - swap(ex_storage_, other.ex_storage_); - swap(raw_storage_, other.raw_storage_); - swap(id_storage_, other.id_storage_); - swap(num_, other.num_); - swap(dim_, other.dim_); - swap(padded_dim_, other.padded_dim_); - swap(num_cluster_, other.num_cluster_); - swap(raw_reranking_, other.raw_reranking_); - swap(ready_, other.ready_); - swap(ex_bits_, other.ex_bits_); - swap(type_, other.type_); - swap(rotator_, other.rotator_); - swap(cluster_lst_, other.cluster_lst_); - swap(metric_type_, other.metric_type_); - swap(ip_func_, other.ip_func_); - } + void swap(IVF& other) noexcept; - void + inline void search_cluster(const Cluster&, const SplitBatchQuery&, buffer::SearchBuffer&, bool, const float*) const; - void + inline void scan_one_batch(const char* batch_data, const char* ex_data, const PID* ids, const SplitBatchQuery& q_obj, buffer::SearchBuffer& knns, size_t num_points, bool, const float*) const; - void validate_search(const float*, size_t, size_t, const PID*) const; + inline void validate_search(const float*, size_t, size_t, const PID*) const; public: - explicit IVF() = default; + explicit IVF() noexcept; + // dim is the original input dimension; total_bits is 1-9 for quantized + // storage (including the sign bit), or 32 for raw-vector reranking. explicit IVF( - size_t, - size_t, - size_t, - size_t, + size_t num_points, + size_t dim, + size_t num_clusters, + size_t total_bits, MetricType metric_type = rabitqlib::METRIC_L2, - RotatorType type = RotatorType::FhtKacRotator, + RotatorType rotator_type = RotatorType::FhtKacRotator, InitializerType initializer = InitializerType::Auto ); ~IVF(); - // Number of stored points, including any removed with `remove`. `add` increases it. - [[nodiscard]] size_t max_elements() const { return num_; } + // Logical capacity, including slots occupied by removed points. add grows it. + [[nodiscard]] size_t max_elements() const { return max_elements_; } + // Unbuilt indexes have no stored or searchable points. + [[nodiscard]] size_t num_points() const { + return ready_ ? max_elements_ - num_removed_ : 0; + } + [[nodiscard]] size_t num_stored_points() const { return ready_ ? max_elements_ : 0; } [[nodiscard]] size_t dimension() const { return dim_; } [[nodiscard]] size_t nbits() const { return raw_reranking_ ? 32 : ex_bits_ + 1; } [[nodiscard]] MetricType metric_type() const { return metric_type_; } - [[nodiscard]] RotatorType rotator_type() const { return type_; } + [[nodiscard]] RotatorType rotator_type() const { return rotator_type_; } [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; } - void construct(const float*, const float*, const PID*, bool, size_t); + // faster selects fixed-scale extra-bit quantization instead of exact scale search, + // trading quantization accuracy for speed. One-bit and raw storage are unaffected. + void construct( + const float* data, + const float* centroids, + const PID* cluster_ids, + bool faster = false, + size_t num_threads = std::numeric_limits::max() + ); /** * @brief Append points to a constructed or loaded index without the original data. @@ -203,7 +164,7 @@ class IVF { * Points are quantized against the centroids and rotator the index already holds, so * the centroids never move: recall can degrade if the added data drifts away from the * data the centroids were trained on. Point `i` of `data` receives PID - * `max_elements() + i`. The index is unchanged if this throws. + * `num_stored_points() + i`. The index is unchanged if this throws. * * The grown index is built next to the old one, so a call needs memory for both and * takes time proportional to the whole index, not just to `n`: prefer a few large @@ -214,7 +175,8 @@ class IVF { * @param n Number of new points * @param cluster_ids Cluster of each new point, or nullptr to route each point to its * nearest centroid the same way a query is routed - * @param faster Same meaning as in construct + * @param faster Use fixed-scale extra-bit quantization for speed at a possible + * loss of accuracy; has no effect on one-bit or raw storage * @param num_threads Threads used for routing and quantization */ void add( @@ -228,7 +190,8 @@ class IVF { /** * @brief Exclude points from all later search results. * - * A removed point keeps its storage, still counts in `max_elements()`, and search may + * A removed point stays in num_stored_points() and occupies capacity, but no longer + * counts in num_points(). Search may * return fewer than k results (unfilled slots use kPidMax). Removal is stored in the * index, so it survives save and load, and is idempotent. Removed points cannot be * restored. Not safe to call while another thread searches or adds. @@ -239,21 +202,36 @@ class IVF { */ size_t remove(const PID* ids_to_remove, size_t n); - void save(const char*) const; + void save(const char* filename) const; - void load(const char*); + void load(const char* filename); // Automatically use HACC for 4-9 quantized bits; use standard FastScan otherwise. - void search(const float*, size_t, size_t, PID*, float* = nullptr) const; + // Distances estimate squared L2 or 1 - dot according to the index metric. + void search( + const float* query, + size_t k, + size_t nprobe, + PID* results, + float* distances = nullptr + ) const; - void search(const float*, size_t, size_t, PID*, bool) const; + void search(const float* query, size_t k, size_t nprobe, PID* results, bool use_hacc) + const; - void search(const float*, size_t, size_t, PID*, float*, bool) const; + void search( + const float* query, + size_t k, + size_t nprobe, + PID* results, + float* distances, + bool use_hacc + ) const; /** * Search contiguous row-major queries using at most num_threads workers. * queries contains num_queries * dimension() floats. results and optional - * dists each hold num_queries * k elements. Empty batches do no work. + * distances each hold num_queries * k elements. Empty batches do no work. * use_hacc defaults to the same automatic policy as search(). * Input and output buffers must not overlap. Concurrent calls are safe if * the index is not being modified and their output buffers are separate. @@ -264,1001 +242,14 @@ class IVF { size_t k, size_t nprobe, PID* results, - float* dists = nullptr, + float* distances = nullptr, std::optional use_hacc = std::nullopt, size_t num_threads = 1 ) const; [[nodiscard]] size_t padded_dim() const { return this->padded_dim_; } - [[nodiscard]] size_t num_clusters() const { return this->num_cluster_; } + [[nodiscard]] size_t num_clusters() const { return this->num_clusters_; } }; -inline IVF::IVF( - size_t n, - size_t dim, - size_t cluster_num, - size_t bits, - MetricType metric_type, - RotatorType type, - InitializerType initializer -) - : initializer_type_(resolve_initializer_type(initializer, cluster_num)) - , num_(n) - , dim_(dim) - , padded_dim_(dim) - , num_cluster_(cluster_num) - , raw_reranking_(bits == 32) - , ex_bits_(bits == 32 ? 0 : bits - 1) - , type_(type) - , metric_type_(metric_type) { - validate_metric_type(metric_type); - if ((bits < 1 || bits > 9) && bits != 32) { - throw std::invalid_argument("IVF bits must be in [1, 9] or 32 for raw reranking"); - }; - if (n == 0 || n > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("IVF point count exceeds the supported ID range"); - } - if (cluster_num == 0 || cluster_num > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("IVF cluster count exceeds the supported ID range"); - } - if (dim == 0 || dim > (std::numeric_limits::max() / 32) - 64) { - throw std::invalid_argument("IVF dimension is invalid or too large"); - } - padded_dim_ = round_up_to_multiple(dim_, 32); - const size_t max_size = std::numeric_limits::max(); - if (exceeds_storage(n) || cluster_num > max_size / sizeof(float) / padded_dim_ || - (type == RotatorType::MatrixRotator && dim > max_size / sizeof(float) / padded_dim_ - )) { - throw std::invalid_argument("IVF configuration exceeds addressable storage"); - } - rotator_.reset(choose_rotator(dim, type, padded_dim_)); - /* check size */ - assert(padded_dim_ % 32 == 0); - assert(padded_dim_ >= dim_); -} - -inline IVF::~IVF() { free_memory(); } - -/** - * @brief Construct clusters in IVF - * - * @param data Data objects (N*DIM) - * @param centroids Centroid vectors (K*DIM) - * @param cluster_ids Cluster ID for each data object - */ -inline void IVF::construct( - const float* data, - const float* centroids, - const PID* cluster_ids, - bool faster = false, - size_t num_threads = std::numeric_limits::max() -) { - if (num_ == 0 || dim_ == 0 || num_cluster_ == 0 || rotator_ == nullptr) { - throw std::logic_error("IVF must be configured before construction"); - } - if (data == nullptr || centroids == nullptr || cluster_ids == nullptr) { - throw std::invalid_argument("IVF construction inputs must not be null"); - } - - // get id list for each cluster - std::vector counts(num_cluster_, 0); - std::vector> id_lists(num_cluster_); - for (size_t i = 0; i < num_; ++i) { - PID cid = cluster_ids[i]; - if (cid >= num_cluster_) { - throw std::invalid_argument("Cluster ID is out of range"); - } - id_lists[cid].push_back(static_cast(i)); - counts[cid] += 1; - } - - allocate_memory(counts); - - // init the cluster list - init_clusters(counts); - - // all rotated centroids - std::vector rotated_centroids(num_cluster_ * padded_dim_); - - quant::RabitqConfig config; - if (faster) { - config = quant::faster_config(padded_dim_, ex_bits_ + 1); - } - - num_threads = std::min(resolve_num_threads(num_threads), num_cluster_); - const size_t scratch_points = - std::min(fastscan::kBatchSize, *std::max_element(counts.begin(), counts.end())); - // Each worker reuses one batch, independently of the largest cluster size. - // Allocate before entering OpenMP so allocation failures reach the caller. - std::vector> rotated_blocks(num_threads); - for (auto& block : rotated_blocks) { - block.resize(scratch_points * padded_dim_); - } - /* Quantize each cluster */ -#pragma omp parallel for schedule(dynamic) num_threads(num_threads) - for (std::ptrdiff_t index = 0; index < static_cast(num_cluster_); - ++index) { - const size_t i = static_cast(index); - const float* cur_centroid = centroids + (i * dim_); - float* cur_rotated_c = &rotated_centroids[i * padded_dim_]; - Cluster& cp = cluster_lst_[i]; - quantize_cluster( - cp, - id_lists[i], - data, - cur_centroid, - cur_rotated_c, - rotated_blocks[static_cast(omp_get_thread_num())].data(), - config - ); - } - - this->initer_->add_vectors(rotated_centroids.data(), num_threads); - ready_ = true; -} - -inline void IVF::allocate_memory(const std::vector& cluster_sizes) { - ready_ = false; - free_memory(); - cluster_lst_.clear(); - - if (initializer_type_ == InitializerType::Flat) { - this->initer_ = - std::make_unique(padded_dim_, num_cluster_, metric_type_); - } else if (initializer_type_ == InitializerType::FlatRaBitQ) { - this->initer_ = std::make_unique( - padded_dim_, num_cluster_, metric_type_ - ); - } else { - this->initer_ = - std::make_unique(padded_dim_, num_cluster_, metric_type_); - } - batch_storage_ = ByteStorage(batch_data_bytes(cluster_sizes)); - if (rerank_vector_bytes() > 0) { - if (raw_reranking_) { - raw_storage_ = FloatStorage(num_ * dim_); - } else { - ex_storage_ = ByteStorage(ex_data_bytes()); - } - } - id_storage_ = IdStorage(num_); - - this->ip_func_ = select_excode_ipfunc(ex_bits_); -} - -/** - * @brief Initialize the cluster list by finding each cluster's storage offsets. - */ -inline void IVF::init_clusters(const std::vector& cluster_sizes) { - this->cluster_lst_.reserve(num_cluster_); - size_t added_vectors = 0; - size_t added_batches = 0; - for (size_t i = 0; i < num_cluster_; ++i) { - // find data location for current cluster - size_t num = cluster_sizes[i]; - size_t num_batches = div_round_up(num, fastscan::kBatchSize); - - char* current_batch_data = - batch_data() + (BatchDataMap::data_bytes(padded_dim_) * added_batches); - char* current_ex_data = rerank_vector_bytes() > 0 - ? ex_data() + (added_vectors * rerank_vector_bytes()) - : nullptr; - PID* cluster_ids = ids() + added_vectors; - - Cluster cur_cluster(num, current_batch_data, current_ex_data, cluster_ids); - this->cluster_lst_.push_back(std::move(cur_cluster)); - - added_vectors += num; - added_batches += num_batches; - } -} - -inline void IVF::quantize_cluster( - Cluster& cp, - const std::vector& IDs, - const float* data, - const float* cur_centroid, - float* rotated_centroid, - float* rotated_data, - const quant::RabitqConfig& config -) { - size_t num_points = IDs.size(); - if (cp.num() != num_points) { - throw std::invalid_argument("Cluster size and ID count differ"); - } - - // copy ids - std::copy(IDs.begin(), IDs.end(), cp.ids()); - - // rotate centroid - this->rotator_->rotate(cur_centroid, rotated_centroid); - - char* batch_data = cp.batch_data(); - char* ex_data = cp.ex_data(); - for (size_t i = 0; i < num_points; i += fastscan::kBatchSize) { - size_t n = std::min(fastscan::kBatchSize, num_points - i); - for (size_t j = 0; j < n; ++j) { - const float* vector = data + (IDs[i + j] * dim_); - rotator_->rotate(vector, rotated_data + (j * padded_dim_)); - if (raw_reranking_) { - std::memcpy( - cp.ex_data() + ((i + j) * rerank_vector_bytes()), - vector, - rerank_vector_bytes() - ); - } - } - - quant::quantize_split_batch( - rotated_data, - rotated_centroid, - n, - padded_dim_, - ex_bits_, - batch_data, - ex_data, - metric_type_, - config - ); - - batch_data += BatchDataMap::data_bytes(padded_dim_); - if (ex_bits_ > 0) { - ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * n; - } - } -} - -inline void IVF::add( - const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads -) { - if (!ready_ || initer_ == nullptr || rotator_ == nullptr || - cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || - id_storage_.size() != num_) { - throw std::logic_error("IVF index must be constructed or loaded before add"); - } - if (n == 0) { - return; - } - if (data == nullptr) { - throw std::invalid_argument("IVF add data must not be null"); - } - if (n > buffer::kSearchBufferMaxPointCount - num_) { - throw std::invalid_argument("IVF point count exceeds the supported ID range"); - } - const size_t new_num = num_ + n; - if (exceeds_storage(new_num)) { - throw std::invalid_argument("IVF configuration exceeds addressable storage"); - } - - /* Assign every new point to a cluster */ - std::vector assigned(n); - if (cluster_ids != nullptr) { - for (size_t i = 0; i < n; ++i) { - if (cluster_ids[i] >= num_cluster_) { - throw std::invalid_argument("Cluster ID is out of range"); - } - assigned[i] = cluster_ids[i]; - } - } else { - parallel_for(0, n, num_threads, [&](size_t row, size_t /*thread_id*/) { - std::vector rotated(padded_dim_); - std::vector> nearest(1); - rotator_->rotate(data + (row * dim_), rotated.data()); - initer_->centroids_distances(rotated.data(), 1, nearest); - assigned[row] = nearest[0].id; - }); - } - - std::vector> rows_of(num_cluster_); - for (size_t i = 0; i < n; ++i) { - rows_of[assigned[i]].push_back(static_cast(i)); - } - std::vector new_sizes(num_cluster_); - for (size_t i = 0; i < num_cluster_; ++i) { - new_sizes[i] = cluster_lst_[i].num() + rows_of[i].size(); - } - - /* Build the grown storage on the side so the index is unchanged on failure */ - ByteStorage new_batch(batch_data_bytes(new_sizes)); - ByteStorage new_ex; - FloatStorage new_raw; - const size_t rerank_bytes = rerank_vector_bytes(); - if (raw_reranking_) { - new_raw = FloatStorage(new_num * dim_); - } else if (rerank_bytes > 0) { - new_ex = ByteStorage(new_num * rerank_bytes); - } - IdStorage new_ids(new_num); - char* new_rerank = raw_reranking_ ? reinterpret_cast(new_raw.data()) - : reinterpret_cast(new_ex.data()); - - std::vector batch_offset(num_cluster_); - std::vector vector_offset(num_cluster_); - size_t added_batches = 0; - size_t added_vectors = 0; - for (size_t i = 0; i < num_cluster_; ++i) { - batch_offset[i] = added_batches; - vector_offset[i] = added_vectors; - added_batches += div_round_up(new_sizes[i], fastscan::kBatchSize); - added_vectors += new_sizes[i]; - } - - quant::RabitqConfig config; - if (faster) { - config = quant::faster_config(padded_dim_, ex_bits_ + 1); - } - - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); - const PID first_pid = static_cast(num_); - // An exception must not leave an OpenMP region, so keep the first one and rethrow it. - std::exception_ptr failure; -#pragma omp parallel for schedule(dynamic) num_threads(resolve_num_threads(num_threads)) - for (std::ptrdiff_t index = 0; index < static_cast(num_cluster_); - ++index) { - const size_t i = static_cast(index); - try { - extend_cluster( - cluster_lst_[i], - rows_of[i], - data, - initer_->centroid(static_cast(i)), - first_pid, - reinterpret_cast(new_batch.data()) + (batch_offset[i] * batch_bytes), - rerank_bytes > 0 ? new_rerank + (vector_offset[i] * rerank_bytes) : nullptr, - new_ids.data() + vector_offset[i], - config - ); - } catch (...) { -#pragma omp critical(rabitq_ivf_add_failure) - if (!failure) { - failure = std::current_exception(); - } - } - } - if (failure) { - std::rethrow_exception(failure); - } - - /* Commit. Rebuilding the cluster list reuses its capacity, so it cannot throw. */ - batch_storage_.swap(new_batch); - ex_storage_.swap(new_ex); - raw_storage_.swap(new_raw); - id_storage_.swap(new_ids); - num_ = new_num; - cluster_lst_.clear(); - init_clusters(new_sizes); -} - -/** - * @brief Copy one cluster into its grown storage and append new points to it. - * - * Full batches, rerank data and IDs are copied unchanged. The old partial tail batch is - * unpacked and repacked together with the new points, so the result is byte-identical to - * quantizing the same points in the same order in one construct call. - */ -inline void IVF::extend_cluster( - const Cluster& old_cluster, - const std::vector& rows, - const float* data, - const float* rotated_centroid, - PID first_pid, - char* batch_dst, - char* rerank_dst, - PID* ids_dst, - const quant::RabitqConfig& config -) const { - const size_t old_num = old_cluster.num(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); - const size_t rerank_bytes = rerank_vector_bytes(); - const size_t cols = padded_dim_ / 8; - - if (rows.empty()) { - const size_t old_batches = div_round_up(old_num, fastscan::kBatchSize); - std::copy_n(old_cluster.batch_data(), old_batches * batch_bytes, batch_dst); - std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); - std::copy_n(old_cluster.ids(), old_num, ids_dst); - return; - } - - const size_t full_batches = old_num / fastscan::kBatchSize; - const size_t kept = old_num - (full_batches * fastscan::kBatchSize); - const size_t tail_num = kept + rows.size(); - - std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); - std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); - std::copy_n(old_cluster.ids(), old_num, ids_dst); - - /* Codes and factors of every point in the tail, old points first */ - std::vector codes(tail_num * cols); - std::vector f_add(tail_num); - std::vector f_rescale(tail_num); - std::vector f_error(tail_num); - if (kept > 0) { - ConstBatchDataMap old_tail( - old_cluster.batch_data() + (full_batches * batch_bytes), padded_dim_ - ); - fastscan::unpack_codes(padded_dim_, old_tail.bin_code(), kept, codes.data()); - old_tail.f_add().copy_to(f_add.data(), kept); - old_tail.f_rescale().copy_to(f_rescale.data(), kept); - old_tail.f_error().copy_to(f_error.data(), kept); - } - - std::vector rotated(padded_dim_); - for (size_t j = 0; j < rows.size(); ++j) { - const float* vector = data + (static_cast(rows[j]) * dim_); - const size_t lane = kept + j; - rotator_->rotate(vector, rotated.data()); - quant::quantize_compact_one_bit( - rotated.data(), - rotated_centroid, - padded_dim_, - codes.data() + (lane * cols), - f_add[lane], - f_rescale[lane], - f_error[lane], - metric_type_ - ); - char* rerank = rerank_dst + ((old_num + j) * rerank_bytes); - if (raw_reranking_) { - std::memcpy(rerank, vector, rerank_bytes); - } else if (ex_bits_ > 0) { - quant::quantize_compact_ex_bits( - rotated.data(), - rotated_centroid, - padded_dim_, - ex_bits_, - rerank, - metric_type_, - config - ); - } - ids_dst[old_num + j] = first_pid + rows[j]; - } - - /* Repack the tail into batches */ - char* batch = batch_dst + (full_batches * batch_bytes); - for (size_t start = 0; start < tail_num; start += fastscan::kBatchSize) { - const size_t count = std::min(fastscan::kBatchSize, tail_num - start); - BatchDataMap cur_batch(batch, padded_dim_); - fastscan::pack_codes( - padded_dim_, codes.data() + (start * cols), count, cur_batch.bin_code() - ); - for (size_t i = 0; i < count; ++i) { - cur_batch.f_add()[i] = f_add[start + i]; - cur_batch.f_rescale()[i] = f_rescale[start + i]; - cur_batch.f_error()[i] = f_error[start + i]; - } - batch += batch_bytes; - } -} - -inline size_t IVF::remove(const PID* ids_to_remove, size_t n) { - if (!ready_ || initer_ == nullptr || rotator_ == nullptr || - cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || - id_storage_.size() != num_) { - throw std::logic_error("IVF index must be constructed or loaded before remove"); - } - if (n == 0) { - return 0; - } - if (ids_to_remove == nullptr) { - throw std::invalid_argument("IVF remove IDs must not be null"); - } - - // Validate every ID before changing anything - std::vector marked(num_, false); - for (size_t i = 0; i < n; ++i) { - if (ids_to_remove[i] >= num_) { - throw std::invalid_argument("IVF remove ID is out of range"); - } - marked[ids_to_remove[i]] = true; - } - - // A point whose f_add is +inf has +inf estimated and lower-bound distances in every - // scan path, so it never enters the result buffer or triggers reranking. - constexpr float kRemoved = std::numeric_limits::infinity(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); - size_t removed = 0; - for (const auto& cluster : cluster_lst_) { - const PID* cluster_ids = cluster.ids(); - for (size_t pos = 0; pos < cluster.num(); ++pos) { - if (!marked[cluster_ids[pos]]) { - continue; - } - BatchDataMap cur_batch( - cluster.batch_data() + ((pos / fastscan::kBatchSize) * batch_bytes), - padded_dim_ - ); - const size_t lane = pos % fastscan::kBatchSize; - if (static_cast(cur_batch.f_add()[lane]) != kRemoved) { - cur_batch.f_add()[lane] = kRemoved; - ++removed; - } - } - } - return removed; -} - -inline void IVF::save(const char* filename) const { - if (!ready_ || initer_ == nullptr || rotator_ == nullptr || - cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || - id_storage_.size() != num_) { - throw std::logic_error("Cannot save an unconstructed IVF index"); - } - if (filename == nullptr || filename[0] == '\0') { - throw std::invalid_argument("IVF save filename must not be empty"); - } - - std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); - output.exceptions(std::ios::failbit | std::ios::badbit); - const uint32_t flags = raw_reranking_ ? 1 : 0; - const uint64_t stored_padded_dim = padded_dim_; - output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); - output.write(reinterpret_cast(&kFormatVersion), sizeof(kFormatVersion)); - output.write(reinterpret_cast(&flags), sizeof(flags)); - output.write( - reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim) - ); - - const auto stored_initializer = static_cast(initializer_type_); - output.write( - reinterpret_cast(&stored_initializer), sizeof(stored_initializer) - ); - - /* Save meta data */ - output.write(reinterpret_cast(&num_), sizeof(size_t)); - output.write(reinterpret_cast(&dim_), sizeof(size_t)); - output.write(reinterpret_cast(&num_cluster_), sizeof(size_t)); - output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); - output.write(reinterpret_cast(&type_), sizeof(type_)); - output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - - /* Save number of vectors of each cluster */ - std::vector cluster_sizes; - cluster_sizes.reserve(num_cluster_); - for (const auto& cur_cluster : cluster_lst_) { - cluster_sizes.push_back(cur_cluster.num()); - } - output.write( - reinterpret_cast(cluster_sizes.data()), - static_cast(sizeof(size_t) * num_cluster_) - ); - - /* Save rotator */ - this->rotator_->save(output); - - /* Save data */ - this->initer_->save(output, filename); - output.write( - batch_data(), static_cast(batch_data_bytes(cluster_sizes)) - ); - if (ex_data_bytes() != 0) { - output.write(ex_data(), static_cast(ex_data_bytes())); - } - output.write( - reinterpret_cast(ids()), static_cast(ids_bytes()) - ); - - output.close(); -} - -inline void IVF::load(const char* filename) { - if (filename == nullptr || filename[0] == '\0') { - throw std::invalid_argument("IVF load filename must not be empty"); - } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); - if (!input.is_open()) { - throw std::runtime_error("Cannot open IVF index file"); - } - - IVF loaded; - input.exceptions(std::ios::failbit | std::ios::badbit); - input.seekg(0, std::ios::end); - const auto file_bytes = static_cast(input.tellg()); - input.seekg(0); - uint64_t magic = 0; - input.read(reinterpret_cast(&magic), sizeof(magic)); - const bool explicit_padding = magic == kFormatMagic; - if (explicit_padding) { - uint32_t version = 0, flags = 0; - uint64_t stored_padded_dim = 0; - input.read(reinterpret_cast(&version), sizeof(version)); - input.read(reinterpret_cast(&flags), sizeof(flags)); - input.read(reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim)); - if (version != 1 && version != kFormatVersion) { - throw std::runtime_error("Unsupported IVF index version"); - } - if (flags > 1 || stored_padded_dim > std::numeric_limits::max()) { - throw std::runtime_error("Invalid IVF index metadata"); - } - loaded.raw_reranking_ = flags == 1; - loaded.padded_dim_ = static_cast(stored_padded_dim); - if (version >= 2) { - uint32_t initializer = 0; - input.read(reinterpret_cast(&initializer), sizeof(initializer)); - if (initializer < static_cast(InitializerType::Flat) || - initializer > static_cast(InitializerType::HNSW)) { - throw std::runtime_error("Invalid IVF initializer type in index file"); - } - loaded.initializer_type_ = static_cast(initializer); - } - } else if (magic == kRawFormatMagic) { - loaded.raw_reranking_ = true; - uint32_t version = 0; - input.read(reinterpret_cast(&version), sizeof(version)); - if (version != kRawFormatVersion) { - throw std::runtime_error("Unsupported raw IVF index version"); - } - } else { - input.seekg(0); - } - - /* Load meta data */ - input.read(reinterpret_cast(&loaded.num_), sizeof(size_t)); - input.read(reinterpret_cast(&loaded.dim_), sizeof(size_t)); - input.read(reinterpret_cast(&loaded.num_cluster_), sizeof(size_t)); - input.read(reinterpret_cast(&loaded.ex_bits_), sizeof(size_t)); - input.read(reinterpret_cast(&loaded.type_), sizeof(loaded.type_)); - input.read(reinterpret_cast(&loaded.metric_type_), sizeof(loaded.metric_type_)); - validate_metric_type(loaded.metric_type_); - // Missing routing metadata means a historical file. Its payload layout must - // follow the original threshold, independently of the current Auto policy. - if (loaded.initializer_type_ == InitializerType::Auto) { - loaded.initializer_type_ = - loaded.num_cluster_ < 20000 ? InitializerType::Flat : InitializerType::HNSW; - } - if (loaded.num_ == 0 || loaded.num_ > buffer::kSearchBufferMaxPointCount || - loaded.dim_ == 0 || loaded.num_cluster_ == 0 || - loaded.num_cluster_ > buffer::kSearchBufferMaxPointCount || loaded.ex_bits_ > 8 || - (loaded.raw_reranking_ && loaded.ex_bits_ != 0) || - loaded.dim_ > (std::numeric_limits::max() / 32) - 64) { - throw std::runtime_error("Invalid IVF index metadata"); - } - if (!explicit_padding) { - // Both historical formats omit padded_dim and always used 64-coordinate blocks. - loaded.padded_dim_ = round_up_to_multiple(loaded.dim_, 64); - } - if (loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || - loaded.padded_dim_ > (std::numeric_limits::max() / 32) - 64) { - throw std::runtime_error("Invalid padded dimension in IVF index file"); - } - // Bound every payload by the actual file before allocating or multiplying sizes. - size_t remaining = file_bytes - static_cast(input.tellg()); - const auto consume = [&remaining](size_t count, size_t bytes) { - if (bytes != 0 && count > remaining / bytes) { - throw std::runtime_error("Invalid or truncated IVF index payload"); - } - remaining -= count * bytes; - }; - consume(loaded.num_cluster_, sizeof(size_t)); - consume(loaded.num_, sizeof(PID)); - consume(loaded.num_, loaded.rerank_vector_bytes()); - if (loaded.type_ == RotatorType::MatrixRotator) { - consume(loaded.dim_, sizeof(float) * loaded.padded_dim_); - } else if (loaded.type_ == RotatorType::FhtKacRotator) { - consume(1, loaded.padded_dim_ / 2); - } else { - throw std::runtime_error("Invalid IVF rotator type"); - } - if (loaded.initializer_type_ != InitializerType::HNSW) { - consume(loaded.num_cluster_, sizeof(float) * loaded.padded_dim_); - } - if (loaded.initializer_type_ == InitializerType::FlatRaBitQ) { - consume(loaded.padded_dim_, sizeof(float)); - consume( - div_round_up(loaded.num_cluster_, fastscan::kBatchSize), - BatchDataMap::data_bytes(loaded.padded_dim_) - ); - } - - /* Load number of vectors of each cluster */ - std::vector cluster_sizes(loaded.num_cluster_, 0); - input.read( - reinterpret_cast(cluster_sizes.data()), - static_cast(sizeof(size_t) * loaded.num_cluster_) - ); - - size_t total = 0; - for (size_t size : cluster_sizes) { - if (size > loaded.num_ - total) { - throw std::runtime_error("Invalid cluster counts in IVF index file"); - } - total += size; - consume( - div_round_up(size, fastscan::kBatchSize), - BatchDataMap::data_bytes(loaded.padded_dim_) - ); - } - if (total != loaded.num_) { - throw std::runtime_error("Invalid cluster counts in IVF index file"); - } - loaded.rotator_.reset( - choose_rotator(loaded.dim_, loaded.type_, loaded.padded_dim_) - ); - - /* Load rotator */ - loaded.rotator_->load(input); - - /* Load data */ - loaded.allocate_memory(cluster_sizes); - loaded.initer_->load(input, filename); - input.read( - loaded.batch_data(), - static_cast(loaded.batch_data_bytes(cluster_sizes)) - ); - if (loaded.ex_data_bytes() != 0) { - input.read(loaded.ex_data(), static_cast(loaded.ex_data_bytes())); - } - input.read( - reinterpret_cast(loaded.ids()), - static_cast(loaded.ids_bytes()) - ); - for (size_t i = 0; i < loaded.num_; ++i) { - if (loaded.ids()[i] >= loaded.num_ || - loaded.ids()[i] >= buffer::kSearchBufferCheckedMask) { - throw std::runtime_error("Invalid point ID in IVF index file"); - } - } - - /* Init each cluster */ - loaded.init_clusters(cluster_sizes); - loaded.ready_ = true; - - input.close(); - swap(loaded); -} - -inline void IVF::search( - const float* query, size_t k, size_t nprobe, PID* results, float* dists -) const { - search(query, k, nprobe, results, dists, !raw_reranking_ && ex_bits_ > 2); -} - -inline void IVF::search( - const float* __restrict__ query, - size_t k, - size_t nprobe, - PID* __restrict__ results, - bool use_hacc -) const { - this->search(query, k, nprobe, results, nullptr, use_hacc); -} - -inline void IVF::search( - const float* __restrict__ query, - size_t k, - size_t nprobe, - PID* __restrict__ results, - float* __restrict__ dists, - bool use_hacc -) const { - validate_search(query, k, nprobe, results); - - nprobe = std::min(nprobe, num_cluster_); // corner case - std::vector rotated_query(padded_dim_); - this->rotator_->rotate(query, rotated_query.data()); - - // use initer to get closest nprobe centroids - std::vector> centroid_dist(nprobe); - this->initer_->centroids_distances(rotated_query.data(), nprobe, centroid_dist); - - buffer::SearchBuffer knns(k); - - SplitBatchQuery q_obj( - rotated_query.data(), padded_dim_, ex_bits_, metric_type_, use_hacc - ); - - for (size_t i = 0; i < nprobe; ++i) { - PID cid = centroid_dist[i].id; - float dist = centroid_dist[i].distance; - const Cluster& cur_cluster = cluster_lst_[cid]; - - if (metric_type_ == METRIC_L2) { - q_obj.set_g_add(dist); - } else if (metric_type_ == METRIC_IP) { - const float residual_norm = std::sqrt( - euclidean_sqr(rotated_query.data(), initer_->centroid(cid), padded_dim_) - ); - auto g_add_ip = dot_product( - rotated_query.data(), initer_->centroid(cid), padded_dim_ - ); - q_obj.set_g_add(residual_norm, g_add_ip); - } else { - // unsupported - throw std::invalid_argument("Quantization only supports L2 and IP metrics"); - } - // q_obj.set_g_add(dist); - search_cluster(cur_cluster, q_obj, knns, use_hacc, query); - } - - const size_t found = knns.size(); - if (dists != nullptr) { - knns.copy_results(results, dists); - std::fill(dists + found, dists + k, std::numeric_limits::infinity()); - } else { - knns.copy_results(results); - } - std::fill(results + found, results + k, kPidMax); -} - -inline void IVF::search_batch( - const float* queries, - size_t num_queries, - size_t k, - size_t nprobe, - PID* results, - float* dists, - std::optional use_hacc, - size_t num_threads -) const { - if (num_queries == 0) { - return; - } - validate_search(queries, k, nprobe, results); - const size_t max_values = std::numeric_limits::max() / sizeof(float); - if (num_queries > max_values / dim_ || num_queries > max_values / k) { - throw std::length_error("IVF query batch is too large"); - } - nprobe = std::min(nprobe, num_cluster_); - const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); - const size_t workers = num_queries == 1 || num_threads == 1 - ? 1 - : std::min(resolve_num_threads(num_threads), num_queries); - if (workers == 1) { - for (size_t i = 0; i < num_queries; ++i) { - search( - queries + i * dim_, - k, - nprobe, - results + i * k, - dists == nullptr ? nullptr : dists + i * k, - hacc - ); - } - return; - } - std::atomic failed{false}; - std::exception_ptr error; - const auto capture_error = [&] { -#pragma omp critical(rabitq_ivf_search_error) - { - if (!error) { - error = std::current_exception(); - } - } - failed.store(true, std::memory_order_relaxed); - }; -#pragma omp parallel for num_threads(workers) schedule(dynamic) - for (std::ptrdiff_t query_index = 0; - query_index < static_cast(num_queries); - ++query_index) { - if (failed.load(std::memory_order_relaxed)) { - continue; - } - const size_t i = static_cast(query_index); - try { - search( - queries + i * dim_, - k, - nprobe, - results + i * k, - dists == nullptr ? nullptr : dists + i * k, - hacc - ); - } catch (...) { capture_error(); } - } - if (error) { - std::rethrow_exception(error); - } -} - -inline void IVF::validate_search( - const float* query, size_t k, size_t nprobe, const PID* results -) const { - if (!ready_ || initer_ == nullptr || rotator_ == nullptr || - cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || - id_storage_.size() != num_) { - throw std::logic_error("IVF index must be constructed or loaded before search"); - } - if (query == nullptr) { - throw std::invalid_argument("IVF search query must not be null"); - } - if (results == nullptr) { - throw std::invalid_argument("IVF search results must not be null"); - } - if (k == 0 || k > num_ || k > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("IVF search k must be between 1 and the point count"); - } - if (nprobe == 0) { - throw std::invalid_argument("IVF search nprobe must be positive"); - } -} - -inline void IVF::search_cluster( - const Cluster& cur_cluster, - const SplitBatchQuery& q_obj, - buffer::SearchBuffer& knns, - bool use_hacc, - const float* query -) const { - size_t iter = cur_cluster.num() / fastscan::kBatchSize; - size_t remain = cur_cluster.num() - (iter * fastscan::kBatchSize); - - const char* batch_data = cur_cluster.batch_data(); - const char* ex_data = cur_cluster.ex_data(); - const PID* ids = cur_cluster.ids(); - - /* Compute distances block by block */ - for (size_t i = 0; i < iter; ++i) { - scan_one_batch( - batch_data, ex_data, ids, q_obj, knns, fastscan::kBatchSize, use_hacc, query - ); - - batch_data += BatchDataMap::data_bytes(padded_dim_); - if (rerank_vector_bytes() > 0) { - ex_data += rerank_vector_bytes() * fastscan::kBatchSize; - } - ids += fastscan::kBatchSize; - } - - if (remain > 0) { - // scan the last block - scan_one_batch(batch_data, ex_data, ids, q_obj, knns, remain, use_hacc, query); - } -} - -inline void IVF::scan_one_batch( - const char* batch_data, - const char* ex_data, - const PID* ids, - const SplitBatchQuery& q_obj, - buffer::SearchBuffer& knns, - size_t num_points, - bool use_hacc, - const float* query -) const { - std::array est_distance; // estimated distance - std::array low_distance; // lower distance - std::array ip_x0_qr; // inner product of the 1st bit - - split_batch_estdist( - batch_data, - q_obj, - padded_dim_, - est_distance.data(), - low_distance.data(), - ip_x0_qr.data(), - use_hacc - ); - - float distk = knns.top_dist(); - - // Without reranking data, return the one-bit estimates directly. - if (ex_bits_ == 0 && !raw_reranking_) { - detail::insert_candidates(knns, ids, est_distance.data(), num_points); - return; - } - - // incremental distance computation - V2 - for (size_t i = 0; i < num_points; ++i) { - float lower_dist = low_distance[i]; - if (lower_dist < distk) { - PID id = ids[i]; - float ex_dist; - if (raw_reranking_) { - const auto* vector = reinterpret_cast(ex_data); - ex_dist = metric_type_ == METRIC_L2 ? euclidean_sqr(query, vector, dim_) - : dot_product_dis(query, vector, dim_); - } else { - ex_dist = split_distance_boosting( - ex_data, ip_func_, q_obj, padded_dim_, ex_bits_, ip_x0_qr[i] - ); - } - knns.insert(id, ex_dist); - distk = knns.top_dist(); - } - ex_data += rerank_vector_bytes(); - } -} } // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/ivf/types.hpp b/include/rabitqlib/index/ivf/types.hpp new file mode 100644 index 00000000..f487e693 --- /dev/null +++ b/include/rabitqlib/index/ivf/types.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include + +namespace rabitqlib::ivf { +enum class InitializerType : uint32_t { Auto = 0, Flat = 1, FlatRaBitQ = 2, HNSW = 3 }; +} // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/lut.hpp b/include/rabitqlib/index/lut.hpp deleted file mode 100644 index a37559d5..00000000 --- a/include/rabitqlib/index/lut.hpp +++ /dev/null @@ -1,107 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include - -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/highacc_fastscan.hpp" -#include "rabitqlib/utils/space.hpp" - -namespace rabitqlib { - -template -class Lut { - static constexpr size_t kNumBits = 8; - static constexpr size_t kNumBitsHacc = 16; - static_assert(std::is_floating_point_v, "T must be a floating-point type in Lut"); - - private: - [[nodiscard]] static size_t checked_table_length(size_t padded_dim) { - if (padded_dim == 0 || padded_dim % 16 != 0) { - throw std::invalid_argument( - "FastScan dimension must be a positive multiple of 16" - ); - } - if (padded_dim > std::numeric_limits::max() / 4) { - throw std::length_error("FastScan lookup table is too large"); - } - return padded_dim * 4; - } - - [[nodiscard]] static size_t checked_storage_length(size_t padded_dim, bool use_hacc) { - const size_t table_length = checked_table_length(padded_dim); - const size_t copies = use_hacc ? 2 : 1; - if (table_length > std::numeric_limits::max() / copies) { - throw std::length_error("FastScan lookup table is too large"); - } - return table_length * copies; - } - - size_t table_length_ = 0; - std::vector lut_; - T delta_ = 0; - T sum_vl_lut_ = 0; - - public: - explicit Lut() = default; - explicit Lut(const T* rotated_query, size_t padded_dim, bool use_hacc = false) - : table_length_(checked_table_length(padded_dim)) - , lut_(checked_storage_length(padded_dim, use_hacc)) { - std::vector lut_float(table_length_); - std::vector lut_u16(use_hacc ? table_length_ : 0); - reset(rotated_query, padded_dim, lut_float.data(), lut_u16.data(), use_hacc); - } - - /** - * Reinitialize the lookup table while reusing its allocation. - * lut_float_scratch must contain space for padded_dim * 4 floats. - * lut_u16_scratch is additionally required when use_hacc is true. - */ - void reset( - const T* rotated_query, - size_t padded_dim, - float* lut_float_scratch, - uint16_t* lut_u16_scratch = nullptr, - bool use_hacc = false - ) { - const size_t required_size = checked_storage_length(padded_dim, use_hacc); - if (rotated_query == nullptr || lut_float_scratch == nullptr) { - throw std::invalid_argument("Lut::reset requires query and float scratch"); - } - if (use_hacc && lut_u16_scratch == nullptr) { - throw std::invalid_argument( - "Lut::reset requires uint16 scratch for high-accuracy lookup tables" - ); - } - table_length_ = checked_table_length(padded_dim); - lut_.resize(required_size); - - fastscan::pack_lut(padded_dim, rotated_query, lut_float_scratch); - T vl_lut; - T vr_lut; - data_range(lut_float_scratch, table_length_, vl_lut, vr_lut); - - if (use_hacc) { - delta_ = (vr_lut - vl_lut) / ((1 << kNumBitsHacc) - 1); - scalar_quantize( - lut_u16_scratch, lut_float_scratch, table_length_, vl_lut, delta_ - ); - fastscan::transfer_lut_hacc(lut_u16_scratch, padded_dim, lut_.data()); - } else { - delta_ = (vr_lut - vl_lut) / ((1 << kNumBits) - 1); - scalar_quantize(lut_.data(), lut_float_scratch, table_length_, vl_lut, delta_); - } - - const size_t num_tables = table_length_ / 16; - sum_vl_lut_ = vl_lut * static_cast(num_tables); - } - - [[nodiscard]] const uint8_t* lut() const { return lut_.data(); }; - [[nodiscard]] T delta() const { return delta_; }; - [[nodiscard]] T sum_vl() const { return sum_vl_lut_; }; -}; -} // namespace rabitqlib diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp deleted file mode 100644 index 51ddf47b..00000000 --- a/include/rabitqlib/index/query.hpp +++ /dev/null @@ -1,232 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/lut.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/space.hpp" - -namespace rabitqlib { -/** - * @brief use an object to store data used for searching on symphonyqg for a given query - */ -template -class BatchQuery { - private: - Lut lookup_table_; - T G_add_ = 0; - T G_k1xSumq_ = 0; // G_k1xSumq - MetricType metric_type_ = METRIC_L2; - - public: - explicit BatchQuery() = default; - - explicit BatchQuery( - const T* rotated_query, size_t padded_dim, MetricType metric_type = METRIC_L2 - ) - : lookup_table_(rotated_query, padded_dim), metric_type_(metric_type) { - float c_1 = -((1 << 1) - 1) / 2.F; - - T sumq = - std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); - - G_k1xSumq_ = sumq * c_1; - } - - [[nodiscard]] T delta() const { return lookup_table_.delta(); } - - [[nodiscard]] T sum_vl_lut() const { return lookup_table_.sum_vl(); } - - [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } - - [[nodiscard]] T g_add() const { return G_add_; } - - void set_g_add(T dist) { - // dist is squared L2 or 1 - dot_product; IP encoder factors already include 1. - G_add_ = metric_type_ == METRIC_IP ? dist - T{1} : dist; - } - - /** - * Reinitialize the query while reusing the lookup-table allocation. - * An optional correction must come from this exact rotated query. - */ - void reset( - const T* rotated_query, - size_t padded_dim, - float* lut_float_scratch, - MetricType metric_type = METRIC_L2, - const T* k1xsumq = nullptr - ) { - metric_type_ = metric_type; - lookup_table_.reset(rotated_query, padded_dim, lut_float_scratch); - - if (k1xsumq != nullptr) { - G_k1xSumq_ = *k1xsumq; - } else { - constexpr float c_1 = -((1 << 1) - 1) / 2.F; - const T sumq = std::accumulate( - rotated_query, rotated_query + padded_dim, static_cast(0) - ); - G_k1xSumq_ = sumq * c_1; - } - G_add_ = 0; - } - - [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } -}; - -template -class SplitBatchQuery { - private: - const T* rotated_query_; - Lut lookup_table_; - T G_add_ = 0; - T G_error_ = 0; - T G_k1xSumq_ = 0; - T G_kbxSumq_ = 0; - MetricType metric_type_ = METRIC_L2; - - public: - explicit SplitBatchQuery( - const T* rotated_query, - size_t padded_dim, - size_t ex_bits, - MetricType metric_type = METRIC_L2, - bool use_hacc = true - ) - : rotated_query_(rotated_query) - , lookup_table_(rotated_query, padded_dim, use_hacc) - , metric_type_((metric_type == METRIC_IP) ? METRIC_IP : METRIC_L2) { - float c_1 = -static_cast((1 << 1) - 1) / 2.F; - float c_b = -static_cast((1 << (ex_bits + 1)) - 1) / 2.F; - T sumq = - std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); - - G_k1xSumq_ = sumq * c_1; - G_kbxSumq_ = sumq * c_b; - } - [[nodiscard]] const T* rotated_query() const { return rotated_query_; } - - [[nodiscard]] T delta() const { return lookup_table_.delta(); } - - [[nodiscard]] T sum_vl_lut() const { return lookup_table_.sum_vl(); } - - [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } - - [[nodiscard]] T kbxsumq() const { return G_kbxSumq_; } - - [[nodiscard]] T g_add() const { return G_add_; } - - [[nodiscard]] T g_error() const { return G_error_; } - - void set_g_add(T norm, T ip = 0) { - if (metric_type_ == METRIC_L2) { - G_add_ = norm * norm; - G_error_ = norm; - } else if (metric_type_ == METRIC_IP) { - G_add_ = -ip; - G_error_ = norm; - } - } - - [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } -}; - -template -class SplitSingleQuery { - private: - const T* rotated_query_; - std::vector QueryBin_; - T G_add_; - T G_k1xSumq_; - T G_kbaseSumq_; - T G_kbxSumq_; - T G_error_; - T delta_; - T vl_; - MetricType metric_type_ = METRIC_L2; - - public: - static constexpr size_t kNumBits = 4; - explicit SplitSingleQuery( - const T* rotated_query, - size_t padded_dim, - size_t ex_bits, - quant::RabitqConfig config, - size_t metric_type = METRIC_L2, - size_t base_bits = 1 - ) - : rotated_query_(rotated_query) - , QueryBin_(div_round_up(padded_dim, 64) * kNumBits, 0) { - if (base_bits < 1 || base_bits > 8 || ex_bits > 8 || base_bits + ex_bits > 9) { - throw std::invalid_argument( - "SplitSingleQuery requires base_bits in [1, 8], ex_bits in [0, 8], " - "and base_bits + ex_bits <= 9" - ); - } - - float c_1 = -static_cast((1 << 1) - 1) / 2.F; - float c_base = -static_cast((1U << base_bits) - 1) / 2.F; - float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.F; - T sumq = - std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); - - G_k1xSumq_ = sumq * c_1; - G_kbaseSumq_ = sumq * c_base; - G_kbxSumq_ = sumq * c_b; - - metric_type_ = (metric_type == METRIC_IP) ? METRIC_IP : METRIC_L2; - - std::vector quant_query(padded_dim); - - // quantize query by rabitq - quant::quantize_scalar( - rotated_query, padded_dim, kNumBits, quant_query.data(), delta_, vl_, config - ); - - // represent quantized query as u64 - rabitqlib::new_transpose_bin_512( - quant_query.data(), QueryBin_.data(), padded_dim, kNumBits - ); - } - - [[nodiscard]] size_t num_bits() const { return kNumBits; } - - [[nodiscard]] const uint64_t* query_bin() const { return QueryBin_.data(); } - - [[nodiscard]] const T* rotated_query() const { return rotated_query_; } - - [[nodiscard]] T delta() const { return delta_; } - - [[nodiscard]] T vl() const { return vl_; } - - [[nodiscard]] T k1xsumq() const { return G_k1xSumq_; } - - // Equals k1xsumq() when base_bits == 1. - [[nodiscard]] T kbase_sumq() const { return G_kbaseSumq_; } - - [[nodiscard]] T kbxsumq() const { return G_kbxSumq_; } - - [[nodiscard]] T g_add() const { return G_add_; } - - [[nodiscard]] T g_error() const { return G_error_; } - - void set_g_add(T norm, T ip = 0) { - if (metric_type_ == METRIC_L2) { - G_add_ = norm * norm; - G_error_ = norm; - } else if (metric_type_ == METRIC_IP) { - G_add_ = -ip; - G_error_ = norm; - } - } - - void set_g_error(T norm) { G_error_ = norm; } -}; - -} // namespace rabitqlib diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp index 0e519d00..859b9dec 100644 --- a/include/rabitqlib/index/symqg/detail/pipnn.hpp +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -1,26 +1,11 @@ #pragma once -#include - -#include -#include #include -#include -#include #include -#include -#include -#include -#include -#include -#include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/simd/matrix_dispatch.hpp" -#include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/tools.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::symqg::detail { @@ -30,463 +15,13 @@ struct InitialGraph { std::vector neighbors; }; -namespace pipnn_impl { -constexpr size_t kLeafSize = 1024; -constexpr size_t kTileSize = 2048; -constexpr size_t kHashBits = 12; -constexpr size_t kLocalDegree = 2; -constexpr float kAlpha = 1.1F; -using Bucket = std::vector; - -struct Candidate { - PID id; - float distance; - uint16_t hash; -}; - -struct BucketJob { - Bucket ids; - size_t depth = 0; - uint32_t seed = 555; -}; - -// Non-owning slices of a single build allocation. Keeping the large matrices out -// of worker heaps avoids retaining their high-water marks after seeding finishes. -struct Workspace { - float* points; - float* distances; - float* norms; - std::vector> candidates; -}; - -struct ScratchPool { - size_t point_size; - size_t distance_size; - size_t stride; - std::unique_ptr storage; - - ScratchPool(size_t count, size_t dim, size_t threads) - : point_size(std::min(count, kTileSize) * dim) - , distance_size(std::max( - std::min(count, kTileSize) * std::min(count, size_t{512}), - std::min(count, kLeafSize) * std::min(count, kLeafSize) - )) - , stride(point_size + distance_size + std::min(count, kTileSize)) - // Workers initialize their own slices before reading, preserving NUMA - // first-touch placement without a serial zero-fill of the entire pool. - , storage(new float[threads * stride]) {} - - Workspace worker(size_t id) { - float* base = storage.get() + id * stride; - return {base, base + point_size, base + point_size + distance_size, {}}; - } -}; - -inline void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { - for (size_t i = 0; i < ids.size(); ++i) { - std::copy_n(data + ids[i] * dim, dim, work.points + i * dim); - } -} - -inline void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) { - simd::pairwise_distances_lower( - work.points, work.distances, work.norms, size, dim, metric == METRIC_IP - ); -} - -inline std::vector partition( - const float* data, - size_t dim, - const BucketJob& job, - MetricType metric, - size_t threads, - ScratchPool& scratch -) { - const auto& ids = job.ids; - const size_t leader_count = std::min( - ids.size(), std::clamp(job.depth == 0 ? 512 : ids.size() / 200, 3, 512) - ); - const size_t fanout = std::min( - leader_count, - job.depth == 0 ? size_t{10} - : job.depth == 1 ? size_t{3} - : size_t{1} - ); - std::mt19937 random(job.seed); - Bucket leaders; - leaders.reserve(leader_count); - std::sample(ids.begin(), ids.end(), std::back_inserter(leaders), leader_count, random); - RowMajorMatrix leader_data(leader_count, dim); - for (size_t i = 0; i < leader_count; ++i) { - std::copy_n(data + leaders[i] * dim, dim, leader_data.data() + i * dim); - } - Vector leader_norms(leader_count); - simd::row_norms(leader_data.data(), leader_norms.data(), leader_count, dim); - std::vector assignments(ids.size() * fanout); - const auto assign_tile = [&](size_t begin, Workspace& work) { - const size_t count = std::min(kTileSize, ids.size() - begin); - RowMajorMatrixMap points( - work.points, static_cast(count), static_cast(dim) - ); - RowMajorMatrixMap distances( - work.distances, - static_cast(count), - static_cast(leader_count) - ); - VectorMap norms(work.norms, static_cast(count)); - for (size_t i = 0; i < count; ++i) { - std::copy_n(data + ids[begin + i] * dim, dim, points.data() + i * dim); - } - simd::matrix_product_transposed( - points.data(), leader_data.data(), distances.data(), count, dim, leader_count - ); - simd::row_norms(points.data(), norms.data(), count, dim); - work.candidates.resize(leader_count); - for (size_t i = 0; i < count; ++i) { - for (size_t j = 0; j < leader_count; ++j) { - const float dot = - distances(static_cast(i), static_cast(j)); - const float distance = - metric == METRIC_L2 - ? std::max(0.0F, norms.data()[i] + leader_norms.data()[j] - 2 * dot) - : -dot; - work.candidates[j] = {distance, static_cast(j)}; - } - std::partial_sort( - work.candidates.begin(), - work.candidates.begin() + static_cast(fanout), - work.candidates.end() - ); - for (size_t j = 0; j < fanout; ++j) { - assignments[(begin + i) * fanout + j] = work.candidates[j].second; - } - } - }; - if (job.depth == 0) { -#pragma omp parallel num_threads(threads) - { - auto work = scratch.worker(static_cast(omp_get_thread_num())); -#pragma omp for schedule(static) - for (std::ptrdiff_t offset = 0; - offset < static_cast(ids.size()); - offset += static_cast(kTileSize)) { - const size_t begin = static_cast(offset); - assign_tile(begin, work); - } - } - } else { - // Already inside the bucket-level team. Eigen must see that team rather - // than a serialized nested region when deciding its own parallelism. - auto work = scratch.worker(static_cast(omp_get_thread_num())); - for (size_t begin = 0; begin < ids.size(); begin += kTileSize) { - assign_tile(begin, work); - } - } - std::vector children(leader_count); - for (size_t i = 0; i < ids.size(); ++i) { - for (size_t j = 0; j < fanout; ++j) { - children[assignments[i * fanout + j]].push_back(ids[i]); - } - } - return children; -} - -inline std::vector cluster( - const float* data, - size_t count, - size_t dim, - MetricType metric, - size_t threads, - ScratchPool& scratch -) { - Bucket all(count); - std::iota(all.begin(), all.end(), PID{0}); - // The root parallelizes tiles; deeper levels parallelize independent buckets. - std::vector active; - active.push_back({std::move(all), 0, 555}); - std::vector leaves; - while (!active.empty()) { - std::vector> next(active.size()); - std::vector> finished(active.size()); - const auto process_bucket = [&](size_t i) { - auto& job = active[i]; - if (job.ids.size() <= kLeafSize) { - finished[i].push_back(std::move(job.ids)); - return; - } - auto children = partition(data, dim, job, metric, threads, scratch); - std::sort(children.begin(), children.end(), [](const auto& a, const auto& b) { - return a.size() < b.size(); - }); - Bucket small; - for (size_t j = 0; j < children.size(); ++j) { - auto& child = children[j]; - if (child.empty()) { - continue; - } - if (child.size() < kLeafSize / 10) { - small.insert(small.end(), child.begin(), child.end()); - if (small.size() < kLeafSize / 10) { - continue; - } - std::sort(small.begin(), small.end()); - small.erase(std::unique(small.begin(), small.end()), small.end()); - finished[i].push_back(std::move(small)); - small = Bucket(); - } else if (child.size() <= kLeafSize) { - finished[i].push_back(std::move(child)); - } else if (job.depth >= 4 || (job.depth >= 2 && child.size() * 5 > job.ids.size() * 4)) { - // Duplicates must not cause unbounded recursion/replication. - std::mt19937 random(job.seed + static_cast(j)); - std::shuffle(child.begin(), child.end(), random); - for (size_t begin = 0; begin < child.size(); begin += kLeafSize) { - finished[i].emplace_back( - child.begin() + static_cast(begin), - child.begin() + static_cast( - std::min(begin + kLeafSize, child.size()) - ) - ); - } - } else { - next[i].push_back( - {std::move(child), - job.depth + 1, - job.seed * 1664525U + static_cast(j) + 1013904223U} - ); - } - } - if (!small.empty()) { - std::sort(small.begin(), small.end()); - small.erase(std::unique(small.begin(), small.end()), small.end()); - finished[i].push_back(std::move(small)); - } - Bucket().swap(job.ids); - }; - if (active.front().depth == 0) { - // A serialized outer OpenMP region would make the root's tile team - // nested, forcing the runtime to create workers again on every reset. - process_bucket(0); - } else { -#pragma omp parallel for num_threads(threads) schedule(dynamic) - for (std::ptrdiff_t i = 0; i < static_cast(active.size()); - ++i) { - process_bucket(static_cast(i)); - } - } - std::vector new_active; - for (size_t i = 0; i < active.size(); ++i) { - for (auto& leaf : finished[i]) { - leaves.push_back(std::move(leaf)); - } - for (auto& job : next[i]) { - new_active.push_back(std::move(job)); - } - } - active = std::move(new_active); - } - return leaves; -} -} // namespace pipnn_impl - -// Native PiPNN-style seeding: overlapping leader partitions, local dense kNN, -// directional hash reservoirs, and alpha pruning. Uses existing Eigen/OpenMP. -inline InitialGraph build_initial_graph( +InitialGraph build_initial_graph( const float* data, size_t count, size_t dim, size_t degree, MetricType metric = METRIC_L2, size_t num_threads = std::numeric_limits::max() -) { - validate_metric_type(metric); - if (data == nullptr || dim == 0 || degree == 0 || degree >= count || - count >= buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { - throw std::invalid_argument( - "PiPNN seed requires data, positive dim, and a degree multiple of 32 below " - "count" - ); - } - using namespace pipnn_impl; - const size_t threads = resolve_num_threads(num_threads); - ScratchPool scratch(count, dim, threads); - auto leaves = cluster(data, count, dim, metric, threads, scratch); - RowMajorMatrix projections(dim, kHashBits); - std::mt19937 random(555); - std::normal_distribution normal; - for (Eigen::Index i = 0; i < projections.size(); ++i) { - projections.data()[i] = normal(random); - } - RowMajorMatrix sketches(count, kHashBits); -#pragma omp parallel for num_threads(threads) schedule(static) - for (std::ptrdiff_t offset = 0; offset < static_cast(count); - offset += static_cast(kTileSize)) { - const size_t begin = static_cast(offset); - const size_t rows = std::min(kTileSize, count - begin); - simd::matrix_product( - data + begin * dim, - projections.data(), - sketches.data() + begin * kHashBits, - rows, - dim, - kHashBits - ); - } - const size_t capacity = degree * 5 / 2; - std::vector table(count * capacity); - std::vector sizes(count, 0), worst(count, 0); - std::vector locks(count); - const auto merge = [&](PID source, PID target, float distance) { - uint16_t hash = 0; - for (size_t bit = 0; bit < kHashBits; ++bit) { - hash = static_cast( - (hash << 1) | (sketches(target, static_cast(bit)) > - sketches(source, static_cast(bit))) - ); - } - std::lock_guard lock(locks[source]); - auto* row = table.data() + source * capacity; - auto& size = sizes[source]; - auto& furthest = worst[source]; - if (size == capacity && distance >= row[furthest].distance) { - return; - } - auto* pos = std::lower_bound( - row, - row + size, - hash, - [](const Candidate& candidate, uint16_t key) { return candidate.hash < key; } - ); - const size_t slot = static_cast(pos - row); - if (slot < size && pos->hash == hash) { - if (distance >= pos->distance) { - return; - } - *pos = {target, distance, hash}; - } else if (size < capacity) { - std::move_backward(pos, row + size, row + size + 1); - *pos = {target, distance, hash}; - ++size; - } else { - const size_t insertion = slot > furthest ? slot - 1 : slot; - if (slot > furthest) { - std::move(row + furthest + 1, row + slot, row + furthest); - } else { - std::move_backward(row + slot, row + furthest, row + furthest + 1); - } - row[insertion] = {target, distance, hash}; - } - // Bounded rows keep this scan small; retain full-precision scores. - furthest = static_cast( - std::max_element( - row, - row + size, - [](const Candidate& a, const Candidate& b) { - return a.distance < b.distance; - } - ) - - row - ); - }; -#pragma omp parallel num_threads(threads) - { - auto work = scratch.worker(static_cast(omp_get_thread_num())); -#pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(leaves.size()); - ++index) { - const size_t leaf = static_cast(index); - const auto& ids = leaves[leaf]; - if (ids.size() < 2) { - continue; - } - gather(data, dim, ids, work); - pairwise(work, ids.size(), dim, metric); - for (size_t i = 0; i < ids.size(); ++i) { - work.candidates.clear(); - for (size_t j = 0; j < ids.size(); ++j) { - if (i != j) { - work.candidates.emplace_back( - work.distances[std::max(i, j) * ids.size() + std::min(i, j)], - ids[j] - ); - } - } - const size_t keep = std::min(kLocalDegree, work.candidates.size()); - std::partial_sort( - work.candidates.begin(), - work.candidates.begin() + static_cast(keep), - work.candidates.end() - ); - for (size_t j = 0; j < keep; ++j) { - const auto [distance, target] = work.candidates[j]; - merge(ids[i], target, distance); - merge(target, ids[i], distance); - } - } - } - } - std::vector().swap(leaves); - scratch.storage.reset(); - sketches.resize(0, 0); - InitialGraph result; - result.offsets.resize(count + 1, 0); - // Prune in-place: retain selected IDs in the front of each reservoir row. -#pragma omp parallel num_threads(threads) - { - std::vector candidates; - Bucket selected; -#pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(count); - ++index) { - const size_t i = static_cast(index); - const auto* row = table.data() + i * capacity; - candidates.assign(row, row + sizes[i]); - std::sort( - candidates.begin(), - candidates.end(), - [](const Candidate& a, const Candidate& b) { - return a.distance < b.distance || - (a.distance == b.distance && a.id < b.id); - } - ); - selected.clear(); - for (const auto& candidate : candidates) { - bool occluded = false; - for (PID accepted : selected) { - const float distance = - metric == METRIC_L2 - ? euclidean_sqr( - data + accepted * dim, data + candidate.id * dim, dim - ) - : dot_product_dis( - data + accepted * dim, data + candidate.id * dim, dim - ) - 1.0F; - if (distance <= candidate.distance / kAlpha) { - occluded = true; - break; - } - } - if (!occluded) { - selected.push_back(candidate.id); - } - if (selected.size() == degree) { - break; - } - } - for (size_t j = 0; j < selected.size(); ++j) { - table[i * capacity + j].id = selected[j]; - } - result.offsets[i + 1] = selected.size(); - } - } - std::partial_sum(result.offsets.begin(), result.offsets.end(), result.offsets.begin()); - result.neighbors.resize(result.offsets.back()); -#pragma omp parallel for num_threads(threads) - for (std::ptrdiff_t index = 0; index < static_cast(count); ++index) { - const size_t i = static_cast(index); - for (size_t j = 0; j < result.offsets[i + 1] - result.offsets[i]; ++j) { - result.neighbors[result.offsets[i] + j] = table[i * capacity + j].id; - } - } - return result; -} +); + } // namespace rabitqlib::symqg::detail diff --git a/include/rabitqlib/index/symqg/detail/search_scratch.hpp b/include/rabitqlib/index/symqg/detail/search_scratch.hpp new file mode 100644 index 00000000..a6fd2785 --- /dev/null +++ b/include/rabitqlib/index/symqg/detail/search_scratch.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include +#include + +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" + +namespace rabitqlib::symqg::detail { + +// Call-only view; storage and visited-set lifetimes are chosen by the caller. +struct SearchScratchView { + float* rotated_query; + float* estimated_distances; + float* lookup_table; + BatchQuery& batch_query; + buffer::SearchBuffer& search_pool; + buffer::SearchBuffer& result_pool; + VisitedSet& visited; +}; + +// Common owned buffers. Visited storage remains separate so single-query and +// clustering callers can retain their worker-local cache across calls. +struct SearchScratch { + std::vector rotated_query; + std::vector estimated_distances; + std::vector lookup_table; + BatchQuery batch_query; + buffer::SearchBuffer search_pool; + buffer::SearchBuffer result_pool; + + SearchScratch(size_t padded_dim, size_t degree, size_t ef, size_t k) + : rotated_query(padded_dim) + , estimated_distances(degree) + , lookup_table(padded_dim * 4) + , search_pool(ef) + , result_pool(k) {} + + SearchScratchView view(VisitedSet& visited) { + return { + rotated_query.data(), + estimated_distances.data(), + lookup_table.data(), + batch_query, + search_pool, + result_pool, + visited}; + } +}; + +} // namespace rabitqlib::symqg::detail diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index cc7460c1..2e99abbf 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -3,18 +3,26 @@ #include #include #include -#include #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/core/metric.hpp" // IWYU pragma: export +#include "rabitqlib/core/types.hpp" // IWYU pragma: export +#include "rabitqlib/index/detail/visited_set_fwd.hpp" // IWYU pragma: export +#include "rabitqlib/index/types.hpp" // IWYU pragma: export #include "rabitqlib/utils/memory.hpp" -#include "rabitqlib/utils/rotator.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/utils/packed.hpp" + +namespace rabitqlib { +template +class Rotator; +template +class BatchQuery; +namespace buffer { +template +class SearchBuffer; +} +} // namespace rabitqlib namespace rabitqlib::symqg { @@ -40,24 +48,43 @@ class QuantizedQuery { template class QuantizedGraph; +template <> +class QuantizedGraph; + +namespace detail { +struct SearchScratchView; + +void search_with_workspace( + QuantizedGraph& graph, + const float* query, + uint32_t knn, + uint32_t* results, + float* dists, + SearchScratchView scratch, + PID hint = kPidMax, + const float* k1xsumq = nullptr +); +} // namespace detail + template <> class QuantizedGraph { friend class QGBuilder; + friend void detail:: + search_with_workspace(QuantizedGraph&, const float*, uint32_t, uint32_t*, float*, detail::SearchScratchView, PID, const float*); friend struct QGConstructionTestAccess; private: - size_t num_points_ = 0; // num points - size_t degree_bound_ = 0; // degree bound - size_t dim_ = 0; // dimension - size_t padded_dim_ = 0; // padded dimension - // Raw-vector distance. + size_t num_points_ = 0; + size_t degree_bound_ = 0; + size_t dim_ = 0; + size_t padded_dim_ = 0; float (*raw_dist_func_)(const float*, const float*, size_t) = nullptr; - PID entry_point_ = 0; // Entry point of graph + PID entry_point_ = 0; MetricType metric_type_ = MetricType::METRIC_L2; RotatorType rotator_type_ = RotatorType::FhtKacRotator; size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors std::vector centroid_; // rotated global centroid for qg-quant - ex_ipfunc quantized_ip_func_ = nullptr; + float (*quantized_ip_func_)(const float*, const uint8_t*, size_t) = nullptr; using RowStorage = std::vector>; @@ -68,24 +95,18 @@ class QuantizedGraph { // Builders and load() fill rows before use; leave scalars uninitialized so // construction workers first-touch their own pages instead of zeroing serially. RowStorage data_; - std::unique_ptr> rotator_; // data rotator + std::unique_ptr> rotator_; // Position of row data (raw vector or packed qg-quant vector), neighbor // quantization data, and neighbor IDs. Since every degree equals degree_bound_ // (a multiple of 32), the degree does not need to be stored per vertex. - size_t batch_data_offset_ = 0; // offset of qg batch data - size_t neighbor_offset_ = 0; // offset of neighbors - size_t row_offset_ = 0; // length of entire row + size_t batch_data_offset_ = 0; + size_t neighbor_offset_ = 0; + size_t row_offset_ = 0; // length of entire row size_t ef_ = 0; bool ready_ = false; uint32_t seed_ = 42; - // Owned by QGBuilder and never persisted. - struct BuildCache { - std::vector rotated_vectors; - std::vector queries; - }; - [[nodiscard]] static size_t checked_add(size_t lhs, size_t rhs); [[nodiscard]] static size_t checked_multiply(size_t lhs, size_t rhs); @@ -96,7 +117,7 @@ class QuantizedGraph { void initialize_layout(); - void initialize(); + void initialize(bool randomize_rotation = true); void copy_vectors(const float*, size_t); @@ -130,14 +151,7 @@ class QuantizedGraph { void search_impl(const float*, uint32_t, uint32_t*, float*, size_t); void - search_with_scratch_impl(const float*, uint32_t, uint32_t*, float*, float*, float*, float*, BatchQuery&, buffer::SearchBuffer&, buffer::SearchBuffer&, VisitedSet&, PID, const float*); - - const float* prepare_build_query( - PID, - std::vector&, - std::optional&, - const BuildCache* = nullptr - ) const; + search_with_scratch_impl(const float*, uint32_t, uint32_t*, float*, detail::SearchScratchView, PID, const float*); float point_distance(const float*, const QuantizedQuery*, PID) const; @@ -166,26 +180,19 @@ class QuantizedGraph { ); } - void find_candidates( + // Inputs are in this graph's rotated, padded domain. The centroid is the + // source vertex; data contains one vector per neighbor, in ID order. + void encode_neighbors( PID, - size_t, - std::vector>&, - VisitedSet&, - const std::vector&, - const BuildCache* = nullptr - ) const; - - void update_qg( - PID, - const std::vector>&, - const BuildCache* = nullptr, - std::vector* = nullptr + const std::vector&, + const float* rotated_data, + const float* rotated_centroid ); void update_results(buffer::SearchBuffer&, VisitedSet&, const float*, const QuantizedQuery*); - inline void scan_neighbors( + void scan_neighbors( const BatchQuery&, PID, float*, @@ -214,6 +221,13 @@ class QuantizedGraph { QuantizedGraph(QuantizedGraph&&) noexcept; QuantizedGraph& operator=(QuantizedGraph&&) noexcept; + // Capacity, nonremoved points, and stored records, respectively. Until building + // completes, both point counts are zero. This index has no remove operation. + [[nodiscard]] size_t max_elements() const; + [[nodiscard]] size_t num_points() const; + [[nodiscard]] size_t num_stored_points() const; + + // Configured graph size, also used by builders before rows are searchable. [[nodiscard]] size_t num_vertices() const; [[nodiscard]] size_t dimension() const; @@ -236,7 +250,6 @@ class QuantizedGraph { void set_ef(size_t); - /* search and copy results to KNN */ void search( const float* __restrict__ query, uint32_t knn, @@ -277,12 +290,14 @@ class QuantizedGraph { /** * Search using caller-owned scratch buffers. * - * This avoids per-query allocations and can use a previous result as a - * warm-start hint. The search and result buffers must be sized for ef and - * knn respectively. rotated_query_scratch, est_dist_scratch, and - * lut_float_scratch require padded_dim(), degree_bound(), and - * padded_dim() * 4 elements respectively. The visited set must be - * initialized before the first search; an uninitialized set raises + * Reuses scratch and can use a previous result as a warm-start hint. + * BatchQuery may allocate LUT storage on its first use or when its shape changes; + * this is not the stable core v1 zero-allocation interface. Search and result + * buffers must be sized for ef and knn respectively. rotated_query_scratch and + * lut_float_scratch need D and 4 * D floats, where D is the graph's stored padded + * dimension; this advanced caller must already know D. dimension() is the original + * input dimension. est_dist_scratch needs degree_bound() floats. The visited set must + * be initialized before the first search; an uninitialized set raises * std::invalid_argument. Concurrent searches require separate scratch * buffers, including visited sets. * diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index cab3ecda..c0d561db 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -1,27 +1,15 @@ #pragma once -#include - -#include -#include -#include #include #include #include -#include -#include #include #include -#include -#include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/symqg/detail/pipnn.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" #include "rabitqlib/index/symqg/qg.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/tools.hpp" -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::symqg { constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar @@ -33,60 +21,91 @@ class QGBuilder { friend struct QGConstructionTestAccess; private: + // Owned by QGBuilder and never persisted. + struct BuildCache { + std::vector rotated_vectors; + std::vector queries; + }; + + // Reset before scoring. The graph and optional cache must outlive this query. + // Owns reconstruction/rotation scratch only when no cache can supply it. + class BuildQuery { + const QuantizedGraph& graph_; + const BuildCache* cache_; + std::vector storage_; + const float* raw_ = nullptr; + const float* rotated_ = nullptr; + std::optional prepared_; + + public: + explicit BuildQuery( + const QuantizedGraph& graph, const BuildCache* cache = nullptr + ); + // Prepared pointers can borrow storage_; never copy their owner. + BuildQuery(const BuildQuery&) = delete; + BuildQuery& operator=(const BuildQuery&) = delete; + + void reset(PID); + [[nodiscard]] float distance_to(PID) const; + // Raw scoring needs no rotation. Prepare it lazily for FastScan only. + [[nodiscard]] const float* rotated_query(); + }; + QuantizedGraph& qg_; - size_t ef_build_; // size of search pool for indexing - size_t num_threads_; // number of threads used for indexing - size_t num_nodes_; // num of data points - size_t dim_; // dimension of data + size_t ef_build_; // size of search pool for indexing + size_t num_threads_; + size_t num_nodes_; + size_t dim_; size_t degree_bound_; // degree bound for qg, multiple of 32 bool seeded_ = true; - static constexpr size_t kMaxCandidatePoolSize = - 750; // max num of candidates for indexing - static constexpr size_t kMaxPrunedSize = - 300; // max number of recorded pruned candidates + static constexpr size_t kMaxCandidatePoolSize = 750; + static constexpr size_t kMaxPrunedSize = 300; std::vector new_neighbors_; // new neighbors for current iteration - std::vector pruned_neighbors_; // recorded pruned neighbors - std::vector visited_list_; // per-thread visited sets - std::vector degrees_; // record degree of qg - uint32_t seed_; // base seed for graph initialization - std::optional::BuildCache> cache_; - std::vector candidates_; + std::vector pruned_neighbors_; + std::vector degrees_; + uint32_t seed_; // base seed for graph initialization + std::optional cache_; std::vector reverse_buffer_; - std::vector> rotated_neighbors_; - - const QuantizedGraph::BuildCache* build_cache() const { - return cache_ ? &*cache_ : nullptr; - } + struct BuildWorker { + VisitedSet visited; + CandidateList candidates; + std::vector rotated_neighbors; + std::vector neighbor_ids; + + BuildWorker(size_t num_nodes, size_t visited_capacity); + }; + std::vector workers_; + + const BuildCache* build_cache() const; + void find_candidates(PID, CandidateList&, VisitedSet&) const; void random_init(); - void search_new_neighbors(bool refine); - void heuristic_prune(PID, CandidateList&, CandidateList&, bool); - void add_reverse_edges(bool); + void search_new_neighbors(bool record_pruned); + void heuristic_prune(PID, CandidateList&, CandidateList&, bool record_pruned); + void add_reverse_edges(bool record_pruned); void add_pruned_edges( const CandidateList&, const CandidateList&, CandidateList&, float ); - void graph_refine(); - void iter(bool); + void complete_degrees(); + void run_pass(bool final_pass); + void run_passes(size_t); + void write_row(PID, BuildWorker&); + void write_row(PID, const CandidateList&, BuildWorker&); void initialize_storage(const float* data); + void initialize_seed( + const float* data, + const std::vector& offsets, + const std::vector& neighbors + ); + QGBuilder( QuantizedGraph& index, uint32_t ef_build, size_t num_threads, uint32_t seed = std::random_device{}(), bool cache_vectors = false - ) - : qg_{index} - , ef_build_{ef_build} - , num_threads_{resolve_num_threads(num_threads)} - , num_nodes_{qg_.num_vertices()} - , dim_{qg_.dimension()} - , degree_bound_(qg_.degree_bound()) - , seed_(seed) { - if (cache_vectors) { - cache_.emplace(); - } - } + ); public: explicit QGBuilder( @@ -97,14 +116,7 @@ class QGBuilder { QGInitialization init = QGInitialization::PiPNN, uint32_t seed = std::random_device{}(), bool cache_vectors = false - ) - : QGBuilder(index, ef_build, num_threads, seed, cache_vectors) { - if (init != QGInitialization::PiPNN && init != QGInitialization::Random) { - throw std::invalid_argument("Unknown QG initialization"); - } - seeded_ = init == QGInitialization::PiPNN; - reset(data); - } + ); // Preserve the seeded random-initialization overload used before PiPNN. QGBuilder( @@ -114,617 +126,25 @@ class QGBuilder { size_t num_threads, uint32_t seed, bool cache_vectors = false - ) - : QGBuilder( - index, - ef_build, - data, - num_threads, - QGInitialization::Random, - seed, - cache_vectors - ) {} + ); /** Replace vectors and restart the selected initializer, retaining workspace. * Cached vectors and queries are owned; input may be released after reset. */ - void reset(const float* data) { - if (data == nullptr) { - throw std::invalid_argument("QGBuilder data must not be null"); - } - if (seeded_) { - auto initial = detail::build_initial_graph( - data, num_nodes_, dim_, degree_bound_, qg_.metric_type_, num_threads_ - ); - initialize_seed(data, initial.offsets, initial.neighbors); - } else { - initialize_storage(data); - random_init(); - } - } + void reset(const float* data); - private: - void initialize_seed( - const float* data, - const std::vector& offsets, - const std::vector& neighbors - ) { - if (offsets.size() != num_nodes_ + 1 || offsets.front() != 0 || - offsets.back() != neighbors.size()) { - throw std::invalid_argument("Seed graph offsets must delimit every vertex"); - } - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || - offsets[i + 1] - offsets[i] > degree_bound_) { - throw std::invalid_argument( - "Seed graph row exceeds degree bound or has invalid offsets" - ); - } - for (size_t j = offsets[i]; j < offsets[i + 1]; ++j) { - if (neighbors[j] >= num_nodes_ || neighbors[j] == i) { - throw std::invalid_argument( - "Seed graph IDs must be in range and exclude self" - ); - } - } - } - initialize_storage(data); -#pragma omp parallel num_threads(num_threads_) - { - CandidateList row; - row.reserve(degree_bound_); - std::vector ids; - ids.reserve(degree_bound_); -#pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - row.clear(); - ids.assign( - neighbors.begin() + static_cast(offsets[i]), - neighbors.begin() + static_cast(offsets[i + 1]) - ); - std::sort(ids.begin(), ids.end()); - ids.erase(std::unique(ids.begin(), ids.end()), ids.end()); - for (PID id : ids) { - // Encoding only consumes IDs. Score surviving seed candidates - // during search instead of retaining a second graph of scores. - row.emplace_back(id, 0.0F); - } - degrees_[i] = row.size(); - qg_.update_qg( - i, - row, - build_cache(), - &rotated_neighbors_[static_cast(omp_get_thread_num())] - ); - } - } - } - - public: // One complete search/prune/reverse-edge/degree-completion iteration. Call // before serving a seeded graph: search() requires full final rows. - void refine() { - qg_.ready_ = false; - iter(true); - qg_.ready_ = true; - } + void refine(); // One refinement for PiPNN initialization, three passes for random init. - void build() { - if (seeded_) { - refine(); - } else { - build(3); - } - } + void build(); - void build(size_t num_iter) { - if (num_iter < 2) { - throw std::invalid_argument( - "The number of QG build iterations must be at least 2" - ); - } - qg_.ready_ = false; - // for first iterations, we do not need to refine the graph structure - for (size_t i = 0; i < num_iter - 1; ++i) { - iter(false); - } - iter(true); - qg_.ready_ = true; - } + void build(size_t num_iter); - [[nodiscard]] bool check_dup() const { - std::atomic flag(false); -#pragma omp parallel for num_threads(num_threads_) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - std::unordered_set edges; - for (auto nei : new_neighbors_[i]) { - if (edges.find(nei.id) != edges.end()) { - flag = true; - } - edges.emplace(nei.id); - } - } - return flag; - } + [[nodiscard]] bool check_dup() const; - [[nodiscard]] float avg_degree() const { - size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), 0U); - return static_cast(degrees) / static_cast(num_nodes_); - } + [[nodiscard]] float avg_degree() const; }; -inline void QGBuilder::initialize_storage(const float* data) { - if (data == nullptr) { - throw std::invalid_argument("QGBuilder data must not be null"); - } - // Allocate refinement scratch after PiPNN releases its dense workspace. - // Clear old graph-dependent state on reset while retaining vector capacities. - new_neighbors_.resize(num_nodes_); - pruned_neighbors_.resize(num_nodes_); - reverse_buffer_.resize(num_nodes_); - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - new_neighbors_[i].clear(); - pruned_neighbors_[i].clear(); - reverse_buffer_[i].clear(); - } - if (visited_list_.empty()) { - visited_list_ = std::vector( - num_threads_, - VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) - ); - } - degrees_.assign(num_nodes_, degree_bound_); - candidates_.resize(num_threads_); - rotated_neighbors_.resize(num_threads_); - std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); - - qg_.ready_ = false; - qg_.set_quantization_centroid(centroid.data()); - qg_.copy_vectors(data, num_threads_); - - if (cache_) { - const size_t padded_dim = qg_.padded_dim_; - cache_->rotated_vectors.resize(num_nodes_ * padded_dim); -#pragma omp parallel for num_threads(num_threads_) schedule(static) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t id = static_cast(index); - float* row = cache_->rotated_vectors.data() + id * padded_dim; - if (qg_.is_quantized()) { - qg_.reconstruct_quantized_vector(id, row); - } else { - qg_.rotator_->rotate(qg_.get_vector(id), row); - } - } - cache_->queries.clear(); - if (qg_.is_quantized()) { - cache_->queries.reserve(num_nodes_); - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t id = static_cast(index); - cache_->queries.emplace_back( - cache_->rotated_vectors.data() + id * padded_dim, - qg_.centroid_.data(), - padded_dim, - qg_.metric_type_ - ); - } - } - } - - PID entry_point = 0; - if (qg_.is_quantized()) { - QuantizedQuery query( - qg_.centroid_.data(), qg_.centroid_.data(), qg_.padded_dim_, qg_.metric_type_ - ); - float best = std::numeric_limits::max(); - for (PID id = 0; id < num_nodes_; ++id) { - const float distance = qg_.quantized_distance(query, id); - if (distance < best) { - best = distance; - entry_point = id; - } - } - } else { - entry_point = exact_nn( - data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ - ); - } - - qg_.set_ep(entry_point); -} - -inline void QGBuilder::add_pruned_edges( - const CandidateList& result, - const CandidateList& pruned_list, - CandidateList& new_result, - float threshold -) { - size_t start = 0; - new_result.clear(); - new_result = result; - - std::unordered_set nei_set; - nei_set.reserve(degree_bound_); - for (const auto& nei : result) { - nei_set.emplace(nei.id); - } - - std::vector reconstructed; - std::optional prepared; - while (new_result.size() < degree_bound_ && start < pruned_list.size()) { - const auto& cur = pruned_list[start]; - bool occlude = false; - const float* cur_data = - qg_.prepare_build_query(cur.id, reconstructed, prepared, build_cache()); - float dik_sqr = cur.distance; - - if (nei_set.find(cur.id) != nei_set.end()) { - break; - } - - for (auto& nei : new_result) { - float dij_sqr = nei.distance; - if (dij_sqr > dik_sqr) { - break; - } - float djk_sqr = - qg_.point_distance(cur_data, prepared ? &*prepared : nullptr, nei.id); - float cosine = - (dik_sqr + dij_sqr - djk_sqr) / (2 * std::sqrt(dij_sqr * dik_sqr)); - if (cosine > threshold) { - occlude = true; - break; - } - } - - if (!occlude) { - new_result.emplace_back(cur); - nei_set.emplace(cur.id); - std::sort(new_result.begin(), new_result.end()); - } - - ++start; - } -} - -inline void QGBuilder::heuristic_prune( - PID cur_id, CandidateList& pool, CandidateList& pruned_results, bool refine -) { - if (pool.empty()) { - return; - } - pruned_results.clear(); - size_t poolsize = pool.size(); - - // if we dont have enough candidates, just keep all neighbors - if (poolsize <= degree_bound_) { - pruned_results = pool; - return; - } - - std::vector pruned( - poolsize, false - ); // bool vector to record if this neighbor is pruned - size_t start = 0; // start position - - std::vector reconstructed; - std::optional prepared; - while (pruned_results.size() < degree_bound_ && start < poolsize) { - auto candidate_id = pool[start].id; - - // if already pruned, move to next - if (pruned[start]) { - ++start; - continue; - } - - pruned_results.emplace_back(pool[start]); // add current candidate to result - const float* data_j = - qg_.prepare_build_query(candidate_id, reconstructed, prepared, build_cache()); - - // i : current vertex - // j : neighbor added in this iter - // k : remained unpruned candidate neighbor - for (size_t k = start + 1; k < poolsize; ++k) { - if (pruned[k]) { - continue; - } - float dik = pool[k].distance; - auto djk = - qg_.point_distance(data_j, prepared ? &*prepared : nullptr, pool[k].id); - - if (djk < dik) { - if (refine && pruned_neighbors_[cur_id].size() < kMaxPrunedSize) { - pruned_neighbors_[cur_id].emplace_back(pool[k]); - } - pruned[k] = true; - } - } - - ++start; - } -} - -/** - * @brief search for new neighbor in qg - * - * @param refine refine = true means recording pruned candidates - */ -inline void QGBuilder::search_new_neighbors(bool refine) { -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - PID cur_id = i; - auto tid = omp_get_thread_num(); - CandidateList& candidates = candidates_[tid]; - candidates.clear(); - VisitedSet& vis = visited_list_[tid]; - candidates.reserve(2 * kMaxCandidatePoolSize); - vis.clear(); - qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_, build_cache()); - - // Seeded construction keeps its initial edges only in QG. Materialize - // their scores on demand, after the caller can release the input/CSR. - if (new_neighbors_[cur_id].empty() && degrees_[cur_id] != 0) { - std::vector reconstructed; - std::optional prepared; - const float* source = - qg_.prepare_build_query(cur_id, reconstructed, prepared, build_cache()); - const auto ids = qg_.get_neighbors(cur_id); - for (size_t j = 0; j < degrees_[cur_id]; ++j) { - if (ids[j] != cur_id && !vis.get(ids[j])) { - candidates.emplace_back( - ids[j], - qg_.point_distance(source, prepared ? &*prepared : nullptr, ids[j]) - ); - } - } - } - - // Add current neighbors retained by preceding iterations. - for (auto& nei : new_neighbors_[cur_id]) { - auto neighbor_id = nei.id; - if (neighbor_id != cur_id && !vis.get(neighbor_id)) { - candidates.emplace_back(nei); - } - } - - size_t min_size = std::min(candidates.size(), kMaxCandidatePoolSize); - std::partial_sort( - candidates.begin(), - candidates.begin() + static_cast(min_size), - candidates.end() - ); - candidates.resize(min_size); - - // prune and update qg - new_neighbors_[cur_id].reserve(degree_bound_); - heuristic_prune(cur_id, candidates, new_neighbors_[cur_id], refine); - } -} - -inline void QGBuilder::add_reverse_edges(bool refine) { - std::vector locks(num_nodes_); - for (auto& row : reverse_buffer_) { - row.clear(); - } - - // Keep new_neighbors_ read-only while reverse candidates are collected. Mutating a - // destination row here races with another worker reading that row as its source. -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const PID data_id = static_cast(index); - for (const auto& nei : new_neighbors_[data_id]) { - const PID destination = nei.id; - const CandidateList& destination_neighbors = new_neighbors_[destination]; - const bool reciprocal = std::any_of( - destination_neighbors.begin(), - destination_neighbors.end(), - [&](const auto& destination_neighbor) { - return destination_neighbor.id == data_id; - } - ); - if (reciprocal) { - continue; - } - - std::lock_guard lock(locks[destination]); - const size_t missing_slots = - degree_bound_ - std::min(degree_bound_, destination_neighbors.size()); - if (reverse_buffer_[destination].size() < - kMaxCandidatePoolSize + missing_slots) { - reverse_buffer_[destination].emplace_back(data_id, nei.distance); - } - } - } - -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const PID data_id = static_cast(index); - CandidateList& tmp_pool = reverse_buffer_[data_id]; - if (qg_.is_quantized() && !tmp_pool.empty()) { - // RaBitQ estimates are directional: score destination -> source afresh. - std::vector reconstructed; - std::optional prepared; - qg_.prepare_build_query(data_id, reconstructed, prepared, build_cache()); - for (auto& candidate : tmp_pool) { - candidate.distance = qg_.quantized_distance(*prepared, candidate.id); - } - } - tmp_pool.reserve(tmp_pool.size() + degree_bound_); - tmp_pool.insert( - tmp_pool.end(), new_neighbors_[data_id].begin(), new_neighbors_[data_id].end() - ); - std::sort(tmp_pool.begin(), tmp_pool.end()); - heuristic_prune(data_id, tmp_pool, new_neighbors_[data_id], refine); - } -} - -inline void QGBuilder::random_init() { - const PID min_id = 0; - const PID max_id = num_nodes_ - 1; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - new_neighbors_[i].clear(); - pruned_neighbors_[i].clear(); - std::mt19937 rng(seed_ ^ static_cast(i)); - std::uniform_int_distribution distribution(min_id, max_id); - std::unordered_set neighbor_set; - neighbor_set.reserve(degree_bound_); - while (neighbor_set.size() < degree_bound_) { - const PID rand_id = distribution(rng); - if (rand_id != i) { - neighbor_set.emplace(rand_id); - } - } - - std::vector reconstructed; - std::optional prepared; - const float* cur_data = - qg_.prepare_build_query(i, reconstructed, prepared, build_cache()); - new_neighbors_[i].reserve(degree_bound_); - for (PID cur_neigh : neighbor_set) { - new_neighbors_[i].emplace_back( - cur_neigh, - qg_.point_distance(cur_data, prepared ? &*prepared : nullptr, cur_neigh) - ); - } - - degrees_[i] = new_neighbors_[i].size(); - qg_.update_qg( - i, - new_neighbors_[i], - build_cache(), - &rotated_neighbors_[static_cast(omp_get_thread_num())] - ); - } -} - -/** - * @brief refine the graph structure, make sure the degree for each vertex in qg equals the - * degree bound (multiple of 32) - * - */ -inline void QGBuilder::graph_refine() { -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - CandidateList& cur_neighbors = new_neighbors_[i]; - size_t cur_degree = cur_neighbors.size(); - - // skip vertices with enough neighbors - if (cur_degree >= degree_bound_) { - continue; - } - - CandidateList& pruned_list = pruned_neighbors_[i]; - CandidateList new_result; - new_result.reserve(degree_bound_); - - std::sort(pruned_list.begin(), pruned_list.end()); - - // use binary search to get refined results - CandidateList successful_result; - successful_result.reserve(degree_bound_); - float left = 0.5; - float right = 1.0; - size_t iter = 0; - while (iter++ < kMaxBsIter) { - float mid = (left + right) / 2; - add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); - if (new_result.size() < degree_bound_) { - left = mid; - } else { - right = mid; - successful_result = new_result; - } - } - - // Every successful right endpoint was already evaluated at that exact threshold. - if (!successful_result.empty()) { - new_result.swap(successful_result); - } else { - add_pruned_edges(cur_neighbors, pruned_list, new_result, right); - } - - // if the vertex still doesn't have enough neighbors, use random vertices - if (new_result.size() < degree_bound_) { - std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); - std::uniform_int_distribution distribution( - 0, static_cast(num_nodes_) - 1 - ); - std::unordered_set ids; - ids.reserve(degree_bound_); - for (auto& neighbor : new_result) { - ids.emplace(neighbor.id); - } - std::vector reconstructed; - std::optional prepared; - const float* source = - qg_.prepare_build_query(i, reconstructed, prepared, build_cache()); - while (new_result.size() < degree_bound_) { - const PID rand_id = distribution(rng); - if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { - new_result.emplace_back( - rand_id, - qg_.point_distance(source, prepared ? &*prepared : nullptr, rand_id) - ); - ids.emplace(rand_id); - } - } - } - - cur_neighbors = new_result; - } -} - -inline void QGBuilder::iter(bool refine) { - if (refine) { - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - pruned_neighbors_[i].clear(); - pruned_neighbors_[i].reserve(kMaxPrunedSize); - } - } - - search_new_neighbors(refine); - - add_reverse_edges(refine); - - // Use pruned edges to refine graph - if (refine) { - graph_refine(); - } - - // update qg -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); - ++index) { - const size_t i = static_cast(index); - qg_.update_qg( - i, - new_neighbors_[i], - build_cache(), - &rotated_neighbors_[static_cast(omp_get_thread_num())] - ); - degrees_[i] = new_neighbors_[i].size(); - } -} } // namespace rabitqlib::symqg diff --git a/include/rabitqlib/index/types.hpp b/include/rabitqlib/index/types.hpp new file mode 100644 index 00000000..f421a3b4 --- /dev/null +++ b/include/rabitqlib/index/types.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include +#include + +namespace rabitqlib { +using PID = uint32_t; +constexpr uint32_t kPidMax = 0xFFFFFFFF; + +template +struct AnnCandidate { + TP id = 0; + T distance = std::numeric_limits::max(); + + AnnCandidate() = default; + explicit AnnCandidate(TP vec_id, T dis) : id(vec_id), distance(dis) {} + + friend bool operator<(const AnnCandidate& first, const AnnCandidate& second) { + return first.distance < second.distance; + } + friend bool operator>(const AnnCandidate& first, const AnnCandidate& second) { + return first.distance > second.distance; + } + friend bool operator>=(const AnnCandidate& first, const AnnCandidate& second) { + return first.distance >= second.distance; + } + friend bool operator<=(const AnnCandidate& first, const AnnCandidate& second) { + return first.distance <= second.distance; + } +}; + +} // namespace rabitqlib diff --git a/include/rabitqlib/quantization/data_layout.hpp b/include/rabitqlib/quantization/data_layout.hpp deleted file mode 100644 index 15231b10..00000000 --- a/include/rabitqlib/quantization/data_layout.hpp +++ /dev/null @@ -1,397 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include "rabitqlib/fastscan/fastscan.hpp" - -namespace rabitqlib { -namespace detail { - -// Factor blocks are part of a packed byte layout and therefore need not be aligned for T. -// Access them through memcpy so the representation remains valid C++17 storage rather than -// manufacturing T objects inside a char buffer with reinterpret_cast. -template -class PackedValueRef { - static_assert( - std::is_trivially_copyable_v, "packed values must be trivially copyable" - ); - - public: - explicit PackedValueRef(char* data) : data_(data) {} - PackedValueRef(const PackedValueRef&) = default; - - PackedValueRef& operator=(T value) { - std::memcpy(data_, &value, sizeof(value)); - return *this; - } - - PackedValueRef& operator=(const PackedValueRef& other) { - return *this = static_cast(other); - } - - operator T() const { - T value; - std::memcpy(&value, data_, sizeof(value)); - return value; - } - - private: - char* data_; -}; - -template -class PackedArrayView { - static_assert( - std::is_trivially_copyable_v, "packed values must be trivially copyable" - ); - - public: - using difference_type = std::ptrdiff_t; - using value_type = T; - using pointer = void; - using reference = PackedValueRef; - using iterator_category = std::forward_iterator_tag; - - explicit PackedArrayView(char* data) : data_(data) {} - - [[nodiscard]] reference operator*() const { return reference(data_); } - [[nodiscard]] reference operator[](size_t index) const { - return reference(data_ + (index * sizeof(T))); - } - - PackedArrayView& operator++() { - data_ += sizeof(T); - return *this; - } - PackedArrayView operator++(int) { - PackedArrayView previous = *this; - ++*this; - return previous; - } - PackedArrayView& operator--() { - data_ -= sizeof(T); - return *this; - } - PackedArrayView& operator+=(difference_type count) { - data_ += count * static_cast(sizeof(T)); - return *this; - } - PackedArrayView& operator-=(difference_type count) { return *this += -count; } - - friend PackedArrayView operator+(PackedArrayView view, difference_type count) { - view += count; - return view; - } - friend PackedArrayView operator+(difference_type count, PackedArrayView view) { - return view + count; - } - friend PackedArrayView operator-(PackedArrayView view, difference_type count) { - view -= count; - return view; - } - friend difference_type operator-(PackedArrayView lhs, PackedArrayView rhs) { - return (lhs.data_ - rhs.data_) / static_cast(sizeof(T)); - } - friend bool operator==(PackedArrayView lhs, PackedArrayView rhs) { - return lhs.data_ == rhs.data_; - } - friend bool operator!=(PackedArrayView lhs, PackedArrayView rhs) { - return !(lhs == rhs); - } - - void copy_to(T* output, size_t count) const { - std::memcpy(output, data_, count * sizeof(T)); - } - - private: - char* data_; -}; - -template -class ConstPackedArrayView { - static_assert( - std::is_trivially_copyable_v, "packed values must be trivially copyable" - ); - - public: - explicit ConstPackedArrayView(const char* data) : data_(data) {} - - [[nodiscard]] T operator[](size_t index) const { - T value; - std::memcpy(&value, data_ + (index * sizeof(T)), sizeof(value)); - return value; - } - - void copy_to(T* output, size_t count) const { - std::memcpy(output, data_, count * sizeof(T)); - } - - private: - const char* data_; -}; - -template -[[nodiscard]] inline T load_packed_value(const char* data) { - static_assert( - std::is_trivially_copyable_v, "packed values must be trivially copyable" - ); - T value; - std::memcpy(&value, data, sizeof(value)); - return value; -} - -} // namespace detail - -template -struct BatchDataMap { - public: - explicit BatchDataMap(char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * fastscan::kBatchSize / 8)) // 1 bit code - , f_rescale_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize) - ) - , f_error_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize * 2) - ) {} - - [[nodiscard]] uint8_t* bin_code() { return batch_bin_code_; } - [[nodiscard]] detail::PackedArrayView f_add() { return f_add_; } - [[nodiscard]] detail::PackedArrayView f_rescale() { return f_rescale_; } - [[nodiscard]] detail::PackedArrayView f_error() { return f_error_; } - - static size_t data_bytes(size_t padded_dim) { - return (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize * 3); - } - - private: - uint8_t* batch_bin_code_; - detail::PackedArrayView f_add_; - detail::PackedArrayView f_rescale_; - detail::PackedArrayView f_error_; -}; - -template -struct ConstBatchDataMap { - public: - explicit ConstBatchDataMap(const char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * fastscan::kBatchSize / 8)) // 1 bit code - , f_rescale_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize) - ) - , f_error_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize * 2) - ) {} - - [[nodiscard]] const uint8_t* bin_code() const { return batch_bin_code_; } - [[nodiscard]] detail::ConstPackedArrayView f_add() const { return f_add_; } - [[nodiscard]] detail::ConstPackedArrayView f_rescale() const { return f_rescale_; } - [[nodiscard]] detail::ConstPackedArrayView f_error() const { return f_error_; } - - private: - const uint8_t* batch_bin_code_; - detail::ConstPackedArrayView f_add_; - detail::ConstPackedArrayView f_rescale_; - detail::ConstPackedArrayView f_error_; -}; - -template -struct QGBatchDataMap { - public: - explicit QGBatchDataMap(char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * fastscan::kBatchSize / 8)) // 1 bit code - , f_rescale_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize) - ) {} - - [[nodiscard]] uint8_t* bin_code() { return batch_bin_code_; } - [[nodiscard]] detail::PackedArrayView f_add() { return f_add_; } - [[nodiscard]] detail::PackedArrayView f_rescale() { return f_rescale_; } - - static size_t data_bytes(size_t padded_dim) { - return (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize * 2); - } - - private: - uint8_t* batch_bin_code_; - detail::PackedArrayView f_add_; - detail::PackedArrayView f_rescale_; -}; - -template -struct ConstQGBatchDataMap { - public: - explicit ConstQGBatchDataMap(const char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * fastscan::kBatchSize / 8)) // 1 bit code - , f_rescale_( - data + (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize) - ) {} - - [[nodiscard]] const uint8_t* bin_code() const { return batch_bin_code_; } - [[nodiscard]] detail::ConstPackedArrayView f_add() const { return f_add_; } - [[nodiscard]] detail::ConstPackedArrayView f_rescale() const { return f_rescale_; } - - static size_t data_bytes(size_t padded_dim) { - return (padded_dim * fastscan::kBatchSize / 8) + - (sizeof(T) * fastscan::kBatchSize * 2); - } - - private: - const uint8_t* batch_bin_code_; - detail::ConstPackedArrayView f_add_; - detail::ConstPackedArrayView f_rescale_; -}; - -template -struct ExDataMap { - public: - explicit ExDataMap(char* data, size_t padded_dim, size_t ex_bits) - : ex_code_(reinterpret_cast(data)) - , f_add_ex_(data + (padded_dim * ex_bits / 8)) - , f_recale_ex_(data + (padded_dim * ex_bits / 8) + sizeof(T)) {} - - static size_t data_bytes(size_t padded_dim, size_t ex_bits) { - return ex_bits > 0 ? (padded_dim * ex_bits / 8) + (sizeof(T) * 2) : 0; - } - - [[nodiscard]] uint8_t* ex_code() { return ex_code_; } - [[nodiscard]] detail::PackedValueRef f_add_ex() { return f_add_ex_; } - [[nodiscard]] detail::PackedValueRef f_rescale_ex() { return f_recale_ex_; } - - private: - uint8_t* ex_code_; - detail::PackedValueRef f_add_ex_; - detail::PackedValueRef f_recale_ex_; -}; - -template -struct ConstExDataMap { - public: - explicit ConstExDataMap(const char* data, size_t padded_dim, size_t ex_bits) - : ex_code_(reinterpret_cast(data)) - , f_add_ex_(data + (padded_dim * ex_bits / 8)) - , f_recale_ex_(data + (padded_dim * ex_bits / 8) + sizeof(T)) {} - - [[nodiscard]] const uint8_t* ex_code() const { return ex_code_; } - [[nodiscard]] T f_add_ex() const { return detail::load_packed_value(f_add_ex_); } - [[nodiscard]] T f_rescale_ex() const { - return detail::load_packed_value(f_recale_ex_); - } - - private: - const uint8_t* ex_code_; - const char* f_add_ex_; - const char* f_recale_ex_; -}; - -template -struct BaseDataMap { - public: - explicit BaseDataMap(char* data, size_t padded_dim, size_t base_bits) - : base_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * base_bits / 8)) - , f_rescale_(data + (padded_dim * base_bits / 8) + sizeof(T)) - , f_error_(data + (padded_dim * base_bits / 8) + (sizeof(T) * 2)) {} - - static size_t data_bytes(size_t padded_dim, size_t base_bits) { - return (padded_dim * base_bits / 8) + (sizeof(T) * 3); - } - - [[nodiscard]] uint8_t* base_code() { return base_code_; } - [[nodiscard]] detail::PackedValueRef f_add() { return f_add_; } - [[nodiscard]] detail::PackedValueRef f_rescale() { return f_rescale_; } - [[nodiscard]] detail::PackedValueRef f_error() { return f_error_; } - - private: - uint8_t* base_code_; - detail::PackedValueRef f_add_; - detail::PackedValueRef f_rescale_; - detail::PackedValueRef f_error_; -}; - -template -struct ConstBaseDataMap { - public: - explicit ConstBaseDataMap(const char* data, size_t padded_dim, size_t base_bits) - : base_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim * base_bits / 8)) - , f_rescale_(data + (padded_dim * base_bits / 8) + sizeof(T)) - , f_error_(data + (padded_dim * base_bits / 8) + (sizeof(T) * 2)) {} - - [[nodiscard]] const uint8_t* base_code() const { return base_code_; } - [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } - [[nodiscard]] T f_rescale() const { return detail::load_packed_value(f_rescale_); } - [[nodiscard]] T f_error() const { return detail::load_packed_value(f_error_); } - - private: - const uint8_t* base_code_; - const char* f_add_; - const char* f_rescale_; - const char* f_error_; -}; - -template -struct BinDataMap { - public: - explicit BinDataMap(char* data, size_t padded_dim) - : bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim / 8)) - , f_rescale_(data + (padded_dim / 8) + sizeof(T)) - , f_error_(data + (padded_dim / 8) + (sizeof(T) * 2)) {} - - [[nodiscard]] uint8_t* bin_code() { return bin_code_; } - [[nodiscard]] detail::PackedValueRef f_add() { return f_add_; } - [[nodiscard]] detail::PackedValueRef f_rescale() { return f_rescale_; } - [[nodiscard]] detail::PackedValueRef f_error() { return f_error_; } - - static size_t data_bytes(size_t padded_dim) { - return (padded_dim / 8) + (sizeof(T) * 3); - } - - private: - uint8_t* bin_code_; - detail::PackedValueRef f_add_; - detail::PackedValueRef f_rescale_; - detail::PackedValueRef f_error_; -}; - -template -struct ConstBinDataMap { - public: - explicit ConstBinDataMap(const char* data, size_t padded_dim) - : bin_code_(reinterpret_cast(data)) - , f_add_(data + (padded_dim / 8)) - , f_rescale_(data + (padded_dim / 8) + sizeof(T)) - , f_error_(data + (padded_dim / 8) + (sizeof(T) * 2)) {} - - [[nodiscard]] const uint8_t* bin_code() const { return bin_code_; } - [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } - [[nodiscard]] T f_rescale() const { return detail::load_packed_value(f_rescale_); } - [[nodiscard]] T f_error() const { return detail::load_packed_value(f_error_); } - - static size_t data_bytes(size_t padded_dim) { - return (padded_dim / 8) + (sizeof(T) * 3); - } - - private: - const uint8_t* bin_code_; - const char* f_add_; - const char* f_rescale_; - const char* f_error_; -}; -} // namespace rabitqlib diff --git a/include/rabitqlib/simd/hnsw_dispatch.hpp b/include/rabitqlib/simd/hnsw_dispatch.hpp deleted file mode 100644 index 0cd23752..00000000 --- a/include/rabitqlib/simd/hnsw_dispatch.hpp +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "rabitqlib/defines.hpp" - -namespace rabitqlib::hnsw { -class HierarchicalNSW; -namespace detail { -std::priority_queue> search_knn( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); - -std::priority_queue> search_knn_neon( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); - -std::priority_queue> search_knn_generic( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); -std::priority_queue> search_knn_avx2( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); - -std::priority_queue> search_knn_avx512_core( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); - -std::priority_queue> search_knn_avx512_popcnt( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 -); - -} // namespace detail -} // namespace rabitqlib::hnsw diff --git a/include/rabitqlib/simd/quantization_dispatch.hpp b/include/rabitqlib/simd/quantization_dispatch.hpp deleted file mode 100644 index d17cfe22..00000000 --- a/include/rabitqlib/simd/quantization_dispatch.hpp +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once - -#include - -namespace rabitqlib::simd { - -// Search normalized, nonnegative float magnitudes with max_code in [0, 255]. -// A negative result requests the event sweep when the bounded search cannot -// certify a winner within its work limit. -double best_rescale_factor( - const float* magnitudes, size_t dim, int max_code, double start, double end -); -double best_rescale_factor_generic( - const float* magnitudes, size_t dim, int max_code, double start, double end -); -double best_rescale_factor_avx2( - const float* magnitudes, size_t dim, int max_code, double start, double end -); -double best_rescale_factor_avx512( - const float* magnitudes, size_t dim, int max_code, double start, double end -); - -} // namespace rabitqlib::simd diff --git a/include/rabitqlib/utils/array.hpp b/include/rabitqlib/utils/array.hpp index 03ea19f1..fa09dd7d 100644 --- a/include/rabitqlib/utils/array.hpp +++ b/include/rabitqlib/utils/array.hpp @@ -19,6 +19,7 @@ #pragma once +#include #include #include #include @@ -108,19 +109,41 @@ class Array { , dims_{std::move(other.dims_)} , allocator_{std::move(other.allocator_)} {} - Array& operator=(Array&& other) noexcept { + // Unequal, non-propagating allocators require allocation and may throw. + // NOLINTBEGIN(performance-noexcept-move-constructor) + Array& operator=(Array&& other) noexcept( + std::is_nothrow_move_assignable_v && + (atraits::propagate_on_container_move_assignment::value + ? std::is_nothrow_move_assignable_v + : atraits::is_always_equal::value) + ) { + if (this == &other) { + return *this; + } + if constexpr (!atraits::propagate_on_container_move_assignment::value && !atraits::is_always_equal::value) { + if (allocator_ != other.allocator_ && other.pointer_ != nullptr) { + // Allocate with the destination's allocator before releasing either + // array. Unequal, non-propagating allocators cannot transfer storage. + Array replacement(other.dims_, allocator_); + std::copy_n(other.pointer_, other.size(), replacement.pointer_); + *this = std::move(replacement); + other.destroy(); + return *this; + } + } if (pointer_ != nullptr) { destroy(); } - if constexpr (atraits::propagate_on_container_move_assignment::value) { allocator_ = std::move(other.allocator_); } - dims_ = std::exchange(other.dims_, Dims()); + dims_ = std::move(other.dims_); pointer_ = std::exchange(other.pointer_, nullptr); return *this; } + // NOLINTEND(performance-noexcept-move-constructor) + [[nodiscard]] pointer data() { return pointer_; } [[nodiscard]] const_pointer data() const { return pointer_; } diff --git a/include/rabitqlib/utils/io.hpp b/include/rabitqlib/utils/io.hpp index 3a61f5a9..e7245c7b 100644 --- a/include/rabitqlib/utils/io.hpp +++ b/include/rabitqlib/utils/io.hpp @@ -51,7 +51,6 @@ void read_value(std::ifstream& input, T& value, const char* description) { } // namespace io_impl -// get num of bytes inline size_t get_filesize(const char* filename) { const auto file_size = std::filesystem::file_size(rabitqlib::io_impl::filesystem_path(filename)); diff --git a/include/rabitqlib/utils/openmp.hpp b/include/rabitqlib/utils/openmp.hpp new file mode 100644 index 00000000..77f6c002 --- /dev/null +++ b/include/rabitqlib/utils/openmp.hpp @@ -0,0 +1,15 @@ +#pragma once + +// CMake targets export an explicit choice. Direct header consumers without a +// target definition infer it from the compiler; match the linked library's setting. +#ifndef RABITQ_USE_OPENMP +#ifdef _OPENMP +#define RABITQ_USE_OPENMP 1 +#else +#define RABITQ_USE_OPENMP 0 +#endif +#endif + +#if RABITQ_USE_OPENMP +#include // IWYU pragma: export +#endif diff --git a/include/rabitqlib/utils/packed.hpp b/include/rabitqlib/utils/packed.hpp new file mode 100644 index 00000000..a01908b4 --- /dev/null +++ b/include/rabitqlib/utils/packed.hpp @@ -0,0 +1,145 @@ +#pragma once + +#include +#include +#include +#include + +namespace rabitqlib { +namespace detail { + +// Factor blocks are part of a packed byte layout and therefore need not be aligned for T. +// Access them through memcpy so the representation remains valid C++17 storage rather than +// manufacturing T objects inside a char buffer with reinterpret_cast. +template +class PackedValueRef { + static_assert( + std::is_trivially_copyable_v, "packed values must be trivially copyable" + ); + + public: + explicit PackedValueRef(char* data) : data_(data) {} + PackedValueRef(const PackedValueRef&) = default; + + PackedValueRef& operator=(T value) { + std::memcpy(data_, &value, sizeof(value)); + return *this; + } + + PackedValueRef& operator=(const PackedValueRef& other) { + return *this = static_cast(other); + } + + operator T() const { + T value; + std::memcpy(&value, data_, sizeof(value)); + return value; + } + + private: + char* data_; +}; + +template +class PackedArrayView { + static_assert( + std::is_trivially_copyable_v, "packed values must be trivially copyable" + ); + + public: + using difference_type = std::ptrdiff_t; + using value_type = T; + using pointer = void; + using reference = PackedValueRef; + using iterator_category = std::forward_iterator_tag; + + explicit PackedArrayView(char* data) : data_(data) {} + + [[nodiscard]] reference operator*() const { return reference(data_); } + [[nodiscard]] reference operator[](size_t index) const { + return reference(data_ + (index * sizeof(T))); + } + + PackedArrayView& operator++() { + data_ += sizeof(T); + return *this; + } + PackedArrayView operator++(int) { + PackedArrayView previous = *this; + ++*this; + return previous; + } + PackedArrayView& operator--() { + data_ -= sizeof(T); + return *this; + } + PackedArrayView& operator+=(difference_type count) { + data_ += count * static_cast(sizeof(T)); + return *this; + } + PackedArrayView& operator-=(difference_type count) { return *this += -count; } + + friend PackedArrayView operator+(PackedArrayView view, difference_type count) { + view += count; + return view; + } + friend PackedArrayView operator+(difference_type count, PackedArrayView view) { + return view + count; + } + friend PackedArrayView operator-(PackedArrayView view, difference_type count) { + view -= count; + return view; + } + friend difference_type operator-(PackedArrayView lhs, PackedArrayView rhs) { + return (lhs.data_ - rhs.data_) / static_cast(sizeof(T)); + } + friend bool operator==(PackedArrayView lhs, PackedArrayView rhs) { + return lhs.data_ == rhs.data_; + } + friend bool operator!=(PackedArrayView lhs, PackedArrayView rhs) { + return !(lhs == rhs); + } + + void copy_to(T* output, size_t count) const { + std::memcpy(output, data_, count * sizeof(T)); + } + + private: + char* data_; +}; + +template +class ConstPackedArrayView { + static_assert( + std::is_trivially_copyable_v, "packed values must be trivially copyable" + ); + + public: + explicit ConstPackedArrayView(const char* data) : data_(data) {} + + [[nodiscard]] T operator[](size_t index) const { + T value; + std::memcpy(&value, data_ + (index * sizeof(T)), sizeof(value)); + return value; + } + + void copy_to(T* output, size_t count) const { + std::memcpy(output, data_, count * sizeof(T)); + } + + private: + const char* data_; +}; + +template +[[nodiscard]] inline T load_packed_value(const char* data) { + static_assert( + std::is_trivially_copyable_v, "packed values must be trivially copyable" + ); + T value; + std::memcpy(&value, data, sizeof(value)); + return value; +} + +} // namespace detail +} // namespace rabitqlib diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 3c2633a8..4430f782 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -1,23 +1,25 @@ #pragma once -#include - #include #include #include #include #include -#include #include #include #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" +#include "rabitqlib/utils/openmp.hpp" // IWYU pragma: keep (configuration macros) #include "rabitqlib/utils/tools.hpp" +#if RABITQ_USE_OPENMP +#include +#endif + namespace rabitqlib { inline uint64_t load_unaligned_u64(const void* data) { @@ -198,6 +200,7 @@ template inline std::vector compute_centroid( const T* data, size_t num_points, size_t dim, size_t num_threads ) { +#if RABITQ_USE_OPENMP const auto thread_count = static_cast(std::max( 1, std::min( @@ -206,13 +209,23 @@ inline std::vector compute_centroid( static_cast(std::numeric_limits::max())} ) )); +#else + (void)num_threads; + constexpr int thread_count = 1; +#endif std::vector> all_results(thread_count, std::vector(dim, 0)); +#if RABITQ_USE_OPENMP #pragma omp parallel for schedule(dynamic) num_threads(thread_count) +#endif for (std::ptrdiff_t index = 0; index < static_cast(num_points); ++index) { const size_t i = static_cast(index); - auto tid = omp_get_thread_num(); +#if RABITQ_USE_OPENMP + const int tid = omp_get_thread_num(); +#else + constexpr int tid = 0; +#endif std::vector& cur_results = all_results[tid]; const T* cur_data = data + (dim * i); for (size_t k = 0; k < dim; ++k) { @@ -235,51 +248,6 @@ inline std::vector compute_centroid( return centroid; } -template -inline PID exact_nn( - const T* data, - const T* query, - size_t num_points, - size_t dim, - size_t num_threads, - T (*dist_func)(const T*, const T*, size_t) -) { - const auto thread_count = static_cast(std::max( - 1, - std::min( - {resolve_num_threads(num_threads), - std::max(num_points, 1), - static_cast(std::numeric_limits::max())} - ) - )); - std::vector> best_entries(thread_count); - -#pragma omp parallel for schedule(dynamic) num_threads(thread_count) - for (std::ptrdiff_t index = 0; index < static_cast(num_points); - ++index) { - const size_t i = static_cast(index); - auto tid = omp_get_thread_num(); - AnnCandidate& cur_entry = best_entries[tid]; - const T* cur_data = data + (dim * i); - - T distance = dist_func(cur_data, query, dim); - if (distance < cur_entry.distance) { - cur_entry.id = static_cast(i); - cur_entry.distance = distance; - } - } - - PID nearest_neighbor = 0; - T min_dist = std::numeric_limits::max(); - for (auto& candi : best_entries) { - if (candi.distance < min_dist) { - nearest_neighbor = candi.id; - min_dist = candi.distance; - } - } - return nearest_neighbor; -} - namespace excode_ipimpl { float ip16_fxu1_avx( diff --git a/include/rabitqlib/utils/warmup_space.hpp b/include/rabitqlib/utils/warmup_space.hpp index 749ac28e..e0b911dd 100644 --- a/include/rabitqlib/utils/warmup_space.hpp +++ b/include/rabitqlib/utils/warmup_space.hpp @@ -32,7 +32,7 @@ inline float warmup_ip_x0_q( float delta, float vl, size_t padded_dim, - [[maybe_unused]] size_t _b_query = 0 // not used + [[maybe_unused]] size_t unused_query_bits = 0 // not used ) { auto num_blk = padded_dim / 64; const auto* it_data = reinterpret_cast(data); @@ -73,8 +73,8 @@ inline float warmup_ip_x0_q( const uint64_t* query, float delta, float vl, - size_t _padded_dim = 0, // not used - size_t _b_query = 0 // not used + size_t unused_padded_dim = 0, // not used + size_t unused_query_bits = 0 // not used ) { return rabitqlib::warmup_ip_x0_q(data, query, delta, vl, padded_dim); } diff --git a/pyproject.toml b/pyproject.toml index 063dd637..836c8ead 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "rabitqlib" -version = "0.6.0" +version = "0.7.0" description = "RaBitQ Python bindings for indexes and clustering" readme = "README.md" requires-python = ">=3.11" @@ -33,7 +33,8 @@ test = ["pytest>=7", "numpy"] # Disable directory auto-discovery so we can map it via CMake's install target wheel.packages = [] -# Python packages need the extension and core library, not C++ sample executables. +# Python wheels include the extension and its core/FastScan/index dependencies; +# C++ sample executables are disabled. cmake.args = [ "-DRABITQ_BUILD_PYTHON_BINDINGS=ON", "-DRABITQ_BUILD_SAMPLES=OFF", diff --git a/python/utils/io.py b/python/utils/io.py index 6316108b..1d23b3af 100644 --- a/python/utils/io.py +++ b/python/utils/io.py @@ -15,21 +15,56 @@ def read_fvecs(filename): return read_ivecs(filename).view("float32") +_WRITE_BLOCK_BYTES = 4 * 1024 * 1024 + + +def _blocks(m, order=None): + rows = max(1, _WRITE_BLOCK_BYTES // ((m.shape[1] + 1) * 4)) + count = len(m) if order is None else len(order) + for start in range(0, count, rows): + yield ( + m[start : start + rows] if order is None else m[order[start : start + rows]] + ) + + +def _write_records(stream, blocks, dim): + for block in blocks: + # Preserve struct.pack's native byte order and signed-int32 validation. + struct.pack("i", dim) + if block.shape[1] != dim: + raise struct.error( + f"pack expected {dim} items for packing (got {block.shape[1]})" + ) + if block.dtype.kind in "iu": + if block.size and (block.min() < -(2**31) or block.max() > 2**31 - 1): + raise struct.error( + "'i' format requires -2147483648 <= number <= 2147483647" + ) + else: + # Object arrays retain struct.pack's __index__ support; floats are rejected. + packed = struct.pack(f"{block.size}i", *block.flat) + block = np.frombuffer(packed, dtype="=i4").reshape(block.shape) + records = np.empty((len(block), dim + 1), dtype="=i4") + records[:, 0] = dim + records[:, 1:] = block + records.tofile(stream) + + def write_ivecs(filename, m): print(f"Writing File - {filename}") - n, d = m.shape - myimt = "i" * d - with open(filename, "wb") as f: - for i in range(n): - f.write(struct.pack("i", d)) - bin = struct.pack(myimt, *m[i]) - f.write(bin) + _, dim = m.shape + with open(filename, "wb") as stream: + _write_records(stream, _blocks(m), dim) print(f"\t{filename} wrote") def write_fvecs(filename, m): - m = m.astype("float32") - write_ivecs(filename, m.view("int32")) + print(f"Writing File - {filename}") + _, dim = m.shape + blocks = (block.astype("float32").view("int32") for block in _blocks(m)) + with open(filename, "wb") as stream: + _write_records(stream, blocks, dim) + print(f"\t{filename} wrote") def read_ibin(filename): @@ -51,8 +86,8 @@ def write_sampled_data(filename, m, num_sampled, seed=42): np.random.seed(seed) sequence = np.random.permutation(num_points) - with open(filename, "wb") as f: - for i in range(num_sampled): - f.write(struct.pack("i", d)) - bin = struct.pack(f"{d}i", *m[sequence[i]]) - f.write(bin) + count = len(range(num_sampled)) + if count > num_points: + raise IndexError("sample count exceeds the number of points") + with open(filename, "wb") as stream: + _write_records(stream, _blocks(m, sequence[:count]), d) diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index 9b459b33..68ce0f27 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -5,13 +5,9 @@ if(PYTHON_EXECUTABLE AND NOT Python_EXECUTABLE) endif() # 2. Locate Python, NumPy, and pybind11 from the active Python environment. -# CMake 3.18+ can request only the module-development files. This avoids -# requiring libpython for embedding, which manylinux images do not provide. -if(CMAKE_VERSION VERSION_LESS 3.18) - find_package(Python REQUIRED COMPONENTS Interpreter Development NumPy) -else() - find_package(Python REQUIRED COMPONENTS Interpreter Development.Module NumPy) -endif() +# Request module-development files without requiring libpython for embedding, +# which manylinux images do not provide. +find_package(Python REQUIRED COMPONENTS Interpreter Development.Module NumPy) find_package(pybind11 CONFIG REQUIRED) # 3. Define the shared pybind11 module. @@ -23,6 +19,8 @@ pybind11_add_module(_rabitqlib kmeans_bindings.cpp ) +add_library(rabitqlib::python ALIAS _rabitqlib) + target_include_directories(_rabitqlib PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ) @@ -30,7 +28,7 @@ rabitq_configure_target(_rabitqlib) # 4. Imported targets provide the matching Python, NumPy, and OpenMP settings. target_link_libraries(_rabitqlib PRIVATE - rabitq_headers + rabitqlib::index Python::NumPy ) @@ -61,3 +59,16 @@ endif() # inside the target wheel, completely bypassing your local layout. install(TARGETS _rabitqlib DESTINATION rabitqlib) install(FILES __init__.py DESTINATION rabitqlib) + +if(RABITQ_BUILD_BENCHMARKS) + # Benchmark the production conversion helper without importing or replacing + # the installed package. This module is never installed in the wheel. + pybind11_add_module(rabitq_ids_benchmark NO_EXTRAS ../benchmarks/hotpaths/ids.cpp) + target_include_directories(rabitq_ids_benchmark PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + target_link_libraries(rabitq_ids_benchmark PRIVATE rabitqlib::index) + rabitq_configure_target(rabitq_ids_benchmark) + set_target_properties(rabitq_ids_benchmark PROPERTIES INTERPROCEDURAL_OPTIMIZATION OFF) + add_test(NAME PythonIdsHotpath COMMAND "${Python_EXECUTABLE}" + "${PROJECT_SOURCE_DIR}/benchmarks/hotpaths/ids.py" + --module-dir "$" --rows 32 --iterations 2 --repetitions 1) +endif() diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index ca99e806..911763d3 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -1,7 +1,7 @@ """Python bindings for RaBitQLib indexes and clustering. -This module re-exports symbols from the compiled `rabitqlib` extension -so code can import using `from python_bindings import ...` if needed. +The installed ``rabitqlib`` package loads its compiled ``._rabitqlib`` extension +and adds the clustering wrappers. Import public classes from ``rabitqlib``. """ from ._rabitqlib import ( diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 9a1f9bca..852b6d44 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -1,16 +1,20 @@ #pragma once #include -#include +#include -#include #include #include +#include #include #include +#include +#include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/packed.hpp" namespace py = pybind11; @@ -127,4 +131,41 @@ ensure_1d_integer_array(py::handle value, const char* name) { return ensure_1d_array(value, name); } +// Validate before narrowing, and own the result for APIs that mutate assignments. +inline std::vector checked_cluster_ids( + py::handle value, + size_t rows, + size_t clusters, + const char* range_error = "cluster_ids must be in [0, num_clusters)" +) { + const auto checked_copy = [&](const auto& ids) { + if (static_cast(ids.shape(0)) != rows) { + throw std::invalid_argument( + "cluster_ids length must match number of rows in data" + ); + } + std::vector result(rows); + using Value = typename std::decay_t::value_type; + // C-contiguous NumPy arrays can still have an unaligned data pointer. + const rabitqlib::detail::ConstPackedArrayView values( + reinterpret_cast(ids.data()) + ); + for (size_t i = 0; i < rows; ++i) { + const int64_t id = values[i]; + if (id < 0 || static_cast(id) >= clusters || + static_cast(id) > std::numeric_limits::max()) { + throw std::invalid_argument(range_error); + } + result[i] = static_cast(id); + } + return result; + }; + // Already-native IDs need no intermediate int64 array. Keep the owned copy + // required by construction APIs, and validate it just like wider integers. + if (py::isinstance>(value)) { + return checked_copy(ensure_1d_array(value, "cluster_ids")); + } + return checked_copy(ensure_1d_integer_array(value, "cluster_ids")); +} + } // namespace rabitqlib::python_bindings diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 1e506d25..1fa8f856 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -6,8 +6,8 @@ #include #include +#include #include -#include #include #include #include @@ -16,8 +16,9 @@ #include #include "bindings_common.hpp" -#include "rabitqlib/defines.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/packed.hpp" namespace py = pybind11; @@ -34,15 +35,14 @@ class HnswIndex { const std::string& metric = "l2", size_t random_seed = 100 ) - : dim_(dim) - , max_elements_(max_elements) - , M_(M) - , ef_construction_(ef_construction) - , nbits_(nbits) - , metric_(metric_from_string(metric)) - , random_seed_(random_seed) - , index_(std::make_unique( - max_elements, dim, nbits, M, ef_construction, random_seed, metric_ + : index_(std::make_unique( + max_elements, + dim, + nbits, + M, + ef_construction, + random_seed, + metric_from_string(metric) )) {} void build( @@ -55,45 +55,28 @@ class HnswIndex { const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); - auto cluster_ids_array = - ensure_1d_array(cluster_ids, "cluster_ids"); - if (static_cast(data_array.shape(1)) != dim_) { + if (static_cast(data_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("data dimension does not match index dim"); } if (data_array.shape(0) == 0 || - static_cast(data_array.shape(0)) > max_elements_) { + static_cast(data_array.shape(0)) > index_->max_elements()) { throw std::invalid_argument( "number of data rows must be between 1 and max_elements" ); } - if (static_cast(centroids_array.shape(1)) != dim_) { + if (static_cast(centroids_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("centroid dimension does not match index dim"); } - if (static_cast(cluster_ids_array.shape(0)) != - static_cast(data_array.shape(0))) { - throw std::invalid_argument( - "cluster_ids length must match number of rows in data" - ); - } - const size_t num_clusters = static_cast(centroids_array.shape(0)); if (num_clusters == 0) { throw std::invalid_argument("at least one centroid is required"); } - for (py::ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { - if (cluster_ids_array.data()[i] >= num_clusters) { - throw std::invalid_argument("cluster_ids contains an out-of-range value"); - } - } - // Ensure cluster_ids are writable for the C++ API by making a copy - std::vector cluster_ids_vec( - static_cast(cluster_ids_array.shape(0)) - ); - std::memcpy( - cluster_ids_vec.data(), - cluster_ids_array.data(), - cluster_ids_vec.size() * sizeof(rabitqlib::PID) + auto cluster_ids_vec = checked_cluster_ids( + cluster_ids, + static_cast(data_array.shape(0)), + num_clusters, + "cluster_ids contains an out-of-range value" ); const auto* data_ptr = data_array.data(); @@ -111,7 +94,6 @@ class HnswIndex { fast_quantization ); } - num_clusters_ = num_clusters; built_ = true; } @@ -126,29 +108,14 @@ class HnswIndex { if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before add"); } - if (static_cast(data_array.shape(1)) != dim_) { + if (static_cast(data_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("data dimension does not match index dim"); } const auto rows = static_cast(data_array.shape(0)); - // Validate as int64 first: casting straight to uint32 would wrap negative - // or oversized values onto valid clusters. std::vector assigned; if (!cluster_ids.is_none()) { - auto cluster_ids_array = ensure_1d_integer_array(cluster_ids, "cluster_ids"); - if (static_cast(cluster_ids_array.shape(0)) != rows) { - throw std::invalid_argument( - "cluster_ids length must match number of rows in data" - ); - } - assigned.reserve(rows); - for (py::ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { - const int64_t id = cluster_ids_array.data()[i]; - if (id < 0 || static_cast(id) >= num_clusters_) { - throw std::invalid_argument("cluster_ids must be in [0, num_clusters)"); - } - assigned.push_back(static_cast(id)); - } + assigned = checked_cluster_ids(cluster_ids, rows, index_->num_clusters()); } // Allocated before the index is touched. @@ -175,7 +142,6 @@ class HnswIndex { } py::gil_scoped_release release; index_->resize(new_max_elements); - max_elements_ = index_->max_elements(); } size_t remove(py::handle ids) { @@ -184,14 +150,17 @@ class HnswIndex { if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before remove"); } - // Labels are dense: construct gives 0..n-1 and add continues from num_points. - const size_t count = index_->num_points(); + // Removed labels remain allocated; validate against the full stored ID range. + const size_t count = index_->num_stored_points(); std::vector to_remove; to_remove.reserve(static_cast(ids_array.shape(0))); + const rabitqlib::detail::ConstPackedArrayView values( + reinterpret_cast(ids_array.data()) + ); for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) { - const int64_t id = ids_array.data()[i]; + const int64_t id = values[static_cast(i)]; if (id < 0 || static_cast(id) >= count) { - throw std::invalid_argument("ids must be in [0, num_points)"); + throw std::invalid_argument("ids must be in [0, num_stored_points)"); } to_remove.push_back(static_cast(id)); } @@ -203,6 +172,10 @@ class HnswIndex { const auto operation = access_.read(); return built_ ? index_->num_points() : 0; } + [[nodiscard]] size_t num_stored_points() const { + const auto operation = access_.read(); + return built_ ? index_->num_stored_points() : 0; + } py::tuple search(py::handle queries, size_t k, size_t ef = 0, size_t num_threads = 1) { const auto operation = access_.read(); @@ -210,13 +183,14 @@ class HnswIndex { if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before search"); } - if (dim_ != 0 && static_cast(query_array.shape(1)) != dim_) { + if (index_->dimension() != 0 && + static_cast(query_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("query dimension does not match index dim"); } if (ef == 0) { ef = std::max(k, 10); } - if (k == 0 || k > max_elements_) { + if (k == 0 || k > index_->max_elements()) { throw std::invalid_argument("k must be between 1 and max_elements"); } @@ -271,28 +245,21 @@ class HnswIndex { py::gil_scoped_release release; wrapper.index_->load(path.c_str()); } - wrapper.dim_ = wrapper.index_->dimension(); - wrapper.max_elements_ = wrapper.index_->max_elements(); - wrapper.M_ = wrapper.index_->M(); - wrapper.ef_construction_ = wrapper.index_->ef_construction(); - wrapper.nbits_ = wrapper.index_->nbits(); - wrapper.num_clusters_ = wrapper.index_->num_clusters(); - wrapper.metric_ = wrapper.index_->metric_type(); wrapper.built_ = true; return wrapper; } [[nodiscard]] size_t dim() const { const auto operation = access_.read(); - return dim_; + return index_->dimension(); } [[nodiscard]] size_t max_elements() const { const auto operation = access_.read(); - return max_elements_; + return index_->max_elements(); } [[nodiscard]] size_t nbits() const { const auto operation = access_.read(); - return nbits_; + return index_->nbits(); } [[nodiscard]] bool is_built() const { const auto operation = access_.read(); @@ -300,24 +267,16 @@ class HnswIndex { } [[nodiscard]] size_t num_clusters() const { const auto operation = access_.read(); - return num_clusters_; + return index_->num_clusters(); } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); - return metric_to_string(metric_); + return metric_to_string(index_->metric_type()); } private: HnswIndex() = default; - size_t dim_ = 0; - size_t max_elements_ = 0; - size_t M_ = 0; - size_t ef_construction_ = 0; - size_t nbits_ = 0; - rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; - size_t random_seed_ = 100; - size_t num_clusters_ = 0; bool built_ = false; mutable IndexAccess access_; std::unique_ptr index_; @@ -371,7 +330,8 @@ void register_hnsw(py::module_& m) { &HnswIndex::remove, py::arg("ids"), "Exclude ids from later search results and return how many were newly " - "removed. Removed points stay in the graph and in num_points, and cannot " + "removed. Removed points stay in the graph and in num_stored_points, and " + "cannot " "be restored." ) .def("save", &HnswIndex::save, py::arg("path")) @@ -379,6 +339,7 @@ void register_hnsw(py::module_& m) { .def_property_readonly("dim", &HnswIndex::dim) .def_property_readonly("max_elements", &HnswIndex::max_elements) .def_property_readonly("num_points", &HnswIndex::num_points) + .def_property_readonly("num_stored_points", &HnswIndex::num_stored_points) .def_property_readonly("nbits", &HnswIndex::nbits) .def_property_readonly("num_clusters", &HnswIndex::num_clusters) .def_property_readonly("is_built", &HnswIndex::is_built) diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 9fe992c8..2ee059d0 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -17,10 +17,11 @@ #include #include "bindings_common.hpp" -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/packed.hpp" namespace py = pybind11; @@ -65,21 +66,18 @@ class IvfIndex { size_t nbits, const std::string& metric = "l2", const std::string& initializer = "auto" - ) - : dim_(dim) - , max_elements_(max_elements) - , num_clusters_(num_clusters) - , nbits_(nbits) - , metric_(metric_from_string(metric)) - , index_(std::make_unique( - max_elements, - dim, - num_clusters, - nbits, - metric_, - rabitqlib::RotatorType::FhtKacRotator, - initializer_from_string(initializer) - )) {} + ) { + const auto metric_type = metric_from_string(metric); + index_ = std::make_unique( + max_elements, + dim, + num_clusters, + nbits, + metric_type, + rabitqlib::RotatorType::FhtKacRotator, + initializer_from_string(initializer) + ); + } void build( py::handle data, @@ -91,34 +89,29 @@ class IvfIndex { const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); - auto cluster_ids_array = - ensure_1d_array(cluster_ids, "cluster_ids"); - if (static_cast(data_array.shape(1)) != dim_) { + if (static_cast(data_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("data dimension does not match index dim"); } - if (static_cast(data_array.shape(0)) != max_elements_) { + if (static_cast(data_array.shape(0)) != index_->max_elements()) { throw std::invalid_argument("number of data rows must match index max_elements" ); } - if (static_cast(centroids_array.shape(1)) != dim_) { + if (static_cast(centroids_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("centroid dimension does not match index dim"); } - if (static_cast(centroids_array.shape(0)) != num_clusters_) { + if (static_cast(centroids_array.shape(0)) != index_->num_clusters()) { throw std::invalid_argument( "number of centroid rows must match index num_clusters" ); } - if (static_cast(cluster_ids_array.shape(0)) != - static_cast(data_array.shape(0))) { - throw std::invalid_argument( - "cluster_ids length must match number of rows in data" - ); - } + auto cluster_ids_vec = checked_cluster_ids( + cluster_ids, static_cast(data_array.shape(0)), index_->num_clusters() + ); const auto* data_ptr = data_array.data(); const auto* centroids_ptr = centroids_array.data(); - const auto* cluster_ids_ptr = cluster_ids_array.data(); + const auto* cluster_ids_ptr = cluster_ids_vec.data(); { py::gil_scoped_release release; index_->construct( @@ -139,32 +132,17 @@ class IvfIndex { if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before add"); } - if (static_cast(data_array.shape(1)) != dim_) { + if (static_cast(data_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("data dimension does not match index dim"); } const auto rows = static_cast(data_array.shape(0)); - // Validate cluster IDs as int64 first: casting straight to uint32 would wrap - // negative or oversized values onto valid clusters. std::vector assigned; if (!cluster_ids.is_none()) { - auto cluster_ids_array = ensure_1d_integer_array(cluster_ids, "cluster_ids"); - if (static_cast(cluster_ids_array.shape(0)) != rows) { - throw std::invalid_argument( - "cluster_ids length must match number of rows in data" - ); - } - assigned.reserve(rows); - for (py::ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { - const int64_t id = cluster_ids_array.data()[i]; - if (id < 0 || static_cast(id) >= num_clusters_) { - throw std::invalid_argument("cluster_ids must be in [0, num_clusters)"); - } - assigned.push_back(static_cast(id)); - } + assigned = checked_cluster_ids(cluster_ids, rows, index_->num_clusters()); } - const size_t first = index_->max_elements(); + const size_t first = index_->num_stored_points(); // Allocate the returned IDs before committing any changes to the index. auto ids = py::array_t(static_cast(rows)); auto* ids_data = ids.mutable_data(); @@ -183,7 +161,6 @@ class IvfIndex { num_threads ); } - max_elements_ = index_->max_elements(); return ids; } @@ -195,10 +172,13 @@ class IvfIndex { } std::vector to_remove; to_remove.reserve(static_cast(ids_array.shape(0))); + const rabitqlib::detail::ConstPackedArrayView values( + reinterpret_cast(ids_array.data()) + ); for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) { - const int64_t id = ids_array.data()[i]; - if (id < 0 || static_cast(id) >= max_elements_) { - throw std::invalid_argument("ids must be in [0, max_elements)"); + const int64_t id = values[static_cast(i)]; + if (id < 0 || static_cast(id) >= index_->num_stored_points()) { + throw std::invalid_argument("ids must be in [0, num_stored_points)"); } to_remove.push_back(static_cast(id)); } @@ -218,10 +198,10 @@ class IvfIndex { if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before search"); } - if (static_cast(query_array.shape(1)) != dim_) { + if (static_cast(query_array.shape(1)) != index_->dimension()) { throw std::invalid_argument("query dimension does not match index dim"); } - if (k == 0 || k > max_elements_) { + if (k == 0 || k > index_->max_elements()) { throw std::invalid_argument("k must be between 1 and max_elements"); } if (nprobe == 0) { @@ -274,30 +254,33 @@ class IvfIndex { py::gil_scoped_release release; wrapper.index_->load(path.c_str()); } - wrapper.dim_ = wrapper.index_->dimension(); - wrapper.max_elements_ = wrapper.index_->max_elements(); - wrapper.num_clusters_ = wrapper.index_->num_clusters(); - wrapper.nbits_ = wrapper.index_->nbits(); - wrapper.metric_ = wrapper.index_->metric_type(); wrapper.built_ = true; return wrapper; } [[nodiscard]] size_t dim() const { const auto operation = access_.read(); - return dim_; + return index_->dimension(); } [[nodiscard]] size_t max_elements() const { const auto operation = access_.read(); - return max_elements_; + return index_->max_elements(); + } + [[nodiscard]] size_t num_points() const { + const auto operation = access_.read(); + return index_->num_points(); + } + [[nodiscard]] size_t num_stored_points() const { + const auto operation = access_.read(); + return index_->num_stored_points(); } [[nodiscard]] size_t num_clusters() const { const auto operation = access_.read(); - return num_clusters_; + return index_->num_clusters(); } [[nodiscard]] size_t nbits() const { const auto operation = access_.read(); - return nbits_; + return index_->nbits(); } [[nodiscard]] bool is_built() const { const auto operation = access_.read(); @@ -309,17 +292,12 @@ class IvfIndex { } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); - return metric_to_string(metric_); + return metric_to_string(index_->metric_type()); } private: IvfIndex() = default; - size_t dim_ = 0; - size_t max_elements_ = 0; - size_t num_clusters_ = 0; - size_t nbits_ = 0; - rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; bool built_ = false; mutable IndexAccess access_; std::unique_ptr index_; @@ -391,6 +369,8 @@ void register_ivf(py::module_& m) { .def_static("load", &IvfIndex::load, py::arg("path")) .def_property_readonly("dim", &IvfIndex::dim) .def_property_readonly("max_elements", &IvfIndex::max_elements) + .def_property_readonly("num_points", &IvfIndex::num_points) + .def_property_readonly("num_stored_points", &IvfIndex::num_stored_points) .def_property_readonly("num_clusters", &IvfIndex::num_clusters) .def_property_readonly("nbits", &IvfIndex::nbits) .def_property_readonly("is_built", &IvfIndex::is_built) diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index 3fbe3db4..e4d0eb08 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -1,8 +1,8 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" #include "rabitqlib/utils/tools.hpp" namespace py = pybind11; diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index e1bbce45..0dd09948 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -15,11 +15,12 @@ #include #include "bindings_common.hpp" -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/index/types.hpp" namespace py = pybind11; @@ -89,7 +90,6 @@ class SymqgIndex { builder.build(); } index_ = std::move(index); - num_points_ = num_points; built_ = true; } @@ -99,7 +99,7 @@ class SymqgIndex { if (!built_) { throw std::runtime_error("SymqgIndex must be built or loaded before search"); } - if (k == 0 || k > num_points_) { + if (k == 0 || k > index_->num_points()) { throw std::invalid_argument("k must be between 1 and num_points"); } if (ef == 0) { @@ -152,7 +152,6 @@ class SymqgIndex { py::gil_scoped_release release; wrapper.index_->load(path.c_str()); } - wrapper.num_points_ = wrapper.index_->num_vertices(); wrapper.dim_ = wrapper.index_->dimension(); wrapper.max_degree_ = wrapper.index_->degree_bound(); wrapper.metric_ = wrapper.index_->metric_type(); @@ -171,7 +170,15 @@ class SymqgIndex { } [[nodiscard]] size_t num_points() const { const auto operation = access_.read(); - return num_points_; + return index_ ? index_->num_points() : 0; + } + [[nodiscard]] size_t num_stored_points() const { + const auto operation = access_.read(); + return index_ ? index_->num_stored_points() : 0; + } + [[nodiscard]] size_t max_elements() const { + const auto operation = access_.read(); + return index_ ? index_->max_elements() : 0; } [[nodiscard]] bool is_built() const { const auto operation = access_.read(); @@ -191,7 +198,6 @@ class SymqgIndex { size_t dim_ = 0; size_t max_degree_ = 0; - size_t num_points_ = 0; rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; size_t quantization_bits_ = 0; bool built_ = false; @@ -242,6 +248,8 @@ void register_symqg(py::module_& m) { .def_property_readonly("dim", &SymqgIndex::dim) .def_property_readonly("max_degree", &SymqgIndex::max_degree) .def_property_readonly("num_points", &SymqgIndex::num_points) + .def_property_readonly("num_stored_points", &SymqgIndex::num_stored_points) + .def_property_readonly("max_elements", &SymqgIndex::max_elements) .def_property_readonly("is_built", &SymqgIndex::is_built) .def_property_readonly("metric", &SymqgIndex::metric) .def_property_readonly("quantization_bits", &SymqgIndex::quantization_bits); diff --git a/sample/cpp/CMakeLists.txt b/sample/cpp/CMakeLists.txt index 060966dc..416d16dd 100644 --- a/sample/cpp/CMakeLists.txt +++ b/sample/cpp/CMakeLists.txt @@ -1,3 +1,21 @@ +add_executable(stable_core_example stable_core/main.cpp) +target_link_libraries(stable_core_example PRIVATE rabitqlib::core) +rabitq_configure_target(stable_core_example) + +add_executable(core_example core/main.cpp) +target_link_libraries(core_example PRIVATE rabitqlib::core) +rabitq_configure_target(core_example) + +if(RABITQ_BUILD_FASTSCAN) + add_executable(fastscan_example fastscan/main.cpp) + target_link_libraries(fastscan_example PRIVATE rabitqlib::fastscan) + rabitq_configure_target(fastscan_example) +endif() + +if(NOT RABITQ_BUILD_INDEX) + return() +endif() + set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin) add_executable(symqg_indexing symqg_indexing.cpp) @@ -21,6 +39,6 @@ foreach(RABITQ_SAMPLE_TARGET hnsw_rabitq_indexing hnsw_rabitq_querying ) - target_link_libraries(${RABITQ_SAMPLE_TARGET} PRIVATE rabitq_headers) + target_link_libraries(${RABITQ_SAMPLE_TARGET} PRIVATE rabitqlib::index) rabitq_configure_target(${RABITQ_SAMPLE_TARGET}) endforeach() diff --git a/sample/cpp/core/CMakeLists.txt b/sample/cpp/core/CMakeLists.txt new file mode 100644 index 00000000..63a0be57 --- /dev/null +++ b/sample/cpp/core/CMakeLists.txt @@ -0,0 +1,9 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitqlib_core_example LANGUAGES CXX) + +find_package(rabitqlib CONFIG REQUIRED COMPONENTS core) +add_executable(core_example main.cpp) +target_link_libraries(core_example PRIVATE rabitqlib::core) + +enable_testing() +add_test(NAME core_example COMMAND core_example) diff --git a/sample/cpp/core/main.cpp b/sample/cpp/core/main.cpp new file mode 100644 index 00000000..30b8d65f --- /dev/null +++ b/sample/cpp/core/main.cpp @@ -0,0 +1,77 @@ +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/utils/space.hpp" + +int main() { + try { + constexpr size_t kDim = 65; + constexpr size_t kTotalBits = 4; + constexpr size_t kExBits = kTotalBits - 1; + std::vector vector(kDim), query(kDim); + for (size_t i = 0; i < kDim; ++i) { + vector[i] = std::sin(static_cast(i) * 0.17F); + query[i] = std::cos(static_cast(i) * 0.11F); + } + + // Data, centroid and query must use the same rotation and padded domain. + std::unique_ptr> rotator(rabitqlib::choose_rotator( + kDim, rabitqlib::RotatorType::FhtKacRotator, 0, 42 + )); + const size_t padded_dim = rotator->size(); + std::vector rotated_vector(padded_dim), rotated_query(padded_dim); + std::vector centroid(padded_dim, 0.0F); + rotator->rotate(vector.data(), rotated_vector.data()); + rotator->rotate(query.data(), rotated_query.data()); + + std::vector bin(rabitqlib::BinDataMap::data_bytes(padded_dim)); + std::vector extra(rabitqlib::ExDataMap::data_bytes(padded_dim, kExBits) + ); + const rabitqlib::quant::RabitqConfig config; + rabitqlib::quant::quantize_split_single( + rotated_vector.data(), + centroid.data(), + padded_dim, + kExBits, + bin.data(), + extra.data(), + rabitqlib::METRIC_L2, + config + ); + + // The query object borrows rotated_query, which stays alive through estimation. + rabitqlib::SplitSingleQuery prepared_query( + rotated_query.data(), padded_dim, kExBits, config + ); + prepared_query.set_g_add( + std::sqrt(rabitqlib::l2norm_sqr(rotated_query.data(), padded_dim)) + ); + const auto estimate = rabitqlib::estimate_split_single_full( + bin.data(), + extra.data(), + rabitqlib::select_excode_ipfunc(kExBits), + prepared_query, + padded_dim, + kExBits, + prepared_query.g_add(), + prepared_query.g_error() + ); + std::cout << "Approximate squared L2: " << estimate.distance << '\n' + << "Exact squared L2: " + << rabitqlib::euclidean_sqr(vector.data(), query.data(), kDim) << '\n'; + return std::isfinite(estimate.distance) ? 0 : 1; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/sample/cpp/fastscan/CMakeLists.txt b/sample/cpp/fastscan/CMakeLists.txt new file mode 100644 index 00000000..b6d58d67 --- /dev/null +++ b/sample/cpp/fastscan/CMakeLists.txt @@ -0,0 +1,9 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitqlib_fastscan_example LANGUAGES CXX) + +find_package(rabitqlib CONFIG REQUIRED COMPONENTS fastscan) +add_executable(fastscan_example main.cpp) +target_link_libraries(fastscan_example PRIVATE rabitqlib::fastscan) + +enable_testing() +add_test(NAME fastscan_example COMMAND fastscan_example) diff --git a/sample/cpp/fastscan/main.cpp b/sample/cpp/fastscan/main.cpp new file mode 100644 index 00000000..7253b3af --- /dev/null +++ b/sample/cpp/fastscan/main.cpp @@ -0,0 +1,107 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/utils/space.hpp" + +int main() { + try { + constexpr size_t kDim = 65, kCount = 49, kExBits = 3; + std::unique_ptr> rotator(rabitqlib::choose_rotator( + kDim, rabitqlib::RotatorType::FhtKacRotator, 0, 42 + )); + const size_t padded_dim = rotator->size(); + std::vector vector(kDim), query(kDim), original_centroid(kDim); + std::vector rotated_data(kCount * padded_dim), rotated_query(padded_dim); + std::vector centroid(padded_dim, 0.0F); + for (size_t i = 0; i < kDim; ++i) { + query[i] = std::cos(static_cast(i) * 0.11F); + original_centroid[i] = 0.1F * std::sin(static_cast(i) * 0.23F); + } + rotator->rotate(query.data(), rotated_query.data()); + rotator->rotate(original_centroid.data(), centroid.data()); + for (size_t row = 0; row < kCount; ++row) { + for (size_t i = 0; i < kDim; ++i) { + vector[i] = std::sin(static_cast(i + row) * 0.17F); + } + rotator->rotate(vector.data(), rotated_data.data() + row * padded_dim); + } + + // A partial batch still owns storage and score outputs for all 32 lanes. + std::vector batch(rabitqlib::BatchDataMap::data_bytes(padded_dim)); + const size_t extra_stride = + rabitqlib::ExDataMap::data_bytes(padded_dim, kExBits); + constexpr size_t kBatchSize = rabitqlib::fastscan::kBatchSize; + std::vector extra(kBatchSize * extra_stride); + std::array estimates{}, bounds{}, products{}; + const auto extra_ip = rabitqlib::select_excode_ipfunc(kExBits); + const float residual_norm = std::sqrt( + rabitqlib::euclidean_sqr(rotated_query.data(), centroid.data(), padded_dim) + ); + const float centroid_ip = + rabitqlib::dot_product(rotated_query.data(), centroid.data(), padded_dim); + for (auto metric : {rabitqlib::METRIC_L2, rabitqlib::METRIC_IP}) { + // This object borrows rotated_query for refinement; keep it alive. + rabitqlib::SplitBatchQuery prepared_query( + rotated_query.data(), padded_dim, kExBits, metric, true + ); + prepared_query.set_g_add(residual_norm, centroid_ip); + for (size_t first = 0; first < kCount; first += kBatchSize) { + const size_t count = std::min(kBatchSize, kCount - first); + rabitqlib::quant::quantize_split_batch( + rotated_data.data() + first * padded_dim, + centroid.data(), + count, + padded_dim, + kExBits, + batch.data(), + extra.data(), + metric + ); + rabitqlib::split_batch_estdist( + batch.data(), + prepared_query, + prepared_query.padded_dim(), + estimates.data(), + bounds.data(), + products.data(), + prepared_query.use_hacc() + ); + // Ignore physical padding lanes in the final, partial batch. + for (size_t row = 0; row < count; ++row) { + const float estimate = rabitqlib::split_distance_boosting( + extra.data() + row * extra_stride, + extra_ip, + prepared_query, + padded_dim, + kExBits, + products[row] + ); + if (!std::isfinite(estimate)) { + return 1; + } + std::cout << (metric == rabitqlib::METRIC_L2 ? "L2 " : "IP ") + << first + row << ": " << estimate << '\n'; + } + } + } + return 0; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/sample/cpp/hnsw_rabitq_indexing.cpp b/sample/cpp/hnsw_rabitq_indexing.cpp index fdd89745..84c8be63 100644 --- a/sample/cpp/hnsw_rabitq_indexing.cpp +++ b/sample/cpp/hnsw_rabitq_indexing.cpp @@ -1,8 +1,14 @@ +#include #include +#include #include #include +#include +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/metric.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" @@ -70,14 +76,14 @@ int run(int argc, char* argv[]) { size_t dim = data.cols(); size_t random_seed = 100; // by default 100 - auto* hnsw = new rabitqlib::hnsw::HierarchicalNSW( + rabitqlib::hnsw::HierarchicalNSW hnsw( num_points, dim, total_bits, m, ef, random_seed, metric_type ); rabitqlib::StopW stopw; stopw.reset(); - hnsw->construct( + hnsw.construct( centroids.rows(), centroids.data(), num_points, @@ -91,7 +97,7 @@ int run(int argc, char* argv[]) { total_time /= 1e6; std::cout << "indexing time = " << total_time << "s" << '\n'; - hnsw->save(index_file); + hnsw.save(index_file); std::cout << "index saved..." << '\n'; diff --git a/sample/cpp/hnsw_rabitq_querying.cpp b/sample/cpp/hnsw_rabitq_querying.cpp index 06487070..45bfdcde 100644 --- a/sample/cpp/hnsw_rabitq_querying.cpp +++ b/sample/cpp/hnsw_rabitq_querying.cpp @@ -1,10 +1,18 @@ +#include +#include +#include #include #include +#include +#include #include +#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" +#include "rabitqlib/utils/tools.hpp" std::vector efs = {10, 20, 40, 50, 60, 80, 100, 150, 170, 190, 200, 250, 300, 400, 500, 600, 700, 800, 1000, 1500, 2000}; diff --git a/sample/cpp/ivf_rabitq_indexing.cpp b/sample/cpp/ivf_rabitq_indexing.cpp index a4fddbb7..f29e6061 100644 --- a/sample/cpp/ivf_rabitq_indexing.cpp +++ b/sample/cpp/ivf_rabitq_indexing.cpp @@ -1,9 +1,14 @@ +#include #include +#include #include #include +#include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/metric.hpp" #include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" diff --git a/sample/cpp/ivf_rabitq_querying.cpp b/sample/cpp/ivf_rabitq_querying.cpp index 39a8fabe..acc21a48 100644 --- a/sample/cpp/ivf_rabitq_querying.cpp +++ b/sample/cpp/ivf_rabitq_querying.cpp @@ -1,10 +1,14 @@ +#include +#include #include #include #include +#include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" #include "rabitqlib/utils/tools.hpp" diff --git a/sample/cpp/kmeans_example.hpp b/sample/cpp/kmeans_example.hpp index 9d9b241b..7a86ad93 100644 --- a/sample/cpp/kmeans_example.hpp +++ b/sample/cpp/kmeans_example.hpp @@ -9,8 +9,9 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/utils/io.hpp" +#include "rabitqlib/utils/path.hpp" namespace kmeans_example { diff --git a/sample/cpp/quantizer.cpp b/sample/cpp/quantizer.cpp deleted file mode 100644 index c0973309..00000000 --- a/sample/cpp/quantizer.cpp +++ /dev/null @@ -1,64 +0,0 @@ -#include -#include - -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/rotator.hpp" - -int run() { - // generate random data - size_t dim = 128; - size_t bit = 4; - - float* data = new float[dim]; - for (size_t i = 0; i < dim; i++) { - data[i] = static_cast(rand()) / static_cast(RAND_MAX); - } - - // choose rotator - rabitqlib::Rotator* rotator = rabitqlib::choose_rotator(dim); - - size_t padded_dim = rotator->size(); - float* rotated_data = new float[padded_dim]; - rotator->rotate(data, rotated_data); - - // print rotated_data - for (size_t i = 0; i < padded_dim; i++) { - std::cout << rotated_data[i] << " "; - } - std::cout << '\n'; - - // quantize - uint8_t* code = new uint8_t[padded_dim]; - float delta = 0; - float vl = 0; - rabitqlib::quant::quantize_scalar(rotated_data, padded_dim, bit, code, delta, vl); - - // [Note: we don't need to store vl as vl = - delta * (2^bit - 1) / 2] - - // reconstruct - float* reconstructed_data = new float[padded_dim]; - rabitqlib::quant::reconstruct_vec(code, delta, vl, padded_dim, reconstructed_data); - - // print reconstructed_data - for (size_t i = 0; i < padded_dim; i++) { - std::cout << reconstructed_data[i] << " "; - } - std::cout << '\n'; - - delete rotator; - delete[] data; - delete[] rotated_data; - delete[] code; - delete[] reconstructed_data; - - return 0; -} - -int main() { - try { - return run(); - } catch (const std::exception& error) { - std::cerr << "Error: " << error.what() << '\n'; - return 1; - } -} diff --git a/sample/cpp/stable_core/CMakeLists.txt b/sample/cpp/stable_core/CMakeLists.txt new file mode 100644 index 00000000..3cd311f2 --- /dev/null +++ b/sample/cpp/stable_core/CMakeLists.txt @@ -0,0 +1,7 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitqlib_stable_core_example LANGUAGES CXX) +find_package(rabitqlib CONFIG REQUIRED COMPONENTS core) +add_executable(stable_core_example main.cpp) +target_link_libraries(stable_core_example PRIVATE rabitqlib::core) +enable_testing() +add_test(NAME stable_core_example COMMAND stable_core_example) diff --git a/sample/cpp/stable_core/main.cpp b/sample/cpp/stable_core/main.cpp new file mode 100644 index 00000000..7a3a5f8f --- /dev/null +++ b/sample/cpp/stable_core/main.cpp @@ -0,0 +1,78 @@ +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/metric.hpp" + +namespace { +constexpr size_t kDim = 65; +using Vector = std::array; + +// Stand-in for the application's existing orthogonal transform. Apply the same +// transform to data, queries and nonzero centroids. The library owns none of it. +Vector external_rotation(const Vector& input) { + Vector output{}; + for (size_t i = 0; i < kDim; ++i) { + output[i] = input[(17 * i) % kDim] * (i % 2 == 0 ? 1.0F : -1.0F); + } + return output; +} +} // namespace + +int main() { + namespace api = rabitqlib::quant::v1; + const api::Parameters parameters{kDim, 4, rabitqlib::METRIC_L2}; + const auto sizes = api::requirements(parameters, api::CodeLayout::UnpackedU16); + // Allocate once during setup, or substitute a block from your memory pool. + std::vector scratch( + (sizes.workspace_bytes + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t) + ); + const api::Buffer workspace{scratch.data(), sizes.workspace_bytes}; + + struct Record { + uint64_t application_id; + std::array code; + }; + std::array records{}; + // Factors can live in separate columns, completely independent of the codes. + std::array add{}, rescale{}, error{}; + Vector raw_query{}; + for (size_t i = 0; i < kDim; ++i) { + raw_query[i] = std::cos(static_cast(i) * 0.11F); + } + const auto query = external_rotation(raw_query); + const auto state = api::prepare_query(parameters, query.data(), nullptr); + // query must stay alive and immutable while state (or any copy) is in use. + for (size_t row = 0; row < records.size(); ++row) { + Vector input{}; + for (size_t i = 0; i < kDim; ++i) { + input[i] = std::sin(static_cast(i + row) * 0.17F); + } + const auto transformed = external_rotation(input); + records[row].application_id = 100 + row; + auto& code = records[row].code; + const auto factors = api::encode( + parameters, + transformed.data(), + nullptr, // A zero centroid; pass a transformed centroid for residual coding. + {api::CodeLayout::UnpackedU16, {code.data(), code.size()}, {}}, + workspace + ); + add[row] = factors.add; + rescale[row] = factors.rescale; + error[row] = factors.error; + const float distance = api::compute( + state, + {api::CodeLayout::UnpackedU16, {code.data(), code.size()}, {}}, + {add[row], rescale[row], error[row]} + ); + if (!std::isfinite(distance)) { + return 1; + } + std::cout << records[row].application_id << ": " << distance << '\n'; + } +} diff --git a/sample/cpp/symqg_indexing.cpp b/sample/cpp/symqg_indexing.cpp index 6853babb..aed9a5ee 100644 --- a/sample/cpp/symqg_indexing.cpp +++ b/sample/cpp/symqg_indexing.cpp @@ -1,15 +1,21 @@ +#include #include +#include #include #include #include #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" +#include "rabitqlib/utils/tools.hpp" using PID = rabitqlib::PID; using index_type = rabitqlib::symqg::QuantizedGraph; diff --git a/sample/cpp/symqg_querying.cpp b/sample/cpp/symqg_querying.cpp index 6cc7c91f..6a09e803 100644 --- a/sample/cpp/symqg_querying.cpp +++ b/sample/cpp/symqg_querying.cpp @@ -1,10 +1,15 @@ +#include +#include #include #include +#include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" #include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/stopw.hpp" +#include "rabitqlib/utils/tools.hpp" using PID = rabitqlib::PID; using index_type = rabitqlib::symqg::QuantizedGraph; diff --git a/sample/python/README.md b/sample/python/README.md index 85977a04..c481dd3a 100644 --- a/sample/python/README.md +++ b/sample/python/README.md @@ -1,10 +1,10 @@ # Python examples Run these commands from the repository root in your project Python environment -with `rabitqlib>=0.5.2` and NumPy installed. Clustering and indexing both use RaBitQ -Library; FAISS is needed only for `compare_with_faiss.py`. When running these -examples from a checkout, [install the library from that same checkout](../../CONTRIBUTING.md#python-changes) -so the examples and native extension use matching APIs. +with `rabitqlib>=0.7.0` and NumPy installed. When working from source, +[install the library from the same checkout](../../CONTRIBUTING.md#python-changes) +so the examples and native extension use matching APIs. Clustering and indexing +both use RaBitQ Library; FAISS is needed only for `compare_with_faiss.py`. ## IVF: cluster, then index @@ -36,7 +36,7 @@ behavior. For inner product, pass `--metric ip` to both stages and provide normalized training vectors. Use queries and ground truth prepared for the same metric. `--num-threads 0` uses the library's available logical CPU count; each stage accepts -its own limit. See [threading](../../docs/docs/quick_start.md#threading-and-file-paths) +its own limit. See [threading](../../docs/docs/index/usage.md#threading-and-file-paths) for affinity and OpenMP placement behavior. ## HNSW diff --git a/sample/python/hnsw_rabitq_querying.py b/sample/python/hnsw_rabitq_querying.py index d605f081..7f276c1c 100644 --- a/sample/python/hnsw_rabitq_querying.py +++ b/sample/python/hnsw_rabitq_querying.py @@ -18,7 +18,7 @@ def main(args=None) -> None: # 1. Load queries and ground truth - queries = read_fvecs(args.query_file) + queries = np.ascontiguousarray(read_fvecs(args.query_file)) gt = read_ivecs(args.gt_file) nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") diff --git a/sample/python/ivf_rabitq_querying.py b/sample/python/ivf_rabitq_querying.py index 709093a1..523692a0 100644 --- a/sample/python/ivf_rabitq_querying.py +++ b/sample/python/ivf_rabitq_querying.py @@ -18,7 +18,7 @@ def main(args=None) -> None: # 1. Load queries and ground truth - queries = read_fvecs(args.query_file) + queries = np.ascontiguousarray(read_fvecs(args.query_file)) gt = read_ivecs(args.gt_file) nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index cd00659e..b837da7f 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -19,7 +19,7 @@ def main(args=None) -> None: # 1. Load queries and ground truth - queries = read_fvecs(args.query_file) + queries = np.ascontiguousarray(read_fvecs(args.query_file)) gt = read_ivecs(args.gt_file) nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") diff --git a/scripts/check-includes.sh b/scripts/check-includes.sh index 9f57e8c2..f75ab33b 100755 --- a/scripts/check-includes.sh +++ b/scripts/check-includes.sh @@ -19,49 +19,190 @@ if [[ ! -f "$include_build/compile_commands.json" ]]; then exit 1 fi -check_file() { - local file="$1" - # Tracked deletions remain in git ls-files until staged. - [[ -f "$file" ]] || return 0 - local status=0 - local report - report="$(mktemp)" - # Isolate this check from the project's general clang-tidy configuration. - # Vendored snapshots lack export annotations on their public headers. - # Do not replace those headers with Eigen/hnswlib implementation includes. - local args=(--config='{CheckOptions: {misc-include-cleaner.IgnoreHeaders: "rabitqlib/third/Eigen/src/.*;rabitqlib/third/hnswlib/(hnswalg|space_l2)[.]h"}}' --checks='-*,misc-include-cleaner' - --warnings-as-errors='misc-include-cleaner' --quiet) - if [[ "$file" == src/*.cpp ]]; then - "$CLANG_TIDY" "${args[@]}" -p "$include_build" "$file" >"$report" 2>&1 || status=1 - else - local flags=(-x c++ -std=c++17 -fopenmp -I"$include_root/include" - -Wno-pragma-once-outside-header) - if [[ "$file" == src/* ]]; then - # Private kernels share AVX2 and AVX-512 implementations. Check both - # preprocessor paths, including the popcount-enabled path. - if [[ "$file" != *avx512* ]]; then - "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" \ - -mavx2 -mfma >"$report" 2>&1 || status=1 - fi - if [[ "$file" != *avx2* ]]; then - "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" \ - -mavx2 -mfma -mavx512f -mavx512bw -mavx512dq \ - -mavx512vpopcntdq >>"$report" 2>&1 || status=1 - fi - else - "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" >"$report" 2>&1 || status=1 - fi - fi - echo "Checking $file" - cat "$report" - rm -f "$report" - return "$status" -} -export -f check_file +# Reuse configured flags for optional dependencies and cross-compilation. Headers +# need to be main files because include-cleaner does not diagnose included files. +"${PYTHON:-python3}" - <<'PYCODE' +import concurrent.futures +import json +import os +import re +from pathlib import Path +import shlex +import subprocess +import sys -# The child shell expands its positional argument. -# shellcheck disable=SC2016 -git ls-files --cached --others --exclude-standard -z -- 'src/*.cpp' 'src/*.hpp' 'include/rabitqlib/*.hpp' \ - ':(exclude)include/rabitqlib/third/**' \ - | sort -zu \ - | xargs -0 -r -n 1 -P "${INCLUDE_JOBS:-2}" bash -c 'check_file "$1"' _ +root = Path(os.environ["include_root"]) +build = Path(os.environ["include_build"]) +tidy = os.environ["CLANG_TIDY"] +entries = {} +for entry in json.loads((build / "compile_commands.json").read_text()): + file = (Path(entry["directory"]) / entry["file"]).resolve() + if file.is_relative_to(root) and not file.is_relative_to( + root / "include/rabitqlib/third" + ): + entries.setdefault(file, entry) +if not entries: + sys.exit("error: no first-party sources in the compilation database") + +# Public vendor entry points lack complete IWYU export annotations. Keep their +# implementation headers out of callers, including PYBIND11_MODULE's provider. +config = json.dumps( + { + "Checks": "-*,misc-include-cleaner,readability-duplicate-include", + "WarningsAsErrors": "misc-include-cleaner,readability-duplicate-include", + "CheckOptions": { + "misc-include-cleaner.IgnoreHeaders": "rabitqlib/third/Eigen/src/.*;" + "rabitqlib/third/hnswlib/(hnswalg|space_l2)[.]h;" + "pybind11/detail/common[.]h" + }, + } +) + + +def flags_for(entry): + args = entry.get("arguments") or shlex.split(entry["command"]) + flags = [] + args = iter(args[1:]) + for arg in args: + if arg == "-o": + next(args) + elif ( + arg != "-c" + and (Path(entry["directory"]) / arg).resolve() + != (Path(entry["directory"]) / entry["file"]).resolve() + ): + flags.append(arg) + return flags + ["-x", "c++-header"] + + +def representative(prefix): + return next((e for f, e in entries.items() if f.is_relative_to(root / prefix)), None) + + +base = entries.get(root / "src/core/api.cpp") +if base is None: + sys.exit("error: configure the library targets before checking includes") +# Private headers inherit the ISA flags of their configured source consumers. +# Follow local includes so shared kernel headers get every configured backend. +private_flags = {} +for source, entry in entries.items(): + if not source.is_relative_to(root / "src"): + continue + pending, seen = [source], set() + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + for include in re.findall( + r'^\s*#\s*include\s*"([^"]+)"', current.read_text(), re.MULTILINE + ): + header = (current.parent / include).resolve() + if header.is_relative_to(root / "src") and header.is_file(): + private_flags.setdefault(header, set()).add(tuple(flags_for(entry))) + pending.append(header) +has_index = representative("src/index") is not None +has_fastscan = representative("src/fastscan") is not None +files = ( + subprocess.check_output( + [ + "git", + "ls-files", + "--cached", + "--others", + "--exclude-standard", + "-z", + "--", + "*.cpp", + "*.cc", + "*.cxx", + "*.h", + "*.hpp", + ":(exclude)include/rabitqlib/third/**", + ], + cwd=root, + ) + .decode() + .split("\0") +) +jobs = [] +skipped = [] +for name in sorted(set(files) - {""}): + file = root / name + if not file.is_file(): + continue + if file.suffix not in {".h", ".hpp"}: + if file in entries: + jobs.append((file, "configured", ["-p", str(build)])) + else: + skipped.append((name, "not in compilation database")) + continue + if ( + name.startswith( + ("include/rabitqlib/index/", "include/rabitqlib/clustering/", "src/index/") + ) + and not has_index + or name.startswith(("include/rabitqlib/fastscan/", "src/fastscan/")) + and not has_fastscan + ): + skipped.append((name, "component disabled")) + continue + if name.startswith("src/"): + contexts = sorted(private_flags.get(file, set())) + if not contexts: + skipped.append((name, "no configured source includes this private header")) + for index, flags in enumerate(contexts): + jobs.append((file, f"private context {index + 1}", ["--"] + list(flags))) + continue + entry = base + for prefix in ("python_bindings", "tests", "sample", "benchmarks"): + if name.startswith(prefix + "/"): + entry = representative(prefix) + break + if entry is None: + skipped.append((name, "owning target not configured")) + continue + flags = flags_for(entry) + jobs.append((file, "header", ["--"] + flags)) + if name.startswith( + ( + "include/rabitqlib/core/", + "include/rabitqlib/fastscan/", + "include/rabitqlib/utils/", + ) + ): + serial = [ + f for f in flags if not f.startswith(("-fopenmp", "-DRABITQ_USE_OPENMP=")) + ] + jobs.append((file, "OpenMP OFF", ["--"] + serial + ["-DRABITQ_USE_OPENMP=0"])) + + +def check(job): + file, mode, flags = job + result = subprocess.run( + [tidy, "--config=" + config, "--quiet", str(file)] + flags, + cwd=root, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + return file, mode, result + + +failed = 0 +with concurrent.futures.ThreadPoolExecutor( + max_workers=int(os.environ.get("INCLUDE_JOBS", "2")) +) as pool: + for file, mode, result in pool.map(check, jobs): + print(f"Checking {file.relative_to(root)} [{mode}]", flush=True) + if result.stdout: + print(result.stdout, end="", flush=True) + failed += result.returncode != 0 +for file, reason in skipped: + print(f"Skipped {file}: {reason}") +print( + f"Checked {len({job[0] for job in jobs})} files in {len(jobs)} contexts; " + f"{failed} failed, {len(skipped)} skipped." +) +sys.exit(1 if failed else 0) +PYCODE diff --git a/src/core/api.cpp b/src/core/api.cpp new file mode 100644 index 00000000..44b375dc --- /dev/null +++ b/src/core/api.cpp @@ -0,0 +1,393 @@ +#include "rabitqlib/core/api.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/detail/scratch.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/utils/bitops.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::quant::v1 { +namespace { +constexpr size_t kWorkspaceAlignment = alignof(std::max_align_t); + +void validate(const Parameters& p) { + if (p.dim < 2) { + throw std::invalid_argument("dim must be at least 2"); + } + if (p.dim > static_cast(std::numeric_limits::max())) { + throw std::length_error("dim exceeds the supported index range"); + } + if (p.total_bits < 1 || p.total_bits > 9) { + throw std::invalid_argument("total_bits must be in [1, 9]"); + } + validate_metric_type(p.metric); + if (!std::isfinite(p.t_const) || p.t_const > std::numeric_limits::max() / 2.0) { + throw std::invalid_argument("t_const must be finite and at most INT_MAX / 2"); + } +} + +size_t multiply(size_t count, size_t width) { + if (width == 0) { + return 0; + } + if (count > std::numeric_limits::max() / width) { + throw std::length_error("Quantization buffer size overflows size_t"); + } + return count * width; +} + +[[noreturn]] void invalid_buffer(const char* message) { + throw std::invalid_argument(message); +} + +// All alignments are internal constants. Keep exception construction out of the +// valid-buffer path so callers can inline the checks and remove byte alignment work. +template +void check_buffer(ConstBuffer buffer, size_t required) { + static_assert(Alignment > 0); + if (required == 0) { + return; + } + if (buffer.data == nullptr || buffer.bytes < required) { + invalid_buffer("Buffer is null or smaller than the required size"); + } + const auto address = reinterpret_cast(buffer.data); + if (address % Alignment != 0) { + invalid_buffer("Buffer does not satisfy the required alignment"); + } + if (required > std::numeric_limits::max() - address) { + invalid_buffer("Buffer address range overflows uintptr_t"); + } +} + +void disjoint(const void* a, size_t a_size, const void* b, size_t b_size) { + if (a_size == 0 || b_size == 0) { + return; + } + const auto first = reinterpret_cast(a); + const auto second = reinterpret_cast(b); + if (first < second + b_size && second < first + a_size) { + 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, multiply(dim, sizeof(float))}, multiply(dim, sizeof(float)) + ); + for (size_t i = 0; i < dim; ++i) { + if (!std::isfinite(input[i])) { + throw std::invalid_argument("Input vectors must contain only finite values"); + } + } +} + +// Parameters and arithmetic ranges were validated by requirements/prepare_query. +std::pair code_sizes(const Parameters& p, CodeLayout layout) { + switch (layout) { + case CodeLayout::UnpackedU16: + return {p.dim * sizeof(uint16_t), 0}; + case CodeLayout::SplitPacked: + if (p.dim % 32 != 0) { + throw std::invalid_argument( + "SplitPacked requires dim to be a multiple of 32" + ); + } + return {p.dim / 8, (p.dim / 8) * (p.total_bits - 1)}; + default: + throw std::invalid_argument("Unsupported code layout"); + } +} + +void check_code(CodeView code, const Requirements& r) { + check_buffer(code.codes, r.code_bytes); + check_buffer(code.extra, r.extra_bytes); +} + +CodeView view(CodeOutput code) { + return { + code.layout, + {code.codes.data, code.codes.bytes}, + {code.extra.data, code.extra.bytes}}; +} + +void check_output(CodeOutput code, const Requirements& r, Buffer workspace) { + check_code(view(code), r); + check_buffer({workspace.data, workspace.bytes}, r.workspace_bytes); + disjoint(code.codes.data, r.code_bytes, code.extra.data, r.extra_bytes); + disjoint(code.codes.data, r.code_bytes, workspace.data, r.workspace_bytes); + disjoint(code.extra.data, r.extra_bytes, workspace.data, r.workspace_bytes); +} + +void check_input_overlap( + const float* input, + size_t dim, + CodeOutput output, + const Requirements& r, + Buffer workspace +) { + if (input == nullptr) { + return; + } + const size_t bytes = multiply(dim, sizeof(float)); + disjoint(input, bytes, output.codes.data, r.code_bytes); + disjoint(input, bytes, output.extra.data, r.extra_bytes); + disjoint(input, bytes, workspace.data, r.workspace_bytes); +} + +// One typed allocation list serves both requirements() and actual storage binding. +class ScratchStorage { + public: + rabitq_impl::Scratch scratch; + uint16_t* codes; + uint8_t* extra; + float* zero; + + explicit ScratchStorage(rabitq_impl::ScratchAllocator& allocator) + : scratch{allocator.array(), allocator.array(), + allocator.array(), allocator.array(), + allocator.array(), allocator.array(), + allocator.array(), allocator.resource()} + , codes(allocator.array()) + , extra(allocator.array()) + , zero(allocator.array()) {} +}; + +uint16_t read_u16(const uint8_t* bytes) { + return static_cast(bytes[0] | (uint16_t{bytes[1]} << 8U)); +} + +void read_codes(const Parameters& p, CodeView input, ScratchStorage& s) { + const auto* bytes = static_cast(input.codes.data); + if (input.layout == CodeLayout::UnpackedU16) { + const auto maximum = static_cast((1U << p.total_bits) - 1); + for (size_t i = 0; i < p.dim; ++i) { + s.codes[i] = read_u16(bytes + 2 * i); + if (s.codes[i] > maximum) { + throw std::invalid_argument("Unpacked code exceeds the configured bit width" + ); + } + } + return; + } + const size_t ex_bits = p.total_bits - 1; + if (ex_bits > 0) { + rabitq_impl::ex_bits::unpacking_rabitqplus_code( + static_cast(input.extra.data), s.extra, p.dim, ex_bits + ); + } + // Sign words retain the existing full-64 / final-32 coordinate convention. + for (size_t start = 0; start < p.dim; start += 64) { + const size_t width = std::min(size_t{64}, p.dim - start); + const uint64_t word = bitops::load_word(bytes + start / 8, width / 8); + for (size_t i = 0; i < width; ++i) { + const auto sign = static_cast((word >> (width - 1 - i)) & 1U); + s.codes[start + i] = static_cast( + (sign << ex_bits) | (ex_bits == 0 ? 0 : s.extra[start + i]) + ); + } + } +} + +void write_codes(const Parameters& p, CodeOutput output, ScratchStorage& s) { + auto* bytes = static_cast(output.codes.data); + if (output.layout == CodeLayout::UnpackedU16) { + for (size_t i = 0; i < p.dim; ++i) { + bytes[2 * i] = static_cast(s.codes[i]); + bytes[2 * i + 1] = static_cast(s.codes[i] >> 8U); + } + return; + } + const size_t ex_bits = p.total_bits - 1; + const auto mask = static_cast((1U << ex_bits) - 1); + for (size_t i = 0; i < p.dim; ++i) { + s.scratch.binary[i] = s.codes[i] >> ex_bits; + s.extra[i] = static_cast(s.codes[i] & mask); + } + pack_binary_to_bytes(s.scratch.binary, bytes, p.dim); + if (ex_bits > 0) { + rabitq_impl::ex_bits::packing_rabitqplus_code( + s.extra, static_cast(output.extra.data), p.dim, ex_bits + ); + } +} +} // namespace + +Requirements requirements(const Parameters& p, CodeLayout layout) { + validate(p); + // Derive the budget once from the same declarations used by encoding and + // exact fallback. No array is allocated in this sizing pass. + static const auto budget = [] { + rabitq_impl::ScratchAllocator sizing; + [[maybe_unused]] const ScratchStorage arrays(sizing); + (void)rabitq_impl::event_sweep_buffers(sizing); + return sizing; + }(); + Requirements r{0, 0, budget.bytes(p.dim), kWorkspaceAlignment}; + const auto sizes = code_sizes(p, layout); + r.code_bytes = sizes.first; + r.extra_bytes = sizes.second; + return r; +} + +Factors encode( + const Parameters& p, + const float* input, + const float* centroid, + CodeOutput output, + Buffer workspace +) { + const Requirements r = requirements(p, output.layout); + check_input(input, p.dim); + if (centroid != nullptr) { + check_input(centroid, p.dim); + } + check_output(output, r, workspace); + check_input_overlap(input, p.dim, output, r, workspace); + check_input_overlap(centroid, p.dim, output, r, workspace); + std::pmr::monotonic_buffer_resource resource( + workspace.data, r.workspace_bytes, std::pmr::null_memory_resource() + ); + rabitq_impl::ScratchAllocator allocator(&resource, p.dim); + ScratchStorage s(allocator); + if (centroid == nullptr) { + std::fill_n(s.zero, p.dim, 0.0F); + centroid = s.zero; + } + // Reject overflow before the quantizer's float-to-integer conversions. + RowMajorArrayMap residual(s.scratch.residual, 1, static_cast(p.dim)); + residual = ConstRowMajorArrayMap(input, 1, static_cast(p.dim)) - + ConstRowMajorArrayMap(centroid, 1, static_cast(p.dim)); + if (!std::isfinite(l2norm_sqr(s.scratch.residual, p.dim))) { + throw std::overflow_error("Residual squared norm is not representable in float32"); + } + Factors factors; + rabitq_impl::total_bits::rabitq_full_impl_with_scratch( + input, + centroid, + p.dim, + p.total_bits, + s.codes, + factors.add, + factors.rescale, + factors.error, + p.metric, + p.t_const, + &s.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"); + } + write_codes(p, output, s); + return factors; +} + +void convert_codes( + const Parameters& p, CodeView source, CodeOutput destination, Buffer workspace +) { + const Requirements input = requirements(p, source.layout); + const Requirements output = requirements(p, destination.layout); + check_code(source, input); + check_output(destination, output, workspace); + for (const auto& range : + {std::pair{source.codes.data, input.code_bytes}, + {source.extra.data, input.extra_bytes}}) { + disjoint(range.first, range.second, destination.codes.data, output.code_bytes); + disjoint(range.first, range.second, destination.extra.data, output.extra_bytes); + disjoint(range.first, range.second, workspace.data, output.workspace_bytes); + } + std::pmr::monotonic_buffer_resource resource( + workspace.data, output.workspace_bytes, std::pmr::null_memory_resource() + ); + rabitq_impl::ScratchAllocator allocator(&resource, p.dim); + ScratchStorage s(allocator); + read_codes(p, source, s); + write_codes(p, destination, s); +} + +QueryState prepare_query(const Parameters& p, const float* query, const float* centroid) { + (void)requirements(p, CodeLayout::UnpackedU16); + check_input(query, p.dim); + if (centroid != nullptr) { + check_input(centroid, p.dim); + } + QueryState state; + state.query_ = query; + state.dim_ = p.dim; + state.total_bits_ = p.total_bits; + state.metric_ = p.metric; + state.extra_ip_ = select_excode_ipfunc(p.total_bits - 1); + state.k1xsumq_ = std::accumulate(query, query + p.dim, 0.0F) * -0.5F; + state.g_add_ = + p.metric == METRIC_L2 + ? (centroid == nullptr ? l2norm_sqr(query, p.dim) + : euclidean_sqr(query, centroid, p.dim)) + : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, p.dim)); + if (!std::isfinite(state.g_add_) || !std::isfinite(state.k1xsumq_)) { + throw std::overflow_error("Query factors are not representable in float32"); + } + return state; +} + +float QueryState::compute_checked( + const QueryState& query, const CodeView& code, const Factors& factors +) { + if (query.query_ == nullptr) { + throw std::invalid_argument("QueryState has not been prepared"); + } + const Parameters p{query.dim_, query.total_bits_, query.metric_}; + const auto sizes = code_sizes(p, code.layout); + check_buffer(code.codes, sizes.first); + check_buffer(code.extra, sizes.second); + if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || + !std::isfinite(factors.error)) { + throw std::invalid_argument("Quantization factors must be finite"); + } + float ip = 0; + if (code.layout == CodeLayout::UnpackedU16) { + const auto* bytes = static_cast(code.codes.data); + const auto maximum = static_cast((1U << p.total_bits) - 1); + for (size_t i = 0; i < p.dim; ++i) { + const uint16_t value = read_u16(bytes + 2 * i); + if (value > maximum) { + throw std::invalid_argument("Unpacked code exceeds the configured bit width" + ); + } + ip += query.query_[i] * static_cast(value); + } + } else { + const size_t ex_bits = p.total_bits - 1; + ip = + static_cast(1U << ex_bits) * + mask_ip_x0_q(query.query_, static_cast(code.codes.data), p.dim); + if (ex_bits > 0) { + ip += query.extra_ip_( + query.query_, static_cast(code.extra.data), p.dim + ); + } + } + const float distance = quant::detail::full_est_dist_from_ip( + ip, p.total_bits, factors.add, factors.rescale, query.g_add_, query.k1xsumq_ + ); + if (!std::isfinite(distance)) { + throw std::overflow_error("Estimated distance is not representable in float32"); + } + return distance; +} +} // namespace rabitqlib::quant::v1 diff --git a/src/simd/estimator_avx2.cpp b/src/fastscan/estimator_avx2.cpp similarity index 70% rename from src/simd/estimator_avx2.cpp rename to src/fastscan/estimator_avx2.cpp index 8dafabff..0be3878e 100644 --- a/src/simd/estimator_avx2.cpp +++ b/src/fastscan/estimator_avx2.cpp @@ -4,7 +4,9 @@ #include #include "estimator_kernels.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" namespace rabitqlib::simd { namespace { @@ -29,7 +31,9 @@ void split_batch_estdist_avx2( float* ip_x0_qr, bool use_hacc ) { - split_batch_estdist_impl( + // Explicit backend entry points also accept query LUTs prepared by the host + // backend. In particular, AVX2 correction can receive AVX-512 HACC bytes. + split_batch_estdist_impl( batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc ); } @@ -40,7 +44,9 @@ void qg_batch_estdist_avx2( size_t padded_dim, float* est_distance ) { - qg_batch_estdist_impl(batch_data, q_obj, padded_dim, est_distance); + qg_batch_estdist_impl( + batch_data, q_obj, padded_dim, est_distance + ); } uint32_t qg_candidate_mask_avx2(const float* distances, float threshold) { @@ -54,7 +60,9 @@ uint32_t qg_batch_estdist_mask_avx2( float* est_distance, float threshold ) { - qg_batch_estdist_impl(batch_data, q_obj, padded_dim, est_distance); + qg_batch_estdist_impl( + batch_data, q_obj, padded_dim, est_distance + ); return candidate_mask_impl(est_distance, threshold); } diff --git a/src/fastscan/estimator_avx512.cpp b/src/fastscan/estimator_avx512.cpp new file mode 100644 index 00000000..32550968 --- /dev/null +++ b/src/fastscan/estimator_avx512.cpp @@ -0,0 +1,47 @@ +#include +#include + +#include "estimator_kernels.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" + +namespace rabitqlib::simd { +void split_batch_estdist_avx512( + const char* batch_data, + const SplitBatchQuery& q_obj, + size_t padded_dim, + float* est_distance, + float* low_distance, + float* ip_x0_qr, + bool use_hacc +) { + if (padded_dim == 0 || padded_dim % 16 != 0) { + throw std::invalid_argument( + use_hacc ? "high-accuracy FastScan dimension must be a positive multiple of 16" + : "FastScan dimension must be a positive multiple of 16" + ); + } + // This backend already requires the highest FastScan ISA. Its LUT format + // matches AVX-512 accumulation, so a second dispatch is unnecessary. + split_batch_estdist_impl< + fastscan::simd::accumulate_avx512, + fastscan::simd::accumulate_hacc_avx512>( + batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + ); +} + +void qg_batch_estdist_avx512( + const char* batch_data, + const BatchQuery& q_obj, + size_t padded_dim, + float* est_distance +) { + if (padded_dim == 0 || padded_dim % 16 != 0) { + throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); + } + qg_batch_estdist_impl( + batch_data, q_obj, padded_dim, est_distance + ); +} + +} // namespace rabitqlib::simd diff --git a/src/simd/estimator_generic.cpp b/src/fastscan/estimator_generic.cpp similarity index 67% rename from src/simd/estimator_generic.cpp rename to src/fastscan/estimator_generic.cpp index 9050e9f8..cf7cf82f 100644 --- a/src/simd/estimator_generic.cpp +++ b/src/fastscan/estimator_generic.cpp @@ -2,7 +2,9 @@ #include #include "estimator_kernels.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" namespace rabitqlib::simd { namespace { @@ -24,7 +26,9 @@ void split_batch_estdist_generic( float* ip_x0_qr, bool use_hacc ) { - split_batch_estdist_impl( + // Explicit backend entry points also accept query LUTs prepared by the host + // backend. In particular, AVX2 correction can receive AVX-512 HACC bytes. + split_batch_estdist_impl( batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc ); } @@ -35,7 +39,9 @@ void qg_batch_estdist_generic( size_t padded_dim, float* est_distance ) { - qg_batch_estdist_impl(batch_data, q_obj, padded_dim, est_distance); + qg_batch_estdist_impl( + batch_data, q_obj, padded_dim, est_distance + ); } uint32_t qg_candidate_mask_generic(const float* distances, float threshold) { @@ -49,7 +55,9 @@ uint32_t qg_batch_estdist_mask_generic( float* est_distance, float threshold ) { - qg_batch_estdist_impl(batch_data, q_obj, padded_dim, est_distance); + qg_batch_estdist_impl( + batch_data, q_obj, padded_dim, est_distance + ); return candidate_mask_impl(est_distance, threshold); } diff --git a/src/simd/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp similarity index 85% rename from src/simd/estimator_kernels.hpp rename to src/fastscan/estimator_kernels.hpp index cc39862a..fb44a6c6 100644 --- a/src/simd/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -4,16 +4,16 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/highacc_fastscan.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" +#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::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. +template inline void split_batch_estdist_impl( const char* batch_data, const SplitBatchQuery& q_obj, @@ -27,13 +27,9 @@ inline void split_batch_estdist_impl( std::array accu_values{}; RowMajorArrayMap accu_arr(accu_values.data(), 1, fastscan::kBatchSize); if (use_hacc) { - fastscan::accumulate_hacc( - cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim - ); + AccumulateHacc(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); } else { - fastscan::accumulate( - cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim - ); + Accumulate(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); } std::array f_add_values; @@ -62,6 +58,7 @@ inline void split_batch_estdist_impl( low_dist_arr = est_dist_arr - f_error_arr * q_obj.g_error(); } +template inline void qg_batch_estdist_impl( const char* batch_data, const BatchQuery& q_obj, @@ -71,7 +68,7 @@ inline void qg_batch_estdist_impl( using T = float; ConstQGBatchDataMap cur_batch(batch_data, padded_dim); std::array accu_values{}; - fastscan::accumulate(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); + Accumulate(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); ConstRowMajorArrayMap ip_arr(accu_values.data(), 1, fastscan::kBatchSize); std::array f_add_values; diff --git a/src/simd/fastscan_avx2.cpp b/src/fastscan/fastscan_avx2.cpp similarity index 99% rename from src/simd/fastscan_avx2.cpp rename to src/fastscan/fastscan_avx2.cpp index 96d30a40..d10db4cc 100644 --- a/src/simd/fastscan_avx2.cpp +++ b/src/fastscan/fastscan_avx2.cpp @@ -7,7 +7,7 @@ #include #include -#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" namespace rabitqlib::fastscan::simd { diff --git a/src/simd/fastscan_avx512.cpp b/src/fastscan/fastscan_avx512.cpp similarity index 91% rename from src/simd/fastscan_avx512.cpp rename to src/fastscan/fastscan_avx512.cpp index af2696bc..0d494419 100644 --- a/src/simd/fastscan_avx512.cpp +++ b/src/fastscan/fastscan_avx512.cpp @@ -7,7 +7,7 @@ #include #include -#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" namespace rabitqlib::fastscan::simd { @@ -66,11 +66,10 @@ void accumulate_chunk_avx512( res_lo = _mm512_shuffle_epi8(lut, lo); // get the target value in lookup table res_hi = _mm512_shuffle_epi8(lut, hi); - // since values in lookup table are represented as i8, we add them as i16 to avoid - // overflow. Since the data order is 0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, - // 7, 15, accu0 accumulates for vec 8 to 15 (the upper 8 bits need to be updated - // since they stored useless info of vec 0 to 7) accu1 accumulates for vec 0 to 7 - // similar for accu2 and accu3 + // Byte pairs follow kPerm0: (0,8), (1,9), ..., (7,15). + // accu0 sums the pairs; accu1 separately sums their high bytes (vectors 8..15). + // Subtracting accu1 << 8 below leaves vectors 0..7 in accu0. + // accu2/accu3 use the same scheme for vectors 16..23 and 24..31. accu0 = _mm512_add_epi16(accu0, res_lo); accu1 = _mm512_add_epi16(accu1, _mm512_srli_epi16(res_lo, 8)); accu2 = _mm512_add_epi16(accu2, res_hi); @@ -80,12 +79,9 @@ void accumulate_chunk_avx512( accu0 = _mm512_sub_epi16(accu0, _mm512_slli_epi16(accu1, 8)); accu2 = _mm512_sub_epi16(accu2, _mm512_slli_epi16(accu3, 8)); - // At this point, we already have the correct accumulating result (accu0: 8-15, accu1: - // 0-7, accu2: 16-23, accu3: 24-31), but we still need to write them back to RAM. Also, - // each accu contains 4 lines of __m128i and we need to sum them together to get the - // final results. 512/16=32, so we can use one __m512i to contain all results. The - // following codes are designed for this purpose. For detailed information, please check - // the SIMD documentation. + // Corrected groups: accu0 = 0..7, accu1 = 8..15, accu2 = 16..23, accu3 = 24..31. + // Sum the four 128-bit partials per group and restore vector order 0..31, + // then widen the 16-bit sums to the int32_t output. __m512i ret1 = _mm512_add_epi16( _mm512_mask_blend_epi64(0b11110000, accu0, accu1), _mm512_shuffle_i64x2(accu0, accu1, 0b01001110) diff --git a/src/fastscan/fastscan_backend.hpp b/src/fastscan/fastscan_backend.hpp new file mode 100644 index 00000000..e71a3b06 --- /dev/null +++ b/src/fastscan/fastscan_backend.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" + +namespace rabitqlib::fastscan::detail { + +// Standard LUTs use consecutive 16-entry tables on every backend. HACC stores +// low/high byte planes in groups of two tables (32 bytes) or four (64 bytes). +// This describes transient query scratch, never persisted codes. +enum class HaccLutLayout { Pair32, Quad64 }; + +// One immutable selection owns the producer/consumer pairing. Keep this private: +// no backend pointers or ISA-specific layouts become public query state. +struct Backend { + const char* name; + HaccLutLayout hacc_layout; + decltype(&simd::pack_lut_generic) pack_lut; + decltype(&simd::transfer_lut_hacc_generic) transfer_hacc; + decltype(&simd::accumulate_generic) accumulate; + decltype(&simd::accumulate_hacc_generic) accumulate_hacc; + decltype(&rabitqlib::simd::qg_batch_estdist_generic) qg; + decltype(&rabitqlib::simd::qg_batch_estdist_mask_generic) qg_mask; + decltype(&rabitqlib::simd::split_batch_estdist_generic) split; +}; + +// Call only during wrapper initialization. Warm wrappers cache the member pointer +// locally, avoiding an extra getter or table lookup on every scored batch. +const Backend& selected_backend(); + +} // namespace rabitqlib::fastscan::detail diff --git a/src/fastscan/fastscan_dispatch.cpp b/src/fastscan/fastscan_dispatch.cpp new file mode 100644 index 00000000..71267d05 --- /dev/null +++ b/src/fastscan/fastscan_dispatch.cpp @@ -0,0 +1,170 @@ +#include +#include +#include + +#include "../simd/dispatch_resolver.hpp" +#include "fastscan_backend.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#if defined(__aarch64__) || defined(_M_ARM64) +#include "rabitqlib/utils/cpu_features.hpp" +#endif + +namespace rabitqlib::fastscan::detail { +namespace { +constexpr Backend kGenericBackend{ + "generic", + HaccLutLayout::Pair32, + simd::pack_lut_generic, + simd::transfer_lut_hacc_generic, + simd::accumulate_unsupported, + simd::accumulate_hacc_generic, + rabitqlib::simd::qg_batch_estdist_generic, + rabitqlib::simd::qg_batch_estdist_mask_generic, + rabitqlib::simd::split_batch_estdist_generic}; + +#if defined(__x86_64__) || defined(_M_X64) +constexpr Backend kAvx2Backend{ + "avx2", + HaccLutLayout::Pair32, + simd::pack_lut_avx2, + simd::transfer_lut_hacc_avx2, + simd::accumulate_avx2, + simd::accumulate_hacc_avx2, + rabitqlib::simd::qg_batch_estdist_avx2, + rabitqlib::simd::qg_batch_estdist_mask_avx2, + rabitqlib::simd::split_batch_estdist_avx2}; +constexpr Backend kAvx512Backend{ + "avx512", + HaccLutLayout::Quad64, + simd::pack_lut_avx512, + simd::transfer_lut_hacc_avx512, + simd::accumulate_avx512, + simd::accumulate_hacc_avx512, + rabitqlib::simd::qg_batch_estdist_avx512, + rabitqlib::simd::qg_batch_estdist_mask_avx2, + rabitqlib::simd::split_batch_estdist_avx512}; +#endif +#if defined(__aarch64__) || defined(_M_ARM64) +constexpr Backend kNeonBackend{ + "neon", + HaccLutLayout::Pair32, + simd::pack_lut_neon, + simd::transfer_lut_hacc_generic, + simd::accumulate_neon, + simd::accumulate_hacc_neon, + rabitqlib::simd::qg_batch_estdist_generic, + rabitqlib::simd::qg_batch_estdist_mask_generic, + rabitqlib::simd::split_batch_estdist_generic}; +#endif +} // namespace + +RABITQ_COLD_INIT const Backend& selected_backend() { +#if defined(__aarch64__) || defined(_M_ARM64) + static const Backend* const backend = rabitqlib::simd::resolve_optional_kernel( + &kNeonBackend, &kGenericBackend, cpu::has_neon() + ); +#else + static const Backend* const backend = + RABITQ_RESOLVE(&kAvx512Backend, &kAvx2Backend, &kGenericBackend); +#endif + return *backend; +} +} // namespace rabitqlib::fastscan::detail + +namespace rabitqlib::simd { + +void qg_batch_estdist( + const char* batch_data, + const BatchQuery& q_obj, + size_t padded_dim, + float* est_distance +) { + static const auto kQgBatchEstdistFn = fastscan::detail::selected_backend().qg; + + kQgBatchEstdistFn(batch_data, q_obj, padded_dim, est_distance); +} + +uint32_t qg_batch_estdist_mask( + const char* batch_data, + const BatchQuery& q_obj, + size_t padded_dim, + float* est_distance, + float threshold +) { + static const auto kQgBatchEstdistMaskFn = fastscan::detail::selected_backend().qg_mask; + + return kQgBatchEstdistMaskFn(batch_data, q_obj, padded_dim, est_distance, threshold); +} + +void split_batch_estdist( + const char* batch_data, + const SplitBatchQuery& q_obj, + size_t padded_dim, + float* est_distance, + float* low_distance, + float* ip_x0_qr, + bool use_hacc +) { + static const auto kSplitBatchEstdistFn = fastscan::detail::selected_backend().split; + + kSplitBatchEstdistFn( + batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + ); +} + +} // namespace rabitqlib::simd + +namespace rabitqlib::fastscan { + +template <> +void pack_lut(size_t dim, const float* __restrict__ query, float* __restrict__ lut) { + static const auto kPackLutFn = detail::selected_backend().pack_lut; + + kPackLutFn(dim, query, lut); +} + +void accumulate( + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lp_table, + int32_t* __restrict__ result, + size_t dim +) { + static const auto kAccumulateFn = detail::selected_backend().accumulate; + + if (dim == 0 || dim % 16 != 0) { + throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); + } + kAccumulateFn(codes, lp_table, result, dim); +} + +void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { + static const auto kTransferLutHaccFn = detail::selected_backend().transfer_hacc; + + if (dim == 0 || dim % 16 != 0) { + throw std::invalid_argument( + "high-accuracy FastScan dimension must be a positive multiple of 16" + ); + } + kTransferLutHaccFn(lut, dim, hc_lut); +} + +void accumulate_hacc( + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ hc_lut, + int32_t* accu_res, + size_t dim +) { + static const auto kAccumulateHaccFn = detail::selected_backend().accumulate_hacc; + + if (dim == 0 || dim % 16 != 0) { + throw std::invalid_argument( + "high-accuracy FastScan dimension must be a positive multiple of 16" + ); + } + kAccumulateHaccFn(codes, hc_lut, accu_res, dim); +} + +} // namespace rabitqlib::fastscan diff --git a/src/simd/fastscan_generic.cpp b/src/fastscan/fastscan_generic.cpp similarity index 93% rename from src/simd/fastscan_generic.cpp rename to src/fastscan/fastscan_generic.cpp index bd4f44ac..a093f5ea 100644 --- a/src/simd/fastscan_generic.cpp +++ b/src/fastscan/fastscan_generic.cpp @@ -1,10 +1,11 @@ #include +#include #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" #include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" namespace rabitqlib::fastscan::simd { void pack_lut_generic(size_t dim, const float* query, float* lut) { @@ -13,7 +14,8 @@ void pack_lut_generic(size_t dim, const float* query, float* lut) { for (size_t group = 0; group < dim / 4; ++group) { lut[0] = zero; for (size_t j = 1; j < 16; ++j) { - lut[j] = lut[j - LOWBIT(j)] + query[kPos[j]]; + const size_t lowbit = j & (size_t{0} - j); + lut[j] = lut[j - lowbit] + query[kPos[j]]; } query += 4; lut += 16; diff --git a/src/simd/fastscan_neon.cpp b/src/fastscan/fastscan_neon.cpp similarity index 71% rename from src/simd/fastscan_neon.cpp rename to src/fastscan/fastscan_neon.cpp index 3a590f9d..0cc93cbd 100644 --- a/src/simd/fastscan_neon.cpp +++ b/src/fastscan/fastscan_neon.cpp @@ -7,8 +7,8 @@ #include #include -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" namespace rabitqlib::fastscan::simd { namespace { @@ -83,3 +83,28 @@ void accumulate_hacc_neon( accumulate_impl(codes, lut, result, dim); } } // namespace rabitqlib::fastscan::simd + +namespace rabitqlib::fastscan::simd { +void pack_lut_neon(size_t dim, const float* query, float* lut) { + // Each table enumerates q0..q3 using bits 3..0. Add in that order, + // including the initial +0, to preserve cancellation and signed zeros. + volatile float zero = 0; + const uint32x4_t select2 = {0, 0, UINT32_MAX, UINT32_MAX}; + const uint32x4_t select3 = {0, UINT32_MAX, 0, UINT32_MAX}; + for (size_t group = 0; group < dim / 4; ++group) { + const float32x4_t z = vdupq_n_f32(zero); + const float32x4_t q0 = vdupq_n_f32(query[0]), q1 = vdupq_n_f32(query[1]), + q2 = vdupq_n_f32(query[2]), q3 = vdupq_n_f32(query[3]); + const float32x4_t high = vaddq_f32(z, q0); + float32x4_t tables[4] = {z, vaddq_f32(z, q1), high, vaddq_f32(high, q1)}; + for (size_t i = 0; i < 4; ++i) { + auto table = tables[i]; + table = vbslq_f32(select2, vaddq_f32(table, q2), table); + table = vbslq_f32(select3, vaddq_f32(table, q3), table); + vst1q_f32(lut + i * 4, table); + } + query += 4; + lut += 16; + } +} +} // namespace rabitqlib::fastscan::simd diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp new file mode 100644 index 00000000..fb0ec28b --- /dev/null +++ b/src/index/hnsw.cpp @@ -0,0 +1,1447 @@ +#include "rabitqlib/index/hnsw/hnsw.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/hnsw_dispatch.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/detail/visited_pool.hpp" +#include "rabitqlib/index/hnsw/hnsw_quant.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/memory.hpp" +#include "rabitqlib/utils/path.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::hnsw { + +HierarchicalNSW::HierarchicalNSW() = default; +HierarchicalNSW::~HierarchicalNSW() = default; +void HierarchicalNSW::AlignedDeleter::operator()(char* pointer) const noexcept { + memory::aligned_deallocate(pointer); +} +void HierarchicalNSW::LinkDeleter::operator()(char* pointer) const noexcept { + std::free(pointer); +} +int HierarchicalNSW::get_random_level(double reverse_size) { + std::uniform_real_distribution distribution(0.0, 1.0); + double r = -log(distribution(level_generator_)) * reverse_size; + return static_cast(r); +} + +HierarchicalNSW::HierarchicalNSW( + size_t max_elements, + size_t dim, + size_t total_bits, + size_t M, + size_t ef_construction, + size_t random_seed, + MetricType metric_type +) + : metric_type_(metric_type) + , label_op_locks_(kMaxLabelOperationLock) + , link_list_locks_(max_elements) + , element_levels_(max_elements) + , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) { + static_assert( + buffer::kSearchBufferMaxPointCount <= kRemovedMask, + "cluster IDs must leave the removal bit free" + ); + validate_metric_type(metric_type); + if (M < 2) { + throw std::invalid_argument("HNSW M must be at least 2"); + } + max_elements_ = max_elements; + dim_ = dim; + rotator_.reset(choose_rotator( + dim, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 32) + )); + padded_dim_ = rotator_->size(); + /* check size */ + assert(padded_dim_ % 32 == 0); + assert(padded_dim_ >= dim_); + ex_bits_ = total_bits - 1; + + if (total_bits < 1 || total_bits > 9) { + throw std::invalid_argument("HNSW quantization bits must be in [1, 9]"); + }; + + assert(padded_dim_ % 32 == 0); + + ip_func_ = select_excode_ipfunc(ex_bits_); + + if (M <= 10000) { + M_ = M; + } else { + M_ = 10000; + } + + maxM_ = M_; + maxM0_ = M_ * 2; + ef_construction_ = std::max(ef_construction, M_); + ef_ = 10; + + size_bin_data_ = BinDataMap::data_bytes(padded_dim_); + size_ex_data_ = ExDataMap::data_bytes(padded_dim_, ex_bits_); + size_links_level0_ = maxM0_ * sizeof(PID) + sizeof(PID); + label_offset_ = + size_links_level0_ + sizeof(PID); // (# of edges + edges) + (cluster_id) + bin_data_offset_ = + label_offset_ + + sizeof(PID); // (# of edges + edges) + (cluster_id) + (external label) + ex_data_offset_ = + bin_data_offset_ + size_bin_data_; // (# of edges + edges) + (cluster_id) + // + (external label) + (BinData) + size_data_per_element_ = ex_data_offset_ + size_ex_data_; + data_level0_memory_.reset( + memory::huge_page_allocate(max_elements_ * size_data_per_element_) + ); + if (data_level0_memory_ == nullptr) { + throw std::runtime_error("Not enough memory"); + } + + level_generator_.seed(random_seed); + + cur_element_count_ = 0; + + visited_list_pool_ = std::make_unique(1, max_elements_); + + // initializations for special treatment of the first node + entry_point_id_ = -1; + max_level_ = -1; + + link_lists_.reserve(max_elements_); + size_links_per_element_ = maxM_ * sizeof(PID) + sizeof(PID); + mult_ = 1 / log(1.0 * static_cast(M_)); + + this->query_config_ = + quant::faster_config(padded_dim_, SplitSingleQuery::kNumBits); +} + +void HierarchicalNSW::save(const char* filename) const { + std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + if (!output.is_open()) { + throw std::runtime_error("HNSW: cannot open index file for writing"); + } + + output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); + output.write(reinterpret_cast(&cur_element_count_), sizeof(size_t)); + + output.write(reinterpret_cast(&dim_), sizeof(size_t)); + output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); + output.write(reinterpret_cast(&num_clusters_), sizeof(size_t)); + output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); + output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); + + output.write(reinterpret_cast(&size_bin_data_), sizeof(size_t)); + output.write(reinterpret_cast(&size_ex_data_), sizeof(size_t)); + output.write(reinterpret_cast(&size_links_level0_), sizeof(size_t)); + output.write(reinterpret_cast(&bin_data_offset_), sizeof(size_t)); + output.write(reinterpret_cast(&ex_data_offset_), sizeof(size_t)); + output.write(reinterpret_cast(&label_offset_), sizeof(PID)); + output.write(reinterpret_cast(&size_data_per_element_), sizeof(size_t)); + output.write(reinterpret_cast(&size_links_per_element_), sizeof(size_t)); + + output.write(reinterpret_cast(&max_level_), sizeof(int)); + output.write(reinterpret_cast(&entry_point_id_), sizeof(PID)); + + output.write(reinterpret_cast(&M_), sizeof(size_t)); + output.write(reinterpret_cast(&maxM_), sizeof(size_t)); + output.write(reinterpret_cast(&maxM0_), sizeof(size_t)); + output.write(reinterpret_cast(&mult_), sizeof(double)); + output.write(reinterpret_cast(&ef_construction_), sizeof(size_t)); + + output.write( + reinterpret_cast(centroids_memory_.get()), + static_cast(num_clusters_ * padded_dim_ * sizeof(float)) + ); + + output.write( + reinterpret_cast(data_level0_memory_.get()), + static_cast(cur_element_count_ * size_data_per_element_) + ); + + for (size_t i = 0; i < cur_element_count_; i++) { + unsigned int link_list_size = + element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; + output.write(reinterpret_cast(&link_list_size), sizeof(unsigned int)); + if (link_list_size != 0) { + output.write(link_lists_[i].get(), link_list_size); + } + } + + rotator_->save(output); + output.close(); + if (!output) { + throw std::runtime_error("HNSW: failed to write index file"); + } +} + +void HierarchicalNSW::load(const char* filename) { + if (filename == nullptr || filename[0] == '\0') { + throw std::invalid_argument("HNSW: index filename must not be empty"); + } + std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + if (!input.is_open()) { + throw std::runtime_error("HNSW: cannot open index file"); + } + + auto invalid_file = []() -> void { + throw std::runtime_error("HNSW: invalid or truncated index file"); + }; + auto read_header = [&](auto& value) { + input.read(reinterpret_cast(&value), sizeof(value)); + if (!input) { + invalid_file(); + } + }; + auto checked_product = [&](size_t left, size_t right) { + if (right != 0 && left > std::numeric_limits::max() / right) { + invalid_file(); + } + return left * right; + }; + + // Keep the current index intact until the complete file has been validated. + HierarchicalNSW loaded; + size_t element_count = 0; + PID label_offset = 0; + read_header(loaded.max_elements_); + read_header(element_count); + read_header(loaded.dim_); + read_header(loaded.padded_dim_); + read_header(loaded.num_clusters_); + read_header(loaded.ex_bits_); + read_header(loaded.metric_type_); + read_header(loaded.size_bin_data_); + read_header(loaded.size_ex_data_); + read_header(loaded.size_links_level0_); + read_header(loaded.bin_data_offset_); + read_header(loaded.ex_data_offset_); + read_header(label_offset); + loaded.label_offset_ = label_offset; + read_header(loaded.size_data_per_element_); + read_header(loaded.size_links_per_element_); + read_header(loaded.max_level_); + read_header(loaded.entry_point_id_); + read_header(loaded.M_); + read_header(loaded.maxM_); + read_header(loaded.maxM0_); + read_header(loaded.mult_); + read_header(loaded.ef_construction_); + + if (loaded.max_elements_ == 0 || + loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || + element_count > loaded.max_elements_ || + loaded.num_clusters_ > buffer::kSearchBufferMaxPointCount || + (element_count != 0 && loaded.num_clusters_ == 0) || loaded.dim_ < 64 || + loaded.dim_ > rotator_impl::FhtKacRotator::kMaxDim || + loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || + loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || + (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || + loaded.M_ < 2 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || + loaded.maxM0_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || + (element_count == 0 && + (loaded.max_level_ != -1 || loaded.entry_point_id_ != kPidMax)) || + (element_count != 0 && + (loaded.max_level_ < 0 || loaded.entry_point_id_ >= element_count))) { + invalid_file(); + } + + const size_t expected_level0 = (loaded.maxM0_ + 1) * sizeof(PID); + const size_t expected_bin = BinDataMap::data_bytes(loaded.padded_dim_); + const size_t expected_ex = + ExDataMap::data_bytes(loaded.padded_dim_, loaded.ex_bits_); + const size_t expected_label = expected_level0 + sizeof(PID); + const size_t expected_bin_offset = expected_label + sizeof(PID); + const size_t expected_ex_offset = expected_bin_offset + expected_bin; + const size_t expected_stride = expected_ex_offset + expected_ex; + if (loaded.size_bin_data_ != expected_bin || loaded.size_ex_data_ != expected_ex || + loaded.size_links_level0_ != expected_level0 || + loaded.label_offset_ != expected_label || + loaded.bin_data_offset_ != expected_bin_offset || + loaded.ex_data_offset_ != expected_ex_offset || + loaded.size_data_per_element_ != expected_stride || + loaded.size_links_per_element_ != (loaded.maxM_ + 1) * sizeof(PID) || + !std::isfinite(loaded.mult_) || loaded.mult_ <= 0.0) { + invalid_file(); + } + + const size_t centroids_bytes = checked_product( + checked_product(loaded.num_clusters_, loaded.padded_dim_), sizeof(float) + ); + const size_t data_capacity_bytes = + checked_product(loaded.max_elements_, loaded.size_data_per_element_); + const size_t data_bytes = checked_product(element_count, loaded.size_data_per_element_); + // Validate pointer-table capacity before allocating it. + (void)checked_product(loaded.max_elements_, sizeof(LinkStorage)); + const size_t link_headers_bytes = checked_product(element_count, sizeof(unsigned int)); + loaded.rotator_ = rotator_impl::make_rotator_for_loading( + loaded.dim_, RotatorType::FhtKacRotator, loaded.padded_dim_ + ); + const size_t rotator_bytes = loaded.rotator_->dump_bytes(); + + const auto payload_start = input.tellg(); + input.seekg(0, std::ios::end); + const auto payload_end = input.tellg(); + if (!input || payload_start < 0 || payload_end < payload_start) { + invalid_file(); + } + size_t remaining = static_cast(payload_end - payload_start); + input.seekg(payload_start); + if (centroids_bytes > remaining || data_bytes > remaining - centroids_bytes || + link_headers_bytes > remaining - centroids_bytes - data_bytes || + rotator_bytes > remaining - centroids_bytes - data_bytes - link_headers_bytes) { + invalid_file(); + } + auto read_payload = [&](char* destination, size_t bytes) { + if (bytes > remaining || + bytes > static_cast(std::numeric_limits::max())) { + invalid_file(); + } + if (bytes != 0) { + input.read(destination, static_cast(bytes)); + if (!input) { + invalid_file(); + } + } + remaining -= bytes; + }; + + if (centroids_bytes != 0) { + loaded.centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); + if (loaded.centroids_memory_ == nullptr) { + throw std::bad_alloc(); + } + read_payload(loaded.centroids_memory_.get(), centroids_bytes); + } + loaded.data_level0_memory_.reset(memory::huge_page_allocate(data_capacity_bytes)); + if (loaded.data_level0_memory_ == nullptr) { + throw std::bad_alloc(); + } + read_payload(loaded.data_level0_memory_.get(), data_bytes); + + loaded.element_levels_ = std::vector(loaded.max_elements_); + loaded.link_list_locks_ = std::vector(loaded.max_elements_); + loaded.label_op_locks_ = std::vector(kMaxLabelOperationLock); + loaded.link_lists_.reserve(loaded.max_elements_); + loaded.link_lists_.resize(element_count); + for (size_t i = 0; i < element_count; ++i) { + const char* row = + loaded.data_level0_memory_.get() + i * loaded.size_data_per_element_; + PID cluster_id = 0; + PID label = 0; + unsigned short level0_degree = 0; + std::memcpy(&level0_degree, row, sizeof(level0_degree)); + std::memcpy(&cluster_id, row + loaded.size_links_level0_, sizeof(cluster_id)); + std::memcpy(&label, row + loaded.label_offset_, sizeof(label)); + if ((cluster_id & ~kRemovedMask) >= loaded.num_clusters_ || + level0_degree > loaded.maxM0_) { + invalid_file(); + } + loaded.num_removed_ += (cluster_id & kRemovedMask) != 0; + for (size_t edge = 0; edge < level0_degree; ++edge) { + PID neighbor = 0; + std::memcpy(&neighbor, row + sizeof(PID) + edge * sizeof(PID), sizeof(PID)); + if (neighbor >= element_count) { + invalid_file(); + } + } + if (!loaded.label_lookup_.emplace(label, static_cast(i)).second) { + invalid_file(); + } + + unsigned int link_list_size = 0; + read_payload(reinterpret_cast(&link_list_size), sizeof(link_list_size)); + if (link_list_size != 0) { + if (link_list_size % loaded.size_links_per_element_ != 0 || + link_list_size / loaded.size_links_per_element_ > + static_cast(loaded.max_level_) || + link_list_size > remaining) { + invalid_file(); + } + loaded.link_lists_[i].reset(static_cast(std::malloc(link_list_size))); + if (loaded.link_lists_[i] == nullptr) { + throw std::bad_alloc(); + } + loaded.element_levels_[i] = + static_cast(link_list_size / loaded.size_links_per_element_); + } + loaded.cur_element_count_ = i + 1; + if (link_list_size != 0) { + read_payload(loaded.link_lists_[i].get(), link_list_size); + for (int level = 0; level < loaded.element_levels_[i]; ++level) { + const char* link = + loaded.link_lists_[i].get() + level * loaded.size_links_per_element_; + unsigned short degree = 0; + std::memcpy(°ree, link, sizeof(degree)); + if (degree > loaded.maxM_) { + invalid_file(); + } + for (size_t edge = 0; edge < degree; ++edge) { + PID neighbor = 0; + std::memcpy( + &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) + ); + if (neighbor >= element_count) { + invalid_file(); + } + } + } + } + } + for (size_t i = 0; i < element_count; ++i) { + for (int level = 1; level <= loaded.element_levels_[i]; ++level) { + const char* link = + loaded.link_lists_[i].get() + + (static_cast(level) - 1) * loaded.size_links_per_element_; + unsigned short degree = 0; + std::memcpy(°ree, link, sizeof(degree)); + for (size_t edge = 0; edge < degree; ++edge) { + PID neighbor = 0; + std::memcpy( + &neighbor, link + sizeof(PID) + edge * sizeof(PID), sizeof(PID) + ); + if (loaded.element_levels_[neighbor] < level) { + invalid_file(); + } + } + } + } + if (element_count != 0 && + loaded.element_levels_[loaded.entry_point_id_] != loaded.max_level_) { + invalid_file(); + } + + if (rotator_bytes > remaining) { + invalid_file(); + } + loaded.rotator_->load(input); + if (!input) { + invalid_file(); + } + remaining -= rotator_bytes; + if (remaining != 0) { + invalid_file(); + } + + loaded.visited_list_pool_ = std::make_unique(1, loaded.max_elements_); + loaded.raw_dist_func_ = + (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; + loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); + loaded.ef_ = 10; + loaded.query_config_ = + quant::faster_config(loaded.padded_dim_, SplitSingleQuery::kNumBits); + + // All swaps below are non-allocating. The temporary releases the old state. + using std::swap; + swap(construction_data_, loaded.construction_data_); + swap(max_elements_, loaded.max_elements_); + const size_t old_count = cur_element_count_.load(); + cur_element_count_ = loaded.cur_element_count_.load(); + loaded.cur_element_count_ = old_count; + swap(num_removed_, loaded.num_removed_); + swap(size_data_per_element_, loaded.size_data_per_element_); + swap(size_links_per_element_, loaded.size_links_per_element_); + swap(M_, loaded.M_); + swap(maxM_, loaded.maxM_); + swap(maxM0_, loaded.maxM0_); + swap(ef_construction_, loaded.ef_construction_); + swap(ef_, loaded.ef_); + swap(metric_type_, loaded.metric_type_); + swap(mult_, loaded.mult_); + swap(max_level_, loaded.max_level_); + label_op_locks_.swap(loaded.label_op_locks_); + link_list_locks_.swap(loaded.link_list_locks_); + swap(entry_point_id_, loaded.entry_point_id_); + swap(size_links_level0_, loaded.size_links_level0_); + swap(bin_data_offset_, loaded.bin_data_offset_); + swap(ex_data_offset_, loaded.ex_data_offset_); + swap(label_offset_, loaded.label_offset_); + swap(size_bin_data_, loaded.size_bin_data_); + swap(size_ex_data_, loaded.size_ex_data_); + swap(ex_bits_, loaded.ex_bits_); + swap(data_level0_memory_, loaded.data_level0_memory_); + swap(link_lists_, loaded.link_lists_); + element_levels_.swap(loaded.element_levels_); + swap(num_clusters_, loaded.num_clusters_); + swap(dim_, loaded.dim_); + swap(padded_dim_, loaded.padded_dim_); + swap(centroids_memory_, loaded.centroids_memory_); + label_lookup_.swap(loaded.label_lookup_); + swap(visited_list_pool_, loaded.visited_list_pool_); + swap(ip_func_, loaded.ip_func_); + swap(rotator_, loaded.rotator_); + swap(query_config_, loaded.query_config_); + swap(raw_dist_func_, loaded.raw_dist_func_); +} + +void HierarchicalNSW::construct( + size_t num_clusters, + const float* centroids, + size_t num_points, + const float* data, + PID* cluster_ids, + size_t num_threads, + bool faster +) { + if (num_clusters_ != 0 || cur_element_count_ != 0) { + throw std::logic_error("HNSW index is already constructed or loaded"); + } + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW cluster count is out of range"); + } + if (num_points == 0 || num_points > max_elements_ || + num_points > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW point count is out of range"); + } + if (centroids == nullptr || data == nullptr || cluster_ids == nullptr) { + throw std::invalid_argument("HNSW construction inputs must not be null"); + } + if (num_clusters > std::numeric_limits::max() / sizeof(float) / padded_dim_) { + throw std::invalid_argument("HNSW centroids exceed addressable storage"); + } + for (size_t i = 0; i < num_points; ++i) { + if (cluster_ids[i] >= num_clusters) { + throw std::invalid_argument("HNSW cluster ID is out of range"); + } + } + + try { + num_clusters_ = num_clusters; + const size_t centroids_bytes = num_clusters_ * padded_dim_ * sizeof(float); + centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); + if (centroids_memory_ == nullptr) { + throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids" + ); + } + + for (size_t i = 0; i < num_clusters; ++i) { + this->rotator_->rotate( + centroids + (i * dim_), + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_) + ); + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + construction_data_ = data; + rabitqlib::index_detail::parallel_for( + 0, + num_points, + num_threads, + [&](size_t idx, size_t /*threadId*/) { + add_point(idx, cluster_ids[idx], config); + } + ); + } catch (...) { + // parallel_for joins every worker before propagating an exception. + // Discard the partial graph, retaining capacity and rotation for a retry. + std::fill_n(element_levels_.begin(), cur_element_count_.load(), 0); + link_lists_.clear(); + cur_element_count_ = 0; + num_removed_ = 0; + label_lookup_.clear(); + entry_point_id_ = kPidMax; + max_level_ = -1; + centroids_memory_.reset(); + num_clusters_ = 0; + construction_data_ = nullptr; + throw; + } + construction_data_ = nullptr; +} + +void HierarchicalNSW::add_point( + PID label, PID cluster_id, const quant::RabitqConfig& config +) { + std::unique_lock lock_label(get_label_op_mutex(label)); + + int node_level = 0; + PID cur_c = 0; + { + std::unique_lock lock_table(label_lookup_lock_); + if (label_lookup_.find(label) != label_lookup_.end()) { + throw std::runtime_error( + "Currently not support replacement of existing elements, only support " + "inserting elements with distinct labels" + ); + } + + if (cur_element_count_ >= max_elements_) { + throw std::runtime_error("The number of elements exceeds the specified limit"); + } + + cur_c = cur_element_count_; + link_lists_.emplace_back(); + cur_element_count_++; + label_lookup_[label] = cur_c; + node_level = get_random_level(mult_); + } + + std::unique_lock lock_el(link_list_locks_[cur_c]); + + element_levels_[cur_c] = node_level; + std::unique_lock templock(global_); + int previous_max_level = max_level_; + PID current_id = entry_point_id_; + if (node_level <= previous_max_level) { + templock.unlock(); + } + + // initialize the current memory. + memset( + data_level0_memory_.get() + (cur_c * size_data_per_element_), + 0, + size_data_per_element_ + ); + + // Initialisation of label and cluster id + memcpy(get_external_label_pt(cur_c), &label, sizeof(PID)); + memcpy(get_clusterid_pt(cur_c), &cluster_id, sizeof(PID)); + + // Quantize raw data and initialize quantized data + std::vector rotated_data(padded_dim_); + rotator_->rotate(construction_data_ + (label * dim_), rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_.get()) + (cluster_id * padded_dim_), + padded_dim_, + ex_bits_, + get_bindata_by_internalid(cur_c), + get_exdata_by_internalid(cur_c), + metric_type_, + config + ); + + // If the current vertex is at level >0, it needs some space to store the extra edges. + if (node_level > 0) { + link_lists_[cur_c].reset( + static_cast(std::malloc((size_links_per_element_ * node_level) + 1)) + ); + if (link_lists_[cur_c] == nullptr) { + throw std::runtime_error( + "Not enough memory: add_point failed to allocate linklist" + ); + } + memset(link_lists_[cur_c].get(), 0, (size_links_per_element_ * node_level) + 1); + } + + if (static_cast(current_id) != -1) { + if (node_level < previous_max_level) { + float curdist = get_data_dist(current_id, cur_c); + for (int level = previous_max_level; level > node_level; level--) { + bool changed = true; + while (changed) { + changed = false; + unsigned int* data; + std::unique_lock lock(link_list_locks_[current_id]); + data = get_linklist(current_id, level); + int size = get_list_count(data); + + auto* datal = static_cast(data + 1); + for (int i = 0; i < size; i++) { + PID cand = datal[i]; + if (cand > max_elements_) { + throw std::runtime_error("cand error"); + } + float d = get_data_dist(cand, cur_c); + if (d < curdist) { + curdist = d; + current_id = cand; + changed = true; + } + } + } + } + } + + for (int level = std::min(node_level, previous_max_level); level >= 0; level--) { + maxheap> top_candidates = + search_base_layer(current_id, cur_c, level); + current_id = mutually_connect_new_element(cur_c, top_candidates, level); + } + } else { + // Do nothing for the first element + entry_point_id_ = 0; + max_level_ = node_level; + } + + // Releasing lock for the maximum level + if (node_level > previous_max_level) { + entry_point_id_ = cur_c; + max_level_ = node_level; + } +} + +maxheap> HierarchicalNSW::search_base_layer( + PID ep_id, PID cur_c, int layer +) { + auto vl = visited_list_pool_->get_free_vislist(); + + maxheap> top_candidates; + minheap> candidate_set; + + float lower_bound = get_data_dist(ep_id, cur_c); + top_candidates.emplace(lower_bound, ep_id); + candidate_set.emplace(lower_bound, ep_id); + vl->set(ep_id); + + while (!candidate_set.empty()) { + std::pair curr_el_pair = candidate_set.top(); + if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { + break; + } + candidate_set.pop(); + + PID cur_node_num = curr_el_pair.second; + + std::unique_lock lock(link_list_locks_[cur_node_num]); + + int* data; + if (layer == 0) { + data = reinterpret_cast(get_linklist0(cur_node_num)); + } else { + data = reinterpret_cast(get_linklist(cur_node_num, layer)); + } + size_t size = get_list_count(reinterpret_cast(data)); + auto* datal = reinterpret_cast(data + 1); + + rabitqlib::memory::mem_prefetch_l1( + reinterpret_cast( + construction_data_ + (get_external_label(*datal) * dim_) + ), + padded_dim_ / 16 + ); + + rabitqlib::memory::mem_prefetch_l1( + reinterpret_cast( + construction_data_ + (get_external_label(*(datal + 1)) * dim_) + ), + padded_dim_ / 16 + ); + + for (size_t j = 0; j < size; j++) { + PID candidate_id = *(datal + j); + if (vl->get(candidate_id)) { + continue; + } + vl->set(candidate_id); + + if (j < size - 1) { + rabitqlib::memory::mem_prefetch_l1( + reinterpret_cast( + construction_data_ + (get_external_label(*(datal + j + 1)) * dim_) + ), + padded_dim_ / 16 + ); + } + + float dist1 = get_data_dist(candidate_id, cur_c); + if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { + candidate_set.emplace(dist1, candidate_id); + top_candidates.emplace(dist1, candidate_id); + if (top_candidates.size() > ef_construction_) { + top_candidates.pop(); + } + if (!top_candidates.empty()) { + lower_bound = top_candidates.top().first; + } + } + } + } + visited_list_pool_->release_vis_list(std::move(vl)); + return top_candidates; +} + +namespace detail { +template +void prune_neighbors( + maxheap>& top_candidates, size_t M, Distance distance +) { + if (top_candidates.size() < M) { + return; + } + + minheap> queue_closest; + std::vector> return_list; + while (top_candidates.size() > 0) { + queue_closest.emplace(top_candidates.top()); + top_candidates.pop(); + } + + while (queue_closest.size() > 0) { + if (return_list.size() >= M) { + break; + } + std::pair current_pair = queue_closest.top(); + float dist_to_query = current_pair.first; + queue_closest.pop(); + bool good = true; + + for (std::pair second_pair : return_list) { + float curdist = distance(second_pair.second, current_pair.second); + if (curdist < dist_to_query) { + good = false; + break; + } + } + if (good) { + return_list.push_back(current_pair); + } + } + + for (std::pair current_pair : return_list) { + top_candidates.emplace(current_pair); + } +} +} // namespace detail + +template +PID HierarchicalNSW::mutually_connect( + PID cur_c, + maxheap>& top_candidates, + int level, + Distance distance, + ReverseEdgeDistance reverse_edge_distance +) { + size_t max_m = level > 0 ? maxM_ : maxM0_; + detail::prune_neighbors(top_candidates, M_, distance); + if (top_candidates.size() > M_) { + throw std::runtime_error( + "Should be not be more than M_ candidates returned by the heuristic" + ); + } + + std::vector selected_neighbors; + selected_neighbors.reserve(M_); + while (top_candidates.size() > 0) { + selected_neighbors.push_back(top_candidates.top().second); + top_candidates.pop(); + } + if (selected_neighbors.empty()) { + throw std::runtime_error("No neighbors selected for the new element"); + } + PID next_closest_entry_point = selected_neighbors.back(); + + { + PID* ll_cur = (level == 0) ? get_linklist0(cur_c) : get_linklist(cur_c, level); + if (*ll_cur > 0) { + throw std::runtime_error( + "The newly inserted element should have blank link list" + ); + } + + set_list_count(ll_cur, selected_neighbors.size()); + auto* data = ll_cur + 1; + for (size_t idx = 0; idx < selected_neighbors.size(); idx++) { + if (data[idx] != 0) { + throw std::runtime_error("Possible memory corruption"); + } + if (level > element_levels_[selected_neighbors[idx]]) { + throw std::runtime_error("Trying to make a link on a non-existent level"); + } + data[idx] = selected_neighbors[idx]; + } + } + + for (auto selected_neighbor : selected_neighbors) { + std::unique_lock lock(link_list_locks_[selected_neighbor]); + + PID* ll_other = (level == 0) ? get_linklist0(selected_neighbor) + : get_linklist(selected_neighbor, level); + size_t sz_link_list_other = get_list_count(ll_other); + + if (sz_link_list_other > max_m) { + throw std::runtime_error("Bad value of sz_link_list_other"); + } + if (selected_neighbor == cur_c) { + throw std::runtime_error("Trying to connect an element to itself"); + } + if (level > element_levels_[selected_neighbor]) { + throw std::runtime_error("Trying to make a link on a non-existent level"); + } + + auto* data = ll_other + 1; + + bool is_cur_c_present = false; + for (size_t j = 0; j < sz_link_list_other; j++) { + if (data[j] == cur_c) { + is_cur_c_present = true; + break; + } + } + + if (!is_cur_c_present) { + if (sz_link_list_other < max_m) { + data[sz_link_list_other] = cur_c; + set_list_count(ll_other, sz_link_list_other + 1); + } else { + float d_max = reverse_edge_distance(selected_neighbor); + maxheap> candidates; + candidates.emplace(d_max, cur_c); + for (size_t j = 0; j < sz_link_list_other; j++) { + candidates.emplace(distance(data[j], selected_neighbor), data[j]); + } + + detail::prune_neighbors(candidates, max_m, distance); + + int indx = 0; + while (candidates.size() > 0) { + data[indx] = candidates.top().second; + candidates.pop(); + indx++; + } + set_list_count(ll_other, indx); + } + } + } + + return next_closest_entry_point; +} + +PID HierarchicalNSW::mutually_connect_new_element( + PID cur_c, maxheap>& top_candidates, int level +) { + return mutually_connect( + cur_c, + top_candidates, + level, + [this](PID target, PID query) { return get_data_dist(target, query); }, + // Preserve the original operand order for a newly proposed reverse edge. + [this, cur_c](PID neighbor) { return get_data_dist(neighbor, cur_c); } + ); +} + +std::vector>> HierarchicalNSW::search( + const float* queries, size_t num_queries, size_t k, size_t ef_search, size_t num_threads +) { + if (k == 0) { + throw std::invalid_argument("HNSW k must be greater than zero"); + } + std::vector>> results(num_queries); + rabitqlib::index_detail::parallel_for( + 0, + num_queries, + num_threads, + [&](size_t idx, size_t /*threadId*/) { + std::vector rotated_query(padded_dim_); + this->rotator_->rotate(queries + (idx * dim_), rotated_query.data()); + results[idx] = detail::search_knn_sorted( + *this, rotated_query.data(), k, std::max(ef_search, k) + ); + } + ); + return results; +} + +size_t HierarchicalNSW::remove(const PID* labels, size_t n) { + if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || + rotator_ == nullptr || num_clusters_ == 0) { + throw std::logic_error("HNSW index must be constructed or loaded before remove"); + } + if (n == 0) { + return 0; + } + if (labels == nullptr) { + throw std::invalid_argument("HNSW remove labels must not be null"); + } + + // Resolve every label before marking any point, so an unknown label removes + // nothing. + std::vector internal_ids(n); + { + std::unique_lock lock_table(label_lookup_lock_); + for (size_t i = 0; i < n; ++i) { + const auto found = label_lookup_.find(labels[i]); + if (found == label_lookup_.end()) { + throw std::invalid_argument("HNSW remove label is not in the index"); + } + internal_ids[i] = found->second; + } + } + + size_t removed = 0; + for (const PID id : internal_ids) { + if (!is_removed(id)) { + const PID marked = get_raw_clusterid(id) | kRemovedMask; + std::memcpy(get_clusterid_pt(id), &marked, sizeof(PID)); + ++removed; + } + } + num_removed_ += removed; + return removed; +} + +// Code per dimension is (sign << ex_bits) + ex, up to nine bits, scaled by +// f_rescale_ex. +void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { + thread_local std::vector combined; + thread_local std::vector ex_codes; + combined.resize(padded_dim_); + + const ConstBinDataMap bin(get_bindata_by_internalid(id), padded_dim_); + const uint8_t* sign_code = bin.bin_code(); + for (size_t word = 0; word * 64 < padded_dim_; ++word) { + uint64_t bits = 0; + const size_t base = word * 64; + const size_t span = std::min(64, padded_dim_ - base); + std::memcpy(&bits, sign_code + (word * sizeof(uint64_t)), span / 8); + for (size_t j = 0; j < span; ++j) { + combined[base + j] = static_cast((bits >> (span - 1 - j)) & 1ULL); + } + } + + const auto* centroid = reinterpret_cast(centroids_memory_.get()) + + (get_clusterid_by_internalid(id) * padded_dim_); + + if (ex_bits_ == 0) { + quant::reconstruct_full_vec( + combined.data(), centroid, padded_dim_, 1, bin.f_rescale(), out, metric_type_ + ); + return; + } + + ex_codes.resize(padded_dim_); + const ConstExDataMap ex(get_exdata_by_internalid(id), padded_dim_, ex_bits_); + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + ex.ex_code(), ex_codes.data(), padded_dim_, ex_bits_ + ); + for (size_t d = 0; d < padded_dim_; ++d) { + combined[d] = static_cast((combined[d] << ex_bits_) + ex_codes[d]); + } + + quant::reconstruct_full_vec( + combined.data(), + centroid, + padded_dim_, + ex_bits_ + 1, + ex.f_rescale_ex(), + out, + metric_type_ + ); +} + +// Only `query` is reconstructed; `target` is scored from its codes, as QGBuilder +// does. Callers hold `query` fixed across a loop, so it is prepared once. +float HierarchicalNSW::get_quant_dist(PID target, PID query, detail::QuantQuery& cached) + const { + if (cached.id != query || !cached.prepared.has_value()) { + cached.rotated.resize(padded_dim_); + reconstruct_rotated(query, cached.rotated.data()); + cached.prepared.emplace( + cached.rotated.data(), padded_dim_, ex_bits_, query_config_, metric_type_ + ); + cached.id = query; + } + + const auto* centroid = reinterpret_cast(centroids_memory_.get()) + + (get_clusterid_by_internalid(target) * padded_dim_); + const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_); + const float g_error = std::sqrt(dist_sqr); + const float g_add = (metric_type_ == METRIC_IP) + ? -dot_product(cached.rotated.data(), centroid, padded_dim_) + : dist_sqr; + + if (ex_bits_ == 0) { + return estimate_split_single( + get_bindata_by_internalid(target), + *cached.prepared, + padded_dim_, + g_add, + g_error + ) + .distance; + } + return estimate_split_single_full( + get_bindata_by_internalid(target), + get_exdata_by_internalid(target), + ip_func_, + *cached.prepared, + padded_dim_, + ex_bits_, + g_add, + g_error + ) + .distance; +} + +maxheap> HierarchicalNSW::search_base_layer_quant( + PID ep_id, PID cur_c, int layer, detail::QuantQuery& query +) { + auto vl = visited_list_pool_->get_free_vislist(); + + maxheap> top_candidates; + minheap> candidate_set; + + float lower_bound = get_quant_dist(ep_id, cur_c, query); + top_candidates.emplace(lower_bound, ep_id); + candidate_set.emplace(lower_bound, ep_id); + vl->set(ep_id); + + while (!candidate_set.empty()) { + std::pair curr_el_pair = candidate_set.top(); + if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { + break; + } + candidate_set.pop(); + + PID cur_node_num = curr_el_pair.second; + std::unique_lock lock(link_list_locks_[cur_node_num]); + + PID* data = + (layer == 0) ? get_linklist0(cur_node_num) : get_linklist(cur_node_num, layer); + size_t size = get_list_count(data); + auto* datal = data + 1; + + for (size_t j = 0; j < size; j++) { + PID candidate_id = datal[j]; + if (candidate_id >= cur_element_count_) { + throw std::runtime_error("cand error"); + } + if (vl->get(candidate_id)) { + continue; + } + vl->set(candidate_id); + + float dist1 = get_quant_dist(candidate_id, cur_c, query); + if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { + candidate_set.emplace(dist1, candidate_id); + top_candidates.emplace(dist1, candidate_id); + if (top_candidates.size() > ef_construction_) { + top_candidates.pop(); + } + if (!top_candidates.empty()) { + lower_bound = top_candidates.top().first; + } + } + } + } + visited_list_pool_->release_vis_list(std::move(vl)); + return top_candidates; +} + +PID HierarchicalNSW::mutually_connect_quant( + PID cur_c, + maxheap>& top_candidates, + int level, + detail::QuantQuery& query +) { + return mutually_connect( + cur_c, + top_candidates, + level, + [this, &query](PID target, PID source) { + return get_quant_dist(target, source, query); + }, + // Quantized scoring is directional: reconstruct the existing neighbor. + [this, cur_c, &query](PID neighbor) { + return get_quant_dist(cur_c, neighbor, query); + } + ); +} + +// The slot, the level and the link list are all decided by the caller, so this +// can run on many threads at once and the labels do not depend on their order. +void HierarchicalNSW::add_point_quant( + PID cur_c, int node_level, LinkStorage& link_list, detail::QuantQuery& query +) { + std::unique_lock lock_el(link_list_locks_[cur_c]); + + std::unique_lock templock(global_); + int previous_max_level = max_level_; + PID current_id = entry_point_id_; + if (node_level <= previous_max_level) { + templock.unlock(); + } + const PID previous_entry_point = current_id; + + link_lists_[cur_c] = std::move(link_list); + element_levels_[cur_c] = node_level; + + try { + if (static_cast(current_id) != -1) { + if (node_level < previous_max_level) { + float curdist = get_quant_dist(current_id, cur_c, query); + for (int level = previous_max_level; level > node_level; level--) { + bool changed = true; + while (changed) { + changed = false; + std::unique_lock lock(link_list_locks_[current_id]); + PID* data = get_linklist(current_id, level); + int size = get_list_count(data); + auto* datal = data + 1; + for (int i = 0; i < size; i++) { + PID cand = datal[i]; + if (cand >= cur_element_count_) { + throw std::runtime_error("cand error"); + } + float d = get_quant_dist(cand, cur_c, query); + if (d < curdist) { + curdist = d; + current_id = cand; + changed = true; + } + } + } + } + } + + for (int level = std::min(node_level, previous_max_level); level >= 0; + level--) { + maxheap> top_candidates = + search_base_layer_quant(current_id, cur_c, level, query); + current_id = mutually_connect_quant(cur_c, top_candidates, level, query); + } + } else { + entry_point_id_ = cur_c; + max_level_ = node_level; + } + + if (node_level > previous_max_level) { + entry_point_id_ = cur_c; + max_level_ = node_level; + } + } catch (...) { + // Upper layers may already expose this node. Keep a route to the old + // graph if base-layer linking was interrupted before choosing neighbors. + PID* base_links = get_linklist0(cur_c); + if (previous_entry_point != kPidMax && get_list_count(base_links) == 0) { + base_links[1] = previous_entry_point; + set_list_count(base_links, 1); + } + // An interrupted promotion has no incoming links above the previous + // maximum level. Discard those levels so save/load sees a valid graph. + if (node_level > previous_max_level) { + element_levels_[cur_c] = std::max(0, previous_max_level); + if (element_levels_[cur_c] == 0) { + link_lists_[cur_c].reset(); + } + } + throw; + } +} + +// A capacity change reallocates the level-0 block, invalidating pointers into it. +// This must not run while another thread searches; equal capacity is a no-op. +void HierarchicalNSW::resize(size_t new_max_elements) { + if (size_data_per_element_ == 0) { + throw std::logic_error("HNSW index must be initialized before resize"); + } + if (new_max_elements < cur_element_count_) { + throw std::invalid_argument("HNSW resize: capacity below the element count"); + } + if (new_max_elements > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW resize: capacity exceeds the supported ID range"); + } + if (new_max_elements == max_elements_) { + return; + } + if (new_max_elements > std::numeric_limits::max() / size_data_per_element_) { + throw std::invalid_argument("HNSW resize: capacity exceeds addressable storage"); + } + + // Everything new is built before anything old is released, so a failure + // leaves the index as it was. realloc is not an option: the level-0 block + // comes from aligned_alloc, which realloc does not keep aligned. + const size_t level0_bytes = new_max_elements * size_data_per_element_; + AlignedStorage level0(memory::huge_page_allocate(level0_bytes)); + if (level0 == nullptr) { + throw std::bad_alloc(); + } + std::memset(level0.get(), 0, level0_bytes); + std::memcpy( + level0.get(), data_level0_memory_.get(), cur_element_count_ * size_data_per_element_ + ); + + std::vector links; + links.reserve(new_max_elements); + + std::vector levels(element_levels_); + levels.resize(new_max_elements); + std::vector locks(new_max_elements); + auto pool = std::make_unique(1, new_max_elements); + + // All potentially throwing work is complete. Moving the live owners into + // reserved slots cannot allocate or throw; unused capacity stays uninitialized. + for (auto& link : link_lists_) { + links.push_back(std::move(link)); + } + data_level0_memory_ = std::move(level0); + link_lists_ = std::move(links); + element_levels_.swap(levels); + link_list_locks_.swap(locks); + visited_list_pool_ = std::move(pool); + max_elements_ = new_max_elements; +} + +std::vector HierarchicalNSW::add( + const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads +) { + if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || + rotator_ == nullptr || num_clusters_ == 0) { + throw std::logic_error("HNSW index must be constructed or loaded before add"); + } + if (n == 0) { + return {}; + } + if (data == nullptr) { + throw std::invalid_argument("HNSW add inputs must not be null"); + } + if (n > max_elements_ - cur_element_count_) { + throw std::invalid_argument("HNSW add: not enough capacity"); + } + if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { + throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); + } + std::vector labels(n); + // Either the caller names the cluster of every point, or each one goes to its + // nearest centroid, ordered the way a query is routed. + std::vector assigned(n, 0); + if (cluster_ids != nullptr) { + for (size_t i = 0; i < n; ++i) { + if (cluster_ids[i] >= num_clusters_) { + throw std::invalid_argument("HNSW cluster ID is out of range"); + } + assigned[i] = cluster_ids[i]; + } + } else { + const auto* centroids = reinterpret_cast(centroids_memory_.get()); + rabitqlib::index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t) { + std::vector rotated(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated.data()); + float best = std::numeric_limits::max(); + for (size_t c = 0; c < num_clusters_; ++c) { + const float dist = raw_dist_func_( + rotated.data(), centroids + (c * padded_dim_), padded_dim_ + ); + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); + } + } + }); + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + // Drawn here because the generator is shared: the inserts below run in + // parallel and must not touch it. + std::vector levels(n); + for (size_t i = 0; i < n; ++i) { + levels[i] = get_random_level(mult_); + } + + // Every link list is allocated before any of them is installed, so running + // out of memory leaves the index untouched. + std::vector link_lists; + link_lists.reserve(n); + for (size_t i = 0; i < n; ++i) { + char* raw = nullptr; + if (levels[i] > 0) { + raw = + static_cast(std::calloc(1, (size_links_per_element_ * levels[i]) + 1) + ); + if (raw == nullptr) { + throw std::runtime_error( + "Not enough memory: HNSW add failed to allocate a link list" + ); + } + } + link_lists.emplace_back(raw); + } + + const size_t workers = std::min(resolve_num_threads(num_threads), n); + std::vector queries(workers); + + // Prepare every code before publishing the slots. If quantization or worker + // creation fails, the graph and point count are unchanged. + { + std::unique_lock lock_table(label_lookup_lock_); + if (cur_element_count_ + n > max_elements_) { + throw std::runtime_error("The number of elements exceeds the specified limit"); + } + const auto first = static_cast(cur_element_count_.load()); + rabitqlib::index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t) { + const auto id = static_cast(first + i); + labels[i] = id; + std::memset( + data_level0_memory_.get() + (id * size_data_per_element_), + 0, + size_data_per_element_ + ); + std::memcpy(get_external_label_pt(id), &id, sizeof(PID)); + std::memcpy(get_clusterid_pt(id), &assigned[i], sizeof(PID)); + std::vector rotated_data(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_.get()) + + (assigned[i] * padded_dim_), + padded_dim_, + ex_bits_, + get_bindata_by_internalid(id), + get_exdata_by_internalid(id), + metric_type_, + config + ); + }); + size_t registered = 0; + try { + for (; registered < n; ++registered) { + label_lookup_[labels[registered]] = labels[registered]; + } + } catch (...) { + for (size_t i = 0; i < registered; ++i) { + label_lookup_.erase(labels[i]); + } + throw; + } + // Unstarted inserts are valid level-0 nodes with owned codes. Their + // upper-layer lists remain under RAII until an insert takes ownership. + link_lists_.resize(cur_element_count_ + n); + for (size_t i = 0; i < n; ++i) { + element_levels_[labels[i]] = 0; + } + cur_element_count_ += n; + } + + rabitqlib::index_detail::parallel_for(0, n, workers, [&](size_t i, size_t worker) { + add_point_quant(labels[i], levels[i], link_lists[i], queries[worker]); + }); + + return labels; +} + +} // namespace rabitqlib::hnsw diff --git a/src/index/hnsw_dispatch.cpp b/src/index/hnsw_dispatch.cpp new file mode 100644 index 00000000..f24f8e5a --- /dev/null +++ b/src/index/hnsw_dispatch.cpp @@ -0,0 +1,44 @@ +#include "rabitqlib/index/detail/hnsw_dispatch.hpp" + +#include +#include +#include +#include + +#include "../simd/dispatch_resolver.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::hnsw::detail { +namespace { +using SearchKnnFn = + std::vector> (*)(HierarchicalNSW&, const float*, size_t, size_t); +// The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. +} // namespace +std::vector> search_knn_sorted( + HierarchicalNSW& index, const float* query, size_t topk, size_t ef +) { +#if defined(__x86_64__) || defined(_M_X64) + static const SearchKnnFn kSearchKnnFn = + cpu::has_avx512_popcnt() ? search_knn_sorted_avx512_popcnt + : RABITQ_RESOLVE( + search_knn_sorted_avx512_core, + search_knn_sorted_avx2, + search_knn_sorted_generic, + cpu::has_avx512_core() && cpu::has_avx2() + ); +#else + static const SearchKnnFn kSearchKnnFn = simd::resolve_optional_kernel( + search_knn_sorted_neon, search_knn_sorted_generic, cpu::has_neon() + ); +#endif + + return kSearchKnnFn(index, query, topk, ef); +} +std::priority_queue> search_knn( + HierarchicalNSW& index, const float* query, size_t topk, size_t ef +) { + auto results = search_knn_sorted(index, query, topk, ef); + return {results.begin(), results.end()}; +} +} // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_avx2.cpp b/src/index/hnsw_search_avx2.cpp index d1de0d63..360cfd99 100644 --- a/src/index/hnsw_search_avx2.cpp +++ b/src/index/hnsw_search_avx2.cpp @@ -1,10 +1,12 @@ #include #include #include +#include #include "hnsw_search_avx2_kernels.hpp" -#include "rabitqlib/defines.hpp" +#include "hnsw_search_impl.hpp" // IWYU pragma: keep (backend template definitions) #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw::detail { @@ -27,10 +29,10 @@ struct HnswAvx2Kernel { } }; -maxheap> search_knn_avx2( +std::vector> search_knn_sorted_avx2( HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct( + return index.search_knn_sorted_direct( rotated_query, topk, ef == 0 ? index.ef_ : ef ); } diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index 4b8f6924..e9bef55d 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -5,7 +5,7 @@ #include "../simd/packed_tail_avx2.hpp" #include "../simd/warmup_kernels.hpp" -#include "rabitqlib/index/query.hpp" +#include "rabitqlib/core/query.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_core.cpp b/src/index/hnsw_search_avx512_core.cpp index a7647c76..5a80e902 100644 --- a/src/index/hnsw_search_avx512_core.cpp +++ b/src/index/hnsw_search_avx512_core.cpp @@ -1,11 +1,13 @@ #include #include #include +#include #include "hnsw_search_avx2_kernels.hpp" #include "hnsw_search_avx512_kernels.hpp" -#include "rabitqlib/defines.hpp" +#include "hnsw_search_impl.hpp" // IWYU pragma: keep (backend template definitions) #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw::detail { @@ -28,10 +30,10 @@ struct HnswAvx512CoreKernel { } }; -maxheap> search_knn_avx512_core( +std::vector> search_knn_sorted_avx512_core( HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct( + return index.search_knn_sorted_direct( rotated_query, topk, ef == 0 ? index.ef_ : ef ); } diff --git a/src/index/hnsw_search_avx512_kernels.hpp b/src/index/hnsw_search_avx512_kernels.hpp index eecef63a..43c26c55 100644 --- a/src/index/hnsw_search_avx512_kernels.hpp +++ b/src/index/hnsw_search_avx512_kernels.hpp @@ -4,8 +4,8 @@ #include #include "../simd/warmup_kernels.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/query.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_popcnt.cpp b/src/index/hnsw_search_avx512_popcnt.cpp index c32adc22..db4c6ddd 100644 --- a/src/index/hnsw_search_avx512_popcnt.cpp +++ b/src/index/hnsw_search_avx512_popcnt.cpp @@ -1,10 +1,12 @@ #include #include #include +#include #include "hnsw_search_avx512_kernels.hpp" -#include "rabitqlib/defines.hpp" +#include "hnsw_search_impl.hpp" // IWYU pragma: keep (backend template definitions) #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw::detail { @@ -27,10 +29,10 @@ struct HnswAvx512PopcntKernel { } }; -maxheap> search_knn_avx512_popcnt( +std::vector> search_knn_sorted_avx512_popcnt( HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct( + return index.search_knn_sorted_direct( rotated_query, topk, ef == 0 ? index.ef_ : ef ); } diff --git a/src/index/hnsw_search_generic.cpp b/src/index/hnsw_search_generic.cpp index a49bbf31..fa7abc84 100644 --- a/src/index/hnsw_search_generic.cpp +++ b/src/index/hnsw_search_generic.cpp @@ -1,11 +1,13 @@ #include #include #include +#include -#include "rabitqlib/defines.hpp" +#include "hnsw_search_impl.hpp" // IWYU pragma: keep (backend template definitions) +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw::detail { @@ -30,10 +32,10 @@ struct HnswGenericKernel { } }; -maxheap> search_knn_generic( +std::vector> search_knn_sorted_generic( HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct( + return index.search_knn_sorted_direct( rotated_query, topk, ef == 0 ? index.ef_ : ef ); } diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp new file mode 100644 index 00000000..17f1e5b9 --- /dev/null +++ b/src/index/hnsw_search_impl.hpp @@ -0,0 +1,297 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/visited_pool.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/memory.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib::hnsw { + +template +inline void HierarchicalNSW::get_bin_est_direct( + std::vector& q_to_centroids, + SplitSingleQuery& query_wrapper, + PID candidate_id, + HierarchicalNSW::EstimateRecord& res +) { + if (metric_type_ == METRIC_IP) { + float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + float error = + q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; + res = estimate_split_single_direct( + get_bindata_by_internalid(candidate_id), + query_wrapper, + padded_dim_, + -norm, + error + ); + } else { + // L2 distance + float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + res = estimate_split_single_direct( + get_bindata_by_internalid(candidate_id), + query_wrapper, + padded_dim_, + norm * norm, + norm + ); + } +} + +template +inline void HierarchicalNSW::get_full_est_direct( + std::vector& q_to_centroids, + SplitSingleQuery& query_wrapper, + PID candidate_id, + HierarchicalNSW::EstimateRecord& res +) const { + if (metric_type_ == METRIC_IP) { + float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + float error = + q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; + res = estimate_split_single_full_direct( + get_bindata_by_internalid(candidate_id), + get_exdata_by_internalid(candidate_id), + ip_func_, + query_wrapper, + padded_dim_, + ex_bits_, + -norm, + error + ); + } else { + // L2 distance + float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + res = estimate_split_single_full_direct( + get_bindata_by_internalid(candidate_id), + get_exdata_by_internalid(candidate_id), + ip_func_, + query_wrapper, + padded_dim_, + ex_bits_, + norm * norm, + norm + ); + } +} + +template +inline std::vector> HierarchicalNSW::search_knn_sorted_direct( + const float* rotated_query, size_t k, size_t ef +) { + std::vector> result; + if (cur_element_count_ == 0) { + return result; + } + + SplitSingleQuery query_wrapper( + rotated_query, padded_dim_, ex_bits_, query_config_, metric_type_ + ); + + // Preprocess - get the distance from query to all centroids + std::vector q_to_centroids(num_clusters_); + + if (metric_type_ == METRIC_L2) { + for (size_t i = 0; i < num_clusters_; i++) { + q_to_centroids[i] = std::sqrt(raw_dist_func_( + rotated_query, + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + padded_dim_ + )); + } + } else if (metric_type_ == METRIC_IP) { + q_to_centroids.resize(2 * num_clusters_); + // first half as g_add, second half as g_error + for (size_t i = 0; i < num_clusters_; i++) { + q_to_centroids[i] = dot_product( + rotated_query, + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + padded_dim_ + ); + q_to_centroids[i + num_clusters_] = std::sqrt(euclidean_sqr( + rotated_query, + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + padded_dim_ + )); + } + } + + PID current_id = entry_point_id_; + EstimateRecord curest; + + get_bin_est_direct(q_to_centroids, query_wrapper, current_id, curest); + + for (int level = max_level_; level > 0; level--) { + bool changed = true; + while (changed) { + changed = false; + unsigned int* data; + + data = static_cast(get_linklist(current_id, level)); + int size = get_list_count(data); + + PID* datal = static_cast(data + 1); + for (int i = 0; i < size; i++) { + PID cand = datal[i]; + if (cand > max_elements_) { + throw std::runtime_error("cand error"); + } + + EstimateRecord candest; + get_bin_est_direct(q_to_centroids, query_wrapper, cand, candest); + + if (candest.distance < curest.distance) { + curest = candest; + current_id = cand; + changed = true; + } + } + } + } + + BoundedKNN nearest_candidates(k); + search_base_layer_direct( + current_id, + std::max(ef, k), + k, + query_wrapper, + q_to_centroids, + rotated_query, + nearest_candidates + ); + result.reserve(nearest_candidates.size()); + for (const auto& candidate : nearest_candidates.candidates()) { + result.emplace_back(candidate.record.est_dist, get_external_label(candidate.id)); + } + // Candidate selection uses distance only. Preserve that selection, then match + // the former heap's (distance, external label) ordering within each tie group. + for (auto first = result.begin(); first != result.end();) { + auto last = first + 1; + while (last != result.end() && last->first == first->first) { + ++last; + } + std::sort(first, last); + first = last; + } + return result; +} + +template +inline void HierarchicalNSW::search_base_layer_direct( + PID ep_id, + size_t ef, + size_t k, + SplitSingleQuery& query_wrapper, + std::vector& q_to_centroids, + [[maybe_unused]] const float* query, + BoundedKNN& nearest_candidates +) { + auto vl = visited_list_pool_->get_free_vislist(); + + // Use our bounded priority queue instead of the maxheap. + buffer::SearchBuffer candidate_set(ef); + + float distk = 1e10; + + EstimateRecord start_estimate_record; + if (ex_bits_ > 0) { + get_full_est_direct( + q_to_centroids, query_wrapper, ep_id, start_estimate_record + ); + } else { + get_bin_est_direct( + q_to_centroids, query_wrapper, ep_id, start_estimate_record + ); + } + float est_dist = start_estimate_record.distance; + float low_dist = start_estimate_record.lower_bound; + + // Insert initial candidate. A removed point is still explored, never returned. + if (!is_removed(ep_id)) { + nearest_candidates.insert({ResultRecord(est_dist, low_dist), ep_id}); + distk = est_dist; + } + candidate_set.insert(ep_id, est_dist); + + vl->set(ep_id); + + const size_t prefetch_size = (((padded_dim_ / 8) + 63) / 64) + 1; + const size_t prefetch_lookahead = 4; // Number of neighbors to prefetch in advance. + + while (candidate_set.has_next()) { + // Step 1 - get the next node to explore. + PID current_node_id = candidate_set.pop(); + int* data = (int*)get_linklist0(current_node_id); + size_t size = get_list_count((PID*)data); + + for (size_t p = 0; p < prefetch_lookahead; ++p) { + rabitqlib::memory::mem_prefetch_l1( + get_bindata_by_internalid(*(data + 1 + p)), prefetch_size + ); + } + // Iterate over neighbors. (List starts at index 1.) + for (size_t j = 1; j <= size; j++) { + int candidate_id = *(data + j); + + if (j + prefetch_lookahead <= size) { + rabitqlib::memory::mem_prefetch_l1( + get_bindata_by_internalid(*(data + j + prefetch_lookahead)), + prefetch_size + ); + } + + if (vl->get(candidate_id)) { + continue; + } + vl->set(candidate_id); + + EstimateRecord candest; + get_bin_est_direct( + q_to_centroids, query_wrapper, candidate_id, candest + ); + + // A removed neighbor keeps its binary estimate and stays in candidate_set. + bool update_results = + !is_removed(static_cast(candidate_id)) && + (nearest_candidates.size() < k || candest.lower_bound < distk); + + if (update_results) { + // Compute the full estimate if promising. + if (ex_bits_ > 0) { + get_full_est_direct( + q_to_centroids, query_wrapper, candidate_id, candest + ); + } + Candidate cand{ + ResultRecord(candest.distance, candest.lower_bound), + static_cast(candidate_id)}; + nearest_candidates.insert(cand); + distk = nearest_candidates.worst().record.est_dist; + } + + if (!candidate_set.is_full(candest.distance)) { + candidate_set.insert(candidate_id, candest.distance); + } + + rabitqlib::memory::mem_prefetch_l2( + (char*)get_linklist0(candidate_set.next_id()), 2 + ); + } + } + + visited_list_pool_->release_vis_list(std::move(vl)); +} + +} // namespace rabitqlib::hnsw diff --git a/src/index/hnsw_search_neon.cpp b/src/index/hnsw_search_neon.cpp index 0d8473f5..c0c8f782 100644 --- a/src/index/hnsw_search_neon.cpp +++ b/src/index/hnsw_search_neon.cpp @@ -1,12 +1,12 @@ #include #include #include +#include #include "../simd/hnsw_neon_kernels.hpp" -#include "rabitqlib/defines.hpp" +#include "hnsw_search_impl.hpp" // IWYU pragma: keep (backend template definitions) #include "rabitqlib/index/hnsw/hnsw.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::hnsw::detail { @@ -33,10 +33,10 @@ struct HnswNeonKernel { } }; -maxheap> search_knn_neon( +std::vector> search_knn_sorted_neon( HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef ) { - return index.search_knn_direct( + return index.search_knn_sorted_direct( rotated_query, topk, ef == 0 ? index.ef_ : ef ); } diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp new file mode 100644 index 00000000..68267a47 --- /dev/null +++ b/src/index/ivf.cpp @@ -0,0 +1,1102 @@ +#include "rabitqlib/index/ivf/ivf.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ivf_search.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/parallel_exception.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/ivf/cluster.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/path.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::ivf { + +IVF::IVF() noexcept = default; +IVF::~IVF() = default; + +bool IVF::exceeds_storage(size_t n) const { + const size_t max_size = std::numeric_limits::max(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const size_t rerank_bytes = rerank_vector_bytes(); + return n > max_size / sizeof(PID) || n > max_size / sizeof(float) / padded_dim_ || + n > max_size / batch_bytes || (rerank_bytes != 0 && n > max_size / rerank_bytes); +} + +size_t IVF::batch_data_bytes(const std::vector& cluster_sizes) const { + assert(cluster_sizes.size() == num_clusters_); + size_t total_blocks = 0; + for (auto size : cluster_sizes) { + total_blocks += div_round_up(size, fastscan::kBatchSize); + } + return total_blocks * BatchDataMap::data_bytes(padded_dim_); +} + +// Keep the original inlining intent inside this translation unit, especially the +// per-vector stride calculation. Public consumers only need the declarations. +inline size_t IVF::rerank_vector_bytes() const { + return raw_reranking_ ? sizeof(float) * dim_ + : ExDataMap::data_bytes(padded_dim_, ex_bits_); +} + +inline size_t IVF::ex_data_bytes() const { return rerank_vector_bytes() * max_elements_; } + +void IVF::free_memory() { + initializer_.reset(); + ByteStorage().swap(batch_storage_); + ByteStorage().swap(ex_storage_); + FloatStorage().swap(raw_storage_); + IdStorage().swap(id_storage_); +} + +void IVF::swap(IVF& other) noexcept { + using std::swap; + swap(initializer_type_, other.initializer_type_); + swap(initializer_, other.initializer_); + swap(batch_storage_, other.batch_storage_); + swap(ex_storage_, other.ex_storage_); + swap(raw_storage_, other.raw_storage_); + swap(id_storage_, other.id_storage_); + swap(max_elements_, other.max_elements_); + swap(num_removed_, other.num_removed_); + swap(dim_, other.dim_); + swap(padded_dim_, other.padded_dim_); + swap(num_clusters_, other.num_clusters_); + swap(raw_reranking_, other.raw_reranking_); + swap(ready_, other.ready_); + swap(ex_bits_, other.ex_bits_); + swap(rotator_type_, other.rotator_type_); + swap(rotator_, other.rotator_); + swap(clusters_, other.clusters_); + swap(metric_type_, other.metric_type_); + swap(ip_func_, other.ip_func_); +} + +IVF::IVF( + size_t num_points, + size_t dim, + size_t num_clusters, + size_t total_bits, + MetricType metric_type, + RotatorType rotator_type, + InitializerType initializer +) + : initializer_type_(resolve_initializer_type(initializer, num_clusters)) + , max_elements_(num_points) + , dim_(dim) + , padded_dim_(dim) + , num_clusters_(num_clusters) + , raw_reranking_(total_bits == 32) + , ex_bits_(total_bits == 32 ? 0 : total_bits - 1) + , rotator_type_(rotator_type) + , metric_type_(metric_type) { + validate_metric_type(metric_type); + if ((total_bits < 1 || total_bits > 9) && total_bits != 32) { + throw std::invalid_argument("IVF bits must be in [1, 9] or 32 for raw reranking"); + }; + if (num_points == 0 || num_points > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("IVF point count exceeds the supported ID range"); + } + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("IVF cluster count exceeds the supported ID range"); + } + if (dim == 0 || dim > (std::numeric_limits::max() / 32) - 64) { + throw std::invalid_argument("IVF dimension is invalid or too large"); + } + padded_dim_ = round_up_to_multiple(dim_, 32); + const size_t max_size = std::numeric_limits::max(); + if (exceeds_storage(num_points) || + num_clusters > max_size / sizeof(float) / padded_dim_ || + (rotator_type == RotatorType::MatrixRotator && + dim > max_size / sizeof(float) / padded_dim_)) { + throw std::invalid_argument("IVF configuration exceeds addressable storage"); + } + rotator_.reset(choose_rotator(dim, rotator_type, padded_dim_)); + /* check size */ + assert(padded_dim_ % 32 == 0); + assert(padded_dim_ >= dim_); +} + +/** + * @brief Construct clusters in IVF + * + * @param data Data objects (N*DIM) + * @param centroids Centroid vectors (K*DIM) + * @param cluster_ids Cluster ID for each data object + */ +void IVF::construct( + const float* data, + const float* centroids, + const PID* cluster_ids, + bool faster, + size_t num_threads +) { + if (max_elements_ == 0 || dim_ == 0 || num_clusters_ == 0 || rotator_ == nullptr) { + throw std::logic_error("IVF must be configured before construction"); + } + if (data == nullptr || centroids == nullptr || cluster_ids == nullptr) { + throw std::invalid_argument("IVF construction inputs must not be null"); + } + + // get id list for each cluster + std::vector counts(num_clusters_, 0); + std::vector> id_lists(num_clusters_); + for (size_t i = 0; i < max_elements_; ++i) { + PID cid = cluster_ids[i]; + if (cid >= num_clusters_) { + throw std::invalid_argument("Cluster ID is out of range"); + } + id_lists[cid].push_back(static_cast(i)); + counts[cid] += 1; + } + + allocate_memory(counts); + + // init the cluster list + init_clusters(counts); + + // all rotated centroids + std::vector rotated_centroids(num_clusters_ * padded_dim_); + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + num_threads = std::min(resolve_num_threads(num_threads), num_clusters_); + const size_t scratch_points = + std::min(fastscan::kBatchSize, *std::max_element(counts.begin(), counts.end())); + // Each worker reuses one batch, independently of the largest cluster size. + // Allocate before entering OpenMP so allocation failures reach the caller. + std::vector> rotated_blocks(num_threads); + for (auto& block : rotated_blocks) { + block.resize(scratch_points * padded_dim_); + } + /* Quantize each cluster */ + { + index_detail::ParallelException failure; +#pragma omp parallel for schedule(dynamic) num_threads(num_threads) + for (std::ptrdiff_t index = 0; index < static_cast(num_clusters_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + const float* cur_centroid = centroids + (i * dim_); + float* cur_rotated_c = &rotated_centroids[i * padded_dim_]; + Cluster& cp = clusters_[i]; + quantize_cluster( + cp, + id_lists[i], + data, + cur_centroid, + cur_rotated_c, + rotated_blocks[static_cast(omp_get_thread_num())].data(), + config + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + + this->initializer_->add_vectors(rotated_centroids.data(), num_threads); + ready_ = true; +} + +void IVF::allocate_memory(const std::vector& cluster_sizes) { + ready_ = false; + num_removed_ = 0; + free_memory(); + clusters_.clear(); + + if (initializer_type_ == InitializerType::Flat) { + this->initializer_ = + std::make_unique(padded_dim_, num_clusters_, metric_type_); + } else if (initializer_type_ == InitializerType::FlatRaBitQ) { + this->initializer_ = std::make_unique( + padded_dim_, num_clusters_, metric_type_ + ); + } else { + this->initializer_ = + std::make_unique(padded_dim_, num_clusters_, metric_type_); + } + batch_storage_ = ByteStorage(batch_data_bytes(cluster_sizes)); + if (rerank_vector_bytes() > 0) { + if (raw_reranking_) { + raw_storage_ = FloatStorage(max_elements_ * dim_); + } else { + ex_storage_ = ByteStorage(ex_data_bytes()); + } + } + id_storage_ = IdStorage(max_elements_); + + this->ip_func_ = select_excode_ipfunc(ex_bits_); +} + +/** + * @brief Initialize the cluster list by finding each cluster's storage offsets. + */ +void IVF::init_clusters(const std::vector& cluster_sizes) { + this->clusters_.reserve(num_clusters_); + size_t added_vectors = 0; + size_t added_batches = 0; + for (size_t i = 0; i < num_clusters_; ++i) { + // find data location for current cluster + size_t num = cluster_sizes[i]; + size_t num_batches = div_round_up(num, fastscan::kBatchSize); + + char* current_batch_data = + batch_data() + (BatchDataMap::data_bytes(padded_dim_) * added_batches); + char* current_ex_data = rerank_vector_bytes() > 0 + ? ex_data() + (added_vectors * rerank_vector_bytes()) + : nullptr; + PID* cluster_ids = ids() + added_vectors; + + this->clusters_.emplace_back(num, current_batch_data, current_ex_data, cluster_ids); + + added_vectors += num; + added_batches += num_batches; + } +} + +void IVF::quantize_cluster( + Cluster& cp, + const std::vector& point_ids, + const float* data, + const float* cur_centroid, + float* rotated_centroid, + float* rotated_data, + const quant::RabitqConfig& config +) { + size_t num_points = point_ids.size(); + if (cp.num() != num_points) { + throw std::invalid_argument("Cluster size and ID count differ"); + } + std::copy(point_ids.begin(), point_ids.end(), cp.ids()); + this->rotator_->rotate(cur_centroid, rotated_centroid); + + char* batch_data = cp.batch_data(); + char* ex_data = cp.ex_data(); + for (size_t i = 0; i < num_points; i += fastscan::kBatchSize) { + size_t n = std::min(fastscan::kBatchSize, num_points - i); + for (size_t j = 0; j < n; ++j) { + const float* vector = data + (point_ids[i + j] * dim_); + rotator_->rotate(vector, rotated_data + (j * padded_dim_)); + if (raw_reranking_) { + std::memcpy( + cp.ex_data() + ((i + j) * rerank_vector_bytes()), + vector, + rerank_vector_bytes() + ); + } + } + + quant::quantize_split_batch( + rotated_data, + rotated_centroid, + n, + padded_dim_, + ex_bits_, + batch_data, + ex_data, + metric_type_, + config + ); + + batch_data += BatchDataMap::data_bytes(padded_dim_); + if (ex_bits_ > 0) { + ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * n; + } + } +} + +void IVF::add( + const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads +) { + if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || + clusters_.size() != num_clusters_ || batch_storage_.empty() || + id_storage_.size() != max_elements_) { + throw std::logic_error("IVF index must be constructed or loaded before add"); + } + if (n == 0) { + return; + } + if (data == nullptr) { + throw std::invalid_argument("IVF add data must not be null"); + } + if (n > buffer::kSearchBufferMaxPointCount - max_elements_) { + throw std::invalid_argument("IVF point count exceeds the supported ID range"); + } + const size_t new_num_points = max_elements_ + n; + if (exceeds_storage(new_num_points)) { + throw std::invalid_argument("IVF configuration exceeds addressable storage"); + } + + /* Assign every new point to a cluster */ + std::vector assigned(n); + if (cluster_ids != nullptr) { + for (size_t i = 0; i < n; ++i) { + if (cluster_ids[i] >= num_clusters_) { + throw std::invalid_argument("Cluster ID is out of range"); + } + assigned[i] = cluster_ids[i]; + } + } else { + rabitqlib::index_detail::parallel_for( + 0, + n, + num_threads, + [&](size_t row, size_t /*thread_id*/) { + std::vector rotated(padded_dim_); + std::vector> nearest(1); + rotator_->rotate(data + (row * dim_), rotated.data()); + initializer_->centroids_distances(rotated.data(), 1, nearest); + assigned[row] = nearest[0].id; + } + ); + } + + std::vector> rows_of(num_clusters_); + for (size_t i = 0; i < n; ++i) { + rows_of[assigned[i]].push_back(static_cast(i)); + } + std::vector new_sizes(num_clusters_); + for (size_t i = 0; i < num_clusters_; ++i) { + new_sizes[i] = clusters_[i].num() + rows_of[i].size(); + } + + /* Build the grown storage on the side so the index is unchanged on failure */ + ByteStorage new_batch(batch_data_bytes(new_sizes)); + ByteStorage new_ex; + FloatStorage new_raw; + const size_t rerank_bytes = rerank_vector_bytes(); + if (raw_reranking_) { + new_raw = FloatStorage(new_num_points * dim_); + } else if (rerank_bytes > 0) { + new_ex = ByteStorage(new_num_points * rerank_bytes); + } + IdStorage new_ids(new_num_points); + char* new_rerank = raw_reranking_ ? reinterpret_cast(new_raw.data()) + : reinterpret_cast(new_ex.data()); + + std::vector batch_offset(num_clusters_); + std::vector vector_offset(num_clusters_); + size_t added_batches = 0; + size_t added_vectors = 0; + for (size_t i = 0; i < num_clusters_; ++i) { + batch_offset[i] = added_batches; + vector_offset[i] = added_vectors; + added_batches += div_round_up(new_sizes[i], fastscan::kBatchSize); + added_vectors += new_sizes[i]; + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const PID first_pid = static_cast(max_elements_); + // An exception must not leave an OpenMP region, so keep the first one and rethrow it. + index_detail::ParallelException failure; +#pragma omp parallel for schedule(dynamic) num_threads(resolve_num_threads(num_threads)) + for (std::ptrdiff_t index = 0; index < static_cast(num_clusters_); + ++index) { + const size_t i = static_cast(index); + try { + extend_cluster( + clusters_[i], + rows_of[i], + data, + initializer_->centroid(static_cast(i)), + first_pid, + reinterpret_cast(new_batch.data()) + (batch_offset[i] * batch_bytes), + rerank_bytes > 0 ? new_rerank + (vector_offset[i] * rerank_bytes) : nullptr, + new_ids.data() + vector_offset[i], + config + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + + /* Commit. Rebuilding the cluster list reuses its capacity, so it cannot throw. */ + batch_storage_.swap(new_batch); + ex_storage_.swap(new_ex); + raw_storage_.swap(new_raw); + id_storage_.swap(new_ids); + max_elements_ = new_num_points; + clusters_.clear(); + init_clusters(new_sizes); +} + +/** + * @brief Copy one cluster into its grown storage and append new points to it. + * + * Full batches, rerank data and IDs are copied unchanged. The old partial tail batch is + * unpacked and repacked together with the new points, so the result is byte-identical to + * quantizing the same points in the same order in one construct call. + */ +void IVF::extend_cluster( + const Cluster& old_cluster, + const std::vector& rows, + const float* data, + const float* rotated_centroid, + PID first_pid, + char* batch_dst, + char* rerank_dst, + PID* ids_dst, + const quant::RabitqConfig& config +) const { + const size_t old_num = old_cluster.num(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const size_t rerank_bytes = rerank_vector_bytes(); + const size_t cols = padded_dim_ / 8; + + if (rows.empty()) { + const size_t old_batches = div_round_up(old_num, fastscan::kBatchSize); + std::copy_n(old_cluster.batch_data(), old_batches * batch_bytes, batch_dst); + std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + std::copy_n(old_cluster.ids(), old_num, ids_dst); + return; + } + + const size_t full_batches = old_num / fastscan::kBatchSize; + const size_t kept = old_num - (full_batches * fastscan::kBatchSize); + const size_t tail_num = kept + rows.size(); + + std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); + std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + std::copy_n(old_cluster.ids(), old_num, ids_dst); + + /* Codes and factors of every point in the tail, old points first */ + std::vector codes(tail_num * cols); + std::vector f_add(tail_num); + std::vector f_rescale(tail_num); + std::vector f_error(tail_num); + if (kept > 0) { + ConstBatchDataMap old_tail( + old_cluster.batch_data() + (full_batches * batch_bytes), padded_dim_ + ); + fastscan::unpack_codes(padded_dim_, old_tail.bin_code(), kept, codes.data()); + old_tail.f_add().copy_to(f_add.data(), kept); + old_tail.f_rescale().copy_to(f_rescale.data(), kept); + old_tail.f_error().copy_to(f_error.data(), kept); + } + + std::vector rotated(padded_dim_); + for (size_t j = 0; j < rows.size(); ++j) { + const float* vector = data + (static_cast(rows[j]) * dim_); + const size_t lane = kept + j; + rotator_->rotate(vector, rotated.data()); + quant::quantize_compact_one_bit( + rotated.data(), + rotated_centroid, + padded_dim_, + codes.data() + (lane * cols), + f_add[lane], + f_rescale[lane], + f_error[lane], + metric_type_ + ); + char* rerank = rerank_dst + ((old_num + j) * rerank_bytes); + if (raw_reranking_) { + std::memcpy(rerank, vector, rerank_bytes); + } else if (ex_bits_ > 0) { + quant::quantize_compact_ex_bits( + rotated.data(), + rotated_centroid, + padded_dim_, + ex_bits_, + rerank, + metric_type_, + config + ); + } + ids_dst[old_num + j] = first_pid + rows[j]; + } + + /* Repack the tail into batches */ + char* batch = batch_dst + (full_batches * batch_bytes); + for (size_t start = 0; start < tail_num; start += fastscan::kBatchSize) { + const size_t count = std::min(fastscan::kBatchSize, tail_num - start); + BatchDataMap cur_batch(batch, padded_dim_); + fastscan::pack_codes( + padded_dim_, codes.data() + (start * cols), count, cur_batch.bin_code() + ); + for (size_t i = 0; i < count; ++i) { + cur_batch.f_add()[i] = f_add[start + i]; + cur_batch.f_rescale()[i] = f_rescale[start + i]; + cur_batch.f_error()[i] = f_error[start + i]; + } + batch += batch_bytes; + } +} + +size_t IVF::remove(const PID* ids_to_remove, size_t n) { + if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || + clusters_.size() != num_clusters_ || batch_storage_.empty() || + id_storage_.size() != max_elements_) { + throw std::logic_error("IVF index must be constructed or loaded before remove"); + } + if (n == 0) { + return 0; + } + if (ids_to_remove == nullptr) { + throw std::invalid_argument("IVF remove IDs must not be null"); + } + + // Validate every ID before changing anything + std::vector marked(max_elements_, false); + for (size_t i = 0; i < n; ++i) { + if (ids_to_remove[i] >= max_elements_) { + throw std::invalid_argument("IVF remove ID is out of range"); + } + marked[ids_to_remove[i]] = true; + } + + // A point whose f_add is +inf has +inf estimated and lower-bound distances in every + // scan path, so it never enters the result buffer or triggers reranking. + constexpr float kRemoved = std::numeric_limits::infinity(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + size_t removed = 0; + for (const auto& cluster : clusters_) { + const PID* cluster_ids = cluster.ids(); + for (size_t pos = 0; pos < cluster.num(); ++pos) { + if (!marked[cluster_ids[pos]]) { + continue; + } + BatchDataMap cur_batch( + cluster.batch_data() + ((pos / fastscan::kBatchSize) * batch_bytes), + padded_dim_ + ); + const size_t lane = pos % fastscan::kBatchSize; + if (static_cast(cur_batch.f_add()[lane]) != kRemoved) { + cur_batch.f_add()[lane] = kRemoved; + ++removed; + } + } + } + num_removed_ += removed; + return removed; +} + +void IVF::save(const char* filename) const { + if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || + clusters_.size() != num_clusters_ || batch_storage_.empty() || + id_storage_.size() != max_elements_) { + throw std::logic_error("Cannot save an unconstructed IVF index"); + } + if (filename == nullptr || filename[0] == '\0') { + throw std::invalid_argument("IVF save filename must not be empty"); + } + + std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + output.exceptions(std::ios::failbit | std::ios::badbit); + const uint32_t flags = raw_reranking_ ? 1 : 0; + const uint64_t stored_padded_dim = padded_dim_; + output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); + output.write(reinterpret_cast(&kFormatVersion), sizeof(kFormatVersion)); + output.write(reinterpret_cast(&flags), sizeof(flags)); + output.write( + reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim) + ); + + const auto stored_initializer = static_cast(initializer_type_); + output.write( + reinterpret_cast(&stored_initializer), sizeof(stored_initializer) + ); + + /* Save meta data */ + output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); + output.write(reinterpret_cast(&dim_), sizeof(size_t)); + output.write(reinterpret_cast(&num_clusters_), sizeof(size_t)); + output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); + output.write(reinterpret_cast(&rotator_type_), sizeof(rotator_type_)); + output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); + + /* Save number of vectors of each cluster */ + std::vector cluster_sizes; + cluster_sizes.reserve(num_clusters_); + for (const auto& cur_cluster : clusters_) { + cluster_sizes.push_back(cur_cluster.num()); + } + output.write( + reinterpret_cast(cluster_sizes.data()), + static_cast(sizeof(size_t) * num_clusters_) + ); + + /* Save rotator */ + this->rotator_->save(output); + + /* Save data */ + this->initializer_->save(output, filename); + output.write( + batch_data(), static_cast(batch_data_bytes(cluster_sizes)) + ); + if (ex_data_bytes() != 0) { + output.write(ex_data(), static_cast(ex_data_bytes())); + } + output.write( + reinterpret_cast(ids()), static_cast(ids_bytes()) + ); + + output.close(); +} + +void IVF::load(const char* filename) { + if (filename == nullptr || filename[0] == '\0') { + throw std::invalid_argument("IVF load filename must not be empty"); + } + std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + if (!input.is_open()) { + throw std::runtime_error("Cannot open IVF index file"); + } + + IVF loaded; + input.exceptions(std::ios::failbit | std::ios::badbit); + input.seekg(0, std::ios::end); + const auto file_bytes = static_cast(input.tellg()); + input.seekg(0); + uint64_t magic = 0; + input.read(reinterpret_cast(&magic), sizeof(magic)); + const bool explicit_padding = magic == kFormatMagic; + if (explicit_padding) { + uint32_t version = 0, flags = 0; + uint64_t stored_padded_dim = 0; + input.read(reinterpret_cast(&version), sizeof(version)); + input.read(reinterpret_cast(&flags), sizeof(flags)); + input.read(reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim)); + if (version != 1 && version != kFormatVersion) { + throw std::runtime_error("Unsupported IVF index version"); + } + if (flags > 1 || stored_padded_dim > std::numeric_limits::max()) { + throw std::runtime_error("Invalid IVF index metadata"); + } + loaded.raw_reranking_ = flags == 1; + loaded.padded_dim_ = static_cast(stored_padded_dim); + if (version >= 2) { + uint32_t initializer = 0; + input.read(reinterpret_cast(&initializer), sizeof(initializer)); + if (initializer < static_cast(InitializerType::Flat) || + initializer > static_cast(InitializerType::HNSW)) { + throw std::runtime_error("Invalid IVF initializer type in index file"); + } + loaded.initializer_type_ = static_cast(initializer); + } + } else if (magic == kRawFormatMagic) { + loaded.raw_reranking_ = true; + uint32_t version = 0; + input.read(reinterpret_cast(&version), sizeof(version)); + if (version != kRawFormatVersion) { + throw std::runtime_error("Unsupported raw IVF index version"); + } + } else { + input.seekg(0); + } + + /* Load meta data */ + input.read(reinterpret_cast(&loaded.max_elements_), sizeof(size_t)); + input.read(reinterpret_cast(&loaded.dim_), sizeof(size_t)); + input.read(reinterpret_cast(&loaded.num_clusters_), sizeof(size_t)); + input.read(reinterpret_cast(&loaded.ex_bits_), sizeof(size_t)); + input.read( + reinterpret_cast(&loaded.rotator_type_), sizeof(loaded.rotator_type_) + ); + input.read(reinterpret_cast(&loaded.metric_type_), sizeof(loaded.metric_type_)); + validate_metric_type(loaded.metric_type_); + // Missing routing metadata means a historical file. Its payload layout must + // follow the original threshold, independently of the current Auto policy. + if (loaded.initializer_type_ == InitializerType::Auto) { + loaded.initializer_type_ = + loaded.num_clusters_ < 20000 ? InitializerType::Flat : InitializerType::HNSW; + } + if (loaded.max_elements_ == 0 || + loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || loaded.dim_ == 0 || + loaded.num_clusters_ == 0 || + loaded.num_clusters_ > buffer::kSearchBufferMaxPointCount || loaded.ex_bits_ > 8 || + (loaded.raw_reranking_ && loaded.ex_bits_ != 0) || + loaded.dim_ > (std::numeric_limits::max() / 32) - 64) { + throw std::runtime_error("Invalid IVF index metadata"); + } + if (!explicit_padding) { + // Both historical formats omit padded_dim and always used 64-coordinate blocks. + loaded.padded_dim_ = round_up_to_multiple(loaded.dim_, 64); + } + if (loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || + loaded.padded_dim_ > (std::numeric_limits::max() / 32) - 64) { + throw std::runtime_error("Invalid padded dimension in IVF index file"); + } + // Bound every payload by the actual file before allocating or multiplying sizes. + size_t remaining = file_bytes - static_cast(input.tellg()); + const auto consume = [&remaining](size_t count, size_t bytes) { + if (bytes != 0 && count > remaining / bytes) { + throw std::runtime_error("Invalid or truncated IVF index payload"); + } + remaining -= count * bytes; + }; + consume(loaded.num_clusters_, sizeof(size_t)); + consume(loaded.max_elements_, sizeof(PID)); + consume(loaded.max_elements_, loaded.rerank_vector_bytes()); + if (loaded.rotator_type_ == RotatorType::MatrixRotator) { + consume(loaded.dim_, sizeof(float) * loaded.padded_dim_); + } else if (loaded.rotator_type_ == RotatorType::FhtKacRotator) { + consume(1, loaded.padded_dim_ / 2); + } else { + throw std::runtime_error("Invalid IVF rotator type"); + } + if (loaded.initializer_type_ != InitializerType::HNSW) { + consume(loaded.num_clusters_, sizeof(float) * loaded.padded_dim_); + } + if (loaded.initializer_type_ == InitializerType::FlatRaBitQ) { + consume(loaded.padded_dim_, sizeof(float)); + consume( + div_round_up(loaded.num_clusters_, fastscan::kBatchSize), + BatchDataMap::data_bytes(loaded.padded_dim_) + ); + } + + /* Load number of vectors of each cluster */ + std::vector cluster_sizes(loaded.num_clusters_, 0); + input.read( + reinterpret_cast(cluster_sizes.data()), + static_cast(sizeof(size_t) * loaded.num_clusters_) + ); + + size_t total = 0; + for (size_t size : cluster_sizes) { + if (size > loaded.max_elements_ - total) { + throw std::runtime_error("Invalid cluster counts in IVF index file"); + } + total += size; + consume( + div_round_up(size, fastscan::kBatchSize), + BatchDataMap::data_bytes(loaded.padded_dim_) + ); + } + if (total != loaded.max_elements_) { + throw std::runtime_error("Invalid cluster counts in IVF index file"); + } + loaded.rotator_ = rotator_impl::make_rotator_for_loading( + loaded.dim_, loaded.rotator_type_, loaded.padded_dim_ + ); + + /* Load rotator */ + loaded.rotator_->load(input); + + /* Load data */ + loaded.allocate_memory(cluster_sizes); + loaded.initializer_->load(input, filename); + input.read( + loaded.batch_data(), + static_cast(loaded.batch_data_bytes(cluster_sizes)) + ); + if (loaded.ex_data_bytes() != 0) { + input.read(loaded.ex_data(), static_cast(loaded.ex_data_bytes())); + } + input.read( + reinterpret_cast(loaded.ids()), + static_cast(loaded.ids_bytes()) + ); + for (size_t i = 0; i < loaded.max_elements_; ++i) { + if (loaded.ids()[i] >= loaded.max_elements_ || + loaded.ids()[i] >= buffer::kSearchBufferCheckedMask) { + throw std::runtime_error("Invalid point ID in IVF index file"); + } + } + + /* Init each cluster */ + loaded.init_clusters(cluster_sizes); + // Removal markers are already persisted, including in historical files. + // Count only logical rows: FastScan padding lanes are not stored points. + const size_t batch_bytes = BatchDataMap::data_bytes(loaded.padded_dim_); + for (const auto& cluster : loaded.clusters_) { + for (size_t first = 0; first < cluster.num(); first += fastscan::kBatchSize) { + ConstBatchDataMap batch( + cluster.batch_data() + (first / fastscan::kBatchSize) * batch_bytes, + loaded.padded_dim_ + ); + const size_t count = std::min(fastscan::kBatchSize, cluster.num() - first); + for (size_t lane = 0; lane < count; ++lane) { + loaded.num_removed_ += + batch.f_add()[lane] == std::numeric_limits::infinity(); + } + } + } + loaded.ready_ = true; + + input.close(); + swap(loaded); +} + +void IVF::search( + const float* query, size_t k, size_t nprobe, PID* results, float* distances +) const { + search(query, k, nprobe, results, distances, !raw_reranking_ && ex_bits_ > 2); +} + +void IVF::search( + const float* __restrict__ query, + size_t k, + size_t nprobe, + PID* __restrict__ results, + bool use_hacc +) const { + this->search(query, k, nprobe, results, nullptr, use_hacc); +} + +void IVF::search( + const float* __restrict__ query, + size_t k, + size_t nprobe, + PID* __restrict__ results, + float* __restrict__ distances, + bool use_hacc +) const { + validate_search(query, k, nprobe, results); + + nprobe = std::min(nprobe, num_clusters_); // corner case + std::vector rotated_query(padded_dim_); + this->rotator_->rotate(query, rotated_query.data()); + + // use initer to get closest nprobe centroids + std::vector> centroid_dist(nprobe); + this->initializer_->centroids_distances(rotated_query.data(), nprobe, centroid_dist); + + buffer::SearchBuffer knns(k); + + SplitBatchQuery q_obj( + rotated_query.data(), padded_dim_, ex_bits_, metric_type_, use_hacc + ); + + for (size_t i = 0; i < nprobe; ++i) { + PID cid = centroid_dist[i].id; + float dist = centroid_dist[i].distance; + const Cluster& cur_cluster = clusters_[cid]; + + if (metric_type_ == METRIC_L2) { + q_obj.set_g_add(dist); + } else if (metric_type_ == METRIC_IP) { + const float residual_norm = std::sqrt(euclidean_sqr( + rotated_query.data(), initializer_->centroid(cid), padded_dim_ + )); + auto g_add_ip = dot_product( + rotated_query.data(), initializer_->centroid(cid), padded_dim_ + ); + q_obj.set_g_add(residual_norm, g_add_ip); + } else { + // unsupported + throw std::invalid_argument("Quantization only supports L2 and IP metrics"); + } + search_cluster(cur_cluster, q_obj, knns, use_hacc, query); + } + + const size_t found = knns.size(); + if (distances != nullptr) { + knns.copy_results(results, distances); + std::fill(distances + found, distances + k, std::numeric_limits::infinity()); + } else { + knns.copy_results(results); + } + std::fill(results + found, results + k, kPidMax); +} + +void IVF::search_batch( + const float* queries, + size_t num_queries, + size_t k, + size_t nprobe, + PID* results, + float* distances, + std::optional use_hacc, + size_t num_threads +) const { + if (num_queries == 0) { + return; + } + validate_search(queries, k, nprobe, results); + const size_t max_values = std::numeric_limits::max() / sizeof(float); + if (num_queries > max_values / dim_ || num_queries > max_values / k) { + throw std::length_error("IVF query batch is too large"); + } + nprobe = std::min(nprobe, num_clusters_); + const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); + const size_t workers = num_queries == 1 || num_threads == 1 + ? 1 + : std::min(resolve_num_threads(num_threads), num_queries); + if (workers == 1) { + for (size_t i = 0; i < num_queries; ++i) { + search( + queries + i * dim_, + k, + nprobe, + results + i * k, + distances == nullptr ? nullptr : distances + i * k, + hacc + ); + } + return; + } + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(workers) schedule(dynamic) + for (std::ptrdiff_t query_index = 0; + query_index < static_cast(num_queries); + ++query_index) { + if (failure.failed()) { + continue; + } + const size_t i = static_cast(query_index); + try { + search( + queries + i * dim_, + k, + nprobe, + results + i * k, + distances == nullptr ? nullptr : distances + i * k, + hacc + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); +} + +inline void IVF::validate_search( + const float* query, size_t k, size_t nprobe, const PID* results +) const { + if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || + clusters_.size() != num_clusters_ || batch_storage_.empty() || + id_storage_.size() != max_elements_) { + throw std::logic_error("IVF index must be constructed or loaded before search"); + } + if (query == nullptr) { + throw std::invalid_argument("IVF search query must not be null"); + } + if (results == nullptr) { + throw std::invalid_argument("IVF search results must not be null"); + } + if (k == 0 || k > max_elements_ || k > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("IVF search k must be between 1 and the point count"); + } + if (nprobe == 0) { + throw std::invalid_argument("IVF search nprobe must be positive"); + } +} + +inline void IVF::search_cluster( + const Cluster& cur_cluster, + const SplitBatchQuery& q_obj, + buffer::SearchBuffer& knns, + bool use_hacc, + const float* query +) const { + size_t iter = cur_cluster.num() / fastscan::kBatchSize; + size_t remain = cur_cluster.num() - (iter * fastscan::kBatchSize); + + const char* batch_data = cur_cluster.batch_data(); + const char* ex_data = cur_cluster.ex_data(); + const PID* ids = cur_cluster.ids(); + + /* Compute distances block by block */ + for (size_t i = 0; i < iter; ++i) { + scan_one_batch( + batch_data, ex_data, ids, q_obj, knns, fastscan::kBatchSize, use_hacc, query + ); + + batch_data += BatchDataMap::data_bytes(padded_dim_); + if (rerank_vector_bytes() > 0) { + ex_data += rerank_vector_bytes() * fastscan::kBatchSize; + } + ids += fastscan::kBatchSize; + } + + if (remain > 0) { + // scan the last block + scan_one_batch(batch_data, ex_data, ids, q_obj, knns, remain, use_hacc, query); + } +} + +inline void IVF::scan_one_batch( + const char* batch_data, + const char* ex_data, + const PID* ids, + const SplitBatchQuery& q_obj, + buffer::SearchBuffer& knns, + size_t num_points, + bool use_hacc, + const float* query +) const { + std::array est_distance; + std::array low_distance; + std::array ip_x0_qr; // inner product of the 1st bit + + split_batch_estdist( + batch_data, + q_obj, + padded_dim_, + est_distance.data(), + low_distance.data(), + ip_x0_qr.data(), + use_hacc + ); + + float distk = knns.top_dist(); + + // Without reranking data, return the one-bit estimates directly. + if (ex_bits_ == 0 && !raw_reranking_) { + detail::insert_candidates(knns, ids, est_distance.data(), num_points); + return; + } + + // Refine candidates passing the probabilistic one-bit filter. Quantized storage + // reuses ip_x0_qr with the extra bits; raw storage scores the original vector. + for (size_t i = 0; i < num_points; ++i) { + float lower_dist = low_distance[i]; + if (lower_dist < distk) { + PID id = ids[i]; + float ex_dist; + if (raw_reranking_) { + const auto* vector = reinterpret_cast(ex_data); + ex_dist = metric_type_ == METRIC_L2 ? euclidean_sqr(query, vector, dim_) + : dot_product_dis(query, vector, dim_); + } else { + ex_dist = split_distance_boosting( + ex_data, ip_func_, q_obj, padded_dim_, ex_bits_, ip_x0_qr[i] + ); + } + knns.insert(id, ex_dist); + distk = knns.top_dist(); + } + ex_data += rerank_vector_bytes(); + } +} + +} // namespace rabitqlib::ivf diff --git a/src/index/ivf_initializer.cpp b/src/index/ivf_initializer.cpp index 3fa0a49c..95356fc3 100644 --- a/src/index/ivf_initializer.cpp +++ b/src/index/ivf_initializer.cpp @@ -3,16 +3,24 @@ #include #include #include +#include +#include #include #include #include -#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" #include "rabitqlib/index/ivf/initializer.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" namespace rabitqlib::ivf { FlatRaBitQInitializer::FlatRaBitQInitializer(size_t d, size_t k, MetricType metric) @@ -47,7 +55,7 @@ void FlatRaBitQInitializer::add_vectors(const float* centroids, size_t num_threa global_centroid_[d] = static_cast(sum / static_cast(num_cluster_)); } const size_t batch_bytes = BatchDataMap::data_bytes(dim_); - parallel_for( + rabitqlib::index_detail::parallel_for( 0, div_round_up(num_cluster_, fastscan::kBatchSize), num_threads, diff --git a/src/index/ivf_search.cpp b/src/index/ivf_search.cpp index 96174f67..eaed4877 100644 --- a/src/index/ivf_search.cpp +++ b/src/index/ivf_search.cpp @@ -1,8 +1,9 @@ +#include "ivf_search.hpp" + #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/ivf/ivf.hpp" -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::ivf::detail { // Keep the scalar top-k loop out of the large inlined IVF search body. This avoids diff --git a/src/index/ivf_search.hpp b/src/index/ivf_search.hpp new file mode 100644 index 00000000..74c302f5 --- /dev/null +++ b/src/index/ivf_search.hpp @@ -0,0 +1,10 @@ +#pragma once + +#include + +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/types.hpp" + +namespace rabitqlib::ivf::detail { +void insert_candidates(buffer::SearchBuffer&, const PID*, const float*, size_t); +} // namespace rabitqlib::ivf::detail diff --git a/src/index/pipnn.cpp b/src/index/pipnn.cpp new file mode 100644 index 00000000..0b50cda8 --- /dev/null +++ b/src/index/pipnn.cpp @@ -0,0 +1,557 @@ +#include "rabitqlib/index/symqg/detail/pipnn.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "pipnn_internal.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/parallel_exception.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::symqg::detail { + +namespace pipnn_impl { +constexpr size_t kTileSize = 2048; +constexpr size_t kHashBits = 12; +constexpr size_t kLocalDegree = 2; +constexpr float kAlpha = 1.1F; + +struct Candidate { + PID id; + float distance; + uint16_t hash; +}; + +// Collection maintains hash-sorted, bounded rows under their own locks. After +// collection finishes, each row is pruned once and reused to hold only selected IDs. +class CandidateTable { + public: + CandidateTable(size_t count, size_t capacity) + : capacity_(capacity) + , table_(count * capacity) + , sizes_(count, 0) + , worst_(count, 0) + , locks_(count) {} + + void merge( + PID source, PID target, float distance, const RowMajorMatrix& sketches + ) { + uint16_t hash = 0; + for (size_t bit = 0; bit < kHashBits; ++bit) { + hash = static_cast( + (hash << 1) | (sketches(target, static_cast(bit)) > + sketches(source, static_cast(bit))) + ); + } + std::lock_guard lock(locks_[source]); + auto* row = table_.data() + source * capacity_; + auto& size = sizes_[source]; + auto& furthest = worst_[source]; + if (size == capacity_ && distance >= row[furthest].distance) { + return; + } + auto* pos = std::lower_bound( + row, + row + size, + hash, + [](const Candidate& candidate, uint16_t key) { return candidate.hash < key; } + ); + const size_t slot = static_cast(pos - row); + if (slot < size && pos->hash == hash) { + if (distance >= pos->distance) { + return; + } + *pos = {target, distance, hash}; + } else if (size < capacity_) { + std::move_backward(pos, row + size, row + size + 1); + *pos = {target, distance, hash}; + ++size; + } else { + const size_t insertion = slot > furthest ? slot - 1 : slot; + if (slot > furthest) { + std::move(row + furthest + 1, row + slot, row + furthest); + } else { + std::move_backward(row + slot, row + furthest, row + furthest + 1); + } + row[insertion] = {target, distance, hash}; + } + // Bounded rows keep this scan small; retain full-precision scores. + furthest = static_cast( + std::max_element( + row, + row + size, + [](const Candidate& a, const Candidate& b) { + return a.distance < b.distance; + } + ) - + row + ); + } + + // Pruning runs after all merges, with one worker owning each row. + void copy_candidates(size_t id, std::vector& candidates) const { + const auto* row = table_.data() + id * capacity_; + candidates.assign(row, row + sizes_[id]); + } + + void retain_ids(size_t id, const Bucket& selected) { + auto* row = table_.data() + id * capacity_; + for (size_t j = 0; j < selected.size(); ++j) { + row[j].id = selected[j]; + } + sizes_[id] = selected.size(); + } + + // Export only after every row has been pruned. + InitialGraph export_graph(size_t threads) const { + InitialGraph result; + result.offsets.resize(sizes_.size() + 1, 0); + for (size_t i = 0; i < sizes_.size(); ++i) { + result.offsets[i + 1] = result.offsets[i] + sizes_[i]; + } + result.neighbors.resize(result.offsets.back()); +#pragma omp parallel for num_threads(threads) + for (std::ptrdiff_t index = 0; index < static_cast(sizes_.size()); + ++index) { + const size_t i = static_cast(index); + const auto* row = table_.data() + i * capacity_; + for (size_t j = 0; j < sizes_[i]; ++j) { + result.neighbors[result.offsets[i] + j] = row[j].id; + } + } + return result; + } + + private: + size_t capacity_; + std::vector table_; + std::vector sizes_; + std::vector worst_; + std::vector locks_; +}; + +struct BucketJob { + Bucket ids; + size_t depth = 0; + uint32_t seed = 555; +}; + +ScratchPool::ScratchPool(size_t count, size_t dim, size_t threads) + : point_size(std::min(count, kTileSize) * dim) + , distance_size(std::max( + std::min(count, kTileSize) * std::min(count, size_t{512}), + std::min(count, kLeafSize) * std::min(count, kLeafSize) + )) + , stride(point_size + distance_size + std::min(count, kTileSize)) + // Preserve uninitialized allocation and worker first-touch placement. + , storage(new float[threads * stride]) {} + +Workspace ScratchPool::worker(size_t id) { + float* base = storage.get() + id * stride; + return {base, base + point_size, base + point_size + distance_size, {}}; +} + +void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { + for (size_t i = 0; i < ids.size(); ++i) { + std::copy_n(data + ids[i] * dim, dim, work.points + i * dim); + } +} + +void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) { + simd::pairwise_distances_lower( + work.points, work.distances, work.norms, size, dim, metric == METRIC_IP + ); +} + +std::vector partition( + const float* data, + size_t dim, + const BucketJob& job, + MetricType metric, + size_t threads, + ScratchPool& scratch +) { + const auto& ids = job.ids; + const size_t leader_count = std::min( + ids.size(), std::clamp(job.depth == 0 ? 512 : ids.size() / 200, 3, 512) + ); + const size_t fanout = std::min( + leader_count, + job.depth == 0 ? size_t{10} + : job.depth == 1 ? size_t{3} + : size_t{1} + ); + std::mt19937 random(job.seed); + Bucket leaders; + leaders.reserve(leader_count); + std::sample(ids.begin(), ids.end(), std::back_inserter(leaders), leader_count, random); + RowMajorMatrix leader_data(leader_count, dim); + for (size_t i = 0; i < leader_count; ++i) { + std::copy_n(data + leaders[i] * dim, dim, leader_data.data() + i * dim); + } + Vector leader_norms(leader_count); + simd::row_norms(leader_data.data(), leader_norms.data(), leader_count, dim); + std::vector assignments(ids.size() * fanout); + const auto assign_tile = [&](size_t begin, Workspace& work) { + const size_t count = std::min(kTileSize, ids.size() - begin); + RowMajorMatrixMap points( + work.points, static_cast(count), static_cast(dim) + ); + RowMajorMatrixMap distances( + work.distances, + static_cast(count), + static_cast(leader_count) + ); + VectorMap norms(work.norms, static_cast(count)); + for (size_t i = 0; i < count; ++i) { + std::copy_n(data + ids[begin + i] * dim, dim, points.data() + i * dim); + } + simd::matrix_product_transposed( + points.data(), leader_data.data(), distances.data(), count, dim, leader_count + ); + simd::row_norms(points.data(), norms.data(), count, dim); + work.candidates.resize(leader_count); + for (size_t i = 0; i < count; ++i) { + for (size_t j = 0; j < leader_count; ++j) { + const float dot = + distances(static_cast(i), static_cast(j)); + const float distance = + metric == METRIC_L2 + ? std::max(0.0F, norms.data()[i] + leader_norms.data()[j] - 2 * dot) + : -dot; + work.candidates[j] = {distance, static_cast(j)}; + } + std::partial_sort( + work.candidates.begin(), + work.candidates.begin() + static_cast(fanout), + work.candidates.end() + ); + for (size_t j = 0; j < fanout; ++j) { + assignments[(begin + i) * fanout + j] = work.candidates[j].second; + } + } + }; + if (job.depth == 0) { + { + index_detail::ParallelException failure; +#pragma omp parallel num_threads(threads) + { + auto work = scratch.worker(static_cast(omp_get_thread_num())); +#pragma omp for schedule(static) + for (std::ptrdiff_t offset = 0; + offset < static_cast(ids.size()); + offset += static_cast(kTileSize)) { + if (failure.failed()) { + continue; + } + try { + const size_t begin = static_cast(offset); + assign_tile(begin, work); + } catch (...) { failure.capture(); } + } + } + failure.rethrow(); + } + } else { + // Already inside the bucket-level team. Eigen must see that team rather + // than a serialized nested region when deciding its own parallelism. + auto work = scratch.worker(static_cast(omp_get_thread_num())); + for (size_t begin = 0; begin < ids.size(); begin += kTileSize) { + assign_tile(begin, work); + } + } + std::vector children(leader_count); + for (size_t i = 0; i < ids.size(); ++i) { + for (size_t j = 0; j < fanout; ++j) { + children[assignments[i * fanout + j]].push_back(ids[i]); + } + } + return children; +} + +std::vector cluster( + const float* data, + size_t count, + size_t dim, + MetricType metric, + size_t threads, + ScratchPool& scratch +) { + Bucket all(count); + std::iota(all.begin(), all.end(), PID{0}); + // The root parallelizes tiles; deeper levels parallelize independent buckets. + std::vector active; + active.push_back({std::move(all), 0, 555}); + std::vector leaves; + while (!active.empty()) { + std::vector> next(active.size()); + std::vector> finished(active.size()); + const auto process_bucket = [&](size_t i) { + auto& job = active[i]; + if (job.ids.size() <= kLeafSize) { + finished[i].push_back(std::move(job.ids)); + return; + } + auto children = partition(data, dim, job, metric, threads, scratch); + std::sort(children.begin(), children.end(), [](const auto& a, const auto& b) { + return a.size() < b.size(); + }); + Bucket small; + for (size_t j = 0; j < children.size(); ++j) { + auto& child = children[j]; + if (child.empty()) { + continue; + } + if (child.size() < kLeafSize / 10) { + small.insert(small.end(), child.begin(), child.end()); + if (small.size() < kLeafSize / 10) { + continue; + } + std::sort(small.begin(), small.end()); + small.erase(std::unique(small.begin(), small.end()), small.end()); + finished[i].push_back(std::move(small)); + small = Bucket(); + } else if (child.size() <= kLeafSize) { + finished[i].push_back(std::move(child)); + } else if (job.depth >= 4 || (job.depth >= 2 && child.size() * 5 > job.ids.size() * 4)) { + // Duplicates must not cause unbounded recursion/replication. + std::mt19937 random(job.seed + static_cast(j)); + std::shuffle(child.begin(), child.end(), random); + for (size_t begin = 0; begin < child.size(); begin += kLeafSize) { + finished[i].emplace_back( + child.begin() + static_cast(begin), + child.begin() + static_cast( + std::min(begin + kLeafSize, child.size()) + ) + ); + } + } else { + next[i].push_back( + {std::move(child), + job.depth + 1, + job.seed * 1664525U + static_cast(j) + 1013904223U} + ); + } + } + if (!small.empty()) { + std::sort(small.begin(), small.end()); + small.erase(std::unique(small.begin(), small.end()), small.end()); + finished[i].push_back(std::move(small)); + } + Bucket().swap(job.ids); + }; + if (active.front().depth == 0) { + // A serialized outer OpenMP region would make the root's tile team + // nested, forcing the runtime to create workers again on every reset. + process_bucket(0); + } else { + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(threads) schedule(dynamic) + for (std::ptrdiff_t i = 0; i < static_cast(active.size()); + ++i) { + if (failure.failed()) { + continue; + } + try { + process_bucket(static_cast(i)); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + } + std::vector new_active; + for (size_t i = 0; i < active.size(); ++i) { + for (auto& leaf : finished[i]) { + leaves.push_back(std::move(leaf)); + } + for (auto& job : next[i]) { + new_active.push_back(std::move(job)); + } + } + active = std::move(new_active); + } + return leaves; +} +} // namespace pipnn_impl + +// Native PiPNN-style seeding: overlapping leader partitions, local dense kNN, +// directional hash reservoirs, and alpha pruning. Uses existing Eigen/OpenMP. +InitialGraph build_initial_graph( + const float* data, + size_t count, + size_t dim, + size_t degree, + MetricType metric, + size_t num_threads +) { + validate_metric_type(metric); + if (data == nullptr || dim == 0 || degree == 0 || degree >= count || + count >= buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { + throw std::invalid_argument( + "PiPNN seed requires data, positive dim, and a degree multiple of 32 below " + "count" + ); + } + using namespace pipnn_impl; + const size_t threads = resolve_num_threads(num_threads); + ScratchPool scratch(count, dim, threads); + auto leaves = cluster(data, count, dim, metric, threads, scratch); + RowMajorMatrix projections(dim, kHashBits); + std::mt19937 random(555); + std::normal_distribution normal; + for (Eigen::Index i = 0; i < projections.size(); ++i) { + projections.data()[i] = normal(random); + } + RowMajorMatrix sketches(count, kHashBits); + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(threads) schedule(static) + for (std::ptrdiff_t offset = 0; offset < static_cast(count); + offset += static_cast(kTileSize)) { + if (failure.failed()) { + continue; + } + try { + const size_t begin = static_cast(offset); + const size_t rows = std::min(kTileSize, count - begin); + simd::matrix_product( + data + begin * dim, + projections.data(), + sketches.data() + begin * kHashBits, + rows, + dim, + kHashBits + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + CandidateTable table(count, degree * 5 / 2); + { + index_detail::ParallelException failure; +#pragma omp parallel num_threads(threads) + { + auto work = scratch.worker(static_cast(omp_get_thread_num())); +#pragma omp for schedule(dynamic) + for (std::ptrdiff_t index = 0; + index < static_cast(leaves.size()); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t leaf = static_cast(index); + const auto& ids = leaves[leaf]; + if (ids.size() < 2) { + continue; + } + gather(data, dim, ids, work); + pairwise(work, ids.size(), dim, metric); + for (size_t i = 0; i < ids.size(); ++i) { + work.candidates.clear(); + for (size_t j = 0; j < ids.size(); ++j) { + if (i != j) { + work.candidates.emplace_back( + work.distances + [std::max(i, j) * ids.size() + std::min(i, j)], + ids[j] + ); + } + } + const size_t keep = std::min(kLocalDegree, work.candidates.size()); + std::partial_sort( + work.candidates.begin(), + work.candidates.begin() + static_cast(keep), + work.candidates.end() + ); + for (size_t j = 0; j < keep; ++j) { + const auto [distance, target] = work.candidates[j]; + table.merge(ids[i], target, distance, sketches); + table.merge(target, ids[i], distance, sketches); + } + } + } catch (...) { failure.capture(); } + } + } + failure.rethrow(); + } + std::vector().swap(leaves); + scratch.storage.reset(); + sketches.resize(0, 0); + // Prune in-place: retain selected IDs in the front of each reservoir row. + { + index_detail::ParallelException failure; +#pragma omp parallel num_threads(threads) + { + std::vector candidates; + Bucket selected; +#pragma omp for schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(count); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + table.copy_candidates(i, candidates); + std::sort( + candidates.begin(), + candidates.end(), + [](const Candidate& a, const Candidate& b) { + return a.distance < b.distance || + (a.distance == b.distance && a.id < b.id); + } + ); + selected.clear(); + for (const auto& candidate : candidates) { + bool occluded = false; + for (PID accepted : selected) { + const float distance = metric == METRIC_L2 + ? euclidean_sqr( + data + accepted * dim, + data + candidate.id * dim, + dim + ) + : dot_product_dis( + data + accepted * dim, + data + candidate.id * dim, + dim + ) - 1.0F; + if (distance <= candidate.distance / kAlpha) { + occluded = true; + break; + } + } + if (!occluded) { + selected.push_back(candidate.id); + } + if (selected.size() == degree) { + break; + } + } + table.retain_ids(i, selected); + } catch (...) { failure.capture(); } + } + } + failure.rethrow(); + } + return table.export_graph(threads); +} +} // namespace rabitqlib::symqg::detail diff --git a/src/index/pipnn_internal.hpp b/src/index/pipnn_internal.hpp new file mode 100644 index 00000000..3634b6d0 --- /dev/null +++ b/src/index/pipnn_internal.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/types.hpp" + +// Private declarations shared only by the implementation and reference tests. +namespace rabitqlib::symqg::detail::pipnn_impl { +constexpr size_t kLeafSize = 1024; +using Bucket = std::vector; + +// Non-owning slices of a build-scoped allocation. Large matrices are released +// after seeding rather than retained in worker heaps across graph resets. +struct Workspace { + float* points; + float* distances; + float* norms; + std::vector> candidates; +}; + +struct ScratchPool { + size_t point_size; + size_t distance_size; + size_t stride; + std::unique_ptr storage; + + ScratchPool(size_t count, size_t dim, size_t threads); + Workspace worker(size_t id); +}; + +void gather(const float*, size_t, const Bucket&, Workspace&); +void pairwise(Workspace&, size_t, size_t, MetricType); +std::vector cluster(const float*, size_t, size_t, MetricType, size_t, ScratchPool&); + +} // namespace rabitqlib::symqg::detail::pipnn_impl diff --git a/src/index/qg.cpp b/src/index/qg.cpp index b88628e4..b1dcaf2a 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -1,14 +1,10 @@ #include "rabitqlib/index/symqg/qg.hpp" -#include - #include -#include #include #include #include #include -#include #include #include #include @@ -19,22 +15,27 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/pack_excode.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/pack_excode.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/parallel_exception.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/symqg/detail/search_scratch.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/bitops.hpp" -#include "rabitqlib/utils/buffer.hpp" #include "rabitqlib/utils/io.hpp" #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/path.hpp" -#include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" -#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { @@ -76,6 +77,12 @@ size_t QuantizedGraph::padded_dimension(size_t dim) { size_t QuantizedGraph::num_vertices() const { return this->num_points_; } +size_t QuantizedGraph::max_elements() const { return num_points_; } + +size_t QuantizedGraph::num_points() const { return ready_ ? num_points_ : 0; } + +size_t QuantizedGraph::num_stored_points() const { return ready_ ? num_points_ : 0; } + size_t QuantizedGraph::dimension() const { return this->dim_; } size_t QuantizedGraph::degree_bound() const { return this->degree_bound_; } @@ -167,38 +174,55 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) if (centroid_.size() != padded_dim_) { throw std::logic_error("qg-quant centroid must be set before copying vectors"); } -#pragma omp parallel num_threads(thread_count) { - std::vector rotated_data(padded_dim_); - std::vector quantized_data(padded_dim_); + index_detail::ParallelException failure; +#pragma omp parallel num_threads(thread_count) + { + std::vector rotated_data; + std::vector quantized_data; + try { + rotated_data.resize(padded_dim_); + quantized_data.resize(padded_dim_); + } catch (...) { failure.capture(); } #pragma omp for schedule(static) - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - const size_t i = static_cast(index); - rotator_->rotate(data + (dim_ * i), rotated_data.data()); - ExDataMap output( - get_quantized_vector(i), padded_dim_, quantization_bits_ - ); - float f_add; - float f_rescale; - float unused_f_error = 0; - quant::quantize_full_single( - rotated_data.data(), - centroid_.data(), - padded_dim_, - quantization_bits_, - quantized_data.data(), - f_add, - f_rescale, - unused_f_error, - metric_type_ - ); - output.f_add_ex() = f_add; - output.f_rescale_ex() = f_rescale; - quant::rabitq_impl::ex_bits::packing_rabitqplus_code( - quantized_data.data(), output.ex_code(), padded_dim_, quantization_bits_ - ); + for (std::ptrdiff_t index = 0; + index < static_cast(num_points_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + rotator_->rotate(data + (dim_ * i), rotated_data.data()); + ExDataMap output( + get_quantized_vector(i), padded_dim_, quantization_bits_ + ); + float f_add; + float f_rescale; + float unused_f_error = 0; + quant::quantize_full_single( + rotated_data.data(), + centroid_.data(), + padded_dim_, + quantization_bits_, + quantized_data.data(), + f_add, + f_rescale, + unused_f_error, + metric_type_ + ); + output.f_add_ex() = f_add; + output.f_rescale_ex() = f_rescale; + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + quantized_data.data(), + output.ex_code(), + padded_dim_, + quantization_bits_ + ); + } catch (...) { failure.capture(); } + } } + failure.rethrow(); } return; } @@ -233,7 +257,8 @@ void QuantizedGraph::save(const char* filename) const { } output.exceptions(std::ios::badbit | std::ios::failbit); - constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; // "QGRABITQ" + // Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". + constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; constexpr uint32_t kFormatVersion = 1; if (quantization_bits_ != 0) { output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); @@ -304,7 +329,8 @@ void QuantizedGraph::load(const char* filename) { read_exact(&value, sizeof(value), field); }; - constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; // "QGRABITQ" + // Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". + constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; constexpr uint32_t kFormatVersion = 1; uint64_t magic = 0; read_value(magic, "format marker"); @@ -370,7 +396,7 @@ void QuantizedGraph::load(const char* filename) { throw std::runtime_error("Invalid QuantizedGraph payload size"); } - loaded.initialize(); + loaded.initialize(false); if (loaded.quantization_bits_ != 0) { loaded.centroid_.resize(loaded.padded_dim_); @@ -421,12 +447,7 @@ void QuantizedGraph::search_impl( const float* query, uint32_t k, uint32_t* results, float* dists, size_t ef ) { validate_search(query, k, results, dists, ef); - std::vector rotated_query(padded_dim_); - std::vector est_dist(degree_bound_); - std::vector lut_float(padded_dim_ * 4); - BatchQuery batch_query; - buffer::SearchBuffer search_pool(ef); - buffer::SearchBuffer result_pool(k); + detail::SearchScratch scratch(padded_dim_, degree_bound_, ef, k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; if (visited_size != num_points_) { @@ -434,19 +455,7 @@ void QuantizedGraph::search_impl( visited_size = num_points_; } search_with_scratch_impl( - query, - k, - results, - dists, - rotated_query.data(), - est_dist.data(), - lut_float.data(), - batch_query, - search_pool, - result_pool, - visited, - kPidMax, - nullptr + query, k, results, dists, scratch.view(visited), kPidMax, nullptr ); } @@ -483,52 +492,26 @@ void QuantizedGraph::search_batch_with_ef( return; } - struct SearchScratch { - std::vector rotated_query; - std::vector estimated_distances; - std::vector lookup_table; - BatchQuery batch_query; - buffer::SearchBuffer search_pool; - buffer::SearchBuffer result_pool; - VisitedSet visited; - - SearchScratch(size_t padded_dim, size_t degree, size_t ef, size_t k, size_t points) - : rotated_query(padded_dim) - , estimated_distances(degree) - , lookup_table(padded_dim * 4) - , search_pool(ef) - , result_pool(k) - , visited(points, points / 10) {} - }; - const auto workers = static_cast(std::min( {resolve_num_threads(num_threads), num_queries, static_cast(std::numeric_limits::max())} )); - std::atomic failed{false}; - std::exception_ptr error; - const auto capture_error = [&] { -#pragma omp critical(rabitq_symqg_search_error) - { - if (!error) { - error = std::current_exception(); - } - } - failed.store(true, std::memory_order_relaxed); - }; + index_detail::ParallelException failure; #pragma omp parallel num_threads(workers) if (workers > 1) { // Allocate and release scratch on its worker, retaining it across queries. - std::optional scratch; + std::optional scratch; + VisitedSet visited; try { - scratch.emplace(padded_dim_, degree_bound_, ef, knn, num_points_); - } catch (...) { capture_error(); } + scratch.emplace(padded_dim_, degree_bound_, ef, knn); + visited.initialize(num_points_, num_points_ / 10); + } catch (...) { failure.capture(); } #pragma omp for schedule(dynamic) for (std::ptrdiff_t query_index = 0; query_index < static_cast(num_queries); ++query_index) { - if (failed.load(std::memory_order_relaxed)) { + if (failure.failed()) { continue; } const size_t i = static_cast(query_index); @@ -539,22 +522,14 @@ void QuantizedGraph::search_batch_with_ef( knn, results + i * knn, dists + i * knn, - local.rotated_query.data(), - local.estimated_distances.data(), - local.lookup_table.data(), - local.batch_query, - local.search_pool, - local.result_pool, - local.visited, + local.view(visited), kPidMax, nullptr ); - } catch (...) { capture_error(); } + } catch (...) { failure.capture(); } } } - if (error) { - std::rethrow_exception(error); - } + failure.rethrow(); } void QuantizedGraph::validate_search( @@ -585,7 +560,7 @@ float QuantizedGraph::point_distance( // Scan a data row and store estimated neighbor distances. The caller scores the current // vertex from either its raw vector (vanilla QG) or its 4/8-bit code (qg-quant). -inline void QuantizedGraph::scan_neighbors( +void QuantizedGraph::scan_neighbors( const BatchQuery& q_obj, PID data_id, float* est_dist, @@ -655,20 +630,14 @@ void QuantizedGraph::update_results( } } -// initialize const offsets & data array void QuantizedGraph::initialize_layout() { if (quantization_bits_ == 0) { batch_data_offset_ = checked_multiply(dim_, sizeof(float)); } else { - const size_t code_bits = checked_multiply(padded_dim_, quantization_bits_); - batch_data_offset_ = checked_add(code_bits / 8, checked_multiply(sizeof(float), 2)); + batch_data_offset_ = ExDataMap::data_bytes(padded_dim_, quantization_bits_); } - const size_t binary_batch_bytes = - checked_multiply(padded_dim_, fastscan::kBatchSize) / 8; - const size_t factor_bytes = - checked_multiply(checked_multiply(sizeof(float), fastscan::kBatchSize), size_t{2}); - const size_t batch_bytes = checked_add(binary_batch_bytes, factor_bytes); + const size_t batch_bytes = QGBatchDataMap::data_bytes(padded_dim_); neighbor_offset_ = checked_add( batch_data_offset_, checked_multiply(batch_bytes, degree_bound_ / fastscan::kBatchSize) @@ -677,8 +646,12 @@ void QuantizedGraph::initialize_layout() { checked_add(neighbor_offset_, checked_multiply(degree_bound_, sizeof(PID))); } -void QuantizedGraph::initialize() { - rotator_.reset(choose_rotator(dim_, rotator_type_, padded_dim_, seed_)); +void QuantizedGraph::initialize(bool randomize_rotation) { + if (randomize_rotation) { + rotator_.reset(choose_rotator(dim_, rotator_type_, padded_dim_, seed_)); + } else { + rotator_ = rotator_impl::make_rotator_for_loading(dim_, rotator_type_, padded_dim_); + } assert(padded_dim_ % 32 == 0); assert(padded_dim_ >= dim_); @@ -718,19 +691,9 @@ void QuantizedGraph::reconstruct_quantized_vector(PID data_id, float* rec get_quantized_vector(data_id), padded_dim_, quantization_bits_ ); std::vector quantized_data(padded_dim_); - if (quantization_bits_ == 8) { - std::copy(data.ex_code(), data.ex_code() + padded_dim_, quantized_data.begin()); - } else { - for (size_t i = 0; i < padded_dim_; i += 16) { - uint64_t packed = 0; - std::memcpy(&packed, data.ex_code() + (i / 2), sizeof(packed)); - for (size_t j = 0; j < 8; ++j) { - const uint8_t pair = static_cast(packed >> (j * 8)); - quantized_data[i + j] = pair & 0x0f; - quantized_data[i + 8 + j] = pair >> 4; - } - } - } + quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( + data.ex_code(), quantized_data.data(), padded_dim_, quantization_bits_ + ); quant::reconstruct_full_vec( quantized_data.data(), centroid_.data(), @@ -742,121 +705,21 @@ void QuantizedGraph::reconstruct_quantized_vector(PID data_id, float* rec ); } -// Construction sources come from owned raw rows or from the existing RaBitQ codes. -// Reconstructed sources are already rotated; never pass them through prepare_query. -const float* QuantizedGraph::prepare_build_query( - PID id, - std::vector& rotated, - std::optional& prepared, - const BuildCache* cache -) const { - if (is_quantized()) { - if (cache != nullptr) { - const float* row = cache->rotated_vectors.data() + id * padded_dim_; - prepared.emplace(row, cache->queries[id]); - return row; - } - rotated.resize(padded_dim_); - reconstruct_quantized_vector(id, rotated.data()); - prepared.emplace(rotated.data(), centroid_.data(), padded_dim_, metric_type_); - return rotated.data(); - } - prepared.reset(); - return get_vector(id); -} - -// find candidate neighbors for cur_id, exclude the vertex itself -void QuantizedGraph::find_candidates( +// Encode a logical neighbor-ID row, including physical storage for tail batches. +void QuantizedGraph::encode_neighbors( PID cur_id, - size_t search_ef, - std::vector>& results, - VisitedSet& vis, - const std::vector& degrees, - const BuildCache* cache -) const { - std::vector rotated_query(cache == nullptr ? padded_dim_ : 0); - std::optional quantized_query; - const float* query = prepare_build_query(cur_id, rotated_query, quantized_query, cache); - if (cache == nullptr && !is_quantized()) { - rotator_->rotate(query, rotated_query.data()); - } - const float* rotated = cache == nullptr - ? rotated_query.data() - : cache->rotated_vectors.data() + cur_id * padded_dim_; - BatchQuery q_obj(rotated, padded_dim_, metric_type_); - - // insert entry point to initialize search buffer - buffer::SearchBuffer tmp_pool(search_ef); - tmp_pool.insert(this->entry_point_, std::numeric_limits::max()); - - /* Current version of fast scan compute 32 distances */ - std::vector est_dist(degree_bound_); // estimated distances - while (tmp_pool.has_next()) { - auto cur_candi = tmp_pool.pop(); - if (vis.get(cur_candi)) { - continue; - } - vis.set(cur_candi); - auto cur_degree = degrees[cur_candi]; - const float vertex_distance = - point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_candi); - q_obj.set_g_add(vertex_distance); - scan_neighbors(q_obj, cur_candi, est_dist.data(), tmp_pool, vis, cur_degree); - if (cur_candi != cur_id) { - results.emplace_back(cur_candi, vertex_distance); - } - } -} - -// based on new neighbor lists to update quantization code and factors -void QuantizedGraph::update_qg( - PID cur_id, - const std::vector>& new_neighbors, - const BuildCache* cache, - std::vector* scratch + const std::vector& new_neighbors, + const float* rotated_data, + const float* rotated_centroid ) { size_t cur_degree = new_neighbors.size(); if (cur_degree == 0) { return; } - // copy neighbors auto neighbor_ptr = get_neighbors(cur_id); for (size_t i = 0; i < cur_degree; ++i) { - neighbor_ptr[i] = new_neighbors[i].id; - } - - // rotated data - std::vector local_scratch; - std::vector& rotated_data = scratch == nullptr ? local_scratch : *scratch; - rotated_data.resize((cur_degree + 1) * padded_dim_); - float* rotated_centroid = rotated_data.data() + cur_degree * padded_dim_; - for (size_t i = 0; i < cur_degree; ++i) { - if (cache != nullptr) { - std::copy_n( - cache->rotated_vectors.data() + new_neighbors[i].id * padded_dim_, - padded_dim_, - rotated_data.data() + i * padded_dim_ - ); - } else if (quantization_bits_ == 0) { - const float* neighbor_vec = get_vector(new_neighbors[i].id); - this->rotator_->rotate(neighbor_vec, &rotated_data[i * padded_dim_]); - } else { - reconstruct_quantized_vector( - new_neighbors[i].id, &rotated_data[i * padded_dim_] - ); - } - } - if (cache != nullptr) { - std::copy_n( - cache->rotated_vectors.data() + cur_id * padded_dim_, - padded_dim_, - rotated_centroid - ); - } else if (quantization_bits_ == 0) { - this->rotator_->rotate(get_vector(cur_id), rotated_centroid); - } else { - reconstruct_quantized_vector(cur_id, rotated_centroid); + neighbor_ptr[i] = new_neighbors[i]; } // quantize batches for current vertex @@ -864,7 +727,7 @@ void QuantizedGraph::update_qg( for (size_t i = 0; i < cur_degree; i += fastscan::kBatchSize) { const size_t batch_size = std::min(cur_degree - i, fastscan::kBatchSize); quant::quantize_qg_batch( - rotated_data.data() + (i * padded_dim_), + rotated_data + (i * padded_dim_), rotated_centroid, batch_size, padded_dim_, @@ -903,90 +766,103 @@ void QuantizedGraph::search_with_scratch( PID hint, const float* k1xsumq ) { - validate_search(query, knn, results, dists, ef_); - search_with_scratch_impl( + detail::search_with_workspace( + *this, query, knn, results, dists, - rotated_query_scratch, - est_dist_scratch, - lut_float_scratch, - batch_query, - search_pool, - result_pool, - visited, + {rotated_query_scratch, + est_dist_scratch, + lut_float_scratch, + batch_query, + search_pool, + result_pool, + visited}, hint, k1xsumq ); } +void detail::search_with_workspace( + QuantizedGraph& graph, + const float* query, + uint32_t knn, + uint32_t* results, + float* dists, + detail::SearchScratchView scratch, + PID hint, + const float* k1xsumq +) { + graph.validate_search(query, knn, results, dists, graph.ef_); + graph.search_with_scratch_impl(query, knn, results, dists, scratch, hint, k1xsumq); +} + void QuantizedGraph::search_with_scratch_impl( const float* __restrict__ query, uint32_t knn, uint32_t* __restrict__ results, float* __restrict__ dists, - float* __restrict__ rotated_query_scratch, - float* __restrict__ est_dist_scratch, - float* __restrict__ lut_float_scratch, - BatchQuery& batch_query, - buffer::SearchBuffer& search_pool, - buffer::SearchBuffer& result_pool, - VisitedSet& visited, + detail::SearchScratchView scratch, PID hint, const float* k1xsumq ) { - if (rotated_query_scratch == nullptr || est_dist_scratch == nullptr || - lut_float_scratch == nullptr) { + if (scratch.rotated_query == nullptr || scratch.estimated_distances == nullptr || + scratch.lookup_table == nullptr) { throw std::invalid_argument("QuantizedGraph search buffers must not be null"); } if (hint != kPidMax && hint >= num_points_) { throw std::invalid_argument("SymphonyQG search hint is outside the graph"); } - if (!visited.initialized()) { + if (!scratch.visited.initialized()) { throw std::invalid_argument("SymphonyQG visited set must be initialized"); } - rotator_->rotate(query, rotated_query_scratch); + rotator_->rotate(query, scratch.rotated_query); std::optional quantized_query; if (quantization_bits_ != 0) { quantized_query.emplace( - rotated_query_scratch, centroid_.data(), padded_dim_, metric_type_ + scratch.rotated_query, centroid_.data(), padded_dim_, metric_type_ ); } const auto* quantized = quantized_query ? &*quantized_query : nullptr; - batch_query.reset( - rotated_query_scratch, padded_dim_, lut_float_scratch, metric_type_, k1xsumq + scratch.batch_query.reset( + scratch.rotated_query, padded_dim_, scratch.lookup_table, metric_type_, k1xsumq ); - search_pool.clear(); - result_pool.clear(); - visited.clear(); + scratch.search_pool.clear(); + scratch.result_pool.clear(); + scratch.visited.clear(); if (hint != kPidMax && hint != entry_point_) { - search_pool.insert(hint, point_distance(query, quantized, hint)); + scratch.search_pool.insert(hint, point_distance(query, quantized, hint)); } - search_pool.insert(entry_point_, std::numeric_limits::max()); + scratch.search_pool.insert(entry_point_, std::numeric_limits::max()); - while (search_pool.has_next()) { - const PID current = search_pool.pop(); - if (visited.get(current)) { + while (scratch.search_pool.has_next()) { + const PID current = scratch.search_pool.pop(); + if (scratch.visited.get(current)) { continue; } - visited.set(current); + scratch.visited.set(current); const float vertex_distance = point_distance(query, quantized, current); - batch_query.set_g_add(vertex_distance); + scratch.batch_query.set_g_add(vertex_distance); scan_neighbors( - batch_query, current, est_dist_scratch, search_pool, visited, degree_bound_ + scratch.batch_query, + current, + scratch.estimated_distances, + scratch.search_pool, + scratch.visited, + degree_bound_ ); - result_pool.insert(current, vertex_distance); + scratch.result_pool.insert(current, vertex_distance); } - update_results(result_pool, visited, query, quantized); - if (result_pool.size() != knn) { + update_results(scratch.result_pool, scratch.visited, query, quantized); + if (scratch.result_pool.size() != knn) { throw std::runtime_error("QuantizedGraph search could not produce k results"); } - result_pool.copy_results(results, dists); + scratch.result_pool.copy_results(results, dists); } } // namespace rabitqlib::symqg diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp new file mode 100644 index 00000000..26f974b3 --- /dev/null +++ b/src/index/qg_builder.cpp @@ -0,0 +1,828 @@ +#include "rabitqlib/index/symqg/qg_builder.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/exact_nn.hpp" +#include "rabitqlib/index/detail/parallel_exception.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/symqg/detail/pipnn.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::symqg { + +QGBuilder::BuildQuery::BuildQuery( + const QuantizedGraph& graph, const BuildCache* cache +) + : graph_(graph), cache_(cache) {} + +QGBuilder::BuildWorker::BuildWorker(size_t num_nodes, size_t visited_capacity) + : visited(num_nodes, visited_capacity) {} +const QGBuilder::BuildCache* QGBuilder::build_cache() const { + return cache_ ? &*cache_ : nullptr; +} + +QGBuilder::QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + size_t num_threads, + uint32_t seed, + bool cache_vectors +) + : qg_{index} + , ef_build_{ef_build} + , num_threads_{resolve_num_threads(num_threads)} + , num_nodes_{qg_.num_vertices()} + , dim_{qg_.dimension()} + , degree_bound_(qg_.degree_bound()) + , seed_(seed) { + if (cache_vectors) { + cache_.emplace(); + } +} + +QGBuilder::QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + const float* data, + size_t num_threads, + QGInitialization init, + uint32_t seed, + bool cache_vectors +) + : QGBuilder(index, ef_build, num_threads, seed, cache_vectors) { + if (init != QGInitialization::PiPNN && init != QGInitialization::Random) { + throw std::invalid_argument("Unknown QG initialization"); + } + seeded_ = init == QGInitialization::PiPNN; + reset(data); +} + +QGBuilder::QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + const float* data, + size_t num_threads, + uint32_t seed, + bool cache_vectors +) + : QGBuilder( + index, ef_build, data, num_threads, QGInitialization::Random, seed, cache_vectors + ) {} + +void QGBuilder::reset(const float* data) { + if (data == nullptr) { + throw std::invalid_argument("QGBuilder data must not be null"); + } + if (seeded_) { + auto initial = detail::build_initial_graph( + data, num_nodes_, dim_, degree_bound_, qg_.metric_type_, num_threads_ + ); + initialize_seed(data, initial.offsets, initial.neighbors); + } else { + initialize_storage(data); + random_init(); + } +} + +void QGBuilder::initialize_seed( + const float* data, const std::vector& offsets, const std::vector& neighbors +) { + if (offsets.size() != num_nodes_ + 1 || offsets.front() != 0 || + offsets.back() != neighbors.size()) { + throw std::invalid_argument("Seed graph offsets must delimit every vertex"); + } + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); + if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || + offsets[i + 1] - offsets[i] > degree_bound_) { + throw std::invalid_argument( + "Seed graph row exceeds degree bound or has invalid offsets" + ); + } + for (size_t j = offsets[i]; j < offsets[i + 1]; ++j) { + if (neighbors[j] >= num_nodes_ || neighbors[j] == i) { + throw std::invalid_argument( + "Seed graph IDs must be in range and exclude self" + ); + } + } + } + initialize_storage(data); + { + index_detail::ParallelException failure; +#pragma omp parallel num_threads(num_threads_) + { + auto& worker = workers_[static_cast(omp_get_thread_num())]; + auto& ids = worker.neighbor_ids; + try { + ids.reserve(degree_bound_); + } catch (...) { failure.capture(); } +#pragma omp for schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + ids.assign( + neighbors.begin() + static_cast(offsets[i]), + neighbors.begin() + static_cast(offsets[i + 1]) + ); + std::sort(ids.begin(), ids.end()); + ids.erase(std::unique(ids.begin(), ids.end()), ids.end()); + // Seed rows need only IDs; scores are materialized during search. + write_row(static_cast(i), worker); + } catch (...) { failure.capture(); } + } + } + failure.rethrow(); + } +} + +void QGBuilder::refine() { run_passes(1); } + +void QGBuilder::build() { run_passes(seeded_ ? 1 : 3); } + +void QGBuilder::build(size_t num_iter) { + if (num_iter < 2) { + throw std::invalid_argument("The number of QG build iterations must be at least 2"); + } + run_passes(num_iter); +} + +bool QGBuilder::check_dup() const { + std::atomic flag(false); + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + std::unordered_set edges; + for (auto nei : new_neighbors_[i]) { + if (edges.find(nei.id) != edges.end()) { + flag = true; + } + edges.emplace(nei.id); + } + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + return flag; +} + +float QGBuilder::avg_degree() const { + size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), 0U); + return static_cast(degrees) / static_cast(num_nodes_); +} +void QGBuilder::run_passes(size_t num_iter) { + qg_.ready_ = false; + // Entry points guarantee a positive count; only the last pass completes degrees. + for (size_t i = 1; i < num_iter; ++i) { + run_pass(/*final_pass=*/false); + } + run_pass(/*final_pass=*/true); + qg_.ready_ = true; +} + +void QGBuilder::write_row(PID id, BuildWorker& worker) { + if (!worker.neighbor_ids.empty()) { + const auto& new_neighbors = worker.neighbor_ids; + const size_t cur_degree = new_neighbors.size(); + const auto* cache = build_cache(); + const size_t padded_dim = qg_.padded_dim_; + auto& rotated_data = worker.rotated_neighbors; + rotated_data.resize((cur_degree + 1) * padded_dim); + float* rotated_centroid = rotated_data.data() + cur_degree * padded_dim; + for (size_t i = 0; i < cur_degree; ++i) { + if (cache != nullptr) { + std::copy_n( + cache->rotated_vectors.data() + new_neighbors[i] * padded_dim, + padded_dim, + rotated_data.data() + i * padded_dim + ); + } else if (qg_.quantization_bits_ == 0) { + const float* neighbor_vec = qg_.get_vector(new_neighbors[i]); + qg_.rotator_->rotate(neighbor_vec, &rotated_data[i * padded_dim]); + } else { + qg_.reconstruct_quantized_vector( + new_neighbors[i], &rotated_data[i * padded_dim] + ); + } + } + if (cache != nullptr) { + std::copy_n( + cache->rotated_vectors.data() + id * padded_dim, + padded_dim, + rotated_centroid + ); + } else if (qg_.quantization_bits_ == 0) { + qg_.rotator_->rotate(qg_.get_vector(id), rotated_centroid); + } else { + qg_.reconstruct_quantized_vector(id, rotated_centroid); + } + qg_.encode_neighbors(id, new_neighbors, rotated_data.data(), rotated_centroid); + } + degrees_[id] = worker.neighbor_ids.size(); +} + +void QGBuilder::write_row(PID id, const CandidateList& neighbors, BuildWorker& worker) { + worker.neighbor_ids.resize(neighbors.size()); + for (size_t i = 0; i < neighbors.size(); ++i) { + worker.neighbor_ids[i] = neighbors[i].id; + } + write_row(id, worker); +} + +void QGBuilder::initialize_storage(const float* data) { + if (data == nullptr) { + throw std::invalid_argument("QGBuilder data must not be null"); + } + // Allocate refinement scratch after PiPNN releases its dense workspace. + // Clear old graph-dependent state on reset while retaining vector capacities. + new_neighbors_.resize(num_nodes_); + pruned_neighbors_.resize(num_nodes_); + reverse_buffer_.resize(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); + new_neighbors_[i].clear(); + pruned_neighbors_[i].clear(); + reverse_buffer_[i].clear(); + } + if (workers_.empty()) { + workers_ = std::vector( + num_threads_, + BuildWorker(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + ); + } + degrees_.assign(num_nodes_, degree_bound_); + std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); + + qg_.ready_ = false; + qg_.set_quantization_centroid(centroid.data()); + qg_.copy_vectors(data, num_threads_); + + if (cache_) { + const size_t padded_dim = qg_.padded_dim_; + cache_->rotated_vectors.resize(num_nodes_ * padded_dim); + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(static) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t id = static_cast(index); + float* row = cache_->rotated_vectors.data() + id * padded_dim; + if (qg_.is_quantized()) { + qg_.reconstruct_quantized_vector(id, row); + } else { + qg_.rotator_->rotate(qg_.get_vector(id), row); + } + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + cache_->queries.clear(); + if (qg_.is_quantized()) { + cache_->queries.reserve(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t id = static_cast(index); + cache_->queries.emplace_back( + cache_->rotated_vectors.data() + id * padded_dim, + qg_.centroid_.data(), + padded_dim, + qg_.metric_type_ + ); + } + } + } + + PID entry_point = 0; + if (qg_.is_quantized()) { + QuantizedQuery query( + qg_.centroid_.data(), qg_.centroid_.data(), qg_.padded_dim_, qg_.metric_type_ + ); + float best = std::numeric_limits::max(); + for (PID id = 0; id < num_nodes_; ++id) { + const float distance = qg_.quantized_distance(query, id); + if (distance < best) { + best = distance; + entry_point = id; + } + } + } else { + entry_point = exact_nn( + data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ + ); + } + + qg_.set_ep(entry_point); +} + +void QGBuilder::add_pruned_edges( + const CandidateList& result, + const CandidateList& pruned_list, + CandidateList& new_result, + float threshold +) { + size_t start = 0; + new_result.clear(); + new_result = result; + + std::unordered_set nei_set; + nei_set.reserve(degree_bound_); + for (const auto& nei : result) { + nei_set.emplace(nei.id); + } + + BuildQuery query(qg_, build_cache()); + while (new_result.size() < degree_bound_ && start < pruned_list.size()) { + const auto& cur = pruned_list[start]; + bool occlude = false; + query.reset(cur.id); + float dik_sqr = cur.distance; + + if (nei_set.find(cur.id) != nei_set.end()) { + break; + } + + for (auto& nei : new_result) { + float dij_sqr = nei.distance; + if (dij_sqr > dik_sqr) { + break; + } + float djk_sqr = query.distance_to(nei.id); + float cosine = + (dik_sqr + dij_sqr - djk_sqr) / (2 * std::sqrt(dij_sqr * dik_sqr)); + if (cosine > threshold) { + occlude = true; + break; + } + } + + if (!occlude) { + new_result.emplace_back(cur); + nei_set.emplace(cur.id); + std::sort(new_result.begin(), new_result.end()); + } + + ++start; + } +} + +void QGBuilder::heuristic_prune( + PID cur_id, CandidateList& pool, CandidateList& pruned_results, bool record_pruned +) { + if (pool.empty()) { + return; + } + pruned_results.clear(); + size_t poolsize = pool.size(); + if (poolsize <= degree_bound_) { + pruned_results = pool; + return; + } + + std::vector pruned(poolsize, false); + size_t start = 0; + + BuildQuery query(qg_, build_cache()); + while (pruned_results.size() < degree_bound_ && start < poolsize) { + auto candidate_id = pool[start].id; + if (pruned[start]) { + ++start; + continue; + } + + pruned_results.emplace_back(pool[start]); + query.reset(candidate_id); + + // i : current vertex + // j : neighbor added in this iter + // k : remained unpruned candidate neighbor + for (size_t k = start + 1; k < poolsize; ++k) { + if (pruned[k]) { + continue; + } + float dik = pool[k].distance; + auto djk = query.distance_to(pool[k].id); + + if (djk < dik) { + if (record_pruned && pruned_neighbors_[cur_id].size() < kMaxPrunedSize) { + pruned_neighbors_[cur_id].emplace_back(pool[k]); + } + pruned[k] = true; + } + } + + ++start; + } +} + +/** + * @brief search for new neighbor in qg + * + * @param record_pruned Record pruned candidates for final degree completion. + */ +void QGBuilder::search_new_neighbors(bool record_pruned) { + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + PID cur_id = i; + auto& worker = workers_[static_cast(omp_get_thread_num())]; + CandidateList& candidates = worker.candidates; + candidates.clear(); + VisitedSet& vis = worker.visited; + candidates.reserve(2 * kMaxCandidatePoolSize); + vis.clear(); + find_candidates(cur_id, candidates, vis); + + // Seeded construction keeps its initial edges only in QG. Materialize + // their scores on demand, after the caller can release the input/CSR. + if (new_neighbors_[cur_id].empty() && degrees_[cur_id] != 0) { + BuildQuery query(qg_, build_cache()); + query.reset(cur_id); + const auto ids = qg_.get_neighbors(cur_id); + for (size_t j = 0; j < degrees_[cur_id]; ++j) { + if (ids[j] != cur_id && !vis.get(ids[j])) { + candidates.emplace_back(ids[j], query.distance_to(ids[j])); + } + } + } + + // Add current neighbors retained by preceding iterations. + for (auto& nei : new_neighbors_[cur_id]) { + auto neighbor_id = nei.id; + if (neighbor_id != cur_id && !vis.get(neighbor_id)) { + candidates.emplace_back(nei); + } + } + + size_t min_size = std::min(candidates.size(), kMaxCandidatePoolSize); + std::partial_sort( + candidates.begin(), + candidates.begin() + static_cast(min_size), + candidates.end() + ); + candidates.resize(min_size); + + // prune and update qg + new_neighbors_[cur_id].reserve(degree_bound_); + heuristic_prune(cur_id, candidates, new_neighbors_[cur_id], record_pruned); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } +} + +void QGBuilder::add_reverse_edges(bool record_pruned) { + std::vector locks(num_nodes_); + for (auto& row : reverse_buffer_) { + row.clear(); + } + + // Keep new_neighbors_ read-only while reverse candidates are collected. Mutating a + // destination row here races with another worker reading that row as its source. + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const PID data_id = static_cast(index); + for (const auto& nei : new_neighbors_[data_id]) { + const PID destination = nei.id; + const CandidateList& destination_neighbors = + new_neighbors_[destination]; + const bool reciprocal = std::any_of( + destination_neighbors.begin(), + destination_neighbors.end(), + [&](const auto& destination_neighbor) { + return destination_neighbor.id == data_id; + } + ); + if (reciprocal) { + continue; + } + + std::lock_guard lock(locks[destination]); + const size_t missing_slots = + degree_bound_ - + std::min(degree_bound_, destination_neighbors.size()); + if (reverse_buffer_[destination].size() < + kMaxCandidatePoolSize + missing_slots) { + reverse_buffer_[destination].emplace_back(data_id, nei.distance); + } + } + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } + + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const PID data_id = static_cast(index); + CandidateList& tmp_pool = reverse_buffer_[data_id]; + if (qg_.is_quantized() && !tmp_pool.empty()) { + // RaBitQ estimates are directional: score destination -> source afresh. + BuildQuery query(qg_, build_cache()); + query.reset(data_id); + for (auto& candidate : tmp_pool) { + candidate.distance = query.distance_to(candidate.id); + } + } + tmp_pool.reserve(tmp_pool.size() + degree_bound_); + tmp_pool.insert( + tmp_pool.end(), + new_neighbors_[data_id].begin(), + new_neighbors_[data_id].end() + ); + std::sort(tmp_pool.begin(), tmp_pool.end()); + heuristic_prune(data_id, tmp_pool, new_neighbors_[data_id], record_pruned); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } +} + +void QGBuilder::random_init() { + const PID min_id = 0; + const PID max_id = num_nodes_ - 1; + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + new_neighbors_[i].clear(); + pruned_neighbors_[i].clear(); + std::mt19937 rng(seed_ ^ static_cast(i)); + std::uniform_int_distribution distribution(min_id, max_id); + std::unordered_set neighbor_set; + neighbor_set.reserve(degree_bound_); + while (neighbor_set.size() < degree_bound_) { + const PID rand_id = distribution(rng); + if (rand_id != i) { + neighbor_set.emplace(rand_id); + } + } + + BuildQuery query(qg_, build_cache()); + query.reset(i); + new_neighbors_[i].reserve(degree_bound_); + for (PID cur_neigh : neighbor_set) { + new_neighbors_[i].emplace_back(cur_neigh, query.distance_to(cur_neigh)); + } + + write_row( + static_cast(i), + new_neighbors_[i], + workers_[static_cast(omp_get_thread_num())] + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } +} + +/** + * @brief Complete each vertex's degree to the degree bound (multiple of 32), using + * pruned candidates and random fallback edges. + * + */ +void QGBuilder::complete_degrees() { + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + CandidateList& cur_neighbors = new_neighbors_[i]; + size_t cur_degree = cur_neighbors.size(); + + // skip vertices with enough neighbors + if (cur_degree >= degree_bound_) { + continue; + } + + CandidateList& pruned_list = pruned_neighbors_[i]; + CandidateList new_result; + new_result.reserve(degree_bound_); + + std::sort(pruned_list.begin(), pruned_list.end()); + + // use binary search to get refined results + CandidateList successful_result; + successful_result.reserve(degree_bound_); + float left = 0.5; + float right = 1.0; + size_t iter = 0; + while (iter++ < kMaxBsIter) { + float mid = (left + right) / 2; + add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); + if (new_result.size() < degree_bound_) { + left = mid; + } else { + right = mid; + successful_result = new_result; + } + } + + // Every successful right endpoint was already evaluated at that exact + // threshold. + if (!successful_result.empty()) { + new_result.swap(successful_result); + } else { + add_pruned_edges(cur_neighbors, pruned_list, new_result, right); + } + + // if the vertex still doesn't have enough neighbors, use random vertices + if (new_result.size() < degree_bound_) { + std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); + std::uniform_int_distribution distribution( + 0, static_cast(num_nodes_) - 1 + ); + std::unordered_set ids; + ids.reserve(degree_bound_); + for (auto& neighbor : new_result) { + ids.emplace(neighbor.id); + } + BuildQuery query(qg_, build_cache()); + query.reset(i); + while (new_result.size() < degree_bound_) { + const PID rand_id = distribution(rng); + if (rand_id != static_cast(i) && + ids.find(rand_id) == ids.end()) { + new_result.emplace_back(rand_id, query.distance_to(rand_id)); + ids.emplace(rand_id); + } + } + } + + cur_neighbors = new_result; + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } +} + +void QGBuilder::run_pass(bool final_pass) { + if (final_pass) { + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + const size_t i = static_cast(index); + pruned_neighbors_[i].clear(); + pruned_neighbors_[i].reserve(kMaxPrunedSize); + } + } + + const bool record_pruned = final_pass; + search_new_neighbors(record_pruned); + + add_reverse_edges(record_pruned); + + // Only the final pass completes degrees; every pass writes updated rows below. + if (final_pass) { + complete_degrees(); + } + + // update qg + { + index_detail::ParallelException failure; +#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + ++index) { + if (failure.failed()) { + continue; + } + try { + const size_t i = static_cast(index); + write_row( + static_cast(i), + new_neighbors_[i], + workers_[static_cast(omp_get_thread_num())] + ); + } catch (...) { failure.capture(); } + } + failure.rethrow(); + } +} +void QGBuilder::BuildQuery::reset(PID id) { + raw_ = nullptr; + rotated_ = nullptr; + prepared_.reset(); + if (cache_ != nullptr) { + rotated_ = cache_->rotated_vectors.data() + id * graph_.padded_dim_; + } + if (graph_.is_quantized()) { + if (cache_ != nullptr) { + prepared_.emplace(rotated_, cache_->queries[id]); + } else { + storage_.resize(graph_.padded_dim_); + // Reconstructed vectors are already rotated; do not rotate them again. + graph_.reconstruct_quantized_vector(id, storage_.data()); + rotated_ = storage_.data(); + prepared_.emplace( + rotated_, graph_.centroid_.data(), graph_.padded_dim_, graph_.metric_type_ + ); + } + } else { + raw_ = graph_.get_vector(id); + } +} + +float QGBuilder::BuildQuery::distance_to(PID target) const { + return graph_.point_distance(raw_, prepared_ ? &*prepared_ : nullptr, target); +} + +const float* QGBuilder::BuildQuery::rotated_query() { + if (rotated_ == nullptr) { + storage_.resize(graph_.padded_dim_); + graph_.rotator_->rotate(raw_, storage_.data()); + rotated_ = storage_.data(); + } + return rotated_; +} + +// find candidate neighbors for cur_id, exclude the vertex itself +void QGBuilder::find_candidates(PID cur_id, CandidateList& results, VisitedSet& vis) const { + BuildQuery query(qg_, build_cache()); + query.reset(cur_id); + BatchQuery q_obj(query.rotated_query(), qg_.padded_dim_, qg_.metric_type_); + buffer::SearchBuffer tmp_pool(ef_build_); + tmp_pool.insert(qg_.entry_point_, std::numeric_limits::max()); + + std::vector est_dist(degree_bound_); + while (tmp_pool.has_next()) { + auto cur_candi = tmp_pool.pop(); + if (vis.get(cur_candi)) { + continue; + } + vis.set(cur_candi); + auto cur_degree = degrees_[cur_candi]; + const float vertex_distance = query.distance_to(cur_candi); + q_obj.set_g_add(vertex_distance); + qg_.scan_neighbors(q_obj, cur_candi, est_dist.data(), tmp_pool, vis, cur_degree); + if (cur_candi != cur_id) { + results.emplace_back(cur_candi, vertex_distance); + } + } +} + +} // namespace rabitqlib::symqg diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index b841d38a..7049b898 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -1,80 +1,20 @@ -#include "rabitqlib/simd/dispatch.hpp" - #include #include -#include #include -#include -#include - -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/highacc_fastscan.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" -#include "rabitqlib/simd/fastscan_dispatch.hpp" -#include "rabitqlib/simd/hnsw_dispatch.hpp" -#include "rabitqlib/simd/matrix_dispatch.hpp" -#include "rabitqlib/simd/pack_excode_dispatch.hpp" -#include "rabitqlib/simd/quantization_dispatch.hpp" -#include "rabitqlib/simd/rotator_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" + +#include "dispatch_resolver.hpp" +#include "excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/warmup_space.hpp" namespace rabitqlib::simd { -namespace { -// Resolve once during static initialization; wrappers never repeat CPU checks. -// The override is used only by kernels needing a stricter AVX-512 subset. -template -Function resolve_kernel( - Function avx512, - Function avx2, - Function fallback, - bool avx512_supported = cpu::has_avx512_core() -) { - if (avx512_supported) { - return avx512; - } - if (cpu::has_avx2()) { - return avx2; - } - return fallback; -} - -template -Function resolve_kernel(Function avx2, Function fallback) { - if (cpu::has_avx2()) { - return avx2; - } - return fallback; -} - -// Use a distinct name: a fallback function pointer can implicitly convert to bool. -// Overloading the x86 resolver would let MSVC select this overload when __restrict -// makes the SIMD and scalar function-pointer types differ. -template -Function resolve_optional_kernel(Function preferred, Function fallback, bool supported) { - return supported ? preferred : fallback; -} - -} // namespace - -// Non-x86 builds must not reference symbols from excluded AVX objects. -#if defined(__x86_64__) || defined(_M_X64) -#define RABITQ_RESOLVE(...) rabitqlib::simd::resolve_kernel(__VA_ARGS__) -#define RABITQ_RESOLVE2(...) rabitqlib::simd::resolve_kernel(__VA_ARGS__) -#else -#define RABITQ_RESOLVE(avx512, avx2, fallback, ...) fallback -#define RABITQ_RESOLVE2(avx2, fallback) fallback -#endif - -const auto kAccumulateClusterSumsFn = RABITQ_RESOLVE( - accumulate_cluster_sums_avx512, - accumulate_cluster_sums_avx2, - accumulate_cluster_sums_generic -); void accumulate_cluster_sums( const float* x, @@ -85,12 +25,15 @@ void accumulate_cluster_sums( double* sums, size_t* counts ) { + static const auto kAccumulateClusterSumsFn = RABITQ_RESOLVE( + accumulate_cluster_sums_avx512, + accumulate_cluster_sums_avx2, + accumulate_cluster_sums_generic + ); + kAccumulateClusterSumsFn(x, point_ids, count, dim, labels, sums, counts); } -const auto kMatrixProductFn = - RABITQ_RESOLVE(matrix_product_avx512, matrix_product_avx2, matrix_product_generic); - void matrix_product( const float* left, const float* right, @@ -99,15 +42,12 @@ void matrix_product( size_t inner, size_t cols ) { + static const auto kMatrixProductFn = + RABITQ_RESOLVE(matrix_product_avx512, matrix_product_avx2, matrix_product_generic); + kMatrixProductFn(left, right, result, rows, inner, cols); } -const auto kMatrixProductTransposedFn = RABITQ_RESOLVE( - matrix_product_transposed_avx512, - matrix_product_transposed_avx2, - matrix_product_transposed_generic -); - void matrix_product_transposed( const float* left, const float* right, @@ -116,22 +56,22 @@ void matrix_product_transposed( size_t inner, size_t cols ) { + static const auto kMatrixProductTransposedFn = RABITQ_RESOLVE( + matrix_product_transposed_avx512, + matrix_product_transposed_avx2, + matrix_product_transposed_generic + ); + kMatrixProductTransposedFn(left, right, result, rows, inner, cols); } -const auto kRowNormsFn = - RABITQ_RESOLVE(row_norms_avx512, row_norms_avx2, row_norms_generic); - void row_norms(const float* data, float* result, size_t rows, size_t dim) { + static const auto kRowNormsFn = + RABITQ_RESOLVE(row_norms_avx512, row_norms_avx2, row_norms_generic); + kRowNormsFn(data, result, rows, dim); } -const auto kPairwiseDistancesLowerFn = RABITQ_RESOLVE( - pairwise_distances_lower_avx512, - pairwise_distances_lower_avx2, - pairwise_distances_lower_generic -); - void pairwise_distances_lower( const float* data, float* result, @@ -140,87 +80,63 @@ void pairwise_distances_lower( size_t dim, bool inner_product ) { - kPairwiseDistancesLowerFn(data, result, norms_data, size, dim, inner_product); -} - -const auto kQgBatchEstdistFn = RABITQ_RESOLVE( - qg_batch_estdist_avx512, qg_batch_estdist_avx2, qg_batch_estdist_generic -); - -void qg_batch_estdist( - const char* batch_data, - const BatchQuery& q_obj, - size_t padded_dim, - float* est_distance -) { - kQgBatchEstdistFn(batch_data, q_obj, padded_dim, est_distance); -} - -const auto kQgBatchEstdistMaskFn = - RABITQ_RESOLVE2(qg_batch_estdist_mask_avx2, qg_batch_estdist_mask_generic); - -uint32_t qg_batch_estdist_mask( - const char* batch_data, - const BatchQuery& q_obj, - size_t padded_dim, - float* est_distance, - float threshold -) { - return kQgBatchEstdistMaskFn(batch_data, q_obj, padded_dim, est_distance, threshold); -} - -const auto kSplitBatchEstdistFn = RABITQ_RESOLVE( - split_batch_estdist_avx512, split_batch_estdist_avx2, split_batch_estdist_generic -); - -void split_batch_estdist( - const char* batch_data, - const SplitBatchQuery& q_obj, - size_t padded_dim, - float* est_distance, - float* low_distance, - float* ip_x0_qr, - bool use_hacc -) { - kSplitBatchEstdistFn( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + static const auto kPairwiseDistancesLowerFn = RABITQ_RESOLVE( + pairwise_distances_lower_avx512, + pairwise_distances_lower_avx2, + pairwise_distances_lower_generic ); + + kPairwiseDistancesLowerFn(data, result, norms_data, size, dim, inner_product); } +float euclidean_sqr(const float* a, const float* b, size_t dim) { #if defined(__aarch64__) -const auto kEuclideanSqrFn = - resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); -const auto kDotProductFn = - resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); -const auto kDotProductDisFn = - resolve_optional_kernel(dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()); -const auto kL2normSqrFn = - resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); + static const auto kEuclideanSqrFn = + resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); #else -const auto kEuclideanSqrFn = - RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); -const auto kDotProductFn = - RABITQ_RESOLVE(dot_product_avx512, dot_product_avx2, dot_product_generic); -const auto kDotProductDisFn = - RABITQ_RESOLVE(dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic); -const auto kL2normSqrFn = - RABITQ_RESOLVE(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic); - + static const auto kEuclideanSqrFn = + RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); #endif -float euclidean_sqr(const float* a, const float* b, size_t dim) { return kEuclideanSqrFn(a, b, dim); } float dot_product(const float* a, const float* b, size_t dim) { +#if defined(__aarch64__) + static const auto kDotProductFn = + resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); +#else + static const auto kDotProductFn = + RABITQ_RESOLVE(dot_product_avx512, dot_product_avx2, dot_product_generic); +#endif + return kDotProductFn(a, b, dim); } float dot_product_dis(const float* a, const float* b, size_t dim) { +#if defined(__aarch64__) + static const auto kDotProductDisFn = resolve_optional_kernel( + dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon() + ); +#else + static const auto kDotProductDisFn = RABITQ_RESOLVE( + dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic + ); +#endif + return kDotProductDisFn(a, b, dim); } -float l2norm_sqr(const float* a, size_t dim) { return kL2normSqrFn(a, dim); } +float l2norm_sqr(const float* a, size_t dim) { +#if defined(__aarch64__) + static const auto kL2normSqrFn = + resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); +#else + static const auto kL2normSqrFn = + RABITQ_RESOLVE(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic); +#endif + return kL2normSqrFn(a, dim); +} // With zero extra bits there is no extra code to contribute to the inner // product, so the ex_bits == 0 slot must be a constant-zero stub rather than @@ -231,23 +147,35 @@ static float ip_fxu0( return 0.0F; } -const auto kBestRescaleFactorFn = RABITQ_RESOLVE( - best_rescale_factor_avx512, best_rescale_factor_avx2, best_rescale_factor_generic -); - double best_rescale_factor( const float* magnitudes, size_t dim, int max_code, double start, double end ) { + static const auto kBestRescaleFactorFn = RABITQ_RESOLVE( + best_rescale_factor_avx512, best_rescale_factor_avx2, best_rescale_factor_generic + ); + return kBestRescaleFactorFn(magnitudes, dim, max_code, start, end); } -#if defined(__aarch64__) || defined(_M_ARM64) -const auto kFhtRotateFn = - resolve_optional_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon()); -#else -const auto kFhtRotateFn = - RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic); -#endif +double best_rescale_factor_with_workspace( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +) { + static const auto kBestRescaleWorkspaceFn = RABITQ_RESOLVE( + best_rescale_factor_with_workspace_avx512, + best_rescale_factor_with_workspace_avx2, + best_rescale_factor_with_workspace_generic + ); + + return kBestRescaleWorkspaceFn( + magnitudes, dim, max_code, start, end, promoted, reciprocals + ); +} void fht_rotate( const float* data, @@ -258,6 +186,14 @@ void fht_rotate( float fac, const uint8_t* flip ) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const auto kFhtRotateFn = + resolve_optional_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon()); +#else + static const auto kFhtRotateFn = + RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic); +#endif + kFhtRotateFn(data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip); } @@ -317,99 +253,110 @@ ExcodeIpTable resolve_excode_ip_table() { #endif } +void flip_sign(const uint8_t* flip, float* data, size_t dim) { #if defined(__aarch64__) || defined(_M_ARM64) -const auto kFlipSignFn = - resolve_optional_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon()); + static const auto kFlipSignFn = + resolve_optional_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon()); #else -const auto kFlipSignFn = - RABITQ_RESOLVE(flip_sign_avx512, flip_sign_avx2, flip_sign_generic); + static const auto kFlipSignFn = + RABITQ_RESOLVE(flip_sign_avx512, flip_sign_avx2, flip_sign_generic); #endif -#if defined(__aarch64__) || defined(_M_ARM64) -const auto kKacsWalkFn = - resolve_optional_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon()); -#else -const auto kKacsWalkFn = - RABITQ_RESOLVE(kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic); -#endif + kFlipSignFn(flip, data, dim); +} +void kacs_walk(float* data, size_t len) { #if defined(__aarch64__) || defined(_M_ARM64) -const auto kScalarQuantizeUint8Fn = resolve_optional_kernel( - scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() -); -const auto kScalarQuantizeUint16Fn = resolve_optional_kernel( - scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon() -); + static const auto kKacsWalkFn = + resolve_optional_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon()); #else -const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE( - scalar_quantize_uint8_avx512, scalar_quantize_uint8_avx2, scalar_quantize_uint8_generic -); - -const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE( - scalar_quantize_uint16_avx512, - scalar_quantize_uint16_avx2, - scalar_quantize_uint16_generic -); + static const auto kKacsWalkFn = + RABITQ_RESOLVE(kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic); #endif - -const auto kPacking2BitExcodeFn = RABITQ_RESOLVE( - packing_2bit_excode_avx512, packing_2bit_excode_avx2, packing_2bit_excode_generic -); -const auto kPacking3BitExcodeFn = RABITQ_RESOLVE( - packing_3bit_excode_avx512, packing_3bit_excode_avx2, packing_3bit_excode_generic -); -const auto kPacking4BitExcodeFn = RABITQ_RESOLVE( - packing_4bit_excode_avx512, packing_4bit_excode_avx2, packing_4bit_excode_generic -); -const auto kPacking5BitExcodeFn = RABITQ_RESOLVE( - packing_5bit_excode_avx512, packing_5bit_excode_avx2, packing_5bit_excode_generic -); -const auto kPacking6BitExcodeFn = RABITQ_RESOLVE( - packing_6bit_excode_avx512, packing_6bit_excode_avx2, packing_6bit_excode_generic -); -const auto kPacking7BitExcodeFn = RABITQ_RESOLVE( - packing_7bit_excode_avx512, packing_7bit_excode_avx2, packing_7bit_excode_generic -); - -void flip_sign(const uint8_t* flip, float* data, size_t dim) { - kFlipSignFn(flip, data, dim); + kKacsWalkFn(data, len); } -void kacs_walk(float* data, size_t len) { kKacsWalkFn(data, len); } - void scalar_quantize_uint8( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const auto kScalarQuantizeUint8Fn = resolve_optional_kernel( + scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() + ); +#else + static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE( + scalar_quantize_uint8_avx512, + scalar_quantize_uint8_avx2, + scalar_quantize_uint8_generic + ); +#endif + kScalarQuantizeUint8Fn(result, vec0, dim, lo, delta); } void scalar_quantize_uint16( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const auto kScalarQuantizeUint16Fn = resolve_optional_kernel( + scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon() + ); +#else + static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE( + scalar_quantize_uint16_avx512, + scalar_quantize_uint16_avx2, + scalar_quantize_uint16_generic + ); +#endif + kScalarQuantizeUint16Fn(result, vec0, dim, lo, delta); } void packing_2bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking2BitExcodeFn = RABITQ_RESOLVE( + packing_2bit_excode_avx512, packing_2bit_excode_avx2, packing_2bit_excode_generic + ); + kPacking2BitExcodeFn(o_raw, o_compact, dim); } void packing_3bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking3BitExcodeFn = RABITQ_RESOLVE( + packing_3bit_excode_avx512, packing_3bit_excode_avx2, packing_3bit_excode_generic + ); + kPacking3BitExcodeFn(o_raw, o_compact, dim); } void packing_4bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking4BitExcodeFn = RABITQ_RESOLVE( + packing_4bit_excode_avx512, packing_4bit_excode_avx2, packing_4bit_excode_generic + ); + kPacking4BitExcodeFn(o_raw, o_compact, dim); } void packing_5bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking5BitExcodeFn = RABITQ_RESOLVE( + packing_5bit_excode_avx512, packing_5bit_excode_avx2, packing_5bit_excode_generic + ); + kPacking5BitExcodeFn(o_raw, o_compact, dim); } void packing_6bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking6BitExcodeFn = RABITQ_RESOLVE( + packing_6bit_excode_avx512, packing_6bit_excode_avx2, packing_6bit_excode_generic + ); + kPacking6BitExcodeFn(o_raw, o_compact, dim); } void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { + static const auto kPacking7BitExcodeFn = RABITQ_RESOLVE( + packing_7bit_excode_avx512, packing_7bit_excode_avx2, packing_7bit_excode_generic + ); + kPacking7BitExcodeFn(o_raw, o_compact, dim); } @@ -417,55 +364,16 @@ void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { namespace rabitqlib { -const simd::ExcodeIpTable kExcodeIpTable = simd::resolve_excode_ip_table(); - -const ex_ipfunc kIp16Fxu1AvxFn = kExcodeIpTable[1]; -const ex_ipfunc kIp64Fxu2AvxFn = kExcodeIpTable[2]; -const ex_ipfunc kIp64Fxu3AvxFn = kExcodeIpTable[3]; -const ex_ipfunc kIp16Fxu4AvxFn = kExcodeIpTable[4]; -const ex_ipfunc kIp64Fxu5AvxFn = kExcodeIpTable[5]; -const ex_ipfunc kIp64Fxu6AvxFn = kExcodeIpTable[6]; -const ex_ipfunc kIp64Fxu7AvxFn = kExcodeIpTable[7]; - -#if defined(__aarch64__) || defined(_M_ARM64) -const auto kNewTransposeBinFn = simd::resolve_optional_kernel( - simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon() -); -const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel( - simd::new_transpose_bin_512_neon, simd::new_transpose_bin_512_generic, cpu::has_neon() -); -#else -const auto kNewTransposeBinFn = RABITQ_RESOLVE( - simd::new_transpose_bin_avx512, - simd::new_transpose_bin_avx2, - simd::new_transpose_bin_generic -); - -const auto kNewTransposeBin512Fn = RABITQ_RESOLVE( - simd::new_transpose_bin_512_avx512, - simd::new_transpose_bin_512_avx2, - simd::new_transpose_bin_512_generic -); -#endif +static const simd::ExcodeIpTable& excode_ip_table() { + static const auto table = simd::resolve_excode_ip_table(); + return table; +} using MaskIpX0QFn = float (*)(const float*, const uint8_t*, size_t); -#if defined(__aarch64__) || defined(_M_ARM64) -const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_optional_kernel( - static_cast(simd::mask_ip_x0_q_neon), - static_cast(simd::mask_ip_x0_q_generic), - cpu::has_neon() -); -#else -const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE( - static_cast(simd::mask_ip_x0_q_avx512), - static_cast(simd::mask_ip_x0_q_avx2), - static_cast(simd::mask_ip_x0_q_generic) -); -#endif ex_ipfunc select_excode_ipfunc(size_t ex_bits) { if (ex_bits <= 8) { - return kExcodeIpTable[ex_bits]; + return excode_ip_table()[ex_bits]; } throw std::invalid_argument("Bad IP function for IVF"); @@ -474,162 +382,109 @@ ex_ipfunc select_excode_ipfunc(size_t ex_bits) { float excode_ipimpl::ip16_fxu1_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp16Fxu1AvxFn(query, compact_code, dim); + return excode_ip_table()[1](query, compact_code, dim); } float excode_ipimpl::ip64_fxu2_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp64Fxu2AvxFn(query, compact_code, dim); + return excode_ip_table()[2](query, compact_code, dim); } float excode_ipimpl::ip64_fxu3_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp64Fxu3AvxFn(query, compact_code, dim); + return excode_ip_table()[3](query, compact_code, dim); } float excode_ipimpl::ip16_fxu4_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp16Fxu4AvxFn(query, compact_code, dim); + return excode_ip_table()[4](query, compact_code, dim); } float excode_ipimpl::ip64_fxu5_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp64Fxu5AvxFn(query, compact_code, dim); + return excode_ip_table()[5](query, compact_code, dim); } float excode_ipimpl::ip64_fxu6_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp64Fxu6AvxFn(query, compact_code, dim); + return excode_ip_table()[6](query, compact_code, dim); } float excode_ipimpl::ip64_fxu7_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return kIp64Fxu7AvxFn(query, compact_code, dim); + return excode_ip_table()[7](query, compact_code, dim); } void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const auto kNewTransposeBinFn = simd::resolve_optional_kernel( + simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon() + ); +#else + static const auto kNewTransposeBinFn = RABITQ_RESOLVE( + simd::new_transpose_bin_avx512, + simd::new_transpose_bin_avx2, + simd::new_transpose_bin_generic + ); +#endif + kNewTransposeBinFn(q, tq, padded_dim, b_query); } void new_transpose_bin_512( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { - kNewTransposeBin512Fn(q, tq, padded_dim, b_query); -} - -float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim) { - return kMaskIpX0QFn(query, data, padded_dim); -} - -float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) { - return mask_ip_x0_q(query, reinterpret_cast(data), padded_dim); -} - -} // namespace rabitqlib - -namespace rabitqlib::fastscan { - #if defined(__aarch64__) || defined(_M_ARM64) -const auto kPackLutFn = rabitqlib::simd::resolve_optional_kernel( - simd::pack_lut_neon, simd::pack_lut_generic, cpu::has_neon() -); + static const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel( + simd::new_transpose_bin_512_neon, + simd::new_transpose_bin_512_generic, + cpu::has_neon() + ); #else -const auto kPackLutFn = - RABITQ_RESOLVE(simd::pack_lut_avx512, simd::pack_lut_avx2, simd::pack_lut_generic); + static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE( + simd::new_transpose_bin_512_avx512, + simd::new_transpose_bin_512_avx2, + simd::new_transpose_bin_512_generic + ); #endif -template <> -void pack_lut(size_t dim, const float* __restrict__ query, float* __restrict__ lut) { - kPackLutFn(dim, query, lut); + kNewTransposeBin512Fn(q, tq, padded_dim, b_query); } -#if defined(__aarch64__) -const auto kAccumulateFn = rabitqlib::simd::resolve_optional_kernel( - simd::accumulate_neon, simd::accumulate_unsupported, cpu::has_neon() -); -#else -const auto kAccumulateFn = RABITQ_RESOLVE( - simd::accumulate_avx512, simd::accumulate_avx2, simd::accumulate_unsupported -); -#endif - -const auto kTransferLutHaccFn = RABITQ_RESOLVE( - simd::transfer_lut_hacc_avx512, - simd::transfer_lut_hacc_avx2, - simd::transfer_lut_hacc_generic -); - -#if defined(__aarch64__) -const auto kAccumulateHaccFn = rabitqlib::simd::resolve_optional_kernel( - simd::accumulate_hacc_neon, simd::accumulate_hacc_generic, cpu::has_neon() -); +float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_optional_kernel( + static_cast(simd::mask_ip_x0_q_neon), + static_cast(simd::mask_ip_x0_q_generic), + cpu::has_neon() + ); #else -const auto kAccumulateHaccFn = RABITQ_RESOLVE( - simd::accumulate_hacc_avx512, simd::accumulate_hacc_avx2, simd::accumulate_hacc_generic -); + static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE( + static_cast(simd::mask_ip_x0_q_avx512), + static_cast(simd::mask_ip_x0_q_avx2), + static_cast(simd::mask_ip_x0_q_generic) + ); #endif -void accumulate( - const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, - int32_t* __restrict__ result, - size_t dim -) { - if (dim == 0 || dim % 16 != 0) { - throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); - } - kAccumulateFn(codes, lp_table, result, dim); -} - -void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { - if (dim == 0 || dim % 16 != 0) { - throw std::invalid_argument( - "high-accuracy FastScan dimension must be a positive multiple of 16" - ); - } - kTransferLutHaccFn(lut, dim, hc_lut); + return kMaskIpX0QFn(query, data, padded_dim); } -void accumulate_hacc( - const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, - size_t dim -) { - if (dim == 0 || dim % 16 != 0) { - throw std::invalid_argument( - "high-accuracy FastScan dimension must be a positive multiple of 16" - ); - } - kAccumulateHaccFn(codes, hc_lut, accu_res, dim); +float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) { + return mask_ip_x0_q(query, reinterpret_cast(data), padded_dim); } -} // namespace rabitqlib::fastscan +} // namespace rabitqlib namespace rabitqlib { using WarmupIpX0Q512Fn = float (*)(const uint8_t*, const uint64_t*, float, float, size_t, size_t); -#if defined(__aarch64__) || defined(_M_ARM64) -const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = simd::resolve_optional_kernel( - static_cast(simd::warmup_ip_x0_q_512_neon), - static_cast(simd::warmup_ip_x0_q_512_generic), - cpu::has_neon() -); -#else -const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE( - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_generic), - cpu::has_avx512_popcnt() -); -#endif float warmup_ip_x0_q_512( const uint8_t* data, @@ -639,6 +494,21 @@ float warmup_ip_x0_q_512( size_t padded_dim, size_t b_query ) { +#if defined(__aarch64__) || defined(_M_ARM64) + static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = simd::resolve_optional_kernel( + static_cast(simd::warmup_ip_x0_q_512_neon), + static_cast(simd::warmup_ip_x0_q_512_generic), + cpu::has_neon() + ); +#else + static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE( + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_generic), + cpu::has_avx512_popcnt() + ); +#endif + return kWarmupIpX0Q512Fn(data, query, delta, vl, padded_dim, b_query); } @@ -656,30 +526,3 @@ float warmup_ip_x0_q_512( } } // namespace rabitqlib - -namespace rabitqlib::hnsw::detail { -namespace { -using SearchKnnFn = std::priority_queue> (*)( - HierarchicalNSW&, const float*, size_t, size_t -); -// The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. -#if defined(__x86_64__) || defined(_M_X64) -const SearchKnnFn kSearchKnnFn = cpu::has_avx512_popcnt() - ? search_knn_avx512_popcnt - : RABITQ_RESOLVE( - search_knn_avx512_core, - search_knn_avx2, - search_knn_generic, - cpu::has_avx512_core() && cpu::has_avx2() - ); -#else -const SearchKnnFn kSearchKnnFn = - simd::resolve_optional_kernel(search_knn_neon, search_knn_generic, cpu::has_neon()); -#endif -} // namespace -std::priority_queue> search_knn( - HierarchicalNSW& index, const float* query, size_t topk, size_t ef -) { - return kSearchKnnFn(index, query, topk, ef); -} -} // namespace rabitqlib::hnsw::detail diff --git a/src/simd/dispatch_resolver.hpp b/src/simd/dispatch_resolver.hpp new file mode 100644 index 00000000..2aa2e2e6 --- /dev/null +++ b/src/simd/dispatch_resolver.hpp @@ -0,0 +1,60 @@ +#pragma once + +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::simd { +namespace { +// Resolve once on first use; function-local caches are safe before main(). +// The override is used only by kernels needing a stricter AVX-512 subset. +template +Function resolve_kernel( + Function avx512, + Function avx2, + Function fallback, + bool avx512_supported = cpu::has_avx512_core() +) { + if (avx512_supported) { + return avx512; + } + if (cpu::has_avx2()) { + return avx2; + } + return fallback; +} + +template +Function resolve_kernel(Function avx2, Function fallback) { + if (cpu::has_avx2()) { + return avx2; + } + return fallback; +} + +// Use a distinct name: a fallback function pointer can implicitly convert to bool. +// Overloading the x86 resolver would let MSVC select this overload when __restrict +// makes the SIMD and scalar function-pointer types differ. +template +Function resolve_optional_kernel(Function preferred, Function fallback, bool supported) { + return supported ? preferred : fallback; +} + +} // namespace +} // namespace rabitqlib::simd + +#if defined(__GNUC__) || defined(__clang__) +#define RABITQ_COLD_INIT __attribute__((cold, noinline)) +#else +#define RABITQ_COLD_INIT +#endif + +// Non-x86 builds must not reference symbols from excluded AVX objects. +#if defined(__x86_64__) || defined(_M_X64) +// Keep feature probes out of the hot wrapper's register-allocation scope. +// The calling function's local static still provides thread-safe initialization. +#define RABITQ_RESOLVE(...) \ + [&]() RABITQ_COLD_INIT { return rabitqlib::simd::resolve_kernel(__VA_ARGS__); }() +#define RABITQ_RESOLVE2(...) RABITQ_RESOLVE(__VA_ARGS__) +#else +#define RABITQ_RESOLVE(avx512, avx2, fallback, ...) fallback +#define RABITQ_RESOLVE2(avx2, fallback) fallback +#endif diff --git a/src/simd/estimator_avx512.cpp b/src/simd/estimator_avx512.cpp deleted file mode 100644 index 15abe7cc..00000000 --- a/src/simd/estimator_avx512.cpp +++ /dev/null @@ -1,30 +0,0 @@ -#include - -#include "estimator_kernels.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" - -namespace rabitqlib::simd { -void split_batch_estdist_avx512( - const char* batch_data, - const SplitBatchQuery& q_obj, - size_t padded_dim, - float* est_distance, - float* low_distance, - float* ip_x0_qr, - bool use_hacc -) { - split_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc - ); -} - -void qg_batch_estdist_avx512( - const char* batch_data, - const BatchQuery& q_obj, - size_t padded_dim, - float* est_distance -) { - qg_batch_estdist_impl(batch_data, q_obj, padded_dim, est_distance); -} - -} // namespace rabitqlib::simd diff --git a/include/rabitqlib/simd/dispatch.hpp b/src/simd/excode_dispatch.hpp similarity index 100% rename from include/rabitqlib/simd/dispatch.hpp rename to src/simd/excode_dispatch.hpp diff --git a/src/simd/hnsw_neon_kernels.hpp b/src/simd/hnsw_neon_kernels.hpp index 4fd800b2..e2600201 100644 --- a/src/simd/hnsw_neon_kernels.hpp +++ b/src/simd/hnsw_neon_kernels.hpp @@ -8,7 +8,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd::detail { diff --git a/src/simd/matrix_avx2.cpp b/src/simd/matrix_avx2.cpp index 38633019..e10d6347 100644 --- a/src/simd/matrix_avx2.cpp +++ b/src/simd/matrix_avx2.cpp @@ -1,6 +1,7 @@ #include +#include -#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" #define RABITQ_MATRIX_EIGEN_NAMESPACE rabitqlib_eigen_avx2 #include "matrix_kernels.hpp" diff --git a/src/simd/matrix_avx512.cpp b/src/simd/matrix_avx512.cpp index 88d75c5e..dd1d6333 100644 --- a/src/simd/matrix_avx512.cpp +++ b/src/simd/matrix_avx512.cpp @@ -1,6 +1,7 @@ #include +#include -#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" #define RABITQ_MATRIX_EIGEN_NAMESPACE rabitqlib_eigen_avx512 #include "matrix_kernels.hpp" diff --git a/src/simd/matrix_generic.cpp b/src/simd/matrix_generic.cpp index 0e9e83a2..aa5ddf13 100644 --- a/src/simd/matrix_generic.cpp +++ b/src/simd/matrix_generic.cpp @@ -1,6 +1,7 @@ #include +#include -#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" #define RABITQ_MATRIX_EIGEN_NAMESPACE rabitqlib_eigen_generic #include "matrix_kernels.hpp" diff --git a/src/simd/matrix_kernels.hpp b/src/simd/matrix_kernels.hpp index 9c34add2..745dcd05 100644 --- a/src/simd/matrix_kernels.hpp +++ b/src/simd/matrix_kernels.hpp @@ -10,7 +10,7 @@ #define RABITQ_MATRIX_EIGEN_NAMESPACE rabitqlib_eigen_generic #endif #define Eigen RABITQ_MATRIX_EIGEN_NAMESPACE -#include "rabitqlib/third/Eigen/Dense" +#include "rabitqlib/third/Eigen/Core" #undef Eigen namespace rabitqlib::simd { diff --git a/src/simd/pack_excode_avx2.cpp b/src/simd/pack_excode_avx2.cpp index 7cf15cf4..861c4743 100644 --- a/src/simd/pack_excode_avx2.cpp +++ b/src/simd/pack_excode_avx2.cpp @@ -2,7 +2,7 @@ #include #include "pack_excode_kernels.hpp" -#include "rabitqlib/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/pack_excode_avx512.cpp b/src/simd/pack_excode_avx512.cpp index 77e3e6cd..a9b0ab23 100644 --- a/src/simd/pack_excode_avx512.cpp +++ b/src/simd/pack_excode_avx512.cpp @@ -2,7 +2,7 @@ #include #include "pack_excode_kernels.hpp" -#include "rabitqlib/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index cbaa6316..a69f8542 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -100,7 +100,7 @@ inline void packing_2bit_excode_intrinsics( // Full blocks preserve the existing byte layout. for (size_t j = 0; j < dim - dim % 64; j += 64) { // pack 64 2-bit codes into 128 bits (16 bytes) - // the lower 2 bits of each byte represent vec00 to vec04... + // Byte i holds codes i, i+16, i+32, i+48 from low to high bits. __m128i vec_00_to_15 = _mm_loadu_si128(reinterpret_cast(o_raw)); __m128i vec_16_to_31 = _mm_loadu_si128(reinterpret_cast(o_raw + 16)); diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index 0f8acf57..55201e50 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -5,10 +5,10 @@ #include #include -#include "rabitqlib/simd/pack_excode_dispatch.hpp" -#include "rabitqlib/simd/rotator_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" namespace rabitqlib::simd { @@ -81,6 +81,21 @@ void pack(const uint8_t* raw, uint8_t* compact, size_t dim) { // Inverse of pack, over the same mapping. Not a search path, so scalar only. template void unpack(const uint8_t* compact, uint8_t* raw, size_t dim) { + if constexpr (Bits == 8) { + std::copy_n(compact, dim, raw); + return; + } + if constexpr (Bits == 4) { + // Each group of eight bytes holds two interleaved groups of eight nibbles. + for (size_t i = 0; i < dim; i += 16) { + for (size_t j = 0; j < 8; ++j) { + const uint8_t pair = compact[i / 2 + j]; + raw[i + j] = pair & 0x0f; + raw[i + 8 + j] = pair >> 4; + } + } + return; + } const size_t full = (Bits == 1 || Bits == 4 || Bits == 8) ? dim : dim - dim % 64; for (size_t i = 0; i < full; ++i) { unsigned code = 0; diff --git a/src/simd/quantization_avx2.cpp b/src/simd/quantization_avx2.cpp index 3502cd1f..c0a137fe 100644 --- a/src/simd/quantization_avx2.cpp +++ b/src/simd/quantization_avx2.cpp @@ -4,7 +4,7 @@ #include #include -#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" #include "rescale_search.hpp" namespace rabitqlib::simd::detail { @@ -165,4 +165,28 @@ double best_rescale_factor_avx2( ? best_t : -1; } +double best_rescale_factor_with_workspace_avx2( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +) { + double best_t = 0; + return detail::try_find_best_scale_by_interval_search( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx2, + best_t, + promoted, + reciprocals + ) + ? best_t + : -1; +} } // namespace rabitqlib::simd diff --git a/src/simd/quantization_avx512.cpp b/src/simd/quantization_avx512.cpp index ab729273..87f0636c 100644 --- a/src/simd/quantization_avx512.cpp +++ b/src/simd/quantization_avx512.cpp @@ -4,7 +4,7 @@ #include #include -#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" #include "rescale_search.hpp" namespace rabitqlib::simd::detail { @@ -156,4 +156,28 @@ double best_rescale_factor_avx512( ? best_t : -1; } +double best_rescale_factor_with_workspace_avx512( + const float* magnitudes, + size_t dim, + int max_code, + double start, + double end, + double* promoted, + double* reciprocals +) { + double best_t = 0; + return detail::try_find_best_scale_by_interval_search( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx512, + best_t, + promoted, + reciprocals + ) + ? best_t + : -1; +} } // namespace rabitqlib::simd diff --git a/src/simd/quantization_generic.cpp b/src/simd/quantization_generic.cpp index bcaeafe6..d3bd4527 100644 --- a/src/simd/quantization_generic.cpp +++ b/src/simd/quantization_generic.cpp @@ -1,11 +1,15 @@ #include -#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rescale_search.hpp" namespace rabitqlib::simd { // Request the existing scalar event sweep; it certifies ambiguous SIMD searches. double best_rescale_factor_generic(const float*, size_t, int, double, double) { return -1; } +double +best_rescale_factor_with_workspace_generic(const float*, size_t, int, double, double, double*, double*) { + return -1; +} namespace detail { RescaleScratch& get_thread_local_rescale_scratch(size_t dim) { diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index 914dbffc..c10b5a8b 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -5,8 +5,7 @@ #include #include -#include "rabitqlib/simd/fastscan_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd { namespace { @@ -123,28 +122,3 @@ void new_transpose_bin_512_neon( } } } // namespace rabitqlib::simd - -namespace rabitqlib::fastscan::simd { -void pack_lut_neon(size_t dim, const float* query, float* lut) { - // Each table enumerates q0..q3 using bits 3..0. Add in that order, - // including the initial +0, to preserve cancellation and signed zeros. - volatile float zero = 0; - const uint32x4_t select2 = {0, 0, UINT32_MAX, UINT32_MAX}; - const uint32x4_t select3 = {0, UINT32_MAX, 0, UINT32_MAX}; - for (size_t group = 0; group < dim / 4; ++group) { - const float32x4_t z = vdupq_n_f32(zero); - const float32x4_t q0 = vdupq_n_f32(query[0]), q1 = vdupq_n_f32(query[1]), - q2 = vdupq_n_f32(query[2]), q3 = vdupq_n_f32(query[3]); - const float32x4_t high = vaddq_f32(z, q0); - float32x4_t tables[4] = {z, vaddq_f32(z, q1), high, vaddq_f32(high, q1)}; - for (size_t i = 0; i < 4; ++i) { - auto table = tables[i]; - table = vbslq_f32(select2, vaddq_f32(table, q2), table); - table = vbslq_f32(select3, vaddq_f32(table, q3), table); - vst1q_f32(lut + i * 4, table); - } - query += 4; - lut += 16; - } -} -} // namespace rabitqlib::fastscan::simd diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index 9e96abdd..500778c2 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -227,7 +227,9 @@ static bool try_find_best_scale_by_interval_search( double t_start, double t_end, Evaluator evaluator, - double& best_t + double& best_t, + double* magnitudes = nullptr, + double* reciprocals = nullptr ) { if (dim == 0 || max_code == 0) { best_t = t_start; @@ -241,9 +243,11 @@ static bool try_find_best_scale_by_interval_search( // Promote float inputs exactly once and cache 1/a_i for repeated candidate // evaluations. The thread-local buffers are reused across vectors; each // call overwrites its active prefix. Zero reciprocals are masked by evaluators. - auto& scratch = get_thread_local_rescale_scratch(dim); - auto& magnitudes = scratch.magnitudes; - auto& reciprocals = scratch.reciprocals; + if (magnitudes == nullptr) { + auto& scratch = get_thread_local_rescale_scratch(dim); + magnitudes = scratch.magnitudes.data(); + reciprocals = scratch.reciprocals.data(); + } for (size_t i = 0; i < dim; ++i) { magnitudes[i] = static_cast(input[i]); reciprocals[i] = input[i] > 0 ? 1.0 / magnitudes[i] : 0; @@ -270,9 +274,7 @@ static bool try_find_best_scale_by_interval_search( if (evaluations == kEvaluationBudget) return false; ++evaluations; - result.state = evaluator( - magnitudes.data(), reciprocals.data(), dim, max_code, t_start, t_end, t - ); + result.state = evaluator(magnitudes, reciprocals, dim, max_code, t_start, t_end, t); // A valid evaluator must return a plateau containing the requested t. // Every t in that plateau produces the same codes and score. if (result.state.first > t || result.state.next <= t) diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index 0f30eb94..82b51838 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -4,7 +4,7 @@ #include #include -#include "rabitqlib/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" #include "rotator_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index c04b0173..f72ba705 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -4,7 +4,7 @@ #include #include -#include "rabitqlib/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" #include "rotator_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/simd/rotator_neon.cpp b/src/simd/rotator_neon.cpp index c06b9506..2c90b3c7 100644 --- a/src/simd/rotator_neon.cpp +++ b/src/simd/rotator_neon.cpp @@ -6,7 +6,7 @@ #include #include -#include "rabitqlib/simd/rotator_dispatch.hpp" +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" namespace rabitqlib::simd { namespace { diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index 8f624b8d..aaf2178d 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -6,8 +6,7 @@ #include #include "packed_tail_avx2.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "space_float_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 9fbf1dbc..71c70710 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -4,8 +4,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "space_float_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 0ac1ac67..8d197f6d 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -4,7 +4,7 @@ #include #include "packed_tail_avx2.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd::excode_ipimpl { @@ -57,11 +57,10 @@ inline float mm256_reduce_add_ps(__m256 v) { return _mm_cvtss_f32(h); } -// ip16: this function is used to compute inner product of -// vectors padded to multiple of 16 -// fxu1: the inner product is computed between float and 1-bit unsigned int (lay out can be -// found rabitq_impl.hpp) -// avx512: only applicable for avx512 +// ip16_fxu1 consumes groups of 16 float query values and packed 1-bit codes. +// This translation unit uses AVX2/FMA; callers should use runtime dispatch. +// Packed extra-bit layouts are owned by core/pack_excode.hpp. + float ip16_fxu1_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { diff --git a/src/simd/space_excode_avx512.cpp b/src/simd/space_excode_avx512.cpp index 31b3233b..f5fdaccb 100644 --- a/src/simd/space_excode_avx512.cpp +++ b/src/simd/space_excode_avx512.cpp @@ -5,7 +5,7 @@ #include #include "packed_tail_avx2.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd::excode_ipimpl { @@ -68,11 +68,10 @@ inline __m512 excode_ip_tail_sum_avx512( } } // namespace -// ip16: this function is used to compute inner product of -// vectors padded to multiple of 16 -// fxu1: the inner product is computed between float and 1-bit unsigned int (lay out can be -// found rabitq_impl.hpp) -// avx512: only applicable for avx512 +// ip16_fxu1 consumes groups of 16 float query values and packed 1-bit codes. +// This translation unit uses AVX-512; callers should use runtime dispatch. +// Packed extra-bit layouts are owned by core/pack_excode.hpp. + float ip16_fxu1_avx512( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { diff --git a/src/simd/space_generic.cpp b/src/simd/space_generic.cpp index f4fd15a9..2a7e3766 100644 --- a/src/simd/space_generic.cpp +++ b/src/simd/space_generic.cpp @@ -1,7 +1,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd { namespace { diff --git a/src/simd/space_neon.cpp b/src/simd/space_neon.cpp index 55b579be..66493b2e 100644 --- a/src/simd/space_neon.cpp +++ b/src/simd/space_neon.cpp @@ -2,7 +2,7 @@ #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd { namespace { diff --git a/src/simd/space_packed_neon.cpp b/src/simd/space_packed_neon.cpp index a39da5fa..7170368e 100644 --- a/src/simd/space_packed_neon.cpp +++ b/src/simd/space_packed_neon.cpp @@ -4,8 +4,8 @@ #include #include "hnsw_neon_kernels.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" namespace rabitqlib::simd { namespace { diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index 2153072d..d0e8ed09 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -2,7 +2,7 @@ #include #include -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "warmup_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/simd/warmup_avx512.cpp b/src/simd/warmup_avx512.cpp index c4b1b9a6..1abfed24 100644 --- a/src/simd/warmup_avx512.cpp +++ b/src/simd/warmup_avx512.cpp @@ -2,7 +2,7 @@ #include #include -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "warmup_kernels.hpp" namespace rabitqlib::simd { diff --git a/src/utils/thread_count.cpp b/src/utils/thread_count.cpp index 2f558393..e2bc56e1 100644 --- a/src/utils/thread_count.cpp +++ b/src/utils/thread_count.cpp @@ -1,11 +1,10 @@ -#include - #include #include #include #include #include +#include "rabitqlib/utils/openmp.hpp" #include "rabitqlib/utils/tools.hpp" #if defined(__linux__) @@ -15,7 +14,7 @@ namespace rabitqlib { size_t total_threads() { -#if defined(_OPENMP) && _OPENMP >= 201511 +#if RABITQ_USE_OPENMP && defined(_OPENMP) && _OPENMP >= 201511 if (omp_get_proc_bind() != omp_proc_bind_false && omp_get_num_places() > 0) { // The runtime may pin the caller to one place. Count its whole partition, // including overlapping places only once, rather than that single mask. @@ -33,7 +32,7 @@ size_t total_threads() { return cpus.size(); } } -#elif defined(_OPENMP) && _OPENMP >= 201307 +#elif RABITQ_USE_OPENMP && defined(_OPENMP) && _OPENMP >= 201307 if (omp_get_proc_bind() != omp_proc_bind_false) { // OpenMP 4.0 has binding but no place-partition query routines. const int count = omp_get_num_procs(); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 1c8f0eee..84ab6e80 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -9,24 +9,37 @@ set(gtest_force_shared_crt ON CACHE BOOL "" FORCE) set(INSTALL_GTEST OFF CACHE BOOL "" FORCE) FetchContent_MakeAvailable(googletest) -# Automatically register all tests under unit and folders +# Automatically register the unit and kernel-reference tests. file(GLOB_RECURSE UNIT_TESTS CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/unit/*_test.cpp ) -file(GLOB_RECURSE INTEG_TESTS CONFIGURE_DEPENDS - ${CMAKE_CURRENT_SOURCE_DIR}/integration/*_test.cpp -) file(GLOB_RECURSE COMMON_SRCS CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/common/*.cpp ) +# Core-only builds must not compile index/clustering tests or their consumers. +set(RABITQ_TEST_TARGET rabitqlib::core) +if(RABITQ_BUILD_FASTSCAN) + set(RABITQ_TEST_TARGET rabitqlib::fastscan) +else() + list(FILTER UNIT_TESTS EXCLUDE REGEX "/fastscan/") +endif() +if(RABITQ_BUILD_INDEX) + set(RABITQ_TEST_TARGET rabitqlib::index) +else() + list(FILTER UNIT_TESTS EXCLUDE REGEX "/(index|clustering)/") + list(FILTER UNIT_TESTS EXCLUDE REGEX "/(neon_hnsw|thread_count|buffer|visited_set)_test[.]cpp$") +endif() + # add executables add_executable(rabitq_tests ${COMMON_SRCS} ${UNIT_TESTS} - ${INTEG_TESTS} ) target_include_directories(rabitq_tests PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/common) +if(RABITQ_BUILD_INDEX) + target_include_directories(rabitq_tests PRIVATE ${PROJECT_SOURCE_DIR}/src/index) +endif() set_source_files_properties(unit/rabitqlib/index/historical_persistence_test.cpp PROPERTIES COMPILE_DEFINITIONS "RABITQ_TEST_FIXTURE_DIR=\"${CMAKE_CURRENT_SOURCE_DIR}/fixtures/persistence/v0.5.2\"") @@ -36,35 +49,112 @@ rabitq_configure_target(rabitq_tests) target_link_libraries(rabitq_tests PRIVATE GTest::gtest_main - rabitq_headers + ${RABITQ_TEST_TARGET} ) # Discover tests for CTest include(GoogleTest) gtest_discover_tests(rabitq_tests DISCOVERY_MODE PRE_TEST) -# Binding is initialized before main; test it in its own process. -add_test(NAME OpenMPBinding COMMAND rabitq_tests - --gtest_filter=ThreadCount.BoundParallelForUsesRuntimePlacementAndPropagatesErrors) -set_tests_properties(OpenMPBinding PROPERTIES - ENVIRONMENT "OMP_PROC_BIND=spread;OMP_PLACES=threads;OMP_NUM_THREADS=4") - -# Substitute CPU predicates only in this executable to exercise real cached dispatch -# on a CPU without SIMD. The static library's detector object is not linked here. -add_executable(fastscan_no_simd_test dispatch/fastscan_no_simd_test.cpp) -target_link_libraries(fastscan_no_simd_test PRIVATE GTest::gtest_main rabitq_headers) -rabitq_configure_target(fastscan_no_simd_test) -gtest_discover_tests(fastscan_no_simd_test DISCOVERY_MODE PRE_TEST) - -# Keep the global allocation-failure override isolated from the main suite. -add_executable(hnsw_add_allocation_test failures/hnsw_add_allocation_test.cpp) -target_link_libraries(hnsw_add_allocation_test PRIVATE GTest::gtest_main rabitq_headers) -rabitq_configure_target(hnsw_add_allocation_test) -gtest_discover_tests(hnsw_add_allocation_test DISCOVERY_MODE PRE_TEST) +add_test(NAME CoreModuleBoundary COMMAND "${CMAKE_COMMAND}" + "-DRABITQ_SOURCE_DIR=${PROJECT_SOURCE_DIR}" + -P "${CMAKE_CURRENT_SOURCE_DIR}/core_consumer/check_core_boundary.cmake") + +add_test(NAME CoreOpenMPFlagCheck COMMAND "${CMAKE_COMMAND}" + "-DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/openmp-flag-check" + -P "${CMAKE_CURRENT_SOURCE_DIR}/core_consumer/check_no_openmp_test.cmake") + +add_executable(core_api_v1_consumer fixtures/core_api_v1/consumer.cpp) +target_link_libraries(core_api_v1_consumer PRIVATE rabitqlib::core) +add_test(NAME CoreApiV1Consumer COMMAND core_api_v1_consumer) + +foreach(mode IN ITEMS dispatch scalar) + set(target core_api_no_alloc_${mode}) + add_executable(${target} allocation/core_api_no_alloc.cpp) + target_link_libraries(${target} PRIVATE rabitqlib::core) + rabitq_configure_target(${target}) + if(mode STREQUAL "scalar") + target_compile_definitions(${target} PRIVATE RABITQ_FORCE_SCALAR) + endif() + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_compile_definitions(${target} PRIVATE RABITQ_WRAP_MALLOC) + target_link_options(${target} PRIVATE + -Wl,--wrap=malloc -Wl,--wrap=calloc -Wl,--wrap=realloc + -Wl,--wrap=aligned_alloc -Wl,--wrap=posix_memalign) + endif() + add_test(NAME ${target} COMMAND ${target}) +endforeach() + +if(RABITQ_BUILD_INDEX) + # Binding is initialized before main; test it in its own process. + add_test(NAME OpenMPBinding COMMAND rabitq_tests + --gtest_filter=ThreadCount.BoundParallelForUsesRuntimePlacementAndPropagatesErrors) + set_tests_properties(OpenMPBinding PROPERTIES + ENVIRONMENT "OMP_PROC_BIND=spread;OMP_PLACES=threads;OMP_NUM_THREADS=4") + + # Substitute CPU predicates only in this executable to exercise real cached dispatch + # on a CPU without SIMD. The static library's detector object is not linked here. + add_executable(fastscan_no_simd_test dispatch/fastscan_no_simd_test.cpp + dispatch/fastscan_backend_test.cpp) + target_include_directories(fastscan_no_simd_test PRIVATE ${PROJECT_SOURCE_DIR}/src/fastscan) + target_link_libraries(fastscan_no_simd_test PRIVATE GTest::gtest_main rabitqlib::index) + rabitq_configure_target(fastscan_no_simd_test) + gtest_discover_tests(fastscan_no_simd_test DISCOVERY_MODE PRE_TEST) + + add_executable(construction_failure_test failures/construction_failure_test.cpp) + target_link_libraries(construction_failure_test PRIVATE GTest::gtest_main rabitqlib::index) + rabitq_configure_target(construction_failure_test) + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_compile_definitions(construction_failure_test PRIVATE RABITQ_WRAP_CONSTRUCTION_ALLOCATIONS) + target_link_options(construction_failure_test PRIVATE + -Wl,--wrap=malloc -Wl,--wrap=aligned_alloc -Wl,--wrap=free) + endif() + gtest_discover_tests(construction_failure_test DISCOVERY_MODE PRE_TEST) + + # Keep the global allocation-failure override isolated from the main suite. + add_executable(hnsw_add_allocation_test failures/hnsw_add_allocation_test.cpp) + target_link_libraries(hnsw_add_allocation_test PRIVATE GTest::gtest_main rabitqlib::index) + rabitq_configure_target(hnsw_add_allocation_test) + gtest_discover_tests(hnsw_add_allocation_test DISCOVERY_MODE PRE_TEST) + +endif() # Message to verify files are found message(STATUS "Discovered test files:") -foreach(TEST_SRC ${UNIT_TESTS} ${INTEG_TESTS}) +foreach(TEST_SRC ${UNIT_TESTS}) file(RELATIVE_PATH REL_PATH ${CMAKE_CURRENT_SOURCE_DIR} ${TEST_SRC}) message(STATUS " - ${REL_PATH}") endforeach() + +add_executable(core_static_init core_consumer/core_static_init.cpp) +target_link_libraries(core_static_init PRIVATE rabitqlib::core) +add_test(NAME CoreStaticInitialization COMMAND core_static_init) + +if(RABITQ_BUILD_FASTSCAN) + add_executable(fastscan_backend_test dispatch/fastscan_backend_test.cpp) + target_include_directories(fastscan_backend_test PRIVATE ${PROJECT_SOURCE_DIR}/src/fastscan) + target_link_libraries(fastscan_backend_test PRIVATE GTest::gtest_main rabitqlib::fastscan) + rabitq_configure_target(fastscan_backend_test) + gtest_discover_tests(fastscan_backend_test DISCOVERY_MODE PRE_TEST TEST_PREFIX "Native.") + + add_executable(modules_static_init dispatch/modules_static_init.cpp) + target_link_libraries(modules_static_init PRIVATE ${RABITQ_TEST_TARGET}) + if(RABITQ_BUILD_INDEX) + target_compile_definitions(modules_static_init PRIVATE RABITQ_TEST_INDEX) + endif() + add_test(NAME ModulesStaticInitialization COMMAND modules_static_init) +endif() + +# Run allocation replacement in a separate process; this also exercises core-only builds. +add_executable(compact_codes_failure_test failures/compact_codes_failure_test.cpp) +target_link_libraries(compact_codes_failure_test PRIVATE GTest::gtest_main rabitqlib::core) +if(RABITQ_BUILD_FASTSCAN) + target_link_libraries(compact_codes_failure_test PRIVATE rabitqlib::fastscan) + target_compile_definitions(compact_codes_failure_test PRIVATE RABITQ_TEST_FASTSCAN) +endif() +rabitq_configure_target(compact_codes_failure_test) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + target_compile_definitions(compact_codes_failure_test PRIVATE RABITQ_WRAP_COMPACT_ALLOCATIONS) + target_link_options(compact_codes_failure_test PRIVATE -Wl,--wrap=malloc -Wl,--wrap=realloc) +endif() +gtest_discover_tests(compact_codes_failure_test DISCOVERY_MODE PRE_TEST) diff --git a/tests/README.md b/tests/README.md index e614c588..edf632df 100644 --- a/tests/README.md +++ b/tests/README.md @@ -17,6 +17,9 @@ and rejection of truncated or unsupported formats. - Git and network access during the first configuration so CMake can download GoogleTest 1.14.0 +OpenMP is optional for core/FastScan tests: configure `RABITQ_BUILD_INDEX=OFF` and +`RABITQ_USE_OPENMP=OFF`. ANN/Python testing retains the full prerequisites above. + On Windows, install Visual Studio 2026 with the Desktop development with C++ workload. Use MSVC for all C++ targets, including SIMD kernels. @@ -194,13 +197,107 @@ and an installed sanitizer build. Native tuning retains its existing default native-compiled library portable. Consumers choose their own optimization flags, including for inline code in the public headers. +## Core-only CMake package tests + +Configure with `RABITQ_BUILD_INDEX=OFF`, `RABITQ_BUILD_FASTSCAN=OFF` and +`RABITQ_BUILD_TESTS=ON` to run core references without batch-scoring or index sources. +Enable FastScan with indexes disabled to add the batch-scoring references. The full configuration +continues to run all tests, including index, clustering, persistence and allocation failures. + +The external consumer runner copies its CMake project and the standalone core example +to separate working directories, then tests both core and FastScan with +`add_subdirectory`, local-source `FetchContent`, and relocated installed packages. +It also runs each producer's reference tests (downloading the existing GoogleTest dependency). Pass absolute paths: + +```text +cmake -DRABITQ_SOURCE_DIR= -DWORK_DIR= -P tests/core_consumer/run.cmake +``` + +On macOS, also pass `-DOpenMP_ROOT=`. On Windows, set +`CMAKE_GENERATOR=Visual Studio 18 2026` and `CMAKE_GENERATOR_PLATFORM=x64` in the +environment. The runner builds Release with two jobs and disables native tuning. +It compiles module headers with higher-module headers deliberately made unavailable, +checks that retired targets and disabled component headers are absent, and compiles every +installed first-party header independently. It checks actual compilation +units with generators supporting compilation databases, +and both archive and executable symbols when a non-MSVC `nm` is available. It also +checks that core-only Python, indexes without FastScan, and missing required +FastScan/index components fail configuration. +CI runs this workflow on Linux x86-64, Linux ARM64, macOS ARM64 and Windows x86-64. + +Core and FastScan consumption must also work from a full installation; the standalone +`sample/cpp/core` and `sample/cpp/fastscan` projects test their respective targets. The original +`tests/consumer` exercises `rabitqlib::index` and checks that HNSW headers compile with +the IVF initializer header deliberately unavailable and index value types compile with +Eigen deliberately unavailable. Numerical consumers forbid the removed `defines.hpp`. +All FastScan batch layout/encoding coverage lives under `unit/rabitqlib/fastscan/`; +index ID and nearest-neighbor helper coverage lives under `unit/rabitqlib/index/`. +`CoreModuleBoundary` also rejects format-specific declarations/constants in core and +numerical utilities, complementing include-path and symbol checks. +Frozen persistence fixture tooling +still targets its pinned historical release and is not migrated to current targets. + +## Stable core API conformance and allocations + +Run the external consumer workflow with `-DRABITQ_USE_OPENMP=OFF` to additionally +forbid OpenMP package discovery, poison `omp.h` in consumer include paths and check +for runtime symbols/compiler flags. The default ON workflow is also retained. Both +variants compile and execute the installed standalone core benchmark smoke tests. +CI covers both choices on its existing platform matrix; local results do not imply +that remote jobs have already passed. + +The normal suite includes `CoreApi.*`, the frozen `CoreApiV1Consumer`, and separate +`core_api_no_alloc_dispatch` / `core_api_no_alloc_scalar` processes. The latter force +and count cold-call encoding, layout conversion, query preparation and estimation; +Linux additionally wraps malloc/calloc/realloc/aligned allocation. The scalar process +substitutes CPU predicates so the exact event heap cannot be skipped by SIMD success. + +The external module runner also builds the unchanged v1 consumer for each integration +method. Its stable header is tested with Eigen and record-layout headers deliberately +unavailable. The suite is independent of ANN targets and remains enabled in pure core. +See [the fixture policy](fixtures/core_api_v1/README.md). + ## Test structure - `unit/`: C++ index, quantization, FastScan, and utility tests -- `integration/`: packing compatibility tests - `common/`: shared C++ test utilities - `python/`: Python binding tests and persistence fixtures - `consumer/`: installed CMake package test -CMake discovers `*_test.cpp` under `unit/` and `integration/`. Python and +CMake discovers `*_test.cpp` under `unit/`. Python and consumer tests run separately using the commands above. + +## Coverage ownership and test size + +C++ reference tests own exhaustive bit-width, SIMD, packing, rounding and algorithm +boundary matrices. Python tests own NumPy conversions, argument/default forwarding, +returned arrays, GIL/access guards and representative end-to-end index operations. +Keep a Python construction/search/persistence case for every exported storage width and +both metrics, but do not cross every width with every original dimension and initializer. + +`test_portable_indexes.py` separates supported storage from representative padded-domain +boundaries. IVF mutation/batching/routing matrices use sign-only, extra-bit and raw-vector +storage; full width coverage remains in the corresponding C++ reference tests. The +SymphonyQG batch test checks both Python names on one built index per storage/metric pair, +including strided/empty queries and bitwise distance equality. + +Do not trim frozen compatibility fixtures, independent numerical oracles, failure +injection or concurrency tests just because their inputs resemble another test's inputs. +Preserve public wrapper calls when consolidating lower-level tests; a backend-only test +does not validate wrapper dispatch. Parameter combinations should correspond to a distinct +branch or regression, not multiply unrelated axes by default. Keep the existing CI entry +points and platform-specific checks. + +## Initialization and construction-failure regressions + +`CoreStaticInitialization` calls the stable core API from a namespace-scope +initializer before `main()`. It also runs in the external numerical consumer +workflow. `ModulesStaticInitialization` covers the enabled FastScan/index dispatch +entry points in the same setting. + +`construction_failure_test` isolates allocation replacement from the normal suite. +It checks OpenMP construction errors with one/four requested threads and, on Linux, +wraps C allocations to sweep HNSW constructor failures and verify ownership. Run +these under ASan/UBSan as well as Release; do not disable allocator-mismatch checking. +`compact_codes_failure_test` covers cold scratch allocation and growth failures in +serial and OpenMP encoding, including the public FastScan batch entry points. diff --git a/tests/allocation/core_api_no_alloc.cpp b/tests/allocation/core_api_no_alloc.cpp new file mode 100644 index 00000000..3efa519b --- /dev/null +++ b/tests/allocation/core_api_no_alloc.cpp @@ -0,0 +1,238 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" + +#ifdef _MSC_VER +#include +#define TEST_NOINLINE __declspec(noinline) +#else +#define TEST_NOINLINE __attribute__((noinline)) +#endif + +namespace { +bool watching = false; +size_t new_calls = 0; +size_t c_calls = 0; +volatile float sink = 0; +struct Watch { + Watch() { watching = true; } + ~Watch() { watching = false; } +}; +} // namespace + +#ifdef RABITQ_WRAP_MALLOC +extern "C" { +void* __real_malloc(size_t); +void* __real_calloc(size_t, size_t); +void* __real_realloc(void*, size_t); +void* __real_aligned_alloc(size_t, size_t); +int __real_posix_memalign(void**, size_t, size_t); +void* __wrap_malloc(size_t n) { + if (watching) + ++c_calls; + return __real_malloc(n); +} +void* __wrap_calloc(size_t n, size_t width) { + if (watching) + ++c_calls; + return __real_calloc(n, width); +} +void* __wrap_realloc(void* p, size_t n) { + if (watching) + ++c_calls; + return __real_realloc(p, n); +} +void* __wrap_aligned_alloc(size_t alignment, size_t n) { + if (watching) + ++c_calls; + return __real_aligned_alloc(alignment, n); +} +int __wrap_posix_memalign(void** p, size_t alignment, size_t n) { + if (watching) + ++c_calls; + return __real_posix_memalign(p, alignment, n); +} +} +#endif + +TEST_NOINLINE void* operator new(size_t n) { + if (watching) + ++new_calls; + if (void* p = std::malloc(n == 0 ? 1 : n)) + return p; + throw std::bad_alloc(); +} +TEST_NOINLINE void* operator new[](size_t n) { return ::operator new(n); } +TEST_NOINLINE void operator delete(void* p) noexcept { std::free(p); } +TEST_NOINLINE void operator delete[](void* p) noexcept { ::operator delete(p); } +TEST_NOINLINE void operator delete(void* p, size_t) noexcept { ::operator delete(p); } +TEST_NOINLINE void operator delete[](void* p, size_t) noexcept { ::operator delete(p); } +TEST_NOINLINE void* operator new(size_t n, std::align_val_t alignment) { + if (watching) + ++new_calls; + const auto a = static_cast(alignment); +#ifdef _MSC_VER + if (void* p = _aligned_malloc(n == 0 ? 1 : n, a)) + return p; +#else + if (void* p = std::aligned_alloc(a, ((n == 0 ? 1 : n) + a - 1) / a * a)) + return p; +#endif + throw std::bad_alloc(); +} +TEST_NOINLINE void* operator new[](size_t n, std::align_val_t a) { + return ::operator new(n, a); +} +TEST_NOINLINE void operator delete(void* p, std::align_val_t) noexcept { +#ifdef _MSC_VER + _aligned_free(p); +#else + std::free(p); +#endif +} +TEST_NOINLINE void operator delete[](void* p, std::align_val_t a) noexcept { + ::operator delete(p, a); +} +TEST_NOINLINE void operator delete(void* p, size_t, std::align_val_t a) noexcept { + ::operator delete(p, a); +} +TEST_NOINLINE void operator delete[](void* p, size_t, std::align_val_t a) noexcept { + ::operator delete(p, a); +} + +#ifdef RABITQ_FORCE_SCALAR +// Isolated executable: force the actual cached dispatch to use scalar kernels. +// This guarantees the exact event-heap fallback is exercised, even on AVX-512 hosts. +namespace rabitqlib::cpu { +bool has_avx2() { return false; } +bool has_avx512_core() { return false; } +bool has_avx512_popcnt() { return false; } +bool has_neon() { return false; } +} // namespace rabitqlib::cpu +#endif + +int main() { + namespace api = rabitqlib::quant::v1; + constexpr size_t max_dim = 4096; + alignas(std::max_align_t) std::array memory{}; + std::array data{}, centroid{}, query{}, rotated{}; + std::array raw{}, restored{}; + std::array signs{}; + std::array extra{}; + for (size_t i = 0; i < max_dim; ++i) { + data[i] = static_cast(static_cast(i % 19) - 9) / 11; + centroid[i] = static_cast(static_cast(i % 7) - 3) / 13; + query[i] = static_cast(static_cast(i % 23) - 11) / 17; + } + try { + for (size_t dim : {32U, 96U, 512U, 4096U}) { + for (size_t bits : {1U, 4U, 9U}) { + for (auto metric : {rabitqlib::METRIC_L2, rabitqlib::METRIC_IP}) { + for (double scale : {-1.0, 7.25}) { + const api::Parameters p{dim, bits, metric, scale}; + api::Requirements r; + { + Watch guard; + r = api::requirements(p, api::CodeLayout::SplitPacked); + } + if (r.workspace_bytes > memory.size()) + return 1; + api::Buffer workspace{memory.data(), r.workspace_bytes}; + api::CodeOutput a{ + api::CodeLayout::UnpackedU16, {raw.data(), raw.size()}, {}}; + api::CodeOutput b{ + api::CodeLayout::SplitPacked, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}; + // No warmup: the first valid call and the first call at each size + // count. + { + Watch guard; + auto q = api::prepare_query(p, query.data(), centroid.data()); + auto f = + api::encode(p, data.data(), centroid.data(), a, workspace); + api::convert_codes( + p, {a.layout, {raw.data(), raw.size()}, {}}, b, workspace + ); + sink = api::compute( + q, + {b.layout, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}, + f + ); + api::convert_codes( + p, + {b.layout, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}, + {a.layout, {restored.data(), restored.size()}, {}}, + workspace + ); + f = api::encode(p, data.data(), nullptr, b, workspace); + q = api::prepare_query(p, query.data(), nullptr); + sink = api::compute( + q, + {b.layout, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}, + f + ); + f = api::encode( + p, centroid.data(), centroid.data(), a, workspace + ); + sink = api::compute( + q, {a.layout, {raw.data(), raw.size()}, {}}, f + ); + } + if (!std::isfinite(sink)) + return 2; + } + } + } + } + for (size_t dim : {2U, 3U, 15U, 33U, 65U}) { + Watch guard; + const api::Parameters p{dim, 9}; + const auto r = api::requirements(p, api::CodeLayout::UnpackedU16); + const auto f = api::encode( + p, + data.data(), + nullptr, + {api::CodeLayout::UnpackedU16, {raw.data(), raw.size()}, {}}, + {memory.data(), r.workspace_bytes} + ); + sink = api::compute( + api::prepare_query(p, query.data(), nullptr), + {api::CodeLayout::UnpackedU16, {raw.data(), raw.size()}, {}}, + f + ); + } + // Rotation remains optional. The existing default FHT path is allocation-free + // after constructing its state, including its first rotate call. + std::unique_ptr> rotation( + rabitqlib::choose_rotator( + 65, rabitqlib::RotatorType::FhtKacRotator, 0, 42 + ) + ); + { + Watch guard; + rotation->rotate(data.data(), rotated.data()); + } + } catch (const std::exception& error) { + std::fprintf(stderr, "Unexpected error: %s\n", error.what()); + return 3; + } + std::printf("C++ allocation calls: %zu; C allocation calls: %zu\n", new_calls, c_calls); + return new_calls == 0 && c_calls == 0 ? 0 : 4; +} diff --git a/tests/common/test_data.cpp b/tests/common/test_data.cpp index 40f5c14b..b2f55679 100644 --- a/tests/common/test_data.cpp +++ b/tests/common/test_data.cpp @@ -1,7 +1,9 @@ #include "test_data.hpp" -#include #include +#include +#include +#include namespace rabitq_test { diff --git a/tests/common/test_data.hpp b/tests/common/test_data.hpp index 96a34abe..673e8b8c 100644 --- a/tests/common/test_data.hpp +++ b/tests/common/test_data.hpp @@ -2,7 +2,6 @@ #define RABITQ_TEST_DATA_HPP #include -#include #include namespace rabitq_test { diff --git a/tests/common/test_helpers.hpp b/tests/common/test_helpers.hpp index cb473dac..df442360 100644 --- a/tests/common/test_helpers.hpp +++ b/tests/common/test_helpers.hpp @@ -4,8 +4,7 @@ #include #include -#include -#include +#include namespace rabitq_test { diff --git a/tests/consumer/CMakeLists.txt b/tests/consumer/CMakeLists.txt index a1c3bd57..db8ba4cb 100644 --- a/tests/consumer/CMakeLists.txt +++ b/tests/consumer/CMakeLists.txt @@ -27,10 +27,47 @@ foreach(forbidden IN ITEMS -O3 -march=native -Wall -Wextra) endforeach() add_executable(rabitqlib_consumer main.cpp) -target_link_libraries(rabitqlib_consumer PRIVATE rabitqlib::rabitqlib) +target_link_libraries(rabitqlib_consumer PRIVATE rabitqlib::index) +if(TARGET rabitqlib::rabitqlib OR TARGET rabitq_headers) + message(FATAL_ERROR "Removed full-library compatibility target is still exported") +endif() + +add_executable(index_header_boundary + index_headers_main.cpp hnsw_headers.cpp ivf_headers.cpp symqg_headers.cpp) +target_link_libraries(index_header_boundary PRIVATE rabitqlib::index) +# Ordinary index declarations must not parse encoding, scoring or persistence implementations. +foreach(header IN ITEMS + core/encoder.hpp core/estimator.hpp core/query.hpp core/rotator.hpp + fastscan/encoder.hpp fastscan/estimator.hpp fastscan/query.hpp + index/hnsw/hnsw_quant.hpp index/ivf/initializer.hpp + index/detail/buffer.hpp index/detail/visited_set.hpp + index/detail/visited_set_hash.hpp index/detail/visited_set_epoch.hpp + index/symqg/detail/search_scratch.hpp + index/detail/parallel_for.hpp third/Eigen/Core third/Eigen/Dense + third/hnswlib/hnswlib.h) + get_filename_component(directory "${header}" DIRECTORY) + file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/forbidden/rabitqlib/${directory}") + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/forbidden/rabitqlib/${header}" + "#error Index public header imports an implementation dependency\n") +endforeach() +target_include_directories(index_header_boundary BEFORE PRIVATE + "${CMAKE_CURRENT_BINARY_DIR}/forbidden") +add_library(hnsw_quant_header_boundary OBJECT hnsw_quant_headers.cpp) +target_link_libraries(hnsw_quant_header_boundary PRIVATE rabitqlib::index) +add_library(symqg_scratch_header_boundary OBJECT symqg_scratch_headers.cpp) +target_link_libraries(symqg_scratch_header_boundary PRIVATE rabitqlib::index) + +add_library(index_types_boundary OBJECT index_types.cpp) +target_link_libraries(index_types_boundary PRIVATE rabitqlib::index) +file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/types-forbidden/rabitqlib/third/Eigen") +file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/types-forbidden/rabitqlib/third/Eigen/Dense" + "#error Index identifiers and candidate values must not depend on Eigen\n") +target_include_directories(index_types_boundary BEFORE PRIVATE + "${CMAKE_CURRENT_BINARY_DIR}/types-forbidden") enable_testing() add_test(NAME rabitqlib_consumer COMMAND rabitqlib_consumer) +add_test(NAME index_header_boundary COMMAND index_header_boundary) # Inspect evaluated, transitive usage requirements rather than just the direct # properties of the public interface target. This works with multi-config builds. diff --git a/tests/consumer/check_options.cmake b/tests/consumer/check_options.cmake index 1c2a3d55..5df01dc5 100644 --- a/tests/consumer/check_options.cmake +++ b/tests/consumer/check_options.cmake @@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.20) include("${OPTIONS_FILE}") # Only public-header requirements may add compiler options. Sanitizer builds -# intentionally propagate instrumentation to keep Eigen allocation consistent. +# intentionally propagate instrumentation across the library/consumer boundary. # Catch optimization, architecture and warnings leaking through any level of # the target graph, including installed exports. set(required_options ${openmp_options}) @@ -23,6 +23,11 @@ endforeach() if(NOT "cxx_std_17" IN_LIST consumer_features) message(FATAL_ERROR "The package must propagate its C++17 requirement") endif() +foreach(definition IN ITEMS EIGEN_MAX_ALIGN_BYTES=64 EIGEN_MAX_STATIC_ALIGN_BYTES=64) + if(NOT definition IN_LIST consumer_definitions) + message(FATAL_ERROR "Missing Eigen ABI definition: ${definition}") + endif() +endforeach() if(expect_msvc) foreach(definition IN ITEMS __restrict__=__restrict EIGEN_DONT_VECTORIZE) if(NOT definition IN_LIST consumer_definitions) diff --git a/tests/consumer/hnsw_headers.cpp b/tests/consumer/hnsw_headers.cpp new file mode 100644 index 00000000..411f6e90 --- /dev/null +++ b/tests/consumer/hnsw_headers.cpp @@ -0,0 +1,15 @@ +#include "rabitqlib/index/hnsw/hnsw.hpp" + +int hnsw_header_smoke() { + rabitqlib::hnsw::HierarchicalNSW empty; + rabitqlib::hnsw::HierarchicalNSW index(8, 64, 4, 2, 10, 42); + float data[8 * 64]{}; + float centroid[64]{}; + rabitqlib::PID assignments[8]{}; + index.construct(1, centroid, 8, data, assignments); + return empty.num_points() == 0 && empty.num_stored_points() == 0 && + index.dimension() == 64 && index.num_points() == 8 && + index.num_stored_points() == 8 && index.max_elements() == 8 + ? 0 + : 1; +} diff --git a/tests/consumer/hnsw_quant_headers.cpp b/tests/consumer/hnsw_quant_headers.cpp new file mode 100644 index 00000000..791494cb --- /dev/null +++ b/tests/consumer/hnsw_quant_headers.cpp @@ -0,0 +1,2 @@ +// Explicit advanced query-cache includes remain self-contained. +#include "rabitqlib/index/hnsw/hnsw_quant.hpp" // IWYU pragma: keep diff --git a/tests/consumer/index_headers_main.cpp b/tests/consumer/index_headers_main.cpp new file mode 100644 index 00000000..545285c9 --- /dev/null +++ b/tests/consumer/index_headers_main.cpp @@ -0,0 +1,5 @@ +int hnsw_header_smoke(); +int ivf_header_smoke(); +int symqg_header_smoke(); + +int main() { return hnsw_header_smoke() + ivf_header_smoke() + symqg_header_smoke(); } diff --git a/tests/consumer/index_types.cpp b/tests/consumer/index_types.cpp new file mode 100644 index 00000000..e829db38 --- /dev/null +++ b/tests/consumer/index_types.cpp @@ -0,0 +1,8 @@ +#include +#include + +#include "rabitqlib/index/types.hpp" + +static_assert(std::is_same_v); +static_assert(rabitqlib::kPidMax == UINT32_MAX); +static_assert(std::is_trivially_copyable_v>); diff --git a/tests/consumer/ivf_headers.cpp b/tests/consumer/ivf_headers.cpp new file mode 100644 index 00000000..3c3fd4e6 --- /dev/null +++ b/tests/consumer/ivf_headers.cpp @@ -0,0 +1,27 @@ +#include + +#include "rabitqlib/index/ivf/ivf.hpp" + +static_assert(std::is_nothrow_default_constructible_v); + +int ivf_header_smoke() { + rabitqlib::ivf::IVF empty; + rabitqlib::ivf::IVF index( + 8, + 64, + 2, + 4, + rabitqlib::METRIC_L2, + rabitqlib::RotatorType::FhtKacRotator, + rabitqlib::ivf::InitializerType::Flat + ); + float data[8 * 64]{}; + float centroids[2 * 64]{}; + rabitqlib::PID assignments[8]{}; + index.construct(data, centroids, assignments); + return empty.max_elements() == 0 && empty.num_points() == 0 && + empty.num_stored_points() == 0 && index.dimension() == 64 && + index.num_points() == 8 && index.num_stored_points() == 8 + ? 0 + : 1; +} diff --git a/tests/consumer/main.cpp b/tests/consumer/main.cpp index a05cd225..13b8696d 100644 --- a/tests/consumer/main.cpp +++ b/tests/consumer/main.cpp @@ -5,13 +5,21 @@ #include #include +#include "rabitqlib/clustering/detail/kmeans.hpp" #include "rabitqlib/clustering/qgkmeans.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +static_assert(EIGEN_DEFAULT_ALIGN_BYTES == 64); +static_assert(EIGEN_MAX_STATIC_ALIGN_BYTES == 64); + namespace { constexpr size_t kDimension = 65; constexpr size_t kPoints = 67; @@ -93,7 +101,7 @@ int main() { return 1; } - // Exercise both header-defined IVF and compiled SymphonyQG with both + // Exercise the compiled IVF and SymphonyQG implementations with both // rotators, including Eigen-backed ownership across the library boundary. for (const auto rotator : {rabitqlib::RotatorType::FhtKacRotator, rabitqlib::RotatorType::MatrixRotator}) { diff --git a/tests/consumer/symqg_headers.cpp b/tests/consumer/symqg_headers.cpp new file mode 100644 index 00000000..5241012f --- /dev/null +++ b/tests/consumer/symqg_headers.cpp @@ -0,0 +1,17 @@ +#include + +#include "rabitqlib/index/symqg/qg.hpp" + +int symqg_header_smoke() { + rabitqlib::symqg::QuantizedGraph empty; + rabitqlib::symqg::QuantizedGraph graph( + 64, 65, 32, rabitqlib::METRIC_L2, rabitqlib::RotatorType::FhtKacRotator, 4, 42 + ); + auto moved = std::move(graph); + empty = std::move(moved); + return empty.dimension() == 65 && empty.padded_dim() == 96 && + empty.max_elements() == 64 && empty.num_points() == 0 && + empty.num_stored_points() == 0 && empty.quantization_bits() == 4 + ? 0 + : 1; +} diff --git a/tests/consumer/symqg_scratch_headers.cpp b/tests/consumer/symqg_scratch_headers.cpp new file mode 100644 index 00000000..8a2fd511 --- /dev/null +++ b/tests/consumer/symqg_scratch_headers.cpp @@ -0,0 +1,33 @@ +#include + +#include "rabitqlib/index/detail/visited_set.hpp" +#include "rabitqlib/index/symqg/detail/search_scratch.hpp" +#include "rabitqlib/index/symqg/qg.hpp" + +void symqg_scratch_header_smoke( + rabitqlib::symqg::QuantizedGraph& graph, + const float* query, + uint32_t* ids, + float* distances +) { + rabitqlib::VisitedSet visited(graph.num_vertices(), graph.num_vertices()); + rabitqlib::symqg::detail::SearchScratch scratch( + graph.padded_dim(), graph.degree_bound(), 10, 1 + ); + graph.search_with_scratch( + query, + 1, + ids, + distances, + scratch.rotated_query.data(), + scratch.estimated_distances.data(), + scratch.lookup_table.data(), + scratch.batch_query, + scratch.search_pool, + scratch.result_pool, + visited + ); + rabitqlib::symqg::detail::search_with_workspace( + graph, query, 1, ids, distances, scratch.view(visited) + ); +} diff --git a/tests/core_consumer/CMakeLists.txt b/tests/core_consumer/CMakeLists.txt new file mode 100644 index 00000000..27fabd6c --- /dev/null +++ b/tests/core_consumer/CMakeLists.txt @@ -0,0 +1,132 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitqlib_core_consumer LANGUAGES CXX) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +set(RABITQ_COMPONENT "core" CACHE STRING "core or fastscan") +if(NOT RABITQ_COMPONENT MATCHES "^(core|fastscan)$") + message(FATAL_ERROR "Unknown module: ${RABITQ_COMPONENT}") +endif() + +set(INTEGRATION "PACKAGE" CACHE STRING "PACKAGE, ADD_SUBDIRECTORY, or FETCHCONTENT") +if(INTEGRATION STREQUAL "PACKAGE") + find_package(rabitqlib CONFIG REQUIRED COMPONENTS ${RABITQ_COMPONENT}) + find_path(component_include rabitqlib/core/api.hpp + PATHS "${CMAKE_PREFIX_PATH}/include" NO_DEFAULT_PATH REQUIRED) + set(forbidden_paths defines.hpp core/pack_codes.hpp quantization simd index clustering third/hnswlib utils/rotator.hpp + utils/buffer.hpp utils/hashset.hpp utils/visited_pool.hpp utils/visited_set.hpp + utils/visited_set_epoch.hpp utils/visited_set_hash.hpp) + if(RABITQ_COMPONENT STREQUAL "core") + list(APPEND forbidden_paths fastscan) + endif() + foreach(path IN LISTS forbidden_paths) + if(EXISTS "${component_include}/rabitqlib/${path}") + message(FATAL_ERROR "Disabled or retired header installed: ${path}") + endif() + endforeach() + # Compile every installed first-party header on its own, without source-tree + # include paths. This verifies the component's complete template dependency closure. + file(GLOB_RECURSE installed_headers RELATIVE "${component_include}" + "${component_include}/rabitqlib/*.hpp") + list(FILTER installed_headers EXCLUDE REGEX "/third/") + set(header_sources) + foreach(header IN LISTS installed_headers) + string(MAKE_C_IDENTIFIER "${header}" identifier) + set(source "${CMAKE_CURRENT_BINARY_DIR}/headers/${identifier}.cpp") + file(WRITE "${source}" "#include <${header}>\n") + list(APPEND header_sources "${source}") + endforeach() + add_library(component_header_check OBJECT ${header_sources}) + target_link_libraries(component_header_check PRIVATE rabitqlib::${RABITQ_COMPONENT}) +else() + # FetchContent writes this path into a generated CMake script. + file(TO_CMAKE_PATH "${RABITQ_SOURCE_DIR}" RABITQ_SOURCE_DIR) + set(RABITQ_BUILD_INDEX OFF CACHE BOOL "" FORCE) + set(RABITQ_BUILD_FASTSCAN OFF CACHE BOOL "" FORCE) + if(RABITQ_COMPONENT STREQUAL "fastscan") + set(RABITQ_BUILD_FASTSCAN ON CACHE BOOL "" FORCE) + endif() + set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) + set(RABITQ_BUILD_TESTS OFF CACHE BOOL "" FORCE) + set(RABITQ_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE) + set(RABITQ_ENABLE_NATIVE_OPTIMIZATION OFF CACHE BOOL "" FORCE) + if(INTEGRATION STREQUAL "ADD_SUBDIRECTORY") + add_subdirectory("${RABITQ_SOURCE_DIR}" rabitqlib) + elseif(INTEGRATION STREQUAL "FETCHCONTENT") + include(FetchContent) + # Test the current checkout without fetching an unrelated published revision. + FetchContent_Declare(rabitqlib SOURCE_DIR "${RABITQ_SOURCE_DIR}") + FetchContent_MakeAvailable(rabitqlib) + else() + message(FATAL_ERROR "Unknown integration: ${INTEGRATION}") + endif() + if(TARGET rabitqlib::index OR TARGET rabitqlib::rabitqlib) + message(FATAL_ERROR "Core-only configuration exposes an index/full target") + endif() +endif() + +if(TARGET rabitqlib::rabitqlib OR TARGET rabitq_headers) + message(FATAL_ERROR "Removed full-library compatibility target is still available") +endif() + +get_target_property(core_links rabitqlib::core INTERFACE_LINK_LIBRARIES) +if(core_links MATCHES "fastscan|index|headers|::rabitqlib") + message(FATAL_ERROR "Core has an ANN dependency: ${core_links}") +endif() + +add_executable(core_consumer main.cpp) +target_link_libraries(core_consumer PRIVATE rabitqlib::${RABITQ_COMPONENT}) +# Canonical module headers must not transitively include a higher module. Poison +# those include paths for this executable. +set(restricted_modules index clustering) +if(RABITQ_COMPONENT STREQUAL "core") + list(APPEND restricted_modules fastscan) +endif() +foreach(module IN LISTS restricted_modules) + file(GLOB_RECURSE restricted_headers RELATIVE "${RABITQ_SOURCE_DIR}/include" + "${RABITQ_SOURCE_DIR}/include/rabitqlib/${module}/*.hpp") + foreach(header IN LISTS restricted_headers) + get_filename_component(directory "${header}" DIRECTORY) + file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/forbidden/${directory}") + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/forbidden/${header}" + "#error Unexpected higher-module header dependency\n") + endforeach() +endforeach() +file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/forbidden/rabitqlib/defines.hpp" + "#error Removed mixed-responsibility header must not be included\n") +target_include_directories(core_consumer BEFORE PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/forbidden") +add_executable(core_api_v1_consumer consumer.cpp) +target_link_libraries(core_api_v1_consumer PRIVATE rabitqlib::core) +if(DEFINED RABITQ_USE_OPENMP AND NOT RABITQ_USE_OPENMP) + file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/no-openmp") + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/no-openmp/omp.h" + "#error OpenMP-free consumer must not require omp.h\n") + foreach(consumer IN ITEMS core_consumer core_api_v1_consumer) + target_include_directories(${consumer} BEFORE PRIVATE + "${CMAKE_CURRENT_BINARY_DIR}/no-openmp") + endforeach() +endif() +# The stable header must be self-contained without exposing Eigen or record maps. +file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/api-forbidden/rabitqlib/third/Eigen") +file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/api-forbidden/rabitqlib/third/Eigen/Dense" + "#error Stable API must not include Eigen\n") +file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/api-forbidden/rabitqlib/core") +file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/api-forbidden/rabitqlib/core/data_layout.hpp" + "#error Stable API must not include record layouts\n") +target_include_directories(core_api_v1_consumer BEFORE PRIVATE + "${CMAKE_CURRENT_BINARY_DIR}/api-forbidden" "${CMAKE_CURRENT_BINARY_DIR}/forbidden") + +enable_testing() +add_test(NAME core_api_v1_consumer COMMAND core_api_v1_consumer) +add_test(NAME core_consumer COMMAND core_consumer) + +if(CMAKE_NM AND NOT MSVC) + add_test(NAME core_symbols COMMAND "${CMAKE_COMMAND}" + "-DNM=${CMAKE_NM}" "-DCORE=$" + "-DCONSUMER=$" "-DCOMPONENT=${RABITQ_COMPONENT}" + "-DUSE_OPENMP=${RABITQ_USE_OPENMP}" + -P "${CMAKE_CURRENT_SOURCE_DIR}/check_symbols.cmake") +endif() + +add_executable(core_static_init core_static_init.cpp) +target_link_libraries(core_static_init PRIVATE rabitqlib::core) +add_test(NAME CoreStaticInitialization COMMAND core_static_init) diff --git a/tests/core_consumer/check_core_boundary.cmake b/tests/core_consumer/check_core_boundary.cmake new file mode 100644 index 00000000..a09a7730 --- /dev/null +++ b/tests/core_consumer/check_core_boundary.cmake @@ -0,0 +1,22 @@ +cmake_minimum_required(VERSION 3.20) +if(NOT RABITQ_SOURCE_DIR) + message(FATAL_ERROR "RABITQ_SOURCE_DIR is required") +endif() + +# Header poisoning alone cannot detect a FastScan declaration copied into core. +# Guard the format-specific vocabulary as well as direct upward includes. +file(GLOB_RECURSE core_files + "${RABITQ_SOURCE_DIR}/include/rabitqlib/core/*.hpp" + "${RABITQ_SOURCE_DIR}/include/rabitqlib/utils/*.hpp" + "${RABITQ_SOURCE_DIR}/src/core/*.cpp" + "${RABITQ_SOURCE_DIR}/src/simd/*.cpp" + "${RABITQ_SOURCE_DIR}/src/simd/*.hpp") +foreach(path IN LISTS core_files) + file(READ "${path}" contents) + if(contents MATCHES "fastscan::|BatchDataMap|kBatchSize|kPerm0|one_bit_batch_code|quantize_(one|split|qg)_batch") + message(FATAL_ERROR "FastScan batch-format dependency in core: ${path}") + endif() + if(contents MATCHES "# *include[^\n]*rabitqlib/(fastscan|index|clustering)/") + message(FATAL_ERROR "Upward module include in core: ${path}") + endif() +endforeach() diff --git a/tests/core_consumer/check_no_openmp.cmake b/tests/core_consumer/check_no_openmp.cmake new file mode 100644 index 00000000..2a6498c4 --- /dev/null +++ b/tests/core_consumer/check_no_openmp.cmake @@ -0,0 +1,18 @@ +cmake_minimum_required(VERSION 3.20) + +# CMake's compilation database stores each invocation in the command field. +# Inspect parsed arguments, not paths or other JSON fields containing /openmp. +file(READ "${COMPILE_COMMANDS}" database) +string(JSON count LENGTH "${database}") +if(count GREATER 0) + math(EXPR last "${count} - 1") + foreach(index RANGE ${last}) + string(JSON command GET "${database}" ${index} command) + separate_arguments(arguments NATIVE_COMMAND "${command}") + foreach(argument IN LISTS arguments) + if(argument MATCHES "^-fopenmp($|[=-])" OR argument MATCHES "^/openmp($|:)") + message(FATAL_ERROR "OpenMP compiler flag in disabled build: ${argument}") + endif() + endforeach() + endforeach() +endif() diff --git a/tests/core_consumer/check_no_openmp_test.cmake b/tests/core_consumer/check_no_openmp_test.cmake new file mode 100644 index 00000000..6ea233d2 --- /dev/null +++ b/tests/core_consumer/check_no_openmp_test.cmake @@ -0,0 +1,36 @@ +cmake_minimum_required(VERSION 3.20) + +file(MAKE_DIRECTORY "${WORK_DIR}") +set(database_path "${WORK_DIR}/compile_commands.json") +set(database_template [=[ +[ + {"directory": "/openmp", "file": "/tmp/-fopenmp/input.cpp", + "command": "c++ -I\"/tmp/openmp-off/include dir\" -c \"/tmp/-fopenmp/input.cpp\""}, + {"directory": "/tmp/openmp-off", "file": "input.cpp", "output": "/openmp", + "command": "c++ -DNAME=/openmp -I/tmp/-fopenmp @option@ -c input.cpp"} +] +]=]) + +function(check_case option expect_failure) + string(CONFIGURE "${database_template}" database @ONLY) + file(WRITE "${database_path}" "${database}") + execute_process(COMMAND "${CMAKE_COMMAND}" + "-DCOMPILE_COMMANDS=${database_path}" + -P "${CMAKE_CURRENT_LIST_DIR}/check_no_openmp.cmake" + RESULT_VARIABLE result OUTPUT_VARIABLE output ERROR_VARIABLE error) + if(expect_failure) + if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "OpenMP compiler flag in disabled build") + message(FATAL_ERROR "Failed to reject ${option}: ${output}${error}") + endif() + elseif(NOT result EQUAL 0) + message(FATAL_ERROR "Rejected paths without OpenMP flags: ${output}${error}") + endif() +endfunction() + +check_case("" FALSE) +foreach(flag IN ITEMS -fopenmp -fopenmp=libomp -fopenmp-simd /openmp /openmp:llvm /openmp:experimental) + check_case("${flag}" TRUE) + # JSON escapes preserve quotes around a single command-line argument. + check_case("\\\"${flag}\\\"" TRUE) +endforeach() +message(STATUS "OpenMP flag detection accepts paths and rejects actual compiler flags") diff --git a/tests/core_consumer/check_symbols.cmake b/tests/core_consumer/check_symbols.cmake new file mode 100644 index 00000000..8ce25310 --- /dev/null +++ b/tests/core_consumer/check_symbols.cmake @@ -0,0 +1,19 @@ +foreach(binary IN ITEMS "${CORE}" "${CONSUMER}") + execute_process(COMMAND "${NM}" -C "${binary}" RESULT_VARIABLE result + OUTPUT_VARIABLE symbols ERROR_VARIABLE error) + if(NOT result EQUAL 0) + message(FATAL_ERROR "Cannot inspect ${binary}: ${error}") + endif() + if(DEFINED USE_OPENMP AND NOT USE_OPENMP AND symbols MATCHES "[ \t](GOMP_|omp_|__kmpc_)[A-Za-z_]") + message(FATAL_ERROR "OpenMP runtime symbol in disabled build: ${binary}") + endif() + # Check definitions AND unresolved references in the archive, plus the linked consumer. + if(symbols MATCHES "rabitqlib::(ivf|hnsw|symqg|qgkmeans|rabitqkmeans)::|hnswlib::") + message(FATAL_ERROR "ANN symbols in ${binary}") + endif() + if(binary STREQUAL CORE OR COMPONENT STREQUAL "core") + if(symbols MATCHES "rabitqlib::fastscan::(simd::)?(accumulate|pack_lut|transfer_lut)|rabitqlib::simd::(qg_batch_estdist|split_batch_estdist)") + message(FATAL_ERROR "FastScan scoring symbols in core: ${binary}") + endif() + endif() +endforeach() diff --git a/tests/core_consumer/core_static_init.cpp b/tests/core_consumer/core_static_init.cpp new file mode 100644 index 00000000..6084c0bd --- /dev/null +++ b/tests/core_consumer/core_static_init.cpp @@ -0,0 +1,58 @@ +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" + +namespace { +// This translation unit is linked before the core archive: this must also work +// before any library translation unit has performed dynamic initialization. +const bool kBeforeMain = [] { + namespace api = rabitqlib::quant::v1; + std::array input{}, rotated{}, query{}; + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(i % 11) - 5.0F; + } + rabitqlib::rotator_impl::FhtKacRotator rotator(64, 64, 42); + rotator.rotate(input.data(), rotated.data()); + query = rotated; + for (const auto metric : {rabitqlib::METRIC_L2, rabitqlib::METRIC_IP}) { + for (size_t bits = 1; bits <= 9; ++bits) { + const api::Parameters params{64, bits, metric}; + const auto req = api::requirements(params, api::CodeLayout::SplitPacked); + std::vector codes(req.code_bytes), extra(req.extra_bytes); + std::vector work( + (req.workspace_bytes + sizeof(std::max_align_t) - 1) / + sizeof(std::max_align_t) + ); + const auto factors = api::encode( + params, + rotated.data(), + nullptr, + {api::CodeLayout::SplitPacked, + {codes.data(), codes.size()}, + {extra.data(), extra.size()}}, + {work.data(), req.workspace_bytes} + ); + const auto prepared = api::prepare_query(params, query.data(), nullptr); + const float distance = api::compute( + prepared, + {api::CodeLayout::SplitPacked, + {codes.data(), codes.size()}, + {extra.data(), extra.size()}}, + factors + ); + if (!std::isfinite(distance)) { + return false; + } + } + } + return true; +}(); +} // namespace + +int main() { return kBeforeMain ? 0 : 1; } diff --git a/tests/core_consumer/run.cmake b/tests/core_consumer/run.cmake new file mode 100644 index 00000000..15b73a28 --- /dev/null +++ b/tests/core_consumer/run.cmake @@ -0,0 +1,141 @@ +# Run with -DRABITQ_SOURCE_DIR= -DWORK_DIR= -P run.cmake. +cmake_minimum_required(VERSION 3.20) +if(NOT RABITQ_SOURCE_DIR OR NOT WORK_DIR) + message(FATAL_ERROR "RABITQ_SOURCE_DIR and WORK_DIR are required") +endif() + +function(run) + execute_process(COMMAND ${ARGV} RESULT_VARIABLE result) + if(NOT result EQUAL 0) + message(FATAL_ERROR "Command failed (${result}): ${ARGV}") + endif() +endfunction() + +set(configure_options -DCMAKE_BUILD_TYPE=Release -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF) +if(NOT DEFINED RABITQ_USE_OPENMP) + set(RABITQ_USE_OPENMP ON) +endif() +list(APPEND configure_options "-DRABITQ_USE_OPENMP=${RABITQ_USE_OPENMP}") +if(NOT RABITQ_USE_OPENMP) + list(APPEND configure_options -DCMAKE_DISABLE_FIND_PACKAGE_OpenMP=TRUE) +endif() +if(DEFINED FETCHCONTENT_SOURCE_DIR_GOOGLETEST) + list(APPEND configure_options + "-DFETCHCONTENT_SOURCE_DIR_GOOGLETEST=${FETCHCONTENT_SOURCE_DIR_GOOGLETEST}") +endif() +if(DEFINED OpenMP_ROOT) + list(APPEND configure_options "-DOpenMP_ROOT=${OpenMP_ROOT}") +endif() +foreach(component IN ITEMS core fastscan) + set(scan OFF) + if(component STREQUAL "fastscan") + set(scan ON) + endif() + set(work "${WORK_DIR}/${component}") + file(MAKE_DIRECTORY "${work}/consumer") + file(COPY "${RABITQ_SOURCE_DIR}/tests/core_consumer/CMakeLists.txt" + "${RABITQ_SOURCE_DIR}/tests/core_consumer/check_symbols.cmake" + "${RABITQ_SOURCE_DIR}/tests/core_consumer/core_static_init.cpp" + "${RABITQ_SOURCE_DIR}/tests/fixtures/core_api_v1/consumer.cpp" + "${RABITQ_SOURCE_DIR}/sample/cpp/${component}/main.cpp" + DESTINATION "${work}/consumer") + run("${CMAKE_COMMAND}" -S "${RABITQ_SOURCE_DIR}" -B "${work}/producer" + ${configure_options} -DRABITQ_BUILD_INDEX=OFF + "-DRABITQ_BUILD_FASTSCAN=${scan}" -DRABITQ_BUILD_TESTS=ON) + run("${CMAKE_COMMAND}" --build "${work}/producer" --config Release --parallel 2) + run("${CMAKE_CTEST_COMMAND}" --test-dir "${work}/producer" -C Release + --output-on-failure --timeout 120) + # Relocate the initial install; subsequent runs update that same test prefix. + if(EXISTS "${work}/relocated") + run("${CMAKE_COMMAND}" --install "${work}/producer" --config Release + --prefix "${work}/relocated") + else() + run("${CMAKE_COMMAND}" --install "${work}/producer" --config Release + --prefix "${work}/install") + file(RENAME "${work}/install" "${work}/relocated") + endif() + foreach(method IN ITEMS ADD_SUBDIRECTORY FETCHCONTENT PACKAGE) + set(build "${work}/${method}") + set(consumer_source "${RABITQ_SOURCE_DIR}") + if(method STREQUAL "FETCHCONTENT") + # Exercise native Windows path spelling on every host. + string(REPLACE "/" "\\" consumer_source "${consumer_source}") + endif() + run("${CMAKE_COMMAND}" -S "${work}/consumer" -B "${build}" + ${configure_options} "-DINTEGRATION=${method}" + "-DRABITQ_COMPONENT=${component}" + "-DRABITQ_SOURCE_DIR=${consumer_source}" + "-DCMAKE_PREFIX_PATH=${work}/relocated") + run("${CMAKE_COMMAND}" --build "${build}" --config Release --parallel 2) + run("${CMAKE_CTEST_COMMAND}" --test-dir "${build}" -C Release --output-on-failure) + endforeach() + foreach(build IN ITEMS producer ADD_SUBDIRECTORY FETCHCONTENT) + if(EXISTS "${work}/${build}/compile_commands.json") + file(READ "${work}/${build}/compile_commands.json" commands) + if(NOT RABITQ_USE_OPENMP) + run("${CMAKE_COMMAND}" + "-DCOMPILE_COMMANDS=${work}/${build}/compile_commands.json" + -P "${CMAKE_CURRENT_LIST_DIR}/check_no_openmp.cmake") + endif() + if(commands MATCHES "src/index/|python_bindings/|clustering/") + message(FATAL_ERROR "ANN/Python source compiled by ${component}/${build}") + endif() + if(component STREQUAL "core" AND commands MATCHES "src/fastscan/|src/simd/(fastscan|estimator)_") + message(FATAL_ERROR "FastScan source compiled by core/${build}") + endif() + endif() + endforeach() +endforeach() + +# Configuration must fail early for combinations that cannot provide full functionality. +execute_process(COMMAND "${CMAKE_COMMAND}" -S "${RABITQ_SOURCE_DIR}" + -B "${WORK_DIR}/invalid-python" -DRABITQ_BUILD_INDEX=OFF + -DRABITQ_BUILD_PYTHON_BINDINGS=ON RESULT_VARIABLE result + OUTPUT_VARIABLE output ERROR_VARIABLE error) +if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "requires RABITQ_BUILD_INDEX=ON") + message(FATAL_ERROR "Core-only Python configuration was not rejected clearly") +endif() + +file(MAKE_DIRECTORY "${WORK_DIR}/require-index") +file(WRITE "${WORK_DIR}/require-index/CMakeLists.txt" + "cmake_minimum_required(VERSION 3.20)\nproject(require_index LANGUAGES CXX)\nfind_package(rabitqlib CONFIG REQUIRED COMPONENTS index)\n") +execute_process(COMMAND "${CMAKE_COMMAND}" -S "${WORK_DIR}/require-index" + -B "${WORK_DIR}/missing-index" ${configure_options} + "-DCMAKE_PREFIX_PATH=${WORK_DIR}/core/relocated" RESULT_VARIABLE result + OUTPUT_VARIABLE output ERROR_VARIABLE error) +if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "rabitqlib") + message(FATAL_ERROR "Core-only installation accepted the missing index component") +endif() + +execute_process(COMMAND "${CMAKE_COMMAND}" -S "${RABITQ_SOURCE_DIR}" + -B "${WORK_DIR}/invalid-index" -DRABITQ_BUILD_FASTSCAN=OFF + -DRABITQ_BUILD_INDEX=ON RESULT_VARIABLE result + OUTPUT_VARIABLE output ERROR_VARIABLE error) +if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "requires RABITQ_BUILD_FASTSCAN=ON") + message(FATAL_ERROR "Index without FastScan was not rejected clearly") +endif() +if(NOT RABITQ_USE_OPENMP) + execute_process(COMMAND "${CMAKE_COMMAND}" -S "${RABITQ_SOURCE_DIR}" + -B "${WORK_DIR}/invalid-openmp" -DRABITQ_BUILD_INDEX=ON + -DRABITQ_BUILD_FASTSCAN=ON -DRABITQ_USE_OPENMP=OFF RESULT_VARIABLE result + OUTPUT_VARIABLE output ERROR_VARIABLE error) + if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "ANN indexes require RABITQ_USE_OPENMP=ON") + message(FATAL_ERROR "Unsupported OpenMP-free ANN build was not rejected clearly") + endif() +endif() +file(WRITE "${WORK_DIR}/require-index/CMakeLists.txt" + "cmake_minimum_required(VERSION 3.20)\nproject(require_fastscan LANGUAGES CXX)\nfind_package(rabitqlib CONFIG REQUIRED COMPONENTS fastscan)\n") +execute_process(COMMAND "${CMAKE_COMMAND}" -S "${WORK_DIR}/require-index" + -B "${WORK_DIR}/missing-fastscan" ${configure_options} + "-DCMAKE_PREFIX_PATH=${WORK_DIR}/core/relocated" RESULT_VARIABLE result + OUTPUT_VARIABLE output ERROR_VARIABLE error) +if(result EQUAL 0 OR NOT "${output}${error}" MATCHES "rabitqlib") + message(FATAL_ERROR "Core-only installation accepted the missing FastScan component") +endif() +run("${CMAKE_COMMAND}" -S "${RABITQ_SOURCE_DIR}/benchmarks/core" + -B "${WORK_DIR}/benchmark" ${configure_options} + "-DCMAKE_PREFIX_PATH=${WORK_DIR}/core/relocated") +run("${CMAKE_COMMAND}" --build "${WORK_DIR}/benchmark" --config Release --parallel 2) +run("${CMAKE_CTEST_COMMAND}" --test-dir "${WORK_DIR}/benchmark" -C Release + --output-on-failure) +message(STATUS "Core and FastScan external consumers and core benchmark passed") diff --git a/tests/dispatch/fastscan_backend_test.cpp b/tests/dispatch/fastscan_backend_test.cpp new file mode 100644 index 00000000..33278018 --- /dev/null +++ b/tests/dispatch/fastscan_backend_test.cpp @@ -0,0 +1,171 @@ +#include "fastscan_backend.hpp" + +#include + +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::fastscan { +namespace { + +TEST(FastScanBackendContract, SelectedHaccLayoutMatchesProducerAndConsumer) { + const auto& backend = detail::selected_backend(); + const size_t tables_per_plane = + backend.hacc_layout == detail::HaccLutLayout::Quad64 ? 4 : 2; + EXPECT_EQ(&backend, &detail::selected_backend()); + for (size_t dim : {16U, 32U, 64U, 128U, 1024U, 1040U}) { + SCOPED_TRACE(::testing::Message() << backend.name << " dim=" << dim); + std::vector lut(dim * 4); + std::vector bytes(dim * 8 + 2, 0xAB); + std::vector expected = bytes; + for (size_t i = 0; i < lut.size(); ++i) { + lut[i] = static_cast(i * 911 + i / 16 * 37); + const size_t table = i / 16; + const size_t low = table / tables_per_plane * tables_per_plane * 32 + + table % tables_per_plane * 16 + i % 16; + expected[1 + low] = static_cast(lut[i]); + expected[1 + low + tables_per_plane * 16] = static_cast(lut[i] >> 8); + } + transfer_lut_hacc(lut.data(), dim, bytes.data() + 1); + EXPECT_EQ(bytes, expected); + backend.transfer_hacc(lut.data(), dim, bytes.data() + 1); + EXPECT_EQ(bytes, expected); + + // Construct codes independently of pack_codes and its lane permutation. + std::vector codes(dim * 4); + std::array reference{}; + for (size_t group = 0; group < dim / 4; ++group) { + for (size_t row = 0; row < 32; ++row) { + const size_t entry = (group * 3 + row * 7) % 16; + const size_t lane = 2 * (row % 8) + (row % 16) / 8; + codes[group * 16 + lane] |= static_cast(entry << (row / 16 * 4)); + reference[row] += lut[group * 16 + entry]; + } + } + std::array result{}; + accumulate_hacc(codes.data(), bytes.data() + 1, result.data(), dim); + EXPECT_EQ(result, reference); + backend.accumulate_hacc(codes.data(), bytes.data() + 1, result.data(), dim); + EXPECT_EQ(result, reference); + } +} + +TEST(FastScanBackendContract, StandardLutIsSharedAcrossSupportedConsumers) { + constexpr size_t kDim = 64; + const auto& backend = detail::selected_backend(); + std::array query{}; + for (size_t i = 0; i < kDim; ++i) { + query[i] = static_cast(i % 7); + } + std::array packed{}, expected{}; + backend.pack_lut(kDim, query.data(), packed.data()); + simd::pack_lut_generic(kDim, query.data(), expected.data()); + EXPECT_EQ(packed, expected); + std::array lut{}, codes{}; + for (size_t i = 0; i < lut.size(); ++i) { + lut[i] = static_cast(packed[i]); + codes[i] = static_cast(i * 73 + i / 7); + } + std::array reference{}, result{}; + simd::accumulate_generic(codes.data(), lut.data(), reference.data(), kDim); + if (cpu::has_avx2() || cpu::has_neon()) { + backend.accumulate(codes.data(), lut.data(), result.data(), kDim); + EXPECT_EQ(result, reference); + } +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + simd::accumulate_avx2(codes.data(), lut.data(), result.data(), kDim); + EXPECT_EQ(result, reference); + } + if (cpu::has_avx512_core()) { + simd::accumulate_avx512(codes.data(), lut.data(), result.data(), kDim); + EXPECT_EQ(result, reference); + } +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + simd::accumulate_neon(codes.data(), lut.data(), result.data(), kDim); + EXPECT_EQ(result, reference); + } +#endif +} + +TEST(FastScanBackendContract, ExplicitCorrectionBackendsConsumeHostLut) { + const auto& selected = detail::selected_backend(); + using Estimator = decltype(selected.split); + std::vector backends{selected.split}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) + backends.push_back(rabitqlib::simd::split_batch_estdist_avx2); + if (cpu::has_avx512_core()) + backends.push_back(rabitqlib::simd::split_batch_estdist_avx512); +#endif + for (size_t dim : {32U, 96U, 1040U}) { + std::vector values(dim); + for (size_t i = 0; i < dim; ++i) + values[i] = static_cast(i % 23) / 7 - 1; + for (bool hacc : {false, true}) { + if (!hacc && !cpu::has_avx2() && !cpu::has_neon()) + continue; + for (auto metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE( + ::testing::Message() + << selected.name << " " << dim << " " << hacc << " " << metric + ); + SplitBatchQuery query(values.data(), dim, 3, metric, hacc); + query.set_g_add(3, 2); + std::vector storage(BatchDataMap::data_bytes(dim)); + BatchDataMap batch(storage.data(), dim); + for (size_t i = 0; i < dim * 4; ++i) + batch.bin_code()[i] = static_cast(i * 73 + i / 7); + for (size_t lane = 0; lane < 32; ++lane) { + batch.f_add()[lane] = static_cast(lane) + 10; + batch.f_rescale()[lane] = lane % 2 == 0 ? 0.125F : -0.25F; + batch.f_error()[lane] = static_cast(lane) / 32; + } + std::array reference_distance{}, reference_lower{}, + reference_ip{}; + rabitqlib::simd::split_batch_estdist_generic( + storage.data(), + query, + dim, + reference_distance.data(), + reference_lower.data(), + reference_ip.data(), + hacc + ); + for (auto score : backends) { + std::array distance{}, lower{}, ip{}; + score( + storage.data(), + query, + dim, + distance.data(), + lower.data(), + ip.data(), + hacc + ); + for (size_t lane = 0; lane < 32; ++lane) { + const double tolerance = 2e-6 * static_cast(dim); + EXPECT_NEAR(ip[lane], reference_ip[lane], tolerance); + EXPECT_NEAR(distance[lane], reference_distance[lane], tolerance); + EXPECT_NEAR(lower[lane], reference_lower[lane], tolerance); + } + } + } + } + } +} + +} // namespace +} // namespace rabitqlib::fastscan diff --git a/tests/dispatch/fastscan_no_simd_test.cpp b/tests/dispatch/fastscan_no_simd_test.cpp index 407d0c79..cce0a827 100644 --- a/tests/dispatch/fastscan_no_simd_test.cpp +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -1,12 +1,15 @@ #include #include +#include #include #include #include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" // Link-time substitutes, isolated from both production detection and the main suite. diff --git a/tests/dispatch/modules_static_init.cpp b/tests/dispatch/modules_static_init.cpp new file mode 100644 index 00000000..633c2d5a --- /dev/null +++ b/tests/dispatch/modules_static_init.cpp @@ -0,0 +1,38 @@ +#include +#include + +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#if defined(RABITQ_TEST_INDEX) +#include "rabitqlib/index/hnsw/hnsw.hpp" +#endif + +namespace { +const bool kBeforeMain = [] { + std::array query{}; + std::array lut{}; + rabitqlib::fastscan::pack_lut(64, query.data(), lut.data()); + std::array integer_lut{}; + std::array packed_lut{}; + std::array codes{}; + std::array scores{}; + rabitqlib::fastscan::transfer_lut_hacc(integer_lut.data(), 64, packed_lut.data()); + rabitqlib::fastscan::accumulate_hacc( + codes.data(), packed_lut.data(), scores.data(), 64 + ); + for (const auto score : scores) { + if (score != 0) { + return false; + } + } +#if defined(RABITQ_TEST_INDEX) + rabitqlib::hnsw::HierarchicalNSW index(2, 64, 4, 2, 16, 42); + const auto result = index.search(query.data(), 1, 1, 10, 1); + if (result.size() != 1 || !result[0].empty()) { + return false; + } +#endif + return true; +}(); +} // namespace +int main() { return kBeforeMain ? 0 : 1; } diff --git a/tests/failures/compact_codes_failure_test.cpp b/tests/failures/compact_codes_failure_test.cpp new file mode 100644 index 00000000..693da78a --- /dev/null +++ b/tests/failures/compact_codes_failure_test.cpp @@ -0,0 +1,396 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/utils/openmp.hpp" +#ifdef RABITQ_TEST_FASTSCAN +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/encoder.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/query.hpp" +#endif + +namespace { +std::atomic watching{false}; +std::atomic allocation_count{0}; +std::atomic deepest_region{0}; +size_t failure_at = std::numeric_limits::max(); + +bool fail_allocation() { + if (!watching) { + return false; + } +#if RABITQ_USE_OPENMP +#if _OPENMP >= 200805 + const int level = omp_get_level(); +#else + // OpenMP 2.0 can identify active teams, but not serialized regions. + const int level = omp_in_parallel(); +#endif + int previous = deepest_region.load(); + while (previous < level && !deepest_region.compare_exchange_weak(previous, level)) {} +#endif + return allocation_count.fetch_add(1) == failure_at; +} +} // namespace + +#ifdef RABITQ_WRAP_COMPACT_ALLOCATIONS +extern "C" void* __real_malloc(size_t); +extern "C" void* __real_realloc(void*, size_t); +extern "C" void* __wrap_malloc(size_t bytes) { + return fail_allocation() ? nullptr : __real_malloc(bytes); +} +extern "C" void* __wrap_realloc(void* ptr, size_t bytes) { + return fail_allocation() ? nullptr : __real_realloc(ptr, bytes); +} +#endif + +GTEST_NO_INLINE_ void* operator new(size_t bytes) { + if (fail_allocation()) { + throw std::bad_alloc(); + } +#ifdef RABITQ_WRAP_COMPACT_ALLOCATIONS + void* ptr = __real_malloc(bytes == 0 ? 1 : bytes); +#else + void* ptr = std::malloc(bytes == 0 ? 1 : bytes); +#endif + if (!ptr) { + throw std::bad_alloc(); + } + return ptr; +} +GTEST_NO_INLINE_ void* operator new[](size_t bytes) { return ::operator new(bytes); } +GTEST_NO_INLINE_ void* operator new(size_t bytes, const std::nothrow_t&) noexcept { + try { + return ::operator new(bytes); + } catch (...) { return nullptr; } +} +GTEST_NO_INLINE_ void* operator new[](size_t bytes, const std::nothrow_t&) noexcept { + try { + return ::operator new[](bytes); + } catch (...) { return nullptr; } +} +GTEST_NO_INLINE_ void operator delete(void* ptr) noexcept { + std::free(ptr); // NOLINT(clang-analyzer-cplusplus.NewDelete) +} +GTEST_NO_INLINE_ void operator delete[](void* ptr) noexcept { ::operator delete(ptr); } +GTEST_NO_INLINE_ void operator delete(void* ptr, size_t) noexcept { + ::operator delete(ptr); +} +GTEST_NO_INLINE_ void operator delete[](void* ptr, size_t) noexcept { + ::operator delete(ptr); +} + +namespace { +template +struct Outcome { + bool caught = false; + bool first_ok = false; + bool retry_ok = false; + size_t allocations = 0; + int level = 0; + bool outer_active = false; + std::vector codes; + std::vector add, rescale, error; + + bool matches( + const uint8_t* actual_codes, + const std::vector& actual_add, + const std::vector& actual_rescale, + const std::vector& actual_error + ) const { + return std::equal(codes.begin(), codes.end(), actual_codes) && add == actual_add && + rescale == actual_rescale && error == actual_error; + } +}; + +template +Outcome run_cold_thread( + bool grow, bool outer_region, size_t fail_at, const Outcome* reference = nullptr +) { + Outcome outcome; + // A new native thread gives both it and its OpenMP team cold TLS buffers. + std::thread worker([&] { +#if RABITQ_USE_OPENMP + omp_set_dynamic(0); + omp_set_num_threads(2); +#endif + constexpr size_t kRows = 4, kDim = 128; + std::vector data(kRows * kDim, T{0.5}), centroid(kDim, T{0}); + std::vector codes((Batch ? 32 : kRows) * kDim / 8); + std::vector add(kRows), rescale(kRows), error(kRows); +#ifdef RABITQ_TEST_FASTSCAN + std::vector batch(rabitqlib::BatchDataMap::data_bytes(kDim)); +#endif + const auto encode = [&](size_t dim) { +#ifdef RABITQ_TEST_FASTSCAN + if constexpr (Public) { + rabitqlib::quant::quantize_one_batch( + data.data(), centroid.data(), kRows, dim, batch.data() + ); + rabitqlib::BatchDataMap view(batch.data(), dim); + view.f_add().copy_to(add.data(), kRows); + view.f_rescale().copy_to(rescale.data(), kRows); + view.f_error().copy_to(error.data(), kRows); + } else if constexpr (Batch) { + rabitqlib::quant::rabitq_impl::one_bit::one_bit_batch_code( + data.data(), + centroid.data(), + kRows, + dim, + codes.data(), + add.data(), + rescale.data(), + error.data() + ); + } else +#endif + { + rabitqlib::quant::rabitq_impl::one_bit:: + one_bit_compact_codes( + data.data(), + centroid.data(), + kRows, + dim, + codes.data(), + add.data(), + rescale.data(), + error.data() + ); + } + }; + if (grow) { + encode(32); + } + const uint8_t* code_data = codes.data(); +#ifdef RABITQ_TEST_FASTSCAN + if constexpr (Public) { + code_data = rabitqlib::ConstBatchDataMap(batch.data(), kDim).bin_code(); + } +#endif + const auto valid_result = [&] { + for (size_t i = 0; i < kRows; ++i) { + if (!std::isfinite(add[i]) || !std::isfinite(rescale[i]) || + !std::isfinite(error[i])) { + return false; + } + } + return reference == nullptr || + reference->matches(code_data, add, rescale, error); + }; + const auto exercise = [&] { +#if RABITQ_USE_OPENMP + outcome.outer_active = omp_in_parallel() != 0; +#endif + failure_at = fail_at; + allocation_count = 0; + deepest_region = 0; + watching = true; + try { + encode(kDim); + } catch (const std::bad_alloc&) { outcome.caught = true; } + watching = false; + outcome.allocations = allocation_count; + outcome.level = deepest_region; + outcome.first_ok = !outcome.caught && valid_result(); + }; +#if RABITQ_USE_OPENMP + if (outer_region) { +#pragma omp parallel num_threads(1) + { exercise(); } + } else +#else + (void)outer_region; +#endif + { + exercise(); + } + // Capture the no-failure probe outside allocation tracking. Each trial + // still starts on a fresh thread and compares against this fixed result. + if (reference == nullptr) { + outcome.codes.assign(code_data, code_data + codes.size()); + outcome.add = add; + outcome.rescale = rescale; + outcome.error = error; + reference = &outcome; + } + encode(kDim); + outcome.retry_ok = valid_result(); + }); + worker.join(); + return outcome; +} + +template +void check_failures(bool outer_region) { + for (bool grow : {false, true}) { + const auto probe = run_cold_thread( + grow, outer_region, std::numeric_limits::max() + ); + ASSERT_FALSE(probe.caught); + ASSERT_TRUE(probe.first_ok); + ASSERT_TRUE(probe.retry_ok); + ASSERT_GT(probe.allocations, Batch ? 1U : 0U); +#if RABITQ_USE_OPENMP +#if _OPENMP >= 200805 + EXPECT_EQ(probe.level, Parallel ? 1 : (outer_region ? 1 : 0)); +#else + EXPECT_EQ(probe.level, Parallel || probe.outer_active ? 1 : 0); +#endif +#else + EXPECT_EQ(probe.level, 0); +#endif + size_t caught_count = 0; + for (size_t point = 0; point < probe.allocations; ++point) { + SCOPED_TRACE(::testing::Message() << "grow=" << grow << " failure=" << point); + const auto result = run_cold_thread( + grow, outer_region, point, &probe + ); + // An OpenMP runtime allocation can fail without stopping encoding. + // In that case the first result must match the no-failure probe. + EXPECT_TRUE(result.caught || result.first_ok); + caught_count += result.caught; + EXPECT_TRUE(result.retry_ok); + } + EXPECT_GT(caught_count, 0U); + } +} + +TEST(CompactCodesFailure, ReferenceRejectsIncompleteAndCorruptOutputs) { + Outcome reference; + reference.codes = {0xAA, 0x55}; + reference.add = {1}; + reference.rescale = {2}; + reference.error = {3}; + auto actual = reference; + const auto matches = [&] { + return reference.matches( + actual.codes.data(), actual.add, actual.rescale, actual.error + ); + }; + EXPECT_TRUE(matches()); + actual.codes[0] ^= 1; + EXPECT_FALSE(matches()); + actual = reference; + for (auto* factors : {&actual.add, &actual.rescale, &actual.error}) { + const float original = factors->front(); + factors->front() = 0; + EXPECT_FALSE(matches()); + factors->front() = original; + } + // Unwritten output is finite, but cannot pass the reference comparison. + std::fill(actual.codes.begin(), actual.codes.end(), 0); + actual.add[0] = actual.rescale[0] = actual.error[0] = 0; + EXPECT_FALSE(matches()); +} + +TEST(CompactCodesFailure, SerialColdScratchAndGrowth) { + check_failures(false); + check_failures(false); +} +TEST(CompactCodesFailure, ParallelColdScratchAndGrowth) { + check_failures(false); + check_failures(false); +} +TEST(CompactCodesFailure, SerialInsideOuterRegion) { + check_failures(true); + check_failures(true); +} +#ifdef RABITQ_TEST_FASTSCAN +TEST(SplitBatchQueryAllocation, ReusesCapacityAcrossQueriesModesAndDimensions) { + std::vector values(128, 0.5F), next(128, 0.25F), scratch(128 * 4); + std::vector hacc_scratch(128 * 4); + rabitqlib::SplitBatchQuery query(values.data(), 128, 3); + const auto* storage = query.lut(); + failure_at = 0; + allocation_count = 0; + bool caught = false; + watching = true; + try { + for (size_t repeat = 0; repeat < 3; ++repeat) { + query.reset( + next.data(), 64, 0, scratch.data(), nullptr, rabitqlib::METRIC_IP, false + ); + query.reset(values.data(), 128, 8, scratch.data(), hacc_scratch.data()); + } + } catch (const std::bad_alloc&) { caught = true; } + watching = false; + EXPECT_FALSE(caught); + EXPECT_EQ(allocation_count, 0U); + EXPECT_EQ(query.lut(), storage); + EXPECT_EQ(query.padded_dim(), 128U); + EXPECT_TRUE(query.use_hacc()); +} + +TEST(SplitBatchQueryAllocation, InitialPreparationAndGrowthFailurePreserveState) { + std::vector values(128, 0.5F), next(128, 0.25F), scratch(128 * 4); + std::vector hacc_scratch(128 * 4); + for (bool prepared : {false, true}) { + rabitqlib::SplitBatchQuery query; + if (prepared) { + query.reset( + values.data(), 32, 3, scratch.data(), nullptr, rabitqlib::METRIC_IP, false + ); + query.set_g_add(3, 2); + } + const auto original = query; + std::vector bytes; + if (prepared) { + bytes.assign(query.lut(), query.lut() + 32 * 4); + } + failure_at = 0; + allocation_count = 0; + bool caught = false; + watching = true; + try { + query.reset(next.data(), 128, 8, scratch.data(), hacc_scratch.data()); + } catch (const std::bad_alloc&) { caught = true; } + watching = false; + EXPECT_TRUE(caught); + EXPECT_EQ(allocation_count, 1U); + EXPECT_EQ(query.rotated_query(), original.rotated_query()); + EXPECT_EQ(query.padded_dim(), original.padded_dim()); + EXPECT_EQ(query.use_hacc(), original.use_hacc()); + EXPECT_EQ(query.delta(), original.delta()); + EXPECT_EQ(query.sum_vl_lut(), original.sum_vl_lut()); + EXPECT_EQ(query.k1xsumq(), original.k1xsumq()); + EXPECT_EQ(query.kbxsumq(), original.kbxsumq()); + EXPECT_EQ(query.g_add(), original.g_add()); + EXPECT_EQ(query.g_error(), original.g_error()); + if (prepared) { + EXPECT_TRUE(std::equal(bytes.begin(), bytes.end(), query.lut())); + query.set_g_add(5, 4); + EXPECT_EQ(query.g_add(), -4); + } + query.reset(next.data(), 128, 8, scratch.data(), hacc_scratch.data()); + EXPECT_EQ(query.rotated_query(), next.data()); + EXPECT_EQ(query.padded_dim(), 128U); + EXPECT_EQ(query.g_add(), 0); + EXPECT_EQ(query.g_error(), 0); + } +} + +TEST(CompactCodesFailure, PublicBatchSerialAndParallelPropagateAndRecover) { + check_failures(true); + check_failures(false); + check_failures(false); +} +TEST(CompactCodesFailure, BatchInnerScratchAfterOuterAllocation) { + check_failures(true); + check_failures(true); + check_failures(false); +} +#endif +} // namespace diff --git a/tests/failures/construction_failure_test.cpp b/tests/failures/construction_failure_test.cpp new file mode 100644 index 00000000..66dcd3f1 --- /dev/null +++ b/tests/failures/construction_failure_test.cpp @@ -0,0 +1,285 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/index/types.hpp" + +namespace { +std::atomic fail_parallel{false}; +std::atomic remaining{0}; +std::atomic fail_constructor{false}; +bool tracking = false; +std::array live{}; +size_t live_count = 0; + +bool should_fail() { +#if _OPENMP >= 200805 + const bool in_region = omp_get_level() > 0; +#else + const bool in_region = omp_in_parallel() != 0; +#endif + if (fail_constructor || (fail_parallel && in_region)) { + if (remaining.fetch_sub(1) == 0) { + fail_constructor = false; + fail_parallel = false; + return true; + } + } + return false; +} +void record(void* ptr) { + if (tracking && ptr != nullptr) { + if (live_count == live.size()) { + std::abort(); + } + live[live_count++] = ptr; + } +} +void forget(void* ptr) { + if (!tracking) { + return; + } + for (size_t i = 0; i < live_count; ++i) { + if (live[i] == ptr) { + live[i] = live[--live_count]; + return; + } + } +} +} // namespace + +#if defined(RABITQ_WRAP_CONSTRUCTION_ALLOCATIONS) +extern "C" void* __real_malloc(size_t); +extern "C" void* __real_aligned_alloc(size_t, size_t); +extern "C" void __real_free(void*); +extern "C" void* __wrap_malloc(size_t size) { + if (fail_constructor && should_fail()) { + return nullptr; + } + void* ptr = __real_malloc(size); + record(ptr); + return ptr; +} +extern "C" void* __wrap_aligned_alloc(size_t alignment, size_t size) { + if (fail_constructor && should_fail()) { + return nullptr; + } + void* ptr = __real_aligned_alloc(alignment, size); + record(ptr); + return ptr; +} +extern "C" void __wrap_free(void* ptr) { + forget(ptr); + __real_free(ptr); +} +#endif + +GTEST_NO_INLINE_ void* operator new(size_t size) { + if (should_fail()) { + throw std::bad_alloc(); + } +#if defined(RABITQ_WRAP_CONSTRUCTION_ALLOCATIONS) + void* ptr = __real_malloc(size == 0 ? 1 : size); +#else + void* ptr = std::malloc(size == 0 ? 1 : size); +#endif + if (ptr == nullptr) { + throw std::bad_alloc(); + } + record(ptr); + return ptr; +} +GTEST_NO_INLINE_ void* operator new[](size_t size) { return ::operator new(size); } +GTEST_NO_INLINE_ void* operator new(size_t size, const std::nothrow_t&) noexcept { + try { + return ::operator new(size); + } catch (...) { return nullptr; } +} +GTEST_NO_INLINE_ void* operator new[](size_t size, const std::nothrow_t&) noexcept { + try { + return ::operator new[](size); + } catch (...) { return nullptr; } +} +GTEST_NO_INLINE_ void operator delete(void* ptr) noexcept { +#if !defined(RABITQ_WRAP_CONSTRUCTION_ALLOCATIONS) + forget(ptr); +#endif + std::free(ptr); // NOLINT(clang-analyzer-cplusplus.NewDelete) +} +GTEST_NO_INLINE_ void operator delete[](void* ptr) noexcept { ::operator delete(ptr); } +GTEST_NO_INLINE_ void operator delete(void* ptr, size_t) noexcept { + ::operator delete(ptr); +} +GTEST_NO_INLINE_ void operator delete[](void* ptr, size_t) noexcept { + ::operator delete(ptr); +} + +namespace { +template +void expect_parallel_failure(Function fn, size_t skip = 0) { + remaining = skip; + fail_parallel = true; + bool caught = false; + try { + fn(); + } catch (const std::bad_alloc&) { caught = true; } + fail_parallel = false; + EXPECT_TRUE(caught); +} + +std::vector vectors() { + std::vector data(64 * 64); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i % 37) / 37.0F; + } + return data; +} + +TEST(ConstructionFailure, HnswConstructorReleasesEveryAllocation) { +#if !defined(RABITQ_WRAP_CONSTRUCTION_ALLOCATIONS) + GTEST_SKIP() << "C allocation ownership tracking requires linker wrapping"; +#else + remaining = std::numeric_limits::max(); + fail_constructor = true; + { rabitqlib::hnsw::HierarchicalNSW index(32, 64, 4, 4, 16, 100); } + fail_constructor = false; + const size_t allocation_count = std::numeric_limits::max() - remaining; + RecordProperty("constructor_allocations", static_cast(allocation_count)); + size_t failures = 0; + bool completed = false; + for (size_t allocation = 0; allocation <= allocation_count; ++allocation) { + SCOPED_TRACE(allocation); + remaining = allocation; + tracking = true; + fail_constructor = true; + bool caught = false; + try { + rabitqlib::hnsw::HierarchicalNSW index(32, 64, 4, 4, 16, 100); + } catch (const std::bad_alloc&) { + caught = true; + } catch (const std::runtime_error&) { caught = true; } + fail_constructor = false; + tracking = false; + ASSERT_EQ(live_count, 0U); + if (!caught) { + completed = true; + break; + } + ++failures; + } + EXPECT_TRUE(completed); + EXPECT_GT(failures, 5U); +#endif +} + +class ConstructionWorkerFailure : public ::testing::TestWithParam { + void SetUp() override { +#if _OPENMP < 200805 + if (GetParam() == 1) { + bool detectable = false; +#pragma omp parallel num_threads(1) + { detectable = omp_in_parallel() != 0; } + if (!detectable) { + GTEST_SKIP() << "This OpenMP runtime cannot identify one-thread regions"; + } + } +#endif + omp_set_dynamic(0); + } +}; + +INSTANTIATE_TEST_SUITE_P(Threads, ConstructionWorkerFailure, ::testing::Values(1U, 4U)); + +TEST_P(ConstructionWorkerFailure, IvfWorkerErrorsReachCallerAndAllowRetry) { + const auto data = vectors(); + std::vector centers(2 * 64, 0.0F); + std::vector ids(64); + for (size_t i = 0; i < ids.size(); ++i) { + ids[i] = i % 2; + } + const size_t threads = GetParam(); + { + for (const size_t bits : {1U, 4U, 32U}) { + rabitqlib::ivf::IVF index(64, 64, 2, bits); + expect_parallel_failure([&] { + index.construct(data.data(), centers.data(), ids.data(), false, threads); + }); + EXPECT_EQ(index.max_elements(), 64U); + EXPECT_EQ(index.num_stored_points(), 0U); + EXPECT_EQ(index.num_points(), 0U); + ASSERT_NO_THROW( + index.construct(data.data(), centers.data(), ids.data(), false, threads) + ); + EXPECT_EQ(index.num_stored_points(), 64U); + EXPECT_EQ(index.num_points(), 64U); + rabitqlib::PID id = 0; + float distance = 0; + ASSERT_NO_THROW(index.search(data.data(), 1, 1, &id, &distance)); + } + } +} + +TEST_P(ConstructionWorkerFailure, QgInitializationWorkerErrorsReachCaller) { + using namespace rabitqlib; + const auto data = vectors(); + const size_t threads = GetParam(); + { + for (const auto init : + {symqg::QGInitialization::Random, symqg::QGInitialization::PiPNN}) { + for (const size_t bits : {0U, 4U, 8U}) { + symqg::QuantizedGraph graph( + 64, 64, 32, METRIC_L2, RotatorType::FhtKacRotator, bits, 42 + ); + expect_parallel_failure([&] { + symqg::QGBuilder builder( + graph, 32, data.data(), threads, init, 42, true + ); + }); + EXPECT_EQ(graph.max_elements(), 64U); + EXPECT_EQ(graph.num_stored_points(), 0U); + EXPECT_EQ(graph.num_points(), 0U); + ASSERT_NO_THROW( + (symqg::QGBuilder(graph, 32, data.data(), threads, init, 42, true)) + ); + } + } + } +} + +TEST_P(ConstructionWorkerFailure, QgRefinementWorkerErrorsReachCaller) { + using namespace rabitqlib; + const auto data = vectors(); + const size_t threads = GetParam(); + { + for (const size_t skip : {0U, 4U, 64U}) { + symqg::QuantizedGraph graph( + 64, 64, 32, METRIC_L2, RotatorType::FhtKacRotator, 4, 42 + ); + symqg::QGBuilder builder( + graph, 32, data.data(), threads, symqg::QGInitialization::Random, 42 + ); + expect_parallel_failure([&] { builder.build(); }, skip); + EXPECT_EQ(graph.num_stored_points(), 0U); + EXPECT_EQ(graph.num_points(), 0U); + ASSERT_NO_THROW(builder.reset(data.data())); + ASSERT_NO_THROW(builder.build()); + EXPECT_EQ(graph.num_stored_points(), 64U); + EXPECT_EQ(graph.num_points(), 64U); + } + } +} +} // namespace diff --git a/tests/failures/hnsw_add_allocation_test.cpp b/tests/failures/hnsw_add_allocation_test.cpp index dbe9af55..5d42823e 100644 --- a/tests/failures/hnsw_add_allocation_test.cpp +++ b/tests/failures/hnsw_add_allocation_test.cpp @@ -1,4 +1,5 @@ #include +#include #include #include @@ -12,6 +13,7 @@ #include #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/types.hpp" namespace { std::atomic fail_enabled{false}; @@ -117,6 +119,8 @@ TEST(HnswConstructionAllocationTest, EveryAllocationFailureLeavesEmptyReusableIn } ++failures; EXPECT_EQ(index.num_points(), 0U); + EXPECT_EQ(index.num_stored_points(), 0U); + EXPECT_EQ(index.max_elements(), kCapacity); EXPECT_EQ(index.num_clusters(), 0U); const auto empty = index.search(fixture.data.data(), 1, 1, kCapacity, 1); ASSERT_EQ(empty.size(), 1U); @@ -143,6 +147,47 @@ TEST(HnswConstructionAllocationTest, EveryAllocationFailureLeavesEmptyReusableIn } } +TEST(HnswResizeAllocationTest, FailurePreservesEveryLinkOwner) { + Fixture fixture; + size_t failures = 0; + bool succeeded = false; + for (size_t allocation = 0; allocation < 32; ++allocation) { + SCOPED_TRACE(allocation); + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); + auto clusters = fixture.clusters; + index.construct( + 1, + fixture.centroid.data(), + kCapacity, + fixture.data.data(), + clusters.data(), + 1, + false + ); + const auto before = index.search(fixture.data.data(), kCapacity, 3, kCapacity, 1); + bool failed = false; + { + AllocationFailure failure(allocation); + try { + index.resize(2 * kCapacity); + } catch (const std::bad_alloc&) { failed = true; } + } + EXPECT_EQ(index.num_points(), kCapacity); + EXPECT_EQ(index.search(fixture.data.data(), kCapacity, 3, kCapacity, 1), before); + if (!failed) { + EXPECT_EQ(index.max_elements(), 2 * kCapacity); + succeeded = true; + break; + } + ++failures; + EXPECT_EQ(index.max_elements(), kCapacity); + ASSERT_NO_THROW(index.resize(2 * kCapacity)); + EXPECT_EQ(index.search(fixture.data.data(), kCapacity, 3, kCapacity, 1), before); + } + EXPECT_GT(failures, 0U); + EXPECT_TRUE(succeeded); +} + TEST(HnswAddAllocationTest, AllocatesReturnedIdsBeforeMutation) { Fixture fixture; HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); @@ -260,7 +305,7 @@ TEST(HnswAddAllocationTest, EveryAllocationFailureLeavesUsableIndex) { ASSERT_NO_THROW(loaded.load(path.c_str())); EXPECT_EQ(loaded.num_points(), index.num_points()); EXPECT_EQ(loaded.search(fixture.data.data(), 1, 1, kCapacity, 1), results); - const auto next_id = static_cast(loaded.num_points()); + const auto next_id = static_cast(loaded.num_stored_points()); EXPECT_EQ( loaded.add( fixture.data.data() + ((1 + kAdded) * kDim), diff --git a/tests/fixtures/core_api_v1/README.md b/tests/fixtures/core_api_v1/README.md new file mode 100644 index 00000000..e6daef1d --- /dev/null +++ b/tests/fixtures/core_api_v1/README.md @@ -0,0 +1,14 @@ +# Core quantization API v1 conformance consumer + +`consumer.cpp` is the initial source-compatibility and numerical-contract fixture for +`rabitqlib/core/api.hpp`. It checks signatures, defaults used by a real consumer, +enum values, separate factor storage, both metrics, logical code meaning and the +one-bit packed format. It is not a C ABI fixture or a release-produced index file. + +Keep this consumer unchanged when extending v1. Add new consumers for additive APIs; +do not edit away failures caused by an incompatible change. Intentional breaking +changes need a new API major and migration instructions. Do not freeze workspace byte +counts or private class sizes: callers must query sizes and recompile on upgrades. + +The C++ suite and external add_subdirectory/FetchContent/installed consumer tests all +compile and execute this same source, including with FastScan and indexes disabled. diff --git a/tests/fixtures/core_api_v1/consumer.cpp b/tests/fixtures/core_api_v1/consumer.cpp new file mode 100644 index 00000000..0fed6318 --- /dev/null +++ b/tests/fixtures/core_api_v1/consumer.cpp @@ -0,0 +1,101 @@ +// Frozen v1 source consumer. Keep this independent of implementation/legacy headers. +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" + +namespace api = rabitqlib::quant::v1; +static_assert(RABITQLIB_QUANT_API_VERSION_MAJOR == 1); +static_assert(RABITQLIB_QUANT_API_VERSION_MINOR >= 0); +static_assert(static_cast(api::CodeLayout::UnpackedU16) == 0); +static_assert(static_cast(api::CodeLayout::SplitPacked) == 1); +static_assert(static_cast(rabitqlib::METRIC_L2) == 0); +static_assert(static_cast(rabitqlib::METRIC_IP) == 1); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_same_v< + decltype(&api::requirements), + api::Requirements (*)(const api::Parameters&, api::CodeLayout)>); +static_assert(std::is_same_v< + decltype(&api::encode), + api::Factors (*)( + const api::Parameters&, + const float*, + const float*, + api::CodeOutput, + api::Buffer + )>); +static_assert(std::is_same_v< + decltype(&api::convert_codes), + void (*)( + const api::Parameters&, api::CodeView, api::CodeOutput, api::Buffer + )>); +static_assert(std::is_same_v< + decltype(&api::prepare_query), + api::QueryState (*)(const api::Parameters&, const float*, const float*)>); +static_assert(std::is_same_v< + decltype(&api::compute), + float (*)(const api::QueryState&, api::CodeView, const api::Factors&)>); + +int main() { + constexpr size_t dim = 32; + std::array input{}, query{}; + for (size_t i = 0; i < dim; ++i) { + input[i] = i % 2 == 0 ? 1.0F : -1.0F; + query[i] = input[i] * 0.5F; + } + // Deliberately not the library's record layout: user metadata, separated factors. + struct Record { + float scale = 0; + std::uint64_t id = 123; + std::array signs{}; + float error = 0; + float add = 0; + } record; + for (auto metric : {rabitqlib::METRIC_L2, rabitqlib::METRIC_IP}) { + api::Parameters p{dim, 1, metric}; + auto sizes = api::requirements(p, api::CodeLayout::SplitPacked); + if (sizes.code_bytes != 4 || sizes.extra_bytes != 0) + return 1; + std::vector scratch( + (sizes.workspace_bytes + sizeof(std::max_align_t) - 1) / + sizeof(std::max_align_t) + ); + api::Buffer work{scratch.data(), sizes.workspace_bytes}; + auto factors = api::encode( + p, + input.data(), + nullptr, + {api::CodeLayout::SplitPacked, {record.signs.data(), record.signs.size()}, {}}, + work + ); + record.add = factors.add; + record.scale = factors.rescale; + record.error = factors.error; + if (record.id != 123 || record.error != 0) + return 2; + for (auto value : record.signs) + if (value != 0xAA) + return 3; + if (record.add != (metric == rabitqlib::METRIC_L2 ? 32 : 1) || + record.scale != (metric == rabitqlib::METRIC_L2 ? -4 : -2)) + return 4; + auto state = api::prepare_query(p, query.data(), nullptr); + api::CodeView packed{ + api::CodeLayout::SplitPacked, {record.signs.data(), record.signs.size()}, {}}; + if (api::compute(state, packed, {record.add, record.scale, record.error}) != + (metric == rabitqlib::METRIC_L2 ? 8 : -15)) + return 5; + std::array raw{}; + api::convert_codes( + p, packed, {api::CodeLayout::UnpackedU16, {raw.data(), raw.size()}, {}}, work + ); + for (size_t i = 0; i < dim; ++i) { + if (raw[2 * i] != (i % 2 == 0 ? 1 : 0) || raw[2 * i + 1] != 0) + return 6; + } + } +} diff --git a/tests/fixtures/persistence/generate.cpp b/tests/fixtures/persistence/generate.cpp index 548aa5ba..f10e98ed 100644 --- a/tests/fixtures/persistence/generate.cpp +++ b/tests/fixtures/persistence/generate.cpp @@ -4,15 +4,18 @@ #include #include #include +#include #include #include #include #include +#include "rabitqlib/defines.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" #include "rabitqlib/index/ivf/ivf.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/utils/rotator.hpp" int main(int argc, char** argv) { if (argc != 2) { diff --git a/tests/integration/bit_pack_unpack_test.cpp b/tests/integration/bit_pack_unpack_test.cpp deleted file mode 100644 index 57bc3849..00000000 --- a/tests/integration/bit_pack_unpack_test.cpp +++ /dev/null @@ -1,135 +0,0 @@ -#include - -#include -#include -#include -#include -#include -#include - -using namespace rabitqlib; - -class BitPackUnpackTest : public ::testing::Test { - protected: - // 1. Shared Constants & Data Structures - const size_t dim = 768; - std::vector query; - std::vector code; - std::vector compact_code; - - // 2. Common Initialization (Runs before EVERY test) - void SetUp() override { - srand(42); - - // Initialize Query (Same for all tests) - query.resize(dim); - for (size_t i = 0; i < dim; ++i) { - query[i] = static_cast((rand() * 100.0) / RAND_MAX); - } - - // Pre-allocate code vector - code.resize(dim); - } - - // 3. Helper to handle bit-specific setup - void PrepareData(size_t bits) { - // Resize compact code based on bits - compact_code.resize(dim * bits / 8 + 1); - - // Generate random codes based on bit depth - for (size_t i = 0; i < dim; ++i) { - code[i] = rand() % (1 << bits); - } - - // Pack the code - rabitqlib::quant::rabitq_impl::ex_bits::packing_rabitqplus_code( - code.data(), compact_code.data(), dim, bits - ); - } - - // 4. Helper for Ground Truth Calculation - double CalculateExpected() const { - double expected_result = 0; - for (size_t i = 0; i < dim; ++i) { - expected_result += static_cast(query[i]) * code[i]; - } - return expected_result; - } - - void ExpectIpNear(float result) const { - double expected = CalculateExpected(); - ASSERT_NEAR(expected, result, std::max(0.1, std::abs(expected) * 1e-6)); - } -}; - -// --- Test Cases --- - -TEST_F(BitPackUnpackTest, ExCode1Bit) { - PrepareData(1); // Set up for 1-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip16_fxu1_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode2Bit) { - PrepareData(2); // Set up for 2-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip64_fxu2_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode3Bit) { - PrepareData(3); // Set up for 3-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip64_fxu3_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode4Bit) { - PrepareData(4); // Set up for 4-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip16_fxu4_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode5Bit) { - PrepareData(5); // Set up for 5-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip64_fxu5_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode6Bit) { - PrepareData(6); // Set up for 6-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip64_fxu6_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} - -TEST_F(BitPackUnpackTest, ExCode7Bit) { - PrepareData(7); // Set up for 7-bit - - // Run the AVX function - float result = - rabitqlib::excode_ipimpl::ip64_fxu7_avx(query.data(), compact_code.data(), dim); - - ExpectIpNear(result); -} diff --git a/tests/python/test_build_cluster_ids.py b/tests/python/test_build_cluster_ids.py new file mode 100644 index 00000000..61d8844d --- /dev/null +++ b/tests/python/test_build_cluster_ids.py @@ -0,0 +1,57 @@ +"""Build must validate original assignments before narrowing them to point IDs.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +@pytest.mark.parametrize( + "bad", + [2**32, -(2**32), 0.5, -1, 1, np.uint64(2**64 - 1), True, "0"], +) +def test_build_rejects_invalid_cluster_ids(kind, bad): + data = np.ones((2, 64), dtype=np.float32) + centroids = np.zeros((1, 64), dtype=np.float32) + index = ( + IvfIndex(64, 2, 1, 4) + if kind == "ivf" + else HnswIndex(64, 2, M=2, ef_construction=10, nbits=4) + ) + ids = np.full(2, bad) + with pytest.raises(ValueError, match="cluster_ids"): + index.build(data, centroids, ids) + assert not index.is_built + index.build(data, centroids, np.zeros(2, dtype=np.int64)) + assert index.is_built + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +@pytest.mark.parametrize("layout", ["native", "strided", "swapped", "unaligned"]) +def test_build_uint32_cluster_ids_validate_and_preserve_input(kind, layout): + data = np.ones((2, 64), dtype=np.float32) + data[1] = 2 + centroids = data.copy() + index = ( + IvfIndex(64, 2, 2, 4) + if kind == "ivf" + else HnswIndex(64, 2, M=2, ef_construction=10, nbits=4) + ) + if layout == "strided": + ids = np.zeros(4, dtype=np.uint32)[::2] + elif layout == "swapped": + ids = np.zeros(2, dtype=np.dtype(np.uint32).newbyteorder("S")) + elif layout == "unaligned": + ids = np.ndarray(2, dtype=np.uint32, buffer=bytearray(9), offset=1) + else: + ids = np.zeros(2, dtype=np.uint32) + ids[1] = np.iinfo(np.uint32).max + with pytest.raises(ValueError, match="cluster_ids"): + index.build(data, centroids, ids, num_threads=1) + assert not index.is_built + ids[1] = 1 + ids.flags.writeable = False + before = ids.copy() + index.build(data, centroids, ids, num_threads=1) + assert index.is_built + np.testing.assert_array_equal(ids, before) diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index 2a6b4ead..0286d556 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -138,17 +138,27 @@ def test_query_timing_handles_unchanged_clock(kind, elapsed, monkeypatch, capsys # Wall-clock timestamps can be identical for a short search on Windows. monkeypatch.setattr(module, "time", lambda: 100.0, raising=False) - ticks = iter([100.0, 100.0 + elapsed]) + ticks = iter([100.0, 100.0 + elapsed] * 2) monkeypatch.setattr(module, "perf_counter", lambda: next(ticks), raising=False) monkeypatch.setattr(module, "NPROBES" if kind == "ivf" else "EFS", [1]) ids = np.array([[0], [1]], dtype=np.uint32) - index = SimpleNamespace( - dim=2, num_clusters=2, search=lambda *a, **kw: (ids, np.zeros((2, 1))) - ) + queries = np.arange(6, dtype=np.float32).reshape(2, 3)[:, 1:] + assert not queries.flags.c_contiguous + searched = [] + + def search(prepared, *args, **kwargs): + assert prepared.flags.c_contiguous + np.testing.assert_array_equal(prepared, queries) + if searched: + assert prepared is searched[0] + searched.append(prepared) + return ids, np.zeros((2, 1)) + + index = SimpleNamespace(dim=2, num_clusters=2, search=search) index.search_batch = index.search index_name = {"ivf": "IvfIndex", "hnsw": "HnswIndex", "symqg": "SymqgIndex"}[kind] monkeypatch.setattr(module, index_name, SimpleNamespace(load=lambda path: index)) - monkeypatch.setattr(module, "read_fvecs", lambda path: np.zeros((2, 2))) + monkeypatch.setattr(module, "read_fvecs", lambda path: queries) monkeypatch.setattr(module, "read_ivecs", lambda path: ids) module.main( SimpleNamespace( @@ -156,12 +166,13 @@ def test_query_timing_handles_unchanged_clock(kind, elapsed, monkeypatch, capsys query_file="queries", gt_file="truth", topk=1, - test_rounds=1, + test_rounds=2, num_threads=1, metric="l2", use_hacc=None, ) ) + assert len(searched) == 2 row = capsys.readouterr().out.splitlines()[-1].split() qps, recall = map(float, row[1:]) if elapsed == 0: diff --git a/tests/python/test_historical_persistence.py b/tests/python/test_historical_persistence.py index 09304b01..c822e218 100644 --- a/tests/python/test_historical_persistence.py +++ b/tests/python/test_historical_persistence.py @@ -49,7 +49,9 @@ def test_historical_load_search_and_resave(name, tmp_path): for index in (original, cls.load(str(destination))): assert index.dim == 65 assert index.metric == name.rsplit("_", 1)[1] - assert (index.max_elements if cls is IvfIndex else index.num_points) == count + assert index.num_stored_points == count + assert index.num_points == count - (0 if cls is SymqgIndex else 1) + assert index.max_elements >= index.num_stored_points actual_ids, actual_distances = index.search(queries, 5, **options) np.testing.assert_array_equal(actual_ids, ids) # Distances may differ across AVX2, AVX-512, NEON, and compiler reductions. @@ -80,7 +82,7 @@ def test_unknown_historical_version_is_rejected(name, tmp_path): @pytest.mark.parametrize("name", ["ivf4_l2", "ivf4_ip", "ivf32_l2", "ivf32_ip"]) def test_ivf_v1_header_uses_historical_payload_and_padding(name, tmp_path): - path, cls, _, queries, ids, distances, options = reference(name) + path, cls, count, queries, ids, distances, options = reference(name) raw = name.startswith("ivf32") # The intermediate IVF v1 layout prepends explicit padding to the same # historical payload. This synthetic header test complements frozen files; @@ -91,6 +93,9 @@ def test_ivf_v1_header_uses_historical_payload_and_padding(name, tmp_path): destination.write_bytes(header + payload) index = cls.load(str(destination)) assert index.initializer == "flat" + assert index.num_stored_points == count + assert index.num_points == count - 1 + assert index.max_elements == count actual_ids, actual_distances = index.search(queries, 5, **options) np.testing.assert_array_equal(actual_ids, ids) np.testing.assert_allclose(actual_distances, distances, rtol=2e-4, atol=2e-4) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index dba5c1a2..27be3e40 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -25,10 +25,6 @@ def built_hnsw(base_data, clusters): # ── construction ────────────────────────────────────────────────────────────── -def test_is_built(built_hnsw): - assert built_hnsw.is_built - - @pytest.mark.parametrize("m", [0, 1]) def test_rejects_too_small_m(m): with pytest.raises(ValueError, match="^HNSW M must be at least 2$"): @@ -57,6 +53,7 @@ def test_rejected_rebuild_preserves_index(base_data, clusters, tmp_path, reload) def test_properties(built_hnsw): + assert built_hnsw.is_built assert built_hnsw.dim == DIM assert built_hnsw.nbits == 4 assert built_hnsw.metric == "l2" @@ -85,32 +82,24 @@ def test_parallel_build_preserves_self_retrieval(base_data, clusters): # ── search output shape and dtype ───────────────────────────────────────────── -def test_search_output_shape(built_hnsw, query_data): +def test_search_output_contract(built_hnsw, query_data): ids, dists = built_hnsw.search(query_data, k=_TOPK, ef=_EF) assert ids.shape == (N_QUERIES, _TOPK) assert dists.shape == (N_QUERIES, _TOPK) - - -def test_search_output_dtype(built_hnsw, query_data): - ids, dists = built_hnsw.search(query_data, k=_TOPK, ef=_EF) - assert np.issubdtype(ids.dtype, np.integer) + assert ids.dtype == np.uint32 assert dists.dtype == np.float32 - - -def test_search_ids_in_range(built_hnsw, query_data): - ids, _ = built_hnsw.search(query_data, k=_TOPK, ef=_EF) assert np.all(ids < N_VECTORS) - - -def test_search_distances_nonneg(built_hnsw, query_data): - _, dists = built_hnsw.search(query_data, k=_TOPK, ef=_EF) assert np.all(dists >= 0) -def test_ef_default(built_hnsw, query_data): +@pytest.mark.parametrize("k", [1, 10, 17]) +def test_ef_default(built_hnsw, query_data, k): """ef=0 should use the internal default (max(k, 10)).""" - ids, dists = built_hnsw.search(query_data, k=_TOPK) - assert ids.shape == (N_QUERIES, _TOPK) + expected_ids, expected_distances = built_hnsw.search(query_data, k=k, ef=max(k, 10)) + for kwargs in ({}, {"ef": 0}): + ids, distances = built_hnsw.search(query_data, k=k, **kwargs) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) def test_single_query(built_hnsw, query_data): @@ -145,7 +134,7 @@ def test_wrong_data_dim_raises(clusters): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^data dimension does not match index dim$"): idx.build(bad_data, centroids, cluster_ids) @@ -160,13 +149,15 @@ def test_out_of_range_cluster_id_raises(base_data, clusters): def test_wrong_query_dim_raises(built_hnsw): bad_queries = np.zeros((5, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^query dimension does not match index dim$"): built_hnsw.search(bad_queries, k=1) def test_search_before_build_raises(query_data): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) - with pytest.raises(Exception): + with pytest.raises( + RuntimeError, match="^HnswIndex must be built or loaded before search$" + ): idx.search(query_data[:1], k=1) @@ -250,12 +241,11 @@ def test_add_into_a_dense_graph_retrieves_every_point(base_data, clusters): np.testing.assert_array_equal(ids[:, 0], np.arange(400, N_VECTORS)) -def test_add_without_cluster_ids_routes(partial_hnsw, base_data, clusters): - _, cluster_ids = clusters +def test_add_without_cluster_ids_preserves_self_retrieval(partial_hnsw, base_data): ids = partial_hnsw.add(base_data[400:]) np.testing.assert_array_equal(ids, np.arange(400, N_VECTORS)) - # Round-robin clusters, so routing should recover the original assignment for - # most points; the centroids are close together and a few land elsewhere. + # Automatically routed points must remain searchable even though nearest-centroid + # routing need not reproduce the round-robin assignments used during construction. found, _ = partial_hnsw.search(base_data[400:], k=1, ef=N_VECTORS) assert np.count_nonzero(found[:, 0] == np.arange(400, N_VECTORS)) >= 0.95 * 100 @@ -382,10 +372,11 @@ def _removable(nbits=4, metric="l2", count=120, dim=64, seed=61): def _everything(index, queries): """With k and ef at the point count, search visits every reachable point.""" - return index.search(queries, k=index.num_points, ef=index.num_points) + return index.search(queries, k=index.num_stored_points, ef=index.num_stored_points) -@pytest.mark.parametrize("nbits", [1, 2, 4, 8, 9]) +# C++ HnswRemoveTest covers intermediate widths; keep both storage endpoints here. +@pytest.mark.parametrize("nbits", [1, 9]) @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_remove_excludes_points_and_survives_reload(tmp_path, nbits, metric): idx, old, new = _removable(nbits=nbits, metric=metric) @@ -399,7 +390,8 @@ def test_remove_excludes_points_and_survives_reload(tmp_path, nbits, metric): assert idx.remove(removed) == 4 assert idx.remove(removed) == 0 assert idx.remove([]) == 0 - assert idx.num_points == count + assert idx.num_points == count - 4 + assert idx.num_stored_points == count def check(index): ids, dists = _everything(index, queries) @@ -429,7 +421,8 @@ def check(index): for index in (idx, loaded, plain): index.resize(count + len(new)) index.add(new) - assert index.num_points == count + len(new) + assert index.num_stored_points == count + len(new) + assert index.num_points == count + len(new) - (0 if index is plain else 4) # The original still supports add after remove, and its resulting graph must # round-trip exactly without any further insertions. diff --git a/tests/python/test_import.py b/tests/python/test_import.py index cf07c84a..f17e344a 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -63,5 +63,5 @@ def test_symqg_construct(): ], ) def test_invalid_metric_raises(cls, args): - with pytest.raises(Exception, match="[Uu]nsupported metric|[Ii]nvalid"): + with pytest.raises(ValueError, match=r"^Unsupported metric\. Use 'l2' or 'ip'\.$"): cls(*args, metric="cosine") diff --git a/tests/python/test_index_counts.py b/tests/python/test_index_counts.py new file mode 100644 index 00000000..8dba7061 --- /dev/null +++ b/tests/python/test_index_counts.py @@ -0,0 +1,95 @@ +"""All index wrappers distinguish capacity, stored records and live points.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex + + +def check_counts(index, capacity, stored, active): + assert (index.max_elements, index.num_stored_points, index.num_points) == ( + capacity, + stored, + active, + ) + assert 0 <= index.num_points <= index.num_stored_points <= index.max_elements + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize( + "kind,bits", + [("ivf", 1), ("ivf", 4), ("ivf", 32), ("hnsw", 1), ("hnsw", 9)] + + [("symqg", bits) for bits in (0, 4, 8)], +) +def test_count_contract_through_lifecycle(tmp_path, kind, bits, metric): + count, dim = 40, 65 + data = np.random.default_rng(42).normal(size=(count + 4, dim)).astype(np.float32) + data /= np.linalg.norm(data, axis=1, keepdims=True) + centroid = np.zeros((1, dim), np.float32) + assignments = np.zeros(count + 4, np.uint32) + if kind == "ivf": + index = IvfIndex(dim, count, 1, bits, metric=metric) + capacity = count + elif kind == "hnsw": + index = HnswIndex(dim, 64, M=8, ef_construction=40, nbits=bits, metric=metric) + capacity = 64 + else: + index = SymqgIndex(dim, max_degree=32, quantization_bits=bits, metric=metric) + capacity = 0 + check_counts(index, capacity, 0, 0) + if kind == "symqg": + index.build(data[:count], ef_construction=40, num_threads=1) + capacity = count + else: + index.build(data[:count], centroid, assignments[:count]) + check_counts(index, capacity, count, count) + + path = tmp_path / "counts.index" + if kind == "symqg": + index.save(str(path)) + check_counts(SymqgIndex.load(str(path)), count, count, count) + return + + assert index.remove([0, 1, 0]) == 2 + assert index.remove([0, 1]) == 0 + check_counts(index, capacity, count, count - 2) + # A live-count bound would incorrectly reject this valid high ID. + assert index.remove([count - 1]) == 1 + check_counts(index, capacity, count, count - 3) + with pytest.raises(ValueError, match="ids"): + index.remove([2, count]) + check_counts(index, capacity, count, count - 3) + if kind == "hnsw": + with pytest.raises(ValueError): + index.resize(count - 3) + index.resize(80) + capacity = 80 + check_counts(index, capacity, count, count - 3) + + ids = index.add(data[count : count + 3], assignments[:3]) + np.testing.assert_array_equal(ids, np.arange(count, count + 3)) + stored = count + 3 + if kind == "ivf": + capacity = stored + check_counts(index, capacity, stored, count) + index.save(str(path)) + index = type(index).load(str(path)) + check_counts(index, capacity, stored, count) + assert index.remove([0, 1, count - 1]) == 0 + check_counts(index, capacity, stored, count) + + assert index.remove(np.arange(stored)) == count + check_counts(index, capacity, stored, 0) + index.save(str(path)) + index = type(index).load(str(path)) + check_counts(index, capacity, stored, 0) + args = {"nprobe": 1} if kind == "ivf" else {"ef": stored} + hits, _ = index.search(data[:1], k=1, **args) + assert np.all(hits == np.iinfo(np.uint32).max) + + np.testing.assert_array_equal(index.add(data[-1:], assignments[:1]), [stored]) + stored += 1 + if kind == "ivf": + capacity = stored + check_counts(index, capacity, stored, 1) + index.save(str(path)) + check_counts(type(index).load(str(path)), capacity, stored, 1) diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 4323b376..163768c2 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -25,11 +25,8 @@ def built_ivf(base_data, clusters): # ── construction ────────────────────────────────────────────────────────────── -def test_is_built(built_ivf): - assert built_ivf.is_built - - def test_properties(built_ivf): + assert built_ivf.is_built assert built_ivf.dim == DIM assert built_ivf.num_clusters == N_CLUSTERS assert built_ivf.nbits == 4 @@ -47,25 +44,13 @@ def test_fast_quantization_builds(base_data, clusters): # ── search output shape and dtype ───────────────────────────────────────────── -def test_search_output_shape(built_ivf, query_data): +def test_search_output_contract(built_ivf, query_data): ids, dists = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL) assert ids.shape == (N_QUERIES, _TOPK) assert dists.shape == (N_QUERIES, _TOPK) - - -def test_search_output_dtype(built_ivf, query_data): - ids, dists = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL) - assert np.issubdtype(ids.dtype, np.integer) + assert ids.dtype == np.uint32 assert dists.dtype == np.float32 - - -def test_search_ids_in_range(built_ivf, query_data): - ids, _ = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL) assert np.all(ids < N_VECTORS) - - -def test_search_distances_nonneg(built_ivf, query_data): - _, dists = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL) assert np.all(dists >= 0) @@ -75,7 +60,9 @@ def test_single_query(built_ivf, query_data): assert dists.shape == (1, 1) -@pytest.mark.parametrize("bits", [1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +# C++ BatchSearchMatchesIndividualQueries covers every width; exercise binding +# batching for sign-only, extra-bit and owned-float storage here. +@pytest.mark.parametrize("bits", [1, 4, 32]) @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_batch_search_matches_individual_queries(bits, metric): rng = np.random.default_rng(73) @@ -170,13 +157,6 @@ def test_inner_product_routing_uses_inner_product(): # ── optional parameters ─────────────────────────────────────────────────────── -def test_high_accuracy_false(built_ivf, query_data): - ids, dists = built_ivf.search( - query_data, k=_TOPK, nprobe=_NPROBE_ALL, high_accuracy=False - ) - assert ids.shape == (N_QUERIES, _TOPK) - - @pytest.mark.parametrize("nbits", [4, 8]) @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_full_quantization_centroid_distances_roundtrip(nbits, metric, tmp_path): @@ -231,7 +211,7 @@ def test_wrong_data_dim_raises(clusters): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) centroids, cluster_ids = clusters bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^data dimension does not match index dim$"): idx.build(bad_data, centroids, cluster_ids) @@ -239,7 +219,9 @@ def test_wrong_centroid_dim_raises(base_data, clusters): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) _, cluster_ids = clusters bad_centroids = np.zeros((N_CLUSTERS, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises( + ValueError, match="^centroid dimension does not match index dim$" + ): idx.build(base_data, bad_centroids, cluster_ids) @@ -247,19 +229,23 @@ def test_wrong_cluster_ids_length_raises(base_data, clusters): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) centroids, _ = clusters bad_ids = np.zeros(N_VECTORS + 5, dtype=np.uint32) - with pytest.raises(Exception): + with pytest.raises( + ValueError, match="^cluster_ids length must match number of rows in data$" + ): idx.build(base_data, centroids, bad_ids) def test_wrong_query_dim_raises(built_ivf): bad_queries = np.zeros((5, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^query dimension does not match index dim$"): built_ivf.search(bad_queries, k=1, nprobe=1) def test_search_before_build_raises(query_data): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) - with pytest.raises(Exception): + with pytest.raises( + RuntimeError, match="^IvfIndex must be built or loaded before search$" + ): idx.search(query_data[:1], k=1, nprobe=1) @@ -282,11 +268,13 @@ def test_save_load_roundtrip(built_ivf, query_data, tmp_path): np.testing.assert_allclose(dists_orig, dists_load, rtol=1e-5) -@pytest.mark.parametrize("dim", [65, 420, 960]) @pytest.mark.parametrize("metric", ["l2", "ip"]) @pytest.mark.parametrize("high_accuracy", [False, True]) @pytest.mark.parametrize("fast_quantization", [False, True]) -def test_raw_reranking(dim, metric, high_accuracy, fast_quantization, tmp_path): +def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): + # Dimension/backend arithmetic is covered in C++; this matrix checks the + # Python flags, original-vector ownership and raw persistence payload. + dim = 65 rng = np.random.default_rng(71) data = rng.standard_normal((65, dim)).astype(np.float32) data[0] = 0 # zero residual, padded dimension, and tail batches @@ -502,18 +490,18 @@ def test_hnsw_centroid_routing_roundtrip(tmp_path): # Explicit routing keeps HNSW coverage independent of the auto thresholds. # One vector per centroid makes nearest-centroid IDs and distances exact. rng = np.random.default_rng(314) - count, dim = 20000, 64 + count, dim = 128, 64 centroids = rng.standard_normal((count, dim)).astype(np.float32) labels = np.arange(count, dtype=np.uint32) idx = IvfIndex(dim, count, count, nbits=1, initializer="hnsw") idx.build(centroids, centroids, labels, num_threads=4) - selected = np.array([0, 17, 1023, count - 1]) + selected = np.array([0, 17, 63, count - 1]) queries = centroids[selected] ids, distances = idx.search(queries, k=1, nprobe=4, num_threads=4) np.testing.assert_array_equal(ids[:, 0], selected) np.testing.assert_allclose(distances, 0, rtol=0, atol=1e-5) - path = str(tmp_path / "large_centroid.index") + path = str(tmp_path / "hnsw_centroid.index") idx.save(path) loaded = IvfIndex.load(path) loaded_ids, loaded_distances = loaded.search(queries, k=1, nprobe=4, num_threads=4) @@ -522,7 +510,7 @@ def test_hnsw_centroid_routing_roundtrip(tmp_path): def test_hnsw_inner_product_routing_roundtrip(tmp_path): - count, dim = 20000, 64 + count, dim = 128, 64 centroids = np.zeros((count, dim), dtype=np.float32) centroids[:, 0] = 0.5 centroids[0, 0] = 100.0 @@ -532,7 +520,7 @@ def test_hnsw_inner_product_routing_roundtrip(tmp_path): idx = IvfIndex(dim, 1, count, nbits=32, metric="ip", initializer="hnsw") idx.build(data, centroids, np.array([0], dtype=np.uint32), num_threads=4) - path = str(tmp_path / "large_ip_centroid.index") + path = str(tmp_path / "hnsw_ip_centroid.index") ids, distances = idx.search(query, k=1, nprobe=1, num_threads=4) assert np.all(ids == 0) np.testing.assert_allclose(distances, -99.0, rtol=0, atol=1e-5) @@ -642,7 +630,8 @@ def test_add_rejects_invalid_arguments(): unbuilt.add(new[:2]) -@pytest.mark.parametrize("nbits", [1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +# IvfRemoveTest owns the full width x metric x scan-path matrix in C++. +@pytest.mark.parametrize("nbits", [1, 4, 32]) @pytest.mark.parametrize("metric", ["l2", "ip"]) @pytest.mark.parametrize("high_accuracy", [False, True]) def test_remove_excludes_points_and_survives_reload( @@ -766,7 +755,7 @@ def test_remove_and_add_accept_any_integer_dtype_and_empty_input(): assert len(idx.add(new[:3], cluster_ids=np.array([0, 1, 2], dtype=np.uint16))) == 3 -@pytest.mark.parametrize("bits", [1, 4, 9, 32]) +@pytest.mark.parametrize("bits", [1, 4, 32]) @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_rebuild_matches_incremental_encoding_with_empty_clusters( tmp_path, bits, metric diff --git a/tests/python/test_ivf_initializers.py b/tests/python/test_ivf_initializers.py index 662051a2..9ebf2b1f 100644 --- a/tests/python/test_ivf_initializers.py +++ b/tests/python/test_ivf_initializers.py @@ -1,5 +1,6 @@ """Explicit centroid routing, persistence, and legacy initializer inference.""" +import struct from pathlib import Path import numpy as np @@ -8,7 +9,48 @@ @pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("bits", list(range(1, 10)) + [32]) +@pytest.mark.parametrize("damage", ["entry", "neighbor", "level"]) +def test_hnsw_rejects_corrupt_sidecar(tmp_path, metric, damage): + count, dim = 64, 64 + data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32) + index = IvfIndex(dim, count, count, 4, metric, initializer="hnsw") + index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=1) + expected = index.search(data[:2], 1, 4, num_threads=1) + path = tmp_path / "centroids.index" + index.save(str(path)) + sidecar = Path(str(path) + ".hnsw") + original = sidecar.read_bytes() + damaged = bytearray(original) + # hnswlib writes native POD fields individually, with no inter-field padding. + size_t_bytes = struct.calcsize("@N") + level_offset = 6 * size_t_bytes + entry_offset = level_offset + struct.calcsize("@i") + if damage == "entry": + capacity = struct.unpack_from("@N", original, size_t_bytes)[0] + struct.pack_into("@I", damaged, entry_offset, capacity) + message = "Invalid HNSW centroid entry point or maximum level" + elif damage == "neighbor": + header_bytes = entry_offset + 4 + 4 * size_t_bytes + 8 + struct.pack_into("@H", damaged, header_bytes, 1) + struct.pack_into("@I", damaged, header_bytes + 4, count) + message = "Invalid HNSW centroid neighbor or target level" + else: + max_level = struct.unpack_from("@i", original, level_offset)[0] + struct.pack_into("@i", damaged, level_offset, max_level + 1) + message = "Invalid HNSW centroid level structure" + sidecar.write_bytes(damaged) + with pytest.raises(RuntimeError, match=f"^{message}$"): + IvfIndex.load(str(path)) + sidecar.write_bytes(original) + loaded = IvfIndex.load(str(path)) + actual = loaded.search(data[:2], 1, 4, num_threads=1) + for left, right in zip(expected, actual): + np.testing.assert_array_equal(left, right) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +# Routing is independent of bit width; retain each distinct reranking storage path. +@pytest.mark.parametrize("bits", [1, 4, 32]) @pytest.mark.parametrize("initializer", ["auto", "flat", "flat_rabitq", "hnsw"]) def test_initializer_roundtrip_and_updates(tmp_path, metric, bits, initializer): rng = np.random.default_rng(773) @@ -69,21 +111,35 @@ def test_auto_initializer_thresholds(clusters, expected): assert IvfIndex(64, 1, clusters, 1, initializer=mode).initializer == mode -@pytest.mark.parametrize("bits", [1, 32]) -@pytest.mark.parametrize("clusters", [3, 5000, 19999, 20000]) -@pytest.mark.parametrize("format_version", ["v1", "legacy"]) -def test_legacy_formats_infer_historical_initializer( - tmp_path, bits, clusters, format_version -): - rng = np.random.default_rng(43) - centroids = rng.standard_normal((clusters, 64)).astype(np.float32) +@pytest.fixture( + scope="module", + params=[ + (bits, clusters) for bits in (1, 32) for clusters in (3, 5000, 19999, 20000) + ], + ids=lambda case: f"{case[0]}bit-{case[1]}clusters", +) +def historical_ivf(request): + # The two format cases only read/search/save this source index. Share the + # expensive historical HNSW construction; damaged files stay per-test. + bits, clusters = request.param + centroids = ( + np.random.default_rng(43).standard_normal((clusters, 64)).astype(np.float32) + ) historical = "flat" if clusters < 20000 else "hnsw" index = IvfIndex(64, 3, clusters, bits, initializer=historical) index.build(centroids[:3], centroids, np.arange(3, dtype=np.uint32), 2) + return index, centroids[:3].copy() + + +@pytest.mark.parametrize("format_version", ["v1", "legacy"]) +def test_legacy_formats_infer_historical_initializer( + tmp_path, historical_ivf, format_version +): + index, queries = historical_ivf path = tmp_path / "v1.index" index.save(str(path)) # v2 must honor the stored type even when today's Auto would choose differently. - assert IvfIndex.load(str(path)).initializer == historical + assert IvfIndex.load(str(path)).initializer == index.initializer payload = bytearray(path.read_bytes()) if format_version == "v1": # v1 has the same payload but no initializer field after padded_dim. @@ -92,14 +148,12 @@ def test_legacy_formats_infer_historical_initializer( else: # Dimension 64 also matches the padding inferred by both older formats. payload = payload[28:] - if bits == 32: + if index.nbits == 32: payload = b"RBFQRAW1" + (1).to_bytes(4, "little") + payload path.write_bytes(payload) loaded = IvfIndex.load(str(path)) assert loaded.initializer == index.initializer - for left, right in zip( - index.search(centroids[:3], 1, 2), loaded.search(centroids[:3], 1, 2) - ): + for left, right in zip(index.search(queries, 1, 2), loaded.search(queries, 1, 2)): np.testing.assert_array_equal(left, right) diff --git a/tests/python/test_portable_indexes.py b/tests/python/test_portable_indexes.py index aa51904a..83896164 100644 --- a/tests/python/test_portable_indexes.py +++ b/tests/python/test_portable_indexes.py @@ -5,15 +5,7 @@ from rabitqlib import HnswIndex, IvfIndex, SymqgIndex -@pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("dim", [64, 65, 81, 97]) -@pytest.mark.parametrize( - "kind,bits", - [(kind, bits) for kind in ("ivf", "hnsw") for bits in range(1, 10)] - + [("ivf", 32)] - + [("qg", bits) for bits in (0, 4, 8)], -) -def test_portable_index_round_trip(tmp_path, metric, dim, kind, bits): +def check_index_round_trip(tmp_path, metric, dim, kind, bits): rng = np.random.default_rng(764) data = rng.standard_normal((96, dim)).astype(np.float32) # Unit vectors make self-query the exact winner for both metrics. @@ -72,40 +64,55 @@ def test_portable_index_round_trip(tmp_path, metric, dim, kind, bits): stored_padded = int.from_bytes(payload[24:32], "little") assert stored_padded == (dim + multiple - 1) // multiple * multiple loaded = type(index).load(str(path)) + assert loaded.dim == dim + assert loaded.metric == metric + assert (loaded.quantization_bits if kind == "qg" else loaded.nbits) == bits actual_ids, actual_distances = loaded.search(queries, k=5, **search_args) np.testing.assert_array_equal(actual_ids, ids) np.testing.assert_array_equal(actual_distances, distances) @pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("bits", range(1, 10)) -def test_hnsw_multi_block_kernel_round_trip(tmp_path, metric, bits): - # 960 coordinates exercise a full 512-bit query block and its 448-bit tail. - test_portable_index_round_trip(tmp_path, metric, 960, "hnsw", bits) - - -@pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("dim", [193, 960]) @pytest.mark.parametrize( "kind,bits", - [("ivf", 5), ("ivf", 32), ("hnsw", 5), ("qg", 0), ("qg", 4), ("qg", 8)], + [(kind, bits) for kind in ("ivf", "hnsw") for bits in range(1, 10)] + + [("ivf", 32)] + + [("qg", bits) for bits in (0, 4, 8)], ) -def test_rotation_padded_domain_round_trip(tmp_path, metric, dim, kind, bits): - test_portable_index_round_trip(tmp_path, metric, dim, kind, bits) +def test_supported_storage_round_trip(tmp_path, metric, kind, bits): + # Every public width crosses the Python boundary once per metric, with padding. + # C++ owns the width x dimension x backend numerical reference matrices. + check_index_round_trip(tmp_path, metric, 65, kind, bits) @pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("dim", [511, 513]) -@pytest.mark.parametrize("kind,bits", [("ivf", 3), ("ivf", 5), ("hnsw", 5), ("qg", 0)]) -def test_query_preparation_block_boundaries(tmp_path, metric, dim, kind, bits): - # Exercise both sides of the 512-coordinate transpose block, standard/HACC - # LUT quantization, and persistence with original dimensions requiring padding. - test_portable_index_round_trip(tmp_path, metric, dim, kind, bits) +@pytest.mark.parametrize( + "kind,bits,dim", + [ + ("ivf", 1, 64), + ("ivf", 3, 511), + ("ivf", 4, 513), + ("ivf", 9, 97), + ("ivf", 32, 960), + ("hnsw", 1, 64), + ("hnsw", 5, 511), + ("hnsw", 5, 513), + ("hnsw", 9, 960), + ("qg", 0, 511), + ("qg", 0, 513), + ("qg", 4, 97), + ("qg", 8, 193), + ], +) +def test_padded_domain_boundaries(tmp_path, metric, dim, kind, bits): + # Representative integration paths: aligned input, compact tails, both sides + # of a 512-coordinate query block, and multi-block rotation/raw reranking. + check_index_round_trip(tmp_path, metric, dim, kind, bits) @pytest.mark.parametrize("kind,bits", [("ivf", 4), ("hnsw", 4), ("qg", 4)]) def test_rejects_development_only_padding(tmp_path, kind, bits): - test_portable_index_round_trip(tmp_path, "l2", 65, kind, bits) + check_index_round_trip(tmp_path, "l2", 65, kind, bits) path = tmp_path / "portable.index" payload = bytearray(path.read_bytes()) offset = {"ivf": 16, "hnsw": 24, "qg": 12 + 3 * np.dtype(np.uintp).itemsize}[kind] diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py index 5c08b757..819422cb 100644 --- a/tests/python/test_qgkmeans.py +++ b/tests/python/test_qgkmeans.py @@ -206,12 +206,6 @@ def test_spherical_mode_returns_normalized_centroids(base_data): np.testing.assert_allclose(clustering.distances, selected_distances, atol=1e-6) -def test_train_rejects_wrong_dimension(): - clustering = qgkmeans() - with pytest.raises(Exception, match="dimension"): - clustering.train(np.zeros((K, DIM + 1), dtype=np.float32)) - - @pytest.mark.parametrize("bits", [0, 4, 8]) @pytest.mark.parametrize("spherical", [False, True]) def test_graph_returns_exact_selected_distances(bits, spherical): @@ -352,9 +346,8 @@ def test_single_empty_cluster_preserves_duplicate_centroids(bits, threads): assert clustering.final_obj == 0 -@pytest.mark.parametrize("bits", [0, 4, 8]) @pytest.mark.parametrize("spherical", [False, True]) -def test_final_objective_does_not_regress_after_last_update(bits, spherical): +def test_final_objective_does_not_regress_after_last_update(spherical): x = np.random.default_rng(739).standard_normal((2048, 65)).astype(np.float32) if spherical: x /= np.linalg.norm(x, axis=1, keepdims=True) @@ -364,7 +357,8 @@ def test_final_objective_does_not_regress_after_last_update(bits, spherical): niter=4, num_threads=2, min_points_per_centroid=1, - quantization_bits=bits, + # C++ checks every storage width with a genuinely approximate graph. + quantization_bits=4, spherical=spherical, ) clustering.train(x) @@ -525,7 +519,9 @@ def test_approximate_final_assignment_keeps_graph_mode_alias(): @pytest.mark.parametrize( "clusters,degree", - [(1, 32), (2, 32), (16, 32), (31, 32), (32, 32), (33, 64), (64, 64)], + # C++ checks the intermediate counts; these exercise below/equal boundaries + # and propagation of a non-default graph degree through Python. + [(1, 32), (32, 32), (64, 64)], ) @pytest.mark.parametrize( "mode", [FinalAssignmentMode.Approximate, FinalAssignmentMode.Exact] diff --git a/tests/python/test_rabitqkmeans.py b/tests/python/test_rabitqkmeans.py index c2f87e1b..ab1c198f 100644 --- a/tests/python/test_rabitqkmeans.py +++ b/tests/python/test_rabitqkmeans.py @@ -70,7 +70,9 @@ def test_flat_native_centroids_before_training_raise(): _ = model.centroids -@pytest.mark.parametrize("clusters", [1, 2, 16, 31, 32, 33, 65]) +# C++ owns the complete cluster-count/batch-tail matrix. Python retains the +# single-cluster path, a partial second batch, both modes and mutable-parameter refit. +@pytest.mark.parametrize("clusters", [1, 33]) @pytest.mark.parametrize("spherical", [False, True]) @pytest.mark.parametrize("exact", [False, True]) def test_flat_train_refit_and_return_exact_selected_distances( diff --git a/tests/python/test_remove_arrays.py b/tests/python/test_remove_arrays.py new file mode 100644 index 00000000..d53b4c6b --- /dev/null +++ b/tests/python/test_remove_arrays.py @@ -0,0 +1,54 @@ +"""Removal accepts contiguous integer views without assuming pointer alignment.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +@pytest.mark.parametrize("offset", [0, 1]) +@pytest.mark.parametrize("readonly", [False, True]) +def test_remove_integer_views_preserve_input_and_validate_atomically( + kind, offset, readonly +): + count, dim = 8, 64 + data = np.random.default_rng(42).normal(size=(count, dim)).astype(np.float32) + index = ( + IvfIndex(dim, count, 1, 4) + if kind == "ivf" + else HnswIndex(dim, count, M=2, ef_construction=10, nbits=4) + ) + index.build(data, np.zeros((1, dim), np.float32), np.zeros(count, np.uint32)) + search_args = {"nprobe": 1} if kind == "ivf" else {"ef": count} + before = index.search(data, k=1, **search_args) + + def make_ids(values): + result = np.ndarray( + len(values), + dtype=np.int64, + buffer=bytearray(8 * len(values) + offset), + offset=offset, + ) + result[:] = values + result.flags.writeable = not readonly + assert result.flags.c_contiguous + assert result.flags.aligned == (offset == 0) + return result + + bad = make_ids([0, count]) + with pytest.raises(ValueError, match=r"^ids must be in \[0, num_stored_points\)$"): + index.remove(bad) + assert index.num_points == count + for actual, expected in zip( + index.search(data, k=1, **search_args), before, strict=True + ): + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(bad, [0, count]) + + ids = make_ids([0, 3, 0]) + assert index.remove(ids) == 2 + assert index.remove(ids) == 0 + assert index.num_points == count - 2 + assert index.num_stored_points == count + assert not np.isin(index.search(data, k=1, **search_args)[0], [0, 3]).any() + np.testing.assert_array_equal(ids, [0, 3, 0]) diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 2d8f3573..d44d6dfb 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -27,11 +27,8 @@ def built_symqg(base_data): # ── construction ────────────────────────────────────────────────────────────── -def test_is_built(built_symqg): - assert built_symqg.is_built - - def test_properties(built_symqg): + assert built_symqg.is_built assert built_symqg.dim == DIM assert built_symqg.max_degree == _MAX_DEGREE assert built_symqg.num_points == N_VECTORS @@ -42,25 +39,13 @@ def test_properties(built_symqg): # ── search output shape and dtype ───────────────────────────────────────────── -def test_search_output_shape(built_symqg, query_data): +def test_search_output_contract(built_symqg, query_data): ids, dists = built_symqg.search(query_data, k=_TOPK, ef=_EF) assert ids.shape == (N_QUERIES, _TOPK) assert dists.shape == (N_QUERIES, _TOPK) - - -def test_search_output_dtype(built_symqg, query_data): - ids, dists = built_symqg.search(query_data, k=_TOPK, ef=_EF) - assert np.issubdtype(ids.dtype, np.integer) + assert ids.dtype == np.uint32 assert dists.dtype == np.float32 - - -def test_search_ids_in_range(built_symqg, query_data): - ids, _ = built_symqg.search(query_data, k=_TOPK, ef=_EF) assert np.all(ids < N_VECTORS) - - -def test_search_distances_nonneg(built_symqg, query_data): - _, dists = built_symqg.search(query_data, k=_TOPK, ef=_EF) assert np.all(dists >= 0) @@ -71,42 +56,47 @@ def test_single_query(built_symqg, query_data): @pytest.mark.parametrize("bits", [0, 4, 8]) @pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("method", ["search", "search_batch"]) -def test_batch_search_matches_individual_queries(tmp_path, bits, metric, method): +def test_batch_search_matches_individual_queries(tmp_path, bits, metric): rng = np.random.default_rng(73) data = rng.standard_normal((97, 65)).astype(np.float32) - queries = rng.standard_normal((19, 65)).astype(np.float32) + queries = rng.standard_normal((19, 130)).astype(np.float32)[:, ::2] + assert not queries.flags.c_contiguous queries[4] = queries[0] index = SymqgIndex(65, 32, metric=metric, quantization_bits=bits) index.build(data, ef_construction=48, num_threads=1) path = str(tmp_path / "batch.index") index.save(path) restored = SymqgIndex.load(path) - search = getattr(restored, method) - for k, ef in [(7, 64), (1, 1), (7, 64)]: - individual = [index.search(query[None, :], k=k, ef=ef) for query in queries] + # C++ owns the full ef/thread matrix. These scenarios retain the two aliases, + # serial/automatic/parallel scheduling, changing windows, and strided inputs + # while sharing a single build for each metric/storage pair. + for k, ef, threads in [ + (7, 64, 1), + (1, 1, 4), + (7, 64, 4), + (10, 16, 0), + (10, 128, 2), + (1, 64, 8), + ]: + individual = [ + index.search(query.copy()[None, :], k=k, ef=ef) for query in queries + ] expected_ids = np.concatenate([result[0] for result in individual]) expected_distances = np.concatenate([result[1] for result in individual]) - for threads in [1, 4]: + for method in ("search", "search_batch"): + search = getattr(restored, method) ids, distances = search(queries, k=k, ef=ef, num_threads=threads) np.testing.assert_array_equal(ids, expected_ids) - np.testing.assert_array_equal(distances, expected_distances) + np.testing.assert_array_equal( + distances.view(np.uint32), expected_distances.view(np.uint32) + ) empty_ids, empty_distances = search( queries[:0], k=k, ef=ef, num_threads=threads ) assert empty_ids.shape == empty_distances.shape == (0, k) -def test_search_batch_accepts_noncontiguous_queries(built_symqg, query_data): - queries = query_data[::2] - assert not queries.flags.c_contiguous - expected = built_symqg.search(queries.copy(), k=5, ef=_EF) - actual = built_symqg.search_batch(queries, k=5, ef=_EF, num_threads=0) - np.testing.assert_array_equal(actual[0], expected[0]) - np.testing.assert_array_equal(actual[1], expected[1]) - - def test_search_batch_validates_inputs(built_symqg, query_data): with pytest.raises(ValueError, match="queries must be a 2D NumPy array"): built_symqg.search_batch(query_data[0], k=1, ef=_EF) @@ -150,20 +140,22 @@ def test_recall_vs_brute_force(built_symqg, base_data, query_data): def test_wrong_data_dim_raises(): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^data dimension does not match index dim$"): idx.build(bad_data, ef_construction=_EF_BUILD) def test_wrong_query_dim_raises(built_symqg): bad_queries = np.zeros((5, DIM + 1), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^query dimension does not match index dim$"): built_symqg.search(bad_queries, k=1, ef=_EF) def test_search_before_build_raises(): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) queries = np.zeros((5, DIM), dtype=np.float32) - with pytest.raises(Exception): + with pytest.raises( + RuntimeError, match="^SymqgIndex must be built or loaded before search$" + ): idx.search(queries, k=1, ef=_EF) @@ -237,12 +229,14 @@ def test_successful_rebuild_replaces_index(base_data, tmp_path, bits): def test_invalid_degree_raises(): - with pytest.raises(Exception): + with pytest.raises( + ValueError, match="^max_degree must be a positive multiple of 32$" + ): SymqgIndex(DIM, max_degree=16) def test_invalid_quantization_bits_raises(): - with pytest.raises(Exception): + with pytest.raises(ValueError, match="^quantization_bits must be 0, 4, or 8$"): SymqgIndex(DIM, max_degree=_MAX_DEGREE, quantization_bits=6) diff --git a/tests/python/test_symqg_batch.py b/tests/python/test_symqg_batch.py deleted file mode 100644 index e9ade38d..00000000 --- a/tests/python/test_symqg_batch.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Batch scheduling must preserve scalar search, including quantized indexes.""" - -import numpy as np -import pytest -from rabitqlib import SymqgIndex - - -@pytest.mark.parametrize("bits", [0, 4, 8]) -@pytest.mark.parametrize("metric", ["l2", "ip"]) -def test_parallel_batch_matches_scalar(bits, metric, tmp_path): - rng = np.random.default_rng(42) - data = rng.normal(size=(257, 65)).astype(np.float32) - queries = rng.normal(size=(19, 65)).astype(np.float32) - idx = SymqgIndex(65, 32, metric, bits) - idx.build(data, 64, 1) - for ef in (16, 64, 128): - for k in (1, 10): - expected_ids, expected_distances = idx.search(queries, k, ef, 1) - for threads in (0, 2, 4, 8): - ids, distances = idx.search(queries, k, ef, threads) - np.testing.assert_array_equal(ids, expected_ids) - np.testing.assert_array_equal( - distances.view(np.uint32), expected_distances.view(np.uint32) - ) - path = str(tmp_path / "batch.index") - idx.save(path) - loaded = SymqgIndex.load(path) - expected = idx.search(queries, 10, 64, 1) - actual = loaded.search(queries, 10, 64, 4) - for a, b in zip(actual, expected, strict=True): - np.testing.assert_array_equal(a, b) - - -def test_empty_and_noncontiguous_batches(base_data, query_data): - built_symqg = SymqgIndex(base_data.shape[1], 32) - built_symqg.build(base_data, 64, 1) - for threads in (0, 1, 4): - ids, distances = built_symqg.search(query_data[:0], 5, 64, threads) - assert ids.shape == distances.shape == (0, 5) - subset = query_data[::2] - actual = built_symqg.search(subset, 5, 64, threads) - expected = built_symqg.search(np.ascontiguousarray(subset), 5, 64, 1) - for a, b in zip(actual, expected, strict=True): - np.testing.assert_array_equal(a, b) diff --git a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp index 51682852..1d079b79 100644 --- a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp +++ b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp @@ -6,6 +6,7 @@ #include "rabitqlib/clustering/detail/kmeans.hpp" #include "rabitqlib/clustering/qgkmeans.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::clustering::detail { namespace { diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index 4dc86a45..aefd6390 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -1,6 +1,7 @@ #include "rabitqlib/clustering/qgkmeans.hpp" #include +#include #include #include @@ -14,16 +15,117 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/query.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::qgkmeans { namespace { +TEST(QGKMeansQuerySumTest, CachePolicyAndSuccessfulFillControlValidity) { + struct Case { + size_t dim; + size_t iterations; + FinalAssignmentMode final_assignment; + bool enabled; + }; + for (const auto configuration : + {Case{960, 25, FinalAssignmentMode::Approximate, false}, + Case{1024, 1, FinalAssignmentMode::Exact, false}, + Case{1024, 1, FinalAssignmentMode::Approximate, true}, + Case{1024, 2, FinalAssignmentMode::Exact, true}}) { + QGKMeansParameters parameters; + parameters.niter = configuration.iterations; + parameters.final_assignment = configuration.final_assignment; + detail::QuerySumCache cache(configuration.dim, parameters); + EXPECT_EQ(cache.cached(0), nullptr); + if (!configuration.enabled) { + // Disabled preparation must not attempt an allocation, even at this size. + EXPECT_NO_THROW(cache.prepare(std::numeric_limits::max())); + cache.store(0, 1.0F); + cache.complete(); + EXPECT_EQ(cache.cached(0), nullptr); + continue; + } + cache.prepare(2); + cache.store(0, 1.0F); + // Partial data remains unavailable until the whole assignment succeeds. + EXPECT_EQ(cache.cached(0), nullptr); + cache.prepare(2); + cache.store(0, 2.0F); + cache.store(1, 3.0F); + cache.complete(); + const float* data = cache.cached(0); + ASSERT_NE(data, nullptr); + EXPECT_FLOAT_EQ(data[0], 2.0F); + EXPECT_FLOAT_EQ(data[1], 3.0F); + cache.prepare(2); + cache.store(0, 99.0F); + EXPECT_EQ(cache.cached(0), data); + EXPECT_FLOAT_EQ(data[0], 2.0F); + } + detail::QuerySumCache disabled(nullptr, true); + disabled.complete(); + EXPECT_EQ(disabled.cached(0), nullptr); +} + +TEST(QGAssignerTest, CallerOwnedQuerySumsMatchUncachedAcrossRebuilds) { + constexpr size_t kClusters = 33, kPoints = 64; + for (size_t dim : {64U, 1024U}) { + std::vector data(kPoints * dim), centroids(kClusters * dim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (uint32_t bits : {0U, 4U, 8U}) { + QGAssigner cached(dim, kClusters, 32, 40, 8, 1, 2, 42, metric, bits); + QGAssigner reference(dim, kClusters, 32, 40, 8, 1, 2, 42, metric, bits); + std::vector sums(kPoints, std::numeric_limits::quiet_NaN()); + std::vector ids(kPoints), expected_ids(kPoints); + std::vector distances(kPoints), expected_distances(kPoints); + std::vector first_sums; + for (size_t iteration = 0; iteration < 3; ++iteration) { + for (size_t i = 0; i < centroids.size(); ++i) { + centroids[i] = std::cos(static_cast(i + iteration) * 0.17F); + } + cached.assign( + centroids.data(), + data.data(), + kPoints, + ids.data(), + distances.data(), + iteration == 2 ? nullptr : sums.data(), + iteration != 0 + ); + reference.assign( + centroids.data(), + data.data(), + kPoints, + expected_ids.data(), + expected_distances.data() + ); + EXPECT_EQ(ids, expected_ids); + EXPECT_EQ(distances, expected_distances); + if (iteration == 0) { + for (float sum : sums) { + EXPECT_TRUE(std::isfinite(sum)); + } + first_sums = sums; + } else { + EXPECT_EQ(sums, first_sums); + } + } + } + } + } +} + TEST(QGKMeansTest, UsesTwentyFiveIterationsByDefault) { EXPECT_EQ(QGKMeansParameters{}.niter, 25U); } diff --git a/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp index 79b972a5..a13583dc 100644 --- a/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp @@ -12,12 +12,13 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::rabitqkmeans { diff --git a/tests/unit/rabitqlib/fastscan/api_test.cpp b/tests/unit/rabitqlib/fastscan/api_test.cpp new file mode 100644 index 00000000..6a8d081c --- /dev/null +++ b/tests/unit/rabitqlib/fastscan/api_test.cpp @@ -0,0 +1,581 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/lut.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib { +namespace { + +template +void expect_invalid(Call call, const char* message) { + try { + call(); + FAIL() << "Expected std::invalid_argument"; + } catch (const std::invalid_argument& error) { EXPECT_STREQ(error.what(), message); } +} + +TEST(FastScanApi, RejectsUnpreparedAndMismatchedQueriesBeforeWriting) { + constexpr size_t kDim = 64; + std::array values{}; + values.fill(1); + std::vector batch(BatchDataMap::data_bytes(kDim)); + std::array distance, lower, ip; + distance.fill(123); + lower.fill(123); + ip.fill(123); + constexpr auto* kShapeError = + "FastScan scoring requires a prepared query with matching dimension"; + constexpr auto* kModeError = "FastScan scoring mode must match the prepared query"; + + SplitBatchQuery empty_split; + EXPECT_EQ(empty_split.rotated_query(), nullptr); + expect_invalid( + [&] { + split_batch_estdist( + batch.data(), + empty_split, + kDim, + distance.data(), + lower.data(), + ip.data(), + true + ); + }, + kShapeError + ); + + BatchQuery query; + expect_invalid( + [&] { qg_batch_estdist(batch.data(), query, kDim, distance.data()); }, kShapeError + ); + std::array scratch{}; + query.reset(values.data(), kDim, scratch.data()); + for (size_t dim : {0U, 16U, 32U, 63U, 128U}) { + expect_invalid( + [&] { qg_batch_estdist(batch.data(), query, dim, distance.data()); }, + kShapeError + ); + } + auto moved_query = std::move(query); + expect_invalid( + [&] { qg_batch_estdist(batch.data(), query, kDim, distance.data()); }, kShapeError + ); + EXPECT_EQ(moved_query.padded_dim(), kDim); + + for (bool hacc : {false, true}) { + SplitBatchQuery split(values.data(), kDim, 3, METRIC_L2, hacc); + expect_invalid( + [&] { + split_batch_estdist( + batch.data(), + split, + kDim, + distance.data(), + lower.data(), + ip.data(), + !hacc + ); + }, + kModeError + ); + for (size_t dim : {0U, 32U, 128U}) { + expect_invalid( + [&] { + split_batch_estdist( + batch.data(), + split, + dim, + distance.data(), + lower.data(), + ip.data(), + hacc + ); + }, + kShapeError + ); + } + auto moved_split = std::move(split); + expect_invalid( + [&] { + split_batch_estdist( + batch.data(), + split, + kDim, + distance.data(), + lower.data(), + ip.data(), + hacc + ); + }, + kShapeError + ); + EXPECT_EQ(moved_split.padded_dim(), kDim); + EXPECT_EQ(moved_split.use_hacc(), hacc); + } + for (size_t i = 0; i < distance.size(); ++i) { + EXPECT_EQ(distance[i], 123); + EXPECT_EQ(lower[i], 123); + EXPECT_EQ(ip[i], 123); + } + EXPECT_NO_THROW(qg_batch_estdist(batch.data(), moved_query, kDim, distance.data())); +} + +TEST(FastScanApi, SplitQueryResetMatchesFreshPreparationAndScoring) { + std::array, 2> queries{}; + std::array scratch{}; + std::array hacc_scratch{}; + SplitBatchQuery reused; + struct Case { + size_t dim; + size_t ex_bits; + MetricType metric; + bool hacc; + }; + const std::array cases{{ + {64, 0, METRIC_L2, false}, + {96, 3, METRIC_IP, true}, + {32, 8, METRIC_IP, false}, + {128, 1, METRIC_L2, true}, + }}; + for (size_t step = 0; step < cases.size(); ++step) { + const auto [dim, ex_bits, metric, hacc] = cases[step]; + SCOPED_TRACE(step); + auto& values = queries[step % queries.size()]; + std::vector data(17 * dim), centroid(dim); + for (size_t i = 0; i < dim; ++i) { + values[i] = std::cos(static_cast(i + step) * 0.13F); + centroid[i] = 0.1F * std::sin(static_cast(i) * 0.23F); + } + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.17F); + } + SplitBatchQuery fresh(values.data(), dim, ex_bits, metric, hacc); + reused.set_g_add(3, 2); + reused.reset( + values.data(), + dim, + ex_bits, + scratch.data(), + hacc ? hacc_scratch.data() : nullptr, + metric, + hacc + ); + EXPECT_EQ(reused.rotated_query(), values.data()); + EXPECT_EQ(reused.padded_dim(), dim); + EXPECT_EQ(reused.use_hacc(), hacc); + EXPECT_EQ(reused.g_add(), 0); + EXPECT_EQ(reused.g_error(), 0); + EXPECT_EQ(reused.delta(), fresh.delta()); + EXPECT_EQ(reused.sum_vl_lut(), fresh.sum_vl_lut()); + EXPECT_EQ(reused.k1xsumq(), fresh.k1xsumq()); + EXPECT_EQ(reused.kbxsumq(), fresh.kbxsumq()); + const size_t bytes = dim * 4 * (hacc ? 2 : 1); + EXPECT_TRUE(std::equal(fresh.lut(), fresh.lut() + bytes, reused.lut())); + scratch.fill(123); + hacc_scratch.fill(123); + + const float norm = std::sqrt(euclidean_sqr(values.data(), centroid.data(), dim)); + const float ip = dot_product(values.data(), centroid.data(), dim); + fresh.set_g_add(norm, ip); + reused.set_g_add(norm, ip); + EXPECT_EQ(reused.g_add(), metric == METRIC_L2 ? norm * norm : -ip); + EXPECT_EQ(reused.g_error(), norm); + std::vector batch(BatchDataMap::data_bytes(dim)); + const size_t stride = ExDataMap::data_bytes(dim, ex_bits); + std::vector extra(32 * stride); + quant::quantize_split_batch( + data.data(), + centroid.data(), + 17, + dim, + ex_bits, + batch.data(), + extra.data(), + metric + ); + std::array fresh_dist{}, fresh_low{}, fresh_ip{}; + std::array reused_dist{}, reused_low{}, reused_ip{}; + split_batch_estdist( + batch.data(), + fresh, + dim, + fresh_dist.data(), + fresh_low.data(), + fresh_ip.data(), + hacc + ); + split_batch_estdist( + batch.data(), + reused, + dim, + reused_dist.data(), + reused_low.data(), + reused_ip.data(), + hacc + ); + for (size_t row = 0; row < 17; ++row) { + EXPECT_EQ(reused_dist[row], fresh_dist[row]); + EXPECT_EQ(reused_low[row], fresh_low[row]); + EXPECT_EQ(reused_ip[row], fresh_ip[row]); + if (ex_bits != 0) { + const auto kernel = select_excode_ipfunc(ex_bits); + EXPECT_EQ( + split_distance_boosting( + extra.data() + row * stride, + kernel, + reused, + dim, + ex_bits, + reused_ip[row] + ), + split_distance_boosting( + extra.data() + row * stride, + kernel, + fresh, + dim, + ex_bits, + fresh_ip[row] + ) + ); + } + } + } +} + +TEST(FastScanApi, SplitQueryInitializationValidatesParametersConsistently) { + std::array values{}; + values.fill(0.5F); + std::array scratch{}; + std::array hacc_scratch{}; + SplitBatchQuery query; + for (bool hacc : {false, true}) { + for (size_t bits : {size_t{9}, std::numeric_limits::max()}) { + const auto* message = "SplitBatchQuery requires ex_bits in [0, 8]"; + expect_invalid( + [&] { SplitBatchQuery invalid(nullptr, 64, bits, METRIC_L2, hacc); }, + message + ); + expect_invalid( + [&] { query.reset(nullptr, 64, bits, nullptr, nullptr, METRIC_L2, hacc); }, + message + ); + } + for (size_t bits : {0U, 8U}) { + const auto invalid_metric = static_cast(7); + const auto* message = "Metric type must be METRIC_L2 or METRIC_IP"; + expect_invalid( + [&] { + SplitBatchQuery invalid(nullptr, 64, bits, invalid_metric, hacc); + }, + message + ); + expect_invalid( + [&] { + query.reset(nullptr, 64, bits, nullptr, nullptr, invalid_metric, hacc); + }, + message + ); + for (auto metric : {METRIC_L2, METRIC_IP}) { + SplitBatchQuery fresh(values.data(), 64, bits, metric, hacc); + query.reset( + values.data(), + 64, + bits, + scratch.data(), + hacc_scratch.data(), + metric, + hacc + ); + EXPECT_EQ(query.kbxsumq(), fresh.kbxsumq()); + EXPECT_EQ(query.use_hacc(), fresh.use_hacc()); + } + } + } +} + +TEST(FastScanApi, RejectedSplitQueryResetPreservesPreparedState) { + std::array values{}, next{}; + values.fill(0.5F); + next.fill(0.25F); + std::array scratch{}; + std::array hacc_scratch{}; + SplitBatchQuery query(values.data(), 64, 3, METRIC_IP, false); + query.set_g_add(3, 2); + const std::vector bytes(query.lut(), query.lut() + 64 * 4); + const auto original = query; + const auto reset = + [&](const float* input, size_t dim, size_t bits, float* lut, uint16_t* u16) { + query.reset(input, dim, bits, lut, u16, METRIC_L2, true); + }; + for (size_t bits : {size_t{9}, std::numeric_limits::max()}) { + expect_invalid( + [&] { reset(next.data(), 64, bits, scratch.data(), hacc_scratch.data()); }, + "SplitBatchQuery requires ex_bits in [0, 8]" + ); + } + expect_invalid( + [&] { + query.reset( + next.data(), + 64, + 3, + scratch.data(), + hacc_scratch.data(), + static_cast(7), + true + ); + }, + "Metric type must be METRIC_L2 or METRIC_IP" + ); + expect_invalid( + [&] { reset(nullptr, 64, 3, scratch.data(), hacc_scratch.data()); }, + "Lut::reset requires query and LUT scratch" + ); + expect_invalid( + [&] { reset(next.data(), 64, 3, nullptr, hacc_scratch.data()); }, + "Lut::reset requires query and LUT scratch" + ); + expect_invalid( + [&] { reset(next.data(), 64, 3, scratch.data(), nullptr); }, + "Lut::reset requires uint16 scratch for high-accuracy lookup tables" + ); + for (size_t dim : {0U, 17U}) { + expect_invalid( + [&] { reset(next.data(), dim, 3, scratch.data(), hacc_scratch.data()); }, + "FastScan dimension must be a positive multiple of 16" + ); + } + EXPECT_THROW( + reset( + next.data(), + std::numeric_limits::max() - 15, + 3, + scratch.data(), + hacc_scratch.data() + ), + std::length_error + ); + EXPECT_EQ(query.rotated_query(), original.rotated_query()); + EXPECT_EQ(query.padded_dim(), original.padded_dim()); + EXPECT_EQ(query.use_hacc(), original.use_hacc()); + EXPECT_EQ(query.delta(), original.delta()); + EXPECT_EQ(query.sum_vl_lut(), original.sum_vl_lut()); + EXPECT_EQ(query.k1xsumq(), original.k1xsumq()); + EXPECT_EQ(query.kbxsumq(), original.kbxsumq()); + EXPECT_EQ(query.g_add(), original.g_add()); + EXPECT_EQ(query.g_error(), original.g_error()); + EXPECT_TRUE(std::equal(bytes.begin(), bytes.end(), query.lut())); + query.set_g_add(5, 4); + EXPECT_EQ(query.g_add(), -4); // Failed reset did not change the metric to L2. + EXPECT_EQ(query.g_error(), 5); + query.reset(next.data(), 64, 8, scratch.data(), hacc_scratch.data()); + EXPECT_EQ(query.rotated_query(), next.data()); + EXPECT_EQ(query.g_add(), 0); +} + +TEST(FastScanApi, SingleBatchEncodersRejectOverflowAndPreserveEmptyOutput) { + constexpr size_t kDim = 64; + std::array data{}; + std::array centroid{}; + std::vector batch(BatchDataMap::data_bytes(kDim), 0x5A); + std::vector qg(QGBatchDataMap::data_bytes(kDim), 0x5A); + std::vector extra(33 * ExDataMap::data_bytes(kDim, 3), 0x5A); + const auto original = batch, original_qg = qg, original_extra = extra; + const auto encode = [&](int api, size_t count) { + switch (api) { + case 0: + quant::quantize_one_batch( + data.data(), centroid.data(), count, kDim, batch.data() + ); + break; + case 1: + quant::quantize_one_batch( + data.data(), count, kDim, batch.data() + ); + break; + case 2: + quant::quantize_qg_batch( + data.data(), centroid.data(), count, kDim, qg.data() + ); + break; + case 3: + quant::quantize_qg_batch(data.data(), count, kDim, qg.data()); + break; + default: + quant::quantize_split_batch( + data.data(), centroid.data(), count, kDim, 3, batch.data(), extra.data() + ); + } + }; + for (int api = 0; api < 5; ++api) { + SCOPED_TRACE(api); + for (size_t count : {size_t{33}, std::numeric_limits::max()}) { + expect_invalid( + [&] { encode(api, count); }, + "FastScan single-batch encoding accepts at most 32 vectors" + ); + } + encode(api, 0); + EXPECT_EQ(batch, original); + EXPECT_EQ(qg, original_qg); + EXPECT_EQ(extra, original_extra); + } +} + +TEST(FastScanApi, PublicSerialAndParallelEncodingMatchIncludingTails) { + constexpr size_t kDim = 96; + std::array data{}; + std::array centroid{}; + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i)); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t count : {1U, 17U, 32U}) { + std::vector serial(BatchDataMap::data_bytes(kDim)); + auto parallel = serial, implicit_centroid = serial; + quant::quantize_one_batch( + data.data(), centroid.data(), count, kDim, serial.data(), metric + ); + quant::quantize_one_batch( + data.data(), centroid.data(), count, kDim, parallel.data(), metric + ); + quant::quantize_one_batch( + data.data(), count, kDim, implicit_centroid.data(), metric + ); + EXPECT_EQ(serial, parallel); + EXPECT_EQ(serial, implicit_centroid); + } + } +} + +TEST(FastScanApi, DoubleEncodingQueryAndScoringWithNonzeroCentroid) { + for (size_t dim : {16U, 32U, 96U}) { + const double dimension = static_cast(dim); + const double value = 0.5 + 1e-10; + std::vector data(17 * dim, value), centroid(dim, 0.125), values(dim, 1); + for (auto metric : {METRIC_L2, METRIC_IP}) { + std::vector batch(QGBatchDataMap::data_bytes(dim)); + quant::quantize_qg_batch( + data.data(), centroid.data(), 17, dim, batch.data(), metric + ); + BatchQuery query(values.data(), dim, metric), reset; + std::vector scratch(dim * 4); + reset.reset(values.data(), dim, scratch.data(), metric); + EXPECT_EQ(reset.delta(), query.delta()); + EXPECT_TRUE(std::equal(query.lut(), query.lut() + dim * 4, reset.lut())); + const double correction = + metric == METRIC_L2 ? dimension * 0.875 * 0.875 : 1 - dimension * 0.125; + query.set_g_add(correction); + auto copy = query; + std::array distance{}; + qg_batch_estdist(batch.data(), copy, dim, distance.data()); + const double expected = metric == METRIC_L2 + ? dimension * (1 - value) * (1 - value) + : 1 - dimension * value; + for (size_t row = 0; row < 17; ++row) { + EXPECT_NEAR(distance[row], expected, 2e-13 * dimension); + } + reset.set_g_add(correction); + reset.reset(values.data(), dim, scratch.data(), metric); + EXPECT_EQ(reset.g_add(), 0); + for (bool hacc : {false, true}) { + SplitBatchQuery split(values.data(), dim, 3, metric, hacc); + EXPECT_EQ(split.padded_dim(), dim); + EXPECT_EQ(split.use_hacc(), hacc); + EXPECT_DOUBLE_EQ(split.delta(), 4.0 / (hacc ? 65535 : 255)); + Lut reused; + std::vector hacc_scratch(dim * 4); + reused.reset(values.data(), dim, scratch.data(), hacc_scratch.data(), hacc); + const size_t bytes = dim * 4 * (hacc ? 2 : 1); + EXPECT_TRUE(std::equal(split.lut(), split.lut() + bytes, reused.lut())); + SplitBatchQuery reset_split; + reset_split.reset( + values.data(), dim, 3, scratch.data(), hacc_scratch.data(), metric, hacc + ); + EXPECT_EQ(reset_split.delta(), split.delta()); + EXPECT_EQ(reset_split.kbxsumq(), split.kbxsumq()); + EXPECT_TRUE(std::equal(split.lut(), split.lut() + bytes, reset_split.lut()) + ); + } + } + } +} + +TEST(FastScanApi, DoubleLutMatchesIndependentSubsetSumQuantization) { + constexpr size_t kDim = 96; + std::array query{}; + for (size_t i = 0; i < kDim; ++i) { + // Avoid exact half-bin ties: the independent division below and the + // production reciprocal multiplication can round on opposite sides. + query[i] = std::sin(static_cast(i) * 0.37) + 1e-10 * static_cast(i); + } + std::array reference{}; + for (size_t group = 0; group < kDim / 4; ++group) { + for (size_t mask = 0; mask < 16; ++mask) { + for (size_t d = 0; d < 4; ++d) { + if ((mask & (8U >> d)) != 0) { + reference[group * 16 + mask] += query[group * 4 + d]; + } + } + } + } + const auto range = std::minmax_element(reference.begin(), reference.end()); + std::vector compact(32 * kDim / 8, 0x5A), packed(compact.size()); + fastscan::pack_codes(kDim, compact.data(), 32, packed.data()); + for (bool hacc : {false, true}) { + const Lut lut(query.data(), kDim, hacc); + const double delta = (*range.second - *range.first) / (hacc ? 65535 : 255); + EXPECT_DOUBLE_EQ(lut.delta(), delta); + EXPECT_DOUBLE_EQ(lut.sum_vl(), *range.first * (kDim / 4)); + int32_t expected = 0; + for (size_t group = 0; group < kDim / 4; ++group) { + const size_t mask = group % 2 == 0 ? 5 : 10; + expected += static_cast( + std::round((reference[group * 16 + mask] - *range.first) / delta) + ); + } + std::array accumulated{}; + if (hacc) { + fastscan::accumulate_hacc(packed.data(), lut.lut(), accumulated.data(), kDim); + } else { + fastscan::accumulate(packed.data(), lut.lut(), accumulated.data(), kDim); + } + for (auto value : accumulated) { + EXPECT_EQ(value, expected); + } + } + query.fill(0); + for (bool hacc : {false, true}) { + const Lut zero(query.data(), kDim, hacc); + EXPECT_EQ(zero.delta(), 0); + EXPECT_EQ(zero.sum_vl(), 0); + EXPECT_TRUE(std::all_of( + zero.lut(), + zero.lut() + kDim * 4 * (hacc ? 2 : 1), + [](uint8_t value) { return value == 0; } + )); + } +} + +} // namespace +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/fastscan/data_layout_test.cpp b/tests/unit/rabitqlib/fastscan/data_layout_test.cpp new file mode 100644 index 00000000..61bdf20e --- /dev/null +++ b/tests/unit/rabitqlib/fastscan/data_layout_test.cpp @@ -0,0 +1,162 @@ +#include "rabitqlib/fastscan/data_layout.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" + +namespace rabitqlib { +namespace { +float load_float(const char* data) { + float value; + std::memcpy(&value, data, sizeof(value)); + return value; +} + +TEST(DataLayoutTest, QgBatchSizeRejectsOverflow) { + const size_t max_dim = std::numeric_limits::max() / fastscan::kBatchSize; + EXPECT_THROW(QGBatchDataMap::data_bytes(max_dim + 1), std::length_error); + EXPECT_THROW(ConstQGBatchDataMap::data_bytes(max_dim + 1), std::length_error); + EXPECT_EQ( + QGBatchDataMap::data_bytes(max_dim), + max_dim * fastscan::kBatchSize / 8 + sizeof(float) * fastscan::kBatchSize * 2 + ); +} + +TEST(DataLayoutTest, BatchFactorsKeepPackedOffsetsOnUnalignedStorage) { + constexpr size_t kDim = 64; + constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; + constexpr size_t kDataBytes = kCodeBytes + (3 * fastscan::kBatchSize * sizeof(float)); + const size_t data_bytes = BatchDataMap::data_bytes(kDim); + ASSERT_EQ(data_bytes, kCodeBytes + (3 * fastscan::kBatchSize * sizeof(float))); + + alignas(float) std::array storage; + storage.fill(char{0x5a}); + char* data = storage.data() + 1; + ASSERT_NE( + reinterpret_cast(data + kCodeBytes) % alignof(float), uintptr_t{0} + ); + BatchDataMap batch(data, kDim); + batch.f_add()[0] = 1.25F; + batch.f_rescale()[0] = -2.5F; + batch.f_error()[fastscan::kBatchSize - 1] = 3.75F; + + ConstBatchDataMap stored(data, kDim); + EXPECT_FLOAT_EQ(stored.f_add()[0], 1.25F); + EXPECT_FLOAT_EQ(stored.f_rescale()[0], -2.5F); + EXPECT_FLOAT_EQ(stored.f_error()[fastscan::kBatchSize - 1], 3.75F); + EXPECT_FLOAT_EQ(load_float(data + kCodeBytes), 1.25F); + EXPECT_FLOAT_EQ( + load_float(data + kCodeBytes + (fastscan::kBatchSize * sizeof(float))), -2.5F + ); + EXPECT_FLOAT_EQ( + load_float(data + kCodeBytes + ((3 * fastscan::kBatchSize - 1) * sizeof(float))), + 3.75F + ); + EXPECT_EQ(storage.front(), char{0x5a}); + EXPECT_EQ(storage.back(), char{0x5a}); +} + +TEST(DataLayoutTest, QgBatchFactorsKeepPackedOffsetsOnUnalignedStorage) { + constexpr size_t kDim = 64; + constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; + constexpr size_t kDataBytes = kCodeBytes + (2 * fastscan::kBatchSize * sizeof(float)); + const size_t data_bytes = QGBatchDataMap::data_bytes(kDim); + ASSERT_EQ(data_bytes, kCodeBytes + (2 * fastscan::kBatchSize * sizeof(float))); + + alignas(float) std::array storage{}; + char* data = storage.data() + 1; + ASSERT_NE( + reinterpret_cast(data + kCodeBytes) % alignof(float), uintptr_t{0} + ); + QGBatchDataMap batch(data, kDim); + std::fill_n(batch.f_add(), fastscan::kBatchSize, 4.25F); + std::fill_n(batch.f_rescale(), fastscan::kBatchSize, -5.5F); + + ConstQGBatchDataMap stored(data, kDim); + EXPECT_FLOAT_EQ(stored.f_add()[fastscan::kBatchSize - 1], 4.25F); + EXPECT_FLOAT_EQ(stored.f_rescale()[fastscan::kBatchSize - 1], -5.5F); + EXPECT_FLOAT_EQ(load_float(data + kCodeBytes), 4.25F); + EXPECT_FLOAT_EQ( + load_float(data + kCodeBytes + (fastscan::kBatchSize * sizeof(float))), -5.5F + ); +} + +template +void check_qg_encoding() { + for (size_t dim : {32U, 64U, 96U}) { + for (size_t count : {1U, 17U, 32U}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(::testing::Message() << dim << ',' << count << ',' << metric); + std::vector data(count * dim), centroid(dim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(static_cast(i % 19) - 9) / T{11}; + } + for (size_t i = 0; i < dim; ++i) { + centroid[i] = static_cast(static_cast(i % 7) - 3) / T{13}; + } + for (bool zero_centroid : {false, true}) { + if (zero_centroid) { + std::fill(centroid.begin(), centroid.end(), T{0}); + } + // The first record has zero residual, including the sign convention. + std::copy(centroid.begin(), centroid.end(), data.begin()); + std::vector reference(BatchDataMap::data_bytes(dim), 0); + const size_t bytes = QGBatchDataMap::data_bytes(dim); + std::vector compact(bytes + 2, char{0x5a}); + quant::quantize_one_batch( + data.data(), centroid.data(), count, dim, reference.data(), metric + ); + if (zero_centroid) { + quant::quantize_qg_batch( + data.data(), count, dim, compact.data() + 1, metric + ); + } else { + quant::quantize_qg_batch( + data.data(), + centroid.data(), + count, + dim, + compact.data() + 1, + metric + ); + } + const ConstBatchDataMap expected(reference.data(), dim); + const ConstQGBatchDataMap actual(compact.data() + 1, dim); + EXPECT_EQ( + std::memcmp( + actual.bin_code(), + expected.bin_code(), + dim * fastscan::kBatchSize / 8 + ), + 0 + ); + for (size_t row = 0; row < count; ++row) { + EXPECT_EQ(actual.f_add()[row], expected.f_add()[row]); + EXPECT_EQ(actual.f_rescale()[row], expected.f_rescale()[row]); + } + EXPECT_EQ(compact.front(), char{0x5a}); + EXPECT_EQ(compact.back(), char{0x5a}); + } + } + } + } +} + +TEST(QgBatchEncoding, PreservesCodesAndFactorsWithExplicitAndZeroCentroids) { + check_qg_encoding(); + check_qg_encoding(); +} + +} // namespace +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/fastscan/encoder_test.cpp b/tests/unit/rabitqlib/fastscan/encoder_test.cpp new file mode 100644 index 00000000..82faa0e4 --- /dev/null +++ b/tests/unit/rabitqlib/fastscan/encoder_test.cpp @@ -0,0 +1,99 @@ +#include "rabitqlib/fastscan/encoder.hpp" // IWYU pragma: keep (public header coverage) + +#include + +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/utils/openmp.hpp" + +namespace rabitqlib::quant { +namespace { +template +void check_one_bit_batches() { + // Alternate dimensions to exercise scratch growth and shrinking on each thread. + for (size_t dim : {64U, 8U, 128U, 16U, 768U, 24U, 192U, 56U, 72U, 64U}) { + for (size_t num : {0U, 1U, 15U, 16U, 17U, 31U, 32U, 33U, 64U, 65U}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE( + ::testing::Message() + << "dim=" << dim << " num=" << num << " metric=" << metric + ); + std::vector centroid(dim); + std::vector data(num * dim); + for (size_t d = 0; d < dim; ++d) { + centroid[d] = static_cast(static_cast(d % 7) - 3) / 8; + for (size_t i = 0; i < num; ++i) { + const T residual = + i % 3 == 0 + ? T{0} + : static_cast(static_cast((i + d) % 11) - 5) / 4; + data[i * dim + d] = centroid[d] + residual; + } + } + // A guard byte/factor also verifies empty inputs and tail writes. + const size_t packed_bytes = ((num + 31) / 32) * 32 * (dim / 8); + std::vector expected(packed_bytes + 1, 0xA5); + std::vector actual = expected; + std::vector compact(num * dim / 8); + std::vector expected_add(num + 1, T{123}); + std::vector expected_rescale = expected_add; + std::vector expected_error = expected_add; + std::vector actual_add = expected_add; + std::vector actual_rescale = expected_add; + std::vector actual_error = expected_add; + for (size_t i = 0; i < num; ++i) { + rabitq_impl::one_bit::one_bit_compact_code( + data.data() + i * dim, + centroid.data(), + dim, + compact.data() + i * dim / 8, + expected_add[i], + expected_rescale[i], + expected_error[i], + metric + ); + } + fastscan::pack_codes(dim, compact.data(), num, expected.data()); + rabitq_impl::one_bit::one_bit_batch_code( + data.data(), + centroid.data(), + num, + dim, + actual.data(), + actual_add.data(), + actual_rescale.data(), + actual_error.data(), + metric + ); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_add, expected_add); + EXPECT_EQ(actual_rescale, expected_rescale); + EXPECT_EQ(actual_error, expected_error); + } + } + } +} + +TEST(RabitqOneBitBatchTest, MatchesIndividualEncoding) { + check_one_bit_batches(); + check_one_bit_batches(); +} + +TEST(RabitqOneBitBatchTest, ParallelMatchesIndividualEncoding) { +#if RABITQ_USE_OPENMP + const int previous_threads = omp_get_max_threads(); + omp_set_num_threads(2); +#endif + check_one_bit_batches(); + check_one_bit_batches(); +#if RABITQ_USE_OPENMP + omp_set_num_threads(previous_threads); +#endif +} + +} // namespace +} // namespace rabitqlib::quant diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 1dcc705e..e7feafd7 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -11,15 +11,16 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/lut.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" -#include "rabitqlib/index/estimator.hpp" -#include "rabitqlib/index/lut.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" -#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" #include "rabitqlib/utils/cpu_features.hpp" namespace rabitqlib::fastscan { @@ -97,6 +98,74 @@ TEST(FastScanLutTest, EverySupportedBackendMatchesOrderedSubsetSumsExactly) { #endif } +TEST(FastScanLutTest, ResetOwnsBytesAndReusesStorageWithoutRetainingScratch) { + constexpr size_t kDim = 64; + std::array query{}; + for (size_t i = 0; i < kDim; ++i) + query[i] = static_cast(i % 17) - 8; + for (bool hacc : {false, true}) { + Lut reused; + std::vector expected; + const size_t bytes = kDim * 4 * (hacc ? 2 : 1); + { + std::array scratch{}; + std::array hacc_scratch{}; + reused.reset(query.data(), kDim, scratch.data(), hacc_scratch.data(), hacc); + expected.assign(reused.lut(), reused.lut() + bytes); + const auto* storage = reused.lut(); + reused.reset(query.data(), kDim, scratch.data(), hacc_scratch.data(), hacc); + EXPECT_EQ(reused.lut(), storage); + scratch.fill(-1234); + hacc_scratch.fill(0xFFFF); + } + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), reused.lut())); + const Lut fresh(query.data(), kDim, hacc); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), fresh.lut())); + EXPECT_EQ(reused.delta(), fresh.delta()); + EXPECT_EQ(reused.sum_vl(), fresh.sum_vl()); + } +} + +TEST(FastScanLutTest, RejectedResetPreservesOwnedBytesAndFactors) { + std::array query{}; + query.fill(1); + std::array scratch{}; + Lut lut(query.data(), query.size()); + const std::vector bytes(lut.lut(), lut.lut() + 256); + const float delta = lut.delta(), sum = lut.sum_vl(); + const auto* storage = lut.lut(); + EXPECT_THROW(lut.reset(nullptr, 64, scratch.data()), std::invalid_argument); + EXPECT_THROW(lut.reset(query.data(), 64, nullptr), std::invalid_argument); + EXPECT_THROW(lut.reset(query.data(), 15, scratch.data()), std::invalid_argument); + EXPECT_THROW( + lut.reset(query.data(), 64, scratch.data(), nullptr, true), std::invalid_argument + ); + const size_t oversized = std::numeric_limits::max() & ~size_t{15}; + EXPECT_THROW(lut.reset(query.data(), oversized, scratch.data()), std::length_error); + EXPECT_EQ(lut.lut(), storage); + EXPECT_TRUE(std::equal(bytes.begin(), bytes.end(), lut.lut())); + EXPECT_EQ(lut.delta(), delta); + EXPECT_EQ(lut.sum_vl(), sum); +} + +TEST(BatchEstimatorTest, RejectedQueryResetPreservesMetricAndCorrections) { + std::array values{}; + values.fill(1); + for (auto metric : {METRIC_L2, METRIC_IP}) { + BatchQuery query(values.data(), values.size(), metric); + const auto other_metric = metric == METRIC_L2 ? METRIC_IP : METRIC_L2; + try { + query.reset(values.data(), values.size(), nullptr, other_metric); + FAIL() << "Null scratch must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "Lut::reset requires query and LUT scratch"); + } + query.set_g_add(9); + EXPECT_EQ(query.g_add(), metric == METRIC_L2 ? 9 : 8); + EXPECT_EQ(query.k1xsumq(), -16.0F); + } +} + TEST(FastScanPackingTest, MatchesBitReferenceIncludingTailsAndUnalignedBuffers) { for (size_t dim : {8U, 16U, 24U, 56U, 64U, 72U, 128U, 768U}) { for (size_t num = 0; num <= 65; ++num) { @@ -470,6 +539,31 @@ TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { } } +#if defined(__x86_64__) || defined(_M_X64) +TEST(BatchEstimatorTest, Avx512RejectsInvalidDimensionsBeforeReadingBuffers) { + if (!cpu::has_avx512_core()) { + GTEST_SKIP() << "AVX-512 is unavailable"; + } + const std::array values{}; + const SplitBatchQuery split_query(values.data(), values.size(), 3); + const BatchQuery query; + for (size_t dim : {0U, 15U, 17U}) { + for (bool hacc : {false, true}) { + EXPECT_THROW( + rabitqlib::simd::split_batch_estdist_avx512( + nullptr, split_query, dim, nullptr, nullptr, nullptr, hacc + ), + std::invalid_argument + ); + } + EXPECT_THROW( + rabitqlib::simd::qg_batch_estdist_avx512(nullptr, query, dim, nullptr), + std::invalid_argument + ); + } +} +#endif + TEST(BatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches) { using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); std::vector backends{ diff --git a/tests/unit/rabitqlib/index/exact_nn_test.cpp b/tests/unit/rabitqlib/index/exact_nn_test.cpp new file mode 100644 index 00000000..9dcc9fef --- /dev/null +++ b/tests/unit/rabitqlib/index/exact_nn_test.cpp @@ -0,0 +1,35 @@ +#include "rabitqlib/index/detail/exact_nn.hpp" + +#include + +#include +#include +#include +#include +#include + +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" + +using namespace rabitqlib; + +TEST(IndexTypes, PreserveIdentifierAndCandidateRepresentation) { + static_assert(std::is_same_v); + EXPECT_EQ(kPidMax, std::numeric_limits::max()); + const AnnCandidate candidate(7, 1.5F); + EXPECT_EQ(candidate.id, 7U); + EXPECT_EQ(candidate.distance, 1.5F); +} + +TEST(ExactNearestNeighbor, SupportsRequestedWorkerCounts) { + const std::array data{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + const std::array query{7, 8, 9}; + for (size_t workers : {0U, 1U, 2U, 64U}) { + EXPECT_EQ( + exact_nn(data.data(), query.data(), 4, 3, workers, euclidean_sqr), 2U + ); + EXPECT_EQ( + exact_nn(data.data(), query.data(), 4, 3, workers, dot_product_dis), 3U + ); + } +} diff --git a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp index 48f7cc3c..03bfd84f 100644 --- a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp +++ b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp @@ -9,13 +9,18 @@ #include #include #include +#include #include #include #include #include +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" namespace rabitqlib::ivf { namespace { diff --git a/tests/unit/rabitqlib/index/historical_persistence_test.cpp b/tests/unit/rabitqlib/index/historical_persistence_test.cpp index da2ab6b0..f9b27684 100644 --- a/tests/unit/rabitqlib/index/historical_persistence_test.cpp +++ b/tests/unit/rabitqlib/index/historical_persistence_test.cpp @@ -9,9 +9,12 @@ #include #include +#include "rabitqlib/core/metric.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" #include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib { namespace { @@ -49,7 +52,14 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail } }; const auto exercise = [&](auto& index, const auto& search) { + const size_t removed = name.find("qg") == 0 ? 0 : 1; + const auto check_counts = [&] { + EXPECT_EQ(index.num_stored_points(), count); + EXPECT_EQ(index.num_points(), count - removed); + EXPECT_GE(index.max_elements(), index.num_stored_points()); + }; index.load(path.c_str()); + check_counts(); EXPECT_EQ(index.dimension(), dim); const auto metric = name.find("_ip") != std::string::npos ? METRIC_IP : METRIC_L2; EXPECT_EQ(index.metric_type(), metric); @@ -57,10 +67,12 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail compare(); index.save(output.string().c_str()); index.load(output.string().c_str()); + check_counts(); search(); compare(); std::filesystem::resize_file(output, std::filesystem::file_size(output) - 1); EXPECT_THROW(index.load(output.string().c_str()), std::runtime_error); + check_counts(); search(); compare(); std::filesystem::remove(output); diff --git a/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp new file mode 100644 index 00000000..be3f8cc4 --- /dev/null +++ b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp @@ -0,0 +1,325 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" + +namespace { +using namespace rabitqlib; +constexpr size_t kCount = 64, kDim = 64; +constexpr size_t kEntryOffset = 6 * sizeof(size_t) + sizeof(int); +constexpr size_t kMaxMOffset = kEntryOffset + sizeof(hnswlib::tableint); +constexpr size_t kHeaderBytes = kMaxMOffset + 4 * sizeof(size_t) + sizeof(double); + +template +T read_at(const std::vector& bytes, size_t offset) { + if (offset > bytes.size() || sizeof(T) > bytes.size() - offset) { + throw std::runtime_error("invalid test fixture offset"); + } + T value; + std::memcpy(&value, bytes.data() + offset, sizeof(value)); + return value; +} +template +void write_at(std::vector& bytes, size_t offset, T value) { + (void)read_at(bytes, offset); + std::memcpy(bytes.data() + offset, &value, sizeof(value)); +} + +struct Fixture { + std::string path; + std::vector data = std::vector(kCount * kDim); + std::vector original; + std::vector upper_offsets, levels; + size_t stride = 0, upper_stride = 0, max_m = 0, max_m0 = 0; + size_t upper_node = 0, base_node = 0; + + explicit Fixture(MetricType metric) { + path = + (std::filesystem::temp_directory_path() / + ("rabitq_corrupt_centroids_" + std::to_string(static_cast(metric)) + "_" + + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()))) + .string(); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i % 113) / 113.0F; + } + std::vector ids(kCount); + for (size_t i = 0; i < kCount; ++i) { + ids[i] = static_cast(i); + } + ivf::IVF index( + kCount, + kDim, + kCount, + 4, + metric, + RotatorType::FhtKacRotator, + ivf::InitializerType::HNSW + ); + index.construct(data.data(), data.data(), ids.data(), false, 1); + index.save(path.c_str()); + std::ifstream file(path + ".hnsw", std::ios::binary); + original.assign(std::istreambuf_iterator(file), {}); + stride = read_at(original, 3 * sizeof(size_t)); + max_m = read_at(original, kMaxMOffset); + max_m0 = read_at(original, kMaxMOffset + sizeof(size_t)); + upper_stride = sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); + size_t pos = kHeaderBytes + kCount * stride; + for (size_t i = 0; i < kCount; ++i) { + upper_offsets.push_back(pos); + const auto bytes = read_at(original, pos); + levels.push_back(bytes / upper_stride); + pos += sizeof(unsigned int) + bytes; + } + if (pos != original.size()) { + throw std::runtime_error("bad fixture length"); + } + auto upper = + std::find_if(levels.begin(), levels.end(), [](size_t n) { return n > 0; }); + auto base = std::find(levels.begin(), levels.end(), 0); + if (upper == levels.end() || base == levels.end()) { + throw std::runtime_error("fixture needs two levels"); + } + upper_node = static_cast(upper - levels.begin()); + base_node = static_cast(base - levels.begin()); + } + ~Fixture() { + std::error_code ignored; + std::filesystem::remove(path, ignored); + std::filesystem::remove(path + ".hnsw", ignored); + } + void save(const std::vector& bytes) const { + std::ofstream file(path + ".hnsw", std::ios::binary | std::ios::trunc); + file.write(bytes.data(), static_cast(bytes.size())); + } +}; + +TEST(HnswSidecarGraph, PublicLoadRejectsCorruptStructureAndCanRetry) { + for (const auto metric : {METRIC_L2, METRIC_IP}) { + Fixture fixture(metric); + const auto reject = [&](const char* name, auto corrupt, const char* message) { + SCOPED_TRACE(name); + auto bytes = fixture.original; + corrupt(bytes); + fixture.save(bytes); + ivf::IVF loaded; + try { + loaded.load(fixture.path.c_str()); + FAIL() << "damaged sidecar was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), message); + } + fixture.save(fixture.original); + ASSERT_NO_THROW(loaded.load(fixture.path.c_str())); + PID id = 0; + float distance = 0; + ASSERT_NO_THROW(loaded.search(fixture.data.data(), 1, 1, &id, &distance)); + }; + const size_t upper = + fixture.upper_offsets[fixture.upper_node] + sizeof(unsigned int); + reject( + "count mismatch", + [](auto& b) { write_at(b, 2 * sizeof(size_t), kCount - 1); }, + "HNSW centroid count does not match IVF cluster count" + ); + reject( + "capacity below count", + [](auto& b) { write_at(b, sizeof(size_t), kCount - 1); }, + "HNSW centroid count does not match IVF cluster count" + ); + reject( + "invalid level multiplier", + [](auto& b) { + write_at( + b, + kMaxMOffset + 3 * sizeof(size_t), + std::numeric_limits::quiet_NaN() + ); + }, + "Invalid HNSW centroid graph parameters" + ); + reject( + "invalid construction breadth", + [](auto& b) { + write_at(b, kMaxMOffset + 3 * sizeof(size_t) + sizeof(double), 0); + }, + "Invalid HNSW centroid graph parameters" + ); + reject( + "truncated base records", + [&](auto& b) { b.resize(kHeaderBytes + kCount * fixture.stride - 1); }, + "Truncated HNSW centroid base records" + ); + reject( + "entry out of range", + [](auto& b) { + write_at(b, kEntryOffset, static_cast(kCount)); + }, + "Invalid HNSW centroid entry point or maximum level" + ); + reject( + "negative level", + [](auto& b) { write_at(b, 6 * sizeof(size_t), -1); }, + "Invalid HNSW centroid entry point or maximum level" + ); + reject( + "wrong maximum level", + [](auto& b) { + write_at(b, 6 * sizeof(size_t), read_at(b, 6 * sizeof(size_t)) + 1); + }, + "Invalid HNSW centroid level structure" + ); + reject( + "entry below maximum level", + [&](auto& b) { + write_at( + b, kEntryOffset, static_cast(fixture.base_node) + ); + }, + "Invalid HNSW centroid level structure" + ); + reject( + "base neighbor", + [](auto& b) { + write_at(b, kHeaderBytes, 1); + write_at( + b, + kHeaderBytes + sizeof(hnswlib::linklistsizeint), + static_cast(kCount) + ); + }, + "Invalid HNSW centroid neighbor or target level" + ); + reject( + "base degree", + [&](auto& b) { + write_at(b, kHeaderBytes, static_cast(fixture.max_m0 + 1)); + }, + "Invalid HNSW centroid neighbor count" + ); + reject( + "upper neighbor", + [&](auto& b) { + write_at(b, upper, 1); + write_at( + b, + upper + sizeof(hnswlib::linklistsizeint), + static_cast(kCount) + ); + }, + "Invalid HNSW centroid neighbor or target level" + ); + reject( + "upper degree", + [&](auto& b) { write_at(b, upper, static_cast(fixture.max_m + 1)); }, + "Invalid HNSW centroid neighbor count" + ); + reject( + "neighbor missing upper layer", + [&](auto& b) { + write_at(b, upper, 1); + write_at( + b, + upper + sizeof(hnswlib::linklistsizeint), + static_cast(fixture.base_node) + ); + }, + "Invalid HNSW centroid neighbor or target level" + ); + reject( + "partial upper layer", + [&](auto& b) { + const auto offset = fixture.upper_offsets[fixture.upper_node]; + write_at(b, offset, read_at(b, offset) - 1); + }, + "Invalid HNSW centroid upper-link length" + ); + reject( + "truncated upper length", + [&](auto& b) { b.resize(fixture.upper_offsets.back() + 2); }, + "Truncated HNSW centroid sidecar" + ); + reject( + "truncated upper links", + [&](auto& b) { + const auto offset = fixture.upper_offsets[fixture.upper_node]; + b.resize(upper + read_at(b, offset) - 1); + }, + "Invalid HNSW centroid upper-link length" + ); + reject( + "invalid graph stride", + [](auto& b) { write_at(b, kMaxMOffset, 0); }, + "Invalid HNSW centroid graph parameters" + ); + reject( + "invalid label offset", + [](auto& b) { + write_at(b, 4 * sizeof(size_t), std::numeric_limits::max()); + }, + "HNSW centroid layout does not match IVF dimension" + ); + reject( + "deleted centroid", + [](auto& b) { b[kHeaderBytes + 2] = 1; }, + "Invalid HNSW centroid entry point or deletion state" + ); + reject( + "truncated header", + [](auto& b) { b.resize(4); }, + "Truncated HNSW centroid sidecar" + ); + reject( + "trailing byte", + [](auto& b) { b.push_back(0); }, + "Invalid HNSW centroid level structure" + ); + } +} + +TEST(HnswSidecarGraph, DisconnectedGraphCannotReadEmptyResults) { + for (const auto metric : {METRIC_L2, METRIC_IP}) { + Fixture fixture(metric); + auto bytes = fixture.original; + for (size_t i = 0; i < kCount; ++i) { + write_at(bytes, kHeaderBytes + i * fixture.stride, 0); + for (size_t level = 0; level < fixture.levels[i]; ++level) { + write_at( + bytes, + fixture.upper_offsets[i] + sizeof(unsigned int) + + level * fixture.upper_stride, + 0 + ); + } + } + fixture.save(bytes); + ivf::IVF loaded; + ASSERT_NO_THROW(loaded.load(fixture.path.c_str())); + PID id = 0; + float distance = 0; + try { + loaded.search(fixture.data.data(), 1, kCount, &id, &distance); + FAIL() << "insufficient candidates must be rejected"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "HNSW centroid graph returned too few candidates"); + } + } +} +} // namespace diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 03554393..04b97f38 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -9,27 +9,38 @@ #include #include #include +#include #include #include #include #include #include +#include +#include #include +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/hnsw_dispatch.hpp" +#include "rabitqlib/index/hnsw/hnsw_quant.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" namespace rabitqlib::hnsw { struct HnswPruningTestAccess { static void quantize(HierarchicalNSW& index, PID id, const float* rotated) { - const auto* centroid = reinterpret_cast(index.centroids_memory_); + const auto* centroid = + reinterpret_cast(index.centroids_memory_.get()); quant::quantize_compact_one_bit( rotated, centroid, index.padded_dim_, index.get_bindata_by_internalid(id) ); } static float distance(const HierarchicalNSW& index, PID target, PID query) { - return index.get_quant_dist(target, query); + detail::QuantQuery cached; + return index.get_quant_dist(target, query, cached); } static std::vector connect(HierarchicalNSW& index, PID incoming) { @@ -39,15 +50,15 @@ struct HnswPruningTestAccess { neighbors[i + 1] = static_cast(i + 1); } std::fill_n(index.get_linklist0(incoming), index.maxM0_ + 1, PID{0}); - detail::quant_query().clear(); + detail::QuantQuery cached; maxheap> candidates; - candidates.emplace(index.get_quant_dist(0, incoming), 0); - index.mutually_connect_quant(incoming, candidates, 0); + candidates.emplace(index.get_quant_dist(0, incoming, cached), 0); + index.mutually_connect_quant(incoming, candidates, 0, cached); return {neighbors + 1, neighbors + 1 + index.get_list_count(neighbors)}; } static PID entry_point_label(const HierarchicalNSW& index) { - return index.get_external_label(index.enterpoint_node_); + return index.get_external_label(index.entry_point_id_); } static PID internal_id(const HierarchicalNSW& index, PID label) { @@ -63,7 +74,7 @@ struct HnswPruningTestAccess { // Give internal slot i the label count - 1 - i, as parallel construction can // when threads take slots out of order. static void reverse_labels(HierarchicalNSW& index) { - const size_t count = index.num_points(); + const size_t count = index.num_stored_points(); index.label_lookup_.clear(); for (size_t i = 0; i < count; ++i) { const auto label = static_cast(count - 1 - i); @@ -74,6 +85,83 @@ struct HnswPruningTestAccess { }; namespace { +TEST(HnswQuantQuery, MovesPreserveOwnedQueryAndLeaveSourceReusable) { + using detail::QuantQuery; + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(std::is_nothrow_move_constructible_v); + static_assert(std::is_nothrow_move_assignable_v); + const auto prepare = [](QuantQuery& query, PID id, float value) { + query.rotated.assign(64, value); + query.prepared.emplace( + query.rotated.data(), 64, 3, quant::RabitqConfig{}, METRIC_IP + ); + query.prepared->set_g_add(3, 2); + query.id = id; + }; + QuantQuery retained; + prepare(retained, 1, 0.25F); + { + QuantQuery original; + prepare(original, 7, 0.5F); + QuantQuery moved(std::move(original)); + EXPECT_EQ(original.id, kPidMax); + EXPECT_FALSE(original.prepared.has_value()); + ASSERT_TRUE(moved.prepared.has_value()); + EXPECT_EQ(moved.prepared->rotated_query(), moved.rotated.data()); + prepare(original, 9, 1.0F); + EXPECT_EQ(moved.rotated.front(), 0.5F); + retained = std::move(moved); + EXPECT_EQ(moved.id, kPidMax); + EXPECT_FALSE(moved.prepared.has_value()); + } + ASSERT_TRUE(retained.prepared.has_value()); + EXPECT_EQ(retained.id, 7U); + EXPECT_EQ(retained.prepared->rotated_query(), retained.rotated.data()); + EXPECT_EQ(retained.prepared->rotated_query()[63], 0.5F); + EXPECT_EQ(retained.prepared->kbxsumq(), -240.0F); + EXPECT_EQ(retained.prepared->g_add(), -2.0F); + auto& alias = retained; + retained = std::move(alias); + ASSERT_TRUE(retained.prepared.has_value()); + EXPECT_EQ(retained.prepared->rotated_query(), retained.rotated.data()); + retained = QuantQuery{}; + EXPECT_EQ(retained.id, kPidMax); + EXPECT_FALSE(retained.prepared.has_value()); + EXPECT_TRUE(retained.rotated.empty()); +} + +TEST(HnswSearchTest, EqualDistancesUseExternalLabelOrder) { + constexpr size_t kCount = 24, kDim = 65; + std::vector data(kCount * kDim, 0), centroid(kDim, 0), queries(3 * kDim, 0); + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1U, 4U, 9U}) { + SCOPED_TRACE(metric); + SCOPED_TRACE(bits); + std::vector clusters(kCount, 0); + HierarchicalNSW index(kCount, kDim, bits, 8, 32, 42, metric); + index.construct( + 1, centroid.data(), kCount, data.data(), clusters.data(), 1, true + ); + HnswPruningTestAccess::reverse_labels(index); + const PID removed = 3; + index.remove(&removed, 1); + for (size_t threads : {1U, 3U}) { + const auto results = index.search(queries.data(), 3, kCount, 32, threads); + for (const auto& row : results) { + // Equal vectors need not make every graph node reachable. + // Verify the returned set's ordering, not exact-search recall. + ASSERT_GE(row.size(), 2U); + EXPECT_TRUE(std::is_sorted(row.begin(), row.end())); + for (const auto& result : row) { + EXPECT_NE(result.second, removed); + EXPECT_EQ(result.first, row.front().first); + } + } + } + } + } +} TEST(HnswPruningTest, ScoresNewAndExistingNeighborsWithTheSameQuery) { constexpr size_t kDim = 64; @@ -95,7 +183,6 @@ TEST(HnswPruningTest, ScoresNewAndExistingNeighborsWithTheSameQuery) { for (PID id = 0; id < kCount; ++id) { HnswPruningTestAccess::quantize(index, id, data.data() + (id * kDim)); } - detail::quant_query().clear(); const float existing_distance = HnswPruningTestAccess::distance(index, 1, 0); EXPECT_LT(HnswPruningTestAccess::distance(index, kIncoming, 0), existing_distance); EXPECT_GT(HnswPruningTestAccess::distance(index, 0, kIncoming), existing_distance); @@ -649,10 +736,10 @@ TEST(HnswAddTest, AddIntoAnEmptyIndexBuildsFromScratch) { } TEST(HnswConstructionTest, RejectsTooSmallM) { - for (const size_t M : {0U, 1U}) { - SCOPED_TRACE(M); + for (const size_t m : {0U, 1U}) { + SCOPED_TRACE(m); try { - HierarchicalNSW index(8, 64, 4, M, 10); + HierarchicalNSW index(8, 64, 4, m, 10); FAIL() << "M below 2 must be rejected"; } catch (const std::invalid_argument& error) { EXPECT_STREQ(error.what(), "HNSW M must be at least 2"); @@ -864,7 +951,7 @@ std::unique_ptr build_for_remove( Results search_everything( HierarchicalNSW& index, const AddFixture& fixture, size_t queries ) { - const size_t count = index.num_points(); + const size_t count = index.num_stored_points(); return index.search(fixture.data.data(), queries, count, count, 1); } @@ -927,7 +1014,8 @@ TEST(HnswRemoveTest, ExcludesRemovedPointsAndKeepsSurvivors) { EXPECT_EQ(index->remove(labels.data(), labels.size()), unique); EXPECT_EQ(index->remove(labels.data(), labels.size()), 0U); - EXPECT_EQ(index->num_points(), kCount); + EXPECT_EQ(index->num_points(), kCount - unique); + EXPECT_EQ(index->num_stored_points(), kCount); // Survivors keep their order and their distances. EXPECT_EQ( @@ -1173,19 +1261,19 @@ TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { using namespace rabitqlib; using namespace rabitqlib::hnsw; constexpr size_t count = 24; - using Search = decltype(&hnsw::detail::search_knn_generic); - std::vector backends{hnsw::detail::search_knn}; + using Search = decltype(&hnsw::detail::search_knn_sorted_generic); + std::vector backends{hnsw::detail::search_knn_sorted}; #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) - backends.push_back(hnsw::detail::search_knn_avx2); + backends.push_back(hnsw::detail::search_knn_sorted_avx2); if (cpu::has_avx512_core()) - backends.push_back(hnsw::detail::search_knn_avx512_core); + backends.push_back(hnsw::detail::search_knn_sorted_avx512_core); if (cpu::has_avx512_popcnt()) - backends.push_back(hnsw::detail::search_knn_avx512_popcnt); + backends.push_back(hnsw::detail::search_knn_sorted_avx512_popcnt); #endif #if defined(__aarch64__) || defined(_M_ARM64) if (cpu::has_neon()) - backends.push_back(hnsw::detail::search_knn_neon); + backends.push_back(hnsw::detail::search_knn_sorted_neon); #endif for (size_t dim : {64U, 96U, 128U, 160U, 256U, 544U, 576U}) { std::mt19937 random(642); @@ -1208,20 +1296,26 @@ TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { // Insertion reconstructs stored sign/extra-bit tails before linking. index.add(data.data() + (count - 4) * dim, 4, clusters.data(), true, 1); const auto reference = - hnsw::detail::search_knn_generic(index, query.data(), 5, count); + hnsw::detail::search_knn_sorted_generic(index, query.data(), 5, count); + const auto ordered = + hnsw::detail::search_knn_sorted(index, query.data(), 5, count); + auto heap = hnsw::detail::search_knn(index, query.data(), 5, count); + ASSERT_EQ(heap.size(), ordered.size()); + for (auto it = ordered.rbegin(); it != ordered.rend(); ++it) { + EXPECT_EQ(heap.top(), *it); + heap.pop(); + } for (auto search : backends) { - auto expected = reference; + const auto& expected = reference; auto actual = search(index, query.data(), 5, count); ASSERT_EQ(actual.size(), expected.size()); - while (!expected.empty()) { - EXPECT_EQ(actual.top().second, expected.top().second); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(actual[i].second, expected[i].second); EXPECT_NEAR( - actual.top().first, - expected.top().first, - 2e-5F * std::max(1.0F, std::abs(expected.top().first)) + actual[i].first, + expected[i].first, + 2e-5F * std::max(1.0F, std::abs(expected[i].first)) ); - actual.pop(); - expected.pop(); } } } @@ -1268,3 +1362,22 @@ TEST(HnswSearchTest, ConcurrentCallsKeepIndependentSearchWindows) { } } } + +TEST(HnswSearchValidation, RejectsZeroKBeforeParallelSearch) { + using namespace rabitqlib; + std::vector data(128, 0.5F), center(64, 0.0F); + PID clusters[2] = {0, 0}; + hnsw::HierarchicalNSW index(2, 64, 4, 2, 16, 100); + index.construct(1, center.data(), 2, data.data(), clusters, 1, false); + for (const size_t threads : {1U, 4U}) { + for (const size_t count : {0U, 1U, 2U}) { + try { + index.search(data.data(), count, 0, 10, threads); + FAIL() << "zero k must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW k must be greater than zero"); + } + } + } + EXPECT_EQ(index.search(data.data(), 1, 1, 10, 1)[0].size(), 1U); +} diff --git a/tests/unit/rabitqlib/index/index_counts_test.cpp b/tests/unit/rabitqlib/index/index_counts_test.cpp new file mode 100644 index 00000000..93989af8 --- /dev/null +++ b/tests/unit/rabitqlib/index/index_counts_test.cpp @@ -0,0 +1,171 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/index/types.hpp" + +namespace rabitqlib { +namespace { +constexpr size_t kCount = 40; +constexpr size_t kDim = 64; + +std::vector count_test_vectors() { + std::vector data((kCount + 4) * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.17F); + } + return data; +} + +template +void expect_counts(const Index& index, size_t capacity, size_t stored, size_t active) { + EXPECT_EQ(index.max_elements(), capacity); + EXPECT_EQ(index.num_stored_points(), stored); + EXPECT_EQ(index.num_points(), active); + EXPECT_LE(index.num_points(), index.num_stored_points()); + EXPECT_LE(index.num_stored_points(), index.max_elements()); +} + +TEST(IndexCounts, IvfTracksRemovalGrowthReloadAndRebuild) { + const auto data = count_test_vectors(); + const std::vector centroid(kDim, 0.0F); + const std::vector clusters(kCount + 4, 0); + const std::array removed{1, kCount - 1, 1}; + const std::string path = ::testing::TempDir() + "rabitq_ivf_counts.index"; + for (const auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1UL, 4UL, 32UL}) { + SCOPED_TRACE(::testing::Message() << "bits=" << bits << " metric=" << metric); + ivf::IVF index(kCount, kDim, 1, bits, metric); + expect_counts(index, kCount, 0, 0); + index.construct(data.data(), centroid.data(), clusters.data(), false, 1); + expect_counts(index, kCount, kCount, kCount); + EXPECT_EQ(index.remove(removed.data(), removed.size()), 2U); + EXPECT_EQ(index.remove(removed.data(), removed.size()), 0U); + expect_counts(index, kCount, kCount, kCount - 2); + EXPECT_THROW(index.add(nullptr, 1), std::invalid_argument); + const std::array invalid{0, kCount}; + EXPECT_THROW( + index.remove(invalid.data(), invalid.size()), std::invalid_argument + ); + expect_counts(index, kCount, kCount, kCount - 2); + + index.add(data.data() + kCount * kDim, 3, clusters.data(), false, 1); + expect_counts(index, kCount + 3, kCount + 3, kCount + 1); + index.save(path.c_str()); + ivf::IVF loaded; + expect_counts(loaded, 0, 0, 0); + loaded.load(path.c_str()); + expect_counts(loaded, kCount + 3, kCount + 3, kCount + 1); + std::filesystem::resize_file(path, std::filesystem::file_size(path) - 1); + EXPECT_THROW(loaded.load(path.c_str()), std::runtime_error); + expect_counts(loaded, kCount + 3, kCount + 3, kCount + 1); + + std::vector all(kCount + 3); + std::iota(all.begin(), all.end(), PID{0}); + EXPECT_EQ(loaded.remove(all.data(), all.size()), kCount + 1); + expect_counts(loaded, kCount + 3, kCount + 3, 0); + loaded.save(path.c_str()); + loaded.load(path.c_str()); + expect_counts(loaded, kCount + 3, kCount + 3, 0); + loaded.add(data.data() + (kCount + 3) * kDim, 1, clusters.data(), false, 1); + expect_counts(loaded, kCount + 4, kCount + 4, 1); + const PID appended_id = kCount + 3; + EXPECT_EQ(loaded.remove(&appended_id, 1), 1U); + expect_counts(loaded, kCount + 4, kCount + 4, 0); + + // Rebuilding repopulates every slot and clears all prior tombstones. + loaded.construct(data.data(), centroid.data(), clusters.data(), false, 1); + expect_counts(loaded, kCount + 4, kCount + 4, kCount + 4); + } + } + std::filesystem::remove(path); +} + +TEST(IndexCounts, HnswSeparatesCapacityStoredLabelsAndActivePoints) { + const auto data = count_test_vectors(); + const std::vector centroid(kDim, 0.0F); + std::vector clusters(kCount, 0); + constexpr size_t kCapacity = 64; + const std::array removed{1, kCount - 1, 1}; + const std::string path = ::testing::TempDir() + "rabitq_hnsw_counts.index"; + hnsw::HierarchicalNSW index(kCapacity, kDim, 4, 8, 40); + expect_counts(index, kCapacity, 0, 0); + index.construct(1, centroid.data(), kCount, data.data(), clusters.data(), 1); + expect_counts(index, kCapacity, kCount, kCount); + EXPECT_EQ(index.remove(removed.data(), removed.size()), 2U); + EXPECT_EQ(index.remove(removed.data(), removed.size()), 0U); + expect_counts(index, kCapacity, kCount, kCount - 2); + EXPECT_THROW(index.resize(kCount - 2), std::invalid_argument); + EXPECT_THROW(index.add(nullptr, 1), std::invalid_argument); + expect_counts(index, kCapacity, kCount, kCount - 2); + index.resize(2 * kCapacity); + expect_counts(index, 2 * kCapacity, kCount, kCount - 2); + EXPECT_EQ( + index.add(data.data() + kCount * kDim, 3, clusters.data(), false, 1), + (std::vector{kCount, kCount + 1, kCount + 2}) + ); + expect_counts(index, 2 * kCapacity, kCount + 3, kCount + 1); + index.save(path.c_str()); + hnsw::HierarchicalNSW loaded; + expect_counts(loaded, 0, 0, 0); + loaded.load(path.c_str()); + expect_counts(loaded, 2 * kCapacity, kCount + 3, kCount + 1); + std::filesystem::resize_file(path, std::filesystem::file_size(path) - 1); + EXPECT_THROW(loaded.load(path.c_str()), std::runtime_error); + expect_counts(loaded, 2 * kCapacity, kCount + 3, kCount + 1); + std::vector all(kCount + 3); + std::iota(all.begin(), all.end(), PID{0}); + EXPECT_EQ(loaded.remove(all.data(), all.size()), kCount + 1); + expect_counts(loaded, 2 * kCapacity, kCount + 3, 0); + loaded.save(path.c_str()); + loaded.load(path.c_str()); + expect_counts(loaded, 2 * kCapacity, kCount + 3, 0); + EXPECT_EQ( + loaded.add(data.data() + (kCount + 3) * kDim, 1, clusters.data(), false, 1), + std::vector{kCount + 3} + ); + expect_counts(loaded, 2 * kCapacity, kCount + 4, 1); + std::filesystem::remove(path); +} + +TEST(IndexCounts, SymphonyCountsOnlyCompletedBuilds) { + const auto data = count_test_vectors(); + const std::string path = ::testing::TempDir() + "rabitq_symqg_counts.index"; + for (size_t bits : {0UL, 4UL, 8UL}) { + symqg::QuantizedGraph index( + kCount, kDim, 32, METRIC_L2, RotatorType::FhtKacRotator, bits, 42 + ); + expect_counts(index, kCount, 0, 0); + symqg::QGBuilder builder(index, kCount, data.data(), 1, 42); + expect_counts(index, kCount, 0, 0); + builder.build(); + expect_counts(index, kCount, kCount, kCount); + index.save(path.c_str()); + symqg::QuantizedGraph loaded; + expect_counts(loaded, 0, 0, 0); + loaded.load(path.c_str()); + expect_counts(loaded, kCount, kCount, kCount); + std::filesystem::resize_file(path, std::filesystem::file_size(path) - 1); + EXPECT_THROW(loaded.load(path.c_str()), std::runtime_error); + expect_counts(loaded, kCount, kCount, kCount); + builder.reset(data.data()); + expect_counts(index, kCount, 0, 0); + builder.refine(); + expect_counts(index, kCount, kCount, kCount); + } + std::filesystem::remove(path); +} +} // namespace +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index 7f4238d8..eba48dcc 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -4,17 +4,25 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include #include #include +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" +#include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + namespace rabitqlib::ivf { namespace { @@ -228,36 +236,41 @@ TEST(HNSWInitializerTest, ConcurrentSearchesKeepRoutingStateSynchronized) { EXPECT_FALSE(invalid_result); } -TEST(ParallelForTest, AutomaticThreadCountProcessesEveryItem) { - std::atomic calls{0}; - parallel_for(0, 100, 0, [&](size_t, size_t) { ++calls; }); - EXPECT_EQ(calls, 100U); -} - -TEST(ParallelForTest, EmptyRangeDoesNotInvokeFunction) { - size_t calls = 0; - parallel_for(5, 5, 0, [&](size_t, size_t) { ++calls; }); - EXPECT_EQ(calls, 0U); -} - TEST(ParallelForTest, JoinsWorkersBeforeRethrowingFunctionException) { - std::atomic active{0}; - EXPECT_THROW( - parallel_for( - 0, - 100, - 4, - [&](size_t id, size_t) { - ++active; - --active; + if (rabitqlib::resolve_num_threads(2) < 2) { + GTEST_SKIP() << "Requires two available workers"; + } + std::promise worker_started, release_worker; + auto started = worker_started.get_future().share(); + auto release = release_worker.get_future(); + std::atomic worker_finished{false}; + auto result = std::async(std::launch::async, [&] { + try { + rabitqlib::index_detail::parallel_for(0, 2, 2, [&](size_t id, size_t) { if (id == 0) { + if (started.wait_for(std::chrono::seconds(5)) != + std::future_status::ready) { + throw std::runtime_error("Second worker did not start"); + } throw std::runtime_error("parallel failure"); } - } - ), - std::runtime_error - ); - EXPECT_EQ(active, 0U); + worker_started.set_value(); + release.wait(); + worker_finished = true; + }); + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "parallel failure"); + return worker_finished.load(); + } + ADD_FAILURE() << "Worker exception was not propagated"; + return false; + }); + EXPECT_EQ(started.wait_for(std::chrono::seconds(5)), std::future_status::ready); + // Keep one callback outstanding while the other throws. The caller must not + // observe that exception until the outstanding callback has completed. + EXPECT_EQ(result.wait_for(std::chrono::milliseconds(20)), std::future_status::timeout); + release_worker.set_value(); + EXPECT_TRUE(result.get()); } } // namespace diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 0d91a70d..715660a0 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -19,9 +20,12 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/buffer.hpp" -#include "rabitqlib/utils/rotator.hpp" +#include "ivf_search.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::ivf { namespace { @@ -331,7 +335,7 @@ TEST(IvfSearchTest, RoutesInnerProductQueriesByInnerProduct) { TEST(IvfSearchTest, RoutesInnerProductQueriesWithHNSWCentroids) { constexpr size_t kDim = 64; - constexpr size_t kCount = 20000; + constexpr size_t kCount = 128; std::array data{}; std::vector centroids(kCount * kDim, 0.0F); std::array query{}; @@ -795,7 +799,7 @@ TEST(IvfAddTest, FasterQuantizationIsSupportedAndDeterministic) { TEST(IvfAddTest, RoutesNewPointsWithHNSWCentroids) { constexpr size_t kDim = 64; - constexpr size_t kCount = 20000; + constexpr size_t kCount = 128; // Centroid i has first coordinate i + 1, so an inner product with e0 is largest for the // last centroid and the routing decision is unambiguous. std::vector centroids(kCount * kDim, 0.0F); @@ -1048,3 +1052,39 @@ TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { } // namespace } // namespace rabitqlib::ivf + +TEST(IVFSidecarValidation, RejectsMismatchedVectorStorage) { + using namespace rabitqlib; + for (const auto metric : {METRIC_L2, METRIC_IP}) { + for (const size_t sidecar_dim : {64U, 128U}) { + const std::string path = ::testing::TempDir() + "rabitq_sidecar_layout"; + std::vector data(96, 1.0F), center(128, 0.0F); + PID cluster = 0; + ivf::IVF original( + 1, 96, 1, 4, metric, RotatorType::FhtKacRotator, ivf::InitializerType::HNSW + ); + original.construct(data.data(), center.data(), &cluster, false, 1); + original.save(path.c_str()); + ivf::HNSWInitializer other(sidecar_dim, 1, metric); + other.add_vectors(center.data(), 1); + std::ofstream ignored; + other.save(ignored, path.c_str()); + ivf::IVF loaded; + try { + loaded.load(path.c_str()); + FAIL() << "mismatched sidecar must be rejected"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), "HNSW centroid layout does not match IVF dimension" + ); + } + original.save(path.c_str()); + ASSERT_NO_THROW(loaded.load(path.c_str())); + PID id = 0; + float distance = 0; + ASSERT_NO_THROW(loaded.search(data.data(), 1, 1, &id, &distance)); + std::filesystem::remove(path); + std::filesystem::remove(path + ".hnsw"); + } + } +} diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp index 4cd9a48e..60001776 100644 --- a/tests/unit/rabitqlib/index/pipnn_test.cpp +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -11,7 +11,9 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "pipnn_internal.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/space.hpp" namespace rabitqlib::symqg::detail { diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 2b5274d1..a2ddeb5d 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -24,21 +23,25 @@ #include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/index/estimator.hpp" -#include "rabitqlib/index/lut.hpp" -#include "rabitqlib/index/query.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/rotator.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/lut.hpp" +#include "rabitqlib/fastscan/encoder.hpp" +#include "rabitqlib/fastscan/estimator.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" +#include "rabitqlib/index/detail/buffer.hpp" +#include "rabitqlib/index/detail/exact_nn.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" #include "rabitqlib/index/symqg/detail/pipnn.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/simd/estimator_dispatch.hpp" -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" -#include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/utils/tools.hpp" namespace rabitqlib::symqg { struct QGConstructionTestAccess { @@ -47,14 +50,9 @@ struct QGConstructionTestAccess { graph.initialize(); } static float build_distance(QGBuilder& builder, PID source_id, PID target_id) { - std::vector reconstructed; - std::optional prepared; - const float* source = builder.qg_.prepare_build_query( - source_id, reconstructed, prepared, builder.build_cache() - ); - return builder.qg_.point_distance( - source, prepared ? &*prepared : nullptr, target_id - ); + QGBuilder::BuildQuery query(builder.qg_, builder.build_cache()); + query.reset(source_id); + return query.distance_to(target_id); } static CandidateList probe_pruned_edges( QGBuilder& builder, @@ -114,7 +112,7 @@ struct QGConstructionTestAccess { } builder.new_neighbors_[0] = current; builder.pruned_neighbors_[0] = pruned; - builder.graph_refine(); + builder.complete_degrees(); const auto& actual = builder.new_neighbors_[0]; EXPECT_EQ(actual.size(), expected.size()); for (size_t i = 0; i < std::min(actual.size(), expected.size()); ++i) { @@ -130,13 +128,13 @@ struct QGConstructionTestAccess { static void disconnect(QuantizedGraph& graph) { for (PID i = 0; i < 66; ++i) { - std::vector> row; + std::vector row; for (PID j = (i / 33) * 33; j < (i / 33 + 1) * 33; ++j) { if (i != j) { - row.emplace_back(j, 0.0F); + row.push_back(j); } } - graph.update_qg(i, row); + update(graph, i, row); } } @@ -214,18 +212,17 @@ struct QGConstructionTestAccess { static void retain_seed_scores(QGBuilder& builder) { auto& graph = builder.qg_; for (PID i = 0; i < builder.num_nodes_; ++i) { - std::vector reconstructed; - std::optional prepared; - const float* source = graph.prepare_build_query(i, reconstructed, prepared); + QGBuilder::BuildQuery query(graph); + query.reset(i); for (size_t j = 0; j < builder.degrees_[i]; ++j) { const PID id = graph.get_neighbors(i)[j]; - builder.new_neighbors_[i].emplace_back( - id, graph.point_distance(source, prepared ? &*prepared : nullptr, id) - ); + builder.new_neighbors_[i].emplace_back(id, query.distance_to(id)); } } } - static void search(QGBuilder& builder) { builder.search_new_neighbors(false); } + static void search(QGBuilder& builder) { + builder.search_new_neighbors(/*record_pruned=*/false); + } static auto encoded_neighbors(const QuantizedGraph& graph, PID id) { return graph.get_neighbors(id); } @@ -240,11 +237,20 @@ struct QGConstructionTestAccess { graph.copy_vectors(data, threads); } static void update( - QuantizedGraph& graph, - PID source, - const std::vector>& neighbors + QuantizedGraph& graph, PID source, const std::vector& neighbors ) { - graph.update_qg(source, neighbors); + const size_t dim = graph.padded_dim_; + std::vector rotated((neighbors.size() + 1) * dim); + QGBuilder::BuildQuery query(graph); + for (size_t i = 0; i < neighbors.size(); ++i) { + query.reset(neighbors[i]); + std::copy_n(query.rotated_query(), dim, rotated.data() + i * dim); + } + query.reset(source); + std::copy_n(query.rotated_query(), dim, rotated.data() + neighbors.size() * dim); + graph.encode_neighbors( + source, neighbors, rotated.data(), rotated.data() + neighbors.size() * dim + ); } static void fill_batch_factors(QuantizedGraph& graph, PID source, float value) { QGBatchDataMap batch(graph.get_batch_data(source), graph.padded_dim_); @@ -999,8 +1005,16 @@ TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { builder.build(); EXPECT_FLOAT_EQ(builder.avg_degree(), 32); EXPECT_FALSE(builder.check_dup()); + graph.set_ef(64); + std::array before{}, after{}; + std::array before_distances{}, after_distances{}; + graph.search(data.data(), 4, before.data(), before_distances.data()); EXPECT_THROW(builder.build(0), std::invalid_argument); EXPECT_THROW(builder.build(1), std::invalid_argument); + // Rejected pass counts must not invalidate an already searchable graph. + graph.search(data.data(), 4, after.data(), after_distances.data()); + EXPECT_EQ(before, after); + EXPECT_EQ(before_distances, after_distances); EXPECT_THROW( (QGBuilder(graph, 64, data.data(), 1, static_cast(255))), std::invalid_argument @@ -1045,9 +1059,7 @@ TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { QGConstructionTestAccess::fill_batch_factors( graph, 0, std::numeric_limits::quiet_NaN() ); - std::vector> neighbors; - neighbors.emplace_back(1, 0.0F); - QGConstructionTestAccess::update(graph, 0, neighbors); + QGConstructionTestAccess::update(graph, 0, {1}); const auto active = QGConstructionTestAccess::batch_factors(graph, 0, 0); EXPECT_TRUE(std::isfinite(active.first)); diff --git a/tests/unit/rabitqlib/quantization/core_api_test.cpp b/tests/unit/rabitqlib/quantization/core_api_test.cpp new file mode 100644 index 00000000..55d450cf --- /dev/null +++ b/tests/unit/rabitqlib/quantization/core_api_test.cpp @@ -0,0 +1,572 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/api.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace { +namespace api = rabitqlib::quant::v1; +using rabitqlib::METRIC_IP; +using rabitqlib::METRIC_L2; + +struct Workspace { + std::vector storage; + size_t bytes; + explicit Workspace(size_t n) + : storage((n + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t)) + , bytes(n) {} + api::Buffer view() { return {storage.data(), bytes}; } +}; + +api::CodeView view(api::CodeOutput output) { + return { + output.layout, + {output.codes.data, output.codes.bytes}, + {output.extra.data, output.extra.bytes}}; +} + +TEST(CoreApi, ExactWorkspaceSizeProtectsBoundaryBytes) { + for (size_t dim : {2U, 3U, 31U, 33U, 65U, 96U}) { + for (size_t bits : {1U, 4U, 9U}) { + SCOPED_TRACE(dim); + SCOPED_TRACE(bits); + const api::Parameters p{dim, bits}; + const auto r = api::requirements(p, api::CodeLayout::UnpackedU16); + std::vector storage( + (r.workspace_bytes + 2 * r.workspace_alignment + sizeof(std::max_align_t) - + 1) / + sizeof(std::max_align_t) + ); + auto* bytes = reinterpret_cast(storage.data()); + const size_t capacity = storage.size() * sizeof(std::max_align_t); + std::fill_n(bytes, capacity, uint8_t{0xA5}); + api::Buffer workspace{bytes + r.workspace_alignment, r.workspace_bytes}; + std::vector data(dim); + for (size_t i = 0; i < dim; ++i) { + data[i] = static_cast(static_cast(i % 17) - 8) / 11; + } + std::vector codes(r.code_bytes), converted(r.code_bytes); + api::CodeOutput out{ + api::CodeLayout::UnpackedU16, {codes.data(), codes.size()}, {}}; + const auto factors = api::encode(p, data.data(), nullptr, out, workspace); + EXPECT_TRUE(std::isfinite(factors.rescale)); + api::convert_codes( + p, + view(out), + {out.layout, {converted.data(), converted.size()}, {}}, + workspace + ); + EXPECT_EQ(codes, converted); + for (size_t i = 0; i < r.workspace_alignment; ++i) { + EXPECT_EQ(bytes[i], 0xA5); + } + for (size_t i = r.workspace_alignment + r.workspace_bytes; i < capacity; ++i) { + EXPECT_EQ(bytes[i], 0xA5); + } + --workspace.bytes; + EXPECT_THROW( + api::encode(p, data.data(), nullptr, out, workspace), std::invalid_argument + ); + } + } +} + +TEST(CoreApi, CodesFactorsAndBothLayoutsMatchExistingEncoder) { + std::mt19937 generator(123); + std::uniform_real_distribution distribution(-1, 1); + for (size_t dim : {32U, 64U, 96U, 128U, 160U, 256U}) { + std::vector data(dim), centroid(dim), query(dim); + for (size_t i = 0; i < dim; ++i) { + centroid[i] = distribution(generator); + data[i] = i % 7 == 0 ? centroid[i] : distribution(generator); + query[i] = distribution(generator); + } + for (size_t bits = 1; bits <= 9; ++bits) { + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (double scale : {-1.0, 7.25}) { + SCOPED_TRACE( + ::testing::Message() + << dim << "," << bits << "," << metric << "," << scale + ); + api::Parameters p{dim, bits, metric, scale}; + auto raw = api::requirements(p, api::CodeLayout::UnpackedU16); + auto packed = api::requirements(p, api::CodeLayout::SplitPacked); + Workspace workspace(raw.workspace_bytes); + // Deliberately unaligned code buffers, with canaries and unrelated + // records. + std::vector code(raw.code_bytes + 2, 0xA5), + roundtrip(raw.code_bytes + 2, 0xA5); + std::vector sign(packed.code_bytes + 2, 0xA5), + extra(packed.extra_bytes + 2, 0xA5); + api::CodeOutput out{ + api::CodeLayout::UnpackedU16, + {code.data() + 1, raw.code_bytes}, + {}}; + auto factors = + api::encode(p, data.data(), centroid.data(), out, workspace.view()); + std::vector reference(dim); + float add = 0, rescale = 0, error = 0; + rabitqlib::quant::RabitqConfig config; + config.t_const = scale; + rabitqlib::quant::quantize_full_single( + data.data(), + centroid.data(), + dim, + bits, + reference.data(), + add, + rescale, + error, + metric, + config + ); + EXPECT_EQ(factors.add, add); + EXPECT_EQ(factors.rescale, rescale); + EXPECT_EQ(factors.error, error); + for (size_t i = 0; i < dim; ++i) { + EXPECT_EQ( + code[1 + 2 * i] | (uint16_t{code[2 + 2 * i]} << 8U), + reference[i] + ); + } + api::CodeOutput split{ + api::CodeLayout::SplitPacked, + {sign.data() + 1, packed.code_bytes}, + {extra.data() + 1, packed.extra_bytes}}; + api::convert_codes(p, view(out), split, workspace.view()); + api::CodeOutput restored{ + api::CodeLayout::UnpackedU16, + {roundtrip.data() + 1, raw.code_bytes}, + {}}; + api::convert_codes(p, view(split), restored, workspace.view()); + EXPECT_EQ(code, roundtrip); + const auto state = api::prepare_query(p, query.data(), centroid.data()); + const float g_add = + metric == METRIC_L2 + ? rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim) + : -rabitqlib::dot_product(query.data(), centroid.data(), dim); + float sum = 0; + for (float x : query) + sum += x; + const float expected = rabitqlib::quant::full_est_dist( + reference.data(), + query.data(), + rabitqlib::excode_ipimpl::ip_fxi, + dim, + bits, + add, + rescale, + g_add, + -0.5F * sum + ); + // The final result may cancel large dot-product terms. Use the + // standard gamma_n bound on absolute terms of both reductions. + double absolute_dot = 0; + for (size_t i = 0; i < dim; ++i) { + absolute_dot += + std::abs(static_cast(query[i]) * reference[i]); + } + const double eps = std::numeric_limits::epsilon(); + const double gamma = (dim + 8) * eps / (1 - (dim + 8) * eps); + const float tolerance = static_cast( + 2 * gamma * std::abs(rescale) * + (absolute_dot + std::abs(sum) * 0.5 * ((1U << bits) - 1)) + + 8 * eps * (1 + std::abs(add) + std::abs(g_add) + std::abs(expected)) + ); + EXPECT_NEAR( + api::compute(state, view(out), factors), expected, tolerance + ); + EXPECT_NEAR( + api::compute(state, view(split), factors), expected, tolerance + ); + if (bits > 1) { + // A stronger oracle uses the existing packed estimator and + // independently computed legacy factors with the same layout. + std::vector bin_data( + rabitqlib::BinDataMap::data_bytes(dim) + ); + std::vector ex_data( + rabitqlib::ExDataMap::data_bytes(dim, bits - 1) + ); + rabitqlib::BinDataMap bin_map(bin_data.data(), dim); + rabitqlib::ExDataMap ex_map(ex_data.data(), dim, bits - 1); + std::copy_n(sign.data() + 1, packed.code_bytes, bin_map.bin_code()); + std::copy_n(extra.data() + 1, packed.extra_bytes, ex_map.ex_code()); + ex_map.f_add_ex() = add; + ex_map.f_rescale_ex() = rescale; + rabitqlib::SplitSingleQuery prepared( + query.data(), dim, bits - 1, config, metric + ); + float legacy_distance = 0, lower = 0, product = 0; + rabitqlib::split_single_fulldist( + bin_data.data(), + ex_data.data(), + rabitqlib::select_excode_ipfunc(bits - 1), + prepared, + dim, + bits - 1, + legacy_distance, + lower, + product, + g_add, + 0 + ); + EXPECT_NEAR( + api::compute(state, view(split), factors), + legacy_distance, + 2 * tolerance + ); + } + const auto direct = api::encode( + p, data.data(), centroid.data(), split, workspace.view() + ); + EXPECT_EQ(direct.add, factors.add); + EXPECT_EQ(direct.rescale, factors.rescale); + EXPECT_NEAR( + api::compute(state, view(split), direct), expected, tolerance + ); + for (const auto* buffer : {&code, &roundtrip, &sign, &extra}) { + EXPECT_EQ(buffer->front(), 0xA5); + EXPECT_EQ(buffer->back(), 0xA5); + } + } + } + } + } +} + +TEST(CoreApi, UnpackedCodesDoNotRequireLibraryPadding) { + for (size_t dim : {2U, 3U, 15U, 33U, 65U}) { + std::vector data(dim, 1.0F), query(dim, 0.5F); + api::Parameters p{dim, 4}; + const auto r = api::requirements(p, api::CodeLayout::UnpackedU16); + Workspace workspace(r.workspace_bytes); + std::vector codes(r.code_bytes); + api::CodeOutput out{api::CodeLayout::UnpackedU16, {codes.data(), codes.size()}, {}}; + auto f = api::encode(p, data.data(), nullptr, out, workspace.view()); + auto state = api::prepare_query(p, query.data(), nullptr); + EXPECT_NEAR( + api::compute(state, view(out), f), static_cast(dim) * 0.25F, 1e-4F + ); + EXPECT_THROW( + api::requirements(p, api::CodeLayout::SplitPacked), std::invalid_argument + ); + } +} + +TEST(CoreApi, NullCentroidZeroResidualAndCallerFactorColumns) { + constexpr size_t dim = 96; + std::array data{}, query{}; + query.fill(0.25F); + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1U, 4U, 9U}) { + api::Parameters p{dim, bits, metric}; + const auto sizes = api::requirements(p, api::CodeLayout::SplitPacked); + Workspace workspace(sizes.workspace_bytes); + std::vector signs(sizes.code_bytes), extra(sizes.extra_bytes); + api::CodeOutput out{ + api::CodeLayout::SplitPacked, + {signs.data(), signs.size()}, + {extra.data(), extra.size()}}; + auto f = api::encode(p, data.data(), nullptr, out, workspace.view()); + // Store factors as separate columns, not adjacent to either code buffer. + std::array add{f.add}, scale{f.rescale}, error{f.error}; + auto q = api::prepare_query(p, query.data(), nullptr); + EXPECT_EQ(q.dimension(), dim); + EXPECT_EQ(q.total_bits(), bits); + EXPECT_EQ(q.metric_type(), metric); + EXPECT_FLOAT_EQ( + api::compute(q, view(out), {add[0], scale[0], error[0]}), + metric == METRIC_L2 ? 6.0F : 1.0F + ); + EXPECT_EQ(f.rescale, 0); + EXPECT_EQ(f.error, 0); + EXPECT_TRUE(std::all_of(signs.begin(), signs.end(), [](uint8_t v) { + return v == 0; + })); + } + } +} + +TEST(CoreApi, RejectsBadParametersAndBuffersBeforeWritingCodes) { + api::Parameters p{64, 4}; + auto r = api::requirements(p, api::CodeLayout::UnpackedU16); + Workspace workspace(r.workspace_bytes); + std::array data{}; + std::vector code(r.code_bytes, 0x5A); + api::CodeOutput out{api::CodeLayout::UnpackedU16, {code.data(), code.size()}, {}}; + for (size_t dim : {0U, 1U}) { + auto bad = p; + bad.dim = dim; + EXPECT_THROW(api::requirements(bad, out.layout), std::invalid_argument); + } + for (size_t bits : {0U, 10U}) { + auto bad = p; + bad.total_bits = bits; + EXPECT_THROW(api::requirements(bad, out.layout), std::invalid_argument); + } + auto bad = p; + bad.metric = static_cast(255); + EXPECT_THROW(api::requirements(bad, out.layout), std::invalid_argument); + bad = p; + bad.t_const = std::numeric_limits::infinity(); + EXPECT_THROW(api::requirements(bad, out.layout), std::invalid_argument); + bad = p; + bad.dim = std::numeric_limits::max() & ~size_t{31}; + EXPECT_THROW(api::requirements(bad, out.layout), std::length_error); + bad.dim = std::numeric_limits::max() / sizeof(float); + EXPECT_THROW(api::requirements(bad, out.layout), std::length_error); + EXPECT_THROW( + api::requirements(p, static_cast(255)), std::invalid_argument + ); + EXPECT_THROW( + api::encode(p, nullptr, nullptr, out, workspace.view()), std::invalid_argument + ); + EXPECT_THROW( + api::encode( + p, data.data(), nullptr, out, {workspace.storage.data(), r.workspace_bytes - 1} + ), + std::invalid_argument + ); + EXPECT_THROW( + api::encode( + p, + data.data(), + nullptr, + out, + {reinterpret_cast(workspace.storage.data()) + 1, r.workspace_bytes} + ), + std::invalid_argument + ); + auto small = out; + small.codes.bytes--; + EXPECT_THROW( + api::encode(p, data.data(), nullptr, small, workspace.view()), std::invalid_argument + ); + auto overlap = out; + overlap.codes = {data.data(), sizeof(data)}; + EXPECT_THROW( + api::encode(p, data.data(), nullptr, overlap, workspace.view()), + std::invalid_argument + ); + overlap.codes = workspace.view(); + EXPECT_THROW( + api::encode(p, data.data(), nullptr, overlap, workspace.view()), + std::invalid_argument + ); + data[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW( + api::encode(p, data.data(), nullptr, out, workspace.view()), std::invalid_argument + ); + data[0] = std::numeric_limits::max(); + EXPECT_THROW( + api::encode(p, data.data(), nullptr, out, workspace.view()), std::overflow_error + ); + EXPECT_TRUE(std::all_of(code.begin(), code.end(), [](uint8_t v) { return v == 0x5A; })); + data[0] = 0; + auto q = api::prepare_query(p, data.data(), nullptr); + EXPECT_THROW(api::compute({}, view(out), {}), std::invalid_argument); + EXPECT_THROW( + api::compute(q, view(out), {}), std::invalid_argument + ); // Invalid logical codes. + EXPECT_THROW( + api::compute(q, view(out), {std::numeric_limits::quiet_NaN(), 0, 0}), + std::invalid_argument + ); + EXPECT_THROW( + api::convert_codes(p, view(out), out, workspace.view()), std::invalid_argument + ); +} + +TEST(CoreApi, ReusedQueryAndWorkspaceStillValidateEachRecord) { + constexpr size_t kDim = 32; + const api::Parameters parameters{kDim, 4, METRIC_L2}; + const auto raw = api::requirements(parameters, api::CodeLayout::UnpackedU16); + const auto packed = api::requirements(parameters, api::CodeLayout::SplitPacked); + Workspace workspace(raw.workspace_bytes); + std::array data{}; + data.fill(0.25F); + std::vector codes(raw.code_bytes), sign(packed.code_bytes), + extra(packed.extra_bytes); + api::CodeOutput output{api::CodeLayout::UnpackedU16, {codes.data(), codes.size()}, {}}; + api::CodeOutput destination{ + api::CodeLayout::SplitPacked, + {sign.data(), sign.size()}, + {extra.data(), extra.size()}}; + const auto factors = + api::encode(parameters, data.data(), nullptr, output, workspace.view()); + const auto query = api::prepare_query(parameters, data.data(), nullptr); + const float expected = api::compute(query, view(output), factors); + api::convert_codes(parameters, view(output), destination, workspace.view()); + const auto expected_sign = sign, expected_extra = extra; + + auto bad_factors = factors; + bad_factors.error = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(api::compute(query, view(output), bad_factors), std::invalid_argument); + const uint8_t first = codes[0]; + codes[0] = 255; + EXPECT_THROW(api::compute(query, view(output), factors), std::invalid_argument); + EXPECT_THROW( + api::convert_codes(parameters, view(output), destination, workspace.view()), + std::invalid_argument + ); + EXPECT_EQ(sign, expected_sign); + EXPECT_EQ(extra, expected_extra); + codes[0] = first; + auto short_destination = destination; + --short_destination.codes.bytes; + EXPECT_THROW( + api::convert_codes(parameters, view(output), short_destination, workspace.view()), + std::invalid_argument + ); + auto invalid_layout = destination; + invalid_layout.layout = static_cast(255); + EXPECT_THROW( + api::convert_codes(parameters, view(output), invalid_layout, workspace.view()), + std::invalid_argument + ); + EXPECT_EQ(api::compute(query, view(output), factors), expected); + api::convert_codes(parameters, view(output), destination, workspace.view()); + EXPECT_EQ(sign, expected_sign); + EXPECT_EQ(extra, expected_extra); +} + +TEST(CoreApi, ScoringRetainsBufferAndNumericalChecks) { + constexpr size_t dim = 64; + std::array values{}; + std::array codes{}, extra{}; + float (*volatile compute_pointer + )(const api::QueryState&, api::CodeView, const api::Factors&) = &api::compute; + alignas(float) std::array misaligned{}; + const auto* bad_alignment = reinterpret_cast(misaligned.data() + 1); + EXPECT_THROW( + api::prepare_query({dim, 4}, bad_alignment, nullptr), std::invalid_argument + ); + EXPECT_THROW( + api::prepare_query({dim, 4}, values.data(), bad_alignment), std::invalid_argument + ); + for (size_t bits : {1U, 4U, 9U}) { + for (auto metric : {METRIC_L2, METRIC_IP}) { + const api::Parameters parameters{dim, bits, metric}; + const auto query = api::prepare_query(parameters, values.data(), nullptr); + for (auto layout : + {api::CodeLayout::SplitPacked, api::CodeLayout::UnpackedU16}) { + const auto sizes = api::requirements(parameters, layout); + const api::CodeView valid{ + layout, + {codes.data(), sizes.code_bytes}, + {extra.data(), sizes.extra_bytes}}; + EXPECT_EQ( + compute_pointer(query, valid, {}), api::compute(query, valid, {}) + ); + auto invalid = valid; + invalid.codes.data = nullptr; + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.codes.bytes--; + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.codes.data = reinterpret_cast( + std::numeric_limits::max() - sizes.code_bytes + 1 + ); + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.layout = static_cast(255); + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.extra.data = nullptr; + if (sizes.extra_bytes != 0) { + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.extra.bytes--; + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + invalid = valid; + invalid.extra.data = reinterpret_cast( + std::numeric_limits::max() - sizes.extra_bytes + 1 + ); + EXPECT_THROW(api::compute(query, invalid, {}), std::invalid_argument); + } else { + EXPECT_NO_THROW(api::compute(query, invalid, {})); + } + for (float bad : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + EXPECT_THROW( + api::compute(query, valid, {bad, 0, 0}), std::invalid_argument + ); + EXPECT_THROW( + api::compute(query, valid, {0, bad, 0}), std::invalid_argument + ); + EXPECT_THROW( + api::compute(query, valid, {0, 0, bad}), std::invalid_argument + ); + } + } + } + } + values.fill(1e18F); + const auto large_query = api::prepare_query({dim, 1}, values.data(), nullptr); + EXPECT_THROW( + api::compute( + large_query, + {api::CodeLayout::SplitPacked, {codes.data(), dim / 8}, {}}, + {std::numeric_limits::max(), 0, 0} + ), + std::overflow_error + ); + std::array unpadded{}; + const auto raw_query = api::prepare_query({65, 4}, unpadded.data(), nullptr); + EXPECT_THROW( + api::compute(raw_query, {api::CodeLayout::SplitPacked, {}, {}}, {}), + std::invalid_argument + ); +} + +TEST(CoreApi, ConcurrentCallsUseOnlyTheirOwnWorkspace) { + constexpr size_t dim = 96; + std::array data{}, query{}; + for (size_t i = 0; i < dim; ++i) { + data[i] = std::sin(static_cast(i)); + query[i] = std::cos(static_cast(i)); + } + api::Parameters p{dim, 9}; + auto r = api::requirements(p, api::CodeLayout::SplitPacked); + auto q = api::prepare_query(p, query.data(), nullptr); + std::array results{}; + std::vector workers; + for (size_t i = 0; i < results.size(); ++i) { + workers.emplace_back([&, i] { + Workspace scratch(r.workspace_bytes); + std::vector sign(r.code_bytes), extra(r.extra_bytes); + api::CodeOutput out{ + api::CodeLayout::SplitPacked, + {sign.data(), sign.size()}, + {extra.data(), extra.size()}}; + for (size_t run = 0; run < 20; ++run) { + auto f = api::encode(p, data.data(), nullptr, out, scratch.view()); + results[i] = api::compute(q, view(out), f); + } + }); + } + for (auto& worker : workers) + worker.join(); + for (float result : results) + EXPECT_EQ(result, results[0]); +} +} // namespace diff --git a/tests/unit/rabitqlib/quantization/data_layout_test.cpp b/tests/unit/rabitqlib/quantization/data_layout_test.cpp index c7b6b046..9a89bcf3 100644 --- a/tests/unit/rabitqlib/quantization/data_layout_test.cpp +++ b/tests/unit/rabitqlib/quantization/data_layout_test.cpp @@ -1,16 +1,16 @@ -#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/core/data_layout.hpp" #include -#include #include #include #include #include +#include +#include #include -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/core/encoder.hpp" namespace rabitqlib { namespace { @@ -21,63 +21,11 @@ float load_float(const char* data) { return value; } -TEST(DataLayoutTest, BatchFactorsKeepPackedOffsetsOnUnalignedStorage) { - constexpr size_t kDim = 64; - constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; - constexpr size_t kDataBytes = kCodeBytes + (3 * fastscan::kBatchSize * sizeof(float)); - const size_t data_bytes = BatchDataMap::data_bytes(kDim); - ASSERT_EQ(data_bytes, kCodeBytes + (3 * fastscan::kBatchSize * sizeof(float))); - - alignas(float) std::array storage; - storage.fill(char{0x5a}); - char* data = storage.data() + 1; - ASSERT_NE( - reinterpret_cast(data + kCodeBytes) % alignof(float), uintptr_t{0} - ); - BatchDataMap batch(data, kDim); - batch.f_add()[0] = 1.25F; - batch.f_rescale()[0] = -2.5F; - batch.f_error()[fastscan::kBatchSize - 1] = 3.75F; - - ConstBatchDataMap stored(data, kDim); - EXPECT_FLOAT_EQ(stored.f_add()[0], 1.25F); - EXPECT_FLOAT_EQ(stored.f_rescale()[0], -2.5F); - EXPECT_FLOAT_EQ(stored.f_error()[fastscan::kBatchSize - 1], 3.75F); - EXPECT_FLOAT_EQ(load_float(data + kCodeBytes), 1.25F); - EXPECT_FLOAT_EQ( - load_float(data + kCodeBytes + (fastscan::kBatchSize * sizeof(float))), -2.5F - ); - EXPECT_FLOAT_EQ( - load_float(data + kCodeBytes + ((3 * fastscan::kBatchSize - 1) * sizeof(float))), - 3.75F - ); - EXPECT_EQ(storage.front(), char{0x5a}); - EXPECT_EQ(storage.back(), char{0x5a}); -} - -TEST(DataLayoutTest, QgBatchFactorsKeepPackedOffsetsOnUnalignedStorage) { - constexpr size_t kDim = 64; - constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; - constexpr size_t kDataBytes = kCodeBytes + (2 * fastscan::kBatchSize * sizeof(float)); - const size_t data_bytes = QGBatchDataMap::data_bytes(kDim); - ASSERT_EQ(data_bytes, kCodeBytes + (2 * fastscan::kBatchSize * sizeof(float))); - - alignas(float) std::array storage{}; - char* data = storage.data() + 1; - ASSERT_NE( - reinterpret_cast(data + kCodeBytes) % alignof(float), uintptr_t{0} - ); - QGBatchDataMap batch(data, kDim); - std::fill_n(batch.f_add(), fastscan::kBatchSize, 4.25F); - std::fill_n(batch.f_rescale(), fastscan::kBatchSize, -5.5F); - - ConstQGBatchDataMap stored(data, kDim); - EXPECT_FLOAT_EQ(stored.f_add()[fastscan::kBatchSize - 1], 4.25F); - EXPECT_FLOAT_EQ(stored.f_rescale()[fastscan::kBatchSize - 1], -5.5F); - EXPECT_FLOAT_EQ(load_float(data + kCodeBytes), 4.25F); - EXPECT_FLOAT_EQ( - load_float(data + kCodeBytes + (fastscan::kBatchSize * sizeof(float))), -5.5F - ); +TEST(DataLayoutTest, ExtraCodeSizeRejectsOverflow) { + const size_t max_size = std::numeric_limits::max(); + EXPECT_EQ(ExDataMap::data_bytes(max_size, 0), 0); + EXPECT_THROW(ExDataMap::data_bytes(max_size / 8 + 1, 8), std::length_error); + EXPECT_EQ(ExDataMap::data_bytes(max_size / 8, 8), max_size / 8 + 8); } TEST(DataLayoutTest, SingleCodeFactorsKeepPackedOffsetsOnUnalignedStorage) { diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index 7a5dca42..20de7373 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -1,98 +1,230 @@ -#include "rabitqlib/quantization/rabitq.hpp" - #include #include #include #include +#include #include #include #include #include #include +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/query.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/warmup_space.hpp" + namespace rabitqlib::quant { namespace { -template -void check_one_bit_batches() { - // Alternate dimensions to exercise scratch growth and shrinking on each thread. - for (size_t dim : {64U, 8U, 128U, 16U, 768U, 24U, 192U, 56U, 72U, 64U}) { - for (size_t num : {0U, 1U, 15U, 16U, 17U, 31U, 32U, 33U, 64U, 65U}) { - for (MetricType metric : {METRIC_L2, METRIC_IP}) { - SCOPED_TRACE( - ::testing::Message() - << "dim=" << dim << " num=" << num << " metric=" << metric +struct DispatchedEstimateKernel { + static constexpr auto mask_ip_x0_q = + static_cast( + &rabitqlib::mask_ip_x0_q + ); + static constexpr auto warmup_ip_x0_q_512 = static_cast< + float (*)(const uint64_t*, const uint64_t*, float, float, size_t, size_t)>( + &rabitqlib::warmup_ip_x0_q_512 + ); +}; + +TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { + const auto check = [](const auto& result, float distance, float lower, float ip) { + EXPECT_EQ(result.distance, distance); + EXPECT_EQ(result.lower_bound, lower); + EXPECT_EQ(result.inner_product, ip); + }; + for (size_t dim : {64U, 96U}) { + std::vector data(dim), centroid(dim), query(dim); + for (size_t i = 0; i < dim; ++i) { + data[i] = static_cast(i % 11) * 0.25F; + centroid[i] = static_cast(i % 3) * 0.125F; + query[i] = static_cast(i % 7) * -0.25F; + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + const float g_add = metric == METRIC_L2 ? 9.0F : -2.0F; + constexpr float g_error = 3.0F; + for (size_t bits = 1; bits <= 9; ++bits) { + const size_t ex_bits = bits - 1; + std::vector bin(BinDataMap::data_bytes(dim)); + std::vector extra(ExDataMap::data_bytes(dim, ex_bits)); + quantize_split_single( + data.data(), + centroid.data(), + dim, + ex_bits, + bin.data(), + extra.data(), + metric ); - std::vector centroid(dim); - std::vector data(num * dim); - for (size_t d = 0; d < dim; ++d) { - centroid[d] = static_cast(static_cast(d % 7) - 3) / 8; - for (size_t i = 0; i < num; ++i) { - const T residual = - i % 3 == 0 - ? T{0} - : static_cast(static_cast((i + d) % 11) - 5) / 4; - data[i * dim + d] = centroid[d] + residual; - } + SplitSingleQuery prepared( + query.data(), dim, ex_bits, RabitqConfig{}, metric + ); + float distance, lower, ip; + split_single_estdist( + bin.data(), prepared, dim, ip, distance, lower, g_add, g_error + ); + check( + estimate_split_single(bin.data(), prepared, dim, g_add, g_error), + distance, + lower, + ip + ); + split_single_estdist_direct( + bin.data(), prepared, dim, ip, distance, lower, g_add, g_error + ); + check( + estimate_split_single_direct( + bin.data(), prepared, dim, g_add, g_error + ), + distance, + lower, + ip + ); + if (ex_bits == 0) { + continue; } - // A guard byte/factor also verifies empty inputs and tail writes. - const size_t packed_bytes = ((num + 31) / 32) * 32 * (dim / 8); - std::vector expected(packed_bytes + 1, 0xA5); - std::vector actual = expected; - std::vector compact(num * dim / 8); - std::vector expected_add(num + 1, T{123}); - std::vector expected_rescale = expected_add; - std::vector expected_error = expected_add; - std::vector actual_add = expected_add; - std::vector actual_rescale = expected_add; - std::vector actual_error = expected_add; - for (size_t i = 0; i < num; ++i) { - rabitq_impl::one_bit::one_bit_compact_code( - data.data() + i * dim, - centroid.data(), + auto ip_func = select_excode_ipfunc(ex_bits); + split_single_fulldist( + bin.data(), + extra.data(), + ip_func, + prepared, + dim, + ex_bits, + distance, + lower, + ip, + g_add, + g_error + ); + check( + estimate_split_single_full( + bin.data(), + extra.data(), + ip_func, + prepared, dim, - compact.data() + i * dim / 8, - expected_add[i], - expected_rescale[i], - expected_error[i], - metric - ); - } - fastscan::pack_codes(dim, compact.data(), num, expected.data()); - rabitq_impl::one_bit::one_bit_batch_code( + ex_bits, + g_add, + g_error + ), + distance, + lower, + ip + ); + split_single_fulldist_direct( + bin.data(), + extra.data(), + ip_func, + prepared, + dim, + ex_bits, + distance, + lower, + ip, + g_add, + g_error + ); + check( + estimate_split_single_full_direct( + bin.data(), + extra.data(), + ip_func, + prepared, + dim, + ex_bits, + g_add, + g_error + ), + distance, + lower, + ip + ); + } + for (auto widths : + {std::pair{1, 0}, {1, 4}, {4, 0}, {4, 4}, {8, 1}}) { + const auto [base_bits, ex_bits] = widths; + std::vector base(BaseDataMap::data_bytes(dim, base_bits)); + std::vector extra(ExDataMap::data_bytes(dim, ex_bits)); + quantize_xy_single( data.data(), centroid.data(), - num, dim, - actual.data(), - actual_add.data(), - actual_rescale.data(), - actual_error.data(), + base_bits, + ex_bits, + base.data(), + extra.data(), metric ); - EXPECT_EQ(actual, expected); - EXPECT_EQ(actual_add, expected_add); - EXPECT_EQ(actual_rescale, expected_rescale); - EXPECT_EQ(actual_error, expected_error); + SplitSingleQuery prepared( + query.data(), dim, ex_bits, RabitqConfig{}, metric, base_bits + ); + auto base_ip = select_excode_ipfunc(base_bits); + auto extra_ip = select_excode_ipfunc(ex_bits); + float distance, lower, ip; + xy_single_base_dist( + base.data(), + base_ip, + prepared, + dim, + base_bits, + ip, + distance, + lower, + g_add, + g_error + ); + check( + estimate_xy_base( + base.data(), base_ip, prepared, dim, base_bits, g_add, g_error + ), + distance, + lower, + ip + ); + xy_single_full_dist( + base.data(), + extra.data(), + base_ip, + extra_ip, + prepared, + dim, + base_bits, + ex_bits, + distance, + lower, + ip, + g_add, + g_error + ); + check( + estimate_xy_full( + base.data(), + extra.data(), + base_ip, + extra_ip, + prepared, + dim, + base_bits, + ex_bits, + g_add, + g_error + ), + distance, + lower, + ip + ); } } } } -TEST(RabitqOneBitBatchTest, MatchesIndividualEncoding) { - check_one_bit_batches(); - check_one_bit_batches(); -} - -TEST(RabitqOneBitBatchTest, ParallelMatchesIndividualEncoding) { - const int previous_threads = omp_get_max_threads(); - omp_set_num_threads(2); - check_one_bit_batches(); - check_one_bit_batches(); - omp_set_num_threads(previous_threads); -} - int level_from_thresholds(double magnitude, double t, int max_code) { int result = 0; if (magnitude > 0) { diff --git a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp index c9a4196b..d13ba642 100644 --- a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp +++ b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp @@ -9,14 +9,16 @@ #include #include -#include "rabitqlib/quantization/rabitq_impl.hpp" -#include "rabitqlib/simd/quantization_dispatch.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" namespace rabitqlib::quant { namespace { using RescaleSearch = double (*)(const float*, size_t, int, double, double); +using WorkspaceRescaleSearch = + double (*)(const float*, size_t, int, double, double, double*, double*); std::vector levels_at_threshold( const std::vector& magnitudes, double scale, int max_code @@ -134,6 +136,15 @@ class RabitqRescaleBackendTest : public ::testing::TestWithParam { : simd::best_rescale_factor_avx2; #else return simd::best_rescale_factor_generic; +#endif + } + + WorkspaceRescaleSearch workspace_search() const { +#if defined(__x86_64__) || defined(_M_X64) + return GetParam() ? simd::best_rescale_factor_with_workspace_avx512 + : simd::best_rescale_factor_with_workspace_avx2; +#else + return simd::best_rescale_factor_with_workspace_generic; #endif } }; @@ -171,6 +182,23 @@ TEST_P(RabitqRescaleBackendTest, MatchesExhaustiveObjectiveAndCanonicalThreshold const auto [start, end] = search_interval(magnitudes, bits); const int max_code = (1 << bits) - 1; const double t = search()(magnitudes.data(), dim, max_code, start, end); + std::vector promoted(dim + 2, -123.0); + std::vector reciprocals(dim + 2, -123.0); + EXPECT_EQ( + workspace_search( + )(magnitudes.data(), + dim, + max_code, + start, + end, + promoted.data() + 1, + reciprocals.data() + 1), + t + ); + EXPECT_EQ(promoted.front(), -123.0); + EXPECT_EQ(promoted.back(), -123.0); + EXPECT_EQ(reciprocals.front(), -123.0); + EXPECT_EQ(reciprocals.back(), -123.0); ASSERT_TRUE(std::isfinite(t)); if (t >= 0) expect_optimal_canonical_scale(magnitudes, max_code, start, end, t); @@ -214,7 +242,9 @@ TEST_P(RabitqRescaleBackendTest, ReusesScratchAcrossDimensionsAndThreads) { } size_t mismatches = 0; +#if RABITQ_USE_OPENMP #pragma omp parallel for num_threads(8) schedule(dynamic, 1) reduction(+ : mismatches) +#endif for (std::ptrdiff_t call = 0; call < 160; ++call) { const size_t index = static_cast(call) % inputs.size(); const auto& magnitudes = inputs[index]; diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp index fa70d43e..2697e637 100644 --- a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -1,15 +1,19 @@ #include #include +#include +#include #include +#include +#include #include -#include "rabitqlib/defines.hpp" -#include "rabitqlib/index/estimator.hpp" -#include "rabitqlib/index/query.hpp" -#include "rabitqlib/quantization/data_layout.hpp" -#include "rabitqlib/quantization/rabitq.hpp" -#include "rabitqlib/quantization/rabitq_impl.hpp" +#include "rabitqlib/core/data_layout.hpp" +#include "rabitqlib/core/detail/rabitq_impl.hpp" +#include "rabitqlib/core/encoder.hpp" +#include "rabitqlib/core/estimator.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/core/query.hpp" #include "rabitqlib/utils/space.hpp" #include "test_helpers.hpp" @@ -27,11 +31,9 @@ std::vector RandomVec(size_t dim, std::mt19937& gen) { return v; } -// The function writes packed codes into the data blocks and no unpacker exists. -// Recover the raw per-dimension codes through the library's own reader instead -// of reimplementing seven SIMD bit layouts: a one-hot query makes the ex-code -// inner product return exactly code[j]. This asserts the codes as the search -// path actually interprets them, which is the property that matters. +// A one-hot query makes the ex-code inner product return exactly code[j]. +// This checks how scoring interprets the packed codes independently of the +// library's unpacker, rather than only round-tripping through paired code paths. std::vector UnpackViaIpKernel(const uint8_t* packed, size_t dim, size_t bits) { std::vector out(dim, 0); if (bits == 0) { @@ -75,8 +77,8 @@ XySplitResult SplitCode( extra_bits > 0 ? ExDataMap::data_bytes(dim, extra_bits) : 1, 0 ); - // Through the public entry point: rabitq_impl now takes raw buffers, and - // rabitq.hpp is what maps them onto BaseDataMap/ExDataMap. + // The public entry point in core/encoder.hpp maps the raw buffers onto + // BaseDataMap/ExDataMap. quant::quantize_xy_single( data.data(), centroid.data(), diff --git a/tests/unit/rabitqlib/utils/array_test.cpp b/tests/unit/rabitqlib/utils/array_test.cpp new file mode 100644 index 00000000..36912b91 --- /dev/null +++ b/tests/unit/rabitqlib/utils/array_test.cpp @@ -0,0 +1,101 @@ +#include "rabitqlib/utils/array.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct Allocations { + std::map owners; + bool mismatch = false; + bool fail = false; +}; + +template +struct TrackedAllocator { + using value_type = int; + using propagate_on_container_move_assignment = std::bool_constant; + using is_always_equal = std::false_type; + Allocations* allocations; + int id; + + int* allocate(size_t count) { + if (allocations->fail) { + throw std::bad_alloc(); + } + auto* ptr = std::allocator{}.allocate(count); + allocations->owners[ptr] = id; + return ptr; + } + void deallocate(int* ptr, size_t count) { + allocations->mismatch |= allocations->owners.at(ptr) != id; + allocations->owners.erase(ptr); + std::allocator{}.deallocate(ptr, count); + } + bool operator!=(const TrackedAllocator& other) const { return id != other.id; } +}; + +template +void check_move(int destination_id) { + Allocations state; + using Array = rabitqlib::Array, TrackedAllocator>; + static_assert(std::is_nothrow_move_assignable_v == Propagate); + { + Array source({3}, {&state, 1}); + Array target({2}, {&state, destination_id}); + source[0] = 17; + source[1] = 23; + source[2] = 42; + auto* previous = source.data(); + target = std::move(source); + EXPECT_EQ(target[0], 17); + EXPECT_EQ(target[1], 23); + EXPECT_EQ(target[2], 42); + // A successful move explicitly leaves the source pointer null. + // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) + EXPECT_EQ(source.data(), nullptr); + EXPECT_EQ(target.data() == previous, Propagate || destination_id == 1); + auto& alias = target; + target = std::move(alias); + EXPECT_EQ(target[2], 42); + } + EXPECT_TRUE(state.owners.empty()); + EXPECT_FALSE(state.mismatch); +} + +TEST(ArrayOwnership, MoveAssignmentHonorsAllocator) { + check_move(1); + check_move(2); + check_move(1); + check_move(2); +} + +TEST(ArrayOwnership, FailedUnequalMovePreservesBothArrays) { + Allocations state; + using Array = rabitqlib::Array, TrackedAllocator>; + { + Array source({3}, {&state, 1}), target({2}, {&state, 2}); + source[0] = 17; + target[0] = 23; + const auto* source_ptr = source.data(); + const auto* target_ptr = target.data(); + state.fail = true; + EXPECT_THROW(target = std::move(source), std::bad_alloc); + state.fail = false; + // The failed allocation must leave the source untouched. + // NOLINTNEXTLINE(bugprone-use-after-move,clang-analyzer-cplusplus.Move) + EXPECT_EQ(source.data(), source_ptr); + EXPECT_EQ(target.data(), target_ptr); + EXPECT_EQ(source[0], 17); + EXPECT_EQ(target[0], 23); + } + EXPECT_TRUE(state.owners.empty()); + EXPECT_FALSE(state.mismatch); +} +} // namespace diff --git a/tests/unit/rabitqlib/utils/buffer_test.cpp b/tests/unit/rabitqlib/utils/buffer_test.cpp index 0222a5f5..ff3fb550 100644 --- a/tests/unit/rabitqlib/utils/buffer_test.cpp +++ b/tests/unit/rabitqlib/utils/buffer_test.cpp @@ -1,4 +1,4 @@ -#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/index/detail/buffer.hpp" #include @@ -6,7 +6,7 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/types.hpp" namespace rabitqlib::buffer { namespace { diff --git a/tests/unit/rabitqlib/utils/excode_ip_test.cpp b/tests/unit/rabitqlib/utils/excode_ip_test.cpp index b3cf394d..874f1c9c 100644 --- a/tests/unit/rabitqlib/utils/excode_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/excode_ip_test.cpp @@ -7,14 +7,23 @@ #include #include -#include "rabitqlib/simd/pack_excode_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { using namespace rabitqlib; + const std::array wrappers{ + excode_ipimpl::ip16_fxu1_avx, + excode_ipimpl::ip64_fxu2_avx, + excode_ipimpl::ip64_fxu3_avx, + excode_ipimpl::ip16_fxu4_avx, + excode_ipimpl::ip64_fxu5_avx, + excode_ipimpl::ip64_fxu6_avx, + excode_ipimpl::ip64_fxu7_avx, + }; #if defined(__x86_64__) || defined(_M_X64) const std::array avx2_functions{ simd::excode_ipimpl::ip16_fxu1_avx2, @@ -62,10 +71,12 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { for (size_t bits = 1; bits <= 8; ++bits) { SCOPED_TRACE(bits); const auto max_code = static_cast((1U << bits) - 1); - for (size_t dim : std::array{ - 0, 64, 96, 128, 160, 192, 256, 544, 576, 960, 1024, 4096}) { + for (size_t dim : std::array{ + 0, 64, 96, 128, 160, 192, 256, 544, 576, 768, 960, 1024, 4096}) { SCOPED_TRACE(dim); - for (size_t pattern = 0; pattern < 4; ++pattern) { + const size_t patterns = + dim == 256 && (bits == 3 || bits == 5 || bits == 7) ? 5 : 4; + for (size_t pattern = 0; pattern < patterns; ++pattern) { SCOPED_TRACE(pattern); std::vector query_storage(dim + 1); auto* query = query_storage.data() + 1; @@ -82,9 +93,13 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { if (pattern == 3) { query[i] = (i % 2 == 0) ? 1.0F : -1.0F; } - codes[i] = pattern >= 2 + codes[i] = pattern == 2 || pattern == 3 ? max_code : static_cast((i * 37U + 19U) & max_code); + // Preserve the odd-width bit-63 regression on otherwise mixed codes. + if (pattern == 4 && i % 64 == 63) { + codes[i] = max_code; + } const double product = static_cast(query[i]) * codes[i]; expected += product; sum_abs += std::abs(product); @@ -126,6 +141,13 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { EXPECT_NEAR( select_excode_ipfunc(bits)(query, compact, dim), expected, tolerance ); + // Named wrappers must route each width correctly too; their names + // do not imply AVX execution. This replaces the old integration fixture. + if (dim == 768 && pattern == 1 && bits <= wrappers.size()) { + EXPECT_NEAR( + wrappers[bits - 1](query, compact, dim), expected, tolerance + ); + } } } } diff --git a/tests/unit/rabitqlib/utils/float_distance_test.cpp b/tests/unit/rabitqlib/utils/float_distance_test.cpp index 1b2fcc0d..3e6fee45 100644 --- a/tests/unit/rabitqlib/utils/float_distance_test.cpp +++ b/tests/unit/rabitqlib/utils/float_distance_test.cpp @@ -5,7 +5,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" diff --git a/tests/unit/rabitqlib/utils/io_test.cpp b/tests/unit/rabitqlib/utils/io_test.cpp index 03a30078..2f1b10ab 100644 --- a/tests/unit/rabitqlib/utils/io_test.cpp +++ b/tests/unit/rabitqlib/utils/io_test.cpp @@ -10,7 +10,8 @@ #include #include -#include "rabitqlib/defines.hpp" +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/utils/path.hpp" namespace rabitqlib { namespace { diff --git a/tests/unit/rabitqlib/utils/mask_ip_test.cpp b/tests/unit/rabitqlib/utils/mask_ip_test.cpp index 52053a68..57e28eb2 100644 --- a/tests/unit/rabitqlib/utils/mask_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/mask_ip_test.cpp @@ -7,7 +7,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" diff --git a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp index e853395e..ceee33c2 100644 --- a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp +++ b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp @@ -1,4 +1,4 @@ -#include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/core/detail/simd/matrix_dispatch.hpp" #include diff --git a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp index 3757dcbd..f41e3cb5 100644 --- a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp +++ b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp @@ -1,3 +1,4 @@ +#if defined(__aarch64__) || defined(_M_ARM64) #include #include @@ -6,15 +7,17 @@ #include #include #include +#include #include +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/hnsw_dispatch.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" -#include "rabitqlib/simd/hnsw_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" -#if defined(__aarch64__) || defined(_M_ARM64) namespace rabitqlib { TEST(NeonHnsw, WarmupHandlesUnalignedCodesTailsAndWideCounts) { if (!cpu::has_neon()) { @@ -104,20 +107,20 @@ TEST(NeonHnsw, SearchMatchesGenericOnSameGraphForBothMetricsAndAllWidths) { value = distribution(random); } for (size_t ef : {10U, 64U}) { - auto expected = - hnsw::detail::search_knn_generic(index, query.data(), 5, ef); - auto actual = - hnsw::detail::search_knn_neon(index, query.data(), 5, ef); + auto expected = hnsw::detail::search_knn_sorted_generic( + index, query.data(), 5, ef + ); + auto actual = hnsw::detail::search_knn_sorted_neon( + index, query.data(), 5, ef + ); ASSERT_EQ(expected.size(), actual.size()); - while (!expected.empty()) { - EXPECT_EQ(actual.top().second, expected.top().second); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(actual[i].second, expected[i].second); EXPECT_NEAR( - actual.top().first, - expected.top().first, - 2e-5F * std::max(1.0F, std::abs(expected.top().first)) + actual[i].first, + expected[i].first, + 2e-5F * std::max(1.0F, std::abs(expected[i].first)) ); - actual.pop(); - expected.pop(); } } } diff --git a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp index d0559c09..fe091d6e 100644 --- a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp +++ b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp @@ -7,10 +7,10 @@ #include #include -#include "rabitqlib/quantization/pack_excode.hpp" -#include "rabitqlib/simd/pack_excode_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/pack_excode.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/warmup_space.hpp" diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp index 6904ebe3..21500ce6 100644 --- a/tests/unit/rabitqlib/utils/query_preparation_test.cpp +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -9,7 +9,7 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" @@ -113,8 +113,9 @@ void check_transpose( TEST(QueryPreparation, QuantizationMatchesHalfAwayRoundingAndTails) { check_quantization(simd::scalar_quantize_uint8_generic); check_quantization(simd::scalar_quantize_uint16_generic); - check_quantization(simd::scalar_quantize_uint8); - check_quantization(simd::scalar_quantize_uint16); + // Exercise the public templates as well as the kernels they dispatch to. + check_quantization(scalar_quantize); + check_quantization(scalar_quantize); #if defined(__aarch64__) || defined(_M_ARM64) if (cpu::has_neon()) { check_quantization(simd::scalar_quantize_uint8_neon); diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 506ed671..4f1a7cb6 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -1,7 +1,8 @@ -#include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/core/rotator.hpp" #include +#include #include #include #include @@ -12,6 +13,7 @@ #include #include +#include "rabitqlib/core/detail/simd/rotator_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/memory.hpp" #include "test_data.hpp" @@ -583,3 +585,31 @@ TEST(FhtDispatchTest, AvxBackendsMatchBitwiseForCompactTails) { } } #endif + +TEST(MatrixRotatorFactory, DoublePreservesNormAndRejectsFloatOnlyRotation) { + const size_t dim = 64; + std::unique_ptr> rotator( + choose_rotator(dim, RotatorType::MatrixRotator) + ); + ASSERT_NE(rotator, nullptr); + std::vector input(dim), output(rotator->size()); + double expected = 0; + for (size_t i = 0; i < dim; ++i) { + input[i] = static_cast(i % 7) - 3; + expected += input[i] * input[i]; + } + rotator->rotate(input.data(), output.data()); + double actual = 0; + for (double x : output) { + actual += x * x; + } + EXPECT_NEAR(actual, expected, 1e-10 * expected); + try { + std::unique_ptr> unsupported( + choose_rotator(dim, RotatorType::FhtKacRotator) + ); + FAIL() << "double FHT must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "FhtKacRotator only supports float"); + } +} diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index aa7ec38a..9748b1cf 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -2,22 +2,35 @@ #include +#include #include #include +#include #include +#include #include #include #include -#include "rabitqlib/simd/pack_excode_dispatch.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" -#include "rabitqlib/simd/warmup_dispatch.hpp" +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/warmup_space.hpp" using namespace rabitqlib; +TEST(SpaceHelpers, CentroidSupportsRequestedWorkerCounts) { + const std::array data{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + for (size_t workers : {0U, 1U, 2U, 64U}) { + const auto centroid = compute_centroid(data.data(), 4, 3, workers); + ASSERT_EQ(centroid.size(), 3U); + EXPECT_FLOAT_EQ(centroid[0], 5.5F); + EXPECT_FLOAT_EQ(centroid[1], 6.5F); + EXPECT_FLOAT_EQ(centroid[2], 7.5F); + } +} + TEST(BitOps, Popcount64MatchesBitReference) { const auto check = [](uint64_t value) { unsigned expected = 0; @@ -316,46 +329,6 @@ TEST(Select_IP_Func, zero_ex_bits_contributes_nothing) { ASSERT_EQ(ip_func(query.data(), codes.data(), dim), 0.0F); } -TEST(ScalarQuantize, Uint8MatchesRoundedScalar) { - constexpr size_t dim = 37; - constexpr float lo = -3.0F; - constexpr float delta = 0.25F; - std::vector input(dim); - std::vector result(dim); - std::vector expected(dim); - - for (size_t i = 0; i < dim; ++i) { - float quantized = static_cast((i * 7) % 251) + - static_cast(static_cast(i % 3) - 1) * 0.2F; - input[i] = lo + delta * quantized; - expected[i] = static_cast(std::round((input[i] - lo) / delta)); - } - - scalar_quantize(result.data(), input.data(), dim, lo, delta); - - ASSERT_EQ(result, expected); -} - -TEST(ScalarQuantize, Uint16MatchesRoundedScalar) { - constexpr size_t dim = 41; - constexpr float lo = 2.0F; - constexpr float delta = 0.125F; - std::vector input(dim); - std::vector result(dim); - std::vector expected(dim); - - for (size_t i = 0; i < dim; ++i) { - float quantized = static_cast(1000 + i * 317) + - static_cast(static_cast(i % 5) - 2) * 0.1F; - input[i] = lo + delta * quantized; - expected[i] = static_cast(std::round((input[i] - lo) / delta)); - } - - scalar_quantize(result.data(), input.data(), dim, lo, delta); - - ASSERT_EQ(result, expected); -} - constexpr size_t kScalarQuantizeTestDim = 33; TEST(ScalarQuantize, HalfValuesMatchScalarAcrossVectorBoundaries) { @@ -444,114 +417,6 @@ TEST(ip64_fxu2_avx, ip_works) { ASSERT_NEAR(rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), expected, 0.1F); } -TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { - constexpr size_t dim = 64 * 4; - std::vector query(dim); - std::vector codes(dim); - - for (size_t bits : std::array{3, 5, 7}) { - const uint8_t max_code = static_cast((1U << bits) - 1U); - for (size_t i = 0; i < dim; ++i) { - query[i] = static_cast(static_cast(i % 23) - 11) / 7.0F; - codes[i] = static_cast((i * 37U + 19U) & max_code); - } - // Exercise the high bit of each packed 64-value block, including bit 63 of the - // scalar word used by the SIMD unpacking path. - for (size_t i = 63; i < dim; i += 64) { - codes[i] = max_code; - } - - std::vector compact(dim * bits / 8); - if (bits == 3) { - simd::packing_3bit_excode(codes.data(), compact.data(), dim); - } else if (bits == 5) { - simd::packing_5bit_excode(codes.data(), compact.data(), dim); - } else { - simd::packing_7bit_excode(codes.data(), compact.data(), dim); - } - - double expected = 0.0; - for (size_t i = 0; i < dim; ++i) { - expected += static_cast(query[i]) * static_cast(codes[i]); - } - const float expected_float = static_cast(expected); - -#if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) { - const std::array avx2_functions{ - nullptr, - simd::excode_ipimpl::ip16_fxu1_avx2, - simd::excode_ipimpl::ip64_fxu2_avx2, - simd::excode_ipimpl::ip64_fxu3_avx2, - simd::excode_ipimpl::ip16_fxu4_avx2, - simd::excode_ipimpl::ip64_fxu5_avx2, - simd::excode_ipimpl::ip64_fxu6_avx2, - simd::excode_ipimpl::ip64_fxu7_avx2, - }; - ASSERT_NEAR( - avx2_functions[bits](query.data(), compact.data(), dim), - expected_float, - 0.1F - ); - } -#endif -#if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx512_core()) { - const std::array avx512_functions{ - nullptr, - simd::excode_ipimpl::ip16_fxu1_avx512, - simd::excode_ipimpl::ip64_fxu2_avx512, - simd::excode_ipimpl::ip64_fxu3_avx512, - simd::excode_ipimpl::ip16_fxu4_avx512, - simd::excode_ipimpl::ip64_fxu5_avx512, - simd::excode_ipimpl::ip64_fxu6_avx512, - simd::excode_ipimpl::ip64_fxu7_avx512, - }; - ASSERT_NEAR( - avx512_functions[bits](query.data(), compact.data(), dim), - expected_float, - 0.1F - ); - } -#endif - } -} - -TEST(ip_fxu8_avx, ip_works) { - constexpr size_t dim = 1024; - std::vector query(dim); - std::vector codes(dim); - double expected = 0.0; - - for (size_t i = 0; i < dim; ++i) { - query[i] = static_cast(i % 97) / 17.0F; - codes[i] = static_cast(i % 251); - expected += static_cast(query[i]) * static_cast(codes[i]); - } - - const float expected_float = static_cast(expected); - ex_ipfunc ip_func = select_excode_ipfunc(8); - ASSERT_NEAR(ip_func(query.data(), codes.data(), dim), expected_float, 0.1F); -#if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) { - ASSERT_NEAR( - simd::excode_ipimpl::ip16_fxu8_avx2(query.data(), codes.data(), dim), - expected_float, - 0.1F - ); - } -#endif -#if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx512_core()) { - ASSERT_NEAR( - simd::excode_ipimpl::ip16_fxu8_avx512(query.data(), codes.data(), dim), - expected_float, - 0.1F - ); - } -#endif -} - TEST(PackBinary, CompactTailWordsMatchBinaryArithmetic) { for (size_t dim : {64U, 96U, 128U, 160U}) { std::vector signs(dim); diff --git a/tests/unit/rabitqlib/utils/thread_count_test.cpp b/tests/unit/rabitqlib/utils/thread_count_test.cpp index 932f0886..5ff7a7fa 100644 --- a/tests/unit/rabitqlib/utils/thread_count_test.cpp +++ b/tests/unit/rabitqlib/utils/thread_count_test.cpp @@ -1,6 +1,8 @@ #include #include +#include + #if defined(__linux__) #include #endif @@ -11,7 +13,7 @@ #include #include -#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/detail/parallel_for.hpp" #include "rabitqlib/utils/tools.hpp" TEST(ThreadCount, AutomaticAndOversizedRequestsUseAvailableCount) { @@ -59,7 +61,7 @@ TEST(ThreadCount, AutomaticAndExplicitRequestsRespectAffinity) { for (size_t requested : {size_t{0}, size_t{1}, std::numeric_limits::max()}) { const size_t expected = requested == 1 ? 1 : available; EXPECT_EQ(rabitqlib::resolve_num_threads(requested), expected); - rabitqlib::ivf::parallel_for(0, 19, requested, [&](size_t, size_t worker) { + rabitqlib::index_detail::parallel_for(0, 19, requested, [&](size_t, size_t worker) { EXPECT_LT(worker, expected); }); } @@ -77,7 +79,7 @@ TEST(ThreadCount, BoundParallelForUsesRuntimePlacementAndPropagatesErrors) { // (for example, LLVM libomp on macOS). const int places = omp_get_num_places(); const size_t threads = std::min(available, 4); - rabitqlib::ivf::parallel_for(0, 257, threads, [&](size_t, size_t worker) { + rabitqlib::index_detail::parallel_for(0, 257, threads, [&](size_t, size_t worker) { EXPECT_LT(worker, threads); if (places > 0) { EXPECT_GE(omp_get_place_num(), 0); @@ -90,7 +92,7 @@ TEST(ThreadCount, BoundParallelForUsesRuntimePlacementAndPropagatesErrors) { } }); EXPECT_THROW( - rabitqlib::ivf::parallel_for( + rabitqlib::index_detail::parallel_for( 0, 257, threads, [](size_t, size_t) { throw std::runtime_error("failed"); } ), std::runtime_error @@ -109,22 +111,39 @@ TEST(ThreadCount, ParallelForVisitsEachItemOnceWithBoundedWorkerIds) { for (auto& visited : visits) { visited.store(0); } - rabitqlib::ivf::parallel_for(0, count, requested, [&](size_t i, size_t worker) { - EXPECT_LT(worker, requested == 1 ? 1U : available); - visits[i].fetch_add(1); - }); + rabitqlib::index_detail::parallel_for( + 0, + count, + requested, + [&](size_t i, size_t worker) { + EXPECT_LT(worker, requested == 1 ? 1U : available); + visits[i].fetch_add(1); + } + ); for (const auto& visited : visits) { EXPECT_EQ(visited.load(), 1U); } size_t calls = 0; - rabitqlib::ivf::parallel_for(7, 8, requested, [&](size_t i, size_t worker) { - EXPECT_EQ(i, 7U); - EXPECT_EQ(worker, 0U); - ++calls; - }); + rabitqlib::index_detail::parallel_for( + 7, + 8, + requested, + [&](size_t i, size_t worker) { + EXPECT_EQ(i, 7U); + EXPECT_EQ(worker, 0U); + ++calls; + } + ); EXPECT_EQ(calls, 1U); - rabitqlib::ivf::parallel_for(0, 0, requested, [](size_t, size_t) { - ADD_FAILURE() << "An empty range must not invoke the callback"; - }); + for (size_t start : {0U, 5U}) { + rabitqlib::index_detail::parallel_for( + start, + start, + requested, + [](size_t, size_t) { + ADD_FAILURE() << "An empty range must not invoke the callback"; + } + ); + } } } diff --git a/tests/unit/rabitqlib/utils/visited_set_test.cpp b/tests/unit/rabitqlib/utils/visited_set_test.cpp index 61573b01..1768369c 100644 --- a/tests/unit/rabitqlib/utils/visited_set_test.cpp +++ b/tests/unit/rabitqlib/utils/visited_set_test.cpp @@ -1,4 +1,4 @@ -#include "rabitqlib/utils/visited_set.hpp" +#include "rabitqlib/index/detail/visited_set.hpp" #include @@ -7,9 +7,14 @@ #include #include #include +#include #include -#include "rabitqlib/utils/hashset.hpp" +#include "rabitqlib/index/detail/hashset.hpp" +#include "rabitqlib/index/detail/visited_pool.hpp" +#include "rabitqlib/index/detail/visited_set_epoch.hpp" +#include "rabitqlib/index/detail/visited_set_hash.hpp" +#include "rabitqlib/index/types.hpp" namespace { @@ -28,6 +33,20 @@ static_assert( "HashBasedBooleanSet must remain available as a compatibility alias" ); +TEST(VisitedListPoolTest, TransfersOwnershipAndClearsReusedStorage) { + rabitqlib::VisitedListPool pool(1, 64); + auto visited = pool.get_free_vislist(); + auto* storage = visited.get(); + visited->set(7); + pool.release_vis_list(std::move(visited)); + EXPECT_EQ(visited, nullptr); + auto reused = pool.get_free_vislist(); + EXPECT_EQ(reused.get(), storage); + EXPECT_FALSE(reused->get(7)); + auto another = pool.get_free_vislist(); + EXPECT_NE(another.get(), reused.get()); +} + template void expect_basic_visited_set_behavior() { constexpr size_t kNumElements = 1024;