From 5b9b0e4600670935bc3906f40a7578b09160e8ea Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 11:23:46 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20Storage=20=E2=80=94=20the=20one?= =?UTF-8?q?=20byte=20owner:=20sealed=20kinds,=20Owner,=20StorageId,=20Stor?= =?UTF-8?q?ageClosedException,=20Heap=20on=20every=20target=20(SKEEP-003?= =?UTF-8?q?=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)")) + } +}