import uk.ac.manchester.tornado.api.ImmutableTaskGraph;
import uk.ac.manchester.tornado.api.TaskGraph;
import uk.ac.manchester.tornado.api.TornadoExecutionPlan;
import uk.ac.manchester.tornado.api.annotations.Parallel;
import uk.ac.manchester.tornado.api.enums.DataTransferMode;
import uk.ac.manchester.tornado.api.types.arrays.FloatArray;

/** Square float matrix multiply as a TornadoVM task graph. Compiled, NOT executed (no GPU/OpenCL in the authoring sandbox). */
public class TornadoMatMul {

    // The two @Parallel loops tell TornadoVM it may map them onto a 2D grid of GPU threads.
    public static void matmul(FloatArray a, FloatArray b, FloatArray c, final int n) {
        for (@Parallel int i = 0; i < n; i++) {
            for (@Parallel int j = 0; j < n; j++) {
                float sum = 0f;
                for (int k = 0; k < n; k++) {
                    sum += a.get(i * n + k) * b.get(k * n + j);
                }
                c.set(i * n + j, sum);
            }
        }
    }

    public static void main(String[] args) throws Exception {
        final int n = args.length > 0 ? Integer.parseInt(args[0]) : 512;
        FloatArray a = new FloatArray(n * n);
        FloatArray b = new FloatArray(n * n);
        FloatArray c = new FloatArray(n * n);
        for (int i = 0; i < n * n; i++) {
            a.set(i, (i % 7) * 0.5f);
            b.set(i, (i % 5) * 0.25f);
        }

        TaskGraph graph = new TaskGraph("mm")
                .transferToDevice(DataTransferMode.FIRST_EXECUTION, a, b)
                .task("t0", TornadoMatMul::matmul, a, b, c, n)
                .transferToHost(DataTransferMode.EVERY_EXECUTION, c);

        ImmutableTaskGraph snapshot = graph.snapshot();
        try (TornadoExecutionPlan plan = new TornadoExecutionPlan(snapshot)) {
            for (int run = 0; run < 5; run++) {
                long t0 = System.nanoTime();
                plan.execute();
                System.out.printf("run %d: %.2f ms%n", run, (System.nanoTime() - t0) / 1e6);
            }
        }
        System.out.println("c[0]=" + c.get(0));
    }
}