From 949e684fec9eb9303c268e19e229cea4198be030 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 10 Aug 2026 10:51:05 +0200 Subject: [PATCH] fix(lang): copyMaterialize handles Aliased and FileBacked; resolver overloads for transfer ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TensorStorage.copyMaterialize() threw UnsupportedOperationException for everything except Owned/Borrowed. That left two gaps: - Aliased: the descriptor level supports views (bounds-checked, mutability delegated to parent) and DefaultBufferResolver can already read them — but materializing one threw. It now resolves through the buffer-accessor path and produces an independent owned copy of the slice. - FileBacked: Placement.MMAP_WEIGHTS is one of the layer's three presets, yet copyToHost() -> copyMaterialize() could not read it — the one transfer the storage layer was designed around. copyMaterialize and copyToHost gain BufferResolver overloads; with a file-backed resolver configured (e.g. skainet-io-core's JVM resolver) mmap weights materialize to the heap. The no-arg forms keep their signatures and delegate to a DefaultBufferResolver, so Aliased now also works without any resolver. DeviceResident stays unsupported with an error that names the reason (no device backend exists yet) instead of a generic ownership message. Tests: aliased materialization (slice content + independence from parent), file-backed materialization and copyToHost through a fake resolver, helpful failure without a resolver, device-resident unchanged. Closes #929 --- CHANGELOG.md | 9 +++ .../lang/tensor/storage/TensorStorage.kt | 39 +++++++++-- .../lang/tensor/storage/TransferOpsTest.kt | 69 ++++++++++++++++++- 3 files changed, 111 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fca59a3c2..c9d47d134 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,15 @@ ## [Unreleased] +### 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) + ## [0.38.0] - 2026-07-30 ### Added 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