Batch native Highway operations to reduce Java/native overhead

[tlwillke]

Goal and changes

Reduce Java/native call overhead and temporary allocation in the Google Highway backend, especially during PQ training and encoding. Builds on the native kernels and provider selection already merged in #753 and #751.

• Batch closest-centroid scans and KMeans initialization; fuse centroid normalization and assignment updates.
• Use lazy array-backed native views and direct byte/scalar access to avoid temporary wrappers and allocations.
• Batch NVQ growth-rate search, use SIMD for NVQ setup/quantization, and keep cosine results packed.

Performance results

Both arms use native Highway AVX-512. Each cell shows main → this PR; bold marks the winning measurement.

PQ training is approximately 1.9–3.2× faster, and encoding 2.3–3.5× faster. Total construction is lower in every comparison. Single-thread query throughput improves in seven of eight dataset/JDK combinations. I confirmed that recall@10 is essentially unchanged.

Construction — seconds, lower is better

Dataset	JDK	PQ training	PQ encoding	Graph/write	Total process
MPNet	22	4.76 → 2.56	3.27 → 0.94	53.01 → 29.89	65.42 → 37.74
MPNet	25	6.93 → 2.41	3.15 → 0.94	32.05 → 32.60	46.52 → 40.15
Ada002	22	12.29 → 4.29	6.20 → 1.93	53.92 → 55.74	79.46 → 69.19
Ada002‡	25	10.24 → 3.47	5.99 → 1.84	47.11 → 45.29	69.56 → 57.18
Cohere†	22	9.31 → 2.87	4.28 → 1.42	40.76 → 43.84	59.82 → 53.78
Cohere	25	7.87 → 2.68	4.07 → 1.22	37.10 → 41.29	54.20 → 50.32
OpenAI3072	22	23.68 → 7.55	13.53 → 5.80	106.51 → 116.92	156.09 → 142.88
OpenAI3072	25	20.42 → 6.85	13.62 → 5.45	119.61 → 118.11	165.65 → 143.05

Queries — QPS, higher is better

Dataset	JDK	1 thread, 1×	1 thread, 2×	8 threads, 1×	8 threads, 2×
MPNet	22	12,180 → 13,136	9,324 → 10,021	97,973 → 108,674	68,327 → 82,580
MPNet	25	12,114 → 13,682	9,393 → 10,480	89,620 → 92,609§	72,653 → 70,871§
Ada002	22	7,020 → 7,186	5,492 → 5,666	54,180 → 60,928	43,811 → 48,857
Ada002	25	6,903 → 7,365	5,443 → 5,812	60,556 → 58,350§	46,388 → 50,278
Cohere	22	10,423 → 9,958	7,905 → 7,723	82,282 → 80,848	61,963 → 62,848
Cohere	25	8,420 → 10,661	6,530 → 8,080	81,661 → 84,463§	61,133 → 66,717
OpenAI3072	22	3,295 → 3,652	2,752 → 2,994	20,433 → 23,676	18,080 → 23,245
OpenAI3072	25	3,295 → 3,768	2,712 → 2,996	25,044 → 24,760	22,230 → 23,384

Graph/write remains slower in several cases. The fresh Cohere/JDK22 repeat shows 7.6% higher graph/write time, with 10.1% lower total construction time; its cause is not established. Cohere/JDK22 single-thread QPS is 2.3–4.5% lower.

Method: BenchYAML, JDK22/25, four nominal-1M datasets, 128 GiB heap, 48 construction workers, deferred batch writer, FusedPQ width8/K256 uncentered/unweighted, NVQ with two 8-bit subvectors, V6 M32/ef100, hierarchy/refinement enabled, pruning and assertions disabled. Every construction uses a fresh codebook and graph. Queries use resident graphs, pinned 1-/8-thread execution, top10 at 1×/2× overquery, and three warmup/measurement rounds.

† Fresh A-B-B-A medians for main eba7bd4e versus latest PR head cebb94f8, replacing the earlier Cohere/JDK22 construction row.

Other rows compare historical main eba7bd4e measurements with the f65c91f8 rerun on the same server. The latest scalar fix received separate correctness/microbenchmark validation and the Cohere construction repeat; the full query sweep was not repeated after it.

Test and methodology

• Full reactor tests passed on JDK 22 and 25: 496 tests per JDK for this PR, zero failures/errors, five skips. The native-module suite also passed with automatic dispatch, AVX2, the scalar tier, and Java fallback on both JDKs. Both original CI test failures were reproduced before applying the fixes. Coverage includes native PQ/NVQ operations, lazy storage, byte reads, and provider/ABI fixtures. Standalone native C++ tests were not run locally because GoogleTest was unavailable. Apache RAT license checks also passed.
• Completed 80 primary benchmark stages plus 12 noisy-query repeat passes, with no stage failures. Separate memory/GC passes were collected as supporting data.
• Standard BenchYAML on the same Xeon 6985P-C server; JDK 22.0.2/25.0.2; independent fresh codebooks/graphs; 128 GiB maximum heap; 48 construction workers; deferred batch writer; FusedPQ width8/K256 uncentered/unweighted; NVQ2; V6 M32/ef100; hierarchy/refinement enabled; pruning disabled; resident query graphs; 1-/8-thread queries pinned to CPU0/CPU0–7; assertions disabled.

[github-actions]

Before you submit for review:

• Does your PR follow guidelines from CONTRIBUTIONS.md?
• Did you summarize what this PR does clearly and concisely?
• Did you include performance data for changes which may be performance impacting?
• Did you include useful docs for any user-facing changes or features?
• Did you include useful javadocs for developer oriented changes, explaining new concepts or key changes?
• Did you rebase your branch onto the latest main for regression testing and PR submission?
• Did you trigger regression testing via Run Bench Main and review results?
• Did you adhere to the code formatting guidelines (TBD)
• Did you group your changes for easy review, providing meaningful descriptions for each commit?
• Did you ensure that all files contain the correct copyright header?
• Did you add documentation for this feature to the release notes directory?

If you did not complete any of these, then please explain below.

[MarkWolters]

Approved, although take the comment into consideration as a possible change before merging.

[jshook]

PR #761 review: Batch native Highway operations

Reviewed at 50d7d50d against origin/main eba7bd4e. The equivalence claims hold. There is
one speed regression on 128-bit vector platforms, plus API and test gaps.

Verification

• Full reactor suite with the native provider forced (JDK 24, avx3_dl): 496 tests, 0 failures,
  5 skipped (unrelated).
• jvector-native suite: 43/43 pass at JVECTOR_MAX_ISA=avx2 + -XX:MaxVectorSize=32, and at
  sse42 + MaxVectorSize=16.

Findings

1. The 8-dimension KMeans++ path is 3.6–5× slower at 128 bits
   (NativeKMeansInitialization.java:37).

   • SPECIES_256 falls back to unaccelerated Java code wherever there's no 256-bit hardware.
     Main used Highway's euclidean_f32 here, which adapts to the ISA.

   • Timing of one seeding pass (32K points, 255 updates), against main's native loop in the
     same JVM. Results were identical in every row.

     Width	PR	main native	PR ÷ main
     512 (avx3_dl)	5.3 ns/dist	14.1	0.38×
     256 (avx2, MaxVectorSize=32)	5.1	14.0	0.37×
     128 (sse42, MaxVectorSize=16)	55–77	15.4	3.6–5.0×

   • The 128-bit row is simulated on x86. NEON, 128-bit SVE2 (Graviton4) and SSE4.2-capped x86
     are expected to match. arm64 CI checks correctness only.

   • Suggested fix: use this path only when SPECIES_PREFERRED.vectorBitSize() >= 256;
     otherwise return false.

   • The comment "Main's exact-width path" describes Panama. For the native provider these
     distances now differ from main, within tolerance.

2. NVQ-only copies of generic operations in the public API (VectorUtil.java:253,
   VectorUtilSupport.java:210).

   • nvqMin, nvqScale and two nvqSub overloads compute the same thing as min, scale
     and sub. They only select a different implementation per call site (see 1899795,
     82cd218).
   • With the try* methods and supportsNvqCosine8bitPacked, that's eight new public entry
     points, most without javadoc.
   • Prefer keeping them internal to the native provider, or choosing by length inside
     NativeVectorUtilSupport. If they stay public, document them.

3. Unneeded volatile and locking in the lazy view (MemorySegmentVectorFloat.java:37,
   :67).

   • Both states reference the same array, so a race only produces two equivalent views. A heap
     segment's state lives in final fields, which should make a plain field safe (worth
     confirming).
   • As written, every get(i)/set(i) pays for a volatile read and a type check that the JIT
     can't move out of loops; the distanceValues workaround in KMeansPlusPlusClusterer
     suggests this cost. synchronized (this) also locks on a public object.
   • After a vector's first native use, scalar reads take getAtIndex plus an extra bounds
     check.

4. encodeTo lost its length check (NVQuantization.java:209). A longer vector is encoded
   silently using only its prefix; a shorter one fails with an index error instead of an
   argument error.

5. Assertions double the closest-centroid cost under -ea (VectorUtil.java:111). This is
   deliberate, but worth noting in the PR description.

Verified correct:

• closest_centroid_f32: same comparison, ties and NaN behavior.
• nvq_quantize_8bit: DemoteTo saturation matches the old clamp, with a bounded tail store.
• nvq_shuffle_query_in_place_8bit: the LoadInterleaved4 order matches the old permutation,
  and the in-place write is safe.
• nvq_growth_rate: same candidates, increments, starting value and comparison as compute.
• subtract_add and scale_to: arithmetic unchanged, overlap checked before any write.
• ABI 1: no released library is affected, and the missing-symbol fallback works.

Test coverage

The correctness oracles are strong: bitwise comparisons with the replaced code, scalar oracles,
ties, non-finite values, tails, aliasing and concurrency. CI runs them, and the C++ tests, at
every ISA cap on both architectures. Gaps:

1. No performance check below 256 bits. All benchmarks ran on one Xeon. A Graviton run of
   PQ training would cover finding 1.
2. KMeansInitializationTest uses the wrong baseline. It compares with Panama, not the
   native path it replaces. Seeding is where codebooks can now differ from main, and only recall
   benchmarks cover that.
3. x86 CI never caps the Java vector width. The Java-side paths run below 256 bits only on
   arm64. Adding -XX:MaxVectorSize=16 to the sse42 job would fix that.
4. No direct test of the packed NVQScorer cosine path against the unpacked one.
5. Minor:
   • No test for the encodeTo length case.
   • KMeansInitializationTest lacks requireNativeProvider(), so it errors rather than skips
     where the native library can't load.

Reproduce the 128-bit timing (a throwaway JUnit class in jvector-native tests, available on
request):
JVECTOR_MAX_ISA=sse42 JAVA_TOOL_OPTIONS=-XX:MaxVectorSize=16 mvn -B test -pl jvector-native -am -Dtest=<class> -Dsurefire.failIfNoSpecifiedTests=false.
JAVA_TOOL_OPTIONS is needed because the module's surefire config sets its own argLine.

[tlwillke]

Thanks @MarkWolters and @jshook. Pushed 404700ba, 2cec3d5c, and 0408be8b addressing:

• Documented the NVQ and optional-operation APIs. Kept NVQ-scoped dispatch to avoid regressing other callers.
• Added dimension validation before encodeTo modifies its destination, with rejection tests.
• Clarified assertion overhead and corrected the misleading Panama/native wording.
• Added native KMeans reference and packed/unpacked/scalar cosine tests.
• Capped Java vector width alongside Highway in SSE4.2/AVX2 CI.

The 11 focused tests passed in all six JDK 22/25 × ISA/width configurations, without skips. The original Java-only KMeans test needs no native-library guard; the new native-reference test has one.

The 128-bit specialization remains unpublished while I verify that it preserves 256/512-bit performance on both JDKs. I’m retaining synchronized lazy-view publication for this PR; removing it warrants separate concurrency and performance validation.

[tlwillke]

Thanks for the detailed reviews @MarkWolters and @jshook. The follow-up commits address the findings:

• 128-bit KMeans: f65c91f8 selects a separately typed 128-bit implementation at provider startup. The existing wide helper is unchanged. JDK22/25 tests show the 128-bit allocation/performance problem resolved, with 256-/512-bit timing preserved.
• Lazy storage: cebb94f8 uses plain storage snapshots for scalar access while retaining synchronized native-view initialization and volatile publication. Scalar reads returned to approximately 0.45 ns/float, without adding per-vector storage.
• API and validation: 404700ba documents the retained NVQ/provider hooks, restores encodeTo dimension validation and tests, corrects the Panama/native wording, and documents the assertion-only second scan.
• Coverage: 2cec3d5c and 0408be8b add the native KMeans baseline, direct packed-versus-unpacked NVQ cosine tests, and Java vector-width caps alongside native ISA caps. Native-unavailable tests skip appropriately.

All 20 CI checks pass, including ARM and capped x86 configurations. Narrow-width performance measurements used capped x86.

[jshook]

All known gaps have been closed