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); +}