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 c405d0e81..29be4c0d9 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 @@ -634,6 +634,37 @@ public final class sk/ainet/java/TrainingResult { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/AllocationSpec { + public static final field Companion Lsk/ainet/lang/memory/AllocationSpec$Companion; + public static final field DEFAULT_ALIGNMENT I + public fun (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZI)V + public synthetic fun (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZIILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Lsk/ainet/lang/memory/Format; + public final fun component2 ()J + public final fun component3 ()Lsk/ainet/lang/tensor/storage/MemoryDomain; + public final fun component4 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component5 ()Z + public final fun component6 ()I + public final fun copy (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZI)Lsk/ainet/lang/memory/AllocationSpec; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/AllocationSpec;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZIILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; + public fun equals (Ljava/lang/Object;)Z + public final fun getAlignment ()I + public final fun getBytes ()J + public final fun getBytesOrNull ()Ljava/lang/Long; + public final fun getDomain ()Lsk/ainet/lang/tensor/storage/MemoryDomain; + public final fun getElementCount ()J + public final fun getFormat ()Lsk/ainet/lang/memory/Format; + public final fun getMutable ()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/AllocationSpec$Companion { + public final fun of (Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZI)Lsk/ainet/lang/memory/AllocationSpec; + public static synthetic fun of$default (Lsk/ainet/lang/memory/AllocationSpec$Companion;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;ZIILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; +} + public abstract interface annotation class sk/ainet/lang/memory/ExperimentalMemoryApi : java/lang/annotation/Annotation { } @@ -663,6 +694,15 @@ 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/ScopeKind : java/lang/Enum { + public static final field AMBIENT Lsk/ainet/lang/memory/ScopeKind; + public static final field FORWARD Lsk/ainet/lang/memory/ScopeKind; + public static final field MODEL Lsk/ainet/lang/memory/ScopeKind; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/ScopeKind; + public static fun values ()[Lsk/ainet/lang/memory/ScopeKind; +} + 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 @@ -5550,6 +5590,7 @@ public final class sk/ainet/lang/tensor/storage/StorageSpec { public final fun getOwnership ()Lsk/ainet/lang/tensor/storage/Ownership; public final fun getPlacement ()Lsk/ainet/lang/tensor/storage/Placement; public fun hashCode ()I + public final fun toAllocationSpec (J)Lsk/ainet/lang/memory/AllocationSpec; public fun toString ()Ljava/lang/String; } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/AllocationSpec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/AllocationSpec.kt new file mode 100644 index 000000000..118614e1a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/AllocationSpec.kt @@ -0,0 +1,74 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.MemoryDomain + +/** + * The lifetime class an allocation belongs to (SKEEP-003 §4.5). `Scope` itself — the object that + * allocates and frees — arrives with milestone M1; this enum is its `kind`, declared now so that + * memory plans and allocation specs can name the lifetime without depending on the allocator. + */ +@ExperimentalMemoryApi +public enum class ScopeKind { + /** Lives until the model is closed: weights, KV-cache backing, embedding tables. */ + MODEL, + /** Recycled every forward pass: activations, attention scratch, adapter outputs. */ + FORWARD, + /** Garbage-collected, no explicit lifetime — the default for notebooks, tests and ad-hoc tensors. */ + AMBIENT, +} + +/** + * What an allocation needs: the [Format] of the elements, how many, and where / for how long the + * bytes should live. Pure description — owns nothing. The single input of the memory plan + * (milestone M0) and, from M1, of `Storage.allocate(spec, scope)`. + * + * Replaces the never-consumed `StorageSpec` (decision recorded in SKEEP-003: "StorageSpec becomes + * the allocation spec"). + * + * @property format dtype + encoding of the elements + * @property elementCount logical number of elements + * @property domain where the bytes should be (heap, off-heap/pinned, mapped file, device …) + * @property scope the lifetime class the allocation belongs to + * @property mutable whether the bytes may be written after allocation + * @property alignment required byte alignment of the start of the buffer (SIMD kernels want 16–64) + */ +@ExperimentalMemoryApi +public data class AllocationSpec( + val format: Format, + val elementCount: Long, + val domain: MemoryDomain = MemoryDomain.HOST_HEAP, + val scope: ScopeKind = ScopeKind.AMBIENT, + val mutable: Boolean = true, + val alignment: Int = DEFAULT_ALIGNMENT, +) { + init { + require(elementCount >= 0) { "elementCount must be >= 0, was $elementCount" } + require(alignment > 0 && (alignment and (alignment - 1)) == 0) { "alignment must be a power of two, was $alignment" } + } + + /** Physical bytes this allocation needs, or `null` when the encoding cannot tell (opaque payloads). */ + val bytesOrNull: Long? get() = format.physicalBytes(elementCount) + + /** + * Physical bytes this allocation needs. + * @throws IllegalStateException for an encoding that cannot compute its size (see [bytesOrNull]) + */ + val bytes: Long + get() = bytesOrNull ?: throw IllegalStateException("Encoding ${format.encoding.name} cannot compute a byte size for $elementCount elements") + + public companion object { + /** 64 bytes: satisfies AVX-512 / NEON / cache-line alignment for every current kernel. */ + public const val DEFAULT_ALIGNMENT: Int = 64 + + /** Spec for a tensor of [shape] in [format]. */ + public fun of( + format: Format, + shape: Shape, + domain: MemoryDomain = MemoryDomain.HOST_HEAP, + scope: ScopeKind = ScopeKind.AMBIENT, + mutable: Boolean = true, + alignment: Int = DEFAULT_ALIGNMENT, + ): AllocationSpec = AllocationSpec(format, shape.volume.toLong(), domain, scope, mutable, alignment) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/StorageSpec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/StorageSpec.kt index f2a40c629..927d9fe2b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/StorageSpec.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/StorageSpec.kt @@ -1,5 +1,9 @@ package sk.ainet.lang.tensor.storage +import sk.ainet.lang.memory.AllocationSpec +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind import sk.ainet.lang.types.DType /** @@ -11,7 +15,15 @@ import sk.ainet.lang.types.DType * [sk.ainet.lang.tensor.data.TensorFactoryRegistry]). Existing dtype-based * lookups remain as a convenience — they build a default [StorageSpec] * with [TensorEncoding.Dense] and [Ownership.OWNED]. + * + * Deprecated (SKEEP-003 Phase 0): never consumed by any factory; the allocation + * description is [sk.ainet.lang.memory.AllocationSpec] (`Format` + element count + + * domain + scope). Use [toAllocationSpec] to convert. Removed at the next major release. */ +@Deprecated( + message = "StorageSpec was never consumed; describe allocations with sk.ainet.lang.memory.AllocationSpec (SKEEP-003).", + replaceWith = ReplaceWith("AllocationSpec", "sk.ainet.lang.memory.AllocationSpec"), +) public data class StorageSpec( val logicalType: LogicalDType, val encoding: TensorEncoding = TensorEncoding.Dense(logicalType.sizeInBytes), @@ -21,8 +33,24 @@ public data class StorageSpec( /** The [DType] of [logicalType] (SKEEP-003 Phase 0 bridge; see [LogicalDType.toDType]). */ val dtype: DType get() = logicalType.toDType() + /** + * The [AllocationSpec] equivalent of this spec for [elementCount] elements: `Format(dtype, + * encoding)`, the placement's memory domain, `MODEL` scope for persistent placements and + * `AMBIENT` otherwise, mutable only when owned. + */ + @OptIn(ExperimentalMemoryApi::class) + public fun toAllocationSpec(elementCount: Long): AllocationSpec = AllocationSpec( + format = Format(dtype, encoding), + elementCount = elementCount, + domain = placement.domain, + scope = if (placement.residency == Residency.PERSISTENT) ScopeKind.MODEL else ScopeKind.AMBIENT, + mutable = ownership == Ownership.OWNED, + ) + + @Suppress("DEPRECATION") // the factories build the deprecated type on purpose public companion object { /** Build a default spec from a legacy DType (dense, owned, CPU heap). */ + @Deprecated("StorageSpec is deprecated; build an AllocationSpec (sk.ainet.lang.memory).") public fun fromDType(dtype: DType): StorageSpec { val logical = dtype.toLogicalDType() return StorageSpec( @@ -34,6 +62,7 @@ public data class StorageSpec( } /** Spec for borrowed dense data. */ + @Deprecated("StorageSpec is deprecated; build an AllocationSpec (sk.ainet.lang.memory).") public fun borrowed(dtype: DType): StorageSpec { val logical = dtype.toLogicalDType() return StorageSpec( @@ -45,6 +74,7 @@ public data class StorageSpec( } /** Spec for Q4_K packed data. */ + @Deprecated("StorageSpec is deprecated; build an AllocationSpec (sk.ainet.lang.memory).") public fun q4k(placement: Placement = Placement.CPU_HEAP): StorageSpec = StorageSpec( logicalType = LogicalDType.FLOAT32, encoding = TensorEncoding.Q4_K, @@ -53,6 +83,7 @@ public data class StorageSpec( ) /** Spec for Q8_0 packed data. */ + @Deprecated("StorageSpec is deprecated; build an AllocationSpec (sk.ainet.lang.memory).") public fun q80(placement: Placement = Placement.CPU_HEAP): StorageSpec = StorageSpec( logicalType = LogicalDType.FLOAT32, encoding = TensorEncoding.Q8_0, @@ -61,6 +92,7 @@ public data class StorageSpec( ) /** Spec for file-backed weights. */ + @Deprecated("StorageSpec is deprecated; build an AllocationSpec (sk.ainet.lang.memory).") public fun mmapWeights(dtype: DType): StorageSpec { val logical = dtype.toLogicalDType() return StorageSpec( diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/AllocationSpecTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/AllocationSpecTest.kt new file mode 100644 index 000000000..75ffa640e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/AllocationSpecTest.kt @@ -0,0 +1,86 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.tensor.storage.Ownership +import sk.ainet.lang.tensor.storage.Placement +import sk.ainet.lang.tensor.storage.StorageSpec +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.BF16 +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.assertNull +import kotlin.test.assertTrue + +/** SKEEP-003 Phase 0: `AllocationSpec` replaces the never-consumed `StorageSpec`. */ +@OptIn(ExperimentalMemoryApi::class) +@Suppress("DEPRECATION") // StorageSpec.toAllocationSpec is the migration path under test +class AllocationSpecTest { + + @Test + fun bytesFollowTheEncoding() { + assertEquals(4L * 1000, AllocationSpec(Format.dense(FP32), 1000).bytes) + assertEquals(2L * 1000, AllocationSpec(Format.dense(BF16), 1000).bytes) + assertEquals(1000L, AllocationSpec(Format.dense(Int8), 1000).bytes) + // Q4_K: 144 bytes per 256 elements + assertEquals(144L * 4, AllocationSpec(Format(FP32, TensorEncoding.Q4_K), 1024).bytes) + assertEquals(144L, AllocationSpec(Format(FP32, TensorEncoding.Q4_K), 1).bytes) // partial block rounds up + // Q8_0: 34 bytes per 32 elements + assertEquals(34L * 2, AllocationSpec.of(Format(FP32, TensorEncoding.Q8_0), Shape(2, 32)).bytes) + // TurboQuant 4-bit, block 128: seed(4) + 4 groups × 2 B scales + 64 B codes = 76 per block + val tq = TensorEncoding.TurboQuantPolar(bitsPerElement = 4, blockSize = 128) + assertEquals(tq.physicalBytes(256), AllocationSpec(Format(FP32, tq), 256).bytes) + } + + @Test + fun opaqueEncodingHasNoComputableSize() { + val spec = AllocationSpec(Format(FP32, TensorEncoding.Opaque("IQ2_XXS", 0)), 64) + // Opaque carries its raw byte count; zero is "unknown" → physicalBytes may be null or 0 depending on the encoding + val b = spec.bytesOrNull + assertTrue(b == null || b == 0L) + } + + @Test + fun defaultsAreAmbientHeapMutableAligned64() { + val s = AllocationSpec(Format.dense(FP32), 8) + assertEquals(MemoryDomain.HOST_HEAP, s.domain) + assertEquals(ScopeKind.AMBIENT, s.scope) + assertTrue(s.mutable) + assertEquals(64, s.alignment) + assertFalse(s.format.isDense.not()) + } + + @Test + fun validation() { + assertFailsWith { AllocationSpec(Format.dense(FP32), -1) } + assertFailsWith { AllocationSpec(Format.dense(FP32), 1, alignment = 48) } + assertFailsWith { AllocationSpec(Format.dense(FP32), 1, alignment = 0) } + } + + @Test + fun storageSpecConvertsToAllocationSpec() { + val weights = StorageSpec.q4k(Placement.MMAP_WEIGHTS).toAllocationSpec(1024) + assertEquals(Format(FP32, TensorEncoding.Q4_K), weights.format) + assertEquals(1024L, weights.elementCount) + assertEquals(MemoryDomain.MMAP_FILE, weights.domain) + assertEquals(ScopeKind.MODEL, weights.scope) // persistent placement → model lifetime + assertFalse(weights.mutable) // borrowed packed bytes + + val owned = StorageSpec.fromDType(BF16).toAllocationSpec(10) + assertEquals(Format.dense(BF16), owned.format) + assertEquals(ScopeKind.AMBIENT, owned.scope) + assertTrue(owned.mutable) + assertEquals(20L, owned.bytes) + assertEquals(Ownership.OWNED, StorageSpec.fromDType(BF16).ownership) + } + + @Test + fun scopeKindHasTheThreeLifetimes() { + assertEquals(listOf(ScopeKind.MODEL, ScopeKind.FORWARD, ScopeKind.AMBIENT), ScopeKind.entries) + assertNull(ScopeKind.entries.firstOrNull { it.name == "DEVICE" }) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/LogicalDTypeBridgeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/LogicalDTypeBridgeTest.kt index 70b09b388..f03767e0a 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/LogicalDTypeBridgeTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/LogicalDTypeBridgeTest.kt @@ -25,6 +25,7 @@ import kotlin.test.assertSame * SKEEP-003 Phase 0, decision #13: the two-way `LogicalDType` <-> `DType` bridge must be total * and bijective so that `LogicalDType` can later be merged into `DType` without a semantic gap. */ +@Suppress("DEPRECATION") // StorageSpec is exercised on purpose: the bridge must keep the legacy descriptor working class LogicalDTypeBridgeTest { private val expectedPairs: List> = listOf(