feat(memory): trace exporters — Perfetto/Chrome JSON, JFR events, android.os.Trace (SKEEP-003 P2, S1.5b)

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -1294,12 +1294,27 @@ public final class sk/ainet/lang/memory/trace/CompositeTraceSink : sk/ainet/lang
 	public fun isEnabled ()Z
 }
 
+public final class sk/ainet/lang/memory/trace/JfrTraceSink : sk/ainet/lang/memory/trace/TraceSink {
+	public fun <init> ()V
+	public final fun anyEventEnabled ()Z
+	public fun emit (Lsk/ainet/lang/memory/trace/TraceEvent;)V
+	public fun isEnabled ()Z
+}
+
 public final class sk/ainet/lang/memory/trace/NoopTraceSink : sk/ainet/lang/memory/trace/TraceSink {
 	public static final field INSTANCE Lsk/ainet/lang/memory/trace/NoopTraceSink;
 	public fun emit (Lsk/ainet/lang/memory/trace/TraceEvent;)V
 	public fun isEnabled ()Z
 }
 
+public final class sk/ainet/lang/memory/trace/PerfettoTraceExporter {
+	public static final field INSTANCE Lsk/ainet/lang/memory/trace/PerfettoTraceExporter;
+	public final fun export (Ljava/util/List;Ljava/lang/String;)Ljava/lang/String;
+	public final fun export (Lsk/ainet/lang/memory/trace/RecordingTraceSink;Ljava/lang/String;)Ljava/lang/String;
+	public static synthetic fun export$default (Lsk/ainet/lang/memory/trace/PerfettoTraceExporter;Ljava/util/List;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/String;
+	public static synthetic fun export$default (Lsk/ainet/lang/memory/trace/PerfettoTraceExporter;Lsk/ainet/lang/memory/trace/RecordingTraceSink;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/String;
+}
+
 public final class sk/ainet/lang/memory/trace/PlanTraceKt {
 	public static final fun emit (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/trace/TraceSink;)V
 }

skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/trace/AndroidTraceSink.kt

@@ -0,0 +1,49 @@
+package sk.ainet.lang.memory.trace
+
+import android.os.Trace
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * A [TraceSink] that forwards SKaiNET's events to `android.os.Trace`, so a memory-architecture
+ * trace lines up with the OS view in Perfetto / systrace on device (SKEEP-003 §4.9, decision #12).
+ *
+ * Phases and kernel runs become async sections (they can nest and overlap across steps); adapters
+ * and scope resets become instants; allocation totals become counters, which is what makes a flat
+ * decode loop visible next to the app's own memory graph.
+ *
+ * Section names are truncated to `android.os.Trace`'s limit (127 characters), and tracing is only
+ * enabled while the OS is capturing, so a shipping build pays one boolean check per event.
+ */
+@ExperimentalMemoryApi
+public class AndroidTraceSink : TraceSink {
+
+    override val isEnabled: Boolean get() = Trace.isEnabled()
+
+    override fun emit(event: TraceEvent) {
+        if (!Trace.isEnabled()) return
+        when (event) {
+            is TraceEvent.PhaseBegin -> Trace.beginAsyncSection(name(event.phase, event.step), cookie(event.phase, event.step))
+            is TraceEvent.PhaseEnd -> Trace.endAsyncSection(name(event.phase, event.step), cookie(event.phase, event.step))
+            is TraceEvent.KernelRun -> {
+                // the kernel already ran; represent it as a zero-length async pair so it shows on the timeline
+                val n = truncate("kernel ${event.op}:${event.kernel}")
+                val c = (event.output?.canonical ?: event.kernel).hashCode()
+                Trace.beginAsyncSection(n, c); Trace.endAsyncSection(n, c)
+            }
+            is TraceEvent.AdapterInserted -> Trace.setCounter(truncate("skainet adapter ${event.kind} bytes"), event.bytes)
+            is TraceEvent.Allocation -> Trace.setCounter(truncate("skainet ${event.scope.name.lowercase()} alloc bytes"), event.bytes)
+            is TraceEvent.Free -> Trace.setCounter(truncate("skainet ${event.scope.name.lowercase()} free bytes"), event.bytes)
+            is TraceEvent.ScopeReset -> Trace.setCounter(truncate("skainet ${event.scope.name.lowercase()} live bytes"), event.liveBytesAfter)
+            is TraceEvent.Counter -> Trace.setCounter(truncate("skainet ${event.name}"), event.value)
+            is TraceEvent.Plan -> Trace.setCounter(truncate("skainet plan total bytes"), event.totalBytes)
+        }
+    }
+
+    private fun name(phase: String, step: Int?): String = truncate(if (step == null) "skainet $phase" else "skainet $phase#$step")
+    private fun cookie(phase: String, step: Int?): Int = phase.hashCode() * 31 + (step ?: 0)
+
+    /** `android.os.Trace` rejects names longer than 127 characters. */
+    private fun truncate(s: String): String = if (s.length <= MAX_SECTION_NAME) s else s.substring(0, MAX_SECTION_NAME)
+
+    private companion object { const val MAX_SECTION_NAME: Int = 127 }
+}

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt

@@ -0,0 +1,143 @@
+package sk.ainet.lang.memory.trace
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.ScopeKind
+
+/**
+ * Renders a [TraceEvent] stream as a Chrome/Perfetto trace (the JSON array format Perfetto and
+ * `chrome://tracing` both read) — SKEEP-003 §4.9, PRD M1-F7 / M1-A7.
+ *
+ * The mapping is the one the design asks for:
+ * - **one track per scope** — phases and kernels on the main thread, allocations on a thread named
+ *   after their [ScopeKind], so `Forward` and `Model` allocations are visually separate;
+ * - **kernel runs and phases as duration slices**, labelled by op and `TensorId`s;
+ * - **adapter insertions as instant events** with their byte cost, so a silent dequantisation shows
+ *   up as a mark on the timeline (the #782 class);
+ * - **live bytes per scope as counter tracks**, which is what makes a flat decode loop *look* flat;
+ * - the memory plan as metadata on the trace.
+ *
+ * Timestamps are microseconds (Perfetto's unit) derived from the events' nanosecond clock.
+ */
+@ExperimentalMemoryApi
+public object PerfettoTraceExporter {
+
+    private const val PID: Int = 1
+    private const val MAIN_TID: Int = 1
+
+    /** Render [events] as a Chrome trace JSON document. */
+    public fun export(events: List<TraceEvent>, processName: String = "skainet"): String {
+        val out = StringBuilder(1024)
+        out.append("{\"traceEvents\":[\n")
+        var first = true
+        fun emit(json: String) {
+            if (!first) out.append(",\n")
+            out.append(json); first = false
+        }
+
+        emit(metadata(PID, MAIN_TID, "process_name", processName))
+        emit(metadata(PID, MAIN_TID, "thread_name", "phases & kernels"))
+        val scopeTids = HashMap<ScopeKind, Int>()
+        fun tidOf(scope: ScopeKind): Int = scopeTids.getOrPut(scope) {
+            val tid = MAIN_TID + 1 + scopeTids.size
+            emit(metadata(PID, tid, "thread_name", "${scope.name.lowercase()} scope"))
+            tid
+        }
+
+        // live bytes per scope, updated as allocations come and go
+        val live = HashMap<ScopeKind, Long>()
+
+        for (e in events) {
+            val ts = e.timeNanos / 1000.0
+            when (e) {
+                is TraceEvent.PhaseBegin -> emit(slice("B", e.phase + (e.step?.let { "#$it" } ?: ""), "phase", ts, MAIN_TID, e.attributes))
+                is TraceEvent.PhaseEnd -> emit(slice("E", e.phase + (e.step?.let { "#$it" } ?: ""), "phase", ts, MAIN_TID, emptyMap()))
+                is TraceEvent.KernelRun -> {
+                    val args = buildMap {
+                        put("kernel", e.kernel)
+                        put("bytesRead", e.bytesRead.toString())
+                        put("bytesWritten", e.bytesWritten.toString())
+                        e.inputs.forEachIndexed { i, id -> if (id != null) put("in$i", id.canonical) }
+                        e.output?.let { put("out", it.canonical) }
+                    }
+                    emit(complete(e.op, "kernel", (e.timeNanos - e.durationNanos) / 1000.0, e.durationNanos / 1000.0, MAIN_TID, args))
+                }
+                is TraceEvent.AdapterInserted -> emit(
+                    instant(
+                        "adapter:${e.kind}", "adapter", ts, MAIN_TID,
+                        mapOf("from" to e.from.toString(), "to" to e.to.toString(), "bytes" to e.bytes.toString(), "target" to (e.target?.canonical ?: "—")),
+                    ),
+                )
+                is TraceEvent.Allocation -> {
+                    val now = (live[e.scope] ?: 0L) + e.bytes
+                    live[e.scope] = now
+                    emit(instant("alloc #${e.storageId}", "alloc", ts, tidOf(e.scope), mapOf("bytes" to e.bytes.toString(), "origin" to (e.origin?.canonical ?: "—"), "site" to (e.site ?: "—"))))
+                    emit(counter("live bytes", ts, mapOf(e.scope.name.lowercase() to now)))
+                }
+                is TraceEvent.Free -> {
+                    val now = ((live[e.scope] ?: 0L) - e.bytes).coerceAtLeast(0L)
+                    live[e.scope] = now
+                    emit(instant("free #${e.storageId}", "alloc", ts, tidOf(e.scope), mapOf("bytes" to e.bytes.toString())))
+                    emit(counter("live bytes", ts, mapOf(e.scope.name.lowercase() to now)))
+                }
+                is TraceEvent.ScopeReset -> {
+                    live[e.scope] = e.liveBytesAfter
+                    emit(instant("reset", "scope", ts, tidOf(e.scope), mapOf("before" to e.liveBytesBefore.toString(), "after" to e.liveBytesAfter.toString())))
+                    emit(counter("live bytes", ts, mapOf(e.scope.name.lowercase() to e.liveBytesAfter)))
+                }
+                is TraceEvent.Counter -> emit(counter(e.name, ts, mapOf(e.unit to e.value)))
+                is TraceEvent.Plan -> emit(
+                    instant(
+                        "plan", "plan", ts, MAIN_TID,
+                        mapOf(
+                            "model" to e.model, "ctx" to e.ctx.toString(),
+                            "weights" to e.weightsBytes.toString(), "kv" to e.kvBytes.toString(),
+                            "forward" to e.forwardBytes.toString(), "headroom" to e.headroomBytes.toString(),
+                            "total" to e.totalBytes.toString(), "budget" to (e.budgetBytes?.toString() ?: "—"),
+                            "fits" to (e.fits?.toString() ?: "—"),
+                        ),
+                    ),
+                )
+            }
+        }
+        out.append("\n],\"displayTimeUnit\":\"ms\"}")
+        return out.toString()
+    }
+
+    /** Export the events a [RecordingTraceSink] has kept. */
+    public fun export(sink: RecordingTraceSink, processName: String = "skainet"): String = export(sink.events(), processName)
+
+    // --- Chrome trace event objects ---
+
+    private fun slice(phase: String, name: String, cat: String, ts: Double, tid: Int, args: Map<String, String>): String =
+        """{"ph":"$phase","name":"${esc(name)}","cat":"$cat","pid":$PID,"tid":$tid,"ts":${fmt(ts)}${argsJson(args)}}"""
+
+    private fun complete(name: String, cat: String, ts: Double, dur: Double, tid: Int, args: Map<String, String>): String =
+        """{"ph":"X","name":"${esc(name)}","cat":"$cat","pid":$PID,"tid":$tid,"ts":${fmt(ts)},"dur":${fmt(dur)}${argsJson(args)}}"""
+
+    private fun instant(name: String, cat: String, ts: Double, tid: Int, args: Map<String, String>): String =
+        """{"ph":"i","name":"${esc(name)}","cat":"$cat","pid":$PID,"tid":$tid,"ts":${fmt(ts)},"s":"t"${argsJson(args)}}"""
+
+    private fun counter(name: String, ts: Double, values: Map<String, Long>): String =
+        """{"ph":"C","name":"${esc(name)}","pid":$PID,"tid":$MAIN_TID,"ts":${fmt(ts)},"args":{${values.entries.joinToString(",") { "\"${esc(it.key)}\":${it.value}" }}}}"""
+
+    private fun metadata(pid: Int, tid: Int, name: String, value: String): String =
+        """{"ph":"M","name":"$name","pid":$pid,"tid":$tid,"args":{"name":"${esc(value)}"}}"""
+
+    private fun argsJson(args: Map<String, String>): String =
+        if (args.isEmpty()) "" else ",\"args\":{${args.entries.joinToString(",") { "\"${esc(it.key)}\":\"${esc(it.value)}\"" }}}"
+
+    /** Microseconds with three decimals, without depending on a platform formatter. */
+    private fun fmt(v: Double): String {
+        val scaled = kotlin.math.round(v * 1000).toLong()
+        val whole = scaled / 1000
+        val frac = (if (scaled < 0) -scaled else scaled) % 1000
+        return "$whole.${frac.toString().padStart(3, '0')}"
+    }
+
+    private fun esc(s: String): String = buildString(s.length) {
+        for (c in s) when (c) {
+            '"' -> append("\\\""); '\\' -> append("\\\\"); '\n' -> append("\\n"); '\r' -> append("\\r"); '\t' -> append("\\t")
+            else -> if (c < ' ') append("\\u").append(c.code.toString(16).padStart(4, '0')) else append(c)
+        }
+    }
+}

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporterTest.kt

@@ -0,0 +1,77 @@
+package sk.ainet.lang.memory.trace
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/** SKEEP-003 §4.9 / PRD M1-A7: one track per scope, kernel spans labelled by TensorId, a live-bytes counter. */
+@OptIn(ExperimentalMemoryApi::class)
+class PerfettoTraceExporterTest {
+
+    private fun decodeStepEvents(): List<TraceEvent> {
+        val w = TensorId.parse("model.layers[0].attn.q_proj.weight")
+        val act = TensorId.parse("model.layers[0].attn.q#step=1")
+        return listOf(
+            TraceEvent.Plan("llama-1b", 2048, 800L shl 20, 64L shl 20, 47L shl 20, 64L shl 20, 1300L shl 20, true, timeNanos = 0),
+            TraceEvent.Allocation(1, ScopeKind.MODEL, 800L shl 20, w, site = "model.gguf", timeNanos = 1_000),
+            TraceEvent.PhaseBegin("decode", 1, mapOf("tokens" to "1"), timeNanos = 2_000),
+            TraceEvent.Allocation(2, ScopeKind.FORWARD, 8192, act, timeNanos = 2_500),
+            TraceEvent.KernelRun("matmul", "scalar-q8_0", listOf(act, w), act, 4096, 64, durationNanos = 500_000, timeNanos = 3_000_000),
+            TraceEvent.AdapterInserted("dequantize", Format(FP32, TensorEncoding.Q6_K), Format.dense(FP32), 96L shl 20, w, timeNanos = 3_100_000),
+            TraceEvent.ScopeReset(ScopeKind.FORWARD, 8192, 0, timeNanos = 3_200_000),
+            TraceEvent.PhaseEnd("decode", 1, durationNanos = 3_198_000, timeNanos = 3_200_000),
+            TraceEvent.Counter("rss", 900L shl 20, timeNanos = 3_300_000),
+        )
+    }
+
+    @Test
+    fun rendersAChromeTraceWithTracksSlicesCountersAndArgs() {
+        val json = PerfettoTraceExporter.export(decodeStepEvents(), processName = "skainet-decode")
+        assertTrue(json.startsWith("{\"traceEvents\":["), "must be a Chrome trace document")
+        assertTrue(json.trimEnd().endsWith("\"displayTimeUnit\":\"ms\"}"))
+        // process / thread naming: one track per scope
+        assertTrue(json.contains("\"name\":\"process_name\"") && json.contains("skainet-decode"))
+        assertTrue(json.contains("\"name\":\"thread_name\"") && json.contains("model scope") && json.contains("forward scope"))
+        // phases as B/E slices, kernels as complete slices with a duration
+        assertTrue(json.contains("\"ph\":\"B\",\"name\":\"decode#1\""), json.take(400))
+        assertTrue(json.contains("\"ph\":\"E\",\"name\":\"decode#1\""))
+        assertTrue(json.contains("\"ph\":\"X\",\"name\":\"matmul\""))
+        assertTrue(json.contains("\"dur\":500.000"), "kernel duration in microseconds")
+        // labelled by TensorId
+        assertTrue(json.contains("model.layers[0].attn.q_proj.weight"))
+        assertTrue(json.contains("\"kernel\":\"scalar-q8_0\""))
+        // adapters visible with their byte cost
+        assertTrue(json.contains("adapter:dequantize") && json.contains("\"bytes\":\"100663296\""))
+        // live-bytes counters per scope
+        assertTrue(json.contains("\"ph\":\"C\",\"name\":\"live bytes\""))
+        assertTrue(json.contains("\"forward\":8192") && json.contains("\"forward\":0"), "the reset must drop the counter back to zero")
+        assertTrue(json.contains("\"model\":838860800"))
+        // plan and platform counters
+        assertTrue(json.contains("\"name\":\"plan\"") && json.contains("\"fits\":\"true\""))
+        assertTrue(json.contains("\"name\":\"rss\""))
+    }
+
+    @Test
+    fun exportsWhatARecordingSinkKept() {
+        val sink = RecordingTraceSink()
+        sink.phase("prefill", 0) { sink.emit(TraceEvent.Counter("tokens", 64, "count")) }
+        val json = PerfettoTraceExporter.export(sink)
+        assertTrue(json.contains("prefill#0"))
+        assertTrue(json.contains("\"name\":\"tokens\""))
+        assertEquals(3, sink.events().size)
+    }
+
+    @Test
+    fun escapesStringsAndFormatsTimestamps() {
+        val json = PerfettoTraceExporter.export(listOf(TraceEvent.PhaseBegin("we\"ird\n", null, mapOf("k" to "a\\b"), timeNanos = 1_500)))
+        assertTrue(json.contains("we\\\"ird\\n"), json)
+        assertTrue(json.contains("\"a\\\\b\""))
+        assertTrue(json.contains("\"ts\":1.500"), "nanoseconds render as microseconds")
+    }
+}