feat(scratch): add ScratchPool SPI for runtime workspace allocation

gradle.properties

@@ -1,5 +1,5 @@
 GROUP=sk.ainet.core
-VERSION_NAME=0.20.0
+VERSION_NAME=0.21.0-SNAPSHOT
 POM_DESCRIPTION=SKaiNET
 
 POM_URL=https://github.com/SKaiNET-developers/skainet/

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/ExecutionContext.kt

@@ -6,6 +6,8 @@ import sk.ainet.lang.tensor.data.TensorData
 import sk.ainet.lang.tensor.data.TensorDataFactory
 import sk.ainet.lang.tensor.operators.OpsBoundTensor
 import sk.ainet.lang.tensor.ops.TensorOps
+import sk.ainet.lang.tensor.scratch.NoopScratchPool
+import sk.ainet.lang.tensor.scratch.ScratchPool
 import sk.ainet.lang.tensor.storage.MemoryPlanner
 import sk.ainet.lang.tensor.storage.MemoryTracker
 import sk.ainet.lang.types.DType
@@ -23,6 +25,20 @@ public interface ExecutionContext {
 
     public val tensorDataFactory: TensorDataFactory
 
+    /**
+     * Workspace allocator for short-lived intermediate buffers (attention
+     * scratch, RoPE tables, KV-cache slice copies, padding scratch, etc.).
+     *
+     * Default is [NoopScratchPool] — every acquire allocates a fresh array,
+     * matching pre-pool behavior. Implementations that want pooling override
+     * this property (or wrap an existing context).
+     *
+     * 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.
+     */
+    public val scratch: ScratchPool get() = NoopScratchPool
+
     // Execution observers for tracing/benchmarking
     public val observers: ExecutionObserverRegistry
 

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/scratch/ScratchPool.kt

@@ -0,0 +1,175 @@
+package sk.ainet.lang.tensor.scratch
+
+/**
+ * Pool of reusable [FloatArray] scratch buffers, scoped to a single forward
+ * (or backward) pass.
+ *
+ * Workspace allocation is generic across nn workloads — attention, RoPE,
+ * convolutions, embedding gathers, training-time gradient buffers all need
+ * short-lived intermediates. Routing those through a pool eliminates per-step
+ * allocation pressure on the GC heap.
+ *
+ * Typical lifecycle:
+ *
+ *  1. The runtime owns a [ScratchPool] (one per forward-running thread).
+ *  2. The runtime opens a [scope] block around each forward pass.
+ *  3. Layers acquire buffers via [acquireFloat] / [acquireFloatZeroed] from
+ *     the [sk.ainet.context.ExecutionContext.scratch] field.
+ *  4. On scope exit every buffer acquired in the scope is recycled to the
+ *     per-size-class free list for the next pass.
+ *
+ * Buffers may be larger than the requested size — callers must always index
+ * within the size they asked for, not `buf.size`. This is the standard
+ * scratch-buffer contract: external `outSize`/`n`/`stride` variables drive
+ * iteration; the array's own `.size` is an implementation detail.
+ *
+ * Acquires made *outside* an active scope are not tracked and become regular
+ * allocations (the buffer is not returned to the pool when the caller drops
+ * it). This is intentional — it lets the no-op pool implementation avoid
+ * any scope bookkeeping.
+ *
+ * Single-threaded by intent: one forward pass at a time per pool.
+ * Concurrent forwards must use separate pools.
+ */
+public interface ScratchPool {
+    /**
+     * Acquire a [FloatArray] with at least [minSize] elements. Contents are
+     * unspecified — callers must overwrite the range they read.
+     */
+    public fun acquireFloat(minSize: Int): FloatArray
+
+    /**
+     * Acquire a [FloatArray] with at least [minSize] elements, with the
+     * range `[0, minSize)` zero-filled. Use when the caller relies on
+     * sparse-write zeros (e.g. padding a smaller block into a larger
+     * destination and leaving the gap implicit).
+     */
+    public fun acquireFloatZeroed(minSize: Int): FloatArray
+
+    /**
+     * Open a forward-pass scope. All buffers acquired inside [block] are
+     * recycled at exit. Scopes may be nested; each open/close balances.
+     */
+    public fun <R> scope(block: () -> R): R
+
+    /** Stats for benchmarking and leak detection. */
+    public fun stats(): ScratchStats
+}
+
+public data class ScratchStats(
+    val acquireCount: Long,
+    val cacheHits: Long,
+    val highWaterBytes: Long,
+    val activeBuffers: Int
+)
+
+/**
+ * No-op pool: every [acquireFloat] / [acquireFloatZeroed] allocates a fresh
+ * array; [scope] just runs the block. Default carrier returned by
+ * [sk.ainet.context.ExecutionContext.scratch] when no pooling is configured —
+ * preserves pre-pool behavior bit-for-bit.
+ */
+public object NoopScratchPool : ScratchPool {
+    override fun acquireFloat(minSize: Int): FloatArray = FloatArray(minSize)
+    override fun acquireFloatZeroed(minSize: Int): FloatArray = FloatArray(minSize)
+    override fun <R> scope(block: () -> R): R = block()
+    override fun stats(): ScratchStats = ScratchStats(0L, 0L, 0L, 0)
+}
+
+/**
+ * Size-classed slab pool with power-of-two buckets starting at 64 floats.
+ *
+ * Sizes round up to the next power of two: `1..64 → 64`, `65..128 → 128`,
+ * ... up to [MAX_CLASSES] classes (`64 * 2^19 ≈ 33M floats` = 128 MB at the
+ * top of the range). Up to [maxBuffersPerClass] buffers are retained per
+ * class; surplus buffers are dropped to GC at scope exit.
+ *
+ * Choice rationale: hotspot sizes in nn workloads are model-derived and
+ * predictable (head_dim, seq_len, n_heads). Power-of-two slabs cap
+ * fragmentation at 2× per buffer with no hash-map churn — strictly better
+ * than a free-list-by-exact-size for these access patterns.
+ */
+public class SizeClassedScratchPool(
+    private val maxBuffersPerClass: Int = 8
+) : ScratchPool {
+
+    private val classes: Array<ArrayDeque<FloatArray>> =
+        Array(MAX_CLASSES) { ArrayDeque() }
+
+    /** Stack of scope frames; each frame is the list of buffers acquired in
+     *  that scope. `addLast` on enter, drain on exit. */
+    private val scopeStack: ArrayDeque<ArrayDeque<FloatArray>> = ArrayDeque()
+
+    private var acquireCount: Long = 0L
+    private var cacheHits: Long = 0L
+    private var highWaterBytes: Long = 0L
+    private var currentBytes: Long = 0L
+
+    override fun acquireFloat(minSize: Int): FloatArray {
+        val cls = sizeClass(minSize)
+        val cache = classes[cls]
+        acquireCount++
+        val buf = if (cache.isNotEmpty()) {
+            cacheHits++
+            cache.removeLast()
+        } else {
+            FloatArray(sizeForClass(cls))
+        }
+        scopeStack.lastOrNull()?.addLast(buf)
+        currentBytes += buf.size.toLong() * 4L
+        if (currentBytes > highWaterBytes) highWaterBytes = currentBytes
+        return buf
+    }
+
+    override fun acquireFloatZeroed(minSize: Int): FloatArray {
+        val buf = acquireFloat(minSize)
+        buf.fill(0f, 0, minSize)
+        return buf
+    }
+
+    override fun <R> scope(block: () -> R): R {
+        val frame: ArrayDeque<FloatArray> = ArrayDeque()
+        scopeStack.addLast(frame)
+        try {
+            return block()
+        } finally {
+            scopeStack.removeLast()
+            for (buf in frame) returnToCache(buf)
+        }
+    }
+
+    override fun stats(): ScratchStats {
+        var active = 0
+        for (frame in scopeStack) active += frame.size
+        return ScratchStats(acquireCount, cacheHits, highWaterBytes, active)
+    }
+
+    private fun returnToCache(buf: FloatArray) {
+        currentBytes -= buf.size.toLong() * 4L
+        val cls = sizeClass(buf.size)
+        val cache = classes[cls]
+        if (cache.size < maxBuffersPerClass) cache.addLast(buf)
+    }
+
+    public companion object {
+        public const val MIN_SIZE: Int = 64
+        public const val LOG_MIN_SIZE: Int = 6
+        public const val MAX_CLASSES: Int = 20
+
+        /**
+         * Bucket index for a request of [minSize] floats. Returns 0 for any
+         * `minSize <= MIN_SIZE`; otherwise the smallest class whose size is
+         * `>= minSize`. Capped at [MAX_CLASSES] - 1 — requests beyond that
+         * still allocate an array of the requested rounded size, but reuse
+         * is bounded by the top class.
+         */
+        public fun sizeClass(minSize: Int): Int {
+            if (minSize <= MIN_SIZE) return 0
+            val cls = (32 - (minSize - 1).countLeadingZeroBits()) - LOG_MIN_SIZE
+            return cls.coerceIn(0, MAX_CLASSES - 1)
+        }
+
+        /** Floats in bucket [cls]. */
+        public fun sizeForClass(cls: Int): Int = MIN_SIZE shl cls
+    }
+}

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/scratch/SizeClassedScratchPoolTest.kt

@@ -0,0 +1,133 @@
+package sk.ainet.lang.tensor.scratch
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotSame
+import kotlin.test.assertSame
+import kotlin.test.assertTrue
+
+class SizeClassedScratchPoolTest {
+
+    @Test
+    fun sizeClassRoundsToNextPowerOfTwo() {
+        assertEquals(0, SizeClassedScratchPool.sizeClass(1))
+        assertEquals(0, SizeClassedScratchPool.sizeClass(64))
+        assertEquals(1, SizeClassedScratchPool.sizeClass(65))
+        assertEquals(1, SizeClassedScratchPool.sizeClass(128))
+        assertEquals(2, SizeClassedScratchPool.sizeClass(129))
+        assertEquals(2, SizeClassedScratchPool.sizeClass(256))
+        assertEquals(3, SizeClassedScratchPool.sizeClass(257))
+    }
+
+    @Test
+    fun sizeForClassReturnsBucketCapacity() {
+        assertEquals(64, SizeClassedScratchPool.sizeForClass(0))
+        assertEquals(128, SizeClassedScratchPool.sizeForClass(1))
+        assertEquals(256, SizeClassedScratchPool.sizeForClass(2))
+    }
+
+    @Test
+    fun bufferIsAtLeastRequestedSize() {
+        val pool = SizeClassedScratchPool()
+        pool.scope {
+            assertTrue(pool.acquireFloat(1).size >= 1)
+            assertTrue(pool.acquireFloat(100).size >= 100)
+            assertTrue(pool.acquireFloat(1000).size >= 1000)
+        }
+    }
+
+    @Test
+    fun scopeRecyclesBuffers() {
+        val pool = SizeClassedScratchPool()
+        val first: FloatArray = pool.scope { pool.acquireFloat(100) }
+        val second: FloatArray = pool.scope { pool.acquireFloat(100) }
+        // Same bucket, recycled buffer comes back.
+        assertSame(first, second)
+    }
+
+    @Test
+    fun differentBucketsDoNotShare() {
+        val pool = SizeClassedScratchPool()
+        val small: FloatArray = pool.scope { pool.acquireFloat(50) }
+        val large: FloatArray = pool.scope { pool.acquireFloat(500) }
+        assertNotSame(small, large)
+    }
+
+    @Test
+    fun acquireZeroedClearsRequestedRange() {
+        val pool = SizeClassedScratchPool()
+        pool.scope {
+            val buf = pool.acquireFloat(100)
+            for (i in 0 until 100) buf[i] = 7f
+        }
+        pool.scope {
+            val buf = pool.acquireFloatZeroed(100)
+            for (i in 0 until 100) assertEquals(0f, buf[i])
+        }
+    }
+
+    @Test
+    fun acquireOutsideScopeStillAllocates() {
+        val pool = SizeClassedScratchPool()
+        val buf = pool.acquireFloat(10)
+        assertTrue(buf.size >= 10)
+    }
+
+    @Test
+    fun nestedScopesBalance() {
+        val pool = SizeClassedScratchPool()
+        var innerRef: FloatArray? = null
+        pool.scope {
+            val outer = pool.acquireFloat(100)
+            pool.scope {
+                innerRef = pool.acquireFloat(100)
+                assertNotSame(outer, innerRef)
+            }
+            // Inner-scope buffer was recycled on inner-scope exit — the next
+            // acquire in the outer scope at the same size class returns it.
+            val recycled = pool.acquireFloat(100)
+            assertSame(innerRef, recycled)
+        }
+    }
+
+    @Test
+    fun statsReportAcquiresAndHits() {
+        val pool = SizeClassedScratchPool()
+        pool.scope { pool.acquireFloat(100) }
+        pool.scope { pool.acquireFloat(100) } // hit
+        pool.scope { pool.acquireFloat(100) } // hit
+        val stats = pool.stats()
+        assertEquals(3L, stats.acquireCount)
+        assertEquals(2L, stats.cacheHits)
+    }
+
+    @Test
+    fun surplusBuffersDropped() {
+        val pool = SizeClassedScratchPool(maxBuffersPerClass = 2)
+        // Acquire 3 buffers in one scope; on exit, only 2 are retained
+        // (the third is dropped because the per-class cap is 2).
+        pool.scope {
+            pool.acquireFloat(100)
+            pool.acquireFloat(100)
+            pool.acquireFloat(100)
+        }
+        // Acquire 3 buffers in a second scope. First 2 hit the cache; the
+        // 3rd misses because the cache only retained 2 from the prior scope.
+        pool.scope {
+            pool.acquireFloat(100)
+            pool.acquireFloat(100)
+            pool.acquireFloat(100)
+        }
+        val stats = pool.stats()
+        assertEquals(6L, stats.acquireCount)
+        assertEquals(2L, stats.cacheHits)
+    }
+
+    @Test
+    fun noopPoolAlwaysAllocates() {
+        val pool = NoopScratchPool
+        val first = pool.scope { pool.acquireFloat(100) }
+        val second = pool.scope { pool.acquireFloat(100) }
+        assertNotSame(first, second)
+    }
+}