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. 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) + } +}