lib/simd/src/matvec.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const bfloat = @import("bfloat.zig");
3 const dot = @import("dot.zig");
4
5 pub 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 }
18
19 pub 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 }
33
34 pub 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 }
47
48 pub 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 }
59
60 pub 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 }
73
74 pub 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 }
85
86 fn 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 }
90
91 fn 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 }
95
96 fn 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 }
101
102 fn 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 }
107
108 fn 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 }
112
113 fn 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 }
146
147 test "Highway same-type matvec and additive forms handle awkward shapes" {
148 inline for (.{ f16, f32, f64 }) |T| try verifySameType(T);
149 }
150
151 test "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 }
205
206 test "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 }
214
215 test "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 }
250
251 fn 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 }