Repository navigation
perf(native-cpu): NEON register tiles for the dense FP32 SGEMM and x·Wᵀ on AArch64 - #1450
Merged
michalharakal merged 4 commits intoOct 25, 2026
Conversation
…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.
aharakal
approved these changes
Oct 25, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_matmulran an outer-product loop that loaded and stored a row of C for every row of A, andskainet_fp32_matmul_wtcomputed 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 sixteenfloat32x4accumulators. Fourpat a time, eight A loads (four consecutivepof each row) and eight B loads feed sixty-four lane-indexed FMAs (vfmaq_laneq_f32takes 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-twonotherwise 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 % 4finished withvfmaq_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 sharedfinish: lanes overp < (k & ~3),vaddvq, then the scalar tail as explicitfmafsteps. The NEON branch of the old row dot is gone (the tile replaces it); NEON is baseline on AArch64, so no probe.pwhichever 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.### Changedentry.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 (munder, at and past a tile;naround eight;karound 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:linuxX64Test56 tests, 0 failures (incl. both new tests);native-cpu:jvmTest181 tests, 0 failures;arm64-v8aandx86_64with the pinned NDK (buildCMakeRelWithDebInfo); the arm64 object has 136 lane-indexedfmlaand no vector spills in the block function.Not measured
No arm64 hardware in this session, and
qemu-aarch64is 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):and
[gemm-x86]/[wt-x86]lines appear inbuild/test-results/jvmTest/*.xml. For a before/after, check outdevelopfor the baseline run.