diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..5503dd8 --- /dev/null +++ b/.clang-format @@ -0,0 +1,33 @@ +--- +BasedOnStyle: Google +IndentWidth: 4 +--- +Language: Cpp +AlignAfterOpenBracket: BlockIndent +AlignEscapedNewlines: Left +AlignTrailingComments: true +AllowAllArgumentsOnNextLine: true +AllowShortBlocksOnASingleLine: Always +AllowShortCaseLabelsOnASingleLine: false +AllowShortEnumsOnASingleLine: true +AllowShortFunctionsOnASingleLine: All +AllowShortIfStatementsOnASingleLine: Never +AllowShortLambdasOnASingleLine: All +AllowShortLoopsOnASingleLine: false +BinPackArguments: false +BinPackParameters: false +BreakConstructorInitializers: BeforeComma +ColumnLimit: 92 +CompactNamespaces: false +DerivePointerAlignment: false +FixNamespaceComments: true +IncludeBlocks: Regroup +IndentCaseLabels: true +KeepEmptyLinesAtTheStartOfBlocks: false +PointerAlignment: Left +ReferenceAlignment: Left +ReflowComments: true +SortIncludes: CaseSensitive +Standard: c++17 +UseTab: Never +--- diff --git a/.clang-tidy b/.clang-tidy new file mode 100644 index 0000000..5ba7be3 --- /dev/null +++ b/.clang-tidy @@ -0,0 +1,30 @@ +Checks: > + -*, + clang-analyzer-*, + bugprone-branch-clone, + bugprone-infinite-loop, + bugprone-macro-parentheses, + bugprone-macro-repeated-side-effects, + bugprone-narrowing-conversions, + bugprone-sizeof-expression, + bugprone-suspicious-*, + bugprone-use-after-move, + misc-redundant-expression, + misc-static-assert, + misc-throw-by-value-catch-by-reference, + modernize-use-noexcept, + modernize-use-nullptr, + modernize-use-override, + modernize-use-using, + performance-*, + portability-*, + readability-duplicate-include, + readability-inconsistent-declaration-parameter-name, + readability-misleading-indentation, + readability-redundant-control-flow, + -portability-simd-intrinsics, + -clang-analyzer-security.insecureAPI.DeprecatedOrUnsafeBufferHandling + +# Keep the selected first-party baseline regression-free. Repository-wide +# naming checks remain deferred because they would change public APIs. +WarningsAsErrors: "*" diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..42b6e56 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,27 @@ +root = true + +[*] +charset = utf-8 +end_of_line = lf +insert_final_newline = true +trim_trailing_whitespace = true + +[*.{c,cc,cpp,cxx,h,hpp,cu,cuh,cmake,py,pyi}] +indent_style = space +indent_size = 4 + +[CMakeLists.txt] +indent_style = space +indent_size = 4 + +[*.{json,yaml,yml}] +indent_style = space +indent_size = 2 + +[*.md] +indent_style = space +indent_size = 2 +trim_trailing_whitespace = false + +[Makefile] +indent_style = tab diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 957a985..1817a4c 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -11,6 +11,21 @@ concurrency: cancel-in-progress: true jobs: + python-quality: + name: Python formatting and linting + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + - name: Install Ruff + run: python -m pip install "ruff==0.16.1" + - name: Check Python + run: ./scripts/check-python.sh + python-build: name: Python ${{ matrix.python-version }} runs-on: ubuntu-latest diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index e04600c..b3b6f1f 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -6,6 +6,58 @@ concurrency: group: ${{ github.head_ref || github.sha }}-${{ github.workflow }} cancel-in-progress: true jobs: + format: + name: C++ formatting + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - name: Install clang-format + run: | + sudo apt-get update -qq + sudo apt-get install -y clang-format-15 + - name: Check formatting + run: ./scripts/check-format.sh + + clang-tidy: + name: C++ static analysis + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install dependencies + run: | + sudo apt-get update -qq + sudo apt-get install -y clang-tidy-15 libomp-15-dev cmake ninja-build + python -m pip install "numpy>=1.23" "pybind11>=2.12" + - name: Locate pybind11 + run: echo "pybind11_DIR=$(python -m pybind11 --cmakedir)" >> "$GITHUB_ENV" + - name: Configure compilation database + run: >- + cmake -S . -B build-tidy -G Ninja + -DRABITQ_BUILD_TESTS=ON + -DRABITQ_BUILD_PYTHON_BINDINGS=ON + -DCMAKE_BUILD_TYPE=Release + -Dpybind11_DIR="$pybind11_DIR" + - name: Run clang-tidy + run: ./scripts/check-tidy.sh build-tidy + + shellcheck: + name: Shell scripts + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - name: Install ShellCheck + run: | + sudo apt-get update -qq + sudo apt-get install -y shellcheck + - name: Check scripts + run: shellcheck scripts/*.sh + ubuntu: name: Ubuntu runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 7ae9cca..80062aa 100644 --- a/.gitignore +++ b/.gitignore @@ -42,6 +42,8 @@ results/ *.pyc __pycache__/ .clang-* +!.clang-format +!.clang-tidy .clangd* # macOS garbages diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..369ae56 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,9 @@ +repos: + - repo: local + hooks: + - id: clang-format + name: clang-format 15 + entry: scripts/apply-format.sh + language: system + files: '\.(c|cc|cpp|cxx|h|hpp|cu|cuh)$' + exclude: '^(include/rabitqlib/third/|include/rabitqlib/utils/fht_avx\.hpp$)' diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..5e6d9bd --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,121 @@ +# Contributing to RaBitQ + +Thank you for contributing to RaBitQ. Before submitting a pull request, build +the library, run the relevant tests, and check the C++ formatting. + +## C++ formatting + +RaBitQ uses the repository's `.clang-format` configuration and clang-format +15. Install that version on Ubuntu or Debian with: + +```bash +sudo apt-get install clang-format-15 +``` + +Format all project-maintained C and C++ files: + +```bash +./scripts/apply-format.sh +``` + +Verify formatting without changing files: + +```bash +./scripts/check-format.sh +``` + +The scripts intentionally exclude vendored Eigen code and the imported FFHT +implementation. To use a nonstandard executable name, set `CLANG_FORMAT`; it +must still identify itself as clang-format 15. + +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, +disable format-on-save and run the repository scripts before submitting. + +To format only lines changed in the staged files, use: + +```bash +./scripts/format-changed.sh --staged +``` + +The complete-file formatter remains useful before the initial formatting pass +or after changing `.clang-format`; CI always checks complete files. + +## Optional pre-commit hook + +Install [pre-commit](https://pre-commit.com/) and enable the repository hook: + +```bash +python -m pip install pre-commit +pre-commit install +``` + +The hook formats only staged C and C++ files. CI runs the read-only formatting +check over the complete project-maintained source tree. + +## Static analysis + +clang-tidy performs semantic checks and is kept separate from clang-format. +The required baseline contains focused correctness, portability, and +performance checks. Install the pinned analyzer and the dependencies needed to +configure every first-party target: + +```bash +sudo apt-get install clang-tidy-15 libomp-15-dev cmake ninja-build +python -m pip install "numpy>=1.23" "pybind11>=2.12" +``` + +Then configure the same tests and Python bindings analyzed by CI and run the +check with the same compiler used by CMake: + +```bash +export pybind11_DIR="$(python -m pybind11 --cmakedir)" +CXX=c++ cmake -S . -B build-tidy -G Ninja \ + -DRABITQ_BUILD_TESTS=ON \ + -DRABITQ_BUILD_PYTHON_BINDINGS=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -Dpybind11_DIR="$pybind11_DIR" +CXX=c++ ./scripts/check-tidy.sh build-tidy +``` + +The wrapper supplies clang-tidy with that compiler's standard-library include +paths and reports diagnostics only for first-party files. Vendored Eigen, +hnswlib, and the imported FFHT implementation are excluded. New checks should +be added incrementally after their existing first-party findings are fixed. + +## Python formatting and linting + +Python sources, examples, and tests use Ruff 0.16.1: + +```bash +python -m pip install "ruff==0.16.1" +./scripts/check-python.sh +``` + +To apply Python formatting and safe automatic lint fixes before running the +check: + +```bash +ruff check --fix python python_bindings sample/python tests/python +ruff format python python_bindings sample/python tests/python +``` + +## Shell scripts + +Run ShellCheck after changing a contributor or automation script: + +```bash +sudo apt-get install shellcheck +shellcheck scripts/*.sh +``` + +## Performance and compatibility + +- Use fixed-width integer types for serialized values and persisted index data. +- 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. +- Benchmark allocations or algorithm changes in search and quantization hot + paths, and include the commands and results in the pull request. +- Avoid unrelated refactoring or formatting in performance-sensitive changes. diff --git a/README.md b/README.md index bbde266..4a6a32c 100644 --- a/README.md +++ b/README.md @@ -127,6 +127,11 @@ GoogleTest is downloaded during test configuration. For a full benchmark on the GIST dataset, see [`example.sh`](example.sh). More detailed API and algorithm guidance is available in the [documentation](docs/docs/index.md). +## Contributing + +Contributions are welcome. See the [contributing guide](CONTRIBUTING.md) for +the build, formatting, pre-commit, and static-analysis workflows. + ## Why RaBitQ? - **High accuracy with tiny codes.** RaBitQ provides state-of-the-art diff --git a/include/rabitqlib/defines.hpp b/include/rabitqlib/defines.hpp index dc61d9e..5ee8178 100644 --- a/include/rabitqlib/defines.hpp +++ b/include/rabitqlib/defines.hpp @@ -4,7 +4,7 @@ #include "rabitqlib/third/Eigen/Dense" -#define BIT_ID(x) (__builtin_popcount((x) - 1)) +#define BIT_ID(x) (__builtin_popcount((x)-1)) #define LOWBIT(x) ((x) & (-(x))) namespace rabitqlib { diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 6ec7ead..6d7e30c 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -37,9 +37,23 @@ constexpr static std::array kPos = { }; // all possible combination for a 4 bit string constexpr static std::array kPerm0 = { - 0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15 -}; // data order of packed quantization code, please refer to code and the link offered - // above for detailed information + 0, + 8, + 1, + 9, + 2, + 10, + 3, + 11, + 4, + 12, + 5, + 13, + 6, + 14, + 7, + 15}; // data order of packed quantization code, please refer to code and the link + // offered above for detailed information template static inline void get_column( diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 7e290cc..7c8eaf4 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -39,9 +39,7 @@ class HierarchicalNSW; namespace detail { -maxheap> search_knn_avx2( - HierarchicalNSW&, const float*, size_t -); +maxheap> search_knn_avx2(HierarchicalNSW&, const float*, size_t); maxheap> search_knn_avx512_core( HierarchicalNSW&, const float*, size_t @@ -55,7 +53,7 @@ maxheap> search_knn_avx512_popcnt( class HierarchicalNSW { public: - explicit HierarchicalNSW() {}; + explicit HierarchicalNSW(){}; explicit HierarchicalNSW( size_t, size_t, size_t, size_t, size_t, size_t = 100, MetricType = METRIC_L2 ); @@ -318,14 +316,13 @@ class HierarchicalNSW { // ANN Search template - void get_bin_est_direct( - std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord& - ); + void + get_bin_est_direct(std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord&); template - void get_full_est_direct( - std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord& - ) const; + void + get_full_est_direct(std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord&) + const; maxheap> search_knn(const float*, size_t); @@ -375,9 +372,7 @@ inline HierarchicalNSW::HierarchicalNSW( , label_op_locks_(kMaxLabelOperationLock) , link_list_locks_(max_elements) , element_levels_(max_elements) - , raw_dist_func_( - (metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr - ) { + , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) { max_elements_ = max_elements; dim_ = dim; rotator_ = choose_rotator( @@ -492,12 +487,12 @@ inline void HierarchicalNSW::save(const char* filename) const { output.write( reinterpret_cast(centroids_memory_), - num_cluster_ * padded_dim_ * sizeof(float) + static_cast(num_cluster_ * padded_dim_ * sizeof(float)) ); output.write( reinterpret_cast(data_level0_memory_), - cur_element_count_ * size_data_per_element_ + static_cast(cur_element_count_ * size_data_per_element_) ); for (size_t i = 0; i < cur_element_count_; i++) { @@ -553,15 +548,25 @@ inline void HierarchicalNSW::load(const char* filename) { input.read(reinterpret_cast(&mult_), sizeof(double)); input.read(reinterpret_cast(&ef_construction_), sizeof(size_t)); - centroids_memory_ = - reinterpret_cast(malloc(num_cluster_ * padded_dim_ * sizeof(float))); + const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); + centroids_memory_ = reinterpret_cast(malloc(centroids_bytes)); + if (centroids_memory_ == nullptr) { + throw std::runtime_error("Not enough memory: loadIndex failed to allocate centroids" + ); + } - input.read(centroids_memory_, num_cluster_ * padded_dim_ * sizeof(float)); + input.read( + centroids_memory_, + static_cast(num_cluster_ * padded_dim_ * sizeof(float)) + ); data_level0_memory_ = reinterpret_cast(malloc(max_elements_ * size_data_per_element_)); - input.read(data_level0_memory_, cur_element_count_ * size_data_per_element_); + input.read( + data_level0_memory_, + static_cast(cur_element_count_ * size_data_per_element_) + ); std::cout << "cur_element_count = " << cur_element_count_ << '\n'; @@ -570,8 +575,7 @@ inline void HierarchicalNSW::load(const char* filename) { linkLists_ = reinterpret_cast(malloc(sizeof(void*) * max_elements_)); if (linkLists_ == nullptr) { - throw std::runtime_error( - "Not enough memory: loadIndex failed to allocate linklists" + throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklists" ); } @@ -624,8 +628,8 @@ inline void HierarchicalNSW::construct( bool faster = false ) { num_cluster_ = cluster_num; - centroids_memory_ = - reinterpret_cast(malloc(num_cluster_ * padded_dim_ * sizeof(float))); + const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); + centroids_memory_ = reinterpret_cast(malloc(centroids_bytes)); if (centroids_memory_ == nullptr) { throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids"); } @@ -1233,7 +1237,9 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( float distk = 1e10; EstimateRecord start_estimate_record; - get_full_est_direct(q_to_centroids, query_wrapper, ep_id, start_estimate_record); + get_full_est_direct( + q_to_centroids, query_wrapper, ep_id, start_estimate_record + ); float est_dist = start_estimate_record.est_dist; float low_dist = start_estimate_record.low_dist; @@ -1270,13 +1276,15 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( ); } - if(vl->get(candidate_id)) { + if (vl->get(candidate_id)) { continue; } vl->set(candidate_id); EstimateRecord candest; - get_bin_est_direct(q_to_centroids, query_wrapper, candidate_id, candest); + get_bin_est_direct( + q_to_centroids, query_wrapper, candidate_id, candest + ); bool flag_update_KNNs = boundedKNN.size() < TOPK || candest.low_dist < distk; @@ -1289,8 +1297,7 @@ inline void HierarchicalNSW::searchBaseLayerST_AdaptiveRerankOptDirect( } Candidate cand{ ResultRecord(candest.est_dist, candest.low_dist), - static_cast(candidate_id) - }; + static_cast(candidate_id)}; boundedKNN.insert(cand); distk = boundedKNN.worst().record.est_dist; } diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 53b39d3..8980bf2 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -9,8 +9,8 @@ #include #include "rabitqlib/defines.hpp" -#include "rabitqlib/utils/space.hpp" #include "rabitqlib/third/hnswlib/hnswlib.h" +#include "rabitqlib/utils/space.hpp" namespace rabitqlib::ivf { template @@ -81,9 +81,8 @@ class Initializer { virtual ~Initializer() = 0; [[nodiscard]] virtual const float* centroid(PID) const = 0; virtual void add_vectors(const float*, size_t) = 0; - virtual void centroids_distances( - const float*, size_t, std::vector>& - ) const = 0; + virtual void + centroids_distances(const float*, size_t, std::vector>&) const = 0; virtual void load(std::ifstream&, const char*) = 0; virtual void save(std::ofstream&, const char*) const = 0; }; diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index a471005..55187f1 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -42,14 +42,8 @@ class IVF { MetricType metric_type_ = rabitqlib::METRIC_L2; // metric type float (*ip_func_)(const float*, const uint8_t*, size_t) = nullptr; - void quantize_cluster( - Cluster&, - const std::vector&, - const float*, - const float*, - float*, - const quant::RabitqConfig& - ); + void + quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, const quant::RabitqConfig&); [[nodiscard]] size_t ids_bytes() const { return sizeof(PID) * num_; } @@ -167,7 +161,10 @@ inline IVF::~IVF() { * @param clustter_ids Cluster ID for each data objects */ inline void IVF::construct( - const float* data, const float* centroids, const PID* cluster_ids, bool faster = false, + const float* data, + const float* centroids, + const PID* cluster_ids, + bool faster = false, size_t num_threads = std::numeric_limits::max() ) { std::cout << "Start IVF construction...\n"; @@ -521,7 +518,6 @@ inline void IVF::scan_one_batch( PID id = ids[i]; float ex_dist = est_distance[i]; knns.insert(id, ex_dist); - distk = knns.top_dist(); } return; } diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 9f91910..06d543a 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -83,8 +83,7 @@ class QuantizedGraph { } [[nodiscard]] PID* get_neighbors(PID data_id) { - return reinterpret_cast( - &data_.at((row_offset_ * data_id) + neighbor_offset_) + return reinterpret_cast(&data_.at((row_offset_ * data_id) + neighbor_offset_) ); } diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 27ffef3..229af3d 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -75,9 +75,8 @@ class QGBuilder { std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); - PID entry_point = exact_nn( - data, centroid.data(), num_nodes_, dim_, num_threads_, euclidean_sqr - ); + PID entry_point = + exact_nn(data, centroid.data(), num_nodes_, dim_, num_threads_, euclidean_sqr); std::cout << "Setting entry_point to " << entry_point << '\n' << std::flush; @@ -143,7 +142,6 @@ inline void QGBuilder::add_pruned_edges( float dik_sqr = cur.distance; if (nei_set.find(cur.id) != nei_set.end()) { - occlude = true; break; } diff --git a/include/rabitqlib/quantization/data_layout.hpp b/include/rabitqlib/quantization/data_layout.hpp index fbb2a07..ddb3b6a 100644 --- a/include/rabitqlib/quantization/data_layout.hpp +++ b/include/rabitqlib/quantization/data_layout.hpp @@ -10,8 +10,7 @@ struct BatchDataMap { public: explicit BatchDataMap(char* data, size_t padded_dim) : batch_bin_code_(reinterpret_cast(data)) - , f_add_( - reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) + , f_add_(reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) ) // 1 bit code , f_rescale_(f_add_ + fastscan::kBatchSize) , f_error_(f_rescale_ + fastscan::kBatchSize) {} @@ -38,8 +37,7 @@ struct ConstBatchDataMap { public: explicit ConstBatchDataMap(const char* data, size_t padded_dim) : batch_bin_code_(reinterpret_cast(data)) - , f_add_( - reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) + , f_add_(reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) ) // 1 bit code , f_rescale_(f_add_ + fastscan::kBatchSize) , f_error_(f_rescale_ + fastscan::kBatchSize) {} @@ -61,8 +59,7 @@ struct QGBatchDataMap { public: explicit QGBatchDataMap(char* data, size_t padded_dim) : batch_bin_code_(reinterpret_cast(data)) - , f_add_( - reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) + , f_add_(reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) ) // 1 bit code , f_rescale_(f_add_ + fastscan::kBatchSize) {} @@ -86,8 +83,7 @@ struct ConstQGBatchDataMap { public: explicit ConstQGBatchDataMap(const char* data, size_t padded_dim) : batch_bin_code_(reinterpret_cast(data)) - , f_add_( - reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) + , f_add_(reinterpret_cast(data + (padded_dim * fastscan::kBatchSize / 8)) ) // 1 bit code , f_rescale_(f_add_ + fastscan::kBatchSize) {} @@ -134,8 +130,7 @@ struct ConstExDataMap { explicit ConstExDataMap(const char* data, size_t padded_dim, size_t ex_bits) : ex_code_(reinterpret_cast(data)) , f_add_ex_(*reinterpret_cast(data + (padded_dim * ex_bits / 8))) - , f_recale_ex_( - *(reinterpret_cast(data + (padded_dim * ex_bits / 8)) + 1) + , f_recale_ex_(*(reinterpret_cast(data + (padded_dim * ex_bits / 8)) + 1) ) {} [[nodiscard]] const uint8_t* ex_code() const { return ex_code_; } diff --git a/include/rabitqlib/simd/space_dispatch.hpp b/include/rabitqlib/simd/space_dispatch.hpp index 43cbf4b..5cf08eb 100644 --- a/include/rabitqlib/simd/space_dispatch.hpp +++ b/include/rabitqlib/simd/space_dispatch.hpp @@ -6,39 +6,91 @@ namespace rabitqlib::simd { namespace excode_ipimpl { -float ip16_fxu1_avx2(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_avx2(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_avx2(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_avx2(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip16_fxu8_avx2(const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim); - -float ip16_fxu1_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu2_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu3_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip16_fxu4_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu5_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu6_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu7_avx512(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip16_fxu8_avx512(const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim); +float ip16_fxu1_avx2( + 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_avx2( + 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_avx2( + 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_avx2( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip16_fxu8_avx2( + const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim +); + +float ip16_fxu1_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu2_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu3_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip16_fxu4_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu5_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu6_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu7_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip16_fxu8_avx512( + const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim +); } // namespace excode_ipimpl -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_avx2(const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query); +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_avx2( + const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query +); float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_dim); -void scalar_quantize_uint8_avx2(uint8_t* result, const float* vec0, 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); +void scalar_quantize_uint8_avx2( + uint8_t* result, const float* vec0, 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 +); -void new_transpose_bin_avx512(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query); -void new_transpose_bin_512_avx512(const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query); +void new_transpose_bin_avx512( + const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query +); +void new_transpose_bin_512_avx512( + const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query +); float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padded_dim); -void scalar_quantize_uint8_avx512(uint8_t* result, const float* vec0, size_t dim, float lo, float delta); -void scalar_quantize_uint16_avx512(uint16_t* result, const float* vec0, size_t dim, float lo, float delta); +void scalar_quantize_uint8_avx512( + uint8_t* result, const float* vec0, size_t dim, float lo, float delta +); +void scalar_quantize_uint16_avx512( + uint16_t* result, const float* vec0, size_t dim, float lo, float delta +); -void scalar_quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, float delta); -void scalar_quantize_uint16(uint16_t* result, const float* vec0, size_t dim, float lo, float delta); +void scalar_quantize_uint8( + uint8_t* result, const float* vec0, size_t dim, float lo, float delta +); +void scalar_quantize_uint16( + uint16_t* result, const float* vec0, size_t dim, float lo, float delta +); } // namespace rabitqlib::simd diff --git a/include/rabitqlib/utils/memory.hpp b/include/rabitqlib/utils/memory.hpp index bee3bd5..fb93ee1 100644 --- a/include/rabitqlib/utils/memory.hpp +++ b/include/rabitqlib/utils/memory.hpp @@ -103,6 +103,8 @@ static inline void prefetch_l2(const void* addr) { } inline void mem_prefetch_l1(const char* ptr, size_t num_lines) { + // The repeated fallthrough branches intentionally unroll up to 20 prefetches. + // NOLINTBEGIN(bugprone-branch-clone) switch (num_lines) { default: [[fallthrough]]; @@ -184,14 +186,16 @@ inline void mem_prefetch_l1(const char* ptr, size_t num_lines) { [[fallthrough]]; case 1: prefetch_l1(ptr); - ptr += 64; [[fallthrough]]; case 0: break; } + // NOLINTEND(bugprone-branch-clone) } inline void mem_prefetch_l2(const char* ptr, size_t num_lines) { + // The repeated fallthrough branches intentionally unroll up to 20 prefetches. + // NOLINTBEGIN(bugprone-branch-clone) switch (num_lines) { default: [[fallthrough]]; @@ -273,10 +277,10 @@ inline void mem_prefetch_l2(const char* ptr, size_t num_lines) { [[fallthrough]]; case 1: prefetch_l2(ptr); - ptr += 64; [[fallthrough]]; case 0: break; } + // NOLINTEND(bugprone-branch-clone) } } // namespace rabitqlib::memory diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 7b6f306..138ba90 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -29,14 +29,14 @@ class Rotator { public: explicit Rotator() = default; - explicit Rotator(size_t dim, size_t padded_dim) : dim_(dim), padded_dim_(padded_dim) {}; + explicit Rotator(size_t dim, size_t padded_dim) : dim_(dim), padded_dim_(padded_dim){}; virtual ~Rotator() = default; virtual void rotate(const T* src, T* dst) const = 0; virtual void load(std::ifstream&) = 0; virtual void save(std::ofstream&) const = 0; // Buffer I/O - virtual void load(const char *data) = 0; - virtual void save(char *data) const = 0; // dump to buffer + virtual void load(const char* data) = 0; + virtual void save(char* data) const = 0; // dump to buffer virtual size_t dump_bytes() const = 0; [[nodiscard]] size_t size() const { return this->padded_dim_; } }; @@ -72,7 +72,7 @@ class MatrixRotator : public Rotator { std::memcpy(&rand_mat_(0, 0), &q_inv(0, 0), sizeof(T) * dim * padded_dim); } MatrixRotator() = default; - ~MatrixRotator() = default; + ~MatrixRotator() override = default; MatrixRotator& operator=(const MatrixRotator& other) { this->dim_ = other.dim_; @@ -95,17 +95,17 @@ class MatrixRotator : public Rotator { ); } - void load(const char *data) override { + void load(const char* data) override { std::memcpy(rand_mat_.data(), data, sizeof(float) * this->dim_ * this->padded_dim_); } - void save(char *data) const override { + void save(char* data) const override { std::memcpy(data, rand_mat_.data(), sizeof(float) * this->dim_ * this->padded_dim_); } size_t dump_bytes() const override { return sizeof(float) * this->dim_ * this->padded_dim_; - } + } void rotate(const T* vec, T* rotated_vec) const override { ConstRowMajorMatrixMap v(vec, 1, this->dim_); @@ -130,8 +130,9 @@ class FhtKacRotator : public Rotator { public: explicit FhtKacRotator(size_t dim, size_t padded_dim) : Rotator(dim, padded_dim), flip_(4 * padded_dim / kByteLen) { - std::random_device rd; // Seed - std::mt19937 gen(rd()); // Mersenne Twister RNG + std::random_device rd; + std::seed_seq seed{rd(), rd(), rd(), rd()}; + std::mt19937 gen(seed); // Uniform distribution in the range [0, 255] std::uniform_int_distribution dist(0, 255); @@ -188,17 +189,15 @@ class FhtKacRotator : public Rotator { ); } - void load(const char *data) override { + void load(const char* data) override { std::memcpy(flip_.data(), data, sizeof(uint8_t) * flip_.size()); } - void save(char *data) const override { + void save(char* data) const override { std::memcpy(data, flip_.data(), sizeof(uint8_t) * flip_.size()); } - size_t dump_bytes() const override { - return sizeof(uint8_t) * flip_.size(); - } + size_t dump_bytes() const override { return sizeof(uint8_t) * flip_.size(); } FhtKacRotator& operator=(const FhtKacRotator& other) { this->dim_ = other.dim_; @@ -210,9 +209,7 @@ class FhtKacRotator : public Rotator { return *this; } - static void kacs_walk(float* data, size_t len) { - simd::kacs_walk(data, len); - } + static void kacs_walk(float* data, size_t len) { simd::kacs_walk(data, len); } void rotate(const float* data, float* rotated_vec) const override { std::memcpy(rotated_vec, data, sizeof(float) * dim_); diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index a174786..c0be95c 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -219,13 +219,27 @@ inline PID exact_nn( namespace excode_ipimpl { -float ip16_fxu1_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu2_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu3_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip16_fxu4_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu5_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu6_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); -float ip64_fxu7_avx(const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim); +float ip16_fxu1_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu2_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu3_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip16_fxu4_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu5_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu6_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); +float ip64_fxu7_avx( + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim +); // inner product between float type and int type vectors template @@ -262,9 +276,7 @@ static inline uint64_t reverse_bits_u64(uint64_t n) { return n; } -void new_transpose_bin( - const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query -); +void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query); void new_transpose_bin_512( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index 2fc26bf..9dc1808 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -26,9 +26,17 @@ inline void assert_floating() { // thread save rand int template inline T rand_integer(T min, T max) { - static thread_local std::mt19937 generator( - std::random_device{}() + std::hash()(std::this_thread::get_id()) - ); + static thread_local std::mt19937 generator = [] { + std::random_device rd; + const size_t thread_hash = std::hash()(std::this_thread::get_id()); + std::seed_seq seed{ + rd(), + rd(), + static_cast(thread_hash), + static_cast(thread_hash >> 32U), + }; + return std::mt19937(seed); + }(); std::uniform_int_distribution distribution(min, max); return distribution(generator); } @@ -130,4 +138,4 @@ std::vector horizontal_avg(const std::vector>& data) { return avg; } -} // namespace rabitqlib \ No newline at end of file +} // namespace rabitqlib diff --git a/pyproject.toml b/pyproject.toml index 078bbc9..9857696 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,3 +25,10 @@ cmake.args = [ [tool.pytest.ini_options] testpaths = ["tests/python"] + +[tool.ruff] +line-length = 88 +target-version = "py39" + +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F", "I"] diff --git a/python/ivf.py b/python/ivf.py index f2b352f..dadc3c6 100644 --- a/python/ivf.py +++ b/python/ivf.py @@ -1,7 +1,8 @@ -import faiss import sys from time import time -from utils.io import read_fvecs, write_fvecs, write_ivecs, read_fbin + +import faiss +from utils.io import read_fvecs, write_fvecs, write_ivecs if __name__ == "__main__": if len(sys.argv) < 5: @@ -12,7 +13,6 @@ print("arg4: path for cluster ids") print("arg5: distance metric") exit(1) - # path data_path = sys.argv[1] @@ -28,9 +28,9 @@ metric = faiss.METRIC_INNER_PRODUCT print("Using InnerProduct metric") else: - raise ValueError("Unsupported distance metric. Use 'l2' or 'ip'.") + raise ValueError("Unsupported distance metric. Use 'l2' or 'ip'.") else: - metric = faiss.METRIC_L2 # by default, L2 metric + metric = faiss.METRIC_L2 # by default, L2 metric print("Using L2 metric by default") X = read_fvecs(data_path) @@ -45,7 +45,7 @@ index.train(X) t2 = time() - print(f"Time for training ivf {t2-t1} secs") + print(f"Time for training ivf {t2 - t1} secs") centroids = index.quantizer.reconstruct_n(0, index.nlist) _, cluster_id = index.quantizer.search(X, 1) diff --git a/python/utils/io.py b/python/utils/io.py index b5f18c8..6316108 100644 --- a/python/utils/io.py +++ b/python/utils/io.py @@ -1,6 +1,7 @@ -import numpy as np import struct +import numpy as np + def read_ivecs(filename): print(f"Reading File - {filename}") diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index 3fff52a..41143fb 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -6,4 +6,4 @@ from ._rabitqlib import HnswIndex, IvfIndex, SymqgIndex -__all__ = ["HnswIndex", "IvfIndex", "SymqgIndex"] \ No newline at end of file +__all__ = ["HnswIndex", "IvfIndex", "SymqgIndex"] diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index f8757bd..6b8a5c6 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -1,14 +1,14 @@ #pragma once +#include +#include + #include #include #include #include #include -#include -#include - #include "rabitqlib/defines.hpp" #include "rabitqlib/utils/rotator.hpp" @@ -42,8 +42,7 @@ inline rabitqlib::RotatorType rotator_from_string(const std::string& method) { template inline py::array_t ensure_2d_array( - py::handle value, - const char* name + py::handle value, const char* name ) { auto array = py::array_t::ensure(value); if (!array) { @@ -57,8 +56,7 @@ inline py::array_t ensure_2d_array template inline py::array_t ensure_1d_array( - py::handle value, - const char* name + py::handle value, const char* name ) { auto array = py::array_t::ensure(value); if (!array) { diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 76c3556..2550eef 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -1,11 +1,11 @@ +#include + #include +#include #include #include -#include #include -#include - #include "bindings_common.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" @@ -32,13 +32,7 @@ class HnswIndex { , metric_(metric_from_string(metric)) , random_seed_(random_seed) , index_(std::make_unique( - max_elements, - dim, - nbits, - M, - ef_construction, - random_seed, - metric_ + max_elements, dim, nbits, M, ef_construction, random_seed, metric_ )) {} void build( @@ -50,7 +44,8 @@ class HnswIndex { ) { auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); - auto cluster_ids_array = ensure_1d_array(cluster_ids, "cluster_ids"); + auto cluster_ids_array = + ensure_1d_array(cluster_ids, "cluster_ids"); if (static_cast(data_array.shape(1)) != dim_) { throw std::invalid_argument("data dimension does not match index dim"); @@ -58,16 +53,25 @@ class HnswIndex { if (static_cast(centroids_array.shape(1)) != dim_) { throw std::invalid_argument("centroid dimension does not match index dim"); } - if (static_cast(cluster_ids_array.shape(0)) != static_cast(data_array.shape(0))) { - throw std::invalid_argument("cluster_ids length must match number of rows in data"); + if (static_cast(cluster_ids_array.shape(0)) != + static_cast(data_array.shape(0))) { + throw std::invalid_argument( + "cluster_ids length must match number of rows in data" + ); } const size_t num_clusters = static_cast(centroids_array.shape(0)); num_clusters_ = num_clusters; // Ensure cluster_ids are writable for the C++ API by making a copy - std::vector cluster_ids_vec(static_cast(cluster_ids_array.shape(0))); - std::memcpy(cluster_ids_vec.data(), cluster_ids_array.data(), cluster_ids_vec.size() * sizeof(rabitqlib::PID)); + std::vector cluster_ids_vec( + static_cast(cluster_ids_array.shape(0)) + ); + std::memcpy( + cluster_ids_vec.data(), + cluster_ids_array.data(), + cluster_ids_vec.size() * sizeof(rabitqlib::PID) + ); index_->construct( num_clusters, @@ -98,21 +102,22 @@ class HnswIndex { auto dists_buf = dists.mutable_unchecked<2>(); std::vector>> results = index_->search( - query_array.data(), - static_cast(query_array.shape(0)), - k, - ef, - num_threads - ); + query_array.data(), + static_cast(query_array.shape(0)), + k, + ef, + num_threads + ); for (ssize_t i = 0; i < static_cast(results.size()); ++i) { - for ( - ssize_t j = 0; - j < static_cast(std::min(k, results[static_cast(i)].size())); - ++j - ) { - ids_buf(i, j) = results[static_cast(i)][static_cast(j)].second; - dists_buf(i, j) = results[static_cast(i)][static_cast(j)].first; + for (ssize_t j = 0; j < static_cast(std::min( + k, results[static_cast(i)].size() + )); + ++j) { + ids_buf(i, j) = + results[static_cast(i)][static_cast(j)].second; + dists_buf(i, j) = + results[static_cast(i)][static_cast(j)].first; } } return py::make_tuple(ids, dists); @@ -165,29 +170,37 @@ class HnswIndex { } // namespace rabitqlib::python_bindings // Register into combined module -void register_hnsw(py::module_ &m) { +void register_hnsw(py::module_& m) { using namespace rabitqlib::python_bindings; py::class_(m, "HnswIndex") - .def(py::init(), - py::arg("dim"), - py::arg("max_elements"), - py::arg("M") = 16, - py::arg("ef_construction") = 200, - py::arg("nbits") = 8, - py::arg("metric") = "l2", - py::arg("random_seed") = 100) - .def("build", &HnswIndex::build, - py::arg("data"), - py::arg("centroids"), - py::arg("cluster_ids"), - py::arg("num_threads") = 1, - py::arg("fast_quantization") = false) - .def("search", &HnswIndex::search, - py::arg("queries"), - py::arg("k"), - py::arg("ef") = 0, - py::arg("num_threads") = 1) + .def( + py::init(), + py::arg("dim"), + py::arg("max_elements"), + py::arg("M") = 16, + py::arg("ef_construction") = 200, + py::arg("nbits") = 8, + py::arg("metric") = "l2", + py::arg("random_seed") = 100 + ) + .def( + "build", + &HnswIndex::build, + py::arg("data"), + py::arg("centroids"), + py::arg("cluster_ids"), + py::arg("num_threads") = 1, + py::arg("fast_quantization") = false + ) + .def( + "search", + &HnswIndex::search, + py::arg("queries"), + py::arg("k"), + py::arg("ef") = 0, + py::arg("num_threads") = 1 + ) .def("save", &HnswIndex::save, py::arg("path")) .def_static("load", &HnswIndex::load, py::arg("path")) .def_property_readonly("dim", &HnswIndex::dim) diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 225bf98..defc06d 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -1,10 +1,10 @@ +#include + #include #include #include #include -#include - #include "bindings_common.hpp" #include "rabitqlib/index/ivf/ivf.hpp" @@ -44,7 +44,8 @@ class IvfIndex { ) { auto data_array = ensure_2d_array(data, "data"); auto centroids_array = ensure_2d_array(centroids, "centroids"); - auto cluster_ids_array = ensure_1d_array(cluster_ids, "cluster_ids"); + auto cluster_ids_array = + ensure_1d_array(cluster_ids, "cluster_ids"); if (static_cast(data_array.shape(1)) != dim_) { throw std::invalid_argument("data dimension does not match index dim"); @@ -52,11 +53,20 @@ class IvfIndex { if (static_cast(centroids_array.shape(1)) != dim_) { throw std::invalid_argument("centroid dimension does not match index dim"); } - if (static_cast(cluster_ids_array.shape(0)) != static_cast(data_array.shape(0))) { - throw std::invalid_argument("cluster_ids length must match number of rows in data"); + if (static_cast(cluster_ids_array.shape(0)) != + static_cast(data_array.shape(0))) { + throw std::invalid_argument( + "cluster_ids length must match number of rows in data" + ); } - index_->construct(data_array.data(), centroids_array.data(), cluster_ids_array.data(), fast_quantization, num_threads); + index_->construct( + data_array.data(), + centroids_array.data(), + cluster_ids_array.data(), + fast_quantization, + num_threads + ); built_ = true; } @@ -73,7 +83,8 @@ class IvfIndex { } const size_t nq = static_cast(query_array.shape(0)); - const auto shape = std::vector{static_cast(nq), static_cast(k)}; + const auto shape = + std::vector{static_cast(nq), static_cast(k)}; auto ids = py::array_t(shape); auto dists = py::array_t(shape); auto ids_buf = ids.mutable_unchecked<2>(); @@ -86,7 +97,7 @@ class IvfIndex { [&](size_t idx, size_t /*threadId*/) { std::vector row_ids(k, 0); std::vector row_dists(k, 0.0F); - + index_->search( query_array.data() + (idx * dim_), k, @@ -95,10 +106,12 @@ class IvfIndex { row_dists.data(), high_accuracy ); - + for (size_t j = 0; j < k; ++j) { - ids_buf(static_cast(idx), static_cast(j)) = row_ids[j]; - dists_buf(static_cast(idx), static_cast(j)) = row_dists[j]; + ids_buf(static_cast(idx), static_cast(j)) = + row_ids[j]; + dists_buf(static_cast(idx), static_cast(j)) = + row_dists[j]; } } ); @@ -148,34 +161,42 @@ class IvfIndex { } // namespace rabitqlib::python_bindings // Register IVF bindings into combined module -void register_ivf(py::module_ &m) { +void register_ivf(py::module_& m) { using namespace rabitqlib::python_bindings; py::class_(m, "IvfIndex") - .def(py::init(), - py::arg("dim"), - py::arg("max_elements"), - py::arg("num_clusters"), - py::arg("nbits"), - py::arg("metric") = "l2") - .def("build", &IvfIndex::build, - py::arg("data"), - py::arg("centroids"), - py::arg("cluster_ids"), - py::arg("num_threads") = 1, - py::arg("fast_quantization") = false) - .def("search", &IvfIndex::search, - py::arg("queries"), - py::arg("k"), - py::arg("nprobe"), - py::arg("high_accuracy") = true, - py::arg("num_threads") = 1) - .def("save", &IvfIndex::save, py::arg("path")) - .def_static("load", &IvfIndex::load, py::arg("path")) - .def_property_readonly("dim", &IvfIndex::dim) - .def_property_readonly("max_elements", &IvfIndex::max_elements) - .def_property_readonly("num_clusters", &IvfIndex::num_clusters) - .def_property_readonly("nbits", &IvfIndex::nbits) - .def_property_readonly("is_built", &IvfIndex::is_built) - .def_property_readonly("metric", &IvfIndex::metric); + .def( + py::init(), + py::arg("dim"), + py::arg("max_elements"), + py::arg("num_clusters"), + py::arg("nbits"), + py::arg("metric") = "l2" + ) + .def( + "build", + &IvfIndex::build, + py::arg("data"), + py::arg("centroids"), + py::arg("cluster_ids"), + py::arg("num_threads") = 1, + py::arg("fast_quantization") = false + ) + .def( + "search", + &IvfIndex::search, + py::arg("queries"), + py::arg("k"), + py::arg("nprobe"), + py::arg("high_accuracy") = true, + py::arg("num_threads") = 1 + ) + .def("save", &IvfIndex::save, py::arg("path")) + .def_static("load", &IvfIndex::load, py::arg("path")) + .def_property_readonly("dim", &IvfIndex::dim) + .def_property_readonly("max_elements", &IvfIndex::max_elements) + .def_property_readonly("num_clusters", &IvfIndex::num_clusters) + .def_property_readonly("nbits", &IvfIndex::nbits) + .def_property_readonly("is_built", &IvfIndex::is_built) + .def_property_readonly("metric", &IvfIndex::metric); } \ No newline at end of file diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index 9c58a77..33a8dd2 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -6,9 +6,9 @@ namespace py = pybind11; // Forward declarations of registration functions implemented in other cpp files -void register_hnsw(py::module_ &m); -void register_ivf(py::module_ &m); -void register_symqg(py::module_ &m); +void register_hnsw(py::module_& m); +void register_ivf(py::module_& m); +void register_symqg(py::module_& m); PYBIND11_MODULE(_rabitqlib, m) { m.doc() = "RabitQ Python bindings combined module"; diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 24691b7..18480a1 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -1,10 +1,10 @@ +#include + #include #include #include #include -#include - #include "bindings_common.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/symqg/qg.hpp" @@ -17,9 +17,7 @@ namespace rabitqlib::python_bindings { class SymqgIndex { public: SymqgIndex(size_t dim, size_t max_degree, const std::string& metric = "l2") - : dim_(dim) - , max_degree_(max_degree) - , metric_(metric_from_string(metric)) {} + : dim_(dim), max_degree_(max_degree), metric_(metric_from_string(metric)) {} void build(py::handle data, size_t ef_construction, size_t num_threads = 1) { auto data_array = ensure_2d_array(data, "data"); @@ -32,7 +30,9 @@ class SymqgIndex { num_points_, dim_, max_degree_, metric_, rabitqlib::RotatorType::FhtKacRotator ); - rabitqlib::symqg::QGBuilder builder(*index_, ef_construction, data_array.data(), num_threads); + rabitqlib::symqg::QGBuilder builder( + *index_, ef_construction, data_array.data(), num_threads + ); builder.build(); built_ = true; } @@ -49,12 +49,13 @@ class SymqgIndex { index_->set_ef(ef); const size_t nq = static_cast(query_array.shape(0)); - const auto shape = std::vector{static_cast(nq), static_cast(k)}; + const auto shape = + std::vector{static_cast(nq), static_cast(k)}; auto ids = py::array_t(shape); auto dists = py::array_t(shape); auto ids_buf = ids.mutable_unchecked<2>(); auto dists_buf = dists.mutable_unchecked<2>(); - + rabitqlib::ivf::parallel_for( 0, nq, @@ -62,10 +63,17 @@ class SymqgIndex { [&](size_t idx, size_t /*threadId*/) { std::vector row_ids(k, 0); std::vector row_dists(k, 0.0F); - index_->search(query_array.data() + (idx * dim_), static_cast(k), row_ids.data(), row_dists.data()); + index_->search( + query_array.data() + (idx * dim_), + static_cast(k), + row_ids.data(), + row_dists.data() + ); for (size_t j = 0; j < k; ++j) { - ids_buf(static_cast(idx), static_cast(j)) = row_ids[j]; - dists_buf(static_cast(idx), static_cast(j)) = row_dists[j]; + ids_buf(static_cast(idx), static_cast(j)) = + row_ids[j]; + dists_buf(static_cast(idx), static_cast(j)) = + row_dists[j]; } } ); @@ -112,28 +120,36 @@ class SymqgIndex { } // namespace rabitqlib::python_bindings // Register Symqg bindings into combined module -void register_symqg(py::module_ &m) { +void register_symqg(py::module_& m) { using namespace rabitqlib::python_bindings; py::class_(m, "SymqgIndex") - .def(py::init(), - py::arg("dim"), - py::arg("max_degree"), - py::arg("metric") = "l2") - .def("build", &SymqgIndex::build, - py::arg("data"), - py::arg("ef_construction"), - py::arg("num_threads") = 1) - .def("search", &SymqgIndex::search, - py::arg("queries"), - py::arg("k"), - py::arg("ef"), - py::arg("num_threads") = 1) - .def("save", &SymqgIndex::save, py::arg("path")) - .def_static("load", &SymqgIndex::load, py::arg("path")) - .def_property_readonly("dim", &SymqgIndex::dim) - .def_property_readonly("max_degree", &SymqgIndex::max_degree) - .def_property_readonly("num_points", &SymqgIndex::num_points) - .def_property_readonly("is_built", &SymqgIndex::is_built) - .def_property_readonly("metric", &SymqgIndex::metric); + .def( + py::init(), + py::arg("dim"), + py::arg("max_degree"), + py::arg("metric") = "l2" + ) + .def( + "build", + &SymqgIndex::build, + py::arg("data"), + py::arg("ef_construction"), + py::arg("num_threads") = 1 + ) + .def( + "search", + &SymqgIndex::search, + py::arg("queries"), + py::arg("k"), + py::arg("ef"), + py::arg("num_threads") = 1 + ) + .def("save", &SymqgIndex::save, py::arg("path")) + .def_static("load", &SymqgIndex::load, py::arg("path")) + .def_property_readonly("dim", &SymqgIndex::dim) + .def_property_readonly("max_degree", &SymqgIndex::max_degree) + .def_property_readonly("num_points", &SymqgIndex::num_points) + .def_property_readonly("is_built", &SymqgIndex::is_built) + .def_property_readonly("metric", &SymqgIndex::metric); } \ No newline at end of file diff --git a/sample/cpp/hnsw_rabitq_querying.cpp b/sample/cpp/hnsw_rabitq_querying.cpp index 0a391db..de8ab3f 100644 --- a/sample/cpp/hnsw_rabitq_querying.cpp +++ b/sample/cpp/hnsw_rabitq_querying.cpp @@ -72,7 +72,8 @@ int main(int argc, char* argv[]) { for (size_t i = 0; i < nq; i++) { for (size_t j = 0; j < topk; j++) { for (size_t k = 0; k < topk; k++) { - if (gt(i, k) == res[i][j].second) { + if (gt(static_cast(i), + static_cast(k)) == res[i][j].second) { total_correct++; break; } @@ -99,4 +100,4 @@ int main(int argc, char* argv[]) { for (size_t i = 0; i < avg_qps.size(); ++i) { std::cout << efs[i] << '\t' << avg_qps[i] << '\t' << avg_recall[i] << '\t' << '\n'; } -} \ No newline at end of file +} diff --git a/sample/cpp/ivf_rabitq_querying.cpp b/sample/cpp/ivf_rabitq_querying.cpp index 5f7f2e3..69a105a 100644 --- a/sample/cpp/ivf_rabitq_querying.cpp +++ b/sample/cpp/ivf_rabitq_querying.cpp @@ -88,18 +88,26 @@ int main(int argc, char** argv) { size_t nprobe = nprobes[l]; if (nprobe > ivf.num_clusters()) { std::cout << "nprobe " << nprobe << " is larger than number of clusters, "; - std::cout << "will use nprobe = num_cluster (" << ivf.num_clusters() << ").\n"; + std::cout << "will use nprobe = num_cluster (" << ivf.num_clusters() + << ").\n"; } size_t total_correct = 0; float total_time = 0; std::vector results(topk); for (size_t i = 0; i < nq; i++) { stopw.reset(); - ivf.search(&query(i, 0), topk, nprobe, results.data(), use_hacc); + ivf.search( + &query(static_cast(i), 0), + topk, + nprobe, + results.data(), + use_hacc + ); total_time += stopw.get_elapsed_micro(); for (size_t j = 0; j < topk; j++) { for (size_t k = 0; k < topk; k++) { - if (gt(i, k) == results[j]) { + if (gt(static_cast(i), + static_cast(k)) == results[j]) { total_correct++; break; } @@ -148,10 +156,13 @@ static std::vector get_nprobes( size_t total_correct = 0; std::vector results(topk); for (size_t i = 0; i < nq; i++) { - ivf.search(&query(i, 0), topk, nprobe, results.data()); + ivf.search( + &query(static_cast(i), 0), topk, nprobe, results.data() + ); for (size_t j = 0; j < topk; j++) { for (size_t k = 0; k < topk; k++) { - if (gt(i, k) == results[j]) { + if (gt(static_cast(i), static_cast(k)) == + results[j]) { total_correct++; break; } @@ -167,4 +178,4 @@ static std::vector get_nprobes( } return nprobes; -} \ No newline at end of file +} diff --git a/sample/cpp/symqg_querying.cpp b/sample/cpp/symqg_querying.cpp index 4c71b04..5ae2047 100644 --- a/sample/cpp/symqg_querying.cpp +++ b/sample/cpp/symqg_querying.cpp @@ -11,8 +11,25 @@ using data_type = rabitqlib::RowMajorArray; using gt_type = rabitqlib::RowMajorArray; std::vector efs = { - 10, 20, 40, 50, 60, 80, 100, 150, 170, 190, 200, 250, 300, 400, 500, 600, 700, 800, 1500 -}; + 10, + 20, + 40, + 50, + 60, + 80, + 100, + 150, + 170, + 190, + 200, + 250, + 300, + 400, + 500, + 600, + 700, + 800, + 1500}; size_t test_round = 3; size_t topk = 10; @@ -53,11 +70,12 @@ int main(int argc, char** argv) { std::vector results(topk); for (size_t z = 0; z < nq; z++) { stopw.reset(); - qg.search(&query(z, 0), topk, results.data()); + qg.search(&query(static_cast(z), 0), topk, results.data()); total_time += stopw.get_elapsed_micro(); for (size_t y = 0; y < topk; y++) { for (size_t k = 0; k < topk; k++) { - if (gt(z, k) == results[y]) { + if (gt(static_cast(z), + static_cast(k)) == results[y]) { total_correct++; break; } @@ -82,4 +100,4 @@ int main(int argc, char** argv) { } return 0; -} \ No newline at end of file +} diff --git a/sample/python/hnsw_rabitq_indexing.py b/sample/python/hnsw_rabitq_indexing.py index b288c01..fef9f4d 100644 --- a/sample/python/hnsw_rabitq_indexing.py +++ b/sample/python/hnsw_rabitq_indexing.py @@ -1,18 +1,20 @@ import argparse from time import time + from rabitqlib import HnswIndex -from utils import read_fvecs, cluster_data + +from utils import cluster_data, 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 -METRIC = "l2" # "l2" or "ip" -FASTER_QUANT = False # use faster quantization -NUM_THREADS = 16 # number of threads for build +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 +METRIC = "l2" # "l2" or "ip" +FASTER_QUANT = False # use faster quantization +NUM_THREADS = 16 # number of threads for build # ────────────────────────────────────────────── @@ -22,13 +24,17 @@ def main(args=None) -> None: print(f"Data shape: {data.shape}") # 2. Cluster with FAISS - centroids, cluster_ids = cluster_data(data, args.num_clusters, args.metric, args.num_threads) + centroids, cluster_ids = cluster_data( + data, args.num_clusters, args.metric, args.num_threads + ) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") # 3. Build HNSW index n, dim = data.shape - print(f"\nBuilding HNSW index: n={n}, dim={dim}, M={args.degree}, " - f"ef={args.ef_construction}, bits={args.total_bits}, metric={args.metric}") + print( + f"\nBuilding HNSW index: n={n}, dim={dim}, M={args.degree}, " + f"ef={args.ef_construction}, bits={args.total_bits}, metric={args.metric}" + ) idx = HnswIndex( dim=dim, @@ -58,14 +64,61 @@ 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") - parser.add_argument("--degree", dest="degree", type=int, metavar="INT", default=M, help="Degree bound for HNSW") - parser.add_argument("--ef-construction", dest="ef_construction", type=int, metavar="INT", default=EF_CONSTRUCTION, help="EF parameter for index construction") - parser.add_argument("--total-bits", dest="total_bits", type=int, metavar="INT", default=TOTAL_BITS, help="Total number of bits for quantization") - parser.add_argument("--metric", dest="metric", type=str, default="l2", choices=["l2", "ip"], help="Distance metric (l2 or ip)") - parser.add_argument("--faster-quant", dest="faster_quant", action="store_true", help="Use faster quantization method") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUM_THREADS, help="Number of threads for building the index") + parser.add_argument( + "--num-clusters", + dest="num_clusters", + type=int, + metavar="INT", + default=256, + help="Number of clusters for quantization", + ) + parser.add_argument( + "--degree", + dest="degree", + type=int, + metavar="INT", + default=M, + help="Degree bound for HNSW", + ) + parser.add_argument( + "--ef-construction", + dest="ef_construction", + type=int, + metavar="INT", + default=EF_CONSTRUCTION, + help="EF parameter for index construction", + ) + parser.add_argument( + "--total-bits", + dest="total_bits", + type=int, + metavar="INT", + default=TOTAL_BITS, + help="Total number of bits for quantization", + ) + parser.add_argument( + "--metric", + dest="metric", + type=str, + default="l2", + choices=["l2", "ip"], + help="Distance metric (l2 or ip)", + ) + parser.add_argument( + "--faster-quant", + dest="faster_quant", + action="store_true", + help="Use faster quantization method", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUM_THREADS, + help="Number of threads for building the index", + ) args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/sample/python/hnsw_rabitq_querying.py b/sample/python/hnsw_rabitq_querying.py index 28a9815..40dbf47 100644 --- a/sample/python/hnsw_rabitq_querying.py +++ b/sample/python/hnsw_rabitq_querying.py @@ -1,15 +1,17 @@ import argparse -import numpy as np from time import time + +import numpy as np from rabitqlib import HnswIndex -from utils import read_fvecs, read_ivecs, compute_recall + +from utils import compute_recall, read_fvecs, read_ivecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -TOPK = 10 # top-k results -NUM_THREADS = 1 #number of threads -EFS = [10, 20, 40, 80, 120, 200, 400, 600, 800, 1000, 1500, 2000] +TOPK = 10 # top-k results +NUM_THREADS = 1 # number of threads +EFS = [10, 20, 40, 80, 120, 200, 400, 600, 800, 1000, 1500, 2000] TEST_ROUNDS = 3 # ────────────────────────────────────────────── @@ -17,8 +19,8 @@ def main(args=None) -> None: # 1. Load queries and ground truth queries = read_fvecs(args.query_file) - gt = read_ivecs(args.gt_file) - nq = queries.shape[0] + gt = read_ivecs(args.gt_file) + nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") # 2. Load index @@ -28,22 +30,24 @@ def main(args=None) -> None: print("\nsearch start >.....\n") - all_qps = np.zeros((args.test_rounds, len(EFS))) + all_qps = np.zeros((args.test_rounds, len(EFS))) all_recall = np.zeros((args.test_rounds, len(EFS))) for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): t0 = time() - ids, _ = idx.search(queries, k=args.topk, ef=ef, num_threads=args.num_threads) + ids, _ = idx.search( + queries, k=args.topk, ef=ef, num_threads=args.num_threads + ) elapsed = time() - t0 # seconds - qps = nq / elapsed + qps = nq / elapsed recall = compute_recall(ids, gt, args.topk) - all_qps[r, i_probe] = qps + all_qps[r, i_probe] = qps all_recall[r, i_probe] = recall - avg_qps = all_qps.mean(axis=0) + avg_qps = all_qps.mean(axis=0) avg_recall = all_recall.mean(axis=0) # 3. Print results table @@ -57,11 +61,32 @@ def main(args=None) -> None: parser = argparse.ArgumentParser(description="RaBitQ HNSW Querying") parser.add_argument("index_file", type=str, help="Path to the HNSW index file") - parser.add_argument("query_file", type=str, help="Path to the query file") + parser.add_argument("query_file", type=str, help="Path to the query file") parser.add_argument("gt_file", type=str, help="Path to the ground truth file") - parser.add_argument("--topk", dest="topk", type=int, metavar="INT", default=TOPK, help="Number of top results to retrieve") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUM_THREADS, help="Number of threads for searching") - parser.add_argument("--test-rounds", dest="test_rounds", type=int, metavar="INT", default=TEST_ROUNDS, help="Number of test rounds for averaging") + parser.add_argument( + "--topk", + dest="topk", + type=int, + metavar="INT", + default=TOPK, + help="Number of top results to retrieve", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUM_THREADS, + help="Number of threads for searching", + ) + parser.add_argument( + "--test-rounds", + dest="test_rounds", + type=int, + metavar="INT", + default=TEST_ROUNDS, + help="Number of test rounds for averaging", + ) args = parser.parse_args() - - main(args) \ No newline at end of file + + main(args) diff --git a/sample/python/ivf_rabitq_indexing.py b/sample/python/ivf_rabitq_indexing.py index 878e309..286d5e1 100644 --- a/sample/python/ivf_rabitq_indexing.py +++ b/sample/python/ivf_rabitq_indexing.py @@ -1,16 +1,18 @@ import argparse from time import time + from rabitqlib import IvfIndex -from utils import read_fvecs, cluster_data + +from utils import cluster_data, 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 -NUM_THREADS = 16 # number of threads for building the index +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 +NUM_THREADS = 16 # number of threads for building the index # ────────────────────────────────────────────── @@ -18,17 +20,21 @@ def main(args=None) -> None: # 1. Load data data = read_fvecs(args.data_file) n, dim = data.shape - print(f"Data loaded") + print("Data loaded") 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) + centroids, cluster_ids = cluster_data( + data, args.num_clusters, args.metric, args.num_threads + ) print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") # 3. Build IVF index - print(f"\nBuilding IVF index: bits={args.total_bits}, metric={args.metric}, " - f"num_threads={args.num_threads}, faster_quant={args.faster_quant}") + print( + f"\nBuilding IVF index: bits={args.total_bits}, metric={args.metric}, " + f"num_threads={args.num_threads}, faster_quant={args.faster_quant}" + ) idx = IvfIndex( dim=dim, @@ -39,7 +45,13 @@ def main(args=None) -> None: ) t0 = time() - idx.build(data, centroids, cluster_ids, num_threads=args.num_threads, fast_quantization=args.faster_quant) + idx.build( + data, + centroids, + cluster_ids, + num_threads=args.num_threads, + fast_quantization=args.faster_quant, + ) elapsed_min = (time() - t0) / 60 print("IVF constructed") @@ -53,10 +65,43 @@ 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)") - parser.add_argument("--total-bits", dest="total_bits", type=int, metavar="INT", default=TOTAL_BITS, help="Total number of bits for quantization") - parser.add_argument("--metric", dest="metric", type=str, default=METRIC, choices=["l2", "ip"], help="Distance metric (l2 or ip)") - parser.add_argument("--faster-quant", dest="faster_quant", action="store_true", help="Use faster quantization method") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUM_THREADS, help="Number of threads for building the index") + parser.add_argument( + "--num-clusters", + dest="num_clusters", + type=int, + metavar="INT", + default=NUM_CLUSTERS, + help="Number of clusters (K for IVF)", + ) + parser.add_argument( + "--total-bits", + dest="total_bits", + type=int, + metavar="INT", + default=TOTAL_BITS, + help="Total number of bits for quantization", + ) + parser.add_argument( + "--metric", + dest="metric", + type=str, + default=METRIC, + choices=["l2", "ip"], + help="Distance metric (l2 or ip)", + ) + parser.add_argument( + "--faster-quant", + dest="faster_quant", + action="store_true", + help="Use faster quantization method", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUM_THREADS, + help="Number of threads for building the index", + ) args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/sample/python/ivf_rabitq_querying.py b/sample/python/ivf_rabitq_querying.py index 98f64a0..6ccddce 100644 --- a/sample/python/ivf_rabitq_querying.py +++ b/sample/python/ivf_rabitq_querying.py @@ -1,16 +1,18 @@ import argparse -import numpy as np from time import time + +import numpy as np from rabitqlib import IvfIndex -from utils import read_fvecs, read_ivecs, compute_recall + +from utils import compute_recall, read_fvecs, read_ivecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -USE_HACC = True # use high accuracy fastscan -TOPK = 100 # top-k results -NUMBER_THREADS = 4 # number of threads for search -TEST_ROUNDS = 3 # number of test rounds +USE_HACC = True # use high accuracy fastscan +TOPK = 100 # top-k results +NUMBER_THREADS = 4 # number of threads for search +TEST_ROUNDS = 3 # number of test rounds NPROBES = [5, 10, 20, 40, 80, 120, 200] # ────────────────────────────────────────────── @@ -18,8 +20,8 @@ def main(args=None) -> None: # 1. Load queries and ground truth queries = read_fvecs(args.query_file) - gt = read_ivecs(args.gt_file) - nq = queries.shape[0] + gt = read_ivecs(args.gt_file) + nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") print(f"TopK: {args.topk}, use_hacc: {args.use_hacc}") @@ -27,24 +29,32 @@ def main(args=None) -> None: idx = IvfIndex.load(args.index_file) print(f"Index loaded — dim={idx.dim}, clusters={idx.num_clusters}") - all_qps = np.zeros((args.test_rounds, len(NPROBES))) + all_qps = np.zeros((args.test_rounds, len(NPROBES))) all_recall = np.zeros((args.test_rounds, len(NPROBES))) # 3. Benchmark for r in range(args.test_rounds): - for l, nprobe in enumerate(NPROBES): + for probe_index, nprobe in enumerate(NPROBES): if nprobe > idx.num_clusters: - print(f"nprobe {nprobe} is larger than number of clusters, " - f"will use nprobe = num_clusters ({idx.num_clusters}).") + print( + f"nprobe {nprobe} is larger than number of clusters, " + f"will use nprobe = num_clusters ({idx.num_clusters})." + ) t0 = time() - ids, _ = idx.search(queries, k=args.topk, nprobe=nprobe, high_accuracy=args.use_hacc, num_threads=args.num_threads) + ids, _ = idx.search( + queries, + k=args.topk, + nprobe=nprobe, + high_accuracy=args.use_hacc, + num_threads=args.num_threads, + ) elapsed = time() - t0 - all_qps[r, l] = nq / elapsed - all_recall[r, l] = compute_recall(ids, gt, args.topk) + all_qps[r, probe_index] = nq / elapsed + all_recall[r, probe_index] = compute_recall(ids, gt, args.topk) - avg_qps = all_qps.mean(axis=0) + avg_qps = all_qps.mean(axis=0) avg_recall = all_recall.mean(axis=0) # 4. Print results table @@ -60,10 +70,36 @@ def main(args=None) -> None: parser.add_argument("index_file", type=str, help="Path to the IVF index file") parser.add_argument("query_file", type=str, help="Path to the query file") parser.add_argument("gt_file", type=str, help="Path to the ground truth file") - parser.add_argument("--topk", dest="topk", type=int, metavar="INT", default=TOPK, help="Number of top results to retrieve") - parser.add_argument("--use-hacc", dest="use_hacc", action="store_true", help="Use high accuracy fastscan method") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUMBER_THREADS, help="Number of threads for search") - parser.add_argument("--test-rounds", dest="test_rounds", type=int, metavar="INT", default=TEST_ROUNDS, help="Number of test rounds for averaging") + parser.add_argument( + "--topk", + dest="topk", + type=int, + metavar="INT", + default=TOPK, + help="Number of top results to retrieve", + ) + parser.add_argument( + "--use-hacc", + dest="use_hacc", + action="store_true", + help="Use high accuracy fastscan method", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUMBER_THREADS, + help="Number of threads for search", + ) + parser.add_argument( + "--test-rounds", + dest="test_rounds", + type=int, + metavar="INT", + default=TEST_ROUNDS, + help="Number of test rounds for averaging", + ) args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/sample/python/symqg_indexing.py b/sample/python/symqg_indexing.py index 4ac421c..7bdf692 100644 --- a/sample/python/symqg_indexing.py +++ b/sample/python/symqg_indexing.py @@ -1,15 +1,17 @@ import argparse from time import time + from rabitqlib import SymqgIndex + from utils import read_fvecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -MAX_DEGREE = 32 # degree bound for SymphonyQG -EF_CONSTRUCTION = 200 # ef for indexing -METRIC = "l2" # "l2" or "ip" -NUM_THREADS = 16 # number of threads for build +MAX_DEGREE = 32 # degree bound for SymphonyQG +EF_CONSTRUCTION = 200 # ef for indexing +METRIC = "l2" # "l2" or "ip" +NUM_THREADS = 16 # number of threads for build # ────────────────────────────────────────────── @@ -20,11 +22,13 @@ def main(args=None) -> None: # 3. Build SymphonyQG index n, dim = data.shape - print(f"\nBuilding SymphonyQG index: n={n}, dim={dim}, MaxDegree={args.max_degree}, " - f"ef={args.ef_construction}, metric={args.metric}") + print( + f"\nBuilding SymphonyQG index: n={n}, dim={dim}, MaxDegree={args.max_degree}, " + f"ef={args.ef_construction}, metric={args.metric}" + ) idx = SymqgIndex(dim=dim, max_degree=args.max_degree, metric=args.metric) - + t0 = time() idx.build(data, ef_construction=args.ef_construction, num_threads=args.num_threads) print(f"Indexing time: {time() - t0:.2f}s") @@ -37,10 +41,38 @@ def main(args=None) -> None: parser = argparse.ArgumentParser(description="RaBitQ SymphonyQG Manager") 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("--max-degree", dest="max_degree", type=int, metavar="INT", default=MAX_DEGREE, help="Degree bound for SymphonyQG") - parser.add_argument("--ef-construction", dest="ef_construction", type=int, metavar="INT", default=EF_CONSTRUCTION, help="EF parameter for index construction") - parser.add_argument("--metric", dest="metric", type=str, default=METRIC, choices=["l2", "ip"], help="Distance metric (l2 or ip)") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUM_THREADS, help="Number of threads for building the index") + parser.add_argument( + "--max-degree", + dest="max_degree", + type=int, + metavar="INT", + default=MAX_DEGREE, + help="Degree bound for SymphonyQG", + ) + parser.add_argument( + "--ef-construction", + dest="ef_construction", + type=int, + metavar="INT", + default=EF_CONSTRUCTION, + help="EF parameter for index construction", + ) + parser.add_argument( + "--metric", + dest="metric", + type=str, + default=METRIC, + choices=["l2", "ip"], + help="Distance metric (l2 or ip)", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUM_THREADS, + help="Number of threads for building the index", + ) args = parser.parse_args() main(args) diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index d4ed5c5..6228064 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -1,16 +1,18 @@ import argparse -import numpy as np from time import time + +import numpy as np from rabitqlib import SymqgIndex -from utils import read_fvecs, read_ivecs, compute_recall + +from utils import compute_recall, read_fvecs, read_ivecs # ────────────────────────────────────────────── # Default configuration # ────────────────────────────────────────────── -METRIC = "l2" # "l2" or "ip" -TOPK = 10 # top-k results -NUM_THREADS = 1 # number of threads -EFS = [10, 20, 40, 80, 120, 200, 400, 600, 800, 1000, 1500, 2000] +METRIC = "l2" # "l2" or "ip" +TOPK = 10 # top-k results +NUM_THREADS = 1 # number of threads +EFS = [10, 20, 40, 80, 120, 200, 400, 600, 800, 1000, 1500, 2000] TEST_ROUNDS = 3 # ────────────────────────────────────────────── @@ -18,8 +20,8 @@ def main(args=None) -> None: # 1. Load queries and ground truth queries = read_fvecs(args.query_file) - gt = read_ivecs(args.gt_file) - nq = queries.shape[0] + gt = read_ivecs(args.gt_file) + nq = queries.shape[0] print(f"Queries: {queries.shape}, GT: {gt.shape}") # 2. Load index @@ -29,22 +31,24 @@ def main(args=None) -> None: print("\nsearch start >.....\n") - all_qps = np.zeros((args.test_rounds, len(EFS))) + all_qps = np.zeros((args.test_rounds, len(EFS))) all_recall = np.zeros((args.test_rounds, len(EFS))) for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): t0 = time() - ids, _ = idx.search(queries, k=args.topk, ef=ef, num_threads=args.num_threads) + ids, _ = idx.search( + queries, k=args.topk, ef=ef, num_threads=args.num_threads + ) elapsed = time() - t0 # seconds - qps = nq / elapsed + qps = nq / elapsed recall = compute_recall(ids, gt, args.topk) - all_qps[r, i_probe] = qps + all_qps[r, i_probe] = qps all_recall[r, i_probe] = recall - avg_qps = all_qps.mean(axis=0) + avg_qps = all_qps.mean(axis=0) avg_recall = all_recall.mean(axis=0) # 3. Print results table @@ -56,12 +60,42 @@ def main(args=None) -> None: if __name__ == "__main__": parser = argparse.ArgumentParser(description="RaBitQ SymphonyQG Querying Manager") - parser.add_argument("index_file", type=str, help="Path to the SymphonyQG index file") + parser.add_argument( + "index_file", type=str, help="Path to the SymphonyQG index file" + ) parser.add_argument("query_file", type=str, help="Path to the query file") parser.add_argument("gt_file", type=str, help="Path to the ground truth file") - parser.add_argument("--metric", dest="metric", type=str, default=METRIC, choices=["l2", "ip"], help="Distance metric (l2 or ip)") - parser.add_argument("--topk", dest="topk", type=int, metavar="INT", default=TOPK, help="Number of top-k results to retrieve") - parser.add_argument("--num-threads", dest="num_threads", type=int, metavar="INT", default=NUM_THREADS, help="Number of threads for searching") - parser.add_argument("--test-rounds", dest="test_rounds", type=int, metavar="INT", default=TEST_ROUNDS, help="Number of test rounds to average results") + parser.add_argument( + "--metric", + dest="metric", + type=str, + default=METRIC, + choices=["l2", "ip"], + help="Distance metric (l2 or ip)", + ) + parser.add_argument( + "--topk", + dest="topk", + type=int, + metavar="INT", + default=TOPK, + help="Number of top-k results to retrieve", + ) + parser.add_argument( + "--num-threads", + dest="num_threads", + type=int, + metavar="INT", + default=NUM_THREADS, + help="Number of threads for searching", + ) + parser.add_argument( + "--test-rounds", + dest="test_rounds", + type=int, + metavar="INT", + default=TEST_ROUNDS, + help="Number of test rounds to average results", + ) args = parser.parse_args() main(args) diff --git a/sample/python/utils.py b/sample/python/utils.py index 9b6adba..c5110c2 100644 --- a/sample/python/utils.py +++ b/sample/python/utils.py @@ -1,12 +1,13 @@ -import numpy as np -import faiss from time import time +import faiss +import numpy as np # ────────────────────────────────────────────── # File I/O # ────────────────────────────────────────────── + def read_ivecs(filename: str) -> np.ndarray: print(f"Reading File - {filename}") a = np.fromfile(filename, dtype="int32") @@ -18,10 +19,12 @@ def read_ivecs(filename: str) -> np.ndarray: def read_fvecs(filename: str) -> np.ndarray: return read_ivecs(filename).view("float32") + # ────────────────────────────────────────────── # Benchmarking utilities # ────────────────────────────────────────────── + def compute_recall(ids: np.ndarray, gt: np.ndarray, topk: int) -> float: """Compute recall@topk: fraction of gt top-k found in returned top-k.""" nq = ids.shape[0] @@ -33,10 +36,12 @@ def compute_recall(ids: np.ndarray, gt: np.ndarray, topk: int) -> float: total_correct += 1 return total_correct / (nq * topk) + # ────────────────────────────────────────────── # Clustering # ────────────────────────────────────────────── + def cluster_data( X: np.ndarray, K: int, @@ -70,8 +75,8 @@ def cluster_data( 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 + 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 - return centroids, cluster_ids \ No newline at end of file + return centroids, cluster_ids diff --git a/scripts/apply-format.sh b/scripts/apply-format.sh new file mode 100755 index 0000000..cab3543 --- /dev/null +++ b/scripts/apply-format.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root +readonly formatter="${CLANG_FORMAT:-clang-format-15}" + +if ! command -v "$formatter" >/dev/null 2>&1; then + echo "error: $formatter was not found; install clang-format 15 or set CLANG_FORMAT" >&2 + exit 1 +fi + +if [[ "$($formatter --version)" != *"version 15."* ]]; then + echo "error: RaBitQ formatting requires clang-format 15" >&2 + echo "found: $($formatter --version)" >&2 + exit 1 +fi + +"$repo_root/scripts/clang-format-files.sh" "$@" \ + | xargs -0 --no-run-if-empty "$formatter" -i diff --git a/scripts/check-format.sh b/scripts/check-format.sh new file mode 100755 index 0000000..0152bc4 --- /dev/null +++ b/scripts/check-format.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root +readonly formatter="${CLANG_FORMAT:-clang-format-15}" + +if ! command -v "$formatter" >/dev/null 2>&1; then + echo "error: $formatter was not found; install clang-format 15 or set CLANG_FORMAT" >&2 + exit 1 +fi + +if [[ "$($formatter --version)" != *"version 15."* ]]; then + echo "error: RaBitQ formatting requires clang-format 15" >&2 + echo "found: $($formatter --version)" >&2 + exit 1 +fi + +mapfile -d '' files < <("$repo_root/scripts/clang-format-files.sh" "$@") + +if printf '%s\n' "${files[@]}" \ + | grep -Eq '/include/rabitqlib/third/|/include/rabitqlib/utils/fht_avx\.hpp$'; then + echo "error: the format file list contains excluded third-party code" >&2 + exit 1 +fi + +if ((${#files[@]} > 0)); then + printf '%s\0' "${files[@]}" \ + | xargs -0 --no-run-if-empty "$formatter" --dry-run --Werror +fi diff --git a/scripts/check-python.sh b/scripts/check-python.sh new file mode 100755 index 0000000..d42b2e3 --- /dev/null +++ b/scripts/check-python.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root +readonly ruff="${RUFF:-ruff}" + +if ! command -v "$ruff" >/dev/null 2>&1; then + echo "error: $ruff was not found; install Ruff 0.16.1 or set RUFF" >&2 + exit 1 +fi + +if [[ "$($ruff --version)" != "ruff 0.16.1" ]]; then + echo "error: RaBitQ Python checks require Ruff 0.16.1" >&2 + echo "found: $($ruff --version)" >&2 + exit 1 +fi + +cd "$repo_root" +"$ruff" format --check python python_bindings sample/python tests/python +"$ruff" check python python_bindings sample/python tests/python diff --git a/scripts/check-tidy.sh b/scripts/check-tidy.sh new file mode 100755 index 0000000..8b2b227 --- /dev/null +++ b/scripts/check-tidy.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root +readonly build_dir="${1:-$repo_root/build}" +readonly tidy_runner="${RUN_CLANG_TIDY:-run-clang-tidy-15}" +readonly compiler="${CXX:-c++}" + +if [[ ! -f "$build_dir/compile_commands.json" ]]; then + echo "error: $build_dir/compile_commands.json was not found" >&2 + echo "configure CMake with CMAKE_EXPORT_COMPILE_COMMANDS=ON first" >&2 + exit 1 +fi + +if ! command -v "$tidy_runner" >/dev/null 2>&1; then + echo "error: $tidy_runner was not found; install clang-tidy 15" >&2 + exit 1 +fi + +if ! command -v "$compiler" >/dev/null 2>&1; then + echo "error: $compiler was not found; set CXX to the compiler used by CMake" >&2 + exit 1 +fi + +mapfile -t system_include_dirs < <( + "$compiler" -E -x c++ - -v &1 \ + | awk ' + /#include <\.\.\.> search starts here:/ { in_search_path = 1; next } + /End of search list/ { in_search_path = 0 } + in_search_path { + sub(/^ /, "") + print + } + ' +) + +if ((${#system_include_dirs[@]} == 0)); then + echo "error: could not determine the system include paths for $compiler" >&2 + exit 1 +fi + +extra_args=("-extra-arg-before=-Wno-unused-command-line-argument") +if [[ -n "${CLANG_RESOURCE_DIR:-}" ]]; then + if [[ ! -d "$CLANG_RESOURCE_DIR" ]]; then + echo "error: CLANG_RESOURCE_DIR is not a directory: $CLANG_RESOURCE_DIR" >&2 + exit 1 + fi + extra_args+=("-extra-arg-before=-resource-dir=$CLANG_RESOURCE_DIR") +fi +for include_dir in "${system_include_dirs[@]}"; do + if [[ "$include_dir" == */lib/gcc/*/include ]]; then + # Keep Clang's intrinsic headers ahead of GCC's, while still making + # compiler-provided headers such as omp.h available. + extra_args+=("-extra-arg-before=-idirafter$include_dir") + else + extra_args+=("-extra-arg-before=-isystem$include_dir") + fi +done + +mapfile -t first_party_headers < <( + git -C "$repo_root" ls-files -- '*.h' '*.hpp' \ + ':(exclude)include/rabitqlib/third/**' \ + ':(exclude)include/rabitqlib/utils/fht_avx.hpp' +) + +header_filter="$repo_root/(" +separator="" +for header in "${first_party_headers[@]}"; do + escaped_header="${header//./\\.}" + header_filter+="$separator$escaped_header" + separator="|" +done +header_filter+=")$" + +tidy_output="$(mktemp)" +if [[ -z "$tidy_output" || ! -f "$tidy_output" ]]; then + echo "error: could not create a temporary clang-tidy output file" >&2 + exit 1 +fi +trap 'rm -f -- "$tidy_output"' EXIT + +if ! "$tidy_runner" \ + -quiet \ + -p "$build_dir" \ + -header-filter="^$header_filter" \ + "${extra_args[@]}" \ + "^$repo_root/(src|python_bindings|sample|tests)/" \ + >"$tidy_output" 2>&1; then + cat "$tidy_output" >&2 + exit 1 +fi + +echo "clang-tidy passed for all first-party translation units" diff --git a/scripts/clang-format-files.sh b/scripts/clang-format-files.sh new file mode 100755 index 0000000..0d51512 --- /dev/null +++ b/scripts/clang-format-files.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root + +if (($# > 0)); then + printf '%s\0' "$@" +else + git -C "$repo_root" ls-files -z -- \ + '*.c' '*.cc' '*.cpp' '*.cxx' '*.h' '*.hpp' '*.cu' '*.cuh' \ + ':(exclude)include/rabitqlib/third/**' \ + ':(exclude)include/rabitqlib/utils/fht_avx.hpp' +fi | while IFS= read -r -d '' file; do + if [[ "$file" == /* ]]; then + file="$(realpath --relative-to="$repo_root" "$file")" + else + file="${file#./}" + fi + + case "$file" in + ../* | */../* | */..) + continue + ;; + include/rabitqlib/third/* | include/rabitqlib/utils/fht_avx.hpp) + continue + ;; + esac + + if [[ -f "$repo_root/$file" ]]; then + printf '%s\0' "$repo_root/$file" + fi +done diff --git a/scripts/format-changed.sh b/scripts/format-changed.sh new file mode 100755 index 0000000..233190b --- /dev/null +++ b/scripts/format-changed.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash + +set -euo pipefail + +repo_root="$(git rev-parse --show-toplevel)" +readonly repo_root +readonly formatter="${CLANG_FORMAT:-clang-format-15}" + +if ! command -v "$formatter" >/dev/null 2>&1; then + echo "error: $formatter was not found; install clang-format 15 or set CLANG_FORMAT" >&2 + exit 1 +fi + +if [[ "$($formatter --version)" != *"version 15."* ]]; then + echo "error: RaBitQ formatting requires clang-format 15" >&2 + echo "found: $($formatter --version)" >&2 + exit 1 +fi + +mapfile -d '' absolute_files < <("$repo_root/scripts/clang-format-files.sh") +files=() +for file in "${absolute_files[@]}"; do + files+=("${file#"$repo_root/"}") +done + +if ((${#files[@]} > 0)); then + git -C "$repo_root" clang-format --binary "$formatter" "$@" -- "${files[@]}" +fi diff --git a/src/index/hnsw_search_avx2.cpp b/src/index/hnsw_search_avx2.cpp index 6fef52c..cb2b920 100644 --- a/src/index/hnsw_search_avx2.cpp +++ b/src/index/hnsw_search_avx2.cpp @@ -1,5 +1,4 @@ #include "hnsw_search_avx2_kernels.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 ee5dae4..4fbe015 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -14,8 +14,38 @@ namespace rabitqlib::hnsw::detail { static inline __m256i hnsw_popcount_avx2(__m256i v) { // Lookup table for population count of 0-15 const __m256i lookup = _mm256_setr_epi8( - 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, - 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4 + 0, + 1, + 1, + 2, + 1, + 2, + 2, + 3, + 1, + 2, + 2, + 3, + 2, + 3, + 3, + 4, + 0, + 1, + 1, + 2, + 1, + 2, + 2, + 3, + 1, + 2, + 2, + 3, + 2, + 3, + 3, + 4 ); const __m256i low_mask = _mm256_set1_epi8(0x0f); @@ -100,7 +130,8 @@ static inline float hnsw_warmup_ip_x0_q_512_avx2( for (; i < dim_end_512; i += 512) { // Load 64 bytes of data using paired 32-byte loads __m256i data_vec_lo = _mm256_loadu_si256(reinterpret_cast(data)); - __m256i data_vec_hi = _mm256_loadu_si256(reinterpret_cast(data + 4)); + __m256i data_vec_hi = + _mm256_loadu_si256(reinterpret_cast(data + 4)); data += 8; // Advance 8 x 64-bit ints (64 bytes) acc_ppc = _mm256_add_epi64(acc_ppc, hnsw_popcount_avx2(data_vec_lo)); diff --git a/src/index/hnsw_search_avx512_core.cpp b/src/index/hnsw_search_avx512_core.cpp index f5c3543..07960f3 100644 --- a/src/index/hnsw_search_avx512_core.cpp +++ b/src/index/hnsw_search_avx512_core.cpp @@ -1,6 +1,5 @@ #include "hnsw_search_avx2_kernels.hpp" #include "hnsw_search_avx512_kernels.hpp" - #include "rabitqlib/index/hnsw/hnsw.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/index/hnsw_search_avx512_popcnt.cpp b/src/index/hnsw_search_avx512_popcnt.cpp index a7dfb18..3af1e92 100644 --- a/src/index/hnsw_search_avx512_popcnt.cpp +++ b/src/index/hnsw_search_avx512_popcnt.cpp @@ -1,5 +1,4 @@ #include "hnsw_search_avx512_kernels.hpp" - #include "rabitqlib/index/hnsw/hnsw.hpp" namespace rabitqlib::hnsw::detail { diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 4aebec3..eab3eb6 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -5,17 +5,19 @@ #include #include -#include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/simd/fastscan_dispatch.hpp" #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/cpu_features.hpp" namespace rabitqlib::simd { [[noreturn]] static void missing_feature(const char* feature_name) { - throw std::runtime_error(std::string(feature_name) + " requires AVX2/FMA or AVX512 support"); + throw std::runtime_error( + std::string(feature_name) + " requires AVX2/FMA or AVX512 support" + ); } ExcodeIpTable resolve_excode_ip_table() { @@ -121,9 +123,7 @@ void flip_sign(const uint8_t* flip, float* data, size_t dim) { kFlipSignFn(flip, data, dim); } -void kacs_walk(float* data, size_t len) { - kKacsWalkFn(data, len); -} +void kacs_walk(float* data, size_t len) { kKacsWalkFn(data, len); } void scalar_quantize_uint8( uint8_t* result, const float* vec0, size_t dim, float lo, float delta @@ -258,9 +258,7 @@ float excode_ipimpl::ip64_fxu7_avx( return kIp64Fxu7AvxFn(query, compact_code, dim); } -void new_transpose_bin( - const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query -) { +void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query) { kNewTransposeBinFn(q, tq, padded_dim, b_query); } @@ -337,7 +335,8 @@ void accumulate_hacc( namespace rabitqlib { -using WarmupIpX0Q512Fn = float (*)(const uint64_t*, const uint64_t*, float, float, size_t, size_t); +using WarmupIpX0Q512Fn = + float (*)(const uint64_t*, const uint64_t*, float, float, size_t, size_t); const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = [] { if (rabitqlib::cpu::has_avx512_popcnt()) { return rabitqlib::simd::warmup_ip_x0_q_512_avx512; diff --git a/src/simd/fastscan_avx512.cpp b/src/simd/fastscan_avx512.cpp index f9411af..747aff6 100644 --- a/src/simd/fastscan_avx512.cpp +++ b/src/simd/fastscan_avx512.cpp @@ -90,9 +90,9 @@ void transfer_lut_hacc_avx512(const uint16_t* lut, size_t dim, uint8_t* hc_lut) hc_lut + (i / kLutPerIter * kCodePerIter) + ((i % kLutPerIter) * kCodePerLine); uint8_t* fill_hi = fill_lo + (kRegBits / kByteBits); - __m512i tmp = _mm512_cvtepi16_epi32( - _mm256_loadu_si256(reinterpret_cast(lut)) - ); + __m512i tmp = + _mm512_cvtepi16_epi32(_mm256_loadu_si256(reinterpret_cast(lut)) + ); __m128i lo = _mm512_cvtepi32_epi8(tmp); __m128i hi = _mm512_cvtepi32_epi8(_mm512_srli_epi32(tmp, 8)); _mm_store_si128(reinterpret_cast<__m128i*>(fill_lo), lo); diff --git a/src/simd/pack_excode_avx2.cpp b/src/simd/pack_excode_avx2.cpp index 9f795d2..1b8c1ef 100644 --- a/src/simd/pack_excode_avx2.cpp +++ b/src/simd/pack_excode_avx2.cpp @@ -1,5 +1,4 @@ #include "pack_excode_kernels.hpp" - #include "rabitqlib/simd/pack_excode_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/pack_excode_avx512.cpp b/src/simd/pack_excode_avx512.cpp index cff3f99..2f2f1b0 100644 --- a/src/simd/pack_excode_avx512.cpp +++ b/src/simd/pack_excode_avx512.cpp @@ -1,5 +1,4 @@ #include "pack_excode_kernels.hpp" - #include "rabitqlib/simd/pack_excode_dispatch.hpp" namespace rabitqlib::simd { diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index f1f8afa..d21d576 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -1,8 +1,9 @@ -#include "rabitqlib/simd/rotator_dispatch.hpp" - #include #include +#include + +#include "rabitqlib/simd/rotator_dispatch.hpp" namespace rabitqlib::simd { @@ -10,10 +11,10 @@ void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim) { // Process 32 floats (4 AVX2 registers) per iteration constexpr size_t kFloatsPerChunk = 32; - const __m256i bit_select = _mm256_setr_epi32( - 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80 - ); - const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000)); + const __m256i bit_select = + _mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80); + const __m256 sign_flip = + _mm256_castsi256_ps(_mm256_set1_epi32(std::numeric_limits::min())); // Utility lambda to create a mask for flipping signs auto create_mask = [&](uint8_t byte_mask) -> __m256 { diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index 5301197..8dc84db 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -1,8 +1,9 @@ -#include "rabitqlib/simd/rotator_dispatch.hpp" - #include #include +#include + +#include "rabitqlib/simd/rotator_dispatch.hpp" namespace rabitqlib::simd { @@ -29,7 +30,8 @@ void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { _cvtu32_mask16(static_cast((mask_bits >> 48) & 0xFFFF)); // Prepare sign-flip constant - const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); + const __m512 sign_flip = + _mm512_castsi512_ps(_mm512_set1_epi32(std::numeric_limits::min())); // Process 16 floats at a time with each mask segment __m512 vec0 = _mm512_loadu_ps(&data[i]); diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index 1cbd349..d533964 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -64,10 +64,11 @@ void new_transpose_bin_avx2( __m256i vec_48_to_63 = _mm256_loadu_si256((__m256i const*)(q + 48)); // the first (16 - b_query) bits are empty - vec_00_to_15 = _mm256_slli_epi32(vec_00_to_15, (16 - b_query)); - vec_16_to_31 = _mm256_slli_epi32(vec_16_to_31, (16 - b_query)); - vec_32_to_47 = _mm256_slli_epi32(vec_32_to_47, (16 - b_query)); - vec_48_to_63 = _mm256_slli_epi32(vec_48_to_63, (16 - b_query)); + const int shift = static_cast(16 - b_query); + vec_00_to_15 = _mm256_slli_epi32(vec_00_to_15, shift); + vec_16_to_31 = _mm256_slli_epi32(vec_16_to_31, shift); + vec_32_to_47 = _mm256_slli_epi32(vec_32_to_47, shift); + vec_48_to_63 = _mm256_slli_epi32(vec_48_to_63, shift); for (size_t j = 0; j < b_query; ++j) { // pack two 16-bit vectors to 8-bit interleaved vectors @@ -118,11 +119,13 @@ void new_transpose_bin_512_avx2( const uint8_t* current_q_lo = q + i + k * 64; const uint8_t* current_q_hi = q + i + k * 64 + 32; - __m256i vec_lo = _mm256_loadu_si256(reinterpret_cast(current_q_lo)); - __m256i vec_hi = _mm256_loadu_si256(reinterpret_cast(current_q_hi)); + __m256i vec_lo = + _mm256_loadu_si256(reinterpret_cast(current_q_lo)); + __m256i vec_hi = + _mm256_loadu_si256(reinterpret_cast(current_q_hi)); for (size_t j = 0; j < b_query; ++j) { - int bit_idx = b_query - 1 - j; + int bit_idx = static_cast(b_query - 1 - j); __m256i mask_vec = _mm256_set1_epi8(static_cast(1 << bit_idx)); // Process lower 32 bytes diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index bd07e06..e164774 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -92,9 +92,14 @@ void new_transpose_bin_512_avx512( __m512i vec = _mm512_loadu_si512(current_q); for (size_t j = 0; j < b_query; ++j) { - int bit_idx = b_query - 1 - j; - __mmask64 m = _mm512_test_epi8_mask(vec, _mm512_set1_epi8(1 << bit_idx)); - tq[(b_query - j - 1) * num_chunks + k] = reverse_bits_u64(static_cast(m)); + int bit_idx = static_cast(b_query - 1 - j); + // The signed byte preserves the intended one-bit mask, including bit 7. + const char bit_mask = static_cast( + 1U << bit_idx + ); // NOLINT(bugprone-narrowing-conversions) + __mmask64 m = _mm512_test_epi8_mask(vec, _mm512_set1_epi8(bit_mask)); + tq[(b_query - j - 1) * num_chunks + k] = + reverse_bits_u64(static_cast(m)); } } diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index df65cd7..665659e 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -138,7 +138,6 @@ float ip64_fxu3_avx2( contribute_ip(vec_16_to_31, &query[i + 16], sum); contribute_ip(vec_32_to_47, &query[i + 32], sum); contribute_ip(vec_48_to_63, &query[i + 48], sum); - } result = mm256_reduce_add_ps(sum); @@ -173,7 +172,6 @@ float ip64_fxu5_avx2( ) { __m256 sum = _mm256_setzero_ps(); - float result = 0.0F; const __m128i mask = _mm_set1_epi8(0b1111); const __m128i top_mask = _mm_set1_epi8(0b10000); @@ -211,7 +209,6 @@ float ip64_fxu5_avx2( contribute_ip(vec_16_to_31, &query[i + 16], sum); contribute_ip(vec_32_to_47, &query[i + 32], sum); contribute_ip(vec_48_to_63, &query[i + 48], sum); - } result = mm256_reduce_add_ps(sum); @@ -249,7 +246,6 @@ float ip64_fxu6_avx2( contribute_ip(vec_16_to_31, &query[i + 16], sum); contribute_ip(vec_32_to_47, &query[i + 32], sum); contribute_ip(vec_48_to_63, &query[i + 48], sum); - } result = mm256_reduce_add_ps(sum); @@ -261,7 +257,6 @@ float ip64_fxu7_avx2( ) { __m256 sum = _mm256_setzero_ps(); - float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); @@ -305,7 +300,6 @@ float ip64_fxu7_avx2( contribute_ip(vec_16_to_31, &query[i + 16], sum); contribute_ip(vec_32_to_47, &query[i + 32], sum); contribute_ip(vec_48_to_63, &query[i + 48], sum); - } result = mm256_reduce_add_ps(sum); diff --git a/src/simd/space_excode_avx512.cpp b/src/simd/space_excode_avx512.cpp index bb75286..3bd05d5 100644 --- a/src/simd/space_excode_avx512.cpp +++ b/src/simd/space_excode_avx512.cpp @@ -128,7 +128,6 @@ float ip64_fxu3_avx512( q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum = _mm512_fmadd_ps(q, cf, sum); - } result = _mm512_reduce_add_ps(sum); @@ -165,7 +164,6 @@ float ip64_fxu5_avx512( ) { __m512 sum = _mm512_setzero_ps(); - float result = 0.0F; const __m128i mask = _mm_set1_epi8(0b1111); const __m128i top_mask = _mm_set1_epi8(0b10000); @@ -217,7 +215,6 @@ float ip64_fxu5_avx512( q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum = _mm512_fmadd_ps(q, cf, sum); - } result = _mm512_reduce_add_ps(sum); @@ -269,7 +266,6 @@ float ip64_fxu6_avx512( q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum = _mm512_fmadd_ps(q, cf, sum); - } result = _mm512_reduce_add_ps(sum); @@ -281,7 +277,6 @@ float ip64_fxu7_avx512( ) { __m512 sum = _mm512_setzero_ps(); - float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); @@ -339,7 +334,6 @@ float ip64_fxu7_avx512( q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); sum = _mm512_fmadd_ps(q, cf, sum); - } result = _mm512_reduce_add_ps(sum); diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index c38d96c..d31e381 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -11,8 +11,38 @@ namespace rabitqlib::simd { static inline __m256i popcount_avx2(__m256i v) { // Lookup table for population count of 0-15 const __m256i lookup = _mm256_setr_epi8( - 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, - 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4 + 0, + 1, + 1, + 2, + 1, + 2, + 2, + 3, + 1, + 2, + 2, + 3, + 2, + 3, + 3, + 4, + 0, + 1, + 1, + 2, + 1, + 2, + 2, + 3, + 1, + 2, + 2, + 3, + 2, + 3, + 3, + 4 ); const __m256i low_mask = _mm256_set1_epi8(0x0f); @@ -57,17 +87,20 @@ float warmup_ip_x0_q_512_avx2( for (; i < dim_end_512; i += 512) { // Load 64 bytes of data using paired 32-byte loads __m256i data_vec_lo = _mm256_loadu_si256(reinterpret_cast(data)); - __m256i data_vec_hi = _mm256_loadu_si256(reinterpret_cast(data + 4)); - data += 8; // Advance 8 x 64-bit ints (64 bytes) + __m256i data_vec_hi = + _mm256_loadu_si256(reinterpret_cast(data + 4)); + data += 8; // Advance 8 x 64-bit ints (64 bytes) acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_lo)); acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_hi)); for (size_t j = 0; j < b_query; ++j) { // Load 64 bytes of transposed query matching the 512-bit block layout - __m256i query_vec_lo = _mm256_loadu_si256(reinterpret_cast(query)); - __m256i query_vec_hi = _mm256_loadu_si256(reinterpret_cast(query + 4)); - query += 8; // Advance 8 x 64-bit ints (64 bytes) + __m256i query_vec_lo = + _mm256_loadu_si256(reinterpret_cast(query)); + __m256i query_vec_hi = + _mm256_loadu_si256(reinterpret_cast(query + 4)); + query += 8; // Advance 8 x 64-bit ints (64 bytes) __m256i pop_lo = popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); __m256i pop_hi = popcount_avx2(_mm256_and_si256(data_vec_hi, query_vec_hi)); @@ -93,21 +126,25 @@ float warmup_ip_x0_q_512_avx2( // If chunks_lo is 3, limit will be [3,3,3,3,3,3,3,3]. // 3 > seq results in [-1, -1, -1, 0, 0, 0, 0, 0], which is the exact mask needed. __m256i limit_lo = _mm256_set1_epi32(static_cast(chunks_lo)); - __m256i mask_lo = _mm256_cmpgt_epi32(limit_lo, sequence); + __m256i mask_lo = _mm256_cmpgt_epi32(limit_lo, sequence); __m256i limit_hi = _mm256_set1_epi32(static_cast(chunks_hi)); - __m256i mask_hi = _mm256_cmpgt_epi32(limit_hi, sequence); + __m256i mask_hi = _mm256_cmpgt_epi32(limit_hi, sequence); // 3. Vectorized execution continues with zero memory latency - __m256i data_vec_lo = _mm256_maskload_epi32(reinterpret_cast(data), mask_lo); - __m256i data_vec_hi = _mm256_maskload_epi32(reinterpret_cast(data + 4), mask_hi); + __m256i data_vec_lo = + _mm256_maskload_epi32(reinterpret_cast(data), mask_lo); + __m256i data_vec_hi = + _mm256_maskload_epi32(reinterpret_cast(data + 4), mask_hi); acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_lo)); acc_ppc = _mm256_add_epi64(acc_ppc, popcount_avx2(data_vec_hi)); for (size_t j = 0; j < b_query; ++j) { - __m256i query_vec_lo = _mm256_maskload_epi32(reinterpret_cast(query), mask_lo); - __m256i query_vec_hi = _mm256_maskload_epi32(reinterpret_cast(query + 4), mask_hi); + __m256i query_vec_lo = + _mm256_maskload_epi32(reinterpret_cast(query), mask_lo); + __m256i query_vec_hi = + _mm256_maskload_epi32(reinterpret_cast(query + 4), mask_hi); query += num_chunks_64; __m256i pop_lo = popcount_avx2(_mm256_and_si256(data_vec_lo, query_vec_lo)); @@ -135,8 +172,6 @@ float warmup_ip_x0_q_512_avx2( ppc_scalar += mm256_reduce_add_epi64(acc_ppc); return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); - } - } // namespace rabitqlib::simd diff --git a/src/simd/warmup_avx512.cpp b/src/simd/warmup_avx512.cpp index 09d2f36..a94117b 100644 --- a/src/simd/warmup_avx512.cpp +++ b/src/simd/warmup_avx512.cpp @@ -70,8 +70,6 @@ float warmup_ip_x0_q_512_avx512( ppc_scalar += static_cast(_mm512_reduce_add_epi64(acc_ppc)); return (delta * static_cast(ip_scalar)) + (vl * static_cast(ppc_scalar)); - } - } // namespace rabitqlib::simd diff --git a/src/utils/cpu_features.cpp b/src/utils/cpu_features.cpp index 4c84b7d..b4cf53a 100644 --- a/src/utils/cpu_features.cpp +++ b/src/utils/cpu_features.cpp @@ -14,7 +14,9 @@ namespace rabitqlib::cpu { namespace { #if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) -void cpuid(uint32_t leaf, uint32_t subleaf, uint32_t* a, uint32_t* b, uint32_t* c, uint32_t* d) { +void cpuid( + uint32_t leaf, uint32_t subleaf, uint32_t* a, uint32_t* b, uint32_t* c, uint32_t* d +) { int info[4]; __cpuidex(info, static_cast(leaf), static_cast(subleaf)); *a = static_cast(info[0]); @@ -23,7 +25,9 @@ void cpuid(uint32_t leaf, uint32_t subleaf, uint32_t* a, uint32_t* b, uint32_t* *d = static_cast(info[3]); } #elif defined(__x86_64__) || defined(__i386__) -void cpuid(uint32_t leaf, uint32_t subleaf, uint32_t* a, uint32_t* b, uint32_t* c, uint32_t* d) { +void cpuid( + uint32_t leaf, uint32_t subleaf, uint32_t* a, uint32_t* b, uint32_t* c, uint32_t* d +) { __cpuid_count(leaf, subleaf, *a, *b, *c, *d); } #else @@ -72,8 +76,6 @@ bool has_avx512_core() { return detected.avx512f && detected.avx512bw && detected.avx512dq; } -bool has_avx512_popcnt() { - return has_avx512_core() && features().avx512vpopcntdq; -} +bool has_avx512_popcnt() { return has_avx512_core() && features().avx512vpopcntdq; } } // namespace rabitqlib::cpu diff --git a/tests/common/test_data.cpp b/tests/common/test_data.cpp index a38ce51..40f5c14 100644 --- a/tests/common/test_data.cpp +++ b/tests/common/test_data.cpp @@ -1,14 +1,12 @@ #include "test_data.hpp" -#include + #include +#include namespace rabitq_test { std::vector TestDataGenerator::GenerateRandomVector( - size_t dim, - float min, - float max, - unsigned int seed + size_t dim, float min, float max, unsigned int seed ) { std::mt19937 rng(seed); std::uniform_real_distribution dist(min, max); @@ -21,8 +19,7 @@ std::vector TestDataGenerator::GenerateRandomVector( } std::vector TestDataGenerator::GenerateNormalizedVector( - size_t dim, - unsigned int seed + size_t dim, unsigned int seed ) { auto vec = GenerateRandomVector(dim, -1.0f, 1.0f, seed); @@ -44,11 +41,7 @@ std::vector TestDataGenerator::GenerateNormalizedVector( } std::vector> TestDataGenerator::GenerateRandomVectors( - size_t num_vectors, - size_t dim, - float min, - float max, - unsigned int seed + size_t num_vectors, size_t dim, float min, float max, unsigned int seed ) { std::vector> vectors; vectors.reserve(num_vectors); @@ -61,10 +54,7 @@ std::vector> TestDataGenerator::GenerateRandomVectors( } std::vector TestDataGenerator::GenerateGaussianVector( - size_t dim, - float mean, - float stddev, - unsigned int seed + size_t dim, float mean, float stddev, unsigned int seed ) { std::mt19937 rng(seed); std::normal_distribution dist(mean, stddev); @@ -100,4 +90,4 @@ std::vector TestDataGenerator::GenerateIncrementalVector(size_t dim) { return vec; } -} // namespace rabitq_test +} // namespace rabitq_test diff --git a/tests/common/test_data.hpp b/tests/common/test_data.hpp index 65f6ef5..96a34ab 100644 --- a/tests/common/test_data.hpp +++ b/tests/common/test_data.hpp @@ -1,27 +1,21 @@ #ifndef RABITQ_TEST_DATA_HPP #define RABITQ_TEST_DATA_HPP -#include -#include #include +#include +#include namespace rabitq_test { class TestDataGenerator { -public: + public: // Generate random float vector with values in [min, max] static std::vector GenerateRandomVector( - size_t dim, - float min = -1.0f, - float max = 1.0f, - unsigned int seed = 42 + size_t dim, float min = -1.0f, float max = 1.0f, unsigned int seed = 42 ); // Generate random normalized vector (unit length) - static std::vector GenerateNormalizedVector( - size_t dim, - unsigned int seed = 42 - ); + static std::vector GenerateNormalizedVector(size_t dim, unsigned int seed = 42); // Generate multiple random vectors static std::vector> GenerateRandomVectors( @@ -34,10 +28,7 @@ class TestDataGenerator { // Generate Gaussian distributed vector static std::vector GenerateGaussianVector( - size_t dim, - float mean = 0.0f, - float stddev = 1.0f, - unsigned int seed = 42 + size_t dim, float mean = 0.0f, float stddev = 1.0f, unsigned int seed = 42 ); // Generate a simple test vector with known values @@ -53,6 +44,6 @@ class TestDataGenerator { static std::vector GenerateIncrementalVector(size_t dim); }; -} // namespace rabitq_test +} // namespace rabitq_test -#endif // RABITQ_TEST_DATA_HPP +#endif // RABITQ_TEST_DATA_HPP diff --git a/tests/common/test_helpers.hpp b/tests/common/test_helpers.hpp index 35d5052..cb473da 100644 --- a/tests/common/test_helpers.hpp +++ b/tests/common/test_helpers.hpp @@ -1,10 +1,11 @@ #ifndef RABITQ_TEST_HELPERS_HPP #define RABITQ_TEST_HELPERS_HPP +#include + #include -#include #include -#include +#include namespace rabitq_test { @@ -18,7 +19,9 @@ inline bool DoubleNearlyEqual(double a, double b, double epsilon = 1e-10) { } // Vector comparison -inline bool VectorsNearlyEqual(const float* a, const float* b, size_t size, float epsilon = 1e-5f) { +inline bool VectorsNearlyEqual( + const float* a, const float* b, size_t size, float epsilon = 1e-5f +) { for (size_t i = 0; i < size; ++i) { if (!FloatNearlyEqual(a[i], b[i], epsilon)) { return false; @@ -42,7 +45,7 @@ inline float MeanSquaredError(const float* a, const float* b, size_t size) { float diff = a[i] - b[i]; mse += diff * diff; } - return mse / size; + return mse / static_cast(size); } // Calculate dot product @@ -65,20 +68,20 @@ inline float L2Distance(const float* a, const float* b, size_t size) { } // Custom assertion macros -#define ASSERT_FLOAT_NEARLY_EQUAL(a, b, epsilon) \ - ASSERT_TRUE(rabitq_test::FloatNearlyEqual(a, b, epsilon)) \ - << "Expected: " << a << " to be nearly equal to " << b \ - << " (epsilon: " << epsilon << "), but difference was " << std::abs(a - b) +#define ASSERT_FLOAT_NEARLY_EQUAL(a, b, epsilon) \ + ASSERT_TRUE(rabitq_test::FloatNearlyEqual((a), (b), (epsilon))) \ + << "Expected: " << (a) << " to be nearly equal to " << (b) \ + << " (epsilon: " << (epsilon) << "), but difference was " << std::abs((a) - (b)) -#define EXPECT_FLOAT_NEARLY_EQUAL(a, b, epsilon) \ - EXPECT_TRUE(rabitq_test::FloatNearlyEqual(a, b, epsilon)) \ - << "Expected: " << a << " to be nearly equal to " << b \ - << " (epsilon: " << epsilon << "), but difference was " << std::abs(a - b) +#define EXPECT_FLOAT_NEARLY_EQUAL(a, b, epsilon) \ + EXPECT_TRUE(rabitq_test::FloatNearlyEqual((a), (b), (epsilon))) \ + << "Expected: " << (a) << " to be nearly equal to " << (b) \ + << " (epsilon: " << (epsilon) << "), but difference was " << std::abs((a) - (b)) -#define ASSERT_VECTORS_NEARLY_EQUAL(a, b, size, epsilon) \ - ASSERT_TRUE(rabitq_test::VectorsNearlyEqual(a, b, size, epsilon)) \ - << "Vectors are not nearly equal (epsilon: " << epsilon << ")" +#define ASSERT_VECTORS_NEARLY_EQUAL(a, b, size, epsilon) \ + ASSERT_TRUE(rabitq_test::VectorsNearlyEqual((a), (b), (size), (epsilon))) \ + << "Vectors are not nearly equal (epsilon: " << (epsilon) << ")" -} // namespace rabitq_test +} // namespace rabitq_test -#endif // RABITQ_TEST_HELPERS_HPP +#endif // RABITQ_TEST_HELPERS_HPP diff --git a/tests/integration/bit_pack_unpack_test.cpp b/tests/integration/bit_pack_unpack_test.cpp index 90ec157..5401ae8 100644 --- a/tests/integration/bit_pack_unpack_test.cpp +++ b/tests/integration/bit_pack_unpack_test.cpp @@ -1,15 +1,16 @@ #include -#include -#include + #include -#include #include #include +#include +#include +#include using namespace rabitqlib; class BitPackUnpackTest : public ::testing::Test { -protected: + protected: // 1. Shared Constants & Data Structures const size_t dim = 768; std::vector query; @@ -19,13 +20,13 @@ class BitPackUnpackTest : public ::testing::Test { // 2. Common Initialization (Runs before EVERY test) void SetUp() override { srand(42); - + // Initialize Query (Same for all tests) query.resize(dim); - for(size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < dim; ++i) { query[i] = static_cast((rand() * 100.0) / RAND_MAX); } - + // Pre-allocate code vector code.resize(dim); } @@ -37,7 +38,7 @@ class BitPackUnpackTest : public ::testing::Test { // Generate random codes based on bit depth for (size_t i = 0; i < dim; ++i) { - code[i] = rand() % (1 << bits); + code[i] = rand() % (1 << bits); } // Pack the code @@ -50,7 +51,7 @@ class BitPackUnpackTest : public ::testing::Test { float CalculateExpected() const { float expected_result = 0.0f; for (size_t i = 0; i < dim; ++i) { - expected_result += query[i] * code[i]; + expected_result += query[i] * static_cast(code[i]); } return expected_result; } @@ -64,79 +65,71 @@ class BitPackUnpackTest : public ::testing::Test { // --- Test Cases --- TEST_F(BitPackUnpackTest, ExCode1Bit) { - PrepareData(1); // Set up for 1-bit + PrepareData(1); // Set up for 1-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip16_fxu1_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip16_fxu1_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode2Bit) { - PrepareData(2); // Set up for 2-bit + PrepareData(2); // Set up for 2-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip64_fxu2_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip64_fxu2_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode3Bit) { - PrepareData(3); // Set up for 3-bit + PrepareData(3); // Set up for 3-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip64_fxu3_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip64_fxu3_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode4Bit) { - PrepareData(4); // Set up for 4-bit + PrepareData(4); // Set up for 4-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip16_fxu4_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip16_fxu4_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode5Bit) { - PrepareData(5); // Set up for 5-bit + PrepareData(5); // Set up for 5-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip64_fxu5_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip64_fxu5_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode6Bit) { - PrepareData(6); // Set up for 6-bit + PrepareData(6); // Set up for 6-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip64_fxu6_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip64_fxu6_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } TEST_F(BitPackUnpackTest, ExCode7Bit) { - PrepareData(7); // Set up for 7-bit + PrepareData(7); // Set up for 7-bit // Run the AVX function - float result = rabitqlib::excode_ipimpl::ip64_fxu7_avx( - query.data(), compact_code.data(), dim - ); + float result = + rabitqlib::excode_ipimpl::ip64_fxu7_avx(query.data(), compact_code.data(), dim); ExpectIpNear(result); } - diff --git a/tests/python/conftest.py b/tests/python/conftest.py index 94dfb80..922668a 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -16,6 +16,7 @@ # ── data fixtures ───────────────────────────────────────────────────────────── + @pytest.fixture(scope="session") def base_data() -> np.ndarray: """500 random float32 vectors of dimension 64.""" @@ -42,6 +43,7 @@ def clusters(base_data: np.ndarray): # ── pure-numpy helpers ──────────────────────────────────────────────────────── + def brute_force_knn(data: np.ndarray, queries: np.ndarray, k: int): """Exact L2 k-NN. Returns (ids, sq_dists) of shape (nq, k). @@ -59,7 +61,5 @@ def brute_force_knn(data: np.ndarray, queries: np.ndarray, k: int): def recall_at_k(approx_ids: np.ndarray, exact_ids: np.ndarray, k: int) -> float: """Fraction of exact top-k neighbors found in the approximate top-k results.""" nq = approx_ids.shape[0] - hits = sum( - len(set(approx_ids[i, :k]) & set(exact_ids[i, :k])) for i in range(nq) - ) + hits = sum(len(set(approx_ids[i, :k]) & set(exact_ids[i, :k])) for i in range(nq)) return hits / (nq * k) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index a4ba251..c9b96b8 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -2,16 +2,16 @@ import numpy as np import pytest +from conftest import DIM, N_CLUSTERS, N_QUERIES, N_VECTORS, brute_force_knn, recall_at_k from rabitqlib import HnswIndex -from conftest import DIM, N_CLUSTERS, N_VECTORS, N_QUERIES, brute_force_knn, recall_at_k - _TOPK = 10 _EF = 50 # generous ef for correctness tests on a small dataset # ── fixtures ────────────────────────────────────────────────────────────────── + @pytest.fixture(scope="module") def built_hnsw(base_data, clusters): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) @@ -22,6 +22,7 @@ def built_hnsw(base_data, clusters): # ── construction ────────────────────────────────────────────────────────────── + def test_is_built(built_hnsw): assert built_hnsw.is_built @@ -43,6 +44,7 @@ def test_fast_quantization_builds(base_data, clusters): # ── search output shape and dtype ───────────────────────────────────────────── + def test_search_output_shape(built_hnsw, query_data): ids, dists = built_hnsw.search(query_data, k=_TOPK, ef=_EF) assert ids.shape == (N_QUERIES, _TOPK) @@ -78,6 +80,7 @@ def test_single_query(built_hnsw, query_data): # ── search correctness ──────────────────────────────────────────────────────── + def test_self_retrieval(built_hnsw, base_data): """Each database vector must be its own nearest neighbor at high ef.""" probes = base_data[:10] @@ -97,6 +100,7 @@ def test_recall_vs_brute_force(built_hnsw, base_data, query_data): # ── error handling ──────────────────────────────────────────────────────────── + def test_wrong_data_dim_raises(clusters): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters @@ -113,6 +117,7 @@ def test_wrong_query_dim_raises(built_hnsw): # ── save / load roundtrip ───────────────────────────────────────────────────── + def test_save_load_roundtrip(built_hnsw, query_data, tmp_path): path = str(tmp_path / "hnsw.index") built_hnsw.save(path) diff --git a/tests/python/test_import.py b/tests/python/test_import.py index f5b8ff6..9508eb0 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -3,9 +3,9 @@ import pytest from rabitqlib import HnswIndex, IvfIndex, SymqgIndex - # ── import ──────────────────────────────────────────────────────────────────── + def test_all_classes_importable(): assert HnswIndex is not None assert IvfIndex is not None @@ -14,6 +14,7 @@ def test_all_classes_importable(): # ── valid construction ──────────────────────────────────────────────────────── + def test_ivf_construct_l2(): idx = IvfIndex(64, 500, 5, nbits=4, metric="l2") assert idx.dim == 64 @@ -35,7 +36,9 @@ def test_hnsw_construct_defaults(): def test_hnsw_construct_explicit(): - idx = HnswIndex(64, 500, M=8, ef_construction=50, nbits=4, metric="ip", random_seed=7) + idx = HnswIndex( + 64, 500, M=8, ef_construction=50, nbits=4, metric="ip", random_seed=7 + ) assert idx.metric == "ip" assert idx.nbits == 4 @@ -49,11 +52,15 @@ def test_symqg_construct(): # ── invalid metric ──────────────────────────────────────────────────────────── -@pytest.mark.parametrize("cls,args", [ - (IvfIndex, (64, 500, 5, 4)), - (HnswIndex, (64, 500)), - (SymqgIndex, (64, 16)), -]) + +@pytest.mark.parametrize( + "cls,args", + [ + (IvfIndex, (64, 500, 5, 4)), + (HnswIndex, (64, 500)), + (SymqgIndex, (64, 16)), + ], +) def test_invalid_metric_raises(cls, args): with pytest.raises(Exception, match="[Uu]nsupported metric|[Ii]nvalid"): cls(*args, metric="cosine") diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index f0cf161..607a6c5 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -2,16 +2,16 @@ import numpy as np import pytest +from conftest import DIM, N_CLUSTERS, N_QUERIES, N_VECTORS, brute_force_knn, recall_at_k from rabitqlib import IvfIndex -from conftest import DIM, N_CLUSTERS, N_VECTORS, N_QUERIES, brute_force_knn, recall_at_k - _TOPK = 10 _NPROBE_ALL = N_CLUSTERS # probe every cluster → deterministic coverage # ── fixtures ────────────────────────────────────────────────────────────────── + @pytest.fixture(scope="module") def built_ivf(base_data, clusters): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) @@ -22,6 +22,7 @@ def built_ivf(base_data, clusters): # ── construction ────────────────────────────────────────────────────────────── + def test_is_built(built_ivf): assert built_ivf.is_built @@ -43,6 +44,7 @@ def test_fast_quantization_builds(base_data, clusters): # ── search output shape and dtype ───────────────────────────────────────────── + def test_search_output_shape(built_ivf, query_data): ids, dists = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL) assert ids.shape == (N_QUERIES, _TOPK) @@ -79,6 +81,7 @@ def test_k_equals_one(built_ivf, query_data): # ── search correctness ──────────────────────────────────────────────────────── + def test_self_retrieval(built_ivf, base_data): """Probing all clusters: each database vector must be its own nearest neighbor.""" probes = base_data[:10] @@ -98,6 +101,7 @@ def test_recall_vs_brute_force(built_ivf, base_data, query_data): # ── optional parameters ─────────────────────────────────────────────────────── + def test_high_accuracy_false(built_ivf, query_data): ids, dists = built_ivf.search( query_data, k=_TOPK, nprobe=_NPROBE_ALL, high_accuracy=False @@ -106,14 +110,19 @@ def test_high_accuracy_false(built_ivf, query_data): def test_multithreaded_search_matches_single(built_ivf, query_data): - ids1, dists1 = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=1) - ids2, dists2 = built_ivf.search(query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=2) + ids1, dists1 = built_ivf.search( + query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=1 + ) + ids2, dists2 = built_ivf.search( + query_data, k=_TOPK, nprobe=_NPROBE_ALL, num_threads=2 + ) np.testing.assert_array_equal(ids1, ids2) np.testing.assert_allclose(dists1, dists2, rtol=1e-5) # ── error handling ──────────────────────────────────────────────────────────── + def test_wrong_data_dim_raises(clusters): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) centroids, cluster_ids = clusters @@ -146,6 +155,7 @@ def test_wrong_query_dim_raises(built_ivf): # ── save / load roundtrip ───────────────────────────────────────────────────── + def test_save_load_roundtrip(built_ivf, query_data, tmp_path): path = str(tmp_path / "ivf.index") built_ivf.save(path) diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index fa28b4c..3603bcb 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -2,10 +2,9 @@ import numpy as np import pytest +from conftest import DIM, N_QUERIES, N_VECTORS, brute_force_knn, recall_at_k from rabitqlib import SymqgIndex -from conftest import DIM, N_VECTORS, N_QUERIES, brute_force_knn, recall_at_k - _TOPK = 10 _EF = 50 _MAX_DEGREE = 32 # must be a multiple of fastscan::kBatchSize (32) @@ -14,6 +13,7 @@ # ── fixtures ────────────────────────────────────────────────────────────────── + @pytest.fixture(scope="module") def built_symqg(base_data): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) @@ -23,6 +23,7 @@ def built_symqg(base_data): # ── construction ────────────────────────────────────────────────────────────── + def test_is_built(built_symqg): assert built_symqg.is_built @@ -36,6 +37,7 @@ def test_properties(built_symqg): # ── search output shape and dtype ───────────────────────────────────────────── + def test_search_output_shape(built_symqg, query_data): ids, dists = built_symqg.search(query_data, k=_TOPK, ef=_EF) assert ids.shape == (N_QUERIES, _TOPK) @@ -65,6 +67,7 @@ def test_single_query(built_symqg, query_data): # ── search correctness ──────────────────────────────────────────────────────── + def test_self_retrieval(built_symqg, base_data): """Each database vector must be its own nearest neighbor at high ef.""" probes = base_data[:10] @@ -84,6 +87,7 @@ def test_recall_vs_brute_force(built_symqg, base_data, query_data): # ── error handling ──────────────────────────────────────────────────────────── + def test_wrong_data_dim_raises(): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) @@ -106,6 +110,7 @@ def test_search_before_build_raises(): # ── save / load roundtrip ───────────────────────────────────────────────────── + def test_save_load_roundtrip(built_symqg, query_data, tmp_path): path = str(tmp_path / "symqg.index") built_symqg.save(path) diff --git a/tests/unit/rabitqlib/utils/cpu_features_test.cpp b/tests/unit/rabitqlib/utils/cpu_features_test.cpp index 58d9b34..974ce60 100644 --- a/tests/unit/rabitqlib/utils/cpu_features_test.cpp +++ b/tests/unit/rabitqlib/utils/cpu_features_test.cpp @@ -1,11 +1,14 @@ -#include - #include "rabitqlib/utils/cpu_features.hpp" +#include + TEST(CpuFeatures, returns_stable_result) { const auto& first = rabitqlib::cpu::features(); const auto& second = rabitqlib::cpu::features(); ASSERT_EQ(&first, &second); - ASSERT_EQ(rabitqlib::cpu::has_avx512_popcnt(), rabitqlib::cpu::has_avx512_core() && first.avx512vpopcntdq); + ASSERT_EQ( + rabitqlib::cpu::has_avx512_popcnt(), + rabitqlib::cpu::has_avx512_core() && first.avx512vpopcntdq + ); } diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 160e49c..ec2d88a 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -1,17 +1,21 @@ -#include #include "rabitqlib/utils/rotator.hpp" -#include "test_helpers.hpp" -#include "test_data.hpp" -#include + +#include + #include -#include #include +#include +#include +#include + +#include "test_data.hpp" +#include "test_helpers.hpp" using namespace rabitqlib; using namespace rabitq_test; class RotatorTest : public ::testing::Test { -protected: + protected: void SetUp() override { dim = 128; test_data = TestDataGenerator::GenerateRandomVector(dim, -1.0f, 1.0f, 42); @@ -28,15 +32,13 @@ class RotatorTest : public ::testing::Test { // Test that FhtKacRotator is chosen by default TEST_F(RotatorTest, DefaultRotatorType) { - Rotator* rotator = choose_rotator(dim); + std::unique_ptr> rotator(choose_rotator(dim)); ASSERT_NE(rotator, nullptr); // FhtKacRotator pads to multiple of 64 size_t padded_dim = rotator->size(); EXPECT_EQ(padded_dim % 64, 0); EXPECT_GE(padded_dim, dim); - - delete rotator; } uint8_t bitreverse8(uint8_t x) { @@ -67,7 +69,7 @@ TEST(FlipSignTest, FlipWorks) { ASSERT_EQ(abs(data[i]), static_cast(i + 1)); int sign = (data[i] < 0) ? 1 : 0; signs = (signs << 1) | sign; - if(i%8 == 7) { + if (i % 8 == 7) { uint8_t expected = flip[i / 8]; signs = bitreverse8(signs); ASSERT_EQ(static_cast(signs & 0xFF), expected); diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index a5449b7..5125a6d 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -1,12 +1,15 @@ +#include "rabitqlib/utils/space.hpp" + #include + +#include +#include + +#include "rabitqlib/defines.hpp" #include "rabitqlib/simd/space_dispatch.hpp" #include "rabitqlib/utils/cpu_features.hpp" -#include "rabitqlib/utils/space.hpp" -#include "rabitqlib/defines.hpp" -#include "test_helpers.hpp" #include "test_data.hpp" -#include -#include +#include "test_helpers.hpp" using namespace rabitqlib; using namespace rabitq_test; @@ -62,7 +65,8 @@ TEST(ScalarQuantize, Uint8MatchesRoundedScalar) { std::vector expected(dim); for (size_t i = 0; i < dim; ++i) { - float quantized = static_cast((i * 7) % 251) + (static_cast(i % 3) - 1) * 0.2F; + float quantized = static_cast((i * 7) % 251) + + static_cast(static_cast(i % 3) - 1) * 0.2F; input[i] = lo + delta * quantized; expected[i] = static_cast(std::round((input[i] - lo) / delta)); } @@ -81,7 +85,8 @@ TEST(ScalarQuantize, Uint16MatchesRoundedScalar) { std::vector expected(dim); for (size_t i = 0; i < dim; ++i) { - float quantized = static_cast(1000 + i * 317) + (static_cast(i % 5) - 2) * 0.1F; + float quantized = static_cast(1000 + i * 317) + + static_cast(static_cast(i % 5) - 2) * 0.1F; input[i] = lo + delta * quantized; expected[i] = static_cast(std::round((input[i] - lo) / delta)); } @@ -95,8 +100,8 @@ TEST(ip16_fxu1_avx, ip_works) { srand(42); size_t dim = 64; float query[dim]; - uint8_t codes[dim/8]; - + uint8_t codes[dim / 8]; + for (size_t i = 0; i < dim; ++i) { query[i] = static_cast(rand()) / RAND_MAX * 1000.0f; } @@ -105,15 +110,17 @@ TEST(ip16_fxu1_avx, ip_works) { codes[i] = static_cast(rand() % 256); } - ASSERT_NEAR(rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, dim), 15055.81f, 0.1f); + ASSERT_NEAR( + rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, dim), 15055.81f, 0.1f + ); } TEST(ip64_fxu2_avx, ip_works) { srand(42); - size_t dim = 64*4; + size_t dim = 64 * 4; float query[dim]; - uint8_t codes[dim/4]; - + uint8_t codes[dim / 4]; + for (size_t i = 0; i < dim; ++i) { query[i] = static_cast(rand()) / RAND_MAX * 1000.0f; } @@ -121,7 +128,9 @@ TEST(ip64_fxu2_avx, ip_works) { for (size_t i = 0; i < dim / 4; ++i) { codes[i] = static_cast(rand() % 256); } - ASSERT_NEAR(rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), 217584.15f, 0.1f); + ASSERT_NEAR( + rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), 217584.15f, 0.1f + ); } TEST(ip_fxu8_avx, ip_works) {