SIMD Matrix Operations¶

const std = @import("std");
const builtin = @import("builtin");

pub const SimdConfig = struct {
    /// Number of f32 elements processed per SIMD iteration.
    vector_len: comptime_int,
    /// Human-readable label for diagnostics.
    label: []const u8,

    pub fn detect() SimdConfig {
        const features = builtin.cpu.features;
        const Tag = std.Target.x86.Feature;

        if (comptime features.isEnabled(@intFromEnum(Tag.avx2)) and
            features.isEnabled(@intFromEnum(Tag.fma)))
        {
            return .{ .vector_len = 8, .label = "AVX2+FMA" };
        }
        if (comptime features.isEnabled(@intFromEnum(Tag.avx))) {
            return .{ .vector_len = 8, .label = "AVX" };
        }
        if (comptime features.isEnabled(@intFromEnum(Tag.sse2))) {
            return .{ .vector_len = 4, .label = "SSE2" };
        }

        // ARM NEON detection
        if (comptime builtin.cpu.arch == .aarch64) {
            return .{ .vector_len = 4, .label = "NEON" };
        }

        // Scalar fallback
        return .{ .vector_len = 1, .label = "scalar" };
    }
};

const config = comptime SimdConfig.detect();

const VEC_LEN = config.vector_len;
const VecF32 = @Vector(VEC_LEN, f32);

/// Compute dot product of two f32 slices using SIMD.
pub fn dotProduct(a: []const f32, b: []const f32) f32 {
    std.debug.assert(a.len == b.len);
    const n = a.len;

    // --- SIMD accumulation ---
    var acc: VecF32 = @splat(0.0);
    var i: usize = 0;

    // Step 1: Process VEC_LEN elements per iteration.
    while (i + VEC_LEN <= n) : (i += VEC_LEN) {
        // Step 2: Load VEC_LEN-wide slices into SIMD registers.
        const va: VecF32 = a[i..][0..VEC_LEN].*;
        const vb: VecF32 = b[i..][0..VEC_LEN].*;

        // Step 3: Fused multiply-add -- acc += va * vb.
        acc = @mulAdd(VecF32, va, vb, acc);
    }

    // Step 4: Horizontal sum of accumulator lanes.
    var sum: f32 = @reduce(.Add, acc);

    // Step 5: Scalar tail for remaining elements.
    while (i < n) : (i += 1) {
        sum += a[i] * b[i];
    }

    return sum;
}

/// SIMD matrix multiply for small matrices.
/// A: M x K, B: K x N, C: M x N (all row-major f32).
pub fn matmulSIMD_f32_simple(
    A: []const f32, B: []const f32, C: []f32,
    M: usize, K: usize, N: usize,
) void {
    for (0..M) |i| {
        // Zero the output row.
        @memset(C[i * N ..][0..N], 0);

        for (0..K) |k| {
            // Broadcast A[i][k] to all SIMD lanes.
            const a_ik: VecF32 = @splat(A[i * K + k]);

            var j: usize = 0;
            while (j + VEC_LEN <= N) : (j += VEC_LEN) {
                // Load a VEC_LEN-wide segment of B[k][j..j+VEC_LEN].
                const b_row: VecF32 = B[k * N + j ..][0..VEC_LEN].*;

                // Load current accumulator for C[i][j..j+VEC_LEN].
                var c_row: VecF32 = C[i * N + j ..][0..VEC_LEN].*;

                // FMA: C[i][j..] += A[i][k] * B[k][j..]
                c_row = @mulAdd(VecF32, a_ik, b_row, c_row);

                // Store back.
                C[i * N + j ..][0..VEC_LEN].* = c_row;
            }

            // Scalar tail.
            while (j < N) : (j += 1) {
                C[i * N + j] += A[i * K + k] * B[k * N + j];
            }
        }
    }
}

const BLOCK_SIZE: usize = 48;

/// Cache-blocked SIMD matrix multiply for large matrices.
pub fn matmulSIMD_f32_blocked(
    A: []const f32, B: []const f32, C: []f32,
    M: usize, K: usize, N: usize,
) void {
    @memset(C[0 .. M * N], 0);

    var ii: usize = 0;
    while (ii < M) : (ii += BLOCK_SIZE) {
        const i_end = @min(ii + BLOCK_SIZE, M);

        var jj: usize = 0;
        while (jj < N) : (jj += BLOCK_SIZE) {
            const j_end = @min(jj + BLOCK_SIZE, N);

            var kk: usize = 0;
            while (kk < K) : (kk += BLOCK_SIZE) {
                const k_end = @min(kk + BLOCK_SIZE, K);

                // Micro-kernel: multiply A[ii..i_end, kk..k_end]
                //                     by B[kk..k_end, jj..j_end]
                //            accumulate into C[ii..i_end, jj..j_end].
                for (ii..i_end) |i| {
                    for (kk..k_end) |k| {
                        const a_ik: VecF32 = @splat(A[i * K + k]);

                        var j: usize = jj;
                        while (j + VEC_LEN <= j_end) : (j += VEC_LEN) {
                            const b_vec: VecF32 = B[k * N + j ..][0..VEC_LEN].*;
                            var c_vec: VecF32 = C[i * N + j ..][0..VEC_LEN].*;
                            c_vec = @mulAdd(VecF32, a_ik, b_vec, c_vec);
                            C[i * N + j ..][0..VEC_LEN].* = c_vec;
                        }

                        // Scalar tail within the block.
                        while (j < j_end) : (j += 1) {
                            C[i * N + j] += A[i * K + k] * B[k * N + j];
                        }
                    }
                }
            }
        }
    }
}

// In src/core/matrix_ops.zig
pub const SimdConfig = struct {
    vector_len: comptime_int,
    label: []const u8,
    pub fn detect() SimdConfig;
};

pub fn dotProduct(a: []const f32, b: []const f32) f32;

pub fn matmulSIMD_f32_simple(
    A: []const f32, B: []const f32, C: []f32,
    M: usize, K: usize, N: usize,
) void;

pub fn matmulSIMD_f32_blocked(
    A: []const f32, B: []const f32, C: []f32,
    M: usize, K: usize, N: usize,
) void;