From d1bae070e90acebcebbf91bd074248e84152bd4b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 11:18:28 +0200 Subject: [PATCH 1/2] =?UTF-8?q?feat(memory):=20TraceSink=20API=20=E2=80=94?= =?UTF-8?q?=20one=20event=20stream=20for=20phases,=20kernels,=20adapters,?= =?UTF-8?q?=20allocations,=20counters=20and=20plans=20(SKEEP-003=20P2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F7 (API half; exporters follow in #1025). - sk.ainet.lang.memory.trace.TraceEvent (sealed): PhaseBegin/PhaseEnd (load · compile · prefill · decode(step) · module spans), KernelRun (op, kernel key, input/output TensorIds, bytes read/written, duration), AdapterInserted (kind, from/to Format, bytes, target, scope), Allocation / Free (storageId, ScopeKind, bytes, origin TensorId, debug site), ScopeReset (live bytes before/after), Counter (platform samples), Plan (the M0 MemoryPlan, for the plan-vs-actual check #1030). TraceClock: monotonic nanoseconds. - TraceSink (isEnabled + emit), NoopTraceSink (the default: one check, no allocation), RecordingTraceSink (ring buffer with emitted/dropped counters; the debugger, tests and #1030 read it), CompositeTraceSink; inline helpers TraceSink.phase { } and TraceSink.kernel { } that short-circuit when disabled and time the block; MemoryPlan.emit(sink). - ExecutionContext.traceSink: default member returning NoopTraceSink — opt-in per context, no change for existing implementations. - TraceSinkTest (ring buffer, ordering, phase/kernel helpers incl. exception path, identity fields on events, composite fan-out, plan event, clock). BCV dumps regenerated (lang-core, backend-cpu, compile-dag gain the inherited getTraceSink()). Closes #1017 Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-backend-cpu.api | 1 + .../api/jvm/skainet-compile-dag.api | 2 + .../api/jvm/skainet-lang-core.api | 258 ++++++++++++++++++ .../sk/ainet/context/ExecutionContext.kt | 8 + .../sk/ainet/lang/memory/trace/PlanTrace.kt | 17 ++ .../sk/ainet/lang/memory/trace/TraceEvent.kt | 90 ++++++ .../sk/ainet/lang/memory/trace/TraceSink.kt | 99 +++++++ .../ainet/lang/memory/trace/TraceSinkTest.kt | 122 +++++++++ 8 files changed, 597 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PlanTrace.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceSink.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/TraceSinkTest.kt diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index a912cfe79..facf4eff7 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -24,6 +24,7 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun getPhase ()Lsk/ainet/context/Phase; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index f98145a5a..1dbc6e43f 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -182,6 +182,7 @@ public final class sk/ainet/lang/graph/DefaultGraphExecutionContext : sk/ainet/l public final fun getSession ()Lsk/ainet/lang/trace/TraceSession; public fun getTapeStack ()Lsk/ainet/tape/TapeStack; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -450,6 +451,7 @@ public final class sk/ainet/lang/graph/exec/GraphExecutionContext$DefaultImpls { public static fun getMemoryPlanner (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; public static fun getMemoryTracker (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; public static fun getScratch (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; + public static fun getTraceSink (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Z public static fun ones (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun placeholder (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 5d22ae6a7..f84ac2224 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -211,6 +211,7 @@ public final class sk/ainet/context/DefaultDataExecutionContext : sk/ainet/conte public fun getPhase ()Lsk/ainet/context/Phase; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -239,6 +240,7 @@ public abstract interface class sk/ainet/context/ExecutionContext { public abstract fun getPhase ()Lsk/ainet/context/Phase; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public abstract fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -261,6 +263,7 @@ public final class sk/ainet/context/ExecutionContext$DefaultImpls { public static fun getMemoryPlanner (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; public static fun getMemoryTracker (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; public static fun getScratch (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; + public static fun getTraceSink (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/context/ExecutionContext;)Z public static fun ones (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun placeholder (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -365,6 +368,7 @@ public final class sk/ainet/context/PhaseOverridingExecutionContext : sk/ainet/c public fun getPhase ()Lsk/ainet/context/Phase; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -409,6 +413,7 @@ public final class sk/ainet/context/TrainingExecutionContext$DefaultImpls { public static fun getMemoryPlanner (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; public static fun getMemoryTracker (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; public static fun getScratch (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; + public static fun getTraceSink (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/context/TrainingExecutionContext;)Z public static fun ones (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun placeholder (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -888,6 +893,257 @@ public final class sk/ainet/lang/memory/plan/Suggestion { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/trace/CompositeTraceSink : sk/ainet/lang/memory/trace/TraceSink { + public fun (Ljava/util/List;)V + public fun ([Lsk/ainet/lang/memory/trace/TraceSink;)V + 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/PlanTraceKt { + public static final fun emit (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/trace/TraceSink;)V +} + +public final class sk/ainet/lang/memory/trace/RecordingTraceSink : sk/ainet/lang/memory/trace/TraceSink { + public static final field Companion Lsk/ainet/lang/memory/trace/RecordingTraceSink$Companion; + public static final field DEFAULT_CAPACITY I + public fun ()V + public fun (I)V + public synthetic fun (IILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun clear ()V + public fun emit (Lsk/ainet/lang/memory/trace/TraceEvent;)V + public final fun events ()Ljava/util/List; + public final fun getCapacity ()I + public final fun getDropped ()J + public final fun getEmitted ()J + public fun isEnabled ()Z +} + +public final class sk/ainet/lang/memory/trace/RecordingTraceSink$Companion { +} + +public final class sk/ainet/lang/memory/trace/TraceClock { + public static final field INSTANCE Lsk/ainet/lang/memory/trace/TraceClock; + public final fun nowNanos ()J +} + +public abstract interface class sk/ainet/lang/memory/trace/TraceEvent { + public abstract fun getTimeNanos ()J +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$AdapterInserted : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;J)V + public synthetic fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Lsk/ainet/lang/memory/Format; + public final fun component3 ()Lsk/ainet/lang/memory/Format; + public final fun component4 ()J + public final fun component5 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component6 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component7 ()J + public final fun copy (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;J)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted;Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytes ()J + public final fun getFrom ()Lsk/ainet/lang/memory/Format; + public final fun getKind ()Ljava/lang/String; + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public final fun getTarget ()Lsk/ainet/lang/tensor/TensorId; + public fun getTimeNanos ()J + public final fun getTo ()Lsk/ainet/lang/memory/Format; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$Allocation : sk/ainet/lang/memory/trace/TraceEvent { + public fun (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;J)V + public synthetic fun (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()J + public final fun component2 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component3 ()J + public final fun component4 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component5 ()Ljava/lang/String; + public final fun component6 ()J + public final fun copy (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;J)Lsk/ainet/lang/memory/trace/TraceEvent$Allocation; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$Allocation;JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$Allocation; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytes ()J + public final fun getOrigin ()Lsk/ainet/lang/tensor/TensorId; + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public final fun getSite ()Ljava/lang/String; + public final fun getStorageId ()J + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$Counter : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;JLjava/lang/String;J)V + public synthetic fun (Ljava/lang/String;JLjava/lang/String;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()J + public final fun component3 ()Ljava/lang/String; + public final fun component4 ()J + public final fun copy (Ljava/lang/String;JLjava/lang/String;J)Lsk/ainet/lang/memory/trace/TraceEvent$Counter; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$Counter;Ljava/lang/String;JLjava/lang/String;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$Counter; + public fun equals (Ljava/lang/Object;)Z + public final fun getName ()Ljava/lang/String; + public fun getTimeNanos ()J + public final fun getUnit ()Ljava/lang/String; + public final fun getValue ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$Free : sk/ainet/lang/memory/trace/TraceEvent { + public fun (JLsk/ainet/lang/memory/ScopeKind;JJ)V + public synthetic fun (JLsk/ainet/lang/memory/ScopeKind;JJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()J + public final fun component2 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component3 ()J + public final fun component4 ()J + public final fun copy (JLsk/ainet/lang/memory/ScopeKind;JJ)Lsk/ainet/lang/memory/trace/TraceEvent$Free; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$Free;JLsk/ainet/lang/memory/ScopeKind;JJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$Free; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytes ()J + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public final fun getStorageId ()J + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$KernelRun : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJJJ)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJJJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()Ljava/util/List; + public final fun component4 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component5 ()J + public final fun component6 ()J + public final fun component7 ()J + public final fun component8 ()J + public final fun copy (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJJJ)Lsk/ainet/lang/memory/trace/TraceEvent$KernelRun; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$KernelRun;Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJJJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$KernelRun; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytesRead ()J + public final fun getBytesWritten ()J + public final fun getDurationNanos ()J + public final fun getInputs ()Ljava/util/List; + public final fun getKernel ()Ljava/lang/String; + public final fun getOp ()Ljava/lang/String; + public final fun getOutput ()Lsk/ainet/lang/tensor/TensorId; + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$PhaseBegin : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;J)V + public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/Integer; + public final fun component3 ()Ljava/util/Map; + public final fun component4 ()J + public final fun copy (Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;J)Lsk/ainet/lang/memory/trace/TraceEvent$PhaseBegin; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$PhaseBegin;Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$PhaseBegin; + public fun equals (Ljava/lang/Object;)Z + public final fun getAttributes ()Ljava/util/Map; + public final fun getPhase ()Ljava/lang/String; + public final fun getStep ()Ljava/lang/Integer; + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$PhaseEnd : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;Ljava/lang/Integer;JJ)V + public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;JJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/Integer; + public final fun component3 ()J + public final fun component4 ()J + public final fun copy (Ljava/lang/String;Ljava/lang/Integer;JJ)Lsk/ainet/lang/memory/trace/TraceEvent$PhaseEnd; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$PhaseEnd;Ljava/lang/String;Ljava/lang/Integer;JJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$PhaseEnd; + public fun equals (Ljava/lang/Object;)Z + public final fun getDurationNanos ()J + public final fun getPhase ()Ljava/lang/String; + public final fun getStep ()Ljava/lang/Integer; + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$Plan : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Ljava/lang/String;IJJJJLjava/lang/Long;Ljava/lang/Boolean;J)V + public synthetic fun (Ljava/lang/String;IJJJJLjava/lang/Long;Ljava/lang/Boolean;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()I + public final fun component3 ()J + public final fun component4 ()J + public final fun component5 ()J + public final fun component6 ()J + public final fun component7 ()Ljava/lang/Long; + public final fun component8 ()Ljava/lang/Boolean; + public final fun component9 ()J + public final fun copy (Ljava/lang/String;IJJJJLjava/lang/Long;Ljava/lang/Boolean;J)Lsk/ainet/lang/memory/trace/TraceEvent$Plan; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$Plan;Ljava/lang/String;IJJJJLjava/lang/Long;Ljava/lang/Boolean;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$Plan; + public fun equals (Ljava/lang/Object;)Z + public final fun getBudgetBytes ()Ljava/lang/Long; + public final fun getCtx ()I + public final fun getFits ()Ljava/lang/Boolean; + public final fun getForwardBytes ()J + public final fun getHeadroomBytes ()J + public final fun getKvBytes ()J + public final fun getModel ()Ljava/lang/String; + public fun getTimeNanos ()J + public final fun getTotalBytes ()J + public final fun getWeightsBytes ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/TraceEvent$ScopeReset : sk/ainet/lang/memory/trace/TraceEvent { + public fun (Lsk/ainet/lang/memory/ScopeKind;JJJ)V + public synthetic fun (Lsk/ainet/lang/memory/ScopeKind;JJJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component2 ()J + public final fun component3 ()J + public final fun component4 ()J + public final fun copy (Lsk/ainet/lang/memory/ScopeKind;JJJ)Lsk/ainet/lang/memory/trace/TraceEvent$ScopeReset; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$ScopeReset;Lsk/ainet/lang/memory/ScopeKind;JJJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$ScopeReset; + public fun equals (Ljava/lang/Object;)Z + public final fun getLiveBytesAfter ()J + public final fun getLiveBytesBefore ()J + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public fun getTimeNanos ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public abstract interface class sk/ainet/lang/memory/trace/TraceSink { + public abstract fun emit (Lsk/ainet/lang/memory/trace/TraceEvent;)V + public fun isEnabled ()Z +} + +public final class sk/ainet/lang/memory/trace/TraceSink$DefaultImpls { + public static fun isEnabled (Lsk/ainet/lang/memory/trace/TraceSink;)Z +} + +public final class sk/ainet/lang/memory/trace/TraceSinkKt { + public static final fun kernel (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJLkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static synthetic fun kernel$default (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/tensor/TensorId;JJLkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object; + public static final fun phase (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static synthetic fun phase$default (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/Integer;Ljava/util/Map;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object; +} + public final class sk/ainet/lang/nn/AvgPool2d : sk/ainet/lang/nn/Module { public fun (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;)V public synthetic fun (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -1018,6 +1274,7 @@ public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ai public fun getPhase ()Lsk/ainet/context/Phase; public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; public fun isRecording ()Z public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; @@ -1349,6 +1606,7 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContext$DefaultImpls { public static fun getMemoryPlanner (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; public static fun getMemoryTracker (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryTracker; public static fun getScratch (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/scratch/ScratchPool; + public static fun getTraceSink (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/memory/trace/TraceSink; public static fun isRecording (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Z public static fun ones (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun placeholder (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt index 042cf722c..39d004f0b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt @@ -14,6 +14,14 @@ import sk.ainet.lang.types.DType import kotlin.reflect.KClass public interface ExecutionContext { + /** + * Where this context's trace events go (SKEEP-003 §4.9): phases, kernel runs, adapter + * insertions, allocations. Default [sk.ainet.lang.memory.trace.NoopTraceSink] — nothing is + * recorded until a context opts in with a recording or exporting sink. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public val traceSink: sk.ainet.lang.memory.trace.TraceSink get() = sk.ainet.lang.memory.trace.NoopTraceSink + public val ops: TensorOps // Optional forward hooks for recording or diagnostics (null → disabled) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PlanTrace.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PlanTrace.kt new file mode 100644 index 000000000..78544c1a0 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PlanTrace.kt @@ -0,0 +1,17 @@ +package sk.ainet.lang.memory.trace + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.MemoryPlan + +/** Emit this plan as a [TraceEvent.Plan] so the plan-vs-actual check (#1030) can find it in the stream. */ +@ExperimentalMemoryApi +public fun MemoryPlan.emit(sink: TraceSink) { + if (!sink.isEnabled) return + sink.emit( + TraceEvent.Plan( + model = input.modelName, ctx = input.ctx, + weightsBytes = weightsBytes, kvBytes = kvBytes, forwardBytes = forwardBytes, headroomBytes = headroomBytes, + budgetBytes = budget?.bytes, fits = fits, + ), + ) +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt new file mode 100644 index 000000000..e7ecf0747 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt @@ -0,0 +1,90 @@ +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 + +/** + * One event of the SKaiNET observability stream (SKEEP-003 §4.9): phases, kernel runs, adapter + * insertions, allocations and platform counters share one event model, keyed by [TensorId], + * storage id and [ScopeKind]. Exporters (Perfetto / JFR / `android.os.Trace`, M1 slice #1025), + * the memory debugger and the benchmark report are consumers of the same stream. + * + * Events are small value objects; [timeNanos] is a monotonic timestamp in nanoseconds + * ([TraceClock.nowNanos]), comparable only within one process. + */ +@ExperimentalMemoryApi +public sealed interface TraceEvent { + public val timeNanos: Long + + /** A phase opened: `load`, `compile`, `prefill`, `decode` (with [step]), `sample`, or a module span like `layers[3].attn`. */ + public data class PhaseBegin(val phase: String, val step: Int? = null, val attributes: Map = emptyMap(), override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + + /** The matching phase closed; [durationNanos] is filled by [TraceSink.phase]. */ + public data class PhaseEnd(val phase: String, val step: Int? = null, val durationNanos: Long = 0L, override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + + /** A kernel ran: which op, which registered kernel (the `KernelKey` string once M1 has it), on which tensors, how many bytes it touched. */ + public data class KernelRun( + val op: String, + val kernel: String, + val inputs: List = emptyList(), + val output: TensorId? = null, + val bytesRead: Long = 0L, + val bytesWritten: Long = 0L, + val durationNanos: Long = 0L, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent + + /** The dispatcher inserted a conversion (dequantize, requantize, gather) — always visible (§5.1). */ + public data class AdapterInserted( + val kind: String, + val from: Format, + val to: Format, + val bytes: Long, + val target: TensorId? = null, + val scope: ScopeKind = ScopeKind.FORWARD, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent + + /** A storage was allocated. [site] is the allocation site in debug mode, [origin] the TensorId it backs. */ + public data class Allocation( + val storageId: Long, + val scope: ScopeKind, + val bytes: Long, + val origin: TensorId? = null, + val site: String? = null, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent + + /** A storage was freed / closed. */ + public data class Free(val storageId: Long, val scope: ScopeKind, val bytes: Long, override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + + /** A `Forward` (or other) scope was reset: how many bytes were live before and after. */ + public data class ScopeReset(val scope: ScopeKind, val liveBytesBefore: Long, val liveBytesAfter: Long, override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + + /** A platform counter sample (RSS, page faults, heap, direct memory …). */ + public data class Counter(val name: String, val value: Long, val unit: String = "bytes", override val timeNanos: Long = TraceClock.nowNanos()) : TraceEvent + + /** A memory plan was computed (M0 `MemoryPlan`): the plan-vs-actual check (#1030) compares this with the allocation events. */ + public data class Plan( + val model: String, + val ctx: Int, + val weightsBytes: Long, + val kvBytes: Long, + val forwardBytes: Long, + val headroomBytes: Long, + val budgetBytes: Long? = null, + val fits: Boolean? = null, + override val timeNanos: Long = TraceClock.nowNanos(), + ) : TraceEvent { + val totalBytes: Long get() = weightsBytes + kvBytes + forwardBytes + headroomBytes + } +} + +/** Monotonic clock for trace timestamps (`kotlin.time.TimeSource.Monotonic`). */ +@ExperimentalMemoryApi +public object TraceClock { + private val start = kotlin.time.TimeSource.Monotonic.markNow() + public fun nowNanos(): Long = start.elapsedNow().inWholeNanoseconds +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceSink.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceSink.kt new file mode 100644 index 000000000..56736e1d5 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceSink.kt @@ -0,0 +1,99 @@ +package sk.ainet.lang.memory.trace + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * Where [TraceEvent]s go. Disabled by default ([NoopTraceSink]) — emitting costs one `isEnabled` + * check — and opt-in per `ExecutionContext` (`traceSink`). Implementations: [RecordingTraceSink] + * (ring buffer, tests and the debugger), [CompositeTraceSink], and the exporters of M1 slice #1025 + * (Perfetto JSON, JFR, `android.os.Trace`). + */ +@ExperimentalMemoryApi +public interface TraceSink { + /** `false` for a sink that drops everything; producers check it before building an event. */ + public val isEnabled: Boolean get() = true + + public fun emit(event: TraceEvent) +} + +/** The default: nothing is recorded, nothing is allocated. */ +@ExperimentalMemoryApi +public object NoopTraceSink : TraceSink { + override val isEnabled: Boolean get() = false + override fun emit(event: TraceEvent) {} +} + +/** + * Keeps the last [capacity] events in a ring buffer. The debugger, tests and the plan-vs-actual + * check read [events]; [clear] starts over. Not thread-safe by design (one sink per context). + */ +@ExperimentalMemoryApi +public class RecordingTraceSink(public val capacity: Int = DEFAULT_CAPACITY) : TraceSink { + init { require(capacity > 0) { "capacity must be > 0" } } + + private val buffer = ArrayDeque(minOf(capacity, 1024)) + /** Total number of events ever emitted (including ones that fell out of the ring). */ + public var emitted: Long = 0L + private set + /** Events dropped from the head because the ring was full. */ + public var dropped: Long = 0L + private set + + override fun emit(event: TraceEvent) { + if (buffer.size == capacity) { buffer.removeFirst(); dropped++ } + buffer.addLast(event); emitted++ + } + + /** The retained events, oldest first. */ + public fun events(): List = buffer.toList() + + public inline fun eventsOf(): List = events().filterIsInstance() + + public fun clear() { buffer.clear(); dropped = 0L; emitted = 0L } + + public companion object { public const val DEFAULT_CAPACITY: Int = 65_536 } +} + +/** Fan-out to several sinks; enabled if any of them is. */ +@ExperimentalMemoryApi +public class CompositeTraceSink(private val sinks: List) : TraceSink { + public constructor(vararg sinks: TraceSink) : this(sinks.toList()) + override val isEnabled: Boolean get() = sinks.any { it.isEnabled } + override fun emit(event: TraceEvent) { for (s in sinks) if (s.isEnabled) s.emit(event) } +} + +/** + * Run [block] inside a phase span: emits [TraceEvent.PhaseBegin], then [TraceEvent.PhaseEnd] with + * the measured duration (also on exception). No events and no allocation when the sink is disabled. + */ +@ExperimentalMemoryApi +public inline fun TraceSink.phase(name: String, step: Int? = null, attributes: Map = emptyMap(), block: () -> T): T { + if (!isEnabled) return block() + val t0 = TraceClock.nowNanos() + emit(TraceEvent.PhaseBegin(name, step, attributes, t0)) + try { + return block() + } finally { + val t1 = TraceClock.nowNanos() + emit(TraceEvent.PhaseEnd(name, step, t1 - t0, t1)) + } +} + +/** Time [block] as a kernel run of [op] on [kernel]; the returned value is the kernel's result. */ +@ExperimentalMemoryApi +public inline fun TraceSink.kernel( + op: String, + kernel: String, + inputs: List = emptyList(), + output: sk.ainet.lang.tensor.TensorId? = null, + bytesRead: Long = 0L, + bytesWritten: Long = 0L, + block: () -> T, +): T { + if (!isEnabled) return block() + val t0 = TraceClock.nowNanos() + val r = block() + val t1 = TraceClock.nowNanos() + emit(TraceEvent.KernelRun(op, kernel, inputs, output, bytesRead, bytesWritten, t1 - t0, t1)) + return r +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/TraceSinkTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/TraceSinkTest.kt new file mode 100644 index 000000000..2695e8b35 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/TraceSinkTest.kt @@ -0,0 +1,122 @@ +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.memory.plan.Budget +import sk.ainet.lang.memory.plan.MemoryPlans +import sk.ainet.lang.memory.plan.PlanInput +import sk.ainet.lang.memory.plan.PlanTensor +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int8 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** SKEEP-003 §4.9 / M1-F7 (API half): one event stream, disabled by default, recorded into a ring buffer. */ +@OptIn(ExperimentalMemoryApi::class) +class TraceSinkTest { + + @Test + fun noopSinkIsDisabledAndDropsEverything() { + assertFalse(NoopTraceSink.isEnabled) + NoopTraceSink.emit(TraceEvent.Counter("rss", 1)) + // helpers short-circuit: the block runs, nothing else happens + var ran = false + val r = NoopTraceSink.phase("prefill", 1) { ran = true; 42 } + assertEquals(42, r); assertTrue(ran) + assertEquals("k", NoopTraceSink.kernel("matmul", "scalar") { "k" }) + } + + @Test + fun recordingSinkKeepsOrderAndRingCapacity() { + val sink = RecordingTraceSink(capacity = 3) + assertTrue(sink.isEnabled) + for (i in 1..5) sink.emit(TraceEvent.Counter("c", i.toLong(), timeNanos = i.toLong())) + val values = sink.eventsOf().map { it.value } + assertEquals(listOf(3L, 4L, 5L), values) + assertEquals(5L, sink.emitted); assertEquals(2L, sink.dropped) + sink.clear() + assertTrue(sink.events().isEmpty()); assertEquals(0L, sink.emitted); assertEquals(0L, sink.dropped) + assertFailsWith { RecordingTraceSink(0) } + } + + @Test + fun phaseHelperEmitsBeginAndEndWithDurationAlsoOnException() { + val sink = RecordingTraceSink() + val r = sink.phase("decode", step = 17, attributes = mapOf("tokens" to "1")) { "ok" } + assertEquals("ok", r) + val ev = sink.events() + assertEquals(2, ev.size) + val begin = assertIs(ev[0]); val end = assertIs(ev[1]) + assertEquals("decode", begin.phase); assertEquals(17, begin.step); assertEquals("1", begin.attributes["tokens"]) + assertEquals("decode", end.phase); assertEquals(17, end.step) + assertTrue(end.durationNanos >= 0); assertTrue(end.timeNanos >= begin.timeNanos) + + sink.clear() + assertFailsWith { sink.phase("load") { throw IllegalStateException("boom") } } + assertEquals(listOf("PhaseBegin", "PhaseEnd"), sink.events().map { it::class.simpleName }) + } + + @Test + fun kernelHelperRecordsOpKernelTensorsAndBytes() { + val sink = RecordingTraceSink() + val w = TensorId.parse("model.layers[3].attn.q_proj.weight") + val out = sink.kernel("matmul", "scalar-q4k", inputs = listOf(null, w), output = TensorId.parse("model.layers[3].attn.q#step=1"), bytesRead = 4096, bytesWritten = 64) { 7 } + assertEquals(7, out) + val k = assertIs(sink.events().single()) + assertEquals("matmul", k.op); assertEquals("scalar-q4k", k.kernel); assertEquals(listOf(null, w), k.inputs) + assertEquals(4096L, k.bytesRead); assertEquals(64L, k.bytesWritten); assertTrue(k.durationNanos >= 0) + } + + @Test + fun allocationAdapterAndScopeEventsCarryIdentityFields() { + val sink = RecordingTraceSink() + val id = TensorId.parse("model.layers[0].mlp.down_proj.weight") + sink.emit(TraceEvent.Allocation(storageId = 412, scope = ScopeKind.MODEL, bytes = 96L shl 20, origin = id, site = "GgufLoader.kt:120")) + sink.emit(TraceEvent.AdapterInserted("dequantize", Format(FP32, TensorEncoding.Q6_K), Format.dense(FP32), 96L shl 20, target = id)) + sink.emit(TraceEvent.ScopeReset(ScopeKind.FORWARD, liveBytesBefore = 8L shl 20, liveBytesAfter = 0)) + sink.emit(TraceEvent.Free(412, ScopeKind.MODEL, 96L shl 20)) + val ev = sink.events() + val a = assertIs(ev[0]); assertEquals(412L, a.storageId); assertEquals(ScopeKind.MODEL, a.scope); assertEquals(id, a.origin); assertEquals("GgufLoader.kt:120", a.site) + val ad = assertIs(ev[1]); assertEquals("dequantize", ad.kind); assertEquals(TensorEncoding.Q6_K, ad.from.encoding); assertTrue(ad.to.isDense); assertEquals(ScopeKind.FORWARD, ad.scope) + val rs = assertIs(ev[2]); assertEquals(8L shl 20, rs.liveBytesBefore); assertEquals(0L, rs.liveBytesAfter) + assertIs(ev[3]) + } + + @Test + fun compositeFansOutToEnabledSinksOnly() { + val a = RecordingTraceSink(); val b = RecordingTraceSink() + val c = CompositeTraceSink(a, NoopTraceSink, b) + assertTrue(c.isEnabled) + c.emit(TraceEvent.Counter("x", 1)) + assertEquals(1, a.events().size); assertEquals(1, b.events().size) + assertFalse(CompositeTraceSink(NoopTraceSink).isEnabled) + } + + @Test + fun memoryPlanEmitsAPlanEvent() { + val f = Format(FP32, TensorEncoding.Q4_K) + val input = PlanInput("m", "llama", listOf(PlanTensor("w", TensorId.parse("model.w"), f, 256, 144)), null, 512) + val plan = MemoryPlans.plan(input, Budget.of(1L shl 30)) + val sink = RecordingTraceSink() + plan.emit(sink) + val p = assertIs(sink.events().single()) + assertEquals("m", p.model); assertEquals(512, p.ctx); assertEquals(144L, p.weightsBytes); assertEquals(plan.totalBytes, p.totalBytes) + assertEquals(1L shl 30, p.budgetBytes); assertEquals(true, p.fits) + plan.emit(NoopTraceSink) // no-op + // an Int8 dense format has a distinct string in events + assertEquals("Int8/Dense(1B)", Format.dense(Int8).toString()) + } + + @Test + fun clockIsMonotonic() { + val t0 = TraceClock.nowNanos(); val t1 = TraceClock.nowNanos() + assertTrue(t1 >= t0) + } +} From 5b9b0e4600670935bc3906f40a7578b09160e8ea Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 11:23:46 +0200 Subject: [PATCH 2/2] =?UTF-8?q?feat(memory):=20Storage=20=E2=80=94=20the?= =?UTF-8?q?=20one=20byte=20owner:=20sealed=20kinds,=20Owner,=20StorageId,?= =?UTF-8?q?=20StorageClosedException,=20Heap=20on=20every=20target=20(SKEE?= =?UTF-8?q?P-003=20P2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F1 (core). SKEEP-003 §4.2–4.4, rules 1–2: exactly one byte owner; ownership is a constructor argument and then enforced. - sk.ainet.lang.memory.Storage (sealed, AutoCloseable): id, sizeBytes, owner, domain, scope (from the owner), debugOrigin TensorId, isAlive, isMutable, checkAlive() → StorageClosedException(id, origin), close() (idempotent; owned = freed once, borrowed = forgotten, alias = detached; emits TraceEvent.Free), slice() → Owner.Alias view that keeps the parent alive, delegates mutability and dies with it. Kinds: Heap (final, common), OffHeap / Mapped / Device (abstract, bound per platform in #1019/#1020; Device is a placeholder). - Storage.Heap: one Kotlin array (floats | ints | bytes) + arrayOffset + sizeBytes; floats()/ints()/bytes() allocate owned storage for a scope kind and emit TraceEvent.Allocation; wrap(array, …) borrows without copying (the #782 copyOf replacement) and is never freed by us. Kernels unwrap once per call (floats/ints/bytes + arrayOffset), per the Phase-2 spike (#1016). - Owner (sealed): Owned(scope) | Borrowed(external) | Alias(parent). StorageId: monotonic per-process value class (kotlin.concurrent.atomics). MemoryDomain gains HOST_OFFHEAP (additive). - StorageTest: owned lifecycle + trace events, ids, borrowed release without touching the lender's bytes, alias lifetime/mutability/ nesting, slice bounds/alignment, rendering. BCV: lang-core jvm dump regenerated (additions only). Closes #1018 Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 126 ++++++++++++ .../kotlin/sk/ainet/lang/memory/Storage.kt | 188 ++++++++++++++++++ .../sk/ainet/lang/tensor/storage/Placement.kt | 2 + .../sk/ainet/lang/memory/StorageTest.kt | 104 ++++++++++ 4 files changed, 420 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/StorageTest.kt diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index f84ac2224..eb71108de 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -705,6 +705,39 @@ public final class sk/ainet/lang/memory/FormatKt { public static final fun getFormatOrNull (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/memory/Format; } +public abstract interface class sk/ainet/lang/memory/Owner { +} + +public final class sk/ainet/lang/memory/Owner$Alias : sk/ainet/lang/memory/Owner { + public fun (Lsk/ainet/lang/memory/Storage;)V + public final fun getParent ()Lsk/ainet/lang/memory/Storage; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/Owner$Borrowed : sk/ainet/lang/memory/Owner { + public fun ()V + public fun (Ljava/lang/Object;)V + public synthetic fun (Ljava/lang/Object;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/Object; + public final fun copy (Ljava/lang/Object;)Lsk/ainet/lang/memory/Owner$Borrowed; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/Owner$Borrowed;Ljava/lang/Object;ILjava/lang/Object;)Lsk/ainet/lang/memory/Owner$Borrowed; + public fun equals (Ljava/lang/Object;)Z + public final fun getExternal ()Ljava/lang/Object; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/Owner$Owned : sk/ainet/lang/memory/Owner { + public fun (Lsk/ainet/lang/memory/ScopeKind;)V + public final fun component1 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun copy (Lsk/ainet/lang/memory/ScopeKind;)Lsk/ainet/lang/memory/Owner$Owned; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/Owner$Owned;Lsk/ainet/lang/memory/ScopeKind;ILjava/lang/Object;)Lsk/ainet/lang/memory/Owner$Owned; + public fun equals (Ljava/lang/Object;)Z + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/memory/ScopeKind : java/lang/Enum { public static final field AMBIENT Lsk/ainet/lang/memory/ScopeKind; public static final field FORWARD Lsk/ainet/lang/memory/ScopeKind; @@ -714,6 +747,98 @@ public final class sk/ainet/lang/memory/ScopeKind : java/lang/Enum { public static fun values ()[Lsk/ainet/lang/memory/ScopeKind; } +public abstract class sk/ainet/lang/memory/Storage : java/lang/AutoCloseable { + public final fun checkAlive ()V + public final fun close ()V + public abstract fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public abstract fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; + public abstract fun getId-TPZW6QE ()J + public abstract fun getOwner ()Lsk/ainet/lang/memory/Owner; + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + protected abstract fun getSink ()Lsk/ainet/lang/memory/trace/TraceSink; + public abstract fun getSizeBytes ()J + public final fun isAlive ()Z + public abstract fun isMutable ()Z + protected fun onClose ()V + public abstract fun slice (JJ)Lsk/ainet/lang/memory/Storage; + public fun toString ()Ljava/lang/String; +} + +public abstract class sk/ainet/lang/memory/Storage$Device : sk/ainet/lang/memory/Storage { + public fun ()V + public fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; +} + +public final class sk/ainet/lang/memory/Storage$Heap : sk/ainet/lang/memory/Storage { + public static final field Companion Lsk/ainet/lang/memory/Storage$Heap$Companion; + public synthetic fun (J[F[I[BIJLsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun getArrayOffset ()I + public final fun getBytes ()[B + public fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; + public final fun getElementBytes ()I + public final fun getElementCount ()I + public final fun getFloats ()[F + public fun getId-TPZW6QE ()J + public final fun getInts ()[I + public fun getOwner ()Lsk/ainet/lang/memory/Owner; + public fun getSizeBytes ()J + public fun isMutable ()Z + public fun slice (JJ)Lsk/ainet/lang/memory/Storage$Heap; + public synthetic fun slice (JJ)Lsk/ainet/lang/memory/Storage; +} + +public final class sk/ainet/lang/memory/Storage$Heap$Companion { + public final fun bytes (ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun bytes$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public final fun floats (ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun floats$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public final fun ints (ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun ints$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public final fun wrap ([BIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public final fun wrap ([FIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public final fun wrap ([IIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun wrap$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;[BIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun wrap$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;[FIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun wrap$default (Lsk/ainet/lang/memory/Storage$Heap$Companion;[IIIZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; +} + +public abstract class sk/ainet/lang/memory/Storage$Mapped : sk/ainet/lang/memory/Storage { + public fun ()V + public fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; +} + +public abstract class sk/ainet/lang/memory/Storage$OffHeap : sk/ainet/lang/memory/Storage { + public fun ()V + public fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; +} + +public final class sk/ainet/lang/memory/StorageClosedException : java/lang/IllegalStateException { + public synthetic fun (JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public synthetic fun (JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Lkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun getOrigin ()Lsk/ainet/lang/tensor/TensorId; + public final fun getStorageId-TPZW6QE ()J +} + +public final class sk/ainet/lang/memory/StorageId { + public static final field Companion Lsk/ainet/lang/memory/StorageId$Companion; + public static final synthetic fun box-impl (J)Lsk/ainet/lang/memory/StorageId; + public static fun constructor-impl (J)J + public fun equals (Ljava/lang/Object;)Z + public static fun equals-impl (JLjava/lang/Object;)Z + public static final fun equals-impl0 (JJ)Z + public final fun getValue ()J + public fun hashCode ()I + public static fun hashCode-impl (J)I + public fun toString ()Ljava/lang/String; + public static fun toString-impl (J)Ljava/lang/String; + public final synthetic fun unbox-impl ()J +} + +public final class sk/ainet/lang/memory/StorageId$Companion { + public final fun next-TPZW6QE ()J +} + public final class sk/ainet/lang/memory/plan/Budget { public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion; public static final field RESERVE_ANDROID_JVM J @@ -5910,6 +6035,7 @@ public final class sk/ainet/lang/tensor/storage/LogicalDType$Companion { public final class sk/ainet/lang/tensor/storage/MemoryDomain : java/lang/Enum { public static final field DEVICE_LOCAL Lsk/ainet/lang/tensor/storage/MemoryDomain; public static final field HOST_HEAP Lsk/ainet/lang/tensor/storage/MemoryDomain; + public static final field HOST_OFFHEAP Lsk/ainet/lang/tensor/storage/MemoryDomain; public static final field HOST_PINNED Lsk/ainet/lang/tensor/storage/MemoryDomain; public static final field MMAP_FILE Lsk/ainet/lang/tensor/storage/MemoryDomain; public static final field UNIFIED Lsk/ainet/lang/tensor/storage/MemoryDomain; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt new file mode 100644 index 000000000..0173d79b6 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt @@ -0,0 +1,188 @@ +@file:OptIn(kotlin.concurrent.atomics.ExperimentalAtomicApi::class) + +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.MemoryDomain +import kotlin.concurrent.atomics.AtomicLong +import kotlin.concurrent.atomics.fetchAndIncrement +import kotlin.jvm.JvmInline + +/** + * Monotonic per-process identity of one allocation (SKEEP-003 §0 *StorageId*): what the memory + * debugger keys on. One `TensorId` maps to many storage ids over time (a `Forward` scope is + * recycled every step); one storage may back many `TensorId`s (views, KV ring). + */ +@ExperimentalMemoryApi +@JvmInline +public value class StorageId(public val value: Long) { + override fun toString(): String = "#$value" + + public companion object { + private val counter = AtomicLong(1L) + /** The next id; thread-safe. */ + public fun next(): StorageId = StorageId(counter.fetchAndIncrement()) + } +} + +/** + * How a [Storage] came to hold its bytes (SKEEP-003 §0 *Owner*, §4.4). Ownership is a + * constructor argument, then it is enforced: [Borrowed] storage cannot be freed through us, + * [Alias] keeps its parent alive and cannot free or resize, [Owned] storage is freed exactly once, + * by its scope. + */ +@ExperimentalMemoryApi +public sealed interface Owner { + /** We allocated the bytes; the scope of [scope] kind frees them. */ + public data class Owned(val scope: ScopeKind) : Owner + /** The caller's array / buffer / segment / mmap; we never free it. [external] identifies the lender for debugging. */ + public data class Borrowed(val external: Any? = null) : Owner + /** A view's storage reference: a strong reference to [parent]; mutability delegated. */ + public class Alias(public val parent: Storage) : Owner { + override fun toString(): String = "Alias(parent=${parent.id})" + } +} + +/** Thrown on any access to a storage after its scope or the storage itself was closed (SKEEP-003 rule 2). */ +@ExperimentalMemoryApi +public class StorageClosedException( + public val storageId: StorageId, + public val origin: TensorId?, + message: String = "Storage ${storageId}${origin?.let { " (" + it.canonical + ")" } ?: ""} is closed", +) : IllegalStateException(message) + +/** + * The one and only owner of bytes (SKEEP-003 §0, §4.2). `TensorView` interprets a storage, `Tensor` + * is the DSL handle over a view — neither owns bytes. Sealed over the four kinds; [Heap] is final + * and common, [OffHeap] / [Mapped] / [Device] are abstract here and bound per platform + * (`MemorySegment` / `FileChannel.map` on the JVM, `malloc` / `mmap` on Native, heap fallbacks on + * JS/Wasm — slices #1019, #1020). + * + * Rules enforced here: exactly one byte owner; closing invalidates every alias; a borrowed storage + * is released (forgotten) but never freed; every access after close throws + * [StorageClosedException] carrying the id and origin — not a JVM crash, not silent corruption. + */ +@ExperimentalMemoryApi +public sealed class Storage : AutoCloseable { + public abstract val id: StorageId + public abstract val sizeBytes: Long + public abstract val owner: Owner + public abstract val domain: MemoryDomain + /** The `TensorId` these bytes back, for diagnostics; `null` for anonymous storage. */ + public abstract val debugOrigin: TensorId? + /** Where trace events about this storage go (allocation, close). */ + protected abstract val sink: TraceSink + + /** The lifetime class: from [Owner.Owned], else the parent's, else `AMBIENT`. */ + public val scope: ScopeKind + get() = when (val o = owner) { + is Owner.Owned -> o.scope + is Owner.Alias -> o.parent.scope + is Owner.Borrowed -> ScopeKind.AMBIENT + } + + private var closed: Boolean = false + + /** `true` until this storage — or, for an alias, its parent — is closed. */ + public val isAlive: Boolean + get() = !closed && ((owner as? Owner.Alias)?.parent?.isAlive ?: true) + + /** Whether writes are allowed: owned and borrowed-mutable storage yes; an alias delegates to its parent. */ + public abstract val isMutable: Boolean + + /** Throws [StorageClosedException] if this storage is no longer alive. Called by every accessor. */ + public fun checkAlive() { if (!isAlive) throw StorageClosedException(id, debugOrigin) } + + /** + * Close: an [Owner.Owned] storage releases its bytes (exactly once); an [Owner.Borrowed] storage + * is forgotten (the lender's bytes are untouched); an [Owner.Alias] is detached (its parent is + * unaffected). Idempotent. + */ + final override fun close() { + if (closed) return + closed = true + onClose() + if (sink.isEnabled && owner !is Owner.Alias) sink.emit(TraceEvent.Free(id.value, scope, sizeBytes)) + } + + /** Release platform resources (owned storage only); default nothing. */ + protected open fun onClose() {} + + /** A zero-copy alias over `[offsetBytes, offsetBytes + lengthBytes)` of this storage. */ + public abstract fun slice(offsetBytes: Long, lengthBytes: Long): Storage + + override fun toString(): String = "${this::class.simpleName}(${id}, ${sizeBytes} B, $owner, $domain${debugOrigin?.let { ", $it" } ?: ""}${if (isAlive) "" else ", closed"})" + + /** + * Heap storage: a Kotlin array on the managed heap — the JIT-friendliest kind, the only kind on + * JS/Wasm, the default for `Ambient` scope. Exactly one of [floats], [ints], [bytes] is non-null; + * [arrayOffset] (in elements of that array) and [sizeBytes] delimit the region. + * + * Kernels unwrap once per call (`floats` / `ints` / `bytes` + [arrayOffset]) — the Phase-2 spike + * showed per-element access through a view is the slow path by design. + */ + public class Heap private constructor( + override val id: StorageId, + public val floats: FloatArray?, + public val ints: IntArray?, + public val bytes: ByteArray?, + public val arrayOffset: Int, + override val sizeBytes: Long, + override val owner: Owner, + override val debugOrigin: TensorId?, + override val sink: TraceSink, + private val mutable: Boolean, + ) : Storage() { + override val domain: MemoryDomain get() = MemoryDomain.HOST_HEAP + override val isMutable: Boolean get() = (owner as? Owner.Alias)?.parent?.isMutable ?: mutable + + /** Bytes per element of the backing array (4 for floats/ints, 1 for bytes). */ + public val elementBytes: Int get() = if (bytes != null) 1 else 4 + /** Number of array elements this storage spans. */ + public val elementCount: Int get() = (sizeBytes / elementBytes).toInt() + + override fun slice(offsetBytes: Long, lengthBytes: Long): Heap { + checkAlive() + require(offsetBytes >= 0 && lengthBytes >= 0 && offsetBytes + lengthBytes <= sizeBytes) { "slice [$offsetBytes, ${offsetBytes + lengthBytes}) outside $sizeBytes bytes" } + require(offsetBytes % elementBytes == 0L && lengthBytes % elementBytes == 0L) { "slice must align to $elementBytes-byte elements" } + return Heap(StorageId.next(), floats, ints, bytes, arrayOffset + (offsetBytes / elementBytes).toInt(), lengthBytes, Owner.Alias(this), debugOrigin, sink, mutable) + } + + public companion object { + private fun create(floats: FloatArray?, ints: IntArray?, bytes: ByteArray?, offset: Int, count: Int, owner: Owner, origin: TensorId?, sink: TraceSink, mutable: Boolean): Heap { + val eb = if (bytes != null) 1 else 4 + val s = Heap(StorageId.next(), floats, ints, bytes, offset, count.toLong() * eb, owner, origin, sink, mutable) + if (sink.isEnabled && owner is Owner.Owned) sink.emit(TraceEvent.Allocation(s.id.value, owner.scope, s.sizeBytes, origin)) + return s + } + + /** Allocate [count] zeroed floats on the heap, owned by a scope of kind [scope]. */ + public fun floats(count: Int, scope: ScopeKind = ScopeKind.AMBIENT, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(FloatArray(count), null, null, 0, count, Owner.Owned(scope), origin, sink, true) + public fun ints(count: Int, scope: ScopeKind = ScopeKind.AMBIENT, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(null, IntArray(count), null, 0, count, Owner.Owned(scope), origin, sink, true) + public fun bytes(count: Int, scope: ScopeKind = ScopeKind.AMBIENT, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(null, null, ByteArray(count), 0, count, Owner.Owned(scope), origin, sink, true) + + /** Wrap the caller's array without copying — never freed by us (the #782 `copyOf` replacement). */ + public fun wrap(array: FloatArray, offset: Int = 0, count: Int = array.size - offset, mutable: Boolean = true, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(array, null, null, offset, count, Owner.Borrowed(array), origin, sink, mutable) + public fun wrap(array: IntArray, offset: Int = 0, count: Int = array.size - offset, mutable: Boolean = true, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(null, array, null, offset, count, Owner.Borrowed(array), origin, sink, mutable) + public fun wrap(array: ByteArray, offset: Int = 0, count: Int = array.size - offset, mutable: Boolean = true, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): Heap = + create(null, null, array, offset, count, Owner.Borrowed(array), origin, sink, mutable) + } + } + + /** Off-heap storage (`MemorySegment` / direct buffer / `malloc`): bound per platform in #1019/#1020. */ + public abstract class OffHeap : Storage() { override val domain: MemoryDomain get() = MemoryDomain.HOST_OFFHEAP } + + /** A mapped file region (`FileChannel.map` / `mmap`): bound per platform in #1019/#1020. */ + public abstract class Mapped : Storage() { override val domain: MemoryDomain get() = MemoryDomain.MMAP_FILE } + + /** An accelerator buffer — placeholder until a device backend is scheduled (PRD non-goal). */ + public abstract class Device : Storage() { override val domain: MemoryDomain get() = MemoryDomain.DEVICE_LOCAL } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/Placement.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/Placement.kt index 2234fec8f..99100ecaf 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/Placement.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/Placement.kt @@ -54,6 +54,8 @@ public enum class DeviceKind { public enum class MemoryDomain { /** Standard JVM / native heap allocation. */ HOST_HEAP, + /** Off-heap host memory: `MemorySegment` / direct `ByteBuffer` / `malloc` — not GC-managed, freed by its scope (SKEEP-003). */ + HOST_OFFHEAP, /** Pinned (non-pageable) host memory for fast DMA transfers. */ HOST_PINNED, /** Memory-mapped file (immutable, OS-paged). */ diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/StorageTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/StorageTest.kt new file mode 100644 index 000000000..3dfbac21e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/StorageTest.kt @@ -0,0 +1,104 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.MemoryDomain +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertIs +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** SKEEP-003 rules 1–2 (one byte owner, ownership enforced) for the common Heap storage. */ +@OptIn(ExperimentalMemoryApi::class) +class StorageTest { + + @Test + fun ownedHeapStorageHasIdSizeScopeAndIsFreedOnce() { + val sink = RecordingTraceSink() + val id = TensorId.parse("model.layers[0].attn.scores#step=1") + val s = Storage.Heap.floats(256, ScopeKind.FORWARD, origin = id, sink = sink) + assertTrue(s.isAlive); assertTrue(s.isMutable) + assertEquals(1024L, s.sizeBytes); assertEquals(256, s.elementCount); assertEquals(4, s.elementBytes) + assertEquals(ScopeKind.FORWARD, s.scope); assertEquals(MemoryDomain.HOST_HEAP, s.domain) + assertIs(s.owner); assertEquals(id, s.debugOrigin) + assertNotNull(s.floats); assertNull(s.ints); assertNull(s.bytes); assertEquals(0, s.arrayOffset) + val alloc = assertIs(sink.events().single()) + assertEquals(s.id.value, alloc.storageId); assertEquals(ScopeKind.FORWARD, alloc.scope); assertEquals(1024L, alloc.bytes); assertEquals(id, alloc.origin) + + s.close(); s.close() // idempotent + assertFalse(s.isAlive) + val free = assertIs(sink.events()[1]); assertEquals(s.id.value, free.storageId); assertEquals(1024L, free.bytes) + assertEquals(2, sink.events().size) + val ex = assertFailsWith { s.checkAlive() } + assertEquals(s.id, ex.storageId); assertEquals(id, ex.origin); assertTrue(ex.message!!.contains("attn.scores")) + } + + @Test + fun storageIdsAreMonotonicAndDistinct() { + val a = Storage.Heap.bytes(4); val b = Storage.Heap.ints(4); val c = Storage.Heap.floats(4) + assertTrue(a.id.value < b.id.value && b.id.value < c.id.value) + assertEquals("#${a.id.value}", a.id.toString()) + assertEquals(4L, a.sizeBytes); assertEquals(16L, b.sizeBytes); assertEquals(1, a.elementBytes) + } + + @Test + fun borrowedStorageIsNeverFreedOnlyReleased() { + val sink = RecordingTraceSink() + val arr = FloatArray(8) { it.toFloat() } + val s = Storage.Heap.wrap(arr, offset = 2, count = 4, mutable = false, sink = sink) + assertIs(s.owner); assertSame(arr, (s.owner as Owner.Borrowed).external) + assertEquals(ScopeKind.AMBIENT, s.scope); assertFalse(s.isMutable) + assertSame(arr, s.floats); assertEquals(2, s.arrayOffset); assertEquals(16L, s.sizeBytes) + assertTrue(sink.events().isEmpty()) // borrowing is not an allocation + s.close() // release: forget, do not touch the caller's bytes + assertFalse(s.isAlive) + assertEquals(listOf(0f, 1f, 2f, 3f, 4f, 5f, 6f, 7f), arr.toList()) + assertFailsWith { s.checkAlive() } + } + + @Test + fun aliasKeepsParentAliveDelegatesMutabilityAndDiesWithIt() { + val parent = Storage.Heap.floats(16, ScopeKind.MODEL) + val view = parent.slice(offsetBytes = 16, lengthBytes = 32) + assertIs(view.owner); assertSame(parent, (view.owner as Owner.Alias).parent) + assertEquals(ScopeKind.MODEL, view.scope) // inherited + assertSame(parent.floats, view.floats); assertEquals(4, view.arrayOffset); assertEquals(32L, view.sizeBytes); assertEquals(8, view.elementCount) + assertTrue(view.isMutable) + val roView = Storage.Heap.wrap(FloatArray(4), mutable = false).slice(0, 8) + assertFalse(roView.isMutable) // delegated + // nested alias + val inner = view.slice(8, 8) + assertEquals(6, inner.arrayOffset); assertEquals(8L, inner.sizeBytes) + // closing an alias detaches it only + inner.close(); assertFalse(inner.isAlive); assertTrue(view.isAlive); assertTrue(parent.isAlive) + // closing the parent invalidates every alias + parent.close() + assertFalse(view.isAlive) + assertFailsWith { view.checkAlive() } + assertFailsWith { parent.slice(0, 4) } + } + + @Test + fun sliceBoundsAndAlignmentAreChecked() { + val s = Storage.Heap.ints(4) + assertFailsWith { s.slice(0, 20) } + assertFailsWith { s.slice(-4, 4) } + assertFailsWith { s.slice(2, 4) } // not 4-byte aligned + assertEquals(0L, s.slice(16, 0).sizeBytes) + assertEquals(2, Storage.Heap.bytes(8).slice(2, 3).arrayOffset) + } + + @Test + fun toStringNamesKindIdSizeOwnerAndState() { + val s = Storage.Heap.floats(2, origin = TensorId.parse("x.w")) + val t = s.toString() + assertTrue(t.startsWith("Heap(#"), t); assertTrue(t.contains("8 B")); assertTrue(t.contains("Owned(scope=AMBIENT)")); assertTrue(t.contains("x.w")) + s.close(); assertTrue(s.toString().endsWith("closed)")) + } +}