Skip to content

Commit 380e7c5

Browse files
michalharakalmichalharakal
authored
Merge pull request #550 from SKaiNET-developers/feature/ISSUE-549-scratch-pool
feat(scratch): add ScratchPool SPI for runtime workspace allocation
2 parents ba6f2ab + 6c76883 commit 380e7c5

4 files changed

Lines changed: 325 additions & 1 deletion

File tree

‎gradle.properties‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
GROUP=sk.ainet.core
2-
VERSION_NAME=0.20.0
2+
VERSION_NAME=0.21.0-SNAPSHOT
33
POM_DESCRIPTION=SKaiNET
44

55
POM_URL=https://github.com/SKaiNET-developers/skainet/

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import sk.ainet.lang.tensor.data.TensorData
66
import sk.ainet.lang.tensor.data.TensorDataFactory
77
import sk.ainet.lang.tensor.operators.OpsBoundTensor
88
import sk.ainet.lang.tensor.ops.TensorOps
9+
import sk.ainet.lang.tensor.scratch.NoopScratchPool
10+
import sk.ainet.lang.tensor.scratch.ScratchPool
911
import sk.ainet.lang.tensor.storage.MemoryPlanner
1012
import sk.ainet.lang.tensor.storage.MemoryTracker
1113
import sk.ainet.lang.types.DType
@@ -23,6 +25,20 @@ public interface ExecutionContext {
2325

2426
public val tensorDataFactory: TensorDataFactory
2527

28+
/**
29+
* Workspace allocator for short-lived intermediate buffers (attention
30+
* scratch, RoPE tables, KV-cache slice copies, padding scratch, etc.).
31+
*
32+
* Default is [NoopScratchPool] — every acquire allocates a fresh array,
33+
* matching pre-pool behavior. Implementations that want pooling override
34+
* this property (or wrap an existing context).
35+
*
36+
* Callers MUST acquire inside an active [ScratchPool.scope] block;
37+
* acquires outside a scope succeed but the buffer is not returned to the
38+
* pool when dropped.
39+
*/
40+
public val scratch: ScratchPool get() = NoopScratchPool
41+
2642
// Execution observers for tracing/benchmarking
2743
public val observers: ExecutionObserverRegistry
2844

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

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
package sk.ainet.lang.tensor.scratch
2+
3+
/**
4+
* Pool of reusable [FloatArray] scratch buffers, scoped to a single forward
5+
* (or backward) pass.
6+
*
7+
* Workspace allocation is generic across nn workloads — attention, RoPE,
8+
* convolutions, embedding gathers, training-time gradient buffers all need
9+
* short-lived intermediates. Routing those through a pool eliminates per-step
10+
* allocation pressure on the GC heap.
11+
*
12+
* Typical lifecycle:
13+
*
14+
* 1. The runtime owns a [ScratchPool] (one per forward-running thread).
15+
* 2. The runtime opens a [scope] block around each forward pass.
16+
* 3. Layers acquire buffers via [acquireFloat] / [acquireFloatZeroed] from
17+
* the [sk.ainet.context.ExecutionContext.scratch] field.
18+
* 4. On scope exit every buffer acquired in the scope is recycled to the
19+
* per-size-class free list for the next pass.
20+
*
21+
* Buffers may be larger than the requested size — callers must always index
22+
* within the size they asked for, not `buf.size`. This is the standard
23+
* scratch-buffer contract: external `outSize`/`n`/`stride` variables drive
24+
* iteration; the array's own `.size` is an implementation detail.
25+
*
26+
* Acquires made *outside* an active scope are not tracked and become regular
27+
* allocations (the buffer is not returned to the pool when the caller drops
28+
* it). This is intentional — it lets the no-op pool implementation avoid
29+
* any scope bookkeeping.
30+
*
31+
* Single-threaded by intent: one forward pass at a time per pool.
32+
* Concurrent forwards must use separate pools.
33+
*/
34+
public interface ScratchPool {
35+
/**
36+
* Acquire a [FloatArray] with at least [minSize] elements. Contents are
37+
* unspecified — callers must overwrite the range they read.
38+
*/
39+
public fun acquireFloat(minSize: Int): FloatArray
40+
41+
/**
42+
* Acquire a [FloatArray] with at least [minSize] elements, with the
43+
* range `[0, minSize)` zero-filled. Use when the caller relies on
44+
* sparse-write zeros (e.g. padding a smaller block into a larger
45+
* destination and leaving the gap implicit).
46+
*/
47+
public fun acquireFloatZeroed(minSize: Int): FloatArray
48+
49+
/**
50+
* Open a forward-pass scope. All buffers acquired inside [block] are
51+
* recycled at exit. Scopes may be nested; each open/close balances.
52+
*/
53+
public fun <R> scope(block: () -> R): R
54+
55+
/** Stats for benchmarking and leak detection. */
56+
public fun stats(): ScratchStats
57+
}
58+
59+
public data class ScratchStats(
60+
val acquireCount: Long,
61+
val cacheHits: Long,
62+
val highWaterBytes: Long,
63+
val activeBuffers: Int
64+
)
65+
66+
/**
67+
* No-op pool: every [acquireFloat] / [acquireFloatZeroed] allocates a fresh
68+
* array; [scope] just runs the block. Default carrier returned by
69+
* [sk.ainet.context.ExecutionContext.scratch] when no pooling is configured —
70+
* preserves pre-pool behavior bit-for-bit.
71+
*/
72+
public object NoopScratchPool : ScratchPool {
73+
override fun acquireFloat(minSize: Int): FloatArray = FloatArray(minSize)
74+
override fun acquireFloatZeroed(minSize: Int): FloatArray = FloatArray(minSize)
75+
override fun <R> scope(block: () -> R): R = block()
76+
override fun stats(): ScratchStats = ScratchStats(0L, 0L, 0L, 0)
77+
}
78+
79+
/**
80+
* Size-classed slab pool with power-of-two buckets starting at 64 floats.
81+
*
82+
* Sizes round up to the next power of two: `1..64 → 64`, `65..128 → 128`,
83+
* ... up to [MAX_CLASSES] classes (`64 * 2^19 ≈ 33M floats` = 128 MB at the
84+
* top of the range). Up to [maxBuffersPerClass] buffers are retained per
85+
* class; surplus buffers are dropped to GC at scope exit.
86+
*
87+
* Choice rationale: hotspot sizes in nn workloads are model-derived and
88+
* predictable (head_dim, seq_len, n_heads). Power-of-two slabs cap
89+
* fragmentation at 2× per buffer with no hash-map churn — strictly better
90+
* than a free-list-by-exact-size for these access patterns.
91+
*/
92+
public class SizeClassedScratchPool(
93+
private val maxBuffersPerClass: Int = 8
94+
) : ScratchPool {
95+
96+
private val classes: Array<ArrayDeque<FloatArray>> =
97+
Array(MAX_CLASSES) { ArrayDeque() }
98+
99+
/** Stack of scope frames; each frame is the list of buffers acquired in
100+
* that scope. `addLast` on enter, drain on exit. */
101+
private val scopeStack: ArrayDeque<ArrayDeque<FloatArray>> = ArrayDeque()
102+
103+
private var acquireCount: Long = 0L
104+
private var cacheHits: Long = 0L
105+
private var highWaterBytes: Long = 0L
106+
private var currentBytes: Long = 0L
107+
108+
override fun acquireFloat(minSize: Int): FloatArray {
109+
val cls = sizeClass(minSize)
110+
val cache = classes[cls]
111+
acquireCount++
112+
val buf = if (cache.isNotEmpty()) {
113+
cacheHits++
114+
cache.removeLast()
115+
} else {
116+
FloatArray(sizeForClass(cls))
117+
}
118+
scopeStack.lastOrNull()?.addLast(buf)
119+
currentBytes += buf.size.toLong() * 4L
120+
if (currentBytes > highWaterBytes) highWaterBytes = currentBytes
121+
return buf
122+
}
123+
124+
override fun acquireFloatZeroed(minSize: Int): FloatArray {
125+
val buf = acquireFloat(minSize)
126+
buf.fill(0f, 0, minSize)
127+
return buf
128+
}
129+
130+
override fun <R> scope(block: () -> R): R {
131+
val frame: ArrayDeque<FloatArray> = ArrayDeque()
132+
scopeStack.addLast(frame)
133+
try {
134+
return block()
135+
} finally {
136+
scopeStack.removeLast()
137+
for (buf in frame) returnToCache(buf)
138+
}
139+
}
140+
141+
override fun stats(): ScratchStats {
142+
var active = 0
143+
for (frame in scopeStack) active += frame.size
144+
return ScratchStats(acquireCount, cacheHits, highWaterBytes, active)
145+
}
146+
147+
private fun returnToCache(buf: FloatArray) {
148+
currentBytes -= buf.size.toLong() * 4L
149+
val cls = sizeClass(buf.size)
150+
val cache = classes[cls]
151+
if (cache.size < maxBuffersPerClass) cache.addLast(buf)
152+
}
153+
154+
public companion object {
155+
public const val MIN_SIZE: Int = 64
156+
public const val LOG_MIN_SIZE: Int = 6
157+
public const val MAX_CLASSES: Int = 20
158+
159+
/**
160+
* Bucket index for a request of [minSize] floats. Returns 0 for any
161+
* `minSize <= MIN_SIZE`; otherwise the smallest class whose size is
162+
* `>= minSize`. Capped at [MAX_CLASSES] - 1 — requests beyond that
163+
* still allocate an array of the requested rounded size, but reuse
164+
* is bounded by the top class.
165+
*/
166+
public fun sizeClass(minSize: Int): Int {
167+
if (minSize <= MIN_SIZE) return 0
168+
val cls = (32 - (minSize - 1).countLeadingZeroBits()) - LOG_MIN_SIZE
169+
return cls.coerceIn(0, MAX_CLASSES - 1)
170+
}
171+
172+
/** Floats in bucket [cls]. */
173+
public fun sizeForClass(cls: Int): Int = MIN_SIZE shl cls
174+
}
175+
}

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

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package sk.ainet.lang.tensor.scratch
2+
3+
import kotlin.test.Test
4+
import kotlin.test.assertEquals
5+
import kotlin.test.assertNotSame
6+
import kotlin.test.assertSame
7+
import kotlin.test.assertTrue
8+
9+
class SizeClassedScratchPoolTest {
10+
11+
@Test
12+
fun sizeClassRoundsToNextPowerOfTwo() {
13+
assertEquals(0, SizeClassedScratchPool.sizeClass(1))
14+
assertEquals(0, SizeClassedScratchPool.sizeClass(64))
15+
assertEquals(1, SizeClassedScratchPool.sizeClass(65))
16+
assertEquals(1, SizeClassedScratchPool.sizeClass(128))
17+
assertEquals(2, SizeClassedScratchPool.sizeClass(129))
18+
assertEquals(2, SizeClassedScratchPool.sizeClass(256))
19+
assertEquals(3, SizeClassedScratchPool.sizeClass(257))
20+
}
21+
22+
@Test
23+
fun sizeForClassReturnsBucketCapacity() {
24+
assertEquals(64, SizeClassedScratchPool.sizeForClass(0))
25+
assertEquals(128, SizeClassedScratchPool.sizeForClass(1))
26+
assertEquals(256, SizeClassedScratchPool.sizeForClass(2))
27+
}
28+
29+
@Test
30+
fun bufferIsAtLeastRequestedSize() {
31+
val pool = SizeClassedScratchPool()
32+
pool.scope {
33+
assertTrue(pool.acquireFloat(1).size >= 1)
34+
assertTrue(pool.acquireFloat(100).size >= 100)
35+
assertTrue(pool.acquireFloat(1000).size >= 1000)
36+
}
37+
}
38+
39+
@Test
40+
fun scopeRecyclesBuffers() {
41+
val pool = SizeClassedScratchPool()
42+
val first: FloatArray = pool.scope { pool.acquireFloat(100) }
43+
val second: FloatArray = pool.scope { pool.acquireFloat(100) }
44+
// Same bucket, recycled buffer comes back.
45+
assertSame(first, second)
46+
}
47+
48+
@Test
49+
fun differentBucketsDoNotShare() {
50+
val pool = SizeClassedScratchPool()
51+
val small: FloatArray = pool.scope { pool.acquireFloat(50) }
52+
val large: FloatArray = pool.scope { pool.acquireFloat(500) }
53+
assertNotSame(small, large)
54+
}
55+
56+
@Test
57+
fun acquireZeroedClearsRequestedRange() {
58+
val pool = SizeClassedScratchPool()
59+
pool.scope {
60+
val buf = pool.acquireFloat(100)
61+
for (i in 0 until 100) buf[i] = 7f
62+
}
63+
pool.scope {
64+
val buf = pool.acquireFloatZeroed(100)
65+
for (i in 0 until 100) assertEquals(0f, buf[i])
66+
}
67+
}
68+
69+
@Test
70+
fun acquireOutsideScopeStillAllocates() {
71+
val pool = SizeClassedScratchPool()
72+
val buf = pool.acquireFloat(10)
73+
assertTrue(buf.size >= 10)
74+
}
75+
76+
@Test
77+
fun nestedScopesBalance() {
78+
val pool = SizeClassedScratchPool()
79+
var innerRef: FloatArray? = null
80+
pool.scope {
81+
val outer = pool.acquireFloat(100)
82+
pool.scope {
83+
innerRef = pool.acquireFloat(100)
84+
assertNotSame(outer, innerRef)
85+
}
86+
// Inner-scope buffer was recycled on inner-scope exit — the next
87+
// acquire in the outer scope at the same size class returns it.
88+
val recycled = pool.acquireFloat(100)
89+
assertSame(innerRef, recycled)
90+
}
91+
}
92+
93+
@Test
94+
fun statsReportAcquiresAndHits() {
95+
val pool = SizeClassedScratchPool()
96+
pool.scope { pool.acquireFloat(100) }
97+
pool.scope { pool.acquireFloat(100) } // hit
98+
pool.scope { pool.acquireFloat(100) } // hit
99+
val stats = pool.stats()
100+
assertEquals(3L, stats.acquireCount)
101+
assertEquals(2L, stats.cacheHits)
102+
}
103+
104+
@Test
105+
fun surplusBuffersDropped() {
106+
val pool = SizeClassedScratchPool(maxBuffersPerClass = 2)
107+
// Acquire 3 buffers in one scope; on exit, only 2 are retained
108+
// (the third is dropped because the per-class cap is 2).
109+
pool.scope {
110+
pool.acquireFloat(100)
111+
pool.acquireFloat(100)
112+
pool.acquireFloat(100)
113+
}
114+
// Acquire 3 buffers in a second scope. First 2 hit the cache; the
115+
// 3rd misses because the cache only retained 2 from the prior scope.
116+
pool.scope {
117+
pool.acquireFloat(100)
118+
pool.acquireFloat(100)
119+
pool.acquireFloat(100)
120+
}
121+
val stats = pool.stats()
122+
assertEquals(6L, stats.acquireCount)
123+
assertEquals(2L, stats.cacheHits)
124+
}
125+
126+
@Test
127+
fun noopPoolAlwaysAllocates() {
128+
val pool = NoopScratchPool
129+
val first = pool.scope { pool.acquireFloat(100) }
130+
val second = pool.scope { pool.acquireFloat(100) }
131+
assertNotSame(first, second)
132+
}
133+
}

0 commit comments

Comments
 (0)