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 eb71108de..875b4ba58 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 @@ -676,6 +676,26 @@ public final class sk/ainet/lang/memory/DescribeKt { public static synthetic fun describe$default (Lsk/ainet/lang/tensor/storage/TensorStorage;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Ljava/lang/String; } +public final class sk/ainet/lang/memory/DirectBufferStorage : sk/ainet/lang/memory/Storage$OffHeap { + public static final field Companion Lsk/ainet/lang/memory/DirectBufferStorage$Companion; + public synthetic fun (JLjava/nio/ByteBuffer;Lsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;Lkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun buffer ()Ljava/nio/ByteBuffer; + public fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public fun getId-TPZW6QE ()J + public fun getOwner ()Lsk/ainet/lang/memory/Owner; + public fun getSizeBytes ()J + public fun isMutable ()Z + public fun slice (JJ)Lsk/ainet/lang/memory/DirectBufferStorage; + public synthetic fun slice (JJ)Lsk/ainet/lang/memory/Storage; +} + +public final class sk/ainet/lang/memory/DirectBufferStorage$Companion { + public final fun allocate (ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/DirectBufferStorage; + public static synthetic fun allocate$default (Lsk/ainet/lang/memory/DirectBufferStorage$Companion;ILsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/DirectBufferStorage; + public final fun borrow (Ljava/nio/ByteBuffer;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/DirectBufferStorage; + public static synthetic fun borrow$default (Lsk/ainet/lang/memory/DirectBufferStorage$Companion;Ljava/nio/ByteBuffer;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/DirectBufferStorage; +} + public abstract interface annotation class sk/ainet/lang/memory/ExperimentalMemoryApi : java/lang/annotation/Annotation { } @@ -705,6 +725,47 @@ public final class sk/ainet/lang/memory/FormatKt { public static final fun getFormatOrNull (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/memory/Format; } +public final class sk/ainet/lang/memory/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 + public final fun buffer ()Ljava/nio/ByteBuffer; + public fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public final fun getFileOffset ()J + public fun getId-TPZW6QE ()J + public fun getOwner ()Lsk/ainet/lang/memory/Owner; + public final fun getPath ()Ljava/nio/file/Path; + public fun getSizeBytes ()J + public fun isMutable ()Z + public fun slice (JJ)Lsk/ainet/lang/memory/MappedBufferStorage; + public synthetic fun slice (JJ)Lsk/ainet/lang/memory/Storage; +} + +public final class sk/ainet/lang/memory/MappedBufferStorage$Companion { + public final fun map (Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/MappedBufferStorage; + public static synthetic fun map$default (Lsk/ainet/lang/memory/MappedBufferStorage$Companion;Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/MappedBufferStorage; +} + +public final class sk/ainet/lang/memory/MappedFileStorage : sk/ainet/lang/memory/Storage$Mapped { + public static final field Companion Lsk/ainet/lang/memory/MappedFileStorage$Companion; + public synthetic fun (JLjava/nio/file/Path;JLjava/lang/foreign/MemorySegment;Lsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/foreign/Arena;Lkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public final fun getFileOffset ()J + public fun getId-TPZW6QE ()J + public fun getOwner ()Lsk/ainet/lang/memory/Owner; + public final fun getPath ()Ljava/nio/file/Path; + public fun getSizeBytes ()J + public fun isMutable ()Z + public final fun segment ()Ljava/lang/foreign/MemorySegment; + public fun slice (JJ)Lsk/ainet/lang/memory/MappedFileStorage; + public synthetic fun slice (JJ)Lsk/ainet/lang/memory/Storage; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/MappedFileStorage$Companion { + public final fun map (Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/MappedFileStorage; + public static synthetic fun map$default (Lsk/ainet/lang/memory/MappedFileStorage$Companion;Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/MappedFileStorage; +} + public abstract interface class sk/ainet/lang/memory/Owner { } @@ -747,6 +808,29 @@ public final class sk/ainet/lang/memory/ScopeKind : java/lang/Enum { public static fun values ()[Lsk/ainet/lang/memory/ScopeKind; } +public final class sk/ainet/lang/memory/SegmentStorage : sk/ainet/lang/memory/Storage$OffHeap { + public static final field Companion Lsk/ainet/lang/memory/SegmentStorage$Companion; + public synthetic fun (JLjava/lang/foreign/MemorySegment;Lsk/ainet/lang/memory/Owner;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/foreign/Arena;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun getDebugOrigin ()Lsk/ainet/lang/tensor/TensorId; + public fun getId-TPZW6QE ()J + public fun getOwner ()Lsk/ainet/lang/memory/Owner; + public fun getSizeBytes ()J + public fun isMutable ()Z + public final fun segment ()Ljava/lang/foreign/MemorySegment; + public fun slice (JJ)Lsk/ainet/lang/memory/SegmentStorage; + public synthetic fun slice (JJ)Lsk/ainet/lang/memory/Storage; +} + +public final class sk/ainet/lang/memory/SegmentStorage$Companion { + public final fun allocate (JLsk/ainet/lang/memory/ScopeKind;Ljava/lang/foreign/Arena;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/SegmentStorage; + public static synthetic fun allocate$default (Lsk/ainet/lang/memory/SegmentStorage$Companion;JLsk/ainet/lang/memory/ScopeKind;Ljava/lang/foreign/Arena;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/SegmentStorage; + public final fun borrow (Ljava/lang/foreign/MemorySegment;ZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/SegmentStorage; + public final fun borrow ([FLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;)Lsk/ainet/lang/memory/SegmentStorage; + public static synthetic fun borrow$default (Lsk/ainet/lang/memory/SegmentStorage$Companion;Ljava/lang/foreign/MemorySegment;ZLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/SegmentStorage; + public static synthetic fun borrow$default (Lsk/ainet/lang/memory/SegmentStorage$Companion;[FLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/SegmentStorage; + public final fun getFLOAT ()Ljava/lang/foreign/ValueLayout$OfFloat; +} + public abstract class sk/ainet/lang/memory/Storage : java/lang/AutoCloseable { public final fun checkAlive ()V public final fun close ()V diff --git a/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/memory/DirectBufferStorage.kt b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/memory/DirectBufferStorage.kt new file mode 100644 index 000000000..6e163108b --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmAndroidMain/kotlin/sk/ainet/lang/memory/DirectBufferStorage.kt @@ -0,0 +1,98 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.tensor.TensorId +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.MappedByteBuffer +import java.nio.channels.FileChannel +import java.nio.file.Path +import java.nio.file.StandardOpenOption + +/** + * Android (and JVM fallback) binding of [Storage.OffHeap]: a direct `ByteBuffer` — outside the ART + * heap and its per-app limit (the root of #922), counted by the OS against the process. Freed when + * the buffer becomes unreachable (direct buffers have no explicit free on Android); [close] marks + * the storage dead so no late access can see it. Use [SegmentStorage] where FFM is available. + */ +@ExperimentalMemoryApi +public class DirectBufferStorage private constructor( + override val id: StorageId, + private val buf: ByteBuffer, + override val owner: Owner, + override val debugOrigin: TensorId?, + override val sink: TraceSink, +) : Storage.OffHeap() { + override val sizeBytes: Long get() = buf.capacity().toLong() + override val isMutable: Boolean get() = !buf.isReadOnly + + /** An independent little-endian duplicate of the buffer (position 0); kernels take it once per call. */ + public fun buffer(): ByteBuffer { checkAlive(); return buf.duplicate().order(ByteOrder.LITTLE_ENDIAN) } + + override fun slice(offsetBytes: Long, lengthBytes: Long): DirectBufferStorage { + checkAlive() + require(offsetBytes >= 0 && lengthBytes >= 0 && offsetBytes + lengthBytes <= sizeBytes) { "slice [$offsetBytes, ${offsetBytes + lengthBytes}) outside $sizeBytes bytes" } + // Java 8 signatures on Android: Buffer.position/limit return Buffer, so keep the ByteBuffer typed. + val d: ByteBuffer = buf.duplicate() + d.position(offsetBytes.toInt()); d.limit((offsetBytes + lengthBytes).toInt()) + val s: ByteBuffer = d.slice().order(ByteOrder.LITTLE_ENDIAN) + return DirectBufferStorage(StorageId.next(), s, Owner.Alias(this), debugOrigin, sink) + } + + public companion object { + /** Allocate [bytes] zeroed direct bytes owned by a scope of kind [scope]. */ + public fun allocate(bytes: Int, scope: ScopeKind = ScopeKind.AMBIENT, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): DirectBufferStorage { + require(bytes >= 0) { "bytes must be >= 0" } + val s = DirectBufferStorage(StorageId.next(), ByteBuffer.allocateDirect(bytes).order(ByteOrder.LITTLE_ENDIAN), Owner.Owned(scope), origin, sink) + if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, scope, bytes.toLong(), origin)) + return s + } + + /** Borrow a caller's buffer (direct or mapped) — never freed by us. */ + public fun borrow(buffer: ByteBuffer, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): DirectBufferStorage = + DirectBufferStorage(StorageId.next(), buffer.duplicate().order(ByteOrder.LITTLE_ENDIAN), Owner.Borrowed(buffer), origin, sink) + } +} + +/** + * Android (and JVM fallback) binding of [Storage.Mapped]: a `MappedByteBuffer` from `FileChannel.map` + * — weights outside ART entirely (SKEEP-002 / #921). Unmapped when the buffer becomes unreachable; + * [close] marks the storage dead. + */ +@ExperimentalMemoryApi +public class MappedBufferStorage private constructor( + override val id: StorageId, + public val path: Path, + public val fileOffset: Long, + private val buf: ByteBuffer, + override val owner: Owner, + override val debugOrigin: TensorId?, + override val sink: TraceSink, +) : Storage.Mapped() { + override val sizeBytes: Long get() = buf.capacity().toLong() + override val isMutable: Boolean get() = false + + public fun buffer(): ByteBuffer { checkAlive(); return buf.duplicate().order(ByteOrder.LITTLE_ENDIAN) } + + override fun slice(offsetBytes: Long, lengthBytes: Long): MappedBufferStorage { + checkAlive() + require(offsetBytes >= 0 && lengthBytes >= 0 && offsetBytes + lengthBytes <= sizeBytes) { "slice [$offsetBytes, ${offsetBytes + lengthBytes}) outside $sizeBytes bytes" } + val d: ByteBuffer = buf.duplicate() + d.position(offsetBytes.toInt()); d.limit((offsetBytes + lengthBytes).toInt()) + val s: ByteBuffer = d.slice().order(ByteOrder.LITTLE_ENDIAN) + return MappedBufferStorage(StorageId.next(), path, fileOffset + offsetBytes, s, Owner.Alias(this), debugOrigin, sink) + } + + public companion object { + /** Map `[fileOffset, fileOffset + length)` of [path] read-only (length ≤ 2 GB per buffer). */ + public fun map(path: Path, fileOffset: Long, length: Long, scope: ScopeKind = ScopeKind.MODEL, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): MappedBufferStorage { + require(fileOffset >= 0 && length in 0..Int.MAX_VALUE.toLong()) { "offset must be >= 0 and length in [0, 2 GB)" } + val mbb: MappedByteBuffer = FileChannel.open(path, StandardOpenOption.READ).use { ch -> ch.map(FileChannel.MapMode.READ_ONLY, fileOffset, length) } + val s = MappedBufferStorage(StorageId.next(), path, fileOffset, mbb.order(ByteOrder.LITTLE_ENDIAN), Owner.Owned(scope), origin, sink) + if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, scope, length, origin, site = path.toString())) + return s + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/JvmStorage.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/JvmStorage.kt new file mode 100644 index 000000000..aca3b77fd --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/JvmStorage.kt @@ -0,0 +1,119 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.tensor.TensorId +import java.lang.foreign.Arena +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.nio.channels.FileChannel +import java.nio.file.Path +import java.nio.file.StandardOpenOption + +/** + * JVM binding of [Storage.OffHeap]: a `MemorySegment` (FFM) that is not scanned or copied by the GC, + * may exceed 2 GB, is handed to IREE/JNI zero-copy, and is freed deterministically by the `Arena` + * that owns it (SKEEP-003 §4.8.1). An owned storage allocated without an explicit arena gets its + * own `Arena.ofShared()` and closes it on [close]; milestone slice #1021 (`Scope`) passes the + * scope's arena instead, so `Forward` becomes a recycled bump slab. + */ +@ExperimentalMemoryApi +public class SegmentStorage private constructor( + override val id: StorageId, + private val seg: MemorySegment, + override val owner: Owner, + override val debugOrigin: TensorId?, + override val sink: TraceSink, + private val ownedArena: Arena?, + private val mutable: Boolean, +) : Storage.OffHeap() { + override val sizeBytes: Long get() = seg.byteSize() + override val isMutable: Boolean get() = (owner as? Owner.Alias)?.parent?.isMutable ?: mutable + + /** The segment — kernels take it once per call (`ByteVector.fromMemorySegment`, `getAtIndex`). Throws when closed. */ + public fun segment(): MemorySegment { checkAlive(); return seg } + + override fun slice(offsetBytes: Long, lengthBytes: Long): SegmentStorage { + checkAlive() + require(offsetBytes >= 0 && lengthBytes >= 0 && offsetBytes + lengthBytes <= sizeBytes) { "slice [$offsetBytes, ${offsetBytes + lengthBytes}) outside $sizeBytes bytes" } + return SegmentStorage(StorageId.next(), seg.asSlice(offsetBytes, lengthBytes), Owner.Alias(this), debugOrigin, sink, null, mutable) + } + + override fun onClose() { ownedArena?.close() } + + public companion object { + /** + * Allocate [bytes] zeroed off-heap bytes (aligned to [alignment]) owned by a scope of kind [scope]. + * With [arena] `null` the storage owns a private `Arena.ofShared()`; pass the scope's arena to + * let the scope free it (the `Forward` slab pattern). + */ + public fun allocate(bytes: Long, scope: ScopeKind = ScopeKind.AMBIENT, arena: Arena? = null, alignment: Long = 64L, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): SegmentStorage { + require(bytes >= 0) { "bytes must be >= 0" } + val owned = arena ?: Arena.ofShared() + val seg = owned.allocate(bytes, alignment) + val s = SegmentStorage(StorageId.next(), seg, Owner.Owned(scope), origin, sink, if (arena == null) owned else null, true) + if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, scope, bytes, origin)) + return s + } + + /** Borrow an existing segment (a loader's, IREE's, a caller's `MemorySegment.ofArray`) — never freed by us. */ + public fun borrow(segment: MemorySegment, mutable: Boolean = !segment.isReadOnly, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): SegmentStorage = + SegmentStorage(StorageId.next(), segment, Owner.Borrowed(segment), origin, sink, null, mutable) + + /** Borrow a heap array zero-copy as a segment (`MemorySegment.ofArray`). */ + public fun borrow(array: FloatArray, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): SegmentStorage = + borrow(MemorySegment.ofArray(array), mutable = true, origin = origin, sink = sink) + + /** Layout helper for float element access on a segment. */ + public val FLOAT: ValueLayout.OfFloat = ValueLayout.JAVA_FLOAT + } +} + +/** + * JVM binding of [Storage.Mapped]: a read-only region of a file mapped with `FileChannel.map` into + * an `Arena.ofShared()`. The OS pages the bytes; resident set = pages touched; the page cache is + * shared across processes; closing the storage unmaps (SKEEP-003 §4.8.1). Packed GGUF weights and + * embedding tables live here. + */ +@ExperimentalMemoryApi +public class MappedFileStorage private constructor( + override val id: StorageId, + public val path: Path, + public val fileOffset: Long, + private val seg: MemorySegment, + override val owner: Owner, + override val debugOrigin: TensorId?, + override val sink: TraceSink, + private val arena: Arena?, +) : Storage.Mapped() { + override val sizeBytes: Long get() = seg.byteSize() + override val isMutable: Boolean get() = false + + public fun segment(): MemorySegment { checkAlive(); return seg } + + override fun slice(offsetBytes: Long, lengthBytes: Long): MappedFileStorage { + checkAlive() + require(offsetBytes >= 0 && lengthBytes >= 0 && offsetBytes + lengthBytes <= sizeBytes) { "slice [$offsetBytes, ${offsetBytes + lengthBytes}) outside $sizeBytes bytes" } + return MappedFileStorage(StorageId.next(), path, fileOffset + offsetBytes, seg.asSlice(offsetBytes, lengthBytes), Owner.Alias(this), debugOrigin, sink, null) + } + + override fun onClose() { arena?.close() } + + override fun toString(): String = "Mapped(${id}, ${sizeBytes} B, $path @0x${fileOffset.toString(16)}${debugOrigin?.let { ", $it" } ?: ""}${if (isAlive) "" else ", closed"})" + + public companion object { + /** + * Map `[fileOffset, fileOffset + length)` of [path] read-only. Owned by a scope of kind [scope] + * (normally `MODEL`): closing the storage — or, from #1021, the scope — unmaps. + */ + public fun map(path: Path, fileOffset: Long, length: Long, scope: ScopeKind = ScopeKind.MODEL, origin: TensorId? = null, sink: TraceSink = NoopTraceSink): MappedFileStorage { + require(fileOffset >= 0 && length >= 0) { "offset/length must be >= 0" } + val arena = Arena.ofShared() + val seg = FileChannel.open(path, StandardOpenOption.READ).use { ch -> ch.map(FileChannel.MapMode.READ_ONLY, fileOffset, length, arena) } + val s = MappedFileStorage(StorageId.next(), path, fileOffset, seg, Owner.Owned(scope), origin, sink, arena) + if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, scope, length, origin, site = path.toString())) + return s + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt new file mode 100644 index 000000000..e7352d4a1 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/memory/JvmStorageTest.kt @@ -0,0 +1,129 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.MemoryDomain +import java.lang.foreign.Arena +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertIs +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** SKEEP-003 §4.8.1/§4.8.2: OffHeap (MemorySegment / direct ByteBuffer) and Mapped (FileChannel.map) storage kinds. */ +@OptIn(ExperimentalMemoryApi::class) +class JvmStorageTest { + + @Test + fun segmentStorageOwnsItsArenaAndFreesOnClose() { + val sink = RecordingTraceSink() + val id = TensorId.parse("model.layers[0].attn.q#step=3") + val s = SegmentStorage.allocate(bytes = 1024, scope = ScopeKind.FORWARD, origin = id, sink = sink) + assertEquals(MemoryDomain.HOST_OFFHEAP, s.domain); assertEquals(ScopeKind.FORWARD, s.scope); assertTrue(s.isMutable) + assertEquals(1024L, s.sizeBytes); assertEquals(0L, s.segment().address() % 64) + s.segment().setAtIndex(ValueLayout.JAVA_FLOAT, 3, 1.5f) + assertEquals(1.5f, s.segment().getAtIndex(ValueLayout.JAVA_FLOAT, 3)) + val alloc = assertIs(sink.events().single()); assertEquals(1024L, alloc.bytes); assertEquals(id, alloc.origin) + s.close() + assertFalse(s.isAlive) + assertFailsWith { s.segment() } + assertIs(sink.events()[1]) + } + + @Test + fun segmentStorageInACallerArenaIsFreedByTheArenaNotByClose() { + Arena.ofShared().use { arena -> + val s = SegmentStorage.allocate(256, ScopeKind.MODEL, arena = arena) + val seg = s.segment() + s.close() // marks dead, does not close the caller's arena + assertTrue(seg.scope().isAlive) + assertFailsWith { s.segment() } + } + } + + @Test + fun segmentSliceIsAnAliasOverTheSameBytes() { + val s = SegmentStorage.allocate(64) + val v = s.slice(16, 32) + assertIs(v.owner); assertEquals(32L, v.sizeBytes) + v.segment().setAtIndex(ValueLayout.JAVA_FLOAT, 0, 9f) + assertEquals(9f, s.segment().getAtIndex(ValueLayout.JAVA_FLOAT, 4)) + assertFailsWith { s.slice(60, 8) } + s.close(); assertFalse(v.isAlive) + } + + @Test + fun borrowedSegmentAndArrayAreNeverFreed() { + val arr = FloatArray(8) { it.toFloat() } + val s = SegmentStorage.borrow(arr) + assertIs(s.owner); assertEquals(32L, s.sizeBytes); assertTrue(s.isMutable) + s.segment().setAtIndex(ValueLayout.JAVA_FLOAT, 1, 42f) + assertEquals(42f, arr[1]) // zero-copy over the caller's array + s.close() + assertEquals(42f, arr[1]) + val ro = SegmentStorage.borrow(MemorySegment.ofArray(IntArray(2)).asReadOnly()) + assertFalse(ro.isMutable) + } + + @Test + fun mappedFileStorageReadsTheFileAndUnmapsOnClose() { + val f = Files.createTempFile("skainet-mapped", ".bin"); f.toFile().deleteOnExit() + val bytes = ByteArray(256) { it.toByte() }; Files.write(f, bytes) + val sink = RecordingTraceSink() + val m = MappedFileStorage.map(f, fileOffset = 16, length = 64, origin = TensorId.parse("model.w"), sink = sink) + assertEquals(MemoryDomain.MMAP_FILE, m.domain); assertEquals(ScopeKind.MODEL, m.scope); assertFalse(m.isMutable) + assertEquals(64L, m.sizeBytes); assertEquals(16L, m.fileOffset) + assertEquals(16.toByte(), m.segment().get(ValueLayout.JAVA_BYTE, 0)); assertEquals(79.toByte(), m.segment().get(ValueLayout.JAVA_BYTE, 63)) + val v = m.slice(8, 8); assertEquals(24L, v.fileOffset); assertEquals(24.toByte(), v.segment().get(ValueLayout.JAVA_BYTE, 0)) + assertTrue(m.toString().startsWith("Mapped(#")); assertTrue(m.toString().contains("@0x10")) + assertEquals(f.toString(), assertIs(sink.events().single()).site) + m.close() + assertFalse(m.isAlive); assertFalse(v.isAlive) + assertFailsWith { m.segment() } + } + + @Test + fun directBufferStorageAllocatesBorrowsAndSlices() { + val sink = RecordingTraceSink() + val d = DirectBufferStorage.allocate(64, ScopeKind.FORWARD, sink = sink) + assertEquals(MemoryDomain.HOST_OFFHEAP, d.domain); assertTrue(d.isMutable); assertTrue(d.buffer().isDirect) + assertEquals(ByteOrder.LITTLE_ENDIAN, d.buffer().order()) + d.buffer().putFloat(8, 2.5f) + assertEquals(2.5f, d.buffer().getFloat(8)) + val v = d.slice(8, 8); assertEquals(2.5f, v.buffer().getFloat(0)); assertIs(v.owner) + assertEquals(64L, assertIs(sink.events().single()).bytes) + val b = DirectBufferStorage.borrow(ByteBuffer.allocate(16).asReadOnlyBuffer()) + assertFalse(b.isMutable); assertIs(b.owner) + d.close(); assertFalse(v.isAlive) + assertFailsWith { d.buffer() } + } + + @Test + fun mappedBufferStorageReadsTheFile() { + val f = Files.createTempFile("skainet-mapped-bb", ".bin"); f.toFile().deleteOnExit() + Files.write(f, ByteArray(128) { (it * 2).toByte() }) + val m = MappedBufferStorage.map(f, 32, 32) + assertEquals(64.toByte(), m.buffer().get(0)); assertEquals(32L, m.sizeBytes); assertFalse(m.isMutable) + assertEquals(ScopeKind.MODEL, m.scope) + val v = m.slice(4, 4); assertEquals(36L, v.fileOffset); assertEquals(72.toByte(), v.buffer().get(0)) + m.close(); assertFailsWith { m.buffer() } + } + + @Test + fun allKindsAreStoragesWithDistinctIds() { + val h = Storage.Heap.floats(1); val s = SegmentStorage.allocate(4); val d = DirectBufferStorage.allocate(4) + val ids = listOf(h.id.value, s.id.value, d.id.value) + assertEquals(3, ids.toSet().size) + assertEquals(Storage.Heap::class, h::class) + assertTrue(s is Storage.OffHeap && d is Storage.OffHeap) + listOf(h, s, d).forEach { it.close() } + } +}