diff --git a/CHANGELOG.md b/CHANGELOG.md index 282f5aef0..78ba9816c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,12 @@ ### Fixed +- **`TensorStorage` transfer API can materialize its own placements.** `copyMaterialize()` + threw for `Aliased` handles (now resolved directly, producing an independent owned copy of + the slice) and for `FileBacked` — meaning `copyToHost()` could not bring `MMAP_WEIGHTS` + storage to the heap, the one transfer the storage layer was designed around. Both methods + gain a `BufferResolver` overload that reads file-backed regions through a configured + resolver; `DeviceResident` remains unsupported with an error that says why. (#929) - **Published Kotlin/Native klibs for `skainet-backend-native-cpu` now carry their machine code.** The static kernel archive was attached via project-local `linkerOpts`, which does not travel with a published klib — downstream K/N consumers of the `-linuxx64`/`-linuxarm64` @@ -40,7 +46,6 @@ linux static archives, and the macOS publish job injects them, so released klibs embed real code. In-repo K/N test binaries link purely from the embedded klib — the consumer-link scenario is what the test suite now exercises. (#941) - - **`TensorData.copyToFloatArray()` default implementation works for rank >= 2.** It used to iterate a single flat index into the vararg `get`, tripping every implementation's one-index-per-dimension arity check — a latent trap for any implementation that didn't diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt index b107b4c20..d50a4e77e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt @@ -72,14 +72,33 @@ public data class TensorStorage( /** * Create a new [TensorStorage] with an owned copy of this storage's data. * The returned storage is independent of the original buffer. + * + * Handles [BufferHandle.Owned], [BufferHandle.Borrowed] and + * [BufferHandle.Aliased] directly. [BufferHandle.FileBacked] needs a + * platform reader — use the [copyMaterialize] overload that takes a + * [BufferResolver] with a file-backed resolver configured. + */ + public fun copyMaterialize(): TensorStorage = copyMaterialize(DefaultBufferResolver()) + + /** + * Create a new [TensorStorage] with an owned copy of this storage's data, + * reading through [resolver] for handles that need platform support + * ([BufferHandle.Aliased] chains, [BufferHandle.FileBacked]). + * + * [BufferHandle.DeviceResident] remains unsupported until a device + * backend exists. */ - public fun copyMaterialize(): TensorStorage { + public fun copyMaterialize(resolver: BufferResolver): TensorStorage { val srcBytes = when (val b = buffer) { is BufferHandle.Owned -> b.data.copyOfRange(b.offset, b.offset + sizeBytes()) is BufferHandle.Borrowed -> b.data.copyOfRange(b.offset, b.offset + sizeBytes()) - else -> throw UnsupportedOperationException( - "copyMaterialize not yet supported for ${buffer.ownership} buffers" + is BufferHandle.DeviceResident -> throw UnsupportedOperationException( + "copyMaterialize cannot read device-resident buffer '${b.deviceId}' on the host — " + + "no device backend exists yet" ) + // Aliased and FileBacked: the resolver knows how to read them. + // readBytes always copies, so the result is a genuinely owned buffer. + else -> resolver.resolve(b).use { it.readBytes(0, sizeBytes()) } } return copy( buffer = BufferHandle.Owned(srcBytes), @@ -90,10 +109,20 @@ public data class TensorStorage( /** * Ensure this storage resides on the host (CPU heap). * If already on host, returns `this`. Otherwise copies to host. + * + * For [BufferHandle.FileBacked] storage (e.g. [Placement.MMAP_WEIGHTS]), + * use the overload that takes a [BufferResolver] with a file-backed + * resolver configured. + */ + public fun copyToHost(): TensorStorage = copyToHost(DefaultBufferResolver()) + + /** + * Ensure this storage resides on the host (CPU heap), reading through + * [resolver] for handles that need platform support. */ - public fun copyToHost(): TensorStorage { + public fun copyToHost(resolver: BufferResolver): TensorStorage { if (placement.device == DeviceKind.CPU && placement.domain == MemoryDomain.HOST_HEAP) return this - return copyMaterialize() + return copyMaterialize(resolver) } /** diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TransferOpsTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TransferOpsTest.kt index 06ef995c1..bf9f004a7 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TransferOpsTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/TransferOpsTest.kt @@ -55,7 +55,53 @@ class TransferOpsTest { } @Test - fun copyMaterialize_fileBackedBuffer_throwsUnsupported() { + fun copyMaterialize_aliasedBuffer_producesOwnedCopyOfTheSlice() { + // Aliased handles are resolvable without any platform support — + // copyMaterialize must not throw for them (#929). + val parentBytes = ByteArray(32) { it.toByte() } + val parent = BufferHandle.Owned(parentBytes) + val storage = TensorStorage( + shape = Shape(4), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.Aliased(parent, byteOffset = 8, sizeInBytes = 16), + placement = Placement.CPU_HEAP + ) + val copy = storage.copyMaterialize() + + assertEquals(Ownership.OWNED, copy.ownership) + val copyData = (copy.buffer as BufferHandle.Owned).data + assertEquals(16, copyData.size) + assertEquals(8, copyData[0]) // slice starts at parent byte 8 + // Independent of the parent after materialization. + parentBytes[8] = 99 + assertEquals(8, copyData[0]) + } + + @Test + fun copyMaterialize_fileBackedBuffer_readsThroughResolver() { + // FileBacked is materializable when the caller supplies a resolver + // that can read the file region — the core transfer #929 restored. + val fileBytes = ByteArray(16) { (it + 100).toByte() } + val resolver = DefaultBufferResolver(fileBackedResolver = { handle -> + ByteArrayAccessor(fileBytes, handle.fileOffset.toInt(), handle.sizeInBytes) + }) + val storage = TensorStorage( + shape = Shape(4), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.FileBacked("/model.bin", 0, 16), + placement = Placement.MMAP_WEIGHTS + ) + val copy = storage.copyMaterialize(resolver) + + assertEquals(Ownership.OWNED, copy.ownership) + assertEquals(MemoryDomain.HOST_HEAP, copy.placement.domain) + assertEquals(100, (copy.buffer as BufferHandle.Owned).data[0]) + } + + @Test + fun copyMaterialize_fileBackedBuffer_withoutResolver_throwsHelpfully() { val storage = TensorStorage( shape = Shape(4), logicalType = LogicalDType.FLOAT32, @@ -81,6 +127,27 @@ class TransferOpsTest { } } + @Test + fun copyToHost_fileBackedMmapWeights_materializesThroughResolver() { + // Placement.MMAP_WEIGHTS is the layer's own preset for file-backed + // weights; copyToHost(resolver) must be able to bring it to the heap. + val fileBytes = ByteArray(16) { it.toByte() } + val resolver = DefaultBufferResolver(fileBackedResolver = { handle -> + ByteArrayAccessor(fileBytes, handle.fileOffset.toInt(), handle.sizeInBytes) + }) + val storage = TensorStorage( + shape = Shape(4), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.FileBacked("/model.bin", 0, 16), + placement = Placement.MMAP_WEIGHTS + ) + val hosted = storage.copyToHost(resolver) + + assertEquals(Ownership.OWNED, hosted.ownership) + assertEquals(MemoryDomain.HOST_HEAP, hosted.placement.domain) + } + // --- copyToHost --- @Test