diff --git a/src/core/quantizer/rotator/matrix_rotator.cc b/src/core/quantizer/rotator/matrix_rotator.cc index dc2053323..f75152eec 100644 --- a/src/core/quantizer/rotator/matrix_rotator.cc +++ b/src/core/quantizer/rotator/matrix_rotator.cc @@ -124,10 +124,16 @@ int MatrixRotator::init_impl(size_t dim) { std::vector Q(dim * dim); householder_qr(rand_mat.data(), Q.data(), dim); - // Store Q^T (transpose) as the rotation matrix + // Store Q^T (transpose) as the rotation matrix. + // + // Interchange so the write matrix_[j * dim + i] is contiguous in the inner i + // loop (stride 1); the original i-outer/j-inner order wrote by column (stride + // `dim`), thrashing cache for large dim. Q is read strided either way, so + // making the stores sequential is the net win. This is a pure copy, so the + // result is identical regardless of loop order. matrix_.resize(dim * dim); - for (size_t i = 0; i < dim; ++i) { - for (size_t j = 0; j < dim; ++j) { + for (size_t j = 0; j < dim; ++j) { + for (size_t i = 0; i < dim; ++i) { matrix_[j * dim + i] = Q[i * dim + j]; } } @@ -137,12 +143,22 @@ int MatrixRotator::init_impl(size_t dim) { void MatrixRotator::rotate(const float *in, float *out) const { const size_t dim = dimension_; // out = in * matrix_ (1 x dim) * (dim x dim) -> (1 x dim) + // + // Accumulate by input index i (outer) so matrix_ is read row-contiguously in + // the inner j loop: matrix_[i * dim + j] steps by 1. The original + // j-outer/i-inner order accessed matrix_ by column (stride `dim`), which + // thrashes cache/TLB for large dim and blocks auto-vectorization. The + // summation over i is performed in the same order, so the result is unchanged + // (any difference is at FMA/vectorization level). for (size_t j = 0; j < dim; ++j) { - float sum = 0.0f; - for (size_t i = 0; i < dim; ++i) { - sum += in[i] * matrix_[i * dim + j]; + out[j] = 0.0f; + } + for (size_t i = 0; i < dim; ++i) { + const float xi = in[i]; + const float *row = &matrix_[i * dim]; + for (size_t j = 0; j < dim; ++j) { + out[j] += xi * row[j]; } - out[j] = sum; } } diff --git a/tests/core/quantizer/matrix_rotator_test.cc b/tests/core/quantizer/matrix_rotator_test.cc new file mode 100644 index 000000000..4c6544caf --- /dev/null +++ b/tests/core/quantizer/matrix_rotator_test.cc @@ -0,0 +1,74 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "quantizer/rotator/rotator.h" + +using zvec::core::Rotator; + +namespace { + +// Independent reference for out = in * matrix (row-major, dim x dim): +// out[j] = sum_i in[i] * matrix[i * dim + j] +// Written in the straightforward j-outer / i-inner order so it does not share +// the loop structure of MatrixRotator::rotate(); it pins the value the rotate +// kernel must produce regardless of how its loops are ordered. +void reference_rotate(const float *in, const float *matrix, size_t dim, + float *out) { + for (size_t j = 0; j < dim; ++j) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + sum += in[i] * matrix[i * dim + j]; + } + out[j] = sum; + } +} + +} // namespace + +// MatrixRotator::rotate() is a cache-friendly (loop-interchanged) matrix-vector +// product. This guards the invariant that the interchange preserves: the output +// must match the plain row-major matvec across a range of dimensions, including +// non-power-of-two and odd sizes that stress the tail of the inner loop. +TEST(MatrixRotatorTest, RotateMatchesReferenceMatvec) { + std::mt19937 gen(0x5eed); + std::normal_distribution dist(0.0f, 1.0f); + + for (size_t dim : {1u, 2u, 3u, 7u, 16u, 31u, 64u, 128u, 257u}) { + std::vector matrix(dim * dim); + for (auto &m : matrix) m = dist(gen); + + std::unique_ptr rotator = Rotator::load_matrix(matrix.data(), dim); + ASSERT_NE(rotator, nullptr) << "dim=" << dim; + ASSERT_EQ(rotator->dimension(), dim); + + std::vector in(dim); + for (auto &x : in) x = dist(gen); + + std::vector out(dim), ref(dim); + rotator->rotate(in.data(), out.data()); + reference_rotate(in.data(), matrix.data(), dim, ref.data()); + + // Difference is only FMA/vectorization rounding; scale tolerance with dim + // since the sums grow with the number of accumulated terms. + const float tol = 1e-4f * static_cast(dim); + for (size_t j = 0; j < dim; ++j) { + EXPECT_NEAR(out[j], ref[j], tol) << "dim=" << dim << " j=" << j; + } + } +} diff --git a/tools/core/CMakeLists.txt b/tools/core/CMakeLists.txt index 86a3c4eac..724c8d536 100644 --- a/tools/core/CMakeLists.txt +++ b/tools/core/CMakeLists.txt @@ -9,6 +9,14 @@ cc_binary( LIBS gflags core_framework zvec_ailego ) +cc_binary( + NAME distance_bench + STRICT PACKED + SRCS distance_bench.cc + INCS ${PROJECT_ROOT_DIR}/src/ + LIBS gflags zvec_ailego + ) + cc_binary( NAME local_builder STRICT PACKED diff --git a/tools/core/distance_bench.cc b/tools/core/distance_bench.cc new file mode 100644 index 000000000..248027758 --- /dev/null +++ b/tools/core/distance_bench.cc @@ -0,0 +1,133 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Opt-in micro-benchmark for the core FP32 SquaredEuclidean distance kernel. +// +// Models a brute-force scan: for each query, the squared-Euclidean distance to +// every database vector is computed via +// zvec::ailego::Distance::SquaredEuclidean (which dispatches to AVX / AVX-512 / +// NEON). We report its throughput against a plain scalar reference across a +// range of dims, and self-validate that the SIMD result agrees with the scalar +// one before trusting any timing. +// +// Standalone executable (built under BUILD_TOOLS); intentionally NOT registered +// as a ctest -- micro-benchmark numbers are machine dependent. It exits +// non-zero only if the SIMD kernel disagrees with the scalar reference. +// +// ./distance_bench --num_db=8192 --dims=1,2,3,7,16,31,64,128,257 --reps=50 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +DEFINE_int32(num_db, 8192, "database vectors scanned per pass"); +DEFINE_int32(reps, 50, "timed repetitions per dim"); +DEFINE_string(dims, "1,2,3,7,16,31,64,128,257", + "comma-separated vector dims to benchmark"); +DEFINE_double(tolerance, 1e-4, + "max relative diff allowed between SIMD and scalar"); + +namespace { + +// Plain scalar reference -- the ground truth the SIMD kernel must match. +float ScalarSquaredEuclidean(const float *a, const float *b, size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float d = a[i] - b[i]; + sum += d * d; + } + return sum; +} + +std::vector ParseDims(const std::string &spec) { + std::vector dims; + std::stringstream ss(spec); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) { + dims.push_back(std::stoi(item)); + } + } + return dims; +} + +} // namespace + +int main(int argc, char *argv[]) { + gflags::SetUsageMessage( + "Micro-benchmark: FP32 SquaredEuclidean SIMD vs scalar reference"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + std::mt19937 rng(12345); // fixed seed: reproducible inputs across runs + std::uniform_real_distribution dist(-1.0f, 1.0f); + + bool ok = true; + std::printf("%6s %12s %12s %9s\n", "dim", "scalar(ms)", "simd(ms)", + "speedup"); + for (int dim : ParseDims(FLAGS_dims)) { + const size_t n = static_cast(FLAGS_num_db); + std::vector db(n * dim); + std::vector query(dim); + for (auto &v : db) v = dist(rng); + for (auto &v : query) v = dist(rng); + + // Correctness gate: every SIMD result must match the scalar reference. + double max_rel = 0.0; + for (size_t i = 0; i < n; ++i) { + const float *dbi = db.data() + i * dim; + float ref = ScalarSquaredEuclidean(dbi, query.data(), dim); + float got = + zvec::ailego::Distance::SquaredEuclidean(dbi, query.data(), dim); + double denom = std::max(1e-6f, std::fabs(ref)); + max_rel = std::max(max_rel, std::fabs(got - ref) / denom); + } + if (max_rel > FLAGS_tolerance) { + std::printf("dim %d: VALIDATION FAILED (max relative diff %.3e > %.1e)\n", + dim, max_rel, FLAGS_tolerance); + ok = false; + continue; + } + + volatile float sink = 0.0f; // keep the loops from being optimized away + + auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < FLAGS_reps; ++r) + for (size_t i = 0; i < n; ++i) + sink += ScalarSquaredEuclidean(db.data() + i * dim, query.data(), dim); + auto t1 = std::chrono::steady_clock::now(); + for (int r = 0; r < FLAGS_reps; ++r) + for (size_t i = 0; i < n; ++i) + sink += zvec::ailego::Distance::SquaredEuclidean(db.data() + i * dim, + query.data(), dim); + auto t2 = std::chrono::steady_clock::now(); + (void)sink; + + double scalar_ms = + std::chrono::duration(t1 - t0).count(); + double simd_ms = std::chrono::duration(t2 - t1).count(); + double speedup = simd_ms > 0.0 ? scalar_ms / simd_ms : 0.0; + std::printf("%6d %12.3f %12.3f %8.2fx\n", dim, scalar_ms, simd_ms, speedup); + } + + gflags::ShutDownCommandLineFlags(); + return ok ? 0 : 1; +}