From 5f7830ed2eb343a85b8d21ad02852c844d961424 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 10:56:49 +0200 Subject: [PATCH] feat(memory): wire ExecutionContext.memoryScope into tensor creation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The answer to #1135: eager needs no graph to reuse buffers — the Scope split already models it. Eager lifetimes follow call structure, so ForwardScope.reset() at step boundaries replaces graph liveness analysis. What was missing was any reader of memoryScope; this is the reader. zeros/ones/full/fromFloatArray consult memoryScope when it is not Scope.Ambient and the dtype is dense FP32: the bytes come from the scope's slab as StorageFloatTensorData, die at reset(), and a use-after-reset is a StorageClosedException naming the storage. Ambient — the default everywhere — short-circuits to the factory path untouched. StorageFloatTensorData deliberately does not implement FloatArrayTensorData: a slab slice has a nonzero arrayOffset and the ops fast paths that unwrap 'buffer' assume offset 0. The backend-cpu parity test pins real CPU ops on slab-backed tensors at nonzero offsets against the Ambient results, across resets. ScratchPool stays, unmerged: intra-kernel workspace vs inter-op activation lifetime are different layers, now documented on both. Op *outputs* still allocate raw arrays; routing them through the scope is #1146. Closes #1145. Co-Authored-By: Claude Fable 5 --- .../tensor/ops/ScopedCreationOpsParityTest.kt | 55 ++++++++++ .../api/jvm/skainet-lang-core.api | 47 ++++++++ .../sk/ainet/context/ExecutionContext.kt | 45 ++++++++ .../ainet/context/ScopedExecutionContext.kt | 64 +++++++++++ .../tensor/data/StorageFloatTensorData.kt | 78 +++++++++++++ .../context/ExecutionContextScopeTest.kt | 103 ++++++++++++++++++ 6 files changed, 392 insertions(+) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedCreationOpsParityTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/StorageFloatTensorData.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/ExecutionContextScopeTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedCreationOpsParityTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedCreationOpsParityTest.kt new file mode 100644 index 000000000..8ee9ead27 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/ScopedCreationOpsParityTest.kt @@ -0,0 +1,55 @@ +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.tensor.Shape +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.assertTrue + +/** + * #1145 guard for the offset-0 trap: a slab-backed tensor has a nonzero `arrayOffset`, and a + * fast path that grabbed the raw buffer would read the whole slab. `StorageFloatTensorData` + * deliberately stays off `FloatArrayTensorData`, so real CPU ops must produce numbers identical + * to the Ambient path — for tensors sliced from anywhere in the slab, across resets. + */ +@OptIn(ExperimentalMemoryApi::class) +class ScopedCreationOpsParityTest { + + @Test + fun opsOnSlabBackedTensorsMatchAmbient() { + val ctx = DirectCpuExecutionContext() + val aVals = FloatArray(6) { (it + 1).toFloat() } // 2×3 + val bVals = FloatArray(12) { (it % 5 - 2).toFloat() } // 3×4 + + val ambientMatmul: FloatArray + val ambientAdd: FloatArray + run { + val a = ctx.fromFloatArray(Shape(2, 3), FP32::class, aVals) + val b = ctx.fromFloatArray(Shape(3, 4), FP32::class, bVals) + ambientMatmul = ctx.ops.matmul(a, b).data.copyToFloatArray() + ambientAdd = ctx.ops.add(a, a).data.copyToFloatArray() + } + + ctx.forwardScope(slabFloats = 64) { scoped, scope -> + repeat(3) { step -> + // A leading allocation pushes the later tensors deeper into the slab, so the + // offsets under test are nonzero and different from the previous step's layout. + scoped.zeros(Shape(1 + step), FP32::class) + val a = scoped.fromFloatArray(Shape(2, 3), FP32::class, aVals) + val b = scoped.fromFloatArray(Shape(3, 4), FP32::class, bVals) + assertTrue(a.data is StorageFloatTensorData<*>, "step $step: creation must draw from the slab") + assertTrue((a.data as StorageFloatTensorData<*>).storage.arrayOffset > 0, "offset under test must be nonzero") + + assertContentEquals(ambientMatmul, ctx.ops.matmul(a, b).data.copyToFloatArray(), "step $step: matmul") + assertContentEquals(ambientAdd, ctx.ops.add(a, a).data.copyToFloatArray(), "step $step: add") + assertEquals(Shape(2, 4), ctx.ops.matmul(a, b).shape) + scope.reset() + } + } + } +} 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 a3a27bcab..a971e8798 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 @@ -396,6 +396,40 @@ public final class sk/ainet/context/ResettableExecutionObserver$DefaultImpls { public static fun onTensorMaterialized (Lsk/ainet/context/ResettableExecutionObserver;Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Tensor;)V } +public final class sk/ainet/context/ScopedExecutionContext : sk/ainet/context/ExecutionContext { + public fun (Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/memory/Scope;)V + public fun fromByteArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[B)Lsk/ainet/lang/tensor/Tensor; + public fun fromData (Lsk/ainet/lang/tensor/data/TensorData;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; + public fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F)Lsk/ainet/lang/tensor/Tensor; + public fun fromIntArray (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[I)Lsk/ainet/lang/tensor/Tensor; + public fun full (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public fun getExecutionStats ()Lsk/ainet/context/ExecutionStats; + public fun getHooks ()Lsk/ainet/lang/nn/hooks/ForwardHooks; + public fun getInTraining ()Z + public fun getMemoryInfo ()Lsk/ainet/context/MemoryInfo; + public fun getMemoryScope ()Lsk/ainet/lang/memory/Scope; + public fun getMemoryTracker ()Lsk/ainet/lang/tensor/storage/MemoryTracker; + public fun getObservers ()Lsk/ainet/context/ExecutionObserverRegistry; + public fun getOps ()Lsk/ainet/lang/tensor/ops/TensorOps; + public fun getPhase ()Lsk/ainet/context/Phase; + public fun getScratch ()Lsk/ainet/lang/tensor/scratch/ScratchPool; + public fun getTensorDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; + public fun getTraceSink ()Lsk/ainet/lang/memory/trace/TraceSink; + public fun isRecording ()Z + public fun ones (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; + 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 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; + public fun zeros (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/Tensor; +} + +public final class sk/ainet/context/ScopedExecutionContextKt { + public static final fun forwardScope (Lsk/ainet/context/ExecutionContext;ILkotlin/jvm/functions/Function2;)Ljava/lang/Object; +} + public abstract interface class sk/ainet/context/TrainingExecutionContext : sk/ainet/context/ExecutionContext { public abstract fun backward (Ljava/util/List;Ljava/util/List;)V public abstract fun startRecording ()V @@ -5361,6 +5395,19 @@ public abstract interface class sk/ainet/lang/tensor/data/RowDequantSource { public abstract fun dequantRow (I)[F } +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 + public fun get ([I)Ljava/lang/Float; + public synthetic fun get ([I)Ljava/lang/Object; + public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public final fun getStorage ()Lsk/ainet/lang/memory/Storage$Heap; + public fun getView ()Lsk/ainet/lang/memory/TensorView; + public fun set ([IF)V + public synthetic fun set ([ILjava/lang/Object;)V +} + public abstract interface class sk/ainet/lang/tensor/data/TensorData : sk/ainet/lang/tensor/data/ItemsAccessor { public fun copyToFloatArray ()[F public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; 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 313f8172c..3a2f572d7 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 @@ -60,6 +60,12 @@ public interface ExecutionContext { * Callers MUST acquire inside an active [ScratchPool.scope] block; * acquires outside a scope succeed but the buffer is not returned to the * pool when dropped. + * + * Boundary with [memoryScope], on purpose (#1135): `scratch` is untyped *intra-kernel* + * workspace — raw arrays inside one op invocation, returned when its block exits. + * [memoryScope] governs *inter-op* activation lifetime — typed tensors that live across ops + * within a step and are recycled by `ForwardScope.reset()`. They are different layers and + * stay separate; neither replaces the other. */ public val scratch: ScratchPool get() = NoopScratchPool @@ -74,16 +80,46 @@ public interface ExecutionContext { observers.unregister(observer) } + /** + * Dense FP32 data drawn from [memoryScope] when a scope other than `Ambient` is active — the + * creation-path reader of the Scope split (#1145). `null` on the Ambient default, so the + * factory path is untouched for every context that never opts in. The region is *not* cleared: + * a slab slice after `reset()` holds old bytes, so callers fill it themselves. + */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + private fun scopedDenseFloats( + shape: Shape, + dtype: KClass, + ): sk.ainet.lang.tensor.data.StorageFloatTensorData? { + val scope = memoryScope + if (scope === sk.ainet.lang.memory.Scope.Ambient || dtype != sk.ainet.lang.types.FP32::class) return null + return sk.ainet.lang.tensor.data.StorageFloatTensorData(shape, scope.allocateFloats(shape.volume)) + } + + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) public fun full(shape: Shape, dtype: KClass, value: Number): Tensor { + scopedDenseFloats(shape, dtype)?.let { scoped -> + val s = scoped.storage + s.floats!!.fill(value.toFloat(), s.arrayOffset, s.arrayOffset + shape.volume) + @Suppress("UNCHECKED_CAST") + return fromData(scoped as TensorData, dtype) + } val data = tensorDataFactory.full(shape, dtype, value) return fromData(data, dtype) } + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) public fun zeros( shape: Shape, dtype: KClass ): Tensor { + scopedDenseFloats(shape, dtype)?.let { scoped -> + val s = scoped.storage + s.floats!!.fill(0f, s.arrayOffset, s.arrayOffset + shape.volume) + @Suppress("UNCHECKED_CAST") + return fromData(scoped as TensorData, dtype) + } val data = tensorDataFactory.zeros(shape, dtype) return fromData(data, dtype) } @@ -102,10 +138,12 @@ public interface ExecutionContext { return fromData(data, dtype) } + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) public fun ones( shape: Shape, dtype: KClass ): Tensor { + if (memoryScope !== sk.ainet.lang.memory.Scope.Ambient) return full(shape, dtype, 1) val data = tensorDataFactory.ones(shape, dtype) return fromData(data, dtype) } @@ -116,11 +154,18 @@ public interface ExecutionContext { ops ) + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) public fun fromFloatArray( shape: Shape, dtype: KClass, data: FloatArray ): Tensor { + scopedDenseFloats(shape, dtype)?.let { scoped -> + val s = scoped.storage + data.copyInto(s.floats!!, s.arrayOffset, 0, shape.volume) + @Suppress("UNCHECKED_CAST") + return fromData(scoped as TensorData, dtype) + } val data = tensorDataFactory.fromFloatArray(shape, dtype, data) return fromData(data, dtype) } 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 new file mode 100644 index 000000000..ba3187afa --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ScopedExecutionContext.kt @@ -0,0 +1,64 @@ +package sk.ainet.context + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.ForwardScope +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType +import kotlin.reflect.KClass + +/** + * [base] with a [memoryScope] — the wiring #1135 asked about (#1145). + * + * The mechanism existed before this class did: [ForwardScope] bump-allocates a pre-sized slab and + * `reset()` recycles it per step; what was missing was any reader of `ExecutionContext.memoryScope`. + * This decorator is that reader's other half: creation methods on [ExecutionContext] consult + * `memoryScope` when it is not [Scope.Ambient], so `zeros`/`full`/`ones`/`fromFloatArray` draw + * from the scope's slab and their tensors die at `reset()`. + * + * `Ambient` remains the default everywhere — nothing changes for a context that never opts in. + * The boundary with [ExecutionContext.scratch] is deliberate: `scratch` is untyped *intra-kernel* + * workspace inside one op invocation; `memoryScope` governs *inter-op* activation lifetime across + * a step. They stay separate. + */ +@ExperimentalMemoryApi +public class ScopedExecutionContext( + private val base: ExecutionContext, + override val memoryScope: Scope, +) : ExecutionContext by base { + + override fun zeros(shape: Shape, dtype: KClass): Tensor = + super.zeros(shape, dtype) + + override fun ones(shape: Shape, dtype: KClass): Tensor = + super.ones(shape, dtype) + + override fun full(shape: Shape, dtype: KClass, value: Number): Tensor = + super.full(shape, dtype, value) + + override fun fromFloatArray(shape: Shape, dtype: KClass, data: FloatArray): Tensor = + super.fromFloatArray(shape, dtype, data) +} + +/** + * Run [block] with a [ForwardScope] of [slabFloats] floats active on this context, closing the + * scope (and everything it handed out) afterwards. Call `reset()` on the scope between steps: + * + * ```kotlin + * ctx.forwardScope(slabFloats = 1 shl 20) { scoped, scope -> + * while (decoding) { + * step(scoped) + * scope.reset() // steady state: zero new slab bytes per step + * } + * } + * ``` + */ +@ExperimentalMemoryApi +public inline fun ExecutionContext.forwardScope( + slabFloats: Int, + block: (ctx: ScopedExecutionContext, scope: ForwardScope) -> R, +): R { + val scope = ForwardScope(slabFloats, traceSink) + return scope.use { block(ScopedExecutionContext(this, scope), scope) } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/StorageFloatTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/StorageFloatTensorData.kt new file mode 100644 index 000000000..29a049def --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/StorageFloatTensorData.kt @@ -0,0 +1,78 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 + +/** + * Dense FP32 tensor data over a [Storage.Heap] region — the creation-path end of the Scope split + * (#1145): what `ExecutionContext.zeros/full/fromFloatArray` hand out when a [sk.ainet.lang.memory.Scope] + * other than `Ambient` is active, so an activation's bytes come from the forward slab and die at + * `reset()` instead of waiting for the GC. + * + * Every access goes through [Storage.checkAlive], so a use-after-reset is a + * [sk.ainet.lang.memory.StorageClosedException] naming the storage — not silent corruption. + * + * **Deliberately not a [FloatArrayTensorData].** A slab slice has a nonzero + * [Storage.Heap.arrayOffset], and the ops fast paths that unwrap `buffer` assume offset 0; exposing + * this data through that interface would hand them the whole slab. Element access and + * [copyToFloatArray] are offset-correct; kernels that want zero-copy take [view], which carries the + * offset properly. + */ +@ExperimentalMemoryApi +public class StorageFloatTensorData( + initialShape: Shape, + public val storage: Storage.Heap, +) : TensorData { + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + private val strides: IntArray = shape.computeStrides() + + init { + requireNotNull(storage.floats) { "StorageFloatTensorData needs float-backed storage" } + require(storage.elementCount >= shape.volume) { + "storage holds ${storage.elementCount} floats, shape $shape needs ${shape.volume}" + } + } + + private fun flatIndex(indices: IntArray): Int { + require(indices.size == shape.dimensions.size) { + "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" + } + var flat = 0 + for (i in indices.indices) { + val idx = indices[i] + require(idx >= 0 && idx < shape.dimensions[i]) { + "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" + } + flat += idx * strides[i] + } + return flat + } + + override fun get(vararg indices: Int): Float { + storage.checkAlive() + return storage.floats!![storage.arrayOffset + flatIndex(indices)] + } + + override fun set(vararg indices: Int, value: Float) { + storage.checkAlive() + storage.floats!![storage.arrayOffset + flatIndex(indices)] = value + } + + override fun copyToFloatArray(): FloatArray { + storage.checkAlive() + val off = storage.arrayOffset + return storage.floats!!.copyOfRange(off, off + shape.volume) + } + + /** A dense view over the same region — the storage carries the offset, nothing is copied. */ + override val view: TensorView + get() { + storage.checkAlive() + return TensorView.dense(storage, shape, FP32) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/ExecutionContextScopeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/ExecutionContextScopeTest.kt new file mode 100644 index 000000000..27d2f97ae --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/context/ExecutionContextScopeTest.kt @@ -0,0 +1,103 @@ +package sk.ainet.context + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.ForwardScope +import sk.ainet.lang.memory.StorageClosedException +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.StorageFloatTensorData +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * #1145: `ExecutionContext.memoryScope` gets its first reader. Ambient is byte-for-byte the old + * path; a ForwardScope makes creation draw from the slab, `reset()` recycles it, and a + * use-after-reset is a loud [StorageClosedException] rather than silent garbage. + */ +@OptIn(ExperimentalMemoryApi::class) +class ExecutionContextScopeTest { + + @Test + fun ambientContextIsUntouched() { + val ctx = DefaultDataExecutionContext() + val t = ctx.zeros(Shape(4), FP32::class) + assertFalse(t.data is StorageFloatTensorData<*>, "Ambient must use the factory path") + } + + @Test + fun scopedCreationDrawsFromTheSlab() { + DefaultDataExecutionContext().forwardScope(slabFloats = 64) { ctx, scope -> + val t = ctx.zeros(Shape(2, 8), FP32::class) + assertTrue(t.data is StorageFloatTensorData<*>, "scoped creation must use the slab") + assertEquals(16, scope.usedFloats) + val u = ctx.full(Shape(8), FP32::class, 3) + assertEquals(24, scope.usedFloats, "second tensor bumps the same slab") + assertEquals(0f, t.data[0, 0]) + assertEquals(3f, u.data[3]) + } + } + + @Test + fun theSlabIsDirtyAndZerosMustZero() { + DefaultDataExecutionContext().forwardScope(slabFloats = 16) { ctx, scope -> + val dirty = ctx.zeros(Shape(16), FP32::class) + for (i in 0 until 16) dirty.data.set(i, value = 7f) + scope.reset() + val clean = ctx.zeros(Shape(16), FP32::class) + for (i in 0 until 16) assertEquals(0f, clean.data[i], "slab byte $i must be re-zeroed") + } + } + + @Test + fun fromFloatArrayCopiesIntoTheSlab() { + DefaultDataExecutionContext().forwardScope(slabFloats = 8) { ctx, _ -> + val t = ctx.fromFloatArray(Shape(4), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)) + assertTrue(t.data is StorageFloatTensorData<*>) + assertEquals(2f, t.data[1]) + assertEquals(4f, t.data[3]) + } + } + + @Test + fun steadyStateReusesTheSlabAndStaleReadsThrow() { + DefaultDataExecutionContext().forwardScope(slabFloats = 32) { ctx, scope -> + val step1 = ctx.ones(Shape(32), FP32::class) + assertEquals(32, scope.peakFloats) + scope.reset() + assertEquals(0, scope.usedFloats, "reset rewinds the slab") + val step2 = ctx.ones(Shape(32), FP32::class) + assertEquals(32, scope.peakFloats, "steady state allocates zero new slab bytes") + assertEquals(1f, step2.data[0]) + assertFailsWith("a step-1 tensor read after reset must be loud") { + step1.data[0] + } + } + } + + @Test + fun retainIsTheSanctionedEscape() { + val ctx = DefaultDataExecutionContext() + lateinit var kept: FloatArray + ctx.forwardScope(slabFloats = 8) { scoped, scope -> + val t = scoped.fromFloatArray(Shape(4), FP32::class, floatArrayOf(9f, 8f, 7f, 6f)) + val data = t.data as StorageFloatTensorData<*> + val retained = scope.retain(data.storage) + scope.reset() + kept = retained.floats!!.copyOfRange(retained.arrayOffset, retained.arrayOffset + 4) + } + assertEquals(listOf(9f, 8f, 7f, 6f), kept.toList()) + } + + @Test + fun overflowBeyondTheSlabStillWorks() { + DefaultDataExecutionContext().forwardScope(slabFloats = 4) { ctx, scope -> + val big = ctx.zeros(Shape(64), FP32::class) + assertTrue(big.data is StorageFloatTensorData<*>) + assertTrue(scope.overflowBytes > 0, "a tensor over the slab size goes to tracked overflow") + assertEquals(0f, big.data[63]) + } + } +}