From aaaff96901a6bd32842b111ec489594d979278ed Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 14:07:24 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20Scope=20=E2=80=94=20Model=20/?= =?UTF-8?q?=20Forward=20/=20Ambient=20lifetimes;=20recycled=20Forward=20sl?= =?UTF-8?q?ab=20with=20reset()=20and=20retain()=20(SKEEP-003=20P2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F3 / M1-F4. SKEEP-003 §4.5: weights and activations must not share a lifetime — the two historical arena failures (a model-lifetime arena pinning tens of GB; per-call arenas leaking per matmul) were violations of exactly this split. - sk.ainet.lang.memory.Scope (sealed): kind, liveBytes, sink, allocate(), allocateFloats(). Scope.Ambient is the GC-managed default: untracked, close() is a no-op, notebook code is unchanged. - ModelScope: tracks every storage it allocates, maps (mapFile) or adopts and closes them all exactly once — weights unmapped, off-heap freed — so Model.load(...).use { } is deterministic. - ForwardScope(slabFloats): one pre-sized heap slab (the Phase-2 spike #1016 says heap activations by default on the JVM), bump-allocated as zero-copy slab views; reset() closes the views handed out, frees overflow and rewinds the offset, so steady-state decode allocates zero slab bytes (asserted in the test: one Allocation event for the slab and none per step); usedFloats / peakFloats / overflowBytes / steps feed the planner and the plan-vs-actual check; retain(storage, to) is the one sanctioned escape (a copy), everything else becomes a StorageClosedException after reset. - ExecutionContext.memoryScope default member → Scope.Ambient (opt-in). - ScopeTest: ambient default, bump/reset cycle with zero steady-state allocations, overflow accounting, retain, model tracking/closing, and the weights-vs-activations separation. BCV dumps regenerated. Closes #1021 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 | 71 +++++++ .../sk/ainet/context/ExecutionContext.kt | 8 + .../kotlin/sk/ainet/lang/memory/Scope.kt | 179 ++++++++++++++++++ .../kotlin/sk/ainet/lang/memory/ScopeTest.kt | 118 ++++++++++++ 6 files changed, 379 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Scope.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/ScopeTest.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 facf4eff7..cedfccd97 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 @@ -18,6 +18,7 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun getInTraining ()Z public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; 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 1dbc6e43f..b314e1db8 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 @@ -173,6 +173,7 @@ public final class sk/ainet/lang/graph/DefaultGraphExecutionContext : sk/ainet/l public fun getInTraining ()Z public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/KspTensorOps; @@ -449,6 +450,7 @@ public final class sk/ainet/lang/graph/exec/GraphExecutionContext$DefaultImpls { public static fun getHooks (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/nn/hooks/ForwardHooks; public static fun getInTraining (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Z public static fun getMemoryPlanner (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public static fun getMemoryScope (Lsk/ainet/lang/graph/exec/GraphExecutionContext;)Lsk/ainet/lang/memory/Scope; 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; 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 25025d56f..7486398e0 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 @@ -205,6 +205,7 @@ public final class sk/ainet/context/DefaultDataExecutionContext : sk/ainet/conte public fun getInTraining ()Z public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; @@ -234,6 +235,7 @@ public abstract interface class sk/ainet/context/ExecutionContext { public fun getInTraining ()Z public abstract fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public abstract fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public abstract fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; @@ -261,6 +263,7 @@ public final class sk/ainet/context/ExecutionContext$DefaultImpls { public static fun getHooks (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/nn/hooks/ForwardHooks; public static fun getInTraining (Lsk/ainet/context/ExecutionContext;)Z public static fun getMemoryPlanner (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public static fun getMemoryScope (Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/memory/Scope; 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; @@ -362,6 +365,7 @@ public final class sk/ainet/context/PhaseOverridingExecutionContext : sk/ainet/c public fun getInTraining ()Z public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; @@ -411,6 +415,7 @@ public final class sk/ainet/context/TrainingExecutionContext$DefaultImpls { public static fun getHooks (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/nn/hooks/ForwardHooks; public static fun getInTraining (Lsk/ainet/context/TrainingExecutionContext;)Z public static fun getMemoryPlanner (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public static fun getMemoryScope (Lsk/ainet/context/TrainingExecutionContext;)Lsk/ainet/lang/memory/Scope; 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; @@ -725,6 +730,27 @@ public final class sk/ainet/lang/memory/FormatKt { public static final fun getFormatOrNull (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/memory/Format; } +public final class sk/ainet/lang/memory/ForwardScope : sk/ainet/lang/memory/Scope { + public fun (ILsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;)V + public synthetic fun (ILsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun allocate (JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage; + public fun allocateFloats (ILsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage$Heap; + public fun close ()V + public fun getKind ()Lsk/ainet/lang/memory/ScopeKind; + public fun getLiveBytes ()J + public final fun getName ()Ljava/lang/String; + public final fun getOverflowBytes ()J + public final fun getPeakFloats ()I + public fun getSink ()Lsk/ainet/lang/memory/trace/TraceSink; + public final fun getSlabFloats ()I + public final fun getSteps ()J + public final fun getUsedFloats ()I + public final fun isClosed ()Z + public final fun reset ()V + public final fun retain (Lsk/ainet/lang/memory/Storage$Heap;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun retain$default (Lsk/ainet/lang/memory/ForwardScope;Lsk/ainet/lang/memory/Storage$Heap;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; +} + public final class sk/ainet/lang/memory/MappedBufferStorage : sk/ainet/lang/memory/Storage$Mapped { public static final field Companion Lsk/ainet/lang/memory/MappedBufferStorage$Companion; public synthetic fun (JLjava/nio/file/Path;JLjava/nio/ByteBuffer;Lsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;Lkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -766,6 +792,24 @@ public final class sk/ainet/lang/memory/MappedFileStorage$Companion { public static synthetic fun map$default (Lsk/ainet/lang/memory/MappedFileStorage$Companion;Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/MappedFileStorage; } +public final class sk/ainet/lang/memory/ModelScope : sk/ainet/lang/memory/Scope { + public fun ()V + public fun (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;)V + public synthetic fun (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun adopt (Lsk/ainet/lang/memory/Storage;)Lsk/ainet/lang/memory/Storage; + public fun allocate (JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage; + public fun allocateFloats (ILsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage$Heap; + public fun close ()V + public fun getKind ()Lsk/ainet/lang/memory/ScopeKind; + public fun getLiveBytes ()J + public final fun getName ()Ljava/lang/String; + public fun getSink ()Lsk/ainet/lang/memory/trace/TraceSink; + public final fun getStorageCount ()I + public final fun isClosed ()Z + public final fun mapFile (Ljava/lang/String;JJLsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage; + public static synthetic fun mapFile$default (Lsk/ainet/lang/memory/ModelScope;Ljava/lang/String;JJLsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage; +} + public abstract interface class sk/ainet/lang/memory/Owner { } @@ -823,6 +867,31 @@ public final class sk/ainet/lang/memory/PlatformStorageInfo { public fun toString ()Ljava/lang/String; } +public abstract interface class sk/ainet/lang/memory/Scope : java/lang/AutoCloseable { + public abstract fun allocate (JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage; + public static synthetic fun allocate$default (Lsk/ainet/lang/memory/Scope;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage; + public abstract fun allocateFloats (ILsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage$Heap; + public static synthetic fun allocateFloats$default (Lsk/ainet/lang/memory/Scope;ILsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; + public abstract fun getKind ()Lsk/ainet/lang/memory/ScopeKind; + public abstract fun getLiveBytes ()J + public abstract fun getSink ()Lsk/ainet/lang/memory/trace/TraceSink; +} + +public final class sk/ainet/lang/memory/Scope$Ambient : sk/ainet/lang/memory/Scope { + public static final field INSTANCE Lsk/ainet/lang/memory/Scope$Ambient; + public fun allocate (JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage; + public fun allocateFloats (ILsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/Storage$Heap; + public fun close ()V + public fun getKind ()Lsk/ainet/lang/memory/ScopeKind; + public fun getLiveBytes ()J + public fun getSink ()Lsk/ainet/lang/memory/trace/TraceSink; +} + +public final class sk/ainet/lang/memory/Scope$DefaultImpls { + public static synthetic fun allocate$default (Lsk/ainet/lang/memory/Scope;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage; + public static synthetic fun allocateFloats$default (Lsk/ainet/lang/memory/Scope;ILsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; +} + 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; @@ -1501,6 +1570,7 @@ public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ai public fun getInTraining ()Z public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; public fun getMemoryPlanner ()Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; @@ -1837,6 +1907,7 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContext$DefaultImpls { public static fun getHooks (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/nn/hooks/ForwardHooks; public static fun getInTraining (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Z public static fun getMemoryPlanner (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/tensor/storage/MemoryPlanner; + public static fun getMemoryScope (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;)Lsk/ainet/lang/memory/Scope; 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; 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 39d004f0b..5ba678385 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 @@ -22,6 +22,14 @@ public interface ExecutionContext { @sk.ainet.lang.memory.ExperimentalMemoryApi public val traceSink: sk.ainet.lang.memory.trace.TraceSink get() = sk.ainet.lang.memory.trace.NoopTraceSink + /** + * The scope new activations and adapter outputs are allocated in (SKEEP-003 §4.5). Default + * [sk.ainet.lang.memory.Scope.Ambient] — GC-managed, today's behaviour; a generation loop opts + * in by providing a `ForwardScope` and calling `reset()` per step. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public val memoryScope: sk.ainet.lang.memory.Scope get() = sk.ainet.lang.memory.Scope.Ambient + 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/Scope.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Scope.kt new file mode 100644 index 000000000..ea3b6d761 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Scope.kt @@ -0,0 +1,179 @@ +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 + +/** + * A lifetime that owned storage belongs to (SKEEP-003 §0 *Scope*, §4.5). Three kinds: + * [Scope.Ambient] (GC-managed — the default, today's behaviour), [ForwardScope] (activations; + * recycled every forward pass by [ForwardScope.reset]), [ModelScope] (weights, KV backing; closed + * on `model.close()`). Closing a scope invalidates every storage it owns (rule 2); scopes are + * opt-in per `ExecutionContext` so `val c = a matMul b` in a notebook keeps working unchanged. + * + * Both historical arena failures were violations of exactly this split: activations in a + * model-lifetime arena (tens of GB pinned) and per-call arenas (leak per matmul). A `Forward` + * scope matches the forward-pass lifetime; a `Model` scope the model's. + */ +@ExperimentalMemoryApi +public sealed interface Scope : AutoCloseable { + public val kind: ScopeKind + /** Bytes currently owned and alive in this scope. */ + public val liveBytes: Long + /** Where this scope's allocation / reset events go. */ + public val sink: TraceSink + + /** Allocate [bytes] zeroed bytes in [domain] (platform-bound), owned by this scope. */ + public fun allocate(bytes: Long, domain: MemoryDomain = MemoryDomain.HOST_HEAP, origin: TensorId? = null): Storage + + /** Allocate [count] zeroed floats on the heap, owned by this scope (the kernel-friendliest form). */ + public fun allocateFloats(count: Int, origin: TensorId? = null): Storage.Heap + + /** The GC-managed default: nothing is tracked, nothing is freed explicitly; [close] and [liveBytes] are no-ops. */ + public object Ambient : Scope { + override val kind: ScopeKind get() = ScopeKind.AMBIENT + override val liveBytes: Long get() = 0L + override val sink: TraceSink get() = NoopTraceSink + override fun allocate(bytes: Long, domain: MemoryDomain, origin: TensorId?): Storage = PlatformStorage.allocate(bytes, domain, ScopeKind.AMBIENT, origin, NoopTraceSink) + override fun allocateFloats(count: Int, origin: TensorId?): Storage.Heap = Storage.Heap.floats(count, ScopeKind.AMBIENT, origin, NoopTraceSink) + override fun close() {} + } +} + +/** + * `Scope.Model`: weights, KV-cache backing, embedding tables — everything that lives until the + * model is closed. Tracks every storage it allocates or maps and closes them all in [close] + * (deterministic: the JVM unmaps the file, native `free`s/`munmap`s). Idiomatic use: + * `ModelScope().use { model -> … }`. + */ +@ExperimentalMemoryApi +public class ModelScope(override val sink: TraceSink = NoopTraceSink, public val name: String = "model") : Scope { + override val kind: ScopeKind get() = ScopeKind.MODEL + private val owned = ArrayList() + private var closed = false + public val isClosed: Boolean get() = closed + + override val liveBytes: Long get() = owned.sumOf { if (it.isAlive) it.sizeBytes else 0L } + /** Number of storages this scope owns (alive or not). */ + public val storageCount: Int get() = owned.size + + private fun track(s: Storage): Storage { check(!closed) { "ModelScope '$name' is closed" }; owned += s; return s } + + override fun allocate(bytes: Long, domain: MemoryDomain, origin: TensorId?): Storage = track(PlatformStorage.allocate(bytes, domain, ScopeKind.MODEL, origin, sink)) + override fun allocateFloats(count: Int, origin: TensorId?): Storage.Heap = track(Storage.Heap.floats(count, ScopeKind.MODEL, origin, sink)) as Storage.Heap + + /** Map a file region (packed weights) into this scope; unmapped when the scope closes. */ + public fun mapFile(path: String, fileOffset: Long, length: Long, origin: TensorId? = null): Storage = track(PlatformStorage.mapFile(path, fileOffset, length, ScopeKind.MODEL, origin, sink)) + + /** Adopt a storage allocated elsewhere (e.g. by a loader) so the scope closes it. */ + public fun adopt(storage: Storage): Storage = track(storage) + + /** Close every owned storage (weights unmapped, off-heap freed) — exactly once. */ + override fun close() { + if (closed) return + closed = true + for (s in owned.asReversed()) s.close() + owned.clear() + } +} + +/** + * `Scope.Forward`: activations, attention scratch, adapter outputs — recycled every forward pass. + * One pre-sized slab ([slabFloats] floats on the heap, per the Phase-2 spike: heap activations on + * the JVM by default) is bump-allocated; [reset] at the end of the step rewinds the offset and + * invalidates the views handed out, so steady-state decode allocates zero slab bytes. If a step + * needs more than the slab, an overflow storage is allocated (tracked, closed at [reset], counted + * in [overflowBytes] so the planner can resize the slab). Outputs that must outlive the step are + * copied out with [retain]. + */ +@ExperimentalMemoryApi +public class ForwardScope( + public val slabFloats: Int, + override val sink: TraceSink = NoopTraceSink, + public val name: String = "forward", +) : Scope { + init { require(slabFloats >= 0) { "slabFloats must be >= 0" } } + + override val kind: ScopeKind get() = ScopeKind.FORWARD + + private val slab: Storage.Heap = Storage.Heap.floats(slabFloats, ScopeKind.FORWARD, TensorId(listOf(name), "slab"), sink) + private var offset: Int = 0 + private val handedOut = ArrayList() + private val overflow = ArrayList() + private var closed = false + + /** Floats allocated from the slab in the current step. */ + public val usedFloats: Int get() = offset + /** High-water mark of slab use across steps (floats). */ + public var peakFloats: Int = 0 + private set + /** Bytes allocated outside the slab in the current step (the planner should grow the slab by this). */ + public val overflowBytes: Long get() = overflow.sumOf { it.sizeBytes } + /** Steps completed ([reset] calls). */ + public var steps: Long = 0L + private set + public val isClosed: Boolean get() = closed + + override val liveBytes: Long get() = offset.toLong() * 4 + overflowBytes + + private fun checkOpen() { check(!closed) { "ForwardScope '$name' is closed" } } + + /** Bump-allocate [count] floats from the slab (or an overflow storage when the slab is exhausted). */ + override fun allocateFloats(count: Int, origin: TensorId?): Storage.Heap { + checkOpen(); require(count >= 0) + if (offset + count <= slabFloats) { + val view = slab.slice(offset.toLong() * 4, count.toLong() * 4) + offset += count + if (offset > peakFloats) peakFloats = offset + handedOut += view + return view + } + val extra = Storage.Heap.floats(count, ScopeKind.FORWARD, origin, sink) + overflow += extra + return extra + } + + override fun allocate(bytes: Long, domain: MemoryDomain, origin: TensorId?): Storage { + checkOpen() + if (domain == MemoryDomain.HOST_HEAP && bytes % 4 == 0L && bytes / 4 <= Int.MAX_VALUE) return allocateFloats((bytes / 4).toInt(), origin) + val s = PlatformStorage.allocate(bytes, domain, ScopeKind.FORWARD, origin, sink) + overflow += s + return s + } + + /** + * Copy a step-scoped storage out into a [to] scope (default [Scope.Ambient]) so it survives + * [reset] — the one sanctioned escape (the rest is a use-after-reset, caught as + * `StorageClosedException`). + */ + public fun retain(storage: Storage.Heap, to: Scope = Scope.Ambient, origin: TensorId? = storage.debugOrigin): Storage.Heap { + storage.checkAlive() + val floats = storage.floats ?: throw IllegalArgumentException("retain() supports float storage in this milestone") + val out = to.allocateFloats(storage.elementCount, origin) + floats.copyInto(out.floats!!, out.arrayOffset, storage.arrayOffset, storage.arrayOffset + storage.elementCount) + return out + } + + /** End of step: invalidate every view handed out, free overflow, rewind the slab. Emits [TraceEvent.ScopeReset]. */ + public fun reset() { + checkOpen() + val before = liveBytes + for (s in handedOut) s.close() + handedOut.clear() + for (s in overflow) s.close() + overflow.clear() + offset = 0 + steps++ + if (sink.isEnabled) sink.emit(TraceEvent.ScopeReset(ScopeKind.FORWARD, before, 0L)) + } + + /** Close the slab itself (end of the model's life). */ + override fun close() { + if (closed) return + reset() + closed = true + slab.close() + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/ScopeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/ScopeTest.kt new file mode 100644 index 000000000..5072c8bde --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/ScopeTest.kt @@ -0,0 +1,118 @@ +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.assertSame +import kotlin.test.assertTrue + +/** SKEEP-003 §4.5: Forward (recycled per step), Model (closed with the model), Ambient (GC) — M1-F3 / M1-F4. */ +@OptIn(ExperimentalMemoryApi::class) +class ScopeTest { + + @Test + fun ambientIsTheDefaultAndUntracked() { + assertEquals(ScopeKind.AMBIENT, Scope.Ambient.kind) + val s = Scope.Ambient.allocateFloats(8) + assertEquals(ScopeKind.AMBIENT, s.scope); assertEquals(0L, Scope.Ambient.liveBytes) + Scope.Ambient.close() // no-op + assertTrue(s.isAlive) + assertIs(Scope.Ambient.allocate(16, MemoryDomain.HOST_HEAP)) + } + + @Test + fun forwardScopeBumpAllocatesAndResetRecyclesWithZeroSlabAllocations() { + val sink = RecordingTraceSink() + val f = ForwardScope(slabFloats = 1024, sink = sink) + val slabAllocs = sink.eventsOf().size // the slab itself + assertEquals(1, slabAllocs) + var last: Storage.Heap? = null + repeat(3) { step -> + val a = f.allocateFloats(256); val b = f.allocateFloats(512) + assertEquals(768, f.usedFloats); assertEquals(768L * 4, f.liveBytes) + assertEquals(ScopeKind.FORWARD, a.scope); assertIs(a.owner) // views over the slab + a.floats!![a.arrayOffset] = step.toFloat(); b.floats!![b.arrayOffset + 511] = 1f + assertSame(a.floats, b.floats) // same slab array + assertEquals(0, a.arrayOffset); assertEquals(256, b.arrayOffset) + last = a + f.reset() + assertEquals(0, f.usedFloats); assertFalse(a.isAlive); assertFalse(b.isAlive) + assertFailsWith { a.checkAlive() } + } + assertEquals(3L, f.steps); assertEquals(768, f.peakFloats) + // steady state: no new Allocation events after the slab — only the per-step ScopeReset events + assertEquals(1, sink.eventsOf().size) + assertEquals(3, sink.eventsOf().size) + assertEquals(768L * 4, sink.eventsOf().first().liveBytesBefore) + f.close(); assertTrue(f.isClosed) + assertFailsWith { f.allocateFloats(1) } + assertFalse(last!!.isAlive) + } + + @Test + fun forwardOverflowIsCountedAndFreedAtReset() { + val sink = RecordingTraceSink() + val f = ForwardScope(slabFloats = 100, sink = sink) + val a = f.allocateFloats(80) // in the slab + val o = f.allocateFloats(50) // does not fit → overflow storage + assertIs(o.owner); assertEquals(200L, f.overflowBytes); assertEquals(80, f.usedFloats) + assertEquals(80L * 4 + 200L, f.liveBytes) + assertEquals(2, sink.eventsOf().size) // slab + overflow + f.reset() + assertEquals(0L, f.overflowBytes); assertFalse(o.isAlive); assertFalse(a.isAlive) + assertTrue(sink.eventsOf().any { it.storageId == o.id.value }) + // non-heap domains go to the platform and are treated as overflow too + val p = f.allocate(64, MemoryDomain.HOST_OFFHEAP) + assertEquals(ScopeKind.FORWARD, p.scope); assertTrue(f.liveBytes >= 64) + f.reset(); assertFalse(p.isAlive) + } + + @Test + fun retainCopiesAStepResultOutOfTheForwardScope() { + val f = ForwardScope(64) + val act = f.allocateFloats(4) + act.floats!![act.arrayOffset + 2] = 7f + val kept = f.retain(act, origin = TensorId.parse("model.logits")) + assertEquals(ScopeKind.AMBIENT, kept.scope); assertEquals(7f, kept.floats!![kept.arrayOffset + 2]); assertEquals("model.logits", kept.debugOrigin!!.canonical) + f.reset() + assertFalse(act.isAlive); assertTrue(kept.isAlive) + val model = ModelScope() + val inModel = f.retain(f.allocateFloats(2), to = model) + assertEquals(ScopeKind.MODEL, inModel.scope) + assertFailsWith { f.retain(act) } // already reset + } + + @Test + fun modelScopeTracksAndClosesEverything() { + val sink = RecordingTraceSink() + val m = ModelScope(sink, "llama") + val w = m.allocateFloats(16, TensorId.parse("model.embed_tokens.weight")) + val o = m.allocate(32, MemoryDomain.HOST_OFFHEAP) + val adopted = m.adopt(Storage.Heap.bytes(8)) + assertEquals(ScopeKind.MODEL, w.scope); assertEquals(ScopeKind.MODEL, o.scope); assertEquals(3, m.storageCount) + assertEquals(64L + 32L + 8L, m.liveBytes) + m.close(); m.close() + assertTrue(m.isClosed); assertEquals(0L, m.liveBytes) + assertFalse(w.isAlive); assertFalse(o.isAlive); assertFalse(adopted.isAlive) + assertFailsWith { m.allocateFloats(1) } + // two of the three emit Free: the adopted storage was created with the default no-op sink + assertEquals(2, sink.eventsOf().size) + } + + @Test + fun scopesSeparateWeightsFromActivations() { + // the two historical arena failures: weights and activations must not share a lifetime + val model = ModelScope(); val fwd = ForwardScope(256) + val weights = model.allocateFloats(64) + repeat(10) { fwd.allocateFloats(128); fwd.allocateFloats(128); fwd.reset() } + assertTrue(weights.isAlive); assertEquals(256, fwd.peakFloats); assertEquals(10L, fwd.steps) + fwd.close(); assertTrue(weights.isAlive) + model.close(); assertFalse(weights.isAlive) + } +}