From 6c76883a7d092dd8d17e111a0ce7f7f7fdca623b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 28 Apr 2026 09:12:03 +0200 Subject: [PATCH] feat(scratch): add ScratchPool SPI for runtime workspace allocation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #549. Adds a generic workspace allocator for short-lived FloatArray buffers used by attention scratch, RoPE tables, KV-cache slice copies, padding scratch, and any other nn workload that allocates intermediates per forward step. * `sk.ainet.lang.tensor.scratch.ScratchPool` — SPI with `acquireFloat`, `acquireFloatZeroed`, `scope { ... }`, and `stats()`. * `NoopScratchPool` — default no-op pool; every acquire allocates fresh. Bit-for-bit equivalent to pre-pool behavior. * `SizeClassedScratchPool` — power-of-two slabs starting at 64 floats; scoped lifetime; per-class cap with surplus drop. Single-threaded by intent (one forward at a time per pool); concurrent forwards use separate pools. * `ExecutionContext.scratch: ScratchPool` — backward-compatible accessor with default `NoopScratchPool`. Existing impls keep working unchanged. Out of scope here: per-thread ambient carrier (only needed where call sites don't take ExecutionContext); direct-memory variants. Bumps version to 0.21.0-SNAPSHOT for downstream composite-build integration. Co-Authored-By: Claude Opus 4.7 (1M context) --- gradle.properties | 2 +- .../sk/ainet/context/ExecutionContext.kt | 16 ++ .../ainet/lang/tensor/scratch/ScratchPool.kt | 175 ++++++++++++++++++ .../scratch/SizeClassedScratchPoolTest.kt | 133 +++++++++++++ 4 files changed, 325 insertions(+), 1 deletion(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/scratch/ScratchPool.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/scratch/SizeClassedScratchPoolTest.kt diff --git a/gradle.properties b/gradle.properties index b638262ba..947dc968c 100644 --- a/gradle.properties +++ b/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/ 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 1ae290718..e5fdcdcd0 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 @@ -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 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/scratch/ScratchPool.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/scratch/ScratchPool.kt new file mode 100644 index 000000000..714b16b00 --- /dev/null +++ b/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 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 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> = + 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() + + 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 scope(block: () -> R): R { + val frame: ArrayDeque = 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 + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/scratch/SizeClassedScratchPoolTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/scratch/SizeClassedScratchPoolTest.kt new file mode 100644 index 000000000..ee0660a41 --- /dev/null +++ b/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) + } +}