lib/simd/src/bfloat.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const tag = @import("tag.zig");
3
4 pub const BFloat16 = extern struct {
5 pub const is_bfloat16 = true;
6 pub const Storage = u16;
7
8 bits: u16,
9
10 pub fn fromBits(bits: u16) BFloat16 {
11 return .{ .bits = bits };
12 }
13
14 pub fn fromF32(value: f32) BFloat16 {
15 return .{ .bits = bitsFromF32(value) };
16 }
17
18 pub fn fromF64(value: f64) BFloat16 {
19 const bits: u64 = @bitCast(value);
20 const rounded_bits = (bits & 0xffff_ffc0_0000_0000) |
21 ((bits +% 0x0000_003f_ffff_ffff) & 0x0000_0040_0000_0000);
22 return fromF32(@floatCast(@as(f64, @bitCast(rounded_bits))));
23 }
24
25 pub fn toF32(self: BFloat16) f32 {
26 return f32FromBits(self.bits);
27 }
28
29 pub fn neg(self: BFloat16) BFloat16 {
30 return fromBits(self.bits ^ 0x8000);
31 }
32
33 pub fn add(self: BFloat16, other: BFloat16) BFloat16 {
34 return fromF32(self.toF32() + other.toF32());
35 }
36
37 pub fn sub(self: BFloat16, other: BFloat16) BFloat16 {
38 return fromF32(self.toF32() - other.toF32());
39 }
40
41 pub fn mul(self: BFloat16, other: BFloat16) BFloat16 {
42 return fromF32(self.toF32() * other.toF32());
43 }
44
45 pub fn div(self: BFloat16, other: BFloat16) BFloat16 {
46 return fromF32(self.toF32() / other.toF32());
47 }
48 };
49
50 pub fn Tag(comptime lanes: usize) type {
51 return tag.Descriptor(BFloat16, lanes, 0);
52 }
53
54 pub fn zero(comptime D: type) D.Vector {
55 requireTag(D);
56 return @splat(0);
57 }
58
59 pub fn set(comptime D: type, value: BFloat16) D.Vector {
60 requireTag(D);
61 return @splat(value.bits);
62 }
63
64 pub fn load(comptime D: type, input: []const BFloat16) D.Vector {
65 requireTag(D);
66 std.debug.assert(input.len >= D.lane_count);
67 var result: D.Vector = undefined;
68 inline for (0..D.lane_count) |index| result[index] = input[index].bits;
69 return result;
70 }
71
72 pub fn store(comptime D: type, value: D.Vector, output: []BFloat16) void {
73 requireTag(D);
74 std.debug.assert(output.len >= D.lane_count);
75 inline for (0..D.lane_count) |index| output[index] = BFloat16.fromBits(value[index]);
76 }
77
78 pub fn dup128VecFromValues(
79 comptime D: type,
80 values: [8]BFloat16,
81 ) D.Vector {
82 requireTag(D);
83 var result: D.Vector = undefined;
84 inline for (0..D.lane_count) |index| result[index] = values[index % 8].bits;
85 return result;
86 }
87
88 pub fn neg(comptime D: type, value: D.Vector) D.Vector {
89 requireTag(D);
90 return value ^ @as(D.Vector, @splat(0x8000));
91 }
92
93 pub fn demoteF32(comptime D: type, value: anytype) D.Vector {
94 requireTag(D);
95 validateF32Vector(@TypeOf(value), D.lane_count);
96 var result: D.Vector = undefined;
97 inline for (0..D.lane_count) |index| result[index] = bitsFromF32(value[index]);
98 return result;
99 }
100
101 pub fn promoteF32(comptime D: type, value: anytype) D.Vector {
102 if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");
103 validateBitsVector(@TypeOf(value), D.lane_count);
104 var result: D.Vector = undefined;
105 inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[index]);
106 return result;
107 }
108
109 pub fn promoteLowerF32(comptime D: type, value: anytype) D.Vector {
110 return promoteSelectedF32(D, value, 0);
111 }
112
113 pub fn promoteUpperF32(comptime D: type, value: anytype) D.Vector {
114 return promoteSelectedF32(D, value, D.lane_count);
115 }
116
117 pub fn promoteEvenF32(comptime D: type, value: anytype) D.Vector {
118 return promoteParityF32(D, value, 0);
119 }
120
121 pub fn promoteOddF32(comptime D: type, value: anytype) D.Vector {
122 return promoteParityF32(D, value, 1);
123 }
124
125 pub fn orderedDemote2F32(comptime D: type, a: anytype, b: @TypeOf(a)) D.Vector {
126 requireTag(D);
127 if (comptime vectorLanes(@TypeOf(a)) * 2 != D.lane_count) {
128 @compileError("ordered bfloat16 demotion requires two half-size f32 vectors");
129 }
130 validateF32Vector(@TypeOf(a), D.lane_count / 2);
131 var result: D.Vector = undefined;
132 inline for (0..D.lane_count / 2) |index| {
133 result[index] = bitsFromF32(a[index]);
134 result[D.lane_count / 2 + index] = bitsFromF32(b[index]);
135 }
136 return result;
137 }
138
139 pub fn mulEvenAdd(
140 comptime D: type,
141 a: anytype,
142 b: @TypeOf(a),
143 addend: D.Vector,
144 ) D.Vector {
145 return mulParityAdd(D, a, b, addend, 0);
146 }
147
148 pub fn mulOddAdd(
149 comptime D: type,
150 a: anytype,
151 b: @TypeOf(a),
152 addend: D.Vector,
153 ) D.Vector {
154 return mulParityAdd(D, a, b, addend, 1);
155 }
156
157 pub fn widenMulPairwiseAdd(comptime D: type, a: anytype, b: @TypeOf(a)) D.Vector {
158 return mulOddAdd(D, a, b, mulEvenAdd(D, a, b, @splat(0)));
159 }
160
161 pub fn maskedWidenMulPairwiseAdd(
162 comptime D: type,
163 mask: D.Mask,
164 a: anytype,
165 b: @TypeOf(a),
166 ) D.Vector {
167 return @select(f32, mask, widenMulPairwiseAdd(D, a, b), @as(D.Vector, @splat(0)));
168 }
169
170 pub fn reorderWidenMulAccumulate(
171 comptime D: type,
172 a: anytype,
173 b: @TypeOf(a),
174 sum0: D.Vector,
175 sum1: *D.Vector,
176 ) D.Vector {
177 sum1.* = mulOddAdd(D, a, b, sum1.*);
178 return mulEvenAdd(D, a, b, sum0);
179 }
180
181 pub fn rearrangeToOddPlusEven(comptime D: type, sum0: D.Vector, sum1: D.Vector) D.Vector {
182 if (comptime D.Lane != f32) @compileError("bfloat16 accumulation requires f32 sums");
183 return sum0 + sum1;
184 }
185
186 pub fn bitsFromF32(value: f32) u16 {
187 const bits: u32 = @bitCast(value);
188 const magnitude = bits & 0x7fff_ffff;
189 const increment: u32 = if (magnitude < 0x7f80_0000)
190 0x7fff + ((bits >> 16) & 1)
191 else
192 0;
193 const quiet_nan: u32 = if (magnitude > 0x7f80_0000) 1 << 6 else 0;
194 return @truncate(quiet_nan | ((bits +% increment) >> 16));
195 }
196
197 pub fn f32FromBits(bits: u16) f32 {
198 return @bitCast(@as(u32, bits) << 16);
199 }
200
201 fn promoteSelectedF32(comptime D: type, value: anytype, comptime offset: usize) D.Vector {
202 if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");
203 validateBitsVector(@TypeOf(value), D.lane_count * 2);
204 var result: D.Vector = undefined;
205 inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[offset + index]);
206 return result;
207 }
208
209 fn promoteParityF32(comptime D: type, value: anytype, comptime parity: usize) D.Vector {
210 if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");
211 validateBitsVector(@TypeOf(value), D.lane_count * 2);
212 var result: D.Vector = undefined;
213 inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[index * 2 + parity]);
214 return result;
215 }
216
217 fn mulParityAdd(
218 comptime D: type,
219 a: anytype,
220 b: @TypeOf(a),
221 addend: D.Vector,
222 comptime parity: usize,
223 ) D.Vector {
224 if (comptime D.Lane != f32) @compileError("bfloat16 multiply-add requires f32 destinations");
225 validateBitsVector(@TypeOf(a), D.lane_count * 2);
226 var result = addend;
227 inline for (0..D.lane_count) |index| {
228 result[index] = @mulAdd(
229 f32,
230 f32FromBits(a[index * 2 + parity]),
231 f32FromBits(b[index * 2 + parity]),
232 addend[index],
233 );
234 }
235 return result;
236 }
237
238 fn requireTag(comptime D: type) void {
239 if (comptime !@hasDecl(D, "is_bfloat16") or !D.is_bfloat16) {
240 @compileError("operation requires a bfloat16 descriptor");
241 }
242 }
243
244 fn validateF32Vector(comptime V: type, comptime lanes: usize) void {
245 if (comptime vectorLane(V) != f32 or vectorLanes(V) != lanes) {
246 @compileError("operation requires the expected number of f32 lanes");
247 }
248 }
249
250 fn validateBitsVector(comptime V: type, comptime lanes: usize) void {
251 if (comptime vectorLane(V) != u16 or vectorLanes(V) != lanes) {
252 @compileError("operation requires the expected number of bfloat16 bit lanes");
253 }
254 }
255
256 fn vectorLane(comptime V: type) type {
257 return switch (@typeInfo(V)) {
258 .vector => |info| info.child,
259 else => @compileError("operation requires a vector"),
260 };
261 }
262
263 fn vectorLanes(comptime V: type) usize {
264 return switch (@typeInfo(V)) {
265 .vector => |info| info.len,
266 else => @compileError("operation requires a vector"),
267 };
268 }
269
270 test "Highway bfloat16 scalar conversion rounds to nearest even and preserves classes" {
271 try std.testing.expectEqual(@as(u16, 0x3f80), BFloat16.fromF32(1).bits);
272 try std.testing.expectEqual(@as(u16, 0xbf80), BFloat16.fromF32(-1).bits);
273 try std.testing.expectEqual(@as(u16, 0x3f80), BFloat16.fromF32(1.00390625).bits);
274 try std.testing.expectEqual(@as(u16, 0x3f82), BFloat16.fromF32(1.01171875).bits);
275 try std.testing.expectEqual(@as(u16, 0x4000), BFloat16.fromF32(1.99609375).bits);
276 try std.testing.expectEqual(@as(u16, 0x8000), BFloat16.fromF32(-0.0).bits);
277 try std.testing.expectEqual(@as(u16, 0x7f80), BFloat16.fromF32(std.math.inf(f32)).bits);
278 try std.testing.expect(std.math.isNan(BFloat16.fromF32(std.math.nan(f32)).toF32()));
279 }
280
281 test "Highway bfloat16 descriptors preserve scalar type relations" {
282 const D = tag.FixedTag(BFloat16, 8);
283 try std.testing.expectEqual(BFloat16, tag.TFromD(D));
284 try std.testing.expectEqual(Tag(8), D);
285 try std.testing.expectEqual(@as(usize, 4), tag.Half(D).lane_count);
286 try std.testing.expectEqual(@as(i8, -1), tag.pow2(tag.Half(D)));
287 try std.testing.expectEqual(Tag(8), tag.BlockDFromD(Tag(16)));
288 try std.testing.expectEqual(tag.FixedTag(i16, 8), tag.RebindToSigned(D));
289 try std.testing.expectEqual(tag.FixedTag(u16, 8), tag.RebindToUnsigned(D));
290 try std.testing.expectEqual(tag.FixedTag(f32, 4), tag.RepartitionToWide(D));
291 try std.testing.expectEqual(u16, tag.MakeUnsigned(BFloat16));
292 try std.testing.expectEqual(i16, tag.MakeSigned(BFloat16));
293 try std.testing.expectEqual(f32, tag.MakeWide(BFloat16));
294 try std.testing.expect(tag.isSpecialFloat(BFloat16));
295 try std.testing.expect(!tag.isFloat(BFloat16));
296 try std.testing.expect(tag.isSigned(BFloat16));
297 try std.testing.expect(!tag.isUnsigned(BFloat16));
298 }
299
300 test "Highway bfloat16 vector promotion demotion and memory preserve bit patterns" {
301 const simd = @import("root.zig");
302 const D = Tag(4);
303 const F = tag.FixedTag(f32, 4);
304 const input: F.Vector = .{ 1, -2, 3.984375, -0.0 };
305 const bits = simd.demoteTo(D, input);
306 try std.testing.expect(@reduce(.And, simd.promoteTo(F, bits) == @as(F.Vector, .{ 1, -2, 3.984375, -0.0 })));
307 var scalar: [4]BFloat16 = undefined;
308 store(D, bits, &scalar);
309 try std.testing.expect(@reduce(.And, load(D, &scalar) == bits));
310 try std.testing.expect(@reduce(.And, neg(D, bits) == (bits ^ @as(D.Vector, @splat(0x8000)))));
311 }
312
313 test "Highway bfloat16 even odd and pairwise multiply-add widen to f32" {
314 const simd = @import("root.zig");
315 const D = tag.FixedTag(f32, 4);
316 const B = Tag(8);
317 const a = demoteF32(B, @as(tag.FixedTag(f32, 8).Vector, .{ 1, 2, 3, 4, 5, 6, 7, 8 }));
318 const b = demoteF32(B, @as(tag.FixedTag(f32, 8).Vector, .{ 8, 7, 6, 5, 4, 3, 2, 1 }));
319 try std.testing.expect(@reduce(.And, mulEvenAdd(D, a, b, @splat(1)) ==
320 @as(D.Vector, .{ 9, 19, 21, 15 })));
321 try std.testing.expect(@reduce(.And, mulOddAdd(D, a, b, @splat(1)) ==
322 @as(D.Vector, .{ 15, 21, 19, 9 })));
323 try std.testing.expect(@reduce(.And, simd.widenMulPairwiseAdd(D, a, b) ==
324 @as(D.Vector, .{ 22, 38, 38, 22 })));
325 var odd: D.Vector = @splat(0);
326 const even = simd.reorderWidenMulAccumulate(D, a, b, @splat(0), &odd);
327 try std.testing.expect(@reduce(.And, simd.rearrangeToOddPlusEven(D, even, odd) ==
328 @as(D.Vector, .{ 22, 38, 38, 22 })));
329 }