feat(kernel): Panama Vector FP32 matmul provider (M5)

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
+        }
+    }
+}

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<Float> = 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
+            }
+        }
+    }
+}

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)
+        }
+    }
+}

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<IllegalArgumentException> {
+            PanamaVectorMatmulKernel.matmul(
+                FloatArray(0), 0, 0,
+                FloatArray(0), 0, 0,
+                FloatArray(0), 0, 0,
+                m = -1, n = 1, k = 1,
+            )
+        }
+    }
+}