tiny.simd.matmul

1const std = @import("std");2const bfloat = @import("bfloat.zig");34pub fn perBlock2x2MatMul(5    comptime D: type,6    a: anytype,7    b: @TypeOf(a),8    c: D.Vector,9) D.Vector {10    validateDescriptor(D);11    if (comptime D.Lane == i32) {12        validateVector(@TypeOf(a), i8, D.lane_count * 4);13        return int8MatMul(D, a, b, c);14    }15    if (comptime D.Lane == f32) {16        validateVector(@TypeOf(a), u16, D.lane_count * 2);17        return bfloatMatMul(D, a, b, c);18    }19    @compileError("2x2 block matrix multiplication requires i32 or f32 accumulators");20}2122fn int8MatMul(comptime D: type, a: anytype, b: @TypeOf(a), c: D.Vector) D.Vector {23    const Input = [D.lane_count * 4]i8;24    const Output = [D.lane_count]i32;25    const input_a: Input = @bitCast(a);26    const input_b: Input = @bitCast(b);27    var result: Output = @bitCast(c);28    for (0..D.lane_count / 4) |block| {29        const input = block * 16;30        const output = block * 4;31        for (0..2) |row| {32            for (0..2) |column| {33                var sum: i32 = 0;34                for (0..8) |inner| {35                    sum += @as(i32, input_a[input + row * 8 + inner]) *36                        @as(i32, input_b[input + column * 8 + inner]);37                }38                const index = output + row * 2 + column;39                result[index] = result[index] +% sum;40            }41        }42    }43    return @bitCast(result);44}4546fn bfloatMatMul(comptime D: type, a: anytype, b: @TypeOf(a), c: D.Vector) D.Vector {47    const Input = [D.lane_count * 2]u16;48    const Output = [D.lane_count]f32;49    const input_a: Input = @bitCast(a);50    const input_b: Input = @bitCast(b);51    var result: Output = @bitCast(c);52    for (0..D.lane_count / 4) |block| {53        const input = block * 8;54        const output = block * 4;55        for (0..2) |row| {56            for (0..2) |column| {57                var sum: f32 = 0;58                for (0..4) |inner| {59                    sum += bfloat.f32FromBits(input_a[input + row * 4 + inner]) *60                        bfloat.f32FromBits(input_b[input + column * 4 + inner]);61                }62                const index = output + row * 2 + column;63                result[index] += sum;64            }65        }66    }67    return @bitCast(result);68}6970fn validateDescriptor(comptime D: type) void {71    if (D.lane_count < 4 or D.lane_count % 4 != 0) {72        @compileError("2x2 block matrix multiplication requires four-lane accumulator blocks");73    }74}7576fn validateVector(comptime V: type, comptime Lane: type, comptime lanes: usize) void {77    switch (@typeInfo(V)) {78        .vector => |info| {79            if (info.child != Lane or info.len != lanes) {80                @compileError("2x2 block matrix multiplication received an incompatible input");81            }82        },83        else => @compileError("2x2 block matrix multiplication requires vector inputs"),84    }85}8687test "Highway int8 per-block 2x2 matrix multiplication matches upstream fixture" {88    const simd = @import("root.zig");89    const D = simd.FixedTag(i32, 4);90    const I = simd.FixedTag(i8, 16);91    const a: I.Vector = .{92        -12, -11, -10, -9, -8, -7, -6, -5,93        -4,  -3,  -2,  -1, 0,  1,  2,  3,94    };95    const b: I.Vector = .{96        -9, -8, -7, -6, -5, -4, -3, -2,97        -1, 0,  1,  2,  3,  4,  5,  6,98    };99    const actual = perBlock2x2MatMul(D, a, b, @as(D.Vector, .{ 10, 11, 12, 13 }));100    try std.testing.expect(@reduce(.And, actual == @as(D.Vector, .{ 426, -117, 76, 45 })));101}102103test "Highway bfloat16 per-block 2x2 matrix multiplication matches upstream fixture" {104    const simd = @import("root.zig");105    const D = simd.FixedTag(f32, 4);106    const B = simd.BFloat16Tag(8);107    const a = bfloat.demoteF32(108        B,109        @as(simd.FixedTag(f32, 8).Vector, .{ 0, 0.5, 1, 1.5, 2, 0, 0.5, 1 }),110    );111    const b = bfloat.demoteF32(112        B,113        @as(simd.FixedTag(f32, 8).Vector, .{ 0, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 0 }),114    );115    const actual = perBlock2x2MatMul(D, a, b, @as(D.Vector, .{ 1, 2, 3, 1 }));116    try std.testing.expect(@reduce(.And, actual == @as(D.Vector, .{ 2.75, 4.125, 4, 3.75 })));117}118119test "Highway int8 per-block 2x2 matrix multiplication wraps accumulators" {120    const simd = @import("root.zig");121    const D = simd.FixedTag(i32, 8);122    const I = simd.FixedTag(i8, 32);123    const actual = perBlock2x2MatMul(124        D,125        @as(I.Vector, @splat(127)),126        @as(I.Vector, @splat(127)),127        @as(D.Vector, @splat(std.math.maxInt(i32))),128    );129    try std.testing.expect(@reduce(.And, actual == @as(D.Vector, @splat(-2_147_354_617))));130}

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