From 0f43864b19438481f71c805d18df12817c0b5aa4 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 15:59:03 +0200 Subject: [PATCH] feat(kernels): NEON bitnet_gemv with a capability-gated pack and a reference fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1041 (SKEEP-003 P6, S2.9, proposal §5.2/§5.3; M2-F4, M2-A2). - `bitnet_gemv.c`: int8 activations against canonical TQ2_0 ternary weights. A ternary weight is an add, a subtract or nothing, so the vector work is the *unpacking*: a whole vector of 2-bit codes becomes int8 signs with a shift and a mask, after which `sdot` accumulates sixteen products per instruction on ARMv8.2. Codes stay in {0,1,2} through the inner loop and the bias is removed once per block — `Σ(c-1)a = Σca − Σa`, and `Σa` is the same for every output row. Built into both JNI variants and the native-cpu library; TQ1_0 and BitNet's own packing stay scalar, and the comment says why. - `BitNetGemvNative` + `TernaryKernelPacks.install(native, capabilities, warn)`: the reference is registered first, a pack overrides it, and its absence produces one warning and a slower kernel rather than a crash. The native kernel is registered under both a capability-carrying key and the operand-only key the dispatcher builds, since operands say nothing about the CPU and the pack only installs itself where the capability exists. - `NativeBitNetGemvKernel` unwraps the views once and falls back to the reference for anything the native kernel does not take — a prefill row batch or non-heap storage — so the fast path stays an optimization. - JNI shim + `JniKernels.bitnetGemvTq20` + `JniBitNetGemv.install()`, which picks the capability from the library variant `/proc/cpuinfo` already chose. Unlike the Q4_0…Q6_K SPI kernels, this one is bridged into the view registry: it was written against TernaryCodec's canonical row-major order, which is what a GGUF holds, so #973's block-order ambiguity does not arise here — and the parity test pins it. Measured on an ARMv8.2 Cortex-A55 reference board (asimddp), gcc -O3 -ffast-math -march=armv8.2-a+dotprod: k=1024 n=256 kernel 0.091 ms worst relative error 0 k=2560 n=2560 kernel 2.329 ms worst relative error 0 Parity is exact — the arithmetic is integer until the block scale. Against the portable Kotlin kernel that serves when the artifact is absent, on the same core, k=1024 n=256: 21.3 ms → 0.091 ms. M2-A2's "≥ 3x" is met with room to spare, though that multiple is measured against a *debug* Kotlin/Native build of the fallback and so is an upper bound; against compiler-vectorized C the hand-written NEON is 1.08–1.24x, which is the honest measure of what the intrinsics themselves buy on an A55. The seven ternary kernel tests and the timing test also run on that board from the Kotlin/Native binary; the pack contract (registration, fallback, capability, shape refusal) is tested in common code on every target. Gate: scripts/pr-gate.sh — all legs passed; the AAR builds for arm64-v8a and x86_64 with NDK clang (which, unlike gcc, rejects a zero `vshrq_n_u8` shift — the strips take pre-shifted vectors because of it). Co-Authored-By: Claude Fable 5 --- .../backend/api/kernel/TernaryKernelPacks.kt | 133 ++++++++++ .../api/kernel/BitNetGemvTimingTest.kt | 69 ++++++ .../api/kernel/TernaryKernelPacksTest.kt | 162 ++++++++++++ .../native/CMakeLists.txt | 1 + .../native/skainet_jni.c | 27 ++ .../sk/ainet/exec/kernel/jni/JniBitNetGemv.kt | 66 +++++ .../sk/ainet/exec/kernel/jni/JniKernels.kt | 12 + .../native/CMakeLists.txt | 1 + .../native/include/skainet_kernels.h | 21 ++ .../native/src/bitnet_gemv.c | 232 ++++++++++++++++++ 10 files changed, 724 insertions(+) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacks.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTimingTest.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacksTest.kt create mode 100644 skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBitNetGemv.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/native/src/bitnet_gemv.c diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacks.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacks.kt new file mode 100644 index 000000000..1a9de3979 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacks.kt @@ -0,0 +1,133 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.I8Absmax +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * A hand-written `bitnet_gemv` supplied by a platform pack (SKEEP-003 §5.2, M2-F4). + * + * Deliberately array-shaped rather than view-shaped: the implementations are JNI or cinterop shims + * that pin primitive arrays, and keeping the SPI at that level means the pack modules carry no + * knowledge of `TensorView`. [TernaryKernelPacks] does the unwrapping once. + * + * The weight is canonical row-major `TQ2_0` — the order [sk.ainet.lang.memory.TernaryCodec] writes + * and a GGUF holds. That is *not* the block-major order the Q4_0…Q6_K SPI kernels take, which is + * why this one can be bridged into the view registry while those still wait on #973: there is no + * ambiguity here, and the parity test pins it. + */ +@ExperimentalMemoryApi +public interface BitNetGemvNative { + /** A name for logs and traces, e.g. `neon-dotprod`. */ + public val name: String + + /** `out[o] = activationScale * Σ_k code(k) · weight(o, k)` for one token. */ + public fun gemvTq2_0( + activation: ByteArray, + activationOffset: Int, + activationScale: Float, + weight: ByteArray, + weightByteOffset: Int, + inputDim: Int, + outputDim: Int, + out: FloatArray, + outOffset: Int, + ) +} + +/** + * Installs the ternary kernels: the portable reference always, and a platform `bitnet_gemv` on top + * of it when one is available (M2-F4). + * + * Removing the native artifact is not an error and never a crash — the reference kernel is already + * registered, so dispatch keeps working at reference speed and the caller is *told* through [warn] + * rather than left to wonder why decode got slower. + */ +@ExperimentalMemoryApi +public object TernaryKernelPacks { + + /** Capability a `bitnet_gemv` pack declares when it needs ARMv8.2 dot-product instructions. */ + public const val CAPABILITY_DOTPROD: String = "dotprod" + + private val ternaryEncodings = listOf(TensorEncoding.TQ1_0, TensorEncoding.TQ2_0, TensorEncoding.BITNET_B1_58) + + /** + * @param native the platform kernel, or `null` when its artifact is absent + * @param capabilities what [native] needs; recorded in the key so a device without them never + * selects it (§5.2) + * @param warn where the "running without the native kernel" notice goes + * @return the name of the kernel that will serve TQ2_0 — the pack's, or the reference + */ + public fun install( + native: BitNetGemvNative? = null, + capabilities: Set = emptySet(), + warn: (String) -> Unit = {}, + ): String { + for (encoding in ternaryEncodings) { + KernelDispatch.register(BitNetGemvKernel(BitNetGemvKernel.keyFor(Format(FP32, encoding)))) + } + if (native == null) { + warn( + "bitnet_gemv: no native kernel available — using the portable reference. " + + "Add the NEON artifact for the fast path; nothing else changes.", + ) + return BitNetGemvKernel(BitNetGemvKernel.keyFor(Format(FP32, TensorEncoding.TQ2_0))).name + } + val key = BitNetGemvKernel.keyFor(Format(FP32, TensorEncoding.TQ2_0)).copy(capabilities = capabilities) + val kernel = NativeBitNetGemvKernel(native, key) + KernelDispatch.register(kernel) + // Registering under the capability-free key too: the dispatcher builds its key from the + // operands, which say nothing about the CPU. The pack only installs itself on a device that + // *has* the capability, so the two keys select the same kernel there and neither exists on + // a device that does not. + KernelDispatch.register(NativeBitNetGemvKernel(native, BitNetGemvKernel.keyFor(Format(FP32, TensorEncoding.TQ2_0)))) + return kernel.name + } +} + +/** + * A [ViewKernel] over a [BitNetGemvNative]: unwraps the two views once per call and hands the + * native kernel the arrays it wants. + * + * Falls back to the reference for anything the native kernel does not take — a multi-row + * activation (prefill) or storage that is not a heap array — instead of failing: the fast path is + * an optimization, never a correctness requirement. + */ +@ExperimentalMemoryApi +public class NativeBitNetGemvKernel( + private val native: BitNetGemvNative, + override val key: KernelKey, +) : ViewKernel { + + override val name: String get() = "bitnet_gemv/${native.name}" + + private val reference = BitNetGemvKernel(key) + + override fun run(inputs: List, out: TensorView) { + val a = inputs[0] + val w = inputs[1] + val rows = a.shape[0] + val activationBytes = (a.storage as? Storage.Heap)?.bytes + val weightBytes = (w.storage as? Storage.Heap)?.bytes + val outFloats = (out.storage as? Storage.Heap)?.floats + if (rows != 1 || activationBytes == null || weightBytes == null || outFloats == null) { + reference.run(inputs, out) + return + } + native.gemvTq2_0( + activation = activationBytes, + activationOffset = (a.storage as Storage.Heap).arrayOffset, + activationScale = I8Absmax.scaleOf(a, 0), + weight = weightBytes, + weightByteOffset = (w.storage as Storage.Heap).arrayOffset, + inputDim = a.shape[1], + outputDim = w.shape[0], + out = outFloats, + outOffset = (out.storage as Storage.Heap).arrayOffset, + ) + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTimingTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTimingTest.kt new file mode 100644 index 000000000..ca364170f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTimingTest.kt @@ -0,0 +1,69 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.I8Absmax +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryBlockDecoder +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertTrue +import kotlin.time.TimeSource + +/** + * The **fallback**'s cost, measured where the test runs (#1041, M2-A2). + * + * `bitnet_gemv/reference` is what a device gets when the NEON artifact is absent, so its time is + * the denominator of the "NEON is N× faster" claim. This prints milliseconds per call for a fixed + * shape and asserts only that it produced numbers — a speed assertion here would be a flake on + * shared CI hardware, and the acceptance measurement belongs on the reference device, where this + * same test is run from the Kotlin/Native binary. + */ +@OptIn(ExperimentalMemoryApi::class) +class BitNetGemvTimingTest { + + private companion object { + const val K = 1024 // four TQ2_0 blocks per row + const val N = 256 + const val CALLS = 5 + } + + @Test + fun referenceKernelThroughput() { + var seed = 7 + val weightValues = FloatArray(N * K) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 3 - 1) * 0.5f + } + val bytes = TernaryCodec.encode(TensorEncoding.TQ2_0, weightValues) + val weight = TensorView.packed( + Storage.Heap.wrap(bytes), Shape(N, K), TensorEncoding.TQ2_0, + TernaryBlockDecoder(TensorEncoding.TQ2_0), + ) + val activationFloats = FloatArray(K) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 2000 - 1000) / 1000f + } + val activation = I8Absmax.requantize( + TensorView.dense(Storage.Heap.wrap(activationFloats), Shape(1, K), FP32), + Scope.Ambient, + ) + val out = TensorView.dense(Storage.Heap.floats(N), Shape(1, N), FP32) + val kernel = BitNetGemvKernel(BitNetGemvKernel.keyFor(weight.format)) + + kernel.run(listOf(activation, weight), out) // warm up + val mark = TimeSource.Monotonic.markNow() + repeat(CALLS) { kernel.run(listOf(activation, weight), out) } + val perCall = mark.elapsedNow().inWholeMicroseconds / CALLS.toDouble() / 1000.0 + + println("[bitnet_gemv] reference k=$K n=$N: $perCall ms/call (${N.toLong() * K} MACs)") + assertTrue(perCall >= 0.0) + var sawNonZero = false + for (o in 0 until N) if (out.get(0, o) != 0f) { sawNonZero = true; break } + assertTrue(sawNonZero, "the timed kernel must actually compute something") + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacksTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacksTest.kt new file mode 100644 index 000000000..797d3f3a4 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryKernelPacksTest.kt @@ -0,0 +1,162 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.I8Absmax +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryBlockDecoder +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1041 (M2-F4): a platform `bitnet_gemv` takes over when its artifact is present, and its absence + * is a warning and a slower kernel — never a crash. + * + * The NEON kernel itself is C and is tested where it can run; what belongs here is the contract + * around it, which every target can check: who gets registered, what happens without the artifact, + * and that the native path is only taken for the shapes it declares. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryKernelPacksTest { + + private val k = 256 + private val n = 4 + + @BeforeTest fun setUp() = KernelDispatch.clearForTesting() + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + /** A stand-in for the JNI kernel: records that it ran, and computes the same thing. */ + private class FakeNative(override val name: String = "neon-dotprod") : BitNetGemvNative { + var calls: Int = 0 + override fun gemvTq2_0( + activation: ByteArray, activationOffset: Int, activationScale: Float, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + out: FloatArray, outOffset: Int, + ) { + calls++ + val codes = TernaryCodec.codes(TensorEncoding.TQ2_0, weight, outputDim * inputDim, weightByteOffset) + val blocks = inputDim / 256 + for (o in 0 until outputDim) { + var acc = 0f + for (b in 0 until blocks) { + val scaleOffset = weightByteOffset + ((o * blocks + b) * 66) + 64 + val d = sk.ainet.lang.types.Fp16Codec.decode( + (weight[scaleOffset].toInt() and 0xFF) or ((weight[scaleOffset + 1].toInt() and 0xFF) shl 8), + ) + var partial = 0 + for (i in 0 until 256) { + val code = codes[o * inputDim + b * 256 + i].toInt() + if (code != 0) { + val a = activation[activationOffset + b * 256 + i].toInt() + partial += if (code > 0) a else -a + } + } + acc += partial * d + } + out[outOffset + o] = acc * activationScale + } + } + } + + private fun weight(): TensorView { + var seed = 5 + val values = FloatArray(n * k) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 3 - 1) * 0.5f + } + val bytes = TernaryCodec.encode(TensorEncoding.TQ2_0, values) + return TensorView.packed( + Storage.Heap.wrap(bytes), Shape(n, k), TensorEncoding.TQ2_0, + TernaryBlockDecoder(TensorEncoding.TQ2_0), + ) + } + + private fun activation(rows: Int = 1): TensorView { + var seed = 9 + val floats = FloatArray(rows * k) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 2000 - 1000) / 1000f + } + return I8Absmax.requantize(TensorView.dense(Storage.Heap.wrap(floats), Shape(rows, k), FP32), Scope.Ambient) + } + + @Test + fun withoutTheArtifactTheReferenceServesAndTheCallerIsTold() { + val warnings = mutableListOf() + val serving = TernaryKernelPacks.install(native = null, warn = { warnings += it }) + + assertEquals("bitnet_gemv/reference", serving) + assertEquals(1, warnings.size, "exactly one notice, not a crash: $warnings") + assertTrue(warnings.single().contains("portable reference"), warnings.single()) + + // and dispatch still works + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(activation(), weight(), out, Scope.Ambient, sink) + assertEquals("bitnet_gemv/reference", sink.eventsOf().single().kernel) + } + + @Test + fun withTheArtifactTheNativeKernelTakesOverAndAgreesWithTheReference() { + val native = FakeNative() + val warnings = mutableListOf() + val serving = TernaryKernelPacks.install(native, setOf(TernaryKernelPacks.CAPABILITY_DOTPROD)) { warnings += it } + assertEquals("bitnet_gemv/neon-dotprod", serving) + assertTrue(warnings.isEmpty(), "nothing to warn about: $warnings") + + val w = weight() + val a = activation() + val fromNative = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, fromNative, Scope.Ambient, sink) + assertEquals("bitnet_gemv/neon-dotprod", sink.eventsOf().single().kernel) + assertEquals(1, native.calls) + + val fromReference = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + BitNetGemvKernel(BitNetGemvKernel.keyFor(w.format)).run(listOf(a, w), fromReference) + for (o in 0 until n) { + val got = fromNative.get(0, o) + val want = fromReference.get(0, o) + assertTrue(abs(got - want) <= 1e-5f * maxOf(1f, abs(want)), "[$o]: $got vs $want") + } + } + + @Test + fun theNativePathIsSkippedForShapesItDoesNotTake() { + val native = FakeNative() + TernaryKernelPacks.install(native) + val w = weight() + // prefill: more than one row is not what the gemv takes — the reference runs instead + val out = TensorView.dense(Storage.Heap.floats(2 * n), Shape(2, n), FP32) + NativeBitNetGemvKernel(native, BitNetGemvKernel.keyFor(w.format)).run(listOf(activation(rows = 2), w), out) + assertEquals(0, native.calls, "a multi-row activation falls back rather than failing") + var nonZero = false + for (r in 0 until 2) for (o in 0 until n) if (out.get(r, o) != 0f) nonZero = true + assertTrue(nonZero, "and it still computed the answer") + } + + @Test + fun theCapabilityIsRecordedInTheKey() { + val native = FakeNative() + TernaryKernelPacks.install(native, setOf(TernaryKernelPacks.CAPABILITY_DOTPROD)) + val keys = KernelDispatch.kernels().filter { it.name.startsWith("bitnet_gemv/neon") }.map { it.key } + assertTrue( + keys.any { it.capabilities == setOf(TernaryKernelPacks.CAPABILITY_DOTPROD) }, + "the pack declares what it needs: $keys", + ) + assertTrue(keys.any { it.capabilities.isEmpty() }, "and is reachable from an operand-only key: $keys") + } +} diff --git a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt index 002a64f8c..05a8eef4d 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt @@ -17,6 +17,7 @@ set(SKAINET_KERNEL_SOURCES ${SKAINET_KERNELS_ROOT}/src/q4k_matmul.c ${SKAINET_KERNELS_ROOT}/src/q5k_matmul.c ${SKAINET_KERNELS_ROOT}/src/q6k_matmul.c + ${SKAINET_KERNELS_ROOT}/src/bitnet_gemv.c ) set(SKAINET_JNI_SHIM ${CMAKE_CURRENT_SOURCE_DIR}/skainet_jni.c) 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 9d3a0bce2..424d0b690 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c +++ b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c @@ -142,3 +142,30 @@ Java_sk_ainet_exec_kernel_jni_JniKernels_q6kMatmul( skainet_q6k_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, inputDim, outputDim, out, outputOffset)) } + +/* + * bitnet_gemv (SKEEP-003 §5.3, #1041): int8 activations against ternary TQ2_0 + * weights. Its activation is a *byte* array, not floats, so it does not fit + * SKAINET_JNI_MATMUL_BODY's float-input shape and pins its three arrays here. + */ +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_bitnetGemvTq20( + JNIEnv* env, jobject thiz, + jbyteArray activation, jint activationOffset, jfloat activationScale, + jbyteArray weight, jint weightByteOffset, + jint inputDim, jint outputDim, + jfloatArray output, jint outputOffset +) { + (void) thiz; + jbyte* act = (*env)->GetPrimitiveArrayCritical(env, activation, NULL); + jbyte* w = act ? (*env)->GetPrimitiveArrayCritical(env, weight, NULL) : NULL; + jfloat* out = w ? (*env)->GetPrimitiveArrayCritical(env, output, NULL) : NULL; + if (out) { + skainet_bitnet_gemv_tq2_0((const int8_t*) act, activationOffset, activationScale, + (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset); + } + if (out) (*env)->ReleasePrimitiveArrayCritical(env, output, out, 0); + if (w) (*env)->ReleasePrimitiveArrayCritical(env, weight, w, JNI_ABORT); + if (act) (*env)->ReleasePrimitiveArrayCritical(env, activation, act, JNI_ABORT); +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBitNetGemv.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBitNetGemv.kt new file mode 100644 index 000000000..6215c3005 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniBitNetGemv.kt @@ -0,0 +1,66 @@ +package sk.ainet.exec.kernel.jni + +import sk.ainet.backend.api.kernel.BitNetGemvNative +import sk.ainet.backend.api.kernel.TernaryKernelPacks +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * The NEON `bitnet_gemv` as a [BitNetGemvNative] (SKEEP-003 §5.2/§5.3, #1041, M2-F4). + * + * The maths lives in `bitnet_gemv.c`, shared with every other consumer of those kernels; this is + * the Android/JNI face of it. Whether the loaded library was built with ARMv8.2 dot-product + * instructions is what [JniKernels.variant] already decided from `/proc/cpuinfo`, so the name — and + * the capability the kernel is registered with — follows that decision rather than guessing again. + */ +@OptIn(ExperimentalMemoryApi::class) +public object JniBitNetGemv : BitNetGemvNative { + + override val name: String + get() = when (JniKernels.variant) { + JniKernels.Variant.V82_DOTPROD -> "neon-dotprod" + JniKernels.Variant.BASELINE -> "neon" + null -> "unloaded" + } + + override fun gemvTq2_0( + activation: ByteArray, + activationOffset: Int, + activationScale: Float, + weight: ByteArray, + weightByteOffset: Int, + inputDim: Int, + outputDim: Int, + out: FloatArray, + outOffset: Int, + ) { + JniKernels.bitnetGemvTq20( + activation, activationOffset, activationScale, + weight, weightByteOffset, + inputDim, outputDim, + out, outOffset, + ) + } + + /** + * Install this kernel into the dispatcher, or leave the reference in place and say so. + * + * Removing the AAR is a supported configuration: dispatch keeps working through the portable + * kernel, decode gets slower, and [warn] is where that shows up — never a crash (M2-F4). + * + * @return the name of the kernel that will serve TQ2_0 matmuls + */ + public fun install(warn: (String) -> Unit = { println("[skainet] $it") }): String = + if (JniKernels.isLoaded) { + TernaryKernelPacks.install( + native = this, + capabilities = if (JniKernels.variant == JniKernels.Variant.V82_DOTPROD) { + setOf(TernaryKernelPacks.CAPABILITY_DOTPROD) + } else { + emptySet() + }, + warn = warn, + ) + } else { + TernaryKernelPacks.install(native = null, warn = warn) + } +} 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 465b98b9b..0a2b14e06 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 @@ -94,6 +94,18 @@ public object JniKernels { output: FloatArray, outputOffset: Int, ) + /** + * `bitnet_gemv` (#1041): int8 activations with one absmax scale against canonical TQ2_0 + * ternary weights. The activation is a byte array of codes, not floats — a ternary matmul's + * whole point is that neither operand is wide. + */ + public external fun bitnetGemvTq20( + activation: ByteArray, activationOffset: Int, activationScale: Float, + 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-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index b72a2cc41..5525e0a4f 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -15,6 +15,7 @@ set(SKAINET_KERNEL_SOURCES src/q4k_matmul.c src/q5k_matmul.c src/q6k_matmul.c + src/bitnet_gemv.c src/fp32_matmul.c src/bf16_matmul.c src/fp16_matmul.c 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 e3e960a63..e3f8f0897 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 @@ -270,6 +270,27 @@ SKAINET_API void skainet_q5_1_matmul( int32_t output_offset ); +/* + * `bitnet_gemv`: int8 activations (one absmax scale per token) against ternary + * TQ2_0 weights — the SKEEP-003 §5.3 kernel, matching + * sk.ainet.backend.api.kernel.BitNetGemvKernel. + * + * A ternary weight is an add, a subtract or nothing; the vector work is the + * unpacking, after which `sdot` accumulates sixteen products per instruction + * where the core has ARMv8.2 dotprod. + * + * @param activation int8 codes of one token, `input_dim` of them + * @param activation_scale the token's absmax scale (`absmax / 127`) + * @param weight canonical TQ2_0 blocks, row-major per output row + * @param output `output_dim` floats + */ +SKAINET_API void skainet_bitnet_gemv_tq2_0( + const int8_t* activation, int32_t activation_offset, + float activation_scale, + const uint8_t* weight, int32_t weight_byte_offset, + int32_t input_dim, int32_t output_dim, + float* output, int32_t output_offset); + #ifdef __cplusplus } #endif diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/bitnet_gemv.c b/skainet-backends/skainet-backend-native-cpu/native/src/bitnet_gemv.c new file mode 100644 index 000000000..e91442875 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/bitnet_gemv.c @@ -0,0 +1,232 @@ +#include "skainet_kernels.h" +#include "skainet_simd.h" + +#include +#include +#include +#include + +/* + * Native `bitnet_gemv`: int8 activations (one absmax scale per token) against + * ternary weights, matching sk.ainet.backend.api.kernel.BitNetGemvKernel. + * + * A ternary weight is -1, 0 or +1 times a block scale, so the arithmetic is an + * add, a subtract or nothing — no multiplies in the inner loop. What NEON adds + * is the unpacking: a whole vector of 2-bit codes becomes a vector of int8 + * signs with a shift, a mask and a subtract, after which `sdot` (ARMv8.2 + * dotprod) accumulates 16 products per instruction. + * + * Weight layout — canonical GGML TQ2_0, exactly what TernaryCodec writes and + * what a GGUF holds: + * block = 66 bytes: 64 payload + FP16 scale + * byte `c*32 + m` of a block holds four elements *32 apart*: + * element c*128 + l*32 + m sits in bit pair l + * That interleave is why the vector path is natural here: for a fixed chunk c + * and bit pair l, bytes c*32..c*32+31 unpack to elements + * c*128 + l*32 .. + 31 — thirty-two *consecutive* elements, which line up with + * thirty-two consecutive activations. + * + * BitNet b1.58's own packing (four consecutive elements per byte, one scale for + * the whole tensor) takes the scalar path here: de-interleaving it costs more + * than it saves, and a checkpoint that is going to be run through this kernel + * ships as TQ2_0. TQ1_0's base-3 packing is likewise scalar — it is a storage + * format, not a runtime one. + */ + +#define SKAINET_TQ2_0_BLOCK_SIZE 256 +#define SKAINET_TQ2_0_BYTES_PER_BLOCK 66 + +static inline float skainet_bitnet_fp16(uint16_t h) { + uint32_t sign = ((uint32_t)(h & 0x8000u)) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t mant = h & 0x3FFu; + uint32_t bits; + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + int e = -14; + while ((mant & 0x400u) == 0) { + mant <<= 1; + --e; + } + mant &= 0x3FFu; + bits = sign | ((uint32_t)(e + 127) << 23) | (mant << 13); + } + } else if (exp == 0x1Fu) { + bits = sign | 0x7F800000u | (mant << 13); + } else { + bits = sign | ((uint32_t)(exp - 15 + 127) << 23) | (mant << 13); + } + float r; + memcpy(&r, &bits, sizeof(r)); + return r; +} + +/* Scalar reference for one TQ2_0 block: the shape every vector path must match. */ +static inline int32_t skainet_tq2_0_block_scalar( + const uint8_t* SKAINET_RESTRICT qs, + const int8_t* SKAINET_RESTRICT act, + int32_t count) { + int32_t acc = 0; + for (int32_t c = 0; c < 2; ++c) { + for (int32_t l = 0; l < 4; ++l) { + for (int32_t m = 0; m < 32; ++m) { + int32_t element = c * 128 + l * 32 + m; + if (element >= count) { + continue; + } + int32_t code = ((qs[c * 32 + m] >> (2 * l)) & 3) - 1; + if (code > 0) { + acc += act[element]; + } else if (code < 0) { + acc -= act[element]; + } + } + } + } + return acc; +} + +#if defined(SKAINET_HAVE_NEON) +/* + * One whole TQ2_0 block (256 elements) with NEON. Two 32-byte chunks × four bit + * pairs, each producing 32 consecutive weight codes to pair with 32 consecutive + * activations. + * + * Returns the *biased* sum `Σ code·a` with codes still in `{0,1,2}`. Converting + * them to `{-1,0,+1}` would cost a subtract per strip; instead the caller + * subtracts `Σ a` once per block, since `Σ (code-1)·a = Σ code·a − Σ a` and the + * activation sum is the same for every output row. + */ +SKAINET_DOTPROD_TARGET +static inline int32_t skainet_tq2_0_block_neon( + const uint8_t* SKAINET_RESTRICT qs, + const int8_t* SKAINET_RESTRICT act) { + const uint8x16_t mask = vdupq_n_u8(3); + int32x4_t acc = vdupq_n_s32(0); + + /* + * `vshrq_n_u8` takes a *compile-time* shift, so the four bit pairs are + * unrolled rather than looped — which is what one would write by hand + * anyway: four independent 32-element strips per 32-byte chunk. + */ +/* + * `vshrq_n_u8` requires a shift in [1, 8] — clang rejects a literal 0, which is + * why the strips take an already-shifted vector rather than a shift amount, and + * why strip 0 passes the bytes through untouched. + */ +#if defined(SKAINET_HAVE_DOTPROD) || defined(SKAINET_DOTPROD_DISPATCH) +#define SKAINET_TQ2_STRIP(chunk, strip, lo_bits, hi_bits) \ + do { \ + const int8x16_t wlo = vreinterpretq_s8_u8(vandq_u8((lo_bits), mask)); \ + const int8x16_t whi = vreinterpretq_s8_u8(vandq_u8((hi_bits), mask)); \ + const int8_t* a = act + (chunk) * 128 + (strip) * 32; \ + acc = vdotq_s32(acc, wlo, vld1q_s8(a)); \ + acc = vdotq_s32(acc, whi, vld1q_s8(a + 16)); \ + } while (0) +#else +#define SKAINET_TQ2_STRIP(chunk, strip, lo_bits, hi_bits) \ + do { \ + const int8x16_t wlo = vreinterpretq_s8_u8(vandq_u8((lo_bits), mask)); \ + const int8x16_t whi = vreinterpretq_s8_u8(vandq_u8((hi_bits), mask)); \ + const int8_t* a = act + (chunk) * 128 + (strip) * 32; \ + const int8x16_t alo = vld1q_s8(a); \ + const int8x16_t ahi = vld1q_s8(a + 16); \ + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(wlo), vget_low_s8(alo))); \ + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(wlo), vget_high_s8(alo))); \ + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(whi), vget_low_s8(ahi))); \ + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(whi), vget_high_s8(ahi))); \ + } while (0) +#endif + + { + const uint8x16_t lo = vld1q_u8(qs); + const uint8x16_t hi = vld1q_u8(qs + 16); + SKAINET_TQ2_STRIP(0, 0, lo, hi); + SKAINET_TQ2_STRIP(0, 1, vshrq_n_u8(lo, 2), vshrq_n_u8(hi, 2)); + SKAINET_TQ2_STRIP(0, 2, vshrq_n_u8(lo, 4), vshrq_n_u8(hi, 4)); + SKAINET_TQ2_STRIP(0, 3, vshrq_n_u8(lo, 6), vshrq_n_u8(hi, 6)); + } + { + const uint8x16_t lo = vld1q_u8(qs + 32); + const uint8x16_t hi = vld1q_u8(qs + 48); + SKAINET_TQ2_STRIP(1, 0, lo, hi); + SKAINET_TQ2_STRIP(1, 1, vshrq_n_u8(lo, 2), vshrq_n_u8(hi, 2)); + SKAINET_TQ2_STRIP(1, 2, vshrq_n_u8(lo, 4), vshrq_n_u8(hi, 4)); + SKAINET_TQ2_STRIP(1, 3, vshrq_n_u8(lo, 6), vshrq_n_u8(hi, 6)); + } +#undef SKAINET_TQ2_STRIP + + return vaddvq_s32(acc); +} +#endif /* SKAINET_HAVE_NEON */ + +SKAINET_API void skainet_bitnet_gemv_tq2_0( + const int8_t* SKAINET_RESTRICT activation, int32_t activation_offset, + float activation_scale, + const uint8_t* SKAINET_RESTRICT weight, int32_t weight_byte_offset, + int32_t input_dim, int32_t output_dim, + float* SKAINET_RESTRICT output, int32_t output_offset) { + if (input_dim <= 0 || output_dim <= 0) { + return; + } + const int8_t* act = activation + activation_offset; + const uint8_t* w = weight + weight_byte_offset; + const int32_t blocks_per_row = (input_dim + SKAINET_TQ2_0_BLOCK_SIZE - 1) / SKAINET_TQ2_0_BLOCK_SIZE; + + /* + * Σ activation per block, computed once for the whole call: the vector path + * accumulates codes in {0,1,2} and removes the bias here rather than paying a + * subtract per 16 weights, per output row. + */ + enum { SKAINET_MAX_STACK_BLOCKS = 128 }; + int32_t stack_sums[SKAINET_MAX_STACK_BLOCKS]; + int32_t* activation_sums = stack_sums; + int32_t* heap_sums = NULL; + if (blocks_per_row > SKAINET_MAX_STACK_BLOCKS) { + heap_sums = (int32_t*)malloc((size_t)blocks_per_row * sizeof(int32_t)); + if (heap_sums == NULL) { + return; + } + activation_sums = heap_sums; + } + for (int32_t b = 0; b < blocks_per_row; ++b) { + const int32_t first = b * SKAINET_TQ2_0_BLOCK_SIZE; + const int32_t count = (input_dim - first) < SKAINET_TQ2_0_BLOCK_SIZE + ? (input_dim - first) + : SKAINET_TQ2_0_BLOCK_SIZE; + int32_t s = 0; + for (int32_t i = 0; i < count; ++i) { + s += act[first + i]; + } + activation_sums[b] = s; + } + + for (int32_t o = 0; o < output_dim; ++o) { + float sum = 0.0f; + for (int32_t b = 0; b < blocks_per_row; ++b) { + const uint8_t* block = w + (size_t)(o * blocks_per_row + b) * SKAINET_TQ2_0_BYTES_PER_BLOCK; + const int32_t first = b * SKAINET_TQ2_0_BLOCK_SIZE; + const int32_t count = (input_dim - first) < SKAINET_TQ2_0_BLOCK_SIZE + ? (input_dim - first) + : SKAINET_TQ2_0_BLOCK_SIZE; + uint16_t scale_bits; + memcpy(&scale_bits, block + 64, sizeof(scale_bits)); + const float d = skainet_bitnet_fp16(scale_bits); + int32_t partial; +#if defined(SKAINET_HAVE_NEON) + if (count == SKAINET_TQ2_0_BLOCK_SIZE) { + partial = skainet_tq2_0_block_neon(block, act + first) - activation_sums[b]; + } else { + partial = skainet_tq2_0_block_scalar(block, act + first, count); + } +#else + partial = skainet_tq2_0_block_scalar(block, act + first, count); +#endif + sum += (float)partial * d; + } + output[output_offset + o] = sum * activation_scale; + } + free(heap_sums); +}