Skip to content

Commit 5814200

Browse files
michalharakalclaude
michalharakal
and
claude
committed
feat(backend-native-cpu): Kotlin/Native cinterop to the C/NEON kernels
The hand-written matmul kernels were JVM-only (consumed via FFM), but the SL2610 board binary is Kotlin/Native — it can't use the FFM wrapper. Add a K/N consumption path via cinterop so the board gets the same C (and, on aarch64, NEON) kernels. - CMake builds a STATIC archive (skainet_kernels_static -> libskainet_kernels.a) alongside the SHARED lib; same sources + flags (incl. the aarch64 NEON march). - cinterop .def (skainet_kernels.h -> sk.ainet.kernels.cinterop). - linuxX64 target on the (previously jvm-only) module, linking the static archive into K/N binaries; link tasks depend on the CMake build. - NativeKnQ5KMatmulKernel (linuxX64Main): calls skainet_q5k_matmul via cinterop with pinned arrays (zero-copy). POC verified on the host (linuxX64): NativeKnQ5KMatmulKernelParityTest — the cinterop kernel matches the commonMain ScalarQ5_KMatmulKernel across 4 shapes (tests=4, failures=0). JVM/FFM path unchanged (jvmTest green). linuxArm64 board target + NEON runtime check are the remaining step. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 0101041 commit 5814200

5 files changed

Lines changed: 220 additions & 39 deletions

File tree

‎skainet-backends/skainet-backend-native-cpu/build.gradle.kts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,30 @@ plugins {
33
alias(libs.plugins.vanniktech.mavenPublish)
44
}
55

6+
// Paths shared by the K/N cinterop (kotlin block) and the CMake tasks below.
7+
val nativeIncludeDir: String = layout.projectDirectory.dir("native/include").asFile.absolutePath
8+
val staticArchivePath: String =
9+
layout.buildDirectory.file("native/cmake-build/libskainet_kernels.a").get().asFile.absolutePath
10+
611
kotlin {
712
explicitApi()
813
jvm()
914

15+
// Kotlin/Native: POC on the host (linuxX64); linuxArm64 is the board target.
16+
// Exposes the hand-written C/NEON kernels to K/N via cinterop to the static
17+
// archive libskainet_kernels.a (CMake `skainet_kernels_static`). This is the
18+
// board-consumption path — the JVM consumes the same kernels via FFM instead.
19+
linuxX64 {
20+
compilations.getByName("main").cinterops.create("skainetKernels") {
21+
defFile(project.file("src/nativeInterop/cinterop/skainet_kernels.def"))
22+
includeDirs(nativeIncludeDir)
23+
}
24+
binaries.all {
25+
// Link the static C archive into every linuxX64 binary (incl. tests).
26+
linkerOpts(staticArchivePath)
27+
}
28+
}
29+
1030
sourceSets {
1131
val jvmMain by getting {
1232
dependencies {
@@ -24,6 +44,18 @@ kotlin {
2444
implementation(libs.kotlinx.coroutines)
2545
}
2646
}
47+
val linuxX64Main by getting {
48+
dependencies {
49+
implementation(project(":skainet-backends:skainet-backend-api"))
50+
}
51+
}
52+
val linuxX64Test by getting {
53+
dependencies {
54+
implementation(libs.kotlin.test)
55+
// ScalarQ5_KMatmulKernel reference for the cinterop parity test.
56+
implementation(project(":skainet-backends:skainet-backend-cpu"))
57+
}
58+
}
2759
}
2860
}
2961

@@ -89,6 +121,12 @@ val buildNativeKernels by tasks.registering(Exec::class) {
89121
)
90122
}
91123

124+
// The linuxX64 (K/N) binaries link libskainet_kernels.a (built by CMake into
125+
// cmakeBuildPath), so the static archive must exist before the K/N link step.
126+
tasks.matching { it.name.startsWith("link") && it.name.endsWith("LinuxX64") }.configureEach {
127+
dependsOn(buildNativeKernels)
128+
}
129+
92130
val packageNativeKernels by tasks.registering(Copy::class) {
93131
group = "build"
94132
description = "Stage the built native kernels library into JVM resources."

‎skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt‎

Lines changed: 48 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ if(NOT CMAKE_BUILD_TYPE)
99
set(CMAKE_BUILD_TYPE Release)
1010
endif()
1111

12-
add_library(skainet_kernels SHARED
12+
set(SKAINET_KERNEL_SOURCES
1313
src/skainet_smoke.c
1414
src/q4k_matmul.c
1515
src/q5k_matmul.c
@@ -19,45 +19,54 @@ add_library(skainet_kernels SHARED
1919
src/q4_0_matmul.c
2020
)
2121

22-
target_include_directories(skainet_kernels PUBLIC
23-
${CMAKE_CURRENT_SOURCE_DIR}/include
24-
)
22+
# SHARED: consumed by the JVM via java.lang.foreign (FFM), bundled as a JAR
23+
# resource (libskainet_kernels.{so,dylib,dll}).
24+
# STATIC: consumed by Kotlin/Native via cinterop, linked into the K/N binary
25+
# (libskainet_kernels.a). Same sources + flags; both keep -fPIC.
26+
add_library(skainet_kernels SHARED ${SKAINET_KERNEL_SOURCES})
27+
add_library(skainet_kernels_static STATIC ${SKAINET_KERNEL_SOURCES})
28+
set_target_properties(skainet_kernels_static PROPERTIES OUTPUT_NAME skainet_kernels)
2529

26-
# Strip the "lib" prefix on Windows so the artifact name is consistent
27-
# with the resource-bundle path skainet_kernels.{dll,so,dylib}.
28-
if(WIN32)
29-
set_target_properties(skainet_kernels PROPERTIES PREFIX "")
30-
endif()
30+
set(SKAINET_KERNEL_TARGETS skainet_kernels skainet_kernels_static)
3131

32-
# Per-compiler tuning. The Q4_K kernel hot loop is straight-line FP
33-
# arithmetic that auto-vectorizes cleanly under aggressive optimization
34-
# (AVX2 on x86_64, NEON on ARM64). Visibility is also handled here on
35-
# ELF / Mach-O; on Windows the SKAINET_API macro adds dllexport so we
36-
# don't need /VISIBILITY flags.
37-
if(CMAKE_C_COMPILER_ID MATCHES "Clang|GNU")
38-
target_compile_options(skainet_kernels PRIVATE
39-
-fvisibility=hidden
40-
-Wall -Wextra
41-
-O3
42-
-ffast-math
43-
-funroll-loops
44-
)
45-
# AArch64: enable the NEON paths guarded by __ARM_NEON in skainet_simd.h.
46-
# The SL2610 is Cortex-A55-class (ARMv8.2-A) — it HAS NEON + fp16 + dotprod
47-
# (asimddp) but NOT i8mm (that is ARMv8.6). Do NOT add +i8mm: it would
48-
# SIGILL on the board. CONFIRM on-device first: `grep Features /proc/cpuinfo`
49-
# must list `asimddp` and `fphp`/`asimdhp`. Explicit -march (not -mcpu=native)
50-
# because this is a cross-compile from an x86 host.
51-
if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
52-
target_compile_options(skainet_kernels PRIVATE
53-
-march=armv8.2-a+fp16+dotprod
32+
foreach(tgt IN LISTS SKAINET_KERNEL_TARGETS)
33+
target_include_directories(${tgt} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
34+
35+
# Strip the "lib" prefix on Windows so the shared artifact name matches
36+
# the resource-bundle path skainet_kernels.{dll,so,dylib}.
37+
if(WIN32)
38+
set_target_properties(${tgt} PROPERTIES PREFIX "")
39+
endif()
40+
41+
# Per-compiler tuning. The matmul hot loops are straight-line FP arithmetic
42+
# that auto-vectorizes under aggressive optimization (AVX2 on x86_64, NEON
43+
# on ARM64). Visibility is handled here on ELF / Mach-O; on Windows the
44+
# SKAINET_API macro adds dllexport.
45+
if(CMAKE_C_COMPILER_ID MATCHES "Clang|GNU")
46+
target_compile_options(${tgt} PRIVATE
47+
-fvisibility=hidden
48+
-Wall -Wextra
49+
-O3
50+
-ffast-math
51+
-funroll-loops
52+
)
53+
# AArch64: enable the NEON paths guarded by __ARM_NEON in skainet_simd.h.
54+
# The SL2610 is Cortex-A55-class (ARMv8.2-A) — it HAS NEON + fp16 +
55+
# dotprod (asimddp) but NOT i8mm (that is ARMv8.6). Do NOT add +i8mm: it
56+
# would SIGILL on the board. CONFIRM on-device first: `grep Features
57+
# /proc/cpuinfo` must list `asimddp` and `fphp`/`asimdhp`. Explicit
58+
# -march (not -mcpu=native) because this is a cross-compile from x86.
59+
if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
60+
target_compile_options(${tgt} PRIVATE
61+
-march=armv8.2-a+fp16+dotprod
62+
)
63+
endif()
64+
set_target_properties(${tgt} PROPERTIES C_VISIBILITY_PRESET hidden)
65+
elseif(CMAKE_C_COMPILER_ID MATCHES "MSVC")
66+
target_compile_options(${tgt} PRIVATE
67+
/O2
68+
/fp:fast
69+
/W3
5470
)
5571
endif()
56-
set_target_properties(skainet_kernels PROPERTIES C_VISIBILITY_PRESET hidden)
57-
elseif(CMAKE_C_COMPILER_ID MATCHES "MSVC")
58-
target_compile_options(skainet_kernels PRIVATE
59-
/O2
60-
/fp:fast
61-
/W3
62-
)
63-
endif()
72+
endforeach()

‎skainet-backends/skainet-backend-native-cpu/src/linuxX64Main/kotlin/sk/ainet/exec/kernel/NativeKnQ5KMatmulKernel.kt‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package sk.ainet.exec.kernel
2+
3+
import kotlinx.cinterop.ExperimentalForeignApi
4+
import kotlinx.cinterop.addressOf
5+
import kotlinx.cinterop.reinterpret
6+
import kotlinx.cinterop.usePinned
7+
import sk.ainet.backend.api.kernel.Q5KMatmulKernel
8+
import sk.ainet.kernels.cinterop.skainet_q5k_matmul
9+
10+
/**
11+
* Kotlin/Native implementation of [Q5KMatmulKernel] that calls the hand-written
12+
* C kernel `skainet_q5k_matmul` (the same `q5k_matmul.c` the JVM consumes via
13+
* FFM) through cinterop, linked from the static archive `libskainet_kernels.a`.
14+
*
15+
* This is the board-consumption path: the SL2610 binary is Kotlin/Native, so it
16+
* cannot use the JVM-FFM wrapper. The arrays are pinned and their base pointers
17+
* passed to C; the C side reads `input + input_offset` etc., so no copy is made.
18+
*
19+
* On `linuxArm64` the linked archive carries the NEON paths
20+
* (`-march=armv8.2-a+fp16+dotprod`); on `linuxX64` (this POC host) it's the
21+
* scalar/auto-vectorized build. Correctness is identical across both.
22+
*/
23+
@OptIn(ExperimentalForeignApi::class)
24+
public object NativeKnQ5KMatmulKernel : Q5KMatmulKernel {
25+
26+
private const val BLOCK_SIZE = 256
27+
28+
override fun matmul(
29+
input: FloatArray, inputOffset: Int,
30+
weight: ByteArray, weightByteOffset: Int,
31+
inputDim: Int, outputDim: Int,
32+
output: FloatArray, outputOffset: Int,
33+
) {
34+
require(inputDim % BLOCK_SIZE == 0) {
35+
"NativeKnQ5KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
36+
}
37+
if (outputDim == 0 || inputDim == 0) return
38+
39+
input.usePinned { inPin ->
40+
weight.usePinned { wPin ->
41+
output.usePinned { outPin ->
42+
skainet_q5k_matmul(
43+
inPin.addressOf(0),
44+
inputOffset,
45+
wPin.addressOf(0).reinterpret(),
46+
weightByteOffset,
47+
inputDim,
48+
outputDim,
49+
outPin.addressOf(0),
50+
outputOffset,
51+
)
52+
}
53+
}
54+
}
55+
}
56+
}

‎skainet-backends/skainet-backend-native-cpu/src/linuxX64Test/kotlin/sk/ainet/exec/kernel/NativeKnQ5KMatmulKernelParityTest.kt‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package sk.ainet.exec.kernel
2+
3+
import kotlin.math.abs
4+
import kotlin.random.Random
5+
import kotlin.test.Test
6+
import kotlin.test.assertTrue
7+
8+
/**
9+
* Proves the Kotlin/Native cinterop path: [NativeKnQ5KMatmulKernel] (calling the
10+
* C `skainet_q5k_matmul` via cinterop, linked from libskainet_kernels.a) must
11+
* agree with the commonMain [ScalarQ5_KMatmulKernel] reference within FMA +
12+
* `-ffast-math` reassociation tolerance.
13+
*
14+
* This is the host (linuxX64) de-risking of the board (linuxArm64) consumption:
15+
* the cinterop mechanism + kernel correctness are verified here; only the NEON
16+
* codegen differs on aarch64 (board-verify-pending).
17+
*/
18+
class NativeKnQ5KMatmulKernelParityTest {
19+
20+
private val blockSize = 256
21+
private val bytesPerBlock = 176
22+
23+
private fun randomQ5KBytes(numBlocks: Int, seed: Int): ByteArray {
24+
val rng = Random(seed)
25+
val bytes = ByteArray(numBlocks * bytesPerBlock)
26+
rng.nextBytes(bytes)
27+
for (block in 0 until numBlocks) {
28+
val base = block * bytesPerBlock
29+
// 0x3C00 == 1.0f16 for d and dMin so dequant stays finite.
30+
bytes[base + 0] = 0x00.toByte()
31+
bytes[base + 1] = 0x3C.toByte()
32+
bytes[base + 2] = 0x00.toByte()
33+
bytes[base + 3] = 0x3C.toByte()
34+
}
35+
return bytes
36+
}
37+
38+
private fun assertParity(inputDim: Int, outputDim: Int, seed: Int, tol: Float) {
39+
val numBlocks = (inputDim / blockSize) * outputDim
40+
val packed = randomQ5KBytes(numBlocks, seed)
41+
val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f }
42+
43+
val refOut = FloatArray(outputDim)
44+
ScalarQ5_KMatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, refOut, 0)
45+
46+
val knOut = FloatArray(outputDim)
47+
NativeKnQ5KMatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, knOut, 0)
48+
49+
for (o in 0 until outputDim) {
50+
val diff = abs(refOut[o] - knOut[o])
51+
val rel = diff / (abs(refOut[o]) + 1e-9f)
52+
assertTrue(
53+
diff <= tol || rel < 1e-4f,
54+
"row $o diverged: scalar=${refOut[o]} cinterop=${knOut[o]} diff=$diff rel=$rel tol=$tol",
55+
)
56+
}
57+
}
58+
59+
@Test
60+
fun single_block_single_row() = assertParity(256, 1, 42, 1e-2f)
61+
62+
@Test
63+
fun single_block_multi_row() = assertParity(256, 16, 7, 1e-2f)
64+
65+
@Test
66+
fun multi_block_multi_row() = assertParity(1024, 64, 123, 5e-2f)
67+
68+
@Test
69+
fun llm_typical_shape() = assertParity(4096, 64, 999, 5e-1f)
70+
}

‎skainet-backends/skainet-backend-native-cpu/src/nativeInterop/cinterop/skainet_kernels.def‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
# Kotlin/Native cinterop binding for the hand-written C matmul kernels
2+
# (skainet_kernels.h). Generates Kotlin bindings for skainet_q5k_matmul etc.;
3+
# the static archive libskainet_kernels.a (built by CMake) is linked into the
4+
# consuming K/N binary via linkerOpts (see build.gradle.kts). includeDirs for
5+
# the header are supplied from the Gradle cinterop block.
6+
headers = skainet_kernels.h
7+
headerFilter = skainet_kernels.h
8+
package = sk.ainet.kernels.cinterop

0 commit comments

Comments
 (0)