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 8debc1862..3a5f25105 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 @@ -854,6 +854,14 @@ public final class sk/ainet/lang/memory/ModelScope : sk/ainet/lang/memory/Scope public static synthetic fun mapFile$default (Lsk/ainet/lang/memory/ModelScope;Ljava/lang/String;JJLsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage; } +public final class sk/ainet/lang/memory/NarrowFloatDecoder : sk/ainet/lang/memory/BlockDecoder { + public fun (Lsk/ainet/lang/types/NarrowFloatCodec;)V + public fun decodeBlock (Lsk/ainet/lang/memory/Storage;J[FI)V + public fun decodeElement (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/memory/Layout;J)F + public fun getBlockSize ()I + public fun getBytesPerBlock ()I +} + public abstract interface class sk/ainet/lang/memory/Owner { } @@ -3861,6 +3869,7 @@ public final class sk/ainet/lang/tensor/data/Bf16TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Bf16TensorData;)[F public static fun getCodec (Lsk/ainet/lang/tensor/data/Bf16TensorData;)Lsk/ainet/lang/types/NarrowFloatCodec; public static fun getEncoding (Lsk/ainet/lang/tensor/data/Bf16TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Bf16TensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Bf16TensorDataKt { @@ -3875,6 +3884,7 @@ public final class sk/ainet/lang/tensor/data/DenseFloatArrayTensorData : sk/aine public fun getBuffer ()[F public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -3887,6 +3897,7 @@ public final class sk/ainet/lang/tensor/data/DenseIntArrayTensorData : sk/ainet/ public fun getBuffer ()[I public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([II)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -3923,11 +3934,13 @@ public final class sk/ainet/lang/tensor/data/DenseTensorDataFactory : sk/ainet/l public abstract interface class sk/ainet/lang/tensor/data/FloatArrayTensorData : sk/ainet/lang/tensor/data/TensorData { public fun copyToFloatArray ()[F public abstract fun getBuffer ()[F + public fun getView ()Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/FloatArrayTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/FloatArrayTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/FloatArrayTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/FloatArrayTensorData;)Lsk/ainet/lang/memory/TensorView; } public abstract interface class sk/ainet/lang/tensor/data/FloatBufferTensorData : sk/ainet/lang/tensor/data/TensorData { @@ -3937,6 +3950,7 @@ public abstract interface class sk/ainet/lang/tensor/data/FloatBufferTensorData public final class sk/ainet/lang/tensor/data/FloatBufferTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/FloatBufferTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/FloatBufferTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/FloatBufferTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData { @@ -3951,11 +3965,13 @@ public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData$Companion { public abstract interface class sk/ainet/lang/tensor/data/IntArrayTensorData : sk/ainet/lang/tensor/data/TensorData { public abstract fun getBuffer ()[I + public fun getView ()Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/IntArrayTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/IntArrayTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/IntArrayTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/IntArrayTensorData;)Lsk/ainet/lang/memory/TensorView; } public abstract interface class sk/ainet/lang/tensor/data/ItemsAccessor { @@ -3971,6 +3987,7 @@ public final class sk/ainet/lang/tensor/data/LazyZeroFloatArrayTensorData : sk/a public fun getBuffer ()[F public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -3983,6 +4000,7 @@ public final class sk/ainet/lang/tensor/data/LazyZeroIntArrayTensorData : sk/ain public fun getBuffer ()[I public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([II)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -4007,6 +4025,7 @@ public final class sk/ainet/lang/tensor/data/MemorySegmentTensorData : sk/ainet/ public fun getSegment ()Ljava/lang/foreign/MemorySegment; public fun getSegmentByteOffset ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public final fun getVolume ()I public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V @@ -4042,6 +4061,7 @@ public final class sk/ainet/lang/tensor/data/MmapFloatTensorData : sk/ainet/lang public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getFloatBuffer ()Ljava/nio/FloatBuffer; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -4070,6 +4090,7 @@ public class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData : sk/ainet/lan public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getPackedData ()[B public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -4088,6 +4109,7 @@ public final class sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData : s public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getPackedData ()[B public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IF)V public synthetic fun set ([ILjava/lang/Object;)V public final fun transposedView ()Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData; @@ -4111,6 +4133,7 @@ public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$Companion { public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/NarrowFloatTensorDataKt { @@ -4141,6 +4164,7 @@ public final class sk/ainet/lang/tensor/data/Q4MemorySegmentTensorData : sk/aine public fun getSegment ()Ljava/lang/foreign/MemorySegment; public fun getSegmentByteOffset ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -4166,6 +4190,7 @@ public final class sk/ainet/lang/tensor/data/Q4_0BlockTensorData : sk/ainet/lang public fun getPackedData ()[B public fun getPhysicalBytes ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4201,6 +4226,7 @@ public final class sk/ainet/lang/tensor/data/Q4_0TensorData$Companion { public final class sk/ainet/lang/tensor/data/Q4_0TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q4_0TensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q4_0TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q4_0TensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q4_0TensorDataKt { @@ -4226,6 +4252,7 @@ public final class sk/ainet/lang/tensor/data/Q4_KBlockTensorData : sk/ainet/lang public fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun getSubBlockMin (II)F public fun getSubBlockScale (II)F + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4262,6 +4289,7 @@ public final class sk/ainet/lang/tensor/data/Q4_KTensorData$Companion { public final class sk/ainet/lang/tensor/data/Q4_KTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q4_KTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q4_KTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q4_KTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q4_KTensorDataKt { @@ -4282,6 +4310,7 @@ public final class sk/ainet/lang/tensor/data/Q5_0BlockTensorData : sk/ainet/lang public fun getPackedData ()[B public fun getPhysicalBytes ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4309,6 +4338,7 @@ public final class sk/ainet/lang/tensor/data/Q5_0TensorData$Companion { public final class sk/ainet/lang/tensor/data/Q5_0TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q5_0TensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q5_0TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q5_0TensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q5_1BlockTensorData : sk/ainet/lang/tensor/data/Q5_1TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { @@ -4325,6 +4355,7 @@ public final class sk/ainet/lang/tensor/data/Q5_1BlockTensorData : sk/ainet/lang public fun getPackedData ()[B public fun getPhysicalBytes ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4352,6 +4383,7 @@ public final class sk/ainet/lang/tensor/data/Q5_1TensorData$Companion { public final class sk/ainet/lang/tensor/data/Q5_1TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q5_1TensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q5_1TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q5_1TensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData : sk/ainet/lang/tensor/data/Q5_KTensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { @@ -4373,6 +4405,7 @@ public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData : sk/ainet/lang public fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun getSubBlockMin (II)F public fun getSubBlockScale (II)F + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4413,6 +4446,7 @@ public final class sk/ainet/lang/tensor/data/Q5_KTensorData$Companion { public final class sk/ainet/lang/tensor/data/Q5_KTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q5_KTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q5_KTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q5_KTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q5_KTensorDataKt { @@ -4436,6 +4470,7 @@ public final class sk/ainet/lang/tensor/data/Q6_KBlockTensorData : sk/ainet/lang public fun getPhysicalBytes ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; public fun getSubBlockScale (II)I + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4470,6 +4505,7 @@ public final class sk/ainet/lang/tensor/data/Q6_KTensorData$Companion { public final class sk/ainet/lang/tensor/data/Q6_KTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q6_KTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q6_KTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q6_KTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q6_KTensorDataKt { @@ -4496,6 +4532,7 @@ public final class sk/ainet/lang/tensor/data/Q8MemorySegmentTensorData : sk/aine public fun getSegment ()Ljava/lang/foreign/MemorySegment; public fun getSegmentByteOffset ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V } @@ -4521,6 +4558,7 @@ public final class sk/ainet/lang/tensor/data/Q8_0BlockTensorData : sk/ainet/lang public fun getPackedData ()[B public fun getPhysicalBytes ()J public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4550,6 +4588,7 @@ public final class sk/ainet/lang/tensor/data/Q8_0TensorData$Companion { public final class sk/ainet/lang/tensor/data/Q8_0TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Q8_0TensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/Q8_0TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/Q8_0TensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/Q8_0TensorDataKt { @@ -4564,11 +4603,13 @@ public abstract interface class sk/ainet/lang/tensor/data/TensorData : sk/ainet/ public fun copyToFloatArray ()[F public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public abstract fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/TensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/TensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/TensorData;)Lsk/ainet/lang/memory/TensorView; } public abstract interface class sk/ainet/lang/tensor/data/TensorDataFactory { @@ -4618,6 +4659,7 @@ public final class sk/ainet/lang/tensor/data/Ternary2BitTensorData : sk/ainet/la public fun getPhysicalBytes ()J public fun getScale ()F public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([IB)V public synthetic fun set ([ILjava/lang/Object;)V public fun toFloatArray ()[F @@ -4640,6 +4682,7 @@ public abstract interface class sk/ainet/lang/tensor/data/TernaryTensorData : sk public final class sk/ainet/lang/tensor/data/TernaryTensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/TernaryTensorData;)[F public static fun getEncoding (Lsk/ainet/lang/tensor/data/TernaryTensorData;)Lsk/ainet/lang/tensor/storage/TensorEncoding; + public static fun getView (Lsk/ainet/lang/tensor/data/TernaryTensorData;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/tensor/data/TernaryTensorDataKt { @@ -4701,6 +4744,7 @@ public final class sk/ainet/lang/tensor/data/views/UnsqueezedTensorData : sk/ain public fun get ([I)Ljava/lang/Object; public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; public fun set ([ILjava/lang/Object;)V } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt index 1338ddd15..da1909b18 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt @@ -89,10 +89,11 @@ public class TensorView( public fun get(vararg indices: Int): Float { storage.checkAlive() require(indices.size == shape.rank) { "expected ${shape.rank} indices, got ${indices.size}" } - if (format.isDense) return readDense(flatDenseIndex(indices)) - val d = decoder ?: throw IllegalStateException("no decoder for ${format.encoding.name}") - val flat = flatLogicalIndex(indices) - return d.decodeElement(storage, layout, flat) + // A decoder wins over the plain path: narrow floats are Dense(2) yet still need decoding. + val d = decoder + if (d != null) return d.decodeElement(storage, layout, flatLogicalIndex(indices)) + check(format.isDense) { "no decoder for ${format.encoding.name}" } + return readDense(flatDenseIndex(indices)) } /** Write [value] at [indices] (dense, mutable views only). */ @@ -115,7 +116,10 @@ public class TensorView( return flat } - /** Logical element index for a packed view: the layout addresses blocks, so the last axis contributes elements. */ + /** + * Logical element index for a decoded view: the layout addresses blocks (one element per block + * for narrow floats), so the last axis contributes both a block step and an offset inside it. + */ private fun flatLogicalIndex(indices: IntArray): Long { val bs = blockSize() val last = indices[indices.size - 1] @@ -220,3 +224,27 @@ public class PackedBlockDecoder(private val packed: PackedBlockStorage) : BlockD packed.dequantizeBlock(blockIndex.toInt(), out, outOffset) } } + +/** + * Decoder for 16-bit narrow floats (FP16 / BF16) held two bytes per element — the "block" is one + * element, so a narrow-float view decodes element by element through its [codec]. + */ +@ExperimentalMemoryApi +public class NarrowFloatDecoder(private val codec: sk.ainet.lang.types.NarrowFloatCodec) : BlockDecoder { + override val blockSize: Int get() = 1 + override val bytesPerBlock: Int get() = codec.bytesPerElement + + override fun decodeBlock(storage: Storage, blockIndex: Long, out: FloatArray, outOffset: Int) { + out[outOffset] = decodeAt(storage, blockIndex) + } + + override fun decodeElement(storage: Storage, layout: Layout, flatElementIndex: Long): Float = decodeAt(storage, flatElementIndex) + + private fun decodeAt(storage: Storage, elementIndex: Long): Float { + val heap = storage as? Storage.Heap ?: throw UnsupportedOperationException("narrow-float views need heap storage in this milestone") + val bytes = heap.bytes ?: throw UnsupportedOperationException("narrow-float views need byte storage") + val off = heap.arrayOffset + (elementIndex * codec.bytesPerElement).toInt() + val bits = (bytes[off].toInt() and 0xFF) or ((bytes[off + 1].toInt() and 0xFF) shl 8) + return codec.decode(bits) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/LazyMaterializationStrategy.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/LazyMaterializationStrategy.kt index d951c0532..b3005f635 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/LazyMaterializationStrategy.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/LazyMaterializationStrategy.kt @@ -98,10 +98,11 @@ public class LazyMaterializationStrategy : MaterializationStrategy * Lazy tensor data implementation with sparse element caching. */ private class LazyMaterializedTensorData( - private val view: TensorView + // named `sourceView`, not `view`: TensorData.view is the memory-model view (SKEEP-003) + private val sourceView: TensorView ) : TensorData { - override val shape: Shape = view.viewShape + override val shape: Shape = sourceView.viewShape // Cache for materialized elements // Using a map to store only accessed elements @@ -113,7 +114,7 @@ public class LazyMaterializationStrategy : MaterializationStrategy // Check if element is already cached return elementCache[cacheKey] ?: run { // Element not cached, fetch from view and cache it - val element = view.data.get(*indices) + val element = sourceView.data.get(*indices) elementCache[cacheKey] = element element } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt index fce7bdc9b..160f1d3a7 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt @@ -62,6 +62,25 @@ public open class NarrowFloatDenseTensorData( /** Physically two bytes per element whatever the declared dtype witness. */ override val encoding: TensorEncoding get() = TensorEncoding.Dense(NarrowFloatTensorData.BYTES_PER_ELEMENT) + /** + * A view over the *same* packed bytes, decoded by this data's [codec] (SKEEP-003 §4.1 façade). + * The dtype is the codec's (FP16 or BF16) and the encoding `Dense(2)`; `view.get()` returns the + * decoded float, exactly like [get]. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + override val view: sk.ainet.lang.memory.TensorView + get() = sk.ainet.lang.memory.TensorView( + shape = shape, + format = sk.ainet.lang.memory.Format(codec.dtype, TensorEncoding.Dense(NarrowFloatTensorData.BYTES_PER_ELEMENT)), + layout = sk.ainet.lang.memory.Layout( + shape = shape, + strides = sk.ainet.lang.memory.Layout.rowMajorStrides(shape), + elementBytes = NarrowFloatTensorData.BYTES_PER_ELEMENT, + ), + storage = sk.ainet.lang.memory.Storage.Heap.wrap(data, mutable = false), + decoder = sk.ainet.lang.memory.NarrowFloatDecoder(codec), + ) + init { val requiredBytes = shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT require(data.size >= requiredBytes) { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorData.kt index b22eeab39..7c6bc55bf 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorData.kt @@ -80,6 +80,20 @@ public interface TensorData : ItemsAccessor { */ public val encoding: sk.ainet.lang.tensor.storage.TensorEncoding? get() = null + /** + * This data as a [sk.ainet.lang.memory.TensorView] — `Shape + Format + Layout + Storage` — or + * `null` when the implementation cannot expose one (SKEEP-003 §4.1: `TensorData` becomes a + * façade over the view; migrated kernels take the view, everything else keeps using this + * interface unchanged). + * + * The view is over the *same* bytes: for array-backed data the storage borrows the array + * (`Storage.Heap.wrap`), so writes through either side are visible on both and nothing is + * copied. Per-element access stays on this interface's own fast path — the Phase-2 spike + * (#1016) showed a view is for unwrapping once per call, not for per-element reads. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + public val view: sk.ainet.lang.memory.TensorView? get() = null + /** * Copies all tensor data to a FloatArray. * @@ -119,6 +133,15 @@ public interface FloatArrayTensorData : TensorData { public val buffer: FloatArray override fun copyToFloatArray(): FloatArray = buffer.copyOf() + + /** A dense FP32 view borrowing [buffer] — zero-copy, the same bytes this data reads and writes. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + override val view: sk.ainet.lang.memory.TensorView + get() = sk.ainet.lang.memory.TensorView.dense( + sk.ainet.lang.memory.Storage.Heap.wrap(buffer), + shape, + sk.ainet.lang.types.FP32, + ) } /** @@ -126,4 +149,13 @@ public interface FloatArrayTensorData : TensorData { */ public interface IntArrayTensorData : TensorData { public val buffer: IntArray + + /** A dense Int32 view borrowing [buffer] — zero-copy, the same bytes this data reads and writes. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + override val view: sk.ainet.lang.memory.TensorView + get() = sk.ainet.lang.memory.TensorView.dense( + sk.ainet.lang.memory.Storage.Heap.wrap(buffer), + shape, + sk.ainet.lang.types.Int32, + ) } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorDataViewTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorDataViewTest.kt new file mode 100644 index 000000000..8ea8dd249 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorDataViewTest.kt @@ -0,0 +1,95 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData +import sk.ainet.lang.tensor.data.DenseIntArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.LazyZeroFloatArrayTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * SKEEP-003 §4.1: `TensorData` becomes a façade over `TensorView` — the view is over the *same* + * bytes (borrowed, zero-copy), reads agree with the data's own accessors, and nothing is copied. + */ +@OptIn(ExperimentalMemoryApi::class) +class TensorDataViewTest { + + @Test + fun denseFloatDataExposesAZeroCopyView() { + val buf = FloatArray(6) { it.toFloat() } + val data = DenseFloatArrayTensorData(Shape(2, 3), buf) + val v = assertNotNull(data.view) + assertEquals(Format.dense(FP32), v.format); assertEquals(Shape(2, 3), v.shape); assertTrue(v.isContiguous) + // same bytes, no copy: the storage borrows the array + assertSame(buf, (v.storage as Storage.Heap).floats) + assertEquals(ScopeKind.AMBIENT, v.storage.scope) + // reads agree + for (i in 0 until 2) for (j in 0 until 3) assertEquals(data.get(i, j), v.get(i, j)) + // writes are visible through both + data.set(1, 2, value = 42f); assertEquals(42f, v.get(1, 2)) + v.set(0, 0, value = -1f); assertEquals(-1f, data.get(0, 0)); assertEquals(-1f, buf[0]) + assertContentEquals(data.copyToFloatArray(), v.toFloatArray()) + } + + @Test + fun denseIntDataExposesAnInt32View() { + val buf = IntArray(4) { it * 10 } + val data = DenseIntArrayTensorData(Shape(4), buf) + val v = assertNotNull(data.view) + assertEquals(Format.dense(Int32), v.format) + assertSame(buf, (v.storage as Storage.Heap).ints) + assertEquals(20f, v.get(2)) + } + + @Test + fun lazyZeroDataMaterializesThroughTheView() { + val data = LazyZeroFloatArrayTensorData(Shape(2, 2)) + val v = assertNotNull(data.view) + assertContentEquals(FloatArray(4), v.toFloatArray()) + data.set(1, 1, value = 5f) + assertEquals(5f, assertNotNull(data.view).get(1, 1)) // the view is over the materialized buffer + } + + @Test + fun narrowFloatDataDecodesThroughTheView() { + // BF16: the high 16 bits of the float + fun bf16(v: Float): Int = (v.toRawBits() ushr 16) and 0xFFFF + val values = floatArrayOf(1f, -2.5f, 0.5f, 100f) + val bytes = ByteArray(values.size * 2) + for ((i, x) in values.withIndex()) { val b = bf16(x); bytes[i * 2] = (b and 0xFF).toByte(); bytes[i * 2 + 1] = ((b ushr 8) and 0xFF).toByte() } + val data = Bf16DenseTensorData(Shape(4), bytes) + val v = assertNotNull(data.view) + assertEquals(BF16, v.format.dtype); assertEquals(2, v.layout.elementBytes) + for (i in values.indices) assertEquals(data.get(i), v.get(i), "element $i") + assertContentEquals(data.copyToFloatArray(), v.toFloatArray()) + assertSame(bytes, (v.storage as Storage.Heap).bytes) + + val fp16 = Fp16DenseTensorData(Shape(2), ByteArray(4) { (it * 17).toByte() }) + val fv = assertNotNull(fp16.view) + assertEquals(FP16, fv.format.dtype) + assertEquals(fp16.get(0), fv.get(0)); assertEquals(fp16.get(1), fv.get(1)) + } + + @Test + fun dataWithoutAViewReportsNull() { + val anonymous = object : TensorData { + override val shape: Shape = Shape(1) + override fun get(vararg indices: Int): Float = 0f + override fun set(vararg indices: Int, value: Float) {} + } + assertNull(anonymous.view) + assertNull(anonymous.encoding) + } +} diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/data/MemorySegmentTensorData.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/data/MemorySegmentTensorData.kt index aafd4eff3..d8b421d6c 100644 --- a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/data/MemorySegmentTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/data/MemorySegmentTensorData.kt @@ -49,6 +49,21 @@ public class MemorySegmentTensorData private constructor( override val segmentByteOffset: Long, private val ownsArena: Boolean, ) : TensorData, MemorySegmentBackedData { + /** + * A dense view over the *same* off-heap bytes (SKEEP-003 §4.1 façade): the storage borrows this + * data's [segment] — nothing is copied and a migrated kernel unwraps it once with + * `SegmentStorage.segment()`. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + override val view: sk.ainet.lang.memory.TensorView + get() = sk.ainet.lang.memory.TensorView.dense( + sk.ainet.lang.memory.SegmentStorage.borrow( + if (segmentByteOffset == 0L) segment else segment.asSlice(segmentByteOffset), + ), + shape, + sk.ainet.lang.types.FP32, + ) + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) private val strides: IntArray = shape.computeStrides()