Skip to content

feat(backend): skainet-backend-jni-cpu — Android JNI bridge for the NEON kernels (#920) - #943

Merged
michalharakal merged 1 commit into
developfrom
feature/jni-kernel-bridge-920
Aug 10, 2026
Merged

michalharakal merged 1 commit into
developfrom
feature/jni-kernel-bridge-920

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

The Android half of #920: ART has no java.lang.foreign, so the priority-100 FFM provider can never run there — Android inference has been sitting on the priority-0 scalar floor (the measured 1 tok/s cliff from the field report). This AAR ships the same C kernel sources (CMake includes them from skainet-backend-native-cpu/native — single source of truth) with thin JNI shims and a priority-100 KernelProvider.

Runtime CPU-feature dispatch — the Android-specific trap

A Play-store .so cannot assume armv8.2: dotprod instructions SIGILL on Cortex-A53-class cores, but leaving dotprod off wastes the vdotq_s32 q4k/q6k paths on modern SoCs. Two libs are built from the same sources, and JniKernels loads exactly one per process after checking /proc/cpuinfo:

lib -march verified by disassembly
libskainet_jni.so armv8-a baseline 38 fmla, 0 udot/sdot — safe on every arm64 core
libskainet_jni_v82.so armv8.2-a+fp16+dotprod 32 udot/sdot sites

Selection-at-load beats symbol renaming / ifunc for simplicity and auditability; the loaded tier is exposed (JniKernelProvider.activeVariant) so field reports name the actual code path.

Details

  • Shims: GetPrimitiveArrayCritical (zero-copy pins on ART), no JNI calls inside the critical window, reverse-order releases, JNI_ABORT for read-only arrays; method names contain no underscores (JNI mangles _ → _1). Q8_0/Q4_0/Q4_K/Q5_K/Q6_K bridged; fp32 GEMM deferred (different ABI shape, not the LLM decode hot path).
  • Discovery: ART supports java.util.ServiceLoader — PlatformCpuOpsFactory.android now installs discovered providers exactly like the JVM does, then the scalar floor. consumer-rules.pro keeps the service entry + native method names through R8 full mode (without it, release builds silently fall back to scalar — the worst failure mode).
  • 16 KB pages: .sos linked with max-page-size=16384 (Android 15+ requirement); LOAD align 0x4000 verified via llvm-readelf.
  • ABIs: arm64-v8a + x86_64 (emulator CI); 32-bit ARM deliberately out of scope.
  • KernelSupportMatrixTest tier list gains native-jni (Android, priority 100).

Verified locally

  • assembleRelease: AAR contains all four .sos (arm64 pair differs — different codegen; x86_64 pair identical by design, no per-ABI loader special case)
  • testDebugUnitTest 4/4 — host graceful-degradation contract (no .so → unavailable, kernels null, nothing throws)
  • assembleDebugAndroidTest compiles the on-device parity suite (vs scalar references)
  • :skainet-backends:skainet-backend-cpu:compileAndroidMain + KernelSupportMatrixTest green

Not yet done (next steps, flagged honestly)

Refs #920

…EON kernels

ART has no java.lang.foreign, so the priority-100 FFM provider can never
run on Android; inference there ran on the priority-0 scalar floor (the
1 tok/s cliff measured in the #920 field report). This AAR ships the
same C kernel sources (single source of truth — the CMakeLists includes
them from skainet-backend-native-cpu/native) via NDK with thin JNI
shims and a priority-100 KernelProvider.

Runtime CPU-feature dispatch, the Android-specific correctness trap: a
Play-store .so cannot assume armv8.2 (dotprod SIGILLs on Cortex-A53).
Two libs are built from the same sources — libskainet_jni.so (baseline
armv8-a: NEON guaranteed on AArch64, 0 dot-product instructions,
verified by llvm-objdump) and libskainet_jni_v82.so
(-march=armv8.2-a+fp16+dotprod: 32 udot/sdot sites) — and JniKernels
loads exactly ONE per process after checking asimddp/fphp in
/proc/cpuinfo. Selection-at-load beats symbol renaming or ifunc for
simplicity and auditability.

Details:
- JNI shims use GetPrimitiveArrayCritical (zero-copy pins on ART), no
  JNI calls inside the critical window, releases in reverse order,
  JNI_ABORT for read-only arrays. Method names contain no underscores
  (JNI mangles _ to _1 — silent-mismatch trap).
- Q8_0/Q4_0/Q4_K/Q5_K/Q6_K bridged; fp32 GEMM deferred (different ABI
  shape, dense fp32 is not the LLM decode hot path).
- ServiceLoader discovery: ART supports java.util.ServiceLoader;
  PlatformCpuOpsFactory.android now installs discovered providers like
  the JVM does, then registers the scalar floor. consumer-rules.pro
  keeps the service entry and native method names through R8 full mode
  (without it, release builds silently fall back to scalar).
- .so files are 16 KB-page aligned (-Wl,-z,max-page-size=16384,
  Android 15+ requirement; LOAD align 0x4000 verified via llvm-readelf).
- Host unit tests pin the graceful-degradation contract (no .so ->
  unavailable, kernels null, nothing throws). On-device androidTest
  parity suite (vs scalar references) compiles and is ready for a
  device run; the KernelSupportMatrix tier list gains native-jni.

Verified locally: assembleRelease green (AAR contains 4 .so files:
arm64 baseline+v82 differ, x86_64 pair identical by design),
testDebugUnitTest 4/4, assembleDebugAndroidTest compiles,
backend-cpu compileAndroidMain green, KernelSupportMatrixTest green.
Device execution of the parity suite + tok/s measurement on the
PromptPong reproducer is the next validation step (physical bench).

Refs #920
@michalharakal
michalharakal requested a review from aharakal August 10, 2026 12:57
@michalharakal
michalharakal merged commit f8d9f5a into develop Aug 10, 2026
16 checks passed
@michalharakal
michalharakal deleted the feature/jni-kernel-bridge-920 branch August 10, 2026 13:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants