diff --git a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c index 424d0b690..2762c39f4 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c +++ b/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)) +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt index aeaa331cb..7b42e66bb 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt +++ b/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) diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt index 0a2b14e06..3cbc510f0 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt +++ b/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, diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniTernaryF32Gemv.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniTernaryF32Gemv.kt new file mode 100644 index 000000000..130ecda80 --- /dev/null +++ b/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) +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32Gemv.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32Gemv.kt new file mode 100644 index 000000000..89302d6ab --- /dev/null +++ b/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) +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32GemvParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnTernaryF32GemvParityTest.kt new file mode 100644 index 000000000..8611e6dd2 --- /dev/null +++ b/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) + } +}