[Java][JVM][JIT] JDK 26 and SIMD operations

SIMD - single instruction, multiple data

If you’re already familiar with SIMD, feel free to skip to the next chapter.

I believe everyone has a good feeling for how traditional (scalar) processing works. The CPU simply executes one operation on one set of data, like a + b. But what can we do to speed it up when we want to do the same operation on multiple data sets? We can rely on CPU pipelining, out-of-order execution and superscalarity, but we can go beyond that. That’s where SIMD comes into play.

Instead of doing one operation on one set of data, we can put multiple data into vectors and apply one instruction to those. That’s the whole magic.

A vector here is not an arbitrary-length list - it is a CPU register of a fixed width. How many numbers fit into it is simply the register width divided by the size of a single element, and each of those slots is called a lane. A 128-bit register holds 4 ints, a 256-bit one holds 8. This is also the reason why the very same Java code can do twice as much work per instruction on one CPU than on another - something we will come back to in the last chapter.

SIMD in Java

Java offers us two ways of using SIMD instructions:

The JIT part is great because our Java application can run faster without any additional work on the programmer’s side. But can it produce code as good as what we can write with the Vector API? That’s what we will test in this article.

Environment

All the benchmarks in this article were run on:

Simple synthetic example

Let’s start with simple number adding:

@Benchmark
public int scalar() {
    int sum = 0;
    for (int i = 0; i < SIZE; i++) {
        sum += array[i];
    }
    return sum;
}

If we want to check if auto-vectorization was applied to our code and measure the speed-up, there is a simple trick. You can disable that optimization with -XX:-UseSuperWord JVM flag. Comparing performance with and without the flag gives you the answers to those questions.

Benchmark results for the scalar method on an M4 Max CPU, for an array of size 1_000_000:

SumBenchmark.scalar                      avgt    6  113,049 ± 3,588  us/op
SumNoSIMDBenchmark.scalar                avgt    6  228,635 ± 2,849  us/op

Auto-vectorization speeds up the method x2. A great result, especially since we don’t have to do anything to get it. Let’s now write the same algorithm using the Vector API. For simple algorithms it’s an easy task; for more complex ones you can help yourself with AI, it’s really good at transforming methods to the Vector API.

private static final VectorSpecies<Integer> SPECIES = IntVector.SPECIES_PREFERRED;

@Benchmark
public int vectorApiOneAccumulator() {
    int lanes = SPECIES.length();
    int bound = SIZE - (SIZE % lanes);
    IntVector acc = IntVector.zero(SPECIES);
    for (int i = 0; i < bound; i += lanes) {
        acc = acc.add(IntVector.fromArray(SPECIES, array, i));
    }
    int sum = acc.reduceLanes(VectorOperators.ADD);
    for (int i = bound; i < SIZE; i++) {
        sum += array[i];
    }
    return sum;
}

The algorithm here is pretty straightforward. We have one accumulator, IntVector acc, filled with zeros, and in the loop we add vectors that are parts of the input array. In the end we need to add the “leftovers” that didn’t fit into a vector. The performance:

SumBenchmark.vectorApiOneAccumulator     avgt    6  113,832 ± 1,114  us/op

is the same as the auto-vectorization version, but can we get more?

If you look at the vector adding loop and unroll it:

acc = acc.add(IntVector.fromArray(SPECIES, array, i));
acc = acc.add(IntVector.fromArray(SPECIES, array, i + lanes));
acc = acc.add(IntVector.fromArray(SPECIES, array, i + (2 * lanes)));
acc = acc.add(IntVector.fromArray(SPECIES, array, i + (3 * lanes)));

you can notice that the CPU cannot speed it up any more with its magic, because we have a data dependency here. Each line needs the result from the previous line to do its work. But nobody is forcing us to use only one accumulator, we can use more. Here is an example with four accumulators:

@Benchmark
public int vectorApiFourAccumulators() {
    int lanes = SPECIES.length();
    int step = lanes * 4;
    int bound = SIZE - (SIZE % step);
    IntVector v0 = IntVector.zero(SPECIES);
    IntVector v1 = IntVector.zero(SPECIES);
    IntVector v2 = IntVector.zero(SPECIES);
    IntVector v3 = IntVector.zero(SPECIES);
    for (int i = 0; i < bound; i += step) {
        v0 = v0.add(IntVector.fromArray(SPECIES, array, i));
        v1 = v1.add(IntVector.fromArray(SPECIES, array, i + lanes));
        v2 = v2.add(IntVector.fromArray(SPECIES, array, i + (2 * lanes)));
        v3 = v3.add(IntVector.fromArray(SPECIES, array, i + (3 * lanes)));
    }
    int sum = v0.add(v1).add(v2).add(v3).reduceLanes(VectorOperators.ADD);
    for (int i = bound; i < SIZE; i++) {
        sum += array[i];
    }
    return sum;
}

A simple idea: four accumulators adding parts of the input array, followed by a sum over all accumulators and a sum of the leftovers. Here are the results for 2, 4 and 8 accumulators:

SumBenchmark.vectorApiTwoAccumulators    avgt    6   56,954 ± 0,963  us/op
SumBenchmark.vectorApiFourAccumulators   avgt    6   34,653 ± 0,732  us/op
SumBenchmark.vectorApiEightAccumulators  avgt    6   34,307 ± 0,464  us/op

The speedups here are as follows:

So with the Vector API we can speed up our algorithm by ~x3.3. Downsides? First of all it’s not automatic, we need to code it ourselves. Secondly, the Vector API version is far less readable than the scalar one. This downside will be even more visible in the next chapter.

Another example

This example is not synthetic at all, although it is very similar to the previous one. Some time ago I had to speed up an algorithm whose first part had to count negative numbers in an array. The scalar code is simple:

@Benchmark
public int branchy() {
    int negativeCount = 0;
    for (int v : array) {
        if (v < 0) {
            negativeCount++;
        }
    }
    return negativeCount;
}

Benchmark results for 1_000_000 random elements with ~50% of negative numbers:

CountBenchmark.branchy                     avgt    6  231,072 ± 1,656  us/op
CountNoSIMDBenchmark.branchy               avgt    6  236,232 ± 2,184  us/op

Auto-vectorization cannot really help here, because we have a branch in our code. Interestingly, JIT is able to get rid of that branch on its own - it compiles if (v < 0) { negativeCount++; } into plain arithmetic, but it cannot remove the branch and vectorize the loop at the same time. We need to get rid of the branch ourselves and figure out a branchless version of this algorithm. It’s also a good time to recommend a great article written by Jaromir Hamala: The Most Expensive Instruction Might Be… cmov.

To make this method branchless we just need one observation: every negative int has its leftmost bit set to 1. So if we shift a number 31 bits to the right, we will get 1 if the number is negative and 0 if the number is >= 0. With that observation we can go with:

@Benchmark
public int branchless() {
    int negativeCount = 0;
    for (int v : array) {
        negativeCount += v >>> 31;
    }
    return negativeCount;
}

Benchmark results:

CountBenchmark.branchless                  avgt    6  168,414 ± 0,250  us/op
CountNoSIMDBenchmark.branchless            avgt    6  228,633 ± 2,312  us/op

Nice speedup. If auto-vectorization can help, it means that we can also try to write the algorithm using the Vector API. Again, AI is a very good assistant here:

@Benchmark
public int vectorApiOneAccumulator() {
    int[] array = this.array;
    int lanes = SPECIES.length();
    int bound = SIZE - (SIZE % lanes);
    IntVector acc = IntVector.zero(SPECIES);
    for (int i = 0; i < bound; i += lanes) {
        acc = acc.add(IntVector.fromArray(SPECIES, array, i).lanewise(VectorOperators.LSHR, 31));
    }
    int negativeCount = acc.reduceLanes(VectorOperators.ADD);
    for (int i = bound; i < SIZE; i++) {
        negativeCount += array[i] >>> 31;
    }
    return negativeCount;
}

The idea is very similar to the previous example. We have one accumulator, we start with zeros and we add a part of the input array shifted by 31 bits. VectorOperators.LSHR here is the vector equivalent of the >>> operation.

This algorithm has the same flaw as the previous single-accumulator one - we have a data dependency when we unroll the loop. Same flaw, so the same trick - multiple accumulators:

@Benchmark
public int vectorApiFourAccumulators() {
    int[] array = this.array;
    int lanes = SPECIES.length();
    int step = lanes * 4;
    int bound = SIZE - (SIZE % step);
    IntVector v0 = IntVector.zero(SPECIES);
    IntVector v1 = IntVector.zero(SPECIES);
    IntVector v2 = IntVector.zero(SPECIES);
    IntVector v3 = IntVector.zero(SPECIES);
    for (int i = 0; i < bound; i += step) {
        v0 = v0.add(IntVector.fromArray(SPECIES, array, i).lanewise(VectorOperators.LSHR, 31));
        v1 = v1.add(IntVector.fromArray(SPECIES, array, i + lanes).lanewise(VectorOperators.LSHR, 31));
        v2 = v2.add(IntVector.fromArray(SPECIES, array, i + (2 * lanes)).lanewise(VectorOperators.LSHR, 31));
        v3 = v3.add(IntVector.fromArray(SPECIES, array, i + (3 * lanes)).lanewise(VectorOperators.LSHR, 31));
    }
    int negativeCount = v0.add(v1).add(v2).add(v3).reduceLanes(VectorOperators.ADD);
    for (int i = bound; i < SIZE; i++) {
        negativeCount += array[i] >>> 31;
    }
    return negativeCount;
}

Results:

CountBenchmark.vectorApiOneAccumulator     avgt    6  168,420 ± 0,403  us/op
CountBenchmark.vectorApiTwoAccumulators    avgt    6   84,183 ± 0,125  us/op
CountBenchmark.vectorApiFourAccumulators   avgt    6   43,338 ± 0,124  us/op
CountBenchmark.vectorApiEightAccumulators  avgt    6   34,263 ± 0,269  us/op

The overall speedup between the auto-vectorization version and the Vector API one is up to x4.92.

CPU matters - a lot

All those results are from my M4 Max machine. The speedup really depends on the type of CPU you have and the size of its vector registers. On my old i7-9850H the results are completely different:

SumBenchmark.scalar                        avgt    6   54,964 ±  3,390  us/op
SumNoSIMDBenchmark.scalar                  avgt    6  228,140 ±  3,094  us/op

SumBenchmark.vectorApiOneAccumulator       avgt    6   57,408 ± 11,323  us/op
SumBenchmark.vectorApiTwoAccumulators      avgt    6   58,190 ±  7,644  us/op
SumBenchmark.vectorApiFourAccumulators     avgt    6   53,996 ±  3,135  us/op
SumBenchmark.vectorApiEightAccumulators    avgt    6   54,350 ±  2,284  us/op

CountBenchmark.branchy                     avgt    6  256,884 ±  6,028  us/op
CountNoSIMDBenchmark.branchy               avgt    6  262,384 ±  2,414  us/op

CountBenchmark.branchless                  avgt    6   57,695 ±  5,557  us/op
CountNoSIMDBenchmark.branchless            avgt    6  232,836 ±  4,617  us/op

CountBenchmark.vectorApiOneAccumulator     avgt    6   62,872 ± 21,866  us/op
CountBenchmark.vectorApiTwoAccumulators    avgt    6   53,855 ±  0,679  us/op
CountBenchmark.vectorApiFourAccumulators   avgt    6   60,530 ± 12,841  us/op
CountBenchmark.vectorApiEightAccumulators  avgt    6   57,576 ±  0,616  us/op

On this CPU there is a massive difference between the auto-vectorization and scalar algorithms, and no difference between the single- and multiple-accumulator versions. The reason for both sits in the architecture of the two CPUs:

  M4 Max (aarch64, ASIMD/NEON) i7-9850H (x86-64, AVX2)
Vector register width 128 bits 256 bits
int lanes per vector 4 8
Vector int add latency 2 cycles? 1 cycle

Two things follow from that table.

The vector width explains the difference in the auto-vectorization speedup. One AVX2 instruction chews through twice as many ints as one NEON instruction, and that is exactly what we measured: x2,02 on the M4 Max against x4,15 on the i7-9850H. You can check the width your JVM picked with IntVector.SPECIES_PREFERRED.

The latency explains the accumulators. On the M4 Max each add needs 2 cycles before its result can feed the next one, so a single accumulator leaves the vector units half idle and adding a second accumulator fills the gap - hence the clean x2. On the i7-9850H the add takes a single cycle, so one accumulator already issues one add per cycle. There is no idle latency left to hide and extra accumulators cannot buy anything, which is exactly what the results above show.

So don’t take any of those numbers for granted and run your benchmarks on your own target (production) machine.

Specs used above:

Here is a similar specification for production grade CPUs:

  Graviton5 (Neoverse V3) Xeon 6 P-core (Granite Rapids) Xeon 6 E-core (Clearwater Forest)
Vector register width 128 bits 512 bits 256 bits
int lanes per vector 4 16 8
Vector int add latency 2 cycles 1 cycle? 1 cycle?

Specs for that table:

Conclusions

Overall, the Vector API is great and can speed up our code significantly. The caveat is that such code is much less readable than the original. If this API can give a boost to a latency-sensitive part of your app, you can/should give it a go. There are also algorithms that auto-vectorization simply misses, and then the Vector API is all you have if you want to use SIMD. JEP 529 names two of them - Arrays::hashCode and lexicographic array comparison.

AI Usage in this article

First of all, check my index page where I explained AI involvement in the content of the articles. Besides polishing it (without any heavy rewrites this time), AI was used to: