lib/simd/src/tag.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const maximum_vector_bytes: usize = 65_536;
4
5 pub fn FixedTag(comptime T: type, comptime lanes: usize) type {
6 return Descriptor(T, lanes, 0);
7 }
8
9 pub fn Descriptor(comptime T: type, comptime lanes: usize, comptime power: i8) type {
10 @setEvalBranchQuota(2_000);
11 validateLane(T);
12 if (power < -3) @compileError("Highway descriptor power must be at least -3");
13 if (lanes == 0) @compileError("SIMD vectors require at least one lane");
14 if (!std.math.isPowerOfTwo(lanes)) {
15 @compileError("fixed SIMD lane counts must be powers of two");
16 }
17 if (lanes > maximum_vector_bytes / @sizeOf(T)) {
18 @compileError("fixed SIMD vector exceeds the Highway maximum byte width");
19 }
20 const Storage = if (isBFloat16(T)) T.Storage else T;
21
22 return struct {
23 pub const Lane = Storage;
24 pub const Scalar = T;
25 pub const lane_count: usize = lanes;
26 pub const byte_count: usize = lanes * @sizeOf(T);
27 pub const descriptor_power: i8 = power;
28 pub const is_bfloat16 = isBFloat16(T);
29 pub const Vector = @Vector(lanes, Storage);
30 pub const Mask = @Vector(lanes, bool);
31
32 pub fn rebind(comptime NewLane: type) type {
33 return Descriptor(
34 NewLane,
35 lanes,
36 power + laneSizePower(NewLane) - laneSizePower(T),
37 );
38 }
39
40 pub fn repartition(comptime NewLane: type) type {
41 validateLane(NewLane);
42 const new_lanes = @max(1, (byte_count + @sizeOf(NewLane) - 1) /
43 @sizeOf(NewLane));
44 return Descriptor(NewLane, new_lanes, power);
45 }
46
47 pub fn half() type {
48 return Descriptor(T, @max(1, lanes / 2), power - 1);
49 }
50
51 pub fn twice() type {
52 return Descriptor(T, lanes * 2, power + 1);
53 }
54
55 pub fn withLanes(comptime new_lanes: usize, comptime new_power: i8) type {
56 return Descriptor(T, new_lanes, new_power);
57 }
58 };
59 }
60
61 pub fn ScalableTag(comptime T: type) type {
62 validateLane(T);
63 return FixedTag(T, suggestedLanes(T));
64 }
65
66 pub fn ScalableTagPow2(comptime T: type, comptime power: comptime_int) type {
67 validateLane(T);
68 const descriptor_power = clampedPower(power);
69 return Descriptor(T, scaleLanes(suggestedLanes(T), descriptor_power), descriptor_power);
70 }
71
72 pub fn CappedTag(comptime T: type, comptime limit: usize) type {
73 validateLane(T);
74 if (limit == 0) @compileError("capped SIMD vectors require a nonzero limit");
75 const native_lanes = suggestedLanes(T);
76 const lanes = std.math.floorPowerOfTwo(usize, @min(limit, native_lanes));
77 return FixedTag(T, @max(1, lanes));
78 }
79
80 pub fn CappedTagPow2(
81 comptime T: type,
82 comptime limit: usize,
83 comptime power: comptime_int,
84 ) type {
85 const descriptor_power = clampedPower(power);
86 return Descriptor(
87 T,
88 scaleLanes(CappedTag(T, limit).lane_count, descriptor_power),
89 descriptor_power,
90 );
91 }
92
93 pub fn CappedTagIfFixed(comptime T: type, comptime limit: usize) type {
94 return CappedTag(T, limit);
95 }
96
97 pub fn CappedTagIfFixedPow2(
98 comptime T: type,
99 comptime limit: usize,
100 comptime power: comptime_int,
101 ) type {
102 return CappedTagPow2(T, limit, power);
103 }
104
105 pub fn Full16(comptime T: type) type {
106 return FullBytes(T, 2);
107 }
108
109 pub fn Full32(comptime T: type) type {
110 return FullBytes(T, 4);
111 }
112
113 pub fn Full64(comptime T: type) type {
114 return FullBytes(T, 8);
115 }
116
117 pub fn Full128(comptime T: type) type {
118 return FullBytes(T, 16);
119 }
120
121 pub fn TFromD(comptime D: type) type {
122 return if (@hasDecl(D, "Scalar")) D.Scalar else D.Lane;
123 }
124
125 pub fn DFromV(comptime V: type) type {
126 return switch (@typeInfo(V)) {
127 .vector => |info| FixedTag(info.child, info.len),
128 else => @compileError("DFromV requires a Zig vector type"),
129 };
130 }
131
132 pub fn TFromV(comptime V: type) type {
133 return switch (@typeInfo(V)) {
134 .vector => |info| info.child,
135 else => @compileError("TFromV requires a Zig vector type"),
136 };
137 }
138
139 pub fn VFromD(comptime D: type) type {
140 return D.Vector;
141 }
142
143 pub fn MFromD(comptime D: type) type {
144 return D.Mask;
145 }
146
147 pub fn Vec(comptime D: type) type {
148 return VFromD(D);
149 }
150
151 pub fn Mask(comptime D: type) type {
152 return MFromD(D);
153 }
154
155 pub fn maxLanes(comptime D: type) usize {
156 return D.lane_count;
157 }
158
159 pub fn laneCount(comptime D: type) usize {
160 return D.lane_count;
161 }
162
163 pub fn maxBytes(comptime D: type) usize {
164 return D.byte_count;
165 }
166
167 pub fn maxBlocks(comptime D: type) usize {
168 return (D.byte_count + 15) / 16;
169 }
170
171 pub fn pow2(comptime D: type) i8 {
172 return D.descriptor_power;
173 }
174
175 pub fn Rebind(comptime T: type, comptime D: type) type {
176 return D.rebind(T);
177 }
178
179 pub fn RebindToSigned(comptime D: type) type {
180 return D.rebind(MakeSigned(TFromD(D)));
181 }
182
183 pub fn RebindToUnsigned(comptime D: type) type {
184 return D.rebind(MakeUnsigned(TFromD(D)));
185 }
186
187 pub fn RebindToFloat(comptime D: type) type {
188 return D.rebind(MakeFloat(TFromD(D)));
189 }
190
191 pub fn Repartition(comptime T: type, comptime D: type) type {
192 return D.repartition(T);
193 }
194
195 pub fn RepartitionToWide(comptime D: type) type {
196 return D.repartition(MakeWide(TFromD(D)));
197 }
198
199 pub fn RepartitionToNarrow(comptime D: type) type {
200 return D.repartition(MakeNarrow(TFromD(D)));
201 }
202
203 pub fn RepartitionToWideX2(comptime D: type) type {
204 return RepartitionToWide(RepartitionToWide(D));
205 }
206
207 pub fn RepartitionToWideX3(comptime D: type) type {
208 return RepartitionToWide(RepartitionToWideX2(D));
209 }
210
211 pub fn Half(comptime D: type) type {
212 return D.half();
213 }
214
215 pub fn Twice(comptime D: type) type {
216 return D.twice();
217 }
218
219 pub fn BlockDFromD(comptime D: type) type {
220 const block_lanes = @max(1, 16 / @sizeOf(TFromD(D)));
221 return D.withLanes(@min(D.lane_count, block_lanes), 0);
222 }
223
224 pub fn MakeUnsigned(comptime T: type) type {
225 if (isBFloat16(T)) return u16;
226 return switch (@typeInfo(T)) {
227 .int => |info| switch (info.bits) {
228 8, 16, 32, 64, 128 => @Int(.unsigned, info.bits),
229 else => @compileError("unsupported Highway integer width"),
230 },
231 .float => |info| switch (info.bits) {
232 16 => u16,
233 32 => u32,
234 64 => u64,
235 else => @compileError("unsupported Highway floating-point width"),
236 },
237 else => @compileError("type has no Highway unsigned relation"),
238 };
239 }
240
241 pub fn MakeSigned(comptime T: type) type {
242 if (isBFloat16(T)) return i16;
243 return switch (@typeInfo(T)) {
244 .int => |info| switch (info.bits) {
245 8, 16, 32, 64 => @Int(.signed, info.bits),
246 else => @compileError("type has no Highway signed relation"),
247 },
248 .float => |info| switch (info.bits) {
249 16 => i16,
250 32 => i32,
251 64 => i64,
252 else => @compileError("unsupported Highway floating-point width"),
253 },
254 else => @compileError("type has no Highway signed relation"),
255 };
256 }
257
258 pub fn MakeFloat(comptime T: type) type {
259 if (isBFloat16(T)) @compileError("type has no Highway floating-point relation");
260 return switch (@typeInfo(T)) {
261 .int => |info| switch (info.bits) {
262 16 => f16,
263 32 => f32,
264 64 => f64,
265 else => @compileError("type has no Highway floating-point relation"),
266 },
267 .float => |info| switch (info.bits) {
268 16 => f16,
269 32 => f32,
270 64 => f64,
271 else => @compileError("type has no Highway floating-point relation"),
272 },
273 else => @compileError("type has no Highway floating-point relation"),
274 };
275 }
276
277 pub fn MakeWide(comptime T: type) type {
278 if (isBFloat16(T)) return f32;
279 return switch (@typeInfo(T)) {
280 .int => |info| switch (info.bits) {
281 8 => @Int(info.signedness, 16),
282 16 => @Int(info.signedness, 32),
283 32 => @Int(info.signedness, 64),
284 64 => if (info.signedness == .unsigned)
285 u128
286 else
287 @compileError("type has no Highway wide relation"),
288 else => @compileError("type has no Highway wide relation"),
289 },
290 .float => |info| switch (info.bits) {
291 16 => f32,
292 32 => f64,
293 else => @compileError("type has no Highway wide relation"),
294 },
295 else => @compileError("type has no Highway wide relation"),
296 };
297 }
298
299 pub fn MakeNarrow(comptime T: type) type {
300 return switch (@typeInfo(T)) {
301 .int => |info| switch (info.bits) {
302 16 => @Int(info.signedness, 8),
303 32 => @Int(info.signedness, 16),
304 64 => @Int(info.signedness, 32),
305 128 => if (info.signedness == .unsigned)
306 u64
307 else
308 @compileError("type has no Highway narrow relation"),
309 else => @compileError("type has no Highway narrow relation"),
310 },
311 .float => |info| switch (info.bits) {
312 32 => f16,
313 64 => f32,
314 else => @compileError("type has no Highway narrow relation"),
315 },
316 else => @compileError("type has no Highway narrow relation"),
317 };
318 }
319
320 pub fn UnsignedFromSize(comptime bytes: usize) type {
321 return switch (bytes) {
322 1 => u8,
323 2 => u16,
324 4 => u32,
325 8 => u64,
326 16 => u128,
327 else => @compileError("unsupported Highway unsigned byte width"),
328 };
329 }
330
331 pub fn SignedFromSize(comptime bytes: usize) type {
332 return switch (bytes) {
333 1 => i8,
334 2 => i16,
335 4 => i32,
336 8 => i64,
337 else => @compileError("unsupported Highway signed byte width"),
338 };
339 }
340
341 pub fn FloatFromSize(comptime bytes: usize) type {
342 return switch (bytes) {
343 2 => f16,
344 4 => f32,
345 8 => f64,
346 else => @compileError("unsupported Highway floating-point byte width"),
347 };
348 }
349
350 pub fn isLane(comptime T: type) bool {
351 if (isBFloat16(T)) return true;
352 return switch (@typeInfo(T)) {
353 .int => |info| info.bits == 8 or info.bits == 16 or
354 info.bits == 32 or info.bits == 64,
355 .float => |info| info.bits == 16 or info.bits == 32 or info.bits == 64,
356 else => false,
357 };
358 }
359
360 pub fn isIntegerLane(comptime T: type) bool {
361 return switch (@typeInfo(T)) {
362 .int => |info| info.bits == 8 or info.bits == 16 or
363 info.bits == 32 or info.bits == 64,
364 else => false,
365 };
366 }
367
368 pub fn isSpecialFloat(comptime T: type) bool {
369 return T == f16 or isBFloat16(T);
370 }
371
372 pub fn isFloat3264(comptime T: type) bool {
373 return T == f32 or T == f64;
374 }
375
376 pub fn isFloat(comptime T: type) bool {
377 return T == f16 or isFloat3264(T);
378 }
379
380 pub fn isSigned(comptime T: type) bool {
381 if (isBFloat16(T)) return true;
382 return switch (@typeInfo(T)) {
383 .int => |info| info.signedness == .signed,
384 .float => true,
385 else => false,
386 };
387 }
388
389 pub fn isUnsigned(comptime T: type) bool {
390 return switch (@typeInfo(T)) {
391 .int => |info| info.signedness == .unsigned,
392 else => false,
393 };
394 }
395
396 fn FullBytes(comptime T: type, comptime bytes: usize) type {
397 validateLane(T);
398 if (@sizeOf(T) > bytes) @compileError("lane type exceeds fixed Highway vector width");
399 return FixedTag(T, bytes / @sizeOf(T));
400 }
401
402 fn scaleLanes(comptime lanes: usize, comptime power: i8) usize {
403 if (power < 0) return @max(1, lanes >> @intCast(-power));
404 return lanes << @intCast(power);
405 }
406
407 fn clampedPower(comptime power: comptime_int) i8 {
408 if (power < -3) @compileError("Highway scalable power must be at least -3");
409 return @intCast(@min(power, 3));
410 }
411
412 fn suggestedLanes(comptime T: type) usize {
413 const Storage = if (isBFloat16(T)) T.Storage else T;
414 return std.simd.suggestVectorLength(Storage) orelse 1;
415 }
416
417 fn laneSizePower(comptime T: type) i8 {
418 validateLane(T);
419 return switch (@sizeOf(T)) {
420 1 => 0,
421 2 => 1,
422 4 => 2,
423 8 => 3,
424 else => @compileError("unsupported Highway lane size"),
425 };
426 }
427
428 fn isBFloat16(comptime T: type) bool {
429 return switch (@typeInfo(T)) {
430 .@"struct" => @hasDecl(T, "is_bfloat16") and T.is_bfloat16,
431 else => false,
432 };
433 }
434
435 fn validateLane(comptime T: type) void {
436 if (!isLane(T)) {
437 @compileError("SIMD lanes must be 8/16/32/64-bit integers or 16/32/64-bit floats");
438 }
439 }
440
441 test "fixed tags preserve lane and byte invariants" {
442 const D = FixedTag(u16, 8);
443 try std.testing.expectEqual(@as(usize, 8), D.lane_count);
444 try std.testing.expectEqual(@as(usize, 16), D.byte_count);
445 try std.testing.expectEqual(@as(i8, 0), D.descriptor_power);
446 try std.testing.expectEqual(@Vector(8, u16), D.Vector);
447 try std.testing.expectEqual(@Vector(8, bool), D.Mask);
448 }
449
450 test "tag transformations match Highway descriptor meanings" {
451 const D = FixedTag(u16, 8);
452 try std.testing.expectEqual(@as(usize, 8), D.rebind(i32).lane_count);
453 try std.testing.expectEqual(@as(i8, 1), D.rebind(i32).descriptor_power);
454 try std.testing.expectEqual(FixedTag(u32, 4), D.repartition(u32));
455 try std.testing.expectEqual(@as(usize, 4), D.half().lane_count);
456 try std.testing.expectEqual(@as(i8, -1), D.half().descriptor_power);
457 try std.testing.expectEqual(@as(usize, 16), D.twice().lane_count);
458 try std.testing.expectEqual(@as(i8, 1), D.twice().descriptor_power);
459 }
460
461 test "capped tags round down and full tags retain 128 bits" {
462 const native = std.simd.suggestVectorLength(u8) orelse 1;
463 const expected = std.math.floorPowerOfTwo(usize, @min(13, native));
464 try std.testing.expectEqual(@max(1, expected), CappedTag(u8, 13).lane_count);
465 try std.testing.expectEqual(@as(usize, 16), Full128(u8).lane_count);
466 try std.testing.expectEqual(@as(usize, 2), Full128(f64).lane_count);
467 }
468
469 test "Highway descriptor aliases preserve fixed widths powers and transforms" {
470 const native = std.simd.suggestVectorLength(u32) orelse 1;
471 try std.testing.expectEqual(native * 2, ScalableTagPow2(u32, 1).lane_count);
472 try std.testing.expectEqual(@max(1, native / 2), ScalableTagPow2(u32, -1).lane_count);
473 try std.testing.expectEqual(native * 8, ScalableTagPow2(u32, 99).lane_count);
474 try std.testing.expectEqual(@as(i8, 3), pow2(ScalableTagPow2(u32, 99)));
475 try std.testing.expectEqual(
476 CappedTag(u32, 3).lane_count * 4,
477 CappedTagPow2(u32, 3, 2).lane_count,
478 );
479 try std.testing.expectEqual(CappedTag(u32, 3), CappedTagIfFixed(u32, 3));
480 try std.testing.expectEqual(@as(usize, 2), Full16(u8).lane_count);
481 try std.testing.expectEqual(@as(usize, 2), Full32(u16).lane_count);
482 try std.testing.expectEqual(@as(usize, 2), Full64(u32).lane_count);
483
484 const D = FixedTag(i16, 8);
485 try std.testing.expectEqual(i16, TFromD(D));
486 try std.testing.expectEqual(D, DFromV(D.Vector));
487 try std.testing.expectEqual(i16, TFromV(D.Vector));
488 try std.testing.expectEqual(D.Vector, VFromD(D));
489 try std.testing.expectEqual(D.Mask, MFromD(D));
490 try std.testing.expectEqual(D.Vector, Vec(D));
491 try std.testing.expectEqual(D.Mask, Mask(D));
492 try std.testing.expectEqual(D.lane_count, maxLanes(D));
493 try std.testing.expectEqual(D.lane_count, laneCount(D));
494 try std.testing.expectEqual(D.byte_count, maxBytes(D));
495 try std.testing.expectEqual(@as(usize, 1), maxBlocks(D));
496 try std.testing.expectEqual(@as(i8, 0), pow2(D));
497 try std.testing.expectEqual(FixedTag(u16, 8), RebindToUnsigned(D));
498 try std.testing.expectEqual(FixedTag(f16, 8), RebindToFloat(D));
499 try std.testing.expectEqual(FixedTag(i32, 4), RepartitionToWide(D));
500 try std.testing.expectEqual(FixedTag(i8, 16), RepartitionToNarrow(D));
501 try std.testing.expectEqual(@as(usize, 4), Half(D).lane_count);
502 try std.testing.expectEqual(@as(i8, -1), pow2(Half(D)));
503 try std.testing.expectEqual(@as(usize, 16), Twice(D).lane_count);
504 try std.testing.expectEqual(@as(i8, 1), pow2(Twice(D)));
505 try std.testing.expectEqual(D, BlockDFromD(FixedTag(i16, 16)));
506 }
507
508 test "Highway lane type relations cover every public scalar family" {
509 try std.testing.expectEqual(u32, MakeUnsigned(f32));
510 try std.testing.expectEqual(i64, MakeSigned(u64));
511 try std.testing.expectEqual(f16, MakeFloat(u16));
512 try std.testing.expectEqual(u128, MakeWide(u64));
513 try std.testing.expectEqual(i16, MakeNarrow(i32));
514 try std.testing.expectEqual(u128, UnsignedFromSize(16));
515 try std.testing.expectEqual(i64, SignedFromSize(8));
516 try std.testing.expectEqual(f64, FloatFromSize(8));
517 try std.testing.expect(isIntegerLane(i8));
518 try std.testing.expect(isSpecialFloat(f16));
519 try std.testing.expect(isFloat3264(f64));
520 try std.testing.expect(isFloat(f16));
521 try std.testing.expect(isSigned(f32));
522 try std.testing.expect(isUnsigned(u32));
523 }
524
525 test "current Highway descriptor differential digest" {
526 var digest: u64 = 0xcbf2_9ce4_8422_2325;
527 const D = FixedTag(i16, 8);
528 descriptorDigestRecord(D, &digest);
529 descriptorDigestRecord(Rebind(i32, D), &digest);
530 descriptorDigestRecord(Repartition(i32, D), &digest);
531 descriptorDigestRecord(Half(D), &digest);
532 descriptorDigestRecord(Twice(D), &digest);
533 descriptorDigestRecord(BlockDFromD(Twice(D)), &digest);
534 descriptorDigestRecord(Full16(u8), &digest);
535 descriptorDigestRecord(Full32(u16), &digest);
536 descriptorDigestRecord(Full64(u32), &digest);
537 descriptorDigestRecord(Full128(u64), &digest);
538 digest = descriptorDigestStep(
539 digest,
540 maxLanes(ScalableTagPow2(u32, 99)) / maxLanes(ScalableTag(u32)),
541 );
542 digest = descriptorDigestStep(digest, @bitCast(@as(i64, pow2(ScalableTagPow2(u32, 99)))));
543 try std.testing.expectEqual(@as(u64, 9_371_378_310_310_814_615), digest);
544 }
545
546 fn descriptorDigestRecord(comptime D: type, digest: *u64) void {
547 digest.* = descriptorDigestStep(digest.*, maxLanes(D));
548 digest.* = descriptorDigestStep(digest.*, maxBytes(D));
549 digest.* = descriptorDigestStep(digest.*, maxBlocks(D));
550 digest.* = descriptorDigestStep(digest.*, @bitCast(@as(i64, pow2(D))));
551 }
552
553 fn descriptorDigestStep(digest: u64, value: u64) u64 {
554 return (digest ^ value) *% 0x0000_0100_0000_01b3;
555 }