diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 1b33c70..5827f3c 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -55,6 +55,44 @@ jobs: - name: Run clang-tidy run: ./scripts/check-tidy.sh build-tidy + includes: + needs: changes + if: ${{ !cancelled() && (needs.changes.result != 'success' || needs.changes.outputs.cpp == 'true') }} + name: Include cleaner (advisory) + runs-on: ubuntu-latest + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@v7 + - name: Install clang-tidy and build dependencies + run: | + sudo apt-get update -qq + sudo apt-get install -y clang-tidy clang libomp-dev cmake ninja-build + - name: Configure compilation database + run: >- + cmake -S . -B build-includes -G Ninja + -DCMAKE_CXX_COMPILER=clang++ + -DRABITQ_BUILD_SAMPLES=OFF + -DRABITQ_BUILD_TESTS=OFF + -DRABITQ_BUILD_PYTHON_BINDINGS=OFF + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF + -DCMAKE_BUILD_TYPE=Release + - name: Analyze library includes + id: analysis + continue-on-error: true + shell: bash + run: | + set -o pipefail + ./scripts/check-includes.sh build-includes 2>&1 | tee include-report.txt + - name: Store include report + if: ${{ !cancelled() && steps.analysis.outcome != 'skipped' }} + uses: actions/upload-artifact@v7 + with: + name: include-report + path: include-report.txt + if-no-files-found: error + shellcheck: needs: changes if: ${{ !cancelled() && (needs.changes.result != 'success' || needs.changes.outputs.shell == 'true') }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index efae1fc..8010697 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -192,6 +192,65 @@ Focused clang-tidy checks on affected code are sufficient during iteration. For first-party C++ changes, run the full check above before merging and report whether validation was focused or complete. +### Include dependency reports + +The **Include cleaner (advisory)** job uses clang-tidy's `misc-include-cleaner` +check to report missing and unused includes. CI uses `ubuntu-latest` and +unversioned distribution packages: + +```bash +sudo apt-get update +sudo apt-get install clang-tidy clang libomp-dev cmake ninja-build +``` + +Clang-tidy 17 or newer is required. On older distributions such as Ubuntu 22.04, +use [LLVM's APT repository](https://apt.llvm.org/) to install a newer release. +For example, install LLVM 22 and its matching analysis and OpenMP packages: + +```bash +wget https://apt.llvm.org/llvm.sh +sudo bash llvm.sh 22 +sudo apt-get install clang-tidy-22 libomp-22-dev cmake ninja-build +export CLANG_TIDY=clang-tidy-22 +export CXX=clang++-22 +``` + +Configure and run the same check locally (use a fresh build directory when +changing compilers): + +```bash +cmake -S . -B build-includes -G Ninja \ + -DCMAKE_CXX_COMPILER="${CXX:-clang++}" \ + -DRABITQ_BUILD_SAMPLES=OFF \ + -DRABITQ_BUILD_TESTS=OFF \ + -DRABITQ_BUILD_PYTHON_BINDINGS=OFF \ + -DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF \ + -DCMAKE_BUILD_TYPE=Release +./scripts/check-includes.sh build-includes +``` + +The script checks library sources using their compilation database and checks +headers as main files, since this clang-tidy check does not report findings in +included headers. Private headers are checked with AVX2 and AVX-512 flags. +Vendored files are excluded. The script also ignores suggestions to include +Eigen and hnswlib implementation headers behind their existing public headers; +these vendor snapshots lack the export annotations needed by include-cleaner. +`INCLUDE_JOBS` controls parallelism (default: 2). +The script fails for findings or analyzer errors; CI keeps this step advisory +and uploads the `include-report` artifact without modifying files. + +Review suggestions before applying them, especially for templates and public +forwarding headers. For a source file, automatic fixes can be applied with: + +```bash +"${CLANG_TIDY:-clang-tidy}" -p build-includes --config='{}' \ + --checks='-*,misc-include-cleaner' --fix src/simd/dispatch.cpp +``` + +Review the diff, run formatting, and rebuild and test affected code after fixes. +Findings can vary between LLVM releases. This include check is independent of +the existing general clang-tidy job. + ### Focused static analysis Use a temporary subset of the compilation database with the same wrapper to diff --git a/include/rabitqlib/defines.hpp b/include/rabitqlib/defines.hpp index 9dc7017..b2d722e 100644 --- a/include/rabitqlib/defines.hpp +++ b/include/rabitqlib/defines.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 7fc285b..8be9b3a 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -5,11 +5,8 @@ #include #include -#include +#include #include -#include -#include -#include #include "rabitqlib/defines.hpp" @@ -82,36 +79,38 @@ static inline void get_column( inline void pack_codes( size_t padded_dim, const uint8_t* quantization_code, size_t num, uint8_t* blocks ) { - size_t num_rd = (num + 31) & ~31; // round up num of vecs to multiple of batch size(32) - // consider codes is a matrix // rows = number of vectors // cols = number of uint8_t of one vector's code size_t cols = padded_dim / 8; - std::array col; // column of a batch of code, 8 bits - std::array col_0; // upper 4 bits - std::array col_1; // lower 4 bits + // Full batches can gather directly in output order, without a temporary column. + size_t row = 0; + for (; num - row >= kBatchSize; row += kBatchSize) { + for (size_t i = 0; i < cols; ++i) { + for (size_t j = 0; j < 16; ++j) { + const size_t first = row + kPerm0[j]; + const uint8_t a = quantization_code[first * cols + i]; + const uint8_t b = quantization_code[(first + 16) * cols + i]; + blocks[j] = (a >> 4) | (b & 0xF0); + blocks[j + 16] = (a & 0x0F) | (b << 4); + } + blocks += 32; + } + } - // pack codes batch by batch - // each batch contain codes for 32 vectors - for (size_t row = 0; row < num_rd; row += kBatchSize) { - // get quantization codes for each column for each batch - // i.e., we get the codes for 8 dims of 32 vectors and reorganize the data layout - // based on the shuffle SIMD instruction used during querying + // Only the final partial batch needs a zero-padded column. + if (row < num) { + std::array col; for (size_t i = 0; i < cols; ++i) { get_column(quantization_code, num, cols, row, i, col); - for (size_t j = 0; j < 32; ++j) { - col_0[j] = col[j] >> 4; - col_1[j] = col[j] & 15; - } for (size_t j = 0; j < 16; ++j) { // the lower 4 bits represent vector 0 to 15 // the upper 4 bits represent vector 16 to 31 - uint8_t val0 = col_0[kPerm0[j]] | (col_0[kPerm0[j] + 16] << 4); - uint8_t val1 = col_1[kPerm0[j]] | (col_1[kPerm0[j] + 16] << 4); - blocks[j] = val0; - blocks[j + 16] = val1; + const uint8_t a = col[kPerm0[j]]; + const uint8_t b = col[kPerm0[j] + 16]; + blocks[j] = (a >> 4) | (b & 0xF0); + blocks[j + 16] = (a & 0x0F) | (b << 4); } blocks += 32; } @@ -140,4 +139,7 @@ inline void pack_lut(size_t dim, const T* __restrict__ query, T* __restrict__ lu query += 4; } } +// Float query tables use runtime SIMD dispatch; other types keep the generic path. +template <> +void pack_lut(size_t dim, const float* __restrict__ query, float* __restrict__ lut); } // namespace rabitqlib::fastscan diff --git a/include/rabitqlib/fastscan/highacc_fastscan.hpp b/include/rabitqlib/fastscan/highacc_fastscan.hpp index d7b05b7..b567424 100644 --- a/include/rabitqlib/fastscan/highacc_fastscan.hpp +++ b/include/rabitqlib/fastscan/highacc_fastscan.hpp @@ -1,10 +1,7 @@ #pragma once -#include +#include #include -#include -#include -#include namespace rabitqlib::fastscan { /** diff --git a/include/rabitqlib/index/estimator.hpp b/include/rabitqlib/index/estimator.hpp index f75e87a..3f75c05 100644 --- a/include/rabitqlib/index/estimator.hpp +++ b/include/rabitqlib/index/estimator.hpp @@ -1,5 +1,7 @@ #pragma once +#include +#include #include #include "rabitqlib/defines.hpp" diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 03049b9..e85f6e9 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -1,16 +1,24 @@ // HNSW is developed from the [HNSW library](https://github.com/nmslib/hnswlib) #pragma once -#include - +#include #include #include +#include #include +#include +#include +#include #include +#include +#include #include #include +#include +#include #include #include +#include #include #include "rabitqlib/defines.hpp" @@ -21,10 +29,12 @@ #include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/utils/buffer.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" #include "rabitqlib/utils/visited_pool.hpp" +#include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::hnsw { diff --git a/include/rabitqlib/index/ivf/cluster.hpp b/include/rabitqlib/index/ivf/cluster.hpp index bd7ae0f..6917b16 100644 --- a/include/rabitqlib/index/ivf/cluster.hpp +++ b/include/rabitqlib/index/ivf/cluster.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "rabitqlib/defines.hpp" diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index dddcaf0..8890692 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -1,11 +1,17 @@ #pragma once +#include +#include #include #include #include +#include #include +#include +#include #include #include +#include #include #include "rabitqlib/defines.hpp" diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 951b851..1cf6d70 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -1,14 +1,15 @@ #pragma once -#include - #include +#include #include -#include #include #include +#include #include #include +#include +#include #include #include #include diff --git a/include/rabitqlib/index/lut.hpp b/include/rabitqlib/index/lut.hpp index 7d76829..7b135c6 100644 --- a/include/rabitqlib/index/lut.hpp +++ b/include/rabitqlib/index/lut.hpp @@ -2,6 +2,8 @@ #include #include +#include +#include #include #include "rabitqlib/fastscan/fastscan.hpp" diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index ad1e4fd..1c31313 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -1,10 +1,11 @@ #pragma once -#include +#include #include #include #include #include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/index/lut.hpp" diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index ede4a81..24b2f5f 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -1,12 +1,13 @@ #pragma once -#include - +#include #include #include #include #include #include +#include +#include #include #include #include @@ -19,6 +20,7 @@ #include "rabitqlib/index/estimator.hpp" #include "rabitqlib/index/query.hpp" #include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/pack_excode.hpp" #include "rabitqlib/quantization/rabitq.hpp" #include "rabitqlib/utils/array.hpp" #include "rabitqlib/utils/buffer.hpp" @@ -26,6 +28,7 @@ #include "rabitqlib/utils/memory.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" #include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index ef20521..19f86a4 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -1,9 +1,13 @@ #pragma once +#include + #include #include -#include +#include +#include #include +#include #include #include #include diff --git a/include/rabitqlib/quantization/data_layout.hpp b/include/rabitqlib/quantization/data_layout.hpp index 83e0ffa..bc12905 100644 --- a/include/rabitqlib/quantization/data_layout.hpp +++ b/include/rabitqlib/quantization/data_layout.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include "rabitqlib/fastscan/fastscan.hpp" diff --git a/include/rabitqlib/quantization/rabitq.hpp b/include/rabitqlib/quantization/rabitq.hpp index 8ef4e23..f3ea9a6 100644 --- a/include/rabitqlib/quantization/rabitq.hpp +++ b/include/rabitqlib/quantization/rabitq.hpp @@ -1,8 +1,5 @@ #pragma once -#include - -#include #include #include diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index ef860ae..b0447da 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include #include @@ -91,13 +89,18 @@ inline void one_bit_code_with_factor( T& f_error, MetricType metric_type = METRIC_L2 ) { - // quantize - RowMajorArray residual_arr = one_bit_code(data, centroid, dim, binary_code); + // Scratch stays private to each thread and is reused across vectors and batches. + thread_local RowMajorArray residual_arr; + thread_local RowMajorArray xu_cb; + ConstRowMajorArrayMap data_arr(data, 1, dim); + ConstRowMajorArrayMap cent_arr(centroid, 1, dim); + residual_arr = data_arr - cent_arr; + RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); + x_u = (residual_arr > 0).template cast(); // xu_cb = x_u + cb, xu_cb has same direction and different length with x_bar float cb = -((1 << 1) - 1) / 2.F; - RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); - RowMajorArray xu_cb = x_u.template cast() + cb; + xu_cb = x_u.template cast() + cb; // distance to centroid T l2_sqr = l2norm_sqr(residual_arr.data(), dim); @@ -186,8 +189,8 @@ inline void one_bit_compact_code( T& f_error, MetricType metric_type = METRIC_L2 ) { - // binary code - std::vector binary_code(padded_dim); + thread_local std::vector binary_code; + binary_code.resize(padded_dim); // get binary code one_bit_code_with_factor( @@ -204,7 +207,7 @@ inline void one_bit_compact_code( pack_binary(binary_code.data(), compact_code, padded_dim); } -// ! padded_dim % 64 == 0 +// Requires a positive padded_dim divisible by sizeof(TC) * 8. template inline void one_bit_compact_codes( const T* data, @@ -234,7 +237,8 @@ inline void one_bit_compact_codes( } } -// ! padded_dim % 64 == 0 +// Encoding requires a positive padded_dim divisible by 8; index pipelines use 64. +// packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail padding. template inline void one_bit_batch_code( const T* data, diff --git a/include/rabitqlib/simd/dispatch.hpp b/include/rabitqlib/simd/dispatch.hpp index 2004e8f..7b6a5d8 100644 --- a/include/rabitqlib/simd/dispatch.hpp +++ b/include/rabitqlib/simd/dispatch.hpp @@ -1,7 +1,6 @@ #pragma once #include -#include #include "rabitqlib/utils/space.hpp" diff --git a/include/rabitqlib/simd/fastscan_dispatch.hpp b/include/rabitqlib/simd/fastscan_dispatch.hpp index 8bf44f5..835d2f0 100644 --- a/include/rabitqlib/simd/fastscan_dispatch.hpp +++ b/include/rabitqlib/simd/fastscan_dispatch.hpp @@ -5,6 +5,10 @@ namespace rabitqlib::fastscan::simd { +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); + void accumulate_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lp_table, diff --git a/include/rabitqlib/utils/array.hpp b/include/rabitqlib/utils/array.hpp index cdc7f42..ee23b4c 100644 --- a/include/rabitqlib/utils/array.hpp +++ b/include/rabitqlib/utils/array.hpp @@ -20,8 +20,12 @@ #pragma once #include +#include #include +#include +#include #include +#include #include "rabitqlib/utils/memory.hpp" diff --git a/include/rabitqlib/utils/buffer.hpp b/include/rabitqlib/utils/buffer.hpp index 786a6fc..f06da97 100644 --- a/include/rabitqlib/utils/buffer.hpp +++ b/include/rabitqlib/utils/buffer.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include diff --git a/include/rabitqlib/utils/io.hpp b/include/rabitqlib/utils/io.hpp index 7646963..bf70e91 100644 --- a/include/rabitqlib/utils/io.hpp +++ b/include/rabitqlib/utils/io.hpp @@ -1,8 +1,11 @@ #pragma once #include +#include +#include #include #include +#include #include #include #include diff --git a/include/rabitqlib/utils/memory.hpp b/include/rabitqlib/utils/memory.hpp index 8dced55..93cd490 100644 --- a/include/rabitqlib/utils/memory.hpp +++ b/include/rabitqlib/utils/memory.hpp @@ -5,8 +5,8 @@ #endif #include +#include #include -#include #include #include #include diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 5622240..f0d6b27 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -8,6 +9,8 @@ #include #include #include +#include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/simd/rotator_dispatch.hpp" diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 8a47d4a..1891fec 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -2,14 +2,16 @@ #include -#include -#include +#include +#include +#include #include #include -#include #include #include +#include #include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/simd/space_dispatch.hpp" diff --git a/include/rabitqlib/utils/stopw.hpp b/include/rabitqlib/utils/stopw.hpp index b0864bb..90f5c8d 100644 --- a/include/rabitqlib/utils/stopw.hpp +++ b/include/rabitqlib/utils/stopw.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include namespace rabitqlib { class StopW { diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index 15ba9a0..57fb527 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -1,10 +1,14 @@ #pragma once -#include +#include +#include +#include +#include #include #include #include #include +#include namespace rabitqlib { template diff --git a/include/rabitqlib/utils/visited_pool.hpp b/include/rabitqlib/utils/visited_pool.hpp index 2926d3d..27abf56 100644 --- a/include/rabitqlib/utils/visited_pool.hpp +++ b/include/rabitqlib/utils/visited_pool.hpp @@ -1,4 +1,5 @@ #pragma once +#include #include #include diff --git a/include/rabitqlib/utils/visited_set_epoch.hpp b/include/rabitqlib/utils/visited_set_epoch.hpp index 04152a9..c32e573 100644 --- a/include/rabitqlib/utils/visited_set_epoch.hpp +++ b/include/rabitqlib/utils/visited_set_epoch.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include diff --git a/include/rabitqlib/utils/visited_set_hash.hpp b/include/rabitqlib/utils/visited_set_hash.hpp index 0692f34..d6aca33 100644 --- a/include/rabitqlib/utils/visited_set_hash.hpp +++ b/include/rabitqlib/utils/visited_set_hash.hpp @@ -20,8 +20,7 @@ #include #include -#include -#include +#include #include #include #include diff --git a/scripts/check-includes.sh b/scripts/check-includes.sh new file mode 100755 index 0000000..1075f9a --- /dev/null +++ b/scripts/check-includes.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +cd "$repo_root" +export CLANG_TIDY="${CLANG_TIDY:-clang-tidy}" +include_build="$(realpath "${1:-build-includes}")" +export include_build +export include_root="$repo_root" + +"$CLANG_TIDY" --version +if ! "$CLANG_TIDY" --checks='-*,misc-include-cleaner' --list-checks \ + | grep -q 'misc-include-cleaner'; then + echo 'error: misc-include-cleaner requires clang-tidy 17 or newer' >&2 + exit 1 +fi +if [[ ! -f "$include_build/compile_commands.json" ]]; then + echo "error: configure a compilation database in $include_build first" >&2 + exit 1 +fi + +check_file() { + local file="$1" + local status=0 + local report + report="$(mktemp)" + # Isolate this check from the project's general clang-tidy configuration. + # Vendored snapshots lack export annotations on their public headers. + # Do not replace those headers with Eigen/hnswlib implementation includes. + local args=(--config='{CheckOptions: {misc-include-cleaner.IgnoreHeaders: "rabitqlib/third/Eigen/src/.*;rabitqlib/third/hnswlib/(hnswalg|space_l2)[.]h"}}' --checks='-*,misc-include-cleaner' + --warnings-as-errors='misc-include-cleaner' --quiet) + if [[ "$file" == src/*.cpp ]]; then + "$CLANG_TIDY" "${args[@]}" -p "$include_build" "$file" >"$report" 2>&1 || status=1 + else + local flags=(-x c++ -std=c++17 -fopenmp -I"$include_root/include" + -Wno-pragma-once-outside-header) + if [[ "$file" == src/* ]]; then + # Private kernels share AVX2 and AVX-512 implementations. Check both + # preprocessor paths, including the popcount-enabled path. + if [[ "$file" != *avx512* ]]; then + "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" \ + -mavx2 -mfma >"$report" 2>&1 || status=1 + fi + if [[ "$file" != *avx2* ]]; then + "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" \ + -mavx2 -mfma -mavx512f -mavx512bw -mavx512dq \ + -mavx512vpopcntdq >>"$report" 2>&1 || status=1 + fi + else + "$CLANG_TIDY" "${args[@]}" "$file" -- "${flags[@]}" >"$report" 2>&1 || status=1 + fi + fi + echo "Checking $file" + cat "$report" + rm -f "$report" + return "$status" +} +export -f check_file + +# The child shell expands its positional argument. +# shellcheck disable=SC2016 +git ls-files -z -- 'src/*.cpp' 'src/*.hpp' 'include/rabitqlib/*.hpp' \ + ':(exclude)include/rabitqlib/third/**' \ + ':(exclude)include/rabitqlib/utils/fht_avx.hpp' \ + | xargs -0 -r -n 1 -P "${INCLUDE_JOBS:-2}" bash -c 'check_file "$1"' _ diff --git a/src/index/hnsw_search_avx2.cpp b/src/index/hnsw_search_avx2.cpp index cb2b920..e5b5873 100644 --- a/src/index/hnsw_search_avx2.cpp +++ b/src/index/hnsw_search_avx2.cpp @@ -1,4 +1,9 @@ +#include +#include +#include + #include "hnsw_search_avx2_kernels.hpp" +#include "rabitqlib/defines.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index 4206857..a1dfb37 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -7,7 +7,6 @@ #include #include "rabitqlib/index/query.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_core.cpp b/src/index/hnsw_search_avx512_core.cpp index 07960f3..25ff182 100644 --- a/src/index/hnsw_search_avx512_core.cpp +++ b/src/index/hnsw_search_avx512_core.cpp @@ -1,5 +1,10 @@ +#include +#include +#include + #include "hnsw_search_avx2_kernels.hpp" #include "hnsw_search_avx512_kernels.hpp" +#include "rabitqlib/defines.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_kernels.hpp b/src/index/hnsw_search_avx512_kernels.hpp index f6282ba..f277357 100644 --- a/src/index/hnsw_search_avx512_kernels.hpp +++ b/src/index/hnsw_search_avx512_kernels.hpp @@ -6,7 +6,7 @@ #include #include "rabitqlib/index/query.hpp" -#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/space.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_popcnt.cpp b/src/index/hnsw_search_avx512_popcnt.cpp index 3af1e92..8330642 100644 --- a/src/index/hnsw_search_avx512_popcnt.cpp +++ b/src/index/hnsw_search_avx512_popcnt.cpp @@ -1,4 +1,9 @@ +#include +#include +#include + #include "hnsw_search_avx512_kernels.hpp" +#include "rabitqlib/defines.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 5e17b8b..5abcef1 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -1,10 +1,13 @@ #include "rabitqlib/simd/dispatch.hpp" -#include -#include +#include +#include #include #include +#include "rabitqlib/defines.hpp" +#include "rabitqlib/fastscan/fastscan.hpp" +#include "rabitqlib/fastscan/highacc_fastscan.hpp" #include "rabitqlib/simd/fastscan_dispatch.hpp" #include "rabitqlib/simd/pack_excode_dispatch.hpp" #include "rabitqlib/simd/quantization_dispatch.hpp" @@ -12,6 +15,8 @@ #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/simd/warmup_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/warmup_space.hpp" #include "rescale_search.hpp" namespace rabitqlib::simd { @@ -415,6 +420,27 @@ float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) namespace rabitqlib::fastscan { +void simd::pack_lut_generic(size_t dim, const float* query, float* lut) { + for (size_t group = 0; group < dim / 4; ++group) { + lut[0] = 0; + for (size_t j = 1; j < 16; ++j) { + lut[j] = lut[j - LOWBIT(j)] + query[kPos[j]]; + } + query += 4; + lut += 16; + } +} + +using PackLutFn = void (*)(size_t, const float*, float*); +const PackLutFn kPackLutFn = cpu::has_avx512_core() ? simd::pack_lut_avx512 + : cpu::has_avx2() ? simd::pack_lut_avx2 + : simd::pack_lut_generic; + +template <> +void pack_lut(size_t dim, const float* __restrict__ query, float* __restrict__ lut) { + kPackLutFn(dim, query, lut); +} + using AccumulateFn = void (*)(const uint8_t*, const uint8_t*, uint16_t*, size_t); const AccumulateFn kAccumulateFn = [] { if (cpu::has_avx512_core()) { diff --git a/src/simd/fastscan_avx2.cpp b/src/simd/fastscan_avx2.cpp index 5d2e5c0..9df68f4 100644 --- a/src/simd/fastscan_avx2.cpp +++ b/src/simd/fastscan_avx2.cpp @@ -1,12 +1,33 @@ #include +#include #include -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/simd/fastscan_dispatch.hpp" namespace rabitqlib::fastscan::simd { +void pack_lut_avx2(size_t dim, const float* query, float* lut) { + for (size_t group = 0; group < dim / 4; ++group) { + __m256 lo = _mm256_setzero_ps(); + __m256 hi = _mm256_add_ps(lo, _mm256_set1_ps(query[0])); + const __m256 q1 = _mm256_set1_ps(query[1]); + const __m256 q2 = _mm256_set1_ps(query[2]); + const __m256 q3 = _mm256_set1_ps(query[3]); + // Select additions by subset bits, in the same order as the generic loop. + lo = _mm256_blend_ps(lo, _mm256_add_ps(lo, q1), 0xF0); + hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q1), 0xF0); + lo = _mm256_blend_ps(lo, _mm256_add_ps(lo, q2), 0xCC); + hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q2), 0xCC); + lo = _mm256_blend_ps(lo, _mm256_add_ps(lo, q3), 0xAA); + hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q3), 0xAA); + _mm256_storeu_ps(lut, lo); + _mm256_storeu_ps(lut + 8, hi); + query += 4; + lut += 16; + } +} + void accumulate_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lp_table, diff --git a/src/simd/fastscan_avx512.cpp b/src/simd/fastscan_avx512.cpp index 19be201..bf01939 100644 --- a/src/simd/fastscan_avx512.cpp +++ b/src/simd/fastscan_avx512.cpp @@ -1,12 +1,27 @@ #include +#include #include -#include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/simd/fastscan_dispatch.hpp" namespace rabitqlib::fastscan::simd { +void pack_lut_avx512(size_t dim, const float* query, float* lut) { + for (size_t group = 0; group < dim / 4; ++group) { + // Lane n represents the subset selected by the four bits of n. + // Masked additions preserve both the coordinate order and the initial +0. + __m512 values = _mm512_setzero_ps(); + values = _mm512_mask_add_ps(values, 0xFF00, values, _mm512_set1_ps(query[0])); + values = _mm512_mask_add_ps(values, 0xF0F0, values, _mm512_set1_ps(query[1])); + values = _mm512_mask_add_ps(values, 0xCCCC, values, _mm512_set1_ps(query[2])); + values = _mm512_mask_add_ps(values, 0xAAAA, values, _mm512_set1_ps(query[3])); + _mm512_storeu_ps(lut, values); + query += 4; + lut += 16; + } +} + void accumulate_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lp_table, diff --git a/src/simd/pack_excode_avx2.cpp b/src/simd/pack_excode_avx2.cpp index 1b8c1ef..6bb88bb 100644 --- a/src/simd/pack_excode_avx2.cpp +++ b/src/simd/pack_excode_avx2.cpp @@ -1,3 +1,6 @@ +#include +#include + #include "pack_excode_kernels.hpp" #include "rabitqlib/simd/pack_excode_dispatch.hpp" diff --git a/src/simd/pack_excode_avx512.cpp b/src/simd/pack_excode_avx512.cpp index 2f2f1b0..a3b6cb9 100644 --- a/src/simd/pack_excode_avx512.cpp +++ b/src/simd/pack_excode_avx512.cpp @@ -1,3 +1,6 @@ +#include +#include + #include "pack_excode_kernels.hpp" #include "rabitqlib/simd/pack_excode_dispatch.hpp" diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index 22dbe19..9e96abd 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -1,7 +1,6 @@ #pragma once #include -#include #include #include #include diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index d21d576..a8902d4 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -1,5 +1,6 @@ #include +#include #include #include diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index 8dc84db..b2feac5 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -1,5 +1,6 @@ #include +#include #include #include diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index be19cb2..45595b9 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -1,9 +1,11 @@ #include #include +#include #include #include +#include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/utils/space.hpp" #include "space_float_kernels.hpp" diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index ac812e6..abdb78e 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -1,8 +1,10 @@ #include #include +#include #include +#include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/utils/space.hpp" #include "space_float_kernels.hpp" diff --git a/src/simd/space_float.cpp b/src/simd/space_float.cpp index 129e94a..64674df 100644 --- a/src/simd/space_float.cpp +++ b/src/simd/space_float.cpp @@ -1,3 +1,5 @@ +#include + #include "rabitqlib/simd/space_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index d31e381..b9d5fd8 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -2,8 +2,8 @@ #include #include -#include -#include + +#include "rabitqlib/simd/warmup_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/warmup_avx512.cpp b/src/simd/warmup_avx512.cpp index a94117b..4333def 100644 --- a/src/simd/warmup_avx512.cpp +++ b/src/simd/warmup_avx512.cpp @@ -2,8 +2,8 @@ #include #include -#include -#include + +#include "rabitqlib/simd/warmup_dispatch.hpp" namespace rabitqlib::simd { diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp new file mode 100644 index 0000000..e176276 --- /dev/null +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -0,0 +1,170 @@ +#include "rabitqlib/fastscan/fastscan.hpp" + +#include + +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/fastscan_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::fastscan { +namespace { + +template +void check_lut_subset_sums(void (*pack_fn)(size_t, const T*, T*) = pack_lut) { + const T large = std::numeric_limits::max() / 4; + const T tiny = std::numeric_limits::denorm_min(); + const std::array, 5> patterns = { + {{1, 2, 4, 8}, + {-0.0, -0.0, -0.0, -0.0}, + {large, tiny, -large, 1}, + {tiny, tiny, -tiny, -tiny}, + {-3.5, 0.25, 2, -0.125}}}; + for (size_t dim : {0U, 4U, 8U, 12U, 16U, 64U, 128U, 768U}) { + for (size_t pattern = 0; pattern < patterns.size(); ++pattern) { + SCOPED_TRACE(::testing::Message() << "dim=" << dim << " pattern=" << pattern); + // Offset by one element to avoid requiring SIMD-aligned pointers. + std::vector query(dim + 1); + for (size_t d = 0; d < dim; ++d) { + query[d + 1] = patterns[(pattern + d / 4) % patterns.size()][d % 4]; + } + const auto original = query; + std::vector expected(dim * 4 + 2, T{123}); + std::vector actual = expected; + for (size_t group = 0; group < dim / 4; ++group) { + for (size_t mask = 0; mask < 16; ++mask) { + T sum = 0; + // Independently enumerate selected coordinates in input order. + for (size_t d = 0; d < 4; ++d) { + if ((mask & (8U >> d)) != 0) { + sum += query[1 + group * 4 + d]; + } + } + expected[1 + group * 16 + mask] = sum; + } + } + pack_fn(dim, query.data() + 1, actual.data() + 1); + // Compare bits to detect signed-zero changes as well as rounding changes. + EXPECT_EQ( + std::memcmp(actual.data(), expected.data(), actual.size() * sizeof(T)), 0 + ); + EXPECT_EQ( + std::memcmp(query.data(), original.data(), query.size() * sizeof(T)), 0 + ); + } + } +} + +TEST(FastScanLutTest, MatchesOrderedSubsetSumsExactly) { + check_lut_subset_sums(); + check_lut_subset_sums(); +} + +TEST(FastScanLutTest, EverySupportedBackendMatchesOrderedSubsetSumsExactly) { + check_lut_subset_sums(simd::pack_lut_generic); + if (cpu::has_avx2()) { + SCOPED_TRACE("AVX2"); + check_lut_subset_sums(simd::pack_lut_avx2); + } + if (cpu::has_avx512_core()) { + SCOPED_TRACE("AVX512"); + check_lut_subset_sums(simd::pack_lut_avx512); + } +} + +TEST(FastScanPackingTest, MatchesBitReferenceIncludingTailsAndUnalignedBuffers) { + for (size_t dim : {8U, 16U, 24U, 56U, 64U, 72U, 128U, 768U}) { + for (size_t num = 0; num <= 65; ++num) { + SCOPED_TRACE(::testing::Message() << "dim=" << dim << " num=" << num); + std::vector input(num * dim / 8 + 1); + uint32_t state = 42; + for (auto& value : input) { + state = state * 1664525U + 1013904223U; + value = static_cast(state >> 24); + } + const auto original = input; + const size_t bytes = ((num + 31) / 32) * 32 * (dim / 8); + std::vector expected(bytes + 2, 0); + expected.front() = expected.back() = 0xA5; + // Scatter individual sign bits to the documented FastScan lane layout. + // This reference neither gathers columns nor uses the packing permutation. + for (size_t row = 0; row < num; ++row) { + const size_t lane = 2 * (row % 8) + (row % 16) / 8; + for (size_t d = 0; d < dim; ++d) { + const uint8_t bit = + (input[1 + row * dim / 8 + d / 8] >> (7 - d % 8)) & 1; + const size_t offset = (row / 32) * (32 * dim / 8) + (d / 4) * 16 + lane; + expected[1 + offset] |= bit << (3 - d % 4 + 4 * ((row % 32) / 16)); + } + } + std::vector actual(bytes + 2, 0xA5); + pack_codes(dim, input.data() + 1, num, actual.data() + 1); + EXPECT_EQ(actual, expected); + EXPECT_EQ(input, original); + } + } +} + +TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) { + constexpr size_t kMaxDim = 768; + for (size_t dim : {16U, 64U, 128U, 768U}) { + for (size_t num : {1U, 17U, 31U, 32U}) { + SCOPED_TRACE(::testing::Message() << "dim=" << dim << " num=" << num); + std::vector input(num * dim / 8); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(i * 73 + i / 7); + } + std::vector packed(32 * dim / 8); + pack_codes(dim, input.data(), num, packed.data()); + std::vector lut(dim * 4); + std::vector lut_hacc(dim * 4); + for (size_t i = 0; i < lut.size(); ++i) { + lut[i] = static_cast((i * 7 + i / 16) % 32); + lut_hacc[i] = static_cast((i * 37 + i / 16) % 4096); + } + std::array expected{}; + std::array expected_hacc{}; + for (size_t row = 0; row < 32; ++row) { + for (size_t group = 0; group < dim / 4; ++group) { + const uint8_t byte = row < num ? input[row * dim / 8 + group / 2] : 0; + const size_t code = (byte >> (group % 2 == 0 ? 4 : 0)) & 15; + expected[row] += lut[group * 16 + code]; + expected_hacc[row] += lut_hacc[group * 16 + code]; + } + } + auto check_backend = [&](auto accumulate_fn, auto transfer_fn, auto hacc_fn) { + std::array actual{}; + std::array actual_hacc{}; + alignas(64) std::array packed_lut{}; + accumulate_fn(packed.data(), lut.data(), actual.data(), dim); + transfer_fn(lut_hacc.data(), dim, packed_lut.data()); + hacc_fn(packed.data(), packed_lut.data(), actual_hacc.data(), dim); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_hacc, expected_hacc); + }; + if (cpu::has_avx2()) { + SCOPED_TRACE("AVX2"); + check_backend( + simd::accumulate_avx2, + simd::transfer_lut_hacc_avx2, + simd::accumulate_hacc_avx2 + ); + } + if (cpu::has_avx512_core()) { + SCOPED_TRACE("AVX512"); + check_backend( + simd::accumulate_avx512, + simd::transfer_lut_hacc_avx512, + simd::accumulate_hacc_avx512 + ); + } + } + } +} + +} // namespace +} // namespace rabitqlib::fastscan diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index d5cd5be..bef869a 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -14,6 +14,85 @@ namespace rabitqlib::quant { namespace { +template +void check_one_bit_batches() { + // Alternate dimensions to exercise scratch growth and shrinking on each thread. + for (size_t dim : {64U, 8U, 128U, 16U, 768U, 24U, 192U, 56U, 72U, 64U}) { + for (size_t num : {0U, 1U, 15U, 16U, 17U, 31U, 32U, 33U, 64U, 65U}) { + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE( + ::testing::Message() + << "dim=" << dim << " num=" << num << " metric=" << metric + ); + std::vector centroid(dim); + std::vector data(num * dim); + for (size_t d = 0; d < dim; ++d) { + centroid[d] = static_cast(static_cast(d % 7) - 3) / 8; + for (size_t i = 0; i < num; ++i) { + const T residual = + i % 3 == 0 + ? T{0} + : static_cast(static_cast((i + d) % 11) - 5) / 4; + data[i * dim + d] = centroid[d] + residual; + } + } + // A guard byte/factor also verifies empty inputs and tail writes. + const size_t packed_bytes = ((num + 31) / 32) * 32 * (dim / 8); + std::vector expected(packed_bytes + 1, 0xA5); + std::vector actual = expected; + std::vector compact(num * dim / 8); + std::vector expected_add(num + 1, T{123}); + std::vector expected_rescale = expected_add; + std::vector expected_error = expected_add; + std::vector actual_add = expected_add; + std::vector actual_rescale = expected_add; + std::vector actual_error = expected_add; + for (size_t i = 0; i < num; ++i) { + rabitq_impl::one_bit::one_bit_compact_code( + data.data() + i * dim, + centroid.data(), + dim, + compact.data() + i * dim / 8, + expected_add[i], + expected_rescale[i], + expected_error[i], + metric + ); + } + fastscan::pack_codes(dim, compact.data(), num, expected.data()); + rabitq_impl::one_bit::one_bit_batch_code( + data.data(), + centroid.data(), + num, + dim, + actual.data(), + actual_add.data(), + actual_rescale.data(), + actual_error.data(), + metric + ); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_add, expected_add); + EXPECT_EQ(actual_rescale, expected_rescale); + EXPECT_EQ(actual_error, expected_error); + } + } + } +} + +TEST(RabitqOneBitBatchTest, MatchesIndividualEncoding) { + check_one_bit_batches(); + check_one_bit_batches(); +} + +TEST(RabitqOneBitBatchTest, ParallelMatchesIndividualEncoding) { + const int previous_threads = omp_get_max_threads(); + omp_set_num_threads(2); + check_one_bit_batches(); + check_one_bit_batches(); + omp_set_num_threads(previous_threads); +} + int level_from_thresholds(double magnitude, double t, int max_code) { int result = 0; if (magnitude > 0) {