From 8fde5e089d1b1574c573e3d6084b49ef136a2ad1 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 14:24:06 +0200 Subject: [PATCH 1/2] test(storage): restore MemoryTrackerTest lost with MemoryPlannerTest.kt's deletion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #1142 deleted MemoryPlannerTest.kt for the dead planner it tested — but the file also housed MemoryTrackerTest, whose four unit tests (trackAndReport, trackCopies, clearResetsState, fileBackedTracking) cover the live MemoryTracker aggregate reports. Restored verbatim in its own file; only the planner-owned imports are dropped. Part of #1146's memory-accounting groundwork. Co-Authored-By: Claude Fable 5 --- .../lang/tensor/storage/MemoryTrackerTest.kt | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/MemoryTrackerTest.kt diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/MemoryTrackerTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/MemoryTrackerTest.kt new file mode 100644 index 000000000..f75b9b90b --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/MemoryTrackerTest.kt @@ -0,0 +1,82 @@ +@file:Suppress("DEPRECATION") // LogicalDType legacy path kept under test until removal (SKEEP-003 #1014) + +package sk.ainet.lang.tensor.storage + +import sk.ainet.lang.tensor.Shape +import kotlin.test.Test +import kotlin.test.assertEquals + +class MemoryTrackerTest { + + @Test + fun trackAndReport() { + val tracker = MemoryTracker() + + val s1 = TensorStorage( + shape = Shape(100), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.Owned(ByteArray(400)) + ) + val s2 = TensorStorage( + shape = Shape(256), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Q4_K, + buffer = BufferHandle.Borrowed(ByteArray(144)) + ) + + tracker.record("weight1", s1) + tracker.record("weight2_q4k", s2) + + val report = tracker.report() + assertEquals(2, report.tensorCount) + assertEquals(1, report.ownedCount) + assertEquals(1, report.borrowedCount) + assertEquals(400L + 1024L, report.totalLogicalBytes) // 100*4 + 256*4 + assertEquals(400L + 144L, report.totalPhysicalBytes) + } + + @Test + fun trackCopies() { + val tracker = MemoryTracker() + tracker.recordCopy("tensor_a", 1024) + tracker.recordCopy("tensor_b", 2048) + + val report = tracker.report() + assertEquals(2L, report.copyCount) + assertEquals(3072L, report.copyBytes) + } + + @Test + fun clearResetsState() { + val tracker = MemoryTracker() + tracker.record("x", TensorStorage( + shape = Shape(10), + logicalType = LogicalDType.FLOAT32, + encoding = TensorEncoding.Dense(4), + buffer = BufferHandle.Owned(ByteArray(40)) + )) + tracker.recordCopy("x", 40) + tracker.clear() + + val report = tracker.report() + assertEquals(0, report.tensorCount) + assertEquals(0L, report.copyCount) + } + + @Test + fun fileBackedTracking() { + val tracker = MemoryTracker() + tracker.record("mmap_weight", TensorStorage( + shape = Shape(1000), + logicalType = LogicalDType.FLOAT16, + encoding = TensorEncoding.Dense(2), + buffer = BufferHandle.FileBacked("/model.bin", 0, 2000), + placement = Placement.MMAP_WEIGHTS + )) + + val report = tracker.report() + assertEquals(1, report.fileBackedCount) + assertEquals(2000L, report.fileBackedBytes) + } +} From 791c3365bca9d637cc5c12c00d252b48fa8d6749 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 14:42:21 +0200 Subject: [PATCH 2/2] feat(memory): op outputs draw from the active scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The creation path learned to consult ExecutionContext.memoryScope in #1145; op outputs still reached the GC because DefaultCpuOps allocates them through its TensorDataFactory. This closes that half (#1146): - TensorDataFactory.adoptFloatArray — the op-output entry point. Unlike wrapFloatArray, the buffer is a fresh, ownership-transferred result, which licenses a factory to relocate it. DenseTensorDataFactory adopts FP32/FP16 zero-copy (one copy FEWER than before on the common path, which paid fromFloatArray's copyOf); narrow dtypes keep their tagged fromFloatArray path. - ScopedTensorDataFactory decorates any factory: under a non-Ambient scope, dense-FP32 zeros/ones/full/init/fromFloatArray/adoptFloatArray land in the slab as StorageFloatTensorData; the region is always fully written because a reset slab is dirty. wrap* stays un-intercepted — its zero-copy caller-owned contract must not die at reset(). - ExecutionContext.withTensorDataFactory + the DirectCpuExecutionContext override let ScopedExecutionContext rebuild the base around the scoped factory and re-bind created tensors to the rebuilt ops, so 'a + b' dispatches into scope-allocating ops. - The FP32 fast paths are offset-aware (floatWindowOf): slab-backed operands stay on the flat primitive loops instead of falling to the boxed generic path (#949's 83%-overhead cliff — live since #1145 for scope-created tensors, now closed). Tail sites read slab data as its exact logical window through floatBufferOf (one copy, still far off the boxed path). DefaultCpuOpsJvm's 18 direct constructions route through floatResult; its Panama vector paths fall back to the common window loops for slab operands (follow-up: offset-aware vectors). - KernelDispatch.matmul receives dispatchScope(), so requantize/prepack adapter allocations land in the slab too. - ScopedOpOutputsTest pins: slab-backed outputs, flat peakFloats across 8 steps with zero overflow (steady-state decode allocates no new slab bytes), bit-identical numerics vs Ambient, loud StorageClosedException on stale reads, and the overflow path. Also restores MemoryTrackerTest (lost with MemoryPlannerTest.kt in #1142). Closes #1146. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-backend-cpu.api | 3 + .../context/DirectCpuExecutionContext.kt | 4 + .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 146 ++++++++++++------ .../exec/tensor/ops/ScopedOpOutputsTest.kt | 115 ++++++++++++++ .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 70 +++------ .../api/jvm/skainet-compile-dag.api | 2 + .../api/jvm/skainet-lang-core.api | 32 ++++ .../sk/ainet/context/ExecutionContext.kt | 11 ++ .../ainet/context/ScopedExecutionContext.kt | 30 ++++ .../tensor/data/DenseTensorDataFactory.kt | 11 ++ .../tensor/data/ScopedTensorDataFactory.kt | 99 ++++++++++++ .../lang/tensor/data/TensorDataFactory.kt | 16 ++ 12 files changed, 437 insertions(+), 102 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedOpOutputsTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ScopedTensorDataFactory.kt diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index c89718f3a..d146896e0 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -30,6 +30,7 @@ public final class sk/ainet/context/DirectCpuExecutionContext : sk/ainet/context public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -249,6 +250,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun convTranspose1d (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IIIII)Lsk/ainet/lang/tensor/Tensor; public fun convert (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/Tensor; public fun cos (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + protected final fun dispatchScope ()Lsk/ainet/lang/memory/Scope; public fun divScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun divide (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; protected final fun elementwise (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/jvm/functions/Function3;)Lsk/ainet/lang/tensor/Tensor; @@ -257,6 +259,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; + protected final fun floatResult (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F[Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun gather (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt index ffc8f16ea..843eae4c2 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/context/DirectCpuExecutionContext.kt @@ -48,4 +48,8 @@ public class DirectCpuExecutionContext @kotlin.jvm.JvmOverloads constructor( override val hooks: sk.ainet.lang.nn.hooks.ForwardHooks? get() = _hooks + + /** A sibling context whose cached ops allocate through [factory] (#1146). */ + override fun withTensorDataFactory(factory: TensorDataFactory): ExecutionContext = + DirectCpuExecutionContext(executionStats, phase, _hooks, factory) } diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 15e2a531a..2f7b15b37 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -191,17 +191,52 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory // `inline` is load-bearing: a non-inlined `(Float, Float) -> Float` lambda // would box through `Function2` and reintroduce the very churn removed. - private fun floatBufferOf(t: Tensor): FloatArray? = - (t.data as? FloatArrayTensorData<*>)?.buffer + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + private fun floatBufferOf(t: Tensor): FloatArray? = when (val d = t.data) { + is FloatArrayTensorData<*> -> d.buffer + // Slab-backed data (#1145/#1146) has a nonzero base offset, so it cannot hand out its raw + // array — but its exact logical window as one copy still beats the boxed generic path by + // orders of magnitude (#949). The zero-copy path for the hot trio is [floatWindowOf]. + is sk.ainet.lang.tensor.data.StorageFloatTensorData<*> -> d.copyToFloatArray() + else -> null + } + + /** Zero-copy dense-FP32 window: the backing array plus the base offset of element 0 (#1146). */ + private class FloatWindow(val arr: FloatArray, val off: Int) - private fun floatResult( + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + private fun floatWindowOf(t: Tensor): FloatWindow? = when (val d = t.data) { + is FloatArrayTensorData<*> -> FloatWindow(d.buffer, 0) + is sk.ainet.lang.tensor.data.StorageFloatTensorData<*> -> { + val s = d.storage + s.checkAlive() + FloatWindow(s.floats!!, s.arrayOffset) + } + else -> null + } + + /** + * The scope the factory is placing outputs in, or `Ambient` — passed to kernel dispatch so + * adapter allocations (requantized activations, prepacked weights) land in the slab too (#1146). + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + protected fun dispatchScope(): sk.ainet.lang.memory.Scope = + (dataFactory as? sk.ainet.lang.tensor.data.ScopedTensorDataFactory)?.currentScope + ?: sk.ainet.lang.memory.Scope.Ambient + + /** + * A freshly computed dense-FP32 result, adopted through the factory — the single funnel a + * scope-aware factory (#1146) intercepts to place op outputs in the active scope. [buf] is + * ops-owned and never touched again; the default factory adopts it zero-copy. + */ + protected fun floatResult( shape: Shape, dtype: kotlin.reflect.KClass, buf: FloatArray, vararg inputs: Tensor, ): Tensor { @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(shape, dtype, buf) + val outData = dataFactory.adoptFloatArray(shape, dtype, buf) as sk.ainet.lang.tensor.data.TensorData return newTensor(outData, dtype, *inputs) } @@ -210,9 +245,12 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory t: Tensor, op: (Float) -> Float, ): Tensor? { - val src = floatBufferOf(t) ?: return null - val out = FloatArray(src.size) - for (i in src.indices) out[i] = op(src[i]) + val src = floatWindowOf(t) ?: return null + val n = t.shape.volume + val sa = src.arr + val so = src.off + val out = FloatArray(n) + for (i in 0 until n) out[i] = op(sa[so + i]) return floatResult(t.shape, t.dtype, out, t) } @@ -222,8 +260,12 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory op: (Float, Float) -> Float, ): Tensor? { if (a.dtype != b.dtype) return null - val ab = floatBufferOf(a) ?: return null - val bb = floatBufferOf(b) ?: return null + val aw = floatWindowOf(a) ?: return null + val bw = floatWindowOf(b) ?: return null + val ab = aw.arr + val ao = aw.off + val bb = bw.arr + val bo = bw.off val outShape = try { broadcastShapes(a.shape, b.shape) } catch (e: IllegalArgumentException) { @@ -232,19 +274,19 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val n = outShape.volume if (a.shape == b.shape) { val out = FloatArray(n) - for (i in 0 until n) out[i] = op(ab[i], bb[i]) + for (i in 0 until n) out[i] = op(ab[ao + i], bb[bo + i]) return floatResult(outShape, a.dtype, out, a, b) } if (a.shape.volume == 1) { - val av = ab[0] + val av = ab[ao] val out = FloatArray(n) - for (i in 0 until n) out[i] = op(av, bb[i]) + for (i in 0 until n) out[i] = op(av, bb[bo + i]) return floatResult(outShape, a.dtype, out, a, b) } if (b.shape.volume == 1) { - val bv = bb[0] + val bv = bb[bo] val out = FloatArray(n) - for (i in 0 until n) out[i] = op(ab[i], bv) + for (i in 0 until n) out[i] = op(ab[ao + i], bv) return floatResult(outShape, a.dtype, out, a, b) } // Last-dim ("bias") broadcast, mirroring DefaultCpuOpsJvm.vectorFloatBinary. @@ -259,7 +301,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val out = FloatArray(n) for (g in 0 until groups) { val off = g * outLast - for (i in 0 until outLast) out[off + i] = op(ab[off + i], bb[i]) + for (i in 0 until outLast) out[off + i] = op(ab[ao + off + i], bb[bo + i]) } return floatResult(outShape, a.dtype, out, a, b) } @@ -267,7 +309,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val out = FloatArray(n) for (g in 0 until groups) { val off = g * outLast - for (i in 0 until outLast) out[off + i] = op(ab[i], bb[off + i]) + for (i in 0 until outLast) out[off + i] = op(ab[ao + i], bb[bo + off + i]) } return floatResult(outShape, a.dtype, out, a, b) } @@ -573,7 +615,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory kernel(bi, 0, packed, 0, inputDim, outputDim, out, batch * outputDim) } @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(Shape(batchSize, outputDim), a.dtype, out) as TensorData + val outData = dataFactory.adoptFloatArray(Shape(batchSize, outputDim), a.dtype, out) as TensorData return newTensor(outData, a.dtype, a, b) } @@ -606,32 +648,36 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory // Packed-quant fast path (FP32 input × packed weight), resolved via KernelRegistry. KernelProfile.timeQuant { chooseQuantizedMatmulHeap(a, b) }?.let { return it } - // Fast path: 2D × 2D with FloatArray backing — direct buffer access, no per-element allocation - if (a.rank == 2 && b.rank == 2 - && (a.dtype == FP32::class) - && a.data is FloatArrayTensorData<*> && b.data is FloatArrayTensorData<*> - ) { - return KernelProfile.timeFp32 { - val aBuf = (a.data as FloatArrayTensorData<*>).buffer - val bBuf = (b.data as FloatArrayTensorData<*>).buffer - val m = a.shape[0] - val k = a.shape[1] - val n = b.shape[1] - require(k == b.shape[0]) { "Matrix multiplication shape mismatch: ${a.shape} vs ${b.shape}" } - val out = FloatArray(m * n) - for (i in 0 until m) { - val aOff = i * k - for (j in 0 until n) { - var sum = 0f - for (p in 0 until k) { - sum += aBuf[aOff + p] * bBuf[p * n + j] + // Fast path: 2D × 2D with dense FP32 backing (array or slab window, #1146) — direct + // buffer access, no per-element allocation + if (a.rank == 2 && b.rank == 2 && (a.dtype == FP32::class)) { + val aWin = floatWindowOf(a) + val bWin = floatWindowOf(b) + if (aWin != null && bWin != null) { + return KernelProfile.timeFp32 { + val aBuf = aWin.arr + val aBase = aWin.off + val bBuf = bWin.arr + val bBase = bWin.off + val m = a.shape[0] + val k = a.shape[1] + val n = b.shape[1] + require(k == b.shape[0]) { "Matrix multiplication shape mismatch: ${a.shape} vs ${b.shape}" } + val out = FloatArray(m * n) + for (i in 0 until m) { + val aOff = aBase + i * k + for (j in 0 until n) { + var sum = 0f + for (p in 0 until k) { + sum += aBuf[aOff + p] * bBuf[bBase + p * n + j] + } + out[i * n + j] = sum } - out[i * n + j] = sum } + @Suppress("UNCHECKED_CAST") + val outData = dataFactory.adoptFloatArray(Shape(m, n), a.dtype, out) as sk.ainet.lang.tensor.data.TensorData + newTensor(outData, a.dtype, a, b) } - @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(Shape(m, n), a.dtype, out) as sk.ainet.lang.tensor.data.TensorData - newTensor(outData, a.dtype, a, b) } } @@ -674,13 +720,13 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory Shape(m, n), FP32, ) - sk.ainet.backend.api.kernel.KernelDispatch.matmul(aNorm, bT, outView) + sk.ainet.backend.api.kernel.KernelDispatch.matmul(aNorm, bT, outView, dispatchScope()) val outShape = when { a.shape.rank == 1 -> Shape(n) // [k] x [k, n] -> [n] leading.isEmpty() -> Shape(m, n) else -> Shape(*(leading + n)) } - val outData = dataFactory.fromFloatArray(outShape, a.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData + val outData = dataFactory.adoptFloatArray(outShape, a.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData return newTensor(outData, a.dtype, a, b) } @@ -957,13 +1003,13 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val outView = sk.ainet.lang.memory.TensorView.dense( sk.ainet.lang.memory.Storage.Heap.wrap(outArray), Shape(m, n), FP32, ) - sk.ainet.backend.api.kernel.KernelDispatch.matmul(xNorm, wView, outView) + sk.ainet.backend.api.kernel.KernelDispatch.matmul(xNorm, wView, outView, dispatchScope()) val outShape = when { x.shape.rank == 1 -> Shape(n) leading.isEmpty() -> Shape(m, n) else -> Shape(*(leading + n)) } - val outData = dataFactory.fromFloatArray(outShape, x.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData + val outData = dataFactory.adoptFloatArray(outShape, x.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData return newTensor(outData, x.dtype, x, weight) } @@ -1141,7 +1187,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory } } @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(Shape(cols, rows), tensor.dtype, out) as sk.ainet.lang.tensor.data.TensorData + val outData = dataFactory.adoptFloatArray(Shape(cols, rows), tensor.dtype, out) as sk.ainet.lang.tensor.data.TensorData return newTensor(outData, tensor.dtype, tensor) } @@ -1213,7 +1259,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory out[flatOut] = srcBuf[flatIn] } @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(outShape, tensor.dtype, out) + val outData = dataFactory.adoptFloatArray(outShape, tensor.dtype, out) as sk.ainet.lang.tensor.data.TensorData return newTensor(outData, tensor.dtype, tensor) } @@ -2699,7 +2745,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory // real i32 tensor — see VoidTensorOps.argMax + ArgMaxOperationsConverter (emits stablehlo i32). val floats = FloatArray(outCount) { indices[it].toFloat() } @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromFloatArray(outShape, tensor.dtype, floats) + val outData = dataFactory.adoptFloatArray(outShape, tensor.dtype, floats) as sk.ainet.lang.tensor.data.TensorData return CpuTensor(outData, this, tensor.dtype, GradState(requiresGrad = false)) } @@ -3343,7 +3389,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory @Suppress("UNCHECKED_CAST") val outData = when (targetClass) { - FP32::class -> dataFactory.fromFloatArray( + FP32::class -> dataFactory.adoptFloatArray( tensor.shape, targetClass, copyTensorValuesAsFloatArray(tensor) @@ -3358,7 +3404,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val codec = if (targetClass == BF16::class) Bf16Codec else Fp16Codec val src = copyTensorValuesAsFloatArray(tensor) val rounded = FloatArray(src.size) { codec.decode(codec.encode(src[it])) } - dataFactory.fromFloatArray( + dataFactory.adoptFloatArray( tensor.shape, targetClass, rounded @@ -3662,7 +3708,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val shape = Shape(batch, heads, seqQ, headDim) @Suppress("UNCHECKED_CAST") - val data = dataFactory.fromFloatArray(shape, query.dtype, outBuf) as sk.ainet.lang.tensor.data.TensorData + val data = dataFactory.adoptFloatArray(shape, query.dtype, outBuf) as sk.ainet.lang.tensor.data.TensorData return newTensor(data, query.dtype, query, key, value) } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedOpOutputsTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedOpOutputsTest.kt new file mode 100644 index 000000000..ee0f3d231 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedOpOutputsTest.kt @@ -0,0 +1,115 @@ +package sk.ainet.sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.forwardScope +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.StorageClosedException +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.StorageFloatTensorData +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 + +/** + * #1146: op *outputs* draw from the active scope. Creation learned this in #1145; these tests pin + * the other half — a `matmul`/`add`/`relu` chain on scope-bound tensors produces slab-backed + * results, steady-state decode reuses the slab exactly (flat `peakFloats`, zero overflow), the + * numbers are bit-identical to Ambient, and a stale read after `reset()` is loud. + */ +@OptIn(ExperimentalMemoryApi::class) +class ScopedOpOutputsTest { + + private val xVals = FloatArray(8) { (it - 3).toFloat() * 0.5f } // [2, 4] + private val wVals = FloatArray(12) { ((it * 7) % 5 - 2).toFloat() } // [4, 3] + private val bVals = floatArrayOf(0.1f, -0.2f, 0.3f) // [3] + + private fun step(ctx: sk.ainet.context.ExecutionContext, x: Tensor, w: Tensor, b: Tensor): Tensor { + val ops = x.ops + return ops.relu(ops.add(ops.matmul(x, w), b)) + } + + private fun ambientResult(): FloatArray { + val ctx = DirectCpuExecutionContext() + val x = ctx.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val w = ctx.fromFloatArray(Shape(4, 3), FP32::class, wVals) + val b = ctx.fromFloatArray(Shape(3), FP32::class, bVals) + return step(ctx, x, w, b).data.copyToFloatArray() + } + + @Test + fun ambientOutputsAreUntouched() { + val ctx = DirectCpuExecutionContext() + val a = ctx.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val out = a.ops.add(a, a) + assertFalse(out.data is StorageFloatTensorData<*>, "Ambient op outputs stay on the plain array path") + } + + @Test + fun opOutputsDrawFromTheSlabAndMatchAmbient() { + val expected = ambientResult() + DirectCpuExecutionContext().forwardScope(slabFloats = 256) { scoped, scope -> + val x = scoped.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val w = scoped.fromFloatArray(Shape(4, 3), FP32::class, wVals) + val b = scoped.fromFloatArray(Shape(3), FP32::class, bVals) + val before = scope.usedFloats + val y = step(scoped, x, w, b) + assertTrue(y.data is StorageFloatTensorData<*>, "op output must be slab-backed under a scope") + assertTrue(scope.usedFloats > before, "the chain must have drawn its outputs from the slab") + assertContentEquals(expected, y.data.copyToFloatArray(), "slab-backed chain must match Ambient") + } + } + + @Test + fun steadyStateDecodeIsAFlatLine() { + val expected = ambientResult() + val base = DirectCpuExecutionContext() + // The real shape of a decode loop: weights persist OUTSIDE the forward scope + // (Ambient here; ModelScope in a real model), activations live inside it. + val w = base.fromFloatArray(Shape(4, 3), FP32::class, wVals) + val b = base.fromFloatArray(Shape(3), FP32::class, bVals) + base.forwardScope(slabFloats = 256) { scoped, scope -> + var peakAfterWarmup = -1 + repeat(8) { stepNo -> + val x = scoped.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val y = step(scoped, x, w, b) + assertTrue(y.data is StorageFloatTensorData<*>, "step $stepNo: output must be slab-backed") + assertContentEquals(expected, y.data.copyToFloatArray(), "step $stepNo numerics") + if (stepNo == 1) peakAfterWarmup = scope.peakFloats + if (stepNo > 1) { + assertEquals(peakAfterWarmup, scope.peakFloats, "step $stepNo: steady state must not grow the slab") + } + assertEquals(0L, scope.overflowBytes, "step $stepNo: nothing may spill past the slab") + scope.reset() + } + assertEquals(8, scope.steps.toInt()) + } + } + + @Test + fun staleOpOutputThrowsAfterReset() { + DirectCpuExecutionContext().forwardScope(slabFloats = 128) { scoped, scope -> + val a = scoped.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val out = a.ops.add(a, a) + scope.reset() + assertFailsWith { out.data[0, 0] } + } + } + + @Test + fun slabOverflowStillComputesCorrectly() { + val expected = ambientResult() + DirectCpuExecutionContext().forwardScope(slabFloats = 4) { scoped, scope -> + val x = scoped.fromFloatArray(Shape(2, 4), FP32::class, xVals) + val w = scoped.fromFloatArray(Shape(4, 3), FP32::class, wVals) + val b = scoped.fromFloatArray(Shape(3), FP32::class, bVals) + val y = step(scoped, x, w, b) + assertTrue(scope.overflowBytes > 0, "a 4-float slab must overflow") + assertContentEquals(expected, y.data.copyToFloatArray(), "overflow path numerics") + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index e62dde190..53d0fa9c0 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -408,9 +408,7 @@ internal class DefaultCpuOpsJvm( } } - val outData = DenseFloatArrayTensorData(Shape(n, cOut, outH, outW), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, input.dtype) + return floatResult(Shape(n, cOut, outH, outW), input.dtype, outBuffer) } override fun conv1d( @@ -463,9 +461,7 @@ internal class DefaultCpuOpsJvm( outBuffer ) - val outData = DenseFloatArrayTensorData(Shape(n, cOut, outL), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, input.dtype) + return floatResult(Shape(n, cOut, outL), input.dtype, outBuffer) } override fun conv3d( @@ -531,9 +527,7 @@ internal class DefaultCpuOpsJvm( outBuffer ) - val outData = DenseFloatArrayTensorData(Shape(n, cOut, outD, outH, outW), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, input.dtype) + return floatResult(Shape(n, cOut, outD, outH, outW), input.dtype, outBuffer) } /** @@ -616,9 +610,7 @@ internal class DefaultCpuOpsJvm( ) } } - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } is Q4_0TensorData -> { val outBuffer = FloatArray(batchSize * outputDim) @@ -632,9 +624,7 @@ internal class DefaultCpuOpsJvm( outBuffer, batch * outputDim, ) } - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } // Q5_1 / Q5_0 dispatch is handled in DefaultCpuOpsBase via the kernel // registry (block-major, shared with Native); not intercepted here. @@ -662,9 +652,7 @@ internal class DefaultCpuOpsJvm( ) } } - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } is NarrowFloatTensorData -> { // Narrow floats are dense (not block-quantized) and the kernel SPI is a @@ -688,9 +676,7 @@ internal class DefaultCpuOpsJvm( outBuffer, 0, outputDim, batchSize, outputDim, inputDim, ) - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } } // Q6_K / Q5_1 / Q5_0 dispatch is handled in DefaultCpuOpsBase via the kernel @@ -750,8 +736,7 @@ internal class DefaultCpuOpsJvm( } else -> return null } - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } /** @@ -781,8 +766,7 @@ internal class DefaultCpuOpsJvm( batch * outputDim, ) } - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(batchSize, outputDim), a.dtype, outBuffer) } override fun relu(tensor: Tensor): Tensor { @@ -803,9 +787,7 @@ internal class DefaultCpuOpsJvm( val x = buf[i] out[i] = x / (1f + kotlin.math.exp(-x)) } - val outData = DenseFloatArrayTensorData(Shape(tensor.shape.dimensions.copyOf()), out) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, tensor.dtype) + return floatResult(Shape(tensor.shape.dimensions.copyOf()), tensor.dtype, out) } override fun sum(tensor: Tensor, dim: Int?): Tensor { @@ -854,9 +836,7 @@ internal class DefaultCpuOpsJvm( // Case 1: exact shape match (fast path) if (a.shape == b.shape) { JvmVectorKernels.binaryFloat(aData.buffer, bData.buffer, outBuffer, outVolume, vectorOp, scalarOp) - val outData = DenseFloatArrayTensorData(Shape(outShape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } // Case 2: scalar broadcast @@ -875,9 +855,7 @@ internal class DefaultCpuOpsJvm( outBuffer[idx] = scalarOp(aval, bData.buffer[idx]) idx++ } - val outData = DenseFloatArrayTensorData(Shape(outShape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } if (bVol == 1) { val bval = bData.buffer[0] @@ -894,9 +872,7 @@ internal class DefaultCpuOpsJvm( outBuffer[idx] = scalarOp(aData.buffer[idx], bval) idx++ } - val outData = DenseFloatArrayTensorData(Shape(outShape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } // Case 3: last-dimension broadcasting (bias add). Supports arbitrary leading dims. @@ -929,9 +905,7 @@ internal class DefaultCpuOpsJvm( idx++ } } - val outData = DenseFloatArrayTensorData(Shape(outShape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } if (aIsBias && bVol == outVolume) { val step = floatSpecies.length() @@ -950,9 +924,7 @@ internal class DefaultCpuOpsJvm( idx++ } } - val outData = DenseFloatArrayTensorData(Shape(outShape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } } @@ -970,9 +942,7 @@ internal class DefaultCpuOpsJvm( val volume = tensor.shape.volume val outBuffer = FloatArray(volume) JvmVectorKernels.unaryFloat(tensorData.buffer, outBuffer, volume, vectorOp, scalarOp) - val outData = DenseFloatArrayTensorData(Shape(tensor.shape.dimensions.copyOf()), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, tensor.dtype) + return floatResult(Shape(tensor.shape.dimensions.copyOf()), tensor.dtype, outBuffer) } private fun supportsFloatOps(a: Tensor, b: Tensor): Boolean { @@ -1049,9 +1019,7 @@ internal class DefaultCpuOpsJvm( if (work >= blasThreshold) { val ok = JvmBlas.sgemmRowMajorNN(m, n, k, 1f, aData.buffer, bData.buffer, outBuffer) if (ok) { - val outData = DenseFloatArrayTensorData(Shape(m, n), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(m, n), a.dtype, outBuffer) } } } @@ -1066,9 +1034,7 @@ internal class DefaultCpuOpsJvm( outBuffer, 0, n, m, n, k, ) - val outData = DenseFloatArrayTensorData(Shape(m, n), outBuffer) - @Suppress("UNCHECKED_CAST") - return CpuTensor(outData as TensorData, this, a.dtype) + return floatResult(Shape(m, n), a.dtype, outBuffer) } private fun vectorFloatReduceAllSum(tensor: Tensor): Tensor? { diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index e8666ad5a..bf3c5b8c5 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -194,6 +194,7 @@ public final class sk/ainet/lang/graph/DefaultGraphExecutionContext : sk/ainet/l public fun stopRecording ()Lsk/ainet/tape/ExecutionTape; public final fun stopRecordingAndGet ()Lsk/ainet/tape/ExecutionTape; public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -457,6 +458,7 @@ public final class sk/ainet/lang/graph/exec/GraphExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withTensorDataFactory (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public static fun wrapIntArray (Lsk/ainet/lang/graph/exec/GraphExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; 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 a971e8798..11566d792 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 @@ -217,6 +217,7 @@ public final class sk/ainet/context/DefaultDataExecutionContext : sk/ainet/conte public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -246,6 +247,7 @@ public abstract interface class sk/ainet/context/ExecutionContext { public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -269,6 +271,7 @@ public final class sk/ainet/context/ExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/context/ExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withTensorDataFactory (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public static fun wrapIntArray (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -374,6 +377,7 @@ public final class sk/ainet/context/PhaseOverridingExecutionContext : sk/ainet/c public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -420,6 +424,7 @@ public final class sk/ainet/context/ScopedExecutionContext : sk/ainet/context/Ex public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -453,6 +458,7 @@ public final class sk/ainet/context/TrainingExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withTensorDataFactory (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public static fun wrapIntArray (Lsk/ainet/context/TrainingExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -2450,6 +2456,7 @@ public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ai public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public fun registerObserver (Lsk/ainet/context/ExecutionObserver;)V public fun unregisterObserver (Lsk/ainet/context/ExecutionObserver;)V + public fun withTensorDataFactory (Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -2782,6 +2789,7 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContext$DefaultImpls { public static fun placeholder (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; public static fun registerObserver (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/context/ExecutionObserver;)V public static fun unregisterObserver (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/context/ExecutionObserver;)V + public static fun withTensorDataFactory (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/data/TensorDataFactory;)Lsk/ainet/context/ExecutionContext; public static fun wrapByteArray (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; public static fun wrapFloatArray (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; public static fun wrapIntArray (Lsk/ainet/lang/nn/NeuralNetworkExecutionContext;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; @@ -4679,6 +4687,7 @@ public final class sk/ainet/lang/tensor/data/DenseIntArrayTensorData : sk/ainet/ public final class sk/ainet/lang/tensor/data/DenseTensorDataFactory : sk/ainet/lang/tensor/data/TensorDataFactory { public fun ()V + public fun adoptFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; public final fun fromByteArray ([BLsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/data/TensorData; public fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; @@ -4811,6 +4820,7 @@ public final class sk/ainet/lang/tensor/data/MemorySegmentTensorDataFactory : ja public fun ()V public fun (J)V public synthetic fun (JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun adoptFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; public fun close ()V public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; public fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; @@ -5395,6 +5405,26 @@ public abstract interface class sk/ainet/lang/tensor/data/RowDequantSource { public abstract fun dequantRow (I)[F } +public final class sk/ainet/lang/tensor/data/ScopedTensorDataFactory : sk/ainet/lang/tensor/data/TensorDataFactory { + public fun (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lkotlin/jvm/functions/Function0;)V + public fun adoptFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; + public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; + public fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; + public fun fromIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/data/TensorData; + public fun full (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;Ljava/lang/Number;)Lsk/ainet/lang/tensor/data/TensorData; + public final fun getCurrentScope ()Lsk/ainet/lang/memory/Scope; + public fun init (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/tensor/data/TensorData; + public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/data/TensorData; + public fun placeholder (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/data/TensorData; + public fun randn (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;FFLkotlin/random/Random;)Lsk/ainet/lang/tensor/data/TensorData; + public fun randomInit (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;Lkotlin/jvm/functions/Function1;Lkotlin/random/Random;)Lsk/ainet/lang/tensor/data/TensorData; + public fun uniform (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;FFLkotlin/random/Random;)Lsk/ainet/lang/tensor/data/TensorData; + public fun wrapByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; + public fun wrapFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; + public fun wrapIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/data/TensorData; + public fun zeros (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/data/TensorData; +} + public final class sk/ainet/lang/tensor/data/StorageFloatTensorData : sk/ainet/lang/tensor/data/TensorData { public fun (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Storage$Heap;)V public fun copyToFloatArray ()[F @@ -5422,6 +5452,7 @@ public final class sk/ainet/lang/tensor/data/TensorData$DefaultImpls { } public abstract interface class sk/ainet/lang/tensor/data/TensorDataFactory { + public fun adoptFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; public abstract fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; public abstract fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; public abstract fun fromIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/data/TensorData; @@ -5439,6 +5470,7 @@ public abstract interface class sk/ainet/lang/tensor/data/TensorDataFactory { } public final class sk/ainet/lang/tensor/data/TensorDataFactory$DefaultImpls { + public static fun adoptFloatArray (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; public static fun placeholder (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/data/TensorData; public static fun wrapByteArray (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/data/TensorData; public static fun wrapFloatArray (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/data/TensorData; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt index 3a2f572d7..721d55c8a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt @@ -49,6 +49,17 @@ public interface ExecutionContext { public val tensorDataFactory: TensorDataFactory + /** + * This context rebuilt around [factory] — the seam that lets a decorator (notably + * `ScopedExecutionContext`, #1146) swap in a scope-aware [TensorDataFactory] and get an + * *ops instance that allocates through it*, since ops are constructed from the factory. + * + * The default returns `this` unchanged: a context that cannot rebuild itself simply keeps + * its own factory, and op outputs stay GC-allocated (correct, just not scope-recycled). + * Contexts that own their ops construction (e.g. `DirectCpuExecutionContext`) override. + */ + public fun withTensorDataFactory(factory: TensorDataFactory): ExecutionContext = this + /** * Workspace allocator for short-lived intermediate buffers (attention * scratch, RoPE tables, KV-cache slice copies, padding scratch, etc.). diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt index ba3187afa..019a5db23 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt @@ -28,6 +28,36 @@ public class ScopedExecutionContext( override val memoryScope: Scope, ) : ExecutionContext by base { + /** + * The base rebuilt around a scope-aware factory, so *op outputs* allocate from the slab too + * (#1146). A base that cannot rebuild itself (default [ExecutionContext.withTensorDataFactory]) + * returns itself — creation still draws from the scope, op outputs stay GC-allocated. + * + * Note the ops-binding rule: `a + b` dispatches through the ops instance that *created* `a`, + * so only tensors created through this context (or explicitly re-bound) produce scoped + * outputs. Tensors made before entering the scope keep their unscoped ops — deliberately, so + * their results do not die at `reset()`. + */ + private val scopedBase: ExecutionContext = + base.withTensorDataFactory( + sk.ainet.lang.tensor.data.ScopedTensorDataFactory(base.tensorDataFactory) { memoryScope }, + ) + + override val tensorDataFactory: sk.ainet.lang.tensor.data.TensorDataFactory + get() = scopedBase.tensorDataFactory + + override val ops: sk.ainet.lang.tensor.ops.TensorOps + get() = scopedBase.ops + + /** + * Bind tensors made through this context to the *scoped* ops — `a + b` dispatches through + * the ops that created `a`, and only the scoped ops allocate outputs from the slab. + */ + override fun fromData( + data: sk.ainet.lang.tensor.data.TensorData, + dtype: KClass, + ): Tensor = sk.ainet.lang.tensor.operators.OpsBoundTensor.fromData(data, dtype, ops) + override fun zeros(shape: Shape, dtype: KClass): Tensor = super.zeros(shape, dtype) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt index 96fb66fe0..e372023a7 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt @@ -704,6 +704,17 @@ public class DenseTensorDataFactory: TensorDataFactory { } } + /** Zero-copy adoption for the dtypes [wrapFloatArray] handles; the narrow-tagged rest go + * through [fromFloatArray] exactly as before (#1146). */ + override fun adoptFloatArray( + shape: Shape, + dtype: KClass, + data: FloatArray + ): TensorData = when (dtype) { + FP32::class, FP16::class -> wrapFloatArray(shape, dtype, data) + else -> fromFloatArray(shape, dtype, data) + } + override fun wrapIntArray( shape: Shape, dtype: KClass, diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ScopedTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ScopedTensorDataFactory.kt new file mode 100644 index 000000000..75fd7c258 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ScopedTensorDataFactory.kt @@ -0,0 +1,99 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 +import kotlin.reflect.KClass + +/** + * [base] with dense-FP32 allocations drawn from the active [Scope] — the op-output half of the + * Scope split (#1146). + * + * The creation path (`ExecutionContext.zeros/…`) learned to consult `memoryScope` in #1145; op + * *outputs* still reached the GC because `DefaultCpuOps` allocates them through its + * [TensorDataFactory]. This decorator is that factory: when [scope] is anything other than + * [Scope.Ambient] and the dtype is dense FP32, `zeros`/`ones`/`full`/`init`/`fromFloatArray` + * return [StorageFloatTensorData] over the scope's slab — recycled at `ForwardScope.reset()`, + * loud after it. Everything else falls through to [base] untouched. + * + * The slab is *not* zeroed between steps, so every intercepted method fully writes its region: + * `zeros`/`ones`/`full` fill, `init` writes each element, `fromFloatArray` copies the source in. + * + * `wrapFloatArray`/`wrapIntArray`/`wrapByteArray` are deliberately **not** intercepted: their + * contract is zero-copy over a caller-owned array (loaders use them for weights), and neither + * copying them into a slab nor letting them die at `reset()` would honour it. + */ +@ExperimentalMemoryApi +public class ScopedTensorDataFactory( + private val base: TensorDataFactory, + private val scope: () -> Scope, +) : TensorDataFactory by base { + + /** The scope currently in effect — lets an ops implementation holding this factory pass the + * same scope to kernel dispatch, so adapter allocations land in the slab too (#1146). */ + public val currentScope: Scope get() = scope() + + private fun slab(shape: Shape, dtype: KClass): StorageFloatTensorData? { + val s = scope() + if (s === Scope.Ambient || dtype != FP32::class) return null + return StorageFloatTensorData(shape, s.allocateFloats(shape.volume)) + } + + @Suppress("UNCHECKED_CAST") + private fun fill(scoped: StorageFloatTensorData, value: Float): TensorData { + val st = scoped.storage + st.floats!!.fill(value, st.arrayOffset, st.arrayOffset + scoped.shape.volume) + return scoped as TensorData + } + + override fun zeros(shape: Shape, dtype: KClass): TensorData = + slab(shape, dtype)?.let { fill(it, 0f) } ?: base.zeros(shape, dtype) + + override fun ones(shape: Shape, dtype: KClass): TensorData = + slab(shape, dtype)?.let { fill(it, 1f) } ?: base.ones(shape, dtype) + + override fun full(shape: Shape, dtype: KClass, value: Number): TensorData = + slab(shape, dtype)?.let { fill(it, value.toFloat()) } ?: base.full(shape, dtype, value) + + @Suppress("UNCHECKED_CAST") + override fun fromFloatArray(shape: Shape, dtype: KClass, data: FloatArray): TensorData { + val scoped = slab(shape, dtype) ?: return base.fromFloatArray(shape, dtype, data) + val st = scoped.storage + data.copyInto(st.floats!!, st.arrayOffset, 0, shape.volume) + return scoped as TensorData + } + + @Suppress("UNCHECKED_CAST") + override fun adoptFloatArray(shape: Shape, dtype: KClass, data: FloatArray): TensorData { + val scoped = slab(shape, dtype) ?: return base.adoptFloatArray(shape, dtype, data) + val st = scoped.storage + data.copyInto(st.floats!!, st.arrayOffset, 0, shape.volume) + return scoped as TensorData + } + + @Suppress("UNCHECKED_CAST") + override fun init( + shape: Shape, + dtype: KClass, + generator: (indices: IntArray) -> V, + ): TensorData { + val scoped = slab(shape, dtype) ?: return base.init(shape, dtype, generator) + val st = scoped.storage + val floats = st.floats!! + val dims = shape.dimensions + val indices = IntArray(dims.size) + val volume = shape.volume + // Row-major walk, same visiting order as the dense factory. + for (flat in 0 until volume) { + var remaining = flat + for (d in dims.indices.reversed()) { + indices[d] = remaining % dims[d] + remaining /= dims[d] + } + floats[st.arrayOffset + flat] = generator(indices) as Float + } + return scoped as TensorData + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorDataFactory.kt index 6525f240d..2ac8c2e22 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/TensorDataFactory.kt @@ -81,6 +81,22 @@ public interface TensorDataFactory { data: ByteArray ): TensorData + /** + * Adopts a freshly computed FloatArray as tensor data — the op-output entry point (#1146). + * + * Unlike [wrapFloatArray], [data] is *owned by nobody else*: the caller computed it, hands it + * over, and never touches it again. That transfer is what licenses a factory to relocate the + * values — a scope-aware factory copies them into its slab so they are recycled at + * `ForwardScope.reset()`; a dense factory adopts the array zero-copy. The interface default + * stays on the copying [fromFloatArray] so every dtype a factory accepts there is accepted + * here; factories with a wider zero-copy wrap override. + */ + public fun adoptFloatArray( + shape: Shape, + dtype: KClass, + data: FloatArray + ): TensorData = fromFloatArray(shape, dtype, data) + /** * Wraps a FloatArray without copying. The caller must ensure the array * is not mutated while the returned TensorData is in use.