perf(native-cpu): AVX2/FMA lane for the dense FP32 x·Wᵀ kernel on x86-64

CHANGELOG.md

@@ -2,13 +2,14 @@
 
 ## [Unreleased]
 
+### Changed
+
+- **Dense FP32 `x · Wᵀ` uses AVX2/FMA on x86-64** (`skainet-backend-native-cpu`): `skainet_fp32_matmul_wt` had SIMD bodies for wasm and NEON only; the x86-64 library is built for baseline x86-64, so the kernel behind the JVM's resident dense weights (SKEEP-007) ran a scalar/SSE2 loop and re-read the weight once per activation row. Hosts with AVX2 and FMA now select, once per call through a cached `cpuid` probe, a lane compiled per function for that target: eight-lane FMA chains over a tile of four weight rows and two activation rows, so the weight streams from memory once per call. Every tile variant reduces an output in the same order, so threaded, sliced (`Schedule`) and single 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, m=63 n=1024 k=1024 on a 6-core laptop: see the pull request.
+
 ### Fixed
 
 - **`CoroutineSchedule.forRange` could call its body with an empty, inverted range** (`skainet-backend-cpu`): the chunk length is rounded up, so for some `n` and task counts the last chunk started past `n` (63 elements on 12 tasks: chunks of 6, the twelfth at 66..63). The engine's own bodies loop over the range and did nothing for it; a body that passes the range length to a kernel saw a negative size. The task count now follows the chunk length, and every range is non-empty and covers `[0, n)` as the `Schedule` contract states.
 
-### Changed
-
-- **Dense FP32 `x · Wᵀ` on the JVM native lane runs under the context schedule** (`skainet-backend-cpu`, SKEEP-005): `DefaultCpuOpsJvm` used to hand the whole product to `skainet_fp32_matmul_wt` in one call and leave threading to the C row pool, which is sized at compile time (four threads) and re-reads the weight for every activation row, so an encoder or prefill pass on a twelve-core host ran on about three cores at the shared cache's bandwidth. A product with several activation rows now splits the weight's rows across the `Schedule` — up to `parallelism` native calls over disjoint column slices, each slice narrow enough for a core's own cache — so `DirectCpuExecutionContext()` (the hardware schedule) uses every core and `withSchedule(Schedule.Sequential)` runs the old single call. Results are bit-identical either way: every output stays the same contiguous dot product in the same order. A single-row product (a decode step) is unchanged, as is a weight in a confined arena, which stays on the calling thread.
 
 ## [0.62.1] - 2026-10-21
 

skainet-backends/skainet-backend-native-cpu/native/include/skainet_cpu_features.h

@@ -14,4 +14,13 @@
  */
 int skainet_cpu_has_dotprod(void);
 
+/*
+ * Returns 1 if this x86-64 CPU and OS execute AVX2 and FMA (and the OS saves
+ * the YMM state). Cached after the first call; 0 on every other architecture.
+ * The shared library is built for baseline x86-64, so kernels that want these
+ * instructions compile the hot body a second time with the avx2+fma target
+ * and pick it per call through this probe (fp32_matmul_wt.c).
+ */
+int skainet_cpu_has_avx2_fma(void);
+
 #endif /* SKAINET_CPU_FEATURES_H */

skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h

@@ -499,9 +499,13 @@ SKAINET_API void skainet_ternary_lmhead_stage1(
  *   out[out_offset + i*n + o] = Σ_p in[in_offset + i*k + p] * w[w_offset + o*k + p]
  *
  * for i in [0, m), o in [0, n). No transposed copy of the weight is made; every
- * output is one contiguous dot product over the weight's own row, accumulated
- * in ascending p — bit-identical to transpose-then-GEMM. Threads over output
- * rows like the packed kernels. m, n, k <= 0 are no-ops (k == 0 zeroes out).
+ * output is one contiguous dot product over the weight's own row. Its
+ * accumulation order is fixed per build (portable loop, four-lane NEON/wasm
+ * sums, or the eight-lane AVX2/FMA tile that x86-64 hosts with AVX2 select at
+ * runtime) and never depends on the thread or row slice that computes it, so
+ * threaded, sliced and single calls are bit-identical to each other. Threads
+ * over output rows like the packed kernels. m, n, k <= 0 are no-ops (k == 0
+ * zeroes out).
  */
 SKAINET_API void skainet_fp32_matmul_wt(
     const float* in, int32_t in_offset,

skainet-backends/skainet-backend-native-cpu/native/src/fp32_matmul_wt.c

@@ -1,6 +1,7 @@
 #include "skainet_kernels.h"
 #include "skainet_simd.h"
 #include "skainet_row_threads.h"
+#include "skainet_cpu_features.h"
 #include <stddef.h>
 #include <stdint.h>
 
@@ -11,14 +12,25 @@
  *
  * Every output is one contiguous dot product over the weight's own row, so no
  * transposed copy of the weight is ever made — the shape a Linear layer stores
- * its weight in is the shape this kernel wants. Four weight rows per tile share
- * one pass over the activation row; each output accumulates its products in
- * ascending p in a single accumulator, so the result is bit-identical to the
- * transpose-then-GEMM path and to the Kotlin `denseRowDotTiled` reference.
+ * 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
+ * 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.
  *
  * Threads over output rows (disjoint out[] slices) when output_dim is large
  * 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.
+ *
+ * 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
+ * with the avx2+fma target (MSVC emits VEX intrinsics without /arch) and
+ * selected once per call by the cached cpuid probe in skainet_cpu_features.h,
+ * the same per-function dispatch the Apple arm64 dotprod kernels use.
  */
 typedef struct {
     const float* in;        /* in + in_offset, m rows of k */
@@ -61,6 +73,122 @@ static void skainet_fp32_wt_rows(void* vctx, int32_t o_start, int32_t o_end) {
     }
 }
 
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(__wasm__)
+#define SKAINET_FP32_WT_AVX2 1
+#include <immintrin.h>
+#if defined(__GNUC__) || defined(__clang__)
+#  define SKAINET_AVX2_TARGET __attribute__((target("avx2,fma")))
+#else
+#  define SKAINET_AVX2_TARGET
+#endif
+
+/*
+ * One output's reduction, shared by every tile variant: the eight lanes of an
+ * FMA chain over p < (k & ~7), then this fixed horizontal sum, then the scalar
+ * tail in ascending p. Keeping the three steps identical in the 4×2, 4×1 and
+ * 1-row variants is what makes the result independent of tile placement.
+ */
+SKAINET_AVX2_TARGET
+static inline float skainet_avx2_finish(__m256 acc, const float* a, const float* b, int32_t k8, int32_t k) {
+    __m128 lo = _mm256_castps256_ps128(acc);
+    const __m128 hi = _mm256_extractf128_ps(acc, 1);
+    lo = _mm_add_ps(lo, hi);
+    lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo));
+    lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 1));
+    float s = _mm_cvtss_f32(lo);
+    for (int32_t p = k8; p < k; ++p) s += a[p] * b[p];
+    return s;
+}
+
+/* Single weight row against one activation row (the o tail of a slice). */
+SKAINET_AVX2_TARGET
+static inline float skainet_avx2_dot(const float* a, const float* b, int32_t k8, int32_t k) {
+    __m256 acc = _mm256_setzero_ps();
+    for (int32_t p = 0; p < k8; p += 8) {
+        acc = _mm256_fmadd_ps(_mm256_loadu_ps(a + p), _mm256_loadu_ps(b + p), acc);
+    }
+    return skainet_avx2_finish(acc, a, b, k8, k);
+}
+
+/*
+ * Four weight rows against activation rows two at a time: per p-step six loads
+ * feed eight FMAs, and the 4 × k floats of weight stay in the core's own cache
+ * while every activation row streams past them — the weight is read from
+ * memory once per call, not once per activation row.
+ */
+SKAINET_AVX2_TARGET
+static void skainet_fp32_wt_rows_avx2(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 k8 = k & ~7;
+    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 + 2 <= m; i += 2) {
+            const float* a0 = c->in + (size_t) i * k;
+            const float* a1 = a0 + k;
+            __m256 s00 = _mm256_setzero_ps(), s01 = _mm256_setzero_ps();
+            __m256 s02 = _mm256_setzero_ps(), s03 = _mm256_setzero_ps();
+            __m256 s10 = _mm256_setzero_ps(), s11 = _mm256_setzero_ps();
+            __m256 s12 = _mm256_setzero_ps(), s13 = _mm256_setzero_ps();
+            for (int32_t p = 0; p < k8; p += 8) {
+                const __m256 x0 = _mm256_loadu_ps(a0 + p);
+                const __m256 x1 = _mm256_loadu_ps(a1 + p);
+                __m256 y = _mm256_loadu_ps(w0 + p);
+                s00 = _mm256_fmadd_ps(x0, y, s00);
+                s10 = _mm256_fmadd_ps(x1, y, s10);
+                y = _mm256_loadu_ps(w1 + p);
+                s01 = _mm256_fmadd_ps(x0, y, s01);
+                s11 = _mm256_fmadd_ps(x1, y, s11);
+                y = _mm256_loadu_ps(w2 + p);
+                s02 = _mm256_fmadd_ps(x0, y, s02);
+                s12 = _mm256_fmadd_ps(x1, y, s12);
+                y = _mm256_loadu_ps(w3 + p);
+                s03 = _mm256_fmadd_ps(x0, y, s03);
+                s13 = _mm256_fmadd_ps(x1, y, s13);
+            }
+            float* out0 = c->out + (size_t) i * n + o;
+            float* out1 = out0 + n;
+            out0[0] = skainet_avx2_finish(s00, a0, w0, k8, k);
+            out0[1] = skainet_avx2_finish(s01, a0, w1, k8, k);
+            out0[2] = skainet_avx2_finish(s02, a0, w2, k8, k);
+            out0[3] = skainet_avx2_finish(s03, a0, w3, k8, k);
+            out1[0] = skainet_avx2_finish(s10, a1, w0, k8, k);
+            out1[1] = skainet_avx2_finish(s11, a1, w1, k8, k);
+            out1[2] = skainet_avx2_finish(s12, a1, w2, k8, k);
+            out1[3] = skainet_avx2_finish(s13, a1, w3, k8, k);
+        }
+        if (i < m) {
+            const float* a0 = c->in + (size_t) i * k;
+            __m256 s0 = _mm256_setzero_ps(), s1 = _mm256_setzero_ps();
+            __m256 s2 = _mm256_setzero_ps(), s3 = _mm256_setzero_ps();
+            for (int32_t p = 0; p < k8; p += 8) {
+                const __m256 x0 = _mm256_loadu_ps(a0 + p);
+                s0 = _mm256_fmadd_ps(x0, _mm256_loadu_ps(w0 + p), s0);
+                s1 = _mm256_fmadd_ps(x0, _mm256_loadu_ps(w1 + p), s1);
+                s2 = _mm256_fmadd_ps(x0, _mm256_loadu_ps(w2 + p), s2);
+                s3 = _mm256_fmadd_ps(x0, _mm256_loadu_ps(w3 + p), s3);
+            }
+            float* out0 = c->out + (size_t) i * n + o;
+            out0[0] = skainet_avx2_finish(s0, a0, w0, k8, k);
+            out0[1] = skainet_avx2_finish(s1, a0, w1, k8, k);
+            out0[2] = skainet_avx2_finish(s2, a0, w2, k8, k);
+            out0[3] = skainet_avx2_finish(s3, a0, w3, k8, 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_avx2_dot(c->in + (size_t) i * k, wo, k8, k);
+        }
+    }
+}
+#endif /* x86-64 */
+
 SKAINET_API void skainet_fp32_matmul_wt(
     const float* in, int32_t in_offset,
     const float* w, int32_t w_offset,
@@ -73,5 +201,9 @@ SKAINET_API void skainet_fp32_matmul_wt(
         return;
     }
     skainet_fp32_wt_ctx ctx = { in + in_offset, w + w_offset, out + out_offset, m, n, k };
-    skainet_run_rows(skainet_fp32_wt_rows, &ctx, n);
+    skainet_row_range_fn rows = skainet_fp32_wt_rows;
+#ifdef SKAINET_FP32_WT_AVX2
+    if (skainet_cpu_has_avx2_fma()) rows = skainet_fp32_wt_rows_avx2;
+#endif
+    skainet_run_rows(rows, &ctx, n);
 }

skainet-backends/skainet-backend-native-cpu/native/src/skainet_cpu_features.c

@@ -34,3 +34,52 @@ int skainet_cpu_has_dotprod(void) {
 }
 
 #endif
+
+#if (defined(__x86_64__) || defined(_M_X64)) && !defined(__wasm__)
+
+#if defined(_MSC_VER) && !defined(__clang__)
+#include <intrin.h>
+
+int skainet_cpu_has_avx2_fma(void) {
+    /* Benign-race cache, as for the dotprod probe above. */
+    static int cached = -1;
+    if (cached < 0) {
+        int ok = 0;
+        int r[4];
+        __cpuid(r, 0);
+        if (r[0] >= 7) {
+            __cpuid(r, 1);
+            const int fma = (r[2] >> 12) & 1;
+            const int osxsave = (r[2] >> 27) & 1;
+            const int avx = (r[2] >> 28) & 1;
+            __cpuidex(r, 7, 0);
+            const int avx2 = (r[1] >> 5) & 1;
+            if (fma && osxsave && avx && avx2) {
+                /* XCR0 bits 1 (XMM) and 2 (YMM): the OS restores the upper halves. */
+                ok = (_xgetbv(0) & 0x6) == 0x6;
+            }
+        }
+        cached = ok;
+    }
+    return cached;
+}
+
+#else /* GCC / Clang: libgcc's probe also checks the OS XSAVE state. */
+
+int skainet_cpu_has_avx2_fma(void) {
+    static int cached = -1;
+    if (cached < 0) {
+        cached = __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma");
+    }
+    return cached;
+}
+
+#endif
+
+#else /* not x86-64 */
+
+int skainet_cpu_has_avx2_fma(void) {
+    return 0;
+}
+
+#endif

skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtKernel.kt

@@ -20,9 +20,12 @@ 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 or
- * auto-vectorised portable body, four rows per tile), threads over output rows through the shared
- * row pool, and accumulates in ascending `p` — bit-identical to transpose-then-GEMM.
+ * 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
+ * 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.
  *
  * Two faces. The heap face ([Fp32MatmulWtKernel]) stages both operands off-heap per call, like
  * [NativeFp32MatmulKernel]; for a decode step that copy is the weight, so it is the slow one. The