From c3f4b81b485374e88929ef6bf058eb5f49f6f88a Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 28 Apr 2026 13:25:15 +0200 Subject: [PATCH] feat(kernel): add KernelProvider SPI for matmul dispatch (Scalar baseline) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #553. Introduces a small SPI between high-level tensor ops (`TensorOps.matmul` et al.) and the actual numeric kernels that do the FLOPs. This is the groundwork that lets a SIMD-accelerated matmul be plugged in without re-implementing the rest of an op-level backend, and lets a hand-written kernel be tested against a scalar reference. What lands: * `sk.ainet.backend.api.kernel.Fp32MatmulKernel` - `C(m, n) = A(m, k) · B(k, n)` row-major - element-stride parameters for caller sub-blocks (no copy needed) - implementations must not mutate inputs / must overwrite the m×n block of out * `sk.ainet.backend.api.kernel.KernelProvider` - `name` / `priority` / `isAvailable()` / per-kernel accessors - per-accessor `null` lets callers fall through to a lower-priority provider when the higher one doesn't ship the kernel * `sk.ainet.backend.api.kernel.KernelRegistry` - process-wide manual registration - `register()` / `find(name)` / `bestAvailable()` / `availableNames()` - `clearForTesting()` for tests - JVM ServiceLoader auto-discovery deferred to a follow-up PR (only one provider ships today; the registry shape supports it without further interface changes) * `sk.ainet.exec.kernel.ScalarMatmulKernel` + `ScalarKernelProvider` (in `skainet-backend-cpu`) - triple-nested-loop reference; honours stride parameters - priority = 0; always available - guaranteed correctness reference and runtime fallback * Tests: - `ScalarMatmulKernelTest`: small / medium / strided sub-blocks on both A and out / zero-m / zero-k / rejects negatives - `KernelRegistryTest`: empty / scalar-only / priority ordering / skip-unavailable / case-insensitive name lookup / re-register no-op Out of scope (separate issues / PRs): * Panama Vector matmul (the actual perf win on JVM). * Native FFM matmul. * Wiring `DefaultCpuOps.matmul` to consult the registry — needs at least one accelerated provider to make the dispatch worth doing. * SDPA kernel API. * Quantized kernels (Q4_K, Q8). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../backend/api/kernel/Fp32MatmulKernel.kt | 50 +++++++ .../backend/api/kernel/KernelProvider.kt | 45 ++++++ .../backend/api/kernel/KernelRegistry.kt | 57 ++++++++ .../ainet/exec/kernel/ScalarKernelProvider.kt | 24 ++++ .../ainet/exec/kernel/ScalarMatmulKernel.kt | 50 +++++++ .../ainet/exec/kernel/KernelRegistryTest.kt | 78 +++++++++++ .../exec/kernel/ScalarMatmulKernelTest.kt | 130 ++++++++++++++++++ 7 files changed, 434 insertions(+) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Fp32MatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/KernelRegistryTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernelTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Fp32MatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Fp32MatmulKernel.kt new file mode 100644 index 000000000..163976136 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Fp32MatmulKernel.kt @@ -0,0 +1,50 @@ +package sk.ainet.backend.api.kernel + +/** + * FP32 matrix multiplication kernel: `C(m, n) = A(m, k) · B(k, n)` in + * row-major layout. + * + * This is a thin SPI between high-level tensor ops and the actual + * numeric kernel that does the FLOPs. It exists so a SIMD-accelerated + * `matmul` can be plugged in without re-implementing the rest of an + * op-level backend, and so a hand-written kernel can be tested against + * a scalar reference. + * + * Strides are in **floats** (not bytes) and let callers pass sub-blocks + * of larger arrays without copying. For a contiguous matrix of shape + * `(m, k)`, `aStride == k`. For a sub-block, `aStride` is the leading + * dimension of the *parent* matrix. + * + * Implementations must NOT mutate `a` or `b`. They MAY assume the + * arrays do not alias each other or `out`. Implementations MUST fully + * overwrite the `m × n` block of `out` they're responsible for — + * accumulator semantics are caller-controlled (e.g. zero `out` first if + * you want C = A·B; pre-fill `out` if you want C += A·B and the kernel + * is fused for that — no fused-accumulate kernel is in scope yet). + */ +public interface Fp32MatmulKernel { + /** + * @param a left operand `(m, k)`, row-major, with stride `aStride` along + * the leading (row) dimension. + * @param aOffset element offset into [a] where the (0, 0) entry lives. + * @param aStride distance in floats between consecutive rows of [a]. + * For a contiguous matrix this equals `k`. + * @param b right operand `(k, n)`, row-major, with stride `bStride`. + * @param bOffset element offset into [b]. + * @param bStride distance in floats between consecutive rows of [b]. + * For a contiguous matrix this equals `n`. + * @param out output `(m, n)`, row-major, with stride `outStride`. + * @param outOffset element offset into [out]. + * @param outStride distance in floats between consecutive rows of [out]. + * For a contiguous matrix this equals `n`. + * @param m number of rows of A and C. + * @param n number of columns of B and C. + * @param k contraction dimension (cols of A == rows of B). + */ + public fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int + ) +} diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt new file mode 100644 index 000000000..ff82197df --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt @@ -0,0 +1,45 @@ +package sk.ainet.backend.api.kernel + +/** + * Provider for a family of related numeric kernels (matmul, SDPA, ...). + * + * A backend (Panama Vector, native FFM, IREE, Metal, ...) bundles its + * kernels behind a single provider so callers only need to know the + * top-level provider name. Providers self-report whether they are + * available on the current platform / runtime — e.g. a Panama Vector + * provider returns `false` from [isAvailable] on a JDK that doesn't + * have the incubator module loaded. + * + * Lookup rules: + * - Higher [priority] wins when multiple providers report + * [isAvailable] = `true`. Providers should rank themselves by + * expected performance: scalar ≈ 0, Panama Vector ≈ 50, hand-tuned + * native ≈ 100. + * - Each per-kernel accessor returns `null` when the provider does not + * carry that kernel, so callers can fall through to a lower-priority + * provider. + */ +public interface KernelProvider { + /** Stable, human-readable identifier. */ + public val name: String + + /** + * Relative ranking versus other providers. Higher = preferred when + * available. The scalar reference uses `0`; SIMD-accelerated + * providers should use a larger value. + */ + public val priority: Int + + /** + * Reports whether this provider's kernels can run in the current + * process. Expensive checks (probing CPU features, loading native + * libraries) should be done once and cached. + */ + public fun isAvailable(): Boolean + + /** + * FP32 matmul kernel exposed by this provider, or `null` if this + * provider does not specialize matmul. + */ + public fun matmulFp32(): Fp32MatmulKernel? +} diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt new file mode 100644 index 000000000..9568b8a15 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelRegistry.kt @@ -0,0 +1,57 @@ +package sk.ainet.backend.api.kernel + +/** + * Process-wide registry of [KernelProvider] instances. + * + * Backends that ship a [KernelProvider] register it via [register] at + * load time. Callers that need a kernel ask [bestAvailable] (or + * [find] with a name) for the highest-priority provider that reports + * itself available, then pull the specific kernel they need from the + * provider's accessors. + * + * The registry is plain manual registration today — JVM auto-discovery + * via `java.util.ServiceLoader` can be layered on in a follow-up PR + * once a second concrete provider exists (Panama Vector). Callers that + * want a guaranteed scalar fallback can pin + * `sk.ainet.exec.kernel.ScalarKernelProvider` directly without going + * through the registry. + * + * Thread safety: [register] is not thread-safe. Call it during + * single-threaded startup or guard with your own lock. + */ +public object KernelRegistry { + private val providers: MutableList = mutableListOf() + + /** + * Register a provider. Re-registering the same instance is a no-op. + */ + public fun register(provider: KernelProvider) { + if (providers.any { it === provider }) return + providers.add(provider) + providers.sortByDescending { it.priority } + } + + /** All registered providers, sorted by priority descending. */ + public fun providers(): List = providers.toList() + + /** Find a provider by name (case-insensitive), or `null`. */ + public fun find(name: String): KernelProvider? = + providers.firstOrNull { it.name.equals(name, ignoreCase = true) } + + /** + * Highest-priority [isAvailable] provider, or `null` if none is + * registered or available. Callers that absolutely need a kernel + * should use the explicit scalar fallback instead. + */ + public fun bestAvailable(): KernelProvider? = + providers.firstOrNull { it.isAvailable() } + + /** Names of all currently-available providers. */ + public fun availableNames(): List = + providers.filter { it.isAvailable() }.map { it.name } + + /** Test/diagnostic helper. Removes all registered providers. */ + public fun clearForTesting() { + providers.clear() + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt new file mode 100644 index 000000000..3418cb11a --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt @@ -0,0 +1,24 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.Fp32MatmulKernel +import sk.ainet.backend.api.kernel.KernelProvider + +/** + * Scalar (non-SIMD) [KernelProvider] — always available, lowest + * priority. Acts as the correctness reference and the guaranteed + * fallback when no accelerated provider is registered. + * + * Callers can pin this provider directly when they want deterministic + * scalar arithmetic without registry interaction (useful in tests): + * + * ```kotlin + * val kernel = ScalarKernelProvider.matmulFp32()!! + * kernel.matmul(a, 0, k, b, 0, n, out, 0, n, m, n, k) + * ``` + */ +public object ScalarKernelProvider : KernelProvider { + override val name: String = "scalar" + override val priority: Int = 0 + override fun isAvailable(): Boolean = true + override fun matmulFp32(): Fp32MatmulKernel = ScalarMatmulKernel +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernel.kt new file mode 100644 index 000000000..c09b8a374 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernel.kt @@ -0,0 +1,50 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.Fp32MatmulKernel + +/** + * Scalar reference implementation of [Fp32MatmulKernel] — three nested + * loops, no SIMD. Always available on every KMP target. Used as: + * + * - The correctness reference that accelerated kernels (Panama, native) + * must match bit-for-bit (within FP order tolerance). + * - A guaranteed fallback when no accelerated provider is registered or + * available. + * + * Performance is modest (no vectorization, no cache-blocking), so + * production code should layer a Panama or native provider on top via + * [sk.ainet.backend.api.kernel.KernelRegistry]. + */ +public object ScalarMatmulKernel : Fp32MatmulKernel { + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int + ) { + require(m >= 0 && n >= 0 && k >= 0) { + "ScalarMatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + // k == 0 → C = 0; the strides may still be > 0 so we need to + // explicitly zero the output block. + if (k == 0) { + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + return + } + for (i in 0 until m) { + val aRowOff = aOffset + i * aStride + val outRowOff = outOffset + i * outStride + for (j in 0 until n) { + var sum = 0f + for (kk in 0 until k) { + sum += a[aRowOff + kk] * b[bOffset + kk * bStride + j] + } + out[outRowOff + j] = sum + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/KernelRegistryTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/KernelRegistryTest.kt new file mode 100644 index 000000000..8b82d676e --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/KernelRegistryTest.kt @@ -0,0 +1,78 @@ +package sk.ainet.exec.kernel + +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.test.assertSame +import sk.ainet.backend.api.kernel.Fp32MatmulKernel +import sk.ainet.backend.api.kernel.KernelProvider +import sk.ainet.backend.api.kernel.KernelRegistry + +class KernelRegistryTest { + + @BeforeTest + fun setUp() = KernelRegistry.clearForTesting() + + @AfterTest + fun tearDown() = KernelRegistry.clearForTesting() + + @Test + fun emptyRegistryHasNoBest() { + assertNull(KernelRegistry.bestAvailable()) + assertEquals(emptyList(), KernelRegistry.availableNames()) + } + + @Test + fun scalarRegistersAndIsBest() { + KernelRegistry.register(ScalarKernelProvider) + assertSame(ScalarKernelProvider, KernelRegistry.bestAvailable()) + assertEquals(listOf("scalar"), KernelRegistry.availableNames()) + } + + @Test + fun higherPriorityWins() { + val fast = object : KernelProvider { + override val name = "fake-fast" + override val priority = 50 + override fun isAvailable() = true + override fun matmulFp32(): Fp32MatmulKernel = ScalarMatmulKernel + } + KernelRegistry.register(ScalarKernelProvider) + KernelRegistry.register(fast) + assertSame(fast, KernelRegistry.bestAvailable()) + } + + @Test + fun unavailableProviderIsSkipped() { + val pretender = object : KernelProvider { + override val name = "pretender" + override val priority = 100 + override fun isAvailable() = false + override fun matmulFp32(): Fp32MatmulKernel = ScalarMatmulKernel + } + KernelRegistry.register(pretender) + KernelRegistry.register(ScalarKernelProvider) + // pretender outranks scalar but isn't available — scalar wins. + assertSame(ScalarKernelProvider, KernelRegistry.bestAvailable()) + // availableNames excludes pretender. + assertEquals(listOf("scalar"), KernelRegistry.availableNames()) + } + + @Test + fun findByNameIsCaseInsensitive() { + KernelRegistry.register(ScalarKernelProvider) + assertSame(ScalarKernelProvider, KernelRegistry.find("scalar")) + assertSame(ScalarKernelProvider, KernelRegistry.find("Scalar")) + assertSame(ScalarKernelProvider, KernelRegistry.find("SCALAR")) + assertNull(KernelRegistry.find("unknown")) + } + + @Test + fun reRegisteringSameInstanceIsNoOp() { + KernelRegistry.register(ScalarKernelProvider) + KernelRegistry.register(ScalarKernelProvider) + assertEquals(1, KernelRegistry.providers().size) + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernelTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernelTest.kt new file mode 100644 index 000000000..55cd3602b --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/kernel/ScalarMatmulKernelTest.kt @@ -0,0 +1,130 @@ +package sk.ainet.exec.kernel + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith + +class ScalarMatmulKernelTest { + + private fun reference(a: FloatArray, b: FloatArray, m: Int, n: Int, k: Int): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) for (j in 0 until n) { + var s = 0f + for (kk in 0 until k) s += a[i * k + kk] * b[kk * n + j] + out[i * n + j] = s + } + return out + } + + private fun assertNearlyEquals(expected: FloatArray, actual: FloatArray, tol: Float = 1e-5f) { + assertEquals(expected.size, actual.size, "length mismatch") + for (i in expected.indices) { + val diff = kotlin.math.abs(expected[i] - actual[i]) + assertEquals(true, diff < tol, "mismatch at $i: expected ${expected[i]} actual ${actual[i]} diff $diff") + } + } + + @Test + fun small_2x3x4_contiguous() { + val a = floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f) // shape [2, 4] + val b = FloatArray(4 * 3) { it.toFloat() } // shape [4, 3] + val out = FloatArray(2 * 3) + ScalarMatmulKernel.matmul( + a, 0, 4, + b, 0, 3, + out, 0, 3, + m = 2, n = 3, k = 4 + ) + assertNearlyEquals(reference(a, b, 2, 3, 4), out) + } + + @Test + fun deterministic_random_8x16x32() { + val rng = kotlin.random.Random(42) + val a = FloatArray(8 * 32) { rng.nextFloat() - 0.5f } + val b = FloatArray(32 * 16) { rng.nextFloat() - 0.5f } + val out = FloatArray(8 * 16) + ScalarMatmulKernel.matmul( + a, 0, 32, + b, 0, 16, + out, 0, 16, + m = 8, n = 16, k = 32 + ) + assertNearlyEquals(reference(a, b, 8, 16, 32), out) + } + + @Test + fun stride_supports_sub_blocks() { + // Parent A is shape [4, 8]. Take rows 1..2 as a 2×8 sub-block. + // aOffset = 1 * 8 = 8, aStride = 8 (parent leading dim). + val parentA = FloatArray(4 * 8) { it.toFloat() } + val b = FloatArray(8 * 3) { (it + 1).toFloat() } + val out = FloatArray(2 * 3) + ScalarMatmulKernel.matmul( + parentA, 8, 8, + b, 0, 3, + out, 0, 3, + m = 2, n = 3, k = 8 + ) + // Compare: extract rows 1..2 of parentA into a contiguous [2, 8] buffer, + // run the reference. + val subA = FloatArray(2 * 8) { idx -> parentA[8 + idx] } + assertNearlyEquals(reference(subA, b, 2, 3, 8), out) + } + + @Test + fun out_stride_supports_partial_writes() { + // Output is a sub-block of a larger 4×6 buffer; write a 2×3 result + // at rows 1..2, cols 1..3. + val a = FloatArray(2 * 5) { (it + 1).toFloat() } + val b = FloatArray(5 * 3) { (it + 1).toFloat() } + val parentOut = FloatArray(4 * 6) + ScalarMatmulKernel.matmul( + a, 0, 5, + b, 0, 3, + parentOut, /* row 1, col 1 */ 1 * 6 + 1, 6, + m = 2, n = 3, k = 5 + ) + val expected = reference(a, b, 2, 3, 5) + for (i in 0 until 2) for (j in 0 until 3) { + assertEquals(expected[i * 3 + j], parentOut[(1 + i) * 6 + (1 + j)], + "output (sub-block) mismatch at parent[${1+i}][${1+j}]") + } + // Outside the written sub-block, parentOut is still zero. + for (i in 0 until 4) for (j in 0 until 6) { + val inside = i in 1..2 && j in 1..3 + if (!inside) { + assertEquals(0f, parentOut[i * 6 + j], "parent[$i][$j] should be untouched") + } + } + } + + @Test + fun zero_m_or_n_no_op() { + val a = FloatArray(0) + val b = FloatArray(0) + val out = FloatArray(5) { 7f } + ScalarMatmulKernel.matmul(a, 0, 0, b, 0, 0, out, 0, 0, m = 0, n = 5, k = 0) + for (v in out) assertEquals(7f, v, "out should be unchanged when m == 0") + } + + @Test + fun zero_k_zeros_output() { + val out = FloatArray(2 * 3) { 9f } + ScalarMatmulKernel.matmul( + FloatArray(0), 0, 0, + FloatArray(0), 0, 0, + out, 0, 3, + m = 2, n = 3, k = 0 + ) + for (v in out) assertEquals(0f, v, "out block should be zeroed when k == 0") + } + + @Test + fun rejects_negative_dimensions() { + val a = FloatArray(0); val b = FloatArray(0); val out = FloatArray(0) + assertFailsWith { + ScalarMatmulKernel.matmul(a, 0, 0, b, 0, 0, out, 0, 0, m = -1, n = 1, k = 1) + } + } +}