diff --git a/benchmarks/DotLLM.Benchmarks/OuterProductDisasmBenchmarks.cs b/benchmarks/DotLLM.Benchmarks/OuterProductDisasmBenchmarks.cs index eec25b86..b15b6826 100644 --- a/benchmarks/DotLLM.Benchmarks/OuterProductDisasmBenchmarks.cs +++ b/benchmarks/DotLLM.Benchmarks/OuterProductDisasmBenchmarks.cs @@ -112,6 +112,24 @@ public void OuterProduct4x3() _output, BlockCount, M); } + /// + /// Outer-product 4×3 via AVX2-VNNI (VPDPBUSD): same output tile as + /// but with the fused multiply-widen-accumulate. + /// Holds 6 live float accumulators (2 rows × 3 tokens) where the AVX2 kernel + /// holds 3 — the register headroom freed by dropping ones+prod. + /// + [Benchmark] + public void OuterProduct4x3_Vnni() + { + byte* groupBase = (byte*)_repackedWeights; + MatMul.OuterProductQ8_0Vnni_4x3( + groupBase, + (byte*)_inputQ8_0, + (byte*)_inputQ8_1, + (byte*)_inputQ8_2, + _output, BlockCount, M); + } + /// /// Full outer-product GEMM: all M rows × 3 tokens. /// 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/MatMul.cs b/src/DotLLM.Cpu/Kernels/MatMul.cs index 741fb812..195849dc 100644 --- a/src/DotLLM.Cpu/Kernels/MatMul.cs +++ b/src/DotLLM.Cpu/Kernels/MatMul.cs @@ -1441,6 +1441,134 @@ internal static void OuterProductQ8_0Avx2_4x3( } } + /// + /// AVX2-VNNI outer-product microkernel for Q8_0 R4 layout — VPDPBUSD-256 fast path. + /// Produces the same 4-row × 3-token output tile as , + /// but replaces the two-instruction maddubs + madd(ones) integer reduction with a + /// single fused AvxVnni.MultiplyWideningAndAdd (VPDPBUSD): unsigned-byte × signed-byte → + /// widening multiply-accumulate into int32, in one instruction. + /// + /// + /// + /// Why this exists (unblocks PR #61). The reverted #61 Q8_0 outer-product tile blew the + /// AVX2 16-YMM budget (~23 YMM) because every (row, token) cell needed a live + /// maddubs product temporary plus a shared ones int16 register on top of the float + /// accumulator. VPDPBUSD eliminates BOTH the ones register and the prod temporary + /// (it fuses the int16-pair multiply and the int32 pairwise add). That freed budget lifts the + /// safe accumulator residency from the AVX2 kernel's conservative 3 (one weight row at a time, + /// see ) up to 6 — two weight rows × three tokens held live + /// — halving the per-row token-vector reload traffic. + /// + /// + /// Register budget (≤ 16 YMM). Processes the 4-row R4 group as two sub-passes of 2 rows: + /// + /// 6 float accumulators (2 rows × 3 tokens) — held across the block loop. + /// 3 token vectors (vx0/vx1/vx2) — held across the inner 2-row loop. + /// Transient per cell: vw, absX (recomputed per row for headroom rather than + /// held), adjW, isum/fsum, scale — staggered lifetimes, peak ≈ 4. + /// + /// Peak ≈ 6 + 3 + 4 = 13 YMM, comfortably inside the 16-YMM file. A full 12-accumulator + /// 4×3 tile (one live accumulator per output cell) remains infeasible at ~18+ YMM even with VNNI; + /// the VNNI win is dropping ones+prod, which is exactly what raises safe residency + /// from 3 → 6. + /// + /// + /// Numerics. Each Q8_0 block carries its own dw·dx scale, so the int32 VPDPBUSD + /// result is folded to float per block (convert → FMA by dx·dw) — identical accumulation + /// order to , so results match to FP rounding. + /// + /// + /// R4-interleaved weight group base (4 rows, blocks interleaved). + /// Token 0 Q8_0 blocks. + /// Token 1 Q8_0 blocks. + /// Token 2 Q8_0 blocks. + /// Output base for this tile; cells written at c[token * cStride + row]. + /// Number of Q8_0 blocks per row (K / 32). + /// Row stride of the output matrix (M). + [SkipLocalsInit] + [MethodImpl(MethodImplOptions.AggressiveOptimization)] + internal static void OuterProductQ8_0Vnni_4x3( + byte* groupBase, byte* x0, byte* x1, byte* x2, + float* c, int blockCount, int cStride) + { + const int wStride = 4 * Q8_0BlockBytes; + + // Two sub-passes over the 4-row R4 group, 2 rows each, holding 6 float + // accumulators (2 rows × 3 tokens) live across the block loop. + for (int rPair = 0; rPair < 4; rPair += 2) + { + int r0 = rPair; + int r1 = rPair + 1; + + // 6 accumulators: a{rowInPair}{token}. + Vector256 a00 = Vector256.Zero, a01 = Vector256.Zero, a02 = Vector256.Zero; + Vector256 a10 = Vector256.Zero, a11 = Vector256.Zero, a12 = Vector256.Zero; + + for (int b = 0; b < blockCount; b++) + { + byte* blockBase = groupBase + b * wStride; + + // Load 3 token blocks (held across the 2-row inner work). + byte* xb0 = x0 + b * Q8_0BlockBytes; + byte* xb1 = x1 + b * Q8_0BlockBytes; + byte* xb2 = x2 + b * Q8_0BlockBytes; + float dx0 = HalfBitsToFloat(xb0); + float dx1 = HalfBitsToFloat(xb1); + float dx2 = HalfBitsToFloat(xb2); + Vector256 vx0 = Unsafe.ReadUnaligned>(xb0 + 2); + Vector256 vx1 = Unsafe.ReadUnaligned>(xb1 + 2); + Vector256 vx2 = Unsafe.ReadUnaligned>(xb2 + 2); + + // Row r0. + { + byte* wBlock = blockBase + r0 * Q8_0BlockBytes; + float dw = HalfBitsToFloat(wBlock); + Vector256 vw = Unsafe.ReadUnaligned>(wBlock + 2); + + // Token 0: VPDPBUSD(0, |x|, sign(x)·w) → int32 partials, fold by dx·dw. + Vector256 isum0 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx0, vx0).AsByte(), Avx2.Sign(vw, vx0)); + a00 = Fma.MultiplyAdd(Vector256.Create(dx0 * dw), Avx.ConvertToVector256Single(isum0), a00); + + Vector256 isum1 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx1, vx1).AsByte(), Avx2.Sign(vw, vx1)); + a01 = Fma.MultiplyAdd(Vector256.Create(dx1 * dw), Avx.ConvertToVector256Single(isum1), a01); + + Vector256 isum2 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx2, vx2).AsByte(), Avx2.Sign(vw, vx2)); + a02 = Fma.MultiplyAdd(Vector256.Create(dx2 * dw), Avx.ConvertToVector256Single(isum2), a02); + } + + // Row r1. + { + byte* wBlock = blockBase + r1 * Q8_0BlockBytes; + float dw = HalfBitsToFloat(wBlock); + Vector256 vw = Unsafe.ReadUnaligned>(wBlock + 2); + + Vector256 isum0 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx0, vx0).AsByte(), Avx2.Sign(vw, vx0)); + a10 = Fma.MultiplyAdd(Vector256.Create(dx0 * dw), Avx.ConvertToVector256Single(isum0), a10); + + Vector256 isum1 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx1, vx1).AsByte(), Avx2.Sign(vw, vx1)); + a11 = Fma.MultiplyAdd(Vector256.Create(dx1 * dw), Avx.ConvertToVector256Single(isum1), a11); + + Vector256 isum2 = AvxVnni.MultiplyWideningAndAdd( + Vector256.Zero, Avx2.Sign(vx2, vx2).AsByte(), Avx2.Sign(vw, vx2)); + a12 = Fma.MultiplyAdd(Vector256.Create(dx2 * dw), Avx.ConvertToVector256Single(isum2), a12); + } + } + + // C[token * cStride + row]. + c[0 * cStride + r0] = HorizontalSumAvx2Float(a00); + c[1 * cStride + r0] = HorizontalSumAvx2Float(a01); + c[2 * cStride + r0] = HorizontalSumAvx2Float(a02); + c[0 * cStride + r1] = HorizontalSumAvx2Float(a10); + c[1 * cStride + r1] = HorizontalSumAvx2Float(a11); + c[2 * cStride + r1] = HorizontalSumAvx2Float(a12); + } + } + /// /// AVX-512 outer-product microkernel for Q8_0 R4 layout. /// Processes 4 weight rows × 6 tokens with 24 ZMM accumulators via dual-block (2 blocks/iteration). @@ -1721,6 +1849,26 @@ internal static void OuterProductGemmQ8_0(byte* repackedWeights, byte* inputQ8, blockCount, c + (long)t * m + baseRow); } } + else if (AvxVnni.IsSupported) + { + // AVX2-VNNI: 4×3 tiles via VPDPBUSD (6 live accumulators, unblocks #61). + int nFull3 = (n / 3) * 3; + for (; t < nFull3; t += 3) + { + OuterProductQ8_0Vnni_4x3( + groupBase, + inputQ8 + (long)t * q8RowBytes, + inputQ8 + (long)(t + 1) * q8RowBytes, + inputQ8 + (long)(t + 2) * q8RowBytes, + c + (long)t * m + baseRow, blockCount, m); + } + // Tail tokens + for (; t < n; t++) + { + VecDotQ8_0Avx2_4RowsR4(groupBase, inputQ8 + (long)t * q8RowBytes, + blockCount, c + (long)t * m + baseRow); + } + } else if (Avx2.IsSupported) { // AVX2: 4×3 tiles @@ -1874,6 +2022,25 @@ private static void OuterProductGemmQ8_0Worker(nint ctxPtr, int threadIdx, int t ctx.BlockCount, ctx.C + (long)t * ctx.M + baseRow); } } + else if (AvxVnni.IsSupported) + { + // AVX2-VNNI: 4×3 tiles via VPDPBUSD (6 live accumulators, unblocks #61). + int nFull3 = (ctx.N / 3) * 3; + for (; t < nFull3; t += 3) + { + OuterProductQ8_0Vnni_4x3( + groupBase, + ctx.InputQ8 + (long)t * q8RowBytes, + ctx.InputQ8 + (long)(t + 1) * q8RowBytes, + ctx.InputQ8 + (long)(t + 2) * q8RowBytes, + ctx.C + (long)t * ctx.M + baseRow, ctx.BlockCount, ctx.M); + } + for (; t < ctx.N; t++) + { + VecDotQ8_0Avx2_4RowsR4(groupBase, ctx.InputQ8 + (long)t * q8RowBytes, + ctx.BlockCount, ctx.C + (long)t * ctx.M + baseRow); + } + } else if (Avx2.IsSupported) { int nFull3 = (ctx.N / 3) * 3; 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..d8b8c544 100644 --- a/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs +++ b/tests/DotLLM.Tests.Unit/Cpu/Kernels/OuterProductGemmTests.cs @@ -144,6 +144,242 @@ public void OuterProductAvx2_4x3_MatchesScalar(int blockCount) } } + // ──────────────────── AVX2-VNNI microkernel ──────────────────── + + [Theory] + [InlineData(1)] + [InlineData(2)] + [InlineData(8)] // K=256 + [InlineData(16)] + [InlineData(18)] // SmolLM-135M: 576/32 + [InlineData(32)] // K=1024 + [InlineData(48)] // K=1536 + [InlineData(128)] // K=4096 + public void OuterProductVnni_4x3_MatchesScalar(int blockCount) + { + if (!AvxVnni.IsSupported) + return; + + var rng = new Random(1234); + int m = 4; + int n = 3; + int rowBytes = blockCount * Q8_0BlockBytes; + + byte* weights = AllocAndFillR4Weights(4, blockCount, rng); + byte*[] xPtrs = new byte*[n]; + for (int t = 0; t < n; t++) + { + xPtrs[t] = (byte*)NativeMemory.AlignedAlloc((nuint)rowBytes, 64); + FillRandomQ8_0Blocks(xPtrs[t], blockCount, rng); + } + + float* cScalar = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + float* cVnni = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + + try + { + MatMul.OuterProductQ8_0Scalar_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], + cScalar, blockCount, m); + + MatMul.OuterProductQ8_0Vnni_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], + cVnni, blockCount, m); + + for (int t = 0; t < n; t++) + for (int r = 0; r < m; r++) + Assert.Equal(cScalar[t * m + r], cVnni[t * m + r], 1e-2f); + } + finally + { + NativeMemory.AlignedFree(weights); + for (int t = 0; t < n; t++) + NativeMemory.AlignedFree(xPtrs[t]); + NativeMemory.AlignedFree(cScalar); + NativeMemory.AlignedFree(cVnni); + } + } + + // VNNI vs the AVX2 maddubs+madd microkernel: both vector paths fold each + // block's int32 sum by dx*dw in the same order, so they should agree to a + // tighter tolerance than either does to the scalar reference. + [Theory] + [InlineData(1)] + [InlineData(16)] + [InlineData(18)] + [InlineData(48)] + [InlineData(128)] + public void OuterProductVnni_4x3_MatchesAvx2(int blockCount) + { + if (!AvxVnni.IsSupported || !Avx2.IsSupported) + return; + + var rng = new Random(777); + int m = 4; + int n = 3; + int rowBytes = blockCount * Q8_0BlockBytes; + + byte* weights = AllocAndFillR4Weights(4, blockCount, rng); + byte*[] xPtrs = new byte*[n]; + for (int t = 0; t < n; t++) + { + xPtrs[t] = (byte*)NativeMemory.AlignedAlloc((nuint)rowBytes, 64); + FillRandomQ8_0Blocks(xPtrs[t], blockCount, rng); + } + + float* cAvx2 = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + float* cVnni = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + + try + { + MatMul.OuterProductQ8_0Avx2_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], + cAvx2, blockCount, m); + + MatMul.OuterProductQ8_0Vnni_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], + cVnni, blockCount, m); + + for (int t = 0; t < n; t++) + for (int r = 0; r < m; r++) + Assert.Equal(cAvx2[t * m + r], cVnni[t * m + r], 1e-3f); + } + finally + { + NativeMemory.AlignedFree(weights); + for (int t = 0; t < n; t++) + NativeMemory.AlignedFree(xPtrs[t]); + NativeMemory.AlignedFree(cAvx2); + NativeMemory.AlignedFree(cVnni); + } + } + + // Discriminating sanity check: the VNNI microkernel parity test must FAIL + // when the kernel output is deliberately perturbed. This guards against a + // vacuous test (e.g. one that compares all-zero buffers, or a tolerance so + // wide that a real tile bug slips through). We corrupt a single (token,row) + // cell of the VNNI result and assert the comparison rejects it. + [Fact] + public void OuterProductVnni_4x3_ParityTestIsDiscriminating() + { + if (!AvxVnni.IsSupported) + return; + + var rng = new Random(31337); + int m = 4; + int n = 3; + int blockCount = 18; + int rowBytes = blockCount * Q8_0BlockBytes; + + byte* weights = AllocAndFillR4Weights(4, blockCount, rng); + byte*[] xPtrs = new byte*[n]; + for (int t = 0; t < n; t++) + { + xPtrs[t] = (byte*)NativeMemory.AlignedAlloc((nuint)rowBytes, 64); + FillRandomQ8_0Blocks(xPtrs[t], blockCount, rng); + } + + float* cScalar = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + float* cVnni = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + + try + { + MatMul.OuterProductQ8_0Scalar_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], cScalar, blockCount, m); + MatMul.OuterProductQ8_0Vnni_4x3( + weights, xPtrs[0], xPtrs[1], xPtrs[2], cVnni, blockCount, m); + + // Unperturbed: every cell must match (and must be meaningfully nonzero, + // so the comparison is not trivially satisfied by zeros). + float maxAbs = 0; + for (int i = 0; i < n * m; i++) + { + Assert.Equal(cScalar[i], cVnni[i], 1e-2f); + maxAbs = MathF.Max(maxAbs, MathF.Abs(cScalar[i])); + } + Assert.True(maxAbs > 1e-3f, "reference output is ~0; parity check would be vacuous"); + + // Perturb one cell (token=1, row=2) by an amount far exceeding tolerance + // and confirm the same equality assertion now fails — i.e. the test + // discriminates broken from correct output. + int idx = 1 * m + 2; + cVnni[idx] += 1.0f; + Assert.ThrowsAny( + () => Assert.Equal(cScalar[idx], cVnni[idx], 1e-2f)); + } + finally + { + NativeMemory.AlignedFree(weights); + for (int t = 0; t < n; t++) + NativeMemory.AlignedFree(xPtrs[t]); + NativeMemory.AlignedFree(cScalar); + NativeMemory.AlignedFree(cVnni); + } + } + + // Full GEMM through the public OuterProductGemmQ8_0 dispatch, which routes + // to the VNNI microkernel on this CPU (AvxVnni.IsSupported). Exercises + // discriminating shapes: multi-block K, full tile, all tail combinations, + // and a large (M>=128, N>=32) shape. + [Theory] + [InlineData(4, 3, 64)] // single full tile, K=64 (2 blocks) + [InlineData(8, 6, 256)] // 2 groups, 2 token-tiles, K=256 (8 blocks) + [InlineData(4, 3, 1024)] // deep K (32 blocks) + [InlineData(7, 5, 64)] // row tail (m%4) + token tail (n%3) + [InlineData(5, 4, 128)] // row tail + token tail, K=128 + [InlineData(13, 11, 256)] // 3 groups + 1 tail row, token tail + [InlineData(128, 32, 64)] // large: 32 groups, 32 tokens + [InlineData(132, 33, 256)] // large with row + token tails, deep K + public void OuterProductGemmVnni_Dispatch_MatchesReference(int m, int n, int k) + { + if (!AvxVnni.IsSupported) + return; + + var rng = new Random(0xBEEF ^ (m * 131 + n) * 17 + k); + int blockCount = k / Q8_0GroupSize; + int q8RowBytes = blockCount * Q8_0BlockBytes; + int fullGroups = m / 4; + int tailRows = m % 4; + + byte* rowMajorWeights = (byte*)NativeMemory.AlignedAlloc((nuint)((long)m * q8RowBytes), 64); + for (int r = 0; r < m; r++) + FillRandomQ8_0Blocks(rowMajorWeights + r * q8RowBytes, blockCount, rng); + + using var repacked = WeightRepacking.RepackR4((nint)rowMajorWeights, QuantizationType.Q8_0, m, k); + + byte* inputQ8 = (byte*)NativeMemory.AlignedAlloc((nuint)((long)n * q8RowBytes), 64); + for (int t = 0; t < n; t++) + FillRandomQ8_0Blocks(inputQ8 + t * q8RowBytes, blockCount, rng); + + float* cOuter = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + float* cRef = (float*)NativeMemory.AlignedAlloc((nuint)(n * m * sizeof(float)), 64); + + try + { + for (int t = 0; t < n; t++) + { + MatMul.ComputeRowsQ8_0Interleaved( + (byte*)repacked.Ptr, inputQ8 + t * q8RowBytes, + cRef + t * m, fullGroups, tailRows, blockCount); + } + + MatMul.OuterProductGemmQ8_0( + (byte*)repacked.Ptr, inputQ8, cOuter, + fullGroups, tailRows, blockCount, m, n); + + for (int t = 0; t < n; t++) + for (int r = 0; r < m; r++) + Assert.Equal(cRef[t * m + r], cOuter[t * m + r], 1e-2f); + } + finally + { + NativeMemory.AlignedFree(rowMajorWeights); + NativeMemory.AlignedFree(inputQ8); + NativeMemory.AlignedFree(cOuter); + NativeMemory.AlignedFree(cRef); + } + } + // ──────────────────── Full GEMM tests ──────────────────── [Theory] @@ -412,4 +648,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); + } + } }