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..7c61d264 100644 --- a/include/rabitqlib/simd/fastscan_dispatch.hpp +++ b/include/rabitqlib/simd/fastscan_dispatch.hpp @@ -5,17 +5,52 @@ namespace rabitqlib::fastscan::simd { +// NEON uses the portable/AVX2 high-accuracy LUT byte order. +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* __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); 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* __restrict__ codes, + const uint8_t* __restrict__ lut, + int32_t* __restrict__ result, + size_t dim +); +void accumulate_unsupported( + 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, 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* __restrict__ codes, + const uint8_t* __restrict__ 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..f14836ac --- /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..1aec6250 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 @@ -56,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__": @@ -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/hnsw_rabitq_querying.py b/sample/python/hnsw_rabitq_querying.py index 40dbf47f..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 @@ -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") @@ -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_indexing.py b/sample/python/ivf_rabitq_indexing.py index 4ba54705..fd7781cc 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, ) @@ -57,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__": @@ -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/ivf_rabitq_querying.py b/sample/python/ivf_rabitq_querying.py index c95cfe6f..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 @@ -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))) @@ -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_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..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 @@ -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") @@ -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/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..b8cbc193 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -51,10 +51,27 @@ Function resolve_kernel(Function avx2, Function fallback) { return fallback; } +// Use a distinct name: a fallback function pointer can implicitly convert to bool. +// Overloading the x86 resolver would let MSVC select this overload when __restrict +// makes the SIMD and scalar function-pointer types differ. +template +Function resolve_optional_kernel(Function preferred, Function fallback, bool supported) { + return supported ? preferred : fallback; +} + } // namespace +// Non-x86 builds must not reference symbols from excluded AVX objects. +#if defined(__x86_64__) || defined(_M_X64) +#define RABITQ_RESOLVE(...) rabitqlib::simd::resolve_kernel(__VA_ARGS__) +#define RABITQ_RESOLVE2(...) rabitqlib::simd::resolve_kernel(__VA_ARGS__) +#else +#define RABITQ_RESOLVE(avx512, avx2, fallback, ...) fallback +#define RABITQ_RESOLVE2(avx2, fallback) fallback +#endif + const auto 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 +84,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 +102,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 +125,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 +139,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 +151,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 +169,26 @@ void split_batch_estdist( ); } +#if defined(__aarch64__) +const auto kEuclideanSqrFn = + resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); +const auto kDotProductFn = + resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); +const auto kDotProductDisFn = + resolve_optional_kernel(dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()); +const auto kL2normSqrFn = + resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); +#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 +204,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 +213,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 +223,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_optional_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon()); +#else +const auto kFhtRotateFn = + RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic); +#endif void fht_rotate( const float* data, @@ -220,65 +243,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() { - return resolve_kernel( + 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_optional_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 +282,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 +293,64 @@ 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_optional_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_avx512, kacs_walk_avx2, missing_kacs_walk); + 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 -const auto kScalarQuantizeUint8Fn = resolve_kernel( - scalar_quantize_uint8_avx512, scalar_quantize_uint8_avx2, missing_scalar_quantize_uint8 +#if defined(__aarch64__) || defined(_M_ARM64) +const auto kScalarQuantizeUint8Fn = resolve_optional_kernel( + scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() +); +const auto kScalarQuantizeUint16Fn = resolve_optional_kernel( + scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon() +); +#else +const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE( + scalar_quantize_uint8_avx512, scalar_quantize_uint8_avx2, scalar_quantize_uint8_generic ); -const auto kScalarQuantizeUint16Fn = resolve_kernel( +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 +409,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_optional_kernel( + simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon() +); +const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel( + simd::new_transpose_bin_512_neon, simd::new_transpose_bin_512_generic, cpu::has_neon() +); +#else +const auto kNewTransposeBinFn = RABITQ_RESOLVE( simd::new_transpose_bin_avx512, simd::new_transpose_bin_avx2, - simd::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_optional_kernel( + static_cast(simd::mask_ip_x0_q_neon), + static_cast(simd::mask_ip_x0_q_generic), + cpu::has_neon() +); +#else +const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE( static_cast(simd::mask_ip_x0_q_avx512), static_cast(simd::mask_ip_x0_q_avx2), - 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 +517,45 @@ float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) namespace rabitqlib::fastscan { -const auto kPackLutFn = rabitqlib::simd::resolve_kernel( - simd::pack_lut_avx512, simd::pack_lut_avx2, simd::pack_lut_generic +#if defined(__aarch64__) || defined(_M_ARM64) +const auto kPackLutFn = rabitqlib::simd::resolve_optional_kernel( + simd::pack_lut_neon, simd::pack_lut_generic, cpu::has_neon() ); +#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); } -const auto kAccumulateFn = rabitqlib::simd::resolve_kernel( - simd::accumulate_avx512, - simd::accumulate_avx2, - rabitqlib::simd::missing_fastscan_accumulate +#if defined(__aarch64__) +const auto kAccumulateFn = rabitqlib::simd::resolve_optional_kernel( + simd::accumulate_neon, simd::accumulate_unsupported, cpu::has_neon() ); +#else +const auto kAccumulateFn = RABITQ_RESOLVE( + simd::accumulate_avx512, simd::accumulate_avx2, simd::accumulate_unsupported +); +#endif -const auto kTransferLutHaccFn = 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 ); -const auto kAccumulateHaccFn = rabitqlib::simd::resolve_kernel( - simd::accumulate_hacc_avx512, - simd::accumulate_hacc_avx2, - rabitqlib::simd::missing_fastscan_accumulate_hacc +#if defined(__aarch64__) +const auto kAccumulateHaccFn = rabitqlib::simd::resolve_optional_kernel( + 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 +598,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_optional_kernel( + static_cast(simd::warmup_ip_x0_q_512_neon), + static_cast(simd::warmup_ip_x0_q_512_generic), + cpu::has_neon() +); +#else +const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE( + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), + static_cast(rabitqlib::simd::warmup_ip_x0_q_512_generic), + cpu::has_avx512_popcnt() +); +#endif float warmup_ip_x0_q_512( const uint8_t* data, @@ -602,20 +643,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_optional_kernel(search_knn_neon, search_knn_generic, cpu::has_neon()); +#endif } // namespace std::priority_queue> search_knn( HierarchicalNSW& index, const float* query, size_t topk diff --git a/src/simd/fastscan_generic.cpp b/src/simd/fastscan_generic.cpp index 6a20fb1e..bd4f44ac 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,71 @@ 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* __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* __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* __restrict__ codes, + const uint8_t* __restrict__ 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..3a590f9d --- /dev/null +++ b/src/simd/fastscan_neon.cpp @@ -0,0 +1,85 @@ +#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* __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* __restrict__ codes, + const uint8_t* __restrict__ 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..7fcf6e05 --- /dev/null +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -0,0 +1,59 @@ +#include + +#include +#include +#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. +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); +} + +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 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..abd132df --- /dev/null +++ b/tests/python/test_examples.py @@ -0,0 +1,266 @@ +"""Example stages keep Faiss and RaBitQ in different Python processes.""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +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() + # 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", + "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, + encoding="cp1252", + 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.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) + 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"; + }); + } +}