feat(backends): JNI + Kotlin/Native bridges for the ternary f32 kernel (#1139)

skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c

@@ -169,3 +169,25 @@ Java_sk_ainet_exec_kernel_jni_JniKernels_bitnetGemvTq20(
     if (w) (*env)->ReleasePrimitiveArrayCritical(env, weight, w, JNI_ABORT);
     if (act) (*env)->ReleasePrimitiveArrayCritical(env, activation, act, JNI_ABORT);
 }
+
+/*
+ * ternary_f32_gemv (#1139): exact FP32 activations against the sequential
+ * BITNET_B1_58 payload — the vendored NeoGPU LUT kernel behind
+ * skainet_ternary_f32_gemv. Float input × byte weight × float output, so the
+ * shared body fits. The kernel threads internally (pthreads) once
+ * outputDim >= 512; the critical-section pins are held for the call's
+ * duration either way, same as every other matmul here.
+ */
+JNIEXPORT void JNICALL
+Java_sk_ainet_exec_kernel_jni_JniKernels_ternaryF32Gemv(
+    JNIEnv* env, jobject thiz,
+    jfloatArray input, jint inputOffset,
+    jbyteArray weight, jint weightByteOffset,
+    jint inputDim, jint outputDim,
+    jfloatArray output, jint outputOffset
+) {
+    (void) thiz;
+    SKAINET_JNI_MATMUL_BODY(
+        skainet_ternary_f32_gemv(in, inputOffset, (const uint8_t*) w, weightByteOffset,
+                                 inputDim, outputDim, out, outputOffset))
+}

skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt

@@ -255,6 +255,39 @@ class JniKernelParityTest {
         reference = Q6_KQ8ActivationReferenceKernel::matmul, jni = JniKernels::q6kMatmul,
     )
 
+    /**
+     * ternary_f32_gemv (#1139): the vendored NeoGPU LUT kernel — exact math,
+     * so parity vs the local decode reference is tight on every tier
+     * (baseline-NEON arm64, v8.2 arm64, x86_64 emulator scalar). The
+     * 1024-row case crosses the kernel's internal pthread threshold (512).
+     */
+    private fun ternaryDecode(b: Byte, lane: Int): Float =
+        (((b.toInt() and 0xFF) shr (lane * 2)) and 3).toFloat() - 1f
+
+    private fun assertTernaryF32Parity(inputDim: Int, outputDim: Int, seed: Int) {
+        val rng = Random(seed)
+        val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f }
+        val rowBytes = inputDim / 4
+        val weight = ByteArray(outputDim * rowBytes).also { rng.nextBytes(it) }
+        val out = FloatArray(outputDim)
+        JniKernels.ternaryF32Gemv(input, 0, weight, 0, inputDim, outputDim, out, 0)
+        for (o in 0 until outputDim) {
+            var want = 0.0
+            for (bi in 0 until rowBytes) {
+                val b = weight[o * rowBytes + bi]
+                for (lane in 0 until 4) want += ternaryDecode(b, lane) * input[bi * 4 + lane]
+            }
+            val diff = abs(want.toFloat() - out[o])
+            assertTrue("[$o]: reference=$want jni=${out[o]} diff=$diff", diff <= 1e-3f)
+        }
+    }
+
+    @Test
+    fun ternary_f32_parity() = assertTernaryF32Parity(inputDim = 2560, outputDim = 64, seed = 21)
+
+    @Test
+    fun ternary_f32_parity_threaded_regime() = assertTernaryF32Parity(inputDim = 256, outputDim = 1024, seed = 22)
+
     @Test
     fun smoke_roundtrip() {
         val input = floatArrayOf(1f, 2f, 3f, 4f)

skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt

@@ -106,6 +106,19 @@ public object JniKernels {
         output: FloatArray, outputOffset: Int,
     )
 
+    /**
+     * `ternary_f32_gemv` (#1139): exact FP32 activations against the sequential
+     * BITNET_B1_58 payload — the vendored NeoGPU LUT kernel (#1137). No scale
+     * is applied; the caller owns the per-tensor scale. Needs only baseline
+     * NEON, so the BASELINE library carries the full SIMD path.
+     */
+    public external fun ternaryF32Gemv(
+        input: FloatArray, inputOffset: Int,
+        weight: ByteArray, weightByteOffset: Int,
+        inputDim: Int, outputDim: Int,
+        output: FloatArray, outputOffset: Int,
+    )
+
     public external fun q40Matmul(
         input: FloatArray, inputOffset: Int,
         weight: ByteArray, weightByteOffset: Int,

skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniTernaryF32Gemv.kt

@@ -0,0 +1,48 @@
+package sk.ainet.exec.kernel.jni
+
+import sk.ainet.backend.api.kernel.TernaryF32GemvNative
+import sk.ainet.backend.api.kernel.TernaryF32KernelPack
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * The vendored NeoGPU LUT kernel as a [TernaryF32GemvNative] (#1139) — the Android/JNI face of
+ * `skainet_ternary_f32_gemv`, sharing its C with the FFM and cinterop consumers.
+ *
+ * Unlike [JniBitNetGemv] there is no capability split: the LUT kernel needs only baseline NEON
+ * (architecturally guaranteed on AArch64), so the BASELINE `libskainet_jni.so` — the one every
+ * arm64 device can load, Cortex-A72/Pi-class included — carries the full SIMD path. Whichever
+ * variant the loader picked, the kernel is the same.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+public object JniTernaryF32Gemv : TernaryF32GemvNative {
+
+    override val name: String get() = if (JniKernels.isLoaded) "neon" else "unloaded"
+
+    override fun gemvPacked(
+        activation: FloatArray,
+        activationOffset: Int,
+        weight: ByteArray,
+        weightByteOffset: Int,
+        inputDim: Int,
+        outputDim: Int,
+        out: FloatArray,
+        outOffset: Int,
+    ) {
+        JniKernels.ternaryF32Gemv(
+            activation, activationOffset,
+            weight, weightByteOffset,
+            inputDim, outputDim,
+            out, outOffset,
+        )
+    }
+
+    /**
+     * Install this kernel into the dispatcher, or leave the int8-requantize path serving and say
+     * so. Removing the AAR is a supported configuration — a notice through [warn], never a crash.
+     *
+     * @return the name of the kernel serving the exact FP32×b1.58 key, or
+     *   [TernaryF32KernelPack.NOT_INSTALLED]
+     */
+    public fun install(warn: (String) -> Unit = { println("[skainet] $it") }): String =
+        TernaryF32KernelPack.install(if (JniKernels.isLoaded) this else null, warn = warn)
+}

skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32Gemv.kt

@@ -0,0 +1,75 @@
+package sk.ainet.exec.kernel
+
+import kotlinx.cinterop.ExperimentalForeignApi
+import kotlinx.cinterop.addressOf
+import kotlinx.cinterop.reinterpret
+import kotlinx.cinterop.usePinned
+import sk.ainet.backend.api.kernel.TernaryF32GemvNative
+import sk.ainet.backend.api.kernel.TernaryF32KernelPack
+import sk.ainet.kernels.cinterop.skainet_ternary_f32_gemv
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * Kotlin/Native face of the vendored NeoGPU LUT kernel (#1139): calls
+ * `skainet_ternary_f32_gemv` through cinterop, linked from the static archive
+ * `libskainet_kernels.a` — the same C the JVM consumes via FFM and Android via
+ * JNI.
+ *
+ * This is the board-consumption path: a linuxArm64 binary on a Pi-4/Cortex-A72
+ * links the archive whose vendored file is pinned to `-march=armv8-a`, so the
+ * NEON LUT path runs on dotprod-less cores — the kernel's whole reason to
+ * exist. The arrays are pinned and base pointers passed; the C side applies
+ * the offsets, no copy is made.
+ */
+@OptIn(ExperimentalForeignApi::class, ExperimentalMemoryApi::class)
+public object NativeKnTernaryF32Gemv : TernaryF32GemvNative {
+
+    override val name: String get() = "cinterop"
+
+    override fun gemvPacked(
+        activation: FloatArray,
+        activationOffset: Int,
+        weight: ByteArray,
+        weightByteOffset: Int,
+        inputDim: Int,
+        outputDim: Int,
+        out: FloatArray,
+        outOffset: Int,
+    ) {
+        require(inputDim % 4 == 0) {
+            "NativeKnTernaryF32Gemv: inputDim must be a multiple of 4; got $inputDim"
+        }
+        if (outputDim == 0) return
+        if (inputDim == 0) {
+            // The C kernel writes 0.0f per row; mirror that without pinning
+            // empty arrays (addressOf(0) needs at least one element).
+            out.fill(0f, outOffset, outOffset + outputDim)
+            return
+        }
+        activation.usePinned { inPin ->
+            weight.usePinned { wPin ->
+                out.usePinned { outPin ->
+                    skainet_ternary_f32_gemv(
+                        inPin.addressOf(0),
+                        activationOffset,
+                        wPin.addressOf(0).reinterpret(),
+                        weightByteOffset,
+                        inputDim,
+                        outputDim,
+                        outPin.addressOf(0),
+                        outOffset,
+                    )
+                }
+            }
+        }
+    }
+
+    /**
+     * Install this kernel into the dispatcher. The archive is linked into the
+     * binary, so unlike the FFM/JNI faces there is no missing-artifact case —
+     * still routed through [TernaryF32KernelPack.install] so the contract
+     * (and the returned serving name) stays uniform across bridges.
+     */
+    public fun install(warn: (String) -> Unit = {}): String =
+        TernaryF32KernelPack.install(this, warn = warn)
+}

skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32GemvParityTest.kt

@@ -0,0 +1,88 @@
+package sk.ainet.exec.kernel
+
+import kotlin.math.abs
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * Proves the Kotlin/Native cinterop path of the vendored NeoGPU LUT kernel
+ * (#1139): [NativeKnTernaryF32Gemv] (calling `skainet_ternary_f32_gemv` from
+ * libskainet_kernels.a) must agree with the local decode reference —
+ * `((byte >> (lane*2)) & 3) - 1`, the sequential BITNET_B1_58 payload rule,
+ * byte code 3 → +2 included.
+ *
+ * Runs on the host target (scalar or Apple-NEON archive) AND on linuxArm64
+ * under qemu (`-PcrossArm64=true`), where the cross-built archive carries the
+ * NEON LUT path pinned to `-march=armv8-a` — the Pi-4/Cortex-A72 consumption
+ * build. The ternary codes-dot is exact; only summation order differs, so
+ * tolerances are tight. The 2048-row case crosses the kernel's internal
+ * pthread threshold (512), pinning the thread partitioning through cinterop.
+ */
+class NativeKnTernaryF32GemvParityTest {
+
+    private fun decode(b: Byte, lane: Int): Float =
+        (((b.toInt() and 0xFF) shr (lane * 2)) and 3).toFloat() - 1f
+
+    private fun reference(
+        input: FloatArray, weight: ByteArray, inputDim: Int, outputDim: Int,
+    ): FloatArray {
+        val rowBytes = inputDim / 4
+        return FloatArray(outputDim) { o ->
+            var acc = 0.0
+            for (bi in 0 until rowBytes) {
+                val b = weight[o * rowBytes + bi]
+                for (lane in 0 until 4) acc += decode(b, lane) * input[bi * 4 + lane]
+            }
+            acc.toFloat()
+        }
+    }
+
+    private fun assertParity(inputDim: Int, outputDim: Int, seed: Int) {
+        val rng = Random(seed)
+        val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f }
+        val weight = ByteArray(outputDim * inputDim / 4).also { rng.nextBytes(it) }
+        val expected = reference(input, weight, inputDim, outputDim)
+        val out = FloatArray(outputDim)
+        NativeKnTernaryF32Gemv.gemvPacked(input, 0, weight, 0, inputDim, outputDim, out, 0)
+        for (o in out.indices) {
+            val diff = abs(expected[o] - out[o])
+            assertTrue(diff <= 1e-3f, "[$o]: reference=${expected[o]} cinterop=${out[o]} diff=$diff")
+        }
+    }
+
+    @Test fun single_row_all_256_byte_values() {
+        // Integer activations keep sums exact; the 256 weight bytes enumerate
+        // the full decode table, pinning code 3 → +2 on this target's branch.
+        val inputDim = 1024
+        val input = FloatArray(inputDim) { ((it % 7) - 3).toFloat() }
+        val weight = ByteArray(256) { it.toByte() }
+        val out = FloatArray(1)
+        NativeKnTernaryF32Gemv.gemvPacked(input, 0, weight, 0, inputDim, 1, out, 0)
+        assertEquals(reference(input, weight, inputDim, 1)[0], out[0])
+    }
+
+    @Test fun projection_shape() = assertParity(inputDim = 2560, outputDim = 64, seed = 7)
+
+    @Test fun threaded_regime_above_512_rows() = assertParity(inputDim = 256, outputDim = 2048, seed = 11)
+
+    @Test fun offsets_are_honoured() {
+        // 0x22 = codes {2,0,2,0} → {+1,-1,+1,-1}; 0x55 = all code 1 → zeros.
+        val pad = 3
+        val input = FloatArray(pad + 8) { if (it < pad) 99f else (it - pad + 1).toFloat() }
+        val weight = ByteArray(5 + 4) { 0x55.toByte() }
+        weight[5] = 0x22
+        weight[6] = 0x22
+        val out = FloatArray(4) { -1f }
+        NativeKnTernaryF32Gemv.gemvPacked(input, pad, weight, 5, 8, 2, out, 2)
+        assertEquals(-1f, out[0]); assertEquals(-1f, out[1])
+        assertEquals(-4f, out[2]); assertEquals(0f, out[3])
+    }
+
+    @Test fun zero_input_dim_zeros_output() {
+        val out = FloatArray(3) { 9f }
+        NativeKnTernaryF32Gemv.gemvPacked(FloatArray(0), 0, ByteArray(0), 0, 0, 3, out, 0)
+        for (v in out) assertEquals(0f, v)
+    }
+}