[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.

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. 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. The branch removal happens too late, when the vectorizer is already done with this loop. So 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 - great writing.

To make the above algorithm branchless we don’t need cmov. 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

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.

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, AI was used to generate all the Vector API code.