Add square-and-add kernel (the auto-vectorized case)

timings/src/main/java/com/ankurm/vectorapi/SquareAdd.java

@@ -0,0 +1,50 @@
+package com.ankurm.vectorapi;
+
+import jdk.incubator.vector.*;
+
+/** c[i] = -(a[i]^2 + b[i]^2) -- the loop C2's SuperWord pass handles well. */
+public final class SquareAdd {
+
+    static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
+
+    static void scalarComputation(float[] a, float[] b, float[] c) {
+        for (int i = 0; i < a.length; i++) {
+            c[i] = -(a[i] * a[i] + b[i] * b[i]);
+        }
+    }
+
+    static void vectorComputation(float[] a, float[] b, float[] c) {
+        int i = 0;
+        int upperBound = SPECIES.loopBound(a.length);   // e.g. 8192 -> 8192, 8190 -> 8184
+
+        // Main loop: process SPECIES.length() elements per iteration
+        for (; i < upperBound; i += SPECIES.length()) {
+            var va = FloatVector.fromArray(SPECIES, a, i);
+            var vb = FloatVector.fromArray(SPECIES, b, i);
+            va.mul(va).add(vb.mul(vb)).neg().intoArray(c, i);
+        }
+
+        // Scalar tail: the last (a.length % laneCount) elements
+        for (; i < a.length; i++) {
+            c[i] = -(a[i] * a[i] + b[i] * b[i]);
+        }
+    }
+
+    static void run(int n) {
+        float[] a = Data.randomFloats(n, 1);
+        float[] b = Data.randomFloats(n, 2);
+        float[] c1 = new float[n], c2 = new float[n];
+
+        scalarComputation(a, b, c1);
+        vectorComputation(a, b, c2);
+        Data.assertSame("square-add", c1, c2);
+
+        System.out.println("--- kernel 1: c[i] = -(a[i]^2 + b[i]^2), " + n + " floats ---");
+        double s = Bench.measure("scalar for-loop", () -> scalarComputation(a, b, c1));
+        double v = Bench.measure("Vector API", () -> vectorComputation(a, b, c2));
+        Bench.speedup("vector advantage", s, v);
+        System.out.println();
+    }
+
+    private SquareAdd() {}
+}