From 611562b5516d017f0212d4d4196d68e4c28e40af Mon Sep 17 00:00:00 2001 From: Raven Date: Thu, 24 Sep 2026 12:52:12 +0800 Subject: [PATCH 1/4] Add native macOS ARM64 support and NEON kernels Build IVF, HNSW, SymphonyQG and Python bindings on ARM64 without x86 translation units. Add portable reference kernels and NEON distance, FastScan, packed-code, rotation and query-preparation implementations while retaining existing metrics, layouts and persisted formats. Cap thread counts at available hardware and treat zero as automatic. Separate Faiss clustering from example indexing processes. Add native Apple Silicon CI and CPython 3.11-3.14 arm64 wheel builds with bundled OpenMP and its license. Keep the package version at 0.3.8. Validation on Apple M5 Pro, macOS 27.0, Apple Clang 21, native tuning off: - C++: 178 passed, 9 skipped; installed CMake consumer passed. - Working Python environment: 391 passed, 3 skipped. - Four repaired installed wheels: 392 passed, 2 skipped per ABI. - Formatting, Ruff, ShellCheck and 21 CI-routing tests passed. - Full configured clang-tidy passed with the LLVM 21 compatibility wrapper. - Query-preparation benchmarks and fixed-index search comparisons retained identical returned IDs and distances for the measured workload. Linux AArch64 and current x86 backend execution remain to be validated on their respective platforms. --- .github/scripts/ci_changes.py | 4 +- .github/scripts/test_ci_changes.py | 8 +- .github/workflows/release.yml | 35 +- .github/workflows/test.yaml | 33 ++ AGENTS.md | 8 +- CMakeLists.txt | 39 +- CONTRIBUTING.md | 30 +- DEVELOPMENT.md | 32 +- README.md | 26 +- cmake/rabitqOpenMP.cmake | 14 + cmake/rabitqlibConfig.cmake.in | 1 + example_py.sh | 8 +- include/rabitqlib/fastscan/fastscan.hpp | 1 + include/rabitqlib/index/hnsw/hnsw.hpp | 6 + include/rabitqlib/index/ivf/initializer.hpp | 6 +- include/rabitqlib/index/ivf/ivf.hpp | 2 +- .../rabitqlib/index/symqg/detail/pipnn.hpp | 2 +- include/rabitqlib/index/symqg/qg_builder.hpp | 2 +- include/rabitqlib/simd/fastscan_dispatch.hpp | 18 + include/rabitqlib/simd/hnsw_dispatch.hpp | 7 + .../rabitqlib/simd/pack_excode_dispatch.hpp | 6 + include/rabitqlib/simd/rotator_dispatch.hpp | 23 ++ include/rabitqlib/simd/space_dispatch.hpp | 67 ++++ include/rabitqlib/simd/warmup_dispatch.hpp | 33 ++ include/rabitqlib/utils/cpu_features.hpp | 2 + include/rabitqlib/utils/space.hpp | 4 +- include/rabitqlib/utils/tools.hpp | 7 + pyproject.toml | 11 +- python_bindings/CMakeLists.txt | 13 + python_bindings/licenses/LLVM-OpenMP.txt | 361 ++++++++++++++++++ python_bindings/symqg_bindings.cpp | 6 +- sample/python/README.md | 53 +++ sample/python/faiss_clustering.py | 66 ++++ sample/python/hnsw_rabitq_indexing.py | 18 +- sample/python/ivf_rabitq_indexing.py | 20 +- sample/python/utils.py | 70 ++-- src/index/hnsw_search_generic.cpp | 39 ++ src/index/hnsw_search_neon.cpp | 42 ++ src/simd/dispatch.cpp | 317 ++++++++------- src/simd/fastscan_generic.cpp | 62 +++ src/simd/fastscan_neon.cpp | 79 ++++ src/simd/hnsw_neon_kernels.hpp | 85 +++++ src/simd/portable_generic.cpp | 260 +++++++++++++ src/simd/query_neon.cpp | 134 +++++++ src/simd/rotator_neon.cpp | 152 ++++++++ src/simd/space_generic.cpp | 49 ++- src/simd/space_neon.cpp | 56 +++ src/simd/space_packed_neon.cpp | 160 ++++++++ src/utils/cpu_features.cpp | 6 + tests/CMakeLists.txt | 7 + tests/README.md | 38 +- tests/dispatch/fastscan_no_simd_test.cpp | 44 +++ tests/integration/bit_pack_unpack_test.cpp | 10 +- tests/python/test_examples.py | 219 +++++++++++ tests/python/test_ivf.py | 8 +- tests/python/test_portable_indexes.py | 90 +++++ tests/python/test_thread_counts.py | 54 +++ tests/python/test_wheel_macos.py | 39 ++ .../unit/rabitqlib/fastscan/fastscan_test.cpp | 58 ++- tests/unit/rabitqlib/index/qg_test.cpp | 33 +- .../quantization/rescale_search_test.cpp | 4 + .../rabitqlib/utils/cpu_features_test.cpp | 11 + tests/unit/rabitqlib/utils/excode_ip_test.cpp | 47 ++- .../rabitqlib/utils/float_distance_test.cpp | 141 +++---- tests/unit/rabitqlib/utils/mask_ip_test.cpp | 25 +- .../rabitqlib/utils/matrix_dispatch_test.cpp | 4 + tests/unit/rabitqlib/utils/neon_hnsw_test.cpp | 126 ++++++ .../rabitqlib/utils/portable_kernel_test.cpp | 139 +++++++ .../utils/query_preparation_test.cpp | 143 +++++++ tests/unit/rabitqlib/utils/rotator_test.cpp | 180 ++++++++- tests/unit/rabitqlib/utils/space_test.cpp | 28 +- .../rabitqlib/utils/thread_count_test.cpp | 51 +++ 72 files changed, 3541 insertions(+), 441 deletions(-) create mode 100644 cmake/rabitqOpenMP.cmake create mode 100644 python_bindings/licenses/LLVM-OpenMP.txt create mode 100644 sample/python/README.md create mode 100644 sample/python/faiss_clustering.py create mode 100644 src/index/hnsw_search_generic.cpp create mode 100644 src/index/hnsw_search_neon.cpp create mode 100644 src/simd/fastscan_neon.cpp create mode 100644 src/simd/hnsw_neon_kernels.hpp create mode 100644 src/simd/portable_generic.cpp create mode 100644 src/simd/query_neon.cpp create mode 100644 src/simd/rotator_neon.cpp create mode 100644 src/simd/space_neon.cpp create mode 100644 src/simd/space_packed_neon.cpp create mode 100644 tests/dispatch/fastscan_no_simd_test.cpp create mode 100644 tests/python/test_examples.py create mode 100644 tests/python/test_portable_indexes.py create mode 100644 tests/python/test_thread_counts.py create mode 100644 tests/python/test_wheel_macos.py create mode 100644 tests/unit/rabitqlib/utils/neon_hnsw_test.cpp create mode 100644 tests/unit/rabitqlib/utils/portable_kernel_test.cpp create mode 100644 tests/unit/rabitqlib/utils/query_preparation_test.cpp create mode 100644 tests/unit/rabitqlib/utils/thread_count_test.cpp diff --git a/.github/scripts/ci_changes.py b/.github/scripts/ci_changes.py index fcb3b34b..496d7b5b 100644 --- a/.github/scripts/ci_changes.py +++ b/.github/scripts/ci_changes.py @@ -44,9 +44,9 @@ def classify(paths): or path == "tests/CMakeLists.txt" ): selected.add("cpp") - elif path.startswith("tests/python/"): + elif path.startswith(("tests/python/", "sample/python/")): selected.update(("python_quality", "python_build", "wheels")) - elif path.startswith(("python/", "sample/python/", ".github/scripts/")): + elif path.startswith(("python/", ".github/scripts/")): selected.add("python_quality") elif path == ".github/workflows/test.yaml": selected.update(("cpp", "shell")) diff --git a/.github/scripts/test_ci_changes.py b/.github/scripts/test_ci_changes.py index f9f4a702..48abb608 100644 --- a/.github/scripts/test_ci_changes.py +++ b/.github/scripts/test_ci_changes.py @@ -45,7 +45,11 @@ def test_language_specific_changes(self): "python_build", "wheels", }, - "sample/python/ivf_rabitq_indexing.py": {"python_quality"}, + "sample/python/ivf_rabitq_indexing.py": { + "python_quality", + "python_build", + "wheels", + }, ".github/scripts/prepare_release.py": {"python_quality"}, "scripts/check-python.sh": {"shell", "python_quality"}, "scripts/check-tidy.sh": {"shell", "cpp"}, @@ -68,7 +72,7 @@ def test_mixed_changes_and_renamed_source(self): self.assertIn("cpp", classify(["docs/new.md", "src/deleted.cpp"])) self.assertEqual( classify(["sample/python/tool.py", "tests/unit/example.cpp"]), - {"cpp", "python_quality"}, + {"cpp", "python_quality", "python_build", "wheels"}, ) def test_unknown_inputs_and_filter_edits_run_everything(self): diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index cb02e8b5..ffd14be9 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -130,9 +130,38 @@ jobs: path: wheelhouse/*.whl if-no-files-found: error + macos-wheels: + name: macOS arm64 wheels (CPython 3.11–3.14) + needs: [prepare, changes] + if: >- + !cancelled() && needs.prepare.result == 'success' && + needs.prepare.outputs.build == 'true' && + (needs.changes.result != 'success' || needs.changes.outputs.wheels == 'true') + runs-on: macos-14 + timeout-minutes: 90 + steps: + - uses: actions/checkout@v7 + with: + ref: ${{ needs.prepare.outputs.sha }} + - uses: actions/setup-python@v7 + with: + python-version: "3.14" + - name: Verify native Apple Silicon runner + run: test "$(uname -m)" = arm64 + # All four ABIs are built, repaired, installed and tested on PRs as well. + - name: Build and test wheels with bundled OpenMP + uses: pypa/cibuildwheel@v4.2.0 + with: + output-dir: wheelhouse + - uses: actions/upload-artifact@v7 + with: + name: macos-arm64-wheels + path: wheelhouse/*.whl + if-no-files-found: error + release-tag: name: Create version tag - needs: [prepare, linux-wheels, windows-wheels] + needs: [prepare, linux-wheels, windows-wheels, macos-wheels] if: >- (github.event_name == 'push' || github.event_name == 'workflow_run') && needs.prepare.outputs.publish == 'true' @@ -161,7 +190,7 @@ jobs: if: >- (github.event_name == 'push' || github.event_name == 'workflow_run') && needs.prepare.outputs.publish == 'true' - needs: [prepare, linux-wheels, windows-wheels, release-tag] + needs: [prepare, linux-wheels, windows-wheels, macos-wheels, release-tag] runs-on: ubuntu-latest environment: name: pypi @@ -188,7 +217,7 @@ jobs: if: >- (github.event_name == 'push' || github.event_name == 'workflow_run') && needs.prepare.outputs.publish == 'true' - needs: [prepare, linux-wheels, windows-wheels, pypi-publish] + needs: [prepare, linux-wheels, windows-wheels, macos-wheels, pypi-publish] runs-on: ubuntu-latest permissions: contents: write diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 3d54f39e..903501ef 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -195,3 +195,36 @@ jobs: cmake --build build-consumer - name: Test downstream consumer run: ctest --test-dir build-consumer --output-on-failure + + macos-arm64: + name: Native Apple Silicon C++ + needs: changes + if: ${{ !cancelled() && (needs.changes.result != 'success' || needs.changes.outputs.cpp == 'true') }} + runs-on: macos-14 + timeout-minutes: 45 + env: + OMP_NUM_THREADS: "2" + steps: + - uses: actions/checkout@v7 + - name: Install build tools and existing OpenMP dependency + run: brew install cmake ninja libomp + - name: Configure portable ARM64 source groups + run: | + test "$(uname -m)" = arm64 + cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_TESTS=ON -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DOpenMP_ROOT="$(brew --prefix libomp)" + ! grep -E -- '-mavx|_avx[25].*\.cpp' build/compile_commands.json + - name: Build and run C++ reference and index tests + run: | + cmake --build build --parallel 3 + ctest --test-dir build --output-on-failure + - name: Test installed CMake package + run: | + cmake --install build --prefix "$RUNNER_TEMP/rabitq-install" + cmake -S tests/consumer -B build-consumer \ + -DCMAKE_PREFIX_PATH="$RUNNER_TEMP/rabitq-install" \ + -DOpenMP_ROOT="$(brew --prefix libomp)" + cmake --build build-consumer --parallel 3 + ctest --test-dir build-consumer --output-on-failure diff --git a/AGENTS.md b/AGENTS.md index 3dcdd66d..c055059e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,7 +7,8 @@ where present. Keep this file current when the paths, commands, or contracts bel RaBitQ-Library is a C++17 library with Python bindings (`rabitqlib`) for compact vector quantization and approximate nearest-neighbor search. It provides one-bit and multi-bit RaBitQ -encoding and IVF, HNSW, and SymphonyQG indexes. It targets x86-64 with runtime AVX2/AVX-512 dispatch. +encoding and IVF, HNSW, and SymphonyQG indexes. It targets x86-64 with runtime AVX2/AVX-512 +dispatch and macOS ARM64 with NEON and portable scalar kernels. Linux AArch64 requires separate platform validation. RaBitQ rotates vectors and quantizes residuals relative to centroids; compact codes and correction factors estimate L2 distance or inner product. @@ -19,7 +20,7 @@ factors estimate L2 distance or inner product. | `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 | -| `src/utils/cpu_features.cpp` | Generic x86 feature detection | +| `src/utils/cpu_features.cpp` | Generic CPU feature detection | | `include/rabitqlib/utils/` | Rotation, allocation, buffers, I/O, and helpers | | `python_bindings/` | pybind11 extension and index wrappers | | `sample/cpp/`, `sample/python/` | Usage examples | @@ -110,6 +111,8 @@ 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. + 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. Keep calculations in backend source files; see the dispatch coverage table in [DEVELOPMENT.md](DEVELOPMENT.md#dispatch-conventions-and-coverage). @@ -117,6 +120,7 @@ Recommended: 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. + Keep portable reference tests enabled on ARM; guard only explicit x86 backend calls. - Native tuning can affect generic code independently of dispatch. Keep native optimization off for portable binaries and wheels; runtime dispatch alone does not make a native build portable. A portable build tested on the current CPU does not establish AVX2 backend coverage. Report diff --git a/CMakeLists.txt b/CMakeLists.txt index 5bca8aa0..6112281c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,6 +21,7 @@ option( ON ) +include(cmake/rabitqOpenMP.cmake) find_package(OpenMP REQUIRED) add_library(rabitq_compile_options INTERFACE) @@ -104,14 +105,39 @@ set(RABITQ_HNSW_AVX512_POPCNT_SOURCES src/index/hnsw_search_avx512_popcnt.cpp ) +# Select source groups from the target architecture, never the build host's ISA. +set(RABITQ_TARGET_PROCESSOR "${CMAKE_SYSTEM_PROCESSOR}") +if(APPLE AND CMAKE_OSX_ARCHITECTURES) + set(RABITQ_TARGET_PROCESSOR "${CMAKE_OSX_ARCHITECTURES}") +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_AVX512_POPCNT_SOURCES} + ${RABITQ_HNSW_AVX2_SOURCES} + ${RABITQ_HNSW_AVX512_CORE_SOURCES} + ${RABITQ_HNSW_AVX512_POPCNT_SOURCES} + ) +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 + ) +else() + message(FATAL_ERROR "Supported targets are x86-64 and ARM64 (single architecture)") +endif() + add_library(rabitq_core STATIC ${RABITQ_COMMON_SOURCES} - ${RABITQ_AVX2_SOURCES} - ${RABITQ_AVX512_SOURCES} - ${RABITQ_AVX512_POPCNT_SOURCES} - ${RABITQ_HNSW_AVX2_SOURCES} - ${RABITQ_HNSW_AVX512_CORE_SOURCES} - ${RABITQ_HNSW_AVX512_POPCNT_SOURCES} + ${RABITQ_ARCH_SOURCES} + src/simd/portable_generic.cpp + src/index/hnsw_search_generic.cpp ) target_include_directories(rabitq_core PUBLIC @@ -194,6 +220,7 @@ install(EXPORT rabitqlibTargets DESTINATION "${rabitqlib_cmake_dir}" ) install(FILES + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/rabitqOpenMP.cmake" "${CMAKE_CURRENT_BINARY_DIR}/rabitqlibConfig.cmake" "${CMAKE_CURRENT_BINARY_DIR}/rabitqlibConfigVersion.cmake" DESTINATION "${rabitqlib_cmake_dir}" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 54d0de5a..8d0387ce 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -38,9 +38,11 @@ Run any examples you add or change, including examples inside Markdown. ## Setup for code changes -Source builds need an x86-64 CPU supported by RaBitQ's AVX2 or AVX-512 dispatch, -CMake 3.20 or newer for the commands below, and a GCC- or Clang-compatible C++17 compiler with OpenMP. -The current build does not support MSVC or AArch64 (including Apple Silicon). +Source builds support x86-64 with AVX2/FMA (AVX-512 optional) 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; Linux AArch64 needs separate platform +validation. See the [platform-specific build instructions](tests/README.md#quick-start) +for Windows and [Apple Silicon](tests/README.md#macos-arm64). Documentation-only contributions do not need this hardware or compiler setup. On Ubuntu or Debian: @@ -75,13 +77,23 @@ python -m pip install ".[test]" python -m pytest tests/python/test_ivf.py::test_search_output_shape -q ``` +On Apple Silicon, pass the OpenMP prefix and disable native tuning when building +the extension in the same environment: + +```bash +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 +``` + This builds the C++ extension, so the source-build prerequisites above apply. Re-run the install command after changing package sources, bindings, or C++ code; this is not an editable install. Tests should use the newly installed package. ### C++ changes -Configure, build, and run the tests: +On Linux, configure, build, and run the tests: ```bash cmake -S . -B build -DRABITQ_BUILD_TESTS=ON \ @@ -93,8 +105,8 @@ ctest --test-dir build --output-on-failure The first configuration needs Git and network access to download GoogleTest. Increase `--parallel 2` if your machine has enough memory for more compiler jobs. -Native CPU tuning is disabled here so the build can run on other supported -x86-64 machines. See [the test guide](tests/README.md) for test layout and options. +Native CPU tuning is disabled here to keep generic code portable. On Windows +or Apple Silicon, use the matching commands in [the test guide](tests/README.md#quick-start). ## Before opening a pull request @@ -137,9 +149,9 @@ for details. | Problem | What to check | | --- | --- | -| CMake cannot find OpenMP or a compiler | Install the source-build prerequisites above. Include the compiler version and CMake error when asking for help. | -| Unsupported CPU or architecture | Code tests require a supported x86-64 CPU; disabling native tuning does not add AArch64 or generic-CPU quantized search support. Report your CPU model in the issue or PR. | -| Tests do not reflect your Python or C++ edits | Re-run `python -m pip install ".[test]"` in the active environment, then check the import paths below. Re-run the regression test that exercises your change. | +| CMake cannot find OpenMP or a compiler | Check the [platform-specific prerequisites](tests/README.md#prerequisites). On Apple Silicon, install `libomp` and pass its prefix as `OpenMP_ROOT`. Include the compiler version and CMake error when asking for help. | +| Unsupported CPU or architecture | Use x86-64 with AVX2/FMA or macOS ARM64 with NEON. Linux AArch64 still needs platform validation; disabling native tuning alone does not establish support. Report your CPU model in the issue or PR. | +| Tests do not reflect your Python or C++ edits | Re-run the platform-specific install command above in the active environment, then check the import paths below. Re-run the regression test that exercises your change. | | Formatting or analysis reports the wrong tool version | Use the pinned versions and executable overrides in the [development reference](DEVELOPMENT.md). | Check which installed package and extension Python imports: diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 1fcc1618..d2f5d368 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -30,7 +30,10 @@ Verify formatting without changing files: The scripts intentionally exclude vendored code in `include/rabitqlib/third/`. To use a nonstandard executable name, set `CLANG_FORMAT`; it must still identify -itself as clang-format 15. +itself as clang-format 15. On macOS, use Bash 4+ on PATH for the shell wrappers. +Apple SDKs newer than clang-tidy 15 may require a matching newer analyzer; set +`RUN_CLANG_TIDY` to its runner and report that version separately from the Linux +clang-tidy 15 baseline. clangd embeds its own formatter, so use clangd 15 in editors such as VS Code if format-on-save must exactly match CI. If another clangd version is required, @@ -222,7 +225,7 @@ This matches CI's script set. Also pass any changed shell scripts outside - Preserve existing public headers, aliases, and index formats unless a change is explicitly documented as breaking. - Add backend-independent tests when introducing or changing SIMD kernels. -- Keep scalar, AVX2, and AVX-512 implementations behaviorally equivalent. +- Keep scalar, NEON, AVX2, and AVX-512 implementations behaviorally equivalent. - Benchmark allocations or algorithm changes in search and quantization hot paths, and substantiate performance claims with benchmarks. Run correctness tests first and include commands, dataset, CPU, compiler, ISA, thread count, @@ -258,11 +261,18 @@ 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,avx2,avx512}.cpp`. HNSW keeps its existing + implementations in `src/simd/*_{generic,neon,avx2,avx512}.cpp`. HNSW keeps its existing `src/index/` implementations and compatibility namespaces. - Select a cached function pointer through `resolve_kernel` in `src/simd/dispatch.cpp`. - The order is AVX-512, AVX2/FMA, then the existing generic implementation or a descriptive - unsupported-operation exception. Public wrappers do not repeat feature checks. + 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. + 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. - Preserve stricter predicates: population-count kernels need AVX512_VPOPCNTDQ, and HNSW's AVX-512 core variant also needs AVX2 for its warmup implementation. Do not infer support from a backend name or from `__AVX*__` macros in a public header. @@ -294,10 +304,10 @@ covers its arithmetic kernels, not every scalar loop in construction and search. | Area | Dispatch coverage / deliberate boundary | | --- | --- | -| Raw float distances, norms, packed-code products | Runtime AVX2/AVX-512 selection; generic raw-float fallback | +| Raw float distances, norms, packed-code products | Runtime AVX2/AVX-512/NEON selection; portable generic fallback | | Quantizer rescale search | SIMD bounded search; the certified scalar event sweep remains the fallback | | Integer scalar quantization, extra-code packing, transpose, sign masks | Existing runtime selection; byte layouts and rounding rules are unchanged | -| Standard and high-accuracy FastScan | Runtime selection; generic LUT construction stays separate from selection | +| Standard and high-accuracy FastScan | Runtime selection, including NEON float LUT construction; HACC LUT byte conversion remains generic on ARM | | IVF and float SymphonyQG batch correction | Complete estimator runs in the selected backend; non-float template paths remain generic | | FHT/Kac rotation | Complete rotation and scaling run in selected ISA translation units; shared AVX intrinsics preserve the FFHT butterfly order | | Float matrix rotation and PiPNN construction | Matrix products, row norms, and lower-triangle pairwise distances use isolated matrix backends | @@ -306,9 +316,11 @@ covers its arithmetic kernels, not every scalar loop in construction and search. | Graph scheduling, candidate queues, I/O, allocation, random initialization | Generic control code; IVF one-bit candidate insertion stays in a small compiled function to avoid inlining-induced register spills; thread scheduling and seeds remain caller-owned | | Example KMeans training | Uses external FAISS, whose build and dispatch are independent of this package | -Portable wheels continue to disable `RABITQ_ENABLE_NATIVE_OPTIMIZATION`. Adding an optimized -backend does not add support for generic-CPU quantized search or AArch64; those require -complete implementations and separate compatibility validation. +Portable wheels disable `RABITQ_ENABLE_NATIVE_OPTIMIZATION`. ARM64 builds exclude all +x86 source groups and use standard AArch64 NEON intrinsics, without Apple-only APIs. +macOS ARM64 has native C++ and installed-wheel CI. Linux ARM64 still requires separate +platform verification. The reference suites also exercise portable scalar kernels on x86, including +byte-for-byte packing parity with the selected x86 backend. ### Change quantization or packing diff --git a/README.md b/README.md index 3b7e2a4c..a8025597 100644 --- a/README.md +++ b/README.md @@ -54,8 +54,10 @@ python -m pip install --upgrade rabitqlib ``` -Wheels: CPython 3.11–3.14 on Linux and Windows x86-64. Requires AVX2 and FMA, -with optional AVX-512 acceleration. +Wheels: CPython 3.11–3.14 on Linux and Windows x86-64, and macOS 14+ ARM64 +(Apple Silicon). x86-64 uses AVX2/FMA with optional AVX-512 acceleration; +ARM64 uses NEON and portable scalar kernels. macOS wheels bundle OpenMP; +Intel Mac and universal2 wheels are not provided. ## Python quick start @@ -89,9 +91,16 @@ print(ids.shape, distances.shape) # (5, 10) (5, 10) print(ids[0]) ``` +For all three indexes, `build` and `search` interpret `num_threads=0` as the +detected hardware thread count. Larger requests are capped at that count; +smaller positive requests are respected. Operations may use fewer workers when +there are fewer work items. If hardware detection is unavailable, one thread is +used. Omitting `num_threads` in Python still defaults to one thread. + Python bindings are also available for `HnswIndex` and `SymqgIndex`. See the -[Python examples](sample/python/) for index construction, querying, and index -persistence. +[Python examples](sample/python/README.md) for index construction, querying, and +index persistence. IVF and HNSW examples run Faiss clustering in a separate process, +then pass saved clusters to RaBitQ indexing so their OpenMP runtimes stay separate. 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. @@ -124,7 +133,7 @@ python -m pip install . | **SymphonyQG** | Fast graph search with a configurable memory/accuracy tradeoff | Uses raw vectors by default, or optional packed 4-bit/8-bit RaBitQ vectors, alongside per-neighborhood quantization data. | IVF and SymphonyQG use [FastScan](https://arxiv.org/abs/1704.07355) for batched -estimates, while HNSW uses single-code AVX2 or AVX-512 kernels. +estimates, while HNSW uses single-code kernels selected for the target architecture. In typical workloads, 4-bit, 5-bit, and 7-bit quantization can achieve roughly 90%, 95%, and 99% recall, respectively, without reranking. Actual results @@ -136,7 +145,7 @@ depend on the dataset, index configuration, and search parameters. | --- | --- | | **Compact by design** | Choose [1-bit](https://doi.org/10.1145/3654970) or [multi-bit](https://doi.org/10.1145/3725413) codes to match your memory and accuracy target. | | **Accurate estimates** | An asymptotically optimal theoretical error bound supports reliable ordering and reranking. | -| **Fast on x86-64** | Dedicated AVX2 and AVX-512 kernels are selected through runtime CPU dispatch. | +| **Native CPU backends** | Runtime AVX2/AVX-512 selection on x86-64; NEON distance, packed-code, FastScan, rotation, and HNSW search kernels on ARM64, with scalar packing and query transposition. | | **Ready for ANN search** | Use the quantizer directly or build complete IVF, HNSW, and [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) indexes. | The library supports Euclidean distance and inner product. Cosine search is @@ -224,7 +233,7 @@ and HNSW implementations, with links to the source code. - CMake 3.20 or newer - a C++17 compiler with OpenMP support -- an x86-64 CPU with AVX2 and FMA +- an x86-64 CPU with AVX2 and FMA, or an Apple Silicon Mac Clone and build the library and example programs: @@ -237,6 +246,9 @@ cmake --build build --parallel ``` For MSVC, follow the [Windows build instructions](tests/README.md#prerequisites). +For Apple Silicon source builds, see the [macOS ARM64 instructions](tests/README.md#macos-arm64). +The ARM kernel sources use portable AArch64 intrinsics, but macOS tests do not +establish Linux ARM64 support. Local GCC/Clang builds enable `-march=native` by default; set `-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF` for portable binaries, as release wheels do. See [CPU dispatch details](DEVELOPMENT.md#dispatch-conventions-and-coverage) diff --git a/cmake/rabitqOpenMP.cmake b/cmake/rabitqOpenMP.cmake new file mode 100644 index 00000000..e235c9e5 --- /dev/null +++ b/cmake/rabitqOpenMP.cmake @@ -0,0 +1,14 @@ +# AppleClang needs an external OpenMP runtime. The same hints work with +# Homebrew libomp and a project environment's llvm-openmp package. +if(APPLE AND CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang") + find_path(RABITQ_OPENMP_INCLUDE omp.h HINTS "${OpenMP_ROOT}" ENV CONDA_PREFIX + PATH_SUFFIXES include) + find_library(RABITQ_OPENMP_LIBRARY omp HINTS "${OpenMP_ROOT}" ENV CONDA_PREFIX + PATH_SUFFIXES lib) + if(RABITQ_OPENMP_INCLUDE AND RABITQ_OPENMP_LIBRARY) + set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp") + set(OpenMP_CXX_INCLUDE_DIR "${RABITQ_OPENMP_INCLUDE}") + set(OpenMP_CXX_LIB_NAMES omp) + set(OpenMP_omp_LIBRARY "${RABITQ_OPENMP_LIBRARY}") + endif() +endif() diff --git a/cmake/rabitqlibConfig.cmake.in b/cmake/rabitqlibConfig.cmake.in index 2ca01da2..8b9840ab 100644 --- a/cmake/rabitqlibConfig.cmake.in +++ b/cmake/rabitqlibConfig.cmake.in @@ -1,6 +1,7 @@ @PACKAGE_INIT@ include(CMakeFindDependencyMacro) +include("${CMAKE_CURRENT_LIST_DIR}/rabitqOpenMP.cmake") find_dependency(OpenMP REQUIRED) include("${CMAKE_CURRENT_LIST_DIR}/rabitqlibTargets.cmake") diff --git a/example_py.sh b/example_py.sh index c9645b01..661e7a2b 100755 --- a/example_py.sh +++ b/example_py.sh @@ -14,11 +14,15 @@ python ./sample/python/symqg_indexing.py --max-degree 32 --ef-construction 400 python ./sample/python/symqg_querying.py ./data/gist/symqg_32.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs # indexing and querying for RabitQ+ with ivf -python ./sample/python/ivf_rabitq_indexing.py --total-bits 5 --num-clusters 4096 ./data/gist/gist_base.fvecs ./data/gist/ivf_4096_5.index +python ./sample/python/faiss_clustering.py --num-clusters 4096 ./data/gist/gist_base.fvecs ./data/gist/clusters_4096_l2.npz + +python ./sample/python/ivf_rabitq_indexing.py --total-bits 5 --clusters ./data/gist/clusters_4096_l2.npz ./data/gist/gist_base.fvecs ./data/gist/ivf_4096_5.index python ./sample/python/ivf_rabitq_querying.py ./data/gist/ivf_4096_5.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs # indexing and querying for RabitQ+ with hnsw -python ./sample/python/hnsw_rabitq_indexing.py --total-bits 5 --num-clusters 16 --degree 16 --ef-construction 200 ./data/gist/gist_base.fvecs ./data/gist/hnsw_5.index +python ./sample/python/faiss_clustering.py --num-clusters 16 ./data/gist/gist_base.fvecs ./data/gist/clusters_16_l2.npz + +python ./sample/python/hnsw_rabitq_indexing.py --total-bits 5 --clusters ./data/gist/clusters_16_l2.npz --degree 16 --ef-construction 200 ./data/gist/gist_base.fvecs ./data/gist/hnsw_5.index python ./sample/python/hnsw_rabitq_querying.py ./data/gist/hnsw_5.index ./data/gist/gist_query.fvecs ./data/gist/gist_groundtruth.ivecs diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 7e531103..78d1253b 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -119,6 +119,7 @@ inline void pack_codes( // 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. void accumulate( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lp_table, diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index bb175619..5c778f46 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -124,6 +124,12 @@ class HierarchicalNSW { }; private: + friend maxheap> detail::search_knn_neon( + HierarchicalNSW&, const float*, size_t + ); + friend maxheap> detail::search_knn_generic( + HierarchicalNSW&, const float*, size_t + ); friend maxheap> detail::search_knn_avx2( HierarchicalNSW&, const float*, size_t ); diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 9a8bfa45..b56a67e7 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -20,6 +20,7 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/third/hnswlib/hnswlib.h" #include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" namespace rabitqlib::ivf { template @@ -27,10 +28,7 @@ inline void parallel_for(size_t start, size_t end, size_t numThreads, Function f if (start >= end) { return; } - if (numThreads == 0) { - numThreads = std::max(1, std::thread::hardware_concurrency()); - } - numThreads = std::min(numThreads, end - start); + numThreads = std::min(resolve_num_threads(numThreads), end - start); if (numThreads == 1) { for (size_t id = start; id < end; id++) { diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 183e76e1..a2514f9e 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -279,7 +279,7 @@ inline void IVF::construct( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - num_threads = std::clamp(num_threads, size_t{1}, rabitqlib::total_threads()); + num_threads = resolve_num_threads(num_threads); /* Quantize each cluster */ #pragma omp parallel for schedule(dynamic) num_threads(num_threads) for (std::ptrdiff_t index = 0; index < static_cast(num_cluster_); diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp index 897700af..7dc86bcb 100644 --- a/include/rabitqlib/index/symqg/detail/pipnn.hpp +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -300,7 +300,7 @@ inline InitialGraph build_initial_graph( ); } using namespace pipnn_impl; - const size_t threads = std::max(1, std::min(num_threads, total_threads())); + 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); diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 55485f13..42bce716 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -62,7 +62,7 @@ class QGBuilder { QGBuilder(QuantizedGraph& index, uint32_t ef_build, size_t num_threads) : qg_{index} , ef_build_{ef_build} - , num_threads_{std::max(1, std::min(num_threads, total_threads()))} + , num_threads_{resolve_num_threads(num_threads)} , num_nodes_{qg_.num_vertices()} , dim_{qg_.dimension()} , degree_bound_(qg_.degree_bound()) {} diff --git a/include/rabitqlib/simd/fastscan_dispatch.hpp b/include/rabitqlib/simd/fastscan_dispatch.hpp index 2117ae41..ce2d29e2 100644 --- a/include/rabitqlib/simd/fastscan_dispatch.hpp +++ b/include/rabitqlib/simd/fastscan_dispatch.hpp @@ -5,17 +5,35 @@ namespace rabitqlib::fastscan::simd { +// NEON uses the portable/AVX2 high-accuracy LUT byte order. +void accumulate_neon(const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim); +void accumulate_hacc_neon( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +); + +void pack_lut_neon(size_t dim, const float* query, float* lut); void pack_lut_generic(size_t dim, const float* query, float* lut); void pack_lut_avx2(size_t dim, const float* query, float* lut); void pack_lut_avx512(size_t dim, const float* query, float* lut); +// Scalar reference for backend correctness tests; never selected by runtime dispatch. +void accumulate_generic( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +); +void accumulate_unsupported( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +); void accumulate_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lp_table, int32_t* __restrict__ result, size_t dim ); +void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut); +void accumulate_hacc_generic( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +); void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, diff --git a/include/rabitqlib/simd/hnsw_dispatch.hpp b/include/rabitqlib/simd/hnsw_dispatch.hpp index 231daff3..f6674955 100644 --- a/include/rabitqlib/simd/hnsw_dispatch.hpp +++ b/include/rabitqlib/simd/hnsw_dispatch.hpp @@ -13,6 +13,13 @@ std::priority_queue> search_knn( HierarchicalNSW&, const float*, size_t ); +std::priority_queue> search_knn_neon( + HierarchicalNSW&, const float*, size_t +); + +std::priority_queue> search_knn_generic( + HierarchicalNSW&, const float*, size_t +); std::priority_queue> search_knn_avx2( HierarchicalNSW&, const float*, size_t ); diff --git a/include/rabitqlib/simd/pack_excode_dispatch.hpp b/include/rabitqlib/simd/pack_excode_dispatch.hpp index f61ed9e3..a1765fb3 100644 --- a/include/rabitqlib/simd/pack_excode_dispatch.hpp +++ b/include/rabitqlib/simd/pack_excode_dispatch.hpp @@ -17,16 +17,22 @@ void packing_6bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_2bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_2bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_3bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_3bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_4bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_4bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_5bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_6bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_7bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_7bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_2bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); diff --git a/include/rabitqlib/simd/rotator_dispatch.hpp b/include/rabitqlib/simd/rotator_dispatch.hpp index 80036277..f6bc0db7 100644 --- a/include/rabitqlib/simd/rotator_dispatch.hpp +++ b/include/rabitqlib/simd/rotator_dispatch.hpp @@ -19,10 +19,33 @@ void fht_rotate( const uint8_t* flip ); +void flip_sign_neon(const uint8_t* flip, float* data, size_t dim); +void kacs_walk_neon(float* data, size_t len); +void fht_rotate_neon( + const float* data, + float* output, + size_t dim, + size_t padded, + size_t trunc, + float factor, + const uint8_t* flip +); + +void flip_sign_generic(const uint8_t* flip, float* data, size_t dim); void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim); +void kacs_walk_generic(float* data, size_t len); void kacs_walk_avx2(float* data, size_t len); +void fht_rotate_generic( + const float* data, + float* output, + size_t dim, + size_t padded, + size_t trunc, + float factor, + const uint8_t* flip +); void fht_rotate_avx2( const float* data, float* rotated_vec, diff --git a/include/rabitqlib/simd/space_dispatch.hpp b/include/rabitqlib/simd/space_dispatch.hpp index 1403e1fb..d79f01c1 100644 --- a/include/rabitqlib/simd/space_dispatch.hpp +++ b/include/rabitqlib/simd/space_dispatch.hpp @@ -10,6 +10,11 @@ float dot_product(const float* a, const float* b, size_t dim); float dot_product_dis(const float* a, const float* b, size_t dim); float l2norm_sqr(const float* a, size_t dim); +float euclidean_sqr_neon(const float* a, const float* b, size_t dim); +float dot_product_neon(const float* a, const float* b, size_t dim); +float dot_product_dis_neon(const float* a, const float* b, size_t dim); +float l2norm_sqr_neon(const float* a, size_t dim); + float euclidean_sqr_generic(const float* a, const float* b, size_t dim); float dot_product_generic(const float* a, const float* b, size_t dim); float dot_product_dis_generic(const float* a, const float* b, size_t dim); @@ -27,27 +32,60 @@ float l2norm_sqr_avx512(const float* a, size_t dim); namespace excode_ipimpl { +float ip16_fxu1_neon(const float* query, const uint8_t* code, size_t dim); +float ip64_fxu2_neon(const float* query, const uint8_t* code, size_t dim); +float ip64_fxu3_neon(const float* query, const uint8_t* code, size_t dim); +float ip16_fxu4_neon(const float* query, const uint8_t* code, size_t dim); +float ip64_fxu5_neon(const float* query, const uint8_t* code, size_t dim); +float ip64_fxu6_neon(const float* query, const uint8_t* code, size_t dim); +float ip64_fxu7_neon(const float* query, const uint8_t* code, size_t dim); +float ip16_fxu8_neon(const float* query, const uint8_t* code, size_t dim); + +float ip16_fxu1_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip16_fxu1_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip64_fxu2_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip64_fxu2_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip64_fxu3_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip64_fxu3_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip16_fxu4_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip16_fxu4_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip64_fxu5_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip64_fxu5_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip64_fxu6_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip64_fxu6_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip64_fxu7_generic( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); float ip64_fxu7_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip16_fxu8_generic( + const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim +); float ip16_fxu8_avx2( const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim ); @@ -79,17 +117,46 @@ float ip16_fxu8_avx512( } // namespace excode_ipimpl +void new_transpose_bin_neon( + const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits +); +void new_transpose_bin_512_neon( + const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits +); +void scalar_quantize_uint8_neon( + uint8_t* result, const float* data, size_t dim, float lo, float delta +); +void scalar_quantize_uint16_neon( + uint16_t* result, const float* data, size_t dim, float lo, float delta +); + +void new_transpose_bin_generic( + const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits +); void new_transpose_bin_avx2( const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ); +void new_transpose_bin_512_generic( + const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits +); void new_transpose_bin_512_avx2( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ); +float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t dim); +float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t dim); +float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t dim); float mask_ip_x0_q_avx2(const float* query, const uint8_t* data, size_t padded_dim); +float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t dim); float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_dim); +void scalar_quantize_uint8_generic( + uint8_t* result, const float* data, size_t dim, float lo, float delta +); void scalar_quantize_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ); +void scalar_quantize_uint16_generic( + uint16_t* result, const float* data, size_t dim, float lo, float delta +); void scalar_quantize_uint16_avx2( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ); diff --git a/include/rabitqlib/simd/warmup_dispatch.hpp b/include/rabitqlib/simd/warmup_dispatch.hpp index bae9358b..b0150a88 100644 --- a/include/rabitqlib/simd/warmup_dispatch.hpp +++ b/include/rabitqlib/simd/warmup_dispatch.hpp @@ -5,6 +5,31 @@ namespace rabitqlib::simd { +float warmup_ip_x0_q_512_neon( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +); +float warmup_ip_x0_q_512_neon( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +); + +float warmup_ip_x0_q_512_generic( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +); float warmup_ip_x0_q_512_avx2( const uint8_t* data, const uint64_t* query, @@ -14,6 +39,14 @@ float warmup_ip_x0_q_512_avx2( size_t b_query ); +float warmup_ip_x0_q_512_generic( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +); float warmup_ip_x0_q_512_avx2( const uint64_t* data, const uint64_t* query, diff --git a/include/rabitqlib/utils/cpu_features.hpp b/include/rabitqlib/utils/cpu_features.hpp index 38605ace..743fa8b0 100644 --- a/include/rabitqlib/utils/cpu_features.hpp +++ b/include/rabitqlib/utils/cpu_features.hpp @@ -13,10 +13,12 @@ struct Features { bool avx512vpopcntdq = false; bool avx512vl = false; bool avx512cd = false; + bool neon = false; }; const Features& features(); bool has_avx2(); +bool has_neon(); bool has_avx512_core(); bool has_avx512_popcnt(); diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 80da7b6a..9ddfa432 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -194,7 +194,7 @@ inline std::vector compute_centroid( const auto thread_count = static_cast(std::max( 1, std::min( - {num_threads, + {resolve_num_threads(num_threads), std::max(num_points, 1), static_cast(std::numeric_limits::max())} ) @@ -240,7 +240,7 @@ inline PID exact_nn( const auto thread_count = static_cast(std::max( 1, std::min( - {num_threads, + {resolve_num_threads(num_threads), std::max(num_points, 1), static_cast(std::numeric_limits::max())} ) diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index 57fb527d..ff6e964b 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -86,6 +87,12 @@ inline size_t total_threads() { return threads == 0 ? 1 : threads; } +// Zero selects all hardware threads; explicit requests cannot exceed that count. +inline size_t resolve_num_threads(size_t requested) { + const size_t available = total_threads(); + return requested == 0 ? available : std::min(requested, available); +} + template float distance_ratio( const T* data, diff --git a/pyproject.toml b/pyproject.toml index 3f22fc70..f39f219a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,12 +40,19 @@ cmake.args = [ ] [tool.cibuildwheel] -build = "cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp314-manylinux_x86_64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64 cp314-win_amd64" +build = "cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp314-manylinux_x86_64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64 cp314-win_amd64 cp311-macosx_arm64 cp312-macosx_arm64 cp313-macosx_arm64 cp314-macosx_arm64" build-verbosity = 1 environment = { CMAKE_ARGS = "-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF" } -test-requires = "pytest" +test-requires = ["pytest", "faiss-cpu"] test-command = "pytest {project}/tests/python --tb=short -q" +[tool.cibuildwheel.macos] +archs = ["arm64"] +before-all = "brew install libomp" +environment = { CMAKE_ARGS = "-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF -DOpenMP_ROOT=$(brew --prefix libomp)", MACOSX_DEPLOYMENT_TARGET = "14.0" } +repair-wheel-command = "delocate-wheel --require-archs arm64 -w {dest_dir} -v {wheel}" +test-environment = { RABITQ_TEST_WHEEL = "1", OMP_NUM_THREADS = "2" } + [tool.cibuildwheel.windows] before-build = "python -m pip install delvewheel" repair-wheel-command = "delvewheel repair -w {dest_dir} {wheel}" diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index 749f4dc7..a1067da1 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -33,6 +33,19 @@ target_link_libraries(_rabitqlib PRIVATE Python::NumPy ) +# Keep the OpenMP runtime discoverable in source installs and during wheel +# repair. Delocate copies it into the wheel and rewrites the loader path. +if(APPLE) + set_property(TARGET _rabitqlib PROPERTY INSTALL_RPATH_USE_LINK_PATH TRUE) + # CMake omits link directories inside the source/build tree from the + # automatic install RPATH, including project-local Python environments. + foreach(runtime IN LISTS OpenMP_CXX_LIBRARIES) + get_filename_component(runtime_dir "${runtime}" DIRECTORY) + set_property(TARGET _rabitqlib APPEND PROPERTY INSTALL_RPATH "${runtime_dir}") + endforeach() + install(FILES licenses/LLVM-OpenMP.txt DESTINATION rabitqlib/licenses) +endif() + # 5. Map files to a 'rabitqlib' directory # inside the target wheel, completely bypassing your local layout. install(TARGETS _rabitqlib DESTINATION rabitqlib) diff --git a/python_bindings/licenses/LLVM-OpenMP.txt b/python_bindings/licenses/LLVM-OpenMP.txt new file mode 100644 index 00000000..99075663 --- /dev/null +++ b/python_bindings/licenses/LLVM-OpenMP.txt @@ -0,0 +1,361 @@ +============================================================================== +The LLVM Project is under the Apache License v2.0 with LLVM Exceptions: +============================================================================== + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +---- LLVM Exceptions to the Apache 2.0 License ---- + +As an exception, if, as a result of your compiling your source code, portions +of this Software are embedded into an Object form of such source code, you +may redistribute such embedded portions in such Object form without complying +with the conditions of Sections 4(a), 4(b) and 4(d) of the License. + +In addition, if you combine or link compiled forms of this Software with +software that is licensed under the GPLv2 ("Combined Software") and if a +court of competent jurisdiction determines that the patent provision (Section +3), the indemnity provision (Section 9) or other Section of the License +conflicts with the conditions of the GPLv2, you may retroactively and +prospectively choose to deem waived or otherwise exclude such Section(s) of +the License, but only in their entirety and only with respect to the Combined +Software. + +============================================================================== +Software from third parties included in the LLVM Project: +============================================================================== +The LLVM Project contains third party software which is under different license +terms. All such code will be identified clearly using at least one of two +mechanisms: +1) It will be in a separate directory tree with its own `LICENSE.txt` or + `LICENSE` file at the top containing the specific license and restrictions + which apply to that software, or +2) It will contain specific license and restriction terms at the top of every + file. + +============================================================================== +Legacy LLVM License (https://llvm.org/docs/DeveloperPolicy.html#legacy): +============================================================================== + +The software contained in this directory tree is dual licensed under both the +University of Illinois "BSD-Like" license and the MIT license. As a user of +this code you may choose to use it under either license. As a contributor, +you agree to allow your code to be used under both. The full text of the +relevant licenses is included below. + +In addition, a license agreement from the copyright/patent holders of the +software contained in this directory tree is included below. + +============================================================================== + +University of Illinois/NCSA +Open Source License + +Copyright (c) 1997-2019 Intel Corporation + +All rights reserved. + +Developed by: + OpenMP Runtime Team + Intel Corporation + http://www.openmprtl.org + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal with +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimers. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimers in the + documentation and/or other materials provided with the distribution. + + * Neither the names of Intel Corporation OpenMP Runtime Team nor the + names of its contributors may be used to endorse or promote products + derived from this Software without specific prior written permission. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE +SOFTWARE. + +============================================================================== + +Copyright (c) 1997-2019 Intel Corporation + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. + +============================================================================== + +Intel Corporation + +Software Grant License Agreement ("Agreement") + +Except for the license granted herein to you, Intel Corporation ("Intel") reserves +all right, title, and interest in and to the Software (defined below). + +Definition + +"Software" means the code and documentation as well as any original work of +authorship, including any modifications or additions to an existing work, that +is intentionally submitted by Intel to llvm.org (http://llvm.org) ("LLVM") for +inclusion in, or documentation of, any of the products owned or managed by LLVM +(the "Work"). For the purposes of this definition, "submitted" means any form of +electronic, verbal, or written communication sent to LLVM or its +representatives, including but not limited to communication on electronic +mailing lists, source code control systems, and issue tracking systems that are +managed by, or on behalf of, LLVM for the purpose of discussing and improving +the Work, but excluding communication that is conspicuously marked otherwise. + +1. Grant of Copyright License. Subject to the terms and conditions of this + Agreement, Intel hereby grants to you and to recipients of the Software + distributed by LLVM a perpetual, worldwide, non-exclusive, no-charge, + royalty-free, irrevocable copyright license to reproduce, prepare derivative + works of, publicly display, publicly perform, sublicense, and distribute the + Software and such derivative works. + +2. Grant of Patent License. Subject to the terms and conditions of this + Agreement, Intel hereby grants you and to recipients of the Software + distributed by LLVM a perpetual, worldwide, non-exclusive, no-charge, + royalty-free, irrevocable (except as stated in this section) patent license + to make, have made, use, offer to sell, sell, import, and otherwise transfer + the Work, where such license applies only to those patent claims licensable + by Intel that are necessarily infringed by Intel's Software alone or by + combination of the Software with the Work to which such Software was + submitted. If any entity institutes patent litigation against Intel or any + other entity (including a cross-claim or counterclaim in a lawsuit) alleging + that Intel's Software, or the Work to which Intel has contributed constitutes + direct or contributory patent infringement, then any patent licenses granted + to that entity under this Agreement for the Software or Work shall terminate + as of the date such litigation is filed. + +Unless required by applicable law or agreed to in writing, the software is +provided on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +either express or implied, including, without limitation, any warranties or +conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A +PARTICULAR PURPOSE. + +============================================================================== diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index edd66e5c..97cceec2 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -13,7 +13,6 @@ #include #include #include -#include #include #include @@ -23,6 +22,7 @@ #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/utils/rotator.hpp" +#include "rabitqlib/utils/tools.hpp" namespace py = pybind11; @@ -120,12 +120,10 @@ class SymqgIndex { auto* dists_data = dists.mutable_data(); const auto* queries_data = query_array.data(); - const size_t requested_threads = - num_threads == 0 ? std::thread::hardware_concurrency() : num_threads; const auto workers = static_cast(std::max( 1, std::min( - {requested_threads, + {resolve_num_threads(num_threads), nq, static_cast(std::numeric_limits::max())} ) diff --git a/sample/python/README.md b/sample/python/README.md new file mode 100644 index 00000000..3f00b37f --- /dev/null +++ b/sample/python/README.md @@ -0,0 +1,53 @@ +# Python examples + +Run these commands from the repository root in your project Python environment. +Install `rabitqlib` and NumPy; the clustering step additionally needs `faiss-cpu`. + +## IVF: cluster, then index + +Faiss clustering and RaBitQ indexing run in separate Python processes. This keeps +their OpenMP runtimes separate, including when both packages come from pip wheels. +The intermediate `.npz` contains float32 `centroids`, uint32 `cluster_ids`, and the +`metric` (`l2` or `ip`). It can be reused for different RaBitQ bit widths. + +```sh +python sample/python/faiss_clustering.py --num-clusters 4096 \ + data/gist/gist_base.fvecs data/gist/clusters_4096_l2.npz + +python sample/python/ivf_rabitq_indexing.py --total-bits 5 \ + --clusters data/gist/clusters_4096_l2.npz \ + data/gist/gist_base.fvecs data/gist/ivf_4096_5.index + +python sample/python/ivf_rabitq_querying.py \ + data/gist/ivf_4096_5.index data/gist/gist_query.fvecs data/gist/gist_groundtruth.ivecs +``` + +The indexing command now takes `--clusters` instead of `--num-clusters`; its cluster +count comes from the saved centroids. Always use the same base vectors in the same +row order for both stages. The loader checks the metric, dimensions, point count, +array types, and cluster ID range; it cannot detect reordered or replaced vectors. + +For inner product, pass `--metric ip` to **both** clustering and indexing, and use +ground truth for that metric when evaluating recall. `--num-threads 0` selects the +hardware thread count; larger requests are capped. Each stage has its own setting. + +## HNSW + +HNSW uses the same saved-cluster format. Choose its cluster count in the Faiss stage: + +```sh +python sample/python/faiss_clustering.py --num-clusters 16 \ + data/gist/gist_base.fvecs data/gist/clusters_16_l2.npz + +python sample/python/hnsw_rabitq_indexing.py --total-bits 5 --degree 16 \ + --ef-construction 200 --clusters data/gist/clusters_16_l2.npz \ + data/gist/gist_base.fvecs data/gist/hnsw_5.index +``` + +HNSW querying is unchanged. SymphonyQG does not need Faiss clustering; run +`symqg_indexing.py` and `symqg_querying.py` directly. All querying scripts and the +shared file/recall utilities are independent of Faiss. + +Run the stages as separate commands, not by importing both libraries into one +Python process or notebook kernel. This separation fixes the example workflow; +it does not change native-runtime coexistence for applications importing both. diff --git a/sample/python/faiss_clustering.py b/sample/python/faiss_clustering.py new file mode 100644 index 00000000..98d3cfdb --- /dev/null +++ b/sample/python/faiss_clustering.py @@ -0,0 +1,66 @@ +"""Run Faiss in its own process and save clusters for the RaBitQ examples.""" + +import argparse +import os +from time import time + +import faiss +import numpy as np + +from utils import read_fvecs + + +def cluster_data( + data: np.ndarray, num_clusters: int, metric: str, num_threads: int +) -> tuple[np.ndarray, np.ndarray]: + if not 1 <= num_clusters <= len(data): + raise ValueError("num_clusters must be between 1 and the number of data points") + if num_threads < 0: + raise ValueError("num_threads must be non-negative") + if metric not in ("l2", "ip"): + raise ValueError("metric must be l2 or ip") + hardware_threads = os.cpu_count() or 1 + threads = ( + hardware_threads if num_threads == 0 else min(num_threads, hardware_threads) + ) + faiss.omp_set_num_threads(threads) + print(f"Clustering metric: {metric.upper()}, threads: {threads}") + faiss_metric = faiss.METRIC_L2 if metric == "l2" else faiss.METRIC_INNER_PRODUCT + index = faiss.index_factory(data.shape[1], f"IVF{num_clusters},Flat", faiss_metric) + index.verbose = True + start = time() + index.train(data) + print(f"IVF training time: {time() - start:.2f}s") + centroids = index.quantizer.reconstruct_n(0, index.nlist) + _, labels = index.quantizer.search(data, 1) + return centroids, labels.ravel().astype(np.uint32) + + +def main(args) -> None: + data = read_fvecs(args.data_file) + centroids, cluster_ids = cluster_data( + data, args.num_clusters, args.metric, args.num_threads + ) + # Open the exact requested path; np.savez otherwise appends .npz automatically. + with open(args.clusters_file, "wb") as output: + np.savez( + output, centroids=centroids, cluster_ids=cluster_ids, metric=args.metric + ) + print(f"Clusters saved → {args.clusters_file}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("data_file", help="Input vectors in fvecs format") + parser.add_argument( + "clusters_file", help="Output clustering file (NumPy npz format)" + ) + parser.add_argument("--num-clusters", type=int, default=256) + parser.add_argument("--metric", choices=["l2", "ip"], default="l2") + parser.add_argument( + "--num-threads", + type=int, + default=0, + help="Threads for clustering (0: hardware count; larger requests are capped)", + ) + main(parser.parse_args()) diff --git a/sample/python/hnsw_rabitq_indexing.py b/sample/python/hnsw_rabitq_indexing.py index fef9f4da..63552d05 100644 --- a/sample/python/hnsw_rabitq_indexing.py +++ b/sample/python/hnsw_rabitq_indexing.py @@ -3,12 +3,11 @@ from rabitqlib import HnswIndex -from utils import cluster_data, read_fvecs +from utils import load_clusters, read_fvecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -NUM_CLUSTERS = 256 # number of clusters M = 16 # degree bound for HNSW EF_CONSTRUCTION = 200 # ef for indexing TOTAL_BITS = 8 # total number of bits for quantization @@ -23,10 +22,8 @@ def main(args=None) -> None: data = read_fvecs(args.data_file) print(f"Data shape: {data.shape}") - # 2. Cluster with FAISS - centroids, cluster_ids = cluster_data( - data, args.num_clusters, args.metric, args.num_threads - ) + # 2. Load clusters computed in a separate Faiss process. + centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") # 3. Build HNSW index @@ -65,12 +62,9 @@ def main(args=None) -> None: parser.add_argument("data_file", type=str, help="Path to the data file") parser.add_argument("index_file", type=str, help="Path to save the index") parser.add_argument( - "--num-clusters", - dest="num_clusters", - type=int, - metavar="INT", - default=256, - help="Number of clusters for quantization", + "--clusters", + required=True, + help="Clustering file produced by faiss_clustering.py for this data and metric", ) parser.add_argument( "--degree", diff --git a/sample/python/ivf_rabitq_indexing.py b/sample/python/ivf_rabitq_indexing.py index 4ba54705..3abf2c1a 100644 --- a/sample/python/ivf_rabitq_indexing.py +++ b/sample/python/ivf_rabitq_indexing.py @@ -3,12 +3,11 @@ from rabitqlib import IvfIndex -from utils import cluster_data, read_fvecs +from utils import load_clusters, read_fvecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -NUM_CLUSTERS = 256 # number of clusters (K for IVF) TOTAL_BITS = 8 # total number of bits for quantization METRIC = "l2" # "l2" or "ip" FASTER_QUANT = True # use faster quantization @@ -24,10 +23,8 @@ def main(args=None) -> None: print(f"\tN: {n}") print(f"\tDIM: {dim}") - # 2. Cluster with FAISS - centroids, cluster_ids = cluster_data( - data, args.num_clusters, args.metric, args.num_threads - ) + # 2. Load clusters computed in a separate Faiss process. + centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") # 3. Build IVF index @@ -39,7 +36,7 @@ def main(args=None) -> None: idx = IvfIndex( dim=dim, max_elements=n, - num_clusters=args.num_clusters, + num_clusters=len(centroids), nbits=args.total_bits, metric=args.metric, ) @@ -66,12 +63,9 @@ def main(args=None) -> None: parser.add_argument("data_file", type=str, help="Path to the data file") parser.add_argument("index_file", type=str, help="Path to save the index") parser.add_argument( - "--num-clusters", - dest="num_clusters", - type=int, - metavar="INT", - default=NUM_CLUSTERS, - help="Number of clusters (K for IVF)", + "--clusters", + required=True, + help="Clustering file produced by faiss_clustering.py for this data and metric", ) parser.add_argument( "--total-bits", diff --git a/sample/python/utils.py b/sample/python/utils.py index c5110c20..30587a96 100644 --- a/sample/python/utils.py +++ b/sample/python/utils.py @@ -1,6 +1,3 @@ -from time import time - -import faiss import numpy as np # ────────────────────────────────────────────── @@ -38,45 +35,38 @@ def compute_recall(ids: np.ndarray, gt: np.ndarray, topk: int) -> float: # ────────────────────────────────────────────── -# Clustering +# Saved clustering results (no Faiss dependency) # ────────────────────────────────────────────── -def cluster_data( - X: np.ndarray, - K: int, - metric_str: str = "l2", - num_threads: int = 0, +def load_clusters( + filename: str, data_shape: tuple[int, int], metric: str ) -> tuple[np.ndarray, np.ndarray]: - """ - Cluster X into K clusters using FAISS IVF. - - Returns - ------- - centroids : np.ndarray, shape (K, dim), float32 - cluster_ids : np.ndarray, shape (n,), uint32 - """ - if num_threads > 0: - faiss.omp_set_num_threads(num_threads) - - dim = X.shape[1] - - if metric_str == "ip": - metric = faiss.METRIC_INNER_PRODUCT - print("Clustering metric: InnerProduct") - else: - metric = faiss.METRIC_L2 - print("Clustering metric: L2") - - index = faiss.index_factory(dim, f"IVF{K},Flat", metric) - index.verbose = True - - t0 = time() - index.train(X) - print(f"IVF training time: {time() - t0:.2f}s") - - centroids = index.quantizer.reconstruct_n(0, index.nlist) # (K, dim) float32 - _, cluster_ids_2d = index.quantizer.search(X, 1) # (n, 1) int64 - cluster_ids = cluster_ids_2d.flatten().astype(np.uint32) # (n,) uint32 - + """Load results from faiss_clustering.py for the same data in the same order.""" + with np.load(filename, allow_pickle=False) as saved: + if not {"centroids", "cluster_ids", "metric"}.issubset(saved.files): + raise ValueError( + "clustering file must contain centroids, cluster_ids and metric" + ) + centroids = saved["centroids"] + cluster_ids = saved["cluster_ids"] + saved_metric = saved["metric"] + if saved_metric.shape != () or saved_metric.item() != metric: + raise ValueError("clustering metric does not match index metric") + num_points, dim = data_shape + if ( + centroids.dtype != np.float32 + or centroids.ndim != 2 + or centroids.shape[1] != dim + or not 1 <= len(centroids) <= num_points + ): + raise ValueError( + "centroids must be float32 with shape (num_clusters, data_dim)" + ) + if not np.isfinite(centroids).all(): + raise ValueError("centroids must contain only finite values") + if cluster_ids.dtype != np.uint32 or cluster_ids.shape != (num_points,): + raise ValueError("cluster_ids must be uint32 with one ID per data point") + if (cluster_ids >= len(centroids)).any(): + raise ValueError("cluster_ids contain an out-of-range cluster ID") return centroids, cluster_ids diff --git a/src/index/hnsw_search_generic.cpp b/src/index/hnsw_search_generic.cpp new file mode 100644 index 00000000..5afc59c0 --- /dev/null +++ b/src/index/hnsw_search_generic.cpp @@ -0,0 +1,39 @@ +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/simd/warmup_dispatch.hpp" + +namespace rabitqlib::hnsw::detail { + +struct HnswGenericKernel { + static inline float warmup_ip_x0_q_512( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t padded_dim, + size_t b_query + ) { + return simd::warmup_ip_x0_q_512_generic( + data, query, delta, vl, padded_dim, b_query + ); + } + + static inline float mask_ip_x0_q( + const float* query, const uint64_t* data, size_t padded_dim + ) { + return simd::mask_ip_x0_q_generic(query, data, padded_dim); + } +}; + +maxheap> search_knn_generic( + HierarchicalNSW& index, const float* rotated_query, size_t topk +) { + return index.search_knn_direct(rotated_query, topk); +} + +} // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_neon.cpp b/src/index/hnsw_search_neon.cpp new file mode 100644 index 00000000..ad0ace15 --- /dev/null +++ b/src/index/hnsw_search_neon.cpp @@ -0,0 +1,42 @@ +#include +#include +#include + +#include "../simd/hnsw_neon_kernels.hpp" +#include "rabitqlib/defines.hpp" +#include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/simd/warmup_dispatch.hpp" + +namespace rabitqlib::hnsw::detail { + +struct HnswNeonKernel { + static inline float warmup_ip_x0_q_512( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t padded_dim, + size_t b_query + ) { + return simd::detail::warmup_ip_neon( + reinterpret_cast(data), query, delta, vl, padded_dim, b_query + ); + } + + static inline float mask_ip_x0_q( + const float* query, const uint64_t* data, size_t padded_dim + ) { + return simd::detail::mask_ip_neon( + query, reinterpret_cast(data), padded_dim + ); + } +}; + +maxheap> search_knn_neon( + HierarchicalNSW& index, const float* rotated_query, size_t topk +) { + return index.search_knn_direct(rotated_query, topk); +} + +} // namespace rabitqlib::hnsw::detail diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 8927f2db..25cf9fd4 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -51,10 +51,25 @@ Function resolve_kernel(Function avx2, Function fallback) { return fallback; } +// Keep ARM selection cached in the same resolver as the x86 tiers. +template +Function resolve_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 kMatrixProductFn = - resolve_kernel(matrix_product_avx512, matrix_product_avx2, matrix_product_generic); + RABITQ_RESOLVE(matrix_product_avx512, matrix_product_avx2, matrix_product_generic); void matrix_product( const float* left, @@ -67,7 +82,7 @@ void matrix_product( kMatrixProductFn(left, right, result, rows, inner, cols); } -const auto kMatrixProductTransposedFn = resolve_kernel( +const auto kMatrixProductTransposedFn = RABITQ_RESOLVE( matrix_product_transposed_avx512, matrix_product_transposed_avx2, matrix_product_transposed_generic @@ -85,13 +100,13 @@ void matrix_product_transposed( } const auto kRowNormsFn = - resolve_kernel(row_norms_avx512, row_norms_avx2, row_norms_generic); + 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) { kRowNormsFn(data, result, rows, dim); } -const auto kPairwiseDistancesLowerFn = resolve_kernel( +const auto kPairwiseDistancesLowerFn = RABITQ_RESOLVE( pairwise_distances_lower_avx512, pairwise_distances_lower_avx2, pairwise_distances_lower_generic @@ -108,7 +123,7 @@ void pairwise_distances_lower( kPairwiseDistancesLowerFn(data, result, norms_data, size, dim, inner_product); } -const auto kQgBatchEstdistFn = resolve_kernel( +const auto kQgBatchEstdistFn = RABITQ_RESOLVE( qg_batch_estdist_avx512, qg_batch_estdist_avx2, qg_batch_estdist_generic ); @@ -122,7 +137,7 @@ void qg_batch_estdist( } const auto kQgBatchEstdistMaskFn = - resolve_kernel(qg_batch_estdist_mask_avx2, qg_batch_estdist_mask_generic); + RABITQ_RESOLVE2(qg_batch_estdist_mask_avx2, qg_batch_estdist_mask_generic); uint32_t qg_batch_estdist_mask( const char* batch_data, @@ -134,7 +149,7 @@ uint32_t qg_batch_estdist_mask( return kQgBatchEstdistMaskFn(batch_data, q_obj, padded_dim, est_distance, threshold); } -const auto kSplitBatchEstdistFn = resolve_kernel( +const auto kSplitBatchEstdistFn = RABITQ_RESOLVE( split_batch_estdist_avx512, split_batch_estdist_avx2, split_batch_estdist_generic ); @@ -152,14 +167,26 @@ void split_batch_estdist( ); } +#if defined(__aarch64__) +const auto kEuclideanSqrFn = + resolve_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); +const auto kDotProductFn = + resolve_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); +const auto kDotProductDisFn = + resolve_kernel(dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()); +const auto kL2normSqrFn = + resolve_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); +#else const auto kEuclideanSqrFn = - resolve_kernel(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); + RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); const auto kDotProductFn = - resolve_kernel(dot_product_avx512, dot_product_avx2, dot_product_generic); + RABITQ_RESOLVE(dot_product_avx512, dot_product_avx2, dot_product_generic); const auto kDotProductDisFn = - resolve_kernel(dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic); + RABITQ_RESOLVE(dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic); const auto kL2normSqrFn = - resolve_kernel(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic); + RABITQ_RESOLVE(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic); + +#endif float euclidean_sqr(const float* a, const float* b, size_t dim) { return kEuclideanSqrFn(a, b, dim); @@ -175,12 +202,6 @@ float dot_product_dis(const float* a, const float* b, size_t dim) { float l2norm_sqr(const float* a, size_t dim) { return kL2normSqrFn(a, dim); } -[[noreturn]] static void missing_feature(const char* feature_name) { - throw std::runtime_error( - std::string(feature_name) + " requires AVX2/FMA or AVX512 support" - ); -} - // 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 // a duplicate of the 1-bit implementation. @@ -190,7 +211,7 @@ static float ip_fxu0( return 0.0F; } -const auto kBestRescaleFactorFn = resolve_kernel( +const auto kBestRescaleFactorFn = RABITQ_RESOLVE( best_rescale_factor_avx512, best_rescale_factor_avx2, best_rescale_factor_generic ); @@ -200,13 +221,13 @@ double best_rescale_factor( return kBestRescaleFactorFn(magnitudes, dim, max_code, start, end); } -static void -missing_fht_rotate(const float*, float*, size_t, size_t, size_t, float, const uint8_t*) { - missing_feature("sign flip"); -} - +#if defined(__aarch64__) || defined(_M_ARM64) const auto kFhtRotateFn = - resolve_kernel(fht_rotate_avx512, fht_rotate_avx2, missing_fht_rotate); + resolve_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 void fht_rotate( const float* data, @@ -220,65 +241,36 @@ void fht_rotate( kFhtRotateFn(data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip); } -static float missing_excode_ip(const float*, const uint8_t*, size_t) { - missing_feature("excode ip functions"); -} - -static void missing_flip_sign(const uint8_t*, float*, size_t) { - missing_feature("sign flip"); -} - -static void missing_kacs_walk(float*, size_t) { missing_feature("FhtKacRotator"); } - -static void missing_scalar_quantize_uint8(uint8_t*, const float*, size_t, float, float) { - missing_feature("uint8 quantize"); -} - -static void missing_scalar_quantize_uint16(uint16_t*, const float*, size_t, float, float) { - missing_feature("uint16 quantize"); -} - -static void missing_pack_excode(const uint8_t*, uint8_t*, size_t) { - missing_feature("excode packing"); -} - -static void missing_new_transpose_bin(const uint16_t*, uint64_t*, size_t, size_t) { - missing_feature("new transpose bin"); -} - -static void missing_new_transpose_bin_512(const uint8_t*, uint64_t*, size_t, size_t) { - missing_feature("new_transpose_bin_512"); -} - -static float missing_mask_ip_x0_q(const float*, const uint8_t*, size_t) { - missing_feature("mask ip x0 q"); -} - -static void missing_fastscan_accumulate( - const uint8_t* __restrict__, const uint8_t* __restrict__, int32_t* __restrict__, size_t -) { - missing_feature("fastscan accumulate"); -} - -static void missing_fastscan_transfer_lut_hacc(const uint16_t*, size_t, uint8_t*) { - missing_feature("fastscan high-accuracy LUT transfer"); -} - -static void missing_fastscan_accumulate_hacc( - const uint8_t* __restrict__, const uint8_t* __restrict__, int32_t*, size_t -) { - missing_feature("fastscan high-accuracy accumulate"); -} - -static float missing_warmup_ip_x0_q_512( - const uint8_t*, const uint64_t*, float, float, size_t, size_t -) { - missing_feature("warmup_ip_x0_q_512"); -} - ExcodeIpTable resolve_excode_ip_table() { + const auto generic = ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_generic, + excode_ipimpl::ip64_fxu2_generic, + excode_ipimpl::ip64_fxu3_generic, + excode_ipimpl::ip16_fxu4_generic, + excode_ipimpl::ip64_fxu5_generic, + excode_ipimpl::ip64_fxu6_generic, + excode_ipimpl::ip64_fxu7_generic, + excode_ipimpl::ip16_fxu8_generic, + }; +#if defined(__aarch64__) || defined(_M_ARM64) return resolve_kernel( ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_neon, + excode_ipimpl::ip64_fxu2_neon, + excode_ipimpl::ip64_fxu3_neon, + excode_ipimpl::ip16_fxu4_neon, + excode_ipimpl::ip64_fxu5_neon, + excode_ipimpl::ip64_fxu6_neon, + excode_ipimpl::ip64_fxu7_neon, + excode_ipimpl::ip16_fxu8_neon}, + generic, + cpu::has_neon() + ); +#else + return RABITQ_RESOLVE( + (ExcodeIpTable{ ip_fxu0, excode_ipimpl::ip16_fxu1_avx512, excode_ipimpl::ip64_fxu2_avx512, @@ -288,8 +280,8 @@ ExcodeIpTable resolve_excode_ip_table() { excode_ipimpl::ip64_fxu6_avx512, excode_ipimpl::ip64_fxu7_avx512, excode_ipimpl::ip16_fxu8_avx512, - }, - ExcodeIpTable{ + }), + (ExcodeIpTable{ ip_fxu0, excode_ipimpl::ip16_fxu1_avx2, excode_ipimpl::ip64_fxu2_avx2, @@ -299,54 +291,62 @@ ExcodeIpTable resolve_excode_ip_table() { excode_ipimpl::ip64_fxu6_avx2, excode_ipimpl::ip64_fxu7_avx2, excode_ipimpl::ip16_fxu8_avx2, - }, - ExcodeIpTable{ - ip_fxu0, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - missing_excode_ip, - } + }), + generic ); +#endif } +#if defined(__aarch64__) || defined(_M_ARM64) +const auto kFlipSignFn = resolve_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon()); +#else const auto kFlipSignFn = - resolve_kernel(flip_sign_avx512, flip_sign_avx2, missing_flip_sign); + RABITQ_RESOLVE(flip_sign_avx512, flip_sign_avx2, flip_sign_generic); +#endif +#if defined(__aarch64__) || defined(_M_ARM64) +const auto kKacsWalkFn = resolve_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon()); +#else const auto kKacsWalkFn = - resolve_kernel(kacs_walk_avx512, kacs_walk_avx2, missing_kacs_walk); + RABITQ_RESOLVE(kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic); +#endif +#if defined(__aarch64__) || defined(_M_ARM64) const auto kScalarQuantizeUint8Fn = resolve_kernel( - scalar_quantize_uint8_avx512, scalar_quantize_uint8_avx2, missing_scalar_quantize_uint8 + scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() ); - const auto kScalarQuantizeUint16Fn = resolve_kernel( + scalar_quantize_uint16_neon, scalar_quantize_uint16_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, - missing_scalar_quantize_uint16 + scalar_quantize_uint16_generic ); +#endif -const auto kPacking2BitExcodeFn = resolve_kernel( - packing_2bit_excode_avx512, packing_2bit_excode_avx2, missing_pack_excode +const auto kPacking2BitExcodeFn = RABITQ_RESOLVE( + packing_2bit_excode_avx512, packing_2bit_excode_avx2, packing_2bit_excode_generic ); -const auto kPacking3BitExcodeFn = resolve_kernel( - packing_3bit_excode_avx512, packing_3bit_excode_avx2, missing_pack_excode +const auto kPacking3BitExcodeFn = RABITQ_RESOLVE( + packing_3bit_excode_avx512, packing_3bit_excode_avx2, packing_3bit_excode_generic ); -const auto kPacking4BitExcodeFn = resolve_kernel( - packing_4bit_excode_avx512, packing_4bit_excode_avx2, missing_pack_excode +const auto kPacking4BitExcodeFn = RABITQ_RESOLVE( + packing_4bit_excode_avx512, packing_4bit_excode_avx2, packing_4bit_excode_generic ); -const auto kPacking5BitExcodeFn = resolve_kernel( - packing_5bit_excode_avx512, packing_5bit_excode_avx2, missing_pack_excode +const auto kPacking5BitExcodeFn = RABITQ_RESOLVE( + packing_5bit_excode_avx512, packing_5bit_excode_avx2, packing_5bit_excode_generic ); -const auto kPacking6BitExcodeFn = resolve_kernel( - packing_6bit_excode_avx512, packing_6bit_excode_avx2, missing_pack_excode +const auto kPacking6BitExcodeFn = RABITQ_RESOLVE( + packing_6bit_excode_avx512, packing_6bit_excode_avx2, packing_6bit_excode_generic ); -const auto kPacking7BitExcodeFn = resolve_kernel( - packing_7bit_excode_avx512, packing_7bit_excode_avx2, missing_pack_excode +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) { @@ -405,24 +405,41 @@ const ex_ipfunc kIp64Fxu5AvxFn = kExcodeIpTable[5]; const ex_ipfunc kIp64Fxu6AvxFn = kExcodeIpTable[6]; const ex_ipfunc kIp64Fxu7AvxFn = kExcodeIpTable[7]; -const auto kNewTransposeBinFn = rabitqlib::simd::resolve_kernel( +#if defined(__aarch64__) || defined(_M_ARM64) +const auto kNewTransposeBinFn = simd::resolve_kernel( + simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon() +); +const auto kNewTransposeBin512Fn = simd::resolve_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::missing_new_transpose_bin + simd::new_transpose_bin_generic ); -const auto kNewTransposeBin512Fn = rabitqlib::simd::resolve_kernel( +const auto kNewTransposeBin512Fn = RABITQ_RESOLVE( simd::new_transpose_bin_512_avx512, simd::new_transpose_bin_512_avx2, - simd::missing_new_transpose_bin_512 + simd::new_transpose_bin_512_generic ); +#endif using MaskIpX0QFn = float (*)(const float*, const uint8_t*, size_t); -const MaskIpX0QFn kMaskIpX0QFn = rabitqlib::simd::resolve_kernel( +#if defined(__aarch64__) || defined(_M_ARM64) +const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_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), - simd::missing_mask_ip_x0_q + static_cast(simd::mask_ip_x0_q_generic) ); +#endif ex_ipfunc select_excode_ipfunc(size_t ex_bits) { if (ex_bits <= 8) { @@ -496,32 +513,45 @@ float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) namespace rabitqlib::fastscan { +#if defined(__aarch64__) || defined(_M_ARM64) const auto kPackLutFn = rabitqlib::simd::resolve_kernel( - simd::pack_lut_avx512, simd::pack_lut_avx2, simd::pack_lut_generic + simd::pack_lut_neon, simd::pack_lut_generic, cpu::has_neon() ); +#else +const auto kPackLutFn = + RABITQ_RESOLVE(simd::pack_lut_avx512, simd::pack_lut_avx2, simd::pack_lut_generic); +#endif template <> void pack_lut(size_t dim, const float* __restrict__ query, float* __restrict__ lut) { kPackLutFn(dim, query, lut); } +#if defined(__aarch64__) const auto kAccumulateFn = rabitqlib::simd::resolve_kernel( - simd::accumulate_avx512, - simd::accumulate_avx2, - rabitqlib::simd::missing_fastscan_accumulate + 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 = rabitqlib::simd::resolve_kernel( +const auto kTransferLutHaccFn = RABITQ_RESOLVE( simd::transfer_lut_hacc_avx512, simd::transfer_lut_hacc_avx2, - rabitqlib::simd::missing_fastscan_transfer_lut_hacc + simd::transfer_lut_hacc_generic ); +#if defined(__aarch64__) const auto kAccumulateHaccFn = rabitqlib::simd::resolve_kernel( - simd::accumulate_hacc_avx512, - simd::accumulate_hacc_avx2, - rabitqlib::simd::missing_fastscan_accumulate_hacc + simd::accumulate_hacc_neon, simd::accumulate_hacc_generic, cpu::has_neon() +); +#else +const auto kAccumulateHaccFn = RABITQ_RESOLVE( + simd::accumulate_hacc_avx512, simd::accumulate_hacc_avx2, simd::accumulate_hacc_generic ); +#endif void accumulate( const uint8_t* __restrict__ codes, @@ -564,13 +594,20 @@ namespace rabitqlib { using WarmupIpX0Q512Fn = float (*)(const uint8_t*, const uint64_t*, float, float, size_t, size_t); -const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = - rabitqlib::simd::resolve_kernel( - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), - rabitqlib::simd::missing_warmup_ip_x0_q_512, - cpu::has_avx512_popcnt() - ); +#if defined(__aarch64__) || defined(_M_ARM64) +const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = simd::resolve_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, @@ -602,20 +639,20 @@ namespace rabitqlib::hnsw::detail { namespace { using SearchKnnFn = std::priority_queue> (*)(HierarchicalNSW&, const float*, size_t); -std::priority_queue> missing_search_knn( - HierarchicalNSW&, const float*, size_t -) { - throw std::runtime_error("HNSW search requires AVX2/FMA or AVX512 support"); -} // 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 - : rabitqlib::simd::resolve_kernel( + : RABITQ_RESOLVE( search_knn_avx512_core, search_knn_avx2, - missing_search_knn, + search_knn_generic, cpu::has_avx512_core() && cpu::has_avx2() ); +#else +const SearchKnnFn kSearchKnnFn = + simd::resolve_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 diff --git a/src/simd/fastscan_generic.cpp b/src/simd/fastscan_generic.cpp index 6a20fb1e..d29c21ee 100644 --- a/src/simd/fastscan_generic.cpp +++ b/src/simd/fastscan_generic.cpp @@ -1,4 +1,6 @@ #include +#include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/fastscan/fastscan.hpp" @@ -18,3 +20,63 @@ void pack_lut_generic(size_t dim, const float* query, float* lut) { } } } // namespace rabitqlib::fastscan::simd + +namespace rabitqlib::fastscan::simd { +void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { + // Match the AVX2 layout: two low-byte tables, then two high-byte tables. + for (size_t group = 0; group < dim / 4; ++group) { + const size_t offset = group / 2 * 64 + group % 2 * 16; + for (size_t i = 0; i < 16; ++i) { + hc_lut[offset + i] = static_cast(lut[group * 16 + i]); + hc_lut[offset + 32 + i] = static_cast(lut[group * 16 + i] >> 8); + } + } +} +namespace { +template +void accumulate_generic_impl( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + int64_t totals[32]{}; + for (size_t group = 0; group < dim / 4; ++group) { + const size_t offset = HighAccuracy ? group / 2 * 64 + group % 2 * 16 : group * 16; + for (size_t lane = 0; lane < 16; ++lane) { + const uint8_t code = codes[group * 16 + lane]; + for (size_t half = 0; half < 2; ++half) { + const size_t entry = (code >> (half * 4)) & 15; + int value = lut[offset + entry]; + if constexpr (HighAccuracy) + value += int{lut[offset + 32 + entry]} * 256; + totals[kPerm0[lane] + half * 16] += value; + } + } + } + for (size_t lane = 0; lane < 32; ++lane) { + if (totals[lane] > std::numeric_limits::max()) { + throw std::overflow_error( + HighAccuracy ? "high-accuracy FastScan result exceeds int32_t" + : "FastScan result exceeds int32_t" + ); + } + } + for (size_t lane = 0; lane < 32; ++lane) + result[lane] = static_cast(totals[lane]); +} +} // namespace +void accumulate_generic( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + accumulate_generic_impl(codes, lut, result, dim); +} +void accumulate_unsupported(const uint8_t*, const uint8_t*, int32_t*, size_t) { + throw std::runtime_error( + "Standard FastScan accumulation requires AVX2/FMA, AVX-512, or ARM NEON; " + "no supported SIMD backend is available" + ); +} +void accumulate_hacc_generic( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + accumulate_generic_impl(codes, lut, result, dim); +} +} // namespace rabitqlib::fastscan::simd diff --git a/src/simd/fastscan_neon.cpp b/src/simd/fastscan_neon.cpp new file mode 100644 index 00000000..202bd996 --- /dev/null +++ b/src/simd/fastscan_neon.cpp @@ -0,0 +1,79 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/simd/fastscan_dispatch.hpp" + +namespace rabitqlib::fastscan::simd { +namespace { +template +void accumulate_impl( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + std::array totals{}; + // Bound each uint32 lane to 256*65535, then widen across chunks. + constexpr size_t kChunkGroups = 256; + for (size_t start = 0; start < dim / 4; start += kChunkGroups) { + uint32x4_t sums[8]; + for (auto& sum : sums) + sum = vdupq_n_u32(0); + const size_t end = std::min(start + kChunkGroups, dim / 4); + for (size_t group = start; group < end; ++group) { + const uint8x16_t packed = vld1q_u8(codes + group * 16); + const uint8x16_t indices[2] = { + vandq_u8(packed, vdupq_n_u8(15)), vshrq_n_u8(packed, 4)}; + const size_t offset = + HighAccuracy ? group / 2 * 64 + group % 2 * 16 : group * 16; + const uint8x16_t low = vld1q_u8(lut + offset); + for (size_t half = 0; half < 2; ++half) { + const uint8x16_t values = vqtbl1q_u8(low, indices[half]); + uint16x8_t lo = vmovl_u8(vget_low_u8(values)); + uint16x8_t hi = vmovl_u8(vget_high_u8(values)); + if constexpr (HighAccuracy) { + const uint8x16_t high = + vqtbl1q_u8(vld1q_u8(lut + offset + 32), indices[half]); + lo = vorrq_u16(lo, vshlq_n_u16(vmovl_u8(vget_low_u8(high)), 8)); + hi = vorrq_u16(hi, vshlq_n_u16(vmovl_u8(vget_high_u8(high)), 8)); + } + sums[half * 4] = vaddw_u16(sums[half * 4], vget_low_u16(lo)); + sums[half * 4 + 1] = vaddw_u16(sums[half * 4 + 1], vget_high_u16(lo)); + sums[half * 4 + 2] = vaddw_u16(sums[half * 4 + 2], vget_low_u16(hi)); + sums[half * 4 + 3] = vaddw_u16(sums[half * 4 + 3], vget_high_u16(hi)); + } + } + uint32_t chunk[32]; + for (size_t i = 0; i < 8; ++i) + vst1q_u32(chunk + i * 4, sums[i]); + for (size_t i = 0; i < 32; ++i) + totals[i] += chunk[i]; + } + for (auto value : totals) { + if (value > std::numeric_limits::max()) { + throw std::overflow_error( + HighAccuracy ? "high-accuracy FastScan result exceeds int32_t" + : "FastScan result exceeds int32_t" + ); + } + } + for (size_t i = 0; i < 32; ++i) { + result[kPerm0[i % 16] + i / 16 * 16] = static_cast(totals[i]); + } +} +} // namespace +void accumulate_neon( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + accumulate_impl(codes, lut, result, dim); +} +void accumulate_hacc_neon( + const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim +) { + accumulate_impl(codes, lut, result, dim); +} +} // namespace rabitqlib::fastscan::simd diff --git a/src/simd/hnsw_neon_kernels.hpp b/src/simd/hnsw_neon_kernels.hpp new file mode 100644 index 00000000..f7cde1d3 --- /dev/null +++ b/src/simd/hnsw_neon_kernels.hpp @@ -0,0 +1,85 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace rabitqlib::simd::detail { + +inline float mask_ip_neon(const float* query, const uint8_t* data, size_t dim) { + float32x4_t sum0 = vdupq_n_f32(0), sum1 = sum0, sum2 = sum0, sum3 = sum0; + const int32x4_t shifts = {-31, -30, -29, -28}; + for (size_t block = 0; block < dim; block += 64) { + uint64_t word; + std::memcpy(&word, data + block / 8, sizeof(word)); + // Sign-code coordinates run from bit 63 to bit 0 in each stored word. + for (size_t half = 0; half < 2; ++half) { + uint32x4_t bits = vdupq_n_u32(static_cast(word >> 32)); + for (size_t group = 0; group < 2; ++group) { + const auto masked = [&](size_t offset, int shift) { + const uint32x4_t selected = vandq_u32( + vshlq_u32(bits, vaddq_s32(shifts, vdupq_n_s32(shift))), + vdupq_n_u32(1) + ); + return vreinterpretq_f32_u32(vandq_u32( + vsubq_u32(vdupq_n_u32(0), selected), + vreinterpretq_u32_f32(vld1q_f32(query + offset)) + )); + }; + sum0 = vaddq_f32(sum0, masked(0, 0)); + sum1 = vaddq_f32(sum1, masked(4, 4)); + sum2 = vaddq_f32(sum2, masked(8, 8)); + sum3 = vaddq_f32(sum3, masked(12, 12)); + bits = vshlq_n_u32(bits, 16); + query += 16; + } + word <<= 32; + } + } + return vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); +} + +inline float warmup_ip_neon( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +) { + if (bits > 8) { + throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); + } + uint64_t ip = 0, count = 0; + for (size_t block = 0; block < dim; block += 512) { + const size_t chunks = std::min(size_t{512}, dim - block) / 64; + size_t chunk = 0; + for (; chunk + 2 <= chunks; chunk += 2) { + const uint8x16_t code = vld1q_u8(data + block / 8 + chunk * 8); + count += vaddlvq_u8(vcntq_u8(code)); + for (size_t b = 0; b < bits; ++b) { + const uint8x16_t plane = + vreinterpretq_u8_u64(vld1q_u64(query + b * chunks + chunk)); + ip += uint64_t{vaddlvq_u8(vcntq_u8(vandq_u8(code, plane)))} << b; + } + } + if (chunk < chunks) { + const uint8x8_t code = vld1_u8(data + block / 8 + chunk * 8); + count += vaddlv_u8(vcnt_u8(code)); + for (size_t b = 0; b < bits; ++b) { + const uint8x8_t plane = + vreinterpret_u8_u64(vld1_u64(query + b * chunks + chunk)); + ip += uint64_t{vaddlv_u8(vcnt_u8(vand_u8(code, plane)))} << b; + } + } + if (bits != 0) { + query += chunks * bits; + } + } + return delta * static_cast(ip) + vl * static_cast(count); +} +} // namespace rabitqlib::simd::detail diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp new file mode 100644 index 00000000..2f78e400 --- /dev/null +++ b/src/simd/portable_generic.cpp @@ -0,0 +1,260 @@ +#include +#include +#include +#include +#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/utils/bitops.hpp" + +namespace rabitqlib::simd { +namespace { +// Coordinate-to-bit mapping of the existing x86 packed format. Odd widths +// store the high bit in eight bytes, with coordinate i at byte i%8, bit i/8. +template +size_t packed_bit(size_t i, size_t bit) { + if constexpr (Bits == 1) { + return i; + } else if constexpr (Bits == 8) { + return i * 8 + bit; + } else if constexpr (Bits == 4) { + return (i / 16) * 64 + (i % 8) * 8 + ((i % 16) / 8) * 4 + bit; + } else { + const size_t block = (i / 64) * Bits * 64; + i %= 64; + if constexpr (Bits % 2 == 1) { + if (bit == Bits - 1) { + return block + (Bits - 1) * 64 + (i % 8) * 8 + i / 8; + } + } + if constexpr (Bits <= 3) { + return block + (i % 16) * 8 + (i / 16) * 2 + bit; + } else if constexpr (Bits == 5) { + return block + (i / 32) * 128 + (i % 16) * 8 + ((i / 16) % 2) * 4 + bit; + } else { + return block + + (i < 48 ? i * 8 + bit : (bit / 2) * 128 + (i % 16) * 8 + 6 + bit % 2); + } + } +} + +template +void pack(const uint8_t* raw, uint8_t* compact, size_t dim) { + std::fill_n(compact, dim * Bits / 8, uint8_t{0}); + for (size_t i = 0; i < dim; ++i) { + for (size_t b = 0; b < Bits; ++b) { + const size_t pos = packed_bit(i, b); + compact[pos / 8] |= ((raw[i] >> b) & 1U) << (pos % 8); + } + } +} + +template +float excode_ip(const float* query, const uint8_t* compact, size_t dim) { + double sum = 0; + for (size_t i = 0; i < dim; ++i) { + unsigned code = 0; + for (size_t b = 0; b < Bits; ++b) { + const size_t pos = packed_bit(i, b); + code |= ((compact[pos / 8] >> (pos % 8)) & 1U) << b; + } + sum += static_cast(query[i]) * code; + } + return static_cast(sum); +} + +template +void quantize(T* result, const float* data, size_t dim, float lo, float delta) { + const float reciprocal = 1.0F / delta; + for (size_t i = 0; i < dim; ++i) { + result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); + } +} + +void fht(float* data, size_t dim) { + for (size_t width = 1; width < dim; width *= 2) { + for (size_t block = 0; block < dim; block += width * 2) { + for (size_t j = 0; j < width; ++j) { + const float a = data[block + j], b = data[block + j + width]; + data[block + j] = a + b; + data[block + j + width] = a - b; + } + } + } +} +} // namespace + +#define RABITQ_PACK(Bits) \ + void packing_##Bits##bit_excode_generic( \ + const uint8_t* raw, uint8_t* compact, size_t dim \ + ) { \ + pack(raw, compact, dim); \ + } +RABITQ_PACK(2) +RABITQ_PACK(3) +RABITQ_PACK(4) +RABITQ_PACK(5) +RABITQ_PACK(6) +RABITQ_PACK(7) +#undef RABITQ_PACK + +namespace excode_ipimpl { +#define RABITQ_IP(Block, Bits) \ + float ip##Block##_fxu##Bits##_generic( \ + const float* query, const uint8_t* code, size_t dim \ + ) { \ + return excode_ip(query, code, dim); \ + } +RABITQ_IP(16, 1) +RABITQ_IP(64, 2) +RABITQ_IP(64, 3) +RABITQ_IP(16, 4) +RABITQ_IP(64, 5) +RABITQ_IP(64, 6) +RABITQ_IP(64, 7) +RABITQ_IP(16, 8) +#undef RABITQ_IP +} // namespace excode_ipimpl + +void scalar_quantize_uint8_generic( + uint8_t* result, const float* data, size_t dim, float lo, float delta +) { + quantize(result, data, dim, lo, delta); +} +void scalar_quantize_uint16_generic( + uint16_t* result, const float* data, size_t dim, float lo, float delta +) { + quantize(result, data, dim, lo, delta); +} + +void new_transpose_bin_generic( + const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits +) { + for (size_t block = 0; block < dim; block += 64) { + for (size_t b = 0; b < bits; ++b) { + uint64_t word = 0; + for (size_t i = 0; i < 64; ++i) { + word |= uint64_t{(query[block + i] >> b) & 1U} << (63 - i); + } + *transposed++ = word; + } + } +} + +void new_transpose_bin_512_generic( + const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits +) { + for (size_t block = 0; block < dim; block += 512) { + const size_t chunks = std::min(size_t{512}, dim - block) / 64; + for (size_t b = 0; b < bits; ++b) { + for (size_t chunk = 0; chunk < chunks; ++chunk) { + uint64_t word = 0; + for (size_t i = 0; i < 64; ++i) { + word |= uint64_t{(query[block + chunk * 64 + i] >> b) & 1U} << (63 - i); + } + *transposed++ = word; + } + } + } +} + +float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t dim) { + double sum = 0; + for (size_t block = 0; block < dim; block += 64) { + uint64_t word; + std::memcpy(&word, data + block / 8, sizeof(word)); + for (size_t i = 0; i < 64; ++i) { + if ((word >> (63 - i)) & 1U) + sum += query[block + i]; + } + } + return static_cast(sum); +} +float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t dim) { + return mask_ip_x0_q_generic(query, reinterpret_cast(data), dim); +} + +float warmup_ip_x0_q_512_generic( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +) { + if (bits > 8) + throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); + uint64_t ip = 0, count = 0; + for (size_t block = 0; block < dim; block += 512) { + const size_t chunks = std::min(size_t{512}, dim - block) / 64; + for (size_t chunk = 0; chunk < chunks; ++chunk) { + uint64_t word; + std::memcpy(&word, data + block / 8 + chunk * 8, sizeof(word)); + count += bitops::popcount64(word); + for (size_t b = 0; b < bits; ++b) { + ip += uint64_t{bitops::popcount64(word & query[b * chunks + chunk])} << b; + } + } + query += chunks * bits; + } + return delta * static_cast(ip) + vl * static_cast(count); +} +float warmup_ip_x0_q_512_generic( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +) { + return warmup_ip_x0_q_512_generic( + reinterpret_cast(data), query, delta, vl, dim, bits + ); +} + +void flip_sign_generic(const uint8_t* flip, float* data, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + if ((flip[i / 8] >> (i % 8)) & 1U) + data[i] = -data[i]; + } +} +void kacs_walk_generic(float* data, size_t len) { + for (size_t i = 0; i < len / 2; ++i) { + const float a = data[i], b = data[i + len / 2]; + data[i] = a + b; + data[i + len / 2] = a - b; + } +} +void fht_rotate_generic( + const float* data, + float* output, + size_t dim, + size_t padded, + size_t trunc, + float factor, + const uint8_t* flip +) { + if (trunc < 64 || trunc > 65536 || (trunc & (trunc - 1)) != 0) { + throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); + } + std::memmove(output, data, dim * sizeof(float)); + std::fill(output + dim, output + padded, 0.0F); + for (size_t pass = 0; pass < 4; ++pass) { + flip_sign_generic(flip + pass * padded / 8, output, padded); + const size_t start = pass % 2 == 0 ? 0 : padded - trunc; + fht(output + start, trunc); + for (size_t i = start; i < start + trunc; ++i) + output[i] *= factor; + if (padded != trunc) + kacs_walk_generic(output, padded); + } + if (padded != trunc) { + for (size_t i = 0; i < padded; ++i) + output[i] *= 0.25F; + } +} +} // namespace rabitqlib::simd diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp new file mode 100644 index 00000000..7d244a1a --- /dev/null +++ b/src/simd/query_neon.cpp @@ -0,0 +1,134 @@ +#include + +#include +#include +#include +#include + +#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" + +namespace rabitqlib::simd { +namespace { +// Convert eight selected byte lanes to MSB-first coordinate bits. +inline uint8_t byte_mask(uint8x8_t selected) { + const uint8x8_t weights = {128, 64, 32, 16, 8, 4, 2, 1}; + return vaddv_u8(vand_u8(selected, weights)); +} + +inline uint16_t byte_mask(uint8x16_t selected) { + return (uint16_t{byte_mask(vget_low_u8(selected))} << 8) | + byte_mask(vget_high_u8(selected)); +} + +inline uint32x4_t quantize4(const float* data, float32x4_t lo, float32x4_t reciprocal) { + // FCVTAU rounds ties away from zero, unlike nearest-even FCVTNU. + // Callers provide finite values whose rounded codes fit the output type. + return vcvtaq_u32_f32(vmulq_f32(vsubq_f32(vld1q_f32(data), lo), reciprocal)); +} +} // namespace + +void scalar_quantize_uint8_neon( + uint8_t* result, const float* data, size_t dim, float lo, float delta +) { + const float reciprocal = 1.0F / delta; + const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal); + size_t i = 0; + for (; i + 8 <= dim; i += 8) { + const uint16x8_t codes = vcombine_u16( + vmovn_u32(quantize4(data + i, low, scale)), + vmovn_u32(quantize4(data + i + 4, low, scale)) + ); + vst1_u8(result + i, vmovn_u16(codes)); + } + for (; i < dim; ++i) { + result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); + } +} + +void scalar_quantize_uint16_neon( + uint16_t* result, const float* data, size_t dim, float lo, float delta +) { + const float reciprocal = 1.0F / delta; + const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal); + size_t i = 0; + for (; i + 8 <= dim; i += 8) { + vst1q_u16( + result + i, + vcombine_u16( + vmovn_u32(quantize4(data + i, low, scale)), + vmovn_u32(quantize4(data + i + 4, low, scale)) + ) + ); + } + for (; i < dim; ++i) { + result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); + } +} + +void new_transpose_bin_neon( + const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits +) { + for (size_t block = 0; block < dim; block += 64) { + uint16x8_t values[8]; + for (size_t c = 0; c < 8; ++c) { + values[c] = vld1q_u16(query + block + c * 8); + } + for (size_t b = 0; b < bits; ++b) { + const uint16x8_t bit = vdupq_n_u16(uint16_t{1} << b); + uint64_t word = 0; + for (auto value : values) { + word = (word << 8) | byte_mask(vmovn_u16(vtstq_u16(value, bit))); + } + *transposed++ = word; + } + } +} + +void new_transpose_bin_512_neon( + const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits +) { + for (size_t block = 0; block < dim; block += 512) { + const size_t chunks = std::min(size_t{512}, dim - block) / 64; + for (size_t c = 0; c < chunks; ++c) { + const uint8_t* row = query + block + c * 64; + const uint8x16_t a = vld1q_u8(row), b = vld1q_u8(row + 16), + d = vld1q_u8(row + 32), e = vld1q_u8(row + 48); + for (size_t plane = 0; plane < bits; ++plane) { + const uint8x16_t bit = vdupq_n_u8(uint8_t{1} << plane); + transposed[plane * chunks + c] = + (uint64_t{byte_mask(vtstq_u8(a, bit))} << 48) | + (uint64_t{byte_mask(vtstq_u8(b, bit))} << 32) | + (uint64_t{byte_mask(vtstq_u8(d, bit))} << 16) | + byte_mask(vtstq_u8(e, bit)); + } + } + transposed += chunks * bits; + } +} +} // 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/rotator_neon.cpp b/src/simd/rotator_neon.cpp new file mode 100644 index 00000000..c06b9506 --- /dev/null +++ b/src/simd/rotator_neon.cpp @@ -0,0 +1,152 @@ +#include + +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/rotator_dispatch.hpp" + +namespace rabitqlib::simd { +namespace { +inline void butterfly(float32x4_t& a, float32x4_t& b) { + const float32x4_t sum = vaddq_f32(a, b); + b = vsubq_f32(a, b); + a = sum; +} + +inline float32x4_t fht4(float32x4_t x) { + // Width-one butterflies, then width-two butterflies. Retain the scalar + // operand order instead of reassociating sums or changing normalization. + const float32x4_t paired = vrev64q_f32(x); + x = vtrn1q_f32(vaddq_f32(x, paired), vsubq_f32(x, paired)); + const float32x2_t lo = vget_low_f32(x), hi = vget_high_f32(x); + return vcombine_f32(vadd_f32(lo, hi), vsub_f32(lo, hi)); +} + +void fht(float* data, size_t dim) { + // Fuse the first four stages in registers, avoiding a memory pass per stage. + for (size_t i = 0; i < dim; i += 16) { + float32x4_t a = fht4(vld1q_f32(data + i)); + float32x4_t b = fht4(vld1q_f32(data + i + 4)); + float32x4_t c = fht4(vld1q_f32(data + i + 8)); + float32x4_t d = fht4(vld1q_f32(data + i + 12)); + butterfly(a, b); + butterfly(c, d); + butterfly(a, c); + butterfly(b, d); + vst1q_f32(data + i, a); + vst1q_f32(data + i + 4, b); + vst1q_f32(data + i + 8, c); + vst1q_f32(data + i + 12, d); + } + // Fuse pairs of larger stages, preserving the same butterfly dependencies. + size_t width = 16; + for (; width * 2 < dim; width *= 4) { + for (size_t block = 0; block < dim; block += width * 4) { + for (size_t i = block; i < block + width; i += 4) { + float32x4_t a = vld1q_f32(data + i); + float32x4_t b = vld1q_f32(data + i + width); + float32x4_t c = vld1q_f32(data + i + 2 * width); + float32x4_t d = vld1q_f32(data + i + 3 * width); + butterfly(a, b); + butterfly(c, d); + butterfly(a, c); + butterfly(b, d); + vst1q_f32(data + i, a); + vst1q_f32(data + i + width, b); + vst1q_f32(data + i + 2 * width, c); + vst1q_f32(data + i + 3 * width, d); + } + } + } + if (width < dim) { + for (size_t i = 0; i < width; i += 4) { + float32x4_t a = vld1q_f32(data + i), b = vld1q_f32(data + i + width); + butterfly(a, b); + vst1q_f32(data + i, a); + vst1q_f32(data + i + width, b); + } + } +} + +void rescale(float* data, size_t dim, float factor) { + const float32x4_t scale = vdupq_n_f32(factor); + for (size_t i = 0; i < dim; i += 4) { + vst1q_f32(data + i, vmulq_f32(vld1q_f32(data + i), scale)); + } +} +} // namespace + +void flip_sign_neon(const uint8_t* flip, float* data, size_t dim) { + const int32x4_t shifts_lo = {31, 30, 29, 28}, shifts_hi = {27, 26, 25, 24}; + const uint32x4_t sign = vdupq_n_u32(0x80000000U); + size_t i = 0; + for (; i + 8 <= dim; i += 8) { + // The stored mask is LSB-first: bit j flips coordinate i+j. + const uint32x4_t bits = vdupq_n_u32(flip[i / 8]); + const uint32x4_t lo = vandq_u32(vshlq_u32(bits, shifts_lo), sign); + const uint32x4_t hi = vandq_u32(vshlq_u32(bits, shifts_hi), sign); + vst1q_f32( + data + i, + vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(vld1q_f32(data + i)), lo)) + ); + vst1q_f32( + data + i + 4, + vreinterpretq_f32_u32( + veorq_u32(vreinterpretq_u32_f32(vld1q_f32(data + i + 4)), hi) + ) + ); + } + for (; i < dim; ++i) { + if ((flip[i / 8] >> (i % 8)) & 1U) { + data[i] = -data[i]; + } + } +} + +void kacs_walk_neon(float* data, size_t len) { + const size_t half = len / 2; + size_t i = 0; + for (; i + 4 <= half; i += 4) { + float32x4_t a = vld1q_f32(data + i), b = vld1q_f32(data + i + half); + butterfly(a, b); + vst1q_f32(data + i, a); + vst1q_f32(data + i + half, b); + } + for (; i < half; ++i) { + const float a = data[i], b = data[i + half]; + data[i] = a + b; + data[i + half] = a - b; + } +} + +void fht_rotate_neon( + const float* data, + float* output, + size_t dim, + size_t padded, + size_t trunc, + float factor, + const uint8_t* flip +) { + if (trunc < 64 || trunc > 65536 || (trunc & (trunc - 1)) != 0) { + throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); + } + std::memmove(output, data, dim * sizeof(float)); + std::fill(output + dim, output + padded, 0.0F); + for (size_t pass = 0; pass < 4; ++pass) { + flip_sign_neon(flip + pass * padded / 8, output, padded); + const size_t start = pass % 2 == 0 ? 0 : padded - trunc; + fht(output + start, trunc); + rescale(output + start, trunc, factor); + if (padded != trunc) { + kacs_walk_neon(output, padded); + } + } + if (padded != trunc) { + rescale(output, padded, 0.25F); + } +} +} // namespace rabitqlib::simd diff --git a/src/simd/space_generic.cpp b/src/simd/space_generic.cpp index 64674dfe..f4fd15a9 100644 --- a/src/simd/space_generic.cpp +++ b/src/simd/space_generic.cpp @@ -1,25 +1,44 @@ +#include #include #include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::simd { +namespace { +// Independent float sums limit reduction depth without widening the arithmetic. +template +float product(const float* a, const float* b, size_t dim) { + std::array sums{}; + auto value = [&](size_t at) { + if constexpr (SquaredDistance) { + const float diff = a[at] - b[at]; + return diff * diff; + } else { + return a[at] * b[at]; + } + }; + size_t i = 0; + for (; dim - i >= sums.size(); i += sums.size()) { + for (size_t lane = 0; lane < sums.size(); ++lane) + sums[lane] += value(i + lane); + } + for (size_t stride = sums.size() / 2; stride != 0; stride /= 2) { + for (size_t lane = 0; lane < stride; ++lane) + sums[lane] += sums[lane + stride]; + } + float result = sums[0]; + for (; i < dim; ++i) + result += value(i); + return result; +} +} // namespace -// Use wider scalar accumulation to limit rounding error on long inputs. float euclidean_sqr_generic(const float* a, const float* b, size_t dim) { - double sum = 0; - for (size_t i = 0; i < dim; ++i) { - const double delta = static_cast(a[i]) - b[i]; - sum += delta * delta; - } - return static_cast(sum); + return product(a, b, dim); } float dot_product_generic(const float* a, const float* b, size_t dim) { - double sum = 0; - for (size_t i = 0; i < dim; ++i) { - sum += static_cast(a[i]) * b[i]; - } - return static_cast(sum); + return product(a, b, dim); } float dot_product_dis_generic(const float* a, const float* b, size_t dim) { @@ -27,11 +46,7 @@ float dot_product_dis_generic(const float* a, const float* b, size_t dim) { } float l2norm_sqr_generic(const float* a, size_t dim) { - double sum = 0; - for (size_t i = 0; i < dim; ++i) { - sum += static_cast(a[i]) * a[i]; - } - return static_cast(sum); + return dot_product_generic(a, a, dim); } } // namespace rabitqlib::simd diff --git a/src/simd/space_neon.cpp b/src/simd/space_neon.cpp new file mode 100644 index 00000000..55b579be --- /dev/null +++ b/src/simd/space_neon.cpp @@ -0,0 +1,56 @@ +#include + +#include + +#include "rabitqlib/simd/space_dispatch.hpp" + +namespace rabitqlib::simd { +namespace { +// Four independent float accumulation chains, matching the x86 kernel strategy. +template +float product(const float* a, const float* b, size_t dim) { + float32x4_t s0 = vdupq_n_f32(0), s1 = vdupq_n_f32(0); + float32x4_t s2 = vdupq_n_f32(0), s3 = vdupq_n_f32(0); + auto step = [&](float32x4_t sum, size_t at) { + const float32x4_t x = vld1q_f32(a + at); + const float32x4_t y = vld1q_f32(b + at); + if constexpr (SquaredDistance) { + const float32x4_t diff = vsubq_f32(x, y); + return vfmaq_f32(sum, diff, diff); + } else { + return vfmaq_f32(sum, x, y); + } + }; + size_t i = 0; + for (; dim - i >= 16; i += 16) { + s0 = step(s0, i); + s1 = step(s1, i + 4); + s2 = step(s2, i + 8); + s3 = step(s3, i + 12); + } + float32x4_t sum = vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)); + for (; dim - i >= 4; i += 4) + sum = step(sum, i); + float result = vaddvq_f32(sum); + for (; i < dim; ++i) { + if constexpr (SquaredDistance) { + const float diff = a[i] - b[i]; + result += diff * diff; + } else { + result += a[i] * b[i]; + } + } + return result; +} +} // namespace +float euclidean_sqr_neon(const float* a, const float* b, size_t dim) { + return product(a, b, dim); +} +float dot_product_neon(const float* a, const float* b, size_t dim) { + return product(a, b, dim); +} +float dot_product_dis_neon(const float* a, const float* b, size_t dim) { + return 1.0F - dot_product_neon(a, b, dim); +} +float l2norm_sqr_neon(const float* a, size_t dim) { return dot_product_neon(a, a, dim); } +} // namespace rabitqlib::simd diff --git a/src/simd/space_packed_neon.cpp b/src/simd/space_packed_neon.cpp new file mode 100644 index 00000000..4142970f --- /dev/null +++ b/src/simd/space_packed_neon.cpp @@ -0,0 +1,160 @@ +#include + +#include +#include + +#include "hnsw_neon_kernels.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/simd/warmup_dispatch.hpp" + +namespace rabitqlib::simd { +namespace { +struct FloatDot { + float32x4_t sum0 = vdupq_n_f32(0), sum1 = sum0, sum2 = sum0, sum3 = sum0; + + void add(uint8x16_t codes, const float* query) { + const uint16x8_t lo = vmovl_u8(vget_low_u8(codes)); + const uint16x8_t hi = vmovl_u8(vget_high_u8(codes)); + sum0 = + vfmaq_f32(sum0, vld1q_f32(query), vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo)))); + sum1 = vfmaq_f32( + sum1, vld1q_f32(query + 4), vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo))) + ); + sum2 = vfmaq_f32( + sum2, vld1q_f32(query + 8), vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi))) + ); + sum3 = vfmaq_f32( + sum3, vld1q_f32(query + 12), vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi))) + ); + } + float result() const { + return vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); + } +}; + +// Odd-width top bits use byte i%8, bit i/8, independently of the low-bit layout. +template +uint8x16_t add_top(uint8x16_t low, const uint8_t* compact) { + const uint8x8_t top = vld1_u8(compact + (Bits - 1) * 8); + const uint8x16_t repeated = vcombine_u8(top, top); + const int8x16_t shift = + vcombine_s8(vdup_n_s8(Bits - 1 - Group * 2), vdup_n_s8(Bits - 2 - Group * 2)); + return vorrq_u8(low, vandq_u8(vshlq_u8(repeated, shift), vdupq_n_u8(1U << (Bits - 1)))); +} + +template +float packed_ip(const float* query, const uint8_t* compact, size_t dim) { + FloatDot sum; + if constexpr (Bits == 1 || Bits == 4 || Bits == 8) { + for (size_t i = 0; i < dim; i += 16) { + uint8x16_t codes; + if constexpr (Bits == 1) { + const int8x8_t shifts = {0, -1, -2, -3, -4, -5, -6, -7}; + codes = vandq_u8( + vcombine_u8( + vshl_u8(vdup_n_u8(compact[0]), shifts), + vshl_u8(vdup_n_u8(compact[1]), shifts) + ), + vdupq_n_u8(1) + ); + } else if constexpr (Bits == 4) { + const uint8x8_t packed = vld1_u8(compact); + codes = vcombine_u8(vand_u8(packed, vdup_n_u8(15)), vshr_n_u8(packed, 4)); + } else { + codes = vld1q_u8(compact); + } + sum.add(codes, query + i); + compact += 2 * Bits; + } + } else { + for (size_t i = 0; i < dim; i += 64) { + uint8x16_t c0, c1, c2, c3; + if constexpr (Bits <= 3) { + const uint8x16_t packed = vld1q_u8(compact), mask = vdupq_n_u8(3); + c0 = vandq_u8(packed, mask); + c1 = vandq_u8(vshrq_n_u8(packed, 2), mask); + c2 = vandq_u8(vshrq_n_u8(packed, 4), mask); + c3 = vshrq_n_u8(packed, 6); + } else if constexpr (Bits == 5) { + const uint8x16_t lo = vld1q_u8(compact), hi = vld1q_u8(compact + 16); + c0 = vandq_u8(lo, vdupq_n_u8(15)); + c1 = vshrq_n_u8(lo, 4); + c2 = vandq_u8(hi, vdupq_n_u8(15)); + c3 = vshrq_n_u8(hi, 4); + } else { + const uint8x16_t p0 = vld1q_u8(compact), p1 = vld1q_u8(compact + 16), + p2 = vld1q_u8(compact + 32); + c0 = vandq_u8(p0, vdupq_n_u8(63)); + c1 = vandq_u8(p1, vdupq_n_u8(63)); + c2 = vandq_u8(p2, vdupq_n_u8(63)); + c3 = vorrq_u8( + vshrq_n_u8(p0, 6), + vorrq_u8( + vshlq_n_u8(vshrq_n_u8(p1, 6), 2), vshlq_n_u8(vshrq_n_u8(p2, 6), 4) + ) + ); + } + if constexpr (Bits % 2 == 1) { + c0 = add_top(c0, compact); + c1 = add_top(c1, compact); + c2 = add_top(c2, compact); + c3 = add_top(c3, compact); + } + sum.add(c0, query + i); + sum.add(c1, query + i + 16); + sum.add(c2, query + i + 32); + sum.add(c3, query + i + 48); + compact += 8 * Bits; + } + } + return sum.result(); +} +} // namespace + +namespace excode_ipimpl { +#define RABITQ_NEON_IP(Block, Bits) \ + float ip##Block##_fxu##Bits##_neon( \ + const float* query, const uint8_t* code, size_t dim \ + ) { \ + return packed_ip(query, code, dim); \ + } +RABITQ_NEON_IP(16, 1) +RABITQ_NEON_IP(64, 2) +RABITQ_NEON_IP(64, 3) +RABITQ_NEON_IP(16, 4) +RABITQ_NEON_IP(64, 5) +RABITQ_NEON_IP(64, 6) +RABITQ_NEON_IP(64, 7) +RABITQ_NEON_IP(16, 8) +#undef RABITQ_NEON_IP +} // namespace excode_ipimpl + +float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t dim) { + return detail::mask_ip_neon(query, data, dim); +} +float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t dim) { + return detail::mask_ip_neon(query, reinterpret_cast(data), dim); +} +float warmup_ip_x0_q_512_neon( + const uint8_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +) { + return detail::warmup_ip_neon(data, query, delta, vl, dim, bits); +} +float warmup_ip_x0_q_512_neon( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t dim, + size_t bits +) { + return detail::warmup_ip_neon( + reinterpret_cast(data), query, delta, vl, dim, bits + ); +} +} // namespace rabitqlib::simd diff --git a/src/utils/cpu_features.cpp b/src/utils/cpu_features.cpp index 6ddba80f..099b6128 100644 --- a/src/utils/cpu_features.cpp +++ b/src/utils/cpu_features.cpp @@ -55,6 +55,10 @@ uint64_t xgetbv(uint32_t) { return 0; } Features detect_features() { Features hardware{}; +#if defined(__aarch64__) + // AdvSIMD is part of the ARM64 ABI on macOS and Linux AArch64. + hardware.neon = true; +#endif if constexpr (!kIsX86) { return hardware; } @@ -148,6 +152,8 @@ const Features& features() { return detected; } +bool has_neon() { return features().neon; } + bool has_avx2() { const Features& detected = features(); return detected.avx2 && detected.fma; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a4a64262..7c51202f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -39,6 +39,13 @@ target_link_libraries(rabitq_tests include(GoogleTest) gtest_discover_tests(rabitq_tests DISCOVERY_MODE PRE_TEST) +# 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_enable_warnings(fastscan_no_simd_test) +gtest_discover_tests(fastscan_no_simd_test DISCOVERY_MODE PRE_TEST) + # Message to verify files are found message(STATUS "Discovered test files:") foreach(TEST_SRC ${UNIT_TESTS} ${INTEG_TESTS}) diff --git a/tests/README.md b/tests/README.md index 24223f1f..e3169d37 100644 --- a/tests/README.md +++ b/tests/README.md @@ -8,7 +8,7 @@ binding tests for RaBitQ Library. - CMake 3.20 or newer for the `ctest --test-dir` commands below (4.2 or newer for the Visual Studio 2026 generator) - A C++17 compiler with OpenMP support (GCC, Clang, or Visual Studio 2026) -- An x86-64 CPU with AVX2 and FMA; AVX-512 is optional +- An x86-64 CPU with AVX2 and FMA (AVX-512 optional), or an Apple Silicon Mac - Git and network access during the first configuration so CMake can download GoogleTest 1.14.0 @@ -45,8 +45,35 @@ cmake --build build --parallel ctest --test-dir build --output-on-failure ``` +#### macOS ARM64 + +Use a native ARM64 project Python environment. Install CMake, Ninja and the existing +OpenMP dependency (AppleClang does not ship an OpenMP runtime): + +```bash +brew install cmake ninja libomp +cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DRABITQ_BUILD_TESTS=ON -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DOpenMP_ROOT="$(brew --prefix libomp)" +cmake --build build --parallel +OMP_NUM_THREADS=2 ctest --test-dir build --output-on-failure +``` + +An environment providing `llvm-openmp` can use its prefix as `OpenMP_ROOT` instead. +CMake excludes AVX source files on ARM64. Reference tests run through generic and +NEON kernels; explicit AVX backend tests are guarded or skipped. Linux AArch64 +needs its own build and execution validation. + +For Python source builds, pass `-Ccmake.define.OpenMP_ROOT=` as well +as the native-optimization setting below. Release CI uses native `macos-14` runners, +builds CPython 3.11–3.14 ARM64 wheels, bundles OpenMP with delocate, then installs and +runs the full Python suite for each wheel. The minimum wheel OS target is macOS 14. +Set `RABITQ_TEST_WHEEL=1` when testing a repaired wheel locally to also check its +ARM64 extension and bundled OpenMP loader path. Intel and universal2 wheels are excluded. + The combined executable is `build/tests/Release/rabitq_tests.exe` on Windows -and `build/tests/rabitq_tests` on Linux. To run a subset, add `-R ` to +and `build/tests/rabitq_tests` on Linux/macOS. To run a subset, add `-R ` to the CTest command. AVX-512 tests skip when the CPU or OS lacks the required features. Tests using @@ -68,6 +95,13 @@ python -m pytest tests/python -ra -q Confirm the printed extension path belongs to the intended environment. The suite includes all three indexes, quantization, persistence, and Unicode paths. +The example integration tests additionally need `faiss-cpu`. Wheel CI installs it +as a test dependency and runs Faiss clustering and RaBitQ indexing/querying in +separate subprocesses, with guards against importing both libraries together. +These tests cover L2/IP and the saved clustering file used by the +[Python examples](../sample/python/README.md). They skip the Faiss-dependent cases +when Faiss is absent from a local test environment. + ## Installed CMake package test After building the library, verify that another project can consume its installed diff --git a/tests/dispatch/fastscan_no_simd_test.cpp b/tests/dispatch/fastscan_no_simd_test.cpp new file mode 100644 index 00000000..f95ebb47 --- /dev/null +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -0,0 +1,44 @@ +#include + +#include +#include +#include + +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +// Link-time substitutes, isolated from both production detection and the main suite. +namespace rabitqlib::cpu { +bool has_avx2() { return false; } +bool has_neon() { return false; } +bool has_avx512_core() { return false; } +bool has_avx512_popcnt() { return false; } +} // namespace rabitqlib::cpu + +namespace rabitqlib::fastscan { +namespace { + +TEST(FastScanDispatchTest, UnsupportedCpuThrowsWithoutAccessingInputsOrOutput) { + std::array result; + result.fill(-123); + const auto original = result; + try { + accumulate(nullptr, nullptr, result.data(), 16); + FAIL() << "Standard FastScan must reject an unsupported SIMD backend"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "Standard FastScan accumulation requires AVX2/FMA, AVX-512, or ARM NEON; " + "no supported SIMD backend is available" + ); + } + EXPECT_EQ(result, original); +} + +TEST(FastScanDispatchTest, InvalidDimensionStillRaisesInvalidArgument) { + EXPECT_THROW(accumulate(nullptr, nullptr, nullptr, 0), std::invalid_argument); + EXPECT_THROW(accumulate(nullptr, nullptr, nullptr, 15), std::invalid_argument); +} + +} // namespace +} // namespace rabitqlib::fastscan diff --git a/tests/integration/bit_pack_unpack_test.cpp b/tests/integration/bit_pack_unpack_test.cpp index 5401ae80..57bc3849 100644 --- a/tests/integration/bit_pack_unpack_test.cpp +++ b/tests/integration/bit_pack_unpack_test.cpp @@ -48,17 +48,17 @@ class BitPackUnpackTest : public ::testing::Test { } // 4. Helper for Ground Truth Calculation - float CalculateExpected() const { - float expected_result = 0.0f; + double CalculateExpected() const { + double expected_result = 0; for (size_t i = 0; i < dim; ++i) { - expected_result += query[i] * static_cast(code[i]); + expected_result += static_cast(query[i]) * code[i]; } return expected_result; } void ExpectIpNear(float result) const { - float expected = CalculateExpected(); - ASSERT_NEAR(expected, result, std::max(0.1f, std::abs(expected) * 1e-6f)); + double expected = CalculateExpected(); + ASSERT_NEAR(expected, result, std::max(0.1, std::abs(expected) * 1e-6)); } }; diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py new file mode 100644 index 00000000..aa94e476 --- /dev/null +++ b/tests/python/test_examples.py @@ -0,0 +1,219 @@ +"""Example stages keep Faiss and RaBitQ in different Python processes.""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex + +EXAMPLES = Path(__file__).resolve().parents[2] / "sample" / "python" +HAS_FAISS = importlib.util.find_spec("faiss") is not None + +# Fail on an accidental import, even when a runtime combination happens to coexist. +BOOTSTRAP = """ +import runpy +import sys +from pathlib import Path + +blocked, script, *args = sys.argv[1:] +class BlockImport: + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] == blocked: + raise AssertionError(f'{script} must not import {blocked}') +sys.meta_path.insert(0, BlockImport()) +sys.path.insert(0, str(Path(script).parent)) +sys.argv = [script, *args] +runpy.run_path(script, run_name='__main__') +""" + + +def run_example(name, *args, blocked="faiss"): + env = os.environ.copy() + for key in ( + "DYLD_LIBRARY_PATH", + "DYLD_INSERT_LIBRARIES", + "LD_LIBRARY_PATH", + "LD_PRELOAD", + "KMP_DUPLICATE_LIB_OK", + ): + env.pop(key, None) + result = subprocess.run( + [ + sys.executable, + "-c", + BOOTSTRAP, + blocked, + str(EXAMPLES / name), + *map(str, args), + ], + env=env, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + return result.stdout + + +def write_vecs(path, values): + words = np.empty((len(values), values.shape[1] + 1), dtype=np.int32) + words[:, 0] = values.shape[1] + words[:, 1:] = values.view(np.int32) + words.tofile(path) + + +@pytest.fixture +def example_data(tmp_path): + data = np.random.default_rng(641).standard_normal((128, 65)).astype(np.float32) + data /= np.linalg.norm(data, axis=1, keepdims=True) + path = tmp_path / "base.fvecs" + write_vecs(path, data) + return path, data + + +def check_saved_index_and_query(tmp_path, kind, metric, data): + index_path = tmp_path / "test.index" + index_type = {"ivf": IvfIndex, "hnsw": HnswIndex, "symqg": SymqgIndex}[kind] + index = index_type.load(str(index_path)) + queries = data[:4].copy() + search_args = {"nprobe": 2} if kind == "ivf" else {"ef": len(data)} + ids, distances = index.search(queries, k=5, num_threads=2, **search_args) + assert index.metric == metric + np.testing.assert_array_equal(ids[:, 0], np.arange(len(queries))) + assert np.isfinite(distances).all() + restored_path = tmp_path / "restored.index" + index.save(str(restored_path)) + actual = index_type.load(str(restored_path)).search( + queries, k=5, num_threads=2, **search_args + ) + np.testing.assert_array_equal(actual[0], ids) + np.testing.assert_array_equal(actual[1], distances) + + query_path = tmp_path / "query.fvecs" + truth_path = tmp_path / "truth.ivecs" + exact = ( + np.sum((queries[:, None] - data) ** 2, axis=2) + if metric == "l2" + else 1 - queries @ data.T + ) + write_vecs(query_path, queries) + write_vecs(truth_path, np.argsort(exact, axis=1)[:, :5].astype(np.int32)) + name = "symqg_querying.py" if kind == "symqg" else f"{kind}_rabitq_querying.py" + extra = ["--metric", metric] if kind == "symqg" else [] + output = run_example( + name, + index_path, + query_path, + truth_path, + "--topk", + 5, + "--test-rounds", + 1, + "--num-threads", + 2, + *extra, + ) + assert "Recall" in output + + +@pytest.mark.skipif( + not HAS_FAISS, reason="Faiss is required only for clustering examples" +) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +def test_separate_clustering_and_indexing(tmp_path, example_data, metric, kind): + data_path, data = example_data + clusters = tmp_path / "clusters.npz" + run_example( + "faiss_clustering.py", + data_path, + clusters, + "--num-clusters", + 2, + "--metric", + metric, + "--num-threads", + 2, + blocked="rabitqlib", + ) + with np.load(clusters, allow_pickle=False) as saved: + assert saved["centroids"].shape == (2, data.shape[1]) + assert saved["centroids"].dtype == np.float32 + assert saved["cluster_ids"].shape == (len(data),) + assert saved["cluster_ids"].dtype == np.uint32 + assert (saved["cluster_ids"] < 2).all() + assert saved["metric"].item() == metric + run_example( + f"{kind}_rabitq_indexing.py", + data_path, + tmp_path / "test.index", + "--clusters", + clusters, + "--metric", + metric, + "--total-bits", + 5, + "--num-threads", + 2, + ) + check_saved_index_and_query(tmp_path, kind, metric, data) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_symqg_examples_do_not_import_faiss(tmp_path, example_data, metric): + data_path, data = example_data + run_example( + "symqg_indexing.py", + data_path, + tmp_path / "test.index", + "--metric", + metric, + "--num-threads", + 2, + "--ef-construction", + len(data), + ) + check_saved_index_and_query(tmp_path, "symqg", metric, data) + + +@pytest.mark.parametrize( + "field,value,message", + [ + ("metric", "ip", "clustering metric does not match index metric"), + ("metric", ["l2"], "clustering metric does not match index metric"), + ("centroids", np.zeros((2, 3), np.float32), "centroids must be float32"), + ("centroids", np.zeros((2, 4), np.float64), "centroids must be float32"), + ( + "centroids", + np.full((2, 4), np.nan, np.float32), + "centroids must contain only finite", + ), + ("cluster_ids", np.zeros(2, np.uint32), "one ID per data point"), + ("cluster_ids", np.zeros(3, np.float32), "cluster_ids must be uint32"), + ("cluster_ids", np.array([0, 1, 2], np.uint32), "out-of-range cluster ID"), + ("metric", None, "must contain centroids, cluster_ids and metric"), + ], +) +def test_reject_mismatched_clustering_file(tmp_path, field, value, message): + spec = importlib.util.spec_from_file_location( + "example_utils", EXAMPLES / "utils.py" + ) + utils = importlib.util.module_from_spec(spec) + spec.loader.exec_module(utils) + saved = { + "centroids": np.zeros((2, 4), np.float32), + "cluster_ids": np.array([0, 1, 0], np.uint32), + "metric": "l2", + } + if value is None: + del saved[field] + else: + saved[field] = value + path = tmp_path / "clusters.npz" + np.savez(path, **saved) + with pytest.raises(ValueError, match=message): + utils.load_clusters(path, (3, 4), "l2") diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 2f189b2a..0f515470 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -265,10 +265,14 @@ def test_raw_reranking(dim, metric, high_accuracy, fast_quantization, tmp_path): fast_quantization=fast_quantization, num_threads=2, ) + # Use a wider oracle: float32 BLAS reduction error near cancellation can + # exceed the tolerance even when the returned float32 distance is accurate. + reference_data = data.astype(np.float64) + reference_queries = queries.astype(np.float64) expected = ( - np.sum((queries[:, None] - data) ** 2, axis=2) + np.sum((reference_queries[:, None] - reference_data) ** 2, axis=2) if metric == "l2" - else 1 - queries @ data.T + else 1 - reference_queries @ reference_data.T ) original = data.copy() data[:] = 1000 # the index must own the raw data diff --git a/tests/python/test_portable_indexes.py b/tests/python/test_portable_indexes.py new file mode 100644 index 00000000..f378dcea --- /dev/null +++ b/tests/python/test_portable_indexes.py @@ -0,0 +1,90 @@ +"""Portable kernels preserve metric, padded-domain, and persistence contracts.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("dim", [64, 65]) +@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): + 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. + data /= np.linalg.norm(data, axis=1, keepdims=True) + if kind == "ivf": + index = IvfIndex(dim, len(data), 1, nbits=bits, metric=metric) + search_args = {"nprobe": 1} + elif kind == "hnsw": + index = HnswIndex( + dim, + len(data), + M=16, + ef_construction=96, + nbits=bits, + metric=metric, + random_seed=42, + ) + search_args = {"ef": 96} + else: + index = SymqgIndex(dim, max_degree=32, quantization_bits=bits, metric=metric) + search_args = {"ef": 96} + if kind == "qg": + index.build(data, ef_construction=96) + else: + index.build( + data, data.mean(axis=0, keepdims=True), np.zeros(len(data), np.uint32) + ) + queries = data[:5].copy() + ids, distances = index.search(queries, k=5, **search_args) + assert ids.shape == distances.shape == (5, 5) + assert np.isfinite(distances).all() + assert ((ids >= 0) & (ids < len(data))).all() + np.testing.assert_array_equal(ids[:, 0], np.arange(5)) + assert (np.diff(distances, axis=1) >= -1e-6).all() + if (kind == "ivf" and bits == 32) or (kind == "qg" and bits == 0): + selected = data[ids] + expected = ( + np.sum((queries[:, None] - selected) ** 2, axis=2) + if metric == "l2" + else 1 - np.einsum("qd,qkd->qk", queries, selected) + ) + np.testing.assert_allclose(distances, expected, atol=2e-5) + path = tmp_path / "portable.index" + index.save(str(path)) + loaded = type(index).load(str(path)) + 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)], +) +def test_rotation_padded_domain_round_trip(tmp_path, metric, dim, kind, bits): + test_portable_index_round_trip(tmp_path, metric, dim, 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) diff --git a/tests/python/test_thread_counts.py b/tests/python/test_thread_counts.py new file mode 100644 index 00000000..c2b5b0c1 --- /dev/null +++ b/tests/python/test_thread_counts.py @@ -0,0 +1,54 @@ +"""Thread requests are safe and preserve search results across all indexes.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex + +MAX_THREADS = int(np.iinfo(np.uintp).max) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("build_threads", [0, 1, MAX_THREADS]) +def test_build_and_search_thread_counts(kind, metric, build_threads): + rng = np.random.default_rng(431) + data = rng.standard_normal((96, 64)).astype(np.float32) + data /= np.linalg.norm(data, axis=1, keepdims=True) + if kind == "ivf": + index = IvfIndex(64, len(data), 1, nbits=8, metric=metric) + search_args = {"nprobe": 1} + elif kind == "hnsw": + index = HnswIndex( + 64, len(data), M=16, ef_construction=96, nbits=8, metric=metric + ) + search_args = {"ef": 96} + else: + index = SymqgIndex(64, max_degree=32, metric=metric) + search_args = {"ef": 96} + if kind == "qg": + index.build(data, ef_construction=96, num_threads=build_threads) + else: + index.build( + data, + data.mean(axis=0, keepdims=True), + np.zeros(len(data), dtype=np.uint32), + num_threads=build_threads, + ) + + # Compare searches on the same graph; parallel construction may change its topology. + expected_ids, expected_distances = index.search( + data, k=5, num_threads=1, **search_args + ) + np.testing.assert_array_equal(expected_ids[:, 0], np.arange(len(data))) + assert np.isfinite(expected_distances).all() + for threads in (0, MAX_THREADS): + ids, distances = index.search(data, k=5, num_threads=threads, **search_args) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + # A small or empty query batch must remain valid with an oversized request. + for count in (0, 1): + ids, distances = index.search( + data[:count], k=5, num_threads=threads, **search_args + ) + np.testing.assert_array_equal(ids, expected_ids[:count]) + np.testing.assert_array_equal(distances, expected_distances[:count]) diff --git a/tests/python/test_wheel_macos.py b/tests/python/test_wheel_macos.py new file mode 100644 index 00000000..379dce1d --- /dev/null +++ b/tests/python/test_wheel_macos.py @@ -0,0 +1,39 @@ +"""Check the installed, repaired macOS wheel, including its OpenMP runtime.""" + +import os +import platform +import subprocess +from pathlib import Path + +import pytest + + +@pytest.mark.skipif( + os.environ.get("RABITQ_TEST_WHEEL") != "1" or platform.system() != "Darwin", + reason="Run with RABITQ_TEST_WHEEL=1 after installing a repaired macOS wheel", +) +def test_native_arm64_wheel_bundles_openmp(): + import rabitqlib._rabitqlib as extension + + assert platform.machine() == "arm64" + path = Path(extension.__file__).resolve() + assert "site-packages" in path.parts + assert ( + subprocess.check_output(["lipo", "-archs", str(path)], text=True).strip() + == "arm64" + ) + assert (path.parent / "licenses" / "LLVM-OpenMP.txt").is_file() + dependencies = subprocess.check_output(["otool", "-L", str(path)], text=True) + omp = [ + line.strip().split(" ")[0] + for line in dependencies.splitlines()[1:] + if "libomp" in line + ] + assert len(omp) == 1, dependencies + assert omp[0].startswith("@loader_path/"), dependencies + runtime = path.parent / omp[0].removeprefix("@loader_path/") + assert runtime.is_file() + assert ( + subprocess.check_output(["lipo", "-archs", str(runtime)], text=True).strip() + == "arm64" + ) diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 6fbb71d7..1375c8c6 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -77,14 +77,24 @@ TEST(FastScanLutTest, MatchesOrderedSubsetSumsExactly) { TEST(FastScanLutTest, EverySupportedBackendMatchesOrderedSubsetSumsExactly) { check_lut_subset_sums(simd::pack_lut_generic); +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + SCOPED_TRACE("NEON"); + check_lut_subset_sums(simd::pack_lut_neon); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); check_lut_subset_sums(simd::pack_lut_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { SCOPED_TRACE("AVX512"); check_lut_subset_sums(simd::pack_lut_avx512); } +#endif } TEST(FastScanPackingTest, MatchesBitReferenceIncludingTailsAndUnalignedBuffers) { @@ -158,6 +168,13 @@ TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) EXPECT_EQ(actual, expected); EXPECT_EQ(actual_hacc, expected_hacc); }; + check_backend(accumulate, transfer_lut_hacc, accumulate_hacc); + check_backend( + simd::accumulate_generic, + simd::transfer_lut_hacc_generic, + simd::accumulate_hacc_generic + ); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); check_backend( @@ -166,6 +183,8 @@ TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) simd::accumulate_hacc_avx2 ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { SCOPED_TRACE("AVX512"); check_backend( @@ -174,6 +193,7 @@ TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) simd::accumulate_hacc_avx512 ); } +#endif } } } @@ -208,14 +228,19 @@ TEST(FastScanPackingTest, AccumulatesLargeDimensionsWithoutLaneOverflow) { EXPECT_EQ(actual, expected); }; check_backend(accumulate); + check_backend(simd::accumulate_generic); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); check_backend(simd::accumulate_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { SCOPED_TRACE("AVX512"); check_backend(simd::accumulate_avx512); } +#endif } } @@ -253,14 +278,20 @@ TEST(FastScanHighAccuracyTest, AccumulatesLargeDimensionsWithoutLaneOverflow) { accumulate_fn(packed.data(), packed_lut.data(), actual.data(), dim); EXPECT_EQ(actual, expected); }; + check_backend(transfer_lut_hacc, accumulate_hacc); + check_backend(simd::transfer_lut_hacc_generic, simd::accumulate_hacc_generic); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); check_backend(simd::transfer_lut_hacc_avx2, simd::accumulate_hacc_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { SCOPED_TRACE("AVX512"); check_backend(simd::transfer_lut_hacc_avx512, simd::accumulate_hacc_avx512); } +#endif } } } @@ -283,14 +314,20 @@ TEST(FastScanHighAccuracyTest, RejectsAccumulationsExceedingInt32Range) { return value == 42; })); }; + check_backend(transfer_lut_hacc, accumulate_hacc); + check_backend(simd::transfer_lut_hacc_generic, simd::accumulate_hacc_generic); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); check_backend(simd::transfer_lut_hacc_avx2, simd::accumulate_hacc_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { SCOPED_TRACE("AVX512"); check_backend(simd::transfer_lut_hacc_avx512, simd::accumulate_hacc_avx512); } +#endif } TEST(FastScanHighAccuracyTest, RejectsDimensionsThatCannotFillASimdBlock) { @@ -318,6 +355,7 @@ TEST(FastScanHighAccuracyTest, RejectsDimensionsThatCannotFillASimdBlock) { } } +#if defined(__x86_64__) || defined(_M_X64) TEST(FastScanHighAccuracyTest, Avx512TransferAcceptsUnalignedOutput) { if (!cpu::has_avx512_core()) { GTEST_SKIP() << "AVX512 is not supported on this CPU"; @@ -352,19 +390,22 @@ TEST(FastScanHighAccuracyTest, Avx512TransferAcceptsUnalignedOutput) { EXPECT_TRUE(std::equal(expected.begin(), expected.end(), storage.begin() + 1)); } +#endif + TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { - if (!cpu::has_avx2() && !cpu::has_avx512_core()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA or AVX512"; - } using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); std::vector backends{ rabitqlib::simd::split_batch_estdist, rabitqlib::simd::split_batch_estdist_generic}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx512); } +#endif constexpr size_t kDim = 512; std::array data; @@ -404,18 +445,19 @@ TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { } TEST(BatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches) { - if (!cpu::has_avx2() && !cpu::has_avx512_core()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA or AVX512"; - } using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); std::vector backends{ rabitqlib::simd::split_batch_estdist, rabitqlib::simd::split_batch_estdist_generic}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx512); } +#endif for (size_t dim : {16U, 64U, 128U, 1024U, 1040U, 4096U}) { for (bool hacc : {false, true}) { for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -499,10 +541,6 @@ TEST(BatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches } TEST(FastScanHighAccuracyTest, AccumulatesAcrossChunks) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "AVX2 is not supported on this CPU"; - } - constexpr size_t kDim = 4096; std::vector query(kDim, 1.0F); SplitBatchQuery q_obj(query.data(), kDim, 3, METRIC_L2, true); diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index f52bd0ab..90155ae5 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -223,9 +223,6 @@ static_assert( ); TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } for (size_t dim : {64U, 1024U, 1088U}) { for (auto metric : {METRIC_L2, METRIC_IP}) { SCOPED_TRACE(::testing::Message() << dim << "/" << metric); @@ -253,10 +250,14 @@ TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { std::array estimates{}; std::vector backends{ simd::qg_batch_estdist_generic, simd::qg_batch_estdist}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) backends.push_back(simd::qg_batch_estdist_avx2); +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) backends.push_back(simd::qg_batch_estdist_avx512); +#endif for (auto backend : backends) { backend(batch.data(), q_obj, dim, estimates.data()); for (size_t i = 0; i < fastscan::kBatchSize; ++i) { @@ -291,9 +292,11 @@ TEST(QGEstimatorTest, CandidateMaskBackendsMatchScalarOrdering) { constexpr float kThreshold = 2.0F; const uint32_t expected = simd::qg_candidate_mask_generic(distances.data(), kThreshold); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_EQ(simd::qg_candidate_mask_avx2(distances.data(), kThreshold), expected); } +#endif } TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { @@ -341,9 +344,6 @@ TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { } TEST(QGConstructionTest, EncodedSeedSearchMatchesRetainedScores) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kCount = 97, kDim = 65, kDegree = 32; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -387,9 +387,6 @@ TEST(QGConstructionTest, EncodedSeedSearchMatchesRetainedScores) { } TEST(QGConstructionTest, DefaultsToPipnnInitialization) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kCount = 97, kDim = 65, kDegree = 32; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -425,9 +422,6 @@ TEST(QGConstructionTest, DefaultsToPipnnInitialization) { } TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } std::vector data(65 * 65, 0.5F); QuantizedGraph graph(65, 65, 32); QGBuilder builder(graph, 64, data.data(), 1, QGInitialization::Random); @@ -456,9 +450,6 @@ TEST(QGConstructionTest, RejectsNullDataForEveryInitializationMode) { } TEST(QGConstructionTest, UsesExplicitThreadCountsWithoutChangingCallerState) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kCount = 33, kDim = 64; std::vector data(kCount * kDim, 0.25F); const int caller_threads = omp_get_max_threads(); @@ -476,9 +467,6 @@ TEST(QGConstructionTest, UsesExplicitThreadCountsWithoutChangingCallerState) { } TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kCount = 33, kDim = 64; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -504,9 +492,6 @@ TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { } TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kCount = 97, kDim = 65, kDegree = 64; constexpr std::array kDegrees{0, 1, 31, 32, 33, 64}; for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -780,9 +765,6 @@ TEST(QuantizedGraphLifecycleTest, RejectsSearchAndSaveBeforeBuild) { } TEST(QuantizedGraphLifecycleTest, BuilderDoesNotPublishPartialGraph) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kNumPoints = 33; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; @@ -920,9 +902,6 @@ TEST(QGQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { } TEST(QGSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FastScan requires AVX2/FMA"; - } constexpr size_t kNumPoints = 33, kDim = 64, kDegree = 32; std::vector data(kNumPoints * kDim); for (size_t i = 0; i < data.size(); ++i) { diff --git a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp index 072af909..c9a4196b 100644 --- a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp +++ b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp @@ -129,8 +129,12 @@ class RabitqRescaleBackendTest : public ::testing::TestWithParam { } RescaleSearch search() const { +#if defined(__x86_64__) || defined(_M_X64) return GetParam() ? simd::best_rescale_factor_avx512 : simd::best_rescale_factor_avx2; +#else + return simd::best_rescale_factor_generic; +#endif } }; diff --git a/tests/unit/rabitqlib/utils/cpu_features_test.cpp b/tests/unit/rabitqlib/utils/cpu_features_test.cpp index 294160c1..cc755699 100644 --- a/tests/unit/rabitqlib/utils/cpu_features_test.cpp +++ b/tests/unit/rabitqlib/utils/cpu_features_test.cpp @@ -125,3 +125,14 @@ TEST(CpuFeatures, Avx512SelectionRequiresEveryCompilerEnabledFeature) { } } // namespace + +TEST(CpuFeatures, NativeArm64UsesNeonWithoutX86Features) { +#if defined(__aarch64__) + EXPECT_TRUE(rabitqlib::cpu::has_neon()); + EXPECT_FALSE(rabitqlib::cpu::has_avx2()); + EXPECT_FALSE(rabitqlib::cpu::has_avx512_core()); + EXPECT_FALSE(rabitqlib::cpu::has_avx512_popcnt()); +#else + EXPECT_FALSE(rabitqlib::cpu::has_neon()); +#endif +} diff --git a/tests/unit/rabitqlib/utils/excode_ip_test.cpp b/tests/unit/rabitqlib/utils/excode_ip_test.cpp index bf4da7f4..e96641c6 100644 --- a/tests/unit/rabitqlib/utils/excode_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/excode_ip_test.cpp @@ -14,10 +14,8 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { using namespace rabitqlib; - if (!cpu::has_avx2()) { - GTEST_SKIP() << "Packed-code tests require AVX2/FMA"; - } +#if defined(__x86_64__) || defined(_M_X64) const std::array avx2_functions{ simd::excode_ipimpl::ip16_fxu1_avx2, simd::excode_ipimpl::ip64_fxu2_avx2, @@ -38,14 +36,27 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { simd::excode_ipimpl::ip64_fxu7_avx512, simd::excode_ipimpl::ip16_fxu8_avx512, }; +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + const std::array neon_functions{ + simd::excode_ipimpl::ip16_fxu1_neon, + simd::excode_ipimpl::ip64_fxu2_neon, + simd::excode_ipimpl::ip64_fxu3_neon, + simd::excode_ipimpl::ip16_fxu4_neon, + simd::excode_ipimpl::ip64_fxu5_neon, + simd::excode_ipimpl::ip64_fxu6_neon, + simd::excode_ipimpl::ip64_fxu7_neon, + simd::excode_ipimpl::ip16_fxu8_neon, + }; +#endif using PackFunction = void (*)(const uint8_t*, uint8_t*, size_t); const std::array pack_functions{ - simd::packing_2bit_excode_avx2, - simd::packing_3bit_excode_avx2, - simd::packing_4bit_excode_avx2, - simd::packing_5bit_excode_avx2, - simd::packing_6bit_excode_avx2, - simd::packing_7bit_excode_avx2, + simd::packing_2bit_excode, + simd::packing_3bit_excode, + simd::packing_4bit_excode, + simd::packing_5bit_excode, + simd::packing_6bit_excode, + simd::packing_7bit_excode, }; for (size_t bits = 1; bits <= 8; ++bits) { @@ -96,14 +107,26 @@ TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { // product itself near zero. Different reduction trees may round // differently. const double tolerance = 2e-6 * std::max(1.0, sum_abs); - EXPECT_NEAR( - avx2_functions[bits - 1](query, compact, dim), expected, tolerance - ); +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + EXPECT_NEAR( + avx2_functions[bits - 1](query, compact, dim), expected, tolerance + ); + } if (cpu::has_avx512_core()) { EXPECT_NEAR( avx512_functions[bits - 1](query, compact, dim), expected, tolerance ); } +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + EXPECT_NEAR( + neon_functions[bits - 1](query, compact, dim), expected, tolerance + ); + EXPECT_EQ(select_excode_ipfunc(bits), neon_functions[bits - 1]); + } +#endif EXPECT_NEAR( select_excode_ipfunc(bits)(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 7db1dd30..1b2fcc0d 100644 --- a/tests/unit/rabitqlib/utils/float_distance_test.cpp +++ b/tests/unit/rabitqlib/utils/float_distance_test.cpp @@ -1,6 +1,5 @@ #include -#include #include #include #include @@ -39,6 +38,7 @@ std::vector backends() { simd::dot_product_generic, simd::dot_product_dis_generic, simd::l2norm_sqr_generic}}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { result.push_back( {"avx2", @@ -48,6 +48,8 @@ std::vector backends() { simd::l2norm_sqr_avx2} ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { result.push_back( {"avx512", @@ -57,88 +59,22 @@ std::vector backends() { simd::l2norm_sqr_avx512} ); } - return result; -} - -void check_reference(const float* a, const float* b, size_t dim) { - long double l2 = 0, dot = 0, norm = 0, absolute_products = 0; - for (size_t i = 0; i < dim; ++i) { - const long double x = a[i], y = b[i]; - l2 += (x - y) * (x - y); - dot += x * y; - norm += x * x; - absolute_products += std::abs(x * y); - } - // Longer float reductions accumulate more rounding error. - const long double tolerance = 2e-6L * std::max(1.0L, dim / 4096.0L); - for (const auto& backend : backends()) { - SCOPED_TRACE(backend.name); - EXPECT_NEAR(backend.l2(a, b, dim), l2, tolerance * std::max(1.0L, l2)); - EXPECT_NEAR( - backend.dot(a, b, dim), dot, tolerance * std::max(1.0L, absolute_products) +#endif +#if defined(__aarch64__) + if (cpu::has_neon()) { + result.push_back( + {"neon", + simd::euclidean_sqr_neon, + simd::dot_product_neon, + simd::dot_product_dis_neon, + simd::l2norm_sqr_neon} ); - EXPECT_NEAR(backend.ip(a, b, dim), 1 - dot, tolerance * (1 + absolute_products)); - EXPECT_NEAR(backend.norm(a, dim), norm, tolerance * std::max(1.0L, norm)); } +#endif + return result; } -} // namespace -TEST(FloatDistance, BackendsMatchReferenceAcrossDimensionsAndAlignments) { - std::vector dimensions; - for (size_t dim = 0; dim <= 80; ++dim) { - dimensions.push_back(dim); - } - for (size_t dim : - {127, - 128, - 129, - 255, - 256, - 257, - 419, - 420, - 421, - 959, - 960, - 961, - 1023, - 1024, - 1025, - 4096, - 65536}) { - dimensions.push_back(dim); - } - for (size_t dim : dimensions) { - SCOPED_TRACE(dim); - for (size_t offset : {0, 1, 7, 15}) { - std::vector a(std::max(size_t{1}, dim + offset)); - std::vector b(std::max(size_t{1}, dim + offset)); - for (int pattern = 0; pattern < 6; ++pattern) { - SCOPED_TRACE(pattern); - for (size_t i = 0; i < dim; ++i) { - float x = static_cast(static_cast(i % 23) - 11) / 7; - float y = static_cast(static_cast(i % 17) - 8) / 9; - if (pattern == 1) { - y = x; - } else if (pattern == 2) { - y = -x; - } else if (pattern == 3) { - x = (i % 4 < 2) ? 4096.0F : -4096.0F; - y = 1; - } else if (pattern == 4) { - x = std::ldexp(x, 30); - y = std::ldexp(y, 30); - } else if (pattern == 5) { - y = std::nextafter(x, 2.0F); - } - a[i + offset] = x; - b[i + offset] = y; - } - check_reference(a.data() + offset, b.data() + offset, dim); - } - } - } -} +} // namespace TEST(FloatDistance, EmptyAndAliasedInputsPreserveDistanceConventions) { const std::array a{1, -2, 3, -4, 5}; @@ -155,6 +91,41 @@ TEST(FloatDistance, EmptyAndAliasedInputsPreserveDistanceConventions) { } } +TEST(FloatDistance, SubtractionUsesFloatPrecisionInVectorBlocksAndTails) { + // In float, 1 - 2^-25 rounds to 1 before squaring. Widening the subtraction + // to double would instead return the float immediately below 1. + for (size_t dim : {1U, 4U, 16U, 17U, 33U}) { + for (size_t at : {size_t{0}, dim - 1}) { + std::vector a(dim, 0), b(dim, 0); + a[at] = 1.0F; + b[at] = std::ldexp(1.0F, -25); + for (const auto& backend : backends()) { + SCOPED_TRACE( + ::testing::Message() << backend.name << " dim=" << dim << " at=" << at + ); + EXPECT_EQ(backend.l2(a.data(), b.data(), dim), 1.0F); + } + } + } +} + +TEST(FloatDistance, AccumulationRetainsFloatPrecision) { + // Each small square is below half an ulp at 1. Place them in the same + // accumulation lane for all backends. Double accumulation would preserve + // their combined 2^-23 contribution and return the float above 1. + std::array a{}, zero{}; + a[0] = 1.0F; + for (size_t i = 64; i < a.size(); i += 64) + a[i] = std::ldexp(1.0F, -13); + for (const auto& backend : backends()) { + SCOPED_TRACE(backend.name); + EXPECT_EQ(backend.l2(a.data(), zero.data(), a.size()), 1.0F); + EXPECT_EQ(backend.dot(a.data(), a.data(), a.size()), 1.0F); + EXPECT_EQ(backend.ip(a.data(), a.data(), a.size()), 0.0F); + EXPECT_EQ(backend.norm(a.data(), a.size()), 1.0F); + } +} + TEST(FloatDistance, DoubleTemplatesRetainDoublePrecision) { const std::array a{1.0000000001, 2.0000000001, 3.0000000001}; const std::array b{1, 2, 3}; @@ -184,10 +155,20 @@ TEST(FloatDistance, TailsDoNotReadPastGuardPage) { SCOPED_TRACE(dim); ASSERT_LE(dim * sizeof(float), page); auto* data = reinterpret_cast(end) - dim; + float expected = 0; for (size_t i = 0; i < dim; ++i) { data[i] = static_cast(i % 7); + expected += data[i] * data[i]; + } + // Small integer inputs make these sums exactly representable in float. + // Keep this a bounds-safety check, independent of precision tolerances. + for (const auto& backend : backends()) { + SCOPED_TRACE(backend.name); + EXPECT_EQ(backend.l2(data, data, dim), 0.0F); + EXPECT_EQ(backend.dot(data, data, dim), expected); + EXPECT_EQ(backend.ip(data, data, dim), 1.0F - expected); + EXPECT_EQ(backend.norm(data, dim), expected); } - check_reference(data, data, dim); } } #endif diff --git a/tests/unit/rabitqlib/utils/mask_ip_test.cpp b/tests/unit/rabitqlib/utils/mask_ip_test.cpp index 3df3ab9b..54476d95 100644 --- a/tests/unit/rabitqlib/utils/mask_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/mask_ip_test.cpp @@ -13,16 +13,26 @@ TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { using namespace rabitqlib; - if (!cpu::has_avx2()) { - GTEST_SKIP() << "Binary dot product tests require AVX2/FMA"; - } using Function = float (*)(const float*, const uint8_t*, size_t); std::vector functions{ - static_cast(simd::mask_ip_x0_q_avx2), + static_cast(simd::mask_ip_x0_q_generic), static_cast(mask_ip_x0_q)}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + functions.push_back(static_cast(simd::mask_ip_x0_q_avx2)); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { functions.push_back(static_cast(simd::mask_ip_x0_q_avx512)); } +#endif + +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + functions.push_back(static_cast(simd::mask_ip_x0_q_neon)); + } +#endif for (size_t dim : {0, 64, 128, 192, 256, 448, 576, 1024, 4096}) { SCOPED_TRACE(dim); @@ -65,11 +75,8 @@ TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { } } -TEST(MaskIpX0Q, Avx2PreservesEveryStoredBitPosition) { +TEST(MaskIpX0Q, DispatchPreservesEveryStoredBitPosition) { using namespace rabitqlib; - if (!cpu::has_avx2()) { - GTEST_SKIP() << "Binary dot product tests require AVX2/FMA"; - } constexpr size_t dim = 192; std::vector query(dim); for (size_t i = 0; i < dim; ++i) { @@ -79,6 +86,6 @@ TEST(MaskIpX0Q, Avx2PreservesEveryStoredBitPosition) { SCOPED_TRACE(bit); std::vector words(dim / 64, 0); words[bit / 64] = uint64_t{1} << (63 - bit % 64); - EXPECT_EQ(simd::mask_ip_x0_q_avx2(query.data(), words.data(), dim), query[bit]); + EXPECT_EQ(mask_ip_x0_q(query.data(), words.data(), dim), query[bit]); } } diff --git a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp index c26f0ec3..6330dcba 100644 --- a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp +++ b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp @@ -25,6 +25,7 @@ TEST(MatrixDispatchTest, AllBackendsMatchScalarForRectangularUnalignedInputs) { matrix_product_transposed_generic, row_norms_generic, pairwise_distances_lower_generic}}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { backends.push_back( {matrix_product_avx2, @@ -33,6 +34,8 @@ TEST(MatrixDispatchTest, AllBackendsMatchScalarForRectangularUnalignedInputs) { pairwise_distances_lower_avx2} ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { backends.push_back( {matrix_product_avx512, @@ -41,6 +44,7 @@ TEST(MatrixDispatchTest, AllBackendsMatchScalarForRectangularUnalignedInputs) { pairwise_distances_lower_avx512} ); } +#endif for (size_t dim : {1U, 7U, 16U, 17U, 65U, 128U}) { for (size_t rows : {1U, 33U}) { constexpr size_t kCols = 19; diff --git a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp new file mode 100644 index 00000000..ce72ff4a --- /dev/null +++ b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp @@ -0,0 +1,126 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +#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/utils/cpu_features.hpp" + +#if defined(__aarch64__) || defined(_M_ARM64) +namespace rabitqlib { +TEST(NeonHnsw, WarmupHandlesUnalignedCodesTailsAndWideCounts) { + if (!cpu::has_neon()) { + GTEST_SKIP(); + } + std::mt19937 random(764); + for (size_t dim : {0U, 64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U, 65536U}) { + for (size_t bits = 0; bits <= 8; ++bits) { + for (bool all_ones : {false, true}) { + SCOPED_TRACE(dim); + SCOPED_TRACE(bits); + SCOPED_TRACE(all_ones); + std::vector query(dim), storage(dim / 8 + 1); + std::vector words(dim / 64), transposed(dim / 64 * bits); + uint64_t ip = 0, count = 0; + for (size_t i = 0; i < dim; ++i) { + query[i] = static_cast( + (all_ones ? 255 : random()) & ((1U << bits) - 1) + ); + if (all_ones || random() % 2 != 0) { + words[i / 64] |= uint64_t{1} << (63 - i % 64); + ip += query[i]; + ++count; + } + } + if (dim != 0) { + std::memcpy(storage.data() + 1, words.data(), dim / 8); + } + simd::new_transpose_bin_512_generic( + query.data(), transposed.data(), dim, bits + ); + const float expected = + 0.25F * static_cast(ip) - 2.0F * static_cast(count); + EXPECT_EQ( + simd::warmup_ip_x0_q_512_neon( + storage.data() + 1, transposed.data(), 0.25F, -2.0F, dim, bits + ), + expected + ); + EXPECT_EQ( + simd::warmup_ip_x0_q_512_neon( + words.data(), transposed.data(), 0.25F, -2.0F, dim, bits + ), + expected + ); + } + } + } + EXPECT_THROW( + simd::warmup_ip_x0_q_512_neon( + static_cast(nullptr), nullptr, 1, 0, 0, 9 + ), + std::invalid_argument + ); +} + +TEST(NeonHnsw, SearchMatchesGenericOnSameGraphForBothMetricsAndAllWidths) { + if (!cpu::has_neon()) { + GTEST_SKIP(); + } + constexpr size_t kCount = 64; + for (size_t dim : {64U, 576U, 960U}) { + std::mt19937 random(42); + std::uniform_real_distribution distribution(-1, 1); + std::vector data(kCount * dim), centroid(dim, 0), query(dim); + for (auto& value : data) { + value = distribution(random); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits = 1; bits <= 9; ++bits) { + SCOPED_TRACE(dim); + SCOPED_TRACE(metric); + SCOPED_TRACE(bits); + std::vector clusters(kCount, 0); + hnsw::HierarchicalNSW index(kCount, dim, bits, 8, 64, 42, metric); + index.construct( + 1, centroid.data(), kCount, data.data(), clusters.data(), 1, true + ); + // The direct entry points accept rotated coordinates. Generate queries + // in that domain and pass the exact same values to both backends. + for (size_t q = 0; q < 3; ++q) { + for (auto& value : query) { + value = distribution(random); + } + for (size_t ef : {10U, 64U}) { + // Configure ef through the public search API. + index.search(query.data(), 1, 5, ef, 1); + auto expected = + hnsw::detail::search_knn_generic(index, query.data(), 5); + auto actual = hnsw::detail::search_knn_neon(index, query.data(), 5); + ASSERT_EQ(expected.size(), actual.size()); + while (!expected.empty()) { + EXPECT_EQ(actual.top().second, expected.top().second); + EXPECT_NEAR( + actual.top().first, + expected.top().first, + 2e-5F * std::max(1.0F, std::abs(expected.top().first)) + ); + actual.pop(); + expected.pop(); + } + } + } + } + } + } +} +} // namespace rabitqlib +#endif diff --git a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp new file mode 100644 index 00000000..7d781ce3 --- /dev/null +++ b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp @@ -0,0 +1,139 @@ +#include + +#include +#include +#include +#include +#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/utils/space.hpp" +#include "rabitqlib/utils/warmup_space.hpp" + +namespace rabitqlib { +TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { + using Pack = void (*)(const uint8_t*, uint8_t*, size_t); + const std::array pack{ + simd::packing_2bit_excode_generic, + simd::packing_3bit_excode_generic, + simd::packing_4bit_excode_generic, + simd::packing_5bit_excode_generic, + simd::packing_6bit_excode_generic, + simd::packing_7bit_excode_generic}; + const std::array ip{ + simd::excode_ipimpl::ip16_fxu1_generic, + simd::excode_ipimpl::ip64_fxu2_generic, + simd::excode_ipimpl::ip64_fxu3_generic, + simd::excode_ipimpl::ip16_fxu4_generic, + simd::excode_ipimpl::ip64_fxu5_generic, + simd::excode_ipimpl::ip64_fxu6_generic, + simd::excode_ipimpl::ip64_fxu7_generic, + simd::excode_ipimpl::ip16_fxu8_generic}; + for (size_t bits = 1; bits <= 8; ++bits) { + for (size_t dim : {64U, 128U, 192U, 576U}) { + SCOPED_TRACE(bits); + SCOPED_TRACE(dim); + std::vector raw(dim), dispatched(dim * bits / 8), + scalar(dispatched.size()); + std::vector query(dim); + double expected = 0; + for (size_t i = 0; i < dim; ++i) { + raw[i] = static_cast((i * 37 + i / 7) & ((1U << bits) - 1)); + query[i] = static_cast(static_cast(i % 17) - 8); + expected += query[i] * raw[i]; + } + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + raw.data(), dispatched.data(), dim, bits + ); + if (bits >= 2 && bits <= 7) { + pack[bits - 2](raw.data(), scalar.data(), dim); + EXPECT_EQ(scalar, dispatched); + } + EXPECT_EQ(ip[bits - 1](query.data(), dispatched.data(), dim), expected); + EXPECT_EQ( + select_excode_ipfunc(bits)(query.data(), dispatched.data(), dim), expected + ); + } + } +} + +TEST(PortableKernels, QueryTransposeAndWarmupMatchIntegerReference) { + for (size_t dim : {64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U}) { + for (size_t bits = 0; bits <= 8; ++bits) { + SCOPED_TRACE(dim); + SCOPED_TRACE(bits); + std::vector query(dim); + std::vector data(dim / 64), transposed(dim / 64 * bits), + scalar(transposed.size()); + uint64_t expected_ip = 0, expected_count = 0; + for (size_t i = 0; i < dim; ++i) { + query[i] = static_cast((i * 19 + 7) & ((1U << bits) - 1)); + if (i % 3 == 0 || i % 63 == 0) { + data[i / 64] |= uint64_t{1} << (63 - i % 64); + expected_ip += query[i]; + ++expected_count; + } + } + new_transpose_bin_512(query.data(), transposed.data(), dim, bits); + simd::new_transpose_bin_512_generic(query.data(), scalar.data(), dim, bits); + EXPECT_EQ(scalar, transposed); + size_t offset = 0; + for (size_t block = 0; block < dim; block += 512) { + const size_t chunks = std::min(size_t{512}, dim - block) / 64; + for (size_t b = 0; b < bits; ++b) { + for (size_t c = 0; c < chunks; ++c) { + for (size_t i = 0; i < 64; ++i) { + EXPECT_EQ( + (transposed[offset + b * chunks + c] >> (63 - i)) & 1U, + (query[block + c * 64 + i] >> b) & 1U + ); + } + } + } + offset += chunks * bits; + } + const float expected = 0.25F * expected_ip - 2.0F * expected_count; + EXPECT_EQ( + simd::warmup_ip_x0_q_512_generic( + data.data(), scalar.data(), 0.25F, -2.0F, dim, bits + ), + expected + ); + EXPECT_EQ( + warmup_ip_x0_q_512(data.data(), transposed.data(), 0.25F, -2.0F, dim, bits), + expected + ); + } + } + EXPECT_THROW( + simd::warmup_ip_x0_q_512_generic( + static_cast(nullptr), nullptr, 1, 0, 0, 9 + ), + std::invalid_argument + ); +} + +TEST(PortableKernels, Uint16TransposePreservesAllBitPlanes) { + constexpr size_t kDim = 192; + std::array query{}; + for (size_t i = 0; i < kDim; ++i) + query[i] = static_cast(i * 977); + for (size_t bits = 1; bits <= 16; ++bits) { + std::vector actual(kDim / 64 * bits), scalar(actual.size()); + new_transpose_bin(query.data(), actual.data(), kDim, bits); + simd::new_transpose_bin_generic(query.data(), scalar.data(), kDim, bits); + EXPECT_EQ(actual, scalar); + for (size_t i = 0; i < kDim; ++i) { + for (size_t b = 0; b < bits; ++b) { + EXPECT_EQ( + (actual[i / 64 * bits + b] >> (63 - i % 64)) & 1U, (query[i] >> b) & 1U + ); + } + } + } +} +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp new file mode 100644 index 00000000..e1f62bff --- /dev/null +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -0,0 +1,143 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" + +namespace rabitqlib { +namespace { +template +void check_quantization(void (*quantize)(T*, const float*, size_t, float, float)) { + const auto max = static_cast(std::numeric_limits::max()); + for (const auto& params : {std::array{0, 1}, {-3.25F, 0.125F}, {-7, 0.1F}}) { + const float lo = params[0], delta = params[1]; + std::vector input(1, 0); // Offset SIMD loads by one float. + input.push_back(lo); + input.push_back(lo + max * delta); + for (size_t code = 0; code < static_cast(max); ++code) { + const float half = lo + (static_cast(code) + 0.5F) * delta; + input.push_back(std::nextafter(half, -std::numeric_limits::infinity())); + input.push_back(half); + input.push_back(std::nextafter(half, std::numeric_limits::infinity())); + } + for (size_t dim : std::array{ + 0UL, + 1UL, + 3UL, + 4UL, + 7UL, + 8UL, + 9UL, + 15UL, + 16UL, + 17UL, + 31UL, + 32UL, + 33UL, + input.size() - 1}) { + SCOPED_TRACE(dim); + SCOPED_TRACE(delta); + std::vector actual(dim + 2, 99), expected(actual); + const float reciprocal = 1.0F / delta; + for (size_t i = 0; i < dim; ++i) { + expected[i + 1] = + static_cast(std::round((input[i + 1] - lo) * reciprocal)); + } + quantize(actual.data() + 1, input.data() + 1, dim, lo, delta); + EXPECT_EQ(actual, expected); + } + } +} + +template +void check_transpose( + void (*transpose)(const T*, uint64_t*, size_t, size_t), bool blocked512 +) { + for (size_t dim : {0U, 64U, 128U, 448U, 512U, 576U, 960U, 1024U, 1088U}) { + std::vector input(dim + 1); + for (size_t i = 0; i < dim; ++i) { + input[i + 1] = + static_cast((i * 977 + i / 7) ^ (1U << (i % (sizeof(T) * 8)))); + } + for (size_t bits = 0; bits <= sizeof(T) * 8; ++bits) { + SCOPED_TRACE(dim); + SCOPED_TRACE(bits); + std::vector actual(dim / 64 * bits + 2, UINT64_MAX); + std::vector expected(actual.size(), 0); + expected.front() = expected.back() = UINT64_MAX; + size_t offset = 1; + const size_t block_size = blocked512 ? 512 : 64; + for (size_t block = 0; block < dim; block += block_size) { + const size_t chunks = std::min(block_size, dim - block) / 64; + for (size_t b = 0; b < bits; ++b) { + for (size_t c = 0; c < chunks; ++c) { + for (size_t i = 0; i < 64; ++i) { + expected[offset + b * chunks + c] |= + uint64_t{(input[block + c * 64 + i + 1] >> b) & 1U} + << (63 - i); + } + } + } + offset += chunks * bits; + } + transpose(input.data() + 1, actual.data() + 1, dim, bits); + EXPECT_EQ(actual, expected); + } + } +} +} // namespace + +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); +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + check_quantization(simd::scalar_quantize_uint8_neon); + check_quantization(simd::scalar_quantize_uint16_neon); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + check_quantization(simd::scalar_quantize_uint8_avx2); + check_quantization(simd::scalar_quantize_uint16_avx2); + } + if (cpu::has_avx512_core()) { + check_quantization(simd::scalar_quantize_uint8_avx512); + check_quantization(simd::scalar_quantize_uint16_avx512); + } +#endif +} + +TEST(QueryPreparation, TranspositionMatchesCoordinateReferenceAndBounds) { + check_transpose(simd::new_transpose_bin_generic, false); + check_transpose(simd::new_transpose_bin_512_generic, true); + check_transpose(new_transpose_bin, false); + check_transpose(new_transpose_bin_512, true); +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + check_transpose(simd::new_transpose_bin_neon, false); + check_transpose(simd::new_transpose_bin_512_neon, true); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + check_transpose(simd::new_transpose_bin_avx2, false); + check_transpose(simd::new_transpose_bin_512_avx2, true); + } + if (cpu::has_avx512_core()) { + check_transpose(simd::new_transpose_bin_avx512, false); + check_transpose(simd::new_transpose_bin_512_avx512, true); + } +#endif +} +} // namespace rabitqlib diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 37fcbf43..576a74ed 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -65,9 +66,6 @@ TEST_F(RotatorTest, RejectsZeroDimension) { } TEST(FhtRotatorTest, LargeDimensionsPreserveNormAndSavedRotation) { - if (!cpu::has_avx2()) { - GTEST_SKIP() << "FHT rotation requires AVX2/FMA"; - } for (size_t dim : {4095U, 4096U, 16384U, 16385U, 65535U, 65536U}) { SCOPED_TRACE(dim); std::unique_ptr> rotator(choose_rotator(dim)); @@ -143,20 +141,26 @@ TEST(FlipSignTest, FlipWorks) { } TEST(FhtDispatchTest, BackendsMatchScalarButterfliesAndPadding) { - std::vector backends; + std::vector backends{ + simd::fht_rotate_generic, simd::fht_rotate}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { backends.push_back(simd::fht_rotate_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { backends.push_back(simd::fht_rotate_avx512); } - if (backends.empty()) { - GTEST_SKIP() << "FHT rotation requires AVX2/FMA or AVX512"; +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + backends.push_back(simd::fht_rotate_neon); } - backends.push_back(simd::fht_rotate); - for (size_t dim : - {64U, 65U, 127U, 128U, 129U, 192U, 256U, 512U, 1024U, 2048U, 2049U, - 4095U, 4096U, 4097U, 8192U, 8193U, 16384U, 16385U, 32768U, 65535U, 65536U}) { +#endif + for (size_t dim : {64U, 65U, 127U, 128U, 129U, 192U, 193U, 256U, + 512U, 960U, 1024U, 2048U, 2049U, 4095U, 4096U, 4097U, + 8192U, 8193U, 16384U, 16385U, 32768U, 65535U, 65536U}) { SCOPED_TRACE(dim); const size_t padded = (dim + 63) / 64 * 64; size_t trunc = 1; @@ -221,16 +225,23 @@ TEST(FhtDispatchTest, BackendsMatchScalarButterfliesAndPadding) { } TEST(FhtDispatchTest, PreservesZeroNormAndInnerProduct) { - std::vector backends; + std::vector backends{ + simd::fht_rotate_generic, simd::fht_rotate}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { backends.push_back(simd::fht_rotate_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { backends.push_back(simd::fht_rotate_avx512); } - if (backends.empty()) { - GTEST_SKIP() << "FHT rotation requires AVX2/FMA or AVX512"; +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + backends.push_back(simd::fht_rotate_neon); } +#endif for (size_t dim : {64U, 65U, 128U, 192U, 256U, 512U, 4097U, 65536U}) { SCOPED_TRACE(dim); const size_t padded = (dim + 63) / 64 * 64; @@ -278,3 +289,146 @@ TEST(MatrixRotatorTest, PreservesOverlappingInputAndOutput) { EXPECT_FLOAT_EQ(data[offset + 2], 42); } } + +TEST(FlipSignTest, BackendsPreserveBitsAndUnalignedBuffers) { + std::vector backends{ + simd::flip_sign_generic, simd::flip_sign}; +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + backends.push_back(simd::flip_sign_neon); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) + backends.push_back(simd::flip_sign_avx2); + if (cpu::has_avx512_core()) + backends.push_back(simd::flip_sign_avx512); +#endif + // Signed zeros, infinities, quiet NaN payloads, subnormals and finite values. + const std::array values{ + 0, 0x80000000U, 0x7f800000U, 0xff800000U, 0x7fc12345U, 0xffc54321U, 1, 0x3f812345U}; + for (unsigned mask = 0; mask < 256; ++mask) { + std::array flip{}; + std::fill(flip.begin(), flip.end(), static_cast(mask)); + for (auto backend : backends) { + std::array data{}; + data.front() = data.back() = 12345; + for (size_t i = 0; i < 64; ++i) { + std::memcpy(&data[i + 1], &values[i % 8], sizeof(float)); + } + backend(flip.data() + 1, data.data() + 1, 64); + for (size_t i = 0; i < 64; ++i) { + uint32_t actual; + std::memcpy(&actual, &data[i + 1], sizeof(actual)); + EXPECT_EQ(actual, values[i % 8] ^ (((mask >> (i % 8)) & 1U) << 31)); + } + EXPECT_EQ(data.front(), 12345); + EXPECT_EQ(data.back(), 12345); + } + } +} + +TEST(KacsWalkTest, BackendsMatchScalarAndRespectBufferBounds) { + std::vector backends{simd::kacs_walk}; +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) + backends.push_back(simd::kacs_walk_neon); +#endif +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) + backends.push_back(simd::kacs_walk_avx2); + if (cpu::has_avx512_core()) + backends.push_back(simd::kacs_walk_avx512); +#endif + for (size_t dim : {64U, 128U, 192U, 960U, 1024U, 65536U}) { + std::vector input(dim + 2, 12345); + for (size_t i = 0; i < dim; ++i) + input[i + 1] = std::sin(static_cast(i)); + auto expected = input; + simd::kacs_walk_generic(expected.data() + 1, dim); + for (auto backend : backends) { + auto actual = input; + backend(actual.data() + 1, dim); + EXPECT_EQ(actual, expected); + } + } +} + +#if defined(__aarch64__) || defined(_M_ARM64) +TEST(FhtNeonTest, MatchesSavedScalarRotationInPlaceAndWithExtraPadding) { + if (!cpu::has_neon()) + GTEST_SKIP(); + for (size_t dim : {64U, 65U, 127U, 128U, 193U, 960U, 1025U, 4097U, 65536U}) { + for (size_t extra : {0U, 64U}) { + const size_t padded = (dim + 63) / 64 * 64 + extra; + if (padded > 65536) + continue; + SCOPED_TRACE(dim); + SCOPED_TRACE(padded); + size_t trunc = 64; + while (trunc * 2 <= dim) + trunc *= 2; + std::vector flip(padded / 2); + for (size_t i = 0; i < flip.size(); ++i) + flip[i] = static_cast(i * 73 + i / 7); + std::vector input(dim), expected(padded); + for (size_t i = 0; i < dim; ++i) + input[i] = std::sin(static_cast(i) * 0.13F); + simd::fht_rotate_generic( + input.data(), + expected.data(), + dim, + padded, + trunc, + 1.0F / std::sqrt(static_cast(trunc)), + flip.data() + ); + rotator_impl::FhtKacRotator rotator(dim, padded); + rotator.load(reinterpret_cast(flip.data())); + for (size_t offset : {0U, 1U}) { + std::vector actual(padded + 2, 12345); + std::copy(input.begin(), input.end(), actual.begin() + 1); + rotator.rotate(actual.data() + 1, actual.data() + offset); + // Preserving the scalar butterfly/normalization order yields identical + // float32 results, including when the input and output overlap. + EXPECT_EQ( + std::memcmp( + actual.data() + offset, expected.data(), padded * sizeof(float) + ), + 0 + ); + EXPECT_EQ(actual.back(), 12345); + } + std::vector saved(flip.size()); + rotator.save(reinterpret_cast(saved.data())); + EXPECT_EQ(saved, flip); + } + } + for (size_t trunc : {0U, 32U, 96U, 131072U}) { + try { + simd::fht_rotate_neon(nullptr, nullptr, 0, 0, trunc, 1, nullptr); + FAIL() << "Invalid transform size must fail before accessing buffers"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "Unsupported dimension for FhtKacRotator"); + } + } +} + +TEST(NeonRotationTest, StandaloneKernelsHandleShortTails) { + if (!cpu::has_neon()) + GTEST_SKIP(); + for (size_t dim = 0; dim <= 33; ++dim) { + std::vector flip((dim + 7) / 8, 0xA5); + std::vector input(dim + 2, 12345); + for (size_t i = 0; i < dim; ++i) + input[i + 1] = static_cast(i) - 16.5F; + auto actual = input, expected = input; + simd::flip_sign_generic(flip.data(), expected.data() + 1, dim); + simd::flip_sign_neon(flip.data(), actual.data() + 1, dim); + EXPECT_EQ(actual, expected); + simd::kacs_walk_generic(expected.data() + 1, dim); + simd::kacs_walk_neon(actual.data() + 1, dim); + EXPECT_EQ(actual, expected); + } +} +#endif diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index ac66889d..e962653c 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -56,12 +56,16 @@ TEST(MaskIpX0Q, SupportsUnalignedCodes) { ASSERT_NE(reinterpret_cast(codes) % alignof(uint64_t), 0U); pack_binary_to_bytes(binary_code.data(), codes, dim); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx2(query.data(), codes, dim), expected); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx512(query.data(), codes, dim), expected); } +#endif EXPECT_FLOAT_EQ(mask_ip_x0_q(query.data(), codes, dim), expected); } @@ -115,6 +119,7 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { const float expected = delta * static_cast(weighted_intersection) + vl * static_cast(data_popcount); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_FLOAT_EQ( simd::warmup_ip_x0_q_512_avx2( @@ -127,6 +132,8 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { expected ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_popcnt()) { EXPECT_FLOAT_EQ( simd::warmup_ip_x0_q_512_avx512( @@ -135,6 +142,7 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { expected ); } +#endif } } } @@ -142,18 +150,22 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { TEST(WarmupIpX0Q, RejectsQueryWidthsBeyondByte) { const std::array data{}; const std::array query{}; +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_THROW( simd::warmup_ip_x0_q_512_avx2(data.data(), query.data(), 1.0F, 0.0F, 64, 9), std::invalid_argument ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_popcnt()) { EXPECT_THROW( simd::warmup_ip_x0_q_512_avx512(data.data(), query.data(), 1.0F, 0.0F, 64, 9), std::invalid_argument ); } +#endif } TEST(Select_IP_Func, returns_stable_function_pointer) { @@ -191,11 +203,13 @@ TEST(Select_IP_Func, returns_stable_function_pointer) { ip_func = select_excode_ipfunc(8); ASSERT_NE(ip_func, nullptr); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { ASSERT_EQ(ip_func, simd::excode_ipimpl::ip16_fxu8_avx512); - } else { + } else if (cpu::has_avx2()) { ASSERT_EQ(ip_func, simd::excode_ipimpl::ip16_fxu8_avx2); } +#endif } TEST(Select_IP_Func, zero_ex_bits_contributes_nothing) { @@ -289,12 +303,16 @@ TEST(ScalarQuantize, HalfValuesMatchScalarAcrossVectorBoundaries) { }; check(simd::scalar_quantize_uint8, simd::scalar_quantize_uint16); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { check(simd::scalar_quantize_uint8_avx2, simd::scalar_quantize_uint16_avx2); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { check(simd::scalar_quantize_uint8_avx512, simd::scalar_quantize_uint16_avx512); } +#endif } TEST(ip16_fxu1_avx, ip_works) { @@ -371,6 +389,7 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { } const float expected_float = static_cast(expected); +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { const std::array avx2_functions{ nullptr, @@ -388,6 +407,8 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { 0.1F ); } +#endif +#if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { const std::array avx512_functions{ nullptr, @@ -405,6 +426,7 @@ TEST(OddBitExcodeIp, MatchesScalarInnerProduct) { 0.1F ); } +#endif } } @@ -423,6 +445,7 @@ TEST(ip_fxu8_avx, ip_works) { 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), @@ -430,6 +453,8 @@ TEST(ip_fxu8_avx, ip_works) { 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), @@ -437,4 +462,5 @@ TEST(ip_fxu8_avx, ip_works) { 0.1F ); } +#endif } diff --git a/tests/unit/rabitqlib/utils/thread_count_test.cpp b/tests/unit/rabitqlib/utils/thread_count_test.cpp new file mode 100644 index 00000000..a45e9f68 --- /dev/null +++ b/tests/unit/rabitqlib/utils/thread_count_test.cpp @@ -0,0 +1,51 @@ +#include + +#include +#include +#include +#include + +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/utils/tools.hpp" + +TEST(ThreadCount, AutomaticAndOversizedRequestsUseHardwareCount) { + const size_t hardware = rabitqlib::total_threads(); + ASSERT_GE(hardware, 1U); + EXPECT_EQ(rabitqlib::resolve_num_threads(0), hardware); + EXPECT_EQ(rabitqlib::resolve_num_threads(1), 1U); + EXPECT_EQ(rabitqlib::resolve_num_threads(hardware), hardware); + EXPECT_EQ(rabitqlib::resolve_num_threads(hardware + 1), hardware); + EXPECT_EQ(rabitqlib::resolve_num_threads(std::numeric_limits::max()), hardware); + if (hardware > 1) { + EXPECT_EQ(rabitqlib::resolve_num_threads(hardware - 1), hardware - 1); + } +} + +TEST(ThreadCount, ParallelForVisitsEachItemOnceWithBoundedWorkerIds) { + const size_t hardware = rabitqlib::total_threads(); + const size_t count = 2 * hardware + 3; + for (size_t requested : {size_t{0}, size_t{1}, std::numeric_limits::max()}) { + SCOPED_TRACE(requested); + std::vector> visits(count); + 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 : hardware); + 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; + }); + 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"; + }); + } +} From c2b0af4ec461ef6bac2190eaca3a0750b8a4cdfb Mon Sep 17 00:00:00 2001 From: gouyt13 Date: Thu, 24 Sep 2026 19:23:33 +0800 Subject: [PATCH 2/4] fix for win --- include/rabitqlib/simd/fastscan_dispatch.hpp | 27 +++++++++--- src/simd/dispatch.cpp | 44 +++++++++++--------- src/simd/fastscan_generic.cpp | 14 +++++-- src/simd/fastscan_neon.cpp | 10 ++++- tests/dispatch/fastscan_no_simd_test.cpp | 15 +++++++ 5 files changed, 80 insertions(+), 30 deletions(-) diff --git a/include/rabitqlib/simd/fastscan_dispatch.hpp b/include/rabitqlib/simd/fastscan_dispatch.hpp index ce2d29e2..7c61d264 100644 --- a/include/rabitqlib/simd/fastscan_dispatch.hpp +++ b/include/rabitqlib/simd/fastscan_dispatch.hpp @@ -6,9 +6,17 @@ namespace rabitqlib::fastscan::simd { // NEON uses the portable/AVX2 high-accuracy LUT byte order. -void accumulate_neon(const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim); +void accumulate_neon( + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim +); void accumulate_hacc_neon( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* result, + size_t dim ); void pack_lut_neon(size_t dim, const float* query, float* lut); @@ -18,10 +26,16 @@ void pack_lut_avx512(size_t dim, const float* query, float* lut); // Scalar reference for backend correctness tests; never selected by runtime dispatch. void accumulate_generic( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim ); void accumulate_unsupported( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim ); void accumulate_avx2( const uint8_t* __restrict__ codes, @@ -32,7 +46,10 @@ void accumulate_avx2( void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc_generic( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* result, + size_t dim ); void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 25cf9fd4..b8cbc193 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -51,9 +51,11 @@ Function resolve_kernel(Function avx2, Function fallback) { return fallback; } -// Keep ARM selection cached in the same resolver as the x86 tiers. +// 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_kernel(Function preferred, Function fallback, bool supported) { +Function resolve_optional_kernel(Function preferred, Function fallback, bool supported) { return supported ? preferred : fallback; } @@ -169,13 +171,13 @@ void split_batch_estdist( #if defined(__aarch64__) const auto kEuclideanSqrFn = - resolve_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); + resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); const auto kDotProductFn = - resolve_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); + resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); const auto kDotProductDisFn = - resolve_kernel(dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()); + resolve_optional_kernel(dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()); const auto kL2normSqrFn = - resolve_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); + resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); #else const auto kEuclideanSqrFn = RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); @@ -223,7 +225,7 @@ double best_rescale_factor( #if defined(__aarch64__) || defined(_M_ARM64) const auto kFhtRotateFn = - resolve_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon()); + 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); @@ -254,7 +256,7 @@ ExcodeIpTable resolve_excode_ip_table() { excode_ipimpl::ip16_fxu8_generic, }; #if defined(__aarch64__) || defined(_M_ARM64) - return resolve_kernel( + return resolve_optional_kernel( ExcodeIpTable{ ip_fxu0, excode_ipimpl::ip16_fxu1_neon, @@ -298,24 +300,26 @@ ExcodeIpTable resolve_excode_ip_table() { } #if defined(__aarch64__) || defined(_M_ARM64) -const auto kFlipSignFn = resolve_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon()); +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); #endif #if defined(__aarch64__) || defined(_M_ARM64) -const auto kKacsWalkFn = resolve_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon()); +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 #if defined(__aarch64__) || defined(_M_ARM64) -const auto kScalarQuantizeUint8Fn = resolve_kernel( +const auto kScalarQuantizeUint8Fn = resolve_optional_kernel( scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() ); -const auto kScalarQuantizeUint16Fn = resolve_kernel( +const auto kScalarQuantizeUint16Fn = resolve_optional_kernel( scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon() ); #else @@ -406,10 +410,10 @@ const ex_ipfunc kIp64Fxu6AvxFn = kExcodeIpTable[6]; const ex_ipfunc kIp64Fxu7AvxFn = kExcodeIpTable[7]; #if defined(__aarch64__) || defined(_M_ARM64) -const auto kNewTransposeBinFn = simd::resolve_kernel( +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_kernel( +const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel( simd::new_transpose_bin_512_neon, simd::new_transpose_bin_512_generic, cpu::has_neon() ); #else @@ -428,7 +432,7 @@ const auto kNewTransposeBin512Fn = RABITQ_RESOLVE( using MaskIpX0QFn = float (*)(const float*, const uint8_t*, size_t); #if defined(__aarch64__) || defined(_M_ARM64) -const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_kernel( +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() @@ -514,7 +518,7 @@ float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) namespace rabitqlib::fastscan { #if defined(__aarch64__) || defined(_M_ARM64) -const auto kPackLutFn = rabitqlib::simd::resolve_kernel( +const auto kPackLutFn = rabitqlib::simd::resolve_optional_kernel( simd::pack_lut_neon, simd::pack_lut_generic, cpu::has_neon() ); #else @@ -528,7 +532,7 @@ void pack_lut(size_t dim, const float* __restrict__ query, float* __restr } #if defined(__aarch64__) -const auto kAccumulateFn = rabitqlib::simd::resolve_kernel( +const auto kAccumulateFn = rabitqlib::simd::resolve_optional_kernel( simd::accumulate_neon, simd::accumulate_unsupported, cpu::has_neon() ); #else @@ -544,7 +548,7 @@ const auto kTransferLutHaccFn = RABITQ_RESOLVE( ); #if defined(__aarch64__) -const auto kAccumulateHaccFn = rabitqlib::simd::resolve_kernel( +const auto kAccumulateHaccFn = rabitqlib::simd::resolve_optional_kernel( simd::accumulate_hacc_neon, simd::accumulate_hacc_generic, cpu::has_neon() ); #else @@ -595,7 +599,7 @@ 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_kernel( +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() @@ -651,7 +655,7 @@ const SearchKnnFn kSearchKnnFn = cpu::has_avx512_popcnt() ); #else const SearchKnnFn kSearchKnnFn = - simd::resolve_kernel(search_knn_neon, search_knn_generic, cpu::has_neon()); + simd::resolve_optional_kernel(search_knn_neon, search_knn_generic, cpu::has_neon()); #endif } // namespace std::priority_queue> search_knn( diff --git a/src/simd/fastscan_generic.cpp b/src/simd/fastscan_generic.cpp index d29c21ee..bd4f44ac 100644 --- a/src/simd/fastscan_generic.cpp +++ b/src/simd/fastscan_generic.cpp @@ -64,18 +64,26 @@ void accumulate_generic_impl( } } // namespace void accumulate_generic( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim ) { accumulate_generic_impl(codes, lut, result, dim); } -void accumulate_unsupported(const uint8_t*, const uint8_t*, int32_t*, size_t) { +void accumulate_unsupported( + const uint8_t* __restrict__, const uint8_t* __restrict__, int32_t* __restrict__, size_t +) { throw std::runtime_error( "Standard FastScan accumulation requires AVX2/FMA, AVX-512, or ARM NEON; " "no supported SIMD backend is available" ); } void accumulate_hacc_generic( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* result, + size_t dim ) { accumulate_generic_impl(codes, lut, result, dim); } diff --git a/src/simd/fastscan_neon.cpp b/src/simd/fastscan_neon.cpp index 202bd996..3a590f9d 100644 --- a/src/simd/fastscan_neon.cpp +++ b/src/simd/fastscan_neon.cpp @@ -67,12 +67,18 @@ void accumulate_impl( } } // namespace void accumulate_neon( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim ) { accumulate_impl(codes, lut, result, dim); } void accumulate_hacc_neon( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* result, + size_t dim ) { accumulate_impl(codes, lut, result, dim); } diff --git a/tests/dispatch/fastscan_no_simd_test.cpp b/tests/dispatch/fastscan_no_simd_test.cpp index f95ebb47..7fcf6e05 100644 --- a/tests/dispatch/fastscan_no_simd_test.cpp +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -5,6 +5,7 @@ #include #include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" #include "rabitqlib/utils/cpu_features.hpp" // Link-time substitutes, isolated from both production detection and the main suite. @@ -40,5 +41,19 @@ TEST(FastScanDispatchTest, InvalidDimensionStillRaisesInvalidArgument) { EXPECT_THROW(accumulate(nullptr, nullptr, nullptr, 15), std::invalid_argument); } +TEST(FastScanDispatchTest, HighAccuracyUsesScalarFallbackWithoutSimd) { + constexpr size_t kDim = 16; + std::array codes{}; + std::array lut; + lut.fill(1000); + std::array high_accuracy_lut{}; + transfer_lut_hacc(lut.data(), kDim, high_accuracy_lut.data()); + std::array result{}; + accumulate_hacc(codes.data(), high_accuracy_lut.data(), result.data(), kDim); + for (const auto value : result) { + EXPECT_EQ(value, 4000); + } +} + } // namespace } // namespace rabitqlib::fastscan From 4111d71992492d9ea6aacff46dcdefd36c53bd76 Mon Sep 17 00:00:00 2001 From: gouyt13 Date: Thu, 24 Sep 2026 20:11:01 +0800 Subject: [PATCH 3/4] fix for win --- sample/python/faiss_clustering.py | 2 +- sample/python/hnsw_rabitq_indexing.py | 2 +- sample/python/hnsw_rabitq_querying.py | 2 +- sample/python/ivf_rabitq_indexing.py | 2 +- sample/python/ivf_rabitq_querying.py | 2 +- sample/python/symqg_indexing.py | 2 +- sample/python/symqg_querying.py | 2 +- tests/python/test_examples.py | 3 +++ 8 files changed, 10 insertions(+), 7 deletions(-) diff --git a/sample/python/faiss_clustering.py b/sample/python/faiss_clustering.py index 98d3cfdb..f14836ac 100644 --- a/sample/python/faiss_clustering.py +++ b/sample/python/faiss_clustering.py @@ -46,7 +46,7 @@ def main(args) -> None: np.savez( output, centroids=centroids, cluster_ids=cluster_ids, metric=args.metric ) - print(f"Clusters saved → {args.clusters_file}") + print(f"Clusters saved: {args.clusters_file}") if __name__ == "__main__": diff --git a/sample/python/hnsw_rabitq_indexing.py b/sample/python/hnsw_rabitq_indexing.py index 63552d05..1aec6250 100644 --- a/sample/python/hnsw_rabitq_indexing.py +++ b/sample/python/hnsw_rabitq_indexing.py @@ -53,7 +53,7 @@ def main(args=None) -> None: print(f"Indexing time: {time() - t0:.2f}s") idx.save(args.index_file) - print(f"Index saved → {args.index_file}") + print(f"Index saved: {args.index_file}") if __name__ == "__main__": diff --git a/sample/python/hnsw_rabitq_querying.py b/sample/python/hnsw_rabitq_querying.py index 40dbf47f..34b7f8ec 100644 --- a/sample/python/hnsw_rabitq_querying.py +++ b/sample/python/hnsw_rabitq_querying.py @@ -25,7 +25,7 @@ def main(args=None) -> None: # 2. Load index idx = HnswIndex.load(args.index_file) - print(f"Index loaded — dim={idx.dim}") + print(f"Index loaded: dim={idx.dim}") print(f"TopK: {args.topk}") print("\nsearch start >.....\n") diff --git a/sample/python/ivf_rabitq_indexing.py b/sample/python/ivf_rabitq_indexing.py index 3abf2c1a..fd7781cc 100644 --- a/sample/python/ivf_rabitq_indexing.py +++ b/sample/python/ivf_rabitq_indexing.py @@ -54,7 +54,7 @@ def main(args=None) -> None: print("IVF constructed") idx.save(args.index_file) print(f"Indexing time: {elapsed_min:.4f} min") - print(f"Index saved → {args.index_file}") + print(f"Index saved: {args.index_file}") if __name__ == "__main__": diff --git a/sample/python/ivf_rabitq_querying.py b/sample/python/ivf_rabitq_querying.py index c95cfe6f..5dea3626 100644 --- a/sample/python/ivf_rabitq_querying.py +++ b/sample/python/ivf_rabitq_querying.py @@ -28,7 +28,7 @@ def main(args=None) -> None: # 2. Load index idx = IvfIndex.load(args.index_file) - print(f"Index loaded — dim={idx.dim}, clusters={idx.num_clusters}") + print(f"Index loaded: dim={idx.dim}, clusters={idx.num_clusters}") all_qps = np.zeros((args.test_rounds, len(NPROBES))) all_recall = np.zeros((args.test_rounds, len(NPROBES))) diff --git a/sample/python/symqg_indexing.py b/sample/python/symqg_indexing.py index 323abac5..4b5dea45 100644 --- a/sample/python/symqg_indexing.py +++ b/sample/python/symqg_indexing.py @@ -47,7 +47,7 @@ def main(args=None) -> None: print(f"Indexing time: {time() - t0:.2f}s") idx.save(args.index_file) - print(f"Index saved → {args.index_file}") + print(f"Index saved: {args.index_file}") if __name__ == "__main__": diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index 6228064e..198b87ab 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -26,7 +26,7 @@ def main(args=None) -> None: # 2. Load index idx = SymqgIndex.load(args.index_file) - print(f"Index loaded — dim={idx.dim}, metric={args.metric}") + print(f"Index loaded: dim={idx.dim}, metric={args.metric}") print(f"TopK: {args.topk}") print("\nsearch start >.....\n") diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index aa94e476..7441c31b 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -33,6 +33,8 @@ def find_spec(self, fullname, path=None, target=None): def run_example(name, *args, blocked="faiss"): env = os.environ.copy() + # Reproduce Windows redirected output even on hosts using UTF-8 by default. + env["PYTHONIOENCODING"] = "cp1252:strict" for key in ( "DYLD_LIBRARY_PATH", "DYLD_INSERT_LIBRARIES", @@ -53,6 +55,7 @@ def run_example(name, *args, blocked="faiss"): env=env, capture_output=True, text=True, + encoding="cp1252", timeout=60, ) assert result.returncode == 0, result.stdout + result.stderr From 4be0f25958e8dc345227387c7c0fd25b611c528d Mon Sep 17 00:00:00 2001 From: gouyt13 Date: Thu, 24 Sep 2026 20:21:10 +0800 Subject: [PATCH 4/4] fix for qps divide by 0 --- sample/python/hnsw_rabitq_querying.py | 9 +++--- sample/python/ivf_rabitq_querying.py | 9 +++--- sample/python/symqg_querying.py | 9 +++--- tests/python/test_examples.py | 44 +++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 12 deletions(-) diff --git a/sample/python/hnsw_rabitq_querying.py b/sample/python/hnsw_rabitq_querying.py index 34b7f8ec..9c61b7ba 100644 --- a/sample/python/hnsw_rabitq_querying.py +++ b/sample/python/hnsw_rabitq_querying.py @@ -1,5 +1,5 @@ import argparse -from time import time +from time import perf_counter import numpy as np from rabitqlib import HnswIndex @@ -35,13 +35,14 @@ def main(args=None) -> None: for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): - t0 = time() + t0 = perf_counter() ids, _ = idx.search( queries, k=args.topk, ef=ef, num_threads=args.num_threads ) - elapsed = time() - t0 # seconds + elapsed = perf_counter() - t0 # seconds - qps = nq / elapsed + # Report an unmeasurable duration instead of inventing a QPS value. + qps = nq / elapsed if elapsed > 0 else float("nan") recall = compute_recall(ids, gt, args.topk) all_qps[r, i_probe] = qps diff --git a/sample/python/ivf_rabitq_querying.py b/sample/python/ivf_rabitq_querying.py index 5dea3626..b607db69 100644 --- a/sample/python/ivf_rabitq_querying.py +++ b/sample/python/ivf_rabitq_querying.py @@ -1,5 +1,5 @@ import argparse -from time import time +from time import perf_counter import numpy as np from rabitqlib import IvfIndex @@ -42,7 +42,7 @@ def main(args=None) -> None: f"will use nprobe = num_clusters ({idx.num_clusters})." ) - t0 = time() + t0 = perf_counter() ids, _ = idx.search( queries, k=args.topk, @@ -50,9 +50,10 @@ def main(args=None) -> None: high_accuracy=args.use_hacc, num_threads=args.num_threads, ) - elapsed = time() - t0 + elapsed = perf_counter() - t0 - all_qps[r, probe_index] = nq / elapsed + # Report an unmeasurable duration instead of inventing a QPS value. + all_qps[r, probe_index] = nq / elapsed if elapsed > 0 else float("nan") all_recall[r, probe_index] = compute_recall(ids, gt, args.topk) avg_qps = all_qps.mean(axis=0) diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index 198b87ab..aab3c9d7 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -1,5 +1,5 @@ import argparse -from time import time +from time import perf_counter import numpy as np from rabitqlib import SymqgIndex @@ -36,13 +36,14 @@ def main(args=None) -> None: for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): - t0 = time() + t0 = perf_counter() ids, _ = idx.search( queries, k=args.topk, ef=ef, num_threads=args.num_threads ) - elapsed = time() - t0 # seconds + elapsed = perf_counter() - t0 # seconds - qps = nq / elapsed + # Report an unmeasurable duration instead of inventing a QPS value. + qps = nq / elapsed if elapsed > 0 else float("nan") recall = compute_recall(ids, gt, args.topk) all_qps[r, i_probe] = qps diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index 7441c31b..abd132df 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -5,6 +5,7 @@ import subprocess import sys from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -69,6 +70,49 @@ def write_vecs(path, values): words.tofile(path) +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "symqg"]) +@pytest.mark.parametrize("elapsed", [0.0, 0.5]) +def test_query_timing_handles_unchanged_clock(kind, elapsed, monkeypatch, capsys): + name = "symqg_querying" if kind == "symqg" else f"{kind}_rabitq_querying" + monkeypatch.syspath_prepend(str(EXAMPLES)) + spec = importlib.util.spec_from_file_location(name, EXAMPLES / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # 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]) + 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))) + ) + 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_ivecs", lambda path: ids) + module.main( + SimpleNamespace( + index_file="index", + query_file="queries", + gt_file="truth", + topk=1, + test_rounds=1, + num_threads=1, + metric="l2", + use_hacc=None, + ) + ) + row = capsys.readouterr().out.splitlines()[-1].split() + qps, recall = map(float, row[1:]) + if elapsed == 0: + assert np.isnan(qps) + else: + assert qps == pytest.approx(4.0) + assert recall == 1.0 + + @pytest.fixture def example_data(tmp_path): data = np.random.default_rng(641).standard_normal((128, 65)).astype(np.float32)