Lab VE4: SIMD Vectorization Microbenchmark

Every dense neural/knn query you ran in the earlier labs (VE1, VE2) bottomed out in the same arithmetic kernel: a distance computation over two float[]s, run thousands of times per query and hundreds of millions of times per graph merge. This lab proves the SIMD speedup that kernel gets. You will write a scalar dot product and a Panama Vector API dot product, benchmark both over many 768-dim vectors, run with and without --add-modules jdk.incubator.vector, measure ns/op, and confirm that Lucene on your machine actually picked PanamaVectorUtilSupport.

This is the hands-on companion to SIMD and the Panama Vector API and the SIMD section of the masterclass index. It is also direct preparation for Capstone Project 04: Contribute an HNSW improvement upstream to Lucene, where this kernel is the thing you optimize.

Background

The Panama Vector API (jdk.incubator.vector) lets Java emit SIMD: you write lane-wise operations on a FloatVector, and the HotSpot JIT compiles them to AVX-512/AVX2 (x86) or NEON (ARM) instructions, or falls back to scalar. A scalar loop does one float per instruction; a SIMD loop does 8 (AVX2) or 16 (AVX-512). On a 768-dim dot product that is a large, measurable win — but only if the module is loaded, the JIT warmed up, and you are not fooled by auto-vectorization. This lab makes all of that concrete and measured.

Why This Matters for Contributors

• "Is SIMD actually on?" is a real operational and contributor question. A node silently running the scalar fallback is 2–8× slower at vector search with no error — an invisible performance bug. You need to be able to prove the kernel is vectorized.
• If you ever touch Lucene's VectorUtil or the k-NN distance path (Capstone 04), this is the measurement harness you will use to justify the change. Distance-kernel changes are high-leverage because they multiply across millions of calls.
• It builds the JIT/warmup/benchmarking discipline that keeps you from reporting bogus "SIMD didn't help" numbers.

Prerequisites

• JDK 21 (the version OpenSearch currently uses; the Vector API incubator ships with it). Confirm: java -version.
• javac/java on the path. No Maven/Gradle needed — single-file.
• You have read SIMD and the Panama Vector API.
• (Optional) a JMH setup if you want rigorous numbers; this lab gives a self-contained harness first, then shows the JMH upgrade.

Note: jdk.incubator.vector is an incubator module — it must be loaded explicitly with --add-modules jdk.incubator.vector. Without it the class will not even load. That on/off switch is the heart of this lab.

Step-by-Step Tasks

Step 1 — Confirm your CPU's preferred lane width

Before benchmarking, find out how wide your SIMD registers are. Write SimdCheck.java:

import jdk.incubator.vector.FloatVector;
import jdk.incubator.vector.VectorSpecies;

public class SimdCheck {
    public static void main(String[] args) {
        VectorSpecies<Float> s = FloatVector.SPECIES_PREFERRED;
        System.out.println("Preferred lane count: " + s.length()
            + "  (vector bits: " + s.vectorBitSize() + ")");
        // 16 => AVX-512, 8 => AVX2, 4 => NEON, 1 => no SIMD / module problem
    }
}

javac --add-modules jdk.incubator.vector SimdCheck.java
java  --add-modules jdk.incubator.vector SimdCheck
# e.g. "Preferred lane count: 8  (vector bits: 256)"  on an AVX2 box
#      "Preferred lane count: 16 (vector bits: 512)"  on AVX-512
#      "Preferred lane count: 4  (vector bits: 128)"  on ARM/Graviton NEON

Record your lane count — it predicts the speedup ceiling (roughly that many floats per instruction in the main loop).

Step 2 — Write the benchmark: scalar vs Panama dot product

Create DotBench.java. It generates many random 768-dim vectors, then times a scalar dot product and a Panama dot product over all of them, with warmup:

import jdk.incubator.vector.FloatVector;
import jdk.incubator.vector.VectorOperators;
import jdk.incubator.vector.VectorSpecies;
import java.util.Random;
import java.util.concurrent.TimeUnit;

public class DotBench {
    static final int DIM = 768;          // typical embedding dimension
    static final int N   = 100_000;      // number of vectors to score against one query
    static final VectorSpecies<Float> SP = FloatVector.SPECIES_PREFERRED;

    /** Scalar dot product: one multiply-add per loop iteration. */
    static float dotScalar(float[] a, float[] b) {
        float sum = 0f;
        for (int i = 0; i < a.length; i++) sum += a[i] * b[i];
        return sum;
    }

    /** Panama dot product: lane-striding main loop with FMA + horizontal reduce + scalar tail. */
    static float dotVector(float[] a, float[] b) {
        FloatVector acc = FloatVector.zero(SP);
        int i = 0, bound = SP.loopBound(a.length);
        for (; i < bound; i += SP.length()) {
            FloatVector va = FloatVector.fromArray(SP, a, i);
            FloatVector vb = FloatVector.fromArray(SP, b, i);
            acc = va.fma(vb, acc);                 // acc += va * vb, one instruction, one rounding
        }
        float sum = acc.reduceLanes(VectorOperators.ADD);
        for (; i < a.length; i++) sum += a[i] * b[i];   // tail for length % laneCount
        return sum;
    }

    public static void main(String[] args) {
        System.out.println("SPECIES lanes=" + SP.length() + " bits=" + SP.vectorBitSize());

        Random rnd = new Random(42);
        float[] query = randomVec(rnd);
        float[][] db = new float[N][];
        for (int i = 0; i < N; i++) db[i] = randomVec(rnd);

        // --- correctness: scalar and vector agree within a float delta (NOT exactly) ---
        float s = dotScalar(query, db[0]);
        float v = dotVector(query, db[0]);
        if (Math.abs(s - v) > 1e-3f * Math.max(1f, Math.abs(s)))
            throw new AssertionError("mismatch beyond tolerance: " + s + " vs " + v);
        System.out.printf("correctness OK (scalar=%.5f vector=%.5f, |diff|=%.2e)%n",
                          s, v, Math.abs(s - v));

        // --- warmup: let the JIT compile + intrinsify both paths before timing ---
        for (int w = 0; w < 20; w++) { runScalar(query, db); runVector(query, db); }

        // --- measure: average over several runs ---
        long scalarNs = timeMany(() -> runScalar(query, db), 10);
        long vectorNs = timeMany(() -> runVector(query, db), 10);

        double scalarPer = (double) scalarNs / N;
        double vectorPer = (double) vectorNs / N;
        System.out.printf("scalar: %8.1f ns/dotproduct%n", scalarPer);
        System.out.printf("vector: %8.1f ns/dotproduct%n", vectorPer);
        System.out.printf("speedup: %.2fx%n", scalarPer / vectorPer);
    }

    // Sum over the DB so the JIT cannot dead-code-eliminate the dot products.
    static float runScalar(float[] q, float[][] db) {
        float acc = 0f; for (float[] d : db) acc += dotScalar(q, d); return acc;
    }
    static float runVector(float[] q, float[][] db) {
        float acc = 0f; for (float[] d : db) acc += dotVector(q, d); return acc;
    }

    static float blackhole;  // sink so results are "used"
    static long timeMany(java.util.function.Supplier<Float> r, int runs) {
        long best = Long.MAX_VALUE;
        for (int i = 0; i < runs; i++) {
            long t0 = System.nanoTime();
            blackhole += r.get();
            long dt = System.nanoTime() - t0;
            best = Math.min(best, dt);     // best-of-N reduces GC/scheduler noise
        }
        return best;
    }

    static float[] randomVec(Random rnd) {
        float[] v = new float[DIM];
        for (int i = 0; i < DIM; i++) v[i] = rnd.nextFloat() * 2f - 1f;
        return v;
    }
}

Key harness choices, each defending against a classic microbench mistake:

• Warmup (20 iterations) before timing, so HotSpot has compiled and intrinsified the Panama ops. Timing cold code shows the JIT, not the algorithm.
• A blackhole sink that consumes the result, so the JIT cannot dead-code-eliminate the whole loop (a dot product whose result is unused can be deleted entirely → fake "infinite" speedup).
• Best-of-N timing to suppress GC pauses and scheduler jitter.
• Correctness within a delta, not exact — floating-point addition is not associative, so scalar and vector sums differ slightly. Asserting exact equality would be a flaky bug.

Step 3 — Run it WITH SIMD

javac --add-modules jdk.incubator.vector DotBench.java
java  --add-modules jdk.incubator.vector DotBench

Expected (an AVX2 laptop; your numbers vary with CPU and JDK):

SPECIES lanes=8 bits=256
correctness OK (scalar=3.21044 vector=3.21043, |diff|=8.34e-06)
scalar:    520.4 ns/dotproduct
vector:    142.7 ns/dotproduct
speedup: 3.65x

The |diff| line confirms the two agree within tolerance, and the speedup line is the prize: ~3–4× on AVX2, larger on AVX-512.

Step 4 — Run it WITHOUT SIMD (the control)

Now remove the module. Because DotBench references FloatVector, it will not load at all without --add-modules — which is itself the lesson (this is exactly why Lucene guards the Panama path):

java DotBench
# Error: ... module jdk.incubator.vector not found / NoClassDefFoundError: FloatVector

To get a clean scalar-only control number, comment out the vector path (or make a copy that only runs dotScalar) and run without the flag. The scalar ns/op should match the scalar: line from Step 3 — proving the speedup came from the vector path, not from the flag changing scalar performance. The takeaway: the flag is load-bearing; absent it, Lucene cannot even instantiate the Panama support and silently uses the scalar fallback.

Step 5 — Defeat the auto-vectorizer caveat

The JIT sometimes auto-vectorizes a simple scalar reduction on its own, which can shrink the apparent gap. Prove the win is real by comparing across -XX flags:

# Force scalar codegen for the scalar method's loop region to see the "true" scalar baseline:
java --add-modules jdk.incubator.vector -XX:-UseSuperWord DotBench
# -XX:-UseSuperWord disables HotSpot's auto-vectorizer (SuperWord). The scalar number should
# get *worse*, widening the gap — confirming part of "scalar" was being auto-vectorized.

This is the subtle part contributors miss: "scalar" Java is not always scalar — the JIT may auto-vectorize some loops. The Vector API's value is reliable, explicit SIMD for cases (like reductions over MemorySegment) the auto-vectorizer will not handle. Always state which baseline you measured against.

Step 6 — Confirm Lucene picked PanamaVectorUtilSupport

The microbench proves your kernel vectorizes. Now confirm Lucene's does, the way it runs inside OpenSearch. Two checks:

# (1) On a running OpenSearch node: is the module on the command line?
ps -ef | grep -i opensearch | grep -o "add-modules [^ ]*"
# expect: add-modules jdk.incubator.vector   (OpenSearch jvm.options enables it on JDK 21)

// (2) Ask Lucene which provider it chose. In a tiny program with lucene-core on the
//     classpath (or add this assertion inside a Lucene unit test):
import org.apache.lucene.internal.vectorization.VectorizationProvider;

public class WhichProvider {
    public static void main(String[] args) {
        System.out.println(VectorizationProvider.getInstance().getClass().getName());
        // -> ...PanamaVectorizationProvider   (SIMD, good)
        // -> ...DefaultVectorizationProvider  (scalar fallback, investigate!)
    }
}

# (the internal package is exported within Lucene; run from a Lucene checkout/test, or:)
java --add-modules jdk.incubator.vector \
     --add-exports java.base/jdk.internal.vm.vector=ALL-UNNAMED \
     -cp lucene-core-*.jar WhichProvider

PanamaVectorizationProvider confirms Lucene's VectorUtil.dotProduct — the exact kernel that scores every neural/knn query — is running SIMD. DefaultVectorizationProvider means the scalar fallback, and your vector search is paying the 2–8× tax.

Step 7 — Tie it back to the HNSW hot loop

Recall the call counts from the index:

Your DotBench measured one dot product at ~140 ns (vector) vs ~520 ns (scalar). Multiply by hundreds of millions for a merge: that ~380 ns/call difference becomes minutes of wall-clock per merge. This is why a change to VectorUtil produced (with an improved graph merger) ~25% indexing speedups in Lucene's nightly benchmarks — and why the distance kernel is the highest-leverage place in the entire vector stack to optimize. That is exactly the target of Capstone Project 04.

Deliverables

• SimdCheck output recording your CPU's preferred lane count.
• DotBench source and its output: scalar ns/op, vector ns/op, speedup, and the within-delta correctness line.
• The three-baseline comparison table (Step 5) with your measured numbers.
• Confirmation of which VectorizationProvider Lucene selected on your machine.

Troubleshooting

Expected Output

A measured, repeatable speedup of the Panama dot product over the scalar dot product (~2× on NEON, ~3–4× on AVX2, more on AVX-512), correctness agreement within a float delta, and confirmation that Lucene selected PanamaVectorizationProvider on your machine — i.e. the exact kernel behind every neural/knn query is running SIMD.

Stretch Goals

• Add squareDistance and cosine (scalar + Panama, from the index); benchmark all three. Cosine has three accumulators — does its speedup differ?
• Rewrite the benchmark as a proper JMH harness (@Benchmark, @Warmup, @Fork, Blackhole) and compare its numbers to your hand-rolled timer. JMH handles DCE and warmup rigorously.
• Add an int8/byte quantized dot product and compare ns/op — this is the kernel behind scalar-quantized k-NN (quantization).
• Load vectors from a MemorySegment instead of float[] and benchmark FloatVector.fromMemorySegment — the zero-copy path modern Lucene uses to score .vec data (SIMD chapter).
• Run on both an x86 and an ARM/Graviton box and tabulate lane width vs speedup; explain the difference from register width alone.

Coding Exercises

DotBench is a hand-rolled timer; these exercises make the measurement rigorous and connect it to the real kernels in Lucene and k-NN. They are Java, compiled against the Vector API (--add-modules jdk.incubator.vector) and, for exercise 4, JMH.

1. (warm-up) Add squareDistance and cosine, scalar + Panama. Extend DotBench with l2Scalar/l2Vector (sum of squared lane differences via va.sub(vb) then fma into the accumulator) and cosScalar/cosVector (three accumulators: dot, |a|², |b|²). Assert each vector path agrees with its scalar twin within a relative delta, then print ns/op for all three. Does cosine's three-accumulator loop show a different speedup than dot product? Explain from register pressure. The exact math is in vector-math-foundations.md.

2. (warm-up) Byte/int8 quantized dot product. Add dotInt8Scalar(byte[],byte[]) and a Panama version using ByteVector widened to ShortVector/IntVector. Benchmark ns/op and the speedup; this is the kernel behind scalar-quantized k-NN (see quantization-and-disk-ann.md). Assert correctness against a scalar reference within an exact integer match (int8 dot is integer arithmetic — no float-associativity excuse here).

3. (core) Benchmark Lucene's own VectorUtil against your kernels. With lucene-core on the classpath, call VectorUtil.dotProduct(a, b) and VectorUtil.squareDistance(a, b) in the same harness as your hand-written Panama versions; assert results agree within delta and tabulate ns/op. Then print VectorizationProvider.getInstance().getClass().getName() (Step 6) so the row is labeled with the provider Lucene actually used. You are now measuring the exact kernel every neural/knn query runs.

4. (core) Promote the harness to a real JMH benchmark. Rewrite DotBench as a JMH project: @State holding the vectors, @Benchmark methods scalar, panama, luceneVectorUtil, with @Warmup/@Measurement/@Fork(1) and a Blackhole parameter consuming each result. Run @BenchmarkMode(AverageTime) in ns/op and compare JMH's numbers to your hand-timer's — JMH handles DCE and warmup rigorously, so any large discrepancy means your hand-timer was lying. Deliverable: the JMH summary table for the three methods.

5. (advanced) Advanced challenge — benchmark the k-NN native SIMD kernel against the JVM Panama kernel. OpenSearch's faiss engine computes distances in native C++ with hand-written SIMD, not the JVM Vector API. In a opensearch-project/k-NN checkout, locate both the Java reference scorers and the native kernels: rg -n "l2Squared\|innerProduct\|cosinesimil" src/main/java/org/opensearch/knn/plugin/script/KNNScoringUtil.java (the Java side — l2Squared, innerProduct, cosinesimil, cosinesimilOptimized), and rg -ln "SIMD\|fp16\|fvec_L2sqr\|fvec_inner_product" jni/src jni/src/simd (the native AVX/NEON kernels). Write a JMH benchmark that times KNNScoringUtil's l2Squared(float[],float[]) and innerProduct(float[],float[]) over 768-dim vectors, and compare ns/op against your Panama and Lucene VectorUtil numbers from exercise 3 and 4. In your write-up, explain why OpenSearch keeps two distance stacks (JVM Panama for the lucene engine, native SIMD for faiss) — the architecture in native-simd-and-faiss-kernels.md. Deliverable: a four-row table (scalar / Panama / Lucene VectorUtil / k-NN KNNScoringUtil) of ns/op for the same 768-dim dot/L2, with one paragraph on which stack serves which engine and why.

Issues to Practice On

The distance kernel spans two repos: the JVM kernel lives in apache/lucene (VectorUtil, PanamaVectorUtilSupport); the native kernel and the Java scorers live in opensearch-project/k-NN (jni/src/simd, KNNScoringUtil). The two use different contribution workflows — Lucene: GitHub issues + PRs + a CHANGES.txt entry, no DCO; k-NN: PRs + a CHANGELOG.md entry + DCO Signed-off-by (git commit -s). Find work with:

# Lucene JVM kernel:
gh label list --repo apache/lucene | grep -iE "good first|core/|performance"   # confirm taxonomy (labels move)
gh issue list --repo apache/lucene --label "good first issue" --state open
gh issue list --repo apache/lucene --search "VectorUtil OR Panama OR SIMD in:title,body" --state open
# k-NN native kernel + scorers:
gh label list --repo opensearch-project/k-NN   # confirm taxonomy
gh issue list --repo opensearch-project/k-NN --label "good first issue" --state open
gh issue list --repo opensearch-project/k-NN --search "SIMD OR AVX OR fp16 OR distance in:title,body" --state open

Representative issue patterns:

• Kernel performance/correctness (Lucene or k-NN). "Square-distance slower than dot product," "ARM NEON path mis-handles the tail," "fp16 distance off." Approach: build a JMH repro (exercises 3–5), locate the kernel (rg "squareDistance\|fvec_L2sqr"), prove the regression with numbers, fix, PR with the right CHANGELOG/CHANGES entry.
• Provider selection / fallback bugs (Lucene). "DefaultVectorizationProvider chosen on a SIMD-capable box." Approach: reproduce with WhichProvider, locate VectorizationProvider.lookup(), add a test asserting the Panama provider on a supported JDK.

Planted-bug drill. In dotVector (Step 2), drop the scalar tail loop (delete for (; i < a.length; i++) sum += a[i]*b[i];). Because 768 is divisible by 8 and 16, the AVX result stays correct — so run it with DIM = 770 (not a multiple of any lane width) and watch your within-delta correctness check go red: the last two components are silently dropped. Restore the tail and confirm green. This is the most common real bug in hand-written SIMD kernels, and the assertion is exactly the guard that catches it — the kind of test every VectorUtil/KNNScoringUtil change must carry.

Etiquette: claim the issue first, reproduce with a benchmark before coding; Apache Lucene PRs need a test + CHANGES.txt (no DCO), OpenSearch k-NN PRs need a test + CHANGELOG.md entry + DCO Signed-off-by (git commit -s). See community-interaction.md and Capstone Project 04, where this kernel is the target.

Validation: Self-check

1. Why is the distance kernel the right place to optimize vector search? Quantify roughly how many times it runs for one query vs one 1M-vector merge.
2. Write, from memory, the Panama dot product: the lane-striding main loop, the fma, the reduceLanes, and the scalar tail. What does fma do that a separate multiply+add does not, and why is the tail mandatory?
3. Why must you assert correctness within a delta rather than exact equality between the scalar and vector results?
4. Name three ways a microbenchmark can lie about SIMD speedup (no warmup, DCE, auto-vectorized scalar baseline) and how your harness defends against each.
5. What single JDK flag gates the Vector API, and what happens to (a) your DotBench and (b) Lucene without it?
6. Give two independent ways to confirm Lucene is using PanamaVectorUtilSupport rather than the scalar fallback.
7. Connect your measured ns/op to the "~25% indexing speedup" claim: why is a distance- kernel change high-leverage?

When you can answer all seven and reproduce a measured speedup, you have proven the SIMD foundation of the entire vector stack from the bottom up. Return to SIMD and the Panama Vector API for the full theory, re-read the masterclass index to see how this kernel powers the neural query at the top, and take it upstream in Capstone Project 04: HNSW improvements to Lucene.