Skip to content

Commit 1b6f7f6

Browse files
michalharakalclaude
michalharakal
and
claude
committed
feat(backend-native-cpu): K/N KernelProvider over the cinterop kernels
The K/N analogue of the JVM NativeKernelProvider (FFM): a KernelProvider (priority 100) exposing the cinterop-backed Q5_K/Q4_K/Q8_0/Q4_0 matmul kernels, plus installNativeKernels() to register it in KernelRegistry — the path the eager runtime's DefaultCpuOps.chooseQuantizedMatmulHeap uses to resolve a kernel. K/N has no ServiceLoader, so registration is an explicit call by the consumer (scalar fallback for Q6_K etc. is registered separately from skainet-backend-cpu). Verified on linuxX64: NativeKnKernelProviderTest — installNativeKernels makes native-cinterop the best-available provider, its Q5_K kernel is the registry-resolved kernel, and it matches the scalar reference (6 K/N tests green total). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 5814200 commit 1b6f7f6

2 files changed

Lines changed: 198 additions & 0 deletions

File tree

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
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.Fp32MatmulKernel
8+
import sk.ainet.backend.api.kernel.KernelProvider
9+
import sk.ainet.backend.api.kernel.KernelRegistry
10+
import sk.ainet.backend.api.kernel.Q4KMatmulKernel
11+
import sk.ainet.backend.api.kernel.Q4_0MatmulKernel
12+
import sk.ainet.backend.api.kernel.Q5KMatmulKernel
13+
import sk.ainet.backend.api.kernel.Q8_0MatmulKernel
14+
import sk.ainet.kernels.cinterop.skainet_q4_0_matmul
15+
import sk.ainet.kernels.cinterop.skainet_q4k_matmul
16+
import sk.ainet.kernels.cinterop.skainet_q8_0_matmul
17+
18+
/**
19+
* Kotlin/Native [KernelProvider] backed by the hand-written C kernels via
20+
* cinterop (static archive `libskainet_kernels.a`) — the K/N analogue of the
21+
* JVM `NativeKernelProvider` (FFM). Priority 100, above the commonMain scalar
22+
* reference (0). On `linuxArm64` the linked archive carries the NEON paths.
23+
*
24+
* **Registration is manual on K/N** (no `ServiceLoader`): a consumer calls
25+
* [installNativeKernels] once at startup. [Q5KMatmulKernel] (the FunctionGemma
26+
* Q5_K_M hot path) plus Q4_K / Q8_0 / Q4_0 are wired; the rest cascade to the
27+
* scalar provider.
28+
*/
29+
@OptIn(ExperimentalForeignApi::class)
30+
public object NativeKnKernelProvider : KernelProvider {
31+
override val name: String = "native-cinterop"
32+
override val priority: Int = 100
33+
34+
// Statically linked — the symbols are always present once the binary links
35+
// libskainet_kernels.a, so the provider is unconditionally available.
36+
override fun isAvailable(): Boolean = true
37+
38+
// Abstract on KernelProvider (no default) — no native FP32 SGEMM wrapper yet.
39+
override fun matmulFp32(): Fp32MatmulKernel? = null
40+
41+
override fun matmulQ5K(): Q5KMatmulKernel = NativeKnQ5KMatmulKernel
42+
override fun matmulQ4K(): Q4KMatmulKernel = NativeKnQ4KMatmulKernel
43+
override fun matmulQ8_0(): Q8_0MatmulKernel = NativeKnQ8_0MatmulKernel
44+
override fun matmulQ4_0(): Q4_0MatmulKernel = NativeKnQ4_0MatmulKernel
45+
}
46+
47+
/**
48+
* Register [NativeKnKernelProvider] in the process-wide [KernelRegistry]. Idempotent
49+
* (re-registering the same instance is a no-op). Call once at startup before any
50+
* `ops.matmul` on quantized weights.
51+
*
52+
* For quant types without a C kernel (e.g. Q6_K) also register the commonMain
53+
* `ScalarKernelProvider` (from `skainet-backend-cpu`) as the fallback — it lives
54+
* in a different module, so the consumer wires it:
55+
* `KernelRegistry.register(ScalarKernelProvider)`.
56+
*/
57+
public fun installNativeKernels() {
58+
KernelRegistry.register(NativeKnKernelProvider)
59+
}
60+
61+
@OptIn(ExperimentalForeignApi::class)
62+
public object NativeKnQ4KMatmulKernel : Q4KMatmulKernel {
63+
private const val BLOCK_SIZE = 256
64+
override fun matmul(
65+
input: FloatArray, inputOffset: Int,
66+
weight: ByteArray, weightByteOffset: Int,
67+
inputDim: Int, outputDim: Int,
68+
output: FloatArray, outputOffset: Int,
69+
) {
70+
require(inputDim % BLOCK_SIZE == 0) {
71+
"NativeKnQ4KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
72+
}
73+
if (outputDim == 0 || inputDim == 0) return
74+
input.usePinned { i -> weight.usePinned { w -> output.usePinned { o ->
75+
skainet_q4k_matmul(
76+
i.addressOf(0), inputOffset,
77+
w.addressOf(0).reinterpret(), weightByteOffset,
78+
inputDim, outputDim,
79+
o.addressOf(0), outputOffset,
80+
)
81+
} } }
82+
}
83+
}
84+
85+
@OptIn(ExperimentalForeignApi::class)
86+
public object NativeKnQ8_0MatmulKernel : Q8_0MatmulKernel {
87+
private const val BLOCK_SIZE = 32
88+
override fun matmul(
89+
input: FloatArray, inputOffset: Int,
90+
weight: ByteArray, weightByteOffset: Int,
91+
inputDim: Int, outputDim: Int,
92+
output: FloatArray, outputOffset: Int,
93+
) {
94+
require(inputDim % BLOCK_SIZE == 0) {
95+
"NativeKnQ8_0MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
96+
}
97+
if (outputDim == 0 || inputDim == 0) return
98+
input.usePinned { i -> weight.usePinned { w -> output.usePinned { o ->
99+
skainet_q8_0_matmul(
100+
i.addressOf(0), inputOffset,
101+
w.addressOf(0).reinterpret(), weightByteOffset,
102+
inputDim, outputDim,
103+
o.addressOf(0), outputOffset,
104+
)
105+
} } }
106+
}
107+
}
108+
109+
@OptIn(ExperimentalForeignApi::class)
110+
public object NativeKnQ4_0MatmulKernel : Q4_0MatmulKernel {
111+
private const val BLOCK_SIZE = 32
112+
override fun matmul(
113+
input: FloatArray, inputOffset: Int,
114+
weight: ByteArray, weightByteOffset: Int,
115+
inputDim: Int, outputDim: Int,
116+
output: FloatArray, outputOffset: Int,
117+
) {
118+
require(inputDim % BLOCK_SIZE == 0) {
119+
"NativeKnQ4_0MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
120+
}
121+
if (outputDim == 0 || inputDim == 0) return
122+
input.usePinned { i -> weight.usePinned { w -> output.usePinned { o ->
123+
skainet_q4_0_matmul(
124+
i.addressOf(0), inputOffset,
125+
w.addressOf(0).reinterpret(), weightByteOffset,
126+
inputDim, outputDim,
127+
o.addressOf(0), outputOffset,
128+
)
129+
} } }
130+
}
131+
}

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package sk.ainet.exec.kernel
2+
3+
import kotlin.math.abs
4+
import kotlin.random.Random
5+
import kotlin.test.AfterTest
6+
import kotlin.test.BeforeTest
7+
import kotlin.test.Test
8+
import kotlin.test.assertEquals
9+
import kotlin.test.assertNotNull
10+
import kotlin.test.assertSame
11+
import kotlin.test.assertTrue
12+
import sk.ainet.backend.api.kernel.KernelRegistry
13+
14+
/**
15+
* Verifies the K/N kernel provider integrates with [KernelRegistry] the way the
16+
* eager runtime's `DefaultCpuOps.chooseQuantizedMatmulHeap` resolves kernels:
17+
* after [installNativeKernels], the highest-priority available provider is the
18+
* cinterop one, and its Q5_K kernel is the registry-resolved kernel that runs.
19+
*/
20+
class NativeKnKernelProviderTest {
21+
22+
@BeforeTest
23+
fun clean() = KernelRegistry.clearForTesting()
24+
25+
@AfterTest
26+
fun reset() = KernelRegistry.clearForTesting()
27+
28+
@Test
29+
fun installs_and_resolves_native_quant_kernels() {
30+
installNativeKernels()
31+
32+
// Priority 100 cinterop beats the scalar (0) fallback.
33+
assertEquals("native-cinterop", KernelRegistry.bestAvailable()?.name)
34+
35+
val provider = KernelRegistry.providers().firstOrNull { it.isAvailable() && it.matmulQ5K() != null }
36+
assertNotNull(provider, "no available provider carries a Q5_K kernel")
37+
assertSame(NativeKnQ5KMatmulKernel, provider.matmulQ5K())
38+
assertTrue(provider.supports("matmul", listOf("Float32", "Q5_K")))
39+
}
40+
41+
@Test
42+
fun registry_resolved_q5k_kernel_is_correct() {
43+
installNativeKernels()
44+
val kernel = KernelRegistry.bestAvailable()!!.matmulQ5K()!!
45+
46+
val inputDim = 1024
47+
val outputDim = 64
48+
val numBlocks = (inputDim / 256) * outputDim
49+
val packed = ByteArray(numBlocks * 176).also { Random(5).nextBytes(it) }
50+
for (b in 0 until numBlocks) {
51+
val base = b * 176
52+
packed[base] = 0x00; packed[base + 1] = 0x3C // d = 1.0f16
53+
packed[base + 2] = 0x00; packed[base + 3] = 0x3C // dMin = 1.0f16
54+
}
55+
val input = FloatArray(inputDim) { Random(it + 1).nextFloat() - 0.5f }
56+
57+
val ref = FloatArray(outputDim)
58+
ScalarQ5_KMatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, ref, 0)
59+
val got = FloatArray(outputDim)
60+
kernel.matmul(input, 0, packed, 0, inputDim, outputDim, got, 0)
61+
62+
for (o in 0 until outputDim) {
63+
val diff = abs(ref[o] - got[o])
64+
assertTrue(diff <= 5e-2f || diff / (abs(ref[o]) + 1e-9f) < 1e-4f, "row $o: ref=${ref[o]} got=${got[o]}")
65+
}
66+
}
67+
}

0 commit comments

Comments
 (0)