From ede22289c6eda891e34372ff46bab34707a41035 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 22 Aug 2026 20:47:26 +0200 Subject: [PATCH] feat(types): DType carries its KClass witness; dtype-first storage constructors and loaders (SKEEP-003 P0) Decision #13 merges LogicalDType into one sealed DType. Second step (milestone M0, #1001): DType becomes self-describing and every storage descriptor can be built from it, additively. - DType: witness: KClass (each of the 14 objects narrows to its own class), isSigned (false for UInt*), sizeInBytes; companion entries (storage-layer order, index-aligned with LogicalDType.entries), fromWitness / fromWitnessOrNull (identity-keyed, no name-based reflection). Companion collections are lazy: with default members on the interface, initializing an object also initializes the interface statics on the JVM, and an eager map would capture a half-initialized object (class-init cycle). - Additive dtype-first API (primary constructors unchanged): TensorStorage(shape, dtype, encoding, buffer, ...), TensorStorageFactory.fromRawBytes / fromRawBytesOwned / fileBacked (dtype overloads), PackedBlockStorage.toTensorStorage(dtype, ...); toTensorData dispatches on storage.dtype with identical semantics. - Internal call sites moved to DType: TensorStorageFactory fromFloatArray/fromIntArray, DefaultKvCacheStore, TurboQuantKvCacheStore, StreamingGGUFReader (ggmlTypeToDType; quantized -> FP32, Opaque fallback kept), StreamingSafeTensorsReader (BOOL -> UInt8, UNKNOWN -> Int8 kept). - Tests: DTypeWitnessTest; dtype twins in TensorStorageContractTest; StorageIntegrationTest asserts storage.dtype === FP32 for F32 and Q8_0. - BCV: lang-core jvm dump regenerated (additions only). Closes #1007 Co-Authored-By: Claude Fable 5 --- .../sk/ainet/io/gguf/StreamingGGUFReader.kt | 25 ++--- .../ainet/io/gguf/StorageIntegrationTest.kt | 6 +- .../safetensors/StreamingSafeTensorsReader.kt | 35 +++---- .../api/jvm/skainet-lang-core.api | 74 +++++++++++++++ .../tensor/storage/DefaultKvCacheStore.kt | 3 +- .../lang/tensor/storage/PackedBlockStorage.kt | 6 ++ .../lang/tensor/storage/TensorStorage.kt | 20 +++- .../tensor/storage/TensorStorageFactory.kt | 46 ++++++++-- .../tensor/storage/TurboQuantKvCacheStore.kt | 3 +- .../kotlin/sk/ainet/lang/types/BF16.kt | 3 + .../kotlin/sk/ainet/lang/types/DType.kt | 68 +++++++++++++- .../kotlin/sk/ainet/lang/types/FP16.kt | 3 + .../kotlin/sk/ainet/lang/types/FP32.kt | 3 + .../kotlin/sk/ainet/lang/types/FP64.kt | 3 + .../kotlin/sk/ainet/lang/types/Int16.kt | 3 + .../kotlin/sk/ainet/lang/types/Int32.kt | 3 + .../kotlin/sk/ainet/lang/types/Int4.kt | 3 + .../kotlin/sk/ainet/lang/types/Int64.kt | 3 + .../kotlin/sk/ainet/lang/types/Int8.kt | 3 + .../kotlin/sk/ainet/lang/types/Ternary.kt | 3 + .../kotlin/sk/ainet/lang/types/UInt16.kt | 4 + .../kotlin/sk/ainet/lang/types/UInt32.kt | 4 + .../kotlin/sk/ainet/lang/types/UInt64.kt | 4 + .../kotlin/sk/ainet/lang/types/UInt8.kt | 4 + .../storage/TensorStorageContractTest.kt | 40 ++++++++ .../sk/ainet/lang/types/DTypeWitnessTest.kt | 91 +++++++++++++++++++ 26 files changed, 418 insertions(+), 45 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeWitnessTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt index 4af9fd3d3..948ce34b5 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt @@ -3,6 +3,7 @@ package sk.ainet.io.gguf import sk.ainet.io.RandomAccessSource import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.storage.* +import sk.ainet.lang.types.* /** * Streaming GGUF reader that parses metadata without loading the entire file. @@ -118,7 +119,7 @@ public class StreamingGGUFReader private constructor( val shape = Shape(*tensor.shape.map { it.toInt() }.toIntArray()) return TensorStorage( shape = shape, - logicalType = ggmlTypeToLogical(tensor.tensorType), + dtype = ggmlTypeToDType(tensor.tensorType), encoding = ggmlTypeToEncoding(tensor.tensorType, tensor.nBytes), buffer = BufferHandle.Borrowed(bytes, isMutable = false), placement = Placement.CPU_HEAP @@ -148,7 +149,7 @@ public class StreamingGGUFReader private constructor( val shape = Shape(*tensor.shape.map { it.toInt() }.toIntArray()) return TensorStorage( shape = shape, - logicalType = ggmlTypeToLogical(tensor.tensorType), + dtype = ggmlTypeToDType(tensor.tensorType), encoding = ggmlTypeToEncoding(tensor.tensorType, tensor.nBytes), buffer = BufferHandle.FileBacked( path = filePath, @@ -159,17 +160,17 @@ public class StreamingGGUFReader private constructor( ) } - private fun ggmlTypeToLogical(type: GGMLQuantizationType): LogicalDType = when (type) { - GGMLQuantizationType.F32 -> LogicalDType.FLOAT32 - GGMLQuantizationType.F16 -> LogicalDType.FLOAT16 - GGMLQuantizationType.BF16 -> LogicalDType.BFLOAT16 - GGMLQuantizationType.F64 -> LogicalDType.FLOAT64 - GGMLQuantizationType.I8 -> LogicalDType.INT8 - GGMLQuantizationType.I16 -> LogicalDType.INT16 - GGMLQuantizationType.I32 -> LogicalDType.INT32 - GGMLQuantizationType.I64 -> LogicalDType.INT64 + private fun ggmlTypeToDType(type: GGMLQuantizationType): DType = when (type) { + GGMLQuantizationType.F32 -> FP32 + GGMLQuantizationType.F16 -> FP16 + GGMLQuantizationType.BF16 -> BF16 + GGMLQuantizationType.F64 -> FP64 + GGMLQuantizationType.I8 -> Int8 + GGMLQuantizationType.I16 -> Int16 + GGMLQuantizationType.I32 -> Int32 + GGMLQuantizationType.I64 -> Int64 // Quantized types logically represent floats - else -> LogicalDType.FLOAT32 + else -> FP32 } private fun ggmlTypeToEncoding(type: GGMLQuantizationType, nBytes: Long): TensorEncoding = when (type) { diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StorageIntegrationTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StorageIntegrationTest.kt index 29cc46c2b..c7b34e2aa 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StorageIntegrationTest.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StorageIntegrationTest.kt @@ -3,12 +3,14 @@ package sk.ainet.io.gguf import org.junit.Test import sk.ainet.io.JvmFileBackedResolver import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.types.FP32 import sk.ainet.lang.tensor.storage.* import java.io.File import java.io.RandomAccessFile import java.nio.ByteBuffer import java.nio.ByteOrder import kotlin.test.assertEquals +import kotlin.test.assertSame import kotlin.test.assertFalse import kotlin.test.assertTrue @@ -106,6 +108,7 @@ class StorageIntegrationTest { // F32 tensor val f32Storage = reader.loadTensorStorage("weight_f32") assertEquals(LogicalDType.FLOAT32, f32Storage.logicalType) + assertSame(FP32, f32Storage.dtype) assertEquals(TensorEncoding.Dense(4), f32Storage.encoding) assertEquals(Ownership.BORROWED, f32Storage.ownership) assertEquals(16L, f32Storage.physicalBytes) @@ -114,7 +117,8 @@ class StorageIntegrationTest { // Q8_0 tensor val q80Storage = reader.loadTensorStorage("weight_q80") - assertEquals(LogicalDType.FLOAT32, q80Storage.logicalType) + assertEquals(LogicalDType.FLOAT32, q80Storage.logicalType) // packed weights are logically FP32 + assertSame(FP32, q80Storage.dtype) assertEquals(TensorEncoding.Q8_0, q80Storage.encoding) assertEquals(Ownership.BORROWED, q80Storage.ownership) assertEquals(34L, q80Storage.physicalBytes) diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingSafeTensorsReader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingSafeTensorsReader.kt index 99246a0ed..d2b2b08de 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingSafeTensorsReader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingSafeTensorsReader.kt @@ -4,6 +4,7 @@ import sk.ainet.io.RandomAccessSource import sk.ainet.io.model.DataType import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.storage.* +import sk.ainet.lang.types.* /** * Streaming SafeTensors reader that parses metadata without loading tensor data. @@ -106,7 +107,7 @@ public class StreamingSafeTensorsReader private constructor( val shape = Shape(*tensor.shape.map { it.toInt() }.toIntArray()) return TensorStorage( shape = shape, - logicalType = safeTensorsTypeToLogical(tensor.dataType), + dtype = safeTensorsTypeToDType(tensor.dataType), encoding = safeTensorsTypeToEncoding(tensor.dataType), buffer = BufferHandle.Borrowed(bytes, isMutable = false), placement = Placement.CPU_HEAP @@ -133,7 +134,7 @@ public class StreamingSafeTensorsReader private constructor( val shape = Shape(*tensor.shape.map { it.toInt() }.toIntArray()) return TensorStorage( shape = shape, - logicalType = safeTensorsTypeToLogical(tensor.dataType), + dtype = safeTensorsTypeToDType(tensor.dataType), encoding = safeTensorsTypeToEncoding(tensor.dataType), buffer = BufferHandle.FileBacked( path = filePath, @@ -144,21 +145,21 @@ public class StreamingSafeTensorsReader private constructor( ) } - private fun safeTensorsTypeToLogical(type: DataType): LogicalDType = when (type) { - DataType.FLOAT32 -> LogicalDType.FLOAT32 - DataType.FLOAT64 -> LogicalDType.FLOAT64 - DataType.FLOAT16 -> LogicalDType.FLOAT16 - DataType.BFLOAT16 -> LogicalDType.BFLOAT16 - DataType.INT8 -> LogicalDType.INT8 - DataType.INT16 -> LogicalDType.INT16 - DataType.INT32 -> LogicalDType.INT32 - DataType.INT64 -> LogicalDType.INT64 - DataType.UINT8 -> LogicalDType.UINT8 - DataType.UINT16 -> LogicalDType.UINT16 - DataType.UINT32 -> LogicalDType.UINT32 - DataType.UINT64 -> LogicalDType.UINT64 - DataType.BOOL -> LogicalDType.UINT8 - else -> LogicalDType.INT8 // fallback for UNKNOWN + private fun safeTensorsTypeToDType(type: DataType): DType = when (type) { + DataType.FLOAT32 -> FP32 + DataType.FLOAT64 -> FP64 + DataType.FLOAT16 -> FP16 + DataType.BFLOAT16 -> BF16 + DataType.INT8 -> Int8 + DataType.INT16 -> Int16 + DataType.INT32 -> Int32 + DataType.INT64 -> Int64 + DataType.UINT8 -> UInt8 + DataType.UINT16 -> UInt16 + DataType.UINT32 -> UInt32 + DataType.UINT64 -> UInt64 + DataType.BOOL -> UInt8 + else -> Int8 // fallback for UNKNOWN } private fun safeTensorsTypeToEncoding(type: DataType): TensorEncoding = when (type) { 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 a04ca4d63..54a621657 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 @@ -3202,6 +3202,7 @@ public final class sk/ainet/lang/tensor/data/Q4_0BlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q4_0BlockTensorData$Companion { @@ -3260,6 +3261,7 @@ public final class sk/ainet/lang/tensor/data/Q4_KBlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q4_KBlockTensorData$Companion { @@ -3314,6 +3316,7 @@ public final class sk/ainet/lang/tensor/data/Q5_0BlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q5_0BlockTensorData$Companion { @@ -3355,6 +3358,7 @@ public final class sk/ainet/lang/tensor/data/Q5_1BlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q5_1BlockTensorData$Companion { @@ -3401,6 +3405,7 @@ public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData$Companion { @@ -3462,6 +3467,7 @@ public final class sk/ainet/lang/tensor/data/Q6_KBlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q6_KBlockTensorData$Companion { @@ -3544,6 +3550,7 @@ public final class sk/ainet/lang/tensor/data/Q8_0BlockTensorData : sk/ainet/lang public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Q8_0BlockTensorData$Companion { @@ -3637,6 +3644,7 @@ public final class sk/ainet/lang/tensor/data/Ternary2BitTensorData : sk/ainet/la public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/data/Ternary2BitTensorData$Companion { @@ -5356,7 +5364,9 @@ public abstract interface class sk/ainet/lang/tensor/storage/PackedBlockStorage public abstract fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun toFloatArray ()[F public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun toTensorStorage$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public static synthetic fun toTensorStorage$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public final class sk/ainet/lang/tensor/storage/PackedBlockStorage$DefaultImpls { @@ -5365,7 +5375,9 @@ public final class sk/ainet/lang/tensor/storage/PackedBlockStorage$DefaultImpls public static fun getPhysicalBytes (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)J public static fun toFloatArray (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)[F public static fun toTensorStorage (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public static fun toTensorStorage (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun toTensorStorage$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public static synthetic fun toTensorStorage$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; } public abstract interface annotation class sk/ainet/lang/tensor/storage/Place : java/lang/annotation/Annotation { @@ -5655,6 +5667,8 @@ public final class sk/ainet/lang/tensor/storage/TensorEncoding$TurboQuantPolarQj public final class sk/ainet/lang/tensor/storage/TensorStorage { public fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferHandle;Lsk/ainet/lang/tensor/storage/Placement;J[JZ)V public synthetic fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferHandle;Lsk/ainet/lang/tensor/storage/Placement;J[JZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferHandle;Lsk/ainet/lang/tensor/storage/Placement;J[JZ)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/tensor/storage/BufferHandle;Lsk/ainet/lang/tensor/storage/Placement;J[JZILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Lsk/ainet/lang/tensor/Shape; public final fun component2 ()Lsk/ainet/lang/tensor/storage/LogicalDType; public final fun component3 ()Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -5697,12 +5711,17 @@ public final class sk/ainet/lang/tensor/storage/TensorStorageFactory { public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorStorageFactory; public final fun borrowFloatArray (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fileBacked (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Ljava/lang/String;JJ)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun fileBacked (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;Ljava/lang/String;JJ)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fromIntArray (Lsk/ainet/lang/tensor/Shape;[I)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun fromRawBytes$default (Lsk/ainet/lang/tensor/storage/TensorStorageFactory;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public static synthetic fun fromRawBytes$default (Lsk/ainet/lang/tensor/storage/TensorStorageFactory;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fromRawBytesOwned (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public final fun fromRawBytesOwned (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; public static synthetic fun fromRawBytesOwned$default (Lsk/ainet/lang/tensor/storage/TensorStorageFactory;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public static synthetic fun fromRawBytesOwned$default (Lsk/ainet/lang/tensor/storage/TensorStorageFactory;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BLsk/ainet/lang/tensor/storage/Placement;ILjava/lang/Object;)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun fromTensorData (Lsk/ainet/lang/tensor/data/TensorData;)Lsk/ainet/lang/tensor/storage/TensorStorage; public final fun toTensorData (Lsk/ainet/lang/tensor/storage/TensorStorage;)Lsk/ainet/lang/tensor/data/TensorData; } @@ -5813,7 +5832,10 @@ public final class sk/ainet/lang/types/BF16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/BF16; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5831,14 +5853,19 @@ public abstract interface class sk/ainet/lang/types/DType { public static fun fp16 ()Lsk/ainet/lang/types/DType; public static fun fp32 ()Lsk/ainet/lang/types/DType; public static fun fp64 ()Lsk/ainet/lang/types/DType; + public static fun fromWitness (Lkotlin/reflect/KClass;)Lsk/ainet/lang/types/DType; + public static fun fromWitnessOrNull (Lkotlin/reflect/KClass;)Lsk/ainet/lang/types/DType; public abstract fun getName ()Ljava/lang/String; public abstract fun getSizeInBits ()I + public fun getSizeInBytes ()I + public abstract fun getWitness ()Lkotlin/reflect/KClass; public static fun int16 ()Lsk/ainet/lang/types/DType; public static fun int32 ()Lsk/ainet/lang/types/DType; public static fun int4 ()Lsk/ainet/lang/types/DType; public static fun int64 ()Lsk/ainet/lang/types/DType; public static fun int8 ()Lsk/ainet/lang/types/DType; public abstract fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public abstract fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; public static fun ternary ()Lsk/ainet/lang/types/DType; public static fun uint16 ()Lsk/ainet/lang/types/DType; @@ -5853,7 +5880,10 @@ public final class sk/ainet/lang/types/DType$Companion { public final fun fp16 ()Lsk/ainet/lang/types/DType; public final fun fp32 ()Lsk/ainet/lang/types/DType; public final fun fp64 ()Lsk/ainet/lang/types/DType; + public final fun fromWitness (Lkotlin/reflect/KClass;)Lsk/ainet/lang/types/DType; + public final fun fromWitnessOrNull (Lkotlin/reflect/KClass;)Lsk/ainet/lang/types/DType; public final fun getAllTypes ()Ljava/util/Map; + public final fun getEntries ()Ljava/util/List; public final fun int16 ()Lsk/ainet/lang/types/DType; public final fun int32 ()Lsk/ainet/lang/types/DType; public final fun int4 ()Lsk/ainet/lang/types/DType; @@ -5866,6 +5896,11 @@ public final class sk/ainet/lang/types/DType$Companion { public final fun uint8 ()Lsk/ainet/lang/types/DType; } +public final class sk/ainet/lang/types/DType$DefaultImpls { + public static fun getSizeInBytes (Lsk/ainet/lang/types/DType;)I + public static fun isSigned (Lsk/ainet/lang/types/DType;)Z +} + public final class sk/ainet/lang/types/DTypeExtensionsKt { public static final fun commonPrecisionWith (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; public static final fun isConvertibleTo (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Z @@ -5937,7 +5972,10 @@ public final class sk/ainet/lang/types/FP16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/FP16; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5945,7 +5983,10 @@ public final class sk/ainet/lang/types/FP32 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/FP32; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5953,7 +5994,10 @@ public final class sk/ainet/lang/types/FP64 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/FP64; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5969,7 +6013,10 @@ public final class sk/ainet/lang/types/Int16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int16; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5977,7 +6024,10 @@ public final class sk/ainet/lang/types/Int32 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int32; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5985,7 +6035,10 @@ public final class sk/ainet/lang/types/Int4 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int4; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -5993,7 +6046,10 @@ public final class sk/ainet/lang/types/Int64 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int64; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6001,7 +6057,10 @@ public final class sk/ainet/lang/types/Int8 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int8; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6028,7 +6087,10 @@ public final class sk/ainet/lang/types/Ternary : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Ternary; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6056,7 +6118,10 @@ public final class sk/ainet/lang/types/UInt16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/UInt16; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6064,7 +6129,10 @@ public final class sk/ainet/lang/types/UInt32 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/UInt32; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6072,7 +6140,10 @@ public final class sk/ainet/lang/types/UInt64 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/UInt64; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } @@ -6080,7 +6151,10 @@ public final class sk/ainet/lang/types/UInt8 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/UInt8; public fun getName ()Ljava/lang/String; public fun getSizeInBits ()I + public fun getSizeInBytes ()I + public fun getWitness ()Lkotlin/reflect/KClass; public fun isCompatible (Lsk/ainet/lang/types/DType;)Z + public fun isSigned ()Z public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt index dfd3f42f8..8c145bf45 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/DefaultKvCacheStore.kt @@ -1,6 +1,7 @@ package sk.ainet.lang.tensor.storage import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 /** * Default KV cache implementation using dense FP32 storage. @@ -169,7 +170,7 @@ public class DefaultKvCacheStore( } return TensorStorage( shape = Shape(numHeads, seqLen, headDim), - logicalType = LogicalDType.FLOAT32, + dtype = FP32, encoding = encoding, buffer = BufferHandle.Owned(bytes), placement = placement diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt index b8b37fc95..1d70f9c21 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt @@ -73,4 +73,10 @@ public interface PackedBlockStorage { buffer = BufferHandle.Borrowed(packedData, isMutable = false), placement = placement ) + + /** Convert this packed storage to a [TensorStorage] descriptor, dtype-first (packed weights are logically [FP32]). */ + public fun toTensorStorage( + dtype: sk.ainet.lang.types.DType, + placement: Placement = Placement.CPU_HEAP + ): TensorStorage = toTensorStorage(dtype.toLogicalDType(), placement) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt index 4572d3a37..4affe5816 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt @@ -1,6 +1,7 @@ package sk.ainet.lang.tensor.storage import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType /** * Runtime descriptor for a tensor's backing memory. @@ -29,8 +30,23 @@ public data class TensorStorage( val strides: LongArray? = null, val isContiguous: Boolean = true ) { - /** The [sk.ainet.lang.types.DType] of [logicalType] (SKEEP-003 Phase 0 bridge). */ - val dtype: sk.ainet.lang.types.DType get() = logicalType.toDType() + /** + * Construct from a [DType] (SKEEP-003 Phase 0): the dtype-first form every new call site uses; + * the [LogicalDType] primary constructor stays for source compatibility until the next major. + */ + public constructor( + shape: Shape, + dtype: DType, + encoding: TensorEncoding, + buffer: BufferHandle, + placement: Placement = Placement.CPU_HEAP, + byteOffset: Long = 0, + strides: LongArray? = null, + isContiguous: Boolean = true, + ) : this(shape, dtype.toLogicalDType(), encoding, buffer, placement, byteOffset, strides, isContiguous) + + /** The [DType] of this storage — what the values mean (SKEEP-003 Phase 0 bridge over [logicalType]). */ + val dtype: DType get() = logicalType.toDType() /** Number of logical elements in this tensor. */ val elementCount: Long get() = shape.volume.toLong() diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt index b1c1d49ba..e0f1dec14 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt @@ -12,6 +12,8 @@ import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.FP32 import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.types.Int32 @@ -29,7 +31,7 @@ public object TensorStorageFactory { public fun fromFloatArray(shape: Shape, data: FloatArray): TensorStorage = TensorStorage( shape = shape, - logicalType = LogicalDType.FLOAT32, + dtype = FP32, encoding = TensorEncoding.Dense(bytesPerElement = 4), buffer = BufferHandleFactory.owned(data) ) @@ -61,7 +63,7 @@ public object TensorStorageFactory { public fun fromIntArray(shape: Shape, data: IntArray): TensorStorage = TensorStorage( shape = shape, - logicalType = LogicalDType.INT32, + dtype = Int32, encoding = TensorEncoding.Dense(bytesPerElement = 4), buffer = BufferHandleFactory.owned(data) ) @@ -84,6 +86,15 @@ public object TensorStorageFactory { placement = placement ) + /** Create storage from raw bytes with explicit encoding, dtype-first (the byte array is borrowed). */ + public fun fromRawBytes( + shape: Shape, + dtype: DType, + encoding: TensorEncoding, + data: ByteArray, + placement: Placement = Placement.CPU_HEAP + ): TensorStorage = fromRawBytes(shape, dtype.toLogicalDType(), encoding, data, placement) + /** * Create storage from raw bytes with explicit encoding (owned copy). */ @@ -101,6 +112,15 @@ public object TensorStorageFactory { placement = placement ) + /** Create storage from raw bytes with explicit encoding (owned copy), dtype-first. */ + public fun fromRawBytesOwned( + shape: Shape, + dtype: DType, + encoding: TensorEncoding, + data: ByteArray, + placement: Placement = Placement.CPU_HEAP + ): TensorStorage = fromRawBytesOwned(shape, dtype.toLogicalDType(), encoding, data, placement) + /** * Create file-backed storage (for memory-mapped model weights). */ @@ -119,6 +139,16 @@ public object TensorStorageFactory { placement = Placement.MMAP_WEIGHTS ) + /** Create file-backed storage (for memory-mapped model weights), dtype-first. */ + public fun fileBacked( + shape: Shape, + dtype: DType, + encoding: TensorEncoding, + path: String, + fileOffset: Long, + sizeInBytes: Long + ): TensorStorage = fileBacked(shape, dtype.toLogicalDType(), encoding, path, fileOffset, sizeInBytes) + /** * Bridge: create a [TensorStorage] descriptor from an existing [TensorData]. * @@ -189,8 +219,8 @@ public object TensorStorageFactory { val bytes = extractBytes(storage) return when (storage.encoding) { - is TensorEncoding.Dense -> when (storage.logicalType) { - LogicalDType.FLOAT32 -> { + is TensorEncoding.Dense -> when (storage.dtype) { + FP32 -> { val floats = bytesToFloatArray(bytes) DenseFloatArrayTensorData(storage.shape, floats) as TensorData } @@ -201,17 +231,17 @@ public object TensorStorageFactory { // ignored it. Widening to f32 here keeps this method's existing contract (it has // always returned float-backed data); preserving 2-byte storage end-to-end is the // loader `KEEP_NATIVE` path, not this one. - LogicalDType.FLOAT16, LogicalDType.BFLOAT16 -> { - val codec = if (storage.logicalType == LogicalDType.FLOAT16) Fp16Codec else Bf16Codec + FP16, BF16 -> { + val codec = if (storage.dtype == FP16) Fp16Codec else Bf16Codec val floats = narrowBytesToFloatArray(bytes, codec) DenseFloatArrayTensorData(storage.shape, floats) as TensorData } - LogicalDType.INT32 -> { + Int32 -> { val ints = bytesToIntArray(bytes) DenseIntArrayTensorData(storage.shape, ints) as TensorData } else -> throw UnsupportedOperationException( - "toTensorData not supported for dense ${storage.logicalType}" + "toTensorData not supported for dense ${storage.dtype.name}" ) } is TensorEncoding.Q4_K -> { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt index 048658f89..972d49901 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TurboQuantKvCacheStore.kt @@ -1,6 +1,7 @@ package sk.ainet.lang.tensor.storage import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 import sk.ainet.lang.tensor.ops.turboquant.TurboQuantBlock import sk.ainet.lang.tensor.ops.turboquant.TurboQuantCodec import sk.ainet.lang.tensor.ops.turboquant.TurboQuantConfig @@ -187,7 +188,7 @@ public class TurboQuantKvCacheStore( } return TensorStorage( shape = Shape(numHeads, seqLen, headDim), - logicalType = LogicalDType.FLOAT32, + dtype = FP32, encoding = encoding, buffer = BufferHandle.Owned(bytes), placement = placement diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/BF16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/BF16.kt index d2f8a8083..945345e1c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/BF16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/BF16.kt @@ -1,11 +1,14 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * Brain Float 16 (BFloat16) type. * 16-bit floating point format with same exponent range as FP32 but reduced mantissa. * Commonly used in machine learning workloads. */ public object BF16 : DType { + override val witness: KClass get() = BF16::class override val sizeInBits: Int = 16 override val name: String = "BFloat16" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DType.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DType.kt index 8f125e228..93ed4de88 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DType.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DType.kt @@ -1,10 +1,35 @@ package sk.ainet.lang.types -// Base marker interface for all dtypes +import kotlin.jvm.JvmStatic +import kotlin.reflect.KClass + +/** + * The logical element type of a tensor: what a value *means* when you read it. + * + * `DType` is sealed with exactly fourteen `object` members, so it is switchable like an enum + * (`when (dtype) { FP32 -> … }`) while each member still serves as the generic witness of + * `Tensor` through [witness]. SKEEP-003 (decision #13): this single type replaces + * the storage-side `LogicalDType` enum; `Format = (DType, TensorEncoding)` has exactly one dtype + * type. + */ public sealed interface DType { public val sizeInBits: Int public val name: String + /** + * The `KClass` witness of this dtype — the value `Tensor.dtype` / `TensorData` carry as + * their type argument. Each member returns its own class (`FP32.witness == FP32::class`), so + * [DType.fromWitness] maps a `Tensor.dtype` back to the `DType` object. (Not to be confused with + * [kotlinClass], which is the KClass of the *value* representation, e.g. `Float::class`.) + */ + public val witness: KClass + + /** Whether the type carries a sign (false only for the unsigned integer types). */ + public val isSigned: Boolean get() = true + + /** Storage width of one element rounded up to whole bytes. */ + public val sizeInBytes: Int get() = (sizeInBits + 7) / 8 + /** * Checks if this data type is compatible with another data type for operations. * @@ -33,10 +58,14 @@ public sealed interface DType { public fun promoteTo(other: DType): DType public companion object { + // All companion collections are lazy on purpose: `DType` now has default members, so on + // the JVM initializing any `object` (e.g. FP32) also initializes this interface's statics; + // an eager map would capture the half-initialized object (class-init cycle → + // ExceptionInInitializerError). Lazy delegates defer the member references until first use. /** * Registry of all available data types. */ - private val typeRegistry: Map = mapOf( + private val typeRegistry: Map by lazy { mapOf( "Ternary" to Ternary, "Int4" to Int4, "Int8" to Int8, @@ -51,7 +80,7 @@ public sealed interface DType { "BFloat16" to BF16, "Float32" to FP32, "Float64" to FP64 - ) + ) } /** * Gets all registered data types. @@ -60,6 +89,39 @@ public sealed interface DType { */ public fun getAllTypes(): Map = typeRegistry + /** + * All fourteen dtypes in a stable order (the storage-layer order: ternary, signed ints, + * unsigned ints, floats) — the enum-like `entries` of this sealed type. + */ + public val entries: List by lazy { + listOf( + Ternary, Int4, Int8, Int16, Int32, Int64, + UInt8, UInt16, UInt32, UInt64, + FP16, BF16, FP32, FP64, + ) + } + + // Identity-keyed (KClass equality), never name-based: stable on JS/Wasm where class + // names may be minified. + private val byWitness: Map, DType> by lazy { entries.associateBy { it.witness } } + + /** + * The [DType] whose [witness] is [kclass], or `null` if [kclass] is not one of the + * fourteen dtype classes (e.g. `DType::class` itself). + */ + @JvmStatic + public fun fromWitnessOrNull(kclass: KClass): DType? = byWitness[kclass] + + /** + * The [DType] whose [witness] is [kclass] — the inverse of [witness], e.g. + * `DType.fromWitness(tensor.dtype)`. + * + * @throws IllegalArgumentException if [kclass] is not a dtype class + */ + @JvmStatic + public fun fromWitness(kclass: KClass): DType = + byWitness[kclass] ?: throw IllegalArgumentException("Not a DType witness: $kclass") + /** * Finds a data type by name. * diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt index e72d2ca12..18356d873 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt @@ -1,6 +1,9 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + public object FP16 : DType { + override val witness: KClass get() = FP16::class override val sizeInBits: Int = 16 override val name: String = "Float16" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt index 0e170e1ed..9294c251b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt @@ -1,6 +1,9 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + public object FP32 : DType { + override val witness: KClass get() = FP32::class override val sizeInBits: Int = 32 override val name: String = "Float32" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP64.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP64.kt index 861547de2..eef550793 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP64.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP64.kt @@ -1,10 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 64-bit floating point type (Double precision). * Highest precision floating point type in SKaiNET. */ public object FP64 : DType { + override val witness: KClass get() = FP64::class override val sizeInBits: Int = 64 override val name: String = "Float64" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int16.kt index 6783ccc9d..91e28a707 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int16.kt @@ -1,9 +1,12 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 16-bit signed integer type. */ public object Int16 : DType { + override val witness: KClass get() = Int16::class override val sizeInBits: Int = 16 override val name: String = "Int16" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int32.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int32.kt index d4dfd7e4a..de0512d99 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int32.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int32.kt @@ -1,6 +1,9 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + public object Int32 : DType { + override val witness: KClass get() = Int32::class override val sizeInBits: Int = 32 override val name: String = "Int32" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int4.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int4.kt index 518ff1cda..c1917a3b3 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int4.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int4.kt @@ -1,6 +1,9 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + public object Int4 : DType { + override val witness: KClass get() = Int4::class override val sizeInBits: Int = 4 override val name: String = "Int4" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int64.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int64.kt index 6c266e957..1186e7ac6 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int64.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int64.kt @@ -1,9 +1,12 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 64-bit signed integer type (Long). */ public object Int64 : DType { + override val witness: KClass get() = Int64::class override val sizeInBits: Int = 64 override val name: String = "Int64" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt index e4254607f..d44a2f43e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt @@ -1,6 +1,9 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + public object Int8 : DType { + override val witness: KClass get() = Int8::class override val sizeInBits: Int = 8 override val name: String = "Int8" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Ternary.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Ternary.kt index bddbb7436..eaf66e310 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Ternary.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Ternary.kt @@ -1,10 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * Custom data type holding 3 values -1,0,1 stored in 2 bits. Used e.g. with BitNet models * https://huggingface.co/microsoft/bitnet-b1.58-2B-4T */ public object Ternary : DType { + override val witness: KClass get() = Ternary::class override val sizeInBits: Int = 2 override val name: String = "Ternary" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt16.kt index 51d477126..5ec9ebfd8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt16.kt @@ -1,9 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 16-bit unsigned integer type. */ public object UInt16 : DType { + override val witness: KClass get() = UInt16::class + override val isSigned: Boolean get() = false override val sizeInBits: Int = 16 override val name: String = "UInt16" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt32.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt32.kt index 031ec808e..6b9a4fb67 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt32.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt32.kt @@ -1,9 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 32-bit unsigned integer type. */ public object UInt32 : DType { + override val witness: KClass get() = UInt32::class + override val isSigned: Boolean get() = false override val sizeInBits: Int = 32 override val name: String = "UInt32" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt64.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt64.kt index 6d3288fd6..c948136f8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt64.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt64.kt @@ -1,9 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 64-bit unsigned integer type. */ public object UInt64 : DType { + override val witness: KClass get() = UInt64::class + override val isSigned: Boolean get() = false override val sizeInBits: Int = 64 override val name: String = "UInt64" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt8.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt8.kt index 9996ef23f..3d09392a0 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt8.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/UInt8.kt @@ -1,9 +1,13 @@ package sk.ainet.lang.types +import kotlin.reflect.KClass + /** * 8-bit unsigned integer type. */ public object UInt8 : DType { + override val witness: KClass get() = UInt8::class + override val isSigned: Boolean get() = false override val sizeInBits: Int = 8 override val name: String = "UInt8" diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TensorStorageContractTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TensorStorageContractTest.kt index fc02605b3..d5b9df3d0 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TensorStorageContractTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TensorStorageContractTest.kt @@ -219,4 +219,44 @@ class TensorStorageContractTest { assertFalse(report.isFileBacked) assertFalse(report.isMutable) } + + // --- dtype-first constructors (SKEEP-003 Phase 0) --- + + @Test + fun dtypeConstructorEqualsLogicalTypeConstructor() { + val shape = Shape(2, 3) + val buffer = BufferHandle.Borrowed(ByteArray(24)) // one instance: BufferHandle equality is identity + val viaLogical = TensorStorage(shape, LogicalDType.FLOAT32, TensorEncoding.Dense(4), buffer) + val viaDType = TensorStorage(shape, FP32, TensorEncoding.Dense(4), buffer) + assertEquals(viaLogical, viaDType) + assertEquals(FP32, viaDType.dtype) + assertEquals(LogicalDType.FLOAT32, viaDType.logicalType) + assertEquals(viaLogical.memoryReport(), viaDType.memoryReport()) + + val bf16 = TensorStorage(shape, BF16, TensorEncoding.Dense(2), BufferHandle.Borrowed(ByteArray(12)), Placement.CPU_HEAP) + assertEquals(LogicalDType.BFLOAT16, bf16.logicalType) + assertEquals(12L, bf16.logicalBytes) + } + + @Test + fun dtypeFactoryOverloadsMatchLogicalTypeOverloads() { + val shape = Shape(4) + val bytes = ByteArray(16) + // BufferHandle subclasses have identity equality, so compare the descriptor fields. + fun sig(s: TensorStorage) = listOf(s.shape, s.logicalType, s.dtype, s.encoding, s.ownership, s.placement, s.physicalBytes) + assertEquals( + sig(TensorStorageFactory.fromRawBytes(shape, LogicalDType.INT32, TensorEncoding.Dense(4), bytes)), + sig(TensorStorageFactory.fromRawBytes(shape, Int32, TensorEncoding.Dense(4), bytes)), + ) + assertEquals( + sig(TensorStorageFactory.fromRawBytesOwned(shape, LogicalDType.FLOAT16, TensorEncoding.Dense(2), ByteArray(8))), + sig(TensorStorageFactory.fromRawBytesOwned(shape, FP16, TensorEncoding.Dense(2), ByteArray(8))), + ) + assertEquals( + sig(TensorStorageFactory.fileBacked(shape, LogicalDType.FLOAT32, TensorEncoding.Dense(4), "/m.gguf", 128L, 16L)), + sig(TensorStorageFactory.fileBacked(shape, FP32, TensorEncoding.Dense(4), "/m.gguf", 128L, 16L)), + ) + assertEquals(FP32, TensorStorageFactory.fromFloatArray(shape, FloatArray(4)).dtype) + assertEquals(Int32, TensorStorageFactory.fromIntArray(shape, IntArray(4)).dtype) + } } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeWitnessTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeWitnessTest.kt new file mode 100644 index 000000000..897416e20 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeWitnessTest.kt @@ -0,0 +1,91 @@ +package sk.ainet.lang.types + +import sk.ainet.lang.tensor.storage.LogicalDType +import sk.ainet.lang.tensor.storage.toLogicalDType +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertNull +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * SKEEP-003 Phase 0, decision #13: `DType` carries its own `KClass` witness so the sealed objects + * serve both as enum-like descriptors and as the `Tensor` type argument. + */ +class DTypeWitnessTest { + + @Test + fun entriesAreTheFourteenRegisteredTypesInStorageOrder() { + assertEquals(14, DType.entries.size) + assertEquals(DType.getAllTypes().values.toSet(), DType.entries.toSet()) + assertEquals( + listOf(Ternary, Int4, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, FP16, BF16, FP32, FP64), + DType.entries, + ) + // same order as LogicalDType, so the two enumerations line up index by index + for ((i, logical) in LogicalDType.entries.withIndex()) { + assertSame(DType.entries[i], logical.toDType(), "entries[$i] vs $logical") + } + } + + @Test + fun eachObjectWitnessesItself() { + // KClass equality (not identity): `X::class` may create a fresh reference per evaluation. + assertEquals(FP32::class, FP32.witness) + assertEquals(FP16::class, FP16.witness) + assertEquals(BF16::class, BF16.witness) + assertEquals(FP64::class, FP64.witness) + assertEquals(Int4::class, Int4.witness) + assertEquals(Int8::class, Int8.witness) + assertEquals(Int16::class, Int16.witness) + assertEquals(Int32::class, Int32.witness) + assertEquals(Int64::class, Int64.witness) + assertEquals(UInt8::class, UInt8.witness) + assertEquals(UInt16::class, UInt16.witness) + assertEquals(UInt32::class, UInt32.witness) + assertEquals(UInt64::class, UInt64.witness) + assertEquals(Ternary::class, Ternary.witness) + } + + @Test + fun fromWitnessIsTheInverseOfWitness() { + for (d in DType.entries) { + assertSame(d, DType.fromWitness(d.witness), "fromWitness(${d.name}.witness)") + assertSame(d, DType.fromWitnessOrNull(d.witness)) + } + assertSame(FP32, DType.fromWitness(FP32::class)) + // all witnesses are distinct + assertEquals(14, DType.entries.map { it.witness }.toSet().size) + } + + @Test + fun fromWitnessRejectsNonDtypeClasses() { + assertNull(DType.fromWitnessOrNull(DType::class)) + assertFailsWith { DType.fromWitness(DType::class) } + } + + @Test + fun signednessAndByteWidthMatchTheStorageEnum() { + for (d in DType.entries) { + val logical = d.toLogicalDType() + assertEquals(logical.isSigned, d.isSigned, "isSigned of ${d.name}") + assertEquals(logical.sizeInBytes, d.sizeInBytes, "sizeInBytes of ${d.name}") + } + assertFalse(UInt8.isSigned); assertFalse(UInt16.isSigned); assertFalse(UInt32.isSigned); assertFalse(UInt64.isSigned) + assertTrue(Int8.isSigned); assertTrue(FP32.isSigned); assertTrue(Ternary.isSigned) + assertEquals(1, Int4.sizeInBytes); assertEquals(1, Ternary.sizeInBytes); assertEquals(2, BF16.sizeInBytes); assertEquals(8, FP64.sizeInBytes) + } + + @Test + fun switchableLikeAnEnum() { + fun label(d: DType): String = when (d) { + FP32, FP16, BF16, FP64 -> "float" + Int4, Int8, Int16, Int32, Int64 -> "int" + UInt8, UInt16, UInt32, UInt64 -> "uint" + Ternary -> "ternary" + } + assertEquals("float", label(BF16)); assertEquals("uint", label(UInt64)); assertEquals("ternary", label(Ternary)) + } +}