diff --git a/.gitignore b/.gitignore index b5dddf1ec..13854d115 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,4 @@ out/ ### Build artifacts ### com/ +.cxx/ diff --git a/CHANGELOG.md b/CHANGELOG.md index b20a28c60..fb951ad32 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,22 @@ ## [Unreleased] +### Added + +- **`skainet-backend-jni-cpu`: Android JNI bridge for the native NEON kernels.** ART has no + `java.lang.foreign`, so the priority-100 FFM provider can never run on Android — until now + Android inference ran on the priority-0 scalar floor. The new AAR ships the same C kernel + sources via NDK/CMake with thin JNI shims (`GetPrimitiveArrayCritical`, zero-copy pins) and + a priority-100 `JniKernelProvider` discovered via `ServiceLoader` (which ART supports; the + Android ops factory now installs discovered providers exactly like the JVM does). Two `.so` + tiers are built from the same sources and selected at load time from `/proc/cpuinfo`: + baseline `armv8-a` (NEON, runs on every arm64 core — 0 dot-product instructions, verified by + disassembly) and `armv8.2-a+fp16+dotprod` (enables the `vdotq_s32` q4k/q6k paths — would + SIGILL on Cortex-A53-class cores, hence the gate). Q8_0/Q4_0/Q4_K/Q5_K/Q6_K bridged; + 16 KB-page-aligned `.so`s (Android 15+); consumer R8 rules keep the ServiceLoader entry in + release builds. On-device parity tests included (`androidTest`, vs the scalar references). + Part of the mobile-kernels effort (#920). + ### Fixed - **`TensorData.copyToFloatArray()` default implementation works for rank >= 2.** It used to diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 350878db8..626e3d2c1 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -56,6 +56,8 @@ npm-brace-expansion = "2.1.4" # GHSA-rgw5-rvv9-x895, GHSA-mh99-v99m-4gvg jackson-databind = { module = "com.fasterxml.jackson.core:jackson-databind", version.ref = "jacksonDatabind" } json-schema-validator = { module = "com.networknt:json-schema-validator", version.ref = "jsonSchemaValidator" } junit = { module = "junit:junit", version.ref = "junit" } +androidx-test-junit = { module = "androidx.test.ext:junit", version = "1.2.1" } +androidx-test-runner = { module = "androidx.test:runner", version = "1.6.2" } junit-jupiter = { module = "org.junit.jupiter:junit-jupiter", version.ref = "junitJupiter" } junit-platform-launcher = { module = "org.junit.platform:junit-platform-launcher" } kotlinx-coroutines = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-core", version.ref = "kotlinxCoroutines" } diff --git a/settings.gradle.kts b/settings.gradle.kts index 5e25ef399..4a67c2b21 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -41,6 +41,7 @@ include("skainet-compile:skainet-compile-minerva") include("skainet-backends:skainet-backend-api") include("skainet-backends:skainet-backend-cpu") include("skainet-backends:skainet-backend-native-cpu") +include("skainet-backends:skainet-backend-jni-cpu") // ====== BENCHMARKS include("skainet-backends:benchmarks:jvm-cpu-jmh") diff --git a/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt b/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt index aa0ed4759..92ceda596 100644 --- a/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt +++ b/skainet-backends/skainet-backend-cpu/src/androidMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.android.kt @@ -1,13 +1,25 @@ package sk.ainet.exec.tensor.ops +import java.util.ServiceLoader +import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.exec.kernel.ScalarKernelProvider import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps { - // Non-JVM has no ServiceLoader; register the scalar packed-quant kernels - // (Q4_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0) so DefaultCpuOpsBase can dispatch them. + // ART supports java.util.ServiceLoader, so Android discovers kernel + // providers the same way the JVM does (#920): modules like + // skainet-backend-jni-cpu ship a META-INF/services entry and are + // picked up here at priority 100. Discovery failures are non-fatal — + // a provider that can't load just doesn't register. + runCatching { + ServiceLoader.load(KernelProvider::class.java).forEach { provider -> + runCatching { KernelRegistry.register(provider) } + } + } + // Scalar reference last: priority 0, always available — the floor the + // registry cascades to when no accelerated provider carries a kernel. KernelRegistry.register(ScalarKernelProvider) return { factory -> DefaultCpuOps(factory) } } diff --git a/skainet-backends/skainet-backend-jni-cpu/build.gradle.kts b/skainet-backends/skainet-backend-jni-cpu/build.gradle.kts new file mode 100644 index 000000000..37d9500ab --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/build.gradle.kts @@ -0,0 +1,75 @@ +plugins { + // AGP 9 ships built-in Kotlin support — applying kotlin("android") is an + // error since 9.0, so this module uses the android plugin alone. + alias(libs.plugins.androidLibrary) + alias(libs.plugins.vanniktech.mavenPublish) +} + +/* + * Android JNI bridge for the hand-written C/NEON matmul kernels + * (skainet-backend-native-cpu/native). ART has no java.lang.foreign, so the + * FFM provider can never run on Android — this module ships the same C + * sources as an AAR (.so via externalNativeBuild/NDK) with thin JNI shims + * and a priority-100 KernelProvider discovered via ServiceLoader (#920). + * + * Two shared libs are built from the same sources (see native/CMakeLists.txt): + * libskainet_jni.so — baseline armv8-a (NEON always present on arm64; + * runs on every device incl. Cortex-A53) + * libskainet_jni_v82.so — -march=armv8.2-a+fp16+dotprod (vdotq_s32 paths + * in q4k/q6k; would SIGILL on armv8.0 cores) + * JniKernels picks ONE at load time from /proc/cpuinfo features — runtime + * dispatch without symbol renaming or ifunc. + */ +android { + namespace = "sk.ainet.exec.kernel.jni" + compileSdk = libs.versions.android.compileSdk.get().toInt() + + defaultConfig { + minSdk = libs.versions.android.minSdk.get().toInt() + consumerProguardFiles("consumer-rules.pro") + ndk { + // arm64 is the target that matters; x86_64 keeps emulator CI and + // desktop AVDs working (scalar C paths). 32-bit ARM is out of + // scope — different intrinsics story, shrinking device share. + abiFilters += listOf("arm64-v8a", "x86_64") + } + externalNativeBuild { + cmake { + arguments += "-DANDROID_STL=none" + } + } + testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner" + } + + externalNativeBuild { + cmake { + path = file("native/CMakeLists.txt") + } + } + + compileOptions { + sourceCompatibility = JavaVersion.VERSION_11 + targetCompatibility = JavaVersion.VERSION_11 + } + + kotlin { + compilerOptions { + jvmTarget.set(org.jetbrains.kotlin.gradle.dsl.JvmTarget.JVM_11) + } + } + + // Publishing variant selection is configured by the vanniktech + // maven-publish plugin (AndroidSingleVariantLibrary "release"). +} + +dependencies { + api(project(":skainet-backends:skainet-backend-api")) + + testImplementation(libs.junit) + testImplementation(libs.kotlin.test) + + androidTestImplementation(libs.androidx.test.junit) + androidTestImplementation(libs.androidx.test.runner) + // Scalar reference kernels for on-device parity checks. + androidTestImplementation(project(":skainet-backends:skainet-backend-cpu")) +} diff --git a/skainet-backends/skainet-backend-jni-cpu/consumer-rules.pro b/skainet-backends/skainet-backend-jni-cpu/consumer-rules.pro new file mode 100644 index 000000000..71d1d9e02 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/consumer-rules.pro @@ -0,0 +1,15 @@ +# ServiceLoader discovery of the JNI kernel provider (#920). +# +# R8 full mode strips classes that are only referenced from +# META-INF/services files. Without these rules a RELEASE build silently +# loses the provider and inference falls back to the 100x-slower scalar +# path with no error — the exact failure mode these rules exist to prevent. +-keep class sk.ainet.exec.kernel.jni.JniKernelProviderFactory { *; } +-keep class sk.ainet.exec.kernel.jni.JniKernels { *; } + +# Native method registration: JNI resolves Java_sk_ainet_exec_kernel_jni_* +# symbols against these exact names — neither class nor method names may be +# renamed or removed. +-keepclasseswithmembers class sk.ainet.exec.kernel.jni.JniKernels { + native ; +} diff --git a/skainet-backends/skainet-backend-jni-cpu/gradle.properties b/skainet-backends/skainet-backend-jni-cpu/gradle.properties new file mode 100644 index 000000000..f2858347b --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/gradle.properties @@ -0,0 +1,2 @@ +POM_ARTIFACT_ID=skainet-backend-jni-cpu +POM_NAME=skainet Android (JNI) CPU kernel provider diff --git a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt new file mode 100644 index 000000000..174f3fafe --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt @@ -0,0 +1,51 @@ +cmake_minimum_required(VERSION 3.22) +project(skainet_jni C) + +# The C kernel sources are shared with skainet-backend-native-cpu — this +# module adds only the JNI shim. Single source of truth for the math. +set(SKAINET_KERNELS_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../skainet-backend-native-cpu/native) + +set(SKAINET_KERNEL_SOURCES + ${SKAINET_KERNELS_ROOT}/src/skainet_smoke.c + ${SKAINET_KERNELS_ROOT}/src/fp32_matmul.c + ${SKAINET_KERNELS_ROOT}/src/bf16_matmul.c + ${SKAINET_KERNELS_ROOT}/src/fp16_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q4_0_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q8_0_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q4k_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q5k_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q6k_matmul.c +) + +set(SKAINET_JNI_SHIM ${CMAKE_CURRENT_SOURCE_DIR}/skainet_jni.c) + +# Optimization flags match the sibling CMakeLists (-O3 -ffast-math) so the +# JNI path performs identically to the FFM/cinterop consumers of the same C. +set(SKAINET_C_FLAGS -O3 -ffast-math -funroll-loops) + +# Android 15+ devices ship 16 KB pages; .so files must be max-page-size +# aligned. Recent NDKs default to this, the explicit flag pins it. +set(SKAINET_LINK_FLAGS -Wl,-z,max-page-size=16384) + +function(skainet_add_jni_lib target) + add_library(${target} SHARED ${SKAINET_JNI_SHIM} ${SKAINET_KERNEL_SOURCES}) + target_include_directories(${target} PRIVATE ${SKAINET_KERNELS_ROOT}/include) + target_compile_options(${target} PRIVATE ${SKAINET_C_FLAGS}) + target_link_options(${target} PRIVATE ${SKAINET_LINK_FLAGS}) +endfunction() + +# Baseline: NDK default -march (armv8-a on arm64-v8a). NEON is architecturally +# guaranteed on AArch64, so fp32/q8_0/q4_0/q5k get their SIMD bodies, while the +# dotprod-guarded q4k/q6k paths fall back to scalar. Runs on EVERY arm64 core, +# including armv8.0 (Cortex-A53/A55-r0) — no SIGILL possible. +skainet_add_jni_lib(skainet_jni) + +# v8.2 variant: enables the vdotq_s32 paths in q4k/q6k (and fp16 storage ops). +# Would SIGILL on armv8.0 cores — JniKernels loads it only after confirming +# `asimddp` (+ `asimdhp`/`fphp`) in /proc/cpuinfo. On non-ARM ABIs (x86_64 +# emulator) the flags don't apply and the two libs are identical; both are +# still packaged so the loader's two-tier logic needs no per-ABI special case. +skainet_add_jni_lib(skainet_jni_v82) +if(ANDROID_ABI STREQUAL "arm64-v8a") + target_compile_options(skainet_jni_v82 PRIVATE -march=armv8.2-a+fp16+dotprod) +endif() diff --git a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c new file mode 100644 index 000000000..84bcc8cac --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c @@ -0,0 +1,116 @@ +#include +#include + +#include "skainet_kernels.h" + +/* + * Thin JNI shims over the shared C matmul kernels (skainet_kernels.h). + * One function per kernel, no logic beyond array pinning + the call. + * + * Array strategy: GetPrimitiveArrayCritical for every array. On ART, heap + * primitive arrays are contiguous, so criticals are zero-copy pins. Rules + * honored here: no JNI calls between Get and Release, releases in reverse + * acquisition order, read-only arrays released with JNI_ABORT (no + * write-back), the output with 0 (write-back + unpin). Kernel calls are + * millisecond-scale — well within acceptable critical-section length. + * + * Method names deliberately contain no underscores (JNI mangles `_` to + * `_1`, which is easy to get wrong silently). + */ + +/* Pin helper: acquire all three arrays, run `CALL`, release in reverse. */ +#define SKAINET_JNI_MATMUL_BODY(CALL) \ + jfloat* in = (*env)->GetPrimitiveArrayCritical(env, input, NULL); \ + jbyte* w = in ? (*env)->GetPrimitiveArrayCritical(env, weight, NULL) : NULL; \ + jfloat* out = w ? (*env)->GetPrimitiveArrayCritical(env, output, NULL) : NULL; \ + if (out) { \ + CALL; \ + } \ + if (out) (*env)->ReleasePrimitiveArrayCritical(env, output, out, 0); \ + if (w) (*env)->ReleasePrimitiveArrayCritical(env, weight, w, JNI_ABORT); \ + if (in) (*env)->ReleasePrimitiveArrayCritical(env, input, in, JNI_ABORT); + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_smoke( + JNIEnv* env, jobject thiz, + jfloatArray input, jfloatArray output, jint length +) { + (void) thiz; + jfloat* in = (*env)->GetPrimitiveArrayCritical(env, input, NULL); + jfloat* out = in ? (*env)->GetPrimitiveArrayCritical(env, output, NULL) : NULL; + if (out) { + skainet_smoke_double(in, out, length); + } + if (out) (*env)->ReleasePrimitiveArrayCritical(env, output, out, 0); + if (in) (*env)->ReleasePrimitiveArrayCritical(env, input, in, JNI_ABORT); +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q80Matmul( + 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_q8_0_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q40Matmul( + 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_q4_0_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q4kMatmul( + 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_q4k_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q5kMatmul( + 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_q5k_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q6kMatmul( + 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_q6k_matmul(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 new file mode 100644 index 000000000..f75565fb4 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt @@ -0,0 +1,105 @@ +package sk.ainet.exec.kernel.jni + +import androidx.test.ext.junit.runners.AndroidJUnit4 +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNotNull +import org.junit.Assert.assertTrue +import org.junit.Test +import org.junit.runner.RunWith +import sk.ainet.exec.kernel.ScalarQ4_0MatmulKernel +import sk.ainet.exec.kernel.ScalarQ4_KMatmulKernel +import sk.ainet.exec.kernel.ScalarQ6_KMatmulKernel +import sk.ainet.exec.kernel.ScalarQ8_0MatmulKernel +import kotlin.math.abs +import kotlin.random.Random + +/** + * On-device parity: the JNI kernels must agree with the commonMain scalar + * references within FMA + `-ffast-math` reassociation tolerance. Runs on + * any arm64-v8a device (NEON baseline or v8.2+dotprod tier, whichever the + * loader picked) and on x86_64 emulators (scalar C paths) — the same + * parity contract the qemu-aarch64 lane enforces for the Kotlin/Native + * consumers of the identical C sources. + */ +@RunWith(AndroidJUnit4::class) +class JniKernelParityTest { + + @Test + fun provider_is_available_on_device() { + assertTrue( + "JNI kernel provider must be available on-device (variant=${JniKernelProvider.activeVariant})", + JniKernelProvider.isAvailable(), + ) + assertNotNull(JniKernels.variant) + } + + private fun randomBlocks(numBlocks: Int, bytesPerBlock: Int, seed: Int): ByteArray { + val rng = Random(seed) + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + bytes[base + 0] = 0x00.toByte() // FP16 scale = 1.0 + bytes[base + 1] = 0x3C.toByte() + } + return bytes + } + + private fun assertParity( + inputDim: Int, outputDim: Int, seed: Int, tol: Float, + blockSize: Int, bytesPerBlock: Int, + reference: (FloatArray, Int, ByteArray, Int, Int, Int, FloatArray, Int) -> Unit, + jni: (FloatArray, Int, ByteArray, Int, Int, Int, FloatArray, Int) -> Unit, + ) { + val numBlocks = (inputDim / blockSize) * outputDim + val packed = randomBlocks(numBlocks, bytesPerBlock, seed) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val refOut = FloatArray(outputDim) + reference(input, 0, packed, 0, inputDim, outputDim, refOut, 0) + val jniOut = FloatArray(outputDim) + jni(input, 0, packed, 0, inputDim, outputDim, jniOut, 0) + + for (o in 0 until outputDim) { + val diff = abs(refOut[o] - jniOut[o]) + val rel = diff / (abs(refOut[o]) + 1e-9f) + assertTrue( + "row $o diverged: scalar=${refOut[o]} jni=${jniOut[o]} diff=$diff (variant=${JniKernels.variant})", + diff <= tol || rel < 1e-4f, + ) + } + } + + @Test + fun q80_parity() = assertParity( + 1024, 64, 42, 2e-1f, 32, 34, + ScalarQ8_0MatmulKernel::matmul, JniKernels::q80Matmul, + ) + + @Test + fun q40_parity() = assertParity( + 1024, 64, 7, 2e-1f, 32, 18, + ScalarQ4_0MatmulKernel::matmul, JniKernels::q40Matmul, + ) + + @Test + fun q4k_parity() = assertParity( + 1024, 64, 123, 2e-1f, 256, 144, + ScalarQ4_KMatmulKernel::matmul, JniKernels::q4kMatmul, + ) + + @Test + fun q6k_parity() = assertParity( + 1024, 64, 999, 2e-1f, 256, 210, + ScalarQ6_KMatmulKernel::matmul, JniKernels::q6kMatmul, + ) + + @Test + fun smoke_roundtrip() { + val input = floatArrayOf(1f, 2f, 3f, 4f) + val output = FloatArray(4) + JniKernels.smoke(input, output, 4) + assertEquals(2f, output[0], 0f) + assertEquals(8f, output[3], 0f) + } +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt new file mode 100644 index 000000000..10781bd02 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt @@ -0,0 +1,122 @@ +package sk.ainet.exec.kernel.jni + +import sk.ainet.backend.api.kernel.KernelProvider +import sk.ainet.backend.api.kernel.Q4KMatmulKernel +import sk.ainet.backend.api.kernel.Q4_0MatmulKernel +import sk.ainet.backend.api.kernel.Q5KMatmulKernel +import sk.ainet.backend.api.kernel.Q6KMatmulKernel +import sk.ainet.backend.api.kernel.Q8_0MatmulKernel +import sk.ainet.backend.api.kernel.Fp32MatmulKernel + +/** + * Priority-100 [KernelProvider] backed by the JNI bridge to the shared C + * kernels — the Android counterpart of the JVM's FFM `native-ffm` provider + * (ART has no `java.lang.foreign`, so FFM can never run there). + * + * Availability is probed once: the two-tier loader in [JniKernels] must have + * loaded a library variant AND the smoke kernel must round-trip correctly. + * Every failure mode is non-fatal — the registry then cascades to the + * scalar provider, exactly like the JVM behaves on hosts where the FFM lib + * doesn't load. + * + * Discovered via `META-INF/services` ([JniKernelProviderFactory]) by the + * ServiceLoader install in `skainet-backend-cpu`'s Android ops factory. + */ +public object JniKernelProvider : KernelProvider { + + override val name: String = "native-jni" + + override val priority: Int = 100 + + /** Which library tier actually loaded — for diagnostics/field reports. */ + public val activeVariant: JniKernels.Variant? get() = JniKernels.variant + + private val available: Boolean by lazy { + if (!JniKernels.isLoaded) return@lazy false + runCatching { + val input = floatArrayOf(1.0f, 2.5f, -3.0f) + val output = FloatArray(3) + JniKernels.smoke(input, output, 3) + output[0] == 2.0f && output[1] == 5.0f && output[2] == -6.0f + }.getOrDefault(false) + } + + override fun isAvailable(): Boolean = available + + override fun matmulFp32(): Fp32MatmulKernel? = null // GEMM shim not bridged yet (#920) + + override fun matmulQ8_0(): Q8_0MatmulKernel? = if (available) JniQ8_0Matmul else null + + override fun matmulQ4_0(): Q4_0MatmulKernel? = if (available) JniQ4_0Matmul else null + + override fun matmulQ4K(): Q4KMatmulKernel? = if (available) JniQ4KMatmul else null + + override fun matmulQ5K(): Q5KMatmulKernel? = if (available) JniQ5KMatmul else null + + override fun matmulQ6K(): Q6KMatmulKernel? = if (available) JniQ6KMatmul else null + + private object JniQ8_0Matmul : Q8_0MatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q80Matmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + + private object JniQ4_0Matmul : Q4_0MatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q40Matmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + + private object JniQ4KMatmul : Q4KMatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q4kMatmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + + private object JniQ5KMatmul : Q5KMatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q5kMatmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + + private object JniQ6KMatmul : Q6KMatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q6kMatmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } +} + +/** + * `ServiceLoader`-friendly wrapper around [JniKernelProvider]: the service + * machinery requires a public no-arg constructor, which a Kotlin `object` + * does not expose. Mirrors `NativeKernelProviderFactory` on the JVM side. + * + * Listed in `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider`; + * `consumer-rules.pro` keeps both the entry and this class through R8. + */ +public class JniKernelProviderFactory : KernelProvider by JniKernelProvider 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 new file mode 100644 index 000000000..9575409be --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt @@ -0,0 +1,116 @@ +package sk.ainet.exec.kernel.jni + +import java.io.File + +/** + * JNI surface over the shared C matmul kernels, plus the two-tier library + * loader. + * + * Two `.so` variants are packaged from the same sources (see + * `native/CMakeLists.txt`): + * + * - `libskainet_jni.so` — baseline `armv8-a`. NEON is architecturally + * guaranteed on AArch64, so this runs on every 64-bit ARM core with the + * NEON bodies for fp32/q8_0/q4_0/q5k (q4k/q6k fall back to scalar — + * their SIMD bodies need the dot-product extension). + * - `libskainet_jni_v82.so` — `-march=armv8.2-a+fp16+dotprod`, enabling the + * `vdotq_s32` paths in q4k/q6k. Executing it on an armv8.0 core + * (Cortex-A53, early A55) would SIGILL, so it is only loaded after + * `/proc/cpuinfo` confirms `asimddp` (+ `asimdhp`/`fphp`). + * + * Exactly ONE variant is loaded per process — both export identical JNI + * symbols, and selection-at-load is simpler and safer than symbol renaming + * or ifunc-style per-call dispatch. + * + * Method names deliberately contain no underscores: JNI mangles `_` to + * `_1` in native symbol names, a silent-mismatch trap. + */ +public object JniKernels { + + /** Loaded library variant, or `null` when no variant could be loaded. */ + public val variant: Variant? by lazy { loadVariant() } + + public enum class Variant(public val libName: String) { + /** armv8-a baseline — every AArch64 device, plus x86_64 emulators. */ + BASELINE("skainet_jni"), + + /** armv8.2+dotprod — vdotq_s32 in q4k/q6k; gated on cpuinfo. */ + V82_DOTPROD("skainet_jni_v82"), + } + + /** Whether a kernel library is loaded and callable. */ + public val isLoaded: Boolean get() = variant != null + + private fun loadVariant(): Variant? { + if (cpuSupportsV82()) { + try { + System.loadLibrary(Variant.V82_DOTPROD.libName) + return Variant.V82_DOTPROD + } catch (_: Throwable) { + // Fall through to baseline — e.g. a packaging that stripped + // the v82 lib. Never fatal. + } + } + return try { + System.loadLibrary(Variant.BASELINE.libName) + Variant.BASELINE + } catch (_: Throwable) { + null + } + } + + /** + * `asimddp` is the AArch64 dot-product hwcap; `asimdhp`/`fphp` cover the + * `+fp16` half of the `-march` the v82 lib was compiled with. Reading + * `/proc/cpuinfo` is the classic NDK-sanctioned detection path and needs + * no JNI (which matters: detection must happen BEFORE choosing which + * library to load). Any read failure means "assume baseline" — never + * SIGILL on a weird device, just slower q4k/q6k. + */ + private fun cpuSupportsV82(): Boolean = runCatching { + val features = File("/proc/cpuinfo").useLines { lines -> + lines.firstOrNull { it.startsWith("Features") } + } ?: return false + "asimddp" in features && ("asimdhp" in features || "fphp" in features) + }.getOrDefault(false) + + // --- JNI entry points (skainet_jni.c) --- + + /** `skainet_smoke_double`: output[i] = 2 * input[i]. Used by the availability probe. */ + public external fun smoke(input: FloatArray, output: FloatArray, length: Int) + + public external fun q80Matmul( + 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, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + + public external fun q4kMatmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + + public external fun q5kMatmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + + public external fun q6kMatmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.KernelProvider b/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.KernelProvider new file mode 100644 index 000000000..bd13bc675 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.KernelProvider @@ -0,0 +1 @@ +sk.ainet.exec.kernel.jni.JniKernelProviderFactory diff --git a/skainet-backends/skainet-backend-jni-cpu/src/test/kotlin/sk/ainet/exec/kernel/jni/JniKernelProviderHostTest.kt b/skainet-backends/skainet-backend-jni-cpu/src/test/kotlin/sk/ainet/exec/kernel/jni/JniKernelProviderHostTest.kt new file mode 100644 index 000000000..78238544a --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/test/kotlin/sk/ainet/exec/kernel/jni/JniKernelProviderHostTest.kt @@ -0,0 +1,48 @@ +package sk.ainet.exec.kernel.jni + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertNull +import org.junit.Test + +/** + * Host-JVM behavior contract: on a machine without the Android `.so` + * (every host unit-test run), the provider must degrade gracefully — + * unavailable, all kernels null, and absolutely no exception escaping. + * This mirrors how `NativeKernelProvider` behaves on JVMs where the FFM + * library doesn't load, and it is what the registry's cascade relies on. + */ +class JniKernelProviderHostTest { + + @Test + fun unavailable_without_native_library_and_never_throws() { + assertFalse(JniKernelProvider.isAvailable()) + assertNull(JniKernelProvider.activeVariant) + assertNull(JniKernelProvider.matmulQ8_0()) + assertNull(JniKernelProvider.matmulQ4_0()) + assertNull(JniKernelProvider.matmulQ4K()) + assertNull(JniKernelProvider.matmulQ5K()) + assertNull(JniKernelProvider.matmulQ6K()) + assertNull(JniKernelProvider.matmulFp32()) + } + + @Test + fun provider_contract_name_and_priority() { + assertEquals("native-jni", JniKernelProvider.name) + assertEquals(100, JniKernelProvider.priority) + } + + @Test + fun supports_reports_false_when_unavailable() { + assertFalse(JniKernelProvider.supports("matmul", listOf("Float32", "Q8_0"))) + assertFalse(JniKernelProvider.supports("matmul", listOf("Float32", "Q4_K"))) + } + + @Test + fun serviceloader_factory_delegates() { + val factory = JniKernelProviderFactory() + assertEquals("native-jni", factory.name) + assertEquals(100, factory.priority) + assertFalse(factory.isAvailable()) + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt index d93a0149a..2326f2890 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt @@ -33,12 +33,16 @@ class KernelSupportMatrixTest { // Source-set -> platforms. commonMain reaches all; backend-cpu jvmMain -> {JVM,Android}; // backend-native-cpu jvmMain -> {JVM} (the native module declares only jvm()). + // native-jni: skainet-backend-jni-cpu AAR — same C kernels via JNI, Android + // only, discovered via ServiceLoader from PlatformCpuOpsFactory.android (#920). private fun tiers(): List = listOf( Tier("scalar", 0, platforms.toSet(), scalarFormats()), Tier("panama-vector", 50, setOf("JVM", "Android"), setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0")), Tier("native-ffm", 100, setOf("JVM"), setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q5_K")), + Tier("native-jni", 100, setOf("Android"), + setOf("Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K")), ) private fun best(fmt: String, platform: String, tiers: List): String? =