diff --git a/CHANGELOG.md b/CHANGELOG.md index 56f1f554..f6e587ae 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 products on the JVM go to the platform BLAS where one is on, and on macOS it is on by default** (`skainet-backend-cpu`): `JvmBlas` never reached a BLAS before — it looked a framework up by its bare name, which `dlopen` cannot resolve, and passed heap arrays to a downcall, which JDK 21 rejects — so the `skainet.cpu.blas.enabled` switch was a no-op. It now resolves `cblas_sgemm` from `skainet.cpu.blas.library` / `SKAINET_CPU_BLAS_LIBRARY`, the Accelerate framework by path on macOS, or the common Linux sonames, and has an off-heap face. Three products use it when `blasEnabled` is true and the product is above a small-work floor: `A · B` on heap arrays (staged off-heap, as the native heap kernel does), `A · B` on two segment-backed operands (read in place), and the direct `x · Wᵀ` on a resident dense weight (`sgemm` with the weight transposed, no copy of it). On macOS the default is now on — Accelerate's SGEMM runs on the matrix units of Apple silicon and threads itself, several times the NEON kernel's throughput — and `-Dskainet.cpu.blas.enabled=false` restores the bundled kernels; elsewhere the default stays off. Results move within rounding against the bundled kernels, as they already do between lanes. ## [0.62.3] - 2026-10-25 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") + } + } +}