From f519b9b9cf8e6ec9c889379407c00e9d2e64b203 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Fri, 23 Oct 2026 20:54:56 +0200 Subject: [PATCH] =?UTF-8?q?perf(native-cpu):=20AVX2/FMA=20lane=20for=20the?= =?UTF-8?q?=20dense=20FP32=20x=C2=B7W=E1=B5=80=20kernel=20on=20x86-64?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 ran an SSE2 loop and re-read the weight once per activation row. Hosts with AVX2 and FMA now select, through a cached cpuid probe (skainet_cpu_has_avx2_fma), a per-function target("avx2,fma") lane: eight-lane FMA chains over a tile of four weight rows by two activation rows, the weight streamed from memory once per call. Every tile variant reduces an output in the same order, so threaded, sliced and single calls stay bit-identical to each other; the header and KDoc now describe that instead of a bit-identity to transpose-then-GEMM that the NEON and wasm lanes never had. Measured single call on a 6-core host, m=63 n=1024 k=1024: 29.6 → 131.7 GFLOP/s; decode m=1: 28.5 → 71.7; one-thread n=256: 9.0 → 56.5. --- CHANGELOG.md | 4 + .../native/include/skainet_cpu_features.h | 9 ++ .../native/include/skainet_kernels.h | 10 +- .../native/src/fp32_matmul_wt.c | 142 +++++++++++++++++- .../native/src/skainet_cpu_features.c | 49 ++++++ .../exec/kernel/NativeFp32MatmulWtKernel.kt | 9 +- .../kernel/NativeFp32MatmulWtX86LaneTest.kt | 131 ++++++++++++++++ 7 files changed, 343 insertions(+), 11 deletions(-) create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtX86LaneTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index a84b75487..08b7b36db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [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. + ## [0.62.1] - 2026-10-21 Headline: **the native Q8_0 / Q4_0 / Q5_0 / Q5_1 kernel bindings no longer zero the earlier rows of a diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_cpu_features.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_cpu_features.h index 48f171872..85e9e530c 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_cpu_features.h +++ b/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 */ diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h index cbfbbf808..101968b14 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h +++ b/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, 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 06ce9ea71..23964070d 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 @@ -1,6 +1,7 @@ #include "skainet_kernels.h" #include "skainet_simd.h" #include "skainet_row_threads.h" +#include "skainet_cpu_features.h" #include #include @@ -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 +#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); } diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/skainet_cpu_features.c b/skainet-backends/skainet-backend-native-cpu/native/src/skainet_cpu_features.c index 0302d0461..9f5d7dfa6 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/src/skainet_cpu_features.c +++ b/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 + +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 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 1ec68d773..79d36d938 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,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 diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtX86LaneTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtX86LaneTest.kt new file mode 100644 index 000000000..e67b59293 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulWtX86LaneTest.kt @@ -0,0 +1,131 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * The tails and the slice independence of `skainet_fp32_matmul_wt`. The x86-64 AVX2/FMA lane tiles + * four weight rows by two activation rows with eight-lane sums, so every edge the tile can fall off + * is a shape here: `k` off a multiple of eight, `n` off a multiple of four, `m` odd and one. The + * reference is a double-accumulated row dot; the kernel's own order (lanes, then a fixed horizontal + * sum, then the scalar tail) is checked to be the same whichever slice an output lands in, which is + * what the row pool and the SKEEP-005 column split rely on. On a host without AVX2 the same tests + * exercise the portable lane. + */ +class NativeFp32MatmulWtX86LaneTest { + + @BeforeTest + fun checkAvailable() { + assertTrue(NativeFp32MatmulWtKernel.isAvailable(), "bundled libskainet_kernels must resolve skainet_fp32_matmul_wt") + } + + 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(11) + val ms = listOf(1, 2, 3, 63) + val ns = listOf(1, 2, 3, 5, 1021) + val ks = listOf(1, 7, 9, 1023) + for (m in ms) for (n in ns) for (k in ks) { + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val out = FloatArray(m * n) + NativeFp32MatmulWtKernel.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(12) + // Shapes under the row-pool threshold, so the full call is one slice; the split calls put + // every output at a different offset inside a tile (3 is not a multiple of 4) and change + // which tile variant computes it (4×2, 4×1 or the single-row tail). + for ((m, n, k) in listOf(Triple(63, 13, 1024), Triple(5, 9, 33), Triple(2, 11, 8), Triple(1, 6, 1023))) { + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val full = FloatArray(m * n) + NativeFp32MatmulWtKernel.matmulWt(x, 0, w, 0, full, 0, m, n, k) + val split = FloatArray(m * n) + val cut = 3 + val left = FloatArray(m * cut) + val right = FloatArray(m * (n - cut)) + NativeFp32MatmulWtKernel.matmulWt(x, 0, w, 0, left, 0, m, cut, k) + NativeFp32MatmulWtKernel.matmulWt(x, 0, w, cut * k, right, 0, m, n - cut, k) + for (i in 0 until m) { + System.arraycopy(left, i * cut, split, i * n, cut) + System.arraycopy(right, i * (n - cut), split, 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 against all rows at once: the 4×1 and single-row + // variants against the 4×2 tile. + val rows = FloatArray(m * n) + for (i in 0 until m) NativeFp32MatmulWtKernel.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") + } + } + + /** + * `-Dskainet.runBench=true`: one native call per shape through the segment face (the weight + * resident off-heap, as the ops layer holds it), so the C row pool is the only threading and + * no weight copy is timed. + */ + @Test + fun singleCallBench() { + if (System.getProperty("skainet.runBench") != "true") return + val shapes = listOf( + Triple(63, 1024, 1024), Triple(63, 2624, 1024), Triple(63, 1024, 2624), + Triple(183, 1024, 1024), Triple(1, 1024, 1024), Triple(63, 256, 1024), + ) + println("[wt-x86] host ${System.getProperty("os.name")} ${System.getProperty("os.arch")} JDK ${System.getProperty("java.version")}") + Arena.ofConfined().use { arena -> + for ((m, n, k) in shapes) { + val rng = Random(m * 31 + n) + val x = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val w = FloatArray(n * k) { rng.nextFloat() - 0.5f } + val wSeg: MemorySegment = arena.allocate(w.size.toLong() * java.lang.Float.BYTES, ValueLayout.JAVA_FLOAT.byteAlignment()) + MemorySegment.copy(w, 0, wSeg, ValueLayout.JAVA_FLOAT, 0L, w.size) + val out = FloatArray(m * n) + val flop = 2.0 * m * n * k + val reps = maxOf(3, (2e9 / flop).toInt().coerceAtMost(400)) + val call = { NativeFp32MatmulWtKernel.matmulWt(x, 0, wSeg, 0L, out, 0, m, n, k) } + repeat(reps / 2 + 1) { call() } + val samples = LongArray(5) + for (s in samples.indices) { + val t0 = System.nanoTime() + repeat(reps) { call() } + samples[s] = (System.nanoTime() - t0) / reps + } + val ns = samples.sorted()[2] + println("[wt-x86] m=$m n=$n k=$k: %.3f ms/call %.1f GFLOP/s".format(ns / 1e6, flop / ns)) + } + } + } +}