I finally got around to investigating the Java Vector API. This feature has been in the incubator for ages. It’s purpose is to optimize vector operations using SIMD instructions, presumably much faster than the performance from scalar operations. Lots of people have reported some pretty fantastic speedups, but unfortunately I’m not seeing them.

The task was simple. I wanted to replace this block of code in the inference phase:

        for (int i=0;i<NN_SIZE_L2;i++) {
            int sum = B1[i];
            for (int j=0;j<(NN_SIZE_L1*2);j++) {
                sum += W1[i * (NN_SIZE_L1*2) + j] * L1[j];
            }
            L2[i] = sum;
        }

With something that looks like this:

VectorSpecies<Integer> INT_SPEC = IntVector.SPECIES_256;
for (int i=0;i<NN_SIZE_L2;i++) {
    IntVector sum32 = IntVector.zero(INT_SPEC);
    for (int j=0;j<(NN_SIZE_L1*2);j+=INT_SPEC.length()) {
        IntVector inp = IntVector.fromArray(INT_SPEC, L1, j);
        IntVector wei = IntVector.fromArray(INT_SPEC, W1, i * (NN_SIZE_L1 * 2) + j);
        IntVector dot = inp.mul(wei);
        sum32 = sum32.add(dot);
    }
    L2[i] = sum32.reduceLanes(VectorOperators.ADD) + B1[i];
}

Notice how the inner loop gets incremented by INT_SPEC.length() , which happens to be 8 (lanes) in this case, since ints in Java are 32 bits, and the entire vector is 256 bits wide. The idea is that we will load up the vector and perform those multiplication operations in parallel, then do an ‘add’ reduction at the end. It works fine, but unfortunately is actually slightly *slower* than the non-vectorized implementation. I tried variations using ByteVectors and ShortVectors, but they were even worse.

I’ll take another look when it comes out of the incubation phase, and/or when I get a machine that supports AVX-512 instructions, but for now I’ll be sticking with the non-vectorized version.