perf(ivf): tiled-GEMM k-means assignment (~1.6× IVF-PQ training)

CHANGELOG.md

@@ -7,6 +7,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Performance (IVF / IVF-PQ training — GEMM k-means assignment)
+- New `rust/vectro_lib/src/index/kmeans.rs` (`assign_nearest`) replaces the
+  Lloyd **assignment** step's `parallel-over-points, serial-over-k` scan (one
+  distance call per centroid per point, `k` separate dot loops, poor
+  centroid-matrix reuse) with a tiled `[chunk, d]·[d, k]` GEMM + per-row argmax.
+  Both metrics reduce to a per-row argmax: `Cosine` on the raw dot (unit-norm
+  vectors), `L2` on `dot − ½‖c‖²` (the standard centroid-norm trick, for the
+  non-unit PQ sub-vector regime). Wired into `IvfPqIndex`'s coarse k-means
+  (`ivf_pq.rs`, cosine) and `IvfFlat`'s (`ivf.rs`, L2). PQ codebook training
+  (`quant/pq.rs`) already uses a LUT/SIMD-across-K assignment and is unchanged.
+- **Measured (this x86_64 host, n=50k · d=128 · n_lists=512, 25 iters):**
+  IVF-PQ train **2.47s → 1.52s (~1.6×)** at unchanged recall (the full
+  index/recall suite passes; `assign_nearest` is validated byte-for-byte against
+  a scalar oracle for both metrics, x86_64 + `qemu-aarch64`).
+- **Rejected — a single monolithic `.dot()`:** the first cut computed one
+  `[n, d]·[d, k]` GEMM and was **2× slower (5.2s)** than the old scalar scan.
+  ndarray's `matrixmultiply` is single-threaded, so a monolithic GEMM loses the
+  old code's `par_iter`-over-points parallelism *and* materialises the full
+  `[n, k]` (102 MB here) each iteration. Tiling across rayon workers — each
+  running its own small GEMM, mirroring `IvfPqIndex::search_batch_flat` — is
+  what recovers the parallelism and caps live memory at `[chunk, k]`.
+
 ### Performance (PQ4 fast-scan — NEON `vqtbl1q_u8` closes the aarch64 gap)
 - `rust/vectro_lib/src/index/pq4.rs` — the PQ4 fast-scan (`IndexPQFastScan`
   analogue, shared by `Pq4FlatIndex` and `IvfPq4Index`) had an AVX2 `pshufb`
@@ -25,9 +47,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   exercises NEON vs the scalar reference (byte-exact `u16` sums) and passes,
   along with `ranking_agrees_with_exact_adc` and the full `ivf_pq4` recall/batch
   suite, cross-compiled to `aarch64-unknown-linux-gnu` and run under
-  `qemu-aarch64`. Throughput on real Apple Silicon is to be measured via the
-  existing `bench-darwin-arm64` harness; the AVX2 twin documents ~22× over the
-  scalar gather on that platform's SIMD.
+  `qemu-aarch64`, and **confirmed on real Apple Silicon** by the green
+  `Rust tests (macos-latest)` (arm64) CI job. A throughput number is still to be
+  captured via the `bench-darwin-arm64` harness; the AVX2 twin documents ~22×
+  over the scalar gather on that platform's SIMD.
+
+### Changed (CI — `konjo-gates` provisions the Rust toolchain)
+- `.github/workflows/konjo-gates.yml` — the kiban gate job installed only Python +
+  kiban, so its `repo:*` Rust gates (`fmt-check`, `clippy`, `cargo-deny`,
+  `cargo-mutants`) had no `cargo` on PATH and no `cargo-deny`/`cargo-mutants`
+  binaries; each shelled out, failed instantly, and was reported as a spurious
+  "net-new findings" (the whole battery finished in ~0.4s, before any real compile
+  could run, returning byte-identical verdicts regardless of the diff). Added a
+  `dtolnay/rust-toolchain@stable` step (rustfmt + clippy), a cargo cache, and
+  `cargo install cargo-deny cargo-mutants` so the gates evaluate the real diff —
+  mirroring the working `konjo-gate.yml` G1/G3 setup. The gate went from failing
+  in ~0.4s to passing in ~3m21s (the runtime is the proof the tools now run).
 
 ### Performance (IVF-PQ — SIMD coarse scan via a shared distance module)
 - New `rust/vectro_lib/src/index/simd.rs` — a single source of truth for the

PLAN.md

@@ -1,7 +1,67 @@
 # Vectro — Plan
 
-> Last updated: 2026-06-18
-> Current version: **5.6.0** (Python) / **8.1.0** (Rust) — INT8 batch path routed through the Rust SIMD kernel with `range_factor` profile parity.
+> Last updated: 2026-07-01
+> Current version: **5.24.0** (Python) / **8.17.0** (Rust) — PQ4 fast-scan gained a NEON `vqtbl1q_u8` path on aarch64, and IVF/IVF-PQ k-means training now uses a tiled-GEMM assignment (~1.6× faster build at high `n_lists`).
+
+---
+
+## GEMM k-means assignment (IVF / IVF-PQ training) ✅ COMPLETE (2026-07-01)
+
+### Summary
+The Lloyd **assignment** step in IVF and IVF-PQ k-means scanned centroids
+serially per point (`parallel-over-points, serial-over-k`; `k` dot loops per
+point) — the CHANGELOG's "~3.5× slower than FAISS at high `n_lists`" gap. New
+`rust/vectro_lib/src/index/kmeans.rs::assign_nearest` replaces it with a tiled
+`[chunk, d]·[d, k]` GEMM + per-row argmax (Cosine = argmax dot; L2 = argmax
+`dot − ½‖c‖²`), reusing the `search_batch_flat` tiling pattern. PQ codebook
+training already used a LUT/SIMD-across-K assignment and was left unchanged.
+
+### Deliverables
+| # | Deliverable | Status |
+|---|-------------|--------|
+| 1 | `index/kmeans.rs` — `assign_nearest` (tiled GEMM, rayon, both metrics) + parity tests vs a scalar oracle | ✅ |
+| 2 | `ivf_pq.rs` / `ivf.rs` `kmeans_lloyd` assignment steps routed through the helper | ✅ |
+| 3 | Honest A/B — monolithic single-`.dot()` GEMM measured 2× *slower*, rejected in favour of the tiled form | ✅ |
+
+### Results
+- IVF-PQ train (this x86_64 host, n=50k · d=128 · n_lists=512, 25 iters):
+  **2.47s → 1.52s (~1.6×)** at unchanged recall.
+- Full index/recall suite green; `assign_nearest` byte-for-byte vs the scalar
+  oracle on both metrics, x86_64 + `qemu-aarch64`.
+
+---
+
+## NEON PQ4 fast-scan (aarch64) ✅ COMPLETE (2026-07-01)
+
+### Summary
+The PQ4 fast-scan (`IndexPQFastScan` analogue in `rust/vectro_lib/src/index/pq4.rs`,
+shared by `Pq4FlatIndex` and `IvfPq4Index`) had an AVX2 `pshufb` kernel on x86_64
+but fell back to the **scalar gather on aarch64** — so Apple Silicon, the flagship
+`bench-darwin-arm64` target, never got the fast-scan win. Added `scan_neon`, the
+aarch64 twin: NEON's `vqtbl1q_u8` is the direct analogue of AVX2 `pshufb`, resolving
+16 candidates' per-subspace distances against the 16-byte LUT in one table lookup.
+Unlike AVX2's lane-crossing `unpack{lo,hi}_epi8` (which needs the `PERM` table to
+undo the permutation), `vqtbl1q_u8` + `vget_{low,high}` preserve candidate order, so
+the four u16 accumulators store straight to the output with no permute. NEON is
+mandatory in the aarch64 base ISA, so the path is unconditional (no runtime
+detection); the scalar reference remains for non-AVX2 x86_64 and other targets.
+
+### Deliverables
+| # | Deliverable | Status |
+|---|-------------|--------|
+| 1 | `scan_neon` NEON `vqtbl1q_u8` fast-scan kernel in `pq4.rs` | ✅ |
+| 2 | `scan` dispatcher routes aarch64 → NEON unconditionally; scalar reference retained for other targets | ✅ |
+| 3 | Konjo gate-hygiene on the diff: rustfmt wrapping, clippy `semicolon_if_nothing_returned`, `// SAFETY:` comments | ✅ (PR #99) |
+| 4 | CI fix: `konjo-gates.yml` provisions the Rust toolchain + `cargo-deny`/`cargo-mutants` so the kiban `repo:*` gates actually run | ✅ (PR #100) |
+
+### Results
+- Byte-exact vs the scalar reference (`scan_simd_matches_scalar`,
+  `ranking_agrees_with_exact_adc`, full `ivf_pq4` recall/batch suite),
+  cross-compiled to `aarch64-unknown-linux-gnu` and run under `qemu-aarch64`.
+- Confirmed on **real Apple Silicon** by the green `Rust tests (macos-latest)`
+  CI job (arm64) on PR #99.
+- Throughput number still pending a `bench-darwin-arm64` run; the AVX2 twin
+  documents ~22× over the scalar gather on that platform's SIMD.
 
 ---
 

rust/vectro_lib/src/index/ivf.rs

@@ -88,23 +88,14 @@ fn kmeans_lloyd(data: &[&[f32]], k: usize, d: usize, max_iter: usize, seed: u64)
     let mut cents = kmeans_pp_init(data, k, d, seed);
     let mut assignments = vec![0usize; n];
 
+    // Flatten the data once into a contiguous [n, d] buffer for the GEMM
+    // assignment (centroids are already flat [k, d]).
+    let data_flat: Vec<f32> = data.iter().flat_map(|v| v.iter().copied()).collect();
+
     for _ in 0..max_iter {
-        let new_asgn: Vec<usize> = data
-            .par_iter()
-            .map(|v| {
-                let mut best = 0;
-                let mut best_d = f32::INFINITY;
-                for ki in 0..k {
-                    let c = &cents[ki * d..(ki + 1) * d];
-                    let dist = l2_sq(v, c);
-                    if dist < best_d {
-                        best_d = dist;
-                        best = ki;
-                    }
-                }
-                best
-            })
-            .collect();
+        // Assignment step — one GEMM + parallel argmax (squared-L2 nearest).
+        let new_asgn =
+            super::kmeans::assign_nearest(&data_flat, &cents, n, k, d, super::kmeans::Metric::L2);
 
         if new_asgn == assignments {
             break;

rust/vectro_lib/src/index/ivf_pq.rs

@@ -124,27 +124,24 @@ pub(crate) fn kmeans_lloyd(
     max_iter: usize,
     seed: u64,
 ) -> Vec<f32> {
-    let _n = data.len();
+    let n = data.len();
     let mut centroids = kmeans_pp_init(data, k, d, seed);
 
+    // Flatten the data once into a contiguous [n, d] buffer for the GEMM
+    // assignment (centroids are already flat [k, d]).
+    let data_flat: Vec<f32> = data.iter().flatten().copied().collect();
+
     for _ in 0..max_iter {
-        // Assignment step — parallelised
-        let assignments: Vec<usize> = data
-            .par_iter()
-            .map(|v| {
-                let mut best_c = 0usize;
-                let mut best_d = f32::MAX;
-                for ci in 0..k {
-                    let cent = &centroids[ci * d..(ci + 1) * d];
-                    let dist = cosine_dist(v, cent);
-                    if dist < best_d {
-                        best_d = dist;
-                        best_c = ci;
-                    }
-                }
-                best_c
-            })
-            .collect();
+        // Assignment step — one GEMM + parallel argmax. Vectors are unit-norm,
+        // so nearest-by-cosine is the max dot product.
+        let assignments = super::kmeans::assign_nearest(
+            &data_flat,
+            &centroids,
+            n,
+            k,
+            d,
+            super::kmeans::Metric::Cosine,
+        );
 
         // Update step
         let mut new_centroids = vec![0.0f32; k * d];