Zig 机器学习推理：张量、GEMM 与 int8 量化

const std = @import("std");

pub const Tensor = struct {
    data: []f32,          // 行主序连续存储
    shape: [4]usize,      // 最多 4 维：N/C/H/W
    strides: [4]usize,    // 每个维度跨过的元素数
    ndim: u8,

    pub fn init(allocator: std.mem.Allocator, shape: []const usize) !Tensor {
        var dims = [_]usize{ 1, 1, 1, 1 };
        @memcpy(dims[0..shape.len], shape);
        var strides = [_]usize{ 1, 1, 1, 1 };
        var acc: usize = 1;
        var i: usize = shape.len;
        while (i > 0) {            // 行主序：最后一维步长为 1，向前累乘
            i -= 1;
            strides[i] = acc;
            acc *= dims[i];
        }
        return .{ .data = try allocator.alloc(f32, acc), .shape = dims, .strides = strides, .ndim = @intCast(shape.len) };
    }
};

const Vec8f = @Vector(8, f32);

/// C = A * B，A 为 M x K，B 为 K x N，均为行主序
pub fn gemm(A: []const f32, B: []const f32, C: []f32, M: usize, K: usize, N: usize) void {
    var i: usize = 0;
    while (i < M) : (i += 1) {
        const a_row = A[i * K ..][0..K];
        const c_row = C[i * N ..][0..N];
        var j: usize = 0;
        while (j + 8 <= N) : (j += 8) {
            var acc: Vec8f = @splat(0);
            for (a_row, 0..) |a_ik, k| {
                const b_vec: Vec8f = B[k * N + j ..][0..8].*;
                acc = @mulAdd(Vec8f, @as(Vec8f, @splat(a_ik)), b_vec, acc);
            }
            c_row[j..][0..8].* = acc;
        }
        while (j < N) : (j += 1) {
            var s: f32 = 0;
            for (a_row, 0..) |a_ik, k| s += a_ik * B[k * N + j];
            c_row[j] = s;
        }
    }
}

pub fn relu(x: []f32) void {
    const V = @Vector(8, f32);
    const zero: V = @splat(0);
    var i: usize = 0;
    while (i + 8 <= x.len) : (i += 8) {
        x[i..][0..8].* = @max(@as(V, x[i..][0..8].*), zero);
    }
    while (i < x.len) : (i += 1) x[i] = @max(x[i], 0);
}

pub fn geluApprox(x: f32) f32 {
    // 0.5 * x * (1 + tanh(0.7978845608 * (x + 0.044715 * x^3)))
    const c: f32 = 0.7978845608;
    return 0.5 * x * (1.0 + std.math.tanh(c * (x + 0.044715 * x * x * x)));
}

pub fn silu(x: f32) f32 { return x / (1.0 + @exp(-x)); }

pub fn softmax(logits: []f32) void {
    var max_v = logits[0];
    for (logits[1..]) |v| max_v = @max(max_v, v);   // 1) 求最大值
    var sum: f32 = 0;
    for (logits) |*v| {                             // 2) 减最大值后取指数
        v.* = @exp(v.* - max_v);
        sum += v.*;
    }
    const inv = 1.0 / sum;
    for (logits) |*v| v.* *= inv;                   // 3) 归一化
}

/// 把 NCHW 输入展开成 (N*OH*OW) x (C*KH*KW) 的矩阵
pub fn im2col(input: []const f32, out: []f32, dims: Dims) void {
    const cols = dims.c * dims.kh * dims.kw;
    const oh = (dims.h + 2 * dims.ph - dims.kh) / dims.sh + 1;
    const ow = (dims.w + 2 * dims.pw - dims.kw) / dims.sw + 1;
    for (0..dims.n) |bi| for (0..oh) |oy| for (0..ow) |ox| {
        const row = ((bi * oh) + oy) * ow + ox;
        var col: usize = 0;
        for (0..dims.c) |ci| for (0..dims.kh) |ky| for (0..dims.kw) |kx| {
            const iy = @as(isize, @intCast(oy * dims.sh + ky)) - @as(isize, @intCast(dims.ph));
            const ix = @as(isize, @intCast(ox * dims.sw + kx)) - @as(isize, @intCast(dims.pw));
            const ok = iy >= 0 and iy < dims.h and ix >= 0 and ix < dims.w;   // 越界补零
            out[row * cols + col] = if (ok)
                input[((bi * dims.c + ci) * dims.h + @as(usize, @intCast(iy))) * dims.w + @as(usize, @intCast(ix))]
            else
                0;
            col += 1;
        };
    };
}

pub const Quantized = struct { data: []i8, scale: f32 };  // 每通道量化时 scale 为 []f32

pub fn quantize(x: []const f32, out: []i8, scale: f32) void {
    for (x, out) |v, *q| q.* = @intFromFloat(std.math.clamp(@round(v / scale), -127.0, 127.0));
}

pub fn dequantize(q: []const i8, out: []f32, scale: f32) void {
    for (q, out) |v, *o| o.* = @as(f32, @floatFromInt(v)) * scale;
}

const I8x16 = @Vector(16, i8);
const I32x16 = @Vector(16, i32);

/// int8 点积：输入 int8，累加到 i32
pub fn dotI8(a: []const i8, b: []const i8) i32 {
    var acc: I32x16 = @splat(0);
    var i: usize = 0;
    while (i + 16 <= a.len) : (i += 16) {
        const wa: @Vector(16, i16) = @as(I8x16, a[i..][0..16].*);  // 先拓宽到 i16
        const wb: @Vector(16, i16) = @as(I8x16, b[i..][0..16].*);
        acc += @as(I32x16, wa) * @as(I32x16, wb);                  // 再累加到 i32
    }
    var sum: i32 = @reduce(.Add, acc);
    while (i < a.len) : (i += 1) sum += @as(i32, a[i]) * @as(i32, b[i]);
    return sum;
}

const ModelHeader = extern struct {
    magic: u32 = 0x5A4D4C31,   // "ZML1"
    version: u32,
    tensor_count: u32,
    node_count: u32,
};

pub fn loadModel(allocator: std.mem.Allocator, path: []const u8) !Model {
    const file = try std.fs.cwd().openFile(path, .{});
    defer file.close();
    var header: ModelHeader = undefined;
    _ = try file.readAll(std.mem.asBytes(&header));   // 直接读进结构体
    if (header.magic != 0x5A4D4C31) return error.BadMagic;
    _ = .{ allocator, header };                       // 再按 tensor_count/node_count 顺序读表与权重
    return .{};
}

const Model = struct {};