From 8595b3d2630839cecd9c103b29d65fb6baf60ba7 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sun, 25 Oct 2026 14:55:00 +0100 Subject: [PATCH 1/2] =?UTF-8?q?perf(native-cpu):=20NEON=20register=20tiles?= =?UTF-8?q?=20for=20the=20dense=20FP32=20SGEMM=20and=20x=C2=B7W=E1=B5=80?= =?UTF-8?q?=20on=20AArch64?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 1 + .../native/src/fp32_matmul.c | 286 ++++++++++++++++-- .../native/src/fp32_matmul_wt.c | 113 ++++++- .../exec/kernel/NativeFp32MatmulKernel.kt | 9 +- .../exec/kernel/NativeFp32MatmulWtKernel.kt | 7 +- .../exec/kernel/NativeKnFp32MatmulTileTest.kt | 137 +++++++++ .../kernel/NativeKnFp32MatmulWtTileTest.kt | 98 ++++++ 7 files changed, 618 insertions(+), 33 deletions(-) create mode 100644 skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulTileTest.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulWtTileTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e268063..4458e2f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Changed +- **Dense FP32 kernels get register tiles on AArch64** (`skainet-backend-native-cpu`, Apple silicon, Android arm64, Linux arm64): the NEON bodies of `skainet_fp32_matmul` and `skainet_fp32_matmul_wt` were the x86 story of a week ago — an outer-product loop that loaded and stored a row of C for every row of A, and a single four-lane chain per `x · Wᵀ` output that left most of a core's FMA pipes idle. The plain SGEMM now runs an 8 × 8 tile in sixteen `float32x4` accumulators with lane-indexed FMAs over a packed 8-column panel of B in a 256-deep k-block (4 × 8 and 1 × 8 tiles for the rows left over, a zero-padded panel and C stripe for the columns left over), and `x · Wᵀ` a tile of four weight rows by three activation rows in twelve chains. 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; against the previous NEON numbers an output moves within rounding. Two Kotlin/Native tests run the tails and the block independence through cinterop on every native target, so the macOS CI leg exercises the lane on Apple silicon. No numbers from this change's author: measured on an arm64 host it is not — see the pull request for how to run the bench. - **Dense FP32 SGEMM uses AVX2/FMA on x86-64** (`skainet-backend-native-cpu`): `skainet_fp32_matmul`, the kernel behind every `A · B` product the JVM hands to the native lane (a weight stored `[in, out]`, attention scores and values, training), had register-tiled bodies for wasm SIMD128 only; the x86-64 library is built for baseline x86-64, so the kernel ran a scalar/SSE2 outer-product loop without FMA that loaded and stored a row of C for every row of A. Hosts with AVX2 and FMA now select, through the cached `cpuid` probe the `x · Wᵀ` kernel introduced, a lane compiled per function for that target: a 6 × 16 tile of C in twelve eight-lane FMA accumulators over a 256-deep k-block, each sixteen-column panel of B packed once so a power-of-two `n` no longer maps every row of the panel to one L1 set, 2 × 16 and 1 × 16 tiles for the rows left over and a lane-masked stripe for the columns left over. Every output takes one fused multiply-add per `p` in ascending `p` whichever tile holds it, so sliced, threaded and whole calls stay bit-identical to each other; against the previous x86 numbers an output moves within rounding, as the NEON and wasm lanes already do. Machines without AVX2, and every other architecture, run the previous code. Measured single call on a six-core laptop, develop → this change: 63 × 1024 × 1024 21 → 281 GFLOP/s, 1024³ 48 → 311, 512³ 59 → 358, 192 × 192 × 64 (one thread) 18 → 99. ### Fixed diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c index 73acde79..8568f3f7 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c +++ b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul.c @@ -16,10 +16,14 @@ * `aStride == k`, `bStride == n`, `cStride == n`. Sub-block scenarios * pass larger strides and the corresponding offsets. * - * Portable and NEON bodies: iteration order is i-p-j (outer product into - * rows of C). The inner loop is `c[j] += a_ip * b[j]` over a contiguous run - * of `n` floats for both b's row and c's row — auto-vectorizes cleanly under - * -O3 -ffast-math into vfmadd231ps / fmla. + * Portable body: iteration order is i-p-j (outer product into rows of C). + * The inner loop is `c[j] += a_ip * b[j]` over a contiguous run of `n` + * floats for both b's row and c's row — auto-vectorizes under -O3 + * -ffast-math. It is what a target without one of the lanes below runs. + * + * AArch64 body (Apple silicon, Android arm64, Linux arm64): an 8×8 register + * tile in sixteen float32x4 accumulators with lane-indexed FMAs, the packed + * 8-column B panel and the tails of the AVX2 lane — see the NEON block. * * SIMD128 body (the Emscripten sidecar): a 4×8 register-blocked micro-kernel. * Eight f32x4 accumulators hold a 4-row × 8-column tile of C for a whole @@ -155,7 +159,266 @@ static void skainet_fp32_block(const skainet_fp32_gemm_ctx* c, int32_t i0, int32 } } -#else /* portable + NEON: i-p-j outer product over the block */ +#elif defined(SKAINET_HAVE_NEON) + +/* + * AArch64 body: the same shape as the AVX2 lane below, sized for NEON's + * 32 four-lane registers. An 8-row × 8-column tile of C lives in sixteen + * float32x4 accumulators; four p-steps 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 straight from a lane, no + * broadcast register). The same 256-deep k-blocking, the same packed + * 8-column B panel (a power-of-two n otherwise maps every panel row to one + * L1 set), 4×8 and 1×8 tiles for the rows left over and a zero-padded panel + * plus C stripe for the columns left over. Every element takes one fused + * step per p in ascending p whichever tile holds it. + */ + +#define SKAINET_FP32_NEON_K_BLOCK 256 + +/* One p-step of the 8×8 tile with the a-values in lane LANE of a0..a7. */ +#define SKAINET_FP32_NEON_8X8_STEP(LANE) \ + bl = vld1q_f32(b_row); bh = vld1q_f32(b_row + 4); b_row += b_stride; \ + s00 = vfmaq_laneq_f32(s00, bl, a0, LANE); s01 = vfmaq_laneq_f32(s01, bh, a0, LANE); \ + s10 = vfmaq_laneq_f32(s10, bl, a1, LANE); s11 = vfmaq_laneq_f32(s11, bh, a1, LANE); \ + s20 = vfmaq_laneq_f32(s20, bl, a2, LANE); s21 = vfmaq_laneq_f32(s21, bh, a2, LANE); \ + s30 = vfmaq_laneq_f32(s30, bl, a3, LANE); s31 = vfmaq_laneq_f32(s31, bh, a3, LANE); \ + s40 = vfmaq_laneq_f32(s40, bl, a4, LANE); s41 = vfmaq_laneq_f32(s41, bh, a4, LANE); \ + s50 = vfmaq_laneq_f32(s50, bl, a5, LANE); s51 = vfmaq_laneq_f32(s51, bh, a5, LANE); \ + s60 = vfmaq_laneq_f32(s60, bl, a6, LANE); s61 = vfmaq_laneq_f32(s61, bh, a6, LANE); \ + s70 = vfmaq_laneq_f32(s70, bl, a7, LANE); s71 = vfmaq_laneq_f32(s71, bh, a7, LANE) + +/* + * C[8 × 8] over p in [p0, p1). `a` is the first of the tile's rows of A, `b_row` + * row p0 of an 8-column panel of B (the packed copy, or B itself with its + * stride), `c0` the first of the tile's rows of C; accumulators are loaded from + * C and stored back, so a tile continues the sum carried from the previous + * k-block. + */ +static inline void skainet_fp32_neon_tile_8x8( + const float* a, int32_t a_stride, const float* b_row, int32_t b_stride, + float* c0, int32_t c_stride, int32_t p0, int32_t p1 +) { + const float* ar0 = a; + const float* ar1 = ar0 + a_stride; + const float* ar2 = ar1 + a_stride; + const float* ar3 = ar2 + a_stride; + const float* ar4 = ar3 + a_stride; + const float* ar5 = ar4 + a_stride; + const float* ar6 = ar5 + a_stride; + const float* ar7 = ar6 + a_stride; + float* c1 = c0 + c_stride; + float* c2 = c1 + c_stride; + float* c3 = c2 + c_stride; + float* c4 = c3 + c_stride; + float* c5 = c4 + c_stride; + float* c6 = c5 + c_stride; + float* c7 = c6 + c_stride; + float32x4_t s00 = vld1q_f32(c0), s01 = vld1q_f32(c0 + 4); + float32x4_t s10 = vld1q_f32(c1), s11 = vld1q_f32(c1 + 4); + float32x4_t s20 = vld1q_f32(c2), s21 = vld1q_f32(c2 + 4); + float32x4_t s30 = vld1q_f32(c3), s31 = vld1q_f32(c3 + 4); + float32x4_t s40 = vld1q_f32(c4), s41 = vld1q_f32(c4 + 4); + float32x4_t s50 = vld1q_f32(c5), s51 = vld1q_f32(c5 + 4); + float32x4_t s60 = vld1q_f32(c6), s61 = vld1q_f32(c6 + 4); + float32x4_t s70 = vld1q_f32(c7), s71 = vld1q_f32(c7 + 4); + float32x4_t bl, bh; + int32_t p = p0; + for (; p + 4 <= p1; p += 4) { + const float32x4_t a0 = vld1q_f32(ar0 + p), a1 = vld1q_f32(ar1 + p); + const float32x4_t a2 = vld1q_f32(ar2 + p), a3 = vld1q_f32(ar3 + p); + const float32x4_t a4 = vld1q_f32(ar4 + p), a5 = vld1q_f32(ar5 + p); + const float32x4_t a6 = vld1q_f32(ar6 + p), a7 = vld1q_f32(ar7 + p); + SKAINET_FP32_NEON_8X8_STEP(0); + SKAINET_FP32_NEON_8X8_STEP(1); + SKAINET_FP32_NEON_8X8_STEP(2); + SKAINET_FP32_NEON_8X8_STEP(3); + } + for (; p < p1; ++p) { + bl = vld1q_f32(b_row); bh = vld1q_f32(b_row + 4); b_row += b_stride; + s00 = vfmaq_n_f32(s00, bl, ar0[p]); s01 = vfmaq_n_f32(s01, bh, ar0[p]); + s10 = vfmaq_n_f32(s10, bl, ar1[p]); s11 = vfmaq_n_f32(s11, bh, ar1[p]); + s20 = vfmaq_n_f32(s20, bl, ar2[p]); s21 = vfmaq_n_f32(s21, bh, ar2[p]); + s30 = vfmaq_n_f32(s30, bl, ar3[p]); s31 = vfmaq_n_f32(s31, bh, ar3[p]); + s40 = vfmaq_n_f32(s40, bl, ar4[p]); s41 = vfmaq_n_f32(s41, bh, ar4[p]); + s50 = vfmaq_n_f32(s50, bl, ar5[p]); s51 = vfmaq_n_f32(s51, bh, ar5[p]); + s60 = vfmaq_n_f32(s60, bl, ar6[p]); s61 = vfmaq_n_f32(s61, bh, ar6[p]); + s70 = vfmaq_n_f32(s70, bl, ar7[p]); s71 = vfmaq_n_f32(s71, bh, ar7[p]); + } + vst1q_f32(c0, s00); vst1q_f32(c0 + 4, s01); + vst1q_f32(c1, s10); vst1q_f32(c1 + 4, s11); + vst1q_f32(c2, s20); vst1q_f32(c2 + 4, s21); + vst1q_f32(c3, s30); vst1q_f32(c3 + 4, s31); + vst1q_f32(c4, s40); vst1q_f32(c4 + 4, s41); + vst1q_f32(c5, s50); vst1q_f32(c5 + 4, s51); + vst1q_f32(c6, s60); vst1q_f32(c6 + 4, s61); + vst1q_f32(c7, s70); vst1q_f32(c7 + 4, s71); +} + +#define SKAINET_FP32_NEON_4X8_STEP(LANE) \ + bl = vld1q_f32(b_row); bh = vld1q_f32(b_row + 4); b_row += b_stride; \ + s00 = vfmaq_laneq_f32(s00, bl, a0, LANE); s01 = vfmaq_laneq_f32(s01, bh, a0, LANE); \ + s10 = vfmaq_laneq_f32(s10, bl, a1, LANE); s11 = vfmaq_laneq_f32(s11, bh, a1, LANE); \ + s20 = vfmaq_laneq_f32(s20, bl, a2, LANE); s21 = vfmaq_laneq_f32(s21, bh, a2, LANE); \ + s30 = vfmaq_laneq_f32(s30, bl, a3, LANE); s31 = vfmaq_laneq_f32(s31, bh, a3, LANE) + +/* C[4 × 8]: the rows left under eight, four at a time. */ +static inline void skainet_fp32_neon_tile_4x8( + const float* a, int32_t a_stride, const float* b_row, int32_t b_stride, + float* c0, int32_t c_stride, int32_t p0, int32_t p1 +) { + const float* ar0 = a; + const float* ar1 = ar0 + a_stride; + const float* ar2 = ar1 + a_stride; + const float* ar3 = ar2 + a_stride; + float* c1 = c0 + c_stride; + float* c2 = c1 + c_stride; + float* c3 = c2 + c_stride; + float32x4_t s00 = vld1q_f32(c0), s01 = vld1q_f32(c0 + 4); + float32x4_t s10 = vld1q_f32(c1), s11 = vld1q_f32(c1 + 4); + float32x4_t s20 = vld1q_f32(c2), s21 = vld1q_f32(c2 + 4); + float32x4_t s30 = vld1q_f32(c3), s31 = vld1q_f32(c3 + 4); + float32x4_t bl, bh; + int32_t p = p0; + for (; p + 4 <= p1; p += 4) { + const float32x4_t a0 = vld1q_f32(ar0 + p), a1 = vld1q_f32(ar1 + p); + const float32x4_t a2 = vld1q_f32(ar2 + p), a3 = vld1q_f32(ar3 + p); + SKAINET_FP32_NEON_4X8_STEP(0); + SKAINET_FP32_NEON_4X8_STEP(1); + SKAINET_FP32_NEON_4X8_STEP(2); + SKAINET_FP32_NEON_4X8_STEP(3); + } + for (; p < p1; ++p) { + bl = vld1q_f32(b_row); bh = vld1q_f32(b_row + 4); b_row += b_stride; + s00 = vfmaq_n_f32(s00, bl, ar0[p]); s01 = vfmaq_n_f32(s01, bh, ar0[p]); + s10 = vfmaq_n_f32(s10, bl, ar1[p]); s11 = vfmaq_n_f32(s11, bh, ar1[p]); + s20 = vfmaq_n_f32(s20, bl, ar2[p]); s21 = vfmaq_n_f32(s21, bh, ar2[p]); + s30 = vfmaq_n_f32(s30, bl, ar3[p]); s31 = vfmaq_n_f32(s31, bh, ar3[p]); + } + vst1q_f32(c0, s00); vst1q_f32(c0 + 4, s01); + vst1q_f32(c1, s10); vst1q_f32(c1 + 4, s11); + vst1q_f32(c2, s20); vst1q_f32(c2 + 4, s21); + vst1q_f32(c3, s30); vst1q_f32(c3 + 4, s31); +} + +/* C[1 × 8]: the last rows, one at a time. */ +static inline void skainet_fp32_neon_tile_1x8( + const float* a, const float* b_row, int32_t b_stride, float* c0, int32_t p0, int32_t p1 +) { + float32x4_t s0 = vld1q_f32(c0), s1 = vld1q_f32(c0 + 4); + int32_t p = p0; + for (; p + 4 <= p1; p += 4) { + const float32x4_t av = vld1q_f32(a + p); + s0 = vfmaq_laneq_f32(s0, vld1q_f32(b_row), av, 0); s1 = vfmaq_laneq_f32(s1, vld1q_f32(b_row + 4), av, 0); b_row += b_stride; + s0 = vfmaq_laneq_f32(s0, vld1q_f32(b_row), av, 1); s1 = vfmaq_laneq_f32(s1, vld1q_f32(b_row + 4), av, 1); b_row += b_stride; + s0 = vfmaq_laneq_f32(s0, vld1q_f32(b_row), av, 2); s1 = vfmaq_laneq_f32(s1, vld1q_f32(b_row + 4), av, 2); b_row += b_stride; + s0 = vfmaq_laneq_f32(s0, vld1q_f32(b_row), av, 3); s1 = vfmaq_laneq_f32(s1, vld1q_f32(b_row + 4), av, 3); b_row += b_stride; + } + for (; p < p1; ++p, b_row += b_stride) { + s0 = vfmaq_n_f32(s0, vld1q_f32(b_row), a[p]); s1 = vfmaq_n_f32(s1, vld1q_f32(b_row + 4), a[p]); + } + vst1q_f32(c0, s0); vst1q_f32(c0 + 4, s1); +} + +/* Rows [0, rows) of A against the panel `b_row` into C rows `c0` (both 8 wide): 8×8 tiles, then 4×8, then 1×8. */ +static inline void skainet_fp32_neon_column( + const float* a, int32_t a_stride, int32_t rows, const float* b_row, int32_t b_stride, + float* c0, int32_t c_stride, int32_t p0, int32_t p1 +) { + int32_t i = 0; + for (; i + 8 <= rows; i += 8) { + skainet_fp32_neon_tile_8x8(a + (size_t) i * a_stride, a_stride, b_row, b_stride, c0 + (size_t) i * c_stride, c_stride, p0, p1); + } + for (; i + 4 <= rows; i += 4) { + skainet_fp32_neon_tile_4x8(a + (size_t) i * a_stride, a_stride, b_row, b_stride, c0 + (size_t) i * c_stride, c_stride, p0, p1); + } + for (; i < rows; ++i) { + skainet_fp32_neon_tile_1x8(a + (size_t) i * a_stride, b_row, b_stride, c0 + (size_t) i * c_stride, p0, p1); + } +} + +/* + * B[p0..p1, j..j+w) copied into a contiguous 8-wide panel, columns at or past + * w (the column tail, w < 8) zero-filled. See the AVX2 lane for why the copy + * pays: in place, a power-of-two n puts every row of the panel in one L1 set. + */ +static inline void skainet_fp32_neon_pack_panel( + float* panel, const float* b, int32_t b_stride, int32_t w, int32_t p0, int32_t p1 +) { + const float* b_row = b + (size_t) p0 * b_stride; + if (w == 8) { + for (int32_t p = p0; p < p1; ++p, b_row += b_stride, panel += 8) { + vst1q_f32(panel, vld1q_f32(b_row)); + vst1q_f32(panel + 4, vld1q_f32(b_row + 4)); + } + } else { + for (int32_t p = p0; p < p1; ++p, b_row += b_stride, panel += 8) { + int32_t j = 0; + for (; j < w; ++j) panel[j] = b_row[j]; + for (; j < 8; ++j) panel[j] = 0.0f; + } + } +} + +/* + * The columns left under eight, from j, for rows [i0, i1): the C block is + * copied into an 8-wide stripe, the full tiles run over the stripe and the + * zero-padded tail panel, and the live columns are stored back — the same + * fused steps a wider block would give them. + */ +static void skainet_fp32_neon_tail( + const skainet_fp32_gemm_ctx* c, int32_t i0, int32_t i1, int32_t j, int32_t w, + const float* panel, int32_t p0, int32_t p1 +) { + float stripe[8 * 8]; + for (int32_t i = i0; i < i1; i += 8) { + const int32_t rows = (i1 - i) < 8 ? (i1 - i) : 8; + float* c_row = c->c + (size_t) i * c->c_stride + j; + for (int32_t r = 0; r < rows; ++r) { + int32_t jj = 0; + for (; jj < w; ++jj) stripe[r * 8 + jj] = c_row[(size_t) r * c->c_stride + jj]; + for (; jj < 8; ++jj) stripe[r * 8 + jj] = 0.0f; + } + skainet_fp32_neon_column(c->a + (size_t) i * c->a_stride, c->a_stride, rows, panel, 8, stripe, 8, p0, p1); + for (int32_t r = 0; r < rows; ++r) { + for (int32_t jj = 0; jj < w; ++jj) c_row[(size_t) r * c->c_stride + jj] = stripe[r * 8 + jj]; + } + } +} + +/* Rows [i0, i1) × columns [j0, j1) of C, k-blocked: the 8-column panels first, the column tail after. */ +static void skainet_fp32_block(const skainet_fp32_gemm_ctx* c, int32_t i0, int32_t i1, int32_t j0, int32_t j1) { + for (int32_t i = i0; i < i1; ++i) { + float* c_row = c->c + (size_t) i * c->c_stride; + for (int32_t j = j0; j < j1; ++j) c_row[j] = 0.0f; + } + if (c->k <= 0) return; + const int32_t rows = i1 - i0; + const int32_t j_tiles_end = j0 + ((j1 - j0) / 8) * 8; + const int32_t j_tail = j1 - j_tiles_end; + /* The packed copy pays for itself from two row tiles up; a thinner block reads B in place. */ + const int pack = rows >= 16; + float panel[SKAINET_FP32_NEON_K_BLOCK * 8]; + const float* a = c->a + (size_t) i0 * c->a_stride; + for (int32_t p0 = 0; p0 < c->k; p0 += SKAINET_FP32_NEON_K_BLOCK) { + const int32_t p1 = (c->k - p0) < SKAINET_FP32_NEON_K_BLOCK ? c->k : p0 + SKAINET_FP32_NEON_K_BLOCK; + for (int32_t j = j0; j < j_tiles_end; j += 8) { + float* c0 = c->c + (size_t) i0 * c->c_stride + j; + if (pack) { + skainet_fp32_neon_pack_panel(panel, c->b + j, c->b_stride, 8, p0, p1); + skainet_fp32_neon_column(a, c->a_stride, rows, panel, 8, c0, c->c_stride, p0, p1); + } else { + skainet_fp32_neon_column(a, c->a_stride, rows, c->b + (size_t) p0 * c->b_stride + j, c->b_stride, c0, c->c_stride, p0, p1); + } + } + if (j_tail > 0) { + skainet_fp32_neon_pack_panel(panel, c->b + j_tiles_end, c->b_stride, j_tail, p0, p1); + skainet_fp32_neon_tail(c, i0, i1, j_tiles_end, j_tail, panel, p0, p1); + } + } +} + +#else /* portable: i-p-j outer product over the block */ static void skainet_fp32_block(const skainet_fp32_gemm_ctx* c, int32_t i0, int32_t i1, int32_t j0, int32_t j1) { for (int32_t i = i0; i < i1; ++i) { @@ -169,22 +432,9 @@ static void skainet_fp32_block(const skainet_fp32_gemm_ctx* c, int32_t i0, int32 for (int32_t p = 0; p < c->k; ++p) { const float a_ip = a_row[p]; const float* SKAINET_RESTRICT b_row = c->b + (size_t) p * c->b_stride; -#ifdef SKAINET_HAVE_NEON - const float32x4_t va = vdupq_n_f32(a_ip); - int32_t j = j0; - for (; j + 4 <= j1; j += 4) { - float32x4_t cv = vld1q_f32(c_row + j); - cv = vfmaq_f32(cv, va, vld1q_f32(b_row + j)); - vst1q_f32(c_row + j, cv); - } - for (; j < j1; ++j) { - c_row[j] += a_ip * b_row[j]; - } -#else for (int32_t j = j0; j < j1; ++j) { c_row[j] += a_ip * b_row[j]; } -#endif } } } diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul_wt.c b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul_wt.c index 23964070..50993ff0 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul_wt.c +++ b/skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul_wt.c @@ -2,6 +2,7 @@ #include "skainet_simd.h" #include "skainet_row_threads.h" #include "skainet_cpu_features.h" +#include #include #include @@ -15,8 +16,8 @@ * its weight in is the shape this kernel wants. * * Numerics: the order in which an output accumulates its products is fixed per - * build (the portable loop, the four-lane wasm/NEON sums, or the eight-lane AVX2 - * tile below) and never depends on which thread or which row slice computes + * build (the portable loop, the four-lane wasm sum, the four-lane NEON tile or + * the eight-lane AVX2 tile below) and never depends on which thread or which row slice computes * it, so a threaded run, a sliced run (SKEEP-005 schedules split the weight's * rows across calls) and a single call are bit-identical to each other. Lanes * differ between builds the way they already do between NEON and wasm. @@ -25,6 +26,10 @@ * enough, through skainet_row_threads.h — on Emscripten that is the pthread * pool when the page is cross-origin isolated and a plain call otherwise. * + * AArch64 (Apple silicon, Android and Linux arm64): skainet_fp32_wt_rows_neon, + * four weight rows by three activation rows in twelve float32x4 chains; NEON + * is baseline there, so no probe. + * * x86-64: the shared library is built without -march (it has to run on any * x86-64), so the portable loop is SSE2 code. Hosts with AVX2+FMA — every * x86 CPU since 2013 — take skainet_fp32_wt_rows_avx2, compiled per function @@ -50,12 +55,6 @@ static inline float skainet_dot_f32(const float* SKAINET_RESTRICT a, const float s = wasm_f32x4_extract_lane(acc, 0) + wasm_f32x4_extract_lane(acc, 1) + wasm_f32x4_extract_lane(acc, 2) + wasm_f32x4_extract_lane(acc, 3); for (; p < k; ++p) s += a[p] * b[p]; -#elif defined(SKAINET_HAVE_NEON) - float32x4_t acc = vdupq_n_f32(0.0f); - int32_t p = 0; - for (; p + 4 <= k; p += 4) acc = vfmaq_f32(acc, vld1q_f32(a + p), vld1q_f32(b + p)); - s = skainet_neon_hadd_f32(acc); - for (; p < k; ++p) s += a[p] * b[p]; #else for (int32_t p = 0; p < k; ++p) s += a[p] * b[p]; #endif @@ -73,6 +72,101 @@ static void skainet_fp32_wt_rows(void* vctx, int32_t o_start, int32_t o_end) { } } +#ifdef SKAINET_HAVE_NEON +/* + * AArch64: four weight rows against activation rows three at a time — twelve + * four-lane FMA chains per tile, so the core's FMA pipes see twelve + * independent sums instead of the one chain a row dot gives them, and the + * 4 × k floats of weight stay in the core's cache while the activation rows + * stream past. Every variant reduces an output the same way: the four-lane + * chain over p < (k & ~3), vaddvq, then the scalar tail as fused steps in + * ascending p — so the result does not depend on the tile an output lands in. + */ +static inline float skainet_neon_finish(float32x4_t acc, const float* a, const float* b, int32_t k4, int32_t k) { + float s = skainet_neon_hadd_f32(acc); + for (int32_t p = k4; p < k; ++p) s = fmaf(a[p], b[p], s); + return s; +} + +static inline float skainet_neon_dot(const float* a, const float* b, int32_t k4, int32_t k) { + float32x4_t acc = vdupq_n_f32(0.0f); + for (int32_t p = 0; p < k4; p += 4) acc = vfmaq_f32(acc, vld1q_f32(a + p), vld1q_f32(b + p)); + return skainet_neon_finish(acc, a, b, k4, k); +} + +static void skainet_fp32_wt_rows_neon(void* vctx, int32_t o_start, int32_t o_end) { + const skainet_fp32_wt_ctx* c = (const skainet_fp32_wt_ctx*) vctx; + const int32_t m = c->m, n = c->n, k = c->k; + const int32_t k4 = k & ~3; + int32_t o = o_start; + for (; o + 4 <= o_end; o += 4) { + const float* w0 = c->w + (size_t) o * k; + const float* w1 = w0 + k; + const float* w2 = w1 + k; + const float* w3 = w2 + k; + int32_t i = 0; + for (; i + 3 <= m; i += 3) { + const float* a0 = c->in + (size_t) i * k; + const float* a1 = a0 + k; + const float* a2 = a1 + k; + float32x4_t s00 = vdupq_n_f32(0.0f), s01 = vdupq_n_f32(0.0f), s02 = vdupq_n_f32(0.0f), s03 = vdupq_n_f32(0.0f); + float32x4_t s10 = vdupq_n_f32(0.0f), s11 = vdupq_n_f32(0.0f), s12 = vdupq_n_f32(0.0f), s13 = vdupq_n_f32(0.0f); + float32x4_t s20 = vdupq_n_f32(0.0f), s21 = vdupq_n_f32(0.0f), s22 = vdupq_n_f32(0.0f), s23 = vdupq_n_f32(0.0f); + for (int32_t p = 0; p < k4; p += 4) { + const float32x4_t x0 = vld1q_f32(a0 + p); + const float32x4_t x1 = vld1q_f32(a1 + p); + const float32x4_t x2 = vld1q_f32(a2 + p); + float32x4_t y = vld1q_f32(w0 + p); + s00 = vfmaq_f32(s00, x0, y); s10 = vfmaq_f32(s10, x1, y); s20 = vfmaq_f32(s20, x2, y); + y = vld1q_f32(w1 + p); + s01 = vfmaq_f32(s01, x0, y); s11 = vfmaq_f32(s11, x1, y); s21 = vfmaq_f32(s21, x2, y); + y = vld1q_f32(w2 + p); + s02 = vfmaq_f32(s02, x0, y); s12 = vfmaq_f32(s12, x1, y); s22 = vfmaq_f32(s22, x2, y); + y = vld1q_f32(w3 + p); + s03 = vfmaq_f32(s03, x0, y); s13 = vfmaq_f32(s13, x1, y); s23 = vfmaq_f32(s23, x2, y); + } + float* out0 = c->out + (size_t) i * n + o; + float* out1 = out0 + n; + float* out2 = out1 + n; + out0[0] = skainet_neon_finish(s00, a0, w0, k4, k); + out0[1] = skainet_neon_finish(s01, a0, w1, k4, k); + out0[2] = skainet_neon_finish(s02, a0, w2, k4, k); + out0[3] = skainet_neon_finish(s03, a0, w3, k4, k); + out1[0] = skainet_neon_finish(s10, a1, w0, k4, k); + out1[1] = skainet_neon_finish(s11, a1, w1, k4, k); + out1[2] = skainet_neon_finish(s12, a1, w2, k4, k); + out1[3] = skainet_neon_finish(s13, a1, w3, k4, k); + out2[0] = skainet_neon_finish(s20, a2, w0, k4, k); + out2[1] = skainet_neon_finish(s21, a2, w1, k4, k); + out2[2] = skainet_neon_finish(s22, a2, w2, k4, k); + out2[3] = skainet_neon_finish(s23, a2, w3, k4, k); + } + for (; i < m; ++i) { + const float* a0 = c->in + (size_t) i * k; + float32x4_t s0 = vdupq_n_f32(0.0f), s1 = vdupq_n_f32(0.0f), s2 = vdupq_n_f32(0.0f), s3 = vdupq_n_f32(0.0f); + for (int32_t p = 0; p < k4; p += 4) { + const float32x4_t x0 = vld1q_f32(a0 + p); + s0 = vfmaq_f32(s0, x0, vld1q_f32(w0 + p)); + s1 = vfmaq_f32(s1, x0, vld1q_f32(w1 + p)); + s2 = vfmaq_f32(s2, x0, vld1q_f32(w2 + p)); + s3 = vfmaq_f32(s3, x0, vld1q_f32(w3 + p)); + } + float* out0 = c->out + (size_t) i * n + o; + out0[0] = skainet_neon_finish(s0, a0, w0, k4, k); + out0[1] = skainet_neon_finish(s1, a0, w1, k4, k); + out0[2] = skainet_neon_finish(s2, a0, w2, k4, k); + out0[3] = skainet_neon_finish(s3, a0, w3, k4, k); + } + } + for (; o < o_end; ++o) { + const float* wo = c->w + (size_t) o * k; + for (int32_t i = 0; i < m; ++i) { + c->out[(size_t) i * n + o] = skainet_neon_dot(c->in + (size_t) i * k, wo, k4, k); + } + } +} +#endif /* NEON */ + #if (defined(__x86_64__) || defined(_M_X64)) && !defined(__wasm__) #define SKAINET_FP32_WT_AVX2 1 #include @@ -202,6 +296,9 @@ SKAINET_API void skainet_fp32_matmul_wt( } skainet_fp32_wt_ctx ctx = { in + in_offset, w + w_offset, out + out_offset, m, n, k }; skainet_row_range_fn rows = skainet_fp32_wt_rows; +#ifdef SKAINET_HAVE_NEON + rows = skainet_fp32_wt_rows_neon; +#endif #ifdef SKAINET_FP32_WT_AVX2 if (skainet_cpu_has_avx2_fma()) rows = skainet_fp32_wt_rows_avx2; #endif diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt index 07511ea5..6b9ceb1f 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt @@ -22,10 +22,11 @@ import sk.ainet.backend.api.kernel.Fp32MatmulKernel * The C kernel has one body per lane: a register-tiled SGEMM on x86-64 hosts * with AVX2 and FMA (a 6 × 16 tile of C in twelve eight-lane accumulators over * a packed sixteen-column panel of B, selected at runtime by a cached `cpuid` - * probe because the library is built for baseline x86-64), a 4 × 8 tile on the - * wasm SIMD128 sidecar, and an i-p-j outer product whose inner `c[j] += a*b[j]` - * loop auto-vectorizes (NEON `fmla`, SSE2 without FMA on x86-64 hosts that lack - * AVX2). Every output accumulates its products in ascending `p` whichever tile, + * probe because the library is built for baseline x86-64), an 8 × 8 tile of + * lane-indexed `fmla` on AArch64 (Apple silicon, Android and Linux arm64), a + * 4 × 8 tile on the wasm SIMD128 sidecar, and an i-p-j outer product whose + * inner `c[j] += a*b[j]` loop auto-vectorizes elsewhere (SSE2 without FMA on + * x86-64 hosts that lack AVX2). Every output accumulates its products in ascending `p` whichever tile, * thread or call computes it, so sliced and whole calls agree bit for bit; * lanes differ from each other within rounding, as NEON and wasm always have * ([NativeFp32MatmulX86LaneTest]). diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtKernel.kt index 79d36d93..98dbedcd 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtKernel.kt @@ -20,9 +20,10 @@ import sk.ainet.backend.api.kernel.Fp32MemSegMatmulWtKernel * * which the wasm sidecar has served since 0.61.0 and the JVM never bound: `DefaultCpuOpsJvm` * transposed every dense weight once into `[k, n]` for the GEMM kernels and cached the copy. The C - * kernel computes every output as one contiguous dot product over the weight's own row (NEON, the - * portable body, or on x86-64 hosts with AVX2+FMA a runtime-selected eight-lane register tile that - * reads the weight once per call instead of once per activation row) and threads over output rows + * kernel computes every output as one contiguous dot product over the weight's own row (the + * portable body, on AArch64 a four-lane NEON tile of four weight rows by three activation rows, or + * on x86-64 hosts with AVX2+FMA a runtime-selected eight-lane register tile — both tiles read the + * weight once per call instead of once per activation row) and threads over output rows * through the shared row pool. The accumulation order of an output is fixed per build and does not * depend on the thread or row slice that computes it, so threaded, sliced and single calls agree * bit for bit; against transpose-then-GEMM the result agrees to rounding. diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulTileTest.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulTileTest.kt new file mode 100644 index 00000000..5d1536f7 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulTileTest.kt @@ -0,0 +1,137 @@ +package sk.ainet.exec.kernel + +import kotlinx.cinterop.ExperimentalForeignApi +import kotlinx.cinterop.addressOf +import kotlinx.cinterop.usePinned +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.kernels.cinterop.skainet_fp32_matmul + +/** + * The tails, the strides and the placement independence of the C `skainet_fp32_matmul`, through + * cinterop so the test runs on every Kotlin/Native target that embeds the archive — on macosArm64 + * (CI) and linuxArm64 that is the NEON lane: an 8×8 tile of lane-indexed FMAs over a packed + * 8-column panel in a 256-deep k-block, 4×8 and 1×8 tiles for the rows left over and a zero-padded + * panel plus C stripe for the columns left over. Shapes put `m` under, at and past a tile, `n` + * around eight, `k` around the k-block and off a multiple of four (the lane-indexed body takes four + * `p` at a time). Every element takes one fused step per `p` in ascending `p` whichever tile holds + * it, so a block of the output computed on its own is checked bit-for-bit against the full product. + * On linuxX64 the same tests exercise the AVX2 lane, on a host without it the portable loop. + */ +@OptIn(ExperimentalForeignApi::class) +class NativeKnFp32MatmulTileTest { + + private fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + c: FloatArray, cOffset: Int, cStride: Int, + m: Int, n: Int, k: Int, + ) { + a.usePinned { aPin -> + b.usePinned { bPin -> + c.usePinned { cPin -> + skainet_fp32_matmul( + aPin.addressOf(0), aOffset, aStride, + bPin.addressOf(0), bOffset, bStride, + cPin.addressOf(0), cOffset, cStride, + m, n, k, + ) + } + } + } + } + + private fun reference( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + m: Int, n: Int, k: Int, + ): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) for (j in 0 until n) { + var acc = 0.0 + for (p in 0 until k) acc += a[aOffset + i * aStride + p].toDouble() * b[bOffset + p * bStride + j] + out[i * n + j] = acc.toFloat() + } + return out + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, what: String) { + assertEquals(expected.size, actual.size, "$what size") + var maxAbs = 0f + for (v in expected) maxAbs = maxOf(maxAbs, abs(v)) + val tol = maxAbs * 1e-5f + 1e-6f + for (i in expected.indices) { + assertTrue(abs(expected[i] - actual[i]) <= tol, "$what[$i]: ${actual[i]} vs ${expected[i]} (tol $tol)") + } + } + + private fun block(c: FloatArray, cOffset: Int, cStride: Int, m: Int, n: Int): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) c.copyInto(out, i * n, cOffset + i * cStride, cOffset + i * cStride + n) + return out + } + + @Test + fun everyTileEdgeMatchesTheDoubleReference() { + val rng = Random(41) + for (m in listOf(1, 3, 4, 5, 8, 9, 13, 17)) for (n in listOf(1, 7, 8, 9, 15, 16, 17, 61)) for (k in listOf(1, 3, 255, 257)) { + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val b = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val out = FloatArray(m * n) + matmul(a, 0, k, b, 0, n, out, 0, n, m, n, k) + assertClose(reference(a, 0, k, b, 0, n, m, n, k), out, "m=$m n=$n k=$k") + } + } + + @Test + fun stridedOperandsAndAnOffsetOutputLeaveThePaddingAlone() { + val rng = Random(42) + for ((m, n, k) in listOf(Triple(17, 21, 300), Triple(8, 8, 256), Triple(9, 9, 257), Triple(1, 1, 1))) { + val aStride = k + 5; val bStride = n + 3; val cStride = n + 7 + val aOffset = 2; val bOffset = 1; val cOffset = 3 + val a = FloatArray(aOffset + m * aStride) { rng.nextFloat() - 0.5f } + val b = FloatArray(bOffset + k * bStride) { rng.nextFloat() - 0.5f } + val c = FloatArray(cOffset + m * cStride) { 123f } + matmul(a, aOffset, aStride, b, bOffset, bStride, c, cOffset, cStride, m, n, k) + assertClose(reference(a, aOffset, aStride, b, bOffset, bStride, m, n, k), block(c, cOffset, cStride, m, n), "m=$m n=$n k=$k") + for (idx in 0 until cOffset) assertEquals(123f, c[idx], "prefix $idx must be untouched") + for (i in 0 until m) for (j in n until cStride) { + if (cOffset + i * cStride + j < c.size) assertEquals(123f, c[cOffset + i * cStride + j], "pad [$i, $j] must be untouched") + } + } + } + + @Test + fun anOutputDoesNotDependOnTheBlockThatComputesIt() { + val rng = Random(43) + // 17 rows: two 8-row tiles and a 1-row tile with packing on; the partial calls move every + // output to another place inside a tile (5 and 3 are not multiples of 8), switch packing + // off (fewer than two tiles) and change which tile variant computes it. + for ((m, n, k) in listOf(Triple(17, 21, 300), Triple(9, 19, 1023), Triple(2, 10, 6), Triple(18, 1030, 64))) { + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val b = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val full = FloatArray(m * n) + matmul(a, 0, k, b, 0, n, full, 0, n, m, n, k) + + val byRows = FloatArray(m * n) + val cut = minOf(5, m - 1) + matmul(a, 0, k, b, 0, n, byRows, 0, n, cut, n, k) + matmul(a, cut * k, k, b, 0, n, byRows, cut * n, n, m - cut, n, k) + assertContentEquals(full, byRows, "m=$m n=$n k=$k: a row boundary must not change a bit") + + val byColumns = FloatArray(m * n) + val col = 3 + matmul(a, 0, k, b, 0, n, byColumns, 0, n, m, col, k) + matmul(a, 0, k, b, col, n, byColumns, col, n, m, n - col, k) + assertContentEquals(full, byColumns, "m=$m n=$n k=$k: a column boundary must not change a bit") + + val rows = FloatArray(m * n) + for (i in 0 until m) matmul(a, i * k, k, b, 0, n, rows, i * n, n, 1, n, k) + assertContentEquals(full, rows, "m=$m n=$n k=$k: a row-at-a-time call must not change a bit") + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulWtTileTest.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulWtTileTest.kt new file mode 100644 index 00000000..28db3cb6 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnFp32MatmulWtTileTest.kt @@ -0,0 +1,98 @@ +package sk.ainet.exec.kernel + +import kotlinx.cinterop.ExperimentalForeignApi +import kotlinx.cinterop.addressOf +import kotlinx.cinterop.usePinned +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.kernels.cinterop.skainet_fp32_matmul_wt + +/** + * The tails and the slice independence of the C `skainet_fp32_matmul_wt`, through cinterop so the + * test runs on every Kotlin/Native target that embeds the archive — on macosArm64 (CI) and + * linuxArm64 that is the NEON lane: four weight rows by three activation rows in twelve four-lane + * chains, a 4×1 tile for the activation rows left over and a single-row dot for the weight rows + * left over. Shapes put `k` off a multiple of four, `n` off a multiple of four, `m` at one, two, + * three and past a tile. Every variant reduces an output the same way (lanes, horizontal sum, + * fused scalar tail), so a slice of the weight's rows computed on its own, and one activation row + * at a time, are checked bit-for-bit against the whole product — what the row pool and the + * SKEEP-005 column split rely on. On linuxX64 the same tests exercise the AVX2 lane. + */ +@OptIn(ExperimentalForeignApi::class) +class NativeKnFp32MatmulWtTileTest { + + private fun matmulWt(x: FloatArray, xOffset: Int, w: FloatArray, wOffset: Int, out: FloatArray, outOffset: Int, m: Int, n: Int, k: Int) { + x.usePinned { xPin -> + w.usePinned { wPin -> + out.usePinned { oPin -> + skainet_fp32_matmul_wt(xPin.addressOf(0), xOffset, wPin.addressOf(0), wOffset, oPin.addressOf(0), outOffset, m, n, k) + } + } + } + } + + private fun reference(x: FloatArray, w: FloatArray, m: Int, n: Int, k: Int): FloatArray { + val out = FloatArray(m * n) + for (i in 0 until m) for (o in 0 until n) { + var acc = 0.0 + for (p in 0 until k) acc += x[i * k + p].toDouble() * w[o * k + p] + out[i * n + o] = acc.toFloat() + } + return out + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, what: String) { + assertEquals(expected.size, actual.size, "$what size") + var maxAbs = 0f + for (v in expected) maxAbs = maxOf(maxAbs, abs(v)) + val tol = maxAbs * 1e-5f + 1e-6f + for (i in expected.indices) { + assertTrue(abs(expected[i] - actual[i]) <= tol, "$what[$i]: ${actual[i]} vs ${expected[i]} (tol $tol)") + } + } + + @Test + fun everyTileEdgeMatchesTheRowDotReference() { + val rng = Random(51) + for (m in listOf(1, 2, 3, 4, 7, 20)) for (n in listOf(1, 2, 3, 4, 5, 9, 61)) for (k in listOf(1, 3, 4, 5, 255, 1023)) { + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val out = FloatArray(m * n) + matmulWt(x, 0, w, 0, out, 0, m, n, k) + assertClose(reference(x, w, m, n, k), out, "m=$m n=$n k=$k") + } + } + + @Test + fun anOutputDoesNotDependOnTheSliceThatComputesIt() { + val rng = Random(52) + for ((m, n, k) in listOf(Triple(20, 13, 1023), Triple(7, 9, 33), Triple(2, 11, 6), Triple(1, 6, 1023), Triple(5, 600, 64))) { + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val full = FloatArray(m * n) + matmulWt(x, 0, w, 0, full, 0, m, n, k) + + // The weight's rows split at 3: every output lands at another offset inside a 4-row tile. + val cut = 3 + val left = FloatArray(m * cut) + val right = FloatArray(m * (n - cut)) + matmulWt(x, 0, w, 0, left, 0, m, cut, k) + matmulWt(x, 0, w, cut * k, right, 0, m, n - cut, k) + val split = FloatArray(m * n) + for (i in 0 until m) { + left.copyInto(split, i * n, i * cut, i * cut + cut) + right.copyInto(split, i * n + cut, i * (n - cut), i * (n - cut) + (n - cut)) + } + assertContentEquals(full, split, "m=$m n=$n k=$k: a slice boundary must not change a bit") + + // One activation row at a time: the 4×1 tile and the single-row dot against the 4×3 tile. + val rows = FloatArray(m * n) + for (i in 0 until m) matmulWt(x, i * k, w, 0, rows, i * n, 1, n, k) + assertContentEquals(full, rows, "m=$m n=$n k=$k: a row-at-a-time call must not change a bit") + } + } +} From 2f1964dbf48a8d5e53c6929c2159439909529369 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sun, 25 Oct 2026 15:12:56 +0100 Subject: [PATCH 2/2] changelog: keep the AArch64 tile entry under Unreleased after the 0.62.3 merge --- CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fa8e7567..bffdeb27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- **Dense FP32 kernels get register tiles on AArch64** (`skainet-backend-native-cpu`, Apple silicon, Android arm64, Linux arm64): the NEON bodies of `skainet_fp32_matmul` and `skainet_fp32_matmul_wt` were the x86 story of a week ago — an outer-product loop that loaded and stored a row of C for every row of A, and a single four-lane chain per `x · Wᵀ` output that left most of a core's FMA pipes idle. The plain SGEMM now runs an 8 × 8 tile in sixteen `float32x4` accumulators with lane-indexed FMAs over a packed 8-column panel of B in a 256-deep k-block (4 × 8 and 1 × 8 tiles for the rows left over, a zero-padded panel and C stripe for the columns left over), and `x · Wᵀ` a tile of four weight rows by three activation rows in twelve chains. 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; against the previous NEON numbers an output moves within rounding. Two Kotlin/Native tests run the tails and the block independence through cinterop on every native target, so the macOS CI leg exercises the lane on Apple silicon. No numbers from this change's author: measured on an arm64 host it is not — see the pull request for how to run the bench. + ## [0.62.3] - 2026-10-25 Headline: **opening a tokenizer GGUF and building its tokenizer is several times cheaper, and the dense FP32 `A · B` @@ -16,7 +20,6 @@ and its heap face no longer zeroes the rest of a sub-block caller's output buffe ### Changed -- **Dense FP32 kernels get register tiles on AArch64** (`skainet-backend-native-cpu`, Apple silicon, Android arm64, Linux arm64): the NEON bodies of `skainet_fp32_matmul` and `skainet_fp32_matmul_wt` were the x86 story of a week ago — an outer-product loop that loaded and stored a row of C for every row of A, and a single four-lane chain per `x · Wᵀ` output that left most of a core's FMA pipes idle. The plain SGEMM now runs an 8 × 8 tile in sixteen `float32x4` accumulators with lane-indexed FMAs over a packed 8-column panel of B in a 256-deep k-block (4 × 8 and 1 × 8 tiles for the rows left over, a zero-padded panel and C stripe for the columns left over), and `x · Wᵀ` a tile of four weight rows by three activation rows in twelve chains. 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; against the previous NEON numbers an output moves within rounding. Two Kotlin/Native tests run the tails and the block independence through cinterop on every native target, so the macOS CI leg exercises the lane on Apple silicon. No numbers from this change's author: measured on an arm64 host it is not — see the pull request for how to run the bench. - **Dense FP32 SGEMM uses AVX2/FMA on x86-64** (`skainet-backend-native-cpu`): `skainet_fp32_matmul`, the kernel behind every `A · B` product the JVM hands to the native lane (a weight stored `[in, out]`, attention scores and values, training), had register-tiled bodies for wasm SIMD128 only; the x86-64 library is built for baseline x86-64, so the kernel ran a scalar/SSE2 outer-product loop without FMA that loaded and stored a row of C for every row of A. Hosts with AVX2 and FMA now select, through the cached `cpuid` probe the `x · Wᵀ` kernel introduced, a lane compiled per function for that target: a 6 × 16 tile of C in twelve eight-lane FMA accumulators over a 256-deep k-block, each sixteen-column panel of B packed once so a power-of-two `n` no longer maps every row of the panel to one L1 set, 2 × 16 and 1 × 16 tiles for the rows left over and a lane-masked stripe for the columns left over. Every output takes one fused multiply-add per `p` in ascending `p` whichever tile holds it, so sliced, threaded and whole calls stay bit-identical to each other; against the previous x86 numbers an output moves within rounding, as the NEON and wasm lanes already do. Machines without AVX2, and every other architecture, run the previous code. Measured single call on a six-core laptop, develop → this change: 63 × 1024 × 1024 21 → 281 GFLOP/s, 1024³ 48 → 311, 512³ 59 → 358, 192 × 192 × 64 (one thread) 18 → 99. - **`StreamingGGUFReader` parses the metadata through a buffered cursor** (`skainet-io-gguf`): the header, KV and tensor-info sections used one `RandomAccessSource.readAt` per value (a fresh byte array each, and on `JvmRandomAccessSource`/`AndroidRandomAccessSource` one `FileChannel` read each), plus a `Pair` and a boxed offset per array element. A tokenizer GGUF has arrays of ~262k strings, floats and ints, so that was over a million source calls. The reader now pulls the source in 64 KB blocks and decodes values from the block; the large arrays are decoded block by block straight into pre-sized lists. `fields` is unchanged: same keys and order, same value and element classes (`UInt` for UINT32, `Float` for FLOAT32, …), still plain lists. A truncated file or an array count larger than the file now fails with an `IllegalArgumentException` before allocating. Gemma-3 tokenizer GGUF (262k tokens), JVM: file source 1,002 → 21–45 ms, in-memory 55 → 21–34 ms; interpreter-only JVM (`-Xint`): file 13.9 s → 0.5 s, in-memory 1.4 → 0.5 s; on a 32-bit ARM Android device (debuggable app): file source 23.3 → 1.2 s, in-memory 2.0 → 1.15 s warm and 4.0–4.7 → 3.0 s on the first parse in a process. - **A GGUF SentencePiece tokenizer is built in one pass over the vocabulary** (`skainet-frontend-text`, `skainet-frontend-core`): `TextFrontEnds.fromGguf` derived the token list, the token types and the special-token map three to four times per build, each a full copy of the ~262k-entry lists, and copied the scores twice. They are now taken once and passed on; the token list is used without a copy when it holds only strings, the scores go straight into a `FloatArray`, and `ArrayVocabulary` fills its map with one operation per piece (first occurrence still wins). Public API unchanged. Token ids are identical to the previous build on 26,803 texts (`GgufTokenizerIdParityTest`, env-gated). Gemma-3 tokenizer: `-Xint` JVM 1,211 → 276 ms; 32-bit ARM Android device (debuggable app): 4.5–5.4 → 3.0–3.3 s on the first build in a process, 2.6–2.8 → 1.4–1.7 s warm.