perf(quantizer): make MatrixRotator::rotate row-contiguous (loop interchange, ~12-15x)

src/core/quantizer/rotator/matrix_rotator.cc

@@ -124,10 +124,16 @@ int MatrixRotator::init_impl(size_t dim) {
   std::vector<float> 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;
   }
 }
 

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 <cmath>
+#include <memory>
+#include <random>
+#include <vector>
+#include <gtest/gtest.h>
+#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<float> dist(0.0f, 1.0f);
+
+  for (size_t dim : {1u, 2u, 3u, 7u, 16u, 31u, 64u, 128u, 257u}) {
+    std::vector<float> matrix(dim * dim);
+    for (auto &m : matrix) m = dist(gen);
+
+    std::unique_ptr<Rotator> rotator = Rotator::load_matrix(matrix.data(), dim);
+    ASSERT_NE(rotator, nullptr) << "dim=" << dim;
+    ASSERT_EQ(rotator->dimension(), dim);
+
+    std::vector<float> in(dim);
+    for (auto &x : in) x = dist(gen);
+
+    std::vector<float> 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<float>(dim);
+    for (size_t j = 0; j < dim; ++j) {
+      EXPECT_NEAR(out[j], ref[j], tol) << "dim=" << dim << " j=" << j;
+    }
+  }
+}

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

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 <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstddef>
+#include <cstdio>
+#include <random>
+#include <sstream>
+#include <string>
+#include <vector>
+#include <ailego/math/distance.h>
+#include <gflags/gflags.h>
+
+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<int> ParseDims(const std::string &spec) {
+  std::vector<int> 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<float> 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<size_t>(FLAGS_num_db);
+    std::vector<float> db(n * dim);
+    std::vector<float> 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<double, std::milli>(t1 - t0).count();
+    double simd_ms = std::chrono::duration<double, std::milli>(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;
+}