tiny.simd.matvec

1const std = @import("std");2const bfloat = @import("bfloat.zig");3const dot = @import("dot.zig");45pub fn compute(6    comptime D: type,7    matrix: []const D.Lane,8    vector: []const D.Lane,9    output: []D.Lane,10) void {11    requireFloat(D.Lane);12    validateShape(matrix.len, vector.len, output.len);13    for (output, 0..) |*result, row| {14        const begin = row * vector.len;15        result.* = dot.compute(D, matrix[begin .. begin + vector.len], vector);16    }17}1819pub fn computeAdd(20    comptime D: type,21    matrix: []const D.Lane,22    vector: []const D.Lane,23    add: []const D.Lane,24    output: []D.Lane,25) void {26    requireFloat(D.Lane);27    std.debug.assert(add.len >= output.len);28    compute(D, matrix, vector, output);29    for (output, add[0..output.len]) |*result, addend| {30        result.* = addScalar(D.Lane, result.*, addend);31    }32}3334pub fn computeBFloatF32(35    comptime D: type,36    matrix: []const bfloat.BFloat16,37    vector: []const f32,38    output: []f32,39) void {40    if (comptime D.Lane != f32) @compileError("mixed matvec requires an f32 descriptor");41    validateShape(matrix.len, vector.len, output.len);42    for (output, 0..) |*result, row| {43        const begin = row * vector.len;44        result.* = dot.computeF32BFloat(D, vector, matrix[begin .. begin + vector.len]);45    }46}4748pub fn computeBFloatF32Add(49    comptime D: type,50    matrix: []const bfloat.BFloat16,51    vector: []const f32,52    add: []const f32,53    output: []f32,54) void {55    std.debug.assert(add.len >= output.len);56    computeBFloatF32(D, matrix, vector, output);57    for (output, add[0..output.len]) |*result, addend| result.* += addend;58}5960pub fn computeBFloat(61    comptime D: type,62    matrix: []const bfloat.BFloat16,63    vector: []const bfloat.BFloat16,64    output: []f32,65) void {66    requireBFloatTag(D);67    validateShape(matrix.len, vector.len, output.len);68    for (output, 0..) |*result, row| {69        const begin = row * vector.len;70        result.* = dot.computeBFloat(D, matrix[begin .. begin + vector.len], vector);71    }72}7374pub fn computeBFloatAdd(75    comptime D: type,76    matrix: []const bfloat.BFloat16,77    vector: []const bfloat.BFloat16,78    add: []const bfloat.BFloat16,79    output: []f32,80) void {81    std.debug.assert(add.len >= output.len);82    computeBFloat(D, matrix, vector, output);83    for (output, add[0..output.len]) |*result, addend| result.* += addend.toF32();84}8586fn validateShape(matrix_len: usize, columns: usize, rows: usize) void {87    const area = std.math.mul(usize, rows, columns) catch @panic("matrix shape exceeds usize");88    std.debug.assert(matrix_len >= area);89}9091fn addScalar(comptime T: type, a: T, b: T) T {92    const sum = @as(f32, @floatCast(a)) + @as(f32, @floatCast(b));93    return @floatCast(sum);94}9596fn requireFloat(comptime T: type) void {97    if (T != f16 and T != f32 and T != f64) {98        @compileError("matvec requires f16/f32/f64 lanes");99    }100}101102fn requireBFloatTag(comptime D: type) void {103    if (comptime !@hasDecl(D, "is_bfloat16") or !D.is_bfloat16) {104        @compileError("bfloat16 matvec requires a bfloat16 descriptor");105    }106}107108fn close(comptime T: type, expected: T, actual: T, scale: T) bool {109    const tolerance = scale * @max(@abs(expected), @as(T, 1));110    return @abs(expected - actual) <= tolerance;111}112113fn verifySameType(comptime T: type) !void {114    const simd = @import("root.zig");115    const D = simd.FixedTag(T, 8);116    const rows: usize = 19;117    const columns: usize = 37;118    var matrix: [rows * columns]T = undefined;119    var vector: [columns]T = undefined;120    var add: [rows]T = undefined;121    for (&matrix, 0..) |*value, index| {122        const integer = @as(i32, @intCast(index % 19)) - 9;123        value.* = @as(T, @floatFromInt(integer)) * @as(T, 0.125);124    }125    for (&vector, 0..) |*value, index| {126        const integer = @as(i32, @intCast(index % 13)) - 6;127        value.* = @as(T, @floatFromInt(integer)) * @as(T, 0.25);128    }129    for (&add, 0..) |*value, index| value.* = @as(T, @floatFromInt(index)) * @as(T, 0.5);130    var output: [rows]T = undefined;131    compute(D, &matrix, &vector, &output);132    for (output, 0..) |actual, row| {133        var expected: T = 0;134        for (matrix[row * columns ..][0..columns], vector) |a, b| {135            expected = @mulAdd(T, a, b, expected);136        }137        try std.testing.expect(close(T, expected, actual, 64 * std.math.floatEps(T)));138    }139    computeAdd(D, &matrix, &vector, &add, &output);140    for (output, 0..) |actual, row| {141        const product = dot.compute(D, matrix[row * columns ..][0..columns], &vector);142        const expected = addScalar(T, product, add[row]);143        try std.testing.expect(close(T, expected, actual, 4 * std.math.floatEps(T)));144    }145}146147test "Highway same-type matvec and additive forms handle awkward shapes" {148    inline for (.{ f16, f32, f64 }) |T| try verifySameType(T);149}150151test "Highway bfloat matrix forms widen outputs to f32" {152    const simd = @import("root.zig");153    const DF = simd.FixedTag(f32, 8);154    const DB = simd.BFloat16Tag(16);155    const rows: usize = 11;156    const columns: usize = 35;157    var matrix: [rows * columns]bfloat.BFloat16 = undefined;158    var vector_f32: [columns]f32 = undefined;159    var vector_bf: [columns]bfloat.BFloat16 = undefined;160    var add_f32: [rows]f32 = undefined;161    var add_bf: [rows]bfloat.BFloat16 = undefined;162    for (&matrix, 0..) |*value, index| {163        const integer = @as(i32, @intCast(index % 17)) - 8;164        value.* = bfloat.BFloat16.fromF32(@as(f32, @floatFromInt(integer)) * 0.125);165    }166    for (&vector_f32, &vector_bf, 0..) |*fv, *bv, index| {167        const integer = @as(i32, @intCast(index % 11)) - 5;168        fv.* = @as(f32, @floatFromInt(integer)) * 0.25;169        bv.* = bfloat.BFloat16.fromF32(fv.*);170    }171    for (&add_f32, &add_bf, 0..) |*fv, *bv, index| {172        fv.* = @as(f32, @floatFromInt(index)) * 0.5;173        bv.* = bfloat.BFloat16.fromF32(fv.*);174    }175    var output: [rows]f32 = undefined;176    computeBFloatF32Add(DF, &matrix, &vector_f32, &add_f32, &output);177    for (output, 0..) |actual, row| {178        const product = dot.computeF32BFloat(179            DF,180            &vector_f32,181            matrix[row * columns ..][0..columns],182        );183        try std.testing.expect(close(184            f32,185            product + add_f32[row],186            actual,187            4 * std.math.floatEps(f32),188        ));189    }190    computeBFloatAdd(DB, &matrix, &vector_bf, &add_bf, &output);191    for (output, 0..) |actual, row| {192        const product = dot.computeBFloat(193            DB,194            matrix[row * columns ..][0..columns],195            &vector_bf,196        );197        try std.testing.expect(close(198            f32,199            product + add_bf[row].toF32(),200            actual,201            4 * std.math.floatEps(f32),202        ));203    }204}205206test "Highway matvec accepts empty dimensions without touching extra output" {207    const simd = @import("root.zig");208    const D = simd.FixedTag(f32, 4);209    var output = [_]f32{ 7, 7, 7 };210    compute(D, &.{}, &.{}, &output);211    try std.testing.expectEqualSlices(f32, &.{ 0, 0, 0 }, &output);212    compute(D, &.{}, &.{}, output[0..0]);213}214215test "Highway AVX2 matvec oracle matches additive forms" {216    const simd = @import("root.zig");217    const DF = simd.FixedTag(f32, 8);218    const DB = simd.BFloat16Tag(16);219    const rows: usize = 5;220    const columns: usize = 35;221    var matrix: [rows * columns]f32 = undefined;222    var vector: [columns]f32 = undefined;223    var add: [rows]f32 = undefined;224    var matrix_bf: [rows * columns]bfloat.BFloat16 = undefined;225    var vector_bf: [columns]bfloat.BFloat16 = undefined;226    var add_bf: [rows]bfloat.BFloat16 = undefined;227    for (&matrix, &matrix_bf, 0..) |*value, *bf, index| {228        const integer = @as(i32, @intCast(index % 23)) - 11;229        value.* = @as(f32, @floatFromInt(integer)) * 0.071 +230            @as(f32, @floatFromInt(index % 3)) * 0.003;231        bf.* = bfloat.BFloat16.fromF32(value.*);232    }233    for (&vector, &vector_bf, 0..) |*value, *bf, index| {234        const integer = @as(i32, @intCast(index % 13)) - 6;235        value.* = @as(f32, @floatFromInt(integer)) * -0.113 + 0.007;236        bf.* = bfloat.BFloat16.fromF32(value.*);237    }238    for (&add, &add_bf, 0..) |*value, *bf, index| {239        value.* = @as(f32, @floatFromInt(index)) * 0.19 - 0.23;240        bf.* = bfloat.BFloat16.fromF32(value.*);241    }242    var output: [rows]f32 = undefined;243    computeAdd(DF, &matrix, &vector, &add, &output);244    try expectBits(&.{ 0xbfbb_9e28, 0x3e15_c7cd, 0xbec5_058a, 0x3fcb_8106, 0x3f03_ffef }, &output);245    computeBFloatF32Add(DF, &matrix_bf, &vector, &add, &output);246    try expectBits(&.{ 0xbfbb_c911, 0x3e14_3edd, 0xbec4_f2da, 0x3fcb_9624, 0x3f03_cb5e }, &output);247    computeBFloatAdd(DB, &matrix_bf, &vector_bf, &add_bf, &output);248    try expectBits(&.{ 0xbfbb_3f80, 0x3e15_7200, 0xbec2_bffc, 0x3fcb_ac8a, 0x3f04_cd00 }, &output);249}250251fn expectBits(expected: []const u32, actual: []const f32) !void {252    try std.testing.expectEqual(expected.len, actual.len);253    for (expected, actual) |bits, value| {254        try std.testing.expect(close(255            f32,256            @bitCast(bits),257            value,258            96 * std.math.floatEps(f32),259        ));260    }261}

1pub const matvec = @import("matvec.zig");