From 08182d8cf1a116da8319b790d47128d3de25ce9a Mon Sep 17 00:00:00 2001 From: James Burton Date: Mon, 8 Jun 2026 22:44:02 +0100 Subject: [PATCH 1/2] kernels(cpu)(matmul): F32 outer-product tiled GEMM kernel for prefill (#312) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new F32 outer-product GEMM kernel as a companion to the existing Q8_0/Q5_0/K-quant outer-product kernels landed in PR #61. The Q8_0 outer-product (4×3 AVX2 tile) is currently blocked on RyuJIT register pressure (23 YMM needed, 16 available) due to the Q8_0-specific overhead — `ones` mask vector, scale extraction, `Half→float` conversion. The F32 path has none of these artifacts: a 4×3 tile reaches exactly 12 + 3 + 1 = 16 YMM (12 accumulators + 3 token vectors + 1 reloaded weight vector), which fits the AVX2 register file naturally. Kernel design - 4 weight rows × 3 input tokens register tile. - Vectorises along K (8 floats/lane via Vector256). - At each K-step, 3 token vectors are held in registers and reused across 4 FMAs per row (weight-vector reused 3× across the tokens). - Horizontal-reduces 12 accumulators into the C tile. - Scalar K-tail, row tail, and token tail handle non-tile-aligned shapes. - AVX2/FMA detection at the public entry point; falls back to a scalar reference implementation otherwise. Tests (tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs) - 25 new test cases extending the existing Q8_0 outer-product test class. - Scalar path: **bit-exact** equality with MatMul.GemmF32Scalar (same accumulation order, no auto-FMA contraction). - Vector path: tolerance ≤ 4e-6·√K vs MatMul.GemmF32 (FMA + horizontal- reduction reorders rounding). - Coverage: tile-aligned shapes, all-tail-combination shapes (row tail, token tail, K tail, all three combined), and edge cases (M=1, N=1, K=1, pure inner product). Benchmark (benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs) - Compares OuterProductGemmF32 vs production MatMul.GemmF32 at three prefill profiles (M ∈ {128, 512, 2048}, K=4096, N=32). - Intel Core Ultra 7 155H (AVX2, no AVX-512), single-threaded, 15 iterations: | M | Baseline | Outer-product | Speedup | |------|-----------|---------------|----------| | 128 | 2.507 ms | 1.238 ms | 2.02× | | 512 | 10.216 ms | 5.588 ms | 1.83× | | 2048 | 47.070 ms | 32.714 ms | 1.44× | Scope - New, independent kernel — no callers re-wired. Production prefill continues to use MatMul.GemmF32 / GemmF32(..., pool). Caller switching (and threaded outer-product dispatch) is a separate PR after broader benchmark validation. - F32 only — Q8_0/Q5_0/K-quant outer-product re-enablement remains blocked on AVX2 register pressure as documented in the roadmap. Co-Authored-By: Claude Opus 4.7 --- .../OuterProductGemmF32Benchmark.cs | 73 ++++ src/DotLLM.Cpu/Kernels/OuterProductGemm.cs | 363 ++++++++++++++++++ .../Cpu/Kernels/OuterProductGemmTests.cs | 145 +++++++ 3 files changed, 581 insertions(+) create mode 100644 benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs create mode 100644 src/DotLLM.Cpu/Kernels/OuterProductGemm.cs diff --git a/benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs b/benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs new file mode 100644 index 00000000..54fc5c2a --- /dev/null +++ b/benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs @@ -0,0 +1,73 @@ +using System.Runtime.InteropServices; +using BenchmarkDotNet.Attributes; +using DotLLM.Cpu.Kernels; + +namespace DotLLM.Benchmarks; + +/// +/// Compares the new against the +/// production +/// at prefill-shaped workloads (multi-token, contraction along K). +/// +/// Convention: C[N,M] = B[N,K] × A[M,K]^T — N is the batch (token count), +/// M is the output dim (e.g. hidden size), K is the contraction dim (e.g. K-proj +/// from the attention block input). +/// +/// Run with: +/// dotnet run -c Release -- --filter '*OuterProductGemmF32Benchmark*' +/// +[MemoryDiagnoser] +[SimpleJob(warmupCount: 5, iterationCount: 15)] +public unsafe class OuterProductGemmF32Benchmark : IDisposable +{ + // Three prefill profiles spanning typical attention-projection shapes. + // K=4096 mirrors Llama-3-8B's hidden_size = 4096 and 32-head q_proj output. + [Params(128, 512, 2048)] + public int M { get; set; } + + public int K { get; set; } = 4096; + + public int N { get; set; } = 32; + + private float* _a; + private float* _b; + private float* _c; + + [GlobalSetup] + public void Setup() + { + var rng = new Random(42); + long aLen = (long)M * K; + long bLen = (long)N * K; + long cLen = (long)N * M; + + _a = (float*)NativeMemory.AlignedAlloc((nuint)(aLen * sizeof(float)), 64); + _b = (float*)NativeMemory.AlignedAlloc((nuint)(bLen * sizeof(float)), 64); + _c = (float*)NativeMemory.AlignedAlloc((nuint)(cLen * sizeof(float)), 64); + + for (long i = 0; i < aLen; i++) _a[i] = rng.NextSingle() * 2f - 1f; + for (long i = 0; i < bLen; i++) _b[i] = rng.NextSingle() * 2f - 1f; + } + + public void Dispose() + { + if (_a != null) { NativeMemory.AlignedFree(_a); _a = null; } + if (_b != null) { NativeMemory.AlignedFree(_b); _b = null; } + if (_c != null) { NativeMemory.AlignedFree(_c); _c = null; } + GC.SuppressFinalize(this); + } + + /// Baseline: production tiled GEMM path used by the engine today. + [Benchmark(Baseline = true)] + public void GemmF32_Baseline() + { + MatMul.GemmF32(_a, _b, _c, M, K, N); + } + + /// Candidate: new outer-product 4×3 AVX2 microkernel. + [Benchmark] + public void OuterProductGemmF32_Avx2() + { + OuterProductGemm.OuterProductGemmF32(_a, _b, _c, M, K, N); + } +} diff --git a/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs b/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs new file mode 100644 index 00000000..22fae4be --- /dev/null +++ b/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs @@ -0,0 +1,363 @@ +using System.Runtime.CompilerServices; +using System.Runtime.Intrinsics; +using System.Runtime.Intrinsics.X86; + +namespace DotLLM.Cpu.Kernels; + +/// +/// F32 outer-product tiled GEMM kernels for prefill (multi-token) workloads. +/// Companion to the Q8_0/Q5_0/K-quant outer-product kernels in . +/// +/// +/// +/// dotLLM convention: C[N,M] = B[N,K] × A[M,K]^T. +/// +/// A is the weight matrix [M,K] in row-major. +/// B is the activation matrix [N,K] in row-major. +/// C is the output matrix [N,M] in row-major. +/// +/// +/// +/// The outer-product formulation processes an M_R × N_R register tile per K-step +/// (here 4 weight rows × 3 input tokens) so each loaded activation vector is reused +/// M_R times and each loaded weight vector is reused N_R times. Reduces +/// the per-FMA memory bandwidth requirement vs the standard inner-product GEMM that +/// computes one (row, token) cell to completion before moving on. +/// +/// +/// AVX2 register accounting for the 4×3 tile (within 16-YMM budget): +/// +/// 12 accumulators (one per (row, token) cell) +/// 3 token vectors (one per input row, reloaded each K-step) +/// 1 weight vector (reloaded for each of the 4 weight rows) +/// +/// Unlike the Q8_0 variant (PR #61 — blocked at 23 YMM due to the dequant artifacts), +/// the F32 path has no ones mask, no scale extraction, and no Half→float +/// conversion, so the 4×3 tile fits the AVX2 register file naturally. +/// +/// +/// Vectorization is along the K dimension (the contraction axis), 8 floats per +/// AVX2 step. The K-tail (last K % 8 elements) is handled scalar. +/// +/// +/// This kernel is a new, independent code path. Callers continue to use +/// until a +/// benchmark-driven decision is made to switch dispatch — that wiring is a separate PR. +/// +/// +public static unsafe class OuterProductGemm +{ + /// Register-tile rows (weight-row count processed per microkernel call). + private const int TileRows = 4; + + /// Register-tile tokens (input-row count processed per microkernel call). + private const int TileTokens = 3; + + /// SIMD lane width for AVX2 (). + private const int LaneAvx2 = 8; + + /// + /// Scalar reference outer-product GEMM. Computes C[N,M] = B[N,K] × A[M,K]^T + /// by iterating M_R × N_R register tiles and accumulating contributions over K. + /// Pure scalar — present as a correctness oracle for the vector variant. + /// + /// Weight matrix [M,K] row-major. + /// Input matrix [N,K] row-major. + /// Output matrix [N,M] row-major. + /// Number of weight rows (output dim). + /// Contraction dim. + /// Number of input tokens (batch dim). + [SkipLocalsInit] + public static void OuterProductGemmF32Scalar(float* a, float* b, float* c, int m, int k, int n) + { + // Every C cell is assigned exactly once across the three nested passes + // (full tiles ∪ row tail ∪ token tail = [0,n) × [0,m), no overlap, no + // gap) — no pre-zero of C is required. + + // Iterate output tiles of size TileTokens × TileRows. Tail rows/tokens + // fall through to a 1×1 inner-product cleanup so every shape is handled. + int mFullEnd = (m / TileRows) * TileRows; + int nFullEnd = (n / TileTokens) * TileTokens; + + for (int tStart = 0; tStart < nFullEnd; tStart += TileTokens) + { + for (int rStart = 0; rStart < mFullEnd; rStart += TileRows) + { + TileScalar4x3(a, b, c, rStart, tStart, k, m); + } + + // Row tail (rStart in [mFullEnd, m)). + for (int rStart = mFullEnd; rStart < m; rStart++) + { + for (int tOff = 0; tOff < TileTokens; tOff++) + { + int t = tStart + tOff; + c[t * m + rStart] = DotScalar(a + (long)rStart * k, b + (long)t * k, k); + } + } + } + + // Token tail (tStart in [nFullEnd, n)) — process remaining tokens as 1×TileRows + // tiles, then per-row scalar for the row tail. + for (int tStart = nFullEnd; tStart < n; tStart++) + { + for (int rStart = 0; rStart < m; rStart++) + { + c[tStart * m + rStart] = DotScalar(a + (long)rStart * k, b + (long)tStart * k, k); + } + } + } + + /// + /// AVX2 outer-product GEMM. Falls back to + /// when AVX2/FMA are unavailable. Bit-exact-modulo-FP-order with the scalar variant + /// to within typical FMA-vs-mul-then-add rounding (≤ a few ULP for typical inputs). + /// + /// Weight matrix [M,K] row-major. + /// Input matrix [N,K] row-major. + /// Output matrix [N,M] row-major. + /// Number of weight rows (output dim). + /// Contraction dim. + /// Number of input tokens (batch dim). + [SkipLocalsInit] + public static void OuterProductGemmF32(float* a, float* b, float* c, int m, int k, int n) + { + if (!Avx2.IsSupported || !Fma.IsSupported) + { + OuterProductGemmF32Scalar(a, b, c, m, k, n); + return; + } + + // Every C cell is assigned exactly once across the three nested passes + // (full tiles ∪ row tail ∪ token tail = [0,n) × [0,m), no overlap, no + // gap) — no pre-zero of C is required. + + int mFullEnd = (m / TileRows) * TileRows; + int nFullEnd = (n / TileTokens) * TileTokens; + int kVec = (k / LaneAvx2) * LaneAvx2; + + for (int tStart = 0; tStart < nFullEnd; tStart += TileTokens) + { + for (int rStart = 0; rStart < mFullEnd; rStart += TileRows) + { + TileAvx2_4x3(a, b, c, rStart, tStart, k, m, kVec); + } + + // Row tail — fall back to standard inner-product per cell. Vectorised + // via Vector256 dot-product, mirroring the K-tail strategy. + for (int rStart = mFullEnd; rStart < m; rStart++) + { + for (int tOff = 0; tOff < TileTokens; tOff++) + { + int t = tStart + tOff; + c[t * m + rStart] = DotAvx2(a + (long)rStart * k, b + (long)t * k, k, kVec); + } + } + } + + // Token tail — remaining tokens (n % TileTokens). Same row-tail vectorised + // dot product path covers full and partial row segments uniformly. + for (int tStart = nFullEnd; tStart < n; tStart++) + { + for (int rStart = 0; rStart < m; rStart++) + { + c[tStart * m + rStart] = DotAvx2(a + (long)rStart * k, b + (long)tStart * k, k, kVec); + } + } + } + + // ──────────────────── Scalar microkernel ──────────────────── + + /// + /// Scalar 4×3 microkernel: accumulates contributions of one weight-tile (4 rows) + /// × one token-tile (3 tokens) over the full K range. + /// + [SkipLocalsInit] + private static void TileScalar4x3( + float* a, float* b, float* c, + int rStart, int tStart, int k, int m) + { + // 12 accumulators (4 rows × 3 tokens). + float acc00 = 0, acc01 = 0, acc02 = 0; + float acc10 = 0, acc11 = 0, acc12 = 0; + float acc20 = 0, acc21 = 0, acc22 = 0; + float acc30 = 0, acc31 = 0, acc32 = 0; + + float* aRow0 = a + (long)(rStart + 0) * k; + float* aRow1 = a + (long)(rStart + 1) * k; + float* aRow2 = a + (long)(rStart + 2) * k; + float* aRow3 = a + (long)(rStart + 3) * k; + + float* bTok0 = b + (long)(tStart + 0) * k; + float* bTok1 = b + (long)(tStart + 1) * k; + float* bTok2 = b + (long)(tStart + 2) * k; + + for (int i = 0; i < k; i++) + { + float b0 = bTok0[i]; + float b1 = bTok1[i]; + float b2 = bTok2[i]; + + float a0 = aRow0[i]; + acc00 += a0 * b0; acc01 += a0 * b1; acc02 += a0 * b2; + + float a1 = aRow1[i]; + acc10 += a1 * b0; acc11 += a1 * b1; acc12 += a1 * b2; + + float a2 = aRow2[i]; + acc20 += a2 * b0; acc21 += a2 * b1; acc22 += a2 * b2; + + float a3 = aRow3[i]; + acc30 += a3 * b0; acc31 += a3 * b1; acc32 += a3 * b2; + } + + // C[t,r] layout: c[t * m + r]. + c[(long)(tStart + 0) * m + rStart + 0] = acc00; + c[(long)(tStart + 0) * m + rStart + 1] = acc10; + c[(long)(tStart + 0) * m + rStart + 2] = acc20; + c[(long)(tStart + 0) * m + rStart + 3] = acc30; + + c[(long)(tStart + 1) * m + rStart + 0] = acc01; + c[(long)(tStart + 1) * m + rStart + 1] = acc11; + c[(long)(tStart + 1) * m + rStart + 2] = acc21; + c[(long)(tStart + 1) * m + rStart + 3] = acc31; + + c[(long)(tStart + 2) * m + rStart + 0] = acc02; + c[(long)(tStart + 2) * m + rStart + 1] = acc12; + c[(long)(tStart + 2) * m + rStart + 2] = acc22; + c[(long)(tStart + 2) * m + rStart + 3] = acc32; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static float DotScalar(float* a, float* b, int k) + { + float sum = 0; + for (int i = 0; i < k; i++) + sum += a[i] * b[i]; + return sum; + } + + // ──────────────────── AVX2 microkernel ──────────────────── + + /// + /// AVX2 4×3 microkernel. Vectorises along K (8 floats/lane) with 12 accumulators + /// (one per (row, token) cell). At each K-step we load 3 token vectors and + /// 4 weight-row vectors, sharing each load across the orthogonal dimension via + /// 12 FMAs. + /// + /// + /// Register inventory (16 YMM): + /// + /// 12 accumulators + /// 3 token vectors (vb0, vb1, vb2) + /// 1 weight vector (reloaded for each of the 4 rows) + /// + /// + [SkipLocalsInit] + [MethodImpl(MethodImplOptions.AggressiveOptimization)] + private static void TileAvx2_4x3( + float* a, float* b, float* c, + int rStart, int tStart, int k, int m, int kVec) + { + Vector256 acc00 = Vector256.Zero, acc01 = Vector256.Zero, acc02 = Vector256.Zero; + Vector256 acc10 = Vector256.Zero, acc11 = Vector256.Zero, acc12 = Vector256.Zero; + Vector256 acc20 = Vector256.Zero, acc21 = Vector256.Zero, acc22 = Vector256.Zero; + Vector256 acc30 = Vector256.Zero, acc31 = Vector256.Zero, acc32 = Vector256.Zero; + + float* aRow0 = a + (long)(rStart + 0) * k; + float* aRow1 = a + (long)(rStart + 1) * k; + float* aRow2 = a + (long)(rStart + 2) * k; + float* aRow3 = a + (long)(rStart + 3) * k; + + float* bTok0 = b + (long)(tStart + 0) * k; + float* bTok1 = b + (long)(tStart + 1) * k; + float* bTok2 = b + (long)(tStart + 2) * k; + + for (int i = 0; i < kVec; i += LaneAvx2) + { + // Load 3 token vectors (held in registers across the row loop). + Vector256 vb0 = Vector256.Load(bTok0 + i); + Vector256 vb1 = Vector256.Load(bTok1 + i); + Vector256 vb2 = Vector256.Load(bTok2 + i); + + // Row 0: load A row vector, do 3 FMAs sharing it across tokens. + Vector256 va = Vector256.Load(aRow0 + i); + acc00 = Fma.MultiplyAdd(va, vb0, acc00); + acc01 = Fma.MultiplyAdd(va, vb1, acc01); + acc02 = Fma.MultiplyAdd(va, vb2, acc02); + + // Row 1. + va = Vector256.Load(aRow1 + i); + acc10 = Fma.MultiplyAdd(va, vb0, acc10); + acc11 = Fma.MultiplyAdd(va, vb1, acc11); + acc12 = Fma.MultiplyAdd(va, vb2, acc12); + + // Row 2. + va = Vector256.Load(aRow2 + i); + acc20 = Fma.MultiplyAdd(va, vb0, acc20); + acc21 = Fma.MultiplyAdd(va, vb1, acc21); + acc22 = Fma.MultiplyAdd(va, vb2, acc22); + + // Row 3. + va = Vector256.Load(aRow3 + i); + acc30 = Fma.MultiplyAdd(va, vb0, acc30); + acc31 = Fma.MultiplyAdd(va, vb1, acc31); + acc32 = Fma.MultiplyAdd(va, vb2, acc32); + } + + // Horizontal-reduce each accumulator into the corresponding C cell. + float s00 = HorizontalSum(acc00), s01 = HorizontalSum(acc01), s02 = HorizontalSum(acc02); + float s10 = HorizontalSum(acc10), s11 = HorizontalSum(acc11), s12 = HorizontalSum(acc12); + float s20 = HorizontalSum(acc20), s21 = HorizontalSum(acc21), s22 = HorizontalSum(acc22); + float s30 = HorizontalSum(acc30), s31 = HorizontalSum(acc31), s32 = HorizontalSum(acc32); + + // K-tail (scalar). + for (int i = kVec; i < k; i++) + { + float b0 = bTok0[i], b1 = bTok1[i], b2 = bTok2[i]; + float a0 = aRow0[i]; s00 += a0 * b0; s01 += a0 * b1; s02 += a0 * b2; + float a1 = aRow1[i]; s10 += a1 * b0; s11 += a1 * b1; s12 += a1 * b2; + float a2 = aRow2[i]; s20 += a2 * b0; s21 += a2 * b1; s22 += a2 * b2; + float a3 = aRow3[i]; s30 += a3 * b0; s31 += a3 * b1; s32 += a3 * b2; + } + + c[(long)(tStart + 0) * m + rStart + 0] = s00; + c[(long)(tStart + 0) * m + rStart + 1] = s10; + c[(long)(tStart + 0) * m + rStart + 2] = s20; + c[(long)(tStart + 0) * m + rStart + 3] = s30; + + c[(long)(tStart + 1) * m + rStart + 0] = s01; + c[(long)(tStart + 1) * m + rStart + 1] = s11; + c[(long)(tStart + 1) * m + rStart + 2] = s21; + c[(long)(tStart + 1) * m + rStart + 3] = s31; + + c[(long)(tStart + 2) * m + rStart + 0] = s02; + c[(long)(tStart + 2) * m + rStart + 1] = s12; + c[(long)(tStart + 2) * m + rStart + 2] = s22; + c[(long)(tStart + 2) * m + rStart + 3] = s32; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static float DotAvx2(float* a, float* b, int k, int kVec) + { + Vector256 acc = Vector256.Zero; + for (int i = 0; i < kVec; i += LaneAvx2) + acc = Fma.MultiplyAdd(Vector256.Load(a + i), Vector256.Load(b + i), acc); + float sum = HorizontalSum(acc); + for (int i = kVec; i < k; i++) + sum += a[i] * b[i]; + return sum; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static float HorizontalSum(Vector256 v) + { + // 8 → 4 via cross-lane add, then 4 → 2 via hadd, then 2 → 1 via hadd. + Vector128 lo = v.GetLower(); + Vector128 hi = v.GetUpper(); + Vector128 sum128 = Sse.Add(lo, hi); + sum128 = Sse3.HorizontalAdd(sum128, sum128); + sum128 = Sse3.HorizontalAdd(sum128, sum128); + return sum128.ToScalar(); + } +} diff --git a/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs b/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs index d783b11d..3986604a 100644 --- a/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs +++ b/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs @@ -412,4 +412,149 @@ private static void FillRandomQ8_0Blocks(byte* ptr, int blockCount, Random rng) ((sbyte*)(block + 2))[i] = (sbyte)rng.Next(-127, 128); } } + + // ──────────────────── F32 outer-product GEMM ──────────────────── + // + // These exercise the new `OuterProductGemm.OuterProductGemmF32` kernel + // (companion to the Q8_0 outer-product kernels above). Parity is checked + // against: + // 1. `MatMul.GemmF32Scalar` — the canonical reference, no FMA. + // 2. `MatMul.GemmF32` — the production tiled path, FMA-enabled. + // 3. `OuterProductGemm.OuterProductGemmF32Scalar` — internal cross-check. + // + // FMA vs separate-mul-add reorders the rounding, so we use a relative + // tolerance that scales with K (the contraction dim accumulates error + // linearly in K for uniformly-distributed inputs). + + [Theory] + [InlineData(4, 3, 8)] // smallest fully-vectorisable tile + [InlineData(4, 3, 64)] // single AVX2 group × 8 + [InlineData(8, 6, 64)] // 2 row-tiles × 2 token-tiles + [InlineData(12, 9, 128)] // 3 row-tiles × 3 token-tiles, K=128 + [InlineData(16, 12, 256)] // 4 row-tiles × 4 token-tiles + [InlineData(32, 16, 512)] // prefill-shaped + public void OuterProductGemmF32_Scalar_MatchesReference(int m, int n, int k) + { + RunF32ParityCase(m, n, k, useVector: false); + } + + [Theory] + [InlineData(4, 3, 8)] + [InlineData(4, 3, 64)] + [InlineData(8, 6, 64)] + [InlineData(12, 9, 128)] + [InlineData(16, 12, 256)] + [InlineData(32, 16, 512)] + [InlineData(128, 32, 1024)] // bigger prefill + public void OuterProductGemmF32_Avx2_MatchesReference(int m, int n, int k) + { + RunF32ParityCase(m, n, k, useVector: true); + } + + // Tail-handling cases: shapes where M, N, or K are not multiples of the + // microkernel tile (4 rows × 3 tokens × 8-wide vector). + + [Theory] + [InlineData(5, 3, 64)] // row tail (m % 4 != 0) + [InlineData(7, 3, 64)] + [InlineData(4, 4, 64)] // token tail (n % 3 != 0) + [InlineData(4, 5, 64)] + [InlineData(7, 5, 64)] // both tails + [InlineData(13, 11, 33)] // K tail (k % 8 != 0) + row + token tails + [InlineData(17, 9, 65)] // K tail with FMA group + public void OuterProductGemmF32_HandlesAllTails(int m, int n, int k) + { + RunF32ParityCase(m, n, k, useVector: true); + RunF32ParityCase(m, n, k, useVector: false); + } + + // Edge cases. + + [Theory] + [InlineData(1, 1, 1)] // degenerate scalar + [InlineData(1, 3, 64)] // single row — all-row-tail + [InlineData(4, 1, 64)] // single token — all-token-tail + [InlineData(1, 1, 4096)] // pure inner product + [InlineData(4, 3, 1)] // K=1 — every tile collapses to single FMA + public void OuterProductGemmF32_EdgeCases(int m, int n, int k) + { + RunF32ParityCase(m, n, k, useVector: true); + RunF32ParityCase(m, n, k, useVector: false); + } + + private static void RunF32ParityCase(int m, int n, int k, bool useVector) + { + if (useVector && (!Avx2.IsSupported || !Fma.IsSupported)) + { + // Vector path falls back to scalar; covered by the scalar case. + return; + } + + var rng = new Random(0xC0FFEE ^ (m * 31 + n) * 31 + k); + long aLen = (long)m * k; + long bLen = (long)n * k; + long cLen = (long)n * m; + + float* a = (float*)NativeMemory.AlignedAlloc((nuint)(aLen * sizeof(float)), 64); + float* b = (float*)NativeMemory.AlignedAlloc((nuint)(bLen * sizeof(float)), 64); + float* cOuter = (float*)NativeMemory.AlignedAlloc((nuint)(cLen * sizeof(float)), 64); + float* cRefScalar = (float*)NativeMemory.AlignedAlloc((nuint)(cLen * sizeof(float)), 64); + float* cRefGemm = (float*)NativeMemory.AlignedAlloc((nuint)(cLen * sizeof(float)), 64); + + try + { + // Use range [-1, 1] — typical of post-normalisation activations and + // weight distributions. Larger magnitudes pile up error sooner. + for (long i = 0; i < aLen; i++) a[i] = rng.NextSingle() * 2f - 1f; + for (long i = 0; i < bLen; i++) b[i] = rng.NextSingle() * 2f - 1f; + + // Reference: `MatMul.GemmF32Scalar` — canonical scalar inner product. + MatMul.GemmF32Scalar(a, b, cRefScalar, m, k, n); + + // Reference: `MatMul.GemmF32` — production path (FMA inside TensorPrimitives). + MatMul.GemmF32(a, b, cRefGemm, m, k, n); + + if (useVector) + OuterProductGemm.OuterProductGemmF32(a, b, cOuter, m, k, n); + else + OuterProductGemm.OuterProductGemmF32Scalar(a, b, cOuter, m, k, n); + + if (!useVector) + { + // Scalar path: accumulates the same terms in the same ascending + // index order as `MatMul.GemmF32Scalar`, and .NET does not + // auto-contract `mul + add` to FMA — so bit-exact equality is + // achievable and required. + for (long i = 0; i < cLen; i++) + { + Assert.Equal(cRefScalar[i], cOuter[i]); + } + } + else + { + // Vector path: AVX2 FMA + horizontal reduction reorders the + // summation vs the scalar inner product. Use an absolute + // tolerance that scales with √K — for unit-magnitude operands + // and float32 (ULP ≈ 1.2e-7) the standard error of an unbiased + // K-term sum grows as √K. + float absTol = 4e-6f * MathF.Sqrt(k); + + for (long i = 0; i < cLen; i++) + { + float diff = MathF.Abs(cOuter[i] - cRefGemm[i]); + Assert.True( + diff <= absTol, + $"shape m={m} n={n} k={k} idx={i}: outer={cOuter[i]:R} ref={cRefGemm[i]:R} diff={diff:R} tol={absTol:R}"); + } + } + } + finally + { + NativeMemory.AlignedFree(a); + NativeMemory.AlignedFree(b); + NativeMemory.AlignedFree(cOuter); + NativeMemory.AlignedFree(cRefScalar); + NativeMemory.AlignedFree(cRefGemm); + } + } } From 7c34e46dcc262566e54181c9341d97c6f7622e30 Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 11:08:11 +0100 Subject: [PATCH 2/2] kernels(cpu)(matmul): address Copilot review on the F32 outer-product GEMM (#312) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - HorizontalSum: drop the Sse.Add/Sse3.HorizontalAdd pair for a cross-platform `Vector128.Sum(lo + hi)` reduction. The Avx2/Fma guard at the entrypoint already implies Sse3 via the .NET ISA hierarchy, so this was not a live throw, but the reduction now carries no ISA requirement of its own and matches the "prefer cross-platform Vector128/Vector256" convention. Parity tests confirm the vector path still matches the reference within tolerance. - Tests: the header comment claimed a *relative* tolerance scaling linearly in K while the code uses an *absolute* tolerance scaling with √K. Comment corrected to describe what the code actually does. - Tests: vector-only AVX2 cases now use [SkippableTheory] + Skip.IfNot instead of an assertion-free early `return`, so a run on a non-AVX2 agent reports "skipped" rather than a false "passed". The mixed tail/edge-case theories keep their unconditional scalar assertions and gate only the vector half. Co-Authored-By: Claude Opus 5 (1M context) --- src/DotLLM.Cpu/Kernels/OuterProductGemm.cs | 12 +++---- .../Cpu/Kernels/OuterProductGemmTests.cs | 33 ++++++++++++------- 2 files changed, 26 insertions(+), 19 deletions(-) diff --git a/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs b/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs index 22fae4be..61d0b10f 100644 --- a/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs +++ b/src/DotLLM.Cpu/Kernels/OuterProductGemm.cs @@ -352,12 +352,10 @@ private static float DotAvx2(float* a, float* b, int k, int kVec) [MethodImpl(MethodImplOptions.AggressiveInlining)] private static float HorizontalSum(Vector256 v) { - // 8 → 4 via cross-lane add, then 4 → 2 via hadd, then 2 → 1 via hadd. - Vector128 lo = v.GetLower(); - Vector128 hi = v.GetUpper(); - Vector128 sum128 = Sse.Add(lo, hi); - sum128 = Sse3.HorizontalAdd(sum128, sum128); - sum128 = Sse3.HorizontalAdd(sum128, sum128); - return sum128.ToScalar(); + // 8 → 4 via a cross-lane add, then a 4-wide reduction. Expressed on the + // cross-platform Vector128 surface (not Sse/Sse3 intrinsics) so the + // reduction needs no ISA guard of its own beyond the Avx2/Fma check the + // public entrypoint already makes. + return Vector128.Sum(v.GetLower() + v.GetUpper()); } } diff --git a/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs b/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs index 3986604a..692d0559 100644 --- a/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs +++ b/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs @@ -422,9 +422,11 @@ private static void FillRandomQ8_0Blocks(byte* ptr, int blockCount, Random rng) // 2. `MatMul.GemmF32` — the production tiled path, FMA-enabled. // 3. `OuterProductGemm.OuterProductGemmF32Scalar` — internal cross-check. // - // FMA vs separate-mul-add reorders the rounding, so we use a relative - // tolerance that scales with K (the contraction dim accumulates error - // linearly in K for uniformly-distributed inputs). + // The scalar path accumulates in the same order as the reference and is + // compared bit-exactly. The vector path uses FMA plus a horizontal + // reduction, which reorders the rounding, so it is compared with an + // *absolute* tolerance that scales with √K — for unit-magnitude operands + // the error of an unbiased K-term float32 sum grows as √K, not linearly. [Theory] [InlineData(4, 3, 8)] // smallest fully-vectorisable tile @@ -438,7 +440,15 @@ public void OuterProductGemmF32_Scalar_MatchesReference(int m, int n, int k) RunF32ParityCase(m, n, k, useVector: false); } - [Theory] + // AVX2/FMA is the only hardware requirement of the vector path. Where it is + // absent the kernel falls back to scalar, so a vector-only case has nothing + // to assert — report it as *skipped* rather than passed, so a run on a + // non-AVX2 agent is visibly distinguishable from a real validation. + private static bool Avx2FmaSupported => Avx2.IsSupported && Fma.IsSupported; + + private const string NoAvx2 = "AVX2/FMA not supported on this machine"; + + [SkippableTheory] [InlineData(4, 3, 8)] [InlineData(4, 3, 64)] [InlineData(8, 6, 64)] @@ -448,6 +458,7 @@ public void OuterProductGemmF32_Scalar_MatchesReference(int m, int n, int k) [InlineData(128, 32, 1024)] // bigger prefill public void OuterProductGemmF32_Avx2_MatchesReference(int m, int n, int k) { + Skip.IfNot(Avx2FmaSupported, NoAvx2); RunF32ParityCase(m, n, k, useVector: true); } @@ -464,7 +475,9 @@ public void OuterProductGemmF32_Avx2_MatchesReference(int m, int n, int k) [InlineData(17, 9, 65)] // K tail with FMA group public void OuterProductGemmF32_HandlesAllTails(int m, int n, int k) { - RunF32ParityCase(m, n, k, useVector: true); + // Scalar half always asserts, so this stays a plain Theory. + if (Avx2FmaSupported) + RunF32ParityCase(m, n, k, useVector: true); RunF32ParityCase(m, n, k, useVector: false); } @@ -478,18 +491,14 @@ public void OuterProductGemmF32_HandlesAllTails(int m, int n, int k) [InlineData(4, 3, 1)] // K=1 — every tile collapses to single FMA public void OuterProductGemmF32_EdgeCases(int m, int n, int k) { - RunF32ParityCase(m, n, k, useVector: true); + // Scalar half always asserts, so this stays a plain Theory. + if (Avx2FmaSupported) + RunF32ParityCase(m, n, k, useVector: true); RunF32ParityCase(m, n, k, useVector: false); } private static void RunF32ParityCase(int m, int n, int k, bool useVector) { - if (useVector && (!Avx2.IsSupported || !Fma.IsSupported)) - { - // Vector path falls back to scalar; covered by the scalar case. - return; - } - var rng = new Random(0xC0FFEE ^ (m * 31 + n) * 31 + k); long aLen = (long)m * k; long bLen = (long)n * k;