diff --git a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index f4672e0a..3c5298b1 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -12,6 +12,7 @@ endif() add_library(skainet_kernels SHARED src/skainet_smoke.c src/q4k_matmul.c + src/fp32_matmul.c ) target_include_directories(skainet_kernels PUBLIC diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h index c4c6a36d..bc614895 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h +++ b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h @@ -60,6 +60,22 @@ SKAINET_API void skainet_q4k_matmul( int32_t output_offset ); +/* + * Row-major FP32 SGEMM: C(m, n) = A(m, k) * B(k, n). + * + * Strides are in floats (not bytes). For a contiguous parent matrix + * `a_stride == k`, `b_stride == n`, `c_stride == n`. The kernel zeros + * the m×n output block before accumulating, so callers always get + * `C = A·B` (not `C += A·B`). `k == 0` zeros the block; `m == 0` + * or `n == 0` is a no-op. + */ +SKAINET_API void skainet_fp32_matmul( + const float* a, int32_t a_offset, int32_t a_stride, + const float* b, int32_t b_offset, int32_t b_stride, + float* c, int32_t c_offset, int32_t c_stride, + int32_t m, int32_t n, int32_t k +); + #ifdef __cplusplus } #endif diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c new file mode 100644 index 00000000..ce88afd8 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c @@ -0,0 +1,61 @@ +#include "skainet_kernels.h" + +#include +#include + +/* + * Native row-major SGEMM matching the + * sk.ainet.backend.api.kernel.Fp32MatmulKernel SPI: + * + * C(m, n) = A(m, k) * B(k, n) + * + * Strides are in floats (not bytes); for a contiguous parent matrix + * `aStride == k`, `bStride == n`, `cStride == n`. Sub-block scenarios + * pass larger strides and the corresponding offsets. + * + * Iteration order is i-p-j (outer product into rows of C). The inner + * loop is `c[j] += a_ip * b[j]` over a contiguous run of `n` floats + * for both b's row and c's row — auto-vectorizes cleanly under + * -O3 -ffast-math into vfmadd231ps / fmla. + * + * Caller contract: + * - C is FULLY OVERWRITTEN in the m×n block (zero-then-accumulate). + * - k == 0 zeros the m×n block. + * - m == 0 || n == 0 is a no-op. + * - Negative m / n / k are caller errors; the Kotlin wrapper rejects + * them. The C kernel still treats negatives as no-op (via the + * `<=` loop bounds) defensively. + */ +SKAINET_API void skainet_fp32_matmul( + const float* SKAINET_RESTRICT a, int32_t a_offset, int32_t a_stride, + const float* SKAINET_RESTRICT b, int32_t b_offset, int32_t b_stride, + float* SKAINET_RESTRICT c, int32_t c_offset, int32_t c_stride, + int32_t m, int32_t n, int32_t k +) { + if (m <= 0 || n <= 0) return; + + /* Zero the output block. Required by the SPI contract for k == 0 + * AND prerequisite for the i-p-j accumulator pattern below. */ + for (int32_t i = 0; i < m; ++i) { + float* SKAINET_RESTRICT c_row = c + c_offset + (size_t) i * c_stride; + for (int32_t j = 0; j < n; ++j) { + c_row[j] = 0.0f; + } + } + if (k <= 0) return; + + /* Outer-product accumulator: streams two contiguous rows on the + * inner loop (b's row and c's row), broadcasts a single A scalar. + * The compiler emits vfmadd231ps with a vbroadcastss for a_ip. */ + for (int32_t i = 0; i < m; ++i) { + const float* SKAINET_RESTRICT a_row = a + a_offset + (size_t) i * a_stride; + float* SKAINET_RESTRICT c_row = c + c_offset + (size_t) i * c_stride; + for (int32_t p = 0; p < k; ++p) { + const float a_ip = a_row[p]; + const float* SKAINET_RESTRICT b_row = b + b_offset + (size_t) p * b_stride; + for (int32_t j = 0; j < n; ++j) { + c_row[j] += a_ip * b_row[j]; + } + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt new file mode 100644 index 00000000..509a00e9 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt @@ -0,0 +1,113 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +import java.lang.foreign.Linker +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.lang.invoke.MethodHandle +import sk.ainet.backend.api.kernel.Fp32MatmulKernel + +/** + * Native (FFM) implementation of [Fp32MatmulKernel]. + * + * Wraps the bundled C symbol + * + * void skainet_fp32_matmul( + * const float* a, int32_t a_offset, int32_t a_stride, + * const float* b, int32_t b_offset, int32_t b_stride, + * float* c, int32_t c_offset, int32_t c_stride, + * int32_t m, int32_t n, int32_t k); + * + * The C kernel is a tight i-p-j outer-product accumulator over rows + * of C; the inner `c[j] += a*b[j]` loop streams two contiguous arrays + * and auto-vectorizes into FMA under -O3 -ffast-math (vfmadd231ps on + * x86_64, fmla on AArch64). + * + * Numerical parity vs [PanamaVectorMatmulKernel] is asserted by + * [NativeFp32MatmulKernelParityTest] within FMA + reordered-reduction + * tolerance (the same `1e-5 * k` bar Panama uses against the scalar + * reference). + * + * PR 5 of the staged native-FFM rollout — wraps the rollout per the + * `native-ffm-plan` asciidoc. Single-threaded, no cache blocking; + * future work could add parallelChunks-style row blocking and B-tile + * packing, but the scalar C path already lands well within the SPI + * contract on host-arch CPUs. + */ +internal object NativeFp32MatmulKernel : Fp32MatmulKernel { + + fun isAvailable(): Boolean = handle != null + + 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) { + "NativeFp32MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + + val mh = handle + ?: error("NativeFp32MatmulKernel.matmul invoked while native library unavailable") + + // Sizes for the off-heap copies. Each of A, B, C uses the + // bytes the kernel actually reaches — for non-contiguous + // strides this can be larger than the matrix's element count + // because the strides skip past unused floats. Allocating to + // the full reach (offset + last-row reach) keeps the kernel + // pointer arithmetic simple and matches Kotlin's bounds. + val aReachFloats = if (m == 0 || k == 0) 0 else aOffset + (m - 1) * aStride + k + val bReachFloats = if (k == 0 || n == 0) 0 else bOffset + (k - 1) * bStride + n + val cReachFloats = outOffset + (m - 1) * outStride + n + + Arena.ofConfined().use { arena -> + val aBytes = aReachFloats.toLong() * java.lang.Float.BYTES + val bBytes = bReachFloats.toLong() * java.lang.Float.BYTES + val cBytes = cReachFloats.toLong() * java.lang.Float.BYTES + val align = ValueLayout.JAVA_FLOAT.byteAlignment() + + val aSeg: MemorySegment = if (aBytes > 0) arena.allocate(aBytes, align) else MemorySegment.NULL + val bSeg: MemorySegment = if (bBytes > 0) arena.allocate(bBytes, align) else MemorySegment.NULL + val cSeg: MemorySegment = arena.allocate(cBytes, align) + + if (aReachFloats > 0) { + MemorySegment.copy(a, 0, aSeg, ValueLayout.JAVA_FLOAT, 0L, aReachFloats) + } + if (bReachFloats > 0) { + MemorySegment.copy(b, 0, bSeg, ValueLayout.JAVA_FLOAT, 0L, bReachFloats) + } + + mh.invoke( + aSeg, aOffset, aStride, + bSeg, bOffset, bStride, + cSeg, outOffset, outStride, + m, n, k, + ) + + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, cReachFloats) + } + } + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_fp32_matmul").orElse(null) ?: return@lazy null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, // a + ValueLayout.JAVA_INT, // a_offset + ValueLayout.JAVA_INT, // a_stride + ValueLayout.ADDRESS, // b + ValueLayout.JAVA_INT, // b_offset + ValueLayout.JAVA_INT, // b_stride + ValueLayout.ADDRESS, // c + ValueLayout.JAVA_INT, // c_offset + ValueLayout.JAVA_INT, // c_stride + ValueLayout.JAVA_INT, // m + ValueLayout.JAVA_INT, // n + ValueLayout.JAVA_INT, // k + ) + runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull() + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt index adb9cabf..f1c3f118 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt @@ -26,8 +26,9 @@ import sk.ainet.backend.api.kernel.Q4KMemSegMatmulKernel * * Staged rollout cursor (see `native-ffm-plan` asciidoc): * - PR 2: real Q4_K matmul wired into the heap SPI. - * - PR 3 (this commit): MemSeg-input zero-copy sibling. - * - Later: native `matmulFp32`, `matmulQ6K`, `matmulQ8_0`. + * - PR 3: MemSeg-input zero-copy sibling. + * - PR 5 (this commit): native FP32 matmul wired into [matmulFp32]. + * - Later: native `matmulQ6K`, `matmulQ8_0` (need new SPI accessors). */ public object NativeKernelProvider : KernelProvider, MemSegKernelProvider { override val name: String = "native-ffm" @@ -35,7 +36,8 @@ public object NativeKernelProvider : KernelProvider, MemSegKernelProvider { override fun isAvailable(): Boolean = NativeQ4KMatmulKernel.isAvailable() - override fun matmulFp32(): Fp32MatmulKernel? = null + override fun matmulFp32(): Fp32MatmulKernel? = + if (NativeFp32MatmulKernel.isAvailable()) NativeFp32MatmulKernel else null override fun matmulQ4K(): Q4KMatmulKernel? = if (NativeQ4KMatmulKernel.isAvailable()) NativeQ4KMatmulKernel else null diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt index 94421c12..baa7a301 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt @@ -32,7 +32,7 @@ class NativeFfmPipelineTest { } @Test - fun `provider exposes Q4_K kernel when the native lib loads`() { + fun `provider exposes Q4_K and FP32 kernels when the native lib loads`() { assertEquals("native-ffm", NativeKernelProvider.name) assertEquals(100, NativeKernelProvider.priority) assertTrue( @@ -40,8 +40,8 @@ class NativeFfmPipelineTest { "Native kernel provider reports unavailable on this host — " + "bundled libskainet_kernels missing or skainet_q4k_matmul unresolved", ) - // FP32 matmul ships in a later PR; Q4_K is wired through PR 2. - assertEquals(null, NativeKernelProvider.matmulFp32()) + // PR 5 wires both Q4_K (PR 2) and FP32 (this PR) through the SPI. + assertNotNull(NativeKernelProvider.matmulFp32()) assertNotNull(NativeKernelProvider.matmulQ4K()) } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernelParityTest.kt new file mode 100644 index 00000000..12b663b1 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernelParityTest.kt @@ -0,0 +1,165 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Numerical parity tests for [NativeFp32MatmulKernel] against + * [PanamaVectorMatmulKernel]. Both kernels follow the same SGEMM + * contract; outputs must agree within FMA + reordered-reduction + * tolerance. + * + * Tolerance scales with the contraction dimension `k`: each summand + * carries up to `eps * |a| * |b|` rounding error, accumulated `k` + * times. The Panama-vs-Scalar bar of `1e-5 * k` swallows native + * `-ffast-math` reassociation comfortably for inputs clamped to + * `[-0.5, 0.5]`. + */ +class NativeFp32MatmulKernelParityTest { + + @BeforeTest + fun checkAvailable() { + assertTrue( + NativeFp32MatmulKernel.isAvailable(), + "Native FP32 kernel must be available — bundled libskainet_kernels missing or " + + "skainet_fp32_matmul symbol unresolved", + ) + } + + private fun assertParity( + m: Int, n: Int, k: Int, + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + outStride: Int, + tolScale: Float = 1e-5f, + ) { + val outPanama = FloatArray(m * outStride) + val outNative = FloatArray(m * outStride) + PanamaVectorMatmulKernel.matmul( + a, aOffset, aStride, + b, bOffset, bStride, + outPanama, 0, outStride, + m, n, k, + ) + NativeFp32MatmulKernel.matmul( + a, aOffset, aStride, + b, bOffset, bStride, + outNative, 0, outStride, + m, n, k, + ) + val tol = (tolScale * k.coerceAtLeast(1)).coerceAtLeast(tolScale) + assertEquals(outPanama.size, outNative.size, "length mismatch") + for (i in outPanama.indices) { + val diff = abs(outPanama[i] - outNative[i]) + assertTrue( + diff <= tol, + "mismatch at $i: panama=${outPanama[i]} native=${outNative[i]} diff=$diff tol=$tol", + ) + } + } + + @Test + fun small_2x3x4_contiguous_matches_panama() { + 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_panama() { + 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() { + 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_panama() { + // 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_panama() { + 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 llm_typical_4096_squared_matches_panama() { + // 256² is the smallest LLM-typical shape; full 4096² matmul + // takes ~hundreds of ms on either kernel and slows the test + // suite. The smaller shape still exercises k-tile loops. + val rng = Random(99) + val m = 256; val n = 256; val k = 256 + 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, tolScale = 5e-5f) + } + + @Test + fun zero_m_or_n_no_op() { + val out = FloatArray(5) { 7f } + NativeFp32MatmulKernel.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 } + NativeFp32MatmulKernel.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 { + NativeFp32MatmulKernel.matmul( + FloatArray(0), 0, 0, + FloatArray(0), 0, 0, + FloatArray(0), 0, 0, + m = -1, n = 1, k = 1, + ) + } + } + + @Test + fun provider_returns_native_fp32_when_available() { + val kernel = NativeKernelProvider.matmulFp32() + assertTrue(kernel === NativeFp32MatmulKernel, "Provider must hand out the native FP32 kernel") + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt index a83b5c13..0b66db5c 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt @@ -69,6 +69,52 @@ class Q4KMatmulMicrobenchTest { return med } + @Test + fun bench_fp32_native_vs_panama() { + if (System.getProperty("skainet.runBench") != "true") { + println("Q4KMatmulMicrobenchTest.fp32 skipped — pass -Dskainet.runBench=true to enable.") + return + } + assertTrue(NativeFp32MatmulKernel.isAvailable(), "Native FP32 kernel must be available for the bench") + + val shapes = listOf( + Triple(256, 256, 256), + Triple(512, 512, 512), + Triple(1024, 1024, 1024), + ) + + println() + println("FP32 SGEMM microbench — Native (FFM, scalar C i-p-j outer-product, -O3 -ffast-math)") + println(" vs Panama Vector (tile-blocked, B-pack, parallelChunks)") + println("Host: ${System.getProperty("os.name")} ${System.getProperty("os.arch")} | JDK ${System.getProperty("java.version")}") + println() + + for ((m, n, k) in shapes) { + val rng = Random(m + n + k) + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val b = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val outNative = FloatArray(m * n) + val outPanama = FloatArray(m * n) + + println("[m=$m n=$n k=$k]") + val nativeNs = benchOne("native", warmup = 5, samples = 9) { + NativeFp32MatmulKernel.matmul(a, 0, k, b, 0, n, outNative, 0, n, m, n, k) + } + val panamaNs = benchOne("panama", warmup = 5, samples = 9) { + PanamaVectorMatmulKernel.matmul(a, 0, k, b, 0, n, outPanama, 0, n, m, n, k) + } + val ratio = panamaNs.toDouble() / nativeNs.toDouble() + println( + " ratio: native is %.2fx panama (%.1f%% %s)".format( + ratio, + abs((ratio - 1.0) * 100.0), + if (ratio >= 1.0) "faster" else "slower", + ), + ) + println() + } + } + @Test fun bench_native_vs_panama_at_llm_shapes() { if (System.getProperty("skainet.runBench") != "true") {