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 7486398e0..8debc1862 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 @@ -675,6 +675,17 @@ public final class sk/ainet/lang/memory/AllocationSpec$Companion { 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 class sk/ainet/lang/memory/BlockDecoder { + public abstract 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 abstract fun getBlockSize ()I + public abstract fun getBytesPerBlock ()I +} + +public final class sk/ainet/lang/memory/BlockDecoder$DefaultImpls { + public static fun decodeElement (Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/memory/Layout;J)F +} + public final class sk/ainet/lang/memory/DescribeKt { public static final fun describe (Lsk/ainet/lang/tensor/Tensor;)Ljava/lang/String; public static final fun describe (Lsk/ainet/lang/tensor/storage/TensorStorage;Lsk/ainet/lang/tensor/TensorId;)Ljava/lang/String; @@ -751,6 +762,39 @@ public final class sk/ainet/lang/memory/ForwardScope : sk/ainet/lang/memory/Scop public static synthetic fun retain$default (Lsk/ainet/lang/memory/ForwardScope;Lsk/ainet/lang/memory/Storage$Heap;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; } +public final class sk/ainet/lang/memory/Layout { + public static final field Companion Lsk/ainet/lang/memory/Layout$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[IJIZ)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[IJIZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun byteOffsetOf ([I)J + public fun equals (Ljava/lang/Object;)Z + public final fun getBlocked ()Z + public final fun getElementBytes ()I + public final fun getElementCount ()J + public final fun getOffsetBytes ()J + public final fun getOffsetElements ()J + public final fun getShape ()Lsk/ainet/lang/tensor/Shape; + public final fun getStrides ()[I + public fun hashCode ()I + public final fun indexOf ([I)J + public final fun isContiguous ()Z + public final fun isRowMajor ()Z + public final fun narrow (III)Lsk/ainet/lang/memory/Layout; + public final fun squeeze (I)Lsk/ainet/lang/memory/Layout; + public fun toString ()Ljava/lang/String; + public final fun transpose (II)Lsk/ainet/lang/memory/Layout; + public static synthetic fun transpose$default (Lsk/ainet/lang/memory/Layout;IIILjava/lang/Object;)Lsk/ainet/lang/memory/Layout; + public final fun unsqueeze (I)Lsk/ainet/lang/memory/Layout; +} + +public final class sk/ainet/lang/memory/Layout$Companion { + public final fun blocked (Lsk/ainet/lang/tensor/Shape;IIJ)Lsk/ainet/lang/memory/Layout; + public static synthetic fun blocked$default (Lsk/ainet/lang/memory/Layout$Companion;Lsk/ainet/lang/tensor/Shape;IIJILjava/lang/Object;)Lsk/ainet/lang/memory/Layout; + public final fun rowMajor (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Format;J)Lsk/ainet/lang/memory/Layout; + public static synthetic fun rowMajor$default (Lsk/ainet/lang/memory/Layout$Companion;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Format;JILjava/lang/Object;)Lsk/ainet/lang/memory/Layout; + public final fun rowMajorStrides (Lsk/ainet/lang/tensor/Shape;)[I +} + public final class sk/ainet/lang/memory/MappedBufferStorage : sk/ainet/lang/memory/Storage$Mapped { public static final field Companion Lsk/ainet/lang/memory/MappedBufferStorage$Companion; public synthetic fun (JLjava/nio/file/Path;JLjava/nio/ByteBuffer;Lsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;Lkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -843,6 +887,14 @@ public final class sk/ainet/lang/memory/Owner$Owned : sk/ainet/lang/memory/Owner public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/PackedBlockDecoder : sk/ainet/lang/memory/BlockDecoder { + public fun (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)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 final class sk/ainet/lang/memory/PlatformStorage { public static final field INSTANCE Lsk/ainet/lang/memory/PlatformStorage; public final fun allocate (JLsk/ainet/lang/tensor/storage/MemoryDomain;Lsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/Storage; @@ -1016,6 +1068,38 @@ public final class sk/ainet/lang/memory/StorageId$Companion { public final fun next-TPZW6QE ()J } +public final class sk/ainet/lang/memory/TensorView { + public static final field Companion Lsk/ainet/lang/memory/TensorView$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Layout;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/BlockDecoder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Layout;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/BlockDecoder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun get ([I)F + public final fun getElementCount ()J + public final fun getFormat ()Lsk/ainet/lang/memory/Format; + public final fun getId ()Lsk/ainet/lang/tensor/TensorId; + public final fun getLayout ()Lsk/ainet/lang/memory/Layout; + public final fun getShape ()Lsk/ainet/lang/tensor/Shape; + public final fun getStorage ()Lsk/ainet/lang/memory/Storage; + public final fun isContiguous ()Z + public final fun isMutable ()Z + public final fun materialize (Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Scope;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun materialize$default (Lsk/ainet/lang/memory/TensorView;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Scope;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun narrow (III)Lsk/ainet/lang/memory/TensorView; + public final fun set ([IF)V + public final fun squeeze (I)Lsk/ainet/lang/memory/TensorView; + public final fun toFloatArray ()[F + public fun toString ()Ljava/lang/String; + public final fun transpose (II)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun transpose$default (Lsk/ainet/lang/memory/TensorView;IIILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun unsqueeze (I)Lsk/ainet/lang/memory/TensorView; +} + +public final class sk/ainet/lang/memory/TensorView$Companion { + public final fun dense (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun dense$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun packed (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun packed$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; +} + public final class sk/ainet/lang/memory/plan/Budget { public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion; public static final field RESERVE_ANDROID_JVM J diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Layout.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Layout.kt new file mode 100644 index 000000000..fc07d4a9c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Layout.kt @@ -0,0 +1,152 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape + +/** + * Strides, byte offset and contiguity of a view over a [Storage] (SKEEP-003 §0 *Layout*). Pure + * metadata: it says *where* the elements of a shape live inside a byte range, never what they mean + * (that is [Format]) and never who owns them (that is [Storage]). + * + * Strides are in **elements** of the layout's unit — for a dense format one element is + * `format.dtype.sizeInBytes`; for a block-packed format ([blocked]) a stride step addresses one + * *block*, which is what makes a packed weight sliceable and transposable without touching bytes + * (rule 5: "a view is a TensorView with the same Storage, different Layout"). + * + * @property shape the extents this layout addresses + * @property strides one stride per dimension, in elements (or blocks for a blocked layout) + * @property offsetElements element (or block) offset of the first element inside the storage + * @property elementBytes bytes per element (dense) or per block (blocked) + */ +@ExperimentalMemoryApi +public class Layout( + public val shape: Shape, + public val strides: IntArray, + public val offsetElements: Long = 0L, + public val elementBytes: Int = 4, + public val blocked: Boolean = false, +) { + init { + require(strides.size == shape.rank) { "strides (${strides.size}) must match rank (${shape.rank})" } + require(offsetElements >= 0) { "offsetElements must be >= 0" } + require(elementBytes > 0) { "elementBytes must be > 0" } + } + + /** Number of elements addressed (the shape's volume). */ + public val elementCount: Long get() = shape.volume.toLong() + + /** Byte offset of the first element. */ + public val offsetBytes: Long get() = offsetElements * elementBytes + + /** Row-major (C-order) strides for [shape] — the canonical layout. */ + public val isRowMajor: Boolean get() = strides.contentEquals(rowMajorStrides(shape)) + + /** Whether the elements occupy one gap-free run, i.e. the layout can be handed to a kernel as a flat range. */ + public val isContiguous: Boolean + get() { + if (shape.rank == 0) return true + var expected = 1 + for (d in shape.rank - 1 downTo 0) { + val extent = shape[d] + if (extent == 1) continue // a unit extent's stride is irrelevant + if (strides[d] != expected) return false + expected *= extent + } + return true + } + + /** Flat element (or block) index of [indices] inside the storage, including [offsetElements]. */ + public fun indexOf(vararg indices: Int): Long { + require(indices.size == shape.rank) { "expected ${shape.rank} indices, got ${indices.size}" } + var flat = offsetElements + for (d in indices.indices) { + val i = indices[d] + require(i in 0 until shape[d]) { "index $i out of range for axis $d (extent ${shape[d]})" } + flat += i.toLong() * strides[d] + } + return flat + } + + /** Byte offset of [indices]. */ + public fun byteOffsetOf(vararg indices: Int): Long = indexOf(*indices) * elementBytes + + /** A layout over `[from, to)` of [axis] — same strides, shifted offset (zero-copy slicing). */ + public fun narrow(axis: Int, from: Int, size: Int): Layout { + require(axis in 0 until shape.rank) { "axis $axis out of range for rank ${shape.rank}" } + require(from >= 0 && size >= 0 && from + size <= shape[axis]) { "narrow($axis, $from, $size) outside extent ${shape[axis]}" } + val dims = shape.dimensions.copyOf(); dims[axis] = size + return Layout(Shape(dims), strides.copyOf(), offsetElements + from.toLong() * strides[axis], elementBytes, blocked) + } + + /** Swap two axes — metadata only; the bytes are untouched (the packed-transpose trick, rule 5). */ + public fun transpose(axis0: Int = shape.rank - 2, axis1: Int = shape.rank - 1): Layout { + require(shape.rank >= 2) { "transpose needs rank >= 2" } + require(axis0 in 0 until shape.rank && axis1 in 0 until shape.rank) { "axes out of range" } + val dims = shape.dimensions.copyOf(); val st = strides.copyOf() + val d = dims[axis0]; dims[axis0] = dims[axis1]; dims[axis1] = d + val s = st[axis0]; st[axis0] = st[axis1]; st[axis1] = s + return Layout(Shape(dims), st, offsetElements, elementBytes, blocked) + } + + /** Insert a unit axis at [axis] (stride 0 — it is never stepped). */ + public fun unsqueeze(axis: Int): Layout { + require(axis in 0..shape.rank) { "axis $axis out of range for rank ${shape.rank}" } + val dims = IntArray(shape.rank + 1); val st = IntArray(shape.rank + 1) + var j = 0 + for (i in 0..shape.rank) { + if (i == axis) { dims[i] = 1; st[i] = 0 } else { dims[i] = shape[j]; st[i] = strides[j]; j++ } + } + return Layout(Shape(dims), st, offsetElements, elementBytes, blocked) + } + + /** Drop the unit axis at [axis]. */ + public fun squeeze(axis: Int): Layout { + require(axis in 0 until shape.rank) { "axis $axis out of range for rank ${shape.rank}" } + require(shape[axis] == 1) { "axis $axis has extent ${shape[axis]}, not 1" } + val dims = ArrayList(shape.rank - 1); val st = ArrayList(shape.rank - 1) + for (i in 0 until shape.rank) if (i != axis) { dims += shape[i]; st += strides[i] } + return Layout(Shape(dims.toIntArray()), st.toIntArray(), offsetElements, elementBytes, blocked) + } + + override fun toString(): String = + "Layout($shape, strides=${strides.joinToString(",", "[", "]")}, offset=$offsetElements${if (blocked) " blocks" else ""}, ${if (isContiguous) "contiguous" else "strided"})" + + override fun equals(other: Any?): Boolean = other is Layout && other.shape == shape && + other.strides.contentEquals(strides) && other.offsetElements == offsetElements && + other.elementBytes == elementBytes && other.blocked == blocked + + override fun hashCode(): Int { + var h = shape.hashCode() + h = 31 * h + strides.contentHashCode(); h = 31 * h + offsetElements.hashCode() + h = 31 * h + elementBytes; h = 31 * h + blocked.hashCode() + return h + } + + public companion object { + /** Row-major strides of [shape], in elements. */ + public fun rowMajorStrides(shape: Shape): IntArray { + val st = IntArray(shape.rank) + var acc = 1 + for (d in shape.rank - 1 downTo 0) { st[d] = acc; acc *= shape[d] } + return st + } + + /** The canonical dense row-major layout of [shape] for [format]. */ + public fun rowMajor(shape: Shape, format: Format, offsetElements: Long = 0L): Layout = + Layout(shape, rowMajorStrides(shape), offsetElements, format.dtype.sizeInBytes, blocked = false) + + /** + * A blocked layout: the last axis is measured in blocks of `blockSize` elements, one + * "element" of the layout being one packed block of `bytesPerBlock` bytes. Used for the + * GGML block formats, where a view addresses whole blocks (rule 5). + */ + public fun blocked(shape: Shape, blockSize: Int, bytesPerBlock: Int, offsetBlocks: Long = 0L): Layout { + require(blockSize > 0 && bytesPerBlock > 0) { "block geometry must be positive" } + require(shape.rank >= 1) { "blocked layout needs rank >= 1" } + require(shape[shape.rank - 1] % blockSize == 0) { "last extent ${shape[shape.rank - 1]} is not a multiple of the block size $blockSize" } + val dims = shape.dimensions.copyOf() + dims[dims.size - 1] = dims[dims.size - 1] / blockSize + val blockShape = Shape(dims) + return Layout(blockShape, rowMajorStrides(blockShape), offsetBlocks, bytesPerBlock, blocked = true) + } + } +} 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 new file mode 100644 index 000000000..1338ddd15 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt @@ -0,0 +1,222 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 + +/** + * `Shape + Format + Layout + Storage` — the interpretation of some bytes as an n-d typed array, and + * **the only thing a kernel receives** (SKEEP-003 §0 *TensorView*, §4.2). A view never owns bytes: + * slicing, transposing or unsqueezing one yields another view over the same [Storage] (rule 5, + * zero-copy, `Owner.Alias`), and [materialize] is the single copy point (rule 6). + * + * Element access follows rule 4: [get] returns the **decoded logical value** — for a packed + * encoding it decodes the containing block; it never returns a raw byte. That is the correct, slow + * reference path; a production kernel unwraps the storage once per call + * ([Storage.Heap.floats] / `SegmentStorage.segment()`), as the Phase-2 spike (#1016) requires. + */ +@ExperimentalMemoryApi +public class TensorView( + public val shape: Shape, + public val format: Format, + public val layout: Layout, + public val storage: Storage, + public val id: TensorId? = null, + /** Decoder for a block-packed [format]; required to read a non-dense view element-wise. */ + private val decoder: BlockDecoder? = null, +) { + init { + require(storage.isAlive) { "cannot build a view over closed storage ${storage.id}" } + if (!format.isDense) require(decoder != null) { "a ${format.encoding.name} view needs a BlockDecoder to decode elements" } + } + + public val elementCount: Long get() = shape.volume.toLong() + /** Whether the layout addresses one gap-free run (a kernel may take it as a flat range). */ + public val isContiguous: Boolean get() = layout.isContiguous + /** Whether the bytes may be written through this view. */ + public val isMutable: Boolean get() = storage.isMutable + + // ---- views (zero-copy) ---- + + /** + * A view over `[from, from + size)` of [axis] — zero-copy. On the block axis of a packed view + * both bounds must be whole blocks (the bytes of a block are indivisible). + */ + public fun narrow(axis: Int, from: Int, size: Int): TensorView { + require(axis in 0 until shape.rank) { "axis $axis out of range for rank ${shape.rank}" } + require(from >= 0 && size >= 0 && from + size <= shape[axis]) { "narrow($axis, $from, $size) outside extent ${shape[axis]}" } + val onBlockAxis = layout.blocked && axis == shape.rank - 1 + val unit = if (onBlockAxis) blockSize() else 1 + if (onBlockAxis) require(from % unit == 0 && size % unit == 0) { "narrowing the block axis must align to the block size $unit" } + return derive(narrowShape(axis, size), layout.narrow(axis, from / unit, size / unit), idSuffix = "$from..${from + size})") + } + + /** A transposed view — metadata only; packed bytes are untouched (the packed-transpose trick). */ + public fun transpose(axis0: Int = shape.rank - 2, axis1: Int = shape.rank - 1): TensorView { + require(!format.isDense.not() || !layout.blocked || axis1 == shape.rank - 1) { "a blocked layout cannot move its block axis" } + val dims = shape.dimensions.copyOf(); val t = dims[axis0]; dims[axis0] = dims[axis1]; dims[axis1] = t + return derive(Shape(dims), layout.transpose(axis0, axis1), idSuffix = "ᵀ") + } + + /** A view with a unit axis inserted at [axis]. */ + public fun unsqueeze(axis: Int): TensorView { + val dims = IntArray(shape.rank + 1); var j = 0 + for (i in 0..shape.rank) { if (i == axis) dims[i] = 1 else { dims[i] = shape[j]; j++ } } + return derive(Shape(dims), layout.unsqueeze(axis)) + } + + /** A view with the unit axis at [axis] removed. */ + public fun squeeze(axis: Int): TensorView { + require(shape[axis] == 1) { "axis $axis has extent ${shape[axis]}, not 1" } + val dims = ArrayList(shape.rank - 1) + for (i in 0 until shape.rank) if (i != axis) dims += shape[i] + return derive(Shape(dims.toIntArray()), layout.squeeze(axis)) + } + + private fun derive(newShape: Shape, newLayout: Layout, idSuffix: String? = null): TensorView = + TensorView(newShape, format, newLayout, storage, if (idSuffix == null) id else id?.view(idSuffix), decoder) + + private fun narrowShape(axis: Int, size: Int): Shape { val d = shape.dimensions.copyOf(); d[axis] = size; return Shape(d) } + + private fun blockSize(): Int = decoder?.blockSize ?: 1 + + // ---- element access (rule 4: decode, never a raw byte) ---- + + /** The decoded logical value at [indices]; `Float` for every float dtype and every packed encoding. */ + 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) + } + + /** Write [value] at [indices] (dense, mutable views only). */ + public fun set(vararg indices: Int, value: Float) { + storage.checkAlive() + check(format.isDense) { "cannot write through a ${format.encoding.name} view; requantize into a dense view instead" } + check(isMutable) { "storage ${storage.id} is read-only" } + val heap = storage as? Storage.Heap ?: throw UnsupportedOperationException("element writes need heap storage in this milestone") + val floats = heap.floats ?: throw UnsupportedOperationException("element writes need float storage") + floats[heap.arrayOffset + flatDenseIndex(indices).toInt()] = value + } + + private fun flatDenseIndex(indices: IntArray): Long { + var flat = layout.offsetElements + for (d in indices.indices) { + val i = indices[d] + require(i in 0 until shape[d]) { "index $i out of range for axis $d (extent ${shape[d]})" } + flat += i.toLong() * layout.strides[d] + } + return flat + } + + /** Logical element index for a packed view: the layout addresses blocks, so the last axis contributes elements. */ + private fun flatLogicalIndex(indices: IntArray): Long { + val bs = blockSize() + val last = indices[indices.size - 1] + val blockIdxWithinRow = last / bs + val within = last % bs + var flat = layout.offsetElements + for (d in 0 until indices.size - 1) { + val i = indices[d] + require(i in 0 until shape[d]) { "index $i out of range for axis $d (extent ${shape[d]})" } + flat += i.toLong() * layout.strides[d] + } + flat += blockIdxWithinRow.toLong() * layout.strides[indices.size - 1] + return flat * bs + within + } + + private fun readDense(flat: Long): Float = when (val s = storage) { + is Storage.Heap -> { + val f = s.floats + if (f != null) f[s.arrayOffset + flat.toInt()] + else { + val ints = s.ints + if (ints != null) ints[s.arrayOffset + flat.toInt()].toFloat() + else throw UnsupportedOperationException("dense element access over byte storage needs a decoder") + } + } + else -> throw UnsupportedOperationException("element access over ${s::class.simpleName} needs a platform reader (use a kernel)") + } + + /** Every element in row-major order, decoded — the reference materialization. */ + public fun toFloatArray(): FloatArray { + storage.checkAlive() + val out = FloatArray(elementCount.toInt()) + val idx = IntArray(shape.rank) + for (flat in out.indices) { + var rem = flat + for (d in shape.rank - 1 downTo 0) { idx[d] = rem % shape[d]; rem /= shape[d] } + out[flat] = get(*idx) + } + return out + } + + /** + * The single copy point (rule 6): decode/convert this view into a new dense view of + * [targetFormat] owned by [scope]. Everything else in the memory model is a view. + */ + public fun materialize(targetFormat: Format = Format.dense(FP32), scope: Scope = Scope.Ambient): TensorView { + require(targetFormat.isDense) { "materialize targets a dense format; re-quantization is an adapter (#1027)" } + require(targetFormat.dtype == FP32) { "only FP32 materialization is implemented in this milestone" } + val values = toFloatArray() + val out = scope.allocateFloats(values.size, id) + values.copyInto(out.floats!!, out.arrayOffset) + return TensorView(shape, targetFormat, Layout.rowMajor(shape, targetFormat), out, id) + } + + override fun toString(): String = "TensorView(${id?.canonical ?: "—"}, $format, $shape, $layout, ${storage.id})" + + public companion object { + /** A dense row-major view of [shape] over [storage]. */ + public fun dense(storage: Storage, shape: Shape, dtype: DType = FP32, id: TensorId? = null): TensorView { + val format = Format.dense(dtype) + return TensorView(shape, format, Layout.rowMajor(shape, format), storage, id) + } + + /** + * A block-packed view: [shape] in logical elements, [storage] holding `blockCount` packed + * blocks in row-major block order, decoded by [decoder]. Slicing and transposing address + * whole blocks — the bytes are never touched. + */ + public fun packed(storage: Storage, shape: Shape, encoding: TensorEncoding, decoder: BlockDecoder, dtype: DType = FP32, id: TensorId? = null): TensorView = + TensorView(shape, Format(dtype, encoding), Layout.blocked(shape, decoder.blockSize, decoder.bytesPerBlock), storage, id, decoder) + } +} + +/** + * Decodes one block of a packed encoding out of a [Storage] — the bridge between the block formats + * (`PackedBlockStorage` implementations today, `Encoding` descriptors in M2) and [TensorView]. + */ +@ExperimentalMemoryApi +public interface BlockDecoder { + public val blockSize: Int + public val bytesPerBlock: Int + + /** Decode the block at [blockIndex] (of the storage's block sequence) into [out] at [outOffset]. */ + public fun decodeBlock(storage: Storage, blockIndex: Long, out: FloatArray, outOffset: Int) + + /** The logical element at [flatElementIndex]; decodes its block by default. */ + public fun decodeElement(storage: Storage, layout: Layout, flatElementIndex: Long): Float { + val block = flatElementIndex / blockSize + val within = (flatElementIndex % blockSize).toInt() + val buf = FloatArray(blockSize) + decodeBlock(storage, block, buf, 0) + return buf[within] + } +} + +/** A [BlockDecoder] backed by an existing [PackedBlockStorage] implementation (the M1 bridge). */ +@ExperimentalMemoryApi +public class PackedBlockDecoder(private val packed: PackedBlockStorage) : BlockDecoder { + override val blockSize: Int get() = packed.blockSize + override val bytesPerBlock: Int get() = (packed.physicalBytes / maxOf(packed.blockCount, 1)).toInt() + override fun decodeBlock(storage: Storage, blockIndex: Long, out: FloatArray, outOffset: Int) { + packed.dequantizeBlock(blockIndex.toInt(), out, outOffset) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/LayoutTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/LayoutTest.kt new file mode 100644 index 000000000..e2e8494f0 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/LayoutTest.kt @@ -0,0 +1,82 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** SKEEP-003 §0 *Layout*: strides, offset and contiguity of a view — metadata only. */ +@OptIn(ExperimentalMemoryApi::class) +class LayoutTest { + + private val f32 = Format.dense(FP32) + + @Test + fun rowMajorStridesOffsetsAndContiguity() { + val l = Layout.rowMajor(Shape(2, 3, 4), f32) + assertContentEquals(intArrayOf(12, 4, 1), l.strides) + assertTrue(l.isRowMajor); assertTrue(l.isContiguous) + assertEquals(4, l.elementBytes); assertEquals(24L, l.elementCount); assertEquals(0L, l.offsetBytes) + assertEquals(17L, l.indexOf(1, 1, 1)); assertEquals(68L, l.byteOffsetOf(1, 1, 1)) + assertFailsWith { l.indexOf(2, 0, 0) } + assertFailsWith { l.indexOf(0, 0) } + } + + @Test + fun narrowKeepsStridesAndShiftsTheOffset() { + val l = Layout.rowMajor(Shape(4, 8), f32) + val n = l.narrow(axis = 0, from = 1, size = 2) + assertEquals(Shape(2, 8), n.shape); assertContentEquals(intArrayOf(8, 1), n.strides) + assertEquals(8L, n.offsetElements); assertEquals(32L, n.offsetBytes) + assertTrue(n.isContiguous) + val cols = l.narrow(axis = 1, from = 2, size = 3) + assertEquals(Shape(4, 3), cols.shape); assertEquals(2L, cols.offsetElements) + assertFalse(cols.isContiguous) // a column slice has gaps + assertFailsWith { l.narrow(1, 6, 4) } + } + + @Test + fun transposeSwapsExtentsAndStridesWithoutTouchingBytes() { + val l = Layout.rowMajor(Shape(2, 3), f32) + val t = l.transpose() + assertEquals(Shape(3, 2), t.shape); assertContentEquals(intArrayOf(1, 3), t.strides) + assertEquals(l.offsetElements, t.offsetElements) + assertFalse(t.isContiguous); assertFalse(t.isRowMajor) + assertEquals(l, t.transpose()) // involution + assertEquals(l.indexOf(1, 2), t.indexOf(2, 1)) + } + + @Test + fun unsqueezeAndSqueezeAreInverse() { + val l = Layout.rowMajor(Shape(4), f32) + val u = l.unsqueeze(0) + assertEquals(Shape(1, 4), u.shape); assertContentEquals(intArrayOf(0, 1), u.strides); assertTrue(u.isContiguous) + assertEquals(l, u.squeeze(0)) + assertEquals(Shape(4, 1), l.unsqueeze(1).shape) + assertFailsWith { Layout.rowMajor(Shape(2, 3), f32).squeeze(0) } + } + + @Test + fun blockedLayoutAddressesBlocksNotElements() { + // Q4_K: 256 elements per block, 144 bytes per block; a [4, 512] weight is [4, 2] blocks + val b = Layout.blocked(Shape(4, 512), blockSize = 256, bytesPerBlock = 144) + assertTrue(b.blocked); assertEquals(Shape(4, 2), b.shape) + assertContentEquals(intArrayOf(2, 1), b.strides); assertEquals(144, b.elementBytes) + assertEquals(144L * 5, b.byteOffsetOf(2, 1)) + val row = b.narrow(0, 1, 1) + assertEquals(2L, row.offsetElements); assertEquals(144L * 2, row.offsetBytes) + assertFailsWith { Layout.blocked(Shape(4, 300), 256, 144) } // not a whole number of blocks + } + + @Test + fun equalityAndRendering() { + val a = Layout.rowMajor(Shape(2, 2), f32); val b = Layout.rowMajor(Shape(2, 2), f32) + assertEquals(a, b); assertEquals(a.hashCode(), b.hashCode()) + assertTrue(a.toString().contains("contiguous")) + assertTrue(a.transpose().toString().contains("strided")) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorViewTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorViewTest.kt new file mode 100644 index 000000000..13e5ff8ba --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TensorViewTest.kt @@ -0,0 +1,130 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * SKEEP-003 rules 4–6: `get()` decodes (never a raw byte), views are zero-copy over the same + * storage, `materialize()` is the only copy point. + */ +@OptIn(ExperimentalMemoryApi::class) +class TensorViewTest { + + private fun denseView(vararg dims: Int): TensorView { + val shape = Shape(*dims) + val s = Storage.Heap.floats(shape.volume) + val f = s.floats!! + for (i in f.indices) f[i] = i.toFloat() + return TensorView.dense(s, shape, FP32, TensorId.parse("model.w")) + } + + @Test + fun denseViewReadsAndWritesElements() { + val v = denseView(2, 3) + assertEquals(Format.dense(FP32), v.format); assertTrue(v.isContiguous); assertTrue(v.isMutable) + assertEquals(0f, v.get(0, 0)); assertEquals(4f, v.get(1, 1)); assertEquals(5f, v.get(1, 2)) + v.set(0, 1, value = 9f) + assertEquals(9f, v.get(0, 1)) + assertContentEquals(floatArrayOf(0f, 9f, 2f, 3f, 4f, 5f), v.toFloatArray()) + assertFailsWith { v.get(2, 0) } + assertTrue(v.toString().contains("model.w")); assertTrue(v.toString().contains("Float32/Dense")) + } + + @Test + fun viewsShareStorageAndAreZeroCopy() { + val v = denseView(4, 4) + val rows = v.narrow(0, 1, 2) + assertEquals(Shape(2, 4), rows.shape); assertSame(v.storage, rows.storage) + assertEquals(4f, rows.get(0, 0)); assertEquals(11f, rows.get(1, 3)) + rows.set(0, 0, value = -1f) + assertEquals(-1f, v.get(1, 0)) // same bytes + val cols = v.narrow(1, 2, 2) + assertEquals(Shape(4, 2), cols.shape); assertFalse(cols.isContiguous) + assertEquals(2f, cols.get(0, 0)); assertEquals(7f, cols.get(1, 1)) + val t = v.transpose() + assertEquals(Shape(4, 4), t.shape); assertEquals(v.get(1, 2), t.get(2, 1)) + val u = v.narrow(0, 0, 1).squeeze(0) + assertEquals(Shape(4), u.shape); assertEquals(-1f, v.get(1, 0)) + assertEquals(Shape(1, 4), u.unsqueeze(0).shape) + assertEquals("model.w[1..3)]", rows.id!!.canonical) // TensorId.view brackets the range + } + + @Test + fun packedViewDecodesElementsNeverRawBytes() { + // a Q8_0 weight: 2 rows x 32 elements = 2 blocks of 34 bytes + val bytes = ByteArray(2 * 34) + fun half(v: Float): Int { val b = v.toRawBits(); val s = (b ushr 16) and 0x8000; val e = ((b ushr 23) and 0xFF) - 127 + 15; val m = b and 0x7FFFFF; return if (e <= 0) s else if (e >= 31) s or 0x7C00 else s or (e shl 10) or (m ushr 13) } + for (blk in 0 until 2) { + val off = blk * 34 + val d = half(0.5f); bytes[off] = (d and 0xFF).toByte(); bytes[off + 1] = ((d ushr 8) and 0xFF).toByte() + for (i in 0 until 32) bytes[off + 2 + i] = (i - 16).toByte() + } + val packed = Q8_0BlockTensorData(Shape(2, 32), bytes) + val storage = Storage.Heap.wrap(bytes) + val v = TensorView.packed(storage, Shape(2, 32), TensorEncoding.Q8_0, PackedBlockDecoder(packed), id = TensorId.parse("model.layers[0].attn.q_proj.weight")) + + assertEquals(Format(FP32, TensorEncoding.Q8_0), v.format) + assertFalse(v.format.isDense) + // rule 4: a decoded float, not the raw code byte + assertEquals(-16 * 0.5f, v.get(0, 0)); assertEquals(15 * 0.5f, v.get(0, 31)); assertEquals(-16 * 0.5f, v.get(1, 0)) + assertContentEquals(packed.toFloatArray(), v.toFloatArray()) + // writing through a packed view is refused + assertFailsWith { v.set(0, 0, value = 1f) } + } + + @Test + fun packedViewsSliceWholeBlocksWithoutTouchingBytes() { + val bytes = ByteArray(4 * 144) // 4 blocks of Q4_K + val packed = Q4_KBlockTensorData(Shape(2, 512), bytes) + val v = TensorView.packed(Storage.Heap.wrap(bytes), Shape(2, 512), TensorEncoding.Q4_K, PackedBlockDecoder(packed)) + assertTrue(v.layout.blocked); assertEquals(Shape(2, 2), v.layout.shape) // 2 rows x 2 blocks + val row = v.narrow(0, 1, 1) + assertEquals(Shape(1, 512), row.shape); assertSame(v.storage, row.storage) + assertEquals(2L, row.layout.offsetElements) // two blocks in + val half = v.narrow(1, 256, 256) // one block wide + assertEquals(Shape(2, 256), half.shape); assertEquals(1L, half.layout.offsetElements) + assertFailsWith { v.narrow(1, 0, 100) } // not a whole block + assertFailsWith { v.narrow(1, 100, 256) } // block-unaligned start + } + + @Test + fun materializeIsTheOnlyCopyPoint() { + val v = denseView(2, 2) + val scope = ForwardScope(64) + val m = v.materialize(scope = scope) + assertTrue(m.format.isDense); assertEquals(ScopeKind.FORWARD, m.storage.scope) + assertContentEquals(v.toFloatArray(), m.toFloatArray()) + m.set(0, 0, value = 42f) + assertEquals(0f, v.get(0, 0)) // a copy, not a view + // packed → dense materialization decodes + val bytes = ByteArray(144) + val packed = Q4_KBlockTensorData(Shape(1, 256), bytes) + val pv = TensorView.packed(Storage.Heap.wrap(bytes), Shape(1, 256), TensorEncoding.Q4_K, PackedBlockDecoder(packed)) + val dense = pv.materialize() + assertTrue(dense.format.isDense) + assertTrue(abs(dense.get(0, 5) - packed.toFloatArray()[5]) < 1e-6f) + scope.close() + } + + @Test + fun viewsOverClosedStorageAreRefused() { + val s = Storage.Heap.floats(4) + val v = TensorView.dense(s, Shape(4)) + s.close() + assertFailsWith { v.get(0) } + assertFailsWith { v.toFloatArray() } + assertFailsWith { TensorView.dense(s, Shape(4)) } + } +}