diff --git a/CHANGELOG.md b/CHANGELOG.md index adc562c5e..63e4a7eef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- **BF16 and FP16 matmul on the wasm tier**: `WasmKernelProvider` now binds the sidecar's `skainet_bf16_matmul` / `skainet_fp16_matmul` exports (`WasmBf16MatmulKernel`, `WasmFp16MatmulKernel`), so 16-bit dense weights no longer fall to the scalar kernels in a browser; a contiguous weight stays resident in the sidecar after its first call. +- **`WasmSimd.features()` / `WasmKernels.features()`**: mode, SIMD128, threads, cross-origin isolation, logical cores and resident weight bytes — the wasm counterpart of the JVM provider banner. `WasmSimd.install()` prints one `[SKaiNET] Using SIMD128 sidecar kernels (…)` line like the JVM factory does. +- **`KernelStrictness` is common code** (was JVM-only): `KernelStrictness.overrideEnabled` is the programmatic switch on every target; the JVM and Android still read `-Dskainet.strict.kernels`. `KernelDispatch` fails fast before the decoding reference fallback and `DefaultCpuOps.matmul` before the generic per-element loop, so strict mode means the same thing on every tier. +- **Kernel support matrix**: a `Float16` row, BF16/FP16 in the wasm column, Q6_K corrected to `native-ffm` on the JVM (the provider has had the kernel; the declared tier had not), and a third table, *Resident serving*, naming the sidecar-resident weights that stand in for mapped serving on JS/WASM. `KernelSupportMatrixTest` now gates the declared Panama and FFM tiers against what the providers expose. +- **`wasm-bench` runs the `engine-*` scenarios** of the JVM publish harness (`engine-fp32-gemm`, `engine-kernel-matmul`, `engine-bf16-matmul`, `engine-q8-matmul`, `engine-q4-gemm`, `engine-ternary-f32-gemv`, `engine-elementwise-add`, `engine-reductions-sum/mean`) with identical shapes and metrics on every tier it measures, so the JVM and wasm numbers sit in one table. + ## [0.61.0] - 2026-10-16 Headline: **the browser lane (SKEEP-007): packed SIMD128 kernels, quantize-at-load, and no more diff --git a/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt b/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt index bb422504f..6964d6e02 100644 --- a/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt +++ b/build-logic/convention/src/main/kotlin/GenerateKernelMatrixTask.kt @@ -49,8 +49,10 @@ abstract class GenerateKernelMatrixTask : DefaultTask() { appendLine( "Each cell is the best (highest-priority) provider that serves " + "`${module.inputDtype} × format` `${module.op}` on that platform: " + - "*native-ffm* (100) → *panama-vector* (50) → *scalar* (0). An empty cell " + - "(`—`) means no provider carries a kernel there (the format is dequant-to-FP32 only).", + "*native-ffm* / *native-jni* / *native-cinterop* (100) → *panama-vector* / " + + "*wasm-simd128* (50) → *scalar* (0). An empty cell " + + "(`—`) means no provider carries a kernel there (the format is dequant-to-FP32 only). " + + "The JS/WASM column assumes `WasmSimd.install()` has run; without the sidecar it is the scalar floor.", ) appendLine("") if (module.formats.isEmpty() || module.platforms.isEmpty()) { @@ -100,6 +102,32 @@ abstract class GenerateKernelMatrixTask : DefaultTask() { appendLine("|===") appendLine("") } + if (module.resident.isNotEmpty()) { + appendLine("== Resident serving (sidecar memory)") + appendLine("") + appendLine( + "The browser has no files to map, so the wasm tier keeps weights *resident* in the " + + "SIMD128 sidecar's linear memory instead: each weight is uploaded once on its first " + + "matmul and every later call reads it in place (`WasmKernels.weightHandle`). This is " + + "the JS/WASM counterpart of mapped serving above — no per-call copy of the weight, " + + "no prepack — bounded by the wasm32 4 GB linear memory rather than the page cache. " + + "The other platforms serve from the heap, off-heap or the mapping and have no cell here.", + ) + appendLine("") + appendLine("[cols=\"$colSpec\", options=\"header\"]") + appendLine("|===") + append("| Weight format ") + module.platforms.forEach { append("| $it ") } + appendLine("") + appendLine("") + module.resident.forEach { fmt -> + append("| `${fmt.name}` ") + module.platforms.forEach { p -> append("| ${fmt.byPlatform[p] ?: "—"} ") } + appendLine("") + } + appendLine("|===") + appendLine("") + } appendLine( "See also the eager backends & kernels mindmap " + "(xref:explanation/eager-execution.adoc[]) for the narrative overview and gaps.", diff --git a/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt b/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt index 1084af35e..90da53887 100644 --- a/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt +++ b/build-logic/convention/src/main/kotlin/models/KernelSupportModels.kt @@ -22,6 +22,7 @@ data class KernelSupportModule( * prepack. Absent/empty in pre-#1189 JSON, so old files stay decodable. */ val mapped: List = emptyList(), + val resident: List = emptyList(), ) @Serializable diff --git a/docs/modules/ROOT/pages/contributing/dtype-model.adoc b/docs/modules/ROOT/pages/contributing/dtype-model.adoc index 43e57e521..8d3e4d13f 100644 --- a/docs/modules/ROOT/pages/contributing/dtype-model.adoc +++ b/docs/modules/ROOT/pages/contributing/dtype-model.adoc @@ -105,11 +105,18 @@ affordance today (W4, shipped): java -Dskainet.strict.kernels=true … ---- +[source,kotlin] +---- +KernelStrictness.overrideEnabled = true // every target; the only switch on JS, Wasm and Kotlin/Native +---- + When set, `DefaultCpuOpsJvm.matmul` raises `NoSuchKernelException` (with the failing dtype pair and the list of currently-registered providers) just before its silent scalar -fallback would have run. Default off — adaptive behaviour is -preserved. +fallback would have run; on every tier `KernelDispatch` raises it +before the decoding reference kernel and `DefaultCpuOps.matmul` +before the generic per-element loop. Default off — adaptive +behaviour is preserved. The constraint-resolution pass (W7) raises the same exception shape at *graph-prep* time, before forward execution can even diff --git a/docs/modules/ROOT/pages/explanation/eager-execution.adoc b/docs/modules/ROOT/pages/explanation/eager-execution.adoc index c2aa8d98f..ff94753b2 100644 --- a/docs/modules/ROOT/pages/explanation/eager-execution.adoc +++ b/docs/modules/ROOT/pages/explanation/eager-execution.adoc @@ -52,28 +52,32 @@ mindmap == Kernel × provider (matmul, FP32 activations) -[cols="2,1,1,1,1,1",options="header"] +[cols="2,1,1,1,1,1,1",options="header"] |=== -| Weight format | Scalar (all targets) | Panama (JVM) | Native FFM (JVM) | Native JNI (Android) | Native cinterop (K/N) -| FP32 | ✅ | ✅ | ✅ | ❌ | ❌ -| BF16 | ✅ | ✅ | ✅ | ❌ | ❌ -| Q8_0 | ✅ | ✅ | ✅ | ✅ | ✅ -| Q4_0 | ✅ | ✅ | ✅ | ✅ | ✅ -| Q4_K | ✅ | ✅ | ✅ | ✅ | ✅ -| Q6_K | ✅ | ✅ | ❌ | ✅ | ✅ -| Q5_K | ✅ | ✅ | ✅ | ✅ | ✅ -| Q5_1 | ✅ | ✅ | ✅ | ✅ | ✅ -| Q5_0 | ✅ | ✅ | ✅ | ✅ | ✅ -| TQ2_0 / BitNet b1.58 (int8 activations) | ✅ | — | — | ✅ NEON | — -| BitNet b1.58 (exact FP32 activations, vendored NeoGPU kernel) | ✅ | ✅ FFM | — | ✅ JNI + fused lm_head | ✅ K/N cinterop -| BITNET_PLANES (multi-plane ternary packing) | ✅ | via kernel pack | — | via kernel pack | via kernel pack -| Q2_K / Q3_K / Q8_K / IQ4 | ❌ (dequant to FP32 only) | ❌ | ❌ | ❌ | ❌ +| Weight format | Scalar (all targets) | Panama (JVM) | Native FFM (JVM) | Native JNI (Android) | Native cinterop (K/N) | Wasm SIMD128 (wasmJs) +| FP32 | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ +| BF16 | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ +| FP16 | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ +| Q8_0 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q4_0 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q4_K | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q6_K | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q5_K | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q5_1 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| Q5_0 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ +| TQ2_0 / BitNet b1.58 (int8 activations) | ✅ | — | — | ✅ NEON | — | — +| BitNet b1.58 (exact FP32 activations, vendored NeoGPU kernel) | ✅ | ✅ FFM | — | ✅ JNI + fused lm_head | ✅ K/N cinterop | ❌ (reference) +| BITNET_PLANES (multi-plane ternary packing) | ✅ | via kernel pack | — | via kernel pack | via kernel pack | ❌ (reference) +| Q2_K / Q3_K / Q8_K / IQ4 | ❌ (dequant to FP32 only) | ❌ | ❌ | ❌ | ❌ | ❌ |=== Resolution is by priority: *Native (100 — whichever of FFM/JNI/cinterop applies to the target) → -Panama (50, JVM only) → Scalar (0)*. The best *available* provider that carries the kernel wins; +Panama / Wasm SIMD128 (50) → Scalar (0)*. The best *available* provider that carries the kernel wins; otherwise it cascades down. At most one native tier is compiled into a given target, so the native -columns are mutually exclusive per platform rather than stacked. +columns are mutually exclusive per platform rather than stacked. The Wasm SIMD128 column is the +sidecar pack (`skainet-backend-wasm-simd`), present only after `WasmSimd.install()`; the explicit +SIMD128 bodies are Q8_0, Q4_0 and dense `x · Wᵀ`, the other formats run the portable C body +auto-vectorised by Emscripten. Packed weights reach a kernel through the registry by declaring their xref:explanation/packed-weight-layout.adoc[block order]; a canonical weight is relayouted once, @@ -96,7 +100,11 @@ Packed-quant matmul works on every target: those formats have a `commonMain` sca == Gaps -* *Native FFM Q6_K* — the only packed format the FFM C kernel set does not cover. +* *Wasm ternary / BitNet* — the ternary gemv and planes lm-head kernels are not in the sidecar +build yet, so BitNet models in a browser run the decoding reference kernel. +* *Wasm beyond matmul* — elementwise, activations, reductions, normalisation and attention are the +common scalar Kotlin paths on wasmJs; Kotlin/Wasm has no SIMD types and no shared-memory threads, so +these follow the matmuls into the sidecar (SKEEP-007 rollout step 3). * *Native JNI dense FP32/BF16* — the Android JNI provider has no GEMM shim, so dense ops fall through to scalar there whichever `.so` tier loaded (SKaiNET#920). * *Dense FP32/BF16 SIMD on Kotlin/Native linux* — still the scalar floor; Apple has Accelerate diff --git a/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc b/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc index f655ab587..a3317ed37 100644 --- a/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/kernel-support-matrix.adoc @@ -3,18 +3,19 @@ Generated from `kernel-support.json` (version `0.61.0`) by `KernelSupportMatrixTest` — registry introspection of the registered `KernelProvider` implementations. Do not edit by hand; run `./gradlew generateKernelMatrix` to refresh. -Each cell is the best (highest-priority) provider that serves `Float32 × format` `matmul` on that platform: *native-ffm* (100) → *panama-vector* (50) → *scalar* (0). An empty cell (`—`) means no provider carries a kernel there (the format is dequant-to-FP32 only). +Each cell is the best (highest-priority) provider that serves `Float32 × format` `matmul` on that platform: *native-ffm* / *native-jni* / *native-cinterop* (100) → *panama-vector* / *wasm-simd128* (50) → *scalar* (0). An empty cell (`—`) means no provider carries a kernel there (the format is dequant-to-FP32 only). The JS/WASM column assumes `WasmSimd.install()` has run; without the sidecar it is the scalar floor. [cols="1,1,1,1,1,1", options="header"] |=== | Weight format | JVM | Android | Native·linux | Native·apple | JS/WASM | `Float32` | native-ffm | scalar | scalar | scalar | wasm-simd128 -| `BFloat16` | native-ffm | scalar | scalar | scalar | scalar +| `BFloat16` | native-ffm | scalar | scalar | scalar | wasm-simd128 +| `Float16` | native-ffm | scalar | scalar | scalar | wasm-simd128 | `Q8_0` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 | `Q4_0` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 | `Q4_K` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 -| `Q6_K` | panama-vector | native-jni | native-cinterop | native-cinterop | wasm-simd128 +| `Q6_K` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 | `Q5_K` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 | `Q5_1` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 | `Q5_0` | native-ffm | native-jni | native-cinterop | native-cinterop | wasm-simd128 @@ -37,4 +38,24 @@ Kernels that read the weight in canonical row-major GGUF file order straight fro | `Q5_0` | ffm-rowmajor | native-jni-direct | — | — | — |=== +== Resident serving (sidecar memory) + +The browser has no files to map, so the wasm tier keeps weights *resident* in the SIMD128 sidecar's linear memory instead: each weight is uploaded once on its first matmul and every later call reads it in place (`WasmKernels.weightHandle`). This is the JS/WASM counterpart of mapped serving above — no per-call copy of the weight, no prepack — bounded by the wasm32 4 GB linear memory rather than the page cache. The other platforms serve from the heap, off-heap or the mapping and have no cell here. + +[cols="1,1,1,1,1,1", options="header"] +|=== +| Weight format | JVM | Android | Native·linux | Native·apple | JS/WASM + +| `Float32` | — | — | — | — | wasm-resident +| `BFloat16` | — | — | — | — | wasm-resident +| `Float16` | — | — | — | — | wasm-resident +| `Q8_0` | — | — | — | — | wasm-resident +| `Q4_0` | — | — | — | — | wasm-resident +| `Q4_K` | — | — | — | — | wasm-resident +| `Q6_K` | — | — | — | — | wasm-resident +| `Q5_K` | — | — | — | — | wasm-resident +| `Q5_1` | — | — | — | — | wasm-resident +| `Q5_0` | — | — | — | — | wasm-resident +|=== + See also the eager backends & kernels mindmap (xref:explanation/eager-execution.adoc[]) for the narrative overview and gaps. diff --git a/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt index 5a61c3071..72c130075 100644 --- a/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt +++ b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Main.kt @@ -21,7 +21,9 @@ import kotlin.time.TimeSource * * `[SKAINET_PERF_WASM] m= n= k= form= ms= gflops=` * - * Every later tier (SIMD128 sidecar, quantize-at-load) is measured against these lines. + * Every later tier (SIMD128 sidecar, quantize-at-load) is measured against these lines. Each pass + * also runs the `engine-*` scenarios of the JVM publish harness ([runEngineScenarios]) so the two + * tiers are compared on identical shapes. */ private data class Case(val name: String, val m: Int, val n: Int, val k: Int) @@ -107,7 +109,8 @@ suspend fun main() { private fun runPass(tier: String) { val ctx = DirectCpuExecutionContext() - println("[SKAINET_PERF_WASM] tier=$tier backend=${ctx.ops::class.simpleName} cases=${cases.size}") + println("[SKAINET_PERF_WASM] tier=$tier backend=${ctx.ops::class.simpleName} cases=${cases.size} features=${WasmKernels.features()}") + runEngineScenarios(tier) for (case in cases) { val x = ctx.fromFloatArray(Shape(case.m, case.k), FP32::class, pseudo(case.m * case.k, 3)) val dense = ctx.fromFloatArray(Shape(case.n, case.k), FP32::class, pseudo(case.n * case.k, 7)) diff --git a/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Scenarios.kt b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Scenarios.kt new file mode 100644 index 000000000..271fc27be --- /dev/null +++ b/skainet-backends/benchmarks/wasm-bench/src/wasmJsMain/kotlin/sk/ainet/bench/wasm/Scenarios.kt @@ -0,0 +1,163 @@ +package sk.ainet.bench.wasm + +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.KernelKey +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.BitNetB158TensorData +import sk.ainet.lang.types.FP32 +import kotlin.math.roundToInt +import kotlin.time.TimeSource + +/** + * The `engine-*` scenarios of `skainet-backends/benchmarks/jvm-cpu-publish`, run on the wasm tier + * with the same shapes, data generators and metrics, so one table compares the two tiers: + * + * `[SKAINET_PERF_WASM] tier= scenario= metric= value= ` + * + * Kernel scenarios take the best registered provider (`scalar` before `WasmSimd.install()`, + * `wasm-simd128` after); ops scenarios go through a fresh `DirectCpuExecutionContext`, exactly as + * the JVM harness does. A scenario the tier cannot serve yet prints `value=skipped reason=…`. + */ +internal class WasmScenario( + val id: String, + val metric: String, + val params: Map, + val reps: Int, + val setup: () -> (() -> Double), +) + +private fun lcg(seed: Int): () -> Float { + var s = seed + return { + s = s * 1103515245 + 12345 + ((s ushr 16) and 0x7FFF) / 32768f + } +} + +internal fun engineScenarios(): List = listOf( + WasmScenario("engine-fp32-gemm", "gflops", mapOf("size" to "1024"), reps = 3) { + val n = 1024 + val ctx = DirectCpuExecutionContext() + val a = ctx.fromFloatArray(Shape(n, n), FP32::class, FloatArray(n * n) { ((it % 251) - 125).toFloat() / 127f }) + val b = ctx.fromFloatArray(Shape(n, n), FP32::class, FloatArray(n * n) { ((it * 13 % 257) - 128).toFloat() / 127f }) + val flops = 2.0 * n * n * n; + { flops / seconds { ctx.ops.matmul(a, b) } / 1e9 } + }, + WasmScenario("engine-kernel-matmul", "gflops", mapOf("size" to "1024"), reps = 3) { + val n = 1024 + DirectCpuExecutionContext() + val kernel = checkNotNull(KernelRegistry.bestAvailable()?.matmulFp32()) { "no FP32 kernel" } + val a = FloatArray(n * n) { ((it % 251) - 125).toFloat() / 127f } + val b = FloatArray(n * n) { ((it * 13 % 257) - 128).toFloat() / 127f } + val out = FloatArray(n * n) + val flops = 2.0 * n * n * n; + { flops / seconds { kernel.matmul(a, 0, n, b, 0, n, out, 0, n, n, n, n) } / 1e9 } + }, + WasmScenario("engine-bf16-matmul", "gflops", mapOf("m" to "1024", "n" to "1024", "k" to "1024"), reps = 3) { + val n = 1024 + DirectCpuExecutionContext() + val kernel = checkNotNull(KernelRegistry.bestAvailable()?.matmulBf16()) { "no BF16 kernel" } + val rnd = lcg(n) + val a = FloatArray(n * n) { rnd() - 0.5f } + val b = ByteArray(n * n * 2) + for (i in 0 until n * n) { + val bits = (rnd() - 0.5f).toRawBits() ushr 16 + b[i * 2] = (bits and 0xFF).toByte() + b[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val out = FloatArray(n * n) + val flops = 2.0 * n * n * n; + { flops / seconds { kernel.matmul(a, 0, n, b, 0, n * 2, out, 0, n, n, n, n) } / 1e9 } + }, + WasmScenario("engine-q8-matmul", "gops", mapOf("input_dim" to "4096", "output_dim" to "4096"), reps = 10) { + val inputDim = 4096; val outputDim = 4096 + DirectCpuExecutionContext() + val kernel = checkNotNull(KernelRegistry.bestAvailable()?.matmulQ8_0()) { "no Q8_0 kernel" } + val numBlocks = inputDim / 32 * outputDim + val rnd = lcg(inputDim + outputDim) + val weight = ByteArray(numBlocks * 34) { (rnd() * 255f).toInt().toByte() } + for (block in 0 until numBlocks) { weight[block * 34] = 0x00; weight[block * 34 + 1] = 0x22 } + val input = FloatArray(inputDim) { rnd() - 0.5f } + val output = FloatArray(outputDim) + val ops = 2.0 * inputDim * outputDim; + { ops / seconds { kernel.matmul(input, 0, weight, 0, inputDim, outputDim, output, 0) } / 1e9 } + }, + WasmScenario("engine-q4-gemm", "gops", mapOf("input_dim" to "4096", "output_dim" to "4096"), reps = 10) { + val inputDim = 4096; val outputDim = 4096 + DirectCpuExecutionContext() + val kernel = checkNotNull(KernelRegistry.bestAvailable()?.matmulQ4K()) { "no Q4_K kernel" } + val numBlocks = inputDim / 256 * outputDim + val rnd = lcg(42) + val weight = ByteArray(numBlocks * 144) { (rnd() * 255f).toInt().toByte() } + for (block in 0 until numBlocks) { + val base = block * 144 + weight[base] = 0x00; weight[base + 1] = 0x3C; weight[base + 2] = 0x00; weight[base + 3] = 0x3C + } + val input = FloatArray(inputDim) { ((it % 251) - 125).toFloat() / 127f } + val output = FloatArray(outputDim) + val ops = 2.0 * inputDim * outputDim; + { ops / seconds { kernel.matmul(input, 0, weight, 0, inputDim, outputDim, output, 0) } / 1e9 } + }, + WasmScenario("engine-ternary-f32-gemv", "gops", mapOf("input_dim" to "2560", "output_dim" to "6912"), reps = 5) { + val inputDim = 2560; val outputDim = 6912 + val rnd = lcg(inputDim + outputDim) + val values = FloatArray(outputDim * inputDim) { ((rnd() * 3f).toInt() - 1) * 0.5f } + val weight = BitNetB158TensorData.fromFloats(Shape(outputDim, inputDim), values).packedView + val activation = TensorView.dense(Storage.Heap.wrap(FloatArray(inputDim) { rnd() - 0.5f }), Shape(1, inputDim), FP32) + // The decoding reference kernel is orders of magnitude slower than a real kernel and + // takes hours at this shape on wasm; measure only when a ternary kernel is registered. + checkNotNull(KernelDispatch.find(KernelKey.matmul(activation, weight))) { "no ternary kernel registered on this tier" } + val out = TensorView.dense(Storage.Heap.floats(outputDim), Shape(1, outputDim), FP32) + val ops = 2.0 * inputDim * outputDim; + { ops / seconds { KernelDispatch.matmul(activation, weight, out, Scope.Ambient) } / 1e9 } + }, + WasmScenario("engine-elementwise-add", "melems_per_sec", mapOf("elements" to "1000000"), reps = 10) { + val n = 1_000_000 + val ctx = DirectCpuExecutionContext() + val a = ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { it.toFloat() * 0.5f }) + val b = ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { 1f }); + { n / seconds { ctx.ops.add(a, b) } / 1e6 } + }, + WasmScenario("engine-reductions-sum", "melems_per_sec", mapOf("elements" to "1000000", "op" to "sum"), reps = 10) { + val n = 1_000_000 + val ctx = DirectCpuExecutionContext() + val a = ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { (it % 1024).toFloat() * 0.25f }); + { n / seconds { ctx.ops.sum(a) } / 1e6 } + }, + WasmScenario("engine-reductions-mean", "melems_per_sec", mapOf("elements" to "1000000", "op" to "mean"), reps = 10) { + val n = 1_000_000 + val ctx = DirectCpuExecutionContext() + val a = ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { (it % 1024).toFloat() * 0.25f }); + { n / seconds { ctx.ops.mean(a) } / 1e6 } + }, +) + +private inline fun seconds(block: () -> Unit): Double { + val t0 = TimeSource.Monotonic.markNow() + block() + return t0.elapsedNow().inWholeNanoseconds / 1e9 +} + +/** Run every scenario once for warm-up, then report the median of [WasmScenario.reps] runs. */ +internal fun runEngineScenarios(tier: String) { + for (s in engineScenarios()) { + val params = s.params.entries.joinToString(" ") { "${it.key}=${it.value}" } + val run = try { + s.setup() + } catch (e: Throwable) { + println("[SKAINET_PERF_WASM] tier=$tier scenario=${s.id} metric=${s.metric} value=skipped reason=${e.message} $params") + continue + } + run() + val samples = DoubleArray(s.reps) { run() } + samples.sort() + val median = samples[s.reps / 2] + val text = ((median * 100).roundToInt() / 100.0).toString() + println("[SKAINET_PERF_WASM] tier=$tier scenario=${s.id} metric=${s.metric} value=$text kernel=${KernelRegistry.bestAvailable()?.name} $params") + } +} diff --git a/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.android.kt b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.android.kt new file mode 100644 index 000000000..ebc6c83e5 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.android.kt @@ -0,0 +1,4 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = + System.getProperty(KernelStrictness.SYSTEM_PROPERTY) == "true" diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt index dfbbb7944..468dce0ff 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt @@ -202,6 +202,13 @@ public object KernelDispatch { // blocked weight (per-element block decode) — a process that lands here on a quantized // weight almost certainly forgot to install a kernel pack. Say so once, loudly, even with // no sink attached: silent fallback is how a 25 s/token regression ships unnoticed. + if (b.layout.blocked) { + KernelStrictness.failIfStrict { + "matmul (activation=${a.format.encoding}, weight=${b.format.encoding}, " + + "order=${b.layout.blockOrder}) has no registered kernel; would fall back to the " + + "decoding reference kernel. Registered: ${kernels().joinToString { it.name }.ifEmpty { "" }}" + } + } if (!warnedReferenceFallback && b.layout.blocked) { warnedReferenceFallback = true println( diff --git a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt similarity index 61% rename from skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt rename to skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt index ad4720485..ac3c4b3b7 100644 --- a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.kt @@ -9,31 +9,36 @@ package sk.ainet.backend.api.kernel * rather than as a silent perf regression at forward time (e.g. a * Q4_K matmul that falls back to a scalar dequant + FP32 matmul). * - * This object provides the smallest possible affordance: a system - * property (`-Dskainet.strict.kernels=true`) flips the runtime into - * fail-fast mode. Dispatch sites call [failIfStrict] right before - * they would otherwise silently fall back; the call is a no-op when - * strict mode is off, preserving the existing adaptive behaviour. + * The switch is common code so every tier honours it the same way: + * [overrideEnabled] is the programmatic setting (the only one on JS, Wasm and + * Kotlin/Native, which have no property store); when it is `null` the platform + * setting is read — `-Dskainet.strict.kernels=true` on the JVM and Android. + * Dispatch sites call [failIfStrict] right before they would otherwise silently + * fall back; the call is a no-op when strict mode is off, preserving the + * existing adaptive behaviour. * * Per-context strict mode (e.g. `DirectCpuExecutionContext.create(strict = true)`) * is a follow-up that requires plumbing the flag through every - * platform's `platformDefaultCpuOpsFactory`. The system-property - * affordance is sufficient for tests, CI, and debugging — the - * primary use cases — and ships with zero cross-platform plumbing. + * platform's `platformDefaultCpuOpsFactory`. */ public object KernelStrictness { - /** System-property name that flips fail-fast mode on. */ + /** System-property name that flips fail-fast mode on (JVM and Android). */ public const val SYSTEM_PROPERTY: String = "skainet.strict.kernels" + /** + * Programmatic switch; `null` means "read the platform setting". This is the + * only switch on targets without system properties. + */ + public var overrideEnabled: Boolean? = null + /** * Returns `true` when the runtime should fail fast instead of - * falling back. Reads the system property on every call so a + * falling back. Reads the platform setting on every call so a * test can toggle it via `System.setProperty(...)` between * cases without restarting the JVM. */ - public fun isEnabled(): Boolean = - System.getProperty(SYSTEM_PROPERTY) == "true" + public fun isEnabled(): Boolean = overrideEnabled ?: platformStrictKernels() /** * Throws [NoSuchKernelException] with the supplied message @@ -48,6 +53,9 @@ public object KernelStrictness { } } +/** Platform reading of [KernelStrictness.SYSTEM_PROPERTY]; `false` where there is no property store. */ +internal expect fun platformStrictKernels(): Boolean + /** * Raised by dispatch sites when [KernelStrictness] is enabled and * no registered kernel matches the requested op/dtype combination. diff --git a/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt new file mode 100644 index 000000000..956e7802f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = false diff --git a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.jvm.kt b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.jvm.kt new file mode 100644 index 000000000..ebc6c83e5 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.jvm.kt @@ -0,0 +1,4 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = + System.getProperty(KernelStrictness.SYSTEM_PROPERTY) == "true" diff --git a/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt new file mode 100644 index 000000000..956e7802f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = false diff --git a/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt new file mode 100644 index 000000000..956e7802f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = false diff --git a/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt new file mode 100644 index 000000000..956e7802f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/KernelStrictness.other.kt @@ -0,0 +1,3 @@ +package sk.ainet.backend.api.kernel + +internal actual fun platformStrictKernels(): Boolean = false diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index db6da6530..f340df8e1 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -711,7 +711,15 @@ public open class DefaultCpuOpsBase( // legacy per-element fallback stays one flag away (skainet.dispatch.registry=false) until the // migration is complete. return KernelProfile.timeGeneric { - dispatchMatmulViaRegistry(a, b) ?: matmulGeneric(a, b) + dispatchMatmulViaRegistry(a, b) ?: run { + // Fail-fast point shared by every tier: the legacy per-element path is the silent + // fallback strict mode exists to catch (the JVM subclass checks its SPI paths first). + sk.ainet.backend.api.kernel.KernelStrictness.failIfStrict { + "matmul (${a.dtype.simpleName} ${a.shape} × ${b.dtype.simpleName} ${b.shape}) is not " + + "expressible for KernelDispatch; would fall back to the generic per-element loop" + } + matmulGeneric(a, b) + } } } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt index 798327990..84bf97d50 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt @@ -28,7 +28,7 @@ class KernelSupportMatrixTest { @AfterTest fun reset() = KernelDispatch.clearForTesting() - private val formats = listOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0") + private val formats = listOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0") // platform key (display) -> the set of providers (by source-set) reaching it. private val platforms = listOf("JVM", "Android", "Native·linux", "Native·apple", "JS/WASM") @@ -48,9 +48,11 @@ class KernelSupportMatrixTest { private fun tiers(): List = listOf( Tier("scalar", 0, platforms.toSet(), scalarFormats()), Tier("panama-vector", 50, setOf("JVM"), - setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0")), + setOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0")), + // native-ffm: every accessor of NativeKernelProvider, Q6_K included (DefaultCpuOps resolves + // `matmulQ6K()` from the best available provider, so the FFM kernel wins over Panama). Tier("native-ffm", 100, setOf("JVM"), - setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q5_1", "Q5_0")), + setOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0")), Tier("native-jni", 100, setOf("Android"), setOf("Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), // native-cinterop: NativeKnKernelProvider (nativeMain) — the same C @@ -65,9 +67,22 @@ class KernelSupportMatrixTest { // kernels compiled by Emscripten into a SIMD128 sidecar module (SKEEP-007 Tier 1), with // pthreads on cross-origin-isolated pages and under Node. Installed by WasmSimd.install(); // without it the JS/WASM column is the scalar floor. The row-major `_rm` kernels are also - // served (WasmSimdKernelPack), heap-resident in the sidecar rather than mapped. + // served (WasmSimdKernelPack), heap-resident in the sidecar rather than mapped. BF16 and FP16 + // bind the sidecar's `skainet_bf16_matmul` / `skainet_fp16_matmul` exports. Tier("wasm-simd128", 50, setOf("JS/WASM"), - setOf("Float32", "Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), + setOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), + ) + + /** + * Resident serving: the wasm tier's counterpart of mapped serving. A browser has no file to map, + * so `WasmKernels.weightHandle` / `denseWeightHandle` upload each weight into the sidecar's linear + * memory once and every later matmul reads it in place — the dense and narrow-float kernels and + * all seven packed `_rm` kernels. Hand-declared like the Android mapped row: nothing on a JVM + * test run can instantiate the sidecar. + */ + private fun residentTiers(): List = listOf( + Tier("wasm-resident", 50, setOf("JS/WASM"), + setOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), ) /** @@ -120,6 +135,16 @@ class KernelSupportMatrixTest { .append(cells.joinToString(", ")).append("}}") sb.append(if (i == mappedFormats.lastIndex) "\n" else ",\n") } + sb.append(" ],\n") + val resident = residentTiers() + val residentFormats = formats.filter { fmt -> resident.any { fmt in it.formats } } + sb.append(" \"resident\": [\n") + residentFormats.forEachIndexed { i, fmt -> + val cells = platforms.mapNotNull { p -> best(fmt, p, resident)?.let { "\"$p\": \"$it\"" } } + sb.append(" {\"name\": \"").append(fmt).append("\", \"byPlatform\": {") + .append(cells.joinToString(", ")).append("}}") + sb.append(if (i == residentFormats.lastIndex) "\n" else ",\n") + } sb.append(" ]\n}\n") return sb.toString() } @@ -131,7 +156,7 @@ class KernelSupportMatrixTest { // Drift gate on the scalar floor (the all-platform baseline): the documented set // below must equal what the scalar provider actually carries. Update both together. assertEquals( - setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0"), + setOf("Float32", "BFloat16", "Float16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0"), scalarFormats(), "ScalarKernelProvider coverage changed — update the declared sets + run ./gradlew generateKernelMatrix", ) @@ -140,6 +165,16 @@ class KernelSupportMatrixTest { val providers: List = listOf(ScalarKernelProvider, PanamaVectorKernelProvider, NativeKernelProvider) assertEquals(3, providers.size) + // Drift gate on the declared JVM tiers, where the environment allows it: a provider whose + // accessors answer (`isAvailable()`) must expose exactly the hand-kept set, so a provider + // gaining or losing a kernel fails here. An unavailable provider (no Vector API module, no + // native library on this runner) is skipped — that is the env gate the KDoc describes. + fun exposed(p: KernelProvider) = formats.filter { p.supports("matmul", listOf("Float32", it)) }.toSet() + for (p in listOf(PanamaVectorKernelProvider, NativeKernelProvider)) { + if (!p.isAvailable()) continue + assertEquals(tiers.first { it.name == p.name }.formats, exposed(p), "${p.name} tier drifted from the provider") + } + // Drift gate on mapped serving (#1193): install the JVM-reachable mapped-capable // kernels and assert what they actually serve against StorageCapabilities' // hand-kept mirror — the two lists can't independently drift without failing this test. diff --git a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernelProvider.kt b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernelProvider.kt index 415103c29..77848742e 100644 --- a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernelProvider.kt +++ b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernelProvider.kt @@ -1,5 +1,7 @@ package sk.ainet.exec.wasm +import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel @@ -21,6 +23,8 @@ public object WasmKernelProvider : KernelProvider { override fun isAvailable(): Boolean = WasmKernels.isReady override fun matmulFp32(): Fp32MatmulKernel = WasmFp32MatmulKernel + override fun matmulBf16(): Bf16MatmulKernel = WasmBf16MatmulKernel + override fun matmulFp16(): Fp16MatmulKernel = WasmFp16MatmulKernel override fun matmulQ8_0(): Q8_0MatmulKernel = object : Q8_0MatmulKernel { override fun matmul(input: FloatArray, inputOffset: Int, weight: ByteArray, weightByteOffset: Int, inputDim: Int, outputDim: Int, output: FloatArray, outputOffset: Int) = @@ -73,6 +77,65 @@ public object WasmKernelProvider : KernelProvider { } } +/** + * `C(m,n) = A(m,k) · B(k,n)` with `B` as 16-bit floats (`skainet_bf16_matmul` / `skainet_fp16_matmul`). + * A contiguous `B` (`bByteStride == 2 * n`, whole array) is a weight and stays resident after the + * first call; a strided or windowed `B` is staged per call, as the activation always is. + */ +internal class WasmNarrowFloatMatmulKernel(private val fn: String) { + fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + require(m >= 0 && n >= 0 && k >= 0) { "m, n, k must be non-negative; got m=$m n=$n k=$k" } + if (m == 0 || n == 0) return + val rowBytes = n * 2 + val resident = bByteStride == rowBytes && bByteOffset == 0 && b.size == k * rowBytes + val aPtr = WasmKernels.alloc(m * k * 4) + val bPtr = if (resident) WasmKernels.weightHandle(b, 0, b.size) else WasmKernels.alloc(k * rowBytes) + val cPtr = WasmKernels.alloc(m * n * 4) + try { + if (aStride == k) WasmKernels.writeFloats(aPtr, a, aOffset, m * k) + else for (i in 0 until m) WasmKernels.writeFloats(aPtr + i * k * 4, a, aOffset + i * aStride, k) + if (!resident) { + if (bByteStride == rowBytes) WasmKernels.writeBytes(bPtr, b, bByteOffset, k * rowBytes) + else for (p in 0 until k) WasmKernels.writeBytes(bPtr + p * rowBytes, b, bByteOffset + p * bByteStride, rowBytes) + } + WasmKernels.narrowFloatMatmul(fn, aPtr, k, bPtr, rowBytes, cPtr, n, m, n, k) + if (outStride == n) WasmKernels.readFloats(cPtr, out, outOffset, m * n) + else for (i in 0 until m) WasmKernels.readFloats(cPtr + i * n * 4, out, outOffset + i * outStride, n) + } finally { + WasmKernels.free(cPtr) + if (!resident) WasmKernels.free(bPtr) + WasmKernels.free(aPtr) + } + } +} + +/** BF16 weights in the sidecar (`skainet_bf16_matmul`). */ +public object WasmBf16MatmulKernel : Bf16MatmulKernel { + private val impl = WasmNarrowFloatMatmulKernel("_skainet_bf16_matmul") + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ): Unit = impl.matmul(a, aOffset, aStride, b, bByteOffset, bByteStride, out, outOffset, outStride, m, n, k) +} + +/** FP16 weights in the sidecar (`skainet_fp16_matmul`). */ +public object WasmFp16MatmulKernel : Fp16MatmulKernel { + private val impl = WasmNarrowFloatMatmulKernel("_skainet_fp16_matmul") + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ): Unit = impl.matmul(a, aOffset, aStride, b, bByteOffset, bByteStride, out, outOffset, outStride, m, n, k) +} + /** `C(m,n) = A(m,k) · B(k,n)` in the sidecar; both operands cross per call (dense weights are not cached). */ public object WasmFp32MatmulKernel : Fp32MatmulKernel { override fun matmul( diff --git a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernels.kt b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernels.kt index 5dcd0d2b5..19b00a64c 100644 --- a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernels.kt +++ b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmKernels.kt @@ -205,6 +205,43 @@ public object WasmKernels { /** `skainet_fp32_matmul_wt`: `out(m,n) = in(m,k) · w(n,k)ᵀ`. */ public fun fp32MatmulWt(inPtr: Int, wPtr: Int, outPtr: Int, m: Int, n: Int, k: Int): Unit = jsFp32MatmulWt(mod(), inPtr, wPtr, outPtr, m, n, k) + + /** + * `skainet_bf16_matmul` / `skainet_fp16_matmul`: `C(m,n) = A(m,k) · B(k,n)` with `B` stored as + * 16-bit floats, strides in floats for `A`/`C` and in bytes for `B`. + */ + public fun narrowFloatMatmul(fnName: String, aPtr: Int, aStride: Int, bPtr: Int, bByteStride: Int, cPtr: Int, cStride: Int, m: Int, n: Int, k: Int): Unit = + jsNarrowFloatMatmul(mod(), fnName, aPtr, aStride, bPtr, bByteStride, cPtr, cStride, m, n, k) + + // ── diagnostics ─────────────────────────────────────────────────────────────────────────────── + + /** + * What the sidecar runs on — the wasm counterpart of the JVM's provider banner and CPU-feature + * probes. [hardwareConcurrency] is the browser's (or Node's) logical core count, [threads] whether + * the pthreads build is loaded (it only is when [crossOriginIsolated] or under Node). + */ + public data class SidecarFeatures( + val mode: String, + val simd128: Boolean, + val threads: Boolean, + val crossOriginIsolated: Boolean, + val hardwareConcurrency: Int, + val residentWeightBytes: Long, + ) { + override fun toString(): String = + "mode=$mode simd128=$simd128 threads=$threads crossOriginIsolated=$crossOriginIsolated " + + "cores=$hardwareConcurrency residentWeightBytes=$residentWeightBytes" + } + + /** Current sidecar capabilities; `mode == "none"` before [initialize]. */ + public fun features(): SidecarFeatures = SidecarFeatures( + mode = modeName, + simd128 = isReady, + threads = modeName == "simd128-mt", + crossOriginIsolated = jsCrossOriginIsolated(), + hardwareConcurrency = jsHardwareConcurrency(), + residentWeightBytes = if (isReady) residentWeightBytes() else 0L, + ) } // ── JS glue ─────────────────────────────────────────────────────────────────────────────────────── @@ -260,3 +297,20 @@ private fun jsFp32Matmul(mod: JsAny, aPtr: Int, aStride: Int, bPtr: Int, bStride private fun jsFp32MatmulWt(mod: JsAny, inPtr: Int, wPtr: Int, outPtr: Int, m: Int, n: Int, k: Int): Unit = js("mod._skainet_fp32_matmul_wt(inPtr, 0, wPtr, 0, outPtr, 0, m, n, k)") + +private fun jsNarrowFloatMatmul(mod: JsAny, fnName: String, aPtr: Int, aStride: Int, bPtr: Int, bByteStride: Int, cPtr: Int, cStride: Int, m: Int, n: Int, k: Int): Unit = + js("mod[fnName](aPtr, 0, aStride, bPtr, 0, bByteStride, cPtr, 0, cStride, m, n, k)") + +private fun jsCrossOriginIsolated(): Boolean = js("typeof crossOriginIsolated !== 'undefined' && crossOriginIsolated === true") + +private fun jsHardwareConcurrency(): Int = js( + """ + (function() { + if (typeof navigator !== 'undefined' && typeof navigator.hardwareConcurrency === 'number') return navigator.hardwareConcurrency; + if (typeof process !== 'undefined' && process.versions != null && process.versions.node != null) { + try { return require('os').cpus().length; } catch (e) { return 0; } + } + return 0; + })() + """, +) diff --git a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmSimd.kt b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmSimd.kt index 9d9ac7ace..6038f5afd 100644 --- a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmSimd.kt +++ b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsMain/kotlin/sk/ainet/exec/wasm/WasmSimd.kt @@ -28,8 +28,13 @@ public object WasmSimd { KernelDispatch.clearForTesting() WasmSimdKernelPack.install() WasmCpuOpsOverride.factory = { factory, schedule -> WasmCpuOps(factory, schedule) } + // The same one-line banner the JVM factory prints, so a console shows which tier runs. + println("[SKaiNET] Using SIMD128 sidecar kernels (${WasmKernels.features()})") } + /** What the installed sidecar runs on; `mode == "none"` when not installed. */ + public fun features(): WasmKernels.SidecarFeatures = WasmKernels.features() + /** Back to the scalar lane; frees every resident weight. Tests and model switching. */ public fun uninstall() { WasmCpuOpsOverride.factory = null diff --git a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsTest/kotlin/sk/ainet/exec/wasm/WasmSimdParityTest.kt b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsTest/kotlin/sk/ainet/exec/wasm/WasmSimdParityTest.kt index 838302d87..ae1587591 100644 --- a/skainet-backends/skainet-backend-wasm-simd/src/wasmJsTest/kotlin/sk/ainet/exec/wasm/WasmSimdParityTest.kt +++ b/skainet-backends/skainet-backend-wasm-simd/src/wasmJsTest/kotlin/sk/ainet/exec/wasm/WasmSimdParityTest.kt @@ -2,8 +2,15 @@ package sk.ainet.exec.wasm import kotlinx.coroutines.test.runTest import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.backend.api.kernel.KernelStrictness +import sk.ainet.backend.api.kernel.NoSuchKernelException import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.exec.kernel.ScalarBf16MatmulKernel +import sk.ainet.exec.kernel.ScalarFp16MatmulKernel import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.Q4_0Quantizer @@ -16,6 +23,7 @@ import kotlin.math.abs import kotlin.test.AfterTest import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertFailsWith import kotlin.test.assertIs import kotlin.test.assertTrue @@ -65,11 +73,70 @@ class WasmSimdParityTest { fun `install swaps the ops and the sidecar reports a SIMD mode`() = runTest { val before = DirectCpuExecutionContext() assertIs(before.ops) + assertEquals("none", WasmKernels.features().mode) WasmSimd.install() assertTrue(WasmSimd.isInstalled) assertTrue(WasmKernels.mode.startsWith("simd128-"), WasmKernels.mode) val after = DirectCpuExecutionContext() assertIs(after.ops) + val f = WasmSimd.features() + assertEquals(WasmKernels.mode, f.mode) + assertTrue(f.simd128) + assertEquals(f.mode == "simd128-mt", f.threads) + assertTrue(f.hardwareConcurrency >= 1, "cores: ${f.hardwareConcurrency}") + } + + @Test + fun `BF16 and FP16 weights run on the sidecar and match the scalar kernels`() = runTest { + WasmSimd.install() + val provider = checkNotNull(KernelRegistry.bestAvailable()) + assertEquals("wasm-simd128", provider.name) + val m = 7; val n = 33; val k = 40 + val a = pseudo(m * k, 21) + val wf = pseudo(k * n, 23) + val bf16 = ByteArray(k * n * 2) + val fp16 = ByteArray(k * n * 2) + for (i in 0 until k * n) { + val b = Bf16Codec.encode(wf[i]); bf16[i * 2] = (b and 0xFF).toByte(); bf16[i * 2 + 1] = ((b ushr 8) and 0xFF).toByte() + val h = Fp16Codec.encode(wf[i]); fp16[i * 2] = (h and 0xFF).toByte(); fp16[i * 2 + 1] = ((h ushr 8) and 0xFF).toByte() + } + val refB = FloatArray(m * n); val refH = FloatArray(m * n) + ScalarBf16MatmulKernel.matmul(a, 0, k, bf16, 0, n * 2, refB, 0, n, m, n, k) + ScalarFp16MatmulKernel.matmul(a, 0, k, fp16, 0, n * 2, refH, 0, n, m, n, k) + val gotB = FloatArray(m * n); val gotH = FloatArray(m * n) + checkNotNull(provider.matmulBf16()).matmul(a, 0, k, bf16, 0, n * 2, gotB, 0, n, m, n, k) + checkNotNull(provider.matmulFp16()).matmul(a, 0, k, fp16, 0, n * 2, gotH, 0, n, m, n, k) + assertClose(refB, gotB, "bf16") + assertClose(refH, gotH, "fp16") + // A contiguous weight is resident after the first call; a strided window is staged per call. + assertTrue(WasmKernels.residentWeightBytes() >= 2L * k * n * 2) + val sub = FloatArray(m * 8) + val refSub = FloatArray(m * 8) + ScalarBf16MatmulKernel.matmul(a, 0, k, bf16, 4, n * 2, refSub, 0, 8, m, 8, k) + checkNotNull(provider.matmulBf16()).matmul(a, 0, k, bf16, 4, n * 2, sub, 0, 8, m, 8, k) + assertClose(refSub, sub, "bf16 window") + } + + @Test + fun `strict mode is honoured on wasm`() = runTest { + KernelStrictness.overrideEnabled = true + try { + assertTrue(KernelStrictness.isEnabled()) + val ctx = DirectCpuExecutionContext() + // Packed weight, no pack installed and the reference fallback forbidden: fail fast. + val w = q8(ctx, 64, 96, 5) + val x = ctx.fromFloatArray(Shape(2, 96), FP32::class, pseudo(2 * 96, 1)) + KernelDispatch.clearForTesting() + // A canonical row-major Q8_0 view has no exact scalar kernel (those are keyed input-block-major) + // and prepack is off by default here, so this is the reference-fallback site. + val out = sk.ainet.lang.memory.TensorView.dense(sk.ainet.lang.memory.Storage.Heap.floats(2 * 64), Shape(2, 64), sk.ainet.lang.types.FP32) + assertFailsWith { + KernelDispatch.matmul(x.data.view!!, w.data.view!!, out, sk.ainet.lang.memory.Scope.Ambient) + } + } finally { + KernelStrictness.overrideEnabled = null + } + assertEquals(false, KernelStrictness.isEnabled()) } @Test