From 3313647ef0fa2ba3785464d463157d1fdea89c30 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 28 Apr 2026 14:21:26 +0200 Subject: [PATCH] feat(kernel): add Panama Vector FP32 matmul provider (priority 50) Implements `PanamaVectorMatmulKernel` (jdk.incubator.vector, FloatVector + fma + reduceLanes) and `PanamaVectorKernelProvider` against the kernel SPI from PR #554. Picks up automatically over `ScalarKernelProvider` once registered, and respects the existing `-Dskainet.cpu.vector.enabled=false` kill switch. Closes the M5 "Panama-first" half of the JVM perf milestone plan. Routing `DefaultCpuOpsJvm.matmul` through the SPI and adding a ServiceLoader-based auto-registration are deferred to follow-ups so this PR stays focused on the kernel itself. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../exec/kernel/PanamaVectorKernelProvider.kt | 62 ++++++++ .../exec/kernel/PanamaVectorMatmulKernel.kt | 86 +++++++++++ .../kernel/PanamaVectorKernelProviderTest.kt | 66 +++++++++ .../kernel/PanamaVectorMatmulKernelTest.kt | 138 ++++++++++++++++++ 4 files changed, 352 insertions(+) create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProviderTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernelTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt new file mode 100644 index 000000000..3b3f75b4b --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt @@ -0,0 +1,62 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.Fp32MatmulKernel +import sk.ainet.backend.api.kernel.KernelProvider +import sk.ainet.exec.tensor.ops.JvmCpuBackendConfig + +/** + * JVM Vector API (`jdk.incubator.vector`) [KernelProvider]. Available + * when the runtime is JDK 21+, the incubator module is loaded + * (`--add-modules jdk.incubator.vector`), and the + * `skainet.cpu.vector.enabled` kill switch hasn't been flipped to + * `false`. + * + * Priority is `50` — above [ScalarKernelProvider] (`0`) and below a + * future hand-tuned native provider (`100`). Concrete kernels are + * exposed via the per-kernel accessors; today only [matmulFp32] is + * specialized — other accessors fall back to `null` so callers can + * cascade to a lower-priority provider. + * + * Registration is **manual** (per the kernel-SPI contract today): the + * runtime that wants this provider must call + * `KernelRegistry.register(PanamaVectorKernelProvider)` at startup. + * Auto-registration via `ServiceLoader` will be layered on once a + * second concrete JVM provider exists. + */ +public object PanamaVectorKernelProvider : KernelProvider { + override val name: String = "panama-vector" + override val priority: Int = 50 + + private val cachedAvailable: Boolean by lazy { + isJdk21Plus() && isVectorApiClassLoaded() + } + + override fun isAvailable(): Boolean = + cachedAvailable && JvmCpuBackendConfig.vectorEnabled + + override fun matmulFp32(): Fp32MatmulKernel? = + if (isAvailable()) PanamaVectorMatmulKernel else null + + private fun isVectorApiClassLoaded(): Boolean = runCatching { + Class.forName("jdk.incubator.vector.FloatVector") + Class.forName("jdk.incubator.vector.VectorSpecies") + true + }.getOrElse { false } + + private fun isJdk21Plus(): Boolean { + val runtimeFeature = runCatching { + val runtimeClass = Class.forName("java.lang.Runtime") + val versionMethod = runtimeClass.getMethod("version") + val versionObj = versionMethod.invoke(Runtime.getRuntime()) + val featureMethod = versionObj.javaClass.getMethod("feature") + featureMethod.invoke(versionObj) as Int + }.getOrNull() + if (runtimeFeature != null) return runtimeFeature >= 21 + + val spec = System.getProperty("java.specification.version") ?: return false + return spec.toIntOrNull()?.let { it >= 21 } ?: run { + val major = spec.split('.', '-').firstOrNull()?.toIntOrNull() ?: return@run false + major >= 21 + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt new file mode 100644 index 000000000..f2fb9f05c --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt @@ -0,0 +1,86 @@ +package sk.ainet.exec.kernel + +import jdk.incubator.vector.FloatVector +import jdk.incubator.vector.VectorOperators +import jdk.incubator.vector.VectorSpecies +import sk.ainet.backend.api.kernel.Fp32MatmulKernel + +/** + * SIMD reference [Fp32MatmulKernel] implemented on the JDK Vector API + * (JEP 338+, `jdk.incubator.vector`). Produces results that match + * [ScalarMatmulKernel] within FP-rounding tolerance. + * + * Strategy: + * - Pack `B` into a transposed buffer `bt` of shape `(n, k)` so the + * inner reduction streams contiguously over `k` for both operands — + * `a[i, kk]` walks one row of `A` and `bt[j, kk]` walks one row of + * the packed transpose. + * - Inner loop is a vector-width FMA accumulator (`v.fma(w, acc)`), + * reduced once per `(i, j)` pair via `reduceLanes(ADD)`. + * - Tail elements that don't fill a vector lane are handled in scalar. + * + * The B-pack is `O(n * k)` floats per call; that's cheap relative to + * the `O(m * n * k)` FLOPs but still allocates each invocation. A + * scratch-pool integration is out of scope for this kernel and lives + * one layer up (see `ScratchPool` SPI in `skainet-lang-core`). + * + * Caller contract is identical to [Fp32MatmulKernel]: strides are in + * floats, `out` is fully overwritten in the `m × n` block, and `k == 0` + * zeros the output block. + */ +public object PanamaVectorMatmulKernel : Fp32MatmulKernel { + private val species: VectorSpecies = FloatVector.SPECIES_PREFERRED + + 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) { + "PanamaVectorMatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + 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 + } + + // Pack B^T: bt[j, kk] = b[kk, j]. + val bt = FloatArray(n * k) + for (kk in 0 until k) { + val src = bOffset + kk * bStride + for (j in 0 until n) { + bt[j * k + kk] = b[src + j] + } + } + + val step = species.length() + val loopBound = species.loopBound(k) + + for (i in 0 until m) { + val aRow = aOffset + i * aStride + val outRow = outOffset + i * outStride + for (j in 0 until n) { + val btRow = j * k + var acc = FloatVector.zero(species) + var idx = 0 + while (idx < loopBound) { + val va = FloatVector.fromArray(species, a, aRow + idx) + val vb = FloatVector.fromArray(species, bt, btRow + idx) + acc = va.fma(vb, acc) + idx += step + } + var sum = acc.reduceLanes(VectorOperators.ADD) + while (idx < k) { + sum += a[aRow + idx] * bt[btRow + idx] + idx++ + } + out[outRow + j] = sum + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProviderTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProviderTest.kt new file mode 100644 index 000000000..b99520a90 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProviderTest.kt @@ -0,0 +1,66 @@ +package sk.ainet.exec.kernel + +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue +import sk.ainet.backend.api.kernel.KernelRegistry + +class PanamaVectorKernelProviderTest { + + @BeforeTest + fun setUp() = KernelRegistry.clearForTesting() + + @AfterTest + fun tearDown() = KernelRegistry.clearForTesting() + + @Test + fun providerHasExpectedNameAndPriority() { + assertEquals("panama-vector", PanamaVectorKernelProvider.name) + assertEquals(50, PanamaVectorKernelProvider.priority) + } + + @Test + fun isAvailableOnTestJdk() { + // The cpu-backend test suite runs on JDK 21+ with the incubator + // module on the module path (see jvm-cpu-jmh build script and the + // project's JDK requirement). Vector should be available here. + assertTrue( + PanamaVectorKernelProvider.isAvailable(), + "expected Panama provider to be available on the test JDK", + ) + } + + @Test + fun matmulFp32IsTheVectorKernelWhenAvailable() { + assertSame(PanamaVectorMatmulKernel, PanamaVectorKernelProvider.matmulFp32()) + } + + @Test + fun beatsScalarInRegistryWhenBothRegistered() { + KernelRegistry.register(ScalarKernelProvider) + KernelRegistry.register(PanamaVectorKernelProvider) + // Higher priority wins. + assertSame(PanamaVectorKernelProvider, KernelRegistry.bestAvailable()) + assertEquals( + listOf("panama-vector", "scalar"), + KernelRegistry.availableNames(), + ) + } + + @Test + fun killSwitchDisablesProvider() { + val key = "skainet.cpu.vector.enabled" + val previous = System.getProperty(key) + try { + System.setProperty(key, "false") + assertEquals(false, PanamaVectorKernelProvider.isAvailable()) + assertEquals(null, PanamaVectorKernelProvider.matmulFp32()) + } finally { + if (previous == null) System.clearProperty(key) + else System.setProperty(key, previous) + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernelTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernelTest.kt new file mode 100644 index 000000000..cdb053bb7 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernelTest.kt @@ -0,0 +1,138 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Parity tests for [PanamaVectorMatmulKernel]. Every case runs the + * Panama kernel and the [ScalarMatmulKernel] reference on the same + * inputs and asserts the outputs agree within FP-rounding tolerance + * (FMA + reordered reduction can differ from a left-to-right scalar + * sum at the last few ULP). + * + * Tolerance scales with the contraction dimension `k`: each summand + * carries up to ~`eps * |a|*|b|` rounding error, and we accumulate `k` + * of them. `1e-5 * k` is comfortable for the inputs used here + * (clamped to `[-0.5, 0.5]`). + */ +class PanamaVectorMatmulKernelTest { + + private fun assertParity( + m: Int, n: Int, k: Int, + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + outStride: Int, + ) { + val outScalar = FloatArray(m * outStride) + val outPanama = FloatArray(m * outStride) + ScalarMatmulKernel.matmul( + a, aOffset, aStride, + b, bOffset, bStride, + outScalar, 0, outStride, + m, n, k, + ) + PanamaVectorMatmulKernel.matmul( + a, aOffset, aStride, + b, bOffset, bStride, + outPanama, 0, outStride, + m, n, k, + ) + val tol = (1e-5f * k.coerceAtLeast(1)).coerceAtLeast(1e-5f) + assertEquals(outScalar.size, outPanama.size, "length mismatch") + for (i in outScalar.indices) { + val diff = abs(outScalar[i] - outPanama[i]) + assertTrue( + diff <= tol, + "mismatch at $i: scalar=${outScalar[i]} panama=${outPanama[i]} diff=$diff tol=$tol", + ) + } + } + + @Test + fun small_2x3x4_contiguous_matches_scalar() { + val a = floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f) // [2, 4] + val b = FloatArray(4 * 3) { it.toFloat() } // [4, 3] + assertParity(m = 2, n = 3, k = 4, a = a, aOffset = 0, aStride = 4, b = b, bOffset = 0, bStride = 3, outStride = 3) + } + + @Test + fun random_8x16x32_matches_scalar() { + val rng = Random(42) + val a = FloatArray(8 * 32) { rng.nextFloat() - 0.5f } + val b = FloatArray(32 * 16) { rng.nextFloat() - 0.5f } + assertParity(m = 8, n = 16, k = 32, a = a, aOffset = 0, aStride = 32, b = b, bOffset = 0, bStride = 16, outStride = 16) + } + + @Test + fun non_aligned_k_exercises_tail_loop() { + // k = 23 is not a multiple of any common vector lane count (4, 8, 16), + // so this forces the scalar tail loop to run. + val rng = Random(1234) + val m = 5; val n = 7; val k = 23 + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val b = FloatArray(k * n) { rng.nextFloat() - 0.5f } + assertParity(m = m, n = n, k = k, a = a, aOffset = 0, aStride = k, b = b, bOffset = 0, bStride = n, outStride = n) + } + + @Test + fun strided_a_sub_block_matches_scalar() { + // Parent A is [4, 8]; take rows 1..2 as a 2×8 sub-block. + val parentA = FloatArray(4 * 8) { it.toFloat() } + val b = FloatArray(8 * 3) { (it + 1).toFloat() } + assertParity( + m = 2, n = 3, k = 8, + a = parentA, aOffset = 1 * 8, aStride = 8, + b = b, bOffset = 0, bStride = 3, + outStride = 3, + ) + } + + @Test + fun large_irregular_31x17x23_matches_scalar() { + val rng = Random(7) + val m = 31; val n = 17; val k = 23 + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val b = FloatArray(k * n) { rng.nextFloat() - 0.5f } + assertParity(m = m, n = n, k = k, a = a, aOffset = 0, aStride = k, b = b, bOffset = 0, bStride = n, outStride = n) + } + + @Test + fun zero_m_or_n_no_op() { + val out = FloatArray(5) { 7f } + PanamaVectorMatmulKernel.matmul( + FloatArray(0), 0, 0, + FloatArray(0), 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 } + PanamaVectorMatmulKernel.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() { + assertFailsWith { + PanamaVectorMatmulKernel.matmul( + FloatArray(0), 0, 0, + FloatArray(0), 0, 0, + FloatArray(0), 0, 0, + m = -1, n = 1, k = 1, + ) + } + } +}