Skip to content

perf(native-cpu): NEON register tiles for the dense FP32 SGEMM and x·Wᵀ on AArch64 - #1450

Merged
michalharakal merged 4 commits into
SKaiNET-developers:developfrom
michalharakal:feat/neon-fp32-tiles
Oct 25, 2026
Merged

michalharakal merged 4 commits into
SKaiNET-developers:developfrom
michalharakal:feat/neon-fp32-tiles

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Summary

The AArch64 bodies of the two dense FP32 kernels get the register tiles x86-64 got in #1442 and #1445. Until now the NEON lane was the x86 story of a week ago: skainet_fp32_matmul ran an outer-product loop that loaded and stored a row of C for every row of A, and skainet_fp32_matmul_wt computed each output as one four-lane chain, which leaves most of an arm64 core's FMA pipes idle. This serves Android arm64 (the JNI library), iOS and macOS Kotlin/Native, Linux arm64, and the JVM on a Mac wherever the product does not go to Accelerate (#1449).

Stacked on #1445 (same files); the first two commits are that PR's and drop out when it merges. C only, no API change.

What changes

  • fp32_matmul.c, NEON block: an 8 × 8 tile of C in sixteen float32x4 accumulators. Four p at a time, eight A loads (four consecutive p of each row) and eight B loads feed sixty-four lane-indexed FMAs (vfmaq_laneq_f32 takes the a-value from a lane, no broadcast register — 136 of them in the object, no vector spills). The same 256-deep k-blocking and packed 8-column B panel as the AVX2 lane (a power-of-two n otherwise maps every panel row to one L1 set), 4 × 8 and 1 × 8 tiles for the rows left under eight, a zero-padded panel plus C stripe for the columns left under eight, k % 4 finished with vfmaq_n_f32. The portable i-p-j body remains for targets with neither lane.
  • fp32_matmul_wt.c, skainet_fp32_wt_rows_neon: four weight rows by three activation rows in twelve four-lane chains, a 4 × 1 tile for the activation rows left over, a single-row dot for the weight rows left over. Every variant reduces an output through one shared finish: lanes over p < (k & ~3), vaddvq, then the scalar tail as explicit fmaf steps. The NEON branch of the old row dot is gone (the tile replaces it); NEON is baseline on AArch64, so no probe.
  • Numerics: every output takes its fused steps in ascending p whichever tile holds it, so sliced, threaded and whole calls stay bit-identical to each other — what the row pool and the SKEEP-005 column split rely on. Against the previous NEON numbers an output moves within rounding.
  • KDoc of both JVM wrappers names the AArch64 tiles; CHANGELOG ### Changed entry.

Tests

Two new Kotlin/Native tests call the C kernels through cinterop, so they run on every native target that embeds the archive — on the macOS CI leg (macosArm64Test) that is Apple silicon running the NEON lane:

  • NativeKnFp32MatmulTileTest: tile edges (m under, at and past a tile; n around eight; k around the k-block and off a multiple of four) against a double reference; strided and offset operands with the padding untouched; block independence bit-for-bit (rows split, columns split, row at a time, packing on vs off).
  • NativeKnFp32MatmulWtTileTest: tile edges against the row-dot reference; slice at weight row 3 and row-at-a-time bit-for-bit against the whole product, including a shape the row pool threads.

On linuxX64 the same tests exercise the AVX2 lane. Run here:

  • native-cpu:linuxX64Test 56 tests, 0 failures (incl. both new tests); native-cpu:jvmTest 181 tests, 0 failures;
  • Android kernel library built for arm64-v8a and x86_64 with the pinned NDK (buildCMakeRelWithDebInfo); the arm64 object has 136 lane-indexed fmla and no vector spills in the block function.

Not measured

No arm64 hardware in this session, and qemu-aarch64 is not installed. Correctness on Apple silicon comes from the macOS CI leg; speed numbers need a Mac or an Android device. On a Mac, the JVM bench prints the native call per shape (the AVX2 names are historical; both tests run whatever lane the host has):

./gradlew :skainet-backends:skainet-backend-native-cpu:jvmTest \
  --tests '*NativeFp32MatmulX86LaneTest*' --tests '*NativeFp32MatmulWtX86LaneTest*' \
  -Dskainet.runBench=true

and [gemm-x86] / [wt-x86] lines appear in build/test-results/jvmTest/*.xml. For a before/after, check out develop for the baseline run.

…Wᵀ on AArch64

The NEON bodies of skainet_fp32_matmul and skainet_fp32_matmul_wt were an
outer-product loop that loaded and stored a row of C for every row of A,
and a single four-lane chain per output that left most of a core's FMA
pipes idle. The SGEMM now runs an 8×8 tile in sixteen float32x4
accumulators with lane-indexed FMAs (vfmaq_laneq_f32) over a packed
8-column panel of B in a 256-deep k-block, with 4×8 and 1×8 tiles for the
rows left over and a zero-padded panel plus C stripe for the columns left
over; x·Wᵀ runs a tile of four weight rows by three activation rows in
twelve chains, every variant reducing through one shared finish. Every
output takes its fused steps in ascending p whichever tile holds it, so
sliced, threaded and whole calls stay bit-identical to each other.

Serves Android arm64, iOS and macOS Kotlin/Native, Linux arm64 and the
JVM on Apple silicon where a product does not go to the platform BLAS.

Two Kotlin/Native tests run the tails, strides and block independence
through cinterop on every native target, so the macOS CI leg exercises
the lane on Apple silicon; on linuxX64 they cover the AVX2 lane.
@michalharakal
michalharakal merged commit 903814e into SKaiNET-developers:develop Oct 25, 2026
17 checks passed
@michalharakal
michalharakal deleted the feat/neon-fp32-tiles branch October 25, 2026 18:40
@michalharakal michalharakal mentioned this pull request Oct 25, 2026
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