From 185f847878477dbfc5dc51e7ac73b5cbe101ef9b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 11 Nov 2026 14:20:59 +0100 Subject: [PATCH 1/3] docs(skeep): SKEEP-003 status matches the code (#1500) An audit at 0.64.0 found the old storage layer still in place beside the new one. The SKEEP goes back to Accepted and gains a status table per acceptance criterion and per open decision, linking the issues that track them. 003a, SKEEP-002, memory-model.adoc and the TensorStorage KDoc lose statements that later work made untrue. Co-Authored-By: Claude Fable 5.2 --- CHANGELOG.md | 4 ++ .../ROOT/pages/explanation/memory-model.adoc | 9 +-- .../002-android-offheap-tensor-storage.adoc | 33 +++++------ .../pages/003-unified-tensor-storage.adoc | 57 ++++++++++++++++++- ...03a-placement-and-planning-resolution.adoc | 8 +-- docs/modules/skeep/pages/index.adoc | 6 +- .../lang/tensor/storage/TensorStorage.kt | 13 +++-- 7 files changed, 95 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ac786a7..4652f488 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- **SKEEP-003 is back at `Accepted`, with an implementation-status table** (docs, #1500): an audit at 0.64.0 found the new storage model in use and the old layer (`TensorStorage`, `BufferHandle`, `TensorData` downcasts in the ops) still in place, so three of the five acceptance criteria are open. The SKEEP now says which, and what tracks each (#932). SKEEP-003a no longer lists #1148 as open, SKEEP-002's status records the Q4_K / Q6_K mapped staging of #1190 and its device measurement, `memory-model.adoc` drops a limit that #1190 removed, and the `TensorStorage` KDoc stops telling new code to target it. + ## [0.64.0] - 2026-11-11 Headline: **every differentiable op has a backward rule, and mixed precision is finished: it trains, exports without diff --git a/docs/modules/ROOT/pages/explanation/memory-model.adoc b/docs/modules/ROOT/pages/explanation/memory-model.adoc index 8aa0e8d7..c6b72d02 100644 --- a/docs/modules/ROOT/pages/explanation/memory-model.adoc +++ b/docs/modules/ROOT/pages/explanation/memory-model.adoc @@ -204,7 +204,7 @@ ARMv8.2 Cortex-A55 reference board. [cols="2,3",options="header"] |=== | Delivered | Notes -| `Storage` / `Scope` / `TensorView` / `Format` / `Layout` | All targets; `TensorData` implementations are façades over views +| `Storage` / `Scope` / `TensorView` / `Format` / `Layout` | All targets; every `TensorData` family used by the kernels exposes a view. Scoped activations (`StorageFloatTensorData`) sit on a `Storage`; the other classes still own their bytes and wrap the view around them (xref:skeep:003-unified-tensor-storage.adoc[SKEEP-003], implementation status) | `KernelKey` dispatch for matmul | Reference kernel plus packed and ternary kernel packs | Ternary encodings and `bitnet_gemv` | `TQ1_0`, `TQ2_0`, BitNet b1.58; reference kernel plus a NEON pack on Android | Memory plan, device profiles, fit check | `skainet-plan`, mobile/desktop/native profiles, two-pool device check @@ -247,9 +247,10 @@ anything is loaded rather than after an OOM. Known limits, stated rather than implied: -* *Packed weights still reach the managed heap under mapped staging.* The packed matmul SPI takes -`ByteArray`s, so a buffer-aware kernel is future work. Mapping therefore lifts the Android heap -ceiling for dense checkpoints, not yet for a Q4_K_M one. +* *Some packed weights still reach the managed heap under mapped staging.* Q4_K and Q6_K weights +stay in the file mapping (`BufferPackedTensorData`, read by row-major kernels through a direct +buffer, #1190), so a Q4_K_M checkpoint larger than the Android heap loads. The other packed +formats are still handed to their kernels as `ByteArray`s. * *Whole-file mapping only.* Files larger than 2 GB are refused rather than mapped in windows. * *Graph-level planning* (allocating a whole forward pass at once) stays in the downstream compiler by decision (#1134). Device placement itself was resolved in 0.49.0: `AllocationResolver` diff --git a/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc b/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc index 40f44773..44cd5988 100644 --- a/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc +++ b/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc @@ -232,14 +232,13 @@ keeps working unchanged and the overload is purely opt-in. Expected consumers: (packed quant blocks) and a synthesized SafeTensors file (dense F32/F16/BF16) produce bit-identical tensors under either placement. -== Implementation status (updated 2026-08-26, for 0.49.0) +== Implementation status (updated 2026-11-11, at 0.64.0) Tracked by SKEEP-003's M2 slice https://github.com/SKaiNET-developers/SKaiNET/issues/1038[#1038]. This SKEEP stays *Draft*: the mechanism is in `develop` (and since 0.49.0 the loader speaks one `WeightForm` — the legacy -`StagingPolicy`/`QuantPolicy`/`WeightOrientation` axes were removed in #1159), the acceptance -criteria that need a physical device are not met yet, and one of them cannot be met until an -unrelated contract is fixed. +`StagingPolicy`/`QuantPolicy`/`WeightOrientation` axes were removed in #1159), one acceptance +criterion has a device measurement behind it, and the others are not recorded. Landed: @@ -257,20 +256,22 @@ Landed: `AndroidGguf.fits(context, path, ctx)` answers from the GGUF header before a byte of payload is read, naming the pool that runs out and what to do about it. +* *Phase 4 for Q4_K and Q6_K* — `BufferPackedTensorData` keeps a packed weight in the file + mapping, and row-major kernels read it through a direct `ByteBuffer` + (https://github.com/SKaiNET-developers/SKaiNET/issues/1189[#1189], PR #1190). Measured on a Pixel 8a with the 256 MB heap cap + (https://github.com/SKaiNET-developers/SKaiNET/issues/1130[#1130]): Qwen2.5-1.5B Q4_K_M, 1.0 GB on disk, loads with 566 KB of + weights on the managed heap and decodes at 153 ms per step; before, it ran out of memory at + load. That is the acceptance criterion "a Q4_K model larger than the heap loads and generates". + Not landed, and why: -* *Packed weights still reach the managed heap.* Mapped staging serves dense F32 tensors as - zero-heap views, but every quantized tensor is still handed to its kernel as a `ByteArray`, - because the packed matmul SPI takes arrays. That is Phase 4, and it is blocked on - https://github.com/SKaiNET-developers/SKaiNET/issues/973[#973] — the packed-quant byte-order - contract — since a buffer-aware kernel first needs an unambiguous answer to *which* byte order it - is reading. Until then the heap ceiling is lifted for dense checkpoints, not for a Q4_K_M one, - which is exactly the acceptance criterion below that remains open. -* *The device numbers.* "≤ 40 MB managed heap for SmolLM2-135M Q8_0" and "a ~600 MB Q4_K model - loads on a 256 MB heap" are measurements on a physical device; the repository's CI has none. They - belong to the `skainet-decode` sample in SKaiNET-transformers, which owns a model and a device - lane. `androidHostTest` covers what can be proven without hardware: the Android compilation - builds, maps, loads, and produces bit-identical tensors under either staging. +* *Other packed formats still reach the managed heap.* Q8_0, Q4_0, Q5_0, Q5_1, Q5_K and the + ternary encodings are handed to their kernels as `ByteArray`s under mapped staging; only Q4_K + and Q6_K have the buffer-reading kernels. +* *The SmolLM2-135M Q8_0 heap figure and the mapped-versus-heap throughput comparison* are not + recorded. Q8_0 is one of the formats above, so the "less than 40 MB" criterion cannot hold yet. +* *Whole-file mapping only.* A file larger than 2 GB is refused instead of being mapped in + windows (https://github.com/SKaiNET-developers/SKaiNET/issues/1422[#1422] asks for the same in the IRPA writer). == Risks diff --git a/docs/modules/skeep/pages/003-unified-tensor-storage.adoc b/docs/modules/skeep/pages/003-unified-tensor-storage.adoc index 86a75a0c..09d7bc47 100644 --- a/docs/modules/skeep/pages/003-unified-tensor-storage.adoc +++ b/docs/modules/skeep/pages/003-unified-tensor-storage.adoc @@ -1,11 +1,64 @@ = SKEEP-003: Unifying the tensor storage model — one byte-owner, enforced ownership, coherent dtype/encoding :description: SKaiNET proposal to converge the TensorData and TensorStorage layers into one storage model with real ownership, a single view mechanism, and dtype/encoding coherence. -Status: Implemented (accepted 2026-08-22; M0/M1/M2 + the P7/P8 resolution shipped in 0.49.0) + +Status: Accepted (2026-08-22; M0/M1/M2 and the placement resolution shipped in 0.49.0; the removal of the old layer is open, see <>) + Audience: SKaiNET maintainers and contributors + Created: 2026-08-10 + Tracking issue: https://github.com/SKaiNET-developers/SKaiNET/issues/932[#932] (umbrella) → https://github.com/SKaiNET-developers/SKaiNET/issues/1001[#1001 M0] · https://github.com/SKaiNET-developers/SKaiNET/issues/1002[#1002 M1] · https://github.com/SKaiNET-developers/SKaiNET/issues/1003[#1003 M2] + -Implemented: the model as built is documented in xref:ROOT:explanation/memory-model.adoc[The memory model] and xref:ROOT:explanation/packed-weight-layout.adoc[Packed weight layout]; the proposal and milestone PRD that led there are in the pull requests and issues linked below (#932, #1001, #1002, #1003). +As built: the model is documented in xref:ROOT:explanation/memory-model.adoc[The memory model] and xref:ROOT:explanation/packed-weight-layout.adoc[Packed weight layout]; the proposal and milestone PRD that led there are in the pull requests and issues linked below (#932, #1001, #1002, #1003). + +[#implementation-status] +== Implementation status (audited 2026-11-11, at 0.64.0) + +The model is built and in use on every target: `Storage`, `TensorView`, `Format`, `Scope`, the +memory plan, the tracer, mapped staging. The old layer it was meant to replace is still in the +code, so this proposal is back at `Accepted` until the open rows below are closed. The umbrella +https://github.com/SKaiNET-developers/SKaiNET/issues/932[#932] lists the work. + +[cols="3,1,3",options="header"] +|=== +| Acceptance criterion | State | Where it stands + +| One abstraction owns bytes, ownership enforced +| Partly +| Enforcement is tested (`StorageTest`, `ScopeTest`, `TensorViewTest`, `JvmStorageTest`). `TensorStorage` and `BufferHandle` still exist beside `Storage`, and the streaming readers' `loadTensorStorage*` return the old type (https://github.com/SKaiNET-developers/SKaiNET/issues/1502[#1502], https://github.com/SKaiNET-developers/SKaiNET/issues/1507[#1507]). + +| One view mechanism +| Partly +| `OneViewMechanismTest` covers slice, transpose and unsqueeze on the new view. The public `slice()` family still returns the deprecated `SlicedTensorView` (https://github.com/SKaiNET-developers/SKaiNET/issues/1502[#1502]). + +| A packed tensor reports FP32 and its encoding; no dispatch site downcasts to a concrete `TensorData` class outside the data package +| Not met +| The first half holds (`FormatTest`). About 130 downcasts remain in the CPU, Accelerate and wasm ops, and matmul reaches `KernelKey` dispatch only as its last fallback (https://github.com/SKaiNET-developers/SKaiNET/issues/1501[#1501] for the guard, https://github.com/SKaiNET-developers/SKaiNET/issues/1507[#1507] for the migration). + +| No kernel regression against a baseline +| Partly +| The benchmarks exist; the recorded baseline was deleted with `docs/design/` in #1106 (https://github.com/SKaiNET-developers/SKaiNET/issues/1501[#1501]). + +| Flat direct memory under a forward scope on the JVM +| No test +| Slab reuse is tested (`ScopedOpOutputsTest`), direct memory only under the GC arena (`MemSegArenaLeakTest`) (https://github.com/SKaiNET-developers/SKaiNET/issues/1501[#1501]). +|=== + +Decisions 1, 2 (naming and package), 6, 9, 11, 12 and 13 are delivered; `LogicalDType` is +deprecated and its removal is https://github.com/SKaiNET-developers/SKaiNET/issues/1502[#1502]. The rest: + +[cols="1,4,3",options="header"] +|=== +| # | Open part | Tracked by + +| 1 | Façades: only `StorageFloatTensorData` sits on a `Storage`; the other `TensorData` classes own their bytes and wrap a view around them on each read | https://github.com/SKaiNET-developers/SKaiNET/issues/1507[#1507] +| 2 | `sk.ainet.lang.graph.Layout` is not deprecated; `BlockOrderLayout` implements it now | SKEEP-004, which owns the compile-side layout +| 3 | `get()` on a packed view decodes; the debug-mode warning on hot paths does not exist | https://github.com/SKaiNET-developers/SKaiNET/issues/1503[#1503] +| 4 | The `(head, tail)` pair and its gather fallback exist (`WindowedKV`); only the reference attention consumes them | https://github.com/SKaiNET-developers/SKaiNET/issues/1504[#1504], outside this SKEEP +| 5 | `BlockSpec.activation` drives the eager dispatcher; the StableHLO export does not read it | https://github.com/SKaiNET-developers/SKaiNET/issues/1505[#1505], outside this SKEEP +| 7 | `SuspendingRandomAccessSource` exists in `skainet-io-core`; there is no remote module | https://github.com/SKaiNET-developers/SKaiNET/issues/1506[#1506], outside this SKEEP +| 8 | Axis labels on `Shape` | https://github.com/SKaiNET-developers/SKaiNET/issues/1503[#1503] +| 10 | Activation ids exist when tracing, not in eager debug mode | https://github.com/SKaiNET-developers/SKaiNET/issues/1503[#1503] +|=== + +Compiled parity (P7 below) was closed as superseded: IREE validation lives in the IREE-tools and +conformity repositories (#1148), and this repository has no eager-versus-compiled parity test. == Decision (2026-08-22) diff --git a/docs/modules/skeep/pages/003a-placement-and-planning-resolution.adoc b/docs/modules/skeep/pages/003a-placement-and-planning-resolution.adoc index 082c289d..9ef9a4e4 100644 --- a/docs/modules/skeep/pages/003a-placement-and-planning-resolution.adoc +++ b/docs/modules/skeep/pages/003a-placement-and-planning-resolution.adoc @@ -98,9 +98,9 @@ The carry work landed as the #1178/#1179/#1180 arc (0.49.0): `TensorRef` carries encoding and block order; `TensorSpec` metadata accessors follow the untyped precedent; the structural `skainet.tensor_layouts` module attribute and `ExternalParameterRef.blockOrder` declare the facts to the consumer; and `LayoutAssignmentPass` runs on the production path via -`HloGenerator.generate(target = …)`, populating the pre-built `ResolvedComputeGraph` seams. Until -the conformity pipeline exists (#1148, the future IREE-validation SKEEP), the carriage tests are -the metadata's reader. +`HloGenerator.generate(target = …)`, populating the pre-built `ResolvedComputeGraph` seams. The +conformity pipeline that reads the metadata lives outside this repository (#1148 was closed as +superseded); here, the carriage tests are the metadata's reader. == Where everything landed @@ -113,7 +113,7 @@ the metadata's reader. | `memoryScope` wired into creation | #1145 | PR #1156 | Op outputs through the scope + offset-aware SIMD | #1146, #1173 | landed (PRs #1168/#1175) | Layout/placement carriage to StableHLO | #1147 (#1178/#1179/#1180) | landed (PRs #1182/#1183/#1184) -| `.vmfb` parity acceptance run | #1148 | open (follow-up) +| `.vmfb` parity acceptance run | #1148 | closed as superseded 2026-09-02: IREE validation lives in the IREE-tools and conformity repositories; no eager-versus-compiled parity test in this repository |=== The SKEEP-003 items "`StorageSpec`, `MemoryPlanner` and the annotations are deleted unless wired" diff --git a/docs/modules/skeep/pages/index.adoc b/docs/modules/skeep/pages/index.adoc index 1083211b..9c2f381d 100644 --- a/docs/modules/skeep/pages/index.adoc +++ b/docs/modules/skeep/pages/index.adoc @@ -116,12 +116,12 @@ Every proposal should include: | Off-heap tensor storage on Android | xref:skeep:003-unified-tensor-storage.adoc[SKEEP-003] -| Implemented +| Accepted | Unifying the tensor storage model | xref:skeep:003a-placement-and-planning-resolution.adoc[SKEEP-003a] -| Resolved -| Placement, eager lifetime and graph-level planning — the P7/P8 answers +| Implemented +| Addendum to SKEEP-003: placement, eager lifetime and graph-level planning | xref:skeep:004-virtual-tensor-layout.adoc[SKEEP-004] | Draft diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt index e6b57530..e4af9e76 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorage.kt @@ -8,9 +8,9 @@ import sk.ainet.lang.types.DType /** * Runtime descriptor for a tensor's backing memory. * - * [TensorStorage] is the main architectural type that replaces ad-hoc - * array passing between loaders, planners, and backends. It carries enough - * information to handle a tensor without inspecting its bytes: + * [TensorStorage] describes a tensor read from a file before it becomes a + * tensor: the streaming GGUF and SafeTensors readers return it. It carries + * enough information to handle a tensor without inspecting its bytes: * * - **What** the values mean: [logicalType] * - **How** they are stored: [encoding] @@ -18,9 +18,10 @@ import sk.ainet.lang.types.DType * - **Layout**: [shape], [byteOffset], [strides], [isContiguous] * - **Ownership**: via [buffer]'s [BufferHandle] subtype * - * Existing [sk.ainet.lang.tensor.data.TensorData] remains as a - * compatibility façade. New loaders, planners, and backends should target - * [TensorStorage] directly. + * This is the older of two storage descriptions. The model new code targets is + * [sk.ainet.lang.memory.Storage] (owns bytes) with + * [sk.ainet.lang.memory.TensorView] (interprets them); SKEEP-003 replaces this + * type with a view once the readers return one (#1507). */ public data class TensorStorage( val shape: Shape, From 84db7cc48b83082138abe56e4992a7034d0e5ea5 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 11 Nov 2026 15:03:05 +0100 Subject: [PATCH 2/3] build: guards for SKEEP-003's acceptance criteria (#1501) checkTensorDataDowncasts counts downcasts to concrete TensorData classes outside the data package against an allowlist that has to match the code exactly: 162 sites in 21 files today. It runs in the stub-markers CI leg. ForwardScopeFlatMemoryTest runs 300 forward passes under a forward scope on the JVM and asserts that direct memory, RSS and the slab stay flat. reference/benchmark-baseline.adoc records the JMH numbers a dispatch or storage change is compared against; scripts/pr-gate.sh points at it. Co-Authored-By: Claude Fable 5.2 --- .github/workflows/build.yml | 7 +- CHANGELOG.md | 4 + .../kotlin/CheckTensorDataDowncastsTask.kt | 140 ++++++++++++++++++ .../src/main/kotlin/DocumentationPlugin.kt | 16 ++ config/tensor-data-downcast-allowlist.txt | 34 +++++ docs/modules/ROOT/nav.adoc | 1 + .../pages/reference/benchmark-baseline.adoc | 77 ++++++++++ scripts/pr-gate.sh | 2 +- .../tensor/ops/ForwardScopeFlatMemoryTest.kt | 98 ++++++++++++ 9 files changed, 375 insertions(+), 4 deletions(-) create mode 100644 build-logic/convention/src/main/kotlin/CheckTensorDataDowncastsTask.kt create mode 100644 config/tensor-data-downcast-allowlist.txt create mode 100644 docs/modules/ROOT/pages/reference/benchmark-baseline.adoc create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 8fec037e..29de43db 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -76,10 +76,11 @@ jobs: # associates with dumps (#1171 merged a stale dump exactly that way; #1174 repaired it). - name: api-compatibility tasks: apiCheck - # Every NotImplementedError / TODO( in production code is either @NotImplemented - # (so it shows in the ops coverage matrix) or listed in config/stub-allowlist.txt. + # Both guards scan sources only, so they share one leg: stubs are either @NotImplemented + # or in config/stub-allowlist.txt, and downcasts to concrete TensorData classes outside + # the data package (SKEEP-003) are in config/tensor-data-downcast-allowlist.txt. - name: stub-markers - tasks: checkUnmarkedStubs + tasks: checkUnmarkedStubs checkTensorDataDowncasts permissions: contents: read diff --git a/CHANGELOG.md b/CHANGELOG.md index 4652f488..25384573 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- **Guards for SKEEP-003's acceptance criteria** (build, `skainet-backend-cpu`, docs, #1501): `./gradlew checkTensorDataDowncasts` counts the downcasts to concrete `TensorData` classes outside the data package (162 in 21 files today) against `config/tensor-data-downcast-allowlist.txt`, so a new one fails the build and the list can only shrink; it runs in the `stub-markers` CI leg. `ForwardScopeFlatMemoryTest` runs 300 forward passes under a forward scope on the JVM and asserts that direct memory, RSS and the slab stay flat. `reference/benchmark-baseline.adoc` records the JMH numbers for matmul, elementwise and reductions that a dispatch or storage change is compared against, replacing the baseline deleted in #1106. + ### Changed - **SKEEP-003 is back at `Accepted`, with an implementation-status table** (docs, #1500): an audit at 0.64.0 found the new storage model in use and the old layer (`TensorStorage`, `BufferHandle`, `TensorData` downcasts in the ops) still in place, so three of the five acceptance criteria are open. The SKEEP now says which, and what tracks each (#932). SKEEP-003a no longer lists #1148 as open, SKEEP-002's status records the Q4_K / Q6_K mapped staging of #1190 and its device measurement, `memory-model.adoc` drops a limit that #1190 removed, and the `TensorStorage` KDoc stops telling new code to target it. diff --git a/build-logic/convention/src/main/kotlin/CheckTensorDataDowncastsTask.kt b/build-logic/convention/src/main/kotlin/CheckTensorDataDowncastsTask.kt new file mode 100644 index 00000000..04179ab2 --- /dev/null +++ b/build-logic/convention/src/main/kotlin/CheckTensorDataDowncastsTask.kt @@ -0,0 +1,140 @@ +import org.gradle.api.DefaultTask +import org.gradle.api.GradleException +import org.gradle.api.file.ConfigurableFileCollection +import org.gradle.api.file.DirectoryProperty +import org.gradle.api.file.RegularFileProperty +import org.gradle.api.tasks.* +import java.io.File + +/** + * Guards a SKEEP-003 acceptance criterion: no dispatch site downcasts to a concrete `TensorData` + * class outside the data package itself. Kernels are meant to be chosen from a view's + * `Format(dtype, encoding)`, not from the Kotlin class that happens to hold the bytes. + * + * A downcast is an `is` / `!is` / `as` / `as?` whose target is a class named `…TensorData` + * (the generic base `TensorData` itself does not count) or one of the storage markers in + * [markerTypes]. Production code only; the packages that define the classes are exempt. + * + * The sites that exist today are listed in [allowlistFile] as `path|count|reason`, where `count` + * is the exact number of downcasts in the file. A new downcast fails the build until the entry is + * raised, which a reviewer sees; removing one fails until the entry is lowered, so the list can + * only describe the code as it is. The goal is an empty list. + */ +@CacheableTask +abstract class CheckTensorDataDowncastsTask : DefaultTask() { + + @get:InputFiles + @get:PathSensitive(PathSensitivity.RELATIVE) + abstract val sources: ConfigurableFileCollection + + /** Repository root that report and allowlist paths are relative to. */ + @get:Internal + abstract val rootDirectory: DirectoryProperty + + @get:InputFile + @get:PathSensitive(PathSensitivity.RELATIVE) + abstract val allowlistFile: RegularFileProperty + + @get:OutputFile + abstract val reportFile: RegularFileProperty + + private val cast = Regex("""(?:!is|\bis|\bas\?|\bas)\s+(?:[A-Za-z_][\w]*\.)*([A-Z]\w*)""") + private val mainSourceDir = Regex("""/src/[A-Za-z0-9]*[Mm]ain/""") + + /** Marker interfaces that stand for a concrete storage class in a dispatch `when`. */ + private val markerTypes = setOf( + "PackedBlockStorage", + "RowDequantSource", + "MemorySegmentBackedData", + "Q4MemorySegmentMarker", + "Q8MemorySegmentMarker", + ) + + /** Packages that define the classes and may inspect them. */ + private val exemptPackages = listOf( + "/sk/ainet/lang/tensor/data/", + "/sk/ainet/lang/tensor/storage/", + "/sk/ainet/lang/memory/", + ) + + @TaskAction + fun check() { + val root = rootDirectory.get().asFile + val allowlist = parseAllowlist(allowlistFile.get().asFile) + val found = mutableMapOf>() + + sources.files + .filter { it.extension == "kt" } + .forEach { file -> + val rel = file.relativeTo(root).invariantSeparatorsPath + if (!mainSourceDir.containsMatchIn("/$rel")) return@forEach + if (exemptPackages.any { it in "/$rel" }) return@forEach + file.readLines().forEachIndexed { index, line -> + val trimmed = line.trimStart() + if (trimmed.startsWith("//") || trimmed.startsWith("*") || trimmed.startsWith("/*")) return@forEachIndexed + val code = stripStrings(line).substringBefore("//") + val hits = cast.findAll(code).count { isConcreteStorageType(it.groupValues[1]) } + repeat(hits) { found.getOrPut(rel) { mutableListOf() } += index + 1 } + } + } + + val problems = mutableListOf() + found.toSortedMap().forEach { (path, hits) -> + val expected = allowlist[path] + when { + expected == null -> hits.forEach { problems += "$path:$it — downcast to a concrete TensorData class" } + expected != hits.size -> problems += "$path — ${hits.size} downcasts, allowlist expects $expected (lines ${hits.joinToString()})" + } + } + allowlist.keys.filter { it !in found }.sorted().forEach { path -> + problems += "$path — allowlisted but has no downcasts any more; remove the entry" + } + + reportFile.get().asFile.apply { parentFile.mkdirs() }.writeText( + found.toSortedMap().entries.joinToString("\n", postfix = "\n") { (path, hits) -> "$path|${hits.size}|${hits.joinToString(",")}" }, + ) + + if (problems.isNotEmpty()) { + throw GradleException( + buildString { + appendLine("Downcasts to concrete TensorData classes outside the data package (SKEEP-003):") + problems.forEach { appendLine(" $it") } + appendLine() + appendLine("Choose the kernel from the view's Format(dtype, encoding) instead of the class. If the") + appendLine("count went down, lower it in ${allowlistFile.get().asFile.relativeTo(root).invariantSeparatorsPath} (path|count|reason).") + }, + ) + } + logger.lifecycle("✅ TensorData downcasts match the allowlist (${found.values.sumOf { it.size }} in ${found.size} files)") + } + + private fun isConcreteStorageType(name: String): Boolean = + name in markerTypes || (name.endsWith("TensorData") && name != "TensorData") + + /** Blank out the contents of `"…"` literals on one line (escape-aware). */ + private fun stripStrings(line: String): String = buildString { + var inString = false + var i = 0 + while (i < line.length) { + val ch = line[i] + when { + inString && ch == '\\' -> i++ + ch == '"' -> { inString = !inString; append(ch) } + !inString -> append(ch) + } + i++ + } + } + + private fun parseAllowlist(file: File): Map = + file.readLines() + .map { it.trim() } + .filter { it.isNotEmpty() && !it.startsWith("#") } + .associate { line -> + val parts = line.split("|").map { it.trim() } + if (parts.size < 3 || parts[2].isEmpty()) { + throw GradleException("Malformed allowlist line (want path|count|reason): $line") + } + parts[0] to (parts[1].toIntOrNull() ?: throw GradleException("Bad count in allowlist line: $line")) + } +} diff --git a/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt b/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt index b2fbd42f..e7bef0eb 100644 --- a/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt +++ b/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt @@ -57,6 +57,22 @@ class DocumentationPlugin : Plugin { } }) + // Guard (SKEEP-003): ops choose kernels from a view's Format, not by downcasting to a + // concrete TensorData class. Today's sites are allowlisted; the list may only shrink. + project.tasks.register("checkTensorDataDowncasts", CheckTensorDataDowncastsTask::class.java, object : Action { + override fun execute(task: CheckTensorDataDowncastsTask) { + task.group = "verification" + task.description = "Fail on downcasts to concrete TensorData classes that are not in the allowlist" + val root = project.rootProject.layout.projectDirectory + project.rootProject.allprojects.forEach { p -> + task.sources.from(p.fileTree("src") { include("**/*.kt") }) + } + task.rootDirectory.set(root) + task.allowlistFile.set(root.file("config/tensor-data-downcast-allowlist.txt")) + task.reportFile.set(project.layout.buildDirectory.file("reports/downcasts/tensor-data-downcasts.txt")) + } + }) + // Register schema validation task in plugin (migrated from skainet-lang-export-ops) val validateTaskProvider = project.tasks.register("validateOperatorSchema", SchemaValidationTask::class.java, object : Action { override fun execute(task: SchemaValidationTask) { diff --git a/config/tensor-data-downcast-allowlist.txt b/config/tensor-data-downcast-allowlist.txt new file mode 100644 index 00000000..cfc51ac5 --- /dev/null +++ b/config/tensor-data-downcast-allowlist.txt @@ -0,0 +1,34 @@ +# Downcasts to concrete TensorData classes outside the data package. Checked by +# `./gradlew checkTensorDataDowncasts`. +# +# SKEEP-003 acceptance criterion: no dispatch site downcasts to a concrete `TensorData` class +# outside the storage/data package itself; a kernel is chosen from the view's +# Format(dtype, encoding). These are the sites that still do. The goal is an empty file. +# +# Format: path|count|issue or reason +# count = exact number of `is` / `!is` / `as` / `as?` whose target is a `…TensorData` class +# or a storage marker (PackedBlockStorage, RowDequantSource, MemorySegmentBackedData, +# Q4/Q8MemorySegmentMarker). Removing a downcast means lowering the count (or deleting the +# line). Do not raise a count: dispatch on the format instead. + +skainet-apps/skainet-grayscale-cli/src/main/kotlin/sk/ainet/apps/grayscale/HloCompiler.kt|1|#1507 +skainet-apps/skainet-grayscale-cli/src/main/kotlin/sk/ainet/apps/grayscale/TensorConversionPipeline.kt|1|#1507 +skainet-apps/skainet-tensor-tools/src/main/kotlin/sk/ainet/apps/tool/quantize/QuantizeCli.kt|3|#1507 +skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt|18|#1507 +skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt|51|#1507 +skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt|48|#1507 +skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmCpuOps.kt|10|#1507 +skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt|3|#1507 +skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt|1|#1507 +skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt|1|#1507 +skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass.kt|1|#1507 +skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt|2|#1507 +skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GGUFWriter.kt|2|#1507 +skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt|1|#1507 +skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/IndexTensor.kt|1|#1507 +skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/QuantizedMatmul.kt|9|#1507 +skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorSpecEncoding.kt|1|#1507 +skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TernaryMatmul.kt|3|#1507 +skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt|1|#1507 +skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/StorageBenchmarks.kt|3|#1507 +skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/TensorAssertions.kt|1|#1507 diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index aabdb2be..7b22053b 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -30,6 +30,7 @@ ** xref:reference/operators/generated/index.adoc[Operator reference] ** xref:reference/ops-status-matrix.adoc[Operator coverage matrix] ** xref:reference/kernel-support-matrix.adoc[Kernel × platform support] +** xref:reference/benchmark-baseline.adoc[Benchmark baseline] ** xref:reference/api.adoc[API reference (Dokka)] * Explanation ** xref:explanation/skainet-for-ai.adoc[SKaiNET for AI/ML] diff --git a/docs/modules/ROOT/pages/reference/benchmark-baseline.adoc b/docs/modules/ROOT/pages/reference/benchmark-baseline.adoc new file mode 100644 index 00000000..c46aba3d --- /dev/null +++ b/docs/modules/ROOT/pages/reference/benchmark-baseline.adoc @@ -0,0 +1,77 @@ += Benchmark baseline +:description: Reference JMH numbers for the matmul, elementwise and reduction hot paths, to compare against when a change touches dispatch or storage. + +These are the numbers a change to kernel dispatch or tensor storage is compared against. +xref:skeep:003-unified-tensor-storage.adoc[SKEEP-003] sets the budget: matmul within noise, +elementwise at most 3 % slower. + +== How to compare + +. Run the suite on the commit before your change and on the commit with it, on the same machine, + with nothing else building: ++ +[source,bash] +---- +./gradlew :skainet-backends:benchmarks:jvm-cpu-jmh:jmh \ + -PjmhIncludes='sk.ainet.bench.(MatmulBench|KernelMatmulBench|QuantizedMatmulBench|ElementwiseAdd1MBench|Reductions1MBench).*' +---- ++ +The results are in `skainet-backends/benchmarks/jvm-cpu-jmh/build/results/jmh/results.txt`. +`scripts/pr-gate.sh --bench` runs the whole suite. +. Put both tables in the pull request. The table below tells you what the numbers looked like on + the reference machine; numbers from another machine are only comparable with each other. +. When a change moves a number on purpose, record a new baseline here in the same pull request. + +Not part of this baseline: `:skainet-lang:skainet-lang-core:jvmBenchmark` (`StorageBenchmarks`, +`TensorBenchmarks`, `TurboQuantBenchmarks`). With the module's default iteration settings a full +run takes hours; run it filtered per class when a change touches the storage hot paths. + +== Baseline of 2026-11-11 + +* Commit: `edadc84b` (`develop`, 0.64.0) +* Machine: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz, 12 threads, Linux 6.8.0-60, OpenJDK 21.0.13 +* Settings: JMH, fork 1, 3 warm-up and 5 measurement iterations, + `--enable-preview --add-modules jdk.incubator.vector` + +[cols="3,3,4,1,2,2,1",options="header"] +|=== +| Benchmark | Class | Parameters | Mode | Score | ± | Unit + +| `add_1M_fp32` | ElementwiseAdd1MBench | vectorEnabled=true | thrpt | 2.061 | 0.023 | ops/ms +| `add_1M_fp32` | ElementwiseAdd1MBench | vectorEnabled=false | thrpt | 1.195 | 0.019 | ops/ms +| `mean_1M_fp32` | Reductions1MBench | vectorEnabled=true | thrpt | 1.027 | 0.004 | ops/ms +| `mean_1M_fp32` | Reductions1MBench | vectorEnabled=false | thrpt | 1.027 | 0.003 | ops/ms +| `sum_1M_fp32` | Reductions1MBench | vectorEnabled=true | thrpt | 1.027 | 0.006 | ops/ms +| `sum_1M_fp32` | Reductions1MBench | vectorEnabled=false | thrpt | 1.030 | 0.005 | ops/ms +| `matmul_fp32_square` | KernelMatmulBench | provider=scalar, size=256 | avgt | 22.468 | 0.182 | ms/op +| `matmul_fp32_square` | KernelMatmulBench | provider=scalar, size=512 | avgt | 292.669 | 0.428 | ms/op +| `matmul_fp32_square` | KernelMatmulBench | provider=scalar, size=1024 | avgt | 1869.206 | 13.642 | ms/op +| `matmul_fp32_square` | KernelMatmulBench | provider=panama, size=256 | avgt | 1.257 | 0.014 | ms/op +| `matmul_fp32_square` | KernelMatmulBench | provider=panama, size=512 | avgt | 9.881 | 0.144 | ms/op +| `matmul_fp32_square` | KernelMatmulBench | provider=panama, size=1024 | avgt | 87.219 | 1.669 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=256, vectorEnabled=true | avgt | 8.548 | 0.072 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=256, vectorEnabled=false | avgt | 7.187 | 0.129 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=512, vectorEnabled=true | avgt | 67.587 | 0.494 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=512, vectorEnabled=false | avgt | 62.912 | 0.486 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=1024, vectorEnabled=true | avgt | 544.217 | 6.110 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=true, size=1024, vectorEnabled=false | avgt | 480.891 | 2.660 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=256, vectorEnabled=true | avgt | 1.269 | 0.015 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=256, vectorEnabled=false | avgt | 7.080 | 0.047 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=512, vectorEnabled=true | avgt | 10.382 | 0.181 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=512, vectorEnabled=false | avgt | 59.932 | 1.181 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=1024, vectorEnabled=true | avgt | 87.704 | 1.994 | ms/op +| `matmul_fp32_square` | MatmulBench | blasEnabled=false, size=1024, vectorEnabled=false | avgt | 480.785 | 6.122 | ms/op +| `matmul_q4k_panama` | QuantizedMatmulBench | shape=1024-1024 | avgt | 0.110 | 0.001 | ms/op +| `matmul_q4k_panama` | QuantizedMatmulBench | shape=4096-1024 | avgt | 0.365 | 0.002 | ms/op +| `matmul_q4k_panama` | QuantizedMatmulBench | shape=4096-4096 | avgt | 1.388 | 0.020 | ms/op +|=== + +`thrpt` is throughput, higher is better; `avgt` is the average time per call, lower is better. + +One thing the table shows that is not a storage matter: on this Linux machine `MatmulBench` with +`blasEnabled=true` is about seven times slower than the Panama kernel at every size +(8.5 ms against 1.3 ms at 256). BLAS routing is off by default outside macOS. + +The baseline recorded on 2026-08-22, before the storage migration, was removed with +`docs/design/` (#1106); it is in the history at commit `3c0ee71d`. It was taken on the same +machine with OpenJDK 25, so the two are not directly comparable. diff --git a/scripts/pr-gate.sh b/scripts/pr-gate.sh index 1ddd5d0b..2dc8463e 100755 --- a/scripts/pr-gate.sh +++ b/scripts/pr-gate.sh @@ -68,7 +68,7 @@ step "Android host tests" "${GRADLE[@]}" testAndroidHostTest if [[ "$mode" == "--bench" ]]; then - step "benchmarks (compare against the committed baseline before/after)" + step "benchmarks (compare with docs/modules/ROOT/pages/reference/benchmark-baseline.adoc)" "${GRADLE[@]}" :skainet-lang:skainet-lang-core:jvmBenchmark "${GRADLE[@]}" :skainet-backends:benchmarks:jvm-cpu-jmh:jmh fi diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt new file mode 100644 index 00000000..e5d6a704 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt @@ -0,0 +1,98 @@ +package sk.ainet.exec.tensor.ops + +import java.lang.management.BufferPoolMXBean +import java.lang.management.ManagementFactory +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.forwardScope +import sk.ainet.lang.memory.MemoryProbe +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 + +/** + * SKEEP-003 acceptance criterion (#1501): repeated forward passes under a forward scope keep + * memory flat on the JVM. This is the failure that forced the retreat to the garbage collector: + * a shared arena grew with every op output, and per-call arenas leaked per matmul. + * + * The loop is a decode step in miniature — weights outside the scope, activations inside it, + * `reset()` at the step boundary — at a size where a leak is unmistakable: one step produces + * about 3 MB of activations, so 300 steps would be about 900 MB if anything accumulated. + */ +class ForwardScopeFlatMemoryTest { + + private val n = 512 + private val steps = 300 + private val warmup = 20 + + private val directPool: BufferPoolMXBean = + ManagementFactory.getPlatformMXBeans(BufferPoolMXBean::class.java).first { it.name == "direct" } + + private fun mb(bytes: Long): String = "%.1f MB".format(bytes / (1024.0 * 1024.0)) + + private fun step(x: Tensor, w: Tensor, b: Tensor): Tensor { + val ops = x.ops + return ops.relu(ops.add(ops.matmul(x, w), b)) + } + + private fun weights(ctx: ExecutionContext): Pair, Tensor> { + val w = ctx.fromFloatArray(Shape(n, n), FP32::class, FloatArray(n * n) { ((it * 31) % 17 - 8) * 0.01f }) + val b = ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { (it % 5 - 2) * 0.1f }) + return w to b + } + + @Test + fun repeatedForwardPassesKeepDirectMemoryAndTheSlabFlat() { + val base = DirectCpuExecutionContext() + val (w, b) = weights(base) + val input = FloatArray(n * n) { ((it * 13) % 11 - 5) * 0.05f } + val stepBytes = 3L * n * n * Float.SIZE_BYTES // input, matmul and add outputs; relu reuses or adds one more + + base.forwardScope(slabFloats = 8 * n * n) { scoped, scope -> + var directAfterWarmup = 0L + var rssAfterWarmup: Long? = null + var peakAfterWarmup = 0 + var directPeak = 0L + + repeat(steps) { stepNo -> + val x = scoped.fromFloatArray(Shape(n, n), FP32::class, input) + val y = step(x, w, b) + assertTrue(y.data is StorageFloatTensorData<*>, "step $stepNo: activations must come from the slab") + if (stepNo == 0) assertTrue(y.data.copyToFloatArray().any { it > 0f }, "the step must compute something") + assertEquals(0L, scope.overflowBytes, "step $stepNo: nothing may spill past the slab") + scope.reset() + + val direct = directPool.memoryUsed + if (stepNo == warmup) { + directAfterWarmup = direct + rssAfterWarmup = MemoryProbe.rssBytes() + peakAfterWarmup = scope.peakFloats + } + if (stepNo > warmup) { + directPeak = maxOf(directPeak, direct) + assertEquals(peakAfterWarmup, scope.peakFloats, "step $stepNo: the slab must not grow in steady state") + } + } + + val directGrowth = directPeak - directAfterWarmup + val rssGrowth = rssAfterWarmup?.let { before -> MemoryProbe.rssBytes()?.minus(before) } + println( + "[forward-scope] ${steps - warmup} steady steps of ${mb(stepBytes)}: direct pool ${mb(directAfterWarmup)} -> peak ${mb(directPeak)}, " + + "rss growth ${rssGrowth?.let(::mb) ?: "n/a"}, slab peak ${scope.peakFloats} floats", + ) + + // Flat: less than one step's worth of direct memory appears over 280 steps. + assertTrue(directGrowth < stepBytes, "direct memory grew by ${mb(directGrowth)} over ${steps - warmup} steps") + // A leak of every activation would be about 840 MB here; allow a tenth of that for the + // JIT, the collector and allocator noise in a shared test JVM. + if (rssGrowth != null) { + val leakIfEverythingStayed = stepBytes * (steps - warmup) + assertTrue(rssGrowth < leakIfEverythingStayed / 10, "RSS grew by ${mb(rssGrowth)} over ${steps - warmup} steps") + } + } + } +} From 9c096ce98cc1d9c237fbda5b592c4be5e88cb8d1 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 11 Nov 2026 15:31:05 +0100 Subject: [PATCH 3/3] test(backend-cpu): assert live heap, not RSS, in ForwardScopeFlatMemoryTest (#1501) RSS grew past the bound on a CI runner although nothing leaked: a JVM commits heap lazily, so RSS follows garbage that is already dead. The test now asserts the live heap after a collection and prints RSS. Co-Authored-By: Claude Fable 5.2 --- CHANGELOG.md | 2 +- .../tensor/ops/ForwardScopeFlatMemoryTest.kt | 25 +++++++++++++------ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 25384573..eb4a5b36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -- **Guards for SKEEP-003's acceptance criteria** (build, `skainet-backend-cpu`, docs, #1501): `./gradlew checkTensorDataDowncasts` counts the downcasts to concrete `TensorData` classes outside the data package (162 in 21 files today) against `config/tensor-data-downcast-allowlist.txt`, so a new one fails the build and the list can only shrink; it runs in the `stub-markers` CI leg. `ForwardScopeFlatMemoryTest` runs 300 forward passes under a forward scope on the JVM and asserts that direct memory, RSS and the slab stay flat. `reference/benchmark-baseline.adoc` records the JMH numbers for matmul, elementwise and reductions that a dispatch or storage change is compared against, replacing the baseline deleted in #1106. +- **Guards for SKEEP-003's acceptance criteria** (build, `skainet-backend-cpu`, docs, #1501): `./gradlew checkTensorDataDowncasts` counts the downcasts to concrete `TensorData` classes outside the data package (162 in 21 files today) against `config/tensor-data-downcast-allowlist.txt`, so a new one fails the build and the list can only shrink; it runs in the `stub-markers` CI leg. `ForwardScopeFlatMemoryTest` runs 300 forward passes under a forward scope on the JVM and asserts that direct memory, the live heap and the slab stay flat. `reference/benchmark-baseline.adoc` records the JMH numbers for matmul, elementwise and reductions that a dispatch or storage change is compared against, replacing the baseline deleted in #1106. ### Changed diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt index e5d6a704..3f8c5a55 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/ForwardScopeFlatMemoryTest.kt @@ -22,6 +22,7 @@ import sk.ainet.lang.types.FP32 * The loop is a decode step in miniature — weights outside the scope, activations inside it, * `reset()` at the step boundary — at a size where a leak is unmistakable: one step produces * about 3 MB of activations, so 300 steps would be about 900 MB if anything accumulated. + * Asserted: the direct buffer pool, the live heap after a collection, and the slab's peak. */ class ForwardScopeFlatMemoryTest { @@ -32,6 +33,13 @@ class ForwardScopeFlatMemoryTest { private val directPool: BufferPoolMXBean = ManagementFactory.getPlatformMXBeans(BufferPoolMXBean::class.java).first { it.name == "direct" } + /** Live heap after a collection: what the loop still holds, not what it has churned through. */ + private fun liveHeap(): Long { + val rt = Runtime.getRuntime() + repeat(3) { System.gc(); Thread.sleep(50) } + return rt.totalMemory() - rt.freeMemory() + } + private fun mb(bytes: Long): String = "%.1f MB".format(bytes / (1024.0 * 1024.0)) private fun step(x: Tensor, w: Tensor, b: Tensor): Tensor { @@ -55,6 +63,7 @@ class ForwardScopeFlatMemoryTest { base.forwardScope(slabFloats = 8 * n * n) { scoped, scope -> var directAfterWarmup = 0L var rssAfterWarmup: Long? = null + var heapAfterWarmup = 0L var peakAfterWarmup = 0 var directPeak = 0L @@ -70,6 +79,7 @@ class ForwardScopeFlatMemoryTest { if (stepNo == warmup) { directAfterWarmup = direct rssAfterWarmup = MemoryProbe.rssBytes() + heapAfterWarmup = liveHeap() peakAfterWarmup = scope.peakFloats } if (stepNo > warmup) { @@ -79,20 +89,21 @@ class ForwardScopeFlatMemoryTest { } val directGrowth = directPeak - directAfterWarmup + val heapGrowth = liveHeap() - heapAfterWarmup val rssGrowth = rssAfterWarmup?.let { before -> MemoryProbe.rssBytes()?.minus(before) } println( "[forward-scope] ${steps - warmup} steady steps of ${mb(stepBytes)}: direct pool ${mb(directAfterWarmup)} -> peak ${mb(directPeak)}, " + - "rss growth ${rssGrowth?.let(::mb) ?: "n/a"}, slab peak ${scope.peakFloats} floats", + "live heap growth ${mb(heapGrowth)}, rss growth ${rssGrowth?.let(::mb) ?: "n/a"}, slab peak ${scope.peakFloats} floats", ) // Flat: less than one step's worth of direct memory appears over 280 steps. assertTrue(directGrowth < stepBytes, "direct memory grew by ${mb(directGrowth)} over ${steps - warmup} steps") - // A leak of every activation would be about 840 MB here; allow a tenth of that for the - // JIT, the collector and allocator noise in a shared test JVM. - if (rssGrowth != null) { - val leakIfEverythingStayed = stepBytes * (steps - warmup) - assertTrue(rssGrowth < leakIfEverythingStayed / 10, "RSS grew by ${mb(rssGrowth)} over ${steps - warmup} steps") - } + // A leak of every activation would be about 840 MB here. Live heap after a collection + // is what the loop still holds; allow a tenth of the leak for whatever else the test + // JVM keeps. RSS is printed, not asserted: a JVM grows its committed heap lazily, so + // RSS moves with garbage that is already dead (seen on a CI runner). + val leakIfEverythingStayed = stepBytes * (steps - warmup) + assertTrue(heapGrowth < leakIfEverythingStayed / 10, "live heap grew by ${mb(heapGrowth)} over ${steps - warmup} steps") } } }