lib/simd/src/matmul.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const bfloat = @import("bfloat.zig");
3
4 pub 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 }
21
22 fn 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 }
45
46 fn 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 }
69
70 fn 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 }
75
76 fn 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 }
86
87 test "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 }
102
103 test "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 }
118
119 test "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 }