Refactor TornadoVM initialization metrics tracking to include finer granularity across plan creation, JIT, and weight transfer stages

Author: orionpapadakis

src/main/java/org/beehive/gpullama3/auxiliary/RunMetrics.java

@@ -29,8 +29,9 @@ public final class RunMetrics {
     private boolean hasPrefillPhase;
 
     // ── TornadoVM-specific metrics (nanoseconds) ──────────────────────────────
-    private long tornadoCompileDurationNs;
-    private long tornadoWarmupDurationNs;
+    private long tornadoPlanCreationNs;
+    private long tornadoJitNs;
+    private long readOnlyWeightsCopyInNs;
 
     // ── Singleton ─────────────────────────────────────────────────────────────
     private static final RunMetrics INSTANCE = new RunMetrics();
@@ -47,12 +48,14 @@ public static void setLoadDuration(long ns) {
     /**
      * Records TornadoVM-specific initialisation durations.
      *
-     * @param compileNs plan-graph construction + JIT compilation ({@code withPreCompilation()})
-     * @param warmupNs  first-execution weight upload ({@code forceCopyInReadOnlyData()})
+     * @param planCreationNs  task-graph construction ({@code createExecutionPlan()})
+     * @param jitNs           JIT compilation ({@code withPreCompilation()})
+     * @param weightCopyNs    first-execution weight upload ({@code forceCopyInReadOnlyData()})
      */
-    public static void setTornadoMetrics(long compileNs, long warmupNs) {
-        INSTANCE.tornadoCompileDurationNs = compileNs;
-        INSTANCE.tornadoWarmupDurationNs  = warmupNs;
+    public static void setTornadoMetrics(long planCreationNs, long jitNs, long weightCopyNs) {
+        INSTANCE.tornadoPlanCreationNs    = planCreationNs;
+        INSTANCE.tornadoJitNs             = jitNs;
+        INSTANCE.readOnlyWeightsCopyInNs  = weightCopyNs;
     }
 
     /**
@@ -85,13 +88,13 @@ public static void setHasPrefillPhase(boolean value) {
 
     // ── Output ────────────────────────────────────────────────────────────────
 
-    /** Prints throughput metrics to {@code stderr}. */
+    /** Prints throughput metrics to {@code stderr}, and TornadoVM init metrics when enabled. */
     public static void printMetrics() {
         RunMetrics m = INSTANCE;
 
-        int    totalTokens  = m.promptEvalCount + m.evalCount;
-        double totalSecs    = m.totalDurationNs / 1e9;
-        double totalRate    = totalSecs > 0 ? totalTokens / totalSecs : 0;
+        int    totalTokens = m.promptEvalCount + m.evalCount;
+        double totalSecs   = m.totalDurationNs / 1e9;
+        double totalRate   = totalSecs > 0 ? totalTokens / totalSecs : 0;
 
         System.err.println();
         System.err.println("==== Performance Metrics ====");
@@ -113,6 +116,17 @@ public static void printMetrics() {
             System.err.printf("achieved tok/s: %.2f. Tokens: %d, seconds: %.2f%n",
                     totalRate, totalTokens, totalSecs);
         }
+
+        if (Boolean.parseBoolean(System.getProperty("llama.EnableTimingForTornadoVMInit", "false"))
+                && m.tornadoPlanCreationNs > 0) {
+            System.err.printf(
+                    "Compilation & CodeGen: %.2f ms%n" +
+                    "Warmup: %.2f ms%n" +
+                    "Read-only weights Copy-in: %.2f ms%n",
+                    m.tornadoPlanCreationNs   / 1_000_000.0,
+                    m.tornadoJitNs            / 1_000_000.0,
+                    m.readOnlyWeightsCopyInNs / 1_000_000.0);
+        }
         System.err.println();
     }
 }

src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanStandard.java

@@ -47,14 +47,7 @@ public TornadoVMMasterPlanStandard(State state, Model model) {
         forceCopyInReadOnlyData();
         long copyTime = System.nanoTime();
 
-        RunMetrics.setTornadoMetrics(warmupTime - startTime, copyTime - warmupTime);
-
-        if (ENABLE_TORNADOVM_INIT_TIME) {
-            System.err.printf("TornadoVM GPU standard execution plan creation: %.2f ms\n", (planCreationTime - startTime) / 1_000_000.0);
-            System.err.printf("Java to GPU JIT compiler warmup: %.2f ms\n", (warmupTime - planCreationTime) / 1_000_000.0);
-            System.err.printf("Transfer read-only weights to GPU: %.2f ms\n", (copyTime - warmupTime) / 1_000_000.0);
-            System.err.printf("Finished TornadoVM initialization...\n \n");
-        }
+        RunMetrics.setTornadoMetrics(planCreationTime - startTime, warmupTime - planCreationTime, copyTime - warmupTime);
     }
 
     /**

src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanWithBatchPrefillDecode.java

@@ -97,14 +97,7 @@ public class TornadoVMMasterPlanWithBatchPrefillDecode implements TornadoVMMaste
         forceCopyInReadOnlyData();
         long copyTime = System.nanoTime();
 
-        RunMetrics.setTornadoMetrics(warmupTime - startTime, copyTime - warmupTime);
-
-        if (ENABLE_TORNADOVM_INIT_TIME) {
-            System.err.printf("TornadoVM GPU batched prefill/decode execution plan creation: %.2f ms\n", (planCreationTime - startTime) / 1_000_000.0);
-            System.err.printf("Java to GPU JIT compiler warmup: %.2f ms\n", (warmupTime - planCreationTime) / 1_000_000.0);
-            System.err.printf("Transfer read-only weights to GPU: %.2f ms\n", (copyTime - warmupTime) / 1_000_000.0);
-            System.err.printf("Finished TornadoVM initialization...\n \n");
-        }
+        RunMetrics.setTornadoMetrics(planCreationTime - startTime, warmupTime - planCreationTime, copyTime - warmupTime);
     }
 
     // ── Batch Prefill Activation graphs ─────────────────────────────────────────────────────

src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanWithPrefillDecode.java

@@ -85,17 +85,7 @@ public class TornadoVMMasterPlanWithPrefillDecode implements TornadoVMMasterPlan
         forceCopyInReadOnlyData();
         long copyTime = System.nanoTime();
 
-        RunMetrics.setTornadoMetrics(warmupTime - startTime, copyTime - warmupTime);
-
-        if (ENABLE_TORNADOVM_INIT_TIME) {
-            System.err.printf("TornadoVM GPU single-token prefill/decode execution plan creation: %.2f ms\n",
-                    (planCreationTime - startTime) / 1_000_000.0);
-            System.err.printf("Java to GPU JIT compiler warmup: %.2f ms\n",
-                    (warmupTime - planCreationTime) / 1_000_000.0);
-            System.err.printf("Transfer read-only weights to GPU: %.2f ms\n",
-                    (copyTime - warmupTime) / 1_000_000.0);
-            System.err.printf("Finished TornadoVM initialization...\n \n");
-        }
+        RunMetrics.setTornadoMetrics(planCreationTime - startTime, warmupTime - planCreationTime, copyTime - warmupTime);
     }
 
     // ── Activation graph ─────────────────────────────────────────────────────