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

import java.util.Arrays;
import java.util.Random;

/**
 * Single-precision n x n matrix multiply three ways on the CPU, one thread, no libraries:
 *   naive   i-j-k loops, B read down a column (cache-hostile)
 *   ikj     i-k-j loops, everything read along rows (cache-friendly, and the JIT may vectorise it itself)
 *   vector  the same i-k-j order written by hand with the Vector API (jdk.incubator.vector)
 *
 *   javac --add-modules jdk.incubator.vector MatMulBench.java
 *   java --add-modules jdk.incubator.vector MatMulBench 512
 */
public class MatMulBench {

    static final VectorSpecies<Float> SP = FloatVector.SPECIES_PREFERRED;

    static void naive(float[] a, float[] b, float[] c, int n) {
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                float sum = 0;
                for (int k = 0; k < n; k++) {
                    sum += a[i * n + k] * b[k * n + j];
                }
                c[i * n + j] = sum;
            }
        }
    }

    static void ikj(float[] a, float[] b, float[] c, int n) {
        Arrays.fill(c, 0f);
        for (int i = 0; i < n; i++) {
            for (int k = 0; k < n; k++) {
                float aik = a[i * n + k];
                for (int j = 0; j < n; j++) {
                    c[i * n + j] += aik * b[k * n + j];
                }
            }
        }
    }

    static void vector(float[] a, float[] b, float[] c, int n) {
        Arrays.fill(c, 0f);
        int upper = SP.loopBound(n);
        for (int i = 0; i < n; i++) {
            for (int k = 0; k < n; k++) {
                float aik = a[i * n + k];
                FloatVector va = FloatVector.broadcast(SP, aik);
                int j = 0;
                for (; j < upper; j += SP.length()) {
                    FloatVector vb = FloatVector.fromArray(SP, b, k * n + j);
                    FloatVector vc = FloatVector.fromArray(SP, c, i * n + j);
                    va.fma(vb, vc).intoArray(c, i * n + j);
                }
                for (; j < n; j++) {
                    c[i * n + j] += aik * b[k * n + j];
                }
            }
        }
    }

    interface Kernel { void run(float[] a, float[] b, float[] c, int n); }

    static double medianMillis(Kernel k, float[] a, float[] b, float[] c, int n, int warmups, int runs) {
        for (int i = 0; i < warmups; i++) {
            k.run(a, b, c, n);
        }
        double[] t = new double[runs];
        for (int i = 0; i < runs; i++) {
            long s = System.nanoTime();
            k.run(a, b, c, n);
            t[i] = (System.nanoTime() - s) / 1e6;
        }
        Arrays.sort(t);
        return t[runs / 2];
    }

    public static void main(String[] args) {
        int n = args.length > 0 ? Integer.parseInt(args[0]) : 512;
        var rnd = new Random(42);
        float[] a = new float[n * n], b = new float[n * n];
        for (int i = 0; i < a.length; i++) {
            a[i] = rnd.nextFloat();
            b[i] = rnd.nextFloat();
        }
        float[] c1 = new float[n * n], c2 = new float[n * n], c3 = new float[n * n];

        System.out.println("n=" + n + " vector species=" + SP + " lanes=" + SP.length());
        double tn = medianMillis(MatMulBench::naive, a, b, c1, n, 3, 7);
        double ti = medianMillis(MatMulBench::ikj, a, b, c2, n, 3, 7);
        double tv = medianMillis(MatMulBench::vector, a, b, c3, n, 3, 7);

        float d12 = 0, d13 = 0;
        for (int i = 0; i < c1.length; i++) {
            d12 = Math.max(d12, Math.abs(c1[i] - c2[i]));
            d13 = Math.max(d13, Math.abs(c1[i] - c3[i]));
        }
        double gflop = 2.0 * n * n * n / 1e9;
        System.out.printf("naive  median %8.1f ms  %6.2f GFLOP/s%n", tn, gflop / (tn / 1000));
        System.out.printf("ikj    median %8.1f ms  %6.2f GFLOP/s  (%.1fx naive)%n", ti, gflop / (ti / 1000), tn / ti);
        System.out.printf("vector median %8.1f ms  %6.2f GFLOP/s  (%.1fx naive, %.2fx ikj)%n", tv, gflop / (tv / 1000), tn / tv, ti / tv);
        System.out.printf("max abs difference vs naive: ikj %.2e, vector %.2e%n", d12, d13);
    }
}