perf(pq): add transposed SIMD encoder with SGEMM fallback

[jerry-024]

What changed

• Replace the primary IVF-PQ blocked SGEMM encoder with a per-call transposed codebook and fused direct-L2 argmin, retaining blocked SGEMM only as the hardware fallback.
• Add specialized AVX2/FMA and NEON kernels for dsub=4, plus generic AVX2/FMA and 16-centroid NEON kernels for other supported shapes.
• Use the transposed path for every batch size when nbits=8, ksub=256, every subvector has at least four dimensions, and the CPU provides a fast fused kernel. x86 CPUs without AVX2+FMA use blocked SGEMM for every batch size; 4-bit and smaller-subvector shapes keep the canonical encoder.
• Add ivf.pq-encoding=auto|canonical. auto is the default; canonical reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend.
• Document and test determinism, batch/thread/split invariance, smallest-index tie breaking, non-finite distance semantics, non-uniform chunks, and scalar/SIMD agreement.

The codebook transpose is temporary and rebuilt for each call. At d=768 it allocates and fills a 768 KiB buffer. This avoids cached derived state becoming stale while ProductQuantizer::centroids remains public.

Why

The previous path used SGEMM for each subquantizer and 512-row block, then materialized and scanned a 512 × 256 distance matrix twice for expanded-form correction and scalar argmin. In the original 32,768-row Cohere profile, SGEMM took 46 ms, while correction plus argmin raised the complete path to 333 ms. The transposed kernel accumulates distances and tracks minima without writing the distance matrix.

Encoding behavior

The automatic backend is selected only from shape and CPU features, never from batch size. On x86 with AVX2+FMA and on AArch64, it computes direct squared L2 distances with the transposed encoder. Equal distances choose the smallest centroid index, NaN distances do not update the minimum, infinite distances lose, and code 0 is returned when no distance is below f32::MAX.

Other CPUs use blocked SGEMM for finite 8-bit codebooks with dsub>=4; non-finite codebooks on that hardware and unsupported shapes use the canonical encoder. Both fallback paths use expanded-form arithmetic (|q|² + |c|² - 2q·c). Codes can differ across backends at floating-point ties, and NaN centroids or high-dynamic-range finite inputs can produce different behavior. Set:

ivf.pq-encoding=canonical

when the canonical per-vector encoder must be reproduced on the same CPU and runtime backend. Canonical mode also uses CPU-dispatched norm reductions, so neither mode promises byte-identical codes across CPU feature sets. This setting affects builds only; the index format and query path are unchanged.

Cohere 10M end-to-end acceptance

Environment: Intel Xeon 6982P-C, 8 cores / 16 threads, AVX2+FMA; object storage with client-side caching disabled. Index shape: cosine, d=768, nlist=4096, m=192, dsub=4, one 10M-row shard. Three paired runs used alternating AB / BA / AB order.

Run	Main SGEMM build	This PR build	Speedup
1	261.786 s	187.200 s	28.49%
2	262.205 s	191.229 s	27.07%
3	258.752 s	191.616 s	25.95%
Median	261.786 s	191.229 s	27.07% paired median

Median index_add_ms fell from 133.447 s to 62.549 s. The observed full-build median is better than the original 205–215 s prediction.

PQ encoding time

Run	Main add.encode	This PR add.encode	Reduction
1	88.372 s	~17.646 s	80.03%
2	88.300 s	~18.095 s	79.51%
3	88.220 s	~16.744 s	81.02%
Median	88.300 s	~17.646 s	80.02% (~5.00×)

Main reports add.encode directly. The PR branch does not contain the benchmark-only phase timer, so its values are derived per pair as candidate index_add - (baseline index_add - baseline add.encode).

PQ encoding by subvector shape

A separate encode-only benchmark used 32,768 synthetic finite vectors with d=768 and 16 Rayon threads. Each number is the median of seven measured runs after one warmup.

Shape	Main blocked SGEMM	This PR auto	Canonical	Auto vs main
m=192, dsub=4	288.628 ms	56.449 ms	638.018 ms	5.11× faster
m=96, dsub=8	152.648 ms	103.389 ms	431.092 ms	1.48× faster
m=48, dsub=16	83.063 ms	79.618 ms	330.097 ms	1.04× faster

These measurements use AVX2+FMA. The specialized dsub=4 kernel targets Cohere/SIFT/GloVe production shapes; dsub=8 and dsub=16 use the generic fused kernel. On CPUs without a fast fused kernel, every finite 8-bit shape with dsub>=4 uses blocked SGEMM instead; there is no batch-size threshold.

A paired encode-only follow-up on an Apple M3 Pro used 32,768 rows, d=768, and 16 Rayon threads. The explicit generic NEON kernel removes the score-table write and scalar argmin from the previous AArch64 path:

Shape	Compiler-vectorized scratch path	Fused generic NEON	Speedup
dsub=8	189.108 ms	79.238 ms	2.39×
dsub=12	156.151 ms	69.463 ms	2.25×
dsub=16	135.517 ms	69.166 ms	1.96×

A Rosetta x86_64 probe, where runtime detection reports AVX2=false and FMA=false, measured the deliberate cost of using SGEMM for every batch size: one row changed from 218.291 µs to 297.875 µs (+36.5%), and seven rows from 470.083 µs to 584.125 µs (+24.3%). These are emulated-host diagnostics rather than native no-FMA hardware claims. The fallback remains independent of batch size so splitting the same input cannot select a different floating-point encoder.

Functional coverage spans the full dsub domain by execution path: dsub=1/2/3 falls back to the unchanged canonical encoder and is tested byte-for-byte against encode_batch; the transposed path covers the specialized dsub=4 kernel and representative generic shapes 5/8/12/16; the SGEMM fallback is tested directly for thread and batch-split invariance. A balanced non-uniform 5/4/4 layout is also covered. Release-mode PQ tests now run on the macOS AArch64 CI runner, and its Rosetta x86_64 target exercises the actual no-AVX2/FMA dispatch. An environment guard fails that job if Rosetta ever starts reporting AVX2+FMA. dsub=1/2/3 is omitted from the performance tables because this PR does not change that path.

Quality used the same 1,000 queries at nprobe=128, top-k=100. Each build trained afresh with the same machine, data, configuration, and deterministic seed; it did not reuse a precomputed shared codebook. All three result files were byte-identical within each encoder.

Metric	Main SGEMM	This PR	Delta / result
Recall@100	0.808990	0.808990	0; bootstrap 95% CI [0, 0]
NDCG@100	0.853580113	0.853579531	-0.000000582; bootstrap 95% CI [-0.000001513, 0.000000020]
Top-100 set changes	—	3 / 1,000	0.3%
Ranking changes	—	20 / 1,000	2.0%
Paired QPS change	—	—	+1.23% median

The 10,000-sample paired-bootstrap lower bounds exceed the -0.002 gate, and paired median QPS does not regress beyond the 3% gate.

A separate 1M-row consistency probe found identical coarse assignments and 377 rows with one changed PQ-code byte each: 377 changed bytes out of 192M (1.96e-6). Among its 1,000 queries, four top-100 sets and 17 rankings changed. These are diagnostic values, not acceptance gates.

Validation

• cargo test --workspace
• cargo clippy --workspace --all-targets -- -D warnings
• x86_64 cross-target compilation
• release-mode PQ tests on the macOS AArch64 CI runner
• no-AVX2/FMA dispatch and SGEMM fallback tests under Rosetta x86_64
• Java API and JNI native validation tests
• Python tests: 28 passed
• Cohere 10M: three paired builds and 1,000-query quality/QPS runs

[shyjsarah]

Code Review Summary

Mode: full
Scope: GitHub PR #90 at b40b193a0fbda18826332eac795fee3ed63cf4d0 against main merge-base 8dcabf208c99707eab3938af0bcc40530fdb4cad
Files Changed: 10 files (+827/-80)
Score: 71/100
Recommendation: Request changes
GAN Stats: 7 raw findings → 6 after dedup; Discriminator accepted 3 / challenged 2 / rejected 1; Arbiter included 5 / adjusted 2 / excluded 1.

Critical Issues

None.

Major Issues

1. core/src/ivfpq.rs:319-320 — canonical does not deliver the documented cross-CPU byte stability.
   The mode delegates to ProductQuantizer::encode_batch, whose expanded-L2 implementation uses target-dependent SGEMM and norm reductions. A fixed-query/fixed-codebook release probe produced code 1 on aarch64-apple-darwin and code 0 on x86_64-apple-darwin, contradicting the new documentation. Implement a fixed-order architecture-independent canonical encoder, or narrow the public guarantee and documentation.

2. core/src/pq.rs:398-399, 707-753 — AArch64 dsub != 4 is routed from blocked SGEMM to the scalar transposed kernel.
   AArch64 unconditionally enables transposed encoding, but only dsub == 4 has a NEON kernel; other shapes fall through to score_argmin_scalar. On this ARM host, d=768/m=48/dsub=16, 32,768 rows regressed from 91.62 ms on the parent to 125.33 ms on the PR (+36.8%). Make backend selection shape-aware and keep SGEMM for ARM shapes without a benchmarked fast kernel.

Minor Issues

1. core/src/pq.rs:484-514 — The transposed codebook is allocated and rebuilt for every add batch.
   This rewrites m * max_dsub * 256 floats per call (768 KiB for d=768), including one-row streaming calls. Cache derived state safely or provide a writer-owned immutable/precomputed transpose.

2. core/src/pq.rs:396-413 — The no-AVX2/FMA SGEMM fallback now runs for every batch size.
   Removing the previous small-row threshold preserves split invariance, but measured tiny-batch latency regressed by about 52% at one row and 17% at seven rows. Add a deterministic low-overhead tiny-batch fallback or document/benchmark the tradeoff.

3. .github/workflows/ci.yml — New AArch64-only unsafe kernels have no behavior test in PR CI.
   The ordinary Rust/FFI/JNI behavior jobs run on Ubuntu x86_64; release build/load smoke tests cannot validate NEON numerical equivalence, tie handling, or non-finite semantics. Add an AArch64 core PQ test job and a forced non-AVX2 fallback job.

Validation

• cargo fmt --all -- --check: passed.
• cargo test --workspace: passed (core 488 passed / 1 ignored; integration, FFI, and JNI tests passed).
• cargo clippy --all-targets --workspace -- -D warnings: passed.
• cargo check -p paimon-vindex-core --target x86_64-apple-darwin: passed.
• Cross-target canonical reproduction independently confirmed: aarch64=1, x86_64=0.
• Security review found no actionable issue.

Positive Observations

• The PR adds extensive correctness tests for ties, non-finite values, batch/thread/split invariance, non-uniform chunks, and scalar/SIMD agreement.
• The specialized dsub=4 fast path is well targeted to the documented production shape and shows a substantial x86 AVX2/FMA speedup.
• The fallback behavior and arithmetic differences are documented more thoroughly than typical performance changes, although the canonical portability claim needs correction.

[JingsongLi]

Could we avoid rebuilding the entire transposed codebook for unamortized tiny batches? build_transposed_codebook allocates and copies d × 256 floats on every call (768 KiB at d=768), even when n=1. The parent used the canonical path below max(32, 4 × threads); on native arm64 with d=768, m=96, dsub=8, an isolated release probe measured auto at 209 µs versus 95 µs for the parent-equivalent canonical path (2.20× slower). At n=7 auto was already faster, so this is limited to row-at-a-time ingestion. A centroid-major direct-L2 single-row kernel with the same dimension-order mul_add, NaN, and tie semantics would preserve split invariance without paying the transpose; alternatively, cache the transpose with reliable invalidation for the publicly mutable centroids.

[JingsongLi]

The documented dispatch is broader than the implementation. Auto uses transposed/SGEMM only for 8-bit, ksub=256 shapes whose every chunk has dsub>=4; dsub=1/2/3 falls back to canonical even on AVX2+FMA/AArch64, and a non-finite codebook on the no-fast-kernel branch also uses canonical. For example, dimension=128 with pq.m=64 is valid but auto already equals canonical, so the arithmetic and performance guidance here is incorrect. Could we qualify this callout—and the matching IVF-PQ and release docs—with the actual gates and canonical fallback?