diff --git a/.clang-format b/.clang-format index 65c1b641..4496ef2b 100644 --- a/.clang-format +++ b/.clang-format @@ -25,6 +25,8 @@ DerivePointerAlignment: false FixNamespaceComments: true IncludeBlocks: Regroup IndentCaseLabels: true +# Brace every control statement body; this was the dominant style before it was enforced. +InsertBraces: true KeepEmptyLinesAtTheStartOfBlocks: false PointerAlignment: Left ReferenceAlignment: Left diff --git a/.clang-tidy b/.clang-tidy index 45f8ae34..6ec0d237 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -31,13 +31,13 @@ Checks: > # LLVM 23's new wildcard matches must not change enum ABI, replace the project's # pragma-once convention, or reject the CPU detector's required inline assembly. -# Keep the established checks blocking. Naming starts as warnings while existing -# internal/parameter names are migrated; do not bulk-fix public interfaces. -WarningsAsErrors: "*,-readability-identifier-naming" +# Every enabled check is blocking, including naming. Retained public spellings are +# listed in the IgnoredRegexp options below; do not bulk-fix public interfaces. +WarningsAsErrors: "*" CheckOptions: - # Check locals, parameters and private fields. Public member/function names, - # types, enums and macros retain their current compatibility contracts. + # Public names that keep their compatibility contracts are listed in the + # IgnoredRegexp options below instead of being renamed. readability-identifier-naming.LocalVariableCase: lower_case readability-identifier-naming.ParameterCase: lower_case # M is the existing HNSW graph-degree parameter, including its public constructor. @@ -50,4 +50,35 @@ CheckOptions: readability-identifier-naming.PrivateMemberSuffix: '_' # Preserve these existing graph parameters and formula coefficient names. # Keep exceptions explicit rather than exempting all capitalized identifiers. - readability-identifier-naming.PrivateMemberIgnoredRegexp: '^(M|maxM|maxM0|G_add|G_error|G_k1xSumq|G_kbaseSumq|G_kbxSumq)_$' + readability-identifier-naming.PrivateMemberIgnoredRegexp: '^(M|G_add|G_error|G_k1xSumq|G_kbaseSumq|G_kbxSumq)_$' + readability-identifier-naming.ProtectedMemberCase: lower_case + readability-identifier-naming.ProtectedMemberSuffix: '_' + readability-identifier-naming.PublicMemberCase: lower_case + # Functions, types, namespaces, constants and macros follow AGENTS.md. The + # ignored names are retained public API or names the standard library requires. + readability-identifier-naming.FunctionCase: lower_case + readability-identifier-naming.MethodCase: lower_case + readability-identifier-naming.MethodIgnoredRegexp: '^M$' + readability-identifier-naming.NamespaceCase: lower_case + readability-identifier-naming.ClassCase: CamelCase + readability-identifier-naming.StructCase: CamelCase + readability-identifier-naming.StructIgnoredRegexp: '^rebind$' + readability-identifier-naming.EnumCase: CamelCase + readability-identifier-naming.EnumConstantCase: CamelCase + readability-identifier-naming.EnumConstantIgnoredRegexp: '^(METRIC_L2|METRIC_IP|RECONSTRUCTION|UNBIASED_ESTIMATION|PLAIN)$' + readability-identifier-naming.TypeAliasCase: CamelCase + readability-identifier-naming.TypeAliasIgnoredRegexp: '^(value_type|pointer|reference|difference_type|iterator_category|is_always_equal|other|maxheap|minheap|ex_ipfunc)$' + readability-identifier-naming.TypedefCase: CamelCase + readability-identifier-naming.TemplateParameterCase: CamelCase + readability-identifier-naming.GlobalConstantCase: CamelCase + readability-identifier-naming.GlobalConstantPrefix: k + readability-identifier-naming.StaticConstantCase: CamelCase + readability-identifier-naming.StaticConstantPrefix: k + readability-identifier-naming.ClassConstantCase: CamelCase + readability-identifier-naming.ClassConstantPrefix: k + readability-identifier-naming.MacroDefinitionCase: UPPER_CASE + # matrix_kernels.hpp temporarily renames the Eigen namespace for each ISA build. + readability-identifier-naming.MacroDefinitionIgnoredRegexp: '^Eigen$' + # Named parameters must match between declarations and definitions, without the + # default prefix/suffix leniency. Unnamed declaration parameters remain allowed. + readability-inconsistent-declaration-parameter-name.Strict: true diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 41d0809a..37c6a232 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -360,10 +360,10 @@ jobs: env: OMP_NUM_THREADS: "2" run: | - simd_filter='CpuFeatures.*:FloatDistance.*:MatrixDispatchTest.*:FastScan*.*:BatchEstimatorTest.*' - simd_filter+=':FhtDispatchTest.*:FlipSignTest.*:KacsWalkTest.*:ExcodeIp.*:MaskIpX0Q.*' - simd_filter+=':PortableKernels.*:QueryPreparation.*:WarmupIpX0Q.*:PackBinary.*' - simd_filter+=':ExplicitSimd/RabitqRescaleBackendTest.*:QGEstimatorTest.*' + simd_filter='CpuFeaturesTest.*:FloatDistanceTest.*:MatrixDispatchTest.*:FastScan*.*' + simd_filter+=':FhtDispatchTest.*:FlipSignTest.*:KacsWalkTest.*:ExcodeIpTest.*:MaskIpX0QTest.*' + simd_filter+=':PortableKernelsTest.*:QueryPreparationTest.*:WarmupIpX0QTest.*:PackBinaryTest.*' + simd_filter+=':ExplicitSimd/RabitqRescaleBackendTest.*:QgEstimatorTest.*' simd_filter+=':QGKMeansTest.TrainsAndReturnsAssignmentsForFinalCentroids' simd_filter+=':RaBitQKMeansTest.FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts' simd_filter+=':IvfPersistenceTest.ReloadsRawAndQuantizedStorage' diff --git a/benchmarks/core/main.cpp b/benchmarks/core/main.cpp index 745e9ffa..25f119de 100644 --- a/benchmarks/core/main.cpp +++ b/benchmarks/core/main.cpp @@ -68,13 +68,13 @@ int run(int argc, char** argv) { const size_t count = argc == 1 ? 4096 : number(argv[3]); const size_t queries = argc == 1 ? 64 : number(argv[4]); const size_t repeats = argc == 1 ? 5 : number(argv[5]); - constexpr size_t k = 10; + constexpr size_t kTopK = 10; if (dim < 64 || dim % 32 != 0) { throw std::invalid_argument( "This FHT/Kac benchmark requires dim >= 64 and divisible by 32" ); } - if (count < k) { + if (count < kTopK) { throw std::invalid_argument("At least ten vectors are needed for recall@10"); } const api::Parameters parameters{dim, bits, rabitqlib::METRIC_L2}; @@ -147,10 +147,10 @@ int run(int argc, char** argv) { squared_error += error * error; max_error = std::max(max_error, std::abs(error)); } - std::partial_sort(exact.begin(), exact.begin() + k, exact.end()); - std::partial_sort(estimated.begin(), estimated.begin() + k, estimated.end()); - for (size_t i = 0; i < k; ++i) { - for (size_t j = 0; j < k; ++j) { + std::partial_sort(exact.begin(), exact.begin() + kTopK, exact.end()); + std::partial_sort(estimated.begin(), estimated.begin() + kTopK, estimated.end()); + for (size_t i = 0; i < kTopK; ++i) { + for (size_t j = 0; j < kTopK; ++j) { hits += estimated[i].second == exact[j].second ? 1 : 0; } } @@ -167,7 +167,7 @@ int run(int argc, char** argv) { << '\n' << "# rmse=" << std::sqrt(squared_error / static_cast(comparisons)) << " max_abs_error=" << max_error << " flat_recall_at_10=" - << static_cast(hits) / static_cast(product(queries, k)) << '\n' + << static_cast(hits) / static_cast(product(queries, kTopK)) << '\n' << "operation,repeat,operations,seconds,ops_per_second,ns_per_operation,checksum\n"; measure("encode_packed", count, repeats, [&] { double checksum = 0; diff --git a/benchmarks/hotpaths/core.cpp b/benchmarks/hotpaths/core.cpp index 9ec22e55..c3800d5b 100644 --- a/benchmarks/hotpaths/core.cpp +++ b/benchmarks/hotpaths/core.cpp @@ -116,8 +116,9 @@ void run(int argc, char** argv) { bench::product(rows, iterations), reps, [&] { - for (size_t i = 0; i < iterations; ++i) + for (size_t i = 0; i < iterations; ++i) { encode(threads > 1); + } }, checksum ); diff --git a/benchmarks/hotpaths/fastscan.cpp b/benchmarks/hotpaths/fastscan.cpp index 5dedb531..ca17b7cf 100644 --- a/benchmarks/hotpaths/fastscan.cpp +++ b/benchmarks/hotpaths/fastscan.cpp @@ -104,8 +104,9 @@ void run(int argc, char** argv) { bench::product(batches, kRows), reps, [&] { - for (size_t batch = 0; batch < batches; ++batch) + for (size_t batch = 0; batch < batches; ++batch) { score(); + } }, [&] { uint64_t h = bench::hash(distance.data(), sizeof(distance)); diff --git a/benchmarks/hotpaths/load.cpp b/benchmarks/hotpaths/load.cpp index 93ff78d0..8e4d0de2 100644 --- a/benchmarks/hotpaths/load.cpp +++ b/benchmarks/hotpaths/load.cpp @@ -28,8 +28,9 @@ void run(int argc, char** argv) { if (mode != "prepare" && mode != "load") { throw std::invalid_argument("Mode must be prepare or load"); } - if (count < 8) + if (count < 8) { throw std::invalid_argument("At least eight centroids are required"); + } const auto data = bench::values(bench::product(count, dim)); if (mode == "prepare") { // Repeated smoke tests reuse their fixture; never replace an existing file. @@ -65,11 +66,12 @@ void run(int argc, char** argv) { h = bench::hash(&candidate.id, sizeof(candidate.id), h); h = bench::hash(&candidate.distance, sizeof(candidate.distance), h); } - if (rep == 0) + if (rep == 0) { expected = h; - else { - if (h != expected) + } else { + if (h != expected) { throw std::runtime_error("Loaded search results changed"); + } bench::row("sidecar_load", rep - 1, 1, seconds, h); } } diff --git a/docs/development.md b/docs/development.md index 34714d56..c5d6832a 100644 --- a/docs/development.md +++ b/docs/development.md @@ -602,7 +602,8 @@ preprocessing objects, `visited` for visited state. **Search windows** -`ef_search` for a per-call graph window, `kDefaultEfSearch` for HNSW's fixed fallback, and +`ef_search` for a per-call graph window, `default_ef_search_` for QG's stored default, +`kDefaultEfSearch` for HNSW's fixed fallback, and `ef_construction` for graph-building search breadth. IVF's `nprobe` counts clusters and remains distinct. @@ -640,7 +641,8 @@ of 32. These settings are not interchangeable. Public names and Python keywords are compatibility contracts. For example, `Cluster::num()`, the protected `Initializer::num_cluster_`, QG's `degree_bound()`/`set_ef()`, Python's `ef`, and `QGKMeansParameters::ef_build` -retain their existing spellings. C++ parameter names may be clarified without +retain their existing spellings, as do the Python vocabulary `random_seed`/`seed`, +`train`/`build`, `obj`, `nsplit`, `final_obj` and `niter`. C++ parameter names may be clarified without changing types, order or defaults. Keep declarations, definitions and comments in agreement; do not rename methods or literals by textual substitution. @@ -666,6 +668,8 @@ public and protected names, Python keywords, field order and numerical operation - Clustering's private state uses `dim_`, `num_points_`, `num_clusters_`, `num_workers_` and `data_`. Keep public `d`, `k`, `num_threads` and graph parameter names compatible. Retain the existing private `niter_` graph iteration name. + Per-point centroid IDs are `assignments` (arrays) and `cluster_id` (scalars); + `label` stays reserved for HNSW's caller-visible labels. - FHT backend arguments use `rotated_data`, `padded_dim`, `trunc_dim` and `scale_factor` consistently across declarations, dispatch and implementations. A lower-level kernel's `dim` still denotes that kernel's own input shape. @@ -674,6 +678,37 @@ Keep short mathematical symbols, loop indices and SIMD register names when their meaning is local and clear. Naming cleanup must not alter formulas, arithmetic order, packed layouts, ISA selection or runtime behavior. +### Library-wide code conventions + +These rules apply to new and changed code; the listed exceptions predate them. + +- **Identifiers.** `.clang-tidy` checks the case of every identifier kind and lists + the public spellings kept for compatibility; extend those lists only for retained + public API. Naming findings fail the clang-tidy check like any other. Template + parameters are PascalCase; `T` is the floating-point element type and `TP` the + code or ID element type. +- **Linkage.** Follow [Implementation linkage](#implementation-linkage) for code + compiled under several ISA flags, including public headers that ISA-flagged sources + include. Other header functions are plain `inline`, and helpers used by one `.cpp` + file go in an anonymous namespace. +- **Namespaces.** Internal code uses `::detail` or an anonymous namespace. + Existing public `*_impl` namespaces are compatibility names, not a pattern. +- **Errors.** New messages read `" "` with the component `HNSW`, + `IVF`, `QuantizedGraph`, `PiPNN`, `FastScan` or `FastScan high-accuracy`, no colon + after it and no trailing period. Core keeps identifier-first messages. Python binding + argument errors start with the lowercase parameter name and state errors with the + wrapper class name. Some older index-file messages (for example `Invalid padded + dimension in IVF index file`) are asserted by tests and keep their wording; changing + a Python-visible message is a behavior change. Allocation failure throws + `std::bad_alloc`. +- **Comments.** Use `//`; keep `/** */` for public API documentation blocks and state + preconditions as sentences (`// Requires ...`). +- **Tests.** Name new suites `Test` (`Hnsw`, `Ivf`, `Qg`, `FastScan`, ...) + with PascalCase test names and no underscores. Older suites such as `ThreadCount` + keep their names because CI filters select them. Test constants use `kDim`, `kCount`, + `kNumQueries` and `kTopK`. Renaming a suite requires updating the GoogleTest filters + in `.github/workflows/test.yaml` and `tests/CMakeLists.txt`. + ## Technical contracts ### Module build boundaries @@ -807,7 +842,7 @@ nothing in an old file is reinterpreted. A file that has removals is rejected by through 0.5.1, whose `load` range-checks every cluster ID; that rejection is the compatibility path. This mark is unrelated to the `SearchBuffer` ID marker. Keep these rules when touching HNSW-RaBitQ code: -- Read cluster IDs only through `get_clusterid_by_internalid`, which masks the bit. Estimators, +- Read cluster IDs only through `cluster_id_of`, which masks the bit. Estimators, routing, and the `add` reconstruction index centroids with it. - A removed point stays in the graph: base-layer search still inserts it into `search_pool`, and `add` may link to it. Only the result buffer skips it, including when it is the entry point. @@ -887,6 +922,8 @@ a short-vector improvement does not justify applying it to wider vectors. Put calculations and scratch-storage helpers outside the dispatcher. Shared implementation headers use internal linkage so independently compiled backends retain their own bodies. +This includes public headers reached from ISA-flagged sources, such as the per-ISA HNSW +search files: keep their non-template helpers `static inline` rather than plain `inline`. The AVX FHT intrinsics in `src/simd/fht_kernels.hpp` use the private kernel namespace for the same reason. Rotation uses this implementation on GCC, Clang, and MSVC; no assembler is required. Keep the FFHT attribution and MIT license in the diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index 40ec3912..26d4db11 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -56,7 +56,7 @@ namespace detail { } // Float vectors accumulated in double, in index order. -[[nodiscard]] inline double squared_norm_f64(const float* vector, size_t d) { +[[nodiscard]] inline double l2norm_sqr_f64(const float* vector, size_t d) { double result = 0.0; for (size_t dim = 0; dim < d; ++dim) { const double value = vector[dim]; @@ -65,7 +65,7 @@ namespace detail { return result; } -[[nodiscard]] inline double squared_l2_f64(const float* a, const float* b, size_t d) { +[[nodiscard]] inline double euclidean_sqr_f64(const float* a, const float* b, size_t d) { double result = 0.0; for (size_t dim = 0; dim < d; ++dim) { const double difference = static_cast(a[dim]) - b[dim]; @@ -85,7 +85,7 @@ inline void exact_assign( size_t d, size_t k, bool spherical, - uint32_t num_threads, + uint32_t num_workers, PID* assignments, float* distances, const PID* previous_assignments = nullptr @@ -95,15 +95,15 @@ inline void exact_assign( // smaller blocks on other backends and require at least two blocks per // worker to avoid doubling the work of a worker handling a small tail. const size_t point_block = - cpu::has_avx512_core() && k >= 256 && n / num_threads >= 2 * kMaxPointBlock + cpu::has_avx512_core() && k >= 256 && n / num_workers >= 2 * kMaxPointBlock ? kMaxPointBlock : 64; constexpr size_t kCentroidBlock = 256; const size_t blocks = 1 + (n - 1) / point_block; - const auto num_workers = static_cast(std::min(num_threads, blocks)); + const auto team_size = static_cast(std::min(num_workers, blocks)); std::vector centroid_norms(k); std::vector> products( - num_workers, std::vector(point_block * std::min(k, kCentroidBlock)) + team_size, std::vector(point_block * std::min(k, kCentroidBlock)) ); index_detail::ParallelException failure; @@ -115,7 +115,7 @@ inline void exact_assign( } distance += 1.0; } else { - distance = squared_l2_f64(vector, centroid, d); + distance = euclidean_sqr_f64(vector, centroid, d); } return distance; }; @@ -128,7 +128,7 @@ inline void exact_assign( const double error_floor = 8.0 * static_cast(d) * std::numeric_limits::min() + 8.0 * std::numeric_limits::epsilon(); -#pragma omp parallel num_threads(num_workers) +#pragma omp parallel num_threads(team_size) { // Set only this OpenMP task's ICV, not Eigen's process-global thread limit. // This also prevents nested GEMM teams when the outer team has one worker. @@ -143,7 +143,7 @@ inline void exact_assign( for (std::ptrdiff_t cluster = 0; cluster < static_cast(k); ++cluster) { centroid_norms[static_cast(cluster)] = - squared_norm_f64(centroids + static_cast(cluster) * d, d); + l2norm_sqr_f64(centroids + static_cast(cluster) * d, d); } #pragma omp for schedule(static) for (std::ptrdiff_t block = 0; block < static_cast(blocks); @@ -156,7 +156,7 @@ inline void exact_assign( const size_t rows = std::min(point_block, n - first_point); for (size_t row = 0; row < rows; ++row) { const float* vector = x + (first_point + row) * d; - point_norms[row] = squared_norm_f64(vector, d); + point_norms[row] = l2norm_sqr_f64(vector, d); best_distances[row] = std::numeric_limits::infinity(); if (previous_assignments != nullptr) { const size_t point = first_point + row; @@ -247,7 +247,7 @@ class LloydKMeans : public Parameters { std::string(Assignment::kName) + " input dimensions overflow ptrdiff_t" ); } - const uint32_t num_workers = effective_num_threads(); + const uint32_t num_workers = resolved_num_workers(); // Leave headroom for products of squared distances in graph pruning, // as well as quantized reconstruction and rotated residuals. const double max_coordinate = std::sqrt( @@ -258,9 +258,9 @@ class LloydKMeans : public Parameters { bool finite = true; #pragma omp parallel for num_threads(num_workers) reduction(&& : finite, in_range) \ schedule(static) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n * d); - ++loop_index) { - const size_t i = static_cast(loop_index); + for (std::ptrdiff_t index = 0; index < static_cast(n * d); + ++index) { + const size_t i = static_cast(index); finite = finite && std::isfinite(x[i]); in_range = in_range && std::abs(static_cast(x[i])) <= max_coordinate; } @@ -379,7 +379,7 @@ class LloydKMeans : public Parameters { std::vector assignments; std::vector distances; std::vector iteration_stats; - /// Objective of the final returned assignments against the final centroids. + // Objective of the final returned assignments against the final centroids. double final_obj = 0.0; private: @@ -449,13 +449,13 @@ class LloydKMeans : public Parameters { for (size_t point = 0; point < n; ++point) { partitions[worker_for_cluster[assignments[point]]].points.push_back(point); } - const uint32_t num_workers = effective_num_threads(); + const uint32_t num_workers = resolved_num_workers(); static_cast(num_workers); #pragma omp parallel for num_threads(num_workers) schedule(static) - for (std::ptrdiff_t loop_worker = 0; - loop_worker < static_cast(partitions.size()); - ++loop_worker) { - const auto& partition = partitions[static_cast(loop_worker)]; + for (std::ptrdiff_t index = 0; + index < static_cast(partitions.size()); + ++index) { + const auto& partition = partitions[static_cast(index)]; // The same interval defines point ownership and accumulator clearing, // including empty clusters, independently of the physical thread team. const size_t first_cluster = partition.first_cluster; @@ -487,19 +487,18 @@ class LloydKMeans : public Parameters { std::vector& cluster_sizes, const std::vector& previous_centroids ) { - const uint32_t num_workers = effective_num_threads(); + const uint32_t num_workers = resolved_num_workers(); // Means, normalization, and displacement share one worker team. Preserve // the per-centroid arithmetic; empty rows are measured after refilling. const auto cluster_shift = [&](size_t cluster) { - return squared_l2_f64( + return euclidean_sqr_f64( centroids.data() + cluster * d, previous_centroids.data() + cluster * d, d ); }; double shift = 0; #pragma omp parallel for num_threads(num_workers) reduction(+ : shift) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); - ++loop_index) { - const size_t cluster = static_cast(loop_index); + for (std::ptrdiff_t index = 0; index < static_cast(k); ++index) { + const size_t cluster = static_cast(index); const double scale = cluster_sizes[cluster] == 0 ? 0.0 : 1.0 / static_cast(cluster_sizes[cluster]); @@ -549,10 +548,9 @@ class LloydKMeans : public Parameters { { size_t best = n; #pragma omp for schedule(static) - for (std::ptrdiff_t loop_index = 0; - loop_index < static_cast(n); - ++loop_index) { - const size_t point = static_cast(loop_index); + for (std::ptrdiff_t index = 0; index < static_cast(n); + ++index) { + const size_t point = static_cast(index); if (cluster_sizes[assignments[point]] <= 1) { continue; } @@ -581,9 +579,8 @@ class LloydKMeans : public Parameters { } #pragma omp parallel for num_threads(num_workers) schedule(static) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); - ++loop_index) { - const size_t point = static_cast(loop_index); + for (std::ptrdiff_t index = 0; index < static_cast(n); ++index) { + const size_t point = static_cast(index); distances[point] = distance_func(x + point * d, centroids.data() + assignments[point] * d, d); } @@ -606,28 +603,27 @@ class LloydKMeans : public Parameters { } void normalize_centroid(float* centroid) const { - const double squared_norm = squared_norm_f64(centroid, d); - if (squared_norm == 0) { + const double norm_sqr = l2norm_sqr_f64(centroid, d); + if (norm_sqr == 0) { return; } - const double scale = 1.0 / std::sqrt(squared_norm); + const double scale = 1.0 / std::sqrt(norm_sqr); for (size_t dim = 0; dim < d; ++dim) { centroid[dim] = static_cast(centroid[dim] * scale); } } void normalize_centroids() { - const uint32_t num_workers = effective_num_threads(); + const uint32_t num_workers = resolved_num_workers(); static_cast(num_workers); #pragma omp parallel for num_threads(num_workers) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); - ++loop_index) { - const size_t cluster = static_cast(loop_index); + for (std::ptrdiff_t index = 0; index < static_cast(k); ++index) { + const size_t cluster = static_cast(index); normalize_centroid(centroids.data() + cluster * d); } } - [[nodiscard]] uint32_t effective_num_threads() const { + [[nodiscard]] uint32_t resolved_num_workers() const { return static_cast(resolve_num_threads(num_threads)); } }; diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index bc6a97a9..6115a498 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -32,22 +32,20 @@ using QGKMeansIterationStats = clustering::KMeansIterationStats; namespace detail { -using clustering::detail::exact_assign; -using clustering::detail::should_terminate_by_objective; class QGAssignment; } // namespace detail struct QGKMeansParameters { - /// Number of clustering iterations. + // Number of clustering iterations. size_t niter = 25; bool verbose = false; - /// Normalize centroids after each update. + // Normalize centroids after each update. bool spherical = false; uint32_t seed = 42; - /// Warn below this number of training vectors per centroid. + // Warn below this number of training vectors per centroid. size_t min_points_per_centroid = 39; - /// Stop when the absolute relative objective (WCSS) change is at most this value. + // Stop when the absolute relative objective (WCSS) change is at most this value. double early_stop_threshold = 0.0; FinalAssignmentMode final_assignment = FinalAssignmentMode::Approximate; @@ -55,10 +53,10 @@ struct QGKMeansParameters { uint32_t graph_degree = 32; uint32_t ef_build = 240; uint32_t ef_search = 16; - /// Graph-building passes after PiPNN initialization; the last refines the graph. + // Graph-building passes after PiPNN initialization; the last refines the graph. uint32_t graph_build_iterations = 1; uint32_t num_threads = 0; - /// Centroid storage in QG: 0 for raw float32, or 4/8-bit RaBitQ. + // Centroid storage in QG: 0 for raw float32, or 4/8-bit RaBitQ. uint32_t quantization_bits = 0; }; @@ -132,10 +130,10 @@ class QGAssigner { QGAssigner( size_t d, size_t k, - uint32_t M, - uint32_t ef_construction, + uint32_t graph_degree, + uint32_t ef_build, uint32_t ef_search, - uint32_t niter, + uint32_t graph_build_iterations, uint32_t num_threads = 0, uint32_t seed = 42, MetricType metric_type = METRIC_L2, @@ -143,19 +141,18 @@ class QGAssigner { ) : dim_(d) , num_clusters_(k) - , max_degree_(M) - , ef_construction_(ef_construction) + , graph_degree_(graph_degree) + , ef_build_(ef_build) , ef_search_(ef_search) - , niter_(niter) - , num_workers_(num_threads) + , graph_build_iterations_(graph_build_iterations) , seed_(seed) { - validate_configuration(); - num_workers_ = static_cast(resolve_num_threads(num_workers_)); + validate_configuration(num_threads); + num_workers_ = static_cast(resolve_num_threads(num_threads)); graph_ = std::make_unique>( num_clusters_, dim_, - max_degree_, + graph_degree_, metric_type, RotatorType::FhtKacRotator, quantization_bits, @@ -167,7 +164,7 @@ class QGAssigner { for (uint32_t thread = 0; thread < num_workers_; ++thread) { scratch_.emplace_back( std::make_unique( - graph_->padded_dim(), max_degree_, ef_search_, 1 + graph_->padded_dim(), graph_degree_, ef_search_, 1 ) ); } @@ -179,16 +176,16 @@ class QGAssigner { const float* centroids, const float* x, size_t n, - PID* labels, + PID* assignments, float* distances, float* query_sums = nullptr, bool sums_ready = false ) { detail::QuerySumCache cache(query_sums, sums_ready); - assign_with_cache(centroids, x, n, labels, distances, cache); + assign_with_cache(centroids, x, n, assignments, distances, cache); } - void clear_hints() { previous_labels_.clear(); } + void clear_hints() { previous_assignments_.clear(); } private: void rebuild_centroid_graph(const float* centroids) { @@ -204,7 +201,7 @@ class QGAssigner { ); builder_ = std::make_unique( *graph_, - ef_construction_, + ef_build_, centroids, build_threads, symqg::QGInitialization::PiPNN, @@ -212,10 +209,10 @@ class QGAssigner { true ); } - if (niter_ == 1) { + if (graph_build_iterations_ == 1) { builder_->build(); } else { - builder_->build(niter_); + builder_->build(graph_build_iterations_); } } @@ -223,11 +220,11 @@ class QGAssigner { const float* centroids, const float* x, size_t n, - PID* labels, + PID* assignments, float* distances, detail::QuerySumCache& query_sums ) { - if (centroids == nullptr || x == nullptr || labels == nullptr || + if (centroids == nullptr || x == nullptr || assignments == nullptr || distances == nullptr) { throw std::invalid_argument( "QGAssigner input and output pointers must not be null" @@ -245,25 +242,24 @@ class QGAssigner { // Stable grouping changes only the order of independent full searches. std::vector query_order; - if (previous_labels_.size() == n && n <= std::numeric_limits::max()) { + if (previous_assignments_.size() == n && + n <= std::numeric_limits::max()) { query_order.resize(n); std::vector offsets(num_clusters_ + 1, 0); - for (PID label : previous_labels_) { - ++offsets[label + 1]; + for (PID cluster_id : previous_assignments_) { + ++offsets[cluster_id + 1]; } std::partial_sum(offsets.begin(), offsets.end(), offsets.begin()); for (size_t point = 0; point < n; ++point) { - query_order[offsets[previous_labels_[point]]++] = + query_order[offsets[previous_assignments_[point]]++] = static_cast(point); } } index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) schedule(dynamic, 64) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); - ++loop_index) { - const size_t i = query_order.empty() - ? static_cast(loop_index) - : query_order[static_cast(loop_index)]; + for (std::ptrdiff_t index = 0; index < static_cast(n); ++index) { + const size_t i = query_order.empty() ? static_cast(index) + : query_order[static_cast(index)]; const int thread_id = omp_get_thread_num(); if (failure.failed()) { continue; @@ -283,12 +279,12 @@ class QGAssigner { visited_capacity = capacity; } const PID hint = - previous_labels_.size() == n ? previous_labels_[i] : kPidMax; + previous_assignments_.size() == n ? previous_assignments_[i] : kPidMax; symqg::detail::search_with_workspace( *graph_, x + i * dim_, 1, - labels + i, + assignments + i, distances + i, scratch.view(visited), hint, @@ -298,20 +294,21 @@ class QGAssigner { if (graph_->is_quantized()) { // Keep objectives and returned distances relative to the original // centroids; graph search scores are quantized estimates. - distances[i] = - distance_func(x + i * dim_, centroids + labels[i] * dim_, dim_); + distances[i] = distance_func( + x + i * dim_, centroids + assignments[i] * dim_, dim_ + ); } } catch (...) { failure.capture(); } } failure.rethrow(); - previous_labels_.assign(labels, labels + n); + previous_assignments_.assign(assignments, assignments + n); query_sums.complete(); } - void validate_configuration() const { - if (num_workers_ > static_cast(std::numeric_limits::max())) { + void validate_configuration(uint32_t num_threads) const { + if (num_threads > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("QGAssigner num_threads must fit in an OpenMP int"); } if (dim_ < 64 || dim_ > rotator_impl::FhtKacRotator::kMaxDim) { @@ -319,36 +316,36 @@ class QGAssigner { "QGKMeans currently requires dimensions in the range [64, 65536]" ); } - if (max_degree_ == 0 || max_degree_ % fastscan::kBatchSize != 0) { + if (graph_degree_ == 0 || graph_degree_ % fastscan::kBatchSize != 0) { throw std::invalid_argument( "QGKMeans graph degree must be a positive multiple of 32" ); } - if (num_clusters_ <= max_degree_) { + if (num_clusters_ <= graph_degree_) { throw std::invalid_argument( "QGKMeans requires more centroids than the graph degree" ); } - if (ef_construction_ == 0 || ef_search_ == 0) { + if (ef_build_ == 0 || ef_search_ == 0) { throw std::invalid_argument("QGKMeans ef_build and ef_search must be positive"); } - if (niter_ == 0) { + if (graph_build_iterations_ == 0) { throw std::invalid_argument("QGKMeans graph_build_iterations must be positive"); } } size_t dim_; size_t num_clusters_; - uint32_t max_degree_; - uint32_t ef_construction_; + uint32_t graph_degree_; + uint32_t ef_build_; uint32_t ef_search_; - uint32_t niter_; - uint32_t num_workers_; + uint32_t graph_build_iterations_; + uint32_t num_workers_ = 0; uint32_t seed_; std::unique_ptr> graph_; std::unique_ptr builder_; std::vector> scratch_; - std::vector previous_labels_; + std::vector previous_assignments_; }; namespace detail { @@ -379,13 +376,13 @@ class QGAssignment { size_t k, const float* x, size_t n, - uint32_t num_threads, + uint32_t num_workers, const QGKMeansParameters& cp ) : dim_(d) , data_(x) , num_points_(n) - , num_workers_(num_threads) + , num_workers_(num_workers) , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) , quantized_(cp.quantization_bits != 0) , assigner_( @@ -402,16 +399,18 @@ class QGAssignment { ) , query_sums_(assigner_.graph_->padded_dim(), cp) {} - void assign(const float* centroids, PID* labels, float* distances, bool final = false) { + void assign( + const float* centroids, PID* assignments, float* distances, bool final = false + ) { query_sums_.prepare(num_points_); - std::vector previous_labels; + std::vector previous_assignments; if (final && quantized_) { - previous_labels.assign(labels, labels + num_points_); + previous_assignments.assign(assignments, assignments + num_points_); } assigner_.assign_with_cache( - centroids, data_, num_points_, labels, distances, query_sums_ + centroids, data_, num_points_, assignments, distances, query_sums_ ); - if (previous_labels.empty()) { + if (previous_assignments.empty()) { return; } @@ -420,18 +419,17 @@ class QGAssignment { const auto distance_func = metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; #pragma omp parallel for num_threads(num_workers_) schedule(static) - for (std::ptrdiff_t loop_index = 0; - loop_index < static_cast(num_points_); - ++loop_index) { - const size_t point = static_cast(loop_index); - const PID previous = previous_labels[point]; - if (previous == labels[point]) { + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); + ++index) { + const size_t point = static_cast(index); + const PID previous = previous_assignments[point]; + if (previous == assignments[point]) { continue; } const float distance = distance_func(data_ + point * dim_, centroids + previous * dim_, dim_); if (distance < distances[point]) { - labels[point] = previous; + assignments[point] = previous; distances[point] = distance; } } diff --git a/include/rabitqlib/clustering/rabitqkmeans.hpp b/include/rabitqlib/clustering/rabitqkmeans.hpp index 40bdcb49..efdb0a52 100644 --- a/include/rabitqlib/clustering/rabitqkmeans.hpp +++ b/include/rabitqlib/clustering/rabitqkmeans.hpp @@ -28,15 +28,15 @@ using clustering::FinalAssignmentMode; using RaBitQKMeansIterationStats = clustering::KMeansIterationStats; struct RaBitQKMeansParameters { - /// Number of clustering iterations. + // Number of clustering iterations. size_t niter = 25; bool verbose = false; - /// Normalize centroids after each update. + // Normalize centroids after each update. bool spherical = false; uint32_t seed = 42; - /// Warn below this number of training vectors per centroid. + // Warn below this number of training vectors per centroid. size_t min_points_per_centroid = 39; - /// Stop when the absolute relative objective (WCSS) change is at most this value. + // Stop when the absolute relative objective (WCSS) change is at most this value. double early_stop_threshold = 0.0; FinalAssignmentMode final_assignment = FinalAssignmentMode::Approximate; uint32_t num_threads = 0; @@ -46,7 +46,7 @@ namespace detail { // Encode immutable training points once; each moving centroid supplies one LUT. // RaBitQKMeans validates inputs and keeps the borrowed x buffer alive through training. -class RaBitQAssigner { +class RaBitQAssignment { size_t dim_; size_t num_clusters_; const float* data_; @@ -73,12 +73,12 @@ class RaBitQAssigner { } } - RaBitQAssigner( + RaBitQAssignment( size_t d, size_t k, const float* x, size_t n, - uint32_t num_threads, + uint32_t num_workers, const RaBitQKMeansParameters& cp ) : dim_(d) @@ -86,7 +86,7 @@ class RaBitQAssigner { , data_(x) , num_points_(n) , num_batches_(1 + (n - 1) / fastscan::kBatchSize) - , num_workers_(static_cast(std::min(num_threads, num_batches_))) + , num_workers_(static_cast(std::min(num_workers, num_batches_))) , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) , padded_dim_(rotator_impl::padding_requirement(d, RotatorType::FhtKacRotator)) , rotator_(d, padded_dim_, cp.seed) @@ -128,15 +128,15 @@ class RaBitQAssigner { ); index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) schedule(static) - for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); - ++batch) { + for (std::ptrdiff_t index = 0; index < static_cast(num_batches_); + ++index) { const auto thread = static_cast(omp_get_thread_num()); if (failure.failed()) { continue; } try { - const auto block = static_cast(batch); - const size_t first = block * fastscan::kBatchSize; + const auto batch = static_cast(index); + const size_t first = batch * fastscan::kBatchSize; const size_t count = std::min(fastscan::kBatchSize, num_points_ - first); auto& rotated = rotated_blocks[thread]; for (size_t point = 0; point < count; ++point) { @@ -150,7 +150,7 @@ class RaBitQAssigner { count, padded_dim_, batch_data_.data() + - block * QGBatchDataMap::data_bytes(padded_dim_), + batch * QGBatchDataMap::data_bytes(padded_dim_), metric_type_ ); } catch (...) { failure.capture(); } @@ -158,7 +158,7 @@ class RaBitQAssigner { failure.rethrow(); } - void assign(const float* centroids, PID* labels, float* distances, bool = false) { + void assign(const float* centroids, PID* assignments, float* distances, bool = false) { const auto distance_func = metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; for (size_t cluster = 0; cluster < num_clusters_; ++cluster) { @@ -173,21 +173,21 @@ class RaBitQAssigner { index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) schedule(static) - for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); - ++batch) { + for (std::ptrdiff_t index = 0; index < static_cast(num_batches_); + ++index) { if (failure.failed()) { continue; } try { - const auto block = static_cast(batch); - const size_t first = block * fastscan::kBatchSize; + const auto batch = static_cast(index); + const size_t first = batch * fastscan::kBatchSize; const size_t count = std::min(fastscan::kBatchSize, num_points_ - first); const char* codes = batch_data_.data() + - block * QGBatchDataMap::data_bytes(padded_dim_); + batch * QGBatchDataMap::data_bytes(padded_dim_); std::array estimated_distances{}; std::array best_distances; best_distances.fill(std::numeric_limits::infinity()); - std::array best_labels{}; + std::array best_cluster_ids{}; for (size_t cluster = 0; cluster < num_clusters_; ++cluster) { qg_batch_estdist( codes, @@ -200,17 +200,17 @@ class RaBitQAssigner { // ties. if (estimated_distances[point] < best_distances[point]) { best_distances[point] = estimated_distances[point]; - best_labels[point] = static_cast(cluster); + best_cluster_ids[point] = static_cast(cluster); } } } for (size_t point = 0; point < count; ++point) { const size_t id = first + point; - labels[id] = best_labels[point]; + assignments[id] = best_cluster_ids[point]; // Scores choose the centroid; objectives retain original-space // distances. distances[id] = distance_func( - data_ + id * dim_, centroids + best_labels[point] * dim_, dim_ + data_ + id * dim_, centroids + best_cluster_ids[point] * dim_, dim_ ); } } catch (...) { failure.capture(); } @@ -223,9 +223,9 @@ class RaBitQAssigner { /** K-means with a graph-free RaBitQ scan, recommended for small cluster counts. */ class RaBitQKMeans : public clustering::detail:: - LloydKMeans { + LloydKMeans { using Base = - clustering::detail::LloydKMeans; + clustering::detail::LloydKMeans; public: RaBitQKMeans(size_t d, size_t k) : RaBitQKMeans(d, k, RaBitQKMeansParameters{}) {} diff --git a/include/rabitqlib/core/api.hpp b/include/rabitqlib/core/api.hpp index a5f551e0..9262a218 100644 --- a/include/rabitqlib/core/api.hpp +++ b/include/rabitqlib/core/api.hpp @@ -12,7 +12,7 @@ namespace rabitqlib::quant::v1 { // Stable C++ source API. This header exposes no index, Eigen, or SIMD implementation. // See docs/docs/rabitq/core_api.md for the complete ownership/numerical contract. -enum class CodeLayout : std::uint8_t { UnpackedU16 = 0, SplitPacked = 1 }; +enum class CodeLayout : uint8_t { UnpackedU16 = 0, SplitPacked = 1 }; struct Parameters { // Working dimension >= 2; SplitPacked additionally requires dim % 32 == 0. @@ -84,7 +84,9 @@ QueryState prepare_query( ); // Full float-query estimation, including at one bit; no quantized-query filter is implied. -inline float compute(const QueryState& query, CodeView code, const Factors& factors); +inline float compute( + const QueryState& query_wrapper, CodeView code, const Factors& factors +); class QueryState { public: @@ -103,17 +105,19 @@ class QueryState { float (*extra_ip_)(const float*, const uint8_t*, size_t) = nullptr; static float compute_checked( - const QueryState& query, const CodeView& code, const Factors& factors + const QueryState& query_wrapper, const CodeView& code, const Factors& factors ); friend QueryState prepare_query(const Parameters&, const float*, const float*); friend float compute(const QueryState&, CodeView, const Factors&); }; -inline float compute(const QueryState& query, CodeView code, const Factors& factors) { +inline float compute( + const QueryState& query_wrapper, CodeView code, const Factors& factors +) { // Keep the public value semantics; the compiled implementation borrows this // local view instead of passing the entire aggregate through its call ABI. - return QueryState::compute_checked(query, code, factors); + return QueryState::compute_checked(query_wrapper, code, factors); } } // namespace rabitqlib::quant::v1 diff --git a/include/rabitqlib/core/data_layout.hpp b/include/rabitqlib/core/data_layout.hpp index 8a1afa30..ab5498bc 100644 --- a/include/rabitqlib/core/data_layout.hpp +++ b/include/rabitqlib/core/data_layout.hpp @@ -13,12 +13,14 @@ namespace detail { // The one-bit layout is the base-code layout with bits == 1. template struct SingleCodeOffsets { - static size_t add(size_t padded_dim, size_t bits = 1) { return padded_dim * bits / 8; } - static size_t rescale(size_t padded_dim, size_t bits = 1) { - return add(padded_dim, bits) + sizeof(T); + static size_t add_offset(size_t padded_dim, size_t bits = 1) { + return padded_dim * bits / 8; } - static size_t error(size_t padded_dim, size_t bits = 1) { - return rescale(padded_dim, bits) + sizeof(T); + static size_t rescale_offset(size_t padded_dim, size_t bits = 1) { + return add_offset(padded_dim, bits) + sizeof(T); + } + static size_t error_offset(size_t padded_dim, size_t bits = 1) { + return rescale_offset(padded_dim, bits) + sizeof(T); } }; } // namespace detail @@ -28,9 +30,10 @@ struct ExDataMap { public: explicit ExDataMap(char* data, size_t padded_dim, size_t ex_bits) : ex_code_(reinterpret_cast(data)) - , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) - , f_rescale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) { - } + , f_add_ex_(data + detail::SingleCodeOffsets::add_offset(padded_dim, ex_bits)) + , f_rescale_ex_( + data + detail::SingleCodeOffsets::rescale_offset(padded_dim, ex_bits) + ) {} static size_t data_bytes(size_t padded_dim, size_t ex_bits) { if (ex_bits == 0) { @@ -40,7 +43,8 @@ struct ExDataMap { if (padded_dim > max_size / ex_bits) { throw std::length_error("ExDataMap storage size exceeds size_t"); } - const size_t code_bytes = detail::SingleCodeOffsets::add(padded_dim, ex_bits); + const size_t code_bytes = + detail::SingleCodeOffsets::add_offset(padded_dim, ex_bits); const size_t factor_bytes = sizeof(T) * 2; if (code_bytes > max_size - factor_bytes) { throw std::length_error("ExDataMap storage size exceeds size_t"); @@ -63,9 +67,10 @@ struct ConstExDataMap { public: explicit ConstExDataMap(const char* data, size_t padded_dim, size_t ex_bits) : ex_code_(reinterpret_cast(data)) - , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) - , f_rescale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) { - } + , f_add_ex_(data + detail::SingleCodeOffsets::add_offset(padded_dim, ex_bits)) + , f_rescale_ex_( + data + detail::SingleCodeOffsets::rescale_offset(padded_dim, ex_bits) + ) {} [[nodiscard]] const uint8_t* ex_code() const { return ex_code_; } [[nodiscard]] T f_add_ex() const { return detail::load_packed_value(f_add_ex_); } @@ -84,12 +89,17 @@ struct BaseDataMap { public: explicit BaseDataMap(char* data, size_t padded_dim, size_t base_bits) : base_code_(reinterpret_cast(data)) - , f_add_(data + detail::SingleCodeOffsets::add(padded_dim, base_bits)) - , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim, base_bits)) - , f_error_(data + detail::SingleCodeOffsets::error(padded_dim, base_bits)) {} + , f_add_(data + detail::SingleCodeOffsets::add_offset(padded_dim, base_bits)) + , f_rescale_( + data + detail::SingleCodeOffsets::rescale_offset(padded_dim, base_bits) + ) + , f_error_( + data + detail::SingleCodeOffsets::error_offset(padded_dim, base_bits) + ) {} static size_t data_bytes(size_t padded_dim, size_t base_bits) { - return detail::SingleCodeOffsets::error(padded_dim, base_bits) + sizeof(T); + return detail::SingleCodeOffsets::error_offset(padded_dim, base_bits) + + sizeof(T); } [[nodiscard]] uint8_t* base_code() { return base_code_; } @@ -109,9 +119,13 @@ struct ConstBaseDataMap { public: explicit ConstBaseDataMap(const char* data, size_t padded_dim, size_t base_bits) : base_code_(reinterpret_cast(data)) - , f_add_(data + detail::SingleCodeOffsets::add(padded_dim, base_bits)) - , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim, base_bits)) - , f_error_(data + detail::SingleCodeOffsets::error(padded_dim, base_bits)) {} + , f_add_(data + detail::SingleCodeOffsets::add_offset(padded_dim, base_bits)) + , f_rescale_( + data + detail::SingleCodeOffsets::rescale_offset(padded_dim, base_bits) + ) + , f_error_( + data + detail::SingleCodeOffsets::error_offset(padded_dim, base_bits) + ) {} [[nodiscard]] const uint8_t* base_code() const { return base_code_; } [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } @@ -130,9 +144,9 @@ struct BinDataMap { public: explicit BinDataMap(char* data, size_t padded_dim) : bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::SingleCodeOffsets::add(padded_dim)) - , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim)) - , f_error_(data + detail::SingleCodeOffsets::error(padded_dim)) {} + , f_add_(data + detail::SingleCodeOffsets::add_offset(padded_dim)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale_offset(padded_dim)) + , f_error_(data + detail::SingleCodeOffsets::error_offset(padded_dim)) {} [[nodiscard]] uint8_t* bin_code() { return bin_code_; } [[nodiscard]] detail::PackedValueRef f_add() { return f_add_; } @@ -140,7 +154,7 @@ struct BinDataMap { [[nodiscard]] detail::PackedValueRef f_error() { return f_error_; } static size_t data_bytes(size_t padded_dim) { - return detail::SingleCodeOffsets::error(padded_dim) + sizeof(T); + return detail::SingleCodeOffsets::error_offset(padded_dim) + sizeof(T); } private: @@ -155,9 +169,9 @@ struct ConstBinDataMap { public: explicit ConstBinDataMap(const char* data, size_t padded_dim) : bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::SingleCodeOffsets::add(padded_dim)) - , f_rescale_(data + detail::SingleCodeOffsets::rescale(padded_dim)) - , f_error_(data + detail::SingleCodeOffsets::error(padded_dim)) {} + , f_add_(data + detail::SingleCodeOffsets::add_offset(padded_dim)) + , f_rescale_(data + detail::SingleCodeOffsets::rescale_offset(padded_dim)) + , f_error_(data + detail::SingleCodeOffsets::error_offset(padded_dim)) {} [[nodiscard]] const uint8_t* bin_code() const { return bin_code_; } [[nodiscard]] T f_add() const { return detail::load_packed_value(f_add_); } diff --git a/include/rabitqlib/core/detail/quantized_level.hpp b/include/rabitqlib/core/detail/quantized_level.hpp index 8ed4615a..46923dc6 100644 --- a/include/rabitqlib/core/detail/quantized_level.hpp +++ b/include/rabitqlib/core/detail/quantized_level.hpp @@ -8,15 +8,17 @@ namespace rabitqlib::quant::detail { // Keep this helper local to each translation unit: callers use different ISA flags. // Search evaluation and final encoding must use the same division-based thresholds. static inline int quantized_level_at_scale(double magnitude, double scale, int max_code) { - if (magnitude == 0) + if (magnitude == 0) { return 0; + } int code = static_cast(std::min(scale * magnitude, static_cast(max_code))); // Multiplication can round to either side of an event; correct against the // exact thresholds rather than adding an epsilon that could cross other events. - if (code < max_code && (code + 1.0) / magnitude <= scale) + if (code < max_code && (code + 1.0) / magnitude <= scale) { ++code; - else if (code > 0 && static_cast(code) / magnitude > scale) + } else if (code > 0 && static_cast(code) / magnitude > scale) { --code; + } return code; } diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index bb1bcd0b..80e40107 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -29,7 +29,7 @@ namespace rabitqlib::quant::rabitq_impl { -constexpr float kConstEpsilon = 1.9; +inline constexpr float kConstEpsilon = 1.9; // Arithmetic only: callers own residual/centered storage and choose rescale arithmetic. template @@ -100,11 +100,9 @@ inline T direct_rescale(T l2_sqr, T residual_code_ip, MetricType metric_type) { : -l2_sqr / residual_code_ip; } -/** - * Inverse of pack_binary_to_bytes: write each coordinate's sign bit (0 or - * 1) to signs. Full 64-coordinate words come first, followed by a final partial word - * of any multiple of 8 coordinates. - */ +// Inverse of pack_binary_to_bytes: write each coordinate's sign bit (0 or +// 1) to signs. Full 64-coordinate words come first, followed by a final partial word +// of any multiple of 8 coordinates. inline void unpack_sign_bits(const uint8_t* packed, size_t dim, uint16_t* signs) { for (size_t start = 0; start < dim; start += 64) { const size_t width = std::min(size_t{64}, dim - start); @@ -151,12 +149,12 @@ inline void one_bit_code_with_factor_into( } else { // Wider vectors benefit from Eigen's packet width rather than the // compiler's narrower vectorization of the fused multi-output loop. - RowMajorArrayMap residual_arr(residual, 1, static_cast(dim)); - RowMajorArrayMap xu_cb(centered, 1, static_cast(dim)); - ConstRowMajorArrayMap data_arr(data, 1, dim); - ConstRowMajorArrayMap cent_arr(centroid, 1, dim); + RowMajorArrayMap residual_arr(residual, 1, static_cast(dim)); + RowMajorArrayMap xu_cb(centered, 1, static_cast(dim)); + ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); + ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); residual_arr = data_arr - cent_arr; - RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); + RowMajorArrayMap x_u(binary_code, 1, static_cast(dim)); x_u = (residual_arr > 0).template cast(); xu_cb = x_u.template cast() + T{-0.5}; } @@ -192,12 +190,12 @@ inline void one_bit_code_with_factor( thread_local RowMajorArray centered; // Eigen's resize frees the old allocation before allocating its replacement. // Allocate first so failure leaves the thread-local owner safe to reuse/destroy. - if (residual.size() != static_cast(dim)) { - RowMajorArray replacement(1, static_cast(dim)); + if (residual.size() != static_cast(dim)) { + RowMajorArray replacement(1, static_cast(dim)); residual.swap(replacement); } - if (centered.size() != static_cast(dim)) { - RowMajorArray replacement(1, static_cast(dim)); + if (centered.size() != static_cast(dim)) { + RowMajorArray replacement(1, static_cast(dim)); centered.swap(replacement); } one_bit_code_with_factor_into( @@ -222,7 +220,7 @@ inline void one_bit_code_with_factor( * Factors have the same meaning as in one_bit_code_with_factor_into, including * the query-independent, probabilistic error factor f_error. */ -template +template inline void one_bit_compact_code_to_bytes( const T* data, const T* centroid, @@ -246,21 +244,21 @@ inline void one_bit_compact_code_to_bytes( metric_type ); - pack_binary_to_bytes(binary_code.data(), compact_code, padded_dim); + pack_binary_to_bytes(binary_code.data(), compact_code, padded_dim); } -template +template inline void one_bit_compact_code( const T* data, const T* centroid, size_t padded_dim, - TC* compact_code, + TP* compact_code, T& f_add, T& f_rescale, T& f_error, MetricType metric_type = METRIC_L2 ) { - one_bit_compact_code_to_bytes( + one_bit_compact_code_to_bytes( data, centroid, padded_dim, @@ -275,7 +273,7 @@ inline void one_bit_compact_code( // Requires a positive padded_dim divisible by 8. template < typename T, - typename TC, + typename TP, bool Parallel = false, typename FAdd, typename FRescale, @@ -283,9 +281,9 @@ template < inline void one_bit_compact_codes( const T* data, const T* centroid, - size_t num, + size_t num_points, size_t padded_dim, - TC* compact_code, + TP* compact_code, FAdd f_add, FRescale f_rescale, FError f_error, @@ -295,7 +293,7 @@ inline void one_bit_compact_codes( T add; T rescale; T error; - one_bit_compact_code_to_bytes( + one_bit_compact_code_to_bytes( data + (padded_dim * i), centroid, padded_dim, @@ -313,7 +311,8 @@ inline void one_bit_compact_codes( if constexpr (Parallel) { std::exception_ptr failure; #pragma omp parallel for - for (std::ptrdiff_t index = 0; index < static_cast(num); ++index) { + for (std::ptrdiff_t index = 0; index < static_cast(num_points); + ++index) { try { encode_one(static_cast(index)); } catch (...) { @@ -332,7 +331,7 @@ inline void one_bit_compact_codes( #endif { // No serialized OpenMP region: an outer caller can catch scratch failures. - for (size_t i = 0; i < num; ++i) { + for (size_t i = 0; i < num_points; ++i) { encode_one(i); } } @@ -342,6 +341,7 @@ inline void one_bit_compact_codes( namespace ex_bits { // Lower endpoint of the scale search as a fraction of t_end, indexed by ex_bits. +// Internal linkage lets the encoder address the table directly instead of through the GOT. constexpr std::array kTightStart = { 0, 0.15, @@ -382,8 +382,9 @@ inline double event_sweep( // Never increment a saturated coordinate; skip zero coordinates. if (magnitude > 0 && cur < max_code) { const double next = static_cast(cur + 1) / magnitude; - if (next < t_end) + if (next < t_end) { next_t.emplace_back(next, i); + } } } std::make_heap(next_t.begin(), next_t.end(), std::greater{}); @@ -404,8 +405,9 @@ inline double event_sweep( numerator += magnitude; Event next{t_end, i}; - if (cur_o_bar[i] < max_code) + if (cur_o_bar[i] < max_code) { next.first = static_cast(cur_o_bar[i] + 1) / magnitude; + } if (next.first >= t_end) { // No further event for this coordinate: fill the root's hole // with the last heap entry, then restore the heap below. @@ -419,10 +421,12 @@ inline double event_sweep( size_t parent = 0; size_t child = 1; while (child < next_t.size()) { - if (child + 1 < next_t.size() && next_t[child + 1] < next_t[child]) + if (child + 1 < next_t.size() && next_t[child + 1] < next_t[child]) { ++child; - if (!(next_t[child] < next)) + } + if (!(next_t[child] < next)) { break; + } next_t[parent] = next_t[child]; parent = child; child = 2 * parent + 1; @@ -448,11 +452,13 @@ inline double best_rescale_factor_with_scratch( if (ex_bits >= kTightStart.size()) { throw std::invalid_argument("ex_bits exceeds the supported search table"); } - if (dim == 0) + if (dim == 0) { return 0; + } const double max_o = static_cast(*std::max_element(o_abs, o_abs + dim)); - if (max_o == 0) + if (max_o == 0) { return 0; + } const int max_code = (1 << ex_bits) - 1; const double t_end = static_cast(max_code + 10) / max_o; @@ -478,8 +484,9 @@ inline double best_rescale_factor_with_scratch( scratch->magnitudes, scratch->reciprocals ); - if (t >= 0) + if (t >= 0) { return t; + } } // Scalar event sweep, also used if SIMD is unavailable or cannot certify a @@ -520,8 +527,9 @@ inline T quantize_ex_with_scratch( code[i] = static_cast(value); ipnorm += (value + 0.5) * magnitude; } - if (ipnorm == 0) + if (ipnorm == 0) { return static_cast(1); + } const T ipnorm_inv = static_cast(1.0 / ipnorm); return std::isnormal(ipnorm_inv) ? ipnorm_inv : static_cast(1); } @@ -534,15 +542,15 @@ inline T quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits) { } // For given dim and ex_bits, use random vectors to get the const rescale factor -inline double get_const_scaling_factors(size_t dim, size_t ex_bits) { - constexpr long kConstNum = 100; +inline double const_scaling_factor(size_t dim, size_t ex_bits) { + constexpr Eigen::Index kConstNum = 100; RowMajorArray rand = random_gaussian_matrix(kConstNum, dim, 42); rand = rand.rowwise().normalized().abs(); double sum = 0; - for (long j = 0; j < kConstNum; ++j) { + for (Eigen::Index j = 0; j < kConstNum; ++j) { sum += best_rescale_factor(&rand(j, 0), dim, ex_bits); } @@ -552,27 +560,22 @@ inline double get_const_scaling_factors(size_t dim, size_t ex_bits) { } template -T faster_quantize_ex(const T* o_abs, TP* code, size_t dim, size_t ex_bits, double t_const) { +inline T faster_quantize_ex( + const T* o_abs, TP* code, size_t dim, size_t ex_bits, double t_const +) { constexpr double kEps = 1e-5; + const int max_code = (1 << ex_bits) - 1; double ipnorm = 0; - - for (size_t i = 0; i < dim; i++) { - int tmp_code = static_cast((t_const * o_abs[i]) + kEps); - if (tmp_code >= (1 << ex_bits)) { - tmp_code = (1 << ex_bits) - 1; - } - code[i] = static_cast(tmp_code); - - // ip * norm = unnormalized ip - ipnorm += (tmp_code + 0.5) * o_abs[i]; - } - - T ipnorm_inv = static_cast(1 / ipnorm); // 1 / (ip*norm) - if (!std::isnormal(ipnorm_inv)) { - ipnorm_inv = 1.F; + for (size_t i = 0; i < dim; ++i) { + const double magnitude = static_cast(o_abs[i]); + const int value = + std::min(static_cast((t_const * magnitude) + kEps), max_code); + code[i] = static_cast(value); + // Unnormalized inner product: ip * norm. + ipnorm += (value + 0.5) * magnitude; } - - return ipnorm_inv; + const T ipnorm_inv = static_cast(1 / ipnorm); + return std::isnormal(ipnorm_inv) ? ipnorm_inv : static_cast(1); } template @@ -585,7 +588,7 @@ inline T ex_bits_code_with_norm( double t_const, Scratch* scratch ) { - ConstRowMajorArrayMap res_arr(residual, 1, dim); + ConstRowMajorArrayMap residual_arr(residual, 1, static_cast(dim)); T ipnorm_inv = 1; if (residual_norm == 0) { // Match code_with_factor_from_residual(): no magnitude bits, so skip the revert. @@ -596,7 +599,7 @@ inline T ex_bits_code_with_norm( auto abs_res = scratch_array( scratch == nullptr ? nullptr : scratch->absolute, dim, owned_absolute ); - abs_res = res_arr.abs() / residual_norm; + abs_res = residual_arr.abs() / residual_norm; if (t_const > 0) { ipnorm_inv = faster_quantize_ex(abs_res.data(), ex_code, dim, ex_bits, t_const); } else { @@ -608,7 +611,7 @@ inline T ex_bits_code_with_norm( // Revert codes for negative or zero dims int32_t mask = (1 << ex_bits) - 1; for (size_t j = 0; j < dim; ++j) { - if (res_arr.data()[j] <= 0) { + if (residual_arr.data()[j] <= 0) { TP tmp = ex_code[j]; ex_code[j] = (~tmp) & mask; } @@ -677,9 +680,11 @@ inline void code_with_factor_from_residual( RowMajorArray owned_total; auto total_code = scratch_array(scratch == nullptr ? nullptr : scratch->total, dim, owned_total); - total_code = RowMajorArrayMap(ex_code, 1, dim).template cast(); + total_code = RowMajorArrayMap(ex_code, 1, static_cast(dim)) + .template cast(); for (size_t i = 0; i < dim; ++i) { - total_code(0, static_cast(i)) += static_cast(residual[i] > 0) << ex_bits; + total_code(0, static_cast(i)) += static_cast(residual[i] > 0) + << ex_bits; } // Full-code factors use the metric-specific derivation in @@ -706,7 +711,7 @@ inline void code_with_factor_from_residual( : ipnorm_inv * -factors.residual_norm; f_error_ex = factors.error; if constexpr (FullCode) { - RowMajorArrayMap(ex_code, 1, static_cast(dim)) = + RowMajorArrayMap(ex_code, 1, static_cast(dim)) = total_code.template cast(); } } @@ -725,8 +730,8 @@ inline void ex_bits_code_with_factor( double t_const = -1 ) { RowMajorArray residual = - ConstRowMajorArrayMap(data, 1, static_cast(dim)) - - ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); + ConstRowMajorArrayMap(data, 1, static_cast(dim)) - + ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); code_with_factor_from_residual( residual.data(), l2norm_sqr(residual.data(), dim), @@ -783,8 +788,8 @@ inline void combined_code( ) { const size_t ex_bits = total_bits - 1; - ConstRowMajorArrayMap res_arr(residual, 1, static_cast(dim)); - RowMajorArrayMap code_arr(total_code, 1, static_cast(dim)); + ConstRowMajorArrayMap residual_arr(residual, 1, static_cast(dim)); + RowMajorArrayMap code_arr(total_code, 1, static_cast(dim)); if (ex_bits > 0) { ex_bits::ex_bits_code(residual, dim, ex_bits, total_code, t_const); @@ -792,9 +797,11 @@ inline void combined_code( code_arr.setZero(); } - code_arr += (res_arr > 0).template cast() * static_cast(1 << ex_bits); + code_arr += (residual_arr > 0).template cast() * static_cast(1 << ex_bits); } +// The per-ISA HNSW search sources include this header, so these entry points keep +// internal linkage (see "Implementation linkage" in docs/development.md). template static inline void rabitq_scalar_impl( const T* data, @@ -809,8 +816,8 @@ static inline void rabitq_scalar_impl( ) { size_t ex_bits = total_bits - 1; - ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); - ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); + ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); + ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); RowMajorArray residual_arr = data_arr - cent_arr; if (l2norm_sqr(residual_arr.data(), dim) == 0) { @@ -823,7 +830,7 @@ static inline void rabitq_scalar_impl( combined_code(residual_arr.data(), dim, total_bits, total_code, t_const); float cb = -(static_cast(1 << ex_bits) - 0.5F); - RowMajorArrayMap total_code_arr(total_code, 1, dim); + RowMajorArrayMap total_code_arr(total_code, 1, static_cast(dim)); RowMajorArray u_cb = total_code_arr.template cast() + cb; T norm_data = std::sqrt(l2norm_sqr(residual_arr.data(), dim)); @@ -862,8 +869,8 @@ static inline void rabitq_full_impl_with_scratch( auto residual = scratch_array( scratch == nullptr ? nullptr : scratch->residual, dim, owned_residual ); - residual = ConstRowMajorArrayMap(data, 1, static_cast(dim)) - - ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); + residual = ConstRowMajorArrayMap(data, 1, static_cast(dim)) - + ConstRowMajorArrayMap(centroid, 1, static_cast(dim)); ex_bits::code_with_factor_from_residual( residual.data(), l2norm_sqr(residual.data(), dim), @@ -908,6 +915,7 @@ static inline void rabitq_full_impl_with_scratch( total_code[i] = static_cast(binary_code[i]); } } + template static inline void rabitq_full_impl( const T* data, @@ -940,7 +948,7 @@ static inline void rabitq_full_impl( namespace xy_bits { -constexpr size_t kMaxCombinedBits = 9; +inline constexpr size_t kMaxCombinedBits = 9; static_assert( ex_bits::kTightStart.size() >= kMaxCombinedBits, @@ -985,7 +993,7 @@ inline void code_factors( T& f_error, MetricType metric_type = METRIC_L2 ) { - ConstRowMajorArrayMap code_arr(code, 1, static_cast(dim)); + ConstRowMajorArrayMap code_arr(code, 1, static_cast(dim)); RowMajorArray xu_cb = code_arr.template cast() + cb; const T l2_sqr = l2norm_sqr(residual, dim); @@ -1017,8 +1025,8 @@ inline void split_code_with_factor( ) { validate_bit_size(base_bits, ex_bits); - ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); - ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); + ConstRowMajorArrayMap data_arr(data, 1, static_cast(dim)); + ConstRowMajorArrayMap cent_arr(centroid, 1, static_cast(dim)); RowMajorArray residual_arr = data_arr - cent_arr; const T* residual = residual_arr.data(); @@ -1053,7 +1061,7 @@ inline void split_code_with_factor( centroid, dim, base_code_int.data(), - -(static_cast((1U << base_bits) - 1) / 2.F), + -(static_cast((1U << base_bits) - 1) / 2.0F), f_add_base, f_rescale_base, f_error_base, @@ -1078,7 +1086,7 @@ inline void split_code_with_factor( centroid, dim, total_code.data(), - -(static_cast((1U << total_bits) - 1) / 2.F), + -(static_cast((1U << total_bits) - 1) / 2.0F), f_add_ex, f_rescale_ex, f_error_ex, diff --git a/include/rabitqlib/core/detail/scratch.hpp b/include/rabitqlib/core/detail/scratch.hpp index 9832ded1..e6968b98 100644 --- a/include/rabitqlib/core/detail/scratch.hpp +++ b/include/rabitqlib/core/detail/scratch.hpp @@ -97,9 +97,9 @@ struct Scratch { template inline RowMajorArrayMap scratch_array(T* external, size_t dim, RowMajorArray& owned) { if (external == nullptr) { - owned.resize(1, static_cast(dim)); + owned.resize(1, static_cast(dim)); external = owned.data(); } - return RowMajorArrayMap(external, 1, static_cast(dim)); + return RowMajorArrayMap(external, 1, static_cast(dim)); } } // namespace rabitqlib::quant::rabitq_impl diff --git a/include/rabitqlib/core/detail/simd/matrix_dispatch.hpp b/include/rabitqlib/core/detail/simd/matrix_dispatch.hpp index 37a6b642..4c37cee7 100644 --- a/include/rabitqlib/core/detail/simd/matrix_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/matrix_dispatch.hpp @@ -125,41 +125,41 @@ void pairwise_distances_lower_avx512( // Accumulate input rows in point_ids order into double cluster sums and counts. // Inputs and outputs must not overlap. Concurrent calls require disjoint clusters. void accumulate_cluster_sums( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ); void accumulate_cluster_sums_generic( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ); void accumulate_cluster_sums_avx2( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ); void accumulate_cluster_sums_avx512( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ); diff --git a/include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp b/include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp index a73117c7..5f493eb5 100644 --- a/include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/pack_excode_dispatch.hpp @@ -6,15 +6,10 @@ namespace rabitqlib::simd { void packing_2bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_3bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_4bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_5bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_6bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); // Inverse of the packers above, for bits in [1, 8]. @@ -23,33 +18,24 @@ void unpacking_excode_generic( ); void packing_2bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); -void packing_2bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_3bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); -void packing_3bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_4bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); -void packing_4bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_5bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); -void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_6bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); -void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_7bit_excode_generic(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); + +void packing_2bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_3bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_4bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); +void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_7bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); void packing_2bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_3bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_4bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_5bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_6bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); - void packing_7bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim); } // namespace rabitqlib::simd diff --git a/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp b/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp index 02d17c08..c835b906 100644 --- a/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp @@ -19,9 +19,9 @@ void fht_rotate( const uint8_t* flip ); -void flip_sign_neon(const uint8_t* flip, float* data, size_t dim); -void kacs_walk_neon(float* data, size_t len); -void fht_rotate_neon( +void flip_sign_generic(const uint8_t* flip, float* data, size_t dim); +void kacs_walk_generic(float* data, size_t len); +void fht_rotate_generic( const float* data, float* rotated_data, size_t dim, @@ -31,13 +31,9 @@ void fht_rotate_neon( const uint8_t* flip ); -void flip_sign_generic(const uint8_t* flip, float* data, size_t dim); void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim); - -void kacs_walk_generic(float* data, size_t len); void kacs_walk_avx2(float* data, size_t len); - -void fht_rotate_generic( +void fht_rotate_avx2( const float* data, float* rotated_data, size_t dim, @@ -46,7 +42,10 @@ void fht_rotate_generic( float scale_factor, const uint8_t* flip ); -void fht_rotate_avx2( + +void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim); +void kacs_walk_avx512(float* data, size_t len); +void fht_rotate_avx512( const float* data, float* rotated_data, size_t dim, @@ -56,11 +55,9 @@ void fht_rotate_avx2( const uint8_t* flip ); -void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim); - -void kacs_walk_avx512(float* data, size_t len); - -void fht_rotate_avx512( +void flip_sign_neon(const uint8_t* flip, float* data, size_t dim); +void kacs_walk_neon(float* data, size_t len); +void fht_rotate_neon( const float* data, float* rotated_data, size_t dim, diff --git a/include/rabitqlib/core/detail/simd/space_dispatch.hpp b/include/rabitqlib/core/detail/simd/space_dispatch.hpp index c3a0dc86..c3576354 100644 --- a/include/rabitqlib/core/detail/simd/space_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/space_dispatch.hpp @@ -10,11 +10,6 @@ float dot_product(const float* a, const float* b, size_t dim); float dot_product_dis(const float* a, const float* b, size_t dim); float l2norm_sqr(const float* a, size_t dim); -float euclidean_sqr_neon(const float* a, const float* b, size_t dim); -float dot_product_neon(const float* a, const float* b, size_t dim); -float dot_product_dis_neon(const float* a, const float* b, size_t dim); -float l2norm_sqr_neon(const float* a, size_t dim); - float euclidean_sqr_generic(const float* a, const float* b, size_t dim); float dot_product_generic(const float* a, const float* b, size_t dim); float dot_product_dis_generic(const float* a, const float* b, size_t dim); @@ -30,64 +25,62 @@ float dot_product_avx512(const float* a, const float* b, size_t dim); float dot_product_dis_avx512(const float* a, const float* b, size_t dim); float l2norm_sqr_avx512(const float* a, size_t dim); -namespace excode_ipimpl { +float euclidean_sqr_neon(const float* a, const float* b, size_t dim); +float dot_product_neon(const float* a, const float* b, size_t dim); +float dot_product_dis_neon(const float* a, const float* b, size_t dim); +float l2norm_sqr_neon(const float* a, size_t dim); -float ip16_fxu1_neon(const float* query, const uint8_t* code, size_t dim); -float ip64_fxu2_neon(const float* query, const uint8_t* code, size_t dim); -float ip64_fxu3_neon(const float* query, const uint8_t* code, size_t dim); -float ip16_fxu4_neon(const float* query, const uint8_t* code, size_t dim); -float ip64_fxu5_neon(const float* query, const uint8_t* code, size_t dim); -float ip64_fxu6_neon(const float* query, const uint8_t* code, size_t dim); -float ip64_fxu7_neon(const float* query, const uint8_t* code, size_t dim); -float ip16_fxu8_neon(const float* query, const uint8_t* code, size_t dim); +// The selector is rabitqlib::select_excode_ipfunc in rabitqlib/utils/space.hpp. +namespace excode_ipimpl { float ip16_fxu1_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip16_fxu1_avx2( +float ip64_fxu2_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu2_generic( +float ip64_fxu3_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu2_avx2( +float ip16_fxu4_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu3_generic( +float ip64_fxu5_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu3_avx2( +float ip64_fxu6_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip16_fxu4_generic( +float ip64_fxu7_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip16_fxu4_avx2( +float ip16_fxu8_generic( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu5_generic( + +float ip16_fxu1_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu5_avx2( +float ip64_fxu2_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu6_generic( +float ip64_fxu3_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu6_avx2( +float ip16_fxu4_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu7_generic( +float ip64_fxu5_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip64_fxu7_avx2( +float ip64_fxu6_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); -float ip16_fxu8_generic( - const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim +float 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 + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); float ip16_fxu1_avx512( @@ -112,84 +105,95 @@ 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 + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); +float ip16_fxu1_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip64_fxu2_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip64_fxu3_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip16_fxu4_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip64_fxu5_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip64_fxu6_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip64_fxu7_neon(const float* query, const uint8_t* compact_code, size_t dim); +float ip16_fxu8_neon(const float* query, const uint8_t* compact_code, size_t dim); + } // namespace excode_ipimpl -void new_transpose_bin_neon( - const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits -); -void new_transpose_bin_512_neon( - const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits -); -void scalar_quantize_uint8_neon( +// The transpose and mask-IP dispatchers are declared in rabitqlib/utils/space.hpp. +void scalar_quantize_uint8( uint8_t* result, const float* data, size_t dim, float lo, float delta ); -void scalar_quantize_uint16_neon( +void scalar_quantize_uint16( uint16_t* result, const float* data, size_t dim, float lo, float delta ); void new_transpose_bin_generic( - const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits -); -void new_transpose_bin_avx2( - const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ); void new_transpose_bin_512_generic( - const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query +); +float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t padded_dim); +float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t padded_dim); +void scalar_quantize_uint8_generic( + uint8_t* result, const float* data, size_t dim, float lo, float delta +); +void scalar_quantize_uint16_generic( + uint16_t* result, const float* data, size_t dim, float lo, float delta +); + +void new_transpose_bin_avx2( + const uint16_t* query, uint64_t* transposed_query, 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 + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ); -float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t dim); -float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t dim); -float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t dim); float mask_ip_x0_q_avx2(const float* query, const uint8_t* data, size_t padded_dim); -float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t dim); float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_dim); -void scalar_quantize_uint8_generic( - uint8_t* result, const float* data, size_t dim, float lo, float delta -); void scalar_quantize_uint8_avx2( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ); // Only for (value - lo) / delta >= 0, such as FastScan LUT entries at least their // minimum; there it matches scalar_quantize_uint8_avx2 with fewer operations. void scalar_quantize_nonnegative_uint8_avx2( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta -); -void scalar_quantize_uint16_generic( - uint16_t* result, const float* data, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ); void scalar_quantize_uint16_avx2( - uint16_t* result, const float* vec0, size_t dim, float lo, float delta + uint16_t* result, const float* data, 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 + const uint16_t* query, uint64_t* transposed_query, 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 + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ); float mask_ip_x0_q_avx512(const float* query, const uint8_t* data, size_t padded_dim); 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 + uint8_t* result, const float* data, size_t dim, float lo, float delta ); // Only for (value - lo) / delta >= 0; see scalar_quantize_nonnegative_uint8_avx2. void scalar_quantize_nonnegative_uint8_avx512( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, 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 + uint16_t* result, const float* data, 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 new_transpose_bin_neon( + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ); -void scalar_quantize_uint16( - uint16_t* result, const float* vec0, size_t dim, float lo, float delta +void new_transpose_bin_512_neon( + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query +); +float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t padded_dim); +float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t padded_dim); +void scalar_quantize_uint8_neon( + uint8_t* result, const float* data, size_t dim, float lo, float delta +); +void scalar_quantize_uint16_neon( + uint16_t* result, const float* data, size_t dim, float lo, float delta ); } // namespace rabitqlib::simd diff --git a/include/rabitqlib/core/detail/simd/warmup_dispatch.hpp b/include/rabitqlib/core/detail/simd/warmup_dispatch.hpp index b0150a88..c3df64c0 100644 --- a/include/rabitqlib/core/detail/simd/warmup_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/warmup_dispatch.hpp @@ -5,33 +5,35 @@ namespace rabitqlib::simd { -float warmup_ip_x0_q_512_neon( +// The dispatchers are declared in rabitqlib/utils/warmup_space.hpp. + +float warmup_ip_x0_q_512_generic( const uint8_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ); -float warmup_ip_x0_q_512_neon( +float warmup_ip_x0_q_512_generic( const uint64_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ); -float warmup_ip_x0_q_512_generic( +float warmup_ip_x0_q_512_avx2( const uint8_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ); float warmup_ip_x0_q_512_avx2( - const uint8_t* data, + const uint64_t* data, const uint64_t* query, float delta, float vl, @@ -39,15 +41,15 @@ float warmup_ip_x0_q_512_avx2( size_t b_query ); -float warmup_ip_x0_q_512_generic( - const uint64_t* data, +float warmup_ip_x0_q_512_avx512( + const uint8_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ); -float warmup_ip_x0_q_512_avx2( +float warmup_ip_x0_q_512_avx512( const uint64_t* data, const uint64_t* query, float delta, @@ -56,7 +58,7 @@ float warmup_ip_x0_q_512_avx2( size_t b_query ); -float warmup_ip_x0_q_512_avx512( +float warmup_ip_x0_q_512_neon( const uint8_t* data, const uint64_t* query, float delta, @@ -64,8 +66,7 @@ float warmup_ip_x0_q_512_avx512( size_t padded_dim, size_t b_query ); - -float warmup_ip_x0_q_512_avx512( +float warmup_ip_x0_q_512_neon( const uint64_t* data, const uint64_t* query, float delta, diff --git a/include/rabitqlib/core/encoder.hpp b/include/rabitqlib/core/encoder.hpp index 50445f61..25a77f87 100644 --- a/include/rabitqlib/core/encoder.hpp +++ b/include/rabitqlib/core/encoder.hpp @@ -17,8 +17,7 @@ namespace rabitqlib::quant { inline RabitqConfig faster_config(size_t dim, size_t total_bits) { RabitqConfig config; if (total_bits > 1) { - config.t_const = - rabitq_impl::ex_bits::get_const_scaling_factors(dim, total_bits - 1); + config.t_const = rabitq_impl::ex_bits::const_scaling_factor(dim, total_bits - 1); } return config; } @@ -51,30 +50,30 @@ inline void quantize_compact_one_bit( cur_bin_data.f_error() = f_error; } -template +template inline void quantize_compact_one_bit( const T* data, const T* centroid, size_t padded_dim, - TC compact_code, + CodePointer compact_code, T& f_add, T& f_rescale, - T& ferror, + T& f_error, MetricType metric_type = METRIC_L2 ) { rabitq_impl::one_bit::one_bit_compact_code( - data, centroid, padded_dim, compact_code, f_add, f_rescale, ferror, metric_type + data, centroid, padded_dim, compact_code, f_add, f_rescale, f_error, metric_type ); } -template +template inline void quantize_compact_one_bit( const T* data, size_t padded_dim, - TC compact_code, + CodePointer compact_code, T& f_add, T& f_rescale, - T& ferror, + T& f_error, MetricType metric_type = METRIC_L2 ) { std::vector centroid(padded_dim, 0); @@ -85,7 +84,7 @@ inline void quantize_compact_one_bit( compact_code, f_add, f_rescale, - ferror, + f_error, metric_type ); } @@ -102,10 +101,10 @@ inline void quantize_compact_ex_bits( ) { ExDataMap cur_ex_data(ex_data, padded_dim, ex_bits); - // we do not use this error factor here + // The ex-code error factor is not stored. T f_add_ex; T f_rescale_ex; - T ex_error; + T f_error_ex; rabitq_impl::ex_bits::ex_bits_compact_code( data, @@ -115,7 +114,7 @@ inline void quantize_compact_ex_bits( cur_ex_data.ex_code(), f_add_ex, f_rescale_ex, - ex_error, + f_error_ex, metric_type, config.t_const ); @@ -265,9 +264,10 @@ template inline void reconstruct_vec( const TP* quantized_vec, T delta, T vl, size_t dim, T* results ) { - RowMajorArrayMap result_arr(results, 1, dim); + const auto cols = static_cast(dim); + RowMajorArrayMap result_arr(results, 1, cols); result_arr = - (ConstRowMajorArrayMap(quantized_vec, 1, dim).template cast() * delta) + vl; + (ConstRowMajorArrayMap(quantized_vec, 1, cols).template cast() * delta) + vl; } template @@ -282,34 +282,35 @@ inline void reconstruct_full_vec( ) { const T code_center = -static_cast((static_cast(1) << bits) - 1) / 2; const T scale = metric_type == METRIC_L2 ? -f_rescale / 2 : -f_rescale; - RowMajorArrayMap result_arr(results, 1, dim); + const auto cols = static_cast(dim); + RowMajorArrayMap result_arr(results, 1, cols); result_arr = - ConstRowMajorArrayMap(centroid, 1, dim) + - scale * (ConstRowMajorArrayMap(quantized_vec, 1, dim).template cast() + + ConstRowMajorArrayMap(centroid, 1, cols) + + scale * (ConstRowMajorArrayMap(quantized_vec, 1, cols).template cast() + code_center); } namespace detail { -template -inline TF full_est_dist_from_ip( - TF ip, size_t bits, TF f_add, TF f_rescale, TF g_add, TF k1xsumq +template +inline T full_est_dist_from_ip( + T ip, size_t bits, T f_add, T f_rescale, T g_add, T k1xsumq ) { return f_add + g_add + (f_rescale * (ip + k1xsumq * static_cast((1 << bits) - 1))); } } // namespace detail -template -inline TF full_est_dist( - const TI* quantized_vec, - const TF* query, - TF (*ip_func)(const TF*, const TI*, size_t), +template +inline T full_est_dist( + const TP* quantized_vec, + const T* query, + T (*ip_func)(const T*, const TP*, size_t), size_t dim, size_t bits, - TF f_add, - TF f_rescale, - TF g_add, - TF k1xsumq + T f_add, + T f_rescale, + T g_add, + T k1xsumq ) { return detail::full_est_dist_from_ip( ip_func(query, quantized_vec, dim), bits, f_add, f_rescale, g_add, k1xsumq diff --git a/include/rabitqlib/core/estimator.hpp b/include/rabitqlib/core/estimator.hpp index fca698c9..b3073241 100644 --- a/include/rabitqlib/core/estimator.hpp +++ b/include/rabitqlib/core/estimator.hpp @@ -27,7 +27,7 @@ struct DistanceEstimate { * same vector and query, as returned by the matching base-distance estimator. * Returns the refined squared-L2 or 1 - dot estimate, without an error margin. */ -template +template inline float split_distance_boosting( const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), @@ -47,7 +47,7 @@ inline float split_distance_boosting( return ex_dist; } -template +template inline void split_single_fulldist_direct( const char* bin_data, const char* ex_data, @@ -137,7 +137,7 @@ inline void split_single_estdist( (cur_bin.f_rescale() * (ip_x0_qr + query_wrapper.k1xsumq())); lower_bound = estimated_distance - (cur_bin.f_error() * g_error); -}; +} template inline void split_single_estdist_direct( @@ -166,7 +166,7 @@ inline void split_single_estdist_direct( (cur_bin.f_rescale() * (ip_x0_qr + query_wrapper.k1xsumq())); lower_bound = estimated_distance - (cur_bin.f_error() * g_error); -}; +} inline void xy_single_base_dist( const char* base_data, @@ -230,7 +230,7 @@ inline void xy_single_full_dist( // Named-result functions keep the numerical implementation and legacy output-reference // interfaces intact. Internal callers do not need to remember output argument order. -template +template inline detail::DistanceEstimate estimate_split_single_full_direct( const char* bin_data, const char* ex_data, diff --git a/include/rabitqlib/core/metric.hpp b/include/rabitqlib/core/metric.hpp index e3d224c8..b65a6f08 100644 --- a/include/rabitqlib/core/metric.hpp +++ b/include/rabitqlib/core/metric.hpp @@ -4,7 +4,7 @@ #include namespace rabitqlib { -enum MetricType : std::uint8_t { METRIC_L2, METRIC_IP }; +enum MetricType : uint8_t { METRIC_L2, METRIC_IP }; inline void validate_metric_type(MetricType metric_type) { if (metric_type != METRIC_L2 && metric_type != METRIC_IP) { diff --git a/include/rabitqlib/core/pack_excode.hpp b/include/rabitqlib/core/pack_excode.hpp index 217ea5b0..186cd48a 100644 --- a/include/rabitqlib/core/pack_excode.hpp +++ b/include/rabitqlib/core/pack_excode.hpp @@ -9,7 +9,7 @@ namespace rabitqlib::quant::rabitq_impl::ex_bits { inline void packing_1bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - // ! require dim % 16 == 0 + // Requires dim % 16 == 0. for (size_t j = 0; j < dim; j += 16) { uint16_t code = 0; for (size_t i = 0; i < 16; ++i) { diff --git a/include/rabitqlib/core/query.hpp b/include/rabitqlib/core/query.hpp index 531e4846..647f6b2a 100644 --- a/include/rabitqlib/core/query.hpp +++ b/include/rabitqlib/core/query.hpp @@ -69,9 +69,9 @@ class SplitSingleQuery { validate_metric_type(metric_type); - float c_1 = -static_cast((1 << 1) - 1) / 2.F; - float c_base = -static_cast((1U << base_bits) - 1) / 2.F; - float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.F; + float c_1 = -static_cast((1 << 1) - 1) / 2.0F; + float c_base = -static_cast((1U << base_bits) - 1) / 2.0F; + float c_b = -static_cast((1U << (base_bits + ex_bits)) - 1) / 2.0F; T sumq = std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); @@ -85,9 +85,7 @@ class SplitSingleQuery { quant::quantize_scalar( rotated_query, padded_dim, kNumBits, quant_query.data(), delta_, vl_, config ); - rabitqlib::new_transpose_bin_512( - quant_query.data(), query_bin_.data(), padded_dim, kNumBits - ); + new_transpose_bin_512(quant_query.data(), query_bin_.data(), padded_dim, kNumBits); } [[nodiscard]] size_t num_bits() const { return kNumBits; } diff --git a/include/rabitqlib/core/rotator.hpp b/include/rabitqlib/core/rotator.hpp index 6dfe31d2..9fe2adb3 100644 --- a/include/rabitqlib/core/rotator.hpp +++ b/include/rabitqlib/core/rotator.hpp @@ -32,16 +32,16 @@ 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 rotate(const T* data, T* rotated_data) 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; - virtual size_t dump_bytes() const = 0; - [[nodiscard]] size_t size() const { return this->padded_dim_; } + [[nodiscard]] virtual size_t dump_bytes() const = 0; + [[nodiscard]] size_t size() const { return padded_dim_; } }; namespace rotator_impl { @@ -131,26 +131,29 @@ class MatrixRotator : public Rotator { std::memcpy(data, rand_mat_.data(), MatrixRotator::dump_bytes()); } - size_t dump_bytes() const override { + [[nodiscard]] size_t dump_bytes() const override { return sizeof(T) * this->dim_ * this->padded_dim_; } - void rotate(const T* vec, T* rotated_vec) const override { + void rotate(const T* data, T* rotated_data) const override { if constexpr (std::is_same_v) { // Match Eigen's alias-safe assignment: callers may rotate in place. std::vector result(this->padded_dim_); simd::matrix_product( - vec, rand_mat_.data(), result.data(), 1, this->dim_, this->padded_dim_ + data, rand_mat_.data(), result.data(), 1, this->dim_, this->padded_dim_ ); - std::copy(result.begin(), result.end(), rotated_vec); + std::copy(result.begin(), result.end(), rotated_data); } else { - ConstRowMajorMatrixMap v(vec, 1, this->dim_); - RowMajorMatrixMap rv(rotated_vec, 1, this->padded_dim_); - rv = v * this->rand_mat_; + ConstRowMajorMatrixMap v(data, 1, static_cast(this->dim_)); + RowMajorMatrixMap rv( + rotated_data, 1, static_cast(this->padded_dim_) + ); + rv = v * rand_mat_; } } }; +// Internal linkage: the per-ISA HNSW search sources include this header. static inline void flip_sign(const uint8_t* flip, float* data, size_t dim) { simd::flip_sign(flip, data, dim); } @@ -212,15 +215,17 @@ class FhtKacRotator : public Rotator { std::memcpy(data, flip_.data(), FhtKacRotator::dump_bytes()); } - size_t dump_bytes() const override { return sizeof(uint8_t) * flip_.size(); } + [[nodiscard]] size_t dump_bytes() const override { + return sizeof(uint8_t) * flip_.size(); + } FhtKacRotator(const FhtKacRotator&) = default; FhtKacRotator& operator=(const FhtKacRotator& other) { // Preserve the old dimensions if copying the sign storage fails. flip_ = other.flip_; - this->dim_ = other.dim_; - this->padded_dim_ = other.padded_dim_; + dim_ = other.dim_; + padded_dim_ = other.padded_dim_; trunc_dim_ = other.trunc_dim_; scale_factor_ = other.scale_factor_; return *this; @@ -228,9 +233,9 @@ class FhtKacRotator : public Rotator { static void kacs_walk(float* data, size_t len) { simd::kacs_walk(data, len); } - void rotate(const float* data, float* rotated_vec) const override { + void rotate(const float* data, float* rotated_data) const override { simd::fht_rotate( - data, rotated_vec, dim_, padded_dim_, trunc_dim_, scale_factor_, flip_.data() + data, rotated_data, dim_, padded_dim_, trunc_dim_, scale_factor_, flip_.data() ); } }; diff --git a/include/rabitqlib/core/types.hpp b/include/rabitqlib/core/types.hpp index ad2fd582..041c60db 100644 --- a/include/rabitqlib/core/types.hpp +++ b/include/rabitqlib/core/types.hpp @@ -3,7 +3,7 @@ #include namespace rabitqlib { -enum ScalarQuantizerType : std::uint8_t { RECONSTRUCTION, UNBIASED_ESTIMATION, PLAIN }; +enum ScalarQuantizerType : uint8_t { RECONSTRUCTION, UNBIASED_ESTIMATION, PLAIN }; enum class RotatorType : uint8_t { MatrixRotator, FhtKacRotator }; namespace quant { diff --git a/include/rabitqlib/fastscan/detail/encoder.hpp b/include/rabitqlib/fastscan/detail/encoder.hpp index f31768a0..b04b2a11 100644 --- a/include/rabitqlib/fastscan/detail/encoder.hpp +++ b/include/rabitqlib/fastscan/detail/encoder.hpp @@ -10,8 +10,8 @@ namespace rabitqlib::quant::rabitq_impl::one_bit { // Encoding requires a positive padded_dim divisible by 8; index pipelines use multiples -// of 32. packed_code needs ceil(num / 32) * 32 * (padded_dim / 8) bytes, including tail -// padding. +// of 32. packed_code needs ceil(num_points / 32) * 32 * (padded_dim / 8) bytes, including +// tail padding. template < typename T, bool Parallel = false, @@ -21,7 +21,7 @@ template < inline void one_bit_batch_code( const T* data, const T* centroid, - size_t num, + size_t num_points, size_t padded_dim, uint8_t* packed_code, FAdd f_add, @@ -29,12 +29,12 @@ inline void one_bit_batch_code( FError f_error, MetricType metric_type = METRIC_L2 ) { - std::vector compact_codes(num * padded_dim / 8); + std::vector compact_codes(num_points * padded_dim / 8); one_bit_compact_codes( data, centroid, - num, + num_points, padded_dim, compact_codes.data(), f_add, @@ -42,7 +42,7 @@ inline void one_bit_batch_code( f_error, metric_type ); - fastscan::pack_codes(padded_dim, compact_codes.data(), num, packed_code); + fastscan::pack_codes(padded_dim, compact_codes.data(), num_points, packed_code); } } // namespace rabitqlib::quant::rabitq_impl::one_bit diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp index 32bf4be8..bd8ecf74 100644 --- a/include/rabitqlib/fastscan/detail/lut.hpp +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -183,8 +183,8 @@ class Lut { [[nodiscard]] size_t padded_dim() const { return lut_.empty() ? 0 : table_length_ / 4; } [[nodiscard]] bool use_hacc() const { return lut_.size() > table_length_; } - [[nodiscard]] const uint8_t* lut() const { return lut_.data(); }; - [[nodiscard]] T delta() const { return delta_; }; - [[nodiscard]] T sum_vl() const { return sum_vl_lut_; }; + [[nodiscard]] const uint8_t* lut() const { return lut_.data(); } + [[nodiscard]] T delta() const { return delta_; } + [[nodiscard]] T sum_vl() const { return sum_vl_lut_; } }; } // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/detail/qg_estimator.hpp b/include/rabitqlib/fastscan/detail/qg_estimator.hpp index b761eda5..f350aaaa 100644 --- a/include/rabitqlib/fastscan/detail/qg_estimator.hpp +++ b/include/rabitqlib/fastscan/detail/qg_estimator.hpp @@ -20,10 +20,14 @@ static inline void qg_batch_estdist_impl( T* estimated_distances ) { ConstQGBatchDataMap cur_batch(batch_data, padded_dim); - std::array accu_values{}; - Accumulate(cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim); + std::array accumulated_values{}; + Accumulate( + cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim + ); - ConstRowMajorArrayMap ip_arr(accu_values.data(), 1, fastscan::kBatchSize); + ConstRowMajorArrayMap ip_arr( + accumulated_values.data(), 1, fastscan::kBatchSize + ); std::array f_add_values; std::array f_rescale_values; cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); @@ -33,10 +37,13 @@ static inline void qg_batch_estdist_impl( f_rescale_values.data(), 1, fastscan::kBatchSize ); - RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); + RowMajorArrayMap estimated_distances_arr( + estimated_distances, 1, fastscan::kBatchSize + ); - est_dist_arr = f_add_arr + query_wrapper.g_add() + - (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + - query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); + estimated_distances_arr = + f_add_arr + query_wrapper.g_add() + + (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + + query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); } } // namespace rabitqlib::fastscan::detail diff --git a/include/rabitqlib/fastscan/detail/simd.hpp b/include/rabitqlib/fastscan/detail/simd.hpp index cedeba0f..abc552ac 100644 --- a/include/rabitqlib/fastscan/detail/simd.hpp +++ b/include/rabitqlib/fastscan/detail/simd.hpp @@ -8,14 +8,14 @@ namespace rabitqlib::fastscan::simd { // NEON uses the portable/AVX2 high-accuracy LUT byte order. void accumulate_neon( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* __restrict__ result, + const uint8_t* __restrict__ lookup_table, + int32_t* __restrict__ accumulated_values, size_t dim ); void accumulate_hacc_neon( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* result, + const uint8_t* __restrict__ lookup_table, + int32_t* accumulated_values, size_t dim ); @@ -41,32 +41,32 @@ void pack_lut_range_neon( // Scalar reference for backend correctness tests; never selected by runtime dispatch. void accumulate_generic( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* __restrict__ result, + const uint8_t* __restrict__ lookup_table, + int32_t* __restrict__ accumulated_values, size_t dim ); void accumulate_unsupported( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* __restrict__ result, + const uint8_t* __restrict__ lookup_table, + int32_t* __restrict__ accumulated_values, size_t dim ); void accumulate_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ); void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc_generic( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* result, + const uint8_t* __restrict__ lookup_table, + int32_t* accumulated_values, size_t dim ); void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ); @@ -74,13 +74,13 @@ void accumulate_hacc_avx2( void accumulate_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ); void transfer_lut_hacc_avx512(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ); diff --git a/include/rabitqlib/fastscan/encoder.hpp b/include/rabitqlib/fastscan/encoder.hpp index 7a11deac..42085f29 100644 --- a/include/rabitqlib/fastscan/encoder.hpp +++ b/include/rabitqlib/fastscan/encoder.hpp @@ -15,23 +15,23 @@ namespace rabitqlib::quant { // Single-batch entry points accept 0..32 rows. Zero rows leave output unchanged. // Allocate full BatchDataMap/QGBatchDataMap storage even for a partial batch; -// only the first num factors/results are meaningful. Inputs are borrowed for the call. -// Parallel=true enables OpenMP when built with it, otherwise encoding remains serial. +// only the first batch_size factors/results are meaningful. Inputs are borrowed for the +// call. Parallel=true enables OpenMP when built with it, otherwise encoding remains serial. template inline void quantize_one_batch( const T* data, const T* centroid, - size_t num, + size_t batch_size, size_t padded_dim, char* batch_data, MetricType metric_type = METRIC_L2 ) { - if (num > fastscan::kBatchSize) { + if (batch_size > fastscan::kBatchSize) { throw std::invalid_argument( "FastScan single-batch encoding accepts at most 32 vectors" ); } - if (num == 0) { + if (batch_size == 0) { return; } BatchDataMap this_batch(batch_data, padded_dim); @@ -39,7 +39,7 @@ inline void quantize_one_batch( rabitq_impl::one_bit::one_bit_batch_code( data, centroid, - num, + batch_size, padded_dim, this_batch.bin_code(), this_batch.f_add(), @@ -52,41 +52,41 @@ inline void quantize_one_batch( template inline void quantize_one_batch( const T* data, - size_t num, + size_t batch_size, size_t padded_dim, char* batch_data, MetricType metric_type = METRIC_L2 ) { std::vector centroid(padded_dim, 0); quantize_one_batch( - data, centroid.data(), num, padded_dim, batch_data, metric_type + data, centroid.data(), batch_size, padded_dim, batch_data, metric_type ); } template -static inline void quantize_qg_batch( +inline void quantize_qg_batch( const T* data, const T* centroid, - size_t num, + size_t batch_size, size_t padded_dim, char* batch_data, MetricType metric_type = METRIC_L2 ) { - if (num > fastscan::kBatchSize) { + if (batch_size > fastscan::kBatchSize) { throw std::invalid_argument( "FastScan single-batch encoding accepts at most 32 vectors" ); } - if (num == 0) { + if (batch_size == 0) { return; } - std::array f_error{}; // we dont need this factor for qg + std::array f_error{}; // QG batches do not store f_error QGBatchDataMap cur_batch(batch_data, padded_dim); rabitq_impl::one_bit::one_bit_batch_code( data, centroid, - num, + batch_size, padded_dim, cur_batch.bin_code(), cur_batch.f_add(), @@ -97,21 +97,23 @@ static inline void quantize_qg_batch( } template -static inline void quantize_qg_batch( +inline void quantize_qg_batch( const T* data, - size_t num, + size_t batch_size, size_t padded_dim, char* batch_data, MetricType metric_type = METRIC_L2 ) { std::vector centroid(padded_dim, 0); - quantize_qg_batch(data, centroid.data(), num, padded_dim, batch_data, metric_type); + quantize_qg_batch( + data, centroid.data(), batch_size, padded_dim, batch_data, metric_type + ); } inline void quantize_split_batch( const float* data, const float* centroid, - size_t num_points, + size_t batch_size, size_t padded_dim, size_t ex_bits, char* batch_data, @@ -119,10 +121,10 @@ inline void quantize_split_batch( MetricType metric_type = METRIC_L2, RabitqConfig config = RabitqConfig() ) { - quantize_one_batch(data, centroid, num_points, padded_dim, batch_data, metric_type); + quantize_one_batch(data, centroid, batch_size, padded_dim, batch_data, metric_type); if (ex_bits > 0) { - for (size_t i = 0; i < num_points; ++i) { + for (size_t i = 0; i < batch_size; ++i) { quantize_compact_ex_bits( data, centroid, padded_dim, ex_bits, ex_data, metric_type, config ); diff --git a/include/rabitqlib/fastscan/estimator.hpp b/include/rabitqlib/fastscan/estimator.hpp index e6c7ff10..6b843b7f 100644 --- a/include/rabitqlib/fastscan/estimator.hpp +++ b/include/rabitqlib/fastscan/estimator.hpp @@ -75,12 +75,11 @@ inline void qg_batch_estdist( } if constexpr (std::is_same_v) { simd::qg_batch_estdist(batch_data, query_wrapper, padded_dim, estimated_distances); - return; + } else { + fastscan::detail::qg_batch_estdist_impl( + batch_data, query_wrapper, padded_dim, estimated_distances + ); } - - fastscan::detail::qg_batch_estdist_impl( - batch_data, query_wrapper, padded_dim, estimated_distances - ); } } // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index bc6546aa..9cb8945c 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -9,7 +9,7 @@ namespace rabitqlib::fastscan { -constexpr static std::array kPos = { +inline constexpr std::array kPos = { 3 /*0000*/, 3 /*0001*/, 2 /*0010*/, @@ -34,18 +34,18 @@ constexpr static std::array kPos = { void accumulate( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ); -// pack lookup table for fastscan, for each 4 dim, we have 16 (2^4) different results -// ! dim % 4 == 0 +// Build one 16-entry table of subset sums for each group of 4 query coordinates. +// Requires dim % 4 == 0. template inline void pack_lut(size_t dim, const T* __restrict__ query, T* __restrict__ lut) { - size_t num_codebook = dim >> 2; + size_t num_groups = dim >> 2; // Preserve +0 when the first selected coordinate is -0. volatile T zero = 0; - for (size_t i = 0; i < num_codebook; ++i) { + for (size_t group = 0; group < num_groups; ++group) { lut[0] = zero; for (size_t j = 1; j < 16; ++j) { const size_t lowbit = j & (size_t{0} - j); diff --git a/include/rabitqlib/fastscan/highacc_fastscan.hpp b/include/rabitqlib/fastscan/highacc_fastscan.hpp index af7d3c38..39ff855f 100644 --- a/include/rabitqlib/fastscan/highacc_fastscan.hpp +++ b/include/rabitqlib/fastscan/highacc_fastscan.hpp @@ -15,7 +15,7 @@ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut); // Accumulate one batch into int32_t results. Throws if a result exceeds int32_t. void accumulate_hacc( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ); diff --git a/include/rabitqlib/fastscan/pack_codes.hpp b/include/rabitqlib/fastscan/pack_codes.hpp index 7d152d10..860ec623 100644 --- a/include/rabitqlib/fastscan/pack_codes.hpp +++ b/include/rabitqlib/fastscan/pack_codes.hpp @@ -10,15 +10,16 @@ namespace rabitqlib::fastscan { -constexpr static size_t kBatchSize = 32; // number of vectors in each batch +inline constexpr size_t kBatchSize = 32; // number of vectors in each batch -// Interleave lanes 0..7 and 8..15 for the 16-bit SIMD accumulators. -constexpr static std::array kPerm0 = { +// Interleave lanes 0..7 and 8..15 for the 16-bit SIMD accumulators. Internal linkage +// lets hot kernels address the table directly instead of through the GOT. +static constexpr std::array kPerm0 = { 0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15}; -template +template static inline void get_column( - const T* src, size_t rows, size_t cols, size_t row, size_t col, TA& dest + const TP* src, size_t rows, size_t cols, size_t row, size_t col, Column& dest ) { size_t k = 0; size_t max_k = std::min(rows - row, dest.size()); @@ -37,21 +38,19 @@ static inline void get_column( * their low nibbles; the first vector always occupies the output's low nibble. * * @param padded_dim dimension of quantized data (i.e., quantization code) - * @param quantization_code num rows of padded_dim / 8 compact sign-code bytes - * @param num number of logical vectors - * @param blocks output, ceil(num / 32) * 32 * padded_dim / 8 bytes + * @param quantization_code num_points rows of padded_dim / 8 compact sign-code bytes + * @param num_points number of logical vectors + * @param blocks output, ceil(num_points / 32) * 32 * padded_dim / 8 bytes */ inline void pack_codes( - size_t padded_dim, const uint8_t* quantization_code, size_t num, uint8_t* blocks + size_t padded_dim, const uint8_t* quantization_code, size_t num_points, uint8_t* blocks ) { - // consider codes is a matrix - // rows = number of vectors - // cols = number of uint8_t of one vector's code + // Treat the codes as a matrix: one row per vector, one column per code byte. size_t cols = padded_dim / 8; // Full batches can gather directly in output order, without a temporary column. size_t row = 0; - for (; num - row >= kBatchSize; row += kBatchSize) { + for (; num_points - row >= kBatchSize; row += kBatchSize) { for (size_t i = 0; i < cols; ++i) { for (size_t j = 0; j < 16; ++j) { const size_t first = row + kPerm0[j]; @@ -65,10 +64,10 @@ inline void pack_codes( } // Only the final partial batch needs a zero-padded column. - if (row < num) { - std::array col; + if (row < num_points) { + std::array col; for (size_t i = 0; i < cols; ++i) { - get_column(quantization_code, num, cols, row, i, col); + get_column(quantization_code, num_points, cols, row, i, col); for (size_t j = 0; j < 16; ++j) { // the lower 4 bits represent vector 0 to 15 // the upper 4 bits represent vector 16 to 31 @@ -83,20 +82,21 @@ inline void pack_codes( } /** - * @brief Inverse of pack_codes: recover the compact quantization codes of the first `num` - * vectors from their packed blocks. Padding lanes of a partial final batch are ignored. + * @brief Inverse of pack_codes: recover the compact quantization codes of the first + * `num_points` vectors from their packed blocks. Padding lanes of a partial final batch + * are ignored. * * @param padded_dim dimension of quantized data (i.e., quantization code) * @param blocks packed quantization code, as produced by pack_codes - * @param num number of quantization code to recover - * @param quantization_code recovered codes, num * padded_dim / 8 bytes + * @param num_points number of quantization codes to recover + * @param quantization_code recovered codes, num_points * padded_dim / 8 bytes */ inline void unpack_codes( - size_t padded_dim, const uint8_t* blocks, size_t num, uint8_t* quantization_code + size_t padded_dim, const uint8_t* blocks, size_t num_points, uint8_t* quantization_code ) { const size_t cols = padded_dim / 8; - for (size_t row = 0; row < num; row += kBatchSize) { + for (size_t row = 0; row < num_points; row += kBatchSize) { for (size_t i = 0; i < cols; ++i) { for (size_t j = 0; j < 16; ++j) { // blocks[j] holds both input vectors' high nibbles; blocks[j + 16] @@ -107,10 +107,10 @@ inline void unpack_codes( ); const auto b = static_cast((blocks[j] & 0xF0) | (blocks[j + 16] >> 4)); - if (first < num) { + if (first < num_points) { quantization_code[first * cols + i] = a; } - if (first + 16 < num) { + if (first + 16 < num_points) { quantization_code[(first + 16) * cols + i] = b; } } diff --git a/include/rabitqlib/fastscan/query.hpp b/include/rabitqlib/fastscan/query.hpp index 65d21c35..417e4b5e 100644 --- a/include/rabitqlib/fastscan/query.hpp +++ b/include/rabitqlib/fastscan/query.hpp @@ -9,6 +9,13 @@ #include "rabitqlib/fastscan/detail/lut.hpp" namespace rabitqlib { +namespace fastscan::detail { +// Offset -(2^bits - 1) / 2 that centers an unsigned bits-bit code. +constexpr float code_center(size_t bits) { + return -static_cast((1 << bits) - 1) / 2.0F; +} +} // namespace fastscan::detail + /** * @brief Own a standard FastScan LUT; the query input is not retained. */ @@ -17,7 +24,7 @@ class BatchQuery { private: Lut lookup_table_; // Owns the prepared bytes; reset scratch is never retained. T G_add_ = 0; - T G_k1xSumq_ = 0; // G_k1xSumq + T G_k1xSumq_ = 0; MetricType metric_type_ = METRIC_L2; static Lut make_lut( @@ -28,7 +35,7 @@ class BatchQuery { } static T query_correction(const T* rotated_query, size_t padded_dim) { - constexpr float c_1 = -((1 << 1) - 1) / 2.F; + constexpr float c_1 = fastscan::detail::code_center(1); const T sumq = std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); return sumq * c_1; @@ -119,8 +126,8 @@ class SplitBatchQuery { MetricType metric_type_ = METRIC_L2; void prepare_coefficients(const T* rotated_query, size_t padded_dim, size_t ex_bits) { - const float c_1 = -static_cast((1 << 1) - 1) / 2.F; - const float c_b = -static_cast((1 << (ex_bits + 1)) - 1) / 2.F; + constexpr float c_1 = fastscan::detail::code_center(1); + const float c_b = fastscan::detail::code_center(ex_bits + 1); const T sumq = std::accumulate(rotated_query, rotated_query + padded_dim, static_cast(0)); diff --git a/include/rabitqlib/index/detail/buffer.hpp b/include/rabitqlib/index/detail/buffer.hpp index c2f64cf5..cdc957d2 100644 --- a/include/rabitqlib/index/detail/buffer.hpp +++ b/include/rabitqlib/index/detail/buffer.hpp @@ -15,16 +15,18 @@ inline constexpr PID kSearchBufferCheckedMask = PID{1} << (std::numeric_limits::digits - 1); inline constexpr size_t kSearchBufferMaxPointCount = static_cast(kSearchBufferCheckedMask); -/** - * @brief sorted linear buffer, used as beam set for graph-based ANN search. In symphonyqg, - * the search buffer may contain duplicate id with different distances - * - */ + +// Sorted linear buffer, used as the beam set for graph-based ANN search. In SymphonyQG, +// the search buffer may contain duplicate IDs with different distances. template class SearchBuffer { private: - std::vector, memory::AlignedAllocator>> data_; - size_t size_ = 0, cur_ = 0, capacity_ = 0; + using Storage = std::vector, memory::AlignedAllocator>>; + + Storage data_; + size_t size_ = 0; + size_t cur_ = 0; + size_t capacity_ = 0; [[nodiscard]] static size_t storage_size(size_t capacity) { if (capacity == std::numeric_limits::max()) { @@ -33,23 +35,22 @@ class SearchBuffer { return capacity + 1; } - [[nodiscard]] auto binary_search(T dist) const { + [[nodiscard]] size_t binary_search(T dist) const { size_t lo = 0; size_t len = size_; - size_t half; while (len > 1) { - half = len >> 1; + const size_t half = len >> 1; len -= half; lo += static_cast(data_[lo + half - 1].distance < dist) * half; } return (lo < size_ && data_[lo].distance < dist) ? lo + 1 : lo; } - // set top bit to 1 as checked - static void set_checked(PID& data_id) { data_id |= kSearchBufferCheckedMask; } + // Mark a candidate as checked by setting the ID's top bit. + static void set_checked(PID& internal_id) { internal_id |= kSearchBufferCheckedMask; } - [[nodiscard]] static auto is_checked(PID data_id) -> bool { - return (data_id & kSearchBufferCheckedMask) != 0; + [[nodiscard]] static bool is_checked(PID internal_id) { + return (internal_id & kSearchBufferCheckedMask) != 0; } public: @@ -58,28 +59,28 @@ class SearchBuffer { explicit SearchBuffer(size_t capacity) : data_(storage_size(capacity)), capacity_(capacity) {} - // insert a data point into buffer - void insert(PID data_id, T dist) { + // Insert a data point into the buffer. + void insert(PID internal_id, T dist) { if (is_full(dist)) { return; } size_t lo = binary_search(dist); std::memmove(&data_[lo + 1], &data_[lo], (size_ - lo) * sizeof(AnnCandidate)); - data_[lo] = AnnCandidate(data_id, dist); + data_[lo] = AnnCandidate(internal_id, dist); size_ += static_cast(size_ < capacity_); cur_ = lo < cur_ ? lo : cur_; } - // get unchecked candidate with minimum distance + // Get the unchecked candidate with the minimum distance. PID pop() { - PID cur_id = data_[cur_].id; + const PID internal_id = data_[cur_].id; set_checked(data_[cur_].id); ++cur_; while (cur_ < size_ && is_checked(data_[cur_].id)) { ++cur_; } - return cur_id; + return internal_id; } void clear() { @@ -87,15 +88,13 @@ class SearchBuffer { cur_ = 0; } - // return candidate id for next pop() - [[nodiscard]] auto next_id() const { return data_[cur_].id; } + // Return the candidate ID for the next pop(). + [[nodiscard]] PID next_id() const { return data_[cur_].id; } - [[nodiscard]] auto has_next() const -> bool { return cur_ < size_; } + [[nodiscard]] bool has_next() const { return cur_ < size_; } void resize(size_t new_size) { - data_ = std::vector, memory::AlignedAllocator>>( - storage_size(new_size) - ); + data_ = Storage(storage_size(new_size)); capacity_ = new_size; clear(); } @@ -115,22 +114,18 @@ class SearchBuffer { } } - T top_dist() const { + [[nodiscard]] T top_dist() const { if (capacity_ == 0) { return std::numeric_limits::lowest(); } return is_full() ? data_[size_ - 1].distance : std::numeric_limits::max(); } - [[nodiscard]] auto is_full() const -> bool { return size_ == capacity_; } + [[nodiscard]] bool is_full() const { return size_ == capacity_; } - // judge if dist can be inserted into buffer - [[nodiscard]] auto is_full(T dist) const -> bool { - return capacity_ == 0 || dist > top_dist(); - } + // Return whether dist cannot be inserted into the buffer. + [[nodiscard]] bool is_full(T dist) const { return capacity_ == 0 || dist > top_dist(); } - const std::vector, memory::AlignedAllocator>>& data() { - return data_; - } + [[nodiscard]] const Storage& data() const { return data_; } }; } // namespace rabitqlib::buffer diff --git a/include/rabitqlib/index/detail/exact_nn.hpp b/include/rabitqlib/index/detail/exact_nn.hpp index 1a6f29c8..792298e3 100644 --- a/include/rabitqlib/index/detail/exact_nn.hpp +++ b/include/rabitqlib/index/detail/exact_nn.hpp @@ -38,22 +38,22 @@ inline PID exact_nn( #else constexpr int tid = 0; #endif - AnnCandidate& cur_entry = best_entries[tid]; - const T* cur_data = data + (dim * i); + AnnCandidate& best_entry = best_entries[tid]; + const T* point = data + (dim * i); - T distance = dist_func(cur_data, query, dim); - if (distance < cur_entry.distance) { - cur_entry.id = static_cast(i); - cur_entry.distance = distance; + T distance = dist_func(point, query, dim); + if (distance < best_entry.distance) { + best_entry.id = static_cast(i); + best_entry.distance = distance; } } PID nearest_neighbor = 0; T min_dist = std::numeric_limits::max(); - for (auto& candi : best_entries) { - if (candi.distance < min_dist) { - nearest_neighbor = candi.id; - min_dist = candi.distance; + for (const auto& entry : best_entries) { + if (entry.distance < min_dist) { + nearest_neighbor = entry.id; + min_dist = entry.distance; } } return nearest_neighbor; diff --git a/include/rabitqlib/index/detail/hnsw_dispatch.hpp b/include/rabitqlib/index/detail/hnsw_dispatch.hpp index dabc82af..d65a2646 100644 --- a/include/rabitqlib/index/detail/hnsw_dispatch.hpp +++ b/include/rabitqlib/index/detail/hnsw_dispatch.hpp @@ -1,7 +1,6 @@ #pragma once #include -#include #include #include @@ -10,29 +9,24 @@ namespace rabitqlib::hnsw { class HierarchicalNSW; namespace detail { -// Compatibility adapter; individual backends return ordered results only. -std::priority_queue> search_knn( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 -); - // Ordered results used by batch search and backend parity tests. std::vector> search_knn_sorted( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); std::vector> search_knn_sorted_generic( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); std::vector> search_knn_sorted_neon( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx2( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx512_core( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx512_popcnt( - HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0 ); } // namespace detail diff --git a/include/rabitqlib/index/detail/parallel_for.hpp b/include/rabitqlib/index/detail/parallel_for.hpp index 576326e2..d1f2caff 100644 --- a/include/rabitqlib/index/detail/parallel_for.hpp +++ b/include/rabitqlib/index/detail/parallel_for.hpp @@ -27,11 +27,11 @@ inline void parallel_for( num_threads = std::min(resolve_num_threads(num_threads), end - start); if (num_threads == 1) { - for (size_t id = start; id < end; id++) { + for (size_t id = start; id < end; ++id) { fn(id, 0); } } else { -#if defined(_OPENMP) && _OPENMP >= 201307 +#if RABITQ_USE_OPENMP && defined(_OPENMP) && _OPENMP >= 201307 if ((reuse_workers && omp_get_level() == 0) || omp_get_proc_bind() != omp_proc_bind_false) { // Independent query batches can reuse workers; graph insertion keeps @@ -86,12 +86,9 @@ inline void parallel_for( last_except_mutex ); last_exception = std::current_exception(); - /* - * This will work even when current is the largest value that - * size_t can fit, because fetch_add returns the previous value - * before the increment (what will result in overflow - * and produce 0 instead of current + 1). - */ + // This works even when current is the largest size_t value: + // fetch_add returns the value before the increment, which + // would otherwise overflow to 0 instead of current + 1. current = end; break; } diff --git a/include/rabitqlib/index/detail/visited_pool.hpp b/include/rabitqlib/index/detail/visited_pool.hpp index 2e3f4240..d8ea4200 100644 --- a/include/rabitqlib/index/detail/visited_pool.hpp +++ b/include/rabitqlib/index/detail/visited_pool.hpp @@ -8,39 +8,39 @@ #include "rabitqlib/index/detail/visited_set.hpp" namespace rabitqlib { -class VisitedListPool { +class VisitedSetPool { std::deque> pool_; - std::mutex poolguard_; - size_t numelements_; + std::mutex pool_mutex_; + size_t num_elements_; public: - VisitedListPool(size_t initpoolsize, size_t max_elements) { - numelements_ = max_elements; - for (size_t i = 0; i < initpoolsize; i++) { - pool_.push_front(std::make_unique(numelements_, numelements_ / 10)); + VisitedSetPool(size_t initial_pool_size, size_t num_elements) + : num_elements_(num_elements) { + for (size_t i = 0; i < initial_pool_size; ++i) { + pool_.push_front( + std::make_unique(num_elements_, num_elements_ / 10) + ); } } - std::unique_ptr get_free_vislist() { - std::unique_ptr rez; + std::unique_ptr acquire() { + std::unique_ptr visited; { - std::unique_lock lock(poolguard_); - if (pool_.size() > 0) { - rez = std::move(pool_.front()); + std::unique_lock lock(pool_mutex_); + if (!pool_.empty()) { + visited = std::move(pool_.front()); pool_.pop_front(); } else { - rez = std::make_unique(numelements_, numelements_ / 10); + visited = std::make_unique(num_elements_, num_elements_ / 10); } } - rez->clear(); - return rez; + visited->clear(); + return visited; } - void release_vis_list(std::unique_ptr vl) { - std::unique_lock lock(poolguard_); - pool_.push_front(std::move(vl)); + void release(std::unique_ptr visited) { + std::unique_lock lock(pool_mutex_); + pool_.push_front(std::move(visited)); } - - ~VisitedListPool() = default; }; } // namespace rabitqlib diff --git a/include/rabitqlib/index/detail/visited_set_epoch.hpp b/include/rabitqlib/index/detail/visited_set_epoch.hpp index 756288ed..79164cb6 100644 --- a/include/rabitqlib/index/detail/visited_set_epoch.hpp +++ b/include/rabitqlib/index/detail/visited_set_epoch.hpp @@ -23,11 +23,11 @@ class EpochBasedVisitedSet { EpochBasedVisitedSet(EpochBasedVisitedSet&&) noexcept = default; EpochBasedVisitedSet& operator=(EpochBasedVisitedSet&&) noexcept = default; - explicit EpochBasedVisitedSet(size_t num_elements, size_t = 0) { + explicit EpochBasedVisitedSet(size_t num_elements, size_t /*size_hint*/ = 0) { initialize(num_elements); } - void initialize(size_t num_elements, size_t = 0) { + void initialize(size_t num_elements, size_t /*size_hint*/ = 0) { stamp_.assign(num_elements, 0); cur_ = 0; clear(); @@ -40,8 +40,8 @@ class EpochBasedVisitedSet { } } - [[nodiscard]] bool get(PID data_id) const { return stamp_[data_id] == cur_; } + [[nodiscard]] bool get(PID internal_id) const { return stamp_[internal_id] == cur_; } - void set(PID data_id) { stamp_[data_id] = cur_; } + void set(PID internal_id) { stamp_[internal_id] = cur_; } }; } // namespace rabitqlib diff --git a/include/rabitqlib/index/detail/visited_set_fwd.hpp b/include/rabitqlib/index/detail/visited_set_fwd.hpp index 3c9addfd..4b40e7f1 100644 --- a/include/rabitqlib/index/detail/visited_set_fwd.hpp +++ b/include/rabitqlib/index/detail/visited_set_fwd.hpp @@ -1,6 +1,11 @@ #pragma once namespace rabitqlib { +// The library names its visited-set backend only here. HashBasedVisitedSet keeps +// memory sublinear in the ID space. EpochBasedVisitedSet is kept as the drop-in +// alternative: it clears in amortized O(1) and never allocates during search, at +// 2 bytes per ID. Both satisfy detail::CheckVisitedSet, and tests keep them +// behaviorally equivalent so this alias can switch backends. class HashBasedVisitedSet; using VisitedSet = HashBasedVisitedSet; } // namespace rabitqlib diff --git a/include/rabitqlib/index/detail/visited_set_hash.hpp b/include/rabitqlib/index/detail/visited_set_hash.hpp index caec2c92..1b4a678d 100644 --- a/include/rabitqlib/index/detail/visited_set_hash.hpp +++ b/include/rabitqlib/index/detail/visited_set_hash.hpp @@ -74,7 +74,7 @@ class HashBasedVisitedSet { size_t bit_size = 0; size_t bit = size; while (bit != 0) { - bit_size++; + ++bit_size; bit >>= 1; } // Limit per-visit spillover allocations while keeping the direct table sublinear. @@ -89,23 +89,23 @@ class HashBasedVisitedSet { [[nodiscard]] bool initialized() const noexcept { return !table_.empty(); } - [[nodiscard]] bool get(PID data_id) const { - PID val = this->table_[hash1(data_id)]; - if (val == data_id) { + [[nodiscard]] bool get(PID internal_id) const { + PID val = table_[hash1(internal_id)]; + if (val == internal_id) { return true; } - return (val != kPidMax && stl_hash_.find(data_id) != stl_hash_.end()); + return (val != kPidMax && stl_hash_.find(internal_id) != stl_hash_.end()); } - void set(PID data_id) { - PID& val = table_[hash1(data_id)]; - if (val == data_id) { + void set(PID internal_id) { + PID& val = table_[hash1(internal_id)]; + if (val == internal_id) { return; } if (val == kPidMax) { - val = data_id; + val = internal_id; } else { - stl_hash_.emplace(data_id); + stl_hash_.emplace(internal_id); } } }; diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 28a3f9c0..97c32896 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -27,7 +27,7 @@ template class Rotator; template class SplitSingleQuery; -class VisitedListPool; +class VisitedSetPool; namespace detail { struct DistanceEstimate; } @@ -74,7 +74,7 @@ class HierarchicalNSW { } [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; } [[nodiscard]] PID cluster_id_of(PID internal_id) const { - return get_clusterid_by_internalid(internal_id); + return raw_cluster_id(internal_id) & ~kRemovedMask; } void save(const char* filename) const; @@ -169,47 +169,37 @@ class HierarchicalNSW { // Borrows construct() input only while its workers run; cleared on success or failure. const float* construction_data_{nullptr}; - struct Candidate { - float est_dist; - PID id; - }; - // Bounded priority queue implemented as a sorted vector. class BoundedKNN { public: explicit BoundedKNN(size_t capacity) : capacity_(capacity) {} - // Insert a candidate in sorted order (ascending by est_dist). - void insert(const Candidate& cand) { - auto it = std::upper_bound( - queue_.begin(), - queue_.end(), - cand, - [](const Candidate& a, const Candidate& b) { - return a.est_dist < b.est_dist; - } - ); - queue_.insert(it, cand); - // If we exceed capacity, drop the worst candidate (largest est_dist). + // Insert a candidate in sorted order (ascending by distance). + void insert(const AnnCandidate& candidate) { + auto it = std::upper_bound(queue_.begin(), queue_.end(), candidate); + queue_.insert(it, candidate); + // If we exceed capacity, drop the worst candidate (largest distance). if (queue_.size() > capacity_) { queue_.pop_back(); } } - // Returns the worst (largest est_dist) candidate. - [[nodiscard]] const Candidate& worst() const { return queue_.back(); } + // Returns the worst (largest distance) candidate. + [[nodiscard]] const AnnCandidate& worst() const { return queue_.back(); } [[nodiscard]] size_t size() const { return queue_.size(); } - [[nodiscard]] const std::vector& candidates() const { return queue_; } + [[nodiscard]] const std::vector>& candidates() const { + return queue_; + } private: size_t capacity_; - // Sorted in ascending order by est_dist so that the worst is at the back. - std::vector queue_; + // Sorted in ascending order by distance so that the worst is at the back. + std::vector> queue_; }; - friend struct HnswPruningTestAccess; + friend struct HnswTestAccess; friend std::vector> detail::search_knn_sorted_neon( HierarchicalNSW&, const float*, size_t, size_t ); @@ -231,8 +221,10 @@ class HierarchicalNSW { static constexpr PID kRemovedMask = 0x80000000U; // Used when a search passes ef_search == 0. static constexpr size_t kDefaultEfSearch = 10; + // The constructor clamps M to this value and load rejects files above it. + static constexpr size_t kMaxM = 10000; size_t max_elements_{0}; - mutable std::atomic num_stored_points_{0}; + std::atomic num_stored_points_{0}; size_t num_removed_{0}; // Recovered from stored cluster-ID markers when loading. size_t level0_record_bytes_{0}; // One complete level-0 record. size_t upper_layer_links_bytes_{0}; // One node's adjacency list at one upper level. @@ -242,16 +234,18 @@ class HierarchicalNSW { size_t ef_construction_{0}; MetricType metric_type_{METRIC_L2}; - double mult_{0.0}; - int max_level_{0}; + double level_multiplier_{0.0}; + // max_level_ == -1 and entry_point_id_ == kPidMax mark an empty graph. + int max_level_{-1}; - std::mutex global_; + // Serializes entry-point and maximum-level updates during insertion. + std::mutex entry_point_mutex_; std::vector link_list_locks_; - PID entry_point_id_{0}; + PID entry_point_id_{kPidMax}; - size_t level0_links_bytes_{ - 0}; // One node's level-0 adjacency list, including its header. + // One node's level-0 adjacency list, including its header. + size_t level0_links_bytes_{0}; size_t bin_data_offset_{0}, ex_data_offset_{0}, label_offset_{0}; size_t bin_data_bytes_{0}, ex_data_bytes_{0}; @@ -279,12 +273,12 @@ class HierarchicalNSW { AlignedStorage centroids_memory_; - mutable std::mutex label_lookup_lock_; + std::mutex label_lookup_lock_; std::unordered_map label_lookup_; std::default_random_engine level_generator_; - std::unique_ptr visited_list_pool_{nullptr}; + std::unique_ptr visited_set_pool_{nullptr}; float (*ip_func_)(const float*, const uint8_t*, size_t){nullptr}; @@ -292,24 +286,24 @@ class HierarchicalNSW { quant::RabitqConfig query_config_; - using EstimateRecord = rabitqlib::detail::DistanceEstimate; + using DistanceEstimate = rabitqlib::detail::DistanceEstimate; float (*raw_dist_func_)(const float* __restrict__, const float* __restrict__, size_t){ nullptr}; - PID get_external_label(PID internal_id) const { - PID return_label; - memcpy( - &return_label, + [[nodiscard]] PID external_label(PID internal_id) const { + PID label; + std::memcpy( + &label, (data_level0_memory_.get() + (internal_id * level0_record_bytes_) + label_offset_), sizeof(PID) ); - return return_label; + return label; } void set_external_label(PID internal_id, PID label) { - memcpy( + std::memcpy( (data_level0_memory_.get() + (internal_id * level0_record_bytes_) + label_offset_), &label, @@ -317,45 +311,41 @@ class HierarchicalNSW { ); } - char* get_bindata_by_internalid(PID internal_id) { + [[nodiscard]] char* bin_data(PID internal_id) { return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + bin_data_offset_; } - const char* get_bindata_by_internalid(PID internal_id) const { + [[nodiscard]] const char* bin_data(PID internal_id) const { return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + bin_data_offset_; } - char* get_exdata_by_internalid(PID internal_id) { + [[nodiscard]] char* ex_data(PID internal_id) { return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + ex_data_offset_; } - const char* get_exdata_by_internalid(PID internal_id) const { + [[nodiscard]] const char* ex_data(PID internal_id) const { return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + ex_data_offset_; } // The stored cluster ID with the removal mark, as save and load see it. - PID get_raw_clusterid(PID internal_id) const { + [[nodiscard]] PID raw_cluster_id(PID internal_id) const { return *(reinterpret_cast( data_level0_memory_.get() + (internal_id * level0_record_bytes_) + level0_links_bytes_ )); } - PID get_clusterid_by_internalid(PID internal_id) const { - return get_raw_clusterid(internal_id) & ~kRemovedMask; - } - - bool is_removed(PID internal_id) const { - return (get_raw_clusterid(internal_id) & kRemovedMask) != 0; + [[nodiscard]] bool is_removed(PID internal_id) const { + return (raw_cluster_id(internal_id) & kRemovedMask) != 0; } // Store a cluster ID, including any removal mark. - void set_raw_clusterid(PID internal_id, PID cluster_id) { - memcpy( + void set_raw_cluster_id(PID internal_id, PID cluster_id) { + std::memcpy( data_level0_memory_.get() + (internal_id * level0_record_bytes_) + level0_links_bytes_, &cluster_id, @@ -379,63 +369,64 @@ class HierarchicalNSW { upper_layer_links_bytes_}; } - int get_random_level(double reverse_size); + int random_level(double multiplier); // Write the index file at exactly this path. void write_file(const char* filename) const; - PID* get_linklist(PID internal_id, int level) { + // A link list is a count header followed by neighbor IDs. + [[nodiscard]] PID* upper_layer_links(PID internal_id, int level) { return reinterpret_cast( link_lists_[internal_id].get() + ((level - 1) * upper_layer_links_bytes_) ); } - const PID* get_linklist(PID internal_id, int level) const { + [[nodiscard]] const PID* upper_layer_links(PID internal_id, int level) const { return reinterpret_cast( link_lists_[internal_id].get() + ((level - 1) * upper_layer_links_bytes_) ); } - PID* get_linklist0(PID internal_id) { + [[nodiscard]] PID* level0_links(PID internal_id) { return reinterpret_cast( data_level0_memory_.get() + (internal_id * level0_record_bytes_) ); } - const PID* get_linklist0(PID internal_id) const { + [[nodiscard]] const PID* level0_links(PID internal_id) const { return reinterpret_cast( data_level0_memory_.get() + (internal_id * level0_record_bytes_) ); } - static unsigned short int get_list_count(const PID* ptr) { - return *(reinterpret_cast(ptr)); + [[nodiscard]] static uint16_t link_count(const PID* links) { + return *(reinterpret_cast(links)); } - static void set_list_count(PID* ptr, unsigned short int size) { - *(reinterpret_cast(ptr)) = size; + static void set_link_count(PID* links, uint16_t count) { + *(reinterpret_cast(links)) = count; } // ANN Search template - void get_bin_est_direct( - const std::vector&, - const SplitSingleQuery&, - PID, - HierarchicalNSW::EstimateRecord& + void estimate_bin_direct( + const std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, + PID candidate_id, + DistanceEstimate& estimate ) const; template - void get_full_est_direct( - const std::vector&, - const SplitSingleQuery&, - PID, - HierarchicalNSW::EstimateRecord& + void estimate_full_direct( + const std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, + PID candidate_id, + DistanceEstimate& estimate ) const; template std::vector> search_knn_sorted_direct( - const float*, size_t, size_t + const float* rotated_query, size_t k, size_t ef_search ) const; template @@ -448,37 +439,52 @@ class HierarchicalNSW { BoundedKNN& nearest_candidates ) const; - // Raw-vector distance used by construct; add scores links with get_quant_dist. - float get_data_dist(PID obj1, PID obj2) const { - PID label1 = get_external_label(obj1); - PID label2 = get_external_label(obj2); + // Raw-vector distance used by construct; add scores links with quantized_distance. + [[nodiscard]] float raw_distance(PID target, PID query) const { + PID target_label = external_label(target); + PID query_label = external_label(query); return raw_dist_func_( - construction_data_ + (label1 * dim_), construction_data_ + (label2 * dim_), dim_ + construction_data_ + (target_label * dim_), + construction_data_ + (query_label * dim_), + dim_ ); } - void add_point(PID, PID, const quant::RabitqConfig&); + void add_point(PID label, PID cluster_id, const quant::RabitqConfig& config); // Insertion steps shared by raw and quantized scoring. distance(id) scores an // existing node against the node being inserted; prefetch(id) may warm its data. template - PID greedy_descend(PID, int from_level, int to_level, Distance); + PID greedy_descend(PID current_id, int from_level, int to_level, Distance distance); template - maxheap> search_layer(PID, int, Distance, Prefetch); + maxheap> search_layer( + PID entry_point_id, int level, Distance distance, Prefetch prefetch + ); - PID mutually_connect_new_element(PID, maxheap>&, int); + PID mutually_connect_new_element( + PID new_node_id, maxheap>& top_candidates, int level + ); template PID mutually_connect( - PID, maxheap>&, int, Distance, ReverseEdgeDistance + PID new_node_id, + maxheap>& top_candidates, + int level, + Distance distance, + ReverseEdgeDistance reverse_edge_distance ); // The same steps once the original data is gone, implemented in hnsw.cpp. - void reconstruct_rotated(PID, float*) const; - float get_quant_dist(PID target, PID query, detail::QuantQuery&) const; - void add_point_quant(PID, int, LinkStorage&, detail::QuantQuery&); + void reconstruct_rotated(PID internal_id, float* rotated) const; + float quantized_distance(PID target, PID query, detail::QuantQuery& cached) const; + void add_point_quant( + PID new_node_id, int node_level, LinkStorage& link_list, detail::QuantQuery& query + ); PID mutually_connect_quant( - PID, maxheap>&, int, detail::QuantQuery& + PID new_node_id, + maxheap>& top_candidates, + int level, + detail::QuantQuery& query ); }; diff --git a/include/rabitqlib/index/ivf/cluster.hpp b/include/rabitqlib/index/ivf/cluster.hpp index f1268a4b..5b2b7514 100644 --- a/include/rabitqlib/index/ivf/cluster.hpp +++ b/include/rabitqlib/index/ivf/cluster.hpp @@ -11,10 +11,10 @@ class Cluster { private: friend class IVF; - size_t num_stored_points_; // Num of vectors in this cluster + size_t num_stored_points_; // Number of points in this cluster char* batch_data_ = nullptr; // RaBitQ code and factors char* ex_data_ = nullptr; // Extra-bit codes/factors or original float32 vectors - PID* ids_ = nullptr; // PID of vectors + PID* ids_ = nullptr; // PID of each point // The owning index refreshes only views whose storage or count changed. void reset( @@ -27,18 +27,18 @@ class Cluster { } public: - explicit Cluster(size_t, char*, char*, PID*); + explicit Cluster(size_t num_stored_points, char* batch_data, char* ex_data, PID* ids); Cluster(const Cluster&) = default; Cluster(Cluster&&) noexcept = default; Cluster& operator=(const Cluster&) = delete; Cluster& operator=(Cluster&&) = delete; ~Cluster() = default; - [[nodiscard]] char* batch_data() const { return this->batch_data_; } + [[nodiscard]] char* batch_data() const { return batch_data_; } [[nodiscard]] char* ex_data() const { return ex_data_; } - [[nodiscard]] PID* ids() const { return this->ids_; } + [[nodiscard]] PID* ids() const { return ids_; } [[nodiscard]] size_t num() const { return num_stored_points_; } }; diff --git a/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp index 179e83bd..32bde339 100644 --- a/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp +++ b/include/rabitqlib/index/ivf/detail/hnsw_sidecar.hpp @@ -16,6 +16,10 @@ namespace rabitqlib::ivf::detail { +// The HNSW initializer stores its graph at the index path plus this suffix. Saving, +// loading and atomic replacement must all agree on it. +inline constexpr char kHnswSidecarSuffix[] = ".hnsw"; + // hnswlib's native sidecar has no version tag. Check the existing writer's // header and record lengths before its loader uses offsets or allocates links. // As with the main index file, callers must not mutate the file during loading. @@ -33,16 +37,16 @@ inline void validate_hnsw_sidecar_file( throw std::runtime_error("Truncated HNSW centroid sidecar"); } }; - size_t offset_level0 = 0, capacity = 0, count = 0, record_bytes = 0; + size_t level0_offset = 0, capacity = 0, count = 0, level0_record_bytes = 0; size_t label_offset = 0, data_offset = 0; int max_level = 0; hnswlib::tableint entry = 0; size_t max_m = 0, max_m0 = 0, m = 0, ef_construction = 0; double mult = 0; - read(offset_level0); + read(level0_offset); read(capacity); read(count); - read(record_bytes); + read(level0_record_bytes); read(label_offset); read(data_offset); read(max_level); @@ -56,29 +60,32 @@ inline void validate_hnsw_sidecar_file( count > std::numeric_limits::max()) { throw std::runtime_error("HNSW centroid count does not match IVF cluster count"); } - if (m < 2 || m > 10000 || max_m != m || max_m0 != 2 * m || !std::isfinite(mult) || + constexpr size_t kMaxM = 10000; // hnswlib caps M at this value. + if (m < 2 || m > kMaxM || max_m != m || max_m0 != 2 * m || !std::isfinite(mult) || mult <= 0 || ef_construction < m) { throw std::runtime_error("Invalid HNSW centroid graph parameters"); } - const size_t upper_stride = + const size_t upper_layer_links_bytes = sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); - const size_t base_bytes = + const size_t level0_links_bytes = sizeof(hnswlib::linklistsizeint) + max_m0 * sizeof(hnswlib::tableint); - if (offset_level0 != 0 || data_offset != base_bytes || data_offset > label_offset || - label_offset > record_bytes || label_offset - data_offset != vector_bytes || - record_bytes - label_offset != sizeof(hnswlib::labeltype)) { + if (level0_offset != 0 || data_offset != level0_links_bytes || + data_offset > label_offset || label_offset > level0_record_bytes || + label_offset - data_offset != vector_bytes || + level0_record_bytes - label_offset != sizeof(hnswlib::labeltype)) { throw std::runtime_error("HNSW centroid layout does not match IVF dimension"); } if (max_level < 0 || entry >= count) { throw std::runtime_error("Invalid HNSW centroid entry point or maximum level"); } auto remaining = end - input.tellg(); - if (remaining < 0 || count > static_cast(remaining) / record_bytes) { + if (remaining < 0 || count > static_cast(remaining) / level0_record_bytes) { throw std::runtime_error("Truncated HNSW centroid base records"); } - const auto base_records_bytes = static_cast(count * record_bytes); - input.seekg(base_records_bytes, std::ios::cur); - remaining -= base_records_bytes; + const auto level0_records_bytes = + static_cast(count * level0_record_bytes); + input.seekg(level0_records_bytes, std::ios::cur); + remaining -= level0_records_bytes; size_t highest_level = 0, entry_level = 0; for (size_t i = 0; i < count; ++i) { unsigned int link_bytes = 0; // Native type written by hnswlib::saveIndex. @@ -87,11 +94,11 @@ inline void validate_hnsw_sidecar_file( } read(link_bytes); remaining -= static_cast(sizeof(link_bytes)); - if (link_bytes % upper_stride != 0 || + if (link_bytes % upper_layer_links_bytes != 0 || link_bytes > static_cast(remaining)) { throw std::runtime_error("Invalid HNSW centroid upper-link length"); } - const size_t levels = link_bytes / upper_stride; + const size_t levels = link_bytes / upper_layer_links_bytes; highest_level = std::max(highest_level, levels); if (i == entry) { entry_level = levels; diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 69d6825a..d1ecdc7d 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -23,13 +23,17 @@ #include "rabitqlib/utils/space.hpp" namespace rabitqlib::ivf { -inline InitializerType resolve_initializer_type(InitializerType type, size_t clusters) { +inline InitializerType resolve_initializer_type(InitializerType type, size_t num_clusters) { + // Auto routes through the cheapest initializer that stays fast at this cluster count. + constexpr size_t kFlatMaxClusters = 5000; + constexpr size_t kFlatRaBitQMaxClusters = 60000; switch (type) { case InitializerType::Auto: - if (clusters < 5000) { + if (num_clusters < kFlatMaxClusters) { return InitializerType::Flat; } - return clusters < 60000 ? InitializerType::FlatRaBitQ : InitializerType::HNSW; + return num_clusters < kFlatRaBitQMaxClusters ? InitializerType::FlatRaBitQ + : InitializerType::HNSW; case InitializerType::Flat: case InitializerType::FlatRaBitQ: case InitializerType::HNSW: @@ -50,17 +54,17 @@ class Initializer { explicit Initializer(size_t dim, size_t num_clusters) : dim_(dim), num_cluster_(num_clusters) {} virtual ~Initializer() = 0; - [[nodiscard]] virtual const float* centroid(PID) const = 0; - virtual void add_vectors(const float*, size_t) = 0; + [[nodiscard]] virtual const float* centroid(PID id) const = 0; + virtual void add_vectors(const float* centroids, size_t num_threads) = 0; // Resize the output to min(nprobe, number of clusters) and fill it with the // nearest centroids in no particular order; the caller need not presize it. virtual void centroids_distances( - const float*, size_t, std::vector>& + const float* query, size_t nprobe, std::vector>& candidates ) const = 0; - virtual void load(std::ifstream&, const char*) = 0; - virtual void save(std::ofstream&, const char*) const = 0; + virtual void load(std::ifstream& input, const char* filename) = 0; + virtual void save(std::ofstream& output, const char* filename) const = 0; }; -inline Initializer::~Initializer() {} +inline Initializer::~Initializer() = default; class FlatInitializer : public Initializer { private: @@ -69,7 +73,7 @@ class FlatInitializer : public Initializer { public: explicit FlatInitializer( - size_t dim, size_t num_clusters, MetricType metric_type = MetricType::METRIC_L2 + size_t dim, size_t num_clusters, MetricType metric_type = METRIC_L2 ) : Initializer(dim, num_clusters) , centroids_(num_cluster_ * dim_) @@ -83,14 +87,14 @@ class FlatInitializer : public Initializer { return ¢roids_[id * dim_]; } - void add_vectors(const float* centroids, size_t) override { + void add_vectors(const float* centroids, size_t /*num_threads*/) override { std::memcpy(centroids_.data(), centroids, sizeof(float) * num_cluster_ * dim_); } void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override { - std::vector> centroid_candidates(this->num_cluster_); + std::vector> centroid_candidates(num_cluster_); for (PID i = 0; i < num_cluster_; ++i) { centroid_candidates[i].id = i; centroid_candidates[i].distance = @@ -110,14 +114,14 @@ class FlatInitializer : public Initializer { } // The flat initializer stores its centroids inline in the index file stream. - void save(std::ofstream& output, const char*) const override { + void save(std::ofstream& output, const char* /*filename*/) const override { output.write( reinterpret_cast(centroids_.data()), static_cast(sizeof(float) * dim_ * num_cluster_) ); } - void load(std::ifstream& input, const char*) override { + void load(std::ifstream& input, const char* /*filename*/) override { input.read( reinterpret_cast(centroids_.data()), static_cast(sizeof(float) * dim_ * num_cluster_) @@ -136,15 +140,15 @@ class FlatRaBitQInitializer : public Initializer { public: explicit FlatRaBitQInitializer( - size_t dim, size_t num_clusters, MetricType metric = METRIC_L2 + size_t dim, size_t num_clusters, MetricType metric_type = METRIC_L2 ); [[nodiscard]] const float* centroid(PID id) const override; void add_vectors(const float* centroids, size_t num_threads) override; void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override; - void save(std::ofstream& output, const char*) const override; - void load(std::ifstream& input, const char*) override; + void save(std::ofstream& output, const char* filename) const override; + void load(std::ifstream& input, const char* filename) override; }; // Keep centroid routing on the same runtime-selected distance kernels as flat IVF. @@ -190,8 +194,11 @@ class CentroidIPSpace : public hnswlib::SpaceInterface { class HNSWInitializer : public Initializer { private: - int M_ = 16; - int ef_construction_ = 400; + // Fixed hnswlib graph parameters for centroid routing. + static constexpr size_t kM = 16; + static constexpr size_t kEfConstruction = 400; + static constexpr size_t kMinSearchEf = 768; + std::unique_ptr> alg_hnsw_; CentroidL2Space l2_space_; CentroidIPSpace ip_space_; @@ -221,7 +228,7 @@ class HNSWInitializer : public Initializer { public: explicit HNSWInitializer( - size_t dim, size_t num_clusters, MetricType metric_type = MetricType::METRIC_L2 + size_t dim, size_t num_clusters, MetricType metric_type = METRIC_L2 ) : Initializer(dim, num_clusters) , l2_space_(dim) @@ -234,15 +241,13 @@ class HNSWInitializer : public Initializer { , metric_type_(metric_type) { validate_metric_type(metric_type_); alg_hnsw_ = std::make_unique>( - space_, num_cluster_, M_, ef_construction_ + space_, num_cluster_, kM, kEfConstruction ); } void add_vectors(const float* centroids, size_t num_threads) override { - size_t start = 0; - size_t rows = num_cluster_; - rabitqlib::index_detail::parallel_for( - start, rows, num_threads, [&](size_t row, size_t /*thread_id*/) { + index_detail::parallel_for( + 0, num_cluster_, num_threads, [&](size_t row, size_t /*thread*/) { alg_hnsw_->addPoint(centroids + (row * dim_), row); } ); @@ -265,7 +270,7 @@ class HNSWInitializer : public Initializer { } std::priority_queue> result = alg_hnsw_->searchKnnWithEf( - query, num_probes, std::max(size_t{768}, 2 * num_probes) + query, num_probes, std::max(kMinSearchEf, 2 * num_probes) ); if (result.size() != num_probes) { throw std::runtime_error("HNSW centroid graph returned too few candidates"); @@ -281,24 +286,24 @@ class HNSWInitializer : public Initializer { } // The HNSW initializer stores its graph in a separate hnswlib sidecar file. - void save(std::ofstream&, const char* filename) const override { - std::string hnsw(filename); - hnsw += ".hnsw"; - alg_hnsw_->saveIndex(hnsw); + void save(std::ofstream& /*output*/, const char* filename) const override { + alg_hnsw_->saveIndex(std::string(filename) + detail::kHnswSidecarSuffix); } - void load(std::ifstream&, const char* filename) override { - std::string hnsw(filename); - hnsw += ".hnsw"; - detail::validate_hnsw_sidecar_file(hnsw, num_cluster_, space_->get_data_size()); - alg_hnsw_->loadIndex(hnsw, space_, num_cluster_); + void load(std::ifstream& /*input*/, const char* filename) override { + const std::string sidecar = std::string(filename) + detail::kHnswSidecarSuffix; + detail::validate_hnsw_sidecar_file(sidecar, num_cluster_, space_->get_data_size()); + alg_hnsw_->loadIndex(sidecar, space_, num_cluster_); const auto& graph = *alg_hnsw_; + // validate_hnsw_sidecar_file already checked these header fields, and loadIndex + // reads them again from a second open of the file. Recheck the loaded values so + // a sidecar replaced in between cannot give centroid() an out-of-record offset. if (graph.offsetData_ > graph.label_offset_ || graph.label_offset_ > graph.size_data_per_element_ || graph.label_offset_ - graph.offsetData_ != space_->get_data_size() || graph.size_data_per_element_ - graph.label_offset_ != sizeof(hnswlib::labeltype)) { - throw std::runtime_error("HNSW centroid layout does not match IVF dimension"); + throw std::runtime_error("HNSW centroid sidecar changed while loading"); } detail::validate_hnsw_centroid_graph(graph); rebuild_label_lookup(); diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index eea36ae7..be504632 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -55,12 +55,13 @@ class IVF { return reinterpret_cast(batch_storage.data()); } - [[nodiscard]] char* ex_data() { + // Raw vectors with raw reranking, extra-bit codes otherwise. + [[nodiscard]] char* rerank_data() { return raw_reranking_ ? reinterpret_cast(raw_storage.data()) : reinterpret_cast(ex_storage.data()); } - [[nodiscard]] const char* ex_data() const { + [[nodiscard]] const char* rerank_data() const { return raw_reranking_ ? reinterpret_cast(raw_storage.data()) : reinterpret_cast(ex_storage.data()); } @@ -87,7 +88,7 @@ class IVF { RotatorType rotator_type_ = RotatorType::FhtKacRotator; std::unique_ptr> rotator_; std::vector clusters_; - MetricType metric_type_ = rabitqlib::METRIC_L2; + MetricType metric_type_ = METRIC_L2; float (*ip_func_)(const float*, const uint8_t*, size_t) = nullptr; static constexpr uint64_t kRawFormatMagic = 0x3157415251464252ULL; @@ -96,26 +97,26 @@ class IVF { static constexpr uint32_t kFormatVersion = 2; void quantize_cluster( - Cluster&, - const std::vector&, - const float*, - const float*, - float*, - float*, - const quant::RabitqConfig& + Cluster& cluster, + const std::vector& point_ids, + const float* data, + const float* centroid, + float* rotated_centroid, + float* rotated_data, + const quant::RabitqConfig& config ); void extend_cluster( - const Cluster&, - const std::vector&, - const float*, - const float*, - PID, - char*, - char*, - PID*, - char*, - const quant::RabitqConfig& + const Cluster& old_cluster, + const std::vector& rows, + const float* data, + const float* rotated_centroid, + PID first_pid, + char* batch_dst, + char* rerank_dst, + PID* ids_dst, + char* tail_dst, + const quant::RabitqConfig& config ) const; // True if `num_points` points cannot be addressed by the storage size computations @@ -125,12 +126,18 @@ class IVF { return sizeof(PID) * cluster_size; } - // get num of bytes used for 1-bit code and corresponding factors + // Bytes of one FastScan batch of 1-bit codes and factors + [[nodiscard]] inline size_t batch_bytes() const; + + // Bytes of the whole batches holding `cluster_size` points' 1-bit codes and factors [[nodiscard]] size_t batch_data_bytes(size_t cluster_size) const; + // True if points carry extra-bit codes or raw vectors for reranking + [[nodiscard]] bool has_rerank_data() const { return raw_reranking_ || ex_bits_ > 0; } + [[nodiscard]] inline size_t rerank_vector_bytes() const; - [[nodiscard]] inline size_t ex_data_bytes(size_t cluster_size) const; + [[nodiscard]] inline size_t rerank_data_bytes(size_t cluster_size) const; // Size every buffer of a cluster's storage for `num_points` points void resize_storage(ClusterStorage& cs, size_t num_points) const; @@ -139,34 +146,38 @@ class IVF { [[nodiscard]] bool has_capacity(const ClusterStorage& cs, size_t num_points) const; - void allocate_memory(const std::vector&); + // Reset the index and allocate a new initializer and per-cluster storage + void allocate_storage(const std::vector& cluster_sizes); - void init_clusters(const std::vector&); + // Create the non-owning cluster views over the allocated storage + void create_cluster_views(const std::vector& cluster_sizes); - void free_memory(); + void free_storage(); void swap(IVF& other) noexcept; inline void search_cluster( - const Cluster&, - const SplitBatchQuery&, - buffer::SearchBuffer&, - bool, - const float* + const Cluster& cluster, + const SplitBatchQuery& query_wrapper, + buffer::SearchBuffer& result_pool, + bool use_hacc, + const float* query ) const; inline void scan_one_batch( const char* batch_data, - const char* ex_data, + const char* rerank_data, const PID* ids, const SplitBatchQuery& query_wrapper, buffer::SearchBuffer& result_pool, - size_t num_points, - bool, - const float* + size_t batch_size, + bool use_hacc, + const float* query ) const; - inline void validate_search(const float*, size_t, size_t, const PID*) const; + inline void validate_search( + const float* query, size_t k, size_t nprobe, const PID* result_ids + ) const; // Write the index file, and any initializer sidecar, at exactly this path. void write_file(const char* filename) const; @@ -180,9 +191,9 @@ class IVF { size_t dim, size_t num_clusters, size_t total_bits, - MetricType metric_type = rabitqlib::METRIC_L2, + MetricType metric_type = METRIC_L2, RotatorType rotator_type = RotatorType::FhtKacRotator, - InitializerType initializer = InitializerType::Auto + InitializerType initializer_type = InitializerType::Auto ); ~IVF(); @@ -200,9 +211,19 @@ class IVF { [[nodiscard]] RotatorType rotator_type() const { return rotator_type_; } [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; } - // fast_quantization selects fixed-scale extra-bit quantization instead of exact scale - // search, trading quantization accuracy for speed. One-bit and raw storage are - // unaffected. + /** + * @brief Quantize every point into its cluster and build the centroid initializer. + * + * Replaces any existing contents of the index. + * + * @param data Data points (max_elements() * dim) + * @param centroids Centroid vectors (num_clusters() * dim) + * @param cluster_ids Cluster ID of each data point + * @param fast_quantization Select fixed-scale extra-bit quantization instead of exact + * scale search, trading quantization accuracy for speed; one-bit and raw storage are + * unaffected + * @param num_threads Threads used for quantization and initializer construction + */ void construct( const float* data, const float* centroids, @@ -317,9 +338,9 @@ class IVF { size_t num_threads = 1 ) const; - [[nodiscard]] size_t padded_dim() const { return this->padded_dim_; } + [[nodiscard]] size_t padded_dim() const { return padded_dim_; } - [[nodiscard]] size_t num_clusters() const { return this->num_clusters_; } + [[nodiscard]] size_t num_clusters() const { return num_clusters_; } }; } // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp index 2d7eeff5..bc085d4e 100644 --- a/include/rabitqlib/index/symqg/detail/pipnn.hpp +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -17,9 +17,9 @@ struct InitialGraph { InitialGraph build_initial_graph( const float* data, - size_t count, + size_t num_points, size_t dim, - size_t degree, + size_t max_degree, MetricType metric_type = METRIC_L2, size_t num_threads = std::numeric_limits::max() ); diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index f0023d20..e9c591cb 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -70,14 +70,14 @@ template <> class QuantizedGraph { friend class QGBuilder; friend void detail::search_with_workspace( - QuantizedGraph&, - const float*, - uint32_t, - uint32_t*, - float*, - detail::SearchScratchView, - PID, - const float* + QuantizedGraph& graph, + const float* query, + uint32_t k, + uint32_t* result_ids, + float* distances, + detail::SearchScratchView scratch, + PID hint, + const float* k1xsumq ); friend struct QGConstructionTestAccess; @@ -88,7 +88,7 @@ class QuantizedGraph { size_t padded_dim_ = 0; float (*raw_dist_func_)(const float*, const float*, size_t) = nullptr; PID entry_point_id_ = 0; - MetricType metric_type_ = MetricType::METRIC_L2; + MetricType metric_type_ = METRIC_L2; RotatorType rotator_type_ = RotatorType::FhtKacRotator; size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors std::vector centroid_; // rotated global centroid for qg-quant @@ -130,96 +130,114 @@ class QuantizedGraph { // Write the graph file at exactly this path. void write_file(const char* filename) const; - void copy_vectors(const float*, size_t); + void copy_vectors(const float* data, size_t num_workers); void set_quantization_centroid(const float* centroid); - [[nodiscard]] char* get_row_data(PID data_id) { - return reinterpret_cast(get_vector(data_id)); + [[nodiscard]] char* get_row_data(PID internal_id) { + return reinterpret_cast(get_vector(internal_id)); } - [[nodiscard]] const char* get_row_data(PID data_id) const { - return reinterpret_cast(get_vector(data_id)); + [[nodiscard]] const char* get_row_data(PID internal_id) const { + return reinterpret_cast(get_vector(internal_id)); } - [[nodiscard]] float* get_vector(PID data_id) { - return data_.data() + ((record_bytes_ / sizeof(float)) * data_id); + [[nodiscard]] float* get_vector(PID internal_id) { + return data_.data() + ((record_bytes_ / sizeof(float)) * internal_id); } - [[nodiscard]] const float* get_vector(PID data_id) const { - return data_.data() + ((record_bytes_ / sizeof(float)) * data_id); + [[nodiscard]] const float* get_vector(PID internal_id) const { + return data_.data() + ((record_bytes_ / sizeof(float)) * internal_id); } - [[nodiscard]] char* get_quantized_vector(PID data_id) { return get_row_data(data_id); } + [[nodiscard]] char* get_quantized_vector(PID internal_id) { + return get_row_data(internal_id); + } - [[nodiscard]] const char* get_quantized_vector(PID data_id) const { - return get_row_data(data_id); + [[nodiscard]] const char* get_quantized_vector(PID internal_id) const { + return get_row_data(internal_id); } void validate_search( - const float*, uint32_t, const uint32_t*, const float*, size_t + const float* query, + uint32_t k, + const uint32_t* result_ids, + const float* distances, + size_t ef_search ) const; - void search_impl(const float*, uint32_t, uint32_t*, float*, size_t); + void search_impl( + const float* query, + uint32_t k, + uint32_t* result_ids, + float* distances, + size_t ef_search + ); void search_with_scratch_impl( - const float*, - uint32_t, - uint32_t*, - float*, - detail::SearchScratchView, - PID, - const float* + const float* query, + uint32_t k, + uint32_t* result_ids, + float* distances, + detail::SearchScratchView scratch, + PID hint, + const float* k1xsumq ); - float point_distance(const float*, const QuantizedQuery*, PID) const; + float point_distance( + const float* raw_query, const QuantizedQuery* quantized_query, PID internal_id + ) const; - float quantized_distance(const QuantizedQuery&, PID) const; + float quantized_distance(const QuantizedQuery& query, PID internal_id) const; - void reconstruct_quantized_vector(PID, float*) const; + void reconstruct_quantized_vector(PID internal_id, float* reconstructed) const; - [[nodiscard]] char* get_batch_data(PID data_id) { - return get_row_data(data_id) + batch_data_offset_; + [[nodiscard]] char* get_batch_data(PID internal_id) { + return get_row_data(internal_id) + batch_data_offset_; } - [[nodiscard]] const char* get_batch_data(PID data_id) const { - return get_row_data(data_id) + batch_data_offset_; + [[nodiscard]] const char* get_batch_data(PID internal_id) const { + return get_row_data(internal_id) + batch_data_offset_; } - [[nodiscard]] rabitqlib::detail::PackedArrayView get_neighbors(PID data_id) { + [[nodiscard]] rabitqlib::detail::PackedArrayView get_neighbors(PID internal_id) { return rabitqlib::detail::PackedArrayView( - get_row_data(data_id) + neighbor_offset_ + get_row_data(internal_id) + neighbor_offset_ ); } [[nodiscard]] rabitqlib::detail::ConstPackedArrayView get_neighbors( - PID data_id + PID internal_id ) const { return rabitqlib::detail::ConstPackedArrayView( - get_row_data(data_id) + neighbor_offset_ + get_row_data(internal_id) + neighbor_offset_ ); } - // Inputs are in this graph's rotated, padded domain. The centroid is the - // source vertex; data contains one vector per neighbor, in ID order. + // Inputs are in this graph's rotated, padded domain. rotated_source is the + // source vertex, used as the quantization centroid; rotated_data contains one + // vector per neighbor, in ID order. void encode_neighbors( - PID, - const std::vector&, + PID internal_id, + const std::vector& new_neighbors, const float* rotated_data, - const float* rotated_centroid + const float* rotated_source ); void update_results( - buffer::SearchBuffer&, VisitedSet&, const float*, const QuantizedQuery* + buffer::SearchBuffer& result_pool, + VisitedSet& visited, + const float* query, + const QuantizedQuery* quantized_query ); void scan_neighbors( - const BatchQuery&, - PID, - float*, - buffer::SearchBuffer&, - VisitedSet&, - size_t + const BatchQuery& query_wrapper, + PID internal_id, + float* estimated_distances, + buffer::SearchBuffer& search_pool, + VisitedSet& visited, + size_t num_neighbors ) const; public: @@ -265,11 +283,11 @@ class QuantizedGraph { void set_ep(PID entry); - void save(const char*) const; + void save(const char* filename) const; - void load(const char*); + void load(const char* filename); - void set_ef(size_t); + void set_ef(size_t ef_search); void search( const float* __restrict__ query, @@ -314,11 +332,11 @@ class QuantizedGraph { * Reuses scratch and can use a previous result as a warm-start hint. * BatchQuery may allocate LUT storage on its first use or when its shape changes; * this is not the stable core v1 zero-allocation interface. Search and result - * buffers must be sized for ef_search and k respectively. rotated_query_scratch and - * lut_float_scratch need D and 4 * D floats, where D is the graph's stored padded + * buffers must be sized for ef_search and k respectively. rotated_query and + * lookup_table need D and 4 * D floats, where D is the graph's stored padded * dimension; this advanced caller must already know D. dimension() is the original - * input dimension. est_dist_scratch needs degree_bound() floats. The visited set must - * be initialized before the first search; an uninitialized set raises + * input dimension. estimated_distances needs degree_bound() floats. The visited set + * must be initialized before the first search; an uninitialized set raises * std::invalid_argument. Concurrent searches require separate scratch * buffers, including visited sets. * @@ -330,9 +348,9 @@ class QuantizedGraph { uint32_t k, uint32_t* __restrict__ result_ids, float* __restrict__ distances, - float* __restrict__ rotated_query_scratch, - float* __restrict__ est_dist_scratch, - float* __restrict__ lut_float_scratch, + float* __restrict__ rotated_query, + float* __restrict__ estimated_distances, + float* __restrict__ lookup_table, BatchQuery& batch_query, buffer::SearchBuffer& search_pool, buffer::SearchBuffer& result_pool, @@ -341,7 +359,7 @@ class QuantizedGraph { const float* k1xsumq = nullptr ); - [[nodiscard]] size_t padded_dim() const { return this->padded_dim_; } + [[nodiscard]] size_t padded_dim() const { return padded_dim_; } }; // Preserve C++17 deduction for callers that omit the float template argument. diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 0b0f9885..e2bf47e7 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -12,7 +12,7 @@ #include "rabitqlib/index/types.hpp" namespace rabitqlib::symqg { -constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar +inline constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar using CandidateList = std::vector>; enum class QGInitialization { PiPNN, Random }; @@ -45,8 +45,8 @@ class QGBuilder { BuildQuery(const BuildQuery&) = delete; BuildQuery& operator=(const BuildQuery&) = delete; - void reset(PID); - [[nodiscard]] float distance_to(PID) const; + void reset(PID internal_id); + [[nodiscard]] float distance_to(PID target) const; // Raw scoring needs no rotation. Prepare it lazily for FastScan only. [[nodiscard]] const float* rotated_query(); }; @@ -57,9 +57,11 @@ class QGBuilder { size_t num_points_; size_t dim_; size_t max_degree_; // degree bound for qg, multiple of 32 - bool seeded_ = true; + QGInitialization initialization_ = QGInitialization::PiPNN; static constexpr size_t kMaxCandidatePoolSize = 750; static constexpr size_t kMaxPrunedSize = 300; + // Striped locks bound reverse-edge collection memory independently of the point count. + static constexpr size_t kLockStripes = 4096; std::vector new_neighbors_; // new neighbors for current iteration std::vector degrees_; uint32_t seed_; // base seed for graph initialization @@ -79,19 +81,26 @@ class QGBuilder { std::vector workers_; const BuildCache* build_cache() const; - void find_candidates(PID, CandidateList&, VisitedSet&) const; + void find_candidates( + PID internal_id, CandidateList& results, VisitedSet& visited + ) const; void random_init(); void search_new_neighbors(); - void heuristic_prune(CandidateList&, CandidateList&, CandidateList* recorded); + void heuristic_prune( + CandidateList& pool, CandidateList& selected, CandidateList* recorded + ); void add_reverse_edges(bool final_pass); void add_pruned_edges( - const CandidateList&, const CandidateList&, CandidateList&, float + const CandidateList& result, + const CandidateList& pruned_list, + CandidateList& new_result, + float threshold ); - void complete_degree(PID, BuildWorker&); + void complete_degree(PID internal_id, BuildWorker& worker); void run_pass(bool final_pass); - void run_passes(size_t); - void write_row(PID, BuildWorker&); - void write_row(PID, const CandidateList&, BuildWorker&); + void run_passes(size_t num_iter); + void write_row(PID internal_id, BuildWorker& worker); + void write_row(PID internal_id, const CandidateList& neighbors, BuildWorker& worker); void initialize_storage(const float* data); diff --git a/include/rabitqlib/index/types.hpp b/include/rabitqlib/index/types.hpp index f421a3b4..cc94266c 100644 --- a/include/rabitqlib/index/types.hpp +++ b/include/rabitqlib/index/types.hpp @@ -5,7 +5,7 @@ namespace rabitqlib { using PID = uint32_t; -constexpr uint32_t kPidMax = 0xFFFFFFFF; +inline constexpr PID kPidMax = 0xFFFFFFFF; template struct AnnCandidate { @@ -13,7 +13,8 @@ struct AnnCandidate { T distance = std::numeric_limits::max(); AnnCandidate() = default; - explicit AnnCandidate(TP vec_id, T dis) : id(vec_id), distance(dis) {} + explicit AnnCandidate(TP candidate_id, T candidate_distance) + : id(candidate_id), distance(candidate_distance) {} friend bool operator<(const AnnCandidate& first, const AnnCandidate& second) { return first.distance < second.distance; diff --git a/include/rabitqlib/utils/io.hpp b/include/rabitqlib/utils/io.hpp index e7245c7b..5a8e9e7c 100644 --- a/include/rabitqlib/utils/io.hpp +++ b/include/rabitqlib/utils/io.hpp @@ -52,8 +52,7 @@ void read_value(std::ifstream& input, T& value, const char* description) { } // namespace io_impl inline size_t get_filesize(const char* filename) { - const auto file_size = - std::filesystem::file_size(rabitqlib::io_impl::filesystem_path(filename)); + const auto file_size = std::filesystem::file_size(io_impl::filesystem_path(filename)); if (file_size > std::numeric_limits::max()) { throw std::length_error("File is too large to address"); } @@ -61,11 +60,11 @@ inline size_t get_filesize(const char* filename) { } inline bool file_exists(const char* filename) { - return std::filesystem::exists(rabitqlib::io_impl::filesystem_path(filename)); + return std::filesystem::exists(io_impl::filesystem_path(filename)); } // load .*vecs file to a matrix (e.g., RowMajorFloatMat) -template +template void load_vecs(const char* filename, M& row_mat) { if (!file_exists(filename)) { throw std::runtime_error("File does not exist: " + std::string(filename)); @@ -77,7 +76,7 @@ void load_vecs(const char* filename, M& row_mat) { if (file_size < sizeof(uint32_t)) { throw std::runtime_error("Vector file is too small to contain a dimension"); } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ifstream input(io_impl::filesystem_path(filename), std::ios::binary); if (!input.is_open()) { throw std::runtime_error("Cannot open vector file: " + std::string(filename)); } @@ -99,7 +98,7 @@ void load_vecs(const char* filename, M& row_mat) { input.seekg(0, std::ifstream::beg); - for (size_t i = 0; i < rows; i++) { + for (size_t i = 0; i < rows; ++i) { uint32_t row_cols = 0; io_impl::read_value(input, row_cols, "row dimension"); if (row_cols != stored_cols) { @@ -114,7 +113,7 @@ void load_vecs(const char* filename, M& row_mat) { } // load .*bin file to a matrix (e.g., RowMajorFloatMat) -template +template void load_bin(const char* filename, M& row_mat) { if (!file_exists(filename)) { throw std::runtime_error("File does not exist: " + std::string(filename)); @@ -126,7 +125,7 @@ void load_bin(const char* filename, M& row_mat) { if (file_size < 2 * sizeof(uint32_t)) { throw std::runtime_error("Binary matrix file is too small to contain its header"); } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ifstream input(io_impl::filesystem_path(filename), std::ios::binary); if (!input.is_open()) { throw std::runtime_error( "Cannot open binary matrix file: " + std::string(filename) diff --git a/include/rabitqlib/utils/memory.hpp b/include/rabitqlib/utils/memory.hpp index 238af43e..d9749dfd 100644 --- a/include/rabitqlib/utils/memory.hpp +++ b/include/rabitqlib/utils/memory.hpp @@ -37,8 +37,8 @@ inline void advise_huge_pages(void* ptr, size_t size) { #if defined(__linux__) madvise(ptr, size, MADV_HUGEPAGE); #else - (void)ptr; - (void)size; + static_cast(ptr); + static_cast(size); #endif } @@ -92,7 +92,7 @@ class AlignedAllocator { using value_type = T; using is_always_equal = std::true_type; - template + template struct rebind { using other = ReboundAllocator; }; @@ -116,15 +116,15 @@ class AlignedAllocator { return false; } - [[nodiscard]] T* allocate(std::size_t n) { - if (n > std::numeric_limits::max() / sizeof(T)) { + [[nodiscard]] T* allocate(size_t n) { + if (n > std::numeric_limits::max() / sizeof(T)) { throw std::bad_array_new_length(); } return align_allocate(n * sizeof(T)); } - void deallocate(T* ptr, [[maybe_unused]] std::size_t n) { aligned_deallocate(ptr); } + void deallocate(T* ptr, [[maybe_unused]] size_t n) { aligned_deallocate(ptr); } }; // Retain alignment and object lifetimes without zero-filling scalar storage. @@ -159,6 +159,8 @@ inline T* huge_page_allocate(size_t nbytes) { return align_allocate(nbytes); } +// The prefetch helpers are also compiled into ISA-specific HNSW search translation +// units, so the whole call chain keeps internal linkage. static inline void prefetch_l1(const void* addr) { #if defined(__SSE2__) || defined(_M_X64) || defined(_M_IX86) _mm_prefetch(static_cast(addr), _MM_HINT_T0); @@ -169,7 +171,7 @@ static inline void prefetch_l1(const void* addr) { static inline void prefetch_l2(const void* addr) { #if defined(__SSE2__) || defined(_M_X64) || defined(_M_IX86) - _mm_prefetch((const char*)addr, _MM_HINT_T1); + _mm_prefetch(static_cast(addr), _MM_HINT_T1); #else __builtin_prefetch(addr, 0, 2); #endif @@ -177,7 +179,7 @@ static inline void prefetch_l2(const void* addr) { // Prefetch num_lines 64-byte lines (at most 20) starting at ptr. template -inline void mem_prefetch_lines(const char* ptr, size_t num_lines) { +static inline void mem_prefetch_lines(const char* ptr, size_t num_lines) { // The repeated fallthrough branches intentionally unroll up to 20 prefetches. // NOLINTBEGIN(bugprone-branch-clone) switch (num_lines) { @@ -268,11 +270,11 @@ inline void mem_prefetch_lines(const char* ptr, size_t num_lines) { // NOLINTEND(bugprone-branch-clone) } -inline void mem_prefetch_l1(const char* ptr, size_t num_lines) { +static inline void mem_prefetch_l1(const char* ptr, size_t num_lines) { mem_prefetch_lines(ptr, num_lines); } -inline void mem_prefetch_l2(const char* ptr, size_t num_lines) { +static inline void mem_prefetch_l2(const char* ptr, size_t num_lines) { mem_prefetch_lines(ptr, num_lines); } } // namespace rabitqlib::memory diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index c74fc471..51c0c7f0 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -29,7 +29,7 @@ namespace scalar_impl { template void scalar_quantize_normal( T* __restrict__ result, - const float* __restrict__ vec0, + const float* __restrict__ data, size_t dim, float lo, float delta @@ -40,7 +40,7 @@ void scalar_quantize_normal( } float one_over_delta = 1.0F / delta; - ConstRowMajorArrayMap v0(vec0, 1, static_cast(dim)); + ConstRowMajorArrayMap v0(data, 1, static_cast(dim)); RowMajorArrayMap res(result, 1, dim); // round to nearest integer, then cast to integer @@ -50,17 +50,17 @@ void scalar_quantize_normal( template void scalar_quantize_optimized( T* __restrict__ result, - const float* __restrict__ vec0, + const float* __restrict__ data, size_t dim, float lo, float delta ) { if constexpr (std::is_same_v) { - simd::scalar_quantize_uint8(result, vec0, dim, lo, delta); + simd::scalar_quantize_uint8(result, data, dim, lo, delta); } else if constexpr (std::is_same_v) { - simd::scalar_quantize_uint16(result, vec0, dim, lo, delta); + simd::scalar_quantize_uint16(result, data, dim, lo, delta); } else { - scalar_quantize_normal(result, vec0, dim, lo, delta); + scalar_quantize_normal(result, data, dim, lo, delta); } } } // namespace scalar_impl @@ -189,7 +189,7 @@ inline std::vector compute_centroid( #if RABITQ_USE_OPENMP const int thread_count = detail::openmp_team_size(num_threads, num_points); #else - (void)num_threads; + static_cast(num_threads); constexpr int thread_count = 1; #endif std::vector> all_results(thread_count, std::vector(dim, 0)); @@ -229,6 +229,8 @@ inline std::vector compute_centroid( namespace excode_ipimpl { +// Compatibility names: despite the _avx suffix, these call the runtime-selected +// kernel for their ex_bits (AVX-512, AVX2, NEON or generic), like select_excode_ipfunc. float ip16_fxu1_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ); @@ -267,6 +269,7 @@ using ex_ipfunc = float (*)(const float*, const uint8_t*, size_t); ex_ipfunc select_excode_ipfunc(size_t ex_bits); +// ISA-flagged sources include this header, so its scalar helpers keep internal linkage. static inline uint32_t reverse_bits(uint32_t n) { n = ((n >> 1) & 0x55555555) | ((n << 1) & 0xaaaaaaaa); n = ((n >> 2) & 0x33333333) | ((n << 2) & 0xcccccccc); @@ -288,10 +291,12 @@ static inline uint64_t reverse_bits_u64(uint64_t n) { // padded_dim is a multiple of 32. Transposed queries need ceil(padded_dim / 64) * b_query // uint64_t words, including a zero-extended final partial word in each bit plane. -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* query, uint64_t* transposed_query, 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 + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ); float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim); @@ -337,7 +342,7 @@ inline uint32_t ip_byte_bin( ) { uint32_t ret = 0; size_t offset = (padded_dim + 63) / 64; - for (size_t i = 0; i < b_query; i++) { + for (size_t i = 0; i < b_query; ++i) { ret += (ip_bin_bin(q, d, padded_dim) << i); q += offset; } diff --git a/include/rabitqlib/utils/stopw.hpp b/include/rabitqlib/utils/stopw.hpp index 90f5c8d3..89c1a1af 100644 --- a/include/rabitqlib/utils/stopw.hpp +++ b/include/rabitqlib/utils/stopw.hpp @@ -5,15 +5,15 @@ namespace rabitqlib { class StopW { - using clock = std::chrono::steady_clock; - clock::time_point time_begin_; + using Clock = std::chrono::steady_clock; + Clock::time_point time_begin_; - [[nodiscard]] clock::duration elapsed() const { return clock::now() - time_begin_; } + [[nodiscard]] Clock::duration elapsed() const { return Clock::now() - time_begin_; } public: - StopW() : time_begin_(clock::now()) {} + StopW() : time_begin_(Clock::now()) {} - void reset() { time_begin_ = clock::now(); } + void reset() { time_begin_ = Clock::now(); } [[nodiscard]] float get_elapsed_sec() const { return std::chrono::duration(elapsed()).count(); diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index cbe4dab5..af34185b 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -69,7 +69,7 @@ constexpr T round_up_to_multiple_of(size_t x, size_t multiple_of) { inline size_t floor_log2(size_t x) { size_t ret = 0; while (x > 1) { - ret++; + ++ret; x >>= 1; } return ret; diff --git a/include/rabitqlib/utils/warmup_space.hpp b/include/rabitqlib/utils/warmup_space.hpp index 050581b3..49bb676b 100644 --- a/include/rabitqlib/utils/warmup_space.hpp +++ b/include/rabitqlib/utils/warmup_space.hpp @@ -27,7 +27,7 @@ float warmup_ip_x0_q_512( namespace detail { // Scalar word-interleaved query planes; distinct from the 512-coordinate layout. -// Keep the shared loop local to each translation unit's ISA configuration. +// Internal linkage: ISA-flagged sources include this header. static inline float ip_x0_q_impl( const uint64_t* data, const uint64_t* query, @@ -44,19 +44,19 @@ static inline float ip_x0_q_impl( size_t ppc = 0; for (size_t i = 0; i < num_blk; ++i) { - const uint64_t x = rabitqlib::bitops::load_word(it_data, sizeof(uint64_t)); + const uint64_t x = bitops::load_word(it_data, sizeof(uint64_t)); ppc += bitops::popcount64(x); for (size_t j = 0; j < b_query; ++j) { - uint64_t y = *static_cast(it_query); + const uint64_t y = *it_query; ip += (bitops::popcount64(x & y) << j); - it_query++; + ++it_query; } it_data += sizeof(uint64_t); } if (padded_dim % 64 != 0) { - const uint64_t x = rabitqlib::bitops::load_word(it_data, (padded_dim % 64) / 8); + const uint64_t x = bitops::load_word(it_data, (padded_dim % 64) / 8); ppc += bitops::popcount64(x); for (size_t j = 0; j < b_query; ++j) { ip += (bitops::popcount64(x & *it_query) << j); @@ -68,7 +68,7 @@ static inline float ip_x0_q_impl( } } // namespace detail -template +template inline float warmup_ip_x0_q( const uint64_t* data, const uint64_t* query, @@ -77,7 +77,7 @@ inline float warmup_ip_x0_q( size_t padded_dim, [[maybe_unused]] size_t unused_query_bits = 0 // not used ) { - return detail::ip_x0_q_impl(data, query, delta, vl, padded_dim, b_query); + return detail::ip_x0_q_impl(data, query, delta, vl, padded_dim, BQuery); } } // namespace rabitqlib diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 123198a6..75cfbafa 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -125,18 +125,18 @@ inline void ensure_dimension(const py::array& array, size_t expected, const char } } -inline rabitqlib::MetricType metric_from_string(const std::string& metric) { +inline MetricType metric_from_string(const std::string& metric) { if (metric == "l2") { - return rabitqlib::METRIC_L2; + return METRIC_L2; } if (metric == "ip" || metric == "innerproduct") { - return rabitqlib::METRIC_IP; + return METRIC_IP; } - throw std::invalid_argument("Unsupported metric. Use 'l2' or 'ip'."); + throw std::invalid_argument("metric must be 'l2' or 'ip'"); } -inline std::string metric_to_string(rabitqlib::MetricType metric) { - return metric == rabitqlib::METRIC_IP ? "ip" : "l2"; +inline std::string metric_to_string(MetricType metric) { + return metric == METRIC_IP ? "ip" : "l2"; } template @@ -239,7 +239,7 @@ inline std::vector checked_removal_ids( ) { std::vector result; result.reserve(static_cast(ids.shape(0))); - const rabitqlib::detail::ConstPackedArrayView values( + const detail::ConstPackedArrayView values( reinterpret_cast(ids.data()) ); for (py::ssize_t i = 0; i < ids.shape(0); ++i) { @@ -269,7 +269,7 @@ inline std::vector checked_cluster_ids( std::vector result(rows); using Value = typename std::decay_t::value_type; // C-contiguous NumPy arrays can still have an unaligned data pointer. - const rabitqlib::detail::ConstPackedArrayView values( + const detail::ConstPackedArrayView values( reinterpret_cast(ids.data()) ); for (size_t i = 0; i < rows; ++i) { diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index a7630ed1..ae47f2c3 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -33,7 +33,7 @@ class HnswIndex { size_t random_seed ) : index_( - std::make_unique( + std::make_unique( max_elements, dim, nbits, @@ -100,7 +100,7 @@ class HnswIndex { built_ = true; } - py::array_t add( + py::array_t add( py::handle data, const py::object& cluster_ids, size_t num_threads, @@ -116,13 +116,13 @@ class HnswIndex { ); const auto rows = static_cast(data_array.shape(0)); - std::vector assigned; + std::vector assigned; if (!cluster_ids.is_none()) { assigned = checked_cluster_ids(cluster_ids, rows, index_->num_clusters()); } // Allocated before the index is touched. - auto ids = py::array_t(static_cast(rows)); + auto ids = py::array_t(static_cast(rows)); const auto* data_ptr = data_array.data(); const auto labels = [&] { py::gil_scoped_release release; @@ -191,7 +191,7 @@ class HnswIndex { const auto* queries_data = query_array.data(); const auto num_queries = static_cast(query_array.shape(0)); - std::vector>> results; + std::vector>> results; { py::gil_scoped_release release; results = index_->search(queries_data, num_queries, k, ef, num_threads); @@ -224,7 +224,7 @@ class HnswIndex { static HnswIndex load(py::handle path) { const auto filename = path_string(path); HnswIndex wrapper; - wrapper.index_ = std::make_unique(); + wrapper.index_ = std::make_unique(); { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); @@ -274,9 +274,11 @@ class HnswIndex { private: HnswIndex() = default; + // The public constructor creates index_; built_ is set only by a successful + // build or load. bool built_ = false; mutable IndexAccess access_; - std::unique_ptr index_; + std::unique_ptr index_; }; } // namespace rabitqlib::python_bindings diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index cecb3acb..e03f8b52 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -26,14 +26,18 @@ namespace rabitqlib::python_bindings { namespace { ivf::InitializerType initializer_from_string(const std::string& value) { - if (value == "auto") + if (value == "auto") { return ivf::InitializerType::Auto; - if (value == "flat") + } + if (value == "flat") { return ivf::InitializerType::Flat; - if (value == "flat_rabitq") + } + if (value == "flat_rabitq") { return ivf::InitializerType::FlatRaBitQ; - if (value == "hnsw") + } + if (value == "hnsw") { return ivf::InitializerType::HNSW; + } throw std::invalid_argument( "initializer must be 'auto', 'flat', 'flat_rabitq', or 'hnsw'" ); @@ -65,13 +69,13 @@ class IvfIndex { const std::string& initializer ) { const auto metric_type = metric_from_string(metric); - index_ = std::make_unique( + index_ = std::make_unique( max_elements, dim, num_clusters, nbits, metric_type, - rabitqlib::RotatorType::FhtKacRotator, + RotatorType::FhtKacRotator, initializer_from_string(initializer) ); } @@ -122,7 +126,7 @@ class IvfIndex { ); } - py::array_t add( + py::array_t add( py::handle data, const py::object& cluster_ids, size_t num_threads, @@ -138,17 +142,17 @@ class IvfIndex { ); const auto rows = static_cast(data_array.shape(0)); - std::vector assigned; + std::vector assigned; if (!cluster_ids.is_none()) { assigned = checked_cluster_ids(cluster_ids, rows, index_->num_clusters()); } const size_t first = index_->num_stored_points(); // Allocate the returned IDs before committing any changes to the index. - auto ids = py::array_t(static_cast(rows)); + auto ids = py::array_t(static_cast(rows)); auto* ids_data = ids.mutable_data(); for (size_t i = 0; i < rows; ++i) { - ids_data[i] = static_cast(first + i); + ids_data[i] = static_cast(first + i); } const auto* data_ptr = data_array.data(); @@ -177,9 +181,7 @@ class IvfIndex { throw std::invalid_argument("num_points must be non-negative"); } py::gil_scoped_release release; - index_->reserve( - static_cast(cluster_id), static_cast(num_points) - ); + index_->reserve(static_cast(cluster_id), static_cast(num_points)); } size_t remove(py::handle ids) { @@ -251,7 +253,7 @@ class IvfIndex { static IvfIndex load(py::handle path) { const auto filename = path_string(path); IvfIndex wrapper; - wrapper.index_ = std::make_unique(); + wrapper.index_ = std::make_unique(); { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); @@ -317,7 +319,7 @@ class IvfIndex { [[nodiscard]] bool built() const { return index_->num_stored_points() != 0; } mutable IndexAccess access_; - std::unique_ptr index_; + std::unique_ptr index_; }; } // namespace rabitqlib::python_bindings diff --git a/python_bindings/kmeans_bindings.cpp b/python_bindings/kmeans_bindings.cpp index e27c9aa3..bab75f8e 100644 --- a/python_bindings/kmeans_bindings.cpp +++ b/python_bindings/kmeans_bindings.cpp @@ -17,6 +17,7 @@ namespace py = pybind11; namespace rabitqlib::python_bindings { +namespace { // Without a base object, pybind11 copies the data into a new array. template @@ -27,7 +28,7 @@ py::array_t vector_to_array(const std::vector& values) { template py::array_t centroids_to_array(const Clustering& clustering, const char* name) { if (clustering.centroids.empty()) { - throw std::logic_error(std::string(name) + " has not been trained"); + throw std::runtime_error(std::string(name) + " has not been trained"); } const auto shape = std::vector{ static_cast(clustering.k), static_cast(clustering.d)}; @@ -129,6 +130,7 @@ void bind_clustering(py::module_& m, const char* native_name, const char* public .def_readonly("iteration_stats", &Clustering::iteration_stats); } +} // namespace } // namespace rabitqlib::python_bindings void register_kmeans(py::module_& m) { diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index a3a298fa..ee1d3fd1 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -35,7 +35,7 @@ class SymqgIndex { , max_degree_(max_degree) , metric_type_(metric_from_string(metric)) , quantization_bits_(quantization_bits) { - if (max_degree == 0 || max_degree % rabitqlib::fastscan::kBatchSize != 0) { + if (max_degree == 0 || max_degree % fastscan::kBatchSize != 0) { throw std::invalid_argument("max_degree must be a positive multiple of 32"); } if (quantization_bits != 0 && quantization_bits != 4 && quantization_bits != 8) { @@ -58,12 +58,12 @@ class SymqgIndex { ensure_dimension(data_array, dim_, "data dimension does not match index dim"); const size_t num_points = static_cast(data_array.shape(0)); - auto index = std::make_unique>( + auto index = std::make_unique>( num_points, dim_, max_degree_, metric_type_, - rabitqlib::RotatorType::FhtKacRotator, + RotatorType::FhtKacRotator, quantization_bits_ ); @@ -87,7 +87,7 @@ class SymqgIndex { py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads) { const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); - if (index_ == nullptr) { + if (!built()) { throw std::runtime_error("SymqgIndex must be built or loaded before search"); } if (k == 0 || k > index_->num_points()) { @@ -104,7 +104,7 @@ class SymqgIndex { const size_t num_queries = static_cast(query_array.shape(0)); const auto shape = std::vector{ static_cast(num_queries), static_cast(k)}; - auto result_ids = py::array_t(shape); + auto result_ids = py::array_t(shape); auto distances = py::array_t(shape); auto* ids_data = result_ids.mutable_data(); auto* distances_data = distances.mutable_data(); @@ -128,8 +128,8 @@ class SymqgIndex { void save(py::handle path) const { const auto filename = path_string(path); const auto operation = access_.read(); - if (index_ == nullptr) { - throw std::runtime_error("SymqgIndex must be built before save"); + if (!built()) { + throw std::runtime_error("SymqgIndex must be built or loaded before save"); } py::gil_scoped_release release; index_->save(filename.c_str()); @@ -138,7 +138,7 @@ class SymqgIndex { static SymqgIndex load(py::handle path) { const auto filename = path_string(path); SymqgIndex wrapper; - wrapper.index_ = std::make_unique>(); + wrapper.index_ = std::make_unique>(); { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); @@ -155,9 +155,9 @@ class SymqgIndex { return "SymqgIndex(dim=" + std::to_string(dim_) + ", metric='" + metric_to_string(metric_type_) + "', quantization_bits=" + std::to_string(quantization_bits_) + - ", num_points=" + std::to_string(index_ ? index_->num_points() : 0) + + ", num_points=" + std::to_string(built() ? index_->num_points() : 0) + ", max_degree=" + std::to_string(max_degree_) + - ", is_built=" + (index_ != nullptr ? "True" : "False") + ")"; + ", is_built=" + (built() ? "True" : "False") + ")"; } [[nodiscard]] size_t dim() const { @@ -170,19 +170,19 @@ class SymqgIndex { } [[nodiscard]] size_t num_points() const { const auto operation = access_.read(); - return index_ ? index_->num_points() : 0; + return built() ? index_->num_points() : 0; } [[nodiscard]] size_t num_stored_points() const { const auto operation = access_.read(); - return index_ ? index_->num_stored_points() : 0; + return built() ? index_->num_stored_points() : 0; } [[nodiscard]] size_t max_elements() const { const auto operation = access_.read(); - return index_ ? index_->max_elements() : 0; + return built() ? index_->max_elements() : 0; } [[nodiscard]] bool is_built() const { const auto operation = access_.read(); - return index_ != nullptr; + return built(); } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); @@ -196,12 +196,15 @@ class SymqgIndex { private: SymqgIndex() = default; + // A failed build keeps the previous index; a failed load discards its wrapper. + [[nodiscard]] bool built() const { return index_ != nullptr; } + size_t dim_ = 0; size_t max_degree_ = 0; - rabitqlib::MetricType metric_type_ = rabitqlib::METRIC_L2; + MetricType metric_type_ = METRIC_L2; size_t quantization_bits_ = 0; mutable IndexAccess access_; - std::unique_ptr> index_; + std::unique_ptr> index_; }; } // namespace rabitqlib::python_bindings diff --git a/src/core/api.cpp b/src/core/api.cpp index 6f89faab..fc8737d8 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -23,23 +23,24 @@ namespace rabitqlib::quant::v1 { namespace { constexpr size_t kWorkspaceAlignment = alignof(std::max_align_t); -void validate(const Parameters& p) { - if (p.dim < 2) { +void validate(const Parameters& parameters) { + if (parameters.dim < 2) { throw std::invalid_argument("dim must be at least 2"); } - if (p.dim > static_cast(std::numeric_limits::max())) { + if (parameters.dim > static_cast(std::numeric_limits::max())) { throw std::length_error("dim exceeds the supported index range"); } - if (p.total_bits < 1 || p.total_bits > 9) { + if (parameters.total_bits < 1 || parameters.total_bits > 9) { throw std::invalid_argument("total_bits must be in [1, 9]"); } - validate_metric_type(p.metric); - if (!std::isfinite(p.t_const) || p.t_const > std::numeric_limits::max() / 2.0) { + validate_metric_type(parameters.metric); + if (!std::isfinite(parameters.t_const) || + parameters.t_const > std::numeric_limits::max() / 2.0) { throw std::invalid_argument("t_const must be finite and at most INT_MAX / 2"); } } -size_t multiply(size_t count, size_t width) { +size_t checked_multiply(size_t count, size_t width) { if (width == 0) { return 0; } @@ -73,13 +74,16 @@ void check_buffer(ConstBuffer buffer, size_t required) { } } -void disjoint(const void* a, size_t a_size, const void* b, size_t b_size) { - if (a_size == 0 || b_size == 0) { +void check_disjoint( + const void* first, size_t first_bytes, const void* second, size_t second_bytes +) { + if (first_bytes == 0 || second_bytes == 0) { return; } - const auto first = reinterpret_cast(a); - const auto second = reinterpret_cast(b); - if (first < second + b_size && second < first + a_size) { + const auto first_address = reinterpret_cast(first); + const auto second_address = reinterpret_cast(second); + if (first_address < second_address + second_bytes && + second_address < first_address + first_bytes) { throw std::invalid_argument( "Writable buffers must not overlap other buffers or inputs" ); @@ -88,7 +92,7 @@ void disjoint(const void* a, size_t a_size, const void* b, size_t b_size) { void check_input(const float* input, size_t dim) { check_buffer( - {input, multiply(dim, sizeof(float))}, multiply(dim, sizeof(float)) + {input, checked_multiply(dim, sizeof(float))}, checked_multiply(dim, sizeof(float)) ); static_assert( sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 @@ -107,25 +111,25 @@ void check_input(const float* input, size_t dim) { } // Parameters and arithmetic ranges were validated by requirements/prepare_query. -std::pair code_sizes(const Parameters& p, CodeLayout layout) { +std::pair code_sizes(const Parameters& parameters, CodeLayout layout) { switch (layout) { case CodeLayout::UnpackedU16: - return {p.dim * sizeof(uint16_t), 0}; + return {parameters.dim * sizeof(uint16_t), 0}; case CodeLayout::SplitPacked: - if (p.dim % 32 != 0) { + if (parameters.dim % 32 != 0) { throw std::invalid_argument( "SplitPacked requires dim to be a multiple of 32" ); } - return {p.dim / 8, (p.dim / 8) * (p.total_bits - 1)}; + return {parameters.dim / 8, (parameters.dim / 8) * (parameters.total_bits - 1)}; default: throw std::invalid_argument("Unsupported code layout"); } } -void check_code(CodeView code, const Requirements& r) { - check_buffer(code.codes, r.code_bytes); - check_buffer(code.extra, r.extra_bytes); +void check_code(CodeView code, const Requirements& required) { + check_buffer(code.codes, required.code_bytes); + check_buffer(code.extra, required.extra_bytes); } CodeView view(CodeOutput code) { @@ -135,33 +139,40 @@ CodeView view(CodeOutput code) { {code.extra.data, code.extra.bytes}}; } -void check_output(CodeOutput code, const Requirements& r, Buffer workspace) { - check_code(view(code), r); - check_buffer({workspace.data, workspace.bytes}, r.workspace_bytes); - disjoint(code.codes.data, r.code_bytes, code.extra.data, r.extra_bytes); - disjoint(code.codes.data, r.code_bytes, workspace.data, r.workspace_bytes); - disjoint(code.extra.data, r.extra_bytes, workspace.data, r.workspace_bytes); +void check_output(CodeOutput code, const Requirements& required, Buffer workspace) { + check_code(view(code), required); + check_buffer( + {workspace.data, workspace.bytes}, required.workspace_bytes + ); + check_disjoint( + code.codes.data, required.code_bytes, code.extra.data, required.extra_bytes + ); + check_disjoint( + code.codes.data, required.code_bytes, workspace.data, required.workspace_bytes + ); + check_disjoint( + code.extra.data, required.extra_bytes, workspace.data, required.workspace_bytes + ); } void check_input_overlap( const float* input, size_t dim, CodeOutput output, - const Requirements& r, + const Requirements& required, Buffer workspace ) { if (input == nullptr) { return; } - const size_t bytes = multiply(dim, sizeof(float)); - disjoint(input, bytes, output.codes.data, r.code_bytes); - disjoint(input, bytes, output.extra.data, r.extra_bytes); - disjoint(input, bytes, workspace.data, r.workspace_bytes); + const size_t bytes = checked_multiply(dim, sizeof(float)); + check_disjoint(input, bytes, output.codes.data, required.code_bytes); + check_disjoint(input, bytes, output.extra.data, required.extra_bytes); + check_disjoint(input, bytes, workspace.data, required.workspace_bytes); } // One typed allocation list serves both requirements() and actual storage binding. -class ScratchStorage { - public: +struct ScratchStorage { rabitq_impl::Scratch scratch; uint16_t* codes; uint8_t* extra; @@ -195,180 +206,206 @@ uint16_t checked_unpacked_code(const uint8_t* bytes, size_t i, uint16_t maximum) return value; } -void read_codes(const Parameters& p, CodeView input, ScratchStorage& s) { +void read_codes(const Parameters& parameters, CodeView input, ScratchStorage& storage) { const auto* bytes = static_cast(input.codes.data); if (input.layout == CodeLayout::UnpackedU16) { - const auto maximum = static_cast((1U << p.total_bits) - 1); - for (size_t i = 0; i < p.dim; ++i) { - s.codes[i] = checked_unpacked_code(bytes, i, maximum); + const auto maximum = static_cast((1U << parameters.total_bits) - 1); + for (size_t i = 0; i < parameters.dim; ++i) { + storage.codes[i] = checked_unpacked_code(bytes, i, maximum); } return; } - rabitq_impl::unpack_sign_bits(bytes, p.dim, s.codes); - const size_t ex_bits = p.total_bits - 1; + rabitq_impl::unpack_sign_bits(bytes, parameters.dim, storage.codes); + const size_t ex_bits = parameters.total_bits - 1; if (ex_bits > 0) { rabitq_impl::ex_bits::unpacking_rabitqplus_code( - static_cast(input.extra.data), s.extra, p.dim, ex_bits + static_cast(input.extra.data), + storage.extra, + parameters.dim, + ex_bits ); - for (size_t i = 0; i < p.dim; ++i) { - s.codes[i] = static_cast((s.codes[i] << ex_bits) | s.extra[i]); + for (size_t i = 0; i < parameters.dim; ++i) { + storage.codes[i] = + static_cast((storage.codes[i] << ex_bits) | storage.extra[i]); } } } -void write_codes(const Parameters& p, CodeOutput output, ScratchStorage& s) { +void write_codes(const Parameters& parameters, CodeOutput output, ScratchStorage& storage) { auto* bytes = static_cast(output.codes.data); if (output.layout == CodeLayout::UnpackedU16) { - for (size_t i = 0; i < p.dim; ++i) { - bytes[2 * i] = static_cast(s.codes[i]); - bytes[2 * i + 1] = static_cast(s.codes[i] >> 8U); + for (size_t i = 0; i < parameters.dim; ++i) { + bytes[2 * i] = static_cast(storage.codes[i]); + bytes[2 * i + 1] = static_cast(storage.codes[i] >> 8U); } return; } - const size_t ex_bits = p.total_bits - 1; + const size_t ex_bits = parameters.total_bits - 1; const auto mask = static_cast((1U << ex_bits) - 1); - for (size_t i = 0; i < p.dim; ++i) { - s.scratch.binary[i] = s.codes[i] >> ex_bits; - s.extra[i] = static_cast(s.codes[i] & mask); + for (size_t i = 0; i < parameters.dim; ++i) { + storage.scratch.binary[i] = storage.codes[i] >> ex_bits; + storage.extra[i] = static_cast(storage.codes[i] & mask); } - pack_binary_to_bytes(s.scratch.binary, bytes, p.dim); + pack_binary_to_bytes(storage.scratch.binary, bytes, parameters.dim); if (ex_bits > 0) { rabitq_impl::ex_bits::packing_rabitqplus_code( - s.extra, static_cast(output.extra.data), p.dim, ex_bits + storage.extra, static_cast(output.extra.data), parameters.dim, ex_bits ); } } } // namespace -Requirements requirements(const Parameters& p, CodeLayout layout) { - validate(p); +Requirements requirements(const Parameters& parameters, CodeLayout layout) { + validate(parameters); // Derive the budget once from the same declarations used by encoding and // exact fallback. No array is allocated in this sizing pass. - static const auto budget = [] { + static const auto kScratchBudget = [] { rabitq_impl::ScratchAllocator sizing; [[maybe_unused]] const ScratchStorage arrays(sizing); (void)rabitq_impl::event_sweep_buffers(sizing); return sizing; }(); - Requirements r{0, 0, budget.bytes(p.dim), kWorkspaceAlignment}; - const auto [code_bytes, extra_bytes] = code_sizes(p, layout); - r.code_bytes = code_bytes; - r.extra_bytes = extra_bytes; - return r; + Requirements required{0, 0, kScratchBudget.bytes(parameters.dim), kWorkspaceAlignment}; + const auto [code_bytes, extra_bytes] = code_sizes(parameters, layout); + required.code_bytes = code_bytes; + required.extra_bytes = extra_bytes; + return required; } Factors encode( - const Parameters& p, + const Parameters& parameters, const float* input, const float* centroid, CodeOutput output, Buffer workspace ) { - const Requirements r = requirements(p, output.layout); - check_input(input, p.dim); + const Requirements required = requirements(parameters, output.layout); + check_input(input, parameters.dim); if (centroid != nullptr) { - check_input(centroid, p.dim); + check_input(centroid, parameters.dim); } - check_output(output, r, workspace); - check_input_overlap(input, p.dim, output, r, workspace); - check_input_overlap(centroid, p.dim, output, r, workspace); + check_output(output, required, workspace); + check_input_overlap(input, parameters.dim, output, required, workspace); + check_input_overlap(centroid, parameters.dim, output, required, workspace); std::pmr::monotonic_buffer_resource resource( - workspace.data, r.workspace_bytes, std::pmr::null_memory_resource() + workspace.data, required.workspace_bytes, std::pmr::null_memory_resource() ); - rabitq_impl::ScratchAllocator allocator(&resource, p.dim); - ScratchStorage s(allocator); + rabitq_impl::ScratchAllocator allocator(&resource, parameters.dim); + ScratchStorage storage(allocator); if (centroid == nullptr) { - std::fill_n(s.zero, p.dim, 0.0F); - centroid = s.zero; + std::fill_n(storage.zero, parameters.dim, 0.0F); + centroid = storage.zero; } // Reject overflow before the quantizer's float-to-integer conversions. - RowMajorArrayMap residual(s.scratch.residual, 1, static_cast(p.dim)); - residual = ConstRowMajorArrayMap(input, 1, static_cast(p.dim)) - - ConstRowMajorArrayMap(centroid, 1, static_cast(p.dim)); - const float residual_squared_norm = l2norm_sqr(s.scratch.residual, p.dim); + RowMajorArrayMap residual( + storage.scratch.residual, 1, static_cast(parameters.dim) + ); + residual = + ConstRowMajorArrayMap(input, 1, static_cast(parameters.dim)) - + ConstRowMajorArrayMap( + centroid, 1, static_cast(parameters.dim) + ); + const float residual_squared_norm = + l2norm_sqr(storage.scratch.residual, parameters.dim); if (!std::isfinite(residual_squared_norm)) { throw std::overflow_error("Residual squared norm is not representable in float32"); } Factors factors; - if (p.total_bits > 1) { + if (parameters.total_bits > 1) { rabitq_impl::ex_bits::code_with_factor_from_residual( - s.scratch.residual, + storage.scratch.residual, residual_squared_norm, centroid, - p.dim, - p.total_bits - 1, - s.codes, + parameters.dim, + parameters.total_bits - 1, + storage.codes, factors.add, factors.rescale, factors.error, - p.metric, - p.t_const, - &s.scratch + parameters.metric, + parameters.t_const, + &storage.scratch ); } else { rabitq_impl::total_bits::rabitq_full_impl_with_scratch( input, centroid, - p.dim, - p.total_bits, - s.codes, + parameters.dim, + parameters.total_bits, + storage.codes, factors.add, factors.rescale, factors.error, - p.metric, - p.t_const, - &s.scratch + parameters.metric, + parameters.t_const, + &storage.scratch ); } if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || !std::isfinite(factors.error)) { throw std::overflow_error("Quantization factors are not representable in float32"); } - write_codes(p, output, s); + write_codes(parameters, output, storage); return factors; } void convert_codes( - const Parameters& p, CodeView source, CodeOutput destination, Buffer workspace + const Parameters& parameters, CodeView source, CodeOutput destination, Buffer workspace ) { - const Requirements input = requirements(p, source.layout); - const Requirements output = requirements(p, destination.layout); - check_code(source, input); - check_output(destination, output, workspace); + const Requirements source_required = requirements(parameters, source.layout); + const Requirements destination_required = requirements(parameters, destination.layout); + check_code(source, source_required); + check_output(destination, destination_required, workspace); for (const auto& range : - {std::pair{source.codes.data, input.code_bytes}, - {source.extra.data, input.extra_bytes}}) { - disjoint(range.first, range.second, destination.codes.data, output.code_bytes); - disjoint(range.first, range.second, destination.extra.data, output.extra_bytes); - disjoint(range.first, range.second, workspace.data, output.workspace_bytes); + {std::pair{source.codes.data, source_required.code_bytes}, + {source.extra.data, source_required.extra_bytes}}) { + check_disjoint( + range.first, + range.second, + destination.codes.data, + destination_required.code_bytes + ); + check_disjoint( + range.first, + range.second, + destination.extra.data, + destination_required.extra_bytes + ); + check_disjoint( + range.first, range.second, workspace.data, destination_required.workspace_bytes + ); } std::pmr::monotonic_buffer_resource resource( - workspace.data, output.workspace_bytes, std::pmr::null_memory_resource() + workspace.data, + destination_required.workspace_bytes, + std::pmr::null_memory_resource() ); - rabitq_impl::ScratchAllocator allocator(&resource, p.dim); - ScratchStorage s(allocator); - read_codes(p, source, s); - write_codes(p, destination, s); + rabitq_impl::ScratchAllocator allocator(&resource, parameters.dim); + ScratchStorage storage(allocator); + read_codes(parameters, source, storage); + write_codes(parameters, destination, storage); } -QueryState prepare_query(const Parameters& p, const float* query, const float* centroid) { - (void)requirements(p, CodeLayout::UnpackedU16); - check_input(query, p.dim); +QueryState prepare_query( + const Parameters& parameters, const float* query, const float* centroid +) { + (void)requirements(parameters, CodeLayout::UnpackedU16); + check_input(query, parameters.dim); if (centroid != nullptr) { - check_input(centroid, p.dim); + check_input(centroid, parameters.dim); } QueryState state; state.query_ = query; - state.dim_ = p.dim; - state.total_bits_ = p.total_bits; - state.metric_ = p.metric; - state.extra_ip_ = select_excode_ipfunc(p.total_bits - 1); - state.k1xsumq_ = std::accumulate(query, query + p.dim, 0.0F) * -0.5F; + state.dim_ = parameters.dim; + state.total_bits_ = parameters.total_bits; + state.metric_ = parameters.metric; + state.extra_ip_ = select_excode_ipfunc(parameters.total_bits - 1); + state.k1xsumq_ = std::accumulate(query, query + parameters.dim, 0.0F) * -0.5F; state.g_add_ = - p.metric == METRIC_L2 - ? (centroid == nullptr ? l2norm_sqr(query, p.dim) - : euclidean_sqr(query, centroid, p.dim)) - : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, p.dim)); + parameters.metric == METRIC_L2 + ? (centroid == nullptr ? l2norm_sqr(query, parameters.dim) + : euclidean_sqr(query, centroid, parameters.dim)) + : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, parameters.dim)); if (!std::isfinite(state.g_add_) || !std::isfinite(state.k1xsumq_)) { throw std::overflow_error("Query factors are not representable in float32"); } @@ -376,13 +413,14 @@ QueryState prepare_query(const Parameters& p, const float* query, const float* c } float QueryState::compute_checked( - const QueryState& query, const CodeView& code, const Factors& factors + const QueryState& query_wrapper, const CodeView& code, const Factors& factors ) { - if (query.query_ == nullptr) { + if (query_wrapper.query_ == nullptr) { throw std::invalid_argument("QueryState has not been prepared"); } - const Parameters p{query.dim_, query.total_bits_, query.metric_}; - const auto [code_bytes, extra_bytes] = code_sizes(p, code.layout); + const Parameters parameters{ + query_wrapper.dim_, query_wrapper.total_bits_, query_wrapper.metric_}; + const auto [code_bytes, extra_bytes] = code_sizes(parameters, code.layout); check_buffer(code.codes, code_bytes); check_buffer(code.extra, extra_bytes); if (!std::isfinite(factors.add) || !std::isfinite(factors.rescale) || @@ -392,24 +430,34 @@ float QueryState::compute_checked( float ip = 0; if (code.layout == CodeLayout::UnpackedU16) { const auto* bytes = static_cast(code.codes.data); - const auto maximum = static_cast((1U << p.total_bits) - 1); - for (size_t i = 0; i < p.dim; ++i) { - ip += query.query_[i] * + const auto maximum = static_cast((1U << parameters.total_bits) - 1); + for (size_t i = 0; i < parameters.dim; ++i) { + ip += query_wrapper.query_[i] * static_cast(checked_unpacked_code(bytes, i, maximum)); } } else { - const size_t ex_bits = p.total_bits - 1; - ip = - static_cast(1U << ex_bits) * - mask_ip_x0_q(query.query_, static_cast(code.codes.data), p.dim); + const size_t ex_bits = parameters.total_bits - 1; + ip = static_cast(1U << ex_bits) * + mask_ip_x0_q( + query_wrapper.query_, + static_cast(code.codes.data), + parameters.dim + ); if (ex_bits > 0) { - ip += query.extra_ip_( - query.query_, static_cast(code.extra.data), p.dim + ip += query_wrapper.extra_ip_( + query_wrapper.query_, + static_cast(code.extra.data), + parameters.dim ); } } const float distance = quant::detail::full_est_dist_from_ip( - ip, p.total_bits, factors.add, factors.rescale, query.g_add_, query.k1xsumq_ + ip, + parameters.total_bits, + factors.add, + factors.rescale, + query_wrapper.g_add_, + query_wrapper.k1xsumq_ ); if (!std::isfinite(distance)) { throw std::overflow_error("Estimated distance is not representable in float32"); diff --git a/src/fastscan/estimator_avx2.cpp b/src/fastscan/estimator_avx2.cpp index 4050fd4c..47e63ff9 100644 --- a/src/fastscan/estimator_avx2.cpp +++ b/src/fastscan/estimator_avx2.cpp @@ -14,7 +14,7 @@ namespace { uint32_t candidate_mask_impl(const float* distances, float threshold) { const __m256 limit = _mm256_set1_ps(threshold); uint32_t mask = 0; - for (size_t lane = 0; lane < 32; lane += 8) { + for (size_t lane = 0; lane < fastscan::kBatchSize; lane += 8) { const __m256 values = _mm256_loadu_ps(distances + lane); const __m256 candidates = _mm256_cmp_ps(values, limit, _CMP_NGT_UQ); mask |= static_cast(_mm256_movemask_ps(candidates)) << lane; diff --git a/src/fastscan/estimator_avx512.cpp b/src/fastscan/estimator_avx512.cpp index 1536b4cc..d791d87a 100644 --- a/src/fastscan/estimator_avx512.cpp +++ b/src/fastscan/estimator_avx512.cpp @@ -18,7 +18,7 @@ void split_batch_estdist_avx512( ) { if (padded_dim == 0 || padded_dim % 16 != 0) { throw std::invalid_argument( - use_hacc ? "high-accuracy FastScan dimension must be a positive multiple of 16" + use_hacc ? "FastScan high-accuracy dimension must be a positive multiple of 16" : "FastScan dimension must be a positive multiple of 16" ); } diff --git a/src/fastscan/estimator_generic.cpp b/src/fastscan/estimator_generic.cpp index d0d3e325..cdf6c796 100644 --- a/src/fastscan/estimator_generic.cpp +++ b/src/fastscan/estimator_generic.cpp @@ -11,7 +11,7 @@ namespace rabitqlib::simd { namespace { uint32_t candidate_mask_impl(const float* distances, float threshold) { uint32_t mask = 0; - for (size_t lane = 0; lane < 32; ++lane) { + for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) { mask |= static_cast(!(distances[lane] > threshold)) << lane; } return mask; diff --git a/src/fastscan/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp index 2454ec26..f046ce7b 100644 --- a/src/fastscan/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -24,15 +24,17 @@ inline void split_batch_estdist_impl( bool use_hacc ) { ConstBatchDataMap cur_batch(batch_data, padded_dim); - std::array accu_values{}; - RowMajorArrayMap accu_arr(accu_values.data(), 1, fastscan::kBatchSize); + std::array accumulated_values{}; + RowMajorArrayMap accumulated_arr( + accumulated_values.data(), 1, fastscan::kBatchSize + ); if (use_hacc) { AccumulateHacc( - cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim + cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim ); } else { Accumulate( - cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim + cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim ); } @@ -50,17 +52,19 @@ inline void split_batch_estdist_impl( f_error_values.data(), 1, fastscan::kBatchSize ); - RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); + RowMajorArrayMap estimated_distances_arr( + estimated_distances, 1, fastscan::kBatchSize + ); RowMajorArrayMap ip_x0_qr_arr(ip_x0_qr, 1, fastscan::kBatchSize); - RowMajorArrayMap low_dist_arr(lower_bounds, 1, fastscan::kBatchSize); + RowMajorArrayMap lower_bounds_arr(lower_bounds, 1, fastscan::kBatchSize); - ip_x0_qr_arr = query_wrapper.delta() * (accu_arr.template cast()) + + ip_x0_qr_arr = query_wrapper.delta() * (accumulated_arr.template cast()) + query_wrapper.sum_vl_lut(); - est_dist_arr = f_add_arr + query_wrapper.g_add() + - f_rescale_arr * (ip_x0_qr_arr + query_wrapper.k1xsumq()); + estimated_distances_arr = f_add_arr + query_wrapper.g_add() + + f_rescale_arr * (ip_x0_qr_arr + query_wrapper.k1xsumq()); - low_dist_arr = est_dist_arr - f_error_arr * query_wrapper.g_error(); + lower_bounds_arr = estimated_distances_arr - f_error_arr * query_wrapper.g_error(); } } // namespace diff --git a/src/fastscan/fastscan_avx2.cpp b/src/fastscan/fastscan_avx2.cpp index 4ddc67b7..03e6fbd2 100644 --- a/src/fastscan/fastscan_avx2.cpp +++ b/src/fastscan/fastscan_avx2.cpp @@ -8,6 +8,7 @@ #include #include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" namespace rabitqlib::fastscan::simd { @@ -22,11 +23,11 @@ void pack_lut_impl( size_t dim, const float* query, float* lut, float& minimum, float& maximum ) { auto min_values = _mm256_set1_ps(kInfinity), max_values = _mm256_set1_ps(-kInfinity); - unsigned nan_mask = 0; + uint32_t nan_mask = 0; const auto update_range = [&](__m256 values) { min_values = _mm256_min_ps(min_values, values); max_values = _mm256_max_ps(max_values, values); - nan_mask |= static_cast( + nan_mask |= static_cast( _mm256_movemask_ps(_mm256_cmp_ps(values, values, _CMP_UNORD_Q)) ); }; @@ -66,8 +67,9 @@ void pack_lut_impl( minimum = min_lanes[lane] < minimum ? min_lanes[lane] : minimum; maximum = max_lanes[lane] > maximum ? max_lanes[lane] : maximum; } - if (nan_mask) + if (nan_mask) { minimum = kQuietNaN; + } } } } // namespace @@ -90,12 +92,11 @@ namespace { // After removing packed high-byte terms, each byte-plane sum is at most // 255 * (1024 / 4) = 65280, below the 16-bit reduction limit. constexpr size_t kChunkDim = 1024; -constexpr size_t kBatchSize = 32; void accumulate_chunk_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ) { size_t code_length = dim << 2; @@ -108,8 +109,8 @@ void accumulate_chunk_avx2( __m256i accu3 = _mm256_setzero_si256(); for (size_t i = 0; i < code_length; i += 64) { - c = _mm256_loadu_si256((__m256i*)&codes[i]); - lut = _mm256_loadu_si256((__m256i*)&lookup_table[i]); + c = _mm256_loadu_si256(reinterpret_cast(&codes[i])); + lut = _mm256_loadu_si256(reinterpret_cast(&lookup_table[i])); lo = _mm256_and_si256(c, low_mask); hi = _mm256_and_si256(_mm256_srli_epi16(c, 4), low_mask); @@ -121,8 +122,8 @@ void accumulate_chunk_avx2( accu2 = _mm256_add_epi16(accu2, res_hi); accu3 = _mm256_add_epi16(accu3, _mm256_srli_epi16(res_hi, 8)); - c = _mm256_loadu_si256((__m256i*)&codes[i + 32]); - lut = _mm256_loadu_si256((__m256i*)&lookup_table[i + 32]); + c = _mm256_loadu_si256(reinterpret_cast(&codes[i + 32])); + lut = _mm256_loadu_si256(reinterpret_cast(&lookup_table[i + 32])); lo = _mm256_and_si256(c, low_mask); hi = _mm256_and_si256(_mm256_srli_epi16(c, 4), low_mask); @@ -141,10 +142,12 @@ void accumulate_chunk_avx2( _mm256_blend_epi32(accu0, accu1, 0xF0) ); _mm256_storeu_si256( - (__m256i*)result, _mm256_cvtepu16_epi32(_mm256_castsi256_si128(dis0)) + reinterpret_cast<__m256i*>(accumulated_values), + _mm256_cvtepu16_epi32(_mm256_castsi256_si128(dis0)) ); _mm256_storeu_si256( - (__m256i*)(result + 8), _mm256_cvtepu16_epi32(_mm256_extracti128_si256(dis0, 1)) + reinterpret_cast<__m256i*>(accumulated_values + 8), + _mm256_cvtepu16_epi32(_mm256_extracti128_si256(dis0, 1)) ); accu2 = _mm256_sub_epi16(accu2, _mm256_slli_epi16(accu3, 8)); @@ -153,10 +156,12 @@ void accumulate_chunk_avx2( _mm256_blend_epi32(accu2, accu3, 0xF0) ); _mm256_storeu_si256( - (__m256i*)(result + 16), _mm256_cvtepu16_epi32(_mm256_castsi256_si128(dis1)) + reinterpret_cast<__m256i*>(accumulated_values + 16), + _mm256_cvtepu16_epi32(_mm256_castsi256_si128(dis1)) ); _mm256_storeu_si256( - (__m256i*)(result + 24), _mm256_cvtepu16_epi32(_mm256_extracti128_si256(dis1, 1)) + reinterpret_cast<__m256i*>(accumulated_values + 24), + _mm256_cvtepu16_epi32(_mm256_extracti128_si256(dis1, 1)) ); } @@ -165,11 +170,11 @@ void accumulate_chunk_avx2( void accumulate_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_chunk_avx2(codes, lookup_table, result, dim); + accumulate_chunk_avx2(codes, lookup_table, accumulated_values, dim); return; } @@ -191,7 +196,7 @@ void accumulate_avx2( } } for (size_t lane = 0; lane < kBatchSize; ++lane) { - result[lane] = static_cast(totals[lane]); + accumulated_values[lane] = static_cast(totals[lane]); } } @@ -199,7 +204,7 @@ namespace { void accumulate_hacc_chunk_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ) { @@ -212,18 +217,19 @@ void accumulate_hacc_chunk_avx2( } } - size_t num_codebook = dim >> 2; + size_t num_groups = dim >> 2; - for (size_t m = 0; m < num_codebook; m += 2) { - __m256i c = _mm256_loadu_si256((__m256i*)codes); + for (size_t group = 0; group < num_groups; group += 2) { + __m256i c = _mm256_loadu_si256(reinterpret_cast(codes)); codes += 32; __m256i lo = _mm256_and_si256(c, low_mask); __m256i hi = _mm256_and_si256(_mm256_srli_epi16(c, 4), low_mask); for (int q = 0; q < 2; ++q) { - __m256i lut = _mm256_loadu_si256((__m256i*)hc_lut); - hc_lut += 32; + __m256i lut = + _mm256_loadu_si256(reinterpret_cast(lookup_table)); + lookup_table += 32; __m256i res_lo = _mm256_shuffle_epi8(lut, lo); __m256i res_hi = _mm256_shuffle_epi8(lut, hi); @@ -239,14 +245,14 @@ void accumulate_hacc_chunk_avx2( __m256i dis0[2]; __m256i dis1[2]; - // lamda function to horizontal sum and combine low/high bytes + // Sum the 128-bit halves of a into the low half and those of b into the high half. auto combine2x2 = [&](__m256i a, __m256i b) -> __m256i { __m256i a1b0 = _mm256_permute2f128_si256(a, b, 0x21); __m256i a0b1 = _mm256_blend_epi32(a, b, 0xF0); return _mm256_add_epi16(a1b0, a0b1); }; - for (size_t i = 0; i < 2; i++) { + for (size_t i = 0; i < 2; ++i) { accu[i][0] = _mm256_sub_epi16(accu[i][0], _mm256_slli_epi16(accu[i][1], 8)); dis0[i] = combine2x2(accu[i][0], accu[i][1]); @@ -267,22 +273,22 @@ void accumulate_hacc_chunk_avx2( add_shiftl8(dis0[0], dis0[1], res[0], res[1]); // res for vec 0 to 15 add_shiftl8(dis1[0], dis1[1], res[2], res[3]); // res for vec 16 to 31 - _mm256_storeu_si256((__m256i*)(accumulated_values), res[0]); - _mm256_storeu_si256((__m256i*)(accumulated_values + 8), res[1]); - _mm256_storeu_si256((__m256i*)(accumulated_values + 16), res[2]); - _mm256_storeu_si256((__m256i*)(accumulated_values + 24), res[3]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(accumulated_values), res[0]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(accumulated_values + 8), res[1]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(accumulated_values + 16), res[2]); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(accumulated_values + 24), res[3]); } } // namespace void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_hacc_chunk_avx2(codes, hc_lut, accumulated_values, dim); + accumulate_hacc_chunk_avx2(codes, lookup_table, accumulated_values, dim); return; } @@ -291,7 +297,7 @@ void accumulate_hacc_avx2( for (size_t offset = 0; offset < dim;) { const size_t chunk_dim = std::min(kChunkDim, dim - offset); accumulate_hacc_chunk_avx2( - codes + offset * 4, hc_lut + offset * 8, chunk_result.data(), chunk_dim + codes + offset * 4, lookup_table + offset * 8, chunk_result.data(), chunk_dim ); for (size_t lane = 0; lane < kBatchSize; ++lane) { totals[lane] += chunk_result[lane]; @@ -300,7 +306,7 @@ void accumulate_hacc_avx2( } for (size_t lane = 0; lane < kBatchSize; ++lane) { if (totals[lane] > std::numeric_limits::max()) { - throw std::overflow_error("high-accuracy FastScan result exceeds int32_t"); + throw std::overflow_error("FastScan high-accuracy result exceeds int32_t"); } } for (size_t lane = 0; lane < kBatchSize; ++lane) { diff --git a/src/fastscan/fastscan_avx512.cpp b/src/fastscan/fastscan_avx512.cpp index 39117296..3a95d466 100644 --- a/src/fastscan/fastscan_avx512.cpp +++ b/src/fastscan/fastscan_avx512.cpp @@ -8,6 +8,7 @@ #include #include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" namespace rabitqlib::fastscan::simd { @@ -22,7 +23,7 @@ void pack_lut_impl( size_t dim, const float* query, float* lut, float& minimum, float& maximum ) { auto min_values = _mm512_set1_ps(kInfinity), max_values = _mm512_set1_ps(-kInfinity); - unsigned nan_mask = 0; + uint32_t nan_mask = 0; // Keep the initial +0 addition observable for signed-zero inputs. volatile float zero = 0; @@ -68,12 +69,11 @@ namespace { // After removing packed high-byte terms, each byte-plane sum is at most // 255 * (1024 / 4) = 65280, below the 16-bit reduction limit. constexpr size_t kChunkDim = 1024; -constexpr size_t kBatchSize = 32; void accumulate_chunk_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ) { size_t code_length = dim << 2; @@ -90,8 +90,7 @@ void accumulate_chunk_avx512( __m512i accu2 = _mm512_setzero_si512(); __m512i accu3 = _mm512_setzero_si512(); - // ! here, we assume the code_length is a multiple of 64, thus the dim must be a - // ! multiple of 16 + // Requires code_length to be a multiple of 64, so dim must be a multiple of 16. for (size_t i = 0; i < code_length; i += 64) { c = _mm512_loadu_si512(&codes[i]); lut = _mm512_loadu_si512(&lookup_table[i]); @@ -130,9 +129,11 @@ void accumulate_chunk_avx512( ret = _mm512_add_epi16(ret, _mm512_shuffle_i64x2(ret1, ret2, 0b10001000)); ret = _mm512_add_epi16(ret, _mm512_shuffle_i64x2(ret1, ret2, 0b11011101)); - _mm512_storeu_si512(result, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(ret))); _mm512_storeu_si512( - result + 16, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(ret, 1)) + accumulated_values, _mm512_cvtepu16_epi32(_mm512_castsi512_si256(ret)) + ); + _mm512_storeu_si512( + accumulated_values + 16, _mm512_cvtepu16_epi32(_mm512_extracti64x4_epi64(ret, 1)) ); } @@ -141,11 +142,11 @@ void accumulate_chunk_avx512( void accumulate_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_chunk_avx512(codes, lookup_table, result, dim); + accumulate_chunk_avx512(codes, lookup_table, accumulated_values, dim); return; } @@ -167,14 +168,14 @@ void accumulate_avx512( } } for (size_t lane = 0; lane < kBatchSize; ++lane) { - result[lane] = static_cast(totals[lane]); + accumulated_values[lane] = static_cast(totals[lane]); } } void transfer_lut_hacc_avx512(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { - size_t num_codebook = dim >> 2; + size_t num_groups = dim >> 2; - for (size_t i = 0; i < num_codebook; i++) { + for (size_t group = 0; group < num_groups; ++group) { constexpr size_t kRegBits = 512; constexpr size_t kLaneBits = 128; constexpr size_t kByteBits = 8; @@ -183,8 +184,8 @@ void transfer_lut_hacc_avx512(const uint16_t* lut, size_t dim, uint8_t* hc_lut) constexpr size_t kCodePerIter = 2 * kRegBits / kByteBits; constexpr size_t kCodePerLine = kLaneBits / kByteBits; - uint8_t* fill_lo = - hc_lut + (i / kLutPerIter * kCodePerIter) + ((i % kLutPerIter) * kCodePerLine); + uint8_t* fill_lo = hc_lut + (group / kLutPerIter * kCodePerIter) + + ((group % kLutPerIter) * kCodePerLine); uint8_t* fill_hi = fill_lo + (kRegBits / kByteBits); __m512i tmp = _mm512_cvtepi16_epi32( @@ -202,7 +203,7 @@ namespace { void accumulate_hacc_chunk_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ) { @@ -215,9 +216,9 @@ void accumulate_hacc_chunk_avx512( } } - size_t num_codebook = dim >> 2; + size_t num_groups = dim >> 2; - for (size_t m = 0; m < num_codebook; m += 4) { + for (size_t group = 0; group < num_groups; group += 4) { __m512i c = _mm512_loadu_si512(codes); __m512i lo = _mm512_and_si512(c, low_mask); __m512i hi = _mm512_and_si512(_mm512_srli_epi16(c, 4), low_mask); @@ -226,7 +227,7 @@ void accumulate_hacc_chunk_avx512( // accu[0][0-3] for lower 8-bit result // accu[1][0-3] for upper 8-bit result for (auto& i : accu) { - __m512i lut = _mm512_loadu_si512(hc_lut); + __m512i lut = _mm512_loadu_si512(lookup_table); __m512i res_lo = _mm512_shuffle_epi8(lut, lo); __m512i res_hi = _mm512_shuffle_epi8(lut, hi); @@ -237,7 +238,7 @@ void accumulate_hacc_chunk_avx512( i[2] = _mm512_add_epi16(i[2], res_hi); i[3] = _mm512_add_epi16(i[3], _mm512_srli_epi16(res_hi, 8)); - hc_lut += 64; + lookup_table += 64; } codes += 64; } @@ -287,12 +288,12 @@ void accumulate_hacc_chunk_avx512( void accumulate_hacc_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_hacc_chunk_avx512(codes, hc_lut, accumulated_values, dim); + accumulate_hacc_chunk_avx512(codes, lookup_table, accumulated_values, dim); return; } @@ -301,7 +302,7 @@ void accumulate_hacc_avx512( for (size_t offset = 0; offset < dim;) { const size_t chunk_dim = std::min(kChunkDim, dim - offset); accumulate_hacc_chunk_avx512( - codes + offset * 4, hc_lut + offset * 8, chunk_result.data(), chunk_dim + codes + offset * 4, lookup_table + offset * 8, chunk_result.data(), chunk_dim ); for (size_t lane = 0; lane < kBatchSize; ++lane) { totals[lane] += chunk_result[lane]; @@ -310,7 +311,7 @@ void accumulate_hacc_avx512( } for (size_t lane = 0; lane < kBatchSize; ++lane) { if (totals[lane] > std::numeric_limits::max()) { - throw std::overflow_error("high-accuracy FastScan result exceeds int32_t"); + throw std::overflow_error("FastScan high-accuracy result exceeds int32_t"); } } for (size_t lane = 0; lane < kBatchSize; ++lane) { diff --git a/src/fastscan/fastscan_backend.hpp b/src/fastscan/fastscan_backend.hpp index 46b06667..dd73875b 100644 --- a/src/fastscan/fastscan_backend.hpp +++ b/src/fastscan/fastscan_backend.hpp @@ -7,7 +7,7 @@ namespace rabitqlib::fastscan::detail { // Standard LUTs use consecutive 16-entry tables on every backend. HACC stores -// low/high byte planes in groups of two tables (32 bytes) or four (64 bytes). +// low/high byte planes for two tables at a time (32 bytes) or four (64 bytes). // This describes transient query scratch, never persisted codes. enum class HaccLutLayout { Pair32, Quad64 }; diff --git a/src/fastscan/fastscan_dispatch.cpp b/src/fastscan/fastscan_dispatch.cpp index b3ba0886..7c9b2aef 100644 --- a/src/fastscan/fastscan_dispatch.cpp +++ b/src/fastscan/fastscan_dispatch.cpp @@ -75,14 +75,14 @@ constexpr Backend kNeonBackend{ RABITQ_COLD_INIT const Backend& selected_backend() { #if defined(__aarch64__) || defined(_M_ARM64) - static const Backend* const backend = rabitqlib::simd::resolve_optional_kernel( + static const Backend* const kBackend = rabitqlib::simd::resolve_optional_kernel( &kNeonBackend, &kGenericBackend, cpu::has_neon() ); #else - static const Backend* const backend = + static const Backend* const kBackend = RABITQ_RESOLVE(&kAvx512Backend, &kAvx2Backend, &kGenericBackend); #endif - return *backend; + return *kBackend; } bool pack_lut_with_range( @@ -114,9 +114,9 @@ void qg_batch_estdist( size_t padded_dim, float* estimated_distances ) { - static const auto kQgBatchEstdistFn = fastscan::detail::selected_backend().qg; + static const auto kQGBatchEstdistFn = fastscan::detail::selected_backend().qg; - kQgBatchEstdistFn(batch_data, query_wrapper, padded_dim, estimated_distances); + kQGBatchEstdistFn(batch_data, query_wrapper, padded_dim, estimated_distances); } uint32_t qg_batch_estdist_mask( @@ -126,9 +126,9 @@ uint32_t qg_batch_estdist_mask( float* estimated_distances, float threshold ) { - static const auto kQgBatchEstdistMaskFn = fastscan::detail::selected_backend().qg_mask; + static const auto kQGBatchEstdistMaskFn = fastscan::detail::selected_backend().qg_mask; - return kQgBatchEstdistMaskFn( + return kQGBatchEstdistMaskFn( batch_data, query_wrapper, padded_dim, estimated_distances, threshold ); } @@ -169,7 +169,7 @@ void pack_lut(size_t dim, const float* __restrict__ query, float* __restr void accumulate( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lookup_table, - int32_t* __restrict__ result, + int32_t* __restrict__ accumulated_values, size_t dim ) { static const auto kAccumulateFn = detail::selected_backend().accumulate; @@ -177,7 +177,7 @@ void accumulate( if (dim == 0 || dim % 16 != 0) { throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); } - kAccumulateFn(codes, lookup_table, result, dim); + kAccumulateFn(codes, lookup_table, accumulated_values, dim); } void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { @@ -185,7 +185,7 @@ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { if (dim == 0 || dim % 16 != 0) { throw std::invalid_argument( - "high-accuracy FastScan dimension must be a positive multiple of 16" + "FastScan high-accuracy dimension must be a positive multiple of 16" ); } kTransferLutHaccFn(lut, dim, hc_lut); @@ -193,7 +193,7 @@ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { void accumulate_hacc( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ hc_lut, + const uint8_t* __restrict__ lookup_table, int32_t* accumulated_values, size_t dim ) { @@ -201,10 +201,10 @@ void accumulate_hacc( if (dim == 0 || dim % 16 != 0) { throw std::invalid_argument( - "high-accuracy FastScan dimension must be a positive multiple of 16" + "FastScan high-accuracy dimension must be a positive multiple of 16" ); } - kAccumulateHaccFn(codes, hc_lut, accumulated_values, dim); + kAccumulateHaccFn(codes, lookup_table, accumulated_values, dim); } } // namespace rabitqlib::fastscan diff --git a/src/fastscan/fastscan_generic.cpp b/src/fastscan/fastscan_generic.cpp index 22ea98d5..a95d357e 100644 --- a/src/fastscan/fastscan_generic.cpp +++ b/src/fastscan/fastscan_generic.cpp @@ -72,41 +72,46 @@ void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut) namespace { template void accumulate_generic_impl( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* codes, + const uint8_t* lookup_table, + int32_t* accumulated_values, + size_t dim ) { - int64_t totals[32]{}; + int64_t totals[kBatchSize]{}; for (size_t group = 0; group < dim / 4; ++group) { const size_t offset = HighAccuracy ? group / 2 * 64 + group % 2 * 16 : group * 16; for (size_t lane = 0; lane < 16; ++lane) { const uint8_t code = codes[group * 16 + lane]; for (size_t half = 0; half < 2; ++half) { const size_t entry = (code >> (half * 4)) & 15; - int value = lut[offset + entry]; - if constexpr (HighAccuracy) - value += int{lut[offset + 32 + entry]} * 256; + int value = lookup_table[offset + entry]; + if constexpr (HighAccuracy) { + value += int{lookup_table[offset + 32 + entry]} * 256; + } totals[kPerm0[lane] + half * 16] += value; } } } - for (size_t lane = 0; lane < 32; ++lane) { + for (size_t lane = 0; lane < kBatchSize; ++lane) { if (totals[lane] > std::numeric_limits::max()) { throw std::overflow_error( - HighAccuracy ? "high-accuracy FastScan result exceeds int32_t" + HighAccuracy ? "FastScan high-accuracy result exceeds int32_t" : "FastScan result exceeds int32_t" ); } } - for (size_t lane = 0; lane < 32; ++lane) - result[lane] = static_cast(totals[lane]); + for (size_t lane = 0; lane < kBatchSize; ++lane) { + accumulated_values[lane] = static_cast(totals[lane]); + } } } // namespace void accumulate_generic( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* __restrict__ result, + const uint8_t* __restrict__ lookup_table, + int32_t* __restrict__ accumulated_values, size_t dim ) { - accumulate_generic_impl(codes, lut, result, dim); + accumulate_generic_impl(codes, lookup_table, accumulated_values, dim); } void accumulate_unsupported( const uint8_t* __restrict__, const uint8_t* __restrict__, int32_t* __restrict__, size_t @@ -118,10 +123,10 @@ void accumulate_unsupported( } void accumulate_hacc_generic( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* result, + const uint8_t* __restrict__ lookup_table, + int32_t* accumulated_values, size_t dim ) { - accumulate_generic_impl(codes, lut, result, dim); + accumulate_generic_impl(codes, lookup_table, accumulated_values, dim); } } // namespace rabitqlib::fastscan::simd diff --git a/src/fastscan/fastscan_neon.cpp b/src/fastscan/fastscan_neon.cpp index ba943aba..b98d2385 100644 --- a/src/fastscan/fastscan_neon.cpp +++ b/src/fastscan/fastscan_neon.cpp @@ -14,15 +14,19 @@ namespace rabitqlib::fastscan::simd { namespace { template void accumulate_impl( - const uint8_t* codes, const uint8_t* lut, int32_t* result, size_t dim + const uint8_t* codes, + const uint8_t* lookup_table, + int32_t* accumulated_values, + size_t dim ) { - std::array totals{}; + std::array totals{}; // Bound each uint32 lane to 256*65535, then widen across chunks. constexpr size_t kChunkGroups = 256; for (size_t start = 0; start < dim / 4; start += kChunkGroups) { uint32x4_t sums[8]; - for (auto& sum : sums) + for (auto& sum : sums) { sum = vdupq_n_u32(0); + } const size_t end = std::min(start + kChunkGroups, dim / 4); for (size_t group = start; group < end; ++group) { const uint8x16_t packed = vld1q_u8(codes + group * 16); @@ -30,14 +34,14 @@ void accumulate_impl( vandq_u8(packed, vdupq_n_u8(15)), vshrq_n_u8(packed, 4)}; const size_t offset = HighAccuracy ? group / 2 * 64 + group % 2 * 16 : group * 16; - const uint8x16_t low = vld1q_u8(lut + offset); + const uint8x16_t low = vld1q_u8(lookup_table + offset); for (size_t half = 0; half < 2; ++half) { const uint8x16_t values = vqtbl1q_u8(low, indices[half]); uint16x8_t lo = vmovl_u8(vget_low_u8(values)); uint16x8_t hi = vmovl_u8(vget_high_u8(values)); if constexpr (HighAccuracy) { const uint8x16_t high = - vqtbl1q_u8(vld1q_u8(lut + offset + 32), indices[half]); + vqtbl1q_u8(vld1q_u8(lookup_table + offset + 32), indices[half]); lo = vorrq_u16(lo, vshlq_n_u16(vmovl_u8(vget_low_u8(high)), 8)); hi = vorrq_u16(hi, vshlq_n_u16(vmovl_u8(vget_high_u8(high)), 8)); } @@ -47,40 +51,42 @@ void accumulate_impl( sums[half * 4 + 3] = vaddw_u16(sums[half * 4 + 3], vget_high_u16(hi)); } } - uint32_t chunk[32]; - for (size_t i = 0; i < 8; ++i) + uint32_t chunk[kBatchSize]; + for (size_t i = 0; i < 8; ++i) { vst1q_u32(chunk + i * 4, sums[i]); - for (size_t i = 0; i < 32; ++i) + } + for (size_t i = 0; i < kBatchSize; ++i) { totals[i] += chunk[i]; + } } for (auto value : totals) { if (value > std::numeric_limits::max()) { throw std::overflow_error( - HighAccuracy ? "high-accuracy FastScan result exceeds int32_t" + HighAccuracy ? "FastScan high-accuracy result exceeds int32_t" : "FastScan result exceeds int32_t" ); } } - for (size_t i = 0; i < 32; ++i) { - result[kPerm0[i % 16] + i / 16 * 16] = static_cast(totals[i]); + for (size_t i = 0; i < kBatchSize; ++i) { + accumulated_values[kPerm0[i % 16] + i / 16 * 16] = static_cast(totals[i]); } } } // namespace void accumulate_neon( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* __restrict__ result, + const uint8_t* __restrict__ lookup_table, + int32_t* __restrict__ accumulated_values, size_t dim ) { - accumulate_impl(codes, lut, result, dim); + accumulate_impl(codes, lookup_table, accumulated_values, dim); } void accumulate_hacc_neon( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lut, - int32_t* result, + const uint8_t* __restrict__ lookup_table, + int32_t* accumulated_values, size_t dim ) { - accumulate_impl(codes, lut, result, dim); + accumulate_impl(codes, lookup_table, accumulated_values, dim); } } // namespace rabitqlib::fastscan::simd diff --git a/src/index/file_replacement.hpp b/src/index/file_replacement.hpp index fe836cbf..56d82aae 100644 --- a/src/index/file_replacement.hpp +++ b/src/index/file_replacement.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -37,15 +38,13 @@ inline void copy_permissions(const std::string& existing, const std::string& rep ); } -/** - * Save through a temporary sibling that replaces `filename` only after `write` succeeds, - * so a failed save leaves any existing file unchanged and removes what it wrote. - * A replaced file keeps its permission bits, but not its owner or other hard links. - * - * `write(path)` writes the main file at `path` and each sidecar at `path + suffix`. - * Sidecars are renamed into place before the main file. If any destination exists but - * is not a regular file, every file is written in place instead, as before. - */ +// Save through a temporary sibling that replaces `filename` only after `write` succeeds, +// so a failed save leaves any existing file unchanged and removes what it wrote. +// A replaced file keeps its permission bits, but not its owner or other hard links. +// +// `write(path)` writes the main file at `path` and each sidecar at `path + suffix`. +// Sidecars are renamed into place before the main file. If any destination exists but +// is not a regular file, every file is written in place instead, as before. template void save_replacing( const std::string& filename, @@ -63,8 +62,7 @@ void save_replacing( // A random name keeps concurrent saves to one destination from sharing a temporary. std::random_device random; - const unsigned long long token = - (static_cast(random()) << 32U) ^ random(); + const uint64_t token = (static_cast(random()) << 32U) ^ random(); const std::string temporary = filename + ".tmp" + std::to_string(token); try { write(temporary); diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 366bb647..17df2a52 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -50,10 +50,10 @@ void HierarchicalNSW::AlignedDeleter::operator()(char* pointer) const noexcept { void HierarchicalNSW::LinkDeleter::operator()(char* pointer) const noexcept { std::free(pointer); } -int HierarchicalNSW::get_random_level(double reverse_size) { +int HierarchicalNSW::random_level(double multiplier) { std::uniform_real_distribution distribution(0.0, 1.0); - double r = -log(distribution(level_generator_)) * reverse_size; - return static_cast(r); + double level = -std::log(distribution(level_generator_)) * multiplier; + return static_cast(level); } HierarchicalNSW::HierarchicalNSW( @@ -93,24 +93,17 @@ HierarchicalNSW::HierarchicalNSW( ) ); padded_dim_ = rotator_->size(); - /* check size */ assert(padded_dim_ % 32 == 0); assert(padded_dim_ >= dim_); ex_bits_ = total_bits - 1; if (total_bits < 1 || total_bits > 9) { throw std::invalid_argument("HNSW quantization bits must be in [1, 9]"); - }; - - assert(padded_dim_ % 32 == 0); + } ip_func_ = select_excode_ipfunc(ex_bits_); - if (M <= 10000) { - M_ = M; - } else { - M_ = 10000; - } + M_ = std::min(M, kMaxM); upper_layer_max_degree_ = M_; base_layer_max_degree_ = M_ * 2; @@ -123,19 +116,12 @@ HierarchicalNSW::HierarchicalNSW( level_generator_.seed(random_seed); - num_stored_points_ = 0; - - visited_list_pool_ = std::make_unique(1, max_elements_); - - // initializations for special treatment of the first node - entry_point_id_ = -1; - max_level_ = -1; + visited_set_pool_ = std::make_unique(1, max_elements_); link_lists_.reserve(max_elements_); - mult_ = 1 / log(1.0 * static_cast(M_)); + level_multiplier_ = 1 / std::log(1.0 * static_cast(M_)); - this->query_config_ = - quant::faster_config(padded_dim_, SplitSingleQuery::kNumBits); + query_config_ = quant::faster_config(padded_dim_, SplitSingleQuery::kNumBits); } void HierarchicalNSW::set_record_layout() { @@ -155,7 +141,7 @@ void HierarchicalNSW::save(const char* filename) const { throw std::logic_error("HNSW index must be initialized before save"); } if (filename == nullptr || filename[0] == '\0') { - throw std::invalid_argument("HNSW: index filename must not be empty"); + throw std::invalid_argument("HNSW index filename must not be empty"); } index_detail::save_replacing(filename, {}, [this](const std::string& path) { write_file(path.c_str()); @@ -163,9 +149,9 @@ void HierarchicalNSW::save(const char* filename) const { } void HierarchicalNSW::write_file(const char* filename) const { - std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { - throw std::runtime_error("HNSW: cannot open index file for writing"); + throw std::runtime_error("HNSW cannot open index file for writing"); } // Serialize one count value rather than the bytes of the atomic object. @@ -194,7 +180,7 @@ void HierarchicalNSW::write_file(const char* filename) const { output.write(reinterpret_cast(&M_), sizeof(size_t)); output.write(reinterpret_cast(&upper_layer_max_degree_), sizeof(size_t)); output.write(reinterpret_cast(&base_layer_max_degree_), sizeof(size_t)); - output.write(reinterpret_cast(&mult_), sizeof(double)); + output.write(reinterpret_cast(&level_multiplier_), sizeof(double)); output.write(reinterpret_cast(&ef_construction_), sizeof(size_t)); output.write( @@ -207,33 +193,33 @@ void HierarchicalNSW::write_file(const char* filename) const { static_cast(num_stored_points * level0_record_bytes_) ); - for (size_t i = 0; i < num_stored_points; i++) { - unsigned int link_list_size = + for (size_t i = 0; i < num_stored_points; ++i) { + uint32_t link_list_bytes = element_levels_[i] > 0 ? upper_layer_links_bytes_ * element_levels_[i] : 0; - output.write(reinterpret_cast(&link_list_size), sizeof(unsigned int)); - if (link_list_size != 0) { - output.write(link_lists_[i].get(), link_list_size); + output.write(reinterpret_cast(&link_list_bytes), sizeof(uint32_t)); + if (link_list_bytes != 0) { + output.write(link_lists_[i].get(), link_list_bytes); } } rotator_->save(output); output.close(); if (!output) { - throw std::runtime_error("HNSW: failed to write index file"); + throw std::runtime_error("HNSW failed to write index file"); } } void HierarchicalNSW::load(const char* filename) { if (filename == nullptr || filename[0] == '\0') { - throw std::invalid_argument("HNSW: index filename must not be empty"); + throw std::invalid_argument("HNSW index filename must not be empty"); } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ifstream input(io_impl::filesystem_path(filename), std::ios::binary); if (!input.is_open()) { - throw std::runtime_error("HNSW: cannot open index file"); + throw std::runtime_error("HNSW cannot open index file"); } auto invalid_file = []() -> void { - throw std::runtime_error("HNSW: invalid or truncated index file"); + throw std::runtime_error("HNSW index file is invalid or truncated"); }; auto read_header = [&](auto& value) { input.read(reinterpret_cast(&value), sizeof(value)); @@ -273,7 +259,7 @@ void HierarchicalNSW::load(const char* filename) { read_header(loaded.M_); read_header(loaded.upper_layer_max_degree_); read_header(loaded.base_layer_max_degree_); - read_header(loaded.mult_); + read_header(loaded.level_multiplier_); read_header(loaded.ef_construction_); if (loaded.max_elements_ == 0 || @@ -285,7 +271,7 @@ void HierarchicalNSW::load(const char* filename) { loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || - loaded.M_ < 2 || loaded.M_ > 10000 || loaded.upper_layer_max_degree_ != loaded.M_ || + loaded.M_ < 2 || loaded.M_ > kMaxM || loaded.upper_layer_max_degree_ != loaded.M_ || loaded.base_layer_max_degree_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || (element_count == 0 && @@ -298,8 +284,8 @@ void HierarchicalNSW::load(const char* filename) { // The stored layout must be the one the validated parameters derive. const auto stored_layout = loaded.record_layout(); loaded.set_record_layout(); - if (loaded.record_layout() != stored_layout || !std::isfinite(loaded.mult_) || - loaded.mult_ <= 0.0) { + if (loaded.record_layout() != stored_layout || + !std::isfinite(loaded.level_multiplier_) || loaded.level_multiplier_ <= 0.0) { invalid_file(); } @@ -311,7 +297,7 @@ void HierarchicalNSW::load(const char* filename) { const size_t data_bytes = checked_product(element_count, loaded.level0_record_bytes_); // Validate pointer-table capacity before allocating it. (void)checked_product(loaded.max_elements_, sizeof(LinkStorage)); - const size_t link_headers_bytes = checked_product(element_count, sizeof(unsigned int)); + const size_t link_headers_bytes = checked_product(element_count, sizeof(uint32_t)); loaded.rotator_ = rotator_impl::make_rotator_for_loading( loaded.dim_, RotatorType::FhtKacRotator, loaded.padded_dim_ ); @@ -360,7 +346,7 @@ void HierarchicalNSW::load(const char* filename) { loaded.data_level0_memory_.get() + i * loaded.level0_record_bytes_; PID cluster_id = 0; PID label = 0; - unsigned short level0_degree = 0; + uint16_t level0_degree = 0; std::memcpy(&level0_degree, row, sizeof(level0_degree)); std::memcpy(&cluster_id, row + loaded.level0_links_bytes_, sizeof(cluster_id)); std::memcpy(&label, row + loaded.label_offset_, sizeof(label)); @@ -380,29 +366,29 @@ void HierarchicalNSW::load(const char* filename) { invalid_file(); } - unsigned int link_list_size = 0; - read_payload(reinterpret_cast(&link_list_size), sizeof(link_list_size)); - if (link_list_size != 0) { - if (link_list_size % loaded.upper_layer_links_bytes_ != 0 || - link_list_size / loaded.upper_layer_links_bytes_ > + uint32_t link_list_bytes = 0; + read_payload(reinterpret_cast(&link_list_bytes), sizeof(link_list_bytes)); + if (link_list_bytes != 0) { + if (link_list_bytes % loaded.upper_layer_links_bytes_ != 0 || + link_list_bytes / loaded.upper_layer_links_bytes_ > static_cast(loaded.max_level_) || - link_list_size > remaining) { + link_list_bytes > remaining) { invalid_file(); } - loaded.link_lists_[i].reset(static_cast(std::malloc(link_list_size))); + loaded.link_lists_[i].reset(static_cast(std::malloc(link_list_bytes))); if (loaded.link_lists_[i] == nullptr) { throw std::bad_alloc(); } loaded.element_levels_[i] = - static_cast(link_list_size / loaded.upper_layer_links_bytes_); + static_cast(link_list_bytes / loaded.upper_layer_links_bytes_); } loaded.num_stored_points_ = i + 1; - if (link_list_size != 0) { - read_payload(loaded.link_lists_[i].get(), link_list_size); + if (link_list_bytes != 0) { + read_payload(loaded.link_lists_[i].get(), link_list_bytes); for (int level = 0; level < loaded.element_levels_[i]; ++level) { const char* link = loaded.link_lists_[i].get() + level * loaded.upper_layer_links_bytes_; - unsigned short degree = 0; + uint16_t degree = 0; std::memcpy(°ree, link, sizeof(degree)); if (degree > loaded.upper_layer_max_degree_) { invalid_file(); @@ -424,7 +410,7 @@ void HierarchicalNSW::load(const char* filename) { const char* link = loaded.link_lists_[i].get() + (static_cast(level) - 1) * loaded.upper_layer_links_bytes_; - unsigned short degree = 0; + uint16_t degree = 0; std::memcpy(°ree, link, sizeof(degree)); for (size_t edge = 0; edge < degree; ++edge) { PID neighbor = 0; @@ -454,7 +440,7 @@ void HierarchicalNSW::load(const char* filename) { invalid_file(); } - loaded.visited_list_pool_ = std::make_unique(1, loaded.max_elements_); + loaded.visited_set_pool_ = std::make_unique(1, loaded.max_elements_); loaded.raw_dist_func_ = (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); @@ -476,7 +462,7 @@ void HierarchicalNSW::load(const char* filename) { swap(base_layer_max_degree_, loaded.base_layer_max_degree_); swap(ef_construction_, loaded.ef_construction_); swap(metric_type_, loaded.metric_type_); - swap(mult_, loaded.mult_); + swap(level_multiplier_, loaded.level_multiplier_); swap(max_level_, loaded.max_level_); link_list_locks_.swap(loaded.link_list_locks_); swap(entry_point_id_, loaded.entry_point_id_); @@ -495,7 +481,7 @@ void HierarchicalNSW::load(const char* filename) { swap(padded_dim_, loaded.padded_dim_); swap(centroids_memory_, loaded.centroids_memory_); label_lookup_.swap(loaded.label_lookup_); - swap(visited_list_pool_, loaded.visited_list_pool_); + swap(visited_set_pool_, loaded.visited_set_pool_); swap(ip_func_, loaded.ip_func_); swap(rotator_, loaded.rotator_); swap(query_config_, loaded.query_config_); @@ -539,7 +525,7 @@ void HierarchicalNSW::construct( centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); for (size_t i = 0; i < num_clusters; ++i) { - this->rotator_->rotate( + rotator_->rotate( centroids + (i * dim_), reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_) ); @@ -551,9 +537,9 @@ void HierarchicalNSW::construct( } construction_data_ = data; - rabitqlib::index_detail::parallel_for( - 0, num_points, num_threads, [&](size_t idx, size_t /*threadId*/) { - add_point(idx, cluster_ids[idx], config); + index_detail::parallel_for( + 0, num_points, num_threads, [&](size_t index, size_t /*thread*/) { + add_point(index, cluster_ids[index], config); } ); } catch (...) { @@ -579,21 +565,21 @@ PID HierarchicalNSW::greedy_descend( PID current_id, int from_level, int to_level, Distance distance ) { float current_distance = distance(current_id); - for (int level = from_level; level > to_level; level--) { + for (int level = from_level; level > to_level; --level) { bool changed = true; while (changed) { changed = false; std::unique_lock lock(link_list_locks_[current_id]); - const PID* data = get_linklist(current_id, level); - const size_t size = get_list_count(data); - const PID* datal = data + 1; - for (size_t i = 0; i < size; i++) { - PID cand = datal[i]; - assert(cand < num_stored_points_); - float d = distance(cand); - if (d < current_distance) { - current_distance = d; - current_id = cand; + const PID* links = upper_layer_links(current_id, level); + const size_t num_neighbors = link_count(links); + const PID* neighbors = links + 1; + for (size_t i = 0; i < num_neighbors; ++i) { + PID candidate_id = neighbors[i]; + assert(candidate_id < num_stored_points_); + float candidate_distance = distance(candidate_id); + if (candidate_distance < current_distance) { + current_distance = candidate_distance; + current_id = candidate_id; changed = true; } } @@ -604,63 +590,62 @@ PID HierarchicalNSW::greedy_descend( template maxheap> HierarchicalNSW::search_layer( - PID entry_point_id, int layer, Distance distance, Prefetch prefetch + PID entry_point_id, int level, Distance distance, Prefetch prefetch ) { - auto visited = visited_list_pool_->get_free_vislist(); + auto visited = visited_set_pool_->acquire(); maxheap> top_candidates; minheap> candidate_set; - float lower_bound = distance(entry_point_id); - top_candidates.emplace(lower_bound, entry_point_id); - candidate_set.emplace(lower_bound, entry_point_id); + float dist_k = distance(entry_point_id); + top_candidates.emplace(dist_k, entry_point_id); + candidate_set.emplace(dist_k, entry_point_id); visited->set(entry_point_id); while (!candidate_set.empty()) { - std::pair curr_el_pair = candidate_set.top(); - if (curr_el_pair.first > lower_bound && top_candidates.size() == ef_construction_) { + std::pair current = candidate_set.top(); + if (current.first > dist_k && top_candidates.size() == ef_construction_) { break; } candidate_set.pop(); - PID current_id = curr_el_pair.second; + PID current_id = current.second; std::unique_lock lock(link_list_locks_[current_id]); - const PID* data = - (layer == 0) ? get_linklist0(current_id) : get_linklist(current_id, layer); - const size_t size = get_list_count(data); - const PID* datal = data + 1; + const PID* links = + (level == 0) ? level0_links(current_id) : upper_layer_links(current_id, level); + const size_t num_neighbors = link_count(links); + const PID* neighbors = links + 1; - prefetch(*datal); - prefetch(*(datal + 1)); + prefetch(*neighbors); + prefetch(*(neighbors + 1)); - for (size_t j = 0; j < size; j++) { - PID candidate_id = datal[j]; + for (size_t j = 0; j < num_neighbors; ++j) { + PID candidate_id = neighbors[j]; assert(candidate_id < num_stored_points_); if (visited->get(candidate_id)) { continue; } visited->set(candidate_id); - if (j < size - 1) { - prefetch(*(datal + j + 1)); + if (j < num_neighbors - 1) { + prefetch(*(neighbors + j + 1)); } float candidate_distance = distance(candidate_id); - if (top_candidates.size() < ef_construction_ || - lower_bound > candidate_distance) { + if (top_candidates.size() < ef_construction_ || dist_k > candidate_distance) { candidate_set.emplace(candidate_distance, candidate_id); top_candidates.emplace(candidate_distance, candidate_id); if (top_candidates.size() > ef_construction_) { top_candidates.pop(); } if (!top_candidates.empty()) { - lower_bound = top_candidates.top().first; + dist_k = top_candidates.top().first; } } } } - visited_list_pool_->release_vis_list(std::move(visited)); + visited_set_pool_->release(std::move(visited)); return top_candidates; } @@ -677,23 +662,23 @@ void HierarchicalNSW::add_point( new_node_id = num_stored_points_; link_lists_.emplace_back(); - num_stored_points_++; + ++num_stored_points_; label_lookup_[label] = new_node_id; - node_level = get_random_level(mult_); + node_level = random_level(level_multiplier_); } std::unique_lock node_lock(link_list_locks_[new_node_id]); element_levels_[new_node_id] = node_level; - std::unique_lock entry_point_lock(global_); + std::unique_lock entry_point_lock(entry_point_mutex_); int previous_max_level = max_level_; PID current_id = entry_point_id_; if (node_level <= previous_max_level) { entry_point_lock.unlock(); } - // initialize the current memory. - memset( + // Initialize the new record. + std::memset( data_level0_memory_.get() + (new_node_id * level0_record_bytes_), 0, level0_record_bytes_ @@ -701,7 +686,7 @@ void HierarchicalNSW::add_point( // Initialisation of label and cluster id set_external_label(new_node_id, label); - set_raw_clusterid(new_node_id, cluster_id); + set_raw_cluster_id(new_node_id, cluster_id); // Quantize raw data and initialize quantized data std::vector rotated_data(padded_dim_); @@ -711,8 +696,8 @@ void HierarchicalNSW::add_point( reinterpret_cast(centroids_memory_.get()) + (cluster_id * padded_dim_), padded_dim_, ex_bits_, - get_bindata_by_internalid(new_node_id), - get_exdata_by_internalid(new_node_id), + bin_data(new_node_id), + ex_data(new_node_id), metric_type_, config ); @@ -725,19 +710,19 @@ void HierarchicalNSW::add_point( if (link_lists_[new_node_id] == nullptr) { throw std::bad_alloc(); } - memset( + std::memset( link_lists_[new_node_id].get(), 0, (upper_layer_links_bytes_ * node_level) + 1 ); } - if (static_cast(current_id) != -1) { + if (current_id != kPidMax) { const auto distance = [this, new_node_id](PID id) { - return get_data_dist(id, new_node_id); + return raw_distance(id, new_node_id); }; const auto prefetch = [this](PID id) { - rabitqlib::memory::mem_prefetch_l1( + memory::mem_prefetch_l1( reinterpret_cast( - construction_data_ + (get_external_label(id) * dim_) + construction_data_ + (external_label(id) * dim_) ), padded_dim_ / 16 ); @@ -747,7 +732,7 @@ void HierarchicalNSW::add_point( greedy_descend(current_id, previous_max_level, node_level, distance); } - for (int level = std::min(node_level, previous_max_level); level >= 0; level--) { + for (int level = std::min(node_level, previous_max_level); level >= 0; --level) { maxheap> top_candidates = search_layer(current_id, level, distance, prefetch); current_id = mutually_connect_new_element(new_node_id, top_candidates, level); @@ -761,7 +746,7 @@ void HierarchicalNSW::add_point( } } -namespace detail { +namespace { template void prune_neighbors( maxheap>& top_candidates, size_t M, Distance distance @@ -770,39 +755,39 @@ void prune_neighbors( return; } - minheap> queue_closest; - std::vector> return_list; - while (top_candidates.size() > 0) { - queue_closest.emplace(top_candidates.top()); + minheap> closest_first; + std::vector> selected; + while (!top_candidates.empty()) { + closest_first.emplace(top_candidates.top()); top_candidates.pop(); } - while (queue_closest.size() > 0) { - if (return_list.size() >= M) { + while (!closest_first.empty()) { + if (selected.size() >= M) { break; } - std::pair current_pair = queue_closest.top(); + std::pair current_pair = closest_first.top(); float query_distance = current_pair.first; - queue_closest.pop(); - bool good = true; + closest_first.pop(); + bool keep = true; - for (std::pair second_pair : return_list) { - float current_distance = distance(second_pair.second, current_pair.second); + for (std::pair selected_pair : selected) { + float current_distance = distance(selected_pair.second, current_pair.second); if (current_distance < query_distance) { - good = false; + keep = false; break; } } - if (good) { - return_list.push_back(current_pair); + if (keep) { + selected.push_back(current_pair); } } - for (std::pair current_pair : return_list) { + for (std::pair current_pair : selected) { top_candidates.emplace(current_pair); } } -} // namespace detail +} // namespace template PID HierarchicalNSW::mutually_connect( @@ -812,95 +797,93 @@ PID HierarchicalNSW::mutually_connect( Distance distance, ReverseEdgeDistance reverse_edge_distance ) { - size_t max_m = level > 0 ? upper_layer_max_degree_ : base_layer_max_degree_; - detail::prune_neighbors(top_candidates, M_, distance); + size_t max_degree = level > 0 ? upper_layer_max_degree_ : base_layer_max_degree_; + prune_neighbors(top_candidates, M_, distance); if (top_candidates.size() > M_) { - throw std::runtime_error( - "Should be not be more than M_ candidates returned by the heuristic" - ); + throw std::runtime_error("HNSW pruning must not return more than M candidates"); } std::vector selected_neighbors; selected_neighbors.reserve(M_); - while (top_candidates.size() > 0) { + while (!top_candidates.empty()) { selected_neighbors.push_back(top_candidates.top().second); top_candidates.pop(); } if (selected_neighbors.empty()) { - throw std::runtime_error("No neighbors selected for the new element"); + throw std::runtime_error("HNSW selected no neighbors for the new element"); } PID next_closest_entry_point = selected_neighbors.back(); { - PID* ll_cur = - (level == 0) ? get_linklist0(new_node_id) : get_linklist(new_node_id, level); - if (*ll_cur > 0) { - throw std::runtime_error( - "The newly inserted element should have blank link list" - ); + PID* new_links = (level == 0) ? level0_links(new_node_id) + : upper_layer_links(new_node_id, level); + if (*new_links > 0) { + throw std::runtime_error("HNSW new element must have an empty link list"); } - set_list_count(ll_cur, selected_neighbors.size()); - auto* data = ll_cur + 1; - for (size_t idx = 0; idx < selected_neighbors.size(); idx++) { - if (data[idx] != 0) { - throw std::runtime_error("Possible memory corruption"); + set_link_count(new_links, selected_neighbors.size()); + auto* neighbors = new_links + 1; + for (size_t i = 0; i < selected_neighbors.size(); ++i) { + if (neighbors[i] != 0) { + throw std::runtime_error("HNSW new element link list is not zeroed"); } - if (level > element_levels_[selected_neighbors[idx]]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); + if (level > element_levels_[selected_neighbors[i]]) { + throw std::runtime_error("HNSW cannot link on a nonexistent level"); } - data[idx] = selected_neighbors[idx]; + neighbors[i] = selected_neighbors[i]; } } for (auto selected_neighbor : selected_neighbors) { std::unique_lock lock(link_list_locks_[selected_neighbor]); - PID* ll_other = (level == 0) ? get_linklist0(selected_neighbor) - : get_linklist(selected_neighbor, level); - size_t sz_link_list_other = get_list_count(ll_other); + PID* neighbor_links = (level == 0) ? level0_links(selected_neighbor) + : upper_layer_links(selected_neighbor, level); + size_t neighbor_degree = link_count(neighbor_links); - if (sz_link_list_other > max_m) { - throw std::runtime_error("Bad value of sz_link_list_other"); + if (neighbor_degree > max_degree) { + throw std::runtime_error("HNSW neighbor degree exceeds the layer maximum"); } if (selected_neighbor == new_node_id) { - throw std::runtime_error("Trying to connect an element to itself"); + throw std::runtime_error("HNSW cannot connect an element to itself"); } if (level > element_levels_[selected_neighbor]) { - throw std::runtime_error("Trying to make a link on a non-existent level"); + throw std::runtime_error("HNSW cannot link on a nonexistent level"); } - auto* data = ll_other + 1; + auto* neighbors = neighbor_links + 1; - bool is_cur_c_present = false; - for (size_t j = 0; j < sz_link_list_other; j++) { - if (data[j] == new_node_id) { - is_cur_c_present = true; + bool already_linked = false; + for (size_t j = 0; j < neighbor_degree; ++j) { + if (neighbors[j] == new_node_id) { + already_linked = true; break; } } - if (!is_cur_c_present) { - if (sz_link_list_other < max_m) { - data[sz_link_list_other] = new_node_id; - set_list_count(ll_other, sz_link_list_other + 1); + if (!already_linked) { + if (neighbor_degree < max_degree) { + neighbors[neighbor_degree] = new_node_id; + set_link_count(neighbor_links, neighbor_degree + 1); } else { float new_neighbor_distance = reverse_edge_distance(selected_neighbor); maxheap> candidates; candidates.emplace(new_neighbor_distance, new_node_id); - for (size_t j = 0; j < sz_link_list_other; j++) { - candidates.emplace(distance(data[j], selected_neighbor), data[j]); + for (size_t j = 0; j < neighbor_degree; ++j) { + candidates.emplace( + distance(neighbors[j], selected_neighbor), neighbors[j] + ); } - detail::prune_neighbors(candidates, max_m, distance); + prune_neighbors(candidates, max_degree, distance); - int indx = 0; - while (candidates.size() > 0) { - data[indx] = candidates.top().second; + int new_degree = 0; + while (!candidates.empty()) { + neighbors[new_degree] = candidates.top().second; candidates.pop(); - indx++; + ++new_degree; } - set_list_count(ll_other, indx); + set_link_count(neighbor_links, new_degree); } } } @@ -915,9 +898,9 @@ PID HierarchicalNSW::mutually_connect_new_element( new_node_id, top_candidates, level, - [this](PID target, PID query) { return get_data_dist(target, query); }, + [this](PID target, PID query) { return raw_distance(target, query); }, // Preserve the original operand order for a newly proposed reverse edge. - [this, new_node_id](PID neighbor) { return get_data_dist(neighbor, new_node_id); } + [this, new_node_id](PID neighbor) { return raw_distance(neighbor, new_node_id); } ); } @@ -931,14 +914,14 @@ std::vector>> HierarchicalNSW::search( throw std::invalid_argument("HNSW k must be greater than zero"); } std::vector>> results(num_queries); - rabitqlib::index_detail::parallel_for( + index_detail::parallel_for( 0, num_queries, num_threads, - [&](size_t idx, size_t /*threadId*/) { + [&](size_t index, size_t /*thread*/) { std::vector rotated_query(padded_dim_); - this->rotator_->rotate(queries + (idx * dim_), rotated_query.data()); - results[idx] = detail::search_knn_sorted( + rotator_->rotate(queries + (index * dim_), rotated_query.data()); + results[index] = detail::search_knn_sorted( *this, rotated_query.data(), k, std::max(ef_search, k) ); }, @@ -974,10 +957,10 @@ size_t HierarchicalNSW::remove(const PID* labels, size_t num_points) { } size_t removed = 0; - for (const PID id : internal_ids) { - if (!is_removed(id)) { - const PID marked = get_raw_clusterid(id) | kRemovedMask; - set_raw_clusterid(id, marked); + for (const PID internal_id : internal_ids) { + if (!is_removed(internal_id)) { + const PID marked = raw_cluster_id(internal_id) | kRemovedMask; + set_raw_cluster_id(internal_id, marked); ++removed; } } @@ -987,26 +970,32 @@ size_t HierarchicalNSW::remove(const PID* labels, size_t num_points) { // Code per dimension is (sign << ex_bits) + ex, up to nine bits, scaled by // f_rescale_ex. -void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { +void HierarchicalNSW::reconstruct_rotated(PID internal_id, float* rotated) const { thread_local std::vector combined; thread_local std::vector ex_codes; combined.resize(padded_dim_); - const ConstBinDataMap bin(get_bindata_by_internalid(id), padded_dim_); + const ConstBinDataMap bin(bin_data(internal_id), padded_dim_); quant::rabitq_impl::unpack_sign_bits(bin.bin_code(), padded_dim_, combined.data()); const auto* centroid = reinterpret_cast(centroids_memory_.get()) + - (get_clusterid_by_internalid(id) * padded_dim_); + (cluster_id_of(internal_id) * padded_dim_); if (ex_bits_ == 0) { quant::reconstruct_full_vec( - combined.data(), centroid, padded_dim_, 1, bin.f_rescale(), out, metric_type_ + combined.data(), + centroid, + padded_dim_, + 1, + bin.f_rescale(), + rotated, + metric_type_ ); return; } ex_codes.resize(padded_dim_); - const ConstExDataMap ex(get_exdata_by_internalid(id), padded_dim_, ex_bits_); + const ConstExDataMap ex(ex_data(internal_id), padded_dim_, ex_bits_); quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( ex.ex_code(), ex_codes.data(), padded_dim_, ex_bits_ ); @@ -1020,14 +1009,14 @@ void HierarchicalNSW::reconstruct_rotated(PID id, float* out) const { padded_dim_, ex_bits_ + 1, ex.f_rescale_ex(), - out, + rotated, metric_type_ ); } // Only `query` is reconstructed; `target` is scored from its codes, as QGBuilder // does. Callers hold `query` fixed across a loop, so it is prepared once. -float HierarchicalNSW::get_quant_dist( +float HierarchicalNSW::quantized_distance( PID target, PID query, detail::QuantQuery& cached ) const { if (cached.id != query || !cached.prepared.has_value()) { @@ -1040,7 +1029,7 @@ float HierarchicalNSW::get_quant_dist( } const auto* centroid = reinterpret_cast(centroids_memory_.get()) + - (get_clusterid_by_internalid(target) * padded_dim_); + (cluster_id_of(target) * padded_dim_); const float dist_sqr = euclidean_sqr(cached.rotated.data(), centroid, padded_dim_); const float g_error = std::sqrt(dist_sqr); const float g_add = (metric_type_ == METRIC_IP) @@ -1049,17 +1038,13 @@ float HierarchicalNSW::get_quant_dist( if (ex_bits_ == 0) { return estimate_split_single( - get_bindata_by_internalid(target), - *cached.prepared, - padded_dim_, - g_add, - g_error + bin_data(target), *cached.prepared, padded_dim_, g_add, g_error ) .distance; } return estimate_split_single_full( - get_bindata_by_internalid(target), - get_exdata_by_internalid(target), + bin_data(target), + ex_data(target), ip_func_, *cached.prepared, padded_dim_, @@ -1081,11 +1066,11 @@ PID HierarchicalNSW::mutually_connect_quant( top_candidates, level, [this, &query](PID target, PID source) { - return get_quant_dist(target, source, query); + return quantized_distance(target, source, query); }, // Quantized scoring is directional: reconstruct the existing neighbor. [this, new_node_id, &query](PID neighbor) { - return get_quant_dist(new_node_id, neighbor, query); + return quantized_distance(new_node_id, neighbor, query); } ); } @@ -1097,7 +1082,7 @@ void HierarchicalNSW::add_point_quant( ) { std::unique_lock node_lock(link_list_locks_[new_node_id]); - std::unique_lock entry_point_lock(global_); + std::unique_lock entry_point_lock(entry_point_mutex_); int previous_max_level = max_level_; PID current_id = entry_point_id_; if (node_level <= previous_max_level) { @@ -1109,9 +1094,9 @@ void HierarchicalNSW::add_point_quant( element_levels_[new_node_id] = node_level; try { - if (static_cast(current_id) != -1) { + if (current_id != kPidMax) { const auto distance = [this, new_node_id, &query](PID id) { - return get_quant_dist(id, new_node_id, query); + return quantized_distance(id, new_node_id, query); }; if (node_level < previous_max_level) { current_id = @@ -1119,7 +1104,7 @@ void HierarchicalNSW::add_point_quant( } for (int level = std::min(node_level, previous_max_level); level >= 0; - level--) { + --level) { maxheap> top_candidates = search_layer(current_id, level, distance, [](PID) {}); current_id = @@ -1134,10 +1119,10 @@ void HierarchicalNSW::add_point_quant( } catch (...) { // Upper layers may already expose this node. Keep a route to the old // graph if base-layer linking was interrupted before choosing neighbors. - PID* base_links = get_linklist0(new_node_id); - if (previous_entry_point != kPidMax && get_list_count(base_links) == 0) { + PID* base_links = level0_links(new_node_id); + if (previous_entry_point != kPidMax && link_count(base_links) == 0) { base_links[1] = previous_entry_point; - set_list_count(base_links, 1); + set_link_count(base_links, 1); } // An interrupted promotion has no incoming links above the previous // maximum level. Discard those levels so save/load sees a valid graph. @@ -1161,16 +1146,16 @@ void HierarchicalNSW::resize(size_t new_max_elements) { throw std::invalid_argument("HNSW resize capacity must be greater than zero"); } if (new_max_elements < num_stored_points_) { - throw std::invalid_argument("HNSW resize: capacity below the element count"); + throw std::invalid_argument("HNSW resize capacity is below the element count"); } if (new_max_elements > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("HNSW resize: capacity exceeds the supported ID range"); + throw std::invalid_argument("HNSW resize capacity exceeds the supported ID range"); } if (new_max_elements == max_elements_) { return; } if (new_max_elements > std::numeric_limits::max() / level0_record_bytes_) { - throw std::invalid_argument("HNSW resize: capacity exceeds addressable storage"); + throw std::invalid_argument("HNSW resize capacity exceeds addressable storage"); } // Everything new is built before anything old is released, so a failure @@ -1189,7 +1174,7 @@ void HierarchicalNSW::resize(size_t new_max_elements) { std::vector levels(element_levels_); levels.resize(new_max_elements); std::vector locks(new_max_elements); - auto pool = std::make_unique(1, new_max_elements); + auto pool = std::make_unique(1, new_max_elements); // All potentially throwing work is complete. Moving the live owners into // reserved slots cannot allocate or throw; unused capacity stays uninitialized. @@ -1200,7 +1185,7 @@ void HierarchicalNSW::resize(size_t new_max_elements) { link_lists_ = std::move(links); element_levels_.swap(levels); link_list_locks_.swap(locks); - visited_list_pool_ = std::move(pool); + visited_set_pool_ = std::move(pool); max_elements_ = new_max_elements; } @@ -1222,10 +1207,10 @@ std::vector HierarchicalNSW::add( throw std::invalid_argument("HNSW add inputs must not be null"); } if (num_points > max_elements_ - num_stored_points_) { - throw std::invalid_argument("HNSW add: not enough capacity"); + throw std::invalid_argument("HNSW add exceeds the remaining capacity"); } if (num_points > buffer::kSearchBufferMaxPointCount - num_stored_points_) { - throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); + throw std::invalid_argument("HNSW add point count exceeds the supported ID range"); } std::vector labels(num_points); // Either the caller names the cluster of every point, or each one goes to its @@ -1240,22 +1225,20 @@ std::vector HierarchicalNSW::add( } } else { const auto* centroids = reinterpret_cast(centroids_memory_.get()); - rabitqlib::index_detail::parallel_for( - 0, num_points, num_threads, [&](size_t i, size_t) { - std::vector rotated(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated.data()); - float best = std::numeric_limits::max(); - for (size_t c = 0; c < num_clusters_; ++c) { - const float dist = raw_dist_func_( - rotated.data(), centroids + (c * padded_dim_), padded_dim_ - ); - if (dist < best) { - best = dist; - assigned[i] = static_cast(c); - } + index_detail::parallel_for(0, num_points, num_threads, [&](size_t i, size_t) { + std::vector rotated(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated.data()); + float best = std::numeric_limits::max(); + for (size_t c = 0; c < num_clusters_; ++c) { + const float dist = raw_dist_func_( + rotated.data(), centroids + (c * padded_dim_), padded_dim_ + ); + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); } } - ); + }); } quant::RabitqConfig config; @@ -1267,7 +1250,7 @@ std::vector HierarchicalNSW::add( // parallel and must not touch it. std::vector levels(num_points); for (size_t i = 0; i < num_points; ++i) { - levels[i] = get_random_level(mult_); + levels[i] = random_level(level_multiplier_); } // Every link list is allocated before any of them is installed, so running @@ -1295,35 +1278,33 @@ std::vector HierarchicalNSW::add( { std::unique_lock label_lookup_lock(label_lookup_lock_); if (num_stored_points_ + num_points > max_elements_) { - throw std::runtime_error("The number of elements exceeds the specified limit"); + throw std::runtime_error("HNSW add exceeds the remaining capacity"); } const auto first = static_cast(num_stored_points_.load()); - rabitqlib::index_detail::parallel_for( - 0, num_points, num_threads, [&](size_t i, size_t) { - const auto id = static_cast(first + i); - labels[i] = id; - std::memset( - data_level0_memory_.get() + (id * level0_record_bytes_), - 0, - level0_record_bytes_ - ); - set_external_label(id, id); - set_raw_clusterid(id, assigned[i]); - std::vector rotated_data(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated_data.data()); - quant::quantize_split_single( - rotated_data.data(), - reinterpret_cast(centroids_memory_.get()) + - (assigned[i] * padded_dim_), - padded_dim_, - ex_bits_, - get_bindata_by_internalid(id), - get_exdata_by_internalid(id), - metric_type_, - config - ); - } - ); + index_detail::parallel_for(0, num_points, num_threads, [&](size_t i, size_t) { + const auto id = static_cast(first + i); + labels[i] = id; + std::memset( + data_level0_memory_.get() + (id * level0_record_bytes_), + 0, + level0_record_bytes_ + ); + set_external_label(id, id); + set_raw_cluster_id(id, assigned[i]); + std::vector rotated_data(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_.get()) + + (assigned[i] * padded_dim_), + padded_dim_, + ex_bits_, + bin_data(id), + ex_data(id), + metric_type_, + config + ); + }); size_t registered = 0; try { for (; registered < num_points; ++registered) { @@ -1344,11 +1325,9 @@ std::vector HierarchicalNSW::add( num_stored_points_ += num_points; } - rabitqlib::index_detail::parallel_for( - 0, num_points, num_workers, [&](size_t i, size_t worker) { - add_point_quant(labels[i], levels[i], link_lists[i], queries[worker]); - } - ); + index_detail::parallel_for(0, num_points, num_workers, [&](size_t i, size_t thread) { + add_point_quant(labels[i], levels[i], link_lists[i], queries[thread]); + }); return labels; } diff --git a/src/index/hnsw_dispatch.cpp b/src/index/hnsw_dispatch.cpp index 5446725b..10de3828 100644 --- a/src/index/hnsw_dispatch.cpp +++ b/src/index/hnsw_dispatch.cpp @@ -1,7 +1,6 @@ #include "rabitqlib/index/detail/hnsw_dispatch.hpp" #include -#include #include #include @@ -15,7 +14,7 @@ using SearchKnnFn = std::vector> (*)(HierarchicalNSW&, const float*, size_t, size_t); } // namespace std::vector> search_knn_sorted( - HierarchicalNSW& index, const float* query, size_t k, size_t ef_search + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { #if defined(__x86_64__) || defined(_M_X64) // The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. @@ -33,12 +32,6 @@ std::vector> search_knn_sorted( ); #endif - return kSearchKnnFn(index, query, k, ef_search); -} -std::priority_queue> search_knn( - HierarchicalNSW& index, const float* query, size_t k, size_t ef_search -) { - auto results = search_knn_sorted(index, query, k, ef_search); - return {results.begin(), results.end()}; + return kSearchKnnFn(index, rotated_query, k, ef_search); } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index 22331535..a7f552d2 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -20,19 +20,18 @@ namespace rabitqlib::hnsw { template -inline void HierarchicalNSW::get_bin_est_direct( +inline void HierarchicalNSW::estimate_bin_direct( const std::vector& q_to_centroids, const SplitSingleQuery& query_wrapper, PID candidate_id, - HierarchicalNSW::EstimateRecord& estimate + DistanceEstimate& estimate ) const { if (metric_type_ == METRIC_IP) { - const float query_centroid_ip = - q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)]; const float residual_norm = - q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; + q_to_centroids[cluster_id_of(candidate_id) + num_clusters_]; estimate = estimate_split_single_direct( - get_bindata_by_internalid(candidate_id), + bin_data(candidate_id), query_wrapper, padded_dim_, -query_centroid_ip, @@ -40,10 +39,9 @@ inline void HierarchicalNSW::get_bin_est_direct( ); } else { // L2 distance - const float residual_norm = - q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)]; estimate = estimate_split_single_direct( - get_bindata_by_internalid(candidate_id), + bin_data(candidate_id), query_wrapper, padded_dim_, residual_norm * residual_norm, @@ -53,20 +51,19 @@ inline void HierarchicalNSW::get_bin_est_direct( } template -inline void HierarchicalNSW::get_full_est_direct( +inline void HierarchicalNSW::estimate_full_direct( const std::vector& q_to_centroids, const SplitSingleQuery& query_wrapper, PID candidate_id, - HierarchicalNSW::EstimateRecord& estimate + DistanceEstimate& estimate ) const { if (metric_type_ == METRIC_IP) { - const float query_centroid_ip = - q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float query_centroid_ip = q_to_centroids[cluster_id_of(candidate_id)]; const float residual_norm = - q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; + q_to_centroids[cluster_id_of(candidate_id) + num_clusters_]; estimate = estimate_split_single_full_direct( - get_bindata_by_internalid(candidate_id), - get_exdata_by_internalid(candidate_id), + bin_data(candidate_id), + ex_data(candidate_id), ip_func_, query_wrapper, padded_dim_, @@ -76,11 +73,10 @@ inline void HierarchicalNSW::get_full_est_direct( ); } else { // L2 distance - const float residual_norm = - q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float residual_norm = q_to_centroids[cluster_id_of(candidate_id)]; estimate = estimate_split_single_full_direct( - get_bindata_by_internalid(candidate_id), - get_exdata_by_internalid(candidate_id), + bin_data(candidate_id), + ex_data(candidate_id), ip_func_, query_wrapper, padded_dim_, @@ -111,7 +107,7 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir std::vector q_to_centroids(num_clusters_); if (metric_type_ == METRIC_L2) { - for (size_t i = 0; i < num_clusters_; i++) { + for (size_t i = 0; i < num_clusters_; ++i) { q_to_centroids[i] = std::sqrt(raw_dist_func_( rotated_query, reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), @@ -121,7 +117,7 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir } else if (metric_type_ == METRIC_IP) { q_to_centroids.resize(2 * num_clusters_); // First half stores dot(query, centroid); second half stores ||query - centroid||. - for (size_t i = 0; i < num_clusters_; i++) { + for (size_t i = 0; i < num_clusters_; ++i) { q_to_centroids[i] = dot_product( rotated_query, reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), @@ -136,30 +132,32 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir } PID current_id = entry_point_id_; - EstimateRecord current_estimate; + DistanceEstimate current_estimate; - get_bin_est_direct(q_to_centroids, query_wrapper, current_id, current_estimate); + estimate_bin_direct( + q_to_centroids, query_wrapper, current_id, current_estimate + ); - for (int level = max_level_; level > 0; level--) { + for (int level = max_level_; level > 0; --level) { bool changed = true; while (changed) { changed = false; - const PID* data = get_linklist(current_id, level); - const size_t size = get_list_count(data); + const PID* links = upper_layer_links(current_id, level); + const size_t num_neighbors = link_count(links); - const PID* datal = data + 1; - for (size_t i = 0; i < size; i++) { - PID cand = datal[i]; - assert(cand < num_stored_points_); + const PID* neighbors = links + 1; + for (size_t i = 0; i < num_neighbors; ++i) { + PID candidate_id = neighbors[i]; + assert(candidate_id < num_stored_points_); - EstimateRecord candidate_estimate; - get_bin_est_direct( - q_to_centroids, query_wrapper, cand, candidate_estimate + DistanceEstimate candidate_estimate; + estimate_bin_direct( + q_to_centroids, query_wrapper, candidate_id, candidate_estimate ); if (candidate_estimate.distance < current_estimate.distance) { current_estimate = candidate_estimate; - current_id = cand; + current_id = candidate_id; changed = true; } } @@ -177,7 +175,7 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir ); result.reserve(nearest_candidates.size()); for (const auto& candidate : nearest_candidates.candidates()) { - result.emplace_back(candidate.est_dist, get_external_label(candidate.id)); + result.emplace_back(candidate.distance, external_label(candidate.id)); } // Candidate selection uses distance only. Preserve that selection, then match // the former heap's (distance, external label) ordering within each tie group. @@ -201,56 +199,53 @@ inline void HierarchicalNSW::search_base_layer_direct( const std::vector& q_to_centroids, BoundedKNN& nearest_candidates ) const { - auto visited = visited_list_pool_->get_free_vislist(); + auto visited = visited_set_pool_->acquire(); // Keep the ef_search best estimates in a bounded sorted buffer. buffer::SearchBuffer search_pool(ef_search); - float dist_k = 1e10; + float dist_k = 1e10F; - EstimateRecord start_estimate_record; + DistanceEstimate start_estimate; if (ex_bits_ > 0) { - get_full_est_direct( - q_to_centroids, query_wrapper, entry_point_id, start_estimate_record + estimate_full_direct( + q_to_centroids, query_wrapper, entry_point_id, start_estimate ); } else { - get_bin_est_direct( - q_to_centroids, query_wrapper, entry_point_id, start_estimate_record + estimate_bin_direct( + q_to_centroids, query_wrapper, entry_point_id, start_estimate ); } - float est_dist = start_estimate_record.distance; + float estimated_distance = start_estimate.distance; // Insert initial candidate. A removed point is still explored, never returned. if (!is_removed(entry_point_id)) { - nearest_candidates.insert({est_dist, entry_point_id}); - dist_k = est_dist; + nearest_candidates.insert(AnnCandidate(entry_point_id, estimated_distance)); + dist_k = estimated_distance; } - search_pool.insert(entry_point_id, est_dist); + search_pool.insert(entry_point_id, estimated_distance); visited->set(entry_point_id); const size_t prefetch_size = (((padded_dim_ / 8) + 63) / 64) + 1; - const size_t prefetch_lookahead = 4; // Number of neighbors to prefetch in advance. + constexpr size_t kPrefetchLookahead = 4; // Number of neighbors to prefetch in advance. while (search_pool.has_next()) { // Step 1 - get the next node to explore. PID current_node_id = search_pool.pop(); - const PID* data = get_linklist0(current_node_id); - const size_t size = get_list_count(data); + const PID* links = level0_links(current_node_id); + const size_t num_neighbors = link_count(links); - for (size_t p = 0; p < prefetch_lookahead; ++p) { - rabitqlib::memory::mem_prefetch_l1( - get_bindata_by_internalid(*(data + 1 + p)), prefetch_size - ); + for (size_t p = 0; p < kPrefetchLookahead; ++p) { + memory::mem_prefetch_l1(bin_data(*(links + 1 + p)), prefetch_size); } // Iterate over neighbors. (List starts at index 1.) - for (size_t j = 1; j <= size; j++) { - const PID candidate_id = *(data + j); + for (size_t j = 1; j <= num_neighbors; ++j) { + const PID candidate_id = *(links + j); - if (j + prefetch_lookahead <= size) { - rabitqlib::memory::mem_prefetch_l1( - get_bindata_by_internalid(*(data + j + prefetch_lookahead)), - prefetch_size + if (j + kPrefetchLookahead <= num_neighbors) { + memory::mem_prefetch_l1( + bin_data(*(links + j + kPrefetchLookahead)), prefetch_size ); } @@ -259,8 +254,8 @@ inline void HierarchicalNSW::search_base_layer_direct( } visited->set(candidate_id); - EstimateRecord candidate_estimate; - get_bin_est_direct( + DistanceEstimate candidate_estimate; + estimate_bin_direct( q_to_centroids, query_wrapper, candidate_id, candidate_estimate ); @@ -272,26 +267,27 @@ inline void HierarchicalNSW::search_base_layer_direct( if (update_results) { // Compute the full estimate if promising. if (ex_bits_ > 0) { - get_full_est_direct( + estimate_full_direct( q_to_centroids, query_wrapper, candidate_id, candidate_estimate ); } - Candidate cand{candidate_estimate.distance, candidate_id}; - nearest_candidates.insert(cand); - dist_k = nearest_candidates.worst().est_dist; + nearest_candidates.insert( + AnnCandidate(candidate_id, candidate_estimate.distance) + ); + dist_k = nearest_candidates.worst().distance; } if (!search_pool.is_full(candidate_estimate.distance)) { search_pool.insert(candidate_id, candidate_estimate.distance); } - rabitqlib::memory::mem_prefetch_l2( - reinterpret_cast(get_linklist0(search_pool.next_id())), 2 + memory::mem_prefetch_l2( + reinterpret_cast(level0_links(search_pool.next_id())), 2 ); } } - visited_list_pool_->release_vis_list(std::move(visited)); + visited_set_pool_->release(std::move(visited)); } } // namespace rabitqlib::hnsw diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 67ccf355..4c7c08ae 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -47,6 +47,8 @@ namespace rabitqlib::ivf { namespace { // remove() marks a point by this f_add; save and load keep the marker. constexpr float kRemovedFAdd = std::numeric_limits::infinity(); +// Files without routing metadata used HNSW routing from this many clusters. +constexpr size_t kHistoricalHnswMinClusters = 20000; } // namespace IVF::IVF() noexcept = default; @@ -54,37 +56,39 @@ IVF::~IVF() = default; bool IVF::exceeds_storage(size_t num_points) const { const size_t max_size = std::numeric_limits::max(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); const size_t rerank_bytes = rerank_vector_bytes(); return num_points > max_size / sizeof(PID) || num_points > max_size / sizeof(float) / padded_dim_ || - num_points > max_size / batch_bytes || + num_points > max_size / batch_bytes() || (rerank_bytes != 0 && num_points > max_size / rerank_bytes); } size_t IVF::batch_data_bytes(size_t cluster_size) const { - return div_round_up(cluster_size, fastscan::kBatchSize) * - BatchDataMap::data_bytes(padded_dim_); + return div_round_up(cluster_size, fastscan::kBatchSize) * batch_bytes(); } // Keep the original inlining intent inside this translation unit, especially the // per-vector stride calculation. Public consumers only need the declarations. +inline size_t IVF::batch_bytes() const { + return BatchDataMap::data_bytes(padded_dim_); +} + inline size_t IVF::rerank_vector_bytes() const { return raw_reranking_ ? sizeof(float) * dim_ : ExDataMap::data_bytes(padded_dim_, ex_bits_); } -inline size_t IVF::ex_data_bytes(size_t cluster_size) const { +inline size_t IVF::rerank_data_bytes(size_t cluster_size) const { return rerank_vector_bytes() * cluster_size; } void IVF::resize_storage(ClusterStorage& cs, size_t num_points) const { cs.batch_storage.resize(batch_data_bytes(num_points)); - if (rerank_vector_bytes() > 0) { + if (has_rerank_data()) { if (raw_reranking_) { cs.raw_storage.resize(num_points * dim_); } else { - cs.ex_storage.resize(ex_data_bytes(num_points)); + cs.ex_storage.resize(rerank_data_bytes(num_points)); } } cs.id_storage.resize(num_points); @@ -92,11 +96,11 @@ void IVF::resize_storage(ClusterStorage& cs, size_t num_points) const { void IVF::reserve_storage(ClusterStorage& cs, size_t num_points) const { cs.batch_storage.reserve(batch_data_bytes(num_points)); - if (rerank_vector_bytes() > 0) { + if (has_rerank_data()) { if (raw_reranking_) { cs.raw_storage.reserve(num_points * dim_); } else { - cs.ex_storage.reserve(ex_data_bytes(num_points)); + cs.ex_storage.reserve(rerank_data_bytes(num_points)); } } cs.id_storage.reserve(num_points); @@ -107,14 +111,14 @@ bool IVF::has_capacity(const ClusterStorage& cs, size_t num_points) const { cs.id_storage.capacity() < num_points) { return false; } - if (rerank_vector_bytes() == 0) { + if (!has_rerank_data()) { return true; } return raw_reranking_ ? cs.raw_storage.capacity() >= num_points * dim_ - : cs.ex_storage.capacity() >= ex_data_bytes(num_points); + : cs.ex_storage.capacity() >= rerank_data_bytes(num_points); } -void IVF::free_memory() { +void IVF::free_storage() { initializer_.reset(); std::vector().swap(cluster_storage_); } @@ -146,9 +150,9 @@ IVF::IVF( size_t total_bits, MetricType metric_type, RotatorType rotator_type, - InitializerType initializer + InitializerType initializer_type ) - : initializer_type_(resolve_initializer_type(initializer, num_clusters)) + : initializer_type_(resolve_initializer_type(initializer_type, num_clusters)) , max_elements_(num_points) , dim_(dim) , padded_dim_(dim) @@ -160,7 +164,7 @@ IVF::IVF( validate_metric_type(metric_type); if ((total_bits < 1 || total_bits > 9) && total_bits != 32) { throw std::invalid_argument("IVF bits must be in [1, 9] or 32 for raw reranking"); - }; + } if (num_points == 0 || num_points > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("IVF point count exceeds the supported ID range"); } @@ -179,18 +183,8 @@ IVF::IVF( throw std::invalid_argument("IVF configuration exceeds addressable storage"); } rotator_.reset(choose_rotator(dim, rotator_type, padded_dim_)); - /* check size */ - assert(padded_dim_ % 32 == 0); - assert(padded_dim_ >= dim_); } -/** - * @brief Construct clusters in IVF - * - * @param data Data objects (N*DIM) - * @param centroids Centroid vectors (K*DIM) - * @param cluster_ids Cluster ID for each data object - */ void IVF::construct( const float* data, const float* centroids, @@ -205,24 +199,22 @@ void IVF::construct( throw std::invalid_argument("IVF construction inputs must not be null"); } - // get id list for each cluster - std::vector counts(num_clusters_, 0); + // Collect the point IDs of each cluster + std::vector cluster_sizes(num_clusters_, 0); std::vector> id_lists(num_clusters_); for (size_t i = 0; i < max_elements_; ++i) { PID cluster_id = cluster_ids[i]; if (cluster_id >= num_clusters_) { - throw std::invalid_argument("Cluster ID is out of range"); + throw std::invalid_argument("IVF cluster ID is out of range"); } id_lists[cluster_id].push_back(static_cast(i)); - counts[cluster_id] += 1; + cluster_sizes[cluster_id] += 1; } - allocate_memory(counts); + allocate_storage(cluster_sizes); + create_cluster_views(cluster_sizes); - // init the cluster list - init_clusters(counts); - - // all rotated centroids + // Rotated centroids, one row per cluster std::vector rotated_centroids(num_clusters_ * padded_dim_); quant::RabitqConfig config; @@ -230,54 +222,48 @@ void IVF::construct( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - num_threads = std::min(resolve_num_threads(num_threads), num_clusters_); - const size_t scratch_points = - std::min(fastscan::kBatchSize, *std::max_element(counts.begin(), counts.end())); + const size_t num_workers = std::min(resolve_num_threads(num_threads), num_clusters_); + const size_t scratch_points = std::min( + fastscan::kBatchSize, *std::max_element(cluster_sizes.begin(), cluster_sizes.end()) + ); // Each worker reuses one batch, independently of the largest cluster size. // Allocate before entering OpenMP so allocation failures reach the caller. - std::vector> rotated_blocks(num_threads); + std::vector> rotated_blocks(num_workers); for (auto& block : rotated_blocks) { block.resize(scratch_points * padded_dim_); } - /* Quantize each cluster */ - { - index_detail::omp_for_dynamic( - num_clusters_, num_threads, [&](size_t i, size_t thread) { - const float* cur_centroid = centroids + (i * dim_); - float* cur_rotated_c = &rotated_centroids[i * padded_dim_]; - Cluster& cp = clusters_[i]; - quantize_cluster( - cp, - id_lists[i], - data, - cur_centroid, - cur_rotated_c, - rotated_blocks[thread].data(), - config - ); - } + // Quantize each cluster + index_detail::omp_for_dynamic(num_clusters_, num_workers, [&](size_t i, size_t thread) { + quantize_cluster( + clusters_[i], + id_lists[i], + data, + centroids + (i * dim_), + &rotated_centroids[i * padded_dim_], + rotated_blocks[thread].data(), + config ); - } + }); - this->initializer_->add_vectors(rotated_centroids.data(), num_threads); + initializer_->add_vectors(rotated_centroids.data(), num_workers); ready_ = true; } -void IVF::allocate_memory(const std::vector& cluster_sizes) { +void IVF::allocate_storage(const std::vector& cluster_sizes) { ready_ = false; num_removed_ = 0; - free_memory(); + free_storage(); clusters_.clear(); if (initializer_type_ == InitializerType::Flat) { - this->initializer_ = + initializer_ = std::make_unique(padded_dim_, num_clusters_, metric_type_); } else if (initializer_type_ == InitializerType::FlatRaBitQ) { - this->initializer_ = std::make_unique( + initializer_ = std::make_unique( padded_dim_, num_clusters_, metric_type_ ); } else { - this->initializer_ = + initializer_ = std::make_unique(padded_dim_, num_clusters_, metric_type_); } cluster_storage_.assign(num_clusters_, ClusterStorage(raw_reranking_)); @@ -285,39 +271,35 @@ void IVF::allocate_memory(const std::vector& cluster_sizes) { resize_storage(cluster_storage_[i], cluster_sizes[i]); } - this->ip_func_ = select_excode_ipfunc(ex_bits_); + ip_func_ = select_excode_ipfunc(ex_bits_); } -/** - * @brief Initialize the cluster list from each cluster's own independent storage. - */ -void IVF::init_clusters(const std::vector& cluster_sizes) { - this->clusters_.reserve(num_clusters_); +// Each cluster view points into that cluster's own independent storage. +void IVF::create_cluster_views(const std::vector& cluster_sizes) { + clusters_.reserve(num_clusters_); for (size_t i = 0; i < num_clusters_; ++i) { ClusterStorage& cs = cluster_storage_[i]; - char* current_ex_data = rerank_vector_bytes() > 0 ? cs.ex_data() : nullptr; - this->clusters_.emplace_back( - cluster_sizes[i], cs.batch_data(), current_ex_data, cs.ids() - ); + char* rerank_data = has_rerank_data() ? cs.rerank_data() : nullptr; + clusters_.emplace_back(cluster_sizes[i], cs.batch_data(), rerank_data, cs.ids()); } } void IVF::quantize_cluster( - Cluster& cp, + Cluster& cluster, const std::vector& point_ids, const float* data, - const float* cur_centroid, + const float* centroid, float* rotated_centroid, float* rotated_data, const quant::RabitqConfig& config ) { size_t num_points = point_ids.size(); - assert(cp.num() == num_points); - std::copy(point_ids.begin(), point_ids.end(), cp.ids()); - this->rotator_->rotate(cur_centroid, rotated_centroid); + assert(cluster.num() == num_points); + std::copy(point_ids.begin(), point_ids.end(), cluster.ids()); + rotator_->rotate(centroid, rotated_centroid); - char* batch_data = cp.batch_data(); - char* ex_data = cp.ex_data(); + char* batch_data = cluster.batch_data(); + char* ex_data = cluster.ex_data(); for (size_t i = 0; i < num_points; i += fastscan::kBatchSize) { size_t batch_size = std::min(fastscan::kBatchSize, num_points - i); for (size_t j = 0; j < batch_size; ++j) { @@ -325,7 +307,7 @@ void IVF::quantize_cluster( rotator_->rotate(vector, rotated_data + (j * padded_dim_)); if (raw_reranking_) { std::memcpy( - cp.ex_data() + ((i + j) * rerank_vector_bytes()), + cluster.ex_data() + ((i + j) * rerank_vector_bytes()), vector, rerank_vector_bytes() ); @@ -344,7 +326,7 @@ void IVF::quantize_cluster( config ); - batch_data += BatchDataMap::data_bytes(padded_dim_); + batch_data += batch_bytes(); if (ex_bits_ > 0) { ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * batch_size; } @@ -375,18 +357,18 @@ void IVF::add( throw std::invalid_argument("IVF configuration exceeds addressable storage"); } - /* Assign every new point to a cluster */ + // Assign every new point to a cluster std::vector assigned(num_points); if (cluster_ids != nullptr) { for (size_t i = 0; i < num_points; ++i) { if (cluster_ids[i] >= num_clusters_) { - throw std::invalid_argument("Cluster ID is out of range"); + throw std::invalid_argument("IVF cluster ID is out of range"); } assigned[i] = cluster_ids[i]; } } else { - rabitqlib::index_detail::parallel_for( - 0, num_points, num_threads, [&](size_t row, size_t /*thread_id*/) { + index_detail::parallel_for( + 0, num_points, num_threads, [&](size_t row, size_t /*thread*/) { std::vector rotated(padded_dim_); std::vector> nearest(1); rotator_->rotate(data + (row * dim_), rotated.data()); @@ -427,17 +409,15 @@ void IVF::add( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - const size_t rerank_bytes = rerank_vector_bytes(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); const PID first_pid = static_cast(max_elements_); - /* Rebuild clusters without reserved capacity on the side; the rest grow in place */ + // Rebuild clusters without reserved capacity on the side; the rest grow in place for (auto& addition : additions) { const PID cluster_id = addition.cluster_id; if (has_capacity(cluster_storage_[cluster_id], addition.new_size)) { const size_t full_batches = clusters_[cluster_id].num() / fastscan::kBatchSize; addition.tail.resize( - batch_data_bytes(addition.new_size) - (full_batches * batch_bytes) + batch_data_bytes(addition.new_size) - (full_batches * batch_bytes()) ); } else { addition.staged.emplace(raw_reranking_); @@ -462,7 +442,7 @@ void IVF::add( ClusterStorage& cs = addition.staged ? *addition.staged : cluster_storage_[cluster_id]; const size_t full_batches = clusters_[cluster_id].num() / fastscan::kBatchSize; - char* tail = addition.staged ? cs.batch_data() + (full_batches * batch_bytes) + char* tail = addition.staged ? cs.batch_data() + (full_batches * batch_bytes()) : reinterpret_cast(addition.tail.data()); extend_cluster( clusters_[cluster_id], @@ -471,7 +451,7 @@ void IVF::add( initializer_->centroid(cluster_id), first_pid, cs.batch_data(), - rerank_bytes > 0 ? cs.ex_data() : nullptr, + has_rerank_data() ? cs.rerank_data() : nullptr, cs.ids(), tail, config @@ -489,7 +469,7 @@ void IVF::add( failure.rethrow(); } - /* Commit. Nothing below can throw. */ + // Commit. Nothing below can throw. for (auto& addition : additions) { const PID cluster_id = addition.cluster_id; auto& cs = cluster_storage_[cluster_id]; @@ -500,27 +480,25 @@ void IVF::add( std::copy_n( addition.tail.data(), addition.tail.size(), - cs.batch_storage.data() + (full_batches * batch_bytes) + cs.batch_storage.data() + (full_batches * batch_bytes()) ); } clusters_[cluster_id].reset( addition.new_size, cs.batch_data(), - rerank_bytes > 0 ? cs.ex_data() : nullptr, + has_rerank_data() ? cs.rerank_data() : nullptr, cs.ids() ); } max_elements_ = new_num_points; } -/** - * @brief Copy one cluster into its grown storage and append new points to it. - * - * Full batches, rerank data and IDs are copied unchanged, unless they are already in place. - * The old partial tail batch is unpacked and repacked together with the new points into - * `tail_dst`, so the result is byte-identical to quantizing the same points in the same - * order in one construct call. - */ +// Copy one cluster into its grown storage and append new points to it. +// +// Full batches, rerank data and IDs are copied unchanged, unless they are already in place. +// The old partial tail batch is unpacked and repacked together with the new points into +// `tail_dst`, so the result is byte-identical to quantizing the same points in the same +// order in one construct call. void IVF::extend_cluster( const Cluster& old_cluster, const std::vector& rows, @@ -534,44 +512,45 @@ void IVF::extend_cluster( const quant::RabitqConfig& config ) const { assert(!rows.empty()); // add() only extends clusters that receive new points - const size_t old_num = old_cluster.num(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const size_t old_size = old_cluster.num(); const size_t rerank_bytes = rerank_vector_bytes(); const size_t cols = padded_dim_ / 8; - const size_t full_batches = old_num / fastscan::kBatchSize; - const size_t kept = old_num - (full_batches * fastscan::kBatchSize); - const size_t tail_num = kept + rows.size(); + const size_t full_batches = old_size / fastscan::kBatchSize; + const size_t old_tail_size = old_size - (full_batches * fastscan::kBatchSize); + const size_t tail_size = old_tail_size + rows.size(); if (batch_dst != old_cluster.batch_data()) { - std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); + std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes(), batch_dst); } if (rerank_dst != old_cluster.ex_data()) { - std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + std::copy_n(old_cluster.ex_data(), old_size * rerank_bytes, rerank_dst); } if (ids_dst != old_cluster.ids()) { - std::copy_n(old_cluster.ids(), old_num, ids_dst); + std::copy_n(old_cluster.ids(), old_size, ids_dst); } - /* Codes and factors of every point in the tail, old points first */ - std::vector codes(tail_num * cols); - std::vector f_add(tail_num); - std::vector f_rescale(tail_num); - std::vector f_error(tail_num); - if (kept > 0) { + // Codes and factors of every point in the tail, old points first + std::vector codes(tail_size * cols); + std::vector f_add(tail_size); + std::vector f_rescale(tail_size); + std::vector f_error(tail_size); + if (old_tail_size > 0) { ConstBatchDataMap old_tail( - old_cluster.batch_data() + (full_batches * batch_bytes), padded_dim_ + old_cluster.batch_data() + (full_batches * batch_bytes()), padded_dim_ + ); + fastscan::unpack_codes( + padded_dim_, old_tail.bin_code(), old_tail_size, codes.data() ); - fastscan::unpack_codes(padded_dim_, old_tail.bin_code(), kept, codes.data()); - old_tail.f_add().copy_to(f_add.data(), kept); - old_tail.f_rescale().copy_to(f_rescale.data(), kept); - old_tail.f_error().copy_to(f_error.data(), kept); + old_tail.f_add().copy_to(f_add.data(), old_tail_size); + old_tail.f_rescale().copy_to(f_rescale.data(), old_tail_size); + old_tail.f_error().copy_to(f_error.data(), old_tail_size); } std::vector rotated(padded_dim_); for (size_t j = 0; j < rows.size(); ++j) { const float* vector = data + (static_cast(rows[j]) * dim_); - const size_t lane = kept + j; + const size_t lane = old_tail_size + j; rotator_->rotate(vector, rotated.data()); quant::quantize_compact_one_bit( rotated.data(), @@ -583,7 +562,7 @@ void IVF::extend_cluster( f_error[lane], metric_type_ ); - char* rerank = rerank_dst + ((old_num + j) * rerank_bytes); + char* rerank = rerank_dst + ((old_size + j) * rerank_bytes); if (raw_reranking_) { std::memcpy(rerank, vector, rerank_bytes); } else if (ex_bits_ > 0) { @@ -597,23 +576,23 @@ void IVF::extend_cluster( config ); } - ids_dst[old_num + j] = first_pid + rows[j]; + ids_dst[old_size + j] = first_pid + rows[j]; } - /* Repack the tail into batches */ + // Repack the tail into batches char* batch = tail_dst; - for (size_t start = 0; start < tail_num; start += fastscan::kBatchSize) { - const size_t count = std::min(fastscan::kBatchSize, tail_num - start); + for (size_t start = 0; start < tail_size; start += fastscan::kBatchSize) { + const size_t batch_size = std::min(fastscan::kBatchSize, tail_size - start); BatchDataMap cur_batch(batch, padded_dim_); fastscan::pack_codes( - padded_dim_, codes.data() + (start * cols), count, cur_batch.bin_code() + padded_dim_, codes.data() + (start * cols), batch_size, cur_batch.bin_code() ); - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < batch_size; ++i) { cur_batch.f_add()[i] = f_add[start + i]; cur_batch.f_rescale()[i] = f_rescale[start + i]; cur_batch.f_error()[i] = f_error[start + i]; } - batch += batch_bytes; + batch += batch_bytes(); } } @@ -622,7 +601,7 @@ void IVF::reserve(PID cluster_id, size_t num_points) { throw std::logic_error("IVF index must be constructed or loaded before reserve"); } if (cluster_id >= num_clusters_) { - throw std::invalid_argument("Cluster ID is out of range"); + throw std::invalid_argument("IVF cluster ID is out of range"); } if (num_points > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("IVF point count exceeds the supported ID range"); @@ -636,20 +615,22 @@ void IVF::reserve(PID cluster_id, size_t num_points) { } const Cluster& old_cluster = clusters_[cluster_id]; - const size_t num = old_cluster.num(); + const size_t old_size = old_cluster.num(); ClusterStorage grown(raw_reranking_); reserve_storage(grown, num_points); - resize_storage(grown, num); - std::copy_n(old_cluster.batch_data(), batch_data_bytes(num), grown.batch_data()); - if (rerank_vector_bytes() > 0) { - std::copy_n(old_cluster.ex_data(), ex_data_bytes(num), grown.ex_data()); + resize_storage(grown, old_size); + std::copy_n(old_cluster.batch_data(), batch_data_bytes(old_size), grown.batch_data()); + if (has_rerank_data()) { + std::copy_n( + old_cluster.ex_data(), rerank_data_bytes(old_size), grown.rerank_data() + ); } - std::copy_n(old_cluster.ids(), num, grown.ids()); + std::copy_n(old_cluster.ids(), old_size, grown.ids()); - /* Commit. Only this cluster's storage moved, so only its view changes. */ + // Commit. Only this cluster's storage moved, so only its view changes. cs = std::move(grown); clusters_[cluster_id].reset( - num, cs.batch_data(), rerank_vector_bytes() > 0 ? cs.ex_data() : nullptr, cs.ids() + old_size, cs.batch_data(), has_rerank_data() ? cs.rerank_data() : nullptr, cs.ids() ); } @@ -675,7 +656,6 @@ size_t IVF::remove(const PID* ids_to_remove, size_t num_points) { // A point whose f_add is +inf has +inf estimated and lower-bound distances in every // scan path, so it never enters the result buffer or triggers reranking. - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); size_t removed = 0; for (const auto& cluster : clusters_) { const PID* cluster_ids = cluster.ids(); @@ -684,7 +664,7 @@ size_t IVF::remove(const PID* ids_to_remove, size_t num_points) { continue; } BatchDataMap cur_batch( - cluster.batch_data() + ((pos / fastscan::kBatchSize) * batch_bytes), + cluster.batch_data() + ((pos / fastscan::kBatchSize) * batch_bytes()), padded_dim_ ); const size_t lane = pos % fastscan::kBatchSize; @@ -700,7 +680,7 @@ size_t IVF::remove(const PID* ids_to_remove, size_t num_points) { void IVF::save(const char* filename) const { if (!ready_) { - throw std::logic_error("Cannot save an unconstructed IVF index"); + throw std::logic_error("IVF index must be constructed or loaded before save"); } if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("IVF save filename must not be empty"); @@ -709,14 +689,14 @@ void IVF::save(const char* filename) const { // An HNSW initializer writes its sidecar next to the path it is given. const auto write = [this](const std::string& path) { write_file(path.c_str()); }; if (initializer_type_ == InitializerType::HNSW) { - index_detail::save_replacing(filename, {".hnsw"}, write); + index_detail::save_replacing(filename, {detail::kHnswSidecarSuffix}, write); } else { index_detail::save_replacing(filename, {}, write); } } void IVF::write_file(const char* filename) const { - std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); output.exceptions(std::ios::failbit | std::ios::badbit); const uint32_t flags = raw_reranking_ ? 1 : 0; const uint64_t stored_padded_dim = padded_dim_; @@ -732,7 +712,7 @@ void IVF::write_file(const char* filename) const { reinterpret_cast(&stored_initializer), sizeof(stored_initializer) ); - /* Save meta data */ + // Save metadata output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); output.write(reinterpret_cast(&dim_), sizeof(size_t)); output.write(reinterpret_cast(&num_clusters_), sizeof(size_t)); @@ -740,7 +720,7 @@ void IVF::write_file(const char* filename) const { output.write(reinterpret_cast(&rotator_type_), sizeof(rotator_type_)); output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - /* Save number of vectors of each cluster */ + // Save the number of points in each cluster std::vector cluster_sizes; cluster_sizes.reserve(num_clusters_); for (const auto& cluster : clusters_) { @@ -751,23 +731,23 @@ void IVF::write_file(const char* filename) const { static_cast(sizeof(size_t) * num_clusters_) ); - /* Save rotator */ - this->rotator_->save(output); + // Save rotator + rotator_->save(output); - /* Save data: initializer payload, then all batch data, all ex data and all IDs, - each in cluster order */ - this->initializer_->save(output, filename); + // Save data: initializer payload, then all batch data, all ex data and all IDs, + // each in cluster order + initializer_->save(output, filename); for (size_t i = 0; i < num_clusters_; ++i) { output.write( cluster_storage_[i].batch_data(), static_cast(batch_data_bytes(cluster_sizes[i])) ); } - if (rerank_vector_bytes() != 0) { + if (has_rerank_data()) { for (size_t i = 0; i < num_clusters_; ++i) { output.write( - cluster_storage_[i].ex_data(), - static_cast(ex_data_bytes(cluster_sizes[i])) + cluster_storage_[i].rerank_data(), + static_cast(rerank_data_bytes(cluster_sizes[i])) ); } } @@ -785,7 +765,7 @@ void IVF::load(const char* filename) { if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("IVF load filename must not be empty"); } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ifstream input(io_impl::filesystem_path(filename), std::ios::binary); if (!input.is_open()) { throw std::runtime_error("Cannot open IVF index file"); } @@ -832,7 +812,7 @@ void IVF::load(const char* filename) { input.seekg(0); } - /* Load meta data */ + // Load metadata input.read(reinterpret_cast(&loaded.max_elements_), sizeof(size_t)); input.read(reinterpret_cast(&loaded.dim_), sizeof(size_t)); input.read(reinterpret_cast(&loaded.num_clusters_), sizeof(size_t)); @@ -845,8 +825,9 @@ void IVF::load(const char* filename) { // Missing routing metadata means a historical file. Its payload layout must // follow the original threshold, independently of the current Auto policy. if (loaded.initializer_type_ == InitializerType::Auto) { - loaded.initializer_type_ = - loaded.num_clusters_ < 20000 ? InitializerType::Flat : InitializerType::HNSW; + loaded.initializer_type_ = loaded.num_clusters_ < kHistoricalHnswMinClusters + ? InitializerType::Flat + : InitializerType::HNSW; } if (loaded.max_elements_ == 0 || loaded.max_elements_ > buffer::kSearchBufferMaxPointCount || loaded.dim_ == 0 || @@ -888,12 +869,11 @@ void IVF::load(const char* filename) { if (loaded.initializer_type_ == InitializerType::FlatRaBitQ) { consume(loaded.padded_dim_, sizeof(float)); consume( - div_round_up(loaded.num_clusters_, fastscan::kBatchSize), - BatchDataMap::data_bytes(loaded.padded_dim_) + div_round_up(loaded.num_clusters_, fastscan::kBatchSize), loaded.batch_bytes() ); } - /* Load number of vectors of each cluster */ + // Load the number of points in each cluster std::vector cluster_sizes(loaded.num_clusters_, 0); input.read( reinterpret_cast(cluster_sizes.data()), @@ -906,10 +886,7 @@ void IVF::load(const char* filename) { throw std::runtime_error("Invalid cluster counts in IVF index file"); } total += size; - consume( - div_round_up(size, fastscan::kBatchSize), - BatchDataMap::data_bytes(loaded.padded_dim_) - ); + consume(div_round_up(size, fastscan::kBatchSize), loaded.batch_bytes()); } if (total != loaded.max_elements_) { throw std::runtime_error("Invalid cluster counts in IVF index file"); @@ -918,12 +895,12 @@ void IVF::load(const char* filename) { loaded.dim_, loaded.rotator_type_, loaded.padded_dim_ ); - /* Load rotator */ + // Load rotator loaded.rotator_->load(input); - /* Load data: initializer payload, then all batch data, all ex data and all IDs, - each in cluster order */ - loaded.allocate_memory(cluster_sizes); + // Load data: initializer payload, then all batch data, all ex data and all IDs, + // each in cluster order + loaded.allocate_storage(cluster_sizes); loaded.initializer_->load(input, filename); for (size_t i = 0; i < loaded.num_clusters_; ++i) { input.read( @@ -931,11 +908,11 @@ void IVF::load(const char* filename) { static_cast(loaded.batch_data_bytes(cluster_sizes[i])) ); } - if (loaded.rerank_vector_bytes() != 0) { + if (loaded.has_rerank_data()) { for (size_t i = 0; i < loaded.num_clusters_; ++i) { input.read( - loaded.cluster_storage_[i].ex_data(), - static_cast(loaded.ex_data_bytes(cluster_sizes[i])) + loaded.cluster_storage_[i].rerank_data(), + static_cast(loaded.rerank_data_bytes(cluster_sizes[i])) ); } } @@ -953,19 +930,18 @@ void IVF::load(const char* filename) { } } - /* Init each cluster */ - loaded.init_clusters(cluster_sizes); + loaded.create_cluster_views(cluster_sizes); // Removal markers are already persisted, including in historical files. // Count only logical rows: FastScan padding lanes are not stored points. - const size_t batch_bytes = BatchDataMap::data_bytes(loaded.padded_dim_); for (const auto& cluster : loaded.clusters_) { for (size_t first = 0; first < cluster.num(); first += fastscan::kBatchSize) { ConstBatchDataMap batch( - cluster.batch_data() + (first / fastscan::kBatchSize) * batch_bytes, + cluster.batch_data() + + (first / fastscan::kBatchSize) * loaded.batch_bytes(), loaded.padded_dim_ ); - const size_t count = std::min(fastscan::kBatchSize, cluster.num() - first); - for (size_t lane = 0; lane < count; ++lane) { + const size_t batch_size = std::min(fastscan::kBatchSize, cluster.num() - first); + for (size_t lane = 0; lane < batch_size; ++lane) { loaded.num_removed_ += batch.f_add()[lane] == kRemovedFAdd; } } @@ -989,7 +965,7 @@ void IVF::search( PID* __restrict__ result_ids, bool use_hacc ) const { - this->search(query, k, nprobe, result_ids, nullptr, use_hacc); + search(query, k, nprobe, result_ids, nullptr, use_hacc); } void IVF::search( @@ -1002,15 +978,13 @@ void IVF::search( ) const { validate_search(query, k, nprobe, result_ids); - nprobe = std::min(nprobe, num_clusters_); // corner case + nprobe = std::min(nprobe, num_clusters_); // Probe at most every cluster std::vector rotated_query(padded_dim_); - this->rotator_->rotate(query, rotated_query.data()); + rotator_->rotate(query, rotated_query.data()); - // use initer to get closest nprobe centroids + // Find the nprobe nearest centroids std::vector> centroid_candidates(nprobe); - this->initializer_->centroids_distances( - rotated_query.data(), nprobe, centroid_candidates - ); + initializer_->centroids_distances(rotated_query.data(), nprobe, centroid_candidates); buffer::SearchBuffer result_pool(k); @@ -1125,18 +1099,18 @@ inline void IVF::search_cluster( bool use_hacc, const float* query ) const { - size_t iter = cluster.num() / fastscan::kBatchSize; - size_t remain = cluster.num() - (iter * fastscan::kBatchSize); + const size_t full_batches = cluster.num() / fastscan::kBatchSize; + const size_t tail_size = cluster.num() - (full_batches * fastscan::kBatchSize); const char* batch_data = cluster.batch_data(); - const char* ex_data = cluster.ex_data(); + const char* rerank_data = cluster.ex_data(); const PID* ids = cluster.ids(); - /* Compute distances block by block */ - for (size_t i = 0; i < iter; ++i) { + // Compute distances batch by batch + for (size_t i = 0; i < full_batches; ++i) { scan_one_batch( batch_data, - ex_data, + rerank_data, ids, query_wrapper, result_pool, @@ -1145,28 +1119,35 @@ inline void IVF::search_cluster( query ); - batch_data += BatchDataMap::data_bytes(padded_dim_); - if (rerank_vector_bytes() > 0) { - ex_data += rerank_vector_bytes() * fastscan::kBatchSize; + batch_data += batch_bytes(); + if (has_rerank_data()) { + rerank_data += rerank_vector_bytes() * fastscan::kBatchSize; } ids += fastscan::kBatchSize; } - if (remain > 0) { - // scan the last block + if (tail_size > 0) { + // Scan the partial last batch scan_one_batch( - batch_data, ex_data, ids, query_wrapper, result_pool, remain, use_hacc, query + batch_data, + rerank_data, + ids, + query_wrapper, + result_pool, + tail_size, + use_hacc, + query ); } } inline void IVF::scan_one_batch( const char* batch_data, - const char* ex_data, + const char* rerank_data, const PID* ids, const SplitBatchQuery& query_wrapper, buffer::SearchBuffer& result_pool, - size_t num_points, + size_t batch_size, bool use_hacc, const float* query ) const { @@ -1187,32 +1168,32 @@ inline void IVF::scan_one_batch( float dist_k = result_pool.top_dist(); // Without reranking data, return the one-bit estimates directly. - if (ex_bits_ == 0 && !raw_reranking_) { - detail::insert_candidates(result_pool, ids, estimated_distances.data(), num_points); + if (!has_rerank_data()) { + detail::insert_candidates(result_pool, ids, estimated_distances.data(), batch_size); return; } // Refine candidates passing the probabilistic one-bit filter. Quantized storage // reuses ip_x0_qr with the extra bits; raw storage scores the original vector. - for (size_t i = 0; i < num_points; ++i) { + for (size_t i = 0; i < batch_size; ++i) { float lower_bound = lower_bounds[i]; if (lower_bound < dist_k) { PID id = ids[i]; float refined_distance; if (raw_reranking_) { - const auto* vector = reinterpret_cast(ex_data); + const auto* vector = reinterpret_cast(rerank_data); refined_distance = metric_type_ == METRIC_L2 ? euclidean_sqr(query, vector, dim_) : dot_product_dis(query, vector, dim_); } else { refined_distance = split_distance_boosting( - ex_data, ip_func_, query_wrapper, padded_dim_, ex_bits_, ip_x0_qr[i] + rerank_data, ip_func_, query_wrapper, padded_dim_, ex_bits_, ip_x0_qr[i] ); } result_pool.insert(id, refined_distance); dist_k = result_pool.top_dist(); } - ex_data += rerank_vector_bytes(); + rerank_data += rerank_vector_bytes(); } } diff --git a/src/index/ivf_initializer.cpp b/src/index/ivf_initializer.cpp index 9aa5e34b..cbaf7e00 100644 --- a/src/index/ivf_initializer.cpp +++ b/src/index/ivf_initializer.cpp @@ -24,24 +24,26 @@ namespace rabitqlib::ivf { FlatRaBitQInitializer::FlatRaBitQInitializer( - size_t dim, size_t num_clusters, MetricType metric + size_t dim, size_t num_clusters, MetricType metric_type ) - : Initializer(dim, num_clusters), metric_type_(metric) { - validate_metric_type(metric); + : Initializer(dim, num_clusters), metric_type_(metric_type) { + validate_metric_type(metric_type); const size_t max_size = std::numeric_limits::max(); if (dim == 0 || dim % 32 != 0 || dim > max_size / 32 || num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount || num_clusters > max_size / sizeof(float) / dim) { - throw std::invalid_argument("Invalid Flat RaBitQ initializer dimensions"); + throw std::invalid_argument("IVF Flat RaBitQ initializer dimensions are invalid"); } - const size_t batches = div_round_up(num_clusters, fastscan::kBatchSize); + const size_t num_batches = div_round_up(num_clusters, fastscan::kBatchSize); const size_t batch_bytes = BatchDataMap::data_bytes(dim); - if (batches > max_size / batch_bytes) { - throw std::invalid_argument("Flat RaBitQ initializer storage overflow"); + if (num_batches > max_size / batch_bytes) { + throw std::invalid_argument( + "IVF Flat RaBitQ initializer storage exceeds addressable memory" + ); } centroids_.resize(num_clusters * dim); global_centroid_.resize(dim); - batch_data_.resize(batches * batch_bytes); + batch_data_.resize(num_batches * batch_bytes); } const float* FlatRaBitQInitializer::centroid(PID id) const { @@ -58,11 +60,11 @@ void FlatRaBitQInitializer::add_vectors(const float* centroids, size_t num_threa global_centroid_[d] = static_cast(sum / static_cast(num_cluster_)); } const size_t batch_bytes = BatchDataMap::data_bytes(dim_); - rabitqlib::index_detail::parallel_for( + index_detail::parallel_for( 0, div_round_up(num_cluster_, fastscan::kBatchSize), num_threads, - [&](size_t batch, size_t) { + [&](size_t batch, size_t /*thread*/) { const size_t first = batch * fastscan::kBatchSize; quant::quantize_split_batch( centroids_.data() + first * dim_, @@ -81,17 +83,17 @@ void FlatRaBitQInitializer::add_vectors(const float* centroids, size_t num_threa void FlatRaBitQInitializer::centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const { - const size_t nprobed = std::min(nprobe, num_cluster_); - candidates.resize(nprobed); - if (nprobed == 0) { + const size_t num_probes = std::min(nprobe, num_cluster_); + candidates.resize(num_probes); + if (num_probes == 0) { return; } const auto exact_distance = [&](PID id) { return metric_type_ == METRIC_L2 ? euclidean_sqr(query, centroid(id), dim_) : dot_product_dis(query, centroid(id), dim_); }; - if (nprobed == num_cluster_) { - for (size_t i = 0; i < nprobed; ++i) { + if (num_probes == num_cluster_) { + for (size_t i = 0; i < num_probes; ++i) { candidates[i] = AnnCandidate(static_cast(i), exact_distance(i)); } std::sort(candidates.begin(), candidates.end()); @@ -103,9 +105,13 @@ void FlatRaBitQInitializer::centroids_distances( metric_type_ == METRIC_IP ? dot_product(query, global_centroid_.data(), dim_) : 0.0F ); - std::vector lower(num_cluster_); - std::vector> upper(num_cluster_); - std::array estimates{}, lows{}, ips{}; + // Use twice the standard RaBitQ error margin for centroid routing: + // missing a centroid also skips all vectors in its IVF list. + constexpr float kRoutingErrorScale = 2.0F; + std::vector cluster_lower_bounds(num_cluster_); + std::vector> cluster_upper_bounds(num_cluster_); + std::array estimated_distances{}, lower_bounds{}, + ip_x0_qr{}; const size_t batch_bytes = BatchDataMap::data_bytes(dim_); for (size_t first = 0; first < num_cluster_; first += fastscan::kBatchSize) { split_batch_estdist( @@ -113,37 +119,42 @@ void FlatRaBitQInitializer::centroids_distances( (first / fastscan::kBatchSize) * batch_bytes, query_wrapper, dim_, - estimates.data(), - lows.data(), - ips.data(), + estimated_distances.data(), + lower_bounds.data(), + ip_x0_qr.data(), false ); - const size_t count = std::min(fastscan::kBatchSize, num_cluster_ - first); - for (size_t j = 0; j < count; ++j) { - // Use twice the standard RaBitQ error margin for centroid routing: - // missing a centroid also skips all vectors in its IVF list. - lower[first + j] = estimates[j] - 2.0F * (estimates[j] - lows[j]); - upper[first + j] = AnnCandidate( - static_cast(first + j), 2 * estimates[j] - lower[first + j] + const size_t batch_size = std::min(fastscan::kBatchSize, num_cluster_ - first); + for (size_t j = 0; j < batch_size; ++j) { + cluster_lower_bounds[first + j] = + estimated_distances[j] - + kRoutingErrorScale * (estimated_distances[j] - lower_bounds[j]); + cluster_upper_bounds[first + j] = AnnCandidate( + static_cast(first + j), + 2 * estimated_distances[j] - cluster_lower_bounds[first + j] ); } } std::nth_element( - upper.begin(), upper.begin() + static_cast(nprobed), upper.end() + cluster_upper_bounds.begin(), + cluster_upper_bounds.begin() + static_cast(num_probes), + cluster_upper_bounds.end() ); - buffer::SearchBuffer result_pool(nprobed); + buffer::SearchBuffer result_pool(num_probes); // Seed with exact distances: an estimated upper bound need not hold for // every vector. Always fill the result before pruning other candidates. - for (size_t i = 0; i < nprobed; ++i) { - result_pool.insert(upper[i].id, exact_distance(upper[i].id)); + for (size_t i = 0; i < num_probes; ++i) { + result_pool.insert( + cluster_upper_bounds[i].id, exact_distance(cluster_upper_bounds[i].id) + ); } - for (size_t i = nprobed; i < num_cluster_; ++i) { - const PID id = upper[i].id; - if (lower[id] <= result_pool.top_dist()) { + for (size_t i = num_probes; i < num_cluster_; ++i) { + const PID id = cluster_upper_bounds[i].id; + if (cluster_lower_bounds[id] <= result_pool.top_dist()) { result_pool.insert(id, exact_distance(id)); } } - std::copy_n(result_pool.data().begin(), nprobed, candidates.begin()); + std::copy_n(result_pool.data().begin(), num_probes, candidates.begin()); } if (metric_type_ == METRIC_L2) { for (auto& candidate : candidates) { @@ -152,7 +163,7 @@ void FlatRaBitQInitializer::centroids_distances( } } -void FlatRaBitQInitializer::save(std::ofstream& output, const char*) const { +void FlatRaBitQInitializer::save(std::ofstream& output, const char* /*filename*/) const { output.write( reinterpret_cast(centroids_.data()), static_cast(centroids_.size() * sizeof(float)) @@ -166,11 +177,11 @@ void FlatRaBitQInitializer::save(std::ofstream& output, const char*) const { static_cast(batch_data_.size()) ); if (!output) { - throw std::runtime_error("Cannot write Flat RaBitQ initializer"); + throw std::runtime_error("IVF Flat RaBitQ initializer cannot be written"); } } -void FlatRaBitQInitializer::load(std::ifstream& input, const char*) { +void FlatRaBitQInitializer::load(std::ifstream& input, const char* /*filename*/) { input.read( reinterpret_cast(centroids_.data()), static_cast(centroids_.size() * sizeof(float)) @@ -184,7 +195,7 @@ void FlatRaBitQInitializer::load(std::ifstream& input, const char*) { static_cast(batch_data_.size()) ); if (!input) { - throw std::runtime_error("Cannot read Flat RaBitQ initializer"); + throw std::runtime_error("IVF Flat RaBitQ initializer cannot be read"); } } } // namespace rabitqlib::ivf diff --git a/src/index/ivf_search.hpp b/src/index/ivf_search.hpp index 74c302f5..bb495671 100644 --- a/src/index/ivf_search.hpp +++ b/src/index/ivf_search.hpp @@ -6,5 +6,10 @@ #include "rabitqlib/index/types.hpp" namespace rabitqlib::ivf::detail { -void insert_candidates(buffer::SearchBuffer&, const PID*, const float*, size_t); +void insert_candidates( + buffer::SearchBuffer& result_pool, + const PID* ids, + const float* distances, + size_t count +); } // namespace rabitqlib::ivf::detail diff --git a/src/index/pipnn.cpp b/src/index/pipnn.cpp index dc8d66fd..82b84224 100644 --- a/src/index/pipnn.cpp +++ b/src/index/pipnn.cpp @@ -28,12 +28,17 @@ namespace rabitqlib::symqg::detail { -namespace pipnn_impl { +namespace { constexpr size_t kTileSize = 2048; constexpr size_t kMinTileSize = 64; constexpr size_t kHashBits = 12; constexpr size_t kLocalDegree = 2; constexpr float kAlpha = 1.1F; +// Root seed for leader sampling, duplicate shuffles and hash projections. +constexpr uint32_t kPipnnSeed = 555; +// Numerical Recipes LCG constants that derive each child bucket's seed. +constexpr uint32_t kLcgMultiplier = 1664525U; +constexpr uint32_t kLcgIncrement = 1013904223U; struct Candidate { PID id; @@ -45,12 +50,12 @@ struct Candidate { // collection finishes, each row is pruned once and reused to hold only selected IDs. class CandidateTable { public: - CandidateTable(size_t count, size_t capacity) + CandidateTable(size_t num_points, size_t capacity) : capacity_(capacity) - , table_(count * capacity) - , sizes_(count, 0) - , worst_(count, 0) - , locks_(count) {} + , table_(num_points * capacity) + , sizes_(num_points, 0) + , worst_(num_points, 0) + , locks_(num_points) {} void merge( PID source, PID target, float distance, const RowMajorMatrix& sketches @@ -151,25 +156,35 @@ class CandidateTable { struct BucketJob { Bucket ids; size_t depth = 0; - uint32_t seed = 555; + uint32_t seed = kPipnnSeed; }; -ScratchPool::ScratchPool(size_t count, size_t dim, size_t num_workers) - : points_buffer_size(std::min(count, kTileSize) * dim) +// Copy the rows named by ids[0..num_points) to consecutive rows of output. +void copy_rows( + const float* data, size_t dim, const PID* ids, size_t num_points, float* output +) { + for (size_t i = 0; i < num_points; ++i) { + std::copy_n(data + ids[i] * dim, dim, output + i * dim); + } +} +} // namespace + +ScratchPool::ScratchPool(size_t num_points, size_t dim, size_t num_workers) + : points_buffer_size(std::min(num_points, kTileSize) * dim) , distances_buffer_size( std::max( - std::min(count, kTileSize) * std::min(count, size_t{512}), - std::min(count, kLeafSize) * std::min(count, kLeafSize) + std::min(num_points, kTileSize) * std::min(num_points, size_t{512}), + std::min(num_points, kLeafSize) * std::min(num_points, kLeafSize) ) ) , worker_buffer_size( - points_buffer_size + distances_buffer_size + std::min(count, kTileSize) + points_buffer_size + distances_buffer_size + std::min(num_points, kTileSize) ) // Preserve uninitialized allocation and worker first-touch placement. , storage(new float[num_workers * worker_buffer_size]) {} -Workspace ScratchPool::worker(size_t id) { - float* base = storage.get() + id * worker_buffer_size; +Workspace ScratchPool::worker(size_t worker_id) { + float* base = storage.get() + worker_id * worker_buffer_size; return { base, base + points_buffer_size, @@ -177,22 +192,13 @@ Workspace ScratchPool::worker(size_t id) { {}}; } -namespace { -// Copy the rows named by ids[0..count) to consecutive rows of output. -void copy_rows(const float* data, size_t dim, const PID* ids, size_t count, float* output) { - for (size_t i = 0; i < count; ++i) { - std::copy_n(data + ids[i] * dim, dim, output + i * dim); - } -} -} // namespace - void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { copy_rows(data, dim, ids.data(), ids.size(), work.points); } -void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric_type) { +void pairwise(Workspace& work, size_t num_points, size_t dim, MetricType metric_type) { simd::pairwise_distances_lower( - work.points, work.distances, work.norms, size, dim, metric_type == METRIC_IP + work.points, work.distances, work.norms, num_points, dim, metric_type == METRIC_IP ); } @@ -223,6 +229,7 @@ void select_local_neighbors(const Bucket& ids, Workspace& work) { } } +namespace { std::vector partition( const float* data, size_t dim, @@ -263,23 +270,30 @@ std::vector partition( // Both callers run inside a PiPNN OpenMP task, including one-thread teams. // Keep Eigen within that worker without changing the caller's thread budget. omp_set_num_threads(1); - const size_t count = std::min(tile_size, ids.size() - begin); + const size_t batch_size = std::min(tile_size, ids.size() - begin); RowMajorMatrixMap points( - work.points, static_cast(count), static_cast(dim) + work.points, + static_cast(batch_size), + static_cast(dim) ); RowMajorMatrixMap distances( work.distances, - static_cast(count), + static_cast(batch_size), static_cast(leader_count) ); - VectorMap norms(work.norms, static_cast(count)); - copy_rows(data, dim, ids.data() + begin, count, points.data()); + VectorMap norms(work.norms, static_cast(batch_size)); + copy_rows(data, dim, ids.data() + begin, batch_size, points.data()); simd::matrix_product_transposed( - points.data(), leader_data.data(), distances.data(), count, dim, leader_count + points.data(), + leader_data.data(), + distances.data(), + batch_size, + dim, + leader_count ); - simd::row_norms(points.data(), norms.data(), count, dim); + simd::row_norms(points.data(), norms.data(), batch_size, dim); work.candidates.resize(leader_count); - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < batch_size; ++i) { for (size_t j = 0; j < leader_count; ++j) { const float dot = distances(static_cast(i), static_cast(j)); @@ -291,7 +305,7 @@ std::vector partition( } std::partial_sort( work.candidates.begin(), - work.candidates.begin() + static_cast(fanout), + work.candidates.begin() + static_cast(fanout), work.candidates.end() ); for (size_t j = 0; j < fanout; ++j) { @@ -300,26 +314,24 @@ std::vector partition( } }; if (job.depth == 0) { - { - index_detail::ParallelException failure; + index_detail::ParallelException failure; #pragma omp parallel num_threads(std::min(num_workers, div_round_up(ids.size(), tile_size))) - { - auto work = scratch.worker(static_cast(omp_get_thread_num())); + { + auto work = scratch.worker(static_cast(omp_get_thread_num())); #pragma omp for schedule(static) - for (std::ptrdiff_t offset = 0; - offset < static_cast(ids.size()); - offset += static_cast(tile_size)) { - if (failure.failed()) { - continue; - } - try { - const size_t begin = static_cast(offset); - assign_tile(begin, work); - } catch (...) { failure.capture(); } + for (std::ptrdiff_t offset = 0; + offset < static_cast(ids.size()); + offset += static_cast(tile_size)) { + if (failure.failed()) { + continue; } + try { + const size_t begin = static_cast(offset); + assign_tile(begin, work); + } catch (...) { failure.capture(); } } - failure.rethrow(); } + failure.rethrow(); } else { // Already inside the bucket-level team. Eigen must see that team rather // than a serialized nested region when deciding its own parallelism. @@ -336,20 +348,21 @@ std::vector partition( } return children; } +} // namespace std::vector cluster( const float* data, - size_t count, + size_t num_points, size_t dim, MetricType metric_type, size_t num_workers, ScratchPool& scratch ) { - Bucket all(count); + Bucket all(num_points); std::iota(all.begin(), all.end(), PID{0}); // The root parallelizes tiles; deeper levels parallelize independent buckets. std::vector active; - active.push_back({std::move(all), 0, 555}); + active.push_back({std::move(all), 0, kPipnnSeed}); std::vector leaves; while (!active.empty()) { std::vector> next(active.size()); @@ -394,8 +407,8 @@ std::vector cluster( std::shuffle(child.begin(), child.end(), random); for (size_t begin = 0; begin < child.size(); begin += kLeafSize) { finished[i].emplace_back( - child.begin() + static_cast(begin), - child.begin() + static_cast( + child.begin() + static_cast(begin), + child.begin() + static_cast( std::min(begin + kLeafSize, child.size()) ) ); @@ -404,7 +417,8 @@ std::vector cluster( next[i].push_back( {std::move(child), job.depth + 1, - job.seed * 1664525U + static_cast(j) + 1013904223U} + job.seed * kLcgMultiplier + static_cast(j) + + kLcgIncrement} ); } } @@ -437,59 +451,57 @@ std::vector cluster( } return leaves; } -} // namespace pipnn_impl // Native PiPNN-style seeding: overlapping leader partitions, local dense kNN, // directional hash reservoirs, and alpha pruning. Uses existing Eigen/OpenMP. InitialGraph build_initial_graph( const float* data, - size_t count, + size_t num_points, size_t dim, - size_t degree, + size_t max_degree, MetricType metric_type, size_t num_threads ) { validate_metric_type(metric_type); - if (data == nullptr || dim == 0 || degree == 0 || degree >= count || - count > buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { + if (data == nullptr || dim == 0 || max_degree == 0 || max_degree >= num_points || + num_points > buffer::kSearchBufferMaxPointCount || max_degree % 32 != 0) { throw std::invalid_argument( "PiPNN seed requires data, positive dim, and a degree multiple of 32 below " "count" ); } - using namespace pipnn_impl; const size_t num_workers = resolve_num_threads(num_threads); - ScratchPool scratch(count, dim, num_workers); - auto leaves = cluster(data, count, dim, metric_type, num_workers, scratch); + ScratchPool scratch(num_points, dim, num_workers); + auto leaves = cluster(data, num_points, dim, metric_type, num_workers, scratch); RowMajorMatrix projections(dim, kHashBits); - std::mt19937 random(555); + std::mt19937 random(kPipnnSeed); std::normal_distribution normal; for (Eigen::Index i = 0; i < projections.size(); ++i) { projections.data()[i] = normal(random); } - RowMajorMatrix sketches(count, kHashBits); + RowMajorMatrix sketches(num_points, kHashBits); const size_t sketch_tile = std::min( - kTileSize, round_up_to_multiple(div_round_up(count, num_workers), kMinTileSize) + kTileSize, round_up_to_multiple(div_round_up(num_points, num_workers), kMinTileSize) ); { index_detail::ParallelException failure; -#pragma omp parallel for num_threads( \ - std::min(num_workers, div_round_up(count, sketch_tile)) \ +#pragma omp parallel for num_threads( \ + std::min(num_workers, div_round_up(num_points, sketch_tile)) \ ) schedule(static) - for (std::ptrdiff_t offset = 0; offset < static_cast(count); + for (std::ptrdiff_t offset = 0; offset < static_cast(num_points); offset += static_cast(sketch_tile)) { if (failure.failed()) { continue; } try { const size_t begin = static_cast(offset); - const size_t rows = std::min(sketch_tile, count - begin); + const size_t batch_size = std::min(sketch_tile, num_points - begin); omp_set_num_threads(1); simd::matrix_product( data + begin * dim, projections.data(), sketches.data() + begin * kHashBits, - rows, + batch_size, dim, kHashBits ); @@ -497,7 +509,7 @@ InitialGraph build_initial_graph( } failure.rethrow(); } - CandidateTable table(count, degree * 5 / 2); + CandidateTable table(num_points, max_degree * 5 / 2); { index_detail::ParallelException failure; #pragma omp parallel num_threads(std::min(num_workers, leaves.size())) @@ -546,7 +558,7 @@ InitialGraph build_initial_graph( std::vector candidates; Bucket selected; #pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(count); + for (std::ptrdiff_t index = 0; index < static_cast(num_points); ++index) { if (failure.failed()) { continue; @@ -585,7 +597,7 @@ InitialGraph build_initial_graph( if (!occluded) { selected.push_back(candidate.id); } - if (selected.size() == degree) { + if (selected.size() == max_degree) { break; } } diff --git a/src/index/pipnn_internal.hpp b/src/index/pipnn_internal.hpp index 99471c97..b4589b0c 100644 --- a/src/index/pipnn_internal.hpp +++ b/src/index/pipnn_internal.hpp @@ -9,8 +9,8 @@ #include "rabitqlib/index/types.hpp" // Private declarations shared only by the implementation and reference tests. -namespace rabitqlib::symqg::detail::pipnn_impl { -constexpr size_t kLeafSize = 1024; +namespace rabitqlib::symqg::detail { +inline constexpr size_t kLeafSize = 1024; using Bucket = std::vector; // Non-owning slices of a build-scoped allocation. Large matrices are released @@ -29,15 +29,22 @@ struct ScratchPool { size_t worker_buffer_size; std::unique_ptr storage; - ScratchPool(size_t count, size_t dim, size_t num_workers); - Workspace worker(size_t id); + ScratchPool(size_t num_points, size_t dim, size_t num_workers); + Workspace worker(size_t worker_id); }; -void gather(const float*, size_t, const Bucket&, Workspace&); -void pairwise(Workspace&, size_t, size_t, MetricType); +void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work); +void pairwise(Workspace& work, size_t num_points, size_t dim, MetricType metric_type); // After pairwise(), write two sorted (distance, ID) candidates per local row. // Excludes self; a missing second neighbor retains kPidMax. -void select_local_neighbors(const Bucket&, Workspace&); -std::vector cluster(const float*, size_t, size_t, MetricType, size_t, ScratchPool&); - -} // namespace rabitqlib::symqg::detail::pipnn_impl +void select_local_neighbors(const Bucket& ids, Workspace& work); +std::vector cluster( + const float* data, + size_t num_points, + size_t dim, + MetricType metric_type, + size_t num_workers, + ScratchPool& scratch +); + +} // namespace rabitqlib::symqg::detail diff --git a/src/index/qg.cpp b/src/index/qg.cpp index b3da8f70..1aeb5551 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -82,7 +82,7 @@ size_t QuantizedGraph::padded_dimension(size_t dim) { return (checked_add(dim, 31) / 32) * 32; } -size_t QuantizedGraph::num_vertices() const { return this->max_elements_; } +size_t QuantizedGraph::num_vertices() const { return max_elements_; } size_t QuantizedGraph::max_elements() const { return max_elements_; } @@ -92,15 +92,15 @@ size_t QuantizedGraph::num_stored_points() const { return ready_ ? max_elements_ : 0; } -size_t QuantizedGraph::dimension() const { return this->dim_; } +size_t QuantizedGraph::dimension() const { return dim_; } -size_t QuantizedGraph::degree_bound() const { return this->max_degree_; } +size_t QuantizedGraph::degree_bound() const { return max_degree_; } -PID QuantizedGraph::entry_point() const { return this->entry_point_id_; } +PID QuantizedGraph::entry_point() const { return entry_point_id_; } -MetricType QuantizedGraph::metric_type() const { return this->metric_type_; } +MetricType QuantizedGraph::metric_type() const { return metric_type_; } -size_t QuantizedGraph::quantization_bits() const { return this->quantization_bits_; } +size_t QuantizedGraph::quantization_bits() const { return quantization_bits_; } bool QuantizedGraph::is_quantized() const { return quantization_bits_ != 0; } @@ -180,58 +180,58 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_workers) const int omp_num_threads = static_cast(num_workers); if (quantization_bits_ != 0) { if (centroid_.size() != padded_dim_) { - throw std::logic_error("qg-quant centroid must be set before copying vectors"); + throw std::logic_error( + "QuantizedGraph quantization centroid must be set before copying vectors" + ); } - { - index_detail::ParallelException failure; + index_detail::ParallelException failure; #pragma omp parallel num_threads(omp_num_threads) - { - std::vector rotated_data; - std::vector quantized_data; - try { - rotated_data.resize(padded_dim_); - quantized_data.resize(padded_dim_); - } catch (...) { failure.capture(); } + { + std::vector rotated_data; + std::vector quantized_data; + try { + rotated_data.resize(padded_dim_); + quantized_data.resize(padded_dim_); + } catch (...) { failure.capture(); } #pragma omp for schedule(static) - for (std::ptrdiff_t index = 0; - index < static_cast(max_elements_); - ++index) { - if (failure.failed()) { - continue; - } - try { - const size_t i = static_cast(index); - rotator_->rotate(data + (dim_ * i), rotated_data.data()); - ExDataMap output( - get_quantized_vector(i), padded_dim_, quantization_bits_ - ); - float f_add; - float f_rescale; - float unused_f_error = 0; - quant::quantize_full_single( - rotated_data.data(), - centroid_.data(), - padded_dim_, - quantization_bits_, - quantized_data.data(), - f_add, - f_rescale, - unused_f_error, - metric_type_ - ); - output.f_add_ex() = f_add; - output.f_rescale_ex() = f_rescale; - quant::rabitq_impl::ex_bits::packing_rabitqplus_code( - quantized_data.data(), - output.ex_code(), - padded_dim_, - quantization_bits_ - ); - } catch (...) { failure.capture(); } + for (std::ptrdiff_t index = 0; + index < static_cast(max_elements_); + ++index) { + if (failure.failed()) { + continue; } + try { + const size_t i = static_cast(index); + rotator_->rotate(data + (dim_ * i), rotated_data.data()); + ExDataMap output( + get_quantized_vector(i), padded_dim_, quantization_bits_ + ); + float f_add; + float f_rescale; + float unused_f_error = 0; + quant::quantize_full_single( + rotated_data.data(), + centroid_.data(), + padded_dim_, + quantization_bits_, + quantized_data.data(), + f_add, + f_rescale, + unused_f_error, + metric_type_ + ); + output.f_add_ex() = f_add; + output.f_rescale_ex() = f_rescale; + quant::rabitq_impl::ex_bits::packing_rabitqplus_code( + quantized_data.data(), + output.ex_code(), + padded_dim_, + quantization_bits_ + ); + } catch (...) { failure.capture(); } } - failure.rethrow(); } + failure.rethrow(); return; } #pragma omp parallel for schedule(static) num_threads(omp_num_threads) @@ -265,9 +265,9 @@ void QuantizedGraph::save(const char* filename) const { } void QuantizedGraph::write_file(const char* filename) const { - std::ofstream output(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ofstream output(io_impl::filesystem_path(filename), std::ios::binary); if (!output.is_open()) { - throw std::runtime_error("Cannot open quantized graph file for writing"); + throw std::runtime_error("QuantizedGraph cannot open file for writing"); } output.exceptions(std::ios::badbit | std::ios::failbit); @@ -278,7 +278,7 @@ void QuantizedGraph::write_file(const char* filename) const { ); } - /* Basic variants */ + // Basic variants output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); output.write(reinterpret_cast(&max_degree_), sizeof(size_t)); output.write(reinterpret_cast(&dim_), sizeof(size_t)); @@ -296,14 +296,14 @@ void QuantizedGraph::write_file(const char* filename) const { ); } - /* Data */ + // Data output.write( get_row_data(0), static_cast(checked_multiply(max_elements_, record_bytes_)) ); - /* Rotator */ - this->rotator_->save(output); + // Rotator + rotator_->save(output); output.flush(); output.close(); @@ -313,14 +313,13 @@ void QuantizedGraph::load(const char* filename) { if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("QuantizedGraph load filename must not be empty"); } - /* Check existence */ if (!file_exists(filename)) { - throw std::runtime_error("Quantized graph file does not exist"); + throw std::runtime_error("QuantizedGraph file does not exist"); } - std::ifstream input(rabitqlib::io_impl::filesystem_path(filename), std::ios::binary); + std::ifstream input(io_impl::filesystem_path(filename), std::ios::binary); if (!input.is_open()) { - throw std::runtime_error("Cannot open quantized graph file"); + throw std::runtime_error("QuantizedGraph cannot open file"); } auto read_exact = [&](void* destination, size_t bytes, const char* field) { @@ -332,7 +331,7 @@ void QuantizedGraph::load(const char* filename) { ); if (!input) { throw std::runtime_error( - std::string("Truncated QuantizedGraph file while reading ") + field + std::string("QuantizedGraph file is truncated while reading ") + field ); } }; @@ -346,7 +345,7 @@ void QuantizedGraph::load(const char* filename) { uint32_t version = 0; read_value(version, "format version"); if (version != kFormatVersion) { - throw std::runtime_error("Unsupported QuantizedGraph file version"); + throw std::runtime_error("QuantizedGraph file version is unsupported"); } } else { // Files produced before qg-quant have no header and always contain raw vectors. @@ -371,7 +370,7 @@ void QuantizedGraph::load(const char* filename) { loaded.validate_configuration(); if (stored_padded_dim < loaded.dim_ || stored_padded_dim % 32 != 0) { - throw std::runtime_error("Invalid padded dimension in quantized graph file"); + throw std::runtime_error("QuantizedGraph file has an invalid padded dimension"); } // Older files use 64-coordinate padding. Their saved rotation and codes must // retain that dimension, including when initialize() allocates the row storage. @@ -393,13 +392,13 @@ void QuantizedGraph::load(const char* filename) { const auto payload_position = input.tellg(); if (payload_position < 0) { - throw std::runtime_error("Cannot determine QuantizedGraph payload position"); + throw std::runtime_error("QuantizedGraph cannot determine payload position"); } const size_t file_size = get_filesize(filename); const size_t payload_offset = static_cast(payload_position); if (payload_offset > file_size || file_size - payload_offset != expected_payload_bytes) { - throw std::runtime_error("Invalid QuantizedGraph payload size"); + throw std::runtime_error("QuantizedGraph payload size is invalid"); } loaded.initialize(false); @@ -415,14 +414,14 @@ void QuantizedGraph::load(const char* filename) { const auto neighbors = loaded.get_neighbors(source); for (size_t i = 0; i < loaded.max_degree_; ++i) { if (neighbors[i] >= loaded.max_elements_) { - throw std::runtime_error("Invalid QuantizedGraph neighbor ID"); + throw std::runtime_error("QuantizedGraph neighbor ID is invalid"); } } } loaded.rotator_->load(input); if (!input) { - throw std::runtime_error("Truncated QuantizedGraph file while reading rotator"); + throw std::runtime_error("QuantizedGraph file is truncated while reading rotator"); } input.close(); @@ -437,7 +436,7 @@ void QuantizedGraph::set_ef(size_t ef_search) { if (ef_search == 0) { throw std::invalid_argument("QuantizedGraph ef must be positive"); } - this->default_ef_search_ = ef_search; + default_ef_search_ = ef_search; } void QuantizedGraph::search( @@ -558,26 +557,26 @@ void QuantizedGraph::validate_search( } float QuantizedGraph::point_distance( - const float* raw_query, const QuantizedQuery* quantized_query, PID data_id + const float* raw_query, const QuantizedQuery* quantized_query, PID internal_id ) const { if (quantized_query != nullptr) { - return quantized_distance(*quantized_query, data_id); + return quantized_distance(*quantized_query, internal_id); } - return raw_dist_func_(raw_query, get_vector(data_id), dim_); + return raw_dist_func_(raw_query, get_vector(internal_id), dim_); } // Scan a data row and store estimated neighbor distances. The caller scores the current // vertex from either its raw vector (vanilla QG) or its 4/8-bit code (qg-quant). void QuantizedGraph::scan_neighbors( const BatchQuery& query_wrapper, - PID data_id, + PID internal_id, float* estimated_distances, buffer::SearchBuffer& search_pool, VisitedSet& visited, size_t num_neighbors ) const { - const auto* batch_data = get_batch_data(data_id); - const auto neighbors = get_neighbors(data_id); + const auto* batch_data = get_batch_data(internal_id); + const auto neighbors = get_neighbors(internal_id); for (size_t begin = 0; begin < num_neighbors; begin += fastscan::kBatchSize) { const float threshold = search_pool.top_dist(); uint32_t candidate_mask = simd::qg_batch_estdist_mask( @@ -619,11 +618,12 @@ void QuantizedGraph::update_results( const auto& pool_data = result_pool.data(); const std::vector> data( - pool_data.begin(), pool_data.begin() + static_cast(result_pool.size()) + pool_data.begin(), + pool_data.begin() + static_cast(result_pool.size()) ); for (const auto& record : data) { auto neighbors = get_neighbors(record.id); - for (uint32_t i = 0; i < this->max_degree_; ++i) { + for (size_t i = 0; i < max_degree_; ++i) { PID neighbor_id = neighbors[i]; if (!visited.get(neighbor_id)) { visited.set(neighbor_id); @@ -678,10 +678,10 @@ void QuantizedGraph::initialize(bool randomize_rotation) { } float QuantizedGraph::quantized_distance( - const QuantizedQuery& query, PID data_id + const QuantizedQuery& query, PID internal_id ) const { ConstExDataMap data( - get_quantized_vector(data_id), padded_dim_, quantization_bits_ + get_quantized_vector(internal_id), padded_dim_, quantization_bits_ ); return quant::full_est_dist( data.ex_code(), @@ -697,10 +697,10 @@ float QuantizedGraph::quantized_distance( } void QuantizedGraph::reconstruct_quantized_vector( - PID data_id, float* reconstructed + PID internal_id, float* reconstructed ) const { ConstExDataMap data( - get_quantized_vector(data_id), padded_dim_, quantization_bits_ + get_quantized_vector(internal_id), padded_dim_, quantization_bits_ ); std::vector quantized_data(padded_dim_); quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code( @@ -719,28 +719,28 @@ void QuantizedGraph::reconstruct_quantized_vector( // Encode a logical neighbor-ID row, including physical storage for tail batches. void QuantizedGraph::encode_neighbors( - PID cur_id, + PID internal_id, const std::vector& new_neighbors, const float* rotated_data, - const float* rotated_centroid + const float* rotated_source ) { size_t num_neighbors = new_neighbors.size(); if (num_neighbors == 0) { return; } - auto neighbor_ptr = get_neighbors(cur_id); + auto neighbor_ptr = get_neighbors(internal_id); for (size_t i = 0; i < num_neighbors; ++i) { neighbor_ptr[i] = new_neighbors[i]; } - // quantize batches for current vertex - auto* batch_data = get_batch_data(cur_id); + // Quantize neighbor batches against the source vertex. + auto* batch_data = get_batch_data(internal_id); for (size_t i = 0; i < num_neighbors; i += fastscan::kBatchSize) { const size_t batch_size = std::min(num_neighbors - i, fastscan::kBatchSize); quant::quantize_qg_batch( rotated_data + (i * padded_dim_), - rotated_centroid, + rotated_source, batch_size, padded_dim_, batch_data, @@ -749,28 +749,29 @@ void QuantizedGraph::encode_neighbors( if (batch_size < fastscan::kBatchSize) { QGBatchDataMap batch(batch_data, padded_dim_); std::fill( - batch.f_add() + static_cast(batch_size), + batch.f_add() + static_cast(batch_size), batch.f_add() + fastscan::kBatchSize, - static_cast(0) + 0.0F ); std::fill( - batch.f_rescale() + static_cast(batch_size), + batch.f_rescale() + static_cast(batch_size), batch.f_rescale() + fastscan::kBatchSize, - static_cast(0) + 0.0F ); } batch_data += QGBatchDataMap::data_bytes(padded_dim_); } } + void QuantizedGraph::search_with_scratch( const float* __restrict__ query, uint32_t k, uint32_t* __restrict__ result_ids, float* __restrict__ distances, - float* __restrict__ rotated_query_scratch, - float* __restrict__ est_dist_scratch, - float* __restrict__ lut_float_scratch, + float* __restrict__ rotated_query, + float* __restrict__ estimated_distances, + float* __restrict__ lookup_table, BatchQuery& batch_query, buffer::SearchBuffer& search_pool, buffer::SearchBuffer& result_pool, @@ -784,9 +785,9 @@ void QuantizedGraph::search_with_scratch( k, result_ids, distances, - {rotated_query_scratch, - est_dist_scratch, - lut_float_scratch, + {rotated_query, + estimated_distances, + lookup_table, batch_query, search_pool, result_pool, @@ -824,10 +825,10 @@ void QuantizedGraph::search_with_scratch_impl( throw std::invalid_argument("QuantizedGraph search buffers must not be null"); } if (hint != kPidMax && hint >= max_elements_) { - throw std::invalid_argument("SymphonyQG search hint is outside the graph"); + throw std::invalid_argument("QuantizedGraph search hint is outside the graph"); } if (!scratch.visited.initialized()) { - throw std::invalid_argument("SymphonyQG visited set must be initialized"); + throw std::invalid_argument("QuantizedGraph visited set must be initialized"); } rotator_->rotate(query, scratch.rotated_query); diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index 2027812e..338a5efd 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -37,6 +37,7 @@ QGBuilder::BuildQuery::BuildQuery( QGBuilder::BuildWorker::BuildWorker(size_t num_points, size_t visited_capacity) : visited(num_points, visited_capacity) {} + const QGBuilder::BuildCache* QGBuilder::build_cache() const { return cache_ ? &*cache_ : nullptr; } @@ -48,11 +49,11 @@ QGBuilder::QGBuilder( uint32_t seed, bool cache_vectors ) - : qg_{index} - , ef_construction_{ef_construction} - , num_workers_{resolve_num_threads(num_threads)} - , num_points_{qg_.num_vertices()} - , dim_{qg_.dimension()} + : qg_(index) + , ef_construction_(ef_construction) + , num_workers_(resolve_num_threads(num_threads)) + , num_points_(qg_.num_vertices()) + , dim_(qg_.dimension()) , max_degree_(qg_.degree_bound()) , seed_(seed) { if (cache_vectors) { @@ -71,9 +72,9 @@ QGBuilder::QGBuilder( ) : QGBuilder(index, ef_construction, num_threads, seed, cache_vectors) { if (init != QGInitialization::PiPNN && init != QGInitialization::Random) { - throw std::invalid_argument("Unknown QG initialization"); + throw std::invalid_argument("QuantizedGraph initialization is invalid"); } - seeded_ = init == QGInitialization::PiPNN; + initialization_ = init; reset(data); } @@ -97,9 +98,9 @@ QGBuilder::QGBuilder( void QGBuilder::reset(const float* data) { if (data == nullptr) { - throw std::invalid_argument("QGBuilder data must not be null"); + throw std::invalid_argument("QuantizedGraph build data must not be null"); } - if (seeded_) { + if (initialization_ == QGInitialization::PiPNN) { auto initial = detail::build_initial_graph( data, num_points_, dim_, max_degree_, qg_.metric_type_, num_workers_ ); @@ -115,92 +116,88 @@ void QGBuilder::initialize_seed( ) { if (offsets.size() != num_points_ + 1 || offsets.front() != 0 || offsets.back() != neighbors.size()) { - throw std::invalid_argument("Seed graph offsets must delimit every vertex"); + throw std::invalid_argument( + "QuantizedGraph seed graph offsets must delimit every vertex" + ); } - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - const size_t i = static_cast(index); + for (size_t i = 0; i < num_points_; ++i) { if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || offsets[i + 1] - offsets[i] > max_degree_) { throw std::invalid_argument( - "Seed graph row exceeds degree bound or has invalid offsets" + "QuantizedGraph seed graph row exceeds degree bound or has invalid offsets" ); } for (size_t j = offsets[i]; j < offsets[i + 1]; ++j) { if (neighbors[j] >= num_points_ || neighbors[j] == i) { throw std::invalid_argument( - "Seed graph IDs must be in range and exclude self" + "QuantizedGraph seed graph IDs must be in range and exclude self" ); } } } initialize_storage(data); - { - index_detail::ParallelException failure; + index_detail::ParallelException failure; #pragma omp parallel num_threads(num_workers_) - { - auto& worker = workers_[static_cast(omp_get_thread_num())]; - auto& ids = worker.neighbor_ids; - try { - ids.reserve(max_degree_); - } catch (...) { failure.capture(); } + { + auto& worker = workers_[static_cast(omp_get_thread_num())]; + auto& ids = worker.neighbor_ids; + try { + ids.reserve(max_degree_); + } catch (...) { failure.capture(); } #pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - if (failure.failed()) { - continue; - } - try { - const size_t i = static_cast(index); - ids.assign( - neighbors.begin() + static_cast(offsets[i]), - neighbors.begin() + static_cast(offsets[i + 1]) - ); - std::sort(ids.begin(), ids.end()); - ids.erase(std::unique(ids.begin(), ids.end()), ids.end()); - // Seed rows need only IDs; scores are materialized during search. - write_row(static_cast(i), worker); - } catch (...) { failure.capture(); } + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); + ++index) { + if (failure.failed()) { + continue; } + try { + const size_t i = static_cast(index); + ids.assign( + neighbors.begin() + static_cast(offsets[i]), + neighbors.begin() + static_cast(offsets[i + 1]) + ); + std::sort(ids.begin(), ids.end()); + ids.erase(std::unique(ids.begin(), ids.end()), ids.end()); + // Seed rows need only IDs; scores are materialized during search. + write_row(static_cast(i), worker); + } catch (...) { failure.capture(); } } - failure.rethrow(); } + failure.rethrow(); } void QGBuilder::refine() { run_passes(1); } -void QGBuilder::build() { run_passes(seeded_ ? 1 : 3); } +void QGBuilder::build() { run_passes(initialization_ == QGInitialization::PiPNN ? 1 : 3); } void QGBuilder::build(size_t num_iter) { if (num_iter < 2) { - throw std::invalid_argument("The number of QG build iterations must be at least 2"); + throw std::invalid_argument("QuantizedGraph build iterations must be at least 2"); } run_passes(num_iter); } bool QGBuilder::check_dup() const { std::atomic flag(false); - { - index_detail::ParallelException failure; + index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - if (failure.failed()) { - continue; - } - try { - const size_t i = static_cast(index); - std::unordered_set edges; - for (auto nei : new_neighbors_[i]) { - if (edges.find(nei.id) != edges.end()) { - flag = true; - } - edges.emplace(nei.id); - } - } catch (...) { failure.capture(); } + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); + ++index) { + if (failure.failed()) { + continue; } - failure.rethrow(); + try { + const size_t i = static_cast(index); + std::unordered_set edges; + for (const auto& neighbor : new_neighbors_[i]) { + if (edges.find(neighbor.id) != edges.end()) { + flag = true; + } + edges.emplace(neighbor.id); + } + } catch (...) { failure.capture(); } } + failure.rethrow(); return flag; } @@ -209,6 +206,7 @@ float QGBuilder::avg_degree() const { size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), size_t{0}); return static_cast(degrees) / static_cast(num_points_); } + void QGBuilder::run_passes(size_t num_iter) { qg_.ready_ = false; // Entry points guarantee a positive count; only the last pass completes degrees. @@ -219,7 +217,7 @@ void QGBuilder::run_passes(size_t num_iter) { qg_.ready_ = true; } -void QGBuilder::write_row(PID id, BuildWorker& worker) { +void QGBuilder::write_row(PID internal_id, BuildWorker& worker) { if (!worker.neighbor_ids.empty()) { const auto& new_neighbors = worker.neighbor_ids; const size_t num_neighbors = new_neighbors.size(); @@ -227,7 +225,7 @@ void QGBuilder::write_row(PID id, BuildWorker& worker) { const size_t padded_dim = qg_.padded_dim_; auto& rotated_data = worker.rotated_neighbors; rotated_data.resize((num_neighbors + 1) * padded_dim); - float* rotated_centroid = rotated_data.data() + num_neighbors * padded_dim; + float* rotated_source = rotated_data.data() + num_neighbors * padded_dim; // Cached rows, rotated raw vectors or reconstructed codes, as configured. const auto rotated_row = [&](PID source, float* row) { if (cache != nullptr) { @@ -243,31 +241,33 @@ void QGBuilder::write_row(PID id, BuildWorker& worker) { for (size_t i = 0; i < num_neighbors; ++i) { rotated_row(new_neighbors[i], rotated_data.data() + i * padded_dim); } - rotated_row(id, rotated_centroid); - qg_.encode_neighbors(id, new_neighbors, rotated_data.data(), rotated_centroid); + rotated_row(internal_id, rotated_source); + qg_.encode_neighbors( + internal_id, new_neighbors, rotated_data.data(), rotated_source + ); } - degrees_[id] = worker.neighbor_ids.size(); + degrees_[internal_id] = worker.neighbor_ids.size(); } -void QGBuilder::write_row(PID id, const CandidateList& neighbors, BuildWorker& worker) { +void QGBuilder::write_row( + PID internal_id, const CandidateList& neighbors, BuildWorker& worker +) { worker.neighbor_ids.resize(neighbors.size()); for (size_t i = 0; i < neighbors.size(); ++i) { worker.neighbor_ids[i] = neighbors[i].id; } - write_row(id, worker); + write_row(internal_id, worker); } void QGBuilder::initialize_storage(const float* data) { if (data == nullptr) { - throw std::invalid_argument("QGBuilder data must not be null"); + throw std::invalid_argument("QuantizedGraph build data must not be null"); } // Allocate refinement scratch after PiPNN releases its dense workspace. // Clear old graph-dependent state on reset while retaining vector capacities. new_neighbors_.resize(num_points_); reverse_buffer_.resize(num_points_); - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - const size_t i = static_cast(index); + for (size_t i = 0; i < num_points_; ++i) { new_neighbors_[i].clear(); reverse_buffer_[i].clear(); } @@ -289,34 +289,30 @@ void QGBuilder::initialize_storage(const float* data) { if (cache_) { const size_t padded_dim = qg_.padded_dim_; cache_->rotated_vectors.resize(num_points_ * padded_dim); - { - index_detail::ParallelException failure; + index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) schedule(static) - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - if (failure.failed()) { - continue; - } - try { - const size_t id = static_cast(index); - float* row = cache_->rotated_vectors.data() + id * padded_dim; - if (qg_.is_quantized()) { - qg_.reconstruct_quantized_vector(id, row); - } else { - qg_.rotator_->rotate(qg_.get_vector(id), row); - } - } catch (...) { failure.capture(); } + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); + ++index) { + if (failure.failed()) { + continue; } - failure.rethrow(); + try { + const auto internal_id = static_cast(index); + float* row = cache_->rotated_vectors.data() + internal_id * padded_dim; + if (qg_.is_quantized()) { + qg_.reconstruct_quantized_vector(internal_id, row); + } else { + qg_.rotator_->rotate(qg_.get_vector(internal_id), row); + } + } catch (...) { failure.capture(); } } + failure.rethrow(); cache_->queries.clear(); if (qg_.is_quantized()) { cache_->queries.reserve(num_points_); - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); - ++index) { - const size_t id = static_cast(index); + for (size_t internal_id = 0; internal_id < num_points_; ++internal_id) { cache_->queries.emplace_back( - cache_->rotated_vectors.data() + id * padded_dim, + cache_->rotated_vectors.data() + internal_id * padded_dim, qg_.centroid_.data(), padded_dim, qg_.metric_type_ @@ -331,11 +327,11 @@ void QGBuilder::initialize_storage(const float* data) { qg_.centroid_.data(), qg_.centroid_.data(), qg_.padded_dim_, qg_.metric_type_ ); float best = std::numeric_limits::max(); - for (PID id = 0; id < num_points_; ++id) { - const float distance = qg_.quantized_distance(query, id); + for (PID internal_id = 0; internal_id < num_points_; ++internal_id) { + const float distance = qg_.quantized_distance(query, internal_id); if (distance < best) { best = distance; - entry_point = id; + entry_point = internal_id; } } } else { @@ -357,29 +353,29 @@ void QGBuilder::add_pruned_edges( new_result.clear(); new_result = result; - std::unordered_set nei_set; - nei_set.reserve(max_degree_); - for (const auto& nei : result) { - nei_set.emplace(nei.id); + std::unordered_set neighbor_ids; + neighbor_ids.reserve(max_degree_); + for (const auto& neighbor : result) { + neighbor_ids.emplace(neighbor.id); } BuildQuery query(qg_, build_cache()); while (new_result.size() < max_degree_ && start < pruned_list.size()) { - const auto& cur = pruned_list[start]; + const auto& candidate = pruned_list[start]; bool occlude = false; - query.reset(cur.id); - float dik_sqr = cur.distance; + query.reset(candidate.id); + float dik_sqr = candidate.distance; - if (nei_set.find(cur.id) != nei_set.end()) { + if (neighbor_ids.find(candidate.id) != neighbor_ids.end()) { break; } - for (auto& nei : new_result) { - float dij_sqr = nei.distance; + for (const auto& neighbor : new_result) { + float dij_sqr = neighbor.distance; if (dij_sqr > dik_sqr) { break; } - float djk_sqr = query.distance_to(nei.id); + float djk_sqr = query.distance_to(neighbor.id); float cosine = (dik_sqr + dij_sqr - djk_sqr) / (2 * std::sqrt(dij_sqr * dik_sqr)); if (cosine > threshold) { @@ -389,8 +385,8 @@ void QGBuilder::add_pruned_edges( } if (!occlude) { - new_result.emplace_back(cur); - nei_set.emplace(cur.id); + new_result.emplace_back(candidate); + neighbor_ids.emplace(candidate.id); std::sort(new_result.begin(), new_result.end()); } @@ -399,36 +395,36 @@ void QGBuilder::add_pruned_edges( } void QGBuilder::heuristic_prune( - CandidateList& pool, CandidateList& pruned_results, CandidateList* recorded + CandidateList& pool, CandidateList& selected, CandidateList* recorded ) { if (pool.empty()) { return; } - pruned_results.clear(); - size_t poolsize = pool.size(); - if (poolsize <= max_degree_) { - pruned_results = pool; + selected.clear(); + const size_t pool_size = pool.size(); + if (pool_size <= max_degree_) { + selected = pool; return; } - std::vector pruned(poolsize, false); + std::vector pruned(pool_size, false); size_t start = 0; BuildQuery query(qg_, build_cache()); - while (pruned_results.size() < max_degree_ && start < poolsize) { + while (selected.size() < max_degree_ && start < pool_size) { auto candidate_id = pool[start].id; if (pruned[start]) { ++start; continue; } - pruned_results.emplace_back(pool[start]); + selected.emplace_back(pool[start]); query.reset(candidate_id); // i : current vertex // j : neighbor added in this iter // k : remained unpruned candidate neighbor - for (size_t k = start + 1; k < poolsize; ++k) { + for (size_t k = start + 1; k < pool_size; ++k) { if (pruned[k]) { continue; } @@ -447,64 +443,55 @@ void QGBuilder::heuristic_prune( } } -/** - * @brief search for new neighbor in qg - * - * Search pruning does not record candidates for final degree completion. - */ +// Search for new neighbors in the graph. Search pruning does not record candidates +// for final degree completion. void QGBuilder::search_new_neighbors() { - { - index_detail::omp_for_dynamic( - num_points_, num_workers_, [&](size_t i, size_t thread) { - PID point_id = i; - auto& worker = workers_[thread]; - CandidateList& candidates = worker.candidates; - candidates.clear(); - VisitedSet& visited = worker.visited; - candidates.reserve(2 * kMaxCandidatePoolSize); - visited.clear(); - find_candidates(point_id, candidates, visited); - - // Seeded construction keeps its initial edges only in QG. Materialize - // their scores on demand, after the caller can release the input/CSR. - if (new_neighbors_[point_id].empty() && degrees_[point_id] != 0) { - BuildQuery query(qg_, build_cache()); - query.reset(point_id); - const auto ids = qg_.get_neighbors(point_id); - for (size_t j = 0; j < degrees_[point_id]; ++j) { - if (ids[j] != point_id && !visited.get(ids[j])) { - candidates.emplace_back(ids[j], query.distance_to(ids[j])); - } - } + index_detail::omp_for_dynamic(num_points_, num_workers_, [&](size_t i, size_t thread) { + const auto internal_id = static_cast(i); + auto& worker = workers_[thread]; + CandidateList& candidates = worker.candidates; + candidates.clear(); + VisitedSet& visited = worker.visited; + candidates.reserve(2 * kMaxCandidatePoolSize); + visited.clear(); + find_candidates(internal_id, candidates, visited); + + // Seeded construction keeps its initial edges only in QG. Materialize + // their scores on demand, after the caller can release the input/CSR. + if (new_neighbors_[internal_id].empty() && degrees_[internal_id] != 0) { + BuildQuery query(qg_, build_cache()); + query.reset(internal_id); + const auto ids = qg_.get_neighbors(internal_id); + for (size_t j = 0; j < degrees_[internal_id]; ++j) { + if (ids[j] != internal_id && !visited.get(ids[j])) { + candidates.emplace_back(ids[j], query.distance_to(ids[j])); } + } + } - // Add current neighbors retained by preceding iterations. - for (auto& nei : new_neighbors_[point_id]) { - auto neighbor_id = nei.id; - if (neighbor_id != point_id && !visited.get(neighbor_id)) { - candidates.emplace_back(nei); - } - } - - size_t min_size = std::min(candidates.size(), kMaxCandidatePoolSize); - std::partial_sort( - candidates.begin(), - candidates.begin() + static_cast(min_size), - candidates.end() - ); - candidates.resize(min_size); - - // prune and update qg - new_neighbors_[point_id].reserve(max_degree_); - heuristic_prune(candidates, new_neighbors_[point_id], /*recorded=*/nullptr); + // Add current neighbors retained by preceding iterations. + for (const auto& neighbor : new_neighbors_[internal_id]) { + const PID neighbor_id = neighbor.id; + if (neighbor_id != internal_id && !visited.get(neighbor_id)) { + candidates.emplace_back(neighbor); } + } + + const size_t min_size = std::min(candidates.size(), kMaxCandidatePoolSize); + std::partial_sort( + candidates.begin(), + candidates.begin() + static_cast(min_size), + candidates.end() ); - } + candidates.resize(min_size); + + // Prune and update the graph. + new_neighbors_[internal_id].reserve(max_degree_); + heuristic_prune(candidates, new_neighbors_[internal_id], /*recorded=*/nullptr); + }); } void QGBuilder::add_reverse_edges(bool final_pass) { - // Striped locks bound this pass's memory independently of the point count. - constexpr size_t kLockStripes = 4096; std::vector locks(std::min(num_points_, kLockStripes)); for (auto& row : reverse_buffer_) { row.clear(); @@ -512,114 +499,103 @@ void QGBuilder::add_reverse_edges(bool final_pass) { // Keep new_neighbors_ read-only while reverse candidates are collected. Mutating a // destination row here races with another worker reading that row as its source. - { - index_detail::omp_for_dynamic( - num_points_, num_workers_, [&](size_t i, size_t /*thread*/) { - const PID data_id = static_cast(i); - for (const auto& nei : new_neighbors_[data_id]) { - const PID destination = nei.id; - const CandidateList& destination_neighbors = - new_neighbors_[destination]; - const bool reciprocal = std::any_of( - destination_neighbors.begin(), - destination_neighbors.end(), - [&](const auto& destination_neighbor) { - return destination_neighbor.id == data_id; - } - ); - if (reciprocal) { - continue; + index_detail::omp_for_dynamic( + num_points_, num_workers_, [&](size_t i, size_t /*thread*/) { + const auto source = static_cast(i); + for (const auto& neighbor : new_neighbors_[source]) { + const PID destination = neighbor.id; + const CandidateList& destination_neighbors = new_neighbors_[destination]; + const bool reciprocal = std::any_of( + destination_neighbors.begin(), + destination_neighbors.end(), + [&](const auto& destination_neighbor) { + return destination_neighbor.id == source; } + ); + if (reciprocal) { + continue; + } - std::lock_guard lock(locks[destination % locks.size()]); - const size_t missing_slots = - max_degree_ - std::min(max_degree_, destination_neighbors.size()); - if (reverse_buffer_[destination].size() < - kMaxCandidatePoolSize + missing_slots) { - reverse_buffer_[destination].emplace_back(data_id, nei.distance); - } + std::lock_guard lock(locks[destination % locks.size()]); + const size_t missing_slots = + max_degree_ - std::min(max_degree_, destination_neighbors.size()); + if (reverse_buffer_[destination].size() < + kMaxCandidatePoolSize + missing_slots) { + reverse_buffer_[destination].emplace_back(source, neighbor.distance); } } - ); - } + } + ); - { - index_detail::omp_for_dynamic( - num_points_, num_workers_, [&](size_t i, size_t thread) { - const PID data_id = static_cast(i); - CandidateList& candidate_pool = reverse_buffer_[data_id]; - if (qg_.is_quantized() && !candidate_pool.empty()) { - // RaBitQ estimates are directional: score destination -> source afresh. - BuildQuery query(qg_, build_cache()); - query.reset(data_id); - for (auto& candidate : candidate_pool) { - candidate.distance = query.distance_to(candidate.id); - } - } - candidate_pool.reserve(candidate_pool.size() + max_degree_); - candidate_pool.insert( - candidate_pool.end(), - new_neighbors_[data_id].begin(), - new_neighbors_[data_id].end() - ); - std::sort(candidate_pool.begin(), candidate_pool.end()); - auto& worker = workers_[thread]; - worker.pruned.clear(); - if (final_pass) { - worker.pruned.reserve(kMaxPrunedSize); - } - heuristic_prune( - candidate_pool, - new_neighbors_[data_id], - final_pass ? &worker.pruned : nullptr - ); - // Reverse collection has finished. Complete this row against the old - // graph before reusing its worker scratch; graph writes stay deferred. - if (final_pass) { - complete_degree(data_id, worker); - } + index_detail::omp_for_dynamic(num_points_, num_workers_, [&](size_t i, size_t thread) { + const auto internal_id = static_cast(i); + CandidateList& candidate_pool = reverse_buffer_[internal_id]; + if (qg_.is_quantized() && !candidate_pool.empty()) { + // RaBitQ estimates are directional: score destination -> source afresh. + BuildQuery query(qg_, build_cache()); + query.reset(internal_id); + for (auto& candidate : candidate_pool) { + candidate.distance = query.distance_to(candidate.id); } + } + candidate_pool.reserve(candidate_pool.size() + max_degree_); + candidate_pool.insert( + candidate_pool.end(), + new_neighbors_[internal_id].begin(), + new_neighbors_[internal_id].end() ); - } + std::sort(candidate_pool.begin(), candidate_pool.end()); + auto& worker = workers_[thread]; + worker.pruned.clear(); + if (final_pass) { + worker.pruned.reserve(kMaxPrunedSize); + } + heuristic_prune( + candidate_pool, + new_neighbors_[internal_id], + final_pass ? &worker.pruned : nullptr + ); + // Reverse collection has finished. Complete this row against the old + // graph before reusing its worker scratch; graph writes stay deferred. + if (final_pass) { + complete_degree(internal_id, worker); + } + }); } void QGBuilder::random_init() { const PID min_id = 0; const PID max_id = num_points_ - 1; - { - index_detail::omp_for_dynamic( - num_points_, num_workers_, [&](size_t i, size_t thread) { - new_neighbors_[i].clear(); - std::mt19937 rng(seed_ ^ static_cast(i)); - std::uniform_int_distribution distribution(min_id, max_id); - std::unordered_set neighbor_set; - neighbor_set.reserve(max_degree_); - while (neighbor_set.size() < max_degree_) { - const PID rand_id = distribution(rng); - if (rand_id != i) { - neighbor_set.emplace(rand_id); - } - } + index_detail::omp_for_dynamic(num_points_, num_workers_, [&](size_t i, size_t thread) { + new_neighbors_[i].clear(); + std::mt19937 rng(seed_ ^ static_cast(i)); + std::uniform_int_distribution distribution(min_id, max_id); + std::unordered_set neighbor_set; + neighbor_set.reserve(max_degree_); + while (neighbor_set.size() < max_degree_) { + const PID rand_id = distribution(rng); + if (rand_id != i) { + neighbor_set.emplace(rand_id); + } + } - BuildQuery query(qg_, build_cache()); - query.reset(i); - new_neighbors_[i].reserve(max_degree_); - for (PID cur_neigh : neighbor_set) { - new_neighbors_[i].emplace_back(cur_neigh, query.distance_to(cur_neigh)); - } + BuildQuery query(qg_, build_cache()); + query.reset(i); + new_neighbors_[i].reserve(max_degree_); + for (PID neighbor_id : neighbor_set) { + new_neighbors_[i].emplace_back(neighbor_id, query.distance_to(neighbor_id)); + } - write_row(static_cast(i), new_neighbors_[i], workers_[thread]); - } - ); - } + write_row(static_cast(i), new_neighbors_[i], workers_[thread]); + }); } // Complete one final-pass row using only this worker's reverse-prune candidates. -void QGBuilder::complete_degree(PID i, BuildWorker& worker) { - CandidateList& cur_neighbors = new_neighbors_[i]; - size_t num_neighbors = cur_neighbors.size(); +void QGBuilder::complete_degree(PID internal_id, BuildWorker& worker) { + CandidateList& current_neighbors = new_neighbors_[internal_id]; + const size_t num_neighbors = current_neighbors.size(); - // skip vertices with enough neighbors + // Skip vertices with enough neighbors. if (num_neighbors >= max_degree_) { return; } @@ -630,16 +606,16 @@ void QGBuilder::complete_degree(PID i, BuildWorker& worker) { std::sort(pruned_list.begin(), pruned_list.end()); - // use binary search to get refined results + // Use binary search to get refined results. CandidateList& successful_result = worker.successful_result; successful_result.clear(); successful_result.reserve(max_degree_); - float left = 0.5; - float right = 1.0; + float left = 0.5F; + float right = 1.0F; size_t iter = 0; while (iter++ < kMaxBsIter) { float mid = (left + right) / 2; - add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); + add_pruned_edges(current_neighbors, pruned_list, new_result, mid); if (new_result.size() < max_degree_) { left = mid; } else { @@ -653,12 +629,12 @@ void QGBuilder::complete_degree(PID i, BuildWorker& worker) { if (!successful_result.empty()) { new_result.swap(successful_result); } else { - add_pruned_edges(cur_neighbors, pruned_list, new_result, right); + add_pruned_edges(current_neighbors, pruned_list, new_result, right); } - // if the vertex still doesn't have enough neighbors, use random vertices + // If the vertex still does not have enough neighbors, use random vertices. if (new_result.size() < max_degree_) { - std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); + std::mt19937 rng(seed_ ^ static_cast(internal_id) ^ 0x80000000U); std::uniform_int_distribution distribution( 0, static_cast(num_points_) - 1 ); @@ -668,53 +644,50 @@ void QGBuilder::complete_degree(PID i, BuildWorker& worker) { ids.emplace(neighbor.id); } BuildQuery query(qg_, build_cache()); - query.reset(i); + query.reset(internal_id); while (new_result.size() < max_degree_) { const PID rand_id = distribution(rng); - if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { + if (rand_id != internal_id && ids.find(rand_id) == ids.end()) { new_result.emplace_back(rand_id, query.distance_to(rand_id)); ids.emplace(rand_id); } } } - cur_neighbors = new_result; + current_neighbors = new_result; } void QGBuilder::run_pass(bool final_pass) { search_new_neighbors(); add_reverse_edges(final_pass); - // update qg - { - index_detail::omp_for_dynamic( - num_points_, num_workers_, [&](size_t i, size_t thread) { - write_row(static_cast(i), new_neighbors_[i], workers_[thread]); - } - ); - } + // Update the graph. + index_detail::omp_for_dynamic(num_points_, num_workers_, [&](size_t i, size_t thread) { + write_row(static_cast(i), new_neighbors_[i], workers_[thread]); + }); } -void QGBuilder::BuildQuery::reset(PID id) { + +void QGBuilder::BuildQuery::reset(PID internal_id) { raw_ = nullptr; rotated_ = nullptr; prepared_.reset(); if (cache_ != nullptr) { - rotated_ = cache_->rotated_vectors.data() + id * graph_.padded_dim_; + rotated_ = cache_->rotated_vectors.data() + internal_id * graph_.padded_dim_; } if (graph_.is_quantized()) { if (cache_ != nullptr) { - prepared_.emplace(rotated_, cache_->queries[id]); + prepared_.emplace(rotated_, cache_->queries[internal_id]); } else { storage_.resize(graph_.padded_dim_); // Reconstructed vectors are already rotated; do not rotate them again. - graph_.reconstruct_quantized_vector(id, storage_.data()); + graph_.reconstruct_quantized_vector(internal_id, storage_.data()); rotated_ = storage_.data(); prepared_.emplace( rotated_, graph_.centroid_.data(), graph_.padded_dim_, graph_.metric_type_ ); } } else { - raw_ = graph_.get_vector(id); + raw_ = graph_.get_vector(internal_id); } } @@ -731,12 +704,12 @@ const float* QGBuilder::BuildQuery::rotated_query() { return rotated_; } -// find candidate neighbors for point_id, exclude the vertex itself +// Find candidate neighbors for internal_id, excluding the vertex itself. void QGBuilder::find_candidates( - PID point_id, CandidateList& results, VisitedSet& visited + PID internal_id, CandidateList& results, VisitedSet& visited ) const { BuildQuery query(qg_, build_cache()); - query.reset(point_id); + query.reset(internal_id); BatchQuery query_wrapper( query.rotated_query(), qg_.padded_dim_, qg_.metric_type_ ); @@ -761,7 +734,7 @@ void QGBuilder::find_candidates( visited, num_neighbors ); - if (candidate_id != point_id) { + if (candidate_id != internal_id) { results.emplace_back(candidate_id, vertex_distance); } } diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 923c3527..26394873 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -17,11 +17,11 @@ namespace rabitqlib::simd { void accumulate_cluster_sums( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ) { @@ -31,7 +31,7 @@ void accumulate_cluster_sums( accumulate_cluster_sums_generic ); - kAccumulateClusterSumsFn(x, point_ids, count, dim, labels, sums, counts); + kAccumulateClusterSumsFn(data, point_ids, count, dim, assignments, sums, counts); } void matrix_product( @@ -120,16 +120,8 @@ float l2norm_sqr(const float* a, size_t dim) { static const auto kL2normSqrFn = RABITQ_RESOLVE_NEON( l2norm_sqr_neon, l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic ); - return kL2normSqrFn(a, dim); -} -// With zero extra bits there is no extra code to contribute to the inner -// product, so the ex_bits == 0 slot must be a constant-zero stub rather than -// a duplicate of the 1-bit implementation. -static float ip_fxu0( - const float* /*query*/, const uint8_t* /*compact_code*/, size_t /*dim*/ -) { - return 0.0F; + return kL2normSqrFn(a, dim); } double best_rescale_factor( @@ -151,13 +143,13 @@ double best_rescale_factor_with_workspace( double* promoted, double* reciprocals ) { - static const auto kBestRescaleWorkspaceFn = RABITQ_RESOLVE( + static const auto kBestRescaleFactorWithWorkspaceFn = RABITQ_RESOLVE( best_rescale_factor_with_workspace_avx512, best_rescale_factor_with_workspace_avx2, best_rescale_factor_with_workspace_generic ); - return kBestRescaleWorkspaceFn( + return kBestRescaleFactorWithWorkspaceFn( magnitudes, dim, max_code, start, end, promoted, reciprocals ); } @@ -178,64 +170,6 @@ void fht_rotate( kFhtRotateFn(data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip); } -using ExcodeIpTable = std::array; - -static ExcodeIpTable resolve_excode_ip_table() { - const auto generic = ExcodeIpTable{ - ip_fxu0, - excode_ipimpl::ip16_fxu1_generic, - excode_ipimpl::ip64_fxu2_generic, - excode_ipimpl::ip64_fxu3_generic, - excode_ipimpl::ip16_fxu4_generic, - excode_ipimpl::ip64_fxu5_generic, - excode_ipimpl::ip64_fxu6_generic, - excode_ipimpl::ip64_fxu7_generic, - excode_ipimpl::ip16_fxu8_generic, - }; -#if defined(__aarch64__) || defined(_M_ARM64) - return resolve_optional_kernel( - ExcodeIpTable{ - ip_fxu0, - excode_ipimpl::ip16_fxu1_neon, - excode_ipimpl::ip64_fxu2_neon, - excode_ipimpl::ip64_fxu3_neon, - excode_ipimpl::ip16_fxu4_neon, - excode_ipimpl::ip64_fxu5_neon, - excode_ipimpl::ip64_fxu6_neon, - excode_ipimpl::ip64_fxu7_neon, - excode_ipimpl::ip16_fxu8_neon}, - generic, - cpu::has_neon() - ); -#else - return RABITQ_RESOLVE( - (ExcodeIpTable{ - ip_fxu0, - excode_ipimpl::ip16_fxu1_avx512, - excode_ipimpl::ip64_fxu2_avx512, - excode_ipimpl::ip64_fxu3_avx512, - excode_ipimpl::ip16_fxu4_avx512, - excode_ipimpl::ip64_fxu5_avx512, - excode_ipimpl::ip64_fxu6_avx512, - excode_ipimpl::ip64_fxu7_avx512, - excode_ipimpl::ip16_fxu8_avx512, - }), - (ExcodeIpTable{ - ip_fxu0, - excode_ipimpl::ip16_fxu1_avx2, - excode_ipimpl::ip64_fxu2_avx2, - excode_ipimpl::ip64_fxu3_avx2, - excode_ipimpl::ip16_fxu4_avx2, - excode_ipimpl::ip64_fxu5_avx2, - excode_ipimpl::ip64_fxu6_avx2, - excode_ipimpl::ip64_fxu7_avx2, - excode_ipimpl::ip16_fxu8_avx2, - }), - generic - ); -#endif -} - void flip_sign(const uint8_t* flip, float* data, size_t dim) { static const auto kFlipSignFn = RABITQ_RESOLVE_NEON( flip_sign_neon, flip_sign_avx512, flip_sign_avx2, flip_sign_generic @@ -248,11 +182,12 @@ void kacs_walk(float* data, size_t len) { static const auto kKacsWalkFn = RABITQ_RESOLVE_NEON( kacs_walk_neon, kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic ); + kKacsWalkFn(data, len); } void scalar_quantize_uint8( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ) { static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE_NEON( scalar_quantize_uint8_neon, @@ -261,11 +196,11 @@ void scalar_quantize_uint8( scalar_quantize_uint8_generic ); - kScalarQuantizeUint8Fn(result, vec0, dim, lo, delta); + kScalarQuantizeUint8Fn(result, data, dim, lo, delta); } void scalar_quantize_uint16( - uint16_t* result, const float* vec0, size_t dim, float lo, float delta + uint16_t* result, const float* data, size_t dim, float lo, float delta ) { static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE_NEON( scalar_quantize_uint16_neon, @@ -274,7 +209,7 @@ void scalar_quantize_uint16( scalar_quantize_uint16_generic ); - kScalarQuantizeUint16Fn(result, vec0, dim, lo, delta); + kScalarQuantizeUint16Fn(result, data, dim, lo, delta); } void packing_2bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { @@ -325,68 +260,131 @@ void packing_7bit_excode(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { kPacking7BitExcodeFn(o_raw, o_compact, dim); } -} // namespace rabitqlib::simd +namespace { +// ex_bits is in [0, kMaxExBits]; the table holds one kernel per value. +constexpr size_t kMaxExBits = 8; +using ExcodeIpTable = std::array; -namespace rabitqlib { +// With zero extra bits there is no extra code to contribute to the inner +// product, so the ex_bits == 0 slot must be a constant-zero stub rather than +// a duplicate of the 1-bit implementation. +float ip_fxu0(const float* /*query*/, const uint8_t* /*compact_code*/, size_t /*dim*/) { + return 0.0F; +} -static const simd::ExcodeIpTable& excode_ip_table() { - static const auto table = simd::resolve_excode_ip_table(); - return table; +const ExcodeIpTable& excode_ip_table() { + static const auto kExcodeIpTable = RABITQ_RESOLVE_NEON( + (ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_neon, + excode_ipimpl::ip64_fxu2_neon, + excode_ipimpl::ip64_fxu3_neon, + excode_ipimpl::ip16_fxu4_neon, + excode_ipimpl::ip64_fxu5_neon, + excode_ipimpl::ip64_fxu6_neon, + excode_ipimpl::ip64_fxu7_neon, + excode_ipimpl::ip16_fxu8_neon, + }), + (ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_avx512, + excode_ipimpl::ip64_fxu2_avx512, + excode_ipimpl::ip64_fxu3_avx512, + excode_ipimpl::ip16_fxu4_avx512, + excode_ipimpl::ip64_fxu5_avx512, + excode_ipimpl::ip64_fxu6_avx512, + excode_ipimpl::ip64_fxu7_avx512, + excode_ipimpl::ip16_fxu8_avx512, + }), + (ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_avx2, + excode_ipimpl::ip64_fxu2_avx2, + excode_ipimpl::ip64_fxu3_avx2, + excode_ipimpl::ip16_fxu4_avx2, + excode_ipimpl::ip64_fxu5_avx2, + excode_ipimpl::ip64_fxu6_avx2, + excode_ipimpl::ip64_fxu7_avx2, + excode_ipimpl::ip16_fxu8_avx2, + }), + (ExcodeIpTable{ + ip_fxu0, + excode_ipimpl::ip16_fxu1_generic, + excode_ipimpl::ip64_fxu2_generic, + excode_ipimpl::ip64_fxu3_generic, + excode_ipimpl::ip16_fxu4_generic, + excode_ipimpl::ip64_fxu5_generic, + excode_ipimpl::ip64_fxu6_generic, + excode_ipimpl::ip64_fxu7_generic, + excode_ipimpl::ip16_fxu8_generic, + }) + ); + + return kExcodeIpTable; } +} // namespace + +} // namespace rabitqlib::simd + +namespace rabitqlib { using MaskIpX0QFn = float (*)(const float*, const uint8_t*, size_t); +using WarmupIpX0Q512Fn = + float (*)(const uint8_t*, const uint64_t*, float, float, size_t, size_t); ex_ipfunc select_excode_ipfunc(size_t ex_bits) { - if (ex_bits <= 8) { - return excode_ip_table()[ex_bits]; + if (ex_bits <= simd::kMaxExBits) { + return simd::excode_ip_table()[ex_bits]; } - throw std::invalid_argument("Bad IP function for IVF"); + throw std::invalid_argument("ex_bits must be in [0, 8]"); } float excode_ipimpl::ip16_fxu1_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[1](query, compact_code, dim); + return simd::excode_ip_table()[1](query, compact_code, dim); } float excode_ipimpl::ip64_fxu2_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[2](query, compact_code, dim); + return simd::excode_ip_table()[2](query, compact_code, dim); } float excode_ipimpl::ip64_fxu3_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[3](query, compact_code, dim); + return simd::excode_ip_table()[3](query, compact_code, dim); } float excode_ipimpl::ip16_fxu4_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[4](query, compact_code, dim); + return simd::excode_ip_table()[4](query, compact_code, dim); } float excode_ipimpl::ip64_fxu5_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[5](query, compact_code, dim); + return simd::excode_ip_table()[5](query, compact_code, dim); } float excode_ipimpl::ip64_fxu6_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[6](query, compact_code, dim); + return simd::excode_ip_table()[6](query, compact_code, dim); } float excode_ipimpl::ip64_fxu7_avx( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - return excode_ip_table()[7](query, compact_code, dim); + return simd::excode_ip_table()[7](query, compact_code, dim); } -void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query) { +void new_transpose_bin( + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query +) { static const auto kNewTransposeBinFn = RABITQ_RESOLVE_NEON( simd::new_transpose_bin_neon, simd::new_transpose_bin_avx512, @@ -394,11 +392,11 @@ void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_ simd::new_transpose_bin_generic ); - kNewTransposeBinFn(q, tq, padded_dim, b_query); + kNewTransposeBinFn(query, transposed_query, padded_dim, b_query); } void new_transpose_bin_512( - const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE_NEON( simd::new_transpose_bin_512_neon, @@ -407,11 +405,11 @@ void new_transpose_bin_512( simd::new_transpose_bin_512_generic ); - kNewTransposeBin512Fn(q, tq, padded_dim, b_query); + kNewTransposeBin512Fn(query, transposed_query, padded_dim, b_query); } float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim) { - static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE_NEON( + static const auto kMaskIpX0QFn = RABITQ_RESOLVE_NEON( static_cast(simd::mask_ip_x0_q_neon), static_cast(simd::mask_ip_x0_q_avx512), static_cast(simd::mask_ip_x0_q_avx2), @@ -425,13 +423,6 @@ float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim) return mask_ip_x0_q(query, reinterpret_cast(data), padded_dim); } -} // namespace rabitqlib - -namespace rabitqlib { - -using WarmupIpX0Q512Fn = - float (*)(const uint8_t*, const uint64_t*, float, float, size_t, size_t); - float warmup_ip_x0_q_512( const uint8_t* data, const uint64_t* query, @@ -440,11 +431,11 @@ float warmup_ip_x0_q_512( size_t padded_dim, size_t b_query ) { - static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE_NEON( + static const auto kWarmupIpX0Q512Fn = RABITQ_RESOLVE_NEON( static_cast(simd::warmup_ip_x0_q_512_neon), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), - static_cast(rabitqlib::simd::warmup_ip_x0_q_512_generic), + static_cast(simd::warmup_ip_x0_q_512_avx512), + static_cast(simd::warmup_ip_x0_q_512_avx2), + static_cast(simd::warmup_ip_x0_q_512_generic), cpu::has_avx512_popcnt() ); diff --git a/src/simd/fht_kernels.hpp b/src/simd/fht_kernels.hpp index 3abcd9fc..20ca0d0b 100644 --- a/src/simd/fht_kernels.hpp +++ b/src/simd/fht_kernels.hpp @@ -14,13 +14,13 @@ namespace { // original upper lane, and adding. XOR changes only the sign bit, replacing // the duplicated-lane permutations, subtraction from zero, and blends. // The masks negate odd lanes, upper lane pairs, then the upper 128-bit half. -template +template inline __m256 fht_lane_stage(__m256 x) { - if constexpr (stage == 0) { + if constexpr (Stage == 0) { const __m256 sign = _mm256_setr_ps(0.0F, -0.0F, 0.0F, -0.0F, 0.0F, -0.0F, 0.0F, -0.0F); return _mm256_add_ps(_mm256_permute_ps(x, 0xB1), _mm256_xor_ps(x, sign)); - } else if constexpr (stage == 1) { + } else if constexpr (Stage == 1) { const __m256 sign = _mm256_setr_ps(0.0F, 0.0F, -0.0F, -0.0F, 0.0F, 0.0F, -0.0F, -0.0F); return _mm256_add_ps(_mm256_permute_ps(x, 0x4E), _mm256_xor_ps(x, sign)); @@ -39,7 +39,7 @@ inline void fht_butterfly(__m256& a, __m256& b) { // Fuse three stages so eight registers are loaded and stored only once. // Stride is eight for the 64-point leaf, or N/8 for a larger transform. -template +template inline void fht_eight_registers(float* data, size_t stride, float factor = 1.0F) { __m256 a = _mm256_loadu_ps(data); __m256 b = _mm256_loadu_ps(data + stride); @@ -49,7 +49,7 @@ inline void fht_eight_registers(float* data, size_t stride, float factor = 1.0F) __m256 f = _mm256_loadu_ps(data + 5 * stride); __m256 g = _mm256_loadu_ps(data + 6 * stride); __m256 h = _mm256_loadu_ps(data + 7 * stride); - if constexpr (leaf) { + if constexpr (Leaf) { // As in FFHT's assembly, finish each stage across all eight registers // before starting the next, exposing independent work to the scheduler. a = fht_lane_stage<0>(a); @@ -89,7 +89,7 @@ inline void fht_eight_registers(float* data, size_t stride, float factor = 1.0F) fht_butterfly(b, f); fht_butterfly(c, g); fht_butterfly(d, h); - if constexpr (normalize) { + if constexpr (Normalize) { const __m256 multiplier = _mm256_set1_ps(factor); a = _mm256_mul_ps(a, multiplier); b = _mm256_mul_ps(b, multiplier); @@ -111,20 +111,20 @@ inline void fht_eight_registers(float* data, size_t stride, float factor = 1.0F) } // Only the outermost merge applies the rotation scale, after every butterfly. -template -void fht_intrinsics(float* data, float factor = 1.0F) { - static_assert(log_dim >= 6 && log_dim <= 16); - constexpr size_t kDim = size_t{1} << log_dim; - if constexpr (log_dim == 6) { - fht_eight_registers(data, 8, factor); - } else if constexpr (log_dim < 9) { - fht_intrinsics(data); - fht_intrinsics(data + kDim / 2); +template +void fht_impl(float* data, float factor = 1.0F) { + static_assert(LogDim >= 6 && LogDim <= 16); + constexpr size_t kDim = size_t{1} << LogDim; + if constexpr (LogDim == 6) { + fht_eight_registers(data, 8, factor); + } else if constexpr (LogDim < 9) { + fht_impl(data); + fht_impl(data + kDim / 2); for (size_t i = 0; i < kDim / 2; i += 8) { __m256 a = _mm256_loadu_ps(data + i); __m256 b = _mm256_loadu_ps(data + kDim / 2 + i); fht_butterfly(a, b); - if constexpr (normalize) { + if constexpr (Normalize) { const __m256 multiplier = _mm256_set1_ps(factor); a = _mm256_mul_ps(a, multiplier); b = _mm256_mul_ps(b, multiplier); @@ -137,10 +137,10 @@ void fht_intrinsics(float* data, float factor = 1.0F) { // set cache-local even for the largest supported rotation. constexpr size_t kStride = kDim / 8; for (size_t block = 0; block < kDim; block += kStride) { - fht_intrinsics(data + block); + fht_impl(data + block); } for (size_t i = 0; i < kStride; i += 8) { - fht_eight_registers(data + i, kStride, factor); + fht_eight_registers(data + i, kStride, factor); } } } diff --git a/src/simd/matrix_avx2.cpp b/src/simd/matrix_avx2.cpp index e10d6347..a095d317 100644 --- a/src/simd/matrix_avx2.cpp +++ b/src/simd/matrix_avx2.cpp @@ -8,15 +8,15 @@ namespace rabitqlib::simd { void accumulate_cluster_sums_avx2( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ) { - accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); + accumulate_cluster_sums_impl(data, point_ids, count, dim, assignments, sums, counts); } void matrix_product_avx2( diff --git a/src/simd/matrix_avx512.cpp b/src/simd/matrix_avx512.cpp index dd1d6333..64885af2 100644 --- a/src/simd/matrix_avx512.cpp +++ b/src/simd/matrix_avx512.cpp @@ -8,15 +8,15 @@ namespace rabitqlib::simd { void accumulate_cluster_sums_avx512( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ) { - accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); + accumulate_cluster_sums_impl(data, point_ids, count, dim, assignments, sums, counts); } void matrix_product_avx512( diff --git a/src/simd/matrix_generic.cpp b/src/simd/matrix_generic.cpp index aa5ddf13..e62f70d1 100644 --- a/src/simd/matrix_generic.cpp +++ b/src/simd/matrix_generic.cpp @@ -8,15 +8,15 @@ namespace rabitqlib::simd { void accumulate_cluster_sums_generic( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ) { - accumulate_cluster_sums_impl(x, point_ids, count, dim, labels, sums, counts); + accumulate_cluster_sums_impl(data, point_ids, count, dim, assignments, sums, counts); } void matrix_product_generic( diff --git a/src/simd/matrix_kernels.hpp b/src/simd/matrix_kernels.hpp index 745dcd05..193b3731 100644 --- a/src/simd/matrix_kernels.hpp +++ b/src/simd/matrix_kernels.hpp @@ -26,19 +26,19 @@ using RowMajorMatrixMap = kernel_eigen::Map; using VectorMap = kernel_eigen::Map>; inline void accumulate_cluster_sums_impl( - const float* x, + const float* data, const size_t* point_ids, size_t count, size_t dim, - const uint32_t* labels, + const uint32_t* assignments, double* sums, size_t* counts ) { for (size_t index = 0; index < count; ++index) { const size_t point = point_ids[index]; - const uint32_t cluster = labels[point]; + const uint32_t cluster = assignments[point]; ++counts[cluster]; - const float* vector = x + point * dim; + const float* vector = data + point * dim; double* sum = sums + static_cast(cluster) * dim; for (size_t coordinate = 0; coordinate < dim; ++coordinate) { sum[coordinate] += static_cast(vector[coordinate]); diff --git a/src/simd/pack_excode_avx2.cpp b/src/simd/pack_excode_avx2.cpp index 861c4743..a4fcc4f8 100644 --- a/src/simd/pack_excode_avx2.cpp +++ b/src/simd/pack_excode_avx2.cpp @@ -7,19 +7,19 @@ namespace rabitqlib::simd { void packing_2bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_2bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_2bit_excode_impl(o_raw, o_compact, dim); } void packing_3bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_3bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_3bit_excode_impl(o_raw, o_compact, dim); } void packing_4bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_4bit_excode_intrinsics<64>(o_raw, o_compact, dim); + detail::packing_4bit_excode_impl<64>(o_raw, o_compact, dim); } void packing_5bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_5bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_5bit_excode_impl(o_raw, o_compact, dim); } void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { @@ -67,7 +67,7 @@ void packing_6bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t d } void packing_7bit_excode_avx2(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_7bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_7bit_excode_impl(o_raw, o_compact, dim); } } // namespace rabitqlib::simd diff --git a/src/simd/pack_excode_avx512.cpp b/src/simd/pack_excode_avx512.cpp index a9b0ab23..575f3fba 100644 --- a/src/simd/pack_excode_avx512.cpp +++ b/src/simd/pack_excode_avx512.cpp @@ -39,7 +39,7 @@ void pack_6bit_tail(const uint8_t* raw, uint8_t* compact, size_t dim) { } // namespace void packing_2bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_2bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_2bit_excode_impl(o_raw, o_compact, dim); } // Reuse the AVX2 odd-width packers: wider code generation adds overhead here. @@ -48,7 +48,7 @@ void packing_3bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t } void packing_4bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_4bit_excode_intrinsics<128>(o_raw, o_compact, dim); + detail::packing_4bit_excode_impl<128>(o_raw, o_compact, dim); } void packing_5bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { @@ -56,7 +56,7 @@ void packing_5bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t } void packing_6bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { - detail::packing_6bit_excode_intrinsics(o_raw, o_compact, dim); + detail::packing_6bit_excode_impl(o_raw, o_compact, dim); } void packing_7bit_excode_avx512(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index 88cdf712..f524f5af 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -13,11 +13,12 @@ namespace rabitqlib::simd::detail { namespace { template -inline void pack_excode_tail_intrinsics(const uint8_t* raw, uint8_t* compact, size_t dim) { +inline void pack_excode_tail_impl(const uint8_t* raw, uint8_t* compact, size_t dim) { static_assert(Bits == 2 || Bits == 3 || Bits == 5 || Bits == 6 || Bits == 7); assert(dim == 0 || dim == 32); - if (dim == 0) + if (dim == 0) { return; + } // A partial block is always 32 coordinates. const __m128i pair_weights = _mm_set1_epi16(static_cast(1U | (1U << (Bits + 8)))); @@ -95,10 +96,8 @@ inline uint64_t pack_excode_high_bit(const uint8_t* raw) { ); } -template > -inline void packing_2bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim -) { +template > +inline void packing_2bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { // Full blocks preserve the existing byte layout. for (size_t j = 0; j < dim - dim % 64; j += 64) { // pack 64 2-bit codes into 128 bits (16 bytes) @@ -121,12 +120,10 @@ inline void packing_2bit_excode_intrinsics( o_raw += 64; o_compact += 16; } - pack_tail(o_raw, o_compact, dim % 64); + PackTail(o_raw, o_compact, dim % 64); } -inline void packing_3bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim -) { +inline void packing_3bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { // Full blocks preserve the existing byte layout. const __m128i mask = _mm_set1_epi8(0b11); for (size_t d = 0; d < dim - dim % 64; d += 64) { @@ -162,7 +159,7 @@ inline void packing_3bit_excode_intrinsics( o_raw += 64; o_compact += 8; } - pack_excode_tail_intrinsics<3>(o_raw, o_compact, dim % 64); + pack_excode_tail_impl<3>(o_raw, o_compact, dim % 64); } // Each 16-coordinate group stores the first eight codes in the low nibbles @@ -205,37 +202,33 @@ inline void pack_4bit_block(const uint8_t* raw, uint8_t* compact) { } template -inline void packing_4bit_excode_intrinsics( - const uint8_t* raw, uint8_t* compact, size_t dim -) { +inline void packing_4bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { static_assert(BlockDim == 64 || BlockDim == 128); for (; dim >= BlockDim; dim -= BlockDim) { - pack_4bit_block(raw, compact); - raw += BlockDim; - compact += BlockDim / 2; + pack_4bit_block(o_raw, o_compact); + o_raw += BlockDim; + o_compact += BlockDim / 2; } if constexpr (BlockDim == 128) { if (dim >= 64) { - pack_4bit_block<64>(raw, compact); - raw += 64; - compact += 32; + pack_4bit_block<64>(o_raw, o_compact); + o_raw += 64; + o_compact += 32; dim -= 64; } } if (dim >= 32) { - pack_4bit_block<32>(raw, compact); - raw += 32; - compact += 16; + pack_4bit_block<32>(o_raw, o_compact); + o_raw += 32; + o_compact += 16; dim -= 32; } if (dim != 0) { - pack_4bit_block<16>(raw, compact); + pack_4bit_block<16>(o_raw, o_compact); } } -inline void packing_5bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim -) { +inline void packing_5bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { // Full blocks preserve the existing byte layout. const __m128i mask = _mm_set1_epi8(0b1111); for (size_t j = 0; j < dim - dim % 64; j += 64) { @@ -268,13 +261,11 @@ inline void packing_5bit_excode_intrinsics( o_raw += 64; o_compact += 8; } - pack_excode_tail_intrinsics<5>(o_raw, o_compact, dim % 64); + pack_excode_tail_impl<5>(o_raw, o_compact, dim % 64); } -template > -inline void packing_6bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim -) { +template > +inline void packing_6bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { // for vec00 to vec47, split code into 6 // for vec48 to vec63, split code into 2 + 2 + 2 // Complementary masks let AVX512 fuse each bit selection into one instruction. @@ -308,12 +299,10 @@ inline void packing_6bit_excode_intrinsics( o_compact += 48; o_raw += 64; } - pack_tail(o_raw, o_compact, dim % 64); + PackTail(o_raw, o_compact, dim % 64); } -inline void packing_7bit_excode_intrinsics( - const uint8_t* o_raw, uint8_t* o_compact, size_t dim -) { +inline void packing_7bit_excode_impl(const uint8_t* o_raw, uint8_t* o_compact, size_t dim) { // for vec00 to vec47, split code into 6 + 1 // for vec48 to vec63, split code into 2 + 2 + 2 + 1 const __m128i mask2 = _mm_set1_epi8(static_cast(0b11000000)); @@ -352,7 +341,7 @@ inline void packing_7bit_excode_intrinsics( o_compact += 8; o_raw += 64; } - pack_excode_tail_intrinsics<7>(o_raw, o_compact, dim % 64); + pack_excode_tail_impl<7>(o_raw, o_compact, dim % 64); } } // namespace diff --git a/src/simd/packed_tail_avx2.hpp b/src/simd/packed_tail_avx2.hpp index 6e7e40f0..ad57f448 100644 --- a/src/simd/packed_tail_avx2.hpp +++ b/src/simd/packed_tail_avx2.hpp @@ -40,8 +40,9 @@ inline void accumulate_excode_tail_avx2( const float* query, const uint8_t* compact, size_t dim, __m256& sum, __m256& sum_hi ) { assert(dim == 0 || dim == 32); - if (dim == 0) + if (dim == 0) { return; + } // The fixed 32-coordinate tail contains two 16-code groups. // Select the two bytes spanning each code, then shift/mask eight codes at once. alignas(16) uint8_t positions[32]; @@ -85,6 +86,15 @@ inline void accumulate_excode_tail_avx2( } } +// Build a vector from two 64-bit words without a signed narrowing conversion. +[[nodiscard]] inline __m128i set_u64x(uint64_t high, uint64_t low) noexcept { + int64_t signed_high = 0; + int64_t signed_low = 0; + std::memcpy(&signed_high, &high, sizeof(signed_high)); + std::memcpy(&signed_low, &low, sizeof(signed_low)); + return _mm_set_epi64x(signed_high, signed_low); +} + // Sum the eight lanes as a tree rather than serially. inline float reduce_add_avx2(__m256 v) { __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index 1c2e54c1..00a5edc0 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -158,11 +158,11 @@ void fht(float* data, size_t dim) { } } // namespace -#define RABITQ_PACK(Bits) \ - void packing_##Bits##bit_excode_generic( \ - const uint8_t* raw, uint8_t* compact, size_t dim \ - ) { \ - pack(raw, compact, dim); \ +#define RABITQ_PACK(Bits) \ + void packing_##Bits##bit_excode_generic( \ + const uint8_t* o_raw, uint8_t* o_compact, size_t dim \ + ) { \ + pack(o_raw, o_compact, dim); \ } RABITQ_PACK(2) RABITQ_PACK(3) @@ -173,32 +173,32 @@ RABITQ_PACK(7) #undef RABITQ_PACK void unpacking_excode_generic( - const uint8_t* compact, uint8_t* raw, size_t dim, size_t bits + const uint8_t* o_compact, uint8_t* o_raw, size_t dim, size_t bits ) { switch (bits) { case 1: - unpack<1>(compact, raw, dim); + unpack<1>(o_compact, o_raw, dim); break; case 2: - unpack<2>(compact, raw, dim); + unpack<2>(o_compact, o_raw, dim); break; case 3: - unpack<3>(compact, raw, dim); + unpack<3>(o_compact, o_raw, dim); break; case 4: - unpack<4>(compact, raw, dim); + unpack<4>(o_compact, o_raw, dim); break; case 5: - unpack<5>(compact, raw, dim); + unpack<5>(o_compact, o_raw, dim); break; case 6: - unpack<6>(compact, raw, dim); + unpack<6>(o_compact, o_raw, dim); break; case 7: - unpack<7>(compact, raw, dim); + unpack<7>(o_compact, o_raw, dim); break; case 8: - unpack<8>(compact, raw, dim); + unpack<8>(o_compact, o_raw, dim); break; default: throw std::invalid_argument("ex_bits must be in [1, 8]"); @@ -206,11 +206,11 @@ void unpacking_excode_generic( } namespace excode_ipimpl { -#define RABITQ_IP(Block, Bits) \ - float ip##Block##_fxu##Bits##_generic( \ - const float* query, const uint8_t* code, size_t dim \ - ) { \ - return excode_ip(query, code, dim); \ +#define RABITQ_IP(Block, Bits) \ + float ip##Block##_fxu##Bits##_generic( \ + const float* query, const uint8_t* compact_code, size_t dim \ + ) { \ + return excode_ip(query, compact_code, dim); \ } RABITQ_IP(16, 1) RABITQ_IP(64, 2) @@ -235,53 +235,54 @@ void scalar_quantize_uint16_generic( } void new_transpose_bin_generic( - const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { - for (size_t block = 0; block < dim; block += 64) { - const size_t width = std::min(size_t{64}, dim - block); - for (size_t b = 0; b < bits; ++b) { + for (size_t block = 0; block < padded_dim; block += 64) { + const size_t width = std::min(size_t{64}, padded_dim - block); + for (size_t b = 0; b < b_query; ++b) { uint64_t word = 0; for (size_t i = 0; i < width; ++i) { word |= uint64_t{(query[block + i] >> b) & 1U} << (width - 1 - i); } - *transposed++ = word; + *transposed_query++ = word; } } } void new_transpose_bin_512_generic( - const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { - for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = (std::min(size_t{512}, dim - block) + 63) / 64; - for (size_t b = 0; b < bits; ++b) { + for (size_t block = 0; block < padded_dim; block += 512) { + const size_t chunks = (std::min(size_t{512}, padded_dim - block) + 63) / 64; + for (size_t b = 0; b < b_query; ++b) { for (size_t chunk = 0; chunk < chunks; ++chunk) { - const size_t width = std::min(size_t{64}, dim - block - chunk * 64); + const size_t width = std::min(size_t{64}, padded_dim - block - chunk * 64); uint64_t word = 0; for (size_t i = 0; i < width; ++i) { word |= uint64_t{(query[block + chunk * 64 + i] >> b) & 1U} << (width - 1 - i); } - *transposed++ = word; + *transposed_query++ = word; } } } } -float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t dim) { +float mask_ip_x0_q_generic(const float* query, const uint8_t* data, size_t padded_dim) { double sum = 0; - for (size_t block = 0; block < dim; block += 64) { - const size_t width = std::min(size_t{64}, dim - block); + for (size_t block = 0; block < padded_dim; block += 64) { + const size_t width = std::min(size_t{64}, padded_dim - block); const uint64_t word = bitops::load_word(data + block / 8, width / 8); for (size_t i = 0; i < width; ++i) { - if ((word >> (width - 1 - i)) & 1U) + if ((word >> (width - 1 - i)) & 1U) { sum += query[block + i]; + } } } return static_cast(sum); } -float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t dim) { - return mask_ip_x0_q_generic(query, reinterpret_cast(data), dim); +float mask_ip_x0_q_generic(const float* query, const uint64_t* data, size_t padded_dim) { + return mask_ip_x0_q_generic(query, reinterpret_cast(data), padded_dim); } float warmup_ip_x0_q_512_generic( @@ -289,24 +290,25 @@ float warmup_ip_x0_q_512_generic( const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ) { - if (bits > 8) + if (b_query > 8) { throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); + } uint64_t ip = 0, count = 0; - for (size_t block = 0; block < dim; block += 512) { - const size_t chunks = (std::min(size_t{512}, dim - block) + 63) / 64; + for (size_t block = 0; block < padded_dim; block += 512) { + const size_t chunks = (std::min(size_t{512}, padded_dim - block) + 63) / 64; for (size_t chunk = 0; chunk < chunks; ++chunk) { - const size_t width = std::min(size_t{64}, dim - block - chunk * 64); + const size_t width = std::min(size_t{64}, padded_dim - block - chunk * 64); const uint64_t word = bitops::load_word(data + block / 8 + chunk * 8, width / 8); count += bitops::popcount64(word); - for (size_t b = 0; b < bits; ++b) { + for (size_t b = 0; b < b_query; ++b) { ip += uint64_t{bitops::popcount64(word & query[b * chunks + chunk])} << b; } } - query += chunks * bits; + query += chunks * b_query; } return delta * static_cast(ip) + vl * static_cast(count); } @@ -315,18 +317,19 @@ float warmup_ip_x0_q_512_generic( const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ) { return warmup_ip_x0_q_512_generic( - reinterpret_cast(data), query, delta, vl, dim, bits + reinterpret_cast(data), query, delta, vl, padded_dim, b_query ); } void flip_sign_generic(const uint8_t* flip, float* data, size_t dim) { for (size_t i = 0; i < dim; ++i) { - if ((flip[i / 8] >> (i % 8)) & 1U) + if ((flip[i / 8] >> (i % 8)) & 1U) { data[i] = -data[i]; + } } } void kacs_walk_generic(float* data, size_t len) { @@ -354,14 +357,17 @@ void fht_rotate_generic( flip_sign_generic(flip + pass * padded_dim / 8, rotated_data, padded_dim); const size_t start = pass % 2 == 0 ? 0 : padded_dim - trunc_dim; fht(rotated_data + start, trunc_dim); - for (size_t i = start; i < start + trunc_dim; ++i) + for (size_t i = start; i < start + trunc_dim; ++i) { rotated_data[i] *= scale_factor; - if (padded_dim != trunc_dim) + } + if (padded_dim != trunc_dim) { kacs_walk_generic(rotated_data, padded_dim); + } } if (padded_dim != trunc_dim) { - for (size_t i = 0; i < padded_dim; ++i) + for (size_t i = 0; i < padded_dim; ++i) { rotated_data[i] *= 0.25F; + } } } } // namespace rabitqlib::simd diff --git a/src/simd/quantization_avx2.cpp b/src/simd/quantization_avx2.cpp index 7d6f065f..7ee1d192 100644 --- a/src/simd/quantization_avx2.cpp +++ b/src/simd/quantization_avx2.cpp @@ -7,7 +7,9 @@ #include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" #include "rescale_search.hpp" + namespace rabitqlib::simd::detail { +namespace { // Evaluate one candidate scale, processing four coordinates per iteration. // For a_i = magnitudes[i], c_i = min(floor(scale*a_i), max_code), q_i = c_i+0.5, @@ -19,7 +21,7 @@ namespace rabitqlib::simd::detail { // unequal comparable level/magnitude ratios differ by at least 2^-33 relatively. // Reciprocal multiplication has much less error, so its extreme indices remain // valid. Recompute those two endpoints by division to retain exact event scales. -static RescaleSearchState evaluate_scale_state_avx2( +RescaleSearchState evaluate_scale_state_avx2( const double* magnitudes, const double* reciprocals, size_t dim, @@ -28,7 +30,7 @@ static RescaleSearchState evaluate_scale_state_avx2( double t_end, double scale ) { - __m256d z = _mm256_setzero_pd(), one = _mm256_set1_pd(1), half = _mm256_set1_pd(.5), + __m256d z = _mm256_setzero_pd(), one = _mm256_set1_pd(1), half = _mm256_set1_pd(0.5), vt = _mm256_set1_pd(scale), vcap = _mm256_set1_pd(max_code); __m256d guard = _mm256_set1_pd( 2 * std::numeric_limits::epsilon() * (max_code + 1.0) @@ -121,6 +123,7 @@ static RescaleSearchState evaluate_scale_state_avx2( return finish_scale_state(s, magnitudes, fi[fk], ni[nk], i, dim, max_code, scale); } +} // namespace } // namespace rabitqlib::simd::detail namespace rabitqlib::simd { diff --git a/src/simd/quantization_avx512.cpp b/src/simd/quantization_avx512.cpp index 67deca86..2ccec1b8 100644 --- a/src/simd/quantization_avx512.cpp +++ b/src/simd/quantization_avx512.cpp @@ -7,7 +7,9 @@ #include "rabitqlib/core/detail/simd/quantization_dispatch.hpp" #include "rabitqlib/utils/bitops.hpp" #include "rescale_search.hpp" + namespace rabitqlib::simd::detail { +namespace { // Evaluate one candidate scale, processing eight coordinates per iteration. // For a_i = magnitudes[i], c_i = min(floor(scale*a_i), max_code), q_i = c_i+0.5, @@ -19,7 +21,7 @@ namespace rabitqlib::simd::detail { // unequal comparable level/magnitude ratios differ by at least 2^-33 relatively. // Reciprocal multiplication has much less error, so its extreme indices remain // valid. Recompute those two endpoints by division to retain exact event scales. -static RescaleSearchState evaluate_scale_state_avx512( +RescaleSearchState evaluate_scale_state_avx512( const double* magnitudes, const double* reciprocals, size_t dim, @@ -28,7 +30,7 @@ static RescaleSearchState evaluate_scale_state_avx512( double t_end, double scale ) { - __m512d z = _mm512_setzero_pd(), one = _mm512_set1_pd(1), half = _mm512_set1_pd(.5), + __m512d z = _mm512_setzero_pd(), one = _mm512_set1_pd(1), half = _mm512_set1_pd(0.5), vt = _mm512_set1_pd(scale), vcap = _mm512_set1_pd(max_code); __m512d guard = _mm512_set1_pd( 2 * std::numeric_limits::epsilon() * (max_code + 1.0) @@ -113,6 +115,7 @@ static RescaleSearchState evaluate_scale_state_avx512( return finish_scale_state(s, magnitudes, fi[fk], ni[nk], i, dim, max_code, scale); } +} // namespace } // namespace rabitqlib::simd::detail namespace rabitqlib::simd { diff --git a/src/simd/quantization_generic.cpp b/src/simd/quantization_generic.cpp index 9b55cff4..72cb2b9d 100644 --- a/src/simd/quantization_generic.cpp +++ b/src/simd/quantization_generic.cpp @@ -18,10 +18,12 @@ RescaleScratch& get_thread_local_rescale_scratch(size_t dim) { // Retain the largest buffers on each worker; smaller vectors overwrite only // their active prefix without shrinking or zero-initializing the storage. thread_local RescaleScratch scratch; - if (scratch.magnitudes.size() < dim) + if (scratch.magnitudes.size() < dim) { scratch.magnitudes.resize(dim); - if (scratch.reciprocals.size() < dim) + } + if (scratch.reciprocals.size() < dim) { scratch.reciprocals.resize(dim); + } return scratch; } diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index a264e4b8..23d1fd04 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -67,41 +67,43 @@ void scalar_quantize_uint16_neon( } void new_transpose_bin_neon( - const uint16_t* query, uint64_t* transposed, size_t dim, size_t bits + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { - for (size_t block = 0; block < dim - dim % 64; block += 64) { + for (size_t block = 0; block < padded_dim - padded_dim % 64; block += 64) { uint16x8_t values[8]; for (size_t c = 0; c < 8; ++c) { values[c] = vld1q_u16(query + block + c * 8); } - for (size_t b = 0; b < bits; ++b) { + for (size_t b = 0; b < b_query; ++b) { const uint16x8_t bit = vdupq_n_u16(uint16_t{1} << b); uint64_t word = 0; for (auto value : values) { word = (word << 8) | byte_mask(vmovn_u16(vtstq_u16(value, bit))); } - *transposed++ = word; + *transposed_query++ = word; } } - if (dim % 64 != 0) { - new_transpose_bin_generic(query + dim - dim % 64, transposed, dim % 64, bits); + if (padded_dim % 64 != 0) { + new_transpose_bin_generic( + query + padded_dim - padded_dim % 64, transposed_query, padded_dim % 64, b_query + ); } } void new_transpose_bin_512_neon( - const uint8_t* query, uint64_t* transposed, size_t dim, size_t bits + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { - for (size_t block = 0; block < dim; block += 512) { - const size_t block_dim = std::min(size_t{512}, dim - block); + for (size_t block = 0; block < padded_dim; block += 512) { + const size_t block_dim = std::min(size_t{512}, padded_dim - block); const size_t full_chunks = block_dim / 64; const size_t chunks = (block_dim + 63) / 64; for (size_t c = 0; c < full_chunks; ++c) { const uint8_t* row = query + block + c * 64; const uint8x16_t a = vld1q_u8(row), b = vld1q_u8(row + 16), d = vld1q_u8(row + 32), e = vld1q_u8(row + 48); - for (size_t plane = 0; plane < bits; ++plane) { + for (size_t plane = 0; plane < b_query; ++plane) { const uint8x16_t bit = vdupq_n_u8(uint8_t{1} << plane); - transposed[plane * chunks + c] = + transposed_query[plane * chunks + c] = (uint64_t{byte_mask(vtstq_u8(a, bit))} << 48) | (uint64_t{byte_mask(vtstq_u8(b, bit))} << 32) | (uint64_t{byte_mask(vtstq_u8(d, bit))} << 16) | @@ -111,15 +113,15 @@ void new_transpose_bin_512_neon( const size_t tail_dim = block_dim % 64; if (tail_dim != 0) { const uint8_t* tail = query + block + full_chunks * 64; - for (size_t plane = 0; plane < bits; ++plane) { + for (size_t plane = 0; plane < b_query; ++plane) { uint64_t word = 0; for (size_t i = 0; i < tail_dim; ++i) { word |= uint64_t{(tail[i] >> plane) & 1U} << (tail_dim - 1 - i); } - transposed[plane * chunks + full_chunks] = word; + transposed_query[plane * chunks + full_chunks] = word; } } - transposed += chunks * bits; + transposed_query += chunks * b_query; } } } // namespace rabitqlib::simd diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index e0a44991..e4d5ffd5 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -36,12 +36,14 @@ struct RescaleSearchState { // Keep executable helpers local to each translation unit. This header is // compiled with different ISA flags; merging external inline/template symbols -// could make the AVX2 backend call an AVX512-compiled helper. +// could make the AVX2 backend call an AVX512-compiled helper. BoundedRescaleState +// is TU-local as well: its default member initializers give it a constructor. +namespace { // Recover exact event scales at the boundary coordinates a vector evaluator selected // (size_t max for none): it compares reciprocal products, while the heap sweep // divides. Then score coordinates [tail, dim) with the scalar definitions. -static inline RescaleSearchState finish_scale_state( +inline RescaleSearchState finish_scale_state( RescaleSearchState s, const double* magnitudes, size_t first_index, @@ -52,39 +54,44 @@ static inline RescaleSearchState finish_scale_state( double scale ) { constexpr size_t kNone = std::numeric_limits::max(); - if (first_index != kNone) + if (first_index != kNone) { s.first = std::max( s.first, static_cast( quantized_level_at_scale(magnitudes[first_index], scale, max_code) ) / magnitudes[first_index] ); - if (next_index != kNone) + } + if (next_index != kNone) { s.next = std::min( s.next, (quantized_level_at_scale(magnitudes[next_index], scale, max_code) + 1.0) / magnitudes[next_index] ); + } for (size_t i = tail; i < dim; ++i) { int c = quantized_level_at_scale(magnitudes[i], scale, max_code); - double q = c + .5; + double q = c + 0.5; s.numerator += magnitudes[i] * q; s.squared_norm += q * q; - if (c > 0) + if (c > 0) { s.first = std::max(s.first, static_cast(c) / magnitudes[i]); - if (c < max_code && magnitudes[i] > 0) + } + if (c < max_code && magnitudes[i] > 0) { s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); + } } return s; } // Finite IEEE-754 binary64 neighbors for outward-rounded bounds. // Adjusting the representation leaves the floating-point rounding mode unchanged. -static inline double next_double_up(double value) { +inline double next_double_up(double value) { static_assert(std::numeric_limits::is_iec559 && sizeof(double) == 8); // Both signed zeros step to the smallest positive subnormal. - if (value == 0) + if (value == 0) { return std::numeric_limits::denorm_min(); + } uint64_t bits = 0; std::memcpy(&bits, &value, sizeof(bits)); // Negative values have reversed bit ordering. @@ -93,10 +100,11 @@ static inline double next_double_up(double value) { return value; } -static inline double next_double_down(double value) { +inline double next_double_down(double value) { // Both signed zeros step to the smallest-magnitude negative subnormal. - if (value == 0) + if (value == 0) { return -std::numeric_limits::denorm_min(); + } uint64_t bits = 0; std::memcpy(&bits, &value, sizeof(bits)); bits += value > 0 ? static_cast(-1) : uint64_t{1}; @@ -104,9 +112,6 @@ static inline double next_double_down(double value) { return value; } -// Its default member initializers give it a constructor, so keep the type -// TU-local as well. -namespace { struct BoundedRescaleState { RescaleSearchState state{}; // Enclose the true dot product despite SIMD summation rounding. S is exact @@ -116,7 +121,6 @@ struct BoundedRescaleState { double score_lower = 0; double score_upper = 0; }; -} // namespace // Upper bound on squared cosine similarity for every code between two evaluated // states. Here a_i = input[i] is a rotated, normalized residual magnitude, and @@ -125,7 +129,7 @@ struct BoundedRescaleState { // so maximizing squared cosine also maximizes cosine, the heap sweep's objective. // This bounds the whole interval; comparing endpoint scores alone cannot // exclude a better interior code because the score need not be monotonic. -static inline double interval_squared_cosine_upper_bound( +inline double interval_squared_cosine_upper_bound( const BoundedRescaleState& left, const BoundedRescaleState& right ) { const double endpoints = std::max(left.score_upper, right.score_upper); @@ -133,8 +137,9 @@ static inline double interval_squared_cosine_upper_bound( // bounds are reversed, every interior numerator is already <= the left // upper bound, while its squared norm is >= the left norm; endpoints suffice. if (left.state.next == right.state.first || - left.numerator_upper >= right.numerator_upper) + left.numerator_upper >= right.numerator_upper) { return endpoints; + } // Derive a bound without enumerating the events inside this interval: // @@ -228,7 +233,7 @@ static inline double interval_squared_cosine_upper_bound( // Return false on ambiguous scores or work limits: the caller discards the // tentative result and runs the exact heap sweep, preserving its tie order. template -static bool try_find_best_scale_by_interval_search( +bool try_find_best_scale_by_interval_search( const float* input, size_t dim, int max_code, @@ -245,8 +250,9 @@ static bool try_find_best_scale_by_interval_search( } // S is a sum of quarter integers. This guard makes every addition exact. const uint64_t largest_odd = static_cast(2 * max_code + 1); - if (dim > (uint64_t{1} << 53) / (largest_odd * largest_odd)) + if (dim > (uint64_t{1} << 53) / (largest_odd * largest_odd)) { return false; + } // Promote float inputs exactly once and cache 1/a_i for repeated candidate // evaluations. The thread-local buffers are reused across vectors; each @@ -279,13 +285,15 @@ static bool try_find_best_scale_by_interval_search( const auto try_evaluate_scale_with_score_bounds = [&](double t, BoundedRescaleState& result) { - if (evaluations == kEvaluationBudget) + if (evaluations == kEvaluationBudget) { return false; + } ++evaluations; result.state = evaluator(magnitudes, reciprocals, dim, max_code, t_start, t_end, t); // Reject evaluator states whose plateau does not contain t. - if (result.state.first > t || result.state.next <= t) + if (result.state.first > t || result.state.next <= t) { return false; + } // If the computed numerator is N_hat and its relative error is at most // delta, then N_hat/(1+delta) <= N <= N_hat/(1-delta). Round outward // again when squaring/dividing so pruning never relies on a rounded-up @@ -322,8 +330,9 @@ static bool try_find_best_scale_by_interval_search( // t_end is excluded by the heap search too. Evaluate the preceding double // so an event exactly at t_end is not accidentally admitted here. if (!try_evaluate_scale_with_score_bounds(t_start, first) || - !try_evaluate_scale_with_score_bounds(next_double_down(t_end), last)) + !try_evaluate_scale_with_score_bounds(next_double_down(t_end), last)) { return false; + } const auto try_search_interval = [&](auto&& self, const BoundedRescaleState& left, @@ -334,10 +343,12 @@ static bool try_find_best_scale_by_interval_search( // mean equal codes. Adjacent plateaus have no intermediate states; // otherwise prune only when the entire interval is provably worse. if (left.state.squared_norm == right.state.squared_norm || - left.state.next == right.state.first || upper_bound < best_lower) + left.state.next == right.state.first || upper_bound < best_lower) { return true; - if (left.state.next > right.state.first || depth == kMaxDepth) + } + if (left.state.next > right.state.first || depth == kMaxDepth) { return false; + } // Split only [left.next, right.first); stay below the right plateau // despite rounding. @@ -345,13 +356,15 @@ static bool try_find_best_scale_by_interval_search( midpoint = std::max( left.state.next, std::min(midpoint, next_double_down(right.state.first)) ); - if (midpoint >= right.state.first) + if (midpoint >= right.state.first) { return false; + } BoundedRescaleState middle; if (!try_evaluate_scale_with_score_bounds(midpoint, middle) || middle.state.squared_norm <= left.state.squared_norm || - middle.state.squared_norm >= right.state.squared_norm) + middle.state.squared_norm >= right.state.squared_norm) { return false; + } // Exclude the middle plateau from both recursive gaps. const double left_upper = interval_squared_cosine_upper_bound(left, middle); @@ -370,8 +383,9 @@ static bool try_find_best_scale_by_interval_search( last, interval_squared_cosine_upper_bound(first, last), 0 - )) + )) { return false; + } // Return the first scale of the winning plateau, not the sampled midpoint, // to preserve the event sweep's canonical scale convention. best_t = best.state.first; @@ -380,7 +394,7 @@ static bool try_find_best_scale_by_interval_search( // The searched scale, or -1 to ask the caller to rerun the exact heap search. template -static double best_scale_or_fallback( +double best_scale_or_fallback( const float* input, size_t dim, int max_code, @@ -406,4 +420,5 @@ static double best_scale_or_fallback( : -1; } +} // namespace } // namespace rabitqlib::simd::detail diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index 191c05dd..ef934222 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -10,7 +10,7 @@ namespace rabitqlib::simd { void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim) { - // Process 32 floats (4 AVX2 registers) per iteration + // Process 32 floats (4 AVX2 registers) per iteration; requires dim % 32 == 0. constexpr size_t kFloatsPerChunk = 32; const __m256i bit_select = @@ -40,7 +40,7 @@ void flip_sign_avx2(const uint8_t* flip, float* data, size_t dim) { } void kacs_walk_avx2(float* data, size_t len) { - // ! len % 32 == 0; + // Requires len % 16 == 0: each half contains whole 8-float registers. for (size_t i = 0; i < len / 2; i += 8) { __m256 x = _mm256_loadu_ps(&data[i]); __m256 y = _mm256_loadu_ps(&data[i + (len / 2)]); diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index 6c0fc5ee..b0d11d00 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -77,6 +77,7 @@ void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { vec3 = _mm512_mask_xor_ps(vec3, mask3, vec3, sign_flip); _mm512_storeu_ps(&data[i + 48], vec3); } + // Requires dim % 32 == 0: a partial chunk holds exactly 32 floats. if (i < dim) { uint32_t bits; std::memcpy(&bits, flip + i / 8, sizeof(bits)); diff --git a/src/simd/rotator_kernels.hpp b/src/simd/rotator_kernels.hpp index 70f9b7af..743c6edc 100644 --- a/src/simd/rotator_kernels.hpp +++ b/src/simd/rotator_kernels.hpp @@ -37,7 +37,7 @@ inline void kacs_walk_rescale(float* data, size_t dim, float factor) { #endif } // Both backends and all compilers use the same AVX butterfly ordering. -template +template void fht_rotate_impl( const float* data, float* rotated_data, @@ -48,9 +48,9 @@ void fht_rotate_impl( const uint8_t* flip ) { void (*fht)(float*, float) = nullptr; -#define RABITQ_FHT_CASE(log_dim) \ - case size_t{1} << (log_dim): \ - fht = fht_intrinsics; \ +#define RABITQ_FHT_CASE(log_dim) \ + case size_t{1} << (log_dim): \ + fht = fht_impl; \ break; switch (trunc_dim) { RABITQ_FHT_CASE(6) @@ -74,23 +74,23 @@ void fht_rotate_impl( if (trunc_dim == padded_dim) { for (size_t round = 0; round < 4; ++round) { - flip_sign(flip + round * (padded_dim / 8), rotated_data, padded_dim); + FlipSign(flip + round * (padded_dim / 8), rotated_data, padded_dim); fht(rotated_data, scale_factor); } return; } const auto walk_and_flip = [&](const uint8_t* signs) { - if constexpr (kacs_flip != nullptr) { - kacs_flip(signs, rotated_data, padded_dim); + if constexpr (KacsFlip != nullptr) { + KacsFlip(signs, rotated_data, padded_dim); } else { - kacs_walk(rotated_data, padded_dim); - flip_sign(signs, rotated_data, padded_dim); + KacsWalk(rotated_data, padded_dim); + FlipSign(signs, rotated_data, padded_dim); } }; size_t start = padded_dim - trunc_dim; - flip_sign(flip, rotated_data, padded_dim); + FlipSign(flip, rotated_data, padded_dim); fht(rotated_data, scale_factor); walk_and_flip(flip + (padded_dim / 8)); fht(rotated_data + start, scale_factor); diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index 0a0fbcaa..80ea31ab 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -49,7 +49,7 @@ namespace { // rounds half away from zero. Signed inputs keep round_away_from_zero. Do not replace // either with adding 0.5f. template -void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta) { +void quantize(Code* result, const float* data, size_t dim, float lo, float delta) { if (delta == 0.0F) { std::fill_n(result, dim, Code{0}); return; @@ -61,7 +61,7 @@ void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta __m256 od256 = _mm256_set1_ps(one_over_delta); for (; i < mul8; i += 8) { - __m256 cur = _mm256_loadu_ps(&vec0[i]); + __m256 cur = _mm256_loadu_ps(&data[i]); cur = _mm256_mul_ps(_mm256_sub_ps(cur, lo256), od256); __m256i i32; if constexpr (NonNegative) { @@ -87,41 +87,41 @@ void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta } } for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); + result[i] = static_cast(std::round((data[i] - lo) * one_over_delta)); } } } // namespace void scalar_quantize_uint8_avx2( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void scalar_quantize_nonnegative_uint8_avx2( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void scalar_quantize_uint16_avx2( - uint16_t* result, const float* vec0, size_t dim, float lo, float delta + uint16_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void new_transpose_bin_avx2( - const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { // Reverse coordinates once; movemask then produces the stored bit order. const __m256i order = _mm256_broadcastsi128_si256( _mm_setr_epi8(14, 15, 12, 13, 10, 11, 8, 9, 6, 7, 4, 5, 2, 3, 0, 1) ); for (size_t i = 0; i < padded_dim - padded_dim % 64; i += 64) { - const __m256i a = _mm256_loadu_si256(reinterpret_cast(q)); - const __m256i b = _mm256_loadu_si256(reinterpret_cast(q + 16)); - const __m256i c = _mm256_loadu_si256(reinterpret_cast(q + 32)); - const __m256i d = _mm256_loadu_si256(reinterpret_cast(q + 48)); + const __m256i a = _mm256_loadu_si256(reinterpret_cast(query)); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(query + 16)); + const __m256i c = _mm256_loadu_si256(reinterpret_cast(query + 32)); + const __m256i d = _mm256_loadu_si256(reinterpret_cast(query + 48)); // Arrange eight-coordinate groups so packs produces the final bit order. __m256i hi_odd = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x20), order); __m256i hi_even = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x31), order); @@ -145,28 +145,29 @@ void new_transpose_bin_avx2( const uint64_t v = uint64_t{m0} | (uint64_t{m1} << 32); - tq[b_query - j - 1] = v; + transposed_query[b_query - j - 1] = v; hi_odd = _mm256_slli_epi16(hi_odd, 1); hi_even = _mm256_slli_epi16(hi_even, 1); lo_odd = _mm256_slli_epi16(lo_odd, 1); lo_even = _mm256_slli_epi16(lo_even, 1); } - tq += b_query; - q += 64; + transposed_query += b_query; + query += 64; } const size_t tail_dim = padded_dim % 64; // The partial word contains exactly 32 coordinates. if (tail_dim != 0) { - const __m256i a = _mm256_loadu_si256(reinterpret_cast(q)); - const __m256i b = _mm256_loadu_si256(reinterpret_cast(q + 16)); + const __m256i a = _mm256_loadu_si256(reinterpret_cast(query)); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(query + 16)); __m256i odd = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x20), order); __m256i even = _mm256_shuffle_epi8(_mm256_permute2x128_si256(b, a, 0x31), order); odd = _mm256_slli_epi16(odd, static_cast(16 - b_query)); even = _mm256_slli_epi16(even, static_cast(16 - b_query)); for (size_t bit = 0; bit < b_query; ++bit) { const __m256i bytes = _mm256_packs_epi16(even, odd); - tq[b_query - bit - 1] = static_cast(_mm256_movemask_epi8(bytes)); + transposed_query[b_query - bit - 1] = + static_cast(_mm256_movemask_epi8(bytes)); odd = _mm256_slli_epi16(odd, 1); even = _mm256_slli_epi16(even, 1); } @@ -174,7 +175,7 @@ void new_transpose_bin_avx2( } void new_transpose_bin_512_avx2( - const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { const __m128i reverse_bytes = _mm_setr_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); @@ -198,8 +199,8 @@ void new_transpose_bin_512_avx2( for (size_t k = 0; k < full_chunks; ++k) { // Load 64 bytes using two sequential 32-byte AVX2 registers - const uint8_t* current_q_lo = q + i + k * 64; - const uint8_t* current_q_hi = q + i + k * 64 + 32; + const uint8_t* current_q_lo = query + i + k * 64; + const uint8_t* current_q_hi = query + i + k * 64 + 32; __m256i vec_lo = _mm256_slli_epi16( load_reversed(current_q_lo), static_cast(8 - b_query) @@ -216,7 +217,7 @@ void new_transpose_bin_512_avx2( uint64_t m = (static_cast(m_lo) << 32) | m_hi; // Write into the 64-bit structured macro-layout - tq[(b_query - j - 1) * num_chunks + k] = m; + transposed_query[(b_query - j - 1) * num_chunks + k] = m; vec_lo = _mm256_slli_epi16(vec_lo, 1); vec_hi = _mm256_slli_epi16(vec_hi, 1); } @@ -224,17 +225,17 @@ void new_transpose_bin_512_avx2( if (tail_dim != 0) { __m256i values = _mm256_slli_epi16( - load_reversed(q + i + full_chunks * 64), static_cast(8 - b_query) + load_reversed(query + i + full_chunks * 64), static_cast(8 - b_query) ); for (size_t bit = 0; bit < b_query; ++bit) { - tq[(b_query - bit - 1) * num_chunks + full_chunks] = + transposed_query[(b_query - bit - 1) * num_chunks + full_chunks] = static_cast(_mm256_movemask_epi8(values)); values = _mm256_slli_epi16(values, 1); } } i += block_size; - tq += num_chunks * b_query; + transposed_query += num_chunks * b_query; } } diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 0594e4ef..f431ec5e 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -49,7 +49,7 @@ namespace { // rounds half away from zero. Signed inputs keep round_away_from_zero. Do not replace // either with adding 0.5f. template -void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta) { +void quantize(Code* result, const float* data, size_t dim, float lo, float delta) { if (delta == 0.0F) { std::fill_n(result, dim, Code{0}); return; @@ -60,7 +60,7 @@ void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta __m512 lo512 = _mm512_set1_ps(lo); __m512 od512 = _mm512_set1_ps(one_over_delta); for (; i < mul16; i += 16) { - __m512 cur = _mm512_loadu_ps(&vec0[i]); + __m512 cur = _mm512_loadu_ps(&data[i]); cur = _mm512_mul_ps(_mm512_sub_ps(cur, lo512), od512); __m512i i32; if constexpr (NonNegative) { @@ -85,31 +85,31 @@ void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta } } for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); + result[i] = static_cast(std::round((data[i] - lo) * one_over_delta)); } } } // namespace void scalar_quantize_uint8_avx512( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void scalar_quantize_nonnegative_uint8_avx512( - uint8_t* result, const float* vec0, size_t dim, float lo, float delta + uint8_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void scalar_quantize_uint16_avx512( - uint16_t* result, const float* vec0, size_t dim, float lo, float delta + uint16_t* result, const float* data, size_t dim, float lo, float delta ) { - quantize(result, vec0, dim, lo, delta); + quantize(result, data, dim, lo, delta); } void new_transpose_bin_avx512( - const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint16_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { // Reverse coordinates once; mask extraction then produces the stored bit order. const auto reverse_words = [](__m512i values) { @@ -121,8 +121,8 @@ void new_transpose_bin_avx512( }; // 512 / 16 = 32 for (size_t i = 0; i < padded_dim - padded_dim % 64; i += 64) { - __m512i vec_00_to_31 = reverse_words(_mm512_loadu_si512(q)); - __m512i vec_32_to_63 = reverse_words(_mm512_loadu_si512(q + 32)); + __m512i vec_00_to_31 = reverse_words(_mm512_loadu_si512(query)); + __m512i vec_32_to_63 = reverse_words(_mm512_loadu_si512(query + 32)); // the first (16 - b_query) bits are empty vec_00_to_31 = _mm512_slli_epi32(vec_00_to_31, (16 - b_query)); @@ -133,26 +133,26 @@ void new_transpose_bin_avx512( uint32_t v1 = _mm512_movepi16_mask(vec_32_to_63); // get most significant bit const uint64_t v = uint64_t{v1} | (uint64_t{v0} << 32); - tq[b_query - j - 1] = v; + transposed_query[b_query - j - 1] = v; vec_00_to_31 = _mm512_slli_epi16(vec_00_to_31, 1); vec_32_to_63 = _mm512_slli_epi16(vec_32_to_63, 1); } - tq += b_query; - q += 64; + transposed_query += b_query; + query += 64; } if (padded_dim % 64 != 0) { - __m512i values = reverse_words(_mm512_loadu_si512(q)); + __m512i values = reverse_words(_mm512_loadu_si512(query)); values = _mm512_slli_epi16(values, static_cast(16 - b_query)); for (size_t bit = 0; bit < b_query; ++bit) { - tq[b_query - bit - 1] = _mm512_movepi16_mask(values); + transposed_query[b_query - bit - 1] = _mm512_movepi16_mask(values); values = _mm512_slli_epi16(values, 1); } } } void new_transpose_bin_512_avx512( - const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query + const uint8_t* query, uint64_t* transposed_query, size_t padded_dim, size_t b_query ) { const auto reverse_bytes = [](__m512i values) { const __m512i order = _mm512_broadcast_i32x4( @@ -173,7 +173,7 @@ void new_transpose_bin_512_avx512( const size_t num_chunks = full_chunks + (tail_dim != 0); for (size_t k = 0; k < full_chunks; ++k) { - const uint8_t* current_q = q + i + k * 64; + const uint8_t* current_q = query + i + k * 64; const __m512i vec = reverse_bytes(_mm512_loadu_si512(current_q)); for (size_t j = 0; j < b_query; ++j) { @@ -183,23 +183,24 @@ void new_transpose_bin_512_avx512( 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] = static_cast(m); + transposed_query[(b_query - j - 1) * num_chunks + k] = + static_cast(m); } } if (tail_dim != 0) { constexpr __mmask64 valid = 0xFFFFFFFFULL; const __m512i values = - reverse_bytes(_mm512_maskz_loadu_epi8(valid, q + i + full_chunks * 64)); + reverse_bytes(_mm512_maskz_loadu_epi8(valid, query + i + full_chunks * 64)); for (size_t bit = 0; bit < b_query; ++bit) { const __m512i mask = _mm512_set1_epi8(static_cast(1U << bit)); const uint64_t selected = _mm512_test_epi8_mask(values, mask); - tq[bit * num_chunks + full_chunks] = selected >> 32; + transposed_query[bit * num_chunks + full_chunks] = selected >> 32; } } i += block_size; - tq += num_chunks * b_query; + transposed_query += num_chunks * b_query; } } diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 270661cd..ff754e64 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -5,24 +5,11 @@ #include "packed_tail_avx2.hpp" #include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/utils/bitops.hpp" namespace rabitqlib::simd::excode_ipimpl { namespace { -[[nodiscard]] inline uint64_t load_u64(const uint8_t* data) noexcept { - uint64_t value = 0; - std::memcpy(&value, data, sizeof(value)); - return value; -} - -[[nodiscard]] inline __m128i set_u64x(uint64_t high, uint64_t low) noexcept { - int64_t signed_high = 0; - int64_t signed_low = 0; - std::memcpy(&signed_high, &high, sizeof(signed_high)); - std::memcpy(&signed_low, &low, sizeof(signed_low)); - return _mm_set_epi64x(signed_high, signed_low); -} - // Accumulate the two halves independently to shorten the FMA dependency chain. inline void contribute_ip( __m128i vec, const float* __restrict__ query, __m256& sum, __m256& sum_hi @@ -101,7 +88,7 @@ float ip64_fxu2_avx2( compact_code += 16; } - simd::detail::accumulate_excode_tail_avx2<2>( + detail::accumulate_excode_tail_avx2<2>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); result = detail::reduce_add_avx2( @@ -124,7 +111,7 @@ float ip64_fxu3_avx2( __m128i compact2 = _mm_loadu_si128(reinterpret_cast(compact_code)); compact_code += 16; - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i vec_00_to_15 = _mm_and_si128(compact2, mask); @@ -133,13 +120,13 @@ float ip64_fxu3_avx2( __m128i vec_48_to_63 = _mm_and_si128(_mm_srli_epi16(compact2, 6), mask); __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit >> 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit >> 0), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit >> 3, top_bit >> 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 3, top_bit >> 2), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 5, top_bit >> 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 5, top_bit >> 4), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -151,7 +138,7 @@ float ip64_fxu3_avx2( contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - simd::detail::accumulate_excode_tail_avx2<3>( + detail::accumulate_excode_tail_avx2<3>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); result = detail::reduce_add_avx2( @@ -201,7 +188,7 @@ float ip64_fxu5_avx2( _mm_loadu_si128(reinterpret_cast(compact_code + 16)); compact_code += 32; - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i vec_00_to_15 = _mm_and_si128(compact4_1, mask); @@ -210,13 +197,13 @@ float ip64_fxu5_avx2( __m128i vec_48_to_63 = _mm_and_si128(_mm_srli_epi16(compact4_2, 4), mask); __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 3, top_bit << 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 3, top_bit << 4), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit >> 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit >> 0), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 3, top_bit >> 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 3, top_bit >> 2), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -228,7 +215,7 @@ float ip64_fxu5_avx2( contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - simd::detail::accumulate_excode_tail_avx2<5>( + detail::accumulate_excode_tail_avx2<5>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); result = detail::reduce_add_avx2( @@ -270,7 +257,7 @@ float ip64_fxu6_avx2( contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - simd::detail::accumulate_excode_tail_avx2<6>( + detail::accumulate_excode_tail_avx2<6>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); result = detail::reduce_add_avx2( @@ -307,17 +294,17 @@ float ip64_fxu7_avx2( _mm_srli_epi16(_mm_and_si128(cpt3, mask2), 2) ); - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 5, top_bit << 6), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 5, top_bit << 6), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit << 3, top_bit << 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 3, top_bit << 4), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit << 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit << 0), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -330,7 +317,7 @@ float ip64_fxu7_avx2( contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - simd::detail::accumulate_excode_tail_avx2<7>( + detail::accumulate_excode_tail_avx2<7>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); result = detail::reduce_add_avx2( @@ -341,14 +328,14 @@ float ip64_fxu7_avx2( } float ip16_fxu8_avx2( - const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { __m256 sum0 = _mm256_setzero_ps(); __m256 sum1 = _mm256_setzero_ps(); for (size_t i = 0; i < dim; i += 16) { - __m128i c8 = _mm_loadu_si128(reinterpret_cast(code)); + __m128i c8 = _mm_loadu_si128(reinterpret_cast(compact_code)); contribute_ip(c8, &query[i], sum0, sum1); - code += 16; + compact_code += 16; } return detail::reduce_add_avx2(_mm256_add_ps(sum0, sum1)); } diff --git a/src/simd/space_excode_avx512.cpp b/src/simd/space_excode_avx512.cpp index f5fdaccb..558d2928 100644 --- a/src/simd/space_excode_avx512.cpp +++ b/src/simd/space_excode_avx512.cpp @@ -6,6 +6,7 @@ #include "packed_tail_avx2.hpp" #include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/utils/bitops.hpp" namespace rabitqlib::simd::excode_ipimpl { @@ -37,7 +38,7 @@ inline __m512 excode_ip_tail_sum_avx512( const __m512i shifts = _mm512_load_si512(kOffsets.data()); const __m512i mask = _mm512_set1_epi32((1U << Bits) - 1); for (size_t i = 0; i < 32; i += 16) { - const __m128i packed = simd::detail::load_excode_tail(compact); + const __m128i packed = detail::load_excode_tail(compact); __m512i words; if constexpr (Bits == 2) { // All 16 codes fit in one word; broadcast it instead of gathering bytes. @@ -52,20 +53,6 @@ inline __m512 excode_ip_tail_sum_avx512( } return sum; } - -[[nodiscard]] inline uint64_t load_u64(const uint8_t* data) noexcept { - uint64_t value = 0; - std::memcpy(&value, data, sizeof(value)); - return value; -} - -[[nodiscard]] inline __m128i set_u64x(uint64_t high, uint64_t low) noexcept { - int64_t signed_high = 0; - int64_t signed_low = 0; - std::memcpy(&signed_high, &high, sizeof(signed_high)); - std::memcpy(&signed_low, &low, sizeof(signed_low)); - return _mm_set_epi64x(signed_high, signed_low); -} } // namespace // ip16_fxu1 consumes groups of 16 float query values and packed 1-bit codes. @@ -152,7 +139,7 @@ float ip64_fxu3_avx512( __m128i compact2 = _mm_loadu_si128(reinterpret_cast(compact_code)); compact_code += 16; - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i vec_00_to_15 = _mm_and_si128(compact2, mask); @@ -161,13 +148,13 @@ float ip64_fxu3_avx512( __m128i vec_48_to_63 = _mm_and_si128(_mm_srli_epi16(compact2, 6), mask); __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit >> 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit >> 0), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit >> 3, top_bit >> 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 3, top_bit >> 2), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 5, top_bit >> 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 5, top_bit >> 4), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -244,7 +231,7 @@ float ip64_fxu5_avx512( _mm_loadu_si128(reinterpret_cast(compact_code + 16)); compact_code += 32; - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i vec_00_to_15 = _mm_and_si128(compact4_1, mask); @@ -253,13 +240,13 @@ float ip64_fxu5_avx512( __m128i vec_48_to_63 = _mm_and_si128(_mm_srli_epi16(compact4_2, 4), mask); __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 3, top_bit << 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 3, top_bit << 4), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit >> 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit >> 0), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 3, top_bit >> 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 3, top_bit >> 2), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -372,17 +359,17 @@ float ip64_fxu7_avx512( _mm_srli_epi16(_mm_and_si128(cpt3, mask2), 2) ); - const uint64_t top_bit = load_u64(compact_code); + const uint64_t top_bit = bitops::load_word(compact_code, sizeof(uint64_t)); compact_code += 8; __m128i top_00_to_15 = - _mm_and_si128(set_u64x(top_bit << 5, top_bit << 6), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 5, top_bit << 6), top_mask); __m128i top_16_to_31 = - _mm_and_si128(set_u64x(top_bit << 3, top_bit << 4), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 3, top_bit << 4), top_mask); __m128i top_32_to_47 = - _mm_and_si128(set_u64x(top_bit << 1, top_bit << 2), top_mask); + _mm_and_si128(detail::set_u64x(top_bit << 1, top_bit << 2), top_mask); __m128i top_48_to_63 = - _mm_and_si128(set_u64x(top_bit >> 1, top_bit << 0), top_mask); + _mm_and_si128(detail::set_u64x(top_bit >> 1, top_bit << 0), top_mask); vec_00_to_15 = _mm_or_si128(top_00_to_15, vec_00_to_15); vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); @@ -418,12 +405,12 @@ float ip64_fxu7_avx512( } float ip16_fxu8_avx512( - const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim + const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; const auto unpack_code = [&](size_t k) { return _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(code + k)) + _mm_loadu_si128(reinterpret_cast(compact_code + k)) )); }; size_t i = 0; diff --git a/src/simd/space_float_kernels.hpp b/src/simd/space_float_kernels.hpp index 95521486..aaec5150 100644 --- a/src/simd/space_float_kernels.hpp +++ b/src/simd/space_float_kernels.hpp @@ -43,44 +43,49 @@ inline Vec tail_load(const float* p, size_t remaining) { enum class FloatOperation { SquaredL2, Dot, InnerProductDistance, SquaredNorm }; template -inline float raw_float(const float* a, const float* b, size_t n) { +inline float raw_float(const float* a, const float* b, size_t dim) { Vec s0 = zero(), s1 = zero(), s2 = zero(), s3 = zero(); size_t i = 0; auto accumulate = [](Vec sum, Vec x, Vec y) { if constexpr (Op == FloatOperation::SquaredL2) { Vec d = sub(x, y); return fma(d, d, sum); - } else if constexpr (Op == FloatOperation::SquaredNorm) + } else if constexpr (Op == FloatOperation::SquaredNorm) { return fma(x, x, sum); - else + } else { return fma(x, y, sum); + } }; auto step = [&](Vec sum, size_t at) { Vec x = load(a + at); - if constexpr (Op == FloatOperation::SquaredNorm) + if constexpr (Op == FloatOperation::SquaredNorm) { return accumulate(sum, x, x); - else + } else { return accumulate(sum, x, load(b + at)); + } }; - for (; n - i >= kWidth * 4; i += kWidth * 4) { + for (; dim - i >= kWidth * 4; i += kWidth * 4) { s0 = step(s0, i); s1 = step(s1, i + kWidth); s2 = step(s2, i + kWidth * 2); s3 = step(s3, i + kWidth * 3); } Vec sum = add(add(s0, s1), add(s2, s3)); - for (; n - i >= kWidth; i += kWidth) + for (; dim - i >= kWidth; i += kWidth) { sum = step(sum, i); - if (i < n) { - Vec x = tail_load(a + i, n - i); - if constexpr (Op == FloatOperation::SquaredNorm) + } + if (i < dim) { + Vec x = tail_load(a + i, dim - i); + if constexpr (Op == FloatOperation::SquaredNorm) { sum = accumulate(sum, x, x); - else - sum = accumulate(sum, x, tail_load(b + i, n - i)); + } else { + sum = accumulate(sum, x, tail_load(b + i, dim - i)); + } } float result = reduce(sum); - if constexpr (Op == FloatOperation::InnerProductDistance) + if constexpr (Op == FloatOperation::InnerProductDistance) { return 1.0F - result; + } return result; } } // namespace diff --git a/src/simd/space_generic.cpp b/src/simd/space_generic.cpp index 2a7e3766..4fceb27b 100644 --- a/src/simd/space_generic.cpp +++ b/src/simd/space_generic.cpp @@ -19,16 +19,19 @@ float product(const float* a, const float* b, size_t dim) { }; size_t i = 0; for (; dim - i >= sums.size(); i += sums.size()) { - for (size_t lane = 0; lane < sums.size(); ++lane) + for (size_t lane = 0; lane < sums.size(); ++lane) { sums[lane] += value(i + lane); + } } for (size_t stride = sums.size() / 2; stride != 0; stride /= 2) { - for (size_t lane = 0; lane < stride; ++lane) + for (size_t lane = 0; lane < stride; ++lane) { sums[lane] += sums[lane + stride]; + } } float result = sums[0]; - for (; i < dim; ++i) + for (; i < dim; ++i) { result += value(i); + } return result; } } // namespace diff --git a/src/simd/space_neon.cpp b/src/simd/space_neon.cpp index 66493b2e..48d5768c 100644 --- a/src/simd/space_neon.cpp +++ b/src/simd/space_neon.cpp @@ -29,8 +29,9 @@ float product(const float* a, const float* b, size_t dim) { s3 = step(s3, i + 12); } float32x4_t sum = vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)); - for (; dim - i >= 4; i += 4) + for (; dim - i >= 4; i += 4) { sum = step(sum, i); + } float result = vaddvq_f32(sum); for (; i < dim; ++i) { if constexpr (SquaredDistance) { diff --git a/src/simd/space_packed_neon.cpp b/src/simd/space_packed_neon.cpp index 7170368e..e79619ee 100644 --- a/src/simd/space_packed_neon.cpp +++ b/src/simd/space_packed_neon.cpp @@ -128,11 +128,11 @@ float packed_ip(const float* query, const uint8_t* compact, size_t dim) { } // namespace namespace excode_ipimpl { -#define RABITQ_NEON_IP(Block, Bits) \ - float ip##Block##_fxu##Bits##_neon( \ - const float* query, const uint8_t* code, size_t dim \ - ) { \ - return packed_ip(query, code, dim); \ +#define RABITQ_NEON_IP(Block, Bits) \ + float ip##Block##_fxu##Bits##_neon( \ + const float* query, const uint8_t* compact_code, size_t dim \ + ) { \ + return packed_ip(query, compact_code, dim); \ } RABITQ_NEON_IP(16, 1) RABITQ_NEON_IP(64, 2) @@ -145,32 +145,32 @@ RABITQ_NEON_IP(16, 8) #undef RABITQ_NEON_IP } // namespace excode_ipimpl -float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t dim) { - return detail::mask_ip_neon(query, data, dim); +float mask_ip_x0_q_neon(const float* query, const uint8_t* data, size_t padded_dim) { + return detail::mask_ip_neon(query, data, padded_dim); } -float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t dim) { - return detail::mask_ip_neon(query, reinterpret_cast(data), dim); +float mask_ip_x0_q_neon(const float* query, const uint64_t* data, size_t padded_dim) { + return detail::mask_ip_neon(query, reinterpret_cast(data), padded_dim); } float warmup_ip_x0_q_512_neon( const uint8_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ) { - return detail::warmup_ip_neon(data, query, delta, vl, dim, bits); + return detail::warmup_ip_neon(data, query, delta, vl, padded_dim, b_query); } float warmup_ip_x0_q_512_neon( const uint64_t* data, const uint64_t* query, float delta, float vl, - size_t dim, - size_t bits + size_t padded_dim, + size_t b_query ) { return detail::warmup_ip_neon( - reinterpret_cast(data), query, delta, vl, dim, bits + reinterpret_cast(data), query, delta, vl, padded_dim, b_query ); } } // namespace rabitqlib::simd diff --git a/src/simd/warmup_avx2.cpp b/src/simd/warmup_avx2.cpp index d0e8ed09..791d1a12 100644 --- a/src/simd/warmup_avx2.cpp +++ b/src/simd/warmup_avx2.cpp @@ -15,11 +15,13 @@ float warmup_ip_x0_q_512_avx2( size_t padded_dim, size_t b_query ) { - if (b_query > 8) { - throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); - } size_t ip_scalar = 0; size_t ppc_scalar = 0; + + constexpr size_t kMaxQueryBits = 8; + if (b_query > kMaxQueryBits) { + throw std::invalid_argument("warmup_ip_x0_q_512 requires at most 8 query bits"); + } if (padded_dim < 512) { detail::accumulate_warmup_tail( data, query, padded_dim, b_query, ip_scalar, ppc_scalar @@ -28,7 +30,9 @@ float warmup_ip_x0_q_512_avx2( (vl * static_cast(ppc_scalar)); } - return detail::warmup_blocks_avx2<8>(data, query, delta, vl, padded_dim, b_query); + return detail::warmup_blocks_avx2( + data, query, delta, vl, padded_dim, b_query + ); } float warmup_ip_x0_q_512_avx2( diff --git a/src/simd/warmup_kernels.hpp b/src/simd/warmup_kernels.hpp index 70bcfe34..77d019a3 100644 --- a/src/simd/warmup_kernels.hpp +++ b/src/simd/warmup_kernels.hpp @@ -9,12 +9,12 @@ #include "rabitqlib/utils/bitops.hpp" namespace rabitqlib::simd::detail { - -// Keep ISA-specific instantiations local to each backend translation unit. +// Internal linkage: backends compiled with different ISA flags keep their own bodies. +namespace { // A partial 512-coordinate block contains at most eight words per query plane. // Scalar POPCNT avoids staging/zeroing two full vectors for each short plane. -static inline void accumulate_warmup_tail( +inline void accumulate_warmup_tail( const uint8_t* data, const uint64_t* query, size_t dim, @@ -43,7 +43,7 @@ static inline void accumulate_warmup_tail( } } -static inline __m256i popcount_avx2(__m256i v) { +inline __m256i popcount_avx2(__m256i v) { // Lookup table for population count of 0-15 const __m256i lookup = _mm256_broadcastsi128_si256( _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4) @@ -66,7 +66,7 @@ static inline __m256i popcount_avx2(__m256i v) { } template -static inline float warmup_blocks_avx2( +inline float warmup_blocks_avx2( const uint8_t* data, const uint64_t* query, float delta, @@ -140,7 +140,7 @@ static inline float warmup_blocks_avx2( #if defined(__AVX512F__) || defined(_MSC_VER) template -static inline float warmup_blocks_avx512( +inline float warmup_blocks_avx512( const uint8_t* data, const uint64_t* query, float delta, @@ -177,7 +177,7 @@ static inline float warmup_blocks_avx512( size_t remaining_dim = padded_dim - i; if (remaining_dim > 0) { const size_t num_chunks = (remaining_dim + 63) / 64; - const auto valid_mask = static_cast<__mmask8>((1u << num_chunks) - 1u); + const auto valid_mask = static_cast<__mmask8>((1U << num_chunks) - 1U); // Byte masking zero-extends the partial last word without staging short // stores through a stack buffer. The query planes already use whole words. const __m512i data_vec = @@ -205,4 +205,5 @@ static inline float warmup_blocks_avx512( } #endif +} // namespace } // namespace rabitqlib::simd::detail diff --git a/src/utils/cpu_features.cpp b/src/utils/cpu_features.cpp index 099b6128..dc769c09 100644 --- a/src/utils/cpu_features.cpp +++ b/src/utils/cpu_features.cpp @@ -148,8 +148,8 @@ bool supports_avx512_core(const Features& detected) { } // namespace detail const Features& features() { - static const Features detected = detect_features(); - return detected; + static const Features kDetected = detect_features(); + return kDetected; } bool has_neon() { return features().neon; } diff --git a/tests/README.md b/tests/README.md index c9f7a06b..6793de65 100644 --- a/tests/README.md +++ b/tests/README.md @@ -269,7 +269,7 @@ variants compile and execute the installed standalone core benchmark smoke tests CI covers both choices on its existing platform matrix; local results do not imply that remote jobs have already passed. -The normal suite includes `CoreApi.*`, the frozen `CoreApiV1Consumer`, and separate +The normal suite includes `CoreApiTest.*`, the frozen `CoreApiV1Consumer`, and separate `core_api_no_alloc_dispatch` / `core_api_no_alloc_scalar` processes. The latter force and count cold-call encoding, layout conversion, query preparation and estimation; Linux additionally wraps malloc/calloc/realloc/aligned allocation. The scalar process diff --git a/tests/allocation/core_api_no_alloc.cpp b/tests/allocation/core_api_no_alloc.cpp index 3efa519b..c861933c 100644 --- a/tests/allocation/core_api_no_alloc.cpp +++ b/tests/allocation/core_api_no_alloc.cpp @@ -38,38 +38,45 @@ void* __real_realloc(void*, size_t); void* __real_aligned_alloc(size_t, size_t); int __real_posix_memalign(void**, size_t, size_t); void* __wrap_malloc(size_t n) { - if (watching) + if (watching) { ++c_calls; + } return __real_malloc(n); } void* __wrap_calloc(size_t n, size_t width) { - if (watching) + if (watching) { ++c_calls; + } return __real_calloc(n, width); } void* __wrap_realloc(void* p, size_t n) { - if (watching) + if (watching) { ++c_calls; + } return __real_realloc(p, n); } void* __wrap_aligned_alloc(size_t alignment, size_t n) { - if (watching) + if (watching) { ++c_calls; + } return __real_aligned_alloc(alignment, n); } int __wrap_posix_memalign(void** p, size_t alignment, size_t n) { - if (watching) + if (watching) { ++c_calls; + } return __real_posix_memalign(p, alignment, n); } } #endif TEST_NOINLINE void* operator new(size_t n) { - if (watching) + if (watching) { ++new_calls; - if (void* p = std::malloc(n == 0 ? 1 : n)) + } + if (void* p = std::malloc(n == 0 ? 1 : n)) { return p; + } throw std::bad_alloc(); } TEST_NOINLINE void* operator new[](size_t n) { return ::operator new(n); } @@ -78,15 +85,18 @@ TEST_NOINLINE void operator delete[](void* p) noexcept { ::operator delete(p); } TEST_NOINLINE void operator delete(void* p, size_t) noexcept { ::operator delete(p); } TEST_NOINLINE void operator delete[](void* p, size_t) noexcept { ::operator delete(p); } TEST_NOINLINE void* operator new(size_t n, std::align_val_t alignment) { - if (watching) + if (watching) { ++new_calls; + } const auto a = static_cast(alignment); #ifdef _MSC_VER - if (void* p = _aligned_malloc(n == 0 ? 1 : n, a)) + if (void* p = _aligned_malloc(n == 0 ? 1 : n, a)) { return p; + } #else - if (void* p = std::aligned_alloc(a, ((n == 0 ? 1 : n) + a - 1) / a * a)) + if (void* p = std::aligned_alloc(a, ((n == 0 ? 1 : n) + a - 1) / a * a)) { return p; + } #endif throw std::bad_alloc(); } @@ -123,13 +133,13 @@ bool has_neon() { return false; } int main() { namespace api = rabitqlib::quant::v1; - constexpr size_t max_dim = 4096; - alignas(std::max_align_t) std::array memory{}; - std::array data{}, centroid{}, query{}, rotated{}; - std::array raw{}, restored{}; - std::array signs{}; - std::array extra{}; - for (size_t i = 0; i < max_dim; ++i) { + constexpr size_t kMaxDim = 4096; + alignas(std::max_align_t) std::array memory{}; + std::array data{}, centroid{}, query{}, rotated{}; + std::array raw{}, restored{}; + std::array signs{}; + std::array extra{}; + for (size_t i = 0; i < kMaxDim; ++i) { data[i] = static_cast(static_cast(i % 19) - 9) / 11; centroid[i] = static_cast(static_cast(i % 7) - 3) / 13; query[i] = static_cast(static_cast(i % 23) - 11) / 17; @@ -145,8 +155,9 @@ int main() { Watch guard; r = api::requirements(p, api::CodeLayout::SplitPacked); } - if (r.workspace_bytes > memory.size()) + if (r.workspace_bytes > memory.size()) { return 1; + } api::Buffer workspace{memory.data(), r.workspace_bytes}; api::CodeOutput a{ api::CodeLayout::UnpackedU16, {raw.data(), raw.size()}, {}}; @@ -195,8 +206,9 @@ int main() { q, {a.layout, {raw.data(), raw.size()}, {}}, f ); } - if (!std::isfinite(sink)) + if (!std::isfinite(sink)) { return 2; + } } } } diff --git a/tests/dispatch/fastscan_backend_test.cpp b/tests/dispatch/fastscan_backend_test.cpp index 33278018..1327514f 100644 --- a/tests/dispatch/fastscan_backend_test.cpp +++ b/tests/dispatch/fastscan_backend_test.cpp @@ -18,7 +18,7 @@ namespace rabitqlib::fastscan { namespace { -TEST(FastScanBackendContract, SelectedHaccLayoutMatchesProducerAndConsumer) { +TEST(FastScanBackendContractTest, SelectedHaccLayoutMatchesProducerAndConsumer) { const auto& backend = detail::selected_backend(); const size_t tables_per_plane = backend.hacc_layout == detail::HaccLutLayout::Quad64 ? 4 : 2; @@ -60,7 +60,7 @@ TEST(FastScanBackendContract, SelectedHaccLayoutMatchesProducerAndConsumer) { } } -TEST(FastScanBackendContract, StandardLutIsSharedAcrossSupportedConsumers) { +TEST(FastScanBackendContractTest, StandardLutIsSharedAcrossSupportedConsumers) { constexpr size_t kDim = 64; const auto& backend = detail::selected_backend(); std::array query{}; @@ -100,23 +100,27 @@ TEST(FastScanBackendContract, StandardLutIsSharedAcrossSupportedConsumers) { #endif } -TEST(FastScanBackendContract, ExplicitCorrectionBackendsConsumeHostLut) { +TEST(FastScanBackendContractTest, ExplicitCorrectionBackendsConsumeHostLut) { const auto& selected = detail::selected_backend(); using Estimator = decltype(selected.split); std::vector backends{selected.split}; #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) + if (cpu::has_avx2()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx2); - if (cpu::has_avx512_core()) + } + if (cpu::has_avx512_core()) { backends.push_back(rabitqlib::simd::split_batch_estdist_avx512); + } #endif for (size_t dim : {32U, 96U, 1040U}) { std::vector values(dim); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { values[i] = static_cast(i % 23) / 7 - 1; + } for (bool hacc : {false, true}) { - if (!hacc && !cpu::has_avx2() && !cpu::has_neon()) + if (!hacc && !cpu::has_avx2() && !cpu::has_neon()) { continue; + } for (auto metric : {METRIC_L2, METRIC_IP}) { SCOPED_TRACE( ::testing::Message() @@ -126,8 +130,9 @@ TEST(FastScanBackendContract, ExplicitCorrectionBackendsConsumeHostLut) { query.set_g_add(3, 2); std::vector storage(BatchDataMap::data_bytes(dim)); BatchDataMap batch(storage.data(), dim); - for (size_t i = 0; i < dim * 4; ++i) + for (size_t i = 0; i < dim * 4; ++i) { batch.bin_code()[i] = static_cast(i * 73 + i / 7); + } for (size_t lane = 0; lane < 32; ++lane) { batch.f_add()[lane] = static_cast(lane) + 10; batch.f_rescale()[lane] = lane % 2 == 0 ? 0.125F : -0.25F; diff --git a/tests/dispatch/fastscan_no_simd_test.cpp b/tests/dispatch/fastscan_no_simd_test.cpp index cce0a827..39bdc1a2 100644 --- a/tests/dispatch/fastscan_no_simd_test.cpp +++ b/tests/dispatch/fastscan_no_simd_test.cpp @@ -48,7 +48,7 @@ TEST(FastScanDispatchTest, PointCodeAssignmentPropagatesUnsupportedCpuFromWorker labels.fill(kPidMax); std::array distances; distances.fill(-1.0F); - rabitqkmeans::detail::RaBitQAssigner assigner( + rabitqkmeans::detail::RaBitQAssignment assigner( kDim, kClusters, data.data(), kPoints, 2, rabitqkmeans::RaBitQKMeansParameters{} ); try { diff --git a/tests/fixtures/persistence/generate.cpp b/tests/fixtures/persistence/generate.cpp index f10e98ed..c0a096f7 100644 --- a/tests/fixtures/persistence/generate.cpp +++ b/tests/fixtures/persistence/generate.cpp @@ -31,13 +31,16 @@ int main(int argc, char** argv) { return static_cast(static_cast((state >> 8) % 2001) - 1000) / 256.0F; }; std::vector data(kCount * kDim), queries(kQueries * kDim), centroid(kDim); - for (auto& value : data) + for (auto& value : data) { value = next(); - for (auto& value : queries) + } + for (auto& value : queries) { value = next(); + } for (size_t i = 0; i < kCount; ++i) { - for (size_t d = 0; d < kDim; ++d) + for (size_t d = 0; d < kDim; ++d) { centroid[d] += data[i * kDim + d] / kCount; + } } std::vector clusters(kCount, 0); for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -89,10 +92,12 @@ int main(int argc, char** argv) { expected.exceptions(std::ios::failbit | std::ios::badbit); expected << std::setprecision(std::numeric_limits::max_digits10); expected << kCount << ' ' << kDim << ' ' << kQueries << ' ' << kTopK << '\n'; - for (size_t i = 0; i < queries.size(); ++i) + for (size_t i = 0; i < queries.size(); ++i) { expected << queries[i] << (i + 1 == queries.size() ? '\n' : ' '); - for (size_t i = 0; i < ids.size(); ++i) + } + for (size_t i = 0; i < ids.size(); ++i) { expected << ids[i] << ' ' << distances[i] << '\n'; + } expected.close(); } } diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp index e27357c9..e8fadda7 100644 --- a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp @@ -30,16 +30,20 @@ int main(int argc, char** argv) { }; std::vector data(kCount * kDim), queries(kQueries * kDim); std::vector centroids(kClusters * kDim); - for (auto& value : data) + for (auto& value : data) { value = next(); - for (auto& value : queries) + } + for (auto& value : queries) { value = next(); - for (auto& value : centroids) + } + for (auto& value : centroids) { value = next(); + } std::vector clusters(kCount, 3); clusters[0] = 1; - for (size_t i = 1; i < 33; ++i) + for (size_t i = 1; i < 33; ++i) { clusters[i] = 2; + } for (const std::string kind : {"ivf4", "ivf32"}) { const std::string name = kind + "_l2"; const auto path = (directory / (name + ".index")).string(); @@ -57,10 +61,12 @@ int main(int argc, char** argv) { expected.exceptions(std::ios::failbit | std::ios::badbit); expected << std::setprecision(std::numeric_limits::max_digits10); expected << kCount << ' ' << kDim << ' ' << kQueries << ' ' << kTopK << '\n'; - for (size_t i = 0; i < queries.size(); ++i) + for (size_t i = 0; i < queries.size(); ++i) { expected << queries[i] << (i + 1 == queries.size() ? '\n' : ' '); - for (size_t i = 0; i < ids.size(); ++i) + } + for (size_t i = 0; i < ids.size(); ++i) { expected << ids[i] << ' ' << distances[i] << '\n'; + } expected.close(); } return 0; diff --git a/tests/python/test_historical_persistence.py b/tests/python/test_historical_persistence.py index 7cb9c006..d6318af8 100644 --- a/tests/python/test_historical_persistence.py +++ b/tests/python/test_historical_persistence.py @@ -88,7 +88,7 @@ def test_unknown_historical_version_is_rejected(name, tmp_path): message = ( "Unsupported raw IVF index version" if cls is IvfIndex - else "Unsupported QuantizedGraph file version" + else "QuantizedGraph file version is unsupported" ) with pytest.raises(RuntimeError, match=f"^{message}$"): cls.load(str(destination)) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index 6af01f38..a5bd35c3 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -199,13 +199,13 @@ def test_search_before_build_raises(query_data): def test_save_rejects_unopenable_destination(built_hnsw, tmp_path): - with pytest.raises(RuntimeError, match="HNSW: cannot open index file for writing"): + with pytest.raises(RuntimeError, match="HNSW cannot open index file for writing"): built_hnsw.save(str(tmp_path / "missing" / "hnsw.index")) @pytest.mark.skipif(not Path("/dev/full").exists(), reason="/dev/full is unavailable") def test_save_reports_write_or_close_failure(built_hnsw): - with pytest.raises(RuntimeError, match="HNSW: failed to write index file"): + with pytest.raises(RuntimeError, match="HNSW failed to write index file"): built_hnsw.save("/dev/full") diff --git a/tests/python/test_import.py b/tests/python/test_import.py index 612cb43b..806da025 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -48,5 +48,5 @@ def test_constructor_contract(cls, args, kwargs, expected): ], ) def test_invalid_metric_raises(cls, args): - with pytest.raises(ValueError, match=r"^Unsupported metric\. Use 'l2' or 'ip'\.$"): + with pytest.raises(ValueError, match=r"^metric must be 'l2' or 'ip'$"): cls(*args, metric="cosine") diff --git a/tests/python/test_portable_indexes.py b/tests/python/test_portable_indexes.py index 83896164..cf9e9db0 100644 --- a/tests/python/test_portable_indexes.py +++ b/tests/python/test_portable_indexes.py @@ -122,8 +122,8 @@ def test_rejects_development_only_padding(tmp_path, kind, bits): cls = {"ivf": IvfIndex, "hnsw": HnswIndex, "qg": SymqgIndex}[kind] message = { "ivf": "Invalid padded dimension in IVF index file", - "hnsw": "HNSW: invalid or truncated index file", - "qg": "Invalid padded dimension in quantized graph file", + "hnsw": "HNSW index file is invalid or truncated", + "qg": "QuantizedGraph file has an invalid padded dimension", }[kind] with pytest.raises(RuntimeError, match=f"^{message}$"): cls.load(str(path)) diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 921e24f8..a558de70 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -161,6 +161,14 @@ def test_wrong_query_dim_raises(built_symqg): built_symqg.search(bad_queries, k=1, ef=_EF) +def test_save_before_build_raises(tmp_path): + idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) + with pytest.raises( + RuntimeError, match="^SymqgIndex must be built or loaded before save$" + ): + idx.save(str(tmp_path / "unbuilt.index")) + + def test_search_before_build_raises(): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) queries = np.zeros((5, DIM), dtype=np.float32) diff --git a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp index 1d079b79..94325d28 100644 --- a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp +++ b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp @@ -36,12 +36,14 @@ class RecordedAssignment { ) : dim_(d), clusters_(k), points_(n), parameters_(parameters) {} - void assign(const float* centroids, PID* labels, float* distances, bool final = false) { + void assign( + const float* centroids, PID* assignments, float* distances, bool final = false + ) { if (!final) { parameters_.initial_centroids->assign(centroids, centroids + clusters_ * dim_); } for (size_t point = 0; point < points_; ++point) { - labels[point] = static_cast(point % parameters_.occupied_clusters); + assignments[point] = static_cast(point % parameters_.occupied_clusters); distances[point] = 0; } } diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index 2d2be63b..8d2401c5 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -148,13 +148,13 @@ TEST(QGKMeansTest, UsesRawGraphStorageByDefaultAndRejectsUnsupportedBits) { } TEST(QGAssignerTest, QuantizedSearchReturnsExactDistancesToOriginalCentroids) { - constexpr size_t kDimension = 65; // Exercise the padded quantized domain. + constexpr size_t kDim = 65; // Exercise the padded quantized domain. constexpr size_t kClusters = 33; - constexpr size_t kQueries = 20; + constexpr size_t kNumQueries = 20; std::mt19937 rng(123); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector centroids(kClusters * kDimension); - std::vector queries(kQueries * kDimension); + std::vector centroids(kClusters * kDim); + std::vector queries(kNumQueries * kDim); std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); std::generate(queries.begin(), queries.end(), [&] { return distribution(rng); }); @@ -162,31 +162,35 @@ TEST(QGAssignerTest, QuantizedSearchReturnsExactDistancesToOriginalCentroids) { SCOPED_TRACE(bits); for (const MetricType metric : {METRIC_L2, METRIC_IP}) { SCOPED_TRACE(static_cast(metric)); - QGAssigner assigner(kDimension, kClusters, 32, 33, 33, 2, 1, 42, metric, bits); - std::vector labels(kQueries); - std::vector distances(kQueries); + QGAssigner assigner(kDim, kClusters, 32, 33, 33, 2, 1, 42, metric, bits); + std::vector assignments(kNumQueries); + std::vector distances(kNumQueries); assigner.assign( - centroids.data(), queries.data(), kQueries, labels.data(), distances.data() + centroids.data(), + queries.data(), + kNumQueries, + assignments.data(), + distances.data() ); symqg::QuantizedGraph graph( - kClusters, kDimension, 32, metric, RotatorType::FhtKacRotator, bits, 42 + kClusters, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 ); symqg::QGBuilder builder(graph, 33, centroids.data(), 1, 42); builder.build(2); graph.set_ef(33); bool saw_estimate_difference = false; - for (size_t i = 0; i < kQueries; ++i) { - PID expected_label = kPidMax; + for (size_t i = 0; i < kNumQueries; ++i) { + PID expected_cluster_id = kPidMax; float estimated_distance = 0; graph.search( - queries.data() + i * kDimension, 1, &expected_label, &estimated_distance + queries.data() + i * kDim, 1, &expected_cluster_id, &estimated_distance ); - ASSERT_EQ(labels[i], expected_label); + ASSERT_EQ(assignments[i], expected_cluster_id); double exact_distance = metric == METRIC_L2 ? 0.0 : 1.0; - for (size_t dim = 0; dim < kDimension; ++dim) { - const double x = queries[i * kDimension + dim]; - const double c = centroids[labels[i] * kDimension + dim]; + for (size_t dim = 0; dim < kDim; ++dim) { + const double x = queries[i * kDim + dim]; + const double c = centroids[assignments[i] * kDim + dim]; exact_distance += metric == METRIC_L2 ? (x - c) * (x - c) : -x * c; } EXPECT_NEAR(distances[i], exact_distance, 1e-5); @@ -199,13 +203,13 @@ TEST(QGAssignerTest, QuantizedSearchReturnsExactDistancesToOriginalCentroids) { } TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) { - constexpr size_t kDimension = 65; + constexpr size_t kDim = 65; constexpr size_t kClusters = 65; - constexpr size_t kQueries = 129; + constexpr size_t kNumQueries = 129; constexpr uint32_t kEf = 8; std::mt19937 rng(823); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector initial(kClusters * kDimension), queries(kQueries * kDimension); + std::vector initial(kClusters * kDim), queries(kNumQueries * kDim); std::generate(initial.begin(), initial.end(), [&] { return distribution(rng); }); std::generate(queries.begin(), queries.end(), [&] { return distribution(rng); }); for (const uint32_t bits : {0U, 4U, 8U}) { @@ -216,10 +220,10 @@ TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) ); auto centroids = initial; QGAssigner assigner( - kDimension, kClusters, 32, 64, kEf, 1, threads, 42, metric, bits + kDim, kClusters, 32, 64, kEf, 1, threads, 42, metric, bits ); symqg::QuantizedGraph graph( - kClusters, kDimension, 32, metric, RotatorType::FhtKacRotator, bits, 42 + kClusters, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 ); symqg::QGBuilder builder( graph, @@ -239,7 +243,7 @@ TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) visited.initialize(kClusters, std::min(size_t{kEf} * kEf, kClusters / 10)); std::vector previous; for (const size_t count : - {kQueries, kQueries, kQueries - 2, kQueries - 2}) { + {kNumQueries, kNumQueries, kNumQueries - 2, kNumQueries - 2}) { for (float& value : centroids) { value += 0.01F; } @@ -260,7 +264,7 @@ TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) const PID hint = previous.size() == count ? previous[point] : kPidMax; graph.search_with_scratch( - queries.data() + point * kDimension, + queries.data() + point * kDim, 1, &expected[point], &expected_distances[point], @@ -275,9 +279,9 @@ TEST(QGAssignerTest, GroupedQueriesPreserveIndependentSearchHintsAcrossRebuilds) ); if (bits != 0) { expected_distances[point] = distance( - queries.data() + point * kDimension, - centroids.data() + expected[point] * kDimension, - kDimension + queries.data() + point * kDim, + centroids.data() + expected[point] * kDim, + kDim ); } } @@ -330,35 +334,35 @@ std::vector separated_points(size_t num_points, size_t dimension) { } TEST(QGAssignerTest, ReusesVisitedStateAcrossQueriesAndAssignerSizes) { - constexpr size_t kDimension = 64; - constexpr size_t kQueries = 256; + constexpr size_t kDim = 64; + constexpr size_t kNumQueries = 256; for (const uint32_t threads : {1U, 4U}) { SCOPED_TRACE(threads); for (const size_t clusters : {33U, 65U, 33U}) { SCOPED_TRACE(clusters); - const auto centroids = separated_points(clusters, kDimension); - std::vector queries(kQueries * kDimension); - std::vector labels(kQueries); - std::vector distances(kQueries); - QGAssigner assigner(kDimension, clusters, 32, 128, 128, 2, threads, 42); + const auto centroids = separated_points(clusters, kDim); + std::vector queries(kNumQueries * kDim); + std::vector assignments(kNumQueries); + std::vector distances(kNumQueries); + QGAssigner assigner(kDim, clusters, 32, 128, 128, 2, threads, 42); for (size_t round = 0; round < 2; ++round) { - for (size_t point = 0; point < kQueries; ++point) { + for (size_t point = 0; point < kNumQueries; ++point) { const size_t cluster = (point + round) % clusters; std::copy_n( - centroids.data() + cluster * kDimension, - kDimension, - queries.data() + point * kDimension + centroids.data() + cluster * kDim, + kDim, + queries.data() + point * kDim ); } assigner.assign( centroids.data(), queries.data(), - kQueries, - labels.data(), + kNumQueries, + assignments.data(), distances.data() ); - for (size_t point = 0; point < kQueries; ++point) { - EXPECT_EQ(labels[point], (point + round) % clusters); + for (size_t point = 0; point < kNumQueries; ++point) { + EXPECT_EQ(assignments[point], (point + round) % clusters); EXPECT_EQ(distances[point], 0.0F); } } @@ -368,15 +372,15 @@ TEST(QGAssignerTest, ReusesVisitedStateAcrossQueriesAndAssignerSizes) { TEST(QGKMeansTest, WarnsWhenThereAreTooFewPointsPerCentroid) { constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - const std::vector data = separated_points(kClusters, kDimension); + constexpr size_t kDim = 64; + const std::vector data = separated_points(kClusters, kDim); QGKMeansParameters parameters = exact_search_parameters(); parameters.niter = 1; parameters.verbose = true; parameters.min_points_per_centroid = 39; testing::internal::CaptureStderr(); - QGKMeans clustering(kDimension, kClusters, parameters); + QGKMeans clustering(kDim, kClusters, parameters); clustering.train(kClusters, data.data()); const std::string warning = testing::internal::GetCapturedStderr(); @@ -388,8 +392,8 @@ TEST(QGKMeansTest, WarnsWhenThereAreTooFewPointsPerCentroid) { TEST(QGKMeansTest, UsesFaissEarlyStopThresholdByDefault) { constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - const std::vector data = separated_points(kClusters, kDimension); + constexpr size_t kDim = 64; + const std::vector data = separated_points(kClusters, kDim); QGKMeansParameters parameters = exact_search_parameters(); EXPECT_DOUBLE_EQ(parameters.early_stop_threshold, 0.0); @@ -397,13 +401,13 @@ TEST(QGKMeansTest, UsesFaissEarlyStopThresholdByDefault) { TEST(QGKMeansTest, StopsAtFaissEarlyStopThreshold) { constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - const std::vector data = separated_points(kClusters * 2, kDimension); + constexpr size_t kDim = 64; + const std::vector data = separated_points(kClusters * 2, kDim); QGKMeansParameters parameters = exact_search_parameters(); parameters.early_stop_threshold = 1.0; const QGKMeans result = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kClusters * 2); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kClusters * 2); ASSERT_EQ(result.iteration_stats.size(), 2U); } @@ -411,8 +415,8 @@ TEST(QGKMeansTest, StopsAtFaissEarlyStopThreshold) { TEST(QGKMeansTest, ExactFinalAssignmentReturnsNearestCentroidForEveryPoint) { constexpr size_t kPoints = 100; constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - std::vector data(kPoints * kDimension); + constexpr size_t kDim = 64; + std::vector data(kPoints * kDim); std::mt19937 rng(123); std::uniform_real_distribution distribution(-10.0F, 10.0F); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); @@ -421,7 +425,7 @@ TEST(QGKMeansTest, ExactFinalAssignmentReturnsNearestCentroidForEveryPoint) { parameters.niter = 2; parameters.final_assignment = FinalAssignmentMode::Exact; const QGKMeans result = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kPoints); EXPECT_DOUBLE_EQ( result.final_obj, @@ -432,10 +436,9 @@ TEST(QGKMeansTest, ExactFinalAssignmentReturnsNearestCentroidForEveryPoint) { double expected_distance = std::numeric_limits::infinity(); for (size_t cluster = 0; cluster < kClusters; ++cluster) { double distance = 0.0; - for (size_t dim = 0; dim < kDimension; ++dim) { - const double difference = - static_cast(data[point * kDimension + dim]) - - result.centroids[cluster * kDimension + dim]; + for (size_t dim = 0; dim < kDim; ++dim) { + const double difference = static_cast(data[point * kDim + dim]) - + result.centroids[cluster * kDim + dim]; distance += difference * difference; } if (distance < expected_distance) { @@ -456,7 +459,7 @@ void expect_exact_assignment_matches_scalar( ) { const size_t points = data.size() / dimension; const size_t clusters = centroids.size() / dimension; - std::vector expected_labels(points); + std::vector expected_assignments(points); std::vector expected_distances(points); for (size_t point = 0; point < points; ++point) { double best = std::numeric_limits::infinity(); @@ -477,7 +480,7 @@ void expect_exact_assignment_matches_scalar( } if (distance < best) { best = distance; - expected_labels[point] = static_cast(cluster); + expected_assignments[point] = static_cast(cluster); } } expected_distances[point] = static_cast(best); @@ -489,13 +492,13 @@ void expect_exact_assignment_matches_scalar( const std::vector previous(points, static_cast(clusters - 1)); for (const int hint_mode : {0, 1, 2}) { SCOPED_TRACE(hint_mode); - std::vector labels = + std::vector assignments = hint_mode == 2 ? previous : std::vector(points, kPidMax); std::vector distances(points, std::numeric_limits::quiet_NaN()); - const PID* hints = hint_mode == 0 - ? nullptr - : (hint_mode == 1 ? previous.data() : labels.data()); - detail::exact_assign( + const PID* hints = + hint_mode == 0 ? nullptr + : (hint_mode == 1 ? previous.data() : assignments.data()); + clustering::detail::exact_assign( data.data(), centroids.data(), points, @@ -503,11 +506,11 @@ void expect_exact_assignment_matches_scalar( clusters, spherical, threads, - labels.data(), + assignments.data(), distances.data(), hints ); - EXPECT_EQ(labels, expected_labels); + EXPECT_EQ(assignments, expected_assignments); EXPECT_EQ(distances, expected_distances); } } @@ -532,30 +535,30 @@ TEST(QGKMeansTest, ExactAssignmentMatchesScalarAcrossMatrixTileTails) { } TEST(QGKMeansTest, ExactAssignmentMatchesScalarAtMaximumDimension) { - constexpr size_t kDimension = 65536; + constexpr size_t kDim = 65536; std::mt19937 rng(963); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector data(3 * kDimension); - std::vector centroids(3 * kDimension); + std::vector data(3 * kDim); + std::vector centroids(3 * kDim); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); for (const bool spherical : {false, true}) { SCOPED_TRACE(spherical); - expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDim, spherical); } } TEST(QGKMeansTest, ExactAssignmentMatchesScalarWithManyPointBlocks) { - constexpr size_t kDimension = 17; + constexpr size_t kDim = 17; std::mt19937 rng(385); std::uniform_real_distribution distribution(-2.0F, 2.0F); - std::vector data(4097 * kDimension); - std::vector centroids(257 * kDimension); + std::vector data(4097 * kDim); + std::vector centroids(257 * kDim); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); for (const bool spherical : {false, true}) { SCOPED_TRACE(spherical); - expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDim, spherical); } } @@ -583,25 +586,25 @@ TEST(QGKMeansTest, ExactAssignmentPreservesCallerOpenMPSettings) { #else omp_set_nested(1); #endif - constexpr size_t kDimension = 65; - const std::vector centroids(3 * kDimension, 1.0F); + constexpr size_t kDim = 65; + const std::vector centroids(3 * kDim, 1.0F); // Cover a single worker, a partial row block, and multiple workers. for (const size_t points : {3U, 65U, 513U}) { SCOPED_TRACE(points); - const std::vector data(points * kDimension, 0.0F); - std::vector labels(points); + const std::vector data(points * kDim, 0.0F); + std::vector assignments(points); std::vector distances(points); for (const uint32_t threads : {1U, 4U}) { SCOPED_TRACE(threads); - detail::exact_assign( + clustering::detail::exact_assign( data.data(), centroids.data(), points, - kDimension, + kDim, 3, false, threads, - labels.data(), + assignments.data(), distances.data() ); EXPECT_EQ(omp_get_max_threads(), 3); @@ -615,92 +618,91 @@ TEST(QGKMeansTest, ExactAssignmentPreservesCallerOpenMPSettings) { } TEST(QGKMeansTest, ExactAssignmentPreservesLowestIndexForTiesAcrossTiles) { - constexpr size_t kDimension = 65; - std::vector data(65 * kDimension, 0.0F); - std::vector centroids(257 * kDimension, 4.0F); + constexpr size_t kDim = 65; + std::vector data(65 * kDim, 0.0F); + std::vector centroids(257 * kDim, 4.0F); // Centroids 3 and 256 tie, while preceding centroids are worse for every row. - std::fill_n(centroids.data() + 3 * kDimension, kDimension, 0.0F); - std::fill_n(centroids.data() + 256 * kDimension, kDimension, 0.0F); - centroids[3 * kDimension] = -1.0F; - centroids[256 * kDimension] = 1.0F; - expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + std::fill_n(centroids.data() + 3 * kDim, kDim, 0.0F); + std::fill_n(centroids.data() + 256 * kDim, kDim, 0.0F); + centroids[3 * kDim] = -1.0F; + centroids[256 * kDim] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDim, false); std::fill(data.begin(), data.end(), 1.0F); std::fill(centroids.begin(), centroids.end(), 0.0F); - centroids[3 * kDimension] = 1.0F; - centroids[256 * kDimension] = 1.0F; - expect_exact_assignment_matches_scalar(data, centroids, kDimension, true); + centroids[3 * kDim] = 1.0F; + centroids[256 * kDim] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDim, true); } TEST(QGKMeansTest, ExactAssignmentRefinesLargeOffsetL2Distances) { - constexpr size_t kDimension = 129; + constexpr size_t kDim = 129; std::mt19937 rng(51); std::uniform_real_distribution distribution(-0.125F, 0.125F); - std::vector data(65 * kDimension); - std::vector centroids(257 * kDimension); + std::vector data(65 * kDim); + std::vector centroids(257 * kDim); auto shifted = [&] { return 100000.0F + distribution(rng); }; std::generate(data.begin(), data.end(), shifted); std::generate(centroids.begin(), centroids.end(), shifted); // Include a zero-distance winner at the tail of the centroid matrix. - std::copy_n(data.data(), kDimension, centroids.data() + 256 * kDimension); - expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + std::copy_n(data.data(), kDim, centroids.data() + 256 * kDim); + expect_exact_assignment_matches_scalar(data, centroids, kDim, false); } TEST(QGKMeansTest, ExactAssignmentResolvesFloatNearTiesAndTinyDistances) { - constexpr size_t kDimension = 65; - std::vector data(65 * kDimension, 0.0F); - std::vector centroids(257 * kDimension, 2.0F); - std::fill_n(centroids.data(), kDimension, 0.0F); - std::fill_n(centroids.data() + 256 * kDimension, kDimension, 0.0F); + constexpr size_t kDim = 65; + std::vector data(65 * kDim, 0.0F); + std::vector centroids(257 * kDim, 2.0F); + std::fill_n(centroids.data(), kDim, 0.0F); + std::fill_n(centroids.data() + 256 * kDim, kDim, 0.0F); // Both squared distances round to 1 in float, but centroid 256 is nearer. centroids[0] = 1.0F; centroids[1] = 0.0001F; - centroids[256 * kDimension] = 1.0F; - expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + centroids[256 * kDim] = 1.0F; + expect_exact_assignment_matches_scalar(data, centroids, kDim, false); for (float& value : centroids) { value *= 1e-22F; } - expect_exact_assignment_matches_scalar(data, centroids, kDimension, false); + expect_exact_assignment_matches_scalar(data, centroids, kDim, false); std::fill(centroids.begin(), centroids.end(), 0.0F); for (const bool spherical : {false, true}) { SCOPED_TRACE(spherical); - expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDim, spherical); } } TEST(QGKMeansTest, ExactAssignmentRefinesInnerProductCancellation) { - constexpr size_t kDimension = 129; - std::vector data(65 * kDimension, 100000.0F); - std::vector centroids(257 * kDimension); + constexpr size_t kDim = 129; + std::vector data(65 * kDim, 100000.0F); + std::vector centroids(257 * kDim); for (size_t point = 0; point < 65; ++point) { - data[point * kDimension + kDimension - 1] = 1.0F; + data[point * kDim + kDim - 1] = 1.0F; } for (size_t cluster = 0; cluster < 257; ++cluster) { - for (size_t dim = 0; dim + 1 < kDimension; ++dim) { - centroids[cluster * kDimension + dim] = dim % 2 == 0 ? 100000.0F : -100000.0F; + for (size_t dim = 0; dim + 1 < kDim; ++dim) { + centroids[cluster * kDim + dim] = dim % 2 == 0 ? 100000.0F : -100000.0F; } - centroids[cluster * kDimension + kDimension - 1] = - 1.0F + static_cast(cluster) / 4096.0F; + centroids[cluster * kDim + kDim - 1] = 1.0F + static_cast(cluster) / 4096.0F; } - expect_exact_assignment_matches_scalar(data, centroids, kDimension, true); + expect_exact_assignment_matches_scalar(data, centroids, kDim, true); } TEST(QGKMeansTest, SphericalModeNormalizesCentroidsAndUsesInnerProduct) { constexpr size_t kPoints = 100; constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - std::vector data(kPoints * kDimension); + constexpr size_t kDim = 64; + std::vector data(kPoints * kDim); std::mt19937 rng(456); std::uniform_real_distribution distribution(-1.0F, 1.0F); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); for (size_t point = 0; point < kPoints; ++point) { double squared_norm = 0.0; - for (size_t dim = 0; dim < kDimension; ++dim) { - squared_norm += static_cast(data[point * kDimension + dim]) * - data[point * kDimension + dim]; + for (size_t dim = 0; dim < kDim; ++dim) { + squared_norm += + static_cast(data[point * kDim + dim]) * data[point * kDim + dim]; } const float scale = 1.0F / std::sqrt(static_cast(squared_norm)); - for (size_t dim = 0; dim < kDimension; ++dim) { - data[point * kDimension + dim] *= scale; + for (size_t dim = 0; dim < kDim; ++dim) { + data[point * kDim + dim] *= scale; } } @@ -709,14 +711,13 @@ TEST(QGKMeansTest, SphericalModeNormalizesCentroidsAndUsesInnerProduct) { parameters.spherical = true; parameters.final_assignment = FinalAssignmentMode::Exact; const QGKMeans result = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kPoints); for (size_t cluster = 0; cluster < kClusters; ++cluster) { double squared_norm = 0.0; - for (size_t dim = 0; dim < kDimension; ++dim) { - squared_norm += - static_cast(result.centroids[cluster * kDimension + dim]) * - result.centroids[cluster * kDimension + dim]; + for (size_t dim = 0; dim < kDim; ++dim) { + squared_norm += static_cast(result.centroids[cluster * kDim + dim]) * + result.centroids[cluster * kDim + dim]; } EXPECT_NEAR(squared_norm, 1.0, 1e-5); } @@ -730,9 +731,9 @@ TEST(QGKMeansTest, SphericalModeNormalizesCentroidsAndUsesInnerProduct) { double expected_distance = std::numeric_limits::infinity(); for (size_t cluster = 0; cluster < kClusters; ++cluster) { double distance = 1.0; - for (size_t dim = 0; dim < kDimension; ++dim) { - distance -= static_cast(data[point * kDimension + dim]) * - result.centroids[cluster * kDimension + dim]; + for (size_t dim = 0; dim < kDim; ++dim) { + distance -= static_cast(data[point * kDim + dim]) * + result.centroids[cluster * kDim + dim]; } if (distance < expected_distance) { expected = static_cast(cluster); @@ -745,20 +746,20 @@ TEST(QGKMeansTest, SphericalModeNormalizesCentroidsAndUsesInnerProduct) { } TEST(QGKMeansTest, ObjectiveTerminationMatchesFaissAbsoluteRelativeChange) { - EXPECT_TRUE(detail::should_terminate_by_objective(100.0, 101.0, 0.01)); - EXPECT_TRUE(detail::should_terminate_by_objective(100.0, 99.0, 0.01)); - EXPECT_FALSE(detail::should_terminate_by_objective(100.0, 98.0, 0.01)); - EXPECT_TRUE(detail::should_terminate_by_objective(0.0, 0.0, 0.0)); - EXPECT_FALSE(detail::should_terminate_by_objective(0.0, 1.0, 1.0)); + EXPECT_TRUE(clustering::detail::should_terminate_by_objective(100.0, 101.0, 0.01)); + EXPECT_TRUE(clustering::detail::should_terminate_by_objective(100.0, 99.0, 0.01)); + EXPECT_FALSE(clustering::detail::should_terminate_by_objective(100.0, 98.0, 0.01)); + EXPECT_TRUE(clustering::detail::should_terminate_by_objective(0.0, 0.0, 0.0)); + EXPECT_FALSE(clustering::detail::should_terminate_by_objective(0.0, 1.0, 1.0)); } TEST(QGKMeansTest, TrainsAndReturnsAssignmentsForFinalCentroids) { constexpr size_t kClusters = 33; - constexpr size_t kDimension = 64; - std::vector data(kClusters * kDimension); + constexpr size_t kDim = 64; + std::vector data(kClusters * kDim); for (size_t point = 0; point < kClusters; ++point) { - for (size_t dim = 0; dim < kDimension; ++dim) { - data[point * kDimension + dim] = + for (size_t dim = 0; dim < kDim; ++dim) { + data[point * kDim + dim] = static_cast(point * 10) + static_cast(dim) / 64.0F; } } @@ -768,9 +769,9 @@ TEST(QGKMeansTest, TrainsAndReturnsAssignmentsForFinalCentroids) { parameters.num_threads = 2; const QGKMeans result = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kClusters); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kClusters); - ASSERT_EQ(result.centroids.size(), kClusters * kDimension); + ASSERT_EQ(result.centroids.size(), kClusters * kDim); ASSERT_EQ(result.assignments.size(), kClusters); ASSERT_EQ(result.distances.size(), kClusters); ASSERT_EQ(result.iteration_stats.size(), 1U); @@ -779,17 +780,17 @@ TEST(QGKMeansTest, TrainsAndReturnsAssignmentsForFinalCentroids) { EXPECT_TRUE(std::isfinite(result.distances[point])); EXPECT_FLOAT_EQ(result.distances[point], 0.0F); const size_t cluster = result.assignments[point]; - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { EXPECT_FLOAT_EQ( - result.centroids[cluster * kDimension + dim], data[point * kDimension + dim] + result.centroids[cluster * kDim + dim], data[point * kDim + dim] ); } } } TEST(QGKMeansTest, RetrainsOnOwnedCentroidsLikeAnIndependentCopy) { - constexpr size_t kDimension = 65, kInitialClusters = 35, kClusters = 33; - const auto data = separated_points(kInitialClusters, kDimension); + constexpr size_t kDim = 65, kInitialClusters = 35, kClusters = 33; + const auto data = separated_points(kInitialClusters, kDim); for (const uint32_t bits : {0U, 4U, 8U}) { for (const auto mode : {FinalAssignmentMode::SymphonyQG, FinalAssignmentMode::Exact}) { @@ -802,18 +803,18 @@ TEST(QGKMeansTest, RetrainsOnOwnedCentroidsLikeAnIndependentCopy) { parameters.niter = 2; parameters.quantization_bits = bits; parameters.final_assignment = mode; - QGKMeans clustering(kDimension, kInitialClusters, parameters); + QGKMeans clustering(kDim, kInitialClusters, parameters); clustering.train(kInitialClusters, data.data()); const auto original = clustering.centroids; - const float* input = original.data() + offset * kDimension; + const float* input = original.data() + offset * kDim; const size_t count = kInitialClusters - offset; - QGKMeans reference(kDimension, kClusters, parameters); + QGKMeans reference(kDim, kClusters, parameters); reference.train(count, input); // Shrink the output while training on its former contents, including // a subrange that does not start at centroids.data(). clustering.k = kClusters; - clustering.train(count, clustering.centroids.data() + offset * kDimension); + clustering.train(count, clustering.centroids.data() + offset * kDim); EXPECT_EQ(clustering.centroids, reference.centroids); EXPECT_EQ(clustering.assignments, reference.assignments); @@ -823,11 +824,11 @@ TEST(QGKMeansTest, RetrainsOnOwnedCentroidsLikeAnIndependentCopy) { for (size_t point = 0; point < count; ++point) { ASSERT_LT(clustering.assignments[point], kClusters); double distance = 0; - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { const double difference = - static_cast(input[point * kDimension + dim]) - - clustering.centroids - [clustering.assignments[point] * kDimension + dim]; + static_cast(input[point * kDim + dim]) - + clustering + .centroids[clustering.assignments[point] * kDim + dim]; distance += difference * difference; } EXPECT_NEAR( @@ -846,38 +847,37 @@ TEST(QGKMeansTest, RetrainsOnOwnedCentroidsLikeAnIndependentCopy) { } TEST(QGKMeansTest, EmptyClusterRecoveryDoesNotPerturbConstantData) { - constexpr size_t kPoints = 100, kDimension = 64; + constexpr size_t kPoints = 100, kDim = 64; auto parameters = exact_search_parameters(); parameters.niter = 10; for (bool spherical : {false, true}) { SCOPED_TRACE(spherical); parameters.spherical = spherical; - std::vector data(kPoints * kDimension, 0.0F); + std::vector data(kPoints * kDim, 0.0F); for (size_t point = 0; point < kPoints; ++point) { - data[point * kDimension] = 1.0F; + data[point * kDim] = 1.0F; } - const auto clustering = - train_qgkmeans(kDimension, 33, parameters, data.data(), kPoints); + const auto clustering = train_qgkmeans(kDim, 33, parameters, data.data(), kPoints); EXPECT_LT(clustering.iteration_stats.size(), parameters.niter); EXPECT_EQ(clustering.final_obj, 0.0); for (size_t cluster = 0; cluster < 33; ++cluster) { - for (size_t dim = 0; dim < kDimension; ++dim) { - EXPECT_EQ(clustering.centroids[cluster * kDimension + dim], data[dim]); + for (size_t dim = 0; dim < kDim; ++dim) { + EXPECT_EQ(clustering.centroids[cluster * kDim + dim], data[dim]); } } } } TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { - constexpr size_t kDimension = 64, kClusters = 33, kPoints = kClusters + 8; + constexpr size_t kDim = 64, kClusters = 33, kPoints = kClusters + 8; const auto normalize = [=](float* row) { double squared_norm = 0; - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { squared_norm += static_cast(row[dim]) * row[dim]; } if (squared_norm != 0) { const double scale = 1.0 / std::sqrt(squared_norm); - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { row[dim] = static_cast(row[dim] * scale); } } @@ -892,13 +892,13 @@ TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { SCOPED_TRACE(::testing::Message() << spherical << "/" << empty_clusters); // Force duplicate initial seeds while all other initial centroids are // distinct basis vectors. Only the duplicate seeds have extra points. - std::vector data(kPoints * kDimension); + std::vector data(kPoints * kDim); for (size_t cluster = 0; cluster < kClusters; ++cluster) { const size_t axis = cluster <= empty_clusters ? 0 : cluster; - data[indices[cluster] * kDimension + axis] = 1.0F; + data[indices[cluster] * kDim + axis] = 1.0F; } for (size_t extra = 0; extra < 8; ++extra) { - float* row = data.data() + indices[kClusters + extra] * kDimension; + float* row = data.data() + indices[kClusters + extra] * kDim; row[0] = 1.0F; row[kClusters] = static_cast(4 - extra / 2) / 8.0F * (extra % 2 == 0 ? 1.0F : -1.0F); @@ -906,63 +906,54 @@ TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { normalize(row); } } - std::vector initial(kClusters * kDimension); + std::vector initial(kClusters * kDim); for (size_t cluster = 0; cluster < kClusters; ++cluster) { std::copy_n( - data.data() + indices[cluster] * kDimension, - kDimension, - initial.data() + cluster * kDimension + data.data() + indices[cluster] * kDim, + kDim, + initial.data() + cluster * kDim ); } QGAssigner reference( - kDimension, - kClusters, - 32, - 33, - 33, - 2, - 1, - 42, - spherical ? METRIC_IP : METRIC_L2 + kDim, kClusters, 32, 33, 33, 2, 1, 42, spherical ? METRIC_IP : METRIC_L2 ); - std::vector labels(kPoints); + std::vector assignments(kPoints); std::vector distances(kPoints); reference.assign( - initial.data(), data.data(), kPoints, labels.data(), distances.data() + initial.data(), data.data(), kPoints, assignments.data(), distances.data() ); // Derive occupied means from the public initial assignment. Empty // centroids receive selected rows without changing these donor means. std::vector counts(kClusters); - std::vector sums(kClusters * kDimension); + std::vector sums(kClusters * kDim); for (size_t point = 0; point < kPoints; ++point) { - ++counts[labels[point]]; - for (size_t dim = 0; dim < kDimension; ++dim) { - sums[labels[point] * kDimension + dim] += - data[point * kDimension + dim]; + ++counts[assignments[point]]; + for (size_t dim = 0; dim < kDim; ++dim) { + sums[assignments[point] * kDim + dim] += data[point * kDim + dim]; } } ASSERT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), empty_clusters); - std::vector means(kClusters * kDimension); + std::vector means(kClusters * kDim); for (size_t cluster = 0; cluster < kClusters; ++cluster) { if (counts[cluster] != 0) { - for (size_t dim = 0; dim < kDimension; ++dim) { - means[cluster * kDimension + dim] = static_cast( - sums[cluster * kDimension + dim] / + for (size_t dim = 0; dim < kDim; ++dim) { + means[cluster * kDim + dim] = static_cast( + sums[cluster * kDim + dim] / static_cast(counts[cluster]) ); } if (spherical) { - normalize(means.data() + cluster * kDimension); + normalize(means.data() + cluster * kDim); } } } const auto distance = spherical ? dot_product_dis : euclidean_sqr; for (size_t point = 0; point < kPoints; ++point) { distances[point] = distance( - data.data() + point * kDimension, - means.data() + labels[point] * kDimension, - kDimension + data.data() + point * kDim, + means.data() + assignments[point] * kDim, + kDim ); } std::vector candidates(kPoints); @@ -979,20 +970,18 @@ TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { if (counts[cluster] != 0) { continue; } - while (counts[labels[candidates[next_candidate]]] <= 1) { + while (counts[assignments[candidates[next_candidate]]] <= 1) { ++next_candidate; ASSERT_LT(next_candidate, candidates.size()); } const size_t point = candidates[next_candidate++]; std::copy_n( - data.data() + point * kDimension, - kDimension, - means.data() + cluster * kDimension + data.data() + point * kDim, kDim, means.data() + cluster * kDim ); if (spherical) { - normalize(means.data() + cluster * kDimension); + normalize(means.data() + cluster * kDim); } - --counts[labels[point]]; + --counts[assignments[point]]; counts[cluster] = 1; } ASSERT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), 0); @@ -1004,27 +993,27 @@ TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { parameters.num_threads = threads; parameters.final_assignment = FinalAssignmentMode::Exact; const auto clustering = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kPoints); ASSERT_EQ(clustering.iteration_stats.front().nsplit, empty_clusters); // Parallel graph construction may choose a different duplicate - // seed label, but occupied means and selected seeds must match. + // seed assignment, but occupied means and selected seeds must match. std::vector> expected_rows, actual_rows; for (size_t cluster = 0; cluster < kClusters; ++cluster) { - const size_t offset = cluster * kDimension; + const size_t offset = cluster * kDim; expected_rows.emplace_back( - means.data() + offset, means.data() + offset + kDimension + means.data() + offset, means.data() + offset + kDim ); actual_rows.emplace_back( clustering.centroids.data() + offset, - clustering.centroids.data() + offset + kDimension + clustering.centroids.data() + offset + kDim ); } std::sort(expected_rows.begin(), expected_rows.end()); std::sort(actual_rows.begin(), actual_rows.end()); EXPECT_EQ(actual_rows, expected_rows); std::fill(counts.begin(), counts.end(), size_t{0}); - for (PID label : clustering.assignments) { - ++counts[label]; + for (PID cluster_id : clustering.assignments) { + ++counts[cluster_id]; } EXPECT_EQ(std::count(counts.begin(), counts.end(), size_t{0}), 0); } @@ -1033,11 +1022,11 @@ TEST(QGKMeansTest, EmptyClusterRecoveryPreservesMeansAndFarthestSeedSelection) { } TEST(QGKMeansTest, RejectsNonfiniteTrainingDataWithoutReplacingResults) { - constexpr size_t kDimension = 64; - const auto data = separated_points(66, kDimension); + constexpr size_t kDim = 64; + const auto data = separated_points(66, kDim); auto parameters = exact_search_parameters(); parameters.niter = 1; - QGKMeans clustering(kDimension, 33, parameters); + QGKMeans clustering(kDim, 33, parameters); clustering.train(66, data.data()); const auto expected_centroids = clustering.centroids; const auto expected_assignments = clustering.assignments; @@ -1059,10 +1048,10 @@ TEST(QGKMeansTest, RejectsNonfiniteTrainingDataWithoutReplacingResults) { } TEST(QGKMeansTest, AccumulatesShiftedClusterMeansAccurately) { - constexpr size_t kDimension = 64, kClusters = 33, kPoints = 10000; + constexpr size_t kDim = 64, kClusters = 33, kPoints = 10000; std::mt19937 rng(123); std::uniform_real_distribution distribution(-1, 1); - std::vector data(kPoints * kDimension); + std::vector data(kPoints * kDim); for (auto& value : data) { value = 10000.0F + distribution(rng); } @@ -1072,32 +1061,32 @@ TEST(QGKMeansTest, AccumulatesShiftedClusterMeansAccurately) { std::iota(indices.begin(), indices.end(), size_t{0}); std::mt19937 initial_rng(42); std::shuffle(indices.begin(), indices.end(), initial_rng); - std::vector initial_centroids(kClusters * kDimension); + std::vector initial_centroids(kClusters * kDim); for (size_t cluster = 0; cluster < kClusters; ++cluster) { std::copy_n( - data.data() + indices[cluster] * kDimension, - kDimension, - initial_centroids.data() + cluster * kDimension + data.data() + indices[cluster] * kDim, + kDim, + initial_centroids.data() + cluster * kDim ); } - QGAssigner assigner(kDimension, kClusters, 32, 33, 33, 2, 1, 42); - std::vector labels(kPoints); + QGAssigner assigner(kDim, kClusters, 32, 33, 33, 2, 1, 42); + std::vector assignments(kPoints); std::vector distances(kPoints); assigner.assign( - initial_centroids.data(), data.data(), kPoints, labels.data(), distances.data() + initial_centroids.data(), data.data(), kPoints, assignments.data(), distances.data() ); - std::vector sums(kClusters * kDimension); + std::vector sums(kClusters * kDim); std::vector counts(kClusters); for (size_t point = 0; point < kPoints; ++point) { - ++counts[labels[point]]; - for (size_t dim = 0; dim < kDimension; ++dim) { - sums[labels[point] * kDimension + dim] += data[point * kDimension + dim]; + ++counts[assignments[point]]; + for (size_t dim = 0; dim < kDim; ++dim) { + sums[assignments[point] * kDim + dim] += data[point * kDim + dim]; } } auto parameters = exact_search_parameters(); parameters.niter = 1; const auto clustering = - train_qgkmeans(kDimension, kClusters, parameters, data.data(), kPoints); + train_qgkmeans(kDim, kClusters, parameters, data.data(), kPoints); EXPECT_EQ( clustering.iteration_stats.front().nsplit, std::count(counts.begin(), counts.end(), size_t{0}) @@ -1106,45 +1095,45 @@ TEST(QGKMeansTest, AccumulatesShiftedClusterMeansAccurately) { if (counts[cluster] == 0) { continue; } - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { EXPECT_FLOAT_EQ( - clustering.centroids[cluster * kDimension + dim], - static_cast(sums[cluster * kDimension + dim] / counts[cluster]) + clustering.centroids[cluster * kDim + dim], + static_cast(sums[cluster * kDim + dim] / counts[cluster]) ); } } } TEST(QGKMeansTest, FinalAssignmentsRetainACloserPreviousCluster) { - constexpr size_t kDimension = 65, kClusters = 256, kPoints = 2048; + constexpr size_t kDim = 65, kClusters = 256, kPoints = 2048; const auto normalize = [=](float* row) { double norm = 0; - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { norm += static_cast(row[dim]) * row[dim]; } const double scale = 1.0 / std::sqrt(norm); - for (size_t dim = 0; dim < kDimension; ++dim) { + for (size_t dim = 0; dim < kDim; ++dim) { row[dim] = static_cast(row[dim] * scale); } }; for (bool spherical : {false, true}) { std::mt19937 rng(123); std::normal_distribution normal; - std::vector data(kPoints * kDimension); + std::vector data(kPoints * kDim); std::generate(data.begin(), data.end(), [&] { return normal(rng); }); if (spherical) { for (size_t point = 0; point < kPoints; ++point) { - normalize(data.data() + point * kDimension); + normalize(data.data() + point * kDim); } } std::vector indices(kPoints); std::iota(indices.begin(), indices.end(), size_t{0}); std::mt19937 init_rng(42); std::shuffle(indices.begin(), indices.end(), init_rng); - std::vector initial(kClusters * kDimension); + std::vector initial(kClusters * kDim); for (size_t cluster = 0; cluster < kClusters; ++cluster) { - float* row = initial.data() + cluster * kDimension; - std::copy_n(data.data() + indices[cluster] * kDimension, kDimension, row); + float* row = initial.data() + cluster * kDim; + std::copy_n(data.data() + indices[cluster] * kDim, kDim, row); if (spherical) { normalize(row); } @@ -1153,7 +1142,7 @@ TEST(QGKMeansTest, FinalAssignmentsRetainACloserPreviousCluster) { for (uint32_t bits : {0U, 4U, 8U}) { SCOPED_TRACE(::testing::Message() << spherical << "/" << bits); QGAssigner reference( - kDimension, + kDim, kClusters, 32, 240, @@ -1175,14 +1164,14 @@ TEST(QGKMeansTest, FinalAssignmentsRetainACloserPreviousCluster) { parameters.min_points_per_centroid = 1; parameters.spherical = spherical; parameters.quantization_bits = bits; - QGKMeans clustering(kDimension, kClusters, parameters); + QGKMeans clustering(kDim, kClusters, parameters); clustering.train(kPoints, data.data()); const auto distance = spherical ? dot_product_dis : euclidean_sqr; for (size_t point = 0; point < kPoints; ++point) { const float previous_distance = distance( - data.data() + point * kDimension, - clustering.centroids.data() + previous[point] * kDimension, - kDimension + data.data() + point * kDim, + clustering.centroids.data() + previous[point] * kDim, + kDim ); EXPECT_LE(clustering.distances[point], previous_distance); } @@ -1302,16 +1291,16 @@ namespace rabitqlib::qgkmeans { TEST(QGAssignerTest, PropagatesSearchFailureFromWorkers) { const std::vector centroids(33 * 64, 0.0F); const std::vector data(128 * 64, 1e19F); - std::vector labels(128, kPidMax); + std::vector assignments(128, kPidMax); std::vector distances(128, -1.0F); QGAssigner assigner(64, 33, 32, 33, 33, 1, 2); EXPECT_THROW( assigner.assign( - centroids.data(), data.data(), 128, labels.data(), distances.data() + centroids.data(), data.data(), 128, assignments.data(), distances.data() ), std::runtime_error ); - EXPECT_EQ(labels.front(), kPidMax); + EXPECT_EQ(assignments.front(), kPidMax); EXPECT_EQ(distances.front(), -1.0F); } @@ -1331,11 +1320,11 @@ TEST(QGKMeansTest, RejectsFiniteOutOfRangeInputBeforeConstruction) { namespace rabitqlib::qgkmeans { TEST(QGKMeansTest, SupportsLargerPaddedDimensions) { - constexpr size_t k = 33; + constexpr size_t kClusters = 33; for (const size_t d : {4097U, 65535U, 65536U}) { SCOPED_TRACE(d); - std::vector data(k * d, 0.0F); - for (size_t i = 0; i < k; ++i) { + std::vector data(kClusters * d, 0.0F); + for (size_t i = 0; i < kClusters; ++i) { data[i * d + i] = 1.0F; data[i * d + d - 1] = static_cast(i) / 64.0F; } @@ -1349,10 +1338,10 @@ TEST(QGKMeansTest, SupportsLargerPaddedDimensions) { parameters.niter = 1; parameters.num_threads = 2; parameters.quantization_bits = bits; - QGKMeans model(d, k, parameters); - model.train(k, data.data()); + QGKMeans model(d, kClusters, parameters); + model.train(kClusters, data.data()); EXPECT_DOUBLE_EQ(model.final_obj, 0.0); - EXPECT_EQ(model.assignments.size(), k); + EXPECT_EQ(model.assignments.size(), kClusters); } } EXPECT_THROW(QGAssigner(65537, 33, 32, 33, 33, 1, 1), std::invalid_argument); @@ -1365,8 +1354,9 @@ void normalize_reference_centroids(std::vector& rows, size_t dim) { for (size_t j = 0; j < dim; ++j) { norm += static_cast(rows[offset + j]) * rows[offset + j]; } - if (norm == 0) + if (norm == 0) { continue; + } const double scale = 1.0 / std::sqrt(norm); for (size_t j = 0; j < dim; ++j) { rows[offset + j] = static_cast(rows[offset + j] * scale); @@ -1375,19 +1365,17 @@ void normalize_reference_centroids(std::vector& rows, size_t dim) { } void check_uncached_reference(size_t clusters, uint32_t ef_search, size_t intrinsic_dim) { - constexpr size_t kDimension = 1024; + constexpr size_t kDim = 1024; constexpr size_t kPoints = 257; constexpr size_t kIterations = 3; std::mt19937 data_rng(823); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector x(kPoints * kDimension); + std::vector x(kPoints * kDim); std::generate(x.begin(), x.end(), [&] { return distribution(data_rng); }); // Low intrinsic dimension makes points change clusters during Lloyd updates. // Padding keeps the scalar-sum cache enabled without concentrating distances. for (size_t point = 0; point < kPoints; ++point) { - std::fill_n( - x.data() + point * kDimension + intrinsic_dim, kDimension - intrinsic_dim, 0.0F - ); + std::fill_n(x.data() + point * kDim + intrinsic_dim, kDim - intrinsic_dim, 0.0F); } for (const uint32_t bits : {0U, 4U, 8U}) { for (const bool spherical : {false, true}) { @@ -1405,7 +1393,7 @@ void check_uncached_reference(size_t clusters, uint32_t ef_search, size_t intrin parameters.spherical = spherical; parameters.ef_build = 128; parameters.ef_search = ef_search; - QGKMeans model(kDimension, clusters, parameters); + QGKMeans model(kDim, clusters, parameters); model.train(kPoints, x.data()); ASSERT_EQ(model.iteration_stats.size(), kIterations); @@ -1413,18 +1401,19 @@ void check_uncached_reference(size_t clusters, uint32_t ef_search, size_t intrin std::iota(initial_ids.begin(), initial_ids.end(), size_t{0}); std::mt19937 initialization_rng(parameters.seed); std::shuffle(initial_ids.begin(), initial_ids.end(), initialization_rng); - std::vector centroids(clusters * kDimension); + std::vector centroids(clusters * kDim); for (size_t cluster = 0; cluster < clusters; ++cluster) { std::copy_n( - x.data() + initial_ids[cluster] * kDimension, - kDimension, - centroids.data() + cluster * kDimension + x.data() + initial_ids[cluster] * kDim, + kDim, + centroids.data() + cluster * kDim ); } - if (spherical) - normalize_reference_centroids(centroids, kDimension); + if (spherical) { + normalize_reference_centroids(centroids, kDim); + } QGAssigner uncached( - kDimension, + kDim, clusters, 32, parameters.ef_build, @@ -1435,40 +1424,47 @@ void check_uncached_reference(size_t clusters, uint32_t ef_search, size_t intrin spherical ? METRIC_IP : METRIC_L2, bits ); - std::vector labels(kPoints); + std::vector assignments(kPoints); std::vector distances(kPoints); for (size_t iteration = 0; iteration < kIterations; ++iteration) { SCOPED_TRACE(iteration); // This public call intentionally never receives a prepared sum. uncached.assign( - centroids.data(), x.data(), kPoints, labels.data(), distances.data() + centroids.data(), + x.data(), + kPoints, + assignments.data(), + distances.data() ); const double objective = std::accumulate(distances.begin(), distances.end(), 0.0); EXPECT_EQ(model.iteration_stats[iteration].obj, objective); if (iteration == 0) { - EXPECT_FALSE(std::is_sorted(labels.begin(), labels.end())); + EXPECT_FALSE( + std::is_sorted(assignments.begin(), assignments.end()) + ); } - std::vector sums(clusters * kDimension, 0.0); + std::vector sums(clusters * kDim, 0.0); std::vector counts(clusters, 0); for (size_t point = 0; point < kPoints; ++point) { - ++counts[labels[point]]; - for (size_t j = 0; j < kDimension; ++j) { - sums[labels[point] * kDimension + j] += - static_cast(x[point * kDimension + j]); + ++counts[assignments[point]]; + for (size_t j = 0; j < kDim; ++j) { + sums[assignments[point] * kDim + j] += + static_cast(x[point * kDim + j]); } } for (size_t cluster = 0; cluster < clusters; ++cluster) { // Keep this an independent Lloyd reference, without empty repair. ASSERT_GT(counts[cluster], 0U); const double scale = 1.0 / static_cast(counts[cluster]); - for (size_t j = 0; j < kDimension; ++j) { - centroids[cluster * kDimension + j] = - static_cast(sums[cluster * kDimension + j] * scale); + for (size_t j = 0; j < kDim; ++j) { + centroids[cluster * kDim + j] = + static_cast(sums[cluster * kDim + j] * scale); } } - if (spherical) - normalize_reference_centroids(centroids, kDimension); + if (spherical) { + normalize_reference_centroids(centroids, kDim); + } } // Compare state before the final assignment, whose safeguard is separate. EXPECT_EQ(model.centroids, centroids); diff --git a/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp index 958bbda0..3054e9df 100644 --- a/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/rabitqkmeans_test.cpp @@ -30,8 +30,9 @@ void normalize_reference_centroids(std::vector& rows, size_t dim) { for (size_t j = 0; j < dim; ++j) { norm += static_cast(rows[offset + j]) * rows[offset + j]; } - if (norm == 0) + if (norm == 0) { continue; + } const double scale = 1.0 / std::sqrt(norm); for (size_t j = 0; j < dim; ++j) { rows[offset + j] = static_cast(rows[offset + j] * scale); @@ -39,7 +40,7 @@ void normalize_reference_centroids(std::vector& rows, size_t dim) { } } -std::vector scalar_rabitq_labels( +std::vector scalar_rabitq_assignments( const std::vector& centroids, const std::vector& data, size_t dim, @@ -79,7 +80,7 @@ std::vector scalar_rabitq_labels( } const auto distance = metric == METRIC_IP ? dot_product_dis : euclidean_sqr; - std::vector labels(points); + std::vector assignments(points); std::vector best(points, std::numeric_limits::infinity()), query(padded); for (size_t cluster = 0; cluster < clusters; ++cluster) { const float* centroid = centroids.data() + cluster * dim; @@ -109,11 +110,11 @@ std::vector scalar_rabitq_labels( batch.f_add()[index] + prepared.g_add() + batch.f_rescale()[index] * ip; if (estimate < best[point]) { best[point] = estimate; - labels[point] = static_cast(cluster); + assignments[point] = static_cast(cluster); } } } - return labels; + return assignments; } double scalar_centroid_distance(const float* x, const float* c, size_t dim, bool ip) { @@ -159,7 +160,7 @@ TEST(RaBitQKMeansTest, ValidatesCommonParametersBeforeAllocation) { ); } -TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroids) { +TEST(RaBitQAssignmentTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroids) { constexpr uint32_t kSeed = 53; std::mt19937 rng(157); std::uniform_real_distribution distribution(-1.0F, 1.0F); @@ -188,23 +189,24 @@ TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroi RaBitQKMeansParameters parameters; parameters.seed = kSeed; parameters.spherical = metric == METRIC_IP; - detail::RaBitQAssigner assigner( + detail::RaBitQAssignment assigner( dim, clusters, data.data(), points, threads, parameters ); for (size_t pass = 0; pass < 2; ++pass) { for (float& value : centers) { value = value * 0.75F + 0.125F; } - const auto expected = - scalar_rabitq_labels(centers, data, dim, kSeed, metric); - std::vector labels(points, kPidMax); + const auto expected = scalar_rabitq_assignments( + centers, data, dim, kSeed, metric + ); + std::vector assignments(points, kPidMax); std::vector distances(points); assigner.assign( - centers.data(), labels.data(), distances.data() + centers.data(), assignments.data(), distances.data() ); - EXPECT_EQ(labels, expected); + EXPECT_EQ(assignments, expected); for (size_t point = 0; point < points; ++point) { - ASSERT_LT(labels[point], clusters); + ASSERT_LT(assignments[point], clusters); const float* vector = data.data() + point * dim; const auto exact_distance = metric == METRIC_IP ? dot_product_dis @@ -212,7 +214,9 @@ TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroi EXPECT_EQ( distances[point], exact_distance( - vector, centers.data() + labels[point] * dim, dim + vector, + centers.data() + assignments[point] * dim, + dim ) ); PID nearest = 0; @@ -229,7 +233,7 @@ TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroi nearest = static_cast(cluster); } } - if (clusters > 32 && labels[point] != nearest) { + if (clusters > 32 && assignments[point] != nearest) { ++approximate_choices_above_batch_size; } } @@ -243,7 +247,7 @@ TEST(RaBitQAssignerTest, MatchesScalarPointCodesAcrossBatchTailsAndMovingCentroi EXPECT_GT(approximate_choices_above_batch_size, 0U); } -TEST(RaBitQAssignerTest, ZeroAndDuplicateResidualsChooseFirstLogicalCentroid) { +TEST(RaBitQAssignmentTest, ZeroAndDuplicateResidualsChooseFirstLogicalCentroid) { constexpr size_t kDim = 65, kPoints = 65; for (const size_t clusters : {1U, 2U, 16U, 31U, 32U, 33U, 65U}) { for (const MetricType metric : {METRIC_L2, METRIC_IP}) { @@ -252,13 +256,13 @@ TEST(RaBitQAssignerTest, ZeroAndDuplicateResidualsChooseFirstLogicalCentroid) { data(kPoints * kDim, value); RaBitQKMeansParameters parameters; parameters.spherical = metric == METRIC_IP; - detail::RaBitQAssigner assigner( + detail::RaBitQAssignment assigner( kDim, clusters, data.data(), kPoints, 2, parameters ); - std::vector labels(kPoints, kPidMax); + std::vector assignments(kPoints, kPidMax); std::vector distances(kPoints); - assigner.assign(centers.data(), labels.data(), distances.data()); - EXPECT_EQ(labels, (std::vector(kPoints, 0))); + assigner.assign(centers.data(), assignments.data(), distances.data()); + EXPECT_EQ(assignments, (std::vector(kPoints, 0))); const double expected = scalar_centroid_distance( data.data(), centers.data(), kDim, metric == METRIC_IP ); @@ -270,7 +274,7 @@ TEST(RaBitQAssignerTest, ZeroAndDuplicateResidualsChooseFirstLogicalCentroid) { } } -TEST(RaBitQAssignerTest, KeepsCentroidDependentCorrectionForZeroResidualPointCodes) { +TEST(RaBitQAssignmentTest, KeepsCentroidDependentCorrectionForZeroResidualPointCodes) { constexpr size_t kDim = 65, kPoints = 33; for (const MetricType metric : {METRIC_L2, METRIC_IP}) { std::vector data(kPoints * kDim), centers(2 * kDim); @@ -285,11 +289,11 @@ TEST(RaBitQAssignerTest, KeepsCentroidDependentCorrectionForZeroResidualPointCod } RaBitQKMeansParameters parameters; parameters.spherical = metric == METRIC_IP; - detail::RaBitQAssigner assigner(kDim, 2, data.data(), kPoints, 2, parameters); - std::vector labels(kPoints, kPidMax); + detail::RaBitQAssignment assigner(kDim, 2, data.data(), kPoints, 2, parameters); + std::vector assignments(kPoints, kPidMax); std::vector distances(kPoints); - assigner.assign(centers.data(), labels.data(), distances.data()); - EXPECT_EQ(labels, (std::vector(kPoints, 1))); + assigner.assign(centers.data(), assignments.data(), distances.data()); + EXPECT_EQ(assignments, (std::vector(kPoints, 1))); EXPECT_EQ(distances, (std::vector(kPoints, 0.0F))); } } @@ -316,7 +320,8 @@ TEST(RaBitQKMeansTest, FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts) { normalize_reference_centroids(initial, kDim); } const auto metric = spherical ? METRIC_IP : METRIC_L2; - const auto expected = scalar_rabitq_labels(initial, data, kDim, 42, metric); + const auto expected = + scalar_rabitq_assignments(initial, data, kDim, 42, metric); const auto distance = spherical ? dot_product_dis : euclidean_sqr; double expected_objective = 0; for (size_t point = 0; point < kPoints; ++point) { @@ -371,7 +376,7 @@ TEST(RaBitQKMeansTest, FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts) { } } const auto approximate = - scalar_rabitq_labels(model.centroids, data, kDim, 42, metric); + scalar_rabitq_assignments(model.centroids, data, kDim, 42, metric); for (size_t point = 0; point < kPoints; ++point) { PID nearest = 0; double best = std::numeric_limits::infinity(); diff --git a/tests/unit/rabitqlib/fastscan/api_test.cpp b/tests/unit/rabitqlib/fastscan/api_test.cpp index 93317ffa..dcba9c26 100644 --- a/tests/unit/rabitqlib/fastscan/api_test.cpp +++ b/tests/unit/rabitqlib/fastscan/api_test.cpp @@ -34,7 +34,7 @@ void expect_invalid(Call call, const char* message) { } catch (const std::invalid_argument& error) { EXPECT_STREQ(error.what(), message); } } -TEST(FastScanApi, RejectsUnpreparedAndMismatchedQueriesBeforeWriting) { +TEST(FastScanApiTest, RejectsUnpreparedAndMismatchedQueriesBeforeWriting) { constexpr size_t kDim = 64; std::array values{}; values.fill(1); @@ -140,7 +140,7 @@ TEST(FastScanApi, RejectsUnpreparedAndMismatchedQueriesBeforeWriting) { EXPECT_NO_THROW(qg_batch_estdist(batch.data(), moved_query, kDim, distance.data())); } -TEST(FastScanApi, SplitQueryResetMatchesFreshPreparationAndScoring) { +TEST(FastScanApiTest, SplitQueryResetMatchesFreshPreparationAndScoring) { std::array, 2> queries{}; std::array scratch{}; std::array hacc_scratch{}; @@ -262,7 +262,7 @@ TEST(FastScanApi, SplitQueryResetMatchesFreshPreparationAndScoring) { } } -TEST(FastScanApi, SplitQueryInitializationValidatesParametersConsistently) { +TEST(FastScanApiTest, SplitQueryInitializationValidatesParametersConsistently) { std::array values{}; values.fill(0.5F); std::array scratch{}; @@ -313,7 +313,7 @@ TEST(FastScanApi, SplitQueryInitializationValidatesParametersConsistently) { } } -TEST(FastScanApi, RejectedSplitQueryResetPreservesPreparedState) { +TEST(FastScanApiTest, RejectedSplitQueryResetPreservesPreparedState) { std::array values{}, next{}; values.fill(0.5F); next.fill(0.25F); @@ -393,7 +393,7 @@ TEST(FastScanApi, RejectedSplitQueryResetPreservesPreparedState) { EXPECT_EQ(query.g_add(), 0); } -TEST(FastScanApi, SingleBatchEncodersRejectOverflowAndPreserveEmptyOutput) { +TEST(FastScanApiTest, SingleBatchEncodersRejectOverflowAndPreserveEmptyOutput) { constexpr size_t kDim = 64; std::array data{}; std::array centroid{}; @@ -442,7 +442,7 @@ TEST(FastScanApi, SingleBatchEncodersRejectOverflowAndPreserveEmptyOutput) { } } -TEST(FastScanApi, PublicSerialAndParallelEncodingMatchIncludingTails) { +TEST(FastScanApiTest, PublicSerialAndParallelEncodingMatchIncludingTails) { constexpr size_t kDim = 96; std::array data{}; std::array centroid{}; @@ -468,7 +468,7 @@ TEST(FastScanApi, PublicSerialAndParallelEncodingMatchIncludingTails) { } } -TEST(FastScanApi, DoubleEncodingQueryAndScoringWithNonzeroCentroid) { +TEST(FastScanApiTest, DoubleEncodingQueryAndScoringWithNonzeroCentroid) { for (size_t dim : {16U, 32U, 96U}) { const double dimension = static_cast(dim); const double value = 0.5 + 1e-10; @@ -522,7 +522,7 @@ TEST(FastScanApi, DoubleEncodingQueryAndScoringWithNonzeroCentroid) { } } -TEST(FastScanApi, DoubleLutMatchesIndependentSubsetSumQuantization) { +TEST(FastScanApiTest, DoubleLutMatchesIndependentSubsetSumQuantization) { constexpr size_t kDim = 96; std::array query{}; for (size_t i = 0; i < kDim; ++i) { diff --git a/tests/unit/rabitqlib/fastscan/data_layout_test.cpp b/tests/unit/rabitqlib/fastscan/data_layout_test.cpp index 61bdf20e..a33e6a3e 100644 --- a/tests/unit/rabitqlib/fastscan/data_layout_test.cpp +++ b/tests/unit/rabitqlib/fastscan/data_layout_test.cpp @@ -23,7 +23,7 @@ float load_float(const char* data) { return value; } -TEST(DataLayoutTest, QgBatchSizeRejectsOverflow) { +TEST(FastScanDataLayoutTest, QgBatchSizeRejectsOverflow) { const size_t max_dim = std::numeric_limits::max() / fastscan::kBatchSize; EXPECT_THROW(QGBatchDataMap::data_bytes(max_dim + 1), std::length_error); EXPECT_THROW(ConstQGBatchDataMap::data_bytes(max_dim + 1), std::length_error); @@ -33,7 +33,7 @@ TEST(DataLayoutTest, QgBatchSizeRejectsOverflow) { ); } -TEST(DataLayoutTest, BatchFactorsKeepPackedOffsetsOnUnalignedStorage) { +TEST(FastScanDataLayoutTest, BatchFactorsKeepPackedOffsetsOnUnalignedStorage) { constexpr size_t kDim = 64; constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; constexpr size_t kDataBytes = kCodeBytes + (3 * fastscan::kBatchSize * sizeof(float)); @@ -67,7 +67,7 @@ TEST(DataLayoutTest, BatchFactorsKeepPackedOffsetsOnUnalignedStorage) { EXPECT_EQ(storage.back(), char{0x5a}); } -TEST(DataLayoutTest, QgBatchFactorsKeepPackedOffsetsOnUnalignedStorage) { +TEST(FastScanDataLayoutTest, QgBatchFactorsKeepPackedOffsetsOnUnalignedStorage) { constexpr size_t kDim = 64; constexpr size_t kCodeBytes = kDim * fastscan::kBatchSize / 8; constexpr size_t kDataBytes = kCodeBytes + (2 * fastscan::kBatchSize * sizeof(float)); @@ -153,7 +153,7 @@ void check_qg_encoding() { } } -TEST(QgBatchEncoding, PreservesCodesAndFactorsWithExplicitAndZeroCentroids) { +TEST(FastScanQgBatchEncodingTest, PreservesCodesAndFactorsWithExplicitAndZeroCentroids) { check_qg_encoding(); check_qg_encoding(); } diff --git a/tests/unit/rabitqlib/fastscan/encoder_test.cpp b/tests/unit/rabitqlib/fastscan/encoder_test.cpp index 82faa0e4..ff565ef5 100644 --- a/tests/unit/rabitqlib/fastscan/encoder_test.cpp +++ b/tests/unit/rabitqlib/fastscan/encoder_test.cpp @@ -78,12 +78,12 @@ void check_one_bit_batches() { } } -TEST(RabitqOneBitBatchTest, MatchesIndividualEncoding) { +TEST(FastScanEncoderTest, MatchesIndividualEncoding) { check_one_bit_batches(); check_one_bit_batches(); } -TEST(RabitqOneBitBatchTest, ParallelMatchesIndividualEncoding) { +TEST(FastScanEncoderTest, ParallelMatchesIndividualEncoding) { #if RABITQ_USE_OPENMP const int previous_threads = omp_get_max_threads(); omp_set_num_threads(2); diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 65dccb1c..25260ff2 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -98,8 +98,9 @@ TEST(FastScanLutTest, EverySupportedBackendMatchesOrderedSubsetSumsExactly) { TEST(FastScanLutTest, ResetOwnsBytesAndReusesStorageWithoutRetainingScratch) { constexpr size_t kDim = 64; std::array query{}; - for (size_t i = 0; i < kDim; ++i) + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast(i % 17) - 8; + } for (bool hacc : {false, true}) { Lut reused; std::vector expected; @@ -171,7 +172,7 @@ TEST(FastScanLutTest, RejectedResetPreservesOwnedBytesAndFactors) { EXPECT_EQ(lut.sum_vl(), sum); } -TEST(BatchEstimatorTest, RejectedQueryResetPreservesMetricAndCorrections) { +TEST(FastScanBatchEstimatorTest, RejectedQueryResetPreservesMetricAndCorrections) { std::array values{}; values.fill(1); for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -486,16 +487,15 @@ TEST(FastScanHighAccuracyTest, Avx512TransferAcceptsUnalignedOutput) { } std::array expected{}; - constexpr size_t kCodebooksPerRegister = 4; + constexpr size_t kGroupsPerRegister = 4; constexpr size_t kBytesPerRegisterPair = 128; - constexpr size_t kBytesPerCodebook = 16; + constexpr size_t kBytesPerGroup = 16; constexpr size_t kHighByteOffset = 64; - for (size_t codebook = 0; codebook < kDim / 4; ++codebook) { - const size_t low_offset = - (codebook / kCodebooksPerRegister * kBytesPerRegisterPair) + - (codebook % kCodebooksPerRegister * kBytesPerCodebook); + for (size_t group = 0; group < kDim / 4; ++group) { + const size_t low_offset = (group / kGroupsPerRegister * kBytesPerRegisterPair) + + (group % kGroupsPerRegister * kBytesPerGroup); for (size_t entry = 0; entry < 16; ++entry) { - const uint16_t value = lut[codebook * 16 + entry]; + const uint16_t value = lut[group * 16 + entry]; expected[low_offset + entry] = static_cast(value); expected[low_offset + kHighByteOffset + entry] = static_cast(value >> 8); @@ -510,7 +510,7 @@ TEST(FastScanHighAccuracyTest, Avx512TransferAcceptsUnalignedOutput) { #endif -TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { +TEST(FastScanBatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); std::vector backends{ rabitqlib::simd::split_batch_estdist, rabitqlib::simd::split_batch_estdist_generic}; @@ -563,7 +563,7 @@ TEST(BatchEstimatorTest, CollinearResidualKeepsFiniteLowerBounds) { } #if defined(__x86_64__) || defined(_M_X64) -TEST(BatchEstimatorTest, Avx512RejectsInvalidDimensionsBeforeReadingBuffers) { +TEST(FastScanBatchEstimatorTest, Avx512RejectsInvalidDimensionsBeforeReadingBuffers) { if (!cpu::has_avx512_core()) { GTEST_SKIP() << "AVX-512 is unavailable"; } @@ -587,7 +587,7 @@ TEST(BatchEstimatorTest, Avx512RejectsInvalidDimensionsBeforeReadingBuffers) { } #endif -TEST(BatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches) { +TEST(FastScanBatchEstimatorTest, BackendsMatchScalarCorrectionAcrossChunksAndTailBatches) { using Estimator = decltype(&rabitqlib::simd::split_batch_estdist); std::vector backends{ rabitqlib::simd::split_batch_estdist, rabitqlib::simd::split_batch_estdist_generic}; diff --git a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp index c7bbcf39..11b664be 100644 --- a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp +++ b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp @@ -77,12 +77,12 @@ struct SavedEnvironment { ~SavedEnvironment() { std::fesetenv(&environment); } }; -template +template void expect_bits_equal(T a, T b) { EXPECT_EQ(std::memcmp(&a, &b, sizeof(T)), 0); } -TEST(LutPreparationTest, FusedBackendsPreserveOrderedTablesAndFiniteExtrema) { +TEST(FastScanLutPreparationTest, FusedBackendsPreserveOrderedTablesAndFiniteExtrema) { const float tiny = std::numeric_limits::denorm_min(); const std::array, 6> patterns{ {{1, 2, 4, 8}, @@ -97,8 +97,9 @@ TEST(LutPreparationTest, FusedBackendsPreserveOrderedTablesAndFiniteExtrema) { SCOPED_TRACE(dim); for (const auto& pattern : patterns) { std::vector query(dim + 1); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { query[i + 1] = pattern[i % 4]; + } const auto original = query; std::vector expected(dim * 4 + 2, 123), actual = expected; backend.pack(dim, query.data() + 1, expected.data() + 1); @@ -125,7 +126,7 @@ TEST(LutPreparationTest, FusedBackendsPreserveOrderedTablesAndFiniteExtrema) { } } -TEST(LutPreparationTest, WrapperPreservesExceptionalAndSignedZeroBounds) { +TEST(FastScanLutPreparationTest, WrapperPreservesExceptionalAndSignedZeroBounds) { const SavedEnvironment saved; for (int rounding : {FE_TONEAREST, FE_DOWNWARD, FE_UPWARD, FE_TOWARDZERO}) { ASSERT_EQ(std::fesetround(rounding), 0); @@ -158,7 +159,7 @@ TEST(LutPreparationTest, WrapperPreservesExceptionalAndSignedZeroBounds) { } } -TEST(LutPreparationTest, Uint8PathPreservesHalfBoundariesFlagsAndZeroRange) { +TEST(FastScanLutPreparationTest, Uint8PathPreservesHalfBoundariesFlagsAndZeroRange) { const SavedEnvironment saved; std::vector boundaries{0, -0.0F, 255}; for (int value = 0; value < 255; ++value) { @@ -173,8 +174,9 @@ TEST(LutPreparationTest, Uint8PathPreservesHalfBoundariesFlagsAndZeroRange) { ASSERT_EQ(std::fesetround(rounding), 0); for (size_t count : {0U, 1U, 7U, 8U, 15U, 16U, 17U, 768U}) { std::vector values(count + 1); - for (size_t i = 0; i < count; ++i) + for (size_t i = 0; i < count; ++i) { values[i + 1] = boundaries[i % boundaries.size()]; + } std::vector expected(count + 2, 71), actual = expected; std::feclearexcept(FE_ALL_EXCEPT); backend.reference_quantize( @@ -238,11 +240,13 @@ void check_prepared_lut() { const T value = delta == 0 ? T{0} : std::round((table[i] - lo) * reciprocal); codes[i] = static_cast(value); - if (!hacc) + if (!hacc) { expected[i] = static_cast(value); + } } - if (hacc) + if (hacc) { transfer_lut_hacc(codes.data(), dim, expected.data()); + } Lut actual(query.data(), dim, hacc); expect_bits_equal(actual.delta(), delta); expect_bits_equal(actual.sum_vl(), lo * static_cast(dim / 4)); @@ -252,10 +256,10 @@ void check_prepared_lut() { } } -TEST(LutPreparationTest, FloatLutPreservesStandardAndHaccCodes) { +TEST(FastScanLutPreparationTest, FloatLutPreservesStandardAndHaccCodes) { check_prepared_lut(); } -TEST(LutPreparationTest, DoubleLutRetainsOriginalPreparation) { +TEST(FastScanLutPreparationTest, DoubleLutRetainsOriginalPreparation) { check_prepared_lut(); } } // namespace diff --git a/tests/unit/rabitqlib/index/exact_nn_test.cpp b/tests/unit/rabitqlib/index/exact_nn_test.cpp index 9dcc9fef..ac3b173e 100644 --- a/tests/unit/rabitqlib/index/exact_nn_test.cpp +++ b/tests/unit/rabitqlib/index/exact_nn_test.cpp @@ -13,7 +13,7 @@ using namespace rabitqlib; -TEST(IndexTypes, PreserveIdentifierAndCandidateRepresentation) { +TEST(IndexTypesTest, PreserveIdentifierAndCandidateRepresentation) { static_assert(std::is_same_v); EXPECT_EQ(kPidMax, std::numeric_limits::max()); const AnnCandidate candidate(7, 1.5F); @@ -21,7 +21,7 @@ TEST(IndexTypes, PreserveIdentifierAndCandidateRepresentation) { EXPECT_EQ(candidate.distance, 1.5F); } -TEST(ExactNearestNeighbor, SupportsRequestedWorkerCounts) { +TEST(ExactNearestNeighborTest, SupportsRequestedWorkerCounts) { const std::array data{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; const std::array query{7, 8, 9}; for (size_t workers : {0U, 1U, 2U, 64U}) { diff --git a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp index aec78798..17daef67 100644 --- a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp +++ b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp @@ -94,15 +94,16 @@ TEST(IvfInitializerTest, MatrixRotationAndExplicitRoutingRoundTrip) { } TEST(IvfInitializerTest, ReloadChangesRoutingAndFailedLoadPreservesExistingIndex) { - constexpr size_t kDim = 65, kCount = 33, kTopk = 5; + constexpr size_t kDim = 65, kCount = 33, kTopK = 5; std::mt19937 generator(9923); std::normal_distribution normal; std::vector data(kCount * kDim); std::vector labels(kCount); for (size_t i = 0; i < kCount; ++i) { labels[i] = static_cast(i); - for (size_t j = 0; j < kDim; ++j) + for (size_t j = 0; j < kDim; ++j) { data[i * kDim + j] = normal(generator); + } } const rabitq_test::TempFile path_file("ivf_initializer_reload.index"); const std::string& path = path_file.path(); @@ -114,15 +115,15 @@ TEST(IvfInitializerTest, ReloadChangesRoutingAndFailedLoadPreservesExistingIndex InitializerType::Flat}) { IVF source(kCount, kDim, kCount, 5, METRIC_L2, RotatorType::FhtKacRotator, type); source.construct(data.data(), data.data(), labels.data(), false, 1); - std::array expected_ids{}, actual_ids{}; - std::array expected_distances{}, actual_distances{}; + std::array expected_ids{}, actual_ids{}; + std::array expected_distances{}, actual_distances{}; source.search( - data.data(), kTopk, 7, expected_ids.data(), expected_distances.data() + data.data(), kTopK, 7, expected_ids.data(), expected_distances.data() ); source.save(path.c_str()); target.load(path.c_str()); EXPECT_EQ(target.initializer_type(), type); - target.search(data.data(), kTopk, 7, actual_ids.data(), actual_distances.data()); + target.search(data.data(), kTopK, 7, actual_ids.data(), actual_distances.data()); EXPECT_EQ(actual_ids, expected_ids); EXPECT_EQ(actual_distances, expected_distances); { @@ -134,13 +135,13 @@ TEST(IvfInitializerTest, ReloadChangesRoutingAndFailedLoadPreservesExistingIndex } EXPECT_THROW(target.load(path.c_str()), std::runtime_error); EXPECT_EQ(target.initializer_type(), type); - target.search(data.data(), kTopk, 7, actual_ids.data(), actual_distances.data()); + target.search(data.data(), kTopK, 7, actual_ids.data(), actual_distances.data()); EXPECT_EQ(actual_ids, expected_ids); EXPECT_EQ(actual_distances, expected_distances); } } -TEST(FlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { +TEST(IvfFlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { std::mt19937 generator(141); std::normal_distribution normal; for (size_t dim : {32U, 64U, 96U, 1056U}) { @@ -152,8 +153,9 @@ TEST(FlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { // Unit vectors make self-query the winner under both metrics. for (size_t i = 0; i < count; ++i) { const float norm = std::sqrt(l2norm_sqr(centroids.data() + i * dim, dim)); - for (size_t j = 0; j < dim; ++j) + for (size_t j = 0; j < dim; ++j) { centroids[i * dim + j] /= norm; + } } for (auto metric : {METRIC_L2, METRIC_IP}) { FlatRaBitQInitializer index(dim, count, metric); @@ -163,8 +165,9 @@ TEST(FlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { std::vector> result; index.centroids_distances(query, nprobe, result); ASSERT_EQ(result.size(), std::min(nprobe, count)); - if (result.empty()) + if (result.empty()) { continue; + } EXPECT_EQ(result[0].id, count - 1); EXPECT_TRUE(std::is_sorted(result.begin(), result.end())); std::vector seen(count); @@ -186,7 +189,7 @@ TEST(FlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { } } -TEST(FlatRaBitQInitializerTest, ZeroResidualsAndInnerProductMagnitude) { +TEST(IvfFlatRaBitQInitializerTest, ZeroResidualsAndInnerProductMagnitude) { constexpr size_t kDim = 96, kCount = 33; std::vector centroids(kDim * kCount, 0.0F), query(kDim, 0.0F); for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -195,11 +198,13 @@ TEST(FlatRaBitQInitializerTest, ZeroResidualsAndInnerProductMagnitude) { std::vector> result; index.centroids_distances(query.data(), 7, result); ASSERT_EQ(result.size(), 7U); - for (const auto& item : result) + for (const auto& item : result) { EXPECT_FLOAT_EQ(item.distance, metric == METRIC_L2 ? 0 : 1); + } } - for (size_t i = 0; i < kCount; ++i) + for (size_t i = 0; i < kCount; ++i) { centroids[i * kDim] = 0.5F; + } centroids[17 * kDim] = 100.0F; query[0] = 1.0F; FlatRaBitQInitializer index(kDim, kCount, METRIC_IP); @@ -211,7 +216,7 @@ TEST(FlatRaBitQInitializerTest, ZeroResidualsAndInnerProductMagnitude) { EXPECT_FLOAT_EQ(result[0].distance, -99.0F); } -TEST(FlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { +TEST(IvfFlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { constexpr size_t kDim = 96, kCount = 97; std::mt19937 generator(731); std::normal_distribution normal; @@ -244,8 +249,9 @@ TEST(FlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { loaded.centroids_distances(queries.data() + i * kDim, 7, actual[i]); }); } - for (auto& thread : threads) + for (auto& thread : threads) { thread.join(); + } for (size_t i = 0; i < 8; ++i) { ASSERT_EQ(actual[i].size(), expected[i].size()); for (size_t j = 0; j < actual[i].size(); ++j) { diff --git a/tests/unit/rabitqlib/index/historical_persistence_test.cpp b/tests/unit/rabitqlib/index/historical_persistence_test.cpp index 645f0214..f456acb1 100644 --- a/tests/unit/rabitqlib/index/historical_persistence_test.cpp +++ b/tests/unit/rabitqlib/index/historical_persistence_test.cpp @@ -58,10 +58,12 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail ASSERT_EQ(k, 5U); std::vector queries(nq * dim), distances(nq * k), reference_distances(nq * k); std::vector ids(nq * k), reference_ids(nq * k); - for (auto& value : queries) + for (auto& value : queries) { expected >> value; - for (size_t i = 0; i < ids.size(); ++i) + } + for (size_t i = 0; i < ids.size(); ++i) { expected >> reference_ids[i] >> reference_distances[i]; + } ASSERT_TRUE(expected.good()); const auto output = std::filesystem::temp_directory_path() / ("rabitq_historical_" + GetParam().directory + "_" + name); diff --git a/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp index be3f8cc4..3937f905 100644 --- a/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_sidecar_test.cpp @@ -48,7 +48,7 @@ struct Fixture { std::vector data = std::vector(kCount * kDim); std::vector original; std::vector upper_offsets, levels; - size_t stride = 0, upper_stride = 0, max_m = 0, max_m0 = 0; + size_t level0_record_bytes = 0, upper_layer_links_bytes = 0, max_m = 0, max_m0 = 0; size_t upper_node = 0, base_node = 0; explicit Fixture(MetricType metric) { @@ -77,15 +77,16 @@ struct Fixture { index.save(path.c_str()); std::ifstream file(path + ".hnsw", std::ios::binary); original.assign(std::istreambuf_iterator(file), {}); - stride = read_at(original, 3 * sizeof(size_t)); + level0_record_bytes = read_at(original, 3 * sizeof(size_t)); max_m = read_at(original, kMaxMOffset); max_m0 = read_at(original, kMaxMOffset + sizeof(size_t)); - upper_stride = sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); - size_t pos = kHeaderBytes + kCount * stride; + upper_layer_links_bytes = + sizeof(hnswlib::linklistsizeint) + max_m * sizeof(hnswlib::tableint); + size_t pos = kHeaderBytes + kCount * level0_record_bytes; for (size_t i = 0; i < kCount; ++i) { upper_offsets.push_back(pos); const auto bytes = read_at(original, pos); - levels.push_back(bytes / upper_stride); + levels.push_back(bytes / upper_layer_links_bytes); pos += sizeof(unsigned int) + bytes; } if (pos != original.size()) { @@ -111,7 +112,7 @@ struct Fixture { } }; -TEST(HnswSidecarGraph, PublicLoadRejectsCorruptStructureAndCanRetry) { +TEST(IvfHnswSidecarGraphTest, PublicLoadRejectsCorruptStructureAndCanRetry) { for (const auto metric : {METRIC_L2, METRIC_IP}) { Fixture fixture(metric); const auto reject = [&](const char* name, auto corrupt, const char* message) { @@ -164,7 +165,9 @@ TEST(HnswSidecarGraph, PublicLoadRejectsCorruptStructureAndCanRetry) { ); reject( "truncated base records", - [&](auto& b) { b.resize(kHeaderBytes + kCount * fixture.stride - 1); }, + [&](auto& b) { + b.resize(kHeaderBytes + kCount * fixture.level0_record_bytes - 1); + }, "Truncated HNSW centroid base records" ); reject( @@ -294,17 +297,17 @@ TEST(HnswSidecarGraph, PublicLoadRejectsCorruptStructureAndCanRetry) { } } -TEST(HnswSidecarGraph, DisconnectedGraphCannotReadEmptyResults) { +TEST(IvfHnswSidecarGraphTest, DisconnectedGraphCannotReadEmptyResults) { for (const auto metric : {METRIC_L2, METRIC_IP}) { Fixture fixture(metric); auto bytes = fixture.original; for (size_t i = 0; i < kCount; ++i) { - write_at(bytes, kHeaderBytes + i * fixture.stride, 0); + write_at(bytes, kHeaderBytes + i * fixture.level0_record_bytes, 0); for (size_t level = 0; level < fixture.levels[i]; ++level) { write_at( bytes, fixture.upper_offsets[i] + sizeof(unsigned int) + - level * fixture.upper_stride, + level * fixture.upper_layer_links_bytes, 0 ); } diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 638096ad..17ec4226 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -30,38 +30,36 @@ namespace rabitqlib::hnsw { -struct HnswPruningTestAccess { +struct HnswTestAccess { static void quantize(HierarchicalNSW& index, PID id, const float* rotated) { const auto* centroid = reinterpret_cast(index.centroids_memory_.get()); quant::quantize_compact_one_bit( - rotated, centroid, index.padded_dim_, index.get_bindata_by_internalid(id) + rotated, centroid, index.padded_dim_, index.bin_data(id) ); } static float distance(const HierarchicalNSW& index, PID target, PID query) { detail::QuantQuery cached; - return index.get_quant_dist(target, query, cached); + return index.quantized_distance(target, query, cached); } static std::vector connect(HierarchicalNSW& index, PID incoming) { - PID* neighbors = index.get_linklist0(0); - index.set_list_count(neighbors, index.base_layer_max_degree_); + PID* neighbors = index.level0_links(0); + index.set_link_count(neighbors, index.base_layer_max_degree_); for (size_t i = 0; i < index.base_layer_max_degree_; ++i) { neighbors[i + 1] = static_cast(i + 1); } - std::fill_n( - index.get_linklist0(incoming), index.base_layer_max_degree_ + 1, PID{0} - ); + std::fill_n(index.level0_links(incoming), index.base_layer_max_degree_ + 1, PID{0}); detail::QuantQuery cached; maxheap> candidates; - candidates.emplace(index.get_quant_dist(0, incoming, cached), 0); + candidates.emplace(index.quantized_distance(0, incoming, cached), 0); index.mutually_connect_quant(incoming, candidates, 0, cached); - return {neighbors + 1, neighbors + 1 + index.get_list_count(neighbors)}; + return {neighbors + 1, neighbors + 1 + index.link_count(neighbors)}; } static PID entry_point_label(const HierarchicalNSW& index) { - return index.get_external_label(index.entry_point_id_); + return index.external_label(index.entry_point_id_); } static PID internal_id(const HierarchicalNSW& index, PID label) { @@ -69,7 +67,7 @@ struct HnswPruningTestAccess { } static void set_raw_cluster_id(HierarchicalNSW& index, PID label, PID value) { - index.set_raw_clusterid(index.label_lookup_.at(label), value); + index.set_raw_cluster_id(index.label_lookup_.at(label), value); } // Give internal slot i the label count - 1 - i, as parallel construction can @@ -86,7 +84,7 @@ struct HnswPruningTestAccess { }; namespace { -TEST(HnswQuantQuery, MovesPreserveOwnedQueryAndLeaveSourceReusable) { +TEST(HnswQuantQueryTest, MovesPreserveOwnedQueryAndLeaveSourceReusable) { using detail::QuantQuery; static_assert(!std::is_copy_constructible_v); static_assert(!std::is_copy_assignable_v); @@ -144,7 +142,7 @@ TEST(HnswSearchTest, EqualDistancesUseExternalLabelOrder) { index.construct( 1, centroid.data(), kCount, data.data(), clusters.data(), 1, true ); - HnswPruningTestAccess::reverse_labels(index); + HnswTestAccess::reverse_labels(index); const PID removed = 3; index.remove(&removed, 1); for (size_t threads : {1U, 3U}) { @@ -182,17 +180,15 @@ TEST(HnswPruningTest, ScoresNewAndExistingNeighborsWithTheSameQuery) { // Set codes directly in the rotated domain to avoid random rotation changing // the asymmetry, then fill node 0's base-layer list to force pruning. for (PID id = 0; id < kCount; ++id) { - HnswPruningTestAccess::quantize(index, id, data.data() + (id * kDim)); + HnswTestAccess::quantize(index, id, data.data() + (id * kDim)); } - const float existing_distance = HnswPruningTestAccess::distance(index, 1, 0); - EXPECT_LT(HnswPruningTestAccess::distance(index, kIncoming, 0), existing_distance); - EXPECT_GT(HnswPruningTestAccess::distance(index, 0, kIncoming), existing_distance); + const float existing_distance = HnswTestAccess::distance(index, 1, 0); + EXPECT_LT(HnswTestAccess::distance(index, kIncoming, 0), existing_distance); + EXPECT_GT(HnswTestAccess::distance(index, 0, kIncoming), existing_distance); // The new point is closest when node 0 is the query. Its duplicate existing // neighbors are pruned; the reversed estimate instead retains an old point. - EXPECT_EQ( - HnswPruningTestAccess::connect(index, kIncoming), std::vector{kIncoming} - ); + EXPECT_EQ(HnswTestAccess::connect(index, kIncoming), std::vector{kIncoming}); } TEST(HnswConfigurationTest, RejectsUnsupportedMetric) { @@ -203,26 +199,26 @@ TEST(HnswConfigurationTest, RejectsUnsupportedMetric) { } TEST(HnswOneBitSearchTest, UsesBinaryEstimateForEntryPoint) { - constexpr size_t dim = 64; - constexpr size_t count = 8; - std::vector data(count * dim); - std::vector centroid(dim); - std::vector query(dim); + constexpr size_t kDim = 64; + constexpr size_t kCount = 8; + std::vector data(kCount * kDim); + std::vector centroid(kDim); + std::vector query(kDim); - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { centroid[i] = static_cast(static_cast(i % 13) - 6) / 5.0F; query[i] = centroid[i]; } - for (size_t point = 0; point < count; ++point) { - for (size_t i = 0; i < dim; ++i) { - data[(point * dim) + i] = + for (size_t point = 0; point < kCount; ++point) { + for (size_t i = 0; i < kDim; ++i) { + data[(point * kDim) + i] = centroid[i] + static_cast((point + 1) * ((i % 11) + 1)); } } - std::vector cluster_ids(count, 0); - HierarchicalNSW index(count, dim, 1, 2, 10, 100, METRIC_L2); - index.construct(1, centroid.data(), count, data.data(), cluster_ids.data(), 1, false); + std::vector cluster_ids(kCount, 0); + HierarchicalNSW index(kCount, kDim, 1, 2, 10, 100, METRIC_L2); + index.construct(1, centroid.data(), kCount, data.data(), cluster_ids.data(), 1, false); const auto results = index.search(query.data(), 1, 1, 10, 1); @@ -230,14 +226,14 @@ TEST(HnswOneBitSearchTest, UsesBinaryEstimateForEntryPoint) { ASSERT_EQ(results[0].size(), 1U); ASSERT_EQ(results[0][0].second, 0U); EXPECT_TRUE(std::isfinite(results[0][0].first)); - const float exact_distance = euclidean_sqr(query.data(), data.data(), dim); + const float exact_distance = euclidean_sqr(query.data(), data.data(), kDim); EXPECT_NEAR(results[0][0].first, exact_distance, exact_distance * 0.1F); } TEST(HnswConstructionTest, ParallelConstructionProducesSearchableIndex) { constexpr size_t kDim = 64; constexpr size_t kCount = 512; - constexpr size_t kQueries = 16; + constexpr size_t kNumQueries = 16; std::mt19937 generator(42); std::uniform_real_distribution distribution(-1, 1); std::vector data(kCount * kDim); @@ -249,9 +245,9 @@ TEST(HnswConstructionTest, ParallelConstructionProducesSearchableIndex) { HierarchicalNSW index(kCount, kDim, 4, 8, 50); index.construct(1, centroid.data(), kCount, data.data(), cluster_ids.data(), 8, false); - const auto results = index.search(data.data(), kQueries, 1, kCount, 1); - ASSERT_EQ(results.size(), kQueries); - for (size_t i = 0; i < kQueries; ++i) { + const auto results = index.search(data.data(), kNumQueries, 1, kCount, 1); + ASSERT_EQ(results.size(), kNumQueries); + for (size_t i = 0; i < kNumQueries; ++i) { ASSERT_EQ(results[i].size(), 1U); EXPECT_EQ(results[i][0].second, i); EXPECT_TRUE(std::isfinite(results[i][0].first)); @@ -845,7 +841,7 @@ TEST_F(HnswSaveTest, RejectsUnopenableDestination) { index_.save((path_ + "/index").c_str()); FAIL() << "Saving to a missing directory must fail"; } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "HNSW: cannot open index file for writing"); + EXPECT_STREQ(error.what(), "HNSW cannot open index file for writing"); } } @@ -891,7 +887,7 @@ TEST_F(HnswSaveTest, ReportsWriteOrCloseFailure) { index_.save("/dev/full"); FAIL() << "Saving to a full destination must fail"; } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "HNSW: failed to write index file"); + EXPECT_STREQ(error.what(), "HNSW failed to write index file"); } } @@ -1062,7 +1058,7 @@ TEST(HnswConstructionTest, SeedReproducesRotationCodesAndSerialSearch) { TEST(HnswRemoveTest, ExcludesRemovedPointsAndKeepsSurvivors) { constexpr size_t kCount = 120; - constexpr size_t kQueries = 8; + constexpr size_t kNumQueries = 8; for (const auto& config : remove_cases()) { SCOPED_TRACE( ::testing::Message() @@ -1070,7 +1066,7 @@ TEST(HnswRemoveTest, ExcludesRemovedPointsAndKeepsSurvivors) { ); AddFixture fixture(kCount, 59); auto index = build_for_remove(fixture, kCount, kCount, config); - const Results before = search_everything(*index, fixture, kQueries); + const Results before = search_everything(*index, fixture, kNumQueries); // Every third point, the last one, and a repeated label. std::vector labels; @@ -1092,11 +1088,12 @@ TEST(HnswRemoveTest, ExcludesRemovedPointsAndKeepsSurvivors) { // Survivors keep their order and their distances. EXPECT_EQ( - search_everything(*index, fixture, kQueries), without_removed(before, removed) + search_everything(*index, fixture, kNumQueries), + without_removed(before, removed) ); // A small ef still fills k with live points only. - const auto top = index->search(fixture.data.data(), kQueries, 10, 10, 1); + const auto top = index->search(fixture.data.data(), kNumQueries, 10, 10, 1); for (const auto& hits : top) { EXPECT_EQ(hits.size(), 10U); for (const auto& hit : hits) { @@ -1117,7 +1114,7 @@ TEST(HnswRemoveTest, RemovingTheEntryPointKeepsSearchWorking) { AddFixture fixture(kCount, 61); auto index = build_for_remove(fixture, kCount, kCount, config); const Results before = search_everything(*index, fixture, 1); - const PID entry = HnswPruningTestAccess::entry_point_label(*index); + const PID entry = HnswTestAccess::entry_point_label(*index); std::vector removed(kCount, false); removed[entry] = true; ASSERT_EQ(index->remove(&entry, 1), 1U); @@ -1171,7 +1168,7 @@ TEST(HnswRemoveTest, UsesLabelsNotInternalIds) { const RemoveCase config{METRIC_L2, 4}; AddFixture fixture(kCount, 71); auto index = build_for_remove(fixture, kCount, kCount, config); - HnswPruningTestAccess::reverse_labels(*index); + HnswTestAccess::reverse_labels(*index); const Results before = search_everything(*index, fixture, 1); // Every fifth label below kCount / 2: no removed label shares a slot with @@ -1317,13 +1314,13 @@ TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { ASSERT_NO_THROW(loaded.load(path.c_str())); // Cluster 1 does not exist in a one-cluster index, marked or not. - HnswPruningTestAccess::set_raw_cluster_id(*index, one, 0x80000001U); + HnswTestAccess::set_raw_cluster_id(*index, one, 0x80000001U); index->save(path.c_str()); try { loaded.load(path.c_str()); FAIL() << "A marked out-of-range cluster ID must be rejected"; } catch (const std::runtime_error& error) { - EXPECT_STREQ(error.what(), "HNSW: invalid or truncated index file"); + EXPECT_STREQ(error.what(), "HNSW index file is invalid or truncated"); } } @@ -1333,54 +1330,53 @@ TEST(HnswRemoveTest, LoadRejectsAnOutOfRangeClusterEvenWithTheMark) { TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { using namespace rabitqlib; using namespace rabitqlib::hnsw; - constexpr size_t count = 24; + constexpr size_t kCount = 24; using Search = decltype(&hnsw::detail::search_knn_sorted_generic); std::vector backends{hnsw::detail::search_knn_sorted}; #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) + if (cpu::has_avx2()) { backends.push_back(hnsw::detail::search_knn_sorted_avx2); - if (cpu::has_avx512_core()) + } + if (cpu::has_avx512_core()) { backends.push_back(hnsw::detail::search_knn_sorted_avx512_core); - if (cpu::has_avx512_popcnt()) + } + if (cpu::has_avx512_popcnt()) { backends.push_back(hnsw::detail::search_knn_sorted_avx512_popcnt); + } #endif #if defined(__aarch64__) || defined(_M_ARM64) - if (cpu::has_neon()) + if (cpu::has_neon()) { backends.push_back(hnsw::detail::search_knn_sorted_neon); + } #endif for (size_t dim : {64U, 96U, 128U, 160U, 256U, 544U, 576U}) { std::mt19937 random(642); std::uniform_real_distribution value(-1, 1); - std::vector data(count * dim), centroid(dim, 0), query(dim); - for (auto& x : data) + std::vector data(kCount * dim), centroid(dim, 0), query(dim); + for (auto& x : data) { x = value(random); - for (auto& x : query) + } + for (auto& x : query) { x = value(random); + } for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits = 1; bits <= 9; ++bits) { SCOPED_TRACE(dim); SCOPED_TRACE(bits); SCOPED_TRACE(metric); - std::vector clusters(count, 0); - HierarchicalNSW index(count, dim, bits, 8, count, 42, metric); + std::vector clusters(kCount, 0); + HierarchicalNSW index(kCount, dim, bits, 8, kCount, 42, metric); index.construct( - 1, centroid.data(), count - 4, data.data(), clusters.data(), 1, true + 1, centroid.data(), kCount - 4, data.data(), clusters.data(), 1, true ); // Insertion reconstructs stored sign/extra-bit tails before linking. - index.add(data.data() + (count - 4) * dim, 4, clusters.data(), true, 1); + index.add(data.data() + (kCount - 4) * dim, 4, clusters.data(), true, 1); + // backends[0] is the dispatched entry point that batch search uses. const auto reference = - hnsw::detail::search_knn_sorted_generic(index, query.data(), 5, count); - const auto ordered = - hnsw::detail::search_knn_sorted(index, query.data(), 5, count); - auto heap = hnsw::detail::search_knn(index, query.data(), 5, count); - ASSERT_EQ(heap.size(), ordered.size()); - for (auto it = ordered.rbegin(); it != ordered.rend(); ++it) { - EXPECT_EQ(heap.top(), *it); - heap.pop(); - } + hnsw::detail::search_knn_sorted_generic(index, query.data(), 5, kCount); for (auto search : backends) { const auto& expected = reference; - auto actual = search(index, query.data(), 5, count); + auto actual = search(index, query.data(), 5, kCount); ASSERT_EQ(actual.size(), expected.size()); for (size_t i = 0; i < expected.size(); ++i) { EXPECT_EQ(actual[i].second, expected[i].second); @@ -1398,14 +1394,16 @@ TEST(HnswPaddingTest, TailDimensionsMatchEveryAvailableSearchBackend) { TEST(HnswSearchTest, ConcurrentCallsKeepIndependentSearchWindows) { using namespace rabitqlib; - constexpr size_t kCount = 129, kDim = 65, kQueries = 13, k = 5; + constexpr size_t kCount = 129, kDim = 65, kNumQueries = 13, kTopK = 5; std::mt19937 random(314); std::normal_distribution value; - std::vector data(kCount * kDim), queries(kQueries * kDim), centroid(kDim, 0); - for (auto& x : data) + std::vector data(kCount * kDim), queries(kNumQueries * kDim), centroid(kDim, 0); + for (auto& x : data) { x = value(random); - for (auto& x : queries) + } + for (auto& x : queries) { x = value(random); + } for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {1U, 4U, 9U}) { SCOPED_TRACE(metric); @@ -1415,19 +1413,22 @@ TEST(HnswSearchTest, ConcurrentCallsKeepIndependentSearchWindows) { index.construct( 1, centroid.data(), kCount, data.data(), clusters.data(), 1, true ); - const auto narrow = index.search(queries.data(), kQueries, k, k, 1); - const auto wide = index.search(queries.data(), kQueries, k, 96, 1); + const auto narrow = index.search(queries.data(), kNumQueries, kTopK, kTopK, 1); + const auto wide = index.search(queries.data(), kNumQueries, kTopK, 96, 1); std::promise start; const auto ready = start.get_future().share(); const auto run = [&](size_t ef, size_t threads, const auto& expected) { ready.wait(); for (size_t i = 0; i < 16; ++i) { - if (index.search(queries.data(), kQueries, k, ef, threads) != expected) + if (index.search(queries.data(), kNumQueries, kTopK, ef, threads) != + expected) { return false; + } } return true; }; - auto first = std::async(std::launch::async, [&] { return run(k, 1, narrow); }); + auto first = + std::async(std::launch::async, [&] { return run(kTopK, 1, narrow); }); auto second = std::async(std::launch::async, [&] { return run(96, 2, wide); }); start.set_value(); EXPECT_TRUE(first.get()); @@ -1436,7 +1437,7 @@ TEST(HnswSearchTest, ConcurrentCallsKeepIndependentSearchWindows) { } } -TEST(HnswSearchValidation, RejectsZeroKBeforeParallelSearch) { +TEST(HnswSearchValidationTest, RejectsZeroKBeforeParallelSearch) { using namespace rabitqlib; std::vector data(128, 0.5F), center(64, 0.0F); PID clusters[2] = {0, 0}; diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index 3355beeb..5abe7820 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -29,7 +29,7 @@ namespace rabitqlib::ivf { namespace { -TEST(HNSWInitializerTest, PersistsUtf8SidecarPaths) { +TEST(IvfHnswInitializerTest, PersistsUtf8SidecarPaths) { const auto directory = std::filesystem::temp_directory_path() / std::filesystem::u8path(u8"rabitq_\u6d4b\u8bd5_\U0001f680"); std::filesystem::create_directories(directory); @@ -51,7 +51,7 @@ TEST(HNSWInitializerTest, PersistsUtf8SidecarPaths) { std::filesystem::remove(directory); } -TEST(HNSWInitializerTest, ReportsSidecarOpenFailureAndCanRetry) { +TEST(IvfHnswInitializerTest, ReportsSidecarOpenFailureAndCanRetry) { const std::string path = ::testing::TempDir() + "rabitq_hnsw_save_failure"; const std::filesystem::path sidecar(path + ".hnsw"); std::filesystem::create_directory(sidecar); @@ -78,7 +78,7 @@ TEST(HNSWInitializerTest, ReportsSidecarOpenFailureAndCanRetry) { std::filesystem::remove(sidecar); } -TEST(HNSWInitializerTest, ReportsBufferedWriteFailure) { +TEST(IvfHnswInitializerTest, ReportsBufferedWriteFailure) { if (!std::filesystem::exists("/dev/full")) { GTEST_SKIP() << "Requires a device that rejects writes"; } @@ -94,7 +94,7 @@ TEST(HNSWInitializerTest, ReportsBufferedWriteFailure) { EXPECT_FLOAT_EQ(result.top().first, 0.0F); } -TEST(CentroidL2SpaceTest, MatchesDispatchedDistanceForUnalignedInputsAndTails) { +TEST(IvfCentroidL2SpaceTest, MatchesDispatchedDistanceForUnalignedInputsAndTails) { for (size_t dim : {1U, 7U, 8U, 15U, 16U, 17U, 63U, 64U, 129U}) { CentroidL2Space space(dim); std::vector a(dim + 1), b(dim + 1); @@ -113,7 +113,7 @@ TEST(CentroidL2SpaceTest, MatchesDispatchedDistanceForUnalignedInputsAndTails) { } } -TEST(HNSWInitializerTest, RoutesToExactCentroidsWithEuclideanDistances) { +TEST(IvfHnswInitializerTest, RoutesToExactCentroidsWithEuclideanDistances) { constexpr size_t kDim = 17; constexpr size_t kCount = 32; std::vector centroids(kCount * kDim); @@ -134,7 +134,7 @@ TEST(HNSWInitializerTest, RoutesToExactCentroidsWithEuclideanDistances) { } } -TEST(HNSWInitializerTest, RoutesByInnerProductWithoutTakingSquareRoot) { +TEST(IvfHnswInitializerTest, RoutesByInnerProductWithoutTakingSquareRoot) { constexpr size_t kDim = 17; constexpr size_t kCount = 32; std::vector centroids(kCount * kDim, 0.0F); @@ -184,7 +184,7 @@ TEST(IvfInitializerTest, CentroidSearchSizesOutputAndClampsProbeCount) { } } -TEST(HNSWInitializerTest, ParallelInsertionPreservesRandomLevelDraws) { +TEST(IvfHnswInitializerTest, ParallelInsertionPreservesRandomLevelDraws) { constexpr size_t kDim = 17; constexpr size_t kCount = 2048; constexpr size_t kSeed = 100; @@ -213,7 +213,7 @@ TEST(HNSWInitializerTest, ParallelInsertionPreservesRandomLevelDraws) { } } -TEST(HNSWInitializerTest, ParallelInsertionPreservesCentroidLabelsAfterReload) { +TEST(IvfHnswInitializerTest, ParallelInsertionPreservesCentroidLabelsAfterReload) { constexpr size_t kDim = 17; constexpr size_t kCount = 1024; std::vector centroids(kCount * kDim); @@ -245,7 +245,7 @@ TEST(HNSWInitializerTest, ParallelInsertionPreservesCentroidLabelsAfterReload) { } } -TEST(HNSWInitializerTest, ReloadsInnerProductGraphForRouting) { +TEST(IvfHnswInitializerTest, ReloadsInnerProductGraphForRouting) { constexpr size_t kDim = 17; constexpr size_t kCount = 32; std::vector centroids(kCount * kDim, 0.0F); @@ -269,7 +269,7 @@ TEST(HNSWInitializerTest, ReloadsInnerProductGraphForRouting) { EXPECT_FLOAT_EQ(candidates[0].distance, -99.0F); } -TEST(HNSWInitializerTest, PerCallSearchWindowsPreserveDefaultAndLegacyResults) { +TEST(IvfHnswInitializerTest, PerCallSearchWindowsPreserveDefaultAndLegacyResults) { constexpr size_t kDim = 17; constexpr size_t kCount = 1024; std::mt19937 rng(42); @@ -306,7 +306,7 @@ TEST(HNSWInitializerTest, PerCallSearchWindowsPreserveDefaultAndLegacyResults) { } } -TEST(HNSWInitializerTest, ConcurrentSearchesMatchSerialWithDifferentWindows) { +TEST(IvfHnswInitializerTest, ConcurrentSearchesMatchSerialWithDifferentWindows) { constexpr size_t kDim = 17; constexpr size_t kCount = 1024; const std::array probes{1, 16, 400, 500}; diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index c1eaf3fe..6bcc640a 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -49,45 +49,45 @@ struct IvfTestAccess { namespace { TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { - constexpr size_t kNum = 97, kDim = 65, kClusters = 4, kQueries = 13, k = 7; + constexpr size_t kCount = 97, kDim = 65, kNumClusters = 4, kNumQueries = 13, kTopK = 7; std::mt19937 rng(73); std::uniform_real_distribution random(-1, 1); - std::vector data(kNum * kDim), centroids(kClusters * kDim); - std::vector queries(kQueries * kDim); - std::vector labels(kNum); + std::vector data(kCount * kDim), centroids(kNumClusters * kDim); + std::vector queries(kNumQueries * kDim); + std::vector labels(kCount); std::generate(data.begin(), data.end(), [&] { return random(rng); }); std::generate(centroids.begin(), centroids.end(), [&] { return random(rng); }); std::generate(queries.begin(), queries.end(), [&] { return random(rng); }); - for (size_t i = 0; i < kNum; ++i) { - labels[i] = static_cast(i % (kClusters - 1)); + for (size_t i = 0; i < kCount; ++i) { + labels[i] = static_cast(i % (kNumClusters - 1)); } for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U, 32U}) { - IVF index(kNum, kDim, kClusters, bits, metric); + IVF index(kCount, kDim, kNumClusters, bits, metric); EXPECT_EQ(index.padded_dim(), 96U); index.construct(data.data(), centroids.data(), labels.data(), false, 1); - std::array expected{}, actual{}, ids_only{}; - std::array expected_distances{}, actual_distances{}; + std::array expected{}, actual{}, ids_only{}; + std::array expected_distances{}, actual_distances{}; for (size_t nprobe : {1U, 99U}) { for (std::optional hacc : {std::optional{}, std::optional(false), std::optional(true)}) { - for (size_t i = 0; i < kQueries; ++i) { + for (size_t i = 0; i < kNumQueries; ++i) { index.search( queries.data() + i * kDim, - k, + kTopK, nprobe, - expected.data() + i * k, - expected_distances.data() + i * k, + expected.data() + i * kTopK, + expected_distances.data() + i * kTopK, hacc.value_or(bits >= 4 && bits <= 9) ); } for (size_t threads : {1U, 4U}) { index.search_batch( queries.data(), - kQueries, - k, + kNumQueries, + kTopK, nprobe, actual.data(), actual_distances.data(), @@ -96,8 +96,8 @@ TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { ); index.search_batch( queries.data(), - kQueries, - k, + kNumQueries, + kTopK, nprobe, ids_only.data(), nullptr, @@ -110,34 +110,40 @@ TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { } } } - EXPECT_NO_THROW(index.search_batch(nullptr, 0, k, 1, nullptr)); + EXPECT_NO_THROW(index.search_batch(nullptr, 0, kTopK, 1, nullptr)); EXPECT_THROW( index.search_batch( - queries.data(), std::numeric_limits::max(), k, 1, actual.data() + queries.data(), + std::numeric_limits::max(), + kTopK, + 1, + actual.data() ), std::length_error ); EXPECT_THROW( - index.search_batch(nullptr, 1, k, 1, actual.data()), std::invalid_argument + index.search_batch(nullptr, 1, kTopK, 1, actual.data()), + std::invalid_argument ); EXPECT_THROW( index.search_batch(queries.data(), 1, 0, 1, actual.data()), std::invalid_argument ); EXPECT_THROW( - index.search_batch(queries.data(), 1, k, 0, actual.data()), + index.search_batch(queries.data(), 1, kTopK, 0, actual.data()), std::invalid_argument ); EXPECT_THROW( - index.search_batch(queries.data(), 1, k, 1, nullptr), std::invalid_argument + index.search_batch(queries.data(), 1, kTopK, 1, nullptr), + std::invalid_argument ); - std::vector removed(kNum); + std::vector removed(kCount); std::iota(removed.begin(), removed.end(), PID{0}); - index.remove(removed.data(), kNum); + index.remove(removed.data(), kCount); index.search_batch( queries.data(), - kQueries, - k, + kNumQueries, + kTopK, 4, actual.data(), actual_distances.data(), @@ -211,26 +217,26 @@ TEST(IvfConfigurationTest, RejectsUnsupportedCountsAndNullConstructionInputs) { } TEST(IvfSearchTest, RawRerankingUsesOriginalCoordinates) { - constexpr size_t kNum = 33; + constexpr size_t kCount = 33; constexpr size_t kDim = 65; - std::vector data(kNum * kDim); + std::vector data(kCount * kDim); std::vector centroid(kDim, 0); std::vector query(kDim, 0.25F); - std::vector clusters(kNum, 0); + std::vector clusters(kCount, 0); for (size_t i = 0; i < data.size(); ++i) { data[i] = static_cast(i % 71) / 71.0F; } for (auto metric : {METRIC_L2, METRIC_IP}) { for (auto rotator : {RotatorType::FhtKacRotator, RotatorType::MatrixRotator}) { - IVF index(kNum, kDim, 1, 32, metric, rotator); + IVF index(kCount, kDim, 1, 32, metric, rotator); EXPECT_EQ(index.nbits(), 32); index.construct(data.data(), centroid.data(), clusters.data(), false, 1); for (bool hacc : {false, true}) { - std::vector ids(kNum); - std::vector distances(kNum); - index.search(query.data(), kNum, 1, ids.data(), distances.data(), hacc); - for (size_t i = 0; i < kNum; ++i) { - ASSERT_LT(ids[i], kNum); + std::vector ids(kCount); + std::vector distances(kCount); + index.search(query.data(), kCount, 1, ids.data(), distances.data(), hacc); + for (size_t i = 0; i < kCount; ++i) { + ASSERT_LT(ids[i], kCount); float expected = metric == METRIC_L2 ? 0.0F : 1.0F; for (size_t j = 0; j < kDim; ++j) { float value = data[ids[i] * kDim + j]; @@ -240,8 +246,8 @@ TEST(IvfSearchTest, RawRerankingUsesOriginalCoordinates) { EXPECT_NEAR(distances[i], expected, 2e-5F); } EXPECT_TRUE(std::is_sorted(distances.begin(), distances.end())); - std::vector ids_only(kNum); - index.search(query.data(), kNum, 1, ids_only.data(), hacc); + std::vector ids_only(kCount); + index.search(query.data(), kCount, 1, ids_only.data(), hacc); EXPECT_EQ(ids, ids_only); } } @@ -249,12 +255,12 @@ TEST(IvfSearchTest, RawRerankingUsesOriginalCoordinates) { } TEST(IvfSearchTest, AutomaticHighAccuracyMatchesBitPolicy) { - constexpr size_t kNum = 33; + constexpr size_t kCount = 33; constexpr size_t kDim = 65; - std::vector data(kNum * kDim); + std::vector data(kCount * kDim); std::vector centroid(kDim, 0); std::vector query(kDim); - std::vector clusters(kNum, 0); + std::vector clusters(kCount, 0); for (size_t i = 0; i < data.size(); ++i) { data[i] = static_cast(i % 71) / 71.0F; } @@ -262,7 +268,7 @@ TEST(IvfSearchTest, AutomaticHighAccuracyMatchesBitPolicy) { query[i] = static_cast(i % 13) / 13.0F; } for (size_t bits : {1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL, 32UL}) { - IVF index(kNum, kDim, 1, bits); + IVF index(kCount, kDim, 1, bits); index.construct(data.data(), centroid.data(), clusters.data(), false, 1); std::array ids{}, expected_ids{}, ids_only{}; std::array distances{}, expected_distances{}; @@ -311,16 +317,16 @@ TEST(IvfSearchTest, RejectsUnbuiltStateAndInvalidArguments) { } TEST(IvfSearchTest, FillsSentinelsWhenProbedClustersCannotFillK) { - constexpr size_t kNum = 2; + constexpr size_t kCount = 2; constexpr size_t kDim = 64; - std::array data{}; - std::array centroids{}; + std::array data{}; + std::array centroids{}; std::array query{}; - std::array clusters{1, 1}; + std::array clusters{1, 1}; std::fill(data.begin(), data.end(), 1.0F); std::fill(centroids.begin() + static_cast(kDim), centroids.end(), 1.0F); - IVF index(kNum, kDim, kNum, 32, MetricType::METRIC_L2, RotatorType::MatrixRotator); + IVF index(kCount, kDim, kCount, 32, MetricType::METRIC_L2, RotatorType::MatrixRotator); index.construct(data.data(), centroids.data(), clusters.data(), false, 1); PID result = kPidMax; @@ -331,17 +337,17 @@ TEST(IvfSearchTest, FillsSentinelsWhenProbedClustersCannotFillK) { } TEST(IvfSearchTest, RoutesInnerProductQueriesByInnerProduct) { - constexpr size_t kNum = 2; + constexpr size_t kCount = 2; constexpr size_t kDim = 64; - std::vector data(kNum * kDim, 0.0F); - std::vector centroids(kNum * kDim, 0.0F); - std::array clusters{0, 1}; + std::vector data(kCount * kDim, 0.0F); + std::vector centroids(kCount * kDim, 0.0F); + std::array clusters{0, 1}; std::array query{}; data[0] = centroids[0] = 0.9F; data[kDim] = centroids[kDim] = 100.0F; query[0] = 1.0F; - IVF index(kNum, kDim, kNum, 32, METRIC_IP); + IVF index(kCount, kDim, kCount, 32, METRIC_IP); index.construct(data.data(), centroids.data(), clusters.data(), false, 1); PID result = kPidMax; @@ -376,18 +382,18 @@ TEST(IvfSearchTest, RoutesInnerProductQueriesWithHNSWCentroids) { } TEST(IvfSearchTest, InnerProductRawRerankingUsesResidualNormForPruning) { - constexpr size_t kNum = 65; + constexpr size_t kCount = 65; constexpr size_t kDim = 64; constexpr size_t kTopK = 1; - std::vector data(kNum * kDim); + std::vector data(kCount * kDim); std::vector centroids(2 * kDim, 0.0F); - std::vector clusters(kNum, 1); + std::vector clusters(kCount, 1); std::array query{}; query[0] = 1.0F; centroids[0] = 100.0F; centroids[kDim] = 50.0F; - for (size_t id = 0; id < kNum; ++id) { + for (size_t id = 0; id < kCount; ++id) { clusters[id] = id < 32 ? 0 : 1; data[id * kDim] = id < 32 ? 10.0F + static_cast(id) * 0.01F : 0.0F; for (size_t d = 1; d < kDim; ++d) { @@ -397,7 +403,7 @@ TEST(IvfSearchTest, InnerProductRawRerankingUsesResidualNormForPruning) { } data[32 * kDim] = 1000.0F; - IVF index(kNum, kDim, 2, 32, METRIC_IP); + IVF index(kCount, kDim, 2, 32, METRIC_IP); index.construct(data.data(), centroids.data(), clusters.data(), false, 1); for (bool use_hacc : {false, true}) { std::array ids{}; @@ -409,30 +415,30 @@ TEST(IvfSearchTest, InnerProductRawRerankingUsesResidualNormForPruning) { } TEST(IvfPersistenceTest, ReloadsRawAndQuantizedStorage) { - constexpr size_t kNum = 33; + constexpr size_t kCount = 33; constexpr size_t kDim = 65; - std::vector data(kNum * kDim, 1.0F); + std::vector data(kCount * kDim, 1.0F); std::vector centroid(kDim, 0.0F); - std::vector clusters(kNum, 0); + std::vector clusters(kCount, 0); const rabitq_test::TempFile path_file("rabitq_ivf_storage.index"); const std::string& path = path_file.path(); IVF loaded; for (auto rotator : {RotatorType::FhtKacRotator, RotatorType::MatrixRotator}) { for (size_t bits : {32UL, 1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL}) { - IVF index(kNum, kDim, 1, bits, METRIC_L2, rotator); + IVF index(kCount, kDim, 1, bits, METRIC_L2, rotator); index.construct(data.data(), centroid.data(), clusters.data(), false, 1); index.save(path.c_str()); loaded.load(path.c_str()); EXPECT_EQ(loaded.nbits(), bits); EXPECT_EQ(loaded.padded_dim(), index.padded_dim()); EXPECT_EQ(loaded.rotator_type(), rotator); - std::vector ids(kNum); - std::vector loaded_ids(kNum); - std::vector distances(kNum); - std::vector loaded_distances(kNum); - index.search(centroid.data(), kNum, 1, ids.data(), distances.data(), true); + std::vector ids(kCount); + std::vector loaded_ids(kCount); + std::vector distances(kCount); + std::vector loaded_distances(kCount); + index.search(centroid.data(), kCount, 1, ids.data(), distances.data(), true); loaded.search( - centroid.data(), kNum, 1, loaded_ids.data(), loaded_distances.data(), true + centroid.data(), kCount, 1, loaded_ids.data(), loaded_distances.data(), true ); EXPECT_EQ(ids, loaded_ids); EXPECT_EQ(distances, loaded_distances); @@ -441,37 +447,37 @@ TEST(IvfPersistenceTest, ReloadsRawAndQuantizedStorage) { } TEST(IvfPersistenceTest, FailedLoadPreservesExistingIndex) { - constexpr size_t kTargetNum = 2; + constexpr size_t kTargetCount = 2; constexpr size_t kTargetDim = 64; - std::vector target_data(kTargetNum * kTargetDim, 1.0F); + std::vector target_data(kTargetCount * kTargetDim, 1.0F); std::vector target_centroid(kTargetDim, 0.0F); - std::array target_clusters{0, 0}; + std::array target_clusters{0, 0}; std::fill( target_data.begin() + static_cast(kTargetDim), target_data.end(), 2.0F ); const rabitq_test::TempFile path_file("rabitq_ivf_truncated.index"); const std::string& path = path_file.path(); - IVF target(kTargetNum, kTargetDim, 1, 32, METRIC_IP, RotatorType::MatrixRotator); + IVF target(kTargetCount, kTargetDim, 1, 32, METRIC_IP, RotatorType::MatrixRotator); target.construct( target_data.data(), target_centroid.data(), target_clusters.data(), false, 1 ); - std::array expected_ids{}; - std::array expected_distances{}; + std::array expected_ids{}; + std::array expected_distances{}; target.search( target_centroid.data(), - kTargetNum, + kTargetCount, 1, expected_ids.data(), expected_distances.data() ); - constexpr size_t kSourceNum = 33; + constexpr size_t kSourceCount = 33; constexpr size_t kSourceDim = 65; - std::vector source_data(kSourceNum * kSourceDim, 0.25F); + std::vector source_data(kSourceCount * kSourceDim, 0.25F); std::vector source_centroid(kSourceDim, 0.5F); - std::vector source_clusters(kSourceNum, 0); - IVF source(kSourceNum, kSourceDim, 1, 4, METRIC_L2, RotatorType::FhtKacRotator); + std::vector source_clusters(kSourceCount, 0); + IVF source(kSourceCount, kSourceDim, 1, 4, METRIC_L2, RotatorType::FhtKacRotator); source.construct( source_data.data(), source_centroid.data(), source_clusters.data(), false, 1 ); @@ -488,32 +494,32 @@ TEST(IvfPersistenceTest, FailedLoadPreservesExistingIndex) { output.close(); EXPECT_ANY_THROW(target.load(path.c_str())); - EXPECT_EQ(target.max_elements(), kTargetNum); + EXPECT_EQ(target.max_elements(), kTargetCount); EXPECT_EQ(target.dimension(), kTargetDim); EXPECT_EQ(target.num_clusters(), 1U); EXPECT_EQ(target.nbits(), 32U); EXPECT_EQ(target.metric_type(), METRIC_IP); EXPECT_EQ(target.rotator_type(), RotatorType::MatrixRotator); - std::array actual_ids{}; - std::array actual_distances{}; + std::array actual_ids{}; + std::array actual_distances{}; target.search( - target_centroid.data(), kTargetNum, 1, actual_ids.data(), actual_distances.data() + target_centroid.data(), kTargetCount, 1, actual_ids.data(), actual_distances.data() ); EXPECT_EQ(actual_ids, expected_ids); EXPECT_EQ(actual_distances, expected_distances); } TEST(IvfPersistenceTest, RejectsPointIdsUsingTheSearchBufferMarker) { - constexpr size_t kNum = 33; + constexpr size_t kCount = 33; constexpr size_t kDim = 64; - std::vector data(kNum * kDim, 0.25F); + std::vector data(kCount * kDim, 0.25F); std::vector centroid(kDim, 0.0F); - std::vector clusters(kNum, 0); + std::vector clusters(kCount, 0); const rabitq_test::TempFile path_file("rabitq_ivf_bad_id.index"); const std::string& path = path_file.path(); - IVF source(kNum, kDim, 1, 4); + IVF source(kCount, kDim, 1, 4); source.construct(data.data(), centroid.data(), clusters.data(), false, 1); source.save(path.c_str()); @@ -531,7 +537,7 @@ TEST(IvfPersistenceTest, RejectsPointIdsUsingTheSearchBufferMarker) { // ── add / remove ──────────────────────────────────────────────────────────────────────── constexpr size_t kDynDim = 65; // Exercise padding and partially occupied FastScan batches. -constexpr size_t kDynClusters = 4; +constexpr size_t kDynNumClusters = 4; std::vector random_rows(size_t rows, uint32_t seed) { std::mt19937 gen(seed); @@ -550,7 +556,7 @@ std::vector nearest_centroids( std::vector assigned(count); for (size_t i = 0; i < count; ++i) { double best = std::numeric_limits::infinity(); - for (size_t c = 0; c < kDynClusters; ++c) { + for (size_t c = 0; c < kDynNumClusters; ++c) { double dist = 0; for (size_t d = 0; d < kDynDim; ++d) { const double x = rows[i * kDynDim + d]; @@ -599,9 +605,11 @@ Hits search_everything( const size_t k = index.max_elements(); Hits hits{std::vector(k), std::vector(k)}; if (hacc.has_value()) { - index.search(query, k, kDynClusters, hits.ids.data(), hits.distances.data(), *hacc); + index.search( + query, k, kDynNumClusters, hits.ids.data(), hits.distances.data(), *hacc + ); } else { - index.search(query, k, kDynClusters, hits.ids.data(), hits.distances.data()); + index.search(query, k, kDynNumClusters, hits.ids.data(), hits.distances.data()); } return hits; } @@ -692,7 +700,7 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { constexpr size_t kOld = 200; constexpr size_t kNew = 200; const auto points = random_rows(kOld + kNew, 11); - const auto centroids = random_rows(kDynClusters, 12); + const auto centroids = random_rows(kDynNumClusters, 12); const float* extra = points.data() + (kOld * kDynDim); const rabitq_test::TempFile base_path_file("rabitq_ivf_add_base.index"); const std::string& base_path = base_path_file.path(); @@ -707,7 +715,7 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits ); - IVF base(kOld, kDynDim, kDynClusters, bits, metric); + IVF base(kOld, kDynDim, kDynNumClusters, bits, metric); base.construct(points.data(), centroids.data(), assigned.data(), false, 1); base.save(base_path.c_str()); @@ -739,11 +747,11 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { // Even clusters are reserved and grow in place; odd ones are rebuilt each time. IVF reserved; reserved.load(base_path.c_str()); - std::vector final_sizes(kDynClusters, 0); + std::vector final_sizes(kDynNumClusters, 0); for (PID cluster : assigned) { ++final_sizes[cluster]; } - for (size_t i = 0; i < kDynClusters; i += 2) { + for (size_t i = 0; i < kDynNumClusters; i += 2) { reserved.reserve(static_cast(i), final_sizes[i]); } const Cluster reserved_view = IvfTestAccess::cluster_view(reserved, 0); @@ -772,17 +780,17 @@ TEST(IvfAddTest, AddingToOneClusterLeavesOtherClustersUntouched) { // Cluster sizes 0, 1, 32, 97 cover empty, single, full-batch and tail clusters. constexpr size_t kOld = 130; const auto points = random_rows(kOld, 741); - const auto centroids = random_rows(kDynClusters, 742); + const auto centroids = random_rows(kDynNumClusters, 742); std::vector assigned(kOld, 3); assigned[0] = 1; std::fill(assigned.begin() + 1, assigned.begin() + 33, PID{2}); - IVF index(kOld, kDynDim, kDynClusters, 4, METRIC_L2); + IVF index(kOld, kDynDim, kDynNumClusters, 4, METRIC_L2); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); std::vector before; - before.reserve(kDynClusters); - for (size_t i = 0; i < kDynClusters; ++i) { + before.reserve(kDynNumClusters); + for (size_t i = 0; i < kDynNumClusters; ++i) { before.push_back(IvfTestAccess::cluster_view(index, i)); } @@ -792,7 +800,7 @@ TEST(IvfAddTest, AddingToOneClusterLeavesOtherClustersUntouched) { const std::vector new_cluster_ids(kNew, 0); index.add(extra.data(), kNew, new_cluster_ids.data(), false, 1); - for (size_t i = 0; i < kDynClusters; ++i) { + for (size_t i = 0; i < kDynNumClusters; ++i) { SCOPED_TRACE(::testing::Message() << "cluster=" << i); const Cluster after = IvfTestAccess::cluster_view(index, i); if (i == 0) { @@ -808,15 +816,15 @@ TEST(IvfAddTest, AddingToOneClusterLeavesOtherClustersUntouched) { } TEST(IvfAddTest, SparseInterleavedClustersPreserveRowOrderAndMatchRebuild) { - constexpr size_t kOld = 3, kCount = 9, kClusters = 257; + constexpr size_t kOld = 3, kCount = 9, kNumClusters = 257; const auto points = random_rows(kCount, 744); - const auto centroids = random_rows(kClusters, 745); + const auto centroids = random_rows(kNumClusters, 745); const std::array assigned{0, 128, 256, 256, 0, 256, 128, 0, 256}; const rabitq_test::TempFile path_file("rabitq_ivf_sparse_add.index"); const std::string& path = path_file.path(); for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {1U, 4U, 32U}) { - IVF index(kOld, kDynDim, kClusters, bits, metric); + IVF index(kOld, kDynDim, kNumClusters, bits, metric); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); index.reserve(256, 32); const Cluster reserved = IvfTestAccess::cluster_view(index, 256); @@ -828,7 +836,7 @@ TEST(IvfAddTest, SparseInterleavedClustersPreserveRowOrderAndMatchRebuild) { 4 ); EXPECT_EQ(IvfTestAccess::cluster_view(index, 256).ids(), reserved.ids()); - for (size_t i = 0; i < kClusters; ++i) { + for (size_t i = 0; i < kNumClusters; ++i) { std::vector expected; for (size_t row = 0; row < kCount; ++row) { if (assigned[row] == i) { @@ -853,12 +861,12 @@ TEST(IvfAddTest, SparseInterleavedClustersPreserveRowOrderAndMatchRebuild) { TEST(IvfAddTest, UnreservedClustersGrowToExactlyTheirNewSize) { constexpr size_t kOld = 40; const auto points = random_rows(kOld + 1, 851); - const auto centroids = random_rows(kDynClusters, 852); + const auto centroids = random_rows(kDynNumClusters, 852); const std::vector assigned(kOld, 0); for (size_t bits : {1UL, 4UL, 32UL}) { SCOPED_TRACE(::testing::Message() << "bits=" << bits); - IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + IVF index(kOld, kDynDim, kDynNumClusters, bits, METRIC_L2); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); const PID route_to_zero = 0; index.add(points.data() + (kOld * kDynDim), 1, &route_to_zero, false, 1); @@ -871,13 +879,13 @@ TEST(IvfReserveTest, ReservedClusterGrowsInPlaceUntilItsCapacityIsUsed) { constexpr size_t kOld = 40; constexpr size_t kReserved = kOld + 2; const auto points = random_rows(kReserved + 1, 861); - const auto centroids = random_rows(kDynClusters, 862); + const auto centroids = random_rows(kDynNumClusters, 862); const std::vector assigned(kOld, 0); const PID route_to_zero = 0; for (size_t bits : {1UL, 4UL, 32UL}) { SCOPED_TRACE(::testing::Message() << "bits=" << bits); - IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + IVF index(kOld, kDynDim, kDynNumClusters, bits, METRIC_L2); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); index.reserve(0, kReserved); const Cluster reserved = IvfTestAccess::cluster_view(index, 0); @@ -904,7 +912,7 @@ TEST(IvfReserveTest, ReservedClusterGrowsInPlaceUntilItsCapacityIsUsed) { TEST(IvfReserveTest, RejectsBadInputAndLeavesSavedBytesUnchanged) { constexpr size_t kOld = 40; const auto points = random_rows(kOld, 871); - const auto centroids = random_rows(kDynClusters, 872); + const auto centroids = random_rows(kDynNumClusters, 872); const std::vector assigned(kOld, 0); const rabitq_test::TempFile path_file("rabitq_ivf_reserve.index"); const std::string& path = path_file.path(); @@ -915,22 +923,24 @@ TEST(IvfReserveTest, RejectsBadInputAndLeavesSavedBytesUnchanged) { } catch (const std::exception& error) { EXPECT_STREQ(error.what(), message); } }; - IVF unbuilt(kOld, kDynDim, kDynClusters, 4); + IVF unbuilt(kOld, kDynDim, kDynNumClusters, 4); EXPECT_THROW(unbuilt.reserve(0, 1), std::logic_error); expect_message( [&] { unbuilt.reserve(0, 1); }, "IVF index must be constructed or loaded before reserve" ); - IVF index(kOld, kDynDim, kDynClusters, 4); + IVF index(kOld, kDynDim, kDynNumClusters, 4); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); const std::string expected = saved_bytes(index, path); const Cluster constructed = IvfTestAccess::cluster_view(index, 0); - EXPECT_THROW(index.reserve(static_cast(kDynClusters), 1), std::invalid_argument); + EXPECT_THROW( + index.reserve(static_cast(kDynNumClusters), 1), std::invalid_argument + ); expect_message( - [&] { index.reserve(static_cast(kDynClusters), 1); }, - "Cluster ID is out of range" + [&] { index.reserve(static_cast(kDynNumClusters), 1); }, + "IVF cluster ID is out of range" ); EXPECT_THROW( index.reserve(0, buffer::kSearchBufferMaxPointCount + 1), std::invalid_argument @@ -952,7 +962,7 @@ TEST(IvfReserveTest, RejectsBadInputAndLeavesSavedBytesUnchanged) { EXPECT_TRUE(saved_bytes(index, path) == expected); IVF loaded; loaded.load(path.c_str()); - for (size_t i = 0; i < kDynClusters; ++i) { + for (size_t i = 0; i < kDynNumClusters; ++i) { EXPECT_FALSE(IvfTestAccess::has_spare_capacity(loaded, i)); } } @@ -963,11 +973,11 @@ class FailingRotator : public Rotator { FailingRotator(std::unique_ptr> inner, size_t fail_at) : Rotator(0, inner->size()), inner_(std::move(inner)), fail_at_(fail_at) {} - void rotate(const float* src, float* dst) const override { + void rotate(const float* data, float* rotated_data) const override { if (calls_++ == fail_at_) { throw std::runtime_error("injected rotation failure"); } - inner_->rotate(src, dst); + inner_->rotate(data, rotated_data); } void load(std::ifstream& input) override { inner_->load(input); } void save(std::ofstream& output) const override { inner_->save(output); } @@ -988,7 +998,7 @@ TEST(IvfAddTest, FailedAddLeavesClustersGrowingInPlaceUnchanged) { constexpr size_t kOld = 70; constexpr size_t kInPlace = 3; const auto points = random_rows(kOld + 6, 881); - const auto centroids = random_rows(kDynClusters, 882); + const auto centroids = random_rows(kDynNumClusters, 882); std::vector assigned(kOld, 1); std::fill(assigned.begin(), assigned.begin() + 40, PID{0}); const std::array new_clusters{0, 0, 0, 1, 1, 1}; @@ -1000,7 +1010,7 @@ TEST(IvfAddTest, FailedAddLeavesClustersGrowingInPlaceUnchanged) { for (size_t bits : {1UL, 4UL, 32UL}) { SCOPED_TRACE(::testing::Message() << "bits=" << bits); - IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + IVF index(kOld, kDynDim, kDynNumClusters, bits, METRIC_L2); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); index.save(base_path.c_str()); index.reserve(0, 64); @@ -1033,7 +1043,7 @@ TEST(IvfConstructionTest, RebuildMatchesIncrementalEncodingWithSkewedClustersAnd // Cluster sizes 0, 1, 32, 97 cover empty rows and multiple physical tail batches. constexpr size_t kCount = 130; const auto points = random_rows(kCount, 731); - const auto centroids = random_rows(kDynClusters, 732); + const auto centroids = random_rows(kDynNumClusters, 732); std::vector assigned(kCount, 3); assigned[0] = 1; std::fill(assigned.begin() + 1, assigned.begin() + 33, PID{2}); @@ -1045,7 +1055,7 @@ TEST(IvfConstructionTest, RebuildMatchesIncrementalEncodingWithSkewedClustersAnd SCOPED_TRACE( ::testing::Message() << "bits=" << bits << " metric=" << metric ); - IVF incremental(1, kDynDim, kDynClusters, bits, metric, rotation); + IVF incremental(1, kDynDim, kDynNumClusters, bits, metric, rotation); incremental.construct( points.data(), centroids.data(), assigned.data(), false, 1 ); @@ -1074,10 +1084,10 @@ TEST(IvfConstructionTest, RebuildMatchesIncrementalEncodingWithSkewedClustersAnd TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { constexpr size_t kOld = 150; constexpr size_t kNew = 150; - constexpr size_t kQueries = 8; + constexpr size_t kNumQueries = 8; const auto points = random_rows(kOld + kNew, 21); - const auto centroids = random_rows(kDynClusters, 22); - const auto queries = random_rows(kQueries, 23); + const auto centroids = random_rows(kDynNumClusters, 22); + const auto queries = random_rows(kNumQueries, 23); for (auto metric : {METRIC_L2, METRIC_IP}) { const auto assigned = @@ -1087,10 +1097,10 @@ TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits ); - IVF index(kOld, kDynDim, kDynClusters, bits, metric); + IVF index(kOld, kDynDim, kDynNumClusters, bits, metric); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); std::vector before; - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { before.push_back(search_everything(index, queries.data() + (q * kDynDim))); } @@ -1098,7 +1108,7 @@ TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { double old_error = 0; double new_error = 0; - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { const float* query = queries.data() + (q * kDynDim); const auto after = search_everything(index, query); std::vector by_id(kOld + kNew, std::nanf("")); @@ -1118,8 +1128,8 @@ TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { (id < kOld ? old_error : new_error) += error; } } - old_error /= kQueries * kOld; - new_error /= kQueries * kNew; + old_error /= kNumQueries * kOld; + new_error /= kNumQueries * kNew; // Mis-encoded points would be off by the spread of the data, not by a factor. EXPECT_LT(new_error, 3 * old_error + 1e-4) << old_error << " vs " << new_error; } @@ -1130,10 +1140,10 @@ TEST(IvfAddTest, AddedRawVectorsAreRerankedExactly) { constexpr size_t kOld = 40; constexpr size_t kNew = 60; const auto points = random_rows(kOld + kNew, 31); - const auto centroids = random_rows(kDynClusters, 32); + const auto centroids = random_rows(kDynNumClusters, 32); const auto queries = random_rows(3, 33); for (auto metric : {METRIC_L2, METRIC_IP}) { - IVF index(kOld, kDynDim, kDynClusters, 32, metric); + IVF index(kOld, kDynDim, kDynNumClusters, 32, metric); const auto assigned = nearest_centroids(centroids, points.data(), kOld, metric); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); index.add(points.data() + (kOld * kDynDim), kNew, nullptr, false, 1); @@ -1156,7 +1166,7 @@ TEST(IvfAddTest, FasterQuantizationIsSupportedAndDeterministic) { constexpr size_t kOld = 100; constexpr size_t kNew = 50; const auto points = random_rows(kOld + kNew, 41); - const auto centroids = random_rows(kDynClusters, 42); + const auto centroids = random_rows(kDynNumClusters, 42); const auto assigned = nearest_centroids(centroids, points.data(), kOld + kNew, METRIC_L2); const rabitq_test::TempFile base_path_file("rabitq_ivf_add_faster_base.index"); @@ -1164,7 +1174,7 @@ TEST(IvfAddTest, FasterQuantizationIsSupportedAndDeterministic) { const rabitq_test::TempFile path_file("rabitq_ivf_add_faster.index"); const std::string& path = path_file.path(); - IVF base(kOld, kDynDim, kDynClusters, 5); + IVF base(kOld, kDynDim, kDynNumClusters, 5); base.construct(points.data(), centroids.data(), assigned.data(), true, 1); base.save(base_path.c_str()); IVF first; @@ -1213,23 +1223,23 @@ TEST(IvfAddTest, RoutesNewPointsWithHNSWCentroids) { } TEST(IvfAddTest, RejectsBadInputAndLeavesTheIndexUnchanged) { - constexpr size_t kNum = 70; - const auto points = random_rows(kNum + 1, 51); - const auto centroids = random_rows(kDynClusters, 52); - const auto assigned = nearest_centroids(centroids, points.data(), kNum, METRIC_L2); + constexpr size_t kCount = 70; + const auto points = random_rows(kCount + 1, 51); + const auto centroids = random_rows(kDynNumClusters, 52); + const auto assigned = nearest_centroids(centroids, points.data(), kCount, METRIC_L2); const rabitq_test::TempFile path_file("rabitq_ivf_add_reject.index"); const std::string& path = path_file.path(); - IVF unbuilt(kNum, kDynDim, kDynClusters, 4); + IVF unbuilt(kCount, kDynDim, kDynNumClusters, 4); EXPECT_THROW(unbuilt.add(points.data(), 1), std::logic_error); IVF unconfigured; EXPECT_THROW(unconfigured.add(points.data(), 1), std::logic_error); - IVF index(kNum, kDynDim, kDynClusters, 4); + IVF index(kCount, kDynDim, kDynNumClusters, 4); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); const std::string expected = saved_bytes(index, path); - const std::array bad_clusters{0, static_cast(kDynClusters)}; + const std::array bad_clusters{0, static_cast(kDynNumClusters)}; EXPECT_THROW(index.add(points.data(), 2, bad_clusters.data()), std::invalid_argument); EXPECT_THROW(index.add(nullptr, 1), std::invalid_argument); // The ID check happens before the data is read. @@ -1238,31 +1248,31 @@ TEST(IvfAddTest, RejectsBadInputAndLeavesTheIndexUnchanged) { ); EXPECT_NO_THROW(index.add(nullptr, 0)); EXPECT_NO_THROW(index.add(points.data(), 0)); - EXPECT_EQ(index.max_elements(), kNum); + EXPECT_EQ(index.max_elements(), kCount); EXPECT_TRUE(saved_bytes(index, path) == expected); } TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { - constexpr size_t kNum = 200; - constexpr size_t kQueries = 5; - const auto points = random_rows(kNum, 61); - const auto centroids = random_rows(kDynClusters, 62); - const auto queries = random_rows(kQueries, 63); + constexpr size_t kCount = 200; + constexpr size_t kNumQueries = 5; + const auto points = random_rows(kCount, 61); + const auto centroids = random_rows(kDynNumClusters, 62); + const auto queries = random_rows(kNumQueries, 63); std::vector removed_ids; - std::vector is_removed(kNum, false); - for (PID id = 0; id < kNum; id += 3) { + std::vector is_removed(kCount, false); + for (PID id = 0; id < kCount; id += 3) { removed_ids.push_back(id); is_removed[id] = true; } - removed_ids.push_back(kNum - 1); // the last point sits in a partial batch - is_removed[kNum - 1] = true; + removed_ids.push_back(kCount - 1); // the last point sits in a partial batch + is_removed[kCount - 1] = true; removed_ids.push_back(0); // repeated IDs count once const size_t unique_removed = static_cast(std::count(is_removed.begin(), is_removed.end(), true)); for (auto metric : {METRIC_L2, METRIC_IP}) { - const auto assigned = nearest_centroids(centroids, points.data(), kNum, metric); + const auto assigned = nearest_centroids(centroids, points.data(), kCount, metric); // 1-3 bits use standard FastScan, 4-9 bits and raw storage exercise other paths too for (size_t bits : {1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL, 32UL}) { for (bool hacc : {false, true}) { @@ -1270,10 +1280,10 @@ TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits << " hacc=" << hacc ); - IVF index(kNum, kDynDim, kDynClusters, bits, metric); + IVF index(kCount, kDynDim, kDynNumClusters, bits, metric); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); std::vector before; - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { before.push_back( search_everything(index, queries.data() + (q * kDynDim), hacc) ); @@ -1283,14 +1293,14 @@ TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { index.remove(removed_ids.data(), removed_ids.size()), unique_removed ); EXPECT_EQ(index.remove(removed_ids.data(), removed_ids.size()), 0U); - EXPECT_EQ(index.max_elements(), kNum); + EXPECT_EQ(index.max_elements(), kCount); - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { const auto after = search_everything(index, queries.data() + (q * kDynDim), hacc); // Survivors keep their order and distances; the rest are sentinels. size_t next = 0; - for (size_t r = 0; r < kNum; ++r) { + for (size_t r = 0; r < kCount; ++r) { if (is_removed[before[q].ids[r]]) { continue; } @@ -1298,8 +1308,8 @@ TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { ASSERT_EQ(after.distances[next], before[q].distances[r]); ++next; } - ASSERT_EQ(next, kNum - unique_removed); - for (size_t r = next; r < kNum; ++r) { + ASSERT_EQ(next, kCount - unique_removed); + for (size_t r = next; r < kCount; ++r) { EXPECT_EQ(after.ids[r], kPidMax); EXPECT_EQ( after.distances[r], std::numeric_limits::infinity() @@ -1312,13 +1322,13 @@ TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { index.search( queries.data() + (q * kDynDim), top.size(), - kDynClusters, + kDynNumClusters, top.data(), top_distances.data(), hacc ); for (size_t r = 0; r < top.size(); ++r) { - ASSERT_LT(top[r], kNum); + ASSERT_LT(top[r], kCount); EXPECT_FALSE(is_removed[top[r]]); EXPECT_TRUE(std::isfinite(top_distances[r])); } @@ -1331,10 +1341,10 @@ TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { constexpr size_t kOld = 100; constexpr size_t kNew = 60; - constexpr size_t kQueries = 4; + constexpr size_t kNumQueries = 4; const auto points = random_rows(kOld + kNew, 71); - const auto centroids = random_rows(kDynClusters, 72); - const auto queries = random_rows(kQueries, 73); + const auto centroids = random_rows(kDynNumClusters, 72); + const auto queries = random_rows(kNumQueries, 73); const std::vector removed{1, 33, 34, 98, 99}; const rabitq_test::TempFile path_file("rabitq_ivf_remove.index"); const std::string& path = path_file.path(); @@ -1363,7 +1373,7 @@ TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits << " hacc=" << hacc ); - IVF index(kOld, kDynDim, kDynClusters, bits, metric); + IVF index(kOld, kDynDim, kDynNumClusters, bits, metric); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); ASSERT_EQ(index.remove(removed.data(), removed.size()), removed.size()); index.save(path.c_str()); @@ -1371,7 +1381,7 @@ TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { IVF loaded; loaded.load(path.c_str()); ASSERT_EQ(loaded.max_elements(), kOld); - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { const float* query = queries.data() + (q * kDynDim); const auto expected = search_everything(index, query, hacc); const auto actual = search_everything(loaded, query, hacc); @@ -1387,7 +1397,7 @@ TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { index.add(added, kNew, nullptr, false, 1); loaded.add(added, kNew, nullptr, false, 1); ASSERT_EQ(loaded.max_elements(), kOld + kNew); - for (size_t q = 0; q < kQueries; ++q) { + for (size_t q = 0; q < kNumQueries; ++q) { const float* query = queries.data() + (q * kDynDim); const auto expected = search_everything(index, query, hacc); const auto actual = search_everything(loaded, query, hacc); @@ -1401,35 +1411,35 @@ TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { } TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { - constexpr size_t kNum = 70; - const auto points = random_rows(kNum, 81); - const auto centroids = random_rows(kDynClusters, 82); - const auto assigned = nearest_centroids(centroids, points.data(), kNum, METRIC_L2); + constexpr size_t kCount = 70; + const auto points = random_rows(kCount, 81); + const auto centroids = random_rows(kDynNumClusters, 82); + const auto assigned = nearest_centroids(centroids, points.data(), kCount, METRIC_L2); const rabitq_test::TempFile path_file("rabitq_ivf_remove_reject.index"); const std::string& path = path_file.path(); const PID id = 0; - IVF unbuilt(kNum, kDynDim, kDynClusters, 4); + IVF unbuilt(kCount, kDynDim, kDynNumClusters, 4); EXPECT_THROW(unbuilt.remove(&id, 1), std::logic_error); - IVF index(kNum, kDynDim, kDynClusters, 4); + IVF index(kCount, kDynDim, kDynNumClusters, 4); index.construct(points.data(), centroids.data(), assigned.data(), false, 1); const std::string expected = saved_bytes(index, path); - const std::array out_of_range{3, static_cast(kNum)}; + const std::array out_of_range{3, static_cast(kCount)}; EXPECT_THROW(index.remove(out_of_range.data(), 2), std::invalid_argument); EXPECT_THROW(index.remove(nullptr, 1), std::invalid_argument); EXPECT_EQ(index.remove(nullptr, 0), 0U); EXPECT_TRUE(saved_bytes(index, path) == expected); // Removing everything leaves a searchable index that returns only sentinels. - std::vector everything(kNum); - for (size_t i = 0; i < kNum; ++i) { + std::vector everything(kCount); + for (size_t i = 0; i < kCount; ++i) { everything[i] = static_cast(i); } - EXPECT_EQ(index.remove(everything.data(), kNum), kNum); + EXPECT_EQ(index.remove(everything.data(), kCount), kCount); const auto hits = search_everything(index, points.data()); - for (size_t r = 0; r < kNum; ++r) { + for (size_t r = 0; r < kCount; ++r) { EXPECT_EQ(hits.ids[r], kPidMax); } } @@ -1437,7 +1447,7 @@ TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { } // namespace } // namespace rabitqlib::ivf -TEST(IVFSidecarValidation, RejectsMismatchedVectorStorage) { +TEST(IvfSidecarValidationTest, RejectsMismatchedVectorStorage) { using namespace rabitqlib; for (const auto metric : {METRIC_L2, METRIC_IP}) { for (const size_t sidecar_dim : {64U, 128U}) { diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp index 1a82ad0b..a4f8d8b4 100644 --- a/tests/unit/rabitqlib/index/pipnn_test.cpp +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -20,31 +20,31 @@ namespace rabitqlib::symqg::detail { namespace { -TEST(PipnnTest, DenseDistancesMatchReference) { +TEST(QgPipnnTest, DenseDistancesMatchReference) { constexpr size_t kCount = 1024, kDim = 65; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.17F); } - pipnn_impl::ScratchPool scratch(kCount, kDim, 2); + ScratchPool scratch(kCount, kDim, 2); // Reuse each worker's slice with changing matrix shapes, including the // largest leaf, and check both distance conventions against scalar doubles. for (size_t worker : {0U, 1U}) { auto work = scratch.worker(worker); - for (size_t count : {33U, 1024U, 257U}) { - pipnn_impl::Bucket ids(count); + for (size_t num_points : {33U, 1024U, 257U}) { + Bucket ids(num_points); std::iota(ids.begin(), ids.end(), PID{0}); - pipnn_impl::gather(data.data(), kDim, ids, work); + gather(data.data(), kDim, ids, work); for (auto metric : {METRIC_L2, METRIC_IP}) { - pipnn_impl::pairwise(work, count, kDim, metric); - for (size_t i = 0; i < count; ++i) { + pairwise(work, num_points, kDim, metric); + for (size_t i = 0; i < num_points; ++i) { for (size_t j = 0; j <= i; ++j) { double expected = 0; for (size_t d = 0; d < kDim; ++d) { const double a = data[i * kDim + d], b = data[j * kDim + d]; expected += metric == METRIC_L2 ? (a - b) * (a - b) : -a * b; } - EXPECT_NEAR(work.distances[i * count + j], expected, 5e-5); + EXPECT_NEAR(work.distances[i * num_points + j], expected, 5e-5); } } } @@ -52,7 +52,7 @@ TEST(PipnnTest, DenseDistancesMatchReference) { } } -TEST(PipnnTest, LocalSeedRetainsNearestNeighborsAndValidIds) { +TEST(QgPipnnTest, LocalSeedRetainsNearestNeighborsAndValidIds) { constexpr size_t kCount = 97, kDim = 65; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -98,31 +98,33 @@ TEST(PipnnTest, LocalSeedRetainsNearestNeighborsAndValidIds) { } } -TEST(PipnnTest, LocalNeighborsMatchExactReferenceAndGlobalIdTies) { +TEST(QgPipnnTest, LocalNeighborsMatchExactReferenceAndGlobalIdTies) { constexpr size_t kDim = 65; std::mt19937 generator(719); std::uniform_int_distribution coordinate(-8, 8); - for (size_t count : {2U, 3U, 17U, 65U, 1024U}) { - std::vector data(count * kDim); - pipnn_impl::Bucket ids(count); + for (size_t num_points : {2U, 3U, 17U, 65U, 1024U}) { + std::vector data(num_points * kDim); + Bucket ids(num_points); std::iota(ids.begin(), ids.end(), PID{0}); std::shuffle(ids.begin(), ids.end(), generator); - pipnn_impl::ScratchPool scratch(count, kDim, 1); + ScratchPool scratch(num_points, kDim, 1); auto work = scratch.worker(0); for (bool ties : {false, true}) { std::generate(data.begin(), data.end(), [&] { return ties ? 1.0F : static_cast(coordinate(generator)); }); for (auto metric : {METRIC_L2, METRIC_IP}) { - SCOPED_TRACE(::testing::Message() << count << "/" << ties << "/" << metric); - pipnn_impl::gather(data.data(), kDim, ids, work); - pipnn_impl::pairwise(work, count, kDim, metric); - pipnn_impl::select_local_neighbors(ids, work); - ASSERT_EQ(work.candidates.size(), 2 * count); - for (size_t i = 0; i < count; ++i) { + SCOPED_TRACE( + ::testing::Message() << num_points << "/" << ties << "/" << metric + ); + gather(data.data(), kDim, ids, work); + pairwise(work, num_points, kDim, metric); + select_local_neighbors(ids, work); + ASSERT_EQ(work.candidates.size(), 2 * num_points); + for (size_t i = 0; i < num_points; ++i) { std::vector> reference; - reference.reserve(count - 1); - for (PID target = 0; target < count; ++target) { + reference.reserve(num_points - 1); + for (PID target = 0; target < num_points; ++target) { if (target == ids[i]) { continue; } @@ -138,7 +140,7 @@ TEST(PipnnTest, LocalNeighborsMatchExactReferenceAndGlobalIdTies) { for (size_t j = 0; j < std::min(size_t{2}, reference.size()); ++j) { EXPECT_EQ(work.candidates[2 * i + j], reference[j]); } - if (count == 2) { + if (num_points == 2) { EXPECT_EQ(work.candidates[2 * i + 1].second, kPidMax); } } @@ -147,14 +149,14 @@ TEST(PipnnTest, LocalNeighborsMatchExactReferenceAndGlobalIdTies) { } } -TEST(PipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { +TEST(QgPipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { constexpr size_t kCount = 1100, kDim = 7; std::vector data(kCount * kDim, 0.5F); - pipnn_impl::ScratchPool scratch(kCount, kDim, 2); - auto leaves = pipnn_impl::cluster(data.data(), kCount, kDim, METRIC_L2, 2, scratch); + ScratchPool scratch(kCount, kDim, 2); + auto leaves = cluster(data.data(), kCount, kDim, METRIC_L2, 2, scratch); std::vector memberships(kCount, 0); for (auto& leaf : leaves) { - EXPECT_LE(leaf.size(), pipnn_impl::kLeafSize); + EXPECT_LE(leaf.size(), kLeafSize); std::sort(leaf.begin(), leaf.end()); EXPECT_EQ(std::unique(leaf.begin(), leaf.end()), leaf.end()); for (PID id : leaf) { @@ -168,7 +170,7 @@ TEST(PipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { } } -TEST(PipnnTest, PreservesCallerThreadCount) { +TEST(QgPipnnTest, PreservesCallerThreadCount) { const int previous = omp_get_max_threads(); std::vector data(65 * 65, 0.5F); const size_t requested = previous == 1 ? 2 : 1; @@ -176,7 +178,7 @@ TEST(PipnnTest, PreservesCallerThreadCount) { EXPECT_EQ(omp_get_max_threads(), previous); } -TEST(PipnnTest, RejectsInvalidConfiguration) { +TEST(QgPipnnTest, RejectsInvalidConfiguration) { std::vector data(33 * 7, 0); EXPECT_THROW(build_initial_graph(nullptr, 33, 7, 32), std::invalid_argument); EXPECT_THROW(build_initial_graph(data.data(), 33, 0, 32), std::invalid_argument); @@ -184,16 +186,15 @@ TEST(PipnnTest, RejectsInvalidConfiguration) { EXPECT_THROW(build_initial_graph(data.data(), 32, 7, 32), std::invalid_argument); } -TEST(PipnnTest, RootTileTailAndRecursiveBucketsCoverDuplicateInputs) { +TEST(QgPipnnTest, RootTileTailAndRecursiveBucketsCoverDuplicateInputs) { constexpr size_t kCount = 4097, kDim = 7; std::vector data(kCount * kDim, 0.5F); for (size_t threads : {1U, 2U}) { - pipnn_impl::ScratchPool scratch(kCount, kDim, threads); - auto leaves = - pipnn_impl::cluster(data.data(), kCount, kDim, METRIC_L2, threads, scratch); + ScratchPool scratch(kCount, kDim, threads); + auto leaves = cluster(data.data(), kCount, kDim, METRIC_L2, threads, scratch); std::vector memberships(kCount, 0); for (auto& leaf : leaves) { - EXPECT_LE(leaf.size(), pipnn_impl::kLeafSize); + EXPECT_LE(leaf.size(), kLeafSize); std::sort(leaf.begin(), leaf.end()); EXPECT_EQ(std::unique(leaf.begin(), leaf.end()), leaf.end()); for (PID id : leaf) { @@ -208,28 +209,28 @@ TEST(PipnnTest, RootTileTailAndRecursiveBucketsCoverDuplicateInputs) { } } -TEST(PipnnTest, AdaptiveRootTilesPreservePartitionsAcrossThreadCounts) { +TEST(QgPipnnTest, AdaptiveRootTilesPreservePartitionsAcrossThreadCounts) { constexpr size_t kDim = 65; std::mt19937 generator(817); // Integer coordinates keep these dot products exact on every SIMD backend. std::uniform_int_distribution uniform(-8, 8); // Exercise partial tiles on either side of the original 2048-row boundary. - for (size_t count : {1025U, 2049U, 4097U}) { - std::vector data(count * kDim); + for (size_t num_points : {1025U, 2049U, 4097U}) { + std::vector data(num_points * kDim); std::generate(data.begin(), data.end(), [&] { return static_cast(uniform(generator)); }); for (auto metric : {METRIC_L2, METRIC_IP}) { - pipnn_impl::ScratchPool serial_scratch(count, kDim, 1); + ScratchPool serial_scratch(num_points, kDim, 1); const auto expected = - pipnn_impl::cluster(data.data(), count, kDim, metric, 1, serial_scratch); + cluster(data.data(), num_points, kDim, metric, 1, serial_scratch); for (size_t threads : {2U, 4U}) { SCOPED_TRACE( - ::testing::Message() << count << "/" << metric << "/" << threads + ::testing::Message() << num_points << "/" << metric << "/" << threads ); - pipnn_impl::ScratchPool scratch(count, kDim, threads); + ScratchPool scratch(num_points, kDim, threads); const auto actual = - pipnn_impl::cluster(data.data(), count, kDim, metric, threads, scratch); + cluster(data.data(), num_points, kDim, metric, threads, scratch); EXPECT_EQ(actual, expected); } } diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 6be998bb..2359a5ed 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -333,7 +333,7 @@ struct QGConstructionTestAccess { namespace { -TEST(QGConstructionTest, OnlyFinalReversePruningCompletesRowsWithoutWritingGraph) { +TEST(QgConstructionTest, OnlyFinalReversePruningCompletesRowsWithoutWritingGraph) { constexpr size_t kCount = 400; constexpr size_t kDim = 64; std::vector data(kCount * kDim, 0.0F); @@ -385,7 +385,7 @@ TEST(QGConstructionTest, OnlyFinalReversePruningCompletesRowsWithoutWritingGraph } } -TEST(QGConstructionTest, RefinementMatchesOriginalFinalEndpointAndSeededFallback) { +TEST(QgConstructionTest, RefinementMatchesOriginalFinalEndpointAndSeededFallback) { constexpr size_t kDim = 65; constexpr size_t kCount = 65; std::vector data(kCount * kDim, 0.0F); @@ -478,18 +478,18 @@ TEST(BatchQueryTest, PreparedSumPreservesLookupTableAndMetric) { } } -TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { - for (const size_t kDim : {64U, 65U, 1025U, 4097U}) { - SCOPED_TRACE(kDim); +TEST(QgSearchTest, ScratchSearchMatchesRegularSearch) { + for (const size_t dim : {64U, 65U, 1025U, 4097U}) { + SCOPED_TRACE(dim); // Cover IDs beyond 255 at small dimensions; wide vectors need only // a non-complete graph with the same degree and search window. - const size_t kNumPoints = kDim <= 65 ? 257 : 65; + const size_t num_points = dim <= 65 ? 257 : 65; constexpr size_t kDegree = 32; constexpr size_t kEfSearch = 1; - const size_t kNumQueries = kDim == 64 ? 256 : 16; + const size_t num_queries = dim == 64 ? 256 : 16; std::mt19937 rng(123); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector data(kNumPoints * kDim); + std::vector data(num_points * dim); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); for (const size_t bits : {0U, 4U, 8U}) { @@ -497,7 +497,7 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { for (const MetricType metric : {METRIC_L2, METRIC_IP}) { SCOPED_TRACE(static_cast(metric)); QuantizedGraph graph( - kNumPoints, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits, 42 + num_points, dim, kDegree, metric, RotatorType::FhtKacRotator, bits, 42 ); // Graph construction is setup here; cache vectors to avoid repeated // reconstruction. CachedResetMatchesFreshConstruction checks parity. @@ -510,7 +510,7 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { BatchQuery batch_query; buffer::SearchBuffer search_pool(kEfSearch); buffer::SearchBuffer result_pool(1); - VisitedSet visited(kNumPoints, kNumPoints / 10); + VisitedSet visited(num_points, num_points / 10); PID label = kPidMax; float distance = 0.0F; @@ -532,9 +532,9 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { ); builder.build(2); - for (size_t query_index = 0; query_index < kNumQueries; ++query_index) { + for (size_t query_index = 0; query_index < num_queries; ++query_index) { SCOPED_TRACE(query_index); - std::vector query(kDim); + std::vector query(dim); std::generate(query.begin(), query.end(), [&] { return distribution(rng); }); @@ -563,10 +563,10 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { const float expected_distance = metric == METRIC_IP ? dot_product_dis( - query.data(), data.data() + scratch_label * kDim, kDim + query.data(), data.data() + scratch_label * dim, dim ) : euclidean_sqr( - query.data(), data.data() + scratch_label * kDim, kDim + query.data(), data.data() + scratch_label * dim, dim ); if (bits == 0) { EXPECT_NEAR(scratch_distance, expected_distance, 1e-5F); @@ -574,11 +574,11 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { } if (metric == METRIC_L2 && bits == 0) { - const PID hint = static_cast(kNumPoints - 1); + const PID hint = static_cast(num_points - 1); PID hinted_label = kPidMax; float hinted_distance = 1.0F; graph.search_with_scratch( - data.data() + hint * kDim, + data.data() + hint * dim, 1, &hinted_label, &hinted_distance, @@ -599,7 +599,7 @@ TEST(QuantizedGraphSearchTest, ScratchSearchMatchesRegularSearch) { } } -TEST(QuantizedGraphSearchTest, SingleResultMatchesTopKWithTiesAndHints) { +TEST(QgSearchTest, SingleResultMatchesTopKWithTiesAndHints) { constexpr size_t kCount = 65; for (size_t dim : {64U, 65U}) { for (size_t bits : {0U, 4U, 8U}) { @@ -672,17 +672,17 @@ TEST(QuantizedGraphSearchTest, SingleResultMatchesTopKWithTiesAndHints) { } } -TEST(QuantizedGraphSearchTest, ScratchSearchRejectsDefaultVisitedSet) { - constexpr size_t kNumPoints = 66; +TEST(QgSearchTest, ScratchSearchRejectsDefaultVisitedSet) { + constexpr size_t kCount = 66; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; constexpr size_t kEfSearch = 32; std::mt19937 rng(123); std::uniform_real_distribution distribution(-1.0F, 1.0F); - std::vector data(kNumPoints * kDim); + std::vector data(kCount * kDim); std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); - QuantizedGraph graph(kNumPoints, kDim, kDegree); + QuantizedGraph graph(kCount, kDim, kDegree); QGBuilder builder(graph, kDegree, data.data(), 1, 42); builder.build(2); graph.set_ef(kEfSearch); @@ -713,9 +713,9 @@ TEST(QuantizedGraphSearchTest, ScratchSearchRejectsDefaultVisitedSet) { ); FAIL() << "Expected an uninitialized visited set to be rejected"; } catch (const std::invalid_argument& error) { - EXPECT_STREQ(error.what(), "SymphonyQG visited set must be initialized"); + EXPECT_STREQ(error.what(), "QuantizedGraph visited set must be initialized"); } - visited.initialize(kNumPoints, kNumPoints / 10); + visited.initialize(kCount, kCount / 10); for (size_t query_index = 0; query_index < 2; ++query_index) { const float* query = data.data() + query_index * kDim; @@ -743,7 +743,7 @@ TEST(QuantizedGraphSearchTest, ScratchSearchRejectsDefaultVisitedSet) { } } -TEST(QuantizedGraphSearchTest, LookupTableResetMatchesFreshAcrossSizesAndModes) { +TEST(QgSearchTest, LookupTableResetMatchesFreshAcrossSizesAndModes) { Lut reused; for (const size_t dim : {64U, 4160U, 128U, 64U}) { SCOPED_TRACE(dim); @@ -768,7 +768,7 @@ TEST(QuantizedGraphSearchTest, LookupTableResetMatchesFreshAcrossSizesAndModes) } } -TEST(QGConstructionTest, CopiedCacheOwnsItsQueryStorage) { +TEST(QgConstructionTest, CopiedCacheOwnsItsQueryStorage) { constexpr size_t kCount = 65, kDim = 65; for (size_t bits : {4U, 8U}) { std::vector data(kCount * kDim); @@ -793,7 +793,7 @@ TEST(QGConstructionTest, CopiedCacheOwnsItsQueryStorage) { } } -TEST(QGConstructionTest, CachedResetMatchesFreshConstruction) { +TEST(QgConstructionTest, CachedResetMatchesFreshConstruction) { // Keep more candidates than the degree bound so construction still prunes, // without making every reset repeat a large graph build. constexpr size_t kCount = 41; @@ -839,7 +839,7 @@ TEST(QGConstructionTest, CachedResetMatchesFreshConstruction) { } } -TEST(QuantizedGraphCompatibilityTest, SupportsExplicitAndDeducedFloatTypes) { +TEST(QgCompatibilityTest, SupportsExplicitAndDeducedFloatTypes) { using Graph = rabitqlib::symqg::QuantizedGraph; static_assert(std::is_same_v>); static_assert(!std::is_copy_constructible_v); @@ -859,7 +859,7 @@ TEST(QuantizedGraphCompatibilityTest, SupportsExplicitAndDeducedFloatTypes) { EXPECT_EQ(explicit_graph.dimension(), 64); } -TEST(QuantizedGraphLayoutTest, KeepsVectorsCodesAndNeighborsInContiguousRows) { +TEST(QgLayoutTest, KeepsVectorsCodesAndNeighborsInContiguousRows) { for (size_t bits : {0U, 4U, 8U}) { SCOPED_TRACE(bits); QuantizedGraph graph( @@ -884,7 +884,7 @@ static_assert( "Intermediate graph input must not be part of the public builder API" ); -TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { +TEST(QgEstimatorTest, MatchesExactDistancesForCollinearResiduals) { for (size_t dim : {64U, 1024U, 1088U}) { for (auto metric : {METRIC_L2, METRIC_IP}) { SCOPED_TRACE(::testing::Message() << dim << "/" << metric); @@ -913,12 +913,14 @@ TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { std::vector backends{ simd::qg_batch_estdist_generic, simd::qg_batch_estdist}; #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) + if (cpu::has_avx2()) { backends.push_back(simd::qg_batch_estdist_avx2); + } #endif #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx512_core()) + if (cpu::has_avx512_core()) { backends.push_back(simd::qg_batch_estdist_avx512); + } #endif for (auto backend : backends) { backend(batch.data(), q_obj, dim, estimates.data()); @@ -943,7 +945,7 @@ TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { } } -TEST(QGEstimatorTest, CandidateMaskBackendsMatchScalarOrdering) { +TEST(QgEstimatorTest, CandidateMaskBackendsMatchScalarOrdering) { std::array distances{}; for (size_t lane = 0; lane < distances.size(); ++lane) { distances[lane] = static_cast(lane % 7) - 2.0F; @@ -955,16 +957,16 @@ TEST(QGEstimatorTest, CandidateMaskBackendsMatchScalarOrdering) { constexpr float kThreshold = 2.0F; // Keep residues 0..4 in each group of seven, except +inf at lane 17. // Unordered NaN at lane 3 and -inf at lane 29 remain candidates. - constexpr uint32_t expected = 0xF3E5CF9FU; - EXPECT_EQ(simd::qg_candidate_mask_generic(distances.data(), kThreshold), expected); + constexpr uint32_t kExpected = 0xF3E5CF9FU; + EXPECT_EQ(simd::qg_candidate_mask_generic(distances.data(), kThreshold), kExpected); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { - EXPECT_EQ(simd::qg_candidate_mask_avx2(distances.data(), kThreshold), expected); + EXPECT_EQ(simd::qg_candidate_mask_avx2(distances.data(), kThreshold), kExpected); } #endif } -TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { +TEST(QgConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { constexpr size_t kCount = 65, kDim = 65; for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {4U, 8U}) { @@ -1008,7 +1010,7 @@ TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { } } -TEST(QGConstructionTest, EncodedSeedSearchMatchesRetainedScores) { +TEST(QgConstructionTest, EncodedSeedSearchMatchesRetainedScores) { constexpr size_t kCount = 97, kDim = 65, kDegree = 32; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -1051,7 +1053,7 @@ TEST(QGConstructionTest, EncodedSeedSearchMatchesRetainedScores) { } } -TEST(QGConstructionTest, DefaultsToPipnnInitialization) { +TEST(QgConstructionTest, DefaultsToPipnnInitialization) { constexpr size_t kCount = 97, kDim = 65, kDegree = 32; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -1086,7 +1088,7 @@ TEST(QGConstructionTest, DefaultsToPipnnInitialization) { } } -TEST(QGConstructionTest, AverageDegreeDoesNotWrapAboveUint32Range) { +TEST(QgConstructionTest, AverageDegreeDoesNotWrapAboveUint32Range) { std::vector data(65 * 65, 0.5F); QuantizedGraph graph(65, 65, 32); QGBuilder builder(graph, 64, data.data(), 1, QGInitialization::Random); @@ -1095,7 +1097,7 @@ TEST(QGConstructionTest, AverageDegreeDoesNotWrapAboveUint32Range) { EXPECT_FLOAT_EQ(builder.avg_degree(), 2147483648.0F); } -TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { +TEST(QgConstructionTest, SupportsExplicitRandomInitialization) { std::vector data(65 * 65, 0.5F); QuantizedGraph graph(65, 65, 32); QGBuilder builder(graph, 64, data.data(), 1, QGInitialization::Random); @@ -1121,7 +1123,7 @@ TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { ); } -TEST(QGConstructionTest, RejectsNullDataForEveryInitializationMode) { +TEST(QgConstructionTest, RejectsNullDataForEveryInitializationMode) { QuantizedGraph graph(33, 64, 32); EXPECT_THROW( (QGBuilder(graph, 32, nullptr, 1, QGInitialization::PiPNN)), std::invalid_argument @@ -1131,7 +1133,7 @@ TEST(QGConstructionTest, RejectsNullDataForEveryInitializationMode) { ); } -TEST(QGConstructionTest, UsesExplicitThreadCountsWithoutChangingCallerState) { +TEST(QgConstructionTest, UsesExplicitThreadCountsWithoutChangingCallerState) { constexpr size_t kCount = 33, kDim = 64; std::vector data(kCount * kDim, 0.25F); const int caller_threads = omp_get_max_threads(); @@ -1148,7 +1150,7 @@ TEST(QGConstructionTest, UsesExplicitThreadCountsWithoutChangingCallerState) { EXPECT_EQ(omp_get_max_threads(), caller_threads); } -TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { +TEST(QgConstructionTest, InitializesUnusedPartialBatchFactors) { constexpr size_t kCount = 33, kDim = 64; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -1171,7 +1173,7 @@ TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { } } -TEST(QGConstructionTest, FullyInitializesRowsBeforeSearchAndSave) { +TEST(QgConstructionTest, FullyInitializesRowsBeforeSearchAndSave) { constexpr size_t kCount = 65, kDim = 65; std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { @@ -1207,7 +1209,7 @@ TEST(QGConstructionTest, FullyInitializesRowsBeforeSearchAndSave) { } } -TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { +TEST(QgConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { constexpr size_t kCount = 97, kDim = 65, kDegree = 64; constexpr std::array kDegrees{0, 1, 31, 32, 33, 64}; for (auto metric : {METRIC_L2, METRIC_IP}) { @@ -1302,7 +1304,7 @@ TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { } } -TEST(QGConstructionTest, RejectsInvalidSeedStructureAndIds) { +TEST(QgConstructionTest, RejectsInvalidSeedStructureAndIds) { constexpr size_t kCount = 33, kDim = 64; std::vector data(kCount * kDim, 0.1F); QuantizedGraph graph(kCount, kDim, 32); @@ -1318,39 +1320,45 @@ TEST(QGConstructionTest, RejectsInvalidSeedStructureAndIds) { EXPECT_STREQ(error.what(), message); } }; - reject({}, {}, "Seed graph offsets must delimit every vertex"); + reject({}, {}, "QuantizedGraph seed graph offsets must delimit every vertex"); std::vector offsets(kCount + 1, 0); offsets.back() = 1; - reject(offsets, {}, "Seed graph offsets must delimit every vertex"); + reject(offsets, {}, "QuantizedGraph seed graph offsets must delimit every vertex"); offsets.back() = 0; offsets[1] = 1; - reject(offsets, {}, "Seed graph row exceeds degree bound or has invalid offsets"); + reject( + offsets, + {}, + "QuantizedGraph seed graph row exceeds degree bound or has invalid offsets" + ); std::fill(offsets.begin() + 1, offsets.end(), 1); - reject(offsets, {kCount}, "Seed graph IDs must be in range and exclude self"); - reject(offsets, {0}, "Seed graph IDs must be in range and exclude self"); + reject( + offsets, {kCount}, "QuantizedGraph seed graph IDs must be in range and exclude self" + ); + reject(offsets, {0}, "QuantizedGraph seed graph IDs must be in range and exclude self"); std::fill(offsets.begin() + 1, offsets.end(), 33); reject( offsets, std::vector(33, 1), - "Seed graph row exceeds degree bound or has invalid offsets" + "QuantizedGraph seed graph row exceeds degree bound or has invalid offsets" ); } -TEST(QuantizedGraphConfigurationTest, RejectsDegreeNotAlignedForFastScan) { +TEST(QgConfigurationTest, RejectsDegreeNotAlignedForFastScan) { EXPECT_THROW( (QuantizedGraph(64, 64, 16, METRIC_L2, RotatorType::MatrixRotator)), std::invalid_argument ); } -TEST(QuantizedGraphConfigurationTest, RejectsDegreeThatCannotExcludeSelf) { +TEST(QgConfigurationTest, RejectsDegreeThatCannotExcludeSelf) { EXPECT_THROW( (QuantizedGraph(32, 64, 32, METRIC_L2, RotatorType::MatrixRotator)), std::invalid_argument ); } -TEST(QuantizedGraphConfigurationTest, AcceptsOnlySupportedVectorQuantizationBits) { +TEST(QgConfigurationTest, AcceptsOnlySupportedVectorQuantizationBits) { EXPECT_NO_THROW( (QuantizedGraph(33, 64, 32, METRIC_L2, RotatorType::MatrixRotator, 0)) ); @@ -1366,7 +1374,7 @@ TEST(QuantizedGraphConfigurationTest, AcceptsOnlySupportedVectorQuantizationBits ); } -TEST(QuantizedGraphConfigurationTest, RejectsUnsupportedMetric) { +TEST(QgConfigurationTest, RejectsUnsupportedMetric) { EXPECT_THROW( (QuantizedGraph( 33, 64, 32, static_cast(255), RotatorType::MatrixRotator @@ -1375,27 +1383,27 @@ TEST(QuantizedGraphConfigurationTest, RejectsUnsupportedMetric) { ); } -TEST(QuantizedGraphConfigurationTest, RejectsZeroDimension) { +TEST(QgConfigurationTest, RejectsZeroDimension) { EXPECT_THROW( (QuantizedGraph(33, 0, 32, METRIC_L2, RotatorType::MatrixRotator)), std::invalid_argument ); } -TEST(QuantizedGraphConfigurationTest, RejectsOutOfRangeEntryPoint) { +TEST(QgConfigurationTest, RejectsOutOfRangeEntryPoint) { QuantizedGraph graph(33, 64, 32); EXPECT_NO_THROW(graph.set_ep(32)); EXPECT_THROW(graph.set_ep(33), std::invalid_argument); EXPECT_EQ(graph.entry_point(), 32U); } -TEST(QuantizedGraphPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget) { - constexpr size_t kNumPoints = 33; +TEST(QgPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget) { + constexpr size_t kCount = 33; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; - std::vector data(kNumPoints * kDim, 0.25F); + std::vector data(kCount * kDim, 0.25F); QuantizedGraph source( - kNumPoints, kDim, kDegree, METRIC_L2, RotatorType::MatrixRotator, 4 + kCount, kDim, kDegree, METRIC_L2, RotatorType::MatrixRotator, 4 ); QGBuilder builder(source, kDegree, data.data(), 1); builder.build(2); @@ -1447,16 +1455,16 @@ TEST(QuantizedGraphPersistenceTest, RejectsMalformedPayloadWithoutChangingTarget EXPECT_THROW(source.save("/dev/full"), std::ios_base::failure); } - QGConstructionTestAccess::set_encoded_neighbor(source, 0, 0, kNumPoints); + QGConstructionTestAccess::set_encoded_neighbor(source, 0, 0, kCount); source.save(path.c_str()); EXPECT_THROW(target.load(path.c_str()), std::runtime_error); EXPECT_EQ(target.num_vertices(), 65U); EXPECT_EQ(target.metric_type(), METRIC_IP); } -TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { - constexpr size_t kNum = 33, kDim = 65, kDegree = 32; - std::vector data(kNum * kDim); +TEST(QgPersistenceTest, PreservesNewAndLegacyPaddedDomains) { + constexpr size_t kCount = 33, kDim = 65, kDegree = 32; + std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.13F); } @@ -1468,7 +1476,7 @@ TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { for (bool legacy : {false, true}) { SCOPED_TRACE(::testing::Message() << bits << '/' << legacy); QuantizedGraph source( - kNum, kDim, kDegree, metric, rotator, bits + kCount, kDim, kDegree, metric, rotator, bits ); EXPECT_EQ(source.padded_dim(), 96U); if (legacy) { @@ -1478,7 +1486,7 @@ TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { } QGBuilder builder(source, kDegree, data.data(), 1); builder.build(2); - source.set_ef(kNum); + source.set_ef(kCount); std::array expected_ids{}, actual_ids{}; std::array expected_distances{}, actual_distances{}; source.search( @@ -1492,7 +1500,7 @@ TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { QGConstructionTestAccess::rows(source), QGConstructionTestAccess::rows(loaded) ); - loaded.set_ef(kNum); + loaded.set_ef(kCount); loaded.search( data.data(), 5, actual_ids.data(), actual_distances.data() ); @@ -1507,12 +1515,12 @@ TEST(QuantizedGraphPersistenceTest, PreservesNewAndLegacyPaddedDomains) { } } -TEST(QuantizedGraphLifecycleTest, DestroysConcreteRotatorThroughBasePointer) { +TEST(QgLifecycleTest, DestroysConcreteRotatorThroughBasePointer) { QuantizedGraph graph(33, 64, 32, METRIC_L2, RotatorType::MatrixRotator); EXPECT_EQ(graph.num_vertices(), 33U); } -TEST(QuantizedGraphLifecycleTest, RejectsSearchAndSaveBeforeBuild) { +TEST(QgLifecycleTest, RejectsSearchAndSaveBeforeBuild) { const rabitq_test::TempFile path_file("rabitq_qg_unbuilt.index"); const std::string& path = path_file.path(); std::array query{}; @@ -1534,11 +1542,11 @@ TEST(QuantizedGraphLifecycleTest, RejectsSearchAndSaveBeforeBuild) { EXPECT_FALSE(std::filesystem::exists(path)); } -TEST(QuantizedGraphLifecycleTest, BuilderDoesNotPublishPartialGraph) { - constexpr size_t kNumPoints = 33; +TEST(QgLifecycleTest, BuilderDoesNotPublishPartialGraph) { + constexpr size_t kCount = 33; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; - std::vector data(kNumPoints * kDim, 0.25F); + std::vector data(kCount * kDim, 0.25F); std::array query{}; PID result = 0; float distance = 0; @@ -1547,7 +1555,7 @@ TEST(QuantizedGraphLifecycleTest, BuilderDoesNotPublishPartialGraph) { SCOPED_TRACE(static_cast(init)); const rabitq_test::TempFile path_file("rabitq_qg_partial.index"); const std::string& path = path_file.path(); - QuantizedGraph graph(kNumPoints, kDim, kDegree); + QuantizedGraph graph(kCount, kDim, kDegree); QGBuilder builder(graph, kDegree, data.data(), 1, init); graph.set_ef(1); EXPECT_THROW(graph.save(path.c_str()), std::logic_error); @@ -1561,30 +1569,30 @@ TEST(QuantizedGraphLifecycleTest, BuilderDoesNotPublishPartialGraph) { } } -TEST(QGBuilderMetricTest, UsesInnerProductDistanceToChooseEntryPoint) { - constexpr size_t kNumPoints = 33; +TEST(QgBuilderMetricTest, UsesInnerProductDistanceToChooseEntryPoint) { + constexpr size_t kCount = 33; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; - std::vector data(kNumPoints * kDim, 0.0F); + std::vector data(kCount * kDim, 0.0F); data[0] = 100.0F; - const std::vector centroid = compute_centroid(data.data(), kNumPoints, kDim, 1); + const std::vector centroid = compute_centroid(data.data(), kCount, kDim, 1); const PID expected = - exact_nn(data.data(), centroid.data(), kNumPoints, kDim, 1, dot_product_dis); + exact_nn(data.data(), centroid.data(), kCount, kDim, 1, dot_product_dis); const PID euclidean_entry = - exact_nn(data.data(), centroid.data(), kNumPoints, kDim, 1, euclidean_sqr); + exact_nn(data.data(), centroid.data(), kCount, kDim, 1, euclidean_sqr); ASSERT_EQ(expected, 0U); ASSERT_NE(euclidean_entry, expected); QuantizedGraph graph( - kNumPoints, kDim, kDegree, METRIC_IP, RotatorType::MatrixRotator + kCount, kDim, kDegree, METRIC_IP, RotatorType::MatrixRotator ); QGBuilder builder(graph, kDegree, data.data(), 1); EXPECT_EQ(graph.entry_point(), expected); } -TEST(QGEstimatorTest, AccumulatesBeyondUint16Range) { +TEST(QgEstimatorTest, AccumulatesBeyondUint16Range) { constexpr std::array kDimensions = {1024, 1088, 2048}; for (size_t padded_dim : kDimensions) { @@ -1621,11 +1629,11 @@ TEST(QGEstimatorTest, AccumulatesBeyondUint16Range) { } } -TEST(QGQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { - constexpr size_t kNumPoints = 33; +TEST(QgQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { + constexpr size_t kCount = 33; constexpr size_t kDim = 64; constexpr size_t kDegree = 32; - std::vector data(kNumPoints * kDim); + std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.13F) + std::cos(static_cast(i) * 0.07F); @@ -1634,19 +1642,19 @@ TEST(QGQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { for (size_t bits : {4U, 8U}) { SCOPED_TRACE(bits); QuantizedGraph graph( - kNumPoints, kDim, kDegree, METRIC_L2, RotatorType::MatrixRotator, bits + kCount, kDim, kDegree, METRIC_L2, RotatorType::MatrixRotator, bits ); { QGBuilder builder(graph, kDegree, data.data(), 1); builder.build(2); } - graph.set_ef(kNumPoints); + graph.set_ef(kCount); std::array ids{}; std::array distances{}; graph.search(data.data(), ids.size(), ids.data(), distances.data()); for (size_t i = 0; i < ids.size(); ++i) { - EXPECT_LT(ids[i], kNumPoints); + EXPECT_LT(ids[i], kCount); EXPECT_TRUE(std::isfinite(distances[i])); } @@ -1656,7 +1664,7 @@ TEST(QGQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { const std::string& path = path_file.path(); graph.save(path.c_str()); QuantizedGraph loaded; - loaded.set_ef(kNumPoints); + loaded.set_ef(kCount); loaded.load(path.c_str()); EXPECT_TRUE(loaded.is_quantized()); EXPECT_EQ(loaded.quantization_bits(), bits); @@ -1671,30 +1679,30 @@ TEST(QGQuantTest, SearchesAndRoundTripsFourAndEightBitIndexes) { } } -TEST(QGSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { - constexpr size_t kNumPoints = 33, kDim = 64, kDegree = 32; - std::vector data(kNumPoints * kDim); +TEST(QgSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { + constexpr size_t kCount = 33, kDim = 64, kDegree = 32; + std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.11F); } - QuantizedGraph graph(kNumPoints, kDim, kDegree); + QuantizedGraph graph(kCount, kDim, kDegree); QGBuilder builder(graph, kDegree, data.data(), 1, QGInitialization::Random); builder.build(2); - std::array ids{}; - std::array distances{}; + std::array ids{}; + std::array distances{}; EXPECT_THROW(graph.set_ef(0), std::invalid_argument); graph.set_ef(4); EXPECT_THROW( graph.search(data.data(), 5, ids.data(), distances.data()), std::invalid_argument ); - graph.set_ef(kNumPoints); + graph.set_ef(kCount); EXPECT_THROW( graph.search(data.data(), 0, ids.data(), distances.data()), std::invalid_argument ); EXPECT_THROW( graph.search( - data.data(), static_cast(kNumPoints + 1), ids.data(), distances.data() + data.data(), static_cast(kCount + 1), ids.data(), distances.data() ), std::invalid_argument ); @@ -1706,15 +1714,13 @@ TEST(QGSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { ); EXPECT_THROW(graph.search(data.data(), 1, ids.data(), nullptr), std::invalid_argument); - graph.search( - data.data(), static_cast(kNumPoints), ids.data(), distances.data() - ); - for (size_t i = 0; i < kNumPoints; ++i) { - EXPECT_LT(ids[i], kNumPoints); + graph.search(data.data(), static_cast(kCount), ids.data(), distances.data()); + for (size_t i = 0; i < kCount; ++i) { + EXPECT_LT(ids[i], kCount); EXPECT_TRUE(std::isfinite(distances[i])); } - for (PID source = 0; source < kNumPoints; ++source) { + for (PID source = 0; source < kCount; ++source) { for (size_t lane = 0; lane < kDegree; ++lane) { QGConstructionTestAccess::set_encoded_neighbor(graph, source, lane, 0); } @@ -1731,7 +1737,7 @@ TEST(QGSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { })); } -TEST(QGSearchTest, BatchAndParallelQueriesMatchSerialAcrossIndexesAndSettings) { +TEST(QgSearchTest, BatchAndParallelQueriesMatchSerialAcrossIndexesAndSettings) { constexpr size_t kNumQueries = 9; constexpr size_t kTopK = 5; for (size_t bits : {0U, 4U, 8U}) { @@ -1812,7 +1818,7 @@ TEST(QGSearchTest, BatchAndParallelQueriesMatchSerialAcrossIndexesAndSettings) { } } -TEST(QGSearchTest, BatchSearchValidatesBuffersSizesAndState) { +TEST(QgSearchTest, BatchSearchValidatesBuffersSizesAndState) { QuantizedGraph graph(97, 65, 32, METRIC_L2, RotatorType::FhtKacRotator, 0, 42); std::vector data(97 * 65, 0.1F), distances(97); std::vector ids(97); @@ -1905,40 +1911,40 @@ TEST(QGSearchTest, BatchSearchValidatesBuffersSizesAndState) { ); } -TEST(QGSearchTest, ConcurrentBatchCallsMatchIndividualQueries) { - constexpr size_t kDimension = 65, kPoints = 65, kQueries = 9, k = 3; - std::vector data(kPoints * kDimension); +TEST(QgSearchTest, ConcurrentBatchCallsMatchIndividualQueries) { + constexpr size_t kDim = 65, kCount = 65, kNumQueries = 9, kTopK = 3; + std::vector data(kCount * kDim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.19F); } QuantizedGraph graph( - kPoints, kDimension, 32, METRIC_IP, RotatorType::FhtKacRotator, 8, 42 + kCount, kDim, 32, METRIC_IP, RotatorType::FhtKacRotator, 8, 42 ); QGBuilder builder(graph, 64, data.data(), 1); builder.build(); graph.set_ef(48); - std::vector expected(2 * kQueries * k), actual(expected.size(), kPidMax); + std::vector expected(2 * kNumQueries * kTopK), actual(expected.size(), kPidMax); std::vector expected_distances(expected.size()), distances(expected.size()); - for (size_t i = 0; i < 2 * kQueries; ++i) { + for (size_t i = 0; i < 2 * kNumQueries; ++i) { graph.search( - data.data() + i * kDimension, - k, - expected.data() + i * k, - expected_distances.data() + i * k + data.data() + i * kDim, + kTopK, + expected.data() + i * kTopK, + expected_distances.data() + i * kTopK ); } const auto search = [&](size_t start) { graph.search_batch( - data.data() + start * kDimension, - kQueries, - k, - actual.data() + start * k, - distances.data() + start * k, + data.data() + start * kDim, + kNumQueries, + kTopK, + actual.data() + start * kTopK, + distances.data() + start * kTopK, 2 ); }; auto first = std::async(std::launch::async, search, 0); - auto second = std::async(std::launch::async, search, kQueries); + auto second = std::async(std::launch::async, search, kNumQueries); first.get(); second.get(); EXPECT_EQ(actual, expected); @@ -1949,7 +1955,7 @@ TEST(QGSearchTest, ConcurrentBatchCallsMatchIndividualQueries) { } // namespace rabitqlib::symqg namespace rabitqlib::symqg { -TEST(QuantizedGraphSearchTest, ScratchSearchRejectsIncompleteResults) { +TEST(QgSearchTest, ScratchSearchRejectsIncompleteResults) { QuantizedGraph graph(66, 64, 32, METRIC_L2, RotatorType::FhtKacRotator, 0, 42); std::vector data(66 * 64, 0.0F); QGBuilder builder(graph, 66, data.data(), 1, 42); @@ -2000,15 +2006,17 @@ TEST(QuantizedGraphSearchTest, ScratchSearchRejectsIncompleteResults) { } // namespace rabitqlib::symqg namespace rabitqlib::symqg { -TEST(QGSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { - constexpr size_t kCount = 129, kDim = 65, kQueries = 13, k = 5; +TEST(QgSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { + constexpr size_t kCount = 129, kDim = 65, kNumQueries = 13, kTopK = 5; std::mt19937 random(314); std::normal_distribution value; - std::vector data(kCount * kDim), queries(kQueries * kDim); - for (auto& x : data) + std::vector data(kCount * kDim), queries(kNumQueries * kDim); + for (auto& x : data) { x = value(random); - for (auto& x : queries) + } + for (auto& x : queries) { x = value(random); + } for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {0U, 4U, 8U}) { SCOPED_TRACE(metric); @@ -2018,15 +2026,15 @@ TEST(QGSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { ); QGBuilder builder(graph, 64, data.data(), 1); builder.build(); - std::vector narrow(kQueries * k), wide(narrow.size()); + std::vector narrow(kNumQueries * kTopK), wide(narrow.size()); std::vector narrow_distances(narrow.size()), wide_distances(wide.size()); - graph.set_ef(k); + graph.set_ef(kTopK); graph.search_batch( - queries.data(), kQueries, k, narrow.data(), narrow_distances.data() + queries.data(), kNumQueries, kTopK, narrow.data(), narrow_distances.data() ); graph.set_ef(96); graph.search_batch( - queries.data(), kQueries, k, wide.data(), wide_distances.data() + queries.data(), kNumQueries, kTopK, wide.data(), wide_distances.data() ); graph.set_ef(1); std::promise start; @@ -2035,36 +2043,46 @@ TEST(QGSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { size_t threads, const auto& expected, const auto& expected_distances) { - std::vector actual(kQueries * k); + std::vector actual(kNumQueries * kTopK); std::vector distances(actual.size()); ready.wait(); for (size_t i = 0; i < 16; ++i) { graph.search_batch_with_ef( queries.data(), - kQueries, - k, + kNumQueries, + kTopK, actual.data(), distances.data(), ef, threads ); - if (actual != expected || distances != expected_distances) + if (actual != expected || distances != expected_distances) { return false; + } graph.search_batch_with_ef( - queries.data(), 1, k, actual.data(), distances.data(), ef, threads + queries.data(), + 1, + kTopK, + actual.data(), + distances.data(), + ef, + threads ); - if (!std::equal(actual.begin(), actual.begin() + k, expected.begin()) || + if (!std::equal( + actual.begin(), actual.begin() + kTopK, expected.begin() + ) || !std::equal( distances.begin(), - distances.begin() + k, + distances.begin() + kTopK, expected_distances.begin() - )) + )) { return false; + } } return true; }; auto first = std::async(std::launch::async, [&] { - return run(k, 1, narrow, narrow_distances); + return run(kTopK, 1, narrow, narrow_distances); }); auto second = std::async(std::launch::async, [&] { return run(96, 2, wide, wide_distances); @@ -2074,13 +2092,22 @@ TEST(QGSearchTest, ExplicitWindowsAreIndependentOfDefaultAndConcurrentCalls) { EXPECT_TRUE(second.get()); EXPECT_THROW( graph.search_batch( - queries.data(), kQueries, k, narrow.data(), narrow_distances.data() + queries.data(), + kNumQueries, + kTopK, + narrow.data(), + narrow_distances.data() ), std::invalid_argument ); EXPECT_THROW( graph.search_batch_with_ef( - queries.data(), kQueries, k, narrow.data(), narrow_distances.data(), 0 + queries.data(), + kNumQueries, + kTopK, + narrow.data(), + narrow_distances.data(), + 0 ), std::invalid_argument ); diff --git a/tests/unit/rabitqlib/quantization/core_api_test.cpp b/tests/unit/rabitqlib/quantization/core_api_test.cpp index 42980bcc..696db310 100644 --- a/tests/unit/rabitqlib/quantization/core_api_test.cpp +++ b/tests/unit/rabitqlib/quantization/core_api_test.cpp @@ -42,7 +42,7 @@ api::CodeView view(api::CodeOutput output) { {output.extra.data, output.extra.bytes}}; } -TEST(CoreApi, RejectsNonFiniteValuesAtEveryInputPosition) { +TEST(CoreApiTest, RejectsNonFiniteValuesAtEveryInputPosition) { for (size_t dim : {2U, 3U, 15U, 16U, 17U, 31U, 32U, 33U, 65U}) { const api::Parameters parameters{dim, 1}; const auto sizes = api::requirements(parameters, api::CodeLayout::UnpackedU16); @@ -92,7 +92,7 @@ TEST(CoreApi, RejectsNonFiniteValuesAtEveryInputPosition) { } } -TEST(CoreApi, ExactWorkspaceSizeProtectsBoundaryBytes) { +TEST(CoreApiTest, ExactWorkspaceSizeProtectsBoundaryBytes) { for (size_t dim : {2U, 3U, 31U, 33U, 65U, 96U}) { for (size_t bits : {1U, 4U, 9U}) { SCOPED_TRACE(dim); @@ -138,7 +138,7 @@ TEST(CoreApi, ExactWorkspaceSizeProtectsBoundaryBytes) { } } -TEST(CoreApi, CodesFactorsAndBothLayoutsMatchExistingEncoder) { +TEST(CoreApiTest, CodesFactorsAndBothLayoutsMatchExistingEncoder) { std::mt19937 generator(123); std::uniform_real_distribution distribution(-1, 1); for (size_t dim : {32U, 64U, 96U, 128U, 160U, 256U}) { @@ -213,8 +213,9 @@ TEST(CoreApi, CodesFactorsAndBothLayoutsMatchExistingEncoder) { ? rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim) : -rabitqlib::dot_product(query.data(), centroid.data(), dim); float sum = 0; - for (float x : query) + for (float x : query) { sum += x; + } const float expected = rabitqlib::quant::full_est_dist( reference.data(), query.data(), @@ -302,7 +303,7 @@ TEST(CoreApi, CodesFactorsAndBothLayoutsMatchExistingEncoder) { } } -TEST(CoreApi, UnpackedCodesDoNotRequireLibraryPadding) { +TEST(CoreApiTest, UnpackedCodesDoNotRequireLibraryPadding) { for (size_t dim : {2U, 3U, 15U, 33U, 65U}) { std::vector data(dim, 1.0F), query(dim, 0.5F); api::Parameters p{dim, 4}; @@ -321,13 +322,13 @@ TEST(CoreApi, UnpackedCodesDoNotRequireLibraryPadding) { } } -TEST(CoreApi, NullCentroidZeroResidualAndCallerFactorColumns) { - constexpr size_t dim = 96; - std::array data{}, query{}; +TEST(CoreApiTest, NullCentroidZeroResidualAndCallerFactorColumns) { + constexpr size_t kDim = 96; + std::array data{}, query{}; query.fill(0.25F); for (auto metric : {METRIC_L2, METRIC_IP}) { for (size_t bits : {1U, 4U, 9U}) { - api::Parameters p{dim, bits, metric}; + api::Parameters p{kDim, bits, metric}; const auto sizes = api::requirements(p, api::CodeLayout::SplitPacked); Workspace workspace(sizes.workspace_bytes); std::vector signs(sizes.code_bytes), extra(sizes.extra_bytes); @@ -339,7 +340,7 @@ TEST(CoreApi, NullCentroidZeroResidualAndCallerFactorColumns) { // Store factors as separate columns, not adjacent to either code buffer. std::array add{f.add}, scale{f.rescale}, error{f.error}; auto q = api::prepare_query(p, query.data(), nullptr); - EXPECT_EQ(q.dimension(), dim); + EXPECT_EQ(q.dimension(), kDim); EXPECT_EQ(q.total_bits(), bits); EXPECT_EQ(q.metric_type(), metric); EXPECT_FLOAT_EQ( @@ -355,7 +356,7 @@ TEST(CoreApi, NullCentroidZeroResidualAndCallerFactorColumns) { } } -TEST(CoreApi, RejectsBadParametersAndBuffersBeforeWritingCodes) { +TEST(CoreApiTest, RejectsBadParametersAndBuffersBeforeWritingCodes) { api::Parameters p{64, 4}; auto r = api::requirements(p, api::CodeLayout::UnpackedU16); Workspace workspace(r.workspace_bytes); @@ -445,7 +446,7 @@ TEST(CoreApi, RejectsBadParametersAndBuffersBeforeWritingCodes) { ); } -TEST(CoreApi, ReusedQueryAndWorkspaceStillValidateEachRecord) { +TEST(CoreApiTest, ReusedQueryAndWorkspaceStillValidateEachRecord) { constexpr size_t kDim = 32; const api::Parameters parameters{kDim, 4, METRIC_L2}; const auto raw = api::requirements(parameters, api::CodeLayout::UnpackedU16); @@ -498,24 +499,24 @@ TEST(CoreApi, ReusedQueryAndWorkspaceStillValidateEachRecord) { EXPECT_EQ(extra, expected_extra); } -TEST(CoreApi, ScoringRetainsBufferAndNumericalChecks) { - constexpr size_t dim = 64; - std::array values{}; - std::array codes{}, extra{}; +TEST(CoreApiTest, ScoringRetainsBufferAndNumericalChecks) { + constexpr size_t kDim = 64; + std::array values{}; + std::array codes{}, extra{}; float (*volatile compute_pointer)( const api::QueryState&, api::CodeView, const api::Factors& ) = &api::compute; - alignas(float) std::array misaligned{}; + alignas(float) std::array misaligned{}; const auto* bad_alignment = reinterpret_cast(misaligned.data() + 1); EXPECT_THROW( - api::prepare_query({dim, 4}, bad_alignment, nullptr), std::invalid_argument + api::prepare_query({kDim, 4}, bad_alignment, nullptr), std::invalid_argument ); EXPECT_THROW( - api::prepare_query({dim, 4}, values.data(), bad_alignment), std::invalid_argument + api::prepare_query({kDim, 4}, values.data(), bad_alignment), std::invalid_argument ); for (size_t bits : {1U, 4U, 9U}) { for (auto metric : {METRIC_L2, METRIC_IP}) { - const api::Parameters parameters{dim, bits, metric}; + const api::Parameters parameters{kDim, bits, metric}; const auto query = api::prepare_query(parameters, values.data(), nullptr); for (auto layout : {api::CodeLayout::SplitPacked, api::CodeLayout::UnpackedU16}) { @@ -574,11 +575,11 @@ TEST(CoreApi, ScoringRetainsBufferAndNumericalChecks) { } } values.fill(1e18F); - const auto large_query = api::prepare_query({dim, 1}, values.data(), nullptr); + const auto large_query = api::prepare_query({kDim, 1}, values.data(), nullptr); EXPECT_THROW( api::compute( large_query, - {api::CodeLayout::SplitPacked, {codes.data(), dim / 8}, {}}, + {api::CodeLayout::SplitPacked, {codes.data(), kDim / 8}, {}}, {std::numeric_limits::max(), 0, 0} ), std::overflow_error @@ -591,7 +592,7 @@ TEST(CoreApi, ScoringRetainsBufferAndNumericalChecks) { ); } -TEST(CoreApi, ConcurrentCallsUseOnlyTheirOwnWorkspace) { +TEST(CoreApiTest, ConcurrentCallsUseOnlyTheirOwnWorkspace) { struct Case { api::Parameters parameters; std::vector data, query; @@ -664,27 +665,32 @@ TEST(CoreApi, ConcurrentCallsUseOnlyTheirOwnWorkspace) { api::prepare_query(c.parameters, c.query.data(), nullptr); if (signs != c.signs || extra != c.extra || f.add != c.factors.add || f.rescale != c.factors.rescale || f.error != c.factors.error || - api::compute(query, view(out), f) != c.score) + api::compute(query, view(out), f) != c.score) { return false; + } api::CodeOutput raw{ api::CodeLayout::UnpackedU16, {unpacked.data(), unpacked.size()}, {}}; api::convert_codes(c.parameters, view(out), raw, scratch.view()); - if (unpacked != c.unpacked) + if (unpacked != c.unpacked) { return false; + } std::fill(signs.begin(), signs.end(), 0xA5); std::fill(extra.begin(), extra.end(), 0xA5); api::convert_codes(c.parameters, view(raw), out, scratch.view()); if (signs != c.signs || extra != c.extra || - api::compute(query, view(out), f) != c.score) + api::compute(query, view(out), f) != c.score) { return false; + } } return true; })); } - while (ready.load() != cases.size()) + while (ready.load() != cases.size()) { std::this_thread::yield(); + } start.set_value(); - for (auto& worker : workers) + for (auto& worker : workers) { EXPECT_TRUE(worker.get()); + } } } // namespace diff --git a/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp b/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp index bf682567..2e1ea808 100644 --- a/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp +++ b/tests/unit/rabitqlib/quantization/full_code_preparation_test.cpp @@ -11,7 +11,7 @@ namespace rabitqlib::quant { namespace { -TEST(FullCodePreparation, SquaredNormUnderflowPreservesSigns) { +TEST(FullCodePreparationTest, SquaredNormUnderflowPreservesSigns) { namespace api = rabitqlib::quant::v1; for (const size_t dim : {32U, 64U, 96U}) { std::vector data(dim), centroid(dim, 0); diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index d5b13a87..598e944c 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -48,7 +48,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { } for (auto metric : {METRIC_L2, METRIC_IP}) { const float g_add = metric == METRIC_L2 ? 9.0F : -2.0F; - constexpr float g_error = 3.0F; + constexpr float kGError = 3.0F; for (size_t bits = 1; bits <= 9; ++bits) { const size_t ex_bits = bits - 1; std::vector bin(BinDataMap::data_bytes(dim)); @@ -67,20 +67,20 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { ); float distance, lower, ip; split_single_estdist( - bin.data(), prepared, dim, ip, distance, lower, g_add, g_error + bin.data(), prepared, dim, ip, distance, lower, g_add, kGError ); check( - estimate_split_single(bin.data(), prepared, dim, g_add, g_error), + estimate_split_single(bin.data(), prepared, dim, g_add, kGError), distance, lower, ip ); split_single_estdist_direct( - bin.data(), prepared, dim, ip, distance, lower, g_add, g_error + bin.data(), prepared, dim, ip, distance, lower, g_add, kGError ); check( estimate_split_single_direct( - bin.data(), prepared, dim, g_add, g_error + bin.data(), prepared, dim, g_add, kGError ), distance, lower, @@ -101,7 +101,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { lower, ip, g_add, - g_error + kGError ); check( estimate_split_single_full( @@ -112,7 +112,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { dim, ex_bits, g_add, - g_error + kGError ), distance, lower, @@ -129,7 +129,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { lower, ip, g_add, - g_error + kGError ); check( estimate_split_single_full_direct( @@ -140,7 +140,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { dim, ex_bits, g_add, - g_error + kGError ), distance, lower, @@ -178,11 +178,11 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { distance, lower, g_add, - g_error + kGError ); check( estimate_xy_base( - base.data(), base_ip, prepared, dim, base_bits, g_add, g_error + base.data(), base_ip, prepared, dim, base_bits, g_add, kGError ), distance, lower, @@ -201,7 +201,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { lower, ip, g_add, - g_error + kGError ); check( estimate_xy_full( @@ -214,7 +214,7 @@ TEST(SingleEstimatorTest, NamedResultsPreserveLegacyOutputOrder) { base_bits, ex_bits, g_add, - g_error + kGError ), distance, lower, @@ -229,8 +229,9 @@ int level_from_thresholds(double magnitude, double t, int max_code) { int result = 0; if (magnitude > 0) { for (int level = 1; level <= max_code; ++level) { - if (static_cast(level) / magnitude <= t) + if (static_cast(level) / magnitude <= t) { result = level; + } } } return result; @@ -238,11 +239,13 @@ int level_from_thresholds(double magnitude, double t, int max_code) { template double sorted_event_rescale_factor(const std::vector& magnitudes, size_t bits) { - if (magnitudes.empty()) + if (magnitudes.empty()) { return 0; + } const double max_o = *std::max_element(magnitudes.begin(), magnitudes.end()); - if (max_o == 0) + if (max_o == 0) { return 0; + } const int max_code = (1 << bits) - 1; const double end = static_cast(max_code + 10) / max_o; @@ -260,8 +263,9 @@ double sorted_event_rescale_factor(const std::vector& magnitudes, size_t bits if (magnitude > 0) { for (int next = level + 1; next <= max_code; ++next) { const double threshold = static_cast(next) / magnitude; - if (threshold < end) + if (threshold < end) { events.emplace_back(threshold, i); + } } } } @@ -294,11 +298,11 @@ void check_rescale_search_against_sorted_events() { double norm_sq = 0; for (size_t i = 0; i < dim; ++i) { double value = 0; - if (pattern == 1) + if (pattern == 1) { value = i == dim / 2 ? 1 : 0; - else if (pattern == 2) + } else if (pattern == 2) { value = static_cast(i % 4); - else if (pattern == 3) { + } else if (pattern == 3) { state = state * 1664525U + 1013904223U; value = i % 5 == 0 ? 0 : 1.0 + (state >> 16U); } @@ -307,8 +311,9 @@ void check_rescale_search_against_sorted_events() { } if (norm_sq > 0) { const double norm = std::sqrt(norm_sq); - for (T& magnitude : magnitudes) + for (T& magnitude : magnitudes) { magnitude = static_cast(static_cast(magnitude) / norm); + } } for (size_t bits = 0; bits <= 8; ++bits) { SCOPED_TRACE( @@ -330,8 +335,9 @@ void check_rescale_search_against_sorted_events() { ipnorm += (level + 0.5) * magnitude; } T expected_factor = ipnorm == 0 ? T{1} : static_cast(1.0 / ipnorm); - if (!std::isnormal(expected_factor)) + if (!std::isnormal(expected_factor)) { expected_factor = T{1}; + } std::vector code(dim, 0xFF); const T factor = rabitq_impl::ex_bits::quantize_ex( magnitudes.data(), code.data(), dim, bits @@ -672,18 +678,20 @@ TEST(RabitqSignConventionTest, ZeroResidualDoesNotInflateReconstructionNorm) { } // namespace } // namespace rabitqlib::quant -TEST(RabitqCompactCodes, WordTailsUseCompactRowStrides) { +TEST(RabitqCompactCodesTest, WordTailsUseCompactRowStrides) { using namespace rabitqlib; for (size_t dim : {96U, 160U}) { - constexpr size_t count = 3; - std::vector data(count * dim), centroid(dim, 0); - for (size_t i = 0; i < data.size(); ++i) + constexpr size_t kCount = 3; + std::vector data(kCount * dim), centroid(dim, 0); + for (size_t i = 0; i < data.size(); ++i) { data[i] = static_cast(i % 17) - 8; + } for (auto metric : {METRIC_L2, METRIC_IP}) { - std::vector expected(count * dim / 8 + 1, 0xA5), actual = expected; - std::vector add(count), rescale(count), error(count); - std::vector batch_add(count), batch_rescale(count), batch_error(count); - for (size_t i = 0; i < count; ++i) { + std::vector expected(kCount * dim / 8 + 1, 0xA5), actual = expected; + std::vector add(kCount), rescale(kCount), error(kCount); + std::vector batch_add(kCount), batch_rescale(kCount), + batch_error(kCount); + for (size_t i = 0; i < kCount; ++i) { quant::rabitq_impl::one_bit::one_bit_compact_code_to_bytes( data.data() + i * dim, centroid.data(), @@ -698,7 +706,7 @@ TEST(RabitqCompactCodes, WordTailsUseCompactRowStrides) { quant::rabitq_impl::one_bit::one_bit_compact_codes( data.data(), centroid.data(), - count, + kCount, dim, reinterpret_cast(actual.data()), batch_add.data(), diff --git a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp index c0ff1064..77d3fb7a 100644 --- a/tests/unit/rabitqlib/quantization/rescale_search_test.cpp +++ b/tests/unit/rabitqlib/quantization/rescale_search_test.cpp @@ -25,11 +25,13 @@ std::vector levels_at_threshold( ) { std::vector levels(magnitudes.size()); for (size_t i = 0; i < magnitudes.size(); ++i) { - if (magnitudes[i] == 0) + if (magnitudes[i] == 0) { continue; + } for (int level = 1; level <= max_code; ++level) { - if (static_cast(level) / magnitudes[i] <= scale) + if (static_cast(level) / magnitudes[i] <= scale) { levels[i] = level; + } } } return levels; @@ -57,12 +59,14 @@ long double exhaustive_best_squared_cosine( std::vector levels = levels_at_threshold(magnitudes, start, max_code); std::vector> events; for (size_t i = 0; i < magnitudes.size(); ++i) { - if (magnitudes[i] == 0) + if (magnitudes[i] == 0) { continue; + } for (int level = levels[i] + 1; level <= max_code; ++level) { const double threshold = static_cast(level) / magnitudes[i]; - if (threshold < end) + if (threshold < end) { events.emplace_back(threshold, i); + } } } std::sort(events.begin(), events.end()); @@ -81,21 +85,25 @@ long double exhaustive_best_squared_cosine( void normalize_magnitudes(std::vector& magnitudes) { long double norm_squared = 0; - for (float magnitude : magnitudes) + for (float magnitude : magnitudes) { norm_squared += static_cast(magnitude) * magnitude; - if (norm_squared == 0) + } + if (norm_squared == 0) { return; + } const long double norm = std::sqrt(norm_squared); - for (float& magnitude : magnitudes) + for (float& magnitude : magnitudes) { magnitude = static_cast(magnitude / norm); + } } std::pair search_interval( const std::vector& magnitudes, size_t bits ) { const double maximum = *std::max_element(magnitudes.begin(), magnitudes.end()); - if (maximum == 0) + if (maximum == 0) { return {0, 1}; + } const double end = static_cast((1 << bits) + 9) / maximum; return {end * rabitq_impl::ex_bits::kTightStart[bits], end}; } @@ -109,8 +117,9 @@ void expect_optimal_canonical_scale( const auto levels = levels_at_threshold(magnitudes, t, max_code); double first = start; for (size_t i = 0; i < magnitudes.size(); ++i) { - if (levels[i] > 0) + if (levels[i] > 0) { first = std::max(first, static_cast(levels[i]) / magnitudes[i]); + } } EXPECT_EQ(t, first); const long double actual = squared_cosine(magnitudes, levels); @@ -126,8 +135,9 @@ void expect_optimal_canonical_scale( class RabitqRescaleBackendTest : public ::testing::TestWithParam { protected: void SetUp() override { - if (GetParam() ? !cpu::has_avx512_core() : !cpu::has_avx2()) + if (GetParam() ? !cpu::has_avx512_core() : !cpu::has_avx2()) { GTEST_SKIP() << "Requested SIMD backend is unavailable"; + } } RescaleSearch search() const { @@ -157,21 +167,22 @@ TEST_P(RabitqRescaleBackendTest, MatchesExhaustiveObjectiveAndCanonicalThreshold std::mt19937 random(8143); std::normal_distribution gaussian; for (size_t i = 0; i < dim; ++i) { - if (pattern == 1) + if (pattern == 1) { magnitudes[i] = i == dim / 2 ? 1.0F : 0; - else if (pattern == 2) + } else if (pattern == 2) { magnitudes[i] = static_cast(i % 4); - else if (pattern == 3) + } else if (pattern == 3) { magnitudes[i] = 1.0F + static_cast(i % 7) * std::numeric_limits::epsilon(); - else if (pattern == 4) + } else if (pattern == 4) { magnitudes[i] = i % 5 == 0 ? 0 : std::abs(gaussian(random)); - else if (pattern == 5) + } else if (pattern == 5) { magnitudes[i] = i + 1 == dim ? std::numeric_limits::denorm_min() : std::ldexp(1.0F, -static_cast(i % 150)); - else if (pattern == 6) + } else if (pattern == 6) { magnitudes[i] = std::abs(gaussian(random)); + } } normalize_magnitudes(magnitudes); for (size_t bits = 1; bits <= 8; ++bits) { @@ -201,8 +212,9 @@ TEST_P(RabitqRescaleBackendTest, MatchesExhaustiveObjectiveAndCanonicalThreshold EXPECT_EQ(reciprocals.front(), -123.0); EXPECT_EQ(reciprocals.back(), -123.0); ASSERT_TRUE(std::isfinite(t)); - if (t >= 0) + if (t >= 0) { expect_optimal_canonical_scale(magnitudes, max_code, start, end, t); + } } } } @@ -212,8 +224,9 @@ TEST_P(RabitqRescaleBackendTest, CompletesGaussianSearchWithoutFallback) { std::vector magnitudes(128); std::mt19937 random(42); std::normal_distribution gaussian; - for (float& magnitude : magnitudes) + for (float& magnitude : magnitudes) { magnitude = std::abs(gaussian(random)); + } normalize_magnitudes(magnitudes); const auto [start, end] = search_interval(magnitudes, 7); const double t = search()(magnitudes.data(), magnitudes.size(), 127, start, end); @@ -230,8 +243,9 @@ TEST_P(RabitqRescaleBackendTest, ReusesScratchAcrossDimensionsAndThreads) { std::normal_distribution gaussian; for (size_t dim : {129U, 7U, 257U, 1U, 65U, 513U, 3U, 128U, 0U, 17U}) { std::vector magnitudes(dim); - for (float& magnitude : magnitudes) + for (float& magnitude : magnitudes) { magnitude = dim == 17 ? 0 : std::abs(gaussian(random)); + } normalize_magnitudes(magnitudes); const auto interval = dim == 0 ? std::pair{0, 1} : search_interval(magnitudes, 7); @@ -259,9 +273,10 @@ TEST_P(RabitqRescaleBackendTest, ReusesScratchAcrossDimensionsAndThreads) { TEST_P(RabitqRescaleBackendTest, FallsBackWhenNearlyCollinearWinnersAreAmbiguous) { std::vector magnitudes(129); - for (size_t i = 0; i < magnitudes.size(); ++i) + for (size_t i = 0; i < magnitudes.size(); ++i) { magnitudes[i] = 1.0F + static_cast(i) * std::numeric_limits::epsilon(); + } normalize_magnitudes(magnitudes); const auto [start, end] = search_interval(magnitudes, 8); const double t = search()(magnitudes.data(), magnitudes.size(), 255, start, end); diff --git a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp index 70fdcf5d..867a635d 100644 --- a/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp +++ b/tests/unit/rabitqlib/quantization/xy_quantization_test.cpp @@ -20,7 +20,7 @@ using namespace rabitqlib; namespace { -std::vector RandomVec(size_t dim, std::mt19937& gen) { +std::vector random_vec(size_t dim, std::mt19937& gen) { std::normal_distribution dist(0.0F, 1.0F); std::vector v(dim); for (auto& x : v) { @@ -32,7 +32,7 @@ std::vector RandomVec(size_t dim, std::mt19937& gen) { // A one-hot query makes the ex-code inner product return exactly code[j]. // This checks how scoring interprets the packed codes independently of the // library's unpacker, rather than only round-tripping through paired code paths. -std::vector UnpackViaIpKernel(const uint8_t* packed, size_t dim, size_t bits) { +std::vector unpack_via_ip_kernel(const uint8_t* packed, size_t dim, size_t bits) { std::vector out(dim, 0); if (bits == 0) { return out; @@ -58,12 +58,12 @@ struct XySplitResult { float f_rescale_full = 0; }; -XySplitResult SplitCode( +XySplitResult split_code( const std::vector& data, const std::vector& centroid, size_t dim, size_t base_bits, - size_t extra_bits, + size_t ex_bits, MetricType metric = METRIC_L2 ) { XySplitResult res; @@ -72,7 +72,7 @@ XySplitResult SplitCode( std::vector base_block(BaseDataMap::data_bytes(dim, base_bits), 0); std::vector extra_block( - extra_bits > 0 ? ExDataMap::data_bytes(dim, extra_bits) : 1, 0 + ex_bits > 0 ? ExDataMap::data_bytes(dim, ex_bits) : 1, 0 ); // The public entry point in core/encoder.hpp maps the raw buffers onto @@ -82,9 +82,9 @@ XySplitResult SplitCode( centroid.data(), dim, base_bits, - extra_bits, + ex_bits, base_block.data(), - extra_bits > 0 ? extra_block.data() : nullptr, + ex_bits > 0 ? extra_block.data() : nullptr, metric ); @@ -92,13 +92,13 @@ XySplitResult SplitCode( res.f_add_base = base_map.f_add(); res.f_rescale_base = base_map.f_rescale(); res.f_error_base = base_map.f_error(); - res.base_code = UnpackViaIpKernel(base_map.base_code(), dim, base_bits); + res.base_code = unpack_via_ip_kernel(base_map.base_code(), dim, base_bits); - if (extra_bits > 0) { - ExDataMap extra_map(extra_block.data(), dim, extra_bits); + if (ex_bits > 0) { + ExDataMap extra_map(extra_block.data(), dim, ex_bits); res.f_add_full = extra_map.f_add_ex(); res.f_rescale_full = extra_map.f_rescale_ex(); - res.extra_code = UnpackViaIpKernel(extra_map.ex_code(), dim, extra_bits); + res.extra_code = unpack_via_ip_kernel(extra_map.ex_code(), dim, ex_bits); } else { res.f_add_full = res.f_add_base; res.f_rescale_full = res.f_rescale_base; @@ -112,12 +112,12 @@ XySplitResult SplitCode( // match today's one_bit_code_with_factor + ex_bits_code_with_factor -- codes // *and* both factor triplets. The base layer's factors are what BinDataMap // holds in the 1+y path, the full layer's are what ExDataMap holds. -TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { +TEST(XyQuantizationTest, BaseBitsOneMatchesExistingOneBitPlusExBits) { constexpr size_t kDim = 64; - constexpr size_t kExtraBits = 3; + constexpr size_t kExBits = 3; std::mt19937 gen(42); - std::vector data = RandomVec(kDim, gen); + std::vector data = random_vec(kDim, gen); std::vector centroid(kDim, 0.0F); // Reference: today's separate one_bit_code_with_factor (sign bit + the @@ -145,7 +145,7 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { data.data(), centroid.data(), kDim, - kExtraBits, + kExBits, ref_ex_code.data(), ref_f_add, ref_f_rescale, @@ -153,8 +153,8 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { METRIC_L2 ); - // New path: base_bits=1, extra_bits=kExtraBits. - XySplitResult res = SplitCode(data, centroid, kDim, /*base_bits=*/1, kExtraBits); + // New path: base_bits=1, ex_bits=kExBits. + XySplitResult res = split_code(data, centroid, kDim, /*base_bits=*/1, kExBits); for (size_t i = 0; i < kDim; ++i) { EXPECT_EQ(res.base_code[i], static_cast(ref_binary_code[i])) @@ -172,7 +172,7 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { // asserted here; ref_f_error stays as documentation of the reference value. } -// ip(base)*2^extra_bits + ip(extra) must equal ip(total_code) -- the +// ip(base)*2^ex_bits + ip(extra) must equal ip(total_code) -- the // linearity the boosting step relies on. // The sign bit must be assigned the same way here as in the 1-bit path, // including for an exactly-zero residual component. one_bit_code_with_factor() uses @@ -182,9 +182,9 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingOneBitPlusExBits) { // Random gaussian data never produces an exact zero, so the other tests here // cannot catch a divergence -- this one constructs the case directly, by // putting the centroid exactly on the data value in half the dimensions. -TEST(XyQuantization, SignBitMatchesOneBitPathOnZeroResiduals) { +TEST(XyQuantizationTest, SignBitMatchesOneBitPathOnZeroResiduals) { constexpr size_t kDim = 64; - constexpr size_t kExtraBits = 3; + constexpr size_t kExBits = 3; std::vector centroid(kDim, 1.0F); std::vector data(kDim); @@ -214,20 +214,20 @@ TEST(XyQuantization, SignBitMatchesOneBitPathOnZeroResiduals) { std::vector total_code(kDim, -1); quant::rabitq_impl::total_bits::combined_code( - residual.data(), kDim, 1 + kExtraBits, total_code.data() + residual.data(), kDim, 1 + kExBits, total_code.data() ); for (size_t j = 0; j < kDim; ++j) { - EXPECT_EQ(total_code[j] >> kExtraBits, ref_bin[j]) << "dim " << j; + EXPECT_EQ(total_code[j] >> kExBits, ref_bin[j]) << "dim " << j; } } // A residual whose squared norm is zero takes the zero-residual path of the other // encoders; the split code must reassemble to the same full code. The second case // has nonzero components whose float squares underflow to zero. -TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { +TEST(XyQuantizationTest, ZeroNormResidualMatchesFullCode) { constexpr size_t kDim = 64; - constexpr size_t kExtraBits = 3; + constexpr size_t kExBits = 3; std::vector data(kDim, 1.0F); std::vector tiny(kDim, 0.0F); @@ -244,7 +244,7 @@ TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { SCOPED_TRACE(base_bits); SCOPED_TRACE(metric); XySplitResult res = - SplitCode(input, centroid, kDim, base_bits, kExtraBits, metric); + split_code(input, centroid, kDim, base_bits, kExBits, metric); std::vector full_code(kDim); float full_add = 0, full_rescale = 0, full_error = 0; @@ -252,7 +252,7 @@ TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { input.data(), centroid.data(), kDim, - base_bits + kExtraBits, + base_bits + kExBits, full_code.data(), full_add, full_rescale, @@ -261,7 +261,7 @@ TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { ); for (size_t j = 0; j < kDim; ++j) { const uint32_t total = - (static_cast(res.base_code[j]) << kExtraBits) | + (static_cast(res.base_code[j]) << kExBits) | static_cast(res.extra_code[j]); EXPECT_EQ(total, full_code[j]) << "dim " << j; } @@ -271,9 +271,9 @@ TEST(XyQuantization, ZeroNormResidualMatchesFullCode) { } } -TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) { +TEST(XyQuantizationTest, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) { constexpr size_t kDim = 64; - constexpr size_t kExtraBits = 3; + constexpr size_t kExBits = 3; std::vector centroid(kDim, 1.0F); std::vector data(kDim); @@ -308,7 +308,7 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) data.data(), centroid.data(), kDim, - kExtraBits, + kExBits, reference_extra_code.data(), reference_full_add, reference_full_rescale, @@ -317,7 +317,7 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) ); XySplitResult result = - SplitCode(data, centroid, kDim, /*base_bits=*/1, kExtraBits, metric); + split_code(data, centroid, kDim, /*base_bits=*/1, kExBits, metric); for (size_t dim = 0; dim < kDim; ++dim) { EXPECT_EQ(result.base_code[dim], reference_base_code[dim]); @@ -340,7 +340,7 @@ TEST(XyQuantization, BaseBitsOneMatchesExistingFactorsOnZeroResidualComponents) // Goes through quantize_xy_single rather than code_factors directly: the // combined_code-level test above cannot reach this, which is how it was // missed. -TEST(XyQuantization, ZeroResidualGivesFiniteFactors) { +TEST(XyQuantizationTest, ZeroResidualGivesFiniteFactors) { constexpr size_t kDim = 64; std::vector centroid(kDim, 0.7F); std::vector data = centroid; // residual is exactly zero @@ -384,7 +384,7 @@ TEST(XyQuantization, ZeroResidualGivesFiniteFactors) { } } -TEST(XyQuantization, CollinearResidualGivesFiniteErrorFactor) { +TEST(XyQuantizationTest, CollinearResidualGivesFiniteErrorFactor) { constexpr size_t kDim = 64; constexpr size_t kBaseBits = 2; @@ -419,18 +419,18 @@ TEST(XyQuantization, CollinearResidualGivesFiniteErrorFactor) { } // The base code must be stored exactly once: total storage is -// (base_bits + extra_bits) bits per dimension plus 5 floats (3 base factors + -// 2 full factors), never (2*base_bits + extra_bits) as an independently +// (base_bits + ex_bits) bits per dimension plus 5 floats (3 base factors + +// 2 full factors), never (2*base_bits + ex_bits) as an independently // quantized second layer would cost. -TEST(XyQuantization, StorageCostsBasePlusExtraBitsOnly) { +TEST(XyQuantizationTest, StorageCostsBasePlusExtraBitsOnly) { constexpr size_t kDim = 128; constexpr size_t kBaseBits = 3; - constexpr size_t kExtraBits = 4; + constexpr size_t kExBits = 4; size_t base_bytes = BaseDataMap::data_bytes(kDim, kBaseBits); - size_t extra_bytes = ExDataMap::data_bytes(kDim, kExtraBits); + size_t extra_bytes = ExDataMap::data_bytes(kDim, kExBits); - EXPECT_EQ(base_bytes + extra_bytes, (kDim * (kBaseBits + kExtraBits) / 8) + (5 * 4)); + EXPECT_EQ(base_bytes + extra_bytes, (kDim * (kBaseBits + kExBits) / 8) + (5 * 4)); EXPECT_EQ(ExDataMap::data_bytes(kDim, 0), 0U); } @@ -440,14 +440,14 @@ TEST(XyQuantization, StorageCostsBasePlusExtraBitsOnly) { // plain reference dot product in place of the popcount kernel for ip_x0_qr. // xy_single_full_dist must land on the same estimate and, like // split_single_fulldist_direct, bound it with f_error_base / 2^ex_bits. -TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { +TEST(XyQuantizationTest, EndToEndMatchesExistingFormulaAtBaseBitsOne) { constexpr size_t kDim = 128; - constexpr size_t kExtraBits = 5; + constexpr size_t kExBits = 5; std::mt19937 gen(123); - std::vector data = RandomVec(kDim, gen); - std::vector centroid = RandomVec(kDim, gen); - std::vector query = RandomVec(kDim, gen); + std::vector data = random_vec(kDim, gen); + std::vector centroid = random_vec(kDim, gen); + std::vector query = random_vec(kDim, gen); constexpr float kGAdd = 3.5F; constexpr float kGError = 0.75F; @@ -477,7 +477,7 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { data.data(), centroid.data(), kDim, - kExtraBits, + kExBits, ref_ex_code.data(), ref_f_add_ex, ref_f_rescale_ex, @@ -493,7 +493,7 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { ex_ip += static_cast(query[i]) * static_cast(ref_ex_code[i]); sumq += query[i]; } - double ref_c_b = -(static_cast(1 << (kExtraBits + 1)) - 1) / 2.0; + double ref_c_b = -(static_cast(1 << (kExBits + 1)) - 1) / 2.0; double ref_kbxsumq = sumq * ref_c_b; // 1-bit stage, as split_single_estdist_direct computes it. @@ -505,28 +505,28 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { float ref_full_est = static_cast( ref_f_add_ex + kGAdd + (ref_f_rescale_ex * - (static_cast(1 << kExtraBits) * ip_x0_qr + ex_ip + ref_kbxsumq)) + (static_cast(1 << kExBits) * ip_x0_qr + ex_ip + ref_kbxsumq)) ); // --- New path --- std::vector base_data(BaseDataMap::data_bytes(kDim, 1)); - std::vector extra_data(ExDataMap::data_bytes(kDim, kExtraBits)); + std::vector extra_data(ExDataMap::data_bytes(kDim, kExBits)); quant::quantize_xy_single( data.data(), centroid.data(), kDim, /*base_bits=*/1, - kExtraBits, + kExBits, base_data.data(), extra_data.data() ); SplitSingleQuery q_obj( - query.data(), kDim, kExtraBits, quant::RabitqConfig(), METRIC_L2, /*base_bits=*/1 + query.data(), kDim, kExBits, quant::RabitqConfig(), METRIC_L2, /*base_bits=*/1 ); auto base_ip_func = select_excode_ipfunc(1); - auto extra_ip_func = select_excode_ipfunc(kExtraBits); + auto extra_ip_func = select_excode_ipfunc(kExBits); float ip_base = 0; float base_est = 0; @@ -552,7 +552,7 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { // it as a parameter; there is no x+y-specific boosting function. q_obj.set_g_add(std::sqrt(kGAdd)); float boosted_est = split_distance_boosting( - extra_data.data(), extra_ip_func, q_obj, kDim, kExtraBits, ip_base + extra_data.data(), extra_ip_func, q_obj, kDim, kExBits, ip_base ); EXPECT_NEAR(boosted_est, ref_full_est, 1e-2F); @@ -570,7 +570,7 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { q_obj, kDim, /*base_bits=*/1, - kExtraBits, + kExBits, full_est, full_low, full_ip_base, @@ -582,25 +582,25 @@ TEST(XyQuantization, EndToEndMatchesExistingFormulaAtBaseBitsOne) { EXPECT_FLOAT_EQ(full_est, boosted_est); EXPECT_NEAR( full_low, - full_est - (ref_f_error_bin * kGError / static_cast(1 << kExtraBits)), + full_est - (ref_f_error_bin * kGError / static_cast(1 << kExBits)), 1e-4F ); } // The boosted estimate must equal the estimate you would get by reading the -// whole (base_bits+extra_bits)-bit code directly -- i.e. reusing the filter +// whole (base_bits+ex_bits)-bit code directly -- i.e. reusing the filter // stage's base inner product loses nothing. This is the property the old // two-layer implementation could not have, since its two layers held // different codes. -TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { +TEST(XyQuantizationTest, BoostedEstimateMatchesUnsplitCombinedCode) { constexpr size_t kDim = 128; constexpr size_t kBaseBits = 3; - constexpr size_t kExtraBits = 4; + constexpr size_t kExBits = 4; std::mt19937 gen(2024); - std::vector data = RandomVec(kDim, gen); - std::vector centroid = RandomVec(kDim, gen); - std::vector query = RandomVec(kDim, gen); + std::vector data = random_vec(kDim, gen); + std::vector centroid = random_vec(kDim, gen); + std::vector query = random_vec(kDim, gen); // g_add / g_error as the L2 search path actually supplies them // (||q - centroid||^2 and ||q - centroid||), so the filter's lower bound @@ -610,10 +610,10 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { float d = query[i] - centroid[i]; q_to_cent_sqr += d * d; } - const float kGError = std::sqrt(q_to_cent_sqr); - const float kGAdd = q_to_cent_sqr; + const float g_error = std::sqrt(q_to_cent_sqr); + const float g_add = q_to_cent_sqr; - XySplitResult res = SplitCode(data, centroid, kDim, kBaseBits, kExtraBits); + XySplitResult res = split_code(data, centroid, kDim, kBaseBits, kExBits); // Encode the unsplit reference independently of XY splitting and packed reads. std::vector full_code(kDim); @@ -622,7 +622,7 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { data.data(), centroid.data(), kDim, - kBaseBits + kExtraBits, + kBaseBits + kExBits, full_code.data(), full_add, full_rescale, @@ -632,30 +632,30 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { double total_ip = 0; double sumq = 0; for (size_t i = 0; i < kDim; ++i) { - uint32_t total = (static_cast(res.base_code[i]) << kExtraBits) | + uint32_t total = (static_cast(res.base_code[i]) << kExBits) | static_cast(res.extra_code[i]); EXPECT_EQ(total, full_code[i]) << "dim " << i; total_ip += static_cast(query[i]) * full_code[i]; sumq += query[i]; } - double c_b = -(static_cast(1 << (kBaseBits + kExtraBits)) - 1) / 2.0; + double c_b = -(static_cast(1 << (kBaseBits + kExBits)) - 1) / 2.0; float ref_est = - static_cast(full_add + kGAdd + (full_rescale * (total_ip + (sumq * c_b)))); + static_cast(full_add + g_add + (full_rescale * (total_ip + (sumq * c_b)))); std::vector base_data(BaseDataMap::data_bytes(kDim, kBaseBits)); - std::vector extra_data(ExDataMap::data_bytes(kDim, kExtraBits)); + std::vector extra_data(ExDataMap::data_bytes(kDim, kExBits)); quant::quantize_xy_single( data.data(), centroid.data(), kDim, kBaseBits, - kExtraBits, + kExBits, base_data.data(), extra_data.data() ); SplitSingleQuery q_obj( - query.data(), kDim, kExtraBits, quant::RabitqConfig(), METRIC_L2, kBaseBits + query.data(), kDim, kExBits, quant::RabitqConfig(), METRIC_L2, kBaseBits ); float ip_base = 0; @@ -670,18 +670,13 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { ip_base, base_est, base_low, - kGAdd, - kGError + g_add, + g_error ); - q_obj.set_g_add(std::sqrt(kGAdd)); + q_obj.set_g_add(std::sqrt(g_add)); float full_est = split_distance_boosting( - extra_data.data(), - select_excode_ipfunc(kExtraBits), - q_obj, - kDim, - kExtraBits, - ip_base + extra_data.data(), select_excode_ipfunc(kExBits), q_obj, kDim, kExBits, ip_base ); EXPECT_NEAR(full_est, ref_est, 1e-2F); @@ -691,19 +686,19 @@ TEST(XyQuantization, BoostedEstimateMatchesUnsplitCombinedCode) { EXPECT_LE(base_low, full_est); } -TEST(XyQuantization, CombinedBitsBeyondCapThrows) { +TEST(XyQuantizationTest, CombinedBitsBeyondCapThrows) { constexpr size_t kDim = 64; std::mt19937 gen(1); - std::vector data = RandomVec(kDim, gen); + std::vector data = random_vec(kDim, gen); std::vector centroid(kDim, 0.0F); auto call_with_bad_bits = [&]() { - SplitCode(data, centroid, kDim, /*base_bits=*/8, /*extra_bits=*/8); + split_code(data, centroid, kDim, /*base_bits=*/8, /*ex_bits=*/8); }; EXPECT_THROW(call_with_bad_bits(), std::invalid_argument); } -TEST(XyQuantization, SplitSingleQueryRejectsInvalidBitWidths) { +TEST(XyQuantizationTest, SplitSingleQueryRejectsInvalidBitWidths) { constexpr size_t kDim = 64; std::vector query(kDim, 1.0F); diff --git a/tests/unit/rabitqlib/utils/bitops_test.cpp b/tests/unit/rabitqlib/utils/bitops_test.cpp index c02d2427..71c909fb 100644 --- a/tests/unit/rabitqlib/utils/bitops_test.cpp +++ b/tests/unit/rabitqlib/utils/bitops_test.cpp @@ -2,7 +2,10 @@ #include +#include #include +#include +#include namespace rabitqlib::bitops { namespace { @@ -16,5 +19,44 @@ TEST(BitopsTest, CountsBitsAndTrailingZeros) { } } +TEST(BitopsTest, Popcount64MatchesBitReference) { + const auto check = [](uint64_t value) { + unsigned expected = 0; + for (size_t bit = 0; bit < 64; ++bit) { + expected += static_cast((value >> bit) & 1U); + } + EXPECT_EQ(popcount64(value), expected) << value; + }; + check(0); + check(std::numeric_limits::max()); + check(0x8000000000000001ULL); + for (size_t bit = 0; bit < 64; ++bit) { + check(uint64_t{1} << bit); + check(~(uint64_t{1} << bit)); + } + uint64_t value = 17; + for (size_t i = 0; i < 1024; ++i) { + value = value * 6364136223846793005ULL + 1442695040888963407ULL; + check(value); + } +} + +TEST(BitopsTest, CompactWordLoadsZeroExtendWithoutReadingPastInput) { + for (size_t bytes = 1; bytes <= 8; ++bytes) { + for (size_t offset : {0U, 1U}) { + SCOPED_TRACE(bytes); + SCOPED_TRACE(offset); + std::vector storage(bytes + offset); + uint64_t expected = 0; + for (size_t i = 0; i < bytes; ++i) { + const auto value = static_cast(0x80U + i * 17U); + storage[offset + i] = value; + expected |= uint64_t{value} << (8 * i); + } + EXPECT_EQ(load_word(storage.data() + offset, bytes), expected); + } + } +} + } // namespace } // namespace rabitqlib::bitops diff --git a/tests/unit/rabitqlib/utils/cpu_features_test.cpp b/tests/unit/rabitqlib/utils/cpu_features_test.cpp index 1beeb51a..b33bd68e 100644 --- a/tests/unit/rabitqlib/utils/cpu_features_test.cpp +++ b/tests/unit/rabitqlib/utils/cpu_features_test.cpp @@ -22,7 +22,7 @@ rabitqlib::cpu::Features all_hardware_features() { }; } -TEST(CpuFeatures, ReturnsStableResult) { +TEST(CpuFeaturesTest, ReturnsStableResult) { const auto& first = rabitqlib::cpu::features(); const auto& second = rabitqlib::cpu::features(); @@ -34,7 +34,7 @@ TEST(CpuFeatures, ReturnsStableResult) { } // The emulated CI runs must fail instead of silently skipping optional backends. -TEST(CpuFeatures, MatchesRequiredEmulatedCpu) { +TEST(CpuFeaturesTest, MatchesRequiredEmulatedCpu) { const bool avx2 = rabitqlib::cpu::has_avx2(); const bool avx512 = rabitqlib::cpu::has_avx512_core(); const bool popcnt = rabitqlib::cpu::has_avx512_popcnt(); @@ -55,7 +55,7 @@ TEST(CpuFeatures, MatchesRequiredEmulatedCpu) { EXPECT_EQ(popcnt, require_avx512); } -TEST(CpuFeatures, RequiresOsManagedAvxState) { +TEST(CpuFeaturesTest, RequiresOsManagedAvxState) { constexpr uint64_t kAvxState = 0x6; const auto hardware = all_hardware_features(); @@ -77,7 +77,7 @@ TEST(CpuFeatures, RequiresOsManagedAvxState) { EXPECT_FALSE(usable.avx512f); } -TEST(CpuFeatures, EnablesAvx2WithoutAvx512State) { +TEST(CpuFeaturesTest, EnablesAvx2WithoutAvx512State) { constexpr uint64_t kAvxState = 0x6; const auto usable = rabitqlib::cpu::detail::filter_usable_features( all_hardware_features(), true, true, kAvxState @@ -93,7 +93,7 @@ TEST(CpuFeatures, EnablesAvx2WithoutAvx512State) { EXPECT_FALSE(usable.avx512cd); } -TEST(CpuFeatures, EnablesAvx512OnlyWithCompleteAvx512State) { +TEST(CpuFeaturesTest, EnablesAvx512OnlyWithCompleteAvx512State) { constexpr uint64_t kAvx512State = 0xe6; const auto usable = rabitqlib::cpu::detail::filter_usable_features( all_hardware_features(), true, true, kAvx512State @@ -124,7 +124,7 @@ TEST(CpuFeatures, EnablesAvx512OnlyWithCompleteAvx512State) { } } -TEST(CpuFeatures, Avx512SelectionRequiresEveryCompilerEnabledFeature) { +TEST(CpuFeaturesTest, Avx512SelectionRequiresEveryCompilerEnabledFeature) { using rabitqlib::cpu::Features; using rabitqlib::cpu::detail::supports_avx512_core; EXPECT_TRUE(supports_avx512_core(all_hardware_features())); @@ -151,7 +151,7 @@ TEST(CpuFeatures, Avx512SelectionRequiresEveryCompilerEnabledFeature) { } // namespace -TEST(CpuFeatures, NativeArm64UsesNeonWithoutX86Features) { +TEST(CpuFeaturesTest, NativeArm64UsesNeonWithoutX86Features) { #if defined(__aarch64__) EXPECT_TRUE(rabitqlib::cpu::has_neon()); EXPECT_FALSE(rabitqlib::cpu::has_avx2()); diff --git a/tests/unit/rabitqlib/utils/excode_ip_test.cpp b/tests/unit/rabitqlib/utils/excode_ip_test.cpp index 874f1c9c..7e5de84a 100644 --- a/tests/unit/rabitqlib/utils/excode_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/excode_ip_test.cpp @@ -12,7 +12,7 @@ #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" -TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { +TEST(ExcodeIpTest, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { using namespace rabitqlib; const std::array wrappers{ diff --git a/tests/unit/rabitqlib/utils/float_distance_test.cpp b/tests/unit/rabitqlib/utils/float_distance_test.cpp index 3e6fee45..68ebea49 100644 --- a/tests/unit/rabitqlib/utils/float_distance_test.cpp +++ b/tests/unit/rabitqlib/utils/float_distance_test.cpp @@ -76,7 +76,7 @@ std::vector backends() { } // namespace -TEST(FloatDistance, EmptyAndAliasedInputsPreserveDistanceConventions) { +TEST(FloatDistanceTest, EmptyAndAliasedInputsPreserveDistanceConventions) { const std::array a{1, -2, 3, -4, 5}; for (const auto& backend : backends()) { SCOPED_TRACE(backend.name); @@ -91,7 +91,7 @@ TEST(FloatDistance, EmptyAndAliasedInputsPreserveDistanceConventions) { } } -TEST(FloatDistance, SubtractionUsesFloatPrecisionInVectorBlocksAndTails) { +TEST(FloatDistanceTest, SubtractionUsesFloatPrecisionInVectorBlocksAndTails) { // In float, 1 - 2^-25 rounds to 1 before squaring. Widening the subtraction // to double would instead return the float immediately below 1. for (size_t dim : {1U, 4U, 16U, 17U, 33U}) { @@ -109,14 +109,15 @@ TEST(FloatDistance, SubtractionUsesFloatPrecisionInVectorBlocksAndTails) { } } -TEST(FloatDistance, AccumulationRetainsFloatPrecision) { +TEST(FloatDistanceTest, AccumulationRetainsFloatPrecision) { // Each small square is below half an ulp at 1. Place them in the same // accumulation lane for all backends. Double accumulation would preserve // their combined 2^-23 contribution and return the float above 1. std::array a{}, zero{}; a[0] = 1.0F; - for (size_t i = 64; i < a.size(); i += 64) + for (size_t i = 64; i < a.size(); i += 64) { a[i] = std::ldexp(1.0F, -13); + } for (const auto& backend : backends()) { SCOPED_TRACE(backend.name); EXPECT_EQ(backend.l2(a.data(), zero.data(), a.size()), 1.0F); @@ -126,7 +127,7 @@ TEST(FloatDistance, AccumulationRetainsFloatPrecision) { } } -TEST(FloatDistance, DoubleTemplatesRetainDoublePrecision) { +TEST(FloatDistanceTest, DoubleTemplatesRetainDoublePrecision) { const std::array a{1.0000000001, 2.0000000001, 3.0000000001}; const std::array b{1, 2, 3}; EXPECT_NEAR(rabitqlib::euclidean_sqr(a.data(), b.data(), a.size()), 3e-20, 1e-26); @@ -138,7 +139,7 @@ TEST(FloatDistance, DoubleTemplatesRetainDoublePrecision) { } #if defined(__linux__) -TEST(FloatDistance, TailsDoNotReadPastGuardPage) { +TEST(FloatDistanceTest, TailsDoNotReadPastGuardPage) { const size_t page = static_cast(sysconf(_SC_PAGESIZE)); void* memory = mmap(nullptr, page * 2, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); diff --git a/tests/unit/rabitqlib/utils/mask_ip_test.cpp b/tests/unit/rabitqlib/utils/mask_ip_test.cpp index 57e28eb2..ab02c994 100644 --- a/tests/unit/rabitqlib/utils/mask_ip_test.cpp +++ b/tests/unit/rabitqlib/utils/mask_ip_test.cpp @@ -11,7 +11,7 @@ #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/space.hpp" -TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { +TEST(MaskIpX0QTest, BackendsMatchScalarAcrossBlocksAndAlignments) { using namespace rabitqlib; using Function = float (*)(const float*, const uint8_t*, size_t); std::vector functions{ @@ -78,17 +78,17 @@ TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { } } -TEST(MaskIpX0Q, DispatchPreservesEveryStoredBitPosition) { +TEST(MaskIpX0QTest, DispatchPreservesEveryStoredBitPosition) { using namespace rabitqlib; - constexpr size_t dim = 192; - std::vector query(dim); - for (size_t i = 0; i < dim; ++i) { + constexpr size_t kDim = 192; + std::vector query(kDim); + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast(i + 1); } - for (size_t bit = 0; bit < dim; ++bit) { + for (size_t bit = 0; bit < kDim; ++bit) { SCOPED_TRACE(bit); - std::vector words((dim + 63) / 64, 0); + std::vector words((kDim + 63) / 64, 0); words[bit / 64] = uint64_t{1} << (63 - bit % 64); - EXPECT_EQ(mask_ip_x0_q(query.data(), words.data(), dim), query[bit]); + EXPECT_EQ(mask_ip_x0_q(query.data(), words.data(), kDim), query[bit]); } } diff --git a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp index ceee33c2..928c9c1b 100644 --- a/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp +++ b/tests/unit/rabitqlib/utils/matrix_dispatch_test.cpp @@ -234,24 +234,24 @@ TEST(MatrixDispatchTest, IndexedAccumulationMatchesOrderedReferenceWithUnaligned } TEST(MatrixDispatchTest, IndexedAccumulationPreservesPointOrderUnderCancellation) { - constexpr size_t kDimension = 65; + constexpr size_t kDim = 65; const std::vector values{1e7F, 1e7F, 1e-10F, -1e7F, -1e7F, 1e-10F}; const std::vector labels{0, 1, 0, 1, 0, 1}; - std::vector data(values.size() * kDimension); + std::vector data(values.size() * kDim); std::vector point_ids(values.size()); std::iota(point_ids.begin(), point_ids.end(), size_t{0}); for (size_t point = 0; point < values.size(); ++point) { - std::fill_n(data.data() + point * kDimension, kDimension, values[point]); + std::fill_n(data.data() + point * kDim, kDim, values[point]); } for (const auto& backend : accumulation_backends()) { SCOPED_TRACE(backend.first); - std::vector sums(2 * kDimension); + std::vector sums(2 * kDim); std::vector counts(2); backend.second( data.data(), point_ids.data(), 3, - kDimension, + kDim, labels.data(), sums.data(), counts.data() @@ -260,15 +260,15 @@ TEST(MatrixDispatchTest, IndexedAccumulationPreservesPointOrderUnderCancellation data.data(), point_ids.data() + 3, 3, - kDimension, + kDim, labels.data(), sums.data(), counts.data() ); EXPECT_EQ(counts, (std::vector{3, 3})); - for (size_t coordinate = 0; coordinate < kDimension; ++coordinate) { + for (size_t coordinate = 0; coordinate < kDim; ++coordinate) { EXPECT_DOUBLE_EQ(sums[coordinate], 0.0); - EXPECT_DOUBLE_EQ(sums[kDimension + coordinate], static_cast(1e-10F)); + EXPECT_DOUBLE_EQ(sums[kDim + coordinate], static_cast(1e-10F)); } } } diff --git a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp index f41e3cb5..7a210fc4 100644 --- a/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp +++ b/tests/unit/rabitqlib/utils/neon_hnsw_test.cpp @@ -19,7 +19,7 @@ #include "rabitqlib/utils/cpu_features.hpp" namespace rabitqlib { -TEST(NeonHnsw, WarmupHandlesUnalignedCodesTailsAndWideCounts) { +TEST(HnswNeonTest, WarmupHandlesUnalignedCodesTailsAndWideCounts) { if (!cpu::has_neon()) { GTEST_SKIP(); } @@ -78,7 +78,7 @@ TEST(NeonHnsw, WarmupHandlesUnalignedCodesTailsAndWideCounts) { ); } -TEST(NeonHnsw, SearchMatchesGenericOnSameGraphForBothMetricsAndAllWidths) { +TEST(HnswNeonTest, SearchMatchesGenericOnSameGraphForBothMetricsAndAllWidths) { if (!cpu::has_neon()) { GTEST_SKIP(); } diff --git a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp index fe091d6e..3ffd1716 100644 --- a/tests/unit/rabitqlib/utils/portable_kernel_test.cpp +++ b/tests/unit/rabitqlib/utils/portable_kernel_test.cpp @@ -16,7 +16,7 @@ #include "rabitqlib/utils/warmup_space.hpp" namespace rabitqlib { -TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { +TEST(PortableKernelsTest, PackingMatchesDispatchedBytesAndDotProducts) { using Pack = void (*)(const uint8_t*, uint8_t*, size_t); const std::array pack{ simd::packing_2bit_excode_generic, @@ -70,7 +70,7 @@ TEST(PortableKernels, PackingMatchesDispatchedBytesAndDotProducts) { } } -TEST(PortableKernels, QueryTransposeAndWarmupMatchIntegerReference) { +TEST(PortableKernelsTest, QueryTransposeAndWarmupMatchIntegerReference) { for (size_t dim : {64U, 96U, 128U, 448U, 512U, 544U, 576U, 960U, 1024U, 1088U}) { for (size_t bits = 0; bits <= 8; ++bits) { SCOPED_TRACE(dim); @@ -130,11 +130,12 @@ TEST(PortableKernels, QueryTransposeAndWarmupMatchIntegerReference) { ); } -TEST(PortableKernels, Uint16TransposePreservesAllBitPlanes) { +TEST(PortableKernelsTest, Uint16TransposePreservesAllBitPlanes) { constexpr size_t kDim = 192; std::array query{}; - for (size_t i = 0; i < kDim; ++i) + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast(i * 977); + } for (size_t bits = 1; bits <= 16; ++bits) { std::vector actual(kDim / 64 * bits), scalar(actual.size()); new_transpose_bin(query.data(), actual.data(), kDim, bits); @@ -151,15 +152,16 @@ TEST(PortableKernels, Uint16TransposePreservesAllBitPlanes) { } } // namespace rabitqlib -TEST(PortableKernels, ExtraBitTailsPreserveFullBlockBytesAndUseCompactBitStream) { +TEST(PortableKernelsTest, ExtraBitTailsPreserveFullBlockBytesAndUseCompactBitStream) { using namespace rabitqlib; for (size_t bits : {2U, 3U, 5U, 6U, 7U}) { - constexpr size_t tail = 32; - const size_t dim = 128 + tail; + constexpr size_t kTail = 32; + const size_t dim = 128 + kTail; std::vector raw(dim), packed(dim * bits / 8 + 2, 0xA5); std::vector full(128 * bits / 8); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { raw[i] = (i * 37 + 3) & ((1U << bits) - 1); + } quant::rabitq_impl::ex_bits::packing_rabitqplus_code( raw.data(), full.data(), 128, bits ); @@ -171,7 +173,7 @@ TEST(PortableKernels, ExtraBitTailsPreserveFullBlockBytesAndUseCompactBitStream) EXPECT_TRUE(std::equal(full.begin(), full.end(), packed.begin() + 1)); // Independent byte-by-byte definition: bit j holds bit j % bits // of coordinate j / bits, with no gaps or rounded word storage. - for (size_t byte = 0; byte < tail * bits / 8; ++byte) { + for (size_t byte = 0; byte < kTail * bits / 8; ++byte) { uint8_t expected = 0; for (size_t b = 0; b < 8; ++b) { const size_t pos = byte * 8 + b; @@ -182,7 +184,7 @@ TEST(PortableKernels, ExtraBitTailsPreserveFullBlockBytesAndUseCompactBitStream) } } -TEST(PortableKernels, PackingMatchesReferenceForMaximumAndIsolatedCodes) { +TEST(PortableKernelsTest, PackingMatchesReferenceForMaximumAndIsolatedCodes) { using namespace rabitqlib; using Pack = void (*)(const uint8_t*, uint8_t*, size_t); const std::array generic{ diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp index f581cf64..5d2ac21b 100644 --- a/tests/unit/rabitqlib/utils/query_preparation_test.cpp +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -151,7 +151,7 @@ void check_transpose( } } // namespace -TEST(QueryPreparation, QuantizationMatchesHalfAwayRoundingAndTails) { +TEST(QueryPreparationTest, QuantizationMatchesHalfAwayRoundingAndTails) { check_quantization(simd::scalar_quantize_uint8_generic); check_quantization(simd::scalar_quantize_uint16_generic); // Exercise the public templates as well as the kernels they dispatch to. @@ -175,7 +175,7 @@ TEST(QueryPreparation, QuantizationMatchesHalfAwayRoundingAndTails) { #endif } -TEST(QueryPreparation, ZeroRangeQuantizationMatchesAcrossBackendsAndTails) { +TEST(QueryPreparationTest, ZeroRangeQuantizationMatchesAcrossBackendsAndTails) { check_zero_range(simd::scalar_quantize_uint8_generic); check_zero_range(simd::scalar_quantize_uint16_generic); check_zero_range(scalar_quantize); @@ -201,7 +201,7 @@ TEST(QueryPreparation, ZeroRangeQuantizationMatchesAcrossBackendsAndTails) { #endif } -TEST(QueryPreparation, TranspositionMatchesCoordinateReferenceAndBounds) { +TEST(QueryPreparationTest, TranspositionMatchesCoordinateReferenceAndBounds) { check_transpose(simd::new_transpose_bin_generic, false); check_transpose(simd::new_transpose_bin_512_generic, true); check_transpose(new_transpose_bin, false); diff --git a/tests/unit/rabitqlib/utils/rotator_test.cpp b/tests/unit/rabitqlib/utils/rotator_test.cpp index 5eb2ef5c..663e210a 100644 --- a/tests/unit/rabitqlib/utils/rotator_test.cpp +++ b/tests/unit/rabitqlib/utils/rotator_test.cpp @@ -384,10 +384,12 @@ TEST(FlipSignTest, BackendsPreserveBitsAndUnalignedBuffers) { } #endif #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) + if (cpu::has_avx2()) { backends.push_back(simd::flip_sign_avx2); - if (cpu::has_avx512_core()) + } + if (cpu::has_avx512_core()) { backends.push_back(simd::flip_sign_avx512); + } #endif // Signed zeros, infinities, quiet NaN payloads, subnormals and finite values. const std::array values{ @@ -429,14 +431,17 @@ TEST(FlipSignTest, BackendsPreserveBitsAndUnalignedBuffers) { TEST(KacsWalkTest, BackendsMatchScalarAndRespectBufferBounds) { std::vector backends{simd::kacs_walk}; #if defined(__aarch64__) || defined(_M_ARM64) - if (cpu::has_neon()) + if (cpu::has_neon()) { backends.push_back(simd::kacs_walk_neon); + } #endif #if defined(__x86_64__) || defined(_M_X64) - if (cpu::has_avx2()) + if (cpu::has_avx2()) { backends.push_back(simd::kacs_walk_avx2); - if (cpu::has_avx512_core()) + } + if (cpu::has_avx512_core()) { backends.push_back(simd::kacs_walk_avx512); + } #endif for (size_t dim : {64U, @@ -463,8 +468,9 @@ TEST(KacsWalkTest, BackendsMatchScalarAndRespectBufferBounds) { std::vector> input( dim + offset + 1, 12345 ); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { input[i + offset] = std::sin(static_cast(i)); + } auto expected = input; simd::kacs_walk_generic(expected.data() + offset, dim); for (auto backend : backends) { @@ -478,24 +484,29 @@ TEST(KacsWalkTest, BackendsMatchScalarAndRespectBufferBounds) { #if defined(__aarch64__) || defined(_M_ARM64) TEST(FhtNeonTest, MatchesSavedScalarRotationInPlaceAndWithExtraPadding) { - if (!cpu::has_neon()) + if (!cpu::has_neon()) { GTEST_SKIP(); + } for (size_t dim : {64U, 65U, 127U, 128U, 193U, 960U, 1025U, 4097U, 65536U}) { for (size_t extra : {0U, 32U, 64U}) { const size_t padded = (dim + 31) / 32 * 32 + extra; - if (padded > 65536) + if (padded > 65536) { continue; + } SCOPED_TRACE(dim); SCOPED_TRACE(padded); size_t trunc = 64; - while (trunc * 2 <= dim) + while (trunc * 2 <= dim) { trunc *= 2; + } std::vector flip(padded / 2); - for (size_t i = 0; i < flip.size(); ++i) + for (size_t i = 0; i < flip.size(); ++i) { flip[i] = static_cast(i * 73 + i / 7); + } std::vector input(dim), expected(padded); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { input[i] = std::sin(static_cast(i) * 0.13F); + } simd::fht_rotate_generic( input.data(), expected.data(), @@ -537,13 +548,15 @@ TEST(FhtNeonTest, MatchesSavedScalarRotationInPlaceAndWithExtraPadding) { } TEST(NeonRotationTest, StandaloneKernelsHandleShortTails) { - if (!cpu::has_neon()) + if (!cpu::has_neon()) { GTEST_SKIP(); + } for (size_t dim = 0; dim <= 33; ++dim) { std::vector flip((dim + 7) / 8, 0xA5); std::vector input(dim + 2, 12345); - for (size_t i = 0; i < dim; ++i) + for (size_t i = 0; i < dim; ++i) { input[i + 1] = static_cast(i) - 16.5F; + } auto actual = input, expected = input; simd::flip_sign_generic(flip.data(), expected.data() + 1, dim); simd::flip_sign_neon(flip.data(), actual.data() + 1, dim); @@ -632,7 +645,7 @@ TEST(FhtDispatchTest, AvxBackendsMatchBitwiseForCompactTails) { } #endif -TEST(MatrixRotatorFactory, DoublePreservesNormAndRejectsFloatOnlyRotation) { +TEST(MatrixRotatorFactoryTest, DoublePreservesNormAndRejectsFloatOnlyRotation) { const size_t dim = 64; std::unique_ptr> rotator( choose_rotator(dim, RotatorType::MatrixRotator) diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index 19a6ce49..4005d4f3 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -7,19 +7,17 @@ #include #include #include -#include #include #include #include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/core/detail/simd/warmup_dispatch.hpp" -#include "rabitqlib/utils/bitops.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/warmup_space.hpp" using namespace rabitqlib; -TEST(SpaceHelpers, CentroidSupportsRequestedWorkerCounts) { +TEST(SpaceHelpersTest, CentroidSupportsRequestedWorkerCounts) { const std::array data{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; for (size_t workers : {0U, 1U, 2U, 64U}) { const auto centroid = compute_centroid(data.data(), 4, 3, workers); @@ -30,69 +28,30 @@ TEST(SpaceHelpers, CentroidSupportsRequestedWorkerCounts) { } } -TEST(BitOps, Popcount64MatchesBitReference) { - const auto check = [](uint64_t value) { - unsigned expected = 0; - for (size_t bit = 0; bit < 64; ++bit) { - expected += static_cast((value >> bit) & 1U); - } - EXPECT_EQ(bitops::popcount64(value), expected) << value; - }; - check(0); - check(std::numeric_limits::max()); - check(0x8000000000000001ULL); - for (size_t bit = 0; bit < 64; ++bit) { - check(uint64_t{1} << bit); - check(~(uint64_t{1} << bit)); - } - uint64_t value = 17; - for (size_t i = 0; i < 1024; ++i) { - value = value * 6364136223846793005ULL + 1442695040888963407ULL; - check(value); - } -} - -TEST(BitOps, CompactWordLoadsZeroExtendWithoutReadingPastInput) { - for (size_t bytes = 1; bytes <= 8; ++bytes) { - for (size_t offset : {0U, 1U}) { - SCOPED_TRACE(bytes); - SCOPED_TRACE(offset); - std::vector storage(bytes + offset); - uint64_t expected = 0; - for (size_t i = 0; i < bytes; ++i) { - const auto value = static_cast(0x80U + i * 17U); - storage[offset + i] = value; - expected |= uint64_t{value} << (8 * i); - } - EXPECT_EQ(bitops::load_word(storage.data() + offset, bytes), expected); - } - } -} - -TEST(PackBinary, SupportsUnalignedOutput) { - constexpr size_t dim = 128; - std::array binary_code{}; - for (size_t i = 0; i < dim; ++i) { +TEST(PackBinaryTest, SupportsUnalignedOutput) { + constexpr size_t kDim = 128; + std::array binary_code{}; + for (size_t i = 0; i < kDim; ++i) { binary_code[i] = (i % 3) == 0; } alignas(uint64_t) std::array storage{}; auto* output = storage.data() + 1; ASSERT_NE(reinterpret_cast(output) % alignof(uint64_t), 0U); - pack_binary_to_bytes(binary_code.data(), output, dim); + pack_binary_to_bytes(binary_code.data(), output, kDim); const std::array packed{ load_unaligned_u64(storage.data() + 1), load_unaligned_u64(storage.data() + 1 + sizeof(uint64_t)), }; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { const size_t word = i / 64; const auto bit = static_cast((packed[word] >> (63 - (i % 64))) & 1U); EXPECT_EQ(bit, binary_code[i]) << "bit " << i; } } -TEST(PackBinary, AllWordSizesPreserveFullAndPartialWordLayout) { +TEST(PackBinaryTest, AllWordSizesPreserveFullAndPartialWordLayout) { const auto check = [](auto word_type) { using Word = decltype(word_type); constexpr size_t kWordBits = sizeof(Word) * 8; @@ -126,12 +85,12 @@ TEST(PackBinary, AllWordSizesPreserveFullAndPartialWordLayout) { check(uint64_t{}); } -TEST(MaskIpX0Q, SupportsUnalignedCodes) { - constexpr size_t dim = 128; - std::array binary_code{}; - std::array query{}; +TEST(MaskIpX0QTest, SupportsUnalignedCodes) { + constexpr size_t kDim = 128; + std::array binary_code{}; + std::array query{}; float expected = 0; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { binary_code[i] = (i % 3) == 0; query[i] = static_cast(i + 1); expected += binary_code[i] != 0 ? query[i] : 0; @@ -140,24 +99,24 @@ TEST(MaskIpX0Q, SupportsUnalignedCodes) { alignas(uint64_t) std::array storage{}; auto* codes = storage.data() + 1; ASSERT_NE(reinterpret_cast(codes) % alignof(uint64_t), 0U); - pack_binary_to_bytes(binary_code.data(), codes, dim); + pack_binary_to_bytes(binary_code.data(), codes, kDim); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { - EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx2(query.data(), codes, dim), expected); + EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx2(query.data(), codes, kDim), expected); } #endif #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx512_core()) { - EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx512(query.data(), codes, dim), expected); + EXPECT_FLOAT_EQ(simd::mask_ip_x0_q_avx512(query.data(), codes, kDim), expected); } #endif - EXPECT_FLOAT_EQ(mask_ip_x0_q(query.data(), codes, dim), expected); + EXPECT_FLOAT_EQ(mask_ip_x0_q(query.data(), codes, kDim), expected); } -TEST(WarmupIpX0Q, SupportsUnalignedCodes) { - constexpr float delta = 0.5F; - constexpr float vl = -0.25F; +TEST(WarmupIpX0QTest, SupportsUnalignedCodes) { + constexpr float kDelta = 0.5F; + constexpr float kVl = -0.25F; for (size_t dim : {0UL, 64UL, @@ -217,16 +176,16 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { query_offset += chunks * b_query; } - const float expected = delta * static_cast(weighted_intersection) + - vl * static_cast(data_popcount); + const float expected = kDelta * static_cast(weighted_intersection) + + kVl * static_cast(data_popcount); EXPECT_FLOAT_EQ( - warmup_ip_x0_q_512(codes, query.data(), delta, vl, dim, b_query), expected + warmup_ip_x0_q_512(codes, query.data(), kDelta, kVl, dim, b_query), expected ); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_FLOAT_EQ( simd::warmup_ip_x0_q_512_avx2( - codes, query.data(), delta, vl, dim, b_query + codes, query.data(), kDelta, kVl, dim, b_query ), expected ); @@ -236,7 +195,7 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { if (cpu::has_avx512_popcnt()) { EXPECT_FLOAT_EQ( simd::warmup_ip_x0_q_512_avx512( - codes, query.data(), delta, vl, dim, b_query + codes, query.data(), kDelta, kVl, dim, b_query ), expected ); @@ -246,7 +205,7 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { } } -TEST(WarmupIpX0Q, RejectsQueryWidthsBeyondByte) { +TEST(WarmupIpX0QTest, RejectsQueryWidthsBeyondByte) { const std::array data{}; const std::array query{}; EXPECT_THROW( @@ -271,7 +230,7 @@ TEST(WarmupIpX0Q, RejectsQueryWidthsBeyondByte) { #endif } -TEST(Select_IP_Func, returns_stable_function_pointer) { +TEST(SelectExcodeIpFuncTest, ReturnsStableFunctionPointer) { for (size_t bits = 0; bits <= 8; ++bits) { SCOPED_TRACE(bits); const auto ip_func = select_excode_ipfunc(bits); @@ -289,12 +248,12 @@ TEST(Select_IP_Func, returns_stable_function_pointer) { } } -TEST(Select_IP_Func, zero_ex_bits_contributes_nothing) { - constexpr size_t dim = 64; - std::vector query(dim); - std::vector codes(dim / 8); +TEST(SelectExcodeIpFuncTest, ZeroExBitsContributesNothing) { + constexpr size_t kDim = 64; + std::vector query(kDim); + std::vector codes(kDim / 8); - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast(i) + 1.0F; } for (size_t i = 0; i < codes.size(); ++i) { @@ -303,52 +262,65 @@ TEST(Select_IP_Func, zero_ex_bits_contributes_nothing) { ex_ipfunc ip_func = select_excode_ipfunc(0); ASSERT_NE(ip_func, nullptr); - ASSERT_EQ(ip_func(query.data(), codes.data(), dim), 0.0F); + ASSERT_EQ(ip_func(query.data(), codes.data(), kDim), 0.0F); +} + +TEST(SelectExcodeIpFuncTest, RejectsExBitsAboveEight) { + try { + select_excode_ipfunc(9); + FAIL() << "ex_bits above 8 must be rejected"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "ex_bits must be in [0, 8]"); + } } -TEST(ip16_fxu1_avx, ip_works) { - constexpr size_t dim = 64; - float query[dim]; - uint8_t codes[dim / 8]; +TEST(LegacyExcodeIpTest, Ip16Fxu1AvxMatchesReference) { + constexpr size_t kDim = 64; + float query[kDim]; + uint8_t codes[kDim / 8]; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast((i * 37 + 11) % 101) * 0.25F; } - for (size_t i = 0; i < dim / 8; ++i) { + for (size_t i = 0; i < kDim / 8; ++i) { codes[i] = static_cast(i * 53 + 19); } float expected = 0.0F; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { expected += query[i] * static_cast((codes[i / 8] >> (i % 8)) & 1U); } - ASSERT_NEAR(rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, dim), expected, 0.1F); + ASSERT_NEAR( + rabitqlib::excode_ipimpl::ip16_fxu1_avx(query, codes, kDim), expected, 0.1F + ); } -TEST(ip64_fxu2_avx, ip_works) { - constexpr size_t dim = 64 * 4; - float query[dim]; - uint8_t codes[dim / 4]; +TEST(LegacyExcodeIpTest, Ip64Fxu2AvxMatchesReference) { + constexpr size_t kDim = 64 * 4; + float query[kDim]; + uint8_t codes[kDim / 4]; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { query[i] = static_cast((i * 37 + 11) % 101) * 0.25F; } - for (size_t i = 0; i < dim / 4; ++i) { + for (size_t i = 0; i < kDim / 4; ++i) { codes[i] = static_cast(i * 53 + 19); } float expected = 0.0F; - for (size_t i = 0; i < dim; ++i) { + for (size_t i = 0; i < kDim; ++i) { const uint8_t packed = codes[(i / 64) * 16 + (i % 16)]; const auto code = static_cast((packed >> (2 * ((i % 64) / 16))) & 3U); expected += query[i] * static_cast(code); } - ASSERT_NEAR(rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, dim), expected, 0.1F); + ASSERT_NEAR( + rabitqlib::excode_ipimpl::ip64_fxu2_avx(query, codes, kDim), expected, 0.1F + ); } -TEST(PackBinary, CompactTailWordsMatchBinaryArithmetic) { +TEST(PackBinaryTest, CompactTailWordsMatchBinaryArithmetic) { for (size_t dim : {64U, 96U, 128U, 160U}) { std::vector signs(dim); std::vector values(dim); diff --git a/tests/unit/rabitqlib/utils/visited_set_test.cpp b/tests/unit/rabitqlib/utils/visited_set_test.cpp index 1768369c..d3968b38 100644 --- a/tests/unit/rabitqlib/utils/visited_set_test.cpp +++ b/tests/unit/rabitqlib/utils/visited_set_test.cpp @@ -33,17 +33,17 @@ static_assert( "HashBasedBooleanSet must remain available as a compatibility alias" ); -TEST(VisitedListPoolTest, TransfersOwnershipAndClearsReusedStorage) { - rabitqlib::VisitedListPool pool(1, 64); - auto visited = pool.get_free_vislist(); +TEST(VisitedSetPoolTest, TransfersOwnershipAndClearsReusedStorage) { + rabitqlib::VisitedSetPool pool(1, 64); + auto visited = pool.acquire(); auto* storage = visited.get(); visited->set(7); - pool.release_vis_list(std::move(visited)); + pool.release(std::move(visited)); EXPECT_EQ(visited, nullptr); - auto reused = pool.get_free_vislist(); + auto reused = pool.acquire(); EXPECT_EQ(reused.get(), storage); EXPECT_FALSE(reused->get(7)); - auto another = pool.get_free_vislist(); + auto another = pool.acquire(); EXPECT_NE(another.get(), reused.get()); } @@ -75,15 +75,15 @@ void expect_basic_visited_set_behavior() { EXPECT_TRUE(visited.get(23)); } -TEST(VisitedSet, HashBackendSatisfiesContract) { +TEST(VisitedSetTest, HashBackendSatisfiesContract) { expect_basic_visited_set_behavior(); } -TEST(VisitedSet, EpochBackendSatisfiesContract) { +TEST(VisitedSetTest, EpochBackendSatisfiesContract) { expect_basic_visited_set_behavior(); } -TEST(VisitedSet, BackendsRemainBehaviorallyEquivalent) { +TEST(VisitedSetTest, BackendsRemainBehaviorallyEquivalent) { constexpr size_t kNumElements = 4096; constexpr size_t kNumOperations = 20000; @@ -117,7 +117,7 @@ TEST(VisitedSet, BackendsRemainBehaviorallyEquivalent) { } } -TEST(HashBasedVisitedSet, HandlesCollidingIds) { +TEST(HashBasedVisitedSetTest, HandlesCollidingIds) { HashBasedVisitedSet visited(1024, 1); const std::vector colliding_ids{0, 32, 64, 96, 128, 1023}; @@ -134,7 +134,7 @@ TEST(HashBasedVisitedSet, HandlesCollidingIds) { } } -TEST(EpochBasedVisitedSet, ClearsStaleValuesWhenEpochWraps) { +TEST(EpochBasedVisitedSetTest, ClearsStaleValuesWhenEpochWraps) { EpochBasedVisitedSet visited(4); visited.set(0);